feat: Forward state delta from all events to parent session instead of just the last event

PiperOrigin-RevId: 900231745

Author: google-genai-bot

core/src/main/java/com/google/adk/tools/AgentTool.java

@@ -166,6 +166,14 @@ public Single<Map<String, Object>> runAsync(Map<String, Object> args, ToolContex
         .sessionService()
         .createSession(toolContext.agentName(), "tmp-user", toolContext.state(), null)
         .flatMapPublisher(session -> runner.runAsync(session.userId(), session.id(), content))
+        .doOnNext(
+            event -> {
+              if (event.actions() != null
+                  && event.actions().stateDelta() != null
+                  && !event.actions().stateDelta().isEmpty()) {
+                updateState(event.actions().stateDelta(), toolContext.state());
+              }
+            })
         .lastElement()
         .map(Optional::of)
         .defaultIfEmpty(Optional.empty())
@@ -177,13 +185,6 @@ 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()) {
-                updateState(lastEvent.actions().stateDelta(), toolContext.state());
-              }
-
               if (outputText.isEmpty()) {
                 return ImmutableMap.of();
               }

core/src/test/java/com/google/adk/tools/AgentToolTest.java

@@ -495,6 +495,46 @@ public void call_withSkipSummarizationAndStateDelta_propagatesStateAndSetsSkipSu
     assertThat(toolContext.actions().skipSummarization()).hasValue(true);
   }
 
+  @Test
+  public void call_withMultipleStateDeltasInResponse_propagatesAllStateDeltas() throws Exception {
+    AfterAgentCallback firstCallback =
+        (callbackContext) -> {
+          callbackContext.state().put("key1", "val1");
+          return Maybe.empty();
+        };
+    AfterAgentCallback secondCallback =
+        (callbackContext) -> {
+          callbackContext.state().put("key2", "val2");
+          return Maybe.empty();
+        };
+    LlmAgent firstAgent =
+        createTestAgentBuilder(createTestLlm(LlmResponse.builder().build()))
+            .name("first_agent")
+            .afterAgentCallback(firstCallback)
+            .build();
+    LlmAgent secondAgent =
+        createTestAgentBuilder(createTestLlm(LlmResponse.builder().build()))
+            .name("second_agent")
+            .afterAgentCallback(secondCallback)
+            .build();
+    SequentialAgent sequentialAgent =
+        SequentialAgent.builder()
+            .name("sequence")
+            .description("Process the query through multiple steps")
+            .subAgents(ImmutableList.of(firstAgent, secondAgent))
+            .build();
+    ToolContext toolContext = createToolContext(sequentialAgent);
+    assertThat(toolContext.state()).isEmpty();
+
+    Map<String, Object> unused =
+        AgentTool.create(sequentialAgent)
+            .runAsync(ImmutableMap.of("request", "test"), toolContext)
+            .blockingGet();
+
+    assertThat(toolContext.state()).containsEntry("key1", "val1");
+    assertThat(toolContext.state()).containsEntry("key2", "val2");
+  }
+
   @Test
   public void
       declaration_sequentialAgentWithFirstSubAgentInputSchema_returnsDeclarationWithSchema() {