From 1063fa532cad59d8e9f7421ce2f523724d49d541 Mon Sep 17 00:00:00 2001 From: Shangjie Chen Date: Mon, 26 Jan 2026 14:18:42 -0800 Subject: [PATCH] fix: Reload the stale session in database_session_service when the update_time of storage is later than the current in memory session object This changes the behavior of how we handle stale session, instead of reject the transaction entirely, we aggresively refresh the session. Removed synchronous inspect(self).session calls within StorageSession.update_timestamp_tz and _dialect_name. This introspection was causing deadlocks/hangs when used with sqlalchemy.ext.asyncio in Python 3.13. Closes issue: https://github.com/google/adk-python/issues/1733 Co-authored-by: Shangjie Chen PiperOrigin-RevId: 861353058 --- .../adk/sessions/database_session_service.py | 61 ++++++++++++----- src/google/adk/sessions/schemas/v0.py | 14 ++-- src/google/adk/sessions/schemas/v1.py | 14 ++-- .../sessions/test_session_service.py | 67 +++++++++++++++---- 4 files changed, 104 insertions(+), 52 deletions(-) diff --git a/src/google/adk/sessions/database_session_service.py b/src/google/adk/sessions/database_session_service.py index 84977294..afc09b3e 100644 --- a/src/google/adk/sessions/database_session_service.py +++ b/src/google/adk/sessions/database_session_service.py @@ -31,7 +31,6 @@ from sqlalchemy.ext.asyncio import async_sessionmaker from sqlalchemy.ext.asyncio import AsyncEngine from sqlalchemy.ext.asyncio import AsyncSession as DatabaseSessionFactory from sqlalchemy.ext.asyncio import create_async_engine -from sqlalchemy.inspection import inspect from sqlalchemy.pool import StaticPool from typing_extensions import override from tzlocal import get_localzone @@ -275,22 +274,29 @@ class DatabaseSessionService(BaseSessionService): storage_user_state.state = storage_user_state.state | user_state_delta # Store the session + now = datetime.now(timezone.utc) + is_sqlite = self.db_engine.dialect.name == "sqlite" + if is_sqlite: + now = now.replace(tzinfo=None) + storage_session = schema.StorageSession( app_name=app_name, user_id=user_id, id=session_id, state=session_state, + create_time=now, + update_time=now, ) sql_session.add(storage_session) await sql_session.commit() - await sql_session.refresh(storage_session) - # Merge states for response merged_state = _merge_state( storage_app_state.state, storage_user_state.state, session_state ) - session = storage_session.to_session(state=merged_state) + session = storage_session.to_session( + state=merged_state, is_sqlite=is_sqlite + ) return session @override @@ -350,7 +356,10 @@ class DatabaseSessionService(BaseSessionService): # Convert storage session to session events = [e.to_event() for e in reversed(storage_events)] - session = storage_session.to_session(state=merged_state, events=events) + is_sqlite = self.db_engine.dialect.name == "sqlite" + session = storage_session.to_session( + state=merged_state, events=events, is_sqlite=is_sqlite + ) return session @override @@ -393,11 +402,14 @@ class DatabaseSessionService(BaseSessionService): user_states_map[storage_user_state.user_id] = storage_user_state.state sessions = [] + is_sqlite = self.db_engine.dialect.name == "sqlite" for storage_session in results: session_state = storage_session.state user_state = user_states_map.get(storage_session.user_id, {}) merged_state = _merge_state(app_state, user_state, session_state) - sessions.append(storage_session.to_session(state=merged_state)) + sessions.append( + storage_session.to_session(state=merged_state, is_sqlite=is_sqlite) + ) return ListSessionsResponse(sessions=sessions) @override @@ -433,15 +445,6 @@ class DatabaseSessionService(BaseSessionService): schema.StorageSession, (session.app_name, session.user_id, session.id) ) - if storage_session.update_timestamp_tz > session.last_update_time: - raise ValueError( - "The last_update_time provided in the session object" - f" {datetime.fromtimestamp(session.last_update_time):'%Y-%m-%d %H:%M:%S'} is" - " earlier than the update_time in the storage_session" - f" {datetime.fromtimestamp(storage_session.update_timestamp_tz):'%Y-%m-%d %H:%M:%S'}." - " Please check if it is a stale session." - ) - # Fetch states from storage storage_app_state = await sql_session.get( schema.StorageAppState, (session.app_name) @@ -450,6 +453,29 @@ class DatabaseSessionService(BaseSessionService): schema.StorageUserState, (session.app_name, session.user_id) ) + is_sqlite = self.db_engine.dialect.name == "sqlite" + if ( + storage_session.get_update_timestamp(is_sqlite) + > session.last_update_time + ): + # Reload the session from storage if it has been updated since it was + # loaded. + app_state = storage_app_state.state if storage_app_state else {} + user_state = storage_user_state.state if storage_user_state else {} + session_state = storage_session.state + session.state = _merge_state(app_state, user_state, session_state) + + stmt = ( + select(schema.StorageEvent) + .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.asc()) + ) + result = await sql_session.stream_scalars(stmt) + storage_events = [e async for e in result] + session.events = [e.to_event() for e in storage_events] + # Extract state delta if event.actions and event.actions.state_delta: state_deltas = _session_util.extract_state_delta( @@ -466,7 +492,7 @@ class DatabaseSessionService(BaseSessionService): if session_state_delta: storage_session.state = storage_session.state | session_state_delta - if storage_session._dialect_name == "sqlite": + if is_sqlite: update_time = datetime.fromtimestamp( event.timestamp, timezone.utc ).replace(tzinfo=None) @@ -476,10 +502,9 @@ class DatabaseSessionService(BaseSessionService): sql_session.add(schema.StorageEvent.from_event(session, event)) await sql_session.commit() - await sql_session.refresh(storage_session) # Update timestamp with commit time - session.last_update_time = storage_session.update_timestamp_tz + session.last_update_time = storage_session.get_update_timestamp(is_sqlite) # Also update the in-memory session await super().append_event(session=session, event=event) diff --git a/src/google/adk/sessions/schemas/v0.py b/src/google/adk/sessions/schemas/v0.py index bbdd66a2..8603ed81 100644 --- a/src/google/adk/sessions/schemas/v0.py +++ b/src/google/adk/sessions/schemas/v0.py @@ -41,7 +41,6 @@ from sqlalchemy import func from sqlalchemy import Text from sqlalchemy.dialects import mysql from sqlalchemy.ext.mutable import MutableDict -from sqlalchemy.inspection import inspect from sqlalchemy.orm import DeclarativeBase from sqlalchemy.orm import Mapped from sqlalchemy.orm import mapped_column @@ -131,15 +130,9 @@ class StorageSession(Base): def __repr__(self): return f"" - @property - def _dialect_name(self) -> Optional[str]: - session = inspect(self).session - return session.bind.dialect.name if session else None - - @property - def update_timestamp_tz(self) -> datetime: + def get_update_timestamp(self, is_sqlite: bool) -> float: """Returns the time zone aware update timestamp.""" - if self._dialect_name == "sqlite": + if is_sqlite: # SQLite does not support timezone. SQLAlchemy returns a naive datetime # object without timezone information. We need to convert it to UTC # manually. @@ -150,6 +143,7 @@ class StorageSession(Base): self, state: dict[str, Any] | None = None, events: list[Event] | None = None, + is_sqlite: bool = False, ) -> Session: """Converts the storage session to a session object.""" if state is None: @@ -163,7 +157,7 @@ class StorageSession(Base): id=self.id, state=state, events=events, - last_update_time=self.update_timestamp_tz, + last_update_time=self.get_update_timestamp(is_sqlite=is_sqlite), ) diff --git a/src/google/adk/sessions/schemas/v1.py b/src/google/adk/sessions/schemas/v1.py index 34b84dc1..a7c429b2 100644 --- a/src/google/adk/sessions/schemas/v1.py +++ b/src/google/adk/sessions/schemas/v1.py @@ -32,7 +32,6 @@ import uuid from sqlalchemy import ForeignKeyConstraint from sqlalchemy import func from sqlalchemy.ext.mutable import MutableDict -from sqlalchemy.inspection import inspect from sqlalchemy.orm import DeclarativeBase from sqlalchemy.orm import Mapped from sqlalchemy.orm import mapped_column @@ -106,15 +105,9 @@ class StorageSession(Base): def __repr__(self): return f"" - @property - def _dialect_name(self) -> Optional[str]: - session = inspect(self).session - return session.bind.dialect.name if session else None - - @property - def update_timestamp_tz(self) -> datetime: + def get_update_timestamp(self, is_sqlite: bool) -> float: """Returns the time zone aware update timestamp.""" - if self._dialect_name == "sqlite": + if is_sqlite: # SQLite does not support timezone. SQLAlchemy returns a naive datetime # object without timezone information. We need to convert it to UTC # manually. @@ -125,6 +118,7 @@ class StorageSession(Base): self, state: dict[str, Any] | None = None, events: list[Event] | None = None, + is_sqlite: bool = False, ) -> Session: """Converts the storage session to a session object.""" if state is None: @@ -138,7 +132,7 @@ class StorageSession(Base): id=self.id, state=state, events=events, - last_update_time=self.update_timestamp_tz, + last_update_time=self.get_update_timestamp(is_sqlite=is_sqlite), ) diff --git a/tests/unittests/sessions/test_session_service.py b/tests/unittests/sessions/test_session_service.py index c4c86575..30eb1567 100644 --- a/tests/unittests/sessions/test_session_service.py +++ b/tests/unittests/sessions/test_session_service.py @@ -504,26 +504,65 @@ async def test_session_last_update_time_updates_on_event(session_service): @pytest.mark.asyncio -async def test_get_session_with_config(session_service): +async def test_append_event_to_stale_session(): + session_service = get_session_service( + service_type=SessionServiceType.DATABASE + ) + app_name = 'my_app' user_id = 'user' + current_time = datetime.now().astimezone(timezone.utc).timestamp() - session = await session_service.create_session( + original_session = await session_service.create_session( app_name=app_name, user_id=user_id ) - original_update_time = session.last_update_time - - event = Event(invocation_id='invocation', author='user') - await session_service.append_event(session=session, event=event) - - assert session.last_update_time >= event.timestamp - - refreshed_session = await session_service.get_session( - app_name=app_name, user_id=user_id, session_id=session.id + event1 = Event( + invocation_id='inv1', + author='user', + timestamp=current_time + 1, + actions=EventActions(state_delta={'sk1': 'v1'}), ) - assert refreshed_session is not None - assert refreshed_session.last_update_time >= event.timestamp - assert refreshed_session.last_update_time > original_update_time + await session_service.append_event(original_session, event1) + + updated_session = await session_service.get_session( + app_name=app_name, user_id=user_id, session_id=original_session.id + ) + event2 = Event( + invocation_id='inv2', + author='user', + timestamp=current_time + 2, + actions=EventActions(state_delta={'sk2': 'v2'}), + ) + await session_service.append_event(updated_session, event2) + + # original_session is now stale + assert original_session.last_update_time < updated_session.last_update_time + assert len(original_session.events) == 1 + assert 'sk2' not in original_session.state + + # Appending another event to stale original_session + event3 = Event( + invocation_id='inv3', + author='user', + timestamp=current_time + 3, + actions=EventActions(state_delta={'sk3': 'v3'}), + ) + await session_service.append_event(original_session, event3) + + # If we fetch session from DB, it should contain all 3 events and all state + # changes. + session_final = await session_service.get_session( + app_name=app_name, user_id=user_id, session_id=original_session.id + ) + assert len(session_final.events) == 3 + assert session_final.state.get('sk1') == 'v1' + assert session_final.state.get('sk2') == 'v2' + assert session_final.state.get('sk3') == 'v3' + assert [e.invocation_id for e in session_final.events] == [ + 'inv1', + 'inv2', + 'inv3', + ] @pytest.mark.asyncio