Merge pull request #1587 from DABH:deterministic-scheduler-seam

PiperOrigin-RevId: 993621169

Author: copybara-github

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

@@ -39,6 +39,8 @@
 import com.google.genai.types.Content;
 import com.google.genai.types.FunctionCall;
 import com.google.genai.types.FunctionResponse;
+import io.reactivex.rxjava3.core.Scheduler;
+import io.reactivex.rxjava3.schedulers.Schedulers;
 import java.util.ArrayList;
 import java.util.Collections;
 import java.util.HashSet;
@@ -74,6 +76,7 @@ public class InvocationContext {
   // Shared by reference so a sub-agent's checkpoint is visible to its parent and the runner.
   private final Map<String, Map<String, Object>> agentStates;
   private final Map<String, Boolean> endOfAgents;
+  @Nullable private final Scheduler scheduler;
 
   @Nullable private String branch;
   private BaseAgent agent;
@@ -103,6 +106,7 @@ protected InvocationContext(Builder builder) {
     this.callbackContextData = builder.callbackContextData;
     this.agentStates = builder.agentStates;
     this.endOfAgents = builder.endOfAgents;
+    this.scheduler = builder.scheduler;
   }
 
   /** Returns a new {@link Builder} for creating {@link InvocationContext} instances. */
@@ -288,6 +292,17 @@ public String userId() {
     return session.userId();
   }
 
+  /**
+   * Returns the {@link Scheduler} on which this invocation runs the work ADK hands to worker
+   * threads: the sub-agents of a {@code ParallelAgent}, tool calls in {@code
+   * ToolExecutionMode.PARALLEL_SUBSCRIBE} mode, and the hop that starts the live send loop. An
+   * agent-level setting ({@code ParallelAgent.Builder#scheduler}, {@code
+   * LlmAgent.Builder#executor}) takes precedence over it. Defaults to {@link Schedulers#io()}.
+   */
+  public Scheduler scheduler() {
+    return scheduler != null ? scheduler : Schedulers.io();
+  }
+
   /** Generates a new unique ID for an invocation context. */
   public static String newInvocationContextId() {
     return "e-" + UUID.randomUUID();
@@ -620,6 +635,7 @@ private Builder(InvocationContext context) {
       // Shared by reference, not copied: the checkpoints belong to the invocation, not a context.
       this.agentStates = context.agentStates;
       this.endOfAgents = context.endOfAgents;
+      this.scheduler = context.scheduler;
     }
 
     private BaseSessionService sessionService;
@@ -642,6 +658,7 @@ private Builder(InvocationContext context) {
     private Map<String, Object> callbackContextData = new ConcurrentHashMap<>();
     private Map<String, Map<String, Object>> agentStates = new ConcurrentHashMap<>();
     private Map<String, Boolean> endOfAgents = new ConcurrentHashMap<>();
+    @Nullable private Scheduler scheduler = null;
 
     /**
      * Sets the session service for managing session state.
@@ -823,6 +840,19 @@ public Builder resumabilityConfig(@Nullable ResumabilityConfig resumabilityConfi
       return this;
     }
 
+    /**
+     * Sets the scheduler for the invocation's internal fan-out; see {@link
+     * InvocationContext#scheduler()}. {@code null} (the default) selects {@link Schedulers#io()}.
+     *
+     * @param scheduler the scheduler, or {@code null} for the default.
+     * @return this builder instance for chaining.
+     */
+    @CanIgnoreReturnValue
+    public Builder scheduler(@Nullable Scheduler scheduler) {
+      this.scheduler = scheduler;
+      return this;
+    }
+
     /**
      * Sets the callback context data for the invocation.
      *
@@ -892,7 +922,8 @@ public boolean equals(Object o) {
         && Objects.equals(contextCacheConfig, that.contextCacheConfig)
         && Objects.equals(resumabilityConfig, that.resumabilityConfig)
         && Objects.equals(invocationCostManager, that.invocationCostManager)
-        && Objects.equals(callbackContextData, that.callbackContextData);
+        && Objects.equals(callbackContextData, that.callbackContextData)
+        && Objects.equals(scheduler, that.scheduler);
   }
 
   @Override
@@ -915,6 +946,7 @@ public int hashCode() {
         contextCacheConfig,
         resumabilityConfig,
         invocationCostManager,
-        callbackContextData);
+        callbackContextData,
+        scheduler);
   }
 }

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

@@ -553,6 +553,11 @@ public Builder outputSchema(Schema outputSchema) {
       return this;
     }
 
+    /**
+     * Sets the executor for this agent's offloaded work: tool calls in {@code
+     * ToolExecutionMode.PARALLEL_SUBSCRIBE} mode and the live request loop. When unset, that work
+     * runs on the invocation's scheduler ({@link InvocationContext#scheduler()}).
+     */
     @CanIgnoreReturnValue
     public Builder executor(Executor executor) {
       this.executor = executor;

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

@@ -24,12 +24,12 @@
 import com.google.errorprone.annotations.CanIgnoreReturnValue;
 import io.reactivex.rxjava3.core.Flowable;
 import io.reactivex.rxjava3.core.Scheduler;
-import io.reactivex.rxjava3.schedulers.Schedulers;
 import java.util.ArrayList;
 import java.util.List;
 import java.util.Objects;
 import java.util.Set;
 import java.util.concurrent.atomic.AtomicBoolean;
+import org.jspecify.annotations.Nullable;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -75,7 +75,7 @@
 public class ParallelAgent extends BaseAgent {
 
   private static final Logger logger = LoggerFactory.getLogger(ParallelAgent.class);
-  private final Scheduler scheduler;
+  @Nullable private final Scheduler scheduler;
 
   /**
    * Constructor for ParallelAgent.
@@ -85,27 +85,38 @@ public class ParallelAgent extends BaseAgent {
    * @param subAgents The list of sub-agents to run in parallel.
    * @param beforeAgentCallback Optional callback before the agent runs.
    * @param afterAgentCallback Optional callback after the agent runs.
-   * @param scheduler The scheduler to use for parallel execution.
+   * @param scheduler The scheduler to use for parallel execution, or {@code null} to use the
+   *     invocation's scheduler.
    */
   private ParallelAgent(
       String name,
       String description,
       List<? extends BaseAgent> subAgents,
       List<Callbacks.BeforeAgentCallback> beforeAgentCallback,
       List<Callbacks.AfterAgentCallback> afterAgentCallback,
-      Scheduler scheduler) {
+      @Nullable Scheduler scheduler) {
 
     super(name, description, subAgents, beforeAgentCallback, afterAgentCallback);
     this.scheduler = scheduler;
   }
 
+  /** Returns the agent's own scheduler when configured, otherwise the invocation's. */
+  private Scheduler schedulerFor(InvocationContext invocationContext) {
+    return scheduler != null ? scheduler : invocationContext.scheduler();
+  }
+
   /** Builder for {@link ParallelAgent}. */
   public static class Builder extends BaseAgent.Builder<Builder> {
 
-    private Scheduler scheduler = Schedulers.io();
+    @Nullable private Scheduler scheduler = null;
 
+    /**
+     * Sets the scheduler the sub-agents are subscribed on, or {@code null} (the default) to inherit
+     * the invocation's scheduler ({@link InvocationContext#scheduler()}: {@code Schedulers.io()}
+     * unless the {@code Runner} is configured with another one).
+     */
     @CanIgnoreReturnValue
-    public Builder scheduler(Scheduler scheduler) {
+    public Builder scheduler(@Nullable Scheduler scheduler) {
       this.scheduler = scheduler;
       return this;
     }
@@ -186,11 +197,13 @@ protected Flowable<Event> runAsyncImpl(InvocationContext invocationContext) {
             .filter(Objects::nonNull)
             .collect(ImmutableSet.toImmutableSet());
 
+    Scheduler fanOutScheduler = schedulerFor(invocationContext);
     if (!invocationContext.isResumable()) {
       var updatedInvocationContext = setBranchForCurrentAgent(this, invocationContext);
       List<Flowable<Event>> agentFlowables = new ArrayList<>();
       for (BaseAgent subAgent : currentSubAgents) {
-        agentFlowables.add(subAgent.runAsync(updatedInvocationContext).subscribeOn(scheduler));
+        agentFlowables.add(
+            subAgent.runAsync(updatedInvocationContext).subscribeOn(fanOutScheduler));
       }
       return Flowable.merge(agentFlowables)
           .takeUntil((Event event) -> asksThisAgentToExit(event, directSubAgentNames));
@@ -219,7 +232,7 @@ protected Flowable<Event> runAsyncImpl(InvocationContext invocationContext) {
             agentFlowables.add(
                 subAgent
                     .runAsync(updatedInvocationContext)
-                    .subscribeOn(scheduler)
+                    .subscribeOn(fanOutScheduler)
                     .doOnNext(
                         event -> {
                           if (invocationContext.shouldPauseInvocation(event)) {

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

@@ -58,7 +58,8 @@ public enum StreamingMode {
    *
    * <p>PARALLEL_SUBSCRIBE: like {@code PARALLEL}, but every tool is additionally subscribed on a
    * worker thread, so blocking tools also run concurrently. Tool implementations must be
-   * thread-safe. The worker is the agent's executor when set, otherwise the RxJava IO scheduler.
+   * thread-safe. The worker is the agent's executor when set, otherwise the invocation's scheduler
+   * ({@code Runner.Builder#scheduler}; the RxJava IO scheduler by default).
    */
   public enum ToolExecutionMode {
     NONE,

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

@@ -55,7 +55,6 @@
 import io.reactivex.rxjava3.core.Single;
 import io.reactivex.rxjava3.disposables.Disposable;
 import io.reactivex.rxjava3.observers.DisposableCompletableObserver;
-import io.reactivex.rxjava3.schedulers.Schedulers;
 import java.util.ArrayList;
 import java.util.List;
 import java.util.Optional;
@@ -636,7 +635,7 @@ public Flowable<Event> runLive(InvocationContext invocationContext) {
                           });
               Disposable sendTask =
                   historySent
-                      .observeOn(agent.executor().map(Schedulers::from).orElse(Schedulers.io()))
+                      .observeOn(Functions.resolveWorkerScheduler(invocationContext))
                       .andThen(
                           liveRequests
                               .onBackpressureBuffer()

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

@@ -263,7 +263,7 @@ private static Observable<Event> buildToolExecutionObservable(
       return Observable.fromIterable(validFunctionCalls).concatMapMaybe(functionCallMapper);
     }
     if (mode == ToolExecutionMode.PARALLEL_SUBSCRIBE) {
-      Scheduler scheduler = resolveToolExecutionScheduler(invocationContext);
+      Scheduler scheduler = resolveWorkerScheduler(invocationContext);
       return Observable.fromIterable(validFunctionCalls)
           .concatMapEager(
               call -> functionCallMapper.apply(call).toObservable().subscribeOn(scheduler));
@@ -274,12 +274,16 @@ private static Observable<Event> buildToolExecutionObservable(
         .concatMapEager(call -> functionCallMapper.apply(call).toObservable());
   }
 
-  /** Agent executor if set, otherwise the IO scheduler. */
-  private static Scheduler resolveToolExecutionScheduler(InvocationContext invocationContext) {
+  /**
+   * Returns the scheduler for the work this invocation hands to a worker (tool calls in {@code
+   * ToolExecutionMode.PARALLEL_SUBSCRIBE} mode and the live send loop): the agent's executor when
+   * set, otherwise the invocation's scheduler.
+   */
+  static Scheduler resolveWorkerScheduler(InvocationContext invocationContext) {
     if (invocationContext.agent() instanceof LlmAgent llmAgent) {
-      return llmAgent.executor().map(Schedulers::from).orElse(Schedulers.io());
+      return llmAgent.executor().map(Schedulers::from).orElseGet(invocationContext::scheduler);
     }
-    return Schedulers.io();
+    return invocationContext.scheduler();
   }
 
   private static Function<FunctionCall, Maybe<Event>> getFunctionCallMapper(

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

@@ -71,6 +71,7 @@
 import io.reactivex.rxjava3.core.Completable;
 import io.reactivex.rxjava3.core.Flowable;
 import io.reactivex.rxjava3.core.Maybe;
+import io.reactivex.rxjava3.core.Scheduler;
 import io.reactivex.rxjava3.core.Single;
 import io.reactivex.rxjava3.subjects.CompletableSubject;
 import java.util.ArrayList;
@@ -96,6 +97,7 @@ public class Runner {
   @Nullable private final EventsCompactionConfig eventsCompactionConfig;
   @Nullable private final ContextCacheConfig contextCacheConfig;
   private final @Nullable ResumabilityConfig resumabilityConfig;
+  @Nullable private final Scheduler scheduler;
   private final ConcurrentMap<String, Completable> activeSessionCompletables =
       new MapMaker().weakValues().makeMap();
 
@@ -108,6 +110,7 @@ public static class Builder {
     private BaseSessionService sessionService = new InMemorySessionService();
     @Nullable private BaseMemoryService memoryService = null;
     private List<? extends Plugin> plugins = ImmutableList.of();
+    @Nullable private Scheduler scheduler = null;
 
     @CanIgnoreReturnValue
     public Builder app(App app) {
@@ -162,6 +165,25 @@ public Builder plugins(Plugin... plugins) {
       return this;
     }
 
+    /**
+     * Sets the {@link Scheduler} for the work ADK hands off to worker threads in this runner's
+     * invocations: the sub-agents of a {@code ParallelAgent}, tool calls in {@link
+     * RunConfig.ToolExecutionMode#PARALLEL_SUBSCRIBE} mode, and the hop that starts the live send
+     * loop. The runs that an {@code AgentTool} nests inherit it. An agent-level {@code
+     * ParallelAgent.Builder#scheduler} or {@code LlmAgent.Builder#executor} takes precedence.
+     * {@code null} (the default) selects {@code Schedulers.io()}.
+     *
+     * <p>Only these hand-offs are affected; work that completes asynchronously elsewhere (model
+     * clients, session services, MCP, plugins) stays on its own threads. {@code
+     * Schedulers.trampoline()} runs the hand-offs inline, so sub-agents and tools are subscribed
+     * sequentially, but anything they start asynchronously can still overlap.
+     */
+    @CanIgnoreReturnValue
+    public Builder scheduler(@Nullable Scheduler scheduler) {
+      this.scheduler = scheduler;
+      return this;
+    }
+
     public Runner build() {
       BaseAgent buildAgent;
       String buildAppName;
@@ -213,7 +235,8 @@ public Runner build() {
           buildPlugins,
           buildEventsCompactionConfig,
           buildContextCacheConfig,
-          buildResumabilityConfig);
+          buildResumabilityConfig,
+          scheduler);
     }
   }
 
@@ -295,6 +318,30 @@ protected Runner(
       @Nullable EventsCompactionConfig eventsCompactionConfig,
       @Nullable ContextCacheConfig contextCacheConfig,
       @Nullable ResumabilityConfig resumabilityConfig) {
+    this(
+        agent,
+        appName,
+        artifactService,
+        sessionService,
+        memoryService,
+        plugins,
+        eventsCompactionConfig,
+        contextCacheConfig,
+        resumabilityConfig,
+        /* scheduler= */ null);
+  }
+
+  private Runner(
+      BaseAgent agent,
+      String appName,
+      BaseArtifactService artifactService,
+      BaseSessionService sessionService,
+      @Nullable BaseMemoryService memoryService,
+      List<? extends Plugin> plugins,
+      @Nullable EventsCompactionConfig eventsCompactionConfig,
+      @Nullable ContextCacheConfig contextCacheConfig,
+      @Nullable ResumabilityConfig resumabilityConfig,
+      @Nullable Scheduler scheduler) {
     this.agent = agent;
     this.appName = appName;
     this.artifactService = artifactService;
@@ -304,6 +351,7 @@ protected Runner(
     this.eventsCompactionConfig = createEventsCompactionConfig(agent, eventsCompactionConfig);
     this.contextCacheConfig = contextCacheConfig;
     this.resumabilityConfig = resumabilityConfig;
+    this.scheduler = scheduler;
   }
 
   /**
@@ -1217,6 +1265,7 @@ private InvocationContext.Builder newInvocationContextBuilder(Session session, B
         .eventsCompactionConfig(this.eventsCompactionConfig)
         .contextCacheConfig(this.contextCacheConfig)
         .resumabilityConfig(this.resumabilityConfig)
+        .scheduler(this.scheduler)
         .agent(agent);
   }
 

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

@@ -28,8 +28,8 @@
 import com.google.adk.agents.RunConfig;
 import com.google.adk.agents.RunConfig.StreamingMode;
 import com.google.adk.events.Event;
+import com.google.adk.memory.InMemoryMemoryService;
 import com.google.adk.plugins.Plugin;
-import com.google.adk.runner.InMemoryRunner;
 import com.google.adk.runner.Runner;
 import com.google.adk.sessions.State;
 import com.google.common.annotations.VisibleForTesting;
@@ -187,7 +187,15 @@ public Single<Map<String, Object>> runAsync(Map<String, Object> args, ToolContex
         this.includePlugins
             ? ImmutableList.of(toolContext.invocationContext().pluginManager())
             : ImmutableList.of();
-    Runner runner = new InMemoryRunner(this.agent, toolContext.agentName(), plugins);
+    // The nested run inherits the caller's scheduler so its fan-out stays on the same scheduler.
+    Runner runner =
+        Runner.builder()
+            .agent(this.agent)
+            .appName(toolContext.agentName())
+            .plugins(plugins)
+            .memoryService(new InMemoryMemoryService())
+            .scheduler(toolContext.invocationContext().scheduler())
+            .build();
     // Follow the caller's RunConfig but run unary: only the last event becomes the result.
     RunConfig callerRunConfig = toolContext.invocationContext().runConfig();
     RunConfig runConfig =