Clean up remote runtime subscription listeners

This commit is contained in:
Neil 2026-05-30 12:40:22 -07:00 committed by GitHub
parent ccb21eb405
commit 2da9c85647
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 76 additions and 9 deletions

View File

@ -57,6 +57,39 @@ describe('subscribeRemoteRuntimeRequest', () => {
expect(onError).not.toHaveBeenCalled()
subscription.close()
})
it('detaches subscription socket listeners after close', async () => {
const offSpy = vi.spyOn(WebSocketClient.prototype, 'off')
try {
const server = await createSubscriptionServer()
const onResponse = vi.fn()
const onError = vi.fn()
const onClose = vi.fn()
const subscription = await subscribeRemoteRuntimeRequest(
server.pairing,
'terminal.subscribe',
{ terminal: 't1' },
1000,
{
onResponse,
onError,
onClose
}
)
await vi.waitFor(() => expect(onResponse).toHaveBeenCalled())
subscription.close()
await vi.waitFor(() => expect(onClose).toHaveBeenCalledOnce())
const removedEvents = offSpy.mock.calls.map(([event]) => event)
expect(removedEvents).toEqual(expect.arrayContaining(['open', 'error', 'close', 'message']))
expect(subscription.sendBinary(new Uint8Array([9]))).toBe(false)
expect(onError).not.toHaveBeenCalled()
} finally {
offSpy.mockRestore()
}
})
})
describe('sendRemoteRuntimeRequest', () => {

View File

@ -341,6 +341,33 @@ export async function subscribeRemoteRuntimeRequest<TResult>(
let settled = false
let ws: WebSocket | null = null
const cleanupSocketListeners = (): WebSocket | null => {
const socket = ws
if (!socket) {
return null
}
socket.off('open', onOpen)
socket.off('error', onError)
socket.off('close', onClose)
socket.off('message', onMessage)
ws = null
// Why: startup failures detach Orca callbacks before closing the ws,
// but ws can still emit a late transport error while close is in flight.
if (socket.readyState !== WebSocket.CLOSED) {
socket.on('error', ignoreSettledRemoteRuntimeSocketError)
}
return socket
}
const closeSocketAfterCleanup = (): void => {
const socket = cleanupSocketListeners()
try {
socket?.close()
} catch {
// ignore best-effort close
}
}
const timeout = setTimeout(() => {
fail(
new RemoteRuntimeClientError(
@ -379,7 +406,7 @@ export async function subscribeRemoteRuntimeRequest<TResult>(
if (!settled) {
settled = true
clearTimeout(timeout)
close()
closeSocketAfterCleanup()
reject(error)
return
}
@ -394,27 +421,29 @@ export async function subscribeRemoteRuntimeRequest<TResult>(
return
}
ws.once('open', () => {
function onOpen(): void {
ws?.send(
JSON.stringify({
type: 'e2ee_hello',
publicKeyB64: publicKeyToBase64(keyPair.publicKey)
})
)
})
}
ws.once('error', () => {
function onError(): void {
fail(
new RemoteRuntimeClientError(
'remote_runtime_unavailable',
'Could not connect to the remote Orca runtime.'
)
)
})
}
ws.on('close', () => {
function onClose(): void {
clearTimeout(timeout)
cleanupSocketListeners()
if (!settled) {
settled = true
reject(
new RemoteRuntimeClientError(
'remote_runtime_unavailable',
@ -424,9 +453,9 @@ export async function subscribeRemoteRuntimeRequest<TResult>(
return
}
callbacks.onClose?.()
})
}
ws.on('message', (data, isBinary) => {
function onMessage(data: WebSocket.RawData, isBinary: boolean): void {
if (isBinary) {
handleBinaryFrame(new Uint8Array(data as Buffer))
return
@ -455,7 +484,12 @@ export async function subscribeRemoteRuntimeRequest<TResult>(
}
handleRpcFrame(plaintext)
})
}
ws.once('open', onOpen)
ws.once('error', onError)
ws.on('close', onClose)
ws.on('message', onMessage)
function handleReadyFrame(frame: string): void {
let ready: unknown