generated from coulomb/repo-seed
feat(token-tracking): three-tier token recording on task done
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:
@@ -1,8 +1,12 @@
|
||||
"""
|
||||
Token passthrough test: update_task_status with tokens_in/tokens_out
|
||||
creates a token event automatically.
|
||||
Token passthrough test: update_task_status creates a token event on done.
|
||||
|
||||
Tests the API-level behaviour (the MCP tool delegates to the same endpoints).
|
||||
Three-tier logic:
|
||||
Tier 1 — exact tokens_in/tokens_out provided
|
||||
Tier 2 — workplan_tokens_in/out provided → prorated by task count (note="workplan")
|
||||
Tier 3 — no token args, status=done → heuristic 1000/500 (note="heuristic")
|
||||
|
||||
Non-done status changes never create a token event.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
@@ -27,22 +31,21 @@ async def _create_workstream(client, topic_id):
|
||||
return r.json()
|
||||
|
||||
|
||||
async def _create_task(client, workstream_id):
|
||||
r = await client.post("/tasks/", json={"workstream_id": workstream_id, "title": "my task"})
|
||||
async def _create_task(client, workstream_id, title="my task"):
|
||||
r = await client.post("/tasks/", json={"workstream_id": workstream_id, "title": title})
|
||||
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."""
|
||||
async def test_tier1_exact_tokens(self, client):
|
||||
"""Tier 1: exact tokens_in/tokens_out → used as-is, no note."""
|
||||
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,
|
||||
@@ -53,10 +56,7 @@ class TestTokenPassthrough:
|
||||
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()
|
||||
events = (await client.get("/token-events/", params={"task_id": task["id"]})).json()
|
||||
assert len(events) == 1
|
||||
ev = events[0]
|
||||
assert ev["tokens_in"] == 1200
|
||||
@@ -65,9 +65,51 @@ class TestTokenPassthrough:
|
||||
assert ev["model"] == "claude-sonnet-4-6"
|
||||
assert ev["agent"] == "custodian"
|
||||
assert ev["workstream_id"] == ws["id"]
|
||||
assert ev["note"] is None
|
||||
|
||||
async def test_update_status_without_tokens_creates_no_event(self, client):
|
||||
"""PATCH /tasks/{id} without token fields creates no token_event."""
|
||||
async def test_tier2_workplan_prorated(self, client):
|
||||
"""Tier 2: workplan totals prorated across 4 tasks → 250/125 each, note='workplan'."""
|
||||
await _create_domain(client)
|
||||
topic = await _create_topic(client)
|
||||
ws = await _create_workstream(client, topic["id"])
|
||||
# Create 4 tasks; mark the first done with workplan totals
|
||||
task = await _create_task(client, ws["id"], "T1")
|
||||
for title in ["T2", "T3", "T4"]:
|
||||
await _create_task(client, ws["id"], title)
|
||||
|
||||
r = await client.patch(f"/tasks/{task['id']}", json={
|
||||
"status": "done",
|
||||
"workplan_tokens_in": 1000,
|
||||
"workplan_tokens_out": 500,
|
||||
})
|
||||
assert r.status_code == 200
|
||||
|
||||
events = (await client.get("/token-events/", params={"task_id": task["id"]})).json()
|
||||
assert len(events) == 1
|
||||
ev = events[0]
|
||||
assert ev["tokens_in"] == 250 # 1000 // 4
|
||||
assert ev["tokens_out"] == 125 # 500 // 4
|
||||
assert ev["note"] == "workplan"
|
||||
|
||||
async def test_tier3_heuristic_fallback(self, client):
|
||||
"""Tier 3: status=done with no token args → heuristic 1000/500, note='heuristic'."""
|
||||
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": "done"})
|
||||
assert r.status_code == 200
|
||||
|
||||
events = (await client.get("/token-events/", params={"task_id": task["id"]})).json()
|
||||
assert len(events) == 1
|
||||
ev = events[0]
|
||||
assert ev["tokens_in"] == 1000
|
||||
assert ev["tokens_out"] == 500
|
||||
assert ev["note"] == "heuristic"
|
||||
|
||||
async def test_non_done_status_creates_no_event(self, client):
|
||||
"""Non-done status updates never create a token event."""
|
||||
await _create_domain(client)
|
||||
topic = await _create_topic(client)
|
||||
ws = await _create_workstream(client, topic["id"])
|
||||
@@ -76,6 +118,5 @@ class TestTokenPassthrough:
|
||||
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() == []
|
||||
events = (await client.get("/token-events/", params={"task_id": task["id"]})).json()
|
||||
assert events == []
|
||||
|
||||
Reference in New Issue
Block a user