From 9d6bc218bf4debd918ea49ad44eb8e7789fba6e4 Mon Sep 17 00:00:00 2001 From: mnajafian-nv Date: Mon, 19 Oct 2026 15:28:16 -0700 Subject: [PATCH 01/26] fix: count every Deep Agents model call in token usage The Deep Agents adapter summed usage_metadata from the AI messages that reached graph state through the updates stream. Three cases came out wrong: - Deep Agents adds SummarizationMiddleware to every agent and subagent. Its summary call returns only text, so the summary model's message and usage never enter graph state, and an invocation that summarized under-reported input, output and total tokens. - When the invocation raised, including GraphRecursionError, the adapter returned no messages, so usage was None even though model calls had finished. - A count was summed over the calls that reported it and returned as known. ChatOpenAI also turns an omitted Chat Completions count into 0 and derives a missing total, so a partial provider report came out as measured zeros. The adapter now records usage from LangChain's chat-model callbacks, which fire for every chat-model call that runs in the adapter process: the main agent, declarative subagents and summarization. Agent Protocol subagents run remotely and are not included. A new recorder is created for each invocation, so earlier turns are never counted again. It outlives an exception, so a failed invocation still reports the fields that every counted call reported. Each token count and the cost are known only when every counted call in the invocation reported them, and an explicit zero stays a known zero. A request the provider rejects with a 4xx status (other than 408) contributes no usage. A 5xx response, timeout, cancellation or unfinished call leaves the invocation's counts and cost unknown, because the provider may have done work without reporting usage. OpenAI SDK retries happen inside one LangChain call, so a 5xx the SDK retries to success is not visible to the recorder, and only the final response's reported usage counts. When a response carries OpenAI-shaped raw usage, the raw counts decide what was reported: an omitted count stays unknown, the provider's total is used when present, and a derived total counts only when input and output were both reported. The graph tests run the real Deep Agents graph and ChatOpenAI against a scripted Chat Completions endpoint for summarization, parallel task subagents, 400, 408, 503 and timeout failures, the recursion limit, partial reports and explicit zeros. Separate ChatOpenAI.with_retry tests cover 429 then success and 503 then success. The existing usage tests now produce their messages through a LangChain chat-model call so the callbacks fire. Signed-off-by: mnajafian-nv --- adapters/python/deepagents/README.md | 28 +- .../deepagents/adapter.py | 254 +++++---- tests/adapters/test_deepagents.py | 491 +++++++++++++++--- tests/e2e/test_deepagents.py | 1 - tests/integrations/test_relay_scope_leak.py | 4 +- 5 files changed, 618 insertions(+), 160 deletions(-) diff --git a/adapters/python/deepagents/README.md b/adapters/python/deepagents/README.md index 2090da4cd..78e94ce51 100644 --- a/adapters/python/deepagents/README.md +++ b/adapters/python/deepagents/README.md @@ -227,10 +227,30 @@ service. The normalized result includes the final response, buffered messages and per-step events, LangGraph thread id, token usage (and cost when the provider -reports it), and errors. Usage aggregates the current turn across the main agent -and any delegated subagents (streamed with `subgraphs=True`). Configuration and -preflight failures (a missing credential, an absent `deepagents` package, or an -invalid MCP server) fail runtime start before an invocation is accepted. +reports it), and errors. Configuration and preflight failures (a missing +credential, an absent `deepagents` package, or an invalid MCP server) fail +runtime start before an invocation is accepted. + +Usage sums every chat-model call that runs in the adapter process for the +invocation: the main agent, declarative subagents, and Deep Agents +summarization. Agent Protocol subagents (entries with `graph_id`) run in their +remote service, and their tokens are not included. A token count or cost is +reported only when every counted call reported it; otherwise it stays unknown. +This also holds for a failed invocation. A request the provider rejects with a +4xx status (other than 408) contributes no usage. A 5xx response, timeout, +cancellation or unfinished call leaves the invocation's counts and cost +unknown, because the provider may have done work without reporting usage. +OpenAI SDK retries happen inside one LangChain call, so a 5xx the SDK retries +to success is not visible to the adapter, and only the final response's +reported usage counts. Usage replayed from a configured LangChain model cache +would be counted as reported; Fabric configures no such cache. + +When a response carries OpenAI-shaped usage in +`response_metadata.token_usage`, as `ChatOpenAI` responses without streaming +do, the raw counts decide what was reported: an omitted count stays unknown, +and the provider's total is used when present. Other providers and streamed +responses use LangChain's `usage_metadata`, which some integrations fill with 0 +for an omitted count. ## Runtime Lifecycle diff --git a/adapters/python/deepagents/src/nemo_fabric_adapters/deepagents/adapter.py b/adapters/python/deepagents/src/nemo_fabric_adapters/deepagents/adapter.py index 618957f1f..37771f419 100644 --- a/adapters/python/deepagents/src/nemo_fabric_adapters/deepagents/adapter.py +++ b/adapters/python/deepagents/src/nemo_fabric_adapters/deepagents/adapter.py @@ -21,9 +21,12 @@ from pathlib import Path from typing import Any from typing import NamedTuple +from uuid import UUID from langchain.agents.middleware import AgentMiddleware +from langchain_core.callbacks import BaseCallbackHandler from langchain_core.messages import ToolMessage +from langchain_core.outputs import LLMResult from langgraph.errors import GraphRecursionError from nemo_fabric_adapter_contract.models import AgentConfig from nemo_fabric_adapter_contract.models import AgentMcpServerConfig @@ -750,7 +753,7 @@ async def invoke( resumed=resumed, result_state=outcome.result_state, events=outcome.events or [], - turn_messages=outcome.turn_messages or [], + model_calls=outcome.model_calls or [], error=outcome.error, telemetry_runtime=telemetry_runtime, relay_artifacts=relay_artifacts, @@ -874,24 +877,26 @@ async def _invoke_agent( ) -> TurnOutcome: """Run one agent turn, normalizing an invocation failure into an error string.""" + model_usage = ModelCallUsage() try: - result_state, events, turn_messages = await invoke_compiled_agent( + result_state, events = await invoke_compiled_agent( self._agent, user_message, self._thread_id, - callbacks=callbacks, + callbacks=[*(callbacks or []), model_usage], ) except GraphRecursionError: return TurnOutcome( + model_calls=model_usage.calls(), error="Deep Agents reached its graph recursion limit.", error_code="deepagents_recursion_limit_reached", ) except Exception as exc: # normalized adapter failure - return TurnOutcome(error=_error_text(exc)) + return TurnOutcome(model_calls=model_usage.calls(), error=_error_text(exc)) return TurnOutcome( result_state=result_state, events=events, - turn_messages=turn_messages, + model_calls=model_usage.calls(), ) def _telemetry_output( @@ -966,18 +971,65 @@ def _apply_callbacks( return config +class ModelCallUsage(BaseCallbackHandler): + """Record what each chat-model call of one invocation reports. + + LangChain fires these callbacks for every chat-model call in the run, + including calls whose response never enters graph state, such as the Deep + Agents summarization call. The records live on the handler, so calls that + finished before an invocation raised are still reported. + + A call is recorded as unreported when it starts, and its usage replaces that + record when it ends. A request the provider rejects with a 4xx status (other + than 408) contributes no usage and is dropped. A 5xx response, timeout, + cancellation or unfinished call stays unreported and makes the invocation's + counts and cost unknown, because the provider may have done work without + reporting usage. + """ + + run_inline = True + + def __init__(self) -> None: + super().__init__() + self._calls: dict[UUID, dict[str, Any]] = {} + + def on_chat_model_start( + self, serialized: Any, messages: Any, *, run_id: UUID, **kwargs: Any + ) -> None: + self._calls[run_id] = {} + + def on_llm_start( + self, serialized: Any, prompts: Any, *, run_id: UUID, **kwargs: Any + ) -> None: + self._calls[run_id] = {} + + def on_llm_end(self, response: LLMResult, *, run_id: UUID, **kwargs: Any) -> None: + generations = response.generations[0] if response.generations else [] + message = getattr(generations[0], "message", None) if generations else None + self._calls[run_id] = {} if message is None else _message_to_dict(message) + + def on_llm_error( + self, error: BaseException, *, run_id: UUID, **kwargs: Any + ) -> None: + # Provider SDKs (OpenAI, Anthropic) raise errors carrying the HTTP status + # code when the provider answered with an error response. A 4xx other + # than 408 (request timeout) means the provider rejected the request. + status_code = getattr(error, "status_code", None) + if isinstance(status_code, int) and 400 <= status_code < 500: + if status_code != 408: + self._calls.pop(run_id, None) + + def calls(self) -> list[dict[str, Any]]: + return list(self._calls.values()) + + async def invoke_compiled_agent( agent: Any, user_message: str, thread_id: str | None, callbacks: list[Any] | None = None, -) -> tuple[Any, list[dict[str, Any]], list[dict[str, Any]]]: - """Run a compiled agent and return state, events, and current-turn messages. - - On a later runtime invocation the final ``values`` state also replays - prior-turn messages, so usage/cost must be aggregated from the messages - emitted *this* turn — the ``updates`` deltas — rather than the full final - state. +) -> tuple[Any, list[dict[str, Any]]]: + """Run a compiled agent and return its final state and per-step events. ``callbacks`` (e.g. the NeMo Relay callback handler) are appended to the LangGraph run config; any callbacks already present are preserved rather than @@ -992,12 +1044,10 @@ async def invoke_compiled_agent( config = config or None events: list[dict[str, Any]] = [] - turn_messages: list[Any] = [] final: Any = None # subgraphs=True surfaces subagent (delegated ``task``) steps as 3-tuples whose - # namespace identifies the subgraph. Folding their message deltas into this turn - # keeps usage/cost accurate for delegating runs; the final ``values`` state is - # taken from the main graph only (empty namespace). + # namespace identifies the subgraph, so they appear in ``events``; the final + # ``values`` state is taken from the main graph only (empty namespace). stream = ( agent.astream(inputs, config, stream_mode=["updates", "values"], subgraphs=True) if config @@ -1009,16 +1059,9 @@ async def invoke_compiled_agent( if namespace: event["subgraph"] = "/".join(str(part) for part in namespace) events.append(event) - for node_output in chunk.values(): - if isinstance(node_output, dict): - new = node_output.get("messages") - if isinstance(new, list): - turn_messages.extend(new) - elif new is not None: - turn_messages.append(new) elif mode == "values" and not namespace: final = chunk - return final, events, _dedup_messages(_messages_to_dicts(turn_messages)) + return final, events class Observability(NamedTuple): @@ -1034,7 +1077,7 @@ class TurnOutcome(NamedTuple): result_state: Any = None events: list[dict[str, Any]] | None = None - turn_messages: list[dict[str, Any]] | None = None + model_calls: list[dict[str, Any]] | None = None error: str | None = None error_code: str | None = None telemetry_error: str | None = None @@ -1151,7 +1194,7 @@ def normalize_output( resumed: bool, result_state: Any, events: list[dict[str, Any]], - turn_messages: list[dict[str, Any]], + model_calls: list[dict[str, Any]], error: str | None, telemetry_runtime: dict[str, Any] | None, relay_artifacts: list[dict[str, str]] | None, @@ -1160,9 +1203,9 @@ def normalize_output( ) -> dict[str, Any]: messages = _extract_messages(result_state) response = _final_response(messages) - # Aggregate usage/cost from this turn's messages only; the replayed final state - # replays prior-turn messages that must not be re-counted. - usage = _aggregate_usage(turn_messages) + # Usage comes from this invocation's model calls, not from the final state, + # which also replays earlier turns and omits internal calls like summarization. + usage = _aggregate_usage(model_calls) output: dict[str, Any] = { "harness": HARNESS, @@ -1210,26 +1253,6 @@ def _messages_to_dicts(raw: Any) -> list[dict[str, Any]]: return [item if isinstance(item, dict) else _message_to_dict(item) for item in raw] -def _dedup_messages(messages: list[dict[str, Any]]) -> list[dict[str, Any]]: - """Drop duplicate messages by id, preserving first-seen order. - - With ``subgraphs=True`` the same message can surface in both a subgraph's - updates and the parent stream; de-duplicating by id keeps usage aggregation - from counting a message twice. Messages without an id are always kept. - """ - - seen: set[str] = set() - unique: list[dict[str, Any]] = [] - for message in messages: - message_id = message.get("id") - if isinstance(message_id, str): - if message_id in seen: - continue - seen.add(message_id) - unique.append(message) - return unique - - def _message_to_dict(message: Any) -> dict[str, Any]: role = ( getattr(message, "type", None) or getattr(message, "role", None) or "assistant" @@ -1259,56 +1282,103 @@ def _final_response(messages: list[dict[str, Any]]) -> Any: return messages[-1].get("content") if messages else None -def _aggregate_usage(messages: list[dict[str, Any]]) -> dict[str, Any] | None: +def _aggregate_usage(calls: list[dict[str, Any]]) -> dict[str, Any] | None: + """Sum model-call usage; a count is known only when every call reported it. + + One call without a count makes the invocation's sum for that count a + partial value, so it stays unknown rather than under-reported. + """ + + if not calls: + return None totals: dict[str, Any] = {} - cost = 0.0 - has_cost = False - for message in messages: - usage = message.get("usage") if isinstance(message.get("usage"), dict) else {} - for source, target in ( - ("input_tokens", "prompt_tokens"), - ("output_tokens", "completion_tokens"), - ("total_tokens", "total_tokens"), - ): - value = usage.get(source) - if ( - isinstance(value, int) - and not isinstance(value, bool) - and 0 <= value <= (1 << 64) - 1 - ): - totals[target] = int(totals.get(target, 0)) + value - # Cost is not part of LangChain's UsageMetadata; surface it only when a - # model/provider reports it on the usage or response metadata. - candidate = ( - usage.get("total_cost") or usage.get("cost") or _metadata_cost(message) - ) - if ( - isinstance(candidate, (int, float)) - and not isinstance(candidate, bool) - and math.isfinite(candidate) - and candidate >= 0 - ): - cost += float(candidate) - has_cost = True - if has_cost: - totals["cost"] = cost + for source, target in ( + ("input_tokens", "prompt_tokens"), + ("output_tokens", "completion_tokens"), + ("total_tokens", "total_tokens"), + ): + values = [_reported_tokens(call, source) for call in calls] + if all(value is not None for value in values): + totals[target] = sum(values) + costs = [_reported_cost(call) for call in calls] + if all(cost is not None for cost in costs): + totals["cost"] = math.fsum(costs) return totals or None -def _metadata_cost(message: dict[str, Any]) -> float | None: - metadata = message.get("response_metadata") - if isinstance(metadata, dict): - value = metadata.get("cost") - if ( - isinstance(value, (int, float)) - and not isinstance(value, bool) - and math.isfinite(value) - and value >= 0 - ): - return float(value) +# Raw OpenAI Chat Completions usage keys behind each LangChain usage counter. +_OPENAI_USAGE_KEYS = { + "input_tokens": "prompt_tokens", + "output_tokens": "completion_tokens", + "total_tokens": "total_tokens", +} + + +def _reported_tokens(call: dict[str, Any], name: str) -> int | None: + """Return one call's token count, or ``None`` when the provider omitted it. + + ``ChatOpenAI`` turns an omitted Chat Completions count into 0 and derives a + missing or zero total from input plus output. When a response carries + OpenAI-shaped raw usage in ``response_metadata.token_usage``, the raw counts + are used: an omitted count is unknown, and a derived total counts only when + both of its parts were reported. Other responses use ``usage_metadata``. + """ + + usage = call.get("usage") + value = usage.get(name) if isinstance(usage, dict) else None + if not _is_token_count(value): + return None + metadata = call.get("response_metadata") + raw = metadata.get("token_usage") if isinstance(metadata, dict) else None + if not isinstance(raw, dict): + return value + reported = raw.get(_OPENAI_USAGE_KEYS[name]) + if _is_token_count(reported): + return reported + if name == "total_tokens" and reported is None: + parts = (raw.get("prompt_tokens"), raw.get("completion_tokens")) + if all(_is_token_count(part) for part in parts): + return value return None +def _is_token_count(value: Any) -> bool: + return ( + isinstance(value, int) + and not isinstance(value, bool) + and 0 <= value <= (1 << 64) - 1 + ) + + +def _reported_cost(call: dict[str, Any]) -> float | None: + """Return a provider-reported cost for one call, or ``None`` when absent. + + Cost is not part of LangChain's UsageMetadata; some providers report it on the + usage or response metadata. + """ + + usage = call.get("usage") if isinstance(call.get("usage"), dict) else {} + metadata = call.get("response_metadata") + candidates = ( + usage.get("total_cost"), + usage.get("cost"), + metadata.get("cost") if isinstance(metadata, dict) else None, + ) + for value in candidates: + if value is not None: + return float(value) if _is_cost(value) else None + return None + + +def _is_cost(value: Any) -> bool: + return ( + isinstance(value, (int, float)) + and not isinstance(value, bool) + and math.isfinite(value) + and value >= 0 + ) + + def _supported_kwargs(callable_obj: Any, kwargs: dict[str, Any]) -> dict[str, Any]: try: signature = inspect.signature(callable_obj) diff --git a/tests/adapters/test_deepagents.py b/tests/adapters/test_deepagents.py index 23304be4d..c916cb2a3 100644 --- a/tests/adapters/test_deepagents.py +++ b/tests/adapters/test_deepagents.py @@ -25,6 +25,9 @@ from unittest.mock import MagicMock import pytest +from langchain_core.callbacks import BaseCallbackHandler +from langchain_core.language_models.fake_chat_models import GenericFakeChatModel +from langchain_core.messages import AIMessage from langchain_openai import ChatOpenAI as InstalledChatOpenAI from langgraph.errors import GraphRecursionError from nemo_fabric_adapter_contract.codec import ContractValidationError @@ -34,6 +37,7 @@ from nemo_fabric_adapter_contract.models import AgentRunRequest from nemo_fabric_adapter_contract.models import AgentRunResult from nemo_fabric_adapter_contract.models import AgentRunStatus +from nemo_fabric_adapter_contract.models import AgentUsage from nemo_fabric_adapter_contract.models import McpOAuth2Config from nemo_fabric_adapter_contract.models import McpServiceAccountConfig from nemo_fabric_adapter_contract.models import RuntimeContext @@ -49,6 +53,23 @@ def _activation(*diagnostics: dict[str, object]) -> types.SimpleNamespace: ROOT = Path(__file__).resolve().parents[2] +async def call_model(config: dict[str, Any] | None, message: AIMessage) -> AIMessage: + """Return ``message`` from a real LangChain chat-model call under ``config``. + + The call fires the run's LangChain callbacks the way a Deep Agents model call + does, which is where the adapter reads usage. Test stand-ins for the Relay + handler are not LangChain handlers, so they are left out. + """ + + callbacks = [ + callback + for callback in (config or {}).get("callbacks") or [] + if isinstance(callback, BaseCallbackHandler) + ] + model = GenericFakeChatModel(messages=iter([message])) + return await model.ainvoke("prompt", config={"callbacks": callbacks}) + + def test_descriptor_declares_supported_normalized_config(): descriptor = json.loads( (ROOT / "adapters/python/deepagents/deepagents.fabric-adapter.json").read_text( @@ -192,11 +213,17 @@ async def stream( recorder["subgraphs"] = subgraphs recorder["checkpointer"] = kwargs.get("checkpointer") user = inputs["messages"][-1]["content"] - ai = { - "role": "ai", - "content": f"reply to {user}", - "usage": {"input_tokens": 5, "output_tokens": 7, "total_tokens": 12}, - } + ai = await call_model( + config, + AIMessage( + content=f"reply to {user}", + usage_metadata={ + "input_tokens": 5, + "output_tokens": 7, + "total_tokens": 12, + }, + ), + ) # subgraphs=True yields 3-tuples ``(namespace, mode, chunk)``; the main # graph has an empty namespace. ``updates`` carries the message produced # this turn; ``values`` is the full replayed state. @@ -1445,7 +1472,7 @@ def test_native_telemetry_requires_mapping(tmp_path, make_payload): ) -async def test_relay_disabled_adds_no_scope_or_callbacks( +async def test_relay_disabled_adds_no_scope_or_relay_callbacks( tmp_path, make_payload, fake_sdks ): # With telemetry disabled the invocation runs without a Relay scope, callback @@ -1454,7 +1481,9 @@ async def test_relay_disabled_adds_no_scope_or_callbacks( assert output["completed"] is True assert "telemetry" not in output - assert (fake_sdks["config"] or {}).get("callbacks") is None + # The only run callback is the adapter's own model-usage recorder. + callbacks = fake_sdks["config"]["callbacks"] + assert [type(callback) for callback in callbacks] == [adapter.ModelCallUsage] assert "relay-mw" not in (fake_sdks["create_kwargs"].get("middleware") or []) @@ -2070,14 +2099,19 @@ async def test_cost_is_extracted_from_response_metadata( ): import deepagents - message = { - "role": "ai", - "content": "done", - "usage": {"input_tokens": 1, "output_tokens": 1, "total_tokens": 2}, - "response_metadata": {"cost": 0.0025}, - } - async def astream(inputs, config=None, *, stream_mode=None, subgraphs=False): + message = await call_model( + config, + AIMessage( + content="done", + usage_metadata={ + "input_tokens": 1, + "output_tokens": 1, + "total_tokens": 2, + }, + response_metadata={"cost": 0.0025}, + ), + ) yield ((), "updates", {"agent": {"messages": [message]}}) yield ((), "values", {"messages": [message]}) @@ -2108,6 +2142,23 @@ def test_aggregate_usage_discards_tokens_larger_than_uint64(): assert adapter._aggregate_usage([{"usage": {"input_tokens": 1 << 64}}]) is None +@pytest.mark.parametrize( + ("costs", "expected"), + [ + pytest.param([0.0, 0.0], 0.0, id="explicit-zero"), + pytest.param([0.5, 0.25], 0.75, id="every-call"), + pytest.param([0.5, None], None, id="one-call-without-cost"), + ], +) +def test_aggregate_usage_reports_cost_only_when_every_call_did(costs, expected): + calls = [ + {"usage": {"input_tokens": 1}, "response_metadata": {"cost": cost}} + for cost in costs + ] + + assert adapter._aggregate_usage(calls).get("cost") == expected + + async def test_replayed_state_usage_counts_current_turn_only( tmp_path, make_payload, monkeypatch ): @@ -2115,20 +2166,25 @@ async def test_replayed_state_usage_counts_current_turn_only( # cost must reflect only the message emitted this turn, not the replayed one. import deepagents - prior = { - "role": "ai", - "content": "prior", - "usage": {"input_tokens": 1, "output_tokens": 1, "total_tokens": 2}, - "response_metadata": {"cost": 0.001}, - } - current = { - "role": "ai", - "content": "now", - "usage": {"input_tokens": 2, "output_tokens": 2, "total_tokens": 4}, - "response_metadata": {"cost": 0.002}, - } + prior = AIMessage( + content="prior", + usage_metadata={"input_tokens": 1, "output_tokens": 1, "total_tokens": 2}, + response_metadata={"cost": 0.001}, + ) async def astream(inputs, config=None, *, stream_mode=None, subgraphs=False): + current = await call_model( + config, + AIMessage( + content="now", + usage_metadata={ + "input_tokens": 2, + "output_tokens": 2, + "total_tokens": 4, + }, + response_metadata={"cost": 0.002}, + ), + ) # Only the current turn's message is emitted as an update... yield ((), "updates", {"agent": {"messages": [current]}}) # ...but the final state also replays the prior turn. @@ -2427,47 +2483,360 @@ def test_deepagents_passthrough_rejects_fabric_owned_option(): adapter._validated_deepagents_settings({"model": "untrusted-model"}) -async def test_subagent_usage_folded_from_subgraph(tmp_path, make_payload, monkeypatch): - # Usage from a delegated subagent is emitted under a subgraph namespace; folding - # it into this turn keeps usage/cost accurate. Duplicate ids are counted once. - import deepagents - - sub_ai = { - "role": "ai", - "content": "subagent work", - "id": "sub-1", - "usage": {"input_tokens": 10, "output_tokens": 20, "total_tokens": 30}, - } - main_ai = { - "role": "ai", - "content": "final", - "id": "main-1", - "usage": {"input_tokens": 2, "output_tokens": 3, "total_tokens": 5}, +def chat_completion( + usage: dict[str, Any] | None, + *, + tool_calls: list[dict[str, Any]] | None = None, +) -> dict[str, Any]: + """Return an OpenAI Chat Completions response body with the given usage.""" + + message: dict[str, Any] = {"role": "assistant", "content": "done"} + if tool_calls: + message["content"] = None + message["tool_calls"] = tool_calls + body: dict[str, Any] = { + "id": f"chatcmpl-{uuid.uuid4().hex}", + "object": "chat.completion", + "created": 0, + "model": "nvidia/nemotron-3-nano-omni-30b-a3b-reasoning", + "choices": [ + { + "index": 0, + "message": message, + "finish_reason": "tool_calls" if tool_calls else "stop", + } + ], } + if usage is not None: + body["usage"] = usage + return body - async def astream(inputs, config=None, *, stream_mode=None, subgraphs=False): - assert subgraphs is True - # subagent step under a subgraph namespace, emitted twice (dedup by id) - yield (("task:researcher",), "updates", {"agent": {"messages": [sub_ai]}}) - yield (("task:researcher",), "updates", {"agent": {"messages": [sub_ai]}}) - # main graph step + replayed final state - yield ((), "updates", {"agent": {"messages": [main_ai]}}) - yield ((), "values", {"messages": [main_ai]}) - agent = MagicMock() - agent.astream = astream - monkeypatch.setattr(deepagents, "create_deep_agent", MagicMock(return_value=agent)) +LS_TOOL_CALL = { + "id": "call-1", + "type": "function", + "function": {"name": "ls", "arguments": '{"path": "/"}'}, +} - output = await invoke_once(make_payload(tmp_path)) - # subagent (10/20/30) + main (2/3/5), the duplicate subagent message counted once - assert output["usage"] == { - "prompt_tokens": 12, - "completion_tokens": 23, - "total_tokens": 35, +@pytest.fixture(name="openai_endpoint") +def openai_endpoint_fixture(fake_sdks, monkeypatch) -> dict[str, Any]: + """Run the real Deep Agents graph and ChatOpenAI against a scripted endpoint. + + Tests append response bodies, ``httpx.Response`` objects, or exceptions to + raise to ``replies``; each model request takes the next one, and its JSON + body is kept in ``requests``. ``profile`` sets the chat model profile, which Deep Agents uses + to choose its summarization trigger. + """ + + import httpx + + for name in ( + "deepagents", + "deepagents.backends", + "deepagents.middleware", + "deepagents.middleware.subagents", + "langchain.chat_models", + "langgraph", + "langgraph.checkpoint", + "langgraph.checkpoint.sqlite", + "langgraph.checkpoint.sqlite.aio", + ): + monkeypatch.delitem(sys.modules, name, raising=False) + endpoint: dict[str, Any] = {"replies": [], "requests": [], "profile": None} + + def respond(request: httpx.Request) -> httpx.Response: + endpoint["requests"].append(json.loads(request.content)) + reply = endpoint["replies"].pop(0) + if isinstance(reply, Exception): + raise reply + if isinstance(reply, httpx.Response): + return reply + return httpx.Response(200, json=reply) + + def build_chat_openai(**kwargs: Any) -> InstalledChatOpenAI: + return InstalledChatOpenAI( + **kwargs, + max_retries=0, + profile=endpoint["profile"], + http_async_client=httpx.AsyncClient(transport=httpx.MockTransport(respond)), + ) + + fake_sdks["chat_openai"].side_effect = build_chat_openai + return endpoint + + +async def run_invocations(payload: dict[str, Any], count: int) -> list[AgentRunResult]: + runtime = adapter.DeepAgentsRuntime() + await runtime.start(lifecycle_start_payload(payload)) + try: + return [ + await runtime.invoke(*lifecycle_invocation(payload)) for _ in range(count) + ] + finally: + await runtime.stop() + + +async def test_summarization_model_call_usage_is_counted( + tmp_path, make_payload, openai_endpoint +): + # Deep Agents summarizes older history with a separate model call whose response + # never enters graph state. Its usage belongs to the invocation that made it. + openai_endpoint["profile"] = {"max_input_tokens": 1000} + openai_endpoint["replies"] += [ + chat_completion( + {"prompt_tokens": 1500, "completion_tokens": 10, "total_tokens": 1510} + ), + # Second invocation: the summary call, then the agent's model call. + chat_completion( + {"prompt_tokens": 1400, "completion_tokens": 60, "total_tokens": 1460} + ), + chat_completion( + {"prompt_tokens": 300, "completion_tokens": 5, "total_tokens": 305} + ), + ] + payload = make_payload(tmp_path) + payload["request"]["input"] = "word " * 6000 + + first, second = await run_invocations(payload, 2) + + summary_request, agent_request = openai_endpoint["requests"][1:] + assert "tools" not in summary_request + assert "tools" in agent_request + assert first.usage == AgentUsage( + input_tokens=1500, output_tokens=10, total_tokens=1510 + ) + assert second.status is AgentRunStatus.SUCCEEDED + assert second.usage == AgentUsage( + input_tokens=1700, output_tokens=65, total_tokens=1765 + ) + + +async def test_subagent_model_calls_are_counted( + tmp_path, make_payload, openai_endpoint +): + # Two parallel ``task`` calls run the general-purpose subagent twice; each + # subagent's model call belongs to this invocation. + task_calls = [ + { + "id": f"task-{index}", + "type": "function", + "function": { + "name": "task", + "arguments": json.dumps( + { + "description": f"Echo the word {index}.", + "subagent_type": "general-purpose", + } + ), + }, + } + for index in (1, 2) + ] + openai_endpoint["replies"] += [ + chat_completion( + {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}, + tool_calls=task_calls, + ), + chat_completion( + {"prompt_tokens": 10, "completion_tokens": 10, "total_tokens": 20} + ), + chat_completion( + {"prompt_tokens": 100, "completion_tokens": 100, "total_tokens": 200} + ), + chat_completion( + {"prompt_tokens": 1000, "completion_tokens": 1000, "total_tokens": 2000} + ), + ] + + (result,) = await run_invocations(make_payload(tmp_path), 1) + + assert result.status is AgentRunStatus.SUCCEEDED + assert len(openai_endpoint["requests"]) == 4 + subgraphs = { + event["subgraph"] for event in result.output["events"] if "subgraph" in event } - # the subgraph step is recorded with its namespace label - assert any(evt.get("subgraph") == "task:researcher" for evt in output["events"]) + assert len(subgraphs) == 2 + assert result.usage == AgentUsage( + input_tokens=1111, output_tokens=1111, total_tokens=2222 + ) + + +async def test_recursion_limit_failure_keeps_usage_of_completed_model_calls( + tmp_path, make_payload, openai_endpoint +): + tool_call_reply = chat_completion( + {"prompt_tokens": 10, "completion_tokens": 1, "total_tokens": 11}, + tool_calls=[LS_TOOL_CALL], + ) + openai_endpoint["replies"] += [tool_call_reply] * 20 + payload = make_payload(tmp_path) + payload["config"]["runtime"] = {"max_turns": 6} + + (result,) = await run_invocations(payload, 1) + + calls = len(openai_endpoint["requests"]) + assert calls > 1 + assert result.status is AgentRunStatus.FAILED + assert result.error.code == "deepagents_recursion_limit_reached" + assert result.usage == AgentUsage( + input_tokens=10 * calls, output_tokens=calls, total_tokens=11 * calls + ) + + +def provider_failure(kind: str) -> Any: + """Return a scripted endpoint reply that fails a model call.""" + + import httpx + + if kind == "timeout": + return httpx.ReadTimeout("provider did not answer in time") + return httpx.Response(int(kind), json={"error": {"message": "request failed"}}) + + +@pytest.mark.parametrize( + ("failure", "expected"), + [ + # A 4xx rejection other than 408 spends nothing, so the finished call's + # usage stays known. + pytest.param( + "400", + AgentUsage(input_tokens=120, output_tokens=8, total_tokens=128), + id="400-rejection", + ), + # The provider may have done work without reporting usage. + pytest.param("408", None, id="408-request-timeout"), + pytest.param("503", None, id="503-server-error"), + pytest.param("timeout", None, id="read-timeout"), + ], +) +async def test_failed_model_call_usage_follows_provider_response( + tmp_path, make_payload, openai_endpoint, failure, expected +): + openai_endpoint["replies"] += [ + chat_completion( + {"prompt_tokens": 120, "completion_tokens": 8, "total_tokens": 128}, + tool_calls=[LS_TOOL_CALL], + ), + provider_failure(failure), + ] + + (result,) = await run_invocations(make_payload(tmp_path), 1) + + assert result.status is AgentRunStatus.FAILED + assert result.usage == expected + + +@pytest.mark.parametrize( + ("failure", "expected"), + [ + pytest.param( + "429", + {"prompt_tokens": 30, "completion_tokens": 2, "total_tokens": 32}, + id="429-then-success", + ), + pytest.param("503", None, id="503-then-success"), + ], +) +async def test_retried_model_call_usage_follows_failed_attempt( + fake_sdks, openai_endpoint, failure, expected +): + # LangChain's with_retry runs each attempt as its own chat-model call, as the + # Deep Agents summarization model does. + openai_endpoint["replies"] += [ + provider_failure(failure), + chat_completion( + {"prompt_tokens": 30, "completion_tokens": 2, "total_tokens": 32} + ), + ] + model = fake_sdks["chat_openai"]( + model="nvidia/nemotron-3-nano-omni-30b-a3b-reasoning", + api_key="test123", + base_url="https://integrate.api.nvidia.com/v1", + ) + model_usage = adapter.ModelCallUsage() + + await model.with_retry(stop_after_attempt=2, wait_exponential_jitter=False).ainvoke( + "summarize", config={"callbacks": [model_usage]} + ) + + assert adapter._aggregate_usage(model_usage.calls()) == expected + + +async def test_recorder_keeps_a_cancelled_call_unreported(): + # Recorder unit test: a cancelled call ends without on_llm_end and, when + # LangChain gathers it, without on_llm_error; its start record keeps the + # invocation's counts unknown. + model_usage = adapter.ModelCallUsage() + await call_model( + {"callbacks": [model_usage]}, + AIMessage( + content="done", + usage_metadata={"input_tokens": 1, "output_tokens": 1, "total_tokens": 2}, + ), + ) + model_usage.on_chat_model_start({}, [], run_id=uuid.uuid4()) + + assert adapter._aggregate_usage(model_usage.calls()) is None + + +@pytest.mark.parametrize( + ("usages", "expected"), + [ + pytest.param( + [{"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0}], + AgentUsage(input_tokens=0, output_tokens=0, total_tokens=0), + id="explicit-zero", + ), + pytest.param( + [{"prompt_tokens": 100}], + AgentUsage(input_tokens=100), + id="input-only", + ), + pytest.param( + [{"completion_tokens": 5}], + AgentUsage(output_tokens=5), + id="output-only", + ), + pytest.param( + [{"prompt_tokens": 100, "completion_tokens": 5}], + AgentUsage(input_tokens=100, output_tokens=5, total_tokens=105), + id="total-omitted", + ), + pytest.param( + [{"completion_tokens": 5, "total_tokens": 105}], + AgentUsage(output_tokens=5, total_tokens=105), + id="provider-total-with-input-omitted", + ), + pytest.param( + [{"prompt_tokens": 7, "completion_tokens": 5, "total_tokens": 0}], + AgentUsage(input_tokens=7, output_tokens=5, total_tokens=0), + id="provider-total-zero-kept", + ), + pytest.param( + [ + {"prompt_tokens": 100, "completion_tokens": 5, "total_tokens": 105}, + None, + ], + None, + id="one-call-without-usage", + ), + ], +) +async def test_partial_usage_reports_keep_missing_counts_unknown( + tmp_path, make_payload, openai_endpoint, usages, expected +): + # ChatOpenAI turns an omitted count into 0; a count is known only when every + # model call in the invocation reported it. + openai_endpoint["replies"] += [ + *(chat_completion(usage, tool_calls=[LS_TOOL_CALL]) for usage in usages[:-1]), + chat_completion(usages[-1]), + ] + + (result,) = await run_invocations(make_payload(tmp_path), 1) + + assert result.status is AgentRunStatus.SUCCEEDED + assert result.usage == expected async def test_bad_mcp_transport_fails_runtime_start(tmp_path, make_payload): diff --git a/tests/e2e/test_deepagents.py b/tests/e2e/test_deepagents.py index 30b07e9ef..08155598c 100644 --- a/tests/e2e/test_deepagents.py +++ b/tests/e2e/test_deepagents.py @@ -486,7 +486,6 @@ async def test_deepagents_builtin_subagent_delegation(): call["args"]["subagent_type"] == "general-purpose" for call in task_calls ), result.to_mapping() assert any("subgraph" in event for event in output["events"]), result.to_mapping() - # delegated steps are folded into this turn's usage aggregation assert output["resumed"] is False, result.to_mapping() diff --git a/tests/integrations/test_relay_scope_leak.py b/tests/integrations/test_relay_scope_leak.py index 9dabfdcd5..5d5dbc58a 100644 --- a/tests/integrations/test_relay_scope_leak.py +++ b/tests/integrations/test_relay_scope_leak.py @@ -168,7 +168,7 @@ async def test_uuid_request_id_seeds_real_relay_parent( monkeypatch, make_runtime_context ): async def fake_invoke(agent, user_message, thread_id, callbacks=None): - return {"messages": []}, [], [] + return {"messages": []}, [] monkeypatch.setattr(adapter, "invoke_compiled_agent", fake_invoke) @@ -210,7 +210,7 @@ async def test_overlapping_turns_stay_telemetry_clean( async def fake_invoke(agent, user_message, thread_id, callbacks=None): if callbacks: await _overlapping_chain_runs(callbacks[0]) - return {"messages": []}, [], [] + return {"messages": []}, [] monkeypatch.setattr(adapter, "invoke_compiled_agent", fake_invoke) recording_scope = _RecordingScope() From cdc0636a647d562a02a72110f9921b8943d3bd04 Mon Sep 17 00:00:00 2001 From: mnajafian-nv Date: Mon, 19 Oct 2026 17:00:16 -0700 Subject: [PATCH 02/26] docs: name NeMo Fabric in the Deep Agents usage limits Signed-off-by: mnajafian-nv --- adapters/python/deepagents/README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/adapters/python/deepagents/README.md b/adapters/python/deepagents/README.md index 78e94ce51..66a44eece 100644 --- a/adapters/python/deepagents/README.md +++ b/adapters/python/deepagents/README.md @@ -243,7 +243,7 @@ unknown, because the provider may have done work without reporting usage. OpenAI SDK retries happen inside one LangChain call, so a 5xx the SDK retries to success is not visible to the adapter, and only the final response's reported usage counts. Usage replayed from a configured LangChain model cache -would be counted as reported; Fabric configures no such cache. +would be counted as reported; NeMo Fabric configures no such cache. When a response carries OpenAI-shaped usage in `response_metadata.token_usage`, as `ChatOpenAI` responses without streaming From fd184bd2f5b6a56ec0ad44583ec23d7fef44661c Mon Sep 17 00:00:00 2001 From: mnajafian-nv Date: Mon, 19 Oct 2026 22:23:56 -0700 Subject: [PATCH 03/26] fix: report Droid cache tokens in the standard usage fields The Droid adapter put cache reads and writes only in adapter-private usage extension keys, so standard usage consumers saw no cache counts. Droid also stores a count that the provider omits as 0, and the adapter copied those zeros, so a provider that sent no output count produced 0 output tokens. Cache reads now go in cached_input_tokens, and cache writes go in the nemo_fabric.detailed_token_usage.v1 extension as cache_write_input_tokens. The descriptor and the adapter catalog declare that extension, and the private cache keys are removed because nothing reads them. Droid's input count does not follow one cache convention. Its Gemini handler keeps Google's prompt count, which includes cache reads. Its Responses and Chat Completions handlers subtract cache reads. Its Anthropic Messages and Bedrock Converse handlers keep provider input that excludes cache reads and writes. Droid picks the handler from the model's provider. The adapter reads the session's current model from the SDK, picks the convention from the model ID or a custom model's provider in Droid's settings, and reports inclusive input with input_tokens_include_cache set to true and a total. Some turns cannot be normalized with one convention. Droid adds Task subagent usage to the turn's counters without saying which model each subagent used, and a model can change during a turn. The adapter watches the session's notifications for the whole invocation: a turn whose completion carries child usage, or that touches models with different conventions, keeps Droid's counters and leaves the flag and, when the turn has cache counts, the total unset. So do the auto routers, models without a known convention, and Responses turns with cache writes. A 0 count is unknown, and so is a total built from it. The exception is a 0 cache read count for Anthropic, OpenAI Responses, Bedrock Converse and Gemini models, whose usage reports carry the count or, for Gemini, omit it only when it is 0. Zero Factory credits are omitted, because Droid also uses 0 for missing credits. The thinking-token count and Factory credits stay as adapter extensions, because Droid estimates thinking tokens from the thinking text for Anthropic, Google and Bedrock models, and credits have no standard field. Signed-off-by: mnajafian-nv --- .../droid/droid.fabric-adapter.json | 14 +- adapters/typescript/droid/src/droid-sdk.ts | 337 +++++++++++++-- .../typescript/droid/test/droid-sdk.test.mjs | 387 +++++++++++++++++- docs/integrations/harness/droid.mdx | 67 ++- .../nemo_fabric_adapter_catalog/catalog.json | 38 +- 5 files changed, 784 insertions(+), 59 deletions(-) diff --git a/adapters/typescript/droid/droid.fabric-adapter.json b/adapters/typescript/droid/droid.fabric-adapter.json index 0c11ce867..c95bd9f14 100644 --- a/adapters/typescript/droid/droid.fabric-adapter.json +++ b/adapters/typescript/droid/droid.fabric-adapter.json @@ -55,8 +55,18 @@ "$schema": "https://json-schema.org/draft/2020-12/schema", "type": "object", "properties": { - "cache_read_tokens": {"type": "integer", "minimum": 0}, - "cache_write_tokens": {"type": "integer", "minimum": 0}, + "nemo_fabric.detailed_token_usage.v1": { + "type": "object", + "properties": { + "input_tokens": {"type": "integer", "minimum": 0}, + "cached_input_tokens": {"type": "integer", "minimum": 0}, + "cache_write_input_tokens": {"type": "integer", "minimum": 0}, + "input_tokens_include_cache": {"type": "boolean"}, + "output_tokens": {"type": "integer", "minimum": 0}, + "output_tokens_include_reasoning": {"type": "boolean"} + }, + "additionalProperties": false + }, "thinking_tokens": {"type": "integer", "minimum": 0}, "factory_credits": {"type": "number", "minimum": 0} }, diff --git a/adapters/typescript/droid/src/droid-sdk.ts b/adapters/typescript/droid/src/droid-sdk.ts index 4bc657e02..7044313f4 100644 --- a/adapters/typescript/droid/src/droid-sdk.ts +++ b/adapters/typescript/droid/src/droid-sdk.ts @@ -17,6 +17,7 @@ import { LifecycleError, type AdapterStartInput } from "nemo-fabric-adapters-com import type { DroidPromptOutcome, DroidSessionFactory, DroidSessionHandle } from "./runtime.js"; const DROID_SDK_INSTALL_COMMAND = "npm install @factory/droid-sdk@0.9.1"; +const DETAILED_TOKEN_USAGE_KEY = "nemo_fabric.detailed_token_usage.v1"; const MCP_STARTUP_TIMEOUT_MS = 10_000; const MCP_STATUS_POLL_INTERVAL_MS = 50; const INHERITED_ENVIRONMENT_NAMES = new Set([ @@ -115,9 +116,14 @@ interface DroidMcpNetworkConfig { type DroidMcpServerConfig = DroidMcpStdioConfig | DroidMcpNetworkConfig; +type DroidNotificationCallback = (notification: Record) => void; + interface DroidSdkSession { readonly id: string; + /** Kept current by the SDK from settings_updated notifications. */ + readonly settings?: { modelId?: unknown }; stream(prompt: string): AsyncGenerator; + onNotification?(callback: DroidNotificationCallback, filter?: { type?: string }): () => void; close(): Promise; listTools(): Promise; updateSettings(settings: { disabledToolIds: string[] }): Promise; @@ -329,6 +335,7 @@ interface RuntimeProfile { home?: string; entrypoints: string[]; settingsEnvironment: Record; + settingsFiles: DroidSettings[]; } function settingsHome(input: AdapterStartInput, parentEnvironment: NodeJS.ProcessEnv): string | undefined { @@ -343,29 +350,20 @@ function settingsHome(input: AdapterStartInput, parentEnvironment: NodeJS.Proces ); } -async function customModelSettingsEnvironment( - input: AdapterStartInput, - model: string, - sourceHome: string | undefined, - parentEnvironment: NodeJS.ProcessEnv, -): Promise> { - if (!model.startsWith("custom:") || sourceHome === undefined) return {}; - const modelId = model.slice("custom:".length); - let apiKey: unknown; +type CustomModelSettings = Record; + +type DroidSettings = Record; + +/** Reads the Droid settings files in the order Droid layers them. */ +async function readDroidSettings(sourceHome: string | undefined): Promise { + if (sourceHome === undefined) return []; + const files: DroidSettings[] = []; for (const filename of ["settings.json", "settings.local.json"]) { try { const settings: unknown = JSON.parse(await readFile(join(sourceHome, ".factory", filename), "utf8")); - if (typeof settings !== "object" || settings === null || !("customModels" in settings)) continue; - const customModels = settings.customModels; - if (!Array.isArray(customModels)) continue; - const configured = customModels.find( - (candidate) => - typeof candidate === "object" && - candidate !== null && - "model" in candidate && - candidate.model === modelId, - ); - if (configured !== undefined && "apiKey" in configured) apiKey = configured.apiKey; + if (typeof settings === "object" && settings !== null && !Array.isArray(settings)) { + files.push(settings as DroidSettings); + } } catch (error) { if (typeof error === "object" && error !== null && "code" in error && error.code === "ENOENT") continue; if (error instanceof SyntaxError) continue; // Droid reports malformed settings during session creation. @@ -375,6 +373,38 @@ async function customModelSettingsEnvironment( ); } } + return files; +} + +/** Returns the first custom model entry for `model` in each settings file. */ +function customModelSettings(settingsFiles: DroidSettings[], model: string): CustomModelSettings[] { + if (!model.startsWith("custom:")) return []; + const modelId = model.slice("custom:".length); + const entries: CustomModelSettings[] = []; + for (const settings of settingsFiles) { + const customModels = settings.customModels; + if (!Array.isArray(customModels)) continue; + const configured: unknown = customModels.find( + (candidate) => + typeof candidate === "object" && + candidate !== null && + "model" in candidate && + candidate.model === modelId, + ); + if (typeof configured === "object" && configured !== null) entries.push(configured as CustomModelSettings); + } + return entries; +} + +function customModelSettingsEnvironment( + input: AdapterStartInput, + customModels: CustomModelSettings[], + parentEnvironment: NodeJS.ProcessEnv, +): Record { + let apiKey: unknown; + for (const configured of customModels) { + if ("apiKey" in configured) apiKey = configured.apiKey; + } if (typeof apiKey !== "string") return {}; const match = /^\$\{([A-Za-z_][A-Za-z0-9_]*)\}$/.exec(apiKey); if (match === null) return {}; @@ -387,6 +417,106 @@ async function customModelSettingsEnvironment( return { [name]: value }; } +/** + * How Droid 0.233.0 fills a turn's inputTokens on the provider path it uses for + * the configured model. Droid routes by model provider: Google models use the + * Gemini handler, Anthropic and Factory-hosted models use the Anthropic + * Messages or Chat Completions handler, OpenAI and xAI models use the Responses + * handler, and Bedrock Converse custom models use the Converse handler. + */ +type DroidInputConvention = + /** Gemini: Droid keeps promptTokenCount, which includes cached content. */ + | "includes_cache_reads" + /** Responses and Chat Completions: Droid subtracts cache reads from the provider input. */ + | "excludes_cache_reads" + /** Anthropic Messages and Bedrock Converse: the provider input excludes cache reads and writes. */ + | "excludes_cache"; + +interface DroidUsageConvention { + input: DroidInputConvention; + /** + * The provider sends a cache read count, or omits it only when it is 0, on + * every response. Droid's 0 then means no cache reads rather than no report. + */ + reportsCacheReads: boolean; +} + +// Anthropic and Bedrock define total input as input plus cache reads and +// writes, and the OpenAI Responses usage object requires cached_tokens. Gemini +// JSON omits a zero cachedContentTokenCount. Factory-hosted, xAI, and custom +// endpoints are not shown to report cache reads on every response. +const GEMINI: DroidUsageConvention = { input: "includes_cache_reads", reportsCacheReads: true }; +const ANTHROPIC_OR_BEDROCK: DroidUsageConvention = { input: "excludes_cache", reportsCacheReads: true }; +const OPENAI_RESPONSES: DroidUsageConvention = { input: "excludes_cache_reads", reportsCacheReads: true }; +const EXCLUDES_CACHE: DroidUsageConvention = { input: "excludes_cache", reportsCacheReads: false }; +const EXCLUDES_CACHE_READS: DroidUsageConvention = { input: "excludes_cache_reads", reportsCacheReads: false }; + +const BUILT_IN_MODEL_CONVENTIONS: ReadonlyArray = [ + ["claude-", ANTHROPIC_OR_BEDROCK], + ["gemini-", GEMINI], + ["gpt-", OPENAI_RESPONSES], + ["grok-", EXCLUDES_CACHE_READS], + // Factory-hosted models use the Anthropic Messages or Chat Completions handler; + // the Chat Completions handler never records cache writes. + ["deepseek-", EXCLUDES_CACHE], + ["glm-", EXCLUDES_CACHE], + ["kimi-", EXCLUDES_CACHE], + ["minimax-", EXCLUDES_CACHE], + ["mistral-", EXCLUDES_CACHE], + ["nemotron-", EXCLUDES_CACHE], + ["qwen", EXCLUDES_CACHE], +]; + +const GOOGLE_HINT = /gemini|garnet|generativelanguage\.googleapis\.com|aiplatform\.googleapis\.com/i; + +function customModelConvention(settings: CustomModelSettings): DroidUsageConvention | undefined { + switch (settings.provider) { + case "bedrock-converse": + return ANTHROPIC_OR_BEDROCK; + case "anthropic": + return EXCLUDES_CACHE; + case "openai": + return EXCLUDES_CACHE_READS; + case "generic-chat-completion-api": { + // Droid sends a chat-completions model to its Gemini handler when the + // model, base model, or base URL resolves to Google. + const hints = [settings.model, settings.baseModelId, settings.baseUrl].filter((value) => typeof value === "string"); + return hints.some((value) => GOOGLE_HINT.test(value)) ? undefined : EXCLUDES_CACHE_READS; + } + default: + return undefined; + } +} + +function customConvention(customModels: CustomModelSettings[]): DroidUsageConvention | undefined { + const conventions = new Set(customModels.map(customModelConvention)); + return conventions.size === 1 ? [...conventions][0] : undefined; +} + +// Droid's router models pick a model per request, so their convention is unknown. +const ROUTER_MODELS = new Set(["auto", "auto-fast"]); + +/** Selects the usage convention for whichever model the Droid session uses. */ +class DroidUsageConventions { + private readonly settingsFiles: DroidSettings[]; + private readonly resolved = new Map(); + + constructor(settingsFiles: DroidSettings[]) { + this.settingsFiles = settingsFiles; + } + + forModel(model: string): DroidUsageConvention | undefined { + if (!this.resolved.has(model)) this.resolved.set(model, this.resolve(model)); + return this.resolved.get(model); + } + + private resolve(model: string): DroidUsageConvention | undefined { + if (ROUTER_MODELS.has(model)) return undefined; + if (model.startsWith("custom:")) return customConvention(customModelSettings(this.settingsFiles, model)); + return BUILT_IN_MODEL_CONVENTIONS.find(([prefix]) => model.startsWith(prefix))?.[1]; + } +} + async function prepareRuntimeProfile( input: AdapterStartInput, hasMcpServers: boolean, @@ -398,8 +528,15 @@ async function prepareRuntimeProfile( } const configured = input.config.skills?.paths ?? []; const sourceHome = settingsHome(input, parentEnvironment); - const settingsEnvironment = await customModelSettingsEnvironment(input, model, sourceHome, parentEnvironment); - if (configured.length === 0 && !hasMcpServers) return { entrypoints: [], settingsEnvironment }; + const settingsFiles = await readDroidSettings(sourceHome); + const settingsEnvironment = customModelSettingsEnvironment( + input, + customModelSettings(settingsFiles, model), + parentEnvironment, + ); + if (configured.length === 0 && !hasMcpServers) { + return { entrypoints: [], settingsEnvironment, settingsFiles }; + } const root = await mkdtemp(join(tmpdir(), "nemo-fabric-droid-")); const home = join(root, "home"); @@ -452,7 +589,7 @@ async function prepareRuntimeProfile( await cp(source, destination, { recursive: true, errorOnExist: true }); entrypoints.push(await realpath(join(destination, "SKILL.md"))); } - return { root, home, entrypoints, settingsEnvironment }; + return { root, home, entrypoints, settingsEnvironment, settingsFiles }; } catch (error) { await rm(root, { recursive: true, force: true }); throw error; @@ -491,7 +628,34 @@ function validTokenCount(value: unknown): value is number { return Number.isSafeInteger(value) && Number(value) >= 0; } -function normalizeUsage(value: unknown): AgentUsage | undefined { +/** + * Returns the turn's input including cache reads and writes, or undefined when + * the provider path or Droid's input count does not establish it. Droid fills + * a missing input count with 0, so 0 input stays unknown. + */ +function inputIncludingCache( + convention: DroidInputConvention | undefined, + inputTokens: number, + cacheReadTokens: number, + cacheWriteTokens: number, +): number | undefined { + if (inputTokens === 0) return undefined; + switch (convention) { + case "excludes_cache": + return inputTokens + cacheReadTokens + cacheWriteTokens; + // The Gemini and Chat Completions handlers never record cache writes. The + // Responses handler records a nonstandard write count without subtracting + // it, so whether input already includes those writes is not established. + case "excludes_cache_reads": + return cacheWriteTokens === 0 ? inputTokens + cacheReadTokens : undefined; + case "includes_cache_reads": + return cacheWriteTokens === 0 ? inputTokens : undefined; + default: + return undefined; + } +} + +function normalizeUsage(value: unknown, convention: DroidUsageConvention | undefined): AgentUsage | undefined { if (value === null || value === undefined) return undefined; if (typeof value !== "object") { throw new LifecycleError("droid_malformed_result", "Droid returned malformed usage metadata"); @@ -507,24 +671,47 @@ function normalizeUsage(value: unknown): AgentUsage | undefined { if (!counts.every(validTokenCount) || (usage.factoryCredits !== undefined && !validNonNegativeNumber(usage.factoryCredits))) { throw new LifecycleError("droid_malformed_result", "Droid returned malformed usage metadata"); } - const total = usage.inputTokens! + usage.outputTokens!; - if (!Number.isSafeInteger(total)) { + const { inputTokens, outputTokens, cacheCreationTokens, cacheReadTokens, thinkingTokens, factoryCredits } = + usage as DroidTokenUsage; + // Droid stores a count the provider omitted as 0, so a 0 count is unknown. + // Output includes reasoning on every provider path. + const output = outputTokens > 0 ? outputTokens : undefined; + const inclusiveInput = inputIncludingCache(convention?.input, inputTokens, cacheReadTokens, cacheCreationTokens); + const cacheReadsKnown = cacheReadTokens > 0 || (inclusiveInput !== undefined && convention?.reportsCacheReads === true); + // Without an inclusive count, report Droid's own input count and leave its + // cache convention unset. A turn with cache counts then has no total. + const hasCache = cacheReadTokens > 0 || cacheCreationTokens > 0; + const input = inclusiveInput ?? (inputTokens > 0 ? inputTokens : undefined); + const totalKnown = inclusiveInput !== undefined || !hasCache; + const total = !totalKnown || input === undefined || output === undefined ? undefined : input + output; + if ((input !== undefined && !Number.isSafeInteger(input)) || (total !== undefined && !Number.isSafeInteger(total))) { throw new LifecycleError("droid_malformed_result", "Droid returned malformed usage metadata"); } + const inputUsage: JsonObject = { + ...(input === undefined ? {} : { input_tokens: input }), + ...(cacheReadsKnown ? { cached_input_tokens: cacheReadTokens } : {}), + ...(inclusiveInput === undefined ? {} : { input_tokens_include_cache: true }), + }; + const detailed: JsonObject = { + ...inputUsage, + ...(cacheCreationTokens > 0 ? { cache_write_input_tokens: cacheCreationTokens } : {}), + ...(output === undefined ? {} : { output_tokens: output, output_tokens_include_reasoning: true }), + }; + const extensions: JsonObject = { + ...(Object.keys(detailed).length === 0 ? {} : { [DETAILED_TOKEN_USAGE_KEY]: detailed }), + ...(thinkingTokens > 0 ? { thinking_tokens: thinkingTokens } : {}), + ...(factoryCredits === undefined || factoryCredits === 0 ? {} : { factory_credits: factoryCredits }), + }; + if (Object.keys(extensions).length === 0) return undefined; return { - input_tokens: usage.inputTokens!, - output_tokens: usage.outputTokens!, - total_tokens: total, - extensions: { - cache_read_tokens: usage.cacheReadTokens!, - cache_write_tokens: usage.cacheCreationTokens!, - thinking_tokens: usage.thinkingTokens!, - ...(usage.factoryCredits === undefined ? {} : { factory_credits: usage.factoryCredits }), - }, + ...inputUsage, + ...(output === undefined ? {} : { output_tokens: output }), + ...(total === undefined ? {} : { total_tokens: total }), + extensions, }; } -function normalizeResult(value: unknown): DroidPromptOutcome { +function normalizeResult(value: unknown, convention: DroidUsageConvention | undefined): DroidPromptOutcome { if (typeof value !== "object" || value === null) { throw new LifecycleError("droid_malformed_result", "Droid did not return a terminal result"); } @@ -542,7 +729,7 @@ function normalizeResult(value: unknown): DroidPromptOutcome { ) { throw new LifecycleError("droid_malformed_result", "Droid returned a malformed terminal result"); } - const usage = normalizeUsage(result.tokenUsage); + const usage = normalizeUsage(result.tokenUsage, convention); const extensions: JsonObject = { session_id: result.sessionId, finish_reason: String(result.subtype), @@ -677,20 +864,77 @@ async function validateMcpTools( } } +function hasTokenCounts(value: unknown): boolean { + if (typeof value !== "object" || value === null) return false; + const usage = value as Record; + return ["inputTokens", "outputTokens", "cacheCreationTokens", "cacheReadTokens", "thinkingTokens"].some( + (key) => typeof usage[key] === "number" && Number(usage[key]) > 0, + ); +} + +/** What the session reported during one invocation that decides its usage convention. */ +interface TurnObservation { + /** Models the session used during the turn, from settings_updated notifications. */ + models: Set; + /** The turn's completion reported usage from Task subagents. */ + childUsage: boolean; +} + class DroidSdkSessionHandle implements DroidSessionHandle { private readonly session: DroidSdkSession; private readonly runtimeRoot?: string; + private readonly configuredModel: string; + private readonly conventions: DroidUsageConventions; private stopped = false; - constructor(session: DroidSdkSession, runtimeRoot?: string) { + constructor( + session: DroidSdkSession, + configuredModel: string, + conventions: DroidUsageConventions, + runtimeRoot?: string, + ) { this.session = session; + this.configuredModel = configuredModel; + this.conventions = conventions; this.runtimeRoot = runtimeRoot; } + private currentModel(): string { + const modelId = this.session.settings?.modelId; + return typeof modelId === "string" && modelId.length > 0 ? modelId : this.configuredModel; + } + + /** + * Watches this session's notifications for the whole invocation: every + * model it switches to and whether the turn's completion carries child usage. + * Without notifications the turn cannot be checked, so it counts as having + * child usage. + */ + private observeTurn(): { observation: TurnObservation; stop(): void } { + const observation: TurnObservation = { models: new Set([this.currentModel()]), childUsage: false }; + if (typeof this.session.onNotification !== "function") { + observation.childUsage = true; + return { observation, stop: () => undefined }; + } + const stop = this.session.onNotification((notification) => { + const params = notification.params as { sessionId?: unknown; notification?: Record } | undefined; + const inner = params?.notification; + if (inner === undefined || (params?.sessionId !== undefined && params.sessionId !== this.session.id)) return; + if (inner.type === "settings_updated") { + const modelId = (inner.settings as { modelId?: unknown } | undefined)?.modelId; + if (typeof modelId === "string" && modelId.length > 0) observation.models.add(modelId); + } else if (inner.type === "agent_turn_completed" && hasTokenCounts(inner.childTokenUsage)) { + observation.childUsage = true; + } + }); + return { observation, stop }; + } + async prompt(text: string): Promise { if (this.stopped) { throw new LifecycleError("droid_session_stopped", "The Droid session is already stopped"); } + const turn = this.observeTurn(); let terminal: unknown; try { for await (const message of this.session.stream(text)) { @@ -700,13 +944,21 @@ class DroidSdkSessionHandle implements DroidSessionHandle { } } catch (error) { throw new LifecycleError("droid_session_failed", sdkErrorMessage(error), { retryable: true }); + } finally { + turn.stop(); } if (terminal === undefined) { throw new LifecycleError("droid_session_failed", "Droid session ended without a terminal result", { retryable: true, }); } - return normalizeResult(terminal); + // Droid adds subagent usage to the turn's counters, and a subagent can use + // another model; a model change mixes conventions too. Such a turn keeps + // the reported counters without a cache convention. + turn.observation.models.add(this.currentModel()); + const conventions = new Set([...turn.observation.models].map((model) => this.conventions.forModel(model))); + const convention = conventions.size === 1 && !turn.observation.childUsage ? [...conventions][0] : undefined; + return normalizeResult(terminal, convention); } async stop(): Promise { @@ -804,7 +1056,12 @@ export class DroidSdkSessionFactory implements DroidSessionFactory { await applyToolPolicy(session, input.config, mcpServers); await validateMcpTools(session, mcpServers, this.mcpStartupTimeoutMs); await validateSkills(session, runtimeProfile.entrypoints); - return new DroidSdkSessionHandle(session, runtimeProfile.root); + return new DroidSdkSessionHandle( + session, + selected.model, + new DroidUsageConventions(runtimeProfile.settingsFiles), + runtimeProfile.root, + ); } catch (error) { await session?.close().catch(() => undefined); if (runtimeProfile.root !== undefined) { diff --git a/adapters/typescript/droid/test/droid-sdk.test.mjs b/adapters/typescript/droid/test/droid-sdk.test.mjs index cd4c98e73..8421dd814 100644 --- a/adapters/typescript/droid/test/droid-sdk.test.mjs +++ b/adapters/typescript/droid/test/droid-sdk.test.mjs @@ -66,14 +66,42 @@ function fakeSdk(options = {}) { create: [], createEnvironment: [], stream: [], close: 0, listTools: 0, updateSettings: [], listMcpTools: 0, listSkills: 0, }; const results = [...(options.results ?? [result(), result({ text: "second", turnCount: 2 })])]; + const observers = new Set(); + let turn = 0; const session = { id: "droid-session-1", + settings: { modelId: options.modelId }, async *stream(prompt) { calls.stream.push(prompt); if (options.streamError) throw options.streamError; + const index = turn++; yield { type: "assistant", text: "intermediate" }; + options.duringTurn?.(index, session); + // Droid's SDK delivers the turn's completion notification before the result. + const notification = { + type: "agent_turn_completed", + reason: "completed", + turnId: `turn-${index}`, + tokenUsage: results[0]?.tokenUsage, + ...(options.childUsage?.[index] === undefined ? {} : { childTokenUsage: options.childUsage[index] }), + }; + for (const observer of observers) { + observer({ jsonrpc: "2.0", method: "droid.session_notification", params: { sessionId: session.id, notification } }); + } yield results.shift(); }, + onNotification(callback) { + observers.add(callback); + return () => observers.delete(callback); + }, + // Droid's SDK applies settings_updated to session.settings and passes it to observers. + changeModel(modelId) { + session.settings.modelId = modelId; + const notification = { type: "settings_updated", settings: { modelId } }; + for (const observer of observers) { + observer({ jsonrpc: "2.0", method: "droid.session_notification", params: { sessionId: session.id, notification } }); + } + }, async close() { calls.close += 1; }, async listTools() { calls.listTools += 1; @@ -98,6 +126,8 @@ function fakeSdk(options = {}) { }; return { calls, + observers, + session, loader: async () => ({ async createSession(value) { calls.create.push(value); @@ -127,11 +157,16 @@ test("maps model, credential, append instructions, and an exact built-in tool po assert.deepEqual(sdk.calls.stream, ["one", "two"]); assert.deepEqual(first.usage, { input_tokens: 4, + cached_input_tokens: 2, output_tokens: 3, - total_tokens: 7, extensions: { - cache_read_tokens: 2, - cache_write_tokens: 1, + "nemo_fabric.detailed_token_usage.v1": { + input_tokens: 4, + cached_input_tokens: 2, + cache_write_input_tokens: 1, + output_tokens: 3, + output_tokens_include_reasoning: true, + }, thinking_tokens: 5, factory_credits: 0.25, }, @@ -417,11 +452,16 @@ test("normalizes interrupted and failed terminal results", async () => { errorMessage: "provider failed", usage: { input_tokens: 4, + cached_input_tokens: 2, output_tokens: 3, - total_tokens: 7, extensions: { - cache_read_tokens: 2, - cache_write_tokens: 1, + "nemo_fabric.detailed_token_usage.v1": { + input_tokens: 4, + cached_input_tokens: 2, + cache_write_input_tokens: 1, + output_tokens: 3, + output_tokens_include_reasoning: true, + }, thinking_tokens: 5, factory_credits: 0.25, }, @@ -436,6 +476,341 @@ test("normalizes interrupted and failed terminal results", async () => { await handle.stop(); }); +async function usageForModel(t, model, tokenUsage, customModel) { + const configured = input({ tools: undefined }); + configured.config.models.default.model = model; + if (customModel !== undefined) { + const home = await mkdtemp(join(tmpdir(), "droid-usage-home-")); + t.after(() => rm(home, { recursive: true, force: true })); + await mkdir(join(home, ".factory"), { recursive: true }); + await writeFile(join(home, ".factory", "settings.json"), JSON.stringify({ customModels: [customModel] })); + configured.runtimeContext.environment.env.FACTORY_HOME_OVERRIDE = home; + } + const sdk = fakeSdk({ results: [result({ tokenUsage: { thinkingTokens: 0, ...tokenUsage } })] }); + const handle = await new DroidSdkSessionFactory(sdk.loader, 10_000, {}).create(configured); + try { + return (await handle.prompt("one")).usage; + } finally { + await handle.stop(); + } +} + +function inclusiveUsage({ input, cached, written, output }) { + const cache = { + ...(cached === undefined ? {} : { cached_input_tokens: cached }), + }; + return { + input_tokens: input, + ...cache, + input_tokens_include_cache: true, + output_tokens: output, + total_tokens: input + output, + extensions: { + "nemo_fabric.detailed_token_usage.v1": { + input_tokens: input, + ...cache, + ...(written === undefined ? {} : { cache_write_input_tokens: written }), + input_tokens_include_cache: true, + output_tokens: output, + output_tokens_include_reasoning: true, + }, + }, + }; +} + +const chatModel = { + model: "nvidia/nemotron-3-super-120b-a12b", + baseUrl: "https://integrate.api.nvidia.com/v1", + provider: "generic-chat-completion-api", +}; + +test("normalizes input to include cache on each established Droid provider path", async (t) => { + // Gemini: Droid persists promptTokenCount, which already includes the cache reads. + assert.deepEqual( + await usageForModel(t, "gemini-3-pro-preview", { inputTokens: 100, outputTokens: 5, cacheCreationTokens: 0, cacheReadTokens: 80 }), + inclusiveUsage({ input: 100, cached: 80, output: 5 }), + ); + assert.deepEqual( + await usageForModel(t, "gemini-3-pro-preview", { inputTokens: 100, outputTokens: 5, cacheCreationTokens: 0, cacheReadTokens: 0 }), + inclusiveUsage({ input: 100, cached: 0, output: 5 }), + ); + // Anthropic Messages: input excludes cache reads and writes. + assert.deepEqual( + await usageForModel(t, "claude-sonnet-4-5-20250929", { inputTokens: 10, outputTokens: 5, cacheCreationTokens: 20, cacheReadTokens: 80 }), + inclusiveUsage({ input: 110, cached: 80, written: 20, output: 5 }), + ); + // OpenAI Responses: Droid subtracts the cache reads from input_tokens. + assert.deepEqual( + await usageForModel(t, "gpt-5.1", { inputTokens: 20, outputTokens: 5, cacheCreationTokens: 0, cacheReadTokens: 80 }), + inclusiveUsage({ input: 100, cached: 80, output: 5 }), + ); + // Factory-hosted model on the Chat Completions handler: Droid subtracts the cache reads. + assert.deepEqual( + await usageForModel(t, "glm-5", { inputTokens: 20, outputTokens: 5, cacheCreationTokens: 0, cacheReadTokens: 80 }), + inclusiveUsage({ input: 100, cached: 80, output: 5 }), + ); + // Custom Chat Completions model: Droid subtracts the cache reads. + assert.deepEqual( + await usageForModel(t, "custom:nvidia/nemotron-3-super-120b-a12b", { inputTokens: 20, outputTokens: 5, cacheCreationTokens: 0, cacheReadTokens: 80 }, chatModel), + inclusiveUsage({ input: 100, cached: 80, output: 5 }), + ); + // Custom Bedrock Converse model: input excludes cache reads and writes. + assert.deepEqual( + await usageForModel( + t, + "custom:anthropic.claude-sonnet-4-5", + { inputTokens: 10, outputTokens: 5, cacheCreationTokens: 20, cacheReadTokens: 80 }, + { model: "anthropic.claude-sonnet-4-5", provider: "bedrock-converse" }, + ), + inclusiveUsage({ input: 110, cached: 80, written: 20, output: 5 }), + ); +}); + +async function usageForRun(model, turns) { + const configured = input({ tools: undefined }); + configured.config.models.default.model = model; + const sdk = fakeSdk({ results: turns.map((tokenUsage) => result({ tokenUsage: { thinkingTokens: 0, ...tokenUsage } })) }); + const handle = await new DroidSdkSessionFactory(sdk.loader, 10_000, {}).create(configured); + try { + const usage = []; + for (const [index] of turns.entries()) usage.push((await handle.prompt(`turn ${index}`)).usage); + return usage; + } finally { + await handle.stop(); + } +} + +test("keeps every count known across a normal cached Claude and Gemini run", async () => { + // First turn writes the prompt cache, later turns read it. + assert.deepEqual( + await usageForRun("claude-sonnet-4-5-20250929", [ + { inputTokens: 5_000, outputTokens: 200, cacheCreationTokens: 4_000, cacheReadTokens: 0 }, + { inputTokens: 300, outputTokens: 150, cacheCreationTokens: 500, cacheReadTokens: 9_000 }, + ]), + [ + inclusiveUsage({ input: 9_000, cached: 0, written: 4_000, output: 200 }), + inclusiveUsage({ input: 9_800, cached: 9_000, written: 500, output: 150 }), + ], + ); + assert.deepEqual( + await usageForRun("gemini-3-pro-preview", [ + { inputTokens: 6_000, outputTokens: 120, cacheCreationTokens: 0, cacheReadTokens: 0 }, + { inputTokens: 6_400, outputTokens: 90, cacheCreationTokens: 0, cacheReadTokens: 5_800 }, + ]), + [ + inclusiveUsage({ input: 6_000, cached: 0, output: 120 }), + inclusiveUsage({ input: 6_400, cached: 5_800, output: 90 }), + ], + ); +}); + +async function usageForSession(t, { model, droidDefinition, turns, beforeTurn, ...sdkOptions }) { + const root = await mkdtemp(join(tmpdir(), "droid-usage-session-")); + t.after(() => rm(root, { recursive: true, force: true })); + const home = join(root, "home"); + const workspace = join(root, "workspace"); + await mkdir(join(home, ".factory"), { recursive: true }); + await mkdir(workspace, { recursive: true }); + if (droidDefinition !== undefined) { + await mkdir(join(workspace, ".factory", "droids"), { recursive: true }); + await writeFile(join(workspace, ".factory", "droids", "reviewer.md"), droidDefinition); + } + const configured = input({ tools: undefined }); + configured.config.models.default.model = model; + configured.runtimeContext.environment.workspace = workspace; + configured.runtimeContext.environment.env.FACTORY_HOME_OVERRIDE = home; + const sdk = fakeSdk({ + ...sdkOptions, + results: turns.map((tokenUsage) => result({ tokenUsage: { thinkingTokens: 0, ...tokenUsage } })), + }); + const handle = await new DroidSdkSessionFactory(sdk.loader, 10_000, {}).create(configured); + try { + const usage = []; + for (const [index] of turns.entries()) { + beforeTurn?.(index, sdk.session); + usage.push((await handle.prompt(`turn ${index}`)).usage); + } + return usage; + } finally { + await handle.stop(); + } +} + +function reportedUsage({ input, cached, written, output }) { + return { + input_tokens: input, + cached_input_tokens: cached, + output_tokens: output, + extensions: { + "nemo_fabric.detailed_token_usage.v1": { + input_tokens: input, + cached_input_tokens: cached, + ...(written === undefined ? {} : { cache_write_input_tokens: written }), + output_tokens: output, + output_tokens_include_reasoning: true, + }, + }, + }; +} + +// Droid sums a turn's parent and child usage: parent Claude 10 uncached plus 20 +// cached, child Gemini 100 inclusive with 80 cached, outputs 3 and 5. +const MIXED_TURN = { inputTokens: 110, outputTokens: 8, cacheCreationTokens: 0, cacheReadTokens: 100 }; +const GEMINI_CHILD = { inputTokens: 100, outputTokens: 5, cacheCreationTokens: 0, cacheReadTokens: 80, thinkingTokens: 0 }; + +test("keeps the reported counters of a turn with subagent usage", async (t) => { + // Droid definitions can name a child model in YAML forms a simple scan misses; + // the adapter does not read them and keeps every turn with child usage raw. + for (const droidDefinition of [ + '---\nname: worker\ndescription: Worker\n"model": gemini-3-pro-preview\n---\nWork.\n', + "---\n name: worker\n description: Worker\n model: gemini-3-pro-preview\n---\nWork.\n", + ]) { + assert.deepEqual( + await usageForSession(t, { + model: "claude-sonnet-4-5-20250929", + droidDefinition, + turns: [MIXED_TURN], + childUsage: [GEMINI_CHILD], + }), + [reportedUsage({ input: 110, cached: 100, output: 8 })], + ); + } + // A subagent on the session's own model is kept raw too: the result does not + // say which model the child used. + assert.deepEqual( + await usageForSession(t, { + model: "claude-sonnet-4-5-20250929", + turns: [{ inputTokens: 40, outputTokens: 8, cacheCreationTokens: 0, cacheReadTokens: 60 }], + childUsage: [{ inputTokens: 30, outputTokens: 5, cacheCreationTokens: 0, cacheReadTokens: 40, thinkingTokens: 0 }], + }), + [reportedUsage({ input: 40, cached: 60, output: 8 })], + ); + // A turn without child usage normalizes as usual. + assert.deepEqual( + await usageForSession(t, { + model: "claude-sonnet-4-5-20250929", + turns: [{ inputTokens: 10, outputTokens: 3, cacheCreationTokens: 0, cacheReadTokens: 20 }], + childUsage: [{ inputTokens: 0, outputTokens: 0, cacheCreationTokens: 0, cacheReadTokens: 0, thinkingTokens: 0 }], + }), + [inclusiveUsage({ input: 30, cached: 20, output: 3 })], + ); +}); + +test("follows the session's effective model and model changes", async (t) => { + const googleTurn = { inputTokens: 100, outputTokens: 5, cacheCreationTokens: 0, cacheReadTokens: 80 }; + const claudeTurn = { inputTokens: 10, outputTokens: 5, cacheCreationTokens: 0, cacheReadTokens: 20 }; + // Droid reports Gemini as the session model although Claude was configured. + assert.deepEqual( + await usageForSession(t, { model: "claude-sonnet-4-5-20250929", modelId: "gemini-3-pro-preview", turns: [googleTurn] }), + [inclusiveUsage({ input: 100, cached: 80, output: 5 })], + ); + // The model changes between turns, then from Claude to Gemini and back within one turn. + assert.deepEqual( + await usageForSession(t, { + model: "claude-sonnet-4-5-20250929", + turns: [claudeTurn, googleTurn, MIXED_TURN], + beforeTurn(index, session) { + if (index === 1) session.changeModel("gemini-3-pro-preview"); + if (index === 2) session.changeModel("claude-sonnet-4-5-20250929"); + }, + duringTurn(index, session) { + if (index !== 2) return; + session.changeModel("gemini-3-pro-preview"); + session.changeModel("claude-sonnet-4-5-20250929"); + }, + }), + [ + inclusiveUsage({ input: 30, cached: 20, output: 5 }), + inclusiveUsage({ input: 100, cached: 80, output: 5 }), + reportedUsage({ input: 110, cached: 100, output: 8 }), + ], + ); +}); + +test("keeps the reported counters for Droid's router models", async (t) => { + for (const model of ["auto", "auto-fast"]) { + assert.deepEqual( + await usageForSession(t, { model, turns: [{ inputTokens: 100, outputTokens: 5, cacheCreationTokens: 0, cacheReadTokens: 80 }] }), + [reportedUsage({ input: 100, cached: 80, output: 5 })], + ); + } +}); + +test("keeps Droid input as reported when the provider convention is not established", async (t) => { + const reported = (input, cached, written) => ({ + input_tokens: input, + cached_input_tokens: cached, + output_tokens: 5, + extensions: { + "nemo_fabric.detailed_token_usage.v1": { + input_tokens: input, + cached_input_tokens: cached, + ...(written === undefined ? {} : { cache_write_input_tokens: written }), + output_tokens: 5, + output_tokens_include_reasoning: true, + }, + }, + }); + // Responses cache writes are a nonstandard field Droid does not subtract. + assert.deepEqual( + await usageForModel(t, "gpt-5.1", { inputTokens: 20, outputTokens: 5, cacheCreationTokens: 10, cacheReadTokens: 80 }), + reported(20, 80, 10), + ); + // Droid can route a Gemini model configured as Chat Completions to its Gemini handler. + assert.deepEqual( + await usageForModel( + t, + "custom:google/gemini-3-pro", + { inputTokens: 100, outputTokens: 5, cacheCreationTokens: 0, cacheReadTokens: 80 }, + { ...chatModel, model: "google/gemini-3-pro" }, + ), + reported(100, 80), + ); + // Without a matching settings entry, the custom model's provider is unknown. + assert.deepEqual( + await usageForModel(t, "custom:nvidia/nemotron-3-super-120b-a12b", { inputTokens: 100, outputTokens: 5, cacheCreationTokens: 0, cacheReadTokens: 80 }), + reported(100, 80), + ); + // Droid selects the model for auto. + assert.deepEqual( + await usageForModel(t, "auto", { inputTokens: 100, outputTokens: 5, cacheCreationTokens: 0, cacheReadTokens: 80 }), + reported(100, 80), + ); +}); + +test("leaves token counts that Droid filled with zero unknown", async () => { + const zeroUsage = { inputTokens: 0, outputTokens: 0, cacheCreationTokens: 0, cacheReadTokens: 0, thinkingTokens: 0 }; + const sdk = fakeSdk({ results: [ + result({ tokenUsage: { ...zeroUsage, inputTokens: 9 } }), + result({ tokenUsage: { ...zeroUsage, outputTokens: 2, cacheReadTokens: 6 } }), + result({ tokenUsage: { ...zeroUsage, thinkingTokens: 3, factoryCredits: 0.5 } }), + result({ tokenUsage: { ...zeroUsage, factoryCredits: 0 } }), + ] }); + const handle = await new DroidSdkSessionFactory(sdk.loader).create(input({ tools: undefined })); + assert.deepEqual((await handle.prompt("one")).usage, { + input_tokens: 9, + extensions: { + "nemo_fabric.detailed_token_usage.v1": { input_tokens: 9 }, + }, + }); + assert.deepEqual((await handle.prompt("two")).usage, { + cached_input_tokens: 6, + output_tokens: 2, + extensions: { + "nemo_fabric.detailed_token_usage.v1": { + cached_input_tokens: 6, + output_tokens: 2, + output_tokens_include_reasoning: true, + }, + }, + }); + assert.deepEqual((await handle.prompt("three")).usage, { + extensions: { thinking_tokens: 3, factory_credits: 0.5 }, + }); + assert.equal((await handle.prompt("four")).usage, undefined); + await handle.stop(); +}); + test("rejects unknown tools and cleans up after partial startup", async () => { const sdk = fakeSdk(); await assert.rejects( diff --git a/docs/integrations/harness/droid.mdx b/docs/integrations/harness/droid.mdx index 1c8065d2c..e61c6c73c 100644 --- a/docs/integrations/harness/droid.mdx +++ b/docs/integrations/harness/droid.mdx @@ -200,8 +200,71 @@ invocation streams one prompt through the same session, and runtime shutdown closes the session and its CLI subprocess. Independent runtimes do not share session history. -Successful results include normalized response text and token usage. Usage -extensions preserve cache, thinking-token, and Factory-credit metadata. +Successful results include normalized response text and token usage for the +invocation. `cached_input_tokens` holds the cache reads, and `output_tokens` +includes reasoning tokens. The `nemo_fabric.detailed_token_usage.v1` usage +extension adds the cache writes as `cache_write_input_tokens`. Usage extensions +also keep Droid's `thinking_tokens` count and its Factory credits as +`factory_credits`. Droid estimates thinking tokens from the length of the +thinking text for Anthropic, Google, and Amazon Bedrock models, so +`thinking_tokens` is not a provider count. + +In Droid CLI 0.233.0, the input count follows the cache convention of the +provider handler that serves the session's model. The adapter reads the +session's current model from the Droid SDK, so a model that Droid resolved +differently from the configuration, or a model change between invocations, +selects its own convention. The model ID, or the `provider` of a `custom:` +model in Droid's settings, selects the convention: + +| Model | Droid's input count | `input_tokens` | +| --- | --- | --- | +| `gemini-*` | Includes cache reads | Droid's input count | +| `claude-*`, the Factory-hosted `deepseek-*`, `glm-*`, `kimi-*`, `minimax-*`, `mistral-*`, `nemotron-*`, and `qwen*` models, and `anthropic` or `bedrock-converse` custom models | Excludes cache reads and cache writes | Input plus cache reads and cache writes | +| `gpt-*`, `grok-*`, and `openai` or `generic-chat-completion-api` custom models | Excludes cache reads | Input plus cache reads | + +For these models, when Droid reports an input count, `input_tokens_include_cache` +is `true` and `total_tokens` adds input and output. When a turn includes child +usage whose convention cannot be established, or crosses input conventions, +the adapter preserves reported counters without asserting cache inclusion. +Fallback turns preserve Droid's reported input and leave +`input_tokens_include_cache` unset. They omit `total_tokens` when either +recorded cache counter is positive. When both cache counters are zero and input +and output are positive, the total is their sum. This applies in the following +cases: + +- The model has no known convention, including Droid's `auto` and `auto-fast` + router models. +- A `generic-chat-completion-api` custom model's model, base model, or base URL + names Gemini or Google, because Droid can send it to its Gemini handler. +- An OpenAI Responses invocation reports cache writes, because Droid does not + subtract those writes. +- The session switches during the invocation to a model with a different + convention, even if it switches back. +- The invocation includes usage from a Task subagent. Droid adds subagent + usage to the invocation's counts, and its result does not say which model + each subagent used. + +Droid records a count that the provider omits as 0, so the adapter treats a 0 +count as unknown and leaves it out, along with a total built from it. The +exception is a 0 cache read count for the models listed below. Droid also uses +0 for missing Factory credits, so 0 credits are left out too. When every count +and the Factory credits are 0, the result has no usage. + +Droid zero-fills counts before aggregating model calls. A positive turn total +can therefore omit usage from another call. The adapter cannot establish +completeness from the terminal counters. + +Anthropic, OpenAI Responses, and Bedrock Converse report a cache read count on +every response. In Google's proto3 usage schema, an omitted +`cachedContentTokenCount` decodes as zero. This does not establish that every +model call contributed a complete usage report. For `claude-*`, `gemini-*`, +and `gpt-*` models and `bedrock-converse` custom models, `cached_input_tokens` +is therefore 0 when Droid reports no cache reads. For +other models, a missing `cached_input_tokens` can mean zero cache reads or an +unreported count. An exact aggregate cache-hit rate cannot be determined from +these results. Report cache-count coverage alongside any rate calculated from +known counts. + Result extensions preserve the Droid session ID, finish reason, duration, and turn count. diff --git a/sdk/python/nemo-fabric-adapter-catalog/src/nemo_fabric_adapter_catalog/catalog.json b/sdk/python/nemo-fabric-adapter-catalog/src/nemo_fabric_adapter_catalog/catalog.json index 4c15e1112..4fd1f7e05 100644 --- a/sdk/python/nemo-fabric-adapter-catalog/src/nemo_fabric_adapter_catalog/catalog.json +++ b/sdk/python/nemo-fabric-adapter-catalog/src/nemo_fabric_adapter_catalog/catalog.json @@ -597,18 +597,38 @@ "$schema": "https://json-schema.org/draft/2020-12/schema", "additionalProperties": false, "properties": { - "cache_read_tokens": { - "minimum": 0, - "type": "integer" - }, - "cache_write_tokens": { - "minimum": 0, - "type": "integer" - }, "factory_credits": { "minimum": 0, "type": "number" }, + "nemo_fabric.detailed_token_usage.v1": { + "additionalProperties": false, + "properties": { + "cache_write_input_tokens": { + "minimum": 0, + "type": "integer" + }, + "cached_input_tokens": { + "minimum": 0, + "type": "integer" + }, + "input_tokens": { + "minimum": 0, + "type": "integer" + }, + "input_tokens_include_cache": { + "type": "boolean" + }, + "output_tokens": { + "minimum": 0, + "type": "integer" + }, + "output_tokens_include_reasoning": { + "type": "boolean" + } + }, + "type": "object" + }, "thinking_tokens": { "minimum": 0, "type": "integer" @@ -660,7 +680,7 @@ "name": "nemo-fabric-adapters-droid", "version": "0.6.0" }, - "sha256": "6d1b3e67bf4d7b803c292057a27aa6c1c272f1160118cf0bb349ee273b067f14", + "sha256": "1bf85e37dc31e8accec221b88e7fd8881f89c9b556c8fb9fc76fdae022310730", "source": "adapters/typescript/droid/droid.fabric-adapter.json" }, "nvidia.fabric.hermes": { From e23cf97149191821b742c53c576d4a72270437ff Mon Sep 17 00:00:00 2001 From: mnajafian-nv Date: Tue, 20 Oct 2026 02:00:50 -0700 Subject: [PATCH 04/26] fix: leave Droid totals unknown when cache reads are unreported Droid stores an omitted cache read count as 0. For Grok, Factory-hosted models, and anthropic, openai and generic-chat-completion-api custom models, the endpoint is not shown to report cache reads on every response, so that 0 may hide cache hits that belong in the input. The adapter already withheld cached_input_tokens there, but still added the 0 to the input, set input_tokens_include_cache and derived a total, so a glm-5 turn with omitted cache hits reported a known input and total that were too small. Inclusive input now needs a known cache read count: a positive count, or a provider that reports cache reads on every response. Otherwise the turn keeps Droid's input count, leaves input_tokens_include_cache unset and derives no total. Claude, GPT, Bedrock Converse and Gemini turns with 0 cache reads stay fully known. The other fallback turns (the auto routers, turns with subagent usage and turns that switch conventions) derived a total from Droid's input when both cache counters were 0. That 0 can be an omission too, so those turns now derive no total either. Signed-off-by: mnajafian-nv --- adapters/typescript/droid/src/droid-sdk.ts | 18 +++++----- .../typescript/droid/test/droid-sdk.test.mjs | 33 +++++++++++++++++++ docs/integrations/harness/droid.mdx | 23 ++++++++----- 3 files changed, 57 insertions(+), 17 deletions(-) diff --git a/adapters/typescript/droid/src/droid-sdk.ts b/adapters/typescript/droid/src/droid-sdk.ts index 7044313f4..ceb568060 100644 --- a/adapters/typescript/droid/src/droid-sdk.ts +++ b/adapters/typescript/droid/src/droid-sdk.ts @@ -676,20 +676,22 @@ function normalizeUsage(value: unknown, convention: DroidUsageConvention | undef // Droid stores a count the provider omitted as 0, so a 0 count is unknown. // Output includes reasoning on every provider path. const output = outputTokens > 0 ? outputTokens : undefined; - const inclusiveInput = inputIncludingCache(convention?.input, inputTokens, cacheReadTokens, cacheCreationTokens); - const cacheReadsKnown = cacheReadTokens > 0 || (inclusiveInput !== undefined && convention?.reportsCacheReads === true); - // Without an inclusive count, report Droid's own input count and leave its - // cache convention unset. A turn with cache counts then has no total. - const hasCache = cacheReadTokens > 0 || cacheCreationTokens > 0; + // A 0 cache read count is a real 0 only on paths that report cache reads on + // every response; elsewhere it may hide cache hits that belong in the input. + const cacheReadsKnown = cacheReadTokens > 0 || convention?.reportsCacheReads === true; + const inclusiveInput = cacheReadsKnown + ? inputIncludingCache(convention?.input, inputTokens, cacheReadTokens, cacheCreationTokens) + : undefined; + // Without an inclusive count, report Droid's own input count, leave its + // cache convention unset, and derive no total: the input may leave out cache. const input = inclusiveInput ?? (inputTokens > 0 ? inputTokens : undefined); - const totalKnown = inclusiveInput !== undefined || !hasCache; - const total = !totalKnown || input === undefined || output === undefined ? undefined : input + output; + const total = inclusiveInput === undefined || output === undefined ? undefined : inclusiveInput + output; if ((input !== undefined && !Number.isSafeInteger(input)) || (total !== undefined && !Number.isSafeInteger(total))) { throw new LifecycleError("droid_malformed_result", "Droid returned malformed usage metadata"); } const inputUsage: JsonObject = { ...(input === undefined ? {} : { input_tokens: input }), - ...(cacheReadsKnown ? { cached_input_tokens: cacheReadTokens } : {}), + ...(cacheReadTokens > 0 || inclusiveInput !== undefined ? { cached_input_tokens: cacheReadTokens } : {}), ...(inclusiveInput === undefined ? {} : { input_tokens_include_cache: true }), }; const detailed: JsonObject = { diff --git a/adapters/typescript/droid/test/droid-sdk.test.mjs b/adapters/typescript/droid/test/droid-sdk.test.mjs index 8421dd814..91b6a1599 100644 --- a/adapters/typescript/droid/test/droid-sdk.test.mjs +++ b/adapters/typescript/droid/test/droid-sdk.test.mjs @@ -736,6 +736,39 @@ test("keeps the reported counters for Droid's router models", async (t) => { } }); +test("leaves inclusive input unknown when a 0 cache read may be an omission", async (t) => { + const uncachedTurn = { inputTokens: 20, outputTokens: 5, cacheCreationTokens: 0, cacheReadTokens: 0 }; + // Factory-hosted endpoints are not shown to report cache reads on every + // response, so Droid's 0 may hide cache hits that belong in the input. + assert.deepEqual(await usageForModel(t, "glm-5", uncachedTurn), { + input_tokens: 20, + output_tokens: 5, + extensions: { + "nemo_fabric.detailed_token_usage.v1": { input_tokens: 20, output_tokens: 5, output_tokens_include_reasoning: true }, + }, + }); + // Anthropic reports cache reads on every response, so its 0 is a real 0. + assert.deepEqual( + await usageForModel(t, "claude-sonnet-4-5-20250929", uncachedTurn), + inclusiveUsage({ input: 20, cached: 0, output: 5 }), + ); + // A fallback turn derives no total either: its 0 cache reads are not established. + assert.deepEqual( + await usageForSession(t, { + model: "glm-5", + turns: [uncachedTurn], + childUsage: [{ inputTokens: 8, outputTokens: 2, cacheCreationTokens: 0, cacheReadTokens: 0, thinkingTokens: 0 }], + }), + [{ + input_tokens: 20, + output_tokens: 5, + extensions: { + "nemo_fabric.detailed_token_usage.v1": { input_tokens: 20, output_tokens: 5, output_tokens_include_reasoning: true }, + }, + }], + ); +}); + test("keeps Droid input as reported when the provider convention is not established", async (t) => { const reported = (input, cached, written) => ({ input_tokens: input, diff --git a/docs/integrations/harness/droid.mdx b/docs/integrations/harness/droid.mdx index e61c6c73c..2f32cbfab 100644 --- a/docs/integrations/harness/droid.mdx +++ b/docs/integrations/harness/droid.mdx @@ -222,15 +222,16 @@ model in Droid's settings, selects the convention: | `claude-*`, the Factory-hosted `deepseek-*`, `glm-*`, `kimi-*`, `minimax-*`, `mistral-*`, `nemotron-*`, and `qwen*` models, and `anthropic` or `bedrock-converse` custom models | Excludes cache reads and cache writes | Input plus cache reads and cache writes | | `gpt-*`, `grok-*`, and `openai` or `generic-chat-completion-api` custom models | Excludes cache reads | Input plus cache reads | -For these models, when Droid reports an input count, `input_tokens_include_cache` -is `true` and `total_tokens` adds input and output. When a turn includes child +For these models, when Droid reports an input count and the cache read count +is known, `input_tokens_include_cache` is `true` and `total_tokens` adds input +and output. The cache read count is known when it is positive, or when the +provider reports cache reads on every response, as described below. When a turn includes child usage whose convention cannot be established, or crosses input conventions, the adapter preserves reported counters without asserting cache inclusion. Fallback turns preserve Droid's reported input and leave -`input_tokens_include_cache` unset. They omit `total_tokens` when either -recorded cache counter is positive. When both cache counters are zero and input -and output are positive, the total is their sum. This applies in the following -cases: +`input_tokens_include_cache` and `total_tokens` unset, because Droid's input +may leave out cache, including cache reads it recorded as 0. This applies in +the following cases: - The model has no known convention, including Droid's `auto` and `auto-fast` router models. @@ -259,9 +260,13 @@ every response. In Google's proto3 usage schema, an omitted `cachedContentTokenCount` decodes as zero. This does not establish that every model call contributed a complete usage report. For `claude-*`, `gemini-*`, and `gpt-*` models and `bedrock-converse` custom models, `cached_input_tokens` -is therefore 0 when Droid reports no cache reads. For -other models, a missing `cached_input_tokens` can mean zero cache reads or an -unreported count. An exact aggregate cache-hit rate cannot be determined from +is therefore 0 when Droid reports no cache reads. For other models, a missing +`cached_input_tokens` can mean zero cache reads or an unreported count. When a +`grok-*` model, a Factory-hosted model from the table, or an `anthropic`, +`openai`, or `generic-chat-completion-api` custom model reports no cache reads, +`input_tokens` is Droid's input count, and `input_tokens_include_cache` and +`total_tokens` are unset, because unreported cache reads would belong in the +input. An exact aggregate cache-hit rate cannot be determined from these results. Report cache-count coverage alongside any rate calculated from known counts. From 4585ac7ee03f0d055f71fc8dd46eef91778fbd97 Mon Sep 17 00:00:00 2001 From: mnajafian-nv Date: Tue, 20 Oct 2026 08:51:15 -0700 Subject: [PATCH 05/26] fix: report Deep Agents cached input tokens LangChain carries cache reads per model call in usage_metadata.input_token_details.cache_read, and OpenAI raw usage carries them in prompt_tokens_details.cached_tokens, but the adapter never reported cached_input_tokens. Sum cached input per call with the same rule as the other counts: one call without a cache count leaves the sum unknown, and a reported 0 stays 0. Raw OpenAI usage is read first because ChatOpenAI renames the cache key for priority and flex service tiers. Report input_tokens_include_cache=True whenever input tokens are known. LangChain's UsageMetadata.input_tokens counts every input token type; langchain-anthropic adds cache reads and writes back into it, and OpenAI's prompt_tokens already includes cached tokens. Signed-off-by: mnajafian-nv --- adapters/python/deepagents/README.md | 7 ++ .../deepagents/adapter.py | 30 +++++ tests/adapters/test_deepagents.py | 113 ++++++++++++++++-- 3 files changed, 141 insertions(+), 9 deletions(-) diff --git a/adapters/python/deepagents/README.md b/adapters/python/deepagents/README.md index 66a44eece..9e41e8f53 100644 --- a/adapters/python/deepagents/README.md +++ b/adapters/python/deepagents/README.md @@ -252,6 +252,13 @@ and the provider's total is used when present. Other providers and streamed responses use LangChain's `usage_metadata`, which some integrations fill with 0 for an omitted count. +Cached input tokens come from `prompt_tokens_details.cached_tokens` in +OpenAI-shaped raw usage and from `usage_metadata.input_token_details.cache_read` +otherwise. They follow the same rule: one call without a cache count leaves the +invocation's cached count unknown, and a reported 0 stays 0. When input tokens +are known, the adapter reports that they include cached tokens, because +LangChain's `UsageMetadata.input_tokens` counts every input token type. + ## Runtime Lifecycle NeMo Fabric starts one local adapter host for every runtime. During runtime start, diff --git a/adapters/python/deepagents/src/nemo_fabric_adapters/deepagents/adapter.py b/adapters/python/deepagents/src/nemo_fabric_adapters/deepagents/adapter.py index 37771f419..ba304d6b3 100644 --- a/adapters/python/deepagents/src/nemo_fabric_adapters/deepagents/adapter.py +++ b/adapters/python/deepagents/src/nemo_fabric_adapters/deepagents/adapter.py @@ -770,6 +770,7 @@ async def invoke( "input_tokens": usage.get("prompt_tokens"), "output_tokens": usage.get("completion_tokens"), "total_tokens": usage.get("total_tokens"), + "cached_input_tokens": usage.get("cached_tokens"), } tokens = { name: value @@ -778,6 +779,11 @@ async def invoke( and not isinstance(value, bool) and 0 <= value <= (1 << 64) - 1 } + if "input_tokens" in tokens: + # LangChain's UsageMetadata defines input_tokens as the sum of every + # input token type, cache reads included; OpenAI's raw prompt_tokens + # includes cached tokens too. + tokens["input_tokens_include_cache"] = True cost = usage.get("cost") if ( not isinstance(cost, (int, float)) @@ -1300,6 +1306,9 @@ def _aggregate_usage(calls: list[dict[str, Any]]) -> dict[str, Any] | None: values = [_reported_tokens(call, source) for call in calls] if all(value is not None for value in values): totals[target] = sum(values) + cached = [_reported_cached_tokens(call) for call in calls] + if all(value is not None for value in cached): + totals["cached_tokens"] = sum(cached) costs = [_reported_cost(call) for call in calls] if all(cost is not None for cost in costs): totals["cost"] = math.fsum(costs) @@ -1342,6 +1351,27 @@ def _reported_tokens(call: dict[str, Any], name: str) -> int | None: return None +def _reported_cached_tokens(call: dict[str, Any]) -> int | None: + """Return one call's cache-read input tokens, or ``None`` when not reported. + + OpenAI-shaped raw usage reports them as ``prompt_tokens_details.cached_tokens``; + ``ChatOpenAI`` renames that key for priority and flex service tiers, so the raw + count is read when present. Other responses use LangChain's + ``usage_metadata.input_token_details.cache_read``. + """ + + metadata = call.get("response_metadata") + raw = metadata.get("token_usage") if isinstance(metadata, dict) else None + if isinstance(raw, dict): + details, key = raw.get("prompt_tokens_details"), "cached_tokens" + else: + usage = call.get("usage") + details = usage.get("input_token_details") if isinstance(usage, dict) else None + key = "cache_read" + value = details.get(key) if isinstance(details, dict) else None + return value if _is_token_count(value) else None + + def _is_token_count(value: Any) -> bool: return ( isinstance(value, int) diff --git a/tests/adapters/test_deepagents.py b/tests/adapters/test_deepagents.py index c916cb2a3..5268f3537 100644 --- a/tests/adapters/test_deepagents.py +++ b/tests/adapters/test_deepagents.py @@ -2604,11 +2604,17 @@ async def test_summarization_model_call_usage_is_counted( assert "tools" not in summary_request assert "tools" in agent_request assert first.usage == AgentUsage( - input_tokens=1500, output_tokens=10, total_tokens=1510 + input_tokens=1500, + output_tokens=10, + total_tokens=1510, + input_tokens_include_cache=True, ) assert second.status is AgentRunStatus.SUCCEEDED assert second.usage == AgentUsage( - input_tokens=1700, output_tokens=65, total_tokens=1765 + input_tokens=1700, + output_tokens=65, + total_tokens=1765, + input_tokens_include_cache=True, ) @@ -2658,7 +2664,10 @@ async def test_subagent_model_calls_are_counted( } assert len(subgraphs) == 2 assert result.usage == AgentUsage( - input_tokens=1111, output_tokens=1111, total_tokens=2222 + input_tokens=1111, + output_tokens=1111, + total_tokens=2222, + input_tokens_include_cache=True, ) @@ -2680,7 +2689,10 @@ async def test_recursion_limit_failure_keeps_usage_of_completed_model_calls( assert result.status is AgentRunStatus.FAILED assert result.error.code == "deepagents_recursion_limit_reached" assert result.usage == AgentUsage( - input_tokens=10 * calls, output_tokens=calls, total_tokens=11 * calls + input_tokens=10 * calls, + output_tokens=calls, + total_tokens=11 * calls, + input_tokens_include_cache=True, ) @@ -2701,7 +2713,12 @@ def provider_failure(kind: str) -> Any: # usage stays known. pytest.param( "400", - AgentUsage(input_tokens=120, output_tokens=8, total_tokens=128), + AgentUsage( + input_tokens=120, + output_tokens=8, + total_tokens=128, + input_tokens_include_cache=True, + ), id="400-rejection", ), # The provider may have done work without reporting usage. @@ -2785,12 +2802,17 @@ async def test_recorder_keeps_a_cancelled_call_unreported(): [ pytest.param( [{"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0}], - AgentUsage(input_tokens=0, output_tokens=0, total_tokens=0), + AgentUsage( + input_tokens=0, + output_tokens=0, + total_tokens=0, + input_tokens_include_cache=True, + ), id="explicit-zero", ), pytest.param( [{"prompt_tokens": 100}], - AgentUsage(input_tokens=100), + AgentUsage(input_tokens=100, input_tokens_include_cache=True), id="input-only", ), pytest.param( @@ -2800,7 +2822,12 @@ async def test_recorder_keeps_a_cancelled_call_unreported(): ), pytest.param( [{"prompt_tokens": 100, "completion_tokens": 5}], - AgentUsage(input_tokens=100, output_tokens=5, total_tokens=105), + AgentUsage( + input_tokens=100, + output_tokens=5, + total_tokens=105, + input_tokens_include_cache=True, + ), id="total-omitted", ), pytest.param( @@ -2810,7 +2837,12 @@ async def test_recorder_keeps_a_cancelled_call_unreported(): ), pytest.param( [{"prompt_tokens": 7, "completion_tokens": 5, "total_tokens": 0}], - AgentUsage(input_tokens=7, output_tokens=5, total_tokens=0), + AgentUsage( + input_tokens=7, + output_tokens=5, + total_tokens=0, + input_tokens_include_cache=True, + ), id="provider-total-zero-kept", ), pytest.param( @@ -2839,6 +2871,69 @@ async def test_partial_usage_reports_keep_missing_counts_unknown( assert result.usage == expected +@pytest.mark.parametrize( + ("cached", "expected"), + [ + pytest.param([40, 60], 100, id="every-call-reports-cache"), + pytest.param([40, None], None, id="one-call-omits-cache"), + pytest.param([0, 0], 0, id="explicit-zero"), + ], +) +async def test_cached_input_tokens_are_known_only_when_every_call_reported_them( + tmp_path, make_payload, openai_endpoint, cached, expected +): + def usage(cached_tokens: int | None) -> dict[str, Any]: + reported: dict[str, Any] = { + "prompt_tokens": 100, + "completion_tokens": 5, + "total_tokens": 105, + } + if cached_tokens is not None: + reported["prompt_tokens_details"] = {"cached_tokens": cached_tokens} + return reported + + openai_endpoint["replies"] += [ + chat_completion(usage(cached[0]), tool_calls=[LS_TOOL_CALL]), + chat_completion(usage(cached[1])), + ] + + (result,) = await run_invocations(make_payload(tmp_path), 1) + + assert result.status is AgentRunStatus.SUCCEEDED + assert result.usage == AgentUsage( + input_tokens=200, + cached_input_tokens=expected, + input_tokens_include_cache=True, + output_tokens=10, + total_tokens=210, + ) + + +@pytest.mark.parametrize( + ("cache_read", "expected"), + [ + pytest.param(30, 30, id="reported"), + pytest.param(0, 0, id="explicit-zero"), + pytest.param(None, None, id="omitted"), + ], +) +def test_anthropic_cache_reads_are_counted_as_included_input(cache_read, expected): + # langchain-anthropic adds cache reads and writes to input_tokens and keeps an + # omitted cache count out of input_token_details. + from anthropic.types import Usage + from langchain_anthropic.chat_models import _create_usage_metadata + + usage_metadata = _create_usage_metadata( + Usage(input_tokens=10, output_tokens=5, cache_read_input_tokens=cache_read) + ) + totals = adapter._aggregate_usage( + [adapter._message_to_dict(AIMessage("done", usage_metadata=usage_metadata))] + ) + + assert totals.get("cached_tokens") == expected + assert totals["prompt_tokens"] == 10 + (cache_read or 0) + + async def test_bad_mcp_transport_fails_runtime_start(tmp_path, make_payload): # A misconfigured MCP server must fail loudly, not be silently dropped. payload = make_payload(tmp_path) From fee4c9287832d1fd24664bfa1848ac784af66bad Mon Sep 17 00:00:00 2001 From: mnajafian-nv Date: Tue, 20 Oct 2026 09:30:54 -0700 Subject: [PATCH 06/26] fix: report zero Deep Agents usage when no model call ran The adapter drops a model call that the provider rejects with a 4xx status other than 408, because the provider spent no tokens on it. When every call was rejected, or no call started, the call list was empty and the usage was reported as unknown, although the invocation provably spent no tokens. The adapter now reports 0 for every token count and the cost in that case. Signed-off-by: mnajafian-nv --- adapters/python/deepagents/README.md | 15 +++++----- .../deepagents/adapter.py | 6 ++-- tests/adapters/test_deepagents.py | 29 +++++++++++++++++++ 3 files changed, 40 insertions(+), 10 deletions(-) diff --git a/adapters/python/deepagents/README.md b/adapters/python/deepagents/README.md index 9e41e8f53..e88b41c04 100644 --- a/adapters/python/deepagents/README.md +++ b/adapters/python/deepagents/README.md @@ -237,13 +237,14 @@ summarization. Agent Protocol subagents (entries with `graph_id`) run in their remote service, and their tokens are not included. A token count or cost is reported only when every counted call reported it; otherwise it stays unknown. This also holds for a failed invocation. A request the provider rejects with a -4xx status (other than 408) contributes no usage. A 5xx response, timeout, -cancellation or unfinished call leaves the invocation's counts and cost -unknown, because the provider may have done work without reporting usage. -OpenAI SDK retries happen inside one LangChain call, so a 5xx the SDK retries -to success is not visible to the adapter, and only the final response's -reported usage counts. Usage replayed from a configured LangChain model cache -would be counted as reported; NeMo Fabric configures no such cache. +4xx status (other than 408) contributes no usage, so an invocation that made no +model call, or whose every call was rejected, reports 0 tokens and 0 cost. A +5xx response, timeout, cancellation or unfinished call leaves the invocation's +counts and cost unknown, because the provider may have done work without +reporting usage. OpenAI SDK retries happen inside one LangChain call, so a 5xx +the SDK retries to success is not visible to the adapter, and only the final +response's reported usage counts. Usage replayed from a configured LangChain +model cache would be counted as reported; NeMo Fabric configures no such cache. When a response carries OpenAI-shaped usage in `response_metadata.token_usage`, as `ChatOpenAI` responses without streaming diff --git a/adapters/python/deepagents/src/nemo_fabric_adapters/deepagents/adapter.py b/adapters/python/deepagents/src/nemo_fabric_adapters/deepagents/adapter.py index ba304d6b3..e5561f215 100644 --- a/adapters/python/deepagents/src/nemo_fabric_adapters/deepagents/adapter.py +++ b/adapters/python/deepagents/src/nemo_fabric_adapters/deepagents/adapter.py @@ -1292,11 +1292,11 @@ def _aggregate_usage(calls: list[dict[str, Any]]) -> dict[str, Any] | None: """Sum model-call usage; a count is known only when every call reported it. One call without a count makes the invocation's sum for that count a - partial value, so it stays unknown rather than under-reported. + partial value, so it stays unknown rather than under-reported. Without a + call, because none started or the provider rejected every one, the + invocation spent nothing, so every count and the cost are 0. """ - if not calls: - return None totals: dict[str, Any] = {} for source, target in ( ("input_tokens", "prompt_tokens"), diff --git a/tests/adapters/test_deepagents.py b/tests/adapters/test_deepagents.py index 5268f3537..6cdb0f466 100644 --- a/tests/adapters/test_deepagents.py +++ b/tests/adapters/test_deepagents.py @@ -2744,6 +2744,35 @@ async def test_failed_model_call_usage_follows_provider_response( assert result.usage == expected +async def test_rejected_only_model_call_reports_zero_usage( + tmp_path, make_payload, openai_endpoint +): + # The provider rejected the only model call, so the invocation spent nothing. + openai_endpoint["replies"].append(provider_failure("400")) + + (result,) = await run_invocations(make_payload(tmp_path), 1) + + assert result.status is AgentRunStatus.FAILED + assert result.usage == AgentUsage( + input_tokens=0, + cached_input_tokens=0, + output_tokens=0, + total_tokens=0, + input_tokens_include_cache=True, + cost_usd=0.0, + ) + + +def test_aggregate_usage_without_model_calls_is_zero(): + assert adapter._aggregate_usage([]) == { + "prompt_tokens": 0, + "completion_tokens": 0, + "total_tokens": 0, + "cached_tokens": 0, + "cost": 0.0, + } + + @pytest.mark.parametrize( ("failure", "expected"), [ From 2ee128d7485f15db1a5b0c66f98cbfdb04b30277 Mon Sep 17 00:00:00 2001 From: mnajafian-nv Date: Tue, 20 Oct 2026 09:31:49 -0700 Subject: [PATCH 07/26] fix: report unknown Deep Agents usage with an Agent Protocol subagent Deep Agents routes a subagent spec with graph_id to its async subagent middleware, which starts and updates tasks on a remote Agent Protocol server. Those model calls never reach the adapter's callbacks, so the reported counts left them out while appearing complete. The adapter blocked such subagents only when a tools policy was configured. When the configured subagents include one with graph_id, the adapter now reports the invocation's token counts and cost as unknown. The local sums stay in output.usage. Signed-off-by: mnajafian-nv --- adapters/python/deepagents/README.md | 24 ++++++++-------- .../deepagents/adapter.py | 25 ++++++++++++++++- tests/adapters/test_deepagents.py | 28 +++++++++++++++++++ 3 files changed, 65 insertions(+), 12 deletions(-) diff --git a/adapters/python/deepagents/README.md b/adapters/python/deepagents/README.md index e88b41c04..56f93989d 100644 --- a/adapters/python/deepagents/README.md +++ b/adapters/python/deepagents/README.md @@ -234,17 +234,19 @@ runtime start before an invocation is accepted. Usage sums every chat-model call that runs in the adapter process for the invocation: the main agent, declarative subagents, and Deep Agents summarization. Agent Protocol subagents (entries with `graph_id`) run in their -remote service, and their tokens are not included. A token count or cost is -reported only when every counted call reported it; otherwise it stays unknown. -This also holds for a failed invocation. A request the provider rejects with a -4xx status (other than 408) contributes no usage, so an invocation that made no -model call, or whose every call was rejected, reports 0 tokens and 0 cost. A -5xx response, timeout, cancellation or unfinished call leaves the invocation's -counts and cost unknown, because the provider may have done work without -reporting usage. OpenAI SDK retries happen inside one LangChain call, so a 5xx -the SDK retries to success is not visible to the adapter, and only the final -response's reported usage counts. Usage replayed from a configured LangChain -model cache would be counted as reported; NeMo Fabric configures no such cache. +remote service, where the adapter cannot count their calls, so a runtime +configured with one reports unknown token counts and cost. A token count or +cost is reported only when every counted call reported it; otherwise it stays +unknown. This also holds for a failed invocation. A request the provider +rejects with a 4xx status (other than 408) contributes no usage, so an +invocation that made no model call, or whose every call was rejected, reports 0 +tokens and 0 cost. A 5xx response, timeout, cancellation or unfinished call +leaves the invocation's counts and cost unknown, because the provider may have +done work without reporting usage. OpenAI SDK retries happen inside one +LangChain call, so a 5xx the SDK retries to success is not visible to the +adapter, and only the final response's reported usage counts. Usage replayed +from a configured LangChain model cache would be counted as reported; NeMo +Fabric configures no such cache. When a response carries OpenAI-shaped usage in `response_metadata.token_usage`, as `ChatOpenAI` responses without streaming diff --git a/adapters/python/deepagents/src/nemo_fabric_adapters/deepagents/adapter.py b/adapters/python/deepagents/src/nemo_fabric_adapters/deepagents/adapter.py index e5561f215..38a0eacbf 100644 --- a/adapters/python/deepagents/src/nemo_fabric_adapters/deepagents/adapter.py +++ b/adapters/python/deepagents/src/nemo_fabric_adapters/deepagents/adapter.py @@ -541,6 +541,18 @@ def _validated_deepagents_settings(extra: Any) -> dict[str, Any]: return dict(extra) +def _has_agent_protocol_subagent(subagents: Any) -> bool: + """Whether ``create_deep_agent`` runs a subagent on an Agent Protocol server. + + Deep Agents routes a subagent spec with ``graph_id`` to its async subagent + middleware, which starts and updates tasks on a remote server. + """ + return isinstance(subagents, list) and any( + isinstance(subagent, dict) and "graph_id" in subagent + for subagent in subagents + ) + + def _gate_subagent( subagent: dict[str, Any], enabled: set[str] | None, blocked: set[str] ) -> dict[str, Any]: @@ -612,6 +624,7 @@ def __init__(self) -> None: self._callback_handler_type: Any = None self._telemetry_quarantine: str | None = None self._telemetry_quarantine_cause: str | None = None + self._has_remote_subagent = False async def start(self, payload: dict[str, Any]) -> None: if self._started: @@ -662,6 +675,9 @@ async def start(self, payload: dict[str, Any]) -> None: model, settings, ) + self._has_remote_subagent = _has_agent_protocol_subagent( + agent_kwargs.get("subagents") + ) self._checkpointer = await open_checkpointer( checkpointer_path(runtime_context, base_dir, runtime_id) ) @@ -765,7 +781,14 @@ async def invoke( failed = bool(output.pop("failed", False)) reported_error = output.pop("error", None) raw_usage = output.get("usage") - usage = raw_usage if isinstance(raw_usage, dict) else {} + # An Agent Protocol subagent's model calls run on its remote server and + # never reach this process, so the local sums may undercount the + # invocation and its counts and cost stay unknown. + usage = ( + raw_usage + if isinstance(raw_usage, dict) and not self._has_remote_subagent + else {} + ) token_fields = { "input_tokens": usage.get("prompt_tokens"), "output_tokens": usage.get("completion_tokens"), diff --git a/tests/adapters/test_deepagents.py b/tests/adapters/test_deepagents.py index 6cdb0f466..873126845 100644 --- a/tests/adapters/test_deepagents.py +++ b/tests/adapters/test_deepagents.py @@ -2763,6 +2763,34 @@ async def test_rejected_only_model_call_reports_zero_usage( ) +async def test_agent_protocol_subagent_makes_usage_unknown( + tmp_path, make_payload, openai_endpoint +): + # An Agent Protocol subagent runs its model calls on a remote server, which + # the adapter cannot count, so local sums would undercount the invocation. + payload = make_payload(tmp_path) + payload["config"]["harness"]["settings"]["deepagents"] = { + "subagents": [ + { + "name": "researcher", + "description": "Researches on a remote server.", + "graph_id": "research_agent", + "url": "https://agents.example.test", + } + ] + } + openai_endpoint["replies"].append( + chat_completion( + {"prompt_tokens": 120, "completion_tokens": 8, "total_tokens": 128} + ) + ) + + (result,) = await run_invocations(payload, 1) + + assert result.status is AgentRunStatus.SUCCEEDED + assert result.usage is None + + def test_aggregate_usage_without_model_calls_is_zero(): assert adapter._aggregate_usage([]) == { "prompt_tokens": 0, From 06de4c0281439b64a9d05340ebca771278fe5743 Mon Sep 17 00:00:00 2001 From: mnajafian-nv Date: Tue, 20 Oct 2026 09:32:07 -0700 Subject: [PATCH 08/26] docs: say a failed Deep Agents call may have been billed The usage notes said that after a 5xx response, timeout, or cancelled call the provider may have done work. The reason the counts stay unknown is narrower: the provider may have billed the call without reporting usage. The README, the recorder docstring, and a test comment now say that. Signed-off-by: mnajafian-nv --- adapters/python/deepagents/README.md | 2 +- .../deepagents/src/nemo_fabric_adapters/deepagents/adapter.py | 4 ++-- tests/adapters/test_deepagents.py | 2 +- 3 files changed, 4 insertions(+), 4 deletions(-) diff --git a/adapters/python/deepagents/README.md b/adapters/python/deepagents/README.md index 56f93989d..c1a5a8c98 100644 --- a/adapters/python/deepagents/README.md +++ b/adapters/python/deepagents/README.md @@ -242,7 +242,7 @@ rejects with a 4xx status (other than 408) contributes no usage, so an invocation that made no model call, or whose every call was rejected, reports 0 tokens and 0 cost. A 5xx response, timeout, cancellation or unfinished call leaves the invocation's counts and cost unknown, because the provider may have -done work without reporting usage. OpenAI SDK retries happen inside one +billed the call without reporting usage. OpenAI SDK retries happen inside one LangChain call, so a 5xx the SDK retries to success is not visible to the adapter, and only the final response's reported usage counts. Usage replayed from a configured LangChain model cache would be counted as reported; NeMo diff --git a/adapters/python/deepagents/src/nemo_fabric_adapters/deepagents/adapter.py b/adapters/python/deepagents/src/nemo_fabric_adapters/deepagents/adapter.py index 38a0eacbf..f886155c9 100644 --- a/adapters/python/deepagents/src/nemo_fabric_adapters/deepagents/adapter.py +++ b/adapters/python/deepagents/src/nemo_fabric_adapters/deepagents/adapter.py @@ -1012,8 +1012,8 @@ class ModelCallUsage(BaseCallbackHandler): record when it ends. A request the provider rejects with a 4xx status (other than 408) contributes no usage and is dropped. A 5xx response, timeout, cancellation or unfinished call stays unreported and makes the invocation's - counts and cost unknown, because the provider may have done work without - reporting usage. + counts and cost unknown, because the provider may have billed the call + without reporting usage. """ run_inline = True diff --git a/tests/adapters/test_deepagents.py b/tests/adapters/test_deepagents.py index 873126845..eec895d6a 100644 --- a/tests/adapters/test_deepagents.py +++ b/tests/adapters/test_deepagents.py @@ -2721,7 +2721,7 @@ def provider_failure(kind: str) -> Any: ), id="400-rejection", ), - # The provider may have done work without reporting usage. + # The provider may have billed the call without reporting usage. pytest.param("408", None, id="408-request-timeout"), pytest.param("503", None, id="503-server-error"), pytest.param("timeout", None, id="read-timeout"), From e535051318613392a86fcb75cd4a469cedd21b18 Mon Sep 17 00:00:00 2001 From: mnajafian-nv Date: Tue, 20 Oct 2026 09:38:36 -0700 Subject: [PATCH 09/26] style: format the Deep Agents remote subagent check Apply ruff format to the check added in the previous commit. Signed-off-by: mnajafian-nv --- .../deepagents/src/nemo_fabric_adapters/deepagents/adapter.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/adapters/python/deepagents/src/nemo_fabric_adapters/deepagents/adapter.py b/adapters/python/deepagents/src/nemo_fabric_adapters/deepagents/adapter.py index f886155c9..9a51bd817 100644 --- a/adapters/python/deepagents/src/nemo_fabric_adapters/deepagents/adapter.py +++ b/adapters/python/deepagents/src/nemo_fabric_adapters/deepagents/adapter.py @@ -548,8 +548,7 @@ def _has_agent_protocol_subagent(subagents: Any) -> bool: middleware, which starts and updates tasks on a remote server. """ return isinstance(subagents, list) and any( - isinstance(subagent, dict) and "graph_id" in subagent - for subagent in subagents + isinstance(subagent, dict) and "graph_id" in subagent for subagent in subagents ) From cb42bfe7d37d76beb93d1c0d8aa7f6bf61edda53 Mon Sep 17 00:00:00 2001 From: mnajafian-nv Date: Tue, 20 Oct 2026 09:45:54 -0700 Subject: [PATCH 10/26] test: keep zero observed Deep Agents calls unknown with a remote subagent The adapter reports 0 when its callbacks observed no billed model call. That 0 holds for the invocation only when every model call runs through those callbacks. Deep Agents runs declarative subagents and summarization inside the graph, where the run config's callbacks reach them, but an Agent Protocol subagent (graph_id) runs on a remote server that the callbacks never see. The adapter already reports unknown usage for a runtime with such a subagent; this test pins that for an invocation whose only observed call was rejected, and the docstring and README state the condition. Signed-off-by: mnajafian-nv --- adapters/python/deepagents/README.md | 18 +++++++------- .../deepagents/adapter.py | 5 +++- tests/adapters/test_deepagents.py | 24 +++++++++++++++++++ 3 files changed, 37 insertions(+), 10 deletions(-) diff --git a/adapters/python/deepagents/README.md b/adapters/python/deepagents/README.md index c1a5a8c98..6340b4ce3 100644 --- a/adapters/python/deepagents/README.md +++ b/adapters/python/deepagents/README.md @@ -238,15 +238,15 @@ remote service, where the adapter cannot count their calls, so a runtime configured with one reports unknown token counts and cost. A token count or cost is reported only when every counted call reported it; otherwise it stays unknown. This also holds for a failed invocation. A request the provider -rejects with a 4xx status (other than 408) contributes no usage, so an -invocation that made no model call, or whose every call was rejected, reports 0 -tokens and 0 cost. A 5xx response, timeout, cancellation or unfinished call -leaves the invocation's counts and cost unknown, because the provider may have -billed the call without reporting usage. OpenAI SDK retries happen inside one -LangChain call, so a 5xx the SDK retries to success is not visible to the -adapter, and only the final response's reported usage counts. Usage replayed -from a configured LangChain model cache would be counted as reported; NeMo -Fabric configures no such cache. +rejects with a 4xx status (other than 408) contributes no usage, so without an +Agent Protocol subagent an invocation that made no model call, or whose every +call was rejected, reports 0 tokens and 0 cost. A 5xx response, timeout, +cancellation or unfinished call leaves the invocation's counts and cost +unknown, because the provider may have billed the call without reporting usage. +OpenAI SDK retries happen inside one LangChain call, so a 5xx the SDK retries +to success is not visible to the adapter, and only the final response's +reported usage counts. Usage replayed from a configured LangChain model cache +would be counted as reported; NeMo Fabric configures no such cache. When a response carries OpenAI-shaped usage in `response_metadata.token_usage`, as `ChatOpenAI` responses without streaming diff --git a/adapters/python/deepagents/src/nemo_fabric_adapters/deepagents/adapter.py b/adapters/python/deepagents/src/nemo_fabric_adapters/deepagents/adapter.py index 9a51bd817..bfda396dc 100644 --- a/adapters/python/deepagents/src/nemo_fabric_adapters/deepagents/adapter.py +++ b/adapters/python/deepagents/src/nemo_fabric_adapters/deepagents/adapter.py @@ -1316,7 +1316,10 @@ def _aggregate_usage(calls: list[dict[str, Any]]) -> dict[str, Any] | None: One call without a count makes the invocation's sum for that count a partial value, so it stays unknown rather than under-reported. Without a call, because none started or the provider rejected every one, the - invocation spent nothing, so every count and the cost are 0. + observed calls spent nothing, so every count and the cost are 0. That 0 + covers the invocation only when every model call runs through these + callbacks; the caller makes usage unknown when an Agent Protocol subagent + can run models on a remote server. """ totals: dict[str, Any] = {} diff --git a/tests/adapters/test_deepagents.py b/tests/adapters/test_deepagents.py index eec895d6a..816ccb1ec 100644 --- a/tests/adapters/test_deepagents.py +++ b/tests/adapters/test_deepagents.py @@ -2791,6 +2791,30 @@ async def test_agent_protocol_subagent_makes_usage_unknown( assert result.usage is None +async def test_rejected_only_call_with_agent_protocol_subagent_is_unknown( + tmp_path, make_payload, openai_endpoint +): + # Zero observed calls prove nothing when a remote subagent can run models + # that the adapter's callbacks never see. + payload = make_payload(tmp_path) + payload["config"]["harness"]["settings"]["deepagents"] = { + "subagents": [ + { + "name": "researcher", + "description": "Researches on a remote server.", + "graph_id": "research_agent", + "url": "https://agents.example.test", + } + ] + } + openai_endpoint["replies"].append(provider_failure("400")) + + (result,) = await run_invocations(payload, 1) + + assert result.status is AgentRunStatus.FAILED + assert result.usage is None + + def test_aggregate_usage_without_model_calls_is_zero(): assert adapter._aggregate_usage([]) == { "prompt_tokens": 0, From ce744150ea1e26be5b6b8a42a26d95b8a533d77f Mon Sep 17 00:00:00 2001 From: mnajafian-nv Date: Thu, 22 Oct 2026 13:04:44 -0700 Subject: [PATCH 11/26] fix: preserve empty OpenShell environment arguments Signed-off-by: mnajafian-nv --- examples/openshell/run-demo.sh | 4 ++- tests/scripts/test_openshell_launcher.py | 44 +++++++++++++++++++++++- 2 files changed, 46 insertions(+), 2 deletions(-) diff --git a/examples/openshell/run-demo.sh b/examples/openshell/run-demo.sh index 56d2b534f..4fe36156e 100644 --- a/examples/openshell/run-demo.sh +++ b/examples/openshell/run-demo.sh @@ -272,9 +272,11 @@ run_consumer() { if [[ "${EXAMPLE_MODE}" == caller_owned || "${EXAMPLE_MODE}" == both ]]; then CALLER_SANDBOX_ATTEMPTED=1 + # Guard empty arrays for Bash 3.2's set -u behavior without adding empty args. "${OPENSHELL_CLI}" --gateway-endpoint "http://127.0.0.1:${GATEWAY_PORT}" \ --workspace default sandbox create --name "${SANDBOX_NAME}" --from "${RUNTIME_IMAGE}" \ - --policy "${POLICY_FILE}" "${MODEL_ENV_ARGS[@]}" "${SERVER_ENV_ARGS[@]}" \ + --policy "${POLICY_FILE}" ${MODEL_ENV_ARGS[@]+"${MODEL_ENV_ARGS[@]}"} \ + ${SERVER_ENV_ARGS[@]+"${SERVER_ENV_ARGS[@]}"} \ --detach --no-tty -- fabric-agent-ctl serve SANDBOX_JSON="$("${OPENSHELL_CLI}" --gateway-endpoint "http://127.0.0.1:${GATEWAY_PORT}" \ --workspace default sandbox get "${SANDBOX_NAME}" --output json)" diff --git a/tests/scripts/test_openshell_launcher.py b/tests/scripts/test_openshell_launcher.py index 317b8052d..15ab68caa 100644 --- a/tests/scripts/test_openshell_launcher.py +++ b/tests/scripts/test_openshell_launcher.py @@ -12,6 +12,7 @@ import json import os from pathlib import Path +import shlex import shutil import socket import subprocess @@ -172,7 +173,22 @@ def _run(launcher, workload, env): ) -def test_generic_workload_gets_ownership_and_consumer_args(launcher): +@pytest.mark.parametrize( + ("model_env_args", "server_env_args"), + [ + pytest.param([], None, id="default-server-env"), + pytest.param([], [], id="no-env"), + pytest.param(["--env", "MODEL_HINT=words with spaces"], [], id="no-server-env"), + pytest.param( + ["--env", "MODEL_HINT=words with spaces"], + ["--env", "SERVER_HINT=words with spaces"], + id="quoted-env", + ), + ], +) +def test_generic_workload_gets_ownership_and_consumer_args( + launcher, model_env_args, server_env_args +): _, root, _ = launcher workload = root / "third-workload" workload.mkdir() @@ -185,6 +201,12 @@ def test_generic_workload_gets_ownership_and_consumer_args(launcher): stage_workload_image() { :; } run_workload_consumer() { local mode="$1"; shift; uv run --no-sync consumer "$mode" "$@"; } """ + + f"MODEL_ENV_ARGS=({shlex.join(model_env_args)})\n" + + ( + f"SERVER_ENV_ARGS=({shlex.join(server_env_args)})\n" + if server_env_args is not None + else "" + ) ) completed = _run(launcher, hook, {"OPENSHELL_EXAMPLE_MODE": "both"}) @@ -195,6 +217,26 @@ def test_generic_workload_gets_ownership_and_consumer_args(launcher): assert all(call["args"][-2:] == ["--input", "words with spaces"] for call in calls) assert all(call["provider"].endswith("/fabric-environment-openshell") for call in calls) assert not (root / "sandbox-state.json").exists() + create = next( + operation["args"] + for operation in _operations(root) + if operation["command"] == "openshell" + and "sandbox" in operation["args"] + and "create" in operation["args"] + ) + expected_server_env = ( + server_env_args + if server_env_args is not None + else [ + "--env", + "ADAPTER_PYTHON=/opt/fabric/bin/python", + "--env", + "TMPDIR=/sandbox/tmp", + ] + ) + assert create[create.index("--policy") + 2 : create.index("--detach")] == ( + model_env_args + expected_server_env + ) @pytest.mark.parametrize("mode", ["caller_owned", "fabric_owned"]) From 790f3693aaebf9465bc224bd76da41fb3d81d420 Mon Sep 17 00:00:00 2001 From: mnajafian-nv Date: Thu, 22 Oct 2026 13:10:57 -0700 Subject: [PATCH 12/26] fix: avoid premature Relay gateway startup failures Signed-off-by: mnajafian-nv --- adapters/python/claude/README.md | 2 +- adapters/python/codex/README.md | 4 +- adapters/python/common/README.md | 14 +++++ .../common/relay_gateway.py | 21 +++++-- tests/_utils/utils.py | 22 +++++-- tests/adapters/test_adapaters_common_utils.py | 4 +- .../test_adapters_common_relay_gateway.py | 63 +++++++++++++++++++ 7 files changed, 117 insertions(+), 13 deletions(-) diff --git a/adapters/python/claude/README.md b/adapters/python/claude/README.md index a6b8ac271..bfa13847c 100644 --- a/adapters/python/claude/README.md +++ b/adapters/python/claude/README.md @@ -169,7 +169,7 @@ config.enable_relay( ``` For each Relay-enabled Claude runtime, NeMo Fabric starts one `nemo-relay` gateway, -waits for its health endpoint, and stops it with the runtime. NeMo Fabric passes the +waits up to 30 seconds for its health endpoint, and stops it with the runtime. NeMo Fabric passes the gateway URL to the connected Claude Code process through `ANTHROPIC_BASE_URL` and `NEMO_RELAY_GATEWAY_URL`, and passes the selected explicit model endpoint to the gateway as its Anthropic upstream. It also stages a runtime-scoped Claude diff --git a/adapters/python/codex/README.md b/adapters/python/codex/README.md index 3cf1de779..ba7eab1b7 100644 --- a/adapters/python/codex/README.md +++ b/adapters/python/codex/README.md @@ -154,7 +154,9 @@ Codex adapter does not provide a separate `relay` extra; its `harness` and Enable Relay with `FabricConfig.enable_relay(...)`. The adapter starts the installed `nemo-relay` CLI as a supervised sidecar; do not start the gateway -separately. +separately. Gateway readiness has a 30-second budget covering configuration +validation and required worker initialization. Startup fails on early process exit +or an expired budget, and the gateway log is retained for diagnosis. NeMo Fabric routes the selected Responses-compatible provider through the gateway and passes its explicit `base_url` to Relay as the upstream endpoint. diff --git a/adapters/python/common/README.md b/adapters/python/common/README.md index 17e537cd5..fb8400337 100644 --- a/adapters/python/common/README.md +++ b/adapters/python/common/README.md @@ -122,6 +122,20 @@ state during `start`. Adapter stdout is reserved for the protocol; diagnostics are redirected to stderr. A host crash or protocol timeout terminates that runtime. +## Relay Gateway Readiness + +The Claude and Codex adapters wait up to 30 seconds for the supervised Relay +gateway's HTTP health endpoint. This budget includes configuration validation +and required worker initialization. Early process exit fails startup immediately; +a gateway that never becomes healthy fails when the budget expires and is stopped. +The adapters retain its log for diagnosis. Local Python hosts separately enforce +a 90-second adapter startup limit. + +Relay 0.10.3 includes an +[upstream startup-memory fix](https://github.com/NVIDIA/NeMo-Relay/pull/1375). +Use that release to obtain the fix on 0.10 installations. Fabric continues to +accept the supported `>=0.9,<0.11` CLI range. + ## Relay Request Correlation In-process Relay SDK adapters that own their Agent scope can use diff --git a/adapters/python/common/src/nemo_fabric_adapters/common/relay_gateway.py b/adapters/python/common/src/nemo_fabric_adapters/common/relay_gateway.py index 45cf04faa..959f54fa7 100644 --- a/adapters/python/common/src/nemo_fabric_adapters/common/relay_gateway.py +++ b/adapters/python/common/src/nemo_fabric_adapters/common/relay_gateway.py @@ -20,7 +20,10 @@ from nemo_fabric_adapters.common.relay_version import is_supported_relay_version -RELAY_HEALTH_TIMEOUT_SECONDS = 10.0 +# Readiness includes plugin preflight and worker startup before health is served. +# This operational budget is separate from Relay's inner request timeouts and +# Fabric's outer adapter-start limit. +RELAY_STARTUP_TIMEOUT_SECONDS = 30.0 RELAY_STOP_TIMEOUT_SECONDS = 5.0 RELAY_VERSION_TIMEOUT_SECONDS = 5.0 @@ -110,24 +113,32 @@ def wait_for_relay_gateway( process: subprocess.Popen[Any], health_url: str, *, - timeout: float = RELAY_HEALTH_TIMEOUT_SECONDS, + timeout: float = RELAY_STARTUP_TIMEOUT_SECONDS, ) -> None: """Wait until the Relay health endpoint succeeds or startup fails.""" deadline = time.monotonic() + timeout - while time.monotonic() < deadline: + while True: + remaining = deadline - time.monotonic() + if remaining <= 0: + break returncode = process.poll() if returncode is not None: raise RelayGatewayError( f"NeMo Relay gateway exited with status {returncode} before becoming ready" ) try: - with urllib.request.urlopen(health_url, timeout=1) as response: + with urllib.request.urlopen( + health_url, timeout=min(1.0, remaining) + ) as response: if 200 <= response.status < 300: return except (OSError, urllib.error.URLError): pass - time.sleep(0.1) + remaining = deadline - time.monotonic() + if remaining <= 0: + break + time.sleep(min(0.1, remaining)) raise RelayGatewayError(f"NeMo Relay gateway did not become ready at {health_url}") diff --git a/tests/_utils/utils.py b/tests/_utils/utils.py index 9b7ca63be..aa5521ae8 100644 --- a/tests/_utils/utils.py +++ b/tests/_utils/utils.py @@ -22,11 +22,17 @@ def assert_atof_skill_selection( output: Mapping[str, Any], expected_skill: str | None ) -> None: records = atof_records(output) - tool_records = [record for record in records if record["category"] == "tool"] + tool_records = [ + record + for record in records + if record["kind"] == "scope" and record["category"] == "tool" + ] llm_end_records = [ record for record in records - if record["category"] == "llm" and record["scope_category"] == "end" + if record["kind"] == "scope" + and record["category"] == "llm" + and record["scope_category"] == "end" ] serialized_tool_records = json.dumps(tool_records).replace("\\\\", "/") serialized_skill_calls = json.dumps(tool_records + llm_end_records).replace( @@ -55,7 +61,9 @@ def assert_atof_model(output: Mapping[str, Any], expected_model: str) -> None: llm_starts = [ record for record in events - if record["category"] == "llm" and record["scope_category"] == "start" + if record["kind"] == "scope" + and record["category"] == "llm" + and record["scope_category"] == "start" ] assert llm_starts, events assert {record["data"]["content"]["model"] for record in llm_starts} == { @@ -81,7 +89,9 @@ def assert_semantic_relay_artifacts( llm_starts = [ event for event in events - if event.get("category") == "llm" and event.get("scope_category") == "start" + if event["kind"] == "scope" + and event["category"] == "llm" + and event["scope_category"] == "start" ] assert llm_starts, events assert all( @@ -92,7 +102,9 @@ def assert_semantic_relay_artifacts( llm_ends = [ event for event in events - if event.get("category") == "llm" and event.get("scope_category") == "end" + if event["kind"] == "scope" + and event["category"] == "llm" + and event["scope_category"] == "end" ] assert llm_ends, events assert any(_relay_event_total_tokens(event) > 0 for event in llm_ends) diff --git a/tests/adapters/test_adapaters_common_utils.py b/tests/adapters/test_adapaters_common_utils.py index d1805f12b..0be298052 100644 --- a/tests/adapters/test_adapaters_common_utils.py +++ b/tests/adapters/test_adapaters_common_utils.py @@ -1332,8 +1332,10 @@ async def test_relay_initializes_v4_atof_atif_config( } common_utils.validate_relay_observability(plugin_config) assert plugin.validate_exact(plugin_config)["config"]["diagnostics"] == [] + # Activation includes discovered host plugins, so compare the same layers. + expected_diagnostics = plugin.validate(plugin_config)["config"]["diagnostics"] async with plugin.activate(plugin_config) as activation: - assert activation.report["config"]["diagnostics"] == [] + assert activation.report["config"]["diagnostics"] == expected_diagnostics def test_relay_validates_all_v4_otlp_fields(): diff --git a/tests/adapters/test_adapters_common_relay_gateway.py b/tests/adapters/test_adapters_common_relay_gateway.py index c307254ae..e478f85de 100644 --- a/tests/adapters/test_adapters_common_relay_gateway.py +++ b/tests/adapters/test_adapters_common_relay_gateway.py @@ -271,3 +271,66 @@ def test_stop_relay_gateway_is_idempotent_for_exited_process(): process.terminate.assert_not_called() process.kill.assert_not_called() + + +@pytest.fixture(name="gateway_clock") +def gateway_clock_fixture(monkeypatch: pytest.MonkeyPatch) -> MagicMock: + mock_clock = MagicMock(spec=relay_gateway.time) + mock_clock.monotonic.return_value = 0.0 + + def advance(seconds: float) -> None: + mock_clock.monotonic.return_value += seconds + + mock_clock.sleep.side_effect = advance + monkeypatch.setattr(relay_gateway, "time", mock_clock) + return mock_clock + + +def test_wait_for_relay_gateway_allows_preflight_before_worker_readiness( + monkeypatch, gateway_clock +): + process = MagicMock() + process.poll.return_value = None + response = MagicMock() + response.__enter__.return_value.status = 200 + + def probe_health(url, *, timeout): + gateway_clock.monotonic.return_value += timeout + if gateway_clock.monotonic() < 14.0: + raise relay_gateway.urllib.error.URLError("still initializing") + return response + + monkeypatch.setattr(relay_gateway.urllib.request, "urlopen", probe_health) + + relay_gateway.wait_for_relay_gateway(process, "http://127.0.0.1:43210/healthz") + + assert 14.0 <= gateway_clock.monotonic() < 30.0 + + +@pytest.mark.parametrize( + ("probe_elapsed", "expected_sleep"), [(0.25, None), (0.20, 0.05)] +) +def test_wait_for_relay_gateway_limits_each_probe_to_remaining_startup_budget( + monkeypatch, gateway_clock, probe_elapsed, expected_sleep +): + process = MagicMock() + process.poll.return_value = None + + def probe_health(url, *, timeout): + gateway_clock.monotonic.return_value += probe_elapsed + raise relay_gateway.urllib.error.URLError("not ready") + + mock_probe = MagicMock(side_effect=probe_health) + monkeypatch.setattr(relay_gateway.urllib.request, "urlopen", mock_probe) + + with pytest.raises(relay_gateway.RelayGatewayError, match="did not become ready"): + relay_gateway.wait_for_relay_gateway( + process, "http://127.0.0.1:43210/healthz", timeout=0.25 + ) + + mock_probe.assert_called_once_with("http://127.0.0.1:43210/healthz", timeout=0.25) + if expected_sleep is None: + gateway_clock.sleep.assert_not_called() + else: + gateway_clock.sleep.assert_called_once_with(pytest.approx(expected_sleep)) + assert gateway_clock.monotonic() == pytest.approx(0.25) From 5db49175c27b37f704350cfbfa6fc51bdba1fd0e Mon Sep 17 00:00:00 2001 From: mnajafian-nv Date: Thu, 22 Oct 2026 13:41:46 -0700 Subject: [PATCH 13/26] docs: use the NeMo Fabric product name Signed-off-by: mnajafian-nv --- adapters/python/common/README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/adapters/python/common/README.md b/adapters/python/common/README.md index fb8400337..919bdd233 100644 --- a/adapters/python/common/README.md +++ b/adapters/python/common/README.md @@ -133,7 +133,7 @@ a 90-second adapter startup limit. Relay 0.10.3 includes an [upstream startup-memory fix](https://github.com/NVIDIA/NeMo-Relay/pull/1375). -Use that release to obtain the fix on 0.10 installations. Fabric continues to +Use that release to obtain the fix on 0.10 installations. NeMo Fabric continues to accept the supported `>=0.9,<0.11` CLI range. ## Relay Request Correlation From cbb06993ec9db38d0b7e2dfdd072c94eecfd16ae Mon Sep 17 00:00:00 2001 From: mnajafian-nv Date: Thu, 22 Oct 2026 14:59:52 -0700 Subject: [PATCH 14/26] fix: validate recovered usage and optional pricing rates Signed-off-by: mnajafian-nv --- .../generated/usage-accounting-vectors.ts | 180 ++++++++++++++++++ crates/fabric-core/src/runtime.rs | 7 +- crates/fabric-core/src/usage_normalization.rs | 23 +++ docs/adapter-contract/tutorials/results.md | 2 +- .../nemo_fabric.types.md | 1 + docs/sdk/python.mdx | 3 - schemas/SCHEMA.md | 9 +- .../src/nemo_fabric/types.py | 7 +- tests/adapters/test_deepagents_accounting.py | 4 +- tests/docs/test_python_api_docs.py | 4 +- .../usage_accounting/conformance.json | 180 ++++++++++++++++++ tests/python/test_model_usage.py | 22 ++- 12 files changed, 426 insertions(+), 16 deletions(-) diff --git a/adapter-contract/typescript/test/generated/usage-accounting-vectors.ts b/adapter-contract/typescript/test/generated/usage-accounting-vectors.ts index 05a913834..cbc90c193 100644 --- a/adapter-contract/typescript/test/generated/usage-accounting-vectors.ts +++ b/adapter-contract/typescript/test/generated/usage-accounting-vectors.ts @@ -493,6 +493,186 @@ export const accountingVectors = [ "observed": 20 } } + }, + { + "name": "aggregate_recovered_cache_cannot_exceed_complete_input", + "usage": { + "input_tokens": 100, + "input_tokens_include_cache": true + }, + "accounting": { + "models_exhaustive": true, + "gaps": [], + "models": [ + { + "model": "cache-model", + "gaps": [], + "reported": { + "cached_input_tokens": 200, + "cache_write_input_tokens": 0, + "output_tokens": 5, + "output_tokens_include_reasoning": true + } + } + ] + }, + "expected": { + "input_tokens": { + "complete": null, + "observed": null + }, + "cached_input_tokens": { + "complete": 200, + "observed": 200 + }, + "cache_write_input_tokens": { + "complete": 0, + "observed": 0 + }, + "output_tokens": { + "complete": 5, + "observed": 5 + }, + "total_tokens": { + "complete": null, + "observed": null + } + } + }, + { + "name": "model_recovered_request_cache_cannot_exceed_complete_input", + "usage": {}, + "accounting": { + "models_exhaustive": true, + "gaps": [], + "models": [ + { + "model": "cache-model", + "gaps": [], + "reported": { + "input_tokens": 100, + "input_tokens_include_cache": true, + "output_tokens": 5, + "output_tokens_include_reasoning": true + }, + "requests_exhaustive": true, + "requests": [ + { + "usage": { + "cached_input_tokens": 200, + "cache_write_input_tokens": 0 + } + } + ] + } + ] + }, + "expected": { + "input_tokens": { + "complete": null, + "observed": null + }, + "cached_input_tokens": { + "complete": 200, + "observed": 200 + }, + "cache_write_input_tokens": { + "complete": 0, + "observed": 0 + }, + "output_tokens": { + "complete": 5, + "observed": 5 + }, + "total_tokens": { + "complete": null, + "observed": null + } + } + }, + { + "name": "model_recovered_request_peak_cannot_exceed_complete_input", + "usage": {}, + "accounting": { + "models_exhaustive": true, + "gaps": [], + "models": [ + { + "model": "peak-model", + "gaps": [], + "reported": { + "input_tokens": 100, + "input_tokens_include_cache": true, + "output_tokens": 5, + "output_tokens_include_reasoning": true + }, + "requests_exhaustive": true, + "requests": [ + { + "usage": { + "peak_request_input_tokens": 200 + } + } + ] + } + ] + }, + "expected": { + "input_tokens": { + "complete": null, + "observed": null + }, + "output_tokens": { + "complete": 5, + "observed": 5 + }, + "total_tokens": { + "complete": null, + "observed": null + } + } + }, + { + "name": "model_recovered_request_reasoning_cannot_exceed_complete_output", + "usage": {}, + "accounting": { + "models_exhaustive": true, + "gaps": [], + "models": [ + { + "model": "reasoning-model", + "gaps": [], + "reported": { + "input_tokens": 100, + "input_tokens_include_cache": true, + "output_tokens": 5, + "output_tokens_include_reasoning": true + }, + "requests_exhaustive": true, + "requests": [ + { + "usage": { + "reasoning_tokens": 10 + } + } + ] + } + ] + }, + "expected": { + "input_tokens": { + "complete": 100, + "observed": 100 + }, + "output_tokens": { + "complete": null, + "observed": null + }, + "total_tokens": { + "complete": null, + "observed": null + } + } } ] satisfies { name: string; diff --git a/crates/fabric-core/src/runtime.rs b/crates/fabric-core/src/runtime.rs index 2b7a780b4..b03484eee 100644 --- a/crates/fabric-core/src/runtime.rs +++ b/crates/fabric-core/src/runtime.rs @@ -2324,7 +2324,7 @@ fn run_local_host_invocation_with_timeout( let model_usage = match agent_result .usage .as_ref() - .map(|raw| crate::model_usage::normalize_model_usage(raw)) + .map(crate::model_usage::normalize_model_usage) .transpose() { Ok(usage) => usage.flatten(), @@ -6114,7 +6114,7 @@ for line in sys.stdin: #[test] fn local_host_returns_flat_supported_accounting_with_the_normalized_result() { - let (_root, plan) = local_host_plan("accounted_usage_result"); + let (root, plan) = local_host_plan("accounted_usage_result"); let result = run_plan(&plan, RunRequest::text("accounted")).expect("normalized result"); assert_eq!(result.status, RunStatus::Succeeded); let usage = result.model_usage.expect("supported accounting"); @@ -6129,6 +6129,7 @@ for line in sys.stdin: assert_eq!(usage.cost.estimated.observed_nano_usd, None); assert!(usage.cost.pricing.is_none()); assert!(usage.cost.applied_rates.is_empty()); + let _ = fs::remove_dir_all(root); } #[test] @@ -6688,7 +6689,7 @@ for line in sys.stdin: let diagnostic = serde_json::to_string(events[0]).unwrap(); assert!(!diagnostic.contains("private-accounting-identity")); assert!(!diagnostic.contains("request totals")); - let _ = root; + let _ = fs::remove_dir_all(root); } } } diff --git a/crates/fabric-core/src/usage_normalization.rs b/crates/fabric-core/src/usage_normalization.rs index f236d7623..f8924207e 100644 --- a/crates/fabric-core/src/usage_normalization.rs +++ b/crates/fabric-core/src/usage_normalization.rs @@ -348,6 +348,7 @@ pub fn describe_model_usage(usage: &AgentUsage) -> Result Result Result input) || peak.is_ok_and(|peak| peak > input) + }) { + counts[INPUT] = Err("inconsistent_usage"); + observed[INPUT] = Err("inconsistent_usage"); + } + let reasoning = counts[UsageDimension::ReasoningTokens as usize]; + if counts[OUTPUT].is_ok_and(|output| reasoning.is_ok_and(|reasoning| reasoning > output)) { + counts[OUTPUT] = Err("inconsistent_usage"); + observed[OUTPUT] = Err("inconsistent_usage"); + } +} + // An unconvertible aggregate must not erase a qualified model/request subset. // Contradictions and overflow stay errors; these reasons only lack a usable // normalized observation. Recovery never certifies a complete invocation. diff --git a/docs/adapter-contract/tutorials/results.md b/docs/adapter-contract/tutorials/results.md index 42cd74b99..f361e00ba 100644 --- a/docs/adapter-contract/tutorials/results.md +++ b/docs/adapter-contract/tutorials/results.md @@ -108,7 +108,7 @@ A model aggregate is not a single request. No adapter should fabricate call boundaries, hide execution features or change retry settings to make usage look complete. Costs computed from a price table never become `cost_usd`. -## 2. Separate Failure Classes +### 2. Separate Failure Classes Use these failure classes consistently: diff --git a/docs/reference/api/python-library-reference/nemo_fabric.types.md b/docs/reference/api/python-library-reference/nemo_fabric.types.md index 57ba50589..83cdfab2b 100644 --- a/docs/reference/api/python-library-reference/nemo_fabric.types.md +++ b/docs/reference/api/python-library-reference/nemo_fabric.types.md @@ -2634,6 +2634,7 @@ The model is both attribute-accessible and mapping-compatible. A harness failure - `output`: Object-shaped adapter output as ``RunOutput``; non-object values are preserved as-is. - `error`: Structured failure, or ``None`` on success. - `usage`: Normalized invocation usage, or ``None`` when unavailable. + - `model_usage`: Complete counts, observations and accounting gaps, or ``None`` when supported accounting evidence is unavailable. - `resource_snapshot`: Runtime resource snapshot read once after the adapter returned this result, or ``None`` when the runtime's resources could not be observed. - `system_stats`: Invocation duration and endpoint observations, or None for older results or execution backends without collection. - `artifacts`: Normalized artifact manifest. diff --git a/docs/sdk/python.mdx b/docs/sdk/python.mdx index d337ed45c..f55fe8ba9 100644 --- a/docs/sdk/python.mdx +++ b/docs/sdk/python.mdx @@ -1194,9 +1194,6 @@ the SDK raises a `FabricError` subclass instead of returning a partial ### Complete counts and useful observations -The bundled Deep Agents adapter supplies this evidence for its documented collection scope. - - `result.model_usage` is present when the adapter supplies the declared `nemo_fabric.usage_accounting.v1` evidence. Its token fields are independently nullable counts: input includes cache reads and writes, output includes diff --git a/schemas/SCHEMA.md b/schemas/SCHEMA.md index 6e90b6cec..8f43fc0a8 100644 --- a/schemas/SCHEMA.md +++ b/schemas/SCHEMA.md @@ -26,6 +26,8 @@ schemas/ │ ├── agent.schema.json │ ├── run-plan.schema.json │ └── ... +├── extensions/ # Descriptor-declared extension payloads +│ └── usage-accounting.schema.json └── adapter-contract/ # Southbound adapter-facing contract ├── adapter-descriptor.schema.json ├── adapter-target-descriptor.schema.json @@ -38,9 +40,10 @@ schemas/ └── openai-stream-record.schema.json ``` -An adapter author can treat `adapter-contract/` as the complete schema entry -point. Contract versions remain explicit fields in the wire format rather than -directory names. +`adapter-contract/` defines the base adapter wire objects. Adapters returning +supported extension payloads must also declare the applicable schema from +`extensions/`. Contract versions remain explicit fields in the wire format +rather than directory names. Start with the [adapter contract overview](https://github.com/NVIDIA/NeMo-Fabric/blob/main/docs/adapter-contract/README.md) diff --git a/sdk/python/nemo-fabric-runtime/src/nemo_fabric/types.py b/sdk/python/nemo-fabric-runtime/src/nemo_fabric/types.py index 72b8ed64d..bea167f4a 100644 --- a/sdk/python/nemo-fabric-runtime/src/nemo_fabric/types.py +++ b/sdk/python/nemo-fabric-runtime/src/nemo_fabric/types.py @@ -2587,8 +2587,11 @@ def _normalize(cls, data: dict[str, Any]) -> dict[str, Any]: if re.fullmatch(r"sha256:[0-9a-f]{64}", data["catalog_sha256"]) is None: raise FabricConfigError("invalid applied pricing fingerprint") rates = _mapping(data["rates"], "pricing rates") - if set(rates) != {"input", "output", "cache_read", "cache_write"}: + allowed = {"input", "output", "cache_read", "cache_write"} + if not {"input", "output"} <= set(rates) <= allowed: raise FabricConfigError("incomplete pricing rates") + rates.setdefault("cache_read", None) + rates.setdefault("cache_write", None) for key, value in rates.items(): if value is not None or key in {"input", "output"}: _usage_counter(value, "pricing rate") @@ -2902,6 +2905,8 @@ class RunResult(FabricMapping): are preserved as-is. error: Structured failure, or ``None`` on success. usage: Normalized invocation usage, or ``None`` when unavailable. + model_usage: Complete counts, observations and accounting gaps, or + ``None`` when supported accounting evidence is unavailable. resource_snapshot: Runtime resource snapshot read once after the adapter returned this result, or ``None`` when the runtime's resources could not be observed. diff --git a/tests/adapters/test_deepagents_accounting.py b/tests/adapters/test_deepagents_accounting.py index 5ba450330..dccaf29a3 100644 --- a/tests/adapters/test_deepagents_accounting.py +++ b/tests/adapters/test_deepagents_accounting.py @@ -365,8 +365,8 @@ async def test_non_model_failure_keeps_completed_call_usage(make_model): @pytest.fixture(name="make_runtime") async def make_runtime_fixture(tmp_path, monkeypatch): runtimes = [] - os.environ["SCRIPTED_MODEL_KEY"] = "unused-scripted" - os.environ.pop("OPENAI_PROXY", None) + monkeypatch.setenv("SCRIPTED_MODEL_KEY", "unused-scripted") + monkeypatch.delenv("OPENAI_PROXY", raising=False) async def make(replies, *, identity="runtime", profile=None, barrier=None): ledger = [] diff --git a/tests/docs/test_python_api_docs.py b/tests/docs/test_python_api_docs.py index 73a2c51ac..6fbcbfec2 100644 --- a/tests/docs/test_python_api_docs.py +++ b/tests/docs/test_python_api_docs.py @@ -177,7 +177,7 @@ def test_python_reference_exactly_covers_exported_sdk_classes() -> None: for module, names in functions.items(): page = REFERENCE_DIR / f"{module}.md" documented = { - match.group(1) for match in FUNCTION_SIGNATURE.finditer(page.read_text()) + match.group(1) for match in FUNCTION_SIGNATURE.finditer(page.read_text(encoding="utf-8")) } assert names <= documented assert f'slug: "{MODULE_SLUGS[module]}"' in page.read_text(encoding="utf-8") @@ -274,7 +274,7 @@ def test_generated_signatures_preserve_python_callable_semantics(): page = REFERENCE_DIR / f"{module_name}.md" rendered = { match.group(1): match.group(2) - for match in FUNCTION_SIGNATURE.finditer(page.read_text()) + for match in FUNCTION_SIGNATURE.finditer(page.read_text(encoding="utf-8")) } for function_name in function_names: function = getattr(module, function_name) diff --git a/tests/fixtures/usage_accounting/conformance.json b/tests/fixtures/usage_accounting/conformance.json index dc9d5cf47..b5e026d61 100644 --- a/tests/fixtures/usage_accounting/conformance.json +++ b/tests/fixtures/usage_accounting/conformance.json @@ -484,5 +484,185 @@ "observed": 20 } } + }, + { + "name": "aggregate_recovered_cache_cannot_exceed_complete_input", + "usage": { + "input_tokens": 100, + "input_tokens_include_cache": true + }, + "accounting": { + "models_exhaustive": true, + "gaps": [], + "models": [ + { + "model": "cache-model", + "gaps": [], + "reported": { + "cached_input_tokens": 200, + "cache_write_input_tokens": 0, + "output_tokens": 5, + "output_tokens_include_reasoning": true + } + } + ] + }, + "expected": { + "input_tokens": { + "complete": null, + "observed": null + }, + "cached_input_tokens": { + "complete": 200, + "observed": 200 + }, + "cache_write_input_tokens": { + "complete": 0, + "observed": 0 + }, + "output_tokens": { + "complete": 5, + "observed": 5 + }, + "total_tokens": { + "complete": null, + "observed": null + } + } + }, + { + "name": "model_recovered_request_cache_cannot_exceed_complete_input", + "usage": {}, + "accounting": { + "models_exhaustive": true, + "gaps": [], + "models": [ + { + "model": "cache-model", + "gaps": [], + "reported": { + "input_tokens": 100, + "input_tokens_include_cache": true, + "output_tokens": 5, + "output_tokens_include_reasoning": true + }, + "requests_exhaustive": true, + "requests": [ + { + "usage": { + "cached_input_tokens": 200, + "cache_write_input_tokens": 0 + } + } + ] + } + ] + }, + "expected": { + "input_tokens": { + "complete": null, + "observed": null + }, + "cached_input_tokens": { + "complete": 200, + "observed": 200 + }, + "cache_write_input_tokens": { + "complete": 0, + "observed": 0 + }, + "output_tokens": { + "complete": 5, + "observed": 5 + }, + "total_tokens": { + "complete": null, + "observed": null + } + } + }, + { + "name": "model_recovered_request_peak_cannot_exceed_complete_input", + "usage": {}, + "accounting": { + "models_exhaustive": true, + "gaps": [], + "models": [ + { + "model": "peak-model", + "gaps": [], + "reported": { + "input_tokens": 100, + "input_tokens_include_cache": true, + "output_tokens": 5, + "output_tokens_include_reasoning": true + }, + "requests_exhaustive": true, + "requests": [ + { + "usage": { + "peak_request_input_tokens": 200 + } + } + ] + } + ] + }, + "expected": { + "input_tokens": { + "complete": null, + "observed": null + }, + "output_tokens": { + "complete": 5, + "observed": 5 + }, + "total_tokens": { + "complete": null, + "observed": null + } + } + }, + { + "name": "model_recovered_request_reasoning_cannot_exceed_complete_output", + "usage": {}, + "accounting": { + "models_exhaustive": true, + "gaps": [], + "models": [ + { + "model": "reasoning-model", + "gaps": [], + "reported": { + "input_tokens": 100, + "input_tokens_include_cache": true, + "output_tokens": 5, + "output_tokens_include_reasoning": true + }, + "requests_exhaustive": true, + "requests": [ + { + "usage": { + "reasoning_tokens": 10 + } + } + ] + } + ] + }, + "expected": { + "input_tokens": { + "complete": 100, + "observed": 100 + }, + "output_tokens": { + "complete": null, + "observed": null + }, + "total_tokens": { + "complete": null, + "observed": null + } + } } ] diff --git a/tests/python/test_model_usage.py b/tests/python/test_model_usage.py index 93d31dd63..fb018ffaa 100644 --- a/tests/python/test_model_usage.py +++ b/tests/python/test_model_usage.py @@ -10,7 +10,7 @@ from nemo_fabric.accounting import summarize_usage from nemo_fabric.errors import FabricConfigError -from nemo_fabric.types import ModelUsage, PerModelUsage, TokenCounts +from nemo_fabric.types import AppliedPricing, ModelUsage, PerModelUsage, TokenCounts def test_per_model_complete_counts_have_the_public_token_types(): @@ -472,3 +472,23 @@ def test_independently_rounded_provider_cost_is_not_a_token_reconciliation_error per_model_usage["cost"]["reported"] = {"amount_nano_usd": 9} usage["models"] = [per_model_usage] assert ModelUsage.from_mapping(usage).cost.reported.amount_nano_usd == 10 + + +@pytest.mark.parametrize("cache_rates", [{}, {"cache_read": 0}, {"cache_write": 0}]) +def test_applied_pricing_accepts_schema_optional_cache_rates(cache_rates): + parsed = AppliedPricing.from_mapping( + { + "catalog_sha256": "sha256:" + "a" * 64, + "provider": "fixture-provider", + "model_id": "fixture-model", + "pricing_as_of": "2026-10-22", + "pricing_source": "fixture-catalog", + "rates": {"input": 10, "output": 20, **cache_rates}, + } + ) + assert parsed.rates == { + "input": 10, + "output": 20, + "cache_read": cache_rates.get("cache_read"), + "cache_write": cache_rates.get("cache_write"), + } From 559ce0a558bb9faaa258b7ec17b1ac870b9d3248 Mon Sep 17 00:00:00 2001 From: mnajafian-nv Date: Thu, 22 Oct 2026 15:05:43 -0700 Subject: [PATCH 15/26] fix: preserve fractional cost precision across usage reducers Signed-off-by: mnajafian-nv --- adapter-contract/typescript/src/usage.ts | 38 ++++++++++++++++++- .../usage_accounting/reducer-cases.json | 36 ++++++++++++++++++ 2 files changed, 73 insertions(+), 1 deletion(-) diff --git a/adapter-contract/typescript/src/usage.ts b/adapter-contract/typescript/src/usage.ts index 726b5aff0..cb8f0dd9f 100644 --- a/adapter-contract/typescript/src/usage.ts +++ b/adapter-contract/typescript/src/usage.ts @@ -44,6 +44,42 @@ function validate(record: UsageObservation): void { } } +// Keep low-order cost contributions until the final round, matching math.fsum. +// Inputs have already passed the finite, nonnegative cost validation. +function sumCosts(values: readonly number[]): number { + const partials: number[] = []; + for (let high of values) { + let retained = 0; + for (let low of partials) { + if (Math.abs(high) < Math.abs(low)) [high, low] = [low, high]; + const sum = high + low; + const residual = low - (sum - high); + if (residual !== 0) partials[retained++] = residual; + high = sum; + } + partials.length = retained; + partials.push(high); + } + let index = partials.length; + let total = index === 0 ? 0 : partials[--index]!; + let residual = 0; + while (index > 0) { + const high = total; + const low = partials[--index]!; + total = high + low; + residual = low - (total - high); + if (residual !== 0) break; + } + // Preserve round-to-even when the remaining low-order terms share a sign. + if (index > 0 && ((residual < 0 && partials[index - 1]! < 0) || + (residual > 0 && partials[index - 1]! > 0))) { + const correction = residual * 2; + const rounded = total + correction; + if (correction === rounded - total) total = rounded; + } + return total; +} + function reduce(records: readonly UsageObservation[]): { reported: ReportedUsage; observed: ReportedUsage; gaps: UsageGap[] } { const reported: ReportedUsage = {}; const observed: ReportedUsage = {}; @@ -58,7 +94,7 @@ function reduce(records: readonly UsageObservation[]): { reported: ReportedUsage let value: number | undefined; if (values.length) { if (key === "cost_usd") { - const total = values.reduce((sum, item) => sum + item, 0); + const total = sumCosts(values); if (Number.isFinite(total) && total * 1e9 < 2 ** 64) value = total; } else if (key === "peak_request_input_tokens") { value = values.reduce((maximum, item) => Math.max(maximum, item)); diff --git a/tests/fixtures/usage_accounting/reducer-cases.json b/tests/fixtures/usage_accounting/reducer-cases.json index e79697c63..cba63db2e 100644 --- a/tests/fixtures/usage_accounting/reducer-cases.json +++ b/tests/fixtures/usage_accounting/reducer-cases.json @@ -437,5 +437,41 @@ "observed_input": 100, "exhaustive": true } + }, + { + "name": "reported_cost_keeps_fractional_sum_precision", + "records": [ + { + "model": "fixture", + "provider": "fixture", + "tokens_aligned": true, + "reported": { + "cost_usd": 0.1 + } + }, + { + "model": "fixture", + "provider": "fixture", + "tokens_aligned": true, + "reported": { + "cost_usd": 0.2000000005 + } + }, + { + "model": "fixture", + "provider": "fixture", + "tokens_aligned": true, + "reported": { + "cost_usd": 0.1 + } + } + ], + "expected": { + "input_tokens": null, + "cost_usd": 0.4000000005, + "observed_input": null, + "observed_output": null, + "exhaustive": true + } } ] From a6318d4dede15289bbf31485296c6bec8bc524b0 Mon Sep 17 00:00:00 2001 From: mnajafian-nv Date: Thu, 22 Oct 2026 16:24:17 -0700 Subject: [PATCH 16/26] chore: check packaged usage schema declarations Keep packaged adapter declarations in the canonical usage-schema synchronization and drift check. Prevent the check from passing with a stale packaged schema. Signed-off-by: mnajafian-nv --- scripts/contracts/generate_usage_accounting.py | 1 + 1 file changed, 1 insertion(+) diff --git a/scripts/contracts/generate_usage_accounting.py b/scripts/contracts/generate_usage_accounting.py index 49bb52bef..3b5cbd19c 100644 --- a/scripts/contracts/generate_usage_accounting.py +++ b/scripts/contracts/generate_usage_accounting.py @@ -22,6 +22,7 @@ def synchronize(*, check: bool) -> None: descriptors = [ *ROOT.glob("adapters/python/*/*.fabric-adapter.json"), *ROOT.glob("adapters/typescript/*/*.fabric-adapter.json"), + *ROOT.glob("crates/fabric-cli/assets/adapters/*/*.fabric-adapter.json"), ] for path in descriptors: original = path.read_text() From 630820f93384bf5ce9729edd1e2c9c18d63f1999 Mon Sep 17 00:00:00 2001 From: mnajafian-nv Date: Thu, 22 Oct 2026 16:44:48 -0700 Subject: [PATCH 17/26] fix: return detached applied pricing rates in Python Return the documented JSON-compatible rate snapshot instead of exposing a frozen internal mapping. Preserve optional cache rates and snapshot isolation. Signed-off-by: mnajafian-nv --- sdk/python/nemo-fabric-runtime/src/nemo_fabric/types.py | 1 + tests/python/test_model_usage.py | 6 +++++- 2 files changed, 6 insertions(+), 1 deletion(-) diff --git a/sdk/python/nemo-fabric-runtime/src/nemo_fabric/types.py b/sdk/python/nemo-fabric-runtime/src/nemo_fabric/types.py index bea167f4a..f5a52cc4c 100644 --- a/sdk/python/nemo-fabric-runtime/src/nemo_fabric/types.py +++ b/sdk/python/nemo-fabric-runtime/src/nemo_fabric/types.py @@ -2577,6 +2577,7 @@ class AppliedPricing(FabricMapping): "rates", } ) + _json_fields = frozenset({"rates"}) @classmethod def _normalize(cls, data: dict[str, Any]) -> dict[str, Any]: diff --git a/tests/python/test_model_usage.py b/tests/python/test_model_usage.py index fb018ffaa..805263033 100644 --- a/tests/python/test_model_usage.py +++ b/tests/python/test_model_usage.py @@ -486,9 +486,13 @@ def test_applied_pricing_accepts_schema_optional_cache_rates(cache_rates): "rates": {"input": 10, "output": 20, **cache_rates}, } ) - assert parsed.rates == { + expected = { "input": 10, "output": 20, "cache_read": cache_rates.get("cache_read"), "cache_write": cache_rates.get("cache_write"), } + rates = parsed.rates + assert json.loads(json.dumps(rates)) == expected + rates["input"] = 999 + assert parsed.rates == expected From 93b00ecdbf82a618e8b6947692b9ea333f59f38b Mon Sep 17 00:00:00 2001 From: mnajafian-nv Date: Thu, 22 Oct 2026 17:23:32 -0700 Subject: [PATCH 18/26] fix: expose omitted optional SDK usage attributes Signed-off-by: mnajafian-nv --- .../nemo_fabric.types.md | 16 +++++++- .../src/nemo_fabric/types.py | 12 +++++- tests/python/test_model_usage.py | 39 +++++++++++++++++++ 3 files changed, 63 insertions(+), 4 deletions(-) diff --git a/docs/reference/api/python-library-reference/nemo_fabric.types.md b/docs/reference/api/python-library-reference/nemo_fabric.types.md index 83cdfab2b..b930369db 100644 --- a/docs/reference/api/python-library-reference/nemo_fabric.types.md +++ b/docs/reference/api/python-library-reference/nemo_fabric.types.md @@ -2067,7 +2067,6 @@ The mapping exposes the following typed fields: | `normalization_version` | `str` | | `assumption_policy` | `str` | | `bundled_catalog_sha256` | `str` | -| `user_catalog_sha256` | `str \| None` | ### method `__init__` @@ -2088,6 +2087,13 @@ Return an immutable view of preserved extension fields. +--- + +### property user_catalog_sha256 + +Return the user catalog fingerprint when supplied, otherwise None. + + --- ### classmethod `from_mapping` @@ -2475,7 +2481,6 @@ The mapping exposes the following typed fields: | Field | Type | | --- | --- | | `model` | `str` | -| `provider` | `str \| None` | | `input_tokens` | `int \| None` | | `cached_input_tokens` | `int \| None` | | `cache_write_input_tokens` | `int \| None` | @@ -2504,6 +2509,13 @@ Return an immutable view of preserved extension fields. +--- + +### property provider + +Return the accounting provider when reported, otherwise None. + + --- ### classmethod `from_mapping` diff --git a/sdk/python/nemo-fabric-runtime/src/nemo_fabric/types.py b/sdk/python/nemo-fabric-runtime/src/nemo_fabric/types.py index f5a52cc4c..2dfcd259b 100644 --- a/sdk/python/nemo-fabric-runtime/src/nemo_fabric/types.py +++ b/sdk/python/nemo-fabric-runtime/src/nemo_fabric/types.py @@ -2527,7 +2527,6 @@ class PricingBasis(FabricMapping): normalization_version: str assumption_policy: str bundled_catalog_sha256: str - user_catalog_sha256: str | None _fields = frozenset( { "basis", @@ -2539,6 +2538,11 @@ class PricingBasis(FabricMapping): } ) + @property + def user_catalog_sha256(self) -> str | None: + """Return the user catalog fingerprint when supplied, otherwise None.""" + return self._data.get("user_catalog_sha256") + @classmethod def _normalize(cls, data: dict[str, Any]) -> dict[str, Any]: if set(data).difference(cls._fields) or any( @@ -2796,7 +2800,6 @@ class PerModelUsage(FabricMapping): """ model: str - provider: str | None input_tokens: int | None cached_input_tokens: int | None cache_write_input_tokens: int | None @@ -2809,6 +2812,11 @@ class PerModelUsage(FabricMapping): {*_MODEL_USAGE_TOKEN_FIELDS, "model", "provider", "observed", "gaps", "cost"} ) + @property + def provider(self) -> str | None: + """Return the accounting provider when reported, otherwise None.""" + return self._data.get("provider") + @classmethod def _normalize(cls, data: dict[str, Any]) -> dict[str, Any]: data["model"] = _required_text(data.get("model"), "model") diff --git a/tests/python/test_model_usage.py b/tests/python/test_model_usage.py index 805263033..4c731b546 100644 --- a/tests/python/test_model_usage.py +++ b/tests/python/test_model_usage.py @@ -496,3 +496,42 @@ def test_applied_pricing_accepts_schema_optional_cache_rates(cache_rates): assert json.loads(json.dumps(rates)) == expected rates["input"] = 999 assert parsed.rates == expected + + +@pytest.mark.parametrize( + "provider", [{}, {"provider": None}, {"provider": "fixture-provider"}] +) +def test_per_model_optional_provider_attribute_preserves_source_presence( + per_model_usage, provider +): + from nemo_fabric.types import PerModelUsage + + per_model_usage.update(provider) + parsed = PerModelUsage.from_mapping(per_model_usage) + assert parsed.provider == provider.get("provider") + serialized = parsed.to_mapping() + assert ("provider" in serialized) == ("provider" in provider) + if "provider" in provider: + assert serialized["provider"] == provider["provider"] + + +@pytest.mark.parametrize( + "user_catalog", + [{}, {"user_catalog_sha256": None}, {"user_catalog_sha256": "sha256:" + "b" * 64}], +) +def test_pricing_basis_optional_user_catalog_attribute_preserves_source_presence( + user_catalog, +): + from nemo_fabric.types import PricingBasis + + source = { + "basis": "catalog", + "currency": "USD", + "normalization_version": "cache_reasoning_inclusive_v1", + "assumption_policy": "applicable_tariff_no_missing_counts_v1", + "bundled_catalog_sha256": "sha256:" + "a" * 64, + **user_catalog, + } + parsed = PricingBasis.from_mapping(source) + assert parsed.user_catalog_sha256 == user_catalog.get("user_catalog_sha256") + assert parsed.to_mapping() == source From 101bce6206a4ce6663e1fc876f99c35404de6fe8 Mon Sep 17 00:00:00 2001 From: mnajafian-nv Date: Thu, 22 Oct 2026 20:53:44 -0700 Subject: [PATCH 19/26] fix: reject Relay exit during successful health probe Signed-off-by: mnajafian-nv --- adapters/python/common/README.md | 6 ++++-- .../nemo_fabric_adapters/common/relay_gateway.py | 5 +++++ .../adapters/test_adapters_common_relay_gateway.py | 13 +++++++++++++ 3 files changed, 22 insertions(+), 2 deletions(-) diff --git a/adapters/python/common/README.md b/adapters/python/common/README.md index 919bdd233..df9e8aa93 100644 --- a/adapters/python/common/README.md +++ b/adapters/python/common/README.md @@ -126,8 +126,10 @@ runtime. The Claude and Codex adapters wait up to 30 seconds for the supervised Relay gateway's HTTP health endpoint. This budget includes configuration validation -and required worker initialization. Early process exit fails startup immediately; -a gateway that never becomes healthy fails when the budget expires and is stopped. +and required worker initialization. Process exit is checked before each health +probe and again before accepting a successful response. An exit during an +unsuccessful probe is detected at the next poll. A gateway that never becomes +healthy fails when the budget expires and is stopped. The adapters retain its log for diagnosis. Local Python hosts separately enforce a 90-second adapter startup limit. diff --git a/adapters/python/common/src/nemo_fabric_adapters/common/relay_gateway.py b/adapters/python/common/src/nemo_fabric_adapters/common/relay_gateway.py index 959f54fa7..1a57c5fe4 100644 --- a/adapters/python/common/src/nemo_fabric_adapters/common/relay_gateway.py +++ b/adapters/python/common/src/nemo_fabric_adapters/common/relay_gateway.py @@ -132,6 +132,11 @@ def wait_for_relay_gateway( health_url, timeout=min(1.0, remaining) ) as response: if 200 <= response.status < 300: + returncode = process.poll() + if returncode is not None: + raise RelayGatewayError( + f"NeMo Relay gateway exited with status {returncode} before becoming ready" + ) return except (OSError, urllib.error.URLError): pass diff --git a/tests/adapters/test_adapters_common_relay_gateway.py b/tests/adapters/test_adapters_common_relay_gateway.py index e478f85de..fbd9ca893 100644 --- a/tests/adapters/test_adapters_common_relay_gateway.py +++ b/tests/adapters/test_adapters_common_relay_gateway.py @@ -239,6 +239,19 @@ def test_wait_for_relay_gateway_reports_early_exit(): ) +def test_wait_for_relay_gateway_rejects_exit_during_successful_probe(monkeypatch): + process = MagicMock() + process.poll.side_effect = [None, 17] + response = MagicMock() + response.__enter__.return_value.status = 200 + monkeypatch.setattr( + relay_gateway.urllib.request, "urlopen", MagicMock(return_value=response) + ) + + with pytest.raises(relay_gateway.RelayGatewayError, match="status 17"): + relay_gateway.wait_for_relay_gateway(process, "http://127.0.0.1:43210/healthz") + + def test_wait_for_relay_gateway_times_out(): process = MagicMock() process.poll.return_value = None From 2014cbe68af1259c837ab3da256883ff3dcd4a74 Mon Sep 17 00:00:00 2001 From: mnajafian-nv Date: Thu, 22 Oct 2026 21:28:53 -0700 Subject: [PATCH 20/26] fix: exclude unknown keys from request usage Signed-off-by: mnajafian-nv --- adapter-contract/typescript/src/usage.ts | 7 ++++++- adapter-contract/typescript/test/usage-reducer.test.js | 8 ++++++++ adapters/python/deepagents/README.md | 7 +++++++ docs/sdk/python.mdx | 2 +- 4 files changed, 22 insertions(+), 2 deletions(-) diff --git a/adapter-contract/typescript/src/usage.ts b/adapter-contract/typescript/src/usage.ts index cb8f0dd9f..16c8cc5ae 100644 --- a/adapter-contract/typescript/src/usage.ts +++ b/adapter-contract/typescript/src/usage.ts @@ -144,7 +144,12 @@ function aligned(records: readonly UsageObservation[]): boolean { function requestUsage(record: UsageObservation): ReportedUsage { // Retain uncertain literal counts in model observations, not complete calls. - const fields = { ...record.reported }; + const fields: ReportedUsage = {}; + const keys = [...dimensions, "input_tokens_include_cache", "output_tokens_include_reasoning"] as const; + for (const key of keys) { + const value = record.reported[key]; + if (value !== undefined) (fields as Record)[key] = value; + } for (const gap of record.gaps ?? []) for (const dimension of gap.dimensions) delete fields[dimension]; return fields; } diff --git a/adapter-contract/typescript/test/usage-reducer.test.js b/adapter-contract/typescript/test/usage-reducer.test.js index e83f1ad8b..de7fbdfb2 100644 --- a/adapter-contract/typescript/test/usage-reducer.test.js +++ b/adapter-contract/typescript/test/usage-reducer.test.js @@ -46,3 +46,11 @@ test("scoped request gaps retain independent fields without certifying cache", ( assert.equal(model.requests[0].gap, undefined); assert.equal(model.requests[0].tokens_aligned, true); }); + +test("request usage omits provider-specific keys outside the wire schema", () => { + const usage = usageFromObservations([{ model: "m", is_model_call: true, + reported: { input_tokens: 7, input_tokens_include_cache: true, prompt_tokens: 7 }, + }]); + const request = usage.extensions[USAGE_ACCOUNTING_KEY].models[0].requests[0].usage; + assert.deepEqual(request, { input_tokens: 7, input_tokens_include_cache: true }); +}); diff --git a/adapters/python/deepagents/README.md b/adapters/python/deepagents/README.md index fdfb09f2e..af2f9fc34 100644 --- a/adapters/python/deepagents/README.md +++ b/adapters/python/deepagents/README.md @@ -241,6 +241,13 @@ The descriptor declares `nemo_fabric.usage_accounting.v1` for complete dimension observations and gaps, and `nemo_fabric.detailed_token_usage.v1` for raw cache and reasoning conventions. An omitted count stays unknown; a reported zero stays zero. +**Upgrade note:** For providers whose usage is visible only in LangGraph +messages, `RunResult.usage.input_tokens` and `output_tokens` can now be `null`. +The previous graph-message subtotal remains in `RunResult.output.usage`, labeled +`usage_source: graph_state_messages`. It is a diagnostic subtotal, not a +complete invocation count; consumers that previously read the top-level fields +should handle unknown counts rather than substituting this subtotal. + For adapter-created OpenAI-compatible clients, the collector observes sync and async calls and their HTTP attempts without changing requests or retries. Normal reported calls stay known. An actual failed or unfinished attempt without usable diff --git a/docs/sdk/python.mdx b/docs/sdk/python.mdx index f55fe8ba9..d47a7bdb1 100644 --- a/docs/sdk/python.mdx +++ b/docs/sdk/python.mdx @@ -1229,7 +1229,7 @@ accounting remains available in `result.usage` and is not certified complete by `summarize_usage`. `cost.reported` contains only provider-reported amounts. Complete and observed -amounts use integer nano-USD. Catalog estimates remain unavailable in this slice; +amounts use integer nano-USD. Catalog estimates are not available; `cost.estimated` has no amounts and no typed pricing setting is exposed. The stable result shape retains optional pricing-basis fields for persisted evidence. From 639df488f5cdbf1b2226b56b624efda6380de2b6 Mon Sep 17 00:00:00 2001 From: mnajafian-nv Date: Thu, 22 Oct 2026 23:34:28 -0700 Subject: [PATCH 21/26] docs: name NeMo Fabric in Droid usage Signed-off-by: mnajafian-nv --- adapters/typescript/droid/README.md | 2 +- docs/integrations/harness/droid.mdx | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/adapters/typescript/droid/README.md b/adapters/typescript/droid/README.md index f9624321a..d5e59c0f4 100644 --- a/adapters/typescript/droid/README.md +++ b/adapters/typescript/droid/README.md @@ -125,7 +125,7 @@ retained in `nemo_fabric.usage_accounting.v1` with a `scope_unverified` gap. Droid's debounced usage notifications have no stable model-call identity and cannot establish exhaustive per-call reporting. -The adapter preserves the persisted provider input convention for Fabric core +The adapter preserves the persisted provider input convention for NeMo Fabric core to normalize. Anthropic and Bedrock input excludes cache reads and writes; Gemini input includes cache reads. Responses subtracts reads without subtracting nonstandard cache writes. Whether input includes those writes is not established, so the shared binary diff --git a/docs/integrations/harness/droid.mdx b/docs/integrations/harness/droid.mdx index 29ef349b0..669d3f703 100644 --- a/docs/integrations/harness/droid.mdx +++ b/docs/integrations/harness/droid.mdx @@ -224,7 +224,7 @@ retained in `nemo_fabric.usage_accounting.v1` with a `scope_unverified` gap. Droid's debounced usage notifications have no stable model-call identity and cannot establish exhaustive per-call reporting. -The adapter preserves the persisted provider input convention for Fabric core +The adapter preserves the persisted provider input convention for NeMo Fabric core to normalize. Anthropic and Bedrock input excludes cache reads and writes; Gemini input includes cache reads. Responses subtracts reads without subtracting nonstandard cache writes. Whether input includes those writes is not established, so the shared binary From e3a2355b585e1eb231a7395cb1f4e7945c76d12b Mon Sep 17 00:00:00 2001 From: mnajafian-nv Date: Fri, 23 Oct 2026 01:41:33 -0700 Subject: [PATCH 22/26] fix: classify empty usage coverage as unverified Signed-off-by: mnajafian-nv --- .../python/src/nemo_fabric_adapter_contract/usage.py | 2 +- adapter-contract/typescript/src/usage.ts | 2 +- adapter-contract/typescript/test/usage-reducer.test.js | 4 ++++ tests/fixtures/usage_accounting/reducer-cases.json | 10 ++++++++++ tests/python/test_usage_reducer_parity.py | 3 +++ 5 files changed, 19 insertions(+), 2 deletions(-) diff --git a/adapter-contract/python/src/nemo_fabric_adapter_contract/usage.py b/adapter-contract/python/src/nemo_fabric_adapter_contract/usage.py index d28998d36..f296d5d21 100644 --- a/adapter-contract/python/src/nemo_fabric_adapter_contract/usage.py +++ b/adapter-contract/python/src/nemo_fabric_adapter_contract/usage.py @@ -79,7 +79,7 @@ def _reduce( if total is not None and not causes: complete[key] = total elif not causes: - causes.add(GapReason.MISSING_USAGE) + causes.add(GapReason.SCOPE_UNVERIFIED if not records else GapReason.MISSING_USAGE) for cause in sorted(causes): gaps[cause].append(dimension) for flag, counter in ( diff --git a/adapter-contract/typescript/src/usage.ts b/adapter-contract/typescript/src/usage.ts index 16c8cc5ae..bf75221c4 100644 --- a/adapter-contract/typescript/src/usage.ts +++ b/adapter-contract/typescript/src/usage.ts @@ -111,7 +111,7 @@ function reduce(records: readonly UsageObservation[]): { reported: ReportedUsage if (values.length !== records.length) causes.add("missing_usage"); if (values.length && value === undefined) causes.add("scope_unverified"); if (value !== undefined && !causes.size) reported[key] = value; - else if (!causes.size) causes.add("missing_usage"); + else if (!causes.size) causes.add(records.length ? "missing_usage" : "scope_unverified"); for (const reason of causes) gap(reason, key); } for (const [flag, counter] of [["input_tokens_include_cache", "input_tokens"], ["output_tokens_include_reasoning", "output_tokens"]] as const) { diff --git a/adapter-contract/typescript/test/usage-reducer.test.js b/adapter-contract/typescript/test/usage-reducer.test.js index de7fbdfb2..4c6e7040c 100644 --- a/adapter-contract/typescript/test/usage-reducer.test.js +++ b/adapter-contract/typescript/test/usage-reducer.test.js @@ -12,6 +12,10 @@ for (const vector of cases) { const usage = usageFromObservations(vector.records); const accounting = usage.extensions[USAGE_ACCOUNTING_KEY]; for (const [field, value] of Object.entries(vector.expected)) { + if (field === "gap_reason") { + assert.deepEqual(accounting.gaps.map(gap => gap.reason), [value]); + continue; + } const actual = field === "exhaustive" ? accounting.models_exhaustive : field.startsWith("observed_") ? accounting.observed[field.slice(9) + "_tokens"] : field === "peak_request_input_tokens" ? usage.extensions["nemo_fabric.detailed_token_usage.v1"][field] diff --git a/tests/fixtures/usage_accounting/reducer-cases.json b/tests/fixtures/usage_accounting/reducer-cases.json index cba63db2e..39b4c4959 100644 --- a/tests/fixtures/usage_accounting/reducer-cases.json +++ b/tests/fixtures/usage_accounting/reducer-cases.json @@ -1,4 +1,14 @@ [ + { + "name": "empty_observations", + "records": [], + "expected": { + "input_tokens": null, + "output_tokens": null, + "exhaustive": false, + "gap_reason": "scope_unverified" + } + }, { "name": "complete_calls", "records": [ diff --git a/tests/python/test_usage_reducer_parity.py b/tests/python/test_usage_reducer_parity.py index c305b4b33..ccdf1628c 100644 --- a/tests/python/test_usage_reducer_parity.py +++ b/tests/python/test_usage_reducer_parity.py @@ -51,6 +51,9 @@ def test_python_reducer_matches_shared_independent_expectations(reduced): for vector, usage in zip(VECTORS, reduced, strict=True): accounting = usage["extensions"][KEY] for field, value in vector["expected"].items(): + if field == "gap_reason": + assert [gap["reason"] for gap in accounting["gaps"]] == [value] + continue if field == "exhaustive": actual = accounting["models_exhaustive"] elif field.startswith("observed_"): From a8f6fd807b9a2dc5cd3ae3f1b01538aec4820a7b Mon Sep 17 00:00:00 2001 From: mnajafian-nv Date: Fri, 23 Oct 2026 02:05:52 -0700 Subject: [PATCH 23/26] test: account for unverified empty usage coverage Signed-off-by: mnajafian-nv --- tests/adapters/test_deepagents_accounting.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tests/adapters/test_deepagents_accounting.py b/tests/adapters/test_deepagents_accounting.py index dccaf29a3..4b674815a 100644 --- a/tests/adapters/test_deepagents_accounting.py +++ b/tests/adapters/test_deepagents_accounting.py @@ -210,7 +210,8 @@ async def test_remote_tool_callback_dispatch_marks_new_work_only_unit( dispatched = operation in {"start_async_task", "update_async_task"} assert result.input_tokens == (100 if foreground_call and not dispatched else None) reasons = {gap["reason"] for gap in result.extensions[KEY]["gaps"]} - assert ("scope_unverified" in reasons) is dispatched + # With no foreground call, an empty callback list cannot prove zero usage. + assert ("scope_unverified" in reasons) is (dispatched or not foreground_call) if foreground_call: assert result.extensions[KEY]["observed"]["input_tokens"] == 100 elif dispatched: From 83d8adace3ca3dd886c8769ba1cdc1a01d3a7c3b Mon Sep 17 00:00:00 2001 From: mnajafian-nv Date: Fri, 23 Oct 2026 08:55:34 -0700 Subject: [PATCH 24/26] docs: keep adapter usage guidance within current contract Signed-off-by: mnajafian-nv --- skills/nemo-fabric-build-adapter/SKILL.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/skills/nemo-fabric-build-adapter/SKILL.md b/skills/nemo-fabric-build-adapter/SKILL.md index 064429506..7b014b3a7 100644 --- a/skills/nemo-fabric-build-adapter/SKILL.md +++ b/skills/nemo-fabric-build-adapter/SKILL.md @@ -231,7 +231,7 @@ never share mutable continuation state. Normalize usage per invocation in `AgentUsage`. If the target reports cumulative session totals, retain a session-local baseline and difference successive observed counters rather than summing cumulative snapshots or using only the last model response. Missing or reset counters remain unknown. Keep the closed `fabric.adapter/v1alpha2` usage fields unchanged. Put cache-write, reasoning, peak request, and per-model detail in `usage.extensions["nemo_fabric.detailed_token_usage.v1"]`, and declare that object with `extension_schemas.usage` in the adapter descriptor. The detailed object repeats the normalized token counters so its inclusion flags are self-contained; `input_tokens_include_cache` covers both cached reads and cache writes there. NeMo Fabric preserves the object as `RunResult.usage.metadata["nemo_fabric.detailed_token_usage.v1"]`. Consumers treat any detailed field for an input or output dimension as selecting the detailed counters for that dimension without falling back to its top-level counters. Report the primary counter, inclusion flag, and every excluded component together so consumers can calculate an inclusive total. Preserve available usage on unsuccessful terminal results. Do not infer missing cost or promote estimates to `cost_usd`. -Preserve field presence: an omitted count stays unknown and a provider-reported zero stays zero. If the target fills missing fields with zero, recover original report presence before certifying a complete count. Capture each call's invocation owner when the call starts, including descendant work. Freeze its usage snapshot when `invoke` returns; late reports cannot amend that result or transfer into a later invocation. A failed execution alone is not an accounting gap: keep settled usage, and distinguish valid observations from complete totals when owned work is missing or unfinished. Use the shared Python `usage_from_observations` or TypeScript `usageFromObservations` reducer with the declared `nemo_fabric.usage_accounting.v1` extension for this evidence. Keep source-specific parsing and lifecycle ownership in the adapter; let core normalize counts and price them. See the [results contract](https://github.com/NVIDIA/NeMo-Fabric/blob/main/docs/adapter-contract/tutorials/results.md) for the extension and compatibility boundary. +Preserve field presence: an omitted count stays unknown and a provider-reported zero stays zero. If the target fills missing fields with zero, recover original report presence before certifying a complete count. Capture each call's invocation owner when the call starts, including descendant work. Freeze its usage snapshot when `invoke` returns; late reports cannot amend that result or transfer into a later invocation. A failed execution alone is not an accounting gap: keep settled usage, and distinguish valid observations from complete totals when owned work is missing or unfinished. Use the shared Python `usage_from_observations` or TypeScript `usageFromObservations` reducer with the declared `nemo_fabric.usage_accounting.v1` extension for this evidence. Keep source-specific parsing and lifecycle ownership in the adapter; let core normalize counts. See the [results contract](https://github.com/NVIDIA/NeMo-Fabric/blob/main/docs/adapter-contract/tutorials/results.md) for the extension and compatibility boundary. Test observable continuation rather than merely calling `invoke` twice: make the second result depend on the first turn without caller-side replay, then From ebcfe97c48bcc90e93f2ae74078f1b82e9049132 Mon Sep 17 00:00:00 2001 From: mnajafian-nv Date: Fri, 23 Oct 2026 10:26:13 -0700 Subject: [PATCH 25/26] test: run OpenShell launcher checks with system Bash Use the system Bash path so macOS tests exercise Bash 3.2 even when PATH contains a newer Bash. This keeps the empty-array regression meaningful in CI. Signed-off-by: mnajafian-nv --- tests/scripts/test_openshell_launcher.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/scripts/test_openshell_launcher.py b/tests/scripts/test_openshell_launcher.py index 15ab68caa..f54157496 100644 --- a/tests/scripts/test_openshell_launcher.py +++ b/tests/scripts/test_openshell_launcher.py @@ -165,7 +165,7 @@ def _operations(root: Path) -> list[dict]: def _run(launcher, workload, env): script, _, base_env = launcher return subprocess.run( - ["bash", str(script), str(workload), "--input", "words with spaces"], + ["/bin/bash", str(script), str(workload), "--input", "words with spaces"], env={**base_env, **env}, text=True, capture_output=True, @@ -385,7 +385,7 @@ def test_invalid_platform_or_mode_stops_before_building(launcher, environment, m "stage_workload_image() { :; }\nrun_workload_consumer() { :; }\n" ) completed = subprocess.run( - ["bash", str(script), str(workload)], + ["/bin/bash", str(script), str(workload)], env={**base_env, **environment}, text=True, capture_output=True, From 83b5e76b8b5f3629e7ef9215bd5be505e3a13298 Mon Sep 17 00:00:00 2001 From: mnajafian-nv Date: Fri, 23 Oct 2026 15:07:21 -0700 Subject: [PATCH 26/26] fix: reject Relay readiness after the startup deadline Signed-off-by: mnajafian-nv --- .../common/relay_gateway.py | 2 ++ .../test_adapters_common_relay_gateway.py | 20 +++++++++++++++++++ 2 files changed, 22 insertions(+) diff --git a/adapters/python/common/src/nemo_fabric_adapters/common/relay_gateway.py b/adapters/python/common/src/nemo_fabric_adapters/common/relay_gateway.py index 1a57c5fe4..f0e6d5cd3 100644 --- a/adapters/python/common/src/nemo_fabric_adapters/common/relay_gateway.py +++ b/adapters/python/common/src/nemo_fabric_adapters/common/relay_gateway.py @@ -137,6 +137,8 @@ def wait_for_relay_gateway( raise RelayGatewayError( f"NeMo Relay gateway exited with status {returncode} before becoming ready" ) + if time.monotonic() >= deadline: + break return except (OSError, urllib.error.URLError): pass diff --git a/tests/adapters/test_adapters_common_relay_gateway.py b/tests/adapters/test_adapters_common_relay_gateway.py index fbd9ca893..86c23ede2 100644 --- a/tests/adapters/test_adapters_common_relay_gateway.py +++ b/tests/adapters/test_adapters_common_relay_gateway.py @@ -252,6 +252,26 @@ def test_wait_for_relay_gateway_rejects_exit_during_successful_probe(monkeypatch relay_gateway.wait_for_relay_gateway(process, "http://127.0.0.1:43210/healthz") +def test_wait_for_relay_gateway_rejects_success_after_deadline( + monkeypatch, gateway_clock +): + process = MagicMock() + process.poll.return_value = None + response = MagicMock() + response.__enter__.return_value.status = 200 + + def probe_health(url, *, timeout): + gateway_clock.monotonic.return_value += 0.26 + return response + + monkeypatch.setattr(relay_gateway.urllib.request, "urlopen", probe_health) + + with pytest.raises(relay_gateway.RelayGatewayError, match="did not become ready"): + relay_gateway.wait_for_relay_gateway( + process, "http://127.0.0.1:43210/healthz", timeout=0.25 + ) + + def test_wait_for_relay_gateway_times_out(): process = MagicMock() process.poll.return_value = None