fix(web): recover from stale background websocket (#1451)

.changeset/bg-ws-stale-recovery.md

@@ -0,0 +1,5 @@
+---
+"@moonshot-ai/kimi-code": patch
+---
+
+web: Recover chat streaming after a stale background-tab WebSocket instead of requiring a page refresh.

apps/kimi-web/src/api/daemon/client.ts

@@ -1397,6 +1397,12 @@ export class DaemonKimiWebApi implements KimiWebApi {
       markSideChannelAgent(agentId: string): void {
         projector.markSideChannelAgent(agentId);
       },
+      health(): { connected: boolean; open: boolean; stale: boolean } {
+        return socket.health();
+      },
+      reconnect(): void {
+        socket.reconnect();
+      },
       close(): void {
         socket.close();
       },

apps/kimi-web/src/api/daemon/ws.ts

@@ -13,6 +13,11 @@ import type { WireEvent, WireServerFrame } from './wire';
 // Sec-WebSocket-Protocol subprotocol instead.
 const WS_BEARER_PROTOCOL_PREFIX = 'kimi-code.bearer.';
 
+// A socket with no incoming frames for this long is presumed half-open even if
+// the browser still reports OPEN (no onclose fired). Derived as 2x the server
+// heartbeat, with a floor so a misconfigured tiny heartbeat can't thrash.
+const STALE_SOCKET_FLOOR_MS = 30_000;
+
 // ---------------------------------------------------------------------------
 // Handler interface
 // ---------------------------------------------------------------------------
@@ -84,6 +89,14 @@ export class DaemonEventSocket {
   private reconnectAttempts = 0;
   private reconnectTimer: ReturnType<typeof setTimeout> | null = null;
 
+  /** Server-advertised heartbeat interval (ms); falls back to the daemon default. */
+  private heartbeatMs = 30_000;
+  /**
+   * Epoch ms of the most recent frame (or the connect attempt). Used to detect
+   * a silent-half-open socket that the browser never fires `onclose` for.
+   */
+  private lastActivityAt = 0;
+
   constructor(
     private readonly wsUrl: string,
     private readonly clientId: string,
@@ -94,6 +107,7 @@ export class DaemonEventSocket {
   connect(): void {
     if (this.ws !== null || this.closed) return;
 
+    this.lastActivityAt = Date.now();
     traceWsLifecycle('connect', { url: this.wsUrl, attempt: this.reconnectAttempts });
     const credential = getCredential();
     const protocols =
@@ -107,6 +121,8 @@ export class DaemonEventSocket {
     };
 
     ws.onmessage = (ev: MessageEvent) => {
+      // Any received frame proves the link is alive; reset the stale detector.
+      this.lastActivityAt = Date.now();
       try {
         const frame = JSON.parse(String(ev.data)) as WireServerFrame;
         traceWsIn(frame);
@@ -256,6 +272,61 @@ export class DaemonEventSocket {
     }
   }
 
+  /**
+   * Snapshot the socket's health. `stale` is true when no frame has arrived for
+   * longer than 2x the server heartbeat (floored at {@link STALE_SOCKET_FLOOR_MS}).
+   * The browser may still report OPEN on a half-open connection that no longer
+   * delivers data, so foreground recovery keys on the staleness signal rather
+   * than the raw readyState.
+   */
+  health(): { connected: boolean; open: boolean; stale: boolean } {
+    const open = this.ws !== null && this.ws.readyState === WebSocket.OPEN;
+    const threshold = Math.max(this.heartbeatMs * 2, STALE_SOCKET_FLOOR_MS);
+    const stale = this.lastActivityAt > 0 && Date.now() - this.lastActivityAt > threshold;
+    return { connected: this.connected, open, stale };
+  }
+
+  /**
+   * Force a clean reconnect. Used to recover from a silent-half-open socket
+   * (e.g. after the browser froze a background tab) where `onclose` never
+   * fires, so the automatic backoff reconnect wired into `onclose` is never
+   * triggered.
+   *
+   * Tears down the current socket without waiting for `onclose`, resets the
+   * handshake state, and opens a fresh socket immediately; `onServerHello`
+   * re-sends every subscription at the last durable cursor. No-op after
+   * {@link close()}.
+   */
+  reconnect(): void {
+    if (this.closed) return;
+    // Cancel any pending automatic reconnect — we're reconnecting synchronously.
+    if (this.reconnectTimer !== null) {
+      clearTimeout(this.reconnectTimer);
+      this.reconnectTimer = null;
+    }
+    const old = this.ws;
+    if (old !== null) {
+      // Detach before closing so the old socket's `onclose` doesn't race our
+      // fresh connect (it would call scheduleReconnect and clobber `this.ws`).
+      old.onopen = null;
+      old.onmessage = null;
+      old.onerror = null;
+      old.onclose = null;
+      try {
+        old.close(1000, 'reconnect');
+      } catch {
+        // Ignore — the socket may already be closing.
+      }
+    }
+    const wasConnected = this.connected;
+    this.ws = null;
+    this.connected = false;
+    if (wasConnected) {
+      this.handlers.onConnectionState(false);
+    }
+    this.connect();
+  }
+
   // ---------------------------------------------------------------------------
   // Private helpers
   // ---------------------------------------------------------------------------
@@ -266,9 +337,12 @@ export class DaemonEventSocket {
     // eslint-disable-next-line @typescript-eslint/no-explicit-any
     const frame = rawFrame as any;
     switch ((rawFrame as { type: string }).type) {
-      case 'server_hello':
+      case 'server_hello': {
+        const hb = (frame.payload as { heartbeat_ms?: unknown } | undefined)?.heartbeat_ms;
+        if (typeof hb === 'number' && hb > 0) this.heartbeatMs = hb;
         this.onServerHello();
         break;
+      }
 
       case 'ping':
         this.send({ type: 'pong', payload: { nonce: frame.payload.nonce } });

apps/kimi-web/src/api/types.ts

@@ -539,6 +539,19 @@ export interface KimiEventConnection {
    * instead of dropping them like background subagents.
    */
   markSideChannelAgent(agentId: string): void;
+  /**
+   * Report the underlying socket's health. Used to detect a silent-half-open
+   * connection after the tab was frozen in the background: the browser still
+   * reports OPEN (so no auto-reconnect) yet no frames have arrived for a while.
+   */
+  health(): { connected: boolean; open: boolean; stale: boolean };
+  /**
+   * Force a clean reconnect of the underlying socket. Used to recover from a
+   * silent-half-open (background-tab freeze) where onclose never fires. The
+   * reconnect handshake re-subscribes at the last durable cursor. No-op after
+   * close().
+   */
+  reconnect(): void;
   close(): void;
 }
 

apps/kimi-web/src/composables/useKimiWebClient.ts

@@ -4,6 +4,7 @@
 
 import { computed, reactive, ref, watch } from 'vue';
 import { i18n } from '../i18n';
+import { traceClientEvent } from '../debug/trace';
 import { getKimiWebApi } from '../api';
 import { isDaemonApiError, isDaemonNetworkError } from '../api/errors';
 import {
@@ -467,10 +468,34 @@ if (typeof window !== 'undefined') {
   });
 }
 
+/**
+ * When the tab returns to the foreground, the WebSocket may be a silent
+ * half-open: the browser still reports OPEN (so no auto-reconnect) yet no
+ * frames have arrived for a while (frozen background tab, dropped NAT mapping,
+ * daemon restart). On such a socket live streaming tokens freeze mid-turn with
+ * no recovery short of a full page reload.
+ *
+ * If the socket looks stale, force a clean reconnect — the handshake
+ * re-subscribes at the last durable cursor — then refresh the active session
+ * from its authoritative snapshot to re-seed the volatile streaming tokens lost
+ * during the gap.
+ */
+function recoverStaleConnection(): void {
+  if (eventConn === null) return;
+  if (!eventConn.health().stale) return;
+  traceClientEvent('ws: stale socket on focus, reconnecting', {
+    activeSessionId: rawState.activeSessionId,
+  });
+  eventConn.reconnect();
+  const active = rawState.activeSessionId;
+  if (active) snapshotSyncRunner.request(active);
+}
+
 if (typeof document !== 'undefined') {
   document.addEventListener('visibilitychange', () => {
     if (document.visibilityState === 'visible') {
       clearActiveUnread();
+      recoverStaleConnection();
     }
   });
 }

apps/kimi-web/test/ws-lifecycle.test.ts

@@ -0,0 +1,146 @@
+// apps/kimi-web/test/ws-lifecycle.test.ts
+// Focused coverage of DaemonEventSocket reconnect + staleness detection, the
+// foreground recovery path added so a frozen/backgrounded tab can recover
+// without a full page reload.
+
+import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest';
+
+import { DaemonEventSocket, type DaemonEventSocketHandlers } from '../src/api/daemon/ws';
+
+class FakeWebSocket {
+  static readonly CONNECTING = 0;
+  static readonly OPEN = 1;
+  static readonly CLOSING = 2;
+  static readonly CLOSED = 3;
+  static instances: FakeWebSocket[] = [];
+
+  readyState = FakeWebSocket.OPEN;
+  onopen: (() => void) | null = null;
+  onmessage: ((ev: { data: unknown }) => void) | null = null;
+  onerror: (() => void) | null = null;
+  onclose: ((ev?: { code: number; reason: string; wasClean: boolean }) => void) | null = null;
+  sent: string[] = [];
+  closeCalls: Array<{ code?: number; reason?: string }> = [];
+
+  constructor(
+    public readonly url: string,
+    public readonly protocols?: string | string[],
+  ) {
+    FakeWebSocket.instances.push(this);
+  }
+
+  send(data: string): void {
+    this.sent.push(data);
+  }
+
+  close(code?: number, reason?: string): void {
+    this.closeCalls.push({ code, reason });
+    this.readyState = FakeWebSocket.CLOSED;
+  }
+
+  emitMessage(frame: unknown): void {
+    this.onmessage?.({ data: JSON.stringify(frame) });
+  }
+}
+
+function makeHandlers(): DaemonEventSocketHandlers & { states: boolean[] } {
+  const states: boolean[] = [];
+  return {
+    states,
+    onWireEvent: () => {},
+    onResync: () => {},
+    onConnectionState: (connected) => states.push(connected),
+    onError: () => {},
+  };
+}
+
+const WS_URL = 'ws://example.test/ws';
+const CLIENT_ID = 'client_test';
+
+// Frames the socket understands; only `type` (+ friends) matters here.
+const SERVER_HELLO = {
+  type: 'server_hello',
+  payload: {
+    ws_connection_id: 'conn_1',
+    protocol_version: 1,
+    heartbeat_ms: 30_000,
+    max_event_buffer_size: 1000,
+    capabilities: {},
+  },
+};
+
+describe('DaemonEventSocket reconnect + staleness', () => {
+  let originalWebSocket: typeof globalThis.WebSocket;
+
+  beforeEach(() => {
+    FakeWebSocket.instances = [];
+    originalWebSocket = globalThis.WebSocket;
+    globalThis.WebSocket = FakeWebSocket as unknown as typeof WebSocket;
+  });
+
+  afterEach(() => {
+    globalThis.WebSocket = originalWebSocket;
+    vi.useRealTimers();
+  });
+
+  it('reconnect() closes the old socket, detaches it, and opens a new one', () => {
+    const handlers = makeHandlers();
+    const socket = new DaemonEventSocket(WS_URL, CLIENT_ID, handlers);
+    socket.connect();
+    const first = FakeWebSocket.instances[0]!;
+    first.emitMessage(SERVER_HELLO);
+    expect(handlers.states).toEqual([true]);
+
+    socket.reconnect();
+
+    expect(first.closeCalls).toEqual([{ code: 1000, reason: 'reconnect' }]);
+    // Old socket is fully detached so its late onclose cannot clobber the new.
+    expect(first.onclose).toBeNull();
+    expect(first.onmessage).toBeNull();
+    // A fresh socket was created, and we reported the transient disconnect.
+    expect(FakeWebSocket.instances).toHaveLength(2);
+    expect(handlers.states).toEqual([true, false]);
+
+    // A late onclose from the stale socket must NOT schedule another connect.
+    first.onclose?.({ code: 1000, reason: 'reconnect', wasClean: true });
+    expect(FakeWebSocket.instances).toHaveLength(2);
+  });
+
+  it('reconnect() is a no-op after close()', () => {
+    const handlers = makeHandlers();
+    const socket = new DaemonEventSocket(WS_URL, CLIENT_ID, handlers);
+    socket.connect();
+    socket.close();
+    socket.reconnect();
+    expect(FakeWebSocket.instances).toHaveLength(1);
+  });
+
+  it('health() flips stale after a long silence and clears on the next frame', () => {
+    vi.useFakeTimers();
+    const handlers = makeHandlers();
+    const socket = new DaemonEventSocket(WS_URL, CLIENT_ID, handlers);
+    socket.connect();
+    const first = FakeWebSocket.instances[0]!;
+    first.emitMessage(SERVER_HELLO);
+
+    // Threshold = max(2 * 30_000, 30_000 floor) = 60s.
+    expect(socket.health().stale).toBe(false);
+
+    vi.advanceTimersByTime(61_000);
+    expect(socket.health().stale).toBe(true);
+
+    // Any received frame (e.g. a server ping) proves the link is alive again.
+    first.emitMessage({ type: 'ping', payload: { nonce: 'n1' } });
+    expect(socket.health().stale).toBe(false);
+  });
+
+  it('health().open reflects the underlying readyState', () => {
+    const handlers = makeHandlers();
+    const socket = new DaemonEventSocket(WS_URL, CLIENT_ID, handlers);
+    socket.connect();
+    const first = FakeWebSocket.instances[0]!;
+    expect(socket.health().open).toBe(true);
+    first.readyState = FakeWebSocket.CLOSING;
+    expect(socket.health().open).toBe(false);
+  });
+});