Stop replayed turn events from re-typing themselves on session switch #272

New issue

Closed

opened 2026-11-13 12:46:52 +01:00 by omp-deck · 0 comments

omp-deck commented

2026-11-13 12:46:52 +01:00

Collaborator

Problem

Switching to a session that was running in the background re-renders its
already-produced assistant text character by character, as if it were being
typed live. This is most visible on sessions that are parked on an ask
dialog: the model finished writing long ago, yet attaching to the session
replays the whole turn at reading speed.

There is no typewriter animation in the client — nothing in apps/web/src
schedules characters. The cause is the replay path:

InProcessSessionHandle.emit() (apps/server/src/bridge/in-process.ts:1422-1443)
buffers every raw SDK event of the in-flight turn into currentTurnBuffer,
from turn_start until turn_end/agent_end, whether or not a client is attached.
A session blocked on an ask has not emitted turn_end — the agent is suspended
inside the tool call — so its buffer holds the entire turn indefinitely. These are
exactly the sessions with the largest buffers.
On subscribe, WsHub.handleSubscribe sends the snapshot and then loops the
buffer out as one WebSocket frame per event
(apps/server/src/ws.ts:283-285).
The client applies one set() per frame (apps/web/src/lib/store.ts:759-766),
so React paints between frames. That sequence of paints is the "typing".

It is also expensive. message_update carries the cumulative message, not a
delta (MessageUpdateEvent.message: AgentMessage,
node_modules/@oh-my-pi/pi-coding-agent/src/extensibility/extensions/types.ts:1203-1207),
and the reducer replaces the assistant blocks wholesale from it
(apps/web/src/lib/reducer.ts:149-152 → updateAssistantMessage,
reducer.ts:593-625). A turn with N text deltas therefore buffers and replays N
copies of a growing message: O(N²) bytes across the socket on every attach,
every reconnect, and every additional browser tab — plus unbounded buffer growth
in the server while the turn stays open.

Live streaming is correct and must stay exactly as it is; only the retroactive
replay is wrong.

Plan

1. Add a batched replay frame to the protocol

In packages/protocol/src/index.ts, add to the ServerFrame union (next to the
existing session_event member at line 1140):

/**
 * Replay of the in-flight turn's buffered events, sent once after `subscribed`.
 * Clients MUST fold these through the same reducer as `session_event`, in order,
 * in a single state update — applying them one at a time re-renders already
 * produced output as if it were streaming live.
 */
| { type: "session_events"; sessionId: string; events: AgentSessionEventJson[] }

AgentSessionEventJson (index.ts:853-855) is already an untyped passthrough;
no new event typing is needed.

Old browser tabs loaded before this change ignore unknown frames — handleFrame
ends in case "pong": default: return; (apps/web/src/lib/store.ts:982-984) —
so a stale tab shows the in-flight turn from the snapshot only, and recovers on
the next live event or a reload. No compatibility shim; this is a loopback tool
that serves its own bundle.

2. Coalesce superseded cumulative events as they enter the buffer

In apps/server/src/bridge/in-process.ts, add an exported pure helper above the
InProcessSessionHandle class (the class begins around line 1340; currentTurnBuffer
is declared at line 1347):

/**
 * Append `event` to a turn-replay buffer, dropping entries it fully supersedes.
 *
 * `message_update` and `tool_execution_update` both carry cumulative payloads —
 * the whole message so far, the whole partial result so far — so when two of
 * them are adjacent the earlier one contributes nothing a late subscriber can
 * observe. Adjacency is the safe predicate: the SDK cannot switch to a different
 * assistant message without emitting `message_start`/`message_end` in between,
 * and a parallel tool's update carries a different `toolCallId`, so in both
 * cases the run is broken and nothing is merged across it.
 *
 * Without this, a turn with N streamed deltas buffers N copies of a growing
 * message and replays all of them to every late subscriber.
 */
export function appendTurnEvent(buffer: AgentSessionEventJson[], event: AgentSessionEventJson): void;

Implementation:

const prev = buffer[buffer.length - 1]; — if absent, buffer.push(event) and return.
Read const type = (event as { type?: string }).type; and the same for prev.
If both are "message_update": buffer[buffer.length - 1] = event; and return.
If both are "tool_execution_update" and
String((prev as { toolCallId?: unknown }).toolCallId ?? "") === String((event as { toolCallId?: unknown }).toolCallId ?? ""):
replace in place and return.
Otherwise buffer.push(event).

Then in emit() (in-process.ts:1429-1431) replace

if (this.turnBufferActive) {
    this.currentTurnBuffer.push(event);
}

with if (this.turnBufferActive) appendTurnEvent(this.currentTurnBuffer, event);.

Nothing else in emit() changes: turn_start still resets the buffer,
turn_end/agent_end still clear it and drop turnBufferActive, and the
listener fan-out at in-process.ts:1436-1442 is untouched — live subscribers
keep receiving every single event. Only the replay buffer is condensed.

getCurrentTurnEvents() (in-process.ts:1445-1447) and
getCurrentTurnReplayEvents() (in-process.ts:1283-1287) keep their signatures.

3. Send the replay as one frame

In apps/server/src/ws.ts, replace the per-event loop at lines 283-285:

const replayEvents = this.bridge.getCurrentTurnReplayEvents(sessionId);
if (replayEvents.length > 0) {
    send(ws, { type: "session_events", sessionId, events: replayEvents });
}

Leave the three preceding replay loops (getUiReplayFrames,
getPlanModeReplayFrames, getIdeaModeReplayFrames, lines 274-282) and their
ordering alone — they emit a handful of frames and carry no streaming content.
The live listener registered at ws.ts:220-222 still sends individual
session_event frames, and because send is synchronous and the buffer is
copied before the first await-free send, an event arriving during attach is
delivered after the batch, preserving order.

The early-return branch for an already-subscribed connection (ws.ts:189-195)
re-sends only the snapshot today; leave it unchanged.

4. Fold the batch in a single client store update

In apps/web/src/lib/store.ts, add a case next to session_event
(lines 759-767). applyEvent is already imported and is pure:

case "session_events": {
    set((s) => {
        const prev = s.sessionsById[frame.sessionId];
        if (!prev) return {};
        let next = prev;
        for (const event of frame.events) next = applyEvent(next, event);
        return { sessionsById: { ...s.sessionsById, [frame.sessionId]: next } };
    });
    return;
}

One set() for the whole replay means one render: the turn appears complete,
instantly. Do not touch the session_event case — live events must keep
rendering per frame.

Critical files & anchors

apps/server/src/bridge/in-process.ts:1422-1447 — emit(), currentTurnBuffer,
getCurrentTurnEvents(); the only place the buffer is written.
apps/server/src/ws.ts:262-286 — ordered attach sequence; the comment block at
268-273 explains why subscribed must precede replay. Keep that ordering.
apps/web/src/lib/store.ts:759-767 — session_event handler to mirror.
apps/web/src/lib/reducer.ts:149-152,593-625 — proof that message_update is
applied as a wholesale overwrite, which is what makes coalescing lossless.
apps/server/src/bridge/in-process-queue.test.ts:1-60 — the established pattern
for exercising InProcessSessionHandle against a duck-typed stub session
without pulling in the SDK.

Verification

Run from the repo root.

New behaviour, server. Add apps/server/src/bridge/turn-replay.test.ts
driving appendTurnEvent directly (pure function, no session needed):
- Feeding turn_start, message_start, then message_update with
  message.content = [{type:"text",text:"a"}], "ab", "abc" yields a buffer
  of 3 entries whose last is the "abc" update — the two intermediates are gone.
- Interleaving message_update("a"), tool_execution_start,
  message_update("ab") keeps all three: the tool event breaks the run.
- tool_execution_update{toolCallId:"t1"}, tool_execution_update{toolCallId:"t2"},
  tool_execution_update{toolCallId:"t1"} keeps all three (parallel tools must
  not collapse into each other), while two consecutive t1 updates collapse to one.
  Run: bun test apps/server/src/bridge/turn-replay.test.ts.
New behaviour, client. Add a case to apps/web/src/lib/reducer.test.ts
(or a sibling test if that file is strictly queue-scoped) asserting that folding
[message_start(assistant), message_update("abc")] through applyEvent in
sequence produces the same SessionUi as the current per-event path: one
assistant message, blocks text "abc", isStreaming: true. This pins the
batched fold as equivalent to the frame-by-frame fold.
Run: bun test apps/web/src/lib.
Regression suite + types. bun test apps/server/src/bridge and
bun run typecheck (the latter covers the new ServerFrame member reaching
both server and web).
End-to-end, the reported symptom. bun run dev, open the web UI, start a
session whose prompt makes the agent produce a long answer and then call ask
(e.g. a planning prompt). Once the ask dialog is pending, navigate to another
session, then back. Expected: the full assistant text and tool cards appear in
a single paint, with the ask dialog restored; no progressive reveal. Compare
against git stash of the change to confirm the before/after difference.
While still attached and generating, confirm live text still streams token by
token — that path must be unaffected.
No duplicate bubbles. In the same manual run, confirm the replayed turn
renders exactly one assistant bubble for the in-flight message. Whether the
SDK's session.messages used by snapshot() (in-process.ts:1509) already
contains the in-flight partial assistant message is unverified; this change
does not alter which distinct events reach the reducer, so the rendering must
match pre-change behaviour either way. If a duplicate appears only after the
change, the fold order in step 4 is wrong — re-check that events are applied
in array order.

Assumptions & contingencies

Coalescing is invisible to the deck UI. The web client reads only
event.message from message_update and only event.partialResult from
tool_execution_update; it ignores assistantMessageEvent (the per-token delta).
If a future consumer needs the individual deltas of a past turn, it must read
them from the live stream, not the replay buffer — the buffer is explicitly a
"state to catch up to", not a log.
Only apps/web consumes session_event. Verified by grep: the server emits
it in ws.ts and the web store handles it; nothing else parses it. If an
external WebSocket consumer turns out to exist, it will ignore the unknown
session_events frame and simply miss the in-flight turn until the next live
event — acceptable, do not add a duplicate per-event replay for it.
If the batch frame turns out to be too large for a pathological turn (many
tool calls with large partial results), do not reintroduce per-event sending:
cap the batch by dropping the oldest tool_execution_update entries whose tool
already has a tool_execution_end later in the buffer, which are equally
superseded.

## Problem Switching to a session that was running in the background re-renders its already-produced assistant text character by character, as if it were being typed live. This is most visible on sessions that are *parked on an `ask` dialog*: the model finished writing long ago, yet attaching to the session replays the whole turn at reading speed. There is no typewriter animation in the client — nothing in `apps/web/src` schedules characters. The cause is the replay path: - `InProcessSessionHandle.emit()` (`apps/server/src/bridge/in-process.ts:1422-1443`) buffers **every** raw SDK event of the in-flight turn into `currentTurnBuffer`, from `turn_start` until `turn_end`/`agent_end`, whether or not a client is attached. - A session blocked on an `ask` has not emitted `turn_end` — the agent is suspended inside the tool call — so its buffer holds the entire turn indefinitely. These are exactly the sessions with the largest buffers. - On `subscribe`, `WsHub.handleSubscribe` sends the snapshot and then loops the buffer out as one WebSocket frame per event (`apps/server/src/ws.ts:283-285`). - The client applies **one `set()` per frame** (`apps/web/src/lib/store.ts:759-766`), so React paints between frames. That sequence of paints is the "typing". It is also expensive. `message_update` carries the *cumulative* message, not a delta (`MessageUpdateEvent.message: AgentMessage`, `node_modules/@oh-my-pi/pi-coding-agent/src/extensibility/extensions/types.ts:1203-1207`), and the reducer replaces the assistant blocks wholesale from it (`apps/web/src/lib/reducer.ts:149-152` → `updateAssistantMessage`, `reducer.ts:593-625`). A turn with N text deltas therefore buffers and replays N copies of a growing message: O(N²) bytes across the socket on every attach, every reconnect, and every additional browser tab — plus unbounded buffer growth in the server while the turn stays open. Live streaming is correct and must stay exactly as it is; only the retroactive replay is wrong. ## Plan ### 1. Add a batched replay frame to the protocol In `packages/protocol/src/index.ts`, add to the `ServerFrame` union (next to the existing `session_event` member at line 1140): ```ts /** * Replay of the in-flight turn's buffered events, sent once after `subscribed`. * Clients MUST fold these through the same reducer as `session_event`, in order, * in a single state update — applying them one at a time re-renders already * produced output as if it were streaming live. */ | { type: "session_events"; sessionId: string; events: AgentSessionEventJson[] } ``` `AgentSessionEventJson` (`index.ts:853-855`) is already an untyped passthrough; no new event typing is needed. Old browser tabs loaded before this change ignore unknown frames — `handleFrame` ends in `case "pong": default: return;` (`apps/web/src/lib/store.ts:982-984`) — so a stale tab shows the in-flight turn from the snapshot only, and recovers on the next live event or a reload. No compatibility shim; this is a loopback tool that serves its own bundle. ### 2. Coalesce superseded cumulative events as they enter the buffer In `apps/server/src/bridge/in-process.ts`, add an exported pure helper above the `InProcessSessionHandle` class (the class begins around line 1340; `currentTurnBuffer` is declared at line 1347): ```ts /** * Append `event` to a turn-replay buffer, dropping entries it fully supersedes. * * `message_update` and `tool_execution_update` both carry cumulative payloads — * the whole message so far, the whole partial result so far — so when two of * them are adjacent the earlier one contributes nothing a late subscriber can * observe. Adjacency is the safe predicate: the SDK cannot switch to a different * assistant message without emitting `message_start`/`message_end` in between, * and a parallel tool's update carries a different `toolCallId`, so in both * cases the run is broken and nothing is merged across it. * * Without this, a turn with N streamed deltas buffers N copies of a growing * message and replays all of them to every late subscriber. */ export function appendTurnEvent(buffer: AgentSessionEventJson[], event: AgentSessionEventJson): void; ``` Implementation: - `const prev = buffer[buffer.length - 1];` — if absent, `buffer.push(event)` and return. - Read `const type = (event as { type?: string }).type;` and the same for `prev`. - If both are `"message_update"`: `buffer[buffer.length - 1] = event;` and return. - If both are `"tool_execution_update"` **and** `String((prev as { toolCallId?: unknown }).toolCallId ?? "") === String((event as { toolCallId?: unknown }).toolCallId ?? "")`: replace in place and return. - Otherwise `buffer.push(event)`. Then in `emit()` (`in-process.ts:1429-1431`) replace ```ts if (this.turnBufferActive) { this.currentTurnBuffer.push(event); } ``` with `if (this.turnBufferActive) appendTurnEvent(this.currentTurnBuffer, event);`. Nothing else in `emit()` changes: `turn_start` still resets the buffer, `turn_end`/`agent_end` still clear it and drop `turnBufferActive`, and the listener fan-out at `in-process.ts:1436-1442` is untouched — **live subscribers keep receiving every single event**. Only the replay buffer is condensed. `getCurrentTurnEvents()` (`in-process.ts:1445-1447`) and `getCurrentTurnReplayEvents()` (`in-process.ts:1283-1287`) keep their signatures. ### 3. Send the replay as one frame In `apps/server/src/ws.ts`, replace the per-event loop at lines 283-285: ```ts const replayEvents = this.bridge.getCurrentTurnReplayEvents(sessionId); if (replayEvents.length > 0) { send(ws, { type: "session_events", sessionId, events: replayEvents }); } ``` Leave the three preceding replay loops (`getUiReplayFrames`, `getPlanModeReplayFrames`, `getIdeaModeReplayFrames`, lines 274-282) and their ordering alone — they emit a handful of frames and carry no streaming content. The live listener registered at `ws.ts:220-222` still sends individual `session_event` frames, and because `send` is synchronous and the buffer is copied before the first `await`-free send, an event arriving during attach is delivered after the batch, preserving order. The early-return branch for an already-subscribed connection (`ws.ts:189-195`) re-sends only the snapshot today; leave it unchanged. ### 4. Fold the batch in a single client store update In `apps/web/src/lib/store.ts`, add a case next to `session_event` (lines 759-767). `applyEvent` is already imported and is pure: ```ts case "session_events": { set((s) => { const prev = s.sessionsById[frame.sessionId]; if (!prev) return {}; let next = prev; for (const event of frame.events) next = applyEvent(next, event); return { sessionsById: { ...s.sessionsById, [frame.sessionId]: next } }; }); return; } ``` One `set()` for the whole replay means one render: the turn appears complete, instantly. Do not touch the `session_event` case — live events must keep rendering per frame. ## Critical files & anchors - `apps/server/src/bridge/in-process.ts:1422-1447` — `emit()`, `currentTurnBuffer`, `getCurrentTurnEvents()`; the only place the buffer is written. - `apps/server/src/ws.ts:262-286` — ordered attach sequence; the comment block at 268-273 explains why `subscribed` must precede replay. Keep that ordering. - `apps/web/src/lib/store.ts:759-767` — `session_event` handler to mirror. - `apps/web/src/lib/reducer.ts:149-152,593-625` — proof that `message_update` is applied as a wholesale overwrite, which is what makes coalescing lossless. - `apps/server/src/bridge/in-process-queue.test.ts:1-60` — the established pattern for exercising `InProcessSessionHandle` against a duck-typed stub session without pulling in the SDK. ## Verification Run from the repo root. 1. **New behaviour, server.** Add `apps/server/src/bridge/turn-replay.test.ts` driving `appendTurnEvent` directly (pure function, no session needed): - Feeding `turn_start`, `message_start`, then `message_update` with `message.content = [{type:"text",text:"a"}]`, `"ab"`, `"abc"` yields a buffer of 3 entries whose last is the `"abc"` update — the two intermediates are gone. - Interleaving `message_update("a")`, `tool_execution_start`, `message_update("ab")` keeps all three: the tool event breaks the run. - `tool_execution_update{toolCallId:"t1"}`, `tool_execution_update{toolCallId:"t2"}`, `tool_execution_update{toolCallId:"t1"}` keeps all three (parallel tools must not collapse into each other), while two consecutive `t1` updates collapse to one. Run: `bun test apps/server/src/bridge/turn-replay.test.ts`. 2. **New behaviour, client.** Add a case to `apps/web/src/lib/reducer.test.ts` (or a sibling test if that file is strictly queue-scoped) asserting that folding `[message_start(assistant), message_update("abc")]` through `applyEvent` in sequence produces the same `SessionUi` as the current per-event path: one assistant message, `blocks` text `"abc"`, `isStreaming: true`. This pins the batched fold as equivalent to the frame-by-frame fold. Run: `bun test apps/web/src/lib`. 3. **Regression suite + types.** `bun test apps/server/src/bridge` and `bun run typecheck` (the latter covers the new `ServerFrame` member reaching both server and web). 4. **End-to-end, the reported symptom.** `bun run dev`, open the web UI, start a session whose prompt makes the agent produce a long answer and then call `ask` (e.g. a planning prompt). Once the ask dialog is pending, navigate to another session, then back. Expected: the full assistant text and tool cards appear in a single paint, with the ask dialog restored; no progressive reveal. Compare against `git stash` of the change to confirm the before/after difference. While still attached and generating, confirm live text still streams token by token — that path must be unaffected. 5. **No duplicate bubbles.** In the same manual run, confirm the replayed turn renders exactly one assistant bubble for the in-flight message. Whether the SDK's `session.messages` used by `snapshot()` (`in-process.ts:1509`) already contains the in-flight partial assistant message is `unverified`; this change does not alter which distinct events reach the reducer, so the rendering must match pre-change behaviour either way. If a duplicate appears *only after* the change, the fold order in step 4 is wrong — re-check that events are applied in array order. ## Assumptions & contingencies - **Coalescing is invisible to the deck UI.** The web client reads only `event.message` from `message_update` and only `event.partialResult` from `tool_execution_update`; it ignores `assistantMessageEvent` (the per-token delta). If a future consumer needs the individual deltas of a *past* turn, it must read them from the live stream, not the replay buffer — the buffer is explicitly a "state to catch up to", not a log. - **Only `apps/web` consumes `session_event`.** Verified by grep: the server emits it in `ws.ts` and the web store handles it; nothing else parses it. If an external WebSocket consumer turns out to exist, it will ignore the unknown `session_events` frame and simply miss the in-flight turn until the next live event — acceptable, do not add a duplicate per-event replay for it. - **If the batch frame turns out to be too large for a pathological turn** (many tool calls with large partial results), do not reintroduce per-event sending: cap the batch by dropping the oldest `tool_execution_update` entries whose tool already has a `tool_execution_end` later in the buffer, which are equally superseded.