From 913a2edad938c36b55e9fe6d11b4bba5da422d9b Mon Sep 17 00:00:00 2001
From: Sathish Gangichetty <datasciencemonkey@gmail.com>
Date: Wed, 11 Mar 2026 19:35:42 -0400
Subject: [PATCH] fix: check actual WebSocket transport before stopping
 poll-worker (#54)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Socket.IO reports connected=true even when falling back to HTTP
long-polling through the Databricks Apps reverse proxy. The app was
prematurely stopping the poll-worker, leaving users with no data
transport when true WebSocket wasn't available.

Now checks socket.io.engine.transport.name before deciding:
- 'websocket' → stop poll-worker, use WS as primary
- 'polling' → keep poll-worker active as primary transport
- Listen for late 'upgrade' event if transport upgrades later

Cherry-picked from PR #52 (dgokeeffe).

Fixes #54

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
---
 static/index.html | 51 +++++++++++++++++++++++++++++++++++++----------
 1 file changed, 41 insertions(+), 10 deletions(-)

diff --git a/static/index.html b/static/index.html
index ac87db1..0df0b68 100644
--- a/static/index.html
+++ b/static/index.html
@@ -893,21 +893,52 @@ <h3>General</h3>
       socket = io({ transports: ['websocket', 'polling'] });
 
       socket.on('connect', () => {
-        console.log('[ws] Connected');
-        wsConnected = true;
-        // Re-join rooms for all active panes and stop their HTTP polling (AC-16)
+        // Check actual transport — Socket.IO reports connected=true even on long-polling
+        // through Databricks proxy. Only stop poll-worker for true WebSocket.
+        const transport = socket.io.engine.transport.name;
+        const isTrueWS = transport === 'websocket';
+        console.log(`[ws] Connected (transport: ${transport}, trueWS: ${isTrueWS})`);
+
+        // Always join rooms regardless of transport
         getAllPanes().forEach(p => {
           if (p.sessionId) {
             socket.emit('join_session', { session_id: p.sessionId });
-            pollWorker.postMessage({ type: 'stop_poll', paneId: p.id });
           }
         });
-        // Start WS heartbeat — keeps idle sessions alive (replaces poll-worker keepalive)
-        if (wsHeartbeatTimer) clearInterval(wsHeartbeatTimer);
-        wsHeartbeatTimer = setInterval(() => {
-          const sids = getAllPanes().map(p => p.sessionId).filter(Boolean);
-          if (sids.length > 0) socket.emit('heartbeat', { session_ids: sids });
-        }, WS_HEARTBEAT_INTERVAL);
+
+        if (isTrueWS) {
+          wsConnected = true;
+          // Only stop poll-worker when we have a real WebSocket
+          getAllPanes().forEach(p => {
+            if (p.id) pollWorker.postMessage({ type: 'stop_poll', paneId: p.id });
+          });
+          // Start WS heartbeat — keeps idle sessions alive (replaces poll-worker keepalive)
+          if (wsHeartbeatTimer) clearInterval(wsHeartbeatTimer);
+          wsHeartbeatTimer = setInterval(() => {
+            const sids = getAllPanes().map(p => p.sessionId).filter(Boolean);
+            if (sids.length > 0) socket.emit('heartbeat', { session_ids: sids });
+          }, WS_HEARTBEAT_INTERVAL);
+        } else {
+          console.log('[ws] Connected via polling — keeping poll-worker active');
+          // Don't set wsConnected — poll-worker stays active as primary transport
+        }
+
+        // Listen for late upgrade from polling → websocket
+        socket.io.engine.on('upgrade', (transport) => {
+          console.log(`[ws] Transport upgraded to: ${transport.name}`);
+          if (transport.name === 'websocket') {
+            wsConnected = true;
+            getAllPanes().forEach(p => {
+              if (p.id) pollWorker.postMessage({ type: 'stop_poll', paneId: p.id });
+            });
+            if (!wsHeartbeatTimer) {
+              wsHeartbeatTimer = setInterval(() => {
+                const sids = getAllPanes().map(p => p.sessionId).filter(Boolean);
+                if (sids.length > 0) socket.emit('heartbeat', { session_ids: sids });
+              }, WS_HEARTBEAT_INTERVAL);
+            }
+          }
+        });
       });
 
       socket.on('disconnect', (reason) => {