fix: Stabilize websocket connection (#166)

Author: christianalfoni

package.json

@@ -90,12 +90,12 @@
     "tslib": "^2.1.0",
     "typescript": "^5.7.2",
     "why-is-node-running": "^2.3.0",
-    "isomorphic-ws": "^5.0.0"
+    "isomorphic-ws": "^5.0.0",
+    "@msgpack/msgpack": "^2.7.1"
   },
   "dependencies": {
     "@hey-api/client-fetch": "^0.7.3",
     "@inkjs/ui": "^2.0.0",
-    "@msgpack/msgpack": "^3.1.0",
     "@opentelemetry/api": "^1.9.0",
     "@xterm/addon-serialize": "^0.13.0",
     "@xterm/headless": "^5.5.0",

src/AgentClient/AgentConnection.ts

@@ -93,8 +93,15 @@ export class AgentConnection {
     });
 
     connection.onMissingHeartbeat(() => {
+      // Be more conservative about disconnection - only disconnect if we have no activity
+      // and no pending messages, indicating a truly dead connection
       if (this.pendingMessages.size === 0) {
-        this.state = "DISCONNECTED";
+        // Add a small delay to allow for network recovery before declaring disconnection
+        setTimeout(() => {
+          if (this.pendingMessages.size === 0 && this.state === "CONNECTED") {
+            this.state = "DISCONNECTED";
+          }
+        }, 1000);
       }
     });
   }

src/AgentClient/WebSocketClient.ts

@@ -235,6 +235,9 @@ export class WebSocketClient extends Disposable {
       );
     }
 
+    // Update lastActivity on send to prevent heartbeat suppression
+    this.lastActivity = Date.now();
+
     // This is an async operation in Node, but to avoid wrapping every send in a promise, we
     // rely on the error listener to deal with any errors. Any unsent messages will be timed out
     // by our PendingMessage logic

src/AgentClient/index.ts

@@ -25,7 +25,8 @@ import { SandboxClient } from "../SandboxClient";
 import { InitStatus } from "../pitcher-protocol/messages/system";
 
 // Timeout for detecting a pong response, leading to a forced disconnect
-let PONG_DETECTION_TIMEOUT = 15_000;
+// Increased from 15s to 30s to be more tolerant of network latency
+let PONG_DETECTION_TIMEOUT = 30_000;
 
 // When focusing the app we do a lower timeout to more quickly detect a potential disconnect
 const FOCUS_PONG_DETECTION_TIMEOUT = 5_000;

src/SandboxClient/index.ts

@@ -195,6 +195,8 @@ export class SandboxClient {
           this.attemptAutoReconnect();
         }
       } else if (state === "CONNECTED") {
+        // Reset keep-alive failures on successful connection
+        this.keepAliveFailures = 0;
         if (this.shouldKeepAlive) {
           this.keepActiveWhileConnected(true);
         }
@@ -389,6 +391,8 @@ export class SandboxClient {
   private keepAliveInterval: NodeJS.Timeout | null = null;
   private shouldKeepAlive = false;
   private isExplicitlyDisconnected = false;
+  private keepAliveFailures = 0;
+  private maxKeepAliveFailures = 3;
   /**
    * If enabled, we will keep the sandbox from hibernating as long as the SDK is connected to it.
    */
@@ -399,9 +403,31 @@ export class SandboxClient {
     if (enabled) {
       if (!this.keepAliveInterval) {
         this.keepAliveInterval = setInterval(() => {
-          this.agentClient.system.update().catch((error) => {
-            console.warn("Unable to keep active while connected", error);
-          });
+          this.agentClient.system.update()
+            .then(() => {
+              // Reset failure count on success
+              this.keepAliveFailures = 0;
+            })
+            .catch((error) => {
+              this.keepAliveFailures++;
+              console.warn(`Keep-alive failed (${this.keepAliveFailures}/${this.maxKeepAliveFailures}):`, error);
+              
+              // If we've hit max failures, stop aggressive keep-alive to prevent connection thrashing
+              if (this.keepAliveFailures >= this.maxKeepAliveFailures) {
+                console.warn("Max keep-alive failures reached, reducing frequency to prevent connection issues");
+                if (this.keepAliveInterval) {
+                  clearInterval(this.keepAliveInterval);
+                  this.keepAliveInterval = null;
+                }
+                // Restart with longer interval after failures
+                setTimeout(() => {
+                  if (this.shouldKeepAlive && !this.keepAliveInterval) {
+                    this.keepActiveWhileConnected(true);
+                    this.keepAliveFailures = 0; // Reset for retry
+                  }
+                }, 60000); // Wait 1 minute before retrying
+              }
+            });
         }, 1000 * 10);
       }
     } else {

src/utils/api.ts

@@ -68,6 +68,15 @@ export type HandledResponse<D, E> = {
 export function getStartOptions(opts: StartSandboxOpts | undefined) {
   if (!opts) return {};
 
+  // Warn about hibernation timeouts that are too short and may cause connection issues
+  if (opts.hibernationTimeoutSeconds !== undefined && opts.hibernationTimeoutSeconds < 60) {
+    console.warn(
+      `Warning: hibernationTimeoutSeconds (${opts.hibernationTimeoutSeconds}s) is less than 60 seconds. ` +
+      `This may cause connection instability and frequent disconnections. ` +
+      `Consider using at least 60 seconds for stable websocket connections.`
+    );
+  }
+
   return {
     ipcountry: opts.ipcountry,
     tier: opts.vmTier?.name,