Repository navigation
fix: report NOOA BenchAgent usage from model responses - #441
Open
mnajafian-nv wants to merge 105 commits into
Open
mnajafian-nv wants to merge 105 commits into
mnajafian-nv wants to merge 105 commits into
Conversation
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 <mnajafian@nvidia.com>
NOOA 0.0.10 emits token_usage events only for CompletionClient. ResponsesClient returns usage on each response but does not feed the task token accumulator, so Fabric reported input, output and total tokens as zero for Responses client runs. When the selected model is a ResponsesClient, wrap its synchronous and asynchronous calls and collect each returned usage report in an invocation-scoped ContextVar. Sum provider-reported input, cached-input, output and total tokens. If any returned report omits a counter, leave that dimension unknown. Preserve explicit zeros and mark input tokens as including cached tokens. The CompletionClient path is unchanged. Tests drive the real ResponsesClient directly with queued LiteLLM responses while leaving NOOA's accumulator unpatched. They cover multi-call sums, explicit zeros, missing usage, realistic missing cached-token details, and a defensive summing case for an omitted output count. Signed-off-by: mnajafian-nv <mnajafian@nvidia.com>
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
WalkthroughAdapters collect invocation-scoped model usage evidence. Shared reducers and Core distinguish complete counts from observations, preserve dimension-specific gaps, and expose normalized usage in run results. The Python SDK adds validated usage models and coverage-aware summaries against an expected invocation inventory. ChangesInvocation Usage Accounting
Priority: Estimated code review effort: 5 (Critical) | ~90 minutes Sequence Diagram(s)sequenceDiagram
participant Adapter
participant UsageAccounting
participant FabricCore
participant RunResult
Adapter->>UsageAccounting: Supply invocation usage evidence
UsageAccounting->>FabricCore: Provide reported values, observations, and gaps
FabricCore->>FabricCore: Validate and normalize supported evidence
FabricCore->>RunResult: Set optional model_usage
Merge Risk:
 |
| Check name | Status | Explanation | Resolution |
|---|---|---|---|
| Docstring Coverage |  |
Docstring coverage is 17.94% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 301 functions across 46 files. (1 skipped… | Write docstrings for the functions missing them to satisfy the coverage threshold. |
-
Check name Status Explanation Title check The title follows Conventional Commits format, uses the allowed lowercase type fix, clearly describes the NOOA BenchAgent usage change, and is 54 characters without a trailing period.Linked Issues check Check skipped because no linked issues were found for this pull request. Out of Scope Changes check Check skipped because no linked issues were found for this pull request. Description check The description follows the repository template. It has Overview, Where should the reviewer start, and Related Issues sections, and the checklist is ticked. It explains the change: NOOA BenchAgent usa… Full details: Description check
-
Explanation
The description follows the repository template. It has Overview, Where should the reviewer start, and Related Issues sections, and the checklist is ticked. It explains the change: NOOA BenchAgent usage is read from model responses before LiteLLM fills defaults, and it covers how unknown, zero, and unfinished-work cases are handled. It also gives a dependency on
#458and#442, the validation commands, and the scope exclusions for streaming, unmanaged clients, and custom-endpoint pricing. The Details section is written as bullet-free prose, and the Validation section is brief but usable. Minor gaps: the 'Related Issues' line uses 'Relates to#489' but the dependency on#442/#458 is only described in Overview, not linked as a related issue, and the Validation section does not list the 54-test or 379-test counts that the objectives mention.
-
Full details: Docstring Coverage
-
Explanation
Docstring coverage is 17.94% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 301 functions across 46 files. (1 skipped: 1 unsupported.)
Resolution
Write docstrings for the functions missing them to satisfy the coverage threshold.
✨ Finishing Touches: 💡 1
-
🧪 Generate unit tests (beta)
-
- Create a new PR
🛠️ Fix failing CI checks: 💡
-
- Commit to this branch
- Create a new PR
-
- Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts
Comment @coderabbitai help to get the list of available commands.
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
This comment was marked as resolved.
This comment was marked as resolved.
Sorry, something went wrong.
|
Fern docs preview: https://nvidia-preview-pull-request-441.docs.buildwithfern.com/nemo/fabric |
mnajafian-nv
changed the title
A RunResult says which agent and adapter ran, but not which Fabric version, models or MCP servers were behind it, so a rerun weeks later cannot tell what changed. Add an optional, versioned what_ran object. Each fact holds a value, or a state and a reason, plus how Fabric learned it: its source, collector and method. Core fills the facts it holds before it sends the request: its own version, the adapter ID, the requested models, the configured MCP servers with secret-free endpoints, and the environment provider. Core reads no adapter report yet, so adapter-reported facts read unsupported with reason key_not_declared. On the OpenShell host path, the host relabels the task core's facts as task_core, records its own version as host_core_version and sets the environment provider to openshell. A task result without a readable what_ran keeps what_ran absent. The Rust and Python readers apply the same structural checks, tested against one shared vector file. A malformed what_ran reads as absent and never fails the enclosing result. Core gains a direct dependency on sha2 0.10, already in the lockfile, for the string length rule. Rust code that builds RunResult with a struct literal must add the what_ran field. Python's RunResult.to_mapping() now always includes it. Signed-off-by: mnajafian-nv <mnajafian@nvidia.com>
Explain the fact envelope, the four no-value states, the endpoint rule, the OpenShell attribution and the structural read, and keep the integration skill's RunResult table in step with the guide. Signed-off-by: mnajafian-nv <mnajafian@nvidia.com>
Regenerate the Rust library reference with Rust 1.94.0 for the new what_ran module and the RunResult field, and the Python types reference for RunResult.what_ran. Only the pages this change affects are included. Signed-off-by: mnajafian-nv <mnajafian@nvidia.com>
A plan with many model roles or MCP servers, or very long names, could make what_ran arbitrarily large, and every RunResult carries it. Keep it bounded without ever failing a run. Core keeps the first 16 model roles and 32 MCP servers in sorted key order, and drops entries whose key is longer than 128 characters. While the compact JSON of what_ran is larger than 128 KiB, it drops MCP servers from the end of the sorted order, then model roles. If plan identifiers alone still exceed the budget, what_ran becomes a minimal object with only the schema name, the version and what_ran_truncated. A marker next to each map counts the dropped entries and names the last rule applied. The OpenShell host applies the same bounds again after it attributes the task core's facts, so they hold for a task core of any version. The six what_ran sections become optional in the Rust type so the minimal object can be represented. Every section is still required unless what_ran_truncated is present, and no section may be null. The Rust reader, the generated schema and the Python reader check this against the shared vectors. Signed-off-by: mnajafian-nv <mnajafian@nvidia.com>
Explain the entry caps, the key length rule, the byte budget, the truncation markers and the minimal object, and tell readers to check what_ran_truncated before reading a section. Keep the integration skill's RunResult notes in step. Signed-off-by: mnajafian-nv <mnajafian@nvidia.com>
Regenerate the Rust library reference pages of the what_ran module for TruncatedEntries and the optional sections, and update the Python types reference for RunResult.what_ran. Only the pages this change affects are included. Signed-off-by: mnajafian-nv <mnajafian@nvidia.com>
Signed-off-by: mnajafian-nv <mnajafian@nvidia.com>
The Responses usage tests drove only acall, and each one started a fresh runtime. The adapter also wraps the synchronous call, and a persistent runtime serves many invocations with the same wrapped client, so neither path was pinned down. The BenchAgent double now also drains a queue of synchronous replies through ResponsesClient.call, and a new test checks that a sync call's counters reach the invocation usage. A second test invokes one runtime twice with different replies and checks each result's usage on its own, so usage from the first invocation cannot carry over into the second. Signed-off-by: mnajafian-nv <mnajafian@nvidia.com>
Run a local host invocation with two model roles and three MCP servers (npx with arguments, streamable-http and sse) and no environment section, and assert the exact what_ran a user gets, including the origin of every unsupported adapter-reported fact and that no configured secret appears. Signed-off-by: mnajafian-nv <mnajafian@nvidia.com>
Resolve a plan from each of the 16 bundled descriptors with a model and, where the adapter accepts MCP, three servers, and check that what_ran passes the schema, carries no truncation marker, and reads every adapter-reported fact as unsupported/key_not_declared. Also show that 16 roles and 32 servers with 128-character keys fit the byte budget, and that stdio placeholders read executable_unparsed. Signed-off-by: mnajafian-nv <mnajafian@nvidia.com>
2 tasks done
Loop over whatever the catalog lists without pinning its size, so adding or removing a bundled adapter doesn't break the test. Name the size test after what it builds: ASCII keys and model names at their maximum length. Signed-off-by: mnajafian-nv <mnajafian@nvidia.com>
2 tasks done
Adapters had no way to report their own facts, so the adapter package version, the harness name and version, the served model and the MCP runtime facts always read unsupported with reason key_not_declared. Core now reads nemo_fabric.versions.v1 and nemo_fabric.mcp.v1 from a result's extensions after the descriptor's extension check, when the descriptor's run_result schema names the key under properties. Only adapter-owned facts with source reported and collector adapter are kept; a malformed or forged fact reads invalid_report, an omitted one fact_omitted, an omitted key key_omitted, and a key over 64 KiB too_large. A reported server the plan does not configure becomes a harness entry. tool_names keeps 64 sorted names and counts the rest in a new tool_names_truncated marker, added to the schema and both readers. Ingestion never changes the run's status or error. Signed-off-by: mnajafian-nv <mnajafian@nvidia.com>
A what_ran could not say which harness build ran. Six adapters now declare nemo_fabric.versions.v1 and report their package version and the harness name and version from a source that names the build that ran: - Claude Code: claude_code_version from the CLI's init message. - Codex: the thread's cli_version from one bounded thread/read request. - Hermes: get_version_info(), a blocking call that runs git, on a daemon thread with a 5 second limit. - OpenClaw: the --version check, or an attached gateway's /startupz version, both already read at startup. - Deep Agents and mini-swe-agent: package metadata, accepted only when the distribution installed the module the adapter loaded. A shared helper builds the report, applies core's string length rule, and runs blocking calls with at most one outstanding probe per runtime; stopping a runtime never waits for it. A source that times out, fails or answers nothing reads a state with a reason, never a substitute value. The adapter catalog is regenerated for the six changed descriptors. Signed-off-by: mnajafian-nv <mnajafian@nvidia.com>
what_ran.models.<role>.served stayed unsupported for Claude Code even though each assistant message names the model that answered. The Claude Code adapter now reports, for the role it runs (default, or the sole configured role), the one distinct model named by this invocation's main-thread assistant messages. Messages with a parent_tool_use_id belong to subagents and are skipped, and so are Claude Code's <synthetic> API-error messages. Two distinct models read not_attributed, none reads empty_response, a failed stream reads collection_failed, and every invocation starts from its own messages. Tests feed saved Claude Code CLI streams through the SDK's parser. Signed-off-by: mnajafian-nv <mnajafian@nvidia.com>
The SDK guide and the integration reference said every adapter-reported fact reads key_not_declared. They now name the six adapters that report harness versions, the Claude Code served-model rule, the states a report key can produce and the tool_names cap. The adapter tutorial and the build-adapter skill show adapter authors how to declare and fill nemo_fabric.versions.v1. Signed-off-by: mnajafian-nv <mnajafian@nvidia.com>
Regenerated with just docs for the new report-key constants and the tool_names_truncated field. Signed-off-by: mnajafian-nv <mnajafian@nvidia.com>
NOOA, NOOA Bench, OpenHands and Remote Agent left adapter.package_version at key_not_declared, so a run with them could not say which adapter build ran. Each now declares nemo_fabric.versions.v1 and attaches its installed package version to every result. These adapters have no harness version or served-model source yet, so those facts read unsupported with reason no_source. The catalog is regenerated. Signed-off-by: mnajafian-nv <mnajafian@nvidia.com>
Cline, Droid, Kilo Code, OpenCode, Pi and Qwen Code left adapter.package_version at key_not_declared. A shared withPackageVersion wrapper in the TypeScript common package reads the adapter's own package.json synchronously at start and adds nemo_fabric.versions.v1 to every result, with the harness and served model facts unsupported/no_source until their sources land. Each CLI wraps its runtime, each descriptor declares the key, and the catalog is regenerated. A catalog test checks that all bundled descriptors declare the key. Signed-off-by: mnajafian-nv <mnajafian@nvidia.com>
The SDK guide and the integration reference named only the six adapters with a harness version source. The TypeScript common README documents withPackageVersion. Signed-off-by: mnajafian-nv <mnajafian@nvidia.com>
Signed-off-by: mnajafian-nv <mnajafian@nvidia.com>
Signed-off-by: mnajafian-nv <mnajafian@nvidia.com>
Signed-off-by: mnajafian-nv <mnajafian@nvidia.com>
Signed-off-by: mnajafian-nv <mnajafian@nvidia.com>
Signed-off-by: mnajafian-nv <mnajafian@nvidia.com>
Signed-off-by: mnajafian-nv <mnajafian@nvidia.com>
mnajafian-nv
marked this pull request as ready for review
October 23, 2026 16:37
mnajafian-nv
marked this pull request as draft
October 23, 2026 16:45
Signed-off-by: mnajafian-nv <mnajafian@nvidia.com>
mnajafian-nv
marked this pull request as ready for review
October 23, 2026 17:11
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 <mnajafian@nvidia.com>
Signed-off-by: mnajafian-nv <mnajafian@nvidia.com>
Signed-off-by: mnajafian-nv <mnajafian@nvidia.com>
Signed-off-by: mnajafian-nv <mnajafian@nvidia.com>
Signed-off-by: mnajafian-nv <mnajafian@nvidia.com>
Signed-off-by: mnajafian-nv <mnajafian@nvidia.com>
Signed-off-by: mnajafian-nv <mnajafian@nvidia.com>
Signed-off-by: mnajafian-nv <mnajafian@nvidia.com>
Signed-off-by: mnajafian-nv <mnajafian@nvidia.com>
Signed-off-by: mnajafian-nv <mnajafian@nvidia.com>
Signed-off-by: mnajafian-nv <mnajafian@nvidia.com>
Signed-off-by: mnajafian-nv <mnajafian@nvidia.com>
Signed-off-by: mnajafian-nv <mnajafian@nvidia.com>
Signed-off-by: mnajafian-nv <mnajafian@nvidia.com>
This branch was previously deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Overview
NOOA BenchAgent can complete model calls before its task fails. Its final task counters can lose those calls or turn missing provider fields into zeros. That makes failed tasks look cheaper than the available evidence supports.
Read model responses before LiteLLM inserts default values. Preserve usable counts on both successful and failed tasks. Expose normalized observations and complete totals through
RunResult.model_usage. When a call has no usable report or work remains unfinished, keep the available counts without claiming that they are the invocation's complete total.Dependency: #458 supplies adapter version reports and includes the shared usage contract and normalizer from #442. This change adds the harness collector. The diff against
mainincludes those unmerged prerequisites.Details
Observe synchronous and asynchronous model clients constructed by NVIDIA NeMo Fabric. For NOOA 0.0.10's managed JSON HTTP paths, collect response fields, request-level model selection and retry attempts. The shared normalizer handles cache and token conventions.
Each call belongs to the invocation that starts it. Its returned usage snapshot cannot change after return. Repeated response IDs within one managed call contribute once; conflicting reports leave an accounting gap. An unsuccessful HTTP response can retain useful usage, but its status alone never proves zero spend.
Track summaries and queue jobs started during the invocation. Unfinished work leaves complete totals unknown; a verified task with no model calls and no unfinished work reports zero tokens. Collection does not wait for or cancel background jobs.
Streaming requests and clients outside the managed transport cannot establish complete coverage. A tool or target that constructs its own client is outside this collector. Provider cost is not inferred, and a custom endpoint does not establish native-provider pricing identity.
Validation
Focused NOOA adapter, accounting, Relay compatibility and catalog tests passed. Tests distinguish omitted fields from reported zero and cover retries, parser failures, pending jobs and late responses.
Use NOOA's supported Python 3.12 or 3.13 environment to run the collector checks:
The Python CI passed. The shared TypeScript workspace audit remains blocked; #489 contains its proposed correction, which needs maintainer approval before integration.
Where should the reviewer start?
Start with
InvocationUsage.finishandinstall_model_observerinadapters/python/nooa/src/nemo_fabric_adapters/nooa/usage.py. Check how raw field presence, retries and unfinished calls affect complete totals. Then review the client wiring inmodel_support.pyand the managed-response cases intests/adapters/test_nooa_usage_accounting.py.Related Issues: (use one of the action keywords Closes / Fixes / Resolves / Relates to)
Relates to chore: update dependencies and resource sampling checks #489 (shared CI prerequisite).
I confirm this contribution is my own work, or I have the right to submit it under this project's license.
I searched existing issues and open pull requests, and this does not duplicate existing work.