1package agent
2
3import (
4 "context"
5 "encoding/json"
6 "net/http"
7 "net/http/httptest"
8 "os"
9 "path/filepath"
10 "slices"
11 "strings"
12 "testing"
13
14 "github.com/bnema/zut/packages/agent/mcp"
15 "github.com/bnema/zut/packages/agent/subagents"
16 "github.com/bnema/zut/packages/agent/tools"
17 "github.com/bnema/zut/packages/core"
18)
19
20// fakeMCPHTTPServer answers the minimal streamable-HTTP MCP surface with one
21// read-only tool ("lookup") and one mutating tool ("deploy").
22func fakeMCPHTTPServer(t *testing.T) *httptest.Server {
23 t.Helper()
24 srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
25 if r.Method != http.MethodPost {
26 w.WriteHeader(http.StatusMethodNotAllowed)
27 return
28 }
29 var req struct {
30 ID json.RawMessage `json:"id"`
31 Method string `json:"method"`
32 Params json.RawMessage `json:"params"`
33 }
34 if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
35 w.WriteHeader(http.StatusBadRequest)
36 return
37 }
38 if len(req.ID) == 0 {
39 w.WriteHeader(http.StatusAccepted)
40 return
41 }
42 var result any
43 switch req.Method {
44 case "initialize":
45 result = map[string]any{
46 "protocolVersion": "2025-06-18",
47 "capabilities": map[string]any{"tools": map[string]any{}},
48 "serverInfo": map[string]any{"name": "fake", "version": "1"},
49 "instructions": "Use lookup before deploy.",
50 }
51 case "tools/list":
52 result = map[string]any{"tools": []any{
53 map[string]any{"name": "lookup", "description": "Look something up.", "inputSchema": map[string]any{"type": "object"}, "annotations": map[string]any{"readOnlyHint": true}},
54 map[string]any{"name": "deploy", "description": "Deploy.", "inputSchema": map[string]any{"type": "object"}},
55 }}
56 case "tools/call":
57 result = map[string]any{"content": []any{map[string]any{"type": "text", "text": "called " + string(req.Params)}}}
58 default:
59 result = map[string]any{}
60 }
61 w.Header().Set("Content-Type", "application/json")
62 _ = json.NewEncoder(w).Encode(map[string]any{"jsonrpc": "2.0", "id": req.ID, "result": result})
63 }))
64 t.Cleanup(srv.Close)
65 return srv
66}
67
68func setupMCPHome(t *testing.T, url string) (project string) {
69 t.Helper()
70 root := t.TempDir()
71 zutHome := filepath.Join(root, "zut-home")
72 home := filepath.Join(root, "home")
73 project = filepath.Join(root, "project")
74 t.Setenv("ZUT_HOME", zutHome)
75 t.Setenv("HOME", home)
76 t.Setenv("USERPROFILE", home)
77 for _, dir := range []string{zutHome, home, project} {
78 if err := os.MkdirAll(dir, 0o755); err != nil {
79 t.Fatal(err)
80 }
81 }
82 cfg := mcp.Config{Servers: map[string]mcp.ServerConfig{"fake": {Type: "http", URL: url}}}
83 if err := mcp.WriteConfigFile(mcp.GlobalConfigPath(zutHome), cfg); err != nil {
84 t.Fatal(err)
85 }
86 return project
87}
88
89func startTestMCP(t *testing.T, args *Args) {
90 t.Helper()
91 stop, err := startMCP(context.Background(), args, "test")
92 if err != nil {
93 t.Fatal(err)
94 }
95 t.Cleanup(stop)
96 if args.mcpManager == nil {
97 t.Fatalf("mcp manager not started: %v", args.mcpStartupErrors)
98 }
99}
100
101func TestResolveAddsMCPToolsAndPrompt(t *testing.T) {
102 project := setupMCPHome(t, fakeMCPHTTPServer(t).URL)
103 args := Args{Provider: "ollama", Model: "any-local-model", CWD: project, NoLSP: true, NoSkill: true, Mode: ModePrint}
104 startTestMCP(t, &args)
105
106 r, err := Resolve(args, false)
107 if err != nil {
108 t.Fatal(err)
109 }
110 for _, name := range []string{"mcp__fake__lookup", "mcp__fake__deploy"} {
111 if _, ok := r.ToolRegistry[name]; !ok {
112 t.Fatalf("registry missing %s", name)
113 }
114 }
115 if !strings.Contains(r.SystemPrompt, "[mcp_servers]") || !strings.Contains(r.SystemPrompt, "Use lookup before deploy.") {
116 t.Fatalf("system prompt lacks MCP addendum:\n%s", r.SystemPrompt)
117 }
118 res, err := r.ToolRegistry["mcp__fake__lookup"].Execute(context.Background(), json.RawMessage(`{"q":"x"}`), nil)
119 if err != nil || res.IsError {
120 t.Fatalf("execute = %+v, %v", res, err)
121 }
122}
123
124func TestMCPToolSelection(t *testing.T) {
125 project := setupMCPHome(t, fakeMCPHTTPServer(t).URL)
126 base := Args{Provider: "ollama", Model: "any-local-model", CWD: project, NoLSP: true, NoSkill: true, Mode: ModePrint}
127 startTestMCP(t, &base)
128
129 for _, tc := range []struct {
130 name string
131 edit func(*Args)
132 wants []string
133 drops []string
134 }{
135 {name: "selector", edit: func(a *Args) { a.Tools, a.ToolsSet = []string{"read", "mcp"}, true }, wants: []string{"mcp__fake__lookup", "mcp__fake__deploy"}},
136 {name: "exact", edit: func(a *Args) { a.Tools, a.ToolsSet = []string{"mcp__fake__lookup"}, true }, wants: []string{"mcp__fake__lookup"}, drops: []string{"mcp__fake__deploy"}},
137 {name: "not listed", edit: func(a *Args) { a.Tools, a.ToolsSet = []string{"read"}, true }, drops: []string{"mcp__fake__lookup", "mcp__fake__deploy"}},
138 {name: "no-mcp", edit: func(a *Args) { a.NoMCP = true }, drops: []string{"mcp__fake__lookup"}},
139 {name: "packaged agent", edit: func(a *Args) { a.PermissionSet = &tools.PermissionSet{} }, drops: []string{"mcp__fake__lookup"}},
140 } {
141 t.Run(tc.name, func(t *testing.T) {
142 args := base
143 tc.edit(&args)
144 reg := core.Registry{}
145 addMCPTools(reg, args)
146 for _, name := range tc.wants {
147 if _, ok := reg[name]; !ok {
148 t.Fatalf("missing %s in %v", name, keys(reg))
149 }
150 }
151 for _, name := range tc.drops {
152 if _, ok := reg[name]; ok {
153 t.Fatalf("unexpected %s", name)
154 }
155 }
156 })
157 }
158}
159
160func TestOrchestratorHoldsReadOnlyMCPToolsAndWorkersGetAll(t *testing.T) {
161 project := setupMCPHome(t, fakeMCPHTTPServer(t).URL)
162 args := Args{Provider: "ollama", Model: "any-local-model", CWD: project, NoLSP: true, NoSkill: true, Mode: ModePrint, Orchestrate: true}
163 startTestMCP(t, &args)
164
165 r, err := Resolve(args, false)
166 if err != nil {
167 t.Fatal(err)
168 }
169 if _, ok := r.ToolRegistry["mcp__fake__lookup"]; !ok {
170 t.Fatal("orchestrator lacks read-only MCP tool")
171 }
172 if _, ok := r.ToolRegistry["mcp__fake__deploy"]; ok {
173 t.Fatal("orchestrator holds a mutating MCP tool")
174 }
175 narrowed := args
176 narrowed.Tools, narrowed.ToolsSet = []string{"subagent", "mcp__fake__deploy"}, true
177 narrowReg := core.Registry{}
178 addMCPTools(narrowReg, narrowed)
179 if len(narrowReg) != 0 {
180 t.Fatalf("orchestrator --tools selection ignored: %v", keys(narrowReg))
181 }
182 if !strings.Contains(r.SystemPrompt, "delegate MCP actions") {
183 t.Fatalf("orchestrator prompt lacks MCP delegation guidance:\n%s", r.SystemPrompt)
184 }
185
186 runtime := newSubagentRuntime(subagentRuntimeConfig{
187 Args: args, Root: t.TempDir(), RepoRoot: project, Provider: "ollama", Model: "any-local-model",
188 })
189 t.Cleanup(func() { _ = runtime.Close(context.Background()) })
190 spec, err := runtime.buildResidentChildSpec(context.Background(), tools.ResidentSpawnRequest{Task: "deploy"}, r.ToolRegistry)
191 if err != nil {
192 t.Fatal(err)
193 }
194 for _, name := range []string{"mcp__fake__lookup", "mcp__fake__deploy"} {
195 if !slices.Contains(spec.Tools, name) {
196 t.Fatalf("worker tools %v missing %s", spec.Tools, name)
197 }
198 }
199
200 childArgs := residentChildArgs(args, "ollama", spec)
201 child, err := Resolve(childArgs, false)
202 if err != nil {
203 t.Fatal(err)
204 }
205 reg, err := residentChildRegistry(child.ToolRegistry, spec.Tools)
206 if err != nil {
207 t.Fatal(err)
208 }
209 if _, ok := reg["mcp__fake__deploy"]; !ok {
210 t.Fatal("worker registry lacks mutating MCP tool")
211 }
212}
213
214func TestResidentChildPolicyMCPSelector(t *testing.T) {
215 project := setupMCPHome(t, fakeMCPHTTPServer(t).URL)
216 args := Args{Provider: "ollama", Model: "any-local-model", CWD: project, NoLSP: true, NoSkill: true}
217 startTestMCP(t, &args)
218 r, err := Resolve(args, false)
219 if err != nil {
220 t.Fatal(err)
221 }
222 runtime := newSubagentRuntime(subagentRuntimeConfig{
223 Args: args, Root: t.TempDir(), RepoRoot: project, Provider: "ollama", Model: "any-local-model",
224 Policy: subagents.SubagentPolicy{AllowedTools: []string{"read", "mcp"}},
225 })
226 t.Cleanup(func() { _ = runtime.Close(context.Background()) })
227 spec, err := runtime.buildResidentChildSpec(context.Background(), tools.ResidentSpawnRequest{Task: "x"}, r.ToolRegistry)
228 if err != nil {
229 t.Fatal(err)
230 }
231 if !slices.Contains(spec.Tools, "mcp__fake__deploy") || slices.Contains(spec.Tools, "write") {
232 t.Fatalf("policy-filtered child tools = %v", spec.Tools)
233 }
234}
235
236func TestResidentChildRegistrySkipsMissingMCPTool(t *testing.T) {
237 reg, err := residentChildRegistry(core.Registry{"read": &tools.ReadTool{}}, []string{"read", "mcp__gone__tool"})
238 if err != nil {
239 t.Fatal(err)
240 }
241 if len(reg) != 1 {
242 t.Fatalf("registry = %v", keys(reg))
243 }
244 if _, err := residentChildRegistry(core.Registry{}, []string{"write"}); err == nil {
245 t.Fatal("missing non-MCP tool must still fail")
246 }
247}
248
249func TestRequiredMCPServerFailureAbortsStart(t *testing.T) {
250 root := t.TempDir()
251 t.Setenv("ZUT_HOME", root)
252 cfg := mcp.Config{Servers: map[string]mcp.ServerConfig{"broken": {Command: filepath.Join(root, "missing-binary"), Required: true}}}
253 if err := mcp.WriteConfigFile(mcp.GlobalConfigPath(root), cfg); err != nil {
254 t.Fatal(err)
255 }
256 args := Args{CWD: root, Mode: ModePrint}
257 if _, err := startMCP(context.Background(), &args, "test"); err == nil || !strings.Contains(err.Error(), "required MCP server") {
258 t.Fatalf("startMCP err = %v", err)
259 }
260 cfg.Servers["broken"] = mcp.ServerConfig{Command: filepath.Join(root, "missing-binary")}
261 if err := mcp.WriteConfigFile(mcp.GlobalConfigPath(root), cfg); err != nil {
262 t.Fatal(err)
263 }
264 args = Args{CWD: root, Mode: ModeInteractive}
265 stop, err := startMCP(context.Background(), &args, "test")
266 if err != nil {
267 t.Fatal(err)
268 }
269 stop()
270 if len(args.mcpStartupErrors) != 1 {
271 t.Fatalf("startup errors = %v", args.mcpStartupErrors)
272 }
273}
274
275func TestMCPCommandAddListRemove(t *testing.T) {
276 root := t.TempDir()
277 t.Setenv("ZUT_HOME", root)
278 t.Setenv("FAKE_TOKEN", "secret-value")
279 var out strings.Builder
280 run := func(args ...string) error {
281 out.Reset()
282 _, err := runMCPCommand(context.Background(), append([]string{"mcp"}, args...), "test", &out)
283 return err
284 }
285 if err := run("add", "local", "--env", "KEY=secret-value", "--", "server-bin", "--flag"); err != nil {
286 t.Fatal(err)
287 }
288 if err := run("add", "remote", "--url", "https://example.test/mcp", "--bearer-token-env-var", "FAKE_TOKEN"); err != nil {
289 t.Fatal(err)
290 }
291 if err := run("add", "bad", "--url", "https://x.test", "--", "cmd"); err == nil {
292 t.Fatal("mixed transport accepted")
293 }
294 if err := run("list"); err != nil {
295 t.Fatal(err)
296 }
297 if !strings.Contains(out.String(), "local\tstdio\tserver-bin --flag") || !strings.Contains(out.String(), "remote\thttp\thttps://example.test/mcp") {
298 t.Fatalf("list output:\n%s", out.String())
299 }
300 if err := run("get", "local"); err != nil {
301 t.Fatal(err)
302 }
303 if strings.Contains(out.String(), "secret-value") {
304 t.Fatalf("get leaked a secret:\n%s", out.String())
305 }
306 if err := run("remove", "local"); err != nil {
307 t.Fatal(err)
308 }
309 file, err := mcp.ReadConfigFile(mcp.GlobalConfigPath(root))
310 if err != nil {
311 t.Fatal(err)
312 }
313 if _, ok := file.Servers["local"]; ok || len(file.Servers) != 1 {
314 t.Fatalf("servers after remove = %v", file.Servers)
315 }
316 if handled, _ := runMCPCommand(context.Background(), []string{"chat"}, "test", &out); handled {
317 t.Fatal("non-mcp args handled")
318 }
319}
320
321func TestMCPCommandListRedactsAndKeepsDisabledEntries(t *testing.T) {
322 root := t.TempDir()
323 t.Setenv("ZUT_HOME", root)
324 t.Setenv("REVIEW_SECRET", "synthetic-review-secret")
325 off := false
326 cfg := mcp.Config{Servers: map[string]mcp.ServerConfig{
327 "local": {Command: "server", Args: []string{"--token", "${REVIEW_SECRET}"}},
328 "remote": {Type: "http", URL: "https://user:pw@example.test/mcp?key=${REVIEW_SECRET}"},
329 "off": {Command: "server", Enabled: &off},
330 }}
331 if err := mcp.WriteConfigFile(mcp.GlobalConfigPath(root), cfg); err != nil {
332 t.Fatal(err)
333 }
334 for _, args := range [][]string{{"list"}, {"list", "--json"}, {"get", "remote"}, {"get", "local"}} {
335 var out strings.Builder
336 if _, err := runMCPCommand(context.Background(), append([]string{"mcp"}, args...), "test", &out); err != nil {
337 t.Fatalf("%v: %v", args, err)
338 }
339 if strings.Contains(out.String(), "synthetic-review-secret") || strings.Contains(out.String(), "pw@") {
340 t.Fatalf("%v leaked a secret:\n%s", args, out.String())
341 }
342 if args[0] == "list" && !strings.Contains(out.String(), "off") {
343 t.Fatalf("%v hides disabled entry:\n%s", args, out.String())
344 }
345 }
346}
347
348func TestRedactURL(t *testing.T) {
349 for in, want := range map[string]string{
350 "https://example.test/mcp": "https://example.test/mcp",
351 "https://u:p@example.test/mcp?key=s": "https://*****@example.test/mcp?*****",
352 "https://${U}:${P}@example.test/mcp?key=literal-secret#frag": "https://*****@example.test/mcp?*****",
353 "${BASE}/mcp?token=literal-secret": "${BASE}/mcp?*****",
354 "https://example.test/a@b": "https://example.test/a@b",
355 } {
356 if got := redactURL(in); got != want {
357 t.Errorf("redactURL(%q) = %q, want %q", in, got, want)
358 }
359 }
360}
361
362func TestMCPCommandFailsOnMalformedConfig(t *testing.T) {
363 root := t.TempDir()
364 t.Setenv("ZUT_HOME", root)
365 if err := os.WriteFile(mcp.GlobalConfigPath(root), []byte("{not json"), 0o600); err != nil {
366 t.Fatal(err)
367 }
368 for _, args := range [][]string{{"list"}, {"list", "--json"}, {"status"}} {
369 var out strings.Builder
370 if _, err := runMCPCommand(context.Background(), append([]string{"mcp"}, args...), "test", &out); err == nil {
371 t.Fatalf("%v succeeded on malformed config:\n%s", args, out.String())
372 }
373 }
374}
375
376func TestParseArgsNoMCP(t *testing.T) {
377 args, err := ParseArgs([]string{"--no-mcp", "-p", "hi"})
378 if err != nil {
379 t.Fatal(err)
380 }
381 if !args.NoMCP || mcpAllowed(args) {
382 t.Fatalf("--no-mcp not applied: %+v", args.NoMCP)
383 }
384}
385
386func keys(reg core.Registry) []string {
387 out := make([]string, 0, len(reg))
388 for name := range reg {
389 out = append(out, name)
390 }
391 slices.Sort(out)
392 return out
393}
394
395// A signal during the MCP handshake must cancel startup (not wait for the
396// startup timeout) so the caller can close the servers already launched.
397func TestStartMCPInterruptibleCancelsHandshake(t *testing.T) {
398 reached := make(chan struct{}, 1)
399 release := make(chan struct{})
400 srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
401 select {
402 case reached <- struct{}{}:
403 default:
404 }
405 // Stall every request until the client gives up or the test ends.
406 select {
407 case <-r.Context().Done():
408 case <-release:
409 }
410 }))
411 t.Cleanup(srv.Close)
412 t.Cleanup(func() { close(release) })
413 cfg := mcp.Config{Servers: map[string]mcp.ServerConfig{"slow": {Type: "http", URL: srv.URL}}}
414 manager := mcp.NewManager(cfg, t.TempDir(), mcp.ClientInfo{Name: "zut"}, srv.Client())
415 t.Cleanup(func() { _ = manager.Close() })
416
417 sigCh := make(chan os.Signal, 1)
418 result := make(chan os.Signal, 1)
419 go func() {
420 _, sig, _ := startMCPInterruptible(context.Background(), manager, sigCh)
421 result <- sig
422 }()
423 <-reached
424 sigCh <- os.Interrupt
425 if sig := <-result; sig != os.Interrupt {
426 t.Fatalf("signal = %v, want interrupt", sig)
427 }
428}