Token tracking by alfredjimmy-redbus · Pull Request #174 · redbus-labs/adk-java

alfredjimmy-redbus · 2026-06-23T14:00:14Z

ADK-Java changes summary

Branch / context

Branch: token_tracking (merge of claude/live-bidi-token-tracking)
Purpose: Enable live BIDI (Gemini Live) token tracking via plugins, and fix live WebSocket message handling so usage events are not masked by spurious errors.

1. `Runner.runLive` — plugin callbacks for live sessions

File: core/src/main/java/com/google/adk/runner/Runner.java

Problem: Non-live runs already invoked onEventCallback / afterRunCallback; runLive did not. Plugins like RedBusTrackingPlugin and LiveTokenTrackingPlugin never saw live events, so BIDI token usage was never tracked.

Change: In runLiveImpl, for each live event:

pluginManager.onEventCallback(invocationContext, event) — per event (e.g. each usageMetadata turn)
pluginManager.afterRunCallback(invocationContext) — once when the live stream completes

Mirrors the existing non-live runImpl path. Session persistence (appendEvent) is unchanged; plugin callbacks are observational only.

2. `GeminiLlmConnection` — handle new Live API control messages

File: core/src/main/java/com/google/adk/models/GeminiLlmConnection.java

Problem: google-genai 1.58+ delivers control messages (sessionResumptionUpdate, goAway, voice-activity signals) that ADK did not recognize. They fell through to the “unknown server message” branch, producing error LlmResponses and fake ADK events — observed in RAE as event #123 with hasContent=false right after:

Received unknown or empty server message: {"sessionResumptionUpdate":{...}}

Change: Treat these as no-op control messages (return Optional.empty(), debug log only):

sessionResumptionUpdate
goAway
voiceActivityDetectionSignal
voiceActivity

Unchanged: Standalone usageMetadata messages still become LlmResponses with usage attached (existing behavior at lines 189–197).

3. Test coverage

File: core/src/test/java/com/google/adk/models/GeminiLlmConnectionTest.java

Added convertToServerResponse_withSessionResumptionUpdate_returnsEmpty — confirms resumption updates do not emit error events.

Existing tests still cover:

Usage-only messages → response with usageMetadata
Content + usage in same message
Unknown empty message → error response

Why this matters for RAE / BIDI voice

Dependency: Requires google-genai ≥ 1.58.0 (RAE already on 1.58.0).

Downstream (not in ADK): genai-sdk RedBusTrackingPlugin persists per-event usage to rae_tokens; LiveAudioSession logs USAGE_METADATA for confirmation.

Wire onEventCallback and afterRunCallback into Runner.runLiveImpl so the plugin lifecycle runs for live/bidi sessions, matching the non-live path. Previously live events were appended to the session but never passed through any plugin hook, so token usage emitted by Gemini Live (including audio modality breakdowns) could not be captured by plugins. Add LiveTokenTrackingPlugin which observes usageMetadata on each event and logs per-invocation totals on run completion. Live reports cumulative counts, so the plugin keeps the latest value per field rather than summing. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01FWJMsUqoQgsxLScGRk7uB2

Runnable example that drives a Gemini Live BIDI session with LiveTokenTrackingPlugin registered, demonstrating that usageMetadata (including audio modality breakdown) reaches onEventCallback and that afterRunCallback receives the aggregated per-session totals. Requires GOOGLE_API_KEY; model selectable via -Dmodel. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01FWJMsUqoQgsxLScGRk7uB2

Log every usageMetadata event with a sequence number and run two turns in one session. Empirically Gemini Live emits exactly one usageMetadata event per turn, and each event reports that turn's own totals (total = prompt + candidates) rather than a session-cumulative running total. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01FWJMsUqoQgsxLScGRk7uB2

A live two-turn session showed Gemini Live emits one usageMetadata event per turn, each reporting that turn's own usage (total = prompt + candidates) rather than a session-cumulative running total. Keeping the latest value therefore dropped all earlier turns and undercounted multi-turn sessions. Accumulate by summing scalar counts and per-modality token counts across events so the aggregate reflects true session totals. Verified end-to-end against a live session: per-turn totals 892 and 1044 sum to 1936. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01FWJMsUqoQgsxLScGRk7uB2

…vaidyanath-b/adk-java into token_tracking

….com/TrueGeometry/adk-ollama-mapdb.git

…s/adk-java into token_tracking

Sandeep-BA

@alfredjimmy-redbus : Pls validate and update the review comments.

Sandeep-BA · 2026-06-23T14:12:51Z

+      return Optional.empty();
+    } else if (message.goAway().isPresent()) {
+      logger.debug("Received go away: {}", message.goAway().get());
+      return Optional.empty();


Should we be really sending optional.empty on goAway or voiceactivity detectionSignal()?

Yes, keeping them as Optional.empty() is the right approach for this specific layer (convertToServerResponse).

This method's job is to map server events into client-facing responses (like content, tools, and usage). Because goAway and voiceActivity are infrastructure/connection-level signals, they don't have a representation in the client response model.

If we don't return Optional.empty(), we'd have to map them to an error event, which we previously found was incorrect and causing issues.

Sandeep-BA · 2026-06-23T14:15:29Z

+        () -> {
+          Usage usage = usageByInvocation.remove(invocationContext.invocationId());
+          if (usage == null) {
+            return;


Can we check what this return and return type is completable?

afterRunCallback always returns a Completable.
The inner return just means “skip logging” — e.g. live session ended with no usageMetadata events, so nothing was accumulated.

Sandeep-BA · 2026-06-23T14:16:55Z

+
+    @Override
+    public synchronized String toString() {
+      return "Usage{"


Should we really create a new class Usage?

Removed LiveTokenTrackingPlugin since this plugin can only log the usage, and moved plugin to sdk for persisting in DB

vaidyanath.b and others added 20 commits December 19, 2025 13:07

add : version file

871d849

Merge remote-tracking branch 'upstream/main'

a4fe90e

Merge remote-tracking branch 'upstream/main'

093480b

removed logs

b0d8f0d

Merge remote-tracking branch 'upstream/main'

4daa678

Merge remote-tracking branch 'upstream/main'

a2187f8

Merge branch 'claude/live-bidi-token-tracking' of https://github.com/…

1159479

…vaidyanath-b/adk-java into token_tracking

token tracking for bidi

3d49867

Merge branch 'main' into token_tracking

1fa262e

added MapDB Support and Ollama basic integration basis https://github…

cfe33b9

….com/TrueGeometry/adk-ollama-mapdb.git

Merge branch 'temp-br' into token_tracking

f8110ec

Merge branch 'main' of https://github.com/alfredjimmy-redbus/adk-java

35a42f0

Merge branch 'main' into token_tracking

797ab08

Merge branch 'token_tracking' of https://github.com/alfredjimmy-redbu…

3020499

…s/adk-java into token_tracking

Merge branch 'main' of https://github.com/alfredjimmy-redbus/adk-java

172f747

Merge branch 'main' into token_tracking

46a96ea

Sandeep-BA reviewed Jun 23, 2026

View reviewed changes

removed liveTokenTrackingPlugin since it was folded inside sdk

bcb703f

alfredjimmy-redbus force-pushed the token_tracking branch from c4c91b3 to bcb703f Compare June 26, 2026 10:02

alfredjimmy-redbus added 2 commits June 29, 2026 15:29

removed appending events

628dbc2

readded appending events

651e934

Sandeep-BA merged commit 6d86a3d into redbus-labs:main Jul 8, 2026
0 of 3 checks passed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

Token tracking - #174

Token tracking#174
Sandeep-BA merged 23 commits into
redbus-labs:mainfrom
alfredjimmy-redbus:token_tracking

alfredjimmy-redbus commented Jun 23, 2026 •

edited

Loading

Uh oh!

Sandeep-BA left a comment

Uh oh!

Sandeep-BA Jun 23, 2026

Uh oh!

alfredjimmy-redbus Jun 24, 2026

Uh oh!

Sandeep-BA Jun 23, 2026

Uh oh!

alfredjimmy-redbus Jun 24, 2026

Uh oh!

Sandeep-BA Jun 23, 2026

Uh oh!

alfredjimmy-redbus Jun 25, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Uh oh!

Conversation

alfredjimmy-redbus commented Jun 23, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

ADK-Java changes summary

Branch / context

1. Runner.runLive — plugin callbacks for live sessions

2. GeminiLlmConnection — handle new Live API control messages

3. Test coverage

Why this matters for RAE / BIDI voice

Uh oh!

Sandeep-BA left a comment

Choose a reason for hiding this comment

Uh oh!

Sandeep-BA Jun 23, 2026

Choose a reason for hiding this comment

Uh oh!

alfredjimmy-redbus Jun 24, 2026

Choose a reason for hiding this comment

Uh oh!

Sandeep-BA Jun 23, 2026

Choose a reason for hiding this comment

Uh oh!

alfredjimmy-redbus Jun 24, 2026

Choose a reason for hiding this comment

Uh oh!

Sandeep-BA Jun 23, 2026

Choose a reason for hiding this comment

Uh oh!

alfredjimmy-redbus Jun 25, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

alfredjimmy-redbus commented Jun 23, 2026 •

edited

Loading

1. `Runner.runLive` — plugin callbacks for live sessions

2. `GeminiLlmConnection` — handle new Live API control messages