From a1610094728b67abac7ed77833d086bb5cad45ea Mon Sep 17 00:00:00 2001 From: Jaehee Shin Date: Sat, 10 Oct 2026 07:14:35 +0900 Subject: [PATCH] fix: stop losing event fields on Firestore session reloads FirestoreSessionService stored only an event's author, timestamp, and text, function call, function response and file data parts, and rebuilt events from those fields alone. Reloaded events lost their ID, invocation ID, branch, actions, thought flags and signatures, and function call and response IDs, and user events came back with the user ID as their author. Function call args were stored as they were, so a tool confirmation request, whose args hold a FunctionCall and a ToolConfirmation, could not be saved. - Events are also stored as JSON in `rawEvent` and read back from it. Events stored before this change still load from the other fields, which FirestoreMemoryService keeps reading. - Tool call args and results are stored as plain values. --- .../adk/sessions/FirestoreSessionService.java | 31 +- .../sessions/FirestoreSessionServiceTest.java | 271 +++++++++++++++++- 2 files changed, 299 insertions(+), 3 deletions(-) diff --git a/contrib/firestore-session-service/src/main/java/com/google/adk/sessions/FirestoreSessionService.java b/contrib/firestore-session-service/src/main/java/com/google/adk/sessions/FirestoreSessionService.java index 54f91a858..5c8dbb8c5 100644 --- a/contrib/firestore-session-service/src/main/java/com/google/adk/sessions/FirestoreSessionService.java +++ b/contrib/firestore-session-service/src/main/java/com/google/adk/sessions/FirestoreSessionService.java @@ -15,6 +15,8 @@ */ package com.google.adk.sessions; +import com.fasterxml.jackson.core.type.TypeReference; +import com.google.adk.JsonBaseModel; import com.google.adk.events.Event; import com.google.adk.events.EventActions; import com.google.adk.platform.UuidProvider; @@ -31,6 +33,7 @@ import com.google.cloud.firestore.QueryDocumentSnapshot; import com.google.cloud.firestore.WriteBatch; import com.google.cloud.firestore.WriteResult; +import com.google.common.base.Throwables; import com.google.common.util.concurrent.MoreExecutors; import com.google.errorprone.annotations.CanIgnoreReturnValue; import com.google.genai.types.Content; @@ -83,6 +86,9 @@ public class FirestoreSessionService implements BaseSessionService { /** Random token each create stores, so a retried create can detect its own write. */ static final String CREATE_TOKEN_KEY = "createToken"; + /** The event as a JSON string, not a map as in ADK Python, to avoid Firestore's map limits. */ + static final String RAW_EVENT_KEY = "rawEvent"; + /** Constructor for FirestoreSessionService. */ public FirestoreSessionService(Firestore firestore) { this(firestore, InstantSource.system(), UuidProvider.SYSTEM); @@ -280,6 +286,9 @@ public Maybe getSession( /** * Reconstructs an Event object from a Map retrieved from Firestore. * + *

A readable {@code rawEvent} takes precedence over the other fields, which remain the + * fallback when it is missing or unreadable. + * * @param data The map representation of the event. * @return An Event object, or null if the data is malformed. */ @@ -288,6 +297,15 @@ private Event eventFromMap(Map data, String sessionUserId) { return null; } try { + if (data.get(RAW_EVENT_KEY) instanceof String rawEvent) { + try { + return Event.fromJson(rawEvent); + } catch (IllegalStateException e) { + logger.warn( + "Ignoring unreadable rawEvent ({}), using the stored event fields", + Throwables.getRootCause(e).getClass().getSimpleName()); + } + } String author = safeCast(data.get("author"), String.class, "author"); String timestampStr = safeCast(data.get(TIMESTAMP_KEY), String.class, "timestamp"); Map contentMap = safeCast(data.get("content"), Map.class, "content"); @@ -443,7 +461,7 @@ private Map eventToMap(Session session, Event event) { fc -> { Map fcMap = new HashMap<>(); fc.name().ifPresent(name -> fcMap.put("name", name)); - fc.args().ifPresent(args -> fcMap.put("args", args)); + fc.args().ifPresent(args -> fcMap.put("args", toPlainValues(args))); if (!fcMap.isEmpty()) { partData.put("functionCall", fcMap); } @@ -453,7 +471,9 @@ private Map eventToMap(Session session, Event event) { fr -> { Map frMap = new HashMap<>(); fr.name().ifPresent(name -> frMap.put("name", name)); - fr.response().ifPresent(response -> frMap.put("response", response)); + fr.response() + .ifPresent( + response -> frMap.put("response", toPlainValues(response))); if (!frMap.isEmpty()) { partData.put("functionResponse", frMap); } @@ -483,10 +503,17 @@ private Map eventToMap(Session session, Event event) { if (!keywords.isEmpty()) { data.put("keywords", new ArrayList<>(keywords)); // Firestore works well with Lists } + data.put(RAW_EVENT_KEY, event.toJson()); return data; } + /** Converts tool args or results to plain values, since Firestore cannot encode some objects. */ + private static Map toPlainValues(Map value) { + return JsonBaseModel.getMapper() + .convertValue(value, new TypeReference>() {}); + } + /** Lists all sessions for a given appName and userId. */ @Override public Single listSessions(String appName, String userId) { 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 ba95dc84c..9d08b0b50 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 @@ -29,6 +29,7 @@ import com.google.adk.events.Event; import com.google.adk.events.EventActions; +import com.google.adk.events.ToolConfirmation; import com.google.adk.utils.Constants; import com.google.api.core.ApiFutures; import com.google.api.gax.grpc.GrpcStatusCode; @@ -48,11 +49,15 @@ import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableMap; import com.google.genai.types.Content; +import com.google.genai.types.FunctionCall; +import com.google.genai.types.FunctionResponse; import com.google.genai.types.Part; import io.grpc.Status; import io.reactivex.rxjava3.observers.TestObserver; import java.time.Instant; +import java.util.ArrayList; import java.util.Collections; +import java.util.HashMap; import java.util.List; import java.util.Map; import java.util.Optional; @@ -622,7 +627,9 @@ void appendAndGet_withAllPartTypes_serializesAndDeserializesCorrectly() { // ARRANGE (Deserialization part) ArgumentCaptor> eventDataCaptor = ArgumentCaptor.forClass(Map.class); verify(mockEventDocRef).set(eventDataCaptor.capture()); - Map savedEventData = eventDataCaptor.getValue(); + // Read back as an event stored before rawEvent, from the other fields. + Map savedEventData = new HashMap<>(eventDataCaptor.getValue()); + savedEventData.remove(FirestoreSessionService.RAW_EVENT_KEY); QueryDocumentSnapshot mockComplexEventSnapshot = mock(QueryDocumentSnapshot.class); when(mockComplexEventSnapshot.getData()).thenReturn(savedEventData); @@ -670,6 +677,222 @@ void appendAndGet_withAllPartTypes_serializesAndDeserializesCorrectly() { }); } + /** Tests that getSession returns appended events with all their fields. */ + @Test + @SuppressWarnings("unchecked") + void appendAndGet_returnsEventsAsAppended() { + Session session = + Session.builder(SESSION_ID) + .appName(APP_NAME) + .userId(USER_ID) + .state(new ConcurrentHashMap<>()) + .build(); + Event question = + Event.builder() + .id("event-1") + .invocationId("invocation-1") + .author(Constants.KEY_USER) + .timestamp(1_000L) + .content(Content.fromParts(Part.fromText("What's the weather in Seoul?"))) + .build(); + Event call = + Event.builder() + .id("event-2") + .invocationId("invocation-1") + .author("weather_agent") + .branch("weather_agent") + .timestamp(2_000L) + .content( + Content.builder() + .role("model") + .parts( + Part.builder() + .text("The user wants the weather.") + .thought(true) + .thoughtSignature(new byte[] {1, 2, 3}) + .build(), + Part.builder() + .functionCall( + FunctionCall.builder() + .id("call-1") + .name("get_weather") + .args(ImmutableMap.of("city", "Seoul")) + .build()) + .thoughtSignature(new byte[] {4, 5, 6}) + .build()) + .build()) + .build(); + Event result = + Event.builder() + .id("event-3") + .invocationId("invocation-1") + .author("weather_agent") + .timestamp(3_000L) + .content( + Content.builder() + .role("user") + .parts( + Part.builder() + .functionResponse( + FunctionResponse.builder() + .id("call-1") + .name("get_weather") + .response(ImmutableMap.of("temp_c", 21)) + .build()) + .build()) + .build()) + .actions(EventActions.builder().skipSummarization(true).build()) + .build(); + ImmutableList appended = ImmutableList.of(question, call, result); + when(mockSessionsCollection.document(SESSION_ID)).thenReturn(mockSessionDocRef); + when(mockEventsCollection.document()).thenReturn(mockEventDocRef); + when(mockEventDocRef.getId()).thenReturn(EVENT_ID); + when(mockEventsCollection.document(EVENT_ID)).thenReturn(mockEventDocRef); + for (Event event : appended) { + sessionService.appendEvent(session, event).test().assertComplete(); + } + + // Serve the stored event documents back to getSession. + ArgumentCaptor> eventDataCaptor = ArgumentCaptor.forClass(Map.class); + verify(mockEventDocRef, times(3)).set(eventDataCaptor.capture()); + List eventDocs = new ArrayList<>(); + for (Map eventData : eventDataCaptor.getAllValues()) { + QueryDocumentSnapshot eventDoc = mock(QueryDocumentSnapshot.class); + when(eventDoc.getData()).thenReturn(eventData); + eventDocs.add(eventDoc); + } + when(mockSessionDocRef.get()).thenReturn(ApiFutures.immediateFuture(mockSessionSnapshot)); + when(mockSessionSnapshot.exists()).thenReturn(true); + when(mockSessionSnapshot.getReference()).thenReturn(mockSessionDocRef); + when(mockSessionSnapshot.getData()) + .thenReturn( + ImmutableMap.of( + "id", + SESSION_ID, + "appName", + APP_NAME, + "userId", + USER_ID, + "updateTime", + NOW.toString(), + "state", + Collections.emptyMap())); + when(mockQuery.get()).thenReturn(ApiFutures.immediateFuture(mockQuerySnapshot)); + when(mockQuerySnapshot.getDocuments()).thenReturn(eventDocs); + + List reloaded = + sessionService + .getSession(APP_NAME, USER_ID, SESSION_ID, Optional.empty()) + .blockingGet() + .events(); + + // Compared as JSON because Part.equals compares thought signatures by reference. + assertThat(reloaded.stream().map(Event::toJson).toList()) + .containsExactlyElementsIn(appended.stream().map(Event::toJson).toList()) + .inOrder(); + } + + /** Tests that a tool confirmation request is stored with plain values in its args. */ + @Test + @SuppressWarnings("unchecked") + void appendEvent_withToolConfirmationRequest_storesArgsAsPlainValues() { + Session session = + Session.builder(SESSION_ID) + .appName(APP_NAME) + .userId(USER_ID) + .state(new ConcurrentHashMap<>()) + .build(); + FunctionCall toolCall = + FunctionCall.builder() + .id("call-1") + .name("get_weather") + .args(ImmutableMap.of("city", "Seoul")) + .build(); + // Shaped like the event Functions.generateRequestConfirmationEvent creates. + Event confirmationRequest = + Event.builder() + .author("weather_agent") + .content( + Content.fromParts( + Part.builder() + .functionCall( + FunctionCall.builder() + .id("call-2") + .name("adk_request_confirmation") + .args( + ImmutableMap.of( + "originalFunctionCall", + toolCall, + "toolConfirmation", + ToolConfirmation.builder().hint("Approve?").build())) + .build()) + .build())) + .build(); + when(mockSessionsCollection.document(SESSION_ID)).thenReturn(mockSessionDocRef); + when(mockEventsCollection.document()).thenReturn(mockEventDocRef); + when(mockEventDocRef.getId()).thenReturn(EVENT_ID); + when(mockEventsCollection.document(EVENT_ID)).thenReturn(mockEventDocRef); + + sessionService.appendEvent(session, confirmationRequest).test().assertComplete(); + + // Firestore cannot encode FunctionCall or ToolConfirmation objects. + ArgumentCaptor> eventDataCaptor = ArgumentCaptor.forClass(Map.class); + verify(mockEventDocRef).set(eventDataCaptor.capture()); + Map content = (Map) eventDataCaptor.getValue().get("content"); + Map part = ((List>) content.get("parts")).get(0); + Map args = + (Map) ((Map) part.get("functionCall")).get("args"); + assertThat(args.get("originalFunctionCall")) + .isEqualTo( + ImmutableMap.of( + "id", "call-1", "name", "get_weather", "args", ImmutableMap.of("city", "Seoul"))); + assertThat((Map) args.get("toolConfirmation")) + .containsEntry("hint", "Approve?"); + } + + /** A tool return value that Firestore cannot encode, because of its Optional. */ + record Alert(String city, Optional warning) {} + + /** Tests that objects in a tool result are stored as plain values. */ + @Test + @SuppressWarnings("unchecked") + void appendEvent_withObjectsInToolResult_storesResultAsPlainValues() { + Session session = + Session.builder(SESSION_ID) + .appName(APP_NAME) + .userId(USER_ID) + .state(new ConcurrentHashMap<>()) + .build(); + // FunctionTool returns {"result": value} for a value it cannot convert to a map, like a list. + Event toolResult = + Event.builder() + .author("weather_agent") + .content( + Content.fromParts( + Part.fromFunctionResponse( + "get_alerts", + ImmutableMap.of( + "result", + ImmutableList.of(new Alert("Seoul", Optional.of("Heavy rain"))))))) + .build(); + when(mockSessionsCollection.document(SESSION_ID)).thenReturn(mockSessionDocRef); + when(mockEventsCollection.document()).thenReturn(mockEventDocRef); + when(mockEventDocRef.getId()).thenReturn(EVENT_ID); + when(mockEventsCollection.document(EVENT_ID)).thenReturn(mockEventDocRef); + + sessionService.appendEvent(session, toolResult).test().assertComplete(); + + ArgumentCaptor> eventDataCaptor = ArgumentCaptor.forClass(Map.class); + verify(mockEventDocRef).set(eventDataCaptor.capture()); + Map content = (Map) eventDataCaptor.getValue().get("content"); + Map part = ((List>) content.get("parts")).get(0); + assertThat(((Map) part.get("functionResponse")).get("response")) + .isEqualTo( + ImmutableMap.of( + "result", + ImmutableList.of(ImmutableMap.of("city", "Seoul", "warning", "Heavy rain")))); + } + /** Tests that appendEvent with only app state deltas updates the correct stores. */ @Test void appendEvent_withAppOnlyStateDeltas_updatesCorrectStores() { @@ -913,6 +1136,52 @@ void getSession_withUnparseableEvent_triggersCatchBlockAndSkipsEvent() { .assertValue(session -> session.events().isEmpty()); // The bad event should be skipped } + /** Tests that an event whose rawEvent cannot be read is loaded from its other fields. */ + @Test + void getSession_withUnreadableRawEvent_loadsEventFromOtherFields() { + QueryDocumentSnapshot eventDoc = mock(QueryDocumentSnapshot.class); + when(eventDoc.getData()) + .thenReturn( + ImmutableMap.of( + Constants.KEY_AUTHOR, + "weather_agent", + Constants.KEY_TIMESTAMP, + Instant.ofEpochMilli(1_000L).toString(), + "content", + ImmutableMap.of("parts", ImmutableList.of(ImmutableMap.of("text", "It is sunny."))), + FirestoreSessionService.RAW_EVENT_KEY, + "{")); + when(mockSessionsCollection.document(SESSION_ID)).thenReturn(mockSessionDocRef); + when(mockSessionDocRef.get()).thenReturn(ApiFutures.immediateFuture(mockSessionSnapshot)); + when(mockSessionSnapshot.exists()).thenReturn(true); + when(mockSessionSnapshot.getReference()).thenReturn(mockSessionDocRef); + when(mockSessionSnapshot.getData()) + .thenReturn( + ImmutableMap.of( + "id", + SESSION_ID, + "appName", + APP_NAME, + "userId", + USER_ID, + "updateTime", + NOW.toString(), + "state", + Collections.emptyMap())); + when(mockQuery.get()).thenReturn(ApiFutures.immediateFuture(mockQuerySnapshot)); + when(mockQuerySnapshot.getDocuments()).thenReturn(ImmutableList.of(eventDoc)); + + List events = + sessionService + .getSession(APP_NAME, USER_ID, SESSION_ID, Optional.empty()) + .blockingGet() + .events(); + + assertThat(events).hasSize(1); + assertThat(events.get(0).author()).isEqualTo("weather_agent"); + assertThat(events.get(0).content().get().parts().get().get(0).text()).hasValue("It is sunny."); + } + // --- listEvents Tests --- /** Tests that listEvents returns the expected events when the session exists. */