| 1 | package agent |
| 2 | |
| 3 | import ( |
| 4 | "context" |
| 5 | "encoding/json" |
| 6 | "net/http" |
| 7 | "net/http/httptest" |
| 8 | "sync/atomic" |
| 9 | "testing" |
| 10 | "time" |
| 11 | |
| 12 | "reasonix/internal/capability" |
| 13 | "reasonix/internal/config" |
| 14 | "reasonix/internal/plugin" |
| 15 | "reasonix/internal/provider" |
| 16 | "reasonix/internal/tool" |
| 17 | ) |
| 18 | |
| 19 | func TestPartitionToolCallsParallelisesCapabilityDiscovery(t *testing.T) { |
| 20 | reg := tool.NewRegistry() |
| 21 | reg.Add(NewUseCapabilityTool(t.Context(), nil, nil, reg, nil, nil, nil)) |
| 22 | reg.Add(fakeTool{name: "read_file", readOnly: true}) |
| 23 | calls := []provider.ToolCall{ |
| 24 | {ID: "1", Name: "use_capability", Arguments: `{"action":"search","query":"github"}`}, |
| 25 | {ID: "2", Name: "use_capability", Arguments: `{"action":"inspect","capability_id":"mcp-server:github"}`}, |
| 26 | {ID: "3", Name: "read_file", Arguments: `{"path":"a.go"}`}, |
| 27 | } |
| 28 | got := partitionToolCalls(reg, calls) |
| 29 | if len(got) != 1 || !got[0].parallel || got[0].end != 3 { |
| 30 | t.Fatalf("partition = %+v, want one parallel batch of 3", got) |
| 31 | } |
| 32 | } |
| 33 | |
| 34 | func TestPartitionToolCallsKeepsUnknownCapabilitySerial(t *testing.T) { |
| 35 | reg := tool.NewRegistry() |
| 36 | reg.Add(NewUseCapabilityTool(t.Context(), nil, nil, reg, nil, nil, nil)) |
| 37 | calls := []provider.ToolCall{ |
| 38 | {ID: "1", Name: "use_capability", Arguments: `{"action":"call","capability_id":"mcp-tool:unknown/write","arguments":{"x":1}}`}, |
| 39 | } |
| 40 | got := partitionToolCalls(reg, calls) |
| 41 | if len(got) != 1 || got[0].parallel { |
| 42 | t.Fatalf("unknown MCP call must stay serial: %+v", got) |
| 43 | } |
| 44 | } |
| 45 | |
| 46 | func TestCapabilityWrappersForwardPureBatchClassification(t *testing.T) { |
| 47 | inner := NewUseCapabilityTool(t.Context(), nil, nil, nil, nil, nil, nil) |
| 48 | pathBound := pathBoundCapabilityProxy{inner: inner, resolver: inner} |
| 49 | restricted := &restrictedCapabilityProxy{ |
| 50 | Tool: inner, resolver: inner, |
| 51 | allowed: map[string]bool{"mcp-server:github": true}, servers: map[string]bool{"github": true}, |
| 52 | } |
| 53 | for name, wrapped := range map[string]tool.Tool{"path-bound": pathBound, "restricted": restricted} { |
| 54 | t.Run(name, func(t *testing.T) { |
| 55 | reg := tool.NewRegistry() |
| 56 | reg.Add(wrapped) |
| 57 | got := partitionToolCalls(reg, []provider.ToolCall{{ |
| 58 | ID: "1", Name: "use_capability", Arguments: `{"action":"inspect","capability_id":"mcp-server:github"}`, |
| 59 | }}) |
| 60 | if len(got) != 1 || !got[0].parallel { |
| 61 | t.Fatalf("wrapped inspect must remain parallel: %+v", got) |
| 62 | } |
| 63 | serial := partitionToolCalls(reg, []provider.ToolCall{{ |
| 64 | ID: "2", Name: "use_capability", Arguments: `{"action":"call","capability_id":"mcp-tool:github/write","arguments":{}}`, |
| 65 | }}) |
| 66 | if len(serial) != 1 || serial[0].parallel { |
| 67 | t.Fatalf("unknown writer must remain serial: %+v", serial) |
| 68 | } |
| 69 | }) |
| 70 | } |
| 71 | } |
| 72 | |
| 73 | func TestClassifyCallSearchIsReadOnlyParallel(t *testing.T) { |
| 74 | proxy := NewUseCapabilityTool(t.Context(), nil, nil, nil, nil, nil, nil) |
| 75 | class := proxy.ClassifyCall(json.RawMessage(`{"action":"search","query":"x"}`)) |
| 76 | if !class.Known || !class.ReadOnly || !class.ParallelSafe { |
| 77 | t.Fatalf("search class = %+v", class) |
| 78 | } |
| 79 | } |
| 80 | |
| 81 | func TestPartitionIndependentReadOnlyMCPCallsAreParallel(t *testing.T) { |
| 82 | t.Setenv("REASONIX_CACHE_HOME", t.TempDir()) |
| 83 | var calls atomic.Int32 |
| 84 | alpha := readonlyMCPServer(t, "alpha", &calls) |
| 85 | beta := readonlyMCPServer(t, "beta", &calls) |
| 86 | defer alpha.Close() |
| 87 | defer beta.Close() |
| 88 | |
| 89 | ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) |
| 90 | defer cancel() |
| 91 | host := plugin.NewHost() |
| 92 | defer host.Close() |
| 93 | specs := []plugin.Spec{ |
| 94 | {Name: "alpha", Type: "http", URL: alpha.URL, Authorized: true}, |
| 95 | {Name: "beta", Type: "http", URL: beta.URL, Authorized: true}, |
| 96 | } |
| 97 | if _, err := host.Add(ctx, specs[0]); err != nil { |
| 98 | t.Fatalf("connect alpha: %v", err) |
| 99 | } |
| 100 | if _, err := host.Add(ctx, specs[1]); err != nil { |
| 101 | t.Fatalf("connect beta: %v", err) |
| 102 | } |
| 103 | runtime := NewMCPCapabilityRuntime(ctx, host, specs, tool.NewRegistry(), nil) |
| 104 | runtime.ConfigureServers([]config.PluginEntry{{Name: "alpha"}, {Name: "beta"}}, specs, map[string]bool{"alpha": true, "beta": true}) |
| 105 | proxy := runtime.NewFrontend(capability.NewLedger(), nil) |
| 106 | reg := tool.NewRegistry() |
| 107 | reg.Add(proxy) |
| 108 | got := partitionToolCalls(reg, []provider.ToolCall{ |
| 109 | {ID: "1", Name: "use_capability", Arguments: `{"action":"call","capability_id":"mcp-tool:alpha/search","arguments":{}}`}, |
| 110 | {ID: "2", Name: "use_capability", Arguments: `{"action":"call","capability_id":"mcp-tool:beta/search","arguments":{}}`}, |
| 111 | }) |
| 112 | if len(got) != 1 || !got[0].parallel || got[0].end != 2 { |
| 113 | t.Fatalf("independent read-only MCP partition = %+v, want one parallel batch", got) |
| 114 | } |
| 115 | |
| 116 | serial := partitionToolCalls(reg, []provider.ToolCall{ |
| 117 | {ID: "w", Name: "use_capability", Arguments: `{"action":"call","capability_id":"mcp-tool:unknown/write","arguments":{}}`}, |
| 118 | }) |
| 119 | if len(serial) != 1 || serial[0].parallel { |
| 120 | t.Fatalf("unknown/write MCP must stay serial: %+v", serial) |
| 121 | } |
| 122 | } |
| 123 | |
| 124 | func TestPartitionStatefulBrowserMCPStaysSerial(t *testing.T) { |
| 125 | proxy := NewUseCapabilityTool(t.Context(), nil, nil, nil, nil, nil, nil) |
| 126 | proxy.runtime = &MCPCapabilityRuntime{servers: map[string]mcpRuntimeServer{ |
| 127 | "chrome-devtools": {entry: config.PluginEntry{Name: "chrome-devtools"}}, |
| 128 | }} |
| 129 | class := proxy.ClassifyCall(json.RawMessage(`{"action":"call","capability_id":"mcp-tool:chrome-devtools/navigate","arguments":{}}`)) |
| 130 | if class.ParallelSafe { |
| 131 | t.Fatalf("stateful browser MCP must not be parallel-safe: %+v", class) |
| 132 | } |
| 133 | } |
| 134 | |
| 135 | func TestOnDemandConnectEmitsOneSessionRemoteToolsListObservation(t *testing.T) { |
| 136 | t.Setenv("REASONIX_CACHE_HOME", t.TempDir()) |
| 137 | var toolCalls atomic.Int32 |
| 138 | server := readonlyMCPServer(t, "observed", &toolCalls) |
| 139 | defer server.Close() |
| 140 | host := plugin.NewHost() |
| 141 | defer host.Close() |
| 142 | spec := plugin.Spec{Name: "observed", Type: "http", URL: server.URL, Authorized: true} |
| 143 | proxy := NewUseCapabilityTool(t.Context(), host, []plugin.Spec{spec}, tool.NewRegistry(), nil, nil, nil) |
| 144 | var observations []mcpListObservation |
| 145 | proxy.bindMCPListObserver(func(observation mcpListObservation) { observations = append(observations, observation) }) |
| 146 | if _, err := proxy.ensureServerToolsForSpec(t.Context(), spec.Name, spec); err != nil { |
| 147 | t.Fatalf("first connect: %v", err) |
| 148 | } |
| 149 | if len(observations) != 1 || observations[0].Source != "remote" || observations[0].Trigger != "connect" || !observations[0].NetworkCall || observations[0].ToolCount != 1 { |
| 150 | t.Fatalf("observations = %+v", observations) |
| 151 | } |
| 152 | if _, err := proxy.ensureServerToolsForSpec(t.Context(), spec.Name, spec); err != nil { |
| 153 | t.Fatalf("shared-host reuse: %v", err) |
| 154 | } |
| 155 | if len(observations) != 1 { |
| 156 | t.Fatalf("shared-host reuse emitted a remote list: %+v", observations) |
| 157 | } |
| 158 | } |
| 159 | |
| 160 | func TestListChangedIsAttributedOnlyToActiveRuntimeFrontend(t *testing.T) { |
| 161 | runtime := NewMCPCapabilityRuntime(t.Context(), nil, nil, tool.NewRegistry(), nil) |
| 162 | audit := &capability.Audit{} |
| 163 | frontend := runtime.NewFrontend(nil, audit) |
| 164 | var observations []mcpListObservation |
| 165 | frontend.bindMCPListObserver(func(observation mcpListObservation) { observations = append(observations, observation) }) |
| 166 | release := frontend.activateMCPListObserver() |
| 167 | runtime.notifyToolListChanged("svc", []tool.Tool{fakeTool{name: "mcp__svc__read", readOnly: true}}) |
| 168 | if len(observations) != 1 || observations[0].Trigger != "list_changed" || !observations[0].NetworkCall { |
| 169 | t.Fatalf("observations = %+v", observations) |
| 170 | } |
| 171 | snapshot := audit.Snapshot() |
| 172 | if snapshot.MCPLists.Remote != 1 || snapshot.MCPLists.Triggers["list_changed"] != 1 { |
| 173 | t.Fatalf("MCP list audit = %+v", snapshot.MCPLists) |
| 174 | } |
| 175 | release() |
| 176 | runtime.notifyToolListChanged("svc", []tool.Tool{fakeTool{name: "mcp__svc__read", readOnly: true}}) |
| 177 | if len(observations) != 1 || audit.Snapshot().MCPLists.Remote != 1 { |
| 178 | t.Fatalf("inactive frontend retained list_changed attribution: observations=%+v audit=%+v", observations, audit.Snapshot().MCPLists) |
| 179 | } |
| 180 | } |
| 181 | |
| 182 | func readonlyMCPServer(t *testing.T, name string, calls *atomic.Int32) *httptest.Server { |
| 183 | t.Helper() |
| 184 | return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { |
| 185 | var request struct { |
| 186 | ID *int `json:"id"` |
| 187 | Method string `json:"method"` |
| 188 | } |
| 189 | if err := json.NewDecoder(r.Body).Decode(&request); err != nil { |
| 190 | http.Error(w, "bad request", http.StatusBadRequest) |
| 191 | return |
| 192 | } |
| 193 | if request.ID == nil { |
| 194 | w.WriteHeader(http.StatusAccepted) |
| 195 | return |
| 196 | } |
| 197 | var result any |
| 198 | switch request.Method { |
| 199 | case "initialize": |
| 200 | result = map[string]any{"protocolVersion": "2024-11-05", "serverInfo": map[string]any{"name": name, "version": "1"}} |
| 201 | case "tools/list": |
| 202 | result = map[string]any{"tools": []map[string]any{{ |
| 203 | "name": "search", "description": "search", |
| 204 | "inputSchema": map[string]any{"type": "object"}, |
| 205 | "annotations": map[string]any{"readOnlyHint": true}, |
| 206 | }}} |
| 207 | case "tools/call": |
| 208 | calls.Add(1) |
| 209 | result = map[string]any{"content": []map[string]any{{"type": "text", "text": "ok"}}} |
| 210 | } |
| 211 | w.Header().Set("Content-Type", "application/json") |
| 212 | _ = json.NewEncoder(w).Encode(map[string]any{"jsonrpc": "2.0", "id": *request.ID, "result": result}) |
| 213 | })) |
| 214 | } |
| 215 |