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, List, 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

# Annotation-only here (no isinstance dispatch), so the runtime symbol can stay `object` and the SDK chain stays off the boot import graph.
from typing import TYPE_CHECKING

if TYPE_CHECKING:
    from claude_agent_sdk import ResultMessage
else:
    ResultMessage = object

logger = logging.getLogger(__name__)


class TurnResultError(Exception):
    """The CLI's ResultMessage reported the turn ended in an error state (is_error, an
    error_* subtype, or a max_tokens/refusal stop). Raised after the turn's token/cost
    accounting so the run loop's existing error-card path owns the failure instead of the
    turn being consumed as a silent success."""


@typechecked
def p_turn_result_error_text(message: ResultMessage, subtype: str, stop_reason: Optional[str]) -> str:
    parts: List[str] = [str(x).strip() for x in (getattr(message, "errors", None) or []) if str(x).strip()]
    if not parts:
        result_text = getattr(message, "result", None)
        if isinstance(result_text, str) and result_text.strip():
            parts.append(result_text.strip())
    denials = getattr(message, "permission_denials", None) or []
    denied_tools = [str(d.get("tool_name")) for d in denials if isinstance(d, dict) and d.get("tool_name")]
    if denied_tools:
        parts.append("denied tools: " + ", ".join(denied_tools))
    if stop_reason == "max_tokens":
        headline = "The model hit its maximum output length before finishing"
    elif stop_reason == "refusal":
        headline = "The model refused to continue this turn"
    else:
        headline = "The agent runtime reported this turn failed"
    label = subtype if subtype and subtype != "success" else (stop_reason or "unknown")
    detail = "; ".join(parts)
    return f"{headline} ({label})." + (f" {detail}" if detail else "")


@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
        # The result's input usage is summed across every inference step of the turn, which is BILLING; live context is the last step's request size. On a 9-step audit turn the sum read 589K while the real context was 70K, and the meter (plus the compaction trigger) believed it.
        p_ctx_input = turn.last_step_input if turn.last_step_input > 0 else total_input
        session.tokens["input"] = p_ctx_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(p_ctx_input / ctx_window, 4) if p_ctx_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": p_ctx_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")

    # An error-shaped result used to be consumed as a normal end-of-turn: the user got preamble, then silence. Raise AFTER the accounting above so the failure surfaces as a real error card.
    p_subtype = str(getattr(message, "subtype", "") or "")
    p_stop_reason = getattr(message, "stop_reason", None)
    if (
        bool(getattr(message, "is_error", False))
        or p_subtype.startswith("error")
        or p_stop_reason in ("max_tokens", "refusal")
    ):
        raise TurnResultError(p_turn_result_error_text(message, p_subtype, p_stop_reason))