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

agent/connection_manager.go

@@ -8,6 +8,7 @@ import (
 	"os"
 	"os/signal"
 	"strings"
+	"sync"
 	"syscall"
 	"time"
 
@@ -20,7 +21,10 @@ import (
 // It handles both WebSocket and SSH connections, automatically switching between
 // them based on availability and managing reconnection attempts.
 type ConnectionManager struct {
-	agent          *Agent               // Reference to the parent agent
+	agent *Agent // Reference to the parent agent
+	// mu guards State and isConnecting, which are read and written from both
+	// the main event loop and the goroutine spawned by connect().
+	mu             sync.Mutex
 	State          ConnectionState      // Current connection state
 	eventChan      chan ConnectionEvent // Channel for connection events
 	wsClient       *WebSocketClient     // WebSocket client for hub communication
@@ -78,6 +82,29 @@ func (c *ConnectionManager) stopWsTicker() {
 	}
 }
 
+// getState returns the current connection state.
+func (c *ConnectionManager) getState() ConnectionState {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	return c.State
+}
+
+// setConnecting sets the isConnecting flag and reports its previous value.
+func (c *ConnectionManager) setConnecting(v bool) (previous bool) {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	previous = c.isConnecting
+	c.isConnecting = v
+	return previous
+}
+
+// isConnectingNow reports whether a reconnection attempt is currently in flight.
+func (c *ConnectionManager) isConnectingNow() bool {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	return c.isConnecting
+}
+
 // Start begins connection attempts and enters the main event loop.
 // It handles connection events, periodic health updates, and graceful shutdown.
 func (c *ConnectionManager) Start(serverOptions ServerOptions) error {
@@ -122,7 +149,10 @@ func (c *ConnectionManager) Start(serverOptions ServerOptions) error {
 		case connectionEvent := <-c.eventChan:
 			c.handleEvent(connectionEvent)
 		case <-c.wsTicker.C:
-			_ = c.startWebSocketConnection()
+			// skip if connect() is still running its own attempt
+			if !c.isConnectingNow() {
+				_ = c.startWebSocketConnection()
+			}
 		case <-healthTicker:
 			_ = health.Update()
 		case <-sigCtx.Done():
@@ -165,15 +195,15 @@ func (c *ConnectionManager) handleEvent(event ConnectionEvent) {
 	case WebSocketConnect:
 		c.handleStateChange(WebSocketConnected)
 	case SSHConnect:
-		if c.State == Disconnected {
+		if c.getState() == Disconnected {
 			c.handleStateChange(SSHConnected)
 		}
 	case WebSocketDisconnect:
-		if c.State == WebSocketConnected {
+		if c.getState() == WebSocketConnected {
 			c.handleStateChange(Disconnected)
 		}
 	case SSHDisconnect:
-		if c.State == SSHConnected {
+		if c.getState() == SSHConnected {
 			c.handleStateChange(Disconnected)
 		}
 	}
@@ -182,30 +212,40 @@ func (c *ConnectionManager) handleEvent(event ConnectionEvent) {
 // handleStateChange updates the connection state and performs necessary actions
 // based on the new state, including stopping services and initiating reconnections.
 func (c *ConnectionManager) handleStateChange(newState ConnectionState) {
+	c.mu.Lock()
 	if c.State == newState {
+		c.mu.Unlock()
 		return
 	}
 	c.State = newState
+	c.mu.Unlock()
+
 	switch newState {
 	case WebSocketConnected:
 		slog.Info("WebSocket connected", "host", c.wsClient.hubURL.Host)
 		c.ConnectionType = system.ConnectionTypeWebSocket
 		c.stopWsTicker()
 		_ = c.agent.StopServer()
-		c.isConnecting = false
+		c.setConnecting(false)
 	case SSHConnected:
 		// stop new ws connection attempts
 		slog.Info("SSH connection established")
 		c.ConnectionType = system.ConnectionTypeSSH
 		c.stopWsTicker()
-		c.isConnecting = false
+		c.setConnecting(false)
 	case Disconnected:
 		c.ConnectionType = system.ConnectionTypeNone
-		if c.isConnecting {
+		// Always keep the ticker running while disconnected. A pending WebSocket
+		// handshake started by connect() can fail asynchronously (e.g. the hub
+		// closes the socket, or the deadline set in OnOpen expires) after
+		// connect() has already returned with a nil error, in which case the
+		// ticker would otherwise never get re-armed and the agent would stop
+		// retrying entirely (#2326).
+		c.startWsTicker()
+		if c.setConnecting(true) {
 			// Already handling reconnection, avoid duplicate attempts
 			return
 		}
-		c.isConnecting = true
 		slog.Warn("Disconnected from hub")
 		// make sure old ws connection is closed
 		c.closeWebSocket()
@@ -217,10 +257,8 @@ func (c *ConnectionManager) handleStateChange(newState ConnectionState) {
 // connect handles the connection logic with proper delays and priority.
 // It attempts WebSocket connection first, falling back to SSH server if needed.
 func (c *ConnectionManager) connect() {
-	c.isConnecting = true
+	c.setConnecting(true)
-	defer func() {
+	defer c.setConnecting(false)
-		c.isConnecting = false
-	}()
 
 	if c.wsClient != nil && time.Since(c.wsClient.lastConnectAttempt) < 5*time.Second {
 		time.Sleep(5 * time.Second)
@@ -234,7 +272,7 @@ func (c *ConnectionManager) connect() {
 			_ = c.stop()
 			os.Exit(1)
 		}
-		if c.State == Disconnected {
+		if c.getState() == Disconnected {
 			c.startSSHServer()
 			c.startWsTicker()
 		}
@@ -243,7 +281,7 @@ func (c *ConnectionManager) connect() {
 
 // startWebSocketConnection attempts to establish a WebSocket connection to the hub.
 func (c *ConnectionManager) startWebSocketConnection() error {
-	if c.State != Disconnected {
+	if c.getState() != Disconnected {
 		return errors.New("already connected")
 	}
 	if c.wsClient == nil {
@@ -263,7 +301,7 @@ func (c *ConnectionManager) startWebSocketConnection() error {
 
 // startSSHServer starts the SSH server if the agent is currently disconnected.
 func (c *ConnectionManager) startSSHServer() {
-	if c.State == Disconnected {
+	if c.getState() == Disconnected {
 		go c.agent.StartServer(c.serverOptions)
 	}
 }

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