Clean up remote runtime websocket listeners (#3846)

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<typeof vi.fn>
   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 {

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<string, PendingRequest<unknown>>()
   private readonly readyWaiters: ReadyWaiter[] = []
   private idleCloseTimer: ReturnType<typeof setTimeout> | 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'
   }
 

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()
+  })
+})

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
 ):