fix: Fix ADK Runner race condition for sequential tool execution

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

@@ -565,24 +565,29 @@ private Flowable<Event> runAgentWithUpdatedSession(
                         .build());
 
     // Agent execution
-    Flowable<Event> agentEvents =
-        contextWithUpdatedSession
-            .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());
+    Flowable<Event> agentEvents;
+    if (contextWithUpdatedSession.agent() instanceof LlmAgent) {
+      agentEvents = contextWithUpdatedSession.agent().runAsync(contextWithUpdatedSession);
+    } else {
+      agentEvents =
+          contextWithUpdatedSession
+              .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());
+    }
 
     // If beforeRunCallback returns content, emit it and skip agent
     Context capturedContext = Context.current();

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

@@ -46,6 +46,7 @@
 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;
@@ -1686,4 +1687,67 @@ 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();
+
+    Runner runner =
+        Runner.builder().app(App.builder().name("test").rootAgent(agent).build()).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
+    List<LlmRequest> requests = 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 = false;
+    for (Content content : history) {
+      for (Part part : content.parts().get()) {
+        if (part.functionResponse().isPresent()
+            && part.functionResponse().get().name().isPresent()
+            && "tool1".equals(part.functionResponse().get().name().get())) {
+          foundToolResponse = true;
+          assertThat(part.functionResponse().get().response().isPresent()).isTrue();
+          assertThat(part.functionResponse().get().response().get())
+              .isEqualTo(ImmutableMap.of("result", "result_value1"));
+        }
+      }
+    }
+    assertThat(foundToolResponse).isTrue();
+  }
+
+  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);
+    }
+  }
 }