From 4a98c9ff7a1b3e09ab02cb90d8fb57f0405fe9c7 Mon Sep 17 00:00:00 2001
From: Jinwoo-H <jinwoo0825@gmail.com>
Date: Sat, 30 May 2026 08:33:44 -0700
Subject: [PATCH] test: cover daemon hello close races

Co-authored-by: Orca <help@stably.ai>
---
 src/main/daemon/client.test.ts | 27 +++++++++++++++++++++++++++
 1 file changed, 27 insertions(+)

diff --git a/src/main/daemon/client.test.ts b/src/main/daemon/client.test.ts
index 7ef2a3da11a..a339694cc7b 100644
--- a/src/main/daemon/client.test.ts
+++ b/src/main/daemon/client.test.ts
@@ -48,6 +48,8 @@ describe('DaemonClient', () => {
   })
 
   function startMockDaemon(opts?: {
+    closeOnConnect?: boolean
+    closeOnHello?: boolean
     onControlMessage?: (msg: unknown) => string | null
     onHello?: (msg: HelloMessage) => void
     onStreamHello?: (msg: HelloMessage) => void
@@ -56,6 +58,11 @@ describe('DaemonClient', () => {
   }): Promise<void> {
     return new Promise((resolve) => {
       server = createServer((socket) => {
+        if (opts?.closeOnConnect) {
+          socket.destroy()
+          return
+        }
+
         let buffer = ''
         socket.on('data', (chunk) => {
           buffer += chunk.toString()
@@ -72,6 +79,10 @@ describe('DaemonClient', () => {
             if (msg.type === 'hello') {
               const hello = msg as HelloMessage
               opts?.onHello?.(hello)
+              if (opts?.closeOnHello) {
+                socket.destroy()
+                return
+              }
               if (opts?.suppressHelloResponse) {
                 return
               }
@@ -164,6 +175,22 @@ describe('DaemonClient', () => {
         vi.useRealTimers()
       }
     })
+
+    it('rejects when the daemon closes before hello completes', async () => {
+      await startMockDaemon({ closeOnHello: true })
+
+      client = new DaemonClient({ socketPath, tokenPath })
+      await expect(client.ensureConnected()).rejects.toThrow(
+        'Connection closed before hello response'
+      )
+    })
+
+    it('rejects when the daemon closes immediately after connect', async () => {
+      await startMockDaemon({ closeOnConnect: true })
+
+      client = new DaemonClient({ socketPath, tokenPath })
+      await expect(client.ensureConnected()).rejects.toThrow()
+    })
   })
 
   describe('RPC', () => {