From 9d6bc218bf4debd918ea49ad44eb8e7789fba6e4 Mon Sep 17 00:00:00 2001 From: mnajafian-nv Date: Mon, 19 Oct 2026 15:28:16 -0700 Subject: [PATCH 01/29] 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/29] 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 78f4b88236a065563756c72ebbfb673142e3dbac Mon Sep 17 00:00:00 2001 From: mnajafian-nv Date: Tue, 20 Oct 2026 00:39:41 -0700 Subject: [PATCH 03/29] fix: sum Kilo token usage across a prompt's model calls Kilo Code records each step of a prompt as its own assistant message, and the prompt response returns only the last one. The adapter took usage from that message, so a prompt that used tools reported only its final model call. It also paired Kilo's total, which includes cache and reasoning tokens, with input and output counts that exclude them. Find the prompt's user message after the newest pre-prompt message and attribute each assistant message through Kilo's parentID to the prompt that started its turn. Compaction messages inside the turn belong to it; when a background turn that has not produced a message yet could be the one compacting, ownership is ambiguous and usage is unknown. Sum the step-finish counts of each owned call, add cache reads and writes back into input and reasoning back into output, mark input as including cache, and report cache reads separately. Derive the total from input plus output when both are known. Kilo's own total is the AI SDK's input plus output with an omitted side counted as 0, so it would turn a missing output into a total that looks complete. Add the calls of foreground task subagents from their child sessions, inside each task call's time window, one level deep. Two task calls on one session that start at the same time leave usage unknown, because their windows cannot be told apart. Kilo 7.7.12 enables background subagents by default, and the subagent and Kilo's reply to its result can run after the prompt returns. A prompt that starts one reports no usage, and the reply, which Kilo prompts with a synthetic background result, is never counted toward a later prompt. Kilo zero-fills every component counter a provider omits, so a call's input or output that adds up to 0 is unknown rather than measured, and a finished call with neither input nor output makes the usage unknown. Kilo zero-fills missing cache details, so a zero cache-read count leaves the cached input unknown while input, output and total stay known; a cache write does not show that reads were reported. A sum that leaves the safe integer range is unknown too, because the result schema accepts only safe integers. Create the session with an explicit title. Kilo generates a title for a session that keeps its default one with a separate model call that no session message records. Stop reporting cost_usd. Kilo computes most message costs from its model catalog and does not mark the ones a provider reported, and the contract keeps that field for reported cost only. Signed-off-by: mnajafian-nv --- adapters/typescript/kilo/README.md | 37 +- adapters/typescript/kilo/src/kilo-sdk.ts | 342 +++++++++++++++- .../typescript/kilo/test/kilo-sdk.test.mjs | 375 +++++++++++++++++- docs/integrations/harness/kilo.mdx | 36 +- 4 files changed, 773 insertions(+), 17 deletions(-) diff --git a/adapters/typescript/kilo/README.md b/adapters/typescript/kilo/README.md index cb6e313dc..96bb37661 100644 --- a/adapters/typescript/kilo/README.md +++ b/adapters/typescript/kilo/README.md @@ -74,7 +74,42 @@ authentication and per-server tool filters are unsupported. ### Runtime Behavior The adapter accepts plain-text input and returns the terminal assistant text in -`output.response`. Ordered invocations reuse one warm Kilo Code session and its +`output.response`. + +Usage sums the model calls that Kilo Code recorded for the invocation's prompt, +including in-turn compaction and the calls of foreground `task` subagents in +their child sessions. Kilo Code 7.7.12 can run a subagent in the background when +the model asks for it. An invocation that starts a background subagent reports +no usage, because the subagent and Kilo Code's reply to its result can run after +the invocation returns. That reply is never counted toward a later invocation. +The adapter sets the session title so that Kilo Code does not make an unrecorded +model call to generate one. + +`input_tokens` includes cache reads and cache writes, `cached_input_tokens` +holds the cache reads, and `output_tokens` includes reasoning tokens. +`total_tokens` is `input_tokens` plus `output_tokens`, so it is unknown when +either of them is unknown. The adapter does not use Kilo Code's own total, +because Kilo Code 7.7.12 computes it as input plus output with an omitted side +counted as 0. + +Kilo Code 7.7.12 records an omitted input, output, reasoning, or cache count as +0 before the adapter sees it. The adapter therefore treats a call's input or +output that adds up to 0 as unknown. Kilo zero-fills missing cache details, so a +zero cache-read count remains unknown unless its presence can be established for +that call. When any model call leaves a field unknown, or the field's sum is +larger than JavaScript can represent exactly, the result omits that field. The +result has no usage when a finished model call reported neither input nor +output, or when an assistant message of the invocation never finished. It also +has no usage when a subagent's child session cannot be read or starts its own +subagent, or when a new model call cannot be traced to a prompt. Kilo Code does +not record an attempt that it retries before the call finishes, so those tokens +are not counted. + +Kilo Code computes most cost figures from its model catalog's rates, and its +messages do not say which costs a provider reported, so the adapter does not +report them as `cost_usd`. + +Ordered invocations reuse one warm Kilo Code session and its conversation history. A prompt has a 30-minute deadline; on timeout the adapter asks Kilo Code to abort the prompt and reports `kilo_prompt_timeout`. diff --git a/adapters/typescript/kilo/src/kilo-sdk.ts b/adapters/typescript/kilo/src/kilo-sdk.ts index e949f0346..896657cbb 100644 --- a/adapters/typescript/kilo/src/kilo-sdk.ts +++ b/adapters/typescript/kilo/src/kilo-sdk.ts @@ -26,6 +26,7 @@ interface KiloClient { }; session: { create(parameters: Record, options: KiloRequestOptions): Promise<{data: unknown}>; + messages(parameters: Record, options: KiloRequestOptions): Promise<{data: unknown}>; prompt(parameters: Record, options: KiloRequestOptions): Promise<{data: unknown}>; abort(parameters: Record, options: KiloRequestOptions): Promise; delete(parameters: Record, options: KiloRequestOptions): Promise; @@ -40,6 +41,11 @@ type PortFactory = () => Promise; const START_TIMEOUT_MS = 30_000; const STOP_TIMEOUT_MS = 5_000; const PROMPT_TIMEOUT_MS = 30 * 60_000; +const LIST_TIMEOUT_MS = 30_000; +// Kilo Code generates a title with a separate model call for a session that +// still has its default title, and that call never appears in the session's +// messages. An explicit title skips it. +const SESSION_TITLE = "NeMo Fabric"; const ENVIRONMENT_VARIABLE = /\$\{([A-Za-z_][A-Za-z0-9_]*)\}/gu; const HTTP_HEADER_NAME = /^[!#$%&'*+\-.^_`|~0-9A-Za-z]+$/u; const INHERITED_ENVIRONMENT_NAMES = new Set([ @@ -322,20 +328,304 @@ export async function startKiloServer(workspace: string, environment: NodeJS.Pro }); } +/** The terminal text and error state of the last assistant message that a Kilo Code prompt returns. */ export function extractPromptOutcome(data: unknown): KiloPromptOutcome { if (!isRecord(data)) return {}; const info = isRecord(data.info) ? data.info : {}; - const tokens = isRecord(info.tokens) ? info.tokens : {}; - const usage: AgentUsage = {}; - const inputTokens = typeof tokens.input === "number" ? tokens.input : undefined; - const outputTokens = typeof tokens.output === "number" ? tokens.output : undefined; - if (inputTokens !== undefined) usage.input_tokens = inputTokens; - if (outputTokens !== undefined) usage.output_tokens = outputTokens; - if (typeof tokens.total === "number") usage.total_tokens = tokens.total; - else if (inputTokens !== undefined && outputTokens !== undefined) usage.total_tokens = inputTokens + outputTokens; - if (typeof info.cost === "number") usage.cost_usd = info.cost; const text = Array.isArray(data.parts) ? data.parts.filter(isRecord).filter((part) => part.type === "text" && typeof part.text === "string").map((part) => part.text).join("") : ""; - return { ...(text.length === 0 ? {} : {text}), ...(info.error === undefined ? {} : {error: true}), ...(Object.keys(usage).length === 0 ? {} : {usage}) }; + return { ...(text.length === 0 ? {} : {text}), ...(info.error === undefined ? {} : {error: true}) }; +} + +/** Token counts of one or more model calls; a count is absent when any call left it unknown. */ +interface CallUsage { input?: number; cacheRead?: number; output?: number } + +function tokenCount(value: unknown): number | undefined { + return typeof value === "number" && Number.isSafeInteger(value) && value >= 0 ? value : undefined; +} + +/** A sum of two counts, unknown when either is unknown or the sum is not a safe integer. */ +function addCount(left: number | undefined, right: number | undefined): number | undefined { + if (left === undefined || right === undefined) return undefined; + const sum = left + right; + return Number.isSafeInteger(sum) ? sum : undefined; +} + +/** A count that Kilo Code may have zero-filled: 0 cannot be told apart from an omitted count. */ +function reportedCount(value: number | undefined): number | undefined { + return value === 0 ? undefined : value; +} + +/** + * One model call's counts from its Kilo Code step-finish tokens. Kilo Code + * reports input without cache reads and writes and output without reasoning, + * so both are added back. Kilo Code writes 0 for any count the provider + * omitted, so a call's input or output that adds up to 0 is unknown, and so + * is a 0 cache read: Kilo Code zero-fills missing cache details, and nothing + * in the step shows whether that call reported its cache reads. + */ +function stepUsage(tokens: unknown): CallUsage { + const counts = isRecord(tokens) ? tokens : {}; + const cache = isRecord(counts.cache) ? counts.cache : {}; + const cacheRead = tokenCount(cache.read); + const input = reportedCount(addCount(addCount(tokenCount(counts.input), cacheRead), tokenCount(cache.write))); + const output = reportedCount(addCount(tokenCount(counts.output), tokenCount(counts.reasoning))); + return {input, cacheRead: reportedCount(cacheRead), output}; +} + +type Message = Record; + +function messageInfo(message: Message): Record { + return isRecord(message.info) ? message.info : {}; +} + +function messageId(message: unknown): string | undefined { + const id = isRecord(message) ? messageInfo(message).id : undefined; + return typeof id === "string" ? id : undefined; +} + +function messageParts(message: Message): Message[] { + return Array.isArray(message.parts) ? message.parts.filter(isRecord) : []; +} + +/** Whether Kilo Code finished processing an assistant message, normally or with an error. */ +function assistantDone(info: Record): boolean { + return (isRecord(info.time) && info.time.completed !== undefined) || info.error !== undefined; +} + +/** A `task` tool call: its status, the child session it runs in, when it started, and whether it runs in the background. */ +interface TaskCall { status?: string; sessionId?: string; start?: number; background: boolean } + +function taskCalls(message: Message): TaskCall[] { + return messageParts(message).filter((part) => part.type === "tool" && part.tool === "task").map((part) => { + const state = isRecord(part.state) ? part.state : {}; + const metadata = isRecord(state.metadata) ? state.metadata : {}; + const input = isRecord(state.input) ? state.input : {}; + const time = isRecord(state.time) ? state.time : {}; + return { + ...(typeof state.status === "string" ? {status: state.status} : {}), + ...(typeof metadata.sessionId === "string" ? {sessionId: metadata.sessionId} : {}), + ...(typeof time.start === "number" ? {start: time.start} : {}), + background: metadata.background === true || input.background === true, + }; + }); +} + +/** + * The child session whose result a user message delivers, when Kilo Code's + * task tool added it for a finished background subagent. Kilo Code marks its + * text part as synthetic with `metadata.background` and renders the child + * session ID as ``. + */ +function backgroundResultSession(message: Message): string | undefined { + if (messageInfo(message).role !== "user") return undefined; + for (const part of messageParts(message)) { + if (part.type !== "text" || part.synthetic !== true || !isRecord(part.metadata) || part.metadata.background !== true) continue; + const match = typeof part.text === "string" ? / turn.owner === newest.owner) || candidates.every((turn) => turn.adapterPrompt)) return newest; + return {owner: AMBIGUOUS, adapterPrompt: false}; +} + +function sessionOwners(messages: Message[], prompts: ReadonlySet): Map { + const users = new Map(); + const tasks = new Map(); + let latestAssistant: Turn | undefined; + // Prompts and background results whose turn has no assistant message yet. + let openRoots: Turn[] = []; + for (const message of messages) { + const info = messageInfo(message); + const id = messageId(message); + if (info.role === "user" && id !== undefined) { + let turn: Turn; + const delivered = backgroundResultSession(message); + if (prompts.has(id)) { + turn = {owner: id, adapterPrompt: true}; + openRoots.push(turn); + } else if (delivered !== undefined) { + turn = {owner: tasks.get(delivered) ?? AMBIGUOUS, adapterPrompt: false}; + openRoots.push(turn); + } else { + turn = runningTurn(latestAssistant, openRoots); + } + users.set(id, turn); + } else if (info.role === "assistant") { + const turn = typeof info.parentID === "string" ? users.get(info.parentID) : undefined; + latestAssistant = turn ?? {owner: AMBIGUOUS, adapterPrompt: false}; + openRoots = openRoots.filter((root) => root !== turn); + for (const task of taskCalls(message)) { + if (task.sessionId !== undefined) tasks.set(task.sessionId, latestAssistant.owner); + } + } + } + return new Map([...users].map(([id, turn]) => [id, turn.owner])); +} + +/** The assistant messages of `prompt`'s turn, or undefined when a message after the prompt cannot be attributed. */ +function ownedAssistants(messages: Message[], prompt: string, prompts: ReadonlySet): Message[] | undefined { + const owners = sessionOwners(messages, prompts); + const promptIndex = messages.findIndex((message) => messageId(message) === prompt); + if (promptIndex < 0) return undefined; + const owned: Message[] = []; + // The prompt's work starts at its message, so only later messages can be its own. + for (const message of messages.slice(promptIndex + 1)) { + const info = messageInfo(message); + if (info.role === "user" && owners.get(messageId(message) ?? "") === AMBIGUOUS) return undefined; + if (info.role !== "assistant") continue; + const owner = typeof info.parentID === "string" ? owners.get(info.parentID) : undefined; + if (owner === undefined || owner === AMBIGUOUS) return undefined; + if (owner === prompt) owned.push(message); + } + return owned; +} + +/** + * The adapter's prompt among the messages after `lastMessageIdBefore`, the + * newest message before the prompt (undefined when the session was empty): + * the first user message with a text part that Kilo Code did not add itself. + * Kilo Code marks the text of background results, compaction replays and + * compaction follow-ups as synthetic. + */ +export function findPrompt(messages: unknown, lastMessageIdBefore: string | undefined): string | undefined { + if (!Array.isArray(messages)) return undefined; + const start = lastMessageIdBefore === undefined ? 0 : messages.findIndex((message) => messageId(message) === lastMessageIdBefore) + 1; + if (start === 0 && lastMessageIdBefore !== undefined) return undefined; + const prompt = messages.slice(start).filter(isRecord).find((message) => messageInfo(message).role === "user" && messageParts(message).some((part) => part.type === "text" && part.synthetic !== true)); + return prompt === undefined ? undefined : messageId(prompt); +} + +/** The foreground `task` calls of `prompt`'s turn whose child sessions hold model calls to add. */ +export function ownedChildSessions(messages: unknown, prompt: string, prompts: ReadonlySet): string[] { + if (!Array.isArray(messages)) return []; + const owned = ownedAssistants(messages.filter(isRecord), prompt, prompts) ?? []; + return [...new Set(owned.flatMap(taskCalls).filter((task) => !task.background).flatMap((task) => task.sessionId === undefined ? [] : [task.sessionId]))]; +} + +/** + * The step-finish tokens of a turn's finished assistant messages, or + * undefined when one of them is still running or finished without a + * step-finish part, because Kilo Code keeps no counts for a call it did not + * finish. + */ +function finishedSteps(assistants: Message[]): unknown[] | undefined { + const steps: unknown[] = []; + for (const message of assistants) { + const parts = messageParts(message).filter((part) => part.type === "step-finish"); + if (!assistantDone(messageInfo(message)) || parts.length === 0) return undefined; + steps.push(...parts.map((part) => part.tokens)); + } + return steps; +} + +/** + * The model calls a foreground `task` call made in its child session: the + * assistant messages created from the call's start until the next `task` call + * that resumed the same session. Undefined when the child could not be read, + * when another `task` call on the session started at the same time, so the + * windows cannot be told apart, when a message has no creation time, or when + * the child started its own subagent, which Kilo Code 7.7.12 allows only with + * a raised subagent depth. + */ +function childSteps(task: TaskCall, resumes: number[], children: ReadonlyMap): unknown[] | undefined { + if (task.sessionId === undefined || task.start === undefined) return undefined; + const child = children.get(task.sessionId); + if (child === undefined) return undefined; + const start = task.start; + if (resumes.filter((value) => value === start).length > 1) return undefined; + const end = Math.min(...resumes.filter((value) => value > start), Number.POSITIVE_INFINITY); + const window: Message[] = []; + for (const message of child.filter(isRecord)) { + const time = messageInfo(message).time; + const created = isRecord(time) && typeof time.created === "number" ? time.created : undefined; + if (created === undefined) return undefined; + if (created >= start && created < end && messageInfo(message).role === "assistant") window.push(message); + } + if (window.some((message) => taskCalls(message).length > 0)) return undefined; + return finishedSteps(window); +} + +/** + * The usage of the model calls that `prompt`'s turn made, from the adapter's + * session messages, oldest first, and the child sessions of its foreground + * subagents. `prompts` holds every prompt the adapter submitted in the + * session. Usage is undefined when a call cannot be attributed, when an owned + * call is unfinished or reported neither input nor output, and when the turn + * started a background subagent, because the subagent and Kilo Code's reply + * to its result can run after the prompt returns. That reply belongs to this + * prompt, so no later prompt counts it. + * + * Cached input is known only when every counted call reported a nonzero + * cache read. Kilo Code computes most message costs from its model catalog and does not + * mark the ones a provider reported, so the usage has no `cost_usd`. + */ +export function extractInvocationUsage(messages: unknown, prompt: string, prompts: ReadonlySet, children: ReadonlyMap = new Map()): AgentUsage | undefined { + if (!Array.isArray(messages)) return undefined; + const records = messages.filter(isRecord); + const owned = ownedAssistants(records, prompt, prompts); + if (owned === undefined || owned.length === 0) return undefined; + const steps = finishedSteps(owned); + if (steps === undefined) return undefined; + const resumes = new Map(); + for (const task of records.flatMap(taskCalls)) { + if (task.sessionId !== undefined && task.start !== undefined) resumes.set(task.sessionId, [...(resumes.get(task.sessionId) ?? []), task.start]); + } + for (const task of owned.flatMap(taskCalls)) { + // A task call that failed before Kilo Code created its child session made no subagent calls. + if (task.sessionId === undefined && task.status === "error") continue; + if (task.background) return undefined; + const child = childSteps(task, resumes.get(task.sessionId ?? "") ?? [], children); + if (child === undefined) return undefined; + steps.push(...child); + } + let sum: CallUsage = {input: 0, cacheRead: 0, output: 0}; + for (const tokens of steps) { + const call = stepUsage(tokens); + if (call.input === undefined && call.output === undefined) return undefined; + sum = {input: addCount(sum.input, call.input), cacheRead: addCount(sum.cacheRead, call.cacheRead), output: addCount(sum.output, call.output)}; + } + const cachedInput = sum.input !== undefined ? sum.cacheRead : undefined; + // Kilo Code's step total is the AI SDK's input plus output with an omitted + // side counted as 0, not a provider total, so the total is derived here. + const total = addCount(sum.input, sum.output); + const usage: AgentUsage = { + ...(sum.input === undefined ? {} : {input_tokens: sum.input, input_tokens_include_cache: true}), + ...(cachedInput === undefined ? {} : {cached_input_tokens: cachedInput}), + ...(sum.output === undefined ? {} : {output_tokens: sum.output}), + ...(total === undefined ? {} : {total_tokens: total}), + }; + return Object.keys(usage).length === 0 ? undefined : usage; } async function loadClientFactory(): Promise { @@ -359,13 +649,43 @@ async function resolveKiloExecutable(): Promise { class KiloSdkSessionHandle implements KiloSessionHandle { private stopped = false; + /** The IDs of the user messages of the adapter's prompts, in this session. */ + private readonly prompts = new Set(); constructor(readonly id: string, private readonly workspace: string, private readonly model: ReturnType, private readonly client: KiloClient, private readonly server: KiloServer, private readonly profile: string, private readonly promptTimeoutMs: number) {} private async abortPrompt(): Promise { try { await this.client.session.abort({sessionID: this.id, directory: this.workspace}, {throwOnError: true, signal: AbortSignal.timeout(STOP_TIMEOUT_MS)}); } catch { /* Preserve the timeout failure. */ } } + /** A session's messages, oldest first, or the newest `limit` of them; undefined when listing fails. */ + private async listMessages(sessionID: string, limit?: number): Promise { + try { + const {data} = await this.client.session.messages({sessionID, directory: this.workspace, ...(limit === undefined ? {} : {limit})}, {throwOnError: true, signal: AbortSignal.timeout(LIST_TIMEOUT_MS)}); + return Array.isArray(data) ? data : undefined; + } catch { + return undefined; + } + } async prompt(text: string): Promise { + const before = await this.listMessages(this.id, 1); + const outcome = await this.sendPrompt(text); + if (before === undefined) return outcome; + const lastId = messageId(before[0]); + if (before.length > 0 && lastId === undefined) return outcome; + const messages = await this.listMessages(this.id); + const prompt = findPrompt(messages, lastId); + if (prompt === undefined) return outcome; + this.prompts.add(prompt); + const children = new Map(); + for (const sessionID of ownedChildSessions(messages, prompt, this.prompts)) { + const child = await this.listMessages(sessionID); + if (child === undefined) return outcome; + children.set(sessionID, child); + } + const usage = extractInvocationUsage(messages, prompt, this.prompts, children); + return usage === undefined ? outcome : {...outcome, usage}; + } + private async sendPrompt(text: string): Promise { const controller = new AbortController(); let timedOut = false; let abortPromise: Promise | undefined; @@ -419,7 +739,7 @@ export class KiloSdkSessionFactory implements KiloSessionFactory { server = await this.serverFactory(workspace, childEnvironment(input, profile, configured.content, model.apiKeyEnv, configured.credential), executable); const client = clientFactory({baseUrl: server.url, directory: workspace}); await verifySkills(client, workspace, skills); - const created = await client.session.create({directory: workspace, agent: "build", model: {providerID: model.provider, id: model.model}}, {throwOnError: true}); + const created = await client.session.create({directory: workspace, title: SESSION_TITLE, agent: "build", model: {providerID: model.provider, id: model.model}}, {throwOnError: true}); if (!isRecord(created.data) || typeof created.data.id !== "string") throw new Error("missing session id"); await verifyMcp(client, workspace, mcpNames); return new KiloSdkSessionHandle(created.data.id, workspace, model, client, server, profile, this.promptTimeoutMs); diff --git a/adapters/typescript/kilo/test/kilo-sdk.test.mjs b/adapters/typescript/kilo/test/kilo-sdk.test.mjs index 38f23e89c..7e5bb4883 100644 --- a/adapters/typescript/kilo/test/kilo-sdk.test.mjs +++ b/adapters/typescript/kilo/test/kilo-sdk.test.mjs @@ -9,13 +9,380 @@ import {tmpdir} from "node:os"; import {join} from "node:path"; import {PassThrough} from "node:stream"; -import {classifySpawnError, closeChild, extractPromptOutcome, KiloSdkSessionFactory, startKiloServer} from "../dist/kilo-sdk.js"; +import {classifySpawnError, closeChild, extractInvocationUsage, extractPromptOutcome, findPrompt, KiloSdkSessionFactory, startKiloServer} from "../dist/kilo-sdk.js"; -test("extracts Kilo Code text, usage, and cost", () => { +let clock = 1000; + +function step(tokens, cost = 0) { + return {type: "step-finish", reason: "stop", cost, tokens}; +} + +// Kilo Code overwrites a message's tokens with each finished step, so the message mirrors its last step. +// Kilo Code sets each assistant message's parentID to the user message of its turn and records when it finished. +function assistant(id, parts, parentID = "user-1") { + clock += 1; + const tokens = parts.findLast((part) => part.type === "step-finish")?.tokens ?? {input: 0, output: 0, reasoning: 0, cache: {read: 0, write: 0}}; + return {info: {id, role: "assistant", parentID, cost: 0, tokens, time: {created: clock, completed: clock}}, parts}; +} + +function user(id, parts = [{type: "text", text: "do it"}]) { + clock += 1; + return {info: {id, role: "user", time: {created: clock}}, parts}; +} + +function taskCall(sessionId, start, background = false) { + return {type: "tool", tool: "task", state: {status: "completed", input: {background}, metadata: {sessionId, ...(background ? {background: true} : {})}, time: {start}}}; +} + +// Kilo Code's task tool prompts the parent session with a background subagent's result as synthetic text. +function backgroundResult(id, sessionId) { + return user(id, [{type: "text", synthetic: true, metadata: {background: true}, text: `done`}]); +} + +const call = (input, output, read = 0, write = 0, reasoning = 0) => ({input, output, reasoning, cache: {read, write}}); +const userPrompt = user("user-1"); +const usageOf = (messages, prompt = "user-1", prompts = [prompt], children = new Map()) => extractInvocationUsage(messages, prompt, new Set(prompts), children); + +test("extracts Kilo Code terminal text without taking usage from the last message", () => { assert.deepEqual(extractPromptOutcome({ info: {tokens: {input: 4, output: 6, total: 10}, cost: 0.05}, parts: [{type: "text", text: "hello"}, {type: "tool", name: "read"}, {type: "text", text: " world"}], - }), {text: "hello world", usage: {input_tokens: 4, output_tokens: 6, total_tokens: 10, cost_usd: 0.05}}); + }), {text: "hello world"}); +}); + +test("sums every Kilo Code model call of the prompt with cache and reasoning", () => { + const messages = [ + user("old-user"), + assistant("old-assistant", [step(call(999, 999))], "old-user"), + userPrompt, + assistant("read-file", [step(call(1200, 80, 0, 1000))]), + assistant("edit-file", [step(call(150, 40, 2200, 0, 20))]), + assistant("answer", [step(call(90, 200, 2350)), {type: "text", text: "done"}]), + ]; + // The first call's cache read is 0, which Kilo Code also stores for omitted cache details. + assert.deepEqual(usageOf(messages, "user-1", ["old-user", "user-1"]), { + input_tokens: 6990, + input_tokens_include_cache: true, + output_tokens: 340, + total_tokens: 7330, + }); +}); + +test("sums every finished step that Kilo Code recorded in one message", () => { + const messages = [userPrompt, assistant("a", [step(call(4, 6)), step(call(8, 12))])]; + assert.deepEqual(usageOf(messages), {input_tokens: 12, input_tokens_include_cache: true, output_tokens: 18, total_tokens: 30}); +}); + +test("counts Kilo Code compaction calls inside the prompt's turn", () => { + const compaction = [ + userPrompt, + assistant("a", [step(call(100, 10, 50))]), + user("compaction", [{type: "compaction", auto: true}]), + assistant("summary", [step(call(300, 40))], "compaction"), + user("continue", [{type: "text", synthetic: true, metadata: {compaction_continue: true}, text: "Continue"}]), + assistant("b", [step(call(60, 5, 20))], "continue"), + ]; + assert.deepEqual(usageOf(compaction), {input_tokens: 530, input_tokens_include_cache: true, output_tokens: 55, total_tokens: 585}); + // Compaction before the prompt's first call follows the previous prompt's last call. + const preflight = [ + user("user-0"), + assistant("earlier", [step(call(10, 1))], "user-0"), + userPrompt, + user("compaction", [{type: "compaction", auto: true}]), + assistant("summary", [step(call(300, 40))], "compaction"), + assistant("answer", [step(call(20, 2))]), + ]; + assert.deepEqual(usageOf(preflight, "user-1", ["user-0", "user-1"]), {input_tokens: 320, input_tokens_include_cache: true, output_tokens: 42, total_tokens: 362}); +}); + +test("adds the calls of a Kilo Code foreground subagent from its child session", () => { + const messages = [userPrompt, assistant("delegate", [taskCall("child-1", 2000), step(call(1500, 15))])]; + const child = [ + {info: {id: "c-user", role: "user", time: {created: 2001}}, parts: [{type: "text", text: "sub task"}]}, + {info: {id: "c-a", role: "assistant", parentID: "c-user", time: {created: 2002, completed: 2003}}, parts: [step(call(5000, 50))]}, + ]; + assert.deepEqual(usageOf(messages, "user-1", ["user-1"], new Map([["child-1", child]])), {input_tokens: 6500, input_tokens_include_cache: true, output_tokens: 65, total_tokens: 6565}); + assert.equal(usageOf(messages), undefined, "an unread child session leaves usage unknown"); +}); + +test("adds only the calls of the resumed Kilo Code subagent window", () => { + const messages = [ + user("user-0"), + assistant("first-delegate", [taskCall("child-1", 2000), step(call(10, 1))], "user-0"), + userPrompt, + assistant("resume", [taskCall("child-1", 3000), step(call(20, 2))]), + ]; + const child = [ + {info: {id: "c-user-0", role: "user", time: {created: 2001}}, parts: []}, + {info: {id: "c-a-0", role: "assistant", parentID: "c-user-0", time: {created: 2002, completed: 2003}}, parts: [step(call(700, 7))]}, + {info: {id: "c-user-1", role: "user", time: {created: 3001}}, parts: []}, + {info: {id: "c-a-1", role: "assistant", parentID: "c-user-1", time: {created: 3002, completed: 3003}}, parts: [step(call(400, 4))]}, + ]; + assert.deepEqual(usageOf(messages, "user-1", ["user-0", "user-1"], new Map([["child-1", child]])), {input_tokens: 420, input_tokens_include_cache: true, output_tokens: 6, total_tokens: 426}); +}); + +test("reports no Kilo Code usage when a background turn may own a compaction after a newer prompt", () => { + // P0 starts a background subagent and returns. Its result arrives, and P1 is stored while that background turn runs. + // The background turn compacts before its first assistant message, then P1's own call runs. + const messages = [ + user("user-0"), + assistant("delegate", [taskCall("child-1", 2000, true), step(call(10, 1))], "user-0"), + backgroundResult("result", "child-1"), + userPrompt, + user("compaction", [{type: "compaction", auto: true}]), + assistant("summary", [step(call(800, 17))], "compaction"), + user("continue", [{type: "text", synthetic: true, metadata: {compaction_continue: true}, text: "Continue"}]), + assistant("answer", [step(call(3, 1))]), + ]; + assert.equal(usageOf(messages, "user-1", ["user-0", "user-1"]), undefined); + // The same when the background result was stored while P0's turn was still running. + const queued = [ + user("user-x"), + assistant("delegate", [taskCall("child-1", 2000, true), step(call(10, 1))], "user-x"), + user("user-0"), + backgroundResult("result", "child-1"), + assistant("p0-answer", [step(call(5, 1))], "user-0"), + userPrompt, + user("compaction", [{type: "compaction", auto: true}]), + assistant("summary", [step(call(800, 17))], "compaction"), + assistant("answer", [step(call(3, 1))]), + ]; + assert.equal(usageOf(queued, "user-1", ["user-x", "user-0", "user-1"]), undefined); +}); + +test("reports no Kilo Code usage when two subagent task windows start at the same time", () => { + const messages = [userPrompt, assistant("delegate", [taskCall("child-1", 2000), taskCall("child-1", 2000), step(call(10, 1))])]; + const child = [ + {info: {id: "c-user", role: "user", time: {created: 2001}}, parts: []}, + {info: {id: "c-a", role: "assistant", parentID: "c-user", time: {created: 2002, completed: 2003}}, parts: [step(call(500, 5))]}, + ]; + assert.equal(usageOf(messages, "user-1", ["user-1"], new Map([["child-1", child]])), undefined); +}); + +test("reports no Kilo Code usage when a subagent starts its own subagent", () => { + const messages = [userPrompt, assistant("delegate", [taskCall("child-1", 2000), step(call(10, 1))])]; + const child = [{info: {id: "c-a", role: "assistant", parentID: "c-user", time: {created: 2002, completed: 2003}}, parts: [taskCall("grandchild", 2002), step(call(5, 1))]}]; + assert.equal(usageOf(messages, "user-1", ["user-1"], new Map([["child-1", child]])), undefined); +}); + +test("reports no Kilo Code usage when the prompt starts a background subagent", () => { + const messages = [userPrompt, assistant("delegate", [taskCall("child-1", 2000, true), step(call(10, 1))]), assistant("answer", [step(call(20, 2))])]; + assert.equal(usageOf(messages), undefined); +}); + +test("leaves a Kilo Code reply to an earlier background subagent out of the prompt that it ran during", () => { + const messages = [ + user("user-0"), + assistant("delegate", [taskCall("child-1", 2000, true), step(call(10, 1))], "user-0"), + userPrompt, + backgroundResult("result", "child-1"), + assistant("background-reply", [step(call(7000, 70))], "result"), + assistant("answer", [step(call(300, 3))]), + ]; + assert.deepEqual(usageOf(messages, "user-1", ["user-0", "user-1"]), {input_tokens: 300, input_tokens_include_cache: true, output_tokens: 3, total_tokens: 303}); + assert.equal(usageOf(messages, "user-0", ["user-0", "user-1"]), undefined); +}); + +test("reports no Kilo Code usage when a new call cannot be attributed", () => { + const unknownResult = [user("user-0"), userPrompt, backgroundResult("result", "missing"), assistant("reply", [step(call(4, 6))], "result"), assistant("answer", [step(call(4, 6))])]; + assert.equal(usageOf(unknownResult, "user-1", ["user-0", "user-1"]), undefined); + const orphan = assistant("orphan", [step(call(4, 6))]); + delete orphan.info.parentID; + assert.equal(usageOf([userPrompt, orphan]), undefined); + assert.equal(usageOf([assistant("a", [step(call(4, 6))])]), undefined, "a missing prompt message leaves usage unknown"); +}); + +test("finds the Kilo Code prompt after the newest message before it", () => { + const messages = [user("user-0"), assistant("a", [step(call(4, 6))], "user-0"), backgroundResult("result", "child-1"), user("user-1")]; + assert.equal(findPrompt(messages, "a"), "user-1"); + assert.equal(findPrompt(messages, undefined), "user-0"); + assert.equal(findPrompt(messages, "deleted"), undefined); + assert.equal(findPrompt(undefined, undefined), undefined); +}); + +test("reports no Kilo Code usage when a model call never finished", () => { + const running = assistant("running", [{type: "text", text: "partial"}]); + delete running.info.time.completed; + assert.equal(usageOf([userPrompt, assistant("a", [step(call(4, 6))]), running]), undefined); + assert.equal(usageOf([userPrompt, assistant("a", [step(call(4, 6))]), assistant("aborted", [{type: "text", text: "partial"}])]), undefined); + assert.equal(usageOf([], "user-1"), undefined); + assert.equal(usageOf(undefined, "user-1"), undefined); +}); + +test("derives the Kilo Code total from input and output", () => { + const messages = [userPrompt, assistant("a", [step({total: 28, ...call(10, 15, 3)})]), assistant("b", [step(call(7, 3, 1, 0, 2))])]; + assert.deepEqual(usageOf(messages), {input_tokens: 21, cached_input_tokens: 4, input_tokens_include_cache: true, output_tokens: 20, total_tokens: 41}); +}); + +test("leaves Kilo Code counts unknown when any call omitted them", () => { + const missingOutput = [userPrompt, assistant("a", [step({total: 10, ...call(4, 6)})]), assistant("b", [step({input: 4, cache: {read: 1, write: 0}})])]; + assert.deepEqual(usageOf(missingOutput), {input_tokens: 9, input_tokens_include_cache: true}); + assert.deepEqual(usageOf([userPrompt, assistant("a", [step({input: 4, output: 6, reasoning: 0})])]), {output_tokens: 6}); +}); + +// Kilo Code replaces every count a provider omitted with 0 and keeps an omitted total undefined. +const zeroFilled = call(0, 0); + +test("reports no Kilo Code usage when a finished call reported no counts", () => { + assert.equal(usageOf([userPrompt, assistant("a", [step(zeroFilled)])]), undefined); + assert.equal(usageOf([userPrompt, assistant("a", [step({...zeroFilled, total: 0})])]), undefined); + assert.equal(usageOf([userPrompt, assistant("a", [step(call(4, 6, 2))]), assistant("b", [step(zeroFilled)])]), undefined); +}); + +test("leaves a Kilo Code count unknown when a call zero-filled it", () => { + const reported = assistant("a", [step(call(10, 15, 0, 0, 5))]); + const zeroOutput = call(7, 0, 2); + assert.deepEqual(usageOf([userPrompt, reported, assistant("b", [step(zeroOutput)])]), {input_tokens: 19, input_tokens_include_cache: true}); + // Kilo Code's total is the AI SDK's input plus output with the omitted output counted as 0. + assert.deepEqual(usageOf([userPrompt, reported, assistant("b", [step({...zeroOutput, total: 9})])]), {input_tokens: 19, input_tokens_include_cache: true}); +}); + +test("leaves Kilo Code cached input unknown when a call's cache read was zero-filled", () => { + // An earlier call read 20 tokens from the cache, then a call omitted its cache details, which Kilo Code stores as 0. + const earlier = [user("user-0"), assistant("earlier", [step(call(100, 10, 20))], "user-0")]; + assert.deepEqual(usageOf([...earlier, userPrompt, assistant("a", [step(call(10, 20))])], "user-1", ["user-0", "user-1"]), {input_tokens: 10, input_tokens_include_cache: true, output_tokens: 20, total_tokens: 30}); + // A cache write does not show that the call reported its cache reads. + assert.deepEqual(usageOf([userPrompt, assistant("a", [step(call(10, 20, 0, 500))])]), {input_tokens: 510, input_tokens_include_cache: true, output_tokens: 20, total_tokens: 530}); + // Every call read from the cache, so the cached input is known. + assert.deepEqual(usageOf([userPrompt, assistant("a", [step(call(10, 20, 30))]), assistant("b", [step(call(5, 5, 40))])]), {input_tokens: 85, cached_input_tokens: 70, input_tokens_include_cache: true, output_tokens: 25, total_tokens: 110}); +}); + +// Step tokens recorded by a real Kilo Code 7.7.12 server for a provider usage report with one side missing. +test("does not take Kilo Code's total when the provider omitted one side", () => { + assert.deepEqual(usageOf([userPrompt, assistant("a", [step({total: 100, ...call(100, 0)})])]), {input_tokens: 100, input_tokens_include_cache: true}); + assert.deepEqual(usageOf([userPrompt, assistant("a", [step({total: 7, ...call(0, 7)})])]), {output_tokens: 7}); + assert.equal(usageOf([userPrompt, assistant("a", [step({total: 0, ...call(0, 0)})])]), undefined); +}); + +test("leaves a Kilo Code sum unknown when it passes the safe integer range", () => { + const half = 2 ** 52; + const big = {total: half + 1, ...call(half, 1)}; + assert.deepEqual(usageOf([userPrompt, assistant("a", [step(big)]), assistant("b", [step(big)])]), {output_tokens: 2}); +}); + +const sessionInput = (workspace) => ({ + agentName: "reviewer", + baseDir: workspace, + config: {models: {default: {provider: "nvidia", model: "nemotron", api_key_env: "MODEL_KEY"}}}, + runtimeContext: {runtime_id: "r", invocation_id: "i", request_id: "q", artifacts: {}, environment: {environment_id: "e", provider: "local", ownership: "caller_owned", control_location: "external_control", workspace, env: {MODEL_KEY: "secret"}}}, +}); + +// A Kilo Code session fake that lists messages oldest first, or only the newest `limit` of them. +function sessionClient(history, respond, listings = [], children = new Map()) { + return { + app: {skills: async () => ({data: []})}, + mcp: {status: async () => ({data: {}})}, + session: { + create: async () => ({data: {id: "session-1"}}), + messages: async (parameters) => { + const result = listings.shift(); + if (result instanceof Error) throw result; + const source = parameters.sessionID === "session-1" ? history : children.get(parameters.sessionID) ?? []; + const listed = parameters.limit === undefined ? source : source.slice(-parameters.limit); + return {data: listed.map((message) => structuredClone(message))}; + }, + prompt: async () => ({data: respond()}), + delete: async () => {}, + }, + }; +} + +async function withSession(client, check) { + const workspace = await mkdtemp(join(tmpdir(), "kilo-adapter-test-")); + const factory = new KiloSdkSessionFactory(async () => () => client, async () => ({url: "http://127.0.0.1:4123", close: async () => {}})); + try { + const handle = await factory.create(sessionInput(workspace)); + await check(handle); + await handle.stop(); + } finally { + await rm(workspace, {recursive: true, force: true}); + } +} + +test("reports each Kilo Code prompt's own usage on a reused session", async () => { + const history = []; + let turn = 0; + const client = sessionClient(history, () => { + turn += 1; + history.push(user(`user-${turn}`)); + history.push(assistant(`tool-${turn}`, [step({total: 12, ...call(10, 2)})], `user-${turn}`)); + const final = assistant(`final-${turn}`, [step({total: 4, ...call(3, 1)}), {type: "text", text: `answer ${turn}`}], `user-${turn}`); + history.push(final); + return final; + }); + await withSession(client, async (handle) => { + const expected = {input_tokens: 13, input_tokens_include_cache: true, output_tokens: 3, total_tokens: 16}; + assert.deepEqual(await handle.prompt("one"), {text: "answer 1", usage: expected}); + assert.deepEqual(await handle.prompt("two"), {text: "answer 2", usage: expected}); + }); +}); + +test("reads a Kilo Code foreground subagent's child session for the prompt's usage", async () => { + const history = []; + const children = new Map([["child-1", [ + {info: {id: "c-user", role: "user", time: {created: 5001}}, parts: [{type: "text", text: "sub task"}]}, + {info: {id: "c-a", role: "assistant", parentID: "c-user", time: {created: 5002, completed: 5003}}, parts: [step(call(5000, 50))]}, + ]]]); + const client = sessionClient(history, () => { + history.push(user("user-1")); + const final = assistant("final", [taskCall("child-1", 5000), step(call(1500, 15)), {type: "text", text: "answer"}]); + history.push(final); + return final; + }, [], children); + await withSession(client, async (handle) => { + assert.deepEqual(await handle.prompt("one"), {text: "answer", usage: {input_tokens: 6500, input_tokens_include_cache: true, output_tokens: 65, total_tokens: 6565}}); + }); +}); + +test("keeps a late Kilo Code background reply out of the next prompt's usage", async () => { + const history = []; + let turn = 0; + const client = sessionClient(history, () => { + turn += 1; + history.push(user(`user-${turn}`)); + if (turn === 2) { + // The background subagent from the first prompt finishes, and Kilo Code prompts the session with its result. + history.push(backgroundResult("result", "child-1")); + history.push(assistant("background-reply", [step(call(100, 10, 5))], "result")); + } + const parts = turn === 1 ? [taskCall("child-1", 5000, true)] : []; + const final = assistant(`final-${turn}`, [...parts, step(call(3, 1, 1)), {type: "text", text: `answer ${turn}`}], `user-${turn}`); + history.push(final); + return final; + }); + await withSession(client, async (handle) => { + assert.deepEqual(await handle.prompt("one"), {text: "answer 1"}); + assert.deepEqual(await handle.prompt("two"), { + text: "answer 2", + usage: {input_tokens: 4, cached_input_tokens: 1, input_tokens_include_cache: true, output_tokens: 1, total_tokens: 5}, + }); + }); +}); + +test("returns the Kilo Code answer without usage when listing messages fails after the prompt", async () => { + const history = []; + const client = sessionClient(history, () => { + history.push(user("user-1")); + const final = assistant("final", [step(call(3, 1)), {type: "text", text: "answer"}]); + history.push(final); + return final; + }, [undefined, new Error("listing failed")]); + await withSession(client, async (handle) => { + assert.deepEqual(await handle.prompt("one"), {text: "answer"}); + }); +}); + +test("reports a Kilo Code model error without usage when the call never finished", async () => { + const history = []; + const client = sessionClient(history, () => { + history.push(user("user-1")); + const failed = {info: {id: "failed", role: "assistant", parentID: "user-1", cost: 0, tokens: call(0, 0), time: {created: 1, completed: 2}, error: {name: "APIError", data: {message: "secret"}}}, parts: []}; + history.push(failed); + return failed; + }); + await withSession(client, async (handle) => { + assert.deepEqual(await handle.prompt("one"), {error: true}); + }); }); test("marks upstream assistant errors without exposing their content", () => { @@ -131,7 +498,7 @@ test("projects normalized configuration into an isolated Kilo Code server", asyn assert.equal(serverEnvironment.KILO_DISABLE_PROJECT_CONFIG, "1"); assert.equal(serverEnvironment.MODEL_KEY, "secret"); assert.equal(serverEnvironment.NVIDIA_API_KEY, undefined); - assert.deepEqual(createParameters, {directory: workspace, agent: "build", model: {providerID: "nvidia", id: "nemotron"}}); + assert.deepEqual(createParameters, {directory: workspace, title: "NeMo Fabric", agent: "build", model: {providerID: "nvidia", id: "nemotron"}}); await handle.stop(); assert.deepEqual(deleteParameters, {sessionID: "session-1", directory: workspace}); assert.ok(deleteSignal instanceof AbortSignal); diff --git a/docs/integrations/harness/kilo.mdx b/docs/integrations/harness/kilo.mdx index 315f34af1..7bbedeab4 100644 --- a/docs/integrations/harness/kilo.mdx +++ b/docs/integrations/harness/kilo.mdx @@ -118,7 +118,41 @@ filters are unsupported. ### Runtime Behavior The adapter accepts plain-text input and returns the terminal assistant text in -`output.response`. Usage and cost are included when Kilo Code reports them. +`output.response`. + +Usage sums the model calls that Kilo Code recorded for the invocation's prompt, +including in-turn compaction and the calls of foreground `task` subagents in +their child sessions. Kilo Code 7.7.12 can run a subagent in the background when +the model asks for it. An invocation that starts a background subagent reports +no usage, because the subagent and Kilo Code's reply to its result can run after +the invocation returns. That reply is never counted toward a later invocation. +The adapter sets the session title so that Kilo Code does not make an unrecorded +model call to generate one. + +`input_tokens` includes cache reads and cache writes, `cached_input_tokens` +holds the cache reads, and `output_tokens` includes reasoning tokens. +`total_tokens` is `input_tokens` plus `output_tokens`, so it is unknown when +either of them is unknown. The adapter does not use Kilo Code's own total, +because Kilo Code 7.7.12 computes it as input plus output with an omitted side +counted as 0. + +Kilo Code 7.7.12 records an omitted input, output, reasoning, or cache count as +0 before the adapter sees it. The adapter therefore treats a call's input or +output that adds up to 0 as unknown. Kilo zero-fills missing cache details, so a +zero cache-read count remains unknown unless its presence can be established for +that call. When any model call leaves a field unknown, or the field's sum is +larger than JavaScript can represent exactly, the result omits that field. The +result has no usage when a finished model call reported neither input nor +output, or when an assistant message of the invocation never finished. It also +has no usage when a subagent's child session cannot be read or starts its own +subagent, or when a new model call cannot be traced to a prompt. Kilo Code does +not record an attempt that it retries before the call finishes, so those tokens +are not counted. + +Kilo Code computes most cost figures from its model catalog's rates, and its +messages do not say which costs a provider reported, so the adapter does not +report them as `cost_usd`. + Ordered invocations reuse one warm session and its conversation history. A prompt has a 30-minute deadline; on timeout the adapter asks Kilo Code to abort the prompt and reports `kilo_prompt_timeout`. From fb57f68dcef9358b2ea806796e3b09f4ef8fb412 Mon Sep 17 00:00:00 2001 From: mnajafian-nv Date: Tue, 20 Oct 2026 01:25:02 -0700 Subject: [PATCH 04/29] docs: state the Kilo Code cache-read rule exactly A cache read of 0 is always unknown because Kilo Code zero-fills missing cache details, so cached input is reported only when every counted call has a nonzero cache read. The previous wording suggested presence could sometimes be established, which the adapter never does. Signed-off-by: mnajafian-nv --- adapters/typescript/kilo/README.md | 6 +++--- docs/integrations/harness/kilo.mdx | 6 +++--- 2 files changed, 6 insertions(+), 6 deletions(-) diff --git a/adapters/typescript/kilo/README.md b/adapters/typescript/kilo/README.md index 96bb37661..c12a1910b 100644 --- a/adapters/typescript/kilo/README.md +++ b/adapters/typescript/kilo/README.md @@ -94,9 +94,9 @@ counted as 0. Kilo Code 7.7.12 records an omitted input, output, reasoning, or cache count as 0 before the adapter sees it. The adapter therefore treats a call's input or -output that adds up to 0 as unknown. Kilo zero-fills missing cache details, so a -zero cache-read count remains unknown unless its presence can be established for -that call. When any model call leaves a field unknown, or the field's sum is +output that adds up to 0 as unknown. Kilo Code zero-fills missing cache details, so +a cache read of 0 is unknown, and `cached_input_tokens` is reported only when +every counted call has a nonzero cache read. When any model call leaves a field unknown, or the field's sum is larger than JavaScript can represent exactly, the result omits that field. The result has no usage when a finished model call reported neither input nor output, or when an assistant message of the invocation never finished. It also diff --git a/docs/integrations/harness/kilo.mdx b/docs/integrations/harness/kilo.mdx index 7bbedeab4..994200569 100644 --- a/docs/integrations/harness/kilo.mdx +++ b/docs/integrations/harness/kilo.mdx @@ -138,9 +138,9 @@ counted as 0. Kilo Code 7.7.12 records an omitted input, output, reasoning, or cache count as 0 before the adapter sees it. The adapter therefore treats a call's input or -output that adds up to 0 as unknown. Kilo zero-fills missing cache details, so a -zero cache-read count remains unknown unless its presence can be established for -that call. When any model call leaves a field unknown, or the field's sum is +output that adds up to 0 as unknown. Kilo Code zero-fills missing cache details, so +a cache read of 0 is unknown, and `cached_input_tokens` is reported only when +every counted call has a nonzero cache read. When any model call leaves a field unknown, or the field's sum is larger than JavaScript can represent exactly, the result omits that field. The result has no usage when a finished model call reported neither input nor output, or when an assistant message of the invocation never finished. It also From 1a3421229d921766a4ae13acba390b6e3f5d191f Mon Sep 17 00:00:00 2001 From: mnajafian-nv Date: Tue, 20 Oct 2026 02:00:41 -0700 Subject: [PATCH 05/29] docs: describe when a Kilo Code call counts as finished The usage docs said the result has no usage when an assistant message never finished. The adapter treats a message as finished when it has a completion time or an error, and counts the step-finish parts of a message that ended with an error. Usage is unknown only when a message has no step-finish part or is still running, with neither a completion time nor an error. State that condition in the README and the docs page. Signed-off-by: mnajafian-nv --- adapters/typescript/kilo/README.md | 6 ++++-- docs/integrations/harness/kilo.mdx | 6 ++++-- 2 files changed, 8 insertions(+), 4 deletions(-) diff --git a/adapters/typescript/kilo/README.md b/adapters/typescript/kilo/README.md index c12a1910b..efb1a4b3c 100644 --- a/adapters/typescript/kilo/README.md +++ b/adapters/typescript/kilo/README.md @@ -98,8 +98,10 @@ output that adds up to 0 as unknown. Kilo Code zero-fills missing cache details, a cache read of 0 is unknown, and `cached_input_tokens` is reported only when every counted call has a nonzero cache read. When any model call leaves a field unknown, or the field's sum is larger than JavaScript can represent exactly, the result omits that field. The -result has no usage when a finished model call reported neither input nor -output, or when an assistant message of the invocation never finished. It also +result has no usage when a counted model call reported neither input nor output, +or when an assistant message of the invocation has no step-finish part or is +still running, meaning it has neither a completion time nor an error. A message +that ended with an error still counts the step-finish parts it recorded. It also has no usage when a subagent's child session cannot be read or starts its own subagent, or when a new model call cannot be traced to a prompt. Kilo Code does not record an attempt that it retries before the call finishes, so those tokens diff --git a/docs/integrations/harness/kilo.mdx b/docs/integrations/harness/kilo.mdx index 994200569..9e8312176 100644 --- a/docs/integrations/harness/kilo.mdx +++ b/docs/integrations/harness/kilo.mdx @@ -142,8 +142,10 @@ output that adds up to 0 as unknown. Kilo Code zero-fills missing cache details, a cache read of 0 is unknown, and `cached_input_tokens` is reported only when every counted call has a nonzero cache read. When any model call leaves a field unknown, or the field's sum is larger than JavaScript can represent exactly, the result omits that field. The -result has no usage when a finished model call reported neither input nor -output, or when an assistant message of the invocation never finished. It also +result has no usage when a counted model call reported neither input nor output, +or when an assistant message of the invocation has no step-finish part or is +still running, meaning it has neither a completion time nor an error. A message +that ended with an error still counts the step-finish parts it recorded. It also has no usage when a subagent's child session cannot be read or starts its own subagent, or when a new model call cannot be traced to a prompt. Kilo Code does not record an attempt that it retries before the call finishes, so those tokens From 4585ac7ee03f0d055f71fc8dd46eef91778fbd97 Mon Sep 17 00:00:00 2001 From: mnajafian-nv Date: Tue, 20 Oct 2026 08:51:15 -0700 Subject: [PATCH 06/29] 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 07/29] 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 08/29] 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 09/29] 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 10/29] 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 11/29] 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 12/29] 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 13/29] 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 14/29] 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 15/29] 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 16/29] 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 baf8be519a93db6af185b0363260aa185e71fbcf Mon Sep 17 00:00:00 2001 From: mnajafian-nv Date: Thu, 22 Oct 2026 15:58:12 -0700 Subject: [PATCH 17/29] test: clean up Kilo Code SDK temporary workspaces Keep SDK fixtures in their unique temporary workspace and use the neighboring tests' cleanup pattern. Avoid writing to the user home directory or renaming across filesystems. Signed-off-by: mnajafian-nv --- adapters/typescript/kilo/test/usage-sdk.test.mjs | 10 ++++------ 1 file changed, 4 insertions(+), 6 deletions(-) diff --git a/adapters/typescript/kilo/test/usage-sdk.test.mjs b/adapters/typescript/kilo/test/usage-sdk.test.mjs index 54f805e1a..0282e98b9 100644 --- a/adapters/typescript/kilo/test/usage-sdk.test.mjs +++ b/adapters/typescript/kilo/test/usage-sdk.test.mjs @@ -3,9 +3,9 @@ import assert from "node:assert/strict"; import {createServer} from "node:http"; -import {mkdir, mkdtemp, realpath, rename} from "node:fs/promises"; -import {homedir, tmpdir} from "node:os"; -import {basename, join} from "node:path"; +import {mkdtemp, realpath, rm} from "node:fs/promises"; +import {tmpdir} from "node:os"; +import {join} from "node:path"; import test from "node:test"; import {createKiloClient} from "@kilocode/sdk/v2"; import {KiloSdkSessionFactory} from "../dist/kilo-sdk.js"; @@ -73,9 +73,7 @@ for (const mode of ["normal", "retry", "lost-events"]) { events?.end(); server.closeAllConnections(); await new Promise(resolve => server.close(resolve)); - const trash = join(homedir(), "trash", "nemo-fabric"); - await mkdir(trash, {recursive: true}); - await rename(workspace, join(trash, `${basename(workspace)}-${Date.now()}`)); + await rm(workspace, {recursive: true, force: true}); } }); } From a6318d4dede15289bbf31485296c6bec8bc524b0 Mon Sep 17 00:00:00 2001 From: mnajafian-nv Date: Thu, 22 Oct 2026 16:24:17 -0700 Subject: [PATCH 18/29] 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 19/29] 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 19224dc9f932c935576dd6008482d94bf28a1446 Mon Sep 17 00:00:00 2001 From: mnajafian-nv Date: Thu, 22 Oct 2026 17:12:11 -0700 Subject: [PATCH 20/29] docs: clarify collector pricing scope Signed-off-by: mnajafian-nv --- adapters/typescript/kilo/README.md | 4 ++-- docs/integrations/harness/kilo.mdx | 4 ++-- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/adapters/typescript/kilo/README.md b/adapters/typescript/kilo/README.md index 0c694897c..82743937f 100644 --- a/adapters/typescript/kilo/README.md +++ b/adapters/typescript/kilo/README.md @@ -117,8 +117,8 @@ is unavailable, no subtotal is invented. The prompt is never replayed. Each assistant message supplies its own accounting model/provider identity. A custom configured endpoint does not establish native provider pricing. -Fabric can separately estimate cost when the source identity, usage and price -catalog establish an applicable basis. +The adapter retains source-established model and provider evidence for separate +pricing; it does not compute catalog estimates. Ordered invocations reuse one warm Kilo Code session and its conversation history. A prompt has a 30-minute deadline; on timeout the adapter diff --git a/docs/integrations/harness/kilo.mdx b/docs/integrations/harness/kilo.mdx index c3caa77d6..e3eff02ec 100644 --- a/docs/integrations/harness/kilo.mdx +++ b/docs/integrations/harness/kilo.mdx @@ -161,8 +161,8 @@ is unavailable, no subtotal is invented. The prompt is never replayed. Each assistant message supplies its own accounting model/provider identity. A custom configured endpoint does not establish native provider pricing. -Fabric can separately estimate cost when the source identity, usage and price -catalog establish an applicable basis. +The adapter retains source-established model and provider evidence for separate +pricing; it does not compute catalog estimates. Ordered invocations reuse one warm session and its conversation history. A prompt has a 30-minute deadline; on timeout the adapter asks Kilo Code to abort From 93b00ecdbf82a618e8b6947692b9ea333f59f38b Mon Sep 17 00:00:00 2001 From: mnajafian-nv Date: Thu, 22 Oct 2026 17:23:32 -0700 Subject: [PATCH 21/29] 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 22/29] 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 23/29] 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 ebe4cdc84428e4bdad1f7705a330c3c9ca0e49f1 Mon Sep 17 00:00:00 2001 From: mnajafian-nv Date: Thu, 22 Oct 2026 23:34:28 -0700 Subject: [PATCH 24/29] docs: clarify Kilo Code usage terminology Signed-off-by: mnajafian-nv --- adapters/typescript/kilo/README.md | 10 +++++----- docs/integrations/harness/kilo.mdx | 10 +++++----- 2 files changed, 10 insertions(+), 10 deletions(-) diff --git a/adapters/typescript/kilo/README.md b/adapters/typescript/kilo/README.md index 82743937f..ac590f24a 100644 --- a/adapters/typescript/kilo/README.md +++ b/adapters/typescript/kilo/README.md @@ -79,17 +79,17 @@ The adapter accepts plain-text input and returns the terminal assistant text in Usage comes from each retained `step-finish` record owned by the invocation's prompt, including in-turn compaction and foreground `task` child-session windows. The adapter sends source observations through -`nemo_fabric.usage_accounting.v1`; Fabric's shared accounting code owns sums, +`nemo_fabric.usage_accounting.v1`; NeMo Fabric's shared accounting code owns sums, complete totals and pricing. Missing evidence preserves useful observations without publishing incomplete counts as complete. Kilo Code 7.7.12 subtracts cache reads/writes from its stored input and reasoning from stored output, then clamps at zero. The adapter reverses that source transform only when a positive adjusted counter proves the clamp lost no -information. Fabric's normalized input includes cache and output includes -reasoning. Kilo zero-fills omitted dimensions, so a stored cache or reasoning +information. NeMo Fabric's normalized input includes cache and output includes +reasoning. Kilo Code zero-fills omitted dimensions, so a stored cache or reasoning zero cannot establish provider field presence. A positive cache count stays -independently useful even when input is unknown. Kilo's derived `total` and +independently useful even when input is unknown. Kilo Code's derived `total` and mixed provider/catalog `cost` do not become reported totals or `cost_usd`. A transient event observer detects actual retries hidden by the final message @@ -106,7 +106,7 @@ Actual background tasks leave complete totals unknown because their child work and injected parent reply can outlive the invocation. Settled owned steps remain observed. Late work never moves into a later invocation. Unread, ambiguous or deeper child sessions also leave gaps. Merely enabling delegation -does not. The adapter sets a fixed session title to avoid Kilo's separate title +does not. The adapter sets a fixed session title to avoid Kilo Code's separate title model call. When prompt communication fails, the adapter uses the same bounded history diff --git a/docs/integrations/harness/kilo.mdx b/docs/integrations/harness/kilo.mdx index e3eff02ec..5660c1275 100644 --- a/docs/integrations/harness/kilo.mdx +++ b/docs/integrations/harness/kilo.mdx @@ -123,17 +123,17 @@ The adapter accepts plain-text input and returns the terminal assistant text in Usage comes from each retained `step-finish` record owned by the invocation's prompt, including in-turn compaction and foreground `task` child-session windows. The adapter sends source observations through -`nemo_fabric.usage_accounting.v1`; Fabric's shared accounting code owns sums, +`nemo_fabric.usage_accounting.v1`; NeMo Fabric's shared accounting code owns sums, complete totals and pricing. Missing evidence preserves useful observations without publishing incomplete counts as complete. Kilo Code 7.7.12 subtracts cache reads/writes from its stored input and reasoning from stored output, then clamps at zero. The adapter reverses that source transform only when a positive adjusted counter proves the clamp lost no -information. Fabric's normalized input includes cache and output includes -reasoning. Kilo zero-fills omitted dimensions, so a stored cache or reasoning +information. NeMo Fabric's normalized input includes cache and output includes +reasoning. Kilo Code zero-fills omitted dimensions, so a stored cache or reasoning zero cannot establish provider field presence. A positive cache count stays -independently useful even when input is unknown. Kilo's derived `total` and +independently useful even when input is unknown. Kilo Code's derived `total` and mixed provider/catalog `cost` do not become reported totals or `cost_usd`. A transient event observer detects actual retries hidden by the final message @@ -150,7 +150,7 @@ Actual background tasks leave complete totals unknown because their child work and injected parent reply can outlive the invocation. Settled owned steps remain observed. Late work never moves into a later invocation. Unread, ambiguous or deeper child sessions also leave gaps. Merely enabling delegation -does not. The adapter sets a fixed session title to avoid Kilo's separate title +does not. The adapter sets a fixed session title to avoid Kilo Code's separate title model call. When prompt communication fails, the adapter uses the same bounded history From e3a2355b585e1eb231a7395cb1f4e7945c76d12b Mon Sep 17 00:00:00 2001 From: mnajafian-nv Date: Fri, 23 Oct 2026 01:41:33 -0700 Subject: [PATCH 25/29] 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 26/29] 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 27/29] 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 28/29] 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 29/29] 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