feat(a2a): report task token usage through a versioned A2A extension

[QuentinBisson]

Fixes #3061

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

A versioned profile extension, https://kagent.dev/extensions/usage/v1, specified in docs/architecture/usage-extension.md. The status update that ends an execution (completed, input-required, failed, canceled) carries the task's token usage in its metadata under the extension URI:

{
  "inputTokens": 300, "outputTokens": 50, "reasoningTokens": 0, "cachedInputTokens": 0, "totalTokens": 350,
  "models": [{"model": "gpt-4o-2024-11-20", "inputTokens": 300, "outputTokens": 50, "reasoningTokens": 0, "cachedInputTokens": 0, "totalTokens": 350}]
}

• Field names are provider-neutral, and counts are attributed per model instead of naming the last model on a combined total.
• The value covers the task lifetime. Resumed tasks (HITL cycles, follow-up messages) seed from the stored task, so consumers replace earlier values with the latest one and never sum them.
• The kagent ADK runtimes advertise the extension on their Agent Card with required: false (the runtime card and the controller-compiled kagent card; Codex and Claude do not emit usage and do not advertise it). The Go runtime echoes A2A-Extensions on activation. The Python A2A SDK has no server-side activation API, same as HITL today.
• No kagent client reads the payload yet, so none of them activate the extension. A client that activates HITL and usage together sends one comma-joined header value, and a2a-go matches raw header values exactly, so the Go HITL and usage activation checks split the header on commas.
• v1 is emitted whether or not the client activated it. Gating emission, and moving the per-event kagent.dev/a2a/usage key into the extension, are left for a later version of the URI.

A call that reports no total gets one derived from input + output + reasoning when it is added, so tasks mixing providers (Anthropic reports no total) and resumed tasks whose persisted total was derived stay correct. Partial streaming events are skipped.

Aggregation runs from an ADK runner plugin in both languages (plugin.Config.OnEventCallback in Go, on_event_callback in Python). The A2A after-event callback misses the event whose long-running function call is stripped before conversion, and that event carries the usage of the call that paused the task. Go stamps the payload from AfterExecuteCallback.

A cancel request is answered by a separate executor call, and A2A discards the interrupted execution's own final event. Both executors therefore index running executions' usage by task ID, so the canceled status carries the calls completed before the cancel. A task parked in input-required repeats its persisted total. A call cut off mid-flight was never reported by the provider and is not counted.

Known limits

• A client that disconnects mid-stream does not reach the terminal write, so the tokens of that in-flight execution are not added. Earlier executions keep theirs.
• toolUsePromptTokenCount is not broken out. It is inside totalTokens, so on Vertex the breakdown fields do not sum to the total.
• Calls whose response names no model are counted in the totals only.

[Copilot]

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

[QuentinBisson]

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.

[QuentinBisson]

@supreme-gg-gg this is clean against main and green. Could you take a look when you have a moment?

[EItanya]

🤖 AI-generated review.

Reviewed 013356fc. I found one correctness issue and a simpler way to integrate this with ADK:

1. Derive missing totals per call, before accumulation. In both Go and Python, the fallback only runs when the accumulated total is zero. An Anthropic-style execution derives and persists 120 tokens; after resuming and spending another 230 without a reported total, prompt/output counts reach 350 but totalTokenCount stays at 120. Mixing calls with and without reported totals also undercounts. Reproduced both cases in both languages.

2. Use Go's existing plugin.Config.OnEventCallback. It observes ADK events before A2A conversion, including the paused tool-call event. Registering it through RunnerConfig.PluginConfig lets us retain the default provider and remove the custom runner/provider and reflection test. I tested that replacement: the Go A2A package passes, including the HITL pause/resume test. The Python implementation already uses this approach.

3. Upstream direction: both ADKs expose per-call usage/hooks; current Python main also has private telemetry aggregation, though it is absent from this PR's locked version and isn't a public A2A task-total API. A public invocation-usage summary consumed by ADK's A2A adapters would be a useful upstream contribution, leaving a thin kagent metadata mapping. No need to wait for that to use the existing plugin hooks here.

Reproduction sources, exact plugin experiment, commands, results and coverage limits. Baseline Go package tests and 17 focused Python tests passed; four added regression probes fail as described. No live-model or database persistence tests were run.

[EItanya]

Could we expose this through an optional, versioned A2A usage extension?

A2A doesn't currently define a built-in token-usage schema. Upstream discussed this exact question in a2aproject/A2A#1155 (comment) and recommended an opt-in extension. There's also a newer cost/budget proposal (a2aproject/A2A#2121), but it remains a proposal rather than an adopted standard.

The existing metadata["kagent_usage_total"] is a valid use of A2A metadata. An extension would still use metadata, while making the schema and behavior discoverable and versioned. I'd suggest:

• Define a versioned kagent usage extension, advertise it in AgentCard.capabilities.extensions with required: false, and support A2A's client opt-in mechanism.
• Document the cumulative task-lifetime semantics: consumers replace earlier totals with the latest value, plus the behavior on resume, failure, and cancellation.
• Prefer provider-neutral fields such as inputTokens and outputTokens over making Google's UsageMetadata the public contract. If model attribution is included, associate counts with each model rather than attaching the last model's name to totals across multiple models.

The aggregation work here would still be useful; this would formalize its public output. A2A extension guidance: https://a2a-protocol.org/latest/topics/extensions/

[QuentinBisson]

Agreed, let's make it an extension, following the same pattern as the HITL one: a versioned kagent.dev URI, advertised on the Agent Card as optional, and a doc page describing the schema and the cumulative task-lifetime semantics. Provider-neutral field names, and counts attributed per model rather than the last model's name on a combined total.

There are two ways to do it, and I'd like to confirm which you want.

The first is a data-only extension. The schema becomes discoverable and versioned, the value keeps being emitted as it is today, and no existing consumer breaks. This covers what you asked for.

The second also gates emission on client activation. That is what the upstream thread had in mind (a2aproject/A2A#1155 (comment)): an agent that does not want to share its token counts can decline. The catch is that it only means something if the per-event kagent_usage_metadata keys move inside the extension too, since they carry strictly more information than the aggregate. That gives one way to read token usage instead of three, which is the better end state, but it moves every current reader at once: the UI, remote A2A sub-agent calls, and anything reading the 0.x adk_usage_metadata name.

I'd go with the first for now. Nobody has asked to withhold token counts, and if that need appears, or if an upstream usage extension gets adopted, gating can come as a v2 of the URI. To keep that door open cheaply I'd have our own clients send the activation header from the start, so a later gate is a no-op inside kagent and only affects third parties. The v1 spec would state plainly that the extension is data-only and always emitted.

One detail either way: the aggregate has to keep being computed regardless of activation, because a resumed task reads the running total back from the stored task.

Happy to do the second instead if you'd rather take the break now while the API version is already breaking.

[supreme-gg-gg]

I think the earlier comments said we want this as an extension instead of just metadata? Otherwise this is looking good

[QuentinBisson]

Yes you're right and this makes total sense :)

I did not do the change yet as I was waiting for feedback on whether you prefer a data extension Always on or it if needs to be gated which imply the existing metadata like kagent_usage_metadata move to the extension or are fully removed

[QuentinBisson]

Going with the always-on version for v1 unless you object: a versioned kagent.dev profile extension, advertised with required: false. The agent echoes A2A-Extensions when a client activates it, and our own clients send the header, but emission isn't gated, so current readers keep working. Gating, and moving kagent_usage_metadata into the extension, can come as v2 of the URI.

[QuentinBisson]

Small correction: only the Go runtime echoes A2A-Extensions. The Python A2A SDK has no server-side activation API, which is the same gap HITL has today. Emission isn't gated, so it doesn't change behavior. I can send the Python echo for both extensions as a follow-up.

[QuentinBisson]

@EItanya I fixed go-lint and I'm fixing the DCO

[QuentinBisson]

Thanks for checking this PR :D

[EItanya]

I'm trying haha, such bad luck!

[moezdil]

what reads this payload, ui only sends the header

[QuentinBisson]

Fair point, nothing on the client side reads it. I removed the activation from the UI, and also from the MCP invoke tool and the Go and Python remote tools, which had the same issue. Emission doesn't depend on activation in v1, and a gated version would get a new URI, so sending it now gains nothing. Fixed in 4100b9c.