| 1 | package tool |
| 2 | |
| 3 | import ( |
| 4 | "context" |
| 5 | "encoding/json" |
| 6 | "testing" |
| 7 | ) |
| 8 | |
| 9 | // stubTool is a minimal Tool for registry tests. |
| 10 | type stubTool struct { |
| 11 | name string |
| 12 | schema json.RawMessage |
| 13 | server string |
| 14 | raw string |
| 15 | visible string |
| 16 | pkg string |
| 17 | } |
| 18 | |
| 19 | func (s stubTool) Name() string { return s.name } |
| 20 | func (s stubTool) Description() string { return s.name + " desc" } |
| 21 | func (s stubTool) Schema() json.RawMessage { |
| 22 | if len(s.schema) > 0 { |
| 23 | return s.schema |
| 24 | } |
| 25 | return json.RawMessage(`{"type":"object"}`) |
| 26 | } |
| 27 | func (s stubTool) Execute(context.Context, json.RawMessage) (string, error) { return "", nil } |
| 28 | func (s stubTool) ReadOnly() bool { return true } |
| 29 | func (s stubTool) MCPServerName() string { return s.server } |
| 30 | func (s stubTool) MCPRawToolName() string { return s.raw } |
| 31 | func (s stubTool) MCPVisibleToolName() string { return s.visible } |
| 32 | func (s stubTool) MCPPackageName() string { return s.pkg } |
| 33 | |
| 34 | func TestRegistryResolvesPortableMCPReferencesOnlyWhenUnique(t *testing.T) { |
| 35 | r := NewRegistry() |
| 36 | first := stubTool{name: "mcp__figma__get_design_context", server: "figma", raw: "figma_get_design_context", visible: "get_design_context", pkg: "figma"} |
| 37 | r.Add(first) |
| 38 | |
| 39 | refs := []string{ |
| 40 | "get_design_context", |
| 41 | "figma_get_design_context", |
| 42 | "figma/get_design_context", |
| 43 | "mcp-tool:figma/figma_get_design_context", |
| 44 | "mcp__plugin_figma_figma__get_design_context", |
| 45 | } |
| 46 | for _, ref := range refs { |
| 47 | got, canonical, ambiguous := r.ResolveCall(ref) |
| 48 | if got == nil || canonical != first.name || len(ambiguous) != 0 { |
| 49 | t.Errorf("ResolveCall(%q) = (%v, %q, %v), want %q", ref, got, canonical, ambiguous, first.name) |
| 50 | } |
| 51 | } |
| 52 | |
| 53 | // Exact registered names always win, even if they are also another MCP |
| 54 | // tool's short/raw alias. |
| 55 | r.Add(stubTool{name: "get_design_context"}) |
| 56 | got, canonical, ambiguous := r.ResolveCall("get_design_context") |
| 57 | if got == nil || canonical != "get_design_context" || len(ambiguous) != 0 { |
| 58 | t.Fatalf("exact name did not win: (%v, %q, %v)", got, canonical, ambiguous) |
| 59 | } |
| 60 | |
| 61 | r.Add(stubTool{name: "mcp__other__get_design_context", server: "other", raw: "get_design_context", visible: "get_design_context"}) |
| 62 | _, _, ambiguous = r.ResolveCall("figma_get_design_context") |
| 63 | if len(ambiguous) != 0 { |
| 64 | t.Fatalf("distinct raw name became ambiguous: %v", ambiguous) |
| 65 | } |
| 66 | _, _, ambiguous = r.ResolveCall("get_design_context") |
| 67 | if len(ambiguous) != 0 { // exact builtin-style name still wins |
| 68 | t.Fatalf("exact registered name should suppress alias ambiguity: %v", ambiguous) |
| 69 | } |
| 70 | r.RemovePrefix("get_design_context") |
| 71 | _, canonical, ambiguous = r.ResolveCall("get_design_context") |
| 72 | if canonical != "" || len(ambiguous) != 2 { |
| 73 | t.Fatalf("ambiguous short reference = canonical %q candidates %v, want two candidates", canonical, ambiguous) |
| 74 | } |
| 75 | } |
| 76 | |
| 77 | func TestRegistryPortableAliasesDoNotChangeProviderSchemas(t *testing.T) { |
| 78 | r := NewRegistry() |
| 79 | r.Add(stubTool{name: "mcp__my_server_deadbeef__do_thing_deadbeef", server: "my.server", raw: "do.thing", visible: "do.thing"}) |
| 80 | before := r.Schemas() |
| 81 | if got, canonical, ambiguous := r.ResolveCall("mcp__my_server__do_thing"); got == nil || canonical == "" || len(ambiguous) != 0 { |
| 82 | t.Fatalf("portable normalized reference did not resolve: (%v, %q, %v)", got, canonical, ambiguous) |
| 83 | } |
| 84 | after := r.Schemas() |
| 85 | if len(before) != 1 || len(after) != 1 || before[0].Name != after[0].Name { |
| 86 | t.Fatalf("alias resolution changed provider schemas: before=%v after=%v", before, after) |
| 87 | } |
| 88 | } |
| 89 | |
| 90 | // TestRegistryRemovePrefix proves an MCP server's namespaced tools are dropped as |
| 91 | // a group on disconnect, leaving built-ins and other servers' tools — and their |
| 92 | // insertion order — intact. |
| 93 | func TestRegistryRemovePrefix(t *testing.T) { |
| 94 | r := NewRegistry() |
| 95 | r.Add(stubTool{name: "bash"}) |
| 96 | r.Add(stubTool{name: "mcp__fs__read"}) |
| 97 | r.Add(stubTool{name: "mcp__fs__write"}) |
| 98 | r.Add(stubTool{name: "mcp__stripe__charge"}) |
| 99 | |
| 100 | if got := r.RemovePrefix("mcp__fs__"); got != 2 { |
| 101 | t.Fatalf("RemovePrefix returned %d, want 2", got) |
| 102 | } |
| 103 | if r.Len() != 2 { |
| 104 | t.Fatalf("registry has %d tools after removal, want 2", r.Len()) |
| 105 | } |
| 106 | if _, ok := r.Get("mcp__fs__read"); ok { |
| 107 | t.Errorf("mcp__fs__read should be gone") |
| 108 | } |
| 109 | if _, ok := r.Get("mcp__stripe__charge"); !ok { |
| 110 | t.Errorf("another server's tool should survive") |
| 111 | } |
| 112 | want := []string{"bash", "mcp__stripe__charge"} |
| 113 | got := r.Names() |
| 114 | if len(got) != len(want) { |
| 115 | t.Fatalf("names = %v, want %v", got, want) |
| 116 | } |
| 117 | for i := range want { |
| 118 | if got[i] != want[i] { |
| 119 | t.Fatalf("names = %v, want %v (order preserved)", got, want) |
| 120 | } |
| 121 | } |
| 122 | |
| 123 | // Removing a prefix that matches nothing is a no-op. |
| 124 | if got := r.RemovePrefix("mcp__nope__"); got != 0 { |
| 125 | t.Errorf("RemovePrefix on absent prefix returned %d, want 0", got) |
| 126 | } |
| 127 | } |
| 128 | |
| 129 | func TestRegistrySuspendPrefixBlocksLateAddsUntilResume(t *testing.T) { |
| 130 | r := NewRegistry() |
| 131 | r.Add(stubTool{name: "bash"}) |
| 132 | r.Add(stubTool{name: "mcp__fs__connect"}) |
| 133 | |
| 134 | if got := r.SuspendPrefix("mcp__fs__"); got != 1 { |
| 135 | t.Fatalf("SuspendPrefix returned %d, want 1", got) |
| 136 | } |
| 137 | r.Add(stubTool{name: "mcp__fs__read"}) |
| 138 | if _, ok := r.Get("mcp__fs__read"); ok { |
| 139 | t.Fatalf("suspended prefix accepted a late tool add; names=%v", r.Names()) |
| 140 | } |
| 141 | if _, ok := r.Get("bash"); !ok { |
| 142 | t.Fatal("suspending an MCP prefix removed unrelated tools") |
| 143 | } |
| 144 | |
| 145 | r.ResumePrefix("mcp__fs__") |
| 146 | r.Add(stubTool{name: "mcp__fs__read"}) |
| 147 | if _, ok := r.Get("mcp__fs__read"); !ok { |
| 148 | t.Fatalf("resumed prefix did not accept tool add; names=%v", r.Names()) |
| 149 | } |
| 150 | } |
| 151 | |
| 152 | // TestRegistrySchemasSorted proves Schemas() emits tool definitions in |
| 153 | // deterministic alphabetical order regardless of insertion order, so a logically |
| 154 | // identical tool set produces a stable provider-facing request prefix (prompt |
| 155 | // cache reuse). Names() must stay in insertion order — only the provider export |
| 156 | // is sorted. |
| 157 | func TestRegistrySchemasSorted(t *testing.T) { |
| 158 | r := NewRegistry() |
| 159 | // Add deliberately out of alphabetical order. |
| 160 | insertion := []string{"write", "bash", "read", "apply_patch"} |
| 161 | for _, n := range insertion { |
| 162 | r.Add(stubTool{name: n}) |
| 163 | } |
| 164 | |
| 165 | var got []string |
| 166 | for _, s := range r.Schemas() { |
| 167 | got = append(got, s.Name) |
| 168 | } |
| 169 | want := []string{"apply_patch", "bash", "read", "write"} |
| 170 | if len(got) != len(want) { |
| 171 | t.Fatalf("Schemas() names = %v, want %v", got, want) |
| 172 | } |
| 173 | for i := range want { |
| 174 | if got[i] != want[i] { |
| 175 | t.Fatalf("Schemas() names = %v, want %v (alphabetical)", got, want) |
| 176 | } |
| 177 | } |
| 178 | |
| 179 | // The sort must not leak into Names(): display order stays insertion order. |
| 180 | gotNames := r.Names() |
| 181 | for i := range insertion { |
| 182 | if gotNames[i] != insertion[i] { |
| 183 | t.Fatalf("Names() = %v, want %v (insertion order)", gotNames, insertion) |
| 184 | } |
| 185 | } |
| 186 | } |
| 187 | |
| 188 | func TestRegistrySchemasStableAndCanonical(t *testing.T) { |
| 189 | r := NewRegistry() |
| 190 | r.Add(stubTool{ |
| 191 | name: "zeta", |
| 192 | schema: json.RawMessage(`{"type":"object","required":["b","a"],"properties":{"b":{"type":"string"},"a":{"type":"string"}}}`), |
| 193 | }) |
| 194 | r.Add(stubTool{ |
| 195 | name: "alpha", |
| 196 | schema: json.RawMessage(`{"required":["y","x"],"type":"object"}`), |
| 197 | }) |
| 198 | |
| 199 | schemas := r.Schemas() |
| 200 | if len(schemas) != 2 { |
| 201 | t.Fatalf("Schemas returned %d entries, want 2", len(schemas)) |
| 202 | } |
| 203 | if schemas[0].Name != "alpha" || schemas[1].Name != "zeta" { |
| 204 | t.Fatalf("Schemas order = %q, %q; want alpha, zeta", schemas[0].Name, schemas[1].Name) |
| 205 | } |
| 206 | if got, want := string(schemas[0].Parameters), `{"properties":{},"required":["x","y"],"type":"object"}`; got != want { |
| 207 | t.Fatalf("alpha schema = %s, want %s", got, want) |
| 208 | } |
| 209 | if got, want := string(schemas[1].Parameters), `{"properties":{"a":{"type":"string"},"b":{"type":"string"}},"required":["a","b"],"type":"object"}`; got != want { |
| 210 | t.Fatalf("zeta schema = %s, want %s", got, want) |
| 211 | } |
| 212 | } |
| 213 | |
| 214 | func TestRegistrySchemasCanonicalizesEquivalentOrdering(t *testing.T) { |
| 215 | first := NewRegistry() |
| 216 | first.Add(stubTool{ |
| 217 | name: "same", |
| 218 | schema: json.RawMessage(`{"type":"object","required":["b","a"],"properties":{"b":{"description":"bee","type":"string"},"a":{"type":"integer"}}}`), |
| 219 | }) |
| 220 | |
| 221 | second := NewRegistry() |
| 222 | second.Add(stubTool{ |
| 223 | name: "same", |
| 224 | schema: json.RawMessage(`{"properties":{"a":{"type":"integer"},"b":{"type":"string","description":"bee"}},"required":["a","b"],"type":"object"}`), |
| 225 | }) |
| 226 | |
| 227 | firstSchemas := first.Schemas() |
| 228 | secondSchemas := second.Schemas() |
| 229 | if got, want := string(firstSchemas[0].Parameters), string(secondSchemas[0].Parameters); got != want { |
| 230 | t.Fatalf("equivalent schemas canonicalized differently:\n first: %s\n second: %s", got, want) |
| 231 | } |
| 232 | } |
| 233 |