fix: Fix ADK Runner race condition for sequential tool execution

BaseLlmFlow.run() now appends each event synchronously inside the per-step
concatMap, so the next runOneStep sees the previous step's events. Without
this, a tool relying on prior events (e.g. a BeforeToolCallback producing a
function response) could see stale history and re-call the tool or
hallucinate its result.

PiperOrigin-RevId: 921989444

Author: damianmomotgoogle

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

@@ -654,7 +654,9 @@ private static boolean isThought(Part part) {
 
   @Override
   protected Flowable<Event> runAsyncImpl(InvocationContext invocationContext) {
-    return llmFlow.run(invocationContext).doOnNext(this::maybeSaveOutputToState);
+    // maybeSaveOutputToState runs as a pre-persist finalizer so the outputKey stateDelta is
+    // part of the persisted append performed inside BaseLlmFlow.run.
+    return llmFlow.run(invocationContext, this::maybeSaveOutputToState);
   }
 
   @Override

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

@@ -37,6 +37,7 @@
 import com.google.adk.models.LlmRegistry;
 import com.google.adk.models.LlmRequest;
 import com.google.adk.models.LlmResponse;
+import com.google.adk.sessions.SessionUtils;
 import com.google.adk.telemetry.Tracing;
 import com.google.adk.tools.BaseTool;
 import com.google.adk.tools.BaseToolset;
@@ -60,8 +61,10 @@
 import java.util.List;
 import java.util.Optional;
 import java.util.Set;
+import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.atomic.AtomicReference;
 import java.util.function.BiFunction;
+import java.util.function.Consumer;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -498,12 +501,54 @@ private Flowable<Event> runOneStep(Context spanContext, InvocationContext contex
    */
   @Override
   public Flowable<Event> run(InvocationContext invocationContext) {
-    return run(Context.current(), invocationContext, 0);
+    return run(invocationContext, event -> {});
+  }
+
+  /**
+   * Same as {@link #run(InvocationContext)} but invokes {@code eventPreFinalize} on each event
+   * authored by this agent immediately before it is appended, so any mutation (e.g. {@code
+   * stateDelta} updates from {@link LlmAgent}'s {@code outputKey}) is part of the persisted append.
+   */
+  public Flowable<Event> run(
+      InvocationContext invocationContext, Consumer<Event> eventPreFinalize) {
+    return run(Context.current(), invocationContext, 0, eventPreFinalize);
   }
 
   private Flowable<Event> run(
-      Context spanContext, InvocationContext invocationContext, int stepsCompleted) {
-    Flowable<Event> currentStepEvents = runOneStep(spanContext, invocationContext).cache();
+      Context spanContext,
+      InvocationContext invocationContext,
+      int stepsCompleted,
+      Consumer<Event> eventPreFinalize) {
+    // Append each event synchronously so the next runOneStep sees prior events (avoids racing the
+    // Runner's append). Record each appended id so the Runner skips re-appending it; skip ids
+    // already recorded (e.g. bubbling up from a sub-agent's flow). Emit the original event, not the
+    // service's return (which may be a mock sentinel).
+    String thisAgentName = invocationContext.agent().name();
+    Flowable<Event> currentStepEvents =
+        runOneStep(spanContext, invocationContext)
+            .concatMap(
+                event -> {
+                  String eventId = event.id();
+                  if (eventId != null
+                      && inFlightAppendedEventIds(invocationContext).contains(eventId)) {
+                    return Flowable.just(event);
+                  }
+                  if (thisAgentName != null && thisAgentName.equals(event.author())) {
+                    eventPreFinalize.accept(event);
+                  }
+                  return SessionUtils.safeAppendEvent(
+                          invocationContext.sessionService(), invocationContext.session(), event)
+                      .ignoreElement()
+                      .andThen(
+                          Flowable.fromCallable(
+                              () -> {
+                                if (eventId != null) {
+                                  inFlightAppendedEventIds(invocationContext).add(eventId);
+                                }
+                                return event;
+                              }));
+                })
+            .cache();
     if (stepsCompleted + 1 >= maxSteps) {
       logger.debug("Ending flow execution because max steps reached.");
       return currentStepEvents;
@@ -523,11 +568,30 @@ private Flowable<Event> run(
                     return Flowable.empty();
                   } else {
                     logger.debug("Continuing to next step of the flow.");
-                    return run(spanContext, invocationContext, stepsCompleted + 1);
+                    return run(
+                        spanContext, invocationContext, stepsCompleted + 1, eventPreFinalize);
                   }
                 }));
   }
 
+  private static final String IN_FLIGHT_APPENDED_EVENT_IDS_KEY =
+      "com.google.adk.internal.inFlightAppendedEventIds";
+
+  /**
+   * Returns the transient, per-invocation set of event ids appended by the flow but not yet
+   * consumed by the Runner, lazily creating it. Ids are added here on append and removed by the
+   * Runner on consume, so this is hand-off state -- not a record of all persisted events.
+   */
+  @SuppressWarnings("unchecked")
+  private static Set<String> inFlightAppendedEventIds(InvocationContext invocationContext) {
+    return (Set<String>)
+        invocationContext
+            .callbackContextData()
+            .computeIfAbsent(
+                IN_FLIGHT_APPENDED_EVENT_IDS_KEY,
+                unusedKey -> ConcurrentHashMap.<String>newKeySet());
+  }
+
   /**
    * Executes the LLM flow in streaming mode.
    *

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

@@ -38,6 +38,7 @@
 import com.google.adk.sessions.InMemorySessionService;
 import com.google.adk.sessions.Session;
 import com.google.adk.sessions.SessionKey;
+import com.google.adk.sessions.SessionUtils;
 import com.google.adk.summarizer.EventsCompactionConfig;
 import com.google.adk.summarizer.LlmEventSummarizer;
 import com.google.adk.summarizer.SlidingWindowEventCompactor;
@@ -70,6 +71,7 @@
 import java.util.Map;
 import java.util.Objects;
 import java.util.Optional;
+import java.util.Set;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.ConcurrentMap;
 import org.jspecify.annotations.Nullable;
@@ -581,19 +583,32 @@ private Flowable<Event> runAgentWithUpdatedSession(
             .agent()
             .runAsync(contextWithUpdatedSession)
             .concatMap(
-                agentEvent ->
-                    this.sessionService
-                        .appendEvent(updatedSession, agentEvent)
-                        .flatMap(
-                            registeredEvent -> {
-                              // TODO: remove this hack after deprecating runAsync with Session.
-                              copySessionStates(updatedSession, initialContext.session());
-                              return contextWithUpdatedSession
+                agentEvent -> {
+                  // TODO: remove this hack after deprecating runAsync with Session.
+                  copySessionStates(updatedSession, initialContext.session());
+                  // Skip events already appended by BaseLlmFlow.run (id recorded in the shared
+                  // set),
+                  // removing the id to keep the set bounded. Append everything else (agent-callback
+                  // events, non-LlmAgent leaves) here, so each event is appended exactly once.
+                  String agentEventId = agentEvent.id();
+                  boolean alreadyPersisted =
+                      agentEventId != null
+                          && inFlightAppendedEventIds(contextWithUpdatedSession)
+                              .remove(agentEventId);
+                  Single<Event> appendResult =
+                      alreadyPersisted
+                          ? Single.just(agentEvent)
+                          : SessionUtils.safeAppendEvent(
+                              this.sessionService, updatedSession, agentEvent);
+                  return appendResult
+                      .flatMap(
+                          registeredEvent ->
+                              contextWithUpdatedSession
                                   .pluginManager()
                                   .onEventCallback(contextWithUpdatedSession, registeredEvent)
-                                  .defaultIfEmpty(registeredEvent);
-                            })
-                        .toFlowable());
+                                  .defaultIfEmpty(registeredEvent))
+                      .toFlowable();
+                });
 
     // If beforeRunCallback returns content, emit it and skip agent
     Context capturedContext = Context.current();
@@ -619,6 +634,24 @@ private void copySessionStates(Session source, Session target) {
     target.state().putAll(source.state());
   }
 
+  private static final String IN_FLIGHT_APPENDED_EVENT_IDS_KEY =
+      "com.google.adk.internal.inFlightAppendedEventIds";
+
+  /**
+   * Returns the transient, per-invocation set of event ids appended by the flow but not yet
+   * consumed by the Runner, lazily creating it. Ids are added by the flow on append and removed
+   * here on consume, so this is hand-off state -- not a record of all persisted events.
+   */
+  @SuppressWarnings("unchecked")
+  private static Set<String> inFlightAppendedEventIds(InvocationContext invocationContext) {
+    return (Set<String>)
+        invocationContext
+            .callbackContextData()
+            .computeIfAbsent(
+                IN_FLIGHT_APPENDED_EVENT_IDS_KEY,
+                unusedKey -> ConcurrentHashMap.<String>newKeySet());
+  }
+
   /**
    * Creates an {@link InvocationContext} for a live (streaming) run.
    *

core/src/main/java/com/google/adk/sessions/SessionUtils.java

@@ -16,10 +16,12 @@
 
 package com.google.adk.sessions;
 
+import com.google.adk.events.Event;
 import com.google.common.collect.ImmutableList;
 import com.google.genai.types.Blob;
 import com.google.genai.types.Content;
 import com.google.genai.types.Part;
+import io.reactivex.rxjava3.core.Single;
 import java.util.ArrayList;
 import java.util.Base64;
 import java.util.List;
@@ -31,6 +33,27 @@ public final class SessionUtils {
 
   public SessionUtils() {}
 
+  /**
+   * Appends {@code event} via {@code service}, or just to {@code session.events()} when the session
+   * is partial (no {@code appName}). The partial-session bypass exists for unit tests that build
+   * {@code Session.builder(id).build()} and bypass {@code Runner}; most production services
+   * (including {@link InMemorySessionService}) {@code requireNonNull(appName)}. Production callers
+   * always pass fully-formed sessions and hit the unchanged {@code service.appendEvent} path.
+   */
+  public static Single<Event> safeAppendEvent(
+      BaseSessionService service, Session session, Event event) {
+    if (session.appName() == null) {
+      List<Event> events = session.events();
+      if (events != null) {
+        synchronized (events) {
+          events.add(event);
+        }
+      }
+      return Single.just(event);
+    }
+    return service.appendEvent(session, event);
+  }
+
   /** Base64-encodes inline blobs in content. */
   public static Content encodeContent(Content content) {
     List<Part> encodedParts = new ArrayList<>();

core/src/test/java/com/google/adk/agents/LlmAgentTest.java

@@ -39,7 +39,11 @@
 import com.google.adk.models.LlmRequest;
 import com.google.adk.models.LlmResponse;
 import com.google.adk.models.Model;
+import com.google.adk.sessions.BaseSessionService;
+import com.google.adk.sessions.GetSessionConfig;
 import com.google.adk.sessions.InMemorySessionService;
+import com.google.adk.sessions.ListEventsResponse;
+import com.google.adk.sessions.ListSessionsResponse;
 import com.google.adk.sessions.Session;
 import com.google.adk.telemetry.Tracing;
 import com.google.adk.testing.TestLlm;
@@ -58,11 +62,14 @@
 import io.opentelemetry.api.trace.Tracer;
 import io.opentelemetry.sdk.testing.junit4.OpenTelemetryRule;
 import io.opentelemetry.sdk.trace.data.SpanData;
+import io.reactivex.rxjava3.core.Completable;
 import io.reactivex.rxjava3.core.Flowable;
 import io.reactivex.rxjava3.core.Maybe;
 import io.reactivex.rxjava3.core.Single;
 import java.util.List;
+import java.util.Optional;
 import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ConcurrentMap;
 import java.util.concurrent.atomic.AtomicBoolean;
 import org.junit.After;
 import org.junit.Before;
@@ -184,6 +191,113 @@ public void testRun_withoutOutputKey_doesNotSaveState() {
     assertThat(events.get(0).actions().stateDelta()).isEmpty();
   }
 
+  /**
+   * Partial-session bypass: tests that build {@code Session.builder(id).build()} and call {@code
+   * agent.runAsync} directly (bypassing Runner) must not trip the {@code requireNonNull(appName)}
+   * in {@code InMemorySessionService.appendEvent}. The event is tracked in {@code session.events()}
+   * so subsequent steps see prior events. Surfaced originally by orcas {@code LlmAgentActionTest},
+   * dataworkeragent and asterix small tests.
+   */
+  @Test
+  public void testRun_partialSessionWithoutAppName_doesNotThrow() {
+    Content modelContent = Content.fromParts(Part.fromText("Agent Response"));
+    TestLlm testLlm = createTestLlm(createLlmResponse(modelContent));
+    LlmAgent agent = createTestAgentBuilder(testLlm).build();
+
+    Session partialSession = Session.builder("session-id").build();
+    InvocationContext invocationContext =
+        InvocationContext.builder()
+            .sessionService(new InMemorySessionService())
+            .agent(agent)
+            .session(partialSession)
+            .invocationId("invocation-id")
+            .runConfig(RunConfig.builder().build())
+            .userContent(Content.fromParts(Part.fromText("hello")))
+            .build();
+
+    List<Event> events = agent.runAsync(invocationContext).toList().blockingGet();
+
+    assertThat(events).hasSize(1);
+    assertThat(events.get(0).content()).hasValue(modelContent);
+    // Event tracked in session.events() so subsequent steps see prior events.
+    assertThat(partialSession.events()).hasSize(1);
+    assertThat(partialSession.events().get(0).id()).isEqualTo(events.get(0).id());
+  }
+
+  /**
+   * Mirrors ads-ux researchagent {@code GenerateReportFromSourcesActionTest}: an action drives
+   * {@code agent.runAsync(...).collect { ... }} directly (bypassing Runner) with a {@link
+   * BaseSessionService} stub whose {@code appendEvent} returns a sentinel empty Event. The original
+   * LLM-derived event must flow downstream with its content intact, not be swapped for the
+   * service's return value.
+   */
+  @Test
+  public void testRun_appendEventReturnsSentinel_originalEventFlowsDownstream() {
+    Content modelContent = Content.fromParts(Part.fromText("generated report content"));
+    TestLlm testLlm = createTestLlm(createLlmResponse(modelContent));
+    LlmAgent agent = createTestAgentBuilder(testLlm).build();
+
+    BaseSessionService sentinelReturningSessionService = new SentinelReturningSessionService();
+    // appName set so safeAppendEvent calls the service (we want to verify its return is ignored).
+    Session sessionWithAppName =
+        Session.builder("session-id").appName("test").userId("user").build();
+    InvocationContext invocationContext =
+        InvocationContext.builder()
+            .sessionService(sentinelReturningSessionService)
+            .agent(agent)
+            .session(sessionWithAppName)
+            .invocationId("invocation-id")
+            .runConfig(RunConfig.builder().build())
+            .userContent(Content.fromParts(Part.fromText("hello")))
+            .build();
+
+    List<Event> events = agent.runAsync(invocationContext).toList().blockingGet();
+
+    assertThat(events).hasSize(1);
+    // Must be the original LLM-derived event, not the sentinel returned by appendEvent.
+    assertThat(events.get(0).content()).hasValue(modelContent);
+  }
+
+  /**
+   * Stub returning an empty Event from {@code appendEvent}, mirroring the shape used by ads-ux
+   * researchagent and nbu paisa tests ({@code
+   * Mockito.when(...).thenReturn(Event.builder().build())}).
+   */
+  private static final class SentinelReturningSessionService implements BaseSessionService {
+    @Override
+    public Single<Session> createSession(
+        String appName, String userId, ConcurrentMap<String, Object> state, String sessionId) {
+      return Single.just(Session.builder("session-id").build());
+    }
+
+    @Override
+    public Maybe<Session> getSession(
+        String appName, String userId, String sessionId, Optional<GetSessionConfig> config) {
+      return Maybe.just(Session.builder(sessionId).build());
+    }
+
+    @Override
+    public Single<ListSessionsResponse> listSessions(String appName, String userId) {
+      return Single.just(ListSessionsResponse.builder().build());
+    }
+
+    @Override
+    public Completable deleteSession(String appName, String userId, String sessionId) {
+      return Completable.complete();
+    }
+
+    @Override
+    public Single<ListEventsResponse> listEvents(String appName, String userId, String sessionId) {
+      return Single.just(ListEventsResponse.builder().build());
+    }
+
+    @Override
+    public Single<Event> appendEvent(Session session, Event event) {
+      // Sentinel return value, mirroring downstream test mocks.
+      return Single.just(Event.builder().build());
+    }
+  }
+
   @Test
   public void run_withToolsAndMaxSteps_stopsAfterMaxSteps() {
     ImmutableMap<String, Object> echoArgs = ImmutableMap.of("arg", "value");