diff --git a/.changeset/ws-heartbeat-keepalive.md b/.changeset/ws-heartbeat-keepalive.md new file mode 100644 index 0000000000..cd18b42ff6 --- /dev/null +++ b/.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. diff --git a/packages/kap-server/src/protocol/ws-control.ts b/packages/kap-server/src/protocol/ws-control.ts index 4fef9de57a..536a1dccfe 100644 --- a/packages/kap-server/src/protocol/ws-control.ts +++ b/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(), diff --git a/packages/kap-server/src/transport/ws/v1/protocol.ts b/packages/kap-server/src/transport/ws/v1/protocol.ts index 2363ee8920..b86c4ddab4 100644 --- a/packages/kap-server/src/transport/ws/v1/protocol.ts +++ b/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

{ type: 'ack'; id: string; diff --git a/packages/kap-server/src/transport/ws/v1/registerWsV1.ts b/packages/kap-server/src/transport/ws/v1/registerWsV1.ts index a6b39c1081..c4cc6580b1 100644 --- a/packages/kap-server/src/transport/ws/v1/registerWsV1.ts +++ b/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)); }); diff --git a/packages/kap-server/src/transport/ws/v1/wsConnectionV1.ts b/packages/kap-server/src/transport/ws/v1/wsConnectionV1.ts index 2545faf5b8..38d8af3b62 100644 --- a/packages/kap-server/src/transport/ws/v1/wsConnectionV1.ts +++ b/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; + /** 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 { 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); diff --git a/packages/kap-server/test/wsConnectionV1.test.ts b/packages/kap-server/test/wsConnectionV1.test.ts index 1dd1707b77..0b0c8a5b98 100644 --- a/packages/kap-server/test/wsConnectionV1.test.ts +++ b/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 // ---------------------------------------------------------------------------