openswarm/backend/apps/agents/manager/streaming/handle_stream_event.py

"""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 typing import Dict
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
from backend.apps.agents.manager.streaming.PartialReply import PartialReply

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[str, PartialReply],
) -> 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] = PartialReply(
                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", "")
            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,
            })