git.iko.soy/crush
on main help sign in
fix: prevent goroutine orphaning in mcp.Close() and shell.KillAll() (#2159)
commit3b3a3d78c8a329ab9274fe95b47deada3f8a133c authorM1xA <AnyCPU@users.noreply.github.com> dateMon 9 Feb 2026, 14:40 (8 months ago) parent16ac973
4 files changed+59−36
compare with main
unifiedsplit
internal/agent/tools/mcp/init.go+16−13 internal/app/app.go+5−3 internal/shell/background.go+13−19 internal/shell/background_test.go+25−1
internal/agent/tools/mcp/init.go+16−13
@@ -124,26 +124,29 @@ func GetState(name string) (ClientInfo, bool) {
124124
125125// Close closes all MCP clients. This should be called during application shutdown.
126126func Close() error {
127+ ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
128+ defer cancel()
129+
127130 var wg sync.WaitGroup
128− done := make(chan struct{}, 1)
129− go func() {
130− for name, session := range sessions.Seq2() {
131− wg.Go(func() {
132− if err := session.Close(); err != nil &&
131+ for name, session := range sessions.Seq2() {
132+ wg.Go(func() {
133+ done := make(chan error, 1)
134+ go func() {
135+ done <- session.Close()
136+ }()
137+ select {
138+ case err := <-done:
139+ if err != nil &&
133140 !errors.Is(err, io.EOF) &&
134141 !errors.Is(err, context.Canceled) &&
135142 err.Error() != "signal: killed" {
136143 slog.Warn("Failed to shutdown MCP client", "name", name, "error", err)
137144 }
138− })
139− }
140− wg.Wait()
141− done <- struct{}{}
142− }()
143− select {
144− case <-done:
145− case <-time.After(5 * time.Second):
145+ case <-ctx.Done():
146+ }
147+ })
146148 }
149+ wg.Wait()
147150 broker.Shutdown()
148151 return nil
149152}
internal/app/app.go+5−3
@@ -540,14 +540,16 @@ func (app *App) Shutdown() {
540540 // Now run remaining cleanup tasks in parallel.
541541 var wg sync.WaitGroup
542542
543+ // Shared shutdown context for all timeout-bounded cleanup.
544+ shutdownCtx, cancel := context.WithTimeout(app.globalCtx, 5*time.Second)
545+ defer cancel()
546+
543547 // Kill all background shells.
544548 wg.Go(func() {
545− shell.GetBackgroundShellManager().KillAll()
549+ shell.GetBackgroundShellManager().KillAll(shutdownCtx)
546550 })
547551
548552 // Shutdown all LSP clients.
549− shutdownCtx, cancel := context.WithTimeout(app.globalCtx, 5*time.Second)
550− defer cancel()
551553 wg.Go(func() {
552554 app.LSPManager.StopAll(shutdownCtx)
553555 })
internal/shell/background.go+13−19
@@ -191,29 +191,23 @@ func (m *BackgroundShellManager) Cleanup() int {
191191 return len(toRemove)
192192}
193193
194−// KillAll terminates all background shells.
195−func (m *BackgroundShellManager) KillAll() {
194+// KillAll terminates all background shells. The provided context bounds how
195+// long the function waits for each shell to exit.
196+func (m *BackgroundShellManager) KillAll(ctx context.Context) {
196197 shells := slices.Collect(m.shells.Seq())
197198 m.shells.Reset(map[string]*BackgroundShell{})
198− done := make(chan struct{}, 1)
199− go func() {
200− var wg sync.WaitGroup
201− for _, shell := range shells {
202− wg.Go(func() {
203− shell.cancel()
204− <-shell.done
205− })
206− }
207− wg.Wait()
208− done <- struct{}{}
209− }()
210199
211− select {
212− case <-done:
213− return
214− case <-time.After(time.Second * 5):
215− return
200+ var wg sync.WaitGroup
201+ for _, shell := range shells {
202+ wg.Go(func() {
203+ shell.cancel()
204+ select {
205+ case <-shell.done:
206+ case <-ctx.Done():
207+ }
208+ })
216209 }
210+ wg.Wait()
217211}
218212
219213// GetOutput returns the current output of a background shell.
internal/shell/background_test.go+25−1
@@ -6,6 +6,8 @@ import (
66 "strings"
77 "testing"
88 "time"
9+
10+ "github.com/stretchr/testify/require"
911)
1012
1113func TestBackgroundShellManager_Start(t *testing.T) {
@@ -248,7 +250,7 @@ func TestBackgroundShellManager_KillAll(t *testing.T) {
248250 }
249251
250252 // Kill all shells
251− manager.KillAll()
253+ manager.KillAll(context.Background())
252254
253255 // Verify all shells are done
254256 if !shell1.IsDone() {
@@ -280,3 +282,25 @@ func TestBackgroundShellManager_KillAll(t *testing.T) {
280282 }
281283 }
282284}
285+
286+func TestBackgroundShellManager_KillAll_Timeout(t *testing.T) {
287+ t.Parallel()
288+
289+ workingDir := t.TempDir()
290+ manager := newBackgroundShellManager()
291+
292+ // Start a shell that traps signals and ignores cancellation.
293+ _, err := manager.Start(context.Background(), workingDir, nil, "trap '' TERM INT; sleep 60", "")
294+ require.NoError(t, err)
295+
296+ // Short timeout to test the timeout path.
297+ ctx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond)
298+ defer cancel()
299+
300+ start := time.Now()
301+ manager.KillAll(ctx)
302+ elapsed := time.Since(start)
303+
304+ // Must return promptly after timeout, not hang for 60 seconds.
305+ require.Less(t, elapsed, 2*time.Second)
306+}