fix(mcp): describe dynamic outputs with explicit schemas (#1626)

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)
+	}
+}

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

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 {