coulomb/state-hub
2
0
Fork 0
generated from coulomb/repo-seed
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat(token-tracking): three-tier token recording on task done

Browse Source 
Token events are now always created when update_task_status is called
with status="done", using the best available data:

  Tier 1 (best): exact tokens_in + tokens_out passed by agent
  Tier 2:        workplan_tokens_in + workplan_tokens_out prorated
                 across workstream task count (note="workplan")
  Tier 3 (fallback): heuristic 1000 in / 500 out (note="heuristic")

Non-done status changes never create a token event.
MCP tool updated with workplan_tokens_in/out params and tiered docs.
Ralph-workplan skill files updated with the three-tier guidance.
184 tests pass.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
Bernd Worsch
2026-03-29 18:28:18 +02:00
parent 58e1bafce9
commit fdfd4365cd
4 changed files with 119 additions and 47 deletions
Download Patch File Download Diff File Expand all files Collapse all files

53
mcp_server/server.py
View File

@@ -428,29 +428,51 @@ def update_task_status(
blocking_reason: Optional[str] = None,
tokens_in: Optional[int] = None,
tokens_out: Optional[int] = None,
workplan_tokens_in: Optional[int] = None,
workplan_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.
When status='done', always records a token event using the best available data:
Tier 1 (best): pass tokens_in + tokens_out — exact counts from the session
Tier 2: pass workplan_tokens_in + workplan_tokens_out — total workplan
effort prorated across task count (note="workplan")
Tier 3 (fallback): no token args — heuristic 1000 in / 500 out (note="heuristic")
Best practice: read tokens from the Claude Code status bar and pass exact counts.
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
tokens_in: exact input token count for this task (Tier 1)
tokens_out: exact output token count for this task (Tier 1)
workplan_tokens_in: total input tokens for the whole workplan (Tier 2)
workplan_tokens_out: total output tokens for the whole workplan (Tier 2)
model: model identifier, e.g. 'claude-sonnet-4-6'
agent: agent name, e.g. 'custodian', 'ralph'
session_id: agent session identifier
"""
body: dict[str, Any] = {"status": status}
body: dict[str, Any] = {
"status": status,
"model": model,
"agent": agent,
"session_id": session_id,
}
if blocking_reason:
body["blocking_reason"] = blocking_reason
if tokens_in is not None:
body["tokens_in"] = tokens_in
if tokens_out is not None:
body["tokens_out"] = tokens_out
if workplan_tokens_in is not None:
body["workplan_tokens_in"] = workplan_tokens_in
if workplan_tokens_out is not None:
body["workplan_tokens_out"] = workplan_tokens_out
task = _patch(f"/tasks/{task_id}", body)
_post("/progress", {
"task_id": task_id,
@@ -461,19 +483,6 @@ def update_task_status(
"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)