From 2b8de26a7d2e70d7308f971a7f3ddea3188e3ce6 Mon Sep 17 00:00:00 2001 From: Levi Date: Wed, 28 Oct 2026 17:05:11 +0800 Subject: [PATCH] fix(mcp): describe dynamic outputs with explicit schemas --- pkg/mcpserver/output_schema_test.go | 148 ++++++++++++++++++++++++++++ pkg/mcpserver/server.go | 60 ++++++++--- pkg/mcpserver/server_test.go | 30 ++++++ 3 files changed, 225 insertions(+), 13 deletions(-) create mode 100644 pkg/mcpserver/output_schema_test.go diff --git a/pkg/mcpserver/output_schema_test.go b/pkg/mcpserver/output_schema_test.go new file mode 100644 index 000000000..7a073c8bb --- /dev/null +++ b/pkg/mcpserver/output_schema_test.go @@ -0,0 +1,148 @@ +package mcpserver + +import ( + "bytes" + "encoding/json" + "fmt" + "net/http" + "net/http/httptest" + "testing" + + "github.com/GopeedLab/gopeed/internal/fetcher" + "github.com/GopeedLab/gopeed/pkg/base" + "github.com/GopeedLab/gopeed/pkg/download" + "github.com/GopeedLab/gopeed/pkg/protocol/bt" + "github.com/GopeedLab/gopeed/pkg/protocol/ed2k" + phttp "github.com/GopeedLab/gopeed/pkg/protocol/http" + "github.com/google/jsonschema-go/jsonschema" +) + +func TestToolSchemaPropertiesAreObjects(t *testing.T) { + downloader := download.NewDownloader(&download.DownloaderConfig{Storage: download.NewMemStorage()}) + req := httptest.NewRequest(http.MethodPost, "http://localhost/mcp", bytes.NewBufferString( + `{"jsonrpc":"2.0","id":1,"method":"tools/list","params":{}}`, + )) + req.Header.Set("Content-Type", "application/json") + req.Header.Set("Accept", "application/json, text/event-stream") + recorder := httptest.NewRecorder() + NewHandler(downloader).ServeHTTP(recorder, req) + if recorder.Code != http.StatusOK { + t.Fatalf("tools/list status = %d: %s", recorder.Code, recorder.Body.String()) + } + var response struct { + Result struct { + Tools []struct { + Name string `json:"name"` + InputSchema map[string]any `json:"inputSchema"` + OutputSchema map[string]any `json:"outputSchema"` + } `json:"tools"` + } `json:"result"` + } + if err := json.Unmarshal(recorder.Body.Bytes(), &response); err != nil { + t.Fatal(err) + } + if len(response.Result.Tools) == 0 { + t.Fatal("tools/list returned no tools") + } + for _, tool := range response.Result.Tools { + t.Run(tool.Name, func(t *testing.T) { + checkSchemaProperties(t, tool.InputSchema, "inputSchema") + checkSchemaProperties(t, tool.OutputSchema, "outputSchema") + }) + } +} + +// Check the serialized definitions, including nested file requests and union +// branches. Boolean additionalProperties remains valid; boolean entries inside +// properties are what clients with object-only schema parsers cannot consume. +func checkSchemaProperties(t *testing.T, value any, path string) { + t.Helper() + switch value := value.(type) { + case map[string]any: + for key, child := range value { + if key == "properties" { + properties, ok := child.(map[string]any) + if !ok { + t.Fatalf("%s.properties = %#v, want object", path, child) + } + for name, property := range properties { + if _, ok := property.(map[string]any); !ok { + t.Errorf("%s.properties.%s = %#v, want schema object", path, name, property) + } + } + } + checkSchemaProperties(t, child, path+"."+key) + } + case []any: + for i, child := range value { + checkSchemaProperties(t, child, fmt.Sprintf("%s[%d]", path, i)) + } + } +} + +func TestToolOutputSchemasPreserveDynamicValues(t *testing.T) { + statsSchema := toolOutputSchema[getTaskStatsOutput]() + for _, test := range []struct { + name string + stats any + }{ + {"unavailable", nil}, + {"typed nil", (*fetcher.Stats)(nil)}, + {"http", &fetcher.Stats{Snapshot: &phttp.Stats{Connections: []*phttp.StatsConnection{{Downloaded: 123}}}}}, + {"bittorrent", &fetcher.Stats{Snapshot: &bt.StatsSnapshot{}, Runtime: &bt.StatsRuntime{}}}, + {"ed2k", &fetcher.Stats{Snapshot: &ed2k.StatsSnapshot{}, Runtime: &ed2k.StatsRuntime{}}}, + {"restored snapshot", &fetcher.Stats{Snapshot: json.RawMessage(`{"segmentsTotal":3,"segmentsDone":2,"segmentsFailed":0}`)}}, + {"empty object", map[string]any{}}, + {"nested JSON", map[string]any{"values": []any{nil, true, "text", 1.5, map[string]any{"key": "value"}}}}, + {"array", []any{1, "two", false}}, + {"string", "protocol value"}, + {"number", 1.5}, + {"boolean", true}, + } { + t.Run(test.name, func(t *testing.T) { + validateToolOutput(t, statsSchema, &getTaskStatsOutput{Stats: test.stats}) + }) + } + + // Output extras can be empty, null, or extension-specific. In particular, + // they need not satisfy the required settings in the download input schema. + for _, extra := range []any{nil, map[string]any{}, &phttp.ReqExtra{}, &bt.ReqExtra{}, + map[string]any{"custom": []any{true, nil, 42}}, "custom value"} { + req := &base.Request{URL: "https://example.com/file", Extra: extra} + res := &base.Resource{Files: []*base.FileInfo{{Name: "file", Req: req}}} + validateToolOutput(t, toolOutputSchema[resolveTaskOutput](), &resolveTaskOutput{Res: res}) + validateToolOutput(t, toolOutputSchema[getTaskOutput](), &getTaskOutput{Task: taskDetails{ + Req: req, Res: res, Opts: &base.Options{Extra: extra}, + }}) + } + + // Supplying an explicit schema must retain validation of the output envelope. + resolved, err := statsSchema.Resolve(nil) + if err != nil { + t.Fatal(err) + } + for _, invalid := range []map[string]any{{}, {"stats": nil, "unexpected": true}} { + if err := resolved.Validate(invalid); err == nil { + t.Errorf("accepted invalid output envelope: %#v", invalid) + } + } +} + +func validateToolOutput(t *testing.T, schema *jsonschema.Schema, output any) { + t.Helper() + data, err := json.Marshal(output) + if err != nil { + t.Fatal(err) + } + var value any + if err := json.Unmarshal(data, &value); err != nil { + t.Fatal(err) + } + resolved, err := schema.Resolve(nil) + if err != nil { + t.Fatal(err) + } + if err := resolved.Validate(value); err != nil { + t.Fatalf("output %s does not match schema: %v", data, err) + } +} diff --git a/pkg/mcpserver/server.go b/pkg/mcpserver/server.go index e14d57c6b..e67b1cceb 100644 --- a/pkg/mcpserver/server.go +++ b/pkg/mcpserver/server.go @@ -190,6 +190,37 @@ func downloadInputSchema[T any]() *jsonschema.Schema { return schema } +// toolOutputSchema preserves arbitrary JSON values in protocol statistics and +// task metadata. Unlike the boolean schema inferred for any, these object-form +// schemas can also be parsed by clients that do not support boolean schemas in +// properties. Output metadata may come from REST, extensions, or persisted tasks, +// so it must not inherit the narrower download input restrictions. +func toolOutputSchema[T any]() *jsonschema.Schema { + schema, err := jsonschema.For[T](&jsonschema.ForOptions{ + TypeSchemas: map[reflect.Type]*jsonschema.Schema{ + // Request.Labels is serialized as null when the map is nil. + reflect.TypeFor[map[string]string](): { + Types: []string{"object", "null"}, + AdditionalProperties: &jsonschema.Schema{Type: "string"}, + }, + reflect.TypeFor[any](): { + AnyOf: []*jsonschema.Schema{ + {Type: "object"}, + {Type: "array"}, + {Type: "string"}, + {Type: "number"}, + {Type: "boolean"}, + {Type: "null"}, + }, + }, + }, + }) + if err != nil { + panic(fmt.Sprintf("mcpserver: build tool output schema: %v", err)) + } + return schema +} + func (o *downloadOptions) baseOptions() *base.Options { if o == nil { return nil @@ -319,11 +350,12 @@ func registerTools(server *mcp.Server, downloader *download.Downloader) { destructive := true mcp.AddTool(server, &mcp.Tool{ - Name: "resolve_task", - Title: "Resolve task", - Description: "Resolve an HTTP, HTTPS, BitTorrent, magnet, or ED2K download request and return its resource metadata and files. Use this before create_task when the user wants to inspect the resource or select files.", - InputSchema: downloadInputSchema[resolveTaskInput](), - Annotations: &mcp.ToolAnnotations{ReadOnlyHint: true, OpenWorldHint: &openWorld}, + Name: "resolve_task", + Title: "Resolve task", + Description: "Resolve an HTTP, HTTPS, BitTorrent, magnet, or ED2K download request and return its resource metadata and files. Use this before create_task when the user wants to inspect the resource or select files.", + InputSchema: downloadInputSchema[resolveTaskInput](), + OutputSchema: toolOutputSchema[resolveTaskOutput](), + Annotations: &mcp.ToolAnnotations{ReadOnlyHint: true, OpenWorldHint: &openWorld}, }, func(_ context.Context, _ *mcp.CallToolRequest, input *resolveTaskInput) (*mcp.CallToolResult, *resolveTaskOutput, error) { if input.Req == nil { return nil, nil, fmt.Errorf("req is required") @@ -403,10 +435,11 @@ func registerTools(server *mcp.Server, downloader *download.Downloader) { }) mcp.AddTool(server, &mcp.Tool{ - Name: "get_task", - Title: "Get task", - Description: "Get one Gopeed task including its request, resource, options, and current progress.", - Annotations: &mcp.ToolAnnotations{ReadOnlyHint: true, OpenWorldHint: &closedWorld}, + Name: "get_task", + Title: "Get task", + Description: "Get one Gopeed task including its request, resource, options, and current progress.", + OutputSchema: toolOutputSchema[getTaskOutput](), + Annotations: &mcp.ToolAnnotations{ReadOnlyHint: true, OpenWorldHint: &closedWorld}, }, func(_ context.Context, _ *mcp.CallToolRequest, input *taskIDInput) (*mcp.CallToolResult, *getTaskOutput, error) { task, err := requireTask(downloader, input.ID) if err != nil { @@ -442,10 +475,11 @@ func registerTools(server *mcp.Server, downloader *download.Downloader) { }) mcp.AddTool(server, &mcp.Tool{ - Name: "get_task_stats", - Title: "Get task statistics", - Description: "Get protocol-specific statistics for one Gopeed task, such as HTTP connections or BitTorrent peers and seeding data.", - Annotations: &mcp.ToolAnnotations{ReadOnlyHint: true, OpenWorldHint: &closedWorld}, + Name: "get_task_stats", + Title: "Get task statistics", + Description: "Get protocol-specific statistics for one Gopeed task, such as HTTP connections or BitTorrent peers and seeding data.", + OutputSchema: toolOutputSchema[getTaskStatsOutput](), + Annotations: &mcp.ToolAnnotations{ReadOnlyHint: true, OpenWorldHint: &closedWorld}, }, func(_ context.Context, _ *mcp.CallToolRequest, input *taskIDInput) (*mcp.CallToolResult, *getTaskStatsOutput, error) { if err := validateTaskID(input.ID); err != nil { return nil, nil, err diff --git a/pkg/mcpserver/server_test.go b/pkg/mcpserver/server_test.go index 097fe841a..0a11e3e57 100644 --- a/pkg/mcpserver/server_test.go +++ b/pkg/mcpserver/server_test.go @@ -192,6 +192,36 @@ func TestStreamableHTTPTools(t *testing.T) { if tasks, ok := structured["tasks"].([]any); !ok || len(tasks) != 0 { t.Fatalf("tasks = %#v, want empty array", structured["tasks"]) } + + result, err = session.CallTool(context.Background(), &mcp.CallToolParams{ + Name: "create_task", + Arguments: map[string]any{ + "req": map[string]any{ + "url": downloadServer.URL, + "extra": map[string]any{"header": map[string]any{"X-Gopeed-Test": "mcp-object"}}, + }, + "opts": map[string]any{"path": t.TempDir()}, + }, + }) + if err != nil { + t.Fatal(err) + } + if result.IsError { + t.Fatalf("create_task returned tool error: %+v", result.Content) + } + created := result.StructuredContent.(map[string]any) + for _, name := range []string{"get_task", "get_task_stats"} { + result, err := session.CallTool(context.Background(), &mcp.CallToolParams{ + Name: name, + Arguments: map[string]any{"id": created["id"]}, + }) + if err != nil { + t.Fatalf("%s: %v", name, err) + } + if result.IsError { + t.Fatalf("%s returned tool error: %+v", name, result.Content) + } + } } func nestedSchemaProperty(t *testing.T, schema map[string]any, path ...string) map[string]any {