"""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 import logging 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 from backend.apps.agents.manager.context_budget import maybe_break_midturn if maybe_break_midturn(session, turn, msg_usage): logging.getLogger(__name__).warning( f"[context-break] session {session_id}: mid-turn input " f"{session.tokens.get('input')} crossed the compact trigger; breaking at the " "next message boundary and continuing on a fresh compacted session" ) try: from backend.apps.service.client import submit_diagnostic submit_diagnostic({ "kind": "context_midturn_break", "session_id": session_id, "model": session.model, "input_tokens": session.tokens.get("input"), "context_window": session.context_window, }) except Exception: pass 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 = {}