From 1a6dd29e41c6f21f81b8fbef051a330afd2da6e3 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 >=. - An afterTimestamp later than any event can have (event timestamps are epoch milliseconds in a long) returns no events, as InMemorySessionService does, instead of overflowing when rounded up near Instant.MAX. - 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 | 48 +++- .../sessions/FirestoreSessionServiceTest.java | 215 +++++++++++++++++- 2 files changed, 247 insertions(+), 16 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..862c37813 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,17 @@ 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; + /** + * Formats stored event timestamps with exactly three fraction digits, so text order is time + * order. The field stays text, unlike the timestamp in ADK Python, because Firestore orders + * timestamp values before strings, so switching types would misorder existing sessions. + */ + private static final DateTimeFormatter TIMESTAMP_FORMATTER = + new DateTimeFormatterBuilder().appendInstant(3).toFormatter(); + + /** The latest timestamp an event can have, since {@link Event#timestamp()} is epoch millis. */ + private static final Instant LATEST_EVENT_TIMESTAMP = Instant.ofEpochMilli(Long.MAX_VALUE); + /** Random token each create stores, so a retried create can detect its own write. */ static final String CREATE_TOKEN_KEY = "createToken"; @@ -233,10 +248,21 @@ public Maybe getSession( document.getReference().collection(EVENTS_SUBCOLLECTION_NAME); Query eventsQuery = eventsCollection.orderBy(TIMESTAMP_KEY); - if (config.afterTimestamp().isPresent()) { + // No event can be at or after a later afterTimestamp, and rounding it up below could + // overflow near Instant.MAX. + boolean afterLatestEvent = + config + .afterTimestamp() + .filter(t -> t.isAfter(LATEST_EVENT_TIMESTAMP)) + .isPresent(); + + if (config.afterTimestamp().isPresent() && !afterLatestEvent) { + // 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()) { @@ -244,10 +270,12 @@ public Maybe getSession( } ApiFuture> eventsFuture = - ApiFutures.transform( - eventsQuery.get(), - com.google.cloud.firestore.QuerySnapshot::getDocuments, - MoreExecutors.directExecutor()); + afterLatestEvent + ? ApiFutures.immediateFuture(List.of()) + : ApiFutures.transform( + eventsQuery.get(), + com.google.cloud.firestore.QuerySnapshot::getDocuments, + MoreExecutors.directExecutor()); return ApiFutureUtils.toSingle(eventsFuture) .map( @@ -259,6 +287,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 +438,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 +653,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..f341918c0 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,28 +270,147 @@ 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 from the same second for a whole-second afterTimestamp. */ @Test - void getSession_withAfterTimestamp_appliesFilterToQuery() { + void getSession_withAfterTimestampOnWholeSecond_queriesFromThatSecond() { + assertAfterTimestampQueriesFrom( + Instant.parse("2026-10-10T05:00:05Z"), "2026-10-10T05:00:05.000Z"); + } + + /** Tests that getSession queries from the same millisecond for a whole-millisecond cursor. */ + @Test + void getSession_withAfterTimestampOnMillisecond_queriesFromThatMillisecond() { + assertAfterTimestampQueriesFrom( + Instant.parse("2026-10-10T05:00:05.200Z"), "2026-10-10T05:00:05.200Z"); + } + + /** Tests that getSession rounds a sub-millisecond afterTimestamp up to the next millisecond. */ + @Test + void getSession_withAfterTimestampWithinMillisecond_queriesFromNextMillisecond() { + assertAfterTimestampQueriesFrom( + Instant.parse("2026-10-10T05:00:05.200300Z"), "2026-10-10T05:00:05.201Z"); + } + + /** Tests that getSession still queries events at the latest timestamp an event can have. */ + @Test + void getSession_withAfterTimestampAtLatestEventTimestamp_queriesEvents() { + // Arrange + stubSessionDocument(); + stubEventsQuery(); + GetSessionConfig config = + GetSessionConfig.builder().afterTimestamp(Instant.ofEpochMilli(Long.MAX_VALUE)).build(); + + // Act + TestObserver testObserver = + sessionService.getSession(APP_NAME, USER_ID, SESSION_ID, Optional.of(config)).test(); + + // Assert + testObserver.assertNoErrors(); + testObserver.assertComplete(); + verify(mockQuery).get(); + } + + /** Tests that getSession returns no events for an afterTimestamp later than any event. */ + @Test + void getSession_withAfterTimestampPastLatestEventTimestamp_returnsSessionWithoutEvents() { + assertAfterTimestampReturnsNoEvents(Instant.ofEpochMilli(Long.MAX_VALUE).plusNanos(1)); + } + + /** Tests that getSession returns no events, like InMemorySessionService, for Instant.MAX. */ + @Test + void getSession_withAfterTimestampAtInstantMax_returnsSessionWithoutEvents() { + assertAfterTimestampReturnsNoEvents(Instant.MAX); + } + + private void assertAfterTimestampQueriesFrom(Instant afterTimestamp, String storedCursor) { + // Arrange + stubSessionDocument(); + stubEventsQuery(); + GetSessionConfig config = GetSessionConfig.builder().afterTimestamp(afterTimestamp).build(); + + // Act + TestObserver testObserver = + sessionService.getSession(APP_NAME, USER_ID, SESSION_ID, Optional.of(config)).test(); + + // Assert + testObserver.assertNoErrors(); + testObserver.assertComplete(); + testObserver.assertValue(session -> session.id().equals(SESSION_ID)); + verify(mockQuery).whereGreaterThanOrEqualTo(Constants.KEY_TIMESTAMP, storedCursor); + } + + private void assertAfterTimestampReturnsNoEvents(Instant afterTimestamp) { // Arrange - Instant timestamp = Instant.now().minusSeconds(60); - GetSessionConfig config = GetSessionConfig.builder().afterTimestamp(timestamp).build(); + stubSessionDocument(); + GetSessionConfig config = GetSessionConfig.builder().afterTimestamp(afterTimestamp).build(); + + // Act + TestObserver testObserver = + sessionService.getSession(APP_NAME, USER_ID, SESSION_ID, Optional.of(config)).test(); + + // Assert + testObserver.assertNoErrors(); + testObserver.assertComplete(); + testObserver.assertValue( + session -> session.id().equals(SESSION_ID) && session.events().isEmpty()); + verify(mockQuery, never()).get(); + } + private void stubSessionDocument() { 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(Constants.KEY_APP_NAME, APP_NAME, "state", Map.of())); - when(mockQuery.whereGreaterThan(Constants.KEY_TIMESTAMP, timestamp.toString())) + .thenReturn( + ImmutableMap.of( + "id", SESSION_ID, + "appName", APP_NAME, + "userId", USER_ID, + "updateTime", NOW.toString(), + "state", Collections.emptyMap())); + } + + private void stubEventsQuery() { + when(mockQuery.whereGreaterThanOrEqualTo(eq(Constants.KEY_TIMESTAMP), anyString())) .thenReturn(mockQuery); when(mockQuery.get()).thenReturn(ApiFutures.immediateFuture(mockQuerySnapshot)); + when(mockQuerySnapshot.getDocuments()).thenReturn(Collections.emptyList()); + } + + /** 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 +693,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 +1107,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 +1313,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; + } }