mirror of
https://github.com/rennf93/roboco.git
synced 2026-08-03 07:23:24 +02:00
[499f9eb1] Token Usage & Cost Analytics — Full-Stack Instrumentation, Persistence, and Visualization (#90)
* [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
co-authored by
Frontend Developer 1
Backend Developer 1
Renn F
Frontend Developer 1
Backend Developer 1
Renn F
parent
93c6ef8a57
commit
b3057628b0
@@ -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]
|
||||
Reference in New Issue
Block a user