fix(acp): preserve streamed usage on timeout

src/benchflow/acp/session.py

@@ -233,11 +233,21 @@ def record_agent_timeout(
 
     def record_prompt_usage(self, usage: object | None) -> None:
         """Record cumulative ACP token usage returned by session/prompt."""
+        if self._record_usage_snapshot(usage):
+            self._notify_change()
+
+    def _record_usage_snapshot(self, usage: object | None) -> bool:
         snapshot = normalize_acp_usage(usage)
         if snapshot is None:
-            return
+            return False
         self.usage_snapshots.append(snapshot)
-        self._notify_change()
+        return True
+
+    def _record_update_usage(self, update: dict) -> bool:
+        snapshot_source = update.get("usage")
+        if snapshot_source is not None and self._record_usage_snapshot(snapshot_source):
+            return True
+        return self._record_usage_snapshot(update)
 
     def latest_usage_totals(self) -> ACPUsageSnapshot | None:
         """Return the latest cumulative ACP usage snapshot, if any."""
@@ -274,12 +284,15 @@ def handle_update(self, update: dict) -> None:
         """Process a session/update notification."""
         self._events_active = True
         update_type = update.get("sessionUpdate")
+        usage_recorded = self._record_update_usage(update)
         # Unknown update types (future ACP versions, agent-specific
         # extensions) mutate no state and must not trigger a no-op
         # snapshot. Mark events_active so the snapshot path stays on
-        # the modern branch, but skip _notify_change for unrecognized
-        # types.
+        # the modern branch, but only notify for unrecognized types when they
+        # carried usage telemetry.
         if update_type not in self._RECOGNIZED_UPDATE_TYPES:
+            if usage_recorded:
+                self._notify_change()
             return
 
         if update_type == "tool_call":

src/benchflow/rollout/__init__.py

@@ -1886,6 +1886,7 @@ async def cleanup(self) -> None:
             finally:
                 self._usage_runtime = None
 
+        self._collect_native_acp_usage()
         self._finalize_usage_metrics()
         self._enforce_required_usage_tracking()
 

tests/test_native_acp_usage.py

@@ -93,6 +93,97 @@ def test_rollout_native_acp_usage_uses_cumulative_deltas():
     }
 
 
+def test_acp_session_records_usage_from_session_update():
+    """Guards #933: ACP timeout accounting can use pre-cancel update usage."""
+    from benchflow.acp.session import ACPSession
+
+    session = ACPSession("session-1")
+
+    session.handle_update(
+        {
+            "sessionUpdate": "agent_message_chunk",
+            "content": {"type": "text", "text": "working"},
+            "usage": {
+                "inputTokens": 10,
+                "outputTokens": 4,
+                "totalTokens": 16,
+                "cachedReadTokens": 2,
+                "cachedWriteTokens": 1,
+                "thoughtTokens": 1,
+            },
+        }
+    )
+
+    assert session.latest_usage_totals() == {
+        "input_tokens": 10,
+        "output_tokens": 4,
+        "total_tokens": 16,
+        "cached_read_tokens": 2,
+        "cached_write_tokens": 1,
+        "thought_tokens": 1,
+    }
+
+
+@pytest.mark.asyncio
+async def test_cleanup_collects_update_usage_after_timed_out_prompt(tmp_path):
+    """Guards #933: timed-out ACP prompts do not lose streamed token usage."""
+    from benchflow.acp.session import ACPSession
+    from benchflow.rollout import Rollout, RolloutConfig
+    from benchflow.rollout._usage import _zero_native_acp_usage_metrics
+
+    session = ACPSession("session-1")
+    session.handle_update(
+        {
+            "sessionUpdate": "usage",
+            "usage": {
+                "inputTokens": 20,
+                "outputTokens": 7,
+                "totalTokens": 30,
+                "cachedReadTokens": 1,
+                "cachedWriteTokens": 2,
+                "thoughtTokens": 3,
+            },
+        }
+    )
+
+    rollout = Rollout.__new__(Rollout)
+    rollout._session = session
+    rollout._acp_client = SimpleNamespace(session=session)
+    rollout._trajectory = []
+    rollout._session_traj_count = 0
+    rollout._session_tool_count = 0
+    rollout._terminal_timeout = True
+    rollout._partial_trajectory = False
+    rollout._trajectory_source = "none"
+    rollout._n_tool_calls = 0
+    rollout._usage_runtime = None
+    rollout._usage_metrics = {"usage_source": "unavailable"}
+    rollout._native_usage_checkpoint = None
+    rollout._native_usage_metrics = _zero_native_acp_usage_metrics()
+    rollout._env = None
+    rollout._environment = None
+    rollout._config = RolloutConfig(task_path=tmp_path)
+
+    async def fake_disconnect():
+        return None
+
+    rollout.disconnect = fake_disconnect
+
+    await rollout.cleanup()
+
+    assert rollout._usage_metrics == {
+        "n_input_tokens": 20,
+        "n_output_tokens": 7,
+        "n_cache_read_tokens": 1,
+        "n_cache_creation_tokens": 2,
+        "total_tokens": 30,
+        "cost_usd": None,
+        "usage_source": "agent_native_acp",
+        "price_source": None,
+        "usage_details": {"thought_tokens": 3},
+    }
+
+
 def test_rollout_provider_usage_wins_over_native_acp_usage():
     """Guards PR #613 follow-up: LiteLLM provider telemetry remains authoritative."""
     from benchflow.rollout import Rollout