fix: Fix race condition and stale session in ADK Runner

PiperOrigin-RevId: 896045084

Author: google-genai-bot

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

@@ -45,6 +45,7 @@
 import com.google.adk.utils.CollectionUtils;
 import com.google.common.base.Preconditions;
 import com.google.common.collect.ImmutableList;
+import com.google.common.collect.MapMaker;
 import com.google.errorprone.annotations.CanIgnoreReturnValue;
 import com.google.genai.types.AudioTranscriptionConfig;
 import com.google.genai.types.Content;
@@ -57,13 +58,16 @@
 import io.reactivex.rxjava3.core.Flowable;
 import io.reactivex.rxjava3.core.Maybe;
 import io.reactivex.rxjava3.core.Single;
+import io.reactivex.rxjava3.subjects.CompletableSubject;
+import java.time.Instant;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Collections;
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
 import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ConcurrentMap;
 import org.jspecify.annotations.Nullable;
 
 /** The main class for the GenAI Agents runner. */
@@ -76,6 +80,8 @@ public class Runner {
   private final PluginManager pluginManager;
   @Nullable private final EventsCompactionConfig eventsCompactionConfig;
   @Nullable private final ContextCacheConfig contextCacheConfig;
+  private final ConcurrentMap<String, Completable> activeSessionCompletables =
+      new MapMaker().weakValues().makeMap();
 
   /** Builder for {@link Runner}. */
   public static class Builder {
@@ -380,25 +386,56 @@ public Flowable<Event> runAsync(
       Content newMessage,
       RunConfig runConfig,
       @Nullable Map<String, Object> stateDelta) {
+    Flowable<Event> result =
+        Flowable.defer(
+                () ->
+                    this.sessionService
+                        .getSession(appName, userId, sessionId, Optional.empty())
+                        .switchIfEmpty(
+                            Single.defer(
+                                () -> {
+                                  if (runConfig.autoCreateSession()) {
+                                    return this.sessionService.createSession(
+                                        appName, userId, (Map<String, Object>) null, sessionId);
+                                  }
+                                  return Single.error(
+                                      new IllegalArgumentException(
+                                          String.format(
+                                              "Session not found: %s for user %s",
+                                              sessionId, userId)));
+                                }))
+                        .flatMapPublisher(
+                            session ->
+                                this.runAsyncImpl(session, newMessage, runConfig, stateDelta)))
+            .compose(Tracing.trace("invocation"));
+
     return Flowable.defer(
-            () ->
-                this.sessionService
-                    .getSession(appName, userId, sessionId, Optional.empty())
-                    .switchIfEmpty(
-                        Single.defer(
-                            () -> {
-                              if (runConfig.autoCreateSession()) {
-                                return this.sessionService.createSession(
-                                    appName, userId, (Map<String, Object>) null, sessionId);
-                              }
-                              return Single.error(
-                                  new IllegalArgumentException(
-                                      String.format(
-                                          "Session not found: %s for user %s", sessionId, userId)));
-                            }))
-                    .flatMapPublisher(
-                        session -> this.runAsyncImpl(session, newMessage, runConfig, stateDelta)))
-        .compose(Tracing.trace("invocation"));
+        () -> {
+          CompletableSubject requestCompletion = CompletableSubject.create();
+
+          Completable[] previousHolder = new Completable[1];
+
+          activeSessionCompletables.compute(
+              sessionId,
+              (key, current) -> {
+                previousHolder[0] = current;
+                return requestCompletion;
+              });
+
+          Completable previous = previousHolder[0];
+
+          Flowable<Event> execution =
+              result.doFinally(
+                  () -> {
+                    requestCompletion.onComplete();
+                    activeSessionCompletables.remove(sessionId, requestCompletion);
+                  });
+
+          if (previous == null) {
+            return execution;
+          }
+          return previous.onErrorComplete().andThen(execution);
+        });
   }
 
   /** See {@link #runAsync(String, String, Content, RunConfig, Map)}. */
@@ -540,6 +577,8 @@ private Flowable<Event> runAgentWithUpdatedSession(
                             registeredEvent -> {
                               // TODO: remove this hack after deprecating runAsync with Session.
                               copySessionStates(updatedSession, initialContext.session());
+                              updatedSession.lastUpdateTime(
+                                  Instant.ofEpochMilli(registeredEvent.timestamp()));
                               return contextWithUpdatedSession
                                   .pluginManager()
                                   .onEventCallback(contextWithUpdatedSession, registeredEvent)
@@ -740,6 +779,9 @@ private BaseAgent findAgentToRun(Session session, BaseAgent rootAgent) {
 
     for (Event event : events) {
       String author = event.author();
+      if (author == null) {
+        continue;
+      }
       if (author.equals("user")) {
         continue;
       }

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

@@ -123,7 +123,7 @@ public Builder userId(String userId) {
     @CanIgnoreReturnValue
     @JsonProperty("events")
     public Builder events(List<Event> events) {
-      this.events = Collections.synchronizedList(events);
+      this.events = Collections.synchronizedList(new ArrayList<>(events));
       return this;
     }