From cd94ca67e7b3e0a21126e22733643796863516b8 Mon Sep 17 00:00:00 2001
From: Google Team Member <genai-sdk-bot@google.com>
Date: Wed, 6 May 2026 10:53:26 -0700
Subject: [PATCH] fix: Account for null values in State constructor parameters

PiperOrigin-RevId: 911429082
---
 .../java/com/google/adk/sessions/State.java   | 26 ++++++++++++-------
 .../com/google/adk/sessions/StateTest.java    | 16 ++++++++++--
 2 files changed, 30 insertions(+), 12 deletions(-)

diff --git a/core/src/main/java/com/google/adk/sessions/State.java b/core/src/main/java/com/google/adk/sessions/State.java
index 577559f85..d86aae9bf 100644
--- a/core/src/main/java/com/google/adk/sessions/State.java
+++ b/core/src/main/java/com/google/adk/sessions/State.java
@@ -46,16 +46,22 @@ public State(Map<String, Object> state) {
 
   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);
+    this.state = toConcurrentMap(state);
+    this.delta = toConcurrentMap(delta);
+  }
+
+  private static ConcurrentMap<String, Object> toConcurrentMap(Map<String, Object> map) {
+    if (map == null) {
+      return new ConcurrentHashMap<>();
+    }
+    if (map instanceof ConcurrentMap) {
+      return (ConcurrentMap<String, Object>) map;
+    }
+    ConcurrentMap<String, Object> concurrentMap = new ConcurrentHashMap<>();
+    for (Map.Entry<String, Object> entry : map.entrySet()) {
+      concurrentMap.put(entry.getKey(), entry.getValue() == null ? REMOVED : entry.getValue());
+    }
+    return concurrentMap;
   }
 
   @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
index e1fcaeadc..305812335 100644
--- a/core/src/test/java/com/google/adk/sessions/StateTest.java
+++ b/core/src/test/java/com/google/adk/sessions/StateTest.java
@@ -23,8 +23,10 @@ public void constructor_nullDelta_createsEmptyConcurrentHashMap() {
   }
 
   @Test
-  public void constructor_nullState_throwsException() {
-    Assert.assertThrows(NullPointerException.class, () -> new State(null, new HashMap<>()));
+  public void constructor_nullState_createsEmptyConcurrentHashMap() {
+    State state = new State(null, new HashMap<>());
+    assertThat(state.isEmpty()).isTrue();
+    assertThat(state.hasDelta()).isFalse();
   }
 
   @Test
@@ -47,4 +49,14 @@ public void constructor_singleArgument() {
     state.put("key", "value");
     assertThat(state.hasDelta()).isTrue();
   }
+
+  @Test
+  public void constructor_stateMapWithNullValues_replacesWithRemoved() {
+    Map<String, Object> stateMap = new HashMap<>();
+    stateMap.put("key1", "value1");
+    stateMap.put("key2", null);
+    State state = new State(stateMap);
+    assertThat(state).containsEntry("key1", "value1");
+    assertThat(state).containsEntry("key2", State.REMOVED);
+  }
 }