From 79e4e64257ffef65bf36f3e5244cadccccabc231 Mon Sep 17 00:00:00 2001 From: Neil <4138956+nwparker@users.noreply.github.com> Date: Sat, 30 May 2026 12:11:39 -0700 Subject: [PATCH] Clean up remote runtime websocket listeners (#3846) --- ...e-runtime-request-connection-stale.test.ts | 27 +++++++++++- .../remote-runtime-request-connection.ts | 5 +++ .../remote-runtime-request-websocket.test.ts | 42 +++++++++++++++++++ .../remote-runtime-request-websocket.ts | 40 ++++++++++++++---- 4 files changed, 105 insertions(+), 9 deletions(-) create mode 100644 src/shared/remote-runtime-request-websocket.test.ts diff --git a/src/shared/remote-runtime-request-connection-stale.test.ts b/src/shared/remote-runtime-request-connection-stale.test.ts index 39afac27b..765840b8e 100644 --- a/src/shared/remote-runtime-request-connection-stale.test.ts +++ b/src/shared/remote-runtime-request-connection-stale.test.ts @@ -13,13 +13,17 @@ vi.mock('./remote-runtime-request-websocket', () => ({ ) => { const socket = createFakeOpenedSocket(callbacks) opens.push(socket) - return { ok: true, socket: { ws: socket.ws, sharedKey: socket.sharedKey } } + return { + ok: true, + socket: { ws: socket.ws, sharedKey: socket.sharedKey, cleanup: socket.cleanup } + } } })) type FakeOpenedSocket = { ws: WebSocket sharedKey: Uint8Array + cleanup: ReturnType sent: string[] callbacks: RemoteRuntimeWebSocketCallbacks } @@ -36,6 +40,7 @@ function createFakeOpenedSocket(callbacks: RemoteRuntimeWebSocketCallbacks): Fak return { ws, sharedKey: new Uint8Array(32).fill(opens.length + 1), + cleanup: vi.fn(), sent, callbacks } @@ -62,6 +67,26 @@ describe('RemoteRuntimeRequestConnection stale socket callbacks', () => { opens.splice(0) }) + it('runs socket cleanup when the cached connection closes', async () => { + const { RemoteRuntimeRequestConnection } = + await import('./remote-runtime-request-connection.js') + const connection = new RemoteRuntimeRequestConnection({ + v: 2, + endpoint: 'ws://127.0.0.1:6768', + deviceToken: 'device-token', + publicKeyB64: Buffer.from(new Uint8Array(32).fill(9)).toString('base64') + }) + + const request = connection.request('status.get', undefined, 1000) + const socket = opens[0]! + connection.close() + connection.close() + + await expect(request).rejects.toThrow('Remote Orca runtime closed the connection.') + expect(socket.cleanup).toHaveBeenCalledTimes(1) + expect(socket.ws.close).toHaveBeenCalledTimes(1) + }) + it('ignores stale socket errors and text frames after a replacement socket opens', async () => { vi.useFakeTimers() try { diff --git a/src/shared/remote-runtime-request-connection.ts b/src/shared/remote-runtime-request-connection.ts index 17f6605ab..ac0a59329 100644 --- a/src/shared/remote-runtime-request-connection.ts +++ b/src/shared/remote-runtime-request-connection.ts @@ -34,6 +34,7 @@ export class RemoteRuntimeRequestConnection { private state: ConnectionState = 'closed' private ws: WebSocket | null = null private sharedKey: Uint8Array | null = null + private socketCleanup: (() => void) | null = null private readonly pendingRequests = new Map>() private readonly readyWaiters: ReadyWaiter[] = [] private idleCloseTimer: ReturnType | null = null @@ -75,8 +76,10 @@ export class RemoteRuntimeRequestConnection { close(error?: Error): void { const ws = this.ws + const cleanup = this.socketCleanup this.ws = null this.sharedKey = null + this.socketCleanup = null this.state = 'closed' this.clearIdleCloseTimer() @@ -89,6 +92,7 @@ export class RemoteRuntimeRequestConnection { } try { + cleanup?.() ws?.close() } catch { // Best-effort shutdown for a cached remote control connection. @@ -136,6 +140,7 @@ export class RemoteRuntimeRequestConnection { } this.ws = opened.socket.ws this.sharedKey = opened.socket.sharedKey + this.socketCleanup = opened.socket.cleanup this.state = 'awaiting_ready' } diff --git a/src/shared/remote-runtime-request-websocket.test.ts b/src/shared/remote-runtime-request-websocket.test.ts new file mode 100644 index 000000000..0458febd6 --- /dev/null +++ b/src/shared/remote-runtime-request-websocket.test.ts @@ -0,0 +1,42 @@ +import type { EventEmitter } from 'node:events' +import { describe, expect, it, vi } from 'vitest' +import { generateKeyPair, publicKeyToBase64 } from './e2ee-crypto' +import { openRemoteRuntimeWebSocket } from './remote-runtime-request-websocket' + +describe('openRemoteRuntimeWebSocket', () => { + it('detaches Orca callback listeners when cleaned up', () => { + const keyPair = generateKeyPair() + const opened = openRemoteRuntimeWebSocket( + { + v: 2, + endpoint: 'ws://127.0.0.1:1', + deviceToken: 'device-token', + publicKeyB64: publicKeyToBase64(keyPair.publicKey) + }, + { + onClose: vi.fn(), + onError: vi.fn(), + onTextFrame: vi.fn() + } + ) + if (!opened.ok) { + throw opened.error + } + + const socketEvents = opened.socket.ws as unknown as EventEmitter + expect(socketEvents.listenerCount('open')).toBe(1) + expect(socketEvents.listenerCount('close')).toBe(1) + expect(socketEvents.listenerCount('message')).toBe(1) + expect(socketEvents.listenerCount('error')).toBe(1) + + opened.socket.cleanup() + opened.socket.cleanup() + + expect(socketEvents.listenerCount('open')).toBe(0) + expect(socketEvents.listenerCount('close')).toBe(0) + expect(socketEvents.listenerCount('message')).toBe(0) + expect(socketEvents.listenerCount('error')).toBe(1) + expect(() => socketEvents.emit('error', new Error('late socket error'))).not.toThrow() + opened.socket.ws.terminate() + }) +}) diff --git a/src/shared/remote-runtime-request-websocket.ts b/src/shared/remote-runtime-request-websocket.ts index 569ef9e27..2972bc7e3 100644 --- a/src/shared/remote-runtime-request-websocket.ts +++ b/src/shared/remote-runtime-request-websocket.ts @@ -15,6 +15,7 @@ import { export type RemoteRuntimeWebSocket = { ws: WebSocket sharedKey: Uint8Array + cleanup: () => void } export type RemoteRuntimeWebSocketCallbacks = { @@ -35,22 +36,23 @@ export function openRemoteRuntimeWebSocket( const serverPublicKey = publicKeyFromBase64(pairing.publicKeyB64) const sharedKey = deriveSharedKey(keyPair.secretKey, serverPublicKey) - ws.once('open', () => { + let cleanedUp = false + const onOpen = (): void => { ws.send( JSON.stringify({ type: 'e2ee_hello', publicKeyB64: publicKeyToBase64(keyPair.publicKey) }) ) - }) - ws.on('error', () => { + } + const onError = (): void => { callbacks.onError( ws, remoteRuntimeUnavailableError('Could not connect to the remote Orca runtime.') ) - }) - ws.on('close', () => callbacks.onClose(ws)) - ws.on('message', (data, isBinary) => { + } + const onClose = (): void => callbacks.onClose(ws) + const onMessage = (data: WebSocket.RawData, isBinary: boolean): void => { if (isBinary) { callbacks.onError( ws, @@ -61,10 +63,32 @@ export function openRemoteRuntimeWebSocket( return } callbacks.onTextFrame(ws, data.toString()) - }) - return { ok: true, socket: { ws, sharedKey } } + } + const cleanup = (): void => { + if (cleanedUp) { + return + } + cleanedUp = true + ws.off('open', onOpen) + ws.off('error', onError) + ws.off('close', onClose) + ws.off('message', onMessage) + // Why: a manually closed ws can still emit a late transport error; keep + // that from becoming an unhandled EventEmitter error after detaching Orca. + if (ws.readyState !== WebSocket.CLOSED) { + ws.on('error', ignoreLateSocketError) + } + } + + ws.once('open', onOpen) + ws.on('error', onError) + ws.on('close', onClose) + ws.on('message', onMessage) + return { ok: true, socket: { ws, sharedKey, cleanup } } } +function ignoreLateSocketError(): void {} + function createSocket( pairing: PairingOffer ):