From d79872441d23ed7d43b1159d5644acecf0455e95 Mon Sep 17 00:00:00 2001 From: Mohit Sharma Date: Fri, 13 Nov 2026 14:10:35 +0530 Subject: [PATCH 1/2] fix(sessions): preserve event append order on tied timestamps (#8051) --- .../adk/sessions/database_session_service.py | 82 +++++++++++++--- .../migrate_from_sqlalchemy_pickle.py | 5 + src/google/adk/sessions/schemas/v1.py | 10 +- .../sessions/test_session_service.py | 95 ++++++++++++++++++- 4 files changed, 178 insertions(+), 14 deletions(-) diff --git a/src/google/adk/sessions/database_session_service.py b/src/google/adk/sessions/database_session_service.py index 14442cca6c5..c2b1bd34b62 100644 --- a/src/google/adk/sessions/database_session_service.py +++ b/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,40 @@ 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) + _ensure_schema_columns_exist(connection) _ensure_schema_indexes_exist(connection, metadata) @@ -677,14 +709,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 +874,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 +1121,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 diff --git a/src/google/adk/sessions/migration/migrate_from_sqlalchemy_pickle.py b/src/google/adk/sessions/migration/migrate_from_sqlalchemy_pickle.py index e86db798f3d..7f238d12d06 100644 --- a/src/google/adk/sessions/migration/migrate_from_sqlalchemy_pickle.py +++ b/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 diff --git a/src/google/adk/sessions/schemas/v1.py b/src/google/adk/sessions/schemas/v1.py index 18a16b25583..74e4426a4ff 100644 --- a/src/google/adk/sessions/schemas/v1.py +++ b/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: diff --git a/tests/unittests/sessions/test_session_service.py b/tests/unittests/sessions/test_session_service.py index 2d748f384e9..b41f5fadcd8 100644 --- a/tests/unittests/sessions/test_session_service.py +++ b/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 From 4c14fbc72d2a056799a6e6dadc7ff08327d38682 Mon Sep 17 00:00:00 2001 From: Aman Agnihotri Date: Fri, 13 Nov 2026 13:27:44 +0000 Subject: [PATCH 2/2] Skip event_index column migration for V0 schema --- .../adk/sessions/database_session_service.py | 3 ++- .../sessions/test_session_service.py | 24 +++++++++++++++++++ 2 files changed, 26 insertions(+), 1 deletion(-) diff --git a/src/google/adk/sessions/database_session_service.py b/src/google/adk/sessions/database_session_service.py index c2b1bd34b62..6dcd5c4e049 100644 --- a/src/google/adk/sessions/database_session_service.py +++ b/src/google/adk/sessions/database_session_service.py @@ -320,7 +320,8 @@ def _ensure_schema_columns_exist(connection: Connection) -> None: def _setup_database_schema(connection: Connection, metadata: MetaData) -> None: """Ensures tables and indexes declared in metadata exist.""" metadata.create_all(bind=connection) - _ensure_schema_columns_exist(connection) + if "event_index" in metadata.tables["events"].c: + _ensure_schema_columns_exist(connection) _ensure_schema_indexes_exist(connection, metadata) diff --git a/tests/unittests/sessions/test_session_service.py b/tests/unittests/sessions/test_session_service.py index b41f5fadcd8..ec29047a26c 100644 --- a/tests/unittests/sessions/test_session_service.py +++ b/tests/unittests/sessions/test_session_service.py @@ -5082,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