fix: wait for the Runner to persist a step's events before the ADK flow's next step (sequential-tool-execution race)

`BaseLlmFlow.run()` builds each step's request from the session, but the
`Runner` persists events asynchronously downstream of the flow. A slow
`appendEvent` could let the next step start from a stale session missing the
prior step's events (e.g. a tool's function response), making the model re-call
the tool or hallucinate its result.

The `Runner` stays the sole `appendEvent` caller and the flow waits: it calls
`PersistBarrier.markPersisted(id)` after each append, and the flow calls
`PersistBarrier.awaitPersisted(stepEvents)` between steps. The barrier is a
reactive per-event signal in the shared `InvocationContext.callbackContextData`
and never blocks a thread; `Contents` is unchanged. `PersistBarrier.enable()`,
called by the `Runner`, keeps `awaitPersisted` a no-op when the flow runs without
a `Runner`.

The barrier self-cleans: a persisted id is recorded in a lightweight set, while a
`CompletableSubject` is kept only for an id awaited before it is persisted and
dropped once persisted -- so the subject map does not grow with invocation length
(an id can be awaited at several flow levels, e.g. across an agent transfer). It
is thread-safe (concurrent maps plus a register-then-recheck handshake), since
`markPersisted` may run off-thread when an async `appendEvent` completes.

PiperOrigin-RevId: 921989444

Author: damianmomotgoogle

core/src/main/java/com/google/adk/flows/llmflows/BaseLlmFlow.java

@@ -523,7 +523,10 @@ private Flowable<Event> run(
                     return Flowable.empty();
                   } else {
                     logger.debug("Continuing to next step of the flow.");
-                    return run(spanContext, invocationContext, stepsCompleted + 1);
+                    // Wait until the Runner has persisted this step's events so the next step's
+                    // request is not built from a stale session (see PersistBarrier).
+                    return PersistBarrier.awaitPersisted(invocationContext, eventList)
+                        .andThen(run(spanContext, invocationContext, stepsCompleted + 1));
                   }
                 }));
   }

core/src/main/java/com/google/adk/flows/llmflows/PersistBarrier.java

@@ -0,0 +1,132 @@
+/*
+ * Copyright 2025 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.flows.llmflows;
+
+import com.google.adk.agents.InvocationContext;
+import com.google.adk.events.Event;
+import com.google.common.annotations.VisibleForTesting;
+import io.reactivex.rxjava3.core.Completable;
+import io.reactivex.rxjava3.subjects.CompletableSubject;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+import java.util.concurrent.ConcurrentHashMap;
+
+/**
+ * Lets {@link BaseLlmFlow}'s multi-step loop wait until the {@code Runner} -- the sole event
+ * 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} 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.
+ *
+ * <p>An event may be awaited at several flow levels (an agent transfer concatenates a sub-agent's
+ * events into the parent's step) but is persisted once, so persisted ids are recorded in a {@link
+ * Set} and a {@link CompletableSubject} is kept only for an id awaited before it is persisted, then
+ * dropped -- the subject map self-drains instead of growing with the invocation.
+ *
+ * <p>Thread-safe: {@code markPersisted} may run off-thread (async {@code appendEvent}) concurrently
+ * with {@code awaitPersisted}; the concurrent maps plus the register-then-recheck in {@link
+ * #awaitOne} close that race without locking.
+ */
+public final class PersistBarrier {
+
+  private static final String ENABLED_KEY = "com.google.adk.flows.llmflows.persistBarrier.enabled";
+  private static final String PERSISTED_KEY =
+      "com.google.adk.flows.llmflows.persistBarrier.persisted";
+  private static final String PENDING_KEY = "com.google.adk.flows.llmflows.persistBarrier.pending";
+
+  private PersistBarrier() {}
+
+  /**
+   * Marks that a {@code Runner} is driving this invocation and will {@link #markPersisted} each
+   * appended event. Otherwise (flow run directly, e.g. unit tests) {@link #awaitPersisted} is a
+   * no-op, avoiding a deadlock waiting for a signal that never comes.
+   */
+  public static void enable(InvocationContext context) {
+    context.callbackContextData().put(ENABLED_KEY, true);
+  }
+
+  /**
+   * Completes once every event in {@code events} has been {@link #markPersisted}, or immediately if
+   * the barrier was never {@link #enable}d. Already-persisted events complete immediately, so the
+   * order of {@code awaitPersisted}/{@code markPersisted} does not matter.
+   */
+  public static Completable awaitPersisted(InvocationContext context, List<Event> events) {
+    if (!Boolean.TRUE.equals(context.callbackContextData().get(ENABLED_KEY))) {
+      return Completable.complete();
+    }
+    Completable result = Completable.complete();
+    for (Event event : events) {
+      String eventId = event.id();
+      if (eventId != null) {
+        result = result.andThen(awaitOne(context, eventId));
+      }
+    }
+    return result;
+  }
+
+  private static Completable awaitOne(InvocationContext context, String eventId) {
+    Set<String> persisted = persistedIds(context);
+    if (persisted.contains(eventId)) {
+      return Completable.complete();
+    }
+    CompletableSubject subject =
+        pending(context).computeIfAbsent(eventId, unusedKey -> CompletableSubject.create());
+    // Re-check after registering so a markPersisted that raced in between is not lost.
+    if (persisted.contains(eventId)) {
+      pending(context).remove(eventId);
+      return Completable.complete();
+    }
+    return subject;
+  }
+
+  /** Signals that the {@code Runner} has finished persisting the event with the given id. */
+  public static void markPersisted(InvocationContext context, String eventId) {
+    if (eventId == null) {
+      return;
+    }
+    persistedIds(context).add(eventId);
+    CompletableSubject subject = pending(context).remove(eventId);
+    if (subject != null) {
+      subject.onComplete();
+    }
+  }
+
+  /** Awaited-but-not-yet-persisted events; drains to 0 once a step's events are persisted. */
+  @VisibleForTesting
+  static int pendingCount(InvocationContext context) {
+    return pending(context).size();
+  }
+
+  @SuppressWarnings("unchecked")
+  private static Set<String> persistedIds(InvocationContext context) {
+    return (Set<String>)
+        context
+            .callbackContextData()
+            .computeIfAbsent(PERSISTED_KEY, unusedKey -> ConcurrentHashMap.<String>newKeySet());
+  }
+
+  @SuppressWarnings("unchecked")
+  private static Map<String, CompletableSubject> pending(InvocationContext context) {
+    return (Map<String, CompletableSubject>)
+        context
+            .callbackContextData()
+            .computeIfAbsent(
+                PENDING_KEY, unusedKey -> new ConcurrentHashMap<String, CompletableSubject>());
+  }
+}

core/src/main/java/com/google/adk/runner/Runner.java

@@ -30,6 +30,7 @@
 import com.google.adk.artifacts.InMemoryArtifactService;
 import com.google.adk.events.Event;
 import com.google.adk.events.EventActions;
+import com.google.adk.flows.llmflows.PersistBarrier;
 import com.google.adk.memory.BaseMemoryService;
 import com.google.adk.models.Model;
 import com.google.adk.plugins.Plugin;
@@ -575,6 +576,9 @@ private Flowable<Event> runAgentWithUpdatedSession(
                         .content(content)
                         .build());
 
+    // Let BaseLlmFlow block each step until this Runner has persisted the prior step's events.
+    PersistBarrier.enable(contextWithUpdatedSession);
+
     // Agent execution
     Flowable<Event> agentEvents =
         contextWithUpdatedSession
@@ -584,6 +588,12 @@ private Flowable<Event> runAgentWithUpdatedSession(
                 agentEvent ->
                     this.sessionService
                         .appendEvent(updatedSession, agentEvent)
+                        // Signal persistence so BaseLlmFlow's loop can release the next step; the
+                        // Runner stays the sole appendEvent caller (see PersistBarrier).
+                        .doOnSuccess(
+                            unusedEvent ->
+                                PersistBarrier.markPersisted(
+                                    contextWithUpdatedSession, agentEvent.id()))
                         .flatMap(
                             registeredEvent -> {
                               // TODO: remove this hack after deprecating runAsync with Session.

core/src/test/java/com/google/adk/flows/llmflows/PersistBarrierTest.java

@@ -0,0 +1,173 @@
+/*
+ * Copyright 2025 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.flows.llmflows;
+
+import static com.google.common.truth.Truth.assertThat;
+import static org.mockito.Mockito.mock;
+
+import com.google.adk.agents.BaseAgent;
+import com.google.adk.agents.InvocationContext;
+import com.google.adk.events.Event;
+import com.google.adk.sessions.BaseSessionService;
+import com.google.adk.sessions.Session;
+import com.google.common.collect.ImmutableList;
+import io.reactivex.rxjava3.observers.TestObserver;
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.List;
+import java.util.concurrent.CountDownLatch;
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+
+@RunWith(JUnit4.class)
+public final class PersistBarrierTest {
+
+  private InvocationContext context;
+
+  @Before
+  public void setUp() {
+    context =
+        InvocationContext.builder()
+            .sessionService(mock(BaseSessionService.class))
+            .invocationId("inv-1")
+            .agent(mock(BaseAgent.class))
+            .session(Session.builder("s").build())
+            .build();
+  }
+
+  private static Event event(String id) {
+    return Event.builder().id(id).author("agent").build();
+  }
+
+  @Test
+  public void awaitBeforeMark_completesOnMark_andDrainsPending() {
+    PersistBarrier.enable(context);
+
+    TestObserver<Void> observer =
+        PersistBarrier.awaitPersisted(context, ImmutableList.of(event("e1"))).test();
+
+    observer.assertNotComplete();
+    assertThat(PersistBarrier.pendingCount(context)).isEqualTo(1);
+
+    PersistBarrier.markPersisted(context, "e1");
+
+    observer.assertComplete();
+    assertThat(PersistBarrier.pendingCount(context)).isEqualTo(0);
+  }
+
+  @Test
+  public void markBeforeAwait_completesImmediately_noPending() {
+    PersistBarrier.enable(context);
+
+    PersistBarrier.markPersisted(context, "e1");
+    PersistBarrier.awaitPersisted(context, ImmutableList.of(event("e1"))).test().assertComplete();
+
+    assertThat(PersistBarrier.pendingCount(context)).isEqualTo(0);
+  }
+
+  @Test
+  public void sameEventAwaitedTwice_secondAwaitStillCompletes_andNothingLingers() {
+    // Mirrors an agent transfer: a sub-agent event is awaited by both the sub-agent and parent
+    // flows but persisted once; the second await must still complete.
+    PersistBarrier.enable(context);
+
+    TestObserver<Void> subLevel =
+        PersistBarrier.awaitPersisted(context, ImmutableList.of(event("e1"))).test();
+    PersistBarrier.markPersisted(context, "e1");
+    subLevel.assertComplete();
+    assertThat(PersistBarrier.pendingCount(context)).isEqualTo(0);
+
+    PersistBarrier.awaitPersisted(context, ImmutableList.of(event("e1"))).test().assertComplete();
+    assertThat(PersistBarrier.pendingCount(context)).isEqualTo(0);
+  }
+
+  @Test
+  public void multiEventStep_completesOnlyAfterAllMarked() {
+    PersistBarrier.enable(context);
+
+    TestObserver<Void> observer =
+        PersistBarrier.awaitPersisted(context, ImmutableList.of(event("e1"), event("e2"))).test();
+
+    PersistBarrier.markPersisted(context, "e1");
+    observer.assertNotComplete();
+
+    PersistBarrier.markPersisted(context, "e2");
+    observer.assertComplete();
+    assertThat(PersistBarrier.pendingCount(context)).isEqualTo(0);
+  }
+
+  @Test
+  public void notEnabled_awaitIsNoOp() {
+    // No enable(): flow runs without a Runner, so await must not block forever.
+    PersistBarrier.awaitPersisted(context, ImmutableList.of(event("e1"))).test().assertComplete();
+    assertThat(PersistBarrier.pendingCount(context)).isEqualTo(0);
+  }
+
+  @Test
+  public void concurrentAwaitAndMark_allComplete_andDrain() throws Exception {
+    // awaitPersisted (flow thread) and markPersisted (async appendEvent thread) race on each id;
+    // none may be stranded and every subject must be dropped.
+    PersistBarrier.enable(context);
+    int eventCount = 1000;
+    List<String> ids = new ArrayList<>();
+    for (int i = 0; i < eventCount; i++) {
+      ids.add("e" + i);
+    }
+    List<TestObserver<Void>> observers = Collections.synchronizedList(new ArrayList<>());
+    CountDownLatch start = new CountDownLatch(1);
+
+    Thread awaiter =
+        new Thread(
+            () -> {
+              awaitQuietly(start);
+              for (String id : ids) {
+                observers.add(
+                    PersistBarrier.awaitPersisted(context, ImmutableList.of(event(id))).test());
+              }
+            });
+    Thread marker =
+        new Thread(
+            () -> {
+              awaitQuietly(start);
+              for (String id : ids) {
+                PersistBarrier.markPersisted(context, id);
+              }
+            });
+
+    awaiter.start();
+    marker.start();
+    start.countDown();
+    awaiter.join();
+    marker.join();
+
+    for (TestObserver<Void> observer : observers) {
+      observer.assertComplete();
+    }
+    assertThat(PersistBarrier.pendingCount(context)).isEqualTo(0);
+  }
+
+  private static void awaitQuietly(CountDownLatch latch) {
+    try {
+      latch.await();
+    } catch (InterruptedException e) {
+      Thread.currentThread().interrupt();
+      throw new AssertionError(e);
+    }
+  }
+}