[DABH]

Link to Issue or Description of Change

Problem:
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 against rate-limited backends, to dispatch it onto an existing worker pool, or to keep every ADK task on a dispatcher it controls, which a durable-execution host must do to stay deterministic.

adk-go exposes exactly this control point as platform.WithTaskRunner / platform.RunTasks (google/adk-go#1050), next to its time/UUID provider seams. This is the Java analog, and the companion to #1255 (time and UUID providers); the two are independent and either can land first. On alignment with adk-python: its asyncio model has no equivalent seam (parallel tool calls run as tasks on the caller's event loop, and the optional tool thread pool is opt-in via RunConfig), so adk-go's platform.WithTaskRunner is the closest analog to mirror.

Solution:
Carry an RxJava Scheduler on InvocationContext, configured once on the Runner:

• Runner.builder().scheduler(@Nullable Scheduler) — the scheduler for the runner's invocations; null (the default) selects Schedulers.io(). Every InvocationContext the runner creates (new, resumed, and live invocations) carries it, and toBuilder() copies it to sub-agent and branch contexts.
• InvocationContext.scheduler() — returns the configured scheduler, or Schedulers.io() when none is set, so the default behavior is unchanged. InvocationContext.Builder.scheduler(@Nullable Scheduler) is the plumbing.
• Consumers, in precedence order "agent-level setting, then invocation, then IO":
  • ParallelAgent uses the invocation's scheduler unless its own Builder.scheduler(...) was set. Its builder default moves from an eager Schedulers.io() to "unset, inherit", and scheduler(null) means unset, so existing callers keep working; the effective default is the same. (One benign change: on main, scheduler(null) stored the null and a run with sub-agents then failed with an NPE from subscribeOn(null); now such an agent inherits the invocation's scheduler.)
  • The "agent executor, then invocation scheduler" rule lives in one package-private helper, Functions.resolveWorkerScheduler, used for the PARALLEL_SUBSCRIBE tool fan-out and for the hop that starts the live send loop in BaseLlmFlow (both previously fell back to Schedulers.io()).
  • AgentTool builds the runner it nests with Runner.builder() (same in-memory artifact, session, and memory services, plugins, app name, and null compaction/cache/resumability configs as the InMemoryRunner it replaces, and the caller's RunConfig handling is unchanged) and passes the caller's invocation scheduler, so nested fan-out inherits it.
• InvocationContext.equals()/hashCode() include the configured scheduler field as-is: a context always equals its toBuilder() copy (no RxJavaPlugins handler runs during a comparison), and an unset scheduler is not equal to an explicit Schedulers.io().

Why an RxJava Scheduler rather than a Go-style task-runner function: every Java hop site already takes a Scheduler, and a Scheduler can also express the delayed scheduling RxJava's time-based operators need, which a plain Executor cannot (Schedulers.from(Executor) needs a helper thread to time delays). A host that wants ADK's own hops removed installs Schedulers.trampoline(), which runs them inline on the thread that would otherwise hand off (sub-agents and tools are then subscribed sequentially, though anything they start asynchronously can still overlap); a host with its own dispatcher implements Scheduler once. A single-threaded scheduler must not be blocked from its own thread (for example by calling blockingGet() on a run from that thread, or from a tool), or the queued fan-out tasks never run.

Not changed here: hops outside the Runner → agent → flow path, for example the MCP retry timers (Flowable.timer / delay on the computation scheduler, tool-side), the Vertex session client's retry timer, the model clients' HTTP executors, the agent-analytics plugin's own executors, and the dev server's own Schedulers.io() hops. Also observed while testing, pre-existing and left alone: ParallelAgent does not re-activate the OpenTelemetry context inside its subscribeOn hop (as Functions does for tool calls), so sub-agent spans are parented to the root under any non-inline scheduler, including today's default; a Tracing.withContext(parentContext) inside the hop fixes it and I'm happy to send that separately.

Testing Plan

Unit Tests:

• I have added or updated unit tests for my change.
• All unit tests pass locally.

New tests (all under core):

• InvocationContextTest — scheduler() defaults to Schedulers.io(), a custom scheduler is carried and copied by toBuilder(), null selects the default, and equality compares the configured scheduler (a context equals its toBuilder() copy; unset differs from an explicit Schedulers.io()).
• ParallelAgentTest — without an agent-level scheduler the sub-agents run on the invocation's scheduler (non-resumable and resumable paths; the resumable one also pins the exact event sequence under the trampoline); an agent-level scheduler takes precedence; scheduler(null) inherits.
• FunctionsTest — PARALLEL_SUBSCRIBE tool calls run on the invocation's scheduler (for an LlmAgent without an executor and for a non-LLM agent); an LlmAgent.executor() takes precedence. Shared body in one helper that returns the worker count.
• RunnerTest — Runner.builder().scheduler(...) reaches a ParallelAgent run; scheduler(null) falls back to the default; the runLive send loop runs on the runner's scheduler (trampoline-backed, and the test model is asserted to have received the request) and an LlmAgent.executor() takes precedence on that hop.
• AgentToolTest — the run an AgentTool nests inherits the caller's scheduler.
• testing/RecordingScheduler — a test Scheduler that delegates to another scheduler (the trampoline in these tests, so they are single-threaded and assert exact worker counts) and counts the workers ADK requested from it.

Local results (./mvnw -Prelease clean package, OpenJDK 17 (Homebrew 17.0.20.1), all modules): BUILD SUCCESS; core: Tests run: 2063, Failures: 0, Errors: 0, Skipped: 24 (the skips are pre-existing); every other module passes with 0 failures.
git diff --exit-code after the build: clean (formatter made no changes).

Manual End-to-End (E2E) Tests:

N/A; this is a low-level concurrency seam with no runner UI surface. The default path (Schedulers.io() when no scheduler is installed) is unchanged, and the installed-scheduler path is exercised by the unit tests above.

Checklist

• I have read the CONTRIBUTING.md document.
• My pull request contains a single commit.
• I have performed a self-review of my own code.
• I have commented my code, particularly in hard-to-understand areas.
• I have added tests that prove my fix is effective or that my feature works.
• New and existing unit tests pass locally with my changes.
• I have manually tested my changes end-to-end. (N/A — low-level seam; see above.)
• Any dependent changes have been merged and published in downstream modules. (N/A.)

Additional context

Motivated by parity with adk-go (platform.WithTaskRunner, google/adk-go#1050) and by the durable-execution use case discussed on #1255, where a host must be able to keep ADK's fan-out on threads it controls. #1255 and this PR touch the same spots in InvocationContext and Runner (a new field, builder setter, and constructor parameter each) and the same runner-construction block in AgentTool, so whichever lands second needs a small, mechanical rebase; I'll do it promptly.

[DABH]

Thanks @damianmomotgoogle for the review! All six threads are addressed in 1b89d23 (single commit, rebased onto current main); replies inline. In short: AgentTool now builds its nested runner with Runner.builder() and passes the caller's scheduler (with a test), scheduler(null) is accepted again on both builders (null = inherit / default), equals/hashCode compare the configured field, the executor-then-invocation rule lives in one package-private helper used by the tool fan-out and the live send loop, the live tests pin the request reaching the model on a trampoline-backed recorder, and the FunctionsTest cases use toolExecutionMode(...) through a shared helper.

Full ./mvnw -Prelease clean package is green locally on JDK 17 (core: 2063 tests, 0 failures); the PR description is updated. Thanks again!