External/roboco
1
0
Fork 0
mirror of https://github.com/rennf93/roboco.git synced 2026-08-03 07:23:24 +02:00
Code Issues Packages Projects Releases Wiki Activity

[499f9eb1] Token Usage & Cost Analytics — Full-Stack Instrumentation, Persistence, and Visualization (#90)

Browse Source 
* [cd2bf666] feat(usage): add token usage types, API client, hooks, and UI components (#87) (#88)

- Append 5 TypeScript interfaces to src/types/index.ts: TokenUsageSnapshot, AgentUsageRow, UsageSession, UsageTimePoint, ModelUsageSlice
- Create src/lib/api/usage.ts: Axios singleton + isMockMode guards for getUsageSnapshot, getUsageTimeSeries, getAgentUsage, getUsageSessions, getModelUsage
- Create src/hooks/use-usage.ts: usageKeys factory + useUsageSnapshot, useUsageTimeSeries, useAgentUsage, useUsageSessions, useModelUsage hooks
- Create UsageOverviewPanel (dashboard/usage-overview-panel.tsx): 6 metric rows with Skeleton loading state; week-over-week trend arrow for cost
- Update CommandCenter: Metrics+Alerts row expanded from 2-col to 3-col grid adding UsageOverviewPanel
- Create src/components/metrics/ folder: UsageTimeSeriesChart (recharts stacked AreaChart with var(--chart-1/2/3)), ModelUsageDonut (PieChart), AgentUsageChart and TeamUsageChart (BarChart), SessionsTable (sortable columns + 10-row Prev/Next pagination)
- Update Metrics page: Token Usage & Costs section with 5 rows (summary cards, time series+donut, agent+team bar charts, projection+cache efficiency, sessions table)
- Add usage mini-bar to AgentCard: token count + cost + progress bar; AgentGrid and Agents page pass agentUsageMap through
- Install recharts 3.8.1
- Export all new symbols through their barrel index.ts files

Co-authored-by: Frontend Developer 1 <fe-dev-1@agents.roboco.dev>

* [10372f0f] Implement full token usage instrumentation: DB migration, SDK endpoints, orchestrator hooks, analytics API, WebSocket events, dashboard integration (#86) (#89)

* [10372f0f] feat(token-usage): add Alembic migration 026 for token usage tables

Create agent_spawn_sessions, token_usage_snapshots, and daily_usage_rollups
tables with correct BIGINT columns, indexes, and unique constraint.
Chain: 025_agentrole_prompter → 026_token_usage_tables.

* [10372f0f] feat(token-usage): add ORM table classes for token usage instrumentation

Add AgentSpawnSessionTable, TokenUsageSnapshotTable, DailyUsageRollupTable
to db/tables.py. Import BigInteger and Date from SQLAlchemy. All columns
match the migration schema with BIGINT token counts and proper indexes.

* [10372f0f] feat(billing): add pricing module with calculate_cost() function

Create roboco/billing/__init__.py and roboco/billing/pricing.py with
calculate_cost() supporting Claude opus/sonnet/haiku models with
input/output/cache pricing. Unknown models return 0.0 without raising.

* [10372f0f] feat(sdk): add POST /usage/report and GET /usage/status endpoints to agent SDK

Extend _SessionState with token counters. Add TokenReportRequest and
TokenUsageStatus models. POST /usage/report additively accumulates token
counts; GET /usage/status returns current session totals for sweeper polling.

* [10372f0f] feat(orchestrator): add token usage instrumentation hooks

- _launch_spawn() calls _record_spawn_session() after successful container spawn
- stop_agent() calls _finalize_spawn_session() before container removal
- _run_sweep() calls _sweep_token_snapshots() and _sweep_daily_rollup() each tick
- New methods: _record_spawn_session, _finalize_spawn_session,
  _sweep_token_snapshots, _sweep_daily_rollup in TOKEN USAGE section

* [10372f0f] feat(api): add token usage analytics API with 7 endpoints

Create roboco/services/usage.py (UsageService) and roboco/api/routes/usage.py.
Endpoints: GET /api/usage/summary, /time-series, /by-agent, /by-team,
/by-model, /projection, /cache-efficiency. Register in app.py.

* [10372f0f] feat(dashboard): add usage_summary field to CEO dashboard

Add UsageSummary schema (tokens_today, cost_today_usd) to dashboard schemas.
Add usage_summary: UsageSummary | None to CEOOverview. Update
get_ceo_overview() to populate usage_summary from daily_usage_rollups.

* [10372f0f] fix(billing/tests): remove dead except block in _sweep_daily_rollup, add unit tests for pricing.py and services/usage.py

- Remove unreachable `except Exception as e` block in orchestrator.py
  _sweep_daily_rollup() (lines 3376-3381) which referenced undefined
  `agent_id` and was copy-pasted from _sweep_token_snapshots by mistake
- Add tests/unit/billing/test_pricing.py: 31 tests covering opus/sonnet/
  haiku tiers with all 4 token types, unknown model → 0.0, empty string
  → 0.0, and substring-match priority (longer fragment wins)
- Add tests/unit/services/test_usage.py: 25 tests covering get_summary
  trend_pct edge cases (prev=0, both=0, prev>0), get_by_agent/team/model
  pct_of_total summing to 100%, get_projection formula (avg_daily×30),
  and get_cache_efficiency hit-rate and cost_saved arithmetic
- pricing.py: 100% coverage; services/usage.py: 83% coverage (>80% target)

* [10372f0f] fix(usage): include cache tokens in time-series total_tokens to fix AC9 consistency violation

get_time_series() previously computed total_tokens as tokens_input +
tokens_output only. get_summary() includes all 4 token types (input +
output + cache_read + cache_write). AC9 requires both endpoints to agree
on their totals for the same period.

Fix: add tokens_cache_read and tokens_cache_write to the SELECT query in
get_time_series() and include them in the total_tokens calculation.

Also adds 4 new unit tests in TestGetTimeSeries covering:
- total_tokens includes cache_read and cache_write (the AC9 guard)
- zero cache tokens still produces correct total
- empty result returns empty list
- required fields are present in each point

* [10372f0f] fix(usage): remove unused imports and include cache tokens in breakdown totals (AC10)

- Remove import math (F401 — never used)
- Remove text from sqlalchemy import (F401 — never used)
- Remove unused local calculate_cost import inside get_cache_efficiency (F401)
- Add tokens_cache_read and tokens_cache_write to SELECT in get_by_agent,
  get_by_team, and get_by_model; update grand_total and per-item total to
  include all 4 token types so totals match get_summary() (AC10 fix)
- Update test mock rows to include explicit tokens_cache_read=0 and
  tokens_cache_write=0 so they work with the fixed code
- Add new test cases: test_cache_tokens_included_in_total_tokens and
  test_pct_of_total_sums_to_100_with_cache_tokens for each breakdown class

---------

Co-authored-by: Backend Developer 1 <be-dev-1@agents.roboco.dev>

* [44b9eb1f] feat(usage): align frontend API client, TS types, and chart components to real backend contract (#92) (#94)

Update all usage-related frontend code to match the actual FastAPI backend
response shapes and endpoint paths:

- panel/src/lib/api/usage.ts: rewrite all 7 API functions to use correct
  endpoint paths (/usage/summary, /usage/by-agent, /usage/by-model,
  /usage/by-team, /usage/time-series, /usage/projection,
  /usage/cache-efficiency); send period query param (24h/7d/30d not hours);
  mock generators produce data matching real backend shapes exactly;
  getUsageSessions returns [] in prod (no /usage/sessions endpoint exists)

- panel/src/types/index.ts: replace TokenUsageSnapshot with UsageSummary
  (tokens_input/tokens_output/total_cost_usd/trend_pct); update AgentUsageRow
  to use agent_slug/total_tokens/cost_usd/pct_of_total; add TeamUsageRow,
  UsageProjection, CacheEfficiencyResponse; update UsageTimePoint to use
  bucket field; update UsageSession to use agent_slug

- panel/src/hooks/use-usage.ts: rewrite all hooks to match new API and types;
  add useTeamUsage, useUsageProjection, useCacheEfficiency hooks

- panel/src/components/metrics/usage-time-series-chart.tsx: use bucket field
  (not timestamp) for axis labels
- panel/src/components/metrics/agent-usage-chart.tsx: use agent_slug and
  total_tokens (not agent_name/tokens_today)
- panel/src/components/metrics/team-usage-chart.tsx: rewrite to accept
  TeamUsageRow[] from API directly
- panel/src/components/metrics/model-usage-donut.tsx: use total_tokens,
  cost_usd, pct_of_total (not tokens/cost/percentage)
- panel/src/components/metrics/sessions-table.tsx: use agent_slug, sort keys
  updated
- panel/src/components/dashboard/usage-overview-panel.tsx: use useUsageSummary
  with tokens_input/tokens_output/total_cost_usd/trend_pct
- panel/src/app/(dashboard)/metrics/page.tsx: wire all new hooks, add
  TeamUsageChart, ProjectionCard, CacheEfficiencyCard with correct types
- panel/src/app/(dashboard)/agents/page.tsx: key agentUsageMap by agent_slug
- panel/src/components/agents/agent-card.tsx: use total_tokens and cost_usd

Co-authored-by: Frontend Developer 1 <fe-dev-1@agents.roboco.dev>

* [2161b832] fix: SDK_PORT constant, stop_agent lock refactor, usage_session_id binding, rollup 7-day window (#93) (#95)

- Add SDK_PORT = 9000 module-level constant to orchestrator.py; replace
  hardcoded 9000 in _sweep_budget_exceeded URL with SDK_PORT
- Add UUID to TYPE_CHECKING imports to satisfy ruff F821
- Refactor stop_agent: call _finalize_spawn_session BEFORE acquiring
  self._lock so the SDK HTTP round-trip does not hold the lock
- Add usage_session_id: UUID | None field to AgentInstance dataclass
- Change _record_spawn_session to return UUID | None; wire return value
  back to instance.usage_session_id in _launch_spawn
- Update _finalize_spawn_session to use WHERE id=usage_session_id for
  direct session row lookup when usage_session_id is not None
- Add started_at >= (now_utc - 7 days) filter to _sweep_daily_rollup
  aggregate query to avoid re-aggregating all-time history each sweep

Co-authored-by: Backend Developer 1 <be-dev-1@agents.roboco.dev>

* [2e0759e1] fix: pricing accuracy, import ordering, session-id binding, rollup cleanup, write-hook tests (#97) (#98)

- pricing.py: correct claude-opus-4 prices (5/25/0.50/6.25 not 15/75/1.5/3.75)
  and haiku family prices (1/5/0.10/1.25 not 0.8/4/0.08/0.20); add Ollama
  zero-cost early-return; add structlog warning for unmatched model names
- app.py: move usage_router import before routes.v1 block (ruff isort fix)
- orchestrator.py _sweep_daily_rollup: remove unused calculate_cost import;
  add blank line between stdlib (uuid4) and third-party (sqlalchemy) imports
- orchestrator.py _sweep_token_snapshots: prefer direct lookup by
  instance.usage_session_id; fall back to agent_slug heuristic only when None
- tests: add test_sweep_daily_rollup_inserts_new_row and
  test_stop_agent_finalizes_before_lock to test_orchestrator_write_hooks.py
- usage.py, routes/usage.py, stream_bus.py, test files: ruff format/lint fixes

Co-authored-by: Backend Developer 1 <be-dev-1@agents.roboco.dev>

* Mypy compliance

* fix(migrations,tests): linearize forked migration chain + correct ceo_reject coordination-root expectation

The master merge brought in 026_completed_dependency_ids alongside the rework's
026_token_usage_tables — both off 025, forking the alembic head and breaking
the enum-parity test. Rebase token-usage onto 026_completed_dependency_ids
(linear chain, single head).

Also: test_ceo_reject_routes_coordination_task_to_main_pm asserted the old
NEEDS_REVISION behavior; the lifecycle fix correctly routes a coordination root
to PENDING (Main PM's claim source). Update the assertion.

---------

Co-authored-by: Frontend Developer 1 <fe-dev-1@agents.roboco.dev>
Co-authored-by: Backend Developer 1 <be-dev-1@agents.roboco.dev>
Co-authored-by: Renn F <rennf93@users.noreply.github.com>
This commit is contained in:
Renzo F
2026-06-10 14:38:44 +02:00
committed by GitHub
co-authored by Frontend Developer 1 Backend Developer 1 Renn F
parent 93c6ef8a57
commit b3057628b0
40 changed files with 5039 additions and 19 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/unit/billing/__init__.py
View File
tests/unit/billing/test_pricing.py
+299
View File
@@ -0,0 +1,299 @@
"""
Unit tests for roboco.billing.pricing — calculate_cost().
Covers:
- Each model tier (opus, sonnet, haiku) with all 4 token types.
- Unknown model name returns 0.0 without raising.
- Empty model string returns 0.0 without raising.
- Substring match correctness: longer fragment wins
(e.g. 'claude-sonnet-4-6' matches 'claude-sonnet-4' not bare 'sonnet').
"""
from __future__ import annotations
import pytest
from roboco.billing.pricing import calculate_cost
# ---------------------------------------------------------------------------
# Named constants (ruff PLR2004: magic values in comparisons must be named).
# ---------------------------------------------------------------------------
# Token counts
_M = 1_000_000 # 1 million tokens
# Pricing — per-1M USD, matches the _PRICING table in pricing.py
_OPUS_INPUT = 5.00
_OPUS_OUTPUT = 25.00
_OPUS_CACHE_READ = 0.50
_OPUS_CACHE_WRITE = 6.25
_SONNET_INPUT = 3.00
_SONNET_OUTPUT = 15.00
_SONNET_CACHE_READ = 0.30
_SONNET_CACHE_WRITE = 0.75
_HAIKU_INPUT = 1.00
_HAIKU_OUTPUT = 5.00
_HAIKU_CACHE_READ = 0.10
_HAIKU_CACHE_WRITE = 1.25
_HAIKU3_INPUT = 0.25 # claude-haiku-3 is cheaper than haiku-3-5 / haiku-4
# Tolerance for floating-point comparisons
_TOL = 1e-4
# ---------------------------------------------------------------------------
# Opus tier
# ---------------------------------------------------------------------------
class TestOpusTier:
"""claude-opus-4 family pricing."""
def test_input_only(self) -> None:
cost = calculate_cost("claude-opus-4-5", tokens_input=_M, tokens_output=0)
assert abs(cost - _OPUS_INPUT) < _TOL
def test_output_only(self) -> None:
cost = calculate_cost("claude-opus-4-5", tokens_input=0, tokens_output=_M)
assert abs(cost - _OPUS_OUTPUT) < _TOL
def test_cache_read_only(self) -> None:
cost = calculate_cost(
"claude-opus-4-5",
tokens_input=0,
tokens_output=0,
tokens_cache_read=_M,
)
assert abs(cost - _OPUS_CACHE_READ) < _TOL
def test_cache_write_only(self) -> None:
cost = calculate_cost(
"claude-opus-4-5",
tokens_input=0,
tokens_output=0,
tokens_cache_write=_M,
)
assert abs(cost - _OPUS_CACHE_WRITE) < _TOL
def test_all_token_types(self) -> None:
cost = calculate_cost(
"claude-opus-4-5",
tokens_input=_M,
tokens_output=_M,
tokens_cache_read=_M,
tokens_cache_write=_M,
)
expected = _OPUS_INPUT + _OPUS_OUTPUT + _OPUS_CACHE_READ + _OPUS_CACHE_WRITE
assert abs(cost - expected) < _TOL
def test_short_alias(self) -> None:
"""Bare 'opus' alias resolves to the opus tier."""
cost = calculate_cost("opus", tokens_input=_M, tokens_output=0)
assert abs(cost - _OPUS_INPUT) < _TOL
def test_returns_float(self) -> None:
cost = calculate_cost("claude-opus-4", tokens_input=100, tokens_output=50)
assert isinstance(cost, float)
# ---------------------------------------------------------------------------
# Sonnet tier
# ---------------------------------------------------------------------------
class TestSonnetTier:
"""claude-sonnet-4 family pricing."""
def test_input_only(self) -> None:
cost = calculate_cost("claude-sonnet-4-6", tokens_input=_M, tokens_output=0)
assert abs(cost - _SONNET_INPUT) < _TOL
def test_output_only(self) -> None:
cost = calculate_cost("claude-sonnet-4-6", tokens_input=0, tokens_output=_M)
assert abs(cost - _SONNET_OUTPUT) < _TOL
def test_cache_read_only(self) -> None:
cost = calculate_cost(
"claude-sonnet-4-6",
tokens_input=0,
tokens_output=0,
tokens_cache_read=_M,
)
assert abs(cost - _SONNET_CACHE_READ) < _TOL
def test_cache_write_only(self) -> None:
cost = calculate_cost(
"claude-sonnet-4-6",
tokens_input=0,
tokens_output=0,
tokens_cache_write=_M,
)
assert abs(cost - _SONNET_CACHE_WRITE) < _TOL
def test_all_token_types(self) -> None:
cost = calculate_cost(
"claude-sonnet-4-6",
tokens_input=_M,
tokens_output=_M,
tokens_cache_read=_M,
tokens_cache_write=_M,
)
expected = (
_SONNET_INPUT + _SONNET_OUTPUT + _SONNET_CACHE_READ + _SONNET_CACHE_WRITE
)
assert abs(cost - expected) < _TOL
def test_short_alias(self) -> None:
"""Bare 'sonnet' alias resolves to the sonnet tier."""
cost = calculate_cost("sonnet", tokens_input=_M, tokens_output=0)
assert abs(cost - _SONNET_INPUT) < _TOL
def test_35_variant(self) -> None:
"""claude-3-5-sonnet resolves to sonnet tier."""
cost = calculate_cost(
"claude-3-5-sonnet-20241022", tokens_input=_M, tokens_output=0
)
assert abs(cost - _SONNET_INPUT) < _TOL
# ---------------------------------------------------------------------------
# Haiku tier
# ---------------------------------------------------------------------------
class TestHaikuTier:
"""claude-haiku family pricing."""
def test_input_only(self) -> None:
cost = calculate_cost("claude-haiku-4-5", tokens_input=_M, tokens_output=0)
assert abs(cost - _HAIKU_INPUT) < _TOL
def test_output_only(self) -> None:
cost = calculate_cost("claude-haiku-4-5", tokens_input=0, tokens_output=_M)
assert abs(cost - _HAIKU_OUTPUT) < _TOL
def test_cache_read_only(self) -> None:
cost = calculate_cost(
"claude-haiku-4-5",
tokens_input=0,
tokens_output=0,
tokens_cache_read=_M,
)
assert abs(cost - _HAIKU_CACHE_READ) < _TOL
def test_cache_write_only(self) -> None:
cost = calculate_cost(
"claude-haiku-4-5",
tokens_input=0,
tokens_output=0,
tokens_cache_write=_M,
)
assert abs(cost - _HAIKU_CACHE_WRITE) < _TOL
def test_all_token_types(self) -> None:
cost = calculate_cost(
"claude-haiku-4-5",
tokens_input=_M,
tokens_output=_M,
tokens_cache_read=_M,
tokens_cache_write=_M,
)
expected = _HAIKU_INPUT + _HAIKU_OUTPUT + _HAIKU_CACHE_READ + _HAIKU_CACHE_WRITE
assert abs(cost - expected) < _TOL
def test_short_alias(self) -> None:
"""Bare 'haiku' alias resolves to the haiku tier."""
cost = calculate_cost("haiku", tokens_input=_M, tokens_output=0)
assert abs(cost - _HAIKU_INPUT) < _TOL
def test_haiku3_variant(self) -> None:
"""claude-haiku-3 has lower pricing than haiku-3-5."""
cost = calculate_cost("claude-haiku-3", tokens_input=_M, tokens_output=0)
assert abs(cost - _HAIKU3_INPUT) < _TOL
# ---------------------------------------------------------------------------
# Unknown / edge cases — must return 0.0 without raising
# ---------------------------------------------------------------------------
class TestUnknownModels:
def test_unknown_model_name_returns_zero(self) -> None:
cost = calculate_cost("gpt-4o", tokens_input=_M, tokens_output=_M)
assert cost == 0.0
def test_empty_string_returns_zero(self) -> None:
cost = calculate_cost("", tokens_input=_M, tokens_output=_M)
assert cost == 0.0
def test_gibberish_returns_zero(self) -> None:
cost = calculate_cost(
"totally-unknown-model-xyz", tokens_input=100, tokens_output=100
)
assert cost == 0.0
def test_zero_tokens_with_unknown_model_returns_zero(self) -> None:
cost = calculate_cost("unknown", tokens_input=0, tokens_output=0)
assert cost == 0.0
def test_does_not_raise_on_unknown_model(self) -> None:
"""Must not raise regardless of token counts."""
try:
calculate_cost(
"not-a-claude-model",
tokens_input=999_999,
tokens_output=999_999,
)
except Exception as exc:
pytest.fail(f"calculate_cost raised unexpectedly: {exc}")
# ---------------------------------------------------------------------------
# Substring match correctness
# ---------------------------------------------------------------------------
# Named constants for the comparison floor/ceiling used in these tests.
_ZERO_COST = 0.0
_SONNET_CHEAPER_THAN_OPUS = True # structural assertion in the test below
class TestSubstringMatchPriority:
def test_claude_sonnet_4_resolves_non_zero(self) -> None:
"""'claude-sonnet-4-6' must find a match (non-zero cost)."""
cost = calculate_cost("claude-sonnet-4-6", tokens_input=_M, tokens_output=0)
assert cost > _ZERO_COST
def test_haiku3_cheaper_than_haiku4(self) -> None:
"""claude-haiku-3 is cheaper than claude-haiku-4 — longest-match wins."""
haiku3_cost = calculate_cost("claude-haiku-3", tokens_input=_M, tokens_output=0)
haiku4_cost = calculate_cost("claude-haiku-4", tokens_input=_M, tokens_output=0)
# haiku-3 ($0.25/1M) < haiku-4 ($1.00/1M)
assert haiku3_cost < haiku4_cost
def test_non_claude_model_returns_zero(self) -> None:
"""A random non-Claude model must not match any Claude pricing entry."""
non_opus_cost = calculate_cost("llama-3-70b", tokens_input=_M, tokens_output=0)
assert non_opus_cost == _ZERO_COST
def test_opus_model_non_zero(self) -> None:
"""Claude opus model resolves to non-zero cost."""
opus_cost = calculate_cost("claude-opus-4", tokens_input=_M, tokens_output=0)
assert opus_cost > _ZERO_COST
def test_zero_tokens_returns_zero_for_known_model(self) -> None:
"""Known model with 0 tokens has 0 cost."""
cost = calculate_cost("claude-opus-4", tokens_input=0, tokens_output=0)
assert cost == _ZERO_COST
def test_case_insensitive_matching(self) -> None:
"""Model name matching is case-insensitive."""
lower_cost = calculate_cost(
"claude-sonnet-4-6", tokens_input=1000, tokens_output=1000
)
upper_cost = calculate_cost(
"CLAUDE-SONNET-4-6", tokens_input=1000, tokens_output=1000
)
assert lower_cost == upper_cost
assert lower_cost > _ZERO_COST
tests/unit/runtime/test_orchestrator_write_hooks.py
+543
View File
@@ -0,0 +1,543 @@
"""
Unit tests for orchestrator write-hooks:
_finalize_spawn_session — closes the agent_spawn_sessions DB row on stop
_sweep_token_snapshots — polls active agents and upserts token snapshots
These tests mock the httpx transport and the SQLAlchemy session factory so no
real network or database is required.
Coverage:
1. _finalize_spawn_session success — SDK returns token data → DB update
carries those exact values to calculate_cost and the UPDATE statement.
2. _finalize_spawn_session HTTP error — SDK unreachable → DB update proceeds
with all-zero token counts (finalization must not raise).
3. _sweep_token_snapshots active agent — non-zero tokens → snapshot row
inserted and session row updated.
4. _sweep_token_snapshots per-agent HTTP error — ConnectError on one agent
is caught; the sweep continues and the next agent is still processed.
"""
from __future__ import annotations
from contextlib import asynccontextmanager
from pathlib import Path
from typing import Any
from unittest.mock import AsyncMock, MagicMock, patch
from uuid import UUID, uuid4
import httpx
from roboco.models.runtime import (
AgentInstance,
OrchestratorAgentConfig,
OrchestratorAgentState,
)
from roboco.runtime.orchestrator import AgentOrchestrator
# ---------------------------------------------------------------------------
# Module-level constants (ruff PLR2004: no magic values in comparisons)
# ---------------------------------------------------------------------------
_AGENT_ID = "be-dev-1"
_AGENT_ID_2 = "be-dev-2"
# Token counts used in success-path assertions
_TI = 111 # tokens_input
_TO = 222 # tokens_output
_TCR = 33 # tokens_cache_read
_TCW = 44 # tokens_cache_write
# Token counts for the snapshot test
_SNAP_TI = 50
_SNAP_TO = 100
_SNAP_TCR = 10
_SNAP_TCW = 5
# Token counts for the loop-continues test (agent-2)
_LOOP_TI = 25
_LOOP_TO = 75
# Expected number of DB execute() calls for a normal finalize (SELECT + UPDATE)
_FINALIZE_EXEC_CALLS = 2
# ---------------------------------------------------------------------------
# Test helpers
# ---------------------------------------------------------------------------
def _make_orchestrator() -> AgentOrchestrator:
"""Minimal AgentOrchestrator — no background tasks, no real DB."""
return AgentOrchestrator(mcp_config_dir=Path("/tmp"), project_root=Path("/tmp"))
def _make_instance(
agent_id: str = _AGENT_ID,
usage_session_id: UUID | None = None,
) -> AgentInstance:
"""Return an ACTIVE AgentInstance with a running container."""
return AgentInstance(
agent_id=agent_id,
state=OrchestratorAgentState.ACTIVE,
container_id="abc123def456",
config=OrchestratorAgentConfig(
agent_id=agent_id,
blueprint_path=Path("/tmp/blueprint.md"),
model="sonnet",
),
usage_session_id=usage_session_id,
)
def _mock_response(
status: int = 200,
json_data: dict[str, Any] | None = None,
) -> MagicMock:
"""Build a mock httpx.Response."""
resp = MagicMock(spec=httpx.Response)
resp.status_code = status
resp.json = MagicMock(return_value=json_data or {})
return resp
def _make_db_factory(
session_row: Any = None,
add_list: list[Any] | None = None,
execute_list: list[Any] | None = None,
) -> Any:
"""Return a callable that acts like get_session_factory().
The returned callable, when called with no arguments, returns an async
context manager yielding a mock AsyncSession whose execute() returns a
result whose scalar_one_or_none() returns *session_row*.
"""
@asynccontextmanager
async def _db_context() -> Any:
db = MagicMock()
result = MagicMock()
result.scalar_one_or_none = MagicMock(return_value=session_row)
async def _exec(stmt: Any) -> MagicMock:
if execute_list is not None:
execute_list.append(stmt)
return result
db.execute = AsyncMock(side_effect=_exec)
db.commit = AsyncMock()
def _add(obj: Any) -> None:
if add_list is not None:
add_list.append(obj)
db.add = _add if add_list is not None else MagicMock()
yield db
return _db_context
class _FakeHTTPClient:
"""Drop-in replacement for ``httpx.AsyncClient`` in tests.
Accepts a *handler* callable ``(url: str) -> httpx.Response | raises``
that is invoked by ``get()``. Supports the ``async with`` protocol.
"""
def __init__(self, handler: Any, **_: Any) -> None:
self._handler = handler
async def __aenter__(self) -> _FakeHTTPClient:
return self
async def __aexit__(self, *_: Any) -> None:
pass
async def get(self, url: str, **_: Any) -> Any:
return self._handler(url)
# ---------------------------------------------------------------------------
# _finalize_spawn_session — success path
# ---------------------------------------------------------------------------
async def test_finalize_spawn_session_success_calls_calculate_cost() -> None:
"""Token values returned by the SDK /usage/status are passed to calculate_cost.
This verifies the full data-flow: SDK response → token vars → cost calc.
"""
orch = _make_orchestrator()
session_uuid = uuid4()
orch._instances[_AGENT_ID] = _make_instance(usage_session_id=session_uuid)
token_data = {
"tokens_input": _TI,
"tokens_output": _TO,
"tokens_cache_read": _TCR,
"tokens_cache_write": _TCW,
}
def _handler(_url: str) -> Any:
return _mock_response(200, token_data)
session_row = MagicMock()
session_row.id = session_uuid
db_factory = _make_db_factory(session_row=session_row)
def _client_cls(**_kw: Any) -> _FakeHTTPClient:
return _FakeHTTPClient(_handler)
with (
patch("roboco.runtime.orchestrator.httpx.AsyncClient", _client_cls),
patch("roboco.db.base.get_session_factory", return_value=db_factory),
patch("roboco.billing.pricing.calculate_cost", return_value=0.001) as mock_cost,
):
await orch._finalize_spawn_session(_AGENT_ID, exit_reason="stopped")
mock_cost.assert_called_once_with(
model="sonnet",
tokens_input=_TI,
tokens_output=_TO,
tokens_cache_read=_TCR,
tokens_cache_write=_TCW,
)
async def test_finalize_spawn_session_success_executes_select_and_update() -> None:
"""When a session row exists the function calls execute() twice: SELECT + UPDATE."""
orch = _make_orchestrator()
session_uuid = uuid4()
orch._instances[_AGENT_ID] = _make_instance(usage_session_id=session_uuid)
def _handler(_url: str) -> Any:
return _mock_response(
200,
{
"tokens_input": 10,
"tokens_output": 20,
"tokens_cache_read": 0,
"tokens_cache_write": 0,
},
)
session_row = MagicMock()
session_row.id = session_uuid
execute_calls: list[Any] = []
db_factory = _make_db_factory(session_row=session_row, execute_list=execute_calls)
def _client_cls(**_kw: Any) -> _FakeHTTPClient:
return _FakeHTTPClient(_handler)
with (
patch("roboco.runtime.orchestrator.httpx.AsyncClient", _client_cls),
patch("roboco.db.base.get_session_factory", return_value=db_factory),
patch("roboco.billing.pricing.calculate_cost", return_value=0.0),
):
await orch._finalize_spawn_session(_AGENT_ID, exit_reason="completed")
# SELECT (find the row) + UPDATE (write the values) = 2 execute() calls
assert len(execute_calls) == _FINALIZE_EXEC_CALLS
# ---------------------------------------------------------------------------
# _finalize_spawn_session — HTTP-error path
# ---------------------------------------------------------------------------
async def test_finalize_spawn_session_http_error_uses_zero_tokens() -> None:
"""When the SDK endpoint is unreachable, finalization uses zero tokens.
The function must not raise; cost must be calculated with all-zero counts.
"""
orch = _make_orchestrator()
session_uuid = uuid4()
orch._instances[_AGENT_ID] = _make_instance(usage_session_id=session_uuid)
def _boom(_url: str) -> Any:
raise httpx.ConnectError("container not reachable")
session_row = MagicMock()
session_row.id = session_uuid
db_factory = _make_db_factory(session_row=session_row)
def _client_cls(**_kw: Any) -> _FakeHTTPClient:
return _FakeHTTPClient(_boom)
with (
patch("roboco.runtime.orchestrator.httpx.AsyncClient", _client_cls),
patch("roboco.db.base.get_session_factory", return_value=db_factory),
patch("roboco.billing.pricing.calculate_cost", return_value=0.0) as mock_cost,
):
# Must not raise even though the SDK is unreachable
await orch._finalize_spawn_session(_AGENT_ID, exit_reason="stopped")
mock_cost.assert_called_once_with(
model="sonnet",
tokens_input=0,
tokens_output=0,
tokens_cache_read=0,
tokens_cache_write=0,
)
async def test_finalize_spawn_session_non_200_uses_zero_tokens() -> None:
"""A non-200 SDK response results in zero-token finalization, no exception."""
orch = _make_orchestrator()
session_uuid = uuid4()
orch._instances[_AGENT_ID] = _make_instance(usage_session_id=session_uuid)
def _handler(_url: str) -> Any:
return _mock_response(503)
session_row = MagicMock()
session_row.id = session_uuid
db_factory = _make_db_factory(session_row=session_row)
def _client_cls(**_kw: Any) -> _FakeHTTPClient:
return _FakeHTTPClient(_handler)
with (
patch("roboco.runtime.orchestrator.httpx.AsyncClient", _client_cls),
patch("roboco.db.base.get_session_factory", return_value=db_factory),
patch("roboco.billing.pricing.calculate_cost", return_value=0.0) as mock_cost,
):
await orch._finalize_spawn_session(_AGENT_ID, exit_reason="stopped")
mock_cost.assert_called_once_with(
model="sonnet",
tokens_input=0,
tokens_output=0,
tokens_cache_read=0,
tokens_cache_write=0,
)
# ---------------------------------------------------------------------------
# _sweep_token_snapshots — active agent
# ---------------------------------------------------------------------------
async def test_sweep_token_snapshots_inserts_snapshot_for_active_agent() -> None:
"""An active agent with non-zero tokens gets a snapshot row added to the DB."""
orch = _make_orchestrator()
instance = _make_instance(_AGENT_ID)
instance.state = OrchestratorAgentState.ACTIVE
orch._instances[_AGENT_ID] = instance
token_data = {
"tokens_input": _SNAP_TI,
"tokens_output": _SNAP_TO,
"tokens_cache_read": _SNAP_TCR,
"tokens_cache_write": _SNAP_TCW,
}
def _handler(_url: str) -> Any:
return _mock_response(200, token_data)
session_row = MagicMock()
session_row.id = uuid4()
added: list[Any] = []
db_factory = _make_db_factory(session_row=session_row, add_list=added)
def _client_cls(**_kw: Any) -> _FakeHTTPClient:
return _FakeHTTPClient(_handler)
with (
patch("roboco.runtime.orchestrator.httpx.AsyncClient", _client_cls),
patch("roboco.db.base.get_session_factory", return_value=db_factory),
):
await orch._sweep_token_snapshots()
# Exactly one snapshot row must have been passed to db.add()
assert len(added) == 1
snap = added[0]
assert snap.tokens_input == _SNAP_TI
assert snap.tokens_output == _SNAP_TO
assert snap.tokens_cache_read == _SNAP_TCR
assert snap.tokens_cache_write == _SNAP_TCW
async def test_sweep_token_snapshots_skips_zero_token_agents() -> None:
"""An agent whose SDK reports all-zero tokens is skipped (no DB writes)."""
orch = _make_orchestrator()
instance = _make_instance(_AGENT_ID)
instance.state = OrchestratorAgentState.ACTIVE
orch._instances[_AGENT_ID] = instance
def _handler(_url: str) -> Any:
return _mock_response(
200,
{
"tokens_input": 0,
"tokens_output": 0,
"tokens_cache_read": 0,
"tokens_cache_write": 0,
},
)
added: list[Any] = []
db_factory = _make_db_factory(add_list=added)
def _client_cls(**_kw: Any) -> _FakeHTTPClient:
return _FakeHTTPClient(_handler)
with (
patch("roboco.runtime.orchestrator.httpx.AsyncClient", _client_cls),
patch("roboco.db.base.get_session_factory", return_value=db_factory),
):
await orch._sweep_token_snapshots()
assert added == []
async def test_sweep_token_snapshots_per_agent_error_does_not_abort_loop() -> None:
"""A ConnectError for one agent is caught; the next agent is still processed."""
orch = _make_orchestrator()
# Agent 1: HTTP error
inst1 = _make_instance(_AGENT_ID)
inst1.state = OrchestratorAgentState.ACTIVE
orch._instances[_AGENT_ID] = inst1
# Agent 2: success with non-zero tokens
inst2 = _make_instance(_AGENT_ID_2)
inst2.state = OrchestratorAgentState.ACTIVE
orch._instances[_AGENT_ID_2] = inst2
def _handler(url: str) -> Any:
if _AGENT_ID in url and _AGENT_ID_2 not in url:
raise httpx.ConnectError("agent-1 unreachable")
return _mock_response(
200,
{
"tokens_input": _LOOP_TI,
"tokens_output": _LOOP_TO,
"tokens_cache_read": 0,
"tokens_cache_write": 0,
},
)
session_row = MagicMock()
session_row.id = uuid4()
added: list[Any] = []
db_factory = _make_db_factory(session_row=session_row, add_list=added)
def _client_cls(**_kw: Any) -> _FakeHTTPClient:
return _FakeHTTPClient(_handler)
with (
patch("roboco.runtime.orchestrator.httpx.AsyncClient", _client_cls),
patch("roboco.db.base.get_session_factory", return_value=db_factory),
):
await orch._sweep_token_snapshots()
# Only agent-2's snapshot should be present; agent-1's error was caught.
assert len(added) == 1
assert added[0].tokens_input == _LOOP_TI
assert added[0].tokens_output == _LOOP_TO
# ---------------------------------------------------------------------------
# _sweep_daily_rollup — inserts new row when none exists
# ---------------------------------------------------------------------------
# Token counts for the rollup test
_ROLLUP_TI = 200
_ROLLUP_TO = 300
_ROLLUP_TCR = 20
_ROLLUP_TCW = 10
async def test_sweep_daily_rollup_inserts_new_row() -> None:
"""When no existing DailyUsageRollupTable row exists, db.add() is called
with the correct aggregated token values."""
orch = _make_orchestrator()
# Build a fake aggregate result row
agg_row = MagicMock()
agg_row.date = "2026-06-10"
agg_row.agent_slug = _AGENT_ID
agg_row.team = "backend"
agg_row.model = "sonnet"
agg_row.tokens_input = _ROLLUP_TI
agg_row.tokens_output = _ROLLUP_TO
agg_row.tokens_cache_read = _ROLLUP_TCR
agg_row.tokens_cache_write = _ROLLUP_TCW
agg_row.total_cost_usd = 0.0
agg_row.session_count = 1
added: list[Any] = []
call_count = 0
@asynccontextmanager
async def _db_context() -> Any:
nonlocal call_count
db = MagicMock()
db.commit = AsyncMock()
def _add(obj: Any) -> None:
added.append(obj)
db.add = _add
async def _exec(_stmt: Any) -> MagicMock:
nonlocal call_count
call_count += 1
result = MagicMock()
if call_count == 1:
# First call: aggregate SELECT — return one agg_row via fetchall()
result.fetchall = MagicMock(return_value=[agg_row])
result.scalar_one_or_none = MagicMock(return_value=None)
else:
# Second call: lookup SELECT for existing row — return None
result.fetchall = MagicMock(return_value=[])
result.scalar_one_or_none = MagicMock(return_value=None)
return result
db.execute = AsyncMock(side_effect=_exec)
yield db
with patch("roboco.db.base.get_session_factory", return_value=_db_context):
await orch._sweep_daily_rollup()
# Exactly one new DailyUsageRollupTable row must have been added
assert len(added) == 1
row = added[0]
assert row.tokens_input == _ROLLUP_TI
assert row.tokens_output == _ROLLUP_TO
assert row.tokens_cache_read == _ROLLUP_TCR
assert row.tokens_cache_write == _ROLLUP_TCW
# ---------------------------------------------------------------------------
# stop_agent — _finalize_spawn_session is awaited before acquiring the lock
# ---------------------------------------------------------------------------
async def test_stop_agent_finalizes_before_lock() -> None:
"""stop_agent awaits _finalize_spawn_session when the instance has a
running container_id (the finalization must happen before the lock)."""
orch = _make_orchestrator()
instance = _make_instance(_AGENT_ID)
instance.container_id = "abc123def456" # non-None → finalize must be called
orch._instances[_AGENT_ID] = instance
finalized: list[str] = []
async def _fake_finalize(agent_id: str, exit_reason: str = "stopped") -> None: # noqa: ARG001
finalized.append(agent_id)
# Stub out the Docker subprocess so stop_agent doesn't actually run Docker
mock_proc = MagicMock()
mock_proc.wait = AsyncMock()
with (
patch.object(orch, "_finalize_spawn_session", side_effect=_fake_finalize),
patch("asyncio.create_subprocess_exec", AsyncMock(return_value=mock_proc)),
patch.object(orch, "_remove_container", AsyncMock()),
):
await orch.stop_agent(_AGENT_ID, graceful=True)
# _finalize_spawn_session must have been called exactly once with our agent id
assert finalized == [_AGENT_ID]
tests/unit/services/test_usage.py
+761
View File
@@ -0,0 +1,761 @@
"""
Unit tests for roboco.services.usage — UsageService analytics methods.
These tests mock the SQLAlchemy AsyncSession.execute() boundary and
verify the arithmetic / logic of each analytics method:
- get_summary: trend_pct edge cases (prev=0, curr=0, both=0, prev>0)
- get_by_agent/team/model: pct_of_total sums to 100%
- get_projection: projected_monthly = avg_daily * 30
- get_cache_efficiency: cache_hit_rate and cost_saved arithmetic
"""
from __future__ import annotations
import datetime
from unittest.mock import AsyncMock, MagicMock
import pytest
from roboco.services.usage import UsageService
# ---------------------------------------------------------------------------
# Named constants (ruff PLR2004: magic values in comparisons must be named).
# ---------------------------------------------------------------------------
# Tolerance for floating-point arithmetic comparisons.
_TOL = 0.001
# Tolerance for percentage-sum assertions (rounding in pct_of_total).
_PCT_TOL = 0.1
# token count helpers
_ZERO = 0
_M = 1_000_000
# Expected values for projection tests
_COST_7D = 70.0
_EXPECTED_AVG_DAILY = 10.0 # 70 / 7
_EXPECTED_MONTHLY = 300.0 # 10 * 30
_DAYS_BASIS = 7
# Expected values for cache efficiency tests
_CACHE_READ_TOKENS = 400
_INPUT_TOKENS = 600
_EXPECTED_HIT_RATE = 0.4 # 400 / (600 + 400)
_FULL_INPUT_PRICE = 3.00 # sonnet baseline USD/1M
_CACHE_READ_PRICE = 0.30
_EXPECTED_COST_SAVED = _FULL_INPUT_PRICE - _CACHE_READ_PRICE # = 2.70 per 1M
# Expected trend_pct values
_TREND_NONE = 0.0
_TREND_NEW = 100.0 # curr > 0, prev == 0
_TREND_DOUBLED = 200.0 # curr / prev = 3.0x → +200 %
_TREND_HALVED = -50.0 # curr / prev = 0.5x → -50 %
# Expected total_tokens when cache tokens are included
_TOTAL_WITH_CACHE = 300 # 100+100+50+50
# pct_of_total checks
_FULL_PCT = 100.0
# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------
def _make_row(**kwargs: object) -> MagicMock:
"""Return a MagicMock that mimics a SQLAlchemy Row with named attributes."""
row = MagicMock()
for k, v in kwargs.items():
setattr(row, k, v)
return row
def _result_one(row: MagicMock) -> MagicMock:
"""Return a mock execute() result whose .one() returns `row`."""
result = MagicMock()
result.one = MagicMock(return_value=row)
return result
def _result_fetchall(rows: list[MagicMock]) -> MagicMock:
"""Return a mock execute() result whose .fetchall() returns `rows`."""
result = MagicMock()
result.fetchall = MagicMock(return_value=rows)
return result
def _service_with_execute(*return_values: object) -> UsageService:
"""Build a UsageService whose session.execute() returns the provided
values in sequence (one per call)."""
session = MagicMock()
session.execute = AsyncMock(side_effect=list(return_values))
return UsageService(session)
# ---------------------------------------------------------------------------
# get_summary — trend_pct arithmetic
# ---------------------------------------------------------------------------
class TestGetSummaryTrendPct:
@pytest.mark.asyncio
async def test_both_zero_returns_zero_trend(self) -> None:
"""When current and previous totals are both 0, trend_pct must be 0.0."""
current_row = _make_row(
tokens_input=_ZERO,
tokens_output=_ZERO,
tokens_cache_read=_ZERO,
tokens_cache_write=_ZERO,
total_cost_usd=0.0,
)
prev_row = _make_row(total=_ZERO)
svc = _service_with_execute(_result_one(current_row), _result_one(prev_row))
result = await svc.get_summary("24h")
assert result["trend_pct"] == _TREND_NONE
@pytest.mark.asyncio
async def test_prev_zero_curr_positive_returns_100(self) -> None:
"""When prev period is 0 but current is positive, trend_pct = 100.0."""
current_row = _make_row(
tokens_input=500,
tokens_output=500,
tokens_cache_read=_ZERO,
tokens_cache_write=_ZERO,
total_cost_usd=0.01,
)
prev_row = _make_row(total=_ZERO)
svc = _service_with_execute(_result_one(current_row), _result_one(prev_row))
result = await svc.get_summary("24h")
assert result["trend_pct"] == _TREND_NEW
@pytest.mark.asyncio
async def test_positive_trend_calculation(self) -> None:
"""trend_pct = (current - previous) / previous * 100 when prev > 0.
current = 1500 input + 1500 output = 3000; prev = 1000
→ (3000 - 1000) / 1000 * 100 = 200.0
"""
current_row = _make_row(
tokens_input=1500,
tokens_output=1500,
tokens_cache_read=_ZERO,
tokens_cache_write=_ZERO,
total_cost_usd=0.1,
)
prev_row = _make_row(total=1000)
svc = _service_with_execute(_result_one(current_row), _result_one(prev_row))
result = await svc.get_summary("24h")
assert abs(result["trend_pct"] - _TREND_DOUBLED) < _TOL
@pytest.mark.asyncio
async def test_negative_trend_calculation(self) -> None:
"""Negative trend when usage drops.
current = 250 + 250 = 500; prev = 1000
→ (500 - 1000) / 1000 * 100 = -50.0
"""
current_row = _make_row(
tokens_input=250,
tokens_output=250,
tokens_cache_read=_ZERO,
tokens_cache_write=_ZERO,
total_cost_usd=0.01,
)
prev_row = _make_row(total=1000)
svc = _service_with_execute(_result_one(current_row), _result_one(prev_row))
result = await svc.get_summary("24h")
assert abs(result["trend_pct"] - _TREND_HALVED) < _TOL
@pytest.mark.asyncio
async def test_cache_tokens_included_in_total(self) -> None:
"""total_tokens includes cache_read and cache_write tokens."""
current_row = _make_row(
tokens_input=100,
tokens_output=100,
tokens_cache_read=50,
tokens_cache_write=50,
total_cost_usd=0.005,
)
prev_row = _make_row(total=_ZERO)
svc = _service_with_execute(_result_one(current_row), _result_one(prev_row))
result = await svc.get_summary("24h")
assert result["total_tokens"] == _TOTAL_WITH_CACHE
@pytest.mark.asyncio
async def test_summary_contains_required_fields(self) -> None:
"""Response dict must include all required summary fields."""
current_row = _make_row(
tokens_input=_ZERO,
tokens_output=_ZERO,
tokens_cache_read=_ZERO,
tokens_cache_write=_ZERO,
total_cost_usd=0.0,
)
prev_row = _make_row(total=_ZERO)
svc = _service_with_execute(_result_one(current_row), _result_one(prev_row))
result = await svc.get_summary("24h")
for field in ("tokens_input", "tokens_output", "total_cost_usd", "trend_pct"):
assert field in result, f"Missing field: {field}"
# ---------------------------------------------------------------------------
# get_time_series — total_tokens includes all 4 token types (summary consistency)
# ---------------------------------------------------------------------------
# Named constants for time-series tests
_TS_INPUT = 100
_TS_OUTPUT = 200
_TS_CACHE_READ = 50
_TS_CACHE_WRITE = 30
# total = 100 + 200 + 50 + 30 = 380
_TS_TOTAL_WITH_CACHE = 380
# Without cache tokens (the old wrong formula): 100 + 200 = 300
_TS_TOTAL_WITHOUT_CACHE = 300
class TestGetTimeSeries:
@pytest.mark.asyncio
async def test_total_tokens_includes_cache_read_and_write(self) -> None:
"""total_tokens in each time-series point must include cache tokens.
This is the time-series / summary consistency requirement: time-series
total_tokens must sum to the same value as get_summary()'s total_tokens
for the same period. The old implementation used ti + to_ (without
cache), which violated this constraint whenever cache tokens were non-zero.
"""
bucket_dt = datetime.datetime(2026, 6, 9, 12, 0, 0, tzinfo=datetime.UTC)
row = _make_row(
bucket=bucket_dt,
tokens_input=_TS_INPUT,
tokens_output=_TS_OUTPUT,
tokens_cache_read=_TS_CACHE_READ,
tokens_cache_write=_TS_CACHE_WRITE,
cost_usd=0.01,
)
svc = _service_with_execute(_result_fetchall([row]))
result = await svc.get_time_series("24h")
assert len(result) == 1
assert result[0]["total_tokens"] == _TS_TOTAL_WITH_CACHE
@pytest.mark.asyncio
async def test_total_tokens_without_cache_still_correct(self) -> None:
"""When cache tokens are zero, total_tokens == tokens_input + tokens_output."""
bucket_dt = datetime.datetime(2026, 6, 9, 12, 0, 0, tzinfo=datetime.UTC)
row = _make_row(
bucket=bucket_dt,
tokens_input=_TS_INPUT,
tokens_output=_TS_OUTPUT,
tokens_cache_read=_ZERO,
tokens_cache_write=_ZERO,
cost_usd=0.01,
)
svc = _service_with_execute(_result_fetchall([row]))
result = await svc.get_time_series("24h")
assert result[0]["total_tokens"] == _TS_INPUT + _TS_OUTPUT
@pytest.mark.asyncio
async def test_empty_result_returns_empty_list(self) -> None:
svc = _service_with_execute(_result_fetchall([]))
result = await svc.get_time_series("24h")
assert result == []
@pytest.mark.asyncio
async def test_point_contains_required_fields(self) -> None:
"""Each time-series point must have bucket, tokens_input, tokens_output,
total_tokens, and cost_usd fields."""
bucket_dt = datetime.datetime(2026, 6, 9, 12, 0, 0, tzinfo=datetime.UTC)
row = _make_row(
bucket=bucket_dt,
tokens_input=100,
tokens_output=100,
tokens_cache_read=_ZERO,
tokens_cache_write=_ZERO,
cost_usd=0.01,
)
svc = _service_with_execute(_result_fetchall([row]))
result = await svc.get_time_series("24h")
assert len(result) == 1
point = result[0]
for field in (
"bucket",
"tokens_input",
"tokens_output",
"total_tokens",
"cost_usd",
):
assert field in point, f"Missing field: {field}"
# ---------------------------------------------------------------------------
# get_by_agent — pct_of_total sums to 100%
# ---------------------------------------------------------------------------
class TestGetByAgent:
@pytest.mark.asyncio
async def test_pct_of_total_sums_to_100(self) -> None:
rows = [
_make_row(
agent_slug="be-dev-1",
tokens_input=600,
tokens_output=400,
tokens_cache_read=0,
tokens_cache_write=0,
cost_usd=0.05,
),
_make_row(
agent_slug="be-dev-2",
tokens_input=300,
tokens_output=200,
tokens_cache_read=0,
tokens_cache_write=0,
cost_usd=0.02,
),
_make_row(
agent_slug="be-qa",
tokens_input=100,
tokens_output=100,
tokens_cache_read=0,
tokens_cache_write=0,
cost_usd=0.01,
),
]
svc = _service_with_execute(_result_fetchall(rows))
result = await svc.get_by_agent("24h")
total_pct = sum(item["pct_of_total"] for item in result)
assert abs(total_pct - _FULL_PCT) < _PCT_TOL
@pytest.mark.asyncio
async def test_empty_result_returns_empty_list(self) -> None:
svc = _service_with_execute(_result_fetchall([]))
result = await svc.get_by_agent("24h")
assert result == []
@pytest.mark.asyncio
async def test_single_agent_has_100_pct(self) -> None:
rows = [
_make_row(
agent_slug="be-dev-1",
tokens_input=1000,
tokens_output=500,
tokens_cache_read=0,
tokens_cache_write=0,
cost_usd=0.1,
)
]
svc = _service_with_execute(_result_fetchall(rows))
result = await svc.get_by_agent("24h")
assert len(result) == 1
assert result[_ZERO]["pct_of_total"] == _FULL_PCT
@pytest.mark.asyncio
async def test_result_contains_agent_slug_field(self) -> None:
rows = [
_make_row(
agent_slug="be-dev-1",
tokens_input=100,
tokens_output=100,
tokens_cache_read=0,
tokens_cache_write=0,
cost_usd=0.01,
)
]
svc = _service_with_execute(_result_fetchall(rows))
result = await svc.get_by_agent()
assert result[_ZERO]["agent_slug"] == "be-dev-1"
@pytest.mark.asyncio
async def test_cache_tokens_included_in_total_tokens(self) -> None:
"""total_tokens must include cache_read and cache_write.
Without the fix, total would be 500+300=800 (input+output only).
With the fix, total = 500+300+100+100 = 1000.
"""
_cache_read = 100
_cache_write = 100
_expected_total = 500 + 300 + _cache_read + _cache_write # 1000
rows = [
_make_row(
agent_slug="be-dev-1",
tokens_input=500,
tokens_output=300,
tokens_cache_read=_cache_read,
tokens_cache_write=_cache_write,
cost_usd=0.05,
)
]
svc = _service_with_execute(_result_fetchall(rows))
result = await svc.get_by_agent("24h")
assert result[_ZERO]["total_tokens"] == _expected_total
@pytest.mark.asyncio
async def test_pct_of_total_sums_to_100_with_cache_tokens(self) -> None:
"""pct_of_total still sums to 100% when agents have cache tokens."""
rows = [
_make_row(
agent_slug="be-dev-1",
tokens_input=400,
tokens_output=200,
tokens_cache_read=150,
tokens_cache_write=50,
cost_usd=0.05,
),
_make_row(
agent_slug="be-dev-2",
tokens_input=200,
tokens_output=100,
tokens_cache_read=75,
tokens_cache_write=25,
cost_usd=0.02,
),
]
svc = _service_with_execute(_result_fetchall(rows))
result = await svc.get_by_agent("24h")
total_pct = sum(item["pct_of_total"] for item in result)
assert abs(total_pct - _FULL_PCT) < _PCT_TOL
# ---------------------------------------------------------------------------
# get_by_team — pct_of_total sums to 100%
# ---------------------------------------------------------------------------
class TestGetByTeam:
@pytest.mark.asyncio
async def test_pct_of_total_sums_to_100(self) -> None:
rows = [
_make_row(
team="backend",
tokens_input=700,
tokens_output=300,
tokens_cache_read=0,
tokens_cache_write=0,
cost_usd=0.05,
),
_make_row(
team="frontend",
tokens_input=200,
tokens_output=200,
tokens_cache_read=0,
tokens_cache_write=0,
cost_usd=0.02,
),
_make_row(
team="uxui",
tokens_input=100,
tokens_output=100,
tokens_cache_read=0,
tokens_cache_write=0,
cost_usd=0.01,
),
]
svc = _service_with_execute(_result_fetchall(rows))
result = await svc.get_by_team("24h")
total_pct = sum(item["pct_of_total"] for item in result)
assert abs(total_pct - _FULL_PCT) < _PCT_TOL
@pytest.mark.asyncio
async def test_result_contains_team_field(self) -> None:
rows = [
_make_row(
team="backend",
tokens_input=100,
tokens_output=100,
tokens_cache_read=0,
tokens_cache_write=0,
cost_usd=0.01,
)
]
svc = _service_with_execute(_result_fetchall(rows))
result = await svc.get_by_team()
assert result[_ZERO]["team"] == "backend"
@pytest.mark.asyncio
async def test_cache_tokens_included_in_total_tokens(self) -> None:
"""total_tokens must include cache_read and cache_write."""
_cache_read = 200
_cache_write = 100
_expected_total = 700 + 300 + _cache_read + _cache_write # 1300
rows = [
_make_row(
team="backend",
tokens_input=700,
tokens_output=300,
tokens_cache_read=_cache_read,
tokens_cache_write=_cache_write,
cost_usd=0.05,
)
]
svc = _service_with_execute(_result_fetchall(rows))
result = await svc.get_by_team("24h")
assert result[_ZERO]["total_tokens"] == _expected_total
@pytest.mark.asyncio
async def test_pct_of_total_sums_to_100_with_cache_tokens(self) -> None:
"""pct_of_total still sums to 100% when teams have cache tokens."""
rows = [
_make_row(
team="backend",
tokens_input=600,
tokens_output=200,
tokens_cache_read=120,
tokens_cache_write=80,
cost_usd=0.05,
),
_make_row(
team="frontend",
tokens_input=300,
tokens_output=100,
tokens_cache_read=60,
tokens_cache_write=40,
cost_usd=0.02,
),
]
svc = _service_with_execute(_result_fetchall(rows))
result = await svc.get_by_team("24h")
total_pct = sum(item["pct_of_total"] for item in result)
assert abs(total_pct - _FULL_PCT) < _PCT_TOL
# ---------------------------------------------------------------------------
# get_by_model — pct_of_total sums to 100%
# ---------------------------------------------------------------------------
class TestGetByModel:
@pytest.mark.asyncio
async def test_pct_of_total_sums_to_100(self) -> None:
rows = [
_make_row(
model="claude-sonnet-4-6",
tokens_input=600,
tokens_output=600,
tokens_cache_read=0,
tokens_cache_write=0,
cost_usd=0.1,
),
_make_row(
model="claude-haiku-4-5",
tokens_input=300,
tokens_output=300,
tokens_cache_read=0,
tokens_cache_write=0,
cost_usd=0.02,
),
_make_row(
model="claude-opus-4-5",
tokens_input=100,
tokens_output=100,
tokens_cache_read=0,
tokens_cache_write=0,
cost_usd=0.04,
),
]
svc = _service_with_execute(_result_fetchall(rows))
result = await svc.get_by_model("24h")
total_pct = sum(item["pct_of_total"] for item in result)
assert abs(total_pct - _FULL_PCT) < _PCT_TOL
@pytest.mark.asyncio
async def test_result_contains_model_field(self) -> None:
rows = [
_make_row(
model="claude-sonnet-4-6",
tokens_input=100,
tokens_output=100,
tokens_cache_read=0,
tokens_cache_write=0,
cost_usd=0.01,
)
]
svc = _service_with_execute(_result_fetchall(rows))
result = await svc.get_by_model()
assert result[_ZERO]["model"] == "claude-sonnet-4-6"
@pytest.mark.asyncio
async def test_cache_tokens_included_in_total_tokens(self) -> None:
"""total_tokens must include cache_read and cache_write."""
_cache_read = 300
_cache_write = 100
_expected_total = 600 + 600 + _cache_read + _cache_write # 1600
rows = [
_make_row(
model="claude-sonnet-4-6",
tokens_input=600,
tokens_output=600,
tokens_cache_read=_cache_read,
tokens_cache_write=_cache_write,
cost_usd=0.1,
)
]
svc = _service_with_execute(_result_fetchall(rows))
result = await svc.get_by_model("24h")
assert result[_ZERO]["total_tokens"] == _expected_total
@pytest.mark.asyncio
async def test_pct_of_total_sums_to_100_with_cache_tokens(self) -> None:
"""pct_of_total still sums to 100% when models have cache tokens."""
rows = [
_make_row(
model="claude-sonnet-4-6",
tokens_input=500,
tokens_output=500,
tokens_cache_read=200,
tokens_cache_write=100,
cost_usd=0.1,
),
_make_row(
model="claude-haiku-4-5",
tokens_input=250,
tokens_output=250,
tokens_cache_read=100,
tokens_cache_write=50,
cost_usd=0.02,
),
]
svc = _service_with_execute(_result_fetchall(rows))
result = await svc.get_by_model("24h")
total_pct = sum(item["pct_of_total"] for item in result)
assert abs(total_pct - _FULL_PCT) < _PCT_TOL
# ---------------------------------------------------------------------------
# get_projection — formula: projected_monthly = (total_7d / 7) * 30
# ---------------------------------------------------------------------------
class TestGetProjection:
@pytest.mark.asyncio
async def test_projection_formula_30_day_extrapolation(self) -> None:
"""projected_monthly_cost_usd = (total_cost_7d / 7) * 30."""
row = _make_row(total_cost_7d=_COST_7D, session_count=10)
svc = _service_with_execute(_result_one(row))
result = await svc.get_projection()
assert abs(result["projected_monthly_cost_usd"] - _EXPECTED_MONTHLY) < _TOL
@pytest.mark.asyncio
async def test_zero_cost_7d_gives_zero_projection(self) -> None:
row = _make_row(total_cost_7d=0.0, session_count=_ZERO)
svc = _service_with_execute(_result_one(row))
result = await svc.get_projection()
assert result["projected_monthly_cost_usd"] == 0.0
@pytest.mark.asyncio
async def test_avg_daily_cost_equals_total_over_7(self) -> None:
"""avg_daily = total_7d / 7."""
row = _make_row(total_cost_7d=21.0, session_count=5)
svc = _service_with_execute(_result_one(row))
result = await svc.get_projection()
# 21 / 7 = 3.0 avg daily cost
_avg_daily_21 = 3.0
assert abs(result["avg_daily_cost_usd"] - _avg_daily_21) < _TOL
@pytest.mark.asyncio
async def test_projection_contains_required_fields(self) -> None:
row = _make_row(total_cost_7d=7.0, session_count=3)
svc = _service_with_execute(_result_one(row))
result = await svc.get_projection()
for field in (
"total_cost_7d",
"avg_daily_cost_usd",
"projected_monthly_cost_usd",
"basis_days",
):
assert field in result, f"Missing field: {field}"
assert result["basis_days"] == _DAYS_BASIS
# ---------------------------------------------------------------------------
# get_cache_efficiency — hit rate and cost_saved arithmetic
# ---------------------------------------------------------------------------
class TestGetCacheEfficiency:
@pytest.mark.asyncio
async def test_cache_hit_rate_formula(self) -> None:
"""cache_hit_rate = cache_read / (input + cache_read).
400 cache reads out of 400+600 total = 0.4
"""
row = _make_row(
tokens_input=_INPUT_TOKENS,
tokens_output=_ZERO,
tokens_cache_read=_CACHE_READ_TOKENS,
tokens_cache_write=_ZERO,
)
svc = _service_with_execute(_result_one(row))
result = await svc.get_cache_efficiency("24h")
assert abs(result["cache_hit_rate"] - _EXPECTED_HIT_RATE) < _TOL
@pytest.mark.asyncio
async def test_zero_input_tokens_gives_zero_hit_rate(self) -> None:
"""When no input or cache_read tokens, hit rate is 0.0."""
row = _make_row(
tokens_input=_ZERO,
tokens_output=_ZERO,
tokens_cache_read=_ZERO,
tokens_cache_write=_ZERO,
)
svc = _service_with_execute(_result_one(row))
result = await svc.get_cache_efficiency("24h")
assert result["cache_hit_rate"] == 0.0
@pytest.mark.asyncio
async def test_full_cache_hit_gives_rate_of_1(self) -> None:
"""When all input-like tokens are cache reads, hit rate = 1.0."""
_full_rate = 1.0
row = _make_row(
tokens_input=_ZERO,
tokens_output=_ZERO,
tokens_cache_read=1000,
tokens_cache_write=_ZERO,
)
svc = _service_with_execute(_result_one(row))
result = await svc.get_cache_efficiency("24h")
assert abs(result["cache_hit_rate"] - _full_rate) < _TOL
@pytest.mark.asyncio
async def test_cost_saved_arithmetic(self) -> None:
"""cost_saved = cache_read * (full_input_price - cache_read_price) / 1M.
Sonnet baseline: full=$3.00/1M, cache_read=$0.30/1M.
For 1M cache-read tokens: saved = 3.00 - 0.30 = 2.70.
"""
row = _make_row(
tokens_input=_ZERO,
tokens_output=_ZERO,
tokens_cache_read=_M,
tokens_cache_write=_ZERO,
)
svc = _service_with_execute(_result_one(row))
result = await svc.get_cache_efficiency("24h")
assert abs(result["cost_saved_by_cache_usd"] - _EXPECTED_COST_SAVED) < _TOL
@pytest.mark.asyncio
async def test_zero_cache_reads_gives_zero_savings(self) -> None:
row = _make_row(
tokens_input=1000,
tokens_output=500,
tokens_cache_read=_ZERO,
tokens_cache_write=_ZERO,
)
svc = _service_with_execute(_result_one(row))
result = await svc.get_cache_efficiency("24h")
assert result["cost_saved_by_cache_usd"] == 0.0
@pytest.mark.asyncio
async def test_cache_efficiency_contains_required_fields(self) -> None:
row = _make_row(
tokens_input=_ZERO,
tokens_output=_ZERO,
tokens_cache_read=_ZERO,
tokens_cache_write=_ZERO,
)
svc = _service_with_execute(_result_one(row))
result = await svc.get_cache_efficiency("24h")
for field in ("cache_hit_rate", "cost_saved_by_cache_usd"):
assert field in result, f"Missing field: {field}"