From 48d8b35923f438c95ee93dd5be5325ff46841836 Mon Sep 17 00:00:00 2001
From: Neil <4138956+nwparker@users.noreply.github.com>
Date: Sun, 25 Oct 2026 03:46:55 -0700
Subject: [PATCH] fix(daemon): admit history shutdowns before connection
 preparation (#33504)

---
 ...adapter-history-shutdown-admission.test.ts | 79 +++++++++++++++++++
 src/main/daemon/daemon-pty-session-control.ts | 14 ++--
 2 files changed, 86 insertions(+), 7 deletions(-)
 create mode 100644 src/main/daemon/daemon-pty-adapter-history-shutdown-admission.test.ts

diff --git a/src/main/daemon/daemon-pty-adapter-history-shutdown-admission.test.ts b/src/main/daemon/daemon-pty-adapter-history-shutdown-admission.test.ts
new file mode 100644
index 00000000000..6df79da759a
--- /dev/null
+++ b/src/main/daemon/daemon-pty-adapter-history-shutdown-admission.test.ts
@@ -0,0 +1,79 @@
+import './mock-descendant-sweep'
+import { expect, it, vi } from 'vitest'
+import { readFileSync, rmSync } from 'node:fs'
+import { join } from 'node:path'
+import { DaemonPtyAdapter } from './daemon-pty-adapter'
+import { createMockSubprocess, startDaemonAdapterHarness } from './daemon-pty-adapter-test-harness'
+import { getHistorySessionDirName } from './history-paths'
+import type * as DaemonHealthModule from './daemon-health'
+
+vi.mock('./daemon-health', async (importOriginal) => {
+  const actual = await importOriginal<typeof DaemonHealthModule>()
+  return {
+    ...actual,
+    getMacDaemonSystemResolverHealth: vi.fn(async () => 'unknown')
+  }
+})
+
+it('drains keepHistory shutdowns admitted before connection preparation completes', async () => {
+  const subprocess = createMockSubprocess()
+  const harness = await startDaemonAdapterHarness(() => subprocess)
+  const historyDir = join(harness.dir, 'history')
+  const adapter = new DaemonPtyAdapter({
+    socketPath: harness.socketPath,
+    tokenPath: harness.tokenPath,
+    historyPath: historyDir
+  })
+  try {
+    const { id } = await adapter.spawn({
+      cols: 80,
+      rows: 24,
+      sessionId: 'sleep-admitted-before-preparation'
+    })
+    subprocess._simulateData('admitted sleep output\r\n')
+    // oxlint-disable-next-line typescript/consistent-type-assertions -- SAFETY: protected members match the adapter declarations; only connection preparation is gated.
+    const internals = adapter as unknown as {
+      ensureConnected(deadlineMs?: number): Promise<void>
+      keepHistoryShutdowns: Set<Promise<void>>
+    }
+    const originalEnsureConnected = internals.ensureConnected.bind(adapter)
+    let releasePreparation!: () => void
+    const preparationRelease = new Promise<void>((resolve) => {
+      releasePreparation = resolve
+    })
+    vi.spyOn(internals, 'ensureConnected').mockImplementationOnce(async (deadlineMs) => {
+      await preparationRelease
+      await originalEnsureConnected(deadlineMs)
+    })
+
+    const shuttingDown = adapter.shutdown(id, { immediate: true, keepHistory: true })
+    const admittedBeforeYield = internals.keepHistoryShutdowns.size
+    const disconnecting = adapter.disconnectOnly()
+    const operations = Promise.allSettled([shuttingDown, disconnecting])
+    try {
+      expect(admittedBeforeYield).toBe(1)
+      await expect(adapter.shutdown(id, { immediate: true, keepHistory: true })).rejects.toThrow(
+        'Cannot keep history after daemon disconnect has started'
+      )
+    } finally {
+      releasePreparation()
+      await operations
+    }
+
+    expect(await operations).toEqual([
+      { status: 'fulfilled', value: undefined },
+      { status: 'fulfilled', value: undefined }
+    ])
+    expect(internals.keepHistoryShutdowns.size).toBe(0)
+    expect(adapter.getActiveSessionIds()).not.toContain(id)
+    const sessionDir = join(historyDir, getHistorySessionDirName(id))
+    const checkpoint = JSON.parse(readFileSync(join(sessionDir, 'checkpoint.json'), 'utf-8'))
+    expect(checkpoint.snapshotAnsi).toContain('admitted sleep output')
+    expect(JSON.parse(readFileSync(join(sessionDir, 'meta.json'), 'utf-8')).endedAt).toBeNull()
+  } finally {
+    adapter.dispose()
+    harness.adapter.dispose()
+    await harness.server.shutdown()
+    rmSync(harness.dir, { recursive: true, force: true })
+  }
+})
diff --git a/src/main/daemon/daemon-pty-session-control.ts b/src/main/daemon/daemon-pty-session-control.ts
index 13fb079a709..60db90030d5 100644
--- a/src/main/daemon/daemon-pty-session-control.ts
+++ b/src/main/daemon/daemon-pty-session-control.ts
@@ -98,16 +98,16 @@ export abstract class DaemonPtySessionControl extends DaemonPtySessionInput {
     if (opts.keepHistory && this.disconnectOnlyPromise) {
       throw new Error('Cannot keep history after daemon disconnect has started')
     }
-    if (opts.keepHistory) {
+    if (!opts.keepHistory) {
+      await this.withHistorySpawnLock(id, () => this.shutdownWithHistoryLock(id, opts))
+      return
+    }
+    const shutdown = Promise.resolve().then(async () => {
       // Outside the spawn lock, which adopting the session's history takes itself.
       await this.ensureConnected(opts.deadlineMs)
       await this.adoptListedSessionForSleep(id, opts.deadlineMs).catch(() => {})
-    }
-    const shutdown = this.withHistorySpawnLock(id, () => this.shutdownWithHistoryLock(id, opts))
-    if (!opts.keepHistory) {
-      await shutdown
-      return
-    }
+      await this.withHistorySpawnLock(id, () => this.shutdownWithHistoryLock(id, opts))
+    })
     this.keepHistoryShutdowns.add(shutdown)
     try {
       await shutdown