Files
the-custodian/state-hub/tests/test_token_passthrough.py
tegwick 68943684e1 feat(token-tracking): introduce token note taxonomy (measured/userbased/workplan/heuristic)
Tier 1 (exact counts) now defaults to note="measured" instead of null,
signalling the counts were read from the Claude Code status bar.
Callers can pass note="userbased" when a human provided the numbers.

  measured  — agent read exact counts from the Claude Code status bar
  userbased — counts provided by a human
  workplan  — prorated from workplan total across task count
  heuristic — server fallback, 1000/500, no agent input

Added token_note field to TaskUpdate schema and exposed note param on
update_task_status and record_interactive_task MCP tools.
TOOLS.md documents the full taxonomy. 185 tests pass.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
2026-03-29 18:47:40 +02:00

141 lines
5.3 KiB
Python

"""
Token passthrough test: update_task_status creates a token event on done.
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
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, 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_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"])
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"
events = (await client.get("/token-events/", params={"task_id": task["id"]})).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"]
assert ev["note"] == "measured"
async def test_tier1_userbased_note_override(self, client):
"""Tier 1 with note='userbased' records that note instead of 'measured'."""
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",
"tokens_in": 500,
"tokens_out": 200,
"token_note": "userbased",
})
assert r.status_code == 200
events = (await client.get("/token-events/", params={"task_id": task["id"]})).json()
assert events[0]["note"] == "userbased"
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"])
task = await _create_task(client, ws["id"])
r = await client.patch(f"/tasks/{task['id']}", json={"status": "in_progress"})
assert r.status_code == 200
events = (await client.get("/token-events/", params={"task_id": task["id"]})).json()
assert events == []