diff --git a/mcp/server20260728/server.go b/mcp/server20260728/server.go index 310047c..5b50fbd 100644 --- a/mcp/server20260728/server.go +++ b/mcp/server20260728/server.go @@ -7,14 +7,19 @@ package mcpserver import ( "context" + "encoding/base64" "encoding/json" "errors" "fmt" "io" + "log" "net/http" + "net/url" + "runtime/debug" "slices" "strings" "sync" + "unicode" mcptypes "github.com/therootcompany/golib/mcp/types20260728" ) @@ -46,18 +51,30 @@ type ReadResourceRequest struct { Params mcptypes.ReadResourceParams } type ToolHandlerFunc func(context.Context, CallToolRequest) (mcptypes.CallToolResult, error) +type TypedToolHandler[T mcptypes.Validatable] func(context.Context, Headers, T) (mcptypes.CallToolResult, error) +type TypedToolMiddleware[T mcptypes.Validatable] func(TypedToolHandler[T]) TypedToolHandler[T] type ResourceHandlerFunc func(context.Context, ReadResourceRequest) (mcptypes.ReadResourceResult, error) // 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 mcptypes.Validatable](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 TypedToolHandler[T]) { + RegisterTypedToolWithMiddleware(s, definition, handler) +} + +// RegisterTypedToolWithMiddleware decodes and validates arguments before +// running middleware. Middleware receives the concrete T value. +func RegisterTypedToolWithMiddleware[T mcptypes.Validatable](s *MCPServer, definition mcptypes.Tool, handler TypedToolHandler[T], middleware ...TypedToolMiddleware[T]) { + wrapped := handler + for i := len(middleware) - 1; i >= 0; i-- { + wrapped = middleware[i](wrapped) + } s.RegisterTool(definition, 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", definition.Name, err) } - return handler(ctx, request.Headers, args) + return wrapped(ctx, request.Headers, args) }) } @@ -98,6 +115,9 @@ func (s *MCPServer) RegisterTool(definition mcptypes.Tool, handler ToolHandlerFu s.toolByName[definition.Name] = registered } func (s *MCPServer) RegisterResource(definition mcptypes.Resource, handler ResourceHandlerFunc) { + if err := validateResourceURI(definition.URI); err != nil { + panic(fmt.Sprintf("mcp: invalid resource URI %q: %v", definition.URI, err)) + } s.mu.Lock() defer s.mu.Unlock() for _, resource := range s.resources { @@ -186,6 +206,18 @@ func (s *MCPServer) discover(id *mcptypes.RequestID) any { } return s.result(id, raw) } +func isAppOnly(tool mcptypes.Tool) bool { + if tool.Meta == nil || tool.Meta.UI == nil || len(tool.Meta.UI.Visibility) == 0 { + return false + } + for _, visibility := range tool.Meta.UI.Visibility { + if visibility != mcptypes.VisibilityApp { + return false + } + } + return true +} + func (s *MCPServer) capabilities() mcptypes.ServerCapabilities { s.mu.RLock() defer s.mu.RUnlock() @@ -200,9 +232,12 @@ func (s *MCPServer) capabilities() mcptypes.ServerCapabilities { } func (s *MCPServer) toolsList(id *mcptypes.RequestID) any { s.mu.RLock() - tools := make([]mcptypes.Tool, len(s.tools)) - for i := range s.tools { - tools[i] = s.tools[i].definition + tools := make([]mcptypes.Tool, 0, len(s.tools)) + for _, registered := range s.tools { + if isAppOnly(registered.definition) { + continue + } + tools = append(tools, registered.definition) } s.mu.RUnlock() raw, err := json.Marshal(mcptypes.ToolsListResult{ResultType: "complete", Tools: tools, CacheScope: mcptypes.CachePrivate}) @@ -226,7 +261,7 @@ func (s *MCPServer) toolsCall(ctx context.Context, headers Headers, id *mcptypes if params.Name != headers.Name { return s.error(id, mcptypes.ErrHeaderMismatch, "Mcp-Name does not match params.name") } - value, err := registered.handler(ctx, CallToolRequest{headers, params}) + value, err := s.callToolWithRecovery(ctx, registered.handler, CallToolRequest{headers, params}) if err != nil { if errors.Is(err, mcptypes.ErrInvalidArguments) { return s.error(id, mcptypes.ErrInvalidParams, err.Error()) @@ -257,6 +292,9 @@ func (s *MCPServer) resourceRead(ctx context.Context, headers Headers, id *mcpty if err != nil || params.URI == "" { return s.error(id, mcptypes.ErrInvalidParams, "Invalid params for resources/read") } + if err := validateResourceURI(params.URI); err != nil { + return s.error(id, mcptypes.ErrInvalidParams, "Invalid resource URI") + } s.mu.RLock() var handler ResourceHandlerFunc for _, resource := range s.resources { @@ -273,6 +311,16 @@ func (s *MCPServer) resourceRead(ctx context.Context, headers Headers, id *mcpty if err != nil { return s.error(id, mcptypes.ErrInternal, "Resource read failed") } + for _, content := range value.Contents { + if err := validateResourceURI(content.URI); err != nil { + return s.error(id, mcptypes.ErrInternal, "Resource returned invalid URI") + } + if content.Blob != "" { + if _, err := base64.StdEncoding.DecodeString(content.Blob); err != nil { + return s.error(id, mcptypes.ErrInternal, "Resource returned invalid binary data") + } + } + } encoded, encodeErr := json.Marshal(value) if encodeErr != nil { return s.error(id, mcptypes.ErrInternal, "Failed to encode result") @@ -280,6 +328,28 @@ func (s *MCPServer) resourceRead(ctx context.Context, headers Headers, id *mcpty return s.result(id, encoded) } +func (s *MCPServer) callToolWithRecovery(ctx context.Context, handler ToolHandlerFunc, request CallToolRequest) (result mcptypes.CallToolResult, err error) { + defer func() { + if recovered := recover(); recovered != nil { + log.Printf("[mcp] panic in tool %q: %v\n%s", request.Params.Name, recovered, debug.Stack()) + result = mcptypes.NewToolResultError("Tool execution failed") + err = nil + } + }() + return handler(ctx, request) +} + +func validateResourceURI(raw string) error { + if raw == "" || strings.IndexFunc(raw, func(r rune) bool { return unicode.IsControl(r) || unicode.IsSpace(r) }) >= 0 { + return errors.New("URI is empty or contains invalid characters") + } + parsed, err := url.Parse(raw) + if err != nil || parsed.Scheme == "" || parsed.Fragment != "" { + return errors.New("URI must be absolute and must not contain a fragment") + } + return nil +} + func (s *MCPServer) result(id *mcptypes.RequestID, value json.RawMessage) mcptypes.Response { return mcptypes.Response{JSONRPC: mcptypes.JSONRPCVersion, ID: id, Result: value} } diff --git a/mcp/server20260728/server_test.go b/mcp/server20260728/server_test.go index 01d7e74..f926d1a 100644 --- a/mcp/server20260728/server_test.go +++ b/mcp/server20260728/server_test.go @@ -49,6 +49,51 @@ func TestServeHTTPUsesHeadersBeforeTypedDecode(t *testing.T) { } } +func TestTypedMiddlewareReceivesValidatedArguments(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) + } + called = true + return next(ctx, headers, args) + } + } + RegisterTypedToolWithMiddleware(srv, mcptypes.Tool{Name: "echo"}, func(context.Context, Headers, echoArgs) (mcptypes.CallToolResult, error) { + return mcptypes.NewToolResultText("ok"), nil + }, middleware) + 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") + } +} + +func TestServeHTTPContainsToolPanic(t *testing.T) { + srv := &MCPServer{} + srv.RegisterTool(mcptypes.Tool{Name: "panic"}, func(context.Context, CallToolRequest) (mcptypes.CallToolResult, error) { + panic("boom") + }) + req := httptest.NewRequest(http.MethodPost, "/mcp", strings.NewReader(`{"jsonrpc":"2.0","id":1,"method":"tools/call","params":{"name":"panic"}}`)) + req.Header.Set("Mcp-Method", "tools/call") + req.Header.Set("Mcp-Name", "panic") + res := httptest.NewRecorder() + srv.ServeHTTP(res, req) + var response struct { + Result mcptypes.CallToolResult `json:"result"` + } + if err := json.NewDecoder(res.Body).Decode(&response); err != nil { + t.Fatal(err) + } + if !response.Result.IsError { + t.Fatal("panic did not become an in-band tool error") + } +} + func TestServeHTTPRejectsInvalidTypedArguments(t *testing.T) { srv := &MCPServer{} RegisterTypedTool(srv, mcptypes.Tool{Name: "echo"}, func(context.Context, Headers, echoArgs) (mcptypes.CallToolResult, error) { @@ -68,6 +113,27 @@ func TestServeHTTPRejectsInvalidTypedArguments(t *testing.T) { } } +func TestRegisterResourceRejectsInvalidURI(t *testing.T) { + defer func() { + if recover() == nil { + t.Fatal("invalid resource URI did not panic") + } + }() + (&MCPServer{}).RegisterResource(mcptypes.Resource{URI: "not absolute"}, nil) +} + +func TestServeHTTPRejectsAppOnlyToolsFromList(t *testing.T) { + srv := &MCPServer{} + srv.RegisterTool(mcptypes.Tool{Name: "hidden", Meta: &mcptypes.Meta{UI: &mcptypes.AppUI{Visibility: []mcptypes.Visibility{mcptypes.VisibilityApp}}}}, nil) + req := httptest.NewRequest(http.MethodPost, "/mcp", strings.NewReader(`{"jsonrpc":"2.0","id":1,"method":"tools/list"}`)) + req.Header.Set("Mcp-Method", "tools/list") + res := httptest.NewRecorder() + srv.ServeHTTP(res, req) + if strings.Contains(res.Body.String(), "hidden") { + t.Fatalf("app-only tool leaked into list: %s", res.Body) + } +} + func TestServeHTTPRejectsHeaderNameMismatch(t *testing.T) { srv := &MCPServer{} srv.RegisterTool(mcptypes.Tool{Name: "echo"}, func(context.Context, CallToolRequest) (mcptypes.CallToolResult, error) { diff --git a/mcp/types20260728/types.go b/mcp/types20260728/types.go index 131ccff..69985d2 100644 --- a/mcp/types20260728/types.go +++ b/mcp/types20260728/types.go @@ -10,6 +10,7 @@ import ( "encoding/json" "errors" "fmt" + "maps" ) const ( @@ -88,9 +89,125 @@ type RequestMeta struct { } type ResultMeta struct { - ServerInfo *ServerInfo `json:"io.modelcontextprotocol/serverInfo,omitempty"` - SubscriptionID string `json:"io.modelcontextprotocol/subscriptionId,omitempty"` - Extensions map[string]json.RawMessage `json:"-"` + ServerInfo *ServerInfo + SubscriptionID string + UI *AppUI + Extensions map[string]json.RawMessage +} + +func (m ResultMeta) MarshalJSON() ([]byte, error) { + values := make(map[string]json.RawMessage, len(m.Extensions)+3) + maps.Copy(values, m.Extensions) + if m.ServerInfo != nil { + value, err := json.Marshal(m.ServerInfo) + if err != nil { + return nil, err + } + values["io.modelcontextprotocol/serverInfo"] = value + } + if m.SubscriptionID != "" { + value, err := json.Marshal(m.SubscriptionID) + if err != nil { + return nil, err + } + values["io.modelcontextprotocol/subscriptionId"] = value + } + if m.UI != nil { + value, err := json.Marshal(m.UI) + if err != nil { + return nil, err + } + values["ui"] = value + } + return json.Marshal(values) +} + +func (m *ResultMeta) UnmarshalJSON(data []byte) error { + var values map[string]json.RawMessage + if err := json.Unmarshal(data, &values); err != nil { + return err + } + m.Extensions = make(map[string]json.RawMessage, len(values)) + for key, value := range values { + switch key { + case "io.modelcontextprotocol/serverInfo": + if err := json.Unmarshal(value, &m.ServerInfo); err != nil { + return err + } + case "io.modelcontextprotocol/subscriptionId": + if err := json.Unmarshal(value, &m.SubscriptionID); err != nil { + return err + } + case "ui": + if err := json.Unmarshal(value, &m.UI); err != nil { + return err + } + default: + m.Extensions[key] = value + } + } + return nil +} + +// Meta contains typed MCP Apps metadata plus preserved extension fields. +type Meta struct { + UI *AppUI + Extensions map[string]json.RawMessage +} + +type AppUI struct { + ResourceURI string `json:"resourceUri,omitempty"` + Visibility []Visibility `json:"visibility,omitempty"` + CSP *AppCSP `json:"csp,omitempty"` + Permissions *AppPermissions `json:"permissions,omitempty"` +} +type Visibility string + +const ( + VisibilityModel Visibility = "model" + VisibilityApp Visibility = "app" +) + +type AppCSP struct { + ConnectDomains []string `json:"connectDomains,omitempty"` + ResourceDomains []string `json:"resourceDomains,omitempty"` + FrameDomains []string `json:"frameDomains,omitempty"` + BaseURIDomains []string `json:"baseUriDomains,omitempty"` +} +type AppPermission struct{} +type AppPermissions struct { + ClipboardWrite *AppPermission `json:"clipboardWrite,omitempty"` + Camera *AppPermission `json:"camera,omitempty"` + Microphone *AppPermission `json:"microphone,omitempty"` + Geolocation *AppPermission `json:"geolocation,omitempty"` +} + +func (m Meta) MarshalJSON() ([]byte, error) { + values := make(map[string]json.RawMessage, len(m.Extensions)+1) + maps.Copy(values, m.Extensions) + if m.UI != nil { + value, err := json.Marshal(m.UI) + if err != nil { + return nil, err + } + values["ui"] = value + } + return json.Marshal(values) +} +func (m *Meta) UnmarshalJSON(data []byte) error { + var values map[string]json.RawMessage + if err := json.Unmarshal(data, &values); err != nil { + return err + } + m.Extensions = make(map[string]json.RawMessage, len(values)) + if raw, ok := values["ui"]; ok { + if err := json.Unmarshal(raw, &m.UI); err != nil { + return err + } + delete(values, "ui") + } + maps.Copy(m.Extensions, values) + return nil } type Response struct { @@ -209,7 +326,7 @@ const ( ) type Tool struct { - Meta json.RawMessage `json:"_meta,omitempty"` + Meta *Meta `json:"_meta,omitempty"` Name string `json:"name"` Title string `json:"title,omitempty"` Description string `json:"description,omitempty"` @@ -225,20 +342,20 @@ type ToolAnnotations struct { OpenWorldHint *bool `json:"openWorldHint,omitempty"` } type Resource struct { - Meta json.RawMessage `json:"_meta,omitempty"` - URI string `json:"uri"` - Name string `json:"name"` - Title string `json:"title,omitempty"` - Description string `json:"description,omitempty"` - MIMEType string `json:"mimeType,omitempty"` - Size *int64 `json:"size,omitempty"` + Meta *Meta `json:"_meta,omitempty"` + URI string `json:"uri"` + Name string `json:"name"` + Title string `json:"title,omitempty"` + Description string `json:"description,omitempty"` + MIMEType string `json:"mimeType,omitempty"` + Size *int64 `json:"size,omitempty"` } type ResourceContents struct { - Meta json.RawMessage `json:"_meta,omitempty"` - URI string `json:"uri"` - MIMEType string `json:"mimeType,omitempty"` - Text string `json:"text,omitempty"` - Blob string `json:"blob,omitempty"` + Meta *Meta `json:"_meta,omitempty"` + URI string `json:"uri"` + MIMEType string `json:"mimeType,omitempty"` + Text string `json:"text,omitempty"` + Blob string `json:"blob,omitempty"` } type Content struct { Type string `json:"type"` diff --git a/mcp/types20260728/types_test.go b/mcp/types20260728/types_test.go index 80ec3af..fac1624 100644 --- a/mcp/types20260728/types_test.go +++ b/mcp/types20260728/types_test.go @@ -10,6 +10,34 @@ type typedResult struct { Count int `json:"count"` } +func TestAppsMetadataUsesEmptyPermissionObjects(t *testing.T) { + data, err := json.Marshal(Meta{UI: &AppUI{Permissions: &AppPermissions{ClipboardWrite: &AppPermission{}}}}) + if err != nil { + t.Fatal(err) + } + if string(data) != `{"ui":{"permissions":{"clipboardWrite":{}}}}` { + t.Fatalf("metadata = %s", data) + } +} + +func TestResultMetadataPreservesAppsUI(t *testing.T) { + original := ResultMeta{UI: &AppUI{ResourceURI: "ui://reports"}, Extensions: map[string]json.RawMessage{"example/trace": json.RawMessage(`"abc"`)}} + data, err := json.Marshal(original) + if err != nil { + t.Fatal(err) + } + var decoded ResultMeta + if err := json.Unmarshal(data, &decoded); err != nil { + t.Fatal(err) + } + if decoded.UI == nil || decoded.UI.ResourceURI != "ui://reports" { + t.Fatalf("UI = %#v", decoded.UI) + } + if string(decoded.Extensions["example/trace"]) != `"abc"` { + t.Fatalf("extensions = %#v", decoded.Extensions) + } +} + func TestNewToolResultStructuredUsesTypedValue(t *testing.T) { result, err := NewToolResultStructured(typedResult{Step: "done", Count: 2}, "finished") if err != nil {