From 1a7a1a63f4ff696e08b7c2d4e453df0713585751 Mon Sep 17 00:00:00 2001 From: Eric Lee Date: Wed, 22 Jul 2026 01:18:13 -0700 Subject: [PATCH] feat: per-step trajectory token metrics + list-price cost estimate on subscription MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Closes the two documented trajectory-parity gaps vs the built-in claude-code agent. FOLLOWUP 1 — per-step token metrics. claude-code's trajectory carries per-turn token usage on every agent step; clawcodex's didn't, because the persisted conversation dropped it. AssistantMessage already holds the turn's usage (query.py sets usage=response.usage), but: - headless _persist called add_message(role, content), dropping usage; - message_to_dict never serialized usage (message_from_dict already loaded it). Fix (production src): thread usage through add_message → create_message → create_assistant_message (which already accepts it), have _persist pass msg.usage, and serialize usage in message_to_dict. Round-trips now. The Harbor adapter builds per-step ATIF Metrics (prompt=input+cache_read+ cache_creation, cached=cache_read, completion=output) from it. FOLLOWUP 2 — cost on subscription runs. record_api_usage zeroes billed cost under a subscription (plan allowance, not metered credits — correct for the live /cost display), so the session cost block reported $0 while claude-code reports a list-price cost. Fix: build_cost_block now also emits estimated_cost_usd — the always-computed list-price figure (via compute_cost, ignoring billing_mode) — purely additive, never touching /cost or the budget gate. The adapter uses it for the leaderboard/ trajectory cost column when billed cost is 0. Verified: fresh headless run persists per-message usage → 4/5 steps get Metrics (user step none, correct); subscription opus tokens → billed $0 / estimated $1.10 (== compute_cost); metered → billed wins. 1348 tests pass across message/conversation/cost/headless/session + new round-trip and estimated-cost tests. Co-Authored-By: Claude Fable 5 --- eval/harbor/clawcodex_agent.py | 41 ++++++++++++++++++++++++-- src/agent/conversation.py | 10 +++++-- src/entrypoints/headless.py | 10 ++++++- src/services/cost_restore.py | 27 ++++++++++++++++- src/types/messages.py | 10 ++++++- tests/test_ch03_state_round4.py | 52 +++++++++++++++++++++++++++++++++ tests/test_message_types.py | 22 ++++++++++++++ 7 files changed, 165 insertions(+), 7 deletions(-) diff --git a/eval/harbor/clawcodex_agent.py b/eval/harbor/clawcodex_agent.py index 254e47a43..7ff61ae0e 100644 --- a/eval/harbor/clawcodex_agent.py +++ b/eval/harbor/clawcodex_agent.py @@ -79,6 +79,7 @@ from harbor.models.trajectories import ( Agent, FinalMetrics, + Metrics, Observation, ObservationResult, Step, @@ -686,12 +687,27 @@ def _billing_totals_from_cost_block( prompt += inp + cread + ccreate cached += cread completion += usage.get("output_tokens") or 0 - total_cost = cost.get("total_cost_usd") + # Cost: prefer the billed ``total_cost_usd``; when it's 0 (a + # subscription run consumes plan allowance, not metered credits, so + # the billed cost is $0) fall back to ``estimated_cost_usd`` — the + # always-computed list-price figure — so the leaderboard/trajectory + # cost column is populated and comparable with the claude-code + # agent, which reports a list-price cost on subscription runs too. + billed = cost.get("total_cost_usd") + estimated = cost.get("estimated_cost_usd") + if isinstance(billed, (int, float)) and billed > 0: + cost_usd: float | None = float(billed) + elif isinstance(estimated, (int, float)) and estimated > 0: + cost_usd = float(estimated) + elif isinstance(billed, (int, float)): + cost_usd = float(billed) + else: + cost_usd = None return { "prompt": prompt, "completion": completion, "cached": cached, - "cost": float(total_cost) if isinstance(total_cost, (int, float)) else None, + "cost": cost_usd, } def _load_session_data( @@ -811,6 +827,26 @@ def _split_content(content: Any) -> tuple[str, str | None, list[dict[str, Any]], reasoning = "\n\n".join(p.strip() for p in reasoning_parts if p and p.strip()) return text, (reasoning or None), tool_uses, tool_results + @staticmethod + def _metrics_from_usage(usage: Any) -> Metrics | None: + """Per-step ATIF ``Metrics`` from a persisted turn's ``usage`` dict + (present since clawcodex persists per-message usage). ``prompt_tokens`` + follows the billing convention (input + cache_read + cache_creation); + ``cached_tokens`` is cache_read. Returns ``None`` when no usage.""" + if not isinstance(usage, dict) or not usage: + return None + inp = usage.get("input_tokens") or 0 + cread = usage.get("cache_read_input_tokens") or 0 + ccreate = usage.get("cache_creation_input_tokens") or 0 + out = usage.get("output_tokens") or 0 + if not any((inp, cread, ccreate, out)): + return None + return Metrics( + prompt_tokens=inp + cread + ccreate, + completion_tokens=out, + cached_tokens=cread, + ) + @staticmethod def _stringify(value: Any) -> str: if isinstance(value, str): @@ -859,6 +895,7 @@ def _steps_from_conversation( message=text, reasoning_content=reasoning, tool_calls=calls or None, + metrics=self._metrics_from_usage(msg.get("usage")), llm_call_count=1, ) steps.append(step) diff --git a/src/agent/conversation.py b/src/agent/conversation.py index 21b352403..dd0a7a44e 100644 --- a/src/agent/conversation.py +++ b/src/agent/conversation.py @@ -41,12 +41,18 @@ class Conversation: # this cap is just a memory-safety backstop. max_history: int = 2000 - def add_message(self, role: str, content: MessageContent): + def add_message( + self, + role: str, + content: MessageContent, + *, + usage: dict[str, Any] | None = None, + ): if len(self.messages) >= self.max_history: self.messages.pop(0) normalized_content = _normalize_message_content(content) - self.messages.append(create_message(role, normalized_content)) + self.messages.append(create_message(role, normalized_content, usage=usage)) def add_user_message(self, content: MessageContent): # ``MessageContent = str | list[ContentBlock]``: ``add_message`` -> diff --git a/src/entrypoints/headless.py b/src/entrypoints/headless.py index a1f32c496..fa6eeaa79 100644 --- a/src/entrypoints/headless.py +++ b/src/entrypoints/headless.py @@ -492,7 +492,15 @@ def _persist(msg: Any) -> None: # the next API call will reject it. Better # to surface now than to debug a 400 later. try: - session.conversation.add_message(msg.role, msg.content) + # Preserve the turn's token usage (assistant + # messages carry it) so the persisted session + # — and the Harbor trajectory's per-step + # Metrics — can attribute tokens per turn. + session.conversation.add_message( + msg.role, + msg.content, + usage=getattr(msg, "usage", None), + ) except Exception: import logging logging.getLogger(__name__).exception( diff --git a/src/services/cost_restore.py b/src/services/cost_restore.py index ce70980a1..662604eef 100644 --- a/src/services/cost_restore.py +++ b/src/services/cost_restore.py @@ -58,8 +58,33 @@ def build_cost_block() -> dict[str, Any]: get_total_tool_duration, ) + model_usage = get_model_usage() + # List-price cost estimate, ALWAYS computed (even under a subscription, + # where ``cost_usd`` is $0 because plan allowance is consumed rather than + # metered credits — see cost_tracker.record_api_usage's billing_mode + # gate). This is an OBSERVABILITY figure — what the tokens would have + # cost at metered API rates — mirroring Claude Code, which reports a + # non-zero cost on subscription runs. It never feeds the live ``/cost`` + # display or budget gate (those keep reading the billed ``total_cost_usd`` + # / ``cost_usd``); it exists so downstream trajectory/leaderboard tooling + # has a comparable cost column. + from src.services.pricing import compute_cost + + estimated_cost_usd = 0.0 + for model, u in model_usage.items(): + estimated_cost_usd += compute_cost( + model, + { + "input_tokens": u.input_tokens, + "output_tokens": u.output_tokens, + "cache_creation_input_tokens": u.cache_creation_input_tokens, + "cache_read_input_tokens": u.cache_read_input_tokens, + }, + ) + return { "total_cost_usd": get_total_cost_usd(), + "estimated_cost_usd": estimated_cost_usd, "total_api_duration": get_total_api_duration(), "total_api_duration_without_retries": get_total_api_duration_without_retries(), @@ -78,7 +103,7 @@ def build_cost_block() -> dict[str, Any]: "cache_read_input_tokens": u.cache_read_input_tokens, "cost_usd": u.cost_usd, } - for model, u in get_model_usage().items() + for model, u in model_usage.items() }, } diff --git a/src/types/messages.py b/src/types/messages.py index b58d9a7db..fc950ee96 100644 --- a/src/types/messages.py +++ b/src/types/messages.py @@ -297,12 +297,15 @@ def create_message( *, timestamp: str | None = None, isMeta: bool = False, + usage: dict[str, Any] | None = None, ) -> Message: ts = timestamp or datetime.now().isoformat() if role == "user": return create_user_message(content, isMeta=isMeta, timestamp=ts) if role == "assistant": - return create_assistant_message(content) + # ``usage`` is the turn's token accounting (input/output/cache); + # only assistant messages carry it. + return create_assistant_message(content, usage=usage) if role == "system": return SystemMessage(content=content, timestamp=ts, isMeta=isMeta) return Message(role=role, content=content, timestamp=ts, isMeta=isMeta) @@ -588,6 +591,11 @@ def message_to_dict(message: MessageLike) -> dict[str, Any]: "isApiErrorMessage", "apiError", "error", "errorDetails", "model", "origin", "toolUseID", "parentToolUseID", "data", "imagePasteIds", "summarizeMetadata", "preventContinuation", + # Per-turn token usage carried by AssistantMessage. Round-trips: + # message_from_dict already loads it. Persisting it lets resumed + # sessions and downstream tooling (the Harbor trajectory builder's + # per-step Metrics) attribute tokens to the turn that spent them. + "usage", ): val = _get_field(message, attr, None) if val is not None and val is not False: diff --git a/tests/test_ch03_state_round4.py b/tests/test_ch03_state_round4.py index 20a98da71..8f862a522 100644 --- a/tests/test_ch03_state_round4.py +++ b/tests/test_ch03_state_round4.py @@ -341,5 +341,57 @@ def test_two_cwds_get_independent_snapshots(self) -> None: self.assertIsNot(asyncio.run(collect_git_context(str(repo_a))), snap_a) +class TestEstimatedCostOnSubscription(unittest.TestCase): + """``build_cost_block`` carries a list-price ``estimated_cost_usd`` even + when the billed ``total_cost_usd`` is $0 under a subscription — an + observability figure for downstream trajectory/leaderboard tooling that + never feeds the live ``/cost`` display.""" + + def setUp(self) -> None: + _reset_all() + + def tearDown(self) -> None: + _reset_all() + + def test_subscription_run_has_zero_billed_but_nonzero_estimate(self) -> None: + from src.cost_tracker import record_api_usage + from src.services.cost_restore import build_cost_block + from src.services.pricing import compute_cost + + tokens = { + "input_tokens": 165812, + "output_tokens": 9246, + "cache_read_input_tokens": 82608, + "cache_creation_input_tokens": 0, + } + record_api_usage("claude-opus-4-8", {**tokens, "billing_mode": "subscription"}) + block = build_cost_block() + + # Billed cost stays $0 (subscription consumes plan allowance). + self.assertEqual(block["total_cost_usd"], 0.0) + self.assertEqual(block["model_usage"]["claude-opus-4-8"]["cost_usd"], 0.0) + # The list-price estimate equals compute_cost for the same tokens. + self.assertGreater(block["estimated_cost_usd"], 0.0) + self.assertAlmostEqual( + block["estimated_cost_usd"], + compute_cost("claude-opus-4-8", tokens), + places=6, + ) + + def test_metered_run_estimate_equals_billed(self) -> None: + from src.cost_tracker import record_api_usage + from src.services.cost_restore import build_cost_block + + record_api_usage( + "claude-opus-4-8", + {"input_tokens": 1000, "output_tokens": 500}, # no billing_mode → metered + ) + block = build_cost_block() + self.assertGreater(block["total_cost_usd"], 0.0) + self.assertAlmostEqual( + block["estimated_cost_usd"], block["total_cost_usd"], places=6 + ) + + if __name__ == "__main__": unittest.main() diff --git a/tests/test_message_types.py b/tests/test_message_types.py index ad502d649..43530e00d 100644 --- a/tests/test_message_types.py +++ b/tests/test_message_types.py @@ -106,6 +106,28 @@ def test_message_round_trip_serialization(self): self.assertEqual(getattr(restored_messages[2], "subtype", None), "api_error") self.assertEqual(getattr(restored_messages[4], "attachments", [])[0]["name"], "diagram.png") + def test_assistant_usage_round_trips(self): + # Per-turn token usage must survive persistence so resumed sessions + # and the Harbor trajectory's per-step Metrics can attribute tokens. + from src.types.messages import create_assistant_message + + usage = { + "input_tokens": 10, + "output_tokens": 5, + "cache_read_input_tokens": 100, + "cache_creation_input_tokens": 2, + } + msg = create_assistant_message("hi", usage=usage) + payload = message_to_dict(msg) + self.assertEqual(payload.get("usage"), usage) + restored = message_from_dict(payload) + self.assertEqual(getattr(restored, "usage", None), usage) + + def test_no_usage_key_when_absent(self): + # A message without usage must not emit a null/empty usage key. + msg = create_user_message("test") + self.assertNotIn("usage", message_to_dict(msg)) + def test_isMeta_field(self): msg = create_user_message("test", isMeta=True) self.assertTrue(msg.isMeta)