[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
This commit is contained in:
ciregenz
2026-09-03 00:18:22 -07:00
co-authored by Claude Fable 5.1
parent a510334f23
commit 2afbdc5b05
6 changed files with 130 additions and 2 deletions
@@ -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,
}
+10
View File
@@ -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({
+30
View File
@@ -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
@@ -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');
});
+14 -1
View File
@@ -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. */
+8 -1
View File
@@ -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;