mirror of
https://github.com/openswarm-ai/openswarm.git
synced 2026-09-26 19:44:51 +02:00
[eric] agents: extract ResultMessage handler (token/cost/reset) into streaming/result_message + tests
This commit is contained in:
ciregenz
@@ -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:
|
||||
|
||||
@@ -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")
|
||||
@@ -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 == {}
|
||||
Reference in New Issue
Block a user