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: #1733

Co-authored-by: Shangjie Chen <deanchen@google.com>
PiperOrigin-RevId: 861353058

Author: DeanChensj

src/google/adk/sessions/database_session_service.py

@@ -31,7 +31,6 @@
 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 @@ async def create_session(
         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 @@ async def get_session(
 
       # 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 @@ async def list_sessions(
           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 @@ async def append_event(self, session: Session, event: Event) -> Event:
           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 @@ async def append_event(self, session: Session, event: Event) -> Event:
           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 @@ async def append_event(self, session: Session, event: Event) -> Event:
         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 @@ async def append_event(self, session: Session, event: Event) -> Event:
       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)

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

@@ -41,7 +41,6 @@
 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"<StorageSession(id={self.id}, update_time={self.update_time})>"
 
-  @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 @@ def to_session(
       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 @@ def to_session(
         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),
     )
 
 

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

@@ -32,7 +32,6 @@
 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"<StorageSession(id={self.id}, update_time={self.update_time})>"
 
-  @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 @@ def to_session(
       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 @@ def to_session(
         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),
     )
 
 

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
+  event1 = Event(
+      invocation_id='inv1',
+      author='user',
+      timestamp=current_time + 1,
+      actions=EventActions(state_delta={'sk1': 'v1'}),
+  )
+  await session_service.append_event(original_session, event1)
 
-  event = Event(invocation_id='invocation', author='user')
-  await session_service.append_event(session=session, event=event)
+  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)
 
-  assert session.last_update_time >= event.timestamp
+  # 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
 
-  refreshed_session = await session_service.get_session(
-      app_name=app_name, user_id=user_id, session_id=session.id
-  )
-  assert refreshed_session is not None
-  assert refreshed_session.last_update_time >= event.timestamp
-  assert refreshed_session.last_update_time > original_update_time
+  # 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