Skip to content

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

Open
QuentinBisson wants to merge 1 commit 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 1 commit into
kagent-dev:mainfrom
QuentinBisson:feat/a2a-usage-total-v2

Conversation

@QuentinBisson

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 its 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 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. A Go test asserts that equality directly.

a2a-go merges terminal status-update metadata into the stored task, which is both 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). Stamping uses the new AfterExecuteCallback hook, which covers all three terminal states in one place.
  • Python: the executor already keeps an _ExecutionState with an after-event interceptor and stamps metadata onto the final status update in _after_agent, so the accumulator plugs straight into those. One ADK event can be converted into several A2A events there, and the interceptor runs per A2A event, so events already counted are tracked by id.
  • 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.

Verification

  • go build ./adk/... ./core/..., go test ./adk/... ./core/... (1843 passed), go vet, golangci-lint
  • uv run pytest packages/kagent-adk/tests/unittests — 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.

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