feat: per-step trajectory token metrics + list-price cost estimate on subscription

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)

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`` ->

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(

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()
         },
     }
 

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:

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()

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)