fix: Handle JSON deserialized State.REMOVED sentinel

The State.REMOVED sentinel object is serialized to a specific string in JSON. When deserialized, this becomes a String object, not the original sentinel instance. This change updates the logic to correctly identify removals when the value is either the State.REMOVED object or its string representation.

Additionally, when applying state deltas from app/user state, removals are now handled using a new removeWithoutDelta method on the State object to avoid re-recording these removals as new deltas.

PiperOrigin-RevId: 881521527

Author: google-genai-bot

core/src/main/java/com/google/adk/sessions/BaseSessionService.java

@@ -236,7 +236,7 @@ default Single<Event> appendEvent(Session session, Event event) {
         stateDelta.forEach(
             (key, value) -> {
               if (!key.startsWith(State.TEMP_PREFIX)) {
-                if (value == State.REMOVED) {
+                if (State.isRemoved(value)) {
                   sessionState.remove(key);
                 } else {
                   sessionState.put(key, value);

core/src/main/java/com/google/adk/sessions/InMemorySessionService.java

@@ -244,7 +244,7 @@ public Single<Event> appendEvent(Session session, Event event) {
             (key, value) -> {
               if (key.startsWith(State.APP_PREFIX)) {
                 String appStateKey = key.substring(State.APP_PREFIX.length());
-                if (value == State.REMOVED) {
+                if (State.isRemoved(value)) {
                   appState
                       .computeIfAbsent(appName, unused -> new ConcurrentHashMap<>())
                       .remove(appStateKey);
@@ -255,7 +255,7 @@ public Single<Event> appendEvent(Session session, Event event) {
                 }
               } else if (key.startsWith(State.USER_PREFIX)) {
                 String userStateKey = key.substring(State.USER_PREFIX.length());
-                if (value == State.REMOVED) {
+                if (State.isRemoved(value)) {
                   userState
                       .computeIfAbsent(appName, unused -> new ConcurrentHashMap<>())
                       .computeIfAbsent(userId, unused -> new ConcurrentHashMap<>())
@@ -267,8 +267,13 @@ public Single<Event> appendEvent(Session session, Event event) {
                       .put(userStateKey, value);
                 }
               } else {
-                if (value == State.REMOVED) {
-                  session.state().remove(key);
+                if (State.isRemoved(value)) {
+                  Map<String, Object> s = session.state();
+                  if (s instanceof State state) {
+                    state.removeWithoutDelta(key);
+                  } else {
+                    s.remove(key);
+                  }
                 } else {
                   session.state().put(key, value);
                 }
@@ -333,12 +338,34 @@ private Session mergeWithGlobalState(String appName, String userId, Session sess
     // Merge App State directly into the session's state map
     appState
         .computeIfAbsent(appName, unused -> new ConcurrentHashMap<>())
-        .forEach((key, value) -> sessionState.put(State.APP_PREFIX + key, value));
+        .forEach(
+            (key, value) -> {
+              if (State.isRemoved(value)) {
+                if (sessionState instanceof State state) {
+                  state.removeWithoutDelta(State.APP_PREFIX + key);
+                } else {
+                  sessionState.remove(State.APP_PREFIX + key);
+                }
+              } else {
+                sessionState.put(State.APP_PREFIX + key, value);
+              }
+            });
 
     userState
         .computeIfAbsent(appName, unused -> new ConcurrentHashMap<>())
         .computeIfAbsent(userId, unused -> new ConcurrentHashMap<>())
-        .forEach((key, value) -> sessionState.put(State.USER_PREFIX + key, value));
+        .forEach(
+            (key, value) -> {
+              if (State.isRemoved(value)) {
+                if (sessionState instanceof State state) {
+                  state.removeWithoutDelta(State.USER_PREFIX + key);
+                } else {
+                  sessionState.remove(State.USER_PREFIX + key);
+                }
+              } else {
+                sessionState.put(State.USER_PREFIX + key, value);
+              }
+            });
 
     return session;
   }

core/src/main/java/com/google/adk/sessions/State.java

@@ -17,6 +17,7 @@
 package com.google.adk.sessions;
 
 import com.fasterxml.jackson.annotation.JsonValue;
+import com.google.errorprone.annotations.CanIgnoreReturnValue;
 import java.util.Collection;
 import java.util.Map;
 import java.util.Map.Entry;
@@ -33,6 +34,8 @@ public final class State implements ConcurrentMap<String, Object> {
   public static final String USER_PREFIX = "user:";
   public static final String TEMP_PREFIX = "temp:";
 
+  public static final String REMOVED_SENTINEL_STRING = "__ADK_SENTINEL_REMOVED__";
+
   /** Sentinel object to mark removed entries in the delta map. */
   public static final Object REMOVED = RemovedSentinel.INSTANCE;
 
@@ -129,6 +132,19 @@ public Object remove(Object key) {
     return state.remove(key);
   }
 
+  /**
+   * Removes a key from the state map without recording the removal in the delta map. This is
+   * intended for internal use when rebuilding state from an event stream where the removal is
+   * already known and doesn't need to be represented as a new change.
+   *
+   * @param key The key to remove.
+   * @return The previous value associated with key, or null if there was no mapping for key.
+   */
+  @CanIgnoreReturnValue
+  public Object removeWithoutDelta(Object key) {
+    return state.remove(key);
+  }
+
   @Override
   public boolean remove(Object key, Object value) {
     boolean removed = state.remove(key, value);
@@ -170,6 +186,16 @@ public boolean hasDelta() {
     return !delta.isEmpty();
   }
 
+  /**
+   * Checks if a value represents a removed state entry, accounting for deserialization from JSON.
+   *
+   * @param value The value to check.
+   * @return True if the value indicates removal, false otherwise.
+   */
+  public static boolean isRemoved(Object value) {
+    return value == REMOVED || Objects.equals(value, REMOVED_SENTINEL_STRING);
+  }
+
   private static final class RemovedSentinel {
     public static final RemovedSentinel INSTANCE = new RemovedSentinel();
 
@@ -179,7 +205,7 @@ private RemovedSentinel() {
 
     @JsonValue
     public String toJson() {
-      return "__ADK_SENTINEL_REMOVED__";
+      return REMOVED_SENTINEL_STRING;
     }
   }
 }

core/src/test/java/com/google/adk/sessions/InMemorySessionServiceTest.java

@@ -214,6 +214,55 @@ public void appendEvent_removesState() {
     assertThat(retrievedSessionRemove.state()).doesNotContainKey("temp:tempKey");
   }
 
+  @Test
+  public void appendEvent_removesStateFromJsonDeserializedSentinel() {
+    InMemorySessionService sessionService = new InMemorySessionService();
+    Session session =
+        sessionService.createSession("app", "user", new HashMap<>(), "session1").blockingGet();
+
+    ConcurrentMap<String, Object> stateDeltaAdd = new ConcurrentHashMap<>();
+    stateDeltaAdd.put("sessionKey", "sessionValue");
+    stateDeltaAdd.put("_app_appKey", "appValue");
+    stateDeltaAdd.put("_user_userKey", "userValue");
+
+    Event eventAdd =
+        Event.builder().actions(EventActions.builder().stateDelta(stateDeltaAdd).build()).build();
+
+    var unused = sessionService.appendEvent(session, eventAdd).blockingGet();
+
+    // Verify state is added
+    Session retrievedSessionAdd =
+        sessionService
+            .getSession(session.appName(), session.userId(), session.id(), Optional.empty())
+            .blockingGet();
+    assertThat(retrievedSessionAdd.state()).containsEntry("sessionKey", "sessionValue");
+    assertThat(retrievedSessionAdd.state()).containsEntry("_app_appKey", "appValue");
+    assertThat(retrievedSessionAdd.state()).containsEntry("_user_userKey", "userValue");
+
+    // Prepare and append event to remove state using the String representation of the sentinel
+    // to simulate Jackson JSON deserialization.
+    ConcurrentMap<String, Object> stateDeltaRemove = new ConcurrentHashMap<>();
+    stateDeltaRemove.put("sessionKey", State.REMOVED_SENTINEL_STRING);
+    stateDeltaRemove.put("_app_appKey", State.REMOVED_SENTINEL_STRING);
+    stateDeltaRemove.put("_user_userKey", State.REMOVED_SENTINEL_STRING);
+
+    Event eventRemove =
+        Event.builder()
+            .actions(EventActions.builder().stateDelta(stateDeltaRemove).build())
+            .build();
+
+    unused = sessionService.appendEvent(session, eventRemove).blockingGet();
+
+    // Verify state is removed despite being a String instead of the State.REMOVED object
+    Session retrievedSessionRemove =
+        sessionService
+            .getSession(session.appName(), session.userId(), session.id(), Optional.empty())
+            .blockingGet();
+    assertThat(retrievedSessionRemove.state()).doesNotContainKey("sessionKey");
+    assertThat(retrievedSessionRemove.state()).doesNotContainKey("_app_appKey");
+    assertThat(retrievedSessionRemove.state()).doesNotContainKey("_user_userKey");
+  }
+
   @Test
   public void sequentialAgents_shareTempState() {
     InMemorySessionService sessionService = new InMemorySessionService();