fix:Refactor State.java to correctly merge state and delta for map interface methods

PiperOrigin-RevId: 871435115

Author: google-genai-bot

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

@@ -51,68 +51,135 @@ public State(ConcurrentMap<String, Object> state, ConcurrentMap<String, Object>
   @Override
   public void clear() {
     state.clear();
+    // Delta should likely be cleared too if we are clearing the state,
+    // or we might want to mark everything as removed in delta.
+    // Given the Python implementation doesn't have clear, and this is a local view,
+    // clearing both seems appropriate to reset the object.
+    delta.clear();
   }
 
   @Override
   public boolean containsKey(Object key) {
+    if (delta.containsKey(key)) {
+      return delta.get(key) != REMOVED;
+    }
     return state.containsKey(key);
   }
 
   @Override
   public boolean containsValue(Object value) {
-    return state.containsValue(value);
+    // This is expensive but necessary for correctness with the merged view.
+    return values().contains(value);
   }
 
   @Override
   public Set<Entry<String, Object>> entrySet() {
-    return state.entrySet();
+    // This provides a snapshot, not a live view backed by the map, which differs from standard Map
+    // contract.
+    // However, given the complexity of merging two concurrent maps, this is a reasonable compromise
+    // for this specific implementation.
+    // TODO: Consider implementing a live view if needed.
+    Map<String, Object> merged = new ConcurrentHashMap<>(state);
+    for (Entry<String, Object> entry : delta.entrySet()) {
+      if (entry.getValue() == REMOVED) {
+        merged.remove(entry.getKey());
+      } else {
+        merged.put(entry.getKey(), entry.getValue());
+      }
+    }
+    return merged.entrySet();
   }
 
   @Override
   public boolean equals(Object o) {
     if (o == this) {
       return true;
     }
-    if (!(o instanceof State other)) {
+    if (!(o instanceof Map)) {
       return false;
     }
-    return state.equals(other.state);
+    Map<?, ?> other = (Map<?, ?>) o;
+    // We can't easily rely on state.equals() because our "content" is merged.
+    // Validating equality against another Map requires checking the merged view.
+    if (size() != other.size()) {
+      return false;
+    }
+    try {
+      for (Entry<String, Object> e : entrySet()) {
+        String key = e.getKey();
+        Object value = e.getValue();
+        if (value == null) {
+          if (!(other.get(key) == null && other.containsKey(key))) return false;
+        } else {
+          if (!value.equals(other.get(key))) return false;
+        }
+      }
+    } catch (ClassCastException | NullPointerException unused) {
+      return false;
+    }
+    return true;
   }
 
   @Override
   public Object get(Object key) {
+    if (delta.containsKey(key)) {
+      Object value = delta.get(key);
+      return value == REMOVED ? null : value;
+    }
     return state.get(key);
   }
 
   @Override
   public int hashCode() {
-    return state.hashCode();
+    // Similar to equals, we need to calculate hash code based on the merged entry set.
+    int h = 0;
+    for (Entry<String, Object> entry : entrySet()) {
+      h += entry.hashCode();
+    }
+    return h;
   }
 
   @Override
   public boolean isEmpty() {
-    return state.isEmpty();
+    if (delta.isEmpty()) {
+      return state.isEmpty();
+    }
+    // If delta is not empty, we need to check if it effectively removes everything from state
+    // or adds something.
+    return size() == 0;
   }
 
   @Override
   public Set<String> keySet() {
-    return state.keySet();
+    // Snapshot view
+    Map<String, Object> merged = new ConcurrentHashMap<>(state);
+    for (Entry<String, Object> entry : delta.entrySet()) {
+      if (entry.getValue() == REMOVED) {
+        merged.remove(entry.getKey());
+      } else {
+        merged.put(entry.getKey(), entry.getValue());
+      }
+    }
+    return merged.keySet();
   }
 
   @Override
   public Object put(String key, Object value) {
-    Object oldValue = state.put(key, value);
+    // Current value logic needs to check delta first to return correct "oldValue"
+    Object oldValue = get(key);
+    state.put(key, value);
     delta.put(key, value);
     return oldValue;
   }
 
   @Override
   public Object putIfAbsent(String key, Object value) {
-    Object existingValue = state.putIfAbsent(key, value);
-    if (existingValue == null) {
-      delta.put(key, value);
+    Object currentValue = get(key);
+    if (currentValue == null) {
+      put(key, value);
+      return null;
     }
-    return existingValue;
+    return currentValue;
   }
 
   @Override
@@ -123,47 +190,130 @@ public void putAll(Map<? extends String, ? extends Object> m) {
 
   @Override
   public Object remove(Object key) {
-    if (state.containsKey(key)) {
+    Object oldValue = get(key);
+    // Explicitly check for containment in the *merged* view (via get != null or containsKey)
+    // before marking as removed, though strictly speaking marking as removed is safe even if not
+    // present.
+    // But we need to return the correct oldValue.
+
+    if (state.containsKey(key) || (delta.containsKey(key) && delta.get(key) != REMOVED)) {
       delta.put((String) key, REMOVED);
     }
-    return state.remove(key);
+    // We should probably NOT remove from state to keep "state" as the committed version?
+    // The original code did:
+    // delta.put((String) key, REMOVED);
+    // return state.remove(key);
+    //
+    // If we want state to represent the "base" and delta the "changes", we should probably
+    // ONLY update delta. However, the original code updated BOTH.
+    // "A State object that also keeps track of the changes to the state."
+    // If it updates both, then "state" is presumably the "current live state" AND "delta" is the
+    // log.
+    // PROPOSAL: Keep updating both to minimize regression risk on the write path,
+    // but ensure read path prioritizes delta (which might be redundant if state is always updated,
+    // UNLESS delta has entries that state doesn't? Or if we want to support a mode where only delta
+    // is updated?)
+    //
+    // Wait, the original code:
+    // Object oldValue = state.put(key, value);
+    // delta.put(key, value);
+    //
+    // This implies `state` IS updated.
+    //
+    // If `state` is updated, why did we need to check `delta` for reads?
+    // "State.java seems to ignore the delta map for the most part. This looks wrong."
+    //
+    // SCENARIO 1: `state` is the "committed" state from DB. `delta` accumulates changes.
+    // If we update `state` in place, then `get(key)` on `state` returns the new value.
+    // So why claim it ignores delta?
+    //
+    // Maybe `State` is initialized with a `state` map that is SHARED or comes from a source
+    // that shouldn't be mutated?
+    //
+    // OR, maybe the `state` map passed in constructor is NOT updated by some other code path?
+    //
+    // Let's re-read the Python code.
+    // Python:
+    // self._value[key] = value
+    // self._delta[key] = value
+    //
+    // It updates BOTH.
+    //
+    // Python `__getitem__`:
+    // if key in self._delta: return self._delta[key]
+    // return self._value[key]
+    //
+    // If both are updated, `self._value[key]` should correspond to `self._delta[key]`.
+    //
+    // UNLESS `delta` is populated with something that IS NOT in `value`.
+    // Example: We load a session. State is {a:1}. We create a `State` object.
+    // We put b:2. state={a:1, b:2}, delta={b:2}.
+    // get(b) -> state.get(b) -> 2. Correct.
+    //
+    // SCENARIO 2: Replay / Deferred execution.
+    // Maybe someone populates `delta` directly?
+    // Constructor: `public State(ConcurrentMap<String, Object> state, ConcurrentMap<String, Object>
+    // delta)`
+    //
+    // If I pass a `delta` map that already has stuff, but `state` doesn't.
+    // State({a:1}, {a:2})
+    // Java get(a) -> state.get(a) -> 1. WRONG. Should be 2.
+    //
+    // Correct. This is the scenario. The `delta` might contain provisional changes not yet applied
+    // to `state`.
+
+    state.remove(key);
+    return oldValue;
   }
 
   @Override
   public boolean remove(Object key, Object value) {
-    boolean removed = state.remove(key, value);
-    if (removed) {
-      delta.put((String) key, REMOVED);
+    Object currentValue = get(key);
+    if (Objects.equals(currentValue, value) && (currentValue != null || containsKey(key))) {
+      remove(key);
+      return true;
     }
-    return removed;
+    return false;
   }
 
   @Override
   public boolean replace(String key, Object oldValue, Object newValue) {
-    boolean replaced = state.replace(key, oldValue, newValue);
-    if (replaced) {
-      delta.put(key, newValue);
+    Object currentValue = get(key);
+    if (Objects.equals(currentValue, oldValue) && (currentValue != null || containsKey(key))) {
+      put(key, newValue);
+      return true;
     }
-    return replaced;
+    return false;
   }
 
   @Override
   public Object replace(String key, Object value) {
-    Object oldValue = state.replace(key, value);
-    if (oldValue != null) {
-      delta.put(key, value);
+    Object currentValue = get(key);
+    if (currentValue != null || containsKey(key)) {
+      put(key, value);
+      return currentValue;
     }
-    return oldValue;
+    return null;
   }
 
   @Override
   public int size() {
-    return state.size();
+    // Expensive, but accurate merged size.
+    return entrySet().size();
   }
 
   @Override
   public Collection<Object> values() {
-    return state.values();
+    // Snapshot view
+    Map<String, Object> merged = new ConcurrentHashMap<>(state);
+    for (Entry<String, Object> entry : delta.entrySet()) {
+      if (entry.getValue() == REMOVED) {
+        merged.remove(entry.getKey());
+      } else {
+        merged.put(entry.getKey(), entry.getValue());
+      }
+    }
+    return merged.values();
   }
 
   public boolean hasDelta() {