Repository navigation
fix: isolate Droid token usage between invocations #451
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Open
mnajafian-nv
wants to merge
40
commits into
NVIDIA:main
Choose a base branch
from
mnajafian-nv:fix/droid-token-usage
base: main
Could not load branches
Branch not found: {{ refName }}
Loading
Could not load tags
Nothing to show
Loading
Are you sure you want to change the base?
Some commits from the old base branch may be removed from the timeline,
and old review comments may become outdated.
Open
Changes from all commits
Commits
Show all changes
40 commits
Select commit
Hold shift + click to select a range
9d6bc21
fix: count every Deep Agents model call in token usage
mnajafian-nv cdc0636
docs: name NeMo Fabric in the Deep Agents usage limits
mnajafian-nv fd184bd
fix: report Droid cache tokens in the standard usage fields
mnajafian-nv e23cf97
fix: leave Droid totals unknown when cache reads are unreported
mnajafian-nv 4585ac7
fix: report Deep Agents cached input tokens
mnajafian-nv fee4c92
fix: report zero Deep Agents usage when no model call ran
mnajafian-nv 2ee128d
fix: report unknown Deep Agents usage with an Agent Protocol subagent
mnajafian-nv 06de4c0
docs: say a failed Deep Agents call may have been billed
mnajafian-nv e535051
style: format the Deep Agents remote subagent check
mnajafian-nv cb42bfe
test: keep zero observed Deep Agents calls unknown with a remote suba…
mnajafian-nv ce74415
fix: preserve empty OpenShell environment arguments
mnajafian-nv 790f369
fix: avoid premature Relay gateway startup failures
mnajafian-nv 5db4917
docs: use the NeMo Fabric product name
mnajafian-nv 4bea319
feat: report complete and observed model usage
mnajafian-nv be32c6e
docs: use the full product name in Relay guidance
mnajafian-nv be56402
fix: isolate Droid token usage between invocations
mnajafian-nv cbb0699
fix: validate recovered usage and optional pricing rates
mnajafian-nv 559ce0a
fix: preserve fractional cost precision across usage reducers
mnajafian-nv 01b5b53
fix: apply shared usage accounting corrections
mnajafian-nv a6318d4
chore: check packaged usage schema declarations
mnajafian-nv 630820f
fix: return detached applied pricing rates in Python
mnajafian-nv 93b00ec
fix: expose omitted optional SDK usage attributes
mnajafian-nv 101bce6
fix: reject Relay exit during successful health probe
mnajafian-nv 2014cbe
fix: exclude unknown keys from request usage
mnajafian-nv 616b824
chore: incorporate Relay readiness prerequisite
mnajafian-nv 639df48
docs: name NeMo Fabric in Droid usage
mnajafian-nv a9fd48c
Merge shared accounting update into Droid usage
mnajafian-nv e3a2355
fix: classify empty usage coverage as unverified
mnajafian-nv c7b9589
chore: incorporate current main into OpenShell launcher fix
mnajafian-nv 8cc5c23
chore: incorporate updated OpenShell prerequisite
mnajafian-nv a8f6fd8
test: account for unverified empty usage coverage
mnajafian-nv 75d579a
chore: incorporate updated Relay prerequisite
mnajafian-nv 99407ae
chore: incorporate shared accounting corrections
mnajafian-nv 83d8ada
docs: keep adapter usage guidance within current contract
mnajafian-nv bedaf1c
chore: incorporate corrected usage guidance
mnajafian-nv ebcfe97
test: run OpenShell launcher checks with system Bash
mnajafian-nv 4286931
chore: include system Bash test prerequisite
mnajafian-nv 83b5e76
fix: reject Relay readiness after the startup deadline
mnajafian-nv 5702ead
chore: sync Relay startup prerequisite
mnajafian-nv 9e53d5d
chore: sync Relay startup prerequisite
mnajafian-nv File filter
Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
There are no files selected for viewing
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
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
238 changes: 238 additions & 0 deletions
238
adapter-contract/python/src/nemo_fabric_adapter_contract/usage.py
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
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,238 @@ | ||
| # SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. | ||
| # SPDX-License-Identifier: Apache-2.0 | ||
|
|
||
| """Shared raw-evidence reduction for Python collectors. | ||
|
|
||
| Normalization, aggregate/model reconciliation and pricing belong to core. | ||
| Collectors supply settled snapshots and explicit coverage gaps, never inferred | ||
| zero usage from HTTP status. One record represents one deduplicated source unit. | ||
| """ | ||
|
|
||
| from __future__ import annotations | ||
|
|
||
| from collections import defaultdict | ||
| from collections.abc import Iterable | ||
| from dataclasses import dataclass | ||
| import math | ||
|
|
||
| from nemo_fabric_adapter_contract.models import AccountedModel | ||
| from nemo_fabric_adapter_contract.models import AgentUsage | ||
| from nemo_fabric_adapter_contract.models import GapReason | ||
| from nemo_fabric_adapter_contract.models import ReportedUsage | ||
| from nemo_fabric_adapter_contract.models import UsageAccounting | ||
| from nemo_fabric_adapter_contract.models import UsageDimension | ||
| from nemo_fabric_adapter_contract.models import UsageGap | ||
| from nemo_fabric_adapter_contract.models import UsageRequest | ||
|
|
||
| ACCOUNTING_KEY = "nemo_fabric.usage_accounting.v1" | ||
| DETAIL_KEY = "nemo_fabric.detailed_token_usage.v1" | ||
|
|
||
|
|
||
| @dataclass(frozen=True, kw_only=True) | ||
| class UsageObservation: | ||
| """One source unit's deduplicated raw snapshot and completeness boundary.""" | ||
|
|
||
| reported: ReportedUsage | ||
| model: str | None = None | ||
| provider: str | None = None | ||
| gap: GapReason | None = None | ||
| tokens_aligned: bool = False | ||
| is_model_call: bool = False | ||
| gaps: tuple[UsageGap, ...] = () | ||
|
|
||
|
|
||
| def _reduce( | ||
| records: list[UsageObservation], | ||
| ) -> tuple[ReportedUsage, ReportedUsage, list[UsageGap]]: | ||
| complete: dict[str, object] = {} | ||
| observed: dict[str, object] = {} | ||
| gaps: dict[GapReason, list[UsageDimension]] = defaultdict(list) | ||
| for dimension in UsageDimension: | ||
| key = dimension.value | ||
| values = [getattr(record.reported, key) for record in records] | ||
| known = [value for value in values if value is not None] | ||
| total = None | ||
| if known: | ||
| if dimension is UsageDimension.COST_USD: | ||
| amount = math.fsum(known) | ||
| if math.isfinite(amount) and 0 <= amount * 1e9 < 1 << 64: | ||
| total = amount | ||
| elif dimension is UsageDimension.PEAK_REQUEST_INPUT_TOKENS: | ||
| total = max(known) | ||
| else: | ||
| amount = sum(known) | ||
| if amount <= (1 << 64) - 1: | ||
| total = amount | ||
| if total is not None: | ||
| observed[key] = total | ||
| causes = {record.gap for record in records if record.gap is not None} | ||
| causes.update( | ||
| gap.reason | ||
| for record in records | ||
| for gap in record.gaps | ||
| if dimension in gap.dimensions | ||
| ) | ||
| if len(known) != len(records): | ||
| causes.add(GapReason.MISSING_USAGE) | ||
| if known and total is None: | ||
| causes.add(GapReason.SCOPE_UNVERIFIED) | ||
| if total is not None and not causes: | ||
| complete[key] = total | ||
| elif not causes: | ||
| 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 ( | ||
| ("input_tokens_include_cache", "input_tokens"), | ||
| ("output_tokens_include_reasoning", "output_tokens"), | ||
| ): | ||
| flags = { | ||
| getattr(record.reported, flag) | ||
| for record in records | ||
| if getattr(record.reported, counter) is not None | ||
| } | ||
| if len(flags) == 1 and None not in flags: | ||
| observed[flag] = next(iter(flags)) | ||
| if counter in complete: | ||
| complete[flag] = observed[flag] | ||
| elif len( | ||
| [ | ||
| record | ||
| for record in records | ||
| if getattr(record.reported, counter) is not None | ||
| ] | ||
| ) == 1 and flags == {None}: | ||
| # One raw source remains useful even when its convention is unknown. | ||
| # Core withholds normalized values; do not discard literal evidence. | ||
| pass | ||
| elif counter in observed: | ||
| # Mixed conventions cannot be combined as one raw input/output count. | ||
| complete.pop(counter, None) | ||
| observed.pop(counter, None) | ||
| dimension = UsageDimension(counter) | ||
| if not all( | ||
| record.is_model_call and record.model is not None for record in records | ||
| ): | ||
| if dimension not in gaps[GapReason.SCOPE_UNVERIFIED]: | ||
| gaps[GapReason.SCOPE_UNVERIFIED].append(dimension) | ||
| return ( | ||
| ReportedUsage.from_mapping(complete), | ||
| ReportedUsage.from_mapping(observed), | ||
| [ | ||
| UsageGap(dimensions=dimensions, reason=reason) | ||
| for reason, dimensions in sorted(gaps.items()) | ||
| ], | ||
| ) | ||
|
|
||
|
|
||
| def _observed_tokens_aligned(records: list[UsageObservation]) -> bool: | ||
| # Shape alone never proves scope. The collector must first identify a | ||
| # coherent source snapshot; matching presence then keeps reductions aligned. | ||
| fields = [ | ||
| "input_tokens", | ||
| "cached_input_tokens", | ||
| "cache_write_input_tokens", | ||
| "output_tokens", | ||
| "peak_request_input_tokens", | ||
| ] | ||
| if any(r.reported.output_tokens_include_reasoning is False for r in records): | ||
| fields.append("reasoning_tokens") | ||
| relevant = [ | ||
| r | ||
| for r in records | ||
| if any(getattr(r.reported, field) is not None for field in fields) | ||
| ] | ||
| return ( | ||
| bool(relevant) | ||
| and all(r.tokens_aligned and not r.gaps for r in relevant) | ||
| and all( | ||
| len({getattr(r.reported, field) is not None for r in relevant}) == 1 | ||
| for field in fields | ||
| ) | ||
| ) | ||
|
|
||
|
|
||
| def _request_usage(record: UsageObservation) -> ReportedUsage: | ||
| # A scoped observation is not a complete per-call count for that dimension. | ||
| # Keep the literal value in model observations without allowing requests to | ||
| # re-certify it or use it as a complete pricing component. | ||
| fields = record.reported.to_mapping() | ||
| for gap in record.gaps: | ||
| for dimension in gap.dimensions: | ||
| fields.pop(dimension.value, None) | ||
| return ReportedUsage.from_mapping(fields) | ||
|
|
||
|
|
||
| def usage_from_observations(records: Iterable[UsageObservation]) -> AgentUsage: | ||
| """Reduce snapshots once, preserving field presence and independently known values. | ||
|
|
||
| Empty source coverage is unverified rather than assumed to mean zero calls. | ||
| A collector that proves no model work occurred can report an explicit zero | ||
| snapshot. This helper does not change scheduling or wait for late work. | ||
| """ | ||
|
|
||
| records = list(records) | ||
| reported, observed, gaps = _reduce(records) | ||
| groups: dict[tuple[str | None, str], list[UsageObservation]] = defaultdict(list) | ||
| for record in records: | ||
| if record.model is not None: | ||
| groups[(record.provider, record.model)].append(record) | ||
| models = [] | ||
| for (provider, model), group in sorted( | ||
| groups.items(), key=lambda item: (item[0][0] or "", item[0][1]) | ||
| ): | ||
| complete_model, observed_model, model_gaps = _reduce(group) | ||
| models.append( | ||
| AccountedModel( | ||
| model=model, | ||
| provider=provider, | ||
| reported=complete_model, | ||
| observed=observed_model, | ||
| observed_tokens_aligned=_observed_tokens_aligned(group), | ||
| gaps=model_gaps, | ||
| requests=[ | ||
| UsageRequest( | ||
| usage=_request_usage(record), | ||
| gap=record.gap, | ||
| tokens_aligned=record.tokens_aligned, | ||
| ) | ||
| for record in group | ||
| ] | ||
| if all(record.is_model_call for record in group) | ||
| else [], | ||
| requests_exhaustive=all(record.is_model_call for record in group), | ||
| ) | ||
| ) | ||
| accounting = UsageAccounting( | ||
| models_exhaustive=bool(records) | ||
| and all(record.model is not None and record.gap is None for record in records), | ||
| gaps=gaps, | ||
| observed=observed, | ||
| models=models, | ||
| ) | ||
| base_keys = { | ||
| "input_tokens", | ||
| "cached_input_tokens", | ||
| "input_tokens_include_cache", | ||
| "output_tokens", | ||
| "total_tokens", | ||
| "cost_usd", | ||
| } | ||
| complete_mapping = reported.to_mapping() | ||
| detail_keys = { | ||
| "input_tokens", | ||
| "cached_input_tokens", | ||
| "cache_write_input_tokens", | ||
| "input_tokens_include_cache", | ||
| "output_tokens", | ||
| "reasoning_tokens", | ||
| "output_tokens_include_reasoning", | ||
| "peak_request_input_tokens", | ||
| } | ||
| detail = { | ||
| key: value for key, value in complete_mapping.items() if key in detail_keys | ||
| } | ||
| return AgentUsage( | ||
| **{key: value for key, value in complete_mapping.items() if key in base_keys}, | ||
| extensions={ACCOUNTING_KEY: accounting.to_mapping(), DETAIL_KEY: detail}, | ||
| ) | ||
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
Oops, something went wrong.
Oops, something went wrong.
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.
Uh oh!
There was an error while loading. Please reload this page.