"""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 logger = logging.getLogger(__name__) # The block types drive isinstance DISPATCH, so they must be real at runtime; imported inside the handler because by stream time the SDK is already resident (the turn's presence check imported it), keeping the 350ms sdk+mcp chain off the boot graph. from typing import TYPE_CHECKING if TYPE_CHECKING: from claude_agent_sdk import AssistantMessage else: AssistantMessage = object @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: from claude_agent_sdk.types import ThinkingBlock, TextBlock, ToolUseBlock 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)) ) # One door for everything the provider says, so classify the class here instead of adding a fifth phrasing above; measured 14/14 caught, 0 false positives on 2035 real assistant messages. from backend.apps.agents.manager.streaming.provider_error_speech import ( AUTH as P_ERR_AUTH, POLICY as P_ERR_POLICY, classify_provider_error, is_transient, user_facing_sentence, ) p_provider_error = None if not looks_like_router_auth_error: p_provider_error = classify_provider_error(asst_text) if p_provider_error is not None and p_provider_error.kind == P_ERR_POLICY: # The policy filter answered in the assistant's place. The run-error path owns this failure (recap ratchet, honest card, telemetry); carding "retrying automatically" here while nothing retried was the lie Alex read before every bricked turn. from backend.apps.agents.manager.streaming.handle_result_message import TurnResultError raise TurnResultError(asst_text) if p_provider_error is not None and p_provider_error.kind == P_ERR_AUTH: # Same failure the phrase list was written for, so it goes to the same healer; two mechanisms for one condition is the ENG-252 mistake. looks_like_router_auth_error = True p_provider_error = None if looks_like_router_auth_error: from backend.apps.agents.manager.streaming.auth_retry import try_auth_self_heal p_is_codex = "codex/" in lower_text or "[codex" in lower_text # First expiry in this ask heals silently (fresh CLI + hidden retry); the banner is # reserved for the second failure, when the credential is genuinely dead (ENG-294). # Codex retries wait ~75s so they land AFTER the 1-2 minute rotation window instead of # inside it (an instant retry re-fails and burns the one-shot budget). p_healed = try_auth_self_heal(session, delay_s=75 if p_is_codex else 5) if p_healed and p_is_codex: p_notice = Message( id=uuid4().hex, role="system", content="GPT subscription token just rotated (automatic, every couple minutes). Retrying your request automatically in about a minute, no action needed.", branch_id=session.active_branch_id, ) session.messages.append(p_notice) await ws_manager.send_to_session(session_id, "agent:message", { "session_id": session_id, "message": p_notice.model_dump(mode="json"), }) if not p_healed: if p_is_codex: friendly = ( "GPT subscription token is still refreshing. This usually clears on " "its own; wait a minute and send your message again. If it keeps " "happening, open Settings → Models and click Reconnect on the " "OpenAI / GPT row." ) 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"), }) elif p_provider_error is not None: # The provider failed, the agent never spoke; say what happens next and, when waiting fixes it, do the waiting for them. from backend.apps.agents.manager.streaming.auth_retry import try_transient_self_heal p_delay = 0 p_healed = False # A verdict waiting cannot change ends the ask's recovery ladder. Without this, the # retries kept running after we had already told the user to switch models, and each # one added a card contradicting that advice (packaged drill: seven cards, three # mutually exclusive instructions, twice). if not is_transient(p_provider_error): session.provider_verdict_final = True elif not session.provider_verdict_final: p_delay = min(int(p_provider_error.reset_seconds or 0), 900) p_healed = try_transient_self_heal(session, delay_s=p_delay) if p_healed: p_copy = user_facing_sentence(p_provider_error, session.model or "") else: p_copy = ( user_facing_sentence(p_provider_error, session.model or "") if not is_transient(p_provider_error) else ( "The model provider kept returning a temporary error, so this step could " "not finish. Send your message again, or switch this agent to another " "model." ) ) p_card = Message( id=uuid4().hex, role="system", content=p_copy, branch_id=session.active_branch_id, ) # Dedup by KIND, not by exact string. The identical-card absorber never engaged here # because consecutive cards alternated wording (spent plan / rate limit / spent plan), # so the same underlying failure stacked a wall. One card per kind per ask; a repeat # rewrites it in place so the newest wording wins without adding a row. p_prev_kind = getattr(session, "last_provider_error_kind", "") p_existing = None if p_prev_kind == p_provider_error.kind: for m in reversed(session.messages): if m.role == "system" and not getattr(m, "hidden", False): p_existing = m break if p_existing is not None: p_existing.content = p_copy p_existing.timestamp = p_card.timestamp p_card = p_existing else: from backend.apps.agents.manager.run.handle_run_error import absorb_repeat_card absorb_repeat_card(session, p_card) session.last_provider_error_kind = p_provider_error.kind logger.warning( f"Agent {session_id}: provider returned {p_provider_error.kind} " f"(status={p_provider_error.status}, lane={p_provider_error.lane}) as assistant " f"text; surfaced as a card instead of the agent's own words " f"(retry_queued={p_healed}, delay={p_delay}s)" ) await ws_manager.send_to_session(session_id, "agent:message", { "session_id": session_id, "message": p_card.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 = {}