fix: emit multiple LlmResponses in GeminiLlmConnection

A single LiveServerMessage is now converted to a series of LlmResponse
messages each corresponding to a different part of the
LiveServerMessage, notably the UsageMetadata field is now converted to
a GenerateResponseUsageMetadata and emitted downstream.

PiperOrigin-RevId: 866010045

Author: tilgalas

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

@@ -36,6 +36,7 @@
 import com.google.genai.types.Part;
 import io.reactivex.rxjava3.core.Completable;
 import io.reactivex.rxjava3.core.Flowable;
+import io.reactivex.rxjava3.core.Observable;
 import io.reactivex.rxjava3.processors.PublishProcessor;
 import java.net.SocketException;
 import java.util.List;
@@ -120,53 +121,103 @@ private void handleServerMessage(LiveServerMessage message) {
 
     logger.debug("Received server message: {}", message.toJson());
 
-    Optional<LlmResponse> llmResponse = convertToServerResponse(message);
-    llmResponse.ifPresent(responseProcessor::onNext);
+    Observable<LlmResponse> llmResponse = convertToServerResponse(message);
+    llmResponse.subscribe(responseProcessor::onNext, responseProcessor::onError);
   }
 
   /** Converts a server message into the standardized LlmResponse format. */
-  static Optional<LlmResponse> convertToServerResponse(LiveServerMessage message) {
+  static Observable<LlmResponse> convertToServerResponse(LiveServerMessage message) {
+    return Observable.create(
+        emitter -> {
+          // AtomicBoolean is used to modify state from within lambdas, which
+          // require captured variables to be effectively final.
+          final AtomicBoolean handled = new AtomicBoolean(false);
+          message
+              .serverContent()
+              .ifPresent(
+                  serverContent -> {
+                    emitter.onNext(createServerContentResponse(serverContent));
+                    handled.set(true);
+                  });
+          message
+              .toolCall()
+              .ifPresent(
+                  toolCall -> {
+                    emitter.onNext(createToolCallResponse(toolCall));
+                    handled.set(true);
+                  });
+          message
+              .usageMetadata()
+              .ifPresent(
+                  usageMetadata -> {
+                    logger.debug("Received usage metadata: {}", usageMetadata);
+                    emitter.onNext(createUsageMetadataResponse(usageMetadata));
+                    handled.set(true);
+                  });
+          message
+              .toolCallCancellation()
+              .ifPresent(
+                  toolCallCancellation -> {
+                    logger.debug("Received tool call cancellation: {}", toolCallCancellation);
+                    // TODO: implement proper CFC and thus tool call cancellation handling.
+                    handled.set(true);
+                  });
+          message
+              .setupComplete()
+              .ifPresent(
+                  setupComplete -> {
+                    logger.debug("Received setup complete.");
+                    handled.set(true);
+                  });
+
+          if (!handled.get()) {
+            logger.warn("Received unknown or empty server message: {}", message.toJson());
+            emitter.onNext(createUnknownMessageResponse());
+          }
+          emitter.onComplete();
+        });
+  }
+
+  private static LlmResponse createServerContentResponse(LiveServerContent serverContent) {
     LlmResponse.Builder builder = LlmResponse.builder();
+    serverContent.modelTurn().ifPresent(builder::content);
+    builder
+        .partial(serverContent.turnComplete().map(completed -> !completed).orElse(false))
+        .turnComplete(serverContent.turnComplete().orElse(false))
+        .interrupted(serverContent.interrupted());
+    return builder.build();
+  }
 
-    if (message.serverContent().isPresent()) {
-      LiveServerContent serverContent = message.serverContent().get();
-      serverContent.modelTurn().ifPresent(builder::content);
-      builder
-          .partial(serverContent.turnComplete().map(completed -> !completed).orElse(false))
-          .turnComplete(serverContent.turnComplete().orElse(false))
-          .interrupted(serverContent.interrupted());
-    } else if (message.toolCall().isPresent()) {
-      LiveServerToolCall toolCall = message.toolCall().get();
-      toolCall
-          .functionCalls()
-          .ifPresent(
-              calls -> {
-                for (FunctionCall call : calls) {
-                  builder.content(
-                      Content.builder()
-                          .parts(ImmutableList.of(Part.builder().functionCall(call).build()))
-                          .build());
-                }
-              });
-      builder.partial(false).turnComplete(false);
-    } else if (message.usageMetadata().isPresent()) {
-      logger.debug("Received usage metadata: {}", message.usageMetadata().get());
-      return Optional.empty();
-    } else if (message.toolCallCancellation().isPresent()) {
-      logger.debug("Received tool call cancellation: {}", message.toolCallCancellation().get());
-      // TODO: implement proper CFC and thus tool call cancellation handling.
-      return Optional.empty();
-    } else if (message.setupComplete().isPresent()) {
-      logger.debug("Received setup complete.");
-      return Optional.empty();
-    } else {
-      logger.warn("Received unknown or empty server message: {}", message.toJson());
-      builder
-          .errorCode(new FinishReason("Unknown server message."))
-          .errorMessage("Received unknown server message.");
-    }
+  private static LlmResponse createToolCallResponse(LiveServerToolCall toolCall) {
+    LlmResponse.Builder builder = LlmResponse.builder();
+    toolCall
+        .functionCalls()
+        .ifPresent(
+            calls -> {
+              for (FunctionCall call : calls) {
+                builder.content(
+                    Content.builder()
+                        .parts(ImmutableList.of(Part.builder().functionCall(call).build()))
+                        .build());
+              }
+            });
+    builder.partial(false).turnComplete(false);
+    return builder.build();
+  }
+
+  private static LlmResponse createUsageMetadataResponse(
+      com.google.genai.types.UsageMetadata usageMetadata) {
+    LlmResponse.Builder builder = LlmResponse.builder();
+    builder.usageMetadata(GeminiUtil.toGenerateContentResponseUsageMetadata(usageMetadata));
+    return builder.build();
+  }
 
-    return Optional.of(builder.build());
+  private static LlmResponse createUnknownMessageResponse() {
+    LlmResponse.Builder builder = LlmResponse.builder();
+    builder
+        .errorCode(new FinishReason("Unknown server message."))
+        .errorMessage("Received unknown server message.");
+    return builder.build();
   }
 
   /** Handles errors that occur *during* the initial connection attempt. */

core/src/main/java/com/google/adk/models/GeminiUtil.java

@@ -24,7 +24,9 @@
 import com.google.genai.types.Blob;
 import com.google.genai.types.Content;
 import com.google.genai.types.FileData;
+import com.google.genai.types.GenerateContentResponseUsageMetadata;
 import com.google.genai.types.Part;
+import com.google.genai.types.UsageMetadata;
 import java.util.List;
 import java.util.Optional;
 import java.util.stream.Stream;
@@ -224,4 +226,22 @@ public static List<Content> stripThoughts(List<Content> originalContents) {
             })
         .collect(toImmutableList());
   }
+
+  public static GenerateContentResponseUsageMetadata toGenerateContentResponseUsageMetadata(
+      UsageMetadata usageMetadata) {
+    GenerateContentResponseUsageMetadata.Builder builder =
+        GenerateContentResponseUsageMetadata.builder();
+    usageMetadata.promptTokenCount().ifPresent(builder::promptTokenCount);
+    usageMetadata.cachedContentTokenCount().ifPresent(builder::cachedContentTokenCount);
+    usageMetadata.responseTokenCount().ifPresent(builder::candidatesTokenCount);
+    usageMetadata.toolUsePromptTokenCount().ifPresent(builder::toolUsePromptTokenCount);
+    usageMetadata.thoughtsTokenCount().ifPresent(builder::thoughtsTokenCount);
+    usageMetadata.totalTokenCount().ifPresent(builder::totalTokenCount);
+    usageMetadata.promptTokensDetails().ifPresent(builder::promptTokensDetails);
+    usageMetadata.cacheTokensDetails().ifPresent(builder::cacheTokensDetails);
+    usageMetadata.responseTokensDetails().ifPresent(builder::candidatesTokensDetails);
+    usageMetadata.toolUsePromptTokensDetails().ifPresent(builder::toolUsePromptTokensDetails);
+    usageMetadata.trafficType().ifPresent(builder::trafficType);
+    return builder.build();
+  }
 }