fix(agent): keep reconnecting after an async WebSocket handshake failure (#2329, #2326)

agent/connection_manager_test.go

@@ -9,6 +9,7 @@ import (
 	"net"
 	"net/url"
 	"testing"
+	"testing/synctest"
 	"time"
 
 	"github.com/stretchr/testify/assert"
@@ -77,6 +78,10 @@ func TestConnectionManager_StateTransitions(t *testing.T) {
 	cm.handleStateChange(SSHConnected)
 	assert.Equal(t, SSHConnected, cm.State, "State should change to SSHConnected")
 
+	// Prevent handleStateChange from spawning its async reconnect goroutine:
+	// this test only checks the synchronous state machine, and the goroutine
+	// would otherwise race with the direct field writes below.
+	cm.setConnecting(true)
 	cm.handleStateChange(Disconnected)
 	assert.Equal(t, Disconnected, cm.State, "State should change to Disconnected")
 
@@ -95,7 +100,6 @@ func TestConnectionManager_EventHandling(t *testing.T) {
 			Host: "localhost:8080",
 		},
 	}
-
 	testCases := []struct {
 		name          string
 		initialState  ConnectionState
@@ -148,6 +152,11 @@ func TestConnectionManager_EventHandling(t *testing.T) {
 
 	for _, tc := range testCases {
 		t.Run(tc.name, func(t *testing.T) {
+			// Prevent handleStateChange from spawning its async reconnect
+			// goroutine: this test only checks the synchronous state machine,
+			// and the goroutine would otherwise race with the direct field
+			// writes here and in later subtests.
+			cm.setConnecting(true)
 			cm.State = tc.initialState
 			cm.handleEvent(tc.event)
 			assert.Equal(t, tc.expectedState, cm.State, "State should match expected after event")
@@ -221,12 +230,56 @@ func TestConnectionManager_ReconnectionLogic(t *testing.T) {
 	// Test that isConnecting flag prevents duplicate reconnection attempts
 	// Start from connected state, then simulate disconnect
 	cm.State = WebSocketConnected
-	cm.isConnecting = false
+	cm.setConnecting(false)
 
 	// First disconnect should trigger reconnection logic
 	cm.handleStateChange(Disconnected)
 	assert.Equal(t, Disconnected, cm.State, "Should change to disconnected")
-	assert.True(t, cm.isConnecting, "Should set isConnecting flag")
+	assert.True(t, cm.isConnectingNow(), "Should set isConnecting flag")
+}
+
+// TestConnectionManager_TickerSurvivesStaleDisconnect reproduces the freeze from
+// https://github.com/henrygd/beszel/issues/2326: a reconnect attempt's handshake
+// can fail asynchronously (after connect() already returned with a nil error)
+// while the manager is still in the Disconnected state. Previously the ticker
+// was only re-armed from connect()'s synchronous error branch, so once that
+// window was missed, the agent stopped retrying forever. The ticker must keep
+// running any time the manager transitions into Disconnected, regardless of
+// what happens to the in-flight handshake afterwards.
+func TestConnectionManager_TickerSurvivesStaleDisconnect(t *testing.T) {
+	agent := createTestAgent(t)
+	cm := agent.connectionManager
+	cm.eventChan = make(chan ConnectionEvent, 1)
+
+	// Run on synctest's fake clock so the ticker fires without waiting a real
+	// wsTickerInterval. The ticker must be created inside the bubble.
+	synctest.Test(t, func(t *testing.T) {
+		// Simulate a healthy WebSocket connection, then a disconnect - mirroring
+		// handleStateChange's own Disconnected branch, but without launching the
+		// real async connect() goroutine so the ticker state can be asserted
+		// deterministically.
+		cm.State = WebSocketConnected
+		cm.stopWsTicker()
+		cm.setConnecting(true)
+		cm.handleStateChange(Disconnected)
+		require.NotNil(t, cm.wsTicker, "ticker must be armed as soon as the manager becomes Disconnected")
+		defer cm.stopWsTicker()
+
+		// Now simulate connect()'s in-flight handshake dying asynchronously with the
+		// manager still Disconnected (e.g. a late OnClose on an unauthenticated
+		// connection). This event is dropped by handleEvent since State is not
+		// WebSocketConnected, but the ticker armed above must still be running so
+		// the manager keeps retrying.
+		cm.setConnecting(false)
+		cm.handleEvent(WebSocketDisconnect)
+		assert.Equal(t, Disconnected, cm.State)
+
+		select {
+		case <-cm.wsTicker.C:
+		case <-time.After(wsTickerInterval + 2*time.Second):
+			t.Fatal("ticker did not fire after a stale disconnect event - agent would freeze forever")
+		}
+	})
 }
 
 // TestConnectionManager_ConnectWithRateLimit tests connection rate limiting