fix: resolve race conditions in stream error handling and test flakiness

- Make debugTriggerStreamError async and await processingPromise to ensure
  error handling completes before returning (eliminates duplicate logic)
- Wait for in-flight partial writes before writing error state to prevent
  inconsistent partial.json
- Pass toolPolicy on resume to match original message (tools disabled)
- Add delay after resume to let history update complete before reading

Author: ThomasK33

src/node/services/aiService.ts

@@ -1270,7 +1270,7 @@ export class AIService extends EventEmitter {
   debugTriggerStreamError(
     workspaceId: string,
     errorMessage = "Test-triggered stream error"
-  ): boolean {
+  ): Promise<boolean> {
     return this.streamManager.debugTriggerStreamError(workspaceId, errorMessage);
   }
 

src/node/services/streamManager.ts

@@ -1177,8 +1177,14 @@ export class StreamManager extends EventEmitter {
         },
         parts: streamInfo.parts,
       };
-      // Write error state to disk (fire-and-forget to not block error emission)
-      void this.partialService.writePartial(workspaceId as string, errorPartialMessage);
+      // Wait for any in-flight partial write to complete before writing error state
+      // This prevents race conditions where the error write and a throttled flush
+      // write at the same time, causing inconsistent partial.json state
+      if (streamInfo.partialWritePromise) {
+        await streamInfo.partialWritePromise;
+      }
+      // Write error state to disk - await to ensure consistent state before any resume
+      await this.partialService.writePartial(workspaceId as string, errorPartialMessage);
 
       // Emit error event
       this.emit("error", {
@@ -1683,11 +1689,12 @@ export class StreamManager extends EventEmitter {
   /**
    * DEBUG ONLY: Trigger an artificial stream error for testing
    * This method allows integration tests to simulate stream errors without
-   * mocking the AI SDK or network layer. It triggers the same error handling
-   * path as genuine stream errors by aborting the stream and manually triggering
-   * the error event (since abort alone doesn't throw, it just sets a flag).
+   * mocking the AI SDK or network layer. It aborts the stream with an error,
+   * which causes the for-await loop to throw. The catch block in
+   * processStreamWithCleanup handles writing partial.json and emitting the
+   * error event - we just need to abort and wait.
    */
-  debugTriggerStreamError(workspaceId: string, errorMessage: string): boolean {
+  async debugTriggerStreamError(workspaceId: string, errorMessage: string): Promise<boolean> {
     const typedWorkspaceId = workspaceId as WorkspaceId;
     const streamInfo = this.workspaceStreams.get(typedWorkspaceId);
 
@@ -1699,41 +1706,15 @@ export class StreamManager extends EventEmitter {
       return false;
     }
 
-    // Abort the stream first
+    // Abort the stream with an error - this causes the for-await loop to throw
+    // when it's blocked waiting for the next iteration. The catch block in
+    // processStreamWithCleanup will handle:
+    // 1. Writing error state to partial.json (with all accumulated parts)
+    // 2. Emitting the error event
     streamInfo.abortController.abort(new Error(errorMessage));
 
-    // Update streamInfo metadata with error (so subsequent flushes preserve it)
-    streamInfo.initialMetadata = {
-      ...streamInfo.initialMetadata,
-      error: errorMessage,
-      errorType: "network",
-    };
-
-    // Write error state to partial.json (same as real error handling)
-    const errorPartialMessage: MuxMessage = {
-      id: streamInfo.messageId,
-      role: "assistant",
-      metadata: {
-        historySequence: streamInfo.historySequence,
-        timestamp: streamInfo.startTime,
-        model: streamInfo.model,
-        partial: true,
-        error: errorMessage,
-        errorType: "network", // Test errors are network-like
-        ...streamInfo.initialMetadata,
-      },
-      parts: streamInfo.parts,
-    };
-    void this.partialService.writePartial(workspaceId, errorPartialMessage);
-
-    // Emit error event (same as real error handling)
-    this.emit("error", {
-      type: "error",
-      workspaceId,
-      messageId: streamInfo.messageId,
-      error: errorMessage,
-      errorType: "network",
-    } as ErrorEvent);
+    // Wait for the stream processing to complete (includes error handling)
+    await streamInfo.processingPromise;
 
     return true;
   }

src/node/services/workspaceService.ts

@@ -108,7 +108,7 @@ export class WorkspaceService extends EventEmitter {
    * This is used by integration tests to simulate network errors mid-stream.
    * @returns true if an active stream was found and error was triggered
    */
-  debugTriggerStreamError(workspaceId: string, errorMessage?: string): boolean {
+  debugTriggerStreamError(workspaceId: string, errorMessage?: string): Promise<boolean> {
     return this.aiService.debugTriggerStreamError(workspaceId, errorMessage);
   }
 

tests/integration/streamErrorRecovery.test.ts

@@ -87,27 +87,28 @@ function truncateToLastCompleteMarker(text: string, nonce: string): string {
   return text.substring(0, endIndex);
 }
 
-import type { OrpcSource } from "./helpers";
 import type { OrpcTestClient } from "./orpcTestClient";
 
 /**
  * Helper: Resume stream and wait for successful completion
  * Uses StreamCollector for ORPC-native event handling
  */
 async function resumeAndWaitForSuccess(
-  source: OrpcSource,
   workspaceId: string,
   client: OrpcTestClient,
   model: string,
-  timeoutMs = 15000
+  timeoutMs = 15000,
+  options?: {
+    toolPolicy?: Array<{ regex_match: string; action: "enable" | "disable" | "require" }>;
+  }
 ): Promise<void> {
   const collector = createStreamCollector(client, workspaceId);
   collector.start();
 
   try {
     const resumeResult = await client.workspace.resumeStream({
       workspaceId,
-      options: { model },
+      options: { model, toolPolicy: options?.toolPolicy },
     });
 
     if (!resumeResult.success) {
@@ -261,11 +262,19 @@ IMPORTANT: Do not add any other text. Start immediately with ${nonce}-1: one. If
         await new Promise((resolve) => setTimeout(resolve, 500));
 
         // Resume and wait for completion
-        await resumeAndWaitForSuccess(env, workspaceId, client, `${PROVIDER}:${MODEL}`);
+        // Disable all tools (same as original message) so model outputs text, not tool calls
+        await resumeAndWaitForSuccess(workspaceId, client, `${PROVIDER}:${MODEL}`, 15000, {
+          toolPolicy: [{ regex_match: ".*", action: "disable" }],
+        });
+
+        // Small delay to let history update complete after stream-end
+        // stream-end is emitted before updateHistory completes
+        await new Promise((resolve) => setTimeout(resolve, 500));
 
         // Read final assistant message from history
         const history = await readChatHistory(env.tempDir, workspaceId);
         const assistantMessages = history.filter((m) => m.role === "assistant");
+
         const finalText = assistantMessages
           .flatMap((m) => m.parts)
           .filter((p) => p.type === "text")