feat(sessions): add merge_state for race-free server-side atomic state merges

Add a state-only write path, merge_state(app_name, user_id, session_id, delta),
that performs a server-side atomic JSON merge of the routed delta into the
app/user/session state row with no events row and without advancing the session
optimistic-concurrency (OCC) marker.

Today the only way to persist state is append_event, which couples a commutative
state merge to event-log append + whole-session OCC. Two workers updating
independent state keys on the same session therefore race on the session-level
marker, and one spuriously fails with "modified in storage" and must retry, even
though the keys never collide.

merge_state routes the delta via the existing app:/user:/unprefixed convention
(temp: rejected) and merges each scope server-side:
  - PostgreSQL: state || CAST(:delta AS JSONB)
  - MySQL/MariaDB: JSON_MERGE_PATCH(state, CAST(:delta AS JSON))
  - SQLite: json_patch(state, :delta)
The session-row merge is issued via sqlalchemy.text touching only the state
column, so update_time's onupdate never fires and the OCC marker is preserved.
App/user rows are auto-created via the existing _get_or_create_state; a missing
session row raises for session-scoped keys. BaseSessionService and
VertexAiSessionService raise NotImplementedError, mirroring get_user_state;
InMemorySessionService and SqliteSessionService implement it natively.

Adds parameterized tests across all backends plus DatabaseSessionService tests
proving the OCC marker is not bumped, no events row is written, and concurrent
merges to independent keys do not lose updates.

Author: lwangverizon

src/google/adk/sessions/base_session_service.py

@@ -151,6 +151,67 @@ async def get_user_state(
         'call get_session on each result to access the merged state.'
     )
 
+  async def merge_state(
+      self,
+      *,
+      app_name: str,
+      user_id: str,
+      session_id: str,
+      delta: dict[str, Any],
+  ) -> None:
+    """Atomically merges a state delta without appending an event.
+
+    This is a state-only write path that bypasses the event log and the
+    whole-session optimistic-concurrency (OCC) check that ``append_event``
+    performs.  It is intended for *commutative* updates to independent state
+    keys (counters, flags, per-user balances, feature toggles), where coupling
+    the write to event-log append + session-level OCC would force unrelated
+    writers to serialize and retry on spurious "stale session" errors.
+
+    The ``delta`` uses the same prefix convention as ``create_session`` and
+    event ``state_delta``; keys are routed to the scope-appropriate storage
+    row:
+
+    * ``app:`` -> app-scoped state (keyed by ``app_name``).
+    * ``user:`` -> user-scoped state (keyed by ``(app_name, user_id)``).
+    * no prefix -> session-scoped state (keyed by
+      ``(app_name, user_id, session_id)``).
+
+    ``temp:`` keys are rejected with ``ValueError`` because temp state is never
+    persisted.
+
+    Guarantees (for backends that implement this method):
+
+    * No ``events`` row is written.
+    * The session's OCC revision marker is **not** advanced, so a concurrently
+      held in-memory ``Session`` does not become stale and its next
+      ``append_event`` still succeeds.
+    * App- and user-scoped merges do not require a pre-existing session and are
+      fully decoupled from any session's revision.  ``session_id`` is used only
+      to route session-scoped keys.
+
+    Merge semantics for non-scalar values are backend-dependent; see the
+    concrete implementations (notably ``DatabaseSessionService.merge_state``)
+    for the dialect-specific caveats around nested objects and ``None`` values.
+    For flat scalar values all backends behave identically.
+
+    Args:
+      app_name: The name of the app.
+      user_id: The ID of the user.
+      session_id: The ID of the session. Required for routing session-scoped
+        keys; need not exist when the delta has no session-scoped keys.
+      delta: The state delta to merge. An empty or ``None`` delta is a no-op.
+
+    Raises:
+      ValueError: If ``delta`` contains ``temp:``-prefixed keys, or if a
+        session-scoped key is supplied for a session that does not exist.
+      NotImplementedError: When the concrete ``BaseSessionService``
+        implementation does not support a server-side state merge.
+    """
+    raise NotImplementedError(
+        f'{type(self).__name__} does not support merge_state.'
+    )
+
   async def append_event(self, session: Session, event: Event) -> Event:
     """Appends an event to a session object."""
     if event.partial:

src/google/adk/sessions/database_session_service.py

@@ -18,6 +18,7 @@
 import copy
 from datetime import datetime
 from datetime import timezone
+import json
 import logging
 from typing import Any
 from typing import AsyncIterator
@@ -34,6 +35,7 @@
   from sqlalchemy import event
   from sqlalchemy import MetaData
   from sqlalchemy import select
+  from sqlalchemy import text
   from sqlalchemy.engine import Connection
   from sqlalchemy.engine import make_url
   from sqlalchemy.exc import ArgumentError
@@ -187,6 +189,30 @@ def _merge_state(
   return merged_state
 
 
+def _json_merge_set_clause(dialect_name: str) -> Optional[str]:
+  """Returns a 'state = <merge>' SET clause for an atomic server-side JSON merge.
+
+  The returned SQL fragment merges a ``:delta`` bind parameter into the
+  existing ``state`` column in place. Only the ``state`` column is referenced,
+  so when issued via ``sqlalchemy.text`` no column ``onupdate`` callable (such
+  as the ``update_time`` revision marker) is triggered.
+
+  Returns None for dialects that do not provide a known server-side JSON merge
+  function; the caller raises ``NotImplementedError`` in that case.
+  """
+  if dialect_name == _POSTGRESQL_DIALECT:
+    # JSONB concatenation: shallow top-level merge; an explicit JSON null is
+    # stored (key kept), matching append_event's Python dict union.
+    return "state = state || CAST(:delta AS JSONB)"
+  if dialect_name in (_MYSQL_DIALECT, _MARIADB_DIALECT):
+    # RFC 7396: recursive merge; a JSON null deletes the key.
+    return "state = JSON_MERGE_PATCH(state, CAST(:delta AS JSON))"
+  if dialect_name == _SQLITE_DIALECT:
+    # RFC 7396 semantics, like the native SqliteSessionService.
+    return "state = json_patch(state, :delta)"
+  return None
+
+
 class _SchemaClasses:
   """A helper class to hold schema classes based on version."""
 
@@ -737,6 +763,152 @@ async def get_user_state(
         return {}
       return dict(storage_user_state.state or {})
 
+  @override
+  async def merge_state(
+      self,
+      *,
+      app_name: str,
+      user_id: str,
+      session_id: str,
+      delta: dict[str, Any],
+  ) -> None:
+    """Atomically merges a state delta server-side without appending an event.
+
+    Each scoped sub-delta is merged into its storage row with a single atomic
+    SQL statement (PostgreSQL ``state || delta``, MySQL/MariaDB
+    ``JSON_MERGE_PATCH``, SQLite ``json_patch``) issued via ``text`` so that
+    only the ``state`` column is written. As a result:
+
+    * No ``events`` row is created.
+    * ``sessions.update_time`` (the optimistic-concurrency revision marker) is
+      never advanced, so a concurrently held ``Session`` does not go stale.
+    * Independent keys merged concurrently do not lose updates, because the
+      merge happens inside the database rather than via Python
+      read-modify-write.
+
+    Merge semantics for non-scalar values differ by dialect. On PostgreSQL the
+    top-level keys are merged shallowly and an explicit ``None`` is stored as
+    JSON ``null`` (the key is kept), matching ``append_event``'s Python ``dict``
+    union. On MySQL, MariaDB and SQLite the merge follows RFC 7396: nested
+    objects are merged recursively and a ``None`` value deletes the key. For
+    flat scalar values (counters, flags, balances -- the intended use case) all
+    dialects behave identically. Use ``append_event`` if you need uniform
+    Python-``dict``-union semantics for nested objects or ``None`` overwrites.
+
+    Args:
+      app_name: The name of the app.
+      user_id: The ID of the user.
+      session_id: The ID of the session. Used to route session-scoped (no
+        prefix) keys; need not exist when ``delta`` has no session-scoped keys.
+      delta: The state delta to merge, using the ``app:``/``user:``/no-prefix
+        convention. An empty or ``None`` delta is a no-op.
+
+    Raises:
+      ValueError: If ``delta`` contains ``temp:`` keys, or if a session-scoped
+        key targets a session that does not exist.
+      NotImplementedError: If the active SQL dialect provides no known
+        server-side JSON merge function.
+    """
+    await self.prepare_tables()
+    if not delta:
+      return
+    if any(key.startswith(State.TEMP_PREFIX) for key in delta):
+      raise ValueError(
+          "merge_state does not support temp: keys; temp state is never"
+          " persisted."
+      )
+
+    merge_clause = _json_merge_set_clause(self.db_engine.dialect.name)
+    if merge_clause is None:
+      raise ValueError(
+          "merge_state is not supported for dialect"
+          f" {self.db_engine.dialect.name!r}: no server-side JSON merge"
+          " function is available."
+      )
+
+    state_deltas = _session_util.extract_state_delta(delta)
+    app_delta = state_deltas["app"]
+    user_delta = state_deltas["user"]
+    session_delta = state_deltas["session"]
+
+    schema = self._get_schema_classes()
+    app_table = schema.StorageAppState.__tablename__
+    user_table = schema.StorageUserState.__tablename__
+    session_table = schema.StorageSession.__tablename__
+    use_row_level_locking = self._supports_row_level_locking()
+
+    async with self._with_session_lock(
+        app_name=app_name,
+        user_id=user_id,
+        session_id=session_id,
+    ):
+      async with self._rollback_on_exception_session() as sql_session:
+        if session_delta:
+          # A session-scoped key requires the session row to exist; never
+          # auto-create it (that is create_session's responsibility).
+          session_exists_stmt = (
+              select(schema.StorageSession.id)
+              .filter(schema.StorageSession.app_name == app_name)
+              .filter(schema.StorageSession.user_id == user_id)
+              .filter(schema.StorageSession.id == session_id)
+          )
+          if use_row_level_locking:
+            session_exists_stmt = session_exists_stmt.with_for_update()
+          session_exists = await sql_session.execute(session_exists_stmt)
+          if session_exists.scalar_one_or_none() is None:
+            raise ValueError(f"Session {session_id} not found.")
+
+        if app_delta:
+          # Ensure the row exists (handles concurrent inserts), then merge
+          # server-side. App/user rows have their own update_time which is not
+          # an OCC marker, so merging here is independent of any session.
+          await _get_or_create_state(
+              sql_session=sql_session,
+              state_model=schema.StorageAppState,
+              primary_key=app_name,
+              defaults={"app_name": app_name, "state": {}},
+          )
+          await sql_session.execute(
+              text(
+                  f"UPDATE {app_table} SET {merge_clause} WHERE"
+                  " app_name = :app_name"
+              ),
+              {"app_name": app_name, "delta": json.dumps(app_delta)},
+          )
+        if user_delta:
+          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": {}},
+          )
+          await sql_session.execute(
+              text(
+                  f"UPDATE {user_table} SET {merge_clause} WHERE"
+                  " app_name = :app_name AND user_id = :user_id"
+              ),
+              {
+                  "app_name": app_name,
+                  "user_id": user_id,
+                  "delta": json.dumps(user_delta),
+              },
+          )
+        if session_delta:
+          await sql_session.execute(
+              text(
+                  f"UPDATE {session_table} SET {merge_clause} WHERE"
+                  " app_name = :app_name AND user_id = :user_id"
+                  " AND id = :session_id"
+              ),
+              {
+                  "app_name": app_name,
+                  "user_id": user_id,
+                  "session_id": session_id,
+                  "delta": json.dumps(session_delta),
+              },
+          )
+        await sql_session.commit()
+
   @override
   async def append_event(self, session: Session, event: Event) -> Event:
     await self.prepare_tables()

src/google/adk/sessions/in_memory_session_service.py

@@ -318,6 +318,46 @@ async def get_user_state(
   ) -> dict[str, Any]:
     return dict(self.user_state.get(app_name, {}).get(user_id, {}))
 
+  @override
+  async def merge_state(
+      self,
+      *,
+      app_name: str,
+      user_id: str,
+      session_id: str,
+      delta: dict[str, Any],
+  ) -> None:
+    if not delta:
+      return
+    if any(key.startswith(State.TEMP_PREFIX) for key in delta):
+      raise ValueError(
+          'merge_state does not support temp: keys; temp state is never'
+          ' persisted.'
+      )
+
+    state_deltas = _session_util.extract_state_delta(delta)
+    app_state_delta = state_deltas['app']
+    user_state_delta = state_deltas['user']
+    session_state_delta = state_deltas['session']
+
+    if session_state_delta:
+      storage_session = (
+          self.sessions.get(app_name, {}).get(user_id, {}).get(session_id)
+      )
+      if storage_session is None:
+        raise ValueError(f'Session {session_id} not found.')
+
+    if app_state_delta:
+      self.app_state.setdefault(app_name, {}).update(app_state_delta)
+    if user_state_delta:
+      self.user_state.setdefault(app_name, {}).setdefault(user_id, {}).update(
+          user_state_delta
+      )
+    if session_state_delta:
+      # Merge into the stored session's state without bumping
+      # last_update_time, so a concurrently held session does not go stale.
+      storage_session.state.update(session_state_delta)
+
   @override
   async def append_event(self, session: Session, event: Event) -> Event:
     if event.partial: