fix(kap-server): add WebSocket heartbeat to survive proxy idle timeouts

.changeset/ws-heartbeat-keepalive.md

@@ -0,0 +1,5 @@
+---
+"@moonshot-ai/kimi-code": patch
+---
+
+Fix the web UI repeatedly losing its realtime connection every ~30 seconds when the server runs behind a reverse proxy or gateway with an idle connection timeout; the server now sends a WebSocket heartbeat and only closes connections that stop responding entirely.

packages/kap-server/src/protocol/ws-control.ts

@@ -77,8 +77,9 @@ export const serverHelloPayloadSchema = z.object({
   ws_connection_id: z.string(),
   protocol_version: z.number().int().positive(),
   /**
-   * Legacy servers advertise their ping interval here. kap-server dropped the
-   * server-initiated heartbeat and omits this field — clients must treat it as
+   * Server heartbeat interval. kap-server sends an application-level `ping`
+   * at this cadence and closes the connection after two silent cycles; older
+   * servers omit the field and send no heartbeat, so clients must treat it as
    * advisory and not require it.
    */
   heartbeat_ms: z.number().int().positive().optional(),

packages/kap-server/src/transport/ws/v1/protocol.ts

@@ -9,6 +9,8 @@
 export interface ServerHelloPayload {
   ws_connection_id: string;
   protocol_version: number;
+  /** Server heartbeat cadence — a `ping` frame arrives at least this often. */
+  heartbeat_ms: number;
   max_event_buffer_size: number;
   capabilities: {
     event_batching: boolean;
@@ -26,6 +28,16 @@ export function buildServerHello(payload: ServerHelloPayload): ServerHelloFrame
   return { type: 'server_hello', timestamp: new Date().toISOString(), payload };
 }
 
+export interface PingFrame {
+  type: 'ping';
+  timestamp: string;
+  payload: { nonce: string };
+}
+
+export function buildPing(nonce: string): PingFrame {
+  return { type: 'ping', timestamp: new Date().toISOString(), payload: { nonce } };
+}
+
 export interface AckFrame<P = unknown> {
   type: 'ack';
   id: string;

packages/kap-server/src/transport/ws/v1/registerWsV1.ts

@@ -32,6 +32,8 @@ export interface RegisterWsV1Options {
   readonly flushIntervalMs?: number;
   readonly maxBatchSize?: number;
   readonly highWaterMarkBytes?: number;
+  /** Heartbeat ping cadence override — tests inject small values. */
+  readonly heartbeatIntervalMs?: number;
 }
 
 export function registerWsV1(core: Scope, opts: RegisterWsV1Options): WebSocketServer {
@@ -53,6 +55,7 @@ export function registerWsV1(core: Scope, opts: RegisterWsV1Options): WebSocketS
       flushIntervalMs: opts.flushIntervalMs,
       maxBatchSize: opts.maxBatchSize,
       highWaterMarkBytes: opts.highWaterMarkBytes,
+      heartbeatIntervalMs: opts.heartbeatIntervalMs,
     });
     socket.on('close', () => registry.remove(conn.id));
   });

packages/kap-server/src/transport/ws/v1/wsConnectionV1.ts

@@ -12,10 +12,16 @@
  * them to the same shared attach path (`attachSession`). Transcript grade
  * subscriptions are a separate concern carried ONLY by `subscribe_v2`.
  *
- * The server never initiates a disconnect: unlike v1's `WsConnection`
- * (`packages/server/src/ws/connection.ts`) there is no ping/pong heartbeat —
- * a connection stays open until the client closes it or the process shuts
- * down.
+ * Heartbeat: the server sends an application-level `ping` frame every
+ * {@link DEFAULT_HEARTBEAT_INTERVAL_MS} (advertised as `heartbeat_ms` in
+ * `server_hello`). Protocol-level WS ping/pong is NOT used because browser
+ * clients cannot observe it from JS — an application-level frame is what
+ * feeds the client's stale-socket detector. Any inbound frame (a `pong`,
+ * but also ordinary control traffic) proves the peer is alive; after two
+ * full silent cycles the connection is presumed half-open (laptop asleep,
+ * network silently gone) and closed with 1001. Besides liveness this keeps
+ * intermediaries (reverse proxies with ~30s idle timeouts) from dropping
+ * idle connections.
  */
 
 import {
@@ -39,6 +45,7 @@ import {
 } from './sessionEventJournal';
 import {
   buildAck,
+  buildPing,
   buildResyncRequired,
   buildServerHello,
 } from './protocol';
@@ -54,6 +61,15 @@ import { FsWatchBridge } from './fsWatchBridge';
 
 const DEFAULT_MAX_BUFFER_SIZE = 1000;
 
+/**
+ * Application-level heartbeat cadence. 10s keeps connections alive through
+ * intermediaries with ~30s idle timeouts (3x headroom) and bounds how long a
+ * half-open connection goes unnoticed.
+ */
+const DEFAULT_HEARTBEAT_INTERVAL_MS = 10_000;
+/** Close the connection once no inbound frame has arrived for this many cycles. */
+const HEARTBEAT_MISS_LIMIT = 2;
+
 /** Per-session subscription state held by the connection (see `TargetSubscription`). */
 type SessionSubscription = TargetSubscription;
 
@@ -96,6 +112,8 @@ export interface WsConnectionV1Options {
   readonly maxBatchSize?: number;
   /** `socket.bufferedAmount` above which flushing is deferred (backpressure). */
   readonly highWaterMarkBytes?: number;
+  /** Heartbeat ping cadence; advertised as `heartbeat_ms` in `server_hello`. */
+  readonly heartbeatIntervalMs?: number;
 }
 
 export class WsConnectionV1 implements BroadcastTarget {
@@ -112,6 +130,7 @@ export class WsConnectionV1 implements BroadcastTarget {
   private readonly flushIntervalMs: number;
   private readonly maxBatchSize: number;
   private readonly highWaterMarkBytes: number;
+  private readonly heartbeatIntervalMs: number;
   private readonly logger?: JournalLogger;
 
   private closed = false;
@@ -134,6 +153,10 @@ export class WsConnectionV1 implements BroadcastTarget {
   /** Epoch ms when the current backpressure deferral started; caps the wait. */
   private backpressureSince?: number;
 
+  private heartbeatTimer?: ReturnType<typeof setInterval>;
+  /** Epoch ms of the most recent inbound frame — any frame proves the peer is alive. */
+  private lastInboundAt = Date.now();
+
   constructor(opts: WsConnectionV1Options) {
     this.id = `conn_${ulid()}`;
     this.connectedAt = new Date().toISOString();
@@ -148,6 +171,7 @@ export class WsConnectionV1 implements BroadcastTarget {
     this.flushIntervalMs = opts.flushIntervalMs ?? DEFAULT_FLUSH_INTERVAL_MS;
     this.maxBatchSize = opts.maxBatchSize ?? DEFAULT_MAX_BATCH_SIZE;
     this.highWaterMarkBytes = opts.highWaterMarkBytes ?? DEFAULT_HIGH_WATER_MARK_BYTES;
+    this.heartbeatIntervalMs = opts.heartbeatIntervalMs ?? DEFAULT_HEARTBEAT_INTERVAL_MS;
 
     this.socket.on('message', (data: RawData) => this.onMessage(data));
     this.socket.on('close', () => this.onClose());
@@ -162,10 +186,15 @@ export class WsConnectionV1 implements BroadcastTarget {
       buildServerHello({
         ws_connection_id: this.id,
         protocol_version: WS_PROTOCOL_VERSION,
+        heartbeat_ms: this.heartbeatIntervalMs,
         max_event_buffer_size: this.maxBufferSize,
         capabilities: { event_batching: false, compression: false },
       }),
     );
+    this.heartbeatTimer = setInterval(() => {
+      this.onHeartbeat();
+    }, this.heartbeatIntervalMs);
+    this.heartbeatTimer.unref?.();
   }
 
   get hasClientHello(): boolean {
@@ -191,8 +220,13 @@ export class WsConnectionV1 implements BroadcastTarget {
       return; // non-JSON frame — drop
     }
     if (typeof frame?.type !== 'string') return;
+    // Any well-formed inbound frame — pongs included — proves the peer is alive.
+    this.lastInboundAt = Date.now();
 
     switch (frame.type) {
+      case 'pong':
+        // Heartbeat reply; the liveness timestamp above is all it needs to do.
+        return;
       case 'client_hello':
         this.enqueueControl(() => this.onClientHello(frame));
         return;
@@ -227,6 +261,19 @@ export class WsConnectionV1 implements BroadcastTarget {
     });
   }
 
+  /**
+   * Heartbeat tick: reap first, ping second. A peer silent for two full cycles
+   * (no pong, no control traffic at all) is half-open — close it rather than
+   * ping a dead pipe. The close also fires the client's reconnect path.
+   */
+  private onHeartbeat(): void {
+    if (Date.now() - this.lastInboundAt >= this.heartbeatIntervalMs * HEARTBEAT_MISS_LIMIT) {
+      this.close(1001, 'heartbeat timeout');
+      return;
+    }
+    this.sendImmediateFrame(buildPing(ulid()));
+  }
+
   private async onClientHello(frame: InboundFrame): Promise<void> {
     if (!(await this.authorize(frame))) return;
     this.gotClientHello = true;
@@ -616,6 +663,7 @@ export class WsConnectionV1 implements BroadcastTarget {
     this.closed = true;
     if (this.flushTimer !== undefined) clearTimeout(this.flushTimer);
     if (this.backpressureRetryTimer !== undefined) clearTimeout(this.backpressureRetryTimer);
+    if (this.heartbeatTimer !== undefined) clearInterval(this.heartbeatTimer);
     this.outbound = [];
     this.broadcaster.removeGlobalTarget(this);
     for (const sid of this.subscriptions.keys()) this.broadcaster.unsubscribe(sid, this);

packages/kap-server/test/wsConnectionV1.test.ts

@@ -765,6 +765,116 @@ describe('WsConnectionV1 outbound buffer', () => {
   });
 });
 
+// ---------------------------------------------------------------------------
+// WsConnectionV1 — heartbeat
+// ---------------------------------------------------------------------------
+
+describe('WsConnectionV1 heartbeat', () => {
+  beforeEach(() => {
+    vi.useFakeTimers();
+  });
+  afterEach(() => {
+    vi.useRealTimers();
+  });
+
+  function sentTypes(socket: FakeSocket): string[] {
+    return socket.frames().map((f) => (f as { type: string }).type);
+  }
+
+  function sentPings(socket: FakeSocket): Array<{ type: string; payload: { nonce: string } }> {
+    return socket.frames() as Array<{ type: string; payload: { nonce: string } }>;
+  }
+
+  it('advertises the heartbeat interval in server_hello', () => {
+    const socket = new FakeSocket();
+    const conn = makeConn(socket, { heartbeatIntervalMs: 10 });
+    const hello = socket.frames()[0] as { type: string; payload: { heartbeat_ms?: number } };
+    expect(hello.type).toBe('server_hello');
+    expect(hello.payload.heartbeat_ms).toBe(10);
+    conn.close();
+  });
+
+  it('defaults to a 10s heartbeat interval', () => {
+    const socket = new FakeSocket();
+    const conn = makeConn(socket);
+    const hello = socket.frames()[0] as { payload: { heartbeat_ms?: number } };
+    expect(hello.payload.heartbeat_ms).toBe(10_000);
+    conn.close();
+  });
+
+  it('sends a ping every interval while the peer keeps answering', () => {
+    const socket = new FakeSocket();
+    const conn = makeConn(socket, { heartbeatIntervalMs: 10 });
+    socket.sent = [];
+
+    for (let i = 0; i < 3; i++) {
+      vi.advanceTimersByTime(10);
+      expect(sentTypes(socket)).toHaveLength(i + 1);
+      socket.emit('message', JSON.stringify({ type: 'pong', payload: { nonce: 'n' } }));
+    }
+
+    const pings = sentPings(socket);
+    expect(pings.every((f) => f.type === 'ping')).toBe(true);
+    expect(typeof pings[0]!.payload.nonce).toBe('string');
+    expect(new Set(pings.map((f) => f.payload.nonce)).size).toBe(3);
+    expect(socket.closeCalls).toHaveLength(0);
+    conn.close();
+  });
+
+  it('reaps the connection after two silent cycles', () => {
+    const socket = new FakeSocket();
+    const conn = makeConn(socket, { heartbeatIntervalMs: 10 });
+    socket.sent = [];
+
+    vi.advanceTimersByTime(10);
+    expect(sentTypes(socket)).toEqual(['ping']);
+    expect(socket.closeCalls).toHaveLength(0);
+
+    // Second silent cycle: the tick closes instead of pinging again.
+    vi.advanceTimersByTime(10);
+    expect(socket.closeCalls).toEqual([{ code: 1001, reason: 'heartbeat timeout' }]);
+    expect(sentTypes(socket)).toEqual(['ping']);
+
+    // The heartbeat stops with the connection.
+    vi.advanceTimersByTime(100);
+    expect(sentTypes(socket)).toEqual(['ping']);
+    expect(socket.closeCalls).toHaveLength(1);
+  });
+
+  it('treats any inbound frame — not just pong — as proof of life', () => {
+    const socket = new FakeSocket();
+    const conn = makeConn(socket, { heartbeatIntervalMs: 10 });
+    socket.sent = [];
+
+    // t=10: ping. t=15: an unknown control frame still resets the window.
+    vi.advanceTimersByTime(15);
+    socket.emit('message', JSON.stringify({ type: 'some_future_frame', payload: {} }));
+
+    // t=20 (silence 5) and t=30 (silence 15): pings, no reap.
+    vi.advanceTimersByTime(20);
+    expect(sentTypes(socket)).toEqual(['ping', 'ping', 'ping']);
+    expect(socket.closeCalls).toHaveLength(0);
+
+    // t=40: silence 25 ≥ 2 cycles — reaped.
+    vi.advanceTimersByTime(5);
+    expect(socket.closeCalls).toEqual([{ code: 1001, reason: 'heartbeat timeout' }]);
+  });
+
+  it('stops heartbeating once the socket closes on its own', () => {
+    const socket = new FakeSocket();
+    makeConn(socket, { heartbeatIntervalMs: 10 });
+    socket.sent = [];
+
+    vi.advanceTimersByTime(10);
+    expect(sentTypes(socket)).toEqual(['ping']);
+
+    socket.terminate();
+    vi.advanceTimersByTime(100);
+    expect(sentTypes(socket)).toEqual(['ping']);
+    expect(socket.closeCalls).toHaveLength(0);
+  });
+});
+
 // ---------------------------------------------------------------------------
 // WsConnectionV1 — global-event registration lifecycle
 // ---------------------------------------------------------------------------