From e534f12bd5c7cadb8a6100b00ac2ae771a868ab0 Mon Sep 17 00:00:00 2001 From: Google Team Member Date: Thu, 19 Mar 2026 05:59:19 -0700 Subject: [PATCH] refactor: Update map handling in EventActions to always use defensive copy and add null handling for `artifactDelta` in the Builder PiperOrigin-RevId: 886130618 --- .../sessions/FirestoreSessionServiceTest.java | 99 ------------------- .../com/google/adk/events/EventActions.java | 17 ++-- .../google/adk/events/EventActionsTest.java | 11 --- 3 files changed, 6 insertions(+), 121 deletions(-) diff --git a/contrib/firestore-session-service/src/test/java/com/google/adk/sessions/FirestoreSessionServiceTest.java b/contrib/firestore-session-service/src/test/java/com/google/adk/sessions/FirestoreSessionServiceTest.java index ffcbcf8f5..43ca6889f 100644 --- a/contrib/firestore-session-service/src/test/java/com/google/adk/sessions/FirestoreSessionServiceTest.java +++ b/contrib/firestore-session-service/src/test/java/com/google/adk/sessions/FirestoreSessionServiceTest.java @@ -47,15 +47,11 @@ import com.google.genai.types.Part; import io.reactivex.rxjava3.observers.TestObserver; import java.time.Instant; -import java.util.AbstractMap; import java.util.Collections; -import java.util.HashMap; import java.util.List; import java.util.Map; import java.util.Optional; -import java.util.Set; import java.util.concurrent.ConcurrentHashMap; -import java.util.concurrent.ConcurrentMap; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; import org.junit.jupiter.api.extension.ExtendWith; @@ -530,44 +526,6 @@ void appendAndGet_withAllPartTypes_serializesAndDeserializesCorrectly() { }); } - /** - * A wrapper class that implements ConcurrentMap but delegates to a HashMap. This is a workaround - * to allow putting null values, which ConcurrentHashMap forbids, for testing state removal logic. - */ - private static class HashMapAsConcurrentMap extends AbstractMap - implements ConcurrentMap { - private final HashMap map; - - public HashMapAsConcurrentMap(Map map) { - this.map = new HashMap<>(map); - } - - @Override - public Set> entrySet() { - return map.entrySet(); - } - - @Override - public V putIfAbsent(K key, V value) { - return map.putIfAbsent(key, value); - } - - @Override - public boolean remove(Object key, Object value) { - return map.remove(key, value); - } - - @Override - public boolean replace(K key, V oldValue, V newValue) { - return map.replace(key, oldValue, newValue); - } - - @Override - public V replace(K key, V value) { - return map.replace(key, value); - } - } - /** Tests that appendEvent with only app state deltas updates the correct stores. */ @Test void appendEvent_withAppOnlyStateDeltas_updatesCorrectStores() { @@ -662,63 +620,6 @@ void appendEvent_withUserOnlyStateDeltas_updatesCorrectStores() { verify(mockSessionDocRef, never()).update(eq(Constants.KEY_STATE), any()); } - /** - * Tests that appendEvent with all types of state deltas updates the correct stores and session - * state. - */ - @Test - void appendEvent_withAllStateDeltas_updatesCorrectStores() { - // Arrange - Session session = - Session.builder(SESSION_ID) - .appName(APP_NAME) - .userId(USER_ID) - .state(new ConcurrentHashMap<>()) // The session state itself must be concurrent - .build(); - session.state().put("keyToRemove", "someValue"); - - Map stateDeltaMap = new HashMap<>(); - stateDeltaMap.put("sessionKey", "sessionValue"); - stateDeltaMap.put("_app_appKey", "appValue"); - stateDeltaMap.put("_user_userKey", "userValue"); - stateDeltaMap.put("keyToRemove", null); - - // Use the wrapper to satisfy the ConcurrentMap interface for the builder - EventActions actions = - EventActions.builder().stateDelta(new HashMapAsConcurrentMap<>(stateDeltaMap)).build(); - - Event event = - Event.builder() - .author("model") - .content(Content.builder().parts(List.of(Part.fromText("..."))).build()) - .actions(actions) - .build(); - - when(mockSessionsCollection.document(SESSION_ID)).thenReturn(mockSessionDocRef); - when(mockEventsCollection.document()).thenReturn(mockEventDocRef); - when(mockEventDocRef.getId()).thenReturn(EVENT_ID); - // THIS IS THE MISSING MOCK: Stub the call to get the document by its specific ID. - when(mockEventsCollection.document(EVENT_ID)).thenReturn(mockEventDocRef); - // Add the missing mock for the final session update call - when(mockSessionDocRef.update(anyMap())) - .thenReturn(ApiFutures.immediateFuture(mockWriteResult)); - - // Act - sessionService.appendEvent(session, event).test().assertComplete(); - - // Assert - assertThat(session.state()).containsEntry("sessionKey", "sessionValue"); - assertThat(session.state()).doesNotContainKey("keyToRemove"); - - ArgumentCaptor> appStateCaptor = ArgumentCaptor.forClass(Map.class); - verify(mockAppStateDocRef).set(appStateCaptor.capture(), any(SetOptions.class)); - assertThat(appStateCaptor.getValue()).containsEntry("appKey", "appValue"); - - ArgumentCaptor> userStateCaptor = ArgumentCaptor.forClass(Map.class); - verify(mockUserStateUserDocRef).set(userStateCaptor.capture(), any(SetOptions.class)); - assertThat(userStateCaptor.getValue()).containsEntry("userKey", "userValue"); - } - /** Tests that getSession skips malformed events and returns only the well-formed ones. */ @Test @SuppressWarnings("unchecked") 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 3565c3e99..83fd60e54 100644 --- a/core/src/main/java/com/google/adk/events/EventActions.java +++ b/core/src/main/java/com/google/adk/events/EventActions.java @@ -157,9 +157,6 @@ public void setRequestedToolConfirmations( Map requestedToolConfirmations) { if (requestedToolConfirmations == null) { this.requestedToolConfirmations = new ConcurrentHashMap<>(); - } else if (requestedToolConfirmations instanceof ConcurrentMap) { - this.requestedToolConfirmations = - (ConcurrentMap) requestedToolConfirmations; } else { this.requestedToolConfirmations = new ConcurrentHashMap<>(requestedToolConfirmations); } @@ -290,8 +287,6 @@ public Builder skipSummarization(boolean skipSummarization) { public Builder stateDelta(@Nullable Map value) { if (value == null) { this.stateDelta = new ConcurrentHashMap<>(); - } else if (value instanceof ConcurrentMap) { - this.stateDelta = (ConcurrentMap) value; } else { this.stateDelta = new ConcurrentHashMap<>(value); } @@ -300,8 +295,12 @@ public Builder stateDelta(@Nullable Map value) { @CanIgnoreReturnValue @JsonProperty("artifactDelta") - public Builder artifactDelta(Map value) { - this.artifactDelta = new ConcurrentHashMap<>(value); + public Builder artifactDelta(@Nullable Map value) { + if (value == null) { + this.artifactDelta = new ConcurrentHashMap<>(); + } else { + this.artifactDelta = new ConcurrentHashMap<>(value); + } return this; } @@ -339,10 +338,6 @@ public Builder requestedAuthConfigs( public Builder requestedToolConfirmations(@Nullable Map value) { if (value == null) { this.requestedToolConfirmations = new ConcurrentHashMap<>(); - return this; - } - if (value instanceof ConcurrentMap) { - this.requestedToolConfirmations = (ConcurrentMap) value; } else { this.requestedToolConfirmations = new ConcurrentHashMap<>(value); } 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 22bb94e64..b1e645e1a 100644 --- a/core/src/test/java/com/google/adk/events/EventActionsTest.java +++ b/core/src/test/java/com/google/adk/events/EventActionsTest.java @@ -177,17 +177,6 @@ public void merge_failsOnMismatchedKeyTypesNestedInStateDelta() { IllegalArgumentException.class, () -> eventActions1.toBuilder().merge(eventActions2)); } - @Test - public void setRequestedToolConfirmations_withConcurrentMap_usesSameInstance() { - ConcurrentHashMap map = new ConcurrentHashMap<>(); - map.put("tool", TOOL_CONFIRMATION); - - EventActions actions = new EventActions(); - actions.setRequestedToolConfirmations(map); - - assertThat(actions.requestedToolConfirmations()).isSameInstanceAs(map); - } - @Test public void setRequestedToolConfirmations_withRegularMap_createsConcurrentMap() { ImmutableMap map = ImmutableMap.of("tool", TOOL_CONFIRMATION);