fix(mcp): close servers on signals during startup; unblock orchestrator stdin on SIGTERM

commit: 16815dd0844458b201765bbc3a29bae284b5d14b
author: brice · 7d ago · 2026-10-15 14:03 UTC
parent: e6d4fe0
signature: GPG VerifiedCD88F9E564E0102646FB6077A1A07C7C7C3E349D · arch-gpg Verified for the committer email

5 files · +147 -10

M packages/agent/cli.go +28 -2


@@ -3246,14 +3246,40 @@ func combinePromptInput(stdin, positional string) string {
3246	3246		}
3247	3247
3248	3248		func readAllStdin() (string, error) {
3249		-	fi, err := os.Stdin.Stat()
	3249	+	return readAllPipe(os.Stdin)
	3250	+	}
	3251	+
	3252	+	// readAllPipe reads f to EOF unless it is a terminal.
	3253	+	func readAllPipe(f *os.File) (string, error) {
	3254	+	fi, err := f.Stat()
3250	3255		if err != nil {
3251	3256		return "", err
3252	3257		}
3253	3258		if (fi.Mode() & os.ModeCharDevice) != 0 {
3254	3259		return "", nil
3255	3260		}
3256		-	return readAllStdinFrom(os.Stdin)
	3261	+	return readAllStdinFrom(f)
	3262	+	}
	3263	+
	3264	+	// readAllStdinContext is readAllStdin that returns ctx.Err() when ctx ends
	3265	+	// first. The abandoned read finishes or dies with the process.
	3266	+	func readAllStdinContext(ctx context.Context) (string, error) {
	3267	+	type result struct {
	3268	+	text string
	3269	+	err error
	3270	+	}
	3271	+	done := make(chan result, 1)
	3272	+	stdin := os.Stdin
	3273	+	go func() {
	3274	+	text, err := readAllPipe(stdin)
	3275	+	done <- result{text, err}
	3276	+	}()
	3277	+	select {
	3278	+	case r := <-done:
	3279	+	return r.text, r.err
	3280	+	case <-ctx.Done():
	3281	+	return "", ctx.Err()
	3282	+	}
3257	3283		}
3258	3284
3259	3285		func readAllStdinFrom(r io.Reader) (string, error) {


@@ -3246,14 +3246,40 @@ func combinePromptInput(stdin, positional string) string {
3246		}	3246		}
3247			3247
3248		func readAllStdin() (string, error) {	3248		func readAllStdin() (string, error) {
3249	-	fi, err := os.Stdin.Stat()	3249	+	return readAllPipe(os.Stdin)
			3250	+	}
			3251	+
			3252	+	// readAllPipe reads f to EOF unless it is a terminal.
			3253	+	func readAllPipe(f *os.File) (string, error) {
			3254	+	fi, err := f.Stat()
3250		if err != nil {	3255		if err != nil {
3251		return "", err	3256		return "", err
3252		}	3257		}
3253		if (fi.Mode() & os.ModeCharDevice) != 0 {	3258		if (fi.Mode() & os.ModeCharDevice) != 0 {
3254		return "", nil	3259		return "", nil
3255		}	3260		}
3256	-	return readAllStdinFrom(os.Stdin)	3261	+	return readAllStdinFrom(f)
			3262	+	}
			3263	+
			3264	+	// readAllStdinContext is readAllStdin that returns ctx.Err() when ctx ends
			3265	+	// first. The abandoned read finishes or dies with the process.
			3266	+	func readAllStdinContext(ctx context.Context) (string, error) {
			3267	+	type result struct {
			3268	+	text string
			3269	+	err error
			3270	+	}
			3271	+	done := make(chan result, 1)
			3272	+	stdin := os.Stdin
			3273	+	go func() {
			3274	+	text, err := readAllPipe(stdin)
			3275	+	done <- result{text, err}
			3276	+	}()
			3277	+	select {
			3278	+	case r := <-done:
			3279	+	return r.text, r.err
			3280	+	case <-ctx.Done():
			3281	+	return "", ctx.Err()
			3282	+	}
3257		}	3283		}
3258			3284
3259		func readAllStdinFrom(r io.Reader) (string, error) {	3285		func readAllStdinFrom(r io.Reader) (string, error) {

M packages/agent/mcp_runtime.go +50 -7


@@ -56,7 +56,17 @@ func startMCP(ctx context.Context, args *Args, version string) (func(), error) {
56	56		}
57	57		// MCP endpoints always verify TLS; --insecure is scoped to inference endpoints.
58	58		manager := mcp.NewManager(cfg, args.CWD, mcp.ClientInfo{Name: "zut", Version: version}, provider.NewHTTPClient(false))
59		-	failures, requiredErr := manager.Start(ctx)
	59	+	sigCh := make(chan os.Signal, 1)
	60	+	signal.Notify(sigCh, mcpSignals...)
	61	+	failures, requiredErr, sig := startMCPInterruptible(ctx, manager, sigCh)
	62	+	signal.Stop(sigCh)
	63	+	if sig != nil {
	64	+	// Interrupted during startup: no mode owns signals yet, so stop the
	65	+	// servers already launched and terminate like an unhandled signal.
	66	+	_ = manager.Close()
	67	+	mcpSignalExit(signalExitCode(sig))
	68	+	return noop, fmt.Errorf("mcp startup interrupted by %v", sig)
	69	+	}
60	70		for _, err := range failures {
61	71		args.mcpStartupErrors = append(args.mcpStartupErrors, err.Error())
62	72		}

@@ -69,6 +79,43 @@ func startMCP(ctx context.Context, args *Args, version string) (func(), error) {
69	79		return func() { _ = manager.Close() }, nil
70	80		}
71	81
	82	+	var mcpSignals = []os.Signal{os.Interrupt, syscall.SIGTERM, syscall.SIGHUP}
	83	+
	84	+	// mcpSignalExit terminates the process after a signal; tests replace it.
	85	+	var mcpSignalExit = os.Exit
	86	+
	87	+	func signalExitCode(sig os.Signal) int {
	88	+	if s, ok := sig.(syscall.Signal); ok {
	89	+	return 128 + int(s)
	90	+	}
	91	+	return 1
	92	+	}
	93	+
	94	+	// startMCPInterruptible runs manager.Start and cancels the handshakes when a
	95	+	// signal arrives on sigCh, waiting until every connection attempt has closed
	96	+	// its transport. It returns the signal that interrupted startup, if any.
	97	+	func startMCPInterruptible(ctx context.Context, manager *mcp.Manager, sigCh <-chan os.Signal) ([]error, error, os.Signal) {
	98	+	startCtx, cancel := context.WithCancel(ctx)
	99	+	defer cancel()
	100	+	type result struct {
	101	+	failures []error
	102	+	required error
	103	+	}
	104	+	done := make(chan result, 1)
	105	+	go func() {
	106	+	failures, required := manager.Start(startCtx)
	107	+	done <- result{failures, required}
	108	+	}()
	109	+	select {
	110	+	case r := <-done:
	111	+	return r.failures, r.required, nil
	112	+	case sig := <-sigCh:
	113	+	cancel()
	114	+	<-done
	115	+	return nil, nil, sig
	116	+	}
	117	+	}
	118	+
72	119		// exitAfterMCPCleanupOnSignal keeps the default "terminate on signal"
73	120		// behavior of headless modes while closing MCP servers first. Default
74	121		// handling is restored before cleanup, so a second signal still kills zut

@@ -76,17 +123,13 @@ func startMCP(ctx context.Context, args *Args, version string) (func(), error) {
76	123		func exitAfterMCPCleanupOnSignal(manager *mcp.Manager) func() {
77	124		sigCh := make(chan os.Signal, 1)
78	125		done := make(chan struct{})
79		-	signal.Notify(sigCh, os.Interrupt, syscall.SIGTERM, syscall.SIGHUP)
	126	+	signal.Notify(sigCh, mcpSignals...)
80	127		go func() {
81	128		select {
82	129		case sig := <-sigCh:
83	130		signal.Stop(sigCh)
84	131		_ = manager.Close()
85		-	code := 1
86		-	if s, ok := sig.(syscall.Signal); ok {
87		-	code = 128 + int(s)
88		-	}
89		-	os.Exit(code)
	132	+	mcpSignalExit(signalExitCode(sig))
90	133		case <-done:
91	134		}
92	135		}()


@@ -56,7 +56,17 @@ func startMCP(ctx context.Context, args *Args, version string) (func(), error) {
56		}	56		}
57		// MCP endpoints always verify TLS; --insecure is scoped to inference endpoints.	57		// MCP endpoints always verify TLS; --insecure is scoped to inference endpoints.
58		manager := mcp.NewManager(cfg, args.CWD, mcp.ClientInfo{Name: "zut", Version: version}, provider.NewHTTPClient(false))	58		manager := mcp.NewManager(cfg, args.CWD, mcp.ClientInfo{Name: "zut", Version: version}, provider.NewHTTPClient(false))
59	-	failures, requiredErr := manager.Start(ctx)	59	+	sigCh := make(chan os.Signal, 1)
			60	+	signal.Notify(sigCh, mcpSignals...)
			61	+	failures, requiredErr, sig := startMCPInterruptible(ctx, manager, sigCh)
			62	+	signal.Stop(sigCh)
			63	+	if sig != nil {
			64	+	// Interrupted during startup: no mode owns signals yet, so stop the
			65	+	// servers already launched and terminate like an unhandled signal.
			66	+	_ = manager.Close()
			67	+	mcpSignalExit(signalExitCode(sig))
			68	+	return noop, fmt.Errorf("mcp startup interrupted by %v", sig)
			69	+	}
60		for _, err := range failures {	70		for _, err := range failures {
61		args.mcpStartupErrors = append(args.mcpStartupErrors, err.Error())	71		args.mcpStartupErrors = append(args.mcpStartupErrors, err.Error())
62		}	72		}

@@ -69,6 +79,43 @@ func startMCP(ctx context.Context, args *Args, version string) (func(), error) {
69		return func() { _ = manager.Close() }, nil	79		return func() { _ = manager.Close() }, nil
70		}	80		}
71			81
			82	+	var mcpSignals = []os.Signal{os.Interrupt, syscall.SIGTERM, syscall.SIGHUP}
			83	+
			84	+	// mcpSignalExit terminates the process after a signal; tests replace it.
			85	+	var mcpSignalExit = os.Exit
			86	+
			87	+	func signalExitCode(sig os.Signal) int {
			88	+	if s, ok := sig.(syscall.Signal); ok {
			89	+	return 128 + int(s)
			90	+	}
			91	+	return 1
			92	+	}
			93	+
			94	+	// startMCPInterruptible runs manager.Start and cancels the handshakes when a
			95	+	// signal arrives on sigCh, waiting until every connection attempt has closed
			96	+	// its transport. It returns the signal that interrupted startup, if any.
			97	+	func startMCPInterruptible(ctx context.Context, manager *mcp.Manager, sigCh <-chan os.Signal) ([]error, error, os.Signal) {
			98	+	startCtx, cancel := context.WithCancel(ctx)
			99	+	defer cancel()
			100	+	type result struct {
			101	+	failures []error
			102	+	required error
			103	+	}
			104	+	done := make(chan result, 1)
			105	+	go func() {
			106	+	failures, required := manager.Start(startCtx)
			107	+	done <- result{failures, required}
			108	+	}()
			109	+	select {
			110	+	case r := <-done:
			111	+	return r.failures, r.required, nil
			112	+	case sig := <-sigCh:
			113	+	cancel()
			114	+	<-done
			115	+	return nil, nil, sig
			116	+	}
			117	+	}
			118	+
72		// exitAfterMCPCleanupOnSignal keeps the default "terminate on signal"	119		// exitAfterMCPCleanupOnSignal keeps the default "terminate on signal"
73		// behavior of headless modes while closing MCP servers first. Default	120		// behavior of headless modes while closing MCP servers first. Default
74		// handling is restored before cleanup, so a second signal still kills zut	121		// handling is restored before cleanup, so a second signal still kills zut

@@ -76,17 +123,13 @@ func startMCP(ctx context.Context, args *Args, version string) (func(), error) {
76		func exitAfterMCPCleanupOnSignal(manager *mcp.Manager) func() {	123		func exitAfterMCPCleanupOnSignal(manager *mcp.Manager) func() {
77		sigCh := make(chan os.Signal, 1)	124		sigCh := make(chan os.Signal, 1)
78		done := make(chan struct{})	125		done := make(chan struct{})
79	-	signal.Notify(sigCh, os.Interrupt, syscall.SIGTERM, syscall.SIGHUP)	126	+	signal.Notify(sigCh, mcpSignals...)
80		go func() {	127		go func() {
81		select {	128		select {
82		case sig := <-sigCh:	129		case sig := <-sigCh:
83		signal.Stop(sigCh)	130		signal.Stop(sigCh)
84		_ = manager.Close()	131		_ = manager.Close()
85	-	code := 1	132	+	mcpSignalExit(signalExitCode(sig))
86	-	if s, ok := sig.(syscall.Signal); ok {
87	-	code = 128 + int(s)
88	-	}
89	-	os.Exit(code)
90		case <-done:	133		case <-done:
91		}	134		}
92		}()	135		}()

M packages/agent/mcp_runtime_test.go +35 -0


@@ -391,3 +391,38 @@ func keys(reg core.Registry) []string {
391	391		slices.Sort(out)
392	392		return out
393	393		}
	394	+
	395	+	// A signal during the MCP handshake must cancel startup (not wait for the
	396	+	// startup timeout) so the caller can close the servers already launched.
	397	+	func TestStartMCPInterruptibleCancelsHandshake(t *testing.T) {
	398	+	reached := make(chan struct{}, 1)
	399	+	release := make(chan struct{})
	400	+	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
	401	+	select {
	402	+	case reached <- struct{}{}:
	403	+	default:
	404	+	}
	405	+	// Stall every request until the client gives up or the test ends.
	406	+	select {
	407	+	case <-r.Context().Done():
	408	+	case <-release:
	409	+	}
	410	+	}))
	411	+	t.Cleanup(srv.Close)
	412	+	t.Cleanup(func() { close(release) })
	413	+	cfg := mcp.Config{Servers: map[string]mcp.ServerConfig{"slow": {Type: "http", URL: srv.URL}}}
	414	+	manager := mcp.NewManager(cfg, t.TempDir(), mcp.ClientInfo{Name: "zut"}, srv.Client())
	415	+	t.Cleanup(func() { _ = manager.Close() })
	416	+
	417	+	sigCh := make(chan os.Signal, 1)
	418	+	result := make(chan os.Signal, 1)
	419	+	go func() {
	420	+	_, _, sig := startMCPInterruptible(context.Background(), manager, sigCh)
	421	+	result <- sig
	422	+	}()
	423	+	<-reached
	424	+	sigCh <- os.Interrupt
	425	+	if sig := <-result; sig != os.Interrupt {
	426	+	t.Fatalf("signal = %v, want interrupt", sig)
	427	+	}
	428	+	}


@@ -391,3 +391,38 @@ func keys(reg core.Registry) []string {
391	slices.Sort(out)	391		slices.Sort(out)
392	return out	392		return out
393	}	393		}
		394	+
		395	+	// A signal during the MCP handshake must cancel startup (not wait for the
		396	+	// startup timeout) so the caller can close the servers already launched.
		397	+	func TestStartMCPInterruptibleCancelsHandshake(t *testing.T) {
		398	+	reached := make(chan struct{}, 1)
		399	+	release := make(chan struct{})
		400	+	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		401	+	select {
		402	+	case reached <- struct{}{}:
		403	+	default:
		404	+	}
		405	+	// Stall every request until the client gives up or the test ends.
		406	+	select {
		407	+	case <-r.Context().Done():
		408	+	case <-release:
		409	+	}
		410	+	}))
		411	+	t.Cleanup(srv.Close)
		412	+	t.Cleanup(func() { close(release) })
		413	+	cfg := mcp.Config{Servers: map[string]mcp.ServerConfig{"slow": {Type: "http", URL: srv.URL}}}
		414	+	manager := mcp.NewManager(cfg, t.TempDir(), mcp.ClientInfo{Name: "zut"}, srv.Client())
		415	+	t.Cleanup(func() { _ = manager.Close() })
		416	+
		417	+	sigCh := make(chan os.Signal, 1)
		418	+	result := make(chan os.Signal, 1)
		419	+	go func() {
		420	+	_, _, sig := startMCPInterruptible(context.Background(), manager, sigCh)
		421	+	result <- sig
		422	+	}()
		423	+	<-reached
		424	+	sigCh <- os.Interrupt
		425	+	if sig := <-result; sig != os.Interrupt {
		426	+	t.Fatalf("signal = %v, want interrupt", sig)
		427	+	}
		428	+	}

M packages/agent/orchestration_modes.go +7 -1


@@ -35,6 +35,12 @@ type orchestratedModeHooks struct {
35	35		func runOrchestratedMode(parentCtx context.Context, args Args, version string, hooks orchestratedModeHooks) (runErr error) {
36	36		ctx, stopSignal := signal.NotifyContext(parentCtx, os.Interrupt, syscall.SIGTERM, syscall.SIGHUP)
37	37		defer stopSignal()
	38	+	// After the first signal cancels ctx, restore default handling so a
	39	+	// second signal terminates a run whose cleanup is stuck.
	40	+	go func() {
	41	+	<-ctx.Done()
	42	+	stopSignal()
	43	+	}()
38	44
39	45		if args.NoYolo {
40	46		fmt.Fprintf(os.Stderr, "warning: --no-yolo has no effect in %s mode (no interactive prompt available); tools will run without confirmation\n", hooks.promptLabel)

@@ -49,7 +55,7 @@ func runOrchestratedMode(parentCtx context.Context, args Args, version string, h
49	55		}
50	56		prompt := args.Prompt
51	57		if prompt == "" {
52		-	piped, readErr := readAllStdin()
	58	+	piped, readErr := readAllStdinContext(ctx)
53	59		if readErr != nil {
54	60		return fmt.Errorf("read %s prompt from stdin: %w", hooks.promptLabel, readErr)
55	61		}

 func runOrchestratedMode(parentCtx context.Context, args Args, version string, hooks orchestratedModeHooks) (runErr error) {
 	ctx, stopSignal := signal.NotifyContext(parentCtx, os.Interrupt, syscall.SIGTERM, syscall.SIGHUP)
 	defer stopSignal()
+	// After the first signal cancels ctx, restore default handling so a
+	// second signal terminates a run whose cleanup is stuck.
+	go func() {
+		<-ctx.Done()
+		stopSignal()
+	}()
 	if args.NoYolo {
 		fmt.Fprintf(os.Stderr, "warning: --no-yolo has no effect in %s mode (no interactive prompt available); tools will run without confirmation\n", hooks.promptLabel)
 	}
 	prompt := args.Prompt
 	if prompt == "" {
-		piped, readErr := readAllStdin()
+		piped, readErr := readAllStdinContext(ctx)
 		if readErr != nil {
 			return fmt.Errorf("read %s prompt from stdin: %w", hooks.promptLabel, readErr)
 		}

A packages/agent/stdin_context_test.go +27 -0

@@ -0,0 +1,27 @@
	1	+	package agent
	2	+
	3	+	import (
	4	+	"context"
	5	+	"errors"
	6	+	"os"
	7	+	"testing"
	8	+	)
	9	+
	10	+	// An orchestrator waiting for a piped prompt must stop when its signal
	11	+	// context ends instead of blocking on an open pipe.
	12	+	func TestReadAllStdinContextReturnsOnCancel(t *testing.T) {
	13	+	r, w, err := os.Pipe()
	14	+	if err != nil {
	15	+	t.Fatal(err)
	16	+	}
	17	+	t.Cleanup(func() { _ = w.Close(); _ = r.Close() })
	18	+	old := os.Stdin
	19	+	os.Stdin = r
	20	+	t.Cleanup(func() { os.Stdin = old })
	21	+
	22	+	ctx, cancel := context.WithCancel(context.Background())
	23	+	cancel()
	24	+	if _, err := readAllStdinContext(ctx); !errors.Is(err, context.Canceled) {
	25	+	t.Fatalf("err = %v, want context.Canceled", err)
	26	+	}
	27	+	}

@@ -0,0 +1,27 @@
	1	+	package agent
	2	+
	3	+	import (
	4	+	"context"
	5	+	"errors"
	6	+	"os"
	7	+	"testing"
	8	+	)
	9	+
	10	+	// An orchestrator waiting for a piped prompt must stop when its signal
	11	+	// context ends instead of blocking on an open pipe.
	12	+	func TestReadAllStdinContextReturnsOnCancel(t *testing.T) {
	13	+	r, w, err := os.Pipe()
	14	+	if err != nil {
	15	+	t.Fatal(err)
	16	+	}
	17	+	t.Cleanup(func() { _ = w.Close(); _ = r.Close() })
	18	+	old := os.Stdin
	19	+	os.Stdin = r
	20	+	t.Cleanup(func() { os.Stdin = old })
	21	+
	22	+	ctx, cancel := context.WithCancel(context.Background())
	23	+	cancel()
	24	+	if _, err := readAllStdinContext(ctx); !errors.Is(err, context.Canceled) {
	25	+	t.Fatalf("err = %v, want context.Canceled", err)
	26	+	}
	27	+	}