From 5012ad98c79ef4818415c7218ef1622fbdd60106 Mon Sep 17 00:00:00 2001
From: Damian Momot <damianmomot@google.com>
Date: Tue, 26 May 2026 03:11:30 -0700
Subject: [PATCH] fix: pre-merge stateDelta before onUserMessageCallback in
 Runner

Previously, stateDelta passed to runAsync was merged into the session
only after pluginManager.onUserMessageCallback was invoked, so plugins
observed a stale session state during the user-message callback.

PiperOrigin-RevId: 921350398
---
 .../java/com/google/adk/runner/Runner.java    |  6 ++++
 .../com/google/adk/runner/RunnerTest.java     | 36 +++++++++++++++++++
 2 files changed, 42 insertions(+)

diff --git a/core/src/main/java/com/google/adk/runner/Runner.java b/core/src/main/java/com/google/adk/runner/Runner.java
index 44a281f72..26a523fdd 100644
--- a/core/src/main/java/com/google/adk/runner/Runner.java
+++ b/core/src/main/java/com/google/adk/runner/Runner.java
@@ -485,6 +485,12 @@ protected Flowable<Event> runAsyncImpl(
               BaseAgent rootAgent = this.agent;
               String invocationId = InvocationContext.newInvocationContextId();
 
+              // Pre-merge stateDelta so onUserMessageCallback can access it.
+              // Safe: session is a copy; persistence still happens via appendNewMessageToSession.
+              if (stateDelta != null && !stateDelta.isEmpty()) {
+                stateDelta.forEach((key, value) -> session.state().put(key, value));
+              }
+
               // Create initial context
               InvocationContext initialContext =
                   newInvocationContextBuilder(session)
diff --git a/core/src/test/java/com/google/adk/runner/RunnerTest.java b/core/src/test/java/com/google/adk/runner/RunnerTest.java
index ff75c97b0..3abfbdc20 100644
--- a/core/src/test/java/com/google/adk/runner/RunnerTest.java
+++ b/core/src/test/java/com/google/adk/runner/RunnerTest.java
@@ -1122,6 +1122,42 @@ public void beforeRunCallback_withStateDelta_seesMergedState() {
     assertThat(sessionInCallback.state()).containsEntry("number", 123);
   }
 
+  @Test
+  public void onUserMessageCallback_withStateDelta_seesMergedState() {
+    // Snapshot the session state *inside* the callback, otherwise the assertion would
+    // observe the post-runAsync state which is mutated by appendEvent regardless of whether
+    // the pre-merge in Runner is applied.
+    AtomicReference<ConcurrentHashMap<String, Object>> stateInCallback = new AtomicReference<>();
+    when(plugin.onUserMessageCallback(any(), any()))
+        .thenAnswer(
+            invocation -> {
+              InvocationContext ctx = invocation.getArgument(0);
+              stateInCallback.set(new ConcurrentHashMap<>(ctx.session().state()));
+              return Maybe.empty();
+            });
+
+    ImmutableMap<String, Object> stateDelta =
+        ImmutableMap.of("callback_key", "callback_value", "number", 123);
+
+    var unused =
+        runner
+            .runAsync(
+                "user",
+                session.id(),
+                createContent("test with state"),
+                RunConfig.builder().build(),
+                stateDelta)
+            .toList()
+            .blockingGet();
+
+    // Verify onUserMessageCallback was called
+    verify(plugin).onUserMessageCallback(any(), any());
+
+    // Verify state delta was merged before onUserMessageCallback was invoked
+    assertThat(stateInCallback.get()).containsEntry("callback_key", "callback_value");
+    assertThat(stateInCallback.get()).containsEntry("number", 123);
+  }
+
   @Test
   public void runAsync_ensureEventsAreAppendedInOrder() throws Exception {
     Event event1 = TestUtils.createEvent("1");