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

PiperOrigin-RevId: 781378937

Author: google-genai-bot

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

@@ -18,7 +18,6 @@
 
 import com.google.adk.events.EventActions;
 import com.google.adk.sessions.State;
-import com.google.genai.types.Content;
 import com.google.genai.types.Part;
 import io.reactivex.rxjava3.core.Maybe;
 import java.util.Optional;
@@ -47,11 +46,6 @@ public State state() {
     return state;
   }
 
-  /** Returns the user content that initiated this invocation. */
-  public Optional<Content> userContent() {
-    return invocationContext.userContent();
-  }
-
   /** Returns the EventActions associated with this context. */
   public EventActions eventActions() {
     return eventActions;

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

@@ -16,6 +16,7 @@
 
 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;
@@ -47,6 +48,7 @@
 import com.google.adk.models.BaseLlm;
 import com.google.adk.models.Model;
 import com.google.adk.tools.BaseTool;
+import com.google.adk.tools.BaseToolset;
 import com.google.common.base.Preconditions;
 import com.google.common.collect.ImmutableList;
 import com.google.errorprone.annotations.CanIgnoreReturnValue;
@@ -56,6 +58,7 @@
 import io.reactivex.rxjava3.core.Flowable;
 import io.reactivex.rxjava3.core.Maybe;
 import io.reactivex.rxjava3.core.Single;
+import java.util.ArrayList;
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
@@ -81,7 +84,7 @@ public enum IncludeContents {
   private final Optional<Model> model;
   private final Instruction instruction;
   private final Instruction globalInstruction;
-  private final List<BaseTool> tools;
+  private final List<Object> toolsUnion;
   private final Optional<GenerateContentConfig> generateContentConfig;
   private final Optional<BaseExampleProvider> exampleProvider;
   private final IncludeContents includeContents;
@@ -130,7 +133,7 @@ protected LlmAgent(Builder builder) {
     this.outputSchema = Optional.ofNullable(builder.outputSchema);
     this.executor = Optional.ofNullable(builder.executor);
     this.outputKey = Optional.ofNullable(builder.outputKey);
-    this.tools = builder.tools != null ? builder.tools : ImmutableList.of();
+    this.toolsUnion = builder.toolsUnion != null ? builder.toolsUnion : ImmutableList.of();
 
     this.llmFlow = determineLlmFlow();
 
@@ -153,7 +156,7 @@ public static class Builder {
     private Instruction instruction;
     private Instruction globalInstruction;
     private ImmutableList<BaseAgent> subAgents;
-    private ImmutableList<BaseTool> tools;
+    private ImmutableList<Object> toolsUnion;
     private GenerateContentConfig generateContentConfig;
     private BaseExampleProvider exampleProvider;
     private IncludeContents includeContents;
@@ -234,14 +237,14 @@ public Builder subAgents(BaseAgent... subAgents) {
     }
 
     @CanIgnoreReturnValue
-    public Builder tools(List<? extends BaseTool> tools) {
-      this.tools = ImmutableList.copyOf(tools);
+    public Builder tools(List<?> tools) {
+      this.toolsUnion = ImmutableList.copyOf(tools);
       return this;
     }
 
     @CanIgnoreReturnValue
-    public Builder tools(BaseTool... tools) {
-      this.tools = ImmutableList.copyOf(tools);
+    public Builder tools(Object... tools) {
+      this.toolsUnion = ImmutableList.copyOf(tools);
       return this;
     }
 
@@ -580,7 +583,7 @@ protected void validate() {
                   + ": if outputSchema is set, subAgents must be empty to disable agent"
                   + " transfer.");
         }
-        if (this.tools != null && !this.tools.isEmpty()) {
+        if (this.toolsUnion != null && !this.toolsUnion.isEmpty()) {
           throw new IllegalArgumentException(
               "Invalid config for agent "
                   + this.name
@@ -687,6 +690,42 @@ public Single<String> canonicalGlobalInstruction(ReadonlyContext context) {
     throw new IllegalStateException("Unknown Instruction subtype: " + instruction.getClass());
   }
 
+  /**
+   * Constructs the list of tools for this agent based on the {@link #tools} field.
+   *
+   * <p>This method is only for use by Agent Development Kit.
+   *
+   * @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<>();
+    for (Object toolOrToolset : toolsUnion) {
+      if (toolOrToolset instanceof BaseTool baseTool) {
+        toolSingles.add(Single.just(ImmutableList.of(baseTool)));
+      } else if (toolOrToolset instanceof BaseToolset baseToolset) {
+        toolSingles.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()));
+  }
+
+  /** Overload of canonicalTools that defaults to an empty context. */
+  public Single<List<BaseTool>> canonicalTools() {
+    return canonicalTools(Optional.empty());
+  }
+
+  /** Convenience overload of canonicalTools that accepts a non-optional ReadonlyContext. */
+  public Single<List<BaseTool>> canonicalTools(ReadonlyContext context) {
+    return canonicalTools(Optional.ofNullable(context));
+  }
+
   public Instruction instruction() {
     return instruction;
   }
@@ -719,8 +758,8 @@ public IncludeContents includeContents() {
     return includeContents;
   }
 
-  public List<BaseTool> tools() {
-    return tools;
+  public List<Object> tools() {
+    return toolsUnion;
   }
 
   public boolean disallowTransferToParent() {

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

@@ -17,6 +17,7 @@
 package com.google.adk.agents;
 
 import com.google.common.collect.ImmutableMap;
+import com.google.genai.types.Content;
 import java.util.Map;
 import java.util.Optional;
 
@@ -29,6 +30,11 @@ public ReadonlyContext(InvocationContext invocationContext) {
     this.invocationContext = invocationContext;
   }
 
+  /** Returns the user content that initiated this invocation. */
+  public Optional<Content> userContent() {
+    return invocationContext.userContent();
+  }
+
   /** Returns the ID of the current invocation. */
   public String invocationId() {
     return invocationContext.invocationId();

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

@@ -24,6 +24,7 @@
 import com.google.adk.agents.InvocationContext;
 import com.google.adk.agents.LiveRequest;
 import com.google.adk.agents.LlmAgent;
+import com.google.adk.agents.ReadonlyContext;
 import com.google.adk.agents.RunConfig.StreamingMode;
 import com.google.adk.events.Event;
 import com.google.adk.exceptions.LlmCallsLimitExceededException;
@@ -112,20 +113,22 @@ protected Single<RequestProcessingResult> preprocess(
         processedRequest -> {
           LlmRequest.Builder updatedRequestBuilder = processedRequest.toBuilder();
 
-          Completable toolProcessingCompletable =
-              Flowable.fromIterable(agent.tools())
-                  .concatMapCompletable(
-                      tool ->
-                          tool.processLlmRequest(
-                              updatedRequestBuilder, ToolContext.builder(context).build()));
-
-          return toolProcessingCompletable.andThen(
-              Single.fromCallable(
-                  () -> {
-                    Iterable<Event> combinedEvents = Iterables.concat(eventIterables);
-                    return RequestProcessingResult.create(
-                        updatedRequestBuilder.build(), combinedEvents);
-                  }));
+          return agent
+              .canonicalTools(new ReadonlyContext(context))
+              .flatMapCompletable(
+                  tools ->
+                      Flowable.fromIterable(tools)
+                          .concatMapCompletable(
+                              tool ->
+                                  tool.processLlmRequest(
+                                      updatedRequestBuilder, ToolContext.builder(context).build())))
+              .andThen(
+                  Single.fromCallable(
+                      () -> {
+                        Iterable<Event> combinedEvents = Iterables.concat(eventIterables);
+                        return RequestProcessingResult.create(
+                            updatedRequestBuilder.build(), combinedEvents);
+                      }));
         });
   }
 

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

@@ -0,0 +1,54 @@
+package com.google.adk.tools;
+
+import com.google.adk.agents.ReadonlyContext;
+import io.reactivex.rxjava3.core.Single;
+import java.util.List;
+import java.util.Optional;
+
+/** Base interface for toolsets. */
+public interface BaseToolset extends AutoCloseable {
+
+  /**
+   * Return all tools in the toolset based on the provided context.
+   *
+   * @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);
+
+  /**
+   * Performs cleanup and releases resources held by the toolset.
+   *
+   * <p>NOTE: This method is invoked, for example, at the end of an agent server's lifecycle or when
+   * the toolset is no longer needed. Implementations should ensure that any open connections,
+   * files, or other managed resources are properly released to prevent leaks.
+   */
+  @Override
+  void close() throws Exception;
+
+  /**
+   * Helper method to be used by implementers that returns true if the given tool is in the provided
+   * list of tools of if testing against the given ToolPredicate returns true (otherwise false).
+   *
+   * @param tool The tool to check.
+   * @param toolFilter An Optional containing either a ToolPredicate or a List of tool names.
+   * @param readonlyContext The current context.
+   * @return true if the tool is selected.
+   */
+  default boolean isToolSelected(
+      BaseTool tool, Optional<Object> toolFilter, Optional<ReadonlyContext> readonlyContext) {
+    if (toolFilter.isEmpty()) {
+      return true;
+    }
+    Object filter = toolFilter.get();
+    if (filter instanceof ToolPredicate toolPredicate) {
+      return toolPredicate.test(tool, readonlyContext);
+    }
+    if (filter instanceof List) {
+      @SuppressWarnings("unchecked")
+      List<String> toolNames = (List<String>) filter;
+      return toolNames.contains(tool.name());
+    }
+    return false;
+  }
+}

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

@@ -0,0 +1,20 @@
+package com.google.adk.tools;
+
+import com.google.adk.agents.ReadonlyContext;
+import java.util.Optional;
+
+/**
+ * Functional interface to decide whether a tool should be exposed to the LLM based on the current
+ * context.
+ */
+@FunctionalInterface
+public interface ToolPredicate {
+  /**
+   * Decides if the given tool is selected.
+   *
+   * @param tool The tool to check.
+   * @param readonlyContext The current context.
+   * @return true if the tool should be selected, false otherwise.
+   */
+  boolean test(BaseTool tool, Optional<ReadonlyContext> readonlyContext);
+}