fix(agent): improve WebSocket and SSH fallback handling (#2441)

Make WebSocket reconnect and SSH fallback transitions reliable across asynchronous disconnects, stale callbacks, and overlapping connections. Keep the SSH listener available while disconnected and allow a verified WebSocket connection to take precedence when it recovers.

Co-authored-by: henrygd <hank@henrygd.me>

agent/connection_manager_test.go

@@ -12,6 +12,7 @@ import (
 	"testing/synctest"
 	"time"
 
+	"github.com/lxzan/gws"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 	"golang.org/x/crypto/ssh"
@@ -95,6 +96,7 @@ func TestConnectionManager_StateTransitions(t *testing.T) {
 func TestConnectionManager_EventHandling(t *testing.T) {
 	agent := createTestAgent(t)
 	cm := agent.connectionManager
+	t.Setenv("BESZEL_AGENT_DISABLE_SSH", "true")
 	cm.wsClient = &WebSocketClient{
 		hubURL: &url.URL{
 			Host: "localhost:8080",
@@ -112,6 +114,12 @@ func TestConnectionManager_EventHandling(t *testing.T) {
 			event:         WebSocketConnect,
 			expectedState: WebSocketConnected,
 		},
+		{
+			name:          "WebSocket connect from SSH connected",
+			initialState:  SSHConnected,
+			event:         WebSocketConnect,
+			expectedState: WebSocketConnected,
+		},
 		{
 			name:          "SSH connect from disconnected",
 			initialState:  Disconnected,
@@ -157,6 +165,20 @@ func TestConnectionManager_EventHandling(t *testing.T) {
 			// and the goroutine would otherwise race with the direct field
 			// writes here and in later subtests.
 			cm.setConnecting(true)
+			cm.mu.Lock()
+			cm.sshConnections = 0
+			if tc.event == SSHConnect {
+				cm.sshConnections = 1
+			}
+			cm.mu.Unlock()
+			cm.wsClient.connMu.Lock()
+			cm.wsClient.Conn = nil
+			cm.wsClient.hubVerified = false
+			if tc.event == WebSocketConnect {
+				cm.wsClient.Conn = &gws.Conn{}
+				cm.wsClient.hubVerified = true
+			}
+			cm.wsClient.connMu.Unlock()
 			cm.State = tc.initialState
 			cm.handleEvent(tc.event)
 			assert.Equal(t, tc.expectedState, cm.State, "State should match expected after event")
@@ -238,6 +260,24 @@ func TestConnectionManager_ReconnectionLogic(t *testing.T) {
 	assert.True(t, cm.isConnectingNow(), "Should set isConnecting flag")
 }
 
+func TestWebSocketDisconnectStartsSSH(t *testing.T) {
+	t.Setenv("BESZEL_AGENT_DISABLE_SSH", "false")
+	agent := createTestAgent(t)
+	cm := agent.connectionManager
+	cm.serverOptions = createTestServerOptions(t)
+	cm.State = WebSocketConnected
+	cm.setConnecting(true) // keep this test focused on the synchronous fallback
+	defer cm.stopWsTicker()
+
+	cm.handleEvent(WebSocketDisconnect)
+	require.Equal(t, Disconnected, cm.getState())
+	agent.serverMu.Lock()
+	listener := agent.serverListener
+	agent.serverMu.Unlock()
+	require.NotNil(t, listener, "SSH should be ready as soon as an established WS closes")
+	require.NoError(t, agent.StopServer())
+}
+
 // 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)