diff --git a/mcp/server20260728/mux.go b/mcp/server20260728/mux.go new file mode 100644 index 0000000..f7c114f --- /dev/null +++ b/mcp/server20260728/mux.go @@ -0,0 +1,100 @@ +package mcpserver + +import ( + "context" + "fmt" + "strings" + + mcptypes "github.com/therootcompany/golib/mcp/types20260728" +) + +// ToolExecution is a decoded tool invocation. Middleware may inspect the +// request headers and context, then call next to run the typed handler. +type ToolExecution func(context.Context, Headers) (mcptypes.CallToolResult, error) + +// ToolMiddleware wraps decoded tool execution. It is suitable for middleware +// such as authentication that does not need the tool's argument type. +type ToolMiddleware func(ToolExecution) ToolExecution + +// Registration is implemented by package-provided registration values such as +// ToolHandler[T]. Its method is intentionally private to keep registration +// behavior under package control. +type Registration interface { + register(*MCPServer, []ToolMiddleware) +} + +// ToolHandler pairs a tool definition with a validated Go argument type and +// typed handler. Use ToolHandler[T].With for middleware that needs T. +type ToolHandler[T mcptypes.Validatable] struct { + Tool mcptypes.Tool + Handle TypedToolHandler[T] + Middleware []TypedToolMiddleware[T] +} + +// With returns a copy with typed middleware appended. +func (h ToolHandler[T]) With(middleware ...TypedToolMiddleware[T]) ToolHandler[T] { + h.Middleware = append(append([]TypedToolMiddleware[T](nil), h.Middleware...), middleware...) + return h +} + +func (h ToolHandler[T]) toolName() string { return h.Tool.Name } + +func (h ToolHandler[T]) register(server *MCPServer, global []ToolMiddleware) { + typed := h.Handle + for i := len(h.Middleware) - 1; i >= 0; i-- { + typed = h.Middleware[i](typed) + } + server.RegisterTool(h.Tool, func(ctx context.Context, request CallToolRequest) (mcptypes.CallToolResult, error) { + args, err := mcptypes.DecodeValidated[T](request.Params.Arguments) + if err != nil { + return mcptypes.CallToolResult{}, fmt.Errorf("decode %s arguments: %w", h.Tool.Name, err) + } + execution := ToolExecution(func(ctx context.Context, headers Headers) (mcptypes.CallToolResult, error) { + return typed(ctx, headers, args) + }) + for i := len(global) - 1; i >= 0; i-- { + execution = global[i](execution) + } + return execution(ctx, request.Headers) + }) +} + +// Mux groups registrations on an MCPServer and creates middleware scopes. +type Mux struct { + server *MCPServer + middleware []ToolMiddleware +} + +func NewMux(server *MCPServer) *Mux { return &Mux{server: server} } + +// With returns a new mux with additional middleware appended. +func (m *Mux) With(middleware ...ToolMiddleware) *Mux { + combined := make([]ToolMiddleware, 0, len(m.middleware)+len(middleware)) + combined = append(combined, m.middleware...) + combined = append(combined, middleware...) + return &Mux{server: m.server, middleware: combined} +} + +// Register registers a complete typed tool handler. The method string is a +// wiring label and must be tools/, matching the name on the wire. +func (m *Mux) Register(method string, handler Registration) { + if !strings.HasPrefix(method, "tools/") { + panic("mcp: tool registered without tools/ prefix: " + method) + } + name := strings.TrimPrefix(method, "tools/") + if tool, ok := handler.(interface{ toolName() string }); ok && tool.toolName() != name { + panic(fmt.Sprintf("mcp: tool name mismatch: Register(%q) but Tool.Name is %q", method, tool.toolName())) + } + handler.register(m.server, m.middleware) +} + +// RegisterTool registers a raw handler for protocol-level use. Typed tools +// should use Register with ToolHandler[T]. +func (m *Mux) RegisterTool(definition mcptypes.Tool, handler ToolHandlerFunc) { + m.server.RegisterTool(definition, handler) +} + +// RegisterResource registers a resource read handler. +func (m *Mux) RegisterResource(definition mcptypes.Resource, handler ResourceHandlerFunc) { + m.server.RegisterResource(definition, handler) +} diff --git a/mcp/server20260728/server_test.go b/mcp/server20260728/server_test.go index f926d1a..4773a80 100644 --- a/mcp/server20260728/server_test.go +++ b/mcp/server20260728/server_test.go @@ -27,7 +27,7 @@ func TestServeHTTPUsesHeadersBeforeTypedDecode(t *testing.T) { var got Headers var gotArgs echoArgs srv := &MCPServer{Name: "test", Version: "1"} - RegisterTypedTool(srv, mcptypes.Tool{Name: "echo", InputSchema: json.RawMessage(`{"type":"object"}`)}, func(_ context.Context, headers Headers, args echoArgs) (mcptypes.CallToolResult, error) { + RegisterTypedTool(srv, mcptypes.Tool{Name: "echo", InputSchema: mcptypes.ToolInputSchema{Type: "object"}}, func(_ context.Context, headers Headers, args echoArgs) (mcptypes.CallToolResult, error) { got = headers gotArgs = args return mcptypes.CallToolResult{ResultType: "complete", Content: []mcptypes.Content{{Type: "text", Text: args.Message}}}, nil @@ -49,27 +49,27 @@ func TestServeHTTPUsesHeadersBeforeTypedDecode(t *testing.T) { } } -func TestTypedMiddlewareReceivesValidatedArguments(t *testing.T) { +func TestMuxWithRunsMiddlewareAfterTypedValidation(t *testing.T) { srv := &MCPServer{} called := false - middleware := func(next TypedToolHandler[echoArgs]) TypedToolHandler[echoArgs] { - return func(ctx context.Context, headers Headers, args echoArgs) (mcptypes.CallToolResult, error) { - if args.Message != "hello" { - t.Fatalf("middleware args = %#v", args) - } + auth := func(next ToolExecution) ToolExecution { + return func(ctx context.Context, headers Headers) (mcptypes.CallToolResult, error) { called = true - return next(ctx, headers, args) + return next(ctx, headers) } } - RegisterTypedToolWithMiddleware(srv, mcptypes.Tool{Name: "echo"}, func(context.Context, Headers, echoArgs) (mcptypes.CallToolResult, error) { - return mcptypes.NewToolResultText("ok"), nil - }, middleware) + NewMux(srv).With(auth).Register("tools/echo", ToolHandler[echoArgs]{ + Tool: mcptypes.Tool{Name: "echo"}, + Handle: func(context.Context, Headers, echoArgs) (mcptypes.CallToolResult, error) { + return mcptypes.NewToolResultText("ok"), nil + }, + }) req := httptest.NewRequest(http.MethodPost, "/mcp", strings.NewReader(`{"jsonrpc":"2.0","id":1,"method":"tools/call","params":{"name":"echo","arguments":{"message":"hello"}}}`)) req.Header.Set("Mcp-Method", "tools/call") req.Header.Set("Mcp-Name", "echo") srv.ServeHTTP(httptest.NewRecorder(), req) if !called { - t.Fatal("typed middleware was not called") + t.Fatal("middleware was not called") } } diff --git a/mcp/types20260728/types.go b/mcp/types20260728/types.go index 69985d2..bd376b1 100644 --- a/mcp/types20260728/types.go +++ b/mcp/types20260728/types.go @@ -330,10 +330,29 @@ type Tool struct { Name string `json:"name"` Title string `json:"title,omitempty"` Description string `json:"description,omitempty"` - InputSchema json.RawMessage `json:"inputSchema"` - OutputSchema json.RawMessage `json:"outputSchema,omitempty"` + InputSchema ToolInputSchema `json:"inputSchema"` + OutputSchema *ToolInputSchema `json:"outputSchema,omitempty"` Annotations *ToolAnnotations `json:"annotations,omitempty"` } + +// ToolInputSchema describes the JSON shape advertised to MCP clients. It is +// documentation for clients; DecodeValidated and the Go input type enforce +// server-side correctness. +type ToolInputSchema struct { + Type string `json:"type"` + Description string `json:"description,omitempty"` + Properties map[string]Property `json:"properties,omitempty"` + Required []string `json:"required,omitempty"` + Items *ToolInputSchema `json:"items,omitempty"` + AdditionalProperties *bool `json:"additionalProperties,omitempty"` +} + +type Property struct { + Type string `json:"type"` + Description string `json:"description,omitempty"` + Enum []string `json:"enum,omitempty"` + Items *ToolInputSchema `json:"items,omitempty"` +} type ToolAnnotations struct { Title string `json:"title,omitempty"` ReadOnlyHint *bool `json:"readOnlyHint,omitempty"`