From 6a18456f2a780b7707c8bb5b9e2b681a58606a98 Mon Sep 17 00:00:00 2001
From: Google Team Member <genai-sdk-bot@google.com>
Date: Sun, 28 Dec 2025 10:03:13 -0800
Subject: [PATCH] feat: Forward state delta to parent session

PiperOrigin-RevId: 849775771
---
 .../java/com/google/adk/tools/AgentTool.java  |  7 +++++
 .../com/google/adk/tools/AgentToolTest.java   | 31 +++++++++++++++++++
 2 files changed, 38 insertions(+)

diff --git a/core/src/main/java/com/google/adk/tools/AgentTool.java b/core/src/main/java/com/google/adk/tools/AgentTool.java
index c31715fef..653939009 100644
--- a/core/src/main/java/com/google/adk/tools/AgentTool.java
+++ b/core/src/main/java/com/google/adk/tools/AgentTool.java
@@ -144,6 +144,13 @@ public Single<Map<String, Object>> runAsync(Map<String, Object> args, ToolContex
               Event lastEvent = optionalLastEvent.get();
               Optional<String> outputText = lastEvent.content().map(Content::text);
 
+              // Forward state delta to parent session.
+              if (lastEvent.actions() != null
+                  && lastEvent.actions().stateDelta() != null
+                  && !lastEvent.actions().stateDelta().isEmpty()) {
+                toolContext.state().putAll(lastEvent.actions().stateDelta());
+              }
+
               if (outputText.isEmpty()) {
                 return ImmutableMap.of();
               }
diff --git a/core/src/test/java/com/google/adk/tools/AgentToolTest.java b/core/src/test/java/com/google/adk/tools/AgentToolTest.java
index d0fb5cbf6..e5318f87f 100644
--- a/core/src/test/java/com/google/adk/tools/AgentToolTest.java
+++ b/core/src/test/java/com/google/adk/tools/AgentToolTest.java
@@ -21,6 +21,7 @@
 import static com.google.common.truth.Truth.assertThat;
 import static org.junit.Assert.assertThrows;
 
+import com.google.adk.agents.Callbacks.AfterAgentCallback;
 import com.google.adk.agents.ConfigAgentUtils.ConfigurationException;
 import com.google.adk.agents.InvocationContext;
 import com.google.adk.agents.LlmAgent;
@@ -35,6 +36,7 @@
 import com.google.genai.types.Part;
 import com.google.genai.types.Schema;
 import io.reactivex.rxjava3.core.Flowable;
+import io.reactivex.rxjava3.core.Maybe;
 import java.util.Map;
 import java.util.Optional;
 import org.junit.Test;
@@ -421,6 +423,35 @@ public void call_withoutInputSchema_requestIsSentToAgent() throws Exception {
         .containsExactly(Content.fromParts(Part.fromText("magic")));
   }
 
+  @Test
+  public void call_withStateDeltaInResponse_propagatesStateDelta() throws Exception {
+    AfterAgentCallback afterAgentCallback =
+        (callbackContext) -> {
+          callbackContext.state().put("test_key", "test_value");
+          return Maybe.empty();
+        };
+    TestLlm testLlm =
+        createTestLlm(
+            LlmResponse.builder()
+                .content(Content.fromParts(Part.fromText("test response")))
+                .build());
+    LlmAgent testAgent =
+        createTestAgentBuilder(testLlm)
+            .name("agent name")
+            .description("agent description")
+            .afterAgentCallback(afterAgentCallback)
+            .build();
+    AgentTool agentTool = AgentTool.create(testAgent);
+    ToolContext toolContext = createToolContext(testAgent);
+
+    assertThat(toolContext.state()).doesNotContainKey("test_key");
+
+    Map<String, Object> unused =
+        agentTool.runAsync(ImmutableMap.of("request", "magic"), toolContext).blockingGet();
+
+    assertThat(toolContext.state()).containsEntry("test_key", "test_value");
+  }
+
   private static ToolContext createToolContext(LlmAgent agent) {
     return ToolContext.builder(
             new InvocationContext(