fix(sessions): stop concurrent writes losing app: and user: state

src/google/adk/integrations/redis/_redis_session_service.py

@@ -139,6 +139,45 @@ def _extract_session_only_state(state: dict[str, Any]) -> dict[str, Any]:
         and not k.startswith(State.TEMP_PREFIX)
     }
 
+  async def _apply_scoped_state_delta(
+      self, client: Any, key: str, delta: dict[str, Any]
+  ) -> dict[str, Any]:
+    """Merges an app or user state delta into the state stored at `key`.
+
+    Every session of an app, or of a user, writes the same key. The read and
+    the write therefore run in one watched transaction, retried when another
+    writer changes the key in between: merging into a copy read earlier would
+    overwrite the keys that writer added.
+
+    Args:
+      client: The Redis client.
+      key: The app- or user-state key.
+      delta: The unprefixed keys to merge in. When empty, nothing is written.
+
+    Returns:
+      The state stored at `key` once the delta is applied.
+    """
+    if not delta:
+      raw = await client.get(key)
+      return cast(dict[str, Any], json.loads(raw)) if raw else {}
+
+    merged: dict[str, Any] = {}
+
+    async def merge(pipe: Any) -> None:
+      nonlocal merged
+      raw = await pipe.get(key)
+      merged = json.loads(raw) if raw else {}
+      merged.update(delta)
+      pipe.multi()
+      pipe.set(
+          key,
+          json.dumps(merged),
+          ex=self.config.ttl_seconds if self.config.ttl_seconds > 0 else None,
+      )
+
+    await client.transaction(merge, key)
+    return merged
+
   async def create_session(
       self,
       *,
@@ -199,28 +238,10 @@ async def create_session(
           f"Session {sid} already exists for user {user_id} in app {app_name}."
       )
 
-    app_raw, user_raw = await asyncio.gather(
-        client.get(app_key),
-        client.get(user_key),
+    current_app, current_user = await asyncio.gather(
+        self._apply_scoped_state_delta(client, app_key, app_state_delta),
+        self._apply_scoped_state_delta(client, user_key, user_state_delta),
     )
-    current_app = json.loads(app_raw) if app_raw else {}
-    current_user = json.loads(user_raw) if user_raw else {}
-
-    if app_state_delta:
-      current_app.update(app_state_delta)
-      await client.set(
-          app_key,
-          json.dumps(current_app),
-          ex=self.config.ttl_seconds if self.config.ttl_seconds > 0 else None,
-      )
-
-    if user_state_delta:
-      current_user.update(user_state_delta)
-      await client.set(
-          user_key,
-          json.dumps(current_user),
-          ex=self.config.ttl_seconds if self.config.ttl_seconds > 0 else None,
-      )
 
     # Return session with dynamically merged state (including app/user/temp state)
     merged_state = self._merge_state(current_app, current_user, session_state)

src/google/adk/sessions/database_session_service.py

@@ -16,6 +16,7 @@
 import asyncio
 from collections.abc import Callable
 from contextlib import asynccontextmanager
+from contextlib import AsyncExitStack
 import copy
 from datetime import datetime
 from datetime import timezone
@@ -113,9 +114,11 @@
     _MYSQL_DIALECT: "aiomysql",
     _MARIADB_DIALECT: "asyncmy",
 }
-# Tuple key order for in-process per-session lock maps:
-# (app_name, user_id, session_id).
-_SessionLockKey: TypeAlias = tuple[str, str, str]
+# Keys of the in-process lock map. The leading tag keeps the locks of one
+# session, of one user's state and of one app's state apart:
+# ("session", app_name, user_id, session_id), ("user", app_name, user_id) and
+# ("app", app_name).
+_LockKey: TypeAlias = tuple[str, ...]
 _StorageState: TypeAlias = (
     StorageAppStateV0
     | StorageAppStateV1
@@ -207,14 +210,23 @@ async def _get_or_create_state(
     state_model: type[_StorageStateT],
     primary_key: Any,
     defaults: dict[str, Any],
+    use_row_level_locking: bool = False,
 ) -> _StorageStateT:
   """Returns an existing state row or creates one, handling concurrent inserts.
 
   Uses a SAVEPOINT so that an IntegrityError from a racing INSERT does not
-  invalidate the outer transaction.
+  invalidate the outer transaction. With `use_row_level_locking`, the row is
+  returned locked `FOR UPDATE` (a row this transaction inserts is locked
+  already), so the caller can merge a delta into it and write it back without
+  overwriting a write that another transaction commits meanwhile.
   """
   row: _StorageStateT | None = await sql_session.get(state_model, primary_key)
   if row is not None:
+    if use_row_level_locking:
+      # Lock the row only once it is known to exist: on MySQL a locking read
+      # of a missing key takes a gap lock, and two transactions that each hold
+      # one deadlock when they both insert the key.
+      await sql_session.refresh(row, with_for_update=True)
     return row
   try:
     async with sql_session.begin_nested():
@@ -225,7 +237,11 @@ async def _get_or_create_state(
   except IntegrityError:
     # Another concurrent caller inserted the row first.
     # The savepoint was rolled back, so re-fetch the winner's row.
-    row = await sql_session.get(state_model, primary_key)
+    row = await sql_session.get(
+        state_model,
+        primary_key,
+        with_for_update=True if use_row_level_locking else None,
+    )
     if row is None:
       raise
     return row
@@ -448,10 +464,12 @@ def __init__(
     # The current database schema version in use, "None" if not yet checked
     self._db_schema_version: str | None = None
 
-    # Per-session locks used to serialize append_event calls in this process.
-    self._session_locks: dict[_SessionLockKey, asyncio.Lock] = {}
-    self._session_lock_ref_count: dict[_SessionLockKey, int] = {}
-    self._session_locks_guard = asyncio.Lock()
+    # Locks that serialize, within this process, appends to one session and
+    # writes to the app- and user-state rows every session of an app or user
+    # shares.
+    self._locks: dict[_LockKey, asyncio.Lock] = {}
+    self._lock_ref_count: dict[_LockKey, int] = {}
+    self._locks_guard = asyncio.Lock()
 
   def _get_schema_classes(self) -> _SchemaClasses:
     return _SchemaClasses(self._db_schema_version)
@@ -499,36 +517,77 @@ def _uses_naive_datetime(self) -> bool:
     return self.db_engine.dialect.name in _NAIVE_DATETIME_DIALECTS
 
   @asynccontextmanager
-  async def _with_session_lock(
-      self, *, app_name: str, user_id: str, session_id: str
-  ) -> AsyncIterator[None]:
-    """Serializes event appends for the same session within this process."""
-    # Use one lock per logical ADK session to prevent concurrent append_event
-    # writes from racing in the same process.
-    lock_key = (app_name, user_id, session_id)
-    async with self._session_locks_guard:
-      lock = self._session_locks.get(lock_key)
+  async def _with_lock(self, lock_key: _LockKey) -> AsyncIterator[None]:
+    """Holds the in-process lock for `lock_key`, creating it on first use."""
+    async with self._locks_guard:
+      lock = self._locks.get(lock_key)
       if lock is None:
         lock = asyncio.Lock()
-        self._session_locks[lock_key] = lock
+        self._locks[lock_key] = lock
       # Reference counting keeps lock objects alive while they are in use by
       # concurrent tasks and allows cleanup once all waiters complete.
-      self._session_lock_ref_count[lock_key] = (
-          self._session_lock_ref_count.get(lock_key, 0) + 1
-      )
+      self._lock_ref_count[lock_key] = self._lock_ref_count.get(lock_key, 0) + 1
 
     try:
       async with lock:
         yield
     finally:
-      async with self._session_locks_guard:
-        remaining = self._session_lock_ref_count.get(lock_key, 0) - 1
+      async with self._locks_guard:
+        remaining = self._lock_ref_count.get(lock_key, 0) - 1
         # Remove lock bookkeeping after the last waiter exits.
         if remaining <= 0 and not lock.locked():
-          self._session_lock_ref_count.pop(lock_key, None)
-          self._session_locks.pop(lock_key, None)
+          self._lock_ref_count.pop(lock_key, None)
+          self._locks.pop(lock_key, None)
         else:
-          self._session_lock_ref_count[lock_key] = remaining
+          self._lock_ref_count[lock_key] = remaining
+
+  @asynccontextmanager
+  async def _with_session_lock(
+      self, *, app_name: str, user_id: str, session_id: str
+  ) -> AsyncIterator[None]:
+    """Serializes event appends for the same session within this process."""
+    # Use one lock per logical ADK session to prevent concurrent append_event
+    # writes from racing in the same process.
+    async with self._with_lock(("session", app_name, user_id, session_id)):
+      yield
+
+  @asynccontextmanager
+  async def _with_shared_state_locks(
+      self,
+      *,
+      app_name: str,
+      user_id: str,
+      app_state: bool,
+      user_state: bool,
+  ) -> AsyncIterator[None]:
+    """Serializes writes to the shared app- and user-state rows in this process.
+
+    Every session of an app writes the same app-state row, and every session
+    of a user the same user-state row. A write reads the row, merges its delta
+    into the JSON in Python and writes the whole document back, so two
+    sessions that write at once would each overwrite the other's keys. The
+    session lock cannot order them. On PostgreSQL and MySQL the row is also
+    read `FOR UPDATE`, which orders writers across processes. SQLAlchemy
+    renders no `FOR UPDATE` for SQLite or SQL Server, so there this lock is
+    all that orders them.
+
+    Callers take these after any session lock, and the app lock always comes
+    before the user lock, so two writers never wait on each other in a cycle.
+
+    Args:
+      app_name: The app whose state rows are written.
+      user_id: The user whose state row is written.
+      app_state: Whether the write changes the app-state row.
+      user_state: Whether the write changes the user-state row.
+    """
+    async with AsyncExitStack() as stack:
+      if app_state:
+        await stack.enter_async_context(self._with_lock(("app", app_name)))
+      if user_state:
+        await stack.enter_async_context(
+            self._with_lock(("user", app_name, user_id))
+        )
+      yield
 
   async def prepare_tables(self) -> None:
     """Ensure database tables are ready for use.
@@ -659,7 +718,25 @@ async def create_session(
     if session_id is None:
       session_id = platform_uuid.new_uuid()
     schema = self._get_schema_classes()
-    async with self._rollback_on_exception_session() as sql_session:
+    use_row_level_locking = self._supports_row_level_locking()
+
+    # Extract state deltas
+    state_deltas = _session_util.extract_json_safe_state_delta(state or {})
+    app_state_delta = state_deltas["app"]
+    user_state_delta = state_deltas["user"]
+    session_state = state_deltas["session"]
+    has_app_delta = bool(app_state_delta)
+    has_user_delta = bool(user_state_delta)
+
+    async with (
+        self._with_shared_state_locks(
+            app_name=app_name,
+            user_id=user_id,
+            app_state=has_app_delta,
+            user_state=has_user_delta,
+        ),
+        self._rollback_on_exception_session() as sql_session,
+    ):
       if has_user_provided_id:
         existing_session = await sql_session.get(
             schema.StorageSession, (app_name, user_id, session_id)
@@ -668,30 +745,28 @@ async def create_session(
           raise AlreadyExistsError(
               f"Session with id {session_id} already exists."
           )
-      # Get or create state rows, handling concurrent insert races.
+      # Get or create state rows, handling concurrent insert races. A row the
+      # initial state merges into is locked as in append_event, so the merge
+      # cannot drop keys another session commits meanwhile.
       storage_app_state = await _get_or_create_state(
           sql_session=sql_session,
           state_model=schema.StorageAppState,
           primary_key=app_name,
           defaults={"app_name": app_name, "state": {}},
+          use_row_level_locking=use_row_level_locking and has_app_delta,
       )
       storage_user_state = await _get_or_create_state(
           sql_session=sql_session,
           state_model=schema.StorageUserState,
           primary_key=(app_name, user_id),
           defaults={"app_name": app_name, "user_id": user_id, "state": {}},
+          use_row_level_locking=use_row_level_locking and has_user_delta,
       )
 
-      # Extract state deltas
-      state_deltas = _session_util.extract_json_safe_state_delta(state or {})
-      app_state_delta = state_deltas["app"]
-      user_state_delta = state_deltas["user"]
-      session_state = state_deltas["session"]
-
       # Apply state delta
-      if app_state_delta:
+      if has_app_delta:
         storage_app_state.state.update(app_state_delta)
-      if user_state_delta:
+      if has_user_delta:
         storage_user_state.state.update(user_state_delta)
 
       # Store the session
@@ -930,10 +1005,18 @@ async def append_event(self, session: Session, event: Event) -> Event:
     has_app_delta = bool(state_deltas["app"])
     has_user_delta = bool(state_deltas["user"])
 
-    async with self._with_session_lock(
-        app_name=session.app_name,
-        user_id=session.user_id,
-        session_id=session.id,
+    async with (
+        self._with_session_lock(
+            app_name=session.app_name,
+            user_id=session.user_id,
+            session_id=session.id,
+        ),
+        self._with_shared_state_locks(
+            app_name=session.app_name,
+            user_id=session.user_id,
+            app_state=has_app_delta,
+            user_state=has_user_delta,
+        ),
     ):
       async with self._rollback_on_exception_session() as sql_session:
         storage_session_stmt = (

tests/unittests/integrations/redis/test_redis_session_service.py

@@ -810,6 +810,52 @@ async def test_cumulative_user_state_creation(session_service):
   assert user_state == {"theme": "dark", "lang": "en"}
 
 
+@pytest.mark.asyncio
+async def test_create_session_keeps_user_state_written_while_it_merges(
+    session_service, fake_redis, monkeypatch
+):
+  """Initial user state merges with a write landing between read and write.
+
+  Setup: another session of the user appends a user: key right after
+    create_session reads the user's state.
+  Act: create a session whose initial state carries a different user: key.
+  Assert: the user's state keeps both keys.
+  """
+  other = await session_service.create_session(
+      app_name="app1", user_id="u1", session_id="other"
+  )
+  user_key = session_service._user_state_key("app1", "u1")
+  original_get = fake_redis.get
+  wrote_concurrently = False
+
+  async def get_then_write_concurrently(read_key):
+    nonlocal wrote_concurrently
+    raw = await original_get(read_key)
+    if read_key == user_key and not wrote_concurrently:
+      wrote_concurrently = True
+      await session_service.append_event(
+          other,
+          Event(
+              author="user",
+              invocation_id="other",
+              actions=EventActions(state_delta={"user:theme": "dark"}),
+          ),
+      )
+    return raw
+
+  monkeypatch.setattr(fake_redis, "get", get_then_write_concurrently)
+  await session_service.create_session(
+      app_name="app1",
+      user_id="u1",
+      session_id="new",
+      state={"user:lang": "en"},
+  )
+
+  assert await session_service.get_user_state(
+      app_name="app1", user_id="u1"
+  ) == {"theme": "dark", "lang": "en"}
+
+
 @pytest.mark.asyncio
 @pytest.mark.parametrize("partial", [False, True])
 async def test_append_event_to_unknown_session_writes_nothing(