fix: Fix ADK Runner race condition for sequential tool execution

Ensure that events are appended to the session and processed sequentially before proceeding to the next step in BaseLlmFlow.

PiperOrigin-RevId: 896524076

Author: google-genai-bot

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

@@ -461,14 +461,31 @@ public Flowable<Event> run(InvocationContext invocationContext) {
 
   private Flowable<Event> run(
       Context spanContext, InvocationContext invocationContext, int stepsCompleted) {
-    Flowable<Event> currentStepEvents = runOneStep(spanContext, invocationContext).cache();
+    Flowable<Event> currentStepEvents = runOneStep(spanContext, invocationContext);
+
+    Flowable<Event> processedEvents =
+        currentStepEvents
+            .concatMap(
+                event ->
+                    invocationContext
+                        .sessionService()
+                        .appendEvent(invocationContext.session(), event)
+                        .flatMap(
+                            registeredEvent ->
+                                invocationContext
+                                    .pluginManager()
+                                    .onEventCallback(invocationContext, registeredEvent)
+                                    .defaultIfEmpty(registeredEvent))
+                        .toFlowable())
+            .cache();
+
     if (stepsCompleted + 1 >= maxSteps) {
       logger.debug("Ending flow execution because max steps reached.");
-      return currentStepEvents;
+      return processedEvents;
     }
 
-    return currentStepEvents.concatWith(
-        currentStepEvents
+    return processedEvents.concatWith(
+        processedEvents
             .toList()
             .flatMapPublisher(
                 eventList -> {

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

@@ -570,19 +570,29 @@ 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
-                                  .pluginManager()
-                                  .onEventCallback(contextWithUpdatedSession, registeredEvent)
-                                  .defaultIfEmpty(registeredEvent);
-                            })
-                        .toFlowable());
+                agentEvent -> {
+                  // TODO: remove this hack after deprecating runAsync with Session.
+                  copySessionStates(updatedSession, initialContext.session());
+
+                  // TODO: b/502182243 - Investigate if appendEvent should be made idempotent in
+                  // SessionService to avoid this check.
+                  if (agentEvent.id() != null
+                      && updatedSession.events().stream()
+                          .anyMatch(e -> agentEvent.id().equals(e.id()))) {
+                    // Already appended (e.g. by BaseLlmFlow). Still apply the hack.
+                    return Flowable.just(agentEvent);
+                  }
+                  return this.sessionService
+                      .appendEvent(updatedSession, agentEvent)
+                      .flatMap(
+                          registeredEvent -> {
+                            return contextWithUpdatedSession
+                                .pluginManager()
+                                .onEventCallback(contextWithUpdatedSession, registeredEvent)
+                                .defaultIfEmpty(registeredEvent);
+                          })
+                      .toFlowable();
+                });
 
     // If beforeRunCallback returns content, emit it and skip agent
     Context capturedContext = Context.current();

core/src/test/java/com/google/adk/runner/RunnerTest.java

@@ -46,9 +46,12 @@
 import com.google.adk.artifacts.BaseArtifactService;
 import com.google.adk.events.Event;
 import com.google.adk.flows.llmflows.Functions;
+import com.google.adk.models.LlmRequest;
 import com.google.adk.models.LlmResponse;
 import com.google.adk.plugins.BasePlugin;
 import com.google.adk.sessions.BaseSessionService;
+import com.google.adk.sessions.GetSessionConfig;
+import com.google.adk.sessions.InMemorySessionService;
 import com.google.adk.sessions.Session;
 import com.google.adk.sessions.SessionKey;
 import com.google.adk.summarizer.EventsCompactionConfig;
@@ -588,12 +591,22 @@ public void onToolErrorCallback_error() {
   @Test
   public void onEventCallback_success() {
     when(plugin.onEventCallback(any(), any()))
-        .thenReturn(Maybe.just(TestUtils.createEvent("form plugin")));
+        .thenAnswer(
+            invocation -> {
+              Event event = invocation.getArgument(1);
+              return Maybe.just(
+                  Event.builder()
+                      .id(event.id())
+                      .invocationId(event.invocationId())
+                      .author("model")
+                      .content(createContent("from plugin"))
+                      .build());
+            });
 
     List<Event> events =
         runner.runAsync("user", session.id(), createContent("from user")).toList().blockingGet();
 
-    assertThat(simplifyEvents(events)).containsExactly("author: content for event form plugin");
+    assertThat(simplifyEvents(events)).containsExactly("model: from plugin");
 
     verify(plugin).onEventCallback(any(), any());
   }
@@ -1686,4 +1699,105 @@ public void runner_executesSaveArtifactFlow() {
     // agent was run
     assertThat(simplifyEvents(events.values())).containsExactly("test agent: from llm");
   }
+
+  @Test
+  public void runAsync_ensuresSequentialConsistencyForTools() {
+    // Arrange
+    TestLlm testLlm =
+        createTestLlm(
+            createFunctionCallLlmResponse("call_1", "tool1", ImmutableMap.of("arg", "value1")),
+            createTextLlmResponse("Final response"));
+
+    LlmAgent agent =
+        createTestAgentBuilder(testLlm)
+            .tools(
+                ImmutableList.of(
+                    FunctionTool.create(RaceConditionTools.class, "tool1"),
+                    FunctionTool.create(RaceConditionTools.class, "tool2")))
+            .build();
+
+    BaseSessionService delegate = new InMemorySessionService();
+    BaseSessionService delayedSessionService = createDelayedSessionService(delegate, 100);
+
+    Runner runner =
+        Runner.builder()
+            .app(App.builder().name("test").rootAgent(agent).build())
+            .sessionService(delayedSessionService)
+            .build();
+    Session session = runner.sessionService().createSession("test", "user").blockingGet();
+
+    // Act
+    var unused =
+        runner
+            .runAsync("user", session.id(), Content.fromParts(Part.fromText("start")))
+            .toList()
+            .blockingGet();
+
+    // Assert
+    ImmutableList<LlmRequest> requests = ImmutableList.copyOf(testLlm.getRequests());
+    assertThat(requests).hasSize(2);
+
+    // Second request should contain the result of tool1
+    LlmRequest secondRequest = requests.get(1);
+    List<Content> history = secondRequest.contents();
+
+    boolean foundToolResponse =
+        history.stream()
+            .flatMap(content -> content.parts().stream().flatMap(List::stream))
+            .filter(part -> part.functionResponse().isPresent())
+            .map(part -> part.functionResponse().get())
+            .anyMatch(
+                response ->
+                    response.name().orElse("").equals("tool1")
+                        && response
+                            .response()
+                            .orElse(null)
+                            .equals(ImmutableMap.of("result", "result_value1")));
+
+    assertThat(foundToolResponse).isTrue();
+  }
+
+  private static BaseSessionService createDelayedSessionService(
+      BaseSessionService delegate, long delayMs) {
+    BaseSessionService delayedSessionService = mock(BaseSessionService.class);
+    when(delayedSessionService.createSession(anyString(), anyString(), any(), anyString()))
+        .thenAnswer(
+            inv ->
+                delegate.createSession(
+                    inv.getArgument(0),
+                    inv.getArgument(1),
+                    inv.getArgument(2),
+                    inv.getArgument(3)));
+    when(delayedSessionService.createSession(anyString(), anyString()))
+        .thenAnswer(
+            inv ->
+                delegate.createSession((String) inv.getArgument(0), (String) inv.getArgument(1)));
+    when(delayedSessionService.getSession(anyString(), anyString(), anyString(), any()))
+        .thenAnswer(
+            inv ->
+                delegate.getSession(
+                    (String) inv.getArgument(0),
+                    (String) inv.getArgument(1),
+                    (String) inv.getArgument(2),
+                    (Optional<GetSessionConfig>) inv.getArgument(3)));
+    when(delayedSessionService.appendEvent(any(), any()))
+        .thenAnswer(
+            inv ->
+                delegate
+                    .appendEvent(inv.getArgument(0), inv.getArgument(1))
+                    .delay(delayMs, MILLISECONDS));
+    return delayedSessionService;
+  }
+
+  public static class RaceConditionTools {
+    private RaceConditionTools() {}
+
+    public static ImmutableMap<String, Object> tool1(String arg) {
+      return ImmutableMap.of("result", "result_" + arg);
+    }
+
+    public static ImmutableMap<String, Object> tool2(String input) {
+      return ImmutableMap.of("status", "received_" + input);
+    }
+  }
 }