fix: Reject appends to stale sessions in DatabaseSessionService

This change introduces optimistic concurrency control for session updates. Instead of automatically reloading and merging when an append is attempted on a session that has been modified in storage, the service now raises a ValueError.

Close #4751

Co-authored-by: George Weale <gweale@google.com>
PiperOrigin-RevId: 885334444

Author: GWeale

src/google/adk/sessions/database_session_service.py

@@ -61,6 +61,11 @@
 
 logger = logging.getLogger("google_adk." + __name__)
 
+_STALE_SESSION_ERROR_MESSAGE = (
+    "The session has been modified in storage since it was loaded. "
+    "Please reload the session before appending more events."
+)
+
 _SQLITE_DIALECT = "sqlite"
 _MARIADB_DIALECT = "mariadb"
 _MYSQL_DIALECT = "mysql"
@@ -309,6 +314,39 @@ async def _prepare_tables(self):
 
       self._tables_created = True
 
+  async def _session_matches_storage_revision(
+      self,
+      *,
+      sql_session: DatabaseSessionFactory,
+      schema: _SchemaClasses,
+      session: Session,
+  ) -> bool:
+    """Returns whether a marker-less session still matches stored events."""
+    if not session.events:
+      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)
+          .limit(1)
+      )
+      result = await sql_session.execute(stmt)
+      return result.scalar_one_or_none() is None
+
+    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()
+        )
+        .limit(1)
+    )
+    result = await sql_session.execute(stmt)
+    latest_storage_event_id = result.scalar_one_or_none()
+    return latest_storage_event_id == session.events[-1].id
+
   @override
   async def create_session(
       self,
@@ -529,9 +567,9 @@ async def append_event(self, session: Session, event: Event) -> Event:
     # Trim temp state before persisting
     event = self._trim_temp_delta_state(event)
 
-    # 1. Check if timestamp is stale
-    # 2. Update session attributes based on event config
-    # 3. Store event to table
+    # 1. Validate the session has not gone stale.
+    # 2. Update session attributes based on event config.
+    # 3. Store the new event.
     schema = self._get_schema_classes()
     is_sqlite = self.db_engine.dialect.name == _SQLITE_DIALECT
     use_row_level_locking = self._supports_row_level_locking()
@@ -563,6 +601,8 @@ async def append_event(self, session: Session, event: Event) -> Event:
         storage_session = storage_session_result.scalars().one_or_none()
         if storage_session is None:
           raise ValueError(f"Session {session.id} not found.")
+        storage_update_time = storage_session.get_update_timestamp(is_sqlite)
+        storage_update_marker = storage_session.get_update_marker()
 
         storage_app_state = await _select_required_state(
             sql_session=sql_session,
@@ -591,27 +631,27 @@ async def append_event(self, session: Session, event: Event) -> Event:
             ),
         )
 
-        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
-          user_state = storage_user_state.state
-          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]
+        if session._storage_update_marker is not None:
+          # Sessions loaded by DatabaseSessionService carry an exact storage
+          # revision marker, so stale-writer detection can use that marker
+          # instead of relying on rounded timestamps.
+          if session._storage_update_marker != storage_update_marker:
+            raise ValueError(_STALE_SESSION_ERROR_MESSAGE)
+          # Keep the float timestamp synchronized with the exact storage value
+          # so tiny round-trip differences do not trigger false stale checks on
+          # the next append.
+          session.last_update_time = storage_update_time
+        elif storage_update_time > session.last_update_time:
+          # Backward-compatible fallback for marker-less session objects, such
+          # as older in-memory sessions or manually constructed Session values.
+          # Only reject when storage has actually advanced beyond the in-memory
+          # revision represented by session.events.
+          if not await self._session_matches_storage_revision(
+              sql_session=sql_session, schema=schema, session=session
+          ):
+            raise ValueError(_STALE_SESSION_ERROR_MESSAGE)
+          session.last_update_time = storage_update_time
+        session._storage_update_marker = storage_update_marker
 
         # Merge pre-extracted state deltas into storage.
         if has_app_delta:
@@ -642,6 +682,7 @@ async def append_event(self, session: Session, event: Event) -> Event:
         session.last_update_time = storage_session.get_update_timestamp(
             is_sqlite
         )
+        session._storage_update_marker = storage_session.get_update_marker()
 
     # Also update the in-memory session
     await super().append_event(session=session, event=event)

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

@@ -154,6 +154,13 @@ def get_update_timestamp(self, is_sqlite: bool) -> float:
       return self.update_time.replace(tzinfo=timezone.utc).timestamp()
     return self.update_time.timestamp()
 
+  def get_update_marker(self) -> str:
+    """Returns a stable revision marker for optimistic concurrency checks."""
+    update_time = self.update_time
+    if update_time.tzinfo is not None:
+      update_time = update_time.astimezone(timezone.utc)
+    return update_time.isoformat(timespec="microseconds")
+
   def to_session(
       self,
       state: dict[str, Any] | None = None,
@@ -166,14 +173,16 @@ def to_session(
     if events is None:
       events = []
 
-    return Session(
+    session = Session(
         app_name=self.app_name,
         user_id=self.user_id,
         id=self.id,
         state=state,
         events=events,
         last_update_time=self.get_update_timestamp(is_sqlite=is_sqlite),
     )
+    session._storage_update_marker = self.get_update_marker()
+    return session
 
 
 class StorageEvent(Base):

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

@@ -128,6 +128,13 @@ def get_update_timestamp(self, is_sqlite: bool) -> float:
       return self.update_time.replace(tzinfo=timezone.utc).timestamp()
     return self.update_time.timestamp()
 
+  def get_update_marker(self) -> str:
+    """Returns a stable revision marker for optimistic concurrency checks."""
+    update_time = self.update_time
+    if update_time.tzinfo is not None:
+      update_time = update_time.astimezone(timezone.utc)
+    return update_time.isoformat(timespec="microseconds")
+
   def to_session(
       self,
       state: dict[str, Any] | None = None,
@@ -140,14 +147,16 @@ def to_session(
     if events is None:
       events = []
 
-    return Session(
+    session = Session(
         app_name=self.app_name,
         user_id=self.user_id,
         id=self.id,
         state=state,
         events=events,
         last_update_time=self.get_update_timestamp(is_sqlite=is_sqlite),
     )
+    session._storage_update_marker = self.get_update_marker()
+    return session
 
 
 class StorageEvent(Base):

src/google/adk/sessions/session.py

@@ -20,6 +20,7 @@
 from pydantic import BaseModel
 from pydantic import ConfigDict
 from pydantic import Field
+from pydantic import PrivateAttr
 
 from ..events.event import Event
 
@@ -48,3 +49,6 @@ class Session(BaseModel):
   call/response, etc."""
   last_update_time: float = 0.0
   """The last update time of the session."""
+
+  _storage_update_marker: str | None = PrivateAttr(default=None)
+  """Internal storage revision marker used for stale-session detection."""

tests/unittests/sessions/test_session_service.py

@@ -657,28 +657,27 @@ async def test_append_event_to_stale_session():
     assert len(original_session.events) == 1
     assert 'sk2' not in original_session.state
 
-    # Appending another event to stale original_session
+    # Appending another event to stale original_session should be rejected.
     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)
+    with pytest.raises(ValueError, match='modified in storage'):
+      await session_service.append_event(original_session, event3)
 
-    # If we fetch session from DB, it should contain all 3 events and all state
-    # changes.
+    # If we fetch session from DB, it should only contain the committed events.
     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 len(session_final.events) == 2
     assert session_final.state.get('sk1') == 'v1'
     assert session_final.state.get('sk2') == 'v2'
-    assert session_final.state.get('sk3') == 'v3'
+    assert session_final.state.get('sk3') is None
     assert [e.invocation_id for e in session_final.events] == [
         'inv1',
         'inv2',
-        'inv3',
     ]
 
 
@@ -738,7 +737,7 @@ async def test_append_event_raises_if_user_state_row_missing():
 
 
 @pytest.mark.asyncio
-async def test_append_event_concurrent_stale_sessions_preserve_all_state():
+async def test_append_event_concurrent_stale_sessions_reject_stale_writer():
   session_service = get_session_service(
       service_type=SessionServiceType.DATABASE
   )
@@ -771,19 +770,103 @@ async def test_append_event_concurrent_stale_sessions_preserve_all_state():
           actions=EventActions(state_delta={f'sk{i}-2': f'v{i}-2'}),
       )
 
-      await asyncio.gather(
+      results = await asyncio.gather(
           session_service.append_event(stale_session_1, event_1),
           session_service.append_event(stale_session_2, event_2),
+          return_exceptions=True,
       )
+      errors = [result for result in results if isinstance(result, Exception)]
+      successes = [
+          result for result in results if not isinstance(result, Exception)
+      ]
+      assert len(successes) == 1
+      assert len(errors) == 1
+      assert isinstance(errors[0], ValueError)
+      assert 'modified in storage' in str(errors[0])
 
     session_final = await session_service.get_session(
         app_name=app_name, user_id=user_id, session_id=session.id
     )
 
     for i in range(iteration_count):
-      assert session_final.state.get(f'sk{i}-1') == f'v{i}-1'
-      assert session_final.state.get(f'sk{i}-2') == f'v{i}-2'
-    assert len(session_final.events) == iteration_count * 2
+      event_values = {
+          session_final.state.get(f'sk{i}-1'),
+          session_final.state.get(f'sk{i}-2'),
+      }
+      assert event_values & {f'v{i}-1', f'v{i}-2'}
+      assert None in event_values
+    assert len(session_final.events) == iteration_count
+
+
+@pytest.mark.asyncio
+async def test_append_event_allows_timestamp_drift_for_current_session():
+  service = DatabaseSessionService('sqlite+aiosqlite:///:memory:')
+  try:
+    session = await service.create_session(
+        app_name='my_app', user_id='user', session_id='s1'
+    )
+    event1 = Event(
+        invocation_id='inv1',
+        author='user',
+        timestamp=session.last_update_time + 10,
+    )
+    await service.append_event(session, event1)
+
+    # Simulate a float round-trip mismatch without changing the persisted
+    # session revision.
+    session.last_update_time -= 0.0001
+
+    event2 = Event(
+        invocation_id='inv2',
+        author='user',
+        timestamp=event1.timestamp + 10,
+    )
+    await service.append_event(session, event2)
+
+    refreshed_session = await service.get_session(
+        app_name='my_app', user_id='user', session_id=session.id
+    )
+    assert [event.invocation_id for event in refreshed_session.events] == [
+        'inv1',
+        'inv2',
+    ]
+  finally:
+    await service.close()
+
+
+@pytest.mark.asyncio
+async def test_append_event_allows_markerless_current_session():
+  service = DatabaseSessionService('sqlite+aiosqlite:///:memory:')
+  try:
+    session = await service.create_session(
+        app_name='my_app', user_id='user', session_id='s1'
+    )
+    event1 = Event(
+        invocation_id='inv1',
+        author='user',
+        timestamp=session.last_update_time + 10,
+    )
+    await service.append_event(session, event1)
+
+    session._storage_update_marker = None
+    session.last_update_time -= 0.0001
+
+    event2 = Event(
+        invocation_id='inv2',
+        author='user',
+        timestamp=event1.timestamp + 10,
+    )
+    await service.append_event(session, event2)
+
+    refreshed_session = await service.get_session(
+        app_name='my_app', user_id='user', session_id=session.id
+    )
+    assert [event.invocation_id for event in refreshed_session.events] == [
+        'inv1',
+        'inv2',
+    ]
+  finally:
+    await service.close()
 
 
 @pytest.mark.asyncio