From 77aff0f4fbfc2e968c3ebcb266e87578709a4539 Mon Sep 17 00:00:00 2001 From: Mohamed Boudra Date: Wed, 2 Sep 2026 13:37:32 +0200 Subject: [PATCH] fix(app): stabilize Mermaid streaming lifecycle (#4210) Drive the streaming assertion from the actual turn lifecycle so iframe replacement at completion is not mistaken for diagram loss. Preserve the cached timeline range during reload reconciliation so a projected message spanning the cache cursor replaces its partial prefix instead of being appended to it. --- .../app/e2e/browser/mermaid-streaming.spec.ts | 6 +- packages/app/e2e/support/helpers/diagram.ts | 17 +++--- packages/app/src/timeline/replica.test.ts | 58 +++++++++++++++++++ .../app/src/timeline/viewed-timeline-sync.ts | 22 ++++--- 4 files changed, 86 insertions(+), 17 deletions(-) diff --git a/packages/app/e2e/browser/mermaid-streaming.spec.ts b/packages/app/e2e/browser/mermaid-streaming.spec.ts index 89357e0e52..19ade0cb54 100644 --- a/packages/app/e2e/browser/mermaid-streaming.spec.ts +++ b/packages/app/e2e/browser/mermaid-streaming.spec.ts @@ -35,6 +35,7 @@ test("keeps a Mermaid diagram rendered while its message streams, completes, and page, }) => { test.setTimeout(120_000); + let diagramTurnCompleted!: Promise; const agent = await seedMockAgentWorkspace({ repoPrefix: "mermaid-streaming-", @@ -48,15 +49,16 @@ test("keeps a Mermaid diagram rendered while its message streams, completes, and await test.step("Open the conversation and request a diagram", async () => { await openAgentRoute(page, agent); await requestDiagram(agent); + diagramTurnCompleted = waitForDiagramTurnToComplete(agent); }); await test.step("The diagram renders and stays rendered while tokens arrive", async () => { await expectDiagramWithLabels(page, ["Start", "Middle"]); - await expectDiagramRemainsRenderedWhileStreaming(page); + await expectDiagramRemainsRenderedWhileStreaming(page, diagramTurnCompleted); }); await test.step("The completed diagram shows the final streamed content", async () => { - await waitForDiagramTurnToComplete(agent); + await diagramTurnCompleted; await expectCompletedDiagram(page, ["Start", "Done", "Release"]); }); diff --git a/packages/app/e2e/support/helpers/diagram.ts b/packages/app/e2e/support/helpers/diagram.ts index 6cf501c0fd..f320f40a6f 100644 --- a/packages/app/e2e/support/helpers/diagram.ts +++ b/packages/app/e2e/support/helpers/diagram.ts @@ -25,13 +25,16 @@ export async function expectDiagramWithLabels( } } -export async function expectDiagramRemainsRenderedWhileStreaming(page: Page): Promise { - const { diagram, svg } = renderedDiagram(page); - const samples = 80; - for (let sample = 0; sample < samples; sample += 1) { - await expect(diagram).toBeVisible({ timeout: 100 }); - await expect(svg).toBeVisible({ timeout: 500 }); - await page.waitForTimeout(25); +export async function expectDiagramRemainsRenderedWhileStreaming( + page: Page, + completion: Promise, +): Promise { + const diagram = page.getByRole("img", { name: DIAGRAM_NAME }).last(); + const completed = completion.then(() => true); + for (;;) { + expect(await diagram.isVisible()).toBe(true); + const didComplete = await Promise.race([completed, page.waitForTimeout(16).then(() => false)]); + if (didComplete) return completion; } } diff --git a/packages/app/src/timeline/replica.test.ts b/packages/app/src/timeline/replica.test.ts index 72cd01772e..daaadd8c08 100644 --- a/packages/app/src/timeline/replica.test.ts +++ b/packages/app/src/timeline/replica.test.ts @@ -115,6 +115,64 @@ describe("viewed timeline persistence", () => { owner.dispose(); }); + it("reconciles an overlapping projected message against its cached cursor", async () => { + useSessionStore.getState().initializeSession(SERVER_ID, null); + const partial = "```mermaid\nflowchart LR\n Start --> Mid"; + const complete = `${partial}dle\n Middle --> Done\n\`\`\``; + const owner = createOwner({ + readTimeline: async () => ({ + agentId: AGENT_ID, + items: [item("cached", partial, 4)], + range: { epoch: "epoch-1", startSeq: 1, endSeq: 4 }, + hasOlder: false, + }), + commitTimeline: () => undefined, + }); + + owner.replaceVisibleAgentIds("test", [AGENT_ID]); + await expect + .poll(() => + selectAgentTimelineState(useSessionStore.getState().sessions[SERVER_ID], AGENT_ID), + ) + .toMatchObject({ status: "painted" }); + + owner.applyTimelineResponse({ + requestId: "page-after-cache", + agentId: AGENT_ID, + agent: null, + direction: "after", + projection: "projected", + reset: false, + epoch: "epoch-1", + window: { minSeq: 1, maxSeq: 5, nextSeq: 6 }, + startCursor: { epoch: "epoch-1", seq: 5 }, + endCursor: { epoch: "epoch-1", seq: 5 }, + entries: [ + { + provider: "mock", + item: { type: "assistant_message", text: complete }, + timestamp: "2026-08-26T10:00:00.000Z", + seqStart: 2, + seqEnd: 5, + sourceSeqRanges: [{ startSeq: 2, endSeq: 5 }], + collapsed: ["assistant_merge"], + }, + ], + error: null, + hasNewer: false, + hasOlder: false, + staleCursor: false, + gap: false, + }); + + const session = useSessionStore.getState().sessions[SERVER_ID]; + expect([ + ...(session?.agentStreamTail.get(AGENT_ID) ?? []), + ...(session?.agentStreamHead.get(AGENT_ID) ?? []), + ]).toMatchObject([{ kind: "assistant_message", text: complete }]); + owner.dispose(); + }); + it("does not let a late cache read overwrite newer network state", async () => { useSessionStore.getState().initializeSession(SERVER_ID, null); let release!: (value: CachedTimeline) => void; diff --git a/packages/app/src/timeline/viewed-timeline-sync.ts b/packages/app/src/timeline/viewed-timeline-sync.ts index 49a127fbd9..d6bd3b2a20 100644 --- a/packages/app/src/timeline/viewed-timeline-sync.ts +++ b/packages/app/src/timeline/viewed-timeline-sync.ts @@ -90,11 +90,12 @@ async function prepareCachedTimeline(input: { export interface TimelineReplica { prepare(agentId: string): Promise; readCursor(agentId: string): { epoch: string; endSeq: number } | undefined; + readRange(agentId: string): AgentTimelineCursorState | undefined; timelineUpdated(agentId: string): void; } class TimelineReplicaOwner implements TimelineReplica { - private readonly cachedCursors = new Map(); + private readonly cachedRanges = new Map(); private readonly preparations = new Map>(); constructor( @@ -124,22 +125,24 @@ class TimelineReplicaOwner implements TimelineReplica { }); if (!stored) return; if (stored.range) { - this.cachedCursors.set(agentId, { - epoch: stored.range.epoch, - endSeq: stored.range.endSeq, - }); + this.cachedRanges.set(agentId, stored.range); } } readCursor(agentId: string): { epoch: string; endSeq: number } | undefined { - return this.cachedCursors.get(agentId); + const range = this.cachedRanges.get(agentId); + return range ? { epoch: range.epoch, endSeq: range.endSeq } : undefined; + } + + readRange(agentId: string): AgentTimelineCursorState | undefined { + return this.cachedRanges.get(agentId); } timelineUpdated(agentId: string): void { const session = useSessionStore.getState().sessions[this.serverId]; const timeline = selectAgentTimelineState(session, agentId); if (timeline.status !== "synced") return; - this.cachedCursors.delete(agentId); + this.cachedRanges.delete(agentId); this.storage.commitTimeline(this.serverId, agentId, { agentId, items: [...timeline.items, ...(session?.agentStreamHead.get(agentId) ?? [])], @@ -255,6 +258,7 @@ function finalizeProcessedTimeline(input: { function applyAuthoritativeTimelineResponse(input: { serverId: string; payload: TimelineResponsePayload; + cachedCursor?: AgentTimelineCursorState; recoverGap: (agentId: string, cursor: { epoch: string; endSeq: number }) => void; drainQueuedAgentMessage: (agentId: string) => void; }): boolean { @@ -264,7 +268,8 @@ function applyAuthoritativeTimelineResponse(input: { const session = useSessionStore.getState().sessions[serverId]; const timeline = selectAgentTimelineState(session, agentId); const activeInitDeferred = getInitDeferred(initKey); - const currentCursor = timeline.status === "synced" ? (timeline.range ?? undefined) : undefined; + const currentCursor = + timeline.status === "synced" ? (timeline.range ?? undefined) : input.cachedCursor; const result = processTimelineResponse({ payload, currentTail: timeline.status === "cold" ? [] : timeline.items, @@ -371,6 +376,7 @@ export function createViewedTimelineOwner(input: { const accepted = applyAuthoritativeTimelineResponse({ serverId: input.serverId, payload, + cachedCursor: input.replica.readRange(payload.agentId), recoverGap: (agentId, cursor) => sync.recoverGap(agentId, cursor), drainQueuedAgentMessage: input.drainQueuedAgentMessage, });