mirror of
https://github.com/openswarm-ai/openswarm.git
synced 2026-09-05 17:27:42 +02:00
[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:
co-authored by
Claude Fable 5.1
parent
a510334f23
commit
2afbdc5b05
@@ -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,
|
||||
}
|
||||
@@ -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({
|
||||
|
||||
@@ -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');
|
||||
});
|
||||
@@ -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. */
|
||||
|
||||
@@ -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;
|
||||
|
||||
Reference in New Issue
Block a user