diff --git a/internal/mcp/tool_contract_test.go b/internal/mcp/tool_contract_test.go new file mode 100644 index 000000000..4b2805707 --- /dev/null +++ b/internal/mcp/tool_contract_test.go @@ -0,0 +1,329 @@ +package mcp + +import ( + "bytes" + "encoding/json" + "fmt" + "io" + "slices" + "sort" + "strings" + "testing" + + mcpserver "github.com/mark3labs/mcp-go/server" +) + +type mcpToolSchema struct { + Types []string + Properties map[string]mcpToolSchema + Required []string + Items *mcpToolSchema + Enum []string + Additional bool +} + +func TestMCPToolSchemaCodec(t *testing.T) { + for _, tc := range []struct { + name, raw, want string + }{ + {"root duplicate", `{"x":1,"x":2}`, "/input/x"}, + {"properties duplicate", `{"type":"object","properties":{"a/b":{"type":"string","type":"number"}}}`, "/input/properties/a~1b/type"}, + {"items duplicate", `{"type":"array","items":{"type":"string","type":"number"}}`, "/input/items/type"}, + {"array duplicate", `[{"x":1,"x":2}]`, "/input/0/x"}, + {"malformed", `{"type":`, "malformed"}, + {"trailing", `{"type":"string"} {}`, "trailing"}, + } { + t.Run(tc.name, func(t *testing.T) { + _, err := decodeMCPToolSchema([]byte(tc.raw), "/input") + if err == nil || !strings.Contains(err.Error(), tc.want) { + t.Fatalf("error = %v, want %q", err, tc.want) + } + }) + } + + value, err := decodeMCPToolSchema([]byte(`{"enum":[1.25,900719925474099312345]}`), "/input") + if err != nil { + t.Fatal(err) + } + enum := value.(map[string]any)["enum"].([]any) + if enum[0] != json.Number("1.25") || enum[1] != json.Number("900719925474099312345") { + t.Fatalf("numbers = %#v", enum) + } +} + +func TestNormalizeMCPToolContract(t *testing.T) { + for _, tc := range []struct { + name, raw, want string + }{ + {"supported recursion", `{"type":"object","title":"ignored","properties":{"a/b":{"type":"array","items":{"type":["integer","string"],"description":"ignored"}}},"required":["a/b","a/b"],"additionalProperties":false}`, ""}, + {"missing type", `{"properties":{}}`, "/input"}, + {"unknown keyword", `{"type":"string","pattern":"x"}`, "/input/pattern"}, + {"malformed properties", `{"type":"object","properties":[]}`, "/input/properties"}, + {"invalid required", `{"type":"object","properties":{},"required":["missing"]}`, "/input/required"}, + {"tuple items", `{"type":"array","items":[{"type":"string"}]}`, "/input/items"}, + {"schema additional", `{"type":"object","additionalProperties":{"type":"string"}}`, "/input/additionalProperties"}, + {"escaped pointer", `{"type":"object","properties":{"a/b~c":{"type":"string","oneOf":[]}}}`, "/input/properties/a~1b~0c/oneOf"}, + } { + t.Run(tc.name, func(t *testing.T) { + _, err := normalizeMCPToolContract([]byte(tc.raw), "/input") + if tc.want == "" && err != nil || tc.want != "" && (err == nil || !strings.Contains(err.Error(), tc.want)) { + t.Fatalf("error = %v, want %q", err, tc.want) + } + }) + } + + one := mustNormalize(t, `{"type":"string","enum":[900719925474099312345,1.25,1.25]}`) + two := mustNormalize(t, `{"examples":["ignored"],"enum":[1.25,900719925474099312345],"type":["string"]}`) + if strings.Join(one.Enum, ",") != "1.25,900719925474099312345" || strings.Join(one.Enum, ",") != strings.Join(two.Enum, ",") { + t.Fatalf("enum normalization differs: %#v %#v", one.Enum, two.Enum) + } +} + +func TestObserveMCPToolContract(t *testing.T) { + server := NewServer(newMCPTestStore(t)) + live, err := observeMCPToolContract(server.ListTools()) + if err != nil { + t.Fatal(err) + } + if len(live) == 0 { + t.Fatal("default registry is empty") + } + if len(live) != len(server.ListTools()) { + t.Fatalf("observed %d of %d live tools", len(live), len(server.ListTools())) + } + for name := range server.ListTools() { + if _, ok := live[name]; !ok { + t.Fatalf("%s bypassed normalization", name) + } + } +} + +func mustNormalize(t *testing.T, raw string) mcpToolSchema { + t.Helper() + schema, err := normalizeMCPToolContract([]byte(raw), "/input") + if err != nil { + t.Fatal(err) + } + return schema +} + +func decodeMCPToolSchema(raw []byte, path string) (any, error) { + decoder := json.NewDecoder(bytes.NewReader(raw)) + decoder.UseNumber() + value, err := uniqueJSONValue(decoder, path) + if err != nil { + return nil, fmt.Errorf("%s: malformed-schema: %w", path, err) + } + var extra any + if err := decoder.Decode(&extra); err != io.EOF { + if err == nil { + return nil, fmt.Errorf("%s: malformed-schema: trailing JSON value", path) + } + return nil, fmt.Errorf("%s: malformed-schema: %w", path, err) + } + return value, nil +} + +func uniqueJSONValue(decoder *json.Decoder, path string) (any, error) { + token, err := decoder.Token() + if err != nil { + return nil, err + } + switch token { + case json.Delim('{'): + out, seen := map[string]any{}, map[string]bool{} + for decoder.More() { + key, err := decoder.Token() + if err != nil { + return nil, err + } + name := key.(string) + childPath := jsonPointer(path, name) + if seen[name] { + return nil, fmt.Errorf("duplicate member at %s", childPath) + } + seen[name] = true + out[name], err = uniqueJSONValue(decoder, childPath) + if err != nil { + return nil, err + } + } + _, err = decoder.Token() + return out, err + case json.Delim('['): + out := []any{} + for index := 0; decoder.More(); index++ { + value, err := uniqueJSONValue(decoder, fmt.Sprintf("%s/%d", path, index)) + if err != nil { + return nil, err + } + out = append(out, value) + } + _, err = decoder.Token() + return out, err + default: + return token, nil + } +} + +func normalizeMCPToolContract(raw []byte, path string) (mcpToolSchema, error) { + value, err := decodeMCPToolSchema(raw, path) + if err != nil { + return mcpToolSchema{}, err + } + return normalizeMCPToolSchema(value, path) +} + +func normalizeMCPToolSchema(value any, path string) (mcpToolSchema, error) { + object, ok := value.(map[string]any) + if !ok { + return mcpToolSchema{}, schemaError(path, "schema must be an object") + } + schema := mcpToolSchema{Properties: map[string]mcpToolSchema{}, Additional: true} + for _, key := range sortedKeys(object) { + value, keyPath := object[key], jsonPointer(path, key) + if excludedMetadata(key) { + continue + } + switch key { + case "type": + if text, ok := value.(string); ok { + value = []any{text} + } + var err error + schema.Types, err = normalizedStrings(value, keyPath, false) + if err != nil { + return schema, err + } + case "properties": + properties, ok := value.(map[string]any) + if !ok { + return schema, schemaError(keyPath, "properties must be an object") + } + for _, name := range sortedKeys(properties) { + child, err := normalizeMCPToolSchema(properties[name], jsonPointer(keyPath, name)) + if err != nil { + return schema, err + } + schema.Properties[name] = child + } + case "required": + var err error + schema.Required, err = normalizedStrings(value, keyPath, true) + if err != nil { + return schema, err + } + case "items": + child, err := normalizeMCPToolSchema(value, keyPath) + if err != nil { + return schema, err + } + schema.Items = &child + case "enum": + values, ok := value.([]any) + if !ok || len(values) == 0 { + return schema, schemaError(keyPath, "enum must be a non-empty array") + } + for _, value := range values { + encoded, err := json.Marshal(value) + if err != nil { + return schema, schemaError(keyPath, err.Error()) + } + schema.Enum = append(schema.Enum, string(encoded)) + } + schema.Enum = uniqueStrings(schema.Enum) + case "additionalProperties": + additional, ok := value.(bool) + if !ok { + return schema, fmt.Errorf("%s: unsupported-keyword: only boolean values are supported", keyPath) + } + schema.Additional = additional + default: + return schema, fmt.Errorf("%s: unsupported-keyword: %s", keyPath, key) + } + } + if len(schema.Types) == 0 { + return schema, schemaError(path, "type is required") + } + if len(schema.Properties) > 0 && !slices.Contains(schema.Types, "object") || schema.Items != nil && !slices.Contains(schema.Types, "array") { + return schema, schemaError(path, "incompatible recursive shape") + } + for _, name := range schema.Required { + if _, ok := schema.Properties[name]; !ok { + return schema, schemaError(jsonPointer(path, "required"), "required property is absent") + } + } + return schema, nil +} + +func observeMCPToolContract(tools map[string]*mcpserver.ServerTool) (map[string]mcpToolSchema, error) { + observed := make(map[string]mcpToolSchema, len(tools)) + for name, tool := range tools { + raw := tool.Tool.RawInputSchema + if len(raw) == 0 { + var err error + raw, err = json.Marshal(tool.Tool.InputSchema) + if err != nil { + return nil, fmt.Errorf("/tools/%s: marshal schema: %w", escapePointer(name), err) + } + } + if len(bytes.TrimSpace(raw)) == 0 { + return nil, schemaError(jsonPointer("/tools", name), "schema is empty") + } + schema, err := normalizeMCPToolContract(raw, jsonPointer("/tools", name)) + if err != nil { + return nil, err + } + observed[name] = schema + } + return observed, nil +} +func normalizedStrings(value any, path string, emptyOK bool) ([]string, error) { + values, ok := value.([]any) + if !ok || !emptyOK && len(values) == 0 { + return nil, schemaError(path, "must be a non-empty string array") + } + out := make([]string, 0, len(values)) + for _, value := range values { + text, ok := value.(string) + if !ok || !emptyOK && !strings.Contains(",object,array,string,number,integer,boolean,null,", ","+text+",") { + return nil, schemaError(path, "must contain supported type strings") + } + out = append(out, text) + } + return uniqueStrings(out), nil +} + +func schemaError(path, detail string) error { + return fmt.Errorf("%s: malformed-schema: %s", path, detail) +} +func uniqueStrings(values []string) []string { + sort.Strings(values) + out := values[:0] + for _, value := range values { + if len(out) == 0 || out[len(out)-1] != value { + out = append(out, value) + } + } + return out +} +func sortedKeys[V any](values map[string]V) []string { + keys := make([]string, 0, len(values)) + for key := range values { + keys = append(keys, key) + } + sort.Strings(keys) + return keys +} +func excludedMetadata(key string) bool { + return strings.Contains(",description,title,default,examples,$comment,deprecated,readOnly,writeOnly,", ","+key+",") +} +func escapePointer(value string) string { + return strings.ReplaceAll(strings.ReplaceAll(value, "~", "~0"), "/", "~1") +} +func jsonPointer(path, value string) string { + if path == "/" { + return "/" + escapePointer(value) + } + return path + "/" + escapePointer(value) +}