diff --git a/core/src/main/java/com/google/adk/agents/BaseAgent.java b/core/src/main/java/com/google/adk/agents/BaseAgent.java index da43622b1..66b1e3e9c 100644 --- a/core/src/main/java/com/google/adk/agents/BaseAgent.java +++ b/core/src/main/java/com/google/adk/agents/BaseAgent.java @@ -21,12 +21,14 @@ import static com.google.common.base.Strings.nullToEmpty; import static com.google.common.collect.ImmutableList.toImmutableList; import static java.lang.String.format; +import static java.util.concurrent.TimeUnit.MILLISECONDS; import com.google.adk.agents.Callbacks.AfterAgentCallback; import com.google.adk.agents.Callbacks.BeforeAgentCallback; import com.google.adk.annotations.Experimental; import com.google.adk.events.Event; import com.google.adk.events.EventActions; +import com.google.adk.flows.llmflows.PersistBarrier; import com.google.adk.plugins.Plugin; import com.google.adk.telemetry.Instrumentation; import com.google.adk.telemetry.Instrumentation.AgentInvocation; @@ -34,31 +36,44 @@ import com.google.adk.utils.AgentEnums.AgentOrigin; import com.google.adk.workflow.BaseNode; import com.google.adk.workflow.NodeExecutionState; +import com.google.common.annotations.VisibleForTesting; import com.google.common.collect.ImmutableList; import com.google.errorprone.annotations.CanIgnoreReturnValue; import com.google.errorprone.annotations.DoNotCall; import com.google.genai.types.Content; import io.opentelemetry.context.Context; import io.reactivex.rxjava3.core.Completable; +import io.reactivex.rxjava3.core.CompletableObserver; import io.reactivex.rxjava3.core.Flowable; import io.reactivex.rxjava3.core.Maybe; +import io.reactivex.rxjava3.core.Scheduler; +import java.time.Duration; import java.util.ArrayList; import java.util.HashSet; import java.util.List; import java.util.Map; import java.util.Optional; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicReference; import java.util.function.Function; import java.util.regex.Pattern; import java.util.stream.Stream; import org.jspecify.annotations.Nullable; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; /** Base class for all agents. */ public abstract class BaseAgent extends BaseNode { + private static final Logger logger = LoggerFactory.getLogger(BaseAgent.class); + // Pattern for valid agent names. private static final String IDENTIFIER_REGEX = "^_?[a-zA-Z0-9]*([. _-][a-zA-Z0-9]+)*$"; private static final Pattern IDENTIFIER_PATTERN = Pattern.compile(IDENTIFIER_REGEX); + /** How long a workflow agent waits for the {@code Runner} to persist a sub-agent's events. */ + @VisibleForTesting static volatile Duration persistWaitTimeout = Duration.ofSeconds(30); + /** * The parent agent in the agent tree. Note that one agent cannot be added to two different * parents' sub-agents lists. @@ -496,6 +511,84 @@ private Event checkpointEvent(InvocationContext context, EventActions actions) { .build(); } + /** + * Returns a stream that re-emits {@code events} and completes once the {@code Runner} has + * persisted the last of them that has an id and is not partial, so the next sub-agent reads their + * state. It waits only when a {@code Runner} drives the invocation; see {@link PersistBarrier}. + * It waits at most {@link #persistWaitTimeout}, then logs a warning and continues, since the + * {@code Runner} may be slow or an agent above may never pass the event on. + */ + static Flowable completeAfterPersisted( + InvocationContext context, String waitingPoint, Flowable events) { + return Flowable.defer( + () -> { + // An append can complete on its own thread; the next sub-agent must stay in this trace. + Context spanContext = Context.current(); + AtomicReference lastTracked = new AtomicReference<>(); + return events + .doOnNext( + event -> { + if (!event.partial().orElse(false) && event.id() != null) { + lastTracked.set(event); + } + }) + .concatWith( + Completable.defer( + () -> { + Event last = lastTracked.get(); + if (last == null) { + return Completable.complete(); + } + Duration timeout = persistWaitTimeout; + return resumeOnSchedulerIfLate( + PersistBarrier.awaitPersisted(context, ImmutableList.of(last)) + .timeout( + timeout.toMillis(), + MILLISECONDS, + Completable.fromAction( + () -> + logger.warn( + "{} ({}): event {} was not persisted within {}" + + " ms, so the next sub-agent starts" + + " anyway. The Runner may be slow, or an" + + " agent above may drop, re-id or hold" + + " back events.", + context.agent().name(), + waitingPoint, + last.id(), + timeout.toMillis()))) + // A failed append stops the timer; the Runner reports it. + .onErrorResumeWith(Completable.never()), + context.scheduler()); + }) + .compose(Tracing.withContext(spanContext))); + }); + } + + /** + * Returns a completable that completes with {@code barrier}: inline if the barrier completes + * while it is being subscribed, otherwise on {@code scheduler}. A late completion comes on the + * thread that persisted the event, which must not run the next sub-agent's model call. + */ + private static Completable resumeOnSchedulerIfLate(Completable barrier, Scheduler scheduler) { + return new Completable() { + @Override + protected void subscribeActual(CompletableObserver observer) { + AtomicBoolean subscribing = new AtomicBoolean(true); + Thread subscriber = Thread.currentThread(); + barrier + .andThen( + Completable.defer( + () -> + subscribing.get() && Thread.currentThread() == subscriber + ? Completable.complete() + : Completable.complete().observeOn(scheduler))) + .subscribe(observer); + subscribing.set(false); + } + }; + } + /** * Runs this agent as a workflow node, on the node's branch so the nodes after it see its events. * It runs the callbacks and {@link #runAsyncImpl} that {@link #runAsync} runs, but does not call @@ -520,7 +613,10 @@ public Flowable runNode( } /** - * Agent-specific asynchronous logic. + * Agent-specific asynchronous logic. An implementation that runs sub-agents itself should emit + * their events as they arrive, with ids unchanged: multi-step {@link LlmAgent}s wait until the + * {@code Runner} has persisted those events, and workflow agents wait up to 30 seconds, then log + * a warning and continue. * * @param invocationContext Current invocation context. * @return stream of agent-generated events. 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 338eaa00f..f7fbe7b68 100644 --- a/core/src/main/java/com/google/adk/agents/InvocationContext.java +++ b/core/src/main/java/com/google/adk/agents/InvocationContext.java @@ -315,9 +315,10 @@ public String 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, workflow node retries and attempts with a timeout, - * 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()}. + * the hop that starts the live send loop, and the next sub-agent of a {@code SequentialAgent} or + * {@code LoopAgent} when the previous one's events are persisted late. 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(); diff --git a/core/src/main/java/com/google/adk/agents/LoopAgent.java b/core/src/main/java/com/google/adk/agents/LoopAgent.java index 9875b7aee..80f5f6861 100644 --- a/core/src/main/java/com/google/adk/agents/LoopAgent.java +++ b/core/src/main/java/com/google/adk/agents/LoopAgent.java @@ -161,7 +161,10 @@ protected Flowable runAsyncImpl(InvocationContext invocationContext) { if (!invocationContext.isResumable()) { return Flowable.fromIterable(subAgents) - .concatMap(subAgent -> subAgent.runAsync(invocationContext)) + .concatMap( + subAgent -> + completeAfterPersisted( + invocationContext, "loop step", subAgent.runAsync(invocationContext))) .repeat(maxIterations != null ? maxIterations : Integer.MAX_VALUE) .takeUntil(LoopAgent::hasEscalateAction); } @@ -193,14 +196,17 @@ private Flowable runAsyncLegacyResumption( subAgent -> paused.get() ? Flowable.empty() - : subAgent - .runAsync(invocationContext) - .doOnNext( - event -> { - if (WorkflowAgentResumption.hasPendingLongRunningCall(event)) { - paused.set(true); - } - })) + : completeAfterPersisted( + invocationContext, + "loop step", + subAgent + .runAsync(invocationContext) + .doOnNext( + event -> { + if (WorkflowAgentResumption.hasPendingLongRunningCall(event)) { + paused.set(true); + } + }))) .repeatUntil( () -> paused.get() @@ -302,7 +308,8 @@ private Flowable oneIteration( state.paused.set(true); } }); - return checkpoint.concatWith(run); + return completeAfterPersisted( + context, "loop step", checkpoint.concatWith(run)); })); return iteration.concatWith( Flowable.defer( diff --git a/core/src/main/java/com/google/adk/agents/SequentialAgent.java b/core/src/main/java/com/google/adk/agents/SequentialAgent.java index 7ed0d90a7..8caef391a 100644 --- a/core/src/main/java/com/google/adk/agents/SequentialAgent.java +++ b/core/src/main/java/com/google/adk/agents/SequentialAgent.java @@ -114,7 +114,10 @@ protected Flowable runAsyncImpl(InvocationContext invocationContext) { return runAsyncLegacyResumption(invocationContext, subAgents); } return Flowable.fromIterable(subAgents) - .concatMap(subAgent -> subAgent.runAsync(invocationContext)); + .concatMap( + subAgent -> + completeAfterPersisted( + invocationContext, "sequential step", subAgent.runAsync(invocationContext))); } /** @@ -132,14 +135,17 @@ private Flowable runAsyncLegacyResumption( subAgent -> paused.get() ? Flowable.empty() - : subAgent - .runAsync(invocationContext) - .doOnNext( - event -> { - if (WorkflowAgentResumption.hasPendingLongRunningCall(event)) { - paused.set(true); - } - })); + : completeAfterPersisted( + invocationContext, + "sequential step", + subAgent + .runAsync(invocationContext) + .doOnNext( + event -> { + if (WorkflowAgentResumption.hasPendingLongRunningCall(event)) { + paused.set(true); + } + }))); } /** @@ -198,7 +204,8 @@ private Flowable runAsyncResumable( paused.set(true); } }); - return checkpoint.concatWith(run); + return completeAfterPersisted( + invocationContext, "sequential step", checkpoint.concatWith(run)); })) .concatWith( Flowable.defer( diff --git a/core/src/main/java/com/google/adk/flows/llmflows/PersistBarrier.java b/core/src/main/java/com/google/adk/flows/llmflows/PersistBarrier.java index 1645de6da..25735f654 100644 --- a/core/src/main/java/com/google/adk/flows/llmflows/PersistBarrier.java +++ b/core/src/main/java/com/google/adk/flows/llmflows/PersistBarrier.java @@ -31,8 +31,8 @@ * persister -- has appended the current step's events, so the next step's request (built from * {@code session.events()} by {@link Contents}) is not assembled from a stale session. The {@code * Runner} calls {@link #markPersisted} (or {@link #markFailed}) after each append; the flow calls - * {@link #awaitPersisted} between steps. State lives in the per-invocation {@link - * InvocationContext#callbackContextData()} map, shared across the agent tree. + * {@link #awaitPersisted} between steps, and workflow agents between sub-agents. State lives in the + * per-invocation {@link InvocationContext#callbackContextData()} map, shared across the agent tree. * *

Each event id maps to a {@link CompletableSubject}: pending until its append finishes, then * terminally completed or failed. The subject retains its terminal state, so {@code @@ -60,6 +60,8 @@ private PersistBarrier() {} */ public static void enable(InvocationContext context) { context.callbackContextData().put(ENABLED_KEY, true); + // Created now, so a shallow copy of callbackContextData still shares the Runner's marks. + var unused = barriers(context); } /** 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 e25ab3a7d..fc1e51192 100644 --- a/core/src/main/java/com/google/adk/runner/Runner.java +++ b/core/src/main/java/com/google/adk/runner/Runner.java @@ -185,10 +185,11 @@ public Builder plugins(Plugin... plugins) { * 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, workflow node retries and attempts with - * a timeout, 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()}. + * a timeout, the hop that starts the live send loop, and the next sub-agent of a {@code + * SequentialAgent} or {@code LoopAgent} when the previous one's events are persisted late. 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 diff --git a/core/src/test/java/com/google/adk/agents/LoopAgentTest.java b/core/src/test/java/com/google/adk/agents/LoopAgentTest.java index 049f0e0b3..0452b132f 100644 --- a/core/src/test/java/com/google/adk/agents/LoopAgentTest.java +++ b/core/src/test/java/com/google/adk/agents/LoopAgentTest.java @@ -20,15 +20,22 @@ import static com.google.adk.testing.TestUtils.createEvent; import static com.google.adk.testing.TestUtils.createInvocationContext; import static com.google.adk.testing.TestUtils.createResumableInvocationContext; +import static com.google.adk.testing.TestUtils.createStreamingTextLlmResponses; import static com.google.adk.testing.TestUtils.createSubAgent; +import static com.google.adk.testing.TestUtils.createTestAgentBuilder; +import static com.google.adk.testing.TestUtils.createTestLlm; +import static com.google.adk.testing.TestUtils.runRequestingLate; import static com.google.common.collect.ImmutableList.toImmutableList; import static com.google.common.truth.Truth.assertThat; import static org.junit.Assert.assertThrows; import com.google.adk.agents.Callbacks.BeforeAgentCallback; +import com.google.adk.apps.App; import com.google.adk.apps.ResumabilityConfig; import com.google.adk.events.Event; +import com.google.adk.runner.Runner; import com.google.adk.testing.TestBaseAgent; +import com.google.adk.testing.TestLlm; import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableMap; import com.google.common.collect.ImmutableSet; @@ -40,6 +47,7 @@ import io.reactivex.rxjava3.core.Maybe; import java.util.List; import java.util.concurrent.atomic.AtomicInteger; +import org.jspecify.annotations.Nullable; import org.junit.Test; import org.junit.runner.RunWith; import org.junit.runners.JUnit4; @@ -441,4 +449,54 @@ private static InvocationContext createLegacyResumableInvocationContext(LoopAgen ResumabilityConfig.builder().plainTextContinuationAutoResume(true).build()) .build(); } + + // ---- Through a Runner, the next iteration sees the state the previous one saved. ---- + + @Test + public void runAsync_withLateConsumer_nextIterationSeesOutputKey() { + assertThat(secondIterationInstruction("v1", /* resumability= */ null)).contains("Improve: v1"); + } + + @Test + public void runAsync_resumableWithLateConsumer_nextIterationSeesOutputKey() { + ResumabilityConfig resumable = ResumabilityConfig.builder().resumable(true).build(); + + assertThat(secondIterationInstruction("v1", resumable)).contains("Improve: v1"); + } + + @Test + @SuppressWarnings("deprecation") // The legacy resumption flow is deprecated by design. + public void runAsync_legacyResumptionWithLateConsumer_nextIterationSeesOutputKey() { + ResumabilityConfig legacy = + ResumabilityConfig.builder().plainTextContinuationAutoResume(true).build(); + + assertThat(secondIterationInstruction("v1", legacy)).contains("Improve: v1"); + } + + /** + * Loops twice over a writer with instruction "Improve: {draft?}" that streams {@code firstDraft} + * and saves it to "draft", consuming the run late; returns the instruction of the second request. + */ + private static String secondIterationInstruction( + String firstDraft, @Nullable ResumabilityConfig resumability) { + TestLlm writerLlm = + createTestLlm( + createStreamingTextLlmResponses(firstDraft), createStreamingTextLlmResponses("v2")); + LlmAgent writer = + createTestAgentBuilder(writerLlm) + .name("writer") + .instruction("Improve: {draft?}") + .outputKey("draft") + .build(); + LoopAgent loop = LoopAgent.builder().name("loop").subAgents(writer).maxIterations(2).build(); + Runner runner = + Runner.builder() + .app(App.builder().name("app").rootAgent(loop).resumabilityConfig(resumability).build()) + .build(); + + runRequestingLate( + runner, RunConfig.builder().streamingMode(RunConfig.StreamingMode.SSE).build()) + .assertComplete(); + return String.join("\n", writerLlm.getRequests().get(1).getSystemInstructions()); + } } diff --git a/core/src/test/java/com/google/adk/agents/SequentialAgentTest.java b/core/src/test/java/com/google/adk/agents/SequentialAgentTest.java index 2989acbcf..e92c8c226 100644 --- a/core/src/test/java/com/google/adk/agents/SequentialAgentTest.java +++ b/core/src/test/java/com/google/adk/agents/SequentialAgentTest.java @@ -20,21 +20,36 @@ import static com.google.adk.testing.TestUtils.createInvocationContext; import static com.google.adk.testing.TestUtils.createLlmResponse; import static com.google.adk.testing.TestUtils.createResumableInvocationContext; +import static com.google.adk.testing.TestUtils.createStreamingTextLlmResponses; 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; import static com.google.adk.testing.TestUtils.createTextLlmResponse; +import static com.google.adk.testing.TestUtils.runRequestingLate; import static com.google.adk.testing.TestUtils.simplifyEvents; import static com.google.common.collect.Iterables.getOnlyElement; import static com.google.common.truth.Truth.assertThat; +import static java.util.concurrent.TimeUnit.MILLISECONDS; import static org.junit.Assert.assertThrows; - +import static org.mockito.AdditionalAnswers.delegatesTo; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.argThat; +import static org.mockito.Mockito.doAnswer; +import static org.mockito.Mockito.doReturn; +import static org.mockito.Mockito.mock; + +import com.google.adk.agents.Callbacks.BeforeAgentCallback; +import com.google.adk.apps.App; import com.google.adk.apps.ResumabilityConfig; import com.google.adk.artifacts.InMemoryArtifactService; import com.google.adk.events.Event; +import com.google.adk.events.EventActions; +import com.google.adk.runner.Runner; +import com.google.adk.sessions.BaseSessionService; 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.common.collect.ImmutableList; @@ -44,9 +59,23 @@ import com.google.genai.types.FunctionCall; import com.google.genai.types.FunctionResponse; import com.google.genai.types.Part; +import io.reactivex.rxjava3.core.Flowable; +import io.reactivex.rxjava3.core.Maybe; +import io.reactivex.rxjava3.core.Single; +import io.reactivex.rxjava3.functions.Consumer; +import io.reactivex.rxjava3.plugins.RxJavaPlugins; +import io.reactivex.rxjava3.schedulers.Schedulers; +import java.time.Duration; import java.util.HashMap; import java.util.List; import java.util.Map; +import java.util.UUID; +import java.util.concurrent.CopyOnWriteArrayList; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.atomic.AtomicReference; +import java.util.function.UnaryOperator; +import org.jspecify.annotations.Nullable; import org.junit.Test; import org.junit.runner.RunWith; import org.junit.runners.JUnit4; @@ -55,6 +84,11 @@ @RunWith(JUnit4.class) public final class SequentialAgentTest { + private static final String APP_NAME = "app"; + private static final String NOT_RUN = ""; + private static final RunConfig SSE = + RunConfig.builder().streamingMode(RunConfig.StreamingMode.SSE).build(); + @Test public void runAsync_withNoSubAgents_returnsEmptyEvents() { SequentialAgent sequentialAgent = @@ -507,4 +541,393 @@ public void runAsync_legacyResumption_doesNotRerunCompletedSubAgents() { assertThat(agent2.getInvocationCount()).isEqualTo(1); assertThat(agent3.getInvocationCount()).isEqualTo(1); } + + // ---- Through a Runner, the next sub-agent sees the state the previous one saved. ---- + + @Test + public void runAsync_inParallelAgentWithSlowAppend_nextSubAgentSeesOutputKey() { + LlmAgent writer = + createTestAgentBuilder(createTestLlm(createTextLlmResponse("upstream output"))) + .name("writer") + .outputKey("k") + .build(); + AtomicReference<@Nullable Object> seen = new AtomicReference<>(NOT_RUN); + // "silent" covers a sub-agent with no event to wait for. + ParallelAgent root = + ParallelAgent.builder() + .name("parallel") + .subAgents( + SequentialAgent.builder() + .name("sequence") + .subAgents(createSubAgent("silent"), writer, stateReader("k", seen)) + .build()) + .build(); + Runner runner = + Runner.builder() + .app(App.builder().name(APP_NAME).rootAgent(root).build()) + .sessionService(appendDelayingSessionService()) + .build(); + + runRequestingLate(runner, RunConfig.builder().build()).assertComplete(); + + assertThat(seen.get()).isEqualTo("upstream output"); + } + + @Test + public void runAsync_afterParallelAgentWithSlowAppend_nextSubAgentSeesEveryBranchOutputKey() { + TestLlm aggregatorLlm = createTestLlm(createTextLlmResponse("combined")); + Runner runner = + Runner.builder() + .app(App.builder().name(APP_NAME).rootAgent(fanOutThenAggregate(aggregatorLlm)).build()) + .sessionService(appendDelayingSessionService()) + .build(); + + runRequestingLate(runner, RunConfig.builder().build()).assertComplete(); + + assertThat(String.join("\n", aggregatorLlm.getLastRequest().getSystemInstructions())) + .contains("Combine: from a and from b"); + } + + @Test + public void runAsync_withLateConsumer_nextAgentInstructionSeesOutputKey() { + TestLlm editorLlm = createTestLlm(createTextLlmResponse("edited")); + LlmAgent editor = + createTestAgentBuilder(editorLlm).name("editor").instruction("Rewrite: {k}").build(); + SequentialAgent pipeline = + SequentialAgent.builder() + .name("pipeline") + .subAgents(streamingWriter("k", "first draft"), editor) + .build(); + Runner runner = + Runner.builder().app(App.builder().name(APP_NAME).rootAgent(pipeline).build()).build(); + + runRequestingLate(runner, SSE).assertComplete(); + + assertThat(String.join("\n", editorLlm.getLastRequest().getSystemInstructions())) + .contains("Rewrite: first draft"); + } + + @Test + public void runAsync_resumableWithLateConsumer_nextSubAgentSeesOutputKey() { + ResumabilityConfig resumable = ResumabilityConfig.builder().resumable(true).build(); + + assertThat(outputKeySeenByNextSubAgent("saved", resumable)).isEqualTo("saved"); + } + + @Test + @SuppressWarnings("deprecation") // The legacy shim is deprecated by design. + public void runAsync_legacyResumptionWithLateConsumer_nextSubAgentSeesOutputKey() { + ResumabilityConfig legacy = + ResumabilityConfig.builder().plainTextContinuationAutoResume(true).build(); + + assertThat(outputKeySeenByNextSubAgent("saved", legacy)).isEqualTo("saved"); + } + + // A wrapper may drop partial events before they reach the Runner, so the wait must skip them. + @Test + public void runAsync_trailingPartialDroppedBeforeRunner_nextSubAgentStillRuns() { + Event saved = savedEvent(); + Event trailingChunk = createEvent("chunk").toBuilder().author("writer").partial(true).build(); + AtomicReference<@Nullable Object> seen = new AtomicReference<>(NOT_RUN); + SequentialAgent sequence = + SequentialAgent.builder() + .name("sequence") + .subAgents(createSubAgent("writer", saved, trailingChunk), stateReader("k", seen)) + .build(); + Runner runner = + Runner.builder() + .app( + App.builder() + .name(APP_NAME) + .rootAgent( + new TransformingAgent( + sequence, events -> events.filter(e -> !e.partial().orElse(false)))) + .build()) + .build(); + + runRequestingLate(runner, RunConfig.builder().build()).assertComplete(); + + assertThat(seen.get()).isEqualTo("saved value"); + } + + // A wrapper that drops, re-ids or holds back events must not hang the run. + @Test + public void runAsync_underAgentDroppingEvents_stopsWaitingAfterTimeout() { + assertSequenceStopsWaitingUnder(events -> events.ignoreElements().toFlowable()); + } + + @Test + public void runAsync_underAgentReIdingEvents_stopsWaitingAfterTimeout() { + assertSequenceStopsWaitingUnder( + events -> events.map(e -> e.toBuilder().id(UUID.randomUUID().toString()).build())); + } + + @Test + public void runAsync_underAgentBufferingEvents_stopsWaitingAfterTimeout() { + assertSequenceStopsWaitingUnder(events -> events.toList().flattenAsFlowable(list -> list)); + } + + @Test + public void runAsync_appendFails_failsOnceWithoutRunningNextSubAgent() { + RuntimeException failure = new RuntimeException("append failed"); + BaseSessionService failingService = + mock(BaseSessionService.class, delegatesTo(new InMemorySessionService())); + // doReturn, because when(...) would call the delegate's real appendEvent. + doReturn(Single.error(failure)) + .when(failingService) + .appendEvent(any(), argThat(event -> event.author().equals("writer"))); + AtomicReference<@Nullable Object> seen = new AtomicReference<>(NOT_RUN); + SequentialAgent sequence = + SequentialAgent.builder() + .name("sequence") + .subAgents(streamingWriter("k", "upstream output"), stateReader("k", seen)) + .build(); + Runner runner = + Runner.builder() + .app(App.builder().name(APP_NAME).rootAgent(sequence).build()) + .sessionService(failingService) + .build(); + List undelivered = new CopyOnWriteArrayList<>(); + Consumer previousHandler = RxJavaPlugins.getErrorHandler(); + RxJavaPlugins.setErrorHandler(undelivered::add); + try { + runRequestingLate(runner, SSE).assertError(failure); + } finally { + RxJavaPlugins.setErrorHandler(previousHandler); + } + + assertThat(seen.get()).isEqualTo(NOT_RUN); + assertThat(undelivered).isEmpty(); + } + + // An event without an id cannot be tracked, so the wait uses the last event that has one. + @Test + public void runAsync_lastEventHasNoId_nextSubAgentWaitsForLastEventWithId() { + // With no demand yet, "saved" stays queued, unpersisted, behind "first". + Event first = createEvent("first").toBuilder().author("writer").build(); + Event saved = savedEvent(); + Event withoutId = createEvent("unused").toBuilder().id(null).author("writer").build(); + AtomicReference<@Nullable Object> seen = new AtomicReference<>(NOT_RUN); + SequentialAgent sequence = + SequentialAgent.builder() + .name("sequence") + .subAgents(createSubAgent("writer", first, saved, withoutId), stateReader("k", seen)) + .build(); + Runner runner = + Runner.builder().app(App.builder().name(APP_NAME).rootAgent(sequence).build()).build(); + + runRequestingLate(runner, RunConfig.builder().build()).assertComplete(); + + assertThat(seen.get()).isEqualTo("saved value"); + } + + // A late append completes on its own thread, which must not run the next sub-agent's model call. + @Test + public void runAsync_afterSlowAppend_nextSubAgentStartsOnInvocationScheduler() { + ExecutorService executor = + Executors.newSingleThreadExecutor(runnable -> new Thread(runnable, "invocation-scheduler")); + try { + AtomicReference readerThread = new AtomicReference<>(); + BeforeAgentCallback recordThread = + callbackContext -> { + readerThread.set(Thread.currentThread().getName()); + return Maybe.empty(); + }; + SequentialAgent sequence = + SequentialAgent.builder() + .name("sequence") + .subAgents( + createSubAgent("writer", savedEvent()), + new TestBaseAgent( + "reader", + "description", + ImmutableList.of(recordThread), + ImmutableList.of(), + Flowable::empty)) + .build(); + Runner runner = + Runner.builder() + .app(App.builder().name(APP_NAME).rootAgent(sequence).build()) + .sessionService(appendDelayingSessionService()) + .scheduler(Schedulers.from(executor)) + .build(); + + runRequestingLate(runner, RunConfig.builder().build()).assertComplete(); + + assertThat(readerThread.get()).isEqualTo("invocation-scheduler"); + } finally { + executor.shutdownNow(); + } + } + + @Test + public void runAsync_eventAlreadyPersisted_nextSubAgentStartsWithoutSchedulerHop() { + RecordingScheduler scheduler = new RecordingScheduler(Schedulers.trampoline()); + AtomicReference<@Nullable Object> seen = new AtomicReference<>(NOT_RUN); + SequentialAgent sequence = + SequentialAgent.builder() + .name("sequence") + .subAgents(createSubAgent("writer", savedEvent()), stateReader("k", seen)) + .build(); + Runner runner = + Runner.builder() + .app(App.builder().name(APP_NAME).rootAgent(sequence).build()) + .scheduler(scheduler) + .build(); + + runRequestingLate(runner, RunConfig.builder().build()).assertComplete(); + + assertThat(seen.get()).isEqualTo("saved value"); + assertThat(scheduler.workersCreated()).isEqualTo(0); + } + + /** + * Runs a writer streaming {@code text} to state, then a state reader, with {@code resumability}; + * returns what the reader saw. + */ + private static @Nullable Object outputKeySeenByNextSubAgent( + String text, ResumabilityConfig resumability) { + AtomicReference<@Nullable Object> seen = new AtomicReference<>(NOT_RUN); + SequentialAgent sequence = + SequentialAgent.builder() + .name("sequence") + .subAgents(streamingWriter("k", text), stateReader("k", seen)) + .build(); + Runner runner = + Runner.builder() + .app( + App.builder() + .name(APP_NAME) + .rootAgent(sequence) + .resumabilityConfig(resumability) + .build()) + .build(); + + runRequestingLate(runner, SSE).assertComplete(); + return seen.get(); + } + + /** + * Returns a SequentialAgent that runs a ParallelAgent, whose branches save "from a" to "a" and + * "from b" to "b", then an aggregator that sends {@code aggregatorLlm} "Combine: {a} and {b}". + */ + private static SequentialAgent fanOutThenAggregate(TestLlm aggregatorLlm) { + ParallelAgent fanOut = + ParallelAgent.builder() + .name("fan_out") + .subAgents( + createTestAgentBuilder(createTestLlm(createTextLlmResponse("from a"))) + .name("writer_a") + .outputKey("a") + .build(), + createTestAgentBuilder(createTestLlm(createTextLlmResponse("from b"))) + .name("writer_b") + .outputKey("b") + .build()) + .build(); + LlmAgent aggregator = + createTestAgentBuilder(aggregatorLlm) + .name("aggregator") + .instruction("Combine: {a} and {b}") + .build(); + return SequentialAgent.builder().name("sequence").subAgents(fanOut, aggregator).build(); + } + + /** Returns an in-memory session service whose {@code appendEvent} completes 50 ms later. */ + private static BaseSessionService appendDelayingSessionService() { + InMemorySessionService delegate = new InMemorySessionService(); + BaseSessionService service = mock(BaseSessionService.class, delegatesTo(delegate)); + doAnswer( + invocation -> + Single.timer(50, MILLISECONDS) + .flatMap( + unused -> + delegate.appendEvent( + invocation.getArgument(0), invocation.getArgument(1)))) + .when(service) + .appendEvent(any(), any()); + return service; + } + + /** Returns an LLM agent that streams {@code text}, then saves it to {@code key}. */ + private static LlmAgent streamingWriter(String key, String text) { + return createTestAgentBuilder(createTestLlm(createStreamingTextLlmResponses(text))) + .name("writer") + .outputKey(key) + .build(); + } + + /** Returns an agent that records what session state holds under {@code key} when it starts. */ + private static TestBaseAgent stateReader(String key, AtomicReference<@Nullable Object> seen) { + BeforeAgentCallback readState = + callbackContext -> { + seen.set(callbackContext.state().get(key)); + return Maybe.empty(); + }; + return new TestBaseAgent( + "reader", "description", ImmutableList.of(readState), ImmutableList.of(), Flowable::empty); + } + + /** + * Runs a writer then a state reader under an agent that applies {@code transform}, with a 200 ms + * bound on the wait, and checks that the run completes, the reader runs, and the run takes at + * least the bound. + */ + private static void assertSequenceStopsWaitingUnder(UnaryOperator> transform) { + Duration defaultTimeout = BaseAgent.persistWaitTimeout; + BaseAgent.persistWaitTimeout = Duration.ofMillis(200); + try { + Event saved = savedEvent(); + AtomicReference<@Nullable Object> seen = new AtomicReference<>(NOT_RUN); + SequentialAgent sequence = + SequentialAgent.builder() + .name("sequence") + .subAgents(createSubAgent("writer", saved), stateReader("k", seen)) + .build(); + Runner runner = + Runner.builder() + .app( + App.builder() + .name(APP_NAME) + .rootAgent(new TransformingAgent(sequence, transform)) + .build()) + .build(); + long start = System.nanoTime(); + + runRequestingLate(runner, RunConfig.builder().build()).assertComplete(); + + assertThat(seen.get()).isNotEqualTo(NOT_RUN); + assertThat(Duration.ofNanos(System.nanoTime() - start)).isAtLeast(Duration.ofMillis(200)); + } finally { + BaseAgent.persistWaitTimeout = defaultTimeout; + } + } + + /** Returns a "writer" event that saves "saved value" to "k". */ + private static Event savedEvent() { + return createEvent("saved").toBuilder() + .author("writer") + .actions(EventActions.builder().stateDelta(ImmutableMap.of("k", "saved value")).build()) + .build(); + } + + /** An agent that passes its only sub-agent's events through {@code transform}. */ + private static final class TransformingAgent extends BaseAgent { + private final UnaryOperator> transform; + + TransformingAgent(BaseAgent subAgent, UnaryOperator> transform) { + super("wrapper", "description", ImmutableList.of(subAgent), null, null); + this.transform = transform; + } + + @Override + protected Flowable runAsyncImpl(InvocationContext invocationContext) { + return transform.apply(subAgents().get(0).runAsync(invocationContext)); + } + + @Override + protected Flowable runLiveImpl(InvocationContext invocationContext) { + return Flowable.empty(); + } + } } diff --git a/core/src/test/java/com/google/adk/flows/llmflows/PersistBarrierTest.java b/core/src/test/java/com/google/adk/flows/llmflows/PersistBarrierTest.java index 478e6bb9b..44daa748f 100644 --- a/core/src/test/java/com/google/adk/flows/llmflows/PersistBarrierTest.java +++ b/core/src/test/java/com/google/adk/flows/llmflows/PersistBarrierTest.java @@ -29,6 +29,7 @@ import java.util.ArrayList; import java.util.Collections; import java.util.List; +import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.CountDownLatch; import org.junit.Before; import org.junit.Test; @@ -71,6 +72,22 @@ public void awaitBeforeMark_completesOnMark_andDrainsPending() { assertThat(PersistBarrier.pendingCount(context)).isEqualTo(0); } + // As when an agent runs a sub-agent on its own copy of the map, before any event is persisted. + @Test + public void awaitOnShallowCopyMadeAfterEnable_completesOnMarkToOriginal() { + PersistBarrier.enable(context); + InvocationContext copy = + context.toBuilder() + .callbackContextData(new ConcurrentHashMap<>(context.callbackContextData())) + .build(); + + TestObserver observer = + PersistBarrier.awaitPersisted(copy, ImmutableList.of(event("e1"))).test(); + PersistBarrier.markPersisted(context, "e1"); + + observer.assertComplete(); + } + @Test public void markBeforeAwait_completesImmediately_noPending() { PersistBarrier.enable(context); 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 2b9951121..967f12838 100644 --- a/core/src/test/java/com/google/adk/runner/RunnerTest.java +++ b/core/src/test/java/com/google/adk/runner/RunnerTest.java @@ -1214,6 +1214,48 @@ public void runAsync_slowAppendEvent_transferTargetSpanIsChildOfParentAgentSpan( .isEqualTo(parentSpan.get().getSpanContext().getSpanId()); } + /** + * With a slow appendEvent, the next sequential sub-agent starts on another thread; its {@code + * invoke_agent} span must still be a child of the SequentialAgent's, not start a new trace. + */ + @Test + public void runAsync_slowAppendEvent_nextSequentialSubAgentSpanIsChildOfSequentialAgentSpan() { + SequentialAgent sequentialAgent = + SequentialAgent.builder() + .name("sequence") + .subAgents( + createTestAgentBuilder(createTestLlm(createTextLlmResponse("first"))) + .name("first_agent") + .build(), + createTestAgentBuilder(createTestLlm(createTextLlmResponse("second"))) + .name("second_agent") + .build()) + .build(); + Runner runnerForRace = + Runner.builder() + .app(App.builder().name("test").rootAgent(sequentialAgent).build()) + .sessionService(new AppendDelayingSessionService(new InMemorySessionService(), 50)) + .build(); + Session raceSession = + runnerForRace.sessionService().createSession("test", "user").blockingGet(); + + var unused = + runnerForRace + .runAsync("user", raceSession.id(), createContent("start")) + .toList() + .blockingGet(); + + List spans = openTelemetryRule.getSpans(); + Optional parentSpan = + spans.stream().filter(s -> s.getName().equals("invoke_agent sequence")).findFirst(); + Optional secondSpan = + spans.stream().filter(s -> s.getName().equals("invoke_agent second_agent")).findFirst(); + assertThat(parentSpan).isPresent(); + assertThat(secondSpan).isPresent(); + assertThat(secondSpan.get().getParentSpanContext().getSpanId()) + .isEqualTo(parentSpan.get().getSpanContext().getSpanId()); + } + /** * When an LlmAgent transfers to a sub-LlmAgent, the sub-agent's events flow back up through the * parent's flow and must each be appended to the session exactly once. diff --git a/core/src/test/java/com/google/adk/testing/TestUtils.java b/core/src/test/java/com/google/adk/testing/TestUtils.java index 577241e78..af25fd0a8 100644 --- a/core/src/test/java/com/google/adk/testing/TestUtils.java +++ b/core/src/test/java/com/google/adk/testing/TestUtils.java @@ -18,6 +18,7 @@ import static com.google.common.collect.ImmutableList.toImmutableList; import static com.google.common.truth.Truth.assertThat; +import static java.util.concurrent.TimeUnit.SECONDS; import static java.util.stream.Collectors.joining; import com.google.adk.agents.BaseAgent; @@ -31,6 +32,7 @@ import com.google.adk.events.EventCompaction; import com.google.adk.models.BaseLlm; import com.google.adk.models.LlmResponse; +import com.google.adk.runner.Runner; import com.google.adk.sessions.BaseSessionService; import com.google.adk.sessions.InMemorySessionService; import com.google.adk.sessions.Session; @@ -46,6 +48,7 @@ import com.google.genai.types.Part; import io.reactivex.rxjava3.core.Flowable; import io.reactivex.rxjava3.core.Single; +import io.reactivex.rxjava3.subscribers.TestSubscriber; import java.util.Arrays; import java.util.List; import java.util.Map; @@ -224,6 +227,20 @@ private static Part overwriteFunctionIdsInPart(Part part) { return part; } + /** + * Runs {@code runner} on a new session with a subscriber that requests events only after the run + * has started. A run on the calling thread meanwhile leaves its later events queued, unpersisted. + */ + public static TestSubscriber runRequestingLate(Runner runner, RunConfig runConfig) { + Session session = runner.sessionService().createSession(runner.appName(), "user").blockingGet(); + TestSubscriber subscriber = + runner + .runAsync("user", session.id(), Content.fromParts(Part.fromText("go")), runConfig) + .test(0); + subscriber.request(Long.MAX_VALUE); + return subscriber.awaitDone(5, SECONDS); + } + public static TestBaseAgent createRootAgent(BaseAgent... subAgents) { return createRootAgent(Arrays.asList(subAgents)); } @@ -280,6 +297,12 @@ public static LlmResponse createTextLlmResponse(String text) { return createLlmResponse(Content.builder().role("model").parts(Part.fromText(text)).build()); } + /** Returns a model stream that sends {@code text} as one partial chunk, then as the answer. */ + public static Flowable createStreamingTextLlmResponses(String text) { + LlmResponse answer = createTextLlmResponse(text); + return Flowable.just(answer.toBuilder().partial(true).build(), answer); + } + public static LlmResponse createFunctionCallLlmResponse( String id, String functionName, Map args) { Content content =