fix(agents): run the next sub-agent only after the previous one's events are persisted

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

@@ -21,44 +21,59 @@
 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;
 import com.google.adk.telemetry.Tracing;
 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<Event> completeAfterPersisted(
+      InvocationContext context, String waitingPoint, Flowable<Event> 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<Event> 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<Event> 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.

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

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

@@ -161,7 +161,10 @@ protected Flowable<Event> 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<Event> runAsyncLegacyResumption(
             subAgent ->
                 paused.get()
                     ? Flowable.<Event>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<Event> oneIteration(
                                                 state.paused.set(true);
                                               }
                                             });
-                                return checkpoint.concatWith(run);
+                                return completeAfterPersisted(
+                                    context, "loop step", checkpoint.concatWith(run));
                               }));
           return iteration.concatWith(
               Flowable.defer(

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

@@ -114,7 +114,10 @@ protected Flowable<Event> 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<Event> runAsyncLegacyResumption(
             subAgent ->
                 paused.get()
                     ? Flowable.<Event>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<Event> runAsyncResumable(
                                             paused.set(true);
                                           }
                                         });
-                            return checkpoint.concatWith(run);
+                            return completeAfterPersisted(
+                                invocationContext, "sequential step", checkpoint.concatWith(run));
                           }))
               .concatWith(
                   Flowable.defer(

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.
  *
  * <p>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);
   }
 
   /**

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()}.
      *
      * <p>Only these hand-offs are affected; work that completes asynchronously elsewhere (model
      * clients, session services, MCP, plugins) stays on its own threads. {@code

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