story-kit: merge 138_bug_no_heartbeat_to_detect_stale_websocket_connections

frontend/src/api/client.test.ts

@@ -163,7 +163,8 @@ interface MockWsInstance {
 	onmessage: ((e: { data: string }) => void) | null;
 	onerror: (() => void) | null;
 	readyState: number;
-	send: () => void;
+	sentMessages: string[];
+	send: (data: string) => void;
 	close: () => void;
 	simulateClose: () => void;
 	simulateMessage: (data: Record<string, unknown>) => void;
@@ -183,12 +184,15 @@ function makeMockWebSocket() {
 		onmessage: ((e: { data: string }) => void) | null = null;
 		onerror: (() => void) | null = null;
 		readyState = 0;
+		sentMessages: string[] = [];
 
 		constructor(_url: string) {
 			instances.push(this as unknown as MockWsInstance);
 		}
 
-		send() {}
+		send(data: string) {
+			this.sentMessages.push(data);
+		}
 
 		close() {
 			this.readyState = 3;
@@ -330,3 +334,99 @@ describe("ChatWebSocket", () => {
 		expect(instances).toHaveLength(4);
 	});
 });
+
+describe("ChatWebSocket heartbeat", () => {
+	beforeEach(() => {
+		vi.useFakeTimers();
+		const { MockWebSocket } = makeMockWebSocket();
+		vi.stubGlobal("WebSocket", MockWebSocket);
+		(ChatWebSocket as unknown as { sharedSocket: null }).sharedSocket = null;
+		(ChatWebSocket as unknown as { refCount: number }).refCount = 0;
+	});
+
+	afterEach(() => {
+		vi.useRealTimers();
+		vi.restoreAllMocks();
+	});
+
+	it("sends ping after heartbeat interval", () => {
+		const { MockWebSocket, instances } = makeMockWebSocket();
+		vi.stubGlobal("WebSocket", MockWebSocket);
+
+		const ws = new ChatWebSocket();
+		ws.connect({});
+		instances[0].readyState = 1; // OPEN
+		instances[0].onopen?.(); // starts heartbeat
+
+		vi.advanceTimersByTime(29_999);
+		expect(instances[0].sentMessages).toHaveLength(0);
+
+		vi.advanceTimersByTime(1);
+		expect(instances[0].sentMessages).toHaveLength(1);
+		expect(JSON.parse(instances[0].sentMessages[0])).toEqual({ type: "ping" });
+
+		ws.close();
+	});
+
+	it("closes stale connection when pong is not received", () => {
+		const { MockWebSocket, instances } = makeMockWebSocket();
+		vi.stubGlobal("WebSocket", MockWebSocket);
+
+		const ws = new ChatWebSocket();
+		ws.connect({});
+		instances[0].readyState = 1; // OPEN
+		instances[0].onopen?.(); // starts heartbeat
+
+		// Fire heartbeat — sends ping and starts pong timeout
+		vi.advanceTimersByTime(30_000);
+
+		// No pong received; advance past pong timeout → socket closed → reconnect scheduled
+		vi.advanceTimersByTime(5_000);
+
+		// Advance past reconnect delay
+		vi.advanceTimersByTime(1_001);
+
+		expect(instances).toHaveLength(2);
+		ws.close();
+	});
+
+	it("does not close when pong is received before timeout", () => {
+		const { MockWebSocket, instances } = makeMockWebSocket();
+		vi.stubGlobal("WebSocket", MockWebSocket);
+
+		const ws = new ChatWebSocket();
+		ws.connect({});
+		instances[0].readyState = 1; // OPEN
+		instances[0].onopen?.(); // starts heartbeat
+
+		// Fire heartbeat
+		vi.advanceTimersByTime(30_000);
+
+		// Server responds with pong — clears the pong timeout
+		instances[0].simulateMessage({ type: "pong" });
+
+		// Advance past where pong timeout would have fired
+		vi.advanceTimersByTime(5_001);
+
+		// No reconnect triggered
+		expect(instances).toHaveLength(1);
+		ws.close();
+	});
+
+	it("stops sending pings after explicit close", () => {
+		const { MockWebSocket, instances } = makeMockWebSocket();
+		vi.stubGlobal("WebSocket", MockWebSocket);
+
+		const ws = new ChatWebSocket();
+		ws.connect({});
+		instances[0].readyState = 1; // OPEN
+		instances[0].onopen?.(); // starts heartbeat
+
+		ws.close();
+
+		// Advance well past multiple heartbeat intervals
+		vi.advanceTimersByTime(90_000);
+
+		expect(instances[0].sentMessages).toHaveLength(0);
+	});
+});

frontend/src/api/client.ts

@@ -11,7 +11,8 @@ export type WsRequest =
 			type: "permission_response";
 			request_id: string;
 			approved: boolean;
-	  };
+	  }
+	| { type: "ping" };
 
 export interface AgentAssignment {
 	agent_name: string;
@@ -60,7 +61,9 @@ export type WsResponse =
 	  }
 	/** `.story_kit/project.toml` was modified; re-fetch the agent roster. */
 	| { type: "agent_config_changed" }
-	| { type: "tool_activity"; tool_name: string };
+	| { type: "tool_activity"; tool_name: string }
+	/** Heartbeat response confirming the connection is alive. */
+	| { type: "pong" };
 
 export interface ProviderConfig {
 	provider: string;
@@ -283,6 +286,30 @@ export class ChatWebSocket {
 	private reconnectTimer?: number;
 	private reconnectDelay = 1000;
 	private shouldReconnect = false;
+	private heartbeatInterval?: number;
+	private heartbeatTimeout?: number;
+	private static readonly HEARTBEAT_INTERVAL = 30_000;
+	private static readonly HEARTBEAT_TIMEOUT = 5_000;
+
+	private _startHeartbeat(): void {
+		this._stopHeartbeat();
+		this.heartbeatInterval = window.setInterval(() => {
+			if (!this.socket || this.socket.readyState !== WebSocket.OPEN) return;
+			const ping: WsRequest = { type: "ping" };
+			this.socket.send(JSON.stringify(ping));
+			this.heartbeatTimeout = window.setTimeout(() => {
+				// No pong received within timeout; close socket to trigger reconnect.
+				this.socket?.close();
+			}, ChatWebSocket.HEARTBEAT_TIMEOUT);
+		}, ChatWebSocket.HEARTBEAT_INTERVAL);
+	}
+
+	private _stopHeartbeat(): void {
+		window.clearInterval(this.heartbeatInterval);
+		window.clearTimeout(this.heartbeatTimeout);
+		this.heartbeatInterval = undefined;
+		this.heartbeatTimeout = undefined;
+	}
 
 	private _buildWsUrl(): string {
 		const protocol = window.location.protocol === "https:" ? "wss" : "ws";
@@ -298,6 +325,7 @@ export class ChatWebSocket {
 		if (!this.socket) return;
 		this.socket.onopen = () => {
 			this.reconnectDelay = 1000;
+			this._startHeartbeat();
 		};
 		this.socket.onmessage = (event) => {
 			try {
@@ -327,6 +355,10 @@ export class ChatWebSocket {
 						data.message,
 					);
 				if (data.type === "agent_config_changed") this.onAgentConfigChanged?.();
+				if (data.type === "pong") {
+					window.clearTimeout(this.heartbeatTimeout);
+					this.heartbeatTimeout = undefined;
+				}
 			} catch (err) {
 				this.onError?.(String(err));
 			}
@@ -420,6 +452,7 @@ export class ChatWebSocket {
 
 	close() {
 		this.shouldReconnect = false;
+		this._stopHeartbeat();
 		window.clearTimeout(this.reconnectTimer);
 		this.reconnectTimer = undefined;