From 5629f70c1165a54cfca63f6fc5925383e562c68c Mon Sep 17 00:00:00 2001 From: ciregenz Date: Tue, 23 Jun 2026 04:24:18 -0700 Subject: [PATCH] [eric] agents: extract ResultMessage handler (token/cost/reset) into streaming/result_message + tests --- backend/apps/agents/agent_manager.py | 213 +-------------- .../manager/streaming/result_message.py | 247 ++++++++++++++++++ backend/tests/test_result_message.py | 72 +++++ 3 files changed, 323 insertions(+), 209 deletions(-) create mode 100644 backend/apps/agents/manager/streaming/result_message.py create mode 100644 backend/tests/test_result_message.py diff --git a/backend/apps/agents/agent_manager.py b/backend/apps/agents/agent_manager.py index 2fa930b2..a6d816b5 100644 --- a/backend/apps/agents/agent_manager.py +++ b/backend/apps/agents/agent_manager.py @@ -57,6 +57,7 @@ from backend.apps.agents.manager.streaming import tool_result_hook from backend.apps.agents.manager.streaming import stop_hook as stop_hook_mod from backend.apps.agents.manager.streaming import stream_event from backend.apps.agents.manager.streaming import assistant_message +from backend.apps.agents.manager.streaming import result_message from backend.apps.agents.manager.streaming.upsert_message import upsert_message from backend.apps.agents.manager.prompt.system_prompt import compose_turn_system_prompt from backend.apps.agents.tools.web import should_register_web_mcp @@ -1368,216 +1369,10 @@ class AgentManager(SessionLifecycleMixin, MessagingMixin): message, session, session_id, turn, thinking, self._live_partial, self.sessions ) elif isinstance(message, ResultMessage): - # 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/")) + await result_message.handle_result_message( + message, session, session_id, turn, thinking, self.sessions, + resolved_model, api_type, global_settings, ) - if thinking.text_parts or _route_strips_reasoning: - try: - await thinking_mod.emit_consolidated_thinking( - thinking, turn, session, session_id, self.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.total_chars = 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") capacity_retry_attempt = 0 while True: diff --git a/backend/apps/agents/manager/streaming/result_message.py b/backend/apps/agents/manager/streaming/result_message.py new file mode 100644 index 00000000..daee18c5 --- /dev/null +++ b/backend/apps/agents/manager/streaming/result_message.py @@ -0,0 +1,247 @@ +"""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 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, + 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.total_chars = 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") diff --git a/backend/tests/test_result_message.py b/backend/tests/test_result_message.py new file mode 100644 index 00000000..98763457 --- /dev/null +++ b/backend/tests/test_result_message.py @@ -0,0 +1,72 @@ +"""Direct coverage for the extracted ResultMessage handler: it writes the session's token +totals, recomputes cost off-Anthropic-rate routes (free routes zero out), broadcasts the +context-usage update, and resets the per-turn state. The harness covers the happy path; these +pin the token math, the free-route cost rule, and the reset.""" + +import pytest +from unittest.mock import patch, AsyncMock + +from claude_agent_sdk import ResultMessage + +from backend.apps.agents.core.models import AgentSession +from backend.apps.agents.manager.streaming.state import TurnState, ThinkingState +from backend.apps.agents.manager.streaming import result_message +from backend.apps.settings.settings import load_settings + + +def _result(usage=None, cost=None): + m = ResultMessage(subtype="success", duration_ms=100, duration_api_ms=80, is_error=False, + num_turns=1, session_id="sdk-1", + usage=usage or {"input_tokens": 100, "output_tokens": 50}) + if cost is not None: + try: + m.total_cost_usd = cost + except Exception: + object.__setattr__(m, "total_cost_usd", cost) + return m + + +def _fixt(): + return AgentSession(name="t", model="sonnet", dashboard_id="d"), TurnState(), ThinkingState() + + +@pytest.mark.asyncio +async def test_writes_session_tokens_and_emits_context_update(): + session, turn, thinking = _fixt() + events = [] + + async def fake_send(sid, ev, data): + events.append(ev) + + with patch.object(result_message.ws_manager, "send_to_session", new=fake_send): + await result_message.handle_result_message( + _result(usage={"input_tokens": 100, "output_tokens": 50, "cache_read_input_tokens": 20}), + session, session.id, turn, thinking, {}, "sonnet", "anthropic", load_settings()) + assert session.tokens["input"] == 120 # 100 fresh + 0 create + 20 cache-read + assert session.tokens["input_fresh"] == 100 + assert session.tokens["output"] == 50 + assert "agent:context_update" in events + + +@pytest.mark.asyncio +async def test_free_route_zeroes_cost(): + session, turn, thinking = _fixt() + with patch.object(result_message.ws_manager, "send_to_session", new=AsyncMock()): + await result_message.handle_result_message( + _result(cost=9.99), session, session.id, turn, thinking, {}, "cc/opus", "anthropic", load_settings()) + assert session.cost_usd == 0.0 # cc/ is a subscription (server-funded) route, never billed per-token + + +@pytest.mark.asyncio +async def test_resets_per_turn_state_at_completion(): + session, turn, thinking = _fixt() + turn.output_tokens = 999 + turn.tool_count = 5 + thinking.total_ms = 100 # text_parts left empty so no pill emit fires in the test + with patch.object(result_message.ws_manager, "send_to_session", new=AsyncMock()): + await result_message.handle_result_message( + _result(), session, session.id, turn, thinking, {}, "sonnet", "anthropic", load_settings()) + assert turn.output_tokens == 0 + assert turn.tool_count == 0 + assert thinking.total_ms == 0 + assert thinking.block_starts == {}