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