From b548ca1474a2ee957a838a649b958f81d7be8b2e Mon Sep 17 00:00:00 2001 From: AJ ONeal Date: Fri, 11 Sep 2026 17:26:27 -0600 Subject: [PATCH] feat: validate typed MCP tool arguments --- mcp/server20260728/server.go | 41 +++++++++++++++--------- mcp/server20260728/server_test.go | 46 +++++++++++++++++++++++++++ mcp/types20260728/types.go | 52 +++++++++++++++++++++++++++++++ mcp/types20260728/types_test.go | 41 ++++++++++++++++++++++++ 4 files changed, 165 insertions(+), 15 deletions(-) create mode 100644 mcp/types20260728/types_test.go diff --git a/mcp/server20260728/server.go b/mcp/server20260728/server.go index 8680186..310047c 100644 --- a/mcp/server20260728/server.go +++ b/mcp/server20260728/server.go @@ -8,6 +8,7 @@ package mcpserver import ( "context" "encoding/json" + "errors" "fmt" "io" "net/http" @@ -50,9 +51,9 @@ type ResourceHandlerFunc func(context.Context, ReadResourceRequest) (mcptypes.Re // RegisterTypedTool decodes a tool's arguments into T before calling handler. // T is the application Go type for the tool input; no JSON Schema validator is // involved. -func RegisterTypedTool[T any](s *MCPServer, definition mcptypes.Tool, handler func(context.Context, Headers, T) (mcptypes.CallToolResult, error)) { +func RegisterTypedTool[T mcptypes.Validatable](s *MCPServer, definition mcptypes.Tool, handler func(context.Context, Headers, T) (mcptypes.CallToolResult, error)) { s.RegisterTool(definition, func(ctx context.Context, request CallToolRequest) (mcptypes.CallToolResult, error) { - args, err := mcptypes.Decode[T](request.Params.Arguments) + args, err := mcptypes.DecodeValidated[T](request.Params.Arguments) if err != nil { return mcptypes.CallToolResult{}, fmt.Errorf("decode %s arguments: %w", definition.Name, err) } @@ -69,6 +70,7 @@ type MCPServer struct { AllowedOrigins []string mu sync.RWMutex tools []registeredTool + toolByName map[string]registeredTool resources []registeredResource } type registeredTool struct { @@ -88,7 +90,12 @@ func (s *MCPServer) RegisterTool(definition mcptypes.Tool, handler ToolHandlerFu panic("mcp: duplicate tool: " + definition.Name) } } - s.tools = append(s.tools, registeredTool{definition, handler}) + if s.toolByName == nil { + s.toolByName = make(map[string]registeredTool) + } + registered := registeredTool{definition, handler} + s.tools = append(s.tools, registered) + s.toolByName[definition.Name] = registered } func (s *MCPServer) RegisterResource(definition mcptypes.Resource, handler ResourceHandlerFunc) { s.mu.Lock() @@ -145,6 +152,9 @@ func (s *MCPServer) handle(ctx context.Context, headers Headers, body []byte) an if headers.Method != envelope.Method { return s.error(envelope.ID, mcptypes.ErrHeaderMismatch, "Mcp-Method does not match request method") } + if envelope.Method == "tools/call" && headers.Name == "" { + return s.error(envelope.ID, mcptypes.ErrHeaderMismatch, "missing Mcp-Name header") + } var request mcptypes.Request if err := json.Unmarshal(body, &request); err != nil || request.JSONRPC != mcptypes.JSONRPCVersion || request.Method == "" || (request.IDPresent && (request.ID == nil || request.ID.IsNull())) { return s.error(envelope.ID, mcptypes.ErrInvalidRequest, "Invalid JSON-RPC request") @@ -202,24 +212,25 @@ func (s *MCPServer) toolsList(id *mcptypes.RequestID) any { return s.result(id, raw) } func (s *MCPServer) toolsCall(ctx context.Context, headers Headers, id *mcptypes.RequestID, raw json.RawMessage) any { + // Select the handler from Mcp-Name before decoding tool parameters. + s.mu.RLock() + registered, found := s.toolByName[headers.Name] + s.mu.RUnlock() + if !found { + return s.error(id, mcptypes.ErrInvalidParams, "Unknown tool") + } params, err := mcptypes.Decode[mcptypes.CallToolParams](raw) if err != nil || params.Name == "" { return s.error(id, mcptypes.ErrInvalidParams, "Invalid params for tools/call") } - s.mu.RLock() - var handler ToolHandlerFunc - for _, tool := range s.tools { - if tool.definition.Name == params.Name { - handler = tool.handler - break - } + if params.Name != headers.Name { + return s.error(id, mcptypes.ErrHeaderMismatch, "Mcp-Name does not match params.name") } - s.mu.RUnlock() - if handler == nil { - return s.error(id, mcptypes.ErrInvalidParams, "Unknown tool") - } - value, err := handler(ctx, CallToolRequest{headers, params}) + value, err := registered.handler(ctx, CallToolRequest{headers, params}) if err != nil { + if errors.Is(err, mcptypes.ErrInvalidArguments) { + return s.error(id, mcptypes.ErrInvalidParams, err.Error()) + } return s.error(id, mcptypes.ErrInternal, "Tool execution failed") } encoded, encodeErr := json.Marshal(value) diff --git a/mcp/server20260728/server_test.go b/mcp/server20260728/server_test.go index 115d4b2..01d7e74 100644 --- a/mcp/server20260728/server_test.go +++ b/mcp/server20260728/server_test.go @@ -3,6 +3,7 @@ package mcpserver import ( "context" "encoding/json" + "errors" "net/http" "net/http/httptest" "strings" @@ -15,6 +16,13 @@ type echoArgs struct { Message string `json:"message"` } +func (a echoArgs) Validate() error { + if a.Message == "" { + return errors.New("message is required") + } + return nil +} + func TestServeHTTPUsesHeadersBeforeTypedDecode(t *testing.T) { var got Headers var gotArgs echoArgs @@ -41,6 +49,44 @@ func TestServeHTTPUsesHeadersBeforeTypedDecode(t *testing.T) { } } +func TestServeHTTPRejectsInvalidTypedArguments(t *testing.T) { + srv := &MCPServer{} + RegisterTypedTool(srv, mcptypes.Tool{Name: "echo"}, 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":{}}}`)) + req.Header.Set("Mcp-Method", "tools/call") + req.Header.Set("Mcp-Name", "echo") + res := httptest.NewRecorder() + srv.ServeHTTP(res, req) + var response mcptypes.ErrorResponse + if err := json.NewDecoder(res.Body).Decode(&response); err != nil { + t.Fatal(err) + } + if response.Error.Code != mcptypes.ErrInvalidParams { + t.Fatalf("code = %d, want %d", response.Error.Code, mcptypes.ErrInvalidParams) + } +} + +func TestServeHTTPRejectsHeaderNameMismatch(t *testing.T) { + srv := &MCPServer{} + srv.RegisterTool(mcptypes.Tool{Name: "echo"}, func(context.Context, CallToolRequest) (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":"other"}}`)) + req.Header.Set("Mcp-Method", "tools/call") + req.Header.Set("Mcp-Name", "echo") + res := httptest.NewRecorder() + srv.ServeHTTP(res, req) + var response mcptypes.ErrorResponse + if err := json.NewDecoder(res.Body).Decode(&response); err != nil { + t.Fatal(err) + } + if response.Error.Code != mcptypes.ErrHeaderMismatch { + t.Fatalf("code = %d, want %d", response.Error.Code, mcptypes.ErrHeaderMismatch) + } +} + func TestServeHTTPRejectsHeaderMethodMismatch(t *testing.T) { srv := &MCPServer{} req := httptest.NewRequest(http.MethodPost, "/mcp", strings.NewReader(`{"jsonrpc":"2.0","id":1,"method":"tools/list"}`)) diff --git a/mcp/types20260728/types.go b/mcp/types20260728/types.go index a84046a..131ccff 100644 --- a/mcp/types20260728/types.go +++ b/mcp/types20260728/types.go @@ -8,6 +8,7 @@ package mcptypes import ( "bytes" "encoding/json" + "errors" "fmt" ) @@ -108,6 +109,8 @@ type ErrorObject struct { Data json.RawMessage `json:"data,omitempty"` } +var ErrInvalidArguments = errors.New("invalid tool arguments") + const ( ErrParse = -32700 ErrInvalidRequest = -32600 @@ -252,3 +255,52 @@ func Decode[T any](data json.RawMessage) (T, error) { err := json.Unmarshal(data, &value) return value, err } + +// Validatable is implemented by typed tool argument structs that enforce +// required fields and application-level constraints. +type Validatable interface { + Validate() error +} + +// DecodeValidated decodes typed tool arguments and runs their validation. +func DecodeValidated[T Validatable](data json.RawMessage) (T, error) { + value, err := Decode[T](data) + if err != nil { + return value, fmt.Errorf("%w: %v", ErrInvalidArguments, err) + } + if err := value.Validate(); err != nil { + return value, fmt.Errorf("%w: %v", ErrInvalidArguments, err) + } + return value, nil +} + +// NewToolResultText returns a successful text-only tool result. +func NewToolResultText(text string) CallToolResult { + return CallToolResult{ResultType: "complete", Content: []Content{{Type: "text", Text: text}}} +} + +// NewToolResultError returns an in-band tool execution error. +func NewToolResultError(message string) CallToolResult { + result := NewToolResultText(message) + result.IsError = true + return result +} + +// NewToolResultStructured encodes a typed Go value as structuredContent. +func NewToolResultStructured[T any](value T, text string) (CallToolResult, error) { + raw, err := json.Marshal(value) + if err != nil { + return CallToolResult{}, fmt.Errorf("encode structured tool result: %w", err) + } + return CallToolResult{ResultType: "complete", Content: []Content{{Type: "text", Text: text}}, StructuredContent: raw}, nil +} + +// NewToolResultStructuredOnly encodes a typed Go value and uses its JSON form +// as the text fallback for clients that do not support structuredContent. +func NewToolResultStructuredOnly[T any](value T) (CallToolResult, error) { + raw, err := json.Marshal(value) + if err != nil { + return CallToolResult{}, fmt.Errorf("encode structured tool result: %w", err) + } + return CallToolResult{ResultType: "complete", Content: []Content{{Type: "text", Text: string(raw)}}, StructuredContent: raw}, nil +} diff --git a/mcp/types20260728/types_test.go b/mcp/types20260728/types_test.go new file mode 100644 index 0000000..80ec3af --- /dev/null +++ b/mcp/types20260728/types_test.go @@ -0,0 +1,41 @@ +package mcptypes + +import ( + "encoding/json" + "testing" +) + +type typedResult struct { + Step string `json:"step"` + Count int `json:"count"` +} + +func TestNewToolResultStructuredUsesTypedValue(t *testing.T) { + result, err := NewToolResultStructured(typedResult{Step: "done", Count: 2}, "finished") + if err != nil { + t.Fatal(err) + } + var got typedResult + if err := json.Unmarshal(result.StructuredContent, &got); err != nil { + t.Fatal(err) + } + if got.Step != "done" || got.Count != 2 { + t.Fatalf("structured result = %#v", got) + } + if result.ResultType != "complete" || result.Content[0].Text != "finished" { + t.Fatalf("result = %#v", result) + } +} + +func TestNewToolResultStructuredOnlyUsesJSONFallback(t *testing.T) { + result, err := NewToolResultStructuredOnly(typedResult{Step: "done", Count: 2}) + if err != nil { + t.Fatal(err) + } + if result.Content[0].Text != `{"step":"done","count":2}` { + t.Fatalf("text fallback = %q", result.Content[0].Text) + } + if result.IsError { + t.Fatal("structured result marked as error") + } +}