From dfb69976fd3fe9a44694046e1d43c4e20924407b Mon Sep 17 00:00:00 2001
From: Relakkes Yang <relakkes@gmail.com>
Date: Fri, 8 May 2026 23:19:22 +0800
Subject: [PATCH] fix: flush late transcript writes

---
 .../__tests__/sessionStorageFlush.test.ts     | 62 ++++++++++++++++++
 src/utils/sessionStorage.ts                   | 63 ++++++++++++++-----
 2 files changed, 108 insertions(+), 17 deletions(-)
 create mode 100644 src/utils/__tests__/sessionStorageFlush.test.ts

diff --git a/src/utils/__tests__/sessionStorageFlush.test.ts b/src/utils/__tests__/sessionStorageFlush.test.ts
new file mode 100644
index 00000000..9f2615a4
--- /dev/null
+++ b/src/utils/__tests__/sessionStorageFlush.test.ts
@@ -0,0 +1,62 @@
+import { afterEach, beforeEach, describe, expect, it } from 'bun:test'
+import * as fs from 'node:fs/promises'
+import * as os from 'node:os'
+import * as path from 'node:path'
+
+import {
+  enqueueSessionEntryAfterPendingForTesting,
+  flushSessionStorage,
+  resetProjectForTesting,
+} from '../sessionStorage.js'
+import type { CustomTitleMessage } from '../../types/logs.js'
+
+const originalConfigDir = process.env.CLAUDE_CONFIG_DIR
+
+async function createTmpDir(): Promise<string> {
+  const dir = path.join(
+    os.tmpdir(),
+    `session-storage-flush-${Date.now()}-${Math.random().toString(36).slice(2)}`,
+  )
+  await fs.mkdir(dir, { recursive: true })
+  return dir
+}
+
+describe('sessionStorage flush', () => {
+  let tmpDir: string
+
+  beforeEach(async () => {
+    tmpDir = await createTmpDir()
+    process.env.CLAUDE_CONFIG_DIR = tmpDir
+    resetProjectForTesting()
+  })
+
+  afterEach(async () => {
+    resetProjectForTesting()
+    if (originalConfigDir === undefined) {
+      delete process.env.CLAUDE_CONFIG_DIR
+    } else {
+      process.env.CLAUDE_CONFIG_DIR = originalConfigDir
+    }
+    await fs.rm(tmpDir, { recursive: true, force: true }).catch(() => {})
+  })
+
+  it('drains writes that are queued by pending operations during flush', async () => {
+    const transcriptPath = path.join(tmpDir, 'late-enqueue.jsonl')
+    const entry: CustomTitleMessage = {
+      type: 'custom-title',
+      customTitle: 'late enqueue',
+      sessionId: '11111111-1111-4111-8111-111111111111',
+    }
+    const writePromise = enqueueSessionEntryAfterPendingForTesting(
+      transcriptPath,
+      entry,
+      10,
+    )
+
+    await flushSessionStorage()
+    await writePromise
+
+    const content = await fs.readFile(transcriptPath, 'utf-8')
+    expect(content).toContain('"customTitle":"late enqueue"')
+  })
+})
diff --git a/src/utils/sessionStorage.ts b/src/utils/sessionStorage.ts
index 6d775d6c..dad2db70 100644
--- a/src/utils/sessionStorage.ts
+++ b/src/utils/sessionStorage.ts
@@ -487,6 +487,24 @@ export function setSessionFileForTesting(path: string): void {
   getProject().sessionFile = path
 }
 
+/** @internal Test hook for flush races where tracked work enqueues late. */
+export async function enqueueSessionEntryAfterPendingForTesting(
+  path: string,
+  entry: Entry,
+  delayMs = 0,
+): Promise<void> {
+  const projectForTesting = getProject() as unknown as {
+    trackWrite<T>(fn: () => Promise<T>): Promise<T>
+    enqueueWrite(filePath: string, entry: Entry): Promise<void>
+  }
+  await projectForTesting.trackWrite(async () => {
+    if (delayMs > 0) {
+      await new Promise(resolve => setTimeout(resolve, delayMs))
+    }
+    void projectForTesting.enqueueWrite(path, entry)
+  })
+}
+
 type InternalEventWriter = (
   eventType: string,
   payload: Record<string, unknown>,
@@ -839,25 +857,36 @@ class Project {
   }
 
   async flush(): Promise<void> {
-    // Cancel pending timer
-    if (this.flushTimer) {
-      clearTimeout(this.flushTimer)
-      this.flushTimer = null
-    }
-    // Wait for any in-flight drain to finish
-    if (this.activeDrain) {
-      await this.activeDrain
-    }
-    // Drain anything remaining in the queues
-    await this.drainWriteQueue()
+    while (true) {
+      // Cancel pending timer so process shutdown does not wait for the next
+      // interval tick before draining already queued transcript writes.
+      if (this.flushTimer) {
+        clearTimeout(this.flushTimer)
+        this.flushTimer = null
+      }
 
-    // Wait for non-queue tracked operations (e.g. removeMessageByUuid)
-    if (this.pendingWriteCount === 0) {
-      return
+      // Wait for any in-flight drain to finish before taking ownership of the
+      // remaining queue.
+      if (this.activeDrain) {
+        await this.activeDrain
+      }
+
+      await this.drainWriteQueue()
+
+      if (this.pendingWriteCount === 0) {
+        // A tracked writer may have enqueued after the drain above and then
+        // completed. Loop once more so flush() only returns after that late
+        // enqueue has also been written.
+        if (!this.flushTimer && !this.activeDrain && this.writeQueues.size === 0) {
+          return
+        }
+        continue
+      }
+
+      await new Promise<void>(resolve => {
+        this.flushResolvers.push(resolve)
+      })
     }
-    return new Promise<void>(resolve => {
-      this.flushResolvers.push(resolve)
-    })
   }
 
   /**