🤖 fix: avoid stream replay blocking caught-up

Author: ammar-agent

src/node/services/streamManager.test.ts

@@ -320,6 +320,90 @@ describe("StreamManager - Concurrent Stream Prevention", () => {
   });
 });
 
+describe("StreamManager - replayStream", () => {
+  let streamManager: StreamManager;
+  let mockHistoryService: HistoryService;
+  let mockPartialService: PartialService;
+
+  beforeEach(() => {
+    mockHistoryService = createMockHistoryService();
+    mockPartialService = createMockPartialService();
+    streamManager = new StreamManager(mockHistoryService, mockPartialService);
+    // Suppress error events from bubbling up as uncaught exceptions during tests
+    streamManager.on("error", () => undefined);
+  });
+
+  test("does not chase newly appended parts during replay", async () => {
+    const workspaceId = "test-workspace-replay";
+
+    // Mock token tracker so we can pause inside replay.
+    let releaseFirstCountTokens: (() => void) | undefined;
+    let onFirstCountTokensCalled: (() => void) | undefined;
+
+    const firstCountTokensCalled = new Promise<void>((resolve) => {
+      onFirstCountTokensCalled = resolve;
+    });
+
+    let countTokensCalls = 0;
+    const tokenTracker = {
+      setModel: mock(() => Promise.resolve()),
+      countTokens: mock(() => {
+        countTokensCalls += 1;
+        if (countTokensCalls === 1) {
+          onFirstCountTokensCalled?.();
+          return new Promise<number>((resolve) => {
+            releaseFirstCountTokens = () => resolve(1);
+          });
+        }
+        return Promise.resolve(1);
+      }),
+    };
+
+    const replaced = Reflect.set(streamManager, "tokenTracker", tokenTracker);
+    if (!replaced) {
+      throw new Error("Failed to mock StreamManager.tokenTracker");
+    }
+
+    const workspaceStreamsValue = Reflect.get(streamManager, "workspaceStreams") as unknown;
+    if (!(workspaceStreamsValue instanceof Map)) {
+      throw new Error("StreamManager.workspaceStreams is not a Map");
+    }
+
+    const workspaceStreams = workspaceStreamsValue as Map<string, unknown>;
+
+    const streamInfo = {
+      state: "streaming",
+      messageId: "msg-1",
+      model: "anthropic:claude-3-5-sonnet",
+      historySequence: 1,
+      startTime: Date.now(),
+      parts: [{ type: "text", text: "hello", timestamp: Date.now() }],
+      initialMetadata: {},
+    };
+
+    workspaceStreams.set(workspaceId, streamInfo);
+
+    const deltas: string[] = [];
+    streamManager.on("stream-delta", (data: { delta: string }) => {
+      deltas.push(data.delta);
+    });
+
+    const replayPromise = streamManager.replayStream(workspaceId);
+
+    // Wait until replay is actively tokenizing the first part.
+    await firstCountTokensCalled;
+
+    // Mutate parts while replay is in progress. replayStream should snapshot, so this part should
+    // not be emitted by replay.
+    streamInfo.parts.push({ type: "text", text: "world", timestamp: Date.now() });
+
+    releaseFirstCountTokens?.();
+    await replayPromise;
+
+    expect(deltas).toEqual(["hello"]);
+  });
+});
+
 describe("StreamManager - Unavailable Tool Handling", () => {
   let streamManager: StreamManager;
   let mockHistoryService: HistoryService;

src/node/services/streamManager.ts

@@ -1754,9 +1754,14 @@ export class StreamManager extends EventEmitter {
       ...(replayMode && { mode: replayMode }),
     });
 
-    // Replay accumulated parts as events using shared emission logic
-    // This guarantees replay produces identical events to the original stream
-    for (const part of streamInfo.parts) {
+    // Replay accumulated parts as events using shared emission logic.
+    //
+    // IMPORTANT: Snapshot the parts array. During active streaming, `streamInfo.parts` grows as new
+    // deltas arrive. Iterating the live array would "chase" new parts and can prevent this replay
+    // from completing (which delays the "caught-up" signal and leaves the UI stuck on
+    // "Loading workspace..." during refresh).
+    const partsSnapshot = streamInfo.parts.slice();
+    for (const part of partsSnapshot) {
       await this.emitPartAsEvent(typedWorkspaceId, streamInfo.messageId, part);
     }
   }