feat: validate typed MCP tool arguments

This commit is contained in:
AJ ONeal 2026-09-11 17:26:27 -06:00
parent d032a5354e
commit b548ca1474
No known key found for this signature in database
4 changed files with 165 additions and 15 deletions

View File

@ -8,6 +8,7 @@ package mcpserver
import ( import (
"context" "context"
"encoding/json" "encoding/json"
"errors"
"fmt" "fmt"
"io" "io"
"net/http" "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. // 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 // T is the application Go type for the tool input; no JSON Schema validator is
// involved. // 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) { 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 { if err != nil {
return mcptypes.CallToolResult{}, fmt.Errorf("decode %s arguments: %w", definition.Name, err) return mcptypes.CallToolResult{}, fmt.Errorf("decode %s arguments: %w", definition.Name, err)
} }
@ -69,6 +70,7 @@ type MCPServer struct {
AllowedOrigins []string AllowedOrigins []string
mu sync.RWMutex mu sync.RWMutex
tools []registeredTool tools []registeredTool
toolByName map[string]registeredTool
resources []registeredResource resources []registeredResource
} }
type registeredTool struct { type registeredTool struct {
@ -88,7 +90,12 @@ func (s *MCPServer) RegisterTool(definition mcptypes.Tool, handler ToolHandlerFu
panic("mcp: duplicate tool: " + definition.Name) 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) { func (s *MCPServer) RegisterResource(definition mcptypes.Resource, handler ResourceHandlerFunc) {
s.mu.Lock() s.mu.Lock()
@ -145,6 +152,9 @@ func (s *MCPServer) handle(ctx context.Context, headers Headers, body []byte) an
if headers.Method != envelope.Method { if headers.Method != envelope.Method {
return s.error(envelope.ID, mcptypes.ErrHeaderMismatch, "Mcp-Method does not match request 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 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())) { 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") 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) return s.result(id, raw)
} }
func (s *MCPServer) toolsCall(ctx context.Context, headers Headers, id *mcptypes.RequestID, raw json.RawMessage) any { 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) params, err := mcptypes.Decode[mcptypes.CallToolParams](raw)
if err != nil || params.Name == "" { if err != nil || params.Name == "" {
return s.error(id, mcptypes.ErrInvalidParams, "Invalid params for tools/call") return s.error(id, mcptypes.ErrInvalidParams, "Invalid params for tools/call")
} }
s.mu.RLock() if params.Name != headers.Name {
var handler ToolHandlerFunc return s.error(id, mcptypes.ErrHeaderMismatch, "Mcp-Name does not match params.name")
for _, tool := range s.tools {
if tool.definition.Name == params.Name {
handler = tool.handler
break
} }
} value, err := registered.handler(ctx, CallToolRequest{headers, params})
s.mu.RUnlock()
if handler == nil {
return s.error(id, mcptypes.ErrInvalidParams, "Unknown tool")
}
value, err := handler(ctx, CallToolRequest{headers, params})
if err != nil { 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") return s.error(id, mcptypes.ErrInternal, "Tool execution failed")
} }
encoded, encodeErr := json.Marshal(value) encoded, encodeErr := json.Marshal(value)

View File

@ -3,6 +3,7 @@ package mcpserver
import ( import (
"context" "context"
"encoding/json" "encoding/json"
"errors"
"net/http" "net/http"
"net/http/httptest" "net/http/httptest"
"strings" "strings"
@ -15,6 +16,13 @@ type echoArgs struct {
Message string `json:"message"` 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) { func TestServeHTTPUsesHeadersBeforeTypedDecode(t *testing.T) {
var got Headers var got Headers
var gotArgs echoArgs 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) { func TestServeHTTPRejectsHeaderMethodMismatch(t *testing.T) {
srv := &MCPServer{} srv := &MCPServer{}
req := httptest.NewRequest(http.MethodPost, "/mcp", strings.NewReader(`{"jsonrpc":"2.0","id":1,"method":"tools/list"}`)) req := httptest.NewRequest(http.MethodPost, "/mcp", strings.NewReader(`{"jsonrpc":"2.0","id":1,"method":"tools/list"}`))

View File

@ -8,6 +8,7 @@ package mcptypes
import ( import (
"bytes" "bytes"
"encoding/json" "encoding/json"
"errors"
"fmt" "fmt"
) )
@ -108,6 +109,8 @@ type ErrorObject struct {
Data json.RawMessage `json:"data,omitempty"` Data json.RawMessage `json:"data,omitempty"`
} }
var ErrInvalidArguments = errors.New("invalid tool arguments")
const ( const (
ErrParse = -32700 ErrParse = -32700
ErrInvalidRequest = -32600 ErrInvalidRequest = -32600
@ -252,3 +255,52 @@ func Decode[T any](data json.RawMessage) (T, error) {
err := json.Unmarshal(data, &value) err := json.Unmarshal(data, &value)
return value, err 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
}

View File

@ -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")
}
}