Token tracking

contrib/sarvam-ai/pom.xml

@@ -20,7 +20,7 @@
     <parent>
         <groupId>com.google.adk</groupId>
         <artifactId>google-adk-parent</artifactId>
-        <version>1.4.1-SNAPSHOT</version><!-- {x-version-update:google-adk:current} -->
+        <version>1.5.1-SNAPSHOT</version><!-- {x-version-update:google-adk:current} -->
         <relativePath>../../pom.xml</relativePath>
     </parent>
 

core/pom.xml

@@ -310,6 +310,16 @@
       <version>4.12.0</version>
       <scope>test</scope>
     </dependency>
+    <dependency>
+      <groupId>org.json</groupId>
+      <artifactId>json</artifactId>
+      <version>20180813</version>
+    </dependency>
+    <dependency>
+      <groupId>org.mapdb</groupId>
+      <artifactId>mapdb</artifactId>
+      <version>3.0.8</version>
+    </dependency>
   </dependencies>
   <build>
     <resources>

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

@@ -164,6 +164,21 @@ static Optional<LlmResponse> convertToServerResponse(LiveServerMessage message)
     } else if (message.setupComplete().isPresent()) {
       logger.debug("Received setup complete.");
       return Optional.empty();
+    } else if (message.sessionResumptionUpdate().isPresent()) {
+      logger.debug(
+          "Received session resumption update: {}", message.sessionResumptionUpdate().get());
+      return Optional.empty();
+    } else if (message.goAway().isPresent()) {
+      logger.debug("Received go away: {}", message.goAway().get());
+      return Optional.empty();
+    } else if (message.voiceActivityDetectionSignal().isPresent()) {
+      logger.debug(
+          "Received voice activity detection signal: {}",
+          message.voiceActivityDetectionSignal().get());
+      return Optional.empty();
+    } else if (message.voiceActivity().isPresent()) {
+      logger.debug("Received voice activity: {}", message.voiceActivity().get());
+      return Optional.empty();
     } else if (message.usageMetadata().isEmpty()) {
       logger.warn("Received unknown or empty server message: {}", message.toJson());
       builder

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

@@ -749,7 +749,23 @@ protected Flowable<Event> runLiveImpl(
                       updatedInvocationContext
                           .agent()
                           .runLive(updatedInvocationContext)
-                          .doOnNext(event -> this.sessionService.appendEvent(session, event)))
+                          .doOnNext(event -> this.sessionService.appendEvent(session, event)) 
+                          .concatMapSingle(
+                              event -> {
+                                return updatedInvocationContext
+                                    .pluginManager()
+                                    .onEventCallback(updatedInvocationContext, event)
+                                    .defaultIfEmpty(event);
+                              })
+                          // Run afterRunCallback once the live run completes so plugins can flush
+                          // or log aggregates (e.g. total token usage for the session).
+                          .concatWith(
+                              Completable.defer(
+                                  () -> {
+                                    return updatedInvocationContext
+                                        .pluginManager()
+                                        .afterRunCallback(updatedInvocationContext);
+                                  })))
               .doOnError(
                   throwable -> {
                     Span span = Span.current();

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

@@ -190,6 +190,22 @@ public void convertToServerResponse_withSetupComplete_returnsEmpty() {
     assertThat(result.isPresent()).isFalse();
   }
 
+  @Test
+  public void convertToServerResponse_withSessionResumptionUpdate_returnsEmpty() {
+    LiveServerMessage message =
+        LiveServerMessage.builder()
+            .sessionResumptionUpdate(
+                com.google.genai.types.LiveServerSessionResumptionUpdate.builder()
+                    .newHandle("handle-123")
+                    .resumable(true)
+                    .build())
+            .build();
+
+    Optional<LlmResponse> result = GeminiLlmConnection.convertToServerResponse(message);
+
+    assertThat(result.isPresent()).isFalse();
+  }
+
   @Test
   public void convertToServerResponse_withUnknownMessage_returnsErrorResponse() {
     LiveServerMessage message = LiveServerMessage.builder().build();