platform/mcp/server.go

// Package mcp is the shared Model Context Protocol server core: JSON-RPC 2.0
// over streamable HTTP, tool registration, and the tools/list + tools/call
// dispatch. It is transport/protocol plumbing only — the tools themselves
// (their domain logic) stay in each module. A module builds its []Tool and calls
// NewServer; nothing here knows about carts, content, or orders.
//
// See docs/platform-uniforming-plan.md (platform/mcp seam, migration done) and
// platform-event-messaging-sketch.md.
package mcp

import (
	"context"
	"encoding/json"
	"io"
	"net/http"
)

// defaultProtocolVersion is advertised when Info.ProtocolVersion is empty.
const defaultProtocolVersion = "2025-06-18"

const maxBodyBytes = 4 << 20 // 4 MiB request cap, matching the prior per-module servers.

// Info identifies the server in the MCP initialize handshake.
type Info struct {
	Name            string
	Version         string
	ProtocolVersion string // optional; defaults to defaultProtocolVersion
}

// Server is the MCP edge over a set of tools. Build it with NewServer and mount
// Handler() (e.g. at /mcp). The tool slice is read on every request, so AddTool
// is safe after Handler() is in use.
type Server struct {
	info  Info
	tools []Tool
}

// NewServer builds an MCP server advertising the given tools.
func NewServer(info Info, tools ...Tool) *Server {
	if info.ProtocolVersion == "" {
		info.ProtocolVersion = defaultProtocolVersion
	}
	return &Server{info: info, tools: tools}
}

// AddTool appends a tool so it is advertised via tools/list and dispatched via
// tools/call alongside the rest.
func (s *Server) AddTool(t Tool) { s.tools = append(s.tools, t) }

// Handler returns the streamable-HTTP handler to mount (e.g. at /mcp).
func (s *Server) Handler() http.Handler { return http.HandlerFunc(s.serveHTTP) }

func (s *Server) serveHTTP(w http.ResponseWriter, r *http.Request) {
	if r.Method != http.MethodPost {
		w.Header().Set("Allow", "POST")
		http.Error(w, "method not allowed", http.StatusMethodNotAllowed)
		return
	}
	body, err := io.ReadAll(io.LimitReader(r.Body, maxBodyBytes))
	if err != nil {
		writeRPC(w, errorResponse(nil, codeParseError, "read error"))
		return
	}
	ctx := r.Context()
	if isBatch(body) {
		var reqs []rpcRequest
		if err := json.Unmarshal(body, &reqs); err != nil {
			writeRPC(w, errorResponse(nil, codeParseError, "invalid JSON"))
			return
		}
		var out []*rpcResponse
		for i := range reqs {
			if resp := s.dispatch(ctx, &reqs[i]); resp != nil {
				out = append(out, resp)
			}
		}
		if len(out) == 0 {
			w.WriteHeader(http.StatusAccepted)
			return
		}
		writeRPC(w, out)
		return
	}
	var req rpcRequest
	if err := json.Unmarshal(body, &req); err != nil {
		writeRPC(w, errorResponse(nil, codeParseError, "invalid JSON"))
		return
	}
	resp := s.dispatch(ctx, &req)
	if resp == nil {
		w.WriteHeader(http.StatusAccepted)
		return
	}
	writeRPC(w, resp)
}

func (s *Server) dispatch(ctx context.Context, req *rpcRequest) *rpcResponse {
	if req.JSONRPC != "2.0" {
		if req.isNotification() {
			return nil
		}
		return errorResponse(req.ID, codeInvalidRequest, "jsonrpc must be \"2.0\"")
	}
	switch req.Method {
	case "initialize":
		return result(req.ID, s.initialize(req.Params))
	case "ping":
		return result(req.ID, struct{}{})
	case "tools/list":
		return result(req.ID, map[string]any{"tools": s.tools})
	case "tools/call":
		return s.callTool(ctx, req)
	case "notifications/initialized", "notifications/cancelled":
		return nil // notifications: acknowledged by the transport, no reply
	default:
		if req.isNotification() {
			return nil
		}
		return errorResponse(req.ID, codeMethodNotFound, "unknown method: "+req.Method)
	}
}

func (s *Server) initialize(params json.RawMessage) map[string]any {
	pv := s.info.ProtocolVersion
	if len(params) > 0 {
		var p struct {
			ProtocolVersion string `json:"protocolVersion"`
		}
		if json.Unmarshal(params, &p) == nil && p.ProtocolVersion != "" {
			pv = p.ProtocolVersion
		}
	}
	return map[string]any{
		"protocolVersion": pv,
		"capabilities":    map[string]any{"tools": map[string]any{}},
		"serverInfo":      map[string]any{"name": s.info.Name, "version": s.info.Version},
	}
}

func (s *Server) callTool(ctx context.Context, req *rpcRequest) *rpcResponse {
	var p struct {
		Name      string          `json:"name"`
		Arguments json.RawMessage `json:"arguments"`
	}
	if err := json.Unmarshal(req.Params, &p); err != nil {
		return errorResponse(req.ID, codeInvalidParams, "invalid params")
	}
	// Strip (and log) the UCP-Agent meta sibling before handing args to the tool,
	// so it never leaks into the tool's argument namespace. No-op when absent.
	args := extractUCPAgentMeta(p.Arguments)

	for i := range s.tools {
		if s.tools[i].Name == p.Name {
			out, err := s.tools[i].Invoke(ctx, args)
			if err != nil {
				return result(req.ID, ErrorResult(err))
			}
			return result(req.ID, Text(out))
		}
	}
	return errorResponse(req.ID, codeInvalidParams, "unknown tool: "+p.Name)
}

func isBatch(body []byte) bool {
	for _, b := range body {
		switch b {
		case ' ', '\t', '\r', '\n':
			continue
		case '[':
			return true
		default:
			return false
		}
	}
	return false
}

func writeRPC(w http.ResponseWriter, v any) {
	w.Header().Set("Content-Type", "application/json")
	w.WriteHeader(http.StatusOK)
	_ = json.NewEncoder(w).Encode(v)
}