fix(v1): record final/context token usage at write time (#1525)

Author: hallerite

tests/test_v1_runtime_lifecycle.py

@@ -732,7 +732,12 @@ async def test_v1_records_default_metrics_usage_and_timing() -> None:
     state = await harness.run(task)
 
     assert state["metrics"]["num_turns"] == 1.0
-    assert state["token_usage"] == {"input_tokens": 11.0, "output_tokens": 7.0}
+    assert state["token_usage"] == {
+        "input_tokens": 11.0,
+        "output_tokens": 7.0,
+        "final_output_tokens": 7.0,
+        "final_input_tokens": 11.0,
+    }
     assert state["usage"] == state["token_usage"]
     assert state["timing"]["total"] > 0.0
     assert state["timing"]["generation"]["duration"] > 0.0

verifiers/utils/save_utils.py

@@ -243,12 +243,30 @@ def state_to_output(
             "input_tokens": usage.get("input_tokens", 0.0),
             "output_tokens": usage.get("output_tokens", 0.0),
         }
-        # Add context token metrics from trajectory
-        trajectory = state.get("trajectory", [])
-        if isinstance(trajectory, list):
-            from verifiers.utils.usage_utils import compute_context_token_metrics
+        # Context ("final") token metrics. v1 records these at write time from
+        # the live Response (the serialized trajectory can't be re-derived since
+        # responses are plain dicts), so prefer them when present. Classic envs
+        # keep live Response objects in the trajectory, so recompute there.
+        raw_usage = state.get("token_usage")
+        final_output = (
+            raw_usage.get("final_output_tokens")
+            if isinstance(raw_usage, Mapping)
+            else None
+        )
+        final_input = (
+            raw_usage.get("final_input_tokens")
+            if isinstance(raw_usage, Mapping)
+            else None
+        )
+        if final_output is not None and final_input is not None:
+            token_usage["final_output_tokens"] = float(final_output)
+            token_usage["final_input_tokens"] = float(final_input)
+        else:
+            trajectory = state.get("trajectory", [])
+            if isinstance(trajectory, list):
+                from verifiers.utils.usage_utils import compute_context_token_metrics
 
-            token_usage.update(compute_context_token_metrics(trajectory))
+                token_usage.update(compute_context_token_metrics(trajectory))
         output["token_usage"] = token_usage
 
     # sanitize messages (handle None for error cases)

verifiers/v1/utils/usage_utils.py

@@ -18,4 +18,17 @@ def record_response_usage(state: State, response: Response) -> None:
     usage["output_tokens"] = float(usage.get("output_tokens", 0.0)) + float(
         output_tokens
     )
+    # Context ("final") token metrics, accumulated at write time from the live
+    # Response. v1 serializes trajectory responses to plain dicts, so they can't
+    # be recomputed from the trajectory afterward (the isinstance(Response) gate
+    # in compute_context_token_metrics fails). Mirror that helper's formula for a
+    # linear rollout: final_output is the running sum of completions; final_input
+    # is the latest step's full context minus that sum.
+    usage["final_output_tokens"] = float(usage.get("final_output_tokens", 0.0)) + float(
+        output_tokens
+    )
+    last_step_total = float(input_tokens) + float(output_tokens)
+    usage["final_input_tokens"] = max(
+        0.0, last_step_total - usage["final_output_tokens"]
+    )
     state["usage"] = usage