fix(drift): close active connections on shutdown to prevent Ctrl+C hang

[sohil-kshirsagar]

tusk drift run --print hangs on Ctrl+C when the mock server has an active SDK connection.

Root cause

server.Stop() calls ms.wg.Wait(), which blocks until all handleConnection goroutines finish. Those goroutines are stuck on io.ReadFull(conn, ...) — a blocking read with no deadline. Closing the listener prevents new connections but doesn't close existing ones, so io.ReadFull never returns and the process hangs indefinitely.

Changes

• Track active connections in a map and close them all in Stop() before wg.Wait(), so blocked reads return immediately with an error
• Add force-exit on second Ctrl+C as a safety net — if cleanup itself hangs for any reason, users aren't stuck
• Add TestServerStopClosesActiveConnections — verified it fails without the fix (hangs for 2s then times out) and passes with it

[cursor]

Cursor Bugbot has reviewed your changes and found 1 potential issue.

^{❌ Bugbot Autofix is OFF. To automatically fix reported issues with cloud agents, enable autofix in the Cursor dashboard.}

^{Reviewed by Cursor Bugbot for commit 6418f90. Configure here.}

[cubic-dev-ai]

1 issue found across 3 files

Prompt for AI agents (unresolved issues)

Check if these issues are valid — if so, understand the root cause of each and fix them. If appropriate, use sub-agents to investigate and fix each issue separately.


<file name="internal/runner/server.go">

<violation number="1" location="internal/runner/server.go:636">
P1: Shutdown still has a race: a connection accepted right around cancellation can be added after `Stop()` closes active connections, leaving a handler blocked on `io.ReadFull` and potentially hanging `wg.Wait()`.</violation>
</file>

_{Reply with feedback, questions, or to request a fix. Tag @cubic-dev-ai to re-run a review, or fix all with cubic.}

[cubic-dev-ai]

P1: Shutdown still has a race: a connection accepted right around cancellation can be added after Stop() closes active connections, leaving a handler blocked on io.ReadFull and potentially hanging wg.Wait().

Prompt for AI agents

Check if this issue is valid — if so, understand the root cause and fix it. At internal/runner/server.go, line 636:

<comment>Shutdown still has a race: a connection accepted right around cancellation can be added after `Stop()` closes active connections, leaving a handler blocked on `io.ReadFull` and potentially hanging `wg.Wait()`.</comment>

<file context>
@@ -622,6 +632,10 @@ func (ms *Server) acceptConnections() {
 			}
 
+			ms.activeConnsMu.Lock()
+			ms.activeConns[conn] = struct{}{}
+			ms.activeConnsMu.Unlock()
+
</file context>

Suggested change

	ms.activeConns[conn] = struct{}{}
	ms.activeConnsMu.Unlock()
	if ms.ctx.Err() != nil {
	_ = conn.Close()
	return
	}
	ms.activeConnsMu.Lock()
	ms.activeConns[conn] = struct{}{}
	ms.activeConnsMu.Unlock()