From 058d93b8dd83231267cffb9ae457562b2bcb3aad Mon Sep 17 00:00:00 2001
From: ethan <ethan@coder.com>
Date: Mon, 1 Dec 2025 11:46:53 +1100
Subject: [PATCH] =?UTF-8?q?=F0=9F=A4=96=20fix:=20prevent=20queued=20messag?=
 =?UTF-8?q?e=20amnesia=20by=20flushing=20partial=20before=20tool-call-end?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

When a queued message is sent after tool-call-end, the model needs to see
the complete assistant response including tool results. Previously, there
was a race condition where tool-call-end was emitted before the partial
file was written, causing commitToHistory to read stale data.

The fix ensures partial.json is flushed to disk BEFORE emitting the
tool-call-end event, so listeners (like sendQueuedMessages) see the
complete tool result when they read via commitToHistory.

Changes:
- Make completeToolCall async
- Await flushPartialWrite before emitting tool-call-end
- Update callers to await the async method

_Generated with `mux`_
---
 src/node/services/streamManager.ts | 27 ++++++++++++++++-----------
 1 file changed, 16 insertions(+), 11 deletions(-)

diff --git a/src/node/services/streamManager.ts b/src/node/services/streamManager.ts
index 3a10893cc..c09d2ec71 100644
--- a/src/node/services/streamManager.ts
+++ b/src/node/services/streamManager.ts
@@ -566,9 +566,11 @@ export class StreamManager extends EventEmitter {
   }
 
   /**
-   * Complete a tool call by updating its part and emitting tool-call-end event
+   * Complete a tool call by updating its part and emitting tool-call-end event.
+   * CRITICAL: Flushes partial to disk BEFORE emitting event to prevent race conditions
+   * where listeners (e.g., sendQueuedMessages) read stale partial data.
    */
-  private completeToolCall(
+  private async completeToolCall(
     workspaceId: WorkspaceId,
     streamInfo: WorkspaceStreamInfo,
     toolCalls: Map<
@@ -578,7 +580,7 @@ export class StreamManager extends EventEmitter {
     toolCallId: string,
     toolName: string,
     output: unknown
-  ): void {
+  ): Promise<void> {
     // Find and update the existing tool part
     const existingPartIndex = streamInfo.parts.findIndex(
       (p) => p.type === "dynamic-tool" && p.toolCallId === toolCallId
@@ -609,7 +611,13 @@ export class StreamManager extends EventEmitter {
       }
     }
 
-    // Emit tool-call-end event
+    // CRITICAL: Flush partial to disk BEFORE emitting event
+    // This ensures listeners (like sendQueuedMessages) see the tool result when they
+    // read partial.json via commitToHistory. Without this await, there's a race condition
+    // where the partial is read before the tool result is written, causing "amnesia".
+    await this.flushPartialWrite(workspaceId, streamInfo);
+
+    // Emit tool-call-end event (listeners can now safely read partial)
     this.emit("tool-call-end", {
       type: "tool-call-end",
       workspaceId: workspaceId as string,
@@ -618,9 +626,6 @@ export class StreamManager extends EventEmitter {
       toolName,
       result: output,
     } as ToolCallEndEvent);
-
-    // Schedule partial write
-    void this.schedulePartialWrite(workspaceId, streamInfo);
   }
 
   /**
@@ -762,8 +767,8 @@ export class StreamManager extends EventEmitter {
               const strippedOutput = stripEncryptedContent(part.output);
               toolCall.output = strippedOutput;
 
-              // Use shared completion logic
-              this.completeToolCall(
+              // Use shared completion logic (await to ensure partial is flushed before event)
+              await this.completeToolCall(
                 workspaceId,
                 streamInfo,
                 toolCalls,
@@ -799,8 +804,8 @@ export class StreamManager extends EventEmitter {
                     : JSON.stringify(toolErrorPart.error),
             };
 
-            // Use shared completion logic
-            this.completeToolCall(
+            // Use shared completion logic (await to ensure partial is flushed before event)
+            await this.completeToolCall(
               workspaceId,
               streamInfo,
               toolCalls,