| 1 | package agent |
| 2 | |
| 3 | import ( |
| 4 | "bytes" |
| 5 | "context" |
| 6 | "encoding/json" |
| 7 | "fmt" |
| 8 | "net/http" |
| 9 | "net/http/httptest" |
| 10 | "slices" |
| 11 | "strings" |
| 12 | "sync" |
| 13 | "sync/atomic" |
| 14 | "testing" |
| 15 | "time" |
| 16 | |
| 17 | "reasonix/internal/capability" |
| 18 | "reasonix/internal/config" |
| 19 | "reasonix/internal/plugin" |
| 20 | "reasonix/internal/tool" |
| 21 | ) |
| 22 | |
| 23 | func dynamicToolsMCPServer(t *testing.T, loaded *atomic.Bool, dynamicCalls *atomic.Int32) *httptest.Server { |
| 24 | t.Helper() |
| 25 | return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { |
| 26 | var request struct { |
| 27 | ID *int `json:"id"` |
| 28 | Method string `json:"method"` |
| 29 | Params json.RawMessage `json:"params"` |
| 30 | } |
| 31 | if err := json.NewDecoder(r.Body).Decode(&request); err != nil { |
| 32 | http.Error(w, "bad request", http.StatusBadRequest) |
| 33 | return |
| 34 | } |
| 35 | if request.ID == nil { |
| 36 | w.WriteHeader(http.StatusAccepted) |
| 37 | return |
| 38 | } |
| 39 | |
| 40 | var result any |
| 41 | notifyChanged := false |
| 42 | switch request.Method { |
| 43 | case "initialize": |
| 44 | result = map[string]any{ |
| 45 | "protocolVersion": "2024-11-05", |
| 46 | "serverInfo": map[string]any{"name": "dynamic", "version": "1"}, |
| 47 | "capabilities": map[string]any{"tools": map[string]any{"listChanged": true}}, |
| 48 | } |
| 49 | case "tools/list": |
| 50 | tools := []map[string]any{{ |
| 51 | "name": "load_toolset", |
| 52 | "description": "Load the schematic toolset.", |
| 53 | "inputSchema": map[string]any{"type": "object"}, |
| 54 | }} |
| 55 | if loaded.Load() { |
| 56 | tools = append(tools, map[string]any{ |
| 57 | "name": "list_schematic_components", |
| 58 | "description": "List schematic components.", |
| 59 | "inputSchema": map[string]any{"type": "object"}, |
| 60 | }) |
| 61 | } |
| 62 | result = map[string]any{"tools": tools} |
| 63 | case "tools/call": |
| 64 | var params struct { |
| 65 | Name string `json:"name"` |
| 66 | } |
| 67 | _ = json.Unmarshal(request.Params, ¶ms) |
| 68 | switch params.Name { |
| 69 | case "load_toolset": |
| 70 | loaded.Store(true) |
| 71 | notifyChanged = true |
| 72 | result = map[string]any{"content": []map[string]any{{"type": "text", "text": "loaded"}}} |
| 73 | case "list_schematic_components": |
| 74 | dynamicCalls.Add(1) |
| 75 | result = map[string]any{"content": []map[string]any{{"type": "text", "text": "R1"}}} |
| 76 | } |
| 77 | } |
| 78 | |
| 79 | response, _ := json.Marshal(map[string]any{"jsonrpc": "2.0", "id": *request.ID, "result": result}) |
| 80 | if notifyChanged { |
| 81 | w.Header().Set("Content-Type", "text/event-stream") |
| 82 | notification, _ := json.Marshal(map[string]any{ |
| 83 | "jsonrpc": "2.0", "method": "notifications/tools/list_changed", |
| 84 | }) |
| 85 | _, _ = fmt.Fprintf(w, "event: message\ndata: %s\n\nevent: message\ndata: %s\n\n", notification, response) |
| 86 | return |
| 87 | } |
| 88 | w.Header().Set("Content-Type", "application/json") |
| 89 | _, _ = w.Write(response) |
| 90 | })) |
| 91 | } |
| 92 | |
| 93 | func TestMCPCapabilityRuntimeRefreshesDynamicToolsInSession(t *testing.T) { |
| 94 | t.Setenv("REASONIX_CACHE_HOME", t.TempDir()) |
| 95 | ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) |
| 96 | defer cancel() |
| 97 | |
| 98 | var loaded atomic.Bool |
| 99 | var dynamicCalls atomic.Int32 |
| 100 | server := dynamicToolsMCPServer(t, &loaded, &dynamicCalls) |
| 101 | defer server.Close() |
| 102 | |
| 103 | host := plugin.NewHost() |
| 104 | defer host.Close() |
| 105 | registry := tool.NewRegistry() |
| 106 | spec := plugin.Spec{Name: "dynamic", Type: "http", URL: server.URL, Authorized: true} |
| 107 | runtime := NewMCPCapabilityRuntime(ctx, host, []plugin.Spec{spec}, registry, nil) |
| 108 | frontend := runtime.NewFrontend(capability.NewLedger(), nil) |
| 109 | registry.Add(frontend) |
| 110 | initial, err := host.Add(ctx, spec) |
| 111 | if err != nil { |
| 112 | t.Fatalf("Host.Add: %v", err) |
| 113 | } |
| 114 | for _, candidate := range initial { |
| 115 | registry.Add(candidate) |
| 116 | } |
| 117 | registry.SetProviderVisibleTools([]string{"use_capability"}) |
| 118 | providerSchemasBefore, err := json.Marshal(registry.Schemas()) |
| 119 | if err != nil { |
| 120 | t.Fatalf("marshal provider schemas before refresh: %v", err) |
| 121 | } |
| 122 | |
| 123 | if _, err := frontend.Execute(ctx, json.RawMessage(`{"action":"call","capability_id":"mcp-tool:dynamic/load_toolset","arguments":{}}`)); err != nil { |
| 124 | t.Fatalf("load_toolset: %v", err) |
| 125 | } |
| 126 | |
| 127 | wantName := "mcp__dynamic__list_schematic_components" |
| 128 | deadline := time.Now().Add(2 * time.Second) |
| 129 | var live []plugin.CachedTool |
| 130 | for time.Now().Before(deadline) { |
| 131 | live = runtime.ConnectedProxyTools()["dynamic"] |
| 132 | if hasCachedTool(live, "list_schematic_components") { |
| 133 | break |
| 134 | } |
| 135 | time.Sleep(10 * time.Millisecond) |
| 136 | } |
| 137 | if _, ok := registry.Get(wantName); !ok { |
| 138 | t.Fatal("dynamic MCP tool was not registered for use_capability routing") |
| 139 | } |
| 140 | providerSchemasAfter, err := json.Marshal(registry.Schemas()) |
| 141 | if err != nil { |
| 142 | t.Fatalf("marshal provider schemas after refresh: %v", err) |
| 143 | } |
| 144 | if !bytes.Equal(providerSchemasAfter, providerSchemasBefore) { |
| 145 | t.Fatalf("provider-visible schema bytes changed after dynamic MCP refresh: before=%s after=%s", providerSchemasBefore, providerSchemasAfter) |
| 146 | } |
| 147 | if len(live) != 2 || !hasCachedTool(live, "list_schematic_components") { |
| 148 | t.Fatalf("live capability tools = %+v, want refreshed dynamic tool", live) |
| 149 | } |
| 150 | if _, err := frontend.Execute(ctx, json.RawMessage(`{"action":"call","capability_id":"mcp-tool:dynamic/list_schematic_components","arguments":{}}`)); err != nil { |
| 151 | t.Fatalf("dynamic tool call: %v", err) |
| 152 | } |
| 153 | if got := dynamicCalls.Load(); got != 1 { |
| 154 | t.Fatalf("dynamic tools/call count = %d, want 1", got) |
| 155 | } |
| 156 | } |
| 157 | |
| 158 | func TestMCPCapabilityRuntimeReplaysCatalogChangedBeforeSubscription(t *testing.T) { |
| 159 | t.Setenv("REASONIX_CACHE_HOME", t.TempDir()) |
| 160 | ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) |
| 161 | defer cancel() |
| 162 | |
| 163 | var loaded atomic.Bool |
| 164 | var dynamicCalls atomic.Int32 |
| 165 | server := dynamicToolsMCPServer(t, &loaded, &dynamicCalls) |
| 166 | defer server.Close() |
| 167 | |
| 168 | host := plugin.NewHost() |
| 169 | defer host.Close() |
| 170 | registry := tool.NewRegistry() |
| 171 | spec := plugin.Spec{Name: "dynamic", Type: "http", URL: server.URL, Authorized: true} |
| 172 | initial, err := host.Add(ctx, spec) |
| 173 | if err != nil { |
| 174 | t.Fatalf("Host.Add: %v", err) |
| 175 | } |
| 176 | for _, candidate := range initial { |
| 177 | registry.Add(candidate) |
| 178 | } |
| 179 | |
| 180 | changed := make(chan []tool.Tool, 1) |
| 181 | unsubscribe := host.SubscribeToolListChanges(ctx, func(changedSpec plugin.Spec, tools []tool.Tool) { |
| 182 | if plugin.MCPRuntimeSpecMatches(changedSpec, spec) { |
| 183 | changed <- tools |
| 184 | } |
| 185 | }) |
| 186 | loader := initial[0] |
| 187 | if _, err := loader.Execute(ctx, json.RawMessage(`{}`)); err != nil { |
| 188 | t.Fatalf("load_toolset before runtime subscription: %v", err) |
| 189 | } |
| 190 | select { |
| 191 | case refreshed := <-changed: |
| 192 | if findMCPTool(refreshed, "list_schematic_components", "") == nil { |
| 193 | t.Fatalf("refreshed tools missing dynamic tool: %v", refreshed) |
| 194 | } |
| 195 | case <-ctx.Done(): |
| 196 | t.Fatalf("wait for pre-subscription refresh: %v", ctx.Err()) |
| 197 | } |
| 198 | unsubscribe() |
| 199 | |
| 200 | runtime := NewMCPCapabilityRuntime(ctx, host, []plugin.Spec{spec}, registry, nil) |
| 201 | frontend := runtime.NewFrontend(capability.NewLedger(), nil) |
| 202 | registry.Add(frontend) |
| 203 | registry.SetProviderVisibleTools([]string{"use_capability"}) |
| 204 | |
| 205 | wantName := "mcp__dynamic__list_schematic_components" |
| 206 | if _, ok := registry.Get(wantName); !ok { |
| 207 | t.Fatal("late runtime subscription did not replay the current dynamic tool catalog") |
| 208 | } |
| 209 | currentLoader, ok := registry.Get("mcp__dynamic__load_toolset") |
| 210 | if !ok || currentLoader == loader { |
| 211 | t.Fatal("late runtime subscription retained the stale pre-refresh adapter") |
| 212 | } |
| 213 | if live := runtime.ConnectedProxyTools()["dynamic"]; len(live) != 2 || !hasCachedTool(live, "list_schematic_components") { |
| 214 | t.Fatalf("replayed capability tools = %+v, want the complete current catalog", live) |
| 215 | } |
| 216 | if _, err := frontend.Execute(ctx, json.RawMessage(`{"action":"call","capability_id":"mcp-tool:dynamic/list_schematic_components","arguments":{}}`)); err != nil { |
| 217 | t.Fatalf("dynamic tool call after replay: %v", err) |
| 218 | } |
| 219 | if got := dynamicCalls.Load(); got != 1 { |
| 220 | t.Fatalf("dynamic tools/call count = %d, want 1", got) |
| 221 | } |
| 222 | } |
| 223 | |
| 224 | func TestConfiguredDisabledSessionDropsSharedHostReplayAndGenericAlias(t *testing.T) { |
| 225 | t.Setenv("REASONIX_CACHE_HOME", t.TempDir()) |
| 226 | ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) |
| 227 | defer cancel() |
| 228 | |
| 229 | var toolCalls atomic.Int32 |
| 230 | server := explicitReaderMCPServer(t, nil, &toolCalls) |
| 231 | defer server.Close() |
| 232 | spec := plugin.Spec{Name: "shared-disabled", Type: "http", URL: server.URL, Authorized: true} |
| 233 | host := plugin.NewHost() |
| 234 | defer host.Close() |
| 235 | initial, err := host.Add(ctx, spec) |
| 236 | if err != nil { |
| 237 | t.Fatalf("Host.Add: %v", err) |
| 238 | } |
| 239 | |
| 240 | registry := tool.NewRegistry() |
| 241 | runtime := NewMCPCapabilityRuntime(ctx, host, []plugin.Spec{spec}, registry, nil) |
| 242 | modelName := plugin.ModelToolName(spec.Name, "search") |
| 243 | if _, ok := registry.Get(modelName); !ok { |
| 244 | t.Fatal("test requires constructor replay to register the shared Host tool") |
| 245 | } |
| 246 | frontend := runtime.NewFrontend(capability.NewLedger(), nil) |
| 247 | generic := json.RawMessage(fmt.Sprintf(`{"action":"call","capability_id":"tool:%s","arguments":{}}`, modelName)) |
| 248 | resolvedGeneric, err := frontend.ResolveCall(ctx, generic) |
| 249 | if err != nil || resolvedGeneric.Target == nil { |
| 250 | t.Fatalf("resolve generic replayed adapter = %+v, %v", resolvedGeneric, err) |
| 251 | } |
| 252 | runtime.ConfigureServers( |
| 253 | []config.PluginEntry{{Name: spec.Name, Type: spec.Type, URL: spec.URL, Source: config.MCPSourceUserConfig}}, |
| 254 | []plugin.Spec{spec}, |
| 255 | map[string]bool{spec.Name: false}, |
| 256 | ) |
| 257 | if _, ok := registry.Get(modelName); ok { |
| 258 | t.Fatal("disabled session retained the replayed shared Host adapter") |
| 259 | } |
| 260 | if _, err := resolvedGeneric.Target.Execute(ctx, resolvedGeneric.Args); err == nil || !strings.Contains(strings.ToLower(err.Error()), "disabled") { |
| 261 | t.Fatalf("resolved generic adapter after disable error = %v", err) |
| 262 | } |
| 263 | canonical := json.RawMessage(`{"action":"call","capability_id":"mcp-tool:shared-disabled/search","arguments":{}}`) |
| 264 | out, err := frontend.Execute(ctx, canonical) |
| 265 | if detail := strings.ToLower(out + " " + fmt.Sprint(err)); !strings.Contains(detail, "disabled") { |
| 266 | t.Fatalf("canonical disabled call = %q, %v, want disabled refusal", out, err) |
| 267 | } |
| 268 | |
| 269 | // Even if another registry owner retains an adapter, generic tool: routing |
| 270 | // must re-check this runtime's current authorization boundary. |
| 271 | staleRegistry := tool.NewRegistry() |
| 272 | staleRegistry.Add(initial[0]) |
| 273 | staleFrontend := runtime.NewFrontend(capability.NewLedger(), nil) |
| 274 | staleFrontend.registry = staleRegistry |
| 275 | out, err = staleFrontend.Execute(ctx, generic) |
| 276 | if detail := strings.ToLower(out + " " + fmt.Sprint(err)); !strings.Contains(detail, "disabled") { |
| 277 | t.Fatalf("generic disabled call = %q, %v, want disabled refusal", out, err) |
| 278 | } |
| 279 | if got := toolCalls.Load(); got != 0 { |
| 280 | t.Fatalf("disabled session executed tools/call %d times", got) |
| 281 | } |
| 282 | } |
| 283 | |
| 284 | func TestMCPResolveReleasesDispatchLockBeforeCatalogCallback(t *testing.T) { |
| 285 | t.Setenv("REASONIX_CACHE_HOME", t.TempDir()) |
| 286 | ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) |
| 287 | defer cancel() |
| 288 | |
| 289 | catalogEntered := make(chan struct{}) |
| 290 | releaseCatalog := make(chan struct{}) |
| 291 | var releaseOnce sync.Once |
| 292 | release := func() { releaseOnce.Do(func() { close(releaseCatalog) }) } |
| 293 | t.Cleanup(release) |
| 294 | spec := plugin.Spec{Name: "catalog-lock", Type: "http", URL: "http://127.0.0.1:1", Authorized: true} |
| 295 | host := plugin.NewHost() |
| 296 | defer host.Close() |
| 297 | runtime := NewMCPCapabilityRuntime(ctx, host, []plugin.Spec{spec}, tool.NewRegistry(), func() capability.Catalog { |
| 298 | close(catalogEntered) |
| 299 | <-releaseCatalog |
| 300 | return capability.Catalog{} |
| 301 | }) |
| 302 | frontend := runtime.NewFrontend(capability.NewLedger(), nil) |
| 303 | |
| 304 | resolveDone := make(chan struct { |
| 305 | call tool.ResolvedCall |
| 306 | err error |
| 307 | }, 1) |
| 308 | go func() { |
| 309 | resolved, err := frontend.ResolveCall(ctx, json.RawMessage(`{"action":"call","capability_id":"mcp-tool:catalog-lock/search","arguments":{}}`)) |
| 310 | resolveDone <- struct { |
| 311 | call tool.ResolvedCall |
| 312 | err error |
| 313 | }{call: resolved, err: err} |
| 314 | }() |
| 315 | select { |
| 316 | case <-catalogEntered: |
| 317 | case <-ctx.Done(): |
| 318 | t.Fatalf("catalog callback was not entered: %v", ctx.Err()) |
| 319 | } |
| 320 | |
| 321 | disableDone := make(chan bool, 1) |
| 322 | go func() { disableDone <- runtime.SetServerEnabled(spec.Name, false) }() |
| 323 | select { |
| 324 | case ok := <-disableDone: |
| 325 | if !ok { |
| 326 | t.Fatal("disable did not find the configured server") |
| 327 | } |
| 328 | case <-time.After(time.Second): |
| 329 | release() |
| 330 | <-disableDone |
| 331 | t.Fatal("runtime writer blocked behind catalog callback; resolve retained dispatchMu across catalog lookup") |
| 332 | } |
| 333 | release() |
| 334 | resolved := <-resolveDone |
| 335 | if resolved.err != nil || resolved.call.Target == nil { |
| 336 | t.Fatalf("resolve = %+v, %v", resolved.call, resolved.err) |
| 337 | } |
| 338 | if _, err := resolved.call.Target.Execute(ctx, resolved.call.Args); err == nil || !strings.Contains(strings.ToLower(err.Error()), "disabled") { |
| 339 | t.Fatalf("resolved target after concurrent disable error = %v", err) |
| 340 | } |
| 341 | } |
| 342 | |
| 343 | func hasCachedTool(tools []plugin.CachedTool, name string) bool { |
| 344 | return slices.ContainsFunc(tools, func(candidate plugin.CachedTool) bool { |
| 345 | return candidate.Name == name |
| 346 | }) |
| 347 | } |
| 348 |