fix: keep NOOA BenchAgent token usage on failed tasks

adapters/python/nooa/docs/bench-agent.md

@@ -23,7 +23,17 @@ One NeMo Fabric invocation maps to one BenchAgent task:
 - `instructions.system` becomes `task_input.instructions`.
 - The native `response`, `success`, and structured `result` fields become the
   normalized NeMo Fabric result.
-- NOOA's task token accumulator becomes normalized usage.
+- With NOOA's completion client, NOOA's task token accumulator becomes
+  normalized usage.
+- With NOOA's Responses API client, such as with
+  `models.<role>.settings.client_type: responses`, the adapter sums the usage
+  that each Responses API call returns. `input_tokens` includes
+  `cached_input_tokens`. A counter that any call omits stays unknown, so a
+  call that returns no usage leaves the whole usage unknown.
+- When NOOA reports that the task failed, the failed result still carries the
+  usage of the model calls that the task made. When the task raises after a
+  model call starts, the failed result also carries that usage. A task that
+  fails before any model call has no usage.
 
 The adapter accepts exactly one model. Before each task, it closes the shell
 that `_run_evaluation()` replaces. During runtime shutdown, it closes the

adapters/python/nooa/src/nemo_fabric_adapters/nooa/bench_adapter.py

@@ -10,6 +10,9 @@
 import inspect
 import logging
 from collections.abc import Callable
+from contextvars import ContextVar
+from dataclasses import dataclass
+from dataclasses import field
 from pathlib import Path
 from typing import Any
 
@@ -31,6 +34,22 @@
 HARNESS = "nooa-bench"
 MODE = "bench_agent"
 LOGGER = logging.getLogger(__name__)
+_MAX_TOKEN_COUNT = (1 << 64) - 1
+
+
+@dataclass
+class _ModelCalls:
+    """Model calls made during one invocation."""
+
+    # Set when the first model call starts, before any response arrives.
+    started: bool = False
+    # The ``usage`` value from each returned call; read for the Responses API.
+    usages: list[Any] = field(default_factory=list)
+
+
+_model_calls: ContextVar[_ModelCalls | None] = ContextVar(
+    "nooa_bench_model_calls", default=None
+)
 
 
 def main() -> None:
@@ -113,7 +132,11 @@ async def _close_shell(agent: Any) -> None:
         await _await_if_needed(close())
 
 
-def _failure_output(code: str, message: str) -> AgentRunResult:
+def _failure_output(
+    code: str,
+    message: str,
+    usage: AgentUsage | None = None,
+) -> AgentRunResult:
     return AgentRunResult(
         status=AgentRunStatus.FAILED,
         output={
@@ -124,6 +147,7 @@ def _failure_output(code: str, message: str) -> AgentRunResult:
             "completed": False,
         },
         error=AgentRunError(code=code, message=message),
+        usage=usage,
     )
 
 
@@ -148,6 +172,91 @@ def _usage(tokens: Any) -> AgentUsage | None:
     )
 
 
+def _token_count(usage: Any, *path: str) -> int | None:
+    value = usage
+    for key in path:
+        if not isinstance(value, dict):
+            return None
+        value = value.get(key)
+    if isinstance(value, bool) or not isinstance(value, int):
+        return None
+    if not 0 <= value <= _MAX_TOKEN_COUNT:
+        return None
+    return value
+
+
+def _summed_token_count(usages: list[Any], *path: str) -> int | None:
+    """Sum one counter across calls; any call that omits it makes the sum unknown."""
+    total = 0
+    for usage in usages:
+        value = _token_count(usage, *path)
+        if value is None:
+            return None
+        total += value
+    return total if total <= _MAX_TOKEN_COUNT else None
+
+
+def _responses_api_usage(usages: list[Any]) -> AgentUsage | None:
+    """Sum OpenAI Responses API usage, whose input_tokens includes cached tokens."""
+    if not usages:
+        return None
+    input_tokens = _summed_token_count(usages, "input_tokens")
+    cached_input_tokens = _summed_token_count(
+        usages, "input_tokens_details", "cached_tokens"
+    )
+    output_tokens = _summed_token_count(usages, "output_tokens")
+    # Provider-reported, not derived, so it stays known when output_tokens is not.
+    total_tokens = _summed_token_count(usages, "total_tokens")
+    if (
+        input_tokens is None
+        and cached_input_tokens is None
+        and output_tokens is None
+        and total_tokens is None
+    ):
+        return None
+    return AgentUsage(
+        input_tokens=input_tokens,
+        cached_input_tokens=cached_input_tokens,
+        input_tokens_include_cache=True if input_tokens is not None else None,
+        output_tokens=output_tokens,
+        total_tokens=total_tokens,
+    )
+
+
+def _record_model_calls(model: Any) -> None:
+    """Record each call made through a NOOA model client.
+
+    The record shows whether a model call started, so a failed invocation keeps
+    usage only after one did. NOOA's task token accumulator counts only
+    ``CompletionClient`` calls, so a Responses API run also needs the usage
+    each call returns.
+    """
+    call = model.call
+    acall = model.acall
+
+    def start() -> _ModelCalls | None:
+        calls = _model_calls.get()
+        if calls is not None:
+            calls.started = True
+        return calls
+
+    def record(calls: _ModelCalls | None, response: Any) -> Any:
+        if calls is not None:
+            calls.usages.append(getattr(response, "usage", None))
+        return response
+
+    def recorded_call(*args: Any, **kwargs: Any) -> Any:
+        calls = start()
+        return record(calls, call(*args, **kwargs))
+
+    async def recorded_acall(*args: Any, **kwargs: Any) -> Any:
+        calls = start()
+        return record(calls, await acall(*args, **kwargs))
+
+    model.call = recorded_call
+    model.acall = recorded_acall
+
+
 def _success_output(native: dict[str, Any], usage: AgentUsage | None) -> AgentRunResult:
     response = native.get("response")
     structured = native.get("result")
@@ -201,6 +310,7 @@ def __init__(self) -> None:
         self._workspace: Path | None = None
         self._instruction: str | None = None
         self._telemetry: RelayTelemetry | None = None
+        self._uses_responses_api = False
 
     async def start(self, payload: dict[str, Any]) -> None:
         if self._bench_agent_factory is not None:
@@ -219,8 +329,12 @@ async def start(self, payload: dict[str, Any]) -> None:
             base_dir = Path(common_utils.base_dir(payload))
             workspace = _workspace(payload)
             models = await build_models(config)
+            from nooa.unifiedllm import ResponsesClient
             from nooa_bench.bench_agent import BenchAgent
 
+            uses_responses_api = isinstance(models[role], ResponsesClient)
+            _record_model_calls(models[role])
+
             bench_agent_factory = BenchAgent
             telemetry = RelayTelemetry(
                 agent_name=common_utils.agent_name(payload),
@@ -255,6 +369,7 @@ async def start(self, payload: dict[str, Any]) -> None:
         self._workspace = workspace
         self._instruction = _system_instruction(config)
         self._telemetry = telemetry
+        self._uses_responses_api = uses_responses_api
 
     async def invoke(
         self,
@@ -304,22 +419,28 @@ async def invoke_agent() -> AgentRunResult:
             }
             if self._instruction is not None:
                 task_input["instructions"] = self._instruction
+            calls = _ModelCalls()
             try:
-                from nooa.runtime.token_usage import get_task_tokens
                 from nooa.runtime.token_usage import start_task_tokens
 
                 # BenchAgent creates a replacement shell for each task. Close the
                 # previous one first so persistent Fabric runtimes do not leak it.
                 await _close_shell(agent)
                 start_task_tokens()
-                native = await agent._run_evaluation(task_input)
-                usage = _usage(get_task_tokens())
+                calls_token = _model_calls.set(calls)
+                try:
+                    native = await agent._run_evaluation(task_input)
+                finally:
+                    _model_calls.reset(calls_token)
+                usage = self._collected_usage(calls)
                 if not isinstance(native, dict):
                     raise ValueError("BenchAgent result must be an object")
                 if native.get("success") is not True:
+                    # A failed task still spent tokens; keep them for cost accounting.
                     return _failure_output(
                         "nooa_bench_task_failed",
                         "OO Agents BenchAgent reported task failure; inspect adapter stderr for details",
+                        usage,
                     )
                 return _success_output(native, usage)
             except asyncio.CancelledError:
@@ -332,6 +453,7 @@ async def invoke_agent() -> AgentRunResult:
                 return _failure_output(
                     "nooa_bench_invoke_failed",
                     "OO Agents BenchAgent invocation failed; inspect adapter stderr for details",
+                    self._usage_after_error(calls),
                 )
 
         try:
@@ -369,6 +491,28 @@ async def invoke_agent() -> AgentRunResult:
             result = relay_invocation.result
         return _with_telemetry(result, relay_invocation.report)
 
+    def _collected_usage(self, calls: _ModelCalls) -> AgentUsage | None:
+        if self._uses_responses_api:
+            return _responses_api_usage(calls.usages)
+        from nooa.runtime.token_usage import get_task_tokens
+
+        return _usage(get_task_tokens())
+
+    def _usage_after_error(self, calls: _ModelCalls) -> AgentUsage | None:
+        """Return the usage collected before an invocation error, if a model call started."""
+        if not calls.started:
+            return None
+        try:
+            return self._collected_usage(calls)
+        except Exception as error:
+            # Keep the original failure result; its usage just stays unknown.
+            LOGGER.error(
+                "OO Agents BenchAgent usage collection failed after an invocation error "
+                "(error_type=%s)",
+                type(error).__name__,
+            )
+            return None
+
     async def stop(self) -> None:
         models = self._models
         telemetry = self._telemetry
@@ -379,6 +523,7 @@ async def stop(self) -> None:
         self._workspace = None
         self._instruction = None
         self._telemetry = None
+        self._uses_responses_api = False
         primary: BaseException | None = None
         if telemetry is not None:
             try: