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 <deanchen@google.com>
PiperOrigin-RevId: 861353058

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)

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"<StorageSession(id={self.id}, update_time={self.update_time})>"
 
-  @property
+  def get_update_timestamp(self, is_sqlite: bool) -> float:
-  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:
     """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),
     )
 
 

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"<StorageSession(id={self.id}, update_time={self.update_time})>"
 
-  @property
+  def get_update_timestamp(self, is_sqlite: bool) -> float:
-  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:
     """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),
     )
 
 

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