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
roboco/agent_sdk/models.py
+39
View File
@@ -202,3 +202,42 @@ class VerbCircuitStatus(BaseModel):
"next gateway call."
),
)
# =============================================================================
# TOKEN USAGE
# =============================================================================
class TokenReportRequest(BaseModel):
"""Payload for POST /usage/report — reports token usage from a model call.
Counts are *additive*: the SDK accumulates them per session so multiple
report calls sum up correctly across tool invocations.
"""
tokens_input: int = Field(default=0, description="Input / prompt tokens consumed")
tokens_output: int = Field(
default=0, description="Output / completion tokens generated"
)
tokens_cache_read: int = Field(
default=0, description="Prompt-cache read tokens (charged at reduced rate)"
)
tokens_cache_write: int = Field(default=0, description="Prompt-cache write tokens")
class TokenUsageStatus(BaseModel):
"""Current cumulative token usage for this session (GET /usage/status)."""
tokens_input: int = Field(
default=0, description="Total input tokens accumulated this session"
)
tokens_output: int = Field(
default=0, description="Total output tokens accumulated this session"
)
tokens_cache_read: int = Field(
default=0, description="Total cache-read tokens this session"
)
tokens_cache_write: int = Field(
default=0, description="Total cache-write tokens this session"
)
roboco/agent_sdk/server.py
+55
View File
@@ -35,6 +35,8 @@ from roboco.agent_sdk.models import (
SendResponse,
TerminalStatus,
TerminalToolRecordRequest,
TokenReportRequest,
TokenUsageStatus,
VerbAttemptRequest,
VerbCircuitStatus,
)
@@ -420,6 +422,11 @@ class _SessionState:
self.verb_attempts: dict[tuple[str, str | None], deque[float]] = defaultdict(
deque
)
# Cumulative token usage for this session (reported via /usage/report)
self.tokens_input: int = 0
self.tokens_output: int = 0
self.tokens_cache_read: int = 0
self.tokens_cache_write: int = 0
def reset(self) -> None:
self._init_fields()
@@ -680,6 +687,54 @@ def _terminal_snapshot() -> TerminalStatus:
)
# =============================================================================
# TOKEN USAGE REPORTING
# =============================================================================
@app.post("/usage/report", response_model=TokenUsageStatus)
async def usage_report(req: TokenReportRequest) -> TokenUsageStatus:
"""Accumulate token usage counts for the current session.
Called by Claude Code hooks (e.g. PostToolUse) after each API call
to report the tokens consumed by that invocation. Counts are additive
— multiple calls sum up correctly across the session lifetime.
"""
_state.tokens_input += req.tokens_input
_state.tokens_output += req.tokens_output
_state.tokens_cache_read += req.tokens_cache_read
_state.tokens_cache_write += req.tokens_cache_write
logger.debug(
"Token usage reported",
delta_input=req.tokens_input,
delta_output=req.tokens_output,
total_input=_state.tokens_input,
total_output=_state.tokens_output,
)
return _token_usage_snapshot()
@app.get("/usage/status", response_model=TokenUsageStatus)
async def usage_status() -> TokenUsageStatus:
"""Return cumulative token usage totals for the current session.
The orchestrator sweeper calls this endpoint every ~60 s to record
snapshots and to finalize session rows when the container stops.
"""
return _token_usage_snapshot()
def _token_usage_snapshot() -> TokenUsageStatus:
return TokenUsageStatus(
tokens_input=_state.tokens_input,
tokens_output=_state.tokens_output,
tokens_cache_read=_state.tokens_cache_read,
tokens_cache_write=_state.tokens_cache_write,
)
@app.post("/journal/post_mortem")
async def journal_post_mortem(req: PostMortemRequest) -> dict[str, str]:
"""SessionEnd hook submits a post-mortem; we log it and flush to the main API."""
roboco/api/app.py
+8
View File
@@ -36,6 +36,7 @@ from roboco.api.routes.provider import router as provider_router
from roboco.api.routes.sessions import router as sessions_router
from roboco.api.routes.stream import router as stream_router
from roboco.api.routes.tasks import router as tasks_router
from roboco.api.routes.usage import router as usage_router
from roboco.api.routes.v1 import do as do_module
from roboco.api.routes.v1 import flow_auditor as flow_auditor_module
from roboco.api.routes.v1 import flow_board as flow_board_module
@@ -340,6 +341,13 @@ def create_app() -> FastAPI:
tags=["Documentation"],
)
# Token Usage Analytics
app.include_router(
usage_router,
prefix=f"{api_prefix}/usage",
tags=["Usage Analytics"],
)
# API v1 — intent-verb flow endpoints
app.include_router(flow_dev_module.router)
roboco/api/routes/dashboard.py
+14
View File
@@ -21,12 +21,14 @@ from roboco.api.schemas.dashboard import (
CreateReportRequest,
FlagSeverity,
TeamHealth,
UsageSummary,
)
from roboco.models.base import Team
from roboco.models.dashboard import CreateFlagParams
from roboco.services.dashboard import get_dashboard_service
from roboco.services.kanban import get_kanban_service
from roboco.services.metrics import get_metrics_service
from roboco.services.usage import get_usage_service
router = APIRouter()
@@ -277,6 +279,7 @@ async def get_ceo_overview(
- Roadmap progress
"""
service = get_dashboard_service(db)
usage_svc = get_usage_service(db)
health_list = await service.get_team_health_list()
health_status = [
@@ -291,11 +294,22 @@ async def get_ceo_overview(
for h in health_list
]
# Populate usage_summary from daily_usage_rollups for today
try:
today_usage = await usage_svc.get_today_summary()
usage_summary = UsageSummary(
tokens_today=today_usage["tokens_today"],
cost_today_usd=today_usage["cost_today_usd"],
)
except Exception:
usage_summary = UsageSummary(tokens_today=0, cost_today_usd=0.0)
return CEOOverview(
health_status=health_status,
key_metrics=await service.get_key_metrics(),
auditor_alerts=service.get_auditor_alerts(),
roadmap_progress=await service.get_roadmap_progress(),
usage_summary=usage_summary,
)
roboco/api/routes/usage.py
+147
View File
@@ -0,0 +1,147 @@
"""
Token Usage Analytics API
Provides endpoints for querying token usage metrics across agents,
teams, and models. Supports period-based queries (24h, 7d, 30d).
"""
from typing import Annotated, Any, Literal
from fastapi import APIRouter, Query
from roboco.api.deps import DbSession
from roboco.services.usage import get_usage_service
router = APIRouter()
_PeriodType = Literal["24h", "7d", "30d"]
_PeriodQuery = Annotated[
_PeriodType,
Query(description="Time period: 24h, 7d, 30d"),
]
# =============================================================================
# SUMMARY
# =============================================================================
@router.get("/summary")
async def get_usage_summary(
db: DbSession,
period: _PeriodQuery = "24h",
) -> dict[str, Any]:
"""Return aggregated token usage and cost for the given period.
Response includes:
- tokens_input: total prompt tokens consumed
- tokens_output: total completion tokens generated
- total_tokens: sum of all token types
- total_cost_usd: estimated USD cost
- trend_pct: percent change vs. previous equivalent period
"""
svc = get_usage_service(db)
return await svc.get_summary(period)
# =============================================================================
# TIME SERIES
# =============================================================================
@router.get("/time-series")
async def get_usage_time_series(
db: DbSession,
period: _PeriodQuery = "24h",
) -> list[dict[str, Any]]:
"""Return bucketed time-series data points.
- 24h → hourly buckets
- 7d / 30d → daily buckets
Each point has: bucket (ISO timestamp), tokens_input, tokens_output,
total_tokens, cost_usd.
"""
svc = get_usage_service(db)
return await svc.get_time_series(period)
# =============================================================================
# BREAKDOWN ENDPOINTS
# =============================================================================
@router.get("/by-agent")
async def get_usage_by_agent(
db: DbSession,
period: _PeriodQuery = "24h",
) -> list[dict[str, Any]]:
"""Return per-agent token usage with pct_of_total.
pct_of_total fields sum to approximately 100%.
"""
svc = get_usage_service(db)
return await svc.get_by_agent(period)
@router.get("/by-team")
async def get_usage_by_team(
db: DbSession,
period: _PeriodQuery = "24h",
) -> list[dict[str, Any]]:
"""Return per-team token usage with pct_of_total.
pct_of_total fields sum to approximately 100%.
"""
svc = get_usage_service(db)
return await svc.get_by_team(period)
@router.get("/by-model")
async def get_usage_by_model(
db: DbSession,
period: _PeriodQuery = "24h",
) -> list[dict[str, Any]]:
"""Return per-model token usage with pct_of_total.
pct_of_total fields sum to approximately 100%.
"""
svc = get_usage_service(db)
return await svc.get_by_model(period)
# =============================================================================
# PROJECTION
# =============================================================================
@router.get("/projection")
async def get_usage_projection(
db: DbSession,
) -> dict[str, Any]:
"""Return projected monthly cost based on 7-day rolling average.
projected_monthly_cost_usd is computed from avg_daily_cost * 30.
"""
svc = get_usage_service(db)
return await svc.get_projection()
# =============================================================================
# CACHE EFFICIENCY
# =============================================================================
@router.get("/cache-efficiency")
async def get_cache_efficiency(
db: DbSession,
period: _PeriodQuery = "24h",
) -> dict[str, Any]:
"""Return cache hit rate and estimated savings from prompt caching.
- cache_hit_rate: fraction of input-like tokens served from cache
- cost_saved_by_cache_usd: estimated savings vs. full input pricing
"""
svc = get_usage_service(db)
return await svc.get_cache_efficiency(period)
roboco/api/schemas/dashboard.py
+15
View File
@@ -78,6 +78,17 @@ class TeamHealth(BaseModel):
completed_this_week: int
class UsageSummary(BaseModel):
"""Today's token usage summary for the CEO dashboard."""
tokens_today: int = Field(
default=0, description="Total tokens (input + output + cache) used today"
)
cost_today_usd: float = Field(
default=0.0, description="Estimated USD cost for today"
)
class CEOOverview(BaseModel):
"""Complete CEO overview data."""
@@ -85,6 +96,10 @@ class CEOOverview(BaseModel):
key_metrics: dict[str, Any]
auditor_alerts: dict[str, Any]
roadmap_progress: dict[str, Any]
usage_summary: UsageSummary | None = Field(
default=None,
description="Today's token usage and cost from daily_usage_rollups",
)
class CreateFlagRequest(BaseModel):
roboco/billing/__init__.py
+8
View File
@@ -0,0 +1,8 @@
"""Billing utilities for RoboCo.
Provides token-cost calculation for Claude API models.
"""
from roboco.billing.pricing import calculate_cost
__all__ = ["calculate_cost"]
roboco/billing/pricing.py
+106
View File
@@ -0,0 +1,106 @@
"""
Token pricing for Claude API models.
Implements per-model USD cost calculation based on Anthropic's published
pricing. All prices are in USD per 1 million tokens.
Unknown model names return 0.0 without raising so callers don't need to
guard against missing pricing data. Self-hosted Ollama models always
return 0.0 (no API cost) — matched by the ``ollama/`` prefix convention.
"""
from __future__ import annotations
import structlog
logger = structlog.get_logger(__name__)
# ---------------------------------------------------------------------------
# Per-model pricing table
# Format: model_name_fragment → (input_usd_per_1m, output_usd_per_1m,
# cache_read_usd_per_1m, cache_write_usd_per_1m)
#
# Cache read is charged at ~10 % of the input price.
# Cache write is charged at ~25 % of the input price.
#
# Match on *substring* of model name so "claude-opus-4-6" and "opus" both
# resolve to the same tier.
# ---------------------------------------------------------------------------
_PRICING: list[tuple[str, float, float, float, float]] = [
# (fragment, input/1M, output/1M, cache_read/1M, cache_write/1M)
# Opus 4 family
("claude-opus-4", 5.00, 25.00, 0.50, 6.25),
# Sonnet 4 / 3.7 / 3.5 family
("claude-sonnet-4", 3.00, 15.00, 0.30, 0.75),
("claude-3-7-sonnet", 3.00, 15.00, 0.30, 0.75),
("claude-3-5-sonnet", 3.00, 15.00, 0.30, 0.75),
# Haiku family
("claude-haiku-4", 1.00, 5.00, 0.10, 1.25),
("claude-haiku-3-5", 1.00, 5.00, 0.10, 1.25),
("claude-3-5-haiku", 1.00, 5.00, 0.10, 1.25),
("claude-haiku-3", 0.25, 1.25, 0.025, 0.0625),
# Short aliases used in ROLE_MODEL_MAP / MODEL_MAP
("opus", 5.00, 25.00, 0.50, 6.25),
("sonnet", 3.00, 15.00, 0.30, 0.75),
("haiku", 1.00, 5.00, 0.10, 1.25),
]
_MILLION = 1_000_000.0
def calculate_cost(
model: str,
tokens_input: int,
tokens_output: int,
tokens_cache_read: int = 0,
tokens_cache_write: int = 0,
) -> float:
"""Calculate the estimated USD cost for a model invocation.
Matches the model name against the known pricing table using substring
search (longest match wins). Unknown models return 0.0 without raising.
Self-hosted Ollama models (``ollama/`` prefix) always return 0.0.
Args:
model: Model name or short alias (e.g. ``"claude-sonnet-4-6"``,
``"sonnet"``, ``"opus"``).
tokens_input: Number of input tokens (prompt / context).
tokens_output: Number of output tokens (completion).
tokens_cache_read: Prompt-cache read tokens (charged at reduced rate).
tokens_cache_write: Prompt-cache write tokens (charged at reduced rate).
Returns:
Estimated cost in USD as a float. Returns 0.0 for unknown models
rather than raising.
"""
if not model:
return 0.0
lower = model.lower()
# Self-hosted Ollama models have no API cost.
if lower.startswith("ollama/"):
return 0.0
# Find the best (longest fragment) match
best_fragment_len = 0
best_prices: tuple[float, float, float, float] | None = None
for fragment, inp_price, out_price, cr_price, cw_price in _PRICING:
if fragment in lower and len(fragment) > best_fragment_len:
best_fragment_len = len(fragment)
best_prices = (inp_price, out_price, cr_price, cw_price)
if best_prices is None:
logger.warning("No pricing data found for model", model=model)
return 0.0
inp_price, out_price, cr_price, cw_price = best_prices
cost = (
tokens_input * inp_price / _MILLION
+ tokens_output * out_price / _MILLION
+ tokens_cache_read * cr_price / _MILLION
+ tokens_cache_write * cw_price / _MILLION
)
return round(cost, 8)
roboco/db/tables.py
+144
View File
@@ -11,7 +11,9 @@ from uuid import uuid4
from sqlalchemy import (
JSON,
BigInteger,
Boolean,
Date,
DateTime,
Enum,
Float,
@@ -1823,6 +1825,148 @@ class GatewayTriggerTable(Base):
)
# =============================================================================
# TOKEN USAGE TABLES
# =============================================================================
class AgentSpawnSessionTable(Base):
"""Records each agent container spawn lifecycle.
Opened when the orchestrator successfully starts a container; closed
(ended_at set) when stop_agent() finishes. Final token counts are
accumulated from the agent SDK's /usage/status endpoint.
"""
__tablename__ = "agent_spawn_sessions"
id: Mapped[UUID] = mapped_column(
UUID(as_uuid=True), primary_key=True, default=uuid4
)
agent_slug: Mapped[str] = mapped_column(String(100), nullable=False)
team: Mapped[str] = mapped_column(String(50), nullable=False)
role: Mapped[str] = mapped_column(String(50), nullable=False)
model: Mapped[str] = mapped_column(String(100), nullable=False)
task_id: Mapped[str | None] = mapped_column(String(36), nullable=True)
started_at: Mapped[datetime] = mapped_column(
DateTime(timezone=True),
default=lambda: datetime.now(UTC),
nullable=False,
)
ended_at: Mapped[datetime | None] = mapped_column(
DateTime(timezone=True), nullable=True
)
# BIGINT — token counts can exceed INT32 for long sessions
tokens_input: Mapped[int] = mapped_column(BigInteger, nullable=False, default=0)
tokens_output: Mapped[int] = mapped_column(BigInteger, nullable=False, default=0)
tokens_cache_read: Mapped[int] = mapped_column(
BigInteger, nullable=False, default=0
)
tokens_cache_write: Mapped[int] = mapped_column(
BigInteger, nullable=False, default=0
)
exit_reason: Mapped[str | None] = mapped_column(String(100), nullable=True)
estimated_cost_usd: Mapped[float | None] = mapped_column(Float, nullable=True)
# Relationship to snapshots (backref for convenience)
snapshots: Mapped[list["TokenUsageSnapshotTable"]] = relationship(
"TokenUsageSnapshotTable",
back_populates="session",
cascade="all, delete-orphan",
)
__table_args__ = (
Index("ix_agent_spawn_sessions_agent_slug", "agent_slug"),
Index("ix_agent_spawn_sessions_started_at", "started_at"),
Index("ix_agent_spawn_sessions_ended_at", "ended_at"),
Index("ix_agent_spawn_sessions_team", "team"),
)
class TokenUsageSnapshotTable(Base):
"""Periodic (every ~60 s) snapshot of cumulative token usage for an
active agent_spawn_session.
The sweeper inserts one row per active agent per sweep cycle when
token counts are non-zero. Snapshots allow tracking how token usage
grows over session lifetime.
"""
__tablename__ = "token_usage_snapshots"
id: Mapped[UUID] = mapped_column(
UUID(as_uuid=True), primary_key=True, default=uuid4
)
agent_spawn_session_id: Mapped[UUID] = mapped_column(
UUID(as_uuid=True),
ForeignKey("agent_spawn_sessions.id", ondelete="CASCADE"),
nullable=False,
)
snapshotted_at: Mapped[datetime] = mapped_column(
DateTime(timezone=True),
default=lambda: datetime.now(UTC),
nullable=False,
)
tokens_input: Mapped[int] = mapped_column(BigInteger, nullable=False, default=0)
tokens_output: Mapped[int] = mapped_column(BigInteger, nullable=False, default=0)
tokens_cache_read: Mapped[int] = mapped_column(
BigInteger, nullable=False, default=0
)
tokens_cache_write: Mapped[int] = mapped_column(
BigInteger, nullable=False, default=0
)
session: Mapped["AgentSpawnSessionTable"] = relationship(
"AgentSpawnSessionTable", back_populates="snapshots"
)
__table_args__ = (
Index("ix_token_usage_snapshots_session_id", "agent_spawn_session_id"),
Index("ix_token_usage_snapshots_snapshotted_at", "snapshotted_at"),
)
class DailyUsageRollupTable(Base):
"""Pre-aggregated daily token usage per (date, agent_slug, team, model).
Populated by the orchestrator sweeper via an upsert query over
closed agent_spawn_sessions. Unique constraint on the natural key
enables ON CONFLICT DO UPDATE so the sweep is idempotent.
"""
__tablename__ = "daily_usage_rollups"
id: Mapped[UUID] = mapped_column(
UUID(as_uuid=True), primary_key=True, default=uuid4
)
date: Mapped[Any] = mapped_column(Date, nullable=False) # datetime.date
agent_slug: Mapped[str] = mapped_column(String(100), nullable=False)
team: Mapped[str] = mapped_column(String(50), nullable=False)
model: Mapped[str] = mapped_column(String(100), nullable=False)
tokens_input: Mapped[int] = mapped_column(BigInteger, nullable=False, default=0)
tokens_output: Mapped[int] = mapped_column(BigInteger, nullable=False, default=0)
tokens_cache_read: Mapped[int] = mapped_column(
BigInteger, nullable=False, default=0
)
tokens_cache_write: Mapped[int] = mapped_column(
BigInteger, nullable=False, default=0
)
total_cost_usd: Mapped[float] = mapped_column(Float, nullable=False, default=0.0)
session_count: Mapped[int] = mapped_column(Integer, nullable=False, default=0)
__table_args__ = (
UniqueConstraint(
"date",
"agent_slug",
"team",
"model",
name="uq_daily_rollup_date_agent_team_model",
),
Index("ix_daily_rollups_date", "date"),
Index("ix_daily_rollups_agent_slug", "agent_slug"),
)
# =============================================================================
# PROMPTER TABLES
# =============================================================================
roboco/events/stream_bus.py
+2 -2
View File
@@ -376,9 +376,9 @@ class StreamEventBus:
recovered = 0
for msg in pending_details:
if msg["time_since_delivered"] >= idle_time_ms:
if int(msg["time_since_delivered"]) >= idle_time_ms:
recovered += await self._claim_and_handle(
stream, msg["message_id"], idle_time_ms
stream, str(msg["message_id"]), idle_time_ms
)
return recovered
roboco/models/runtime.py
+4
View File
@@ -71,6 +71,10 @@ class AgentInstance:
error_count: int = 0
waiting_for: str | None = None # For WAITING_LONG state
waiting_context: dict[str, Any] = field(default_factory=dict)
# UUID of the agent_spawn_sessions row created at spawn time.
# Used by _finalize_spawn_session for a direct-by-id lookup instead of a
# fragile (agent_slug, ended_at IS NULL) query.
usage_session_id: UUID | None = None
def __post_init__(self) -> None:
if not self.id:
roboco/runtime/orchestrator.py
+437 -3
View File
@@ -27,6 +27,7 @@ import httpx
if TYPE_CHECKING:
from collections.abc import Callable, Coroutine
from uuid import UUID
from roboco.services.llm import AgentRoute
from roboco.services.task import TaskService
@@ -69,6 +70,11 @@ AgentConfig = OrchestratorAgentConfig
AGENT_NETWORK = "roboco_default"
AGENT_BASE_IMAGE = "roboco-agent-base"
# Port on which each agent's Claude Code SDK server listens inside its container.
# Referenced by write-hooks (_finalize_spawn_session, _sweep_token_snapshots,
# _sweep_budget_exceeded) to build the SDK health/usage URL.
SDK_PORT: int = 9000
# The intake (prompter) agent: a single seeded, board-adjacent interviewer.
# Unlike delivery agents it is never dispatched and runs ONE persistent
# container at a time (single CEO → one live chat). See the INTAKE section
@@ -1382,6 +1388,13 @@ class AgentOrchestrator:
"model": config.model,
},
)
# Record a token-usage session row in the DB and bind its UUID to
# the instance so _finalize_spawn_session can look it up directly.
usage_session_id = await self._record_spawn_session(config, task_id)
if usage_session_id is not None:
instance.usage_session_id = usage_session_id
return instance
except Exception as e:
instance.state = AgentState.OFFLINE
@@ -2868,8 +2881,28 @@ class AgentOrchestrator:
# AGENT STOPPING
# =========================================================================
async def stop_agent(self, agent_id: str, graceful: bool = True) -> None:
"""Stop an agent container."""
async def stop_agent(
self,
agent_id: str,
graceful: bool = True,
exit_reason: str = "stopped",
) -> None:
"""Stop an agent container.
Finalization (the HTTP call to the agent SDK's /usage/status endpoint)
is performed BEFORE acquiring self._lock so that the network I/O does
not block other operations that need the lock.
"""
# Finalize the spawn-session row before the container is removed so we
# can still query the SDK's /usage/status endpoint. This must happen
# outside self._lock — the HTTP round-trip would otherwise hold the
# lock for the full network timeout.
instance = self._instances.get(agent_id)
if instance is None:
return
if instance.container_id:
await self._finalize_spawn_session(agent_id, exit_reason=exit_reason)
async with self._lock:
if agent_id not in self._instances:
return
@@ -3013,6 +3046,402 @@ class AgentOrchestrator:
error=str(e),
)
# =========================================================================
# TOKEN USAGE INSTRUMENTATION
# =========================================================================
async def _record_spawn_session(
self,
config: "OrchestratorAgentConfig",
task_id: str | None,
) -> "UUID | None":
"""Insert a row into agent_spawn_sessions after a successful spawn.
Returns the UUID of the created row so the caller can store it on
the AgentInstance for later direct-by-id lookup in
_finalize_spawn_session. Returns None when the insert fails; a
missing session row must never block the spawn path.
"""
try:
from uuid import uuid4 as _uuid4
from roboco.db.base import get_session_factory
from roboco.db.tables import AgentSpawnSessionTable
agent_slug = config.agent_id
team = get_agent_team(agent_slug) or "backend"
role = get_agent_role(agent_slug) or "developer"
session_id = _uuid4()
session_factory = get_session_factory()
async with session_factory() as db:
row = AgentSpawnSessionTable(
id=session_id,
agent_slug=agent_slug,
team=team,
role=role,
model=config.model or "unknown",
task_id=task_id,
started_at=datetime.now(UTC),
)
db.add(row)
await db.commit()
logger.debug(
"Spawn session recorded",
agent_slug=agent_slug,
session_id=str(session_id),
task_id=task_id,
)
return session_id
except Exception as exc:
logger.warning(
"Failed to record spawn session",
agent_slug=config.agent_id,
error=str(exc),
)
return None
async def _finalize_spawn_session(
self,
agent_id: str,
exit_reason: str = "stopped",
) -> None:
"""Close the open agent_spawn_sessions row for this agent.
Fetches final token counts from the agent SDK's /usage/status endpoint,
calculates cost via pricing module, then updates the DB row with
ended_at, token totals, exit_reason, and estimated_cost_usd.
Errors are caught and logged — finalization must never block stop_agent.
"""
try:
from roboco.billing.pricing import calculate_cost
from roboco.db.base import get_session_factory
from roboco.db.tables import AgentSpawnSessionTable
# Fetch final token counts from the agent's SDK
sdk_url = f"http://roboco-agent-{agent_id}:{SDK_PORT}/usage/status"
tokens_input = 0
tokens_output = 0
tokens_cache_read = 0
tokens_cache_write = 0
model = "unknown"
try:
async with httpx.AsyncClient(timeout=3.0) as client:
resp = await client.get(sdk_url)
if resp.status_code == http_status.HTTP_200_OK:
data = resp.json()
tokens_input = data.get("tokens_input", 0)
tokens_output = data.get("tokens_output", 0)
tokens_cache_read = data.get("tokens_cache_read", 0)
tokens_cache_write = data.get("tokens_cache_write", 0)
except Exception as sdk_exc:
logger.debug(
"Could not fetch final token counts from SDK",
agent_id=agent_id,
error=str(sdk_exc),
)
# Look up the model and usage_session_id from the running instance config.
instance = self._instances.get(agent_id)
if instance and instance.config:
model = instance.config.model or "unknown"
usage_session_id = instance.usage_session_id if instance else None
cost = calculate_cost(
model=model,
tokens_input=tokens_input,
tokens_output=tokens_output,
tokens_cache_read=tokens_cache_read,
tokens_cache_write=tokens_cache_write,
)
session_factory = get_session_factory()
async with session_factory() as db:
from sqlalchemy import select, update
# Prefer a direct lookup by the session UUID captured at spawn
# time; fall back to the (agent_slug, ended_at IS NULL) query
# for instances that pre-date the usage_session_id field.
if usage_session_id is not None:
result = await db.execute(
select(AgentSpawnSessionTable).where(
AgentSpawnSessionTable.id == usage_session_id
)
)
else:
result = await db.execute(
select(AgentSpawnSessionTable)
.where(
AgentSpawnSessionTable.agent_slug == agent_id,
AgentSpawnSessionTable.ended_at.is_(None),
)
.order_by(AgentSpawnSessionTable.started_at.desc())
.limit(1)
)
session_row = result.scalar_one_or_none()
if session_row is not None:
await db.execute(
update(AgentSpawnSessionTable)
.where(AgentSpawnSessionTable.id == session_row.id)
.values(
ended_at=datetime.now(UTC),
tokens_input=tokens_input,
tokens_output=tokens_output,
tokens_cache_read=tokens_cache_read,
tokens_cache_write=tokens_cache_write,
exit_reason=exit_reason,
estimated_cost_usd=cost,
)
)
await db.commit()
logger.debug(
"Spawn session finalized",
agent_id=agent_id,
session_id=str(session_row.id),
tokens_input=tokens_input,
tokens_output=tokens_output,
estimated_cost_usd=cost,
)
except Exception as exc:
logger.warning(
"Failed to finalize spawn session",
agent_id=agent_id,
error=str(exc),
)
async def _sweep_token_snapshots(self) -> None:
"""Write a token_usage_snapshots row for each active agent with non-zero tokens.
Called from _run_sweep() every ~60 s. Also updates the cumulative
token counts on the open agent_spawn_sessions row so the DB reflects
current progress without waiting for session close.
Errors per-agent are caught so one bad agent doesn't abort the whole sweep.
"""
if not self._instances:
return
try:
from roboco.db.base import get_session_factory
from roboco.db.tables import AgentSpawnSessionTable, TokenUsageSnapshotTable
except ImportError:
return
session_factory = get_session_factory()
async with httpx.AsyncClient(timeout=3.0) as client:
for agent_id, instance in list(self._instances.items()):
if instance.state not in (
AgentState.ACTIVE,
AgentState.WAITING_SHORT,
):
continue
sdk_url = f"http://roboco-agent-{agent_id}:{SDK_PORT}/usage/status"
try:
resp = await client.get(sdk_url)
if resp.status_code != http_status.HTTP_200_OK:
continue
data = resp.json()
tokens_input = data.get("tokens_input", 0)
tokens_output = data.get("tokens_output", 0)
tokens_cache_read = data.get("tokens_cache_read", 0)
tokens_cache_write = data.get("tokens_cache_write", 0)
# Skip agents with no token usage yet
total = (
tokens_input
+ tokens_output
+ tokens_cache_read
+ tokens_cache_write
)
if total == 0:
continue
async with session_factory() as db:
from sqlalchemy import select, update
# Prefer a direct lookup by the session UUID captured at
# spawn time; fall back to the agent_slug heuristic for
# instances that pre-date the usage_session_id field.
if instance.usage_session_id is not None:
result = await db.execute(
select(AgentSpawnSessionTable).where(
AgentSpawnSessionTable.id
== instance.usage_session_id
)
)
else:
result = await db.execute(
select(AgentSpawnSessionTable)
.where(
AgentSpawnSessionTable.agent_slug == agent_id,
AgentSpawnSessionTable.ended_at.is_(None),
)
.order_by(AgentSpawnSessionTable.started_at.desc())
.limit(1)
)
session_row = result.scalar_one_or_none()
if session_row is None:
continue
# Insert snapshot
from uuid import uuid4 as _uuid4
snapshot = TokenUsageSnapshotTable(
id=_uuid4(),
agent_spawn_session_id=session_row.id,
snapshotted_at=datetime.now(UTC),
tokens_input=tokens_input,
tokens_output=tokens_output,
tokens_cache_read=tokens_cache_read,
tokens_cache_write=tokens_cache_write,
)
db.add(snapshot)
# Update cumulative totals on the session row
await db.execute(
update(AgentSpawnSessionTable)
.where(AgentSpawnSessionTable.id == session_row.id)
.values(
tokens_input=tokens_input,
tokens_output=tokens_output,
tokens_cache_read=tokens_cache_read,
tokens_cache_write=tokens_cache_write,
)
)
await db.commit()
except Exception as agent_exc:
logger.debug(
"Token snapshot failed for agent",
agent_id=agent_id,
error=str(agent_exc),
)
async def _sweep_daily_rollup(self) -> None:
"""Upsert daily_usage_rollups from closed agent_spawn_sessions.
Groups ended sessions by (date, agent_slug, team, model) and sums
their token counts + cost. Uses a Python-side upsert to stay
compatible with asyncpg / SQLAlchemy without raw INSERT ... ON CONFLICT
dialect-specific SQL.
Errors are caught so a bad rollup doesn't abort the sweeper.
"""
try:
from roboco.db.base import get_session_factory
from roboco.db.tables import AgentSpawnSessionTable, DailyUsageRollupTable
except ImportError:
return
try:
from uuid import uuid4 as _uuid4
from sqlalchemy import func, select
session_factory = get_session_factory()
async with session_factory() as db:
# Aggregate closed sessions by (date, agent_slug, team, model).
# Limit to the last 7 days to avoid re-aggregating all-time
# history on every sweep — older days are already stable.
rollup_window_start = datetime.now(UTC) - timedelta(days=7)
result = await db.execute(
select(
func.date(AgentSpawnSessionTable.started_at).label("date"),
AgentSpawnSessionTable.agent_slug,
AgentSpawnSessionTable.team,
AgentSpawnSessionTable.model,
func.sum(AgentSpawnSessionTable.tokens_input).label(
"tokens_input"
),
func.sum(AgentSpawnSessionTable.tokens_output).label(
"tokens_output"
),
func.sum(AgentSpawnSessionTable.tokens_cache_read).label(
"tokens_cache_read"
),
func.sum(AgentSpawnSessionTable.tokens_cache_write).label(
"tokens_cache_write"
),
func.sum(AgentSpawnSessionTable.estimated_cost_usd).label(
"total_cost_usd"
),
func.count(AgentSpawnSessionTable.id).label("session_count"),
)
.where(
AgentSpawnSessionTable.ended_at.isnot(None),
AgentSpawnSessionTable.started_at >= rollup_window_start,
)
.group_by(
func.date(AgentSpawnSessionTable.started_at),
AgentSpawnSessionTable.agent_slug,
AgentSpawnSessionTable.team,
AgentSpawnSessionTable.model,
)
)
rows = result.fetchall()
for row in rows:
date_val = row.date
agent_slug = row.agent_slug
team = row.team
model = row.model
# Look for existing rollup row
existing_result = await db.execute(
select(DailyUsageRollupTable).where(
DailyUsageRollupTable.date == date_val,
DailyUsageRollupTable.agent_slug == agent_slug,
DailyUsageRollupTable.team == team,
DailyUsageRollupTable.model == model,
)
)
existing = existing_result.scalar_one_or_none()
tokens_input = int(row.tokens_input or 0)
tokens_output = int(row.tokens_output or 0)
tokens_cache_read = int(row.tokens_cache_read or 0)
tokens_cache_write = int(row.tokens_cache_write or 0)
total_cost = float(row.total_cost_usd or 0.0)
session_count = int(row.session_count or 0)
if existing is not None:
from sqlalchemy import update
await db.execute(
update(DailyUsageRollupTable)
.where(DailyUsageRollupTable.id == existing.id)
.values(
tokens_input=tokens_input,
tokens_output=tokens_output,
tokens_cache_read=tokens_cache_read,
tokens_cache_write=tokens_cache_write,
total_cost_usd=total_cost,
session_count=session_count,
)
)
else:
new_row = DailyUsageRollupTable(
id=_uuid4(),
date=date_val,
agent_slug=agent_slug,
team=team,
model=model,
tokens_input=tokens_input,
tokens_output=tokens_output,
tokens_cache_read=tokens_cache_read,
tokens_cache_write=tokens_cache_write,
total_cost_usd=total_cost,
session_count=session_count,
)
db.add(new_row)
await db.commit()
logger.debug("Daily usage rollup complete", rows_processed=len(rows))
except Exception as exc:
logger.warning("Daily usage rollup failed", error=str(exc))
async def restore_waiting_records(self) -> int:
"""Load persisted waiting records into memory on orchestrator start.
@@ -3212,6 +3641,11 @@ Start by:
# same session.
await self._sweep_budget_exceeded()
# Token-usage instrumentation: snapshot active agents and roll up
# closed sessions into the daily aggregation table.
await self._sweep_token_snapshots()
await self._sweep_daily_rollup()
@staticmethod
async def _fetch_budget_status(
client: httpx.AsyncClient, url: str, agent_id: str
@@ -3263,7 +3697,7 @@ Start by:
AgentState.WAITING_SHORT,
):
continue
url = f"http://roboco-agent-{agent_id}:9000/budget/status"
url = f"http://roboco-agent-{agent_id}:{SDK_PORT}/budget/status"
data = await self._fetch_budget_status(client, url, agent_id)
if data is None or not data.get("halt"):
continue
roboco/services/usage.py
+554
View File
@@ -0,0 +1,554 @@
"""
Usage Analytics Service
Provides token usage analytics over agent_spawn_sessions and
daily_usage_rollups tables. Supports period-based queries (24h, 7d, 30d)
and aggregation by agent, team, and model.
"""
from __future__ import annotations
from datetime import UTC, datetime, timedelta
from typing import TYPE_CHECKING, Any
from sqlalchemy import func, select
if TYPE_CHECKING:
from sqlalchemy.ext.asyncio import AsyncSession
from roboco.db.tables import AgentSpawnSessionTable, DailyUsageRollupTable
from roboco.services.base import BaseService
def _parse_period(period: str) -> tuple[datetime, int]:
"""Parse period string into (start_dt, hours).
Accepts '24h', '7d', '30d'. Defaults to 24h for unknown values.
Returns (start_datetime_utc, total_hours).
"""
now = datetime.now(UTC)
if period == "7d":
return now - timedelta(days=7), 7 * 24
if period == "30d":
return now - timedelta(days=30), 30 * 24
# default 24h
return now - timedelta(hours=24), 24
class UsageService(BaseService):
"""Analytics service for token usage data."""
# =========================================================================
# SUMMARY
# =========================================================================
async def get_summary(self, period: str = "24h") -> dict[str, Any]:
"""Return aggregated token and cost totals for the given period.
Queries daily_usage_rollups for whole-day periods; falls back to
agent_spawn_sessions for sub-day precision.
Returns dict with: tokens_input, tokens_output, total_tokens,
total_cost_usd, trend_pct.
"""
start_dt, hours = _parse_period(period)
# Current period totals from closed sessions
result = await self.session.execute(
select(
func.coalesce(func.sum(AgentSpawnSessionTable.tokens_input), 0).label(
"tokens_input"
),
func.coalesce(func.sum(AgentSpawnSessionTable.tokens_output), 0).label(
"tokens_output"
),
func.coalesce(
func.sum(AgentSpawnSessionTable.tokens_cache_read), 0
).label("tokens_cache_read"),
func.coalesce(
func.sum(AgentSpawnSessionTable.tokens_cache_write), 0
).label("tokens_cache_write"),
func.coalesce(
func.sum(AgentSpawnSessionTable.estimated_cost_usd), 0.0
).label("total_cost_usd"),
).where(
AgentSpawnSessionTable.started_at >= start_dt,
AgentSpawnSessionTable.ended_at.isnot(None),
)
)
row = result.one()
tokens_input = int(row.tokens_input or 0)
tokens_output = int(row.tokens_output or 0)
total_cost = float(row.total_cost_usd or 0.0)
total_tokens = (
tokens_input
+ tokens_output
+ int(row.tokens_cache_read or 0)
+ int(row.tokens_cache_write or 0)
)
# Previous period for trend calculation.
# Sum all 4 token columns so the comparison is consistent with the
# current-period total_tokens (which also sums all 4 columns).
prev_start = start_dt - timedelta(hours=hours)
prev_result = await self.session.execute(
select(
func.coalesce(
func.sum(
AgentSpawnSessionTable.tokens_input
+ AgentSpawnSessionTable.tokens_output
+ AgentSpawnSessionTable.tokens_cache_read
+ AgentSpawnSessionTable.tokens_cache_write
),
0,
).label("total")
).where(
AgentSpawnSessionTable.started_at >= prev_start,
AgentSpawnSessionTable.started_at < start_dt,
AgentSpawnSessionTable.ended_at.isnot(None),
)
)
prev_row = prev_result.one()
prev_total = int(prev_row.total or 0)
if prev_total > 0:
trend_pct = round((total_tokens - prev_total) / prev_total * 100, 1)
elif total_tokens > 0:
trend_pct = 100.0
else:
trend_pct = 0.0
return {
"tokens_input": tokens_input,
"tokens_output": tokens_output,
"total_tokens": total_tokens,
"total_cost_usd": round(total_cost, 6),
"trend_pct": trend_pct,
"period": period,
}
# =========================================================================
# TIME SERIES
# =========================================================================
async def get_time_series(self, period: str = "24h") -> list[dict[str, Any]]:
"""Return bucketed time-series data points.
- 24h → hourly buckets
- 7d / 30d → daily buckets
Each point has: bucket (ISO string), tokens_input, tokens_output,
total_tokens, cost_usd.
total_tokens includes all 4 token types (input + output + cache_read +
cache_write) so it is consistent with get_summary()'s total_tokens
field — the two sums must match for the same period.
"""
start_dt, _hours = _parse_period(period)
if period == "24h":
# Hourly buckets
trunc_fn = func.date_trunc("hour", AgentSpawnSessionTable.started_at)
else:
# Daily buckets
trunc_fn = func.date_trunc("day", AgentSpawnSessionTable.started_at)
result = await self.session.execute(
select(
trunc_fn.label("bucket"),
func.coalesce(func.sum(AgentSpawnSessionTable.tokens_input), 0).label(
"tokens_input"
),
func.coalesce(func.sum(AgentSpawnSessionTable.tokens_output), 0).label(
"tokens_output"
),
func.coalesce(
func.sum(AgentSpawnSessionTable.tokens_cache_read), 0
).label("tokens_cache_read"),
func.coalesce(
func.sum(AgentSpawnSessionTable.tokens_cache_write), 0
).label("tokens_cache_write"),
func.coalesce(
func.sum(AgentSpawnSessionTable.estimated_cost_usd), 0.0
).label("cost_usd"),
)
.where(
AgentSpawnSessionTable.started_at >= start_dt,
AgentSpawnSessionTable.ended_at.isnot(None),
)
.group_by(trunc_fn)
.order_by(trunc_fn)
)
rows = result.fetchall()
points = []
for r in rows:
ti = int(r.tokens_input or 0)
to_ = int(r.tokens_output or 0)
tcr = int(r.tokens_cache_read or 0)
tcw = int(r.tokens_cache_write or 0)
points.append(
{
"bucket": r.bucket.isoformat() if r.bucket else None,
"tokens_input": ti,
"tokens_output": to_,
"total_tokens": ti + to_ + tcr + tcw,
"cost_usd": round(float(r.cost_usd or 0.0), 6),
}
)
return points
# =========================================================================
# BY-AGENT
# =========================================================================
async def get_by_agent(self, period: str = "24h") -> list[dict[str, Any]]:
"""Return per-agent token usage with pct_of_total."""
start_dt, _ = _parse_period(period)
result = await self.session.execute(
select(
AgentSpawnSessionTable.agent_slug,
func.coalesce(func.sum(AgentSpawnSessionTable.tokens_input), 0).label(
"tokens_input"
),
func.coalesce(func.sum(AgentSpawnSessionTable.tokens_output), 0).label(
"tokens_output"
),
func.coalesce(
func.sum(AgentSpawnSessionTable.tokens_cache_read), 0
).label("tokens_cache_read"),
func.coalesce(
func.sum(AgentSpawnSessionTable.tokens_cache_write), 0
).label("tokens_cache_write"),
func.coalesce(
func.sum(AgentSpawnSessionTable.estimated_cost_usd), 0.0
).label("cost_usd"),
)
.where(
AgentSpawnSessionTable.started_at >= start_dt,
AgentSpawnSessionTable.ended_at.isnot(None),
)
.group_by(AgentSpawnSessionTable.agent_slug)
.order_by(
func.sum(
AgentSpawnSessionTable.tokens_input
+ AgentSpawnSessionTable.tokens_output
).desc()
)
)
rows = result.fetchall()
grand_total = sum(
int(r.tokens_input or 0)
+ int(r.tokens_output or 0)
+ int(r.tokens_cache_read or 0)
+ int(r.tokens_cache_write or 0)
for r in rows
)
items = []
for r in rows:
ti = int(r.tokens_input or 0)
to_ = int(r.tokens_output or 0)
tcr = int(r.tokens_cache_read or 0)
tcw = int(r.tokens_cache_write or 0)
total = ti + to_ + tcr + tcw
items.append(
{
"agent_slug": r.agent_slug,
"tokens_input": ti,
"tokens_output": to_,
"total_tokens": total,
"cost_usd": round(float(r.cost_usd or 0.0), 6),
"pct_of_total": round(total / grand_total * 100, 2)
if grand_total > 0
else 0.0,
}
)
return items
# =========================================================================
# BY-TEAM
# =========================================================================
async def get_by_team(self, period: str = "24h") -> list[dict[str, Any]]:
"""Return per-team token usage with pct_of_total."""
start_dt, _ = _parse_period(period)
result = await self.session.execute(
select(
AgentSpawnSessionTable.team,
func.coalesce(func.sum(AgentSpawnSessionTable.tokens_input), 0).label(
"tokens_input"
),
func.coalesce(func.sum(AgentSpawnSessionTable.tokens_output), 0).label(
"tokens_output"
),
func.coalesce(
func.sum(AgentSpawnSessionTable.tokens_cache_read), 0
).label("tokens_cache_read"),
func.coalesce(
func.sum(AgentSpawnSessionTable.tokens_cache_write), 0
).label("tokens_cache_write"),
func.coalesce(
func.sum(AgentSpawnSessionTable.estimated_cost_usd), 0.0
).label("cost_usd"),
)
.where(
AgentSpawnSessionTable.started_at >= start_dt,
AgentSpawnSessionTable.ended_at.isnot(None),
)
.group_by(AgentSpawnSessionTable.team)
.order_by(
func.sum(
AgentSpawnSessionTable.tokens_input
+ AgentSpawnSessionTable.tokens_output
).desc()
)
)
rows = result.fetchall()
grand_total = sum(
int(r.tokens_input or 0)
+ int(r.tokens_output or 0)
+ int(r.tokens_cache_read or 0)
+ int(r.tokens_cache_write or 0)
for r in rows
)
items = []
for r in rows:
ti = int(r.tokens_input or 0)
to_ = int(r.tokens_output or 0)
tcr = int(r.tokens_cache_read or 0)
tcw = int(r.tokens_cache_write or 0)
total = ti + to_ + tcr + tcw
items.append(
{
"team": r.team,
"tokens_input": ti,
"tokens_output": to_,
"total_tokens": total,
"cost_usd": round(float(r.cost_usd or 0.0), 6),
"pct_of_total": round(total / grand_total * 100, 2)
if grand_total > 0
else 0.0,
}
)
return items
# =========================================================================
# BY-MODEL
# =========================================================================
async def get_by_model(self, period: str = "24h") -> list[dict[str, Any]]:
"""Return per-model token usage with pct_of_total."""
start_dt, _ = _parse_period(period)
result = await self.session.execute(
select(
AgentSpawnSessionTable.model,
func.coalesce(func.sum(AgentSpawnSessionTable.tokens_input), 0).label(
"tokens_input"
),
func.coalesce(func.sum(AgentSpawnSessionTable.tokens_output), 0).label(
"tokens_output"
),
func.coalesce(
func.sum(AgentSpawnSessionTable.tokens_cache_read), 0
).label("tokens_cache_read"),
func.coalesce(
func.sum(AgentSpawnSessionTable.tokens_cache_write), 0
).label("tokens_cache_write"),
func.coalesce(
func.sum(AgentSpawnSessionTable.estimated_cost_usd), 0.0
).label("cost_usd"),
)
.where(
AgentSpawnSessionTable.started_at >= start_dt,
AgentSpawnSessionTable.ended_at.isnot(None),
)
.group_by(AgentSpawnSessionTable.model)
.order_by(
func.sum(
AgentSpawnSessionTable.tokens_input
+ AgentSpawnSessionTable.tokens_output
).desc()
)
)
rows = result.fetchall()
grand_total = sum(
int(r.tokens_input or 0)
+ int(r.tokens_output or 0)
+ int(r.tokens_cache_read or 0)
+ int(r.tokens_cache_write or 0)
for r in rows
)
items = []
for r in rows:
ti = int(r.tokens_input or 0)
to_ = int(r.tokens_output or 0)
tcr = int(r.tokens_cache_read or 0)
tcw = int(r.tokens_cache_write or 0)
total = ti + to_ + tcr + tcw
items.append(
{
"model": r.model,
"tokens_input": ti,
"tokens_output": to_,
"total_tokens": total,
"cost_usd": round(float(r.cost_usd or 0.0), 6),
"pct_of_total": round(total / grand_total * 100, 2)
if grand_total > 0
else 0.0,
}
)
return items
# =========================================================================
# PROJECTION
# =========================================================================
async def get_projection(self) -> dict[str, Any]:
"""Return projected monthly cost based on 7-day rolling average.
Computes the average daily cost over the last 7 days and extrapolates
to 30 days.
"""
seven_days_ago = datetime.now(UTC) - timedelta(days=7)
result = await self.session.execute(
select(
func.coalesce(
func.sum(AgentSpawnSessionTable.estimated_cost_usd), 0.0
).label("total_cost_7d"),
func.coalesce(func.count(AgentSpawnSessionTable.id), 0).label(
"session_count"
),
).where(
AgentSpawnSessionTable.started_at >= seven_days_ago,
AgentSpawnSessionTable.ended_at.isnot(None),
)
)
row = result.one()
total_cost_7d = float(row.total_cost_7d or 0.0)
avg_daily_cost = total_cost_7d / 7.0
projected_monthly = avg_daily_cost * 30.0
return {
"total_cost_7d": round(total_cost_7d, 6),
"avg_daily_cost_usd": round(avg_daily_cost, 6),
"projected_monthly_cost_usd": round(projected_monthly, 4),
"basis_days": 7,
}
# =========================================================================
# CACHE EFFICIENCY
# =========================================================================
async def get_cache_efficiency(self, period: str = "24h") -> dict[str, Any]:
"""Return cache hit rate and estimated savings from prompt caching.
cache_hit_rate = cache_read_tokens / (input_tokens + cache_read_tokens)
cost_saved = what cache reads would have cost at full input price
minus what they actually cost at cache-read price.
"""
start_dt, _ = _parse_period(period)
result = await self.session.execute(
select(
func.coalesce(func.sum(AgentSpawnSessionTable.tokens_input), 0).label(
"tokens_input"
),
func.coalesce(func.sum(AgentSpawnSessionTable.tokens_output), 0).label(
"tokens_output"
),
func.coalesce(
func.sum(AgentSpawnSessionTable.tokens_cache_read), 0
).label("tokens_cache_read"),
func.coalesce(
func.sum(AgentSpawnSessionTable.tokens_cache_write), 0
).label("tokens_cache_write"),
).where(
AgentSpawnSessionTable.started_at >= start_dt,
AgentSpawnSessionTable.ended_at.isnot(None),
)
)
row = result.one()
tokens_input = int(row.tokens_input or 0)
tokens_cache_read = int(row.tokens_cache_read or 0)
tokens_cache_write = int(row.tokens_cache_write or 0)
total_input_like = tokens_input + tokens_cache_read
cache_hit_rate = (
tokens_cache_read / total_input_like if total_input_like > 0 else 0.0
)
# Cost saved = (cache_read_tokens * full_input_rate) - actual_cache_read_cost
# Use sonnet as the baseline (most common model) for the aggregate estimate.
# The full input price for sonnet is $3/1M; cache read is $0.30/1M.
_MILLION = 1_000_000.0
_FULL_INPUT_PRICE = 3.00 # sonnet baseline, USD/1M
_CACHE_READ_PRICE = 0.30 # 10% of input
cost_at_full_price = tokens_cache_read * _FULL_INPUT_PRICE / _MILLION
cost_at_cache_price = tokens_cache_read * _CACHE_READ_PRICE / _MILLION
cost_saved = cost_at_full_price - cost_at_cache_price
return {
"cache_hit_rate": round(cache_hit_rate, 4),
"tokens_cache_read": tokens_cache_read,
"tokens_cache_write": tokens_cache_write,
"tokens_input": tokens_input,
"cost_saved_by_cache_usd": round(cost_saved, 6),
"period": period,
}
# =========================================================================
# TODAY'S USAGE (for CEO dashboard)
# =========================================================================
async def get_today_summary(self) -> dict[str, Any]:
"""Return today's aggregated usage from daily_usage_rollups.
Used by the CEO dashboard to populate tokens_today and cost_today_usd.
Falls back to 0 values when no data exists for today.
"""
today = datetime.now(UTC).date()
result = await self.session.execute(
select(
func.coalesce(func.sum(DailyUsageRollupTable.tokens_input), 0).label(
"tokens_input"
),
func.coalesce(func.sum(DailyUsageRollupTable.tokens_output), 0).label(
"tokens_output"
),
func.coalesce(
func.sum(DailyUsageRollupTable.tokens_cache_read), 0
).label("tokens_cache_read"),
func.coalesce(
func.sum(DailyUsageRollupTable.tokens_cache_write), 0
).label("tokens_cache_write"),
func.coalesce(
func.sum(DailyUsageRollupTable.total_cost_usd), 0.0
).label("total_cost_usd"),
).where(DailyUsageRollupTable.date == today)
)
row = result.one()
tokens_today = (
int(row.tokens_input or 0)
+ int(row.tokens_output or 0)
+ int(row.tokens_cache_read or 0)
+ int(row.tokens_cache_write or 0)
)
return {
"tokens_today": tokens_today,
"cost_today_usd": round(float(row.total_cost_usd or 0.0), 6),
}
def get_usage_service(db: AsyncSession) -> UsageService:
"""Factory function matching the pattern used by other services."""
return UsageService(db)