Skip to content

feat(a2a): emit aggregated turn token usage on the terminal status update - #2461

Open
QuentinBisson wants to merge 2 commits into
kagent-dev:mainfrom
QuentinBisson:feat/a2a-usage-total-v2
Open

feat(a2a): emit aggregated turn token usage on the terminal status update#2461
QuentinBisson wants to merge 2 commits into
kagent-dev:mainfrom
QuentinBisson:feat/a2a-usage-total-v2

Conversation

@QuentinBisson

@QuentinBisson QuentinBisson commented Aug 17, 2026

Copy link
Copy Markdown
Contributor

Supersedes #2332, which the stale bot closed and GitHub refuses to reopen. This is a reimplementation, not a rebase: both executors were rewritten to delegate to the ADK library's own A2A executor, so the manual event loop the original PR hooked into no longer exists.

What this adds

Every terminal status update (completed, input-required, failed) carries a kagent_usage_total metadata entry with the token usage summed across the task:

{
  "promptTokenCount": 300,
  "candidatesTokenCount": 50,
  "totalTokenCount": 350,
  "modelVersion": "gpt-4o-2024-11-20"
}

Partial (streaming chunk) events are skipped, since each LLM call reports its usage on the final non-partial event. Resumed tasks seed the accumulator from Task.metadata, so a task that spans several executions (HITL input-required cycles, follow-up messages) reports a task-lifetime total instead of the last segment only. The value is a running total re-emitted on each terminal update, so consumers take the latest value rather than summing across executions.

The value is serialized with the same genai UsageMetadata mapping as the existing per-event adk_usage_metadata / kagent_usage_metadata entries, plus modelVersion, so consumers can share one parser (with the caveat that modelVersion is not a field of the genai type). Tests in both languages assert that equality directly.

totalTokenCount is derived from prompt + candidates + thoughts when no provider reported one. The Anthropic models report per-call input and output counts without a total, and a zero total next to non-zero counts reads as "no tokens were used".

a2a-go and a2a-python both merge terminal status-update metadata into the stored task, which is how the total is persisted and how the next execution reads it back.

How it differs from #2332

  • Go: aggregation hangs off a RunnerProvider that wraps the ADK runner instead of AfterEventCallback. The callback only fires for ADK events that convert to an A2A artifact, and the input-required processor consumes the paused tool-call event, which is exactly the event carrying the usage of the call that paused. Wrapping the runner keeps the original vantage point: every ADK event, before conversion. The provider mirrors the ADK default provider (it has to, since RunnerProvider replaces RunnerConfig), and a test reflects over adka2a.RunnerConfig so a field added upstream fails the build instead of being silently dropped. Stamping uses the new AfterExecuteCallback hook, which covers all three terminal states in one place.
  • Python: aggregation runs from an ADK runner plugin (on_event_callback), for the same reason. The A2A after-event interceptor only runs for ADK events the converter turns into at least one A2A event, and the event whose long-running function call is stripped before conversion produces none. The plugin sees every ADK event the runner produces, matching the Go vantage point.
  • Test intent is ported rather than the test code; the Go tests now run end-to-end through the executor (including a full HITL pause/resume cycle) instead of unit-testing the accumulator in isolation.

Known limits

  • A cancelled task, or a client that disconnects mid-stream, does not reach the terminal status write, so the tokens of that in-flight execution are not added to the total. Earlier executions of the task keep theirs.
  • toolUsePromptTokenCount is not aggregated. It is inside totalTokenCount, so the total stays correct, but the breakdown fields do not sum to it on Vertex.
  • modelVersion is last-writer-wins over summed counts, so a task that uses more than one model reports a total that cannot be split back into per-model cost.

Verification

  • go build ./adk/... ./core/..., go test ./adk/..., go vet, gofmt
  • uv run pytest packages/kagent-adk/tests/unittests (390 passed) - the 9 failures in models/test_tls_e2e.py are pre-existing in this environment (verified against a clean checkout) and unrelated

@QuentinBisson
QuentinBisson force-pushed the feat/a2a-usage-total-v2 branch from 90d57f0 to c8981c7 Compare August 17, 2026 13:40
@github-actions github-actions Bot added the enhancement New feature or request label Aug 17, 2026
@QuentinBisson
QuentinBisson marked this pull request as ready for review August 17, 2026 14:00
@QuentinBisson
QuentinBisson requested review from a team and supreme-gg-gg as code owners August 17, 2026 14:00
Copilot AI lite review requested due to automatic review settings August 17, 2026 14:00

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Copilot encountered an error and was unable to review this pull request. You can try again by re-requesting a review.

…date

Both runtimes accumulate non-partial per-call token usage across an
execution and stamp it on the terminal status update (completed,
input-required, failed) under kagent_usage_total. Resumed tasks seed the
accumulator from the task metadata, so the total covers the whole task
rather than the last segment.

The value uses the same genai UsageMetadata serialization as the
per-event adk_usage_metadata / kagent_usage_metadata entries, plus
modelVersion.

Signed-off-by: QuentinBisson <quentin@giantswarm.io>
@QuentinBisson
QuentinBisson force-pushed the feat/a2a-usage-total-v2 branch from c8981c7 to a8d0c31 Compare August 17, 2026 19:01
@QuentinBisson

QuentinBisson commented Aug 17, 2026

Copy link
Copy Markdown
Contributor Author

Rebased onto main (bc92b68). Re-verified locally after the rebase: go build, go vet, and the touched Go/Python test packages pass.

The three remaining red checks are not from this branch, they reproduce identically on main at the same commit (run 32034286380):

  • go-unit-tests: the go/api/v1alpha2 and go/api/v1alpha3 CEL tests fail in setup, not in assertions. setup-envtest cannot write the 1.36.2-linux-amd64 assets ("unable to create file kubectl/etcd from archive to disk"), so the control plane never starts. Looks like parallel test packages racing on the same go/bin/k8s store.
  • upgrade-tests (prev-stable) / rolling-upgrade-tests (prev-stable): UPGRADE FAILED: context deadline exceeded in TestUpgrade/upgrade_with_helm.

The Python accumulator ran from the A2A after-event interceptor, which
only fires for ADK events the converter turns into at least one A2A
event. The event carrying the usage of the call that pauses a task for
input has its long-running function call stripped before conversion and
produces no A2A event, so its tokens were never counted. Accumulate from
an ADK runner plugin instead, matching the Go vantage point.

Failures raised outside the upstream executor publish their own terminal
event, which now carries the total like the other terminal states.

Derive totalTokenCount when no provider reported one, so Anthropic
agents stop reporting a zero total next to non-zero counts.
@github-actions github-actions Bot added enhancement New feature or request and removed enhancement New feature or request labels Aug 18, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

enhancement New feature or request

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants