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