From 2f0c101d6a9e75be0d076aa79cf7ed8c7c41398e Mon Sep 17 00:00:00 2001 From: ciregenz Date: Tue, 23 Jun 2026 03:54:45 -0700 Subject: [PATCH] [eric] agents: extract StreamEvent (incremental stream) handler into streaming/stream_event + tests --- backend/apps/agents/agent_manager.py | 141 +-------------- .../agents/manager/streaming/stream_event.py | 167 ++++++++++++++++++ backend/tests/test_stream_event.py | 73 ++++++++ 3 files changed, 244 insertions(+), 137 deletions(-) create mode 100644 backend/apps/agents/manager/streaming/stream_event.py create mode 100644 backend/tests/test_stream_event.py diff --git a/backend/apps/agents/agent_manager.py b/backend/apps/agents/agent_manager.py index a4a72294..7a562c18 100644 --- a/backend/apps/agents/agent_manager.py +++ b/backend/apps/agents/agent_manager.py @@ -55,6 +55,7 @@ from backend.apps.agents.manager.streaming.hook_context import HookContext from backend.apps.agents.manager.streaming import thinking as thinking_mod from backend.apps.agents.manager.streaming import tool_result_hook from backend.apps.agents.manager.streaming import stop_hook as stop_hook_mod +from backend.apps.agents.manager.streaming import stream_event from backend.apps.agents.manager.prompt.system_prompt import compose_turn_system_prompt from backend.apps.agents.tools.web import should_register_web_mcp from backend.apps.agents.manager.session.SessionLifecycleMixin import SessionLifecycleMixin @@ -1356,143 +1357,9 @@ class AgentManager(SessionLifecycleMixin, MessagingMixin): logger.info(f"[MCP-DEBUG] SystemMessage: {raw}") if isinstance(message, StreamEvent): - event = message.event - event_type = event.get("type") - - if event_type == "content_block_start": - # Stamp the first stream event of the session - # so the session list can show "first response - # at HH:MM" on reload. Only the first turn - # sets this; later turns leave it untouched. - if session.first_response_at is None: - session.first_response_at = datetime.now() - - block = event.get("content_block", {}) - index = event.get("index") - block_type = block.get("type") - - if block_type == "text": - 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", { - "session_id": session_id, - "message_id": turn.stream_text_msg_id, - "role": "assistant", - }) - turn.stream_block_index_map[index] = turn.stream_text_msg_id - - elif block_type == "thinking": - # 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 - # 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", { - "session_id": session_id, - "message_id": thinking_msg_id, - "role": "thinking", - }) - - elif block_type == "tool_use": - tool_msg_id = uuid4().hex - turn.stream_tool_msg_ids_ordered.append(tool_msg_id) - turn.stream_block_index_map[index] = tool_msg_id - # Stream-level tool count for the - # consolidated thinking pill. The - # AssistantMessage path (further down) - # ALSO increments turn.tool_count when - # ToolUseBlocks fully arrive, but for - # OpenAI/Gemini through 9Router the - # AssistantMessage envelope is sometimes - # incomplete, so this stream-level count - # is what guarantees the "N tools used" - # segment renders cross-provider. To - # avoid double-counting we DON'T also - # increment on AssistantMessage when - # this code path already fired, see - # the dedupe at the AssistantMessage - # block below. - turn.tool_count += 1 - await ws_manager.send_to_session(session_id, "agent:stream_start", { - "session_id": session_id, - "message_id": tool_msg_id, - "role": "tool_call", - "tool_name": block.get("name", ""), - }) - - elif event_type == "content_block_delta": - index = event.get("index") - delta = event.get("delta", {}) - delta_type = delta.get("type") - msg_id = turn.stream_block_index_map.get(index) - - if msg_id and delta_type == "text_delta": - _text_chunk = delta.get("text", "") - turn.assistant_text_chars += len(_text_chunk) - turn.stream_text_accum += _text_chunk - self._live_partial[session_id] = { - "msg_id": turn.stream_text_msg_id, - "text": turn.stream_text_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, - "delta": _text_chunk, - }) - 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", "") - thinking.total_chars += len(_think_chunk) - await ws_manager.send_to_session(session_id, "agent:stream_delta", { - "session_id": session_id, - "message_id": msg_id, - "delta": _think_chunk, - }) - elif msg_id and delta_type == "input_json_delta": - _json_chunk = delta.get("partial_json", "") - turn.tool_input_chars += len(_json_chunk) - await ws_manager.send_to_session(session_id, "agent:stream_delta", { - "session_id": session_id, - "message_id": msg_id, - "delta": _json_chunk, - }) - - elif event_type == "content_block_stop": - index = event.get("index") - msg_id = turn.stream_block_index_map.get(index) - # 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( - (time.time() - thinking.block_starts.pop(index)) * 1000 - ) - if msg_id and msg_id != turn.stream_text_msg_id: - await ws_manager.send_to_session(session_id, "agent:stream_end", { - "session_id": session_id, - "message_id": msg_id, - }) - - elif event_type == "message_stop": - if turn.stream_text_msg_id: - await ws_manager.send_to_session(session_id, "agent:stream_end", { - "session_id": session_id, - "message_id": turn.stream_text_msg_id, - }) + await stream_event.handle_stream_event( + message, session, session_id, turn, thinking, self._live_partial + ) elif isinstance(message, AssistantMessage): content_parts = [] diff --git a/backend/apps/agents/manager/streaming/stream_event.py b/backend/apps/agents/manager/streaming/stream_event.py new file mode 100644 index 00000000..144b56f9 --- /dev/null +++ b/backend/apps/agents/manager/streaming/stream_event.py @@ -0,0 +1,167 @@ +"""Handle one streaming StreamEvent from the SDK: the incremental content_block_start / +delta / stop / message_stop path that drives live text, thinking, and tool streaming to the UI. +Lifted out of the agent loop; mutates the passed TurnState / ThinkingState by reference and +writes the manager's live-partial mirror, exactly as it did inline.""" + +import time +from datetime import datetime +from uuid import uuid4 + +from typeguard import typechecked + +from backend.apps.agents.core.models import AgentSession +from backend.apps.agents.core.ws_manager import ws_manager +from backend.apps.agents.manager.streaming.state import ThinkingState, TurnState + +try: + from claude_agent_sdk.types import StreamEvent +except ImportError: # the SDK is optional at runtime (mock mode); keep this module importable + StreamEvent = object # type: ignore + + +@typechecked +async def handle_stream_event( + message: StreamEvent, + session: AgentSession, + session_id: str, + turn: TurnState, + thinking: ThinkingState, + live_partial: dict, +) -> None: + event = message.event + event_type = event.get("type") + + if event_type == "content_block_start": + # Stamp the first stream event of the session + # so the session list can show "first response + # at HH:MM" on reload. Only the first turn + # sets this; later turns leave it untouched. + if session.first_response_at is None: + session.first_response_at = datetime.now() + + block = event.get("content_block", {}) + index = event.get("index") + block_type = block.get("type") + + if block_type == "text": + 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", { + "session_id": session_id, + "message_id": turn.stream_text_msg_id, + "role": "assistant", + }) + turn.stream_block_index_map[index] = turn.stream_text_msg_id + + elif block_type == "thinking": + # 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 + # 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", { + "session_id": session_id, + "message_id": thinking_msg_id, + "role": "thinking", + }) + + elif block_type == "tool_use": + tool_msg_id = uuid4().hex + turn.stream_tool_msg_ids_ordered.append(tool_msg_id) + turn.stream_block_index_map[index] = tool_msg_id + # Stream-level tool count for the + # consolidated thinking pill. The + # AssistantMessage path (further down) + # ALSO increments turn.tool_count when + # ToolUseBlocks fully arrive, but for + # OpenAI/Gemini through 9Router the + # AssistantMessage envelope is sometimes + # incomplete, so this stream-level count + # is what guarantees the "N tools used" + # segment renders cross-provider. To + # avoid double-counting we DON'T also + # increment on AssistantMessage when + # this code path already fired, see + # the dedupe at the AssistantMessage + # block below. + turn.tool_count += 1 + await ws_manager.send_to_session(session_id, "agent:stream_start", { + "session_id": session_id, + "message_id": tool_msg_id, + "role": "tool_call", + "tool_name": block.get("name", ""), + }) + + elif event_type == "content_block_delta": + index = event.get("index") + delta = event.get("delta", {}) + delta_type = delta.get("type") + msg_id = turn.stream_block_index_map.get(index) + + if msg_id and delta_type == "text_delta": + text_chunk = delta.get("text", "") + turn.assistant_text_chars += len(text_chunk) + turn.stream_text_accum += text_chunk + live_partial[session_id] = { + "msg_id": turn.stream_text_msg_id, + "text": turn.stream_text_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, + "delta": text_chunk, + }) + 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", "") + thinking.total_chars += len(think_chunk) + await ws_manager.send_to_session(session_id, "agent:stream_delta", { + "session_id": session_id, + "message_id": msg_id, + "delta": think_chunk, + }) + elif msg_id and delta_type == "input_json_delta": + json_chunk = delta.get("partial_json", "") + turn.tool_input_chars += len(json_chunk) + await ws_manager.send_to_session(session_id, "agent:stream_delta", { + "session_id": session_id, + "message_id": msg_id, + "delta": json_chunk, + }) + + elif event_type == "content_block_stop": + index = event.get("index") + msg_id = turn.stream_block_index_map.get(index) + # 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( + (time.time() - thinking.block_starts.pop(index)) * 1000 + ) + if msg_id and msg_id != turn.stream_text_msg_id: + await ws_manager.send_to_session(session_id, "agent:stream_end", { + "session_id": session_id, + "message_id": msg_id, + }) + + elif event_type == "message_stop": + if turn.stream_text_msg_id: + await ws_manager.send_to_session(session_id, "agent:stream_end", { + "session_id": session_id, + "message_id": turn.stream_text_msg_id, + }) diff --git a/backend/tests/test_stream_event.py b/backend/tests/test_stream_event.py new file mode 100644 index 00000000..93737c4b --- /dev/null +++ b/backend/tests/test_stream_event.py @@ -0,0 +1,73 @@ +"""Direct coverage for the extracted StreamEvent handler. The streaming harness yields whole +AssistantMessage/ResultMessage envelopes, never partial StreamEvents, so this is the only test +that drives the incremental content_block_start/delta/stop path the live UI actually uses.""" + +import pytest +from unittest.mock import patch, AsyncMock + +from claude_agent_sdk.types import StreamEvent + +from backend.apps.agents.core.models import AgentSession +from backend.apps.agents.manager.streaming.state import TurnState, ThinkingState +from backend.apps.agents.manager.streaming import stream_event + + +def _ev(event: dict) -> StreamEvent: + return StreamEvent(uuid="u", session_id="s", event=event) + + +def _fixt(): + session = AgentSession(name="t", model="sonnet", dashboard_id="d") + return session, TurnState(), ThinkingState(), {} + + +@pytest.mark.asyncio +async def test_content_block_start_text_inits_stream_message(): + session, turn, thinking, lp = _fixt() + with patch.object(stream_event.ws_manager, "send_to_session", new=AsyncMock()) as send: + await stream_event.handle_stream_event( + _ev({"type": "content_block_start", "index": 0, "content_block": {"type": "text"}}), + session, session.id, turn, thinking, lp) + assert turn.stream_text_msg_id is not None + assert turn.stream_block_index_map[0] == turn.stream_text_msg_id + send.assert_awaited() # agent:stream_start broadcast + + +@pytest.mark.asyncio +async def test_text_delta_accumulates_and_mirrors_live_partial(): + session, turn, thinking, lp = _fixt() + turn.stream_text_msg_id = "m1" + turn.stream_block_index_map[0] = "m1" + with patch.object(stream_event.ws_manager, "send_to_session", new=AsyncMock()): + await stream_event.handle_stream_event( + _ev({"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "Hello"}}), + session, session.id, turn, thinking, lp) + assert turn.stream_text_accum == "Hello" + assert turn.assistant_text_chars == 5 + assert lp[session.id]["text"] == "Hello" # the live-partial mirror the manager reads on resume + + +@pytest.mark.asyncio +async def test_tool_use_start_increments_stream_tool_count(): + session, turn, thinking, lp = _fixt() + with patch.object(stream_event.ws_manager, "send_to_session", new=AsyncMock()): + await stream_event.handle_stream_event( + _ev({"type": "content_block_start", "index": 1, "content_block": {"type": "tool_use", "name": "Read"}}), + session, session.id, turn, thinking, lp) + assert turn.tool_count == 1 + assert 1 in turn.stream_block_index_map + + +@pytest.mark.asyncio +async def test_thinking_block_start_then_stop_pops_and_accumulates(): + session, turn, thinking, lp = _fixt() + with patch.object(stream_event.ws_manager, "send_to_session", new=AsyncMock()): + await stream_event.handle_stream_event( + _ev({"type": "content_block_start", "index": 2, "content_block": {"type": "thinking"}}), + session, session.id, turn, thinking, lp) + assert 2 in thinking.block_starts + await stream_event.handle_stream_event( + _ev({"type": "content_block_stop", "index": 2}), + session, session.id, turn, thinking, lp) + assert 2 not in thinking.block_starts # the start was popped on stop + assert thinking.total_ms >= 0 # elapsed accumulated server-side