fix(mcp): close servers on headless signals and redact templated URLs
- commit
- be3ffc0e3103486ce44062c6783e195fbcfba8ad
- author
- brice · · 2026-10-15 13:39 UTC
- parent
- 6978ba7
- signature
- GPG Verified
CD88F9E564E0102646FB6077A1A07C7C7C3E349D· arch-gpg Verified for the committer email
4 files · +71 -21
M packages/agent/cli.go +6 -13
| @@ -1005,19 +1005,6 @@ 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 | - | } | |
| 1021 | 1008 | // MCP servers start once per process. Every later Resolve, including | |
| 1022 | 1009 | // resident subagents, shares these connections through args. | |
| 1023 | 1010 | stopMCP, err := startMCP(ctx, &args, version) | |
| @@ -1025,6 +1012,12 @@ func runWithArgs(args Args, version string) error { | |||
| 1025 | 1012 | return err | |
| 1026 | 1013 | } | |
| 1027 | 1014 | defer stopMCP() | |
| 1015 | + | if args.mcpManager != nil && args.Mode != ModeInteractive && args.Mode != "" && !args.Orchestrate { | |
| 1016 | + | // Plain headless modes terminate on signals. Stop MCP servers first so | |
| 1017 | + | // their process groups and HTTP sessions do not outlive zut. The TUI | |
| 1018 | + | // and --orchestrate own their signal handling and cancel gracefully. | |
| 1019 | + | defer exitAfterMCPCleanupOnSignal(args.mcpManager)() | |
| 1020 | + | } | |
| 1028 | 1021 | switch args.Mode { | |
| 1029 | 1022 | case ModePrint: | |
| 1030 | 1023 | return runPrintMode(ctx, args, version) | |
| @@ -1005,19 +1005,6 @@ 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 | - | } | |||
| 1021 | // MCP servers start once per process. Every later Resolve, including | 1008 | // MCP servers start once per process. Every later Resolve, including | ||
| 1022 | // resident subagents, shares these connections through args. | 1009 | // resident subagents, shares these connections through args. | ||
| 1023 | stopMCP, err := startMCP(ctx, &args, version) | 1010 | stopMCP, err := startMCP(ctx, &args, version) | ||
| @@ -1025,6 +1012,12 @@ func runWithArgs(args Args, version string) error { | |||||
| 1025 | return err | 1012 | return err | ||
| 1026 | } | 1013 | } | ||
| 1027 | defer stopMCP() | 1014 | defer stopMCP() | ||
| 1015 | + | if args.mcpManager != nil && args.Mode != ModeInteractive && args.Mode != "" && !args.Orchestrate { | |||
| 1016 | + | // Plain headless modes terminate on signals. Stop MCP servers first so | |||
| 1017 | + | // their process groups and HTTP sessions do not outlive zut. The TUI | |||
| 1018 | + | // and --orchestrate own their signal handling and cancel gracefully. | |||
| 1019 | + | defer exitAfterMCPCleanupOnSignal(args.mcpManager)() | |||
| 1020 | + | } | |||
| 1028 | switch args.Mode { | 1021 | switch args.Mode { | ||
| 1029 | case ModePrint: | 1022 | case ModePrint: | ||
| 1030 | return runPrintMode(ctx, args, version) | 1023 | return runPrintMode(ctx, args, version) | ||
M packages/agent/mcp_runtime.go +29 -0
| @@ -4,7 +4,9 @@ import ( | |||
| 4 | 4 | "context" | |
| 5 | 5 | "fmt" | |
| 6 | 6 | "os" | |
| 7 | + | "os/signal" | |
| 7 | 8 | "strings" | |
| 9 | + | "syscall" | |
| 8 | 10 | ||
| 9 | 11 | "github.com/bnema/zut/packages/agent/mcp" | |
| 10 | 12 | "github.com/bnema/zut/packages/core" | |
| @@ -67,6 +69,33 @@ func startMCP(ctx context.Context, args *Args, version string) (func(), error) { | |||
| 67 | 69 | return func() { _ = manager.Close() }, nil | |
| 68 | 70 | } | |
| 69 | 71 | ||
| 72 | + | // exitAfterMCPCleanupOnSignal keeps the default "terminate on signal" | |
| 73 | + | // behavior of headless modes while closing MCP servers first. Default | |
| 74 | + | // handling is restored before cleanup, so a second signal still kills zut | |
| 75 | + | // immediately. The returned func uninstalls the handler. | |
| 76 | + | func exitAfterMCPCleanupOnSignal(manager *mcp.Manager) func() { | |
| 77 | + | sigCh := make(chan os.Signal, 1) | |
| 78 | + | done := make(chan struct{}) | |
| 79 | + | signal.Notify(sigCh, os.Interrupt, syscall.SIGTERM, syscall.SIGHUP) | |
| 80 | + | go func() { | |
| 81 | + | select { | |
| 82 | + | case sig := <-sigCh: | |
| 83 | + | signal.Stop(sigCh) | |
| 84 | + | _ = manager.Close() | |
| 85 | + | code := 1 | |
| 86 | + | if s, ok := sig.(syscall.Signal); ok { | |
| 87 | + | code = 128 + int(s) | |
| 88 | + | } | |
| 89 | + | os.Exit(code) | |
| 90 | + | case <-done: | |
| 91 | + | } | |
| 92 | + | }() | |
| 93 | + | return func() { | |
| 94 | + | signal.Stop(sigCh) | |
| 95 | + | close(done) | |
| 96 | + | } | |
| 97 | + | } | |
| 98 | + | ||
| 70 | 99 | // reportMCPStartup prints startup failures for non-interactive modes. The | |
| 71 | 100 | // interactive TUI shows them with the extension startup errors instead. | |
| 72 | 101 | func reportMCPStartup(args Args) { | |
| @@ -4,7 +4,9 @@ import ( | |||||
| 4 | "context" | 4 | "context" | ||
| 5 | "fmt" | 5 | "fmt" | ||
| 6 | "os" | 6 | "os" | ||
| 7 | + | "os/signal" | |||
| 7 | "strings" | 8 | "strings" | ||
| 9 | + | "syscall" | |||
| 8 | 10 | ||||
| 9 | "github.com/bnema/zut/packages/agent/mcp" | 11 | "github.com/bnema/zut/packages/agent/mcp" | ||
| 10 | "github.com/bnema/zut/packages/core" | 12 | "github.com/bnema/zut/packages/core" | ||
| @@ -67,6 +69,33 @@ func startMCP(ctx context.Context, args *Args, version string) (func(), error) { | |||||
| 67 | return func() { _ = manager.Close() }, nil | 69 | return func() { _ = manager.Close() }, nil | ||
| 68 | } | 70 | } | ||
| 69 | 71 | ||||
| 72 | + | // exitAfterMCPCleanupOnSignal keeps the default "terminate on signal" | |||
| 73 | + | // behavior of headless modes while closing MCP servers first. Default | |||
| 74 | + | // handling is restored before cleanup, so a second signal still kills zut | |||
| 75 | + | // immediately. The returned func uninstalls the handler. | |||
| 76 | + | func exitAfterMCPCleanupOnSignal(manager *mcp.Manager) func() { | |||
| 77 | + | sigCh := make(chan os.Signal, 1) | |||
| 78 | + | done := make(chan struct{}) | |||
| 79 | + | signal.Notify(sigCh, os.Interrupt, syscall.SIGTERM, syscall.SIGHUP) | |||
| 80 | + | go func() { | |||
| 81 | + | select { | |||
| 82 | + | case sig := <-sigCh: | |||
| 83 | + | signal.Stop(sigCh) | |||
| 84 | + | _ = manager.Close() | |||
| 85 | + | code := 1 | |||
| 86 | + | if s, ok := sig.(syscall.Signal); ok { | |||
| 87 | + | code = 128 + int(s) | |||
| 88 | + | } | |||
| 89 | + | os.Exit(code) | |||
| 90 | + | case <-done: | |||
| 91 | + | } | |||
| 92 | + | }() | |||
| 93 | + | return func() { | |||
| 94 | + | signal.Stop(sigCh) | |||
| 95 | + | close(done) | |||
| 96 | + | } | |||
| 97 | + | } | |||
| 98 | + | ||||
| 70 | // reportMCPStartup prints startup failures for non-interactive modes. The | 99 | // reportMCPStartup prints startup failures for non-interactive modes. The | ||
| 71 | // interactive TUI shows them with the extension startup errors instead. | 100 | // interactive TUI shows them with the extension startup errors instead. | ||
| 72 | func reportMCPStartup(args Args) { | 101 | func reportMCPStartup(args Args) { | ||
M packages/agent/mcp_runtime_test.go +14 -0
| @@ -345,6 +345,20 @@ func TestMCPCommandListRedactsAndKeepsDisabledEntries(t *testing.T) { | |||
| 345 | 345 | } | |
| 346 | 346 | } | |
| 347 | 347 | ||
| 348 | + | func TestRedactURL(t *testing.T) { | |
| 349 | + | for in, want := range map[string]string{ | |
| 350 | + | "https://example.test/mcp": "https://example.test/mcp", | |
| 351 | + | "https://u:p@example.test/mcp?key=s": "https://*****@example.test/mcp?*****", | |
| 352 | + | "https://${U}:${P}@example.test/mcp?key=literal-secret#frag": "https://*****@example.test/mcp?*****", | |
| 353 | + | "${BASE}/mcp?token=literal-secret": "${BASE}/mcp?*****", | |
| 354 | + | "https://example.test/a@b": "https://example.test/a@b", | |
| 355 | + | } { | |
| 356 | + | if got := redactURL(in); got != want { | |
| 357 | + | t.Errorf("redactURL(%q) = %q, want %q", in, got, want) | |
| 358 | + | } | |
| 359 | + | } | |
| 360 | + | } | |
| 361 | + | ||
| 348 | 362 | func TestMCPCommandFailsOnMalformedConfig(t *testing.T) { | |
| 349 | 363 | root := t.TempDir() | |
| 350 | 364 | t.Setenv("ZUT_HOME", root) | |
| @@ -345,6 +345,20 @@ func TestMCPCommandListRedactsAndKeepsDisabledEntries(t *testing.T) { | |||||
| 345 | } | 345 | } | ||
| 346 | } | 346 | } | ||
| 347 | 347 | ||||
| 348 | + | func TestRedactURL(t *testing.T) { | |||
| 349 | + | for in, want := range map[string]string{ | |||
| 350 | + | "https://example.test/mcp": "https://example.test/mcp", | |||
| 351 | + | "https://u:p@example.test/mcp?key=s": "https://*****@example.test/mcp?*****", | |||
| 352 | + | "https://${U}:${P}@example.test/mcp?key=literal-secret#frag": "https://*****@example.test/mcp?*****", | |||
| 353 | + | "${BASE}/mcp?token=literal-secret": "${BASE}/mcp?*****", | |||
| 354 | + | "https://example.test/a@b": "https://example.test/a@b", | |||
| 355 | + | } { | |||
| 356 | + | if got := redactURL(in); got != want { | |||
| 357 | + | t.Errorf("redactURL(%q) = %q, want %q", in, got, want) | |||
| 358 | + | } | |||
| 359 | + | } | |||
| 360 | + | } | |||
| 361 | + | ||||
| 348 | func TestMCPCommandFailsOnMalformedConfig(t *testing.T) { | 362 | func TestMCPCommandFailsOnMalformedConfig(t *testing.T) { | ||
| 349 | root := t.TempDir() | 363 | root := t.TempDir() | ||
| 350 | t.Setenv("ZUT_HOME", root) | 364 | t.Setenv("ZUT_HOME", root) | ||
M packages/agent/mcpcmd.go +22 -8
| @@ -7,7 +7,6 @@ import ( | |||
| 7 | 7 | "fmt" | |
| 8 | 8 | "io" | |
| 9 | 9 | "io/fs" | |
| 10 | - | "net/url" | |
| 11 | 10 | "os" | |
| 12 | 11 | "sort" | |
| 13 | 12 | "strconv" | |
| @@ -144,13 +143,7 @@ func mcpList(out io.Writer, cwd string, args []string) error { | |||
| 144 | 143 | // values, URL credentials and query strings. Arguments are shown as written | |
| 145 | 144 | // in the file (${VAR} references stay unexpanded). | |
| 146 | 145 | 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 | - | } | |
| 146 | + | srv.URL = redactURL(srv.URL) | |
| 154 | 147 | hide := func(m map[string]string) map[string]string { | |
| 155 | 148 | if len(m) == 0 { | |
| 156 | 149 | return m | |
| @@ -166,6 +159,27 @@ func redactServer(srv mcp.ServerConfig) mcp.ServerConfig { | |||
| 166 | 159 | return srv | |
| 167 | 160 | } | |
| 168 | 161 | ||
| 162 | + | // redactURL hides userinfo, query, and fragment with string operations so an | |
| 163 | + | // unexpanded ${VAR} template that url.Parse rejects cannot fail open. | |
| 164 | + | func redactURL(raw string) string { | |
| 165 | + | if i := strings.IndexAny(raw, "?#"); i >= 0 { | |
| 166 | + | raw = raw[:i] + "?*****" | |
| 167 | + | } | |
| 168 | + | schemeEnd := strings.Index(raw, "://") | |
| 169 | + | start := 0 | |
| 170 | + | if schemeEnd >= 0 { | |
| 171 | + | start = schemeEnd + 3 | |
| 172 | + | } | |
| 173 | + | authorityEnd := len(raw) | |
| 174 | + | if i := strings.IndexByte(raw[start:], '/'); i >= 0 { | |
| 175 | + | authorityEnd = start + i | |
| 176 | + | } | |
| 177 | + | if at := strings.LastIndexByte(raw[start:authorityEnd], '@'); at >= 0 { | |
| 178 | + | raw = raw[:start] + "*****@" + raw[start+at+1:] | |
| 179 | + | } | |
| 180 | + | return raw | |
| 181 | + | } | |
| 182 | + | ||
| 169 | 183 | func mcpGet(out io.Writer, cwd string, args []string) error { | |
| 170 | 184 | if len(args) == 0 { | |
| 171 | 185 | return errors.New("usage: zut mcp get NAME [--json]") | |
| @@ -7,7 +7,6 @@ import ( | |||||
| 7 | "fmt" | 7 | "fmt" | ||
| 8 | "io" | 8 | "io" | ||
| 9 | "io/fs" | 9 | "io/fs" | ||
| 10 | - | "net/url" | |||
| 11 | "os" | 10 | "os" | ||
| 12 | "sort" | 11 | "sort" | ||
| 13 | "strconv" | 12 | "strconv" | ||
| @@ -144,13 +143,7 @@ func mcpList(out io.Writer, cwd string, args []string) error { | |||||
| 144 | // values, URL credentials and query strings. Arguments are shown as written | 143 | // values, URL credentials and query strings. Arguments are shown as written | ||
| 145 | // in the file (${VAR} references stay unexpanded). | 144 | // in the file (${VAR} references stay unexpanded). | ||
| 146 | func redactServer(srv mcp.ServerConfig) mcp.ServerConfig { | 145 | func redactServer(srv mcp.ServerConfig) mcp.ServerConfig { | ||
| 147 | - | if u, err := url.Parse(srv.URL); err == nil && (u.User != nil || u.RawQuery != "") { | 146 | + | srv.URL = redactURL(srv.URL) |
| 148 | - | u.User = nil | |||
| 149 | - | if u.RawQuery != "" { | |||
| 150 | - | u.RawQuery = "*****" | |||
| 151 | - | } | |||
| 152 | - | srv.URL = u.String() | |||
| 153 | - | } | |||
| 154 | hide := func(m map[string]string) map[string]string { | 147 | hide := func(m map[string]string) map[string]string { | ||
| 155 | if len(m) == 0 { | 148 | if len(m) == 0 { | ||
| 156 | return m | 149 | return m | ||
| @@ -166,6 +159,27 @@ func redactServer(srv mcp.ServerConfig) mcp.ServerConfig { | |||||
| 166 | return srv | 159 | return srv | ||
| 167 | } | 160 | } | ||
| 168 | 161 | ||||
| 162 | + | // redactURL hides userinfo, query, and fragment with string operations so an | |||
| 163 | + | // unexpanded ${VAR} template that url.Parse rejects cannot fail open. | |||
| 164 | + | func redactURL(raw string) string { | |||
| 165 | + | if i := strings.IndexAny(raw, "?#"); i >= 0 { | |||
| 166 | + | raw = raw[:i] + "?*****" | |||
| 167 | + | } | |||
| 168 | + | schemeEnd := strings.Index(raw, "://") | |||
| 169 | + | start := 0 | |||
| 170 | + | if schemeEnd >= 0 { | |||
| 171 | + | start = schemeEnd + 3 | |||
| 172 | + | } | |||
| 173 | + | authorityEnd := len(raw) | |||
| 174 | + | if i := strings.IndexByte(raw[start:], '/'); i >= 0 { | |||
| 175 | + | authorityEnd = start + i | |||
| 176 | + | } | |||
| 177 | + | if at := strings.LastIndexByte(raw[start:authorityEnd], '@'); at >= 0 { | |||
| 178 | + | raw = raw[:start] + "*****@" + raw[start+at+1:] | |||
| 179 | + | } | |||
| 180 | + | return raw | |||
| 181 | + | } | |||
| 182 | + | ||||
| 169 | func mcpGet(out io.Writer, cwd string, args []string) error { | 183 | func mcpGet(out io.Writer, cwd string, args []string) error { | ||
| 170 | if len(args) == 0 { | 184 | if len(args) == 0 { | ||
| 171 | return errors.New("usage: zut mcp get NAME [--json]") | 185 | return errors.New("usage: zut mcp get NAME [--json]") | ||