From 3beeb0f34309848748bf620e21ac398c20108a9b Mon Sep 17 00:00:00 2001 From: mnajafian-nv Date: Mon, 19 Oct 2026 14:55:32 -0700 Subject: [PATCH 1/4] fix: report NOOA BenchAgent usage for the Responses API client NOOA 0.0.10 emits token_usage events only for CompletionClient. ResponsesClient returns usage on each response but does not feed the task token accumulator, so Fabric reported input, output and total tokens as zero for Responses client runs. When the selected model is a ResponsesClient, wrap its synchronous and asynchronous calls and collect each returned usage report in an invocation-scoped ContextVar. Sum provider-reported input, cached-input, output and total tokens. If any returned report omits a counter, leave that dimension unknown. Preserve explicit zeros and mark input tokens as including cached tokens. The CompletionClient path is unchanged. Tests drive the real ResponsesClient directly with queued LiteLLM responses while leaving NOOA's accumulator unpatched. They cover multi-call sums, explicit zeros, missing usage, realistic missing cached-token details, and a defensive summing case for an omitted output count. Signed-off-by: mnajafian-nv --- adapters/python/nooa/docs/bench-agent.md | 8 +- .../nooa/bench_adapter.py | 103 +++++++- tests/adapters/test_nooa_bench_adapter.py | 242 ++++++++++++++++++ .../src/nooa/unifiedllm/__init__.py | 4 + 4 files changed, 354 insertions(+), 3 deletions(-) diff --git a/adapters/python/nooa/docs/bench-agent.md b/adapters/python/nooa/docs/bench-agent.md index c0b22a1fa..4af4f3743 100644 --- a/adapters/python/nooa/docs/bench-agent.md +++ b/adapters/python/nooa/docs/bench-agent.md @@ -23,7 +23,13 @@ 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..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. The adapter accepts exactly one model. Before each task, it closes the shell that `_run_evaluation()` replaces. During runtime shutdown, it closes the diff --git a/adapters/python/nooa/src/nemo_fabric_adapters/nooa/bench_adapter.py b/adapters/python/nooa/src/nemo_fabric_adapters/nooa/bench_adapter.py index ff103932b..727012060 100644 --- a/adapters/python/nooa/src/nemo_fabric_adapters/nooa/bench_adapter.py +++ b/adapters/python/nooa/src/nemo_fabric_adapters/nooa/bench_adapter.py @@ -10,6 +10,7 @@ import inspect import logging from collections.abc import Callable +from contextvars import ContextVar from pathlib import Path from typing import Any @@ -31,6 +32,11 @@ HARNESS = "nooa-bench" MODE = "bench_agent" LOGGER = logging.getLogger(__name__) +_MAX_TOKEN_COUNT = (1 << 64) - 1 +# Per-invocation list of the ``usage`` value from each Responses API call. +_responses_usage: ContextVar[list[Any] | None] = ContextVar( + "nooa_bench_responses_usage", default=None +) def main() -> None: @@ -148,6 +154,82 @@ 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_responses_usage(model: Any) -> None: + """Record the usage of each call made through a NOOA ``ResponsesClient``. + + NOOA's task token accumulator counts only ``CompletionClient`` calls, so a + Responses API run would otherwise report zero tokens. + """ + call = model.call + acall = model.acall + + def record(response: Any) -> Any: + usages = _responses_usage.get() + if usages is not None: + usages.append(getattr(response, "usage", None)) + return response + + def recorded_call(*args: Any, **kwargs: Any) -> Any: + return record(call(*args, **kwargs)) + + async def recorded_acall(*args: Any, **kwargs: Any) -> Any: + return record(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 +283,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 +302,13 @@ 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) + if uses_responses_api: + _record_responses_usage(models[role]) + bench_agent_factory = BenchAgent telemetry = RelayTelemetry( agent_name=common_utils.agent_name(payload), @@ -255,6 +343,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, @@ -312,8 +401,17 @@ async def invoke_agent() -> AgentRunResult: # 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()) + responses_usages: list[Any] = [] + responses_token = _responses_usage.set(responses_usages) + try: + native = await agent._run_evaluation(task_input) + finally: + _responses_usage.reset(responses_token) + usage = ( + _responses_api_usage(responses_usages) + if self._uses_responses_api + else _usage(get_task_tokens()) + ) if not isinstance(native, dict): raise ValueError("BenchAgent result must be an object") if native.get("success") is not True: @@ -379,6 +477,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: diff --git a/tests/adapters/test_nooa_bench_adapter.py b/tests/adapters/test_nooa_bench_adapter.py index a0e3dd6f2..aadf30b60 100644 --- a/tests/adapters/test_nooa_bench_adapter.py +++ b/tests/adapters/test_nooa_bench_adapter.py @@ -96,6 +96,7 @@ def bench_dependencies_fixture(monkeypatch: pytest.MonkeyPatch): mock_get_llm_client = MagicMock(return_value=mock_model) unifiedllm = types.ModuleType("nooa.unifiedllm") unifiedllm.get_llm_client = mock_get_llm_client + unifiedllm.ResponsesClient = type("ResponsesClient", (), {}) monkeypatch.setitem(sys.modules, unifiedllm.__name__, unifiedllm) def agent_double(index: int) -> tuple[MagicMock, MagicMock, MagicMock]: @@ -176,6 +177,247 @@ async def invoke_telemetry(*, call: Any, **_kwargs: Any): } +@pytest.fixture(name="responses_api_bench") +def responses_api_bench_fixture(monkeypatch: pytest.MonkeyPatch) -> list[Any]: + """Run a real NOOA ResponsesClient against queued Responses API replies. + + The BenchAgent double calls the real ``ResponsesClient`` directly, once per + queued reply. NOOA's task token accumulator is left in place, unpatched. + """ + pytest.importorskip("nooa.unifiedllm") + import litellm + + replies: list[Any] = [] + + async def aresponses(**_kwargs: Any) -> Any: + return replies.pop(0) + + monkeypatch.setattr(litellm, "aresponses", aresponses) + + class BenchAgentDouble: + def __init__(self, llm: Any) -> None: + self.llm = llm + self.shell = None + self.event_manager = object() + + async def _run_evaluation(self, task_input: dict[str, Any]) -> dict[str, Any]: + while replies: + await self.llm.acall( + [{"role": "user", "content": task_input["user_message"]}] + ) + return {"response": "done", "success": True, "result": None} + + bench_module = types.ModuleType("nooa_bench.bench_agent") + bench_module.BenchAgent = BenchAgentDouble + bench_package = types.ModuleType("nooa_bench") + bench_package.__path__ = [] # type: ignore[attr-defined] + monkeypatch.setitem(sys.modules, bench_package.__name__, bench_package) + monkeypatch.setitem(sys.modules, bench_module.__name__, bench_module) + + telemetry = MagicMock(name="relay_telemetry") + + async def invoke_telemetry(*, call: Any, **_kwargs: Any): + return nooa_telemetry.RelayInvocation( + called=True, + result=await call(), + report=None, + ) + + telemetry.invoke = AsyncMock(side_effect=invoke_telemetry) + telemetry.close = AsyncMock() + monkeypatch.setattr( + bench_adapter, "RelayTelemetry", MagicMock(return_value=telemetry) + ) + monkeypatch.setenv("OPENAI_API_KEY", "test-key") + return replies + + +def _responses_api_reply(usage: Any) -> Any: + from litellm.types.llms.openai import ResponsesAPIResponse + + reply = ResponsesAPIResponse( + id="resp-1", + created_at=0, + model="fixture-model", + object="response", + status="completed", + output=[ + { + "type": "message", + "id": "msg-1", + "role": "assistant", + "status": "completed", + "content": [{"type": "output_text", "text": "done", "annotations": []}], + } + ], + parallel_tool_calls=False, + tool_choice="auto", + tools=[], + ) + return reply.model_copy(update={"usage": usage}) + + +def _usage_report(**counters: Any) -> Any: + from litellm.types.llms.openai import ResponseAPIUsage + + return ResponseAPIUsage(**counters) + + +async def _invoke_responses_api_bench(tmp_path: Path) -> Any: + payload = _start_payload(tmp_path) + payload["config"] = AgentConfig.from_mapping( + { + "models": { + "default": { + "provider": "openai", + "model": "openai/fixture-model", + "base_url": "https://model.example.test/v1", + "settings": {"client_type": "responses"}, + } + } + } + ) + runtime = bench_adapter.BenchRuntime() + await runtime.start(payload) + try: + return await runtime.invoke( + AgentRunRequest.from_mapping({"input": "Fix the failing test."}), + _runtime_context(), + ) + finally: + await runtime.stop() + + +async def test_bench_responses_api_usage_sums_every_call( + tmp_path: Path, + responses_api_bench: list[Any], +): + responses_api_bench.extend( + [ + _responses_api_reply( + _usage_report( + input_tokens=100, + input_tokens_details={"cached_tokens": 40}, + output_tokens=20, + output_tokens_details={"reasoning_tokens": 5}, + total_tokens=120, + ) + ), + _responses_api_reply( + _usage_report( + input_tokens=150, + input_tokens_details={"cached_tokens": 90}, + output_tokens=30, + output_tokens_details={"reasoning_tokens": 0}, + total_tokens=180, + ) + ), + ] + ) + + result = await _invoke_responses_api_bench(tmp_path) + + assert result.status is AgentRunStatus.SUCCEEDED + assert result.usage is not None + assert result.usage.to_mapping() == { + "input_tokens": 250, + "cached_input_tokens": 130, + "input_tokens_include_cache": True, + "output_tokens": 50, + "total_tokens": 300, + } + + +async def test_bench_responses_api_reported_zero_usage_stays_zero( + tmp_path: Path, + responses_api_bench: list[Any], +): + responses_api_bench.append( + _responses_api_reply( + _usage_report( + input_tokens=0, + input_tokens_details={"cached_tokens": 0}, + output_tokens=0, + total_tokens=0, + ) + ) + ) + + result = await _invoke_responses_api_bench(tmp_path) + + assert result.usage is not None + assert result.usage.to_mapping() == { + "input_tokens": 0, + "cached_input_tokens": 0, + "input_tokens_include_cache": True, + "output_tokens": 0, + "total_tokens": 0, + } + + +async def test_bench_responses_api_call_without_usage_leaves_usage_unknown( + tmp_path: Path, + responses_api_bench: list[Any], +): + responses_api_bench.extend( + [ + _responses_api_reply( + _usage_report( + input_tokens=100, + input_tokens_details={"cached_tokens": 40}, + output_tokens=20, + total_tokens=120, + ) + ), + _responses_api_reply(None), + ] + ) + + result = await _invoke_responses_api_bench(tmp_path) + + assert result.status is AgentRunStatus.SUCCEEDED + assert result.usage is None + + +async def test_bench_responses_api_partial_usage_keeps_omitted_counters_unknown( + tmp_path: Path, + responses_api_bench: list[Any], +): + from litellm.types.llms.openai import ResponseAPIUsage + + responses_api_bench.extend( + [ + _responses_api_reply( + _usage_report( + input_tokens=100, + input_tokens_details={"cached_tokens": 40}, + output_tokens=20, + total_tokens=120, + ) + ), + # No input_tokens_details, so the cached share of this call is unknown. + _responses_api_reply( + _usage_report(input_tokens=60, output_tokens=10, total_tokens=70) + ), + # Defensive: a usage mapping without output_tokens. litellm 1.100.1 + # cannot hand NOOA 0.0.10 such a reply (NOOA raises first), so this + # leg covers the summing rule only. + _responses_api_reply( + ResponseAPIUsage.model_construct(input_tokens=30, total_tokens=30) + ), + ] + ) + + result = await _invoke_responses_api_bench(tmp_path) + + assert result.usage is not None + assert result.usage.to_mapping() == { + "input_tokens": 190, + "input_tokens_include_cache": True, + "total_tokens": 220, + } + + def test_bench_descriptor_is_a_closed_harness_adapter(tmp_path: Path): descriptor = json.loads( (NOOA_ROOT / "nooa-bench.fabric-adapter.json").read_text(encoding="utf-8") diff --git a/tests/fixtures/nooa_bench/src/nooa/unifiedllm/__init__.py b/tests/fixtures/nooa_bench/src/nooa/unifiedllm/__init__.py index c6116218d..51af21534 100644 --- a/tests/fixtures/nooa_bench/src/nooa/unifiedllm/__init__.py +++ b/tests/fixtures/nooa_bench/src/nooa/unifiedllm/__init__.py @@ -17,6 +17,10 @@ async def aclose(self) -> None: return None +class ResponsesClient: + """Stands in for NOOA's Responses API client type; the fixture never builds one.""" + + def get_llm_client( name: str, *, From 93bf3b1ce59bea41b654c0b0c502bb04aebea45c Mon Sep 17 00:00:00 2001 From: mnajafian-nv Date: Mon, 19 Oct 2026 17:02:38 -0700 Subject: [PATCH 2/4] test: cover the sync Responses call and a reused runtime The Responses usage tests drove only acall, and each one started a fresh runtime. The adapter also wraps the synchronous call, and a persistent runtime serves many invocations with the same wrapped client, so neither path was pinned down. The BenchAgent double now also drains a queue of synchronous replies through ResponsesClient.call, and a new test checks that a sync call's counters reach the invocation usage. A second test invokes one runtime twice with different replies and checks each result's usage on its own, so usage from the first invocation cannot carry over into the second. Signed-off-by: mnajafian-nv --- tests/adapters/test_nooa_bench_adapter.py | 117 ++++++++++++++++++++-- 1 file changed, 111 insertions(+), 6 deletions(-) diff --git a/tests/adapters/test_nooa_bench_adapter.py b/tests/adapters/test_nooa_bench_adapter.py index aadf30b60..97f31b452 100644 --- a/tests/adapters/test_nooa_bench_adapter.py +++ b/tests/adapters/test_nooa_bench_adapter.py @@ -177,12 +177,23 @@ async def invoke_telemetry(*, call: Any, **_kwargs: Any): } +@pytest.fixture(name="sync_responses_api_replies") +def sync_responses_api_replies_fixture() -> list[Any]: + """Replies that the BenchAgent double requests through the synchronous ``call``.""" + return [] + + @pytest.fixture(name="responses_api_bench") -def responses_api_bench_fixture(monkeypatch: pytest.MonkeyPatch) -> list[Any]: +def responses_api_bench_fixture( + monkeypatch: pytest.MonkeyPatch, + sync_responses_api_replies: list[Any], +) -> list[Any]: """Run a real NOOA ResponsesClient against queued Responses API replies. The BenchAgent double calls the real ``ResponsesClient`` directly, once per - queued reply. NOOA's task token accumulator is left in place, unpatched. + queued reply: ``acall`` for this fixture's queue, then ``call`` for + ``sync_responses_api_replies``. NOOA's task token accumulator is left in + place, unpatched. """ pytest.importorskip("nooa.unifiedllm") import litellm @@ -192,7 +203,11 @@ def responses_api_bench_fixture(monkeypatch: pytest.MonkeyPatch) -> list[Any]: async def aresponses(**_kwargs: Any) -> Any: return replies.pop(0) + def responses(**_kwargs: Any) -> Any: + return sync_responses_api_replies.pop(0) + monkeypatch.setattr(litellm, "aresponses", aresponses) + monkeypatch.setattr(litellm, "responses", responses) class BenchAgentDouble: def __init__(self, llm: Any) -> None: @@ -201,10 +216,11 @@ def __init__(self, llm: Any) -> None: self.event_manager = object() async def _run_evaluation(self, task_input: dict[str, Any]) -> dict[str, Any]: + messages = [{"role": "user", "content": task_input["user_message"]}] while replies: - await self.llm.acall( - [{"role": "user", "content": task_input["user_message"]}] - ) + await self.llm.acall(messages) + while sync_responses_api_replies: + self.llm.call(messages) return {"response": "done", "success": True, "result": None} bench_module = types.ModuleType("nooa_bench.bench_agent") @@ -263,7 +279,7 @@ def _usage_report(**counters: Any) -> Any: return ResponseAPIUsage(**counters) -async def _invoke_responses_api_bench(tmp_path: Path) -> Any: +async def _start_responses_api_bench(tmp_path: Path) -> Any: payload = _start_payload(tmp_path) payload["config"] = AgentConfig.from_mapping( { @@ -279,6 +295,11 @@ async def _invoke_responses_api_bench(tmp_path: Path) -> Any: ) runtime = bench_adapter.BenchRuntime() await runtime.start(payload) + return runtime + + +async def _invoke_responses_api_bench(tmp_path: Path) -> Any: + runtime = await _start_responses_api_bench(tmp_path) try: return await runtime.invoke( AgentRunRequest.from_mapping({"input": "Fix the failing test."}), @@ -328,6 +349,90 @@ async def test_bench_responses_api_usage_sums_every_call( } +async def test_bench_responses_api_usage_counts_sync_calls( + tmp_path: Path, + responses_api_bench: list[Any], + sync_responses_api_replies: list[Any], +): + sync_responses_api_replies.append( + _responses_api_reply( + _usage_report( + input_tokens=80, + input_tokens_details={"cached_tokens": 30}, + output_tokens=15, + total_tokens=95, + ) + ) + ) + + result = await _invoke_responses_api_bench(tmp_path) + + assert not responses_api_bench + assert result.usage is not None + assert result.usage.to_mapping() == { + "input_tokens": 80, + "cached_input_tokens": 30, + "input_tokens_include_cache": True, + "output_tokens": 15, + "total_tokens": 95, + } + + +async def test_bench_responses_api_usage_is_per_invocation_on_a_reused_runtime( + tmp_path: Path, + responses_api_bench: list[Any], +): + runtime = await _start_responses_api_bench(tmp_path) + try: + responses_api_bench.append( + _responses_api_reply( + _usage_report( + input_tokens=100, + input_tokens_details={"cached_tokens": 40}, + output_tokens=20, + total_tokens=120, + ) + ) + ) + first = await runtime.invoke( + AgentRunRequest.from_mapping({"input": "Fix the first task."}), + _runtime_context(), + ) + responses_api_bench.append( + _responses_api_reply( + _usage_report( + input_tokens=7, + input_tokens_details={"cached_tokens": 0}, + output_tokens=3, + total_tokens=10, + ) + ) + ) + second = await runtime.invoke( + AgentRunRequest.from_mapping({"input": "Fix the second task."}), + _runtime_context(invocation_id="invocation-2", request_id="request-2"), + ) + finally: + await runtime.stop() + + assert first.usage is not None + assert first.usage.to_mapping() == { + "input_tokens": 100, + "cached_input_tokens": 40, + "input_tokens_include_cache": True, + "output_tokens": 20, + "total_tokens": 120, + } + assert second.usage is not None + assert second.usage.to_mapping() == { + "input_tokens": 7, + "cached_input_tokens": 0, + "input_tokens_include_cache": True, + "output_tokens": 3, + "total_tokens": 10, + } + + async def test_bench_responses_api_reported_zero_usage_stays_zero( tmp_path: Path, responses_api_bench: list[Any], From 6c2810e72a0573a4013e3481818c2b7441400dc4 Mon Sep 17 00:00:00 2001 From: mnajafian-nv Date: Mon, 19 Oct 2026 20:24:39 -0700 Subject: [PATCH 3/4] fix: keep NOOA BenchAgent token usage on failed tasks When BenchAgent reported a task as failed, the adapter computed the invocation's usage and then returned a failed result without it. The tokens a failed task spent dropped out of run totals, so cost per successful task read too low. The adapter contract asks adapters to keep available usage on unsuccessful results. Attach the same usage to the failed result for both the completion and Responses API clients. Counting rules and successful results are unchanged, and each invocation still reports only its own calls. Signed-off-by: mnajafian-nv --- adapters/python/nooa/docs/bench-agent.md | 2 + .../nooa/bench_adapter.py | 9 +- tests/adapters/test_nooa_bench_adapter.py | 121 +++++++++++++++++- 3 files changed, 129 insertions(+), 3 deletions(-) diff --git a/adapters/python/nooa/docs/bench-agent.md b/adapters/python/nooa/docs/bench-agent.md index 4af4f3743..01a7d1b0e 100644 --- a/adapters/python/nooa/docs/bench-agent.md +++ b/adapters/python/nooa/docs/bench-agent.md @@ -30,6 +30,8 @@ One NeMo Fabric invocation maps to one BenchAgent task: 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. The adapter accepts exactly one model. Before each task, it closes the shell that `_run_evaluation()` replaces. During runtime shutdown, it closes the diff --git a/adapters/python/nooa/src/nemo_fabric_adapters/nooa/bench_adapter.py b/adapters/python/nooa/src/nemo_fabric_adapters/nooa/bench_adapter.py index 727012060..5043a19e2 100644 --- a/adapters/python/nooa/src/nemo_fabric_adapters/nooa/bench_adapter.py +++ b/adapters/python/nooa/src/nemo_fabric_adapters/nooa/bench_adapter.py @@ -119,7 +119,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={ @@ -130,6 +134,7 @@ def _failure_output(code: str, message: str) -> AgentRunResult: "completed": False, }, error=AgentRunError(code=code, message=message), + usage=usage, ) @@ -415,9 +420,11 @@ async def invoke_agent() -> AgentRunResult: 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: diff --git a/tests/adapters/test_nooa_bench_adapter.py b/tests/adapters/test_nooa_bench_adapter.py index 97f31b452..1ab286292 100644 --- a/tests/adapters/test_nooa_bench_adapter.py +++ b/tests/adapters/test_nooa_bench_adapter.py @@ -170,6 +170,7 @@ async def invoke_telemetry(*, call: Any, **_kwargs: Any): "initial_shells": mock_initial_shells, "model": mock_model, "relay_telemetry": mock_relay_telemetry, + "get_task_tokens": mock_get_task_tokens, "start_task_tokens": mock_start_task_tokens, "task_shell": mock_task_shells[0], "task_shells": mock_task_shells, @@ -183,16 +184,24 @@ def sync_responses_api_replies_fixture() -> list[Any]: return [] +@pytest.fixture(name="failing_responses_api_tasks") +def failing_responses_api_tasks_fixture() -> set[str]: + """User messages whose task the BenchAgent double reports as failed.""" + return set() + + @pytest.fixture(name="responses_api_bench") def responses_api_bench_fixture( monkeypatch: pytest.MonkeyPatch, sync_responses_api_replies: list[Any], + failing_responses_api_tasks: set[str], ) -> list[Any]: """Run a real NOOA ResponsesClient against queued Responses API replies. The BenchAgent double calls the real ``ResponsesClient`` directly, once per queued reply: ``acall`` for this fixture's queue, then ``call`` for - ``sync_responses_api_replies``. NOOA's task token accumulator is left in + ``sync_responses_api_replies``. It reports failure for a user message in + ``failing_responses_api_tasks``. NOOA's task token accumulator is left in place, unpatched. """ pytest.importorskip("nooa.unifiedllm") @@ -221,7 +230,8 @@ async def _run_evaluation(self, task_input: dict[str, Any]) -> dict[str, Any]: await self.llm.acall(messages) while sync_responses_api_replies: self.llm.call(messages) - return {"response": "done", "success": True, "result": None} + success = task_input["user_message"] not in failing_responses_api_tasks + return {"response": "done", "success": success, "result": None} bench_module = types.ModuleType("nooa_bench.bench_agent") bench_module.BenchAgent = BenchAgentDouble @@ -523,6 +533,67 @@ async def test_bench_responses_api_partial_usage_keeps_omitted_counters_unknown( } +async def test_bench_responses_api_failed_task_keeps_its_own_usage( + tmp_path: Path, + responses_api_bench: list[Any], + failing_responses_api_tasks: set[str], +): + failing_responses_api_tasks.add("Fix the first task.") + runtime = await _start_responses_api_bench(tmp_path) + try: + responses_api_bench.append( + _responses_api_reply( + _usage_report( + input_tokens=100, + input_tokens_details={"cached_tokens": 40}, + output_tokens=20, + total_tokens=120, + ) + ) + ) + failed = await runtime.invoke( + AgentRunRequest.from_mapping({"input": "Fix the first task."}), + _runtime_context(), + ) + responses_api_bench.append( + _responses_api_reply( + _usage_report( + input_tokens=7, + input_tokens_details={"cached_tokens": 0}, + output_tokens=3, + total_tokens=10, + ) + ) + ) + succeeded = await runtime.invoke( + AgentRunRequest.from_mapping({"input": "Fix the second task."}), + _runtime_context(invocation_id="invocation-2", request_id="request-2"), + ) + finally: + await runtime.stop() + + assert failed.status is AgentRunStatus.FAILED + assert failed.error is not None + assert failed.error.code == "nooa_bench_task_failed" + assert failed.usage is not None + assert failed.usage.to_mapping() == { + "input_tokens": 100, + "cached_input_tokens": 40, + "input_tokens_include_cache": True, + "output_tokens": 20, + "total_tokens": 120, + } + assert succeeded.status is AgentRunStatus.SUCCEEDED + assert succeeded.usage is not None + assert succeeded.usage.to_mapping() == { + "input_tokens": 7, + "cached_input_tokens": 0, + "input_tokens_include_cache": True, + "output_tokens": 3, + "total_tokens": 10, + } + + def test_bench_descriptor_is_a_closed_harness_adapter(tmp_path: Path): descriptor = json.loads( (NOOA_ROOT / "nooa-bench.fabric-adapter.json").read_text(encoding="utf-8") @@ -776,3 +847,49 @@ def test_bench_persistent_host_runs_start_invoke_stop(tmp_path: Path): "total_tokens": 16, } assert (workspace / "bench-agent-result.txt").is_file() + + +async def test_bench_runtime_failed_task_keeps_its_own_usage( + tmp_path: Path, + bench_dependencies: dict[str, MagicMock], +): + os.environ["OPENAI_API_KEY"] = "test-key" + bench_dependencies["agents"][0]._run_evaluation.side_effect = None + bench_dependencies["agents"][0]._run_evaluation.return_value = { + "response": "", + "success": False, + } + bench_dependencies["get_task_tokens"].side_effect = [ + {"n_input_tokens": 30, "n_output_tokens": 9}, + {"n_input_tokens": 5, "n_output_tokens": 2}, + ] + runtime = bench_adapter.BenchRuntime() + await runtime.start(_start_payload(tmp_path)) + + failed = await runtime.invoke( + AgentRunRequest.from_mapping({"input": "Fix the first task."}), + _runtime_context(), + ) + succeeded = await runtime.invoke( + AgentRunRequest.from_mapping({"input": "Fix the second task."}), + _runtime_context(invocation_id="invocation-2", request_id="request-2"), + ) + await runtime.stop() + + assert bench_dependencies["start_task_tokens"].call_count == 2 + assert failed.status is AgentRunStatus.FAILED + assert failed.error is not None + assert failed.error.code == "nooa_bench_task_failed" + assert failed.usage is not None + assert failed.usage.to_mapping() == { + "input_tokens": 30, + "output_tokens": 9, + "total_tokens": 39, + } + assert succeeded.status is AgentRunStatus.SUCCEEDED + assert succeeded.usage is not None + assert succeeded.usage.to_mapping() == { + "input_tokens": 5, + "output_tokens": 2, + "total_tokens": 7, + } From 2b0f42530740007779b043e6f2da88303a185a4c Mon Sep 17 00:00:00 2001 From: mnajafian-nv Date: Tue, 20 Oct 2026 01:25:18 -0700 Subject: [PATCH 4/4] fix: keep NOOA BenchAgent usage when a task raises The adapter collected usage only after BenchAgent returned. When the task raised after a model call, or returned a malformed result, the failed result lost the tokens already spent. Mark the start of each model call on the selected client, for both the completion and Responses API clients. After an invocation error, keep the collected usage once a model call has started, and leave it unknown before then. A usage collection error keeps the original failure result. The subprocess fixture model gains the call methods that NOOA clients have. Signed-off-by: mnajafian-nv --- adapters/python/nooa/docs/bench-agent.md | 4 +- .../nooa/bench_adapter.py | 87 +++++++++++---- tests/adapters/test_nooa_bench_adapter.py | 105 +++++++++++++++++- .../src/nooa/unifiedllm/__init__.py | 7 ++ 4 files changed, 177 insertions(+), 26 deletions(-) diff --git a/adapters/python/nooa/docs/bench-agent.md b/adapters/python/nooa/docs/bench-agent.md index 01a7d1b0e..f0a5b8fa5 100644 --- a/adapters/python/nooa/docs/bench-agent.md +++ b/adapters/python/nooa/docs/bench-agent.md @@ -31,7 +31,9 @@ One NeMo Fabric invocation maps to one BenchAgent task: `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. + 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 diff --git a/adapters/python/nooa/src/nemo_fabric_adapters/nooa/bench_adapter.py b/adapters/python/nooa/src/nemo_fabric_adapters/nooa/bench_adapter.py index 5043a19e2..69f92af6c 100644 --- a/adapters/python/nooa/src/nemo_fabric_adapters/nooa/bench_adapter.py +++ b/adapters/python/nooa/src/nemo_fabric_adapters/nooa/bench_adapter.py @@ -11,6 +11,8 @@ 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 @@ -33,9 +35,20 @@ MODE = "bench_agent" LOGGER = logging.getLogger(__name__) _MAX_TOKEN_COUNT = (1 << 64) - 1 -# Per-invocation list of the ``usage`` value from each Responses API call. -_responses_usage: ContextVar[list[Any] | None] = ContextVar( - "nooa_bench_responses_usage", default=None + + +@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 ) @@ -210,26 +223,35 @@ def _responses_api_usage(usages: list[Any]) -> AgentUsage | None: ) -def _record_responses_usage(model: Any) -> None: - """Record the usage of each call made through a NOOA ``ResponsesClient``. +def _record_model_calls(model: Any) -> None: + """Record each call made through a NOOA model client. - NOOA's task token accumulator counts only ``CompletionClient`` calls, so a - Responses API run would otherwise report zero tokens. + 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 record(response: Any) -> Any: - usages = _responses_usage.get() - if usages is not None: - usages.append(getattr(response, "usage", None)) + 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: - return record(call(*args, **kwargs)) + calls = start() + return record(calls, call(*args, **kwargs)) async def recorded_acall(*args: Any, **kwargs: Any) -> Any: - return record(await acall(*args, **kwargs)) + calls = start() + return record(calls, await acall(*args, **kwargs)) model.call = recorded_call model.acall = recorded_acall @@ -311,8 +333,7 @@ async def start(self, payload: dict[str, Any]) -> None: from nooa_bench.bench_agent import BenchAgent uses_responses_api = isinstance(models[role], ResponsesClient) - if uses_responses_api: - _record_responses_usage(models[role]) + _record_model_calls(models[role]) bench_agent_factory = BenchAgent telemetry = RelayTelemetry( @@ -398,25 +419,20 @@ 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() - responses_usages: list[Any] = [] - responses_token = _responses_usage.set(responses_usages) + calls_token = _model_calls.set(calls) try: native = await agent._run_evaluation(task_input) finally: - _responses_usage.reset(responses_token) - usage = ( - _responses_api_usage(responses_usages) - if self._uses_responses_api - else _usage(get_task_tokens()) - ) + _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: @@ -437,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: @@ -474,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 diff --git a/tests/adapters/test_nooa_bench_adapter.py b/tests/adapters/test_nooa_bench_adapter.py index 1ab286292..c387a9f62 100644 --- a/tests/adapters/test_nooa_bench_adapter.py +++ b/tests/adapters/test_nooa_bench_adapter.py @@ -93,6 +93,7 @@ def _runtime_context( def bench_dependencies_fixture(monkeypatch: pytest.MonkeyPatch): mock_model = MagicMock(name="model") mock_model.aclose = AsyncMock() + mock_model.acall = AsyncMock() mock_get_llm_client = MagicMock(return_value=mock_model) unifiedllm = types.ModuleType("nooa.unifiedllm") unifiedllm.get_llm_client = mock_get_llm_client @@ -190,18 +191,26 @@ def failing_responses_api_tasks_fixture() -> set[str]: return set() +@pytest.fixture(name="raising_responses_api_tasks") +def raising_responses_api_tasks_fixture() -> set[str]: + """User messages whose task the BenchAgent double raises on after its calls.""" + return set() + + @pytest.fixture(name="responses_api_bench") def responses_api_bench_fixture( monkeypatch: pytest.MonkeyPatch, sync_responses_api_replies: list[Any], failing_responses_api_tasks: set[str], + raising_responses_api_tasks: set[str], ) -> list[Any]: """Run a real NOOA ResponsesClient against queued Responses API replies. The BenchAgent double calls the real ``ResponsesClient`` directly, once per queued reply: ``acall`` for this fixture's queue, then ``call`` for ``sync_responses_api_replies``. It reports failure for a user message in - ``failing_responses_api_tasks``. NOOA's task token accumulator is left in + ``failing_responses_api_tasks`` and raises for one in + ``raising_responses_api_tasks``. NOOA's task token accumulator is left in place, unpatched. """ pytest.importorskip("nooa.unifiedllm") @@ -230,6 +239,8 @@ async def _run_evaluation(self, task_input: dict[str, Any]) -> dict[str, Any]: await self.llm.acall(messages) while sync_responses_api_replies: self.llm.call(messages) + if task_input["user_message"] in raising_responses_api_tasks: + raise RuntimeError("BenchAgent double failed after its model calls") success = task_input["user_message"] not in failing_responses_api_tasks return {"response": "done", "success": success, "result": None} @@ -594,6 +605,38 @@ async def test_bench_responses_api_failed_task_keeps_its_own_usage( } +async def test_bench_responses_api_invoke_error_keeps_collected_usage( + tmp_path: Path, + responses_api_bench: list[Any], + raising_responses_api_tasks: set[str], +): + raising_responses_api_tasks.add("Fix the failing test.") + responses_api_bench.append( + _responses_api_reply( + _usage_report( + input_tokens=100, + input_tokens_details={"cached_tokens": 40}, + output_tokens=20, + total_tokens=120, + ) + ) + ) + + result = await _invoke_responses_api_bench(tmp_path) + + assert result.status is AgentRunStatus.FAILED + assert result.error is not None + assert result.error.code == "nooa_bench_invoke_failed" + assert result.usage is not None + assert result.usage.to_mapping() == { + "input_tokens": 100, + "cached_input_tokens": 40, + "input_tokens_include_cache": True, + "output_tokens": 20, + "total_tokens": 120, + } + + def test_bench_descriptor_is_a_closed_harness_adapter(tmp_path: Path): descriptor = json.loads( (NOOA_ROOT / "nooa-bench.fabric-adapter.json").read_text(encoding="utf-8") @@ -893,3 +936,63 @@ async def test_bench_runtime_failed_task_keeps_its_own_usage( "output_tokens": 2, "total_tokens": 7, } + + +@pytest.mark.parametrize("native_result", [RuntimeError("model error"), "not a dict"]) +async def test_bench_runtime_invoke_error_after_model_call_keeps_usage( + tmp_path: Path, + bench_dependencies: dict[str, MagicMock], + native_result: Any, +): + os.environ["OPENAI_API_KEY"] = "test-key" + model = bench_dependencies["model"] + + async def run_evaluation(_task_input: dict[str, Any]) -> Any: + # The BenchAgent receives this model as ``llm``. + await model.acall([{"role": "user", "content": "Fix it."}]) + if isinstance(native_result, Exception): + raise native_result + return native_result + + bench_dependencies["agent"]._run_evaluation.side_effect = run_evaluation + runtime = bench_adapter.BenchRuntime() + await runtime.start(_start_payload(tmp_path)) + + result = await runtime.invoke( + AgentRunRequest.from_mapping({"input": "Fix it."}), + _runtime_context(), + ) + await runtime.stop() + + assert result.status is AgentRunStatus.FAILED + assert result.error is not None + assert result.error.code == "nooa_bench_invoke_failed" + assert result.usage is not None + assert result.usage.to_mapping() == { + "input_tokens": 12, + "output_tokens": 4, + "total_tokens": 16, + } + + +async def test_bench_runtime_invoke_error_before_model_call_has_no_usage( + tmp_path: Path, + bench_dependencies: dict[str, MagicMock], +): + os.environ["OPENAI_API_KEY"] = "test-key" + bench_dependencies["agent"]._run_evaluation.side_effect = RuntimeError( + "shell setup failed" + ) + runtime = bench_adapter.BenchRuntime() + await runtime.start(_start_payload(tmp_path)) + + result = await runtime.invoke( + AgentRunRequest.from_mapping({"input": "Fix it."}), + _runtime_context(), + ) + await runtime.stop() + + assert result.status is AgentRunStatus.FAILED + assert result.error is not None + assert result.error.code == "nooa_bench_invoke_failed" + assert result.usage is None diff --git a/tests/fixtures/nooa_bench/src/nooa/unifiedllm/__init__.py b/tests/fixtures/nooa_bench/src/nooa/unifiedllm/__init__.py index 51af21534..d1a391bdb 100644 --- a/tests/fixtures/nooa_bench/src/nooa/unifiedllm/__init__.py +++ b/tests/fixtures/nooa_bench/src/nooa/unifiedllm/__init__.py @@ -13,6 +13,13 @@ def __init__(self, name: str, **settings: Any) -> None: self.name = name self.settings = settings + # The adapter wraps NOOA's model-call methods; the fixture BenchAgent never calls them. + def call(self, *_args: Any, **_kwargs: Any) -> None: + return None + + async def acall(self, *_args: Any, **_kwargs: Any) -> None: + return None + async def aclose(self) -> None: return None