From bcb230b7268e3f7cefa64155ea839999ef5f2da2 Mon Sep 17 00:00:00 2001 From: mnajafian-nv Date: Thu, 15 Oct 2026 15:15:05 -0700 Subject: [PATCH 01/68] feat: report per-invocation token usage from Hermes The Hermes adapter returned no usage, so every Hermes run read unknown tokens. Hermes keeps token counters per session, and the adapter keeps one session per runtime, so each turn's result holds session totals. Each invocation now reports the change since the previous invocation's totals: input_tokens with cache reads and cache writes included, cached_input_tokens, output_tokens and total_tokens. The baseline is 0 for a new Hermes agent, is cleared while an invocation runs, and is set again only from a returned result. After an invocation that raised, was cancelled or timed out, or returned no counters, the next invocation reports no counters instead of a carried-over or whole-session total. A counter that went down is unknown for that invocation. This follows the rule the Codex adapter already uses for thread totals. Signed-off-by: mnajafian-nv --- adapters/python/hermes/README.md | 13 +++ .../nemo_fabric_adapters/hermes/adapter.py | 69 ++++++++++++ tests/adapters/test_hermes_adapter.py | 104 ++++++++++++++++++ tests/e2e/test_hermes_e2e.py | 24 ++++ 4 files changed, 210 insertions(+) diff --git a/adapters/python/hermes/README.md b/adapters/python/hermes/README.md index 3c0fc0fec..824d73c84 100644 --- a/adapters/python/hermes/README.md +++ b/adapters/python/hermes/README.md @@ -112,6 +112,19 @@ a child under `runtimes/`, so invocations in one NeMo Fabric runtime share state without sharing config or the session database with another runtime. +## Token Usage + +Hermes counts tokens per session, and the adapter keeps one session per +runtime, so each result's `usage` is the change since the previous +invocation: `input_tokens` with cache reads and cache writes included, +`cached_input_tokens`, `output_tokens`, and `total_tokens`. A counter that +Hermes does not report, or that went down because its session counters were +reset, is left out for that invocation. After an invocation that raised, was +cancelled, or timed out, or whose result held no counters, the adapter has no +baseline. The next invocation then reports no token counters rather than a +carried-over or whole-session total, and the one after it reports its own +change again. + ## Execution Model Each NeMo Fabric runtime starts one local adapter host, constructs one Hermes Agent diff --git a/adapters/python/hermes/src/nemo_fabric_adapters/hermes/adapter.py b/adapters/python/hermes/src/nemo_fabric_adapters/hermes/adapter.py index 653304993..888d700b5 100755 --- a/adapters/python/hermes/src/nemo_fabric_adapters/hermes/adapter.py +++ b/adapters/python/hermes/src/nemo_fabric_adapters/hermes/adapter.py @@ -27,6 +27,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 RuntimeContext from nemo_fabric_adapters.common import instructions as common_instructions from nemo_fabric_adapters.common import lifecycle @@ -66,6 +67,9 @@ def __init__(self) -> None: self._relay_plugin_config: dict[str, Any] | None = None self._relay_plugin_config_path: Path | None = None self._active_invoke_task: asyncio.Task[tuple[dict[str, Any], str]] | None = None + # Hermes reports session totals; each invocation reports the change + # from this baseline. A counter missing from it has no baseline. + self._session_token_totals: dict[str, int] = _fresh_session_totals() async def start(self, payload: dict[str, Any]) -> None: if self._started: @@ -217,6 +221,7 @@ async def start(self, payload: dict[str, Any]) -> None: session_db=self._session_db, ) ) + self._session_token_totals = _fresh_session_totals() self._agent_config = agent_config self._started = True except BaseException: @@ -284,6 +289,10 @@ def run_hermes_turn() -> tuple[dict[str, Any], str]: "hermes_invocation_in_progress", "Hermes runtime already has an active invocation", ) + # Until this invocation's result arrives the baseline is unknown, so a + # failed, cancelled or timed-out invocation never hands its tokens to + # the next one. + previous_totals, self._session_token_totals = self._session_token_totals, {} invoke_task = asyncio.create_task(asyncio.to_thread(run_hermes_turn)) self._active_invoke_task = invoke_task @@ -368,10 +377,12 @@ def clear_active_invoke_task( code="hermes_reported_failure", message="Hermes did not complete the invocation", ) + usage, self._session_token_totals = _invocation_usage(result, previous_totals) return AgentRunResult( status=AgentRunStatus.FAILED if failed else AgentRunStatus.SUCCEEDED, output=output, error=error, + usage=usage, ) async def _authenticate_mcp_servers(self) -> None: @@ -584,6 +595,64 @@ def _invoke_hermes_turn( return result, hermes_stdout.getvalue() +# Hermes' per-session token counters in a turn's result. Its input_tokens +# excludes cache reads and cache writes. +_SESSION_TOKEN_KEYS = ( + "input_tokens", + "cache_read_tokens", + "cache_write_tokens", + "output_tokens", +) + + +def _fresh_session_totals() -> dict[str, int]: + """Return the baseline of a new Hermes agent, whose counters start at 0.""" + return dict.fromkeys(_SESSION_TOKEN_KEYS, 0) + + +def _invocation_usage( + result: dict[str, Any], previous_totals: dict[str, int] +) -> tuple[AgentUsage | None, dict[str, int]]: + """Return this invocation's usage and the session totals to keep as baseline. + + Hermes keeps one session across a runtime's invocations, so each counter is + the change since ``previous_totals``. A counter that is missing from the + result or from ``previous_totals`` (no baseline: the previous invocation + failed or its result held no counters), or that went down (the session + counters were reset), is unknown for this invocation. The usage therefore + never holds a whole-session total or another invocation's tokens. + """ + totals = { + key: value + for key in _SESSION_TOKEN_KEYS + if isinstance((value := result.get(key)), int) + and not isinstance(value, bool) + and value >= 0 + } + change = { + key: value - previous_totals[key] + for key, value in totals.items() + if key in previous_totals and value >= previous_totals[key] + } + fields: dict[str, Any] = {} + if all( + key in change + for key in ("input_tokens", "cache_read_tokens", "cache_write_tokens") + ): + fields["input_tokens"] = ( + change["input_tokens"] + + change["cache_read_tokens"] + + change["cache_write_tokens"] + ) + fields["cached_input_tokens"] = change["cache_read_tokens"] + fields["input_tokens_include_cache"] = True + if "output_tokens" in change: + fields["output_tokens"] = change["output_tokens"] + if "input_tokens" in fields and "output_tokens" in fields: + fields["total_tokens"] = fields["input_tokens"] + fields["output_tokens"] + return (AgentUsage(**fields) if fields else None), totals + + def filter_supported_kwargs(callable_obj: Any, **kwargs: Any) -> dict[str, Any]: signature = inspect.signature(callable_obj.__init__) supported = set(signature.parameters) diff --git a/tests/adapters/test_hermes_adapter.py b/tests/adapters/test_hermes_adapter.py index 9a466f0b1..5c2a11a8d 100644 --- a/tests/adapters/test_hermes_adapter.py +++ b/tests/adapters/test_hermes_adapter.py @@ -1371,12 +1371,20 @@ async def test_persistent_runtime_reuses_hermes_agent_session_and_history( "completed": True, "failed": False, "messages": first_messages, + "input_tokens": 100, + "cache_read_tokens": 0, + "cache_write_tokens": 50, + "output_tokens": 20, }, { "response": "second response", "completed": True, "failed": False, "messages": second_messages, + "input_tokens": 130, + "cache_read_tokens": 50, + "cache_write_tokens": 50, + "output_tokens": 35, }, ] mock_ai_agent_type = MagicMock(spec=AIAgent, return_value=mock_ai_agent) @@ -1506,6 +1514,14 @@ async def test_persistent_runtime_reuses_hermes_agent_session_and_history( assert first.output["response"] == "first response" assert second.output["response"] == "second response" assert "session_id" not in second.output + # Hermes reports session totals; each invocation reports its own change, + # with input counting its cache reads and writes. + assert first.usage is not None and second.usage is not None + assert (first.usage.input_tokens, first.usage.cached_input_tokens) == (150, 0) + assert first.usage.output_tokens == 20 + assert (second.usage.input_tokens, second.usage.cached_input_tokens) == (80, 50) + assert second.usage.input_tokens_include_cache is True + assert (second.usage.output_tokens, second.usage.total_tokens) == (15, 95) assert Path(second.output["hermes_home"]) == ( tmp_path / "artifacts" @@ -1516,6 +1532,94 @@ async def test_persistent_runtime_reuses_hermes_agent_session_and_history( ) +async def test_hermes_invoke_after_failed_or_counterless_turn_reports_no_usage( + started_hermes_runtime, +): + runtime, mock_turn = started_hermes_runtime + + def turn(**counters: int) -> tuple[dict[str, object], str]: + return ( + {"final_response": "done", "completed": True, "messages": [], **counters}, + "", + ) + + def totals(tokens: int) -> dict[str, int]: + return { + "input_tokens": tokens, + "cache_read_tokens": 0, + "cache_write_tokens": 0, + "output_tokens": tokens, + } + + mock_turn.side_effect = [ + turn(**totals(100)), + turn(), + turn(**totals(300)), + RuntimeError("turn failed"), + turn(**totals(700)), + turn(**totals(800)), + ] + invocation = _invocation( + { + "runtime_context": _runtime_context( + runtime_id="runtime-incomplete" + ).to_mapping(), + "request": {"input": "next"}, + } + ) + + first = await runtime.invoke(*invocation) + counterless = await runtime.invoke(*invocation) + after_counterless = await runtime.invoke(*invocation) + with pytest.raises(RuntimeError): + await runtime.invoke(*invocation) + after_failure = await runtime.invoke(*invocation) + recovered = await runtime.invoke(*invocation) + + assert first.usage is not None + assert (first.usage.input_tokens, first.usage.output_tokens) == (100, 100) + # No baseline: never the session total or a carried-over change. + assert counterless.usage is None + assert after_counterless.usage is None + assert after_failure.usage is None + assert recovered.usage is not None + assert (recovered.usage.input_tokens, recovered.usage.output_tokens) == (100, 100) + + +def test_invocation_usage_with_lower_or_missing_counters_leaves_them_unknown(): + usage, totals = adapter._invocation_usage( + { + "input_tokens": 10, + "cache_read_tokens": 0, + "cache_write_tokens": 0, + "output_tokens": 4, + }, + { + "input_tokens": 50, + "cache_read_tokens": 0, + "cache_write_tokens": 0, + "output_tokens": 1, + }, + ) + missing, _ = adapter._invocation_usage({"completed": True}, {}) + no_baseline, _ = adapter._invocation_usage( + { + "input_tokens": 10, + "cache_read_tokens": 0, + "cache_write_tokens": 0, + "output_tokens": 4, + }, + {}, + ) + + assert usage is not None + assert usage.input_tokens is None + assert usage.output_tokens == 3 + assert totals["input_tokens"] == 10 + assert missing is None + assert no_baseline is None + + @pytest.mark.parametrize( ("request_input", "expected_message"), [ diff --git a/tests/e2e/test_hermes_e2e.py b/tests/e2e/test_hermes_e2e.py index 20472ff62..dcb063615 100644 --- a/tests/e2e/test_hermes_e2e.py +++ b/tests/e2e/test_hermes_e2e.py @@ -112,6 +112,30 @@ async def test_hermes_persistent_host_with_relay( ) +@pytest.mark.usefixtures("mock_nvidia_api_key") +async def test_hermes_persistent_host_reports_usage_for_each_invocation( + code_review_agent_dir: Path, + api_server: str, +): + os.environ["ADAPTER_PYTHON"] = sys.executable + config = hermes_config() + config.models["default"].base_url = f"{api_server}/v1" + + async with await Fabric().start_runtime( + config, base_dir=code_review_agent_dir + ) as runtime: + first = await runtime.invoke(input="first") + second = await runtime.invoke(input="second") + + # The mock endpoint bills 0 tokens per call, so each invocation's change + # since the previous one is a measured 0, not unknown. + for turn in (first, second): + assert turn["status"] == "succeeded", turn.to_mapping() + assert turn.usage is not None, turn.to_mapping() + assert turn.usage.input_tokens == 0, turn.to_mapping() + assert turn.usage.output_tokens == 0, turn.to_mapping() + + @pytest.mark.usefixtures("mock_nvidia_api_key") async def test_env_secrets_in_headers( code_review_agent_dir: Path, From 02c2fc6af7f9b661ffa53387163279fbc0642355 Mon Sep 17 00:00:00 2001 From: mnajafian-nv Date: Thu, 15 Oct 2026 15:15:06 -0700 Subject: [PATCH 02/68] feat: report per-invocation token usage from Pi The Pi adapter returned no usage although Pi reports token counts for every call it bills. Pi reports usage per call, so the adapter now sums what Pi bills during the invocation, the way Pi's own usage totals count it: assistant messages, tool results that report usage, finished compaction summaries and usage entries. The result holds input_tokens with cache reads and cache writes included, cached_input_tokens, output_tokens and total_tokens. A call without a complete usage report, or a compaction that was aborted or finished without usage, leaves the invocation without usage, so a partial sum is never reported as the invocation's total. Pi's cost figures come from its model catalog's rates, so they are not reported as cost_usd. Signed-off-by: mnajafian-nv --- adapters/typescript/pi/README.md | 11 ++++ adapters/typescript/pi/src/pi-sdk.ts | 69 +++++++++++++++++++- adapters/typescript/pi/src/runtime.ts | 44 ++++++++++++- adapters/typescript/pi/test/pi-sdk.test.mjs | 36 ++++++++++ adapters/typescript/pi/test/runtime.test.mjs | 26 ++++++++ 5 files changed, 184 insertions(+), 2 deletions(-) diff --git a/adapters/typescript/pi/README.md b/adapters/typescript/pi/README.md index 0dbafbca1..c3348731c 100644 --- a/adapters/typescript/pi/README.md +++ b/adapters/typescript/pi/README.md @@ -212,6 +212,17 @@ This Relay-backed path runs the adapter's ordinary `invoke` operation. It is independent of native OpenAI streaming, so the adapter descriptor's `capabilities.streaming` value remains `false`. +## Token Usage + +Each result's `usage` sums the tokens Pi bills during the invocation, counted +the way Pi's own usage totals count them: assistant messages, tool results that +report usage, finished compaction summaries, and usage entries such as cache +warm-ups. It holds `input_tokens` with cache reads and cache writes included, +`cached_input_tokens`, `output_tokens`, and `total_tokens`. When one of those +calls carries no complete usage, or a compaction was aborted or reported no +usage, the result has no `usage`. Pi computes its cost figures from the model +catalog's rates, so the adapter does not report them as `cost_usd`. + ## Custom Tool Modules The adapter accepts custom tools with `kind: "module"`. The `ref` is a diff --git a/adapters/typescript/pi/src/pi-sdk.ts b/adapters/typescript/pi/src/pi-sdk.ts index 18214ac00..20e146961 100644 --- a/adapters/typescript/pi/src/pi-sdk.ts +++ b/adapters/typescript/pi/src/pi-sdk.ts @@ -28,7 +28,7 @@ import { type PiRelayControllerFactory, type PiRelayRuntime, } from "./relay.js"; -import type { PiPromptOutcome, PiSessionFactory, PiSessionHandle } from "./runtime.js"; +import type { PiPromptOutcome, PiSessionFactory, PiSessionHandle, PiTokenUsage } from "./runtime.js"; // ProviderConfigInput isn't re-exported from pi's public entrypoint; derive it from // registerProvider's param type (ModelRuntime's constructor is private, so no InstanceType). @@ -69,6 +69,70 @@ interface PiToolFactoryContext { workspace: string; } +/** + * Add one billed call's usage to a prompt's running total. A call without a + * complete usage report makes the whole total unknown. + */ +export function addUsage(total: PiTokenUsage | undefined, usage: unknown): PiTokenUsage | undefined { + if (total === undefined || typeof usage !== "object" || usage === null) { + return undefined; + } + const counts = usage as Record; + const count = (name: keyof PiTokenUsage): number | undefined => { + const value = counts[name]; + return typeof value === "number" && Number.isSafeInteger(value) && value >= 0 ? value : undefined; + }; + const input = count("input"); + const output = count("output"); + const cacheRead = count("cacheRead"); + const cacheWrite = count("cacheWrite"); + if (input === undefined || output === undefined || cacheRead === undefined || cacheWrite === undefined) { + return undefined; + } + return { + input: total.input + input, + output: total.output + output, + cacheRead: total.cacheRead + cacheRead, + cacheWrite: total.cacheWrite + cacheWrite, + }; +} + +/** + * Add the tokens one session event billed to a prompt's running total, as + * Pi's own usage totals count them: an assistant message, a tool result that + * reports usage, a finished compaction's summary call, and a usage entry such + * as a cache warm-up. A compaction that was aborted or finished without a + * usage report makes the total unknown, because its summary call may have + * been billed. + */ +export function addEventUsage(total: PiTokenUsage | undefined, event: unknown): PiTokenUsage | undefined { + const fields = (typeof event === "object" && event !== null ? event : {}) as { + type?: unknown; + message?: { role?: unknown; usage?: unknown }; + aborted?: unknown; + result?: { usage?: unknown }; + entry?: { type?: unknown; usage?: unknown }; + }; + switch (fields.type) { + case "message_end": + if (fields.message?.role === "assistant") { + return addUsage(total, fields.message.usage); + } + if (fields.message?.role === "toolResult" && fields.message.usage !== undefined) { + return addUsage(total, fields.message.usage); + } + return total; + case "compaction_end": + return fields.aborted === false && fields.result !== undefined + ? addUsage(total, fields.result.usage) + : undefined; + case "entry_appended": + return fields.entry?.type === "usage" ? addUsage(total, fields.entry.usage) : total; + default: + return total; + } +} + export function withCustomBaseUrl( catalogModel: T, baseUrl: string | null | undefined, @@ -892,6 +956,7 @@ class PiSdkSessionHandle implements PiSessionHandle { let finalAssistant: | { role: "assistant"; content: unknown; stopReason: string; errorMessage?: string } | undefined; + let usage: PiTokenUsage | undefined = { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 }; const unsubscribe = this.session.subscribe((event) => { if (event.type === "turn_start") { turnStarted = true; @@ -899,6 +964,7 @@ class PiSdkSessionHandle implements PiSessionHandle { if (event.type === "message_end" && event.message.role === "assistant") { finalAssistant = event.message; } + usage = addEventUsage(usage, event); }); try { await this.session.prompt(text, { @@ -919,6 +985,7 @@ class PiSdkSessionHandle implements PiSessionHandle { stopReason: finalAssistant?.stopReason, errorMessage: finalAssistant?.errorMessage, shutdownRequested: this.state.shutdownRequested, + usage: finalAssistant === undefined ? undefined : usage, }; } diff --git a/adapters/typescript/pi/src/runtime.ts b/adapters/typescript/pi/src/runtime.ts index 5ce0838c9..febac8dad 100644 --- a/adapters/typescript/pi/src/runtime.ts +++ b/adapters/typescript/pi/src/runtime.ts @@ -5,7 +5,13 @@ // to a PiSessionHandle and normalized results while keeping SDK construction // behind a factory for focused lifecycle testing. -import type { AgentRunRequest, AgentRunResult, JsonObject, RuntimeContext } from "nemo-fabric-adapter-contract"; +import type { + AgentRunRequest, + AgentRunResult, + AgentUsage, + JsonObject, + RuntimeContext, +} from "nemo-fabric-adapter-contract"; import { LifecycleError, type AdapterRuntime, type AdapterStartInput } from "nemo-fabric-adapters-common"; import { collectRelayArtifacts, type RelayArtifact } from "./relay-config.js"; @@ -13,6 +19,17 @@ import type { PiRelayRuntime } from "./relay.js"; export type PiStopReason = "stop" | "length" | "toolUse" | "error" | "aborted" | string; +/** + * Token counts Pi billed during one prompt. Pi's `input` excludes cache reads + * and writes. + */ +export interface PiTokenUsage { + input: number; + output: number; + cacheRead: number; + cacheWrite: number; +} + export interface PiPromptOutcome { accepted: boolean; turnStarted: boolean; @@ -21,6 +38,8 @@ export interface PiPromptOutcome { stopReason?: PiStopReason; errorMessage?: string; shutdownRequested?: boolean; + /** Absent when any billed call of the prompt carried no complete usage. */ + usage?: PiTokenUsage; } export interface PiSessionHandle { @@ -42,14 +61,31 @@ function failed(code: string, message: string): AgentRunResult { }; } +/** The invocation's usage, with input counting its cache reads and writes. */ +function agentUsage(usage: PiTokenUsage | undefined): AgentUsage | undefined { + if (usage === undefined) { + return undefined; + } + const input = usage.input + usage.cacheRead + usage.cacheWrite; + return { + input_tokens: input, + cached_input_tokens: usage.cacheRead, + input_tokens_include_cache: true, + output_tokens: usage.output, + total_tokens: input + usage.output, + }; +} + async function withRelayOutput( result: AgentRunResult, relay: PiRelayRuntime | undefined, turnStarted: boolean, turnCount: number, + usage?: AgentUsage, ): Promise { const annotated = { ...result, + ...(usage === undefined ? {} : { usage }), extensions: { ...result.extensions, pi_turn_count: turnCount, pi_turn_started: turnStarted }, }; if (relay === undefined) { @@ -113,6 +149,7 @@ export class PiAdapterRuntime implements AdapterRuntime { const relay = this.session.relay; const outcome = await this.session.prompt(request.input); + const usage = agentUsage(outcome.usage); if ( !Number.isInteger(outcome.turnCount) || outcome.turnCount < this.turnCount || @@ -127,6 +164,7 @@ export class PiAdapterRuntime implements AdapterRuntime { relay, outcome.turnStarted, outcome.turnCount, + usage, ); } if (outcome.shutdownRequested || outcome.stopReason === "aborted") { @@ -148,6 +186,7 @@ export class PiAdapterRuntime implements AdapterRuntime { relay, outcome.turnStarted, outcome.turnCount, + usage, ); } if (outcome.stopReason === "error") { @@ -156,6 +195,7 @@ export class PiAdapterRuntime implements AdapterRuntime { relay, outcome.turnStarted, outcome.turnCount, + usage, ); } if (outcome.text === undefined || outcome.text.length === 0) { @@ -164,6 +204,7 @@ export class PiAdapterRuntime implements AdapterRuntime { relay, outcome.turnStarted, outcome.turnCount, + usage, ); } return withRelayOutput( @@ -171,6 +212,7 @@ export class PiAdapterRuntime implements AdapterRuntime { relay, outcome.turnStarted, outcome.turnCount, + usage, ); } diff --git a/adapters/typescript/pi/test/pi-sdk.test.mjs b/adapters/typescript/pi/test/pi-sdk.test.mjs index 9095722ac..0a0c58ce9 100644 --- a/adapters/typescript/pi/test/pi-sdk.test.mjs +++ b/adapters/typescript/pi/test/pi-sdk.test.mjs @@ -9,6 +9,8 @@ import test from "node:test"; import { createServer } from "node:http"; import { + addEventUsage, + addUsage, buildCatalogModel, modelAwareCompactionReserveTokens, PiSdkSessionFactory, @@ -1440,3 +1442,37 @@ test("accepts a well-formed extensions.cost", async () => { await rm(workspace, { recursive: true, force: true }); } }); + +test("sums each assistant message's usage and drops an incomplete one", () => { + const zero = { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 }; + const first = addUsage(zero, { input: 10, output: 2, cacheRead: 5, cacheWrite: 1, totalTokens: 18 }); + const second = addUsage(first, { input: 3, output: 4, cacheRead: 0, cacheWrite: 0, reasoning: 2 }); + + assert.deepEqual(second, { input: 13, output: 6, cacheRead: 5, cacheWrite: 1 }); + assert.equal(addUsage(second, { input: 1, output: 1 }), undefined); + assert.equal(addUsage(undefined, { input: 1, output: 1, cacheRead: 0, cacheWrite: 0 }), undefined); +}); + +test("counts the tokens Pi bills outside assistant messages", () => { + const usage = (tokens) => ({ input: tokens, output: 1, cacheRead: 0, cacheWrite: 0 }); + const events = [ + { type: "message_end", message: { role: "assistant", usage: usage(10) } }, + { type: "compaction_end", aborted: false, result: { summary: "s", usage: usage(100) } }, + { type: "message_end", message: { role: "toolResult", usage: usage(20) } }, + { type: "message_end", message: { role: "toolResult" } }, + { type: "entry_appended", entry: { type: "usage", usage: usage(5) } }, + { type: "entry_appended", entry: { type: "custom" } }, + { type: "turn_start" }, + ]; + + const total = events.reduce(addEventUsage, { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 }); + + assert.deepEqual(total, { input: 135, output: 4, cacheRead: 0, cacheWrite: 0 }); +}); + +test("a compaction without a usage report leaves the prompt's tokens unknown", () => { + const zero = { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 }; + + assert.equal(addEventUsage(zero, { type: "compaction_end", aborted: false, result: { summary: "s" } }), undefined); + assert.equal(addEventUsage(zero, { type: "compaction_end", aborted: true }), undefined); +}); diff --git a/adapters/typescript/pi/test/runtime.test.mjs b/adapters/typescript/pi/test/runtime.test.mjs index 1726e6201..9269ee643 100644 --- a/adapters/typescript/pi/test/runtime.test.mjs +++ b/adapters/typescript/pi/test/runtime.test.mjs @@ -557,3 +557,29 @@ test("retries Relay cleanup when gateway shutdown fails", async () => { assert.equal(sessionStopAttempts, 2); assert.equal(relayStopAttempts, 2); }); + +test("reports the invocation's tokens with cache reads and writes inside the input", async () => { + const usage = { input: 100, output: 40, cacheRead: 300, cacheWrite: 20 }; + const runtime = new PiAdapterRuntime({ + async create() { + return { + async prompt() { + return { accepted: true, turnStarted: true, turnCount: 1, text: "ok", stopReason: "stop", usage }; + }, + async stop() {}, + }; + }, + }); + + await runtime.start(startInput()); + const result = await runtime.invoke({ input: "one" }, context); + await runtime.stop(); + + assert.deepEqual(result.usage, { + input_tokens: 420, + cached_input_tokens: 300, + input_tokens_include_cache: true, + output_tokens: 40, + total_tokens: 460, + }); +}); From e5eae71a5fecebe2cbc4d00fd72e955aab222d00 Mon Sep 17 00:00:00 2001 From: mnajafian-nv Date: Thu, 15 Oct 2026 15:42:19 -0700 Subject: [PATCH 03/68] fix: leave Pi usage unknown after an aborted call Pi keeps no usage for a provider call it stopped, so an aborted assistant message reports zero tokens although its input may already have been billed. Adding that zero made a cancelled invocation report a complete-looking total that left out the aborted call. An aborted assistant message now leaves the invocation without usage, the same rule as an aborted compaction. The comments and README now say "records" instead of "bills", because the adapter sums what Pi recorded, and the README lists every case in which the result has no usage. Signed-off-by: mnajafian-nv --- adapters/typescript/pi/README.md | 9 +++++---- adapters/typescript/pi/src/pi-sdk.ts | 16 ++++++++-------- adapters/typescript/pi/src/runtime.ts | 6 +++--- adapters/typescript/pi/test/pi-sdk.test.mjs | 10 ++++++++++ 4 files changed, 26 insertions(+), 15 deletions(-) diff --git a/adapters/typescript/pi/README.md b/adapters/typescript/pi/README.md index c3348731c..a8fe6bf59 100644 --- a/adapters/typescript/pi/README.md +++ b/adapters/typescript/pi/README.md @@ -214,13 +214,14 @@ independent of native OpenAI streaming, so the adapter descriptor's ## Token Usage -Each result's `usage` sums the tokens Pi bills during the invocation, counted +Each result's `usage` sums the tokens Pi records during the invocation, counted the way Pi's own usage totals count them: assistant messages, tool results that report usage, finished compaction summaries, and usage entries such as cache warm-ups. It holds `input_tokens` with cache reads and cache writes included, -`cached_input_tokens`, `output_tokens`, and `total_tokens`. When one of those -calls carries no complete usage, or a compaction was aborted or reported no -usage, the result has no `usage`. Pi computes its cost figures from the model +`cached_input_tokens`, `output_tokens`, and `total_tokens`. The result has no +`usage` when Pi produced no assistant message, when one of those calls carries +no complete usage, when an assistant message was aborted, or when a compaction +was aborted or reported no usage. Pi computes its cost figures from the model catalog's rates, so the adapter does not report them as `cost_usd`. ## Custom Tool Modules diff --git a/adapters/typescript/pi/src/pi-sdk.ts b/adapters/typescript/pi/src/pi-sdk.ts index 20e146961..1f6d76f71 100644 --- a/adapters/typescript/pi/src/pi-sdk.ts +++ b/adapters/typescript/pi/src/pi-sdk.ts @@ -70,8 +70,8 @@ interface PiToolFactoryContext { } /** - * Add one billed call's usage to a prompt's running total. A call without a - * complete usage report makes the whole total unknown. + * Add one call's usage to a prompt's running total. A call without a complete + * usage report makes the whole total unknown. */ export function addUsage(total: PiTokenUsage | undefined, usage: unknown): PiTokenUsage | undefined { if (total === undefined || typeof usage !== "object" || usage === null) { @@ -98,17 +98,17 @@ export function addUsage(total: PiTokenUsage | undefined, usage: unknown): PiTok } /** - * Add the tokens one session event billed to a prompt's running total, as + * Add the tokens one session event recorded to a prompt's running total, as * Pi's own usage totals count them: an assistant message, a tool result that * reports usage, a finished compaction's summary call, and a usage entry such - * as a cache warm-up. A compaction that was aborted or finished without a - * usage report makes the total unknown, because its summary call may have - * been billed. + * as a cache warm-up. An aborted assistant message, or a compaction that was + * aborted or finished without a usage report, makes the total unknown: Pi + * keeps no usage for a call it stopped, which may still have been billed. */ export function addEventUsage(total: PiTokenUsage | undefined, event: unknown): PiTokenUsage | undefined { const fields = (typeof event === "object" && event !== null ? event : {}) as { type?: unknown; - message?: { role?: unknown; usage?: unknown }; + message?: { role?: unknown; stopReason?: unknown; usage?: unknown }; aborted?: unknown; result?: { usage?: unknown }; entry?: { type?: unknown; usage?: unknown }; @@ -116,7 +116,7 @@ export function addEventUsage(total: PiTokenUsage | undefined, event: unknown): switch (fields.type) { case "message_end": if (fields.message?.role === "assistant") { - return addUsage(total, fields.message.usage); + return fields.message.stopReason === "aborted" ? undefined : addUsage(total, fields.message.usage); } if (fields.message?.role === "toolResult" && fields.message.usage !== undefined) { return addUsage(total, fields.message.usage); diff --git a/adapters/typescript/pi/src/runtime.ts b/adapters/typescript/pi/src/runtime.ts index febac8dad..f1dfd30a5 100644 --- a/adapters/typescript/pi/src/runtime.ts +++ b/adapters/typescript/pi/src/runtime.ts @@ -20,8 +20,8 @@ import type { PiRelayRuntime } from "./relay.js"; export type PiStopReason = "stop" | "length" | "toolUse" | "error" | "aborted" | string; /** - * Token counts Pi billed during one prompt. Pi's `input` excludes cache reads - * and writes. + * Token counts Pi recorded during one prompt. Pi's `input` excludes cache + * reads and writes. */ export interface PiTokenUsage { input: number; @@ -38,7 +38,7 @@ export interface PiPromptOutcome { stopReason?: PiStopReason; errorMessage?: string; shutdownRequested?: boolean; - /** Absent when any billed call of the prompt carried no complete usage. */ + /** Absent when any call of the prompt carried no complete usage. */ usage?: PiTokenUsage; } diff --git a/adapters/typescript/pi/test/pi-sdk.test.mjs b/adapters/typescript/pi/test/pi-sdk.test.mjs index 0a0c58ce9..037843961 100644 --- a/adapters/typescript/pi/test/pi-sdk.test.mjs +++ b/adapters/typescript/pi/test/pi-sdk.test.mjs @@ -1476,3 +1476,13 @@ test("a compaction without a usage report leaves the prompt's tokens unknown", ( assert.equal(addEventUsage(zero, { type: "compaction_end", aborted: false, result: { summary: "s" } }), undefined); assert.equal(addEventUsage(zero, { type: "compaction_end", aborted: true }), undefined); }); + +test("an aborted assistant message leaves the prompt's tokens unknown", () => { + const zero = { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 }; + const usage = { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 }; + + assert.equal( + addEventUsage(zero, { type: "message_end", message: { role: "assistant", stopReason: "aborted", usage } }), + undefined, + ); +}); From b0dd2f770e0d95a02e4e629c8ba92c8bcf2621b0 Mon Sep 17 00:00:00 2001 From: mnajafian-nv Date: Thu, 15 Oct 2026 15:42:21 -0700 Subject: [PATCH 04/68] fix: bound Hermes usage counters and state what they cover A Hermes counter above the 64-bit range would make AgentUsage raise and turn a finished invocation into an error, so such a counter is now ignored like any other invalid value, matching the Codex adapter. Hermes records auxiliary calls such as context-compression summaries outside its session counters, so the README now says the usage covers the main model's calls. A test now checks that the invocation after a cancelled one reports no usage, which is the rule that keeps a cancelled turn's tokens out of the next result. Signed-off-by: mnajafian-nv --- adapters/python/hermes/README.md | 11 ++++++----- .../hermes/src/nemo_fabric_adapters/hermes/adapter.py | 8 ++++---- tests/adapters/test_hermes_adapter.py | 6 ++++++ 3 files changed, 16 insertions(+), 9 deletions(-) diff --git a/adapters/python/hermes/README.md b/adapters/python/hermes/README.md index 824d73c84..69f7bd90a 100644 --- a/adapters/python/hermes/README.md +++ b/adapters/python/hermes/README.md @@ -114,12 +114,13 @@ runtime. ## Token Usage -Hermes counts tokens per session, and the adapter keeps one session per -runtime, so each result's `usage` is the change since the previous +Hermes counts its main model's tokens per session, and the adapter keeps one +session per runtime, so each result's `usage` is the change since the previous invocation: `input_tokens` with cache reads and cache writes included, -`cached_input_tokens`, `output_tokens`, and `total_tokens`. A counter that -Hermes does not report, or that went down because its session counters were -reset, is left out for that invocation. After an invocation that raised, was +`cached_input_tokens`, `output_tokens`, and `total_tokens`. Hermes records +auxiliary calls, such as context-compression summaries, separately from these +counters, so they are not included. A counter that Hermes does not report, or +that went down, is left out for that invocation. After an invocation that raised, was cancelled, or timed out, or whose result held no counters, the adapter has no baseline. The next invocation then reports no token counters rather than a carried-over or whole-session total, and the one after it reports its own diff --git a/adapters/python/hermes/src/nemo_fabric_adapters/hermes/adapter.py b/adapters/python/hermes/src/nemo_fabric_adapters/hermes/adapter.py index 888d700b5..3539d7ce0 100755 --- a/adapters/python/hermes/src/nemo_fabric_adapters/hermes/adapter.py +++ b/adapters/python/hermes/src/nemo_fabric_adapters/hermes/adapter.py @@ -618,16 +618,16 @@ def _invocation_usage( Hermes keeps one session across a runtime's invocations, so each counter is the change since ``previous_totals``. A counter that is missing from the result or from ``previous_totals`` (no baseline: the previous invocation - failed or its result held no counters), or that went down (the session - counters were reset), is unknown for this invocation. The usage therefore - never holds a whole-session total or another invocation's tokens. + failed or its result held no counters), or that went down, is unknown for + this invocation. The usage therefore never holds a whole-session total or + another invocation's tokens. """ totals = { key: value for key in _SESSION_TOKEN_KEYS if isinstance((value := result.get(key)), int) and not isinstance(value, bool) - and value >= 0 + and 0 <= value <= (1 << 64) - 1 } change = { key: value - previous_totals[key] diff --git a/tests/adapters/test_hermes_adapter.py b/tests/adapters/test_hermes_adapter.py index 5c2a11a8d..5d380dee6 100644 --- a/tests/adapters/test_hermes_adapter.py +++ b/tests/adapters/test_hermes_adapter.py @@ -1712,6 +1712,10 @@ def run_turn(**_kwargs): "completed": True, "failed": False, "messages": [], + "input_tokens": calls * 100, + "cache_read_tokens": 0, + "cache_write_tokens": 0, + "output_tokens": calls * 10, }, "", ) @@ -1748,6 +1752,8 @@ def run_turn(**_kwargs): result = await runtime.invoke(*_invocation(invocation)) assert result.output["response"] == "turn-2" + # The cancelled turn left no baseline, so its tokens never reach this result. + assert result.usage is None await runtime.stop() From b9bf5d45b9fc04da1de4d4e462c52e871fe82b55 Mon Sep 17 00:00:00 2001 From: mnajafian-nv Date: Thu, 15 Oct 2026 23:44:46 -0700 Subject: [PATCH 05/68] fix: leave Hermes usage unknown when a provider response omits it Hermes counts a provider response without usage as an API call but adds no tokens to its session counters, so the adapter reported such an invocation as a measured 0, and a mixed invocation as a complete subtotal. Count those responses through the context compressor's note_usage_less_response hook, which both Hermes runtimes call, and report no usage for an invocation that had one. A response that reported zero tokens still reports 0. Signed-off-by: mnajafian-nv --- adapters/python/hermes/README.md | 14 +- .../nemo_fabric_adapters/hermes/adapter.py | 28 ++++ tests/adapters/test_hermes_adapter.py | 126 ++++++++++++++++++ 3 files changed, 162 insertions(+), 6 deletions(-) diff --git a/adapters/python/hermes/README.md b/adapters/python/hermes/README.md index 69f7bd90a..479a101df 100644 --- a/adapters/python/hermes/README.md +++ b/adapters/python/hermes/README.md @@ -119,12 +119,14 @@ session per runtime, so each result's `usage` is the change since the previous invocation: `input_tokens` with cache reads and cache writes included, `cached_input_tokens`, `output_tokens`, and `total_tokens`. Hermes records auxiliary calls, such as context-compression summaries, separately from these -counters, so they are not included. A counter that Hermes does not report, or -that went down, is left out for that invocation. After an invocation that raised, was -cancelled, or timed out, or whose result held no counters, the adapter has no -baseline. The next invocation then reports no token counters rather than a -carried-over or whole-session total, and the one after it reports its own -change again. +counters, so they are not included. Hermes adds no tokens for a provider +response that carried no usage, so an invocation with such a response has no +`usage`. A response that reported zero tokens still counts as 0. A counter that +Hermes does not report, or that went down, is left out for that invocation. +After an invocation that raised, was cancelled, or timed out, or whose result +held no counters, the adapter has no baseline. The next invocation then reports +no token counters rather than a carried-over or whole-session total, and the +one after it reports its own change again. ## Execution Model diff --git a/adapters/python/hermes/src/nemo_fabric_adapters/hermes/adapter.py b/adapters/python/hermes/src/nemo_fabric_adapters/hermes/adapter.py index 3539d7ce0..d2690a0b5 100755 --- a/adapters/python/hermes/src/nemo_fabric_adapters/hermes/adapter.py +++ b/adapters/python/hermes/src/nemo_fabric_adapters/hermes/adapter.py @@ -70,6 +70,7 @@ def __init__(self) -> None: # Hermes reports session totals; each invocation reports the change # from this baseline. A counter missing from it has no baseline. self._session_token_totals: dict[str, int] = _fresh_session_totals() + self._usage_less_responses = 0 async def start(self, payload: dict[str, Any]) -> None: if self._started: @@ -222,6 +223,7 @@ async def start(self, payload: dict[str, Any]) -> None: ) ) self._session_token_totals = _fresh_session_totals() + self._count_usage_less_responses(self._agent) self._agent_config = agent_config self._started = True except BaseException: @@ -293,6 +295,7 @@ def run_hermes_turn() -> tuple[dict[str, Any], str]: # failed, cancelled or timed-out invocation never hands its tokens to # the next one. previous_totals, self._session_token_totals = self._session_token_totals, {} + usage_less_before = self._usage_less_responses invoke_task = asyncio.create_task(asyncio.to_thread(run_hermes_turn)) self._active_invoke_task = invoke_task @@ -378,6 +381,10 @@ def clear_active_invoke_task( message="Hermes did not complete the invocation", ) usage, self._session_token_totals = _invocation_usage(result, previous_totals) + if self._usage_less_responses != usage_less_before: + # A provider response without usage added no tokens, so the + # change in the counters is an undercount, not this invocation's. + usage = None return AgentRunResult( status=AgentRunStatus.FAILED if failed else AgentRunStatus.SUCCEEDED, output=output, @@ -385,6 +392,27 @@ def clear_active_invoke_task( usage=usage, ) + def _count_usage_less_responses(self, agent: Any) -> None: + """Count the agent's provider responses that carried no usage. + + Hermes counts such a response as an API call but leaves its token + counters unchanged, so the counters cannot tell an unreported call + from a reported zero. Hermes reports each of these responses to the + context compressor's ``note_usage_less_response`` hook, so the + adapter wraps that hook on the agent it owns. + """ + compressor = agent.context_compressor + note_usage_less_response = getattr( + compressor, "note_usage_less_response", None + ) + + def count_and_note() -> None: + self._usage_less_responses += 1 + if callable(note_usage_less_response): + note_usage_less_response() + + compressor.note_usage_less_response = count_and_note + async def _authenticate_mcp_servers(self) -> None: if self._mcp_authentication_checked: return diff --git a/tests/adapters/test_hermes_adapter.py b/tests/adapters/test_hermes_adapter.py index 5d380dee6..f972fc1a1 100644 --- a/tests/adapters/test_hermes_adapter.py +++ b/tests/adapters/test_hermes_adapter.py @@ -16,6 +16,7 @@ from io import StringIO from pathlib import Path from types import ModuleType +from types import SimpleNamespace from unittest.mock import MagicMock import pytest @@ -874,6 +875,7 @@ async def test_runtime_start_discovers_mcp_tools_when_configured( mock_session_db = MagicMock(spec=SessionDB) mock_session_db_type = MagicMock(spec=SessionDB, return_value=mock_session_db) mock_ai_agent = MagicMock(spec=AIAgent) + mock_ai_agent.context_compressor = MagicMock() mock_ai_agent_type = MagicMock(spec=AIAgent, return_value=mock_ai_agent) monkeypatch.setattr( mock_ai_agent_type.__init__.__func__, @@ -1352,6 +1354,7 @@ async def test_persistent_runtime_reuses_hermes_agent_session_and_history( mock_session_db_type = MagicMock(spec=SessionDB, return_value=mock_session_db) mock_ai_agent = MagicMock(spec=AIAgent) + mock_ai_agent.context_compressor = MagicMock() mock_ai_agent.session_id = "runtime-fabric-123" mock_ai_agent.model = "test-model" mock_ai_agent.platform = "fabric" @@ -1620,6 +1623,129 @@ def test_invocation_usage_with_lower_or_missing_counters_leaves_them_unknown(): assert no_baseline is None +class _UsageRecordingAgent: + """Hermes agent stand-in whose turns run scripted provider responses + through Hermes' own usage accounting and context compressor.""" + + def __init__(self, turns: list[list[dict[str, int] | None]]) -> None: + from agent.context_compressor import ContextCompressor + + self._turns = turns + self.context_compressor = ContextCompressor( + model="test-model", quiet_mode=True + ) + self.provider = "custom" + self.api_mode = "chat_completions" + self.model = "test-model" + self.base_url = "http://127.0.0.1:1/v1" + self.api_key = "" + self.client = None + self.session_id = "runtime-usage" + self._session_db = None + self.quiet_mode = True + self.verbose_logging = False + for counter in ( + "api_calls", + "prompt_tokens", + "completion_tokens", + "total_tokens", + "input_tokens", + "output_tokens", + "cache_read_tokens", + "cache_write_tokens", + "reasoning_tokens", + "estimated_cost_usd", + ): + setattr(self, f"session_{counter}", 0) + + def run_conversation(self, user_message: str) -> dict[str, object]: + from agent.turn_usage import record_response_usage + + messages = [{"role": "user", "content": user_message}] + for usage in self._turns.pop(0): + record_response_usage( + self, + SimpleNamespace(usage=usage), + messages=messages, + api_call_count=1, + api_duration=0.0, + compression_attempts=0, + max_compression_attempts=3, + ) + return { + "final_response": "done", + "completed": True, + "messages": messages, + **{ + key: getattr(self, f"session_{key}") + for key in ( + "input_tokens", + "cache_read_tokens", + "cache_write_tokens", + "output_tokens", + ) + }, + } + + +async def test_hermes_usage_from_provider_responses_keeps_unreported_tokens_unknown(): + def reported(prompt: int, completion: int) -> dict[str, int]: + return { + "prompt_tokens": prompt, + "completion_tokens": completion, + "total_tokens": prompt + completion, + } + + agent = _UsageRecordingAgent( + [ + [reported(120, 7), reported(200, 3)], + [reported(0, 0)], + [None], + [reported(50, 5), None], + [reported(400, 40)], + ] + ) + agent_config = _agent_config( + {"models": {"default": {"provider": "test", "model": "test-model"}}} + ) + runtime = adapter.HermesRuntime() + runtime._started = True + runtime._agent_config = agent_config + runtime._model_config = agent_config.models["default"] + runtime._runtime_id = "runtime-usage" + runtime._agent = agent + runtime._count_usage_less_responses(agent) + invocation = _invocation( + { + "runtime_context": _runtime_context( + runtime_id="runtime-usage" + ).to_mapping(), + "request": {"input": "next"}, + } + ) + + two_calls = await runtime.invoke(*invocation) + reported_zero = await runtime.invoke(*invocation) + omitted = await runtime.invoke(*invocation) + mixed = await runtime.invoke(*invocation) + after_mixed = await runtime.invoke(*invocation) + + assert agent.session_api_calls == 7 + assert two_calls.usage is not None + assert (two_calls.usage.input_tokens, two_calls.usage.output_tokens) == (320, 10) + assert reported_zero.usage is not None + assert reported_zero.usage.input_tokens == 0 + assert reported_zero.usage.output_tokens == 0 + assert reported_zero.usage.total_tokens == 0 + # Hermes counts a response without usage as an API call but adds no + # tokens, so the invocation's tokens are unknown rather than 0 or partial. + assert omitted.usage is None + assert mixed.usage is None + assert after_mixed.usage is not None + assert after_mixed.usage.input_tokens == 400 + assert after_mixed.usage.output_tokens == 40 + + @pytest.mark.parametrize( ("request_input", "expected_message"), [ From fd68928d0394192123b89366f52a467ad97f4898 Mon Sep 17 00:00:00 2001 From: mnajafian-nv Date: Thu, 15 Oct 2026 23:44:47 -0700 Subject: [PATCH 06/68] fix: leave Pi usage unknown for unreported or retried calls Pi fills a model call's usage with zeros and overwrites it only when the provider reports usage, so a response without usage was summed as a measured 0. Treat a model call whose counts are all zero as unreported. Pi's compaction summary retry also returns only the final attempt's usage, and failed attempts emit no usage, so a retried compaction undercounted. Leave the invocation's usage unknown in both cases. Signed-off-by: mnajafian-nv --- adapters/typescript/pi/README.md | 14 ++- adapters/typescript/pi/src/pi-sdk.ts | 35 ++++-- adapters/typescript/pi/test/pi-sdk.test.mjs | 116 ++++++++++++++++++++ 3 files changed, 153 insertions(+), 12 deletions(-) diff --git a/adapters/typescript/pi/README.md b/adapters/typescript/pi/README.md index a8fe6bf59..0dfb5703d 100644 --- a/adapters/typescript/pi/README.md +++ b/adapters/typescript/pi/README.md @@ -218,11 +218,15 @@ Each result's `usage` sums the tokens Pi records during the invocation, counted the way Pi's own usage totals count them: assistant messages, tool results that report usage, finished compaction summaries, and usage entries such as cache warm-ups. It holds `input_tokens` with cache reads and cache writes included, -`cached_input_tokens`, `output_tokens`, and `total_tokens`. The result has no -`usage` when Pi produced no assistant message, when one of those calls carries -no complete usage, when an assistant message was aborted, or when a compaction -was aborted or reported no usage. Pi computes its cost figures from the model -catalog's rates, so the adapter does not report them as `cost_usd`. +`cached_input_tokens`, `output_tokens`, and `total_tokens`. Pi fills a model +call's usage with zeros until the provider reports it, so the adapter treats a +model call whose counts are all zero as unreported. The result has no `usage` +when Pi produced no assistant message, when one of those calls has no complete +usage report, when an assistant message was aborted, when a compaction was +aborted or reported no usage, or when Pi retried a compaction summary call, +because Pi keeps no usage for the failed attempt. Pi computes its cost figures +from the model catalog's rates, so the adapter does not report them as +`cost_usd`. ## Custom Tool Modules diff --git a/adapters/typescript/pi/src/pi-sdk.ts b/adapters/typescript/pi/src/pi-sdk.ts index 1f6d76f71..5152ac7ba 100644 --- a/adapters/typescript/pi/src/pi-sdk.ts +++ b/adapters/typescript/pi/src/pi-sdk.ts @@ -69,6 +69,8 @@ interface PiToolFactoryContext { workspace: string; } +const NO_TOKENS: PiTokenUsage = { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 }; + /** * Add one call's usage to a prompt's running total. A call without a complete * usage report makes the whole total unknown. @@ -97,13 +99,30 @@ export function addUsage(total: PiTokenUsage | undefined, usage: unknown): PiTok }; } +/** + * Add one model call's usage to a prompt's running total. Pi fills a call's + * usage with zeros before the response arrives and overwrites it only when the + * provider reports usage, so a call whose four counts are all zero carries no + * usage report and makes the total unknown. A model call always sends input + * tokens, so a provider that measured the call reports a nonzero count. + */ +function addModelCallUsage(total: PiTokenUsage | undefined, usage: unknown): PiTokenUsage | undefined { + const call = addUsage(NO_TOKENS, usage); + if (call === undefined || Object.values(call).every((count) => count === 0)) { + return undefined; + } + return addUsage(total, call); +} + /** * Add the tokens one session event recorded to a prompt's running total, as * Pi's own usage totals count them: an assistant message, a tool result that * reports usage, a finished compaction's summary call, and a usage entry such - * as a cache warm-up. An aborted assistant message, or a compaction that was - * aborted or finished without a usage report, makes the total unknown: Pi - * keeps no usage for a call it stopped, which may still have been billed. + * as a cache warm-up. The total becomes unknown when a model call carries no + * usage report, when an assistant message or a compaction was aborted, or when + * a compaction summary call was retried: Pi keeps no usage for a call it + * stopped or a summary attempt that failed, and either may still have been + * billed. */ export function addEventUsage(total: PiTokenUsage | undefined, event: unknown): PiTokenUsage | undefined { const fields = (typeof event === "object" && event !== null ? event : {}) as { @@ -116,7 +135,7 @@ export function addEventUsage(total: PiTokenUsage | undefined, event: unknown): switch (fields.type) { case "message_end": if (fields.message?.role === "assistant") { - return fields.message.stopReason === "aborted" ? undefined : addUsage(total, fields.message.usage); + return fields.message.stopReason === "aborted" ? undefined : addModelCallUsage(total, fields.message.usage); } if (fields.message?.role === "toolResult" && fields.message.usage !== undefined) { return addUsage(total, fields.message.usage); @@ -124,10 +143,12 @@ export function addEventUsage(total: PiTokenUsage | undefined, event: unknown): return total; case "compaction_end": return fields.aborted === false && fields.result !== undefined - ? addUsage(total, fields.result.usage) + ? addModelCallUsage(total, fields.result.usage) : undefined; + case "summarization_retry_scheduled": + return undefined; case "entry_appended": - return fields.entry?.type === "usage" ? addUsage(total, fields.entry.usage) : total; + return fields.entry?.type === "usage" ? addModelCallUsage(total, fields.entry.usage) : total; default: return total; } @@ -956,7 +977,7 @@ class PiSdkSessionHandle implements PiSessionHandle { let finalAssistant: | { role: "assistant"; content: unknown; stopReason: string; errorMessage?: string } | undefined; - let usage: PiTokenUsage | undefined = { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 }; + let usage: PiTokenUsage | undefined = NO_TOKENS; const unsubscribe = this.session.subscribe((event) => { if (event.type === "turn_start") { turnStarted = true; diff --git a/adapters/typescript/pi/test/pi-sdk.test.mjs b/adapters/typescript/pi/test/pi-sdk.test.mjs index 037843961..6fe3978cd 100644 --- a/adapters/typescript/pi/test/pi-sdk.test.mjs +++ b/adapters/typescript/pi/test/pi-sdk.test.mjs @@ -1486,3 +1486,119 @@ test("an aborted assistant message leaves the prompt's tokens unknown", () => { undefined, ); }); + +test("a model call with zero-filled usage leaves the prompt's tokens unknown", () => { + const zero = { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 }; + const unreported = { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, totalTokens: 0 }; + + assert.equal(addEventUsage(zero, { type: "message_end", message: { role: "assistant", usage: unreported } }), undefined); + assert.equal(addEventUsage(zero, { type: "compaction_end", aborted: false, result: { usage: unreported } }), undefined); + assert.equal(addEventUsage(zero, { type: "entry_appended", entry: { type: "usage", usage: unreported } }), undefined); + // A tool reports its own usage, so a zero there is a measured zero. + assert.deepEqual(addEventUsage(zero, { type: "message_end", message: { role: "toolResult", usage: unreported } }), zero); +}); + +test("a retried compaction summary leaves the prompt's tokens unknown", () => { + const usage = { input: 20, output: 1, cacheRead: 0, cacheWrite: 0 }; + const events = [ + { type: "summarization_retry_scheduled", attempt: 1, maxAttempts: 3, delayMs: 0, errorMessage: "stream dropped" }, + { type: "summarization_retry_attempt_start" }, + { type: "summarization_retry_finished" }, + { type: "compaction_end", aborted: false, result: { summary: "s", usage } }, + ]; + + assert.equal(events.reduce(addEventUsage, { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 }), undefined); +}); + +/** + * Run prompts through the real Pi SDK against a local chat-completions endpoint + * that answers each request from `replies` in order. A reply's `usage` is sent + * as the stream's final usage chunk; a reply without one sends no usage. + */ +async function promptsAgainstScriptedEndpoint(replies, prompts) { + const workspace = await realpath(await mkdtemp(join(tmpdir(), "fabric-pi-usage-"))); + await writeFile(join(workspace, "note.txt"), "nonce\n", "utf8"); + let served = 0; + const provider = createServer(async (request, response) => { + for await (const _chunk of request); + const reply = replies[served++]; + const common = { id: `chatcmpl-usage-${served}`, object: "chat.completion.chunk", created: 0, model: "openai/gpt-oss-20b" }; + const chunk = (fields) => `data: ${JSON.stringify({ ...common, ...fields })}\n\n`; + const delta = reply.readNote + ? { role: "assistant", tool_calls: [{ index: 0, id: `call-${served}`, type: "function", function: { name: "read", arguments: JSON.stringify({ path: "note.txt" }) } }] } + : { role: "assistant", content: reply.content }; + response.writeHead(200, { "content-type": "text/event-stream" }); + response.end( + chunk({ choices: [{ index: 0, delta, finish_reason: null }] }) + + chunk({ choices: [{ index: 0, delta: {}, finish_reason: reply.readNote ? "tool_calls" : "stop" }] }) + + (reply.usage === undefined ? "" : chunk({ choices: [], usage: reply.usage })) + + "data: [DONE]\n\n", + ); + }); + const providerUrl = await listen(provider); + let handle; + try { + handle = await new PiSdkSessionFactory().create({ + agentName: "pi-usage-test", + baseDir: workspace, + config: { + models: { + default: { api_key_env: "TEST_API_KEY", base_url: `${providerUrl}/v1`, model: "openai/gpt-oss-20b", provider: "nvidia" }, + }, + tools: { enabled: null }, + }, + runtimeContext: { ...makeRuntimeContext(workspace), telemetry: undefined }, + }); + const outcomes = []; + for (const prompt of prompts) { + outcomes.push(await handle.prompt(prompt)); + } + assert.equal(served, replies.length); + return outcomes; + } finally { + await handle?.stop(); + await close(provider); + await rm(workspace, { recursive: true, force: true }); + } +} + +test( + "reports each prompt's own tokens through the Pi SDK", + { skip: supportsPi ? false : "Pi 1.0 requires Node 22.19 or newer" }, + async () => { + const [first, second] = await promptsAgainstScriptedEndpoint( + [ + { readNote: true, usage: { prompt_tokens: 100, completion_tokens: 10, prompt_tokens_details: { cached_tokens: 40 } } }, + { content: "first", usage: { prompt_tokens: 150, completion_tokens: 0 } }, + { content: "second", usage: { prompt_tokens: 300, completion_tokens: 25 } }, + ], + ["read the note", "again"], + ); + + // A reported zero output count stays a measured zero. + assert.deepEqual(first.usage, { input: 210, output: 10, cacheRead: 40, cacheWrite: 0 }); + assert.deepEqual(second.usage, { input: 300, output: 25, cacheRead: 0, cacheWrite: 0 }); + }, +); + +test( + "leaves a prompt's tokens unknown when any provider response omits usage", + { skip: supportsPi ? false : "Pi 1.0 requires Node 22.19 or newer" }, + async () => { + const [omitted, mixed, reported] = await promptsAgainstScriptedEndpoint( + [ + { content: "no usage" }, + { readNote: true, usage: { prompt_tokens: 100, completion_tokens: 10 } }, + { content: "partial" }, + { content: "reported", usage: { prompt_tokens: 50, completion_tokens: 5 } }, + ], + ["one", "two", "three"], + ); + + assert.equal(omitted.text, "no usage"); + assert.equal(omitted.usage, undefined); + assert.equal(mixed.text, "partial"); + assert.equal(mixed.usage, undefined); + assert.deepEqual(reported.usage, { input: 50, output: 5, cacheRead: 0, cacheWrite: 0 }); + }, +); From f734edffd798e35347db7a380662737ef7b9b78b Mon Sep 17 00:00:00 2001 From: mnajafian-nv Date: Fri, 16 Oct 2026 00:45:48 -0700 Subject: [PATCH 07/68] fix: leave Pi usage unknown after a split-turn compaction When a compaction cuts inside a turn, Pi summarizes the earlier history and the turn prefix in two model calls and reports only their combined usage on compaction_end. A summary call without a usage report is zero-filled, so the combined usage can look complete while one call is missing. Pi exposes the compaction plan, including isSplitTurn, only to extensions through session_before_compact. Register a hidden adapter extension that forwards that event to the running prompt, and leave the prompt's usage unknown unless the plan says the compaction is not split. State in the README and code comments that an all-zero provider report is discarded together with an unreported call. Signed-off-by: mnajafian-nv --- adapters/typescript/pi/README.md | 11 ++-- adapters/typescript/pi/src/pi-sdk.ts | 56 +++++++++++++++--- adapters/typescript/pi/test/pi-sdk.test.mjs | 65 ++++++++++++++++++++- 3 files changed, 118 insertions(+), 14 deletions(-) diff --git a/adapters/typescript/pi/README.md b/adapters/typescript/pi/README.md index 0dfb5703d..3163ba101 100644 --- a/adapters/typescript/pi/README.md +++ b/adapters/typescript/pi/README.md @@ -220,13 +220,16 @@ report usage, finished compaction summaries, and usage entries such as cache warm-ups. It holds `input_tokens` with cache reads and cache writes included, `cached_input_tokens`, `output_tokens`, and `total_tokens`. Pi fills a model call's usage with zeros until the provider reports it, so the adapter treats a -model call whose counts are all zero as unreported. The result has no `usage` +model call whose counts are all zero as unreported. A provider report of all +zeros looks the same, so the adapter discards it too. The result has no `usage` when Pi produced no assistant message, when one of those calls has no complete usage report, when an assistant message was aborted, when a compaction was aborted or reported no usage, or when Pi retried a compaction summary call, -because Pi keeps no usage for the failed attempt. Pi computes its cost figures -from the model catalog's rates, so the adapter does not report them as -`cost_usd`. +because Pi keeps no usage for the failed attempt. It also has no `usage` when a +compaction cuts inside a turn, because Pi then summarizes the earlier history +and the start of that turn in separate calls and reports only their combined +usage. Pi computes its cost figures from the model catalog's rates, so the +adapter does not report them as `cost_usd`. ## Custom Tool Modules diff --git a/adapters/typescript/pi/src/pi-sdk.ts b/adapters/typescript/pi/src/pi-sdk.ts index 5152ac7ba..4bca4ca66 100644 --- a/adapters/typescript/pi/src/pi-sdk.ts +++ b/adapters/typescript/pi/src/pi-sdk.ts @@ -102,9 +102,9 @@ export function addUsage(total: PiTokenUsage | undefined, usage: unknown): PiTok /** * Add one model call's usage to a prompt's running total. Pi fills a call's * usage with zeros before the response arrives and overwrites it only when the - * provider reports usage, so a call whose four counts are all zero carries no - * usage report and makes the total unknown. A model call always sends input - * tokens, so a provider that measured the call reports a nonzero count. + * provider reports usage, so a call whose four counts are all zero is treated + * as unreported and makes the total unknown. A provider report of all zeros + * looks the same and is discarded too. */ function addModelCallUsage(total: PiTokenUsage | undefined, usage: unknown): PiTokenUsage | undefined { const call = addUsage(NO_TOKENS, usage); @@ -122,12 +122,16 @@ function addModelCallUsage(total: PiTokenUsage | undefined, usage: unknown): PiT * usage report, when an assistant message or a compaction was aborted, or when * a compaction summary call was retried: Pi keeps no usage for a call it * stopped or a summary attempt that failed, and either may still have been - * billed. + * billed. It also becomes unknown when Pi prepares a split-turn compaction + * (a `session_before_compact` event whose preparation has `isSplitTurn`): + * Pi summarizes the history and the turn prefix in separate calls and reports + * only their combined usage, so one unreported call cannot be detected. */ export function addEventUsage(total: PiTokenUsage | undefined, event: unknown): PiTokenUsage | undefined { const fields = (typeof event === "object" && event !== null ? event : {}) as { type?: unknown; message?: { role?: unknown; stopReason?: unknown; usage?: unknown }; + preparation?: { isSplitTurn?: unknown }; aborted?: unknown; result?: { usage?: unknown }; entry?: { type?: unknown; usage?: unknown }; @@ -145,6 +149,8 @@ export function addEventUsage(total: PiTokenUsage | undefined, event: unknown): return fields.aborted === false && fields.result !== undefined ? addModelCallUsage(total, fields.result.usage) : undefined; + case "session_before_compact": + return fields.preparation?.isSplitTurn === false ? total : undefined; case "summarization_retry_scheduled": return undefined; case "entry_appended": @@ -581,6 +587,20 @@ function mcpRegistrationExtension(servers: Record): Ext }; } +/** + * Forward each compaction Pi prepares to the running prompts. Pi passes the + * preparation only to extensions, not to session subscribers. + */ +function compactionPreparationExtension(listeners: Set<(event: unknown) => void>): ExtensionFactory { + return (api) => { + api.on("session_before_compact", (event) => { + for (const listener of listeners) { + listener(event); + } + }); + }; +} + function selectModel(config: AgentConfig): AgentModelConfig { const entries = Object.entries(config.models ?? {}); if (entries.length === 0) { @@ -954,11 +974,19 @@ class PiSdkSessionHandle implements PiSessionHandle { private stopped = false; private cumulativeTurnCount = 0; private readonly runtimeDir: string; - - constructor(session: AgentSession, state: { shutdownRequested: boolean }, runtimeDir: string, relay?: PiRelayRuntime) { + private readonly compactionListeners: Set<(event: unknown) => void>; + + constructor( + session: AgentSession, + state: { shutdownRequested: boolean }, + runtimeDir: string, + compactionListeners: Set<(event: unknown) => void>, + relay?: PiRelayRuntime, + ) { this.session = session; this.state = state; this.runtimeDir = runtimeDir; + this.compactionListeners = compactionListeners; this.relay = relay; this.unsubscribeTurnCounter = this.session.subscribe((event) => { if (event.type === "turn_start") { @@ -987,6 +1015,10 @@ class PiSdkSessionHandle implements PiSessionHandle { } usage = addEventUsage(usage, event); }); + const onCompactionPrepared = (event: unknown): void => { + usage = addEventUsage(usage, event); + }; + this.compactionListeners.add(onCompactionPrepared); try { await this.session.prompt(text, { expandPromptTemplates: true, @@ -996,6 +1028,7 @@ class PiSdkSessionHandle implements PiSessionHandle { }, }); } finally { + this.compactionListeners.delete(onCompactionPrepared); unsubscribe(); } return { @@ -1167,9 +1200,16 @@ export class PiSdkSessionFactory implements PiSessionFactory { runtimeDir = await mkdtemp(join(tmpdir(), "nemo-fabric-pi-")); const agentDir = runtimeDir; const mcpConnections = new Map(Object.keys(mcpServers).map((name) => [name, mcpReadiness()])); + const compactionListeners = new Set<(event: unknown) => void>(); + const compactionExtension = { + name: "nemo-fabric-compaction-usage", + hidden: true, + factory: compactionPreparationExtension(compactionListeners), + }; const extensionFactories = Object.keys(mcpServers).length === 0 - ? [] + ? [compactionExtension] : [ + compactionExtension, { name: "nemo-fabric-mcp-servers", hidden: true, factory: mcpRegistrationExtension(mcpServers) }, { name: "nemo-fabric-mcp", @@ -1258,7 +1298,7 @@ export class PiSdkSessionFactory implements PiSessionFactory { tools: enabled === null ? undefined : enabled, excludeTools: blocked, }); - handle = new PiSdkSessionHandle(session, state, runtimeDir, relay); + handle = new PiSdkSessionHandle(session, state, runtimeDir, compactionListeners, relay); const blockedNames = new Set(blocked); const availableNames = new Set(session.getAllTools().map((tool) => tool.name)); const missing = (enabled ?? []).filter((name) => !blockedNames.has(name) && !availableNames.has(name)); diff --git a/adapters/typescript/pi/test/pi-sdk.test.mjs b/adapters/typescript/pi/test/pi-sdk.test.mjs index 6fe3978cd..e82e0807d 100644 --- a/adapters/typescript/pi/test/pi-sdk.test.mjs +++ b/adapters/typescript/pi/test/pi-sdk.test.mjs @@ -1498,6 +1498,15 @@ test("a model call with zero-filled usage leaves the prompt's tokens unknown", ( assert.deepEqual(addEventUsage(zero, { type: "message_end", message: { role: "toolResult", usage: unreported } }), zero); }); +test("a split-turn compaction leaves the prompt's tokens unknown", () => { + const zero = { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 }; + const prepared = (preparation) => ({ type: "session_before_compact", preparation }); + + assert.deepEqual(addEventUsage(zero, prepared({ isSplitTurn: false })), zero); + assert.equal(addEventUsage(zero, prepared({ isSplitTurn: true })), undefined); + assert.equal(addEventUsage(zero, prepared(undefined)), undefined); +}); + test("a retried compaction summary leaves the prompt's tokens unknown", () => { const usage = { input: 20, output: 1, cacheRead: 0, cacheWrite: 0 }; const events = [ @@ -1514,8 +1523,9 @@ test("a retried compaction summary leaves the prompt's tokens unknown", () => { * Run prompts through the real Pi SDK against a local chat-completions endpoint * that answers each request from `replies` in order. A reply's `usage` is sent * as the stream's final usage chunk; a reply without one sends no usage. + * `modelExtensions` overrides the catalog entry, such as its context window. */ -async function promptsAgainstScriptedEndpoint(replies, prompts) { +async function promptsAgainstScriptedEndpoint(replies, prompts, modelExtensions = undefined) { const workspace = await realpath(await mkdtemp(join(tmpdir(), "fabric-pi-usage-"))); await writeFile(join(workspace, "note.txt"), "nonce\n", "utf8"); let served = 0; @@ -1543,7 +1553,13 @@ async function promptsAgainstScriptedEndpoint(replies, prompts) { baseDir: workspace, config: { models: { - default: { api_key_env: "TEST_API_KEY", base_url: `${providerUrl}/v1`, model: "openai/gpt-oss-20b", provider: "nvidia" }, + default: { + api_key_env: "TEST_API_KEY", + base_url: `${providerUrl}/v1`, + model: "openai/gpt-oss-20b", + provider: "nvidia", + ...(modelExtensions === undefined ? {} : { extensions: modelExtensions }), + }, }, tools: { enabled: null }, }, @@ -1602,3 +1618,48 @@ test( assert.deepEqual(reported.usage, { input: 50, output: 5, cacheRead: 0, cacheWrite: 0 }); }, ); + +test( + "counts a compaction that Pi summarizes in one call", + { skip: supportsPi ? false : "Pi 1.0 requires Node 22.19 or newer" }, + async () => { + // The second prompt is longer than Pi's 20,000 recent tokens, so Pi cuts at + // that prompt and summarizes only the earlier history, in one call. + const [, compacted] = await promptsAgainstScriptedEndpoint( + [ + { content: "first", usage: { prompt_tokens: 100, completion_tokens: 5 } }, + { readNote: true, usage: { prompt_tokens: 200, completion_tokens: 5 } }, + { content: "done", usage: { prompt_tokens: 90_000, completion_tokens: 5 } }, + { content: "history summary", usage: { prompt_tokens: 20, completion_tokens: 1 } }, + ], + ["one", "y".repeat(100_000)], + { context_window: 100_000, max_tokens: 1_000 }, + ); + + assert.deepEqual(compacted.usage, { input: 90_220, output: 11, cacheRead: 0, cacheWrite: 0 }); + }, +); + +test( + "leaves a prompt's tokens unknown after a split-turn compaction", + { skip: supportsPi ? false : "Pi 1.0 requires Node 22.19 or newer" }, + async () => { + // The second prompt's final reply reports a context above the compaction + // threshold and is longer than Pi's 20,000 recent tokens, so Pi cuts inside + // that turn and summarizes the history and the turn prefix in two calls. + // The history summary reports no usage; the turn-prefix summary does. + const [, compacted] = await promptsAgainstScriptedEndpoint( + [ + { content: "first", usage: { prompt_tokens: 100, completion_tokens: 5 } }, + { readNote: true, usage: { prompt_tokens: 200, completion_tokens: 5 } }, + { content: "x".repeat(100_000), usage: { prompt_tokens: 90_000, completion_tokens: 5 } }, + { content: "history summary" }, + { content: "turn prefix summary", usage: { prompt_tokens: 20, completion_tokens: 1 } }, + ], + ["one", "read the note"], + { context_window: 100_000, max_tokens: 1_000 }, + ); + + assert.equal(compacted.usage, undefined); + }, +); From a1d62ce92d6f7ada08b3a6966d6d68f6fd8d3e23 Mon Sep 17 00:00:00 2001 From: mnajafian-nv Date: Mon, 19 Oct 2026 14:01:36 -0700 Subject: [PATCH 08/68] fix: leave omitted Hermes and Pi token counts unknown A provider response can report usage without one of its counts, such as prompt_tokens with no completion_tokens. Hermes' normalize_usage and Pi's provider parsers fill the missing count with 0, so both adapters reported it as a measured 0 and derived total_tokens from it. The usage-less checks only caught a response with no usage at all. Both adapters now read the provider's raw usage before the SDK normalizes it and keep only the counts it reported. Hermes wraps turn_response_check.record_response_usage for the agent it owns and counts responses whose usage lacks an input or output count. Pi wraps the agent's stream function, reads each call's raw stream events through onProviderStreamEvent for every supported wire protocol, and keys the result by the call's final usage object, which covers assistant messages and compaction summaries. An omitted input count leaves out input_tokens, cached_input_tokens and total_tokens; an omitted output count leaves out output_tokens and total_tokens. A reported zero stays 0, which also lets Pi keep a provider report of all zeros that it used to discard. Tests cover a prompt-only report, a completion-only report, reported zeros and a partial response inside a multi-call invocation, through Hermes' real usage accounting and the real Pi SDK against a local streaming endpoint. Signed-off-by: mnajafian-nv --- adapters/python/hermes/README.md | 6 +- .../nemo_fabric_adapters/hermes/adapter.py | 78 ++++++- adapters/typescript/pi/README.md | 22 +- adapters/typescript/pi/src/pi-sdk.ts | 210 +++++++++++++++--- adapters/typescript/pi/src/runtime.ts | 39 ++-- adapters/typescript/pi/test/pi-sdk.test.mjs | 129 +++++++++-- adapters/typescript/pi/test/runtime.test.mjs | 32 +++ tests/adapters/test_hermes_adapter.py | 95 ++++++-- 8 files changed, 517 insertions(+), 94 deletions(-) diff --git a/adapters/python/hermes/README.md b/adapters/python/hermes/README.md index 479a101df..4d738c249 100644 --- a/adapters/python/hermes/README.md +++ b/adapters/python/hermes/README.md @@ -121,7 +121,11 @@ invocation: `input_tokens` with cache reads and cache writes included, auxiliary calls, such as context-compression summaries, separately from these counters, so they are not included. Hermes adds no tokens for a provider response that carried no usage, so an invocation with such a response has no -`usage`. A response that reported zero tokens still counts as 0. A counter that +`usage`. A response that reported zero tokens still counts as 0. Hermes also +reads a count that a response's usage left out as 0. When a response left out +its input count, the invocation's `usage` leaves out `input_tokens`, +`cached_input_tokens`, and `total_tokens`. When it left out its output count, +the `usage` leaves out `output_tokens` and `total_tokens`. A counter that Hermes does not report, or that went down, is left out for that invocation. After an invocation that raised, was cancelled, or timed out, or whose result held no counters, the adapter has no baseline. The next invocation then reports diff --git a/adapters/python/hermes/src/nemo_fabric_adapters/hermes/adapter.py b/adapters/python/hermes/src/nemo_fabric_adapters/hermes/adapter.py index d2690a0b5..a1a093fea 100755 --- a/adapters/python/hermes/src/nemo_fabric_adapters/hermes/adapter.py +++ b/adapters/python/hermes/src/nemo_fabric_adapters/hermes/adapter.py @@ -71,6 +71,8 @@ def __init__(self) -> None: # from this baseline. A counter missing from it has no baseline. self._session_token_totals: dict[str, int] = _fresh_session_totals() self._usage_less_responses = 0 + # Provider responses whose usage left out the input or output count. + self._unreported_counts = {"input": 0, "output": 0} async def start(self, payload: dict[str, Any]) -> None: if self._started: @@ -223,7 +225,7 @@ async def start(self, payload: dict[str, Any]) -> None: ) ) self._session_token_totals = _fresh_session_totals() - self._count_usage_less_responses(self._agent) + self._observe_usage_reports(self._agent) self._agent_config = agent_config self._started = True except BaseException: @@ -296,6 +298,7 @@ def run_hermes_turn() -> tuple[dict[str, Any], str]: # the next one. previous_totals, self._session_token_totals = self._session_token_totals, {} usage_less_before = self._usage_less_responses + unreported_before = dict(self._unreported_counts) invoke_task = asyncio.create_task(asyncio.to_thread(run_hermes_turn)) self._active_invoke_task = invoke_task @@ -380,7 +383,14 @@ def clear_active_invoke_task( code="hermes_reported_failure", message="Hermes did not complete the invocation", ) - usage, self._session_token_totals = _invocation_usage(result, previous_totals) + usage, self._session_token_totals = _invocation_usage( + result, + previous_totals, + input_reported=self._unreported_counts["input"] + == unreported_before["input"], + output_reported=self._unreported_counts["output"] + == unreported_before["output"], + ) if self._usage_less_responses != usage_less_before: # A provider response without usage added no tokens, so the # change in the counters is an undercount, not this invocation's. @@ -392,6 +402,11 @@ def clear_active_invoke_task( usage=usage, ) + def _observe_usage_reports(self, agent: Any) -> None: + """Count the agent's provider responses that left usage unreported.""" + self._count_usage_less_responses(agent) + self._count_unreported_counts(agent) + def _count_usage_less_responses(self, agent: Any) -> None: """Count the agent's provider responses that carried no usage. @@ -413,6 +428,36 @@ def count_and_note() -> None: compressor.note_usage_less_response = count_and_note + def _count_unreported_counts(self, agent: Any) -> None: + """Count the agent's provider responses whose usage left out the input + or the output count. + + Hermes reads a missing count as 0, so its counters cannot tell an + omitted count from a reported zero. Hermes' turn loop passes every + response to ``agent.turn_response_check.record_response_usage``, so + the adapter wraps that call and reads the raw usage before Hermes + normalizes it. Responses of other agents, such as delegated + subagents, are not counted: their tokens never reach this agent's + counters. The lifecycle host serves one runtime per process, so the + wrapper is installed once. + """ + from agent import turn_response_check + + record_response_usage = turn_response_check.record_response_usage + + def count_and_record( + usage_agent: Any, response: Any, *args: Any, **kwargs: Any + ) -> Any: + usage = getattr(response, "usage", None) + # A response without usage is counted by the compressor hook. + if usage_agent is agent and usage: + for count, fields in _RAW_COUNT_FIELDS.items(): + if not any(_is_reported_count(usage, field) for field in fields): + self._unreported_counts[count] += 1 + return record_response_usage(usage_agent, response, *args, **kwargs) + + turn_response_check.record_response_usage = count_and_record + async def _authenticate_mcp_servers(self) -> None: if self._mcp_authentication_checked: return @@ -633,13 +678,32 @@ def _invoke_hermes_turn( ) +# The raw usage fields Hermes reads for a response's input and output counts. +# Chat Completions usage may use either spelling; Anthropic Messages and +# Responses usage carry only input_tokens and output_tokens. +_RAW_COUNT_FIELDS = { + "input": ("prompt_tokens", "input_tokens"), + "output": ("completion_tokens", "output_tokens"), +} + + +def _is_reported_count(usage: Any, field: str) -> bool: + """Return whether the provider's raw usage reports ``field`` as a count.""" + value = usage.get(field) if isinstance(usage, dict) else getattr(usage, field, None) + return isinstance(value, int) and not isinstance(value, bool) and value >= 0 + + def _fresh_session_totals() -> dict[str, int]: """Return the baseline of a new Hermes agent, whose counters start at 0.""" return dict.fromkeys(_SESSION_TOKEN_KEYS, 0) def _invocation_usage( - result: dict[str, Any], previous_totals: dict[str, int] + result: dict[str, Any], + previous_totals: dict[str, int], + *, + input_reported: bool = True, + output_reported: bool = True, ) -> tuple[AgentUsage | None, dict[str, int]]: """Return this invocation's usage and the session totals to keep as baseline. @@ -648,7 +712,9 @@ def _invocation_usage( result or from ``previous_totals`` (no baseline: the previous invocation failed or its result held no counters), or that went down, is unknown for this invocation. The usage therefore never holds a whole-session total or - another invocation's tokens. + another invocation's tokens. ``input_reported`` or ``output_reported`` is + False when a provider response of this invocation left that count out, + which Hermes counted as 0, so that count and the total are unknown. """ totals = { key: value @@ -663,7 +729,7 @@ def _invocation_usage( if key in previous_totals and value >= previous_totals[key] } fields: dict[str, Any] = {} - if all( + if input_reported and all( key in change for key in ("input_tokens", "cache_read_tokens", "cache_write_tokens") ): @@ -674,7 +740,7 @@ def _invocation_usage( ) fields["cached_input_tokens"] = change["cache_read_tokens"] fields["input_tokens_include_cache"] = True - if "output_tokens" in change: + if output_reported and "output_tokens" in change: fields["output_tokens"] = change["output_tokens"] if "input_tokens" in fields and "output_tokens" in fields: fields["total_tokens"] = fields["input_tokens"] + fields["output_tokens"] diff --git a/adapters/typescript/pi/README.md b/adapters/typescript/pi/README.md index 353515405..5758b031b 100644 --- a/adapters/typescript/pi/README.md +++ b/adapters/typescript/pi/README.md @@ -218,14 +218,20 @@ Each result's `usage` sums the tokens Pi records during the invocation, counted the way Pi's own usage totals count them: assistant messages, tool results that report usage, finished compaction summaries, and usage entries such as cache warm-ups. It holds `input_tokens` with cache reads and cache writes included, -`cached_input_tokens`, `output_tokens`, and `total_tokens`. Pi fills a model -call's usage with zeros until the provider reports it, so the adapter treats a -model call whose counts are all zero as unreported. A provider report of all -zeros looks the same, so the adapter discards it too. The result has no `usage` -when Pi produced no assistant message, when one of those calls has no complete -usage report, when an assistant message was aborted, when a compaction was -aborted or reported no usage, or when Pi retried a compaction summary call, -because Pi keeps no usage for the failed attempt. It also has no `usage` when a +`cached_input_tokens`, `output_tokens`, and `total_tokens`. Pi fills a count +that the provider did not report with 0, so the adapter reads the provider's +raw stream events for each assistant message and compaction summary call and +keeps only the counts the provider reported. A reported zero counts as 0. When +a call left out its input count, the result leaves out `input_tokens`, +`cached_input_tokens`, and `total_tokens`. When a call left out its output +count, the result leaves out `output_tokens` and `total_tokens`. The result has +no `usage` when Pi produced no assistant message, when a model call reported no +counts, when a tool result has no complete usage report, when an assistant +message was aborted, when a compaction was aborted or reported no usage, or +when Pi retried a compaction summary call, because Pi keeps no usage for the +failed attempt. It also has no `usage` when Pi recorded a cache warm-up during +the invocation, because Pi sends warm-ups outside the agent's stream function, +where the adapter cannot read the raw report. It also has no `usage` when a compaction cuts inside a turn, because Pi then summarizes the earlier history and the start of that turn in separate calls and reports only their combined usage. Pi computes its cost figures from the model catalog's rates, so the diff --git a/adapters/typescript/pi/src/pi-sdk.ts b/adapters/typescript/pi/src/pi-sdk.ts index 4bca4ca66..c57d6d165 100644 --- a/adapters/typescript/pi/src/pi-sdk.ts +++ b/adapters/typescript/pi/src/pi-sdk.ts @@ -71,12 +71,20 @@ interface PiToolFactoryContext { const NO_TOKENS: PiTokenUsage = { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 }; -/** - * Add one call's usage to a prompt's running total. A call without a complete - * usage report makes the whole total unknown. - */ -export function addUsage(total: PiTokenUsage | undefined, usage: unknown): PiTokenUsage | undefined { - if (total === undefined || typeof usage !== "object" || usage === null) { +/** Which counts of one model call the provider reported before Pi filled the rest with 0. */ +export interface ReportedCounts { + input: boolean; + output: boolean; +} + +/** What the provider reported for the model call that produced a usage object, if the adapter saw it. */ +export type ReportedCountsLookup = (usage: unknown) => ReportedCounts | undefined; + +const BOTH_REPORTED: ReportedCounts = { input: true, output: true }; + +/** One call's four token counts, or undefined unless each is a non-negative safe integer. */ +function callCounts(usage: unknown): Required | undefined { + if (typeof usage !== "object" || usage === null) { return undefined; } const counts = usage as Record; @@ -91,43 +99,76 @@ export function addUsage(total: PiTokenUsage | undefined, usage: unknown): PiTok if (input === undefined || output === undefined || cacheRead === undefined || cacheWrite === undefined) { return undefined; } - return { - input: total.input + input, - output: total.output + output, - cacheRead: total.cacheRead + cacheRead, - cacheWrite: total.cacheWrite + cacheWrite, - }; + return { input, output, cacheRead, cacheWrite }; +} + +/** + * Add a call's counts to a prompt's running total. The input side (input with + * its cache reads and writes) and the output stay known only while every call + * reported them; the total is unknown once neither is. + */ +function sumCounts(total: PiTokenUsage, call: Required, reported: ReportedCounts): PiTokenUsage | undefined { + const sum: PiTokenUsage = {}; + if (reported.input && total.input !== undefined && total.cacheRead !== undefined && total.cacheWrite !== undefined) { + sum.input = total.input + call.input; + sum.cacheRead = total.cacheRead + call.cacheRead; + sum.cacheWrite = total.cacheWrite + call.cacheWrite; + } + if (reported.output && total.output !== undefined) { + sum.output = total.output + call.output; + } + return Object.keys(sum).length === 0 ? undefined : sum; +} + +/** + * Add one call's own usage report, such as a tool's, to a prompt's running + * total. A report without all four counts makes the whole total unknown. + */ +export function addUsage(total: PiTokenUsage | undefined, usage: unknown): PiTokenUsage | undefined { + const call = callCounts(usage); + return total === undefined || call === undefined ? undefined : sumCounts(total, call, BOTH_REPORTED); } /** - * Add one model call's usage to a prompt's running total. Pi fills a call's - * usage with zeros before the response arrives and overwrites it only when the - * provider reports usage, so a call whose four counts are all zero is treated - * as unreported and makes the total unknown. A provider report of all zeros - * looks the same and is discarded too. + * Add one model call's usage to a prompt's running total. Pi's provider + * parsers fill a count the provider omitted with 0, so `reportedCounts` says + * which counts the provider actually reported, from the raw stream events the + * session recorded for that call. A count the provider omitted is unknown for + * the whole prompt, and so is every count of a call whose raw events the + * adapter did not see. A reported zero stays a measured zero. */ -function addModelCallUsage(total: PiTokenUsage | undefined, usage: unknown): PiTokenUsage | undefined { - const call = addUsage(NO_TOKENS, usage); - if (call === undefined || Object.values(call).every((count) => count === 0)) { +function addModelCallUsage( + total: PiTokenUsage | undefined, + usage: unknown, + reportedCounts: ReportedCountsLookup, +): PiTokenUsage | undefined { + const call = callCounts(usage); + const reported = reportedCounts(usage); + if (total === undefined || call === undefined || reported === undefined) { return undefined; } - return addUsage(total, call); + return sumCounts(total, call, reported); } /** * Add the tokens one session event recorded to a prompt's running total, as * Pi's own usage totals count them: an assistant message, a tool result that * reports usage, a finished compaction's summary call, and a usage entry such - * as a cache warm-up. The total becomes unknown when a model call carries no - * usage report, when an assistant message or a compaction was aborted, or when - * a compaction summary call was retried: Pi keeps no usage for a call it + * as a cache warm-up. A model call adds only the counts its provider + * reported, as `reportedCounts` records them. The total becomes unknown when a + * model call reports no count, when an assistant message or a compaction was + * aborted, or when a compaction summary call was retried: Pi keeps no usage for a call it * stopped or a summary attempt that failed, and either may still have been * billed. It also becomes unknown when Pi prepares a split-turn compaction * (a `session_before_compact` event whose preparation has `isSplitTurn`): * Pi summarizes the history and the turn prefix in separate calls and reports * only their combined usage, so one unreported call cannot be detected. */ -export function addEventUsage(total: PiTokenUsage | undefined, event: unknown): PiTokenUsage | undefined { +export function addEventUsage( + total: PiTokenUsage | undefined, + event: unknown, + reportedCounts: ReportedCountsLookup, +): PiTokenUsage | undefined { const fields = (typeof event === "object" && event !== null ? event : {}) as { type?: unknown; message?: { role?: unknown; stopReason?: unknown; usage?: unknown }; @@ -139,7 +180,9 @@ export function addEventUsage(total: PiTokenUsage | undefined, event: unknown): switch (fields.type) { case "message_end": if (fields.message?.role === "assistant") { - return fields.message.stopReason === "aborted" ? undefined : addModelCallUsage(total, fields.message.usage); + return fields.message.stopReason === "aborted" + ? undefined + : addModelCallUsage(total, fields.message.usage, reportedCounts); } if (fields.message?.role === "toolResult" && fields.message.usage !== undefined) { return addUsage(total, fields.message.usage); @@ -147,19 +190,122 @@ export function addEventUsage(total: PiTokenUsage | undefined, event: unknown): return total; case "compaction_end": return fields.aborted === false && fields.result !== undefined - ? addModelCallUsage(total, fields.result.usage) + ? addModelCallUsage(total, fields.result.usage, reportedCounts) : undefined; case "session_before_compact": return fields.preparation?.isSplitTurn === false ? total : undefined; case "summarization_retry_scheduled": return undefined; case "entry_appended": - return fields.entry?.type === "usage" ? addModelCallUsage(total, fields.entry.usage) : total; + return fields.entry?.type === "usage" ? addModelCallUsage(total, fields.entry.usage, reportedCounts) : total; default: return total; } } +/** A field of a raw provider event, or undefined when `value` is not an object. */ +function rawField(value: unknown, name: string): unknown { + return typeof value === "object" && value !== null ? (value as Record)[name] : undefined; +} + +/** + * Record which counts one raw provider stream event reports, before Pi's parser + * for the call's wire protocol fills an omitted count with 0. Each case reads + * the fields that parser reads (`@earendil-works/pi-ai` 1.0.3). Anthropic + * Messages keeps the counts of `message_start` and then takes only the counts a + * `message_delta` carries, so its reports accumulate; the other protocols + * replace a call's usage with each usage report. For Google, the output count + * is `candidatesTokenCount`, to which Pi adds any thinking tokens. + */ +export function noteReportedCounts(reported: ReportedCounts, api: string, data: unknown): void { + const hasCount = (usage: unknown, name: string): boolean => { + const value = rawField(usage, name); + return typeof value === "number" && Number.isFinite(value) && value >= 0; + }; + const replace = (usage: unknown, input: string, output: string): void => { + if (usage) { + reported.input = hasCount(usage, input); + reported.output = hasCount(usage, output); + } + }; + const type = rawField(data, "type"); + switch (api) { + case "openai-completions": { + const choices = rawField(data, "choices"); + const choice = Array.isArray(choices) ? choices[0] : undefined; + replace(rawField(data, "usage") || rawField(choice, "usage"), "prompt_tokens", "completion_tokens"); + return; + } + case "mistral-conversations": + replace(rawField(data, "usage"), "prompt_tokens", "completion_tokens"); + return; + case "openai-responses": + case "azure-openai-responses": + case "openai-codex-responses": + if (type === "response.completed" || type === "response.incomplete" || type === "response.done") { + replace(rawField(rawField(data, "response"), "usage"), "input_tokens", "output_tokens"); + } + return; + case "bedrock-converse-stream": + replace(rawField(rawField(data, "metadata"), "usage"), "inputTokens", "outputTokens"); + return; + case "google-generative-ai": + case "google-vertex": + replace(rawField(data, "usageMetadata"), "promptTokenCount", "candidatesTokenCount"); + return; + case "pi-messages": + if (type === "done" || type === "error") { + replace(rawField(data, "usage"), "input", "output"); + } + return; + case "anthropic-messages": + if (type === "message_start") { + replace(rawField(rawField(data, "message"), "usage"), "input_tokens", "output_tokens"); + } else if (type === "message_delta") { + const usage = rawField(data, "usage"); + if (usage) { + reported.input ||= hasCount(usage, "input_tokens"); + reported.output ||= hasCount(usage, "output_tokens"); + } + } + return; + } +} + +/** + * Record which counts each model call's provider reported, keyed by the call's + * final usage object. Pi hands every raw provider stream event to + * `onProviderStreamEvent` before parsing it. The agent loop and compaction + * summaries both call the agent's stream function, so wrapping it covers both, + * and Pi passes the final usage object unchanged to `message_end` and + * `compaction_end`. The record is stored before Pi emits either event. Pi's + * cache warmer replays a request with its options but sends it through its own + * model runtime, so the record is copied when the call ends and a warm-up's + * usage entry has no record. + */ +function recordReportedCounts(agent: AgentSession["agent"], reportedCounts: WeakMap): void { + const streamFunction = agent.streamFunction; + agent.streamFunction = async (model, context, options) => { + const reported: ReportedCounts = { input: false, output: false }; + const response = await streamFunction(model, context, { + ...options, + onProviderStreamEvent: async (data, eventModel) => { + noteReportedCounts(reported, model.api, data); + await options?.onProviderStreamEvent?.(data, eventModel); + }, + }); + void response.result().then( + (message) => { + if (typeof message.usage === "object" && message.usage !== null) { + reportedCounts.set(message.usage, { ...reported }); + } + }, + () => undefined, + ); + return response; + }; +} + export function withCustomBaseUrl( catalogModel: T, baseUrl: string | null | undefined, @@ -975,6 +1121,7 @@ class PiSdkSessionHandle implements PiSessionHandle { private cumulativeTurnCount = 0; private readonly runtimeDir: string; private readonly compactionListeners: Set<(event: unknown) => void>; + private readonly reportedCounts = new WeakMap(); constructor( session: AgentSession, @@ -988,6 +1135,7 @@ class PiSdkSessionHandle implements PiSessionHandle { this.runtimeDir = runtimeDir; this.compactionListeners = compactionListeners; this.relay = relay; + recordReportedCounts(session.agent, this.reportedCounts); this.unsubscribeTurnCounter = this.session.subscribe((event) => { if (event.type === "turn_start") { this.cumulativeTurnCount += 1; @@ -1006,6 +1154,8 @@ class PiSdkSessionHandle implements PiSessionHandle { | { role: "assistant"; content: unknown; stopReason: string; errorMessage?: string } | undefined; let usage: PiTokenUsage | undefined = NO_TOKENS; + const reportedCounts: ReportedCountsLookup = (callUsage) => + typeof callUsage === "object" && callUsage !== null ? this.reportedCounts.get(callUsage) : undefined; const unsubscribe = this.session.subscribe((event) => { if (event.type === "turn_start") { turnStarted = true; @@ -1013,10 +1163,10 @@ class PiSdkSessionHandle implements PiSessionHandle { if (event.type === "message_end" && event.message.role === "assistant") { finalAssistant = event.message; } - usage = addEventUsage(usage, event); + usage = addEventUsage(usage, event, reportedCounts); }); const onCompactionPrepared = (event: unknown): void => { - usage = addEventUsage(usage, event); + usage = addEventUsage(usage, event, reportedCounts); }; this.compactionListeners.add(onCompactionPrepared); try { diff --git a/adapters/typescript/pi/src/runtime.ts b/adapters/typescript/pi/src/runtime.ts index f1dfd30a5..7488d1e87 100644 --- a/adapters/typescript/pi/src/runtime.ts +++ b/adapters/typescript/pi/src/runtime.ts @@ -21,13 +21,15 @@ export type PiStopReason = "stop" | "length" | "toolUse" | "error" | "aborted" | /** * Token counts Pi recorded during one prompt. Pi's `input` excludes cache - * reads and writes. + * reads and writes. `input`, `cacheRead` and `cacheWrite` are absent together + * when a call's provider left its input count unreported, and `output` is + * absent when one left its output count unreported. */ export interface PiTokenUsage { - input: number; - output: number; - cacheRead: number; - cacheWrite: number; + input?: number; + output?: number; + cacheRead?: number; + cacheWrite?: number; } export interface PiPromptOutcome { @@ -38,7 +40,7 @@ export interface PiPromptOutcome { stopReason?: PiStopReason; errorMessage?: string; shutdownRequested?: boolean; - /** Absent when any call of the prompt carried no complete usage. */ + /** Absent when no count of the prompt's calls is known. */ usage?: PiTokenUsage; } @@ -61,18 +63,29 @@ function failed(code: string, message: string): AgentRunResult { }; } -/** The invocation's usage, with input counting its cache reads and writes. */ +/** + * The invocation's usage, with input counting its cache reads and writes. A + * count that a call left unreported stays absent, and so does the total unless + * both input and output are known. + */ function agentUsage(usage: PiTokenUsage | undefined): AgentUsage | undefined { if (usage === undefined) { return undefined; } - const input = usage.input + usage.cacheRead + usage.cacheWrite; + const { output, cacheRead, cacheWrite } = usage; + const input = + usage.input !== undefined && cacheRead !== undefined && cacheWrite !== undefined + ? usage.input + cacheRead + cacheWrite + : undefined; + if (input === undefined && output === undefined) { + return undefined; + } return { - input_tokens: input, - cached_input_tokens: usage.cacheRead, - input_tokens_include_cache: true, - output_tokens: usage.output, - total_tokens: input + usage.output, + ...(input === undefined + ? {} + : { input_tokens: input, cached_input_tokens: cacheRead, input_tokens_include_cache: true }), + ...(output === undefined ? {} : { output_tokens: output }), + ...(input === undefined || output === undefined ? {} : { total_tokens: input + output }), }; } diff --git a/adapters/typescript/pi/test/pi-sdk.test.mjs b/adapters/typescript/pi/test/pi-sdk.test.mjs index e82e0807d..a356e41b9 100644 --- a/adapters/typescript/pi/test/pi-sdk.test.mjs +++ b/adapters/typescript/pi/test/pi-sdk.test.mjs @@ -11,6 +11,7 @@ import { createServer } from "node:http"; import { addEventUsage, addUsage, + noteReportedCounts, buildCatalogModel, modelAwareCompactionReserveTokens, PiSdkSessionFactory, @@ -1443,6 +1444,11 @@ test("accepts a well-formed extensions.cost", async () => { } }); +// Lookups for what a model call's provider reported: every count, or a call +// whose raw stream events the adapter never saw. +const allReported = () => ({ input: true, output: true }); +const unseen = () => undefined; + test("sums each assistant message's usage and drops an incomplete one", () => { const zero = { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 }; const first = addUsage(zero, { input: 10, output: 2, cacheRead: 5, cacheWrite: 1, totalTokens: 18 }); @@ -1465,7 +1471,10 @@ test("counts the tokens Pi bills outside assistant messages", () => { { type: "turn_start" }, ]; - const total = events.reduce(addEventUsage, { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 }); + const total = events.reduce( + (sum, event) => addEventUsage(sum, event, allReported), + { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 }, + ); assert.deepEqual(total, { input: 135, output: 4, cacheRead: 0, cacheWrite: 0 }); }); @@ -1473,8 +1482,11 @@ test("counts the tokens Pi bills outside assistant messages", () => { test("a compaction without a usage report leaves the prompt's tokens unknown", () => { const zero = { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 }; - assert.equal(addEventUsage(zero, { type: "compaction_end", aborted: false, result: { summary: "s" } }), undefined); - assert.equal(addEventUsage(zero, { type: "compaction_end", aborted: true }), undefined); + assert.equal( + addEventUsage(zero, { type: "compaction_end", aborted: false, result: { summary: "s" } }, allReported), + undefined, + ); + assert.equal(addEventUsage(zero, { type: "compaction_end", aborted: true }, allReported), undefined); }); test("an aborted assistant message leaves the prompt's tokens unknown", () => { @@ -1482,29 +1494,86 @@ test("an aborted assistant message leaves the prompt's tokens unknown", () => { const usage = { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 }; assert.equal( - addEventUsage(zero, { type: "message_end", message: { role: "assistant", stopReason: "aborted", usage } }), + addEventUsage( + zero, + { type: "message_end", message: { role: "assistant", stopReason: "aborted", usage } }, + allReported, + ), undefined, ); }); -test("a model call with zero-filled usage leaves the prompt's tokens unknown", () => { +test("a model call's counts are known only as far as its provider reported them", () => { const zero = { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 }; - const unreported = { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, totalTokens: 0 }; - - assert.equal(addEventUsage(zero, { type: "message_end", message: { role: "assistant", usage: unreported } }), undefined); - assert.equal(addEventUsage(zero, { type: "compaction_end", aborted: false, result: { usage: unreported } }), undefined); - assert.equal(addEventUsage(zero, { type: "entry_appended", entry: { type: "usage", usage: unreported } }), undefined); + const zeroFilled = { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, totalTokens: 0 }; + const assistant = { type: "message_end", message: { role: "assistant", usage: zeroFilled } }; + const inputOnly = () => ({ input: true, output: false }); + + // Pi fills an omitted count with 0, so a call the adapter did not see is unknown. + assert.equal(addEventUsage(zero, assistant, unseen), undefined); + assert.equal(addEventUsage(zero, { type: "compaction_end", aborted: false, result: { usage: zeroFilled } }, unseen), undefined); + assert.equal(addEventUsage(zero, { type: "entry_appended", entry: { type: "usage", usage: zeroFilled } }, unseen), undefined); + // A reported zero is a measured zero; an omitted count stays unknown. + assert.deepEqual(addEventUsage(zero, assistant, allReported), zero); + assert.deepEqual(addEventUsage(zero, assistant, inputOnly), { input: 0, cacheRead: 0, cacheWrite: 0 }); + assert.equal(addEventUsage(zero, assistant, () => ({ input: false, output: false })), undefined); // A tool reports its own usage, so a zero there is a measured zero. - assert.deepEqual(addEventUsage(zero, { type: "message_end", message: { role: "toolResult", usage: unreported } }), zero); + assert.deepEqual( + addEventUsage(zero, { type: "message_end", message: { role: "toolResult", usage: zeroFilled } }, unseen), + zero, + ); +}); + +test("reads which counts each wire protocol's raw usage reports", () => { + const counts = (api, events) => { + const reported = { input: false, output: false }; + for (const event of events) { + noteReportedCounts(reported, api, event); + } + return reported; + }; + + // Anthropic keeps message_start's counts and takes the counts a delta carries. + assert.deepEqual( + counts("anthropic-messages", [ + { type: "message_start", message: { usage: { input_tokens: 10, output_tokens: 1 } } }, + { type: "message_delta", usage: { output_tokens: 5 } }, + ]), + { input: true, output: true }, + ); + assert.deepEqual( + counts("anthropic-messages", [ + { type: "message_start", message: { usage: { output_tokens: 1 } } }, + { type: "message_delta", usage: { output_tokens: 5 } }, + ]), + { input: false, output: true }, + ); + // The other protocols replace a call's usage with each report. + assert.deepEqual( + counts("openai-completions", [{ choices: [], usage: { prompt_tokens: 7 } }, { choices: [], usage: { completion_tokens: 2 } }]), + { input: false, output: true }, + ); + assert.deepEqual( + counts("openai-responses", [{ type: "response.completed", response: { usage: { input_tokens: 3 } } }]), + { input: true, output: false }, + ); + assert.deepEqual( + counts("google-generative-ai", [{ usageMetadata: { promptTokenCount: 4, candidatesTokenCount: 0 } }]), + { input: true, output: true }, + ); + assert.deepEqual( + counts("bedrock-converse-stream", [{ metadata: { usage: { outputTokens: 6 } } }]), + { input: false, output: true }, + ); }); test("a split-turn compaction leaves the prompt's tokens unknown", () => { const zero = { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 }; const prepared = (preparation) => ({ type: "session_before_compact", preparation }); - assert.deepEqual(addEventUsage(zero, prepared({ isSplitTurn: false })), zero); - assert.equal(addEventUsage(zero, prepared({ isSplitTurn: true })), undefined); - assert.equal(addEventUsage(zero, prepared(undefined)), undefined); + assert.deepEqual(addEventUsage(zero, prepared({ isSplitTurn: false }), allReported), zero); + assert.equal(addEventUsage(zero, prepared({ isSplitTurn: true }), allReported), undefined); + assert.equal(addEventUsage(zero, prepared(undefined), allReported), undefined); }); test("a retried compaction summary leaves the prompt's tokens unknown", () => { @@ -1516,7 +1585,13 @@ test("a retried compaction summary leaves the prompt's tokens unknown", () => { { type: "compaction_end", aborted: false, result: { summary: "s", usage } }, ]; - assert.equal(events.reduce(addEventUsage, { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 }), undefined); + assert.equal( + events.reduce( + (sum, event) => addEventUsage(sum, event, allReported), + { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 }, + ), + undefined, + ); }); /** @@ -1619,6 +1694,30 @@ test( }, ); +test( + "leaves a count the provider omitted unknown and keeps reported zeros", + { skip: supportsPi ? false : "Pi 1.0 requires Node 22.19 or newer" }, + async () => { + const [inputOnly, outputOnly, zeros, afterPartial] = await promptsAgainstScriptedEndpoint( + [ + { content: "input only", usage: { prompt_tokens: 100 } }, + { content: "output only", usage: { completion_tokens: 5 } }, + { content: "zeros", usage: { prompt_tokens: 0, completion_tokens: 0 } }, + { readNote: true, usage: { prompt_tokens: 60 } }, + { content: "both", usage: { prompt_tokens: 70, completion_tokens: 7 } }, + ], + ["one", "two", "three", "four"], + ); + + // Pi fills an omitted count with 0; the provider's raw report decides. + assert.deepEqual(inputOnly.usage, { input: 100, cacheRead: 0, cacheWrite: 0 }); + assert.deepEqual(outputOnly.usage, { output: 5 }); + assert.deepEqual(zeros.usage, { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 }); + // One call without an output count leaves the prompt's output unknown. + assert.deepEqual(afterPartial.usage, { input: 130, cacheRead: 0, cacheWrite: 0 }); + }, +); + test( "counts a compaction that Pi summarizes in one call", { skip: supportsPi ? false : "Pi 1.0 requires Node 22.19 or newer" }, diff --git a/adapters/typescript/pi/test/runtime.test.mjs b/adapters/typescript/pi/test/runtime.test.mjs index 9269ee643..0573e9c25 100644 --- a/adapters/typescript/pi/test/runtime.test.mjs +++ b/adapters/typescript/pi/test/runtime.test.mjs @@ -583,3 +583,35 @@ test("reports the invocation's tokens with cache reads and writes inside the inp total_tokens: 460, }); }); + +test("reports only the token counts the prompt's provider reported", async () => { + const outcomes = [ + { input: 100, cacheRead: 30, cacheWrite: 0 }, + { output: 5 }, + ]; + const runtime = new PiAdapterRuntime({ + async create() { + let turnCount = 0; + return { + async prompt() { + turnCount += 1; + const usage = outcomes.shift(); + return { accepted: true, turnStarted: true, turnCount, text: "ok", stopReason: "stop", usage }; + }, + async stop() {}, + }; + }, + }); + + await runtime.start(startInput()); + const inputOnly = await runtime.invoke({ input: "one" }, context); + const outputOnly = await runtime.invoke({ input: "two" }, context); + await runtime.stop(); + + assert.deepEqual(inputOnly.usage, { + input_tokens: 130, + cached_input_tokens: 30, + input_tokens_include_cache: true, + }); + assert.deepEqual(outputOnly.usage, { output_tokens: 5 }); +}); diff --git a/tests/adapters/test_hermes_adapter.py b/tests/adapters/test_hermes_adapter.py index f972fc1a1..5d5b824b4 100644 --- a/tests/adapters/test_hermes_adapter.py +++ b/tests/adapters/test_hermes_adapter.py @@ -1625,7 +1625,8 @@ def test_invocation_usage_with_lower_or_missing_counters_leaves_them_unknown(): class _UsageRecordingAgent: """Hermes agent stand-in whose turns run scripted provider responses - through Hermes' own usage accounting and context compressor.""" + through Hermes' own usage accounting and context compressor, called the + way Hermes' turn loop calls it.""" def __init__(self, turns: list[list[dict[str, int] | None]]) -> None: from agent.context_compressor import ContextCompressor @@ -1659,11 +1660,11 @@ def __init__(self, turns: list[list[dict[str, int] | None]]) -> None: setattr(self, f"session_{counter}", 0) def run_conversation(self, user_message: str) -> dict[str, object]: - from agent.turn_usage import record_response_usage + from agent import turn_response_check messages = [{"role": "user", "content": user_message}] for usage in self._turns.pop(0): - record_response_usage( + turn_response_check.record_response_usage( self, SimpleNamespace(usage=usage), messages=messages, @@ -1688,6 +1689,31 @@ def run_conversation(self, user_message: str) -> dict[str, object]: } +def _usage_recording_runtime(agent: _UsageRecordingAgent) -> adapter.HermesRuntime: + agent_config = _agent_config( + {"models": {"default": {"provider": "test", "model": "test-model"}}} + ) + runtime = adapter.HermesRuntime() + runtime._started = True + runtime._agent_config = agent_config + runtime._model_config = agent_config.models["default"] + runtime._runtime_id = "runtime-usage" + runtime._agent = agent + runtime._observe_usage_reports(agent) + return runtime + + +def _usage_invocation() -> tuple[AgentRunRequest, RuntimeContext]: + return _invocation( + { + "runtime_context": _runtime_context( + runtime_id="runtime-usage" + ).to_mapping(), + "request": {"input": "next"}, + } + ) + + async def test_hermes_usage_from_provider_responses_keeps_unreported_tokens_unknown(): def reported(prompt: int, completion: int) -> dict[str, int]: return { @@ -1705,24 +1731,8 @@ def reported(prompt: int, completion: int) -> dict[str, int]: [reported(400, 40)], ] ) - agent_config = _agent_config( - {"models": {"default": {"provider": "test", "model": "test-model"}}} - ) - runtime = adapter.HermesRuntime() - runtime._started = True - runtime._agent_config = agent_config - runtime._model_config = agent_config.models["default"] - runtime._runtime_id = "runtime-usage" - runtime._agent = agent - runtime._count_usage_less_responses(agent) - invocation = _invocation( - { - "runtime_context": _runtime_context( - runtime_id="runtime-usage" - ).to_mapping(), - "request": {"input": "next"}, - } - ) + runtime = _usage_recording_runtime(agent) + invocation = _usage_invocation() two_calls = await runtime.invoke(*invocation) reported_zero = await runtime.invoke(*invocation) @@ -1746,6 +1756,49 @@ def reported(prompt: int, completion: int) -> dict[str, int]: assert after_mixed.usage.output_tokens == 40 +async def test_hermes_usage_with_partial_provider_report_leaves_omitted_count_unknown(): + agent = _UsageRecordingAgent( + [ + [{"prompt_tokens": 100}], + [{"completion_tokens": 5}], + [{"prompt_tokens": 0, "completion_tokens": 0}], + [{"prompt_tokens": 60, "completion_tokens": 6}, {"prompt_tokens": 70}], + [{"prompt_tokens": 400, "completion_tokens": 40}], + ] + ) + runtime = _usage_recording_runtime(agent) + invocation = _usage_invocation() + + input_only = await runtime.invoke(*invocation) + output_only = await runtime.invoke(*invocation) + reported_zero = await runtime.invoke(*invocation) + one_partial = await runtime.invoke(*invocation) + after_partial = await runtime.invoke(*invocation) + + # Hermes normalizes an omitted count to 0; the provider's report decides. + assert input_only.usage is not None + assert input_only.usage.input_tokens == 100 + assert input_only.usage.output_tokens is None + assert input_only.usage.total_tokens is None + assert output_only.usage is not None + assert output_only.usage.input_tokens is None + assert output_only.usage.cached_input_tokens is None + assert output_only.usage.output_tokens == 5 + assert output_only.usage.total_tokens is None + assert reported_zero.usage is not None + assert reported_zero.usage.input_tokens == 0 + assert reported_zero.usage.output_tokens == 0 + assert reported_zero.usage.total_tokens == 0 + # One response without an output count leaves the invocation's output unknown. + assert one_partial.usage is not None + assert one_partial.usage.input_tokens == 130 + assert one_partial.usage.output_tokens is None + assert after_partial.usage is not None + assert after_partial.usage.input_tokens == 400 + assert after_partial.usage.output_tokens == 40 + assert after_partial.usage.total_tokens == 440 + + @pytest.mark.parametrize( ("request_input", "expected_message"), [ From 0bc8fcf6f57afa804fd0bb6e562ceab378688deb Mon Sep 17 00:00:00 2001 From: mnajafian-nv Date: Mon, 19 Oct 2026 14:48:31 -0700 Subject: [PATCH 09/68] fix: check cache fields and the Hermes call site for partial usage Anthropic Messages reports input_tokens without cache reads and writes, and both Hermes and Pi add a missing cache field back as 0. A response with input_tokens but no cache fields therefore still produced a measured input count. Input is now reported only when input_tokens and both cache fields are present, the rule the Claude adapter uses. When a provider sends no cache details, both SDKs count 0 cached tokens. cached_input_tokens is now left out in that case. For the protocols whose input count includes the cache, input_tokens and total_tokens stay known. The Hermes adapter replaced turn_response_check.record_response_usage on every start and never restored it, and nothing showed that Hermes' turn loop still calls that name. The wrapper is now installed once per process and dispatches to observers registered per agent, and stop() removes the agent's observer. Start fails with a lifecycle error unless record_response_usage is still Hermes' turn_usage function and check_api_response still calls it by that name. The mock endpoint takes a usage object per scenario, and the Hermes e2e test sends a partial report through Hermes' real turn loop. For Pi, only a message_delta reports the final Anthropic output count, because message_start carries a placeholder. Google output counts as reported when thoughtsTokenCount is present without candidatesTokenCount. The protocol comes from the model that served the stream event, and a protocol outside Pi's built-in ones reports no counts. PiTokenUsage now holds input with the cache included, so the cache reads can be unknown while input stays known. The READMEs now say which Hermes paths the check covers. Hermes' native Gemini and Bedrock transports, and its Anthropic transport when a stream omits usage, fill missing counts before the adapter sees them, and its Codex app-server runtime records usage elsewhere. They also say that a Pi cache warm-up during an invocation leaves the result without usage: the adapter receives the warm-up's raw events but cannot tie them to the warm-up's usage entry. Signed-off-by: mnajafian-nv --- adapters/python/hermes/README.md | 38 +++- .../nemo_fabric_adapters/hermes/adapter.py | 180 +++++++++++---- adapters/typescript/pi/README.md | 45 ++-- adapters/typescript/pi/src/pi-sdk.ts | 208 +++++++++++++----- adapters/typescript/pi/src/runtime.ts | 27 +-- adapters/typescript/pi/test/pi-sdk.test.mjs | 142 +++++++----- adapters/typescript/pi/test/runtime.test.mjs | 14 +- tests/_utils/mock_api_server.py | 32 +-- tests/adapters/test_hermes_adapter.py | 94 +++++++- tests/e2e/test_hermes_e2e.py | 16 ++ 10 files changed, 566 insertions(+), 230 deletions(-) diff --git a/adapters/python/hermes/README.md b/adapters/python/hermes/README.md index 4d738c249..fa1a57ea2 100644 --- a/adapters/python/hermes/README.md +++ b/adapters/python/hermes/README.md @@ -119,18 +119,32 @@ session per runtime, so each result's `usage` is the change since the previous invocation: `input_tokens` with cache reads and cache writes included, `cached_input_tokens`, `output_tokens`, and `total_tokens`. Hermes records auxiliary calls, such as context-compression summaries, separately from these -counters, so they are not included. Hermes adds no tokens for a provider -response that carried no usage, so an invocation with such a response has no -`usage`. A response that reported zero tokens still counts as 0. Hermes also -reads a count that a response's usage left out as 0. When a response left out -its input count, the invocation's `usage` leaves out `input_tokens`, -`cached_input_tokens`, and `total_tokens`. When it left out its output count, -the `usage` leaves out `output_tokens` and `total_tokens`. A counter that -Hermes does not report, or that went down, is left out for that invocation. -After an invocation that raised, was cancelled, or timed out, or whose result -held no counters, the adapter has no baseline. The next invocation then reports -no token counters rather than a carried-over or whole-session total, and the -one after it reports its own change again. +counters, so they are not included. A counter that Hermes does not report, or +that went down, is left out for that invocation. After an invocation that +raised, was cancelled, or timed out, or whose result held no counters, the +adapter has no baseline. The next invocation then reports no token counters +rather than a carried-over or whole-session total, and the one after it reports +its own change again. + +Hermes adds no tokens for a provider response that carried no usage, so an +invocation with such a response has no `usage`. Hermes also reads a count that +a response's usage left out as 0, so the adapter checks each response's usage +before Hermes normalizes it. A reported zero still counts as 0. When a response +left out its input count, the invocation's `usage` leaves out `input_tokens`, +`cached_input_tokens`, and `total_tokens`. Anthropic Messages usage reports +input without the cache, so its input count also needs both cache fields. When +a response left out its output count, the `usage` leaves out `output_tokens` +and `total_tokens`. When a response sent no cache details, the `usage` leaves +out `cached_input_tokens`. + +This check covers the usage that the OpenAI-compatible, Responses, and +Anthropic clients return. Hermes' native Gemini and Bedrock transports, and its +Anthropic transport when a stream omits usage, fill missing counts with 0 +before the adapter sees them. Hermes' Codex app-server runtime records usage +outside this path. Hermes passes a response's raw usage to no per-agent hook, +so the adapter wraps `agent.turn_response_check.record_response_usage` once per +process. The runtime fails to start if the installed Hermes records usage +through another call. ## Execution Model diff --git a/adapters/python/hermes/src/nemo_fabric_adapters/hermes/adapter.py b/adapters/python/hermes/src/nemo_fabric_adapters/hermes/adapter.py index a1a093fea..57b24e8f3 100755 --- a/adapters/python/hermes/src/nemo_fabric_adapters/hermes/adapter.py +++ b/adapters/python/hermes/src/nemo_fabric_adapters/hermes/adapter.py @@ -15,6 +15,8 @@ import logging import os import sys +import weakref +from collections.abc import Callable from contextlib import redirect_stdout from io import StringIO from pathlib import Path @@ -71,8 +73,8 @@ def __init__(self) -> None: # from this baseline. A counter missing from it has no baseline. self._session_token_totals: dict[str, int] = _fresh_session_totals() self._usage_less_responses = 0 - # Provider responses whose usage left out the input or output count. - self._unreported_counts = {"input": 0, "output": 0} + # Provider responses whose usage left out a count, by count. + self._unreported_counts = dict.fromkeys(_REPORTED_COUNTS, 0) async def start(self, payload: dict[str, Any]) -> None: if self._started: @@ -386,10 +388,11 @@ def clear_active_invoke_task( usage, self._session_token_totals = _invocation_usage( result, previous_totals, - input_reported=self._unreported_counts["input"] - == unreported_before["input"], - output_reported=self._unreported_counts["output"] - == unreported_before["output"], + unreported=frozenset( + count + for count, responses in self._unreported_counts.items() + if responses != unreported_before[count] + ), ) if self._usage_less_responses != usage_less_before: # A provider response without usage added no tokens, so the @@ -429,34 +432,24 @@ def count_and_note() -> None: compressor.note_usage_less_response = count_and_note def _count_unreported_counts(self, agent: Any) -> None: - """Count the agent's provider responses whose usage left out the input - or the output count. + """Count the agent's provider responses whose usage left out a count. Hermes reads a missing count as 0, so its counters cannot tell an - omitted count from a reported zero. Hermes' turn loop passes every - response to ``agent.turn_response_check.record_response_usage``, so - the adapter wraps that call and reads the raw usage before Hermes - normalizes it. Responses of other agents, such as delegated - subagents, are not counted: their tokens never reach this agent's - counters. The lifecycle host serves one runtime per process, so the - wrapper is installed once. + omitted count from a reported zero, and Hermes passes the raw usage to + no per-agent hook. The adapter therefore registers the agent with a + process-wide wrapper around the function Hermes' turn loop calls for + each response; see ``_install_usage_observer``. """ - from agent import turn_response_check - - record_response_usage = turn_response_check.record_response_usage + _install_usage_observer() - def count_and_record( - usage_agent: Any, response: Any, *args: Any, **kwargs: Any - ) -> Any: + def count_unreported(response: Any) -> None: usage = getattr(response, "usage", None) # A response without usage is counted by the compressor hook. - if usage_agent is agent and usage: - for count, fields in _RAW_COUNT_FIELDS.items(): - if not any(_is_reported_count(usage, field) for field in fields): - self._unreported_counts[count] += 1 - return record_response_usage(usage_agent, response, *args, **kwargs) + if usage: + for count in _omitted_counts(agent, usage): + self._unreported_counts[count] += 1 - turn_response_check.record_response_usage = count_and_record + _USAGE_OBSERVERS[agent] = count_unreported async def _authenticate_mcp_servers(self) -> None: if self._mcp_authentication_checked: @@ -573,6 +566,8 @@ async def stop(self) -> None: agent = self._agent session_db = self._session_db + if agent is not None: + _USAGE_OBSERVERS.pop(agent, None) had_mcp_servers = bool(self._hermes_config.get("mcp_servers")) had_relay_plugin = self._relay_plugin_config_path is not None self._agent = None @@ -678,18 +673,114 @@ def _invoke_hermes_turn( ) -# The raw usage fields Hermes reads for a response's input and output counts. -# Chat Completions usage may use either spelling; Anthropic Messages and -# Responses usage carry only input_tokens and output_tokens. -_RAW_COUNT_FIELDS = { - "input": ("prompt_tokens", "input_tokens"), - "output": ("completion_tokens", "output_tokens"), +# Response observers by Hermes agent, called before Hermes records a +# response's usage. Delegated subagents are never registered: their tokens +# never reach the registered agent's counters. +_USAGE_OBSERVERS: weakref.WeakKeyDictionary[Any, Callable[[Any], None]] = ( + weakref.WeakKeyDictionary() +) + + +def _observed_record_response_usage( + agent: Any, response: Any, *args: Any, **kwargs: Any +) -> Any: + """Show ``response`` to the agent's observer, then record it as Hermes does.""" + from agent import turn_usage + + observer = _USAGE_OBSERVERS.get(agent) + if observer is not None: + observer(response) + return turn_usage.record_response_usage(agent, response, *args, **kwargs) + + +def _install_usage_observer() -> None: + """Route Hermes' per-response usage accounting through the observers. + + ``agent.turn_response_check.check_api_response`` calls + ``record_response_usage`` for each provider response, looked up as a + global of its module, and that function is the first place the raw usage + reaches before Hermes normalizes it. The wrapper replaces that module + global once per process and always delegates to Hermes' function. If + Hermes no longer calls usage accounting through that name, start fails + rather than reporting omitted counts as measured zeros. + """ + from agent import turn_response_check + from agent import turn_usage + + current = turn_response_check.record_response_usage + if current is _observed_record_response_usage: + return + if ( + current is not turn_usage.record_response_usage + or "record_response_usage" + not in turn_response_check.check_api_response.__code__.co_names + ): + raise lifecycle.LifecycleError( + "hermes_usage_accounting_unsupported", + "This Hermes version records response usage outside " + "agent.turn_response_check.record_response_usage", + ) + turn_response_check.record_response_usage = _observed_record_response_usage + + +_REPORTED_COUNTS = ("input", "cached", "output") +# The raw usage fields Hermes' normalize_usage reads for each count, by the +# usage shape it picks (agent/usage_pricing.py). A count is reported when every +# field of one alternative is present. Anthropic's input_tokens excludes cache +# reads and writes, so its input needs both cache fields, as in the Claude +# adapter. The other shapes' input count includes the cache, so a missing cache +# detail leaves only the cached count unknown. +_ANTHROPIC_FIELDS = { + "input": ( + ("input_tokens", "cache_read_input_tokens", "cache_creation_input_tokens"), + ), + "cached": (("cache_read_input_tokens",),), + "output": (("output_tokens",),), +} +_RESPONSES_FIELDS = { + "input": (("input_tokens",),), + "cached": (("input_tokens_details.cached_tokens",),), + "output": (("output_tokens",),), +} +_CHAT_FIELDS = { + "input": (("prompt_tokens",), ("input_tokens",)), + "cached": ( + ("prompt_tokens_details.cached_tokens",), + ("cache_read_input_tokens",), + ("prompt_cache_hit_tokens",), + ("cached_tokens",), + ), + "output": (("completion_tokens",), ("output_tokens",)), } -def _is_reported_count(usage: Any, field: str) -> bool: - """Return whether the provider's raw usage reports ``field`` as a count.""" - value = usage.get(field) if isinstance(usage, dict) else getattr(usage, field, None) +def _omitted_counts(agent: Any, usage: Any) -> list[str]: + """Return the counts that the provider's raw ``usage`` left out.""" + mode = str(getattr(agent, "api_mode", None) or "").strip().lower() + provider = str(getattr(agent, "provider", None) or "").strip().lower() + if mode == "anthropic_messages" or provider == "anthropic": + fields = _ANTHROPIC_FIELDS + elif mode == "codex_responses": + fields = _RESPONSES_FIELDS + else: + fields = _CHAT_FIELDS + return [ + count + for count, alternatives in fields.items() + if not any( + all(_is_reported_count(usage, path) for path in alternative) + for alternative in alternatives + ) + ] + + +def _is_reported_count(usage: Any, path: str) -> bool: + """Return whether the raw usage reports the dotted ``path`` as a count.""" + value = usage + for field in path.split("."): + value = ( + value.get(field) if isinstance(value, dict) else getattr(value, field, None) + ) return isinstance(value, int) and not isinstance(value, bool) and value >= 0 @@ -702,8 +793,7 @@ def _invocation_usage( result: dict[str, Any], previous_totals: dict[str, int], *, - input_reported: bool = True, - output_reported: bool = True, + unreported: frozenset[str] = frozenset(), ) -> tuple[AgentUsage | None, dict[str, int]]: """Return this invocation's usage and the session totals to keep as baseline. @@ -712,9 +802,10 @@ def _invocation_usage( result or from ``previous_totals`` (no baseline: the previous invocation failed or its result held no counters), or that went down, is unknown for this invocation. The usage therefore never holds a whole-session total or - another invocation's tokens. ``input_reported`` or ``output_reported`` is - False when a provider response of this invocation left that count out, - which Hermes counted as 0, so that count and the total are unknown. + another invocation's tokens. ``unreported`` names the counts ("input", + "cached", "output") that a provider response of this invocation left out + and Hermes counted as 0. Those counts are unknown, and so is the total when + input or output is. """ totals = { key: value @@ -729,7 +820,7 @@ def _invocation_usage( if key in previous_totals and value >= previous_totals[key] } fields: dict[str, Any] = {} - if input_reported and all( + if "input" not in unreported and all( key in change for key in ("input_tokens", "cache_read_tokens", "cache_write_tokens") ): @@ -738,9 +829,10 @@ def _invocation_usage( + change["cache_read_tokens"] + change["cache_write_tokens"] ) - fields["cached_input_tokens"] = change["cache_read_tokens"] + if "cached" not in unreported: + fields["cached_input_tokens"] = change["cache_read_tokens"] fields["input_tokens_include_cache"] = True - if output_reported and "output_tokens" in change: + if "output" not in unreported and "output_tokens" in change: fields["output_tokens"] = change["output_tokens"] if "input_tokens" in fields and "output_tokens" in fields: fields["total_tokens"] = fields["input_tokens"] + fields["output_tokens"] diff --git a/adapters/typescript/pi/README.md b/adapters/typescript/pi/README.md index 5758b031b..ccb3edeec 100644 --- a/adapters/typescript/pi/README.md +++ b/adapters/typescript/pi/README.md @@ -218,24 +218,33 @@ Each result's `usage` sums the tokens Pi records during the invocation, counted the way Pi's own usage totals count them: assistant messages, tool results that report usage, finished compaction summaries, and usage entries such as cache warm-ups. It holds `input_tokens` with cache reads and cache writes included, -`cached_input_tokens`, `output_tokens`, and `total_tokens`. Pi fills a count -that the provider did not report with 0, so the adapter reads the provider's -raw stream events for each assistant message and compaction summary call and -keeps only the counts the provider reported. A reported zero counts as 0. When -a call left out its input count, the result leaves out `input_tokens`, -`cached_input_tokens`, and `total_tokens`. When a call left out its output -count, the result leaves out `output_tokens` and `total_tokens`. The result has -no `usage` when Pi produced no assistant message, when a model call reported no -counts, when a tool result has no complete usage report, when an assistant -message was aborted, when a compaction was aborted or reported no usage, or -when Pi retried a compaction summary call, because Pi keeps no usage for the -failed attempt. It also has no `usage` when Pi recorded a cache warm-up during -the invocation, because Pi sends warm-ups outside the agent's stream function, -where the adapter cannot read the raw report. It also has no `usage` when a -compaction cuts inside a turn, because Pi then summarizes the earlier history -and the start of that turn in separate calls and reports only their combined -usage. Pi computes its cost figures from the model catalog's rates, so the -adapter does not report them as `cost_usd`. +`cached_input_tokens`, `output_tokens`, and `total_tokens`. Pi computes its +cost figures from the model catalog's rates, so the adapter does not report +them as `cost_usd`. + +Pi fills a count that the provider did not report with 0, so the adapter reads +the provider's raw stream events for each assistant message and compaction +summary call and keeps only the counts the provider reported. A reported zero +counts as 0, including a report of all zeros. When a call left out its input +count, the result leaves out `input_tokens`, `cached_input_tokens`, and +`total_tokens`. Anthropic Messages usage reports input without the cache, so +its input count also needs both cache fields. When a call left out its output +count, the result leaves out `output_tokens` and `total_tokens`; for Anthropic +Messages, only a `message_delta` reports the final output count. When a call +sent no cache details, the result leaves out `cached_input_tokens`. A wire +protocol outside Pi's built-in ones reports no counts the adapter can check, so +its calls leave the result without `usage`. + +The result has no `usage` when Pi produced no assistant message, when a model +call reported none of the counts, when a tool result has no complete usage +report, when an assistant message was aborted, when a compaction was aborted or +reported no usage, or when Pi retried a compaction summary call, because Pi +keeps no usage for the failed attempt. It also has no `usage` when Pi recorded +a cache warm-up during the invocation: Pi replays a warm-up outside the agent's +stream function, so the adapter cannot tie its raw report to the warm-up's +usage entry. It also has no `usage` when a compaction cuts inside a turn, +because Pi then summarizes the earlier history and the start of that turn in +separate calls and reports only their combined usage. ## Custom Tool Modules diff --git a/adapters/typescript/pi/src/pi-sdk.ts b/adapters/typescript/pi/src/pi-sdk.ts index c57d6d165..5294a2e56 100644 --- a/adapters/typescript/pi/src/pi-sdk.ts +++ b/adapters/typescript/pi/src/pi-sdk.ts @@ -69,26 +69,30 @@ interface PiToolFactoryContext { workspace: string; } -const NO_TOKENS: PiTokenUsage = { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 }; +const NO_TOKENS: PiTokenUsage = { input: 0, cacheRead: 0, output: 0 }; /** Which counts of one model call the provider reported before Pi filled the rest with 0. */ export interface ReportedCounts { input: boolean; + cacheRead: boolean; output: boolean; } /** What the provider reported for the model call that produced a usage object, if the adapter saw it. */ export type ReportedCountsLookup = (usage: unknown) => ReportedCounts | undefined; -const BOTH_REPORTED: ReportedCounts = { input: true, output: true }; +const ALL_REPORTED: ReportedCounts = { input: true, cacheRead: true, output: true }; -/** One call's four token counts, or undefined unless each is a non-negative safe integer. */ +/** + * One call's counts from Pi's usage, with input counting its cache reads and + * writes, or undefined unless Pi's four counts are non-negative safe integers. + */ function callCounts(usage: unknown): Required | undefined { if (typeof usage !== "object" || usage === null) { return undefined; } const counts = usage as Record; - const count = (name: keyof PiTokenUsage): number | undefined => { + const count = (name: string): number | undefined => { const value = counts[name]; return typeof value === "number" && Number.isSafeInteger(value) && value >= 0 ? value : undefined; }; @@ -99,25 +103,26 @@ function callCounts(usage: unknown): Required | undefined { if (input === undefined || output === undefined || cacheRead === undefined || cacheWrite === undefined) { return undefined; } - return { input, output, cacheRead, cacheWrite }; + return { input: input + cacheRead + cacheWrite, cacheRead, output }; } /** - * Add a call's counts to a prompt's running total. The input side (input with - * its cache reads and writes) and the output stay known only while every call - * reported them; the total is unknown once neither is. + * Add a call's counts to a prompt's running total. Each count stays known only + * while every call reported it; the cache reads also need the input. The total + * is unknown once neither input nor output is known. */ function sumCounts(total: PiTokenUsage, call: Required, reported: ReportedCounts): PiTokenUsage | undefined { const sum: PiTokenUsage = {}; - if (reported.input && total.input !== undefined && total.cacheRead !== undefined && total.cacheWrite !== undefined) { + if (reported.input && total.input !== undefined) { sum.input = total.input + call.input; - sum.cacheRead = total.cacheRead + call.cacheRead; - sum.cacheWrite = total.cacheWrite + call.cacheWrite; + if (reported.cacheRead && total.cacheRead !== undefined) { + sum.cacheRead = total.cacheRead + call.cacheRead; + } } if (reported.output && total.output !== undefined) { sum.output = total.output + call.output; } - return Object.keys(sum).length === 0 ? undefined : sum; + return sum.input === undefined && sum.output === undefined ? undefined : sum; } /** @@ -126,7 +131,7 @@ function sumCounts(total: PiTokenUsage, call: Required, reported: */ export function addUsage(total: PiTokenUsage | undefined, usage: unknown): PiTokenUsage | undefined { const call = callCounts(usage); - return total === undefined || call === undefined ? undefined : sumCounts(total, call, BOTH_REPORTED); + return total === undefined || call === undefined ? undefined : sumCounts(total, call, ALL_REPORTED); } /** @@ -134,8 +139,8 @@ export function addUsage(total: PiTokenUsage | undefined, usage: unknown): PiTok * parsers fill a count the provider omitted with 0, so `reportedCounts` says * which counts the provider actually reported, from the raw stream events the * session recorded for that call. A count the provider omitted is unknown for - * the whole prompt, and so is every count of a call whose raw events the - * adapter did not see. A reported zero stays a measured zero. + * the whole prompt, and so is every count of a call the session has no + * record for, such as a cache warm-up. A reported zero stays a measured zero. */ function addModelCallUsage( total: PiTokenUsage | undefined, @@ -203,29 +208,106 @@ export function addEventUsage( } } -/** A field of a raw provider event, or undefined when `value` is not an object. */ -function rawField(value: unknown, name: string): unknown { - return typeof value === "object" && value !== null ? (value as Record)[name] : undefined; +/** The value at a dotted `path` of a raw provider event, or undefined. */ +function rawField(value: unknown, path: string): unknown { + let field = value; + for (const name of path.split(".")) { + field = typeof field === "object" && field !== null ? (field as Record)[name] : undefined; + } + return field; } /** - * Record which counts one raw provider stream event reports, before Pi's parser - * for the call's wire protocol fills an omitted count with 0. Each case reads - * the fields that parser reads (`@earendil-works/pi-ai` 1.0.3). Anthropic - * Messages keeps the counts of `message_start` and then takes only the counts a - * `message_delta` carries, so its reports accumulate; the other protocols - * replace a call's usage with each usage report. For Google, the output count - * is `candidatesTokenCount`, to which Pi adds any thinking tokens. + * The raw usage fields Pi's parser reads for each count, by wire protocol + * (`@earendil-works/pi-ai` 1.0.3). Input is reported when every `input` field + * is present, cache reads and output when any of theirs is. Anthropic Messages + * and pi-messages report input without the cache, so their input also needs + * the cache fields, as in the Claude adapter. The other protocols' input count + * includes the cache, so a missing cache detail leaves only the cache reads + * unknown. Google output counts thinking tokens too. */ -export function noteReportedCounts(reported: ReportedCounts, api: string, data: unknown): void { - const hasCount = (usage: unknown, name: string): boolean => { - const value = rawField(usage, name); - return typeof value === "number" && Number.isFinite(value) && value >= 0; +interface UsageFields { + input: string[]; + cacheRead: string[]; + output: string[]; +} +const RESPONSES_FIELDS: UsageFields = { + input: ["input_tokens"], + cacheRead: ["input_tokens_details.cached_tokens"], + output: ["output_tokens"], +}; +const GOOGLE_FIELDS: UsageFields = { + input: ["promptTokenCount"], + cacheRead: ["cachedContentTokenCount"], + output: ["candidatesTokenCount", "thoughtsTokenCount"], +}; +const REPORTED_FIELDS: Record = { + "openai-completions": { + input: ["prompt_tokens"], + cacheRead: ["prompt_tokens_details.cached_tokens", "prompt_cache_hit_tokens", "cached_tokens"], + output: ["completion_tokens"], + }, + "mistral-conversations": { + input: ["prompt_tokens"], + cacheRead: [ + "promptTokensDetails.cachedTokens", + "prompt_tokens_details.cached_tokens", + "promptTokenDetails.cachedTokens", + "prompt_token_details.cached_tokens", + "numCachedTokens", + "num_cached_tokens", + ], + output: ["completion_tokens"], + }, + "openai-responses": RESPONSES_FIELDS, + "azure-openai-responses": RESPONSES_FIELDS, + "openai-codex-responses": RESPONSES_FIELDS, + "bedrock-converse-stream": { + input: ["inputTokens"], + cacheRead: ["cacheReadInputTokens"], + output: ["outputTokens"], + }, + "google-generative-ai": GOOGLE_FIELDS, + "google-vertex": GOOGLE_FIELDS, + "pi-messages": { + input: ["input", "cacheRead", "cacheWrite"], + cacheRead: ["cacheRead"], + output: ["output"], + }, + "anthropic-messages": { + input: ["input_tokens", "cache_read_input_tokens", "cache_creation_input_tokens"], + cacheRead: ["cache_read_input_tokens"], + output: ["output_tokens"], + }, +}; + +/** + * Record the usage fields one raw provider stream event reports, in `seen`, + * before Pi's parser fills an omitted count with 0. Pi reads the usage where + * each protocol sends it. Anthropic Messages keeps the counts of + * `message_start` and then takes only the counts a `message_delta` carries, so + * its fields accumulate; its `message_start` output count is a placeholder + * that the delta replaces, so only a delta reports output. The other + * protocols replace a call's usage with each usage report. + */ +export function noteReportedFields(seen: Set, api: string, data: unknown): void { + const fields = REPORTED_FIELDS[api]; + if (fields === undefined) { + return; + } + const paths = [...fields.input, ...fields.cacheRead, ...fields.output]; + const isCount = (value: unknown): boolean => typeof value === "number" && Number.isFinite(value) && value >= 0; + const add = (usage: unknown, names: string[]): void => { + for (const path of names) { + if (isCount(rawField(usage, path))) { + seen.add(path); + } + } }; - const replace = (usage: unknown, input: string, output: string): void => { + const replace = (usage: unknown, names = paths): void => { if (usage) { - reported.input = hasCount(usage, input); - reported.output = hasCount(usage, output); + seen.clear(); + add(usage, names); } }; const type = rawField(data, "type"); @@ -233,71 +315,83 @@ export function noteReportedCounts(reported: ReportedCounts, api: string, data: case "openai-completions": { const choices = rawField(data, "choices"); const choice = Array.isArray(choices) ? choices[0] : undefined; - replace(rawField(data, "usage") || rawField(choice, "usage"), "prompt_tokens", "completion_tokens"); + replace(rawField(data, "usage") || rawField(choice, "usage")); return; } case "mistral-conversations": - replace(rawField(data, "usage"), "prompt_tokens", "completion_tokens"); + replace(rawField(data, "usage")); return; case "openai-responses": case "azure-openai-responses": case "openai-codex-responses": if (type === "response.completed" || type === "response.incomplete" || type === "response.done") { - replace(rawField(rawField(data, "response"), "usage"), "input_tokens", "output_tokens"); + replace(rawField(data, "response.usage")); } return; case "bedrock-converse-stream": - replace(rawField(rawField(data, "metadata"), "usage"), "inputTokens", "outputTokens"); + replace(rawField(data, "metadata.usage")); return; case "google-generative-ai": case "google-vertex": - replace(rawField(data, "usageMetadata"), "promptTokenCount", "candidatesTokenCount"); + replace(rawField(data, "usageMetadata")); return; case "pi-messages": if (type === "done" || type === "error") { - replace(rawField(data, "usage"), "input", "output"); + replace(rawField(data, "usage")); } return; case "anthropic-messages": if (type === "message_start") { - replace(rawField(rawField(data, "message"), "usage"), "input_tokens", "output_tokens"); - } else if (type === "message_delta") { - const usage = rawField(data, "usage"); - if (usage) { - reported.input ||= hasCount(usage, "input_tokens"); - reported.output ||= hasCount(usage, "output_tokens"); - } + replace(rawField(data, "message.usage"), [...fields.input, ...fields.cacheRead]); + } else if (type === "message_delta" && rawField(data, "usage")) { + add(rawField(data, "usage"), paths); } return; } } +/** Which counts a call's provider reported, from the raw usage fields it sent. */ +export function reportedCountsFrom(api: string, seen: ReadonlySet): ReportedCounts { + const fields = REPORTED_FIELDS[api]; + if (fields === undefined) { + return { input: false, cacheRead: false, output: false }; + } + return { + input: fields.input.every((path) => seen.has(path)), + cacheRead: fields.cacheRead.some((path) => seen.has(path)), + output: fields.output.some((path) => seen.has(path)), + }; +} + /** * Record which counts each model call's provider reported, keyed by the call's - * final usage object. Pi hands every raw provider stream event to - * `onProviderStreamEvent` before parsing it. The agent loop and compaction - * summaries both call the agent's stream function, so wrapping it covers both, - * and Pi passes the final usage object unchanged to `message_end` and - * `compaction_end`. The record is stored before Pi emits either event. Pi's - * cache warmer replays a request with its options but sends it through its own - * model runtime, so the record is copied when the call ends and a warm-up's - * usage entry has no record. + * final usage object. Pi hands every raw provider stream event, with the model + * that served it, to `onProviderStreamEvent` before parsing it. The agent loop + * and compaction summaries both call the agent's stream function, so wrapping + * it covers both, and Pi passes the final usage object unchanged to + * `message_end` and `compaction_end`. The record is stored before Pi emits + * either event. Pi's cache warmer replays a request with the same options, so + * a warm-up's raw events reach this observer too, but the warm-up runs outside + * the agent's stream function: its usage entry has no record and stays unknown. + * The record is copied when the call ends so a later warm-up cannot change it. */ -function recordReportedCounts(agent: AgentSession["agent"], reportedCounts: WeakMap): void { +function recordReportedCounts(agent: AgentSession["agent"], records: WeakMap): void { const streamFunction = agent.streamFunction; agent.streamFunction = async (model, context, options) => { - const reported: ReportedCounts = { input: false, output: false }; + const seen = new Set(); + let api: string = model.api; const response = await streamFunction(model, context, { ...options, onProviderStreamEvent: async (data, eventModel) => { - noteReportedCounts(reported, model.api, data); + api = eventModel?.api ?? model.api; + noteReportedFields(seen, api, data); await options?.onProviderStreamEvent?.(data, eventModel); }, }); void response.result().then( (message) => { if (typeof message.usage === "object" && message.usage !== null) { - reportedCounts.set(message.usage, { ...reported }); + records.set(message.usage, reportedCountsFrom(api, seen)); } }, () => undefined, diff --git a/adapters/typescript/pi/src/runtime.ts b/adapters/typescript/pi/src/runtime.ts index 7488d1e87..fea83d60e 100644 --- a/adapters/typescript/pi/src/runtime.ts +++ b/adapters/typescript/pi/src/runtime.ts @@ -20,16 +20,15 @@ import type { PiRelayRuntime } from "./relay.js"; export type PiStopReason = "stop" | "length" | "toolUse" | "error" | "aborted" | string; /** - * Token counts Pi recorded during one prompt. Pi's `input` excludes cache - * reads and writes. `input`, `cacheRead` and `cacheWrite` are absent together - * when a call's provider left its input count unreported, and `output` is - * absent when one left its output count unreported. + * Token counts Pi recorded during one prompt. Each count is absent when a + * call's provider left it unreported. */ export interface PiTokenUsage { + /** Input tokens, including cache reads and cache writes. */ input?: number; - output?: number; + /** Cache reads, part of `input`. */ cacheRead?: number; - cacheWrite?: number; + output?: number; } export interface PiPromptOutcome { @@ -69,21 +68,13 @@ function failed(code: string, message: string): AgentRunResult { * both input and output are known. */ function agentUsage(usage: PiTokenUsage | undefined): AgentUsage | undefined { - if (usage === undefined) { - return undefined; - } - const { output, cacheRead, cacheWrite } = usage; - const input = - usage.input !== undefined && cacheRead !== undefined && cacheWrite !== undefined - ? usage.input + cacheRead + cacheWrite - : undefined; - if (input === undefined && output === undefined) { + if (usage === undefined || (usage.input === undefined && usage.output === undefined)) { return undefined; } + const { input, cacheRead, output } = usage; return { - ...(input === undefined - ? {} - : { input_tokens: input, cached_input_tokens: cacheRead, input_tokens_include_cache: true }), + ...(input === undefined ? {} : { input_tokens: input, input_tokens_include_cache: true }), + ...(input === undefined || cacheRead === undefined ? {} : { cached_input_tokens: cacheRead }), ...(output === undefined ? {} : { output_tokens: output }), ...(input === undefined || output === undefined ? {} : { total_tokens: input + output }), }; diff --git a/adapters/typescript/pi/test/pi-sdk.test.mjs b/adapters/typescript/pi/test/pi-sdk.test.mjs index a356e41b9..35723f75f 100644 --- a/adapters/typescript/pi/test/pi-sdk.test.mjs +++ b/adapters/typescript/pi/test/pi-sdk.test.mjs @@ -11,10 +11,11 @@ import { createServer } from "node:http"; import { addEventUsage, addUsage, - noteReportedCounts, + noteReportedFields, buildCatalogModel, modelAwareCompactionReserveTokens, PiSdkSessionFactory, + reportedCountsFrom, resolveCustomTools, selectPiMcpServers, withCustomBaseUrl, @@ -1446,15 +1447,16 @@ test("accepts a well-formed extensions.cost", async () => { // Lookups for what a model call's provider reported: every count, or a call // whose raw stream events the adapter never saw. -const allReported = () => ({ input: true, output: true }); +const allReported = () => ({ input: true, cacheRead: true, output: true }); const unseen = () => undefined; +// A prompt's running total before any call, with input counting the cache. +const NO_TOKENS = { input: 0, cacheRead: 0, output: 0 }; test("sums each assistant message's usage and drops an incomplete one", () => { - const zero = { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 }; - const first = addUsage(zero, { input: 10, output: 2, cacheRead: 5, cacheWrite: 1, totalTokens: 18 }); + const first = addUsage(NO_TOKENS, { input: 10, output: 2, cacheRead: 5, cacheWrite: 1, totalTokens: 18 }); const second = addUsage(first, { input: 3, output: 4, cacheRead: 0, cacheWrite: 0, reasoning: 2 }); - assert.deepEqual(second, { input: 13, output: 6, cacheRead: 5, cacheWrite: 1 }); + assert.deepEqual(second, { input: 19, cacheRead: 5, output: 6 }); assert.equal(addUsage(second, { input: 1, output: 1 }), undefined); assert.equal(addUsage(undefined, { input: 1, output: 1, cacheRead: 0, cacheWrite: 0 }), undefined); }); @@ -1471,31 +1473,25 @@ test("counts the tokens Pi bills outside assistant messages", () => { { type: "turn_start" }, ]; - const total = events.reduce( - (sum, event) => addEventUsage(sum, event, allReported), - { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 }, - ); + const total = events.reduce((sum, event) => addEventUsage(sum, event, allReported), NO_TOKENS); - assert.deepEqual(total, { input: 135, output: 4, cacheRead: 0, cacheWrite: 0 }); + assert.deepEqual(total, { input: 135, cacheRead: 0, output: 4 }); }); test("a compaction without a usage report leaves the prompt's tokens unknown", () => { - const zero = { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 }; - assert.equal( - addEventUsage(zero, { type: "compaction_end", aborted: false, result: { summary: "s" } }, allReported), + addEventUsage(NO_TOKENS, { type: "compaction_end", aborted: false, result: { summary: "s" } }, allReported), undefined, ); - assert.equal(addEventUsage(zero, { type: "compaction_end", aborted: true }, allReported), undefined); + assert.equal(addEventUsage(NO_TOKENS, { type: "compaction_end", aborted: true }, allReported), undefined); }); test("an aborted assistant message leaves the prompt's tokens unknown", () => { - const zero = { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 }; const usage = { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 }; assert.equal( addEventUsage( - zero, + NO_TOKENS, { type: "message_end", message: { role: "assistant", stopReason: "aborted", usage } }, allReported, ), @@ -1504,76 +1500,108 @@ test("an aborted assistant message leaves the prompt's tokens unknown", () => { }); test("a model call's counts are known only as far as its provider reported them", () => { - const zero = { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 }; const zeroFilled = { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, totalTokens: 0 }; const assistant = { type: "message_end", message: { role: "assistant", usage: zeroFilled } }; - const inputOnly = () => ({ input: true, output: false }); + const reported = (counts) => () => ({ input: false, cacheRead: false, output: false, ...counts }); // Pi fills an omitted count with 0, so a call the adapter did not see is unknown. - assert.equal(addEventUsage(zero, assistant, unseen), undefined); - assert.equal(addEventUsage(zero, { type: "compaction_end", aborted: false, result: { usage: zeroFilled } }, unseen), undefined); - assert.equal(addEventUsage(zero, { type: "entry_appended", entry: { type: "usage", usage: zeroFilled } }, unseen), undefined); + assert.equal(addEventUsage(NO_TOKENS, assistant, unseen), undefined); + assert.equal( + addEventUsage(NO_TOKENS, { type: "compaction_end", aborted: false, result: { usage: zeroFilled } }, unseen), + undefined, + ); + assert.equal( + addEventUsage(NO_TOKENS, { type: "entry_appended", entry: { type: "usage", usage: zeroFilled } }, unseen), + undefined, + ); // A reported zero is a measured zero; an omitted count stays unknown. - assert.deepEqual(addEventUsage(zero, assistant, allReported), zero); - assert.deepEqual(addEventUsage(zero, assistant, inputOnly), { input: 0, cacheRead: 0, cacheWrite: 0 }); - assert.equal(addEventUsage(zero, assistant, () => ({ input: false, output: false })), undefined); + assert.deepEqual(addEventUsage(NO_TOKENS, assistant, allReported), NO_TOKENS); + assert.deepEqual(addEventUsage(NO_TOKENS, assistant, reported({ input: true, cacheRead: true })), { + input: 0, + cacheRead: 0, + }); + assert.deepEqual(addEventUsage(NO_TOKENS, assistant, reported({ input: true, output: true })), { + input: 0, + output: 0, + }); + assert.equal(addEventUsage(NO_TOKENS, assistant, reported({ cacheRead: true })), undefined); // A tool reports its own usage, so a zero there is a measured zero. assert.deepEqual( - addEventUsage(zero, { type: "message_end", message: { role: "toolResult", usage: zeroFilled } }, unseen), - zero, + addEventUsage(NO_TOKENS, { type: "message_end", message: { role: "toolResult", usage: zeroFilled } }, unseen), + NO_TOKENS, ); }); test("reads which counts each wire protocol's raw usage reports", () => { const counts = (api, events) => { - const reported = { input: false, output: false }; + const seen = new Set(); for (const event of events) { - noteReportedCounts(reported, api, event); + noteReportedFields(seen, api, event); } - return reported; + return reportedCountsFrom(api, seen); }; + const anthropicStart = (usage) => ({ type: "message_start", message: { usage } }); + const anthropicUsage = { input_tokens: 10, cache_read_input_tokens: 4, cache_creation_input_tokens: 0 }; - // Anthropic keeps message_start's counts and takes the counts a delta carries. + // Anthropic input excludes the cache, so it needs both cache fields, and only + // a message_delta reports the output count. assert.deepEqual( counts("anthropic-messages", [ - { type: "message_start", message: { usage: { input_tokens: 10, output_tokens: 1 } } }, + anthropicStart({ ...anthropicUsage, output_tokens: 1 }), { type: "message_delta", usage: { output_tokens: 5 } }, ]), - { input: true, output: true }, + { input: true, cacheRead: true, output: true }, ); assert.deepEqual( counts("anthropic-messages", [ - { type: "message_start", message: { usage: { output_tokens: 1 } } }, + anthropicStart({ input_tokens: 10, output_tokens: 1 }), { type: "message_delta", usage: { output_tokens: 5 } }, ]), - { input: false, output: true }, + { input: false, cacheRead: false, output: true }, + ); + assert.deepEqual( + counts("anthropic-messages", [anthropicStart({ ...anthropicUsage, output_tokens: 1 }), { type: "message_delta" }]), + { input: true, cacheRead: true, output: false }, + ); + // The other protocols replace a call's usage with each report, and their + // input includes the cache, so a missing cache detail leaves only it unknown. + assert.deepEqual( + counts("openai-completions", [ + { choices: [], usage: { prompt_tokens: 7, prompt_tokens_details: { cached_tokens: 2 } } }, + { choices: [], usage: { completion_tokens: 2 } }, + ]), + { input: false, cacheRead: false, output: true }, ); - // The other protocols replace a call's usage with each report. assert.deepEqual( - counts("openai-completions", [{ choices: [], usage: { prompt_tokens: 7 } }, { choices: [], usage: { completion_tokens: 2 } }]), - { input: false, output: true }, + counts("openai-completions", [{ choices: [], usage: { prompt_tokens: 7, prompt_tokens_details: { cached_tokens: 2 } } }]), + { input: true, cacheRead: true, output: false }, ); assert.deepEqual( counts("openai-responses", [{ type: "response.completed", response: { usage: { input_tokens: 3 } } }]), - { input: true, output: false }, + { input: true, cacheRead: false, output: false }, ); assert.deepEqual( - counts("google-generative-ai", [{ usageMetadata: { promptTokenCount: 4, candidatesTokenCount: 0 } }]), - { input: true, output: true }, + counts("google-generative-ai", [{ usageMetadata: { promptTokenCount: 4, thoughtsTokenCount: 3 } }]), + { input: true, cacheRead: false, output: true }, ); assert.deepEqual( counts("bedrock-converse-stream", [{ metadata: { usage: { outputTokens: 6 } } }]), - { input: false, output: true }, + { input: false, cacheRead: false, output: true }, ); + // A protocol the adapter does not know reports nothing it can trust. + assert.deepEqual(counts("custom-api", [{ usage: { prompt_tokens: 1, completion_tokens: 1 } }]), { + input: false, + cacheRead: false, + output: false, + }); }); test("a split-turn compaction leaves the prompt's tokens unknown", () => { - const zero = { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 }; const prepared = (preparation) => ({ type: "session_before_compact", preparation }); - assert.deepEqual(addEventUsage(zero, prepared({ isSplitTurn: false }), allReported), zero); - assert.equal(addEventUsage(zero, prepared({ isSplitTurn: true }), allReported), undefined); - assert.equal(addEventUsage(zero, prepared(undefined), allReported), undefined); + assert.deepEqual(addEventUsage(NO_TOKENS, prepared({ isSplitTurn: false }), allReported), NO_TOKENS); + assert.equal(addEventUsage(NO_TOKENS, prepared({ isSplitTurn: true }), allReported), undefined); + assert.equal(addEventUsage(NO_TOKENS, prepared(undefined), allReported), undefined); }); test("a retried compaction summary leaves the prompt's tokens unknown", () => { @@ -1586,10 +1614,7 @@ test("a retried compaction summary leaves the prompt's tokens unknown", () => { ]; assert.equal( - events.reduce( - (sum, event) => addEventUsage(sum, event, allReported), - { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 }, - ), + events.reduce((sum, event) => addEventUsage(sum, event, allReported), NO_TOKENS), undefined, ); }); @@ -1666,9 +1691,10 @@ test( ["read the note", "again"], ); - // A reported zero output count stays a measured zero. - assert.deepEqual(first.usage, { input: 210, output: 10, cacheRead: 40, cacheWrite: 0 }); - assert.deepEqual(second.usage, { input: 300, output: 25, cacheRead: 0, cacheWrite: 0 }); + // A reported zero output count stays a measured zero. The second call sent + // no cache details, so the prompt's cache reads are unknown. + assert.deepEqual(first.usage, { input: 250, output: 10 }); + assert.deepEqual(second.usage, { input: 300, output: 25 }); }, ); @@ -1690,7 +1716,7 @@ test( assert.equal(omitted.usage, undefined); assert.equal(mixed.text, "partial"); assert.equal(mixed.usage, undefined); - assert.deepEqual(reported.usage, { input: 50, output: 5, cacheRead: 0, cacheWrite: 0 }); + assert.deepEqual(reported.usage, { input: 50, output: 5 }); }, ); @@ -1702,7 +1728,7 @@ test( [ { content: "input only", usage: { prompt_tokens: 100 } }, { content: "output only", usage: { completion_tokens: 5 } }, - { content: "zeros", usage: { prompt_tokens: 0, completion_tokens: 0 } }, + { content: "zeros", usage: { prompt_tokens: 0, completion_tokens: 0, prompt_tokens_details: { cached_tokens: 0 } } }, { readNote: true, usage: { prompt_tokens: 60 } }, { content: "both", usage: { prompt_tokens: 70, completion_tokens: 7 } }, ], @@ -1710,11 +1736,11 @@ test( ); // Pi fills an omitted count with 0; the provider's raw report decides. - assert.deepEqual(inputOnly.usage, { input: 100, cacheRead: 0, cacheWrite: 0 }); + assert.deepEqual(inputOnly.usage, { input: 100 }); assert.deepEqual(outputOnly.usage, { output: 5 }); - assert.deepEqual(zeros.usage, { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 }); + assert.deepEqual(zeros.usage, { input: 0, cacheRead: 0, output: 0 }); // One call without an output count leaves the prompt's output unknown. - assert.deepEqual(afterPartial.usage, { input: 130, cacheRead: 0, cacheWrite: 0 }); + assert.deepEqual(afterPartial.usage, { input: 130 }); }, ); @@ -1735,7 +1761,7 @@ test( { context_window: 100_000, max_tokens: 1_000 }, ); - assert.deepEqual(compacted.usage, { input: 90_220, output: 11, cacheRead: 0, cacheWrite: 0 }); + assert.deepEqual(compacted.usage, { input: 90_220, output: 11 }); }, ); diff --git a/adapters/typescript/pi/test/runtime.test.mjs b/adapters/typescript/pi/test/runtime.test.mjs index 0573e9c25..caa39a100 100644 --- a/adapters/typescript/pi/test/runtime.test.mjs +++ b/adapters/typescript/pi/test/runtime.test.mjs @@ -559,7 +559,7 @@ test("retries Relay cleanup when gateway shutdown fails", async () => { }); test("reports the invocation's tokens with cache reads and writes inside the input", async () => { - const usage = { input: 100, output: 40, cacheRead: 300, cacheWrite: 20 }; + const usage = { input: 420, cacheRead: 300, output: 40 }; const runtime = new PiAdapterRuntime({ async create() { return { @@ -585,10 +585,7 @@ test("reports the invocation's tokens with cache reads and writes inside the inp }); test("reports only the token counts the prompt's provider reported", async () => { - const outcomes = [ - { input: 100, cacheRead: 30, cacheWrite: 0 }, - { output: 5 }, - ]; + const outcomes = [{ input: 130, cacheRead: 30 }, { output: 5 }, { input: 70, output: 7 }]; const runtime = new PiAdapterRuntime({ async create() { let turnCount = 0; @@ -606,6 +603,7 @@ test("reports only the token counts the prompt's provider reported", async () => await runtime.start(startInput()); const inputOnly = await runtime.invoke({ input: "one" }, context); const outputOnly = await runtime.invoke({ input: "two" }, context); + const noCacheDetail = await runtime.invoke({ input: "three" }, context); await runtime.stop(); assert.deepEqual(inputOnly.usage, { @@ -614,4 +612,10 @@ test("reports only the token counts the prompt's provider reported", async () => input_tokens_include_cache: true, }); assert.deepEqual(outputOnly.usage, { output_tokens: 5 }); + assert.deepEqual(noCacheDetail.usage, { + input_tokens: 70, + input_tokens_include_cache: true, + output_tokens: 7, + total_tokens: 77, + }); }); diff --git a/tests/_utils/mock_api_server.py b/tests/_utils/mock_api_server.py index 020f2a341..d3c9ae2ef 100644 --- a/tests/_utils/mock_api_server.py +++ b/tests/_utils/mock_api_server.py @@ -11,6 +11,9 @@ from fastapi.responses import JSONResponse, Response, StreamingResponse import uvicorn +# Chat-completion usage the mock reports unless a scenario sets another one. +_ZERO_USAGE = {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0} + @contextmanager def mock_api_server(port: int) -> Iterator[str]: @@ -18,7 +21,8 @@ def mock_api_server(port: int) -> Iterator[str]: Context manager for a mock API server. Use the /_requests endpoint to inspect captured chat-completion payloads after a test action. - Use the /_scenario endpoint to configure a status code or ordered tool calls for subsequent requests. + Use the /_scenario endpoint to configure a status code, ordered tool calls, or the + chat-completion usage object for subsequent requests. Use the /_reset endpoint to clear all retained request and scenario state. Args: @@ -36,6 +40,7 @@ def reset_state() -> None: app.state.tool_calls = [] app.state.tool_call_index = 0 app.state.mcp_authorization_headers = [] + app.state.usage = dict(_ZERO_USAGE) reset_state() @@ -65,7 +70,7 @@ def requests() -> list[dict[str, object]]: @app.post("/_scenario") async def scenario(request: Request) -> dict[str, object]: - """Configure the status code or ordered tool calls for subsequent requests.""" + """Configure the status code, tool calls or usage for subsequent requests.""" payload = await request.json() app.state.status_code = int(payload.get("status_code", 200)) @@ -74,6 +79,7 @@ async def scenario(request: Request) -> dict[str, object]: [tool_call] if tool_call is not None else [] ) app.state.tool_call_index = 0 + app.state.usage = payload.get("usage", dict(_ZERO_USAGE)) return { "status_code": app.state.status_code, "tool_call": tool_call, @@ -179,7 +185,7 @@ async def chat_completions(request: Request): content = f"echo user_count={len(user_messages)} latest={latest}" if payload.get("stream") is True: return StreamingResponse( - _stream_chat_completion(payload, content), + _stream_chat_completion(payload, content, app.state.usage), media_type="text/event-stream", ) @@ -196,11 +202,7 @@ async def chat_completions(request: Request): "finish_reason": "stop", } ], - "usage": { - "prompt_tokens": 0, - "completion_tokens": 0, - "total_tokens": 0, - }, + "usage": app.state.usage, } ) @@ -278,7 +280,9 @@ async def messages(request: Request): thread.join(timeout=5) -def _stream_chat_completion(payload: dict[str, object], content: str) -> Iterator[str]: +def _stream_chat_completion( + payload: dict[str, object], content: str, usage: dict[str, object] +) -> Iterator[str]: model = payload.get("model", "fabric-echo") chunks = [ { @@ -325,19 +329,15 @@ def _stream_chat_completion(payload: dict[str, object], content: str) -> Iterato for chunk in chunks: yield f"data: {json.dumps(chunk)}\n\n" if payload.get("stream_options") == {"include_usage": True}: - usage = { + usage_chunk = { "id": "chatcmpl-fabric-test", "object": "chat.completion.chunk", "created": 0, "model": model, "choices": [], - "usage": { - "prompt_tokens": 0, - "completion_tokens": 0, - "total_tokens": 0, - }, + "usage": usage, } - yield f"data: {json.dumps(usage)}\n\n" + yield f"data: {json.dumps(usage_chunk)}\n\n" yield "data: [DONE]\n\n" diff --git a/tests/adapters/test_hermes_adapter.py b/tests/adapters/test_hermes_adapter.py index 5d5b824b4..92edade5b 100644 --- a/tests/adapters/test_hermes_adapter.py +++ b/tests/adapters/test_hermes_adapter.py @@ -1628,7 +1628,11 @@ class _UsageRecordingAgent: through Hermes' own usage accounting and context compressor, called the way Hermes' turn loop calls it.""" - def __init__(self, turns: list[list[dict[str, int] | None]]) -> None: + def __init__( + self, + turns: list[list[dict[str, object] | None]], + api_mode: str = "chat_completions", + ) -> None: from agent.context_compressor import ContextCompressor self._turns = turns @@ -1636,7 +1640,7 @@ def __init__(self, turns: list[list[dict[str, int] | None]]) -> None: model="test-model", quiet_mode=True ) self.provider = "custom" - self.api_mode = "chat_completions" + self.api_mode = api_mode self.model = "test-model" self.base_url = "http://127.0.0.1:1/v1" self.api_key = "" @@ -1659,6 +1663,9 @@ def __init__(self, turns: list[list[dict[str, int] | None]]) -> None: ): setattr(self, f"session_{counter}", 0) + def close(self) -> None: + pass + def run_conversation(self, user_message: str) -> dict[str, object]: from agent import turn_response_check @@ -1778,6 +1785,8 @@ async def test_hermes_usage_with_partial_provider_report_leaves_omitted_count_un # Hermes normalizes an omitted count to 0; the provider's report decides. assert input_only.usage is not None assert input_only.usage.input_tokens == 100 + # The provider sent no cache details, so the cached share is unknown. + assert input_only.usage.cached_input_tokens is None assert input_only.usage.output_tokens is None assert input_only.usage.total_tokens is None assert output_only.usage is not None @@ -1799,6 +1808,87 @@ async def test_hermes_usage_with_partial_provider_report_leaves_omitted_count_un assert after_partial.usage.total_tokens == 440 +async def test_hermes_usage_reports_cache_counts_only_when_the_provider_sent_them(): + chat = _UsageRecordingAgent( + [ + [ + { + "prompt_tokens": 50, + "completion_tokens": 5, + "prompt_tokens_details": {"cached_tokens": 20}, + } + ], + ] + ) + # Anthropic's input_tokens excludes the cache, so input needs both cache fields. + anthropic = _UsageRecordingAgent( + [ + [{"input_tokens": 100, "output_tokens": 5}], + [ + { + "input_tokens": 100, + "cache_read_input_tokens": 30, + "cache_creation_input_tokens": 0, + "output_tokens": 5, + } + ], + ], + api_mode="anthropic_messages", + ) + chat_runtime = _usage_recording_runtime(chat) + anthropic_runtime = _usage_recording_runtime(anthropic) + invocation = _usage_invocation() + + cached = await chat_runtime.invoke(*invocation) + without_cache_fields = await anthropic_runtime.invoke(*invocation) + with_cache_fields = await anthropic_runtime.invoke(*invocation) + + assert cached.usage is not None + assert (cached.usage.input_tokens, cached.usage.cached_input_tokens) == (50, 20) + assert without_cache_fields.usage is not None + assert without_cache_fields.usage.input_tokens is None + assert without_cache_fields.usage.cached_input_tokens is None + assert without_cache_fields.usage.output_tokens == 5 + assert without_cache_fields.usage.total_tokens is None + assert with_cache_fields.usage is not None + assert with_cache_fields.usage.input_tokens == 130 + assert with_cache_fields.usage.cached_input_tokens == 30 + assert with_cache_fields.usage.total_tokens == 135 + + +async def test_hermes_stop_unregisters_the_agent_usage_observer(): + agent = _UsageRecordingAgent([]) + runtime = _usage_recording_runtime(agent) + assert agent in adapter._USAGE_OBSERVERS + + await runtime.stop() + + assert agent not in adapter._USAGE_OBSERVERS + + +@pytest.mark.parametrize("drift", ["replaced", "not_called"]) +def test_hermes_usage_observer_rejects_unknown_usage_call_site(monkeypatch, drift): + from agent import turn_response_check + from agent import turn_usage + + monkeypatch.setattr( + turn_response_check, "record_response_usage", turn_usage.record_response_usage + ) + if drift == "replaced": + monkeypatch.setattr( + turn_response_check, "record_response_usage", lambda *args, **kwargs: None + ) + else: + monkeypatch.setattr( + turn_response_check, "check_api_response", lambda *args, **kwargs: None + ) + + with pytest.raises(adapter.lifecycle.LifecycleError) as raised: + adapter._install_usage_observer() + + assert raised.value.code == "hermes_usage_accounting_unsupported" + + @pytest.mark.parametrize( ("request_input", "expected_message"), [ diff --git a/tests/e2e/test_hermes_e2e.py b/tests/e2e/test_hermes_e2e.py index dcb063615..659a2de0e 100644 --- a/tests/e2e/test_hermes_e2e.py +++ b/tests/e2e/test_hermes_e2e.py @@ -126,6 +126,14 @@ async def test_hermes_persistent_host_reports_usage_for_each_invocation( ) as runtime: first = await runtime.invoke(input="first") second = await runtime.invoke(input="second") + # The provider now reports a prompt count but no completion count. + scenario_response = requests.post( + f"{api_server}/_scenario", + json={"usage": {"prompt_tokens": 12}}, + timeout=5, + ) + scenario_response.raise_for_status() + partial = await runtime.invoke(input="third") # The mock endpoint bills 0 tokens per call, so each invocation's change # since the previous one is a measured 0, not unknown. @@ -134,6 +142,14 @@ async def test_hermes_persistent_host_reports_usage_for_each_invocation( assert turn.usage is not None, turn.to_mapping() assert turn.usage.input_tokens == 0, turn.to_mapping() assert turn.usage.output_tokens == 0, turn.to_mapping() + # Through Hermes' real turn loop, the omitted completion count stays + # unknown rather than a measured 0, and so does the total. + assert partial["status"] == "succeeded", partial.to_mapping() + assert partial.to_mapping()["usage"] == { + "input_tokens": 12, + "input_tokens_include_cache": True, + "metadata": {}, + }, partial.to_mapping() @pytest.mark.usefixtures("mock_nvidia_api_key") From 4491ef93747139fe83bca5b172bda9ab68787288 Mon Sep 17 00:00:00 2001 From: mnajafian-nv Date: Mon, 19 Oct 2026 14:55:07 -0700 Subject: [PATCH 10/68] fix: leave Hermes usage unknown after a truncated or refused reply Hermes' check_api_response returns on a reply that ends with finish_reason length or content_filter before it calls record_response_usage, so the provider bills that call but Hermes' session counters never include it. A turn that hit max_tokens and continued therefore reported measured input, output and total tokens that left out the truncated call. The adapter's process-wide wrapper now also covers the two handlers check_api_response calls in those cases, recover_from_truncation and handle_content_policy_refusal, with the same once-per-process install and start-time check. When either handler returns before usage is recorded, the response counts like a response without usage, and the invocation reports no usage. A truncation that falls through to record_response_usage is still counted. The mock endpoint takes a finish reason per scenario, and the Hermes e2e test sends a length reply and a content_filter reply through Hermes' real turn loop. A terminal content_filter result already had no token counters, so that case reported no usage before this change; the fix matters when a configured fallback answers after the refusal. Signed-off-by: mnajafian-nv --- adapters/python/hermes/README.md | 20 +-- .../nemo_fabric_adapters/hermes/adapter.py | 117 +++++++++++------- tests/_utils/mock_api_server.py | 17 ++- tests/e2e/test_hermes_e2e.py | 43 +++++++ 4 files changed, 136 insertions(+), 61 deletions(-) diff --git a/adapters/python/hermes/README.md b/adapters/python/hermes/README.md index fa1a57ea2..0b05bc8f7 100644 --- a/adapters/python/hermes/README.md +++ b/adapters/python/hermes/README.md @@ -127,15 +127,17 @@ rather than a carried-over or whole-session total, and the one after it reports its own change again. Hermes adds no tokens for a provider response that carried no usage, so an -invocation with such a response has no `usage`. Hermes also reads a count that -a response's usage left out as 0, so the adapter checks each response's usage -before Hermes normalizes it. A reported zero still counts as 0. When a response -left out its input count, the invocation's `usage` leaves out `input_tokens`, -`cached_input_tokens`, and `total_tokens`. Anthropic Messages usage reports -input without the cache, so its input count also needs both cache fields. When -a response left out its output count, the `usage` leaves out `output_tokens` -and `total_tokens`. When a response sent no cache details, the `usage` leaves -out `cached_input_tokens`. +invocation with such a response has no `usage`. Hermes also returns before +recording the usage of a reply that hit the token limit or a content filter, so +an invocation with such a reply has no `usage` either. Hermes reads a count +that a response's usage left out as 0, so the adapter checks each response's +usage before Hermes normalizes it. A reported zero still counts as 0. When a +response left out its input count, the invocation's `usage` leaves out +`input_tokens`, `cached_input_tokens`, and `total_tokens`. Anthropic Messages +usage reports input without the cache, so its input count also needs both cache +fields. When a response left out its output count, the `usage` leaves out +`output_tokens` and `total_tokens`. When a response sent no cache details, the +`usage` leaves out `cached_input_tokens`. This check covers the usage that the OpenAI-compatible, Responses, and Anthropic clients return. Hermes' native Gemini and Bedrock transports, and its diff --git a/adapters/python/hermes/src/nemo_fabric_adapters/hermes/adapter.py b/adapters/python/hermes/src/nemo_fabric_adapters/hermes/adapter.py index 57b24e8f3..2ddc08deb 100755 --- a/adapters/python/hermes/src/nemo_fabric_adapters/hermes/adapter.py +++ b/adapters/python/hermes/src/nemo_fabric_adapters/hermes/adapter.py @@ -72,7 +72,7 @@ def __init__(self) -> None: # Hermes reports session totals; each invocation reports the change # from this baseline. A counter missing from it has no baseline. self._session_token_totals: dict[str, int] = _fresh_session_totals() - self._usage_less_responses = 0 + self._uncounted_responses = 0 # Provider responses whose usage left out a count, by count. self._unreported_counts = dict.fromkeys(_REPORTED_COUNTS, 0) @@ -299,7 +299,7 @@ def run_hermes_turn() -> tuple[dict[str, Any], str]: # failed, cancelled or timed-out invocation never hands its tokens to # the next one. previous_totals, self._session_token_totals = self._session_token_totals, {} - usage_less_before = self._usage_less_responses + uncounted_before = self._uncounted_responses unreported_before = dict(self._unreported_counts) invoke_task = asyncio.create_task(asyncio.to_thread(run_hermes_turn)) self._active_invoke_task = invoke_task @@ -394,9 +394,10 @@ def clear_active_invoke_task( if responses != unreported_before[count] ), ) - if self._usage_less_responses != usage_less_before: - # A provider response without usage added no tokens, so the - # change in the counters is an undercount, not this invocation's. + if self._uncounted_responses != uncounted_before: + # A provider response without usage, or one Hermes returned on + # before recording its usage, added no tokens, so the change in + # the counters is an undercount, not this invocation's. usage = None return AgentRunResult( status=AgentRunStatus.FAILED if failed else AgentRunStatus.SUCCEEDED, @@ -425,7 +426,7 @@ def _count_usage_less_responses(self, agent: Any) -> None: ) def count_and_note() -> None: - self._usage_less_responses += 1 + self._uncounted_responses += 1 if callable(note_usage_less_response): note_usage_less_response() @@ -436,13 +437,20 @@ def _count_unreported_counts(self, agent: Any) -> None: Hermes reads a missing count as 0, so its counters cannot tell an omitted count from a reported zero, and Hermes passes the raw usage to - no per-agent hook. The adapter therefore registers the agent with a - process-wide wrapper around the function Hermes' turn loop calls for - each response; see ``_install_usage_observer``. + no per-agent hook. The adapter therefore registers the agent with + process-wide wrappers around the functions Hermes' turn loop calls for + each response; see ``_install_usage_observer``. A refused or truncated + response that Hermes never records counts as a response without usage. """ _install_usage_observer() - def count_unreported(response: Any) -> None: + def count_unreported(handler: str, response: Any, outcome: Any) -> None: + if handler != "record_response_usage": + # Hermes returns on a refused or truncated response before + # recording its usage, unless the truncation falls through. + if getattr(outcome, "action", None) != "fallthrough": + self._uncounted_responses += 1 + return usage = getattr(response, "usage", None) # A response without usage is counted by the compressor hook. if usage: @@ -673,54 +681,69 @@ def _invoke_hermes_turn( ) -# Response observers by Hermes agent, called before Hermes records a -# response's usage. Delegated subagents are never registered: their tokens +# Response observers by Hermes agent, called with each response Hermes' +# turn loop handles. Delegated subagents are never registered: their tokens # never reach the registered agent's counters. -_USAGE_OBSERVERS: weakref.WeakKeyDictionary[Any, Callable[[Any], None]] = ( +_USAGE_OBSERVERS: weakref.WeakKeyDictionary[Any, Callable[[str, Any, Any], None]] = ( weakref.WeakKeyDictionary() ) +# The functions ``agent.turn_response_check.check_api_response`` calls with +# (agent, response, ...) for a provider response, by the module defining +# each. record_response_usage folds the usage; on a content-filter refusal or +# a reply cut off at the token limit, Hermes calls one of the other two and +# returns before recording the usage. +_RESPONSE_HANDLERS = { + "record_response_usage": "agent.turn_usage", + "handle_content_policy_refusal": "agent.turn_truncation", + "recover_from_truncation": "agent.turn_truncation", +} + +def _observed_response_handler(name: str, handler: Callable[..., Any]) -> Any: + """Wrap a Hermes response handler to show each call to the agent's observer.""" -def _observed_record_response_usage( - agent: Any, response: Any, *args: Any, **kwargs: Any -) -> Any: - """Show ``response`` to the agent's observer, then record it as Hermes does.""" - from agent import turn_usage + def observed(agent: Any, response: Any, *args: Any, **kwargs: Any) -> Any: + outcome = handler(agent, response, *args, **kwargs) + observer = _USAGE_OBSERVERS.get(agent) + if observer is not None: + observer(name, response, outcome) + return outcome - observer = _USAGE_OBSERVERS.get(agent) - if observer is not None: - observer(response) - return turn_usage.record_response_usage(agent, response, *args, **kwargs) + observed.nemo_fabric_observed = True # type: ignore[attr-defined] + return observed def _install_usage_observer() -> None: - """Route Hermes' per-response usage accounting through the observers. - - ``agent.turn_response_check.check_api_response`` calls - ``record_response_usage`` for each provider response, looked up as a - global of its module, and that function is the first place the raw usage - reaches before Hermes normalizes it. The wrapper replaces that module - global once per process and always delegates to Hermes' function. If - Hermes no longer calls usage accounting through that name, start fails - rather than reporting omitted counts as measured zeros. + """Route Hermes' per-response handlers through the observers. + + ``check_api_response`` looks up each handler as a global of its module, + and ``record_response_usage`` is the first place the raw usage reaches + before Hermes normalizes it. The wrappers replace those module globals + once per process and always delegate to Hermes' functions. If Hermes no + longer calls a handler through that name, start fails rather than + reporting an omitted or uncounted response as measured tokens. """ + import importlib + from agent import turn_response_check - from agent import turn_usage - - current = turn_response_check.record_response_usage - if current is _observed_record_response_usage: - return - if ( - current is not turn_usage.record_response_usage - or "record_response_usage" - not in turn_response_check.check_api_response.__code__.co_names - ): - raise lifecycle.LifecycleError( - "hermes_usage_accounting_unsupported", - "This Hermes version records response usage outside " - "agent.turn_response_check.record_response_usage", - ) - turn_response_check.record_response_usage = _observed_record_response_usage + + called = turn_response_check.check_api_response.__code__.co_names + handlers = {} + for name, module in _RESPONSE_HANDLERS.items(): + current = getattr(turn_response_check, name, None) + if getattr(current, "nemo_fabric_observed", False): + continue + if name not in called or current is not getattr( + importlib.import_module(module), name, None + ): + raise lifecycle.LifecycleError( + "hermes_usage_accounting_unsupported", + f"This Hermes version handles responses outside " + f"agent.turn_response_check.{name}", + ) + handlers[name] = current + for name, handler in handlers.items(): + setattr(turn_response_check, name, _observed_response_handler(name, handler)) _REPORTED_COUNTS = ("input", "cached", "output") diff --git a/tests/_utils/mock_api_server.py b/tests/_utils/mock_api_server.py index d3c9ae2ef..7db01e615 100644 --- a/tests/_utils/mock_api_server.py +++ b/tests/_utils/mock_api_server.py @@ -22,7 +22,7 @@ def mock_api_server(port: int) -> Iterator[str]: Use the /_requests endpoint to inspect captured chat-completion payloads after a test action. Use the /_scenario endpoint to configure a status code, ordered tool calls, or the - chat-completion usage object for subsequent requests. + chat-completion usage object and finish reason for subsequent requests. Use the /_reset endpoint to clear all retained request and scenario state. Args: @@ -41,6 +41,7 @@ def reset_state() -> None: app.state.tool_call_index = 0 app.state.mcp_authorization_headers = [] app.state.usage = dict(_ZERO_USAGE) + app.state.finish_reason = "stop" reset_state() @@ -80,6 +81,7 @@ async def scenario(request: Request) -> dict[str, object]: ) app.state.tool_call_index = 0 app.state.usage = payload.get("usage", dict(_ZERO_USAGE)) + app.state.finish_reason = payload.get("finish_reason", "stop") return { "status_code": app.state.status_code, "tool_call": tool_call, @@ -185,7 +187,9 @@ async def chat_completions(request: Request): content = f"echo user_count={len(user_messages)} latest={latest}" if payload.get("stream") is True: return StreamingResponse( - _stream_chat_completion(payload, content, app.state.usage), + _stream_chat_completion( + payload, content, app.state.usage, app.state.finish_reason + ), media_type="text/event-stream", ) @@ -199,7 +203,7 @@ async def chat_completions(request: Request): { "index": 0, "message": {"role": "assistant", "content": content}, - "finish_reason": "stop", + "finish_reason": app.state.finish_reason, } ], "usage": app.state.usage, @@ -281,7 +285,10 @@ async def messages(request: Request): def _stream_chat_completion( - payload: dict[str, object], content: str, usage: dict[str, object] + payload: dict[str, object], + content: str, + usage: dict[str, object], + finish_reason: str, ) -> Iterator[str]: model = payload.get("model", "fabric-echo") chunks = [ @@ -320,7 +327,7 @@ def _stream_chat_completion( { "index": 0, "delta": {}, - "finish_reason": "stop", + "finish_reason": finish_reason, } ], }, diff --git a/tests/e2e/test_hermes_e2e.py b/tests/e2e/test_hermes_e2e.py index 659a2de0e..b902b6be0 100644 --- a/tests/e2e/test_hermes_e2e.py +++ b/tests/e2e/test_hermes_e2e.py @@ -152,6 +152,49 @@ async def test_hermes_persistent_host_reports_usage_for_each_invocation( }, partial.to_mapping() +@pytest.mark.usefixtures("mock_nvidia_api_key") +@pytest.mark.parametrize( + ("finish_reason", "usage"), + [ + # A reply cut off at max_tokens bills the whole budget. + ( + "length", + {"prompt_tokens": 40, "completion_tokens": 512, "total_tokens": 552}, + ), + ( + "content_filter", + {"prompt_tokens": 40, "completion_tokens": 0, "total_tokens": 40}, + ), + ], +) +async def test_hermes_usage_is_unknown_when_hermes_skips_a_response( + code_review_agent_dir: Path, + api_server: str, + finish_reason: str, + usage: dict[str, int], +): + os.environ["ADAPTER_PYTHON"] = sys.executable + config = hermes_config() + config.models["default"].base_url = f"{api_server}/v1" + + async with await Fabric().start_runtime( + config, base_dir=code_review_agent_dir + ) as runtime: + counted = await runtime.invoke(input="first") + scenario_response = requests.post( + f"{api_server}/_scenario", + json={"finish_reason": finish_reason, "usage": usage}, + timeout=5, + ) + scenario_response.raise_for_status() + skipped = await runtime.invoke(input="second") + + assert counted.usage is not None, counted.to_mapping() + # Hermes returns on a truncated or refused reply before recording its + # usage, so the invocation's counters miss a billed call. + assert skipped.usage is None, skipped.to_mapping() + + @pytest.mark.usefixtures("mock_nvidia_api_key") async def test_env_secrets_in_headers( code_review_agent_dir: Path, From de7fe86991a22584452faf6e963167aaa0e85e26 Mon Sep 17 00:00:00 2001 From: mnajafian-nv Date: Mon, 19 Oct 2026 15:03:53 -0700 Subject: [PATCH 11/68] docs: explain the Pi Bedrock input rule Bedrock Converse reports input without the cache, like Anthropic Messages, but its API documents cacheReadInputTokens and cacheWriteInputTokens as optional. The adapter therefore takes inputTokens alone as the input count, and the README now says why. Signed-off-by: mnajafian-nv --- adapters/typescript/pi/README.md | 15 +++++++++------ 1 file changed, 9 insertions(+), 6 deletions(-) diff --git a/adapters/typescript/pi/README.md b/adapters/typescript/pi/README.md index ccb3edeec..caa1050d6 100644 --- a/adapters/typescript/pi/README.md +++ b/adapters/typescript/pi/README.md @@ -228,12 +228,15 @@ summary call and keeps only the counts the provider reported. A reported zero counts as 0, including a report of all zeros. When a call left out its input count, the result leaves out `input_tokens`, `cached_input_tokens`, and `total_tokens`. Anthropic Messages usage reports input without the cache, so -its input count also needs both cache fields. When a call left out its output -count, the result leaves out `output_tokens` and `total_tokens`; for Anthropic -Messages, only a `message_delta` reports the final output count. When a call -sent no cache details, the result leaves out `cached_input_tokens`. A wire -protocol outside Pi's built-in ones reports no counts the adapter can check, so -its calls leave the result without `usage`. +its input count also needs both cache fields. Bedrock Converse also reports +input without the cache, but its API documents the cache fields as optional, so +its input count needs only `inputTokens` rather than going unknown on every +call that omits them. When a call left out its output count, the result leaves +out `output_tokens` and `total_tokens`; for Anthropic Messages, only a +`message_delta` reports the final output count. When a call sent no cache +details, the result leaves out `cached_input_tokens`. A wire protocol outside +Pi's built-in ones reports no counts the adapter can check, so its calls leave +the result without `usage`. The result has no `usage` when Pi produced no assistant message, when a model call reported none of the counts, when a tool result has no complete usage From 9d6bc218bf4debd918ea49ad44eb8e7789fba6e4 Mon Sep 17 00:00:00 2001 From: mnajafian-nv Date: Mon, 19 Oct 2026 15:28:16 -0700 Subject: [PATCH 12/68] 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 99ac9dffeb57945bf13cf71225c8c5c20205ebe6 Mon Sep 17 00:00:00 2001 From: mnajafian-nv Date: Mon, 19 Oct 2026 15:59:57 -0700 Subject: [PATCH 13/68] fix: leave Hermes usage unknown after an invalid response Hermes' check_api_response hands a response that fails validation to retry_invalid_response and returns before record_response_usage. The provider can still bill that call, and a successful retry then produced a measured invocation total that left it out. The adapter's process-wide wrapper now covers retry_invalid_response next to the truncation and refusal handlers, with the same start-time check. Unless the handler falls through to usage recording, the response counts like a response without usage and the invocation reports no usage. The start check now also rejects a handler name that the installed Hermes does not define. The Hermes e2e test answers with the router timeout sentinel that Hermes' chat-completions transport rejects in validate_response, with usage that bills the prompt and no completion tokens, followed by a successful retry through Hermes' real turn loop. It fails without the wrapper. The content_filter e2e case is replaced by a unit test of the refusal handler's dispatch with a fallback reply: a terminal refusal result carries no counters, so the e2e case passed without the wrapper, and Fabric's Hermes config enables no fallback. Signed-off-by: mnajafian-nv --- .../nemo_fabric_adapters/hermes/adapter.py | 22 ++++++---- tests/_utils/mock_api_server.py | 13 +++++- tests/adapters/test_hermes_adapter.py | 44 ++++++++++++++++++- tests/e2e/test_hermes_e2e.py | 34 +++++++------- 4 files changed, 87 insertions(+), 26 deletions(-) diff --git a/adapters/python/hermes/src/nemo_fabric_adapters/hermes/adapter.py b/adapters/python/hermes/src/nemo_fabric_adapters/hermes/adapter.py index 2ddc08deb..246e5bb1e 100755 --- a/adapters/python/hermes/src/nemo_fabric_adapters/hermes/adapter.py +++ b/adapters/python/hermes/src/nemo_fabric_adapters/hermes/adapter.py @@ -439,15 +439,16 @@ def _count_unreported_counts(self, agent: Any) -> None: omitted count from a reported zero, and Hermes passes the raw usage to no per-agent hook. The adapter therefore registers the agent with process-wide wrappers around the functions Hermes' turn loop calls for - each response; see ``_install_usage_observer``. A refused or truncated - response that Hermes never records counts as a response without usage. + each response; see ``_install_usage_observer``. An invalid, refused or + truncated response that Hermes never records counts as a response + without usage. """ _install_usage_observer() def count_unreported(handler: str, response: Any, outcome: Any) -> None: if handler != "record_response_usage": - # Hermes returns on a refused or truncated response before - # recording its usage, unless the truncation falls through. + # Hermes returns on an invalid, refused or truncated response + # before recording its usage, unless the handler falls through. if getattr(outcome, "action", None) != "fallthrough": self._uncounted_responses += 1 return @@ -689,11 +690,12 @@ def _invoke_hermes_turn( ) # The functions ``agent.turn_response_check.check_api_response`` calls with # (agent, response, ...) for a provider response, by the module defining -# each. record_response_usage folds the usage; on a content-filter refusal or -# a reply cut off at the token limit, Hermes calls one of the other two and -# returns before recording the usage. +# each. record_response_usage folds the usage; on an invalid response, a +# content-filter refusal or a reply cut off at the token limit, Hermes calls +# one of the other three and returns before recording the usage. _RESPONSE_HANDLERS = { "record_response_usage": "agent.turn_usage", + "retry_invalid_response": "agent.turn_response_check", "handle_content_policy_refusal": "agent.turn_truncation", "recover_from_truncation": "agent.turn_truncation", } @@ -733,8 +735,10 @@ def _install_usage_observer() -> None: current = getattr(turn_response_check, name, None) if getattr(current, "nemo_fabric_observed", False): continue - if name not in called or current is not getattr( - importlib.import_module(module), name, None + if ( + current is None + or name not in called + or current is not getattr(importlib.import_module(module), name, None) ): raise lifecycle.LifecycleError( "hermes_usage_accounting_unsupported", diff --git a/tests/_utils/mock_api_server.py b/tests/_utils/mock_api_server.py index 7db01e615..31007fe49 100644 --- a/tests/_utils/mock_api_server.py +++ b/tests/_utils/mock_api_server.py @@ -11,6 +11,9 @@ from fastapi.responses import JSONResponse, Response, StreamingResponse import uvicorn +# The sentinel an OpenAI-compatible router returns as the assistant message +# when its upstream connection times out. +_ROUTER_TIMEOUT_TEXT = "Connect timeout, please try again later." # Chat-completion usage the mock reports unless a scenario sets another one. _ZERO_USAGE = {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0} @@ -22,7 +25,8 @@ def mock_api_server(port: int) -> Iterator[str]: Use the /_requests endpoint to inspect captured chat-completion payloads after a test action. Use the /_scenario endpoint to configure a status code, ordered tool calls, or the - chat-completion usage object and finish reason for subsequent requests. + chat-completion usage object, finish reason and number of leading invalid responses + for subsequent requests. Use the /_reset endpoint to clear all retained request and scenario state. Args: @@ -42,6 +46,7 @@ def reset_state() -> None: app.state.mcp_authorization_headers = [] app.state.usage = dict(_ZERO_USAGE) app.state.finish_reason = "stop" + app.state.invalid_responses = 0 reset_state() @@ -82,6 +87,7 @@ async def scenario(request: Request) -> dict[str, object]: app.state.tool_call_index = 0 app.state.usage = payload.get("usage", dict(_ZERO_USAGE)) app.state.finish_reason = payload.get("finish_reason", "stop") + app.state.invalid_responses = int(payload.get("invalid_responses", 0)) return { "status_code": app.state.status_code, "tool_call": tool_call, @@ -185,6 +191,11 @@ async def chat_completions(request: Request): ] latest = user_messages[-1].get("content", "") if user_messages else "" content = f"echo user_count={len(user_messages)} latest={latest}" + if app.state.invalid_responses > 0: + # A router timeout reported as a successful completion, which + # OpenAI-compatible clients such as Hermes reject as invalid. + app.state.invalid_responses -= 1 + content = _ROUTER_TIMEOUT_TEXT if payload.get("stream") is True: return StreamingResponse( _stream_chat_completion( diff --git a/tests/adapters/test_hermes_adapter.py b/tests/adapters/test_hermes_adapter.py index 92edade5b..4e2f78b4a 100644 --- a/tests/adapters/test_hermes_adapter.py +++ b/tests/adapters/test_hermes_adapter.py @@ -13,10 +13,12 @@ import sys import threading import tomllib +from collections.abc import Callable from io import StringIO from pathlib import Path from types import ModuleType from types import SimpleNamespace +from typing import Any from unittest.mock import MagicMock import pytest @@ -1630,7 +1632,7 @@ class _UsageRecordingAgent: def __init__( self, - turns: list[list[dict[str, object] | None]], + turns: list[list[dict[str, object] | Callable[[Any], None] | None]], api_mode: str = "chat_completions", ) -> None: from agent.context_compressor import ContextCompressor @@ -1671,6 +1673,9 @@ def run_conversation(self, user_message: str) -> dict[str, object]: messages = [{"role": "user", "content": user_message}] for usage in self._turns.pop(0): + if callable(usage): + usage(self) + continue turn_response_check.record_response_usage( self, SimpleNamespace(usage=usage), @@ -1856,6 +1861,43 @@ async def test_hermes_usage_reports_cache_counts_only_when_the_provider_sent_the assert with_cache_fields.usage.total_tokens == 135 +@pytest.mark.parametrize( + ("action", "counted"), [("break", False), ("fallthrough", True)] +) +async def test_hermes_refused_response_before_a_fallback_reply(action, counted): + # Hermes hands a content-filter refusal to handle_content_policy_refusal; + # with a fallback configured it returns "break" and the fallback model + # answers, so the refused call is billed but never recorded. + refusal = adapter._observed_response_handler( + "handle_content_policy_refusal", + lambda agent, response, *args, **kwargs: SimpleNamespace(action=action), + ) + + def refused(agent: Any) -> None: + refused_usage = {"prompt_tokens": 40, "completion_tokens": 0} + refusal(agent, SimpleNamespace(usage=refused_usage), None) + + agent = _UsageRecordingAgent( + [ + [refused, {"prompt_tokens": 50, "completion_tokens": 5}], + [{"prompt_tokens": 60, "completion_tokens": 6}], + ] + ) + runtime = _usage_recording_runtime(agent) + invocation = _usage_invocation() + + fallback = await runtime.invoke(*invocation) + after = await runtime.invoke(*invocation) + + if counted: + assert fallback.usage is not None + assert fallback.usage.input_tokens == 50 + else: + assert fallback.usage is None + assert after.usage is not None + assert (after.usage.input_tokens, after.usage.output_tokens) == (60, 6) + + async def test_hermes_stop_unregisters_the_agent_usage_observer(): agent = _UsageRecordingAgent([]) runtime = _usage_recording_runtime(agent) diff --git a/tests/e2e/test_hermes_e2e.py b/tests/e2e/test_hermes_e2e.py index b902b6be0..8df6ea084 100644 --- a/tests/e2e/test_hermes_e2e.py +++ b/tests/e2e/test_hermes_e2e.py @@ -154,24 +154,30 @@ async def test_hermes_persistent_host_reports_usage_for_each_invocation( @pytest.mark.usefixtures("mock_nvidia_api_key") @pytest.mark.parametrize( - ("finish_reason", "usage"), + "scenario", [ # A reply cut off at max_tokens bills the whole budget. - ( - "length", - {"prompt_tokens": 40, "completion_tokens": 512, "total_tokens": 552}, - ), - ( - "content_filter", - {"prompt_tokens": 40, "completion_tokens": 0, "total_tokens": 40}, - ), + { + "finish_reason": "length", + "usage": { + "prompt_tokens": 40, + "completion_tokens": 512, + "total_tokens": 552, + }, + }, + # A router timeout reported as a completion that billed the prompt, + # then a successful retry. + { + "invalid_responses": 1, + "usage": {"prompt_tokens": 40, "completion_tokens": 0, "total_tokens": 40}, + }, ], + ids=["length", "invalid_then_retry"], ) async def test_hermes_usage_is_unknown_when_hermes_skips_a_response( code_review_agent_dir: Path, api_server: str, - finish_reason: str, - usage: dict[str, int], + scenario: dict[str, object], ): os.environ["ADAPTER_PYTHON"] = sys.executable config = hermes_config() @@ -182,15 +188,13 @@ async def test_hermes_usage_is_unknown_when_hermes_skips_a_response( ) as runtime: counted = await runtime.invoke(input="first") scenario_response = requests.post( - f"{api_server}/_scenario", - json={"finish_reason": finish_reason, "usage": usage}, - timeout=5, + f"{api_server}/_scenario", json=scenario, timeout=5 ) scenario_response.raise_for_status() skipped = await runtime.invoke(input="second") assert counted.usage is not None, counted.to_mapping() - # Hermes returns on a truncated or refused reply before recording its + # Hermes returns on a truncated or invalid response before recording its # usage, so the invocation's counters miss a billed call. assert skipped.usage is None, skipped.to_mapping() From 1f9af061f9820bb709bdcde71774cb22c0b78939 Mon Sep 17 00:00:00 2001 From: mnajafian-nv Date: Mon, 19 Oct 2026 16:07:44 -0700 Subject: [PATCH 14/68] docs: describe the Hermes response handlers the adapter observes The Hermes README still said the adapter wraps only record_response_usage and listed only truncated and refused replies. The adapter also observes retry_invalid_response, and a streamed call that Hermes retries inside its API call, such as an empty stream, never reaches these handlers, so its billed usage is not observed. The README now names all three handlers and states that limit. Signed-off-by: mnajafian-nv --- adapters/python/hermes/README.md | 32 ++++++++++++++++++-------------- 1 file changed, 18 insertions(+), 14 deletions(-) diff --git a/adapters/python/hermes/README.md b/adapters/python/hermes/README.md index 0b05bc8f7..588e1b1c3 100644 --- a/adapters/python/hermes/README.md +++ b/adapters/python/hermes/README.md @@ -127,26 +127,30 @@ rather than a carried-over or whole-session total, and the one after it reports its own change again. Hermes adds no tokens for a provider response that carried no usage, so an -invocation with such a response has no `usage`. Hermes also returns before -recording the usage of a reply that hit the token limit or a content filter, so -an invocation with such a reply has no `usage` either. Hermes reads a count -that a response's usage left out as 0, so the adapter checks each response's -usage before Hermes normalizes it. A reported zero still counts as 0. When a -response left out its input count, the invocation's `usage` leaves out -`input_tokens`, `cached_input_tokens`, and `total_tokens`. Anthropic Messages -usage reports input without the cache, so its input count also needs both cache -fields. When a response left out its output count, the `usage` leaves out -`output_tokens` and `total_tokens`. When a response sent no cache details, the -`usage` leaves out `cached_input_tokens`. +invocation with such a response has no `usage`. Hermes also skips usage +recording for a response it rejects as invalid, a reply that hit the token +limit, and a reply stopped by a content filter, so an invocation with such a +response has no `usage` either. A streamed call that fails and is retried +inside Hermes' API call, such as an empty stream, never reaches these checks, +so its billed usage is not observed and the invocation can undercount. Hermes +reads a count that a response's usage left out as 0, so the adapter checks each +response's usage before Hermes normalizes it. A reported zero still counts as +0. When a response left out its input count, the invocation's `usage` leaves +out `input_tokens`, `cached_input_tokens`, and `total_tokens`. Anthropic +Messages usage reports input without the cache, so its input count also needs +both cache fields. When a response left out its output count, the `usage` +leaves out `output_tokens` and `total_tokens`. When a response sent no cache +details, the `usage` leaves out `cached_input_tokens`. This check covers the usage that the OpenAI-compatible, Responses, and Anthropic clients return. Hermes' native Gemini and Bedrock transports, and its Anthropic transport when a stream omits usage, fill missing counts with 0 before the adapter sees them. Hermes' Codex app-server runtime records usage outside this path. Hermes passes a response's raw usage to no per-agent hook, -so the adapter wraps `agent.turn_response_check.record_response_usage` once per -process. The runtime fails to start if the installed Hermes records usage -through another call. +so the adapter wraps `record_response_usage` and the three handlers for +invalid, truncated, and refused responses in `agent.turn_response_check` once +per process. The runtime fails to start if the installed Hermes no longer calls +them by those names. ## Execution Model From cdc0636a647d562a02a72110f9921b8943d3bd04 Mon Sep 17 00:00:00 2001 From: mnajafian-nv Date: Mon, 19 Oct 2026 17:00:16 -0700 Subject: [PATCH 15/68] 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 9be221e075cda77e95c3ff46419902913c56d359 Mon Sep 17 00:00:00 2001 From: mnajafian-nv Date: Mon, 19 Oct 2026 17:21:53 -0700 Subject: [PATCH 16/68] fix: treat empty Hermes usage as a report with every count omitted The Hermes adapter's response observer skipped a falsy usage object, so an empty usage dict never reached the omitted-count check. Hermes' record_response_usage also treats an empty usage object as missing and calls note_usage_less_response for it, so the invocation's usage was already unknown through the compressor hook. The observer now checks for None explicitly and counts an empty usage object as a report with input, cached input and output left out, so the result no longer depends on that hook alone. The new test runs an empty usage object through Hermes' real record_response_usage and expects unknown usage, followed by a counted invocation. The Pi README said a tool result without a complete usage report leaves the result without usage. A tool result with no usage field adds nothing; only a usage object with a missing or invalid count makes the invocation's usage unknown. The README now says both. Signed-off-by: mnajafian-nv --- .../nemo_fabric_adapters/hermes/adapter.py | 5 +++-- adapters/typescript/pi/README.md | 19 +++++++++-------- tests/adapters/test_hermes_adapter.py | 21 +++++++++++++++++++ 3 files changed, 34 insertions(+), 11 deletions(-) diff --git a/adapters/python/hermes/src/nemo_fabric_adapters/hermes/adapter.py b/adapters/python/hermes/src/nemo_fabric_adapters/hermes/adapter.py index 246e5bb1e..aa128bea7 100755 --- a/adapters/python/hermes/src/nemo_fabric_adapters/hermes/adapter.py +++ b/adapters/python/hermes/src/nemo_fabric_adapters/hermes/adapter.py @@ -453,8 +453,9 @@ def count_unreported(handler: str, response: Any, outcome: Any) -> None: self._uncounted_responses += 1 return usage = getattr(response, "usage", None) - # A response without usage is counted by the compressor hook. - if usage: + # A response without usage is counted by the compressor hook. An + # empty usage object is a report with every count left out. + if usage is not None: for count in _omitted_counts(agent, usage): self._unreported_counts[count] += 1 diff --git a/adapters/typescript/pi/README.md b/adapters/typescript/pi/README.md index caa1050d6..213419e61 100644 --- a/adapters/typescript/pi/README.md +++ b/adapters/typescript/pi/README.md @@ -239,15 +239,16 @@ Pi's built-in ones reports no counts the adapter can check, so its calls leave the result without `usage`. The result has no `usage` when Pi produced no assistant message, when a model -call reported none of the counts, when a tool result has no complete usage -report, when an assistant message was aborted, when a compaction was aborted or -reported no usage, or when Pi retried a compaction summary call, because Pi -keeps no usage for the failed attempt. It also has no `usage` when Pi recorded -a cache warm-up during the invocation: Pi replays a warm-up outside the agent's -stream function, so the adapter cannot tie its raw report to the warm-up's -usage entry. It also has no `usage` when a compaction cuts inside a turn, -because Pi then summarizes the earlier history and the start of that turn in -separate calls and reports only their combined usage. +call reported none of the counts, when a tool result's `usage` has a missing or +invalid count, when an assistant message was aborted, when a compaction was +aborted or reported no usage, or when Pi retried a compaction summary call, +because Pi keeps no usage for the failed attempt. It also has no `usage` when +Pi recorded a cache warm-up during the invocation: Pi replays a warm-up outside +the agent's stream function, so the adapter cannot tie its raw report to the +warm-up's usage entry. It also has no `usage` when a compaction cuts inside a +turn, because Pi then summarizes the earlier history and the start of that turn +in separate calls and reports only their combined usage. A tool result without +a `usage` field adds nothing and leaves the result's `usage` unchanged. ## Custom Tool Modules diff --git a/tests/adapters/test_hermes_adapter.py b/tests/adapters/test_hermes_adapter.py index 4e2f78b4a..b78ef1803 100644 --- a/tests/adapters/test_hermes_adapter.py +++ b/tests/adapters/test_hermes_adapter.py @@ -1813,6 +1813,27 @@ async def test_hermes_usage_with_partial_provider_report_leaves_omitted_count_un assert after_partial.usage.total_tokens == 440 +async def test_hermes_usage_with_empty_provider_usage_is_unknown(): + # Hermes treats an empty usage object like a missing one; the adapter also + # counts it as a report with every count left out. + agent = _UsageRecordingAgent( + [[{}], [{"prompt_tokens": 400, "completion_tokens": 40}]] + ) + runtime = _usage_recording_runtime(agent) + invocation = _usage_invocation() + + empty = await runtime.invoke(*invocation) + after_empty = await runtime.invoke(*invocation) + + assert agent.session_api_calls == 2 + assert empty.usage is None + assert after_empty.usage is not None + assert (after_empty.usage.input_tokens, after_empty.usage.output_tokens) == ( + 400, + 40, + ) + + async def test_hermes_usage_reports_cache_counts_only_when_the_provider_sent_them(): chat = _UsageRecordingAgent( [ From f47ca12aeda7cc6ed90c7de4d29ca89a1664828b Mon Sep 17 00:00:00 2001 From: mnajafian-nv Date: Mon, 19 Oct 2026 17:39:58 -0700 Subject: [PATCH 17/68] docs: exclude cache warm-ups from the Pi usage list Signed-off-by: mnajafian-nv --- adapters/typescript/pi/README.md | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/adapters/typescript/pi/README.md b/adapters/typescript/pi/README.md index 213419e61..c45669c53 100644 --- a/adapters/typescript/pi/README.md +++ b/adapters/typescript/pi/README.md @@ -216,11 +216,11 @@ independent of native OpenAI streaming, so the adapter descriptor's Each result's `usage` sums the tokens Pi records during the invocation, counted the way Pi's own usage totals count them: assistant messages, tool results that -report usage, finished compaction summaries, and usage entries such as cache -warm-ups. It holds `input_tokens` with cache reads and cache writes included, -`cached_input_tokens`, `output_tokens`, and `total_tokens`. Pi computes its -cost figures from the model catalog's rates, so the adapter does not report -them as `cost_usd`. +report usage, finished compaction summaries, and usage entries other than cache +warm-ups, which leave `usage` unknown (see below). It holds `input_tokens` with +cache reads and cache writes included, `cached_input_tokens`, `output_tokens`, +and `total_tokens`. Pi computes its cost figures from the model catalog's +rates, so the adapter does not report them as `cost_usd`. Pi fills a count that the provider did not report with 0, so the adapter reads the provider's raw stream events for each assistant message and compaction From e2e1c42edfeebb1add56f7de5fe2a4fe37079cbe Mon Sep 17 00:00:00 2001 From: mnajafian-nv Date: Mon, 19 Oct 2026 22:47:39 -0700 Subject: [PATCH 18/68] fix: count Hermes usage from session counters and skipped responses The adapter rebuilt its baseline only from a turn's result, so after a failed, cancelled or counterless invocation the next invocation had no usage. It now reads the agent's session counters right before and after each turn, so every invocation reports its own change, including a failure result. Hermes returns on a truncated, invalid or refused response before adding its usage to the counters. With the 512-token default, truncation is common, and it made the whole invocation unknown. The adapter now adds the counts such a response reported, normalized with Hermes' own normalize_usage for the provider and API mode captured before the handler can switch to a fallback route. A response that falls through to record_response_usage is counted there only. The handler wrapper also passed the response positionally, but Hermes calls retry_invalid_response with a keyword-only response, so every invalid response raised TypeError and failed the turn. The wrapper now forwards the call unchanged and reads the response from the handler's signature. Signed-off-by: mnajafian-nv --- .../nemo_fabric_adapters/hermes/adapter.py | 191 ++++++--- tests/_utils/mock_api_server.py | 23 +- tests/adapters/test_hermes_adapter.py | 375 +++++++++++------- tests/e2e/test_hermes_e2e.py | 87 +++- 4 files changed, 438 insertions(+), 238 deletions(-) diff --git a/adapters/python/hermes/src/nemo_fabric_adapters/hermes/adapter.py b/adapters/python/hermes/src/nemo_fabric_adapters/hermes/adapter.py index aa128bea7..238446418 100755 --- a/adapters/python/hermes/src/nemo_fabric_adapters/hermes/adapter.py +++ b/adapters/python/hermes/src/nemo_fabric_adapters/hermes/adapter.py @@ -69,12 +69,12 @@ def __init__(self) -> None: self._relay_plugin_config: dict[str, Any] | None = None self._relay_plugin_config_path: Path | None = None self._active_invoke_task: asyncio.Task[tuple[dict[str, Any], str]] | None = None - # Hermes reports session totals; each invocation reports the change - # from this baseline. A counter missing from it has no baseline. - self._session_token_totals: dict[str, int] = _fresh_session_totals() self._uncounted_responses = 0 # Provider responses whose usage left out a count, by count. self._unreported_counts = dict.fromkeys(_REPORTED_COUNTS, 0) + # Tokens of provider responses Hermes returned on before adding them to + # its session counters, in the counters' buckets. + self._skipped_response_tokens = dict.fromkeys(_SESSION_TOKEN_KEYS, 0) async def start(self, payload: dict[str, Any]) -> None: if self._started: @@ -226,7 +226,6 @@ async def start(self, payload: dict[str, Any]) -> None: session_db=self._session_db, ) ) - self._session_token_totals = _fresh_session_totals() self._observe_usage_reports(self._agent) self._agent_config = agent_config self._started = True @@ -295,12 +294,12 @@ def run_hermes_turn() -> tuple[dict[str, Any], str]: "hermes_invocation_in_progress", "Hermes runtime already has an active invocation", ) - # Until this invocation's result arrives the baseline is unknown, so a - # failed, cancelled or timed-out invocation never hands its tokens to - # the next one. - previous_totals, self._session_token_totals = self._session_token_totals, {} + # No worker is running here, so the agent's session counters hold every + # earlier invocation's tokens, including a failed or cancelled one's. + previous_totals = _session_counters(self._agent) uncounted_before = self._uncounted_responses unreported_before = dict(self._unreported_counts) + skipped_before = dict(self._skipped_response_tokens) invoke_task = asyncio.create_task(asyncio.to_thread(run_hermes_turn)) self._active_invoke_task = invoke_task @@ -385,9 +384,13 @@ def clear_active_invoke_task( code="hermes_reported_failure", message="Hermes did not complete the invocation", ) - usage, self._session_token_totals = _invocation_usage( - result, + usage = _invocation_usage( + _session_counters(self._agent), previous_totals, + skipped={ + key: tokens - skipped_before[key] + for key, tokens in self._skipped_response_tokens.items() + }, unreported=frozenset( count for count, responses in self._unreported_counts.items() @@ -395,9 +398,8 @@ def clear_active_invoke_task( ), ) if self._uncounted_responses != uncounted_before: - # A provider response without usage, or one Hermes returned on - # before recording its usage, added no tokens, so the change in - # the counters is an undercount, not this invocation's. + # Hermes counted a provider response without usage as 0 tokens, + # so the change in the counters is an undercount. usage = None return AgentRunResult( status=AgentRunStatus.FAILED if failed else AgentRunStatus.SUCCEEDED, @@ -439,28 +441,52 @@ def _count_unreported_counts(self, agent: Any) -> None: omitted count from a reported zero, and Hermes passes the raw usage to no per-agent hook. The adapter therefore registers the agent with process-wide wrappers around the functions Hermes' turn loop calls for - each response; see ``_install_usage_observer``. An invalid, refused or - truncated response that Hermes never records counts as a response - without usage. + each response; see ``_install_usage_observer``. Hermes returns on an + invalid, refused or truncated response before adding its usage to the + session counters, so the adapter adds that response's reported counts + itself. """ _install_usage_observer() - def count_unreported(handler: str, response: Any, outcome: Any) -> None: - if handler != "record_response_usage": - # Hermes returns on an invalid, refused or truncated response - # before recording its usage, unless the handler falls through. - if getattr(outcome, "action", None) != "fallthrough": - self._uncounted_responses += 1 - return + def count_unreported( + handler: str, + response: Any, + outcome: Any, + provider: Any, + api_mode: Any, + ) -> None: usage = getattr(response, "usage", None) - # A response without usage is counted by the compressor hook. An - # empty usage object is a report with every count left out. - if usage is not None: - for count in _omitted_counts(agent, usage): - self._unreported_counts[count] += 1 + if handler == "record_response_usage": + # A response without usage is counted by the compressor hook. + # An empty usage object is a report with every count left out. + if usage is not None: + for count in _omitted_counts(usage, provider, api_mode): + self._unreported_counts[count] += 1 + return + if not _skips_usage_recording(handler, outcome): + return + for count in _omitted_counts(usage, provider, api_mode): + self._unreported_counts[count] += 1 + if usage: + self._add_skipped_response_tokens(usage, provider, api_mode) _USAGE_OBSERVERS[agent] = count_unreported + def _add_skipped_response_tokens( + self, usage: Any, provider: Any, api_mode: Any + ) -> None: + """Add a response's usage to the tokens Hermes left out of its counters. + + Hermes' own normalization splits the usage into its counters' buckets. + It fills an omitted field with 0; every count that needs an omitted + field is already marked unreported. + """ + from agent.usage_pricing import normalize_usage + + canonical = normalize_usage(usage, provider=provider, api_mode=api_mode) + for key in _SESSION_TOKEN_KEYS: + self._skipped_response_tokens[key] += getattr(canonical, key) + async def _authenticate_mcp_servers(self) -> None: if self._mcp_authentication_checked: return @@ -684,16 +710,18 @@ def _invoke_hermes_turn( # Response observers by Hermes agent, called with each response Hermes' -# turn loop handles. Delegated subagents are never registered: their tokens -# never reach the registered agent's counters. -_USAGE_OBSERVERS: weakref.WeakKeyDictionary[Any, Callable[[str, Any, Any], None]] = ( - weakref.WeakKeyDictionary() -) +# turn loop handles, the handler's outcome, and the provider and API mode +# that served the response. Delegated subagents are never registered: their +# tokens never reach the registered agent's counters. +_USAGE_OBSERVERS: weakref.WeakKeyDictionary[ + Any, Callable[[str, Any, Any, Any, Any], None] +] = weakref.WeakKeyDictionary() # The functions ``agent.turn_response_check.check_api_response`` calls with # (agent, response, ...) for a provider response, by the module defining # each. record_response_usage folds the usage; on an invalid response, a # content-filter refusal or a reply cut off at the token limit, Hermes calls -# one of the other three and returns before recording the usage. +# one of the other three and usually returns before recording the usage (see +# _skips_usage_recording). _RESPONSE_HANDLERS = { "record_response_usage": "agent.turn_usage", "retry_invalid_response": "agent.turn_response_check", @@ -703,13 +731,24 @@ def _invoke_hermes_turn( def _observed_response_handler(name: str, handler: Callable[..., Any]) -> Any: - """Wrap a Hermes response handler to show each call to the agent's observer.""" + """Wrap a Hermes response handler to show each call to the agent's observer. - def observed(agent: Any, response: Any, *args: Any, **kwargs: Any) -> Any: - outcome = handler(agent, response, *args, **kwargs) + Hermes passes the response positionally to some handlers and by keyword + to others, so the wrapper forwards the call unchanged and reads the + handler's ``response`` argument from its signature. + """ + signature = inspect.signature(handler) + + def observed(agent: Any, *args: Any, **kwargs: Any) -> Any: + # A handler can switch the agent to a fallback provider, so read the + # route that served this response before the handler runs. + provider = getattr(agent, "provider", None) + api_mode = getattr(agent, "api_mode", None) + outcome = handler(agent, *args, **kwargs) observer = _USAGE_OBSERVERS.get(agent) if observer is not None: - observer(name, response, outcome) + response = signature.bind(agent, *args, **kwargs).arguments["response"] + observer(name, response, outcome, provider, api_mode) return outcome observed.nemo_fabric_observed = True # type: ignore[attr-defined] @@ -740,6 +779,7 @@ def _install_usage_observer() -> None: current is None or name not in called or current is not getattr(importlib.import_module(module), name, None) + or "response" not in inspect.signature(current).parameters ): raise lifecycle.LifecycleError( "hermes_usage_accounting_unsupported", @@ -782,10 +822,28 @@ def _install_usage_observer() -> None: } -def _omitted_counts(agent: Any, usage: Any) -> list[str]: +def _skips_usage_recording(handler: str, outcome: Any) -> bool: + """Return whether ``check_api_response`` returns after ``handler`` without + recording the response's usage. + + It returns after a content-filter refusal, after an invalid response + unless the handler's action is ``fallthrough``, and after a truncated + response when the action is ``return``, ``break`` or ``continue`` + (``agent/turn_response_check.py``). Any other outcome falls through to + ``record_response_usage``, which adds the response's usage itself. + """ + action = getattr(outcome, "action", None) + if handler == "recover_from_truncation": + return action in ("return", "break", "continue") + if handler == "retry_invalid_response": + return action != "fallthrough" + return True + + +def _omitted_counts(usage: Any, provider: Any, api_mode: Any) -> list[str]: """Return the counts that the provider's raw ``usage`` left out.""" - mode = str(getattr(agent, "api_mode", None) or "").strip().lower() - provider = str(getattr(agent, "provider", None) or "").strip().lower() + mode = str(api_mode or "").strip().lower() + provider = str(provider or "").strip().lower() if mode == "anthropic_messages" or provider == "anthropic": fields = _ANTHROPIC_FIELDS elif mode == "codex_responses": @@ -812,38 +870,43 @@ def _is_reported_count(usage: Any, path: str) -> bool: return isinstance(value, int) and not isinstance(value, bool) and value >= 0 -def _fresh_session_totals() -> dict[str, int]: - """Return the baseline of a new Hermes agent, whose counters start at 0.""" - return dict.fromkeys(_SESSION_TOKEN_KEYS, 0) +def _session_counters(agent: Any) -> dict[str, int]: + """Return the agent's session token counters that hold a valid count. + + Hermes stamps exactly these attributes into a finished turn's result, but + its early failure results carry none, so the adapter reads the agent. + """ + return { + key: value + for key in _SESSION_TOKEN_KEYS + if isinstance((value := getattr(agent, f"session_{key}", None)), int) + and not isinstance(value, bool) + and 0 <= value <= (1 << 64) - 1 + } def _invocation_usage( - result: dict[str, Any], + totals: dict[str, int], previous_totals: dict[str, int], *, + skipped: dict[str, int] | None = None, unreported: frozenset[str] = frozenset(), -) -> tuple[AgentUsage | None, dict[str, int]]: - """Return this invocation's usage and the session totals to keep as baseline. +) -> AgentUsage | None: + """Return this invocation's usage from the agent's session counters. Hermes keeps one session across a runtime's invocations, so each counter is - the change since ``previous_totals``. A counter that is missing from the - result or from ``previous_totals`` (no baseline: the previous invocation - failed or its result held no counters), or that went down, is unknown for - this invocation. The usage therefore never holds a whole-session total or - another invocation's tokens. ``unreported`` names the counts ("input", - "cached", "output") that a provider response of this invocation left out - and Hermes counted as 0. Those counts are unknown, and so is the total when - input or output is. + the change from ``previous_totals``, read before the turn, to ``totals``, + read after it, plus the ``skipped`` tokens of responses Hermes left out of + its counters. A counter that is missing from either reading, or that went + down (Hermes reset its session), is unknown for this invocation, so the + usage never holds a whole-session total. ``unreported`` names the counts + ("input", "cached", "output") that a provider response of this invocation + left out and Hermes counted as 0. Those counts are unknown, and so is the + total when input or output is. """ - totals = { - key: value - for key in _SESSION_TOKEN_KEYS - if isinstance((value := result.get(key)), int) - and not isinstance(value, bool) - and 0 <= value <= (1 << 64) - 1 - } + skipped = skipped or {} change = { - key: value - previous_totals[key] + key: value - previous_totals[key] + skipped.get(key, 0) for key, value in totals.items() if key in previous_totals and value >= previous_totals[key] } @@ -864,7 +927,7 @@ def _invocation_usage( fields["output_tokens"] = change["output_tokens"] if "input_tokens" in fields and "output_tokens" in fields: fields["total_tokens"] = fields["input_tokens"] + fields["output_tokens"] - return (AgentUsage(**fields) if fields else None), totals + return AgentUsage(**fields) if fields else None def filter_supported_kwargs(callable_obj: Any, **kwargs: Any) -> dict[str, Any]: diff --git a/tests/_utils/mock_api_server.py b/tests/_utils/mock_api_server.py index 31007fe49..3661775af 100644 --- a/tests/_utils/mock_api_server.py +++ b/tests/_utils/mock_api_server.py @@ -25,8 +25,9 @@ def mock_api_server(port: int) -> Iterator[str]: Use the /_requests endpoint to inspect captured chat-completion payloads after a test action. Use the /_scenario endpoint to configure a status code, ordered tool calls, or the - chat-completion usage object, finish reason and number of leading invalid responses - for subsequent requests. + chat-completion usage object and finish reason for subsequent requests. Its ordered + ``replies`` override the usage or finish reason of the next chat completions, one + reply each, or make one invalid. Use the /_reset endpoint to clear all retained request and scenario state. Args: @@ -46,7 +47,7 @@ def reset_state() -> None: app.state.mcp_authorization_headers = [] app.state.usage = dict(_ZERO_USAGE) app.state.finish_reason = "stop" - app.state.invalid_responses = 0 + app.state.replies = [] reset_state() @@ -87,7 +88,7 @@ async def scenario(request: Request) -> dict[str, object]: app.state.tool_call_index = 0 app.state.usage = payload.get("usage", dict(_ZERO_USAGE)) app.state.finish_reason = payload.get("finish_reason", "stop") - app.state.invalid_responses = int(payload.get("invalid_responses", 0)) + app.state.replies = list(payload.get("replies", [])) return { "status_code": app.state.status_code, "tool_call": tool_call, @@ -191,16 +192,16 @@ async def chat_completions(request: Request): ] latest = user_messages[-1].get("content", "") if user_messages else "" content = f"echo user_count={len(user_messages)} latest={latest}" - if app.state.invalid_responses > 0: + reply = app.state.replies.pop(0) if app.state.replies else {} + usage = reply.get("usage", app.state.usage) + finish_reason = reply.get("finish_reason", app.state.finish_reason) + if reply.get("invalid", False): # A router timeout reported as a successful completion, which # OpenAI-compatible clients such as Hermes reject as invalid. - app.state.invalid_responses -= 1 content = _ROUTER_TIMEOUT_TEXT if payload.get("stream") is True: return StreamingResponse( - _stream_chat_completion( - payload, content, app.state.usage, app.state.finish_reason - ), + _stream_chat_completion(payload, content, usage, finish_reason), media_type="text/event-stream", ) @@ -214,10 +215,10 @@ async def chat_completions(request: Request): { "index": 0, "message": {"role": "assistant", "content": content}, - "finish_reason": app.state.finish_reason, + "finish_reason": finish_reason, } ], - "usage": app.state.usage, + "usage": usage, } ) diff --git a/tests/adapters/test_hermes_adapter.py b/tests/adapters/test_hermes_adapter.py index b78ef1803..fc4708c79 100644 --- a/tests/adapters/test_hermes_adapter.py +++ b/tests/adapters/test_hermes_adapter.py @@ -1370,28 +1370,47 @@ async def test_persistent_runtime_reuses_hermes_agent_session_and_history( *first_messages, {"role": "assistant", "content": "second response"}, ] - mock_ai_agent.run_conversation.side_effect = [ - { - "response": "first response", - "completed": True, - "failed": False, - "messages": first_messages, - "input_tokens": 100, - "cache_read_tokens": 0, - "cache_write_tokens": 50, - "output_tokens": 20, - }, - { - "response": "second response", - "completed": True, - "failed": False, - "messages": second_messages, - "input_tokens": 130, - "cache_read_tokens": 50, - "cache_write_tokens": 50, - "output_tokens": 35, - }, - ] + turns = iter( + [ + { + "response": "first response", + "completed": True, + "failed": False, + "messages": first_messages, + "input_tokens": 100, + "cache_read_tokens": 0, + "cache_write_tokens": 50, + "output_tokens": 20, + }, + { + "response": "second response", + "completed": True, + "failed": False, + "messages": second_messages, + "input_tokens": 130, + "cache_read_tokens": 50, + "cache_write_tokens": 50, + "output_tokens": 35, + }, + ] + ) + session_counters = ( + "input_tokens", + "cache_read_tokens", + "cache_write_tokens", + "output_tokens", + ) + for counter in session_counters: + setattr(mock_ai_agent, f"session_{counter}", 0) + + def run_conversation(*args, **kwargs): + # Hermes stamps its session counters into a finished turn's result. + result = next(turns) + for counter in session_counters: + setattr(mock_ai_agent, f"session_{counter}", result[counter]) + return result + + mock_ai_agent.run_conversation.side_effect = run_conversation mock_ai_agent_type = MagicMock(spec=AIAgent, return_value=mock_ai_agent) monkeypatch.setattr( mock_ai_agent_type.__init__.__func__, @@ -1537,102 +1556,66 @@ async def test_persistent_runtime_reuses_hermes_agent_session_and_history( ) -async def test_hermes_invoke_after_failed_or_counterless_turn_reports_no_usage( - started_hermes_runtime, -): - runtime, mock_turn = started_hermes_runtime +def test_invocation_usage_with_lower_or_missing_counters_leaves_them_unknown(): + counters = { + "input_tokens": 10, + "cache_read_tokens": 0, + "cache_write_tokens": 0, + "output_tokens": 4, + } + usage = adapter._invocation_usage( + counters, {**counters, "input_tokens": 50, "output_tokens": 1} + ) + missing = adapter._invocation_usage({}, counters) + no_baseline = adapter._invocation_usage(counters, {}) - def turn(**counters: int) -> tuple[dict[str, object], str]: - return ( - {"final_response": "done", "completed": True, "messages": [], **counters}, - "", - ) + assert usage is not None + assert usage.input_tokens is None + assert usage.output_tokens == 3 + assert missing is None + assert no_baseline is None - def totals(tokens: int) -> dict[str, int]: - return { - "input_tokens": tokens, - "cache_read_tokens": 0, - "cache_write_tokens": 0, - "output_tokens": tokens, - } - mock_turn.side_effect = [ - turn(**totals(100)), - turn(), - turn(**totals(300)), - RuntimeError("turn failed"), - turn(**totals(700)), - turn(**totals(800)), - ] - invocation = _invocation( - { - "runtime_context": _runtime_context( - runtime_id="runtime-incomplete" - ).to_mapping(), - "request": {"input": "next"}, - } +def test_session_counters_keep_only_valid_counts(): + agent = SimpleNamespace( + session_input_tokens=7, + session_cache_read_tokens=True, + session_cache_write_tokens=-1, + session_output_tokens="3", ) - first = await runtime.invoke(*invocation) - counterless = await runtime.invoke(*invocation) - after_counterless = await runtime.invoke(*invocation) - with pytest.raises(RuntimeError): - await runtime.invoke(*invocation) - after_failure = await runtime.invoke(*invocation) - recovered = await runtime.invoke(*invocation) + assert adapter._session_counters(agent) == {"input_tokens": 7} - assert first.usage is not None - assert (first.usage.input_tokens, first.usage.output_tokens) == (100, 100) - # No baseline: never the session total or a carried-over change. - assert counterless.usage is None - assert after_counterless.usage is None - assert after_failure.usage is None - assert recovered.usage is not None - assert (recovered.usage.input_tokens, recovered.usage.output_tokens) == (100, 100) +# A turn step that ends the turn with a failure result carrying no session +# counters, like Hermes' billing, content-policy and partial-turn results. +_COUNTERLESS_FAILURE = "counterless failure" -def test_invocation_usage_with_lower_or_missing_counters_leaves_them_unknown(): - usage, totals = adapter._invocation_usage( - { - "input_tokens": 10, - "cache_read_tokens": 0, - "cache_write_tokens": 0, - "output_tokens": 4, - }, - { - "input_tokens": 50, - "cache_read_tokens": 0, - "cache_write_tokens": 0, - "output_tokens": 1, - }, - ) - missing, _ = adapter._invocation_usage({"completed": True}, {}) - no_baseline, _ = adapter._invocation_usage( - { - "input_tokens": 10, - "cache_read_tokens": 0, - "cache_write_tokens": 0, - "output_tokens": 4, - }, - {}, - ) - assert usage is not None - assert usage.input_tokens is None - assert usage.output_tokens == 3 - assert totals["input_tokens"] == 10 - assert missing is None - assert no_baseline is None +def _record_response(agent: Any, response: Any) -> None: + """Fold a provider response the way Hermes' turn loop does.""" + from agent import turn_response_check + + turn_response_check.record_response_usage( + agent, + response, + messages=[], + api_call_count=1, + api_duration=0.0, + compression_attempts=0, + max_compression_attempts=3, + ) class _UsageRecordingAgent: """Hermes agent stand-in whose turns run scripted provider responses through Hermes' own usage accounting and context compressor, called the - way Hermes' turn loop calls it.""" + way Hermes' turn loop calls it. A callable step runs against the agent and + may raise; ``_COUNTERLESS_FAILURE`` ends the turn with a failure result.""" def __init__( self, - turns: list[list[dict[str, object] | Callable[[Any], None] | None]], + turns: list[list[dict[str, object] | Callable[[Any], None] | str | None]], api_mode: str = "chat_completions", ) -> None: from agent.context_compressor import ContextCompressor @@ -1669,22 +1652,14 @@ def close(self) -> None: pass def run_conversation(self, user_message: str) -> dict[str, object]: - from agent import turn_response_check - messages = [{"role": "user", "content": user_message}] for usage in self._turns.pop(0): + if usage == _COUNTERLESS_FAILURE: + return {"final_response": "", "failed": True, "messages": messages} if callable(usage): usage(self) continue - turn_response_check.record_response_usage( - self, - SimpleNamespace(usage=usage), - messages=messages, - api_call_count=1, - api_duration=0.0, - compression_attempts=0, - max_compression_attempts=3, - ) + _record_response(self, SimpleNamespace(usage=usage)) return { "final_response": "done", "completed": True, @@ -1882,41 +1857,152 @@ async def test_hermes_usage_reports_cache_counts_only_when_the_provider_sent_the assert with_cache_fields.usage.total_tokens == 135 -@pytest.mark.parametrize( - ("action", "counted"), [("break", False), ("fallthrough", True)] -) -async def test_hermes_refused_response_before_a_fallback_reply(action, counted): - # Hermes hands a content-filter refusal to handle_content_policy_refusal; - # with a fallback configured it returns "break" and the fallback model - # answers, so the refused call is billed but never recorded. - refusal = adapter._observed_response_handler( - "handle_content_policy_refusal", - lambda agent, response, *args, **kwargs: SimpleNamespace(action=action), +def _reported(prompt: int, completion: int) -> dict[str, int]: + return { + "prompt_tokens": prompt, + "completion_tokens": completion, + "total_tokens": prompt + completion, + } + + +async def test_hermes_usage_after_failed_or_counterless_turn_is_each_turns_own(): + def fail_turn(agent: Any) -> None: + raise RuntimeError("turn failed") + + agent = _UsageRecordingAgent( + [ + [_reported(100, 10)], + # Hermes' early failure results carry no session counters. + [_reported(40, 4), _COUNTERLESS_FAILURE], + [_reported(50, 5)], + [_reported(70, 7), fail_turn], + [_reported(60, 6)], + ] ) + runtime = _usage_recording_runtime(agent) + invocation = _usage_invocation() + + first = await runtime.invoke(*invocation) + counterless = await runtime.invoke(*invocation) + after_counterless = await runtime.invoke(*invocation) + with pytest.raises(RuntimeError): + await runtime.invoke(*invocation) + after_exception = await runtime.invoke(*invocation) + + usages = [ + (turn.usage.input_tokens, turn.usage.output_tokens, turn.usage.total_tokens) + for turn in (first, counterless, after_counterless, after_exception) + if turn.usage is not None + ] + # Each invocation reports its own tokens; the turn that raised never + # reaches the next invocation's usage. + assert usages == [(100, 10, 110), (40, 4, 44), (50, 5, 55), (60, 6, 66)] + assert counterless.status == AgentRunStatus.FAILED - def refused(agent: Any) -> None: - refused_usage = {"prompt_tokens": 40, "completion_tokens": 0} - refusal(agent, SimpleNamespace(usage=refused_usage), None) + +async def test_hermes_usage_after_session_counter_reset_is_unknown_for_that_turn(): + def reset_session(agent: Any) -> None: + for counter in ("input_tokens", "cache_read_tokens", "output_tokens"): + setattr(agent, f"session_{counter}", 0) agent = _UsageRecordingAgent( [ - [refused, {"prompt_tokens": 50, "completion_tokens": 5}], - [{"prompt_tokens": 60, "completion_tokens": 6}], + [_reported(100, 10)], + [reset_session, _reported(30, 3)], + [_reported(20, 2)], ] ) runtime = _usage_recording_runtime(agent) invocation = _usage_invocation() - fallback = await runtime.invoke(*invocation) - after = await runtime.invoke(*invocation) + await runtime.invoke(*invocation) + reset = await runtime.invoke(*invocation) + after_reset = await runtime.invoke(*invocation) - if counted: - assert fallback.usage is not None - assert fallback.usage.input_tokens == 50 - else: - assert fallback.usage is None - assert after.usage is not None - assert (after.usage.input_tokens, after.usage.output_tokens) == (60, 6) + assert reset.usage is None + assert after_reset.usage is not None + assert (after_reset.usage.input_tokens, after_reset.usage.output_tokens) == (20, 2) + + +@pytest.mark.parametrize( + ("handler", "action"), + [ + ("handle_content_policy_refusal", "break"), + ("recover_from_truncation", "continue"), + ("recover_from_truncation", "fallthrough"), + ("retry_invalid_response", "continue"), + ("retry_invalid_response", "fallthrough"), + ], +) +async def test_hermes_skipped_response_usage_is_counted_once(handler, action): + # check_api_response returns after these handlers before recording the + # response's usage, unless the handler falls through to + # record_response_usage, which then counts it. + def positional_handler(agent: Any, response: Any, *args: Any) -> Any: + return SimpleNamespace(action=action) + + def keyword_handler(agent: Any, *, response: Any, **kwargs: Any) -> Any: + return SimpleNamespace(action=action) + + # Hermes passes retry_invalid_response the response by keyword. + observed = adapter._observed_response_handler( + handler, + keyword_handler if handler == "retry_invalid_response" else positional_handler, + ) + + def skipped(agent: Any) -> None: + response = SimpleNamespace(usage=_reported(40, 1)) + if handler == "retry_invalid_response": + observed(agent, response=response, error_details=[]) + else: + observed(agent, response, None) + if action == "fallthrough": + _record_response(agent, response) + + agent = _UsageRecordingAgent([[skipped, _reported(50, 5)]]) + runtime = _usage_recording_runtime(agent) + + turn = await runtime.invoke(*_usage_invocation()) + + assert turn.usage is not None + assert (turn.usage.input_tokens, turn.usage.output_tokens) == (90, 6) + + +async def test_hermes_skipped_response_uses_the_route_that_served_it(): + # A refusal handler can switch the agent to a fallback provider before it + # returns; the refused response still has the first route's usage shape. + def refuse_and_fall_back(agent: Any, response: Any, *args: Any) -> Any: + agent.provider = "custom" + agent.api_mode = "chat_completions" + return SimpleNamespace(action="break") + + observed = adapter._observed_response_handler( + "handle_content_policy_refusal", refuse_and_fall_back + ) + + def refused(agent: Any) -> None: + anthropic_usage = { + "input_tokens": 100, + "cache_read_input_tokens": 30, + "cache_creation_input_tokens": 0, + "output_tokens": 0, + } + observed(agent, SimpleNamespace(usage=anthropic_usage), None) + + fallback_usage = { + **_reported(50, 5), + "prompt_tokens_details": {"cached_tokens": 0}, + } + agent = _UsageRecordingAgent( + [[refused, fallback_usage]], api_mode="anthropic_messages" + ) + runtime = _usage_recording_runtime(agent) + + turn = await runtime.invoke(*_usage_invocation()) + + assert turn.usage is not None + assert (turn.usage.input_tokens, turn.usage.cached_input_tokens) == (180, 30) + assert turn.usage.output_tokens == 5 async def test_hermes_stop_unregisters_the_agent_usage_observer(): @@ -2029,12 +2115,19 @@ async def test_runtime_allows_invoke_after_cancelled_worker_finishes(monkeypatch worker_started = threading.Event() worker_release = threading.Event() mock_agent = MagicMock() + mock_agent.session_input_tokens = 0 + mock_agent.session_cache_read_tokens = 0 + mock_agent.session_cache_write_tokens = 0 + mock_agent.session_output_tokens = 0 mock_session_db = MagicMock() calls = 0 def run_turn(**_kwargs): nonlocal calls calls += 1 + # Each turn bills calls * 100 input and calls * 10 output tokens. + mock_agent.session_input_tokens += calls * 100 + mock_agent.session_output_tokens += calls * 10 if calls == 1: worker_started.set() assert worker_release.wait(timeout=1) @@ -2044,10 +2137,6 @@ def run_turn(**_kwargs): "completed": True, "failed": False, "messages": [], - "input_tokens": calls * 100, - "cache_read_tokens": 0, - "cache_write_tokens": 0, - "output_tokens": calls * 10, }, "", ) @@ -2084,8 +2173,10 @@ def run_turn(**_kwargs): result = await runtime.invoke(*_invocation(invocation)) assert result.output["response"] == "turn-2" - # The cancelled turn left no baseline, so its tokens never reach this result. - assert result.usage is None + # The cancelled turn's tokens are in the counters before this turn starts, + # so this result reports only its own. + assert result.usage is not None + assert (result.usage.input_tokens, result.usage.output_tokens) == (200, 20) await runtime.stop() diff --git a/tests/e2e/test_hermes_e2e.py b/tests/e2e/test_hermes_e2e.py index 8df6ea084..74fb674cb 100644 --- a/tests/e2e/test_hermes_e2e.py +++ b/tests/e2e/test_hermes_e2e.py @@ -152,51 +152,96 @@ async def test_hermes_persistent_host_reports_usage_for_each_invocation( }, partial.to_mapping() +def _chat_usage(prompt: int, completion: int) -> dict[str, int]: + return { + "prompt_tokens": prompt, + "completion_tokens": completion, + "total_tokens": prompt + completion, + } + + @pytest.mark.usefixtures("mock_nvidia_api_key") @pytest.mark.parametrize( - "scenario", + ("replies", "expected_status", "expected_usage"), [ - # A reply cut off at max_tokens bills the whole budget. - { - "finish_reason": "length", - "usage": { - "prompt_tokens": 40, - "completion_tokens": 512, - "total_tokens": 552, - }, - }, + # A reply cut off at the 512-token default, then the continuation. + ( + [ + {"finish_reason": "length", "usage": _chat_usage(40, 512)}, + {"usage": _chat_usage(600, 20)}, + ], + "succeeded", + {"input_tokens": 640, "output_tokens": 532, "total_tokens": 1172}, + ), # A router timeout reported as a completion that billed the prompt, # then a successful retry. - { - "invalid_responses": 1, - "usage": {"prompt_tokens": 40, "completion_tokens": 0, "total_tokens": 40}, - }, + ( + [ + {"invalid": True, "usage": _chat_usage(40, 0)}, + {"usage": _chat_usage(50, 5)}, + ], + "succeeded", + {"input_tokens": 90, "output_tokens": 5, "total_tokens": 95}, + ), + # A content-filter refusal ends the turn without a fallback model. + ( + [{"finish_reason": "content_filter", "usage": _chat_usage(30, 0)}], + "failed", + {"input_tokens": 30, "output_tokens": 0, "total_tokens": 30}, + ), + # The truncated reply reports no completion count, so only output and + # the total stay unknown. + ( + [ + {"finish_reason": "length", "usage": {"prompt_tokens": 40}}, + {"usage": _chat_usage(60, 6)}, + ], + "succeeded", + {"input_tokens": 100}, + ), + ( + [ + {"finish_reason": "length", "usage": _chat_usage(0, 0)}, + {"usage": _chat_usage(0, 0)}, + ], + "succeeded", + {"input_tokens": 0, "output_tokens": 0, "total_tokens": 0}, + ), ], - ids=["length", "invalid_then_retry"], + ids=["truncated", "invalid", "refused", "partial", "zeros"], ) -async def test_hermes_usage_is_unknown_when_hermes_skips_a_response( +async def test_hermes_usage_counts_responses_hermes_skips( code_review_agent_dir: Path, api_server: str, - scenario: dict[str, object], + replies: list[dict[str, object]], + expected_status: str, + expected_usage: dict[str, int], ): os.environ["ADAPTER_PYTHON"] = sys.executable config = hermes_config() config.models["default"].base_url = f"{api_server}/v1" + # Leave Hermes room to continue or retry within its iteration budget. + config.runtime.max_turns = 5 async with await Fabric().start_runtime( config, base_dir=code_review_agent_dir ) as runtime: counted = await runtime.invoke(input="first") scenario_response = requests.post( - f"{api_server}/_scenario", json=scenario, timeout=5 + f"{api_server}/_scenario", json={"replies": replies}, timeout=5 ) scenario_response.raise_for_status() skipped = await runtime.invoke(input="second") assert counted.usage is not None, counted.to_mapping() - # Hermes returns on a truncated or invalid response before recording its - # usage, so the invocation's counters miss a billed call. - assert skipped.usage is None, skipped.to_mapping() + # Hermes returns on a truncated, invalid or refused response before adding + # its usage to the session counters; the adapter adds the reported counts. + assert skipped["status"] == expected_status, skipped.to_mapping() + assert skipped.to_mapping()["usage"] == { + **expected_usage, + "input_tokens_include_cache": True, + "metadata": {}, + }, skipped.to_mapping() @pytest.mark.usefixtures("mock_nvidia_api_key") From 718c42a9d7b4f768879410354f76f7c6f40b3f6b Mon Sep 17 00:00:00 2001 From: mnajafian-nv Date: Mon, 19 Oct 2026 22:47:52 -0700 Subject: [PATCH 19/68] fix: keep Pi usage known through rate limits and split compactions Under rate limits Pi's auto-retry is on by default, but the rejected request's error message carries Pi's zero-filled usage with no reported field, so one 429 made the whole invocation unknown. A call that ends in an error with no successful HTTP response and no raw stream event never reached generation, so it now counts as 0 tokens. A stream that failed after it started stays unknown. The first threshold compaction of a single-prompt run cuts inside the turn, and Pi reports only the combined usage of its two summary calls. Both calls already pass the adapter's stream wrapper, so the adapter now sums their own records when the calls it saw are the ones the preparation needs and add up to the combined usage. A tool usage report missing one count made every count unknown; each count now stays known on its own. Records are also keyed by the assistant message, because an extension's message_end handler can replace the message's usage object while Pi keeps the message. Signed-off-by: mnajafian-nv --- adapters/typescript/pi/src/pi-sdk.ts | 245 +++++++++++++++----- adapters/typescript/pi/test/pi-sdk.test.mjs | 149 ++++++++++-- 2 files changed, 318 insertions(+), 76 deletions(-) diff --git a/adapters/typescript/pi/src/pi-sdk.ts b/adapters/typescript/pi/src/pi-sdk.ts index 5294a2e56..dc6780caa 100644 --- a/adapters/typescript/pi/src/pi-sdk.ts +++ b/adapters/typescript/pi/src/pi-sdk.ts @@ -78,32 +78,44 @@ export interface ReportedCounts { output: boolean; } -/** What the provider reported for the model call that produced a usage object, if the adapter saw it. */ -export type ReportedCountsLookup = (usage: unknown) => ReportedCounts | undefined; +/** + * What the provider reported for the model call that produced an assistant + * message or a usage object, if the adapter saw that call. + */ +export type ReportedCountsLookup = (messageOrUsage: unknown) => ReportedCounts | undefined; const ALL_REPORTED: ReportedCounts = { input: true, cacheRead: true, output: true }; +/** Pi's four usage counts, each undefined unless it is a non-negative safe integer. */ +function piCounts(usage: unknown): Record<"input" | "cacheRead" | "cacheWrite" | "output", number | undefined> { + const fields = (typeof usage === "object" && usage !== null ? usage : {}) as Record; + const count = (name: string): number | undefined => { + const value = fields[name]; + return typeof value === "number" && Number.isSafeInteger(value) && value >= 0 ? value : undefined; + }; + return { input: count("input"), cacheRead: count("cacheRead"), cacheWrite: count("cacheWrite"), output: count("output") }; +} + /** * One call's counts from Pi's usage, with input counting its cache reads and - * writes, or undefined unless Pi's four counts are non-negative safe integers. + * writes, and which of them the usage holds. Input needs Pi's `input`, + * `cacheRead` and `cacheWrite`, cache reads need `cacheRead`, and output needs + * `output`; a missing field adds 0 to `counts` and leaves the counts that need + * it unreported. Undefined when the usage is not an object. */ -function callCounts(usage: unknown): Required | undefined { +function callCounts(usage: unknown): { counts: Required; present: ReportedCounts } | undefined { if (typeof usage !== "object" || usage === null) { return undefined; } - const counts = usage as Record; - const count = (name: string): number | undefined => { - const value = counts[name]; - return typeof value === "number" && Number.isSafeInteger(value) && value >= 0 ? value : undefined; + const { input, cacheRead, cacheWrite, output } = piCounts(usage); + return { + counts: { input: (input ?? 0) + (cacheRead ?? 0) + (cacheWrite ?? 0), cacheRead: cacheRead ?? 0, output: output ?? 0 }, + present: { + input: input !== undefined && cacheRead !== undefined && cacheWrite !== undefined, + cacheRead: cacheRead !== undefined, + output: output !== undefined, + }, }; - const input = count("input"); - const output = count("output"); - const cacheRead = count("cacheRead"); - const cacheWrite = count("cacheWrite"); - if (input === undefined || output === undefined || cacheRead === undefined || cacheWrite === undefined) { - return undefined; - } - return { input: input + cacheRead + cacheWrite, cacheRead, output }; } /** @@ -127,17 +139,18 @@ function sumCounts(total: PiTokenUsage, call: Required, reported: /** * Add one call's own usage report, such as a tool's, to a prompt's running - * total. A report without all four counts makes the whole total unknown. + * total. A count the report leaves out is unknown for the whole prompt, and + * the other counts stay known. */ export function addUsage(total: PiTokenUsage | undefined, usage: unknown): PiTokenUsage | undefined { const call = callCounts(usage); - return total === undefined || call === undefined ? undefined : sumCounts(total, call, ALL_REPORTED); + return total === undefined || call === undefined ? undefined : sumCounts(total, call.counts, call.present); } /** * Add one model call's usage to a prompt's running total. Pi's provider - * parsers fill a count the provider omitted with 0, so `reportedCounts` says - * which counts the provider actually reported, from the raw stream events the + * parsers fill a count the provider omitted with 0, so `reported` says which + * counts the provider actually reported, from the raw stream events the * session recorded for that call. A count the provider omitted is unknown for * the whole prompt, and so is every count of a call the session has no * record for, such as a cache warm-up. A reported zero stays a measured zero. @@ -145,39 +158,100 @@ export function addUsage(total: PiTokenUsage | undefined, usage: unknown): PiTok function addModelCallUsage( total: PiTokenUsage | undefined, usage: unknown, - reportedCounts: ReportedCountsLookup, + reported: ReportedCounts | undefined, ): PiTokenUsage | undefined { const call = callCounts(usage); - const reported = reportedCounts(usage); if (total === undefined || call === undefined || reported === undefined) { return undefined; } - return sumCounts(total, call, reported); + return sumCounts(total, call.counts, { + input: reported.input && call.present.input, + cacheRead: reported.cacheRead && call.present.cacheRead, + output: reported.output && call.present.output, + }); +} + +/** + * The summary calls Pi's `compact` makes for a compaction preparation: one for + * the history, and for a split turn one for the turn prefix, without the + * history call when there is no history to summarize. Undefined without a + * preparation. + */ +export function summaryCallCount(preparation: unknown): number | undefined { + if (typeof preparation !== "object" || preparation === null) { + return undefined; + } + const { isSplitTurn, messagesToSummarize, turnPrefixMessages } = preparation as Record; + if (isSplitTurn !== true || !Array.isArray(turnPrefixMessages) || turnPrefixMessages.length === 0) { + return 1; + } + return Array.isArray(messagesToSummarize) && messagesToSummarize.length > 0 ? 2 : 1; +} + +/** The model calls that ran between a compaction's preparation and its end. */ +export interface CompactionCalls { + /** How many summary calls the preparation needs, from `summaryCallCount`. */ + expected: number; + /** The usage objects of the model calls that ended since the preparation. */ + calls: readonly unknown[]; +} + +/** + * Add a split-turn compaction's summary calls one by one. Pi reports only + * their combined usage, so the calls count only when the session saw exactly + * the summary calls the preparation needs and their usage adds up to the + * combined usage; otherwise one call could be missing and the prompt's tokens + * are unknown. + */ +function addSummaryCallsUsage( + total: PiTokenUsage | undefined, + combined: unknown, + compaction: CompactionCalls | undefined, + reportedCounts: ReportedCountsLookup, +): PiTokenUsage | undefined { + if (compaction === undefined || compaction.calls.length !== compaction.expected) { + return undefined; + } + const expected = piCounts(combined); + const names = ["input", "cacheRead", "cacheWrite", "output"] as const; + const addsUp = names.every((name) => { + const sum = compaction.calls.reduce((partial, call) => { + const count = piCounts(call)[name]; + return partial === undefined || count === undefined ? undefined : partial + count; + }, 0); + return sum !== undefined && sum === expected[name]; + }); + if (!addsUp) { + return undefined; + } + return compaction.calls.reduce( + (sum, call) => addModelCallUsage(sum, call, reportedCounts(call)), + total, + ); } /** * Add the tokens one session event recorded to a prompt's running total, as * Pi's own usage totals count them: an assistant message, a tool result that - * reports usage, a finished compaction's summary call, and a usage entry such + * reports usage, a finished compaction's summary calls, and a usage entry such * as a cache warm-up. A model call adds only the counts its provider - * reported, as `reportedCounts` records them. The total becomes unknown when a - * model call reports no count, when an assistant message or a compaction was - * aborted, or when a compaction summary call was retried: Pi keeps no usage for a call it - * stopped or a summary attempt that failed, and either may still have been - * billed. It also becomes unknown when Pi prepares a split-turn compaction - * (a `session_before_compact` event whose preparation has `isSplitTurn`): - * Pi summarizes the history and the turn prefix in separate calls and reports - * only their combined usage, so one unreported call cannot be detected. + * reported, as `reportedCounts` records them by assistant message or usage + * object. The total becomes unknown when a model call reports no count, when + * an assistant message or a compaction was aborted, or when a compaction + * summary call was retried: Pi keeps no usage for a call it stopped or a + * summary attempt that failed, and either may still have been billed. A + * split-turn compaction's combined usage counts through its summary calls, + * which `compaction` lists; see `addSummaryCallsUsage`. */ export function addEventUsage( total: PiTokenUsage | undefined, event: unknown, reportedCounts: ReportedCountsLookup, + compaction?: CompactionCalls, ): PiTokenUsage | undefined { const fields = (typeof event === "object" && event !== null ? event : {}) as { type?: unknown; message?: { role?: unknown; stopReason?: unknown; usage?: unknown }; - preparation?: { isSplitTurn?: unknown }; aborted?: unknown; result?: { usage?: unknown }; entry?: { type?: unknown; usage?: unknown }; @@ -185,24 +259,32 @@ export function addEventUsage( switch (fields.type) { case "message_end": if (fields.message?.role === "assistant") { + // An extension's message_end handler can replace the message's usage + // object, but Pi keeps the message object itself. + const reported = reportedCounts(fields.message) ?? reportedCounts(fields.message.usage); return fields.message.stopReason === "aborted" ? undefined - : addModelCallUsage(total, fields.message.usage, reportedCounts); + : addModelCallUsage(total, fields.message.usage, reported); } if (fields.message?.role === "toolResult" && fields.message.usage !== undefined) { return addUsage(total, fields.message.usage); } return total; - case "compaction_end": - return fields.aborted === false && fields.result !== undefined - ? addModelCallUsage(total, fields.result.usage, reportedCounts) - : undefined; - case "session_before_compact": - return fields.preparation?.isSplitTurn === false ? total : undefined; + case "compaction_end": { + if (fields.aborted !== false || fields.result === undefined) { + return undefined; + } + const reported = reportedCounts(fields.result.usage); + return reported === undefined + ? addSummaryCallsUsage(total, fields.result.usage, compaction, reportedCounts) + : addModelCallUsage(total, fields.result.usage, reported); + } case "summarization_retry_scheduled": return undefined; case "entry_appended": - return fields.entry?.type === "usage" ? addModelCallUsage(total, fields.entry.usage, reportedCounts) : total; + return fields.entry?.type === "usage" + ? addModelCallUsage(total, fields.entry.usage, reportedCounts(fields.entry.usage)) + : total; default: return total; } @@ -363,26 +445,61 @@ export function reportedCountsFrom(api: string, seen: ReadonlySet): Repo }; } +/** + * Whether a model call ended in an error before the provider accepted the + * request: no successful HTTP response and no raw stream event, with Pi's + * zero-filled usage. Pi's protocols call `onResponse` only after the request + * succeeds and raise a rejection such as a rate limit (429), an overload or a + * server error before the first raw event (`@earendil-works/pi-ai` 1.0.3, + * `api/*.js`), so such a call generated nothing and reports 0 tokens. + */ +function rejectedBeforeGeneration(message: { stopReason?: unknown; usage?: unknown }, accepted: boolean, streamed: boolean): boolean { + const counts = piCounts(message.usage); + return ( + message.stopReason === "error" && + !accepted && + !streamed && + counts.input === 0 && + counts.cacheRead === 0 && + counts.cacheWrite === 0 && + counts.output === 0 + ); +} + /** * Record which counts each model call's provider reported, keyed by the call's - * final usage object. Pi hands every raw provider stream event, with the model - * that served it, to `onProviderStreamEvent` before parsing it. The agent loop - * and compaction summaries both call the agent's stream function, so wrapping - * it covers both, and Pi passes the final usage object unchanged to - * `message_end` and `compaction_end`. The record is stored before Pi emits - * either event. Pi's cache warmer replays a request with the same options, so - * a warm-up's raw events reach this observer too, but the warm-up runs outside - * the agent's stream function: its usage entry has no record and stays unknown. - * The record is copied when the call ends so a later warm-up cannot change it. + * final assistant message and its usage object, and pass each ended call's + * usage object to `onCallEnd`. Pi hands every raw provider stream event, with + * the model that served it, to `onProviderStreamEvent` before parsing it. The + * agent loop and compaction summaries both call the agent's stream function, + * so wrapping it covers both, and Pi passes the final message to + * `message_end` and its usage object unchanged to `compaction_end`. The record + * is stored before Pi emits either event. A request the provider rejected + * before generating reports all its counts as 0. Pi's cache warmer replays a + * request with the same options, so a warm-up's raw events reach this + * observer too, but the warm-up runs outside the agent's stream function: its + * usage entry has no record and stays unknown. The record is copied when the + * call ends so a later warm-up cannot change it. */ -function recordReportedCounts(agent: AgentSession["agent"], records: WeakMap): void { +function recordReportedCounts( + agent: AgentSession["agent"], + records: WeakMap, + onCallEnd: (usage: object) => void, +): void { const streamFunction = agent.streamFunction; agent.streamFunction = async (model, context, options) => { const seen = new Set(); let api: string = model.api; + let accepted = false; + let streamed = false; const response = await streamFunction(model, context, { ...options, + onResponse: async (providerResponse, responseModel) => { + accepted ||= providerResponse.status >= 200 && providerResponse.status < 300; + await options?.onResponse?.(providerResponse, responseModel); + }, onProviderStreamEvent: async (data, eventModel) => { + streamed = true; api = eventModel?.api ?? model.api; noteReportedFields(seen, api, data); await options?.onProviderStreamEvent?.(data, eventModel); @@ -391,7 +508,12 @@ function recordReportedCounts(agent: AgentSession["agent"], records: WeakMap { if (typeof message.usage === "object" && message.usage !== null) { - records.set(message.usage, reportedCountsFrom(api, seen)); + const counts = rejectedBeforeGeneration(message, accepted, streamed) + ? ALL_REPORTED + : reportedCountsFrom(api, seen); + records.set(message, counts); + records.set(message.usage, counts); + onCallEnd(message.usage); } }, () => undefined, @@ -1216,6 +1338,8 @@ class PiSdkSessionHandle implements PiSessionHandle { private readonly runtimeDir: string; private readonly compactionListeners: Set<(event: unknown) => void>; private readonly reportedCounts = new WeakMap(); + // The model calls that ended since the current compaction was prepared. + private compaction: { expected: number; calls: object[] } | undefined; constructor( session: AgentSession, @@ -1229,7 +1353,9 @@ class PiSdkSessionHandle implements PiSessionHandle { this.runtimeDir = runtimeDir; this.compactionListeners = compactionListeners; this.relay = relay; - recordReportedCounts(session.agent, this.reportedCounts); + recordReportedCounts(session.agent, this.reportedCounts, (usage) => { + this.compaction?.calls.push(usage); + }); this.unsubscribeTurnCounter = this.session.subscribe((event) => { if (event.type === "turn_start") { this.cumulativeTurnCount += 1; @@ -1248,8 +1374,8 @@ class PiSdkSessionHandle implements PiSessionHandle { | { role: "assistant"; content: unknown; stopReason: string; errorMessage?: string } | undefined; let usage: PiTokenUsage | undefined = NO_TOKENS; - const reportedCounts: ReportedCountsLookup = (callUsage) => - typeof callUsage === "object" && callUsage !== null ? this.reportedCounts.get(callUsage) : undefined; + const reportedCounts: ReportedCountsLookup = (key) => + typeof key === "object" && key !== null ? this.reportedCounts.get(key) : undefined; const unsubscribe = this.session.subscribe((event) => { if (event.type === "turn_start") { turnStarted = true; @@ -1257,10 +1383,14 @@ class PiSdkSessionHandle implements PiSessionHandle { if (event.type === "message_end" && event.message.role === "assistant") { finalAssistant = event.message; } - usage = addEventUsage(usage, event, reportedCounts); + usage = addEventUsage(usage, event, reportedCounts, this.compaction); + if (event.type === "compaction_end") { + this.compaction = undefined; + } }); const onCompactionPrepared = (event: unknown): void => { - usage = addEventUsage(usage, event, reportedCounts); + const expected = summaryCallCount((event as { preparation?: unknown }).preparation); + this.compaction = expected === undefined ? undefined : { expected, calls: [] }; }; this.compactionListeners.add(onCompactionPrepared); try { @@ -1273,6 +1403,7 @@ class PiSdkSessionHandle implements PiSessionHandle { }); } finally { this.compactionListeners.delete(onCompactionPrepared); + this.compaction = undefined; unsubscribe(); } return { diff --git a/adapters/typescript/pi/test/pi-sdk.test.mjs b/adapters/typescript/pi/test/pi-sdk.test.mjs index 35723f75f..db75fd313 100644 --- a/adapters/typescript/pi/test/pi-sdk.test.mjs +++ b/adapters/typescript/pi/test/pi-sdk.test.mjs @@ -17,6 +17,7 @@ import { PiSdkSessionFactory, reportedCountsFrom, resolveCustomTools, + summaryCallCount, selectPiMcpServers, withCustomBaseUrl, } from "../dist/pi-sdk.js"; @@ -1452,15 +1453,37 @@ const unseen = () => undefined; // A prompt's running total before any call, with input counting the cache. const NO_TOKENS = { input: 0, cacheRead: 0, output: 0 }; -test("sums each assistant message's usage and drops an incomplete one", () => { +test("sums each call's own usage report", () => { const first = addUsage(NO_TOKENS, { input: 10, output: 2, cacheRead: 5, cacheWrite: 1, totalTokens: 18 }); const second = addUsage(first, { input: 3, output: 4, cacheRead: 0, cacheWrite: 0, reasoning: 2 }); assert.deepEqual(second, { input: 19, cacheRead: 5, output: 6 }); - assert.equal(addUsage(second, { input: 1, output: 1 }), undefined); assert.equal(addUsage(undefined, { input: 1, output: 1, cacheRead: 0, cacheWrite: 0 }), undefined); }); +// A tool result that reports usage, as a custom tool can. +const toolResult = (usage) => ({ type: "message_end", message: { role: "toolResult", usage } }); + +test("a tool usage report without a cache count keeps its output known", () => { + const total = addEventUsage(NO_TOKENS, toolResult({ input: 10, cacheRead: 2, output: 3 }), allReported); + + // Input counts cache writes, so without cacheWrite the input, the cache + // reads inside it and the total are unknown. + assert.deepEqual(total, { output: 3 }); +}); + +test("a tool usage report with only an output count adds that count", () => { + const total = addEventUsage({ input: 40, cacheRead: 0, output: 1 }, toolResult({ output: 4 }), allReported); + + assert.deepEqual(total, { output: 5 }); +}); + +test("a tool result without usage leaves the prompt's tokens unchanged", () => { + const total = { input: 40, cacheRead: 0, output: 1 }; + + assert.deepEqual(addEventUsage(total, { type: "message_end", message: { role: "toolResult" } }, unseen), total); +}); + test("counts the tokens Pi bills outside assistant messages", () => { const usage = (tokens) => ({ input: tokens, output: 1, cacheRead: 0, cacheWrite: 0 }); const events = [ @@ -1596,12 +1619,30 @@ test("reads which counts each wire protocol's raw usage reports", () => { }); }); -test("a split-turn compaction leaves the prompt's tokens unknown", () => { - const prepared = (preparation) => ({ type: "session_before_compact", preparation }); - - assert.deepEqual(addEventUsage(NO_TOKENS, prepared({ isSplitTurn: false }), allReported), NO_TOKENS); - assert.equal(addEventUsage(NO_TOKENS, prepared({ isSplitTurn: true }), allReported), undefined); - assert.equal(addEventUsage(NO_TOKENS, prepared(undefined), allReported), undefined); +test("a split-turn compaction counts its summary calls only when they add up", () => { + const call = (tokens) => ({ input: tokens, output: 1, cacheRead: 0, cacheWrite: 0 }); + const history = call(20); + const prefix = call(30); + const combined = { input: 50, output: 2, cacheRead: 0, cacheWrite: 0 }; + const ended = { type: "compaction_end", aborted: false, result: { summary: "s", usage: combined } }; + // Only the summary calls have records; Pi reports their combined usage. + const recorded = (key) => (key === history || key === prefix ? allReported() : undefined); + const split = { isSplitTurn: true, messagesToSummarize: [{}], turnPrefixMessages: [{}] }; + + assert.equal(summaryCallCount(split), 2); + assert.equal(summaryCallCount({ ...split, messagesToSummarize: [] }), 1); + assert.equal(summaryCallCount({ isSplitTurn: false, messagesToSummarize: [{}], turnPrefixMessages: [] }), 1); + assert.equal(summaryCallCount(undefined), undefined); + assert.deepEqual(addEventUsage(NO_TOKENS, ended, recorded, { expected: 2, calls: [history, prefix] }), { + input: 50, + cacheRead: 0, + output: 2, + }); + // A summary call the session did not see, or calls that do not add up to the + // combined usage, leave the prompt's tokens unknown. + assert.equal(addEventUsage(NO_TOKENS, ended, recorded, { expected: 2, calls: [prefix] }), undefined); + assert.equal(addEventUsage(NO_TOKENS, ended, recorded, { expected: 2, calls: [history, call(29)] }), undefined); + assert.equal(addEventUsage(NO_TOKENS, ended, recorded), undefined); }); test("a retried compaction summary leaves the prompt's tokens unknown", () => { @@ -1622,22 +1663,39 @@ test("a retried compaction summary leaves the prompt's tokens unknown", () => { /** * Run prompts through the real Pi SDK against a local chat-completions endpoint * that answers each request from `replies` in order. A reply's `usage` is sent - * as the stream's final usage chunk; a reply without one sends no usage. - * `modelExtensions` overrides the catalog entry, such as its context window. + * as the stream's final usage chunk; a reply without one sends no usage. A + * reply with `status` rejects the request with that HTTP status, and one with + * `dropStream` closes the connection after the first chunk. + * `modelExtensions` overrides the catalog entry, such as its context window, + * and `extension` is the source of a Pi extension the session loads. */ -async function promptsAgainstScriptedEndpoint(replies, prompts, modelExtensions = undefined) { +async function promptsAgainstScriptedEndpoint(replies, prompts, modelExtensions = undefined, extension = undefined) { const workspace = await realpath(await mkdtemp(join(tmpdir(), "fabric-pi-usage-"))); await writeFile(join(workspace, "note.txt"), "nonce\n", "utf8"); + if (extension !== undefined) { + await writeFile(join(workspace, "usage-extension.js"), extension, "utf8"); + } let served = 0; const provider = createServer(async (request, response) => { for await (const _chunk of request); const reply = replies[served++]; + if (reply.status !== undefined) { + response.writeHead(reply.status, { "content-type": "application/json" }); + response.end(JSON.stringify({ error: { message: "rate limited", type: "rate_limit_error" } })); + return; + } const common = { id: `chatcmpl-usage-${served}`, object: "chat.completion.chunk", created: 0, model: "openai/gpt-oss-20b" }; const chunk = (fields) => `data: ${JSON.stringify({ ...common, ...fields })}\n\n`; const delta = reply.readNote ? { role: "assistant", tool_calls: [{ index: 0, id: `call-${served}`, type: "function", function: { name: "read", arguments: JSON.stringify({ path: "note.txt" }) } }] } : { role: "assistant", content: reply.content }; response.writeHead(200, { "content-type": "text/event-stream" }); + if (reply.dropStream) { + response.write(chunk({ choices: [{ index: 0, delta, finish_reason: null }] })); + // Let the client read the first chunk before the connection drops. + setTimeout(() => response.destroy(), 100); + return; + } response.end( chunk({ choices: [{ index: 0, delta, finish_reason: null }] }) + chunk({ choices: [{ index: 0, delta: {}, finish_reason: reply.readNote ? "tool_calls" : "stop" }] }) + @@ -1652,6 +1710,7 @@ async function promptsAgainstScriptedEndpoint(replies, prompts, modelExtensions agentName: "pi-usage-test", baseDir: workspace, config: { + ...(extension === undefined ? {} : { harness: { settings: { extensions: ["usage-extension.js"] } } }), models: { default: { api_key_env: "TEST_API_KEY", @@ -1765,26 +1824,78 @@ test( }, ); -test( - "leaves a prompt's tokens unknown after a split-turn compaction", - { skip: supportsPi ? false : "Pi 1.0 requires Node 22.19 or newer" }, - async () => { +for (const [name, historyUsage, expected] of [ + ["counts both summary calls of a split-turn compaction", { prompt_tokens: 20, completion_tokens: 1 }, { input: 90_240, output: 12 }], + ["leaves input unknown when a split-turn summary call omits it", { completion_tokens: 1 }, { output: 12 }], +]) { + test(name, { skip: supportsPi ? false : "Pi 1.0 requires Node 22.19 or newer" }, async () => { // The second prompt's final reply reports a context above the compaction // threshold and is longer than Pi's 20,000 recent tokens, so Pi cuts inside - // that turn and summarizes the history and the turn prefix in two calls. - // The history summary reports no usage; the turn-prefix summary does. + // that turn and summarizes the history and the turn prefix in two calls, + // then reports only their combined usage. const [, compacted] = await promptsAgainstScriptedEndpoint( [ { content: "first", usage: { prompt_tokens: 100, completion_tokens: 5 } }, { readNote: true, usage: { prompt_tokens: 200, completion_tokens: 5 } }, { content: "x".repeat(100_000), usage: { prompt_tokens: 90_000, completion_tokens: 5 } }, - { content: "history summary" }, + { content: "history summary", usage: historyUsage }, { content: "turn prefix summary", usage: { prompt_tokens: 20, completion_tokens: 1 } }, ], ["one", "read the note"], { context_window: 100_000, max_tokens: 1_000 }, ); - assert.equal(compacted.usage, undefined); + assert.deepEqual(compacted.usage, expected); + }); +} + +test( + "counts a rate-limited request as 0 tokens when Pi's retry succeeds", + { skip: supportsPi ? false : "Pi 1.0 requires Node 22.19 or newer" }, + async () => { + const [retried] = await promptsAgainstScriptedEndpoint( + [{ status: 429 }, { content: "after retry", usage: { prompt_tokens: 50, completion_tokens: 5 } }], + ["one"], + ); + + assert.equal(retried.text, "after retry"); + assert.deepEqual(retried.usage, { input: 50, output: 5 }); + }, +); + +test( + "leaves a prompt's tokens unknown when a stream fails after it started", + { skip: supportsPi ? false : "Pi 1.0 requires Node 22.19 or newer" }, + async () => { + // The provider accepted the request and streamed, so it may have billed + // tokens that the dropped stream never reported. + const [dropped] = await promptsAgainstScriptedEndpoint( + [{ content: "partial", dropStream: true }, { content: "after retry", usage: { prompt_tokens: 50, completion_tokens: 5 } }], + ["one"], + ); + + assert.equal(dropped.text, "after retry"); + assert.equal(dropped.usage, undefined); + }, +); + +test( + "keeps usage known when an extension replaces an assistant message's usage object", + { skip: supportsPi ? false : "Pi 1.0 requires Node 22.19 or newer" }, + async () => { + const extension = `export default function (pi) { + pi.on("message_end", (event) => + event.message.role === "assistant" ? { message: { ...event.message, usage: { ...event.message.usage } } } : undefined, + ); +} +`; + const [replaced] = await promptsAgainstScriptedEndpoint( + [{ content: "done", usage: { prompt_tokens: 50, completion_tokens: 5 } }], + ["one"], + undefined, + extension, + ); + + assert.deepEqual(replaced.usage, { input: 50, output: 5 }); }, ); From 4656ff69d9ef214aba000c46c4eb23349ba3b5bf Mon Sep 17 00:00:00 2001 From: mnajafian-nv Date: Mon, 19 Oct 2026 23:07:28 -0700 Subject: [PATCH 20/68] fix: leave Hermes counts above the usage bound unknown AgentUsage holds counts up to 2^64 - 1. Each session counter was checked against that bound, but the input sum, the total and the tokens added from skipped responses were not, so an out-of-range provider count made invoke raise a ContractValidationError instead of returning a result. A count that does not fit is now unknown, like any other count the adapter cannot trust, and a raw report above the bound counts as omitted. Signed-off-by: mnajafian-nv --- .../nemo_fabric_adapters/hermes/adapter.py | 38 +++++++++++++------ tests/adapters/test_hermes_adapter.py | 14 +++++++ 2 files changed, 41 insertions(+), 11 deletions(-) diff --git a/adapters/python/hermes/src/nemo_fabric_adapters/hermes/adapter.py b/adapters/python/hermes/src/nemo_fabric_adapters/hermes/adapter.py index 238446418..9b064b656 100755 --- a/adapters/python/hermes/src/nemo_fabric_adapters/hermes/adapter.py +++ b/adapters/python/hermes/src/nemo_fabric_adapters/hermes/adapter.py @@ -867,7 +867,16 @@ def _is_reported_count(usage: Any, path: str) -> bool: value = ( value.get(field) if isinstance(value, dict) else getattr(value, field, None) ) - return isinstance(value, int) and not isinstance(value, bool) and value >= 0 + return _is_count(value) + + +# The largest count AgentUsage accepts. +_MAX_COUNT = (1 << 64) - 1 + + +def _is_count(value: Any) -> bool: + """Return whether ``value`` is a token count AgentUsage can hold.""" + return isinstance(value, int) and not isinstance(value, bool) and 0 <= value <= _MAX_COUNT def _session_counters(agent: Any) -> dict[str, int]: @@ -879,9 +888,7 @@ def _session_counters(agent: Any) -> dict[str, int]: return { key: value for key in _SESSION_TOKEN_KEYS - if isinstance((value := getattr(agent, f"session_{key}", None)), int) - and not isinstance(value, bool) - and 0 <= value <= (1 << 64) - 1 + if _is_count(value := getattr(agent, f"session_{key}", None)) } @@ -902,7 +909,8 @@ def _invocation_usage( usage never holds a whole-session total. ``unreported`` names the counts ("input", "cached", "output") that a provider response of this invocation left out and Hermes counted as 0. Those counts are unknown, and so is the - total when input or output is. + total when input or output is. A count above what AgentUsage holds is + unknown too. """ skipped = skipped or {} change = { @@ -915,18 +923,26 @@ def _invocation_usage( key in change for key in ("input_tokens", "cache_read_tokens", "cache_write_tokens") ): - fields["input_tokens"] = ( + input_tokens = ( change["input_tokens"] + change["cache_read_tokens"] + change["cache_write_tokens"] ) - if "cached" not in unreported: - fields["cached_input_tokens"] = change["cache_read_tokens"] - fields["input_tokens_include_cache"] = True - if "output" not in unreported and "output_tokens" in change: + if _is_count(input_tokens): + fields["input_tokens"] = input_tokens + if "cached" not in unreported: + fields["cached_input_tokens"] = change["cache_read_tokens"] + fields["input_tokens_include_cache"] = True + if ( + "output" not in unreported + and "output_tokens" in change + and _is_count(change["output_tokens"]) + ): fields["output_tokens"] = change["output_tokens"] if "input_tokens" in fields and "output_tokens" in fields: - fields["total_tokens"] = fields["input_tokens"] + fields["output_tokens"] + total_tokens = fields["input_tokens"] + fields["output_tokens"] + if _is_count(total_tokens): + fields["total_tokens"] = total_tokens return AgentUsage(**fields) if fields else None diff --git a/tests/adapters/test_hermes_adapter.py b/tests/adapters/test_hermes_adapter.py index fc4708c79..80d908bdc 100644 --- a/tests/adapters/test_hermes_adapter.py +++ b/tests/adapters/test_hermes_adapter.py @@ -1576,6 +1576,20 @@ def test_invocation_usage_with_lower_or_missing_counters_leaves_them_unknown(): assert no_baseline is None +def test_invocation_usage_above_the_contract_bound_is_unknown(): + zero = dict.fromkeys(adapter._SESSION_TOKEN_KEYS, 0) + near_max = {**zero, "input_tokens": (1 << 64) - 2, "output_tokens": 3} + + usage = adapter._invocation_usage( + near_max, zero, skipped={"input_tokens": 5, "output_tokens": 1} + ) + + assert usage is not None + assert usage.input_tokens is None + assert usage.total_tokens is None + assert usage.output_tokens == 4 + + def test_session_counters_keep_only_valid_counts(): agent = SimpleNamespace( session_input_tokens=7, From 75c29a3f12c9e4fd169b2100d26725da91867018 Mon Sep 17 00:00:00 2001 From: mnajafian-nv Date: Mon, 19 Oct 2026 23:08:29 -0700 Subject: [PATCH 21/68] fix: apply the Pi rejected-request rule only to built-in protocols The rule that a call with no successful HTTP response and no raw stream event was rejected relies on Pi's built-in protocols calling onResponse and onProviderStreamEvent. A custom protocol may call neither, so any of its failed calls looked rejected and counted as 0 tokens even after the provider generated output. Such calls now keep only the counts their raw events reported. Signed-off-by: mnajafian-nv --- adapters/typescript/pi/src/pi-sdk.ts | 34 +++++++++++++-------- adapters/typescript/pi/test/pi-sdk.test.mjs | 18 +++++++++++ 2 files changed, 39 insertions(+), 13 deletions(-) diff --git a/adapters/typescript/pi/src/pi-sdk.ts b/adapters/typescript/pi/src/pi-sdk.ts index dc6780caa..a9347494b 100644 --- a/adapters/typescript/pi/src/pi-sdk.ts +++ b/adapters/typescript/pi/src/pi-sdk.ts @@ -446,24 +446,34 @@ export function reportedCountsFrom(api: string, seen: ReadonlySet): Repo } /** - * Whether a model call ended in an error before the provider accepted the - * request: no successful HTTP response and no raw stream event, with Pi's - * zero-filled usage. Pi's protocols call `onResponse` only after the request - * succeeds and raise a rejection such as a rate limit (429), an overload or a - * server error before the first raw event (`@earendil-works/pi-ai` 1.0.3, - * `api/*.js`), so such a call generated nothing and reports 0 tokens. + * Which counts an ended model call's provider reported. A call that ended in + * an error before the provider accepted the request, with no successful HTTP + * response, no raw stream event and Pi's zero-filled usage, reports all its + * counts as 0. Pi's built-in protocols call `onResponse` only after the + * request succeeds and raise a rejection such as a rate limit (429), an + * overload or a server error before the first raw event + * (`@earendil-works/pi-ai` 1.0.3, `api/*.js`), so such a call generated + * nothing. A protocol outside those may never call either hook, so its calls + * keep only the counts their raw events reported. */ -function rejectedBeforeGeneration(message: { stopReason?: unknown; usage?: unknown }, accepted: boolean, streamed: boolean): boolean { +export function endedCallCounts( + api: string, + seen: ReadonlySet, + message: { stopReason?: unknown; usage?: unknown }, + accepted: boolean, + streamed: boolean, +): ReportedCounts { const counts = piCounts(message.usage); - return ( + const rejected = + REPORTED_FIELDS[api] !== undefined && message.stopReason === "error" && !accepted && !streamed && counts.input === 0 && counts.cacheRead === 0 && counts.cacheWrite === 0 && - counts.output === 0 - ); + counts.output === 0; + return rejected ? ALL_REPORTED : reportedCountsFrom(api, seen); } /** @@ -508,9 +518,7 @@ function recordReportedCounts( void response.result().then( (message) => { if (typeof message.usage === "object" && message.usage !== null) { - const counts = rejectedBeforeGeneration(message, accepted, streamed) - ? ALL_REPORTED - : reportedCountsFrom(api, seen); + const counts = endedCallCounts(api, seen, message, accepted, streamed); records.set(message, counts); records.set(message.usage, counts); onCallEnd(message.usage); diff --git a/adapters/typescript/pi/test/pi-sdk.test.mjs b/adapters/typescript/pi/test/pi-sdk.test.mjs index db75fd313..980d5e981 100644 --- a/adapters/typescript/pi/test/pi-sdk.test.mjs +++ b/adapters/typescript/pi/test/pi-sdk.test.mjs @@ -13,6 +13,7 @@ import { addUsage, noteReportedFields, buildCatalogModel, + endedCallCounts, modelAwareCompactionReserveTokens, PiSdkSessionFactory, reportedCountsFrom, @@ -1660,6 +1661,23 @@ test("a retried compaction summary leaves the prompt's tokens unknown", () => { ); }); +test("counts a rejected request as 0 tokens only for a built-in protocol", () => { + const rejected = { stopReason: "error", usage: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 } }; + const none = { input: false, cacheRead: false, output: false }; + + assert.deepEqual(endedCallCounts("openai-completions", new Set(), rejected, false, false), allReported()); + // A custom protocol may never report its HTTP response or raw events. + assert.deepEqual(endedCallCounts("custom-api", new Set(), rejected, false, false), none); + // An accepted or streamed request, or one with counts, may have billed tokens. + assert.deepEqual(endedCallCounts("openai-completions", new Set(), rejected, true, false), none); + assert.deepEqual(endedCallCounts("openai-completions", new Set(), rejected, false, true), none); + assert.deepEqual( + endedCallCounts("openai-completions", new Set(), { ...rejected, usage: { ...rejected.usage, input: 3 } }, false, false), + none, + ); + assert.deepEqual(endedCallCounts("openai-completions", new Set(), { ...rejected, stopReason: "aborted" }, false, false), none); +}); + /** * Run prompts through the real Pi SDK against a local chat-completions endpoint * that answers each request from `replies` in order. A reply's `usage` is sent From 744906c6f17df177eed49cab8d95a9c2bf30e8ac Mon Sep 17 00:00:00 2001 From: mnajafian-nv Date: Mon, 19 Oct 2026 23:08:49 -0700 Subject: [PATCH 22/68] docs: describe the corrected Hermes and Pi usage rules The Hermes and Pi READMEs described the old behavior: no usage after a failed or counterless Hermes invocation or after a skipped response, and no Pi usage after a rate limit, a split-turn compaction or a partial tool report. They now state the rules the adapters apply, and the Hermes README names the iteration-limit summary request that Hermes leaves out of its counters. Signed-off-by: mnajafian-nv --- adapters/python/hermes/README.md | 35 ++++++++++++++++-------------- adapters/typescript/pi/README.md | 37 +++++++++++++++++++++++--------- 2 files changed, 46 insertions(+), 26 deletions(-) diff --git a/adapters/python/hermes/README.md b/adapters/python/hermes/README.md index 588e1b1c3..7428917e9 100644 --- a/adapters/python/hermes/README.md +++ b/adapters/python/hermes/README.md @@ -115,24 +115,27 @@ runtime. ## Token Usage Hermes counts its main model's tokens per session, and the adapter keeps one -session per runtime, so each result's `usage` is the change since the previous -invocation: `input_tokens` with cache reads and cache writes included, -`cached_input_tokens`, `output_tokens`, and `total_tokens`. Hermes records -auxiliary calls, such as context-compression summaries, separately from these -counters, so they are not included. A counter that Hermes does not report, or -that went down, is left out for that invocation. After an invocation that -raised, was cancelled, or timed out, or whose result held no counters, the -adapter has no baseline. The next invocation then reports no token counters -rather than a carried-over or whole-session total, and the one after it reports -its own change again. +session per runtime. The adapter reads the agent's session counters right +before each turn runs and again after it, so each result's `usage` is that +turn's change: `input_tokens` with cache reads and cache writes included, +`cached_input_tokens`, `output_tokens`, and `total_tokens`. A failure result +reports its own tokens too, and an invocation that raised or was cancelled +never adds its tokens to the next one. A counter that is missing, that is not +a non-negative integer below 2^64, or that went down because Hermes reset its +session, is left out for that invocation. Hermes records auxiliary calls, such as +context-compression summaries and session titles, separately from these +counters, and it does not count the summary request it makes when a turn +reaches its iteration limit, so those calls are not included. Hermes adds no tokens for a provider response that carried no usage, so an -invocation with such a response has no `usage`. Hermes also skips usage -recording for a response it rejects as invalid, a reply that hit the token -limit, and a reply stopped by a content filter, so an invocation with such a -response has no `usage` either. A streamed call that fails and is retried -inside Hermes' API call, such as an empty stream, never reaches these checks, -so its billed usage is not observed and the invocation can undercount. Hermes +invocation with such a response has no `usage`. For a response it rejects as +invalid, a reply that hit the output token limit, and a reply stopped by a +content filter, Hermes returns before adding the response's usage to its +counters. The adapter adds the counts that such a response reported itself, +using the provider and API mode that served it, and counts it only once when +Hermes records it after all. A streamed call that fails and is retried inside +Hermes' API call, such as an empty stream, never reaches these checks, so its +billed usage is not observed and the invocation can undercount. Hermes reads a count that a response's usage left out as 0, so the adapter checks each response's usage before Hermes normalizes it. A reported zero still counts as 0. When a response left out its input count, the invocation's `usage` leaves diff --git a/adapters/typescript/pi/README.md b/adapters/typescript/pi/README.md index c45669c53..14801c93d 100644 --- a/adapters/typescript/pi/README.md +++ b/adapters/typescript/pi/README.md @@ -238,17 +238,34 @@ details, the result leaves out `cached_input_tokens`. A wire protocol outside Pi's built-in ones reports no counts the adapter can check, so its calls leave the result without `usage`. +With Pi's built-in wire protocols, a request that fails before the provider +sends a successful response or any stream event, such as a rate limit (429), an +overload, a server error, or a connection that fails before any response, +counts as 0 tokens, so a call that Pi retries after such an error keeps the +result's `usage` known. A call that fails after its stream started can have +billed tokens that were never reported, so its counts stay unknown. + +When a compaction cuts inside a turn, Pi summarizes the earlier history and the +start of that turn in separate calls and reports only their combined usage. The +adapter counts each of those calls from its own report, after checking that the +calls it saw are the ones the compaction needed and that they add up to the +combined usage; otherwise the result has no `usage`. + +A tool result's `usage` counts each count it holds: its input needs `input`, +`cacheRead`, and `cacheWrite`, its cache reads need `cacheRead`, and its output +needs `output`. A count the tool left out, or that is not a non-negative +safe integer, is left out of the result. A tool result without a `usage` field adds +nothing and leaves the result's `usage` unchanged. An extension's `message_end` +handler that replaces an assistant message's usage does not change which +counts the provider reported. + The result has no `usage` when Pi produced no assistant message, when a model -call reported none of the counts, when a tool result's `usage` has a missing or -invalid count, when an assistant message was aborted, when a compaction was -aborted or reported no usage, or when Pi retried a compaction summary call, -because Pi keeps no usage for the failed attempt. It also has no `usage` when -Pi recorded a cache warm-up during the invocation: Pi replays a warm-up outside -the agent's stream function, so the adapter cannot tie its raw report to the -warm-up's usage entry. It also has no `usage` when a compaction cuts inside a -turn, because Pi then summarizes the earlier history and the start of that turn -in separate calls and reports only their combined usage. A tool result without -a `usage` field adds nothing and leaves the result's `usage` unchanged. +call reported none of the counts, when an assistant message was aborted, when a +compaction was aborted or reported no usage, or when Pi retried a compaction +summary call, because Pi keeps no usage for the failed attempt. It also has no +`usage` when Pi recorded a cache warm-up during the invocation: Pi replays a +warm-up outside the agent's stream function, so the adapter cannot tie its raw +report to the warm-up's usage entry. ## Custom Tool Modules From cdea9915dfdd7234a04baca0389bec894a5787d3 Mon Sep 17 00:00:00 2001 From: mnajafian-nv Date: Tue, 20 Oct 2026 00:39:31 -0700 Subject: [PATCH 23/68] fix: count the Hermes iteration-limit summary response When a turn reaches max_iterations, Hermes makes one extra summary request and never adds its usage to the session counters, so the invocation undercounted silently. Every summary attempt hands its response to agent.chat_completion_helpers._summary_text, so the adapter wraps that function and counts the response like one Hermes skipped. A summary without usage leaves the invocation's usage unknown. Signed-off-by: mnajafian-nv --- .../nemo_fabric_adapters/hermes/adapter.py | 61 ++++++++++++++++++- tests/adapters/test_hermes_adapter.py | 3 + tests/e2e/test_hermes_e2e.py | 54 ++++++++++++++++ 3 files changed, 117 insertions(+), 1 deletion(-) diff --git a/adapters/python/hermes/src/nemo_fabric_adapters/hermes/adapter.py b/adapters/python/hermes/src/nemo_fabric_adapters/hermes/adapter.py index 9b064b656..67c3246f1 100755 --- a/adapters/python/hermes/src/nemo_fabric_adapters/hermes/adapter.py +++ b/adapters/python/hermes/src/nemo_fabric_adapters/hermes/adapter.py @@ -789,6 +789,64 @@ def _install_usage_observer() -> None: handlers[name] = current for name, handler in handlers.items(): setattr(turn_response_check, name, _observed_response_handler(name, handler)) + _observe_iteration_summary() + + +# The function each of Hermes' iteration-limit summary attempts passes its +# provider response to (agent/chat_completion_helpers.py), and the attempt +# builders that call it. +_ITERATION_SUMMARY_HANDLER = "_summary_text" +_ITERATION_SUMMARY_ATTEMPTS = ( + "_chat_summary_attempt", + "_anthropic_summary_attempt", + "_codex_summary_attempt", +) + + +def _observe_iteration_summary() -> None: + """Route Hermes' iteration-limit summary responses through the observers. + + When a turn reaches ``max_iterations``, ``handle_max_iterations`` makes an + extra summary request and never adds its usage to the session counters. + Every summary attempt hands its response to ``_summary_text``, a global of + ``agent.chat_completion_helpers``, so the adapter wraps that function once + per process and counts the response like one Hermes skipped. Start fails + if an attempt no longer calls it by that name. + """ + from agent import chat_completion_helpers + + current = getattr(chat_completion_helpers, _ITERATION_SUMMARY_HANDLER, None) + if getattr(current, "nemo_fabric_observed", False): + return + if ( + current is None + or not all( + _inner_function_calls( + getattr(chat_completion_helpers, attempt, None), + _ITERATION_SUMMARY_HANDLER, + ) + for attempt in _ITERATION_SUMMARY_ATTEMPTS + ) + or "response" not in inspect.signature(current).parameters + ): + raise lifecycle.LifecycleError( + "hermes_usage_accounting_unsupported", + "This Hermes version makes its iteration-limit summary request " + "outside agent.chat_completion_helpers._summary_text", + ) + setattr( + chat_completion_helpers, + _ITERATION_SUMMARY_HANDLER, + _observed_response_handler(_ITERATION_SUMMARY_HANDLER, current), + ) + + +def _inner_function_calls(function: Any, name: str) -> bool: + """Return whether a function defined inside ``function`` calls global ``name``.""" + code = getattr(function, "__code__", None) + return code is not None and any( + name in getattr(const, "co_names", ()) for const in code.co_consts + ) _REPORTED_COUNTS = ("input", "cached", "output") @@ -830,7 +888,8 @@ def _skips_usage_recording(handler: str, outcome: Any) -> bool: unless the handler's action is ``fallthrough``, and after a truncated response when the action is ``return``, ``break`` or ``continue`` (``agent/turn_response_check.py``). Any other outcome falls through to - ``record_response_usage``, which adds the response's usage itself. + ``record_response_usage``, which adds the response's usage itself. An + iteration-limit summary response is never recorded. """ action = getattr(outcome, "action", None) if handler == "recover_from_truncation": diff --git a/tests/adapters/test_hermes_adapter.py b/tests/adapters/test_hermes_adapter.py index 80d908bdc..a9050810b 100644 --- a/tests/adapters/test_hermes_adapter.py +++ b/tests/adapters/test_hermes_adapter.py @@ -31,6 +31,9 @@ pytestmark = pytest.mark.usefixtures("requires_hermes_agent") if importlib.util.find_spec("run_agent") is not None: + # Load the Hermes module the adapter's usage observer wraps before tests + # replace the tools package with stubs. + import agent.chat_completion_helpers # noqa: F401 from hermes_state import SessionDB from run_agent import AIAgent diff --git a/tests/e2e/test_hermes_e2e.py b/tests/e2e/test_hermes_e2e.py index 74fb674cb..71ec57436 100644 --- a/tests/e2e/test_hermes_e2e.py +++ b/tests/e2e/test_hermes_e2e.py @@ -244,6 +244,60 @@ async def test_hermes_usage_counts_responses_hermes_skips( }, skipped.to_mapping() +@pytest.mark.usefixtures("mock_nvidia_api_key") +@pytest.mark.parametrize( + ("summary_usage", "expected_usage"), + [ + ( + _chat_usage(600, 20), + { + "input_tokens": 640, + "input_tokens_include_cache": True, + "output_tokens": 532, + "total_tokens": 1172, + "metadata": {}, + }, + ), + (None, None), + ], + ids=["reported", "omitted"], +) +async def test_hermes_usage_counts_the_iteration_limit_summary( + code_review_agent_dir: Path, + api_server: str, + summary_usage: dict[str, int] | None, + expected_usage: dict[str, object] | None, +): + os.environ["ADAPTER_PYTHON"] = sys.executable + config = hermes_config() + config.models["default"].base_url = f"{api_server}/v1" + # With one iteration, the truncated reply leaves no iteration for its + # continuation, so Hermes asks for a summary instead. + config.runtime.max_turns = 1 + + async with await Fabric().start_runtime( + config, base_dir=code_review_agent_dir + ) as runtime: + await runtime.invoke(input="first") + scenario_response = requests.post( + f"{api_server}/_scenario", + json={ + "replies": [ + {"finish_reason": "length", "usage": _chat_usage(40, 512)}, + {"usage": summary_usage}, + ] + }, + timeout=5, + ) + scenario_response.raise_for_status() + summarized = await runtime.invoke(input="second") + + assert summarized["status"] == "succeeded", summarized.to_mapping() + assert summarized.to_mapping().get("usage") == expected_usage, ( + summarized.to_mapping() + ) + + @pytest.mark.usefixtures("mock_nvidia_api_key") async def test_env_secrets_in_headers( code_review_agent_dir: Path, From d012300a91fec7b09c3e312a639fb71538770c4a Mon Sep 17 00:00:00 2001 From: mnajafian-nv Date: Tue, 20 Oct 2026 00:39:34 -0700 Subject: [PATCH 24/68] fix: leave Hermes usage unknown when context compression runs Hermes records each context-compression summary call through auxiliary accounting, not the session counters, so an invocation that compressed reported a known but low total. The adapter counts compression calls on its own agent's compressor and leaves input, cached input and output unknown for an invocation in which one ran. Signed-off-by: mnajafian-nv --- .../nemo_fabric_adapters/hermes/adapter.py | 36 +++++++++++++++++ tests/e2e/test_hermes_e2e.py | 39 +++++++++++++++++++ 2 files changed, 75 insertions(+) diff --git a/adapters/python/hermes/src/nemo_fabric_adapters/hermes/adapter.py b/adapters/python/hermes/src/nemo_fabric_adapters/hermes/adapter.py index 67c3246f1..77af05f6e 100755 --- a/adapters/python/hermes/src/nemo_fabric_adapters/hermes/adapter.py +++ b/adapters/python/hermes/src/nemo_fabric_adapters/hermes/adapter.py @@ -75,6 +75,9 @@ def __init__(self) -> None: # Tokens of provider responses Hermes returned on before adding them to # its session counters, in the counters' buckets. self._skipped_response_tokens = dict.fromkeys(_SESSION_TOKEN_KEYS, 0) + # Context-compression summary calls, whose usage Hermes keeps out of + # its session counters. + self._compression_calls = 0 async def start(self, payload: dict[str, Any]) -> None: if self._started: @@ -300,6 +303,7 @@ def run_hermes_turn() -> tuple[dict[str, Any], str]: uncounted_before = self._uncounted_responses unreported_before = dict(self._unreported_counts) skipped_before = dict(self._skipped_response_tokens) + compression_calls_before = self._compression_calls invoke_task = asyncio.create_task(asyncio.to_thread(run_hermes_turn)) self._active_invoke_task = invoke_task @@ -395,6 +399,13 @@ def clear_active_invoke_task( count for count, responses in self._unreported_counts.items() if responses != unreported_before[count] + ) + # A compression summary call's tokens are not in the counters, so + # every count it adds to is an undercount. + | ( + frozenset(_REPORTED_COUNTS) + if self._compression_calls != compression_calls_before + else frozenset() ), ) if self._uncounted_responses != uncounted_before: @@ -412,6 +423,7 @@ def _observe_usage_reports(self, agent: Any) -> None: """Count the agent's provider responses that left usage unreported.""" self._count_usage_less_responses(agent) self._count_unreported_counts(agent) + self._count_compression_calls(agent) def _count_usage_less_responses(self, agent: Any) -> None: """Count the agent's provider responses that carried no usage. @@ -434,6 +446,30 @@ def count_and_note() -> None: compressor.note_usage_less_response = count_and_note + def _count_compression_calls(self, agent: Any) -> None: + """Count the context-compression summary calls of the agent it owns. + + Hermes sends each compression summary through the auxiliary client and + records its usage outside the session counters + (``agent/aux_accounting.py``). The built-in ``ContextCompressor`` + issues every summary call through its ``_call_summary_llm`` method + (``agent/context_compressor.py``); a plugin context engine is counted + at its ``compress`` method, since its calls cannot be seen. + """ + compressor = agent.context_compressor + name = ( + "_call_summary_llm" + if callable(getattr(compressor, "_call_summary_llm", None)) + else "compress" + ) + summarize = getattr(compressor, name) + + def count_and_summarize(*args: Any, **kwargs: Any) -> Any: + self._compression_calls += 1 + return summarize(*args, **kwargs) + + setattr(compressor, name, count_and_summarize) + def _count_unreported_counts(self, agent: Any) -> None: """Count the agent's provider responses whose usage left out a count. diff --git a/tests/e2e/test_hermes_e2e.py b/tests/e2e/test_hermes_e2e.py index 71ec57436..ab36fb2f7 100644 --- a/tests/e2e/test_hermes_e2e.py +++ b/tests/e2e/test_hermes_e2e.py @@ -298,6 +298,45 @@ async def test_hermes_usage_counts_the_iteration_limit_summary( ) +@pytest.mark.usefixtures("mock_nvidia_api_key") +async def test_hermes_usage_is_unknown_when_context_compression_runs( + code_review_agent_dir: Path, + api_server: str, +): + os.environ["ADAPTER_PYTHON"] = sys.executable + config = hermes_config() + config.models["default"].base_url = f"{api_server}/v1" + config.runtime.max_turns = 5 + + async with await Fabric().start_runtime( + config, base_dir=code_review_agent_dir + ) as runtime: + # Build enough history for Hermes to have turns to summarize. + for turn in range(6): + await runtime.invoke(input=f"turn {turn}") + # A prompt above the compression threshold makes Hermes compress + # before its next request. + scenario_response = requests.post( + f"{api_server}/_scenario", + json={"replies": [{"usage": _chat_usage(120_000, 5)}]}, + timeout=5, + ) + scenario_response.raise_for_status() + large = await runtime.invoke(input="large") + requests.post(f"{api_server}/_reset", timeout=5).raise_for_status() + compressed = await runtime.invoke(input="after") + sent = requests.get(f"{api_server}/_requests", timeout=5).json() + + assert large.usage is not None, large.to_mapping() + assert (large.usage.input_tokens, large.usage.output_tokens) == (120_000, 5) + assert sent[0]["messages"][0]["content"].startswith( + "You are a summarization agent" + ), sent + # Hermes keeps the compression summary call out of its session counters. + assert compressed["status"] == "succeeded", compressed.to_mapping() + assert compressed.usage is None, compressed.to_mapping() + + @pytest.mark.usefixtures("mock_nvidia_api_key") async def test_env_secrets_in_headers( code_review_agent_dir: Path, From 9ea4ff821f014ee35019ab071f036e665fb2507a Mon Sep 17 00:00:00 2001 From: mnajafian-nv Date: Tue, 20 Oct 2026 00:39:35 -0700 Subject: [PATCH 25/68] fix: count a failed Pi call as 0 only after an HTTP error status A connection failure before any response, or a Google stream that failed before its first chunk, counted as a known 0 although the provider may have billed the prompt. The adapter now reads each attempt's HTTP status through the request's fetch and counts a failed call as 0 only when every attempt got a 4xx or 5xx status before streaming. Protocols that take no custom fetch, and calls with an unanswered attempt, stay unknown. Signed-off-by: mnajafian-nv --- adapters/typescript/pi/src/pi-sdk.ts | 132 +++++++++++++++----- adapters/typescript/pi/test/pi-sdk.test.mjs | 60 +++++++-- 2 files changed, 150 insertions(+), 42 deletions(-) diff --git a/adapters/typescript/pi/src/pi-sdk.ts b/adapters/typescript/pi/src/pi-sdk.ts index a9347494b..a1367bdfa 100644 --- a/adapters/typescript/pi/src/pi-sdk.ts +++ b/adapters/typescript/pi/src/pi-sdk.ts @@ -5,6 +5,7 @@ // into a controlled in-memory Pi session, including model credentials, skills, // extensions, custom tools, and workspace containment. +import { AsyncLocalStorage } from "node:async_hooks"; import { mkdtemp, realpath, rm, stat } from "node:fs/promises"; import { homedir, tmpdir } from "node:os"; import { extname, isAbsolute, join, relative, resolve, sep } from "node:path"; @@ -85,6 +86,7 @@ export interface ReportedCounts { export type ReportedCountsLookup = (messageOrUsage: unknown) => ReportedCounts | undefined; const ALL_REPORTED: ReportedCounts = { input: true, cacheRead: true, output: true }; +const NONE_REPORTED: ReportedCounts = { input: false, cacheRead: false, output: false }; /** Pi's four usage counts, each undefined unless it is a non-negative safe integer. */ function piCounts(usage: unknown): Record<"input" | "cacheRead" | "cacheWrite" | "output", number | undefined> { @@ -445,30 +447,90 @@ export function reportedCountsFrom(api: string, seen: ReadonlySet): Repo }; } +/** What the adapter observed of one model call's HTTP attempts and stream. */ +export interface CallObservation { + /** The provider answered an attempt with a 2xx status (`onResponse`). */ + accepted: boolean; + /** The provider sent at least one raw stream event. */ + streamed: boolean; + /** The provider answered an attempt with a 4xx or 5xx status. */ + rejectedStatus: boolean; + /** An attempt failed before any HTTP response, such as a connection error. */ + unanswered: boolean; +} + +// Built-in protocols that send their model request through the call's `fetch` +// option (`@earendil-works/pi-ai` 1.0.3, `api/*.js`). The Google protocols +// reject a custom fetch, Bedrock does not use one, and Codex may use a +// WebSocket, which the fetch does not see. +const FETCH_PROTOCOLS = new Set([ + "anthropic-messages", + "azure-openai-responses", + "mistral-conversations", + "openai-codex-responses", + "openai-completions", + "openai-responses", + "pi-messages", +]); + +interface CallAttempts { + delegate: typeof globalThis.fetch; + observation: CallObservation; +} + +const callAttempts = new AsyncLocalStorage(); + +/** + * The fetch passed to every observed model call. It is one function, so Pi's + * clients cached per fetch, such as the Anthropic federation client, keep + * their token cache. It forwards to the call's own fetch and records each + * attempt's outcome on the call that made it. + */ +const observedFetch: typeof globalThis.fetch = async (input, init) => { + const attempts = callAttempts.getStore(); + if (attempts === undefined) { + return globalThis.fetch(input, init); + } + let response: Response; + try { + response = await attempts.delegate(input, init); + } catch (error) { + attempts.observation.unanswered = true; + throw error; + } + attempts.observation.rejectedStatus ||= response.status >= 400; + return response; +}; + /** - * Which counts an ended model call's provider reported. A call that ended in - * an error before the provider accepted the request, with no successful HTTP - * response, no raw stream event and Pi's zero-filled usage, reports all its - * counts as 0. Pi's built-in protocols call `onResponse` only after the - * request succeeds and raise a rejection such as a rate limit (429), an - * overload or a server error before the first raw event - * (`@earendil-works/pi-ai` 1.0.3, `api/*.js`), so such a call generated - * nothing. A protocol outside those may never call either hook, so its calls - * keep only the counts their raw events reported. + * Which counts an ended model call's provider reported. A call with an attempt + * that failed before any HTTP response, such as a connection error or a + * timeout before headers, reports none: the provider may have billed its + * prompt. A call that ended in an error reports all its counts as 0 only when + * its protocol sends requests through the observed fetch, the provider + * answered every attempt with a 4xx or 5xx status, no attempt was accepted, + * no raw event arrived, and Pi's usage is all zeros. Pi's built-in protocols + * raise such a rejection, such as a rate limit (429), an overload or a server + * error, before the first raw event (`@earendil-works/pi-ai` 1.0.3, + * `api/*.js`), so such a call generated nothing. Any other call keeps only the + * counts its raw events reported. */ export function endedCallCounts( api: string, seen: ReadonlySet, message: { stopReason?: unknown; usage?: unknown }, - accepted: boolean, - streamed: boolean, + observation: CallObservation, ): ReportedCounts { + if (observation.unanswered) { + return NONE_REPORTED; + } const counts = piCounts(message.usage); const rejected = - REPORTED_FIELDS[api] !== undefined && + FETCH_PROTOCOLS.has(api) && message.stopReason === "error" && - !accepted && - !streamed && + observation.rejectedStatus && + !observation.accepted && + !observation.streamed && counts.input === 0 && counts.cacheRead === 0 && counts.cacheWrite === 0 && @@ -485,7 +547,8 @@ export function endedCallCounts( * so wrapping it covers both, and Pi passes the final message to * `message_end` and its usage object unchanged to `compaction_end`. The record * is stored before Pi emits either event. A request the provider rejected - * before generating reports all its counts as 0. Pi's cache warmer replays a + * with an HTTP error status before generating reports all its counts as 0, and + * a call with an attempt that got no HTTP response reports none. Pi's cache warmer replays a * request with the same options, so a warm-up's raw events reach this * observer too, but the warm-up runs outside the agent's stream function: its * usage entry has no record and stays unknown. The record is copied when the @@ -500,25 +563,32 @@ function recordReportedCounts( agent.streamFunction = async (model, context, options) => { const seen = new Set(); let api: string = model.api; - let accepted = false; - let streamed = false; - const response = await streamFunction(model, context, { - ...options, - onResponse: async (providerResponse, responseModel) => { - accepted ||= providerResponse.status >= 200 && providerResponse.status < 300; - await options?.onResponse?.(providerResponse, responseModel); - }, - onProviderStreamEvent: async (data, eventModel) => { - streamed = true; - api = eventModel?.api ?? model.api; - noteReportedFields(seen, api, data); - await options?.onProviderStreamEvent?.(data, eventModel); - }, - }); + const observation: CallObservation = { accepted: false, streamed: false, rejectedStatus: false, unanswered: false }; + const ownFetch = options?.fetch; + const attempts: CallAttempts = { + delegate: ownFetch === undefined || ownFetch === observedFetch ? globalThis.fetch : ownFetch, + observation, + }; + const response = await callAttempts.run(attempts, () => + streamFunction(model, context, { + ...options, + ...(FETCH_PROTOCOLS.has(model.api) ? { fetch: observedFetch } : {}), + onResponse: async (providerResponse, responseModel) => { + observation.accepted ||= providerResponse.status >= 200 && providerResponse.status < 300; + await options?.onResponse?.(providerResponse, responseModel); + }, + onProviderStreamEvent: async (data, eventModel) => { + observation.streamed = true; + api = eventModel?.api ?? model.api; + noteReportedFields(seen, api, data); + await options?.onProviderStreamEvent?.(data, eventModel); + }, + }), + ); void response.result().then( (message) => { if (typeof message.usage === "object" && message.usage !== null) { - const counts = endedCallCounts(api, seen, message, accepted, streamed); + const counts = endedCallCounts(api, seen, message, observation); records.set(message, counts); records.set(message.usage, counts); onCallEnd(message.usage); diff --git a/adapters/typescript/pi/test/pi-sdk.test.mjs b/adapters/typescript/pi/test/pi-sdk.test.mjs index 980d5e981..c410d0d56 100644 --- a/adapters/typescript/pi/test/pi-sdk.test.mjs +++ b/adapters/typescript/pi/test/pi-sdk.test.mjs @@ -1661,29 +1661,47 @@ test("a retried compaction summary leaves the prompt's tokens unknown", () => { ); }); -test("counts a rejected request as 0 tokens only for a built-in protocol", () => { +test("counts a request as 0 tokens only when the provider rejected it with an HTTP error status", () => { const rejected = { stopReason: "error", usage: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 } }; + const statusOnly = { accepted: false, streamed: false, rejectedStatus: true, unanswered: false }; const none = { input: false, cacheRead: false, output: false }; - - assert.deepEqual(endedCallCounts("openai-completions", new Set(), rejected, false, false), allReported()); - // A custom protocol may never report its HTTP response or raw events. - assert.deepEqual(endedCallCounts("custom-api", new Set(), rejected, false, false), none); + const counts = (api, message, observation) => endedCallCounts(api, new Set(), message, { ...statusOnly, ...observation }); + + assert.deepEqual(counts("openai-completions", rejected, {}), allReported()); + assert.deepEqual(counts("anthropic-messages", rejected, {}), allReported()); + // No observed error status: a connection failure, or a protocol whose + // requests do not go through the observed fetch, may have billed the prompt. + assert.deepEqual(counts("openai-completions", rejected, { rejectedStatus: false }), none); + assert.deepEqual(counts("openai-completions", rejected, { unanswered: true }), none); + assert.deepEqual(counts("google-generative-ai", rejected, {}), none); + assert.deepEqual(counts("bedrock-converse-stream", rejected, {}), none); + assert.deepEqual(counts("custom-api", rejected, {}), none); // An accepted or streamed request, or one with counts, may have billed tokens. - assert.deepEqual(endedCallCounts("openai-completions", new Set(), rejected, true, false), none); - assert.deepEqual(endedCallCounts("openai-completions", new Set(), rejected, false, true), none); + assert.deepEqual(counts("openai-completions", rejected, { accepted: true }), none); + assert.deepEqual(counts("openai-completions", rejected, { streamed: true }), none); + assert.deepEqual(counts("openai-completions", { ...rejected, usage: { ...rejected.usage, input: 3 } }, {}), none); + assert.deepEqual(counts("openai-completions", { ...rejected, stopReason: "aborted" }, {}), none); + // An attempt without a response leaves even a later successful attempt's + // counts unknown. + const seen = new Set(["prompt_tokens", "completion_tokens"]); assert.deepEqual( - endedCallCounts("openai-completions", new Set(), { ...rejected, usage: { ...rejected.usage, input: 3 } }, false, false), + endedCallCounts("openai-completions", seen, { stopReason: "stop", usage: rejected.usage }, { + accepted: true, + streamed: true, + rejectedStatus: false, + unanswered: true, + }), none, ); - assert.deepEqual(endedCallCounts("openai-completions", new Set(), { ...rejected, stopReason: "aborted" }, false, false), none); }); /** * Run prompts through the real Pi SDK against a local chat-completions endpoint * that answers each request from `replies` in order. A reply's `usage` is sent * as the stream's final usage chunk; a reply without one sends no usage. A - * reply with `status` rejects the request with that HTTP status, and one with - * `dropStream` closes the connection after the first chunk. + * reply with `status` rejects the request with that HTTP status, one with + * `dropStream` closes the connection after the first chunk, and one with + * `closeBeforeHeaders` closes it before sending any response. * `modelExtensions` overrides the catalog entry, such as its context window, * and `extension` is the source of a Pi extension the session loads. */ @@ -1697,6 +1715,10 @@ async function promptsAgainstScriptedEndpoint(replies, prompts, modelExtensions const provider = createServer(async (request, response) => { for await (const _chunk of request); const reply = replies[served++]; + if (reply.closeBeforeHeaders) { + request.socket.destroy(); + return; + } if (reply.status !== undefined) { response.writeHead(reply.status, { "content-type": "application/json" }); response.end(JSON.stringify({ error: { message: "rate limited", type: "rate_limit_error" } })); @@ -1881,6 +1903,22 @@ test( }, ); +test( + "leaves a prompt's tokens unknown when a connection fails before any HTTP response", + { skip: supportsPi ? false : "Pi 1.0 requires Node 22.19 or newer" }, + async () => { + // The provider may have received and billed the prompt before the + // connection failed, so the failed attempt is not a known 0. + const [retried] = await promptsAgainstScriptedEndpoint( + [{ closeBeforeHeaders: true }, { content: "after retry", usage: { prompt_tokens: 50, completion_tokens: 5 } }], + ["one"], + ); + + assert.equal(retried.text, "after retry"); + assert.equal(retried.usage, undefined); + }, +); + test( "leaves a prompt's tokens unknown when a stream fails after it started", { skip: supportsPi ? false : "Pi 1.0 requires Node 22.19 or newer" }, From 17e0735a70feb6146ba7e7396891f31fff042192 Mon Sep 17 00:00:00 2001 From: mnajafian-nv Date: Tue, 20 Oct 2026 00:39:36 -0700 Subject: [PATCH 26/68] docs: describe the Hermes summary and compression and Pi rejection rules State that the Hermes iteration-limit summary is counted, that a compressed invocation has no usage, and the exact HTTP status rule for a failed Pi call. Signed-off-by: mnajafian-nv --- adapters/python/hermes/README.md | 9 +++++++-- adapters/typescript/pi/README.md | 20 ++++++++++++++------ 2 files changed, 21 insertions(+), 8 deletions(-) diff --git a/adapters/python/hermes/README.md b/adapters/python/hermes/README.md index 7428917e9..3af96d2b1 100644 --- a/adapters/python/hermes/README.md +++ b/adapters/python/hermes/README.md @@ -124,8 +124,13 @@ never adds its tokens to the next one. A counter that is missing, that is not a non-negative integer below 2^64, or that went down because Hermes reset its session, is left out for that invocation. Hermes records auxiliary calls, such as context-compression summaries and session titles, separately from these -counters, and it does not count the summary request it makes when a turn -reaches its iteration limit, so those calls are not included. +counters. Session title calls are not included. An invocation in which Hermes +ran a context-compression summary call, or in which a plugin context engine +compressed, has no `usage`, because the compression calls' tokens are missing +from the counters. Hermes also leaves out of its counters the summary request it +makes when a turn reaches its iteration limit. The adapter adds the counts that +this summary response reported, and an invocation whose summary response +carried no usage has no `usage`. Hermes adds no tokens for a provider response that carried no usage, so an invocation with such a response has no `usage`. For a response it rejects as diff --git a/adapters/typescript/pi/README.md b/adapters/typescript/pi/README.md index 14801c93d..aad02a7c1 100644 --- a/adapters/typescript/pi/README.md +++ b/adapters/typescript/pi/README.md @@ -238,12 +238,20 @@ details, the result leaves out `cached_input_tokens`. A wire protocol outside Pi's built-in ones reports no counts the adapter can check, so its calls leave the result without `usage`. -With Pi's built-in wire protocols, a request that fails before the provider -sends a successful response or any stream event, such as a rate limit (429), an -overload, a server error, or a connection that fails before any response, -counts as 0 tokens, so a call that Pi retries after such an error keeps the -result's `usage` known. A call that fails after its stream started can have -billed tokens that were never reported, so its counts stay unknown. +A failed call counts as 0 tokens only when the provider answered each of its +attempts with a 4xx or 5xx HTTP status, such as a rate limit (429), an overload +or a server error, before sending a successful response or any stream event. +A call that Pi retries after such an error keeps the result's `usage` known. +The adapter reads that status through the request's `fetch`, which the +`anthropic-messages`, `azure-openai-responses`, `mistral-conversations`, +`openai-codex-responses`, `openai-completions`, `openai-responses` and +`pi-messages` protocols use. A failed call of any other protocol, including the +Google protocols and Bedrock, or a Codex call sent over a WebSocket, leaves its +counts unknown. A call with an attempt that got no HTTP response, such as a +connection failure or a timeout before headers, leaves its counts unknown even +when a later attempt succeeds, because the provider may have billed the prompt. +A call that fails after its stream started can have billed tokens that were +never reported, so its counts stay unknown too. When a compaction cuts inside a turn, Pi summarizes the earlier history and the start of that turn in separate calls and reports only their combined usage. The From e28bab53bb12260129193bf9a79a70e8cd3866eb Mon Sep 17 00:00:00 2001 From: mnajafian-nv Date: Tue, 20 Oct 2026 01:32:24 -0700 Subject: [PATCH 27/68] fix: keep Pi cache reads known when input is unknown Pi dropped cached_input_tokens whenever the input that includes them was unknown, for example a tool report with cacheRead but no cacheWrite. Cache reads are their own count, so the aggregation, the result mapping and the empty-usage checks now keep them independently of input. Signed-off-by: mnajafian-nv --- adapters/typescript/pi/src/pi-sdk.ts | 12 ++++++------ adapters/typescript/pi/src/runtime.ts | 9 +++++---- adapters/typescript/pi/test/pi-sdk.test.mjs | 13 ++++++++----- adapters/typescript/pi/test/runtime.test.mjs | 7 ++++++- 4 files changed, 25 insertions(+), 16 deletions(-) diff --git a/adapters/typescript/pi/src/pi-sdk.ts b/adapters/typescript/pi/src/pi-sdk.ts index a1367bdfa..062c8d55d 100644 --- a/adapters/typescript/pi/src/pi-sdk.ts +++ b/adapters/typescript/pi/src/pi-sdk.ts @@ -122,21 +122,21 @@ function callCounts(usage: unknown): { counts: Required; present: /** * Add a call's counts to a prompt's running total. Each count stays known only - * while every call reported it; the cache reads also need the input. The total - * is unknown once neither input nor output is known. + * while every call reported it, independently of the others. The total is + * unknown once no count is known. */ function sumCounts(total: PiTokenUsage, call: Required, reported: ReportedCounts): PiTokenUsage | undefined { const sum: PiTokenUsage = {}; if (reported.input && total.input !== undefined) { sum.input = total.input + call.input; - if (reported.cacheRead && total.cacheRead !== undefined) { - sum.cacheRead = total.cacheRead + call.cacheRead; - } + } + if (reported.cacheRead && total.cacheRead !== undefined) { + sum.cacheRead = total.cacheRead + call.cacheRead; } if (reported.output && total.output !== undefined) { sum.output = total.output + call.output; } - return sum.input === undefined && sum.output === undefined ? undefined : sum; + return sum.input === undefined && sum.cacheRead === undefined && sum.output === undefined ? undefined : sum; } /** diff --git a/adapters/typescript/pi/src/runtime.ts b/adapters/typescript/pi/src/runtime.ts index fea83d60e..200308bbf 100644 --- a/adapters/typescript/pi/src/runtime.ts +++ b/adapters/typescript/pi/src/runtime.ts @@ -64,17 +64,18 @@ function failed(code: string, message: string): AgentRunResult { /** * The invocation's usage, with input counting its cache reads and writes. A - * count that a call left unreported stays absent, and so does the total unless - * both input and output are known. + * count that a call left unreported stays absent, cache reads are kept even + * when input is unknown, and the total is absent unless both input and output + * are known. */ function agentUsage(usage: PiTokenUsage | undefined): AgentUsage | undefined { - if (usage === undefined || (usage.input === undefined && usage.output === undefined)) { + if (usage === undefined || (usage.input === undefined && usage.cacheRead === undefined && usage.output === undefined)) { return undefined; } const { input, cacheRead, output } = usage; return { ...(input === undefined ? {} : { input_tokens: input, input_tokens_include_cache: true }), - ...(input === undefined || cacheRead === undefined ? {} : { cached_input_tokens: cacheRead }), + ...(cacheRead === undefined ? {} : { cached_input_tokens: cacheRead }), ...(output === undefined ? {} : { output_tokens: output }), ...(input === undefined || output === undefined ? {} : { total_tokens: input + output }), }; diff --git a/adapters/typescript/pi/test/pi-sdk.test.mjs b/adapters/typescript/pi/test/pi-sdk.test.mjs index c410d0d56..b559226b0 100644 --- a/adapters/typescript/pi/test/pi-sdk.test.mjs +++ b/adapters/typescript/pi/test/pi-sdk.test.mjs @@ -1465,12 +1465,14 @@ test("sums each call's own usage report", () => { // A tool result that reports usage, as a custom tool can. const toolResult = (usage) => ({ type: "message_end", message: { role: "toolResult", usage } }); -test("a tool usage report without a cache count keeps its output known", () => { +test("a tool usage report without a cache count keeps its other counts known", () => { const total = addEventUsage(NO_TOKENS, toolResult({ input: 10, cacheRead: 2, output: 3 }), allReported); - // Input counts cache writes, so without cacheWrite the input, the cache - // reads inside it and the total are unknown. - assert.deepEqual(total, { output: 3 }); + // Input counts cache writes, so without cacheWrite the input and the total + // are unknown; the cache reads and output stay known. + assert.deepEqual(total, { cacheRead: 2, output: 3 }); + // A report with only cache reads keeps them. + assert.deepEqual(addEventUsage(NO_TOKENS, toolResult({ cacheRead: 4 }), allReported), { cacheRead: 4 }); }); test("a tool usage report with only an output count adds that count", () => { @@ -1548,7 +1550,8 @@ test("a model call's counts are known only as far as its provider reported them" input: 0, output: 0, }); - assert.equal(addEventUsage(NO_TOKENS, assistant, reported({ cacheRead: true })), undefined); + // Cache reads stay known without the input that includes them. + assert.deepEqual(addEventUsage(NO_TOKENS, assistant, reported({ cacheRead: true })), { cacheRead: 0 }); // A tool reports its own usage, so a zero there is a measured zero. assert.deepEqual( addEventUsage(NO_TOKENS, { type: "message_end", message: { role: "toolResult", usage: zeroFilled } }, unseen), diff --git a/adapters/typescript/pi/test/runtime.test.mjs b/adapters/typescript/pi/test/runtime.test.mjs index caa39a100..4604fad49 100644 --- a/adapters/typescript/pi/test/runtime.test.mjs +++ b/adapters/typescript/pi/test/runtime.test.mjs @@ -585,7 +585,7 @@ test("reports the invocation's tokens with cache reads and writes inside the inp }); test("reports only the token counts the prompt's provider reported", async () => { - const outcomes = [{ input: 130, cacheRead: 30 }, { output: 5 }, { input: 70, output: 7 }]; + const outcomes = [{ input: 130, cacheRead: 30 }, { output: 5 }, { input: 70, output: 7 }, { cacheRead: 12, output: 3 }, { cacheRead: 4 }]; const runtime = new PiAdapterRuntime({ async create() { let turnCount = 0; @@ -604,6 +604,8 @@ test("reports only the token counts the prompt's provider reported", async () => const inputOnly = await runtime.invoke({ input: "one" }, context); const outputOnly = await runtime.invoke({ input: "two" }, context); const noCacheDetail = await runtime.invoke({ input: "three" }, context); + const noInput = await runtime.invoke({ input: "four" }, context); + const cacheOnly = await runtime.invoke({ input: "five" }, context); await runtime.stop(); assert.deepEqual(inputOnly.usage, { @@ -618,4 +620,7 @@ test("reports only the token counts the prompt's provider reported", async () => output_tokens: 7, total_tokens: 77, }); + // Cache reads stay known when the input that includes them is not. + assert.deepEqual(noInput.usage, { cached_input_tokens: 12, output_tokens: 3 }); + assert.deepEqual(cacheOnly.usage, { cached_input_tokens: 4 }); }); From 4eecbda493c15f39b433bcf280b91e269e576b57 Mon Sep 17 00:00:00 2001 From: mnajafian-nv Date: Tue, 20 Oct 2026 01:32:26 -0700 Subject: [PATCH 28/68] fix: count a failed Pi attempt as 0 only after a rate limit A 408, a 5xx or another error status can arrive after the provider processed the prompt, yet the adapter counted it as 0 and kept usage known when a retry succeeded. The adapter now records each attempt's status and infers 0 only when every failed attempt got a 429. Any other error status, or an attempt without a response, leaves the call's counts unknown even after a successful retry. Signed-off-by: mnajafian-nv --- adapters/typescript/pi/src/pi-sdk.ts | 48 ++++++++++------- adapters/typescript/pi/test/pi-sdk.test.mjs | 59 ++++++++++++++------- 2 files changed, 68 insertions(+), 39 deletions(-) diff --git a/adapters/typescript/pi/src/pi-sdk.ts b/adapters/typescript/pi/src/pi-sdk.ts index 062c8d55d..af3f34b95 100644 --- a/adapters/typescript/pi/src/pi-sdk.ts +++ b/adapters/typescript/pi/src/pi-sdk.ts @@ -453,12 +453,18 @@ export interface CallObservation { accepted: boolean; /** The provider sent at least one raw stream event. */ streamed: boolean; - /** The provider answered an attempt with a 4xx or 5xx status. */ - rejectedStatus: boolean; + /** The HTTP status of each attempt that got a response, in order. */ + statuses: number[]; /** An attempt failed before any HTTP response, such as a connection error. */ unanswered: boolean; } +// The one HTTP status after which a failed attempt is a known 0: a rate limit +// rejects the request before the model runs. A timeout (408) or a server +// error (5xx) can arrive after the provider processed the prompt, and other +// 4xx statuses are not documented as unbilled for every provider. +const UNBILLED_REJECTIONS = new Set([429]); + // Built-in protocols that send their model request through the call's `fetch` // option (`@earendil-works/pi-ai` 1.0.3, `api/*.js`). The Google protocols // reject a custom fetch, Bedrock does not use one, and Codex may use a @@ -498,22 +504,20 @@ const observedFetch: typeof globalThis.fetch = async (input, init) => { attempts.observation.unanswered = true; throw error; } - attempts.observation.rejectedStatus ||= response.status >= 400; + attempts.observation.statuses.push(response.status); return response; }; /** - * Which counts an ended model call's provider reported. A call with an attempt - * that failed before any HTTP response, such as a connection error or a - * timeout before headers, reports none: the provider may have billed its - * prompt. A call that ended in an error reports all its counts as 0 only when - * its protocol sends requests through the observed fetch, the provider - * answered every attempt with a 4xx or 5xx status, no attempt was accepted, - * no raw event arrived, and Pi's usage is all zeros. Pi's built-in protocols - * raise such a rejection, such as a rate limit (429), an overload or a server - * error, before the first raw event (`@earendil-works/pi-ai` 1.0.3, - * `api/*.js`), so such a call generated nothing. Any other call keeps only the - * counts its raw events reported. + * Which counts an ended model call's provider reported. A call with an + * ambiguous attempt reports none, even when a later attempt succeeded: an + * attempt that got no HTTP response, such as a connection error or a timeout + * before headers, or that got any error status other than a rate limit (429), + * such as a 408 or a 5xx, may have been billed. A call that ended in an error + * reports all its counts as 0 only when its protocol sends requests through + * the observed fetch, the provider answered every attempt with a 429, no + * attempt was accepted, no raw event arrived, and Pi's usage is all zeros. + * Any other call keeps only the counts its raw events reported. */ export function endedCallCounts( api: string, @@ -521,14 +525,18 @@ export function endedCallCounts( message: { stopReason?: unknown; usage?: unknown }, observation: CallObservation, ): ReportedCounts { - if (observation.unanswered) { + const ambiguous = + observation.unanswered || + observation.statuses.some((status) => (status < 200 || status >= 300) && !UNBILLED_REJECTIONS.has(status)); + if (ambiguous) { return NONE_REPORTED; } const counts = piCounts(message.usage); const rejected = FETCH_PROTOCOLS.has(api) && message.stopReason === "error" && - observation.rejectedStatus && + observation.statuses.length > 0 && + observation.statuses.every((status) => UNBILLED_REJECTIONS.has(status)) && !observation.accepted && !observation.streamed && counts.input === 0 && @@ -546,9 +554,9 @@ export function endedCallCounts( * agent loop and compaction summaries both call the agent's stream function, * so wrapping it covers both, and Pi passes the final message to * `message_end` and its usage object unchanged to `compaction_end`. The record - * is stored before Pi emits either event. A request the provider rejected - * with an HTTP error status before generating reports all its counts as 0, and - * a call with an attempt that got no HTTP response reports none. Pi's cache warmer replays a + * is stored before Pi emits either event. A request the provider rate-limited + * before generating reports all its counts as 0, and a call with an ambiguous + * attempt reports none (see `endedCallCounts`). Pi's cache warmer replays a * request with the same options, so a warm-up's raw events reach this * observer too, but the warm-up runs outside the agent's stream function: its * usage entry has no record and stays unknown. The record is copied when the @@ -563,7 +571,7 @@ function recordReportedCounts( agent.streamFunction = async (model, context, options) => { const seen = new Set(); let api: string = model.api; - const observation: CallObservation = { accepted: false, streamed: false, rejectedStatus: false, unanswered: false }; + const observation: CallObservation = { accepted: false, streamed: false, statuses: [], unanswered: false }; const ownFetch = options?.fetch; const attempts: CallAttempts = { delegate: ownFetch === undefined || ownFetch === observedFetch ? globalThis.fetch : ownFetch, diff --git a/adapters/typescript/pi/test/pi-sdk.test.mjs b/adapters/typescript/pi/test/pi-sdk.test.mjs index b559226b0..0e33c7b76 100644 --- a/adapters/typescript/pi/test/pi-sdk.test.mjs +++ b/adapters/typescript/pi/test/pi-sdk.test.mjs @@ -1664,18 +1664,22 @@ test("a retried compaction summary leaves the prompt's tokens unknown", () => { ); }); -test("counts a request as 0 tokens only when the provider rejected it with an HTTP error status", () => { +test("counts a failed request as 0 tokens only after rate limits", () => { const rejected = { stopReason: "error", usage: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 } }; - const statusOnly = { accepted: false, streamed: false, rejectedStatus: true, unanswered: false }; const none = { input: false, cacheRead: false, output: false }; - const counts = (api, message, observation) => endedCallCounts(api, new Set(), message, { ...statusOnly, ...observation }); + const counts = (api, message, observation) => + endedCallCounts(api, new Set(), message, { accepted: false, streamed: false, statuses: [429], unanswered: false, ...observation }); assert.deepEqual(counts("openai-completions", rejected, {}), allReported()); - assert.deepEqual(counts("anthropic-messages", rejected, {}), allReported()); - // No observed error status: a connection failure, or a protocol whose - // requests do not go through the observed fetch, may have billed the prompt. - assert.deepEqual(counts("openai-completions", rejected, { rejectedStatus: false }), none); - assert.deepEqual(counts("openai-completions", rejected, { unanswered: true }), none); + assert.deepEqual(counts("anthropic-messages", rejected, { statuses: [429, 429] }), allReported()); + // A timeout, a server error, another client error, or an attempt without + // any response may have been billed. + for (const statuses of [[400], [408], [503], [504], [500], [429, 503]]) { + assert.deepEqual(counts("openai-completions", rejected, { statuses }), none, String(statuses)); + } + assert.deepEqual(counts("openai-completions", rejected, { statuses: [], unanswered: true }), none); + assert.deepEqual(counts("openai-completions", rejected, { statuses: [] }), none); + // Protocols whose requests do not go through the observed fetch. assert.deepEqual(counts("google-generative-ai", rejected, {}), none); assert.deepEqual(counts("bedrock-converse-stream", rejected, {}), none); assert.deepEqual(counts("custom-api", rejected, {}), none); @@ -1684,18 +1688,17 @@ test("counts a request as 0 tokens only when the provider rejected it with an HT assert.deepEqual(counts("openai-completions", rejected, { streamed: true }), none); assert.deepEqual(counts("openai-completions", { ...rejected, usage: { ...rejected.usage, input: 3 } }, {}), none); assert.deepEqual(counts("openai-completions", { ...rejected, stopReason: "aborted" }, {}), none); - // An attempt without a response leaves even a later successful attempt's - // counts unknown. + // A successful attempt keeps its reported counts after a rate limit, but not + // after an ambiguous attempt. const seen = new Set(["prompt_tokens", "completion_tokens"]); - assert.deepEqual( - endedCallCounts("openai-completions", seen, { stopReason: "stop", usage: rejected.usage }, { - accepted: true, - streamed: true, - rejectedStatus: false, - unanswered: true, - }), - none, - ); + const succeeded = { stopReason: "stop", usage: rejected.usage }; + const afterRetry = (observation) => + endedCallCounts("openai-completions", seen, succeeded, { accepted: true, streamed: true, unanswered: false, ...observation }); + assert.deepEqual(afterRetry({ statuses: [429, 200] }), { input: true, cacheRead: false, output: true }); + for (const statuses of [[408, 200], [503, 200], [504, 200]]) { + assert.deepEqual(afterRetry({ statuses }), none, String(statuses)); + } + assert.deepEqual(afterRetry({ statuses: [200], unanswered: true }), none); }); /** @@ -1906,6 +1909,24 @@ test( }, ); +for (const status of [408, 503, 504]) { + test( + `leaves a prompt's tokens unknown when a ${status} precedes a successful retry`, + { skip: supportsPi ? false : "Pi 1.0 requires Node 22.19 or newer" }, + async () => { + // The provider may have processed the prompt before answering with a + // timeout or a server error. + const [retried] = await promptsAgainstScriptedEndpoint( + [{ status }, { content: "after retry", usage: { prompt_tokens: 50, completion_tokens: 5 } }], + ["one"], + ); + + assert.equal(retried.text, "after retry"); + assert.equal(retried.usage, undefined); + }, + ); +} + test( "leaves a prompt's tokens unknown when a connection fails before any HTTP response", { skip: supportsPi ? false : "Pi 1.0 requires Node 22.19 or newer" }, From 54b5332e1b09c7ed6b10bc73598b37d59a03d182 Mon Sep 17 00:00:00 2001 From: mnajafian-nv Date: Tue, 20 Oct 2026 01:33:21 -0700 Subject: [PATCH 29/68] fix: keep Hermes cached input known when input is unknown The adapter set cached_input_tokens only together with input, so an invocation whose input was unknown, such as one with a missing cache-write counter, lost a cache-read count it did have. Cached input is now reported whenever the cache-read counter is known and reported. Signed-off-by: mnajafian-nv --- .../nemo_fabric_adapters/hermes/adapter.py | 13 ++++++---- tests/adapters/test_hermes_adapter.py | 24 +++++++++++++++++++ 2 files changed, 33 insertions(+), 4 deletions(-) diff --git a/adapters/python/hermes/src/nemo_fabric_adapters/hermes/adapter.py b/adapters/python/hermes/src/nemo_fabric_adapters/hermes/adapter.py index 77af05f6e..67673b1e3 100755 --- a/adapters/python/hermes/src/nemo_fabric_adapters/hermes/adapter.py +++ b/adapters/python/hermes/src/nemo_fabric_adapters/hermes/adapter.py @@ -1004,8 +1004,8 @@ def _invocation_usage( usage never holds a whole-session total. ``unreported`` names the counts ("input", "cached", "output") that a provider response of this invocation left out and Hermes counted as 0. Those counts are unknown, and so is the - total when input or output is. A count above what AgentUsage holds is - unknown too. + total when input or output is; cached input stays known without input. A + count above what AgentUsage holds is unknown too. """ skipped = skipped or {} change = { @@ -1025,9 +1025,14 @@ def _invocation_usage( ) if _is_count(input_tokens): fields["input_tokens"] = input_tokens - if "cached" not in unreported: - fields["cached_input_tokens"] = change["cache_read_tokens"] fields["input_tokens_include_cache"] = True + # Cache reads are their own count, known even when the input is not. + if ( + "cached" not in unreported + and "cache_read_tokens" in change + and _is_count(change["cache_read_tokens"]) + ): + fields["cached_input_tokens"] = change["cache_read_tokens"] if ( "output" not in unreported and "output_tokens" in change diff --git a/tests/adapters/test_hermes_adapter.py b/tests/adapters/test_hermes_adapter.py index a9050810b..7b4d2f9be 100644 --- a/tests/adapters/test_hermes_adapter.py +++ b/tests/adapters/test_hermes_adapter.py @@ -1579,6 +1579,30 @@ def test_invocation_usage_with_lower_or_missing_counters_leaves_them_unknown(): assert no_baseline is None +def test_invocation_usage_keeps_cached_input_when_input_is_unknown(): + before = dict.fromkeys(adapter._SESSION_TOKEN_KEYS, 0) + after = {**before, "input_tokens": 10, "cache_read_tokens": 2, "output_tokens": 3} + without_cache_writes = { + key: value for key, value in after.items() if key != "cache_write_tokens" + } + + unreported = adapter._invocation_usage( + after, before, unreported=frozenset({"input"}) + ) + no_write_counter = adapter._invocation_usage(without_cache_writes, before) + cache_only = adapter._invocation_usage( + after, before, unreported=frozenset({"input", "output"}) + ) + + for usage in (unreported, no_write_counter): + assert usage is not None + assert (usage.input_tokens, usage.total_tokens) == (None, None) + assert (usage.cached_input_tokens, usage.output_tokens) == (2, 3) + assert cache_only is not None + assert cache_only.cached_input_tokens == 2 + assert (cache_only.input_tokens, cache_only.output_tokens) == (None, None) + + def test_invocation_usage_above_the_contract_bound_is_unknown(): zero = dict.fromkeys(adapter._SESSION_TOKEN_KEYS, 0) near_max = {**zero, "input_tokens": (1 << 64) - 2, "output_tokens": 3} From f298d46b13fe627d5e313677c7f509773caa0d4d Mon Sep 17 00:00:00 2001 From: mnajafian-nv Date: Tue, 20 Oct 2026 01:42:37 -0700 Subject: [PATCH 30/68] fix: skip the Hermes model title call in Fabric runs Without Relay, Hermes upgrades each session's locally derived title with a model call (auxiliary.title_generation.model_upgrade_enabled, default on), and records its tokens outside the session counters the adapter reports. Fabric runs now keep the derived title and make no title model call; with Relay, titles stay disabled as before. Signed-off-by: mnajafian-nv --- .../hermes/configuration.py | 8 ++++++- tests/adapters/test_hermes_config_builder.py | 21 ++++++++++++++++++ tests/e2e/test_hermes_e2e.py | 22 +++++++++++++++++++ 3 files changed, 50 insertions(+), 1 deletion(-) diff --git a/adapters/python/hermes/src/nemo_fabric_adapters/hermes/configuration.py b/adapters/python/hermes/src/nemo_fabric_adapters/hermes/configuration.py index 39a0244c8..4e7875723 100644 --- a/adapters/python/hermes/src/nemo_fabric_adapters/hermes/configuration.py +++ b/adapters/python/hermes/src/nemo_fabric_adapters/hermes/configuration.py @@ -110,9 +110,15 @@ def build_hermes_config( ), } + # Hermes replaces a session's locally derived title with one from a model + # call unless auxiliary.title_generation.model_upgrade_enabled is false + # (agent/title_generator.py), and records that call's tokens outside the + # session counters the adapter reports. Keep the derived title instead. + title_generation: dict[str, bool] = {"model_upgrade_enabled": False} if relay_enabled: # Hermes makes an extra model call for titles, which skews Relay telemetry. - config["auxiliary"] = {"title_generation": {"enabled": False}} + title_generation["enabled"] = False + config["auxiliary"] = {"title_generation": title_generation} skill_dirs = ( [str(path) for path in agent_config.skills.paths] diff --git a/tests/adapters/test_hermes_config_builder.py b/tests/adapters/test_hermes_config_builder.py index b39eaf4e2..3fb53ecb2 100644 --- a/tests/adapters/test_hermes_config_builder.py +++ b/tests/adapters/test_hermes_config_builder.py @@ -40,6 +40,27 @@ def test_build_hermes_config_omits_unset_values_without_hermes_agent(): "default": "nvidia/test-model", } assert config["agent"] == {} + # The model title call's tokens bypass the counters the adapter reports. + assert config["auxiliary"] == { + "title_generation": {"model_upgrade_enabled": False} + } + + +def test_build_hermes_config_with_relay_disables_titles(): + config = AgentConfig.from_mapping( + { + "harness": {"settings": {}}, + "models": {"default": {"provider": "nvidia", "model": "nvidia/test-model"}}, + } + ) + + config = configuration.build_hermes_config( + config, workspace=".", relay_enabled=True + ) + + assert config["auxiliary"] == { + "title_generation": {"enabled": False, "model_upgrade_enabled": False} + } def test_write_hermes_config_round_trips_without_pyyaml( diff --git a/tests/e2e/test_hermes_e2e.py b/tests/e2e/test_hermes_e2e.py index ab36fb2f7..04a2fe108 100644 --- a/tests/e2e/test_hermes_e2e.py +++ b/tests/e2e/test_hermes_e2e.py @@ -298,6 +298,28 @@ async def test_hermes_usage_counts_the_iteration_limit_summary( ) +@pytest.mark.usefixtures("mock_nvidia_api_key") +async def test_hermes_makes_no_model_call_for_session_titles( + code_review_agent_dir: Path, + api_server: str, +): + os.environ["ADAPTER_PYTHON"] = sys.executable + config = hermes_config() + config.models["default"].base_url = f"{api_server}/v1" + + async with await Fabric().start_runtime( + config, base_dir=code_review_agent_dir + ) as runtime: + first = await runtime.invoke(input="first") + second = await runtime.invoke(input="second") + sent = requests.get(f"{api_server}/_requests", timeout=5).json() + + # A title call's tokens would bypass the session counters, so each + # invocation makes only its own turn's request. + assert (first["status"], second["status"]) == ("succeeded", "succeeded") + assert len(sent) == 2, [request["messages"][0]["content"][:60] for request in sent] + + @pytest.mark.usefixtures("mock_nvidia_api_key") async def test_hermes_usage_is_unknown_when_context_compression_runs( code_review_agent_dir: Path, From ea9dbca132d08974f20f29ae8f752792603111ec Mon Sep 17 00:00:00 2001 From: mnajafian-nv Date: Tue, 20 Oct 2026 01:45:22 -0700 Subject: [PATCH 31/68] test: cover the Hermes model-title setting in the existing config tests The config mapping and title tests still expected no auxiliary block outside Relay runs. They now expect model title upgrades off in every run and titles off with Relay, and the duplicate builder tests from the previous commit are removed. Signed-off-by: mnajafian-nv --- tests/adapters/test_hermes_adapter.py | 13 +++++++++--- tests/adapters/test_hermes_config_builder.py | 21 -------------------- 2 files changed, 10 insertions(+), 24 deletions(-) diff --git a/tests/adapters/test_hermes_adapter.py b/tests/adapters/test_hermes_adapter.py index 7b4d2f9be..66e0ca64c 100644 --- a/tests/adapters/test_hermes_adapter.py +++ b/tests/adapters/test_hermes_adapter.py @@ -578,6 +578,7 @@ def test_build_hermes_config_maps_fabric_config_to_hermes_config(): "cwd": "/workspace/repo", "timeout": 90, }, + "auxiliary": {"title_generation": {"model_upgrade_enabled": False}}, "skills": {"external_dirs": ["skills/review"]}, "mcp_servers": { "github": { @@ -1079,7 +1080,7 @@ async def test_runtime_reports_failed_oauth_mcp_authentication(monkeypatch): assert runtime._mcp_authentication_checked is False -def test_build_hermes_config_disables_session_title_generation_only_for_relay(): +def test_build_hermes_config_disables_model_titles_and_relay_titles(): agent_config = _agent_config( { "harness": {"settings": {}}, @@ -1092,8 +1093,14 @@ def test_build_hermes_config_disables_session_title_generation_only_for_relay(): agent_config, workspace=".", relay_enabled=True ) - assert "auxiliary" not in ordinary_config - assert relay_config["auxiliary"] == {"title_generation": {"enabled": False}} + # The model title call's tokens bypass the counters the adapter reports, + # so Hermes keeps its locally derived title; Relay runs have no titles. + assert ordinary_config["auxiliary"] == { + "title_generation": {"model_upgrade_enabled": False} + } + assert relay_config["auxiliary"] == { + "title_generation": {"enabled": False, "model_upgrade_enabled": False} + } def test_write_hermes_config_writes_file(tmp_path: Path): diff --git a/tests/adapters/test_hermes_config_builder.py b/tests/adapters/test_hermes_config_builder.py index 3fb53ecb2..b39eaf4e2 100644 --- a/tests/adapters/test_hermes_config_builder.py +++ b/tests/adapters/test_hermes_config_builder.py @@ -40,27 +40,6 @@ def test_build_hermes_config_omits_unset_values_without_hermes_agent(): "default": "nvidia/test-model", } assert config["agent"] == {} - # The model title call's tokens bypass the counters the adapter reports. - assert config["auxiliary"] == { - "title_generation": {"model_upgrade_enabled": False} - } - - -def test_build_hermes_config_with_relay_disables_titles(): - config = AgentConfig.from_mapping( - { - "harness": {"settings": {}}, - "models": {"default": {"provider": "nvidia", "model": "nvidia/test-model"}}, - } - ) - - config = configuration.build_hermes_config( - config, workspace=".", relay_enabled=True - ) - - assert config["auxiliary"] == { - "title_generation": {"enabled": False, "model_upgrade_enabled": False} - } def test_write_hermes_config_round_trips_without_pyyaml( From 72ccbb21d0eff1823408084d1fc9209b294af4e6 Mon Sep 17 00:00:00 2001 From: mnajafian-nv Date: Tue, 20 Oct 2026 01:51:37 -0700 Subject: [PATCH 32/68] fix: account for failed Hermes iteration-limit summaries The adapter saw a summary response only after _summary_text returned, so a summary request that raised, or a response whose text Hermes could not read, escaped accounting while Hermes returned its fallback text. The adapter now notes each summary attempt before its request, counts the response's usage before Hermes normalizes it, and leaves the counts unknown when an attempt never produced a response. Signed-off-by: mnajafian-nv --- .../nemo_fabric_adapters/hermes/adapter.py | 99 +++++++++++++++---- tests/_utils/mock_api_server.py | 18 +++- tests/e2e/test_hermes_e2e.py | 59 +++++++++++ 3 files changed, 156 insertions(+), 20 deletions(-) diff --git a/adapters/python/hermes/src/nemo_fabric_adapters/hermes/adapter.py b/adapters/python/hermes/src/nemo_fabric_adapters/hermes/adapter.py index 67673b1e3..1d8b4ea9f 100755 --- a/adapters/python/hermes/src/nemo_fabric_adapters/hermes/adapter.py +++ b/adapters/python/hermes/src/nemo_fabric_adapters/hermes/adapter.py @@ -78,6 +78,9 @@ def __init__(self) -> None: # Context-compression summary calls, whose usage Hermes keeps out of # its session counters. self._compression_calls = 0 + # Iteration-limit summary attempts started minus summary responses + # observed; an attempt whose request raised leaves it above zero. + self._unanswered_summary_attempts = 0 async def start(self, payload: dict[str, Any]) -> None: if self._started: @@ -304,6 +307,7 @@ def run_hermes_turn() -> tuple[dict[str, Any], str]: unreported_before = dict(self._unreported_counts) skipped_before = dict(self._skipped_response_tokens) compression_calls_before = self._compression_calls + unanswered_summaries_before = self._unanswered_summary_attempts invoke_task = asyncio.create_task(asyncio.to_thread(run_hermes_turn)) self._active_invoke_task = invoke_task @@ -400,11 +404,13 @@ def clear_active_invoke_task( for count, responses in self._unreported_counts.items() if responses != unreported_before[count] ) - # A compression summary call's tokens are not in the counters, so - # every count it adds to is an undercount. + # A compression summary call's tokens are not in the counters, and + # a summary request that raised before its response was observed + # may have been billed, so every count is an undercount. | ( frozenset(_REPORTED_COUNTS) if self._compression_calls != compression_calls_before + or self._unanswered_summary_attempts != unanswered_summaries_before else frozenset() ), ) @@ -492,6 +498,8 @@ def count_unreported( api_mode: Any, ) -> None: usage = getattr(response, "usage", None) + if handler == _ITERATION_SUMMARY_HANDLER: + self._unanswered_summary_attempts -= 1 if handler == "record_response_usage": # A response without usage is counted by the compressor hook. # An empty usage object is a report with every count left out. @@ -506,7 +514,11 @@ def count_unreported( if usage: self._add_skipped_response_tokens(usage, provider, api_mode) + def count_summary_attempt() -> None: + self._unanswered_summary_attempts += 1 + _USAGE_OBSERVERS[agent] = count_unreported + _SUMMARY_ATTEMPT_OBSERVERS[agent] = count_summary_attempt def _add_skipped_response_tokens( self, usage: Any, provider: Any, api_mode: Any @@ -752,6 +764,10 @@ def _invoke_hermes_turn( _USAGE_OBSERVERS: weakref.WeakKeyDictionary[ Any, Callable[[str, Any, Any, Any, Any], None] ] = weakref.WeakKeyDictionary() +# Called before each iteration-limit summary request of the agent. +_SUMMARY_ATTEMPT_OBSERVERS: weakref.WeakKeyDictionary[Any, Callable[[], None]] = ( + weakref.WeakKeyDictionary() +) # The functions ``agent.turn_response_check.check_api_response`` calls with # (agent, response, ...) for a provider response, by the module defining # each. record_response_usage folds the usage; on an invalid response, a @@ -766,12 +782,16 @@ def _invoke_hermes_turn( } -def _observed_response_handler(name: str, handler: Callable[..., Any]) -> Any: +def _observed_response_handler( + name: str, handler: Callable[..., Any], *, observe_first: bool = False +) -> Any: """Wrap a Hermes response handler to show each call to the agent's observer. Hermes passes the response positionally to some handlers and by keyword to others, so the wrapper forwards the call unchanged and reads the - handler's ``response`` argument from its signature. + handler's ``response`` argument from its signature. With + ``observe_first``, the observer sees the response, with no outcome, + before the handler runs, so a handler that raises cannot hide its usage. """ signature = inspect.signature(handler) @@ -780,6 +800,12 @@ def observed(agent: Any, *args: Any, **kwargs: Any) -> Any: # route that served this response before the handler runs. provider = getattr(agent, "provider", None) api_mode = getattr(agent, "api_mode", None) + if observe_first: + observer = _USAGE_OBSERVERS.get(agent) + if observer is not None: + response = signature.bind(agent, *args, **kwargs).arguments["response"] + observer(name, response, None, provider, api_mode) + return handler(agent, *args, **kwargs) outcome = handler(agent, *args, **kwargs) observer = _USAGE_OBSERVERS.get(agent) if observer is not None: @@ -829,41 +855,52 @@ def _install_usage_observer() -> None: # The function each of Hermes' iteration-limit summary attempts passes its -# provider response to (agent/chat_completion_helpers.py), and the attempt -# builders that call it. +# provider response to (agent/chat_completion_helpers.py), the attempt +# builders that call it, and where ``handle_max_iterations`` looks them up. _ITERATION_SUMMARY_HANDLER = "_summary_text" _ITERATION_SUMMARY_ATTEMPTS = ( "_chat_summary_attempt", "_anthropic_summary_attempt", "_codex_summary_attempt", ) +_ITERATION_SUMMARY_BUILDERS = "_SUMMARY_ATTEMPT_BUILDERS" +_ITERATION_SUMMARY_DEFAULT_BUILDER = "_chat_summary_attempt" def _observe_iteration_summary() -> None: - """Route Hermes' iteration-limit summary responses through the observers. + """Route Hermes' iteration-limit summary requests through the observers. When a turn reaches ``max_iterations``, ``handle_max_iterations`` makes an extra summary request and never adds its usage to the session counters. Every summary attempt hands its response to ``_summary_text``, a global of ``agent.chat_completion_helpers``, so the adapter wraps that function once - per process and counts the response like one Hermes skipped. Start fails - if an attempt no longer calls it by that name. + per process and counts the response like one Hermes skipped, before + Hermes normalizes it. Hermes picks each attempt's builder from + ``_SUMMARY_ATTEMPT_BUILDERS`` or the ``_chat_summary_attempt`` global, so + the adapter also wraps the builders' attempts to note each request before + it is sent: an attempt that raises before its response reaches + ``_summary_text`` leaves the invocation's counts unknown. Start fails if + Hermes no longer looks them up by those names. """ - from agent import chat_completion_helpers + from agent import chat_completion_helpers as helpers - current = getattr(chat_completion_helpers, _ITERATION_SUMMARY_HANDLER, None) + current = getattr(helpers, _ITERATION_SUMMARY_HANDLER, None) if getattr(current, "nemo_fabric_observed", False): return + builders = getattr(helpers, _ITERATION_SUMMARY_BUILDERS, None) + handle_max_iterations = getattr(helpers, "handle_max_iterations", None) + attempts = {name: getattr(helpers, name, None) for name in _ITERATION_SUMMARY_ATTEMPTS} if ( current is None or not all( - _inner_function_calls( - getattr(chat_completion_helpers, attempt, None), - _ITERATION_SUMMARY_HANDLER, - ) - for attempt in _ITERATION_SUMMARY_ATTEMPTS + _inner_function_calls(attempt, _ITERATION_SUMMARY_HANDLER) + for attempt in attempts.values() ) or "response" not in inspect.signature(current).parameters + or not isinstance(builders, dict) + or not all(builder in attempts.values() for builder in builders.values()) + or not {_ITERATION_SUMMARY_BUILDERS, _ITERATION_SUMMARY_DEFAULT_BUILDER} + <= set(getattr(getattr(handle_max_iterations, "__code__", None), "co_names", ())) ): raise lifecycle.LifecycleError( "hermes_usage_accounting_unsupported", @@ -871,10 +908,36 @@ def _observe_iteration_summary() -> None: "outside agent.chat_completion_helpers._summary_text", ) setattr( - chat_completion_helpers, + helpers, _ITERATION_SUMMARY_HANDLER, - _observed_response_handler(_ITERATION_SUMMARY_HANDLER, current), + _observed_response_handler( + _ITERATION_SUMMARY_HANDLER, current, observe_first=True + ), ) + for mode, builder in list(builders.items()): + builders[mode] = _counted_summary_attempts(builder) + setattr( + helpers, + _ITERATION_SUMMARY_DEFAULT_BUILDER, + _counted_summary_attempts(attempts[_ITERATION_SUMMARY_DEFAULT_BUILDER]), + ) + + +def _counted_summary_attempts(build: Callable[..., Any]) -> Callable[..., Any]: + """Wrap a summary attempt builder so each attempt notes its request first.""" + + def build_counted(agent: Any, *args: Any, **kwargs: Any) -> Any: + attempt = build(agent, *args, **kwargs) + + def counted_attempt(*attempt_args: Any, **attempt_kwargs: Any) -> Any: + note_attempt = _SUMMARY_ATTEMPT_OBSERVERS.get(agent) + if note_attempt is not None: + note_attempt() + return attempt(*attempt_args, **attempt_kwargs) + + return counted_attempt + + return build_counted def _inner_function_calls(function: Any, name: str) -> bool: diff --git a/tests/_utils/mock_api_server.py b/tests/_utils/mock_api_server.py index 3661775af..1e6b339af 100644 --- a/tests/_utils/mock_api_server.py +++ b/tests/_utils/mock_api_server.py @@ -27,7 +27,8 @@ def mock_api_server(port: int) -> Iterator[str]: Use the /_scenario endpoint to configure a status code, ordered tool calls, or the chat-completion usage object and finish reason for subsequent requests. Its ordered ``replies`` override the usage or finish reason of the next chat completions, one - reply each, or make one invalid. + reply each, make one invalid, reject one with its ``status_code``, or send one + with ``no_choices`` (an empty ``choices`` list). Use the /_reset endpoint to clear all retained request and scenario state. Args: @@ -164,6 +165,17 @@ async def mcp(request: Request): async def chat_completions(request: Request): payload = await request.json() app.state.requests.append(payload) + if app.state.replies and "status_code" in app.state.replies[0]: + rejected = app.state.replies.pop(0) + return JSONResponse( + status_code=int(rejected["status_code"]), + content={ + "error": { + "message": f"configured status {rejected['status_code']}", + "type": "invalid_request_error", + } + }, + ) if app.state.status_code != 200: return JSONResponse( status_code=app.state.status_code, @@ -211,7 +223,9 @@ async def chat_completions(request: Request): "object": "chat.completion", "created": 0, "model": payload.get("model", "fabric-echo"), - "choices": [ + "choices": [] + if reply.get("no_choices", False) + else [ { "index": 0, "message": {"role": "assistant", "content": content}, diff --git a/tests/e2e/test_hermes_e2e.py b/tests/e2e/test_hermes_e2e.py index 04a2fe108..7e17c0364 100644 --- a/tests/e2e/test_hermes_e2e.py +++ b/tests/e2e/test_hermes_e2e.py @@ -320,6 +320,65 @@ async def test_hermes_makes_no_model_call_for_session_titles( assert len(sent) == 2, [request["messages"][0]["content"][:60] for request in sent] +@pytest.mark.usefixtures("mock_nvidia_api_key") +@pytest.mark.parametrize( + ("summary_reply", "expected_usage"), + [ + # Hermes cannot read the text of a response without choices, but the + # response reported its usage. + ( + {"usage": _chat_usage(600, 20), "no_choices": True}, + { + "input_tokens": 640, + "input_tokens_include_cache": True, + "output_tokens": 532, + "total_tokens": 1172, + "metadata": {}, + }, + ), + # The request raised before any response reached Hermes' summary text. + ({"status_code": 400}, None), + ], + ids=["unreadable", "rejected"], +) +async def test_hermes_usage_accounts_for_a_failed_iteration_limit_summary( + code_review_agent_dir: Path, + api_server: str, + summary_reply: dict[str, object], + expected_usage: dict[str, object] | None, +): + os.environ["ADAPTER_PYTHON"] = sys.executable + config = hermes_config() + config.models["default"].base_url = f"{api_server}/v1" + config.runtime.max_turns = 1 + + async with await Fabric().start_runtime( + config, base_dir=code_review_agent_dir + ) as runtime: + await runtime.invoke(input="first") + scenario_response = requests.post( + f"{api_server}/_scenario", + json={ + "replies": [ + {"finish_reason": "length", "usage": _chat_usage(40, 512)}, + summary_reply, + ] + }, + timeout=5, + ) + scenario_response.raise_for_status() + summarized = await runtime.invoke(input="second") + sent = requests.get(f"{api_server}/_requests", timeout=5).json() + + # Hermes answers with its fallback text when the summary fails, without + # retrying it. + assert summarized["status"] == "succeeded", summarized.to_mapping() + assert len(sent) == 3, sent + assert summarized.to_mapping().get("usage") == expected_usage, ( + summarized.to_mapping() + ) + + @pytest.mark.usefixtures("mock_nvidia_api_key") async def test_hermes_usage_is_unknown_when_context_compression_runs( code_review_agent_dir: Path, From 5fe84001cbbb4d7bd74514f4624268fee888af52 Mon Sep 17 00:00:00 2001 From: mnajafian-nv Date: Tue, 20 Oct 2026 01:52:14 -0700 Subject: [PATCH 33/68] docs: describe the review fixes to Hermes and Pi usage State that Pi infers 0 only after rate limits, that cache reads stay known without input in both adapters, that failed Hermes iteration-limit summaries are accounted for, and that Fabric runs make no Hermes title model call. Signed-off-by: mnajafian-nv --- adapters/python/hermes/README.md | 18 +++++++++++------ adapters/typescript/pi/README.md | 34 +++++++++++++++++--------------- 2 files changed, 30 insertions(+), 22 deletions(-) diff --git a/adapters/python/hermes/README.md b/adapters/python/hermes/README.md index 3af96d2b1..0e0ee8376 100644 --- a/adapters/python/hermes/README.md +++ b/adapters/python/hermes/README.md @@ -124,13 +124,17 @@ never adds its tokens to the next one. A counter that is missing, that is not a non-negative integer below 2^64, or that went down because Hermes reset its session, is left out for that invocation. Hermes records auxiliary calls, such as context-compression summaries and session titles, separately from these -counters. Session title calls are not included. An invocation in which Hermes +counters. The adapter sets `auxiliary.title_generation.model_upgrade_enabled` to +false, so Hermes keeps its locally derived session title and makes no title +model call; with Relay enabled, title generation is off entirely. An invocation in which Hermes ran a context-compression summary call, or in which a plugin context engine compressed, has no `usage`, because the compression calls' tokens are missing from the counters. Hermes also leaves out of its counters the summary request it makes when a turn reaches its iteration limit. The adapter adds the counts that -this summary response reported, and an invocation whose summary response -carried no usage has no `usage`. +this summary response reported, before Hermes reads its text, so they count +even when Hermes cannot read the response and answers with its fallback text. +An invocation whose summary response carried no usage, or whose summary +request failed before any response, has no `usage`. Hermes adds no tokens for a provider response that carried no usage, so an invocation with such a response has no `usage`. For a response it rejects as @@ -144,7 +148,8 @@ billed usage is not observed and the invocation can undercount. Hermes reads a count that a response's usage left out as 0, so the adapter checks each response's usage before Hermes normalizes it. A reported zero still counts as 0. When a response left out its input count, the invocation's `usage` leaves -out `input_tokens`, `cached_input_tokens`, and `total_tokens`. Anthropic +out `input_tokens` and `total_tokens`, and keeps `cached_input_tokens` when the +cache-read count is known. Anthropic Messages usage reports input without the cache, so its input count also needs both cache fields. When a response left out its output count, the `usage` leaves out `output_tokens` and `total_tokens`. When a response sent no cache @@ -156,8 +161,9 @@ Anthropic transport when a stream omits usage, fill missing counts with 0 before the adapter sees them. Hermes' Codex app-server runtime records usage outside this path. Hermes passes a response's raw usage to no per-agent hook, so the adapter wraps `record_response_usage` and the three handlers for -invalid, truncated, and refused responses in `agent.turn_response_check` once -per process. The runtime fails to start if the installed Hermes no longer calls +invalid, truncated, and refused responses in `agent.turn_response_check`, and +the iteration-limit summary's `_summary_text` and attempt builders in +`agent.chat_completion_helpers`, once per process. The runtime fails to start if the installed Hermes no longer calls them by those names. ## Execution Model diff --git a/adapters/typescript/pi/README.md b/adapters/typescript/pi/README.md index aad02a7c1..5736c6b3b 100644 --- a/adapters/typescript/pi/README.md +++ b/adapters/typescript/pi/README.md @@ -226,8 +226,8 @@ Pi fills a count that the provider did not report with 0, so the adapter reads the provider's raw stream events for each assistant message and compaction summary call and keeps only the counts the provider reported. A reported zero counts as 0, including a report of all zeros. When a call left out its input -count, the result leaves out `input_tokens`, `cached_input_tokens`, and -`total_tokens`. Anthropic Messages usage reports input without the cache, so +count, the result leaves out `input_tokens` and `total_tokens`, and keeps +`cached_input_tokens` when every call reported its cache reads. Anthropic Messages usage reports input without the cache, so its input count also needs both cache fields. Bedrock Converse also reports input without the cache, but its API documents the cache fields as optional, so its input count needs only `inputTokens` rather than going unknown on every @@ -239,19 +239,20 @@ Pi's built-in ones reports no counts the adapter can check, so its calls leave the result without `usage`. A failed call counts as 0 tokens only when the provider answered each of its -attempts with a 4xx or 5xx HTTP status, such as a rate limit (429), an overload -or a server error, before sending a successful response or any stream event. -A call that Pi retries after such an error keeps the result's `usage` known. -The adapter reads that status through the request's `fetch`, which the -`anthropic-messages`, `azure-openai-responses`, `mistral-conversations`, -`openai-codex-responses`, `openai-completions`, `openai-responses` and -`pi-messages` protocols use. A failed call of any other protocol, including the -Google protocols and Bedrock, or a Codex call sent over a WebSocket, leaves its -counts unknown. A call with an attempt that got no HTTP response, such as a -connection failure or a timeout before headers, leaves its counts unknown even -when a later attempt succeeds, because the provider may have billed the prompt. -A call that fails after its stream started can have billed tokens that were -never reported, so its counts stay unknown too. +attempts with a rate limit (429) before sending a successful response or any +stream event, so a call that Pi retries after rate limits keeps the result's +`usage` known. The adapter reads each attempt's status through the request's +`fetch`, which the `anthropic-messages`, `azure-openai-responses`, +`mistral-conversations`, `openai-codex-responses`, `openai-completions`, +`openai-responses` and `pi-messages` protocols use. A failed call of any other +protocol, including the Google protocols and Bedrock, or a Codex call sent over +a WebSocket, leaves its counts unknown. An attempt that got any other error +status, such as a request timeout (408) or a server error (5xx), or that got no +HTTP response, such as a connection failure or a timeout before headers, leaves +its call's counts unknown even when a later attempt succeeds, because the +provider may have processed and billed the prompt. A call that fails after its +stream started can have billed tokens that were never reported, so its counts +stay unknown too. When a compaction cuts inside a turn, Pi summarizes the earlier history and the start of that turn in separate calls and reports only their combined usage. The @@ -262,7 +263,8 @@ combined usage; otherwise the result has no `usage`. A tool result's `usage` counts each count it holds: its input needs `input`, `cacheRead`, and `cacheWrite`, its cache reads need `cacheRead`, and its output needs `output`. A count the tool left out, or that is not a non-negative -safe integer, is left out of the result. A tool result without a `usage` field adds +safe integer, is left out of the result, and the other counts stay, so a +report without `cacheWrite` keeps its cache reads and output. A tool result without a `usage` field adds nothing and leaves the result's `usage` unchanged. An extension's `message_end` handler that replaces an assistant message's usage does not change which counts the provider reported. From c5e9f5eaaf6c0f2af95a6242f3589cc8804c9226 Mon Sep 17 00:00:00 2001 From: mnajafian-nv Date: Tue, 20 Oct 2026 02:39:44 -0700 Subject: [PATCH 34/68] fix: leave a Pi Codex call unknown after a failed WebSocket attempt Pi's Codex transport sends the request over a WebSocket first, which the adapter's fetch never sees. When that attempt fails, Pi records a provider_transport_failure diagnostic and falls back to HTTP, so a successful or rate-limited fallback looked like the call's only attempt and its counts stayed known. A transport-failure diagnostic on the final message now makes the call's counts unknown. Signed-off-by: mnajafian-nv --- adapters/typescript/pi/src/pi-sdk.ts | 23 ++- adapters/typescript/pi/test/pi-sdk.test.mjs | 149 ++++++++++++++++++++ 2 files changed, 169 insertions(+), 3 deletions(-) diff --git a/adapters/typescript/pi/src/pi-sdk.ts b/adapters/typescript/pi/src/pi-sdk.ts index af3f34b95..dc3f131ba 100644 --- a/adapters/typescript/pi/src/pi-sdk.ts +++ b/adapters/typescript/pi/src/pi-sdk.ts @@ -513,7 +513,10 @@ const observedFetch: typeof globalThis.fetch = async (input, init) => { * ambiguous attempt reports none, even when a later attempt succeeded: an * attempt that got no HTTP response, such as a connection error or a timeout * before headers, or that got any error status other than a rate limit (429), - * such as a 408 or a 5xx, may have been billed. A call that ended in an error + * such as a 408 or a 5xx, may have been billed. So may a Codex WebSocket + * attempt that failed, which the fetch never sees: Pi records it on the final + * message as a `provider_transport_failure` diagnostic before falling back to + * HTTP (`@earendil-works/pi-ai` 1.0.3, `api/openai-codex-responses.js`). A call that ended in an error * reports all its counts as 0 only when its protocol sends requests through * the observed fetch, the provider answered every attempt with a 429, no * attempt was accepted, no raw event arrived, and Pi's usage is all zeros. @@ -522,11 +525,12 @@ const observedFetch: typeof globalThis.fetch = async (input, init) => { export function endedCallCounts( api: string, seen: ReadonlySet, - message: { stopReason?: unknown; usage?: unknown }, + message: { stopReason?: unknown; usage?: unknown; diagnostics?: unknown }, observation: CallObservation, ): ReportedCounts { const ambiguous = observation.unanswered || + hasTransportFailure(message.diagnostics) || observation.statuses.some((status) => (status < 200 || status >= 300) && !UNBILLED_REJECTIONS.has(status)); if (ambiguous) { return NONE_REPORTED; @@ -546,6 +550,19 @@ export function endedCallCounts( return rejected ? ALL_REPORTED : reportedCountsFrom(api, seen); } +/** Whether a message's diagnostics record a failed transport attempt. */ +function hasTransportFailure(diagnostics: unknown): boolean { + return ( + Array.isArray(diagnostics) && + diagnostics.some( + (diagnostic) => + typeof diagnostic === "object" && + diagnostic !== null && + (diagnostic as { type?: unknown }).type === "provider_transport_failure", + ) + ); +} + /** * Record which counts each model call's provider reported, keyed by the call's * final assistant message and its usage object, and pass each ended call's @@ -562,7 +579,7 @@ export function endedCallCounts( * usage entry has no record and stays unknown. The record is copied when the * call ends so a later warm-up cannot change it. */ -function recordReportedCounts( +export function recordReportedCounts( agent: AgentSession["agent"], records: WeakMap, onCallEnd: (usage: object) => void, diff --git a/adapters/typescript/pi/test/pi-sdk.test.mjs b/adapters/typescript/pi/test/pi-sdk.test.mjs index 0e33c7b76..a0f79992c 100644 --- a/adapters/typescript/pi/test/pi-sdk.test.mjs +++ b/adapters/typescript/pi/test/pi-sdk.test.mjs @@ -6,6 +6,7 @@ import { access, mkdir, mkdtemp, readFile, realpath, rm, writeFile } from "node: import { tmpdir } from "node:os"; import { join } from "node:path"; import test from "node:test"; +import { createHash, randomUUID } from "node:crypto"; import { createServer } from "node:http"; import { @@ -16,6 +17,7 @@ import { endedCallCounts, modelAwareCompactionReserveTokens, PiSdkSessionFactory, + recordReportedCounts, reportedCountsFrom, resolveCustomTools, summaryCallCount, @@ -1943,6 +1945,153 @@ test( }, ); +/** One unmasked WebSocket text frame, as a server sends it. */ +function webSocketTextFrame(text) { + const payload = Buffer.from(text); + let header; + if (payload.length < 126) { + header = Buffer.from([0x81, payload.length]); + } else if (payload.length < 65_536) { + header = Buffer.from([0x81, 126, payload.length >> 8, payload.length & 0xff]); + } else { + header = Buffer.alloc(10); + header[0] = 0x81; + header[1] = 127; + header.writeBigUInt64BE(BigInt(payload.length), 2); + } + return Buffer.concat([header, payload]); +} + +/** The Responses events of a completed reply that reports `usage`. */ +function codexReplyEvents(usage) { + const part = { type: "output_text", text: "done", annotations: [] }; + const item = { id: "msg_usage", type: "message", role: "assistant", status: "completed", content: [part] }; + const response = { id: "resp_usage", object: "response", created_at: 0, status: "completed", model: "gpt-test", output: [item], usage }; + return [ + { type: "response.created", response: { ...response, status: "in_progress", output: [] } }, + { type: "response.output_item.added", output_index: 0, item: { ...item, status: "in_progress", content: [] } }, + { type: "response.content_part.added", item_id: item.id, output_index: 0, content_index: 0, part: { ...part, text: "" } }, + { type: "response.output_text.delta", item_id: item.id, output_index: 0, content_index: 0, delta: "done" }, + { type: "response.output_text.done", item_id: item.id, output_index: 0, content_index: 0, text: "done" }, + { type: "response.content_part.done", item_id: item.id, output_index: 0, content_index: 0, part }, + { type: "response.output_item.done", output_index: 0, item }, + { type: "response.completed", response }, + ]; +} + +/** + * Run one model call through Pi's Codex transport, wrapped by the adapter's + * recorder, against a local endpoint. `webSocket` is "reply" to answer the + * WebSocket request or "drop" to close the connection after receiving it + * and before any event; `http` is the HTTP fallback's status, with 200 + * replying over server-sent events. Returns the counts the adapter recorded + * for the call and the requests each transport received. + */ +async function codexCallCounts({ webSocket, http }) { + const usage = { input_tokens: 50, input_tokens_details: { cached_tokens: 0 }, output_tokens: 5, total_tokens: 55 }; + const received = { webSocket: 0, http: 0 }; + const sockets = new Set(); + const server = createServer(async (request, response) => { + for await (const _chunk of request); + received.http += 1; + if (http !== 200) { + response.writeHead(http, { "content-type": "application/json" }); + response.end(JSON.stringify({ error: { message: "rate limited", type: "rate_limit_error" } })); + return; + } + response.writeHead(200, { "content-type": "text/event-stream" }); + response.end(codexReplyEvents(usage).map((event) => `data: ${JSON.stringify(event)}\n\n`).join("")); + }); + server.on("upgrade", (request, socket) => { + sockets.add(socket); + const accept = createHash("sha1") + .update(`${request.headers["sec-websocket-key"]}258EAFA5-E914-47DA-95CA-C5AB0DC85B11`) + .digest("base64"); + socket.write( + `HTTP/1.1 101 Switching Protocols\r\nUpgrade: websocket\r\nConnection: Upgrade\r\nSec-WebSocket-Accept: ${accept}\r\n\r\n`, + ); + // The first client frame is Pi's response.create request. + socket.once("data", () => { + received.webSocket += 1; + if (webSocket === "drop") { + socket.destroy(); + return; + } + for (const event of codexReplyEvents(usage)) { + socket.write(webSocketTextFrame(JSON.stringify(event))); + } + }); + }); + const baseUrl = await listen(server); + const sessionId = randomUUID(); + const { closeOpenAICodexWebSocketSessions, stream } = await import("@earendil-works/pi-ai/api/openai-codex-responses"); + try { + const records = new WeakMap(); + const agent = { streamFunction: stream }; + recordReportedCounts(agent, records, () => undefined); + const claims = Buffer.from(JSON.stringify({ "https://api.openai.com/auth": { chatgpt_account_id: "account" } })).toString("base64url"); + const model = { + id: "gpt-test", + name: "gpt-test", + api: "openai-codex-responses", + provider: "openai-codex", + baseUrl, + reasoning: false, + input: ["text"], + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 }, + contextWindow: 100_000, + maxTokens: 1_000, + }; + const response = await agent.streamFunction( + model, + { messages: [{ role: "user", content: "one", timestamp: 0 }] }, + { apiKey: `header.${claims}.signature`, sessionId, transport: "auto" }, + ); + const message = await response.result(); + await new Promise((resolve) => setImmediate(resolve)); + return { counts: records.get(message), message, received }; + } finally { + // Pi keeps a session's WebSocket open for reuse until an idle timer ends. + closeOpenAICodexWebSocketSessions(sessionId); + for (const socket of sockets) { + socket.destroy(); + } + await close(server); + } +} + +test( + "keeps a Codex WebSocket reply's counts known", + { skip: supportsPi ? false : "Pi 1.0 requires Node 22.19 or newer" }, + async () => { + const { counts, message, received } = await codexCallCounts({ webSocket: "reply", http: 200 }); + + assert.deepEqual(received, { webSocket: 1, http: 0 }); + assert.equal(message.stopReason, "stop"); + assert.deepEqual(counts, { input: true, cacheRead: true, output: true }); + }, +); + +for (const [name, http, stopReason] of [ + ["a successful HTTP fallback", 200, "stop"], + ["a rate-limited HTTP fallback", 429, "error"], +]) { + test( + `leaves a Codex call's counts unknown when its WebSocket request failed before ${name}`, + { skip: supportsPi ? false : "Pi 1.0 requires Node 22.19 or newer" }, + async () => { + // The provider received the WebSocket request and may have billed it, + // but the fetch only sees the HTTP fallback. + const { counts, message, received } = await codexCallCounts({ webSocket: "drop", http }); + + assert.deepEqual(received, { webSocket: 1, http: 1 }); + assert.equal(message.stopReason, stopReason); + assert.ok(message.diagnostics.some((diagnostic) => diagnostic.type === "provider_transport_failure")); + assert.deepEqual(counts, { input: false, cacheRead: false, output: false }); + }, + ); +} + test( "leaves a prompt's tokens unknown when a stream fails after it started", { skip: supportsPi ? false : "Pi 1.0 requires Node 22.19 or newer" }, From ce14531d02b92e83702b746520e043f7fda91e03 Mon Sep 17 00:00:00 2001 From: mnajafian-nv Date: Tue, 20 Oct 2026 02:41:25 -0700 Subject: [PATCH 35/68] fix: unregister the Hermes summary observer when the runtime stops stop() removed the agent's usage observer but left its iteration-limit summary observer registered, so the stopped runtime's counter stayed reachable from the process-wide registry until the agent was collected. Signed-off-by: mnajafian-nv --- .../python/hermes/src/nemo_fabric_adapters/hermes/adapter.py | 1 + tests/adapters/test_hermes_adapter.py | 2 ++ 2 files changed, 3 insertions(+) diff --git a/adapters/python/hermes/src/nemo_fabric_adapters/hermes/adapter.py b/adapters/python/hermes/src/nemo_fabric_adapters/hermes/adapter.py index 1d8b4ea9f..8c2d4c636 100755 --- a/adapters/python/hermes/src/nemo_fabric_adapters/hermes/adapter.py +++ b/adapters/python/hermes/src/nemo_fabric_adapters/hermes/adapter.py @@ -652,6 +652,7 @@ async def stop(self) -> None: session_db = self._session_db if agent is not None: _USAGE_OBSERVERS.pop(agent, None) + _SUMMARY_ATTEMPT_OBSERVERS.pop(agent, None) had_mcp_servers = bool(self._hermes_config.get("mcp_servers")) had_relay_plugin = self._relay_plugin_config_path is not None self._agent = None diff --git a/tests/adapters/test_hermes_adapter.py b/tests/adapters/test_hermes_adapter.py index 66e0ca64c..794749619 100644 --- a/tests/adapters/test_hermes_adapter.py +++ b/tests/adapters/test_hermes_adapter.py @@ -2057,10 +2057,12 @@ async def test_hermes_stop_unregisters_the_agent_usage_observer(): agent = _UsageRecordingAgent([]) runtime = _usage_recording_runtime(agent) assert agent in adapter._USAGE_OBSERVERS + assert agent in adapter._SUMMARY_ATTEMPT_OBSERVERS await runtime.stop() assert agent not in adapter._USAGE_OBSERVERS + assert agent not in adapter._SUMMARY_ATTEMPT_OBSERVERS @pytest.mark.parametrize("drift", ["replaced", "not_called"]) From 779f894ecb5867e993ab8b9c191b93c5edba64bb Mon Sep 17 00:00:00 2001 From: mnajafian-nv Date: Tue, 20 Oct 2026 02:41:26 -0700 Subject: [PATCH 36/68] docs: describe the Pi Codex WebSocket fallback rule State that a Codex call whose WebSocket attempt failed before an HTTP fallback keeps its counts unknown. Signed-off-by: mnajafian-nv --- adapters/typescript/pi/README.md | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/adapters/typescript/pi/README.md b/adapters/typescript/pi/README.md index 5736c6b3b..bb9faebde 100644 --- a/adapters/typescript/pi/README.md +++ b/adapters/typescript/pi/README.md @@ -246,7 +246,10 @@ stream event, so a call that Pi retries after rate limits keeps the result's `mistral-conversations`, `openai-codex-responses`, `openai-completions`, `openai-responses` and `pi-messages` protocols use. A failed call of any other protocol, including the Google protocols and Bedrock, or a Codex call sent over -a WebSocket, leaves its counts unknown. An attempt that got any other error +a WebSocket, leaves its counts unknown. When a Codex WebSocket attempt fails and +Pi falls back to HTTP, Pi records the failure on the final message, and the +call's counts stay unknown even when the fallback succeeds, because the +provider may have received and billed the WebSocket request. An attempt that got any other error status, such as a request timeout (408) or a server error (5xx), or that got no HTTP response, such as a connection failure or a timeout before headers, leaves its call's counts unknown even when a later attempt succeeds, because the From b78d66ed031f63bcef79dcbd5028f5acada77930 Mon Sep 17 00:00:00 2001 From: mnajafian-nv Date: Tue, 20 Oct 2026 03:15:11 -0700 Subject: [PATCH 37/68] docs: state the Pi Codex retry and mid-stream usage rules exactly The endedCallCounts comment now says that a failed Codex WebSocket attempt counts as ambiguous whenever Pi records a provider_transport_failure diagnostic, including before the HTTP fallback, and that the two WebSocket retries Pi makes without a diagnostic are provider rejections before the model runs. The README now says that a call failing mid-stream keeps the input the provider already reported. Signed-off-by: mnajafian-nv --- adapters/typescript/pi/README.md | 17 +++++++++-------- adapters/typescript/pi/src/pi-sdk.ts | 11 +++++++---- 2 files changed, 16 insertions(+), 12 deletions(-) diff --git a/adapters/typescript/pi/README.md b/adapters/typescript/pi/README.md index bb9faebde..5e6acb82a 100644 --- a/adapters/typescript/pi/README.md +++ b/adapters/typescript/pi/README.md @@ -248,14 +248,15 @@ stream event, so a call that Pi retries after rate limits keeps the result's protocol, including the Google protocols and Bedrock, or a Codex call sent over a WebSocket, leaves its counts unknown. When a Codex WebSocket attempt fails and Pi falls back to HTTP, Pi records the failure on the final message, and the -call's counts stay unknown even when the fallback succeeds, because the -provider may have received and billed the WebSocket request. An attempt that got any other error -status, such as a request timeout (408) or a server error (5xx), or that got no -HTTP response, such as a connection failure or a timeout before headers, leaves -its call's counts unknown even when a later attempt succeeds, because the -provider may have processed and billed the prompt. A call that fails after its -stream started can have billed tokens that were never reported, so its counts -stay unknown too. +call's counts stay unknown even when the fallback succeeds, because the provider +may have received and billed the WebSocket request. An attempt that got any +other error status, such as a request timeout (408) or a server error (5xx), or +that got no HTTP response, such as a connection failure or a timeout before +headers, leaves its call's counts unknown even when a later attempt succeeds, +because the provider may have processed and billed the prompt. A call that fails +after its stream started can have billed tokens that were never reported, so the +counts it had not reported yet stay unknown; input that the provider already +reported, such as Anthropic's `message_start` usage, stays known. When a compaction cuts inside a turn, Pi summarizes the earlier history and the start of that turn in separate calls and reports only their combined usage. The diff --git a/adapters/typescript/pi/src/pi-sdk.ts b/adapters/typescript/pi/src/pi-sdk.ts index dc3f131ba..d9ae3fc5b 100644 --- a/adapters/typescript/pi/src/pi-sdk.ts +++ b/adapters/typescript/pi/src/pi-sdk.ts @@ -514,10 +514,13 @@ const observedFetch: typeof globalThis.fetch = async (input, init) => { * attempt that got no HTTP response, such as a connection error or a timeout * before headers, or that got any error status other than a rate limit (429), * such as a 408 or a 5xx, may have been billed. So may a Codex WebSocket - * attempt that failed, which the fetch never sees: Pi records it on the final - * message as a `provider_transport_failure` diagnostic before falling back to - * HTTP (`@earendil-works/pi-ai` 1.0.3, `api/openai-codex-responses.js`). A call that ended in an error - * reports all its counts as 0 only when its protocol sends requests through + * attempt that failed and made Pi fall back to HTTP, which the fetch never + * sees: Pi records it on the final message as a `provider_transport_failure` + * diagnostic. Pi retries `previous_response_not_found` and + * `websocket_connection_limit_reached` on the WebSocket without a diagnostic; + * the provider rejects both before the model runs, so the retried call keeps + * its counts (`@earendil-works/pi-ai` 1.0.3, `api/openai-codex-responses.js`). + * A call that ended in an error reports all its counts as 0 only when its protocol sends requests through * the observed fetch, the provider answered every attempt with a 429, no * attempt was accepted, no raw event arrived, and Pi's usage is all zeros. * Any other call keeps only the counts its raw events reported. From 8be9fab4116b57ccaf2fab663648e7f707aaba57 Mon Sep 17 00:00:00 2001 From: mnajafian-nv Date: Tue, 20 Oct 2026 07:29:15 -0700 Subject: [PATCH 38/68] fix: leave Hermes usage unknown for an invocation that delegates Delegation is one of Hermes' default toolsets. A subagent is its own agent with its own session counters, and Hermes runs a top-level delegation in the background, so its calls can finish after the delegating invocation returns. The invocation reported only the parent's tokens as known. The adapter now notes each subagent an observed agent builds, through tools.delegate_tool._build_child_agent, and leaves that invocation's counts unknown. Later invocations stay known and exclude the subagent's tokens. Signed-off-by: mnajafian-nv --- .../nemo_fabric_adapters/hermes/adapter.py | 76 ++++++++++++++++++- tests/_utils/mock_api_server.py | 15 +++- tests/adapters/test_hermes_adapter.py | 38 ++++++++++ tests/e2e/test_hermes_e2e.py | 56 ++++++++++++++ 4 files changed, 181 insertions(+), 4 deletions(-) diff --git a/adapters/python/hermes/src/nemo_fabric_adapters/hermes/adapter.py b/adapters/python/hermes/src/nemo_fabric_adapters/hermes/adapter.py index 8c2d4c636..275c6feff 100755 --- a/adapters/python/hermes/src/nemo_fabric_adapters/hermes/adapter.py +++ b/adapters/python/hermes/src/nemo_fabric_adapters/hermes/adapter.py @@ -81,6 +81,8 @@ def __init__(self) -> None: # Iteration-limit summary attempts started minus summary responses # observed; an attempt whose request raised leaves it above zero. self._unanswered_summary_attempts = 0 + # Subagents the agent built for delegated work; see _observe_delegation. + self._delegated_subagents = 0 async def start(self, payload: dict[str, Any]) -> None: if self._started: @@ -308,6 +310,7 @@ def run_hermes_turn() -> tuple[dict[str, Any], str]: skipped_before = dict(self._skipped_response_tokens) compression_calls_before = self._compression_calls unanswered_summaries_before = self._unanswered_summary_attempts + delegated_before = self._delegated_subagents invoke_task = asyncio.create_task(asyncio.to_thread(run_hermes_turn)) self._active_invoke_task = invoke_task @@ -404,13 +407,16 @@ def clear_active_invoke_task( for count, responses in self._unreported_counts.items() if responses != unreported_before[count] ) - # A compression summary call's tokens are not in the counters, and - # a summary request that raised before its response was observed - # may have been billed, so every count is an undercount. + # A compression summary call's tokens are not in the counters, a + # summary request that raised before its response was observed may + # have been billed, and a subagent's tokens go to its own counters, + # possibly after this invocation returns, so every count is an + # undercount. | ( frozenset(_REPORTED_COUNTS) if self._compression_calls != compression_calls_before or self._unanswered_summary_attempts != unanswered_summaries_before + or self._delegated_subagents != delegated_before else frozenset() ), ) @@ -518,7 +524,11 @@ def count_summary_attempt() -> None: self._unanswered_summary_attempts += 1 _USAGE_OBSERVERS[agent] = count_unreported + def count_delegated_subagent() -> None: + self._delegated_subagents += 1 + _SUMMARY_ATTEMPT_OBSERVERS[agent] = count_summary_attempt + _DELEGATION_OBSERVERS[agent] = count_delegated_subagent def _add_skipped_response_tokens( self, usage: Any, provider: Any, api_mode: Any @@ -653,6 +663,7 @@ async def stop(self) -> None: if agent is not None: _USAGE_OBSERVERS.pop(agent, None) _SUMMARY_ATTEMPT_OBSERVERS.pop(agent, None) + _DELEGATION_OBSERVERS.pop(agent, None) had_mcp_servers = bool(self._hermes_config.get("mcp_servers")) had_relay_plugin = self._relay_plugin_config_path is not None self._agent = None @@ -769,6 +780,10 @@ def _invoke_hermes_turn( _SUMMARY_ATTEMPT_OBSERVERS: weakref.WeakKeyDictionary[Any, Callable[[], None]] = ( weakref.WeakKeyDictionary() ) +# Called when the agent builds a subagent for delegated work. +_DELEGATION_OBSERVERS: weakref.WeakKeyDictionary[Any, Callable[[], None]] = ( + weakref.WeakKeyDictionary() +) # The functions ``agent.turn_response_check.check_api_response`` calls with # (agent, response, ...) for a provider response, by the module defining # each. record_response_usage folds the usage; on an invalid response, a @@ -853,6 +868,7 @@ def _install_usage_observer() -> None: for name, handler in handlers.items(): setattr(turn_response_check, name, _observed_response_handler(name, handler)) _observe_iteration_summary() + _observe_delegation() # The function each of Hermes' iteration-limit summary attempts passes its @@ -941,6 +957,60 @@ def counted_attempt(*attempt_args: Any, **attempt_kwargs: Any) -> Any: return build_counted +def _observe_delegation() -> None: + """Note each subagent an observed agent builds for delegated work. + + A subagent is its own ``AIAgent`` with its own session counters, and + Hermes runs a top-level delegation in the background, so its calls can + finish after the delegating invocation returns. Every subagent is built + by ``tools.delegate_tool._build_child_agent``, which + ``_build_child_preserving_parent_tools`` looks up at each call, so the + adapter wraps it once per process. An invocation whose agent built a + subagent has unknown counts; a subagent's own subagents belong to that + delegated work. Start fails if Hermes no longer builds subagents there. + """ + from tools import delegate_tool, delegate_tool_results + + current = getattr(delegate_tool, _DELEGATION_BUILDER, None) + if getattr(current, "nemo_fabric_observed", False): + return + builder_caller = getattr( + delegate_tool_results, "_build_child_preserving_parent_tools", None + ) + if ( + current is None + or "parent_agent" not in inspect.signature(current).parameters + or _DELEGATION_BUILDER + not in getattr(getattr(builder_caller, "__code__", None), "co_names", ()) + ): + raise lifecycle.LifecycleError( + "hermes_usage_accounting_unsupported", + "This Hermes version builds subagents outside " + "tools.delegate_tool._build_child_agent", + ) + setattr(delegate_tool, _DELEGATION_BUILDER, _observed_subagent_builder(current)) + + +def _observed_subagent_builder(build: Callable[..., Any]) -> Callable[..., Any]: + """Wrap a subagent builder to note the build to its parent's observer.""" + signature = inspect.signature(build) + + def observed_build(*args: Any, **kwargs: Any) -> Any: + parent = signature.bind(*args, **kwargs).arguments.get("parent_agent") + note_delegation = ( + _DELEGATION_OBSERVERS.get(parent) if parent is not None else None + ) + if note_delegation is not None: + note_delegation() + return build(*args, **kwargs) + + observed_build.nemo_fabric_observed = True # type: ignore[attr-defined] + return observed_build + + +_DELEGATION_BUILDER = "_build_child_agent" + + def _inner_function_calls(function: Any, name: str) -> bool: """Return whether a function defined inside ``function`` calls global ``name``.""" code = getattr(function, "__code__", None) diff --git a/tests/_utils/mock_api_server.py b/tests/_utils/mock_api_server.py index 1e6b339af..e4830d7a4 100644 --- a/tests/_utils/mock_api_server.py +++ b/tests/_utils/mock_api_server.py @@ -28,7 +28,8 @@ def mock_api_server(port: int) -> Iterator[str]: chat-completion usage object and finish reason for subsequent requests. Its ordered ``replies`` override the usage or finish reason of the next chat completions, one reply each, make one invalid, reject one with its ``status_code``, or send one - with ``no_choices`` (an empty ``choices`` list). + with ``no_choices`` (an empty ``choices`` list). A tool call with ``usage`` streams + that usage. Use the /_reset endpoint to clear all retained request and scenario state. Args: @@ -441,6 +442,18 @@ def _stream_tool_call_completion( ], }, ] + if "usage" in tool_call: + # A tool call scripted with ``usage`` ends with that usage chunk. + chunks.append( + { + "id": "chatcmpl-fabric-tool-test", + "object": "chat.completion.chunk", + "created": 0, + "model": model, + "choices": [], + "usage": tool_call["usage"], + } + ) for chunk in chunks: yield f"data: {json.dumps(chunk)}\n\n" diff --git a/tests/adapters/test_hermes_adapter.py b/tests/adapters/test_hermes_adapter.py index 794749619..ea59dc9d9 100644 --- a/tests/adapters/test_hermes_adapter.py +++ b/tests/adapters/test_hermes_adapter.py @@ -2058,11 +2058,49 @@ async def test_hermes_stop_unregisters_the_agent_usage_observer(): runtime = _usage_recording_runtime(agent) assert agent in adapter._USAGE_OBSERVERS assert agent in adapter._SUMMARY_ATTEMPT_OBSERVERS + assert agent in adapter._DELEGATION_OBSERVERS await runtime.stop() assert agent not in adapter._USAGE_OBSERVERS assert agent not in adapter._SUMMARY_ATTEMPT_OBSERVERS + assert agent not in adapter._DELEGATION_OBSERVERS + + +def test_hermes_subagent_builder_notes_only_the_observed_parent(): + class Agent: + pass + + noted: list[str] = [] + parent, subagent = Agent(), Agent() + adapter._DELEGATION_OBSERVERS[parent] = lambda: noted.append("parent") + + def build_child_agent(task_index: int, goal: str, parent_agent: object) -> str: + return f"child of {parent_agent is parent}" + + build = adapter._observed_subagent_builder(build_child_agent) + + try: + assert build(0, "goal", parent_agent=parent) == "child of True" + # A subagent's own subagent belongs to its parent's delegated work. + assert build(0, "goal", parent_agent=subagent) == "child of False" + finally: + adapter._DELEGATION_OBSERVERS.pop(parent, None) + + assert noted == ["parent"] + + +def test_hermes_usage_observer_rejects_an_unknown_subagent_builder(monkeypatch): + from tools import delegate_tool + + monkeypatch.setattr( + delegate_tool, "_build_child_agent", lambda *args, **kwargs: None + ) + + with pytest.raises(adapter.lifecycle.LifecycleError) as raised: + adapter._observe_delegation() + + assert raised.value.code == "hermes_usage_accounting_unsupported" @pytest.mark.parametrize("drift", ["replaced", "not_called"]) diff --git a/tests/e2e/test_hermes_e2e.py b/tests/e2e/test_hermes_e2e.py index 7e17c0364..c2a308a73 100644 --- a/tests/e2e/test_hermes_e2e.py +++ b/tests/e2e/test_hermes_e2e.py @@ -26,6 +26,7 @@ RelayAtofConfig, RelayAtofFileSinkConfig, RelayObservabilityConfig, + ToolsConfig, ) pytestmark = pytest.mark.usefixtures("requires_hermes_agent") @@ -298,6 +299,61 @@ async def test_hermes_usage_counts_the_iteration_limit_summary( ) +@pytest.mark.usefixtures("mock_nvidia_api_key") +async def test_hermes_usage_is_unknown_for_an_invocation_that_delegates( + code_review_agent_dir: Path, + api_server: str, +): + os.environ["ADAPTER_PYTHON"] = sys.executable + config = hermes_config() + config.models["default"].base_url = f"{api_server}/v1" + # Delegation is one of Hermes' default toolsets. + config.tools = ToolsConfig(enabled=["delegation"]) + config.runtime.max_turns = 5 + + async with await Fabric().start_runtime( + config, base_dir=code_review_agent_dir + ) as runtime: + await runtime.invoke(input="first") + scenario_response = requests.post( + f"{api_server}/_scenario", + json={ + "tool_calls": [ + { + "name": "delegate_task", + "arguments": {"goal": "count files"}, + "usage": _chat_usage(0, 0), + } + ], + "usage": _chat_usage(10, 2), + }, + timeout=5, + ) + scenario_response.raise_for_status() + delegated = await runtime.invoke(input="delegate") + following = await runtime.invoke(input="after") + sent = requests.get(f"{api_server}/_requests", timeout=5).json() + + # Hermes runs the subagent in the background, so its tokens can be spent + # after the delegating invocation returns and never reach its counters. + assert delegated["status"] == "succeeded", delegated.to_mapping() + assert any( + "count files" in str(message.get("content")) + for request in sent + for message in request["messages"] + if message.get("role") == "user" + ), sent + assert delegated.usage is None, delegated.to_mapping() + # The next invocation reports only its own turn, not the subagent's. + assert following.to_mapping().get("usage") == { + "input_tokens": 10, + "input_tokens_include_cache": True, + "output_tokens": 2, + "total_tokens": 12, + "metadata": {}, + }, following.to_mapping() + + @pytest.mark.usefixtures("mock_nvidia_api_key") async def test_hermes_makes_no_model_call_for_session_titles( code_review_agent_dir: Path, From 70bf92543fec8c2bb5fa038a1d1c996280add682 Mon Sep 17 00:00:00 2001 From: mnajafian-nv Date: Tue, 20 Oct 2026 07:31:39 -0700 Subject: [PATCH 39/68] fix: leave a Pi count unknown when its sum is not a safe integer Pi's counts are safe integers, but the adapter added them without checking the result, so 9007199254740991 + 2 reported 9007199254740992. Every addition now checks its sum: an unsafe input sum, running total or derived total leaves that count unknown and the other counts known. Signed-off-by: mnajafian-nv --- adapters/typescript/pi/src/pi-sdk.ts | 37 ++++++++++++-------- adapters/typescript/pi/src/runtime.ts | 5 +-- adapters/typescript/pi/test/pi-sdk.test.mjs | 20 +++++++++++ adapters/typescript/pi/test/runtime.test.mjs | 25 +++++++++++++ 4 files changed, 71 insertions(+), 16 deletions(-) diff --git a/adapters/typescript/pi/src/pi-sdk.ts b/adapters/typescript/pi/src/pi-sdk.ts index d9ae3fc5b..7b13e2591 100644 --- a/adapters/typescript/pi/src/pi-sdk.ts +++ b/adapters/typescript/pi/src/pi-sdk.ts @@ -98,22 +98,33 @@ function piCounts(usage: unknown): Record<"input" | "cacheRead" | "cacheWrite" | return { input: count("input"), cacheRead: count("cacheRead"), cacheWrite: count("cacheWrite"), output: count("output") }; } +/** + * The sum of token counts, or undefined when it is not a safe integer: a + * larger sum loses precision, so the count it would report is unknown. + */ +function safeSum(...counts: number[]): number | undefined { + const sum = counts.reduce((partial, count) => partial + count, 0); + return Number.isSafeInteger(sum) ? sum : undefined; +} + /** * One call's counts from Pi's usage, with input counting its cache reads and * writes, and which of them the usage holds. Input needs Pi's `input`, * `cacheRead` and `cacheWrite`, cache reads need `cacheRead`, and output needs * `output`; a missing field adds 0 to `counts` and leaves the counts that need - * it unreported. Undefined when the usage is not an object. + * it unreported, and so does an input whose sum is not a safe integer. + * Undefined when the usage is not an object. */ function callCounts(usage: unknown): { counts: Required; present: ReportedCounts } | undefined { if (typeof usage !== "object" || usage === null) { return undefined; } const { input, cacheRead, cacheWrite, output } = piCounts(usage); + const inclusiveInput = safeSum(input ?? 0, cacheRead ?? 0, cacheWrite ?? 0); return { - counts: { input: (input ?? 0) + (cacheRead ?? 0) + (cacheWrite ?? 0), cacheRead: cacheRead ?? 0, output: output ?? 0 }, + counts: { input: inclusiveInput ?? 0, cacheRead: cacheRead ?? 0, output: output ?? 0 }, present: { - input: input !== undefined && cacheRead !== undefined && cacheWrite !== undefined, + input: input !== undefined && cacheRead !== undefined && cacheWrite !== undefined && inclusiveInput !== undefined, cacheRead: cacheRead !== undefined, output: output !== undefined, }, @@ -122,19 +133,17 @@ function callCounts(usage: unknown): { counts: Required; present: /** * Add a call's counts to a prompt's running total. Each count stays known only - * while every call reported it, independently of the others. The total is - * unknown once no count is known. + * while every call reported it and its sum stays a safe integer, independently + * of the others. The total is unknown once no count is known. */ function sumCounts(total: PiTokenUsage, call: Required, reported: ReportedCounts): PiTokenUsage | undefined { const sum: PiTokenUsage = {}; - if (reported.input && total.input !== undefined) { - sum.input = total.input + call.input; - } - if (reported.cacheRead && total.cacheRead !== undefined) { - sum.cacheRead = total.cacheRead + call.cacheRead; - } - if (reported.output && total.output !== undefined) { - sum.output = total.output + call.output; + for (const name of ["input", "cacheRead", "output"] as const) { + const running = total[name]; + const count = reported[name] && running !== undefined ? safeSum(running, call[name]) : undefined; + if (count !== undefined) { + sum[name] = count; + } } return sum.input === undefined && sum.cacheRead === undefined && sum.output === undefined ? undefined : sum; } @@ -219,7 +228,7 @@ function addSummaryCallsUsage( const addsUp = names.every((name) => { const sum = compaction.calls.reduce((partial, call) => { const count = piCounts(call)[name]; - return partial === undefined || count === undefined ? undefined : partial + count; + return partial === undefined || count === undefined ? undefined : safeSum(partial, count); }, 0); return sum !== undefined && sum === expected[name]; }); diff --git a/adapters/typescript/pi/src/runtime.ts b/adapters/typescript/pi/src/runtime.ts index 200308bbf..190cf45c0 100644 --- a/adapters/typescript/pi/src/runtime.ts +++ b/adapters/typescript/pi/src/runtime.ts @@ -66,18 +66,19 @@ function failed(code: string, message: string): AgentRunResult { * The invocation's usage, with input counting its cache reads and writes. A * count that a call left unreported stays absent, cache reads are kept even * when input is unknown, and the total is absent unless both input and output - * are known. + * are known and their sum is a safe integer. */ function agentUsage(usage: PiTokenUsage | undefined): AgentUsage | undefined { if (usage === undefined || (usage.input === undefined && usage.cacheRead === undefined && usage.output === undefined)) { return undefined; } const { input, cacheRead, output } = usage; + const total = input === undefined || output === undefined ? undefined : input + output; return { ...(input === undefined ? {} : { input_tokens: input, input_tokens_include_cache: true }), ...(cacheRead === undefined ? {} : { cached_input_tokens: cacheRead }), ...(output === undefined ? {} : { output_tokens: output }), - ...(input === undefined || output === undefined ? {} : { total_tokens: input + output }), + ...(total === undefined || !Number.isSafeInteger(total) ? {} : { total_tokens: total }), }; } diff --git a/adapters/typescript/pi/test/pi-sdk.test.mjs b/adapters/typescript/pi/test/pi-sdk.test.mjs index a0f79992c..997c80eb1 100644 --- a/adapters/typescript/pi/test/pi-sdk.test.mjs +++ b/adapters/typescript/pi/test/pi-sdk.test.mjs @@ -1464,6 +1464,26 @@ test("sums each call's own usage report", () => { assert.equal(addUsage(undefined, { input: 1, output: 1, cacheRead: 0, cacheWrite: 0 }), undefined); }); +test("a count whose sum is not a safe integer is unknown, and the others stay known", () => { + const max = Number.MAX_SAFE_INTEGER; + + // Input counts cache reads and writes, so its own sum can overflow. + assert.deepEqual(addUsage(NO_TOKENS, { input: max, cacheRead: 2, cacheWrite: 0, output: 1 }), { + cacheRead: 2, + output: 1, + }); + // A running total that reaches the largest safe integer stays exact. + assert.deepEqual(addUsage({ input: max - 1, cacheRead: 0, output: max }, { input: 1, cacheRead: 0, cacheWrite: 0, output: 1 }), { + input: max, + cacheRead: 0, + }); + // Only the cache reads overflow here; input and output stay known. + assert.deepEqual(addUsage({ input: 1, cacheRead: max, output: 1 }, { input: 0, cacheRead: 1, cacheWrite: 0, output: 1 }), { + input: 2, + output: 2, + }); +}); + // A tool result that reports usage, as a custom tool can. const toolResult = (usage) => ({ type: "message_end", message: { role: "toolResult", usage } }); diff --git a/adapters/typescript/pi/test/runtime.test.mjs b/adapters/typescript/pi/test/runtime.test.mjs index 4604fad49..d24e62ce2 100644 --- a/adapters/typescript/pi/test/runtime.test.mjs +++ b/adapters/typescript/pi/test/runtime.test.mjs @@ -584,6 +584,31 @@ test("reports the invocation's tokens with cache reads and writes inside the inp }); }); +test("leaves out a total that is not a safe integer", async () => { + const usage = { input: Number.MAX_SAFE_INTEGER, cacheRead: 0, output: 2 }; + const runtime = new PiAdapterRuntime({ + async create() { + return { + async prompt() { + return { accepted: true, turnStarted: true, turnCount: 1, text: "ok", stopReason: "stop", usage }; + }, + async stop() {}, + }; + }, + }); + + await runtime.start(startInput()); + const result = await runtime.invoke({ input: "one" }, context); + await runtime.stop(); + + assert.deepEqual(result.usage, { + input_tokens: Number.MAX_SAFE_INTEGER, + cached_input_tokens: 0, + input_tokens_include_cache: true, + output_tokens: 2, + }); +}); + test("reports only the token counts the prompt's provider reported", async () => { const outcomes = [{ input: 130, cacheRead: 30 }, { output: 5 }, { input: 70, output: 7 }, { cacheRead: 12, output: 3 }, { cacheRead: 4 }]; const runtime = new PiAdapterRuntime({ From c9e6881cd1c7af21dba32b22051538f4f63bae4b Mon Sep 17 00:00:00 2001 From: mnajafian-nv Date: Tue, 20 Oct 2026 07:36:54 -0700 Subject: [PATCH 40/68] fix: leave a Pi prompt unknown after any cache warm-up attempt The adapter saw a warm-up only through its usage entry, which Pi writes only after a successful refresh. A failed warm-up writes nothing, and Pi aborts and discards an outstanding one when the agent settles, so both left the prompt's usage known. Pi's warmer replays a call's options, and every built-in protocol passes a call's request body to onPayload once before sending, so a later onPayload is a warm-up attempt. A prompt with a warm-up attempt since the previous prompt now has unknown usage. Signed-off-by: mnajafian-nv --- adapters/typescript/pi/src/pi-sdk.ts | 44 ++++++-- adapters/typescript/pi/test/pi-sdk.test.mjs | 114 ++++++++++++++++++++ 2 files changed, 150 insertions(+), 8 deletions(-) diff --git a/adapters/typescript/pi/src/pi-sdk.ts b/adapters/typescript/pi/src/pi-sdk.ts index 7b13e2591..51b501ec0 100644 --- a/adapters/typescript/pi/src/pi-sdk.ts +++ b/adapters/typescript/pi/src/pi-sdk.ts @@ -586,15 +586,20 @@ function hasTransportFailure(diagnostics: unknown): boolean { * is stored before Pi emits either event. A request the provider rate-limited * before generating reports all its counts as 0, and a call with an ambiguous * attempt reports none (see `endedCallCounts`). Pi's cache warmer replays a - * request with the same options, so a warm-up's raw events reach this - * observer too, but the warm-up runs outside the agent's stream function: its - * usage entry has no record and stays unknown. The record is copied when the - * call ends so a later warm-up cannot change it. + * request with the same options outside the agent's stream function, so a + * warm-up's raw events reach this observer too, but its usage entry has no + * record and stays unknown. Every built-in protocol passes each call's request + * body to `onPayload` once, before sending it (`@earendil-works/pi-ai` 1.0.3, + * `api/*.js`), so a later `onPayload` on the same options is a warm-up + * attempt, and `onWarmUp` is called before it is sent, including one that + * fails or that Pi aborts and discards when the agent settles. The record is + * copied when the call ends so a later warm-up cannot change it. */ export function recordReportedCounts( agent: AgentSession["agent"], records: WeakMap, onCallEnd: (usage: object) => void, + onWarmUp: () => void = () => undefined, ): void { const streamFunction = agent.streamFunction; agent.streamFunction = async (model, context, options) => { @@ -606,10 +611,18 @@ export function recordReportedCounts( delegate: ownFetch === undefined || ownFetch === observedFetch ? globalThis.fetch : ownFetch, observation, }; + let payloads = 0; const response = await callAttempts.run(attempts, () => streamFunction(model, context, { ...options, ...(FETCH_PROTOCOLS.has(model.api) ? { fetch: observedFetch } : {}), + onPayload: async (payload, payloadModel) => { + payloads += 1; + if (payloads > 1) { + onWarmUp(); + } + return options?.onPayload?.(payload, payloadModel); + }, onResponse: async (providerResponse, responseModel) => { observation.accepted ||= providerResponse.status >= 200 && providerResponse.status < 300; await options?.onResponse?.(providerResponse, responseModel); @@ -1453,6 +1466,9 @@ class PiSdkSessionHandle implements PiSessionHandle { private readonly runtimeDir: string; private readonly compactionListeners: Set<(event: unknown) => void>; private readonly reportedCounts = new WeakMap(); + // Cache warm-up attempts, and how many of them a prompt has accounted for. + private warmUpAttempts = 0; + private warmUpAttemptsCounted = 0; // The model calls that ended since the current compaction was prepared. private compaction: { expected: number; calls: object[] } | undefined; @@ -1468,9 +1484,16 @@ class PiSdkSessionHandle implements PiSessionHandle { this.runtimeDir = runtimeDir; this.compactionListeners = compactionListeners; this.relay = relay; - recordReportedCounts(session.agent, this.reportedCounts, (usage) => { - this.compaction?.calls.push(usage); - }); + recordReportedCounts( + session.agent, + this.reportedCounts, + (usage) => { + this.compaction?.calls.push(usage); + }, + () => { + this.warmUpAttempts += 1; + }, + ); this.unsubscribeTurnCounter = this.session.subscribe((event) => { if (event.type === "turn_start") { this.cumulativeTurnCount += 1; @@ -1521,6 +1544,11 @@ class PiSdkSessionHandle implements PiSessionHandle { this.compaction = undefined; unsubscribe(); } + // A warm-up's tokens reach no usage entry the adapter can check, and a + // failed or discarded one reaches none at all, so a prompt with a warm-up + // attempt since the previous prompt has unknown tokens. + const warmedUp = this.warmUpAttempts !== this.warmUpAttemptsCounted; + this.warmUpAttemptsCounted = this.warmUpAttempts; return { accepted, turnStarted, @@ -1529,7 +1557,7 @@ class PiSdkSessionHandle implements PiSessionHandle { stopReason: finalAssistant?.stopReason, errorMessage: finalAssistant?.errorMessage, shutdownRequested: this.state.shutdownRequested, - usage: finalAssistant === undefined ? undefined : usage, + usage: finalAssistant === undefined || warmedUp ? undefined : usage, }; } diff --git a/adapters/typescript/pi/test/pi-sdk.test.mjs b/adapters/typescript/pi/test/pi-sdk.test.mjs index 997c80eb1..89fb8508a 100644 --- a/adapters/typescript/pi/test/pi-sdk.test.mjs +++ b/adapters/typescript/pi/test/pi-sdk.test.mjs @@ -1965,6 +1965,120 @@ test( }, ); +/** + * Run one model call through the adapter's recorder while Pi's own cache + * warmer keeps its prompt cache entry warm, against a local endpoint that + * answers the call and, by `warmUp`, fails ("fail"), answers ("answer") or + * holds ("hold") each one-token warm-up request. With `settleFirst`, the + * agent settles before the first refresh is due. Returns how many warm-up + * attempts the recorder noted and how many the endpoint received. + */ +async function warmedCall({ warmUp, settleFirst = false }) { + const received = { call: 0, warmUp: 0 }; + const held = []; + const server = createServer(async (request, response) => { + let body = ""; + for await (const chunk of request) body += chunk; + const payload = JSON.parse(body); + const isWarmUp = (payload.max_completion_tokens ?? payload.max_tokens) === 1; + received[isWarmUp ? "warmUp" : "call"] += 1; + if (isWarmUp && warmUp === "fail") { + response.writeHead(500, { "content-type": "application/json" }); + response.end(JSON.stringify({ error: { message: "overloaded", type: "server_error" } })); + return; + } + if (isWarmUp && warmUp === "hold") { + held.push(response); + return; + } + const common = { id: "chatcmpl-warm", object: "chat.completion.chunk", created: 0, model: "m" }; + const chunk = (fields) => `data: ${JSON.stringify({ ...common, ...fields })}\n\n`; + response.writeHead(200, { "content-type": "text/event-stream" }); + response.end( + chunk({ choices: [{ index: 0, delta: { role: "assistant", content: "done" }, finish_reason: null }] }) + + chunk({ choices: [{ index: 0, delta: {}, finish_reason: "stop" }] }) + + chunk({ choices: [], usage: { prompt_tokens: 100_000, completion_tokens: 1 } }) + + "data: [DONE]\n\n", + ); + }); + const baseUrl = await listen(server); + const { CacheWarmer } = await import( + new URL("../node_modules/@earendil-works/pi-coding-agent/dist/core/cache-warmer.js", import.meta.url) + ); + const { streamSimple } = await import("@earendil-works/pi-ai/api/openai-completions"); + // The warmer prices a refresh from the branch's last prompt size. + const sessionManager = { + getBranch: () => [{ type: "message", message: { role: "assistant", usage: { input: 100_000, cacheRead: 0, cacheWrite: 0 } } }], + appendUsage: (kind, provider, model, usage) => ({ type: "usage", kind, usage }), + }; + const warmer = new CacheWarmer({ streamSimple }, sessionManager, () => "streaming"); + const agent = { + streamFunction: (model, context, options) => { + warmer.start({ model, context, options }, () => true); + return streamSimple(model, context, options); + }, + }; + let warmUps = 0; + recordReportedCounts(agent, new WeakMap(), () => undefined, () => { + warmUps += 1; + }); + // An 11-second prompt cache lifetime schedules the first refresh after 1 second. + const model = { + id: "m", + name: "m", + api: "openai-completions", + provider: "test", + baseUrl: `${baseUrl}/v1`, + reasoning: false, + input: ["text"], + cost: { input: 3, output: 15, cacheRead: 0.3, cacheWrite: 3.75 }, + contextWindow: 200_000, + maxTokens: 1_000, + promptCache: { short: 11 }, + }; + try { + const response = await agent.streamFunction( + model, + { messages: [{ role: "user", content: "one", timestamp: 0 }] }, + { apiKey: "test", cacheRetention: "short" }, + ); + await response.result(); + if (settleFirst) { + warmer.onAgentSettled(); + } + const deadline = Date.now() + 3_000; + while (!settleFirst && received.warmUp === 0 && Date.now() < deadline) { + await new Promise((resolve) => setTimeout(resolve, 50)); + } + // The agent settles: Pi aborts an outstanding warm-up and discards it. + warmer.onAgentSettled(); + await new Promise((resolve) => setTimeout(resolve, 100)); + return { warmUps, received }; + } finally { + warmer.cancel(); + for (const response of held) { + response.destroy(); + } + await close(server); + } +} + +for (const warmUp of ["fail", "hold", "answer"]) { + test(`notes a cache warm-up attempt that Pi's warmer ${warmUp === "hold" ? "aborts" : warmUp === "fail" ? "fails" : "completes"}`, async () => { + const { warmUps, received } = await warmedCall({ warmUp }); + + assert.deepEqual(received, { call: 1, warmUp: 1 }); + assert.equal(warmUps, 1); + }); +} + +test("notes no warm-up when the agent settles before the refresh is due", async () => { + const { warmUps, received } = await warmedCall({ warmUp: "answer", settleFirst: true }); + + assert.deepEqual(received, { call: 1, warmUp: 0 }); + assert.equal(warmUps, 0); +}); + /** One unmasked WebSocket text frame, as a server sends it. */ function webSocketTextFrame(text) { const payload = Buffer.from(text); From 94ce714af3d8b6fcccc8793c12b23dbc82a0e66c Mon Sep 17 00:00:00 2001 From: mnajafian-nv Date: Tue, 20 Oct 2026 07:41:37 -0700 Subject: [PATCH 41/68] fix: keep Hermes Bedrock and Gemini omitted counts unknown Hermes' native Bedrock Converse and Gemini transports fill every count the provider left out with 0 before the observer sees the usage, so a Bedrock reply with only inputTokens reported output 0 and a total as known, and a Converse stream without its metadata event reported zeros. The adapter wraps the functions that read the raw usage and records the counts it reported on the usage object, under a private attribute that Relay payloads skip; those counts decide presence. Signed-off-by: mnajafian-nv --- .../nemo_fabric_adapters/hermes/adapter.py | 163 ++++++++++++++++++ tests/adapters/test_hermes_adapter.py | 105 +++++++++++ 2 files changed, 268 insertions(+) diff --git a/adapters/python/hermes/src/nemo_fabric_adapters/hermes/adapter.py b/adapters/python/hermes/src/nemo_fabric_adapters/hermes/adapter.py index 275c6feff..353a0de44 100755 --- a/adapters/python/hermes/src/nemo_fabric_adapters/hermes/adapter.py +++ b/adapters/python/hermes/src/nemo_fabric_adapters/hermes/adapter.py @@ -869,6 +869,7 @@ def _install_usage_observer() -> None: setattr(turn_response_check, name, _observed_response_handler(name, handler)) _observe_iteration_summary() _observe_delegation() + _observe_native_usage() # The function each of Hermes' iteration-limit summary attempts passes its @@ -1010,6 +1011,165 @@ def observed_build(*args: Any, **kwargs: Any) -> Any: _DELEGATION_BUILDER = "_build_child_agent" +# The attribute that holds the counts a native transport's raw usage reported, +# set on the usage object it builds before Hermes reads it. The leading +# underscore keeps it out of Hermes' Relay payloads (agent/relay_llm.py +# _jsonable skips private attributes). +_NATIVE_REPORTED_COUNTS = "_nemo_fabric_reported_counts" +# The raw usage fields of Hermes' native Bedrock Converse and Gemini +# transports, by the count each reports, in the shape of _CHAT_FIELDS. Bedrock +# documents its cache fields as optional, so its input needs only inputTokens; +# Gemini's output is visible and thinking tokens, either of which it reports. +_BEDROCK_FIELDS = { + "input": (("inputTokens",),), + "cached": (("cacheReadInputTokens",),), + "output": (("outputTokens",),), +} +_GEMINI_FIELDS = { + "input": (("promptTokenCount",),), + "cached": (("cachedContentTokenCount",),), + "output": (("candidatesTokenCount",), ("thoughtsTokenCount",)), +} + + +def _observe_native_usage() -> None: + """Record which counts a native transport's raw usage reported. + + Hermes' Bedrock Converse and Gemini transports build an OpenAI-shaped + usage object and fill every count the provider left out with 0 + (``agent/bedrock_adapter.py``, ``agent/gemini_native_adapter.py``), so the + observer cannot tell an omitted count from a reported zero. The adapter + wraps the functions that read the raw usage once per process and records + the reported counts on the usage object they return; the streamed + Bedrock path reads them from the stream's metadata event, and a stream + without one reports no count. A transport whose module cannot be + imported cannot serve a turn, so it is left as it is. + """ + try: + from agent import bedrock_adapter + except ImportError: + bedrock_adapter = None + try: + from agent import gemini_native_adapter + except ImportError: + gemini_native_adapter = None + if bedrock_adapter is not None: + _wrap_native_usage( + bedrock_adapter, + "normalize_converse_response", + lambda args, kwargs: _raw_usage(_first_argument(args, kwargs, "response")), + _BEDROCK_FIELDS, + ) + _wrap_bedrock_stream(bedrock_adapter) + if gemini_native_adapter is not None: + _wrap_native_usage( + gemini_native_adapter, + "_usage_from_metadata", + lambda args, kwargs: _first_argument(args, kwargs, "usage_meta"), + _GEMINI_FIELDS, + annotate_result=True, + ) + + +def _first_argument(args: tuple[Any, ...], kwargs: dict[str, Any], name: str) -> Any: + return args[0] if args else kwargs.get(name) + + +def _raw_usage(response: Any) -> Any: + return response.get("usage") if isinstance(response, dict) else None + + +def _native_reported_counts(raw: Any, fields: dict[str, Any]) -> frozenset[str]: + """Return the counts that a native transport's raw ``usage`` reported.""" + if not isinstance(raw, dict): + return frozenset() + return frozenset( + count + for count, alternatives in fields.items() + if any( + all(_is_reported_count(raw, path) for path in alternative) + for alternative in alternatives + ) + ) + + +def _annotate_usage(usage: Any, reported: frozenset[str]) -> None: + if usage is not None: + setattr(usage, _NATIVE_REPORTED_COUNTS, reported) + + +def _wrap_native_usage( + module: Any, + name: str, + raw_usage: Callable[[tuple[Any, ...], dict[str, Any]], Any], + fields: dict[str, Any], + *, + annotate_result: bool = False, +) -> None: + """Wrap ``module.name`` to record its raw usage's reported counts. + + The wrapped function returns a response with ``usage``, or with + ``annotate_result`` the usage object itself. + """ + current = getattr(module, name, None) + if getattr(current, "nemo_fabric_observed", False): + return + if not callable(current): + raise lifecycle.LifecycleError( + "hermes_usage_accounting_unsupported", + f"This Hermes version reads native usage outside {module.__name__}.{name}", + ) + + def observed(*args: Any, **kwargs: Any) -> Any: + reported = _native_reported_counts(raw_usage(args, kwargs), fields) + result = current(*args, **kwargs) + _annotate_usage( + result if annotate_result else getattr(result, "usage", None), reported + ) + return result + + observed.nemo_fabric_observed = True # type: ignore[attr-defined] + setattr(module, name, observed) + + +def _wrap_bedrock_stream(bedrock_adapter: Any) -> None: + """Wrap Bedrock's ConverseStream reader to record its metadata usage.""" + name = "stream_converse_with_callbacks" + current = getattr(bedrock_adapter, name, None) + if getattr(current, "nemo_fabric_observed", False): + return + if not callable(current): + raise lifecycle.LifecycleError( + "hermes_usage_accounting_unsupported", + f"This Hermes version reads native usage outside " + f"agent.bedrock_adapter.{name}", + ) + + def observed(event_stream: Any, *args: Any, **kwargs: Any) -> Any: + if not isinstance(event_stream, dict): + return current(event_stream, *args, **kwargs) + raw: dict[str, Any] = {} + + def events() -> Any: + for event in event_stream.get("stream", []): + metadata = event.get("metadata") if isinstance(event, dict) else None + if isinstance(metadata, dict): + usage = metadata.get("usage") + raw.clear() + if isinstance(usage, dict): + raw.update(usage) + yield event + + result = current({**event_stream, "stream": events()}, *args, **kwargs) + _annotate_usage( + getattr(result, "usage", None), + _native_reported_counts(raw, _BEDROCK_FIELDS), + ) + return result + + observed.nemo_fabric_observed = True # type: ignore[attr-defined] + setattr(bedrock_adapter, name, observed) + def _inner_function_calls(function: Any, name: str) -> bool: """Return whether a function defined inside ``function`` calls global ``name``.""" @@ -1071,6 +1231,9 @@ def _skips_usage_recording(handler: str, outcome: Any) -> bool: def _omitted_counts(usage: Any, provider: Any, api_mode: Any) -> list[str]: """Return the counts that the provider's raw ``usage`` left out.""" + reported = getattr(usage, _NATIVE_REPORTED_COUNTS, None) + if isinstance(reported, frozenset): + return [count for count in _REPORTED_COUNTS if count not in reported] mode = str(api_mode or "").strip().lower() provider = str(provider or "").strip().lower() if mode == "anthropic_messages" or provider == "anthropic": diff --git a/tests/adapters/test_hermes_adapter.py b/tests/adapters/test_hermes_adapter.py index ea59dc9d9..d4ed5fb9c 100644 --- a/tests/adapters/test_hermes_adapter.py +++ b/tests/adapters/test_hermes_adapter.py @@ -2053,6 +2053,111 @@ def refused(agent: Any) -> None: assert turn.usage.output_tokens == 5 +def _bedrock_response(usage: dict[str, int]) -> Any: + from agent import bedrock_adapter + + return bedrock_adapter.normalize_converse_response( + {"output": {"message": {"content": [{"text": "done"}]}}, "usage": usage} + ) + + +def _bedrock_streamed_response(usage: dict[str, int] | None) -> Any: + from agent import bedrock_adapter + + events: list[dict[str, Any]] = [ + {"contentBlockDelta": {"contentBlockIndex": 0, "delta": {"text": "done"}}}, + {"messageStop": {"stopReason": "end_turn"}}, + ] + if usage is not None: + events.append({"metadata": {"usage": usage}}) + return bedrock_adapter.stream_converse_with_callbacks({"stream": events}) + + +def _gemini_response(usage: dict[str, int]) -> Any: + from agent import gemini_native_adapter + + return gemini_native_adapter.translate_gemini_response( + { + "candidates": [ + {"content": {"parts": [{"text": "done"}]}, "finishReason": "STOP"} + ], + "usageMetadata": usage, + }, + "gemini-test", + ) + + +@pytest.mark.parametrize( + ("response", "expected"), + [ + # Bedrock fills an omitted count with 0 before Hermes sees it. + ( + lambda: _bedrock_response({"inputTokens": 100}), + {"input_tokens": 100, "cached_input_tokens": None, "output_tokens": None}, + ), + ( + lambda: _bedrock_response( + {"inputTokens": 100, "outputTokens": 5, "cacheReadInputTokens": 20} + ), + {"input_tokens": 120, "cached_input_tokens": 20, "output_tokens": 5}, + ), + # A stream that ends without its metadata event reports no count. + ( + lambda: _bedrock_streamed_response(None), + {"input_tokens": None, "cached_input_tokens": None, "output_tokens": None}, + ), + ( + lambda: _bedrock_streamed_response({"inputTokens": 30, "outputTokens": 3}), + {"input_tokens": 30, "cached_input_tokens": None, "output_tokens": 3}, + ), + ( + lambda: _gemini_response({"promptTokenCount": 40}), + {"input_tokens": 40, "cached_input_tokens": None, "output_tokens": None}, + ), + ( + lambda: _gemini_response( + { + "promptTokenCount": 40, + "cachedContentTokenCount": 10, + "candidatesTokenCount": 4, + "thoughtsTokenCount": 2, + } + ), + {"input_tokens": 40, "cached_input_tokens": 10, "output_tokens": 6}, + ), + ], + ids=[ + "bedrock-input-only", + "bedrock-complete", + "bedrock-stream-without-metadata", + "bedrock-stream", + "gemini-input-only", + "gemini-complete", + ], +) +async def test_hermes_usage_keeps_native_transport_omissions_unknown( + response: Callable[[], Any], expected: dict[str, int | None] +): + agent = _UsageRecordingAgent( + [[lambda agent: _record_response(agent, response())]] + ) + runtime = _usage_recording_runtime(agent) + request, context = _usage_invocation() + + result = await runtime.invoke(request, context) + + usage = result.usage + if not any(expected.values()): + assert usage is None + return + assert usage is not None + assert { + "input_tokens": usage.input_tokens, + "cached_input_tokens": usage.cached_input_tokens, + "output_tokens": usage.output_tokens, + } == expected + + async def test_hermes_stop_unregisters_the_agent_usage_observer(): agent = _UsageRecordingAgent([]) runtime = _usage_recording_runtime(agent) From f43462a696364fd8b286ae1fca8f66d789b2dc27 Mon Sep 17 00:00:00 2001 From: mnajafian-nv Date: Tue, 20 Oct 2026 07:42:10 -0700 Subject: [PATCH 42/68] docs: describe Hermes delegation, native transport and Pi warm-up rules State that an invocation that delegates has no usage, which raw fields decide presence for Hermes' Bedrock and Gemini transports, the remaining Anthropic-compatible gateway gap, that any Pi warm-up attempt leaves the invocation unknown, and that an unsafe Pi sum leaves its count unknown. Signed-off-by: mnajafian-nv --- adapters/python/hermes/README.md | 92 +++++++++++++++++++------------- adapters/typescript/pi/README.md | 15 ++++-- 2 files changed, 66 insertions(+), 41 deletions(-) diff --git a/adapters/python/hermes/README.md b/adapters/python/hermes/README.md index 0e0ee8376..514bcde24 100644 --- a/adapters/python/hermes/README.md +++ b/adapters/python/hermes/README.md @@ -115,26 +115,30 @@ runtime. ## Token Usage Hermes counts its main model's tokens per session, and the adapter keeps one -session per runtime. The adapter reads the agent's session counters right -before each turn runs and again after it, so each result's `usage` is that -turn's change: `input_tokens` with cache reads and cache writes included, +session per runtime. The adapter reads the agent's session counters right before +each turn runs and again after it, so each result's `usage` is that turn's +change: `input_tokens` with cache reads and cache writes included, `cached_input_tokens`, `output_tokens`, and `total_tokens`. A failure result -reports its own tokens too, and an invocation that raised or was cancelled -never adds its tokens to the next one. A counter that is missing, that is not -a non-negative integer below 2^64, or that went down because Hermes reset its -session, is left out for that invocation. Hermes records auxiliary calls, such as -context-compression summaries and session titles, separately from these +reports its own tokens too, and an invocation that raised or was cancelled never +adds its tokens to the next one. A counter that is missing, that is not a +non-negative integer below 2^64, or that went down because Hermes reset its +session, is left out for that invocation. Hermes records auxiliary calls, such +as context-compression summaries and session titles, separately from these counters. The adapter sets `auxiliary.title_generation.model_upgrade_enabled` to false, so Hermes keeps its locally derived session title and makes no title -model call; with Relay enabled, title generation is off entirely. An invocation in which Hermes -ran a context-compression summary call, or in which a plugin context engine -compressed, has no `usage`, because the compression calls' tokens are missing -from the counters. Hermes also leaves out of its counters the summary request it -makes when a turn reaches its iteration limit. The adapter adds the counts that -this summary response reported, before Hermes reads its text, so they count -even when Hermes cannot read the response and answers with its fallback text. -An invocation whose summary response carried no usage, or whose summary -request failed before any response, has no `usage`. +model call; with Relay enabled, title generation is off entirely. An invocation +in which Hermes ran a context-compression summary call, or in which a plugin +context engine compressed, has no `usage`, because the compression calls' tokens +are missing from the counters. Delegation is one of Hermes' default toolsets: a +subagent is its own agent with its own counters, and Hermes runs a top-level +delegation in the background, so the subagent's calls can finish after the +invocation returns. An invocation whose agent built a subagent has no `usage`; +later invocations report only their own turns. Hermes also leaves out of its +counters the summary request it makes when a turn reaches its iteration limit. +The adapter adds the counts that this summary response reported, before Hermes +reads its text, so they count even when Hermes cannot read the response and +answers with its fallback text. An invocation whose summary response carried no +usage, or whose summary request failed before any response, has no `usage`. Hermes adds no tokens for a provider response that carried no usage, so an invocation with such a response has no `usage`. For a response it rejects as @@ -144,27 +148,39 @@ counters. The adapter adds the counts that such a response reported itself, using the provider and API mode that served it, and counts it only once when Hermes records it after all. A streamed call that fails and is retried inside Hermes' API call, such as an empty stream, never reaches these checks, so its -billed usage is not observed and the invocation can undercount. Hermes -reads a count that a response's usage left out as 0, so the adapter checks each -response's usage before Hermes normalizes it. A reported zero still counts as -0. When a response left out its input count, the invocation's `usage` leaves -out `input_tokens` and `total_tokens`, and keeps `cached_input_tokens` when the -cache-read count is known. Anthropic -Messages usage reports input without the cache, so its input count also needs -both cache fields. When a response left out its output count, the `usage` -leaves out `output_tokens` and `total_tokens`. When a response sent no cache -details, the `usage` leaves out `cached_input_tokens`. - -This check covers the usage that the OpenAI-compatible, Responses, and -Anthropic clients return. Hermes' native Gemini and Bedrock transports, and its -Anthropic transport when a stream omits usage, fill missing counts with 0 -before the adapter sees them. Hermes' Codex app-server runtime records usage -outside this path. Hermes passes a response's raw usage to no per-agent hook, -so the adapter wraps `record_response_usage` and the three handlers for -invalid, truncated, and refused responses in `agent.turn_response_check`, and -the iteration-limit summary's `_summary_text` and attempt builders in -`agent.chat_completion_helpers`, once per process. The runtime fails to start if the installed Hermes no longer calls -them by those names. +billed usage is not observed and the invocation can undercount. Hermes reads a +count that a response's usage left out as 0, so the adapter checks each +response's usage before Hermes normalizes it. A reported zero still counts as 0. +When a response left out its input count, the invocation's `usage` leaves out +`input_tokens` and `total_tokens`, and keeps `cached_input_tokens` when the +cache-read count is known. Anthropic Messages usage reports input without the +cache, so its input count also needs both cache fields. When a response left out +its output count, the `usage` leaves out `output_tokens` and `total_tokens`. +When a response sent no cache details, the `usage` leaves out +`cached_input_tokens`. + +This check covers the usage that the OpenAI-compatible, Responses, and Anthropic +clients return. Hermes' native Bedrock Converse and Gemini transports fill a +count the provider left out with 0, so the adapter reads which counts their raw +usage reported before they do: Bedrock's input needs `inputTokens`, its cached +input `cacheReadInputTokens` and its output `outputTokens`, and a Converse +stream without its metadata event reports none; Gemini's input needs +`promptTokenCount`, its cached input `cachedContentTokenCount`, and its output +`candidatesTokenCount` or `thoughtsTokenCount`. Hermes' Anthropic transport +replaces a stream event's `usage: null`, which some Anthropic-compatible +gateways send, with zeros or with the start event's output count before the SDK +builds the final message, so the adapter cannot see it. A null start usage +leaves input and cached input unknown, because Hermes' filled usage has no cache +fields, but a null final usage leaves a wrong output count. The Anthropic API +itself always sends both. Hermes' Codex app-server runtime records usage outside +this path. Hermes passes a response's raw usage to no per-agent hook, so the +adapter wraps `record_response_usage` and the three handlers for invalid, +truncated, and refused responses in `agent.turn_response_check`, and the +iteration-limit summary's `_summary_text` and attempt builders in +`agent.chat_completion_helpers`, `tools.delegate_tool._build_child_agent`, and +the native usage readers in `agent.bedrock_adapter` and +`agent.gemini_native_adapter`, once per process. The runtime fails to start if +the installed Hermes no longer calls them by those names. ## Execution Model diff --git a/adapters/typescript/pi/README.md b/adapters/typescript/pi/README.md index 5e6acb82a..4b8851972 100644 --- a/adapters/typescript/pi/README.md +++ b/adapters/typescript/pi/README.md @@ -277,9 +277,18 @@ The result has no `usage` when Pi produced no assistant message, when a model call reported none of the counts, when an assistant message was aborted, when a compaction was aborted or reported no usage, or when Pi retried a compaction summary call, because Pi keeps no usage for the failed attempt. It also has no -`usage` when Pi recorded a cache warm-up during the invocation: Pi replays a -warm-up outside the agent's stream function, so the adapter cannot tie its raw -report to the warm-up's usage entry. +`usage` when Pi attempted a cache warm-up since the previous invocation. Pi +replays a warm-up outside the agent's stream function, so the adapter cannot tie +its raw report to the warm-up's usage entry, and a warm-up that fails, or that +Pi aborts and discards when the agent settles, writes no usage entry at all. The +adapter notes every warm-up attempt before it is sent: each built-in protocol +passes a call's request body to `onPayload` once, so a later `onPayload` on the +same request options is a warm-up. Pi's model catalog gives a prompt cache +lifetime only to Anthropic models, and in its default `streaming` mode Pi warms +only while a step outlasts nine tenths of that lifetime, such as a tool that +runs longer than 270 seconds. A count whose sum over the invocation's calls is +not a safe integer is left out, and so is a total that would not be one; the +other counts stay. ## Custom Tool Modules From 49c33a6b0966f27b45581180d85ef22fe90a8e3d Mon Sep 17 00:00:00 2001 From: mnajafian-nv Date: Tue, 20 Oct 2026 08:21:12 -0700 Subject: [PATCH 43/68] fix: account for Pi warm-up attempts when a prompt throws The handle compared warm-up attempts only after session.prompt returned, so a prompt that started a warm-up and then threw left the difference for the next prompt, which reported unknown usage although it did not warm. Each prompt now records whether it warmed and advances the counter in its cleanup, so only the prompt that owns a warm-up is unknown. Signed-off-by: mnajafian-nv --- adapters/typescript/pi/src/pi-sdk.ts | 15 +++-- adapters/typescript/pi/test/pi-sdk.test.mjs | 71 +++++++++++++++++++++ 2 files changed, 80 insertions(+), 6 deletions(-) diff --git a/adapters/typescript/pi/src/pi-sdk.ts b/adapters/typescript/pi/src/pi-sdk.ts index 51b501ec0..1fe44b44d 100644 --- a/adapters/typescript/pi/src/pi-sdk.ts +++ b/adapters/typescript/pi/src/pi-sdk.ts @@ -1456,7 +1456,7 @@ function unsupportedSessionAction(name: string): never { throw new LifecycleError("pi_unsupported_session_operation", `Pi session operation ${name} is not supported`); } -class PiSdkSessionHandle implements PiSessionHandle { +export class PiSdkSessionHandle implements PiSessionHandle { readonly relay?: PiRelayRuntime; private readonly session: AgentSession; private readonly state: { shutdownRequested: boolean }; @@ -1531,6 +1531,7 @@ class PiSdkSessionHandle implements PiSessionHandle { this.compaction = expected === undefined ? undefined : { expected, calls: [] }; }; this.compactionListeners.add(onCompactionPrepared); + let warmedUp = false; try { await this.session.prompt(text, { expandPromptTemplates: true, @@ -1543,12 +1544,14 @@ class PiSdkSessionHandle implements PiSessionHandle { this.compactionListeners.delete(onCompactionPrepared); this.compaction = undefined; unsubscribe(); + // A warm-up's tokens reach no usage entry the adapter can check, and a + // failed or discarded one reaches none at all, so a prompt with a + // warm-up attempt since the previous prompt has unknown tokens. Every + // prompt accounts for its attempts, including one that throws, so a + // later prompt does not inherit them. + warmedUp = this.warmUpAttempts !== this.warmUpAttemptsCounted; + this.warmUpAttemptsCounted = this.warmUpAttempts; } - // A warm-up's tokens reach no usage entry the adapter can check, and a - // failed or discarded one reaches none at all, so a prompt with a warm-up - // attempt since the previous prompt has unknown tokens. - const warmedUp = this.warmUpAttempts !== this.warmUpAttemptsCounted; - this.warmUpAttemptsCounted = this.warmUpAttempts; return { accepted, turnStarted, diff --git a/adapters/typescript/pi/test/pi-sdk.test.mjs b/adapters/typescript/pi/test/pi-sdk.test.mjs index 89fb8508a..677506694 100644 --- a/adapters/typescript/pi/test/pi-sdk.test.mjs +++ b/adapters/typescript/pi/test/pi-sdk.test.mjs @@ -17,6 +17,7 @@ import { endedCallCounts, modelAwareCompactionReserveTokens, PiSdkSessionFactory, + PiSdkSessionHandle, recordReportedCounts, reportedCountsFrom, resolveCustomTools, @@ -2063,6 +2064,76 @@ async function warmedCall({ warmUp, settleFirst = false }) { } } +/** + * A Pi session stand-in whose prompts each run one model call through the + * agent's stream function and reports 10 input and 2 output tokens. A + * "warm" entry also replays the call's request body the way Pi's cache + * warmer does, "warm-and-throw" then fails the prompt, and "answer" does + * neither. + */ +function scriptedPiSession(script) { + const listeners = new Set(); + const model = { id: "m", api: "openai-completions", provider: "test" }; + let step; + const agent = { + streamFunction: async (_model, _context, options) => { + await options.onPayload?.({}, model); + if (step === "warm" || step === "warm-and-throw") { + await options.onPayload?.({}, model); + } + await options.onProviderStreamEvent?.( + { choices: [], usage: { prompt_tokens: 10, completion_tokens: 2, prompt_tokens_details: { cached_tokens: 0 } } }, + model, + ); + const message = { + role: "assistant", + content: [{ type: "text", text: "done" }], + stopReason: "stop", + usage: { input: 10, output: 2, cacheRead: 0, cacheWrite: 0 }, + }; + return { result: async () => message }; + }, + }; + return { + agent, + subscribe(listener) { + listeners.add(listener); + return () => listeners.delete(listener); + }, + async prompt(_text, options) { + options.preflightResult?.(); + step = script.shift(); + const response = await agent.streamFunction(model, { messages: [] }, {}); + const message = await response.result(); + await new Promise((resolve) => setImmediate(resolve)); + if (step === "warm-and-throw") { + throw new Error("prompt failed"); + } + for (const listener of listeners) { + listener({ type: "turn_start" }); + listener({ type: "message_end", message }); + } + }, + }; +} + +test("a prompt that throws after a warm-up attempt does not leave the next prompt unknown", async () => { + const session = scriptedPiSession(["warm-and-throw", "answer", "answer", "warm", "answer"]); + const handle = new PiSdkSessionHandle(session, { shutdownRequested: false }, tmpdir(), new Set()); + + await assert.rejects(handle.prompt("one"), /prompt failed/); + const second = await handle.prompt("two"); + const third = await handle.prompt("three"); + const warmed = await handle.prompt("four"); + const fifth = await handle.prompt("five"); + + assert.deepEqual(second.usage, { input: 10, cacheRead: 0, output: 2 }); + assert.deepEqual(third.usage, { input: 10, cacheRead: 0, output: 2 }); + // The prompt that owns a warm-up attempt stays unknown. + assert.equal(warmed.usage, undefined); + assert.deepEqual(fifth.usage, { input: 10, cacheRead: 0, output: 2 }); +}); + for (const warmUp of ["fail", "hold", "answer"]) { test(`notes a cache warm-up attempt that Pi's warmer ${warmUp === "hold" ? "aborts" : warmUp === "fail" ? "fails" : "completes"}`, async () => { const { warmUps, received } = await warmedCall({ warmUp }); From d690432b93838381100c2943d693e183cbde37d8 Mon Sep 17 00:00:00 2001 From: mnajafian-nv Date: Tue, 20 Oct 2026 08:21:35 -0700 Subject: [PATCH 44/68] docs: state what the Hermes adapter checks at startup The README said startup fails whenever Hermes stops calling any wrapped function by name, which overstated the native usage-reader checks. It now says which interception points startup checks and that the checks do not detect every internal Hermes change. Signed-off-by: mnajafian-nv --- adapters/python/hermes/README.md | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/adapters/python/hermes/README.md b/adapters/python/hermes/README.md index 514bcde24..3525fd4b1 100644 --- a/adapters/python/hermes/README.md +++ b/adapters/python/hermes/README.md @@ -179,8 +179,10 @@ truncated, and refused responses in `agent.turn_response_check`, and the iteration-limit summary's `_summary_text` and attempt builders in `agent.chat_completion_helpers`, `tools.delegate_tool._build_child_agent`, and the native usage readers in `agent.bedrock_adapter` and -`agent.gemini_native_adapter`, once per process. The runtime fails to start if -the installed Hermes no longer calls them by those names. +`agent.gemini_native_adapter`, once per process. Startup checks the +response-handler, summary and delegation interception points, and requires the +native usage-reader functions to exist. These checks do not detect every +internal Hermes change. ## Execution Model From 4585ac7ee03f0d055f71fc8dd46eef91778fbd97 Mon Sep 17 00:00:00 2001 From: mnajafian-nv Date: Tue, 20 Oct 2026 08:51:15 -0700 Subject: [PATCH 45/68] 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 46/68] 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 47/68] 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 48/68] 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 49/68] 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 50/68] 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 51/68] 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 52/68] 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 53/68] 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 54/68] 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 55/68] fix: preserve fractional cost precision across usage reducers Signed-off-by: mnajafian-nv --- adapter-contract/typescript/src/usage.ts | 38 ++++++++++++++++++- .../usage_accounting/reducer-cases.json | 36 ++++++++++++++++++ 2 files changed, 73 insertions(+), 1 deletion(-) diff --git a/adapter-contract/typescript/src/usage.ts b/adapter-contract/typescript/src/usage.ts index 726b5aff0..cb8f0dd9f 100644 --- a/adapter-contract/typescript/src/usage.ts +++ b/adapter-contract/typescript/src/usage.ts @@ -44,6 +44,42 @@ function validate(record: UsageObservation): void { } } +// Keep low-order cost contributions until the final round, matching math.fsum. +// Inputs have already passed the finite, nonnegative cost validation. +function sumCosts(values: readonly number[]): number { + const partials: number[] = []; + for (let high of values) { + let retained = 0; + for (let low of partials) { + if (Math.abs(high) < Math.abs(low)) [high, low] = [low, high]; + const sum = high + low; + const residual = low - (sum - high); + if (residual !== 0) partials[retained++] = residual; + high = sum; + } + partials.length = retained; + partials.push(high); + } + let index = partials.length; + let total = index === 0 ? 0 : partials[--index]!; + let residual = 0; + while (index > 0) { + const high = total; + const low = partials[--index]!; + total = high + low; + residual = low - (total - high); + if (residual !== 0) break; + } + // Preserve round-to-even when the remaining low-order terms share a sign. + if (index > 0 && ((residual < 0 && partials[index - 1]! < 0) || + (residual > 0 && partials[index - 1]! > 0))) { + const correction = residual * 2; + const rounded = total + correction; + if (correction === rounded - total) total = rounded; + } + return total; +} + function reduce(records: readonly UsageObservation[]): { reported: ReportedUsage; observed: ReportedUsage; gaps: UsageGap[] } { const reported: ReportedUsage = {}; const observed: ReportedUsage = {}; @@ -58,7 +94,7 @@ function reduce(records: readonly UsageObservation[]): { reported: ReportedUsage let value: number | undefined; if (values.length) { if (key === "cost_usd") { - const total = values.reduce((sum, item) => sum + item, 0); + const total = sumCosts(values); if (Number.isFinite(total) && total * 1e9 < 2 ** 64) value = total; } else if (key === "peak_request_input_tokens") { value = values.reduce((maximum, item) => Math.max(maximum, item)); diff --git a/tests/fixtures/usage_accounting/reducer-cases.json b/tests/fixtures/usage_accounting/reducer-cases.json index e79697c63..cba63db2e 100644 --- a/tests/fixtures/usage_accounting/reducer-cases.json +++ b/tests/fixtures/usage_accounting/reducer-cases.json @@ -437,5 +437,41 @@ "observed_input": 100, "exhaustive": true } + }, + { + "name": "reported_cost_keeps_fractional_sum_precision", + "records": [ + { + "model": "fixture", + "provider": "fixture", + "tokens_aligned": true, + "reported": { + "cost_usd": 0.1 + } + }, + { + "model": "fixture", + "provider": "fixture", + "tokens_aligned": true, + "reported": { + "cost_usd": 0.2000000005 + } + }, + { + "model": "fixture", + "provider": "fixture", + "tokens_aligned": true, + "reported": { + "cost_usd": 0.1 + } + } + ], + "expected": { + "input_tokens": null, + "cost_usd": 0.4000000005, + "observed_input": null, + "observed_output": null, + "exhaustive": true + } } ] From a6318d4dede15289bbf31485296c6bec8bc524b0 Mon Sep 17 00:00:00 2001 From: mnajafian-nv Date: Thu, 22 Oct 2026 16:24:17 -0700 Subject: [PATCH 56/68] 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 57/68] 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 648cbed9be3ef0f53bda49c58929a9ff75af4c86 Mon Sep 17 00:00:00 2001 From: mnajafian-nv Date: Thu, 22 Oct 2026 17:06:44 -0700 Subject: [PATCH 58/68] docs: clarify collector pricing scope Signed-off-by: mnajafian-nv --- adapters/typescript/pi/README.md | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/adapters/typescript/pi/README.md b/adapters/typescript/pi/README.md index 5b71c388c..da073c0d0 100644 --- a/adapters/typescript/pi/README.md +++ b/adapters/typescript/pi/README.md @@ -319,8 +319,8 @@ observation without certifying all attempts. Accounting provider identity comes from a qualified HTTPS endpoint and protocol, not Pi's configured provider label. Native OpenAI, Anthropic, Google Generative AI, and Mistral routes qualify; custom or unverified routes retain their literal model -and counts without inheriting native-provider pricing. Such routes can use an -explicit user pricing catalog. Overrides and cache warm-ups follow the same rule. +and counts without inheriting native-provider pricing. The adapter retains this evidence for separate +pricing; it does not compute catalog estimates. Overrides and cache warm-ups follow the same rule. Redirects and mixed routes cannot restore a native identity after qualification fails. This changes accounting identity only; model execution settings are preserved. From 93b00ecdbf82a618e8b6947692b9ea333f59f38b Mon Sep 17 00:00:00 2001 From: mnajafian-nv Date: Thu, 22 Oct 2026 17:23:32 -0700 Subject: [PATCH 59/68] 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 60/68] 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 61/68] 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 90e7be3123d9a76cc5b14f634b8acafac1813882 Mon Sep 17 00:00:00 2001 From: mnajafian-nv Date: Fri, 23 Oct 2026 00:04:13 -0700 Subject: [PATCH 62/68] fix: preserve custom Pi fetch during cache warming Signed-off-by: mnajafian-nv --- adapters/typescript/pi/src/usage.ts | 19 +++++++++++++++++-- adapters/typescript/pi/test/usage.test.mjs | 22 ++++++++++++++++++++++ 2 files changed, 39 insertions(+), 2 deletions(-) diff --git a/adapters/typescript/pi/src/usage.ts b/adapters/typescript/pi/src/usage.ts index 456a9c2ea..1cc215727 100644 --- a/adapters/typescript/pi/src/usage.ts +++ b/adapters/typescript/pi/src/usage.ts @@ -54,6 +54,20 @@ const observedFetch: typeof globalThis.fetch = async (input, init) => { call.attempts[index] = response.status; return response; }; +const observedFetches = new WeakMap(); +const observedDelegates = new WeakMap(); + +function observedFetchFor(delegate: typeof globalThis.fetch | undefined): typeof globalThis.fetch { + if (delegate === undefined) return observedFetch; + let wrapper = observedFetches.get(delegate); + if (wrapper === undefined) { + wrapper = (input, init) => attemptScope.getStore() === undefined + ? delegate(input, init) : observedFetch(input, init); + observedFetches.set(delegate, wrapper); + observedDelegates.set(wrapper, delegate); + } + return wrapper; +} const text = (value: unknown): string | undefined => typeof value === "string" && value.trim() ? value : undefined; const count = (value: unknown): number | undefined => typeof value === "number" && Number.isSafeInteger(value) && value >= 0 ? value : undefined; const object = (value: unknown): Record => typeof value === "object" && value !== null ? value as Record : {}; @@ -228,11 +242,12 @@ export class PiUsageCollector { let current = initial; let payloads = 0; const ownFetch = options?.fetch; - const scope: AttemptScope = { delegate: ownFetch === undefined || ownFetch === observedFetch ? globalThis.fetch : ownFetch, current: () => current }; + const delegate = ownFetch === undefined || ownFetch === observedFetch ? undefined : observedDelegates.get(ownFetch) ?? ownFetch; + const scope: AttemptScope = { delegate: delegate ?? globalThis.fetch, current: () => current }; try { const response = await attemptScope.run(scope, () => stream(model, context, { ...options, - ...(FETCH_PROTOCOLS.has(model.api) ? { fetch: observedFetch } : {}), + ...(FETCH_PROTOCOLS.has(model.api) ? { fetch: observedFetchFor(delegate) } : {}), onPayload: async (payload, payloadModel) => { if (++payloads > 1) { // Pi replays these options for a cache warm-up outside streamFunction. diff --git a/adapters/typescript/pi/test/usage.test.mjs b/adapters/typescript/pi/test/usage.test.mjs index c2343608f..ca2f896a7 100644 --- a/adapters/typescript/pi/test/usage.test.mjs +++ b/adapters/typescript/pi/test/usage.test.mjs @@ -302,6 +302,28 @@ test("internal retries retain an aligned final-attempt snapshot for observed pri assert.equal(group.requests[0].gap, "scope_unverified"); }); +test("cache-warm replay keeps a stable caller fetch outside the invocation", async () => { + let replayFetch; + let requests = 0; + const ownFetch = async () => { + requests += 1; + return new Response("", { status: 200 }); + }; + const harness = source(async (selected, _context, options) => { + replayFetch = options.fetch; + await replayFetch("data:,initial"); + await options.onPayload({}, selected); + await options.onProviderStreamEvent(raw(), selected); + return { result: async () => message() }; + }); + await harness.collector.prompt(() => harness.call(model, { fetch: ownFetch }), harness.events); + await replayFetch("data:,warm"); + const firstFetch = replayFetch; + await harness.collector.prompt(() => harness.call(model, { fetch: ownFetch }), harness.events); + assert.equal(replayFetch, firstFetch); + assert.equal(requests, 3); +}); + test("Pi's actual resolved result promise is recorded before message_end without a timing wait", async () => { const harness = source(async (selected, _context, options) => { await options.onPayload({}, selected); From e3a2355b585e1eb231a7395cb1f4e7945c76d12b Mon Sep 17 00:00:00 2001 From: mnajafian-nv Date: Fri, 23 Oct 2026 01:41:33 -0700 Subject: [PATCH 63/68] 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 64/68] 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 65/68] 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 66/68] 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 67/68] 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 From 2149bd1a0e98d996f64065724c3658e6ea43f06d Mon Sep 17 00:00:00 2001 From: mnajafian-nv Date: Fri, 23 Oct 2026 16:29:34 -0700 Subject: [PATCH 68/68] fix: bound Pi Relay readiness checks Signed-off-by: mnajafian-nv --- adapters/typescript/pi/src/relay-gateway.ts | 31 +++++++-- adapters/typescript/pi/test/relay.test.mjs | 71 +++++++++++++++++++++ 2 files changed, 95 insertions(+), 7 deletions(-) diff --git a/adapters/typescript/pi/src/relay-gateway.ts b/adapters/typescript/pi/src/relay-gateway.ts index 163bcac09..99446bf14 100644 --- a/adapters/typescript/pi/src/relay-gateway.ts +++ b/adapters/typescript/pi/src/relay-gateway.ts @@ -147,26 +147,43 @@ export async function waitForRelayGateway( const fetchRequest = options.fetch ?? fetch; const pollIntervalMs = options.pollIntervalMs ?? 100; const deadline = performance.now() + timeoutMs; - while (performance.now() < deadline) { + while (true) { + const remaining = deadline - performance.now(); + if (remaining <= 0) { + break; + } if (processExited(child)) { throw new RelayGatewayError( `NeMo Relay gateway exited with status ${child.exitCode ?? child.signalCode} before becoming ready`, ); } + let ready = false; try { const response = await fetchRequest(healthUrl, { method: "HEAD", - signal: AbortSignal.timeout(1_000), + signal: AbortSignal.timeout(Math.min(1_000, Math.ceil(remaining))), }); - const ready = response.ok; await response.body?.cancel(); - if (ready) { - return; - } + ready = response.ok; } catch { // Retry until the gateway exits or the readiness deadline expires. } - await delay(Math.min(pollIntervalMs, Math.max(0, deadline - performance.now()))); + if (ready) { + if (processExited(child)) { + throw new RelayGatewayError( + `NeMo Relay gateway exited with status ${child.exitCode ?? child.signalCode} before becoming ready`, + ); + } + if (performance.now() >= deadline) { + break; + } + return; + } + const remainingAfterProbe = deadline - performance.now(); + if (remainingAfterProbe <= 0) { + break; + } + await delay(Math.min(pollIntervalMs, remainingAfterProbe)); } throw new RelayGatewayError(`NeMo Relay gateway did not become ready at ${healthUrl}`); } diff --git a/adapters/typescript/pi/test/relay.test.mjs b/adapters/typescript/pi/test/relay.test.mjs index ee0d541d5..b77e3c459 100644 --- a/adapters/typescript/pi/test/relay.test.mjs +++ b/adapters/typescript/pi/test/relay.test.mjs @@ -844,6 +844,77 @@ test("drains unsuccessful and successful Relay health responses", async () => { assert.deepEqual(cancelled, [0, 1]); }); +test("rejects a gateway exit during a successful health probe", async (t) => { + for (const boundary of ["fetch", "body cancellation"]) { + await t.test(boundary, async () => { + const child = new MockChild(); + await assert.rejects( + waitForRelayGateway(child, "http://127.0.0.1:41000/healthz", { + async fetch() { + if (boundary === "fetch") child.exitCode = 17; + return { + ok: true, + body: { + async cancel() { + if (boundary === "body cancellation") child.exitCode = 17; + }, + }, + }; + }, + }), + /exited with status 17 before becoming ready/u, + ); + }); + } +}); + +test("rejects a successful health probe at the readiness deadline", async (t) => { + for (const boundary of ["fetch", "body cancellation"]) { + await t.test(boundary, async (subtest) => { + let now = 0; + subtest.mock.method(performance, "now", () => now); + await assert.rejects( + waitForRelayGateway(new MockChild(), "http://127.0.0.1:41000/healthz", { + timeoutMs: 25, + async fetch() { + if (boundary === "fetch") now = 25; + return { + ok: true, + body: { + async cancel() { + if (boundary === "body cancellation") now = 25; + }, + }, + }; + }, + }), + /did not become ready/u, + ); + }); + } +}); + +test("limits a health probe to the remaining readiness budget", async (t) => { + let now = 0; + t.mock.method(performance, "now", () => now); + const timeouts = []; + t.mock.method(AbortSignal, "timeout", (milliseconds) => { + timeouts.push(milliseconds); + return new AbortController().signal; + }); + await assert.rejects( + waitForRelayGateway(new MockChild(), "http://127.0.0.1:41000/healthz", { + timeoutMs: 25, + async fetch() { + now = 25; + throw new Error("still initializing"); + }, + }), + /did not become ready/u, + ); + assert.deepEqual(timeouts, [25]); +}); + test("escalates gateway shutdown from terminate to kill and is safe after exit", async () => { const mockChild = new MockChild(); mockChild.kill = function (signal) {