fix: clean up daemon stream sockets (#4214)

src/main/daemon/daemon-server.test.ts

@@ -91,6 +91,50 @@ describe('DaemonServer', () => {
     return client
   }
 
+  async function connectRawHello(role: 'control' | 'stream', clientId: string): Promise<Socket> {
+    const socket = connect(socketPath)
+    await new Promise<void>((resolve) => socket.once('connect', resolve))
+    socket.write(
+      encodeNdjson({
+        type: 'hello',
+        version: PROTOCOL_VERSION,
+        token: readFileSync(tokenPath, 'utf-8').trim(),
+        clientId,
+        role
+      })
+    )
+    await new Promise<void>((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<void> {
+    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()

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,