feat: Add BaseToolset and update McpToolset to use the new interface

PiperOrigin-RevId: 781532922

Author: google-genai-bot

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

@@ -16,7 +16,6 @@
 
 package com.google.adk.agents;
 
-import static com.google.common.collect.ImmutableList.toImmutableList;
 import static java.util.stream.Collectors.joining;
 
 import com.fasterxml.jackson.core.JsonProcessingException;
@@ -698,31 +697,29 @@ public Single<String> canonicalGlobalInstruction(ReadonlyContext context) {
    * @param context The context to retrieve the session state.
    * @return The resolved list of tools as a {@link Single} wrapped list of {@link BaseTool}.
    */
-  public Single<List<BaseTool>> canonicalTools(Optional<ReadonlyContext> context) {
-    List<Single<List<BaseTool>>> toolSingles = new ArrayList<>();
+  public Flowable<BaseTool> canonicalTools(Optional<ReadonlyContext> context) {
+    List<Flowable<BaseTool>> toolFlowables = new ArrayList<>();
     for (Object toolOrToolset : toolsUnion) {
       if (toolOrToolset instanceof BaseTool baseTool) {
-        toolSingles.add(Single.just(ImmutableList.of(baseTool)));
+        toolFlowables.add(Flowable.just(baseTool));
       } else if (toolOrToolset instanceof BaseToolset baseToolset) {
-        toolSingles.add(baseToolset.getTools(context.orElse(null)));
+        toolFlowables.add(baseToolset.getTools(context.orElse(null)));
       } else {
         throw new IllegalArgumentException(
             "Object in tools list is not of a supported type: "
                 + toolOrToolset.getClass().getName());
       }
     }
-    return Single.concat(toolSingles)
-        .toList()
-        .map(listOfLists -> listOfLists.stream().flatMap(List::stream).collect(toImmutableList()));
+    return Flowable.concat(toolFlowables);
   }
 
   /** Overload of canonicalTools that defaults to an empty context. */
-  public Single<List<BaseTool>> canonicalTools() {
+  public Flowable<BaseTool> canonicalTools() {
     return canonicalTools(Optional.empty());
   }
 
   /** Convenience overload of canonicalTools that accepts a non-optional ReadonlyContext. */
-  public Single<List<BaseTool>> canonicalTools(ReadonlyContext context) {
+  public Flowable<BaseTool> canonicalTools(ReadonlyContext context) {
     return canonicalTools(Optional.ofNullable(context));
   }
 
@@ -758,7 +755,11 @@ public IncludeContents includeContents() {
     return includeContents;
   }
 
-  public List<Object> tools() {
+  public List<BaseTool> tools() {
+    return canonicalTools().toList().blockingGet();
+  }
+
+  public List<Object> toolsUnion() {
     return toolsUnion;
   }
 

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

@@ -115,13 +115,10 @@ protected Single<RequestProcessingResult> preprocess(
 
           return agent
               .canonicalTools(new ReadonlyContext(context))
-              .flatMapCompletable(
-                  tools ->
-                      Flowable.fromIterable(tools)
-                          .concatMapCompletable(
-                              tool ->
-                                  tool.processLlmRequest(
-                                      updatedRequestBuilder, ToolContext.builder(context).build())))
+              .concatMapCompletable(
+                  tool ->
+                      tool.processLlmRequest(
+                          updatedRequestBuilder, ToolContext.builder(context).build()))
               .andThen(
                   Single.fromCallable(
                       () -> {

core/src/main/java/com/google/adk/tools/BaseToolset.java

@@ -1,7 +1,7 @@
 package com.google.adk.tools;
 
 import com.google.adk.agents.ReadonlyContext;
-import io.reactivex.rxjava3.core.Single;
+import io.reactivex.rxjava3.core.Flowable;
 import java.util.List;
 import java.util.Optional;
 
@@ -14,7 +14,7 @@ public interface BaseToolset extends AutoCloseable {
    * @param readonlyContext Context used to filter tools available to the agent.
    * @return A Single emitting a list of tools available under the specified context.
    */
-  Single<List<BaseTool>> getTools(ReadonlyContext readonlyContext);
+  Flowable<BaseTool> getTools(ReadonlyContext readonlyContext);
 
   /**
    * Performs cleanup and releases resources held by the toolset.

core/src/main/java/com/google/adk/tools/mcp/McpToolset.java

@@ -26,8 +26,7 @@
 import io.modelcontextprotocol.client.McpSyncClient;
 import io.modelcontextprotocol.client.transport.ServerParameters;
 import io.modelcontextprotocol.spec.McpSchema.ListToolsResult;
-import io.reactivex.rxjava3.core.Single;
-import java.util.List;
+import io.reactivex.rxjava3.core.Flowable;
 import java.util.Objects;
 import java.util.Optional;
 import org.slf4j.Logger;
@@ -151,69 +150,74 @@ public McpToolset(ServerParameters connectionParams) {
   }
 
   @Override
-  public Single<List<BaseTool>> getTools(ReadonlyContext readonlyContext) {
-    return Single.fromCallable(
-        () -> {
-          for (int i = 0; i < MAX_RETRIES; i++) {
-            try {
-              if (this.mcpSession == null) {
-                logger.info("MCP session is null or closed, initializing (attempt {}).", i + 1);
-                this.mcpSession = this.mcpSessionManager.createSession();
-              }
-
-              ListToolsResult toolsResponse = this.mcpSession.listTools();
-              return toolsResponse.tools().stream()
-                  .map(
-                      tool ->
-                          new McpTool(
-                              tool, this.mcpSession, this.mcpSessionManager, this.objectMapper))
-                  .filter(
-                      tool ->
-                          isToolSelected(tool, toolFilter, Optional.ofNullable(readonlyContext)))
-                  .collect(toImmutableList());
-            } catch (IllegalArgumentException e) {
-              // This could happen if parameters for tool loading are somehow invalid.
-              // This is likely a fatal error and should not be retried.
-              logger.error("Invalid argument encountered during tool loading.", e);
-              throw new McpToolLoadingException(
-                  "Invalid argument encountered during tool loading.", e);
-            } catch (RuntimeException e) { // Catch any other unexpected runtime exceptions
-              logger.error("Unexpected error during tool loading, retry attempt " + (i + 1), e);
-              if (i < MAX_RETRIES - 1) {
-                // For other general exceptions, we might still want to retry if they are
-                // potentially transient, or if we don't have more specific handling. But it's
-                // better to be specific. For now, we'll treat them as potentially retryable but log
-                // them at a higher level.
+  public Flowable<BaseTool> getTools(ReadonlyContext readonlyContext) {
+    return Flowable.fromCallable(
+            () -> {
+              for (int i = 0; i < MAX_RETRIES; i++) {
                 try {
-                  logger.info("Reinitializing MCP session before next retry for unexpected error.");
-                  this.mcpSession = this.mcpSessionManager.createSession();
-                  Thread.sleep(RETRY_DELAY_MILLIS);
-                } catch (InterruptedException ie) {
-                  Thread.currentThread().interrupt();
-                  logger.error(
-                      "Interrupted during retry delay for loadTools (unexpected error).", ie);
+                  if (this.mcpSession == null) {
+                    logger.info("MCP session is null or closed, initializing (attempt {}).", i + 1);
+                    this.mcpSession = this.mcpSessionManager.createSession();
+                  }
+
+                  ListToolsResult toolsResponse = this.mcpSession.listTools();
+                  return toolsResponse.tools().stream()
+                      .map(
+                          tool ->
+                              new McpTool(
+                                  tool, this.mcpSession, this.mcpSessionManager, this.objectMapper))
+                      .filter(
+                          tool ->
+                              isToolSelected(
+                                  tool, toolFilter, Optional.ofNullable(readonlyContext)))
+                      .collect(toImmutableList());
+                } catch (IllegalArgumentException e) {
+                  // This could happen if parameters for tool loading are somehow invalid.
+                  // This is likely a fatal error and should not be retried.
+                  logger.error("Invalid argument encountered during tool loading.", e);
                   throw new McpToolLoadingException(
-                      "Interrupted during retry delay (unexpected error)", ie);
-                } catch (RuntimeException reinitE) {
-                  logger.error(
-                      "Failed to reinitialize session during retry (unexpected error).", reinitE);
-                  throw new McpInitializationException(
-                      "Failed to reinitialize session during tool loading retry (unexpected"
-                          + " error).",
-                      reinitE);
+                      "Invalid argument encountered during tool loading.", e);
+                } catch (RuntimeException e) { // Catch any other unexpected runtime exceptions
+                  logger.error("Unexpected error during tool loading, retry attempt " + (i + 1), e);
+                  if (i < MAX_RETRIES - 1) {
+                    // For other general exceptions, we might still want to retry if they are
+                    // potentially transient, or if we don't have more specific handling. But it's
+                    // better to be specific. For now, we'll treat them as potentially retryable but
+                    // log
+                    // them at a higher level.
+                    try {
+                      logger.info(
+                          "Reinitializing MCP session before next retry for unexpected error.");
+                      this.mcpSession = this.mcpSessionManager.createSession();
+                      Thread.sleep(RETRY_DELAY_MILLIS);
+                    } catch (InterruptedException ie) {
+                      Thread.currentThread().interrupt();
+                      logger.error(
+                          "Interrupted during retry delay for loadTools (unexpected error).", ie);
+                      throw new McpToolLoadingException(
+                          "Interrupted during retry delay (unexpected error)", ie);
+                    } catch (RuntimeException reinitE) {
+                      logger.error(
+                          "Failed to reinitialize session during retry (unexpected error).",
+                          reinitE);
+                      throw new McpInitializationException(
+                          "Failed to reinitialize session during tool loading retry (unexpected"
+                              + " error).",
+                          reinitE);
+                    }
+                  } else {
+                    logger.error(
+                        "Failed to load tools after multiple retries due to unexpected error.", e);
+                    throw new McpToolLoadingException(
+                        "Failed to load tools after multiple retries due to unexpected error.", e);
+                  }
                 }
-              } else {
-                logger.error(
-                    "Failed to load tools after multiple retries due to unexpected error.", e);
-                throw new McpToolLoadingException(
-                    "Failed to load tools after multiple retries due to unexpected error.", e);
               }
-            }
-          }
-          // This line should ideally not be reached if retries are handled correctly or an
-          // exception is always thrown.
-          throw new IllegalStateException("Unexpected state in getTools retry loop");
-        });
+              // This line should ideally not be reached if retries are handled correctly or an
+              // exception is always thrown.
+              throw new IllegalStateException("Unexpected state in getTools retry loop");
+            })
+        .flatMapIterable(tools -> tools);
   }
 
   @Override

core/src/test/java/com/google/adk/tools/BaseToolsetTest.java

@@ -4,8 +4,7 @@
 import static org.mockito.Mockito.mock;
 
 import com.google.adk.agents.ReadonlyContext;
-import com.google.common.collect.ImmutableList;
-import io.reactivex.rxjava3.core.Single;
+import io.reactivex.rxjava3.core.Flowable;
 import java.util.List;
 import org.junit.Test;
 import org.junit.runner.RunWith;
@@ -23,15 +22,15 @@ public void testGetTools() {
     BaseToolset toolset =
         new BaseToolset() {
           @Override
-          public Single<List<BaseTool>> getTools(ReadonlyContext readonlyContext) {
-            return Single.just(ImmutableList.of(mockTool1, mockTool2));
+          public Flowable<BaseTool> getTools(ReadonlyContext readonlyContext) {
+            return Flowable.just(mockTool1, mockTool2);
           }
 
           @Override
           public void close() throws Exception {}
         };
 
-    List<BaseTool> tools = toolset.getTools(mockContext).blockingGet();
+    List<BaseTool> tools = toolset.getTools(mockContext).toList().blockingGet();
     assertThat(tools).containsExactly(mockTool1, mockTool2);
   }
 }

dev/src/main/java/com/google/adk/web/AgentGraphGenerator.java

@@ -105,7 +105,7 @@ private static void buildGraphRecursive(
 
     if (agent instanceof LlmAgent) {
       LlmAgent llmAgent = (LlmAgent) agent;
-      List<BaseTool> tools = llmAgent.canonicalTools().blockingGet();
+      List<BaseTool> tools = llmAgent.canonicalTools().toList().blockingGet();
       if (tools != null) {
         for (BaseTool tool : tools) {
           if (tool != null) {