feat: Add support for streaming tools

PiperOrigin-RevId: 788618741

Author: Poggecci

core/src/main/java/com/google/adk/agents/ActiveStreamingTool.java

@@ -0,0 +1,80 @@
+/*
+ * Copyright 2025 Google LLC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.google.adk.agents;
+
+import io.reactivex.rxjava3.disposables.Disposable;
+import javax.annotation.Nullable;
+
+/** Manages streaming tool related resources during invocation. */
+public class ActiveStreamingTool {
+  @Nullable private Disposable task;
+  @Nullable private LiveRequestQueue stream;
+
+  public ActiveStreamingTool(Disposable task) {
+    this(task, null);
+  }
+
+  public ActiveStreamingTool(LiveRequestQueue stream) {
+    this(null, stream);
+  }
+
+  public ActiveStreamingTool(Disposable task, LiveRequestQueue stream) {
+    this.task = task;
+    this.stream = stream;
+  }
+
+  public ActiveStreamingTool() {}
+
+  /**
+   * Returns the active task of this streaming tool.
+   *
+   * @return The active task.
+   */
+  @Nullable
+  public Disposable task() {
+    return task;
+  }
+
+  /**
+   * Sets the active task of this streaming tool.
+   *
+   * @param task The new active task.
+   */
+  public void task(@Nullable Disposable task) {
+    this.task = task;
+  }
+
+  /**
+   * Returns the active stream of this streaming tool.
+   *
+   * @return The active stream.
+   */
+  @Nullable
+  public LiveRequestQueue stream() {
+    return stream;
+  }
+
+  /**
+   * Sets the active stream of this streaming tool.
+   *
+   * @param stream The new active stream.
+   */
+  public void stream(@Nullable LiveRequestQueue stream) {
+    this.stream = stream;
+  }
+}

core/src/main/java/com/google/adk/agents/InvocationContext.java

@@ -21,9 +21,11 @@
 import com.google.adk.sessions.BaseSessionService;
 import com.google.adk.sessions.Session;
 import com.google.genai.types.Content;
+import java.util.Map;
 import java.util.Objects;
 import java.util.Optional;
 import java.util.UUID;
+import java.util.concurrent.ConcurrentHashMap;
 import javax.annotation.Nullable;
 
 /** The context for an agent invocation. */
@@ -32,6 +34,7 @@ public class InvocationContext {
   private final BaseSessionService sessionService;
   private final BaseArtifactService artifactService;
   private final Optional<LiveRequestQueue> liveRequestQueue;
+  private final Map<String, ActiveStreamingTool> activeStreamingTools = new ConcurrentHashMap<>();
 
   private Optional<String> branch;
   private final String invocationId;
@@ -108,17 +111,20 @@ public static InvocationContext create(
   }
 
   public static InvocationContext copyOf(InvocationContext other) {
-    return new InvocationContext(
-        other.sessionService,
-        other.artifactService,
-        other.liveRequestQueue,
-        other.branch,
-        other.invocationId,
-        other.agent,
-        other.session,
-        other.userContent,
-        other.runConfig,
-        other.endInvocation);
+    InvocationContext newContext =
+        new InvocationContext(
+            other.sessionService,
+            other.artifactService,
+            other.liveRequestQueue,
+            other.branch,
+            other.invocationId,
+            other.agent,
+            other.session,
+            other.userContent,
+            other.runConfig,
+            other.endInvocation);
+    newContext.activeStreamingTools.putAll(other.activeStreamingTools);
+    return newContext;
   }
 
   public BaseSessionService sessionService() {
@@ -129,6 +135,10 @@ public BaseArtifactService artifactService() {
     return artifactService;
   }
 
+  public Map<String, ActiveStreamingTool> activeStreamingTools() {
+    return activeStreamingTools;
+  }
+
   public Optional<LiveRequestQueue> liveRequestQueue() {
     return liveRequestQueue;
   }
@@ -217,6 +227,7 @@ public boolean equals(Object o) {
         && Objects.equals(sessionService, that.sessionService)
         && Objects.equals(artifactService, that.artifactService)
         && Objects.equals(liveRequestQueue, that.liveRequestQueue)
+        && Objects.equals(activeStreamingTools, that.activeStreamingTools)
         && Objects.equals(branch, that.branch)
         && Objects.equals(invocationId, that.invocationId)
         && Objects.equals(agent, that.agent)
@@ -231,6 +242,7 @@ public int hashCode() {
         sessionService,
         artifactService,
         liveRequestQueue,
+        activeStreamingTools,
         branch,
         invocationId,
         agent,

core/src/main/java/com/google/adk/flows/llmflows/BaseLlmFlow.java

@@ -17,6 +17,7 @@
 package com.google.adk.flows.llmflows;
 
 import com.google.adk.Telemetry;
+import com.google.adk.agents.ActiveStreamingTool;
 import com.google.adk.agents.BaseAgent;
 import com.google.adk.agents.CallbackContext;
 import com.google.adk.agents.Callbacks.AfterModelCallback;
@@ -168,11 +169,16 @@ protected Single<ResponseProcessingResult> postprocess(
               buildModelResponseEvent(baseEventForLlmResponse, llmRequest, updatedResponse);
           eventIterables.add(Collections.singleton(modelResponseEvent));
 
-          Maybe<Event> maybeFunctionCallEvent =
-              modelResponseEvent.functionCalls().isEmpty()
-                  ? Maybe.empty()
-                  : Functions.handleFunctionCalls(context, modelResponseEvent, llmRequest.tools());
-
+          Maybe<Event> maybeFunctionCallEvent;
+          if (modelResponseEvent.functionCalls().isEmpty()) {
+            maybeFunctionCallEvent = Maybe.empty();
+          } else if (context.runConfig().streamingMode() == StreamingMode.BIDI) {
+            maybeFunctionCallEvent =
+                Functions.handleFunctionCallsLive(context, modelResponseEvent, llmRequest.tools());
+          } else {
+            maybeFunctionCallEvent =
+                Functions.handleFunctionCalls(context, modelResponseEvent, llmRequest.tools());
+          }
           return maybeFunctionCallEvent
               .map(Optional::of)
               .defaultIfEmpty(Optional.empty())
@@ -497,7 +503,22 @@ public Flowable<Event> runLive(InvocationContext invocationContext) {
                             }
                           });
 
-              Flowable<LiveRequest> liveRequests = invocationContext.liveRequestQueue().get().get();
+              Flowable<LiveRequest> liveRequests =
+                  invocationContext
+                      .liveRequestQueue()
+                      .get()
+                      .get()
+                      .doOnNext(
+                          request -> {
+                            if (!invocationContext.activeStreamingTools().isEmpty()) {
+                              for (ActiveStreamingTool activeStreamingTool :
+                                  invocationContext.activeStreamingTools().values()) {
+                                if (activeStreamingTool.stream() != null) {
+                                  activeStreamingTool.stream().send(request);
+                                }
+                              }
+                            }
+                          });
               Disposable sendTask =
                   historySent
                       .observeOn(agent.executor().map(Schedulers::from).orElse(Schedulers.io()))