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; + } }