From 2afbdc5b05ffbbceb022dc30efc8659797026f04 Mon Sep 17 00:00:00 2001
From: ciregenz <ciregenz@berkeley.edu>
Date: Thu, 3 Sep 2026 00:18:22 -0700
Subject: [PATCH] [eric] streaming: a session socket that connects mid-reply is
 handed the text so far after its resume ack, instead of a static Thinking
 then the whole reply at once

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01C9zwUaHucUgrdxvK8FvjYT
---
 backend/apps/agents/core/stream_snapshot.py   | 27 ++++++++++++
 backend/main.py                               | 10 +++++
 backend/tests/test_stream_snapshot.py         | 30 ++++++++++++++
 .../src/shared/state/streamSnapshot.test.ts   | 41 +++++++++++++++++++
 frontend/src/shared/state/streamingSlice.ts   | 15 ++++++-
 frontend/src/shared/ws/WebSocketManager.ts    |  9 +++-
 6 files changed, 130 insertions(+), 2 deletions(-)
 create mode 100644 backend/apps/agents/core/stream_snapshot.py
 create mode 100644 backend/tests/test_stream_snapshot.py
 create mode 100644 frontend/src/shared/state/streamSnapshot.test.ts

diff --git a/backend/apps/agents/core/stream_snapshot.py b/backend/apps/agents/core/stream_snapshot.py
new file mode 100644
index 00000000..8f6c7e9e
--- /dev/null
+++ b/backend/apps/agents/core/stream_snapshot.py
@@ -0,0 +1,27 @@
+"""The text of an assistant reply that is streaming RIGHT NOW, for a session socket that just connected.
+
+Every delta before the per-session socket connects is lost to that client: the ring replays them, but
+the client drops replayed stream frames on purpose (they predate its resume ack). Until now the
+transcript stayed on a static "Thinking..." and then received the whole reply at once, which is the
+"streams halfway, wipes, then retypes everything fast" Eric kept seeing. The manager already keeps the
+accumulated text per session (`live_partial`, for the crash snapshot); this hands it to the socket.
+"""
+
+from typing import Dict, Optional
+
+from typeguard import typechecked
+
+from backend.apps.agents.manager.streaming.PartialReply import PartialReply
+
+
+@typechecked
+def stream_snapshot_payload(session_id: str, live_partial: Dict[str, PartialReply]) -> Optional[dict]:
+    partial = live_partial.get(session_id)
+    if partial is None or not partial.msg_id or not partial.text:
+        return None
+    return {
+        "session_id": session_id,
+        "message_id": partial.msg_id,
+        "role": "assistant",
+        "text": partial.text,
+    }
diff --git a/backend/main.py b/backend/main.py
index 13d3da71..ea903a79 100644
--- a/backend/main.py
+++ b/backend/main.py
@@ -209,6 +209,16 @@ async def websocket_session(websocket: WebSocket, session_id: str):
                         "ack": ack,
                     },
                 }))
+                # AFTER the ack, never before: the client drops stream frames until it has the ack.
+                from backend.apps.agents.agent_manager import agent_manager as p_am
+                from backend.apps.agents.core.stream_snapshot import stream_snapshot_payload
+                snapshot = stream_snapshot_payload(session_id, p_am.live_partial)
+                if snapshot is not None:
+                    await websocket.send_text(json.dumps({
+                        "event": "agent:stream_snapshot",
+                        "session_id": session_id,
+                        "data": snapshot,
+                    }))
             elif event == "client:ping":
                 # Heartbeat. Cheap, keeps NATs/firewalls from silently dropping the connection. Carry the client's nonce back so it can match pong→ping for round-trip latency tracking if it wants.
                 await websocket.send_text(json.dumps({
diff --git a/backend/tests/test_stream_snapshot.py b/backend/tests/test_stream_snapshot.py
new file mode 100644
index 00000000..53b340e0
--- /dev/null
+++ b/backend/tests/test_stream_snapshot.py
@@ -0,0 +1,30 @@
+"""A session socket that connects mid-reply gets the text so far, once, and nothing when there is none."""
+
+import pathlib
+
+from backend.apps.agents.core.stream_snapshot import stream_snapshot_payload
+from backend.apps.agents.manager.streaming.PartialReply import PartialReply
+
+
+def test_mid_reply_connect_gets_the_accumulated_text():
+    live = {"s1": PartialReply(msg_id="m1", text="The Eiffel Tower is a wrought-iron", branch_id="main")}
+    assert stream_snapshot_payload("s1", live) == {
+        "session_id": "s1", "message_id": "m1", "role": "assistant", "text": "The Eiffel Tower is a wrought-iron",
+    }
+
+
+def test_no_reply_in_flight_means_no_snapshot():
+    assert stream_snapshot_payload("s1", {}) is None
+    assert stream_snapshot_payload("s1", {"s1": PartialReply(msg_id="m1", text="", branch_id="main")}) is None
+    assert stream_snapshot_payload("s1", {"other": PartialReply(msg_id="m1", text="x", branch_id="main")}) is None
+
+
+def test_the_snapshot_is_sent_after_the_hello_ack_not_before():
+    """The client drops every stream frame that arrives before its resume ack, so a snapshot sent
+    earlier would be dropped exactly like the replayed deltas it exists to replace."""
+    src = pathlib.Path("backend/main.py").read_text()
+    hello = src.index('"event": "server:hello"')
+    snapshot = src.index('"event": "agent:stream_snapshot"')
+    assert hello < snapshot
+    handler = src[src.index('if event == "client:hello":'):src.index('elif event == "client:ping":')]
+    assert "stream_snapshot_payload(session_id, p_am.live_partial)" in handler
diff --git a/frontend/src/shared/state/streamSnapshot.test.ts b/frontend/src/shared/state/streamSnapshot.test.ts
new file mode 100644
index 00000000..66b97019
--- /dev/null
+++ b/frontend/src/shared/state/streamSnapshot.test.ts
@@ -0,0 +1,41 @@
+import { test } from 'node:test';
+import assert from 'node:assert/strict';
+import fs from 'node:fs';
+import path from 'node:path';
+import reducer, { streamStart, streamSnapshot, streamDelta, streamEnd } from './streamingSlice';
+
+// Every delta that arrives before the per-session socket connects is lost to that client (the ring
+// replays them, the client drops pre-ack stream frames on purpose), so the transcript sat on a static
+// "Thinking..." and then got the whole reply at once. The server now sends the text so far right after
+// the resume ack; the slice seeds the bubble with it and later deltas append.
+
+test('a snapshot seeds the streaming bubble and later deltas append to it', () => {
+  let s = reducer(undefined, streamSnapshot({ sessionId: 'a', messageId: 'm1', role: 'assistant', text: 'The Eiffel' }));
+  s = reducer(s, streamDelta({ sessionId: 'a', messageId: 'm1', delta: ' Tower' }));
+  assert.equal(s.bySession.a?.content, 'The Eiffel Tower');
+  s = reducer(s, streamEnd({ sessionId: 'a', messageId: 'm1' }));
+  assert.equal(s.bySession.a, undefined);
+});
+
+test('a snapshot never shortens text the client already has for the same message', () => {
+  let s = reducer(undefined, streamStart({ sessionId: 'a', messageId: 'm1', role: 'assistant' }));
+  s = reducer(s, streamDelta({ sessionId: 'a', messageId: 'm1', delta: 'The Eiffel Tower is tall' }));
+  s = reducer(s, streamSnapshot({ sessionId: 'a', messageId: 'm1', role: 'assistant', text: 'The Eiffel' }));
+  assert.equal(s.bySession.a?.content, 'The Eiffel Tower is tall');
+});
+
+test('a snapshot for a newer message replaces a stale bubble', () => {
+  let s = reducer(undefined, streamStart({ sessionId: 'a', messageId: 'old', role: 'assistant' }));
+  s = reducer(s, streamSnapshot({ sessionId: 'a', messageId: 'new', role: 'assistant', text: 'Second reply so far' }));
+  assert.equal(s.bySession.a?.id, 'new');
+  assert.equal(s.bySession.a?.content, 'Second reply so far');
+});
+
+test('the socket handles the snapshot ABOVE the replay-skip guard that drops pre-ack stream frames', () => {
+  const src = fs.readFileSync(path.join(process.cwd(), 'src/shared/ws/WebSocketManager.ts'), 'utf8');
+  const snapshot = src.indexOf("event === 'agent:stream_snapshot'");
+  const guard = src.indexOf('if (!this.resumeAcked) break;');
+  assert.ok(snapshot > 0 && guard > 0 && snapshot < guard, 'the snapshot must be handled before the pre-ack guard');
+  const dashboardSkip = src.indexOf('if (this.skipStreamEvents) {');
+  assert.ok(snapshot < dashboardSkip, 'the snapshot handler decides skipStreamEvents itself, above the generic skip');
+});
diff --git a/frontend/src/shared/state/streamingSlice.ts b/frontend/src/shared/state/streamingSlice.ts
index 95814275..d14695e0 100644
--- a/frontend/src/shared/state/streamingSlice.ts
+++ b/frontend/src/shared/state/streamingSlice.ts
@@ -41,6 +41,19 @@ const streamingSlice = createSlice({
         tool_name: action.payload.toolName,
       };
     },
+    // A socket that connected mid-reply is handed the text so far; every delta before it was lost to this client, so this seeds rather than appends.
+    streamSnapshot(
+      state,
+      action: PayloadAction<{ sessionId: string; messageId: string; role: StreamingMessage['role']; text: string }>,
+    ) {
+      const entry = state.bySession[action.payload.sessionId];
+      if (entry && entry.id === action.payload.messageId && entry.content.length >= action.payload.text.length) return;
+      state.bySession[action.payload.sessionId] = {
+        id: action.payload.messageId,
+        role: action.payload.role,
+        content: action.payload.text,
+      };
+    },
     streamDelta(
       state,
       action: PayloadAction<{ sessionId: string; messageId: string; delta: string }>,
@@ -91,7 +104,7 @@ const streamingSlice = createSlice({
   },
 });
 
-export const { streamStart, streamDelta, streamEnd, clearStreamingForSession } = streamingSlice.actions;
+export const { streamStart, streamSnapshot, streamDelta, streamEnd, clearStreamingForSession } = streamingSlice.actions;
 export default streamingSlice.reducer;
 
 /** Subscribes only to one session's stream entry; null when no stream is active. */
diff --git a/frontend/src/shared/ws/WebSocketManager.ts b/frontend/src/shared/ws/WebSocketManager.ts
index 26afeed0..9f20b435 100644
--- a/frontend/src/shared/ws/WebSocketManager.ts
+++ b/frontend/src/shared/ws/WebSocketManager.ts
@@ -32,7 +32,7 @@ import {
   setQueued,
   clearTurnLabel,
 } from '../state/agentsSlice';
-import { streamStart, streamDelta, streamEnd, clearStreamingForSession } from '../state/streamingSlice';
+import { streamStart, streamSnapshot, streamDelta, streamEnd, clearStreamingForSession } from '../state/streamingSlice';
 import { remountAppPreview } from '../state/outputsSlice';
 import { BackgroundDeltaBuffer } from './BackgroundDeltaBuffer';
 import { interactionActive, installInteractionListeners } from '../interactionPriority';
@@ -493,6 +493,13 @@ class WebSocketManager {
       return;
     }
 
+    // Sent once per (re)connect, after the resume ack, so it must not sit behind the replay-skip guard below that drops pre-ack stream frames.
+    if (event === 'agent:stream_snapshot') {
+      if (session_id && data.message_id && typeof data.text === 'string' && !this.skipStreamEvents) {
+        store.dispatch(streamSnapshot({ sessionId: session_id, messageId: data.message_id, role: data.role, text: data.text }));
+      }
+      return;
+    }
     if (this.skipStreamEvents) {
       if (event === 'agent:stream_start' || event === 'agent:stream_delta' || event === 'agent:stream_end') {
         return;