fix: keep Firestore session events in time order

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<Session> 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<Session> 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<String, Object> 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<String, Object> contentData = new HashMap<>();
@@ -623,6 +636,8 @@ public Single<ListEventsResponse> 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();
         });

contrib/firestore-session-service/src/test/java/com/google/adk/sessions/FirestoreSessionServiceTest.java

@@ -270,28 +270,72 @@ 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<String, String> 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));
     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()))
+    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<QueryDocumentSnapshot> 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<Map<String, Object>> eventCaptor = ArgumentCaptor.forClass(Map.class);
+    verify(mockEventDocRef, times(3)).set(eventCaptor.capture());
+    List<String> 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<QueryDocumentSnapshot> 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;
+  }
 }