From 6309674eefbd234c1ef8827c7d8199813cb3442a Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Fri, 9 Oct 2026 11:43:55 +0000 Subject: [PATCH] fix(mcp): preserve complete tool input schemas The restricted schema conversion drops defaults, compositions, and root keywords. Preserve the official MCP SDK schema as raw JSON instead of rebuilding it through LangChainGo definitions. Keep the deprecated normalization helper for existing callers, but remove it from discovery. Replace the invented-properties test with lossless discovery regressions. Check OpenAI and LocalAI HTTP serialization for completion and streaming requests. --- clients/schema_serialization_test.go | 107 ++++++++++++ mcp.go | 240 +-------------------------- mcp_boolschema_test.go | 3 + mcp_coerce_test.go | 3 + mcp_params_test.go | 34 ---- mcp_schema_legacy.go | 166 ++++++++++++++++++ mcp_schema_test.go | 70 ++++++++ 7 files changed, 355 insertions(+), 268 deletions(-) create mode 100644 clients/schema_serialization_test.go delete mode 100644 mcp_params_test.go create mode 100644 mcp_schema_legacy.go create mode 100644 mcp_schema_test.go diff --git a/clients/schema_serialization_test.go b/clients/schema_serialization_test.go new file mode 100644 index 0000000..4953008 --- /dev/null +++ b/clients/schema_serialization_test.go @@ -0,0 +1,107 @@ +package clients + +import ( + "context" + "encoding/json" + "net/http" + "net/http/httptest" + "reflect" + "testing" + "time" + + "github.com/mudler/cogito" + "github.com/sashabaranov/go-openai" +) + +// MCP discovery supplies RawMessage parameters. Exercise the actual HTTP +// serializers, including LocalAI's extensions and native-media message rewrite. +func TestClientsPreserveRawToolSchema(t *testing.T) { + const schema = `{"type":"object","$defs":{"id":{"type":"integer"}},"properties":{"status":{"anyOf":[{"enum":["pending","shipped","delivered","cancelled"],"type":"string"},{"type":"null"}],"default":null},"limit":{"type":"integer","default":20},"id":{"$ref":"#/$defs/id"},"never":false,"flag":{"type":"boolean","default":false},"tags":{"type":["array","null"],"items":true}},"allOf":[{"required":["id"]}],"additionalProperties":false}` + factories := map[string]func(string) cogito.StreamingLLM{ + "OpenAI": func(url string) cogito.StreamingLLM { return NewOpenAILLM("test-model", "", url) }, + "LocalAI": func(url string) cogito.StreamingLLM { return NewLocalAILLM("test-model", "", url) }, + "LocalAI extensions": func(url string) cogito.StreamingLLM { + c := NewLocalAILLM("test-model", "", url) + c.SetMetadata(map[string]string{"test": "true"}) + return c + }, + "LocalAI native media": func(url string) cogito.StreamingLLM { + c := NewLocalAILLM("test-model", "", url) + c.SetPendingNativeParts([]cogito.NativePart{{Kind: cogito.MediaAudio, Data: "AAAA", Format: "wav"}}) + return c + }, + } + for name, factory := range factories { + for _, stream := range []bool{false, true} { + mode := "completion" + if stream { + mode = "stream" + } + t.Run(name+"/"+mode, func(t *testing.T) { + captured := make(chan json.RawMessage, 1) + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + var request struct { + Tools []struct { + Function struct { + Parameters json.RawMessage `json:"parameters"` + } `json:"function"` + } `json:"tools"` + } + if err := json.NewDecoder(r.Body).Decode(&request); err != nil { + t.Error(err) + http.Error(w, "invalid JSON", 400) + return + } + if len(request.Tools) != 1 { + t.Errorf("got %d tools, want 1", len(request.Tools)) + http.Error(w, "missing tool", 400) + return + } + captured <- request.Tools[0].Function.Parameters + if stream { + w.Header().Set("Content-Type", "text/event-stream") + _, _ = w.Write([]byte("data: {\"choices\":[{\"index\":0,\"delta\":{\"content\":\"ok\"},\"finish_reason\":\"stop\"}]}\n\ndata: [DONE]\n\n")) + } else { + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(`{"choices":[{"index":0,"message":{"role":"assistant","content":"ok"},"finish_reason":"stop"}]}`)) + } + })) + defer srv.Close() + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + client := factory(srv.URL + "/v1") + req := openai.ChatCompletionRequest{Messages: []openai.ChatCompletionMessage{{Role: "user", Content: "test"}}, Tools: []openai.Tool{{Type: openai.ToolTypeFunction, Function: &openai.FunctionDefinition{Name: "example", Parameters: json.RawMessage(schema)}}}} + if stream { + events, err := client.CreateChatCompletionStream(ctx, req) + if err != nil { + t.Fatal(err) + } + for event := range events { + if event.Error != nil { + t.Fatal(event.Error) + } + } + } else { + if _, _, err := client.CreateChatCompletion(ctx, req); err != nil { + t.Fatal(err) + } + } + select { + case got := <-captured: + var wantValue, gotValue any + if err := json.Unmarshal([]byte(schema), &wantValue); err != nil { + t.Fatal(err) + } + if err := json.Unmarshal(got, &gotValue); err != nil { + t.Fatal(err) + } + if !reflect.DeepEqual(wantValue, gotValue) { + t.Fatalf("schema changed on HTTP boundary:\nwant %s\n got %s", schema, got) + } + case <-ctx.Done(): + t.Fatal("no request captured") + } + }) + } + } +} diff --git a/mcp.go b/mcp.go index f05573d..9beaa04 100644 --- a/mcp.go +++ b/mcp.go @@ -9,36 +9,22 @@ import ( "github.com/modelcontextprotocol/go-sdk/mcp" "github.com/mudler/xlog" "github.com/sashabaranov/go-openai" - "github.com/tmc/langchaingo/jsonschema" ) type mcpTool struct { name, description string - inputSchema toolInputSchema + inputSchema json.RawMessage session *mcp.ClientSession ctx context.Context - props map[string]jsonschema.Definition } func (t *mcpTool) Tool() openai.Tool { - // A tool without arguments has no properties, and the unmarshal leaves - // the map nil. jsonschema.Definition encoded that as {}, but this - // struct would send null, which vLLM and others reject. - props := t.props - if props == nil { - props = map[string]jsonschema.Definition{} - } return openai.Tool{ Type: openai.ToolTypeFunction, Function: &openai.FunctionDefinition{ Name: t.name, Description: t.description, - Parameters: mcpToolParameters{ - Type: jsonschema.Object, - Properties: props, - Required: t.inputSchema.Required, - AdditionalProperties: t.inputSchema.AdditionalProperties, - }, + Parameters: t.inputSchema, }, } } @@ -105,194 +91,6 @@ func (t *mcpTool) Close() { } } -type toolInputSchema struct { - Type string `json:"type"` - Properties map[string]interface{} `json:"properties,omitempty"` - Required []string `json:"required,omitempty"` - // AdditionalProperties is kept as the server declared it. A server whose - // schema forbids unknown arguments (the go-sdk infers false for a Go - // struct) rejects a call that has one, so the model should see the same - // rule, and a strict or grammar-constrained backend can enforce it. - AdditionalProperties any `json:"additionalProperties,omitempty"` -} - -// mcpToolParameters is the parameters schema an MCP tool is advertised -// with. It is not a jsonschema.Definition because that type has no -// additionalProperties field, and embedding it would not help: its value -// MarshalJSON would replace this struct's encoding. -type mcpToolParameters struct { - Type jsonschema.DataType `json:"type"` - Properties map[string]jsonschema.Definition `json:"properties"` - Required []string `json:"required,omitempty"` - AdditionalProperties any `json:"additionalProperties,omitempty"` -} - -// CoerceNullableTypes is an exported alias for the same workaround so -// downstream tests can verify their own MCP servers stay compatible -// with the import path. Most callers won't need it. -func CoerceNullableTypes(props map[string]any) { coerceNullableTypes(props) } - -// coerceNullableTypes recursively walks a property bag and normalizes the -// JSON-Schema 2020-12 constructs that langchaingo/jsonschema.Definition -// cannot represent. Both rewrites exist for the same reason: Definition is -// the struct we unmarshal every listed tool's properties into, and ANY -// failure there drops the whole tool from the model's tool list with a -// single log line — the server stays healthy and its tool simply never -// reaches the model. -// -// 1. "type": ["null", "X"] becomes "type": "X". Definition.Type is a single -// string. Picks the first non-null member; falls back to the first -// member if all are null. -// 2. A boolean schema becomes its object equivalent: `true` (allow -// anything) becomes {}, `false` (allow nothing) becomes {"not": {}}. -// 2020-12 permits a boolean wherever a schema is allowed, and -// google/jsonschema-go emits exactly that for a Go `any` — an empty -// schema marshals as `true`, so a [][]any field yields "items": true. -// Definition models nested schemas as *Definition, so a bare bool fails -// the unmarshal. -// -// Recurses into every nested schema location: properties, items, -// oneOf/anyOf/allOf members, prefixItems, $defs/definitions, -// additionalProperties, patternProperties, contains, not, if/then/else, -// propertyNames. -func coerceNullableTypes(props map[string]any) { - if props == nil { - return - } - // Range with the key so a property that is ITSELF a boolean schema can be - // replaced in place; the value-only loop could not rewrite it. - for name, raw := range props { - if b, ok := raw.(bool); ok { - props[name] = boolSchema(b) - continue - } - coerceSchema(raw) - } -} - -// boolSchema returns the object form of a JSON-Schema boolean schema, -// preserving its meaning: `true` allows anything, `false` allows nothing. -// Collapsing both to {} would silently widen a deliberately-closed schema. -func boolSchema(allow bool) map[string]any { - if allow { - return map[string]any{} - } - return map[string]any{"not": map[string]any{}} -} - -// schemaValuedKeys are the keywords whose value is a single schema, so a -// boolean there is a boolean schema rather than an ordinary flag. -var schemaValuedKeys = []string{ - "items", "additionalProperties", "contains", "not", - "if", "then", "else", "propertyNames", -} - -// schemaListKeys are the keywords whose value is an array of schemas. -var schemaListKeys = []string{"oneOf", "anyOf", "allOf", "prefixItems", "items"} - -// normalizeBoolSchemas replaces every boolean schema directly under obj with -// its object equivalent. Nested bags (properties, $defs, …) are reached by -// coerceNullableTypes, which performs the same replacement for their members. -func normalizeBoolSchemas(obj map[string]any) { - for _, key := range schemaValuedKeys { - if b, ok := obj[key].(bool); ok { - obj[key] = boolSchema(b) - } - } - for _, key := range schemaListKeys { - arr, ok := obj[key].([]any) - if !ok { - continue - } - for i, member := range arr { - if b, ok := member.(bool); ok { - arr[i] = boolSchema(b) - } - } - } -} - -// coerceSchema applies the type-array → string and boolean-schema → -// object rewrites to a single schema node, then recurses into every -// nested schema location. -func coerceSchema(node any) { - obj, ok := node.(map[string]any) - if !ok { - return - } - - // Run first, so the objects it creates are recursed into below like any - // other nested schema. - normalizeBoolSchemas(obj) - - if t, ok := obj["type"].([]any); ok { - pick := "" - for _, m := range t { - s, ok := m.(string) - if !ok || s == "null" { - continue - } - pick = s - break - } - if pick == "" && len(t) > 0 { - if s, ok := t[0].(string); ok { - pick = s - } - } - if pick != "" { - obj["type"] = pick - } - } - - // Object properties — map of name → schema. - if nested, ok := obj["properties"].(map[string]any); ok { - coerceNullableTypes(nested) - } - // patternProperties — same shape as properties, just regex-keyed. - if nested, ok := obj["patternProperties"].(map[string]any); ok { - coerceNullableTypes(nested) - } - // $defs / definitions — JSON-Schema named schema bag. - if nested, ok := obj["$defs"].(map[string]any); ok { - coerceNullableTypes(nested) - } - if nested, ok := obj["definitions"].(map[string]any); ok { - coerceNullableTypes(nested) - } - // Single nested schema fields. - coerceSchema(obj["items"]) - coerceSchema(obj["additionalProperties"]) - coerceSchema(obj["contains"]) - coerceSchema(obj["not"]) - coerceSchema(obj["if"]) - coerceSchema(obj["then"]) - coerceSchema(obj["else"]) - coerceSchema(obj["propertyNames"]) - // Composition keywords — arrays of schemas. - for _, key := range []string{"oneOf", "anyOf", "allOf"} { - arr, ok := obj[key].([]any) - if !ok { - continue - } - for _, member := range arr { - coerceSchema(member) - } - } - // "items" can also be an array (tuple validation in older drafts). - if arr, ok := obj["items"].([]any); ok { - for _, member := range arr { - coerceSchema(member) - } - } - // "prefixItems" (2020-12 tuple). - if arr, ok := obj["prefixItems"].([]any); ok { - for _, member := range arr { - coerceSchema(member) - } - } -} - func mcpPromptsFromTransport(ctx context.Context, session *mcp.ClientSession, arguments map[string]string) ([]openai.ChatCompletionMessage, error) { prompts, err := session.ListPrompts(ctx, nil) if err != nil { @@ -340,47 +138,21 @@ func mcpToolsFromTransport(ctx context.Context, session *mcp.ClientSession, filt if filter != nil && !filter(session, tool.Name) { continue } + // Keep the official SDK schema intact, including keywords and boolean + // schemas that a restricted schema struct cannot represent. Do not + // normalize types, add properties, or apply defaults here. dat, err := json.Marshal(tool.InputSchema) if err != nil { xlog.Error("Error marshalling input schema: %v", err) continue } - var inputSchema toolInputSchema - err = json.Unmarshal(dat, &inputSchema) - if err != nil { - xlog.Error("Error unmarshalling input schema: %v", err) - continue - } - - // Some MCP servers (e.g. modelcontextprotocol/go-sdk v1.4+) - // emit JSON Schema 2020-12 "type": ["null", "array"] for - // nullable fields like Go []string slices. langchaingo's - // jsonschema.Definition.Type is a single string, so the - // unmarshal would fail and silently drop the entire tool. - // Coerce any type-array to its non-null string member before - // unmarshalling so those tools stay discoverable. - coerceNullableTypes(inputSchema.Properties) - - props := map[string]jsonschema.Definition{} - dat, err = json.Marshal(inputSchema.Properties) - if err != nil { - xlog.Error("Error marshalling input schema: %v", err) - continue - } - err = json.Unmarshal(dat, &props) - if err != nil { - xlog.Error("Error unmarshalling input schema properties: %v", err) - continue - } - allTools = append(allTools, &mcpTool{ name: tool.Name, description: tool.Description, session: session, ctx: ctx, - props: props, - inputSchema: inputSchema, + inputSchema: json.RawMessage(dat), }) } diff --git a/mcp_boolschema_test.go b/mcp_boolschema_test.go index d444e9c..12d412c 100644 --- a/mcp_boolschema_test.go +++ b/mcp_boolschema_test.go @@ -1,5 +1,8 @@ package cogito +// These tests retain compatibility coverage for the deprecated normalization +// helper. MCP discovery does not use it; mcp_schema_test.go covers preservation. + import ( "encoding/json" "testing" diff --git a/mcp_coerce_test.go b/mcp_coerce_test.go index 6e3ec0c..ec95291 100644 --- a/mcp_coerce_test.go +++ b/mcp_coerce_test.go @@ -1,5 +1,8 @@ package cogito +// These tests retain compatibility coverage for the deprecated normalization +// helper. MCP discovery does not use it; mcp_schema_test.go covers preservation. + import ( "encoding/json" "testing" diff --git a/mcp_params_test.go b/mcp_params_test.go deleted file mode 100644 index 3f5ab13..0000000 --- a/mcp_params_test.go +++ /dev/null @@ -1,34 +0,0 @@ -package cogito - -import ( - "encoding/json" - "testing" -) - -// An MCP tool that takes no arguments has no "properties" in its input -// schema, so the unmarshal leaves props nil. OpenAI-compatible servers -// such as vLLM validate the schema and reject "properties": null, which -// fails the whole request, not only this tool. -func TestMCPToolWithoutArgumentsSendsEmptyProperties(t *testing.T) { - tool := &mcpTool{ - name: "list_things", - inputSchema: toolInputSchema{Type: "object"}, - } - - dat, err := json.Marshal(tool.Tool().Function.Parameters) - if err != nil { - t.Fatalf("marshal parameters: %v", err) - } - - var params map[string]any - if err := json.Unmarshal(dat, ¶ms); err != nil { - t.Fatalf("unmarshal parameters: %v", err) - } - props, ok := params["properties"].(map[string]any) - if !ok { - t.Fatalf("properties must be an object, got %s", dat) - } - if len(props) != 0 { - t.Fatalf("properties must be empty, got %s", dat) - } -} diff --git a/mcp_schema_legacy.go b/mcp_schema_legacy.go new file mode 100644 index 0000000..b08aa14 --- /dev/null +++ b/mcp_schema_legacy.go @@ -0,0 +1,166 @@ +package cogito + +// CoerceNullableTypes applies the legacy lossy schema normalization. +// +// Deprecated: MCP discovery now preserves the complete input schema and does +// not need this workaround. Retained only for existing callers. +func CoerceNullableTypes(props map[string]any) { coerceNullableTypes(props) } + +// coerceNullableTypes recursively walks a property bag and normalizes the +// JSON-Schema 2020-12 constructs that langchaingo/jsonschema.Definition +// cannot represent. Both rewrites exist for the same reason: Definition is +// the struct the old discovery pipeline used. This helper is no longer +// called by discovery because it changes the server's schema semantics. +// +// 1. "type": ["null", "X"] becomes "type": "X". Definition.Type is a single +// string. Picks the first non-null member; falls back to the first +// member if all are null. +// 2. A boolean schema becomes its object equivalent: `true` (allow +// anything) becomes {}, `false` (allow nothing) becomes {"not": {}}. +// 2020-12 permits a boolean wherever a schema is allowed, and +// google/jsonschema-go emits exactly that for a Go `any` — an empty +// schema marshals as `true`, so a [][]any field yields "items": true. +// Definition models nested schemas as *Definition, so a bare bool fails +// the unmarshal. +// +// Recurses into every nested schema location: properties, items, +// oneOf/anyOf/allOf members, prefixItems, $defs/definitions, +// additionalProperties, patternProperties, contains, not, if/then/else, +// propertyNames. +func coerceNullableTypes(props map[string]any) { + if props == nil { + return + } + // Range with the key so a property that is ITSELF a boolean schema can be + // replaced in place; the value-only loop could not rewrite it. + for name, raw := range props { + if b, ok := raw.(bool); ok { + props[name] = boolSchema(b) + continue + } + coerceSchema(raw) + } +} + +// boolSchema returns the object form of a JSON-Schema boolean schema, +// preserving its meaning: `true` allows anything, `false` allows nothing. +// Collapsing both to {} would silently widen a deliberately-closed schema. +func boolSchema(allow bool) map[string]any { + if allow { + return map[string]any{} + } + return map[string]any{"not": map[string]any{}} +} + +// schemaValuedKeys are the keywords whose value is a single schema, so a +// boolean there is a boolean schema rather than an ordinary flag. +var schemaValuedKeys = []string{ + "items", "additionalProperties", "contains", "not", + "if", "then", "else", "propertyNames", +} + +// schemaListKeys are the keywords whose value is an array of schemas. +var schemaListKeys = []string{"oneOf", "anyOf", "allOf", "prefixItems", "items"} + +// normalizeBoolSchemas replaces every boolean schema directly under obj with +// its object equivalent. Nested bags (properties, $defs, …) are reached by +// coerceNullableTypes, which performs the same replacement for their members. +func normalizeBoolSchemas(obj map[string]any) { + for _, key := range schemaValuedKeys { + if b, ok := obj[key].(bool); ok { + obj[key] = boolSchema(b) + } + } + for _, key := range schemaListKeys { + arr, ok := obj[key].([]any) + if !ok { + continue + } + for i, member := range arr { + if b, ok := member.(bool); ok { + arr[i] = boolSchema(b) + } + } + } +} + +// coerceSchema applies the type-array → string and boolean-schema → +// object rewrites to a single schema node, then recurses into every +// nested schema location. +func coerceSchema(node any) { + obj, ok := node.(map[string]any) + if !ok { + return + } + + // Run first, so the objects it creates are recursed into below like any + // other nested schema. + normalizeBoolSchemas(obj) + + if t, ok := obj["type"].([]any); ok { + pick := "" + for _, m := range t { + s, ok := m.(string) + if !ok || s == "null" { + continue + } + pick = s + break + } + if pick == "" && len(t) > 0 { + if s, ok := t[0].(string); ok { + pick = s + } + } + if pick != "" { + obj["type"] = pick + } + } + + // Object properties — map of name → schema. + if nested, ok := obj["properties"].(map[string]any); ok { + coerceNullableTypes(nested) + } + // patternProperties — same shape as properties, just regex-keyed. + if nested, ok := obj["patternProperties"].(map[string]any); ok { + coerceNullableTypes(nested) + } + // $defs / definitions — JSON-Schema named schema bag. + if nested, ok := obj["$defs"].(map[string]any); ok { + coerceNullableTypes(nested) + } + if nested, ok := obj["definitions"].(map[string]any); ok { + coerceNullableTypes(nested) + } + // Single nested schema fields. + coerceSchema(obj["items"]) + coerceSchema(obj["additionalProperties"]) + coerceSchema(obj["contains"]) + coerceSchema(obj["not"]) + coerceSchema(obj["if"]) + coerceSchema(obj["then"]) + coerceSchema(obj["else"]) + coerceSchema(obj["propertyNames"]) + // Composition keywords — arrays of schemas. + for _, key := range []string{"oneOf", "anyOf", "allOf"} { + arr, ok := obj[key].([]any) + if !ok { + continue + } + for _, member := range arr { + coerceSchema(member) + } + } + // "items" can also be an array (tuple validation in older drafts). + if arr, ok := obj["items"].([]any); ok { + for _, member := range arr { + coerceSchema(member) + } + } + // "prefixItems" (2020-12 tuple). + if arr, ok := obj["prefixItems"].([]any); ok { + for _, member := range arr { + coerceSchema(member) + } + } +} diff --git a/mcp_schema_test.go b/mcp_schema_test.go new file mode 100644 index 0000000..6248595 --- /dev/null +++ b/mcp_schema_test.go @@ -0,0 +1,70 @@ +package cogito + +import ( + "context" + "encoding/json" + "reflect" + "testing" + "time" + + "github.com/modelcontextprotocol/go-sdk/mcp" +) + +// Test the discovery boundary, not a hand-built approximation of Parameters. +func TestMCPDiscoveryPreservesInputSchema(t *testing.T) { + cases := map[string]string{ + "nullable enum and defaults": `{"type":"object","properties":{"status":{"anyOf":[{"enum":["pending","shipped","delivered","cancelled"],"type":"string"},{"type":"null"}],"default":null},"limit":{"default":20,"type":"integer"}}}`, + "root references and compositions": `{"type":"object","$schema":"https://json-schema.org/draft/2020-12/schema","$defs":{"args":{"type":"object","properties":{"id":{"type":"integer","minimum":1}}}},"$ref":"#/$defs/args","allOf":[{"required":["id"]}],"anyOf":[{"required":["id"]},false],"oneOf":[{"required":["id"]}],"unevaluatedProperties":false,"title":"Arguments","x-extension":{"enabled":false}}`, + "nullable types and boolean schemas": `{"type":"object","properties":{"tags":{"type":["null","array"],"items":{"type":["string","null"]}},"rows":{"type":"array","items":{"type":"array","items":true}},"never":false,"anything":true,"closed":{"type":"object","additionalProperties":false},"flag":{"type":"boolean","default":false},"tuple":{"prefixItems":[true,false],"items":false}},"additionalProperties":false}`, + "nested constraints": `{"type":"object","properties":{"config":{"type":"object","patternProperties":{"^x":{"type":["string","null"]}},"additionalProperties":{"type":"integer"},"dependentSchemas":{"x":{"required":["y"]}},"propertyNames":{"minLength":1}},"values":{"type":"array","contains":{"const":0},"minContains":1,"maxItems":10,"uniqueItems":true,"items":{"anyOf":[{"type":"integer"},false]}}},"if":{"required":["config"]},"then":{"required":["values"]},"else":{"not":{"required":["values"]}},"required":[]}`, + "no arguments absent properties": `{"type":"object"}`, + "no arguments explicit properties": `{"type":"object","properties":{},"additionalProperties":false}`, + } + for name, schema := range cases { + t.Run(name, func(t *testing.T) { + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + impl := &mcp.Implementation{Name: "schema-test", Version: "1"} + server := mcp.NewServer(impl, nil) + server.AddTool(&mcp.Tool{Name: "example", Description: "Example tool", InputSchema: json.RawMessage(schema)}, func(context.Context, *mcp.CallToolRequest) (*mcp.CallToolResult, error) { + return &mcp.CallToolResult{}, nil + }) + st, ct := mcp.NewInMemoryTransports() + ss, err := server.Connect(ctx, st, nil) + if err != nil { + t.Fatal(err) + } + defer ss.Close() + session, err := mcp.NewClient(impl, nil).Connect(ctx, ct, nil) + if err != nil { + t.Fatal(err) + } + defer session.Close() + tools, err := mcpToolsFromTransport(ctx, session, nil) + if err != nil { + t.Fatal(err) + } + if len(tools) != 1 { + t.Fatalf("discovered %d tools, want 1", len(tools)) + } + tool := tools[0].Tool() + if tool.Function.Name != "example" || tool.Function.Description != "Example tool" { + t.Fatalf("tool metadata changed: %+v", tool) + } + got, err := json.Marshal(tool.Function.Parameters) + if err != nil { + t.Fatal(err) + } + var wantValue, gotValue any + if err := json.Unmarshal([]byte(schema), &wantValue); err != nil { + t.Fatal(err) + } + if err := json.Unmarshal(got, &gotValue); err != nil { + t.Fatal(err) + } + if !reflect.DeepEqual(wantValue, gotValue) { + t.Fatalf("schema changed during discovery:\nwant %s\n got %s", schema, got) + } + }) + } +}