diff --git a/src/main/runtime/rpc/unix-socket-transport.test.ts b/src/main/runtime/rpc/unix-socket-transport.test.ts new file mode 100644 index 00000000000..e17ed9abea2 --- /dev/null +++ b/src/main/runtime/rpc/unix-socket-transport.test.ts @@ -0,0 +1,77 @@ +import { EventEmitter } from 'events' +import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest' +import type { Socket } from 'net' +import { UnixSocketTransport } from './unix-socket-transport' + +class FakeSocket extends EventEmitter { + destroyed = false + writable = true + readonly writes: string[] = [] + + setEncoding(): void {} + setNoDelay(): void {} + setTimeout(): void {} + + write(data: string): boolean { + this.writes.push(data) + return true + } + + destroy(): this { + if (!this.destroyed) { + this.destroyed = true + this.writable = false + this.emit('close') + } + return this + } +} + +type UnixSocketTransportInternals = { + handleConnection(socket: Socket): void +} + +describe('UnixSocketTransport', () => { + beforeEach(() => { + vi.useFakeTimers() + }) + + afterEach(() => { + vi.useRealTimers() + }) + + it('clears request keepalive timers when the socket closes before a reply', () => { + const transport = new UnixSocketTransport({ + endpoint: '/tmp/orca-runtime-rpc-test.sock', + kind: 'unix', + keepaliveIntervalMs: 100 + }) + const socket = new FakeSocket() + let aborted = false + + transport.onMessage((_msg, _reply, context) => { + context?.signal?.addEventListener( + 'abort', + () => { + aborted = true + }, + { once: true } + ) + context?.startKeepalive() + }) + + ;(transport as unknown as UnixSocketTransportInternals).handleConnection( + socket as unknown as Socket + ) + socket.emit('data', '{"id":"pending","method":"wait"}\n') + + vi.advanceTimersByTime(100) + expect(socket.writes).toHaveLength(1) + + socket.destroy() + expect(aborted).toBe(true) + + vi.advanceTimersByTime(500) + expect(socket.writes).toHaveLength(1) + }) +}) diff --git a/src/main/runtime/rpc/unix-socket-transport.ts b/src/main/runtime/rpc/unix-socket-transport.ts index 50c7d9822ca..07866b6ae3d 100644 --- a/src/main/runtime/rpc/unix-socket-transport.ts +++ b/src/main/runtime/rpc/unix-socket-transport.ts @@ -104,7 +104,7 @@ export class UnixSocketTransport implements RpcTransport { // completing one dispatch does not abort any other dispatch still running // on the same socket — future-proofing for a persistent CLI socket that // multiplexes sequential requests. - const inflight = new Set() + const inflight = new Set<() => void>() socket.setEncoding('utf8') socket.setNoDelay(true) @@ -115,8 +115,8 @@ export class UnixSocketTransport implements RpcTransport { socket.destroy() }) socket.on('close', () => { - for (const ctrl of inflight) { - ctrl.abort() + for (const cleanup of inflight) { + cleanup() } inflight.clear() }) @@ -151,29 +151,37 @@ export class UnixSocketTransport implements RpcTransport { // Why: the keepalive timer is opt-in per request via `startKeepalive()`. // Short RPCs never call it and pay no timer overhead; only long-poll // handlers (e.g. orchestration.check --wait) arm it. See §3.1. - private dispatchMessage( - socket: Socket, - rawMessage: string, - inflight: Set - ): void { + private dispatchMessage(socket: Socket, rawMessage: string, inflight: Set<() => void>): void { let replied = false let keepaliveTimer: NodeJS.Timeout | null = null - // Why: per-dispatch AbortController so completing one request does not - // abort any sibling request running on the same connection. The - // connection-level `socket.on('close')` iterates `inflight` to cancel all - // outstanding dispatches at once. + // Why: each dispatch needs its own abort signal and keepalive timer + // cleanup. Socket close runs every cleanup without touching sibling + // dispatches that already replied on the same connection. const abortController = new AbortController() - inflight.add(abortController) + let cleanedUp = false + const cleanupDispatch = (abort: boolean): void => { + if (cleanedUp) { + return + } + cleanedUp = true + if (keepaliveTimer) { + clearInterval(keepaliveTimer) + keepaliveTimer = null + } + if (abort) { + abortController.abort() + } + inflight.delete(abortDispatch) + } + const abortDispatch = (): void => cleanupDispatch(true) + inflight.add(abortDispatch) const reply = (response: string): void => { if (replied) { return } replied = true - inflight.delete(abortController) - if (keepaliveTimer) { - clearInterval(keepaliveTimer) - } + cleanupDispatch(false) if (!socket.destroyed && socket.writable) { socket.write(`${response}\n`) } @@ -185,6 +193,7 @@ export class UnixSocketTransport implements RpcTransport { } keepaliveTimer = setInterval(() => { if (replied || socket.destroyed || !socket.writable) { + cleanupDispatch(socket.destroyed || !socket.writable) return } socket.write('{"_keepalive":true}\n')