From 8837843df17ec78539418f70ab42b73e87c87634 Mon Sep 17 00:00:00 2001 From: William Valentin Date: Mon, 16 Feb 2026 19:42:23 -0800 Subject: [PATCH] fix(companion): reject event waiters on unexpected socket close --- README.md | 2 +- docs/plans/state.json | 13 ++++++++++ src/companion/runtimeClient.test.ts | 40 +++++++++++++++++++++++++++++ src/companion/runtimeClient.ts | 8 +++++- 4 files changed, 61 insertions(+), 2 deletions(-) diff --git a/README.md b/README.md index 4404e24..c0dd68e 100644 --- a/README.md +++ b/README.md @@ -1190,7 +1190,7 @@ Methods: - `system.capabilities` returns gateway protocol and node policy snapshot. Companion runtime helper: -- `src/companion/runtimeClient.ts` provides a typed Node/WebSocket client for companion runtimes (macOS/iOS/Android workers) with wrappers for `node.register`, `node.capabilities.get`, `node.location.set/get`, `node.status.set`, `node.push_token.set`, `system.capabilities`, `system.nodes`, and canvas artifact RPCs (`canvas.put/get/list/delete/clear`), plus convenience helpers (`bootstrapNode`, optional `autoConnect`, `dispose()`) and event helpers (`subscribeEvents()`, `subscribeEvent()`, `subscribeAgentStream()`, `subscribeAgentTyping()`, `subscribeContextWarning()`, `waitForEvent()` with timeout/predicate/abort support and deterministic teardown cancellation, `waitForAnyEvent()`, `waitForAgentStream()`, `waitForAgentTyping()`, `waitForContextWarning()`, `clearEventSubscriptions()`, `listKnownEventNames()`, `eventSubscriptionCount`). +- `src/companion/runtimeClient.ts` provides a typed Node/WebSocket client for companion runtimes (macOS/iOS/Android workers) with wrappers for `node.register`, `node.capabilities.get`, `node.location.set/get`, `node.status.set`, `node.push_token.set`, `system.capabilities`, `system.nodes`, and canvas artifact RPCs (`canvas.put/get/list/delete/clear`), plus convenience helpers (`bootstrapNode`, optional `autoConnect`, `dispose()`) and event helpers (`subscribeEvents()`, `subscribeEvent()`, `subscribeAgentStream()`, `subscribeAgentTyping()`, `subscribeContextWarning()`, `waitForEvent()` with timeout/predicate/abort support and deterministic teardown cancellation (including socket-close rejection), `waitForAnyEvent()`, `waitForAgentStream()`, `waitForAgentTyping()`, `waitForContextWarning()`, `clearEventSubscriptions()`, `listKnownEventNames()`, `eventSubscriptionCount`). - `src/companion/platformClients.ts` provides platform-focused wrappers: - `MacOSCompanionClient` (`platform: "macos"`, APNs push registration) - `IOSCompanionClient` (`platform: "ios"`, APNs push registration) diff --git a/docs/plans/state.json b/docs/plans/state.json index 37d28cf..f9470ff 100644 --- a/docs/plans/state.json +++ b/docs/plans/state.json @@ -772,6 +772,19 @@ ], "test_status": "pnpm test:run src/companion/heartbeatLoop.test.ts src/companion/platformClients.test.ts src/companion/runtimeClient.test.ts src/companion/platformClients.integration.test.ts + pnpm typecheck passing" }, + "companion-runtime-socket-close-waiter-rejection": { + "status": "completed", + "date": "2026-02-17", + "updated": "2026-02-17", + "summary": "Hardened runtime close-path behavior so unexpected WebSocket closure clears stale socket references and rejects pending event waiters with a deterministic `WebSocket closed` error.", + "files_modified": [ + "src/companion/runtimeClient.ts", + "src/companion/runtimeClient.test.ts", + "README.md", + "docs/plans/state.json" + ], + "test_status": "pnpm test:run src/companion/runtimeClient.test.ts src/companion/platformClients.test.ts src/companion/heartbeatLoop.test.ts src/companion/platformClients.integration.test.ts + pnpm typecheck passing" + }, "browser-tools-activation-clarity": { "status": "completed", "date": "2026-02-17", diff --git a/src/companion/runtimeClient.test.ts b/src/companion/runtimeClient.test.ts index bbac87f..7e8bb9a 100644 --- a/src/companion/runtimeClient.test.ts +++ b/src/companion/runtimeClient.test.ts @@ -1,6 +1,8 @@ import { describe, it, expect, beforeAll, afterAll, vi } from 'vitest'; import { resolve } from 'path'; import { createServer } from 'net'; +import { EventEmitter } from 'events'; +import { WebSocket } from 'ws'; import type { GatewayServerConfig } from '../gateway/server.js'; import { GatewayServer } from '../gateway/server.js'; import { CompanionRuntimeClient, GatewayRpcError } from './runtimeClient.js'; @@ -378,6 +380,44 @@ describe('CompanionRuntimeClient', () => { await awaited; }); + it('waitForEvent rejects when websocket closes unexpectedly', async () => { + class FakeWebSocket extends EventEmitter { + readyState: number = WebSocket.CONNECTING; + + constructor() { + super(); + queueMicrotask(() => { + this.readyState = WebSocket.OPEN; + this.emit('open'); + }); + } + + send(_payload: string, callback?: (error?: Error) => void): void { + callback?.(); + } + + close(_code?: number, _reason?: string): void { + this.readyState = WebSocket.CLOSED; + this.emit('close'); + } + } + + const client = new CompanionRuntimeClient({ + url: 'ws://127.0.0.1:1', + websocketFactory: () => new FakeWebSocket() as unknown as WebSocket, + }); + await client.connect(); + + const awaited = expect( + client.waitForEvent('agent.stream', { timeoutMs: 10_000 }), + ).rejects.toThrow('WebSocket closed'); + + const ws = (client as unknown as { ws: WebSocket | null }).ws; + ws?.close(); + await awaited; + expect(client.connected).toBe(false); + }); + it('waitForAgentStream resolves on agent.stream events', async () => { const client = new CompanionRuntimeClient({ url: 'ws://127.0.0.1:1', diff --git a/src/companion/runtimeClient.ts b/src/companion/runtimeClient.ts index 349d53e..7968767 100644 --- a/src/companion/runtimeClient.ts +++ b/src/companion/runtimeClient.ts @@ -330,7 +330,13 @@ export class CompanionRuntimeClient { settled = true; this.ws = ws; this.ws.on('message', (raw) => this.handleMessage(raw.toString())); - this.ws.on('close', () => this.rejectAllPending(new Error('WebSocket closed'))); + this.ws.on('close', () => { + if (this.ws === ws) { + this.ws = null; + } + this.rejectAllPending(new Error('WebSocket closed')); + this.rejectEventWaits(new Error('WebSocket closed')); + }); this.ws.on('error', () => { // close event handles pending rejection });