Merge pull request #2436 from dgageot/fix-sleep

fix: replace time.Sleep with context-aware select in MCP tryRestart

Author: dgageot

pkg/tools/mcp/mcp.go

@@ -309,7 +309,14 @@ func (ts *Toolset) tryRestart(ctx context.Context) bool {
 	for attempt := range maxAttempts {
 		backoff := time.Duration(1<<uint(attempt)) * time.Second
 		slog.Debug("Restarting MCP server", "server", ts.logID, "attempt", attempt+1, "backoff", backoff)
-		time.Sleep(backoff)
+
+		timer := time.NewTimer(backoff)
+		select {
+		case <-timer.C:
+		case <-ctx.Done():
+			timer.Stop()
+			return false
+		}
 
 		ts.mu.Lock()
 		if ts.stopping {

pkg/tools/mcp/mcp_race_test.go

@@ -1,8 +1,12 @@
 package mcp
 
 import (
+	"context"
 	"sync"
 	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
 )
 
 func TestInstructions_Concurrent(t *testing.T) {
@@ -30,3 +34,31 @@ func TestInstructions_Concurrent(t *testing.T) {
 	}
 	wg.Wait()
 }
+
+func TestTryRestart_RespectsContextCancellation(t *testing.T) {
+	t.Parallel()
+
+	ts := &Toolset{
+		logID:     "test",
+		mcpClient: &mockMCPClient{},
+	}
+
+	ctx, cancel := context.WithCancel(t.Context())
+
+	done := make(chan bool, 1)
+	go func() {
+		done <- ts.tryRestart(ctx)
+	}()
+
+	// Cancel almost immediately; tryRestart should return promptly
+	// instead of sleeping through the full backoff.
+	time.Sleep(50 * time.Millisecond)
+	cancel()
+
+	select {
+	case result := <-done:
+		assert.False(t, result, "tryRestart should return false on cancellation")
+	case <-time.After(2 * time.Second):
+		t.Fatal("tryRestart did not return promptly after context cancellation")
+	}
+}