From 153a150a2c214c71c6bcb930c5b30b276e26aca5 Mon Sep 17 00:00:00 2001 From: Jaehee Shin Date: Mon, 12 Oct 2026 22:22:50 +0900 Subject: [PATCH] fix: keep Firestore session events in time order FirestoreSessionService stored event timestamps with Instant.toString(), which leaves out the fraction when it is zero, and ordered and filtered events by comparing those strings. Within one second, an event at .000 sorted after the later events, a whole-second afterTimestamp cursor left out the rest of that second, and afterTimestamp used > where InMemorySessionService, VertexAiSessionService and ADK Python keep events at or after the cursor. - Event timestamps are written with exactly three fraction digits. - afterTimestamp is rounded up to whole milliseconds, written the same way and matched with >=. - getSession and listEvents sort the events they read by timestamp, so events stored in the old format also come back in time order. The field stays a string. ADK Python stores a Firestore timestamp, but Firestore orders timestamps before strings, so sessions with events stored before this change would be misordered. --- .../adk/sessions/FirestoreSessionService.java | 21 ++- .../sessions/FirestoreSessionServiceTest.java | 138 +++++++++++++++++- 2 files changed, 149 insertions(+), 10 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..b4558c36d 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 @@ -40,7 +40,11 @@ import io.reactivex.rxjava3.core.Single; import java.time.Instant; import java.time.InstantSource; +import java.time.format.DateTimeFormatter; +import java.time.format.DateTimeFormatterBuilder; +import java.time.temporal.ChronoUnit; import java.util.ArrayList; +import java.util.Comparator; import java.util.HashMap; import java.util.HashSet; import java.util.List; @@ -80,6 +84,10 @@ public class FirestoreSessionService implements BaseSessionService { private static final String UPDATE_TIME_KEY = Constants.KEY_UPDATE_TIME; private static final String TIMESTAMP_KEY = Constants.KEY_TIMESTAMP; + /** Text, not a timestamp as in ADK Python, with three fraction digits so it sorts by time. */ + private static final DateTimeFormatter TIMESTAMP_FORMATTER = + new DateTimeFormatterBuilder().appendInstant(3).toFormatter(); + /** Random token each create stores, so a retried create can detect its own write. */ static final String CREATE_TOKEN_KEY = "createToken"; @@ -234,9 +242,12 @@ public Maybe getSession( Query eventsQuery = eventsCollection.orderBy(TIMESTAMP_KEY); if (config.afterTimestamp().isPresent()) { + // Round up to whole milliseconds, the precision of stored event timestamps. + Instant cursor = + config.afterTimestamp().get().plusNanos(999_999).truncatedTo(ChronoUnit.MILLIS); eventsQuery = - eventsQuery.whereGreaterThan( - TIMESTAMP_KEY, config.afterTimestamp().get().toString()); + eventsQuery.whereGreaterThanOrEqualTo( + TIMESTAMP_KEY, TIMESTAMP_FORMATTER.format(cursor)); } if (config.numRecentEvents().isPresent()) { @@ -259,6 +270,8 @@ public Maybe getSession( events.add(event); } } + // Events stored before fixed-width timestamps can come back out of order. + events.sort(Comparator.comparingLong(Event::timestamp)); return events; }) .map( @@ -408,7 +421,7 @@ private Map eventToMap(Session session, Event event) { } else { data.put("author", event.author()); } - data.put(TIMESTAMP_KEY, Instant.ofEpochMilli(event.timestamp()).toString()); + data.put(TIMESTAMP_KEY, TIMESTAMP_FORMATTER.format(Instant.ofEpochMilli(event.timestamp()))); data.put(APP_NAME_KEY, session.appName()); // Persist appName with the event Map contentData = new HashMap<>(); @@ -623,6 +636,8 @@ public Single listEvents(String appName, String userId, Stri events.add(event); } } + // Events stored before fixed-width timestamps can come back out of order. + events.sort(Comparator.comparingLong(Event::timestamp)); logger.info("Returning {} events for sessionId: {}", events.size(), sessionId); return ListEventsResponse.builder().events(events).build(); }); 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..84b3dd322 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 @@ -270,12 +270,15 @@ void getSession_withNumRecentEvents_appliesLimitToQuery() { verify(mockQuery).limitToLast(numEvents); } - /** Tests that getSession applies the afterTimestamp filter to the events query when specified. */ + /** Tests that getSession queries at or after afterTimestamp rounded up to whole milliseconds. */ @Test - void getSession_withAfterTimestamp_appliesFilterToQuery() { + void getSession_withAfterTimestamp_queriesAtOrAfterRoundedCursor() { // Arrange - Instant timestamp = Instant.now().minusSeconds(60); - GetSessionConfig config = GetSessionConfig.builder().afterTimestamp(timestamp).build(); + ImmutableMap storedCursorByAfterTimestamp = + ImmutableMap.of( + "2026-10-10T05:00:05Z", "2026-10-10T05:00:05.000Z", + "2026-10-10T05:00:05.200Z", "2026-10-10T05:00:05.200Z", + "2026-10-10T05:00:05.200300Z", "2026-10-10T05:00:05.201Z"); when(mockSessionsCollection.document(SESSION_ID)).thenReturn(mockSessionDocRef); when(mockSessionDocRef.get()).thenReturn(ApiFutures.immediateFuture(mockSessionSnapshot)); @@ -283,15 +286,56 @@ void getSession_withAfterTimestamp_appliesFilterToQuery() { when(mockSessionSnapshot.getReference()).thenReturn(mockSessionDocRef); when(mockSessionSnapshot.getData()) .thenReturn(ImmutableMap.of(Constants.KEY_APP_NAME, APP_NAME, "state", Map.of())); - when(mockQuery.whereGreaterThan(Constants.KEY_TIMESTAMP, timestamp.toString())) + when(mockQuery.whereGreaterThanOrEqualTo(eq(Constants.KEY_TIMESTAMP), anyString())) .thenReturn(mockQuery); when(mockQuery.get()).thenReturn(ApiFutures.immediateFuture(mockQuerySnapshot)); + storedCursorByAfterTimestamp.forEach( + (afterTimestamp, storedCursor) -> { + GetSessionConfig config = + GetSessionConfig.builder().afterTimestamp(Instant.parse(afterTimestamp)).build(); + + // Act + sessionService.getSession(APP_NAME, USER_ID, SESSION_ID, Optional.of(config)).test(); + + // Assert + verify(mockQuery).whereGreaterThanOrEqualTo(Constants.KEY_TIMESTAMP, storedCursor); + }); + } + + /** Tests that getSession returns events in time order even if they were stored out of order. */ + @Test + void getSession_eventsStoredOutOfOrder_returnsEventsInTimeOrder() { + // Arrange + 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)); + // Timestamps written before they had fixed millisecond digits, in Firestore's text order. + ImmutableList eventDocs = + ImmutableList.of( + eventSnapshot("2026-10-10T05:00:05.400Z"), eventSnapshot("2026-10-10T05:00:05Z")); + when(mockQuerySnapshot.getDocuments()).thenReturn(eventDocs); + // Act - sessionService.getSession(APP_NAME, USER_ID, SESSION_ID, Optional.of(config)).test(); + Session session = + sessionService.getSession(APP_NAME, USER_ID, SESSION_ID, Optional.empty()).blockingGet(); // Assert - verify(mockQuery).whereGreaterThan(Constants.KEY_TIMESTAMP, timestamp.toString()); + assertThat(session.events().stream().map(Event::timestamp).toList()) + .containsExactly( + Instant.parse("2026-10-10T05:00:05Z").toEpochMilli(), + Instant.parse("2026-10-10T05:00:05.400Z").toEpochMilli()) + .inOrder(); } // --- createSession Tests --- @@ -574,6 +618,43 @@ void appendEvent_persistsEventAndUpdatesSession() { assertThat(keywords).containsExactly("hello", "world"); } + /** Tests that appendEvent writes three fraction digits, so stored timestamps sort as text. */ + @Test + @SuppressWarnings("unchecked") + void appendEvent_storesTimestampsWithThreeFractionDigits() { + // Arrange + Session session = + Session.builder(SESSION_ID) + .appName(APP_NAME) + .userId(USER_ID) + .state(new ConcurrentHashMap<>()) + .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); + + // Act + for (String time : + List.of("2026-10-10T05:00:05Z", "2026-10-10T05:00:05.400Z", "2026-10-10T05:00:06Z")) { + Event event = + Event.builder().author("agent").timestamp(Instant.parse(time).toEpochMilli()).build(); + sessionService.appendEvent(session, event).blockingGet(); + } + + // Assert + ArgumentCaptor> eventCaptor = ArgumentCaptor.forClass(Map.class); + verify(mockEventDocRef, times(3)).set(eventCaptor.capture()); + List storedTimestamps = + eventCaptor.getAllValues().stream() + .map(data -> (String) data.get(Constants.KEY_TIMESTAMP)) + .toList(); + assertThat(storedTimestamps) + .containsExactly( + "2026-10-10T05:00:05.000Z", "2026-10-10T05:00:05.400Z", "2026-10-10T05:00:06.000Z") + .inOrder(); + } + /** Tests that appendAndGet correctly serializes and deserializes events with all part types. */ @Test void appendAndGet_withAllPartTypes_serializesAndDeserializesCorrectly() { @@ -951,6 +1032,34 @@ void listEvents_sessionExists_returnsEvents() { }); } + /** Tests that listEvents returns events in time order even if they were stored out of order. */ + @Test + void listEvents_eventsStoredOutOfOrder_returnsEventsInTimeOrder() { + // Arrange + when(mockSessionsCollection.document(SESSION_ID)).thenReturn(mockSessionDocRef); + when(mockSessionDocRef.get()).thenReturn(ApiFutures.immediateFuture(mockSessionSnapshot)); + when(mockSessionSnapshot.exists()).thenReturn(true); + when(mockSessionSnapshot.get(Constants.KEY_APP_NAME)).thenReturn(APP_NAME); + when(mockSessionSnapshot.getReference()).thenReturn(mockSessionDocRef); + when(mockQuery.get()).thenReturn(ApiFutures.immediateFuture(mockQuerySnapshot)); + // Timestamps written before they had fixed millisecond digits, in Firestore's text order. + ImmutableList eventDocs = + ImmutableList.of( + eventSnapshot("2026-10-10T05:00:05.400Z"), eventSnapshot("2026-10-10T05:00:05Z")); + when(mockQuerySnapshot.getDocuments()).thenReturn(eventDocs); + + // Act + ListEventsResponse response = + sessionService.listEvents(APP_NAME, USER_ID, SESSION_ID).blockingGet(); + + // Assert + assertThat(response.events().stream().map(Event::timestamp).toList()) + .containsExactly( + Instant.parse("2026-10-10T05:00:05Z").toEpochMilli(), + Instant.parse("2026-10-10T05:00:05.400Z").toEpochMilli()) + .inOrder(); + } + /** Tests that listEvents throws SessionNotFoundException when the session does not exist. */ @Test void listEvents_sessionDoesNotExist_throwsException() { @@ -1129,4 +1238,19 @@ private static AlreadyExistsException alreadyExists() { GrpcStatusCode.of(Status.Code.ALREADY_EXISTS), /* retryable= */ false); } + + /** Returns a stored event document with the given timestamp and no content parts. */ + private static QueryDocumentSnapshot eventSnapshot(String timestamp) { + QueryDocumentSnapshot snapshot = mock(QueryDocumentSnapshot.class); + when(snapshot.getData()) + .thenReturn( + ImmutableMap.of( + Constants.KEY_AUTHOR, + "agent", + Constants.KEY_TIMESTAMP, + timestamp, + "content", + ImmutableMap.of("parts", ImmutableList.of()))); + return snapshot; + } }