package plugin import ( "bufio" "bytes" "context" "encoding/json" "errors" "os" "path/filepath" "runtime" "strings" "sync" "testing" "time" "reasonix/internal/event" "reasonix/internal/tool" ) type countingToolsTransport struct { mu sync.Mutex calls int raw json.RawMessage } func (t *countingToolsTransport) call(ctx context.Context, method string, params any) (json.RawMessage, error) { if method != "tools/list" { return json.RawMessage(`{}`), nil } t.mu.Lock() t.calls++ t.mu.Unlock() if len(t.raw) > 0 { return t.raw, nil } return json.RawMessage(`{"tools":[{"name":"zed","description":"Sorted after echo.","inputSchema":{"type":"object"}},{"name":"echo","description":"Echo back the message.","inputSchema":{"type":"object","properties":{"msg":{"type":"string"}},"required":["z","msg"]},"annotations":{"readOnlyHint":true}}]}`), nil } func (t *countingToolsTransport) notify(ctx context.Context, method string, params any) error { return nil } func (t *countingToolsTransport) close() {} func (t *countingToolsTransport) toolsListCalls() int { t.mu.Lock() defer t.mu.Unlock() return t.calls } type sequenceToolsTransport struct { mu sync.Mutex calls int raws []json.RawMessage } func (t *sequenceToolsTransport) call(ctx context.Context, method string, params any) (json.RawMessage, error) { if method != "tools/list" { return json.RawMessage(`{}`), nil } t.mu.Lock() defer t.mu.Unlock() t.calls++ if len(t.raws) == 0 { return json.RawMessage(`{"tools":[]}`), nil } idx := t.calls - 1 if idx >= len(t.raws) { idx = len(t.raws) - 1 } return t.raws[idx], nil } func (t *sequenceToolsTransport) notify(ctx context.Context, method string, params any) error { return nil } func (t *sequenceToolsTransport) close() {} func (t *sequenceToolsTransport) toolsListCalls() int { t.mu.Lock() defer t.mu.Unlock() return t.calls } type deadlineRecordingTransport struct { mu sync.Mutex deadline []time.Duration methods []string block bool noContext bool } func (t *deadlineRecordingTransport) call(ctx context.Context, method string, params any) (json.RawMessage, error) { if d, ok := ctx.Deadline(); ok { t.mu.Lock() t.deadline = append(t.deadline, time.Until(d)) t.methods = append(t.methods, method) t.mu.Unlock() } else { t.mu.Lock() t.noContext = true t.methods = append(t.methods, method) t.mu.Unlock() } if t.block { <-ctx.Done() return nil, ctx.Err() } return json.RawMessage(`{}`), nil } func (t *deadlineRecordingTransport) notify(ctx context.Context, method string, params any) error { return nil } func (t *deadlineRecordingTransport) close() {} func (t *deadlineRecordingTransport) lastDeadline(tst *testing.T) time.Duration { tst.Helper() t.mu.Lock() defer t.mu.Unlock() if len(t.deadline) == 0 { tst.Fatalf("transport recorded no deadline; methods=%v noContext=%v", t.methods, t.noContext) } return t.deadline[len(t.deadline)-1] } func assertDeadlineNear(t *testing.T, got, want time.Duration) { t.Helper() if got < want-2*time.Second || got > want+2*time.Second { t.Fatalf("deadline = %v, want near %v", got, want) } } func TestClientCallAppliesBuiltInDefaultTimeout(t *testing.T) { for _, transportName := range []string{"stdio", "http"} { t.Run(transportName, func(t *testing.T) { tr := &deadlineRecordingTransport{} c := &Client{name: "maker", t: tr, spec: Spec{Name: "maker"}, transport: transportName} if _, err := c.call(context.Background(), "tools/list", map[string]any{}); err != nil { t.Fatalf("call: %v", err) } assertDeadlineNear(t, tr.lastDeadline(t), defaultCallTimeout) }) } } func TestClientCallTimeoutPrecedence(t *testing.T) { tr := &deadlineRecordingTransport{} c := &Client{ name: "maker", t: tr, spec: Spec{ Name: "maker", DefaultCallTimeout: 300 * time.Second, CallTimeout: 600 * time.Second, ToolTimeouts: map[string]time.Duration{"generate_video": 1800 * time.Second}, }, transport: "stdio", } if _, err := c.call(context.Background(), "tools/call", map[string]any{"name": "generate_video"}); err != nil { t.Fatalf("tool override call: %v", err) } assertDeadlineNear(t, tr.lastDeadline(t), 1800*time.Second) if _, err := c.call(context.Background(), "tools/call", map[string]any{"name": "search"}); err != nil { t.Fatalf("plugin override call: %v", err) } assertDeadlineNear(t, tr.lastDeadline(t), 600*time.Second) if _, err := c.call(context.Background(), "prompts/list", map[string]any{}); err != nil { t.Fatalf("method call: %v", err) } assertDeadlineNear(t, tr.lastDeadline(t), 600*time.Second) } func TestClientCallRespectsParentDeadline(t *testing.T) { tr := &deadlineRecordingTransport{} c := &Client{ name: "maker", t: tr, spec: Spec{ Name: "maker", CallTimeout: 10 * time.Minute, }, transport: "http", } ctx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond) defer cancel() if _, err := c.call(ctx, "tools/call", map[string]any{"name": "generate_video"}); err != nil { t.Fatalf("call: %v", err) } got := tr.lastDeadline(t) if got > 150*time.Millisecond { t.Fatalf("deadline = %v, want caller deadline around 100ms", got) } } func TestClientCallTimeoutErrorNamesToolAndConfig(t *testing.T) { tr := &deadlineRecordingTransport{block: true} c := &Client{ name: "maker", t: tr, spec: Spec{ Name: "maker", CallTimeout: 25 * time.Millisecond, }, transport: "stdio", } _, err := c.call(context.Background(), "tools/call", map[string]any{"name": "generate_video"}) if err == nil { t.Fatal("timed-out call returned nil error") } if !errors.Is(err, context.DeadlineExceeded) { t.Fatalf("error should wrap context deadline exceeded, got %v", err) } msg := err.Error() if !strings.Contains(msg, `MCP tool "maker.generate_video" timed out after 25ms`) || !strings.Contains(msg, "tool_timeout_seconds or call_timeout_seconds") { t.Fatalf("timeout error lacks useful guidance: %v", err) } } // TestStdioEndToEnd drives a real subprocess (this test binary re-invoked in // helper mode) through the full MCP handshake and a tool call, exercising // StartAll, tools/list, and tools/call over stdio JSON-RPC. func TestStdioEndToEnd(t *testing.T) { ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) defer cancel() spec := Spec{ Name: "mock", Command: os.Args[0], Args: []string{"-test.run=TestHelperProcess", "--"}, Env: map[string]string{"GO_WANT_HELPER_PROCESS": "1"}, } host, tools, err := StartAll(ctx, []Spec{spec}) if err != nil { t.Fatalf("StartAll: %v", err) } defer host.Close() if len(tools) != 2 { t.Fatalf("want 2 tools, got %d", len(tools)) } if got := tools[0].Name(); got != "mcp__mock__echo" { t.Fatalf("tool name: want mcp__mock__echo, got %q", got) } if got, want := string(tools[0].Schema()), `{"properties":{"msg":{"type":"string"}},"required":["msg","z"],"type":"object"}`; got != want { t.Fatalf("tool schema = %s, want %s", got, want) } out, err := tools[0].Execute(ctx, json.RawMessage(`{"msg":"hi"}`)) if err != nil { t.Fatalf("Execute: %v", err) } if out != "echo: hi" { t.Fatalf("result: want %q, got %q", "echo: hi", out) } } func TestHostToolsForReusesCachedTools(t *testing.T) { ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) defer cancel() tr := &countingToolsTransport{} host := NewHost() defer host.Close() host.clients = []*Client{{ name: "mock", t: tr, spec: Spec{Name: "mock"}, transport: "stdio", }} first, err := host.ToolsFor(ctx, "mock") if err != nil { t.Fatalf("first ToolsFor: %v", err) } second, err := host.ToolsFor(ctx, "mock") if err != nil { t.Fatalf("second ToolsFor: %v", err) } if got := tr.toolsListCalls(); got != 1 { t.Fatalf("tools/list calls = %d, want 1", got) } if len(first) != 2 || len(second) != 2 { t.Fatalf("ToolsFor lengths = %d and %d, want 2 each", len(first), len(second)) } if got := first[0].Name(); got != "mcp__mock__echo" { t.Fatalf("first tool name = %q, want sorted echo first", got) } if got, want := string(second[0].Schema()), string(first[0].Schema()); got != want { t.Fatalf("cached schema changed:\n first=%s\nsecond=%s", want, got) } if !second[0].ReadOnly() { t.Fatal("cached tool lost readOnlyHint") } statuses := host.Servers() if len(statuses) != 1 || len(statuses[0].ToolList) != 2 { t.Fatalf("server tool status = %+v, want cached tool metadata", statuses) } if statuses[0].ToolList[0].Name != "echo" || !statuses[0].ToolList[0].ReadOnlyHint { t.Fatalf("tool metadata = %+v, want sorted echo with readOnlyHint", statuses[0].ToolList) } } func TestHostToolsForCachesEmptyToolList(t *testing.T) { ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) defer cancel() tr := &countingToolsTransport{raw: json.RawMessage(`{"tools":[]}`)} host := NewHost() defer host.Close() host.clients = []*Client{{ name: "empty", t: tr, spec: Spec{Name: "empty"}, transport: "stdio", }} first, err := host.ToolsFor(ctx, "empty") if err != nil { t.Fatalf("first ToolsFor: %v", err) } second, err := host.ToolsFor(ctx, "empty") if err != nil { t.Fatalf("second ToolsFor: %v", err) } if len(first) != 0 || len(second) != 0 { t.Fatalf("ToolsFor lengths = %d and %d, want 0 each", len(first), len(second)) } if got := tr.toolsListCalls(); got != 1 { t.Fatalf("empty tools/list calls = %d, want 1", got) } } func TestClientListToolsRetriesAdvertisedEmptyToolList(t *testing.T) { ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) defer cancel() tr := &sequenceToolsTransport{raws: []json.RawMessage{ json.RawMessage(`{"tools":[]}`), json.RawMessage(`{"tools":[{"name":"echo","description":"Echo back the message.","inputSchema":{"type":"object"}}]}`), }} c := &Client{ name: "race", t: tr, spec: Spec{Name: "race"}, transport: "stdio", hasTools: true, } tools, err := c.listTools(ctx) if err != nil { t.Fatalf("listTools: %v", err) } if len(tools) != 1 || tools[0].Name() != "mcp__race__echo" { t.Fatalf("tools = %v, want mcp__race__echo", names(tools)) } if got := tr.toolsListCalls(); got != 2 { t.Fatalf("tools/list calls = %d, want 2", got) } } func TestSpecReadOnlyToolNamesMarksUnhintedToolsReadOnly(t *testing.T) { ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) defer cancel() spec := Spec{ Name: "mock", Command: os.Args[0], Args: []string{"-test.run=TestHelperProcess", "--"}, Env: map[string]string{"GO_WANT_HELPER_PROCESS": "1"}, ReadOnlyToolNames: map[string]bool{ "echo": true, }, } host, tools, err := StartAll(ctx, []Spec{spec}) if err != nil { t.Fatalf("StartAll: %v", err) } defer host.Close() byName := map[string]tool.Tool{} for _, tl := range tools { byName[tl.Name()] = tl } echo := byName["mcp__mock__echo"] if echo == nil { t.Fatalf("mcp__mock__echo missing from %v", byName) } if !echo.ReadOnly() { t.Fatal("read-only override did not mark unhinted echo tool read-only") } zed := byName["mcp__mock__zed"] if zed == nil { t.Fatalf("mcp__mock__zed missing from %v", byName) } if zed.ReadOnly() { t.Fatal("read-only override should not mark non-listed tools read-only") } } func TestSpecReadOnlyModelToolNamesMarksVisibleToolsTrusted(t *testing.T) { ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) defer cancel() spec := Spec{ Name: "mock", Command: os.Args[0], Args: []string{"-test.run=TestHelperProcess", "--"}, Env: map[string]string{"GO_WANT_HELPER_PROCESS": "1"}, ReadOnlyModelToolNames: map[string]bool{ "mcp__mock__echo": true, }, } host, tools, err := StartAll(ctx, []Spec{spec}) if err != nil { t.Fatalf("StartAll: %v", err) } defer host.Close() byName := map[string]tool.Tool{} for _, tl := range tools { byName[tl.Name()] = tl } echo := byName["mcp__mock__echo"] if echo == nil { t.Fatalf("mcp__mock__echo missing from %v", byName) } if !echo.ReadOnly() { t.Fatal("model-visible read-only override did not mark echo tool read-only") } if u, ok := echo.(tool.PlanModeUntrustedReadOnly); ok && u.PlanModeUntrustedReadOnly() { t.Fatal("model-visible read-only override should be trusted in plan mode") } zed := byName["mcp__mock__zed"] if zed == nil { t.Fatalf("mcp__mock__zed missing from %v", byName) } if zed.ReadOnly() { t.Fatal("model-visible read-only override should not mark non-listed tools read-only") } } func TestApplyKnownReadOnlyOverridesMarksCodeGraphReadTools(t *testing.T) { got := ApplyKnownReadOnlyOverrides(Spec{Name: "codegraph", ReadOnlyToolNames: map[string]bool{"custom": true}}) for _, name := range []string{"custom", "codegraph_context", "codegraph_search", "context", "search"} { if !got.ReadOnlyToolNames[name] { t.Fatalf("codegraph read-only override missing %q: %+v", name, got.ReadOnlyToolNames) } } other := ApplyKnownReadOnlyOverrides(Spec{Name: "not-codegraph"}) if other.ReadOnlyToolNames["codegraph_context"] { t.Fatalf("non-codegraph spec should not receive codegraph overrides: %+v", other.ReadOnlyToolNames) } } func TestApplyKnownOverridesPinsCodeGraphStdioToWorkspace(t *testing.T) { got := ApplyKnownOverrides(Spec{Name: "codegraph"}, "/workspace") if got.Dir != "/workspace" { t.Fatalf("codegraph stdio Dir = %q, want workspace root", got.Dir) } if !got.ReadOnlyToolNames["codegraph_search"] { t.Fatalf("codegraph read-only override missing: %+v", got.ReadOnlyToolNames) } if got.Env[codeGraphDaemonIdleTimeoutEnv] != codeGraphDaemonIdleTimeoutDefaultMS { t.Fatalf("codegraph daemon idle timeout env = %q, want %s; env=%v", got.Env[codeGraphDaemonIdleTimeoutEnv], codeGraphDaemonIdleTimeoutDefaultMS, got.Env) } preset := ApplyKnownOverrides(Spec{Name: "codegraph", Dir: "/custom"}, "/workspace") if preset.Dir != "/custom" { t.Fatalf("existing Dir should be preserved, got %q", preset.Dir) } httpSpec := ApplyKnownOverrides(Spec{Name: "codegraph", Type: "http"}, "/workspace") if httpSpec.Dir != "" { t.Fatalf("http codegraph should not receive stdio Dir, got %q", httpSpec.Dir) } if _, ok := httpSpec.Env[codeGraphDaemonIdleTimeoutEnv]; ok { t.Fatalf("http codegraph should not receive daemon idle env, got %+v", httpSpec.Env) } other := ApplyKnownOverrides(Spec{Name: "other"}, "/workspace") if other.Dir != "" { t.Fatalf("non-codegraph should not receive Dir, got %q", other.Dir) } if _, ok := other.Env[codeGraphDaemonIdleTimeoutEnv]; ok { t.Fatalf("non-codegraph should not receive daemon idle env, got %+v", other.Env) } } func TestApplyKnownOverridesPinsCodebaseMemoryToWorkspace(t *testing.T) { got := ApplyKnownOverrides(Spec{Name: "codebase-memory-mcp"}, "/workspace") if got.Dir != "/workspace" { t.Fatalf("codebase-memory-mcp stdio Dir = %q, want workspace root", got.Dir) } if !got.LowPriority { t.Fatalf("codebase-memory-mcp should run at low priority") } preset := ApplyKnownOverrides(Spec{Name: "codebase-memory-mcp", Dir: "/custom"}, "/workspace") if preset.Dir != "/custom" { t.Fatalf("existing Dir should be preserved, got %q", preset.Dir) } httpSpec := ApplyKnownOverrides(Spec{Name: "codebase-memory-mcp", Type: "http"}, "/workspace") if httpSpec.Dir != "" { t.Fatalf("http codebase-memory-mcp should not receive stdio Dir, got %q", httpSpec.Dir) } npxSpec := ApplyKnownOverrides(Spec{ Name: "custom", Command: "npx", Args: []string{"-y", "codebase-memory-mcp@latest"}, }, "/workspace") if npxSpec.Dir != "/workspace" || !npxSpec.LowPriority { t.Fatalf("npx codebase-memory-mcp override missing: %+v", npxSpec) } } func TestApplyKnownOverridesPreservesConfiguredCodeGraphDaemonIdleTimeout(t *testing.T) { got := ApplyKnownOverrides(Spec{ Name: "codegraph", Env: map[string]string{codeGraphDaemonIdleTimeoutEnv: "30000"}, }, "/workspace") if got.Env[codeGraphDaemonIdleTimeoutEnv] != "30000" { t.Fatalf("configured codegraph daemon idle timeout was overwritten: %+v", got.Env) } } func TestStartAvailableKeepsGoodServers(t *testing.T) { ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) defer cancel() good := Spec{ Name: "good", Command: os.Args[0], Args: []string{"-test.run=TestHelperProcess", "--"}, Env: map[string]string{"GO_WANT_HELPER_PROCESS": "1"}, } bad := Spec{Name: "bad", Command: "reasonix-missing-mcp-binary"} host, tools := StartAvailable(ctx, []Spec{bad, good}) defer host.Close() if len(tools) != 2 { t.Fatalf("want tools from the good server, got %d", len(tools)) } if got := host.ServerNames(); len(got) != 1 || got[0] != "good" { t.Fatalf("connected servers = %v, want [good]", got) } failures := host.Failures() if len(failures) != 1 || failures[0].Name != "bad" { t.Fatalf("failures = %+v, want bad", failures) } } // TestStartAllAllOrNothingOnFailure pins the strict StartAll contract the // parallel rewrite must preserve: any single plugin failing aborts the whole // set, returns no Host or tools, and tears down every server that did start — // including, under parallel start, a good server whose index sits after the // failing one ([bad, good]). On error the Host is nil, so callers never see a // half-built set; the started servers are closed before StartAll returns. func TestStartAllAllOrNothingOnFailure(t *testing.T) { ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) defer cancel() good := Spec{ Name: "good", Command: os.Args[0], Args: []string{"-test.run=TestHelperProcess", "--"}, Env: map[string]string{"GO_WANT_HELPER_PROCESS": "1"}, } bad := Spec{Name: "bad", Command: "reasonix-missing-mcp-binary"} for _, tc := range []struct { name string specs []Spec }{ {"failure first", []Spec{bad, good}}, {"failure last", []Spec{good, bad}}, } { t.Run(tc.name, func(t *testing.T) { host, tools, err := StartAll(ctx, tc.specs) if err == nil { if host != nil { host.Close() } t.Fatal("StartAll should fail when a plugin can't start") } if host != nil || tools != nil { t.Fatalf("failed StartAll must return nil host/tools, got host=%v tools=%d", host, len(tools)) } }) } } func TestStdioFailureCapturesStderr(t *testing.T) { ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) defer cancel() host, _ := StartAvailable(ctx, []Spec{{ Name: "stderr", Command: os.Args[0], Args: []string{"-test.run=TestHelperProcess", "--"}, Env: map[string]string{"GO_WANT_HELPER_STDERR_EXIT": "1"}, }}) defer host.Close() failures := host.Failures() if len(failures) != 1 { t.Fatalf("failures = %+v, want one", failures) } if !strings.Contains(failures[0].Error, "helper stderr boom") { t.Fatalf("failure should include stderr, got %q", failures[0].Error) } } func TestStdioUsesConfiguredPATHForCommandLookup(t *testing.T) { ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) defer cancel() dir, command := helperLauncher(t, "mock-mcp") t.Setenv("PATH", "") host, tools, err := StartAll(ctx, []Spec{{ Name: "path", Command: command, Args: []string{"-test.run=TestHelperProcess", "--"}, Env: map[string]string{ "GO_WANT_HELPER_PROCESS": "1", "PATH": dir, }, }}) if err != nil { t.Fatalf("StartAll: %v", err) } defer host.Close() if len(tools) != 2 { t.Fatalf("want helper tools, got %d", len(tools)) } } func TestStdioFallsBackToShellPATHForCommandLookup(t *testing.T) { ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) defer cancel() dir, command := helperLauncher(t, "shell-mcp") t.Setenv("PATH", "") old := stdioShellPATH stdioShellPATH = func(context.Context) string { return dir } t.Cleanup(func() { stdioShellPATH = old }) host, tools, err := StartAll(ctx, []Spec{{ Name: "shell-path", Command: command, Args: []string{"-test.run=TestHelperProcess", "--"}, Env: map[string]string{"GO_WANT_HELPER_PROCESS": "1"}, }}) if err != nil { t.Fatalf("StartAll: %v", err) } defer host.Close() if len(tools) != 2 { t.Fatalf("want helper tools, got %d", len(tools)) } } func TestStdioCommandNotFoundSuggestsPATHFix(t *testing.T) { ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) defer cancel() t.Setenv("PATH", "") old := stdioShellPATH stdioShellPATH = func(context.Context) string { return "" } t.Cleanup(func() { stdioShellPATH = old }) host, _ := StartAvailable(ctx, []Spec{{Name: "missing", Command: "reasonix-missing-mcp-binary"}}) defer host.Close() failures := host.Failures() if len(failures) != 1 { t.Fatalf("failures = %+v, want one", failures) } msg := failures[0].Error for _, want := range []string{ `command "reasonix-missing-mcp-binary" not found on PATH`, "absolute command path", "MCP server env", } { if !strings.Contains(msg, want) { t.Fatalf("failure %q missing %q", msg, want) } } } func TestStdioIgnoresRelativePATHEntries(t *testing.T) { dir := t.TempDir() bin := filepath.Join(dir, "bin") if err := os.Mkdir(bin, 0o755); err != nil { t.Fatalf("mkdir bin: %v", err) } name := "mock-mcp" target := filepath.Join(bin, name) env := []string{"PATH=bin"} if runtime.GOOS == "windows" { target += ".cmd" env = append(env, "PATHEXT=.CMD") } if err := os.WriteFile(target, []byte(""), 0o755); err != nil { t.Fatalf("write fake executable: %v", err) } t.Chdir(dir) if exe, ok := lookPathInEnv(name, env); ok { t.Fatalf("relative PATH entry resolved to %q; want no match", exe) } } func helperLauncher(t *testing.T, name string) (dir, command string) { t.Helper() if runtime.GOOS == "windows" { t.Skip("shell launcher fixture is POSIX-only") } dir = t.TempDir() command = name target := filepath.Join(dir, name) script := "#!/bin/sh\nexec " + shellQuote(os.Args[0]) + " \"$@\"\n" if err := os.WriteFile(target, []byte(script), 0o755); err != nil { t.Fatalf("write helper launcher: %v", err) } return dir, command } func shellQuote(s string) string { return "'" + strings.ReplaceAll(s, "'", "'\\''") + "'" } // TestStartPolicyConcurrencyCap verifies the semaphore-style cap: with // Concurrency=1 the handshakes must serialise even though every spec runs // in its own goroutine. We sleep briefly inside each helper's initialize so // the goroutines have a chance to overlap if the cap is broken, then assert // that observed max-in-flight never exceeded 1. func TestStartPolicyConcurrencyCap(t *testing.T) { ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) defer cancel() mk := func(name string) Spec { return Spec{ Name: name, Command: os.Args[0], Args: []string{"-test.run=TestHelperProcess", "--"}, Env: map[string]string{ "GO_WANT_HELPER_PROCESS": "1", "GO_WANT_HELPER_INIT_MS": "50", }, } } specs := []Spec{mk("a"), mk("b"), mk("c"), mk("d")} t0 := time.Now() host, tools, err := Start(ctx, specs, StartPolicy{Concurrency: 1, AbortOnError: true}) if err != nil { t.Fatalf("Start: %v", err) } defer host.Close() elapsed := time.Since(t0) // 4 specs × 50ms init each, serialised. Allow generous slack for CI. if elapsed < 4*50*time.Millisecond { t.Fatalf("with Concurrency=1, total time should be ≥ Σ(per-spec) but was %v", elapsed) } if len(tools) != 4*2 { // helper exposes 2 tools per server t.Fatalf("want %d tools, got %d", 4*2, len(tools)) } } // TestStartPolicyPerPluginTimeout verifies that one slow plugin can't take // down the whole batch in StartAvailable mode: the slow spec times out and // gets recorded as a failure while the fast one connects. func TestStartPolicyPerPluginTimeout(t *testing.T) { ctx, cancel := context.WithTimeout(context.Background(), 15*time.Second) defer cancel() fast := Spec{ Name: "fast", Command: os.Args[0], Args: []string{"-test.run=TestHelperProcess", "--"}, Env: map[string]string{"GO_WANT_HELPER_PROCESS": "1"}, } slow := Spec{ Name: "slow", Command: os.Args[0], Args: []string{"-test.run=TestHelperProcess", "--"}, Env: map[string]string{ "GO_WANT_HELPER_PROCESS": "1", "GO_WANT_HELPER_INIT_MS": "5000", // 5s, well past the 2s budget }, } host, tools, err := Start(ctx, []Spec{fast, slow}, StartPolicy{ PerPluginTimeout: 2 * time.Second, Concurrency: 2, AbortOnError: false, }) if err != nil { t.Fatalf("Start should not return err in record-failure mode: %v", err) } defer host.Close() // Regression: the per-plugin timeout context must NOT bound the long-lived // stdio child. If transport was bound to cctx instead of the parent ctx, the // goroutine's deferred cancel would kill `fast`'s subprocess at handshake // success and this Execute would fail. We invoke it explicitly here so any // future re-introduction of the bug breaks loudly. if len(tools) > 0 { if _, callErr := tools[0].Execute(ctx, json.RawMessage(`{"msg":"hi"}`)); callErr != nil { t.Fatalf("fast plugin's subprocess was killed by deferred timeout cancel: %v", callErr) } } if len(tools) != 2 { // fast contributes 2 tools t.Fatalf("want only fast's 2 tools, got %d", len(tools)) } failures := host.Failures() if len(failures) != 1 || failures[0].Name != "slow" { t.Fatalf("failures = %+v, want [slow]", failures) } } func TestStartRecordsTimeoutStats(t *testing.T) { withTempCache(t) ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) defer cancel() slow := Spec{ Name: "slow-stats", Command: os.Args[0], Args: []string{"-test.run=TestHelperProcess", "--"}, Env: map[string]string{ "GO_WANT_HELPER_PROCESS": "1", "GO_WANT_HELPER_INIT_MS": "300", }, } for i := 0; i < 3; i++ { host, _, err := Start(ctx, []Spec{slow}, StartPolicy{ PerPluginTimeout: 50 * time.Millisecond, Concurrency: 1, AbortOnError: false, }) if err != nil { t.Fatalf("Start #%d: %v", i, err) } host.Close() } deadline := time.Now().Add(2 * time.Second) for { rec := Recommend("slow-stats", 50*time.Millisecond, 3) if rec.Demote { return } if time.Now().After(deadline) { t.Fatalf("timeout samples did not trigger demote; stats=%+v rec=%+v", readStats(t, "slow-stats"), rec) } time.Sleep(10 * time.Millisecond) } } // TestStartPhaseAReturnsBeforePhaseB pins the two-phase handshake contract. // The helper advertises prompts and stalls prompts/list by 200ms; StartAvailable // must return with tools ready while the prompts surface is still empty, and the // prompts must only materialise on Host after StartPhaseB has been called and // drained — proving prompts ride the background phase, not the boot critical path. func TestStartPhaseAReturnsBeforePhaseB(t *testing.T) { ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) defer cancel() spec := Spec{ Name: "mock", Command: os.Args[0], Args: []string{"-test.run=TestHelperProcess", "--"}, Env: map[string]string{ "GO_WANT_HELPER_PROCESS": "1", "GO_WANT_HELPER_PROMPTS": "1", "GO_WANT_HELPER_PROMPT_DELAY_MS": "200", }, } host, tools := StartAvailable(ctx, []Spec{spec}) defer host.Close() if len(tools) == 0 { t.Fatalf("want tools from helper, got 0") } // Phase A returns with tools but the prompts surface must still be empty: // StartAvailable never issues prompts/list (the helper stalls it 200ms), so // prompts can only appear after StartPhaseB drains them below. We assert this // deferral directly instead of timing StartAvailable — subprocess spawn plus // the MCP handshake make a wall-clock threshold flaky on slow CI runners. if got := host.Prompts(); len(got) != 0 { t.Fatalf("phase A must not surface prompts yet, got %d", len(got)) } // Drive phase B and wait for the surface-ready event. Use a buffered channel // sink so the test never blocks the emitter — the event payload itself is // our completion signal. ready := make(chan event.Event, 4) host.StartPhaseB(ctx, event.FuncSink(func(e event.Event) { if e.Kind == event.MCPSurfaceReady { select { case ready <- e: default: } } })) select { case e := <-ready: if !strings.Contains(e.Text, "prompts ready") { t.Fatalf("phase B event text = %q, want it to mention prompts", e.Text) } case <-time.After(3 * time.Second): t.Fatal("phase B never fired MCPSurfaceReady for prompts") } if got := host.Prompts(); len(got) != 1 || got[0].Raw != "hello" { t.Fatalf("after phase B, prompts = %+v, want one named hello", got) } } func TestStartPhaseBDoesNotBlockToolCalls(t *testing.T) { ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) defer cancel() spec := Spec{ Name: "mock", Command: os.Args[0], Args: []string{"-test.run=TestHelperProcess", "--"}, Env: map[string]string{ "GO_WANT_HELPER_PROCESS": "1", "GO_WANT_HELPER_PROMPTS": "1", "GO_WANT_HELPER_PROMPT_DELAY_MS": "1000", }, } host, tools := StartAvailable(ctx, []Spec{spec}) defer host.Close() var echo tool.Tool for _, t := range tools { if t.Name() == "mcp__mock__echo" { echo = t break } } if echo == nil { t.Fatal("missing echo tool") } host.StartPhaseB(ctx, event.Discard) time.Sleep(50 * time.Millisecond) callCtx, callCancel := context.WithTimeout(ctx, 150*time.Millisecond) defer callCancel() out, err := echo.Execute(callCtx, json.RawMessage(`{"msg":"hi"}`)) if err != nil { t.Fatalf("tool call should not be blocked by background prompts/list: %v", err) } if out != "echo: hi" { t.Fatalf("Execute result = %q, want %q", out, "echo: hi") } } // TestHelperProcess is not a real test; it acts as a minimal MCP stdio server // when invoked by TestStdioEndToEnd. It exits before the test framework can // print to stdout, keeping the JSON-RPC channel clean. // // GO_WANT_HELPER_INIT_MS optionally injects a sleep before responding to the // initialize call, used by the timeout / concurrency tests to simulate slow // handshakes without depending on external processes. // GO_WANT_HELPER_PROMPTS advertises the prompts capability and registers a // "hello" prompt; GO_WANT_HELPER_PROMPT_DELAY_MS stalls prompts/list so the // phase-A vs phase-B split can be exercised. func TestHelperProcess(t *testing.T) { if os.Getenv("GO_WANT_HELPER_STDERR_EXIT") == "1" { os.Stderr.WriteString("helper stderr boom\n") os.Exit(2) } if os.Getenv("GO_WANT_HELPER_PROCESS") != "1" { return } defer os.Exit(0) var initDelay time.Duration if ms := os.Getenv("GO_WANT_HELPER_INIT_MS"); ms != "" { if v, err := time.ParseDuration(ms + "ms"); err == nil { initDelay = v } } in := bufio.NewReader(os.Stdin) for { line, err := in.ReadBytes('\n') if err != nil { return } line = bytes.TrimSpace(line) if len(line) == 0 { continue } var req struct { ID *int `json:"id"` Method string `json:"method"` Params json.RawMessage `json:"params"` } if err := json.Unmarshal(line, &req); err != nil { continue } if req.ID == nil { continue // notification: no response } var result any switch req.Method { case "initialize": if initDelay > 0 { time.Sleep(initDelay) } caps := map[string]any{} if os.Getenv("GO_WANT_HELPER_PROMPTS") == "1" { caps["prompts"] = map[string]any{} } result = map[string]any{ "protocolVersion": protocolVersion, "serverInfo": map[string]any{"name": "mock", "version": "0"}, "capabilities": caps, } case "prompts/list": if ms := os.Getenv("GO_WANT_HELPER_PROMPT_DELAY_MS"); ms != "" { if v, err := time.ParseDuration(ms + "ms"); err == nil && v > 0 { time.Sleep(v) } } result = map[string]any{"prompts": []map[string]any{{ "name": "hello", "description": "say hi", "arguments": []map[string]any{}, }}} case "tools/list": result = map[string]any{"tools": []map[string]any{{ "name": "zed", "description": "Sorted after echo.", "inputSchema": map[string]any{"type": "object"}, }, { "name": "echo", "description": "Echo back the message.", "inputSchema": map[string]any{ "type": "object", "properties": map[string]any{"msg": map[string]any{"type": "string"}}, "required": []string{"z", "msg"}, }, }}} case "tools/call": var p struct { Arguments struct { Msg string `json:"msg"` } `json:"arguments"` } _ = json.Unmarshal(req.Params, &p) result = map[string]any{"content": []map[string]any{ {"type": "text", "text": "echo: " + p.Arguments.Msg}, }} } resp := map[string]any{"jsonrpc": "2.0", "id": *req.ID, "result": result} b, _ := json.Marshal(resp) os.Stdout.Write(append(b, '\n')) } } // TestReadOnlyTrustDoesNotChangeModelVisibleSchema locks the cache invariant // behind the MCP trust controls: marking a tool trusted read-only changes its // execution/approval flags (ReadOnly / PlanModeUntrustedReadOnly) but must not // alter the model-visible tool name or input schema. If trust leaked into the // provider-visible tool list/schema, every trust toggle would break the stable // prompt prefix and drop the prompt-cache hit rate. func TestReadOnlyTrustDoesNotChangeModelVisibleSchema(t *testing.T) { startMockEcho := func(spec Spec) (*Host, map[string]tool.Tool) { t.Helper() ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) t.Cleanup(cancel) spec.Name = "mock" spec.Command = os.Args[0] spec.Args = []string{"-test.run=TestHelperProcess", "--"} spec.Env = map[string]string{"GO_WANT_HELPER_PROCESS": "1"} host, tools, err := StartAll(ctx, []Spec{spec}) if err != nil { t.Fatalf("StartAll: %v", err) } t.Cleanup(func() { host.Close() }) byName := map[string]tool.Tool{} for _, tl := range tools { byName[tl.Name()] = tl } return host, byName } _, untrusted := startMockEcho(Spec{}) _, trusted := startMockEcho(Spec{ReadOnlyModelToolNames: map[string]bool{"mcp__mock__echo": true}}) base, ok := untrusted["mcp__mock__echo"] if !ok { t.Fatalf("mcp__mock__echo missing from untrusted tools %v", untrusted) } trustedEcho, ok := trusted["mcp__mock__echo"] if !ok { t.Fatalf("mcp__mock__echo missing from trusted tools %v", trusted) } // The model-visible surface (name + schema bytes) must be byte-identical. if base.Name() != trustedEcho.Name() { t.Fatalf("trust changed model-visible tool name: %q vs %q", base.Name(), trustedEcho.Name()) } if got, want := string(trustedEcho.Schema()), string(base.Schema()); got != want { t.Fatalf("trust changed model-visible schema bytes:\n trusted=%s\n base=%s", got, want) } // Trust only flips the execution/approval flags. if base.ReadOnly() { t.Fatal("untrusted echo should not be read-only without a hint") } if !trustedEcho.ReadOnly() { t.Fatal("trusted echo should be marked read-only") } if u, ok := trustedEcho.(tool.PlanModeUntrustedReadOnly); ok && u.PlanModeUntrustedReadOnly() { t.Fatal("trusted echo should be trusted in plan mode") } }