Skip to content

Commit 00d6d30

Browse files
google-genai-botcopybara-github
authored andcommitted
feat: Forward state delta to parent session
PiperOrigin-RevId: 850221284
1 parent 9434949 commit 00d6d30

File tree

2 files changed

+38
-0
lines changed

2 files changed

+38
-0
lines changed

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

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -144,6 +144,13 @@ public Single<Map<String, Object>> runAsync(Map<String, Object> args, ToolContex
144144
Event lastEvent = optionalLastEvent.get();
145145
Optional<String> outputText = lastEvent.content().map(Content::text);
146146

147+
// Forward state delta to parent session.
148+
if (lastEvent.actions() != null
149+
&& lastEvent.actions().stateDelta() != null
150+
&& !lastEvent.actions().stateDelta().isEmpty()) {
151+
toolContext.state().putAll(lastEvent.actions().stateDelta());
152+
}
153+
147154
if (outputText.isEmpty()) {
148155
return ImmutableMap.of();
149156
}

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

Lines changed: 31 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@
2121
import static com.google.common.truth.Truth.assertThat;
2222
import static org.junit.Assert.assertThrows;
2323

24+
import com.google.adk.agents.Callbacks.AfterAgentCallback;
2425
import com.google.adk.agents.ConfigAgentUtils.ConfigurationException;
2526
import com.google.adk.agents.InvocationContext;
2627
import com.google.adk.agents.LlmAgent;
@@ -35,6 +36,7 @@
3536
import com.google.genai.types.Part;
3637
import com.google.genai.types.Schema;
3738
import io.reactivex.rxjava3.core.Flowable;
39+
import io.reactivex.rxjava3.core.Maybe;
3840
import java.util.Map;
3941
import java.util.Optional;
4042
import org.junit.Test;
@@ -421,6 +423,35 @@ public void call_withoutInputSchema_requestIsSentToAgent() throws Exception {
421423
.containsExactly(Content.fromParts(Part.fromText("magic")));
422424
}
423425

426+
@Test
427+
public void call_withStateDeltaInResponse_propagatesStateDelta() throws Exception {
428+
AfterAgentCallback afterAgentCallback =
429+
(callbackContext) -> {
430+
callbackContext.state().put("test_key", "test_value");
431+
return Maybe.empty();
432+
};
433+
TestLlm testLlm =
434+
createTestLlm(
435+
LlmResponse.builder()
436+
.content(Content.fromParts(Part.fromText("test response")))
437+
.build());
438+
LlmAgent testAgent =
439+
createTestAgentBuilder(testLlm)
440+
.name("agent name")
441+
.description("agent description")
442+
.afterAgentCallback(afterAgentCallback)
443+
.build();
444+
AgentTool agentTool = AgentTool.create(testAgent);
445+
ToolContext toolContext = createToolContext(testAgent);
446+
447+
assertThat(toolContext.state()).doesNotContainKey("test_key");
448+
449+
Map<String, Object> unused =
450+
agentTool.runAsync(ImmutableMap.of("request", "magic"), toolContext).blockingGet();
451+
452+
assertThat(toolContext.state()).containsEntry("test_key", "test_value");
453+
}
454+
424455
private static ToolContext createToolContext(LlmAgent agent) {
425456
return ToolContext.builder(
426457
new InvocationContext(

0 commit comments

Comments
 (0)