generated from coulomb/repo-seed
feat(token-tracking): record AI token consumption per task (CUST-WP-0029)
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:
@@ -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
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
Reference in New Issue
Block a user