Gemini Model not returning Usage Metadata in Streaming responses

[chaskyhoffmann]

** Please make sure you read the contribution guide and file the issues in the rigth place. **
Contribution guide.

Describe the bug
When using streaming mode, the Gemini Model does not return the usageMetadata in the LlmResponse. When stream=true the usageMetadata should return in the final aggregated response.

To Reproduce
Steps to reproduce the behavior:

1. Create an LlmAgent with the Gemini model.
2. In the RunConfig set the streaming model to RunConfig.StreamingMode.SSE
3. observe the LlmResponse in afterModelCallback of a plugin or agent.
4. See that both for partial and final LLM responses, usageMetadata is Optional.Empty

Expected behavior
The final and non-partial LlmResponse should include the usageMetadata.

Desktop (please complete the following information):

• OS: MacOS
• Java version:
• ADK version(see maven dependency): 0.5.0

Additional context
The python version of ADK returns the usageMetadata correctly.
The usageMetadata is returned from the Gemini API in the last chunk.

[hemasekhar-p]

Hi @chaskyhoffmann,

Thanks for reporting this issue. I’ve tried to reproduce it locally using the latest version of the Java ADK and a live-enabled model, but I’m unable to see the behavior you described. I noticed you are using afterModelCallback(). Could you please share a minimal reproducible example showing how you’re using it, along with your agent, model configuration. And also, could you help me understand the frequency of this issue?

Having this information will allow us to reproduce the issue accurately and provide guidance.

[chaskyhoffmann]

Hi @hemasekhar-p

Here is a minimal program that reproduces this issue. It is complete with our Agent & Model configurations.

import com.google.adk.agents.CallbackContext;
import com.google.adk.agents.LlmAgent;
import com.google.adk.agents.RunConfig;
import com.google.adk.events.Event;
import com.google.adk.models.Gemini;
import com.google.adk.models.LlmResponse;
import com.google.adk.plugins.BasePlugin;
import com.google.adk.runner.InMemoryRunner;
import com.google.adk.runner.Runner;
import com.google.genai.Client;
import com.google.genai.types.Content;
import com.google.genai.types.GenerateContentResponseUsageMetadata;
import com.google.genai.types.Part;
import io.reactivex.rxjava3.core.Flowable;
import io.reactivex.rxjava3.core.Maybe;

import java.util.List;
import java.util.Map;
import java.util.concurrent.atomic.AtomicInteger;

/**
 * Minimal reproduction case for ADK Java issue #777:
 * usageMetadata is not populated when streaming is enabled
 *
 * Expected: usageMetadata should be populated in the final aggregated response
 * Actual: usageMetadata remains Optional.empty() in both partial and final responses
 *
 * Issue: https://github.com/google/adk-java/issues/777
 */
public class UsageMetadataReproduction {

    public static void main(String[] args) {
        // IMPORTANT: Set these environment variables or replace with your values
        String project = System.getenv("GOOGLE_CLOUD_PROJECT"); // e.g., "your-project-id"
        String location = System.getenv("GOOGLE_CLOUD_LOCATION"); // e.g., "us-central1"

        if (project == null || location == null) {
            System.err.println("ERROR: Please set GOOGLE_CLOUD_PROJECT and GOOGLE_CLOUD_LOCATION environment variables");
            System.err.println("Example:");
            System.err.println("  export GOOGLE_CLOUD_PROJECT=your-project-id");
            System.err.println("  export GOOGLE_CLOUD_LOCATION=us-central1");
            System.exit(1);
        }

        System.out.println("=== ADK Java Issue #777 Reproduction ===");
        System.out.println("Testing usageMetadata with streaming enabled");
        System.out.println("Project: " + project);
        System.out.println("Location: " + location);
        System.out.println();

        // Create Gemini client with Vertex AI
        Client client = Client.builder()
                .vertexAI(true)
                .project(project)
                .location(location)
                .build();

        Gemini geminiModel = new Gemini("gemini-2.5-flash", client);

        // Counters for tracking callbacks
        AtomicInteger afterModelCallbackCount = new AtomicInteger(0);
        AtomicInteger finalResponseCount = new AtomicInteger(0);
        AtomicInteger partialResponseCount = new AtomicInteger(0);

        // Create a simple plugin to capture usageMetadata
        BasePlugin usagePlugin = new BasePlugin("UsageMetadataPlugin") {
            @Override
            public Maybe<LlmResponse> afterModelCallback(CallbackContext ctx, LlmResponse response) {
                int callbackNum = afterModelCallbackCount.incrementAndGet();
                System.out.println("\n--- afterModelCallback #" + callbackNum + " ---");

                // This is the key issue: usageMetadata is always empty during streaming
                if (response.usageMetadata().isPresent()) {
                    GenerateContentResponseUsageMetadata metadata = response.usageMetadata().get();
                    System.out.println("✓ usageMetadata found:");
                    System.out.println("  - Prompt tokens: " + metadata.promptTokenCount().orElse(0));
                    System.out.println("  - Candidates tokens: " + metadata.candidatesTokenCount().orElse(0));
                    System.out.println("  - Total tokens: " + metadata.totalTokenCount().orElse(0));
                    finalResponseCount.incrementAndGet();
                } else {
                    System.out.println("✗ usageMetadata is EMPTY (this is the bug!)");
                    partialResponseCount.incrementAndGet();
                }

                // Show the raw response for debugging
                String preview = "(no text)";
                if (response.content().isPresent()) {
                    Content content = response.content().get();
                    String text = content.text();
                    if (text != null && !text.isEmpty()) {
                        preview = text.substring(0, Math.min(50, text.length())) + "...";
                    }
                }
                System.out.println("Raw content preview: " + preview);

                return Maybe.empty();
            }
        };

        // Create agent
        LlmAgent agent = LlmAgent.builder()
                .name("TestAgent")
                .model(geminiModel)
                .instruction("You are a helpful assistant. Keep your responses concise.")
                .build();

        System.out.println("=== Sending message with STREAMING enabled ===");
        System.out.println("Input: \"What is 2+2?\"");
        System.out.println();

        // Create runner with in-memory session and plugin
        Runner runner = InMemoryRunner.builder()
                .agent(agent)
                .appName("TestApp")
                .plugins(List.of(usagePlugin))
                .build();

        // Create user content
        Content userContent = Content.builder()
                .role("user")
                .parts(List.of(Part.fromText("What is 2+2?")))
                .build();

        // Run with streaming enabled
        RunConfig config = RunConfig.builder()
                .setStreamingMode(RunConfig.StreamingMode.SSE)  // This is the key: streaming is enabled
                .build();

        var session = runner.sessionService()
                .createSession(runner.appName(), "user-id")
                .blockingGet();

        // Subscribe to the response stream
        Flowable<Event> responseStream = runner.runAsync(
                "user-id",
                session.id(),
                userContent,
                config,
                Map.of()
        );

        StringBuilder fullResponse = new StringBuilder();

        responseStream.blockingSubscribe(
                event -> {
                    if (event.author().equals("TestAgent") && event.content().isPresent() && !event.partial().orElse(false)) {
                        String text = event.content().get().text();
                        if (text != null && !text.isEmpty()) {
                            fullResponse.append(text);
                        }
                    }
                },
                error -> System.err.println("Error: " + error.getMessage()),
                () -> {
                    System.out.println("\n=== Streaming Complete ===");
                    System.out.println("Full Response: " + fullResponse);
                    System.out.println();
                    System.out.println("=== RESULTS ===");
                    System.out.println("Total afterModelCallback invocations: " + afterModelCallbackCount.get());
                    System.out.println("Callbacks with usageMetadata: " + finalResponseCount.get());
                    System.out.println("Callbacks WITHOUT usageMetadata: " + partialResponseCount.get());
                    System.out.println();
                    System.out.println("EXPECTED: usageMetadata should be populated in at least one callback");
                    System.out.println("ACTUAL: usageMetadata was empty in " + partialResponseCount.get() + " callbacks");
                    System.out.println();
                    System.out.println("This reproduces issue: https://github.com/google/adk-java/issues/777");
                }
        );
    }
}

[hemasekhar-p]

@chaskyhoffmann, thank you for sharing the required information. We will reproduce and investigate on this.

[AdamBroadcom]

This bug is also blocking token tracking in my codebase for streaming sessions. I did pull the code in #835 and it resolved my issue. If there is anything I can do to assist with getting this merged so I can use this in production, let me know.

Thanks

[hemasekhar-p]

Hi @chaskyhoffmann, I am able to reproduce issue with this code and during testing, i observed that while the Gemini API includes the usageMetadata object in the final chunk of the SSE stream, the ADK's afterModelCallback receives an LlmResponse where getUsageMetadata() is consistently Optional.empty(). This creates a significant gap in observability for streaming agents.

The changes in this PR correctly target the terminal event of the stream to capture the usage statistics. By ensuring the aggregator maps this metadata into the final response object, it restores the ability to track token consumption and billing metrics

[hemasekhar-p]

@sherryfox, Please have a look into it this and do needful

[OwenDavisBC]

@hemasekhar-p / @sherryfox Is there anything more you may need from my end to help close this out?

[sherryfox]

Thanks for reaching out and apologies for the delay in updating this issue.

We're happy to let you know that the changes related to this have now been merged (PR#835)!
@OwenDavisBC, Thank you very much for your contribution regarding the fix and patience.

We really appreciate you taking the time to help improve the project!