fix(companion): reject event waiters on unexpected socket close

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',

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
         });