"""Handle the SDK ResultMessage that closes a turn: fold in authoritative output tokens, write the session's token + cost totals (recomputing cost off-Anthropic-rate routes), emit the final consolidated thinking pill, broadcast the context-usage update, and reset the per-turn TurnState / ThinkingState. Lifted out of the agent loop; mutates the passed state by reference exactly as inline. resolved_model / api_type / global_settings are the loop's per-run config, threaded in.""" import asyncio import logging from typing import Dict, Optional 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 import thinking as thinking_mod try: from claude_agent_sdk import ResultMessage except ImportError: # the SDK is optional at runtime (mock mode); keep this module importable ResultMessage = object # type: ignore logger = logging.getLogger(__name__) @typechecked async def handle_result_message( message: ResultMessage, session: AgentSession, session_id: str, turn: TurnState, thinking: ThinkingState, sessions: Dict[str, AgentSession], resolved_model: object, api_type: Optional[str], global_settings: object, ) -> None: # ResultMessage carries the AUTHORITATIVE per-turn output_tokens count. Some providers (notably OpenAI/Gemini through 9Router) only populate `usage.output_tokens` here, not on individual AssistantMessages. Fold this into the running turn aggregate BEFORE emitting the final consolidated thinking message, so the bubble's tokens segment reflects ground truth on those providers too. try: result_usage = getattr(message, "usage", None) or {} if isinstance(result_usage, dict): result_out = int(result_usage.get("output_tokens", 0) or 0) # Take the max, if individual AssistantMessages already summed to a larger number we trust that; otherwise ResultMessage's count fills the gap. if result_out > turn.output_tokens: turn.output_tokens = result_out except Exception: pass # Pre-populate session.tokens BEFORE emitting the final consolidated thinking pill. Order matters: emit_consolidated_thinking reads session.tokens["input"]/["output"] for the combined-total stamp on the pill. If we emit first, the pill freezes with input=0 because the ResultMessage hasn't been consumed yet (the writes below at line ~2918 wouldn't land until after the pill is already broadcast). try: pre_usage = getattr(message, "usage", None) or {} if isinstance(pre_usage, dict): pre_in = int(pre_usage.get("input_tokens", 0) or 0) pre_create = int(pre_usage.get("cache_creation_input_tokens", 0) or 0) pre_read = int(pre_usage.get("cache_read_input_tokens", 0) or 0) pre_total_in = pre_in + pre_create + pre_read pre_out = int(pre_usage.get("output_tokens", 0) or 0) if pre_total_in > 0: session.tokens["input"] = pre_total_in # Pill reads the fresh lane: uncached input only, so re-read/cached context doesn't inflate it. session.tokens["input_fresh"] = pre_in if pre_out > 0: session.tokens["output"] = pre_out except Exception: pass # Final consolidated emission with the full duration + authoritative tokens. The frontend bubble freezes on this final value. For routes whose translator strips reasoning content (cx/ for OpenAI, gc/ for Gemini), force-emit a pill even when no text or upstream token count was captured. Without this, GPT/ Gemini turns show no thinking bubble at all because 9Router's translator doesn't carry reasoning_content across the Anthropic-shape round-trip. The frontend's ThinkingBubble detects empty content and renders a friendly "provider doesn't expose reasoning text" explanation instead of a blank panel. route_strips_reasoning = ( isinstance(resolved_model, str) and resolved_model.startswith(("cx/", "gc/", "ag/", "gemini/")) ) if thinking.text_parts or route_strips_reasoning: try: await thinking_mod.emit_consolidated_thinking( thinking, turn, session, session_id, sessions, force_provider_unavailable=route_strips_reasoning, ) except Exception: pass if thinking.ticker_task is not None and not thinking.ticker_task.done(): thinking.ticker_task.cancel() try: await thinking.ticker_task except (asyncio.CancelledError, Exception): pass thinking.ticker_task = None thinking.msg_id = None thinking.text_parts = [] turn.tool_count = 0 turn.started_ts = None turn.total_ms = 0 turn.output_tokens = 0 turn.assistant_text_chars = 0 turn.tool_input_chars = 0 thinking.thought_signature = None turn.baseline_session_in = 0 turn.baseline_session_out = 0 turn.baseline_children_in = 0 turn.baseline_children_out = 0 turn.baseline_captured = False thinking.total_ms = 0 thinking.block_starts = {} session.sdk_session_id = getattr(message, "session_id", None) # Pull usage first; SDK's total_cost_usd is wrong for OR (assumes Anthropic rates) and we recompute below. usage = getattr(message, "usage", None) or {} inp = out = cache_create = cache_read = total_input = 0 if isinstance(usage, dict): inp = usage.get("input_tokens", 0) or 0 out = usage.get("output_tokens", 0) or 0 cache_create = usage.get("cache_creation_input_tokens", 0) or 0 cache_read = usage.get("cache_read_input_tokens", 0) or 0 total_input = inp + cache_create + cache_read session.tokens["input"] = total_input session.tokens["input_fresh"] = inp session.tokens["output"] = out cost = getattr(message, "total_cost_usd", None) if cost is not None: free_route = False if isinstance(resolved_model, str): if resolved_model.startswith(("cc/", "cx/", "gc/", "ag/")): free_route = True elif resolved_model.startswith("openrouter/") and ":free" in resolved_model: free_route = True elif resolved_model.startswith("cp-"): # User-configured custom OpenAI-compatible provider (Ollama Cloud, Together, Groq, local LMs, etc.). Pricing is unknowable without per-provider rate tables that would rot fast, zero out instead of showing the SDK's Anthropic-rate estimate, which is meaningless here. free_route = True if api_type == "anthropic": from backend.apps.settings.credentials import proxy_auth as proxy_auth pa_tok, _ = proxy_auth(global_settings) # Pro and free-trial both run server-funded, so per-token cost to the user is 0. if pa_tok: free_route = True if free_route: cost = 0.0 elif isinstance(resolved_model, str) and resolved_model.startswith("openrouter/"): # SDK assumes Anthropic rates → 50-100× off for OR. from backend.apps.agents.providers.registry import get_openrouter_pricing pricing = get_openrouter_pricing(resolved_model) if pricing: in_rate, out_rate = pricing cost = ( (inp + cache_create + cache_read) * in_rate + out * out_rate ) / 1_000_000 elif api_type in ("openai", "gemini") or ( isinstance(resolved_model, str) and (resolved_model.startswith("cp-openai/") or resolved_model.startswith("cp-gemini/") or resolved_model.startswith("cp-google/")) ): # Direct OpenAI/Gemini API key lane. SDK's total_cost_usd is computed at Anthropic rates (Opus pricing), for GPT-5.4-Mini at $0.25/M input that's a 60x overcount ($30 instead of $0.04 per Mehmet-style 4-PDF turn). Use the published per-model rates instead. from backend.apps.agents.providers.registry import get_direct_pricing pricing = get_direct_pricing(resolved_model) or get_direct_pricing(session.model) if pricing: in_rate, out_rate = pricing cost = ( (inp + cache_create + cache_read) * in_rate + out * out_rate ) / 1_000_000 else: # Unknown model in this family: zero out rather than ship an Anthropic-rate estimate that's wildly wrong. cost = 0.0 session.cost_usd = cost await ws_manager.send_to_session(session_id, "agent:cost_update", { "session_id": session_id, "cost_usd": session.cost_usd, }) if isinstance(usage, dict): # Per-turn context-usage broadcast. Drives the UI status pill and the auto-compact threshold. The denominator is the session's real model cap, populated from registry.get_context_window at session creation, restore, and model-switch (see apply_context_window). max(1, ...) is a belt-and-braces guard against zero/None drift from any future restore-from-disk corner case. ctx_window = max(1, getattr(session, "context_window", 0) or 200_000) ctx_used_pct = round(total_input / ctx_window, 4) if total_input else 0.0 cache_read_pct = round(cache_read / total_input, 4) if total_input else 0.0 try: await ws_manager.send_to_session(session_id, "agent:context_update", { "session_id": session_id, "input_tokens": total_input, "output_tokens": out, "cache_read_tokens": cache_read, "cache_read_pct": cache_read_pct, "ctx_used_pct": ctx_used_pct, "context_window": ctx_window, "framework_overhead_tokens": session.framework_overhead_tokens, "active_mcps": list(session.active_mcps), }) except Exception: logger.exception("Failed to emit agent:context_update")