diff --git a/core/src/main/java/com/google/adk/events/EventActions.java b/core/src/main/java/com/google/adk/events/EventActions.java index cde23c10e..063a56a44 100644 --- a/core/src/main/java/com/google/adk/events/EventActions.java +++ b/core/src/main/java/com/google/adk/events/EventActions.java @@ -21,6 +21,7 @@ import com.google.adk.JsonBaseModel; import com.google.adk.sessions.State; import com.google.errorprone.annotations.CanIgnoreReturnValue; +import java.util.HashMap; import java.util.HashSet; import java.util.Map; import java.util.Objects; @@ -44,6 +45,7 @@ public class EventActions extends JsonBaseModel { private ConcurrentMap> requestedAuthConfigs; private ConcurrentMap requestedToolConfirmations; private boolean endOfAgent; + private @Nullable Map agentState; private @Nullable EventCompaction compaction; /** Default constructor for Jackson. */ @@ -66,6 +68,7 @@ private EventActions(Builder builder) { this.requestedAuthConfigs = builder.requestedAuthConfigs; this.requestedToolConfirmations = builder.requestedToolConfirmations; this.endOfAgent = builder.endOfAgent; + this.agentState = builder.agentState; this.compaction = builder.compaction; } @@ -192,6 +195,19 @@ public void setEndInvocation(boolean endInvocation) { this.endOfAgent = endInvocation; } + /** + * The checkpointed state of the authoring agent at this event, used for session resumability. + * Only set by ADK workflow/agent machinery on resumable invocations. + */ + @JsonProperty("agentState") + public Optional> agentState() { + return Optional.ofNullable(agentState); + } + + public void setAgentState(@Nullable Map agentState) { + this.agentState = agentState; + } + @JsonProperty("compaction") public Optional compaction() { return Optional.ofNullable(compaction); @@ -226,6 +242,7 @@ public boolean equals(Object o) { && Objects.equals(requestedAuthConfigs, that.requestedAuthConfigs) && Objects.equals(requestedToolConfirmations, that.requestedToolConfirmations) && (endOfAgent == that.endOfAgent) + && Objects.equals(agentState, that.agentState) && Objects.equals(compaction, that.compaction); } @@ -241,6 +258,7 @@ public int hashCode() { requestedAuthConfigs, requestedToolConfirmations, endOfAgent, + agentState, compaction); } @@ -255,6 +273,7 @@ public static class Builder { private ConcurrentMap> requestedAuthConfigs; private ConcurrentMap requestedToolConfirmations; private boolean endOfAgent = false; + private @Nullable Map agentState; private @Nullable EventCompaction compaction; public Builder() { @@ -276,6 +295,8 @@ private Builder(EventActions eventActions) { this.requestedToolConfirmations = new ConcurrentHashMap<>(eventActions.requestedToolConfirmations()); this.endOfAgent = eventActions.endOfAgent; + this.agentState = + eventActions.agentState == null ? null : new HashMap<>(eventActions.agentState); this.compaction = eventActions.compaction; } @@ -376,6 +397,13 @@ public Builder endInvocation(boolean endInvocation) { return this; } + @CanIgnoreReturnValue + @JsonProperty("agentState") + public Builder agentState(@Nullable Map agentState) { + this.agentState = agentState; + return this; + } + @CanIgnoreReturnValue @JsonProperty("compaction") public Builder compaction(@Nullable EventCompaction value) { @@ -394,6 +422,7 @@ public Builder merge(EventActions other) { this.requestedAuthConfigs.putAll(other.requestedAuthConfigs()); this.requestedToolConfirmations.putAll(other.requestedToolConfirmations()); this.endOfAgent = this.endOfAgent || other.endOfAgent(); + other.agentState().ifPresent(this::agentState); other.compaction().ifPresent(this::compaction); return this; } diff --git a/core/src/test/java/com/google/adk/events/EventActionsTest.java b/core/src/test/java/com/google/adk/events/EventActionsTest.java index c5949caf7..77fb805d4 100644 --- a/core/src/test/java/com/google/adk/events/EventActionsTest.java +++ b/core/src/test/java/com/google/adk/events/EventActionsTest.java @@ -111,6 +111,61 @@ public void merge_mergesAllFields() { assertThat(merged.compaction()).hasValue(COMPACTION); } + @Test + public void agentState_roundTripsThroughToBuilder() { + EventActions actions = + EventActions.builder().agentState(ImmutableMap.of("current_sub_agent", "b")).build(); + + EventActions rebuilt = actions.toBuilder().build(); + + assertThat(rebuilt).isEqualTo(actions); + assertThat(rebuilt.agentState()).hasValue(ImmutableMap.of("current_sub_agent", "b")); + } + + @Test + public void agentState_roundTripsThroughJson() { + EventActions actions = + EventActions.builder().agentState(ImmutableMap.of("times_looped", 2)).build(); + + EventActions deserialized = EventActions.fromJsonString(actions.toJson(), EventActions.class); + + assertThat(deserialized.agentState()).isPresent(); + assertThat(deserialized.agentState().get()).containsEntry("times_looped", 2); + } + + @Test + public void agentState_absentByDefault_andOmittedFromJson() { + EventActions actions = EventActions.builder().build(); + + assertThat(actions.agentState()).isEmpty(); + // Kept out of the serialized form so pre-existing events stay byte-identical. + assertThat(actions.toJson()).doesNotContain("agentState"); + } + + @Test + public void merge_agentState_lastWins() { + EventActions first = + EventActions.builder().agentState(ImmutableMap.of("current_sub_agent", "a")).build(); + EventActions second = + EventActions.builder().agentState(ImmutableMap.of("current_sub_agent", "b")).build(); + + EventActions merged = first.toBuilder().merge(second).build(); + + assertThat(merged.agentState()).hasValue(ImmutableMap.of("current_sub_agent", "b")); + } + + @Test + public void merge_agentState_disjointKeys_replacesWholeMap() { + // agentState is a single checkpoint payload: merge replaces it wholesale (last-wins) rather + // than deep-merging keys. + EventActions first = EventActions.builder().agentState(ImmutableMap.of("a", 1)).build(); + EventActions second = EventActions.builder().agentState(ImmutableMap.of("b", 2)).build(); + + EventActions merged = first.toBuilder().merge(second).build(); + + assertThat(merged.agentState()).hasValue(ImmutableMap.of("b", 2)); + } + @Test public void merge_endOfAgentIsOrderIndependent() { // A tool that ends the invocation, and one that leaves the flag at its default false. Folding