Skip to content

fix: keep Firestore session events in time order #1653

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
innoprej wants to merge 1 commit into
base: main
Choose a base branch
from
+247 −16
Open

fix: keep Firestore session events in time order #1653

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
48 changes: 40 additions & 8 deletions ...estore-session-service/src/main/java/com/google/adk/sessions/FirestoreSessionService.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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";

Expand Down Expand Up @@ -233,21 +248,34 @@ public Maybe<Session> 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);

@damianmomotgoogle damianmomotgoogle Oct 27, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

plusNanos(999_999) overflows near Instant.MAX: an afterTimestamp of Instant.MAX throws DateTimeException instead of returning no events like InMemorySessionService. This is an edge case, but the overflow is easy to guard against.

@innoprej innoprej Oct 29, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for catching this. I added a guard for cursors later than the maximum event timestamp representable as a long epoch-millisecond value. In that case, getSession returns the session with no events without issuing the event query. The new tests cover both Instant.MAX and a cursor one nanosecond beyond the maximum event timestamp.

eventsQuery =
eventsQuery.whereGreaterThan(
TIMESTAMP_KEY, config.afterTimestamp().get().toString());
eventsQuery.whereGreaterThanOrEqualTo(
TIMESTAMP_KEY, TIMESTAMP_FORMATTER.format(cursor));
}

if (config.numRecentEvents().isPresent()) {
eventsQuery = eventsQuery.limitToLast(config.numRecentEvents().get());
}

ApiFuture<List<QueryDocumentSnapshot>> 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(
Expand All @@ -259,6 +287,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(
Expand Down Expand Up @@ -408,7 +438,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<>();
Expand Down Expand Up @@ -623,6 +653,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();
});
Expand Down
215 changes: 207 additions & 8 deletions ...re-session-service/src/test/java/com/google/adk/sessions/FirestoreSessionServiceTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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<Session> 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<Session> 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<Session> 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<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 ---
Expand Down Expand Up @@ -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<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() {
Expand Down Expand Up @@ -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<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() {
Expand Down Expand Up @@ -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;
}
}
Loading