fix(sessions): preserve event append order on tied timestamps (#8051)

src/google/adk/sessions/database_session_service.py

@@ -37,6 +37,7 @@
 try:
   from sqlalchemy import delete
   from sqlalchemy import event
+  from sqlalchemy import func
   from sqlalchemy import inspect
   from sqlalchemy import MetaData
   from sqlalchemy import select
@@ -286,9 +287,41 @@ def _ensure_schema_indexes_exist(
           raise
 
 
+def _ensure_schema_columns_exist(connection: Connection) -> None:
+  """Ensures columns declared in metadata exist for existing tables."""
+  inspector = inspect(connection)
+  if not inspector.has_table("events"):
+    return
+  existing_columns = {col["name"] for col in inspector.get_columns("events")}
+  if "event_index" not in existing_columns:
+    logger.debug("Adding 'event_index' column to existing 'events' table.")
+    dialect = connection.dialect.name
+    if dialect == "spanner":
+      stmt = "ALTER TABLE events ADD COLUMN event_index INT64"
+    elif dialect in ("mssql", "sqlserver"):
+      stmt = "ALTER TABLE events ADD event_index INT DEFAULT 0 NOT NULL"
+    elif dialect == "mysql":
+      stmt = "ALTER TABLE events ADD COLUMN event_index INT DEFAULT 0 NOT NULL"
+    else:
+      stmt = (
+          "ALTER TABLE events ADD COLUMN event_index INTEGER DEFAULT 0 NOT NULL"
+      )
+    try:
+      with connection.begin_nested():
+        connection.exec_driver_sql(stmt)
+    except (OperationalError, ProgrammingError):
+      refreshed = {
+          col["name"] for col in inspect(connection).get_columns("events")
+      }
+      if "event_index" not in refreshed:
+        raise
+
+
 def _setup_database_schema(connection: Connection, metadata: MetaData) -> None:
   """Ensures tables and indexes declared in metadata exist."""
   metadata.create_all(bind=connection)
+  if "event_index" in metadata.tables["events"].c:
+    _ensure_schema_columns_exist(connection)
   _ensure_schema_indexes_exist(connection, metadata)
 
 
@@ -677,14 +710,17 @@ async def _session_matches_storage_revision(
       result = await sql_session.execute(stmt)
       return result.scalar_one_or_none() is None
 
+    order_by_clauses = [schema.StorageEvent.timestamp.desc()]
+    if hasattr(schema.StorageEvent, "event_index"):
+      order_by_clauses.append(schema.StorageEvent.event_index.desc())
+    order_by_clauses.append(schema.StorageEvent.id.desc())
+
     stmt = (
         select(schema.StorageEvent.id)
         .filter(schema.StorageEvent.app_name == session.app_name)
         .filter(schema.StorageEvent.session_id == session.id)
         .filter(schema.StorageEvent.user_id == session.user_id)
-        .order_by(
-            schema.StorageEvent.timestamp.desc(), schema.StorageEvent.id.desc()
-        )
+        .order_by(*order_by_clauses)
         .limit(1)
     )
     result = await sql_session.execute(stmt)
@@ -839,13 +875,16 @@ async def get_session(
             after_dt = after_dt.replace(tzinfo=None)
           stmt = stmt.filter(schema.StorageEvent.timestamp >= after_dt)
 
-        # Break timestamp ties on id, matching the ordering the stale-session
-        # check uses. Without it the database is free to return tied events in
-        # a different order on every read, so a replayed conversation shuffles
-        # and `num_recent_events` truncates at an arbitrary point in the tie.
-        stmt = stmt.order_by(
-            schema.StorageEvent.timestamp.desc(), schema.StorageEvent.id.desc()
-        )
+        # Break timestamp ties on event_index, matching the ordering the
+        # stale-session check uses. Without it the database is free to return
+        # tied events in a different order on every read, so a replayed
+        # conversation shuffles and `num_recent_events` truncates at an
+        # arbitrary point in the tie.
+        order_by_clauses = [schema.StorageEvent.timestamp.desc()]
+        if hasattr(schema.StorageEvent, "event_index"):
+          order_by_clauses.append(schema.StorageEvent.event_index.desc())
+        order_by_clauses.append(schema.StorageEvent.id.desc())
+        stmt = stmt.order_by(*order_by_clauses)
 
         if config and config.num_recent_events is not None:
           stmt = stmt.limit(config.num_recent_events)
@@ -1083,11 +1122,33 @@ async def append_event(self, session: Session, event: Event) -> Event:
         if state_deltas["session"]:
           storage_session.state.update(state_deltas["session"])
 
+        event_index = 0
+        if hasattr(schema.StorageEvent, "event_index"):
+          next_index_stmt = (
+              select(
+                  func.coalesce(func.max(schema.StorageEvent.event_index), -1)
+                  + 1
+              )
+              .filter(schema.StorageEvent.app_name == session.app_name)
+              .filter(schema.StorageEvent.user_id == session.user_id)
+              .filter(schema.StorageEvent.session_id == session.id)
+          )
+          event_index = (
+              await sql_session.execute(next_index_stmt)
+          ).scalar_one()
+
         update_time = datetime.fromtimestamp(event.timestamp, timezone.utc)
         if self._uses_naive_datetime():
           update_time = update_time.replace(tzinfo=None)
         storage_session.update_time = update_time
-        sql_session.add(schema.StorageEvent.from_event(session, event))
+        if hasattr(schema.StorageEvent, "event_index"):
+          sql_session.add(
+              schema.StorageEvent.from_event(
+                  session, event, event_index=event_index
+              )
+          )
+        else:
+          sql_session.add(schema.StorageEvent.from_event(session, event))
 
         # Read revision fields before commit. Post-commit ORM attribute access
         # can lazy-load expired columns and trigger MissingGreenlet with asyncpg

src/google/adk/sessions/migration/migrate_from_sqlalchemy_pickle.py

@@ -323,6 +323,7 @@ def migrate(
       logger.info("Migrating events...")
       num_rows = 0
       if inspector.has_table("events"):
+        session_event_counters: dict[tuple[str, str, str], int] = {}
         for row in source_session.execute(
             text("SELECT * FROM events")
         ).mappings():
@@ -331,6 +332,9 @@ def migrate(
                 dict(row),
                 allow_unsafe_unpickling=allow_unsafe_unpickling,
             )
+            session_key = (row["app_name"], row["user_id"], row["session_id"])
+            event_index = session_event_counters.get(session_key, 0)
+            session_event_counters[session_key] = event_index + 1
             new_event = v1.StorageEvent(
                 id=event_obj.id,
                 app_name=row["app_name"],
@@ -341,6 +345,7 @@ def migrate(
                     event_obj.timestamp, timezone.utc
                 ).replace(tzinfo=None),
                 event_data=event_obj.model_dump(mode="json", exclude_none=True),
+                event_index=event_index,
             )
             dest_session.merge(new_event)
             num_rows += 1

src/google/adk/sessions/schemas/v1.py

@@ -30,6 +30,8 @@
 from sqlalchemy import ForeignKeyConstraint
 from sqlalchemy import func
 from sqlalchemy import Index
+from sqlalchemy import Integer
+from sqlalchemy import text
 from sqlalchemy.ext.mutable import MutableDict
 from sqlalchemy.orm import DeclarativeBase
 from sqlalchemy.orm import Mapped
@@ -179,6 +181,9 @@ class StorageEvent(Base):
   )
 
   invocation_id: Mapped[str] = mapped_column(String(DEFAULT_MAX_VARCHAR_LENGTH))
+  event_index: Mapped[int] = mapped_column(
+      Integer, default=0, server_default=text("0"), nullable=False
+  )
   timestamp: Mapped[datetime] = mapped_column(
       PreciseTimestamp, default=func.now()
   )
@@ -210,7 +215,9 @@ class StorageEvent(Base):
   )
 
   @classmethod
-  def from_event(cls, session: Session, event: Event) -> StorageEvent:
+  def from_event(
+      cls, session: Session, event: Event, event_index: int = 0
+  ) -> StorageEvent:
     """Creates a StorageEvent from an Event."""
     return StorageEvent(
         id=event.id,
@@ -220,6 +227,7 @@ def from_event(cls, session: Session, event: Event) -> StorageEvent:
         user_id=session.user_id,
         timestamp=timestamp_to_utc_datetime(event.timestamp),
         event_data=event.model_dump(exclude_none=True, mode="json"),
+        event_index=event_index,
     )
 
   def to_event(self) -> Event:

tests/unittests/sessions/test_session_service.py

@@ -4488,7 +4488,7 @@ async def test_get_session_keeps_exact_epoch_across_a_repeated_local_hour(
 async def test_get_session_orders_tied_timestamps_by_id(
     append_ids_in_reverse, tmp_path
 ):
-  """DatabaseSessionService breaks tied timestamps by id, not append order."""
+  """DatabaseSessionService returns tied-timestamp events in append order."""
   app_name = 'my_app'
   user_id = 'user'
   event_ids = ['event_a', 'event_m', 'event_z']
@@ -4512,7 +4512,98 @@ async def test_get_session_orders_tied_timestamps_by_id(
   finally:
     await service.close()
 
-  assert [event.id for event in retrieved_session.events] == event_ids
+  assert [event.id for event in retrieved_session.events] == append_order
+
+
+@pytest.mark.asyncio
+async def test_database_session_service_multiple_events_tied_timestamp_order(
+    tmp_path,
+):
+  """DatabaseSessionService reloads multiple events with equal timestamps in FIFO order."""
+  app_name = 'my_app'
+  user_id = 'user'
+  service = get_session_service(SessionServiceType.DATABASE, tmp_path)
+  try:
+    session = await service.create_session(app_name=app_name, user_id=user_id)
+    shared_timestamp = 1234567.89
+    # Append 8 events with random IDs and identical timestamps
+    event_ids = [
+        'id_7_zeta',
+        'id_3_gamma',
+        'id_5_epsilon',
+        'id_8_theta',
+        'id_2_beta',
+        'id_1_alpha',
+        'id_4_delta',
+        'id_6_eta',
+    ]
+    for event_id in event_ids:
+      await service.append_event(
+          session,
+          Event(author='agent', id=event_id, timestamp=shared_timestamp),
+      )
+
+    retrieved_session = await service.get_session(
+        app_name=app_name, user_id=user_id, session_id=session.id
+    )
+    assert [e.id for e in retrieved_session.events] == event_ids
+  finally:
+    await service.close()
+
+
+@pytest.mark.asyncio
+async def test_database_session_service_migration_adds_event_index_column(
+    tmp_path,
+):
+  """DatabaseSessionService automatically migrates existing tables missing event_index."""
+  db_path = str(tmp_path / 'legacy.db')
+  db_url = f'sqlite+aiosqlite:///{db_path}'
+
+  # 1. Create a legacy database table without event_index
+  import aiosqlite
+
+  async with aiosqlite.connect(db_path) as db:
+    await db.execute("""
+      CREATE TABLE sessions (
+        app_name VARCHAR(128) NOT NULL,
+        user_id VARCHAR(128) NOT NULL,
+        id VARCHAR(128) NOT NULL,
+        state TEXT,
+        create_time TIMESTAMP,
+        update_time TIMESTAMP,
+        PRIMARY KEY (app_name, user_id, id)
+      )
+    """)
+    await db.execute("""
+      CREATE TABLE events (
+        id VARCHAR(128) NOT NULL,
+        app_name VARCHAR(128) NOT NULL,
+        user_id VARCHAR(128) NOT NULL,
+        session_id VARCHAR(128) NOT NULL,
+        invocation_id VARCHAR(256),
+        timestamp TIMESTAMP,
+        event_data TEXT,
+        PRIMARY KEY (id, app_name, user_id, session_id)
+      )
+    """)
+    await db.commit()
+
+  # 2. Open with DatabaseSessionService; prepare_tables should migrate and add event_index
+  service = DatabaseSessionService(db_url)
+  try:
+    session = await service.create_session(app_name='app', user_id='user')
+    stamp = 999.0
+    for eid in ['e3', 'e1', 'e2']:
+      await service.append_event(
+          session, Event(author='user', id=eid, timestamp=stamp)
+      )
+
+    reloaded = await service.get_session(
+        app_name='app', user_id='user', session_id=session.id
+    )
+    assert [e.id for e in reloaded.events] == ['e3', 'e1', 'e2']
+  finally:
+    await service.close()
 
 
 @pytest.mark.asyncio
@@ -4991,3 +5082,27 @@ async def test_in_memory_append_event_copies_given_stored_session():
       'app:config': {'flags': ['v1']},
   }
   assert other_user_session.state == {'app:config': {'flags': ['v1']}}
+
+
+@pytest.mark.asyncio
+async def test_v0_database_is_not_altered_with_event_index(tmp_path):
+  """A legacy V0 events table must not get the V1 only event_index column."""
+  db_path = tmp_path / 'legacy_v0.db'
+  db_url = f'sqlite+aiosqlite:///{db_path}'
+  engine = create_async_engine(db_url)
+  async with engine.begin() as conn:
+    await conn.run_sync(v0.Base.metadata.create_all)
+  await engine.dispose()
+
+  service = DatabaseSessionService(db_url)
+  try:
+    session = await service.create_session(app_name='app', user_id='user')
+    await service.append_event(
+        session, Event(author='user', id='e1', timestamp=1.0)
+    )
+  finally:
+    await service.close()
+
+  with sqlite3.connect(db_path) as conn:
+    columns = {row[1] for row in conn.execute('PRAGMA table_info(events)')}
+  assert 'event_index' not in columns