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

"""Handle a complete AssistantMessage envelope from the SDK: split its blocks into thinking /
text / tool-use, fold the thinking into the consolidated pill, surface a friendly card for a
router auth-expiry that arrived as assistant text, and commit the assistant + tool-call messages.
Lifted out of the agent loop; mutates the passed TurnState / ThinkingState by reference and writes
through the manager's live-partial mirror + session registry, exactly as it did inline."""

import asyncio
from typing import Dict, Optional
from uuid import uuid4

from typeguard import typechecked

from backend.apps.agents.core.models import AgentSession, Message
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.upsert_message import upsert_message
from backend.apps.agents.manager.streaming.PartialReply import PartialReply
from backend.apps.agents.manager.streaming import thinking as thinking_mod

try:
    from claude_agent_sdk import AssistantMessage
    from claude_agent_sdk.types import ThinkingBlock, TextBlock, ToolUseBlock
except ImportError:  # the SDK is optional at runtime (mock mode); keep this module importable
    AssistantMessage = ThinkingBlock = TextBlock = ToolUseBlock = object  # type: ignore


@typechecked
async def handle_assistant_message(
    message: AssistantMessage,
    session: AgentSession,
    session_id: str,
    turn: TurnState,
    thinking: ThinkingState,
    live_partial: Dict[str, PartialReply],
    sessions: Dict[str, AgentSession],
) -> None:
    content_parts = []
    new_thinking_parts = []
    tool_uses = []
    # Capture the latest Gemini thoughtSignature (and Anthropic's signature_delta if present) off any ThinkingBlock in this message. We store it on the turn's consolidated thinking message so it survives session.json serialization, and re-attach it on the next request so Google's continuity check passes.
    new_thought_signature: Optional[str] = None
    for block in message.content:
        if isinstance(block, ThinkingBlock):
            thinking_text = getattr(block, "thinking", None) or getattr(block, "text", None) or ""
            if thinking_text:
                new_thinking_parts.append(thinking_text)
            # Try multiple field-name variants, SDK versions and 9Router translations have used `signature`, `thoughtSignature`, and `thought_signature` over time.
            sig = (
                getattr(block, "signature", None)
                or getattr(block, "thoughtSignature", None)
                or getattr(block, "thought_signature", None)
            )
            if sig:
                new_thought_signature = sig
        elif isinstance(block, TextBlock):
            content_parts.append(block.text)
        elif isinstance(block, ToolUseBlock):
            tool_uses.append({
                "id": block.id,
                "tool": block.name,
                "input": block.input,
            })

    # Accumulate this AssistantMessage's contributions into the turn-level thinking pill. We re-emit the SAME message id each time so the frontend dedupes (addMessage replaces by id) and the bubble updates live as more thought / tools arrive. This is what gives us "Thought for 18s · 412 tokens · 3 tools used" reflecting the whole turn rather than just one think-step. NOTE: tool count is incremented in the content_block_start (block_type=="tool_use") branch above, NOT here. That path fires for both Anthropic and 9Router-translated providers; counting again here would double. If a provider somehow doesn't surface content_block_start for tool blocks but DOES surface them in the AssistantMessage envelope (defensive case), the max() in the consolidated emit will still pick up the higher count.
    if new_thinking_parts:
        thinking.text_parts.extend(new_thinking_parts)
    # Latch the most recent thoughtSignature, Gemini only validates against the LATEST one in the conversation history, so older signatures from earlier think-steps in the same turn are superseded by newer ones.
    if new_thought_signature:
        thinking.thought_signature = new_thought_signature
    # Accumulate this message's total output tokens (SDK populates `usage.output_tokens` with the full output for the inference: thinking text + visible text + tool-call JSON args). Summing across the turn's AssistantMessages gives us "all output the model produced this turn," which is what users intuit when they see a token count.
    try:
        msg_usage = getattr(message, "usage", None) or {}
        if isinstance(msg_usage, dict):
            ot = int(msg_usage.get("output_tokens", 0) or 0)
            if ot > 0:
                turn.output_tokens += ot
    except Exception:
        pass

    # Re-emit the consolidated thinking message on every AssistantMessage (event-driven). The background ticker loop keeps it updating between events too, so the elapsed counter ticks even during tool execution / slow text generation gaps.
    if thinking.text_parts:
        await thinking_mod.emit_consolidated_thinking(thinking, turn, session, session_id, sessions)
        # Start the 1Hz ticker once we have a consolidated message in flight so the bubble keeps updating between SDK events.
        if thinking.ticker_task is None or thinking.ticker_task.done():
            thinking.ticker_task = asyncio.create_task(thinking_mod.ticker_loop(thinking, turn, session, session_id, sessions))

    if content_parts:
        asst_text = "\n".join(content_parts)
        # 9Router sometimes returns upstream 401s as the assistant reply (no SDK exception), so the catch-all auth handler never fires. Match the text pattern and surface a friendly system bubble instead.
        lower_text = asst_text.lower()
        looks_like_router_auth_error = (
            ("failed to authenticate" in lower_text and "401" in lower_text)
            or ("authentication token is expired" in lower_text)
            or ("authentication token has expired" in lower_text)
            or ("provided authentication token" in lower_text and ("401" in lower_text or "expired" in lower_text))
        )
        if looks_like_router_auth_error:
            if "codex/" in lower_text or "[codex" in lower_text:
                friendly = (
                    "GPT subscription token expired. Open Settings → Models and click "
                    "Reconnect on the OpenAI / GPT row to refresh, should take ~10s, "
                    "then send your message again."
                )
                reason = "codex_token_expired"
            elif "gemini-cli/" in lower_text or "[gemini" in lower_text:
                friendly = (
                    "Gemini subscription token expired. Open Settings → Models and click "
                    "Reconnect on the Google / Gemini row, then send your message again."
                )
                reason = "gemini_token_expired"
            else:
                friendly = (
                    "Provider authentication expired. Open Settings → Models and "
                    "reconnect, then send your message again."
                )
                reason = "router_auth_expired"
            err_msg = Message(
                id=uuid4().hex,
                role="system",
                content=friendly,
                branch_id=session.active_branch_id,
            )
            session.messages.append(err_msg)
            await ws_manager.send_to_session(session_id, "agent:auth_error", {
                "session_id": session_id,
                "reason": reason,
                "message": friendly,
                "model": session.model,
            })
            await ws_manager.send_to_session(session_id, "agent:message", {
                "session_id": session_id,
                "message": err_msg.model_dump(mode="json"),
            })
        else:
            asst_msg = Message(
                id=turn.stream_text_msg_id or uuid4().hex,
                role="assistant",
                content=asst_text,
                branch_id=session.active_branch_id,
            )
            upsert_message(session, asst_msg)
            turn.stream_text_accum = ""
            live_partial.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"),
            })

    for i, tu in enumerate(tool_uses):
        msg_id = turn.stream_tool_msg_ids_ordered[i] if i < len(turn.stream_tool_msg_ids_ordered) else uuid4().hex
        tool_msg = Message(id=msg_id, role="tool_call", content=tu, branch_id=session.active_branch_id)
        upsert_message(session, tool_msg)
        await ws_manager.send_to_session(session_id, "agent:message", {
            "session_id": session_id,
            "message": tool_msg.model_dump(mode="json"),
        })

    turn.number += 1

    turn.stream_text_msg_id = None
    turn.stream_tool_msg_ids_ordered = []
    turn.stream_block_index_map = {}