feat(aws-strands): report provider token usage on the terminal run event - #2617
Merged
Conversation
Contributor
Python Preview PackagesVersion
Install with uvAdd the TestPyPI index to your [[tool.uv.index]]
name = "testpypi"
url = "https://test.pypi.org/simple/"
explicit = trueThen install the packages you need: # Core SDK
uv add 'ag-ui-protocol==0.0.0.dev1788439541' --index testpypi
# Integrations (each already depends on the matching ag-ui-protocol preview)
uv add 'ag-ui-langgraph==0.0.0.dev1788439541' --index testpypi
uv add 'ag-ui-crewai==0.0.0.dev1788439541' --index testpypi
# NOTE: ag-ui-agent-spec depends on pyagentspec (git-only, not on PyPI).
# You will need to install pyagentspec separately from its git repo.
uv add 'ag-ui-agent-spec==0.0.0.dev1788439541' --index testpypi
uv add 'ag_ui_adk==0.0.0.dev1788439541' --index testpypi
uv add 'ag_ui_strands==0.0.0.dev1788439541' --index testpypiInstall with pippip install \
--index-url https://test.pypi.org/simple/ \
--extra-index-url https://pypi.org/simple/ \
ag-ui-protocol==0.0.0.dev1788439541
Commit: 8f424c6 |
@ag-ui/a2a-middleware
@ag-ui/a2ui-middleware
@ag-ui/event-throttle-middleware
@ag-ui/mcp-apps-middleware
@ag-ui/mcp-middleware
@ag-ui/a2a
@ag-ui/adk
@ag-ui/ag2
@ag-ui/agno
@ag-ui/aws-strands
@ag-ui/claude-agent-sdk
@ag-ui/claude-managed-agents
@ag-ui/crewai
@ag-ui/langchain
@ag-ui/langgraph
@ag-ui/llamaindex
@ag-ui/mastra
@ag-ui/pydantic-ai
@ag-ui/vercel-ai-sdk
@ag-ui/watsonx
@ag-ui/a2ui-toolkit
create-ag-ui-app
@ag-ui/client
@ag-ui/core
@ag-ui/encoder
@ag-ui/proto
commit: |
AlemTuzlak
previously approved these changes
Sep 3, 2026
RUN_FINISHED.usage and RUN_ERROR.usage are documented on the protocol but this bridge never populated either, so a Strands run reported nothing at all about what it spent. The counts are read per model invocation off the stream's metadata event, accumulated for the length of the run, and folded into one entry per (provider, model) at whichever terminal event ends it. The metadata event is the source rather than AgentResult.metrics.accumulated_usage, which cannot answer the question that matters: Strands seeds that field with inputTokens=0, outputTokens=0, totalTokens=0, so a provider that reported nothing is indistinguishable from one that reported zero. The metadata event carries the provider's own chunk verbatim, ahead of the SDK's own accumulation, so an unreported run leaves the field omitted and a measured zero is reported as zero. Those are different answers, and a consumer showing 0 tokens for an unmeasured run is wrong. Every count passes a guard before it is accepted: a real, finite, non-negative whole number no larger than 2**53 - 1, the ceiling the TypeScript protobuf decoder imposes on every binding. A count that fails is DROPPED and the rest of the entry survives, because TokenUsage validates its bounds inside the producer's own constructor: an unguarded value would raise while BUILDING the terminal event and cost the caller a whole successful run over a token count. Integers settle before any float check, since math.isfinite coerces to float and raises OverflowError on a large int, which would abort the run from inside the guard meant to protect it. cacheWriteInputTokens is dropped rather than folded into a neighbouring count, since AG-UI has no slot for it and folding would overstate the count that received it. Strands reports no reasoning-token count, so reasoning_tokens stays unset. Provider labels come from an explicit class-name table sharing its canonical labels with the TypeScript bridge, not from a derivation: the two SDKs do not name these classes identically, so Python's GeminiModel and the TypeScript SDK's GoogleModel both report "google" rather than splitting one vendor in two across the two bridges. A class the table does not name omits the label rather than guessing, and a test reads the installed SDK's own model modules so a provider Strands adds later cannot go unlabelled. Aggregation uses the published ag_ui.core helper, which is why the ag-ui-protocol pin moves to >=0.1.22; the vendor mapper stays local to the integration so this change needs no SDK release. Both agent paths report. The orchestrator's inner events do surface the metadata event, one wrapper deeper, and node identity is available where they do, so each entry is labelled with the model of the node that spent the tokens and a multi-model Graph keeps its models apart. Accumulators are local to each run's generator, so a second sequential run in one stream cannot inherit the first run's counts. Usage is attached only where a model call could already have reported it. The early-exit validation and idempotent-replay terminals fire before any model runs and carry nothing, rather than claiming a measured nothing.
RUN_FINISHED.usage and RUN_ERROR.usage are part of the protocol but this bridge populated neither, so a Strands run said nothing at all about what it spent. Counts are now read per model invocation off modelMetadataEvent, accumulated for the length of the run, and folded into one entry per (provider, model) at whichever terminal event ends it. The Python bridge implements the same contract from the same spec, so the two emit the same array for the same run. The metadata event is the source rather than AgentResult.metrics .accumulatedUsage, which cannot answer the question that matters: Strands seeds that field with zeros, so a provider that reported nothing looks identical to one that reported zero. The metadata event carries the provider's own chunk, so an unreported run leaves the field omitted and a measured zero is reported as zero. Those are different answers, and a consumer showing 0 tokens for an unmeasured run is showing a number nobody gave it. The event still reaches RAW forwarding afterwards: its latency metrics have no AG-UI equivalent, and reading usage off it must not cost the client the rest of the report. Counts are guarded to real, finite, non-negative whole numbers at or below Number.MAX_SAFE_INTEGER, and a value outside that is DROPPED with the rest of the entry kept, never clamped and never zeroed. The bound is not belt-and-braces: TokenUsageSchema constrains counts to non-negative integers and stops there, so an oversized count validates and then throws inside the protobuf transport's int64 decoder, failing an otherwise successful run at its final event on the binary wire while the SSE wire carries the same run fine. @ag-ui/core's shared num() checks finiteness only, which is why the guard is local rather than a reuse: the full bound has to be applied here for the Python guard to behave identically. Provider labels come from an explicit table keyed on the model class name, not from stripping "Model" off it, because the two SDKs do not name the same provider's class identically and a derived label would drift between the bridges without anyone noticing. This SDK's GoogleModel and the Python SDK's GeminiModel both label "google": one provider, one label. A class the table does not name omits the provider label rather than guessing, which also keeps an integrator's own Model subclass from having its spend attributed to a provider nobody named. cacheWriteInputTokens has no AG-UI slot and is dropped rather than folded into another count; Strands reports no reasoning-token count, so reasoningTokens is never set. Usage rides only the terminals a model call can precede. The normal finish, the interrupt-variant finish (an interrupted run is a finished run, and the calls that raised the interrupt were real), the post-stream reconciliation refusals and both paths' catch-all RUN_ERROR all carry it. The preflight resume gates, the idempotent-replay finish and the media-refusal error fire before the agent stream exists and carry nothing rather than claiming a measured nothing. The orchestrator path reports too, and it was verified rather than assumed: a real Graph surfaces the metadata event nested one wrapper deeper inside nodeStreamUpdateEvent, and node identity is available there. Each node's beforeModelCallEvent arrives before that node's metadata event and carries the Model, which is the only place the pairing exists, so an entry is labelled with the model of the node that spent the tokens and a multi-model Graph keeps its models apart. Only the labels are kept, never the model. Accumulators are locals of each run's generator, so they are seeded per run by construction and a second sequential run on one thread cannot inherit the first's counts. Aggregation goes through the published aggregateTokenUsage rather than a local sum, so every AG-UI producer groups identically; the vendor mapper stays local to the integration, since the Python bridge consumes the published protocol package and a new core mapper would not exist for it until the next SDK release.
The two token-usage commits made ARCHITECTURE.md incomplete against main: the terminal events now carry usage and the document said nothing about it. Adds the entry beside Citations, covering the metadata-event source and why AgentResult.metrics.accumulated_usage is not it, which terminals carry usage and which deliberately do not, the drop-not-clamp count guard and the protobuf int64 reason its upper bound exists, per-node labelling on both orchestrator paths, the provider-label tables the two SDKs force apart, the agent-as-tool spend that goes uncounted, and why no capabilities flag is advertised for it. Amends the RAW row of the recap table: its usage-metadata example implied RAW was the only way to read counts, so it now names the latency metrics that really are RAW-only and points at the mapped channel.
…e fixtures The provider table missed `OpenAIResponsesModel`, which Strands ships as a second OpenAI model class, so a run served by OpenAI's Responses API reported its counts with no provider label. The exhaustiveness test caught it on the newest release; the label is `openai`, the same vendor its sibling class reports. The TypeScript table is unchanged on purpose: that SDK reaches the Responses API through a config on its single `OpenAIModel`, so it has no such class to key on. The token-usage fixtures then get the audit that failure implies. Strands declares `inputTokens`, `outputTokens` and `totalTokens` as `Required` on its own `Usage`, and its metrics accumulation and telemetry subscript them bare, so a payload missing one aborts the run inside the SDK on a version that does not default them. Ten fixtures were minimised to the single count their test talked about. Two failed at the declared floor; the other eight passed only because their assertions held on the RUN_ERROR that the incomplete payload caused, which made them tests of which Strands is installed rather than of this bridge. Every payload on a real model path now carries all three counts, and every one of those tests asserts the terminal event is RUN_FINISHED so the shortcut cannot come back silently. The three shapes that cannot carry all three and still make their point (no `usage` key at all, a counts-free usage object, string counts) move to the scripted core, which is the pattern the file already uses for a non-mapping payload, and each says why in its docstring.
ranst91
force-pushed
the
claude/determined-chatelet-72bd04
branch
from
September 3, 2026 12:45
aed86fa to
1bcfe0b
Compare
AlemTuzlak
approved these changes
Sep 3, 2026
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
2 participants
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.
RUN_FINISHED.usageandRUN_ERROR.usagehave been on the protocol since #2188 andare populated by the TypeScript LangGraph, LangChain and Mastra producers, and since
2026-08-31 by LangGraph Python. Neither AWS Strands bridge wrote them, so a Strands
run reported no token usage at all while the field's presence in the schema read as
support. This adds the producer to both bridges.
Groundwork credit: draft PR #2190 opened this thread, identified that Strands exposes
usage on its own metrics surface, and correctly diagnosed that the Python half was
blocked on an unreleased
ag-ui-protocol. Its LangGraph Python half landedindependently in 47302ba / 1d1f23c / eecc2ed with a different implementation,
which is the reference this change follows. Two of that draft's premises turned out
not to hold, and both changed the design; they are set out under "What the draft got
right and what changed" below.
What lands
One
TokenUsageentry per model invocation, accumulated across the run and foldedinto one entry per
(provider, model)at the terminal event by the shared SDKaggregator (
aggregate_token_usage/aggregateTokenUsage). Numeric counts and theprovider and model labels only: no prompts, completions, message content, thread,
run or user ids, no latency and no traces. This shape feeds anonymous telemetry.
Field mapping, verified against
strands-agents1.18.0 and@strands-agents/sdk1.1.0:UsageinputTokensinput_tokensinputTokensoutputTokensoutput_tokensoutputTokenstotalTokenstotal_tokenstotalTokenscacheReadInputTokenscached_input_tokenscachedInputTokenscacheWriteInputTokenscacheWriteInputTokenshas no AG-UI slot and is dropped rather than folded into aneighbouring count, which would overstate that count. Strands reports no
reasoning-token count, so
reasoning_tokensis never set from this channel.Terminal events that carry usage: the normal
RUN_FINISHED, the interrupt-variantRUN_FINISHED(an interrupted run is a finished run, and the calls that raised theinterrupt were real spend), the post-stream reconciliation refusals, the forced-stop
RUN_ERROR, and the catch-allRUN_ERROR. Pre-stream error paths deliberately carrynone: they fire before the Strands stream is ever opened, so there is nothing to
report, and an empty array there would claim a measured nothing. The field is omitted,
not
[]and not zeros, when the provider reported nothing, so a consumer reads amissing field as "not measured".
Both agent paths are covered. On the multi-agent orchestrator path each entry is
labelled with the model of the node that actually spent the tokens, so a two-model
Graphkeeps its models apart rather than collapsing them.What the draft got right and what changed
The release blocker was half gone, not gone.
TokenUsagereached PyPI inag-ui-protocol0.1.20, but the shared mapper and aggregator module(
ag_ui/core/token_usage.py), and with it the oversized-count guard, first shipped in0.1.22.
integrations/aws-strands/python/pyproject.tomlpinned>=0.1.21, whichgets the type but not the helpers. Verified by unpacking both wheels. The pin moves to
>=0.1.22anduv.lockis refreshed. The TypeScript side needed no manifest change:the published
@ag-ui/core@0.0.59already exportsaggregateTokenUsageand theexisting
>=0.0.59peer range admits it, likewise verified against the publishedtarball.
AgentResult.metrics.accumulated_usageis the wrong source. The draft proposedreading usage off the terminal result. That value is pre-summed AND seeded with
inputTokens=0, outputTokens=0, totalTokens=0, so it cannot distinguish "the providerreported nothing" from "the provider reported zero", and keeping those apart is what
lets the terminal event omit
usageinstead of publishing a measured zero. Bothbridges already receive, and already forward as
RAW, Strands' per-model-call metadataevent, which carries the same counts before accumulation. That per-call channel is what
this reads, matching how the LangGraph Python producer works. The metadata event still
reaches
RAW, because its latency metrics have no AG-UI equivalent.No terminal-result fallback channel was added, and adding one would be wrong: Strands'
process_streamyields a chunk event for every provider chunk unconditionally, beforeits own accumulation, so a provider that reports usage always reaches the metadata
channel.
Where the two bridges differ, and where they do not
The mappers are behaviourally identical. Twelve payloads, including every guard edge
case, were driven through both and produce the same result field for field.
The one real divergence is the provider label table, and it is a difference in the two
Strands SDKs rather than in this change. Every provider both SDKs ship maps to the same
canonical label, Gemini and Google included: Python's
GeminiModeland the TypeScriptSDK's
GoogleModelboth label asgoogle, resolved by hand precisely because derivinga label from the class name is what would have silently split one vendor in two. Python
additionally covers
litellm,llamaapi,llamacpp,mistral,ollama,sagemakerand
writer; TypeScript additionally coversvercel. Those classes do not exist in theother SDK. An unrecognised model class omits the provider label rather than guessing,
which also covers an integrator's own
Modelsubclass.Second difference, in the guard rather than in behaviour: Python must exclude
boolexplicitly, since there it subclasses
int, and must settle the integer case before anyfloat check, because
math.isfinitecoerces to float and raisesOverflowErroron alarge int, which would abort the run from inside the guard that exists to protect it.
TypeScript's
typeofexcludes booleans for free. Both reject the same inputs.The vendor mapper stays local to the integration rather than joining the LangChain and
AI-SDK mappers in the SDK cores. The Python bridge consumes the published
ag-ui-protocol, so a core mapper would not exist for it until the next SDK release,and the two bridges have to ship together. Only the aggregator is shared, because it is
already published.
Dropped, not clamped
An out-of-range count is dropped and the rest of the entry survives. It is not clamped
to the ceiling, which would report a number no provider gave, and not zeroed, which
would claim a measurement never made. This follows the established behaviour of the
LangGraph Python guard rather than the word "clamped" in the original request.
The ceiling is
2**53 - 1, and the reason is worth stating because it iscounterintuitive and was checked rather than assumed.
TokenUsageSchemaconstrainscounts to non-negative integers but sets no upper bound, so an oversized count
passes schema validation:
The real failure is one layer out, on the protobuf transport, exactly as the shared
guard's own rationale says: driving
@ag-ui/protodirectly,2**53throwsValue is larger than Number.MAX_SAFE_INTEGERon decode and1e30throwsinvalid int64on encode, while2**53 - 1round-trips. So an unguarded oversizedcount would break the run on the binary wire and carry fine on SSE. Bounding at the
source is what keeps the two transports reporting the same thing. The TypeScript test
suite asserts that zod currently accepts the oversized value, with a failure message
telling a future reader to revisit the guard if that ever changes, so the guard cannot
be quietly deleted as redundant.
Tests
Both suites, all required cases in both languages: usage on a normal completion, usage
accumulated across several model calls in one run, partial usage on a run that failed
after a model call reported, an oversized count dropped with the rest of the entry
surviving, no
usagefield when the provider reports none, the orchestrator pathdriven against a real
Graphwith two nodes on distinct models, a second sequentialrun not inheriting the first run's counts, and an assertion that nothing
content-bearing rides along.
Baselines before this change were 1270 and 1611, both green, so nothing was already
failing and nothing regressed.
One environment note for anyone reproducing the TypeScript run: the workspace
@ag-ui/coreneeds building first, otherwise most test files fail to resolve it whilestill reporting a passing test count, which reads like a green run.
Known gaps, deliberately not closed here
A tool that wraps another Agent does not have its inner spend counted. Verified
empirically: a generator tool re-yields the inner agent's stream as tool-stream
payloads, which the parent loop routes through a different path that never reaches the
metadata branch. With an inner model reporting 100/50/150 and an outer model reporting
1/1/2 then 2/2/4, the run reports only the outer 3/3/6. That is real spend going
unreported. It is out of the scope agreed for this change, and widening it on one bridge
would make the same run report different totals depending on which bridge served it, so
the current boundary is pinned by a test whose docstring says changing it requires both
bridges at once. Worth a coordinated follow-up.
No capabilities-endpoint flag. A
tokenUsageflag would sit naturally in the servedcapability matrix, but adding one to a served document on one bridge without the other
is exactly the drift the recent error-code and resume-contract work closed. Follow-up
for both bridges together.