diff --git a/src/main/daemon/daemon-server.test.ts b/src/main/daemon/daemon-server.test.ts index 794d5f44737..feb9fe9259e 100644 --- a/src/main/daemon/daemon-server.test.ts +++ b/src/main/daemon/daemon-server.test.ts @@ -91,6 +91,50 @@ describe('DaemonServer', () => { return client } + async function connectRawHello(role: 'control' | 'stream', clientId: string): Promise { + const socket = connect(socketPath) + await new Promise((resolve) => socket.once('connect', resolve)) + socket.write( + encodeNdjson({ + type: 'hello', + version: PROTOCOL_VERSION, + token: readFileSync(tokenPath, 'utf-8').trim(), + clientId, + role + }) + ) + await new Promise((resolve, reject) => { + const cleanup = (): void => { + socket.off('data', onData) + socket.off('error', onError) + } + const onData = (data: Buffer): void => { + cleanup() + const parsed = JSON.parse(data.toString().trim()) as { ok?: boolean; error?: string } + if (parsed.ok) { + resolve() + return + } + reject(new Error(parsed.error ?? 'hello rejected')) + } + const onError = (error: Error): void => { + cleanup() + reject(error) + } + socket.on('data', onData) + socket.on('error', onError) + }) + return socket + } + + async function waitFor(predicate: () => boolean, timeoutMs = 2_000): Promise { + const startedAt = Date.now() + while (!predicate() && Date.now() - startedAt < timeoutMs) { + await new Promise((resolve) => setTimeout(resolve, 20)) + } + expect(predicate()).toBe(true) + } + describe('startup', () => { it('creates token file and starts listening', async () => { await startServer() @@ -380,6 +424,75 @@ describe('DaemonServer', () => { }) }) + describe('stream socket lifecycle', () => { + it('clears the tracked stream socket when it closes', async () => { + await startServer() + const daemon = server as unknown as DaemonServerPrivate + const control = await connectRawHello('control', 'raw-client') + const stream = await connectRawHello('stream', 'raw-client') + + expect(daemon.clients.get('raw-client')?.streamSocket).toBeTruthy() + + stream.destroy() + + await waitFor(() => daemon.clients.get('raw-client')?.streamSocket === null) + control.destroy() + }) + + it('destroys a replaced stream socket for the same client', async () => { + await startServer() + const daemon = server as unknown as DaemonServerPrivate + const control = await connectRawHello('control', 'raw-client') + const firstStream = await connectRawHello('stream', 'raw-client') + let firstClosed = false + firstStream.once('close', () => { + firstClosed = true + }) + + const secondStream = await connectRawHello('stream', 'raw-client') + + await waitFor(() => firstClosed) + expect(daemon.clients.get('raw-client')?.streamSocket).toBeTruthy() + secondStream.destroy() + control.destroy() + }) + + it('destroys previous sockets when a control client id reconnects', async () => { + await startServer() + const daemon = server as unknown as DaemonServerPrivate + const firstControl = await connectRawHello('control', 'raw-client') + const firstStream = await connectRawHello('stream', 'raw-client') + let firstControlClosed = false + let firstStreamClosed = false + firstControl.once('close', () => { + firstControlClosed = true + }) + firstStream.once('close', () => { + firstStreamClosed = true + }) + + const secondControl = await connectRawHello('control', 'raw-client') + + await waitFor(() => firstControlClosed && firstStreamClosed) + expect(daemon.clients.get('raw-client')?.controlSocket).toBeTruthy() + expect(daemon.clients.get('raw-client')?.streamSocket).toBeNull() + secondControl.destroy() + }) + + it('destroys orphan stream sockets without a control client', async () => { + await startServer() + const daemon = server as unknown as DaemonServerPrivate + const stream = await connectRawHello('stream', 'missing-client') + let closed = false + stream.once('close', () => { + closed = true + }) + + await waitFor(() => closed || stream.destroyed) + expect(daemon.clients.has('missing-client')).toBe(false) + }) + }) + describe('shutdown', () => { it('stops accepting connections after shutdown', async () => { await startServer() diff --git a/src/main/daemon/daemon-server.ts b/src/main/daemon/daemon-server.ts index b663f4dfe33..381400e25da 100644 --- a/src/main/daemon/daemon-server.ts +++ b/src/main/daemon/daemon-server.ts @@ -151,6 +151,7 @@ export class DaemonServer { socket.write(encodeNdjson({ type: 'hello', ok: true })) if (hello.role === 'control') { + const previous = this.clients.get(hello.clientId) const client: ConnectedClient = { clientId: hello.clientId, controlSocket: socket, @@ -158,12 +159,22 @@ export class DaemonServer { } this.clients.set(hello.clientId, client) this.setupControlSocket(socket, hello.clientId) + if (previous) { + // Why: a reconnect can reuse a clientId before the old sockets notice + // their close. Tear them down after installing the new owner so stale + // close events cannot delete the replacement client entry. + previous.streamSocket?.destroy() + previous.controlSocket.destroy() + } } else if (hello.role === 'stream') { const client = this.clients.get(hello.clientId) - if (client) { - client.streamSocket = socket + if (!client) { + // Why: stream sockets are only meaningful beside a control socket; an + // orphan stream would otherwise stay open with no tracked owner. + socket.destroy() + return } - // Stream socket is receive-only from daemon's perspective (for events) + this.setupStreamSocket(socket, client) } } @@ -178,11 +189,41 @@ export class DaemonServer { socket.on('data', (chunk) => parser.feed(chunk.toString())) socket.on('close', () => { + const client = this.clients.get(clientId) + if (client?.controlSocket !== socket) { + return + } this.streamDataBatcher.clear(clientId) + client.streamSocket?.destroy() this.clients.delete(clientId) }) } + private setupStreamSocket(socket: Socket, client: ConnectedClient): void { + const previous = client.streamSocket + socket.removeAllListeners('data') + client.streamSocket = socket + + const cleanup = (): void => { + socket.removeListener('close', cleanup) + socket.removeListener('error', cleanup) + if (this.clients.get(client.clientId) !== client || client.streamSocket !== socket) { + return + } + this.streamDataBatcher.clear(client.clientId) + client.streamSocket = null + } + + socket.on('close', cleanup) + socket.on('error', cleanup) + + if (previous && previous !== socket) { + // Why: replacing a stream socket must not leave the old receive-only + // channel alive and untracked. + previous.destroy() + } + } + private async handleRequest( socket: Socket, clientId: string,