feat: update State constructors to accept general Map types

google-genai-bot · copybara-github · commit c6fdb63c92e2 · 2026-03-11T02:34:33.000-07:00
PiperOrigin-RevId: 881890323
diff --git a/core/src/main/java/com/google/adk/sessions/State.java b/core/src/main/java/com/google/adk/sessions/State.java
@@ -24,6 +24,7 @@
 import java.util.Set;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.ConcurrentMap;
+import javax.annotation.Nullable;
 
 /** A {@link State} object that also keeps track of the changes to the state. */
 @SuppressWarnings("ShouldNotSubclass")
@@ -39,13 +40,22 @@ public final class State implements ConcurrentMap<String, Object> {
   private final ConcurrentMap<String, Object> state;
   private final ConcurrentMap<String, Object> delta;
 
-  public State(ConcurrentMap<String, Object> state) {
-    this(state, new ConcurrentHashMap<>());
-  }
-
-  public State(ConcurrentMap<String, Object> state, ConcurrentMap<String, Object> delta) {
-    this.state = Objects.requireNonNull(state);
-    this.delta = delta;
+  public State(Map<String, Object> state) {
+    this(state, null);
+  }
+
+  public State(Map<String, Object> state, @Nullable Map<String, Object> delta) {
+    Objects.requireNonNull(state, "state is null");
+    this.state =
+        state instanceof ConcurrentMap
+            ? (ConcurrentMap<String, Object>) state
+            : new ConcurrentHashMap<>(state);
+    this.delta =
+        delta == null
+            ? new ConcurrentHashMap<>()
+            : delta instanceof ConcurrentMap
+                ? (ConcurrentMap<String, Object>) delta
+                : new ConcurrentHashMap<>(delta);
   }
 
   @Override
diff --git a/core/src/test/java/com/google/adk/sessions/StateTest.java b/core/src/test/java/com/google/adk/sessions/StateTest.java
@@ -0,0 +1,50 @@
+package com.google.adk.sessions;
+
+import static com.google.common.truth.Truth.assertThat;
+
+import java.util.HashMap;
+import java.util.Map;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ConcurrentMap;
+import org.junit.Assert;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+
+@RunWith(JUnit4.class)
+public class StateTest {
+  @Test
+  public void constructor_nullDelta_createsEmptyConcurrentHashMap() {
+    ConcurrentMap<String, Object> stateMap = new ConcurrentHashMap<>();
+    State state = new State(stateMap, null);
+    assertThat(state.hasDelta()).isFalse();
+    state.put("key", "value");
+    assertThat(state.hasDelta()).isTrue();
+  }
+
+  @Test
+  public void constructor_nullState_throwsException() {
+    Assert.assertThrows(NullPointerException.class, () -> new State(null, new HashMap<>()));
+  }
+
+  @Test
+  public void constructor_regularMapState() {
+    Map<String, Object> stateMap = new HashMap<>();
+    stateMap.put("initial", "val");
+    State state = new State(stateMap, null);
+    // It should have copied the contents
+    assertThat(state).containsEntry("initial", "val");
+    state.put("key", "value");
+    // The original map should NOT be updated because a copy was created
+    assertThat(stateMap).doesNotContainKey("key");
+  }
+
+  @Test
+  public void constructor_singleArgument() {
+    ConcurrentMap<String, Object> stateMap = new ConcurrentHashMap<>();
+    State state = new State(stateMap);
+    assertThat(state.hasDelta()).isFalse();
+    state.put("key", "value");
+    assertThat(state.hasDelta()).isTrue();
+  }
+}
