coulomb/the-custodian
2
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat(token-tracking): record AI token consumption per task (CUST-WP-0029)

Browse Source 
Introduces end-to-end token consumption tracking so agent work is
visible as a cost/effort metric alongside tasks and workplans.

- Migration o2j3k4l5m6n7: token_events table with FK indexes on
  task_id, workstream_id, repo_id, created_at
- ORM model, Pydantic schemas (TokenEventCreate, TokenEventRead with
  computed tokens_total, TokenSummary)
- Router: POST /token-events/, GET /token-events/ (7 filters),
  GET /token-events/summary/ (task|workstream|repo|commit|release scope)
- MCP tools: record_token_event, get_token_summary (formatted table)
- update_task_status enriched with optional tokens_in/tokens_out
  passthrough — one call creates status update + token event
- Dashboard token-cost.md page: by-repo bar, by-workplan table,
  by-model bar, top-10 tasks by tokens
- ralph-workplan skill updated with token reporting guidance and
  per-task heuristics for estimating counts
- Tests: test_token_events.py + test_token_passthrough.py (182 pass)

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
Bernd Worsch
2026-03-29 17:46:46 +02:00
parent 280f04c10d
commit d65bc701da
16 changed files with 985 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
state-hub/api/main.py
View File

@@ -7,6 +7,7 @@ from fastapi.middleware.cors import CORSMiddleware
from api.database import engine
from api.routers import decisions, extension_points, progress, state, tasks, technical_debt, topics, workstreams, workstream_dependencies
from api.routers import domains, repos, contributions, sbom, policy, domain_goals, repo_goals, messages, capability_requests, tpsc
from api.routers import token_events
@asynccontextmanager
@@ -49,6 +50,7 @@ app.include_router(sbom.router)
app.include_router(messages.router)
app.include_router(capability_requests.router)
app.include_router(tpsc.router)
app.include_router(token_events.router)
app.include_router(state.router)
app.include_router(policy.router)

2
state-hub/api/models/__init__.py
View File

@@ -19,6 +19,7 @@ from api.models.capability_catalog import CapabilityCatalog
from api.models.capability_request import CapabilityRequest
from api.models.tpsc import TPSCCatalog, TPSCSnapshot, TPSCEntry
from api.models.doi_cache import DOICache
from api.models.token_event import TokenEvent
__all__ = [
"Base",
@@ -42,4 +43,5 @@ __all__ = [
"CapabilityRequest",
"TPSCCatalog", "TPSCSnapshot", "TPSCEntry",
"DOICache",
"TokenEvent",
]

40
state-hub/api/models/token_event.py Normal file
View File

@@ -0,0 +1,40 @@
import uuid
from datetime import datetime
from sqlalchemy import DateTime, ForeignKey, Integer, Text, func
from sqlalchemy.dialects.postgresql import UUID
from sqlalchemy.orm import Mapped, mapped_column, relationship
from api.models.base import Base, new_uuid
class TokenEvent(Base):
__tablename__ = "token_events"
id: Mapped[uuid.UUID] = mapped_column(
UUID(as_uuid=True), primary_key=True, default=new_uuid
)
task_id: Mapped[uuid.UUID | None] = mapped_column(
UUID(as_uuid=True), ForeignKey("tasks.id", ondelete="SET NULL"), nullable=True, index=True
)
workstream_id: Mapped[uuid.UUID | None] = mapped_column(
UUID(as_uuid=True), ForeignKey("workstreams.id", ondelete="SET NULL"), nullable=True, index=True
)
repo_id: Mapped[uuid.UUID | None] = mapped_column(
UUID(as_uuid=True), ForeignKey("managed_repos.id", ondelete="SET NULL"), nullable=True, index=True
)
session_id: Mapped[str | None] = mapped_column(Text, nullable=True)
model: Mapped[str | None] = mapped_column(Text, nullable=True)
tokens_in: Mapped[int] = mapped_column(Integer, nullable=False)
tokens_out: Mapped[int] = mapped_column(Integer, nullable=False)
agent: Mapped[str | None] = mapped_column(Text, nullable=True)
ref_type: Mapped[str | None] = mapped_column(Text, nullable=True)
ref_id: Mapped[str | None] = mapped_column(Text, nullable=True)
note: Mapped[str | None] = mapped_column(Text, nullable=True)
created_at: Mapped[datetime] = mapped_column(
DateTime(timezone=True), server_default=func.now(), nullable=False, index=True
)
task: Mapped["Task | None"] = relationship("Task", lazy="selectin") # noqa: F821
workstream: Mapped["Workstream | None"] = relationship("Workstream", lazy="selectin") # noqa: F821
repo: Mapped["ManagedRepo | None"] = relationship("ManagedRepo", lazy="selectin") # noqa: F821

26
state-hub/api/routers/tasks.py
View File

@@ -7,6 +7,7 @@ from sqlalchemy.ext.asyncio import AsyncSession
from api.database import get_session
from api.models.task import Task, TaskStatus
from api.models.token_event import TokenEvent
from api.schemas.task import TaskCreate, TaskRead, TaskUpdate
router = APIRouter(prefix="/tasks", tags=["tasks"])
@@ -72,10 +73,33 @@ async def update_task(
task = await session.get(Task, task_id)
if task is None:
raise HTTPException(status_code=404, detail="Task not found")
for field, value in body.model_dump(exclude_unset=True).items():
# Separate token fields from task fields
token_fields = {"tokens_in", "tokens_out", "model", "agent", "session_id"}
update_data = body.model_dump(exclude_unset=True)
token_data = {k: update_data.pop(k) for k in list(update_data.keys()) if k in token_fields}
for field, value in update_data.items():
setattr(task, field, value)
await session.commit()
await session.refresh(task)
# Create token event if token passthrough fields provided
if "tokens_in" in token_data and "tokens_out" in token_data:
event = TokenEvent(
task_id=task_id,
workstream_id=task.workstream_id,
tokens_in=token_data["tokens_in"],
tokens_out=token_data["tokens_out"],
model=token_data.get("model"),
agent=token_data.get("agent"),
session_id=token_data.get("session_id"),
ref_type="task",
ref_id=str(task_id),
)
session.add(event)
await session.commit()
return task

122
state-hub/api/routers/token_events.py Normal file
View File

@@ -0,0 +1,122 @@
import uuid
from collections import defaultdict
from fastapi import APIRouter, Depends, HTTPException, Query, status
from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncSession
from api.database import get_session
from api.models.task import Task
from api.models.token_event import TokenEvent
from api.schemas.token_event import TokenEventCreate, TokenEventRead, TokenSummary
router = APIRouter(prefix="/token-events", tags=["token-events"])
@router.post("/", response_model=TokenEventRead, status_code=status.HTTP_201_CREATED)
async def create_token_event(
body: TokenEventCreate,
session: AsyncSession = Depends(get_session),
) -> TokenEvent:
data = body.model_dump()
# Auto-populate workstream_id from task if not provided
if data.get("task_id") and not data.get("workstream_id"):
task = await session.get(Task, data["task_id"])
if task:
data["workstream_id"] = task.workstream_id
event = TokenEvent(**data)
session.add(event)
await session.commit()
await session.refresh(event)
return event
@router.get("/summary/", response_model=TokenSummary)
async def get_token_summary(
scope: str = Query(..., description="task|workstream|repo|commit|release|session"),
id: str = Query(..., description="FK value or ref_id depending on scope"),
session: AsyncSession = Depends(get_session),
) -> TokenSummary:
q = select(TokenEvent)
if scope == "task":
try:
uid = uuid.UUID(id)
except ValueError:
raise HTTPException(status_code=422, detail="id must be a valid UUID for scope=task")
q = q.where(TokenEvent.task_id == uid)
elif scope == "workstream":
try:
uid = uuid.UUID(id)
except ValueError:
raise HTTPException(status_code=422, detail="id must be a valid UUID for scope=workstream")
q = q.where(TokenEvent.workstream_id == uid)
elif scope == "repo":
try:
uid = uuid.UUID(id)
except ValueError:
raise HTTPException(status_code=422, detail="id must be a valid UUID for scope=repo")
q = q.where(TokenEvent.repo_id == uid)
elif scope in ("commit", "release", "session"):
q = q.where(TokenEvent.ref_type == scope, TokenEvent.ref_id == id)
else:
raise HTTPException(status_code=422, detail=f"Unknown scope: {scope!r}")
result = await session.execute(q)
events = list(result.scalars().all())
tokens_in = sum(e.tokens_in for e in events)
tokens_out = sum(e.tokens_out for e in events)
by_model: dict[str, int] = defaultdict(int)
by_agent: dict[str, int] = defaultdict(int)
for e in events:
if e.model:
by_model[e.model] += e.tokens_in + e.tokens_out
if e.agent:
by_agent[e.agent] += e.tokens_in + e.tokens_out
return TokenSummary(
scope=scope,
scope_id=id,
tokens_in=tokens_in,
tokens_out=tokens_out,
tokens_total=tokens_in + tokens_out,
event_count=len(events),
by_model=dict(by_model),
by_agent=dict(by_agent),
)
@router.get("/", response_model=list[TokenEventRead])
async def list_token_events(
task_id: uuid.UUID | None = None,
workstream_id: uuid.UUID | None = None,
repo_id: uuid.UUID | None = None,
ref_type: str | None = None,
ref_id: str | None = None,
model: str | None = None,
agent: str | None = None,
limit: int = Query(100, le=1000),
session: AsyncSession = Depends(get_session),
) -> list[TokenEvent]:
q = select(TokenEvent)
if task_id:
q = q.where(TokenEvent.task_id == task_id)
if workstream_id:
q = q.where(TokenEvent.workstream_id == workstream_id)
if repo_id:
q = q.where(TokenEvent.repo_id == repo_id)
if ref_type:
q = q.where(TokenEvent.ref_type == ref_type)
if ref_id:
q = q.where(TokenEvent.ref_id == ref_id)
if model:
q = q.where(TokenEvent.model == model)
if agent:
q = q.where(TokenEvent.agent == agent)
q = q.order_by(TokenEvent.created_at.desc()).limit(limit)
result = await session.execute(q)
return list(result.scalars().all())

6
state-hub/api/schemas/task.py
View File

@@ -38,6 +38,12 @@ class TaskUpdate(BaseModel):
needs_human: bool | None = None
intervention_note: str | None = None
parent_task_id: uuid.UUID | None = None
# Optional token passthrough — when provided, a token_event is created
tokens_in: int | None = None
tokens_out: int | None = None
model: str | None = None
agent: str | None = None
session_id: str | None = None
@model_validator(mode="after")
def blocking_reason_required_when_blocked(self) -> Self:

52
state-hub/api/schemas/token_event.py Normal file
View File

@@ -0,0 +1,52 @@
import uuid
from datetime import datetime
from pydantic import BaseModel, ConfigDict, computed_field
class TokenEventCreate(BaseModel):
tokens_in: int
tokens_out: int
task_id: uuid.UUID | None = None
workstream_id: uuid.UUID | None = None
repo_id: uuid.UUID | None = None
session_id: str | None = None
model: str | None = None
agent: str | None = None
ref_type: str | None = None
ref_id: str | None = None
note: str | None = None
class TokenEventRead(BaseModel):
model_config = ConfigDict(from_attributes=True)
id: uuid.UUID
tokens_in: int
tokens_out: int
task_id: uuid.UUID | None = None
workstream_id: uuid.UUID | None = None
repo_id: uuid.UUID | None = None
session_id: str | None = None
model: str | None = None
agent: str | None = None
ref_type: str | None = None
ref_id: str | None = None
note: str | None = None
created_at: datetime
@computed_field
@property
def tokens_total(self) -> int:
return self.tokens_in + self.tokens_out
class TokenSummary(BaseModel):
scope: str
scope_id: str
tokens_in: int
tokens_out: int
tokens_total: int
event_count: int
by_model: dict[str, int]
by_agent: dict[str, int]

1
state-hub/dashboard/observablehq.config.js
View File

@@ -25,6 +25,7 @@ export default {
{ name: "Goals", path: "/goals" },
{ name: "Inbox", path: "/inbox" },
{ name: "Progress", path: "/progress" },
{ name: "Token Cost", path: "/token-cost" },
{ name: "Services (TPSC)", path: "/tpsc" },
{ name: "Todo", path: "/todo" },
{ name: "Tools & Apps", path: "/tools" },

80
state-hub/dashboard/src/data/token-summary.json.py Normal file
View File

@@ -0,0 +1,80 @@
#!/usr/bin/env python3
"""Observable data loader: token consumption summary by repo and workstream."""
import json
import os
import urllib.error
import urllib.request
API_BASE = os.environ.get("API_BASE", "http://127.0.0.1:8000").rstrip("/")
def fetch(url: str):
try:
with urllib.request.urlopen(url, timeout=10) as resp:
return json.loads(resp.read())
except urllib.error.URLError:
return None
# Fetch all repos and workstreams for scope resolution
repos = fetch(f"{API_BASE}/repos/") or []
workstreams_raw = fetch(f"{API_BASE}/workstreams/?limit=500") or []
# Fetch all token events (up to 1000) for aggregation
events = fetch(f"{API_BASE}/token-events/?limit=1000") or []
def aggregate(events, key_fn, label_fn):
"""Group token events by a key function and return aggregated records."""
groups: dict = {}
for e in events:
k = key_fn(e)
if not k:
continue
if k not in groups:
groups[k] = {"scope_id": k, "label": label_fn(k), "tokens_in": 0, "tokens_out": 0, "event_count": 0, "by_model": {}}
groups[k]["tokens_in"] += e.get("tokens_in", 0)
groups[k]["tokens_out"] += e.get("tokens_out", 0)
groups[k]["event_count"] += 1
model = e.get("model") or "unknown"
groups[k]["by_model"][model] = groups[k]["by_model"].get(model, 0) + e.get("tokens_in", 0) + e.get("tokens_out", 0)
for v in groups.values():
v["tokens_total"] = v["tokens_in"] + v["tokens_out"]
return sorted(groups.values(), key=lambda x: -x["tokens_total"])
repo_map = {r["id"]: r.get("slug", r["id"]) for r in repos}
ws_map = {w["id"]: w.get("title", w["id"]) for w in workstreams_raw}
by_repo = aggregate(events, lambda e: e.get("repo_id"), lambda k: repo_map.get(k, k))
by_workstream = aggregate(events, lambda e: e.get("workstream_id"), lambda k: ws_map.get(k, k))
# Top 10 tasks by tokens
task_groups: dict = {}
for e in events:
tid = e.get("task_id")
if not tid:
continue
if tid not in task_groups:
task_groups[tid] = {"task_id": tid, "tokens_in": 0, "tokens_out": 0, "event_count": 0}
task_groups[tid]["tokens_in"] += e.get("tokens_in", 0)
task_groups[tid]["tokens_out"] += e.get("tokens_out", 0)
task_groups[tid]["event_count"] += 1
for v in task_groups.values():
v["tokens_total"] = v["tokens_in"] + v["tokens_out"]
top_tasks = sorted(task_groups.values(), key=lambda x: -x["tokens_total"])[:10]
# Model breakdown across all events
model_totals: dict = {}
for e in events:
model = e.get("model") or "unknown"
model_totals[model] = model_totals.get(model, 0) + e.get("tokens_in", 0) + e.get("tokens_out", 0)
by_model = [{"model": k, "tokens_total": v} for k, v in sorted(model_totals.items(), key=lambda x: -x[1])]
print(json.dumps({
"by_repo": by_repo,
"by_workstream": by_workstream,
"top_tasks": top_tasks,
"by_model": by_model,
"total_events": len(events),
}))

170
state-hub/dashboard/src/token-cost.md Normal file
View File

@@ -0,0 +1,170 @@
---
title: Token Cost
---
```js
import {API} from "./components/config.js";
const POLL = 60_000;
```
```js
// Live poll for token data
const tokenState = (async function*() {
while (true) {
let data = {by_repo: [], by_workstream: [], top_tasks: [], by_model: [], total_events: 0}, ok = false;
try {
const r = await fetch(`${API}/token-events/?limit=1000`);
ok = r.ok;
if (ok) {
const events = await r.json();
data = buildSummary(events);
}
} catch {}
yield {data, ok, ts: new Date()};
await new Promise(res => setTimeout(res, POLL));
}
})();
```
```js
function buildSummary(events) {
const byRepo = {}, byWs = {}, byModel = {}, byTask = {};
for (const e of events) {
const tot = (e.tokens_in || 0) + (e.tokens_out || 0);
if (e.repo_id) {
byRepo[e.repo_id] = byRepo[e.repo_id] || {scope_id: e.repo_id, tokens_in: 0, tokens_out: 0, event_count: 0};
byRepo[e.repo_id].tokens_in += e.tokens_in || 0;
byRepo[e.repo_id].tokens_out += e.tokens_out || 0;
byRepo[e.repo_id].event_count++;
}
if (e.workstream_id) {
byWs[e.workstream_id] = byWs[e.workstream_id] || {scope_id: e.workstream_id, tokens_in: 0, tokens_out: 0, event_count: 0};
byWs[e.workstream_id].tokens_in += e.tokens_in || 0;
byWs[e.workstream_id].tokens_out += e.tokens_out || 0;
byWs[e.workstream_id].event_count++;
}
const model = e.model || "unknown";
byModel[model] = (byModel[model] || 0) + tot;
if (e.task_id) {
byTask[e.task_id] = byTask[e.task_id] || {task_id: e.task_id, tokens_in: 0, tokens_out: 0};
byTask[e.task_id].tokens_in += e.tokens_in || 0;
byTask[e.task_id].tokens_out += e.tokens_out || 0;
}
}
const sortDesc = obj => Object.entries(obj)
.map(([k,v]) => typeof v === "number" ? {id: k, tokens_total: v} : {...v, tokens_total: (v.tokens_in||0)+(v.tokens_out||0)})
.sort((a,b) => b.tokens_total - a.tokens_total);
return {
by_repo: sortDesc(byRepo),
by_workstream: sortDesc(byWs),
by_model: Object.entries(byModel).map(([model,tokens_total]) => ({model,tokens_total})).sort((a,b)=>b.tokens_total-a.tokens_total),
top_tasks: sortDesc(byTask).slice(0,10),
total_events: events.length,
};
}
```
```js
const td = tokenState.data ?? {by_repo:[], by_workstream:[], top_tasks:[], by_model:[], total_events:0};
const _ok = tokenState.ok ?? false;
const _ts = tokenState.ts;
```
# Token Cost
```js
const _liveEl = html`<div style="font-size:0.8rem;color:${_ok?'var(--theme-foreground-focus)':'red'}">
● ${_ok ? `Live · ${_ts?.toLocaleTimeString()} · ${td.total_events} events` : "API offline"}
</div>`;
display(_liveEl);
```
## By Repo
```js
if (td.by_repo.length === 0) {
display(html`<p style="color:var(--theme-foreground-muted)">No token events recorded yet.</p>`);
} else {
display(Plot.plot({
title: "Token consumption by repo",
marginLeft: 160,
width: Math.min(900, width),
x: {label: "Tokens", tickFormat: "~s"},
y: {label: null},
color: {legend: true, domain: ["tokens_in", "tokens_out"], range: ["#4e79a7","#f28e2b"]},
marks: [
Plot.barX(
td.by_repo.flatMap(r => [
{repo: r.scope_id.slice(0,8), type: "tokens_in", value: r.tokens_in},
{repo: r.scope_id.slice(0,8), type: "tokens_out", value: r.tokens_out},
]),
{x: "value", y: "repo", fill: "type", tip: true}
),
],
}));
}
```
## By Workplan
```js
const wsRows = td.by_workstream.slice(0, 20);
if (wsRows.length === 0) {
display(html`<p style="color:var(--theme-foreground-muted)">No workstream data yet.</p>`);
} else {
display(Inputs.table(wsRows, {
columns: ["scope_id", "tokens_in", "tokens_out", "tokens_total", "event_count"],
header: {
scope_id: "Workstream ID",
tokens_in: "Tokens In",
tokens_out: "Tokens Out",
tokens_total: "Total",
event_count: "Events",
},
format: {
scope_id: d => d.slice(0,8) + "…",
tokens_in: d => d.toLocaleString(),
tokens_out: d => d.toLocaleString(),
tokens_total: d => d.toLocaleString(),
},
width: {scope_id: 120, tokens_in: 110, tokens_out: 110, tokens_total: 110, event_count: 80},
}));
}
```
## By Model
```js
if (td.by_model.length === 0) {
display(html`<p style="color:var(--theme-foreground-muted)">No model data yet.</p>`);
} else {
display(Plot.plot({
title: "Token consumption by model",
marginLeft: 200,
width: Math.min(700, width),
x: {label: "Total tokens", tickFormat: "~s"},
marks: [
Plot.barX(td.by_model, {x: "tokens_total", y: "model", fill: "#4e79a7", tip: true}),
],
}));
}
```
## Top 10 Tasks by Tokens
```js
if (td.top_tasks.length === 0) {
display(html`<p style="color:var(--theme-foreground-muted)">No task-level data yet.</p>`);
} else {
display(Inputs.table(td.top_tasks, {
columns: ["task_id", "tokens_in", "tokens_out", "tokens_total"],
header: {task_id: "Task ID", tokens_in: "In", tokens_out: "Out", tokens_total: "Total"},
format: {
task_id: d => d.slice(0,8) + "…",
tokens_in: d => d.toLocaleString(),
tokens_out: d => d.toLocaleString(),
tokens_total: d => d.toLocaleString(),
},
}));
}
```

13
state-hub/mcp_server/TOOLS.md
View File

@@ -73,6 +73,19 @@ Use `list_human_interventions()` at session start to see Bernd's action items.
---
## Token Consumption Tools
Record and query AI token usage at task/workstream/repo/commit/release granularity.
Agents should call `record_token_event` (or pass `tokens_in`/`tokens_out` via
`update_task_status`) at task completion.
| Tool | Key Args | Notes |
|------|----------|-------|
| `record_token_event(tokens_in, tokens_out, ...)` | `task_id`?, `workstream_id`?, `repo_id`?, `model`?, `agent`?, `ref_type`?, `ref_id`?, `note`?, `session_id`? | POSTs to `/token-events/`. `workstream_id` auto-filled from task. Returns event id + running total. |
| `get_token_summary(scope, id)` | `scope`: task\|workstream\|repo\|commit\|release\|session; `id`: UUID or ref string | Returns formatted table of tokens_in/out/total, event_count, by_model, by_agent. |
---
## Governance Tools
| Tool | Key Args | When to use |

146
state-hub/mcp_server/server.py
View File

@@ -425,14 +425,28 @@ def create_task(
def update_task_status(
task_id: str,
status: str,
blocking_reason: str | None = None,
blocking_reason: Optional[str] = None,
tokens_in: Optional[int] = None,
tokens_out: Optional[int] = None,
model: Optional[str] = None,
agent: Optional[str] = None,
session_id: Optional[str] = None,
) -> str:
"""Update a task's status. blocking_reason is required when status='blocked'.
Optionally record token consumption in one call by passing tokens_in/tokens_out.
When provided, a token_event is created automatically with workstream_id and
repo_id auto-populated from the task.
Args:
task_id: UUID of the task
status: todo | in_progress | blocked | done | cancelled
blocking_reason: required when status=blocked
tokens_in: optional input token count (triggers token_event creation)
tokens_out: optional output token count (required if tokens_in provided)
model: optional model identifier, e.g. 'claude-sonnet-4-6'
agent: optional agent name, e.g. 'custodian', 'ralph'
session_id: optional agent session identifier
"""
body: dict[str, Any] = {"status": status}
if blocking_reason:
@@ -446,6 +460,20 @@ def update_task_status(
"author": "custodian",
"detail": {"blocking_reason": blocking_reason},
})
if tokens_in is not None and tokens_out is not None:
_post("/token-events", {
"task_id": task_id,
"workstream_id": task.get("workstream_id"),
"tokens_in": tokens_in,
"tokens_out": tokens_out,
"model": model,
"agent": agent,
"session_id": session_id,
"ref_type": "task",
"ref_id": task_id,
})
return json.dumps(task, indent=2)
@@ -2185,6 +2213,122 @@ def get_doi_summary() -> str:
return json.dumps(_get("/repos/doi/summary"), indent=2)
# ---------------------------------------------------------------------------
# Token events
# ---------------------------------------------------------------------------
@mcp.tool()
def record_token_event(
tokens_in: int,
tokens_out: int,
task_id: Optional[str] = None,
workstream_id: Optional[str] = None,
repo_id: Optional[str] = None,
model: Optional[str] = None,
agent: Optional[str] = None,
ref_type: Optional[str] = None,
ref_id: Optional[str] = None,
note: Optional[str] = None,
session_id: Optional[str] = None,
) -> str:
"""Record AI token consumption for a task, workstream, or session.
workstream_id is auto-populated from the task if task_id is provided and
workstream_id is omitted. Returns the created event id and running total
for the task/workstream (if applicable).
Args:
tokens_in: Input token count
tokens_out: Output token count
task_id: UUID of the task (nullable)
workstream_id: UUID of the workstream (nullable; auto-filled from task)
repo_id: UUID of the managed repo (nullable)
model: Model identifier, e.g. 'claude-sonnet-4-6'
agent: Agent name, e.g. 'custodian', 'ralph'
ref_type: 'task'|'workstream'|'commit'|'release'|'session'
ref_id: Commit SHA, release tag, or other reference string
note: Free-text note
session_id: Agent session identifier
"""
body = {
"tokens_in": tokens_in,
"tokens_out": tokens_out,
"task_id": task_id,
"workstream_id": workstream_id,
"repo_id": repo_id,
"model": model,
"agent": agent,
"ref_type": ref_type,
"ref_id": ref_id,
"note": note,
"session_id": session_id,
}
result = _post("/token-events", body)
if "error" in result:
return json.dumps(result)
out = {
"event_id": result.get("id"),
"tokens_total": result.get("tokens_total"),
"tokens_in": result.get("tokens_in"),
"tokens_out": result.get("tokens_out"),
}
# Append running total for the task if available
scope_id = task_id or workstream_id
scope = "task" if task_id else ("workstream" if workstream_id else None)
if scope and scope_id:
summary = _get("/token-events/summary", {"scope": scope, "id": scope_id})
if "error" not in summary:
out["running_total"] = {
"scope": scope,
"scope_id": scope_id,
"tokens_total": summary.get("tokens_total"),
"event_count": summary.get("event_count"),
}
return json.dumps(out, indent=2)
@mcp.tool()
def get_token_summary(scope: str, id: str) -> str:
"""Return token consumption summary for a given scope.
Returns a formatted table of token usage aggregated by scope.
Args:
scope: One of: task | workstream | repo | commit | release | session
id: UUID for task/workstream/repo scopes; ref_id string for commit/release/session
"""
result = _get("/token-events/summary", {"scope": scope, "id": id})
if "error" in result:
return json.dumps(result)
lines = [
f"Token Summary — {scope}: {id}",
f"{'─' * 50}",
f" tokens_in : {result.get('tokens_in', 0):>10,}",
f" tokens_out : {result.get('tokens_out', 0):>10,}",
f" tokens_total: {result.get('tokens_total', 0):>10,}",
f" event_count : {result.get('event_count', 0):>10,}",
]
by_model = result.get("by_model", {})
if by_model:
lines.append("\nBy model:")
for m, t in sorted(by_model.items(), key=lambda x: -x[1]):
lines.append(f" {m:<35} {t:>10,}")
by_agent = result.get("by_agent", {})
if by_agent:
lines.append("\nBy agent:")
for a, t in sorted(by_agent.items(), key=lambda x: -x[1]):
lines.append(f" {a:<35} {t:>10,}")
return "\n".join(lines)
# ---------------------------------------------------------------------------
# Entry point
# ---------------------------------------------------------------------------

46
state-hub/migrations/versions/o2j3k4l5m6n7_add_token_events.py Normal file
View File

@@ -0,0 +1,46 @@
"""add token_events table
Revision ID: o2j3k4l5m6n7
Revises: n1i2j3k4l5m6
Create Date: 2026-03-29
"""
from alembic import op
import sqlalchemy as sa
from sqlalchemy.dialects.postgresql import UUID
revision = "o2j3k4l5m6n7"
down_revision = "n1i2j3k4l5m6"
branch_labels = None
depends_on = None
def upgrade() -> None:
op.create_table(
"token_events",
sa.Column("id", UUID(as_uuid=True), primary_key=True, server_default=sa.text("gen_random_uuid()")),
sa.Column("task_id", UUID(as_uuid=True), sa.ForeignKey("tasks.id", ondelete="SET NULL"), nullable=True),
sa.Column("workstream_id", UUID(as_uuid=True), sa.ForeignKey("workstreams.id", ondelete="SET NULL"), nullable=True),
sa.Column("repo_id", UUID(as_uuid=True), sa.ForeignKey("managed_repos.id", ondelete="SET NULL"), nullable=True),
sa.Column("session_id", sa.Text(), nullable=True),
sa.Column("model", sa.Text(), nullable=True),
sa.Column("tokens_in", sa.Integer(), nullable=False),
sa.Column("tokens_out", sa.Integer(), nullable=False),
sa.Column("agent", sa.Text(), nullable=True),
sa.Column("ref_type", sa.Text(), nullable=True),
sa.Column("ref_id", sa.Text(), nullable=True),
sa.Column("note", sa.Text(), nullable=True),
sa.Column("created_at", sa.TIMESTAMP(timezone=True), server_default=sa.text("now()"), nullable=False),
)
op.create_index("ix_token_events_task_id", "token_events", ["task_id"])
op.create_index("ix_token_events_workstream_id", "token_events", ["workstream_id"])
op.create_index("ix_token_events_repo_id", "token_events", ["repo_id"])
op.create_index("ix_token_events_created_at", "token_events", ["created_at"])
def downgrade() -> None:
op.drop_index("ix_token_events_created_at", table_name="token_events")
op.drop_index("ix_token_events_repo_id", table_name="token_events")
op.drop_index("ix_token_events_workstream_id", table_name="token_events")
op.drop_index("ix_token_events_task_id", table_name="token_events")
op.drop_table("token_events")

198
state-hub/tests/test_token_events.py Normal file
View File

@@ -0,0 +1,198 @@
"""
Token events router tests.
Covers: create event, list with filters, summary aggregation (single task,
cross-workstream rollup, by-model breakdown).
"""
from __future__ import annotations
import pytest
# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------
async def _create_domain(client, slug="testdomain"):
r = await client.post("/domains/", json={"slug": slug, "name": "Test Domain"})
assert r.status_code == 201, r.text
return r.json()
async def _create_topic(client, domain_slug="testdomain"):
r = await client.post("/topics/", json={"slug": "testtopic", "title": "T", "domain": domain_slug})
assert r.status_code == 201, r.text
return r.json()
async def _create_workstream(client, topic_id, slug="ws1"):
r = await client.post("/workstreams/", json={"topic_id": topic_id, "slug": slug, "title": "WS"})
assert r.status_code == 201, r.text
return r.json()
async def _create_task(client, workstream_id):
r = await client.post("/tasks/", json={"workstream_id": workstream_id, "title": "task"})
assert r.status_code == 201, r.text
return r.json()
async def _post_event(client, tokens_in=100, tokens_out=50, **kwargs):
body = {"tokens_in": tokens_in, "tokens_out": tokens_out, **kwargs}
r = await client.post("/token-events/", json=body)
assert r.status_code == 201, r.text
return r.json()
# ---------------------------------------------------------------------------
# Tests
# ---------------------------------------------------------------------------
@pytest.mark.asyncio
class TestTokenEventsCreate:
async def test_create_minimal(self, client):
ev = await _post_event(client, tokens_in=200, tokens_out=100)
assert ev["tokens_in"] == 200
assert ev["tokens_out"] == 100
assert ev["tokens_total"] == 300
assert ev["id"] is not None
async def test_create_with_all_fields(self, client):
await _create_domain(client)
topic = await _create_topic(client)
ws = await _create_workstream(client, topic["id"])
task = await _create_task(client, ws["id"])
ev = await _post_event(
client,
tokens_in=1000,
tokens_out=500,
task_id=task["id"],
model="claude-sonnet-4-6",
agent="custodian",
ref_type="task",
ref_id=task["id"],
note="T01 done",
session_id="ses-abc",
)
assert ev["task_id"] == task["id"]
assert ev["workstream_id"] == ws["id"] # auto-populated from task
assert ev["model"] == "claude-sonnet-4-6"
assert ev["tokens_total"] == 1500
async def test_workstream_auto_populated_from_task(self, client):
await _create_domain(client)
topic = await _create_topic(client)
ws = await _create_workstream(client, topic["id"])
task = await _create_task(client, ws["id"])
ev = await _post_event(client, task_id=task["id"])
assert ev["workstream_id"] == ws["id"]
@pytest.mark.asyncio
class TestTokenEventsList:
async def test_list_empty(self, client):
r = await client.get("/token-events/")
assert r.status_code == 200
assert r.json() == []
async def test_list_returns_events(self, client):
await _post_event(client, tokens_in=100, tokens_out=50)
await _post_event(client, tokens_in=200, tokens_out=100)
r = await client.get("/token-events/")
assert r.status_code == 200
assert len(r.json()) == 2
async def test_filter_by_task_id(self, client):
await _create_domain(client)
topic = await _create_topic(client)
ws = await _create_workstream(client, topic["id"])
task = await _create_task(client, ws["id"])
await _post_event(client, task_id=task["id"], tokens_in=100, tokens_out=50)
await _post_event(client, tokens_in=200, tokens_out=100) # unrelated
r = await client.get("/token-events/", params={"task_id": task["id"]})
assert r.status_code == 200
events = r.json()
assert len(events) == 1
assert events[0]["task_id"] == task["id"]
async def test_filter_by_model(self, client):
await _post_event(client, model="claude-sonnet-4-6", tokens_in=100, tokens_out=50)
await _post_event(client, model="claude-opus-4-6", tokens_in=200, tokens_out=100)
r = await client.get("/token-events/", params={"model": "claude-sonnet-4-6"})
assert r.status_code == 200
events = r.json()
assert len(events) == 1
assert events[0]["model"] == "claude-sonnet-4-6"
@pytest.mark.asyncio
class TestTokenSummary:
async def test_summary_single_task(self, client):
await _create_domain(client)
topic = await _create_topic(client)
ws = await _create_workstream(client, topic["id"])
task = await _create_task(client, ws["id"])
await _post_event(client, task_id=task["id"], tokens_in=500, tokens_out=300, model="model-a")
await _post_event(client, task_id=task["id"], tokens_in=100, tokens_out=50, model="model-a")
r = await client.get("/token-events/summary/", params={"scope": "task", "id": task["id"]})
assert r.status_code == 200
s = r.json()
assert s["scope"] == "task"
assert s["tokens_in"] == 600
assert s["tokens_out"] == 350
assert s["tokens_total"] == 950
assert s["event_count"] == 2
assert "model-a" in s["by_model"]
assert s["by_model"]["model-a"] == 950
async def test_summary_workstream_rollup(self, client):
await _create_domain(client)
topic = await _create_topic(client)
ws = await _create_workstream(client, topic["id"])
task1 = await _create_task(client, ws["id"])
task2 = await _create_task(client, ws["id"])
await _post_event(client, task_id=task1["id"], tokens_in=1000, tokens_out=500)
await _post_event(client, task_id=task2["id"], workstream_id=ws["id"], tokens_in=200, tokens_out=100)
r = await client.get("/token-events/summary/", params={"scope": "workstream", "id": ws["id"]})
assert r.status_code == 200
s = r.json()
# task1 auto-populates workstream_id; task2 explicitly sets it
assert s["tokens_total"] >= 1800
async def test_summary_by_model_breakdown(self, client):
await _post_event(client, model="sonnet", tokens_in=300, tokens_out=200, agent="custodian")
await _post_event(client, model="opus", tokens_in=100, tokens_out=50, agent="ralph")
await _post_event(client, model="sonnet", tokens_in=200, tokens_out=100)
# Use workstream_id scope via events directly tagged with workstream
# Instead, just check the ref_type/ref_id scope path
await _post_event(
client, model="sonnet", tokens_in=50, tokens_out=25,
ref_type="session", ref_id="ses-001",
)
r = await client.get("/token-events/summary/", params={"scope": "session", "id": "ses-001"})
assert r.status_code == 200
s = r.json()
assert s["event_count"] == 1
assert s["tokens_total"] == 75
async def test_summary_unknown_scope_returns_422(self, client):
r = await client.get("/token-events/summary/", params={"scope": "foobar", "id": "x"})
assert r.status_code == 422
async def test_summary_empty_scope_returns_zeros(self, client):
import uuid
r = await client.get("/token-events/summary/", params={"scope": "task", "id": str(uuid.uuid4())})
assert r.status_code == 200
s = r.json()
assert s["tokens_total"] == 0
assert s["event_count"] == 0

81
state-hub/tests/test_token_passthrough.py Normal file
View File

@@ -0,0 +1,81 @@
"""
Token passthrough test: update_task_status with tokens_in/tokens_out
creates a token event automatically.
Tests the API-level behaviour (the MCP tool delegates to the same endpoints).
"""
from __future__ import annotations
import pytest
async def _create_domain(client, slug="td"):
r = await client.post("/domains/", json={"slug": slug, "name": "D"})
assert r.status_code == 201, r.text
return r.json()
async def _create_topic(client, domain_slug="td"):
r = await client.post("/topics/", json={"slug": "tp", "title": "T", "domain": domain_slug})
assert r.status_code == 201, r.text
return r.json()
async def _create_workstream(client, topic_id):
r = await client.post("/workstreams/", json={"topic_id": topic_id, "slug": "ws", "title": "WS"})
assert r.status_code == 201, r.text
return r.json()
async def _create_task(client, workstream_id):
r = await client.post("/tasks/", json={"workstream_id": workstream_id, "title": "my task"})
assert r.status_code == 201, r.text
return r.json()
@pytest.mark.asyncio
class TestTokenPassthrough:
async def test_update_status_with_tokens_creates_event(self, client):
"""PATCH /tasks/{id} with tokens_in/tokens_out creates a token_event."""
await _create_domain(client)
topic = await _create_topic(client)
ws = await _create_workstream(client, topic["id"])
task = await _create_task(client, ws["id"])
# Update task status with token data
r = await client.patch(f"/tasks/{task['id']}", json={
"status": "done",
"tokens_in": 1200,
"tokens_out": 800,
"model": "claude-sonnet-4-6",
"agent": "custodian",
})
assert r.status_code == 200
assert r.json()["status"] == "done"
# Token event should now exist for this task
r2 = await client.get("/token-events/", params={"task_id": task["id"]})
assert r2.status_code == 200
events = r2.json()
assert len(events) == 1
ev = events[0]
assert ev["tokens_in"] == 1200
assert ev["tokens_out"] == 800
assert ev["tokens_total"] == 2000
assert ev["model"] == "claude-sonnet-4-6"
assert ev["agent"] == "custodian"
assert ev["workstream_id"] == ws["id"]
async def test_update_status_without_tokens_creates_no_event(self, client):
"""PATCH /tasks/{id} without token fields creates no token_event."""
await _create_domain(client)
topic = await _create_topic(client)
ws = await _create_workstream(client, topic["id"])
task = await _create_task(client, ws["id"])
r = await client.patch(f"/tasks/{task['id']}", json={"status": "in_progress"})
assert r.status_code == 200
r2 = await client.get("/token-events/", params={"task_id": task["id"]})
assert r2.status_code == 200
assert r2.json() == []

4
workplans/CUST-WP-0029-token-consumption-tracking.md
View File

@@ -4,7 +4,7 @@ type: workplan
title: "Token Consumption Tracking"
domain: custodian
repo: the-custodian
status: active
status: done
owner: custodian
topic_slug: custodian
created: "2026-03-29"
@@ -239,7 +239,7 @@ for an agent running a loop.
```task
id: CUST-WP-0029-T08
status: todo
status: done
priority: high
state_hub_task_id: "a3627144-9d98-4a3b-aa64-3079fd087448"
```