From b43faecdac427400b9db0ee8f591f51ee6ad983d Mon Sep 17 00:00:00 2001 From: ciregenz Date: Thu, 3 Sep 2026 10:40:05 -0700 Subject: [PATCH] [eric] chat: a streamed reply lands whole instead of re-typing from zero, the action bar waits for a reveal to settle, and a socket that connects mid-thought gets the thinking so far Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01C9zwUaHucUgrdxvK8FvjYT --- backend/apps/agents/agent_manager.py | 1 + backend/apps/agents/core/stream_snapshot.py | 13 +++++++++-- .../agents/manager/AgentManagerProtocol.py | 1 + backend/apps/agents/manager/run/TurnRunner.py | 2 +- .../manager/streaming/handle_stream_event.py | 13 ++++++++++- .../apps/agents/manager/streaming/state.py | 3 +++ backend/main.py | 2 +- backend/tests/test_stream_snapshot.py | 9 +++++++- .../app/components/share/ImportEntryPoint.tsx | 4 ++-- .../src/app/components/share/PublishModal.tsx | 4 ++-- .../src/app/components/share/ShareModal.tsx | 4 ++-- .../src/app/components/share/publishApi.ts | 4 ++-- frontend/src/app/components/share/shareApi.ts | 8 +++---- .../src/app/pages/AgentChat/AgentChat.tsx | 15 +++++++++++- .../AgentChat/bubbles/BurstRevealBubble.tsx | 9 +++++++- .../bubbles/actionBarWaitsForSettle.test.ts | 23 +++++++++++++++++++ linter/config/config.json | 1 + 17 files changed, 96 insertions(+), 20 deletions(-) create mode 100644 frontend/src/app/pages/AgentChat/bubbles/actionBarWaitsForSettle.test.ts diff --git a/backend/apps/agents/agent_manager.py b/backend/apps/agents/agent_manager.py index f6e90502..38538210 100644 --- a/backend/apps/agents/agent_manager.py +++ b/backend/apps/agents/agent_manager.py @@ -122,6 +122,7 @@ class AgentManager(SessionLifecycle, SessionHistory, SessionPersistence, Messagi self.tasks: Dict[str, asyncio.Task] = {} # Live mirror of the in-flight streamed assistant text per session, so a stop can persist the partial reply instantly instead of waiting out the multi-second SDK teardown the cancel handler sits behind. self.live_partial: Dict[str, PartialReply] = {} + self.live_thinking: Dict[str, PartialReply] = {} # Per-session cancel signal: the loop stashes its asyncio.Event here so a stop/close can set it. Lives on the manager, not the AgentSession model, so it stays out of serialization (an Event can't be model_dump'd). self.cancel_events: Dict[str, asyncio.Event] = {} # Persistent-client pool (lever A, flag-gated): one live CLI per session, reused across turns. diff --git a/backend/apps/agents/core/stream_snapshot.py b/backend/apps/agents/core/stream_snapshot.py index 8f6c7e9e..343650f6 100644 --- a/backend/apps/agents/core/stream_snapshot.py +++ b/backend/apps/agents/core/stream_snapshot.py @@ -15,13 +15,22 @@ from backend.apps.agents.manager.streaming.PartialReply import PartialReply @typechecked -def stream_snapshot_payload(session_id: str, live_partial: Dict[str, PartialReply]) -> Optional[dict]: +def stream_snapshot_payload( + session_id: str, + live_partial: Dict[str, PartialReply], + live_thinking: Optional[Dict[str, PartialReply]] = None, +) -> Optional[dict]: + # The answer wins over the thought: once text is flowing the thinking block is over. partial = live_partial.get(session_id) + role = "assistant" + if partial is None or not partial.msg_id or not partial.text: + partial = (live_thinking or {}).get(session_id) + role = "thinking" 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", + "role": role, "text": partial.text, } diff --git a/backend/apps/agents/manager/AgentManagerProtocol.py b/backend/apps/agents/manager/AgentManagerProtocol.py index 904f531a..3f183872 100644 --- a/backend/apps/agents/manager/AgentManagerProtocol.py +++ b/backend/apps/agents/manager/AgentManagerProtocol.py @@ -28,6 +28,7 @@ class AgentManagerProtocol: sessions: Dict[str, AgentSession] tasks: Dict[str, asyncio.Task] live_partial: Dict[str, PartialReply] + live_thinking: Dict[str, PartialReply] cancel_events: Dict[str, asyncio.Event] client_pool: Dict[str, ClientHandle] hook_ctxs: Dict[str, HookContext] diff --git a/backend/apps/agents/manager/run/TurnRunner.py b/backend/apps/agents/manager/run/TurnRunner.py index c186656f..49d37ae3 100644 --- a/backend/apps/agents/manager/run/TurnRunner.py +++ b/backend/apps/agents/manager/run/TurnRunner.py @@ -170,7 +170,7 @@ class TurnRunner(AgentManagerProtocol): if isinstance(message, StreamEvent): await handle_stream_event( - message, session, session_id, turn, thinking, self.live_partial + message, session, session_id, turn, thinking, self.live_partial, self.live_thinking ) elif isinstance(message, AssistantMessage): diff --git a/backend/apps/agents/manager/streaming/handle_stream_event.py b/backend/apps/agents/manager/streaming/handle_stream_event.py index 33d91157..4751d16b 100644 --- a/backend/apps/agents/manager/streaming/handle_stream_event.py +++ b/backend/apps/agents/manager/streaming/handle_stream_event.py @@ -5,7 +5,7 @@ writes the manager's live-partial mirror, exactly as it did inline.""" import time from datetime import datetime -from typing import Dict +from typing import Optional, Dict from uuid import uuid4 from typeguard import typechecked @@ -33,6 +33,7 @@ async def handle_stream_event( turn: TurnState, thinking: ThinkingState, live_partial: Dict[str, PartialReply], + live_thinking: Optional[Dict[str, PartialReply]] = None, ) -> None: event = message.event event_type = event.get("type") @@ -60,6 +61,8 @@ async def handle_stream_event( # Reasoning trace from thinking-capable models (GPT-5.3 Codex, Gemini 3 Pro/Flash, Claude with extended thinking). Rendered as a collapsible "thinking" message in the UI via the existing stream infrastructure, the frontend already handles role="thinking" for the DynamicIsland/agent card rendering. thinking_msg_id = uuid4().hex turn.stream_block_index_map[index] = thinking_msg_id + turn.stream_thinking_msg_id = thinking_msg_id + turn.stream_thinking_accum = "" # Server-stamp start so we can accumulate per-turn elapsed_ms across multiple thinking blocks (think → tool → think → answer turns sum correctly). thinking.block_starts[index] = time.time() await ws_manager.send_to_session(session_id, "agent:stream_start", { @@ -104,6 +107,9 @@ async def handle_stream_event( elif msg_id and delta_type == "thinking_delta": # Thinking content streams as thinking_delta with a "thinking" field (not "text") think_chunk = delta.get("thinking", "") + turn.stream_thinking_accum += think_chunk + if live_thinking is not None and turn.stream_thinking_msg_id: + live_thinking[session_id] = PartialReply(msg_id=turn.stream_thinking_msg_id, text=turn.stream_thinking_accum, branch_id=session.active_branch_id) await ws_manager.send_to_session(session_id, "agent:stream_delta", { "session_id": session_id, "message_id": msg_id, @@ -121,6 +127,11 @@ async def handle_stream_event( elif event_type == "content_block_stop": index = event.get("index") msg_id = turn.stream_block_index_map.get(index) + if msg_id and msg_id == turn.stream_thinking_msg_id: + turn.stream_thinking_msg_id = None + turn.stream_thinking_accum = "" + if live_thinking is not None: + live_thinking.pop(session_id, None) # If this was a thinking block, accumulate elapsed_ms server-side. We don't include per-block elapsed/tokens on the WS event, the pill stays in "Thinking…" until the AssistantMessage lands carrying the per-turn aggregate values. if index in thinking.block_starts: thinking.total_ms += int( diff --git a/backend/apps/agents/manager/streaming/state.py b/backend/apps/agents/manager/streaming/state.py index bf1a368d..3b220f17 100644 --- a/backend/apps/agents/manager/streaming/state.py +++ b/backend/apps/agents/manager/streaming/state.py @@ -39,6 +39,9 @@ class TurnState(BaseModel): stream_tool_msg_ids_ordered: List[str] = [] stream_block_index_map: Dict[int, str] = {} stream_text_accum: str = "" + # The thinking block in flight, so a socket that connects mid-thought can be handed it (the text snapshot's twin). + stream_thinking_msg_id: Optional[str] = None + stream_thinking_accum: str = "" current_turn_emitted: bool = False number: int = 0 first_event: bool = True diff --git a/backend/main.py b/backend/main.py index ea903a79..f29c838a 100644 --- a/backend/main.py +++ b/backend/main.py @@ -212,7 +212,7 @@ async def websocket_session(websocket: WebSocket, session_id: str): # 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) + snapshot = stream_snapshot_payload(session_id, p_am.live_partial, p_am.live_thinking) if snapshot is not None: await websocket.send_text(json.dumps({ "event": "agent:stream_snapshot", diff --git a/backend/tests/test_stream_snapshot.py b/backend/tests/test_stream_snapshot.py index 53b340e0..f8914f9f 100644 --- a/backend/tests/test_stream_snapshot.py +++ b/backend/tests/test_stream_snapshot.py @@ -27,4 +27,11 @@ def test_the_snapshot_is_sent_after_the_hello_ack_not_before(): 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 + assert "stream_snapshot_payload(session_id, p_am.live_partial, p_am.live_thinking)" in handler + + +def test_a_socket_that_connects_mid_thought_gets_the_thinking_so_far_and_the_answer_wins_once_it_flows(): + thinking = {"s1": PartialReply(msg_id="t1", text="Considering the four legs", branch_id="main")} + assert stream_snapshot_payload("s1", {}, thinking) == {"session_id": "s1", "message_id": "t1", "role": "thinking", "text": "Considering the four legs"} + text = {"s1": PartialReply(msg_id="m1", text="The Eiffel", branch_id="main")} + assert stream_snapshot_payload("s1", text, thinking)["role"] == "assistant" diff --git a/frontend/src/app/components/share/ImportEntryPoint.tsx b/frontend/src/app/components/share/ImportEntryPoint.tsx index b01c39af..f992c963 100644 --- a/frontend/src/app/components/share/ImportEntryPoint.tsx +++ b/frontend/src/app/components/share/ImportEntryPoint.tsx @@ -67,7 +67,7 @@ const ImportEntryPoint: React.FC = () => { setConfirm(null); confirmRef.current = false; } catch (e: any) { - setToast({ msg: e?.message || "We couldn't finish the import.", sev: 'error' }); + setToast({ msg: e?.message || "Couldn't finish installing.", sev: 'error' }); } finally { setCommitting(false); } @@ -88,7 +88,7 @@ const ImportEntryPoint: React.FC = () => { try { [, pf] = await Promise.all([delay(DIGEST_MS), importPreflight(file)]); } catch (e: any) { - setToast({ msg: e?.message || "We couldn't read this file.", sev: 'error' }); + setToast({ msg: e?.message || "Couldn't read that file.", sev: 'error' }); return; } if (importNeedsConfirm(pf)) { diff --git a/frontend/src/app/components/share/PublishModal.tsx b/frontend/src/app/components/share/PublishModal.tsx index d03589e3..e7ebd792 100644 --- a/frontend/src/app/components/share/PublishModal.tsx +++ b/frontend/src/app/components/share/PublishModal.tsx @@ -86,7 +86,7 @@ const PublishModal: React.FC = ({ outputId, outputName, open, onClose }) let alive = true; publishPreflight(outputId) .then((r) => alive && (setReview(r), setPhase('review'))) - .catch((e) => alive && (setErrorMsg(e?.message || "We couldn't check this app."), setPhase('error'))); + .catch((e) => alive && (setErrorMsg(e?.message || "Couldn't check this app. Try again."), setPhase('error'))); return () => { alive = false; }; @@ -141,7 +141,7 @@ const PublishModal: React.FC = ({ outputId, outputName, open, onClose }) onClose(); } catch (e: any) { setConfirmUnpublish(false); - setToast(e?.message || "We couldn't unpublish."); + setToast(e?.message || "Couldn't unpublish."); } finally { setBusy(false); } diff --git a/frontend/src/app/components/share/ShareModal.tsx b/frontend/src/app/components/share/ShareModal.tsx index d90f4122..ab4de476 100644 --- a/frontend/src/app/components/share/ShareModal.tsx +++ b/frontend/src/app/components/share/ShareModal.tsx @@ -44,7 +44,7 @@ const ShareModal: React.FC = ({ target, open, onClose }) => { let alive = true; exportPreflight(target) .then((pf) => alive && setPreflight(pf)) - .catch((e) => alive && setError(e?.message || "We couldn't read this for sharing.")) + .catch((e) => alive && setError(e?.message || "Couldn't read this. Try again.")) .finally(() => alive && setLoading(false)); return () => { alive = false; @@ -64,7 +64,7 @@ const ShareModal: React.FC = ({ target, open, onClose }) => { setToast(`Saved ${preflight.filename}`); onClose(); } catch (e: any) { - setError(e?.message || "We couldn't build the file."); + setError(e?.message || "Couldn't build the file. Try again."); } finally { setDownloading(false); } diff --git a/frontend/src/app/components/share/publishApi.ts b/frontend/src/app/components/share/publishApi.ts index 2f72248e..143d5d5f 100644 --- a/frontend/src/app/components/share/publishApi.ts +++ b/frontend/src/app/components/share/publishApi.ts @@ -21,7 +21,7 @@ export async function publishPreflight(outputId: string): Promise headers: { 'Content-Type': 'application/json' }, body: JSON.stringify({ output_id: outputId }), }); - if (!res.ok) throw new Error("We couldn't check this app."); + if (!res.ok) throw new Error("Couldn't check this app. Try again."); const data = await res.json(); return data.review as ReviewSummary; } @@ -48,7 +48,7 @@ export async function unpublishApp(outputId: string): Promise { headers: { 'Content-Type': 'application/json' }, body: JSON.stringify({ output_id: outputId }), }); - const fallback = "We couldn't unpublish this app. It may still be live."; + const fallback = "Couldn't unpublish. It may still be live."; if (!res.ok) throw new Error(fallback); let body: { ok?: boolean; error?: string } | null = null; try { diff --git a/frontend/src/app/components/share/shareApi.ts b/frontend/src/app/components/share/shareApi.ts index b2fd4771..158a2559 100644 --- a/frontend/src/app/components/share/shareApi.ts +++ b/frontend/src/app/components/share/shareApi.ts @@ -24,7 +24,7 @@ export async function exportPreflight(target: ShareTarget): Promise { form.append('file', file); // No Content-Type header: the browser sets the multipart boundary itself. const res = await fetch(`${API_BASE}/swarm/import/preflight`, { method: 'POST', body: form }); - if (!res.ok) throw new Error(await _detail(res, "We couldn't read this file.")); + if (!res.ok) throw new Error(await _detail(res, "Couldn't read that file.")); return res.json(); } @@ -64,6 +64,6 @@ export async function importCommit( headers: { 'Content-Type': 'application/json' }, body: JSON.stringify({ staging_token: stagingToken, accept_requirements: acceptRequirements }), }); - if (!res.ok) throw new Error(await _detail(res, "We couldn't finish the import.")); + if (!res.ok) throw new Error(await _detail(res, "Couldn't finish installing.")); return res.json(); } diff --git a/frontend/src/app/pages/AgentChat/AgentChat.tsx b/frontend/src/app/pages/AgentChat/AgentChat.tsx index 0e809c2a..596a3ec3 100644 --- a/frontend/src/app/pages/AgentChat/AgentChat.tsx +++ b/frontend/src/app/pages/AgentChat/AgentChat.tsx @@ -935,8 +935,11 @@ const AgentChat: React.FC = ({ sessionId: sessionIdProp, onClose const prevStreamingIdRef = useRef(null); const wasFollowingRef = useRef(true); const pinAbortRef = useRef(false); + // Read synchronously in render: justStreamedId is set in an effect AFTER the commit render, so the committed reply mounted with animate=true and typed itself out again from zero. + const lastStreamingIdRef = useRef(null); // Keep the follow-intent fresh while streaming so it's accurate at the instant the stream ends (handleScroll updates isAtBottomRef on every real scroll). if (streamingMessageId) wasFollowingRef.current = isAtBottomRef.current; + if (streamingMessageId) lastStreamingIdRef.current = streamingMessageId; useEffect(() => { const prev = prevStreamingIdRef.current; prevStreamingIdRef.current = streamingMessageId; @@ -1121,6 +1124,13 @@ const AgentChat: React.FC = ({ sessionId: sessionIdProp, onClose }, [id, mode, model, dispatch]); const [editingMessageId, setEditingMessageId] = useState(null); + // Assistant ids whose burst reveal has drained; the action bar under a reply mounts only then, so it never sits under the first revealed line. + const [settledIds, setSettledIds] = useState>(() => new Set()); + // Ids whose burst reveal actually started; the burst flag itself flips false one render later (the id is then "seen"), which used to mount the bar under an empty, still-revealing bubble. + const revealingIdsRef = useRef>(new Set()); + const markSettled = useCallback((id: string) => { + setSettledIds((prev) => (prev.has(id) ? prev : new Set(prev).add(id))); + }, []); const handleSaveEdit = useCallback( (messageId: string, newContent: string) => { @@ -2007,7 +2017,9 @@ const AgentChat: React.FC = ({ sessionId: sessionIdProp, onClose !isEditing && !seenMessageIdsRef.current.has(msg.id) && msg.id !== justStreamedId && + msg.id !== lastStreamingIdRef.current && (sessionRunning || awaitingResponse); + if (burstAnimate) revealingIdsRef.current.add(msg.id); seenMessageIdsRef.current.add(msg.id); const siblings = getSiblingBranches(msg.id); const hasBranches = siblings.length > 0; @@ -2046,6 +2058,7 @@ const AgentChat: React.FC = ({ sessionId: sessionIdProp, onClose message={msg} animate={burstAnimate} onGrew={stickToBottomIfNeeded} + onSettled={() => markSettled(msg.id)} viewportHeight={viewportHeight} viewportWidth={viewportWidth} scrollRoot={scrollRoot} @@ -2061,7 +2074,7 @@ const AgentChat: React.FC = ({ sessionId: sessionIdProp, onClose scrollRoot={scrollRoot} /> )} - {!isEditing && (msg.role === 'user' || (msg.role === 'assistant' && lastAssistantIdsInTurn.has(msg.id))) && ( + {!isEditing && (msg.role === 'user' || (msg.role === 'assistant' && lastAssistantIdsInTurn.has(msg.id) && (!revealingIdsRef.current.has(msg.id) || settledIds.has(msg.id)))) && ( void; + // Fires once when the reveal has drained (or at once when nothing animates), so chrome below the bubble can wait for the answer to settle. + onSettled?: () => void; viewportHeight?: number; viewportWidth?: number; scrollRoot?: Element | null; @@ -16,7 +18,7 @@ interface Props { /** Short post-tool answers often COMMIT whole and skip the streaming slice entirely, so they popped while true streams typed. Route fresh commits through the same smooth reveal (assistant-ui's drain pattern), then settle into the identical committed render so the handoff can't flash. */ -function BurstRevealBubble({ message, animate, onGrew, viewportHeight, viewportWidth, scrollRoot }: Props): React.ReactElement { +function BurstRevealBubble({ message, animate, onGrew, onSettled, viewportHeight, viewportWidth, scrollRoot }: Props): React.ReactElement { const [shouldAnimate] = useState(animate); const full = typeof message.content === 'string' ? message.content : ''; const [done, setDone] = useState(!shouldAnimate || full.length === 0); @@ -26,6 +28,11 @@ function BurstRevealBubble({ message, animate, onGrew, viewportHeight, viewportW }, [done, text.length, full.length]); const grewRef = React.useRef(onGrew); grewRef.current = onGrew; + const settledRef = React.useRef(onSettled); + settledRef.current = onSettled; + useEffect(() => { + if (done) settledRef.current?.(); + }, [done]); useEffect(() => { if (!done) grewRef.current?.(); }, [text.length, done]); diff --git a/frontend/src/app/pages/AgentChat/bubbles/actionBarWaitsForSettle.test.ts b/frontend/src/app/pages/AgentChat/bubbles/actionBarWaitsForSettle.test.ts new file mode 100644 index 00000000..e0bff35e --- /dev/null +++ b/frontend/src/app/pages/AgentChat/bubbles/actionBarWaitsForSettle.test.ts @@ -0,0 +1,23 @@ +import { test } from 'node:test'; +import assert from 'node:assert/strict'; +import fs from 'node:fs'; +import path from 'node:path'; + +// Eric's send-to-answer table (2026-09-03): "the action bar and the Done chip are already mounted under the +// first revealed line, while the text is still revealing". The bar under a burst-revealed reply now mounts +// only once the reveal has drained; history and streamed replies (no burst) keep it at once. +test('the action bar under a burst-revealed reply mounts only after the reveal settles', () => { + const chat = fs.readFileSync(path.join(process.cwd(), 'src/app/pages/AgentChat/AgentChat.tsx'), 'utf8'); + assert.match(chat, /onSettled=\{\(\) => markSettled\(msg\.id\)\}/); + assert.match(chat, /lastAssistantIdsInTurn\.has\(msg\.id\) && \(!revealingIdsRef\.current\.has\(msg\.id\) \|\| settledIds\.has\(msg\.id\)\)/); + assert.match(chat, /if \(burstAnimate\) revealingIdsRef\.current\.add\(msg\.id\);/, 'the gate keys on a reveal that STARTED, not on the burst flag, which flips one render later'); + const bubble = fs.readFileSync(path.join(process.cwd(), 'src/app/pages/AgentChat/bubbles/BurstRevealBubble.tsx'), 'utf8'); + assert.match(bubble, /if \(done\) settledRef\.current\?\.\(\);/, 'onSettled must fire from the done flag, once'); +}); + +test('a reply that just streamed lands whole: it never re-types itself from zero after the commit', () => { + const chat = fs.readFileSync(path.join(process.cwd(), 'src/app/pages/AgentChat/AgentChat.tsx'), 'utf8'); + assert.match(chat, /if \(streamingMessageId\) lastStreamingIdRef\.current = streamingMessageId;/, 'remembered during render, not in an effect'); + const burst = chat.indexOf('const burstAnimate ='); + assert.ok(chat.slice(burst, burst + 400).includes('msg.id !== lastStreamingIdRef.current'), 'the burst reveal must exclude the id that was streaming a moment ago'); +}); diff --git a/linter/config/config.json b/linter/config/config.json index 8ea3fbec..d0b4a60e 100644 --- a/linter/config/config.json +++ b/linter/config/config.json @@ -270,6 +270,7 @@ "backend/apps/agents/manager/streaming/PartialReply.py::PartialReply.msg_id", "backend/apps/agents/manager/streaming/state.py::ThinkingState.msg_id", "backend/apps/agents/manager/streaming/state.py::TurnState.stream_text_msg_id", + "backend/apps/agents/manager/streaming/state.py::TurnState.stream_thinking_msg_id", "backend/apps/apps_sdk/apps_sdk.py::GrantResolveRequest.request_id", "backend/apps/apps_sdk/tool_grants.py::PendingGrant.request_id", "backend/apps/dashboards/models.py::BrowserCardPosition.browser_id",