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();