From 1b89d23b78c09b87b0dd7d2d7daac866fd9fc5c0 Mon Sep 17 00:00:00 2001 From: DABH Date: Thu, 1 Oct 2026 02:48:52 -0500 Subject: [PATCH] feat: add an injectable scheduler for ADK's internal fan-out ADK hard-codes the RxJava IO scheduler wherever it leaves the caller's thread: ParallelAgent subscribes each sub-agent on Schedulers.io(), tool calls in ToolExecutionMode.PARALLEL_SUBSCRIBE are subscribed on it unless the LlmAgent has an executor, and the live send loop is started on it. An embedding application has no way to influence that fan-out, for example to cap it, to run it on an existing worker pool, or to keep every ADK task on a dispatcher it controls, which a durable-execution host must. adk-go exposes this control point as platform.WithTaskRunner / platform.RunTasks (google/adk-go#1050). Add the Java analog as an RxJava Scheduler carried on InvocationContext. Runner.builder().scheduler(...) sets it once, every InvocationContext the runner creates carries it (and toBuilder() copies it), and InvocationContext.scheduler() returns it or Schedulers.io() when none is configured, so the default behavior is unchanged. ParallelAgent uses the invocation's scheduler unless its own builder set one (its default moves from an eager Schedulers.io() to "inherit"), Functions resolves the PARALLEL_SUBSCRIBE worker from the LlmAgent executor and then the invocation, and BaseLlmFlow's live send loop goes through the same helper. AgentTool builds the runner it nests with Runner.builder() and hands it the caller's scheduler, so a nested agent fans out on the same threads as its caller. It is an RxJava Scheduler rather than a task-runner function because every hop site already takes one, and because a Scheduler can also express the delayed scheduling RxJava's time-based operators need, which a plain Executor cannot. Agent-level settings keep precedence, so existing ParallelAgent.Builder.scheduler() and LlmAgent.Builder.executor() callers see the same threading. A null scheduler means "inherit the invocation's" on ParallelAgent.Builder and "use the default" on Runner.Builder, and InvocationContext.equals()/hashCode() compare the configured scheduler field, so a context still equals its toBuilder() copy. --- .../google/adk/agents/InvocationContext.java | 36 +++++- .../java/com/google/adk/agents/LlmAgent.java | 5 + .../com/google/adk/agents/ParallelAgent.java | 29 +++-- .../java/com/google/adk/agents/RunConfig.java | 3 +- .../adk/flows/llmflows/BaseLlmFlow.java | 3 +- .../google/adk/flows/llmflows/Functions.java | 14 ++- .../java/com/google/adk/runner/Runner.java | 51 +++++++- .../java/com/google/adk/tools/AgentTool.java | 12 +- .../adk/agents/InvocationContextTest.java | 69 +++++++++++ .../google/adk/agents/ParallelAgentTest.java | 86 +++++++++++++ .../adk/flows/llmflows/FunctionsTest.java | 86 +++++++++++++ .../com/google/adk/runner/RunnerTest.java | 115 ++++++++++++++++++ .../adk/testing/RecordingScheduler.java | 59 +++++++++ .../com/google/adk/tools/AgentToolTest.java | 30 +++++ 14 files changed, 577 insertions(+), 21 deletions(-) create mode 100644 core/src/test/java/com/google/adk/testing/RecordingScheduler.java diff --git a/core/src/main/java/com/google/adk/agents/InvocationContext.java b/core/src/main/java/com/google/adk/agents/InvocationContext.java index 15175d1cd..6651c5720 100644 --- a/core/src/main/java/com/google/adk/agents/InvocationContext.java +++ b/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> agentStates; private final Map 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 callbackContextData = new ConcurrentHashMap<>(); private Map> agentStates = new ConcurrentHashMap<>(); private Map 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); } } diff --git a/core/src/main/java/com/google/adk/agents/LlmAgent.java b/core/src/main/java/com/google/adk/agents/LlmAgent.java index d36a2b2f6..69360943b 100644 --- a/core/src/main/java/com/google/adk/agents/LlmAgent.java +++ b/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; diff --git a/core/src/main/java/com/google/adk/agents/ParallelAgent.java b/core/src/main/java/com/google/adk/agents/ParallelAgent.java index e78d47796..83ebbaf5d 100644 --- a/core/src/main/java/com/google/adk/agents/ParallelAgent.java +++ b/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,7 +85,8 @@ 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, @@ -93,19 +94,29 @@ private ParallelAgent( List subAgents, List beforeAgentCallback, List 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 { - 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 runAsyncImpl(InvocationContext invocationContext) { .filter(Objects::nonNull) .collect(ImmutableSet.toImmutableSet()); + Scheduler fanOutScheduler = schedulerFor(invocationContext); if (!invocationContext.isResumable()) { var updatedInvocationContext = setBranchForCurrentAgent(this, invocationContext); List> 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 runAsyncImpl(InvocationContext invocationContext) { agentFlowables.add( subAgent .runAsync(updatedInvocationContext) - .subscribeOn(scheduler) + .subscribeOn(fanOutScheduler) .doOnNext( event -> { if (invocationContext.shouldPauseInvocation(event)) { diff --git a/core/src/main/java/com/google/adk/agents/RunConfig.java b/core/src/main/java/com/google/adk/agents/RunConfig.java index bd20b6183..4023eaef4 100644 --- a/core/src/main/java/com/google/adk/agents/RunConfig.java +++ b/core/src/main/java/com/google/adk/agents/RunConfig.java @@ -58,7 +58,8 @@ public enum StreamingMode { * *

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, diff --git a/core/src/main/java/com/google/adk/flows/llmflows/BaseLlmFlow.java b/core/src/main/java/com/google/adk/flows/llmflows/BaseLlmFlow.java index f255bbe03..de011f188 100644 --- a/core/src/main/java/com/google/adk/flows/llmflows/BaseLlmFlow.java +++ b/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 runLive(InvocationContext invocationContext) { }); Disposable sendTask = historySent - .observeOn(agent.executor().map(Schedulers::from).orElse(Schedulers.io())) + .observeOn(Functions.resolveWorkerScheduler(invocationContext)) .andThen( liveRequests .onBackpressureBuffer() diff --git a/core/src/main/java/com/google/adk/flows/llmflows/Functions.java b/core/src/main/java/com/google/adk/flows/llmflows/Functions.java index 9e2accdc3..f11028104 100644 --- a/core/src/main/java/com/google/adk/flows/llmflows/Functions.java +++ b/core/src/main/java/com/google/adk/flows/llmflows/Functions.java @@ -263,7 +263,7 @@ private static Observable 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 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> getFunctionCallMapper( diff --git a/core/src/main/java/com/google/adk/runner/Runner.java b/core/src/main/java/com/google/adk/runner/Runner.java index 9a662516c..0ae5a1e2f 100644 --- a/core/src/main/java/com/google/adk/runner/Runner.java +++ b/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 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 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()}. + * + *

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 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); } diff --git a/core/src/main/java/com/google/adk/tools/AgentTool.java b/core/src/main/java/com/google/adk/tools/AgentTool.java index ab821a221..86a5eb44f 100644 --- a/core/src/main/java/com/google/adk/tools/AgentTool.java +++ b/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> runAsync(Map 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 = diff --git a/core/src/test/java/com/google/adk/agents/InvocationContextTest.java b/core/src/test/java/com/google/adk/agents/InvocationContextTest.java index 546e84c95..adcf49ab7 100644 --- a/core/src/test/java/com/google/adk/agents/InvocationContextTest.java +++ b/core/src/test/java/com/google/adk/agents/InvocationContextTest.java @@ -37,6 +37,8 @@ import com.google.genai.types.FunctionCall; import com.google.genai.types.FunctionResponse; import com.google.genai.types.Part; +import io.reactivex.rxjava3.core.Scheduler; +import io.reactivex.rxjava3.schedulers.Schedulers; import java.util.HashMap; import java.util.Map; import java.util.concurrent.ConcurrentHashMap; @@ -305,6 +307,73 @@ public void testNewInvocationContextId() { .matches("^[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}$"); } + @Test + public void scheduler_defaultsToIoScheduler() { + InvocationContext context = + InvocationContext.builder() + .sessionService(mockSessionService) + .artifactService(mockArtifactService) + .agent(mockAgent) + .session(session) + .build(); + + assertThat(context.scheduler()).isSameInstanceAs(Schedulers.io()); + } + + @Test + public void scheduler_customIsCarriedAndCopiedByToBuilder() { + Scheduler scheduler = Schedulers.trampoline(); + InvocationContext context = + InvocationContext.builder() + .sessionService(mockSessionService) + .artifactService(mockArtifactService) + .agent(mockAgent) + .session(session) + .scheduler(scheduler) + .build(); + + assertThat(context.scheduler()).isSameInstanceAs(scheduler); + assertThat(context.toBuilder().build().scheduler()).isSameInstanceAs(scheduler); + } + + @Test + public void scheduler_nullSelectsTheDefault() { + InvocationContext context = + InvocationContext.builder() + .sessionService(mockSessionService) + .artifactService(mockArtifactService) + .agent(mockAgent) + .session(session) + .scheduler(Schedulers.trampoline()) + .scheduler(null) + .build(); + + assertThat(context.scheduler()).isSameInstanceAs(Schedulers.io()); + } + + @Test + public void equals_comparesTheConfiguredScheduler() { + InvocationContext unset = + InvocationContext.builder() + .sessionService(mockSessionService) + .artifactService(mockArtifactService) + .agent(mockAgent) + .session(session) + .invocationId("same-id") + .build(); + InvocationContext trampoline = unset.toBuilder().scheduler(Schedulers.trampoline()).build(); + InvocationContext io = unset.toBuilder().scheduler(Schedulers.io()).build(); + + assertThat(unset.toBuilder().build()).isEqualTo(unset); + assertThat(unset.toBuilder().build().hashCode()).isEqualTo(unset.hashCode()); + assertThat(trampoline.toBuilder().build()).isEqualTo(trampoline); + assertThat(trampoline.toBuilder().build().hashCode()).isEqualTo(trampoline.hashCode()); + assertThat(trampoline).isNotEqualTo(io); + assertThat(trampoline).isNotEqualTo(unset); + // The configured field is compared, so an unset scheduler differs from an explicit io(). + assertThat(io).isNotEqualTo(unset); + } + @Test public void testEquals_sameObject() { InvocationContext context = diff --git a/core/src/test/java/com/google/adk/agents/ParallelAgentTest.java b/core/src/test/java/com/google/adk/agents/ParallelAgentTest.java index 97d107aa9..984fa6e14 100644 --- a/core/src/test/java/com/google/adk/agents/ParallelAgentTest.java +++ b/core/src/test/java/com/google/adk/agents/ParallelAgentTest.java @@ -17,6 +17,7 @@ package com.google.adk.agents; import static com.google.adk.testing.TestUtils.createEscalateEvent; +import static com.google.adk.testing.TestUtils.createEvent; import static com.google.adk.testing.TestUtils.createFunctionCallLlmResponse; import static com.google.adk.testing.TestUtils.createInvocationContext; import static com.google.adk.testing.TestUtils.createResumableInvocationContext; @@ -25,10 +26,12 @@ import static com.google.adk.testing.TestUtils.createTestLlm; import static com.google.adk.testing.TestUtils.createTextLlmResponse; import static com.google.adk.testing.TestUtils.simplifyEvents; +import static com.google.adk.testing.TestUtils.simplifyResumableEvents; import static com.google.common.truth.Truth.assertThat; import static java.util.concurrent.TimeUnit.MILLISECONDS; import com.google.adk.events.Event; +import com.google.adk.testing.RecordingScheduler; import com.google.adk.tools.FunctionTool; import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableMap; @@ -86,6 +89,89 @@ protected Flowable runLiveImpl(InvocationContext invocationContext) { } } + @Test + public void runAsync_withoutAgentScheduler_runsSubAgentsOnInvocationScheduler() { + RecordingScheduler recordingScheduler = new RecordingScheduler(Schedulers.trampoline()); + ParallelAgent parallelAgent = + ParallelAgent.builder() + .name("parallel") + .subAgents( + createSubAgent("a", createEvent("ea")), createSubAgent("b", createEvent("eb"))) + .build(); + InvocationContext invocationContext = + createInvocationContext(parallelAgent).toBuilder().scheduler(recordingScheduler).build(); + + List events = parallelAgent.runAsync(invocationContext).toList().blockingGet(); + + assertThat(events).hasSize(2); + // One subscription per sub-agent went through the invocation's scheduler. + assertThat(recordingScheduler.workersCreated()).isEqualTo(2); + } + + @Test + public void runAsync_resumable_withoutAgentScheduler_runsSubAgentsOnInvocationScheduler() { + RecordingScheduler recordingScheduler = new RecordingScheduler(Schedulers.trampoline()); + ParallelAgent parallelAgent = + ParallelAgent.builder() + .name("parallel") + .subAgents( + createSubAgent("a", createEvent("ea")), createSubAgent("b", createEvent("eb"))) + .build(); + InvocationContext invocationContext = + createResumableInvocationContext(parallelAgent).toBuilder() + .scheduler(recordingScheduler) + .build(); + + List events = parallelAgent.runAsync(invocationContext).toList().blockingGet(); + + assertThat(simplifyResumableEvents(events)) + .containsExactly( + "parallel: agent_state={}", + "author: content for event ea", + "author: content for event eb") + .inOrder(); + assertThat(recordingScheduler.workersCreated()).isEqualTo(2); + } + + @Test + public void builder_scheduler_null_inheritsInvocationScheduler() { + RecordingScheduler recordingScheduler = new RecordingScheduler(Schedulers.trampoline()); + ParallelAgent parallelAgent = + ParallelAgent.builder() + .name("parallel") + .scheduler(Schedulers.trampoline()) + .scheduler(null) + .subAgents( + createSubAgent("a", createEvent("ea")), createSubAgent("b", createEvent("eb"))) + .build(); + InvocationContext invocationContext = + createInvocationContext(parallelAgent).toBuilder().scheduler(recordingScheduler).build(); + + List events = parallelAgent.runAsync(invocationContext).toList().blockingGet(); + + assertThat(events).hasSize(2); + assertThat(recordingScheduler.workersCreated()).isEqualTo(2); + } + + @Test + public void runAsync_withAgentScheduler_ignoresInvocationScheduler() { + RecordingScheduler recordingScheduler = new RecordingScheduler(Schedulers.trampoline()); + ParallelAgent parallelAgent = + ParallelAgent.builder() + .name("parallel") + .scheduler(Schedulers.trampoline()) + .subAgents( + createSubAgent("a", createEvent("ea")), createSubAgent("b", createEvent("eb"))) + .build(); + InvocationContext invocationContext = + createInvocationContext(parallelAgent).toBuilder().scheduler(recordingScheduler).build(); + + List events = parallelAgent.runAsync(invocationContext).toList().blockingGet(); + + assertThat(events).hasSize(2); + assertThat(recordingScheduler.workersCreated()).isEqualTo(0); + } + @Test public void runAsync_subAgentsExecuteInParallel_eventsOrderedByCompletion() { String agent1Name = "test_agent_1_delayed"; diff --git a/core/src/test/java/com/google/adk/flows/llmflows/FunctionsTest.java b/core/src/test/java/com/google/adk/flows/llmflows/FunctionsTest.java index c4aec95c0..0d3d99fbc 100644 --- a/core/src/test/java/com/google/adk/flows/llmflows/FunctionsTest.java +++ b/core/src/test/java/com/google/adk/flows/llmflows/FunctionsTest.java @@ -19,12 +19,18 @@ import static com.google.adk.testing.TestUtils.createEvent; import static com.google.adk.testing.TestUtils.createInvocationContext; import static com.google.adk.testing.TestUtils.createRootAgent; +import static com.google.adk.testing.TestUtils.createTestAgentBuilder; +import static com.google.adk.testing.TestUtils.createTestLlm; +import static com.google.adk.testing.TestUtils.createTextLlmResponse; import static com.google.common.truth.Truth.assertThat; +import com.google.adk.agents.BaseAgent; import com.google.adk.agents.InvocationContext; +import com.google.adk.agents.LlmAgent; import com.google.adk.agents.RunConfig; import com.google.adk.agents.RunConfig.ToolExecutionMode; import com.google.adk.events.Event; +import com.google.adk.testing.RecordingScheduler; import com.google.adk.testing.TestUtils; import com.google.adk.tools.BaseTool; import com.google.adk.tools.ToolContext; @@ -37,6 +43,7 @@ import com.google.genai.types.FunctionResponse; import com.google.genai.types.Part; import io.reactivex.rxjava3.core.Single; +import io.reactivex.rxjava3.schedulers.Schedulers; import java.util.ArrayList; import java.util.LinkedHashMap; import java.util.List; @@ -499,6 +506,85 @@ public void handleFunctionCalls_parallelSubscribe_blockingTools_runConcurrently_ runParallelSubscribeBlockingToolsTest(/* toolCount= */ 5); } + @Test + public void handleFunctionCalls_parallelSubscribe_usesInvocationScheduler() { + // Each tool was subscribed on the invocation's scheduler. + assertThat(workersCreatedForParallelSubscribe(createRootAgent())).isEqualTo(2); + } + + @Test + public void + handleFunctionCalls_parallelSubscribe_llmAgentWithoutExecutor_usesInvocationScheduler() { + LlmAgent agentWithoutExecutor = + createTestAgentBuilder(createTestLlm(createTextLlmResponse("unused"))).build(); + + assertThat(workersCreatedForParallelSubscribe(agentWithoutExecutor)).isEqualTo(2); + } + + @Test + public void handleFunctionCalls_parallelSubscribe_agentExecutorTakesPrecedenceOverScheduler() { + LlmAgent agentWithExecutor = + createTestAgentBuilder(createTestLlm(createTextLlmResponse("unused"))) + .executor(Runnable::run) + .build(); + + assertThat(workersCreatedForParallelSubscribe(agentWithExecutor)).isEqualTo(0); + } + + /** + * Runs two tool calls for {@code agent} in PARALLEL_SUBSCRIBE mode on an invocation whose + * scheduler records its use, asserts that both responses arrived, and returns how many workers + * that scheduler handed out. + */ + private static int workersCreatedForParallelSubscribe(BaseAgent agent) { + RecordingScheduler recordingScheduler = new RecordingScheduler(Schedulers.trampoline()); + InvocationContext invocationContext = + createInvocationContext( + agent, + RunConfig.builder().toolExecutionMode(ToolExecutionMode.PARALLEL_SUBSCRIBE).build()) + .toBuilder() + .scheduler(recordingScheduler) + .build(); + + Event functionResponseEvent = + Functions.handleFunctionCalls( + invocationContext, + twoToolCallsEvent(), + ImmutableMap.of( + "tool_1", + new SleepingTool("tool_1", 0L), + "tool_2", + new SleepingTool("tool_2", 0L))) + .blockingGet(); + + assertThat(functionResponseEvent).isNotNull(); + assertThat(functionResponseEvent.functionResponses()).hasSize(2); + return recordingScheduler.workersCreated(); + } + + private static Event twoToolCallsEvent() { + return createEvent("event").toBuilder() + .content( + Content.fromParts( + Part.builder() + .functionCall( + FunctionCall.builder() + .id("call_1") + .name("tool_1") + .args(ImmutableMap.of()) + .build()) + .build(), + Part.builder() + .functionCall( + FunctionCall.builder() + .id("call_2") + .name("tool_2") + .args(ImmutableMap.of()) + .build()) + .build())) + .build(); + } + /** Single-tool case bypasses the parallel scheduler path; must still return the correct event. */ @Test public void handleFunctionCalls_parallelSubscribe_blockingTool_singleTool() { diff --git a/core/src/test/java/com/google/adk/runner/RunnerTest.java b/core/src/test/java/com/google/adk/runner/RunnerTest.java index 861a797f1..606de85bb 100644 --- a/core/src/test/java/com/google/adk/runner/RunnerTest.java +++ b/core/src/test/java/com/google/adk/runner/RunnerTest.java @@ -16,8 +16,10 @@ package com.google.adk.runner; +import static com.google.adk.testing.TestUtils.createEvent; import static com.google.adk.testing.TestUtils.createFunctionCallLlmResponse; import static com.google.adk.testing.TestUtils.createLlmResponse; +import static com.google.adk.testing.TestUtils.createSubAgent; import static com.google.adk.testing.TestUtils.createTestAgent; import static com.google.adk.testing.TestUtils.createTestAgentBuilder; import static com.google.adk.testing.TestUtils.createTestLlm; @@ -45,6 +47,7 @@ import com.google.adk.agents.InvocationContext; import com.google.adk.agents.LiveRequestQueue; import com.google.adk.agents.LlmAgent; +import com.google.adk.agents.ParallelAgent; import com.google.adk.agents.RunConfig; import com.google.adk.agents.SequentialAgent; import com.google.adk.apps.App; @@ -65,6 +68,7 @@ import com.google.adk.sessions.SessionKey; import com.google.adk.summarizer.EventsCompactionConfig; import com.google.adk.telemetry.Tracing; +import com.google.adk.testing.RecordingScheduler; import com.google.adk.testing.TestLlm; import com.google.adk.testing.TestUtils; import com.google.adk.testing.TestUtils.EchoTool; @@ -91,6 +95,7 @@ import io.reactivex.rxjava3.core.Flowable; import io.reactivex.rxjava3.core.Maybe; import io.reactivex.rxjava3.core.Single; +import io.reactivex.rxjava3.schedulers.Schedulers; import io.reactivex.rxjava3.subjects.PublishSubject; import io.reactivex.rxjava3.subscribers.TestSubscriber; import java.time.Instant; @@ -1841,6 +1846,116 @@ public void runAsync_ensureEventsAreAppendedInOrder() throws Exception { .inOrder(); } + @Test + public void runAsync_parallelAgent_usesRunnerScheduler() { + RecordingScheduler recordingScheduler = new RecordingScheduler(Schedulers.trampoline()); + Runner schedulerRunner = + Runner.builder() + .app(App.builder().name("test").rootAgent(parallelAgentWithTwoSubAgents()).build()) + .scheduler(recordingScheduler) + .build(); + Session schedulerSession = + schedulerRunner.sessionService().createSession("test", "user").blockingGet(); + + List events = + schedulerRunner + .runAsync("user", schedulerSession.id(), createContent("hi")) + .toList() + .blockingGet(); + + assertThat(events).hasSize(2); + assertThat(recordingScheduler.workersCreated()).isEqualTo(2); + } + + @Test + public void builder_scheduler_null_usesDefaultScheduler() { + RecordingScheduler recordingScheduler = new RecordingScheduler(Schedulers.trampoline()); + Runner schedulerRunner = + Runner.builder() + .app(App.builder().name("test").rootAgent(parallelAgentWithTwoSubAgents()).build()) + .scheduler(recordingScheduler) + .scheduler(null) + .build(); + Session schedulerSession = + schedulerRunner.sessionService().createSession("test", "user").blockingGet(); + + List events = + schedulerRunner + .runAsync("user", schedulerSession.id(), createContent("hi")) + .toList() + .blockingGet(); + + assertThat(events).hasSize(2); + assertThat(recordingScheduler.workersCreated()).isEqualTo(0); + } + + @Test + public void runLive_usesRunnerScheduler() throws Exception { + RecordingScheduler recordingScheduler = new RecordingScheduler(Schedulers.trampoline()); + Runner schedulerRunner = + Runner.builder() + .app(App.builder().name("test").rootAgent(agent).build()) + .scheduler(recordingScheduler) + .build(); + Session liveSession = + schedulerRunner.sessionService().createSession("test", "user").blockingGet(); + LiveRequestQueue liveRequestQueue = new LiveRequestQueue(); + TestSubscriber testSubscriber = + schedulerRunner.runLive(liveSession, liveRequestQueue, RunConfig.builder().build()).test(); + + liveRequestQueue.content(createContent("from user")); + liveRequestQueue.close(); + + testSubscriber.await(); + testSubscriber.assertComplete(); + assertThat(simplifyEvents(testSubscriber.values())).containsExactly("test agent: from llm"); + assertThat(liveRequestTexts(testLlm)).contains("from user"); + // The hop that starts the live send loop went through the runner's scheduler. + assertThat(recordingScheduler.workersCreated()).isEqualTo(1); + } + + @Test + public void runLive_agentExecutorTakesPrecedenceOverRunnerScheduler() throws Exception { + RecordingScheduler recordingScheduler = new RecordingScheduler(Schedulers.trampoline()); + TestLlm liveTestLlm = createTestLlm(createLlmResponse(createContent("from llm"))); + LlmAgent agentWithExecutor = + createTestAgentBuilder(liveTestLlm).executor(Runnable::run).build(); + Runner schedulerRunner = + Runner.builder() + .app(App.builder().name("test").rootAgent(agentWithExecutor).build()) + .scheduler(recordingScheduler) + .build(); + Session liveSession = + schedulerRunner.sessionService().createSession("test", "user").blockingGet(); + LiveRequestQueue liveRequestQueue = new LiveRequestQueue(); + TestSubscriber testSubscriber = + schedulerRunner.runLive(liveSession, liveRequestQueue, RunConfig.builder().build()).test(); + + liveRequestQueue.content(createContent("from user")); + liveRequestQueue.close(); + + testSubscriber.await(); + testSubscriber.assertComplete(); + assertThat(simplifyEvents(testSubscriber.values())).containsExactly("test agent: from llm"); + assertThat(liveRequestTexts(liveTestLlm)).contains("from user"); + assertThat(recordingScheduler.workersCreated()).isEqualTo(0); + } + + private static ParallelAgent parallelAgentWithTwoSubAgents() { + return ParallelAgent.builder() + .name("parallel") + .subAgents(createSubAgent("a", createEvent("ea")), createSubAgent("b", createEvent("eb"))) + .build(); + } + + /** Returns the text of each content request that the live connection of {@code llm} received. */ + private static ImmutableList liveRequestTexts(TestLlm llm) { + return llm.getLiveRequestHistory().stream() + .map(request -> request.content().map(Content::text)) + .flatMap(Optional::stream) + .collect(toImmutableList()); + } + private Content createContent(String text) { return Content.builder().parts(Part.builder().text(text).build()).build(); } diff --git a/core/src/test/java/com/google/adk/testing/RecordingScheduler.java b/core/src/test/java/com/google/adk/testing/RecordingScheduler.java new file mode 100644 index 000000000..3344d376d --- /dev/null +++ b/core/src/test/java/com/google/adk/testing/RecordingScheduler.java @@ -0,0 +1,59 @@ +/* + * Copyright 2026 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 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.testing; + +import io.reactivex.rxjava3.core.Scheduler; +import io.reactivex.rxjava3.schedulers.Schedulers; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicInteger; + +/** + * A {@link Scheduler} for tests that delegates to another scheduler and counts how many workers + * were requested from it, so a test can assert that a code path ran its work through it. + */ +public final class RecordingScheduler extends Scheduler { + + private final Scheduler delegate; + private final AtomicInteger workersCreated = new AtomicInteger(); + + /** + * Records on top of {@code delegate}; use {@link Schedulers#trampoline()} to stay + * single-threaded. + */ + public RecordingScheduler(Scheduler delegate) { + this.delegate = delegate; + } + + // Every operator ADK uses on the seam ends up here: Flowable.subscribeOn (ParallelAgent) calls + // createWorker() directly, while Observable.subscribeOn (tool fan-out) and Completable.observeOn + // (live loop) go through scheduleDirect(), whose base implementation calls createWorker(). + @Override + public Worker createWorker() { + workersCreated.incrementAndGet(); + return delegate.createWorker(); + } + + @Override + public long now(TimeUnit unit) { + return delegate.now(unit); + } + + /** Returns how many times {@link #createWorker()} was called. */ + public int workersCreated() { + return workersCreated.get(); + } +} diff --git a/core/src/test/java/com/google/adk/tools/AgentToolTest.java b/core/src/test/java/com/google/adk/tools/AgentToolTest.java index 6a8b82f8d..113236b63 100644 --- a/core/src/test/java/com/google/adk/tools/AgentToolTest.java +++ b/core/src/test/java/com/google/adk/tools/AgentToolTest.java @@ -16,6 +16,7 @@ package com.google.adk.tools; +import static com.google.adk.testing.TestUtils.createEvent; import static com.google.adk.testing.TestUtils.createInvocationContext; import static com.google.adk.testing.TestUtils.createSubAgent; import static com.google.adk.testing.TestUtils.createTestAgentBuilder; @@ -28,6 +29,7 @@ import com.google.adk.agents.ConfigAgentUtils.ConfigurationException; import com.google.adk.agents.InvocationContext; import com.google.adk.agents.LlmAgent; +import com.google.adk.agents.ParallelAgent; import com.google.adk.agents.RunConfig; import com.google.adk.agents.RunConfig.StreamingMode; import com.google.adk.agents.RunConfig.ToolExecutionMode; @@ -37,6 +39,7 @@ import com.google.adk.plugins.PluginManager; import com.google.adk.sessions.InMemorySessionService; import com.google.adk.sessions.Session; +import com.google.adk.testing.RecordingScheduler; import com.google.adk.testing.TestBaseAgent; import com.google.adk.testing.TestLlm; import com.google.adk.utils.ComponentRegistry; @@ -48,6 +51,7 @@ import com.google.genai.types.Schema; import io.reactivex.rxjava3.core.Flowable; import io.reactivex.rxjava3.core.Maybe; +import io.reactivex.rxjava3.schedulers.Schedulers; import java.util.Map; import java.util.Optional; import java.util.concurrent.atomic.AtomicBoolean; @@ -914,6 +918,32 @@ public void call_withStreamingRunConfig_runsAgentWithoutStreaming() throws Excep } } + @Test + public void call_nestedRunInheritsCallerScheduler() throws Exception { + RecordingScheduler recordingScheduler = new RecordingScheduler(Schedulers.trampoline()); + ParallelAgent parallelAgent = + ParallelAgent.builder() + .name("parallel") + .description("runs two sub-agents in parallel") + .subAgents( + createSubAgent("a", createEvent("ea")), createSubAgent("b", createEvent("eb"))) + .build(); + AgentTool agentTool = AgentTool.create(parallelAgent); + ToolContext toolContext = + ToolContext.builder( + createInvocationContext(parallelAgent).toBuilder() + .scheduler(recordingScheduler) + .build()) + .build(); + + Map result = + agentTool.runAsync(ImmutableMap.of("request", "magic"), toolContext).blockingGet(); + + assertThat(result).containsExactly("result", "content for event eb"); + // Each sub-agent of the nested run was subscribed on the caller's scheduler. + assertThat(recordingScheduler.workersCreated()).isEqualTo(2); + } + private ToolContext createToolContext(BaseAgent agent) { Session session = sessionService.createSession("test-app", "test-user", null, "test-session").blockingGet();