mirror of
https://github.com/openswarm-ai/openswarm.git
synced 2026-09-07 02:07:45 +02:00
[eric] agents: extract StreamEvent (incremental stream) handler into streaming/stream_event + tests
This commit is contained in:
ciregenz
@@ -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 = []
|
||||
|
||||
@@ -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,
|
||||
})
|
||||
@@ -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
|
||||
Reference in New Issue
Block a user