fix(mcp): harden transports, process cleanup, diagnostics, and mcp commands
- commit
- 6978ba7257a065cba70b00e95a6aad5c704668d8
- author
- brice · · 2026-10-15 13:35 UTC
- parent
- a124bc3
- signature
- GPG Verified
CD88F9E564E0102646FB6077A1A07C7C7C3E349D· arch-gpg Verified for the committer email
20 files · +2108 -181
M docs/mcp.md +23 -8
| @@ -14,6 +14,12 @@ with the same name: | |||
| 14 | 14 | | `$ZUT_HOME/mcp.json` | Global | | |
| 15 | 15 | | `<cwd>/.zut/mcp.json` | Project | | |
| 16 | 16 | ||
| 17 | + | > **Executable configuration.** zut starts every configured stdio command and | |
| 18 | + | > contacts every HTTP endpoint when it launches, before any model turn. There | |
| 19 | + | > is no trust prompt, `--no-yolo` does not gate startup, and servers are not | |
| 20 | + | > sandboxed. Review a repository's `.zut/mcp.json` before running zut in it, or | |
| 21 | + | > pass `--no-mcp`. | |
| 22 | + | ||
| 17 | 23 | The format matches Claude Desktop and Claude Code: | |
| 18 | 24 | ||
| 19 | 25 | ```json | |
| @@ -84,8 +90,11 @@ startup errors. | |||
| 84 | 90 | ||
| 85 | 91 | - Servers connect in parallel at startup. A failure is reported (stderr in | |
| 86 | 92 | headless modes, the startup panel in the TUI) and zut continues without that | |
| 87 | - | server unless it is `required`. The next call to one of its tools retries the | |
| 88 | - | connection; a crashed stdio server is restarted the same way. | |
| 93 | + | server unless it is `required`. A server that failed at startup exposes no | |
| 94 | + | tools for the rest of the run; fix it (`zut mcp status NAME`) and restart | |
| 95 | + | zut. A server that started and later dropped (a crashed stdio process or an | |
| 96 | + | expired HTTP session) reconnects on the next call to one of its tools; the | |
| 97 | + | failed call itself is not retried. | |
| 89 | 98 | - zut negotiates protocol `2025-06-18` and accepts `2025-03-26` and | |
| 90 | 99 | `2024-11-05`. It answers `ping` and `roots/list` (the zut cwd) and follows | |
| 91 | 100 | `tools/list` pagination. | |
| @@ -104,15 +113,21 @@ startup errors. | |||
| 104 | 113 | `mcp_read_resource` tools are added. | |
| 105 | 114 | - Server-provided `instructions` are added to the system prompt, marked as | |
| 106 | 115 | server-provided. Treat MCP output as untrusted data. | |
| 107 | - | - Cancelling a turn sends `notifications/cancelled`. Exiting zut closes HTTP | |
| 108 | - | sessions and stops stdio server process groups. | |
| 116 | + | - Cancelling a turn sends a best-effort `notifications/cancelled`. Exiting zut | |
| 117 | + | closes HTTP sessions and stops stdio servers with their process group | |
| 118 | + | (Unix) or job object (Windows). A descendant that starts its own session | |
| 119 | + | (`setsid`) on Unix is not stopped. | |
| 120 | + | - Errors and status output strip terminal control sequences, URL credentials | |
| 121 | + | and query strings, and configured `env`, header, and bearer-token values. | |
| 122 | + | Secrets passed through `args` are not redacted. | |
| 109 | 123 | ||
| 110 | 124 | ## Tool selection and confirmation | |
| 111 | 125 | ||
| 112 | 126 | - `--no-mcp` skips MCP for one run. `--no-tools` also disables it. | |
| 113 | 127 | - With `--tools`, add `mcp` for every MCP tool, or list exact names such as | |
| 114 | 128 | `mcp__docs__search`. | |
| 115 | - | - Packaged agents (`zut run`) never start MCP servers. | |
| 129 | + | - Packaged agents (`zut run`), bot runs, and SDK runtimes never start MCP | |
| 130 | + | servers. | |
| 116 | 131 | - `--no-yolo` confirmation covers MCP tools like any other tool. | |
| 117 | 132 | - Tool annotations such as `readOnlyHint` are hints from the server, not a | |
| 118 | 133 | security boundary. | |
| @@ -122,9 +137,9 @@ startup errors. | |||
| 122 | 137 | All agents in one zut process share the same MCP connections. | |
| 123 | 138 | ||
| 124 | 139 | - **Headless `--orchestrate`**: the manager-only orchestrator holds only MCP | |
| 125 | - | tools annotated `readOnlyHint: true` plus the resource tools, for research. | |
| 126 | - | Its workers receive every MCP tool, so actions that change external state | |
| 127 | - | are delegated. | |
| 140 | + | tools annotated `readOnlyHint: true` plus the resource tools, for research, | |
| 141 | + | still limited by any `--tools` selection. Its workers receive every MCP tool, | |
| 142 | + | so actions that change external state are delegated. | |
| 128 | 143 | - **Resident subagents** inherit MCP tools like other child-safe tools. | |
| 129 | 144 | `subagents.allowed_tools` and profile `tools:` lists accept `mcp` to grant | |
| 130 | 145 | every MCP tool, or exact `mcp__<server>__<tool>` names. | |
| @@ -14,6 +14,12 @@ with the same name: | |||||
| 14 | | `$ZUT_HOME/mcp.json` | Global | | 14 | | `$ZUT_HOME/mcp.json` | Global | | ||
| 15 | | `<cwd>/.zut/mcp.json` | Project | | 15 | | `<cwd>/.zut/mcp.json` | Project | | ||
| 16 | 16 | ||||
| 17 | + | > **Executable configuration.** zut starts every configured stdio command and | |||
| 18 | + | > contacts every HTTP endpoint when it launches, before any model turn. There | |||
| 19 | + | > is no trust prompt, `--no-yolo` does not gate startup, and servers are not | |||
| 20 | + | > sandboxed. Review a repository's `.zut/mcp.json` before running zut in it, or | |||
| 21 | + | > pass `--no-mcp`. | |||
| 22 | + | ||||
| 17 | The format matches Claude Desktop and Claude Code: | 23 | The format matches Claude Desktop and Claude Code: | ||
| 18 | 24 | ||||
| 19 | ```json | 25 | ```json | ||
| @@ -84,8 +90,11 @@ startup errors. | |||||
| 84 | 90 | ||||
| 85 | - Servers connect in parallel at startup. A failure is reported (stderr in | 91 | - Servers connect in parallel at startup. A failure is reported (stderr in | ||
| 86 | headless modes, the startup panel in the TUI) and zut continues without that | 92 | headless modes, the startup panel in the TUI) and zut continues without that | ||
| 87 | - | server unless it is `required`. The next call to one of its tools retries the | 93 | + | server unless it is `required`. A server that failed at startup exposes no |
| 88 | - | connection; a crashed stdio server is restarted the same way. | 94 | + | tools for the rest of the run; fix it (`zut mcp status NAME`) and restart |
| 95 | + | zut. A server that started and later dropped (a crashed stdio process or an | |||
| 96 | + | expired HTTP session) reconnects on the next call to one of its tools; the | |||
| 97 | + | failed call itself is not retried. | |||
| 89 | - zut negotiates protocol `2025-06-18` and accepts `2025-03-26` and | 98 | - zut negotiates protocol `2025-06-18` and accepts `2025-03-26` and | ||
| 90 | `2024-11-05`. It answers `ping` and `roots/list` (the zut cwd) and follows | 99 | `2024-11-05`. It answers `ping` and `roots/list` (the zut cwd) and follows | ||
| 91 | `tools/list` pagination. | 100 | `tools/list` pagination. | ||
| @@ -104,15 +113,21 @@ startup errors. | |||||
| 104 | `mcp_read_resource` tools are added. | 113 | `mcp_read_resource` tools are added. | ||
| 105 | - Server-provided `instructions` are added to the system prompt, marked as | 114 | - Server-provided `instructions` are added to the system prompt, marked as | ||
| 106 | server-provided. Treat MCP output as untrusted data. | 115 | server-provided. Treat MCP output as untrusted data. | ||
| 107 | - | - Cancelling a turn sends `notifications/cancelled`. Exiting zut closes HTTP | 116 | + | - Cancelling a turn sends a best-effort `notifications/cancelled`. Exiting zut |
| 108 | - | sessions and stops stdio server process groups. | 117 | + | closes HTTP sessions and stops stdio servers with their process group |
| 118 | + | (Unix) or job object (Windows). A descendant that starts its own session | |||
| 119 | + | (`setsid`) on Unix is not stopped. | |||
| 120 | + | - Errors and status output strip terminal control sequences, URL credentials | |||
| 121 | + | and query strings, and configured `env`, header, and bearer-token values. | |||
| 122 | + | Secrets passed through `args` are not redacted. | |||
| 109 | 123 | ||||
| 110 | ## Tool selection and confirmation | 124 | ## Tool selection and confirmation | ||
| 111 | 125 | ||||
| 112 | - `--no-mcp` skips MCP for one run. `--no-tools` also disables it. | 126 | - `--no-mcp` skips MCP for one run. `--no-tools` also disables it. | ||
| 113 | - With `--tools`, add `mcp` for every MCP tool, or list exact names such as | 127 | - With `--tools`, add `mcp` for every MCP tool, or list exact names such as | ||
| 114 | `mcp__docs__search`. | 128 | `mcp__docs__search`. | ||
| 115 | - | - Packaged agents (`zut run`) never start MCP servers. | 129 | + | - Packaged agents (`zut run`), bot runs, and SDK runtimes never start MCP |
| 130 | + | servers. | |||
| 116 | - `--no-yolo` confirmation covers MCP tools like any other tool. | 131 | - `--no-yolo` confirmation covers MCP tools like any other tool. | ||
| 117 | - Tool annotations such as `readOnlyHint` are hints from the server, not a | 132 | - Tool annotations such as `readOnlyHint` are hints from the server, not a | ||
| 118 | security boundary. | 133 | security boundary. | ||
| @@ -122,9 +137,9 @@ startup errors. | |||||
| 122 | All agents in one zut process share the same MCP connections. | 137 | All agents in one zut process share the same MCP connections. | ||
| 123 | 138 | ||||
| 124 | - **Headless `--orchestrate`**: the manager-only orchestrator holds only MCP | 139 | - **Headless `--orchestrate`**: the manager-only orchestrator holds only MCP | ||
| 125 | - | tools annotated `readOnlyHint: true` plus the resource tools, for research. | 140 | + | tools annotated `readOnlyHint: true` plus the resource tools, for research, |
| 126 | - | Its workers receive every MCP tool, so actions that change external state | 141 | + | still limited by any `--tools` selection. Its workers receive every MCP tool, |
| 127 | - | are delegated. | 142 | + | so actions that change external state are delegated. |
| 128 | - **Resident subagents** inherit MCP tools like other child-safe tools. | 143 | - **Resident subagents** inherit MCP tools like other child-safe tools. | ||
| 129 | `subagents.allowed_tools` and profile `tools:` lists accept `mcp` to grant | 144 | `subagents.allowed_tools` and profile `tools:` lists accept `mcp` to grant | ||
| 130 | every MCP tool, or exact `mcp__<server>__<tool>` names. | 145 | every MCP tool, or exact `mcp__<server>__<tool>` names. | ||
M packages/agent/cli.go +17 -0
| @@ -1005,6 +1005,19 @@ func runWithArgs(args Args, version string) error { | |||
| 1005 | 1005 | return err | |
| 1006 | 1006 | } | |
| 1007 | 1007 | ctx := context.Background() | |
| 1008 | + | if args.Mode != ModeInteractive && args.Mode != "" { | |
| 1009 | + | // Headless runs cancel on the first SIGINT/SIGTERM/SIGHUP so deferred | |
| 1010 | + | // cleanup (MCP server process groups, HTTP sessions) runs. The first | |
| 1011 | + | // signal restores default handling, so a second one still kills zut. | |
| 1012 | + | // The TUI keeps Ctrl+C as a key and has its own SIGTERM handler. | |
| 1013 | + | var stopSignals context.CancelFunc | |
| 1014 | + | ctx, stopSignals = signal.NotifyContext(ctx, os.Interrupt, syscall.SIGTERM, syscall.SIGHUP) | |
| 1015 | + | defer stopSignals() | |
| 1016 | + | go func() { | |
| 1017 | + | <-ctx.Done() | |
| 1018 | + | stopSignals() | |
| 1019 | + | }() | |
| 1020 | + | } | |
| 1008 | 1021 | // MCP servers start once per process. Every later Resolve, including | |
| 1009 | 1022 | // resident subagents, shares these connections through args. | |
| 1010 | 1023 | stopMCP, err := startMCP(ctx, &args, version) | |
| @@ -3083,6 +3096,10 @@ func runInteractive(ctx context.Context, args Args, version string) (runErr erro | |||
| 3083 | 3096 | // cleanup, so wait for reload and close the shared runtime explicitly. | |
| 3084 | 3097 | <-reloadDone | |
| 3085 | 3098 | _ = closeSubagentRuntimeFresh(runtime) | |
| 3099 | + | // MCP servers run in their own process groups; stop them too. | |
| 3100 | + | if args.mcpManager != nil { | |
| 3101 | + | _ = args.mcpManager.Close() | |
| 3102 | + | } | |
| 3086 | 3103 | // Exit cleanly. Re-raising the signal would skip os.Exit's | |
| 3087 | 3104 | // at-exit hooks; explicit exit is fine because we've already | |
| 3088 | 3105 | // flushed the only at-risk state (the session file) and stopped | |
| @@ -1005,6 +1005,19 @@ func runWithArgs(args Args, version string) error { | |||||
| 1005 | return err | 1005 | return err | ||
| 1006 | } | 1006 | } | ||
| 1007 | ctx := context.Background() | 1007 | ctx := context.Background() | ||
| 1008 | + | if args.Mode != ModeInteractive && args.Mode != "" { | |||
| 1009 | + | // Headless runs cancel on the first SIGINT/SIGTERM/SIGHUP so deferred | |||
| 1010 | + | // cleanup (MCP server process groups, HTTP sessions) runs. The first | |||
| 1011 | + | // signal restores default handling, so a second one still kills zut. | |||
| 1012 | + | // The TUI keeps Ctrl+C as a key and has its own SIGTERM handler. | |||
| 1013 | + | var stopSignals context.CancelFunc | |||
| 1014 | + | ctx, stopSignals = signal.NotifyContext(ctx, os.Interrupt, syscall.SIGTERM, syscall.SIGHUP) | |||
| 1015 | + | defer stopSignals() | |||
| 1016 | + | go func() { | |||
| 1017 | + | <-ctx.Done() | |||
| 1018 | + | stopSignals() | |||
| 1019 | + | }() | |||
| 1020 | + | } | |||
| 1008 | // MCP servers start once per process. Every later Resolve, including | 1021 | // MCP servers start once per process. Every later Resolve, including | ||
| 1009 | // resident subagents, shares these connections through args. | 1022 | // resident subagents, shares these connections through args. | ||
| 1010 | stopMCP, err := startMCP(ctx, &args, version) | 1023 | stopMCP, err := startMCP(ctx, &args, version) | ||
| @@ -3083,6 +3096,10 @@ func runInteractive(ctx context.Context, args Args, version string) (runErr erro | |||||
| 3083 | // cleanup, so wait for reload and close the shared runtime explicitly. | 3096 | // cleanup, so wait for reload and close the shared runtime explicitly. | ||
| 3084 | <-reloadDone | 3097 | <-reloadDone | ||
| 3085 | _ = closeSubagentRuntimeFresh(runtime) | 3098 | _ = closeSubagentRuntimeFresh(runtime) | ||
| 3099 | + | // MCP servers run in their own process groups; stop them too. | |||
| 3100 | + | if args.mcpManager != nil { | |||
| 3101 | + | _ = args.mcpManager.Close() | |||
| 3102 | + | } | |||
| 3086 | // Exit cleanly. Re-raising the signal would skip os.Exit's | 3103 | // Exit cleanly. Re-raising the signal would skip os.Exit's | ||
| 3087 | // at-exit hooks; explicit exit is fine because we've already | 3104 | // at-exit hooks; explicit exit is fine because we've already | ||
| 3088 | // flushed the only at-risk state (the session file) and stopped | 3105 | // flushed the only at-risk state (the session file) and stopped | ||
M packages/agent/mcp/client.go +55 -4
| @@ -8,6 +8,8 @@ import ( | |||
| 8 | 8 | "net/http" | |
| 9 | 9 | "net/url" | |
| 10 | 10 | "path/filepath" | |
| 11 | + | "runtime" | |
| 12 | + | "strings" | |
| 11 | 13 | "sync" | |
| 12 | 14 | "sync/atomic" | |
| 13 | 15 | "time" | |
| @@ -124,6 +126,13 @@ type ClientInfo struct { | |||
| 124 | 126 | ||
| 125 | 127 | // Connect starts the transport for cfg and performs the initialize handshake. | |
| 126 | 128 | func Connect(ctx context.Context, name string, cfg ServerConfig, cwd string, info ClientInfo, httpClient *http.Client) (*Client, error) { | |
| 129 | + | return connectWith(ctx, name, cfg, cwd, info, httpClient, nil) | |
| 130 | + | } | |
| 131 | + | ||
| 132 | + | // connectWith is Connect with the tools/list_changed callback installed before | |
| 133 | + | // the handshake starts, so a notification sent right after initialize is not | |
| 134 | + | // lost. | |
| 135 | + | func connectWith(ctx context.Context, name string, cfg ServerConfig, cwd string, info ClientInfo, httpClient *http.Client, onToolsChanged func()) (*Client, error) { | |
| 127 | 136 | var ( | |
| 128 | 137 | t transport | |
| 129 | 138 | err error | |
| @@ -147,14 +156,17 @@ func Connect(ctx context.Context, name string, cfg ServerConfig, cwd string, inf | |||
| 147 | 156 | if err != nil { | |
| 148 | 157 | return nil, err | |
| 149 | 158 | } | |
| 150 | - | c := &Client{name: name, t: t, cwd: cwd} | |
| 159 | + | c := &Client{name: name, t: t, cwd: cwd, onToolsChanged: onToolsChanged} | |
| 151 | 160 | t.setHandler(c.handleIncoming) | |
| 152 | 161 | if err := c.initialize(ctx, info); err != nil { | |
| 162 | + | // Close first: it reaps the process and waits for stderr to drain, so | |
| 163 | + | // the tail is complete (a server that fails fast writes its reason | |
| 164 | + | // to stderr just before exiting). | |
| 165 | + | _ = t.close() | |
| 153 | 166 | tail := "" | |
| 154 | 167 | if st, ok := t.(*stdioTransport); ok { | |
| 155 | 168 | tail = st.stderrTail() | |
| 156 | 169 | } | |
| 157 | - | _ = t.close() | |
| 158 | 170 | if tail != "" { | |
| 159 | 171 | return nil, fmt.Errorf("%w\nserver stderr:\n%s", err, tail) | |
| 160 | 172 | } | |
| @@ -245,8 +257,7 @@ func (c *Client) handleIncoming(msg message) (any, *RPCError) { | |||
| 245 | 257 | roots := []map[string]string{} | |
| 246 | 258 | if c.cwd != "" { | |
| 247 | 259 | if abs, err := filepath.Abs(c.cwd); err == nil { | |
| 248 | - | u := url.URL{Scheme: "file", Path: filepath.ToSlash(abs)} | |
| 249 | - | roots = append(roots, map[string]string{"uri": u.String(), "name": filepath.Base(abs)}) | |
| 260 | + | roots = append(roots, map[string]string{"uri": pathToFileURI(abs), "name": filepath.Base(abs)}) | |
| 250 | 261 | } | |
| 251 | 262 | } | |
| 252 | 263 | return map[string]any{"roots": roots}, nil | |
| @@ -407,3 +418,43 @@ func withTimeout(ctx context.Context, d time.Duration) (context.Context, context | |||
| 407 | 418 | } | |
| 408 | 419 | return context.WithTimeout(ctx, d) | |
| 409 | 420 | } | |
| 421 | + | ||
| 422 | + | // pathToFileURI converts an absolute path of the running platform to a file URI. | |
| 423 | + | func pathToFileURI(path string) string { | |
| 424 | + | return fileURI(path, runtime.GOOS == "windows") | |
| 425 | + | } | |
| 426 | + | ||
| 427 | + | // fileURI builds a file URI from an absolute path. windows selects Windows | |
| 428 | + | // path syntax (drive letters, UNC shares, backslashes) and is a parameter so | |
| 429 | + | // both shapes can be tested on any platform: | |
| 430 | + | // | |
| 431 | + | // C:\Users\me\proj -> file:///C:/Users/me/proj | |
| 432 | + | // \\host\share\dir -> file://host/share/dir | |
| 433 | + | // /home/me/my proj -> file:///home/me/my%20proj | |
| 434 | + | func fileURI(path string, windows bool) string { | |
| 435 | + | u := url.URL{Scheme: "file"} | |
| 436 | + | if !windows { | |
| 437 | + | u.Path = path | |
| 438 | + | return u.String() | |
| 439 | + | } | |
| 440 | + | p := path | |
| 441 | + | // Strip the extended-length prefixes: \\?\C:\x and \\?\UNC\host\share. | |
| 442 | + | switch { | |
| 443 | + | case strings.HasPrefix(p, `\\?\UNC\`): | |
| 444 | + | p = `\\` + p[len(`\\?\UNC\`):] | |
| 445 | + | case strings.HasPrefix(p, `\\?\`): | |
| 446 | + | p = p[len(`\\?\`):] | |
| 447 | + | } | |
| 448 | + | p = strings.ReplaceAll(p, `\`, "/") | |
| 449 | + | if strings.HasPrefix(p, "//") { | |
| 450 | + | host, rest, _ := strings.Cut(strings.TrimLeft(p, "/"), "/") | |
| 451 | + | u.Host = host | |
| 452 | + | u.Path = "/" + rest | |
| 453 | + | return u.String() | |
| 454 | + | } | |
| 455 | + | if !strings.HasPrefix(p, "/") { | |
| 456 | + | p = "/" + p // C:/x -> /C:/x | |
| 457 | + | } | |
| 458 | + | u.Path = p | |
| 459 | + | return u.String() | |
| 460 | + | } | |
| @@ -8,6 +8,8 @@ import ( | |||||
| 8 | "net/http" | 8 | "net/http" | ||
| 9 | "net/url" | 9 | "net/url" | ||
| 10 | "path/filepath" | 10 | "path/filepath" | ||
| 11 | + | "runtime" | |||
| 12 | + | "strings" | |||
| 11 | "sync" | 13 | "sync" | ||
| 12 | "sync/atomic" | 14 | "sync/atomic" | ||
| 13 | "time" | 15 | "time" | ||
| @@ -124,6 +126,13 @@ type ClientInfo struct { | |||||
| 124 | 126 | ||||
| 125 | // Connect starts the transport for cfg and performs the initialize handshake. | 127 | // Connect starts the transport for cfg and performs the initialize handshake. | ||
| 126 | func Connect(ctx context.Context, name string, cfg ServerConfig, cwd string, info ClientInfo, httpClient *http.Client) (*Client, error) { | 128 | func Connect(ctx context.Context, name string, cfg ServerConfig, cwd string, info ClientInfo, httpClient *http.Client) (*Client, error) { | ||
| 129 | + | return connectWith(ctx, name, cfg, cwd, info, httpClient, nil) | |||
| 130 | + | } | |||
| 131 | + | ||||
| 132 | + | // connectWith is Connect with the tools/list_changed callback installed before | |||
| 133 | + | // the handshake starts, so a notification sent right after initialize is not | |||
| 134 | + | // lost. | |||
| 135 | + | func connectWith(ctx context.Context, name string, cfg ServerConfig, cwd string, info ClientInfo, httpClient *http.Client, onToolsChanged func()) (*Client, error) { | |||
| 127 | var ( | 136 | var ( | ||
| 128 | t transport | 137 | t transport | ||
| 129 | err error | 138 | err error | ||
| @@ -147,14 +156,17 @@ func Connect(ctx context.Context, name string, cfg ServerConfig, cwd string, inf | |||||
| 147 | if err != nil { | 156 | if err != nil { | ||
| 148 | return nil, err | 157 | return nil, err | ||
| 149 | } | 158 | } | ||
| 150 | - | c := &Client{name: name, t: t, cwd: cwd} | 159 | + | c := &Client{name: name, t: t, cwd: cwd, onToolsChanged: onToolsChanged} |
| 151 | t.setHandler(c.handleIncoming) | 160 | t.setHandler(c.handleIncoming) | ||
| 152 | if err := c.initialize(ctx, info); err != nil { | 161 | if err := c.initialize(ctx, info); err != nil { | ||
| 162 | + | // Close first: it reaps the process and waits for stderr to drain, so | |||
| 163 | + | // the tail is complete (a server that fails fast writes its reason | |||
| 164 | + | // to stderr just before exiting). | |||
| 165 | + | _ = t.close() | |||
| 153 | tail := "" | 166 | tail := "" | ||
| 154 | if st, ok := t.(*stdioTransport); ok { | 167 | if st, ok := t.(*stdioTransport); ok { | ||
| 155 | tail = st.stderrTail() | 168 | tail = st.stderrTail() | ||
| 156 | } | 169 | } | ||
| 157 | - | _ = t.close() | |||
| 158 | if tail != "" { | 170 | if tail != "" { | ||
| 159 | return nil, fmt.Errorf("%w\nserver stderr:\n%s", err, tail) | 171 | return nil, fmt.Errorf("%w\nserver stderr:\n%s", err, tail) | ||
| 160 | } | 172 | } | ||
| @@ -245,8 +257,7 @@ func (c *Client) handleIncoming(msg message) (any, *RPCError) { | |||||
| 245 | roots := []map[string]string{} | 257 | roots := []map[string]string{} | ||
| 246 | if c.cwd != "" { | 258 | if c.cwd != "" { | ||
| 247 | if abs, err := filepath.Abs(c.cwd); err == nil { | 259 | if abs, err := filepath.Abs(c.cwd); err == nil { | ||
| 248 | - | u := url.URL{Scheme: "file", Path: filepath.ToSlash(abs)} | 260 | + | roots = append(roots, map[string]string{"uri": pathToFileURI(abs), "name": filepath.Base(abs)}) |
| 249 | - | roots = append(roots, map[string]string{"uri": u.String(), "name": filepath.Base(abs)}) | |||
| 250 | } | 261 | } | ||
| 251 | } | 262 | } | ||
| 252 | return map[string]any{"roots": roots}, nil | 263 | return map[string]any{"roots": roots}, nil | ||
| @@ -407,3 +418,43 @@ func withTimeout(ctx context.Context, d time.Duration) (context.Context, context | |||||
| 407 | } | 418 | } | ||
| 408 | return context.WithTimeout(ctx, d) | 419 | return context.WithTimeout(ctx, d) | ||
| 409 | } | 420 | } | ||
| 421 | + | ||||
| 422 | + | // pathToFileURI converts an absolute path of the running platform to a file URI. | |||
| 423 | + | func pathToFileURI(path string) string { | |||
| 424 | + | return fileURI(path, runtime.GOOS == "windows") | |||
| 425 | + | } | |||
| 426 | + | ||||
| 427 | + | // fileURI builds a file URI from an absolute path. windows selects Windows | |||
| 428 | + | // path syntax (drive letters, UNC shares, backslashes) and is a parameter so | |||
| 429 | + | // both shapes can be tested on any platform: | |||
| 430 | + | // | |||
| 431 | + | // C:\Users\me\proj -> file:///C:/Users/me/proj | |||
| 432 | + | // \\host\share\dir -> file://host/share/dir | |||
| 433 | + | // /home/me/my proj -> file:///home/me/my%20proj | |||
| 434 | + | func fileURI(path string, windows bool) string { | |||
| 435 | + | u := url.URL{Scheme: "file"} | |||
| 436 | + | if !windows { | |||
| 437 | + | u.Path = path | |||
| 438 | + | return u.String() | |||
| 439 | + | } | |||
| 440 | + | p := path | |||
| 441 | + | // Strip the extended-length prefixes: \\?\C:\x and \\?\UNC\host\share. | |||
| 442 | + | switch { | |||
| 443 | + | case strings.HasPrefix(p, `\\?\UNC\`): | |||
| 444 | + | p = `\\` + p[len(`\\?\UNC\`):] | |||
| 445 | + | case strings.HasPrefix(p, `\\?\`): | |||
| 446 | + | p = p[len(`\\?\`):] | |||
| 447 | + | } | |||
| 448 | + | p = strings.ReplaceAll(p, `\`, "/") | |||
| 449 | + | if strings.HasPrefix(p, "//") { | |||
| 450 | + | host, rest, _ := strings.Cut(strings.TrimLeft(p, "/"), "/") | |||
| 451 | + | u.Host = host | |||
| 452 | + | u.Path = "/" + rest | |||
| 453 | + | return u.String() | |||
| 454 | + | } | |||
| 455 | + | if !strings.HasPrefix(p, "/") { | |||
| 456 | + | p = "/" + p // C:/x -> /C:/x | |||
| 457 | + | } | |||
| 458 | + | u.Path = p | |||
| 459 | + | return u.String() | |||
| 460 | + | } | |||
M packages/agent/mcp/fakeserver_test.go +135 -8
| @@ -2,6 +2,7 @@ package mcp | |||
| 2 | 2 | ||
| 3 | 3 | import ( | |
| 4 | 4 | "bufio" | |
| 5 | + | "bytes" | |
| 5 | 6 | "context" | |
| 6 | 7 | "encoding/base64" | |
| 7 | 8 | "encoding/json" | |
| @@ -9,11 +10,14 @@ import ( | |||
| 9 | 10 | "net/http" | |
| 10 | 11 | "net/http/httptest" | |
| 11 | 12 | "os" | |
| 13 | + | "os/exec" | |
| 14 | + | "os/signal" | |
| 12 | 15 | "path/filepath" | |
| 13 | 16 | "strconv" | |
| 14 | 17 | "strings" | |
| 15 | 18 | "sync" | |
| 16 | 19 | "sync/atomic" | |
| 20 | + | "syscall" | |
| 17 | 21 | "testing" | |
| 18 | 22 | "time" | |
| 19 | 23 | ) | |
| @@ -24,12 +28,23 @@ import ( | |||
| 24 | 28 | ||
| 25 | 29 | const fakeServerEnv = "ZUT_MCP_FAKE_SERVER" | |
| 26 | 30 | ||
| 31 | + | // fakeToolCount is the number of tools the fake server lists in total. | |
| 32 | + | const fakeToolCount = 17 | |
| 33 | + | ||
| 27 | 34 | var fakeImageBytes = []byte("\x89PNG\r\n\x1a\nfake-image-bytes") | |
| 28 | 35 | ||
| 29 | 36 | func TestMain(m *testing.M) { | |
| 30 | - | if os.Getenv(fakeServerEnv) == "1" { | |
| 37 | + | switch os.Getenv(fakeServerEnv) { | |
| 38 | + | case "1": | |
| 31 | 39 | runFakeStdioServer() | |
| 32 | 40 | os.Exit(0) | |
| 41 | + | case "sleeper", "sleeper-stubborn": | |
| 42 | + | // A descendant of the fake server that outlives it. | |
| 43 | + | if os.Getenv(fakeServerEnv) == "sleeper-stubborn" { | |
| 44 | + | signal.Ignore(syscall.SIGTERM) | |
| 45 | + | } | |
| 46 | + | time.Sleep(2 * time.Minute) | |
| 47 | + | os.Exit(0) | |
| 33 | 48 | } | |
| 34 | 49 | os.Exit(m.Run()) | |
| 35 | 50 | } | |
| @@ -41,14 +56,22 @@ type fakeBackend struct { | |||
| 41 | 56 | protocol string | |
| 42 | 57 | repeatPage bool | |
| 43 | 58 | extraTool atomic.Bool | |
| 59 | + | // manyResources makes resources/list return this many long entries. | |
| 60 | + | manyResources int | |
| 61 | + | events string | |
| 44 | 62 | } | |
| 45 | 63 | ||
| 46 | 64 | func fakeBackendFromEnv() *fakeBackend { | |
| 47 | - | return &fakeBackend{ | |
| 65 | + | b := &fakeBackend{ | |
| 48 | 66 | noResources: os.Getenv("FAKE_NO_RESOURCES") == "1", | |
| 49 | 67 | protocol: os.Getenv("FAKE_PROTO"), | |
| 50 | 68 | repeatPage: os.Getenv("FAKE_BAD_CURSOR") == "1", | |
| 69 | + | events: os.Getenv("FAKE_EVENT_FILE"), | |
| 51 | 70 | } | |
| 71 | + | if n, err := strconv.Atoi(os.Getenv("FAKE_MANY_RESOURCES")); err == nil { | |
| 72 | + | b.manyResources = n | |
| 73 | + | } | |
| 74 | + | return b | |
| 52 | 75 | } | |
| 53 | 76 | ||
| 54 | 77 | func fakeToolDef(name, desc string, readOnly bool) map[string]any { | |
| @@ -92,6 +115,9 @@ func (b *fakeBackend) toolPage(cursor string) map[string]any { | |||
| 92 | 115 | fakeToolDef("add_tool", "Add a tool and notify list_changed", false), | |
| 93 | 116 | fakeToolDef("getenv", "Read an environment variable", false), | |
| 94 | 117 | fakeToolDef("pid", "Return the server process id", false), | |
| 118 | + | fakeToolDef("spawn", "Start a descendant process and return its pid", false), | |
| 119 | + | fakeToolDef("stall", "Stop reading stdin", false), | |
| 120 | + | fakeToolDef("stderr_flood", "Write a huge stderr line", false), | |
| 95 | 121 | fakeToolDef("cwd", "Return the server working directory", false), | |
| 96 | 122 | } | |
| 97 | 123 | if b.extraTool.Load() { | |
| @@ -184,6 +210,16 @@ func (b *fakeBackend) handle(method string, params json.RawMessage) (any, *RPCEr | |||
| 184 | 210 | if b.noResources { | |
| 185 | 211 | return nil, notFound | |
| 186 | 212 | } | |
| 213 | + | fakeEvent(b.events, "list-resources") | |
| 214 | + | if b.manyResources > 0 { | |
| 215 | + | list := make([]any, 0, b.manyResources) | |
| 216 | + | for i := 0; i < b.manyResources; i++ { | |
| 217 | + | list = append(list, map[string]any{ | |
| 218 | + | "uri": fmt.Sprintf("mem://item/%d", i), "name": "item", "description": strings.Repeat("d", 100), | |
| 219 | + | }) | |
| 220 | + | } | |
| 221 | + | return map[string]any{"resources": list}, nil | |
| 222 | + | } | |
| 187 | 223 | return map[string]any{"resources": []any{map[string]any{ | |
| 188 | 224 | "uri": "mem://a", "name": "A", "mimeType": "text/plain", "description": "first", | |
| 189 | 225 | }}}, nil | |
| @@ -191,6 +227,7 @@ func (b *fakeBackend) handle(method string, params json.RawMessage) (any, *RPCEr | |||
| 191 | 227 | if b.noResources { | |
| 192 | 228 | return nil, notFound | |
| 193 | 229 | } | |
| 230 | + | fakeEvent(b.events, "list-templates") | |
| 194 | 231 | return map[string]any{"resourceTemplates": []any{map[string]any{ | |
| 195 | 232 | "uriTemplate": "mem://item/{id}", "name": "Item", | |
| 196 | 233 | }}}, nil | |
| @@ -287,11 +324,15 @@ func runFakeStdioServer() { | |||
| 287 | 324 | time.Sleep(ms) | |
| 288 | 325 | } | |
| 289 | 326 | if os.Getenv("FAKE_INIT_ERR") == "1" { | |
| 290 | - | fmt.Fprintln(os.Stderr, "fatal: missing credentials") | |
| 291 | - | // Give the host's stderr reader time to drain the line | |
| 292 | - | // before the failure is reported. | |
| 293 | - | time.Sleep(250 * time.Millisecond) | |
| 294 | - | respond(msg.ID, nil, &RPCError{Code: -32000, Message: "init refused"}) | |
| 327 | + | line := "fatal: missing credentials" | |
| 328 | + | if echo := os.Getenv("FAKE_ERR_ECHO"); echo != "" { | |
| 329 | + | line += " token=" + echo + " see http://example.invalid/x?k=" + echo | |
| 330 | + | } | |
| 331 | + | if os.Getenv("FAKE_ANSI") == "1" { | |
| 332 | + | line += " \x1b[31mred\x1b[0m\x1b]0;evil-title\x07 end" | |
| 333 | + | } | |
| 334 | + | fmt.Fprintln(os.Stderr, line) | |
| 335 | + | respond(msg.ID, nil, &RPCError{Code: -32000, Message: "init refused " + os.Getenv("FAKE_ERR_ECHO")}) | |
| 295 | 336 | return | |
| 296 | 337 | } | |
| 297 | 338 | } | |
| @@ -313,6 +354,24 @@ func runFakeStdioServer() { | |||
| 313 | 354 | case "exit": | |
| 314 | 355 | fakeEvent(events, "exit-tool") | |
| 315 | 356 | os.Exit(0) | |
| 357 | + | case "spawn": | |
| 358 | + | var in struct { | |
| 359 | + | Mode string `json:"mode"` | |
| 360 | + | } | |
| 361 | + | _ = json.Unmarshal(p.Args, &in) | |
| 362 | + | child := exec.Command(os.Args[0], "-test.run=^$") | |
| 363 | + | child.Env = append(os.Environ(), fakeServerEnv+"="+in.Mode) | |
| 364 | + | if err := child.Start(); err != nil { | |
| 365 | + | respond(msg.ID, nil, &RPCError{Code: -32000, Message: err.Error()}) | |
| 366 | + | return | |
| 367 | + | } | |
| 368 | + | respond(msg.ID, fakeCallResult(fakeText(strconv.Itoa(child.Process.Pid))), nil) | |
| 369 | + | return | |
| 370 | + | case "stderr_flood": | |
| 371 | + | os.Stderr.WriteString(strings.Repeat("x", 3<<20) + "\n") | |
| 372 | + | fmt.Fprintln(os.Stderr, "after-flood") | |
| 373 | + | respond(msg.ID, fakeCallResult(fakeText("flooded")), nil) | |
| 374 | + | return | |
| 316 | 375 | case "ask_roots": | |
| 317 | 376 | roots, ok1 := askClient("srv-roots", "roots/list") | |
| 318 | 377 | ping, ok2 := askClient("srv-ping", "ping") | |
| @@ -331,6 +390,17 @@ func runFakeStdioServer() { | |||
| 331 | 390 | } | |
| 332 | 391 | result, rpcErr := backend.handle(msg.Method, msg.Params) | |
| 333 | 392 | respond(msg.ID, result, rpcErr) | |
| 393 | + | if msg.Method == "tools/list" && os.Getenv("FAKE_NOTIFY_ON_LIST") == "1" { | |
| 394 | + | var p struct { | |
| 395 | + | Cursor string `json:"cursor"` | |
| 396 | + | } | |
| 397 | + | _ = json.Unmarshal(msg.Params, &p) | |
| 398 | + | if p.Cursor == "p2" && !backend.extraTool.Load() { | |
| 399 | + | // The tool set changes right after the initial listing. | |
| 400 | + | backend.extraTool.Store(true) | |
| 401 | + | send(map[string]any{"jsonrpc": "2.0", "method": "notifications/tools/list_changed"}) | |
| 402 | + | } | |
| 403 | + | } | |
| 334 | 404 | } | |
| 335 | 405 | ||
| 336 | 406 | sc := bufio.NewScanner(os.Stdin) | |
| @@ -350,6 +420,11 @@ func runFakeStdioServer() { | |||
| 350 | 420 | ch <- msg | |
| 351 | 421 | } | |
| 352 | 422 | case msg.isRequest(): | |
| 423 | + | if msg.Method == "tools/call" && bytes.Contains(msg.Params, []byte(`"name":"stall"`)) { | |
| 424 | + | // Stop reading stdin so the client's writes back up. | |
| 425 | + | fakeEvent(events, "stalled") | |
| 426 | + | time.Sleep(time.Minute) | |
| 427 | + | } | |
| 353 | 428 | go serve(msg) | |
| 354 | 429 | case msg.Method == "notifications/cancelled": | |
| 355 | 430 | var p struct { | |
| @@ -454,10 +529,23 @@ type httpFake struct { | |||
| 454 | 529 | wantAuth string | |
| 455 | 530 | sessionID string | |
| 456 | 531 | ||
| 532 | + | // expireNextCall makes the next tools/call that carries a session id | |
| 533 | + | // answer 404 (session expired). | |
| 534 | + | expireNextCall atomic.Bool | |
| 535 | + | // hangInit makes initialize block until the client goes away. | |
| 536 | + | hangInit bool | |
| 537 | + | // cancelGate, when set, makes notifications/cancelled block until it is | |
| 538 | + | // closed or the client disconnects. | |
| 539 | + | cancelGate chan struct{} | |
| 540 | + | // errorBody, when set, is returned with a 500 for every tools/list. | |
| 541 | + | errorBody string | |
| 542 | + | ||
| 457 | 543 | mu sync.Mutex | |
| 458 | 544 | records []httpRecord | |
| 459 | 545 | pingReplies chan message | |
| 460 | 546 | deleted chan string | |
| 547 | + | initCount int | |
| 548 | + | callCount int | |
| 461 | 549 | } | |
| 462 | 550 | ||
| 463 | 551 | func newHTTPFake(sse bool) *httpFake { | |
| @@ -476,6 +564,12 @@ func (f *httpFake) start(t *testing.T) *httptest.Server { | |||
| 476 | 564 | return srv | |
| 477 | 565 | } | |
| 478 | 566 | ||
| 567 | + | func (f *httpFake) counts() (inits, calls int) { | |
| 568 | + | f.mu.Lock() | |
| 569 | + | defer f.mu.Unlock() | |
| 570 | + | return f.initCount, f.callCount | |
| 571 | + | } | |
| 572 | + | ||
| 479 | 573 | func (f *httpFake) snapshot() []httpRecord { | |
| 480 | 574 | f.mu.Lock() | |
| 481 | 575 | defer f.mu.Unlock() | |
| @@ -535,6 +629,12 @@ func (f *httpFake) servePost(w http.ResponseWriter, r *http.Request, rec httpRec | |||
| 535 | 629 | return | |
| 536 | 630 | } | |
| 537 | 631 | if !msg.isRequest() { | |
| 632 | + | if msg.Method == "notifications/cancelled" && f.cancelGate != nil { | |
| 633 | + | select { | |
| 634 | + | case <-f.cancelGate: | |
| 635 | + | case <-r.Context().Done(): | |
| 636 | + | } | |
| 637 | + | } | |
| 538 | 638 | w.WriteHeader(http.StatusAccepted) | |
| 539 | 639 | return | |
| 540 | 640 | } | |
| @@ -544,7 +644,34 @@ func (f *httpFake) servePost(w http.ResponseWriter, r *http.Request, rec httpRec | |||
| 544 | 644 | } | |
| 545 | 645 | _ = json.Unmarshal(msg.Params, &p) | |
| 546 | 646 | if msg.Method == "initialize" { | |
| 547 | - | w.Header().Set("Mcp-Session-Id", f.sessionID) | |
| 647 | + | if f.hangInit { | |
| 648 | + | <-r.Context().Done() | |
| 649 | + | return | |
| 650 | + | } | |
| 651 | + | f.mu.Lock() | |
| 652 | + | f.initCount++ | |
| 653 | + | n := f.initCount | |
| 654 | + | f.mu.Unlock() | |
| 655 | + | sid := f.sessionID | |
| 656 | + | if n > 1 { | |
| 657 | + | sid = fmt.Sprintf("%s-%d", f.sessionID, n) | |
| 658 | + | } | |
| 659 | + | w.Header().Set("Mcp-Session-Id", sid) | |
| 660 | + | } | |
| 661 | + | if msg.Method == "tools/call" { | |
| 662 | + | f.mu.Lock() | |
| 663 | + | f.callCount++ | |
| 664 | + | f.mu.Unlock() | |
| 665 | + | if rec.Session != "" && f.expireNextCall.CompareAndSwap(true, false) { | |
| 666 | + | http.Error(w, "session not found", http.StatusNotFound) | |
| 667 | + | return | |
| 668 | + | } | |
| 669 | + | } | |
| 670 | + | if msg.Method == "tools/list" && f.errorBody != "" { | |
| 671 | + | w.Header().Set("Content-Type", "text/plain") | |
| 672 | + | w.WriteHeader(http.StatusInternalServerError) | |
| 673 | + | _, _ = w.Write([]byte(f.errorBody)) | |
| 674 | + | return | |
| 548 | 675 | } | |
| 549 | 676 | ||
| 550 | 677 | if msg.Method == "tools/call" && p.Name == "slow" { | |
| @@ -2,6 +2,7 @@ package mcp | |||||
| 2 | 2 | ||||
| 3 | import ( | 3 | import ( | ||
| 4 | "bufio" | 4 | "bufio" | ||
| 5 | + | "bytes" | |||
| 5 | "context" | 6 | "context" | ||
| 6 | "encoding/base64" | 7 | "encoding/base64" | ||
| 7 | "encoding/json" | 8 | "encoding/json" | ||
| @@ -9,11 +10,14 @@ import ( | |||||
| 9 | "net/http" | 10 | "net/http" | ||
| 10 | "net/http/httptest" | 11 | "net/http/httptest" | ||
| 11 | "os" | 12 | "os" | ||
| 13 | + | "os/exec" | |||
| 14 | + | "os/signal" | |||
| 12 | "path/filepath" | 15 | "path/filepath" | ||
| 13 | "strconv" | 16 | "strconv" | ||
| 14 | "strings" | 17 | "strings" | ||
| 15 | "sync" | 18 | "sync" | ||
| 16 | "sync/atomic" | 19 | "sync/atomic" | ||
| 20 | + | "syscall" | |||
| 17 | "testing" | 21 | "testing" | ||
| 18 | "time" | 22 | "time" | ||
| 19 | ) | 23 | ) | ||
| @@ -24,12 +28,23 @@ import ( | |||||
| 24 | 28 | ||||
| 25 | const fakeServerEnv = "ZUT_MCP_FAKE_SERVER" | 29 | const fakeServerEnv = "ZUT_MCP_FAKE_SERVER" | ||
| 26 | 30 | ||||
| 31 | + | // fakeToolCount is the number of tools the fake server lists in total. | |||
| 32 | + | const fakeToolCount = 17 | |||
| 33 | + | ||||
| 27 | var fakeImageBytes = []byte("\x89PNG\r\n\x1a\nfake-image-bytes") | 34 | var fakeImageBytes = []byte("\x89PNG\r\n\x1a\nfake-image-bytes") | ||
| 28 | 35 | ||||
| 29 | func TestMain(m *testing.M) { | 36 | func TestMain(m *testing.M) { | ||
| 30 | - | if os.Getenv(fakeServerEnv) == "1" { | 37 | + | switch os.Getenv(fakeServerEnv) { |
| 38 | + | case "1": | |||
| 31 | runFakeStdioServer() | 39 | runFakeStdioServer() | ||
| 32 | os.Exit(0) | 40 | os.Exit(0) | ||
| 41 | + | case "sleeper", "sleeper-stubborn": | |||
| 42 | + | // A descendant of the fake server that outlives it. | |||
| 43 | + | if os.Getenv(fakeServerEnv) == "sleeper-stubborn" { | |||
| 44 | + | signal.Ignore(syscall.SIGTERM) | |||
| 45 | + | } | |||
| 46 | + | time.Sleep(2 * time.Minute) | |||
| 47 | + | os.Exit(0) | |||
| 33 | } | 48 | } | ||
| 34 | os.Exit(m.Run()) | 49 | os.Exit(m.Run()) | ||
| 35 | } | 50 | } | ||
| @@ -41,14 +56,22 @@ type fakeBackend struct { | |||||
| 41 | protocol string | 56 | protocol string | ||
| 42 | repeatPage bool | 57 | repeatPage bool | ||
| 43 | extraTool atomic.Bool | 58 | extraTool atomic.Bool | ||
| 59 | + | // manyResources makes resources/list return this many long entries. | |||
| 60 | + | manyResources int | |||
| 61 | + | events string | |||
| 44 | } | 62 | } | ||
| 45 | 63 | ||||
| 46 | func fakeBackendFromEnv() *fakeBackend { | 64 | func fakeBackendFromEnv() *fakeBackend { | ||
| 47 | - | return &fakeBackend{ | 65 | + | b := &fakeBackend{ |
| 48 | noResources: os.Getenv("FAKE_NO_RESOURCES") == "1", | 66 | noResources: os.Getenv("FAKE_NO_RESOURCES") == "1", | ||
| 49 | protocol: os.Getenv("FAKE_PROTO"), | 67 | protocol: os.Getenv("FAKE_PROTO"), | ||
| 50 | repeatPage: os.Getenv("FAKE_BAD_CURSOR") == "1", | 68 | repeatPage: os.Getenv("FAKE_BAD_CURSOR") == "1", | ||
| 69 | + | events: os.Getenv("FAKE_EVENT_FILE"), | |||
| 51 | } | 70 | } | ||
| 71 | + | if n, err := strconv.Atoi(os.Getenv("FAKE_MANY_RESOURCES")); err == nil { | |||
| 72 | + | b.manyResources = n | |||
| 73 | + | } | |||
| 74 | + | return b | |||
| 52 | } | 75 | } | ||
| 53 | 76 | ||||
| 54 | func fakeToolDef(name, desc string, readOnly bool) map[string]any { | 77 | func fakeToolDef(name, desc string, readOnly bool) map[string]any { | ||
| @@ -92,6 +115,9 @@ func (b *fakeBackend) toolPage(cursor string) map[string]any { | |||||
| 92 | fakeToolDef("add_tool", "Add a tool and notify list_changed", false), | 115 | fakeToolDef("add_tool", "Add a tool and notify list_changed", false), | ||
| 93 | fakeToolDef("getenv", "Read an environment variable", false), | 116 | fakeToolDef("getenv", "Read an environment variable", false), | ||
| 94 | fakeToolDef("pid", "Return the server process id", false), | 117 | fakeToolDef("pid", "Return the server process id", false), | ||
| 118 | + | fakeToolDef("spawn", "Start a descendant process and return its pid", false), | |||
| 119 | + | fakeToolDef("stall", "Stop reading stdin", false), | |||
| 120 | + | fakeToolDef("stderr_flood", "Write a huge stderr line", false), | |||
| 95 | fakeToolDef("cwd", "Return the server working directory", false), | 121 | fakeToolDef("cwd", "Return the server working directory", false), | ||
| 96 | } | 122 | } | ||
| 97 | if b.extraTool.Load() { | 123 | if b.extraTool.Load() { | ||
| @@ -184,6 +210,16 @@ func (b *fakeBackend) handle(method string, params json.RawMessage) (any, *RPCEr | |||||
| 184 | if b.noResources { | 210 | if b.noResources { | ||
| 185 | return nil, notFound | 211 | return nil, notFound | ||
| 186 | } | 212 | } | ||
| 213 | + | fakeEvent(b.events, "list-resources") | |||
| 214 | + | if b.manyResources > 0 { | |||
| 215 | + | list := make([]any, 0, b.manyResources) | |||
| 216 | + | for i := 0; i < b.manyResources; i++ { | |||
| 217 | + | list = append(list, map[string]any{ | |||
| 218 | + | "uri": fmt.Sprintf("mem://item/%d", i), "name": "item", "description": strings.Repeat("d", 100), | |||
| 219 | + | }) | |||
| 220 | + | } | |||
| 221 | + | return map[string]any{"resources": list}, nil | |||
| 222 | + | } | |||
| 187 | return map[string]any{"resources": []any{map[string]any{ | 223 | return map[string]any{"resources": []any{map[string]any{ | ||
| 188 | "uri": "mem://a", "name": "A", "mimeType": "text/plain", "description": "first", | 224 | "uri": "mem://a", "name": "A", "mimeType": "text/plain", "description": "first", | ||
| 189 | }}}, nil | 225 | }}}, nil | ||
| @@ -191,6 +227,7 @@ func (b *fakeBackend) handle(method string, params json.RawMessage) (any, *RPCEr | |||||
| 191 | if b.noResources { | 227 | if b.noResources { | ||
| 192 | return nil, notFound | 228 | return nil, notFound | ||
| 193 | } | 229 | } | ||
| 230 | + | fakeEvent(b.events, "list-templates") | |||
| 194 | return map[string]any{"resourceTemplates": []any{map[string]any{ | 231 | return map[string]any{"resourceTemplates": []any{map[string]any{ | ||
| 195 | "uriTemplate": "mem://item/{id}", "name": "Item", | 232 | "uriTemplate": "mem://item/{id}", "name": "Item", | ||
| 196 | }}}, nil | 233 | }}}, nil | ||
| @@ -287,11 +324,15 @@ func runFakeStdioServer() { | |||||
| 287 | time.Sleep(ms) | 324 | time.Sleep(ms) | ||
| 288 | } | 325 | } | ||
| 289 | if os.Getenv("FAKE_INIT_ERR") == "1" { | 326 | if os.Getenv("FAKE_INIT_ERR") == "1" { | ||
| 290 | - | fmt.Fprintln(os.Stderr, "fatal: missing credentials") | 327 | + | line := "fatal: missing credentials" |
| 291 | - | // Give the host's stderr reader time to drain the line | 328 | + | if echo := os.Getenv("FAKE_ERR_ECHO"); echo != "" { |
| 292 | - | // before the failure is reported. | 329 | + | line += " token=" + echo + " see http://example.invalid/x?k=" + echo |
| 293 | - | time.Sleep(250 * time.Millisecond) | 330 | + | } |
| 294 | - | respond(msg.ID, nil, &RPCError{Code: -32000, Message: "init refused"}) | 331 | + | if os.Getenv("FAKE_ANSI") == "1" { |
| 332 | + | line += " \x1b[31mred\x1b[0m\x1b]0;evil-title\x07 end" | |||
| 333 | + | } | |||
| 334 | + | fmt.Fprintln(os.Stderr, line) | |||
| 335 | + | respond(msg.ID, nil, &RPCError{Code: -32000, Message: "init refused " + os.Getenv("FAKE_ERR_ECHO")}) | |||
| 295 | return | 336 | return | ||
| 296 | } | 337 | } | ||
| 297 | } | 338 | } | ||
| @@ -313,6 +354,24 @@ func runFakeStdioServer() { | |||||
| 313 | case "exit": | 354 | case "exit": | ||
| 314 | fakeEvent(events, "exit-tool") | 355 | fakeEvent(events, "exit-tool") | ||
| 315 | os.Exit(0) | 356 | os.Exit(0) | ||
| 357 | + | case "spawn": | |||
| 358 | + | var in struct { | |||
| 359 | + | Mode string `json:"mode"` | |||
| 360 | + | } | |||
| 361 | + | _ = json.Unmarshal(p.Args, &in) | |||
| 362 | + | child := exec.Command(os.Args[0], "-test.run=^$") | |||
| 363 | + | child.Env = append(os.Environ(), fakeServerEnv+"="+in.Mode) | |||
| 364 | + | if err := child.Start(); err != nil { | |||
| 365 | + | respond(msg.ID, nil, &RPCError{Code: -32000, Message: err.Error()}) | |||
| 366 | + | return | |||
| 367 | + | } | |||
| 368 | + | respond(msg.ID, fakeCallResult(fakeText(strconv.Itoa(child.Process.Pid))), nil) | |||
| 369 | + | return | |||
| 370 | + | case "stderr_flood": | |||
| 371 | + | os.Stderr.WriteString(strings.Repeat("x", 3<<20) + "\n") | |||
| 372 | + | fmt.Fprintln(os.Stderr, "after-flood") | |||
| 373 | + | respond(msg.ID, fakeCallResult(fakeText("flooded")), nil) | |||
| 374 | + | return | |||
| 316 | case "ask_roots": | 375 | case "ask_roots": | ||
| 317 | roots, ok1 := askClient("srv-roots", "roots/list") | 376 | roots, ok1 := askClient("srv-roots", "roots/list") | ||
| 318 | ping, ok2 := askClient("srv-ping", "ping") | 377 | ping, ok2 := askClient("srv-ping", "ping") | ||
| @@ -331,6 +390,17 @@ func runFakeStdioServer() { | |||||
| 331 | } | 390 | } | ||
| 332 | result, rpcErr := backend.handle(msg.Method, msg.Params) | 391 | result, rpcErr := backend.handle(msg.Method, msg.Params) | ||
| 333 | respond(msg.ID, result, rpcErr) | 392 | respond(msg.ID, result, rpcErr) | ||
| 393 | + | if msg.Method == "tools/list" && os.Getenv("FAKE_NOTIFY_ON_LIST") == "1" { | |||
| 394 | + | var p struct { | |||
| 395 | + | Cursor string `json:"cursor"` | |||
| 396 | + | } | |||
| 397 | + | _ = json.Unmarshal(msg.Params, &p) | |||
| 398 | + | if p.Cursor == "p2" && !backend.extraTool.Load() { | |||
| 399 | + | // The tool set changes right after the initial listing. | |||
| 400 | + | backend.extraTool.Store(true) | |||
| 401 | + | send(map[string]any{"jsonrpc": "2.0", "method": "notifications/tools/list_changed"}) | |||
| 402 | + | } | |||
| 403 | + | } | |||
| 334 | } | 404 | } | ||
| 335 | 405 | ||||
| 336 | sc := bufio.NewScanner(os.Stdin) | 406 | sc := bufio.NewScanner(os.Stdin) | ||
| @@ -350,6 +420,11 @@ func runFakeStdioServer() { | |||||
| 350 | ch <- msg | 420 | ch <- msg | ||
| 351 | } | 421 | } | ||
| 352 | case msg.isRequest(): | 422 | case msg.isRequest(): | ||
| 423 | + | if msg.Method == "tools/call" && bytes.Contains(msg.Params, []byte(`"name":"stall"`)) { | |||
| 424 | + | // Stop reading stdin so the client's writes back up. | |||
| 425 | + | fakeEvent(events, "stalled") | |||
| 426 | + | time.Sleep(time.Minute) | |||
| 427 | + | } | |||
| 353 | go serve(msg) | 428 | go serve(msg) | ||
| 354 | case msg.Method == "notifications/cancelled": | 429 | case msg.Method == "notifications/cancelled": | ||
| 355 | var p struct { | 430 | var p struct { | ||
| @@ -454,10 +529,23 @@ type httpFake struct { | |||||
| 454 | wantAuth string | 529 | wantAuth string | ||
| 455 | sessionID string | 530 | sessionID string | ||
| 456 | 531 | ||||
| 532 | + | // expireNextCall makes the next tools/call that carries a session id | |||
| 533 | + | // answer 404 (session expired). | |||
| 534 | + | expireNextCall atomic.Bool | |||
| 535 | + | // hangInit makes initialize block until the client goes away. | |||
| 536 | + | hangInit bool | |||
| 537 | + | // cancelGate, when set, makes notifications/cancelled block until it is | |||
| 538 | + | // closed or the client disconnects. | |||
| 539 | + | cancelGate chan struct{} | |||
| 540 | + | // errorBody, when set, is returned with a 500 for every tools/list. | |||
| 541 | + | errorBody string | |||
| 542 | + | ||||
| 457 | mu sync.Mutex | 543 | mu sync.Mutex | ||
| 458 | records []httpRecord | 544 | records []httpRecord | ||
| 459 | pingReplies chan message | 545 | pingReplies chan message | ||
| 460 | deleted chan string | 546 | deleted chan string | ||
| 547 | + | initCount int | |||
| 548 | + | callCount int | |||
| 461 | } | 549 | } | ||
| 462 | 550 | ||||
| 463 | func newHTTPFake(sse bool) *httpFake { | 551 | func newHTTPFake(sse bool) *httpFake { | ||
| @@ -476,6 +564,12 @@ func (f *httpFake) start(t *testing.T) *httptest.Server { | |||||
| 476 | return srv | 564 | return srv | ||
| 477 | } | 565 | } | ||
| 478 | 566 | ||||
| 567 | + | func (f *httpFake) counts() (inits, calls int) { | |||
| 568 | + | f.mu.Lock() | |||
| 569 | + | defer f.mu.Unlock() | |||
| 570 | + | return f.initCount, f.callCount | |||
| 571 | + | } | |||
| 572 | + | ||||
| 479 | func (f *httpFake) snapshot() []httpRecord { | 573 | func (f *httpFake) snapshot() []httpRecord { | ||
| 480 | f.mu.Lock() | 574 | f.mu.Lock() | ||
| 481 | defer f.mu.Unlock() | 575 | defer f.mu.Unlock() | ||
| @@ -535,6 +629,12 @@ func (f *httpFake) servePost(w http.ResponseWriter, r *http.Request, rec httpRec | |||||
| 535 | return | 629 | return | ||
| 536 | } | 630 | } | ||
| 537 | if !msg.isRequest() { | 631 | if !msg.isRequest() { | ||
| 632 | + | if msg.Method == "notifications/cancelled" && f.cancelGate != nil { | |||
| 633 | + | select { | |||
| 634 | + | case <-f.cancelGate: | |||
| 635 | + | case <-r.Context().Done(): | |||
| 636 | + | } | |||
| 637 | + | } | |||
| 538 | w.WriteHeader(http.StatusAccepted) | 638 | w.WriteHeader(http.StatusAccepted) | ||
| 539 | return | 639 | return | ||
| 540 | } | 640 | } | ||
| @@ -544,7 +644,34 @@ func (f *httpFake) servePost(w http.ResponseWriter, r *http.Request, rec httpRec | |||||
| 544 | } | 644 | } | ||
| 545 | _ = json.Unmarshal(msg.Params, &p) | 645 | _ = json.Unmarshal(msg.Params, &p) | ||
| 546 | if msg.Method == "initialize" { | 646 | if msg.Method == "initialize" { | ||
| 547 | - | w.Header().Set("Mcp-Session-Id", f.sessionID) | 647 | + | if f.hangInit { |
| 648 | + | <-r.Context().Done() | |||
| 649 | + | return | |||
| 650 | + | } | |||
| 651 | + | f.mu.Lock() | |||
| 652 | + | f.initCount++ | |||
| 653 | + | n := f.initCount | |||
| 654 | + | f.mu.Unlock() | |||
| 655 | + | sid := f.sessionID | |||
| 656 | + | if n > 1 { | |||
| 657 | + | sid = fmt.Sprintf("%s-%d", f.sessionID, n) | |||
| 658 | + | } | |||
| 659 | + | w.Header().Set("Mcp-Session-Id", sid) | |||
| 660 | + | } | |||
| 661 | + | if msg.Method == "tools/call" { | |||
| 662 | + | f.mu.Lock() | |||
| 663 | + | f.callCount++ | |||
| 664 | + | f.mu.Unlock() | |||
| 665 | + | if rec.Session != "" && f.expireNextCall.CompareAndSwap(true, false) { | |||
| 666 | + | http.Error(w, "session not found", http.StatusNotFound) | |||
| 667 | + | return | |||
| 668 | + | } | |||
| 669 | + | } | |||
| 670 | + | if msg.Method == "tools/list" && f.errorBody != "" { | |||
| 671 | + | w.Header().Set("Content-Type", "text/plain") | |||
| 672 | + | w.WriteHeader(http.StatusInternalServerError) | |||
| 673 | + | _, _ = w.Write([]byte(f.errorBody)) | |||
| 674 | + | return | |||
| 548 | } | 675 | } | ||
| 549 | 676 | ||||
| 550 | if msg.Method == "tools/call" && p.Name == "slow" { | 677 | if msg.Method == "tools/call" && p.Name == "slow" { | ||
A packages/agent/mcp/hardening_test.go +734 -0
| @@ -0,0 +1,734 @@ | |||
| 1 | + | package mcp | |
| 2 | + | ||
| 3 | + | import ( | |
| 4 | + | "context" | |
| 5 | + | "errors" | |
| 6 | + | "net/http" | |
| 7 | + | "net/http/httptest" | |
| 8 | + | "net/url" | |
| 9 | + | "strconv" | |
| 10 | + | "strings" | |
| 11 | + | "sync/atomic" | |
| 12 | + | "testing" | |
| 13 | + | "time" | |
| 14 | + | ) | |
| 15 | + | ||
| 16 | + | // Regression tests for the review findings: write path, SSE limits, session | |
| 17 | + | // expiry, diagnostics hygiene, process cleanup, refresh coalescing, and the | |
| 18 | + | // cancellation rules. | |
| 19 | + | ||
| 20 | + | func connectFake(t *testing.T, extraEnv map[string]string) (*Client, string) { | |
| 21 | + | t.Helper() | |
| 22 | + | cfg, events := fakeStdioConfig(t, extraEnv) | |
| 23 | + | ctx, cancel := context.WithTimeout(context.Background(), 20*time.Second) | |
| 24 | + | defer cancel() | |
| 25 | + | c, err := Connect(ctx, "fake", cfg, "", ClientInfo{Name: "t"}, nil) | |
| 26 | + | if err != nil { | |
| 27 | + | t.Fatalf("Connect: %v", err) | |
| 28 | + | } | |
| 29 | + | t.Cleanup(func() { _ = c.Close() }) | |
| 30 | + | return c, events | |
| 31 | + | } | |
| 32 | + | ||
| 33 | + | func callText(t *testing.T, c *Client, tool, args string) string { | |
| 34 | + | t.Helper() | |
| 35 | + | ctx, cancel := context.WithTimeout(context.Background(), 20*time.Second) | |
| 36 | + | defer cancel() | |
| 37 | + | res, err := c.CallTool(ctx, tool, []byte(args)) | |
| 38 | + | if err != nil { | |
| 39 | + | t.Fatalf("CallTool %s: %v", tool, err) | |
| 40 | + | } | |
| 41 | + | if len(res.Content) == 0 { | |
| 42 | + | t.Fatalf("CallTool %s: no content", tool) | |
| 43 | + | } | |
| 44 | + | return res.Content[0].Text | |
| 45 | + | } | |
| 46 | + | ||
| 47 | + | // ---- 1, 8: cancellable writes, close behind a stuck writer ---- | |
| 48 | + | ||
| 49 | + | func TestStdioWriteStuckChildDoesNotWedgeCancelOrClose(t *testing.T) { | |
| 50 | + | c, events := connectFake(t, nil) | |
| 51 | + | ||
| 52 | + | // The server stops reading stdin after receiving this call. | |
| 53 | + | stallCtx, stallCancel := context.WithCancel(context.Background()) | |
| 54 | + | stallDone := make(chan error, 1) | |
| 55 | + | go func() { | |
| 56 | + | _, err := c.CallTool(stallCtx, "stall", []byte(`{}`)) | |
| 57 | + | stallDone <- err | |
| 58 | + | }() | |
| 59 | + | waitEvent(t, events, "stalled") | |
| 60 | + | ||
| 61 | + | // A frame far larger than the pipe buffer now blocks the writer. | |
| 62 | + | big := `{"text":"` + strings.Repeat("a", 4<<20) + `"}` | |
| 63 | + | ctx, cancel := context.WithTimeout(context.Background(), 300*time.Millisecond) | |
| 64 | + | defer cancel() | |
| 65 | + | start := time.Now() | |
| 66 | + | _, err := c.CallTool(ctx, "echo", []byte(big)) | |
| 67 | + | if !errors.Is(err, context.DeadlineExceeded) { | |
| 68 | + | t.Fatalf("err = %v, want deadline exceeded", err) | |
| 69 | + | } | |
| 70 | + | if d := time.Since(start); d > 3*time.Second { | |
| 71 | + | t.Fatalf("request ignored its context for %v", d) | |
| 72 | + | } | |
| 73 | + | ||
| 74 | + | closed := make(chan struct{}) | |
| 75 | + | go func() { _ = c.Close(); close(closed) }() | |
| 76 | + | select { | |
| 77 | + | case <-closed: | |
| 78 | + | case <-time.After(10 * time.Second): | |
| 79 | + | t.Fatal("Close hung behind a blocked write") | |
| 80 | + | } | |
| 81 | + | select { | |
| 82 | + | case <-c.Done(): | |
| 83 | + | default: | |
| 84 | + | t.Error("Done not closed after Close") | |
| 85 | + | } | |
| 86 | + | stallCancel() | |
| 87 | + | select { | |
| 88 | + | case <-stallDone: | |
| 89 | + | case <-time.After(5 * time.Second): | |
| 90 | + | t.Fatal("pending call not released by Close") | |
| 91 | + | } | |
| 92 | + | } | |
| 93 | + | ||
| 94 | + | // ---- 2: SSE aggregate limit ---- | |
| 95 | + | ||
| 96 | + | func TestReadSSEFramesAggregateLimit(t *testing.T) { | |
| 97 | + | line := "data: " + strings.Repeat("x", 9) + "\n" // 9 bytes of data per line | |
| 98 | + | var delivered int | |
| 99 | + | handle := func(message) bool { delivered++; return false } | |
| 100 | + | ||
| 101 | + | // Many short lines of one event exceed the limit. | |
| 102 | + | err := readSSEFramesLimit(strings.NewReader(strings.Repeat(line, 200)+"\n"), 1000, handle) | |
| 103 | + | var tooLarge errSSEEventTooLarge | |
| 104 | + | if !errors.As(err, &tooLarge) || !strings.Contains(err.Error(), "1000") { | |
| 105 | + | t.Fatalf("err = %v, want errSSEEventTooLarge", err) | |
| 106 | + | } | |
| 107 | + | if delivered != 0 { | |
| 108 | + | t.Errorf("oversized event delivered %d frames", delivered) | |
| 109 | + | } | |
| 110 | + | ||
| 111 | + | // The limit is per event: a blank line resets it. | |
| 112 | + | ok := `data: {"jsonrpc":"2.0","method":"a"}` + "\n\n" | |
| 113 | + | delivered = 0 | |
| 114 | + | if err := readSSEFramesLimit(strings.NewReader(strings.Repeat(ok, 50)), 100, handle); err != nil { | |
| 115 | + | t.Fatal(err) | |
| 116 | + | } | |
| 117 | + | if delivered != 50 { | |
| 118 | + | t.Errorf("delivered %d, want 50", delivered) | |
| 119 | + | } | |
| 120 | + | ||
| 121 | + | // A single line longer than the limit is the same explicit error. | |
| 122 | + | err = readSSEFramesLimit(strings.NewReader("data: "+strings.Repeat("y", 5000)+"\n\n"), 1000, handle) | |
| 123 | + | if !errors.As(err, &tooLarge) { | |
| 124 | + | t.Errorf("long line err = %v", err) | |
| 125 | + | } | |
| 126 | + | } | |
| 127 | + | ||
| 128 | + | func TestSSEEndlessEventStopsAtLimit(t *testing.T) { | |
| 129 | + | // A server that never ends its event must not grow memory without bound. | |
| 130 | + | err := readSSEFramesLimit(&infiniteSSE{}, 64<<10, func(message) bool { return false }) | |
| 131 | + | var tooLarge errSSEEventTooLarge | |
| 132 | + | if !errors.As(err, &tooLarge) { | |
| 133 | + | t.Fatalf("err = %v", err) | |
| 134 | + | } | |
| 135 | + | } | |
| 136 | + | ||
| 137 | + | type infiniteSSE struct{} | |
| 138 | + | ||
| 139 | + | func (*infiniteSSE) Read(p []byte) (int, error) { | |
| 140 | + | line := []byte("data: " + strings.Repeat("q", 100) + "\n") | |
| 141 | + | n := copy(p, line) | |
| 142 | + | return n, nil | |
| 143 | + | } | |
| 144 | + | ||
| 145 | + | // ---- 3: session expiry ---- | |
| 146 | + | ||
| 147 | + | func TestHTTPSessionExpiredReconnects(t *testing.T) { | |
| 148 | + | for _, mode := range []struct { | |
| 149 | + | name string | |
| 150 | + | sse bool | |
| 151 | + | }{{"json", false}, {"sse", true}} { | |
| 152 | + | t.Run(mode.name, func(t *testing.T) { | |
| 153 | + | fake := newHTTPFake(mode.sse) | |
| 154 | + | m, failures := startHTTP(t, fake, nil) | |
| 155 | + | if len(failures) != 0 { | |
| 156 | + | t.Fatal(failures) | |
| 157 | + | } | |
| 158 | + | echo := findTool(t, m, "mcp__web__echo") | |
| 159 | + | fake.expireNextCall.Store(true) | |
| 160 | + | res := execTool(t, echo, `{"text":"one"}`) | |
| 161 | + | if !res.IsError || !strings.Contains(resultText(res), "session expired") { | |
| 162 | + | t.Fatalf("expired call = %q (error=%v)", resultText(res), res.IsError) | |
| 163 | + | } | |
| 164 | + | // The transport is unusable, so the status reflects it at once. | |
| 165 | + | if st := m.Status()[0]; st.Connected { | |
| 166 | + | t.Errorf("still marked connected after 404: %+v", st) | |
| 167 | + | } | |
| 168 | + | // The next call reconnects (fresh initialize) and succeeds. | |
| 169 | + | if got := resultText(execTool(t, echo, `{"text":"two"}`)); got != "echo: two" { | |
| 170 | + | t.Fatalf("after reconnect = %q", got) | |
| 171 | + | } | |
| 172 | + | inits, calls := fake.counts() | |
| 173 | + | if inits != 2 { | |
| 174 | + | t.Errorf("initialize count = %d, want 2", inits) | |
| 175 | + | } | |
| 176 | + | if calls != 2 { | |
| 177 | + | t.Errorf("tools/call count = %d, want 2 (the failed call must not be replayed)", calls) | |
| 178 | + | } | |
| 179 | + | // No DELETE for the dead session. | |
| 180 | + | select { | |
| 181 | + | case sid := <-fake.deleted: | |
| 182 | + | t.Errorf("DELETE sent for expired session %q", sid) | |
| 183 | + | default: | |
| 184 | + | } | |
| 185 | + | // The new session carries its own id. | |
| 186 | + | var last httpRecord | |
| 187 | + | for _, r := range fake.snapshot() { | |
| 188 | + | if r.RPCMethod == "tools/call" { | |
| 189 | + | last = r | |
| 190 | + | } | |
| 191 | + | } | |
| 192 | + | if last.Session != "sess-abc-2" { | |
| 193 | + | t.Errorf("session after reconnect = %q", last.Session) | |
| 194 | + | } | |
| 195 | + | }) | |
| 196 | + | } | |
| 197 | + | } | |
| 198 | + | ||
| 199 | + | func TestHTTPNotFoundWithoutSessionIsOrdinaryError(t *testing.T) { | |
| 200 | + | srv := httptest.NewServer(http.NotFoundHandler()) | |
| 201 | + | defer srv.Close() | |
| 202 | + | tr, err := newHTTPTransport("x", ServerConfig{URL: srv.URL}, nil) | |
| 203 | + | if err != nil { | |
| 204 | + | t.Fatal(err) | |
| 205 | + | } | |
| 206 | + | defer tr.close() | |
| 207 | + | _, err = tr.request(context.Background(), 1, "initialize", nil) | |
| 208 | + | var he *HTTPStatusError | |
| 209 | + | if !errors.As(err, &he) || he.Status != 404 { | |
| 210 | + | t.Fatalf("err = %v", err) | |
| 211 | + | } | |
| 212 | + | select { | |
| 213 | + | case <-tr.done(): | |
| 214 | + | t.Error("a 404 before any session must not poison the transport") | |
| 215 | + | default: | |
| 216 | + | } | |
| 217 | + | } | |
| 218 | + | ||
| 219 | + | // ---- 4, 13: diagnostics hygiene ---- | |
| 220 | + | ||
| 221 | + | func TestScrubURLsAndErrors(t *testing.T) { | |
| 222 | + | red := newRedactor(ServerConfig{}) | |
| 223 | + | uerr := &url.Error{Op: "Post", URL: "https://user:pw@h.example:8443/mcp/v1?token=abc&x=1#frag", Err: errors.New("connection refused")} | |
| 224 | + | got := red.scrubError(uerr).Error() | |
| 225 | + | if want := `Post "https://h.example:8443/mcp/v1": connection refused`; got != want { | |
| 226 | + | t.Errorf("got %q, want %q", got, want) | |
| 227 | + | } | |
| 228 | + | var target *url.Error | |
| 229 | + | if !errors.As(red.scrubError(uerr), &target) { | |
| 230 | + | t.Error("original error must stay reachable via errors.As") | |
| 231 | + | } | |
| 232 | + | for in, want := range map[string]string{ | |
| 233 | + | "see http://a.example/p?k=v for details": "see http://a.example/p for details", | |
| 234 | + | "ws://u@h/x and https://h2/y?q=1": "ws://h/x and https://h2/y", | |
| 235 | + | "no urls here": "no urls here", | |
| 236 | + | "odd-scheme://x:y@z/p?secret=1,next": "odd-scheme://z/p", | |
| 237 | + | "bad scheme:// and http://host?only=query end": "bad scheme:// and http://host end", | |
| 238 | + | } { | |
| 239 | + | if got := scrubURLs(in); got != want { | |
| 240 | + | t.Errorf("scrubURLs(%q) = %q, want %q", in, got, want) | |
| 241 | + | } | |
| 242 | + | } | |
| 243 | + | } | |
| 244 | + | ||
| 245 | + | func TestRedactorValues(t *testing.T) { | |
| 246 | + | t.Setenv("MCP_T_HDR", "hdr-secret-77") | |
| 247 | + | t.Setenv("MCP_T_BEARER2", " bearer-secret-99 ") | |
| 248 | + | t.Setenv("MCP_T_PASS2", "passthru-secret-5") | |
| 249 | + | red := newRedactor(ServerConfig{ | |
| 250 | + | Env: map[string]string{"A": "env-secret-1", "SHORT": "ab"}, | |
| 251 | + | Headers: map[string]string{"Authorization": "Bearer static-secret-2"}, | |
| 252 | + | EnvHeaders: map[string]string{"X": "MCP_T_HDR"}, | |
| 253 | + | BearerTokenEnvVar: "MCP_T_BEARER2", | |
| 254 | + | EnvVars: []string{"MCP_T_PASS2"}, | |
| 255 | + | URL: "https://u:urlpass-3@h.example/p?key=querysecret-4", | |
| 256 | + | }) | |
| 257 | + | in := "env-secret-1 static-secret-2 Bearer static-secret-2 hdr-secret-77 bearer-secret-99 passthru-secret-5 urlpass-3 querysecret-4 ab" | |
| 258 | + | out := red.scrub(in) | |
| 259 | + | for _, secret := range []string{"env-secret-1", "static-secret-2", "hdr-secret-77", "bearer-secret-99", "passthru-secret-5", "urlpass-3", "querysecret-4"} { | |
| 260 | + | if strings.Contains(out, secret) { | |
| 261 | + | t.Errorf("%q leaked in %q", secret, out) | |
| 262 | + | } | |
| 263 | + | } | |
| 264 | + | if !strings.HasSuffix(out, " ab") { | |
| 265 | + | t.Errorf("short values must not be redacted: %q", out) | |
| 266 | + | } | |
| 267 | + | var nilRed *redactor | |
| 268 | + | if got := nilRed.scrub("plain \x1b[1mtext"); got != "plain text" { | |
| 269 | + | t.Errorf("nil redactor = %q", got) | |
| 270 | + | } | |
| 271 | + | } | |
| 272 | + | ||
| 273 | + | func TestSanitizeText(t *testing.T) { | |
| 274 | + | cases := []struct { | |
| 275 | + | name, in string | |
| 276 | + | keepNL bool | |
| 277 | + | want string | |
| 278 | + | }{ | |
| 279 | + | {"csi", "a\x1b[31mred\x1b[0mb", true, "aredb"}, | |
| 280 | + | {"osc bel", "x\x1b]0;title\x07y", true, "xy"}, | |
| 281 | + | {"osc st", "\x1b]8;;http://e\x1b\\link\x1b]8;;\x1b\\", true, "link"}, | |
| 282 | + | {"dcs", "p\x1bPq#0;2\x1b\\q", true, "pq"}, | |
| 283 | + | {"c1 csi", "\u009b31mX", true, "X"}, | |
| 284 | + | {"cr dropped", "a\rb", true, "ab"}, | |
| 285 | + | {"tab", "a\tb", true, "a b"}, | |
| 286 | + | {"newline kept", "l1\nl2", true, "l1\nl2"}, | |
| 287 | + | {"newline folded", "l1\nl2", false, "l1 l2"}, | |
| 288 | + | {"bidi", "\u202eevil\u2066", true, "evil"}, | |
| 289 | + | {"controls", "\x00a\x07b\x7fc", true, "abc"}, | |
| 290 | + | {"lone esc", "end\x1b", true, "end"}, | |
| 291 | + | {"unicode kept", "héllo ✓ 日本", true, "héllo ✓ 日本"}, | |
| 292 | + | {"unterminated osc", "a\x1b]0;never ends", true, "a"}, | |
| 293 | + | {"esc cancels osc", "a\x1b]0;t\x1b[31mred", true, "ared"}, | |
| 294 | + | {"charset", "a\x1b(Bb", true, "ab"}, | |
| 295 | + | } | |
| 296 | + | for _, c := range cases { | |
| 297 | + | if got := sanitizeText(c.in, c.keepNL); got != c.want { | |
| 298 | + | t.Errorf("%s: sanitizeText(%q) = %q, want %q", c.name, c.in, got, c.want) | |
| 299 | + | } | |
| 300 | + | } | |
| 301 | + | } | |
| 302 | + | ||
| 303 | + | func TestFormatStatusStripsControlSequences(t *testing.T) { | |
| 304 | + | text := FormatStatus([]ServerStatus{{ | |
| 305 | + | Name: "ev\x1b]0;pwn\x07il", | |
| 306 | + | Transport: "stdio", | |
| 307 | + | Err: "boom \x1b[31mred\x1b[0m\x1b]52;c;ZXZpbA==\x07\nsecond line", | |
| 308 | + | }}) | |
| 309 | + | if strings.ContainsAny(text, "\x1b\x07") { | |
| 310 | + | t.Errorf("control characters survived: %q", text) | |
| 311 | + | } | |
| 312 | + | if !strings.Contains(text, "evil (stdio): disconnected") || !strings.Contains(text, "error: boom red") || strings.Contains(text, "second line") { | |
| 313 | + | t.Errorf("text = %q", text) | |
| 314 | + | } | |
| 315 | + | } | |
| 316 | + | ||
| 317 | + | func TestStdioDiagnosticsAreScrubbed(t *testing.T) { | |
| 318 | + | const secret = "s3cr3t-value" | |
| 319 | + | cfg, _ := fakeStdioConfig(t, map[string]string{"FAKE_INIT_ERR": "1", "FAKE_ERR_ECHO": secret, "FAKE_ANSI": "1"}) | |
| 320 | + | _, failures, _ := startManager(t, map[string]ServerConfig{"s": cfg}, "") | |
| 321 | + | if len(failures) != 1 { | |
| 322 | + | t.Fatalf("failures = %v", failures) | |
| 323 | + | } | |
| 324 | + | msg := failures[0].Error() | |
| 325 | + | if strings.Contains(msg, secret) { | |
| 326 | + | t.Errorf("secret leaked: %q", msg) | |
| 327 | + | } | |
| 328 | + | if strings.ContainsAny(msg, "\x1b\x07") || strings.Contains(msg, "evil-title") { | |
| 329 | + | t.Errorf("control sequences survived: %q", msg) | |
| 330 | + | } | |
| 331 | + | for _, want := range []string{redactedMarker, "fatal: missing credentials", "http://example.invalid/x", "red end"} { | |
| 332 | + | if !strings.Contains(msg, want) { | |
| 333 | + | t.Errorf("message missing %q: %q", want, msg) | |
| 334 | + | } | |
| 335 | + | } | |
| 336 | + | if strings.Contains(msg, "?k=") { | |
| 337 | + | t.Errorf("URL query survived: %q", msg) | |
| 338 | + | } | |
| 339 | + | } | |
| 340 | + | ||
| 341 | + | func TestHTTPDiagnosticsAreScrubbed(t *testing.T) { | |
| 342 | + | t.Run("error body", func(t *testing.T) { | |
| 343 | + | t.Setenv("MCP_T_BEARER", "tok-abcdef123") | |
| 344 | + | fake := newHTTPFake(false) | |
| 345 | + | fake.errorBody = "denied tok-abcdef123 \x1b[31mred\x1b[0m at http://x.example/p?token=tok-abcdef123" | |
| 346 | + | _, failures := startHTTP(t, fake, func(c *ServerConfig) { c.BearerTokenEnvVar = "MCP_T_BEARER" }) | |
| 347 | + | if len(failures) != 1 { | |
| 348 | + | t.Fatalf("failures = %v", failures) | |
| 349 | + | } | |
| 350 | + | msg := failures[0].Error() | |
| 351 | + | if strings.Contains(msg, "tok-abcdef123") || strings.ContainsRune(msg, 0x1b) || strings.Contains(msg, "?token") { | |
| 352 | + | t.Errorf("unsanitized: %q", msg) | |
| 353 | + | } | |
| 354 | + | if !strings.Contains(msg, redactedMarker) || !strings.Contains(msg, "500") { | |
| 355 | + | t.Errorf("message = %q", msg) | |
| 356 | + | } | |
| 357 | + | }) | |
| 358 | + | t.Run("url error", func(t *testing.T) { | |
| 359 | + | srv := httptest.NewServer(http.NotFoundHandler()) | |
| 360 | + | base := srv.URL | |
| 361 | + | srv.Close() | |
| 362 | + | cfg := ServerConfig{Type: "http", URL: base + "/mcp?access_token=QUERYSECRET"} | |
| 363 | + | _, failures, _ := startManager(t, map[string]ServerConfig{"gone": cfg}, "") | |
| 364 | + | if len(failures) != 1 { | |
| 365 | + | t.Fatalf("failures = %v", failures) | |
| 366 | + | } | |
| 367 | + | msg := failures[0].Error() | |
| 368 | + | if strings.Contains(msg, "QUERYSECRET") || strings.Contains(msg, "access_token") { | |
| 369 | + | t.Errorf("query leaked: %q", msg) | |
| 370 | + | } | |
| 371 | + | if !strings.Contains(msg, "127.0.0.1") { | |
| 372 | + | t.Errorf("host should remain for debugging: %q", msg) | |
| 373 | + | } | |
| 374 | + | }) | |
| 375 | + | } | |
| 376 | + | ||
| 377 | + | // ---- 5: process cleanup ---- | |
| 378 | + | ||
| 379 | + | func TestStdioCloseKillsDescendants(t *testing.T) { | |
| 380 | + | if !canCheckProcess { | |
| 381 | + | t.Skip("process liveness is not observable on this platform") | |
| 382 | + | } | |
| 383 | + | for _, mode := range []string{"sleeper", "sleeper-stubborn"} { | |
| 384 | + | t.Run(mode, func(t *testing.T) { | |
| 385 | + | c, _ := connectFake(t, nil) | |
| 386 | + | pid, err := strconv.Atoi(callText(t, c, "spawn", `{"mode":"`+mode+`"}`)) | |
| 387 | + | if err != nil { | |
| 388 | + | t.Fatal(err) | |
| 389 | + | } | |
| 390 | + | if !processAlive(pid) { | |
| 391 | + | t.Fatal("descendant is not running before Close") | |
| 392 | + | } | |
| 393 | + | // Make sure we never leak it even if the assertion below fails. | |
| 394 | + | t.Cleanup(func() { killPID(pid) }) | |
| 395 | + | _ = c.Close() | |
| 396 | + | eventually(t, "descendant to die", func() bool { return !processAlive(pid) }) | |
| 397 | + | }) | |
| 398 | + | } | |
| 399 | + | } | |
| 400 | + | ||
| 401 | + | // ---- 6, 10: refresh coalescing ---- | |
| 402 | + | ||
| 403 | + | func TestToolRefresherCoalescesAndSerializes(t *testing.T) { | |
| 404 | + | r := newToolRefresher() | |
| 405 | + | var running, maxRunning, runs atomic.Int32 | |
| 406 | + | release := make(chan struct{}) | |
| 407 | + | started := make(chan struct{}, 16) | |
| 408 | + | ||
| 409 | + | r.notify() // before start: remembered, not lost | |
| 410 | + | r.start(func(ctx context.Context) { | |
| 411 | + | n := running.Add(1) | |
| 412 | + | for { | |
| 413 | + | old := maxRunning.Load() | |
| 414 | + | if n <= old || maxRunning.CompareAndSwap(old, n) { | |
| 415 | + | break | |
| 416 | + | } | |
| 417 | + | } | |
| 418 | + | runs.Add(1) | |
| 419 | + | started <- struct{}{} | |
| 420 | + | select { | |
| 421 | + | case <-release: | |
| 422 | + | case <-ctx.Done(): | |
| 423 | + | } | |
| 424 | + | running.Add(-1) | |
| 425 | + | }) | |
| 426 | + | select { | |
| 427 | + | case <-started: | |
| 428 | + | case <-time.After(5 * time.Second): | |
| 429 | + | t.Fatal("notification before start never triggered a listing") | |
| 430 | + | } | |
| 431 | + | for i := 0; i < 100; i++ { | |
| 432 | + | r.notify() // all arrive during the running listing | |
| 433 | + | } | |
| 434 | + | close(release) | |
| 435 | + | eventually(t, "refresher to go idle", func() bool { | |
| 436 | + | r.mu.Lock() | |
| 437 | + | defer r.mu.Unlock() | |
| 438 | + | return !r.running | |
| 439 | + | }) | |
| 440 | + | if got := runs.Load(); got != 2 { | |
| 441 | + | t.Errorf("runs = %d, want exactly 2 (initial + one coalesced re-run)", got) | |
| 442 | + | } | |
| 443 | + | if got := maxRunning.Load(); got != 1 { | |
| 444 | + | t.Errorf("max concurrent listings = %d, want 1", got) | |
| 445 | + | } | |
| 446 | + | } | |
| 447 | + | ||
| 448 | + | func TestToolRefresherStop(t *testing.T) { | |
| 449 | + | r := newToolRefresher() | |
| 450 | + | entered := make(chan struct{}) | |
| 451 | + | finished := make(chan struct{}) | |
| 452 | + | var runs atomic.Int32 | |
| 453 | + | r.start(func(ctx context.Context) { | |
| 454 | + | runs.Add(1) | |
| 455 | + | close(entered) | |
| 456 | + | <-ctx.Done() // only a stop can end this listing | |
| 457 | + | close(finished) | |
| 458 | + | }) | |
| 459 | + | r.notify() | |
| 460 | + | <-entered | |
| 461 | + | r.notify() // dirty again while running | |
| 462 | + | r.stop() | |
| 463 | + | select { | |
| 464 | + | case <-finished: | |
| 465 | + | case <-time.After(5 * time.Second): | |
| 466 | + | t.Fatal("stop did not cancel the running listing") | |
| 467 | + | } | |
| 468 | + | r.notify() | |
| 469 | + | eventually(t, "refresher idle", func() bool { | |
| 470 | + | r.mu.Lock() | |
| 471 | + | defer r.mu.Unlock() | |
| 472 | + | return !r.running | |
| 473 | + | }) | |
| 474 | + | if got := runs.Load(); got != 1 { | |
| 475 | + | t.Errorf("runs after stop = %d, want 1", got) | |
| 476 | + | } | |
| 477 | + | } | |
| 478 | + | ||
| 479 | + | func TestManagerStopsRefresherOnClose(t *testing.T) { | |
| 480 | + | m, _ := startFake(t, nil, nil) | |
| 481 | + | s := m.servers["fake"] | |
| 482 | + | s.mu.Lock() | |
| 483 | + | r := s.refresher | |
| 484 | + | s.mu.Unlock() | |
| 485 | + | if r == nil { | |
| 486 | + | t.Fatal("no refresher installed") | |
| 487 | + | } | |
| 488 | + | _ = m.Close() | |
| 489 | + | r.mu.Lock() | |
| 490 | + | stopped := r.stopped | |
| 491 | + | r.mu.Unlock() | |
| 492 | + | if !stopped { | |
| 493 | + | t.Error("refresher still active after Close") | |
| 494 | + | } | |
| 495 | + | s.mu.Lock() | |
| 496 | + | left := s.refresher | |
| 497 | + | s.mu.Unlock() | |
| 498 | + | if left != nil { | |
| 499 | + | t.Error("refresher reference kept after Close") | |
| 500 | + | } | |
| 501 | + | } | |
| 502 | + | ||
| 503 | + | func TestManagerListChangedDuringInitialListing(t *testing.T) { | |
| 504 | + | // The server announces a changed tool set right after answering the last | |
| 505 | + | // page of the first listing, before Start has published the client. | |
| 506 | + | m, _ := startFake(t, map[string]string{"FAKE_NOTIFY_ON_LIST": "1"}, nil) | |
| 507 | + | eventually(t, "tool added during initial listing to appear", func() bool { | |
| 508 | + | for _, tool := range m.ServerTools("fake") { | |
| 509 | + | if tool.Name == "extra" { | |
| 510 | + | return true | |
| 511 | + | } | |
| 512 | + | } | |
| 513 | + | return false | |
| 514 | + | }) | |
| 515 | + | } | |
| 516 | + | ||
| 517 | + | // ---- 7: stderr draining ---- | |
| 518 | + | ||
| 519 | + | func TestStdioStderrFloodDoesNotBlockChild(t *testing.T) { | |
| 520 | + | c, _ := connectFake(t, nil) | |
| 521 | + | // 3 MiB on one stderr line, then a short one. The child blocks on the | |
| 522 | + | // stderr pipe unless the client keeps draining it. | |
| 523 | + | if got := callText(t, c, "stderr_flood", `{}`); got != "flooded" { | |
| 524 | + | t.Fatalf("got %q", got) | |
| 525 | + | } | |
| 526 | + | st := c.t.(*stdioTransport) | |
| 527 | + | eventually(t, "stderr tail to include the last line", func() bool { | |
| 528 | + | return strings.Contains(st.stderrTail(), "after-flood") | |
| 529 | + | }) | |
| 530 | + | tail := st.stderrTail() | |
| 531 | + | if !strings.Contains(tail, "[line truncated]") { | |
| 532 | + | t.Error("long line should be marked truncated") | |
| 533 | + | } | |
| 534 | + | if len(tail) > stderrTailLines*(maxStderrLine+64) { | |
| 535 | + | t.Errorf("tail not bounded: %d bytes", len(tail)) | |
| 536 | + | } | |
| 537 | + | // The server still works afterwards. | |
| 538 | + | if got := callText(t, c, "echo", `{"text":"ok"}`); got != "echo: ok" { | |
| 539 | + | t.Errorf("after flood: %q", got) | |
| 540 | + | } | |
| 541 | + | } | |
| 542 | + | ||
| 543 | + | // ---- 8, 9: cancellation notifications ---- | |
| 544 | + | ||
| 545 | + | func TestStdioCancelNotificationRules(t *testing.T) { | |
| 546 | + | // initialize never answers within the deadline. | |
| 547 | + | cfg, _ := fakeStdioConfig(t, map[string]string{"FAKE_INIT_DELAY_MS": "5000"}) | |
| 548 | + | st, err := startStdio("x", cfg, "") | |
| 549 | + | if err != nil { | |
| 550 | + | t.Fatal(err) | |
| 551 | + | } | |
| 552 | + | defer st.close() | |
| 553 | + | ctx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond) | |
| 554 | + | defer cancel() | |
| 555 | + | if _, err := st.request(ctx, 1, "initialize", map[string]any{"protocolVersion": LatestProtocolVersion}); !errors.Is(err, context.DeadlineExceeded) { | |
| 556 | + | t.Fatalf("err = %v", err) | |
| 557 | + | } | |
| 558 | + | if got := st.cancelsSent.Load(); got != 0 { | |
| 559 | + | t.Errorf("notifications/cancelled sent for initialize (%d)", got) | |
| 560 | + | } | |
| 561 | + | // Any other method is announced exactly once, without blocking the caller. | |
| 562 | + | ctx2, cancel2 := context.WithCancel(context.Background()) | |
| 563 | + | cancel2() | |
| 564 | + | start := time.Now() | |
| 565 | + | _, err = st.request(ctx2, 2, "tools/call", map[string]any{"name": "slow", "arguments": map[string]any{}}) | |
| 566 | + | if !errors.Is(err, context.Canceled) { | |
| 567 | + | t.Fatalf("err = %v", err) | |
| 568 | + | } | |
| 569 | + | if time.Since(start) > time.Second { | |
| 570 | + | t.Error("caller delayed by the cancel notification") | |
| 571 | + | } | |
| 572 | + | if got := st.cancelsSent.Load(); got != 1 { | |
| 573 | + | t.Errorf("cancel count for tools/call = %d, want 1", got) | |
| 574 | + | } | |
| 575 | + | } | |
| 576 | + | ||
| 577 | + | func TestHTTPInitializeNotCancelledOnTimeout(t *testing.T) { | |
| 578 | + | fake := newHTTPFake(false) | |
| 579 | + | fake.hangInit = true | |
| 580 | + | srv := fake.start(t) | |
| 581 | + | cfg := httpServerConfig(srv) | |
| 582 | + | ctx, cancel := context.WithTimeout(context.Background(), 200*time.Millisecond) | |
| 583 | + | defer cancel() | |
| 584 | + | _, err := Connect(ctx, "web", cfg, "", ClientInfo{}, nil) | |
| 585 | + | if !errors.Is(err, context.DeadlineExceeded) { | |
| 586 | + | t.Fatalf("err = %v, want deadline exceeded", err) | |
| 587 | + | } | |
| 588 | + | for _, r := range fake.snapshot() { | |
| 589 | + | if r.RPCMethod == "notifications/cancelled" { | |
| 590 | + | t.Fatal("notifications/cancelled sent for initialize") | |
| 591 | + | } | |
| 592 | + | } | |
| 593 | + | } | |
| 594 | + | ||
| 595 | + | func TestHTTPCancelNotificationIsBoundedAndOwnedByTransport(t *testing.T) { | |
| 596 | + | fake := newHTTPFake(false) | |
| 597 | + | fake.cancelGate = make(chan struct{}) | |
| 598 | + | t.Cleanup(func() { close(fake.cancelGate) }) | |
| 599 | + | m, failures := startHTTP(t, fake, func(c *ServerConfig) { c.ToolTimeoutSec = 0.2 }) | |
| 600 | + | if len(failures) != 0 { | |
| 601 | + | t.Fatal(failures) | |
| 602 | + | } | |
| 603 | + | start := time.Now() | |
| 604 | + | res := execTool(t, findTool(t, m, "mcp__web__slow"), `{}`) | |
| 605 | + | if !res.IsError || !strings.Contains(resultText(res), "toolTimeoutSec") { | |
| 606 | + | t.Fatalf("got %q", resultText(res)) | |
| 607 | + | } | |
| 608 | + | // The cancellation POST is stuck on the server (gate closed); the caller | |
| 609 | + | // must not wait for it. | |
| 610 | + | if d := time.Since(start); d > 1500*time.Millisecond { | |
| 611 | + | t.Fatalf("caller delayed by the cancel notification: %v", d) | |
| 612 | + | } | |
| 613 | + | eventually(t, "cancel notification to reach the server", func() bool { | |
| 614 | + | for _, r := range fake.snapshot() { | |
| 615 | + | if r.RPCMethod == "notifications/cancelled" { | |
| 616 | + | return true | |
| 617 | + | } | |
| 618 | + | } | |
| 619 | + | return false | |
| 620 | + | }) | |
| 621 | + | // Closing the transport aborts the stuck notification; otherwise the | |
| 622 | + | // httptest server would wait for it at cleanup. | |
| 623 | + | closed := make(chan struct{}) | |
| 624 | + | go func() { _ = m.Close(); close(closed) }() | |
| 625 | + | select { | |
| 626 | + | case <-closed: | |
| 627 | + | case <-time.After(10 * time.Second): | |
| 628 | + | t.Fatal("Close hung behind the cancel notification") | |
| 629 | + | } | |
| 630 | + | } | |
| 631 | + | ||
| 632 | + | // ---- 11: file URIs ---- | |
| 633 | + | ||
| 634 | + | func TestFileURI(t *testing.T) { | |
| 635 | + | cases := []struct { | |
| 636 | + | name, path string | |
| 637 | + | windows bool | |
| 638 | + | want string | |
| 639 | + | }{ | |
| 640 | + | {"unix", "/home/me/proj", false, "file:///home/me/proj"}, | |
| 641 | + | {"unix space", "/home/me/my proj", false, "file:///home/me/my%20proj"}, | |
| 642 | + | {"drive", `C:\Users\me\proj`, true, "file:///C:/Users/me/proj"}, | |
| 643 | + | {"drive space", `D:\My Docs\x`, true, "file:///D:/My%20Docs/x"}, | |
| 644 | + | {"drive root", `C:\`, true, "file:///C:/"}, | |
| 645 | + | {"drive forward slashes", `C:/x/y`, true, "file:///C:/x/y"}, | |
| 646 | + | {"unc", `\\host\share\dir`, true, "file://host/share/dir"}, | |
| 647 | + | {"unc share only", `\\host\share`, true, "file://host/share"}, | |
| 648 | + | {"extended drive", `\\?\C:\x`, true, "file:///C:/x"}, | |
| 649 | + | {"extended unc", `\\?\UNC\host\share\d`, true, "file://host/share/d"}, | |
| 650 | + | } | |
| 651 | + | for _, c := range cases { | |
| 652 | + | if got := fileURI(c.path, c.windows); got != c.want { | |
| 653 | + | t.Errorf("%s: fileURI(%q) = %q, want %q", c.name, c.path, got, c.want) | |
| 654 | + | } | |
| 655 | + | } | |
| 656 | + | } | |
| 657 | + | ||
| 658 | + | // ---- 12: DELETE shape ---- | |
| 659 | + | ||
| 660 | + | func TestHTTPDeleteCarriesProtocolVersionAndHeaders(t *testing.T) { | |
| 661 | + | t.Setenv("MCP_T_BEARER", "tok-delete-1") | |
| 662 | + | fake := newHTTPFake(false) | |
| 663 | + | m, failures := startHTTP(t, fake, func(c *ServerConfig) { | |
| 664 | + | c.BearerTokenEnvVar = "MCP_T_BEARER" | |
| 665 | + | c.Headers = map[string]string{"X-Static": "static"} | |
| 666 | + | }) | |
| 667 | + | if len(failures) != 0 { | |
| 668 | + | t.Fatal(failures) | |
| 669 | + | } | |
| 670 | + | _ = m.Close() | |
| 671 | + | select { | |
| 672 | + | case <-fake.deleted: | |
| 673 | + | case <-time.After(5 * time.Second): | |
| 674 | + | t.Fatal("no DELETE") | |
| 675 | + | } | |
| 676 | + | var del *httpRecord | |
| 677 | + | for _, r := range fake.snapshot() { | |
| 678 | + | if r.Method == http.MethodDelete { | |
| 679 | + | r := r | |
| 680 | + | del = &r | |
| 681 | + | } | |
| 682 | + | } | |
| 683 | + | if del == nil { | |
| 684 | + | t.Fatal("DELETE not recorded") | |
| 685 | + | } | |
| 686 | + | if del.Protocol != "2025-06-18" || del.Session != "sess-abc" || del.Auth != "Bearer tok-delete-1" || del.Static != "static" { | |
| 687 | + | t.Errorf("DELETE = %+v", *del) | |
| 688 | + | } | |
| 689 | + | } | |
| 690 | + | ||
| 691 | + | // ---- 14: resource listing budget ---- | |
| 692 | + | ||
| 693 | + | func TestListResourcesAppliesBudgetWhileFormatting(t *testing.T) { | |
| 694 | + | m, events := startFake(t, map[string]string{"FAKE_MANY_RESOURCES": "2000"}, nil) | |
| 695 | + | text := resultText(execTool(t, findTool(t, m, ListResourcesToolName), `{}`)) | |
| 696 | + | marker := "\n... [truncated at 51200 bytes]" | |
| 697 | + | if !strings.Contains(text, marker) { | |
| 698 | + | t.Fatalf("no truncation marker (len=%d)", len(text)) | |
| 699 | + | } | |
| 700 | + | if len(text) > 50*1024+len(marker) { | |
| 701 | + | t.Errorf("output not bounded: %d bytes", len(text)) | |
| 702 | + | } | |
| 703 | + | if !strings.Contains(text, "mem://item/0") || strings.Contains(text, "mem://item/1999") { | |
| 704 | + | t.Error("expected the head of the list, not the tail") | |
| 705 | + | } | |
| 706 | + | // Once the budget is spent no further requests are made. | |
| 707 | + | if n := countEvents(events, "list-templates"); n != 0 { | |
| 708 | + | t.Errorf("templates were fetched after the budget was exhausted (%d)", n) | |
| 709 | + | } | |
| 710 | + | if n := countEvents(events, "list-resources"); n != 1 { | |
| 711 | + | t.Errorf("list-resources calls = %d", n) | |
| 712 | + | } | |
| 713 | + | } | |
| 714 | + | ||
| 715 | + | func TestBudgetedText(t *testing.T) { | |
| 716 | + | b := &budgetedText{} | |
| 717 | + | b.write(strings.Repeat("a", maxResultTextBytes-1)) | |
| 718 | + | if b.full() { | |
| 719 | + | t.Fatal("full too early") | |
| 720 | + | } | |
| 721 | + | b.write("bcdef") | |
| 722 | + | if !b.full() { | |
| 723 | + | t.Fatal("not full after overflow") | |
| 724 | + | } | |
| 725 | + | before := b.String() | |
| 726 | + | b.write("more") | |
| 727 | + | b.printf("%s", "more") | |
| 728 | + | if b.String() != before { | |
| 729 | + | t.Error("writes accepted after the budget was exhausted") | |
| 730 | + | } | |
| 731 | + | if !strings.HasSuffix(before, "[truncated at 51200 bytes]") { | |
| 732 | + | t.Errorf("tail = %q", before[len(before)-40:]) | |
| 733 | + | } | |
| 734 | + | } | |
| @@ -0,0 +1,734 @@ | |||||
| 1 | + | package mcp | |||
| 2 | + | ||||
| 3 | + | import ( | |||
| 4 | + | "context" | |||
| 5 | + | "errors" | |||
| 6 | + | "net/http" | |||
| 7 | + | "net/http/httptest" | |||
| 8 | + | "net/url" | |||
| 9 | + | "strconv" | |||
| 10 | + | "strings" | |||
| 11 | + | "sync/atomic" | |||
| 12 | + | "testing" | |||
| 13 | + | "time" | |||
| 14 | + | ) | |||
| 15 | + | ||||
| 16 | + | // Regression tests for the review findings: write path, SSE limits, session | |||
| 17 | + | // expiry, diagnostics hygiene, process cleanup, refresh coalescing, and the | |||
| 18 | + | // cancellation rules. | |||
| 19 | + | ||||
| 20 | + | func connectFake(t *testing.T, extraEnv map[string]string) (*Client, string) { | |||
| 21 | + | t.Helper() | |||
| 22 | + | cfg, events := fakeStdioConfig(t, extraEnv) | |||
| 23 | + | ctx, cancel := context.WithTimeout(context.Background(), 20*time.Second) | |||
| 24 | + | defer cancel() | |||
| 25 | + | c, err := Connect(ctx, "fake", cfg, "", ClientInfo{Name: "t"}, nil) | |||
| 26 | + | if err != nil { | |||
| 27 | + | t.Fatalf("Connect: %v", err) | |||
| 28 | + | } | |||
| 29 | + | t.Cleanup(func() { _ = c.Close() }) | |||
| 30 | + | return c, events | |||
| 31 | + | } | |||
| 32 | + | ||||
| 33 | + | func callText(t *testing.T, c *Client, tool, args string) string { | |||
| 34 | + | t.Helper() | |||
| 35 | + | ctx, cancel := context.WithTimeout(context.Background(), 20*time.Second) | |||
| 36 | + | defer cancel() | |||
| 37 | + | res, err := c.CallTool(ctx, tool, []byte(args)) | |||
| 38 | + | if err != nil { | |||
| 39 | + | t.Fatalf("CallTool %s: %v", tool, err) | |||
| 40 | + | } | |||
| 41 | + | if len(res.Content) == 0 { | |||
| 42 | + | t.Fatalf("CallTool %s: no content", tool) | |||
| 43 | + | } | |||
| 44 | + | return res.Content[0].Text | |||
| 45 | + | } | |||
| 46 | + | ||||
| 47 | + | // ---- 1, 8: cancellable writes, close behind a stuck writer ---- | |||
| 48 | + | ||||
| 49 | + | func TestStdioWriteStuckChildDoesNotWedgeCancelOrClose(t *testing.T) { | |||
| 50 | + | c, events := connectFake(t, nil) | |||
| 51 | + | ||||
| 52 | + | // The server stops reading stdin after receiving this call. | |||
| 53 | + | stallCtx, stallCancel := context.WithCancel(context.Background()) | |||
| 54 | + | stallDone := make(chan error, 1) | |||
| 55 | + | go func() { | |||
| 56 | + | _, err := c.CallTool(stallCtx, "stall", []byte(`{}`)) | |||
| 57 | + | stallDone <- err | |||
| 58 | + | }() | |||
| 59 | + | waitEvent(t, events, "stalled") | |||
| 60 | + | ||||
| 61 | + | // A frame far larger than the pipe buffer now blocks the writer. | |||
| 62 | + | big := `{"text":"` + strings.Repeat("a", 4<<20) + `"}` | |||
| 63 | + | ctx, cancel := context.WithTimeout(context.Background(), 300*time.Millisecond) | |||
| 64 | + | defer cancel() | |||
| 65 | + | start := time.Now() | |||
| 66 | + | _, err := c.CallTool(ctx, "echo", []byte(big)) | |||
| 67 | + | if !errors.Is(err, context.DeadlineExceeded) { | |||
| 68 | + | t.Fatalf("err = %v, want deadline exceeded", err) | |||
| 69 | + | } | |||
| 70 | + | if d := time.Since(start); d > 3*time.Second { | |||
| 71 | + | t.Fatalf("request ignored its context for %v", d) | |||
| 72 | + | } | |||
| 73 | + | ||||
| 74 | + | closed := make(chan struct{}) | |||
| 75 | + | go func() { _ = c.Close(); close(closed) }() | |||
| 76 | + | select { | |||
| 77 | + | case <-closed: | |||
| 78 | + | case <-time.After(10 * time.Second): | |||
| 79 | + | t.Fatal("Close hung behind a blocked write") | |||
| 80 | + | } | |||
| 81 | + | select { | |||
| 82 | + | case <-c.Done(): | |||
| 83 | + | default: | |||
| 84 | + | t.Error("Done not closed after Close") | |||
| 85 | + | } | |||
| 86 | + | stallCancel() | |||
| 87 | + | select { | |||
| 88 | + | case <-stallDone: | |||
| 89 | + | case <-time.After(5 * time.Second): | |||
| 90 | + | t.Fatal("pending call not released by Close") | |||
| 91 | + | } | |||
| 92 | + | } | |||
| 93 | + | ||||
| 94 | + | // ---- 2: SSE aggregate limit ---- | |||
| 95 | + | ||||
| 96 | + | func TestReadSSEFramesAggregateLimit(t *testing.T) { | |||
| 97 | + | line := "data: " + strings.Repeat("x", 9) + "\n" // 9 bytes of data per line | |||
| 98 | + | var delivered int | |||
| 99 | + | handle := func(message) bool { delivered++; return false } | |||
| 100 | + | ||||
| 101 | + | // Many short lines of one event exceed the limit. | |||
| 102 | + | err := readSSEFramesLimit(strings.NewReader(strings.Repeat(line, 200)+"\n"), 1000, handle) | |||
| 103 | + | var tooLarge errSSEEventTooLarge | |||
| 104 | + | if !errors.As(err, &tooLarge) || !strings.Contains(err.Error(), "1000") { | |||
| 105 | + | t.Fatalf("err = %v, want errSSEEventTooLarge", err) | |||
| 106 | + | } | |||
| 107 | + | if delivered != 0 { | |||
| 108 | + | t.Errorf("oversized event delivered %d frames", delivered) | |||
| 109 | + | } | |||
| 110 | + | ||||
| 111 | + | // The limit is per event: a blank line resets it. | |||
| 112 | + | ok := `data: {"jsonrpc":"2.0","method":"a"}` + "\n\n" | |||
| 113 | + | delivered = 0 | |||
| 114 | + | if err := readSSEFramesLimit(strings.NewReader(strings.Repeat(ok, 50)), 100, handle); err != nil { | |||
| 115 | + | t.Fatal(err) | |||
| 116 | + | } | |||
| 117 | + | if delivered != 50 { | |||
| 118 | + | t.Errorf("delivered %d, want 50", delivered) | |||
| 119 | + | } | |||
| 120 | + | ||||
| 121 | + | // A single line longer than the limit is the same explicit error. | |||
| 122 | + | err = readSSEFramesLimit(strings.NewReader("data: "+strings.Repeat("y", 5000)+"\n\n"), 1000, handle) | |||
| 123 | + | if !errors.As(err, &tooLarge) { | |||
| 124 | + | t.Errorf("long line err = %v", err) | |||
| 125 | + | } | |||
| 126 | + | } | |||
| 127 | + | ||||
| 128 | + | func TestSSEEndlessEventStopsAtLimit(t *testing.T) { | |||
| 129 | + | // A server that never ends its event must not grow memory without bound. | |||
| 130 | + | err := readSSEFramesLimit(&infiniteSSE{}, 64<<10, func(message) bool { return false }) | |||
| 131 | + | var tooLarge errSSEEventTooLarge | |||
| 132 | + | if !errors.As(err, &tooLarge) { | |||
| 133 | + | t.Fatalf("err = %v", err) | |||
| 134 | + | } | |||
| 135 | + | } | |||
| 136 | + | ||||
| 137 | + | type infiniteSSE struct{} | |||
| 138 | + | ||||
| 139 | + | func (*infiniteSSE) Read(p []byte) (int, error) { | |||
| 140 | + | line := []byte("data: " + strings.Repeat("q", 100) + "\n") | |||
| 141 | + | n := copy(p, line) | |||
| 142 | + | return n, nil | |||
| 143 | + | } | |||
| 144 | + | ||||
| 145 | + | // ---- 3: session expiry ---- | |||
| 146 | + | ||||
| 147 | + | func TestHTTPSessionExpiredReconnects(t *testing.T) { | |||
| 148 | + | for _, mode := range []struct { | |||
| 149 | + | name string | |||
| 150 | + | sse bool | |||
| 151 | + | }{{"json", false}, {"sse", true}} { | |||
| 152 | + | t.Run(mode.name, func(t *testing.T) { | |||
| 153 | + | fake := newHTTPFake(mode.sse) | |||
| 154 | + | m, failures := startHTTP(t, fake, nil) | |||
| 155 | + | if len(failures) != 0 { | |||
| 156 | + | t.Fatal(failures) | |||
| 157 | + | } | |||
| 158 | + | echo := findTool(t, m, "mcp__web__echo") | |||
| 159 | + | fake.expireNextCall.Store(true) | |||
| 160 | + | res := execTool(t, echo, `{"text":"one"}`) | |||
| 161 | + | if !res.IsError || !strings.Contains(resultText(res), "session expired") { | |||
| 162 | + | t.Fatalf("expired call = %q (error=%v)", resultText(res), res.IsError) | |||
| 163 | + | } | |||
| 164 | + | // The transport is unusable, so the status reflects it at once. | |||
| 165 | + | if st := m.Status()[0]; st.Connected { | |||
| 166 | + | t.Errorf("still marked connected after 404: %+v", st) | |||
| 167 | + | } | |||
| 168 | + | // The next call reconnects (fresh initialize) and succeeds. | |||
| 169 | + | if got := resultText(execTool(t, echo, `{"text":"two"}`)); got != "echo: two" { | |||
| 170 | + | t.Fatalf("after reconnect = %q", got) | |||
| 171 | + | } | |||
| 172 | + | inits, calls := fake.counts() | |||
| 173 | + | if inits != 2 { | |||
| 174 | + | t.Errorf("initialize count = %d, want 2", inits) | |||
| 175 | + | } | |||
| 176 | + | if calls != 2 { | |||
| 177 | + | t.Errorf("tools/call count = %d, want 2 (the failed call must not be replayed)", calls) | |||
| 178 | + | } | |||
| 179 | + | // No DELETE for the dead session. | |||
| 180 | + | select { | |||
| 181 | + | case sid := <-fake.deleted: | |||
| 182 | + | t.Errorf("DELETE sent for expired session %q", sid) | |||
| 183 | + | default: | |||
| 184 | + | } | |||
| 185 | + | // The new session carries its own id. | |||
| 186 | + | var last httpRecord | |||
| 187 | + | for _, r := range fake.snapshot() { | |||
| 188 | + | if r.RPCMethod == "tools/call" { | |||
| 189 | + | last = r | |||
| 190 | + | } | |||
| 191 | + | } | |||
| 192 | + | if last.Session != "sess-abc-2" { | |||
| 193 | + | t.Errorf("session after reconnect = %q", last.Session) | |||
| 194 | + | } | |||
| 195 | + | }) | |||
| 196 | + | } | |||
| 197 | + | } | |||
| 198 | + | ||||
| 199 | + | func TestHTTPNotFoundWithoutSessionIsOrdinaryError(t *testing.T) { | |||
| 200 | + | srv := httptest.NewServer(http.NotFoundHandler()) | |||
| 201 | + | defer srv.Close() | |||
| 202 | + | tr, err := newHTTPTransport("x", ServerConfig{URL: srv.URL}, nil) | |||
| 203 | + | if err != nil { | |||
| 204 | + | t.Fatal(err) | |||
| 205 | + | } | |||
| 206 | + | defer tr.close() | |||
| 207 | + | _, err = tr.request(context.Background(), 1, "initialize", nil) | |||
| 208 | + | var he *HTTPStatusError | |||
| 209 | + | if !errors.As(err, &he) || he.Status != 404 { | |||
| 210 | + | t.Fatalf("err = %v", err) | |||
| 211 | + | } | |||
| 212 | + | select { | |||
| 213 | + | case <-tr.done(): | |||
| 214 | + | t.Error("a 404 before any session must not poison the transport") | |||
| 215 | + | default: | |||
| 216 | + | } | |||
| 217 | + | } | |||
| 218 | + | ||||
| 219 | + | // ---- 4, 13: diagnostics hygiene ---- | |||
| 220 | + | ||||
| 221 | + | func TestScrubURLsAndErrors(t *testing.T) { | |||
| 222 | + | red := newRedactor(ServerConfig{}) | |||
| 223 | + | uerr := &url.Error{Op: "Post", URL: "https://user:pw@h.example:8443/mcp/v1?token=abc&x=1#frag", Err: errors.New("connection refused")} | |||
| 224 | + | got := red.scrubError(uerr).Error() | |||
| 225 | + | if want := `Post "https://h.example:8443/mcp/v1": connection refused`; got != want { | |||
| 226 | + | t.Errorf("got %q, want %q", got, want) | |||
| 227 | + | } | |||
| 228 | + | var target *url.Error | |||
| 229 | + | if !errors.As(red.scrubError(uerr), &target) { | |||
| 230 | + | t.Error("original error must stay reachable via errors.As") | |||
| 231 | + | } | |||
| 232 | + | for in, want := range map[string]string{ | |||
| 233 | + | "see http://a.example/p?k=v for details": "see http://a.example/p for details", | |||
| 234 | + | "ws://u@h/x and https://h2/y?q=1": "ws://h/x and https://h2/y", | |||
| 235 | + | "no urls here": "no urls here", | |||
| 236 | + | "odd-scheme://x:y@z/p?secret=1,next": "odd-scheme://z/p", | |||
| 237 | + | "bad scheme:// and http://host?only=query end": "bad scheme:// and http://host end", | |||
| 238 | + | } { | |||
| 239 | + | if got := scrubURLs(in); got != want { | |||
| 240 | + | t.Errorf("scrubURLs(%q) = %q, want %q", in, got, want) | |||
| 241 | + | } | |||
| 242 | + | } | |||
| 243 | + | } | |||
| 244 | + | ||||
| 245 | + | func TestRedactorValues(t *testing.T) { | |||
| 246 | + | t.Setenv("MCP_T_HDR", "hdr-secret-77") | |||
| 247 | + | t.Setenv("MCP_T_BEARER2", " bearer-secret-99 ") | |||
| 248 | + | t.Setenv("MCP_T_PASS2", "passthru-secret-5") | |||
| 249 | + | red := newRedactor(ServerConfig{ | |||
| 250 | + | Env: map[string]string{"A": "env-secret-1", "SHORT": "ab"}, | |||
| 251 | + | Headers: map[string]string{"Authorization": "Bearer static-secret-2"}, | |||
| 252 | + | EnvHeaders: map[string]string{"X": "MCP_T_HDR"}, | |||
| 253 | + | BearerTokenEnvVar: "MCP_T_BEARER2", | |||
| 254 | + | EnvVars: []string{"MCP_T_PASS2"}, | |||
| 255 | + | URL: "https://u:urlpass-3@h.example/p?key=querysecret-4", | |||
| 256 | + | }) | |||
| 257 | + | in := "env-secret-1 static-secret-2 Bearer static-secret-2 hdr-secret-77 bearer-secret-99 passthru-secret-5 urlpass-3 querysecret-4 ab" | |||
| 258 | + | out := red.scrub(in) | |||
| 259 | + | for _, secret := range []string{"env-secret-1", "static-secret-2", "hdr-secret-77", "bearer-secret-99", "passthru-secret-5", "urlpass-3", "querysecret-4"} { | |||
| 260 | + | if strings.Contains(out, secret) { | |||
| 261 | + | t.Errorf("%q leaked in %q", secret, out) | |||
| 262 | + | } | |||
| 263 | + | } | |||
| 264 | + | if !strings.HasSuffix(out, " ab") { | |||
| 265 | + | t.Errorf("short values must not be redacted: %q", out) | |||
| 266 | + | } | |||
| 267 | + | var nilRed *redactor | |||
| 268 | + | if got := nilRed.scrub("plain \x1b[1mtext"); got != "plain text" { | |||
| 269 | + | t.Errorf("nil redactor = %q", got) | |||
| 270 | + | } | |||
| 271 | + | } | |||
| 272 | + | ||||
| 273 | + | func TestSanitizeText(t *testing.T) { | |||
| 274 | + | cases := []struct { | |||
| 275 | + | name, in string | |||
| 276 | + | keepNL bool | |||
| 277 | + | want string | |||
| 278 | + | }{ | |||
| 279 | + | {"csi", "a\x1b[31mred\x1b[0mb", true, "aredb"}, | |||
| 280 | + | {"osc bel", "x\x1b]0;title\x07y", true, "xy"}, | |||
| 281 | + | {"osc st", "\x1b]8;;http://e\x1b\\link\x1b]8;;\x1b\\", true, "link"}, | |||
| 282 | + | {"dcs", "p\x1bPq#0;2\x1b\\q", true, "pq"}, | |||
| 283 | + | {"c1 csi", "\u009b31mX", true, "X"}, | |||
| 284 | + | {"cr dropped", "a\rb", true, "ab"}, | |||
| 285 | + | {"tab", "a\tb", true, "a b"}, | |||
| 286 | + | {"newline kept", "l1\nl2", true, "l1\nl2"}, | |||
| 287 | + | {"newline folded", "l1\nl2", false, "l1 l2"}, | |||
| 288 | + | {"bidi", "\u202eevil\u2066", true, "evil"}, | |||
| 289 | + | {"controls", "\x00a\x07b\x7fc", true, "abc"}, | |||
| 290 | + | {"lone esc", "end\x1b", true, "end"}, | |||
| 291 | + | {"unicode kept", "héllo ✓ 日本", true, "héllo ✓ 日本"}, | |||
| 292 | + | {"unterminated osc", "a\x1b]0;never ends", true, "a"}, | |||
| 293 | + | {"esc cancels osc", "a\x1b]0;t\x1b[31mred", true, "ared"}, | |||
| 294 | + | {"charset", "a\x1b(Bb", true, "ab"}, | |||
| 295 | + | } | |||
| 296 | + | for _, c := range cases { | |||
| 297 | + | if got := sanitizeText(c.in, c.keepNL); got != c.want { | |||
| 298 | + | t.Errorf("%s: sanitizeText(%q) = %q, want %q", c.name, c.in, got, c.want) | |||
| 299 | + | } | |||
| 300 | + | } | |||
| 301 | + | } | |||
| 302 | + | ||||
| 303 | + | func TestFormatStatusStripsControlSequences(t *testing.T) { | |||
| 304 | + | text := FormatStatus([]ServerStatus{{ | |||
| 305 | + | Name: "ev\x1b]0;pwn\x07il", | |||
| 306 | + | Transport: "stdio", | |||
| 307 | + | Err: "boom \x1b[31mred\x1b[0m\x1b]52;c;ZXZpbA==\x07\nsecond line", | |||
| 308 | + | }}) | |||
| 309 | + | if strings.ContainsAny(text, "\x1b\x07") { | |||
| 310 | + | t.Errorf("control characters survived: %q", text) | |||
| 311 | + | } | |||
| 312 | + | if !strings.Contains(text, "evil (stdio): disconnected") || !strings.Contains(text, "error: boom red") || strings.Contains(text, "second line") { | |||
| 313 | + | t.Errorf("text = %q", text) | |||
| 314 | + | } | |||
| 315 | + | } | |||
| 316 | + | ||||
| 317 | + | func TestStdioDiagnosticsAreScrubbed(t *testing.T) { | |||
| 318 | + | const secret = "s3cr3t-value" | |||
| 319 | + | cfg, _ := fakeStdioConfig(t, map[string]string{"FAKE_INIT_ERR": "1", "FAKE_ERR_ECHO": secret, "FAKE_ANSI": "1"}) | |||
| 320 | + | _, failures, _ := startManager(t, map[string]ServerConfig{"s": cfg}, "") | |||
| 321 | + | if len(failures) != 1 { | |||
| 322 | + | t.Fatalf("failures = %v", failures) | |||
| 323 | + | } | |||
| 324 | + | msg := failures[0].Error() | |||
| 325 | + | if strings.Contains(msg, secret) { | |||
| 326 | + | t.Errorf("secret leaked: %q", msg) | |||
| 327 | + | } | |||
| 328 | + | if strings.ContainsAny(msg, "\x1b\x07") || strings.Contains(msg, "evil-title") { | |||
| 329 | + | t.Errorf("control sequences survived: %q", msg) | |||
| 330 | + | } | |||
| 331 | + | for _, want := range []string{redactedMarker, "fatal: missing credentials", "http://example.invalid/x", "red end"} { | |||
| 332 | + | if !strings.Contains(msg, want) { | |||
| 333 | + | t.Errorf("message missing %q: %q", want, msg) | |||
| 334 | + | } | |||
| 335 | + | } | |||
| 336 | + | if strings.Contains(msg, "?k=") { | |||
| 337 | + | t.Errorf("URL query survived: %q", msg) | |||
| 338 | + | } | |||
| 339 | + | } | |||
| 340 | + | ||||
| 341 | + | func TestHTTPDiagnosticsAreScrubbed(t *testing.T) { | |||
| 342 | + | t.Run("error body", func(t *testing.T) { | |||
| 343 | + | t.Setenv("MCP_T_BEARER", "tok-abcdef123") | |||
| 344 | + | fake := newHTTPFake(false) | |||
| 345 | + | fake.errorBody = "denied tok-abcdef123 \x1b[31mred\x1b[0m at http://x.example/p?token=tok-abcdef123" | |||
| 346 | + | _, failures := startHTTP(t, fake, func(c *ServerConfig) { c.BearerTokenEnvVar = "MCP_T_BEARER" }) | |||
| 347 | + | if len(failures) != 1 { | |||
| 348 | + | t.Fatalf("failures = %v", failures) | |||
| 349 | + | } | |||
| 350 | + | msg := failures[0].Error() | |||
| 351 | + | if strings.Contains(msg, "tok-abcdef123") || strings.ContainsRune(msg, 0x1b) || strings.Contains(msg, "?token") { | |||
| 352 | + | t.Errorf("unsanitized: %q", msg) | |||
| 353 | + | } | |||
| 354 | + | if !strings.Contains(msg, redactedMarker) || !strings.Contains(msg, "500") { | |||
| 355 | + | t.Errorf("message = %q", msg) | |||
| 356 | + | } | |||
| 357 | + | }) | |||
| 358 | + | t.Run("url error", func(t *testing.T) { | |||
| 359 | + | srv := httptest.NewServer(http.NotFoundHandler()) | |||
| 360 | + | base := srv.URL | |||
| 361 | + | srv.Close() | |||
| 362 | + | cfg := ServerConfig{Type: "http", URL: base + "/mcp?access_token=QUERYSECRET"} | |||
| 363 | + | _, failures, _ := startManager(t, map[string]ServerConfig{"gone": cfg}, "") | |||
| 364 | + | if len(failures) != 1 { | |||
| 365 | + | t.Fatalf("failures = %v", failures) | |||
| 366 | + | } | |||
| 367 | + | msg := failures[0].Error() | |||
| 368 | + | if strings.Contains(msg, "QUERYSECRET") || strings.Contains(msg, "access_token") { | |||
| 369 | + | t.Errorf("query leaked: %q", msg) | |||
| 370 | + | } | |||
| 371 | + | if !strings.Contains(msg, "127.0.0.1") { | |||
| 372 | + | t.Errorf("host should remain for debugging: %q", msg) | |||
| 373 | + | } | |||
| 374 | + | }) | |||
| 375 | + | } | |||
| 376 | + | ||||
| 377 | + | // ---- 5: process cleanup ---- | |||
| 378 | + | ||||
| 379 | + | func TestStdioCloseKillsDescendants(t *testing.T) { | |||
| 380 | + | if !canCheckProcess { | |||
| 381 | + | t.Skip("process liveness is not observable on this platform") | |||
| 382 | + | } | |||
| 383 | + | for _, mode := range []string{"sleeper", "sleeper-stubborn"} { | |||
| 384 | + | t.Run(mode, func(t *testing.T) { | |||
| 385 | + | c, _ := connectFake(t, nil) | |||
| 386 | + | pid, err := strconv.Atoi(callText(t, c, "spawn", `{"mode":"`+mode+`"}`)) | |||
| 387 | + | if err != nil { | |||
| 388 | + | t.Fatal(err) | |||
| 389 | + | } | |||
| 390 | + | if !processAlive(pid) { | |||
| 391 | + | t.Fatal("descendant is not running before Close") | |||
| 392 | + | } | |||
| 393 | + | // Make sure we never leak it even if the assertion below fails. | |||
| 394 | + | t.Cleanup(func() { killPID(pid) }) | |||
| 395 | + | _ = c.Close() | |||
| 396 | + | eventually(t, "descendant to die", func() bool { return !processAlive(pid) }) | |||
| 397 | + | }) | |||
| 398 | + | } | |||
| 399 | + | } | |||
| 400 | + | ||||
| 401 | + | // ---- 6, 10: refresh coalescing ---- | |||
| 402 | + | ||||
| 403 | + | func TestToolRefresherCoalescesAndSerializes(t *testing.T) { | |||
| 404 | + | r := newToolRefresher() | |||
| 405 | + | var running, maxRunning, runs atomic.Int32 | |||
| 406 | + | release := make(chan struct{}) | |||
| 407 | + | started := make(chan struct{}, 16) | |||
| 408 | + | ||||
| 409 | + | r.notify() // before start: remembered, not lost | |||
| 410 | + | r.start(func(ctx context.Context) { | |||
| 411 | + | n := running.Add(1) | |||
| 412 | + | for { | |||
| 413 | + | old := maxRunning.Load() | |||
| 414 | + | if n <= old || maxRunning.CompareAndSwap(old, n) { | |||
| 415 | + | break | |||
| 416 | + | } | |||
| 417 | + | } | |||
| 418 | + | runs.Add(1) | |||
| 419 | + | started <- struct{}{} | |||
| 420 | + | select { | |||
| 421 | + | case <-release: | |||
| 422 | + | case <-ctx.Done(): | |||
| 423 | + | } | |||
| 424 | + | running.Add(-1) | |||
| 425 | + | }) | |||
| 426 | + | select { | |||
| 427 | + | case <-started: | |||
| 428 | + | case <-time.After(5 * time.Second): | |||
| 429 | + | t.Fatal("notification before start never triggered a listing") | |||
| 430 | + | } | |||
| 431 | + | for i := 0; i < 100; i++ { | |||
| 432 | + | r.notify() // all arrive during the running listing | |||
| 433 | + | } | |||
| 434 | + | close(release) | |||
| 435 | + | eventually(t, "refresher to go idle", func() bool { | |||
| 436 | + | r.mu.Lock() | |||
| 437 | + | defer r.mu.Unlock() | |||
| 438 | + | return !r.running | |||
| 439 | + | }) | |||
| 440 | + | if got := runs.Load(); got != 2 { | |||
| 441 | + | t.Errorf("runs = %d, want exactly 2 (initial + one coalesced re-run)", got) | |||
| 442 | + | } | |||
| 443 | + | if got := maxRunning.Load(); got != 1 { | |||
| 444 | + | t.Errorf("max concurrent listings = %d, want 1", got) | |||
| 445 | + | } | |||
| 446 | + | } | |||
| 447 | + | ||||
| 448 | + | func TestToolRefresherStop(t *testing.T) { | |||
| 449 | + | r := newToolRefresher() | |||
| 450 | + | entered := make(chan struct{}) | |||
| 451 | + | finished := make(chan struct{}) | |||
| 452 | + | var runs atomic.Int32 | |||
| 453 | + | r.start(func(ctx context.Context) { | |||
| 454 | + | runs.Add(1) | |||
| 455 | + | close(entered) | |||
| 456 | + | <-ctx.Done() // only a stop can end this listing | |||
| 457 | + | close(finished) | |||
| 458 | + | }) | |||
| 459 | + | r.notify() | |||
| 460 | + | <-entered | |||
| 461 | + | r.notify() // dirty again while running | |||
| 462 | + | r.stop() | |||
| 463 | + | select { | |||
| 464 | + | case <-finished: | |||
| 465 | + | case <-time.After(5 * time.Second): | |||
| 466 | + | t.Fatal("stop did not cancel the running listing") | |||
| 467 | + | } | |||
| 468 | + | r.notify() | |||
| 469 | + | eventually(t, "refresher idle", func() bool { | |||
| 470 | + | r.mu.Lock() | |||
| 471 | + | defer r.mu.Unlock() | |||
| 472 | + | return !r.running | |||
| 473 | + | }) | |||
| 474 | + | if got := runs.Load(); got != 1 { | |||
| 475 | + | t.Errorf("runs after stop = %d, want 1", got) | |||
| 476 | + | } | |||
| 477 | + | } | |||
| 478 | + | ||||
| 479 | + | func TestManagerStopsRefresherOnClose(t *testing.T) { | |||
| 480 | + | m, _ := startFake(t, nil, nil) | |||
| 481 | + | s := m.servers["fake"] | |||
| 482 | + | s.mu.Lock() | |||
| 483 | + | r := s.refresher | |||
| 484 | + | s.mu.Unlock() | |||
| 485 | + | if r == nil { | |||
| 486 | + | t.Fatal("no refresher installed") | |||
| 487 | + | } | |||
| 488 | + | _ = m.Close() | |||
| 489 | + | r.mu.Lock() | |||
| 490 | + | stopped := r.stopped | |||
| 491 | + | r.mu.Unlock() | |||
| 492 | + | if !stopped { | |||
| 493 | + | t.Error("refresher still active after Close") | |||
| 494 | + | } | |||
| 495 | + | s.mu.Lock() | |||
| 496 | + | left := s.refresher | |||
| 497 | + | s.mu.Unlock() | |||
| 498 | + | if left != nil { | |||
| 499 | + | t.Error("refresher reference kept after Close") | |||
| 500 | + | } | |||
| 501 | + | } | |||
| 502 | + | ||||
| 503 | + | func TestManagerListChangedDuringInitialListing(t *testing.T) { | |||
| 504 | + | // The server announces a changed tool set right after answering the last | |||
| 505 | + | // page of the first listing, before Start has published the client. | |||
| 506 | + | m, _ := startFake(t, map[string]string{"FAKE_NOTIFY_ON_LIST": "1"}, nil) | |||
| 507 | + | eventually(t, "tool added during initial listing to appear", func() bool { | |||
| 508 | + | for _, tool := range m.ServerTools("fake") { | |||
| 509 | + | if tool.Name == "extra" { | |||
| 510 | + | return true | |||
| 511 | + | } | |||
| 512 | + | } | |||
| 513 | + | return false | |||
| 514 | + | }) | |||
| 515 | + | } | |||
| 516 | + | ||||
| 517 | + | // ---- 7: stderr draining ---- | |||
| 518 | + | ||||
| 519 | + | func TestStdioStderrFloodDoesNotBlockChild(t *testing.T) { | |||
| 520 | + | c, _ := connectFake(t, nil) | |||
| 521 | + | // 3 MiB on one stderr line, then a short one. The child blocks on the | |||
| 522 | + | // stderr pipe unless the client keeps draining it. | |||
| 523 | + | if got := callText(t, c, "stderr_flood", `{}`); got != "flooded" { | |||
| 524 | + | t.Fatalf("got %q", got) | |||
| 525 | + | } | |||
| 526 | + | st := c.t.(*stdioTransport) | |||
| 527 | + | eventually(t, "stderr tail to include the last line", func() bool { | |||
| 528 | + | return strings.Contains(st.stderrTail(), "after-flood") | |||
| 529 | + | }) | |||
| 530 | + | tail := st.stderrTail() | |||
| 531 | + | if !strings.Contains(tail, "[line truncated]") { | |||
| 532 | + | t.Error("long line should be marked truncated") | |||
| 533 | + | } | |||
| 534 | + | if len(tail) > stderrTailLines*(maxStderrLine+64) { | |||
| 535 | + | t.Errorf("tail not bounded: %d bytes", len(tail)) | |||
| 536 | + | } | |||
| 537 | + | // The server still works afterwards. | |||
| 538 | + | if got := callText(t, c, "echo", `{"text":"ok"}`); got != "echo: ok" { | |||
| 539 | + | t.Errorf("after flood: %q", got) | |||
| 540 | + | } | |||
| 541 | + | } | |||
| 542 | + | ||||
| 543 | + | // ---- 8, 9: cancellation notifications ---- | |||
| 544 | + | ||||
| 545 | + | func TestStdioCancelNotificationRules(t *testing.T) { | |||
| 546 | + | // initialize never answers within the deadline. | |||
| 547 | + | cfg, _ := fakeStdioConfig(t, map[string]string{"FAKE_INIT_DELAY_MS": "5000"}) | |||
| 548 | + | st, err := startStdio("x", cfg, "") | |||
| 549 | + | if err != nil { | |||
| 550 | + | t.Fatal(err) | |||
| 551 | + | } | |||
| 552 | + | defer st.close() | |||
| 553 | + | ctx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond) | |||
| 554 | + | defer cancel() | |||
| 555 | + | if _, err := st.request(ctx, 1, "initialize", map[string]any{"protocolVersion": LatestProtocolVersion}); !errors.Is(err, context.DeadlineExceeded) { | |||
| 556 | + | t.Fatalf("err = %v", err) | |||
| 557 | + | } | |||
| 558 | + | if got := st.cancelsSent.Load(); got != 0 { | |||
| 559 | + | t.Errorf("notifications/cancelled sent for initialize (%d)", got) | |||
| 560 | + | } | |||
| 561 | + | // Any other method is announced exactly once, without blocking the caller. | |||
| 562 | + | ctx2, cancel2 := context.WithCancel(context.Background()) | |||
| 563 | + | cancel2() | |||
| 564 | + | start := time.Now() | |||
| 565 | + | _, err = st.request(ctx2, 2, "tools/call", map[string]any{"name": "slow", "arguments": map[string]any{}}) | |||
| 566 | + | if !errors.Is(err, context.Canceled) { | |||
| 567 | + | t.Fatalf("err = %v", err) | |||
| 568 | + | } | |||
| 569 | + | if time.Since(start) > time.Second { | |||
| 570 | + | t.Error("caller delayed by the cancel notification") | |||
| 571 | + | } | |||
| 572 | + | if got := st.cancelsSent.Load(); got != 1 { | |||
| 573 | + | t.Errorf("cancel count for tools/call = %d, want 1", got) | |||
| 574 | + | } | |||
| 575 | + | } | |||
| 576 | + | ||||
| 577 | + | func TestHTTPInitializeNotCancelledOnTimeout(t *testing.T) { | |||
| 578 | + | fake := newHTTPFake(false) | |||
| 579 | + | fake.hangInit = true | |||
| 580 | + | srv := fake.start(t) | |||
| 581 | + | cfg := httpServerConfig(srv) | |||
| 582 | + | ctx, cancel := context.WithTimeout(context.Background(), 200*time.Millisecond) | |||
| 583 | + | defer cancel() | |||
| 584 | + | _, err := Connect(ctx, "web", cfg, "", ClientInfo{}, nil) | |||
| 585 | + | if !errors.Is(err, context.DeadlineExceeded) { | |||
| 586 | + | t.Fatalf("err = %v, want deadline exceeded", err) | |||
| 587 | + | } | |||
| 588 | + | for _, r := range fake.snapshot() { | |||
| 589 | + | if r.RPCMethod == "notifications/cancelled" { | |||
| 590 | + | t.Fatal("notifications/cancelled sent for initialize") | |||
| 591 | + | } | |||
| 592 | + | } | |||
| 593 | + | } | |||
| 594 | + | ||||
| 595 | + | func TestHTTPCancelNotificationIsBoundedAndOwnedByTransport(t *testing.T) { | |||
| 596 | + | fake := newHTTPFake(false) | |||
| 597 | + | fake.cancelGate = make(chan struct{}) | |||
| 598 | + | t.Cleanup(func() { close(fake.cancelGate) }) | |||
| 599 | + | m, failures := startHTTP(t, fake, func(c *ServerConfig) { c.ToolTimeoutSec = 0.2 }) | |||
| 600 | + | if len(failures) != 0 { | |||
| 601 | + | t.Fatal(failures) | |||
| 602 | + | } | |||
| 603 | + | start := time.Now() | |||
| 604 | + | res := execTool(t, findTool(t, m, "mcp__web__slow"), `{}`) | |||
| 605 | + | if !res.IsError || !strings.Contains(resultText(res), "toolTimeoutSec") { | |||
| 606 | + | t.Fatalf("got %q", resultText(res)) | |||
| 607 | + | } | |||
| 608 | + | // The cancellation POST is stuck on the server (gate closed); the caller | |||
| 609 | + | // must not wait for it. | |||
| 610 | + | if d := time.Since(start); d > 1500*time.Millisecond { | |||
| 611 | + | t.Fatalf("caller delayed by the cancel notification: %v", d) | |||
| 612 | + | } | |||
| 613 | + | eventually(t, "cancel notification to reach the server", func() bool { | |||
| 614 | + | for _, r := range fake.snapshot() { | |||
| 615 | + | if r.RPCMethod == "notifications/cancelled" { | |||
| 616 | + | return true | |||
| 617 | + | } | |||
| 618 | + | } | |||
| 619 | + | return false | |||
| 620 | + | }) | |||
| 621 | + | // Closing the transport aborts the stuck notification; otherwise the | |||
| 622 | + | // httptest server would wait for it at cleanup. | |||
| 623 | + | closed := make(chan struct{}) | |||
| 624 | + | go func() { _ = m.Close(); close(closed) }() | |||
| 625 | + | select { | |||
| 626 | + | case <-closed: | |||
| 627 | + | case <-time.After(10 * time.Second): | |||
| 628 | + | t.Fatal("Close hung behind the cancel notification") | |||
| 629 | + | } | |||
| 630 | + | } | |||
| 631 | + | ||||
| 632 | + | // ---- 11: file URIs ---- | |||
| 633 | + | ||||
| 634 | + | func TestFileURI(t *testing.T) { | |||
| 635 | + | cases := []struct { | |||
| 636 | + | name, path string | |||
| 637 | + | windows bool | |||
| 638 | + | want string | |||
| 639 | + | }{ | |||
| 640 | + | {"unix", "/home/me/proj", false, "file:///home/me/proj"}, | |||
| 641 | + | {"unix space", "/home/me/my proj", false, "file:///home/me/my%20proj"}, | |||
| 642 | + | {"drive", `C:\Users\me\proj`, true, "file:///C:/Users/me/proj"}, | |||
| 643 | + | {"drive space", `D:\My Docs\x`, true, "file:///D:/My%20Docs/x"}, | |||
| 644 | + | {"drive root", `C:\`, true, "file:///C:/"}, | |||
| 645 | + | {"drive forward slashes", `C:/x/y`, true, "file:///C:/x/y"}, | |||
| 646 | + | {"unc", `\\host\share\dir`, true, "file://host/share/dir"}, | |||
| 647 | + | {"unc share only", `\\host\share`, true, "file://host/share"}, | |||
| 648 | + | {"extended drive", `\\?\C:\x`, true, "file:///C:/x"}, | |||
| 649 | + | {"extended unc", `\\?\UNC\host\share\d`, true, "file://host/share/d"}, | |||
| 650 | + | } | |||
| 651 | + | for _, c := range cases { | |||
| 652 | + | if got := fileURI(c.path, c.windows); got != c.want { | |||
| 653 | + | t.Errorf("%s: fileURI(%q) = %q, want %q", c.name, c.path, got, c.want) | |||
| 654 | + | } | |||
| 655 | + | } | |||
| 656 | + | } | |||
| 657 | + | ||||
| 658 | + | // ---- 12: DELETE shape ---- | |||
| 659 | + | ||||
| 660 | + | func TestHTTPDeleteCarriesProtocolVersionAndHeaders(t *testing.T) { | |||
| 661 | + | t.Setenv("MCP_T_BEARER", "tok-delete-1") | |||
| 662 | + | fake := newHTTPFake(false) | |||
| 663 | + | m, failures := startHTTP(t, fake, func(c *ServerConfig) { | |||
| 664 | + | c.BearerTokenEnvVar = "MCP_T_BEARER" | |||
| 665 | + | c.Headers = map[string]string{"X-Static": "static"} | |||
| 666 | + | }) | |||
| 667 | + | if len(failures) != 0 { | |||
| 668 | + | t.Fatal(failures) | |||
| 669 | + | } | |||
| 670 | + | _ = m.Close() | |||
| 671 | + | select { | |||
| 672 | + | case <-fake.deleted: | |||
| 673 | + | case <-time.After(5 * time.Second): | |||
| 674 | + | t.Fatal("no DELETE") | |||
| 675 | + | } | |||
| 676 | + | var del *httpRecord | |||
| 677 | + | for _, r := range fake.snapshot() { | |||
| 678 | + | if r.Method == http.MethodDelete { | |||
| 679 | + | r := r | |||
| 680 | + | del = &r | |||
| 681 | + | } | |||
| 682 | + | } | |||
| 683 | + | if del == nil { | |||
| 684 | + | t.Fatal("DELETE not recorded") | |||
| 685 | + | } | |||
| 686 | + | if del.Protocol != "2025-06-18" || del.Session != "sess-abc" || del.Auth != "Bearer tok-delete-1" || del.Static != "static" { | |||
| 687 | + | t.Errorf("DELETE = %+v", *del) | |||
| 688 | + | } | |||
| 689 | + | } | |||
| 690 | + | ||||
| 691 | + | // ---- 14: resource listing budget ---- | |||
| 692 | + | ||||
| 693 | + | func TestListResourcesAppliesBudgetWhileFormatting(t *testing.T) { | |||
| 694 | + | m, events := startFake(t, map[string]string{"FAKE_MANY_RESOURCES": "2000"}, nil) | |||
| 695 | + | text := resultText(execTool(t, findTool(t, m, ListResourcesToolName), `{}`)) | |||
| 696 | + | marker := "\n... [truncated at 51200 bytes]" | |||
| 697 | + | if !strings.Contains(text, marker) { | |||
| 698 | + | t.Fatalf("no truncation marker (len=%d)", len(text)) | |||
| 699 | + | } | |||
| 700 | + | if len(text) > 50*1024+len(marker) { | |||
| 701 | + | t.Errorf("output not bounded: %d bytes", len(text)) | |||
| 702 | + | } | |||
| 703 | + | if !strings.Contains(text, "mem://item/0") || strings.Contains(text, "mem://item/1999") { | |||
| 704 | + | t.Error("expected the head of the list, not the tail") | |||
| 705 | + | } | |||
| 706 | + | // Once the budget is spent no further requests are made. | |||
| 707 | + | if n := countEvents(events, "list-templates"); n != 0 { | |||
| 708 | + | t.Errorf("templates were fetched after the budget was exhausted (%d)", n) | |||
| 709 | + | } | |||
| 710 | + | if n := countEvents(events, "list-resources"); n != 1 { | |||
| 711 | + | t.Errorf("list-resources calls = %d", n) | |||
| 712 | + | } | |||
| 713 | + | } | |||
| 714 | + | ||||
| 715 | + | func TestBudgetedText(t *testing.T) { | |||
| 716 | + | b := &budgetedText{} | |||
| 717 | + | b.write(strings.Repeat("a", maxResultTextBytes-1)) | |||
| 718 | + | if b.full() { | |||
| 719 | + | t.Fatal("full too early") | |||
| 720 | + | } | |||
| 721 | + | b.write("bcdef") | |||
| 722 | + | if !b.full() { | |||
| 723 | + | t.Fatal("not full after overflow") | |||
| 724 | + | } | |||
| 725 | + | before := b.String() | |||
| 726 | + | b.write("more") | |||
| 727 | + | b.printf("%s", "more") | |||
| 728 | + | if b.String() != before { | |||
| 729 | + | t.Error("writes accepted after the budget was exhausted") | |||
| 730 | + | } | |||
| 731 | + | if !strings.HasSuffix(before, "[truncated at 51200 bytes]") { | |||
| 732 | + | t.Errorf("tail = %q", before[len(before)-40:]) | |||
| 733 | + | } | |||
| 734 | + | } | |||
M packages/agent/mcp/http.go +136 -27
| @@ -13,7 +13,9 @@ import ( | |||
| 13 | 13 | "os" | |
| 14 | 14 | "strings" | |
| 15 | 15 | "sync" | |
| 16 | + | "sync/atomic" | |
| 16 | 17 | "time" | |
| 18 | + | "unicode/utf8" | |
| 17 | 19 | ) | |
| 18 | 20 | ||
| 19 | 21 | const ( | |
| @@ -31,18 +33,26 @@ type httpTransport struct { | |||
| 31 | 33 | url string | |
| 32 | 34 | headers map[string]string | |
| 33 | 35 | client *http.Client | |
| 36 | + | red *redactor | |
| 34 | 37 | ||
| 35 | 38 | mu sync.Mutex | |
| 36 | 39 | sessionID string | |
| 37 | 40 | protocol string | |
| 38 | 41 | handler incomingHandler | |
| 39 | 42 | err error | |
| 43 | + | // sessionGone is set when the server dropped our session (HTTP 404); the | |
| 44 | + | // DELETE on close is pointless then. | |
| 45 | + | sessionGone bool | |
| 46 | + | deleted bool | |
| 40 | 47 | ||
| 41 | 48 | ctx context.Context | |
| 42 | 49 | cancel context.CancelFunc | |
| 43 | 50 | doneCh chan struct{} | |
| 44 | 51 | closeOnce sync.Once | |
| 45 | 52 | listening bool | |
| 53 | + | ||
| 54 | + | // cancelsSent counts notifications/cancelled sends started. | |
| 55 | + | cancelsSent atomic.Int64 | |
| 46 | 56 | } | |
| 47 | 57 | ||
| 48 | 58 | func newHTTPTransport(name string, cfg ServerConfig, client *http.Client) (*httpTransport, error) { | |
| @@ -71,6 +81,7 @@ func newHTTPTransport(name string, cfg ServerConfig, client *http.Client) (*http | |||
| 71 | 81 | url: cfg.URL, | |
| 72 | 82 | headers: headers, | |
| 73 | 83 | client: client, | |
| 84 | + | red: newRedactor(cfg), | |
| 74 | 85 | ctx: ctx, | |
| 75 | 86 | cancel: cancel, | |
| 76 | 87 | doneCh: make(chan struct{}), | |
| @@ -105,6 +116,17 @@ func (t *httpTransport) newRequest(ctx context.Context, method string, body []by | |||
| 105 | 116 | return req, nil | |
| 106 | 117 | } | |
| 107 | 118 | ||
| 119 | + | // bind returns a context that ends when ctx ends or the transport fails or | |
| 120 | + | // closes, so no request outlives the connection. | |
| 121 | + | func (t *httpTransport) bind(ctx context.Context) (context.Context, context.CancelFunc) { | |
| 122 | + | bound, cancel := context.WithCancel(ctx) | |
| 123 | + | stop := context.AfterFunc(t.ctx, cancel) | |
| 124 | + | return bound, func() { | |
| 125 | + | stop() | |
| 126 | + | cancel() | |
| 127 | + | } | |
| 128 | + | } | |
| 129 | + | ||
| 108 | 130 | func (t *httpTransport) post(ctx context.Context, frame []byte) (*http.Response, error) { | |
| 109 | 131 | req, err := t.newRequest(ctx, http.MethodPost, frame) | |
| 110 | 132 | if err != nil { | |
| @@ -112,7 +134,7 @@ func (t *httpTransport) post(ctx context.Context, frame []byte) (*http.Response, | |||
| 112 | 134 | } | |
| 113 | 135 | resp, err := t.client.Do(req) | |
| 114 | 136 | if err != nil { | |
| 115 | - | return nil, err | |
| 137 | + | return nil, t.red.scrubError(err) | |
| 116 | 138 | } | |
| 117 | 139 | if sid := resp.Header.Get(headerSessionID); sid != "" { | |
| 118 | 140 | t.mu.Lock() | |
| @@ -122,11 +144,27 @@ func (t *httpTransport) post(ctx context.Context, frame []byte) (*http.Response, | |||
| 122 | 144 | if resp.StatusCode >= 300 { | |
| 123 | 145 | defer resp.Body.Close() | |
| 124 | 146 | snippet, _ := io.ReadAll(io.LimitReader(resp.Body, maxErrorBody)) | |
| 125 | - | return nil, httpStatusError(resp.StatusCode, snippet) | |
| 147 | + | return nil, t.statusFailure(req, resp.StatusCode, snippet) | |
| 126 | 148 | } | |
| 127 | 149 | return resp, nil | |
| 128 | 150 | } | |
| 129 | 151 | ||
| 152 | + | // statusFailure builds the error for a non-success status. A 404 for a | |
| 153 | + | // request that carried a session id means the server forgot the session | |
| 154 | + | // (expired or restarted): the transport becomes unusable so the manager | |
| 155 | + | // reconnects with a fresh initialize. The failed call is not replayed. | |
| 156 | + | func (t *httpTransport) statusFailure(req *http.Request, status int, body []byte) error { | |
| 157 | + | err := t.red.statusError(status, body) | |
| 158 | + | if status == http.StatusNotFound && req.Header.Get(headerSessionID) != "" { | |
| 159 | + | err = fmt.Errorf("mcp session expired: %w", err) | |
| 160 | + | t.mu.Lock() | |
| 161 | + | t.sessionGone = true | |
| 162 | + | t.mu.Unlock() | |
| 163 | + | t.fail(err) | |
| 164 | + | } | |
| 165 | + | return err | |
| 166 | + | } | |
| 167 | + | ||
| 130 | 168 | // HTTPStatusError reports a non-success HTTP status from a server. | |
| 131 | 169 | type HTTPStatusError struct { | |
| 132 | 170 | Status int | |
| @@ -145,9 +183,20 @@ func (e *HTTPStatusError) Error() string { | |||
| 145 | 183 | } | |
| 146 | 184 | ||
| 147 | 185 | func httpStatusError(status int, body []byte) error { | |
| 148 | - | text := strings.TrimSpace(string(body)) | |
| 186 | + | return (*redactor)(nil).statusError(status, body) | |
| 187 | + | } | |
| 188 | + | ||
| 189 | + | // statusError builds an HTTPStatusError whose body is redacted, stripped of | |
| 190 | + | // terminal control sequences, and shortened for display. | |
| 191 | + | func (r *redactor) statusError(status int, body []byte) error { | |
| 192 | + | // Redact before truncating so a secret cut at the limit is not half shown. | |
| 193 | + | text := strings.TrimSpace(r.scrubOneLine(string(body))) | |
| 149 | 194 | if len(text) > 300 { | |
| 150 | - | text = text[:300] + "…" | |
| 195 | + | cut := 300 | |
| 196 | + | for cut > 0 && !utf8.RuneStart(text[cut]) { | |
| 197 | + | cut-- | |
| 198 | + | } | |
| 199 | + | text = text[:cut] + "…" | |
| 151 | 200 | } | |
| 152 | 201 | return &HTTPStatusError{Status: status, Body: text} | |
| 153 | 202 | } | |
| @@ -161,18 +210,16 @@ func (t *httpTransport) request(ctx context.Context, id int64, method string, pa | |||
| 161 | 210 | return nil, err | |
| 162 | 211 | } | |
| 163 | 212 | // Bind the request to both the caller and the transport lifetime. | |
| 164 | - | reqCtx, cancel := context.WithCancel(ctx) | |
| 213 | + | reqCtx, cancel := t.bind(ctx) | |
| 165 | 214 | defer cancel() | |
| 166 | - | stop := context.AfterFunc(t.ctx, cancel) | |
| 167 | - | defer stop() | |
| 168 | 215 | resp, err := t.post(reqCtx, frame) | |
| 169 | 216 | if err != nil { | |
| 170 | 217 | if ctx.Err() != nil { | |
| 171 | - | _ = t.notify(context.Background(), "notifications/cancelled", map[string]any{"requestId": id, "reason": ctx.Err().Error()}) | |
| 218 | + | t.cancelRequest(id, method, ctx.Err().Error()) | |
| 172 | 219 | return nil, ctx.Err() | |
| 173 | 220 | } | |
| 174 | 221 | if t.ctx.Err() != nil { | |
| 175 | - | return nil, errors.New("connection closed") | |
| 222 | + | return nil, t.failure() | |
| 176 | 223 | } | |
| 177 | 224 | return nil, err | |
| 178 | 225 | } | |
| @@ -194,7 +241,10 @@ func (t *httpTransport) request(ctx context.Context, id int64, method string, pa | |||
| 194 | 241 | err = readSSEFrames(resp.Body, handle) | |
| 195 | 242 | default: | |
| 196 | 243 | var raw []byte | |
| 197 | - | raw, err = io.ReadAll(io.LimitReader(resp.Body, maxHTTPBody)) | |
| 244 | + | raw, err = io.ReadAll(io.LimitReader(resp.Body, maxHTTPBody+1)) | |
| 245 | + | if err == nil && len(raw) > maxHTTPBody { | |
| 246 | + | err = fmt.Errorf("response exceeds %d bytes", maxHTTPBody) | |
| 247 | + | } | |
| 198 | 248 | if err == nil { | |
| 199 | 249 | var frames []message | |
| 200 | 250 | frames, err = parseFrames(raw) | |
| @@ -207,25 +257,57 @@ func (t *httpTransport) request(ctx context.Context, id int64, method string, pa | |||
| 207 | 257 | } | |
| 208 | 258 | if result == nil { | |
| 209 | 259 | if ctx.Err() != nil { | |
| 210 | - | _ = t.notify(context.Background(), "notifications/cancelled", map[string]any{"requestId": id, "reason": ctx.Err().Error()}) | |
| 260 | + | t.cancelRequest(id, method, ctx.Err().Error()) | |
| 211 | 261 | return nil, ctx.Err() | |
| 212 | 262 | } | |
| 263 | + | if t.ctx.Err() != nil { | |
| 264 | + | return nil, t.failure() | |
| 265 | + | } | |
| 213 | 266 | if err != nil { | |
| 214 | - | return nil, fmt.Errorf("read response: %w", err) | |
| 267 | + | return nil, fmt.Errorf("read response: %w", t.red.scrubError(err)) | |
| 215 | 268 | } | |
| 216 | 269 | return nil, errors.New("server closed the response without a result") | |
| 217 | 270 | } | |
| 218 | 271 | if result.Error != nil { | |
| 219 | - | return nil, result.Error | |
| 272 | + | return nil, t.red.scrubRPCError(result.Error) | |
| 220 | 273 | } | |
| 221 | 274 | return result.Result, nil | |
| 222 | 275 | } | |
| 223 | 276 | ||
| 277 | + | // cancelRequest announces an abandoned request on a best-effort basis. It | |
| 278 | + | // never delays the caller, is bounded by cancelNotifyTimeout, and ends with | |
| 279 | + | // the transport. initialize is never cancelled (MCP lifecycle). | |
| 280 | + | func (t *httpTransport) cancelRequest(id int64, method, reason string) { | |
| 281 | + | if !shouldSendCancel(method) || t.ctx.Err() != nil { | |
| 282 | + | return | |
| 283 | + | } | |
| 284 | + | t.cancelsSent.Add(1) | |
| 285 | + | go func() { | |
| 286 | + | ctx, cancel := context.WithTimeout(t.ctx, cancelNotifyTimeout) | |
| 287 | + | defer cancel() | |
| 288 | + | _ = t.notify(ctx, "notifications/cancelled", map[string]any{"requestId": id, "reason": reason}) | |
| 289 | + | }() | |
| 290 | + | } | |
| 291 | + | ||
| 292 | + | // errSSEEventTooLarge reports an SSE event whose aggregated data exceeds the | |
| 293 | + | // body limit. | |
| 294 | + | type errSSEEventTooLarge struct{ limit int } | |
| 295 | + | ||
| 296 | + | func (e errSSEEventTooLarge) Error() string { | |
| 297 | + | return fmt.Sprintf("sse event data exceeds %d bytes", e.limit) | |
| 298 | + | } | |
| 299 | + | ||
| 224 | 300 | // readSSEFrames delivers each SSE data payload as JSON-RPC frames until | |
| 225 | 301 | // handle returns true or the stream ends. | |
| 226 | 302 | func readSSEFrames(r io.Reader, handle func(message) bool) error { | |
| 303 | + | return readSSEFramesLimit(r, maxHTTPBody, handle) | |
| 304 | + | } | |
| 305 | + | ||
| 306 | + | // readSSEFramesLimit is readSSEFrames with an explicit bound on the data of | |
| 307 | + | // one event, counted across all of its data: lines. | |
| 308 | + | func readSSEFramesLimit(r io.Reader, limit int, handle func(message) bool) error { | |
| 227 | 309 | sc := bufio.NewScanner(r) | |
| 228 | - | sc.Buffer(make([]byte, 0, 64*1024), maxHTTPBody) | |
| 310 | + | sc.Buffer(make([]byte, 0, 64*1024), limit+1) | |
| 229 | 311 | var data strings.Builder | |
| 230 | 312 | flush := func() bool { | |
| 231 | 313 | if data.Len() == 0 { | |
| @@ -257,16 +339,27 @@ func readSSEFrames(r io.Reader, handle func(message) bool) error { | |||
| 257 | 339 | field, value, _ := strings.Cut(line, ":") | |
| 258 | 340 | value = strings.TrimPrefix(value, " ") | |
| 259 | 341 | if field == "data" { | |
| 342 | + | extra := len(value) | |
| 343 | + | if data.Len() > 0 { | |
| 344 | + | extra++ | |
| 345 | + | } | |
| 346 | + | if data.Len()+extra > limit { | |
| 347 | + | return errSSEEventTooLarge{limit: limit} | |
| 348 | + | } | |
| 260 | 349 | if data.Len() > 0 { | |
| 261 | 350 | data.WriteByte('\n') | |
| 262 | 351 | } | |
| 263 | 352 | data.WriteString(value) | |
| 264 | 353 | } | |
| 265 | 354 | } | |
| 266 | - | if flush() { | |
| 267 | - | return nil | |
| 355 | + | if err := sc.Err(); err != nil { | |
| 356 | + | if errors.Is(err, bufio.ErrTooLong) { | |
| 357 | + | return errSSEEventTooLarge{limit: limit} | |
| 358 | + | } | |
| 359 | + | return err | |
| 268 | 360 | } | |
| 269 | - | return sc.Err() | |
| 361 | + | flush() | |
| 362 | + | return nil | |
| 270 | 363 | } | |
| 271 | 364 | ||
| 272 | 365 | func (t *httpTransport) dispatchIncoming(msg message) { | |
| @@ -305,6 +398,8 @@ func (t *httpTransport) notify(ctx context.Context, method string, params any) e | |||
| 305 | 398 | } | |
| 306 | 399 | ctx, cancel := context.WithTimeout(ctx, 30*time.Second) | |
| 307 | 400 | defer cancel() | |
| 401 | + | ctx, stop := t.bind(ctx) | |
| 402 | + | defer stop() | |
| 308 | 403 | resp, err := t.post(ctx, frame) | |
| 309 | 404 | if err != nil { | |
| 310 | 405 | return err | |
| @@ -341,6 +436,12 @@ func (t *httpTransport) listen() { | |||
| 341 | 436 | } | |
| 342 | 437 | if resp.StatusCode != http.StatusOK { | |
| 343 | 438 | resp.Body.Close() | |
| 439 | + | if resp.StatusCode == http.StatusNotFound && req.Header.Get(headerSessionID) != "" { | |
| 440 | + | t.mu.Lock() | |
| 441 | + | t.sessionGone = true | |
| 442 | + | t.mu.Unlock() | |
| 443 | + | t.fail(fmt.Errorf("mcp session expired: %w", t.red.statusError(resp.StatusCode, nil))) | |
| 444 | + | } | |
| 344 | 445 | // 405: no server stream. Other statuses: give up quietly; the | |
| 345 | 446 | // POST path still works. | |
| 346 | 447 | return | |
| @@ -389,24 +490,32 @@ func (t *httpTransport) failure() error { | |||
| 389 | 490 | return t.err | |
| 390 | 491 | } | |
| 391 | 492 | ||
| 392 | - | // close ends the session with DELETE (best effort) and stops all streams. | |
| 393 | - | func (t *httpTransport) close() error { | |
| 493 | + | // fail marks the transport unusable: pending and future requests error out, | |
| 494 | + | // background streams stop, and done() is closed so the manager reconnects. | |
| 495 | + | func (t *httpTransport) fail(err error) { | |
| 394 | 496 | t.mu.Lock() | |
| 395 | 497 | if t.err == nil { | |
| 396 | - | t.err = errors.New("connection closed") | |
| 498 | + | t.err = err | |
| 397 | 499 | } | |
| 398 | - | sid := t.sessionID | |
| 399 | 500 | t.mu.Unlock() | |
| 400 | 501 | t.cancel() | |
| 401 | 502 | t.closeOnce.Do(func() { close(t.doneCh) }) | |
| 402 | - | if sid != "" { | |
| 503 | + | } | |
| 504 | + | ||
| 505 | + | // close ends the session with DELETE (best effort) and stops all streams. | |
| 506 | + | func (t *httpTransport) close() error { | |
| 507 | + | t.fail(errors.New("connection closed")) | |
| 508 | + | t.mu.Lock() | |
| 509 | + | sid, gone := t.sessionID, t.sessionGone | |
| 510 | + | send := !t.deleted && sid != "" && !gone | |
| 511 | + | t.deleted = true | |
| 512 | + | t.mu.Unlock() | |
| 513 | + | if send { | |
| 403 | 514 | ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second) | |
| 404 | 515 | defer cancel() | |
| 405 | - | if req, err := http.NewRequestWithContext(ctx, http.MethodDelete, t.url, nil); err == nil { | |
| 406 | - | for k, v := range t.headers { | |
| 407 | - | req.Header.Set(k, v) | |
| 408 | - | } | |
| 409 | - | req.Header.Set(headerSessionID, sid) | |
| 516 | + | // Built like every other request so it carries the protocol version | |
| 517 | + | // and configured headers as well as the session id. | |
| 518 | + | if req, err := t.newRequest(ctx, http.MethodDelete, nil); err == nil { | |
| 410 | 519 | if resp, err := t.client.Do(req); err == nil { | |
| 411 | 520 | resp.Body.Close() | |
| 412 | 521 | } | |
| @@ -13,7 +13,9 @@ import ( | |||||
| 13 | "os" | 13 | "os" | ||
| 14 | "strings" | 14 | "strings" | ||
| 15 | "sync" | 15 | "sync" | ||
| 16 | + | "sync/atomic" | |||
| 16 | "time" | 17 | "time" | ||
| 18 | + | "unicode/utf8" | |||
| 17 | ) | 19 | ) | ||
| 18 | 20 | ||||
| 19 | const ( | 21 | const ( | ||
| @@ -31,18 +33,26 @@ type httpTransport struct { | |||||
| 31 | url string | 33 | url string | ||
| 32 | headers map[string]string | 34 | headers map[string]string | ||
| 33 | client *http.Client | 35 | client *http.Client | ||
| 36 | + | red *redactor | |||
| 34 | 37 | ||||
| 35 | mu sync.Mutex | 38 | mu sync.Mutex | ||
| 36 | sessionID string | 39 | sessionID string | ||
| 37 | protocol string | 40 | protocol string | ||
| 38 | handler incomingHandler | 41 | handler incomingHandler | ||
| 39 | err error | 42 | err error | ||
| 43 | + | // sessionGone is set when the server dropped our session (HTTP 404); the | |||
| 44 | + | // DELETE on close is pointless then. | |||
| 45 | + | sessionGone bool | |||
| 46 | + | deleted bool | |||
| 40 | 47 | ||||
| 41 | ctx context.Context | 48 | ctx context.Context | ||
| 42 | cancel context.CancelFunc | 49 | cancel context.CancelFunc | ||
| 43 | doneCh chan struct{} | 50 | doneCh chan struct{} | ||
| 44 | closeOnce sync.Once | 51 | closeOnce sync.Once | ||
| 45 | listening bool | 52 | listening bool | ||
| 53 | + | ||||
| 54 | + | // cancelsSent counts notifications/cancelled sends started. | |||
| 55 | + | cancelsSent atomic.Int64 | |||
| 46 | } | 56 | } | ||
| 47 | 57 | ||||
| 48 | func newHTTPTransport(name string, cfg ServerConfig, client *http.Client) (*httpTransport, error) { | 58 | func newHTTPTransport(name string, cfg ServerConfig, client *http.Client) (*httpTransport, error) { | ||
| @@ -71,6 +81,7 @@ func newHTTPTransport(name string, cfg ServerConfig, client *http.Client) (*http | |||||
| 71 | url: cfg.URL, | 81 | url: cfg.URL, | ||
| 72 | headers: headers, | 82 | headers: headers, | ||
| 73 | client: client, | 83 | client: client, | ||
| 84 | + | red: newRedactor(cfg), | |||
| 74 | ctx: ctx, | 85 | ctx: ctx, | ||
| 75 | cancel: cancel, | 86 | cancel: cancel, | ||
| 76 | doneCh: make(chan struct{}), | 87 | doneCh: make(chan struct{}), | ||
| @@ -105,6 +116,17 @@ func (t *httpTransport) newRequest(ctx context.Context, method string, body []by | |||||
| 105 | return req, nil | 116 | return req, nil | ||
| 106 | } | 117 | } | ||
| 107 | 118 | ||||
| 119 | + | // bind returns a context that ends when ctx ends or the transport fails or | |||
| 120 | + | // closes, so no request outlives the connection. | |||
| 121 | + | func (t *httpTransport) bind(ctx context.Context) (context.Context, context.CancelFunc) { | |||
| 122 | + | bound, cancel := context.WithCancel(ctx) | |||
| 123 | + | stop := context.AfterFunc(t.ctx, cancel) | |||
| 124 | + | return bound, func() { | |||
| 125 | + | stop() | |||
| 126 | + | cancel() | |||
| 127 | + | } | |||
| 128 | + | } | |||
| 129 | + | ||||
| 108 | func (t *httpTransport) post(ctx context.Context, frame []byte) (*http.Response, error) { | 130 | func (t *httpTransport) post(ctx context.Context, frame []byte) (*http.Response, error) { | ||
| 109 | req, err := t.newRequest(ctx, http.MethodPost, frame) | 131 | req, err := t.newRequest(ctx, http.MethodPost, frame) | ||
| 110 | if err != nil { | 132 | if err != nil { | ||
| @@ -112,7 +134,7 @@ func (t *httpTransport) post(ctx context.Context, frame []byte) (*http.Response, | |||||
| 112 | } | 134 | } | ||
| 113 | resp, err := t.client.Do(req) | 135 | resp, err := t.client.Do(req) | ||
| 114 | if err != nil { | 136 | if err != nil { | ||
| 115 | - | return nil, err | 137 | + | return nil, t.red.scrubError(err) |
| 116 | } | 138 | } | ||
| 117 | if sid := resp.Header.Get(headerSessionID); sid != "" { | 139 | if sid := resp.Header.Get(headerSessionID); sid != "" { | ||
| 118 | t.mu.Lock() | 140 | t.mu.Lock() | ||
| @@ -122,11 +144,27 @@ func (t *httpTransport) post(ctx context.Context, frame []byte) (*http.Response, | |||||
| 122 | if resp.StatusCode >= 300 { | 144 | if resp.StatusCode >= 300 { | ||
| 123 | defer resp.Body.Close() | 145 | defer resp.Body.Close() | ||
| 124 | snippet, _ := io.ReadAll(io.LimitReader(resp.Body, maxErrorBody)) | 146 | snippet, _ := io.ReadAll(io.LimitReader(resp.Body, maxErrorBody)) | ||
| 125 | - | return nil, httpStatusError(resp.StatusCode, snippet) | 147 | + | return nil, t.statusFailure(req, resp.StatusCode, snippet) |
| 126 | } | 148 | } | ||
| 127 | return resp, nil | 149 | return resp, nil | ||
| 128 | } | 150 | } | ||
| 129 | 151 | ||||
| 152 | + | // statusFailure builds the error for a non-success status. A 404 for a | |||
| 153 | + | // request that carried a session id means the server forgot the session | |||
| 154 | + | // (expired or restarted): the transport becomes unusable so the manager | |||
| 155 | + | // reconnects with a fresh initialize. The failed call is not replayed. | |||
| 156 | + | func (t *httpTransport) statusFailure(req *http.Request, status int, body []byte) error { | |||
| 157 | + | err := t.red.statusError(status, body) | |||
| 158 | + | if status == http.StatusNotFound && req.Header.Get(headerSessionID) != "" { | |||
| 159 | + | err = fmt.Errorf("mcp session expired: %w", err) | |||
| 160 | + | t.mu.Lock() | |||
| 161 | + | t.sessionGone = true | |||
| 162 | + | t.mu.Unlock() | |||
| 163 | + | t.fail(err) | |||
| 164 | + | } | |||
| 165 | + | return err | |||
| 166 | + | } | |||
| 167 | + | ||||
| 130 | // HTTPStatusError reports a non-success HTTP status from a server. | 168 | // HTTPStatusError reports a non-success HTTP status from a server. | ||
| 131 | type HTTPStatusError struct { | 169 | type HTTPStatusError struct { | ||
| 132 | Status int | 170 | Status int | ||
| @@ -145,9 +183,20 @@ func (e *HTTPStatusError) Error() string { | |||||
| 145 | } | 183 | } | ||
| 146 | 184 | ||||
| 147 | func httpStatusError(status int, body []byte) error { | 185 | func httpStatusError(status int, body []byte) error { | ||
| 148 | - | text := strings.TrimSpace(string(body)) | 186 | + | return (*redactor)(nil).statusError(status, body) |
| 187 | + | } | |||
| 188 | + | ||||
| 189 | + | // statusError builds an HTTPStatusError whose body is redacted, stripped of | |||
| 190 | + | // terminal control sequences, and shortened for display. | |||
| 191 | + | func (r *redactor) statusError(status int, body []byte) error { | |||
| 192 | + | // Redact before truncating so a secret cut at the limit is not half shown. | |||
| 193 | + | text := strings.TrimSpace(r.scrubOneLine(string(body))) | |||
| 149 | if len(text) > 300 { | 194 | if len(text) > 300 { | ||
| 150 | - | text = text[:300] + "…" | 195 | + | cut := 300 |
| 196 | + | for cut > 0 && !utf8.RuneStart(text[cut]) { | |||
| 197 | + | cut-- | |||
| 198 | + | } | |||
| 199 | + | text = text[:cut] + "…" | |||
| 151 | } | 200 | } | ||
| 152 | return &HTTPStatusError{Status: status, Body: text} | 201 | return &HTTPStatusError{Status: status, Body: text} | ||
| 153 | } | 202 | } | ||
| @@ -161,18 +210,16 @@ func (t *httpTransport) request(ctx context.Context, id int64, method string, pa | |||||
| 161 | return nil, err | 210 | return nil, err | ||
| 162 | } | 211 | } | ||
| 163 | // Bind the request to both the caller and the transport lifetime. | 212 | // Bind the request to both the caller and the transport lifetime. | ||
| 164 | - | reqCtx, cancel := context.WithCancel(ctx) | 213 | + | reqCtx, cancel := t.bind(ctx) |
| 165 | defer cancel() | 214 | defer cancel() | ||
| 166 | - | stop := context.AfterFunc(t.ctx, cancel) | |||
| 167 | - | defer stop() | |||
| 168 | resp, err := t.post(reqCtx, frame) | 215 | resp, err := t.post(reqCtx, frame) | ||
| 169 | if err != nil { | 216 | if err != nil { | ||
| 170 | if ctx.Err() != nil { | 217 | if ctx.Err() != nil { | ||
| 171 | - | _ = t.notify(context.Background(), "notifications/cancelled", map[string]any{"requestId": id, "reason": ctx.Err().Error()}) | 218 | + | t.cancelRequest(id, method, ctx.Err().Error()) |
| 172 | return nil, ctx.Err() | 219 | return nil, ctx.Err() | ||
| 173 | } | 220 | } | ||
| 174 | if t.ctx.Err() != nil { | 221 | if t.ctx.Err() != nil { | ||
| 175 | - | return nil, errors.New("connection closed") | 222 | + | return nil, t.failure() |
| 176 | } | 223 | } | ||
| 177 | return nil, err | 224 | return nil, err | ||
| 178 | } | 225 | } | ||
| @@ -194,7 +241,10 @@ func (t *httpTransport) request(ctx context.Context, id int64, method string, pa | |||||
| 194 | err = readSSEFrames(resp.Body, handle) | 241 | err = readSSEFrames(resp.Body, handle) | ||
| 195 | default: | 242 | default: | ||
| 196 | var raw []byte | 243 | var raw []byte | ||
| 197 | - | raw, err = io.ReadAll(io.LimitReader(resp.Body, maxHTTPBody)) | 244 | + | raw, err = io.ReadAll(io.LimitReader(resp.Body, maxHTTPBody+1)) |
| 245 | + | if err == nil && len(raw) > maxHTTPBody { | |||
| 246 | + | err = fmt.Errorf("response exceeds %d bytes", maxHTTPBody) | |||
| 247 | + | } | |||
| 198 | if err == nil { | 248 | if err == nil { | ||
| 199 | var frames []message | 249 | var frames []message | ||
| 200 | frames, err = parseFrames(raw) | 250 | frames, err = parseFrames(raw) | ||
| @@ -207,25 +257,57 @@ func (t *httpTransport) request(ctx context.Context, id int64, method string, pa | |||||
| 207 | } | 257 | } | ||
| 208 | if result == nil { | 258 | if result == nil { | ||
| 209 | if ctx.Err() != nil { | 259 | if ctx.Err() != nil { | ||
| 210 | - | _ = t.notify(context.Background(), "notifications/cancelled", map[string]any{"requestId": id, "reason": ctx.Err().Error()}) | 260 | + | t.cancelRequest(id, method, ctx.Err().Error()) |
| 211 | return nil, ctx.Err() | 261 | return nil, ctx.Err() | ||
| 212 | } | 262 | } | ||
| 263 | + | if t.ctx.Err() != nil { | |||
| 264 | + | return nil, t.failure() | |||
| 265 | + | } | |||
| 213 | if err != nil { | 266 | if err != nil { | ||
| 214 | - | return nil, fmt.Errorf("read response: %w", err) | 267 | + | return nil, fmt.Errorf("read response: %w", t.red.scrubError(err)) |
| 215 | } | 268 | } | ||
| 216 | return nil, errors.New("server closed the response without a result") | 269 | return nil, errors.New("server closed the response without a result") | ||
| 217 | } | 270 | } | ||
| 218 | if result.Error != nil { | 271 | if result.Error != nil { | ||
| 219 | - | return nil, result.Error | 272 | + | return nil, t.red.scrubRPCError(result.Error) |
| 220 | } | 273 | } | ||
| 221 | return result.Result, nil | 274 | return result.Result, nil | ||
| 222 | } | 275 | } | ||
| 223 | 276 | ||||
| 277 | + | // cancelRequest announces an abandoned request on a best-effort basis. It | |||
| 278 | + | // never delays the caller, is bounded by cancelNotifyTimeout, and ends with | |||
| 279 | + | // the transport. initialize is never cancelled (MCP lifecycle). | |||
| 280 | + | func (t *httpTransport) cancelRequest(id int64, method, reason string) { | |||
| 281 | + | if !shouldSendCancel(method) || t.ctx.Err() != nil { | |||
| 282 | + | return | |||
| 283 | + | } | |||
| 284 | + | t.cancelsSent.Add(1) | |||
| 285 | + | go func() { | |||
| 286 | + | ctx, cancel := context.WithTimeout(t.ctx, cancelNotifyTimeout) | |||
| 287 | + | defer cancel() | |||
| 288 | + | _ = t.notify(ctx, "notifications/cancelled", map[string]any{"requestId": id, "reason": reason}) | |||
| 289 | + | }() | |||
| 290 | + | } | |||
| 291 | + | ||||
| 292 | + | // errSSEEventTooLarge reports an SSE event whose aggregated data exceeds the | |||
| 293 | + | // body limit. | |||
| 294 | + | type errSSEEventTooLarge struct{ limit int } | |||
| 295 | + | ||||
| 296 | + | func (e errSSEEventTooLarge) Error() string { | |||
| 297 | + | return fmt.Sprintf("sse event data exceeds %d bytes", e.limit) | |||
| 298 | + | } | |||
| 299 | + | ||||
| 224 | // readSSEFrames delivers each SSE data payload as JSON-RPC frames until | 300 | // readSSEFrames delivers each SSE data payload as JSON-RPC frames until | ||
| 225 | // handle returns true or the stream ends. | 301 | // handle returns true or the stream ends. | ||
| 226 | func readSSEFrames(r io.Reader, handle func(message) bool) error { | 302 | func readSSEFrames(r io.Reader, handle func(message) bool) error { | ||
| 303 | + | return readSSEFramesLimit(r, maxHTTPBody, handle) | |||
| 304 | + | } | |||
| 305 | + | ||||
| 306 | + | // readSSEFramesLimit is readSSEFrames with an explicit bound on the data of | |||
| 307 | + | // one event, counted across all of its data: lines. | |||
| 308 | + | func readSSEFramesLimit(r io.Reader, limit int, handle func(message) bool) error { | |||
| 227 | sc := bufio.NewScanner(r) | 309 | sc := bufio.NewScanner(r) | ||
| 228 | - | sc.Buffer(make([]byte, 0, 64*1024), maxHTTPBody) | 310 | + | sc.Buffer(make([]byte, 0, 64*1024), limit+1) |
| 229 | var data strings.Builder | 311 | var data strings.Builder | ||
| 230 | flush := func() bool { | 312 | flush := func() bool { | ||
| 231 | if data.Len() == 0 { | 313 | if data.Len() == 0 { | ||
| @@ -257,16 +339,27 @@ func readSSEFrames(r io.Reader, handle func(message) bool) error { | |||||
| 257 | field, value, _ := strings.Cut(line, ":") | 339 | field, value, _ := strings.Cut(line, ":") | ||
| 258 | value = strings.TrimPrefix(value, " ") | 340 | value = strings.TrimPrefix(value, " ") | ||
| 259 | if field == "data" { | 341 | if field == "data" { | ||
| 342 | + | extra := len(value) | |||
| 343 | + | if data.Len() > 0 { | |||
| 344 | + | extra++ | |||
| 345 | + | } | |||
| 346 | + | if data.Len()+extra > limit { | |||
| 347 | + | return errSSEEventTooLarge{limit: limit} | |||
| 348 | + | } | |||
| 260 | if data.Len() > 0 { | 349 | if data.Len() > 0 { | ||
| 261 | data.WriteByte('\n') | 350 | data.WriteByte('\n') | ||
| 262 | } | 351 | } | ||
| 263 | data.WriteString(value) | 352 | data.WriteString(value) | ||
| 264 | } | 353 | } | ||
| 265 | } | 354 | } | ||
| 266 | - | if flush() { | 355 | + | if err := sc.Err(); err != nil { |
| 267 | - | return nil | 356 | + | if errors.Is(err, bufio.ErrTooLong) { |
| 357 | + | return errSSEEventTooLarge{limit: limit} | |||
| 358 | + | } | |||
| 359 | + | return err | |||
| 268 | } | 360 | } | ||
| 269 | - | return sc.Err() | 361 | + | flush() |
| 362 | + | return nil | |||
| 270 | } | 363 | } | ||
| 271 | 364 | ||||
| 272 | func (t *httpTransport) dispatchIncoming(msg message) { | 365 | func (t *httpTransport) dispatchIncoming(msg message) { | ||
| @@ -305,6 +398,8 @@ func (t *httpTransport) notify(ctx context.Context, method string, params any) e | |||||
| 305 | } | 398 | } | ||
| 306 | ctx, cancel := context.WithTimeout(ctx, 30*time.Second) | 399 | ctx, cancel := context.WithTimeout(ctx, 30*time.Second) | ||
| 307 | defer cancel() | 400 | defer cancel() | ||
| 401 | + | ctx, stop := t.bind(ctx) | |||
| 402 | + | defer stop() | |||
| 308 | resp, err := t.post(ctx, frame) | 403 | resp, err := t.post(ctx, frame) | ||
| 309 | if err != nil { | 404 | if err != nil { | ||
| 310 | return err | 405 | return err | ||
| @@ -341,6 +436,12 @@ func (t *httpTransport) listen() { | |||||
| 341 | } | 436 | } | ||
| 342 | if resp.StatusCode != http.StatusOK { | 437 | if resp.StatusCode != http.StatusOK { | ||
| 343 | resp.Body.Close() | 438 | resp.Body.Close() | ||
| 439 | + | if resp.StatusCode == http.StatusNotFound && req.Header.Get(headerSessionID) != "" { | |||
| 440 | + | t.mu.Lock() | |||
| 441 | + | t.sessionGone = true | |||
| 442 | + | t.mu.Unlock() | |||
| 443 | + | t.fail(fmt.Errorf("mcp session expired: %w", t.red.statusError(resp.StatusCode, nil))) | |||
| 444 | + | } | |||
| 344 | // 405: no server stream. Other statuses: give up quietly; the | 445 | // 405: no server stream. Other statuses: give up quietly; the | ||
| 345 | // POST path still works. | 446 | // POST path still works. | ||
| 346 | return | 447 | return | ||
| @@ -389,24 +490,32 @@ func (t *httpTransport) failure() error { | |||||
| 389 | return t.err | 490 | return t.err | ||
| 390 | } | 491 | } | ||
| 391 | 492 | ||||
| 392 | - | // close ends the session with DELETE (best effort) and stops all streams. | 493 | + | // fail marks the transport unusable: pending and future requests error out, |
| 393 | - | func (t *httpTransport) close() error { | 494 | + | // background streams stop, and done() is closed so the manager reconnects. |
| 495 | + | func (t *httpTransport) fail(err error) { | |||
| 394 | t.mu.Lock() | 496 | t.mu.Lock() | ||
| 395 | if t.err == nil { | 497 | if t.err == nil { | ||
| 396 | - | t.err = errors.New("connection closed") | 498 | + | t.err = err |
| 397 | } | 499 | } | ||
| 398 | - | sid := t.sessionID | |||
| 399 | t.mu.Unlock() | 500 | t.mu.Unlock() | ||
| 400 | t.cancel() | 501 | t.cancel() | ||
| 401 | t.closeOnce.Do(func() { close(t.doneCh) }) | 502 | t.closeOnce.Do(func() { close(t.doneCh) }) | ||
| 402 | - | if sid != "" { | 503 | + | } |
| 504 | + | ||||
| 505 | + | // close ends the session with DELETE (best effort) and stops all streams. | |||
| 506 | + | func (t *httpTransport) close() error { | |||
| 507 | + | t.fail(errors.New("connection closed")) | |||
| 508 | + | t.mu.Lock() | |||
| 509 | + | sid, gone := t.sessionID, t.sessionGone | |||
| 510 | + | send := !t.deleted && sid != "" && !gone | |||
| 511 | + | t.deleted = true | |||
| 512 | + | t.mu.Unlock() | |||
| 513 | + | if send { | |||
| 403 | ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second) | 514 | ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second) | ||
| 404 | defer cancel() | 515 | defer cancel() | ||
| 405 | - | if req, err := http.NewRequestWithContext(ctx, http.MethodDelete, t.url, nil); err == nil { | 516 | + | // Built like every other request so it carries the protocol version |
| 406 | - | for k, v := range t.headers { | 517 | + | // and configured headers as well as the session id. |
| 407 | - | req.Header.Set(k, v) | 518 | + | if req, err := t.newRequest(ctx, http.MethodDelete, nil); err == nil { |
| 408 | - | } | |||
| 409 | - | req.Header.Set(headerSessionID, sid) | |||
| 410 | if resp, err := t.client.Do(req); err == nil { | 519 | if resp, err := t.client.Do(req); err == nil { | ||
| 411 | resp.Body.Close() | 520 | resp.Body.Close() | ||
| 412 | } | 521 | } | ||
M packages/agent/mcp/http_test.go +1 -1
| @@ -44,7 +44,7 @@ func TestManagerHTTPEndToEnd(t *testing.T) { | |||
| 44 | 44 | if len(failures) != 0 { | |
| 45 | 45 | t.Fatalf("failures: %v", failures) | |
| 46 | 46 | } | |
| 47 | - | if st := m.Status()[0]; !st.Connected || st.Transport != "http" || st.Tools != 14 || !st.Resources { | |
| 47 | + | if st := m.Status()[0]; !st.Connected || st.Transport != "http" || st.Tools != fakeToolCount || !st.Resources { | |
| 48 | 48 | t.Fatalf("status = %+v", st) | |
| 49 | 49 | } | |
| 50 | 50 | ||
| @@ -44,7 +44,7 @@ func TestManagerHTTPEndToEnd(t *testing.T) { | |||||
| 44 | if len(failures) != 0 { | 44 | if len(failures) != 0 { | ||
| 45 | t.Fatalf("failures: %v", failures) | 45 | t.Fatalf("failures: %v", failures) | ||
| 46 | } | 46 | } | ||
| 47 | - | if st := m.Status()[0]; !st.Connected || st.Transport != "http" || st.Tools != 14 || !st.Resources { | 47 | + | if st := m.Status()[0]; !st.Connected || st.Transport != "http" || st.Tools != fakeToolCount || !st.Resources { |
| 48 | t.Fatalf("status = %+v", st) | 48 | t.Fatalf("status = %+v", st) | ||
| 49 | } | 49 | } | ||
| 50 | 50 | ||||
M packages/agent/mcp/jsonrpc.go +19 -0
| @@ -4,8 +4,27 @@ import ( | |||
| 4 | 4 | "context" | |
| 5 | 5 | "encoding/json" | |
| 6 | 6 | "fmt" | |
| 7 | + | "time" | |
| 7 | 8 | ) | |
| 8 | 9 | ||
| 10 | + | // cancelNotifyTimeout bounds one best-effort notifications/cancelled send. | |
| 11 | + | const cancelNotifyTimeout = 2 * time.Second | |
| 12 | + | ||
| 13 | + | // shouldSendCancel reports whether abandoning a request of this method may be | |
| 14 | + | // announced with notifications/cancelled. The MCP lifecycle forbids cancelling | |
| 15 | + | // initialize; an abandoned initialize just fails and closes the connection. | |
| 16 | + | func shouldSendCancel(method string) bool { return method != "initialize" } | |
| 17 | + | ||
| 18 | + | // scrubRPCError returns a copy of e whose message is safe to print. | |
| 19 | + | func (r *redactor) scrubRPCError(e *RPCError) *RPCError { | |
| 20 | + | if e == nil { | |
| 21 | + | return nil | |
| 22 | + | } | |
| 23 | + | cp := *e | |
| 24 | + | cp.Message = r.scrubOneLine(e.Message) | |
| 25 | + | return &cp | |
| 26 | + | } | |
| 27 | + | ||
| 9 | 28 | // message is one JSON-RPC 2.0 frame: request, notification, or response. | |
| 10 | 29 | type message struct { | |
| 11 | 30 | JSONRPC string `json:"jsonrpc"` | |
| @@ -4,8 +4,27 @@ import ( | |||||
| 4 | "context" | 4 | "context" | ||
| 5 | "encoding/json" | 5 | "encoding/json" | ||
| 6 | "fmt" | 6 | "fmt" | ||
| 7 | + | "time" | |||
| 7 | ) | 8 | ) | ||
| 8 | 9 | ||||
| 10 | + | // cancelNotifyTimeout bounds one best-effort notifications/cancelled send. | |||
| 11 | + | const cancelNotifyTimeout = 2 * time.Second | |||
| 12 | + | ||||
| 13 | + | // shouldSendCancel reports whether abandoning a request of this method may be | |||
| 14 | + | // announced with notifications/cancelled. The MCP lifecycle forbids cancelling | |||
| 15 | + | // initialize; an abandoned initialize just fails and closes the connection. | |||
| 16 | + | func shouldSendCancel(method string) bool { return method != "initialize" } | |||
| 17 | + | ||||
| 18 | + | // scrubRPCError returns a copy of e whose message is safe to print. | |||
| 19 | + | func (r *redactor) scrubRPCError(e *RPCError) *RPCError { | |||
| 20 | + | if e == nil { | |||
| 21 | + | return nil | |||
| 22 | + | } | |||
| 23 | + | cp := *e | |||
| 24 | + | cp.Message = r.scrubOneLine(e.Message) | |||
| 25 | + | return &cp | |||
| 26 | + | } | |||
| 27 | + | ||||
| 9 | // message is one JSON-RPC 2.0 frame: request, notification, or response. | 28 | // message is one JSON-RPC 2.0 frame: request, notification, or response. | ||
| 10 | type message struct { | 29 | type message struct { | ||
| 11 | JSONRPC string `json:"jsonrpc"` | 30 | JSONRPC string `json:"jsonrpc"` | ||
M packages/agent/mcp/manager.go +106 -8
| @@ -32,6 +32,7 @@ type serverEntry struct { | |||
| 32 | 32 | client *Client | |
| 33 | 33 | tools []Tool | |
| 34 | 34 | err error | |
| 35 | + | refresher *toolRefresher | |
| 35 | 36 | } | |
| 36 | 37 | ||
| 37 | 38 | // Manager owns the MCP connections of one zut process. It is safe for | |
| @@ -108,7 +109,12 @@ func (m *Manager) connect(ctx context.Context, s *serverEntry) (*Client, error) | |||
| 108 | 109 | s.mu.Lock() | |
| 109 | 110 | stale := s.client | |
| 110 | 111 | s.client = nil | |
| 112 | + | staleRefresher := s.refresher | |
| 113 | + | s.refresher = nil | |
| 111 | 114 | s.mu.Unlock() | |
| 115 | + | if staleRefresher != nil { | |
| 116 | + | staleRefresher.stop() | |
| 117 | + | } | |
| 112 | 118 | if stale != nil { | |
| 113 | 119 | _ = stale.Close() | |
| 114 | 120 | } | |
| @@ -120,7 +126,11 @@ func (m *Manager) connect(ctx context.Context, s *serverEntry) (*Client, error) | |||
| 120 | 126 | } | |
| 121 | 127 | startCtx, cancel := withTimeout(ctx, s.cfg.StartupTimeout()) | |
| 122 | 128 | defer cancel() | |
| 123 | - | client, err := Connect(startCtx, s.name, s.cfg, m.cwd, m.info, m.httpClient) | |
| 129 | + | // The refresher exists before the handshake so a list_changed notification | |
| 130 | + | // that arrives while the first listing is in flight is remembered, and | |
| 131 | + | // the listing is repeated once the client is published. | |
| 132 | + | refresher := newToolRefresher() | |
| 133 | + | client, err := connectWith(startCtx, s.name, s.cfg, m.cwd, m.info, m.httpClient, refresher.notify) | |
| 124 | 134 | var tools []Tool | |
| 125 | 135 | if err == nil { | |
| 126 | 136 | tools, err = client.ListTools(startCtx) | |
| @@ -129,6 +139,7 @@ func (m *Manager) connect(ctx context.Context, s *serverEntry) (*Client, error) | |||
| 129 | 139 | } | |
| 130 | 140 | } | |
| 131 | 141 | if err != nil { | |
| 142 | + | refresher.stop() | |
| 132 | 143 | if errors.Is(err, context.DeadlineExceeded) && ctx.Err() == nil { | |
| 133 | 144 | err = fmt.Errorf("timed out after %s; raise startupTimeoutSec if the server is slow to start", s.cfg.StartupTimeout()) | |
| 134 | 145 | } | |
| @@ -138,12 +149,12 @@ func (m *Manager) connect(ctx context.Context, s *serverEntry) (*Client, error) | |||
| 138 | 149 | s.mu.Unlock() | |
| 139 | 150 | return nil, err | |
| 140 | 151 | } | |
| 141 | - | client.SetToolsChanged(func() { m.refreshTools(s, client) }) | |
| 142 | 152 | // Publish under m.mu so Close either sees this client or we see closed. | |
| 143 | 153 | m.mu.Lock() | |
| 144 | 154 | if m.closed { | |
| 145 | 155 | m.mu.Unlock() | |
| 146 | 156 | // Close ran during the handshake; do not leak the new process. | |
| 157 | + | refresher.stop() | |
| 147 | 158 | _ = client.Close() | |
| 148 | 159 | return nil, fmt.Errorf("mcp server %q: manager closed", s.name) | |
| 149 | 160 | } | |
| @@ -151,8 +162,18 @@ func (m *Manager) connect(ctx context.Context, s *serverEntry) (*Client, error) | |||
| 151 | 162 | s.client = client | |
| 152 | 163 | s.tools = filterTools(s.cfg, tools) | |
| 153 | 164 | s.err = nil | |
| 165 | + | old := s.refresher | |
| 166 | + | s.refresher = refresher | |
| 154 | 167 | s.mu.Unlock() | |
| 155 | 168 | m.mu.Unlock() | |
| 169 | + | if old != nil { | |
| 170 | + | old.stop() | |
| 171 | + | } | |
| 172 | + | go func() { | |
| 173 | + | <-client.Done() | |
| 174 | + | refresher.stop() | |
| 175 | + | }() | |
| 176 | + | refresher.start(func(ctx context.Context) { m.refreshTools(ctx, s, client) }) | |
| 156 | 177 | return client, nil | |
| 157 | 178 | } | |
| 158 | 179 | ||
| @@ -172,9 +193,11 @@ func (s *serverEntry) live() *Client { | |||
| 172 | 193 | } | |
| 173 | 194 | ||
| 174 | 195 | // refreshTools re-lists tools after notifications/tools/list_changed. New | |
| 175 | - | // definitions reach the model the next time zut builds a tool registry. | |
| 176 | - | func (m *Manager) refreshTools(s *serverEntry, client *Client) { | |
| 177 | - | ctx, cancel := withTimeout(context.Background(), s.cfg.StartupTimeout()) | |
| 196 | + | // definitions reach the model the next time zut builds a tool registry. It is | |
| 197 | + | // only called from the server's refresher, which serializes and coalesces | |
| 198 | + | // calls. | |
| 199 | + | func (m *Manager) refreshTools(ctx context.Context, s *serverEntry, client *Client) { | |
| 200 | + | ctx, cancel := withTimeout(ctx, s.cfg.StartupTimeout()) | |
| 178 | 201 | defer cancel() | |
| 179 | 202 | tools, err := client.ListTools(ctx) | |
| 180 | 203 | if err != nil { | |
| @@ -187,6 +210,75 @@ func (m *Manager) refreshTools(s *serverEntry, client *Client) { | |||
| 187 | 210 | s.mu.Unlock() | |
| 188 | 211 | } | |
| 189 | 212 | ||
| 213 | + | // toolRefresher coalesces list_changed notifications for one connection into | |
| 214 | + | // serialized re-listings. Any number of notifications that arrive while a | |
| 215 | + | // listing runs cause exactly one more listing afterwards. It does nothing | |
| 216 | + | // until start installs the listing function, so notifications that arrive | |
| 217 | + | // during the initial handshake are remembered rather than raced. | |
| 218 | + | type toolRefresher struct { | |
| 219 | + | mu sync.Mutex | |
| 220 | + | dirty bool | |
| 221 | + | running bool | |
| 222 | + | stopped bool | |
| 223 | + | list func(ctx context.Context) | |
| 224 | + | ctx context.Context | |
| 225 | + | cancel context.CancelFunc | |
| 226 | + | } | |
| 227 | + | ||
| 228 | + | func newToolRefresher() *toolRefresher { | |
| 229 | + | ctx, cancel := context.WithCancel(context.Background()) | |
| 230 | + | return &toolRefresher{ctx: ctx, cancel: cancel} | |
| 231 | + | } | |
| 232 | + | ||
| 233 | + | // notify records that the server's tool list changed. | |
| 234 | + | func (r *toolRefresher) notify() { | |
| 235 | + | r.mu.Lock() | |
| 236 | + | r.dirty = true | |
| 237 | + | r.spawnLocked() | |
| 238 | + | r.mu.Unlock() | |
| 239 | + | } | |
| 240 | + | ||
| 241 | + | // start enables listing with fn and runs it once if a notification arrived | |
| 242 | + | // earlier. | |
| 243 | + | func (r *toolRefresher) start(fn func(ctx context.Context)) { | |
| 244 | + | r.mu.Lock() | |
| 245 | + | r.list = fn | |
| 246 | + | r.spawnLocked() | |
| 247 | + | r.mu.Unlock() | |
| 248 | + | } | |
| 249 | + | ||
| 250 | + | // stop cancels any running listing and ignores later notifications. | |
| 251 | + | func (r *toolRefresher) stop() { | |
| 252 | + | r.mu.Lock() | |
| 253 | + | r.stopped = true | |
| 254 | + | r.dirty = false | |
| 255 | + | r.mu.Unlock() | |
| 256 | + | r.cancel() | |
| 257 | + | } | |
| 258 | + | ||
| 259 | + | func (r *toolRefresher) spawnLocked() { | |
| 260 | + | if r.stopped || r.running || !r.dirty || r.list == nil { | |
| 261 | + | return | |
| 262 | + | } | |
| 263 | + | r.running = true | |
| 264 | + | go r.run() | |
| 265 | + | } | |
| 266 | + | ||
| 267 | + | func (r *toolRefresher) run() { | |
| 268 | + | for { | |
| 269 | + | r.mu.Lock() | |
| 270 | + | if r.stopped || !r.dirty { | |
| 271 | + | r.running = false | |
| 272 | + | r.mu.Unlock() | |
| 273 | + | return | |
| 274 | + | } | |
| 275 | + | r.dirty = false | |
| 276 | + | list := r.list | |
| 277 | + | r.mu.Unlock() | |
| 278 | + | list(r.ctx) | |
| 279 | + | } | |
| 280 | + | } | |
| 281 | + | ||
| 190 | 282 | func filterTools(cfg ServerConfig, tools []Tool) []Tool { | |
| 191 | 283 | out := make([]Tool, 0, len(tools)) | |
| 192 | 284 | seen := map[string]bool{} | |
| @@ -294,7 +386,12 @@ func (m *Manager) Close() error { | |||
| 294 | 386 | s.mu.Lock() | |
| 295 | 387 | client := s.client | |
| 296 | 388 | s.client = nil | |
| 389 | + | refresher := s.refresher | |
| 390 | + | s.refresher = nil | |
| 297 | 391 | s.mu.Unlock() | |
| 392 | + | if refresher != nil { | |
| 393 | + | refresher.stop() | |
| 394 | + | } | |
| 298 | 395 | if client != nil { | |
| 299 | 396 | _ = client.Close() | |
| 300 | 397 | } | |
| @@ -304,7 +401,8 @@ func (m *Manager) Close() error { | |||
| 304 | 401 | return nil | |
| 305 | 402 | } | |
| 306 | 403 | ||
| 307 | - | // FormatStatus renders a short human-readable status table. | |
| 404 | + | // FormatStatus renders a short human-readable status table. Server-derived | |
| 405 | + | // text (names, errors) is stripped of terminal control sequences. | |
| 308 | 406 | func FormatStatus(statuses []ServerStatus) string { | |
| 309 | 407 | if len(statuses) == 0 { | |
| 310 | 408 | return "no MCP servers configured" | |
| @@ -316,12 +414,12 @@ func FormatStatus(statuses []ServerStatus) string { | |||
| 316 | 414 | if st.Connected { | |
| 317 | 415 | state = "connected" | |
| 318 | 416 | } | |
| 319 | - | fmt.Fprintf(&b, "%s (%s): %s, %d tools", st.Name, st.Transport, state, st.Tools) | |
| 417 | + | fmt.Fprintf(&b, "%s (%s): %s, %d tools", sanitizeText(st.Name, false), sanitizeText(st.Transport, false), state, st.Tools) | |
| 320 | 418 | if st.Resources { | |
| 321 | 419 | b.WriteString(", resources") | |
| 322 | 420 | } | |
| 323 | 421 | if st.Err != "" { | |
| 324 | - | b.WriteString("\n error: " + firstLine(st.Err)) | |
| 422 | + | b.WriteString("\n error: " + sanitizeText(firstLine(st.Err), false)) | |
| 325 | 423 | } | |
| 326 | 424 | b.WriteByte('\n') | |
| 327 | 425 | } | |
| @@ -32,6 +32,7 @@ type serverEntry struct { | |||||
| 32 | client *Client | 32 | client *Client | ||
| 33 | tools []Tool | 33 | tools []Tool | ||
| 34 | err error | 34 | err error | ||
| 35 | + | refresher *toolRefresher | |||
| 35 | } | 36 | } | ||
| 36 | 37 | ||||
| 37 | // Manager owns the MCP connections of one zut process. It is safe for | 38 | // Manager owns the MCP connections of one zut process. It is safe for | ||
| @@ -108,7 +109,12 @@ func (m *Manager) connect(ctx context.Context, s *serverEntry) (*Client, error) | |||||
| 108 | s.mu.Lock() | 109 | s.mu.Lock() | ||
| 109 | stale := s.client | 110 | stale := s.client | ||
| 110 | s.client = nil | 111 | s.client = nil | ||
| 112 | + | staleRefresher := s.refresher | |||
| 113 | + | s.refresher = nil | |||
| 111 | s.mu.Unlock() | 114 | s.mu.Unlock() | ||
| 115 | + | if staleRefresher != nil { | |||
| 116 | + | staleRefresher.stop() | |||
| 117 | + | } | |||
| 112 | if stale != nil { | 118 | if stale != nil { | ||
| 113 | _ = stale.Close() | 119 | _ = stale.Close() | ||
| 114 | } | 120 | } | ||
| @@ -120,7 +126,11 @@ func (m *Manager) connect(ctx context.Context, s *serverEntry) (*Client, error) | |||||
| 120 | } | 126 | } | ||
| 121 | startCtx, cancel := withTimeout(ctx, s.cfg.StartupTimeout()) | 127 | startCtx, cancel := withTimeout(ctx, s.cfg.StartupTimeout()) | ||
| 122 | defer cancel() | 128 | defer cancel() | ||
| 123 | - | client, err := Connect(startCtx, s.name, s.cfg, m.cwd, m.info, m.httpClient) | 129 | + | // The refresher exists before the handshake so a list_changed notification |
| 130 | + | // that arrives while the first listing is in flight is remembered, and | |||
| 131 | + | // the listing is repeated once the client is published. | |||
| 132 | + | refresher := newToolRefresher() | |||
| 133 | + | client, err := connectWith(startCtx, s.name, s.cfg, m.cwd, m.info, m.httpClient, refresher.notify) | |||
| 124 | var tools []Tool | 134 | var tools []Tool | ||
| 125 | if err == nil { | 135 | if err == nil { | ||
| 126 | tools, err = client.ListTools(startCtx) | 136 | tools, err = client.ListTools(startCtx) | ||
| @@ -129,6 +139,7 @@ func (m *Manager) connect(ctx context.Context, s *serverEntry) (*Client, error) | |||||
| 129 | } | 139 | } | ||
| 130 | } | 140 | } | ||
| 131 | if err != nil { | 141 | if err != nil { | ||
| 142 | + | refresher.stop() | |||
| 132 | if errors.Is(err, context.DeadlineExceeded) && ctx.Err() == nil { | 143 | if errors.Is(err, context.DeadlineExceeded) && ctx.Err() == nil { | ||
| 133 | err = fmt.Errorf("timed out after %s; raise startupTimeoutSec if the server is slow to start", s.cfg.StartupTimeout()) | 144 | err = fmt.Errorf("timed out after %s; raise startupTimeoutSec if the server is slow to start", s.cfg.StartupTimeout()) | ||
| 134 | } | 145 | } | ||
| @@ -138,12 +149,12 @@ func (m *Manager) connect(ctx context.Context, s *serverEntry) (*Client, error) | |||||
| 138 | s.mu.Unlock() | 149 | s.mu.Unlock() | ||
| 139 | return nil, err | 150 | return nil, err | ||
| 140 | } | 151 | } | ||
| 141 | - | client.SetToolsChanged(func() { m.refreshTools(s, client) }) | |||
| 142 | // Publish under m.mu so Close either sees this client or we see closed. | 152 | // Publish under m.mu so Close either sees this client or we see closed. | ||
| 143 | m.mu.Lock() | 153 | m.mu.Lock() | ||
| 144 | if m.closed { | 154 | if m.closed { | ||
| 145 | m.mu.Unlock() | 155 | m.mu.Unlock() | ||
| 146 | // Close ran during the handshake; do not leak the new process. | 156 | // Close ran during the handshake; do not leak the new process. | ||
| 157 | + | refresher.stop() | |||
| 147 | _ = client.Close() | 158 | _ = client.Close() | ||
| 148 | return nil, fmt.Errorf("mcp server %q: manager closed", s.name) | 159 | return nil, fmt.Errorf("mcp server %q: manager closed", s.name) | ||
| 149 | } | 160 | } | ||
| @@ -151,8 +162,18 @@ func (m *Manager) connect(ctx context.Context, s *serverEntry) (*Client, error) | |||||
| 151 | s.client = client | 162 | s.client = client | ||
| 152 | s.tools = filterTools(s.cfg, tools) | 163 | s.tools = filterTools(s.cfg, tools) | ||
| 153 | s.err = nil | 164 | s.err = nil | ||
| 165 | + | old := s.refresher | |||
| 166 | + | s.refresher = refresher | |||
| 154 | s.mu.Unlock() | 167 | s.mu.Unlock() | ||
| 155 | m.mu.Unlock() | 168 | m.mu.Unlock() | ||
| 169 | + | if old != nil { | |||
| 170 | + | old.stop() | |||
| 171 | + | } | |||
| 172 | + | go func() { | |||
| 173 | + | <-client.Done() | |||
| 174 | + | refresher.stop() | |||
| 175 | + | }() | |||
| 176 | + | refresher.start(func(ctx context.Context) { m.refreshTools(ctx, s, client) }) | |||
| 156 | return client, nil | 177 | return client, nil | ||
| 157 | } | 178 | } | ||
| 158 | 179 | ||||
| @@ -172,9 +193,11 @@ func (s *serverEntry) live() *Client { | |||||
| 172 | } | 193 | } | ||
| 173 | 194 | ||||
| 174 | // refreshTools re-lists tools after notifications/tools/list_changed. New | 195 | // refreshTools re-lists tools after notifications/tools/list_changed. New | ||
| 175 | - | // definitions reach the model the next time zut builds a tool registry. | 196 | + | // definitions reach the model the next time zut builds a tool registry. It is |
| 176 | - | func (m *Manager) refreshTools(s *serverEntry, client *Client) { | 197 | + | // only called from the server's refresher, which serializes and coalesces |
| 177 | - | ctx, cancel := withTimeout(context.Background(), s.cfg.StartupTimeout()) | 198 | + | // calls. |
| 199 | + | func (m *Manager) refreshTools(ctx context.Context, s *serverEntry, client *Client) { | |||
| 200 | + | ctx, cancel := withTimeout(ctx, s.cfg.StartupTimeout()) | |||
| 178 | defer cancel() | 201 | defer cancel() | ||
| 179 | tools, err := client.ListTools(ctx) | 202 | tools, err := client.ListTools(ctx) | ||
| 180 | if err != nil { | 203 | if err != nil { | ||
| @@ -187,6 +210,75 @@ func (m *Manager) refreshTools(s *serverEntry, client *Client) { | |||||
| 187 | s.mu.Unlock() | 210 | s.mu.Unlock() | ||
| 188 | } | 211 | } | ||
| 189 | 212 | ||||
| 213 | + | // toolRefresher coalesces list_changed notifications for one connection into | |||
| 214 | + | // serialized re-listings. Any number of notifications that arrive while a | |||
| 215 | + | // listing runs cause exactly one more listing afterwards. It does nothing | |||
| 216 | + | // until start installs the listing function, so notifications that arrive | |||
| 217 | + | // during the initial handshake are remembered rather than raced. | |||
| 218 | + | type toolRefresher struct { | |||
| 219 | + | mu sync.Mutex | |||
| 220 | + | dirty bool | |||
| 221 | + | running bool | |||
| 222 | + | stopped bool | |||
| 223 | + | list func(ctx context.Context) | |||
| 224 | + | ctx context.Context | |||
| 225 | + | cancel context.CancelFunc | |||
| 226 | + | } | |||
| 227 | + | ||||
| 228 | + | func newToolRefresher() *toolRefresher { | |||
| 229 | + | ctx, cancel := context.WithCancel(context.Background()) | |||
| 230 | + | return &toolRefresher{ctx: ctx, cancel: cancel} | |||
| 231 | + | } | |||
| 232 | + | ||||
| 233 | + | // notify records that the server's tool list changed. | |||
| 234 | + | func (r *toolRefresher) notify() { | |||
| 235 | + | r.mu.Lock() | |||
| 236 | + | r.dirty = true | |||
| 237 | + | r.spawnLocked() | |||
| 238 | + | r.mu.Unlock() | |||
| 239 | + | } | |||
| 240 | + | ||||
| 241 | + | // start enables listing with fn and runs it once if a notification arrived | |||
| 242 | + | // earlier. | |||
| 243 | + | func (r *toolRefresher) start(fn func(ctx context.Context)) { | |||
| 244 | + | r.mu.Lock() | |||
| 245 | + | r.list = fn | |||
| 246 | + | r.spawnLocked() | |||
| 247 | + | r.mu.Unlock() | |||
| 248 | + | } | |||
| 249 | + | ||||
| 250 | + | // stop cancels any running listing and ignores later notifications. | |||
| 251 | + | func (r *toolRefresher) stop() { | |||
| 252 | + | r.mu.Lock() | |||
| 253 | + | r.stopped = true | |||
| 254 | + | r.dirty = false | |||
| 255 | + | r.mu.Unlock() | |||
| 256 | + | r.cancel() | |||
| 257 | + | } | |||
| 258 | + | ||||
| 259 | + | func (r *toolRefresher) spawnLocked() { | |||
| 260 | + | if r.stopped || r.running || !r.dirty || r.list == nil { | |||
| 261 | + | return | |||
| 262 | + | } | |||
| 263 | + | r.running = true | |||
| 264 | + | go r.run() | |||
| 265 | + | } | |||
| 266 | + | ||||
| 267 | + | func (r *toolRefresher) run() { | |||
| 268 | + | for { | |||
| 269 | + | r.mu.Lock() | |||
| 270 | + | if r.stopped || !r.dirty { | |||
| 271 | + | r.running = false | |||
| 272 | + | r.mu.Unlock() | |||
| 273 | + | return | |||
| 274 | + | } | |||
| 275 | + | r.dirty = false | |||
| 276 | + | list := r.list | |||
| 277 | + | r.mu.Unlock() | |||
| 278 | + | list(r.ctx) | |||
| 279 | + | } | |||
| 280 | + | } | |||
| 281 | + | ||||
| 190 | func filterTools(cfg ServerConfig, tools []Tool) []Tool { | 282 | func filterTools(cfg ServerConfig, tools []Tool) []Tool { | ||
| 191 | out := make([]Tool, 0, len(tools)) | 283 | out := make([]Tool, 0, len(tools)) | ||
| 192 | seen := map[string]bool{} | 284 | seen := map[string]bool{} | ||
| @@ -294,7 +386,12 @@ func (m *Manager) Close() error { | |||||
| 294 | s.mu.Lock() | 386 | s.mu.Lock() | ||
| 295 | client := s.client | 387 | client := s.client | ||
| 296 | s.client = nil | 388 | s.client = nil | ||
| 389 | + | refresher := s.refresher | |||
| 390 | + | s.refresher = nil | |||
| 297 | s.mu.Unlock() | 391 | s.mu.Unlock() | ||
| 392 | + | if refresher != nil { | |||
| 393 | + | refresher.stop() | |||
| 394 | + | } | |||
| 298 | if client != nil { | 395 | if client != nil { | ||
| 299 | _ = client.Close() | 396 | _ = client.Close() | ||
| 300 | } | 397 | } | ||
| @@ -304,7 +401,8 @@ func (m *Manager) Close() error { | |||||
| 304 | return nil | 401 | return nil | ||
| 305 | } | 402 | } | ||
| 306 | 403 | ||||
| 307 | - | // FormatStatus renders a short human-readable status table. | 404 | + | // FormatStatus renders a short human-readable status table. Server-derived |
| 405 | + | // text (names, errors) is stripped of terminal control sequences. | |||
| 308 | func FormatStatus(statuses []ServerStatus) string { | 406 | func FormatStatus(statuses []ServerStatus) string { | ||
| 309 | if len(statuses) == 0 { | 407 | if len(statuses) == 0 { | ||
| 310 | return "no MCP servers configured" | 408 | return "no MCP servers configured" | ||
| @@ -316,12 +414,12 @@ func FormatStatus(statuses []ServerStatus) string { | |||||
| 316 | if st.Connected { | 414 | if st.Connected { | ||
| 317 | state = "connected" | 415 | state = "connected" | ||
| 318 | } | 416 | } | ||
| 319 | - | fmt.Fprintf(&b, "%s (%s): %s, %d tools", st.Name, st.Transport, state, st.Tools) | 417 | + | fmt.Fprintf(&b, "%s (%s): %s, %d tools", sanitizeText(st.Name, false), sanitizeText(st.Transport, false), state, st.Tools) |
| 320 | if st.Resources { | 418 | if st.Resources { | ||
| 321 | b.WriteString(", resources") | 419 | b.WriteString(", resources") | ||
| 322 | } | 420 | } | ||
| 323 | if st.Err != "" { | 421 | if st.Err != "" { | ||
| 324 | - | b.WriteString("\n error: " + firstLine(st.Err)) | 422 | + | b.WriteString("\n error: " + sanitizeText(firstLine(st.Err), false)) |
| 325 | } | 423 | } | ||
| 326 | b.WriteByte('\n') | 424 | b.WriteByte('\n') | ||
| 327 | } | 425 | } | ||
M packages/agent/mcp/manager_test.go +7 -7
| @@ -98,7 +98,7 @@ func TestManagerStdioToolsAndDescriptions(t *testing.T) { | |||
| 98 | 98 | t.Fatal("manager should not be empty") | |
| 99 | 99 | } | |
| 100 | 100 | st := m.Status() | |
| 101 | - | if len(st) != 1 || !st[0].Connected || st[0].Tools != 14 || !st[0].Resources || st[0].Transport != "stdio" { | |
| 101 | + | if len(st) != 1 || !st[0].Connected || st[0].Tools != fakeToolCount || !st[0].Resources || st[0].Transport != "stdio" { | |
| 102 | 102 | t.Fatalf("status = %+v", st) | |
| 103 | 103 | } | |
| 104 | 104 | if st[0].Info.ProtocolVersion != "2025-06-18" || !st[0].Info.ToolsChanged || st[0].Info.Name != "fake" { | |
| @@ -106,8 +106,8 @@ func TestManagerStdioToolsAndDescriptions(t *testing.T) { | |||
| 106 | 106 | } | |
| 107 | 107 | ||
| 108 | 108 | tools := m.Tools() | |
| 109 | - | // 14 server tools across two pages + 2 resource helpers. | |
| 110 | - | if len(tools) != 16 { | |
| 109 | + | // All server tools across two pages + 2 resource helpers. | |
| 110 | + | if len(tools) != fakeToolCount+2 { | |
| 111 | 111 | t.Fatalf("got %d tools: %v", len(tools), toolNames(m)) | |
| 112 | 112 | } | |
| 113 | 113 | seen := map[string]bool{} | |
| @@ -160,14 +160,14 @@ func TestManagerStdioToolsAndDescriptions(t *testing.T) { | |||
| 160 | 160 | } | |
| 161 | 161 | ||
| 162 | 162 | text := FormatStatus(m.Status()) | |
| 163 | - | if !strings.Contains(text, "fake (stdio): connected, 14 tools, resources") { | |
| 163 | + | if !strings.Contains(text, "fake (stdio): connected, 17 tools, resources") { | |
| 164 | 164 | t.Errorf("FormatStatus = %q", text) | |
| 165 | 165 | } | |
| 166 | 166 | if FormatStatus(nil) != "no MCP servers configured" { | |
| 167 | 167 | t.Error("empty FormatStatus") | |
| 168 | 168 | } | |
| 169 | 169 | add := m.SystemPromptAddendum() | |
| 170 | - | for _, want := range []string{"[mcp_servers]", "fake: 14 tools", "mcp__fake__<tool>", "Be nice"} { | |
| 170 | + | for _, want := range []string{"[mcp_servers]", "fake: 17 tools", "mcp__fake__<tool>", "Be nice"} { | |
| 171 | 171 | if !strings.Contains(add, want) { | |
| 172 | 172 | t.Errorf("addendum missing %q:\n%s", want, add) | |
| 173 | 173 | } | |
| @@ -325,7 +325,7 @@ func TestManagerToolFiltering(t *testing.T) { | |||
| 325 | 325 | t.Errorf("%s should be filtered", n) | |
| 326 | 326 | } | |
| 327 | 327 | } | |
| 328 | - | if !hasTool(m, "mcp__fake__image") || len(m.ServerTools("fake")) != 11 { | |
| 328 | + | if !hasTool(m, "mcp__fake__image") || len(m.ServerTools("fake")) != fakeToolCount-3 { | |
| 329 | 329 | t.Errorf("remaining tools wrong: %v", toolNames(m)) | |
| 330 | 330 | } | |
| 331 | 331 | }) | |
| @@ -621,7 +621,7 @@ func TestManagerResourceToolsOnlyWhenServerHasResources(t *testing.T) { | |||
| 621 | 621 | if hasTool(m, ListResourcesToolName) || hasTool(m, ReadResourceToolName) { | |
| 622 | 622 | t.Errorf("resource tools should be absent: %v", toolNames(m)) | |
| 623 | 623 | } | |
| 624 | - | if len(m.Tools()) != 14 { | |
| 624 | + | if len(m.Tools()) != fakeToolCount { | |
| 625 | 625 | t.Errorf("tools = %v", toolNames(m)) | |
| 626 | 626 | } | |
| 627 | 627 | if strings.Contains(FormatStatus(m.Status()), "resources") { | |
| @@ -98,7 +98,7 @@ func TestManagerStdioToolsAndDescriptions(t *testing.T) { | |||||
| 98 | t.Fatal("manager should not be empty") | 98 | t.Fatal("manager should not be empty") | ||
| 99 | } | 99 | } | ||
| 100 | st := m.Status() | 100 | st := m.Status() | ||
| 101 | - | if len(st) != 1 || !st[0].Connected || st[0].Tools != 14 || !st[0].Resources || st[0].Transport != "stdio" { | 101 | + | if len(st) != 1 || !st[0].Connected || st[0].Tools != fakeToolCount || !st[0].Resources || st[0].Transport != "stdio" { |
| 102 | t.Fatalf("status = %+v", st) | 102 | t.Fatalf("status = %+v", st) | ||
| 103 | } | 103 | } | ||
| 104 | if st[0].Info.ProtocolVersion != "2025-06-18" || !st[0].Info.ToolsChanged || st[0].Info.Name != "fake" { | 104 | if st[0].Info.ProtocolVersion != "2025-06-18" || !st[0].Info.ToolsChanged || st[0].Info.Name != "fake" { | ||
| @@ -106,8 +106,8 @@ func TestManagerStdioToolsAndDescriptions(t *testing.T) { | |||||
| 106 | } | 106 | } | ||
| 107 | 107 | ||||
| 108 | tools := m.Tools() | 108 | tools := m.Tools() | ||
| 109 | - | // 14 server tools across two pages + 2 resource helpers. | 109 | + | // All server tools across two pages + 2 resource helpers. |
| 110 | - | if len(tools) != 16 { | 110 | + | if len(tools) != fakeToolCount+2 { |
| 111 | t.Fatalf("got %d tools: %v", len(tools), toolNames(m)) | 111 | t.Fatalf("got %d tools: %v", len(tools), toolNames(m)) | ||
| 112 | } | 112 | } | ||
| 113 | seen := map[string]bool{} | 113 | seen := map[string]bool{} | ||
| @@ -160,14 +160,14 @@ func TestManagerStdioToolsAndDescriptions(t *testing.T) { | |||||
| 160 | } | 160 | } | ||
| 161 | 161 | ||||
| 162 | text := FormatStatus(m.Status()) | 162 | text := FormatStatus(m.Status()) | ||
| 163 | - | if !strings.Contains(text, "fake (stdio): connected, 14 tools, resources") { | 163 | + | if !strings.Contains(text, "fake (stdio): connected, 17 tools, resources") { |
| 164 | t.Errorf("FormatStatus = %q", text) | 164 | t.Errorf("FormatStatus = %q", text) | ||
| 165 | } | 165 | } | ||
| 166 | if FormatStatus(nil) != "no MCP servers configured" { | 166 | if FormatStatus(nil) != "no MCP servers configured" { | ||
| 167 | t.Error("empty FormatStatus") | 167 | t.Error("empty FormatStatus") | ||
| 168 | } | 168 | } | ||
| 169 | add := m.SystemPromptAddendum() | 169 | add := m.SystemPromptAddendum() | ||
| 170 | - | for _, want := range []string{"[mcp_servers]", "fake: 14 tools", "mcp__fake__<tool>", "Be nice"} { | 170 | + | for _, want := range []string{"[mcp_servers]", "fake: 17 tools", "mcp__fake__<tool>", "Be nice"} { |
| 171 | if !strings.Contains(add, want) { | 171 | if !strings.Contains(add, want) { | ||
| 172 | t.Errorf("addendum missing %q:\n%s", want, add) | 172 | t.Errorf("addendum missing %q:\n%s", want, add) | ||
| 173 | } | 173 | } | ||
| @@ -325,7 +325,7 @@ func TestManagerToolFiltering(t *testing.T) { | |||||
| 325 | t.Errorf("%s should be filtered", n) | 325 | t.Errorf("%s should be filtered", n) | ||
| 326 | } | 326 | } | ||
| 327 | } | 327 | } | ||
| 328 | - | if !hasTool(m, "mcp__fake__image") || len(m.ServerTools("fake")) != 11 { | 328 | + | if !hasTool(m, "mcp__fake__image") || len(m.ServerTools("fake")) != fakeToolCount-3 { |
| 329 | t.Errorf("remaining tools wrong: %v", toolNames(m)) | 329 | t.Errorf("remaining tools wrong: %v", toolNames(m)) | ||
| 330 | } | 330 | } | ||
| 331 | }) | 331 | }) | ||
| @@ -621,7 +621,7 @@ func TestManagerResourceToolsOnlyWhenServerHasResources(t *testing.T) { | |||||
| 621 | if hasTool(m, ListResourcesToolName) || hasTool(m, ReadResourceToolName) { | 621 | if hasTool(m, ListResourcesToolName) || hasTool(m, ReadResourceToolName) { | ||
| 622 | t.Errorf("resource tools should be absent: %v", toolNames(m)) | 622 | t.Errorf("resource tools should be absent: %v", toolNames(m)) | ||
| 623 | } | 623 | } | ||
| 624 | - | if len(m.Tools()) != 14 { | 624 | + | if len(m.Tools()) != fakeToolCount { |
| 625 | t.Errorf("tools = %v", toolNames(m)) | 625 | t.Errorf("tools = %v", toolNames(m)) | ||
| 626 | } | 626 | } | ||
| 627 | if strings.Contains(FormatStatus(m.Status()), "resources") { | 627 | if strings.Contains(FormatStatus(m.Status()), "resources") { | ||
M packages/agent/mcp/proc_unix.go +28 -4
| @@ -13,15 +13,39 @@ func isolateProcess(cmd *exec.Cmd) { | |||
| 13 | 13 | cmd.SysProcAttr = &syscall.SysProcAttr{Setpgid: true} | |
| 14 | 14 | } | |
| 15 | 15 | ||
| 16 | - | func terminateProcess(cmd *exec.Cmd, force bool) { | |
| 17 | - | if cmd == nil || cmd.Process == nil { | |
| 16 | + | // processGuard tracks everything a stdio server spawned: the process group on | |
| 17 | + | // Unix, a Job Object on Windows. | |
| 18 | + | type processGuard struct { | |
| 19 | + | cmd *exec.Cmd | |
| 20 | + | } | |
| 21 | + | ||
| 22 | + | // newProcessGuard is called right after cmd.Start. | |
| 23 | + | func newProcessGuard(cmd *exec.Cmd) *processGuard { return &processGuard{cmd: cmd} } | |
| 24 | + | ||
| 25 | + | // terminate signals the whole process group, including descendants that | |
| 26 | + | // outlive the direct child. | |
| 27 | + | func (g *processGuard) terminate(force bool) { | |
| 28 | + | if g == nil || g.cmd == nil || g.cmd.Process == nil { | |
| 18 | 29 | return | |
| 19 | 30 | } | |
| 20 | 31 | sig := syscall.SIGTERM | |
| 21 | 32 | if force { | |
| 22 | 33 | sig = syscall.SIGKILL | |
| 23 | 34 | } | |
| 24 | - | if err := syscall.Kill(-cmd.Process.Pid, sig); err != nil { | |
| 25 | - | _ = cmd.Process.Signal(sig) | |
| 35 | + | err := syscall.Kill(-g.cmd.Process.Pid, sig) | |
| 36 | + | if err != nil && err != syscall.ESRCH { | |
| 37 | + | _ = g.cmd.Process.Signal(sig) | |
| 26 | 38 | } | |
| 27 | 39 | } | |
| 40 | + | ||
| 41 | + | // alive reports whether any member of the process group still exists. | |
| 42 | + | func (g *processGuard) alive() bool { | |
| 43 | + | if g == nil || g.cmd == nil || g.cmd.Process == nil { | |
| 44 | + | return false | |
| 45 | + | } | |
| 46 | + | err := syscall.Kill(-g.cmd.Process.Pid, 0) | |
| 47 | + | return err == nil || err == syscall.EPERM | |
| 48 | + | } | |
| 49 | + | ||
| 50 | + | // release frees resources held for the group. Nothing is held on Unix. | |
| 51 | + | func (g *processGuard) release() {} | |
| @@ -13,15 +13,39 @@ func isolateProcess(cmd *exec.Cmd) { | |||||
| 13 | cmd.SysProcAttr = &syscall.SysProcAttr{Setpgid: true} | 13 | cmd.SysProcAttr = &syscall.SysProcAttr{Setpgid: true} | ||
| 14 | } | 14 | } | ||
| 15 | 15 | ||||
| 16 | - | func terminateProcess(cmd *exec.Cmd, force bool) { | 16 | + | // processGuard tracks everything a stdio server spawned: the process group on |
| 17 | - | if cmd == nil || cmd.Process == nil { | 17 | + | // Unix, a Job Object on Windows. |
| 18 | + | type processGuard struct { | |||
| 19 | + | cmd *exec.Cmd | |||
| 20 | + | } | |||
| 21 | + | ||||
| 22 | + | // newProcessGuard is called right after cmd.Start. | |||
| 23 | + | func newProcessGuard(cmd *exec.Cmd) *processGuard { return &processGuard{cmd: cmd} } | |||
| 24 | + | ||||
| 25 | + | // terminate signals the whole process group, including descendants that | |||
| 26 | + | // outlive the direct child. | |||
| 27 | + | func (g *processGuard) terminate(force bool) { | |||
| 28 | + | if g == nil || g.cmd == nil || g.cmd.Process == nil { | |||
| 18 | return | 29 | return | ||
| 19 | } | 30 | } | ||
| 20 | sig := syscall.SIGTERM | 31 | sig := syscall.SIGTERM | ||
| 21 | if force { | 32 | if force { | ||
| 22 | sig = syscall.SIGKILL | 33 | sig = syscall.SIGKILL | ||
| 23 | } | 34 | } | ||
| 24 | - | if err := syscall.Kill(-cmd.Process.Pid, sig); err != nil { | 35 | + | err := syscall.Kill(-g.cmd.Process.Pid, sig) |
| 25 | - | _ = cmd.Process.Signal(sig) | 36 | + | if err != nil && err != syscall.ESRCH { |
| 37 | + | _ = g.cmd.Process.Signal(sig) | |||
| 26 | } | 38 | } | ||
| 27 | } | 39 | } | ||
| 40 | + | ||||
| 41 | + | // alive reports whether any member of the process group still exists. | |||
| 42 | + | func (g *processGuard) alive() bool { | |||
| 43 | + | if g == nil || g.cmd == nil || g.cmd.Process == nil { | |||
| 44 | + | return false | |||
| 45 | + | } | |||
| 46 | + | err := syscall.Kill(-g.cmd.Process.Pid, 0) | |||
| 47 | + | return err == nil || err == syscall.EPERM | |||
| 48 | + | } | |||
| 49 | + | ||||
| 50 | + | // release frees resources held for the group. Nothing is held on Unix. | |||
| 51 | + | func (g *processGuard) release() {} | |||
M packages/agent/mcp/proc_unix_test.go +23 -3
| @@ -2,9 +2,29 @@ | |||
| 2 | 2 | ||
| 3 | 3 | package mcp | |
| 4 | 4 | ||
| 5 | - | import "syscall" | |
| 5 | + | import ( | |
| 6 | + | "os" | |
| 7 | + | "strconv" | |
| 8 | + | "strings" | |
| 9 | + | "syscall" | |
| 10 | + | ) | |
| 6 | 11 | ||
| 7 | - | // processAlive reports whether pid still exists. | |
| 8 | - | func processAlive(pid int) bool { return syscall.Kill(pid, 0) == nil } | |
| 12 | + | // processAlive reports whether pid still runs. A zombie that nobody reaped | |
| 13 | + | // (common when the test runs as PID 1's child in a container) counts as dead. | |
| 14 | + | func processAlive(pid int) bool { | |
| 15 | + | if syscall.Kill(pid, 0) != nil { | |
| 16 | + | return false | |
| 17 | + | } | |
| 18 | + | if raw, err := os.ReadFile("/proc/" + strconv.Itoa(pid) + "/stat"); err == nil { | |
| 19 | + | // Format: pid (comm) S ...; comm may contain spaces and parens. | |
| 20 | + | if i := strings.LastIndexByte(string(raw), ')'); i >= 0 && i+2 < len(raw) && raw[i+2] == 'Z' { | |
| 21 | + | return false | |
| 22 | + | } | |
| 23 | + | } | |
| 24 | + | return true | |
| 25 | + | } | |
| 9 | 26 | ||
| 10 | 27 | const canCheckProcess = true | |
| 28 | + | ||
| 29 | + | // killPID is a best-effort cleanup for processes a test started. | |
| 30 | + | func killPID(pid int) { _ = syscall.Kill(pid, syscall.SIGKILL) } | |
| @@ -2,9 +2,29 @@ | |||||
| 2 | 2 | ||||
| 3 | package mcp | 3 | package mcp | ||
| 4 | 4 | ||||
| 5 | - | import "syscall" | 5 | + | import ( |
| 6 | + | "os" | |||
| 7 | + | "strconv" | |||
| 8 | + | "strings" | |||
| 9 | + | "syscall" | |||
| 10 | + | ) | |||
| 6 | 11 | ||||
| 7 | - | // processAlive reports whether pid still exists. | 12 | + | // processAlive reports whether pid still runs. A zombie that nobody reaped |
| 8 | - | func processAlive(pid int) bool { return syscall.Kill(pid, 0) == nil } | 13 | + | // (common when the test runs as PID 1's child in a container) counts as dead. |
| 14 | + | func processAlive(pid int) bool { | |||
| 15 | + | if syscall.Kill(pid, 0) != nil { | |||
| 16 | + | return false | |||
| 17 | + | } | |||
| 18 | + | if raw, err := os.ReadFile("/proc/" + strconv.Itoa(pid) + "/stat"); err == nil { | |||
| 19 | + | // Format: pid (comm) S ...; comm may contain spaces and parens. | |||
| 20 | + | if i := strings.LastIndexByte(string(raw), ')'); i >= 0 && i+2 < len(raw) && raw[i+2] == 'Z' { | |||
| 21 | + | return false | |||
| 22 | + | } | |||
| 23 | + | } | |||
| 24 | + | return true | |||
| 25 | + | } | |||
| 9 | 26 | ||||
| 10 | const canCheckProcess = true | 27 | const canCheckProcess = true | ||
| 28 | + | ||||
| 29 | + | // killPID is a best-effort cleanup for processes a test started. | |||
| 30 | + | func killPID(pid int) { _ = syscall.Kill(pid, syscall.SIGKILL) } | |||
M packages/agent/mcp/proc_windows.go +83 -6
| @@ -2,14 +2,91 @@ | |||
| 2 | 2 | ||
| 3 | 3 | package mcp | |
| 4 | 4 | ||
| 5 | - | import "os/exec" | |
| 5 | + | import ( | |
| 6 | + | "os/exec" | |
| 7 | + | "sync" | |
| 8 | + | "unsafe" | |
| 9 | + | ||
| 10 | + | "golang.org/x/sys/windows" | |
| 11 | + | ) | |
| 6 | 12 | ||
| 7 | 13 | func isolateProcess(_ *exec.Cmd) {} | |
| 8 | 14 | ||
| 9 | - | // terminateProcess kills the server process. Windows has no process-group | |
| 10 | - | // signal equivalent here, so both phases kill the direct child. | |
| 11 | - | func terminateProcess(cmd *exec.Cmd, _ bool) { | |
| 12 | - | if cmd != nil && cmd.Process != nil { | |
| 13 | - | _ = cmd.Process.Kill() | |
| 15 | + | // processGuard assigns the server to a Job Object with | |
| 16 | + | // JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE, so descendants die with the server (and | |
| 17 | + | // with zut, should zut itself be killed). | |
| 18 | + | // | |
| 19 | + | // Assignment happens right after Start; a server that spawns children in the | |
| 20 | + | // first instants of its life can escape the job. If the job cannot be created | |
| 21 | + | // or assigned, the guard falls back to killing the direct child. | |
| 22 | + | type processGuard struct { | |
| 23 | + | cmd *exec.Cmd | |
| 24 | + | ||
| 25 | + | mu sync.Mutex | |
| 26 | + | job windows.Handle | |
| 27 | + | } | |
| 28 | + | ||
| 29 | + | func newProcessGuard(cmd *exec.Cmd) *processGuard { | |
| 30 | + | g := &processGuard{cmd: cmd} | |
| 31 | + | if cmd == nil || cmd.Process == nil { | |
| 32 | + | return g | |
| 33 | + | } | |
| 34 | + | job, err := windows.CreateJobObject(nil, nil) | |
| 35 | + | if err != nil { | |
| 36 | + | return g | |
| 37 | + | } | |
| 38 | + | var info windows.JOBOBJECT_EXTENDED_LIMIT_INFORMATION | |
| 39 | + | info.BasicLimitInformation.LimitFlags = windows.JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE | |
| 40 | + | if _, err := windows.SetInformationJobObject(job, windows.JobObjectExtendedLimitInformation, | |
| 41 | + | uintptr(unsafe.Pointer(&info)), uint32(unsafe.Sizeof(info))); err != nil { | |
| 42 | + | _ = windows.CloseHandle(job) | |
| 43 | + | return g | |
| 44 | + | } | |
| 45 | + | proc, err := windows.OpenProcess(windows.PROCESS_SET_QUOTA|windows.PROCESS_TERMINATE, false, uint32(cmd.Process.Pid)) | |
| 46 | + | if err != nil { | |
| 47 | + | _ = windows.CloseHandle(job) | |
| 48 | + | return g | |
| 49 | + | } | |
| 50 | + | defer windows.CloseHandle(proc) | |
| 51 | + | if err := windows.AssignProcessToJobObject(job, proc); err != nil { | |
| 52 | + | _ = windows.CloseHandle(job) | |
| 53 | + | return g | |
| 54 | + | } | |
| 55 | + | g.job = job | |
| 56 | + | return g | |
| 57 | + | } | |
| 58 | + | ||
| 59 | + | // terminate kills the job (every process in it) or, without a job, the direct | |
| 60 | + | // child. Windows has no graceful group signal, so force is ignored. | |
| 61 | + | func (g *processGuard) terminate(_ bool) { | |
| 62 | + | if g == nil { | |
| 63 | + | return | |
| 64 | + | } | |
| 65 | + | g.mu.Lock() | |
| 66 | + | job := g.job | |
| 67 | + | g.mu.Unlock() | |
| 68 | + | if job != 0 { | |
| 69 | + | _ = windows.TerminateJobObject(job, 1) | |
| 70 | + | return | |
| 71 | + | } | |
| 72 | + | if g.cmd != nil && g.cmd.Process != nil { | |
| 73 | + | _ = g.cmd.Process.Kill() | |
| 74 | + | } | |
| 75 | + | } | |
| 76 | + | ||
| 77 | + | // alive is always false: terminate is already forceful. | |
| 78 | + | func (g *processGuard) alive() bool { return false } | |
| 79 | + | ||
| 80 | + | // release closes the job handle, which kills any survivor. | |
| 81 | + | func (g *processGuard) release() { | |
| 82 | + | if g == nil { | |
| 83 | + | return | |
| 84 | + | } | |
| 85 | + | g.mu.Lock() | |
| 86 | + | job := g.job | |
| 87 | + | g.job = 0 | |
| 88 | + | g.mu.Unlock() | |
| 89 | + | if job != 0 { | |
| 90 | + | _ = windows.CloseHandle(job) | |
| 14 | 91 | } | |
| 15 | 92 | } | |
| @@ -2,14 +2,91 @@ | |||||
| 2 | 2 | ||||
| 3 | package mcp | 3 | package mcp | ||
| 4 | 4 | ||||
| 5 | - | import "os/exec" | 5 | + | import ( |
| 6 | + | "os/exec" | |||
| 7 | + | "sync" | |||
| 8 | + | "unsafe" | |||
| 9 | + | ||||
| 10 | + | "golang.org/x/sys/windows" | |||
| 11 | + | ) | |||
| 6 | 12 | ||||
| 7 | func isolateProcess(_ *exec.Cmd) {} | 13 | func isolateProcess(_ *exec.Cmd) {} | ||
| 8 | 14 | ||||
| 9 | - | // terminateProcess kills the server process. Windows has no process-group | 15 | + | // processGuard assigns the server to a Job Object with |
| 10 | - | // signal equivalent here, so both phases kill the direct child. | 16 | + | // JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE, so descendants die with the server (and |
| 11 | - | func terminateProcess(cmd *exec.Cmd, _ bool) { | 17 | + | // with zut, should zut itself be killed). |
| 12 | - | if cmd != nil && cmd.Process != nil { | 18 | + | // |
| 13 | - | _ = cmd.Process.Kill() | 19 | + | // Assignment happens right after Start; a server that spawns children in the |
| 20 | + | // first instants of its life can escape the job. If the job cannot be created | |||
| 21 | + | // or assigned, the guard falls back to killing the direct child. | |||
| 22 | + | type processGuard struct { | |||
| 23 | + | cmd *exec.Cmd | |||
| 24 | + | ||||
| 25 | + | mu sync.Mutex | |||
| 26 | + | job windows.Handle | |||
| 27 | + | } | |||
| 28 | + | ||||
| 29 | + | func newProcessGuard(cmd *exec.Cmd) *processGuard { | |||
| 30 | + | g := &processGuard{cmd: cmd} | |||
| 31 | + | if cmd == nil || cmd.Process == nil { | |||
| 32 | + | return g | |||
| 33 | + | } | |||
| 34 | + | job, err := windows.CreateJobObject(nil, nil) | |||
| 35 | + | if err != nil { | |||
| 36 | + | return g | |||
| 37 | + | } | |||
| 38 | + | var info windows.JOBOBJECT_EXTENDED_LIMIT_INFORMATION | |||
| 39 | + | info.BasicLimitInformation.LimitFlags = windows.JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE | |||
| 40 | + | if _, err := windows.SetInformationJobObject(job, windows.JobObjectExtendedLimitInformation, | |||
| 41 | + | uintptr(unsafe.Pointer(&info)), uint32(unsafe.Sizeof(info))); err != nil { | |||
| 42 | + | _ = windows.CloseHandle(job) | |||
| 43 | + | return g | |||
| 44 | + | } | |||
| 45 | + | proc, err := windows.OpenProcess(windows.PROCESS_SET_QUOTA|windows.PROCESS_TERMINATE, false, uint32(cmd.Process.Pid)) | |||
| 46 | + | if err != nil { | |||
| 47 | + | _ = windows.CloseHandle(job) | |||
| 48 | + | return g | |||
| 49 | + | } | |||
| 50 | + | defer windows.CloseHandle(proc) | |||
| 51 | + | if err := windows.AssignProcessToJobObject(job, proc); err != nil { | |||
| 52 | + | _ = windows.CloseHandle(job) | |||
| 53 | + | return g | |||
| 54 | + | } | |||
| 55 | + | g.job = job | |||
| 56 | + | return g | |||
| 57 | + | } | |||
| 58 | + | ||||
| 59 | + | // terminate kills the job (every process in it) or, without a job, the direct | |||
| 60 | + | // child. Windows has no graceful group signal, so force is ignored. | |||
| 61 | + | func (g *processGuard) terminate(_ bool) { | |||
| 62 | + | if g == nil { | |||
| 63 | + | return | |||
| 64 | + | } | |||
| 65 | + | g.mu.Lock() | |||
| 66 | + | job := g.job | |||
| 67 | + | g.mu.Unlock() | |||
| 68 | + | if job != 0 { | |||
| 69 | + | _ = windows.TerminateJobObject(job, 1) | |||
| 70 | + | return | |||
| 71 | + | } | |||
| 72 | + | if g.cmd != nil && g.cmd.Process != nil { | |||
| 73 | + | _ = g.cmd.Process.Kill() | |||
| 74 | + | } | |||
| 75 | + | } | |||
| 76 | + | ||||
| 77 | + | // alive is always false: terminate is already forceful. | |||
| 78 | + | func (g *processGuard) alive() bool { return false } | |||
| 79 | + | ||||
| 80 | + | // release closes the job handle, which kills any survivor. | |||
| 81 | + | func (g *processGuard) release() { | |||
| 82 | + | if g == nil { | |||
| 83 | + | return | |||
| 84 | + | } | |||
| 85 | + | g.mu.Lock() | |||
| 86 | + | job := g.job | |||
| 87 | + | g.job = 0 | |||
| 88 | + | g.mu.Unlock() | |||
| 89 | + | if job != 0 { | |||
| 90 | + | _ = windows.CloseHandle(job) | |||
| 14 | } | 91 | } | ||
| 15 | } | 92 | } | ||
M packages/agent/mcp/proc_windows_test.go +2 -0
| @@ -5,3 +5,5 @@ package mcp | |||
| 5 | 5 | func processAlive(pid int) bool { return false } | |
| 6 | 6 | ||
| 7 | 7 | const canCheckProcess = false | |
| 8 | + | ||
| 9 | + | func killPID(pid int) {} | |
| @@ -5,3 +5,5 @@ package mcp | |||||
| 5 | func processAlive(pid int) bool { return false } | 5 | func processAlive(pid int) bool { return false } | ||
| 6 | 6 | ||||
| 7 | const canCheckProcess = false | 7 | const canCheckProcess = false | ||
| 8 | + | ||||
| 9 | + | func killPID(pid int) {} | |||
A packages/agent/mcp/sanitize.go +296 -0
| @@ -0,0 +1,296 @@ | |||
| 1 | + | package mcp | |
| 2 | + | ||
| 3 | + | import ( | |
| 4 | + | "os" | |
| 5 | + | "regexp" | |
| 6 | + | "sort" | |
| 7 | + | "strings" | |
| 8 | + | "unicode/utf8" | |
| 9 | + | ) | |
| 10 | + | ||
| 11 | + | // This file holds the diagnostics boundary: everything that originates from a | |
| 12 | + | // server (stderr, HTTP error bodies, JSON-RPC error messages) or from the HTTP | |
| 13 | + | // client (url.Error carries the full request URL) is scrubbed here before it | |
| 14 | + | // reaches an error value, the terminal, or a log. | |
| 15 | + | ||
| 16 | + | const ( | |
| 17 | + | redactedMarker = "[redacted]" | |
| 18 | + | // minSecretLen keeps one- and two-character configuration values (flags | |
| 19 | + | // such as "1") from turning every diagnostic into noise. | |
| 20 | + | minSecretLen = 4 | |
| 21 | + | ) | |
| 22 | + | ||
| 23 | + | // redactor removes configured secret values from server-provided text. | |
| 24 | + | // A nil *redactor still strips URLs' credentials and control characters. | |
| 25 | + | type redactor struct { | |
| 26 | + | secrets []string | |
| 27 | + | } | |
| 28 | + | ||
| 29 | + | // newRedactor collects the secret-bearing values of one server configuration: | |
| 30 | + | // env and header values, the bearer token, values passed through from the host | |
| 31 | + | // environment, and credentials embedded in the URL. | |
| 32 | + | func newRedactor(cfg ServerConfig) *redactor { | |
| 33 | + | set := map[string]bool{} | |
| 34 | + | add := func(v string) { | |
| 35 | + | v = strings.TrimSpace(v) | |
| 36 | + | if len(v) < minSecretLen { | |
| 37 | + | return | |
| 38 | + | } | |
| 39 | + | set[v] = true | |
| 40 | + | // "Bearer abc123" is also worth redacting as just "abc123". | |
| 41 | + | if strings.ContainsAny(v, " \t") { | |
| 42 | + | for _, field := range strings.Fields(v) { | |
| 43 | + | if len(field) >= 8 { | |
| 44 | + | set[field] = true | |
| 45 | + | } | |
| 46 | + | } | |
| 47 | + | } | |
| 48 | + | } | |
| 49 | + | for _, v := range cfg.Env { | |
| 50 | + | add(v) | |
| 51 | + | } | |
| 52 | + | for _, v := range cfg.Headers { | |
| 53 | + | add(v) | |
| 54 | + | } | |
| 55 | + | for _, name := range cfg.EnvHeaders { | |
| 56 | + | if v, ok := os.LookupEnv(name); ok { | |
| 57 | + | add(v) | |
| 58 | + | } | |
| 59 | + | } | |
| 60 | + | for _, name := range cfg.EnvVars { | |
| 61 | + | if v, ok := os.LookupEnv(name); ok { | |
| 62 | + | add(v) | |
| 63 | + | } | |
| 64 | + | } | |
| 65 | + | if cfg.BearerTokenEnvVar != "" { | |
| 66 | + | if v, ok := os.LookupEnv(cfg.BearerTokenEnvVar); ok { | |
| 67 | + | add(v) | |
| 68 | + | } | |
| 69 | + | } | |
| 70 | + | if cfg.URL != "" { | |
| 71 | + | for _, v := range urlSecrets(cfg.URL) { | |
| 72 | + | add(v) | |
| 73 | + | } | |
| 74 | + | } | |
| 75 | + | r := &redactor{} | |
| 76 | + | for v := range set { | |
| 77 | + | r.secrets = append(r.secrets, v) | |
| 78 | + | } | |
| 79 | + | // Longest first so a secret containing another is replaced whole. | |
| 80 | + | sort.Slice(r.secrets, func(i, j int) bool { | |
| 81 | + | if len(r.secrets[i]) != len(r.secrets[j]) { | |
| 82 | + | return len(r.secrets[i]) > len(r.secrets[j]) | |
| 83 | + | } | |
| 84 | + | return r.secrets[i] < r.secrets[j] | |
| 85 | + | }) | |
| 86 | + | return r | |
| 87 | + | } | |
| 88 | + | ||
| 89 | + | // urlSecrets returns the userinfo password/user and query values of raw. | |
| 90 | + | func urlSecrets(raw string) []string { | |
| 91 | + | rest := raw | |
| 92 | + | if i := strings.Index(rest, "://"); i >= 0 { | |
| 93 | + | rest = rest[i+3:] | |
| 94 | + | } | |
| 95 | + | var out []string | |
| 96 | + | authority := rest | |
| 97 | + | if i := strings.IndexAny(authority, "/?#"); i >= 0 { | |
| 98 | + | authority = authority[:i] | |
| 99 | + | } | |
| 100 | + | if i := strings.LastIndex(authority, "@"); i >= 0 { | |
| 101 | + | for _, part := range strings.SplitN(authority[:i], ":", 2) { | |
| 102 | + | out = append(out, part) | |
| 103 | + | } | |
| 104 | + | } | |
| 105 | + | if i := strings.Index(rest, "?"); i >= 0 { | |
| 106 | + | query := rest[i+1:] | |
| 107 | + | if j := strings.Index(query, "#"); j >= 0 { | |
| 108 | + | query = query[:j] | |
| 109 | + | } | |
| 110 | + | for _, pair := range strings.Split(query, "&") { | |
| 111 | + | if _, v, ok := strings.Cut(pair, "="); ok { | |
| 112 | + | out = append(out, v) | |
| 113 | + | } | |
| 114 | + | } | |
| 115 | + | } | |
| 116 | + | return out | |
| 117 | + | } | |
| 118 | + | ||
| 119 | + | // scrub redacts configured secrets, strips credentials and query strings from | |
| 120 | + | // URLs, and removes terminal control sequences. Newlines are kept. | |
| 121 | + | func (r *redactor) scrub(s string) string { | |
| 122 | + | if s == "" { | |
| 123 | + | return s | |
| 124 | + | } | |
| 125 | + | if r != nil { | |
| 126 | + | for _, secret := range r.secrets { | |
| 127 | + | s = strings.ReplaceAll(s, secret, redactedMarker) | |
| 128 | + | } | |
| 129 | + | } | |
| 130 | + | s = scrubURLs(s) | |
| 131 | + | return sanitizeText(s, true) | |
| 132 | + | } | |
| 133 | + | ||
| 134 | + | // scrubOneLine is scrub with newlines collapsed to spaces. | |
| 135 | + | func (r *redactor) scrubOneLine(s string) string { | |
| 136 | + | return sanitizeText(r.scrub(s), false) | |
| 137 | + | } | |
| 138 | + | ||
| 139 | + | var urlInText = regexp.MustCompile("(?i)\\b[a-z][a-z0-9+.-]*://[^\\s\"'<>`]+") | |
| 140 | + | ||
| 141 | + | // scrubURLs rewrites every URL in s to scheme://host/path, dropping userinfo, | |
| 142 | + | // query, and fragment. | |
| 143 | + | func scrubURLs(s string) string { | |
| 144 | + | if !strings.Contains(s, "://") { | |
| 145 | + | return s | |
| 146 | + | } | |
| 147 | + | return urlInText.ReplaceAllStringFunc(s, stripURL) | |
| 148 | + | } | |
| 149 | + | ||
| 150 | + | func stripURL(raw string) string { | |
| 151 | + | schemeEnd := strings.Index(raw, "://") | |
| 152 | + | if schemeEnd < 0 { | |
| 153 | + | return raw | |
| 154 | + | } | |
| 155 | + | scheme, rest := raw[:schemeEnd+3], raw[schemeEnd+3:] | |
| 156 | + | end := len(rest) | |
| 157 | + | if i := strings.IndexAny(rest, "/?#"); i >= 0 { | |
| 158 | + | end = i | |
| 159 | + | } | |
| 160 | + | authority, tail := rest[:end], rest[end:] | |
| 161 | + | if i := strings.LastIndex(authority, "@"); i >= 0 { | |
| 162 | + | authority = authority[i+1:] | |
| 163 | + | } | |
| 164 | + | if i := strings.IndexAny(tail, "?#"); i >= 0 { | |
| 165 | + | tail = tail[:i] | |
| 166 | + | } | |
| 167 | + | return scheme + authority + tail | |
| 168 | + | } | |
| 169 | + | ||
| 170 | + | // scrubbedError carries a sanitized message while keeping the original error | |
| 171 | + | // reachable through errors.Is/As. | |
| 172 | + | type scrubbedError struct { | |
| 173 | + | msg string | |
| 174 | + | err error | |
| 175 | + | } | |
| 176 | + | ||
| 177 | + | func (e *scrubbedError) Error() string { return e.msg } | |
| 178 | + | func (e *scrubbedError) Unwrap() error { return e.err } | |
| 179 | + | ||
| 180 | + | // scrubError returns err with a scrubbed message, or nil for nil. Context | |
| 181 | + | // errors are returned unchanged so callers keep comparing them directly. | |
| 182 | + | func (r *redactor) scrubError(err error) error { | |
| 183 | + | if err == nil { | |
| 184 | + | return nil | |
| 185 | + | } | |
| 186 | + | msg := r.scrub(err.Error()) | |
| 187 | + | if msg == err.Error() { | |
| 188 | + | return err | |
| 189 | + | } | |
| 190 | + | return &scrubbedError{msg: msg, err: err} | |
| 191 | + | } | |
| 192 | + | ||
| 193 | + | // sanitizeText removes terminal control sequences (CSI, OSC, DCS/APC/PM/SOS | |
| 194 | + | // strings, other escapes), C0/C1 control characters, and bidirectional | |
| 195 | + | // overrides. Tabs become spaces. Newlines are kept when keepNewlines is true | |
| 196 | + | // and become spaces otherwise; carriage returns are dropped. | |
| 197 | + | func sanitizeText(s string, keepNewlines bool) string { | |
| 198 | + | clean := true | |
| 199 | + | for i := 0; i < len(s); i++ { | |
| 200 | + | c := s[i] | |
| 201 | + | if c < 0x20 && c != '\n' || c == 0x7f || c >= 0x80 || (c == '\n' && !keepNewlines) { | |
| 202 | + | clean = false | |
| 203 | + | break | |
| 204 | + | } | |
| 205 | + | } | |
| 206 | + | if clean { | |
| 207 | + | return s | |
| 208 | + | } | |
| 209 | + | var b strings.Builder | |
| 210 | + | b.Grow(len(s)) | |
| 211 | + | for i := 0; i < len(s); { | |
| 212 | + | r, size := utf8.DecodeRuneInString(s[i:]) | |
| 213 | + | i += size | |
| 214 | + | switch { | |
| 215 | + | case r == 0x1b: | |
| 216 | + | i = skipEscape(s, i) | |
| 217 | + | case r == 0x9b: | |
| 218 | + | i = skipCSI(s, i) | |
| 219 | + | case r == 0x9d || r == 0x90 || r == 0x98 || r == 0x9e || r == 0x9f: | |
| 220 | + | i = skipControlString(s, i) | |
| 221 | + | case r == '\n': | |
| 222 | + | if keepNewlines { | |
| 223 | + | b.WriteByte('\n') | |
| 224 | + | } else { | |
| 225 | + | b.WriteByte(' ') | |
| 226 | + | } | |
| 227 | + | case r == '\t': | |
| 228 | + | b.WriteByte(' ') | |
| 229 | + | case r < 0x20 || r == 0x7f || (r >= 0x80 && r <= 0x9f): | |
| 230 | + | // dropped (includes \r) | |
| 231 | + | case (r >= 0x202a && r <= 0x202e) || (r >= 0x2066 && r <= 0x2069) || r == 0x200e || r == 0x200f: | |
| 232 | + | // bidirectional controls can reorder terminal text | |
| 233 | + | default: | |
| 234 | + | b.WriteRune(r) | |
| 235 | + | } | |
| 236 | + | } | |
| 237 | + | return b.String() | |
| 238 | + | } | |
| 239 | + | ||
| 240 | + | // skipEscape returns the index after the escape sequence whose ESC byte was | |
| 241 | + | // just consumed (i points at the byte after ESC). | |
| 242 | + | func skipEscape(s string, i int) int { | |
| 243 | + | if i >= len(s) { | |
| 244 | + | return i | |
| 245 | + | } | |
| 246 | + | switch c := s[i]; { | |
| 247 | + | case c == '[': | |
| 248 | + | return skipCSI(s, i+1) | |
| 249 | + | case c == ']' || c == 'P' || c == 'X' || c == '^' || c == '_': | |
| 250 | + | return skipControlString(s, i+1) | |
| 251 | + | case c >= 0x20 && c <= 0x2f: | |
| 252 | + | for i < len(s) && s[i] >= 0x20 && s[i] <= 0x2f { | |
| 253 | + | i++ | |
| 254 | + | } | |
| 255 | + | if i < len(s) { | |
| 256 | + | i++ // final byte | |
| 257 | + | } | |
| 258 | + | return i | |
| 259 | + | default: | |
| 260 | + | return i + 1 | |
| 261 | + | } | |
| 262 | + | } | |
| 263 | + | ||
| 264 | + | func skipCSI(s string, i int) int { | |
| 265 | + | for i < len(s) && s[i] >= 0x30 && s[i] <= 0x3f { | |
| 266 | + | i++ | |
| 267 | + | } | |
| 268 | + | for i < len(s) && s[i] >= 0x20 && s[i] <= 0x2f { | |
| 269 | + | i++ | |
| 270 | + | } | |
| 271 | + | if i < len(s) && s[i] >= 0x40 && s[i] <= 0x7e { | |
| 272 | + | i++ | |
| 273 | + | } | |
| 274 | + | return i | |
| 275 | + | } | |
| 276 | + | ||
| 277 | + | // skipControlString consumes an OSC/DCS/APC/PM/SOS payload up to and | |
| 278 | + | // including its terminator (BEL, ESC \, or C1 ST). An ESC that does not start | |
| 279 | + | // ST cancels the string without being consumed. | |
| 280 | + | func skipControlString(s string, i int) int { | |
| 281 | + | for i < len(s) { | |
| 282 | + | switch { | |
| 283 | + | case s[i] == 0x07: | |
| 284 | + | return i + 1 | |
| 285 | + | case s[i] == 0x1b: | |
| 286 | + | if i+1 < len(s) && s[i+1] == '\\' { | |
| 287 | + | return i + 2 | |
| 288 | + | } | |
| 289 | + | return i | |
| 290 | + | case s[i] == 0xc2 && i+1 < len(s) && s[i+1] == 0x9c: | |
| 291 | + | return i + 2 | |
| 292 | + | } | |
| 293 | + | i++ | |
| 294 | + | } | |
| 295 | + | return i | |
| 296 | + | } | |
| @@ -0,0 +1,296 @@ | |||||
| 1 | + | package mcp | |||
| 2 | + | ||||
| 3 | + | import ( | |||
| 4 | + | "os" | |||
| 5 | + | "regexp" | |||
| 6 | + | "sort" | |||
| 7 | + | "strings" | |||
| 8 | + | "unicode/utf8" | |||
| 9 | + | ) | |||
| 10 | + | ||||
| 11 | + | // This file holds the diagnostics boundary: everything that originates from a | |||
| 12 | + | // server (stderr, HTTP error bodies, JSON-RPC error messages) or from the HTTP | |||
| 13 | + | // client (url.Error carries the full request URL) is scrubbed here before it | |||
| 14 | + | // reaches an error value, the terminal, or a log. | |||
| 15 | + | ||||
| 16 | + | const ( | |||
| 17 | + | redactedMarker = "[redacted]" | |||
| 18 | + | // minSecretLen keeps one- and two-character configuration values (flags | |||
| 19 | + | // such as "1") from turning every diagnostic into noise. | |||
| 20 | + | minSecretLen = 4 | |||
| 21 | + | ) | |||
| 22 | + | ||||
| 23 | + | // redactor removes configured secret values from server-provided text. | |||
| 24 | + | // A nil *redactor still strips URLs' credentials and control characters. | |||
| 25 | + | type redactor struct { | |||
| 26 | + | secrets []string | |||
| 27 | + | } | |||
| 28 | + | ||||
| 29 | + | // newRedactor collects the secret-bearing values of one server configuration: | |||
| 30 | + | // env and header values, the bearer token, values passed through from the host | |||
| 31 | + | // environment, and credentials embedded in the URL. | |||
| 32 | + | func newRedactor(cfg ServerConfig) *redactor { | |||
| 33 | + | set := map[string]bool{} | |||
| 34 | + | add := func(v string) { | |||
| 35 | + | v = strings.TrimSpace(v) | |||
| 36 | + | if len(v) < minSecretLen { | |||
| 37 | + | return | |||
| 38 | + | } | |||
| 39 | + | set[v] = true | |||
| 40 | + | // "Bearer abc123" is also worth redacting as just "abc123". | |||
| 41 | + | if strings.ContainsAny(v, " \t") { | |||
| 42 | + | for _, field := range strings.Fields(v) { | |||
| 43 | + | if len(field) >= 8 { | |||
| 44 | + | set[field] = true | |||
| 45 | + | } | |||
| 46 | + | } | |||
| 47 | + | } | |||
| 48 | + | } | |||
| 49 | + | for _, v := range cfg.Env { | |||
| 50 | + | add(v) | |||
| 51 | + | } | |||
| 52 | + | for _, v := range cfg.Headers { | |||
| 53 | + | add(v) | |||
| 54 | + | } | |||
| 55 | + | for _, name := range cfg.EnvHeaders { | |||
| 56 | + | if v, ok := os.LookupEnv(name); ok { | |||
| 57 | + | add(v) | |||
| 58 | + | } | |||
| 59 | + | } | |||
| 60 | + | for _, name := range cfg.EnvVars { | |||
| 61 | + | if v, ok := os.LookupEnv(name); ok { | |||
| 62 | + | add(v) | |||
| 63 | + | } | |||
| 64 | + | } | |||
| 65 | + | if cfg.BearerTokenEnvVar != "" { | |||
| 66 | + | if v, ok := os.LookupEnv(cfg.BearerTokenEnvVar); ok { | |||
| 67 | + | add(v) | |||
| 68 | + | } | |||
| 69 | + | } | |||
| 70 | + | if cfg.URL != "" { | |||
| 71 | + | for _, v := range urlSecrets(cfg.URL) { | |||
| 72 | + | add(v) | |||
| 73 | + | } | |||
| 74 | + | } | |||
| 75 | + | r := &redactor{} | |||
| 76 | + | for v := range set { | |||
| 77 | + | r.secrets = append(r.secrets, v) | |||
| 78 | + | } | |||
| 79 | + | // Longest first so a secret containing another is replaced whole. | |||
| 80 | + | sort.Slice(r.secrets, func(i, j int) bool { | |||
| 81 | + | if len(r.secrets[i]) != len(r.secrets[j]) { | |||
| 82 | + | return len(r.secrets[i]) > len(r.secrets[j]) | |||
| 83 | + | } | |||
| 84 | + | return r.secrets[i] < r.secrets[j] | |||
| 85 | + | }) | |||
| 86 | + | return r | |||
| 87 | + | } | |||
| 88 | + | ||||
| 89 | + | // urlSecrets returns the userinfo password/user and query values of raw. | |||
| 90 | + | func urlSecrets(raw string) []string { | |||
| 91 | + | rest := raw | |||
| 92 | + | if i := strings.Index(rest, "://"); i >= 0 { | |||
| 93 | + | rest = rest[i+3:] | |||
| 94 | + | } | |||
| 95 | + | var out []string | |||
| 96 | + | authority := rest | |||
| 97 | + | if i := strings.IndexAny(authority, "/?#"); i >= 0 { | |||
| 98 | + | authority = authority[:i] | |||
| 99 | + | } | |||
| 100 | + | if i := strings.LastIndex(authority, "@"); i >= 0 { | |||
| 101 | + | for _, part := range strings.SplitN(authority[:i], ":", 2) { | |||
| 102 | + | out = append(out, part) | |||
| 103 | + | } | |||
| 104 | + | } | |||
| 105 | + | if i := strings.Index(rest, "?"); i >= 0 { | |||
| 106 | + | query := rest[i+1:] | |||
| 107 | + | if j := strings.Index(query, "#"); j >= 0 { | |||
| 108 | + | query = query[:j] | |||
| 109 | + | } | |||
| 110 | + | for _, pair := range strings.Split(query, "&") { | |||
| 111 | + | if _, v, ok := strings.Cut(pair, "="); ok { | |||
| 112 | + | out = append(out, v) | |||
| 113 | + | } | |||
| 114 | + | } | |||
| 115 | + | } | |||
| 116 | + | return out | |||
| 117 | + | } | |||
| 118 | + | ||||
| 119 | + | // scrub redacts configured secrets, strips credentials and query strings from | |||
| 120 | + | // URLs, and removes terminal control sequences. Newlines are kept. | |||
| 121 | + | func (r *redactor) scrub(s string) string { | |||
| 122 | + | if s == "" { | |||
| 123 | + | return s | |||
| 124 | + | } | |||
| 125 | + | if r != nil { | |||
| 126 | + | for _, secret := range r.secrets { | |||
| 127 | + | s = strings.ReplaceAll(s, secret, redactedMarker) | |||
| 128 | + | } | |||
| 129 | + | } | |||
| 130 | + | s = scrubURLs(s) | |||
| 131 | + | return sanitizeText(s, true) | |||
| 132 | + | } | |||
| 133 | + | ||||
| 134 | + | // scrubOneLine is scrub with newlines collapsed to spaces. | |||
| 135 | + | func (r *redactor) scrubOneLine(s string) string { | |||
| 136 | + | return sanitizeText(r.scrub(s), false) | |||
| 137 | + | } | |||
| 138 | + | ||||
| 139 | + | var urlInText = regexp.MustCompile("(?i)\\b[a-z][a-z0-9+.-]*://[^\\s\"'<>`]+") | |||
| 140 | + | ||||
| 141 | + | // scrubURLs rewrites every URL in s to scheme://host/path, dropping userinfo, | |||
| 142 | + | // query, and fragment. | |||
| 143 | + | func scrubURLs(s string) string { | |||
| 144 | + | if !strings.Contains(s, "://") { | |||
| 145 | + | return s | |||
| 146 | + | } | |||
| 147 | + | return urlInText.ReplaceAllStringFunc(s, stripURL) | |||
| 148 | + | } | |||
| 149 | + | ||||
| 150 | + | func stripURL(raw string) string { | |||
| 151 | + | schemeEnd := strings.Index(raw, "://") | |||
| 152 | + | if schemeEnd < 0 { | |||
| 153 | + | return raw | |||
| 154 | + | } | |||
| 155 | + | scheme, rest := raw[:schemeEnd+3], raw[schemeEnd+3:] | |||
| 156 | + | end := len(rest) | |||
| 157 | + | if i := strings.IndexAny(rest, "/?#"); i >= 0 { | |||
| 158 | + | end = i | |||
| 159 | + | } | |||
| 160 | + | authority, tail := rest[:end], rest[end:] | |||
| 161 | + | if i := strings.LastIndex(authority, "@"); i >= 0 { | |||
| 162 | + | authority = authority[i+1:] | |||
| 163 | + | } | |||
| 164 | + | if i := strings.IndexAny(tail, "?#"); i >= 0 { | |||
| 165 | + | tail = tail[:i] | |||
| 166 | + | } | |||
| 167 | + | return scheme + authority + tail | |||
| 168 | + | } | |||
| 169 | + | ||||
| 170 | + | // scrubbedError carries a sanitized message while keeping the original error | |||
| 171 | + | // reachable through errors.Is/As. | |||
| 172 | + | type scrubbedError struct { | |||
| 173 | + | msg string | |||
| 174 | + | err error | |||
| 175 | + | } | |||
| 176 | + | ||||
| 177 | + | func (e *scrubbedError) Error() string { return e.msg } | |||
| 178 | + | func (e *scrubbedError) Unwrap() error { return e.err } | |||
| 179 | + | ||||
| 180 | + | // scrubError returns err with a scrubbed message, or nil for nil. Context | |||
| 181 | + | // errors are returned unchanged so callers keep comparing them directly. | |||
| 182 | + | func (r *redactor) scrubError(err error) error { | |||
| 183 | + | if err == nil { | |||
| 184 | + | return nil | |||
| 185 | + | } | |||
| 186 | + | msg := r.scrub(err.Error()) | |||
| 187 | + | if msg == err.Error() { | |||
| 188 | + | return err | |||
| 189 | + | } | |||
| 190 | + | return &scrubbedError{msg: msg, err: err} | |||
| 191 | + | } | |||
| 192 | + | ||||
| 193 | + | // sanitizeText removes terminal control sequences (CSI, OSC, DCS/APC/PM/SOS | |||
| 194 | + | // strings, other escapes), C0/C1 control characters, and bidirectional | |||
| 195 | + | // overrides. Tabs become spaces. Newlines are kept when keepNewlines is true | |||
| 196 | + | // and become spaces otherwise; carriage returns are dropped. | |||
| 197 | + | func sanitizeText(s string, keepNewlines bool) string { | |||
| 198 | + | clean := true | |||
| 199 | + | for i := 0; i < len(s); i++ { | |||
| 200 | + | c := s[i] | |||
| 201 | + | if c < 0x20 && c != '\n' || c == 0x7f || c >= 0x80 || (c == '\n' && !keepNewlines) { | |||
| 202 | + | clean = false | |||
| 203 | + | break | |||
| 204 | + | } | |||
| 205 | + | } | |||
| 206 | + | if clean { | |||
| 207 | + | return s | |||
| 208 | + | } | |||
| 209 | + | var b strings.Builder | |||
| 210 | + | b.Grow(len(s)) | |||
| 211 | + | for i := 0; i < len(s); { | |||
| 212 | + | r, size := utf8.DecodeRuneInString(s[i:]) | |||
| 213 | + | i += size | |||
| 214 | + | switch { | |||
| 215 | + | case r == 0x1b: | |||
| 216 | + | i = skipEscape(s, i) | |||
| 217 | + | case r == 0x9b: | |||
| 218 | + | i = skipCSI(s, i) | |||
| 219 | + | case r == 0x9d || r == 0x90 || r == 0x98 || r == 0x9e || r == 0x9f: | |||
| 220 | + | i = skipControlString(s, i) | |||
| 221 | + | case r == '\n': | |||
| 222 | + | if keepNewlines { | |||
| 223 | + | b.WriteByte('\n') | |||
| 224 | + | } else { | |||
| 225 | + | b.WriteByte(' ') | |||
| 226 | + | } | |||
| 227 | + | case r == '\t': | |||
| 228 | + | b.WriteByte(' ') | |||
| 229 | + | case r < 0x20 || r == 0x7f || (r >= 0x80 && r <= 0x9f): | |||
| 230 | + | // dropped (includes \r) | |||
| 231 | + | case (r >= 0x202a && r <= 0x202e) || (r >= 0x2066 && r <= 0x2069) || r == 0x200e || r == 0x200f: | |||
| 232 | + | // bidirectional controls can reorder terminal text | |||
| 233 | + | default: | |||
| 234 | + | b.WriteRune(r) | |||
| 235 | + | } | |||
| 236 | + | } | |||
| 237 | + | return b.String() | |||
| 238 | + | } | |||
| 239 | + | ||||
| 240 | + | // skipEscape returns the index after the escape sequence whose ESC byte was | |||
| 241 | + | // just consumed (i points at the byte after ESC). | |||
| 242 | + | func skipEscape(s string, i int) int { | |||
| 243 | + | if i >= len(s) { | |||
| 244 | + | return i | |||
| 245 | + | } | |||
| 246 | + | switch c := s[i]; { | |||
| 247 | + | case c == '[': | |||
| 248 | + | return skipCSI(s, i+1) | |||
| 249 | + | case c == ']' || c == 'P' || c == 'X' || c == '^' || c == '_': | |||
| 250 | + | return skipControlString(s, i+1) | |||
| 251 | + | case c >= 0x20 && c <= 0x2f: | |||
| 252 | + | for i < len(s) && s[i] >= 0x20 && s[i] <= 0x2f { | |||
| 253 | + | i++ | |||
| 254 | + | } | |||
| 255 | + | if i < len(s) { | |||
| 256 | + | i++ // final byte | |||
| 257 | + | } | |||
| 258 | + | return i | |||
| 259 | + | default: | |||
| 260 | + | return i + 1 | |||
| 261 | + | } | |||
| 262 | + | } | |||
| 263 | + | ||||
| 264 | + | func skipCSI(s string, i int) int { | |||
| 265 | + | for i < len(s) && s[i] >= 0x30 && s[i] <= 0x3f { | |||
| 266 | + | i++ | |||
| 267 | + | } | |||
| 268 | + | for i < len(s) && s[i] >= 0x20 && s[i] <= 0x2f { | |||
| 269 | + | i++ | |||
| 270 | + | } | |||
| 271 | + | if i < len(s) && s[i] >= 0x40 && s[i] <= 0x7e { | |||
| 272 | + | i++ | |||
| 273 | + | } | |||
| 274 | + | return i | |||
| 275 | + | } | |||
| 276 | + | ||||
| 277 | + | // skipControlString consumes an OSC/DCS/APC/PM/SOS payload up to and | |||
| 278 | + | // including its terminator (BEL, ESC \, or C1 ST). An ESC that does not start | |||
| 279 | + | // ST cancels the string without being consumed. | |||
| 280 | + | func skipControlString(s string, i int) int { | |||
| 281 | + | for i < len(s) { | |||
| 282 | + | switch { | |||
| 283 | + | case s[i] == 0x07: | |||
| 284 | + | return i + 1 | |||
| 285 | + | case s[i] == 0x1b: | |||
| 286 | + | if i+1 < len(s) && s[i+1] == '\\' { | |||
| 287 | + | return i + 2 | |||
| 288 | + | } | |||
| 289 | + | return i | |||
| 290 | + | case s[i] == 0xc2 && i+1 < len(s) && s[i+1] == 0x9c: | |||
| 291 | + | return i + 2 | |||
| 292 | + | } | |||
| 293 | + | i++ | |||
| 294 | + | } | |||
| 295 | + | return i | |||
| 296 | + | } | |||
M packages/agent/mcp/stdio.go +268 -60
| @@ -13,14 +13,24 @@ import ( | |||
| 13 | 13 | "runtime" | |
| 14 | 14 | "strings" | |
| 15 | 15 | "sync" | |
| 16 | + | "sync/atomic" | |
| 16 | 17 | "time" | |
| 17 | 18 | ) | |
| 18 | 19 | ||
| 19 | 20 | // maxStdioFrame bounds one newline-delimited JSON-RPC message from a server. | |
| 20 | 21 | const maxStdioFrame = 64 << 20 | |
| 21 | 22 | ||
| 22 | - | // stderrTailLines is how many stderr lines are kept for startup diagnostics. | |
| 23 | - | const stderrTailLines = 20 | |
| 23 | + | const ( | |
| 24 | + | // stderrTailLines is how many stderr lines are kept for diagnostics. | |
| 25 | + | stderrTailLines = 20 | |
| 26 | + | // maxStderrLine truncates one retained stderr line; the rest of a longer | |
| 27 | + | // line is still drained from the pipe, just not kept. | |
| 28 | + | maxStderrLine = 4096 | |
| 29 | + | // writeQueueSize bounds frames waiting for the writer goroutine. | |
| 30 | + | writeQueueSize = 64 | |
| 31 | + | // responseWriteTimeout bounds answering one server-initiated request. | |
| 32 | + | responseWriteTimeout = 30 * time.Second | |
| 33 | + | ) | |
| 24 | 34 | ||
| 25 | 35 | // defaultEnvVars are inherited by stdio servers. Everything else from the host | |
| 26 | 36 | // environment, including provider API keys, is withheld unless listed in | |
| @@ -42,12 +52,28 @@ var defaultEnvVars = func() []string { | |||
| 42 | 52 | } | |
| 43 | 53 | }() | |
| 44 | 54 | ||
| 55 | + | // writeReq is one frame queued for the writer goroutine. | |
| 56 | + | type writeReq struct { | |
| 57 | + | ctx context.Context | |
| 58 | + | frame []byte | |
| 59 | + | res chan error | |
| 60 | + | } | |
| 61 | + | ||
| 45 | 62 | type stdioTransport struct { | |
| 46 | - | name string | |
| 47 | - | cmd *exec.Cmd | |
| 63 | + | name string | |
| 64 | + | cmd *exec.Cmd | |
| 65 | + | guard *processGuard | |
| 66 | + | red *redactor | |
| 48 | 67 | ||
| 49 | - | writeMu sync.Mutex | |
| 50 | - | stdin io.WriteCloser | |
| 68 | + | // ctx spans the transport's lifetime; it is cancelled when the | |
| 69 | + | // connection fails or closes and bounds every background send. | |
| 70 | + | ctx context.Context | |
| 71 | + | cancel context.CancelFunc | |
| 72 | + | ||
| 73 | + | stdin io.WriteCloser | |
| 74 | + | outR *os.File | |
| 75 | + | errR *os.File | |
| 76 | + | writeQ chan *writeReq | |
| 51 | 77 | ||
| 52 | 78 | mu sync.Mutex | |
| 53 | 79 | pending map[string]chan message | |
| @@ -57,9 +83,15 @@ type stdioTransport struct { | |||
| 57 | 83 | stderrMu sync.Mutex | |
| 58 | 84 | stderr []string | |
| 59 | 85 | ||
| 60 | - | doneCh chan struct{} | |
| 61 | - | closeOnce sync.Once | |
| 62 | - | exited chan struct{} | |
| 86 | + | doneCh chan struct{} | |
| 87 | + | closeOnce sync.Once | |
| 88 | + | exited chan struct{} | |
| 89 | + | readDone chan struct{} | |
| 90 | + | stderrDone chan struct{} | |
| 91 | + | ||
| 92 | + | // cancelsSent counts notifications/cancelled sends started (diagnostics | |
| 93 | + | // and tests). | |
| 94 | + | cancelsSent atomic.Int64 | |
| 63 | 95 | } | |
| 64 | 96 | ||
| 65 | 97 | // buildStdioEnv returns the restricted child environment. | |
| @@ -96,60 +128,122 @@ func startStdio(name string, cfg ServerConfig, cwd string) (*stdioTransport, err | |||
| 96 | 128 | if err != nil { | |
| 97 | 129 | return nil, err | |
| 98 | 130 | } | |
| 99 | - | stdout, err := cmd.StdoutPipe() | |
| 131 | + | // Plain OS pipes instead of StdoutPipe/StderrPipe: Wait closes those as | |
| 132 | + | // soon as the child exits, which can discard output the child wrote just | |
| 133 | + | // before exiting (a final response, a fatal stderr line). | |
| 134 | + | outR, outW, err := os.Pipe() | |
| 100 | 135 | if err != nil { | |
| 101 | 136 | return nil, err | |
| 102 | 137 | } | |
| 103 | - | stderr, err := cmd.StderrPipe() | |
| 138 | + | errR, errW, err := os.Pipe() | |
| 104 | 139 | if err != nil { | |
| 140 | + | outR.Close() | |
| 141 | + | outW.Close() | |
| 105 | 142 | return nil, err | |
| 106 | 143 | } | |
| 144 | + | cmd.Stdout = outW | |
| 145 | + | cmd.Stderr = errW | |
| 146 | + | red := newRedactor(cfg) | |
| 107 | 147 | if err := cmd.Start(); err != nil { | |
| 108 | - | return nil, fmt.Errorf("start %s: %w", cfg.Command, err) | |
| 148 | + | outR.Close() | |
| 149 | + | outW.Close() | |
| 150 | + | errR.Close() | |
| 151 | + | errW.Close() | |
| 152 | + | return nil, fmt.Errorf("start %s: %w", cfg.Command, red.scrubError(err)) | |
| 109 | 153 | } | |
| 154 | + | outW.Close() | |
| 155 | + | errW.Close() | |
| 156 | + | ctx, cancel := context.WithCancel(context.Background()) | |
| 110 | 157 | t := &stdioTransport{ | |
| 111 | - | name: name, | |
| 112 | - | cmd: cmd, | |
| 113 | - | stdin: stdin, | |
| 114 | - | pending: map[string]chan message{}, | |
| 115 | - | doneCh: make(chan struct{}), | |
| 116 | - | exited: make(chan struct{}), | |
| 117 | - | } | |
| 118 | - | go t.readStderr(stderr) | |
| 119 | - | go t.readLoop(stdout) | |
| 158 | + | name: name, | |
| 159 | + | cmd: cmd, | |
| 160 | + | guard: newProcessGuard(cmd), | |
| 161 | + | red: red, | |
| 162 | + | ctx: ctx, | |
| 163 | + | cancel: cancel, | |
| 164 | + | stdin: stdin, | |
| 165 | + | outR: outR, | |
| 166 | + | errR: errR, | |
| 167 | + | writeQ: make(chan *writeReq, writeQueueSize), | |
| 168 | + | pending: map[string]chan message{}, | |
| 169 | + | doneCh: make(chan struct{}), | |
| 170 | + | exited: make(chan struct{}), | |
| 171 | + | readDone: make(chan struct{}), | |
| 172 | + | stderrDone: make(chan struct{}), | |
| 173 | + | } | |
| 174 | + | go t.writeLoop() | |
| 175 | + | go t.readStderr(errR) | |
| 176 | + | go t.readLoop(outR) | |
| 120 | 177 | go func() { | |
| 121 | 178 | waitErr := cmd.Wait() | |
| 122 | 179 | close(t.exited) | |
| 123 | 180 | if waitErr == nil { | |
| 124 | 181 | waitErr = errors.New("process exited") | |
| 125 | 182 | } | |
| 183 | + | // Let the reader hand over whatever the server wrote before exiting, | |
| 184 | + | // so a response sitting in the pipe wins over the failure. | |
| 185 | + | select { | |
| 186 | + | case <-t.readDone: | |
| 187 | + | case <-time.After(250 * time.Millisecond): | |
| 188 | + | } | |
| 126 | 189 | t.fail(fmt.Errorf("server process ended: %w", waitErr)) | |
| 127 | 190 | }() | |
| 128 | 191 | return t, nil | |
| 129 | 192 | } | |
| 130 | 193 | ||
| 194 | + | // readStderr drains the child's stderr for its whole life, keeping a bounded | |
| 195 | + | // tail of bounded lines. It never stops early: a child blocked on a full | |
| 196 | + | // stderr pipe would hang. | |
| 131 | 197 | func (t *stdioTransport) readStderr(r io.Reader) { | |
| 132 | - | scanner := bufio.NewScanner(r) | |
| 133 | - | scanner.Buffer(make([]byte, 0, 4096), 1<<20) | |
| 134 | - | for scanner.Scan() { | |
| 135 | - | line := strings.TrimRight(scanner.Text(), "\r") | |
| 198 | + | defer close(t.stderrDone) | |
| 199 | + | reader := bufio.NewReaderSize(r, 4096) | |
| 200 | + | var line []byte | |
| 201 | + | truncated := false | |
| 202 | + | flush := func() { | |
| 203 | + | text := strings.TrimRight(string(line), "\r") | |
| 204 | + | if truncated { | |
| 205 | + | text += "…[line truncated]" | |
| 206 | + | } | |
| 207 | + | line, truncated = nil, false | |
| 136 | 208 | t.stderrMu.Lock() | |
| 137 | - | t.stderr = append(t.stderr, line) | |
| 209 | + | t.stderr = append(t.stderr, text) | |
| 138 | 210 | if len(t.stderr) > stderrTailLines { | |
| 139 | 211 | t.stderr = t.stderr[len(t.stderr)-stderrTailLines:] | |
| 140 | 212 | } | |
| 141 | 213 | t.stderrMu.Unlock() | |
| 142 | 214 | } | |
| 215 | + | for { | |
| 216 | + | chunk, isPrefix, err := reader.ReadLine() | |
| 217 | + | if room := maxStderrLine - len(line); len(chunk) > room { | |
| 218 | + | if room > 0 { | |
| 219 | + | line = append(line, chunk[:room]...) | |
| 220 | + | } | |
| 221 | + | truncated = true | |
| 222 | + | } else { | |
| 223 | + | line = append(line, chunk...) | |
| 224 | + | } | |
| 225 | + | if err != nil { | |
| 226 | + | if len(line) > 0 || truncated { | |
| 227 | + | flush() | |
| 228 | + | } | |
| 229 | + | return | |
| 230 | + | } | |
| 231 | + | if !isPrefix { | |
| 232 | + | flush() | |
| 233 | + | } | |
| 234 | + | } | |
| 143 | 235 | } | |
| 144 | 236 | ||
| 145 | - | // stderrTail returns recent stderr output for diagnostics. | |
| 237 | + | // stderrTail returns recent stderr output, scrubbed for display. | |
| 146 | 238 | func (t *stdioTransport) stderrTail() string { | |
| 147 | 239 | t.stderrMu.Lock() | |
| 148 | - | defer t.stderrMu.Unlock() | |
| 149 | - | return strings.Join(t.stderr, "\n") | |
| 240 | + | lines := append([]string(nil), t.stderr...) | |
| 241 | + | t.stderrMu.Unlock() | |
| 242 | + | return t.red.scrub(strings.Join(lines, "\n")) | |
| 150 | 243 | } | |
| 151 | 244 | ||
| 152 | 245 | func (t *stdioTransport) readLoop(r io.Reader) { | |
| 246 | + | defer close(t.readDone) | |
| 153 | 247 | reader := bufio.NewReaderSize(r, 64*1024) | |
| 154 | 248 | for { | |
| 155 | 249 | line, err := readLine(reader, maxStdioFrame) | |
| @@ -163,6 +257,17 @@ func (t *stdioTransport) readLoop(r io.Reader) { | |||
| 163 | 257 | // Non-JSON stdout lines (stray logging) are ignored. | |
| 164 | 258 | } | |
| 165 | 259 | if err != nil { | |
| 260 | + | if errors.Is(err, io.EOF) || errors.Is(err, os.ErrClosed) { | |
| 261 | + | // Output closed. If the process is gone the wait goroutine | |
| 262 | + | // reports its exit status; otherwise give it a moment. | |
| 263 | + | select { | |
| 264 | + | case <-t.exited: | |
| 265 | + | return | |
| 266 | + | case <-t.doneCh: | |
| 267 | + | return | |
| 268 | + | case <-time.After(time.Second): | |
| 269 | + | } | |
| 270 | + | } | |
| 166 | 271 | t.fail(fmt.Errorf("read server output: %w", err)) | |
| 167 | 272 | return | |
| 168 | 273 | } | |
| @@ -208,7 +313,9 @@ func (t *stdioTransport) dispatch(msg message) { | |||
| 208 | 313 | result, rpcErr = handler(msg) | |
| 209 | 314 | } | |
| 210 | 315 | if frame, err := encodeResponse(msg.ID, result, rpcErr); err == nil { | |
| 211 | - | _ = t.write(frame) | |
| 316 | + | ctx, cancel := context.WithTimeout(t.ctx, responseWriteTimeout) | |
| 317 | + | defer cancel() | |
| 318 | + | _ = t.write(ctx, frame) | |
| 212 | 319 | } | |
| 213 | 320 | }() | |
| 214 | 321 | return | |
| @@ -218,16 +325,55 @@ func (t *stdioTransport) dispatch(msg message) { | |||
| 218 | 325 | } | |
| 219 | 326 | } | |
| 220 | 327 | ||
| 221 | - | func (t *stdioTransport) write(frame []byte) error { | |
| 222 | - | t.writeMu.Lock() | |
| 223 | - | defer t.writeMu.Unlock() | |
| 328 | + | // writeLoop owns all writes to the child's stdin. Callers never block on the | |
| 329 | + | // pipe themselves: they queue a frame and wait for the result or their | |
| 330 | + | // context, so a child that stops reading cannot wedge a request, a timeout, or | |
| 331 | + | // close(). | |
| 332 | + | func (t *stdioTransport) writeLoop() { | |
| 333 | + | for { | |
| 334 | + | select { | |
| 335 | + | case <-t.doneCh: | |
| 336 | + | return | |
| 337 | + | case req := <-t.writeQ: | |
| 338 | + | if err := req.ctx.Err(); err != nil { | |
| 339 | + | req.res <- err // never started: nothing partial on the wire | |
| 340 | + | continue | |
| 341 | + | } | |
| 342 | + | _, err := t.stdin.Write(req.frame) | |
| 343 | + | req.res <- err | |
| 344 | + | if err != nil { | |
| 345 | + | // A failed write may have left half a frame behind. | |
| 346 | + | t.fail(fmt.Errorf("write to server: %w", err)) | |
| 347 | + | return | |
| 348 | + | } | |
| 349 | + | } | |
| 350 | + | } | |
| 351 | + | } | |
| 352 | + | ||
| 353 | + | // write queues one frame (newline appended) and waits until it has been | |
| 354 | + | // written, ctx ends, or the connection fails. | |
| 355 | + | func (t *stdioTransport) write(ctx context.Context, frame []byte) error { | |
| 356 | + | req := &writeReq{ctx: ctx, frame: append(append(make([]byte, 0, len(frame)+1), frame...), '\n'), res: make(chan error, 1)} | |
| 224 | 357 | select { | |
| 358 | + | case t.writeQ <- req: | |
| 359 | + | case <-ctx.Done(): | |
| 360 | + | return ctx.Err() | |
| 225 | 361 | case <-t.doneCh: | |
| 226 | 362 | return t.failure() | |
| 227 | - | default: | |
| 228 | 363 | } | |
| 229 | - | _, err := t.stdin.Write(append(frame, '\n')) | |
| 230 | - | return err | |
| 364 | + | select { | |
| 365 | + | case err := <-req.res: | |
| 366 | + | return err | |
| 367 | + | case <-ctx.Done(): | |
| 368 | + | return ctx.Err() | |
| 369 | + | case <-t.doneCh: | |
| 370 | + | select { | |
| 371 | + | case err := <-req.res: | |
| 372 | + | return err | |
| 373 | + | default: | |
| 374 | + | } | |
| 375 | + | return t.failure() | |
| 376 | + | } | |
| 231 | 377 | } | |
| 232 | 378 | ||
| 233 | 379 | func (t *stdioTransport) request(ctx context.Context, id int64, method string, params any) (json.RawMessage, error) { | |
| @@ -250,32 +396,60 @@ func (t *stdioTransport) request(ctx context.Context, id int64, method string, p | |||
| 250 | 396 | delete(t.pending, key) | |
| 251 | 397 | t.mu.Unlock() | |
| 252 | 398 | } | |
| 253 | - | if err := t.write(frame); err != nil { | |
| 399 | + | if err := t.write(ctx, frame); err != nil { | |
| 254 | 400 | cleanup() | |
| 401 | + | if ctx.Err() != nil { | |
| 402 | + | t.cancelRequest(id, method, ctx.Err().Error()) | |
| 403 | + | return nil, ctx.Err() | |
| 404 | + | } | |
| 255 | 405 | return nil, err | |
| 256 | 406 | } | |
| 257 | - | select { | |
| 258 | - | case msg := <-ch: | |
| 407 | + | result := func(msg message) (json.RawMessage, error) { | |
| 259 | 408 | if msg.Error != nil { | |
| 260 | - | return nil, msg.Error | |
| 409 | + | return nil, t.red.scrubRPCError(msg.Error) | |
| 261 | 410 | } | |
| 262 | 411 | return msg.Result, nil | |
| 412 | + | } | |
| 413 | + | select { | |
| 414 | + | case msg := <-ch: | |
| 415 | + | return result(msg) | |
| 263 | 416 | case <-ctx.Done(): | |
| 264 | 417 | cleanup() | |
| 265 | - | _ = t.notify(context.Background(), "notifications/cancelled", map[string]any{"requestId": id, "reason": ctx.Err().Error()}) | |
| 418 | + | t.cancelRequest(id, method, ctx.Err().Error()) | |
| 266 | 419 | return nil, ctx.Err() | |
| 267 | 420 | case <-t.doneCh: | |
| 421 | + | // A response that arrived together with the failure still counts. | |
| 422 | + | select { | |
| 423 | + | case msg := <-ch: | |
| 424 | + | return result(msg) | |
| 425 | + | default: | |
| 426 | + | } | |
| 268 | 427 | cleanup() | |
| 269 | 428 | return nil, t.failure() | |
| 270 | 429 | } | |
| 271 | 430 | } | |
| 272 | 431 | ||
| 273 | - | func (t *stdioTransport) notify(_ context.Context, method string, params any) error { | |
| 432 | + | // cancelRequest announces an abandoned request on a best-effort basis. It | |
| 433 | + | // never blocks the caller, is bounded by cancelNotifyTimeout, and dies with | |
| 434 | + | // the transport. | |
| 435 | + | func (t *stdioTransport) cancelRequest(id int64, method, reason string) { | |
| 436 | + | if !shouldSendCancel(method) { | |
| 437 | + | return | |
| 438 | + | } | |
| 439 | + | t.cancelsSent.Add(1) | |
| 440 | + | go func() { | |
| 441 | + | ctx, cancel := context.WithTimeout(t.ctx, cancelNotifyTimeout) | |
| 442 | + | defer cancel() | |
| 443 | + | _ = t.notify(ctx, "notifications/cancelled", map[string]any{"requestId": id, "reason": reason}) | |
| 444 | + | }() | |
| 445 | + | } | |
| 446 | + | ||
| 447 | + | func (t *stdioTransport) notify(ctx context.Context, method string, params any) error { | |
| 274 | 448 | frame, err := encodeNotification(method, params) | |
| 275 | 449 | if err != nil { | |
| 276 | 450 | return err | |
| 277 | 451 | } | |
| 278 | - | return t.write(frame) | |
| 452 | + | return t.write(ctx, frame) | |
| 279 | 453 | } | |
| 280 | 454 | ||
| 281 | 455 | func (t *stdioTransport) setHandler(h incomingHandler) { | |
| @@ -303,31 +477,65 @@ func (t *stdioTransport) fail(err error) { | |||
| 303 | 477 | t.err = err | |
| 304 | 478 | } | |
| 305 | 479 | t.mu.Unlock() | |
| 306 | - | t.closeOnce.Do(func() { close(t.doneCh) }) | |
| 480 | + | t.closeOnce.Do(func() { | |
| 481 | + | t.cancel() | |
| 482 | + | close(t.doneCh) | |
| 483 | + | }) | |
| 307 | 484 | } | |
| 308 | 485 | ||
| 309 | - | // close shuts the server down: close stdin (the MCP stdio shutdown signal), | |
| 310 | - | // wait briefly, then terminate the whole process group. | |
| 311 | - | func (t *stdioTransport) close() error { | |
| 312 | - | t.fail(errors.New("connection closed")) | |
| 313 | - | t.writeMu.Lock() | |
| 314 | - | _ = t.stdin.Close() | |
| 315 | - | t.writeMu.Unlock() | |
| 486 | + | func (t *stdioTransport) waitExit(d time.Duration) bool { | |
| 316 | 487 | select { | |
| 317 | 488 | case <-t.exited: | |
| 318 | - | return nil | |
| 319 | - | case <-time.After(500 * time.Millisecond): | |
| 489 | + | return true | |
| 490 | + | case <-time.After(d): | |
| 491 | + | return false | |
| 320 | 492 | } | |
| 321 | - | terminateProcess(t.cmd, false) | |
| 322 | - | select { | |
| 323 | - | case <-t.exited: | |
| 324 | - | return nil | |
| 325 | - | case <-time.After(1500 * time.Millisecond): | |
| 493 | + | } | |
| 494 | + | ||
| 495 | + | // close shuts the server down: close stdin (the MCP stdio shutdown signal) | |
| 496 | + | // without waiting behind any blocked write, wait briefly, then terminate the | |
| 497 | + | // process group. Descendants that outlive the direct child are signalled too. | |
| 498 | + | func (t *stdioTransport) close() error { | |
| 499 | + | t.fail(errors.New("connection closed")) | |
| 500 | + | _ = t.stdin.Close() // also unblocks a writer stuck on a full pipe | |
| 501 | + | if !t.waitExit(500 * time.Millisecond) { | |
| 502 | + | t.guard.terminate(false) | |
| 503 | + | if !t.waitExit(1500 * time.Millisecond) { | |
| 504 | + | t.guard.terminate(true) | |
| 505 | + | t.waitExit(time.Second) | |
| 506 | + | } | |
| 326 | 507 | } | |
| 327 | - | terminateProcess(t.cmd, true) | |
| 508 | + | t.reapGroup() | |
| 509 | + | // With the child gone its stderr hits EOF once buffered data is read; give | |
| 510 | + | // the reader a moment so diagnostics are complete. A descendant holding the | |
| 511 | + | // pipe open only costs the timeout. | |
| 328 | 512 | select { | |
| 329 | 513 | case <-t.exited: | |
| 330 | - | case <-time.After(time.Second): | |
| 514 | + | select { | |
| 515 | + | case <-t.stderrDone: | |
| 516 | + | case <-time.After(500 * time.Millisecond): | |
| 517 | + | } | |
| 518 | + | default: | |
| 331 | 519 | } | |
| 520 | + | _ = t.outR.Close() | |
| 521 | + | _ = t.errR.Close() | |
| 522 | + | t.guard.release() | |
| 332 | 523 | return nil | |
| 333 | 524 | } | |
| 525 | + | ||
| 526 | + | // reapGroup stops processes the direct child left behind. The process group | |
| 527 | + | // outlives its leader, so this is meaningful even after the leader has been | |
| 528 | + | // reaped (the group id cannot be reused while any member remains). | |
| 529 | + | func (t *stdioTransport) reapGroup() { | |
| 530 | + | if !t.guard.alive() { | |
| 531 | + | return | |
| 532 | + | } | |
| 533 | + | t.guard.terminate(false) | |
| 534 | + | deadline := time.Now().Add(300 * time.Millisecond) | |
| 535 | + | for t.guard.alive() && time.Now().Before(deadline) { | |
| 536 | + | time.Sleep(10 * time.Millisecond) | |
| 537 | + | } | |
| 538 | + | if t.guard.alive() { | |
| 539 | + | t.guard.terminate(true) | |
| 540 | + | } | |
| 541 | + | } | |
| @@ -13,14 +13,24 @@ import ( | |||||
| 13 | "runtime" | 13 | "runtime" | ||
| 14 | "strings" | 14 | "strings" | ||
| 15 | "sync" | 15 | "sync" | ||
| 16 | + | "sync/atomic" | |||
| 16 | "time" | 17 | "time" | ||
| 17 | ) | 18 | ) | ||
| 18 | 19 | ||||
| 19 | // maxStdioFrame bounds one newline-delimited JSON-RPC message from a server. | 20 | // maxStdioFrame bounds one newline-delimited JSON-RPC message from a server. | ||
| 20 | const maxStdioFrame = 64 << 20 | 21 | const maxStdioFrame = 64 << 20 | ||
| 21 | 22 | ||||
| 22 | - | // stderrTailLines is how many stderr lines are kept for startup diagnostics. | 23 | + | const ( |
| 23 | - | const stderrTailLines = 20 | 24 | + | // stderrTailLines is how many stderr lines are kept for diagnostics. |
| 25 | + | stderrTailLines = 20 | |||
| 26 | + | // maxStderrLine truncates one retained stderr line; the rest of a longer | |||
| 27 | + | // line is still drained from the pipe, just not kept. | |||
| 28 | + | maxStderrLine = 4096 | |||
| 29 | + | // writeQueueSize bounds frames waiting for the writer goroutine. | |||
| 30 | + | writeQueueSize = 64 | |||
| 31 | + | // responseWriteTimeout bounds answering one server-initiated request. | |||
| 32 | + | responseWriteTimeout = 30 * time.Second | |||
| 33 | + | ) | |||
| 24 | 34 | ||||
| 25 | // defaultEnvVars are inherited by stdio servers. Everything else from the host | 35 | // defaultEnvVars are inherited by stdio servers. Everything else from the host | ||
| 26 | // environment, including provider API keys, is withheld unless listed in | 36 | // environment, including provider API keys, is withheld unless listed in | ||
| @@ -42,12 +52,28 @@ var defaultEnvVars = func() []string { | |||||
| 42 | } | 52 | } | ||
| 43 | }() | 53 | }() | ||
| 44 | 54 | ||||
| 55 | + | // writeReq is one frame queued for the writer goroutine. | |||
| 56 | + | type writeReq struct { | |||
| 57 | + | ctx context.Context | |||
| 58 | + | frame []byte | |||
| 59 | + | res chan error | |||
| 60 | + | } | |||
| 61 | + | ||||
| 45 | type stdioTransport struct { | 62 | type stdioTransport struct { | ||
| 46 | - | name string | 63 | + | name string |
| 47 | - | cmd *exec.Cmd | 64 | + | cmd *exec.Cmd |
| 65 | + | guard *processGuard | |||
| 66 | + | red *redactor | |||
| 48 | 67 | ||||
| 49 | - | writeMu sync.Mutex | 68 | + | // ctx spans the transport's lifetime; it is cancelled when the |
| 50 | - | stdin io.WriteCloser | 69 | + | // connection fails or closes and bounds every background send. |
| 70 | + | ctx context.Context | |||
| 71 | + | cancel context.CancelFunc | |||
| 72 | + | ||||
| 73 | + | stdin io.WriteCloser | |||
| 74 | + | outR *os.File | |||
| 75 | + | errR *os.File | |||
| 76 | + | writeQ chan *writeReq | |||
| 51 | 77 | ||||
| 52 | mu sync.Mutex | 78 | mu sync.Mutex | ||
| 53 | pending map[string]chan message | 79 | pending map[string]chan message | ||
| @@ -57,9 +83,15 @@ type stdioTransport struct { | |||||
| 57 | stderrMu sync.Mutex | 83 | stderrMu sync.Mutex | ||
| 58 | stderr []string | 84 | stderr []string | ||
| 59 | 85 | ||||
| 60 | - | doneCh chan struct{} | 86 | + | doneCh chan struct{} |
| 61 | - | closeOnce sync.Once | 87 | + | closeOnce sync.Once |
| 62 | - | exited chan struct{} | 88 | + | exited chan struct{} |
| 89 | + | readDone chan struct{} | |||
| 90 | + | stderrDone chan struct{} | |||
| 91 | + | ||||
| 92 | + | // cancelsSent counts notifications/cancelled sends started (diagnostics | |||
| 93 | + | // and tests). | |||
| 94 | + | cancelsSent atomic.Int64 | |||
| 63 | } | 95 | } | ||
| 64 | 96 | ||||
| 65 | // buildStdioEnv returns the restricted child environment. | 97 | // buildStdioEnv returns the restricted child environment. | ||
| @@ -96,60 +128,122 @@ func startStdio(name string, cfg ServerConfig, cwd string) (*stdioTransport, err | |||||
| 96 | if err != nil { | 128 | if err != nil { | ||
| 97 | return nil, err | 129 | return nil, err | ||
| 98 | } | 130 | } | ||
| 99 | - | stdout, err := cmd.StdoutPipe() | 131 | + | // Plain OS pipes instead of StdoutPipe/StderrPipe: Wait closes those as |
| 132 | + | // soon as the child exits, which can discard output the child wrote just | |||
| 133 | + | // before exiting (a final response, a fatal stderr line). | |||
| 134 | + | outR, outW, err := os.Pipe() | |||
| 100 | if err != nil { | 135 | if err != nil { | ||
| 101 | return nil, err | 136 | return nil, err | ||
| 102 | } | 137 | } | ||
| 103 | - | stderr, err := cmd.StderrPipe() | 138 | + | errR, errW, err := os.Pipe() |
| 104 | if err != nil { | 139 | if err != nil { | ||
| 140 | + | outR.Close() | |||
| 141 | + | outW.Close() | |||
| 105 | return nil, err | 142 | return nil, err | ||
| 106 | } | 143 | } | ||
| 144 | + | cmd.Stdout = outW | |||
| 145 | + | cmd.Stderr = errW | |||
| 146 | + | red := newRedactor(cfg) | |||
| 107 | if err := cmd.Start(); err != nil { | 147 | if err := cmd.Start(); err != nil { | ||
| 108 | - | return nil, fmt.Errorf("start %s: %w", cfg.Command, err) | 148 | + | outR.Close() |
| 149 | + | outW.Close() | |||
| 150 | + | errR.Close() | |||
| 151 | + | errW.Close() | |||
| 152 | + | return nil, fmt.Errorf("start %s: %w", cfg.Command, red.scrubError(err)) | |||
| 109 | } | 153 | } | ||
| 154 | + | outW.Close() | |||
| 155 | + | errW.Close() | |||
| 156 | + | ctx, cancel := context.WithCancel(context.Background()) | |||
| 110 | t := &stdioTransport{ | 157 | t := &stdioTransport{ | ||
| 111 | - | name: name, | 158 | + | name: name, |
| 112 | - | cmd: cmd, | 159 | + | cmd: cmd, |
| 113 | - | stdin: stdin, | 160 | + | guard: newProcessGuard(cmd), |
| 114 | - | pending: map[string]chan message{}, | 161 | + | red: red, |
| 115 | - | doneCh: make(chan struct{}), | 162 | + | ctx: ctx, |
| 116 | - | exited: make(chan struct{}), | 163 | + | cancel: cancel, |
| 117 | - | } | 164 | + | stdin: stdin, |
| 118 | - | go t.readStderr(stderr) | 165 | + | outR: outR, |
| 119 | - | go t.readLoop(stdout) | 166 | + | errR: errR, |
| 167 | + | writeQ: make(chan *writeReq, writeQueueSize), | |||
| 168 | + | pending: map[string]chan message{}, | |||
| 169 | + | doneCh: make(chan struct{}), | |||
| 170 | + | exited: make(chan struct{}), | |||
| 171 | + | readDone: make(chan struct{}), | |||
| 172 | + | stderrDone: make(chan struct{}), | |||
| 173 | + | } | |||
| 174 | + | go t.writeLoop() | |||
| 175 | + | go t.readStderr(errR) | |||
| 176 | + | go t.readLoop(outR) | |||
| 120 | go func() { | 177 | go func() { | ||
| 121 | waitErr := cmd.Wait() | 178 | waitErr := cmd.Wait() | ||
| 122 | close(t.exited) | 179 | close(t.exited) | ||
| 123 | if waitErr == nil { | 180 | if waitErr == nil { | ||
| 124 | waitErr = errors.New("process exited") | 181 | waitErr = errors.New("process exited") | ||
| 125 | } | 182 | } | ||
| 183 | + | // Let the reader hand over whatever the server wrote before exiting, | |||
| 184 | + | // so a response sitting in the pipe wins over the failure. | |||
| 185 | + | select { | |||
| 186 | + | case <-t.readDone: | |||
| 187 | + | case <-time.After(250 * time.Millisecond): | |||
| 188 | + | } | |||
| 126 | t.fail(fmt.Errorf("server process ended: %w", waitErr)) | 189 | t.fail(fmt.Errorf("server process ended: %w", waitErr)) | ||
| 127 | }() | 190 | }() | ||
| 128 | return t, nil | 191 | return t, nil | ||
| 129 | } | 192 | } | ||
| 130 | 193 | ||||
| 194 | + | // readStderr drains the child's stderr for its whole life, keeping a bounded | |||
| 195 | + | // tail of bounded lines. It never stops early: a child blocked on a full | |||
| 196 | + | // stderr pipe would hang. | |||
| 131 | func (t *stdioTransport) readStderr(r io.Reader) { | 197 | func (t *stdioTransport) readStderr(r io.Reader) { | ||
| 132 | - | scanner := bufio.NewScanner(r) | 198 | + | defer close(t.stderrDone) |
| 133 | - | scanner.Buffer(make([]byte, 0, 4096), 1<<20) | 199 | + | reader := bufio.NewReaderSize(r, 4096) |
| 134 | - | for scanner.Scan() { | 200 | + | var line []byte |
| 135 | - | line := strings.TrimRight(scanner.Text(), "\r") | 201 | + | truncated := false |
| 202 | + | flush := func() { | |||
| 203 | + | text := strings.TrimRight(string(line), "\r") | |||
| 204 | + | if truncated { | |||
| 205 | + | text += "…[line truncated]" | |||
| 206 | + | } | |||
| 207 | + | line, truncated = nil, false | |||
| 136 | t.stderrMu.Lock() | 208 | t.stderrMu.Lock() | ||
| 137 | - | t.stderr = append(t.stderr, line) | 209 | + | t.stderr = append(t.stderr, text) |
| 138 | if len(t.stderr) > stderrTailLines { | 210 | if len(t.stderr) > stderrTailLines { | ||
| 139 | t.stderr = t.stderr[len(t.stderr)-stderrTailLines:] | 211 | t.stderr = t.stderr[len(t.stderr)-stderrTailLines:] | ||
| 140 | } | 212 | } | ||
| 141 | t.stderrMu.Unlock() | 213 | t.stderrMu.Unlock() | ||
| 142 | } | 214 | } | ||
| 215 | + | for { | |||
| 216 | + | chunk, isPrefix, err := reader.ReadLine() | |||
| 217 | + | if room := maxStderrLine - len(line); len(chunk) > room { | |||
| 218 | + | if room > 0 { | |||
| 219 | + | line = append(line, chunk[:room]...) | |||
| 220 | + | } | |||
| 221 | + | truncated = true | |||
| 222 | + | } else { | |||
| 223 | + | line = append(line, chunk...) | |||
| 224 | + | } | |||
| 225 | + | if err != nil { | |||
| 226 | + | if len(line) > 0 || truncated { | |||
| 227 | + | flush() | |||
| 228 | + | } | |||
| 229 | + | return | |||
| 230 | + | } | |||
| 231 | + | if !isPrefix { | |||
| 232 | + | flush() | |||
| 233 | + | } | |||
| 234 | + | } | |||
| 143 | } | 235 | } | ||
| 144 | 236 | ||||
| 145 | - | // stderrTail returns recent stderr output for diagnostics. | 237 | + | // stderrTail returns recent stderr output, scrubbed for display. |
| 146 | func (t *stdioTransport) stderrTail() string { | 238 | func (t *stdioTransport) stderrTail() string { | ||
| 147 | t.stderrMu.Lock() | 239 | t.stderrMu.Lock() | ||
| 148 | - | defer t.stderrMu.Unlock() | 240 | + | lines := append([]string(nil), t.stderr...) |
| 149 | - | return strings.Join(t.stderr, "\n") | 241 | + | t.stderrMu.Unlock() |
| 242 | + | return t.red.scrub(strings.Join(lines, "\n")) | |||
| 150 | } | 243 | } | ||
| 151 | 244 | ||||
| 152 | func (t *stdioTransport) readLoop(r io.Reader) { | 245 | func (t *stdioTransport) readLoop(r io.Reader) { | ||
| 246 | + | defer close(t.readDone) | |||
| 153 | reader := bufio.NewReaderSize(r, 64*1024) | 247 | reader := bufio.NewReaderSize(r, 64*1024) | ||
| 154 | for { | 248 | for { | ||
| 155 | line, err := readLine(reader, maxStdioFrame) | 249 | line, err := readLine(reader, maxStdioFrame) | ||
| @@ -163,6 +257,17 @@ func (t *stdioTransport) readLoop(r io.Reader) { | |||||
| 163 | // Non-JSON stdout lines (stray logging) are ignored. | 257 | // Non-JSON stdout lines (stray logging) are ignored. | ||
| 164 | } | 258 | } | ||
| 165 | if err != nil { | 259 | if err != nil { | ||
| 260 | + | if errors.Is(err, io.EOF) || errors.Is(err, os.ErrClosed) { | |||
| 261 | + | // Output closed. If the process is gone the wait goroutine | |||
| 262 | + | // reports its exit status; otherwise give it a moment. | |||
| 263 | + | select { | |||
| 264 | + | case <-t.exited: | |||
| 265 | + | return | |||
| 266 | + | case <-t.doneCh: | |||
| 267 | + | return | |||
| 268 | + | case <-time.After(time.Second): | |||
| 269 | + | } | |||
| 270 | + | } | |||
| 166 | t.fail(fmt.Errorf("read server output: %w", err)) | 271 | t.fail(fmt.Errorf("read server output: %w", err)) | ||
| 167 | return | 272 | return | ||
| 168 | } | 273 | } | ||
| @@ -208,7 +313,9 @@ func (t *stdioTransport) dispatch(msg message) { | |||||
| 208 | result, rpcErr = handler(msg) | 313 | result, rpcErr = handler(msg) | ||
| 209 | } | 314 | } | ||
| 210 | if frame, err := encodeResponse(msg.ID, result, rpcErr); err == nil { | 315 | if frame, err := encodeResponse(msg.ID, result, rpcErr); err == nil { | ||
| 211 | - | _ = t.write(frame) | 316 | + | ctx, cancel := context.WithTimeout(t.ctx, responseWriteTimeout) |
| 317 | + | defer cancel() | |||
| 318 | + | _ = t.write(ctx, frame) | |||
| 212 | } | 319 | } | ||
| 213 | }() | 320 | }() | ||
| 214 | return | 321 | return | ||
| @@ -218,16 +325,55 @@ func (t *stdioTransport) dispatch(msg message) { | |||||
| 218 | } | 325 | } | ||
| 219 | } | 326 | } | ||
| 220 | 327 | ||||
| 221 | - | func (t *stdioTransport) write(frame []byte) error { | 328 | + | // writeLoop owns all writes to the child's stdin. Callers never block on the |
| 222 | - | t.writeMu.Lock() | 329 | + | // pipe themselves: they queue a frame and wait for the result or their |
| 223 | - | defer t.writeMu.Unlock() | 330 | + | // context, so a child that stops reading cannot wedge a request, a timeout, or |
| 331 | + | // close(). | |||
| 332 | + | func (t *stdioTransport) writeLoop() { | |||
| 333 | + | for { | |||
| 334 | + | select { | |||
| 335 | + | case <-t.doneCh: | |||
| 336 | + | return | |||
| 337 | + | case req := <-t.writeQ: | |||
| 338 | + | if err := req.ctx.Err(); err != nil { | |||
| 339 | + | req.res <- err // never started: nothing partial on the wire | |||
| 340 | + | continue | |||
| 341 | + | } | |||
| 342 | + | _, err := t.stdin.Write(req.frame) | |||
| 343 | + | req.res <- err | |||
| 344 | + | if err != nil { | |||
| 345 | + | // A failed write may have left half a frame behind. | |||
| 346 | + | t.fail(fmt.Errorf("write to server: %w", err)) | |||
| 347 | + | return | |||
| 348 | + | } | |||
| 349 | + | } | |||
| 350 | + | } | |||
| 351 | + | } | |||
| 352 | + | ||||
| 353 | + | // write queues one frame (newline appended) and waits until it has been | |||
| 354 | + | // written, ctx ends, or the connection fails. | |||
| 355 | + | func (t *stdioTransport) write(ctx context.Context, frame []byte) error { | |||
| 356 | + | req := &writeReq{ctx: ctx, frame: append(append(make([]byte, 0, len(frame)+1), frame...), '\n'), res: make(chan error, 1)} | |||
| 224 | select { | 357 | select { | ||
| 358 | + | case t.writeQ <- req: | |||
| 359 | + | case <-ctx.Done(): | |||
| 360 | + | return ctx.Err() | |||
| 225 | case <-t.doneCh: | 361 | case <-t.doneCh: | ||
| 226 | return t.failure() | 362 | return t.failure() | ||
| 227 | - | default: | |||
| 228 | } | 363 | } | ||
| 229 | - | _, err := t.stdin.Write(append(frame, '\n')) | 364 | + | select { |
| 230 | - | return err | 365 | + | case err := <-req.res: |
| 366 | + | return err | |||
| 367 | + | case <-ctx.Done(): | |||
| 368 | + | return ctx.Err() | |||
| 369 | + | case <-t.doneCh: | |||
| 370 | + | select { | |||
| 371 | + | case err := <-req.res: | |||
| 372 | + | return err | |||
| 373 | + | default: | |||
| 374 | + | } | |||
| 375 | + | return t.failure() | |||
| 376 | + | } | |||
| 231 | } | 377 | } | ||
| 232 | 378 | ||||
| 233 | func (t *stdioTransport) request(ctx context.Context, id int64, method string, params any) (json.RawMessage, error) { | 379 | func (t *stdioTransport) request(ctx context.Context, id int64, method string, params any) (json.RawMessage, error) { | ||
| @@ -250,32 +396,60 @@ func (t *stdioTransport) request(ctx context.Context, id int64, method string, p | |||||
| 250 | delete(t.pending, key) | 396 | delete(t.pending, key) | ||
| 251 | t.mu.Unlock() | 397 | t.mu.Unlock() | ||
| 252 | } | 398 | } | ||
| 253 | - | if err := t.write(frame); err != nil { | 399 | + | if err := t.write(ctx, frame); err != nil { |
| 254 | cleanup() | 400 | cleanup() | ||
| 401 | + | if ctx.Err() != nil { | |||
| 402 | + | t.cancelRequest(id, method, ctx.Err().Error()) | |||
| 403 | + | return nil, ctx.Err() | |||
| 404 | + | } | |||
| 255 | return nil, err | 405 | return nil, err | ||
| 256 | } | 406 | } | ||
| 257 | - | select { | 407 | + | result := func(msg message) (json.RawMessage, error) { |
| 258 | - | case msg := <-ch: | |||
| 259 | if msg.Error != nil { | 408 | if msg.Error != nil { | ||
| 260 | - | return nil, msg.Error | 409 | + | return nil, t.red.scrubRPCError(msg.Error) |
| 261 | } | 410 | } | ||
| 262 | return msg.Result, nil | 411 | return msg.Result, nil | ||
| 412 | + | } | |||
| 413 | + | select { | |||
| 414 | + | case msg := <-ch: | |||
| 415 | + | return result(msg) | |||
| 263 | case <-ctx.Done(): | 416 | case <-ctx.Done(): | ||
| 264 | cleanup() | 417 | cleanup() | ||
| 265 | - | _ = t.notify(context.Background(), "notifications/cancelled", map[string]any{"requestId": id, "reason": ctx.Err().Error()}) | 418 | + | t.cancelRequest(id, method, ctx.Err().Error()) |
| 266 | return nil, ctx.Err() | 419 | return nil, ctx.Err() | ||
| 267 | case <-t.doneCh: | 420 | case <-t.doneCh: | ||
| 421 | + | // A response that arrived together with the failure still counts. | |||
| 422 | + | select { | |||
| 423 | + | case msg := <-ch: | |||
| 424 | + | return result(msg) | |||
| 425 | + | default: | |||
| 426 | + | } | |||
| 268 | cleanup() | 427 | cleanup() | ||
| 269 | return nil, t.failure() | 428 | return nil, t.failure() | ||
| 270 | } | 429 | } | ||
| 271 | } | 430 | } | ||
| 272 | 431 | ||||
| 273 | - | func (t *stdioTransport) notify(_ context.Context, method string, params any) error { | 432 | + | // cancelRequest announces an abandoned request on a best-effort basis. It |
| 433 | + | // never blocks the caller, is bounded by cancelNotifyTimeout, and dies with | |||
| 434 | + | // the transport. | |||
| 435 | + | func (t *stdioTransport) cancelRequest(id int64, method, reason string) { | |||
| 436 | + | if !shouldSendCancel(method) { | |||
| 437 | + | return | |||
| 438 | + | } | |||
| 439 | + | t.cancelsSent.Add(1) | |||
| 440 | + | go func() { | |||
| 441 | + | ctx, cancel := context.WithTimeout(t.ctx, cancelNotifyTimeout) | |||
| 442 | + | defer cancel() | |||
| 443 | + | _ = t.notify(ctx, "notifications/cancelled", map[string]any{"requestId": id, "reason": reason}) | |||
| 444 | + | }() | |||
| 445 | + | } | |||
| 446 | + | ||||
| 447 | + | func (t *stdioTransport) notify(ctx context.Context, method string, params any) error { | |||
| 274 | frame, err := encodeNotification(method, params) | 448 | frame, err := encodeNotification(method, params) | ||
| 275 | if err != nil { | 449 | if err != nil { | ||
| 276 | return err | 450 | return err | ||
| 277 | } | 451 | } | ||
| 278 | - | return t.write(frame) | 452 | + | return t.write(ctx, frame) |
| 279 | } | 453 | } | ||
| 280 | 454 | ||||
| 281 | func (t *stdioTransport) setHandler(h incomingHandler) { | 455 | func (t *stdioTransport) setHandler(h incomingHandler) { | ||
| @@ -303,31 +477,65 @@ func (t *stdioTransport) fail(err error) { | |||||
| 303 | t.err = err | 477 | t.err = err | ||
| 304 | } | 478 | } | ||
| 305 | t.mu.Unlock() | 479 | t.mu.Unlock() | ||
| 306 | - | t.closeOnce.Do(func() { close(t.doneCh) }) | 480 | + | t.closeOnce.Do(func() { |
| 481 | + | t.cancel() | |||
| 482 | + | close(t.doneCh) | |||
| 483 | + | }) | |||
| 307 | } | 484 | } | ||
| 308 | 485 | ||||
| 309 | - | // close shuts the server down: close stdin (the MCP stdio shutdown signal), | 486 | + | func (t *stdioTransport) waitExit(d time.Duration) bool { |
| 310 | - | // wait briefly, then terminate the whole process group. | |||
| 311 | - | func (t *stdioTransport) close() error { | |||
| 312 | - | t.fail(errors.New("connection closed")) | |||
| 313 | - | t.writeMu.Lock() | |||
| 314 | - | _ = t.stdin.Close() | |||
| 315 | - | t.writeMu.Unlock() | |||
| 316 | select { | 487 | select { | ||
| 317 | case <-t.exited: | 488 | case <-t.exited: | ||
| 318 | - | return nil | 489 | + | return true |
| 319 | - | case <-time.After(500 * time.Millisecond): | 490 | + | case <-time.After(d): |
| 491 | + | return false | |||
| 320 | } | 492 | } | ||
| 321 | - | terminateProcess(t.cmd, false) | 493 | + | } |
| 322 | - | select { | 494 | + | |
| 323 | - | case <-t.exited: | 495 | + | // close shuts the server down: close stdin (the MCP stdio shutdown signal) |
| 324 | - | return nil | 496 | + | // without waiting behind any blocked write, wait briefly, then terminate the |
| 325 | - | case <-time.After(1500 * time.Millisecond): | 497 | + | // process group. Descendants that outlive the direct child are signalled too. |
| 498 | + | func (t *stdioTransport) close() error { | |||
| 499 | + | t.fail(errors.New("connection closed")) | |||
| 500 | + | _ = t.stdin.Close() // also unblocks a writer stuck on a full pipe | |||
| 501 | + | if !t.waitExit(500 * time.Millisecond) { | |||
| 502 | + | t.guard.terminate(false) | |||
| 503 | + | if !t.waitExit(1500 * time.Millisecond) { | |||
| 504 | + | t.guard.terminate(true) | |||
| 505 | + | t.waitExit(time.Second) | |||
| 506 | + | } | |||
| 326 | } | 507 | } | ||
| 327 | - | terminateProcess(t.cmd, true) | 508 | + | t.reapGroup() |
| 509 | + | // With the child gone its stderr hits EOF once buffered data is read; give | |||
| 510 | + | // the reader a moment so diagnostics are complete. A descendant holding the | |||
| 511 | + | // pipe open only costs the timeout. | |||
| 328 | select { | 512 | select { | ||
| 329 | case <-t.exited: | 513 | case <-t.exited: | ||
| 330 | - | case <-time.After(time.Second): | 514 | + | select { |
| 515 | + | case <-t.stderrDone: | |||
| 516 | + | case <-time.After(500 * time.Millisecond): | |||
| 517 | + | } | |||
| 518 | + | default: | |||
| 331 | } | 519 | } | ||
| 520 | + | _ = t.outR.Close() | |||
| 521 | + | _ = t.errR.Close() | |||
| 522 | + | t.guard.release() | |||
| 332 | return nil | 523 | return nil | ||
| 333 | } | 524 | } | ||
| 525 | + | ||||
| 526 | + | // reapGroup stops processes the direct child left behind. The process group | |||
| 527 | + | // outlives its leader, so this is meaningful even after the leader has been | |||
| 528 | + | // reaped (the group id cannot be reused while any member remains). | |||
| 529 | + | func (t *stdioTransport) reapGroup() { | |||
| 530 | + | if !t.guard.alive() { | |||
| 531 | + | return | |||
| 532 | + | } | |||
| 533 | + | t.guard.terminate(false) | |||
| 534 | + | deadline := time.Now().Add(300 * time.Millisecond) | |||
| 535 | + | for t.guard.alive() && time.Now().Before(deadline) { | |||
| 536 | + | time.Sleep(10 * time.Millisecond) | |||
| 537 | + | } | |||
| 538 | + | if t.guard.alive() { | |||
| 539 | + | t.guard.terminate(true) | |||
| 540 | + | } | |||
| 541 | + | } | |||
M packages/agent/mcp/tools.go +70 -22
| @@ -81,6 +81,9 @@ func (t *serverTool) Execute(ctx context.Context, args json.RawMessage, _ func(s | |||
| 81 | 81 | } | |
| 82 | 82 | ||
| 83 | 83 | func errorResult(text string) core.ToolResult { | |
| 84 | + | // Error text may embed server-controlled strings; keep it free of | |
| 85 | + | // terminal control sequences. | |
| 86 | + | text = sanitizeText(text, true) | |
| 84 | 87 | return core.ToolResult{IsError: true, Content: []provider.Content{provider.TextBlock{Text: text}}} | |
| 85 | 88 | } | |
| 86 | 89 | ||
| @@ -271,66 +274,110 @@ func (t *listResourcesTool) Execute(ctx context.Context, args json.RawMessage, _ | |||
| 271 | 274 | if in.Server != "" { | |
| 272 | 275 | servers = []string{in.Server} | |
| 273 | 276 | } | |
| 274 | - | var b strings.Builder | |
| 277 | + | // Output is budgeted while it is formatted: once maxResultTextBytes is | |
| 278 | + | // reached nothing more is built and no further server is contacted. | |
| 279 | + | out := &budgetedText{} | |
| 275 | 280 | for _, server := range servers { | |
| 281 | + | if out.full() { | |
| 282 | + | break | |
| 283 | + | } | |
| 276 | 284 | client, cfg, err := t.manager.client(ctx, server) | |
| 277 | 285 | if err != nil { | |
| 278 | - | fmt.Fprintf(&b, "## %s\nunavailable: %s\n\n", server, firstLine(err.Error())) | |
| 286 | + | out.printf("## %s\nunavailable: %s\n\n", sanitizeText(server, false), sanitizeText(firstLine(err.Error()), false)) | |
| 279 | 287 | continue | |
| 280 | 288 | } | |
| 281 | 289 | if !client.Info().HasResources { | |
| 282 | 290 | if in.Server != "" { | |
| 283 | - | fmt.Fprintf(&b, "## %s\nthis server does not offer resources\n\n", server) | |
| 291 | + | out.printf("## %s\nthis server does not offer resources\n\n", sanitizeText(server, false)) | |
| 284 | 292 | } | |
| 285 | 293 | continue | |
| 286 | 294 | } | |
| 287 | 295 | reqCtx, cancel := withTimeout(ctx, cfg.ToolTimeout()) | |
| 288 | 296 | resources, rErr := client.ListResources(reqCtx) | |
| 289 | - | templates, tErr := client.ListResourceTemplates(reqCtx) | |
| 290 | - | cancel() | |
| 291 | - | fmt.Fprintf(&b, "## %s\n", server) | |
| 297 | + | out.printf("## %s\n", sanitizeText(server, false)) | |
| 292 | 298 | if rErr != nil { | |
| 293 | - | fmt.Fprintf(&b, "resources error: %v\n", rErr) | |
| 299 | + | out.printf("resources error: %s\n", sanitizeText(rErr.Error(), false)) | |
| 294 | 300 | } | |
| 295 | 301 | for _, r := range resources { | |
| 296 | - | fmt.Fprintf(&b, "- %s", r.URI) | |
| 302 | + | if out.full() { | |
| 303 | + | break | |
| 304 | + | } | |
| 305 | + | var line strings.Builder | |
| 306 | + | fmt.Fprintf(&line, "- %s", sanitizeText(r.URI, false)) | |
| 297 | 307 | if name := firstNonEmpty(r.Title, r.Name); name != "" { | |
| 298 | - | fmt.Fprintf(&b, " — %s", name) | |
| 308 | + | fmt.Fprintf(&line, " — %s", sanitizeText(name, false)) | |
| 299 | 309 | } | |
| 300 | 310 | if r.MimeType != "" { | |
| 301 | - | fmt.Fprintf(&b, " (%s)", r.MimeType) | |
| 311 | + | fmt.Fprintf(&line, " (%s)", sanitizeText(r.MimeType, false)) | |
| 302 | 312 | } | |
| 303 | 313 | if r.Description != "" { | |
| 304 | - | fmt.Fprintf(&b, ": %s", r.Description) | |
| 314 | + | fmt.Fprintf(&line, ": %s", sanitizeText(r.Description, false)) | |
| 305 | 315 | } | |
| 306 | - | b.WriteByte('\n') | |
| 316 | + | line.WriteByte('\n') | |
| 317 | + | out.write(line.String()) | |
| 318 | + | } | |
| 319 | + | // Templates are only fetched while there is budget left for them. | |
| 320 | + | var templates []ResourceTemplate | |
| 321 | + | var tErr error | |
| 322 | + | if !out.full() { | |
| 323 | + | templates, tErr = client.ListResourceTemplates(reqCtx) | |
| 307 | 324 | } | |
| 325 | + | cancel() | |
| 308 | 326 | if tErr != nil && !isMethodNotFound(tErr) { | |
| 309 | - | fmt.Fprintf(&b, "templates error: %v\n", tErr) | |
| 327 | + | out.printf("templates error: %s\n", sanitizeText(tErr.Error(), false)) | |
| 310 | 328 | } | |
| 311 | 329 | for _, tpl := range templates { | |
| 312 | - | fmt.Fprintf(&b, "- template %s", tpl.URITemplate) | |
| 330 | + | if out.full() { | |
| 331 | + | break | |
| 332 | + | } | |
| 333 | + | var line strings.Builder | |
| 334 | + | fmt.Fprintf(&line, "- template %s", sanitizeText(tpl.URITemplate, false)) | |
| 313 | 335 | if name := firstNonEmpty(tpl.Title, tpl.Name); name != "" { | |
| 314 | - | fmt.Fprintf(&b, " — %s", name) | |
| 336 | + | fmt.Fprintf(&line, " — %s", sanitizeText(name, false)) | |
| 315 | 337 | } | |
| 316 | 338 | if tpl.Description != "" { | |
| 317 | - | fmt.Fprintf(&b, ": %s", tpl.Description) | |
| 339 | + | fmt.Fprintf(&line, ": %s", sanitizeText(tpl.Description, false)) | |
| 318 | 340 | } | |
| 319 | - | b.WriteByte('\n') | |
| 341 | + | line.WriteByte('\n') | |
| 342 | + | out.write(line.String()) | |
| 320 | 343 | } | |
| 321 | 344 | if len(resources) == 0 && len(templates) == 0 && rErr == nil { | |
| 322 | - | b.WriteString("(no resources)\n") | |
| 345 | + | out.write("(no resources)\n") | |
| 323 | 346 | } | |
| 324 | - | b.WriteByte('\n') | |
| 347 | + | out.write("\n") | |
| 325 | 348 | } | |
| 326 | - | text := strings.TrimSpace(b.String()) | |
| 349 | + | text := strings.TrimSpace(out.String()) | |
| 327 | 350 | if text == "" { | |
| 328 | 351 | text = "no connected MCP server offers resources" | |
| 329 | 352 | } | |
| 330 | - | budget := &textBudget{} | |
| 331 | - | return core.ToolResult{Content: []provider.Content{provider.TextBlock{Text: budget.take(text)}}}, nil | |
| 353 | + | return core.ToolResult{Content: []provider.Content{provider.TextBlock{Text: text}}}, nil | |
| 332 | 354 | } | |
| 333 | 355 | ||
| 356 | + | // budgetedText accumulates text up to maxResultTextBytes, then ignores | |
| 357 | + | // further writes and appends a truncation marker once. | |
| 358 | + | type budgetedText struct { | |
| 359 | + | b strings.Builder | |
| 360 | + | budget textBudget | |
| 361 | + | } | |
| 362 | + | ||
| 363 | + | func (t *budgetedText) full() bool { return t.budget.truncated } | |
| 364 | + | ||
| 365 | + | func (t *budgetedText) write(s string) { | |
| 366 | + | if t.budget.truncated { | |
| 367 | + | return | |
| 368 | + | } | |
| 369 | + | t.b.WriteString(t.budget.take(s)) | |
| 370 | + | } | |
| 371 | + | ||
| 372 | + | func (t *budgetedText) printf(format string, args ...any) { | |
| 373 | + | if t.budget.truncated { | |
| 374 | + | return | |
| 375 | + | } | |
| 376 | + | t.write(fmt.Sprintf(format, args...)) | |
| 377 | + | } | |
| 378 | + | ||
| 379 | + | func (t *budgetedText) String() string { return t.b.String() } | |
| 380 | + | ||
| 334 | 381 | func isMethodNotFound(err error) bool { | |
| 335 | 382 | var rpcErr *RPCError | |
| 336 | 383 | return errors.As(err, &rpcErr) && rpcErr.Code == codeMethodNotFound | |
| @@ -401,6 +448,7 @@ func (m *Manager) SystemPromptAddendum() string { | |||
| 401 | 448 | } | |
| 402 | 449 | b.WriteByte('\n') | |
| 403 | 450 | if instr := strings.TrimSpace(st.Info.Instructions); instr != "" { | |
| 451 | + | instr = sanitizeText(instr, true) | |
| 404 | 452 | if len(instr) > maxDescriptionBytes { | |
| 405 | 453 | instr = instr[:maxDescriptionBytes] + "…" | |
| 406 | 454 | } | |
| @@ -81,6 +81,9 @@ func (t *serverTool) Execute(ctx context.Context, args json.RawMessage, _ func(s | |||||
| 81 | } | 81 | } | ||
| 82 | 82 | ||||
| 83 | func errorResult(text string) core.ToolResult { | 83 | func errorResult(text string) core.ToolResult { | ||
| 84 | + | // Error text may embed server-controlled strings; keep it free of | |||
| 85 | + | // terminal control sequences. | |||
| 86 | + | text = sanitizeText(text, true) | |||
| 84 | return core.ToolResult{IsError: true, Content: []provider.Content{provider.TextBlock{Text: text}}} | 87 | return core.ToolResult{IsError: true, Content: []provider.Content{provider.TextBlock{Text: text}}} | ||
| 85 | } | 88 | } | ||
| 86 | 89 | ||||
| @@ -271,66 +274,110 @@ func (t *listResourcesTool) Execute(ctx context.Context, args json.RawMessage, _ | |||||
| 271 | if in.Server != "" { | 274 | if in.Server != "" { | ||
| 272 | servers = []string{in.Server} | 275 | servers = []string{in.Server} | ||
| 273 | } | 276 | } | ||
| 274 | - | var b strings.Builder | 277 | + | // Output is budgeted while it is formatted: once maxResultTextBytes is |
| 278 | + | // reached nothing more is built and no further server is contacted. | |||
| 279 | + | out := &budgetedText{} | |||
| 275 | for _, server := range servers { | 280 | for _, server := range servers { | ||
| 281 | + | if out.full() { | |||
| 282 | + | break | |||
| 283 | + | } | |||
| 276 | client, cfg, err := t.manager.client(ctx, server) | 284 | client, cfg, err := t.manager.client(ctx, server) | ||
| 277 | if err != nil { | 285 | if err != nil { | ||
| 278 | - | fmt.Fprintf(&b, "## %s\nunavailable: %s\n\n", server, firstLine(err.Error())) | 286 | + | out.printf("## %s\nunavailable: %s\n\n", sanitizeText(server, false), sanitizeText(firstLine(err.Error()), false)) |
| 279 | continue | 287 | continue | ||
| 280 | } | 288 | } | ||
| 281 | if !client.Info().HasResources { | 289 | if !client.Info().HasResources { | ||
| 282 | if in.Server != "" { | 290 | if in.Server != "" { | ||
| 283 | - | fmt.Fprintf(&b, "## %s\nthis server does not offer resources\n\n", server) | 291 | + | out.printf("## %s\nthis server does not offer resources\n\n", sanitizeText(server, false)) |
| 284 | } | 292 | } | ||
| 285 | continue | 293 | continue | ||
| 286 | } | 294 | } | ||
| 287 | reqCtx, cancel := withTimeout(ctx, cfg.ToolTimeout()) | 295 | reqCtx, cancel := withTimeout(ctx, cfg.ToolTimeout()) | ||
| 288 | resources, rErr := client.ListResources(reqCtx) | 296 | resources, rErr := client.ListResources(reqCtx) | ||
| 289 | - | templates, tErr := client.ListResourceTemplates(reqCtx) | 297 | + | out.printf("## %s\n", sanitizeText(server, false)) |
| 290 | - | cancel() | |||
| 291 | - | fmt.Fprintf(&b, "## %s\n", server) | |||
| 292 | if rErr != nil { | 298 | if rErr != nil { | ||
| 293 | - | fmt.Fprintf(&b, "resources error: %v\n", rErr) | 299 | + | out.printf("resources error: %s\n", sanitizeText(rErr.Error(), false)) |
| 294 | } | 300 | } | ||
| 295 | for _, r := range resources { | 301 | for _, r := range resources { | ||
| 296 | - | fmt.Fprintf(&b, "- %s", r.URI) | 302 | + | if out.full() { |
| 303 | + | break | |||
| 304 | + | } | |||
| 305 | + | var line strings.Builder | |||
| 306 | + | fmt.Fprintf(&line, "- %s", sanitizeText(r.URI, false)) | |||
| 297 | if name := firstNonEmpty(r.Title, r.Name); name != "" { | 307 | if name := firstNonEmpty(r.Title, r.Name); name != "" { | ||
| 298 | - | fmt.Fprintf(&b, " — %s", name) | 308 | + | fmt.Fprintf(&line, " — %s", sanitizeText(name, false)) |
| 299 | } | 309 | } | ||
| 300 | if r.MimeType != "" { | 310 | if r.MimeType != "" { | ||
| 301 | - | fmt.Fprintf(&b, " (%s)", r.MimeType) | 311 | + | fmt.Fprintf(&line, " (%s)", sanitizeText(r.MimeType, false)) |
| 302 | } | 312 | } | ||
| 303 | if r.Description != "" { | 313 | if r.Description != "" { | ||
| 304 | - | fmt.Fprintf(&b, ": %s", r.Description) | 314 | + | fmt.Fprintf(&line, ": %s", sanitizeText(r.Description, false)) |
| 305 | } | 315 | } | ||
| 306 | - | b.WriteByte('\n') | 316 | + | line.WriteByte('\n') |
| 317 | + | out.write(line.String()) | |||
| 318 | + | } | |||
| 319 | + | // Templates are only fetched while there is budget left for them. | |||
| 320 | + | var templates []ResourceTemplate | |||
| 321 | + | var tErr error | |||
| 322 | + | if !out.full() { | |||
| 323 | + | templates, tErr = client.ListResourceTemplates(reqCtx) | |||
| 307 | } | 324 | } | ||
| 325 | + | cancel() | |||
| 308 | if tErr != nil && !isMethodNotFound(tErr) { | 326 | if tErr != nil && !isMethodNotFound(tErr) { | ||
| 309 | - | fmt.Fprintf(&b, "templates error: %v\n", tErr) | 327 | + | out.printf("templates error: %s\n", sanitizeText(tErr.Error(), false)) |
| 310 | } | 328 | } | ||
| 311 | for _, tpl := range templates { | 329 | for _, tpl := range templates { | ||
| 312 | - | fmt.Fprintf(&b, "- template %s", tpl.URITemplate) | 330 | + | if out.full() { |
| 331 | + | break | |||
| 332 | + | } | |||
| 333 | + | var line strings.Builder | |||
| 334 | + | fmt.Fprintf(&line, "- template %s", sanitizeText(tpl.URITemplate, false)) | |||
| 313 | if name := firstNonEmpty(tpl.Title, tpl.Name); name != "" { | 335 | if name := firstNonEmpty(tpl.Title, tpl.Name); name != "" { | ||
| 314 | - | fmt.Fprintf(&b, " — %s", name) | 336 | + | fmt.Fprintf(&line, " — %s", sanitizeText(name, false)) |
| 315 | } | 337 | } | ||
| 316 | if tpl.Description != "" { | 338 | if tpl.Description != "" { | ||
| 317 | - | fmt.Fprintf(&b, ": %s", tpl.Description) | 339 | + | fmt.Fprintf(&line, ": %s", sanitizeText(tpl.Description, false)) |
| 318 | } | 340 | } | ||
| 319 | - | b.WriteByte('\n') | 341 | + | line.WriteByte('\n') |
| 342 | + | out.write(line.String()) | |||
| 320 | } | 343 | } | ||
| 321 | if len(resources) == 0 && len(templates) == 0 && rErr == nil { | 344 | if len(resources) == 0 && len(templates) == 0 && rErr == nil { | ||
| 322 | - | b.WriteString("(no resources)\n") | 345 | + | out.write("(no resources)\n") |
| 323 | } | 346 | } | ||
| 324 | - | b.WriteByte('\n') | 347 | + | out.write("\n") |
| 325 | } | 348 | } | ||
| 326 | - | text := strings.TrimSpace(b.String()) | 349 | + | text := strings.TrimSpace(out.String()) |
| 327 | if text == "" { | 350 | if text == "" { | ||
| 328 | text = "no connected MCP server offers resources" | 351 | text = "no connected MCP server offers resources" | ||
| 329 | } | 352 | } | ||
| 330 | - | budget := &textBudget{} | 353 | + | return core.ToolResult{Content: []provider.Content{provider.TextBlock{Text: text}}}, nil |
| 331 | - | return core.ToolResult{Content: []provider.Content{provider.TextBlock{Text: budget.take(text)}}}, nil | |||
| 332 | } | 354 | } | ||
| 333 | 355 | ||||
| 356 | + | // budgetedText accumulates text up to maxResultTextBytes, then ignores | |||
| 357 | + | // further writes and appends a truncation marker once. | |||
| 358 | + | type budgetedText struct { | |||
| 359 | + | b strings.Builder | |||
| 360 | + | budget textBudget | |||
| 361 | + | } | |||
| 362 | + | ||||
| 363 | + | func (t *budgetedText) full() bool { return t.budget.truncated } | |||
| 364 | + | ||||
| 365 | + | func (t *budgetedText) write(s string) { | |||
| 366 | + | if t.budget.truncated { | |||
| 367 | + | return | |||
| 368 | + | } | |||
| 369 | + | t.b.WriteString(t.budget.take(s)) | |||
| 370 | + | } | |||
| 371 | + | ||||
| 372 | + | func (t *budgetedText) printf(format string, args ...any) { | |||
| 373 | + | if t.budget.truncated { | |||
| 374 | + | return | |||
| 375 | + | } | |||
| 376 | + | t.write(fmt.Sprintf(format, args...)) | |||
| 377 | + | } | |||
| 378 | + | ||||
| 379 | + | func (t *budgetedText) String() string { return t.b.String() } | |||
| 380 | + | ||||
| 334 | func isMethodNotFound(err error) bool { | 381 | func isMethodNotFound(err error) bool { | ||
| 335 | var rpcErr *RPCError | 382 | var rpcErr *RPCError | ||
| 336 | return errors.As(err, &rpcErr) && rpcErr.Code == codeMethodNotFound | 383 | return errors.As(err, &rpcErr) && rpcErr.Code == codeMethodNotFound | ||
| @@ -401,6 +448,7 @@ func (m *Manager) SystemPromptAddendum() string { | |||||
| 401 | } | 448 | } | ||
| 402 | b.WriteByte('\n') | 449 | b.WriteByte('\n') | ||
| 403 | if instr := strings.TrimSpace(st.Info.Instructions); instr != "" { | 450 | if instr := strings.TrimSpace(st.Info.Instructions); instr != "" { | ||
| 451 | + | instr = sanitizeText(instr, true) | |||
| 404 | if len(instr) > maxDescriptionBytes { | 452 | if len(instr) > maxDescriptionBytes { | ||
| 405 | instr = instr[:maxDescriptionBytes] + "…" | 453 | instr = instr[:maxDescriptionBytes] + "…" | ||
| 406 | } | 454 | } | ||
M packages/agent/mcp_runtime.go +2 -6
| @@ -97,14 +97,10 @@ func addMCPTools(reg core.Registry, args Args) { | |||
| 97 | 97 | strictOrchestrator := args.Orchestrate && !args.ResidentChild | |
| 98 | 98 | for _, tool := range args.mcpManager.Tools() { | |
| 99 | 99 | name := tool.Name() | |
| 100 | - | if _, exists := reg[name]; exists { | |
| 100 | + | if _, exists := reg[name]; exists || !mcpToolSelected(args, name) { | |
| 101 | 101 | continue | |
| 102 | 102 | } | |
| 103 | - | if strictOrchestrator { | |
| 104 | - | if !mcp.IsReadOnlyTool(tool) { | |
| 105 | - | continue | |
| 106 | - | } | |
| 107 | - | } else if !mcpToolSelected(args, name) { | |
| 103 | + | if strictOrchestrator && !mcp.IsReadOnlyTool(tool) { | |
| 108 | 104 | continue | |
| 109 | 105 | } | |
| 110 | 106 | reg[name] = tool | |
| @@ -97,14 +97,10 @@ func addMCPTools(reg core.Registry, args Args) { | |||||
| 97 | strictOrchestrator := args.Orchestrate && !args.ResidentChild | 97 | strictOrchestrator := args.Orchestrate && !args.ResidentChild | ||
| 98 | for _, tool := range args.mcpManager.Tools() { | 98 | for _, tool := range args.mcpManager.Tools() { | ||
| 99 | name := tool.Name() | 99 | name := tool.Name() | ||
| 100 | - | if _, exists := reg[name]; exists { | 100 | + | if _, exists := reg[name]; exists || !mcpToolSelected(args, name) { |
| 101 | continue | 101 | continue | ||
| 102 | } | 102 | } | ||
| 103 | - | if strictOrchestrator { | 103 | + | if strictOrchestrator && !mcp.IsReadOnlyTool(tool) { |
| 104 | - | if !mcp.IsReadOnlyTool(tool) { | |||
| 105 | - | continue | |||
| 106 | - | } | |||
| 107 | - | } else if !mcpToolSelected(args, name) { | |||
| 108 | continue | 104 | continue | ||
| 109 | } | 105 | } | ||
| 110 | reg[name] = tool | 106 | reg[name] = tool | ||
M packages/agent/mcp_runtime_test.go +48 -0
| @@ -172,6 +172,13 @@ func TestOrchestratorHoldsReadOnlyMCPToolsAndWorkersGetAll(t *testing.T) { | |||
| 172 | 172 | if _, ok := r.ToolRegistry["mcp__fake__deploy"]; ok { | |
| 173 | 173 | t.Fatal("orchestrator holds a mutating MCP tool") | |
| 174 | 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 | + | } | |
| 175 | 182 | if !strings.Contains(r.SystemPrompt, "delegate MCP actions") { | |
| 176 | 183 | t.Fatalf("orchestrator prompt lacks MCP delegation guidance:\n%s", r.SystemPrompt) | |
| 177 | 184 | } | |
| @@ -311,6 +318,47 @@ func TestMCPCommandAddListRemove(t *testing.T) { | |||
| 311 | 318 | } | |
| 312 | 319 | } | |
| 313 | 320 | ||
| 321 | + | func 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 | + | ||
| 348 | + | func TestMCPCommandFailsOnMalformedConfig(t *testing.T) { | |
| 349 | + | root := t.TempDir() | |
| 350 | + | t.Setenv("ZUT_HOME", root) | |
| 351 | + | if err := os.WriteFile(mcp.GlobalConfigPath(root), []byte("{not json"), 0o600); err != nil { | |
| 352 | + | t.Fatal(err) | |
| 353 | + | } | |
| 354 | + | for _, args := range [][]string{{"list"}, {"list", "--json"}, {"status"}} { | |
| 355 | + | var out strings.Builder | |
| 356 | + | if _, err := runMCPCommand(context.Background(), append([]string{"mcp"}, args...), "test", &out); err == nil { | |
| 357 | + | t.Fatalf("%v succeeded on malformed config:\n%s", args, out.String()) | |
| 358 | + | } | |
| 359 | + | } | |
| 360 | + | } | |
| 361 | + | ||
| 314 | 362 | func TestParseArgsNoMCP(t *testing.T) { | |
| 315 | 363 | args, err := ParseArgs([]string{"--no-mcp", "-p", "hi"}) | |
| 316 | 364 | if err != nil { | |
| @@ -172,6 +172,13 @@ func TestOrchestratorHoldsReadOnlyMCPToolsAndWorkersGetAll(t *testing.T) { | |||||
| 172 | if _, ok := r.ToolRegistry["mcp__fake__deploy"]; ok { | 172 | if _, ok := r.ToolRegistry["mcp__fake__deploy"]; ok { | ||
| 173 | t.Fatal("orchestrator holds a mutating MCP tool") | 173 | t.Fatal("orchestrator holds a mutating MCP tool") | ||
| 174 | } | 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 | + | } | |||
| 175 | if !strings.Contains(r.SystemPrompt, "delegate MCP actions") { | 182 | if !strings.Contains(r.SystemPrompt, "delegate MCP actions") { | ||
| 176 | t.Fatalf("orchestrator prompt lacks MCP delegation guidance:\n%s", r.SystemPrompt) | 183 | t.Fatalf("orchestrator prompt lacks MCP delegation guidance:\n%s", r.SystemPrompt) | ||
| 177 | } | 184 | } | ||
| @@ -311,6 +318,47 @@ func TestMCPCommandAddListRemove(t *testing.T) { | |||||
| 311 | } | 318 | } | ||
| 312 | } | 319 | } | ||
| 313 | 320 | ||||
| 321 | + | func 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 | + | ||||
| 348 | + | func TestMCPCommandFailsOnMalformedConfig(t *testing.T) { | |||
| 349 | + | root := t.TempDir() | |||
| 350 | + | t.Setenv("ZUT_HOME", root) | |||
| 351 | + | if err := os.WriteFile(mcp.GlobalConfigPath(root), []byte("{not json"), 0o600); err != nil { | |||
| 352 | + | t.Fatal(err) | |||
| 353 | + | } | |||
| 354 | + | for _, args := range [][]string{{"list"}, {"list", "--json"}, {"status"}} { | |||
| 355 | + | var out strings.Builder | |||
| 356 | + | if _, err := runMCPCommand(context.Background(), append([]string{"mcp"}, args...), "test", &out); err == nil { | |||
| 357 | + | t.Fatalf("%v succeeded on malformed config:\n%s", args, out.String()) | |||
| 358 | + | } | |||
| 359 | + | } | |||
| 360 | + | } | |||
| 361 | + | ||||
| 314 | func TestParseArgsNoMCP(t *testing.T) { | 362 | func TestParseArgsNoMCP(t *testing.T) { | ||
| 315 | args, err := ParseArgs([]string{"--no-mcp", "-p", "hi"}) | 363 | args, err := ParseArgs([]string{"--no-mcp", "-p", "hi"}) | ||
| 316 | if err != nil { | 364 | if err != nil { | ||
M packages/agent/mcpcmd.go +55 -17
| @@ -7,6 +7,7 @@ import ( | |||
| 7 | 7 | "fmt" | |
| 8 | 8 | "io" | |
| 9 | 9 | "io/fs" | |
| 10 | + | "net/url" | |
| 10 | 11 | "os" | |
| 11 | 12 | "sort" | |
| 12 | 13 | "strconv" | |
| @@ -80,9 +81,33 @@ func mcpConfigPath(cwd string, project bool) string { | |||
| 80 | 81 | return mcp.GlobalConfigPath(ZutHome()) | |
| 81 | 82 | } | |
| 82 | 83 | ||
| 84 | + | // rawMCPConfig merges the global and project files without ${VAR} expansion | |
| 85 | + | // or validation, so list/get show disabled and currently broken entries and | |
| 86 | + | // never print expanded secret values. | |
| 87 | + | func rawMCPConfig(cwd string) (mcp.Config, error) { | |
| 88 | + | cfg := mcp.Config{Servers: map[string]mcp.ServerConfig{}} | |
| 89 | + | var errs []error | |
| 90 | + | for _, path := range []string{mcp.GlobalConfigPath(ZutHome()), mcp.ProjectConfigPath(cwd)} { | |
| 91 | + | file, err := mcp.ReadConfigFile(path) | |
| 92 | + | if err != nil { | |
| 93 | + | if !errors.Is(err, fs.ErrNotExist) { | |
| 94 | + | errs = append(errs, err) | |
| 95 | + | } | |
| 96 | + | continue | |
| 97 | + | } | |
| 98 | + | for name, srv := range file.Servers { | |
| 99 | + | cfg.Servers[name] = srv | |
| 100 | + | } | |
| 101 | + | } | |
| 102 | + | return cfg, errors.Join(errs...) | |
| 103 | + | } | |
| 104 | + | ||
| 83 | 105 | func mcpList(out io.Writer, cwd string, args []string) error { | |
| 84 | 106 | asJSON := len(args) > 0 && args[0] == "--json" | |
| 85 | - | cfg, errs := mcp.LoadConfig(ZutHome(), cwd) | |
| 107 | + | cfg, cfgErr := rawMCPConfig(cwd) | |
| 108 | + | if cfgErr != nil { | |
| 109 | + | return cfgErr | |
| 110 | + | } | |
| 86 | 111 | if asJSON { | |
| 87 | 112 | redacted := map[string]mcp.ServerConfig{} | |
| 88 | 113 | for name, srv := range cfg.Servers { | |
| @@ -101,21 +126,31 @@ func mcpList(out io.Writer, cwd string, args []string) error { | |||
| 101 | 126 | fmt.Fprintln(out, "no MCP servers configured") | |
| 102 | 127 | } | |
| 103 | 128 | for _, name := range names { | |
| 104 | - | srv := cfg.Servers[name] | |
| 129 | + | srv := redactServer(cfg.Servers[name]) | |
| 105 | 130 | target := srv.URL | |
| 106 | 131 | if srv.Kind() == "stdio" { | |
| 107 | 132 | target = strings.TrimSpace(srv.Command + " " + strings.Join(srv.Args, " ")) | |
| 108 | 133 | } | |
| 109 | - | fmt.Fprintf(out, "%s\t%s\t%s\n", name, srv.Kind(), target) | |
| 110 | - | } | |
| 111 | - | for _, err := range errs { | |
| 112 | - | fmt.Fprintln(os.Stderr, "mcp:", err) | |
| 134 | + | state := "" | |
| 135 | + | if srv.Enabled != nil && !*srv.Enabled { | |
| 136 | + | state = "\tdisabled" | |
| 137 | + | } | |
| 138 | + | fmt.Fprintf(out, "%s\t%s\t%s%s\n", name, srv.Kind(), target, state) | |
| 113 | 139 | } | |
| 114 | 140 | return nil | |
| 115 | 141 | } | |
| 116 | 142 | ||
| 117 | - | // redactServer hides values that commonly carry secrets. | |
| 143 | + | // redactServer hides values that commonly carry secrets: env and header | |
| 144 | + | // values, URL credentials and query strings. Arguments are shown as written | |
| 145 | + | // in the file (${VAR} references stay unexpanded). | |
| 118 | 146 | func redactServer(srv mcp.ServerConfig) mcp.ServerConfig { | |
| 147 | + | if u, err := url.Parse(srv.URL); err == nil && (u.User != nil || u.RawQuery != "") { | |
| 148 | + | u.User = nil | |
| 149 | + | if u.RawQuery != "" { | |
| 150 | + | u.RawQuery = "*****" | |
| 151 | + | } | |
| 152 | + | srv.URL = u.String() | |
| 153 | + | } | |
| 119 | 154 | hide := func(m map[string]string) map[string]string { | |
| 120 | 155 | if len(m) == 0 { | |
| 121 | 156 | return m | |
| @@ -135,10 +170,13 @@ func mcpGet(out io.Writer, cwd string, args []string) error { | |||
| 135 | 170 | if len(args) == 0 { | |
| 136 | 171 | return errors.New("usage: zut mcp get NAME [--json]") | |
| 137 | 172 | } | |
| 138 | - | cfg, _ := mcp.LoadConfig(ZutHome(), cwd) | |
| 173 | + | cfg, err := rawMCPConfig(cwd) | |
| 174 | + | if err != nil { | |
| 175 | + | return err | |
| 176 | + | } | |
| 139 | 177 | srv, ok := cfg.Servers[args[0]] | |
| 140 | 178 | if !ok { | |
| 141 | - | return fmt.Errorf("mcp server %q is not configured or is invalid", args[0]) | |
| 179 | + | return fmt.Errorf("mcp server %q is not configured", args[0]) | |
| 142 | 180 | } | |
| 143 | 181 | enc := json.NewEncoder(out) | |
| 144 | 182 | enc.SetIndent("", " ") | |
| @@ -293,12 +331,12 @@ func mcpRemove(out io.Writer, cwd string, args []string) error { | |||
| 293 | 331 | func mcpTest(ctx context.Context, out io.Writer, cwd, version string, args []string) error { | |
| 294 | 332 | cfg, errs := mcp.LoadConfig(ZutHome(), cwd) | |
| 295 | 333 | for _, err := range errs { | |
| 296 | - | fmt.Fprintln(out, "config error:", err) | |
| 334 | + | fmt.Fprintln(os.Stderr, "mcp config error:", err) | |
| 297 | 335 | } | |
| 298 | 336 | if len(args) > 0 { | |
| 299 | 337 | srv, ok := cfg.Servers[args[0]] | |
| 300 | 338 | if !ok { | |
| 301 | - | return fmt.Errorf("mcp server %q is not configured or is invalid", args[0]) | |
| 339 | + | return fmt.Errorf("mcp server %q is not configured, disabled, or invalid", args[0]) | |
| 302 | 340 | } | |
| 303 | 341 | cfg.Servers = map[string]mcp.ServerConfig{args[0]: srv} | |
| 304 | 342 | } | |
| @@ -306,13 +344,13 @@ func mcpTest(ctx context.Context, out io.Writer, cwd, version string, args []str | |||
| 306 | 344 | defer manager.Close() | |
| 307 | 345 | failures, _ := manager.Start(ctx) | |
| 308 | 346 | fmt.Fprintln(out, mcp.FormatStatus(manager.Status())) | |
| 309 | - | for _, server := range manager.Servers() { | |
| 310 | - | for _, tool := range manager.ServerTools(server) { | |
| 311 | - | fmt.Fprintf(out, " %s\n", mcp.QualifiedName(server, tool.Name)) | |
| 312 | - | } | |
| 347 | + | // Names come from the registry builder so collision suffixes match what | |
| 348 | + | // --tools and the model see. | |
| 349 | + | for _, tool := range manager.Tools() { | |
| 350 | + | fmt.Fprintf(out, " %s\n", tool.Name()) | |
| 313 | 351 | } | |
| 314 | - | if len(failures) > 0 { | |
| 315 | - | return fmt.Errorf("%d MCP server(s) failed to start", len(failures)) | |
| 352 | + | if n := len(failures) + len(errs); n > 0 { | |
| 353 | + | return fmt.Errorf("%d MCP configuration or startup problem(s)", n) | |
| 316 | 354 | } | |
| 317 | 355 | return nil | |
| 318 | 356 | } | |
| @@ -7,6 +7,7 @@ import ( | |||||
| 7 | "fmt" | 7 | "fmt" | ||
| 8 | "io" | 8 | "io" | ||
| 9 | "io/fs" | 9 | "io/fs" | ||
| 10 | + | "net/url" | |||
| 10 | "os" | 11 | "os" | ||
| 11 | "sort" | 12 | "sort" | ||
| 12 | "strconv" | 13 | "strconv" | ||
| @@ -80,9 +81,33 @@ func mcpConfigPath(cwd string, project bool) string { | |||||
| 80 | return mcp.GlobalConfigPath(ZutHome()) | 81 | return mcp.GlobalConfigPath(ZutHome()) | ||
| 81 | } | 82 | } | ||
| 82 | 83 | ||||
| 84 | + | // rawMCPConfig merges the global and project files without ${VAR} expansion | |||
| 85 | + | // or validation, so list/get show disabled and currently broken entries and | |||
| 86 | + | // never print expanded secret values. | |||
| 87 | + | func rawMCPConfig(cwd string) (mcp.Config, error) { | |||
| 88 | + | cfg := mcp.Config{Servers: map[string]mcp.ServerConfig{}} | |||
| 89 | + | var errs []error | |||
| 90 | + | for _, path := range []string{mcp.GlobalConfigPath(ZutHome()), mcp.ProjectConfigPath(cwd)} { | |||
| 91 | + | file, err := mcp.ReadConfigFile(path) | |||
| 92 | + | if err != nil { | |||
| 93 | + | if !errors.Is(err, fs.ErrNotExist) { | |||
| 94 | + | errs = append(errs, err) | |||
| 95 | + | } | |||
| 96 | + | continue | |||
| 97 | + | } | |||
| 98 | + | for name, srv := range file.Servers { | |||
| 99 | + | cfg.Servers[name] = srv | |||
| 100 | + | } | |||
| 101 | + | } | |||
| 102 | + | return cfg, errors.Join(errs...) | |||
| 103 | + | } | |||
| 104 | + | ||||
| 83 | func mcpList(out io.Writer, cwd string, args []string) error { | 105 | func mcpList(out io.Writer, cwd string, args []string) error { | ||
| 84 | asJSON := len(args) > 0 && args[0] == "--json" | 106 | asJSON := len(args) > 0 && args[0] == "--json" | ||
| 85 | - | cfg, errs := mcp.LoadConfig(ZutHome(), cwd) | 107 | + | cfg, cfgErr := rawMCPConfig(cwd) |
| 108 | + | if cfgErr != nil { | |||
| 109 | + | return cfgErr | |||
| 110 | + | } | |||
| 86 | if asJSON { | 111 | if asJSON { | ||
| 87 | redacted := map[string]mcp.ServerConfig{} | 112 | redacted := map[string]mcp.ServerConfig{} | ||
| 88 | for name, srv := range cfg.Servers { | 113 | for name, srv := range cfg.Servers { | ||
| @@ -101,21 +126,31 @@ func mcpList(out io.Writer, cwd string, args []string) error { | |||||
| 101 | fmt.Fprintln(out, "no MCP servers configured") | 126 | fmt.Fprintln(out, "no MCP servers configured") | ||
| 102 | } | 127 | } | ||
| 103 | for _, name := range names { | 128 | for _, name := range names { | ||
| 104 | - | srv := cfg.Servers[name] | 129 | + | srv := redactServer(cfg.Servers[name]) |
| 105 | target := srv.URL | 130 | target := srv.URL | ||
| 106 | if srv.Kind() == "stdio" { | 131 | if srv.Kind() == "stdio" { | ||
| 107 | target = strings.TrimSpace(srv.Command + " " + strings.Join(srv.Args, " ")) | 132 | target = strings.TrimSpace(srv.Command + " " + strings.Join(srv.Args, " ")) | ||
| 108 | } | 133 | } | ||
| 109 | - | fmt.Fprintf(out, "%s\t%s\t%s\n", name, srv.Kind(), target) | 134 | + | state := "" |
| 110 | - | } | 135 | + | if srv.Enabled != nil && !*srv.Enabled { |
| 111 | - | for _, err := range errs { | 136 | + | state = "\tdisabled" |
| 112 | - | fmt.Fprintln(os.Stderr, "mcp:", err) | 137 | + | } |
| 138 | + | fmt.Fprintf(out, "%s\t%s\t%s%s\n", name, srv.Kind(), target, state) | |||
| 113 | } | 139 | } | ||
| 114 | return nil | 140 | return nil | ||
| 115 | } | 141 | } | ||
| 116 | 142 | ||||
| 117 | - | // redactServer hides values that commonly carry secrets. | 143 | + | // redactServer hides values that commonly carry secrets: env and header |
| 144 | + | // values, URL credentials and query strings. Arguments are shown as written | |||
| 145 | + | // in the file (${VAR} references stay unexpanded). | |||
| 118 | func redactServer(srv mcp.ServerConfig) mcp.ServerConfig { | 146 | func redactServer(srv mcp.ServerConfig) mcp.ServerConfig { | ||
| 147 | + | if u, err := url.Parse(srv.URL); err == nil && (u.User != nil || u.RawQuery != "") { | |||
| 148 | + | u.User = nil | |||
| 149 | + | if u.RawQuery != "" { | |||
| 150 | + | u.RawQuery = "*****" | |||
| 151 | + | } | |||
| 152 | + | srv.URL = u.String() | |||
| 153 | + | } | |||
| 119 | hide := func(m map[string]string) map[string]string { | 154 | hide := func(m map[string]string) map[string]string { | ||
| 120 | if len(m) == 0 { | 155 | if len(m) == 0 { | ||
| 121 | return m | 156 | return m | ||
| @@ -135,10 +170,13 @@ func mcpGet(out io.Writer, cwd string, args []string) error { | |||||
| 135 | if len(args) == 0 { | 170 | if len(args) == 0 { | ||
| 136 | return errors.New("usage: zut mcp get NAME [--json]") | 171 | return errors.New("usage: zut mcp get NAME [--json]") | ||
| 137 | } | 172 | } | ||
| 138 | - | cfg, _ := mcp.LoadConfig(ZutHome(), cwd) | 173 | + | cfg, err := rawMCPConfig(cwd) |
| 174 | + | if err != nil { | |||
| 175 | + | return err | |||
| 176 | + | } | |||
| 139 | srv, ok := cfg.Servers[args[0]] | 177 | srv, ok := cfg.Servers[args[0]] | ||
| 140 | if !ok { | 178 | if !ok { | ||
| 141 | - | return fmt.Errorf("mcp server %q is not configured or is invalid", args[0]) | 179 | + | return fmt.Errorf("mcp server %q is not configured", args[0]) |
| 142 | } | 180 | } | ||
| 143 | enc := json.NewEncoder(out) | 181 | enc := json.NewEncoder(out) | ||
| 144 | enc.SetIndent("", " ") | 182 | enc.SetIndent("", " ") | ||
| @@ -293,12 +331,12 @@ func mcpRemove(out io.Writer, cwd string, args []string) error { | |||||
| 293 | func mcpTest(ctx context.Context, out io.Writer, cwd, version string, args []string) error { | 331 | func mcpTest(ctx context.Context, out io.Writer, cwd, version string, args []string) error { | ||
| 294 | cfg, errs := mcp.LoadConfig(ZutHome(), cwd) | 332 | cfg, errs := mcp.LoadConfig(ZutHome(), cwd) | ||
| 295 | for _, err := range errs { | 333 | for _, err := range errs { | ||
| 296 | - | fmt.Fprintln(out, "config error:", err) | 334 | + | fmt.Fprintln(os.Stderr, "mcp config error:", err) |
| 297 | } | 335 | } | ||
| 298 | if len(args) > 0 { | 336 | if len(args) > 0 { | ||
| 299 | srv, ok := cfg.Servers[args[0]] | 337 | srv, ok := cfg.Servers[args[0]] | ||
| 300 | if !ok { | 338 | if !ok { | ||
| 301 | - | return fmt.Errorf("mcp server %q is not configured or is invalid", args[0]) | 339 | + | return fmt.Errorf("mcp server %q is not configured, disabled, or invalid", args[0]) |
| 302 | } | 340 | } | ||
| 303 | cfg.Servers = map[string]mcp.ServerConfig{args[0]: srv} | 341 | cfg.Servers = map[string]mcp.ServerConfig{args[0]: srv} | ||
| 304 | } | 342 | } | ||
| @@ -306,13 +344,13 @@ func mcpTest(ctx context.Context, out io.Writer, cwd, version string, args []str | |||||
| 306 | defer manager.Close() | 344 | defer manager.Close() | ||
| 307 | failures, _ := manager.Start(ctx) | 345 | failures, _ := manager.Start(ctx) | ||
| 308 | fmt.Fprintln(out, mcp.FormatStatus(manager.Status())) | 346 | fmt.Fprintln(out, mcp.FormatStatus(manager.Status())) | ||
| 309 | - | for _, server := range manager.Servers() { | 347 | + | // Names come from the registry builder so collision suffixes match what |
| 310 | - | for _, tool := range manager.ServerTools(server) { | 348 | + | // --tools and the model see. |
| 311 | - | fmt.Fprintf(out, " %s\n", mcp.QualifiedName(server, tool.Name)) | 349 | + | for _, tool := range manager.Tools() { |
| 312 | - | } | 350 | + | fmt.Fprintf(out, " %s\n", tool.Name()) |
| 313 | } | 351 | } | ||
| 314 | - | if len(failures) > 0 { | 352 | + | if n := len(failures) + len(errs); n > 0 { |
| 315 | - | return fmt.Errorf("%d MCP server(s) failed to start", len(failures)) | 353 | + | return fmt.Errorf("%d MCP configuration or startup problem(s)", n) |
| 316 | } | 354 | } | ||
| 317 | return nil | 355 | return nil | ||
| 318 | } | 356 | } | ||