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

"""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")