diff --git a/cmd/whale-acp/main.go b/cmd/whale-acp/main.go index 5b7d6012..109b37bf 100644 --- a/cmd/whale-acp/main.go +++ b/cmd/whale-acp/main.go @@ -8,8 +8,10 @@ import ( "math" "os" "path/filepath" + "sort" "strconv" "strings" + "sync" "github.com/BurntSushi/toml" "github.com/usewhale/whale/internal/acp" @@ -400,6 +402,29 @@ func mergePermissions(dst, src *permFile) { } } +// mcpConsentRefusedMarker is the standardized stderr token that +// whale-ecosystem MCP servers emit when spawned without explicit consent +// (see the --allow-spawn spec). whale-acp matches it to enrich logs and to +// treat the failure as a deterministic refusal. +const mcpConsentRefusedMarker = "MCP-CONSENT-REFUSED" + +// failedMCPServers is a per-process negative cache of servers that failed to +// start, keyed by identity. Without it, a session-heavy host would re-spawn a +// broken server (e.g. a consent-refusing codemap) on every session/new. +var failedMCPServers sync.Map // key: mcpServerKey -> struct{} + +func mcpServerKey(name, command string, args []string) string { + return name + "\x00" + command + "\x00" + strings.Join(args, "\x00") +} + +// isConsentRefused reports whether a server failure carries the standardized +// "refused to start without explicit consent" marker (see the --allow-spawn +// spec), matched anywhere in the joined error/stderr text. It only affects log +// classification; the failure is negative-cached either way. +func isConsentRefused(err string) bool { + return strings.Contains(err, mcpConsentRefusedMarker) +} + // wireMCPServers loads the session MCP config, connects the configured // servers, and configures the toolset's deferred MCP catalog so the agent can // discover and promote mcp____ tools via tool_search. MCP tools @@ -414,6 +439,15 @@ func wireMCPServers(ts *tools.Toolset, dataDir, cwd string, mcps []acp.MCPServer acp.Logger.Printf("mcp config: %v", err) mcpCfg = whalemcp.Config{Servers: map[string]whalemcp.ServerConfig{}} } + // Skip servers that failed to start in an earlier session of this process: + // don't re-spawn something known-broken on every session/new. + for _, name := range sortedKeys(mcpCfg.Servers) { + srv := mcpCfg.Servers[name] + if _, bad := failedMCPServers.Load(mcpServerKey(name, srv.Command, srv.Args)); bad { + delete(mcpCfg.Servers, name) + acp.Logger.Printf("mcp server %s: previously failed to start, skipping", sanitizeLogName(name)) + } + } var mcpManager *whalemcp.Manager var reg *core.ToolRegistry if len(mcpCfg.Servers) > 0 { @@ -422,9 +456,16 @@ func wireMCPServers(ts *tools.Toolset, dataDir, cwd string, mcps []acp.MCPServer for _, st := range mcpManager.States() { switch st.Status { case whalemcp.StatusFailed, whalemcp.StatusCancelled: - acp.Logger.Printf("mcp server %s: %s (%s)", st.Name, st.Status, st.Error) + if srv, ok := mcpCfg.Servers[st.Name]; ok { + failedMCPServers.Store(mcpServerKey(st.Name, srv.Command, srv.Args), struct{}{}) + } + if isConsentRefused(st.Error) { + acp.Logger.Printf("mcp server %s: refused to start without explicit consent (%s)", sanitizeLogName(st.Name), sanitizeLogName(st.Error)) + } else { + acp.Logger.Printf("mcp server %s: %s (%s)", sanitizeLogName(st.Name), st.Status, sanitizeLogName(st.Error)) + } default: - acp.Logger.Printf("mcp server %s: %s (%d tools)", st.Name, st.Status, len(st.ToolNames)) + acp.Logger.Printf("mcp server %s: %s (%d tools)", sanitizeLogName(st.Name), st.Status, len(st.ToolNames)) } } if catalog := mcpManager.BuildDeferredCatalog(); catalog != nil && !catalog.Empty() { @@ -464,6 +505,19 @@ func mcpConfigForSession(dataDir string, mcps []acp.MCPServer) (whalemcp.Config, if err != nil { return cfg, err } + if len(mcps) > 0 { + // Client-supplied servers are arbitrary stdio processes spawned with + // the whale-acp user's privileges. That is the ACP trust model (the + // host is fully trusted), but make it visible in the log. + acp.Logger.Printf("connecting %d MCP server(s) supplied by the ACP client", len(mcps)) + } + for _, name := range sortedKeys(cfg.Servers) { + if srv := cfg.Servers[name]; strings.TrimSpace(srv.URL) != "" { + // The local baseline is passed through unchanged (matching the main + // app), but http transport is outside the stdio-only advertisement. + acp.Logger.Printf("baseline mcp server %s uses url transport (%s); whale-acp advertises stdio only", sanitizeLogName(name), sanitizeLogName(srv.URL)) + } + } for _, m := range mcps { name := strings.TrimSpace(m.Name) if name == "" { @@ -471,7 +525,7 @@ func mcpConfigForSession(dataDir string, mcps []acp.MCPServer) (whalemcp.Config, } if kind := clientMCPServerTransport(m); kind != "stdio" { // We advertise mcpCapabilities {http:false, sse:false} — stdio only. - acp.Logger.Printf("mcp server %s: unsupported transport %q (stdio only), skipping", name, kind) + acp.Logger.Printf("mcp server %s: unsupported transport %q (stdio only), skipping", sanitizeLogName(name), kind) continue } cfg.Servers[name] = whalemcp.ServerConfig{ @@ -501,6 +555,30 @@ func envVariableMap(envs []acp.EnvVariable) map[string]string { return out } +// sanitizeLogName strips control characters (log-injection defense) before a +// value is written to the log: newline/CR/tab and the rest of the C0 controls +// plus DEL are replaced with spaces so a client-supplied server name or error +// cannot forge log lines or inject terminal sequences. +func sanitizeLogName(s string) string { + return strings.Map(func(r rune) rune { + if r < ' ' || r == 0x7f { + return ' ' + } + return r + }, s) +} + +// sortedKeys returns the sorted keys of a string-keyed map, for deterministic +// iteration order in logs. +func sortedKeys[V any](m map[string]V) []string { + out := make([]string, 0, len(m)) + for k := range m { + out = append(out, k) + } + sort.Strings(out) + return out +} + // clientMCPServerTransport returns the transport kind of a client-supplied // MCP server. whale-acp advertises mcpCapabilities {http:false, sse:false}, // so only stdio servers (command + args + env) are accepted. diff --git a/cmd/whale-acp/main_test.go b/cmd/whale-acp/main_test.go index ee77aad4..667dd68c 100644 --- a/cmd/whale-acp/main_test.go +++ b/cmd/whale-acp/main_test.go @@ -1,13 +1,17 @@ package main import ( + "bytes" "context" "encoding/json" "fmt" + "log" "os" "path/filepath" "strings" + "sync" "testing" + "time" sdk "github.com/modelcontextprotocol/go-sdk/mcp" "github.com/usewhale/whale/internal/acp" @@ -505,3 +509,275 @@ func TestContextWindowFromEnv(t *testing.T) { }) } } + +// TestSanitizeLogName verifies control characters are stripped before server +// names/errors hit the log (log-injection defense), including the C0 controls +// and DEL beyond newline/CR/tab. +func TestSanitizeLogName(t *testing.T) { + cases := []struct{ in, want string }{ + {"evil\nserver\r\tname", "evil server name"}, + {"plain-name", "plain-name"}, + {"", ""}, + {"nul\x00esc\x1bdel\x7f", "nul esc del "}, + {"tab\there", "tab here"}, + } + for _, c := range cases { + if got := sanitizeLogName(c.in); got != c.want { + t.Errorf("sanitizeLogName(%q) = %q, want %q", c.in, got, c.want) + } + } +} + +// TestSortedKeys verifies deterministic iteration order for logs. +func TestSortedKeys(t *testing.T) { + if got := sortedKeys(map[string]int{}); len(got) != 0 { + t.Fatalf("empty map: %v", got) + } + got := sortedKeys(map[string]int{"z": 1, "a": 2, "m": 3}) + if strings.Join(got, ",") != "a,m,z" { + t.Fatalf("sortedKeys = %v, want a,m,z", got) + } +} + +// TestWireMCPServersNegativeCacheSkipsFailedServer verifies a server that +// failed to start is not re-spawned on a subsequent session (per-process +// negative cache), preventing repeated spawn attempts on session-heavy hosts. +func TestWireMCPServersNegativeCacheSkipsFailedServer(t *testing.T) { + name := "boom-" + fmt.Sprintf("%d", time.Now().UnixNano()) // unique per run + dataDir := t.TempDir() + cfg := map[string]any{ + "mcpServers": map[string]any{ + name: map[string]any{ + "command": "/nonexistent/definitely-not-a-binary", + "timeout": 2, + }, + }, + } + b, err := json.Marshal(cfg) + if err != nil { + t.Fatal(err) + } + if err := os.WriteFile(filepath.Join(dataDir, "mcp.json"), b, 0o600); err != nil { + t.Fatal(err) + } + ts, err := tools.NewToolset(t.TempDir()) + if err != nil { + t.Fatal(err) + } + mgr1, _ := wireMCPServers(ts, dataDir, t.TempDir(), nil) + if mgr1 == nil { + t.Fatal("first session should attempt the server") + } + _ = mgr1.Close() + mgr2, _ := wireMCPServers(ts, dataDir, t.TempDir(), nil) + if mgr2 != nil { + t.Fatal("failed server must not be re-spawned in a later session") + } +} + +// TestMCPServerKey verifies the negative-cache key separates name, command, +// and args so distinct servers never collide and identical servers always do. +func TestMCPServerKey(t *testing.T) { + a := mcpServerKey("srv", "/bin/true", []string{"-a"}) + if a != mcpServerKey("srv", "/bin/true", []string{"-a"}) { + t.Fatal("identical servers must share a key") + } + for name, other := range map[string]func() string{ + "name differs": func() string { return mcpServerKey("other", "/bin/true", []string{"-a"}) }, + "command differs": func() string { return mcpServerKey("srv", "/bin/false", []string{"-a"}) }, + "arg differs": func() string { return mcpServerKey("srv", "/bin/true", []string{"-b"}) }, + "arg order": func() string { return mcpServerKey("srv", "/bin/true", []string{"-b", "-a"}) }, + } { + if other() == a { + t.Errorf("%s: keys must differ", name) + } + } +} + +// TestIsConsentRefused verifies the consent-refusal marker classification used +// for log enrichment. +func TestIsConsentRefused(t *testing.T) { + for _, in := range []string{ + "MCP-CONSENT-REFUSED", + "error: MCP-CONSENT-REFUSED: no consent", + "exit status 1: MCP-CONSENT-REFUSED", + "prefix MCP-CONSENT-REFUSED suffix", + } { + if !isConsentRefused(in) { + t.Errorf("isConsentRefused(%q) = false, want true", in) + } + } + for _, in := range []string{"", "boom", "mcp-consent-refused", "MCP-CONSENT-REFUSE"} { + if isConsentRefused(in) { + t.Errorf("isConsentRefused(%q) = true, want false", in) + } + } +} + +// writeMCPConfig writes an mcp.json baseline into dir. +func writeMCPConfig(t *testing.T, dir string, servers map[string]any) { + t.Helper() + b, err := json.Marshal(map[string]any{"mcpServers": servers}) + if err != nil { + t.Fatal(err) + } + if err := os.WriteFile(filepath.Join(dir, "mcp.json"), b, 0o600); err != nil { + t.Fatal(err) + } +} + +// TestMCPConfigForSessionLogLines verifies the changed log sites in +// mcpConfigForSession: the client-supplied server count and the baseline http +// url-transport warning are both emitted and sanitized. +func TestMCPConfigForSessionLogLines(t *testing.T) { + var logs bytes.Buffer + prev := acp.Logger + acp.Logger = log.New(&logs, "", 0) + defer func() { acp.Logger = prev }() + + dataDir := t.TempDir() + writeMCPConfig(t, dataDir, map[string]any{ + "baseline-http": map[string]any{"url": "http://127.0.0.1:9999"}, + }) + client := []acp.MCPServer{ + {Name: "client-a", Command: "/bin/echo", Args: []string{"hi"}}, + {Name: "client-b", Command: "/bin/echo", Args: []string{"bye"}}, + } + cfg, err := mcpConfigForSession(dataDir, client) + if err != nil { + t.Fatal(err) + } + if got := cfg.Servers["client-a"].Command; got != "/bin/echo" { + t.Fatalf("client server not merged into config: %+v", cfg.Servers["client-a"]) + } + out := logs.String() + if !strings.Contains(out, "connecting 2 MCP server(s) supplied by the ACP client") { + t.Errorf("missing client-count log line, got:\n%s", out) + } + if !strings.Contains(out, "baseline mcp server baseline-http uses url transport") { + t.Errorf("missing baseline url-transport warning, got:\n%s", out) + } +} + +// TestWireMCPServersNegativeCacheKeyedByCommand: the negative cache is keyed by +// name+command+args — re-entering the same name with a different command must +// be attempted again, while the original failing key stays cached. +func TestWireMCPServersNegativeCacheKeyedByCommand(t *testing.T) { + name := "keyed-" + fmt.Sprintf("%d", time.Now().UnixNano()) + dataDir := t.TempDir() + write := func(command string) { + writeMCPConfig(t, dataDir, map[string]any{ + name: map[string]any{"command": command, "timeout": 2}, + }) + } + ts, err := tools.NewToolset(t.TempDir()) + if err != nil { + t.Fatal(err) + } + write("/nonexistent/bin-one") + mgr1, _ := wireMCPServers(ts, dataDir, t.TempDir(), nil) + if mgr1 == nil { + t.Fatal("first session should attempt the server") + } + _ = mgr1.Close() + + write("/nonexistent/bin-two") // same name, different command + mgr2, _ := wireMCPServers(ts, dataDir, t.TempDir(), nil) + if mgr2 == nil { + t.Fatal("different command must not be skipped by the cache") + } + _ = mgr2.Close() + + write("/nonexistent/bin-one") // original key is still cached + mgr3, _ := wireMCPServers(ts, dataDir, t.TempDir(), nil) + if mgr3 != nil { + t.Fatal("original failing key must still be skipped") + } +} + +// TestWireMCPServersNegativeCacheConsentRefused: a server that refuses to +// start without explicit consent (the MCP-CONSENT-REFUSED marker on stderr) is +// classified in the log and negative-cached like any other startup failure. +func TestWireMCPServersNegativeCacheConsentRefused(t *testing.T) { + name := "consent-" + fmt.Sprintf("%d", time.Now().UnixNano()) + dataDir := t.TempDir() + writeMCPConfig(t, dataDir, map[string]any{ + name: map[string]any{ + "command": "/bin/sh", + "args": []string{"-c", "echo MCP-CONSENT-REFUSED >&2; exit 1"}, + "timeout": 5, + }, + }) + var logs bytes.Buffer + prev := acp.Logger + acp.Logger = log.New(&logs, "", 0) + defer func() { acp.Logger = prev }() + + ts, err := tools.NewToolset(t.TempDir()) + if err != nil { + t.Fatal(err) + } + mgr1, _ := wireMCPServers(ts, dataDir, t.TempDir(), nil) + if mgr1 == nil { + t.Fatal("first session should attempt the server") + } + _ = mgr1.Close() + if !strings.Contains(logs.String(), "refused to start without explicit consent") { + t.Fatalf("expected consent-refusal log line, got:\n%s", logs.String()) + } + mgr2, _ := wireMCPServers(ts, dataDir, t.TempDir(), nil) + if mgr2 != nil { + t.Fatal("consent-refused server must be negative-cached") + } +} + +// TestWireMCPServersNegativeCacheConcurrent exercises the shared negative +// cache from concurrent sessions (whale-acp runs one runtime per session): +// concurrent reads of a cached failure must all skip, and concurrent stores of +// distinct failures must not race (run with -race). +func TestWireMCPServersNegativeCacheConcurrent(t *testing.T) { + ts, err := tools.NewToolset(t.TempDir()) + if err != nil { + t.Fatal(err) + } + primeDir := t.TempDir() + prime := "prime-" + fmt.Sprintf("%d", time.Now().UnixNano()) + writeMCPConfig(t, primeDir, map[string]any{ + prime: map[string]any{"command": "/nonexistent/bin-prime", "timeout": 2}, + }) + mgr, _ := wireMCPServers(ts, primeDir, t.TempDir(), nil) + if mgr == nil { + t.Fatal("prime should be attempted") + } + _ = mgr.Close() + + const n = 8 + var wg sync.WaitGroup + for i := 0; i < n; i++ { + i := i + wg.Add(1) + go func() { + defer wg.Done() + if i%2 == 0 { + // Concurrent reads of the primed (cached) key. + if m, _ := wireMCPServers(ts, primeDir, t.TempDir(), nil); m != nil { + t.Error("cached failing server must be skipped concurrently") + _ = m.Close() + } + } else { + // Concurrent stores of unique failing keys. + name := fmt.Sprintf("conc-%d-%d", time.Now().UnixNano(), i) + dir := t.TempDir() + writeMCPConfig(t, dir, map[string]any{ + name: map[string]any{"command": "/nonexistent/bin-" + name, "timeout": 2}, + }) + if m, _ := wireMCPServers(ts, dir, t.TempDir(), nil); m == nil { + t.Error("unique failing server must be attempted") + } else { + _ = m.Close() + } + } + }() + } + wg.Wait() +} diff --git a/internal/app/app_accessors.go b/internal/app/app_accessors.go index 5e6fc729..742c6077 100644 --- a/internal/app/app_accessors.go +++ b/internal/app/app_accessors.go @@ -10,8 +10,28 @@ import ( "strings" ) -func (a *App) SessionID() string { return a.sessionID } -func (a *App) SessionsDir() string { return a.sessionsDir } +func (a *App) SessionID() string { + a.sessionMu.Lock() + defer a.sessionMu.Unlock() + return a.sessionID +} +func (a *App) SessionsDir() string { return a.sessionsDir } + +// setSessionID is the only mutating path to sessionID. All writers (resume, +// session-new, fork) and the async MCP restore readers serialize on sessionMu. +func (a *App) setSessionID(id string) { + a.sessionMu.Lock() + defer a.sessionMu.Unlock() + a.sessionID = id +} + +// sessionPath returns a consistent snapshot of the session location for paths +// built by goroutines outside the dispatch loop (MCP startup restore). +func (a *App) sessionPath() (sessionsDir, sessionID string) { + a.sessionMu.Lock() + defer a.sessionMu.Unlock() + return a.sessionsDir, a.sessionID +} func (a *App) CurrentMode() session.Mode { return a.currentMode } func (a *App) PermissionDefault() policy.PermissionAction { return a.permissionPolicy.Default } func (a *App) AutoAcceptPermissions() bool { diff --git a/internal/app/app_types.go b/internal/app/app_types.go index 96214315..f4b96091 100644 --- a/internal/app/app_types.go +++ b/internal/app/app_types.go @@ -170,6 +170,11 @@ type App struct { // the entire refreshMCPTools body so concurrent refreshes serialize and // the last one always observes the latest pluginTools. toolMu sync.Mutex + // sessionMu guards sessionID against the MCP startup goroutine, which reads + // it during RestorePromotedTools while the dispatch goroutine switches + // sessions (resume / session-new / fork). sessionsDir is immutable after + // construction and does not need the lock. + sessionMu sync.Mutex a *agent.Agent apiKey string diff --git a/internal/app/command_handlers.go b/internal/app/command_handlers.go index be83734c..cea05d96 100644 --- a/internal/app/command_handlers.go +++ b/internal/app/command_handlers.go @@ -111,7 +111,7 @@ func (a *App) ExecuteSlash(line string) (CommandExecution, error) { oldMsgCount = len(msgs) } } - a.sessionID = cmdResult.SessionID + a.setSessionID(cmdResult.SessionID) if isNewCommand { modeState, err := session.LoadModeState(a.sessionsDir, a.sessionID) if err != nil { diff --git a/internal/app/fork.go b/internal/app/fork.go index d3b02e65..311db6a1 100644 --- a/internal/app/fork.go +++ b/internal/app/fork.go @@ -78,7 +78,7 @@ func (a *App) forkCurrentSession(name string) (forkSessionResult, error) { } } - a.sessionID = nextID + a.setSessionID(nextID) a.resetAgent() resume := resumeCommand(a.workspaceRoot, sourceID) msg := fmt.Sprintf("Forked conversation %q. You are now in the fork.\nTo resume the original: %s", title, resume) diff --git a/internal/app/mcp_runtime.go b/internal/app/mcp_runtime.go index 778fdc40..cc2fc087 100644 --- a/internal/app/mcp_runtime.go +++ b/internal/app/mcp_runtime.go @@ -264,14 +264,18 @@ type promotedToolState struct { } func (a *App) writePromotedToolState(state promotedToolState) error { - if a == nil || a.sessionsDir == "" || a.sessionID == "" { + if a == nil { + return nil + } + sessionsDir, sessionID := a.sessionPath() + if sessionsDir == "" || sessionID == "" { return nil } b, err := json.Marshal(state) if err != nil { return err } - path := filepath.Join(a.sessionsDir, a.sessionID, "promoted_tools.json") + path := filepath.Join(sessionsDir, sessionID, "promoted_tools.json") if err := os.MkdirAll(filepath.Dir(path), 0o755); err != nil { return err } @@ -279,10 +283,14 @@ func (a *App) writePromotedToolState(state promotedToolState) error { } func (a *App) loadPromotedToolState() ([]string, error) { - if a == nil || a.sessionsDir == "" || a.sessionID == "" { + if a == nil { + return nil, nil + } + sessionsDir, sessionID := a.sessionPath() + if sessionsDir == "" || sessionID == "" { return nil, nil } - path := filepath.Join(a.sessionsDir, a.sessionID, "promoted_tools.json") + path := filepath.Join(sessionsDir, sessionID, "promoted_tools.json") data, err := os.ReadFile(path) if err != nil { if os.IsNotExist(err) { @@ -322,15 +330,18 @@ func (a *App) RestorePromotedTools() error { if err := a.baseToolRegistry.AddTools(built); err != nil { return err } - if err := a.rebuildToolRegistriesLocked(); err != nil { - return err - } + // Track promoted tools BEFORE rebuild so collectPromotedToolsLocked sees + // them (mirrors promoteToolsLocked); otherwise the rebuild replaces the + // registry without the restored tools and they are pruned as stale. if a.promotedTools == nil { a.promotedTools = make(map[string]bool) } for _, t := range built { a.promotedTools[t.Name()] = true } + if err := a.rebuildToolRegistriesLocked(); err != nil { + return err + } return nil } diff --git a/internal/app/promoted_state_test.go b/internal/app/promoted_state_test.go new file mode 100644 index 00000000..d87dcd68 --- /dev/null +++ b/internal/app/promoted_state_test.go @@ -0,0 +1,184 @@ +package app + +import ( + "encoding/json" + "os" + "path/filepath" + "testing" +) + +// Unit coverage for the promoted-tool state persistence and the sessionID +// synchronization added with the sessionMu fix: load/write/restore branches +// plus the setSessionID/SessionID/sessionPath accessor round-trip. + +func TestLoadPromotedToolState(t *testing.T) { + mgr := newMCPRuntimeTestManager(t, "echoes a message") + app := newMCPRuntimeTestApp(mgr) + dir := t.TempDir() + app.sessionsDir = dir + app.sessionID = "sess" + sessionDir := filepath.Join(dir, "sess") + if err := os.MkdirAll(sessionDir, 0755); err != nil { + t.Fatalf("mkdir: %v", err) + } + catalog := mgr.BuildDeferredCatalog() + + writeState := func(t *testing.T, state promotedToolState) { + t.Helper() + b, err := json.Marshal(state) + if err != nil { + t.Fatalf("marshal: %v", err) + } + if err := os.WriteFile(filepath.Join(sessionDir, "promoted_tools.json"), b, 0644); err != nil { + t.Fatalf("write state: %v", err) + } + } + + t.Run("missing file returns no tools", func(t *testing.T) { + names, err := app.loadPromotedToolState() + if err != nil || names != nil { + t.Fatalf("load = %v, %v; want nil, nil", names, err) + } + }) + + t.Run("stale catalog hash returns no tools", func(t *testing.T) { + writeState(t, promotedToolState{CatalogHash: "stale-hash", ToolNames: []string{"mcp__runtime__echo"}}) + names, err := app.loadPromotedToolState() + if err != nil || names != nil { + t.Fatalf("load = %v, %v; want nil, nil (stale)", names, err) + } + }) + + t.Run("valid state returns names and records hash", func(t *testing.T) { + writeState(t, promotedToolState{CatalogHash: catalog.Hash(), ToolNames: []string{"mcp__runtime__echo"}}) + names, err := app.loadPromotedToolState() + if err != nil || len(names) != 1 || names[0] != "mcp__runtime__echo" { + t.Fatalf("load = %v, %v; want [mcp__runtime__echo], nil", names, err) + } + if app.promotedCatalogHash != catalog.Hash() { + t.Fatalf("promotedCatalogHash = %q, want %q", app.promotedCatalogHash, catalog.Hash()) + } + }) + + t.Run("malformed json returns error", func(t *testing.T) { + if err := os.WriteFile(filepath.Join(sessionDir, "promoted_tools.json"), []byte("not json"), 0644); err != nil { + t.Fatalf("write: %v", err) + } + if _, err := app.loadPromotedToolState(); err == nil { + t.Fatal("load = nil error; want error for malformed JSON") + } + }) + + t.Run("nil app returns no tools", func(t *testing.T) { + var nilApp *App + if names, err := nilApp.loadPromotedToolState(); err != nil || names != nil { + t.Fatalf("nil app load = %v, %v; want nil, nil", names, err) + } + }) +} + +func TestWritePromotedToolState(t *testing.T) { + app := &App{} + dir := t.TempDir() + app.sessionsDir = dir + app.sessionID = "sess" + + t.Run("nil app is a no-op", func(t *testing.T) { + var nilApp *App + if err := nilApp.writePromotedToolState(promotedToolState{}); err != nil { + t.Fatalf("nil app write: %v", err) + } + }) + + t.Run("empty session path is a no-op", func(t *testing.T) { + empty := &App{} + if err := empty.writePromotedToolState(promotedToolState{CatalogHash: "h", ToolNames: []string{"x"}}); err != nil { + t.Fatalf("empty path write: %v", err) + } + }) + + t.Run("writes json to session path", func(t *testing.T) { + state := promotedToolState{CatalogHash: "hash-123", ToolNames: []string{"mcp__runtime__echo", "tool_search"}} + if err := app.writePromotedToolState(state); err != nil { + t.Fatalf("write: %v", err) + } + b, err := os.ReadFile(filepath.Join(dir, "sess", "promoted_tools.json")) + if err != nil { + t.Fatalf("read back: %v", err) + } + var got promotedToolState + if err := json.Unmarshal(b, &got); err != nil { + t.Fatalf("unmarshal back: %v", err) + } + if got.CatalogHash != state.CatalogHash || len(got.ToolNames) != len(state.ToolNames) { + t.Fatalf("round-trip = %+v, want %+v", got, state) + } + }) +} + +func TestRestorePromotedToolsNoState(t *testing.T) { + mgr := newMCPRuntimeTestManager(t, "echoes a message") + app := newMCPRuntimeTestApp(mgr) + app.sessionsDir = t.TempDir() + app.sessionID = "sess" + + if err := app.RestorePromotedTools(); err != nil { + t.Fatalf("restore without state: %v", err) + } + + t.Run("nil app and nil manager are no-ops", func(t *testing.T) { + var nilApp *App + if err := nilApp.RestorePromotedTools(); err != nil { + t.Fatalf("nil app restore: %v", err) + } + noMgr := newMCPRuntimeTestApp(nil) + if err := noMgr.RestorePromotedTools(); err != nil { + t.Fatalf("nil manager restore: %v", err) + } + }) +} + +func TestRestorePromotedToolsSuccess(t *testing.T) { + mgr := newMCPRuntimeTestManager(t, "echoes a message") + app := newMCPRuntimeTestApp(mgr) + dir := t.TempDir() + app.sessionsDir = dir + app.sessionID = "sess" + sessionDir := filepath.Join(dir, "sess") + if err := os.MkdirAll(sessionDir, 0755); err != nil { + t.Fatalf("mkdir: %v", err) + } + catalog := mgr.BuildDeferredCatalog() + state := promotedToolState{CatalogHash: catalog.Hash(), ToolNames: []string{"mcp__runtime__echo"}} + b, err := json.Marshal(state) + if err != nil { + t.Fatalf("marshal: %v", err) + } + if err := os.WriteFile(filepath.Join(sessionDir, "promoted_tools.json"), b, 0644); err != nil { + t.Fatalf("write state: %v", err) + } + + if err := app.RestorePromotedTools(); err != nil { + t.Fatalf("restore: %v", err) + } + if app.toolRegistry.Get("mcp__runtime__echo") == nil && app.baseToolRegistry.Get("mcp__runtime__echo") == nil { + t.Fatal("promoted tool not registered after restore") + } + if !app.promotedTools["mcp__runtime__echo"] { + t.Fatalf("promotedTools = %v, want mcp__runtime__echo=true", app.promotedTools) + } +} + +func TestSessionIDAccessors(t *testing.T) { + dir := t.TempDir() + a := &App{sessionsDir: dir, sessionID: "old"} + + a.setSessionID("new") + if got := a.SessionID(); got != "new" { + t.Fatalf("SessionID() = %q, want %q", got, "new") + } + gotDir, gotID := a.sessionPath() + if gotDir != dir || gotID != "new" { + t.Fatalf("sessionPath() = (%q, %q), want (%q, %q)", gotDir, gotID, dir, "new") + } +} diff --git a/internal/app/resume.go b/internal/app/resume.go index ccb4de22..fb340955 100644 --- a/internal/app/resume.go +++ b/internal/app/resume.go @@ -358,7 +358,7 @@ func (a *App) ApplyResumeChoice(choice string) (ResumeApplyResult, error) { } else if blocked { return ResumeApplyResult{Message: msg}, nil } - a.sessionID = next + a.setSessionID(next) modeState, err := session.LoadModeState(a.sessionsDir, a.sessionID) if err != nil { return ResumeApplyResult{}, err diff --git a/internal/app/session_id_race_test.go b/internal/app/session_id_race_test.go new file mode 100644 index 00000000..1c5c2942 --- /dev/null +++ b/internal/app/session_id_race_test.go @@ -0,0 +1,51 @@ +package app + +import ( + "fmt" + "sync" + "testing" +) + +// TestSessionIDConcurrentReadWrite guards the sessionID synchronization used +// by the MCP startup goroutine: RestorePromotedTools (and the promoted-tool +// persistence path) read the current session while the dispatch goroutine +// switches sessions via resume / session-new / fork. Readers must never see a +// torn write; run with -race. +func TestSessionIDConcurrentReadWrite(t *testing.T) { + a := &App{sessionsDir: t.TempDir(), sessionID: "s0"} + + var wg sync.WaitGroup + // Readers mirror the async MCP startup path (sessionPath snapshot + the + // public SessionID accessor used by other goroutines). + for i := 0; i < 4; i++ { + wg.Add(1) + go func() { + defer wg.Done() + for j := 0; j < 1000; j++ { + dir, id := a.sessionPath() + if dir == "" || id == "" { + t.Errorf("sessionPath returned empty snapshot") + return + } + if a.SessionID() == "" { + t.Errorf("SessionID returned empty") + return + } + } + }() + } + // Writers mirror the dispatch paths that switch sessions. + for i := 0; i < 4; i++ { + wg.Add(1) + go func(n int) { + defer wg.Done() + for j := 0; j < 1000; j++ { + a.setSessionID(fmt.Sprintf("s%d-%d", n, j%8)) + } + }(i) + } + wg.Wait() + if a.SessionID() == "" { + t.Fatalf("sessionID lost after concurrent writes") + } +} diff --git a/internal/lsp/client.go b/internal/lsp/client.go index 51656f09..e50f4ed7 100644 --- a/internal/lsp/client.go +++ b/internal/lsp/client.go @@ -105,6 +105,14 @@ func (c *Client) Start(ctx context.Context) error { } c.readyCh = make(chan struct{}) readyCh := c.readyCh + // Publish procDone up front so Close (and its test) can always observe a + // process-tracking channel once a start attempt is in flight. The starting + // CAS above is visible to isStarting() before the spawn block below stores + // procDone, so Close must never see procDone == nil for a start that has + // begun; the reaper goroutine closes the same channel when cmd.Wait() + // returns. + procDone := make(chan struct{}) + c.procDone = procDone c.mu.Unlock() defer c.starting.Store(false) // Release concurrent waiters even when startup fails (they check @@ -168,8 +176,8 @@ func (c *Client) Start(ctx context.Context) error { // Wait for the process in the background to release OS resources. // procDone signals that the process has actually been reaped — on // Windows an unreaped process still holds handles on its working - // directory. - procDone := make(chan struct{}) + // directory. The channel was published before the spawn so Close can + // always wait on it; the goroutine closes it once cmd.Wait() returns. go func() { defer close(procDone) if err := cmd.Wait(); err != nil { diff --git a/internal/mcp/manager.go b/internal/mcp/manager.go index 9c4d8d46..30bfeebd 100644 --- a/internal/mcp/manager.go +++ b/internal/mcp/manager.go @@ -4,7 +4,6 @@ import ( "context" "errors" "fmt" - "io" "net/http" "net/url" "os" @@ -235,8 +234,13 @@ func (m *Manager) startServer(ctx context.Context, srv ServerConfig) (*clientSes if isContextTimeout(timeoutCtx, err) { return nil, nil, nil, startupTimeoutErr(srv, "connect") } - if errors.Is(err, io.EOF) && stdioCmd != nil { - err = maybeStdioErr(err, stdioCmd) + if stdioCmd != nil { + // Enrich ANY stdio connect failure with the captured stderr, not + // just io.EOF: a server that exits before the handshake write + // surfaces as EPIPE on Linux ("write |1: broken pipe"), which + // errors.Is(err, io.EOF) misses — the diagnostic stderr was then + // silently dropped. + err = maybeStdioErr(err, stdioCmd, transport) } return nil, nil, nil, startupErr(srv, "connect", err, httpDiag) } @@ -399,9 +403,12 @@ func createTransport(ctx context.Context, kind string, srv ServerConfig) (sdk.Tr cmd := exec.CommandContext(ctx, expandStdioCommand(srv.Command), expandStdioArgs(srv.Args)...) cmd.Env = append(os.Environ(), env...) shell.ConfigureCommand(cmd) + stderrBuf := &boundedStderr{} + cmd.Stderr = stderrBuf // the SDK only touches stdout/stdin transport := &stdioProcessTransport{ - base: &sdk.CommandTransport{Command: cmd}, - cmd: cmd, + base: &sdk.CommandTransport{Command: cmd}, + cmd: cmd, + stderr: stderrBuf, } return transport, cmd, nil, nil case "http": @@ -743,7 +750,15 @@ func expandWindowsPercentEnv(value string, getenv func(string) string) string { return out.String() } -func maybeStdioErr(err error, cmd *exec.Cmd) error { +func maybeStdioErr(err error, cmd *exec.Cmd, transport sdk.Transport) error { + // Prefer the captured stderr from the original spawn — no re-spawn needed. + if st, ok := transport.(*stdioProcessTransport); ok && st.stderr != nil { + if out := strings.TrimSpace(st.stderr.String()); out != "" { + return errors.Join(err, fmt.Errorf("%s", out)) + } + } + // Fallback for empty captures (e.g. the copy goroutine hadn't flushed): + // re-run once to collect diagnostics. checkErr := stdioCheck(cmd) if checkErr == nil { return err diff --git a/internal/mcp/manager_test.go b/internal/mcp/manager_test.go index 143a35b6..5cd79b0c 100644 --- a/internal/mcp/manager_test.go +++ b/internal/mcp/manager_test.go @@ -1,6 +1,7 @@ package mcp import ( + "bytes" "context" "fmt" "io" @@ -670,3 +671,135 @@ func singleFailedStateError(t *testing.T, mgr *Manager) string { } return states[0].Error } + +// TestManagerCapturesStdioStderrOnFastFailure verifies the server's stderr is +// captured from the original spawn (no re-spawn needed to diagnose). +func TestManagerCapturesStdioStderrOnFastFailure(t *testing.T) { + mgr := NewManager(Config{ + Servers: map[string]ServerConfig{ + "boom": { + Command: "/bin/sh", + Args: []string{"-c", "echo boom-stderr >&2; exit 1"}, + Timeout: 5, + }, + }, + }) + mgr.Initialize(context.Background()) + t.Cleanup(func() { _ = mgr.Close() }) + + states := mgr.States() + if len(states) != 1 { + t.Fatalf("states: %+v", states) + } + if states[0].Status != StatusFailed { + t.Fatalf("expected failed state, got %+v", states[0]) + } + if !strings.Contains(states[0].Error, "boom-stderr") { + t.Fatalf("expected captured stderr in error, got %q", states[0].Error) + } +} + +// TestBoundedStderrTail verifies the stderr capture is bounded and keeps the +// tail: the diagnostic value of stderr is the last lines, and a server may +// live (and chatter) for the whole process lifetime. +func TestBoundedStderrTail(t *testing.T) { + var b boundedStderr + // 201KB of input against a 64KB cap: the head must be dropped, the tail kept. + if _, err := b.Write(bytes.Repeat([]byte("H"), 1000)); err != nil { + t.Fatal(err) + } + for i := 0; i < 200; i++ { + if _, err := b.Write(bytes.Repeat([]byte("T"), 1000)); err != nil { + t.Fatal(err) + } + } + got := b.String() + if len(got) > stdioStderrCap { + t.Fatalf("captured %d bytes, cap %d", len(got), stdioStderrCap) + } + if strings.Contains(got, "H") { + t.Fatal("head bytes must be dropped once the cap is reached") + } + if !strings.HasSuffix(got, strings.Repeat("T", 1000)) { + t.Fatal("tail must be preserved") + } + if len(got) != stdioStderrCap { + t.Fatalf("full stream should saturate the cap: %d != %d", len(got), stdioStderrCap) + } + + // A single write larger than the cap keeps only its tail. + var b2 boundedStderr + big := bytes.Repeat([]byte("B"), stdioStderrCap+50) + if _, err := b2.Write(big); err != nil { + t.Fatal(err) + } + if got := b2.String(); len(got) != stdioStderrCap || !strings.HasSuffix(got, strings.Repeat("B", 50)) { + t.Fatalf("oversized write: len=%d, tail=%v", len(got), strings.HasSuffix(got, strings.Repeat("B", 50))) + } +} + +// TestManagerCapturesStdioStderrTailWhenLarge verifies a server that floods +// stderr before failing is diagnosed with the tail (the lines that explain the +// failure) rather than truncated head. Run with -race: the capture buffer is +// written by os/exec's copy goroutine and read concurrently by maybeStdioErr, +// so this also exercises the concurrent read/write path. +func TestManagerCapturesStdioStderrTailWhenLarge(t *testing.T) { + mgr := NewManager(Config{ + Servers: map[string]ServerConfig{ + "flood": { + Command: "/bin/sh", + Args: []string{"-c", + "i=0; while [ $i -lt 20000 ]; do echo 'xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx' >&2; i=$((i+1)); done; echo TAIL-MARKER >&2; exit 1"}, + Timeout: 10, + }, + }, + }) + mgr.Initialize(context.Background()) + t.Cleanup(func() { _ = mgr.Close() }) + + states := mgr.States() + if len(states) != 1 { + t.Fatalf("states: %+v", states) + } + if states[0].Status != StatusFailed { + t.Fatalf("expected failed state, got %+v", states[0]) + } + if !strings.Contains(states[0].Error, "TAIL-MARKER") { + t.Fatalf("expected tail marker in error, got %q", states[0].Error) + } +} + +// TestBoundedStderrConcurrent hammers Write/String from multiple goroutines: +// os/exec's internal copy goroutine writes the capture buffer while the parent +// reads it from another goroutine (e.g. maybeStdioErr after a failed Connect), +// so the sink must be safe under concurrent use. Run with -race. +func TestBoundedStderrConcurrent(t *testing.T) { + var b boundedStderr + const writers, perWriter = 4, 500 + done := make(chan struct{}, writers) + for i := 0; i < writers; i++ { + go func(seed byte) { + defer func() { done <- struct{}{} }() + for j := 0; j < perWriter; j++ { + if _, err := b.Write([]byte{seed}); err != nil { + t.Errorf("write: %v", err) + return + } + } + }(byte('A' + i)) + } + readerDone := make(chan struct{}) + go func() { + defer close(readerDone) + for j := 0; j < 1000; j++ { + _ = b.String() // must never race with the writers + } + }() + for i := 0; i < writers; i++ { + <-done + } + <-readerDone + if got := b.String(); len(got) > stdioStderrCap { + t.Fatalf("captured %d bytes, cap %d", len(got), stdioStderrCap) + } +} diff --git a/internal/mcp/stdio_transport.go b/internal/mcp/stdio_transport.go index aad21e52..f7cd30de 100644 --- a/internal/mcp/stdio_transport.go +++ b/internal/mcp/stdio_transport.go @@ -1,11 +1,13 @@ package mcp import ( + "bytes" "context" "errors" "fmt" "os" "os/exec" + "sync" "time" sdk "github.com/modelcontextprotocol/go-sdk/mcp" @@ -19,8 +21,47 @@ const ( ) type stdioProcessTransport struct { - base *sdk.CommandTransport - cmd *exec.Cmd + base *sdk.CommandTransport + cmd *exec.Cmd + stderr *boundedStderr // captures the server's stderr so failures can be + // diagnosed without re-spawning the command +} + +// stdioStderrCap bounds how much of a server's stderr is retained for +// diagnostics. A stdio server may live for the whole process lifetime and emit +// arbitrary output, so the capture must be bounded; keeping the tail preserves +// the last lines, which are the ones that explain a failure. +const stdioStderrCap = 64 << 10 + +// boundedStderr is a concurrency-safe, bounded stderr sink. os/exec's internal +// copy goroutine writes to it while the parent may read it concurrently from +// another goroutine (e.g. maybeStdioErr after a failed Connect) — bytes.Buffer +// is not safe for that. It keeps the tail of the stream, capped at +// stdioStderrCap bytes. +type boundedStderr struct { + mu sync.Mutex + buf bytes.Buffer +} + +func (b *boundedStderr) Write(p []byte) (int, error) { + b.mu.Lock() + defer b.mu.Unlock() + if len(p) > stdioStderrCap { + p = p[len(p)-stdioStderrCap:] + } + if b.buf.Len()+len(p) > stdioStderrCap { + // Make room by dropping the oldest bytes — the tail (the newest + // writes) is what explains a failure. + b.buf.Next(b.buf.Len() + len(p) - stdioStderrCap) + } + return b.buf.Write(p) +} + +// String returns the retained tail, safe for concurrent use. +func (b *boundedStderr) String() string { + b.mu.Lock() + defer b.mu.Unlock() + return b.buf.String() } func (t *stdioProcessTransport) Connect(ctx context.Context) (sdk.Connection, error) {