diff --git a/backend/apps/agents/agent_manager.py b/backend/apps/agents/agent_manager.py index 38538210..409efe46 100644 --- a/backend/apps/agents/agent_manager.py +++ b/backend/apps/agents/agent_manager.py @@ -361,6 +361,7 @@ class AgentManager(SessionLifecycle, SessionHistory, SessionPersistence, Messagi "message_id": p_tool_msg_id, }) self.live_partial.pop(session_id, None) + self.live_thinking.pop(session_id, None) # Tell the user we self-healed instead of retrying in silence: the frontend renders this as a muted transient pill (same language as the rate-limit pill), not an error card. try: await ws_manager.send_to_session(session_id, "agent:context_recovered", { @@ -428,6 +429,7 @@ class AgentManager(SessionLifecycle, SessionHistory, SessionPersistence, Messagi p_is_live_task = self.tasks.get(session_id) is asyncio.current_task() if p_is_live_task: self.live_partial.pop(session_id, None) + self.live_thinking.pop(session_id, None) # The floor: every terminal path funnels through here, so "turn ended with nothing readable" stops being representable instead of being caught shape by shape. try: from backend.apps.agents.manager.run.turn_spoke import ensure_turn_spoke diff --git a/backend/apps/agents/core/stream_snapshot.py b/backend/apps/agents/core/stream_snapshot.py index 343650f6..2ca29c0b 100644 --- a/backend/apps/agents/core/stream_snapshot.py +++ b/backend/apps/agents/core/stream_snapshot.py @@ -19,8 +19,12 @@ def stream_snapshot_payload( session_id: str, live_partial: Dict[str, PartialReply], live_thinking: Optional[Dict[str, PartialReply]] = None, + in_flight: bool = True, ) -> Optional[dict]: # The answer wins over the thought: once text is flowing the thinking block is over. + # A finished turn has nothing live by definition; whatever the dicts still hold is a leak, never a reply. + if not in_flight: + return None partial = live_partial.get(session_id) role = "assistant" if partial is None or not partial.msg_id or not partial.text: diff --git a/backend/apps/agents/manager/RunSupport.py b/backend/apps/agents/manager/RunSupport.py index bc51c5dc..a0a3e8b3 100644 --- a/backend/apps/agents/manager/RunSupport.py +++ b/backend/apps/agents/manager/RunSupport.py @@ -200,6 +200,7 @@ class RunSupport(AgentManagerProtocol): instantly instead of waiting out the SDK teardown the cancel handler sits behind. Returns True if it committed something.""" live = self.live_partial.pop(session.id, None) + self.live_thinking.pop(session.id, None) if not live: return False text = live.text or "" diff --git a/backend/apps/agents/manager/run/TurnRunner.py b/backend/apps/agents/manager/run/TurnRunner.py index 49d37ae3..3d409794 100644 --- a/backend/apps/agents/manager/run/TurnRunner.py +++ b/backend/apps/agents/manager/run/TurnRunner.py @@ -176,7 +176,7 @@ class TurnRunner(AgentManagerProtocol): elif isinstance(message, AssistantMessage): flight_recorder.crumb(session_id, "assistant-msg") await handle_assistant_message( - message, session, session_id, turn, thinking, self.live_partial, self.sessions + message, session, session_id, turn, thinking, self.live_partial, self.sessions, self.live_thinking ) elif isinstance(message, ResultMessage): flight_recorder.crumb(session_id, "result-msg", subtype=str(getattr(message, "subtype", ""))) @@ -226,6 +226,7 @@ class TurnRunner(AgentManagerProtocol): turn.stream_text_msg_id = None turn.stream_text_accum = "" self.live_partial.pop(session_id, None) + self.live_thinking.pop(session_id, None) for p_tool_msg_id in turn.stream_tool_msg_ids_ordered: await ws_manager.send_to_session(session_id, "agent:stream_end", { "session_id": session_id, diff --git a/backend/apps/agents/manager/session/SessionLifecycle.py b/backend/apps/agents/manager/session/SessionLifecycle.py index e117c970..bfaec573 100644 --- a/backend/apps/agents/manager/session/SessionLifecycle.py +++ b/backend/apps/agents/manager/session/SessionLifecycle.py @@ -100,6 +100,7 @@ class SessionLifecycle(AgentManagerProtocol): self.sessions.pop(session_id, None) self.tasks.pop(session_id, None) self.live_partial.pop(session_id, None) + self.live_thinking.pop(session_id, None) self.cancel_events.pop(session_id, None) self.pending_messages.pop(session_id, None) view_builder_render_retry_counts.pop(session_id, None) diff --git a/backend/apps/agents/manager/streaming/handle_assistant_message.py b/backend/apps/agents/manager/streaming/handle_assistant_message.py index 9e4ddd89..2612d87c 100644 --- a/backend/apps/agents/manager/streaming/handle_assistant_message.py +++ b/backend/apps/agents/manager/streaming/handle_assistant_message.py @@ -38,6 +38,7 @@ async def handle_assistant_message( thinking: ThinkingState, live_partial: Dict[str, PartialReply], sessions: Dict[str, AgentSession], + live_thinking: Optional[Dict[str, PartialReply]] = None, ) -> None: from claude_agent_sdk.types import ThinkingBlock, TextBlock, ToolUseBlock @@ -266,6 +267,8 @@ async def handle_assistant_message( if p_fault_armed("empty_finish"): turn.stream_text_accum = "" live_partial.pop(session_id, None) + if live_thinking is not None: + live_thinking.pop(session_id, None) return asst_msg = Message( id=turn.stream_text_msg_id or uuid4().hex, @@ -276,6 +279,8 @@ async def handle_assistant_message( upsert_message(session, asst_msg) turn.stream_text_accum = "" live_partial.pop(session_id, None) + if live_thinking is not None: + live_thinking.pop(session_id, None) await ws_manager.send_to_session(session_id, "agent:message", { "session_id": session_id, "message": asst_msg.model_dump(mode="json"), diff --git a/backend/apps/agents/manager/streaming/handle_stream_event.py b/backend/apps/agents/manager/streaming/handle_stream_event.py index 4751d16b..6865cf6f 100644 --- a/backend/apps/agents/manager/streaming/handle_stream_event.py +++ b/backend/apps/agents/manager/streaming/handle_stream_event.py @@ -48,6 +48,10 @@ async def handle_stream_event( block_type = block.get("type") if block_type == "text": + # The answer has begun, so the thought is over whether or not the CLI ever sent the block's stop + # (it does not for Haiku's thinking; a socket connecting later was handed a "still thinking" bubble). + if live_thinking is not None: + live_thinking.pop(session_id, None) if turn.stream_text_msg_id is None: turn.stream_text_msg_id = uuid4().hex await ws_manager.send_to_session(session_id, "agent:stream_start", { @@ -144,6 +148,8 @@ async def handle_stream_event( }) elif event_type == "message_stop": + if live_thinking is not None: + live_thinking.pop(session_id, None) if turn.stream_text_msg_id: await ws_manager.send_to_session(session_id, "agent:stream_end", { "session_id": session_id, diff --git a/backend/main.py b/backend/main.py index f29c838a..e116c739 100644 --- a/backend/main.py +++ b/backend/main.py @@ -212,7 +212,9 @@ 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, p_am.live_thinking) + p_sess = p_am.sessions.get(session_id) + p_in_flight = bool(p_sess is not None and p_sess.status in ("running", "waiting_approval")) + snapshot = stream_snapshot_payload(session_id, p_am.live_partial, p_am.live_thinking, in_flight=p_in_flight) 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 f8914f9f..bf1c9742 100644 --- a/backend/tests/test_stream_snapshot.py +++ b/backend/tests/test_stream_snapshot.py @@ -27,7 +27,7 @@ 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, p_am.live_thinking)" in handler + assert "stream_snapshot_payload(session_id, p_am.live_partial, p_am.live_thinking, in_flight=p_in_flight)" in handler def test_a_socket_that_connects_mid_thought_gets_the_thinking_so_far_and_the_answer_wins_once_it_flows(): @@ -35,3 +35,65 @@ def test_a_socket_that_connects_mid_thought_gets_the_thinking_so_far_and_the_ans 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" + + +def test_a_finished_turn_never_snapshots_even_if_a_dict_still_holds_text(): + """Eric, 2026-09-03, on the packaged exp.4: every card opened after a Haiku reply showed a "still + thinking" bubble with a caret, because the thinking text was never dropped and the snapshot had no + notion of whether the turn was over. The door refuses first; the leak is sealed separately.""" + thinking = {"s1": PartialReply(msg_id="t1", text="leftover thought", branch_id="main")} + text = {"s1": PartialReply(msg_id="m1", text="leftover reply", branch_id="main")} + assert stream_snapshot_payload("s1", text, thinking, in_flight=False) is None + assert stream_snapshot_payload("s1", text, thinking, in_flight=True) is not None + + +def test_the_thought_is_dropped_when_the_answer_starts_even_without_a_block_stop(): + import asyncio + from claude_agent_sdk.types import StreamEvent + from backend.apps.agents.core.models import AgentSession + from backend.apps.agents.manager.streaming import handle_stream_event as mod + from backend.apps.agents.manager.streaming.state import ThinkingState, TurnState + + class P_WS: + async def send_to_session(self, sid, event, data): + return None + + def ev(payload): + return StreamEvent(uuid="u", session_id="s1", event=payload, parent_tool_use_id=None) + + session = AgentSession(name="probe", model="opus", cwd="/tmp") + turn, thinking_state = TurnState(), ThinkingState() + live_partial, live_thinking = {}, {} + + async def run(): + p_orig = mod.ws_manager + mod.ws_manager = P_WS() + try: + await mod.handle_stream_event(ev({"type": "content_block_start", "index": 0, "content_block": {"type": "thinking"}}), session, "s1", turn, thinking_state, live_partial, live_thinking) + await mod.handle_stream_event(ev({"type": "content_block_delta", "index": 0, "delta": {"type": "thinking_delta", "thinking": "hmm"}}), session, "s1", turn, thinking_state, live_partial, live_thinking) + assert "s1" in live_thinking, "the thought is live while it streams" + # No content_block_stop for index 0 (Haiku), straight to the answer. + await mod.handle_stream_event(ev({"type": "content_block_start", "index": 1, "content_block": {"type": "text"}}), session, "s1", turn, thinking_state, live_partial, live_thinking) + assert "s1" not in live_thinking, "the thought must be gone once the answer begins" + live_thinking["s1"] = PartialReply(msg_id="t9", text="late", branch_id="main") + await mod.handle_stream_event(ev({"type": "message_stop"}), session, "s1", turn, thinking_state, live_partial, live_thinking) + assert "s1" not in live_thinking, "message_stop drops it too" + finally: + mod.ws_manager = p_orig + asyncio.run(run()) + + +def test_every_site_that_drops_the_live_reply_drops_the_live_thought(): + import os, re + root = os.path.join(os.path.dirname(__file__), "..", "apps", "agents") + hits = 0 + for dirpath, _d, files in os.walk(root): + for name in files: + if not name.endswith(".py"): + continue + src = open(os.path.join(dirpath, name)).read() + for m in re.finditer(r"live_partial\.pop\((session\.id|session_id), None\)", src): + hits += 1 + window = src[m.end(): m.end() + 200] + assert "live_thinking.pop(" in window, f"{name}: live_partial popped without live_thinking at offset {m.start()}" + assert hits >= 4, f"expected the four turn-end sites, found {hits}" diff --git a/frontend/src/shared/state/streamSnapshot.test.ts b/frontend/src/shared/state/streamSnapshot.test.ts index 66b97019..cc204335 100644 --- a/frontend/src/shared/state/streamSnapshot.test.ts +++ b/frontend/src/shared/state/streamSnapshot.test.ts @@ -39,3 +39,11 @@ test('the socket handles the snapshot ABOVE the replay-skip guard that drops pre const dashboardSkip = src.indexOf('if (this.skipStreamEvents) {'); assert.ok(snapshot < dashboardSkip, 'the snapshot handler decides skipStreamEvents itself, above the generic skip'); }); + +test('the socket refuses a snapshot for a session it already knows is finished', () => { + const src = fs.readFileSync(path.join(process.cwd(), 'src/shared/ws/WebSocketManager.ts'), 'utf8'); + const i = src.indexOf("event === 'agent:stream_snapshot'"); + const block = src.slice(i, i + 900); + assert.match(block, /agents\.sessions\[session_id\]\?\.status/); + assert.match(block, /!finished/); +}); diff --git a/frontend/src/shared/ws/WebSocketManager.ts b/frontend/src/shared/ws/WebSocketManager.ts index 9f20b435..51766354 100644 --- a/frontend/src/shared/ws/WebSocketManager.ts +++ b/frontend/src/shared/ws/WebSocketManager.ts @@ -495,7 +495,10 @@ class WebSocketManager { // 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) { + // A snapshot is only ever for a turn in flight; a finished session that still gets one is a leak, not a reply. + const snapStatus = session_id ? store.getState().agents.sessions[session_id]?.status : undefined; + const finished = snapStatus !== undefined && snapStatus !== 'running' && snapStatus !== 'waiting_approval'; + if (session_id && data.message_id && typeof data.text === 'string' && !this.skipStreamEvents && !finished) { store.dispatch(streamSnapshot({ sessionId: session_id, messageId: data.message_id, role: data.role, text: data.text })); } return;