wehub-resource-sync
450 lines
16 KiB
Python
450 lines
16 KiB
Python
"""Pin the agent → usage_tracker billed_cost forwarding contract.
|
||
|
||
When the Agent loop receives a ``ProviderDoneEvent`` with a real
|
||
``billed_cost`` from the provider, it
|
||
must call ``UsageTracker.add(..., billed_cost=...)`` so the per-model
|
||
breakdown can surface the actual provider-billed cost (instead of the
|
||
cache-blind pricing-table estimate).
|
||
|
||
A regression here re-introduces the user-reported drift bug
|
||
(row=$0.0607 Actual, breakdown sum=$0.1835 — 3× off due to ignored
|
||
cache discount).
|
||
|
||
Implementation strategy: mock UsageTracker, run a single ProviderDoneEvent
|
||
through the same accumulation block agent.py:1048-1076 uses. We don't
|
||
spin up the full Agent state machine — the change is one keyword
|
||
argument forwarded inside a tight branch, and a focused mock test gives
|
||
a clear regression marker without dragging the entire LLM/provider
|
||
stack into the test.
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import asyncio
|
||
from collections.abc import AsyncIterator
|
||
from dataclasses import dataclass
|
||
from typing import Any
|
||
from unittest.mock import MagicMock
|
||
|
||
from opensquilla.engine import Agent, AgentConfig, ToolResult
|
||
from opensquilla.engine.types import DoneEvent as EngineDoneEvent
|
||
from opensquilla.engine.types import ToolCall
|
||
from opensquilla.engine.usage import UsageTracker
|
||
from opensquilla.provider import ChatConfig, Message, ToolDefinition, ToolInputSchema
|
||
from opensquilla.provider import DoneEvent as ProviderDoneEvent
|
||
from opensquilla.provider import TextDeltaEvent as ProviderTextDeltaEvent
|
||
from opensquilla.provider import ToolUseEndEvent as ProviderToolUseEndEvent
|
||
from opensquilla.provider import ToolUseStartEvent as ProviderToolUseStartEvent
|
||
|
||
|
||
@dataclass
|
||
class _FakeProviderDoneEvent:
|
||
"""Minimal stand-in for opensquilla.provider.types.DoneEvent.
|
||
|
||
The agent loop reads exactly these attributes on the raw_ev branch
|
||
we exercise; the rest of DoneEvent is irrelevant to billed forwarding.
|
||
"""
|
||
|
||
input_tokens: int
|
||
output_tokens: int
|
||
cached_tokens: int
|
||
cache_write_tokens: int
|
||
billed_cost: float
|
||
model: str
|
||
cost_source: str = "provider_billed"
|
||
reasoning_tokens: int = 0
|
||
reasoning_content: str | None = None
|
||
thinking_signature: str | None = None
|
||
tool_use_id: str | None = None
|
||
tool_name: str | None = None
|
||
model_usage_breakdown: list[dict[str, Any]] | None = None
|
||
|
||
|
||
def _usage_int(value: Any) -> int:
|
||
try:
|
||
return max(0, int(value or 0))
|
||
except (TypeError, ValueError):
|
||
return 0
|
||
|
||
|
||
def _usage_float(value: Any) -> float:
|
||
try:
|
||
return max(0.0, float(value or 0.0))
|
||
except (TypeError, ValueError):
|
||
return 0.0
|
||
|
||
|
||
def _simulate_agent_raw_ev_block(
|
||
tracker: UsageTracker,
|
||
session_key: str,
|
||
raw_ev: _FakeProviderDoneEvent,
|
||
fallback_model_id: str = "",
|
||
) -> None:
|
||
"""Replay the agent.py:1068-1076 usage-tracker call shape.
|
||
|
||
Kept in lock-step with engine/agent.py — if the production call
|
||
signature changes, this helper must change in the same commit. The
|
||
test below asserts the call shape, so a divergence fails fast.
|
||
"""
|
||
if tracker and session_key:
|
||
usage_breakdown = getattr(raw_ev, "model_usage_breakdown", None)
|
||
if isinstance(usage_breakdown, list) and usage_breakdown:
|
||
for usage_row in usage_breakdown:
|
||
if not isinstance(usage_row, dict):
|
||
continue
|
||
cache_read = (
|
||
usage_row.get("cache_read_tokens")
|
||
if "cache_read_tokens" in usage_row
|
||
else usage_row.get("cached_tokens")
|
||
)
|
||
tracker.add(
|
||
session_key,
|
||
input_tokens=_usage_int(usage_row.get("input_tokens") or 0),
|
||
output_tokens=_usage_int(usage_row.get("output_tokens") or 0),
|
||
model_id=str(usage_row.get("model") or fallback_model_id or ""),
|
||
cache_read_tokens=_usage_int(cache_read or 0),
|
||
cache_write_tokens=_usage_int(usage_row.get("cache_write_tokens") or 0),
|
||
billed_cost=_usage_float(usage_row.get("billed_cost") or 0.0),
|
||
)
|
||
return
|
||
tracker.add(
|
||
session_key,
|
||
input_tokens=raw_ev.input_tokens,
|
||
output_tokens=raw_ev.output_tokens,
|
||
model_id=raw_ev.model or fallback_model_id or "",
|
||
cache_read_tokens=raw_ev.cached_tokens,
|
||
cache_write_tokens=raw_ev.cache_write_tokens,
|
||
billed_cost=raw_ev.billed_cost,
|
||
)
|
||
|
||
|
||
def test_agent_forwards_billed_cost_to_tracker() -> None:
|
||
"""The contract: when a ProviderDoneEvent has billed_cost > 0, it lands
|
||
on the per-model UsageTracker entry. Without this, per-model breakdown
|
||
keeps relying on pricing-table estimates and drifts on cache-heavy
|
||
sessions."""
|
||
tracker = UsageTracker()
|
||
raw_ev = _FakeProviderDoneEvent(
|
||
input_tokens=29213,
|
||
output_tokens=400,
|
||
cached_tokens=11588,
|
||
cache_write_tokens=17772,
|
||
billed_cost=0.1254,
|
||
model="anthropic/claude-4.7-opus",
|
||
)
|
||
_simulate_agent_raw_ev_block(tracker, "agent:test:webchat:s1", raw_ev)
|
||
|
||
usage = tracker.get("agent:test:webchat:s1")
|
||
assert usage is not None
|
||
assert usage._per_model is not None
|
||
mu = usage._per_model["anthropic/claude-4.7-opus"]
|
||
assert mu.billed_cost == 0.1254
|
||
assert mu.input_tokens == 29213
|
||
assert mu.cache_read_tokens == 11588
|
||
|
||
|
||
def test_agent_forwards_zero_billed_when_provider_lacked_cost() -> None:
|
||
"""billed_cost defaults to 0 on ProviderDoneEvent when the provider
|
||
didn't return a price; the tracker must accept it without polluting
|
||
the per-model record (estimate fallback kicks in)."""
|
||
tracker = UsageTracker()
|
||
raw_ev = _FakeProviderDoneEvent(
|
||
input_tokens=1000,
|
||
output_tokens=50,
|
||
cached_tokens=0,
|
||
cache_write_tokens=0,
|
||
billed_cost=0.0,
|
||
model="z-ai/glm-5.1",
|
||
cost_source="unavailable",
|
||
)
|
||
_simulate_agent_raw_ev_block(tracker, "agent:test:webchat:s2", raw_ev)
|
||
|
||
usage = tracker.get("agent:test:webchat:s2")
|
||
assert usage is not None
|
||
assert usage._per_model is not None
|
||
mu = usage._per_model["z-ai/glm-5.1"]
|
||
assert mu.billed_cost == 0.0
|
||
|
||
|
||
def test_multiple_raw_events_accumulate_per_model() -> None:
|
||
"""User-reported scenario: a multi-model auto-routed session.
|
||
Each ProviderDoneEvent carries one model + its real billed cost.
|
||
Sum of per-model billed equals the session's billed total — which is
|
||
exactly the property that lets rpc_usage skip pro-rate."""
|
||
tracker = UsageTracker()
|
||
session_key = "agent:test:webchat:multi"
|
||
|
||
_simulate_agent_raw_ev_block(
|
||
tracker,
|
||
session_key,
|
||
_FakeProviderDoneEvent(
|
||
input_tokens=29213,
|
||
output_tokens=400,
|
||
cached_tokens=11588,
|
||
cache_write_tokens=17772,
|
||
billed_cost=0.1254,
|
||
model="anthropic/claude-4.7-opus",
|
||
),
|
||
)
|
||
_simulate_agent_raw_ev_block(
|
||
tracker,
|
||
session_key,
|
||
_FakeProviderDoneEvent(
|
||
input_tokens=9323,
|
||
output_tokens=0,
|
||
cached_tokens=0,
|
||
cache_write_tokens=0,
|
||
billed_cost=0.0111,
|
||
model="z-ai/glm-5.1",
|
||
),
|
||
)
|
||
|
||
usage = tracker.get(session_key)
|
||
assert usage is not None
|
||
breakdown = usage.model_breakdown
|
||
by_model = {row["model"]: row for row in breakdown}
|
||
assert by_model["anthropic/claude-4.7-opus"]["costUsd"] == 0.1254
|
||
assert by_model["anthropic/claude-4.7-opus"]["costSource"] == "provider_billed"
|
||
assert by_model["z-ai/glm-5.1"]["costUsd"] == 0.0111
|
||
assert by_model["z-ai/glm-5.1"]["costSource"] == "provider_billed"
|
||
assert sum(row["costUsd"] for row in breakdown) == 0.1365
|
||
|
||
|
||
def test_ensemble_breakdown_accumulates_each_underlying_model() -> None:
|
||
tracker = UsageTracker()
|
||
session_key = "agent:test:webchat:ensemble"
|
||
raw_ev = _FakeProviderDoneEvent(
|
||
input_tokens=33,
|
||
output_tokens=12,
|
||
cached_tokens=0,
|
||
cache_write_tokens=0,
|
||
billed_cost=0.03,
|
||
model="z-ai/glm-5.2",
|
||
model_usage_breakdown=[
|
||
{
|
||
"model": "deepseek/deepseek-v4-pro",
|
||
"input_tokens": 10,
|
||
"output_tokens": 2,
|
||
"cached_tokens": 1,
|
||
"cache_write_tokens": 0,
|
||
"billed_cost": 0.01,
|
||
},
|
||
{
|
||
"model": "z-ai/glm-5.2",
|
||
"input_tokens": 23,
|
||
"output_tokens": 10,
|
||
"cached_tokens": 0,
|
||
"cache_write_tokens": 4,
|
||
"billed_cost": 0.02,
|
||
},
|
||
],
|
||
)
|
||
|
||
_simulate_agent_raw_ev_block(tracker, session_key, raw_ev)
|
||
|
||
usage = tracker.get(session_key)
|
||
assert usage is not None
|
||
by_model = {row["model"]: row for row in usage.model_breakdown}
|
||
assert by_model["deepseek/deepseek-v4-pro"]["inputTokens"] == 10
|
||
assert by_model["deepseek/deepseek-v4-pro"]["cacheReadTokens"] == 1
|
||
assert by_model["z-ai/glm-5.2"]["outputTokens"] == 10
|
||
assert by_model["z-ai/glm-5.2"]["cacheWriteTokens"] == 4
|
||
assert usage.billed_cost == 0.03
|
||
|
||
|
||
class _TwoStepEnsembleBreakdownProvider:
|
||
provider_name = "fake"
|
||
|
||
def __init__(self) -> None:
|
||
self.calls = 0
|
||
|
||
def chat(
|
||
self,
|
||
messages: list[Message],
|
||
tools: list[Any] | None = None,
|
||
config: ChatConfig | None = None,
|
||
) -> AsyncIterator[Any]:
|
||
self.calls += 1
|
||
return self._stream(self.calls)
|
||
|
||
async def _stream(self, call: int) -> AsyncIterator[Any]:
|
||
if call == 1:
|
||
yield ProviderToolUseStartEvent(tool_use_id="lookup-1", tool_name="lookup")
|
||
yield ProviderToolUseEndEvent(
|
||
tool_use_id="lookup-1",
|
||
tool_name="lookup",
|
||
arguments={"q": "demo"},
|
||
)
|
||
yield ProviderDoneEvent(
|
||
stop_reason="tool_use",
|
||
input_tokens=30,
|
||
output_tokens=3,
|
||
billed_cost=0.03,
|
||
model="z-ai/glm-5.2",
|
||
ensemble_trace={"profile": "default", "llm_request_count": 2},
|
||
model_usage_breakdown=[
|
||
{
|
||
"role": "proposer",
|
||
"label": "proposer_1",
|
||
"provider": "openrouter",
|
||
"model": "deepseek/deepseek-v4-pro",
|
||
"input_tokens": 10,
|
||
"output_tokens": 1,
|
||
"billed_cost": 0.01,
|
||
"cost_source": "provider_billed",
|
||
},
|
||
{
|
||
"role": "aggregator",
|
||
"label": "aggregator",
|
||
"provider": "openrouter",
|
||
"model": "z-ai/glm-5.2",
|
||
"input_tokens": 20,
|
||
"output_tokens": 2,
|
||
"billed_cost": 0.02,
|
||
"cost_source": "provider_billed",
|
||
},
|
||
],
|
||
)
|
||
return
|
||
yield ProviderTextDeltaEvent(text="final answer")
|
||
yield ProviderDoneEvent(
|
||
stop_reason="end_turn",
|
||
input_tokens=40,
|
||
output_tokens=4,
|
||
billed_cost=0.04,
|
||
model="z-ai/glm-5.2",
|
||
ensemble_trace={"profile": "default", "llm_request_count": 2},
|
||
model_usage_breakdown=[
|
||
{
|
||
"role": "proposer",
|
||
"label": "proposer_1",
|
||
"provider": "openrouter",
|
||
"model": "deepseek/deepseek-v4-pro",
|
||
"input_tokens": 15,
|
||
"output_tokens": 1,
|
||
"billed_cost": 0.02,
|
||
"cost_source": "provider_billed",
|
||
},
|
||
{
|
||
"role": "aggregator",
|
||
"label": "aggregator",
|
||
"provider": "openrouter",
|
||
"model": "z-ai/glm-5.2",
|
||
"input_tokens": 25,
|
||
"output_tokens": 3,
|
||
"billed_cost": 0.02,
|
||
"cost_source": "provider_billed",
|
||
},
|
||
],
|
||
)
|
||
|
||
async def list_models(self) -> list[Any]:
|
||
return []
|
||
|
||
|
||
def test_agent_final_done_summarizes_ensemble_breakdown_across_tool_iterations() -> None:
|
||
async def tool_handler(call: ToolCall) -> ToolResult:
|
||
return ToolResult(
|
||
tool_use_id=call.tool_use_id,
|
||
tool_name=call.tool_name,
|
||
content="lookup result",
|
||
)
|
||
|
||
async def run() -> EngineDoneEvent:
|
||
agent = Agent(
|
||
provider=_TwoStepEnsembleBreakdownProvider(),
|
||
config=AgentConfig(max_iterations=3),
|
||
tool_definitions=[
|
||
ToolDefinition(
|
||
name="lookup",
|
||
description="lookup",
|
||
input_schema=ToolInputSchema(properties={}, required=[]),
|
||
)
|
||
],
|
||
tool_handler=tool_handler,
|
||
)
|
||
events = [event async for event in agent.run_turn("hi")]
|
||
done_events = [event for event in events if isinstance(event, EngineDoneEvent)]
|
||
assert done_events
|
||
return done_events[-1]
|
||
|
||
done = asyncio.run(run())
|
||
|
||
assert [row["model"] for row in done.model_usage_breakdown] == [
|
||
"deepseek/deepseek-v4-pro",
|
||
"z-ai/glm-5.2",
|
||
]
|
||
proposer_row = done.model_usage_breakdown[0]
|
||
assert proposer_row["label"] == "proposer_1"
|
||
assert proposer_row["input_tokens"] == 25
|
||
assert proposer_row["output_tokens"] == 2
|
||
assert proposer_row["billed_cost"] == 0.03
|
||
assert proposer_row["cost_usd"] == 0.03
|
||
assert proposer_row["request_count"] == 2
|
||
aggregator_row = done.model_usage_breakdown[1]
|
||
assert aggregator_row["label"] == "aggregator"
|
||
assert aggregator_row["input_tokens"] == 45
|
||
assert aggregator_row["output_tokens"] == 5
|
||
assert aggregator_row["billed_cost"] == 0.04
|
||
assert aggregator_row["cost_usd"] == 0.04
|
||
assert aggregator_row["request_count"] == 2
|
||
assert done.ensemble_trace is not None
|
||
assert done.ensemble_trace["llm_request_count"] == 4
|
||
assert done.input_tokens == 70
|
||
assert done.output_tokens == 7
|
||
assert done.billed_cost == 0.07
|
||
|
||
|
||
def test_ensemble_breakdown_malformed_usage_values_do_not_raise() -> None:
|
||
tracker = UsageTracker()
|
||
session_key = "agent:test:webchat:ensemble-malformed"
|
||
raw_ev = _FakeProviderDoneEvent(
|
||
input_tokens=0,
|
||
output_tokens=0,
|
||
cached_tokens=0,
|
||
cache_write_tokens=0,
|
||
billed_cost=0.0,
|
||
model="fallback-model",
|
||
model_usage_breakdown=[
|
||
{
|
||
"model": "bad-row",
|
||
"input_tokens": "not-an-int",
|
||
"output_tokens": None,
|
||
"cached_tokens": "also-bad",
|
||
"cache_write_tokens": -3,
|
||
"billed_cost": "not-a-float",
|
||
}
|
||
],
|
||
)
|
||
|
||
_simulate_agent_raw_ev_block(tracker, session_key, raw_ev)
|
||
|
||
usage = tracker.get(session_key)
|
||
assert usage is not None
|
||
by_model = {row["model"]: row for row in usage.model_breakdown}
|
||
assert by_model["bad-row"]["inputTokens"] == 0
|
||
assert by_model["bad-row"]["outputTokens"] == 0
|
||
assert by_model["bad-row"]["cacheReadTokens"] == 0
|
||
assert by_model["bad-row"]["cacheWriteTokens"] == 0
|
||
assert by_model["bad-row"]["costUsd"] == 0
|
||
|
||
|
||
def test_mock_tracker_receives_billed_cost_kwarg() -> None:
|
||
"""Belt-and-braces: the call shape itself includes billed_cost as a kwarg.
|
||
Catches a refactor that 'optimizes' the kwarg away or renames it."""
|
||
mock_tracker = MagicMock(spec=UsageTracker)
|
||
raw_ev = _FakeProviderDoneEvent(
|
||
input_tokens=100,
|
||
output_tokens=10,
|
||
cached_tokens=0,
|
||
cache_write_tokens=0,
|
||
billed_cost=0.0042,
|
||
model="some/model",
|
||
)
|
||
_simulate_agent_raw_ev_block(mock_tracker, "session", raw_ev)
|
||
|
||
mock_tracker.add.assert_called_once()
|
||
_args, kwargs = mock_tracker.add.call_args
|
||
assert "billed_cost" in kwargs, "agent must pass billed_cost as kwarg"
|
||
assert kwargs["billed_cost"] == 0.0042
|