fix: Handle concurrent creation of app/user state rows in DatabaseSessionService

Introduces a new helper function, `_get_or_create_state`, which uses a nested transaction (SAVEPOINT) to safely attempt creating a state row. If a concurrent transaction has already inserted the row, the inner insert will fail with an IntegrityError, which is caught, and the already-existing row is then fetched

Close #4954

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

Author: GWeale

src/google/adk/sessions/database_session_service.py

@@ -33,6 +33,7 @@
 from sqlalchemy.engine import Connection
 from sqlalchemy.engine import make_url
 from sqlalchemy.exc import ArgumentError
+from sqlalchemy.exc import IntegrityError
 from sqlalchemy.ext.asyncio import async_sessionmaker
 from sqlalchemy.ext.asyncio import AsyncEngine
 from sqlalchemy.ext.asyncio import AsyncSession as DatabaseSessionFactory
@@ -103,6 +104,35 @@ async def _select_required_state(
   return state_row
 
 
+async def _get_or_create_state(
+    *,
+    sql_session: DatabaseSessionFactory,
+    state_model: type[_StorageStateT],
+    primary_key: Any,
+    defaults: dict[str, Any],
+) -> _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.
+  """
+  row = await sql_session.get(state_model, primary_key)
+  if row is not None:
+    return row
+  try:
+    async with sql_session.begin_nested():
+      row = state_model(**defaults)
+      sql_session.add(row)
+    return row
+  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)
+    if row is None:
+      raise
+    return row
+
+
 def _set_sqlite_pragma(dbapi_connection, connection_record):
   cursor = dbapi_connection.cursor()
   cursor.execute("PRAGMA foreign_keys=ON")
@@ -401,24 +431,20 @@ async def create_session(
         raise AlreadyExistsError(
             f"Session with id {session_id} already exists."
         )
-      # Fetch app and user states from storage
-      storage_app_state = await sql_session.get(
-          schema.StorageAppState, (app_name)
+      # Get or create state rows, handling concurrent insert races.
+      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": {}},
       )
-      storage_user_state = await sql_session.get(
-          schema.StorageUserState, (app_name, user_id)
+      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": {}},
       )
 
-      # Create state tables if not exist
-      if not storage_app_state:
-        storage_app_state = schema.StorageAppState(app_name=app_name, state={})
-        sql_session.add(storage_app_state)
-      if not storage_user_state:
-        storage_user_state = schema.StorageUserState(
-            app_name=app_name, user_id=user_id, state={}
-        )
-        sql_session.add(storage_user_state)
-
       # Extract state deltas
       state_deltas = _session_util.extract_state_delta(state)
       app_state_delta = state_deltas["app"]

tests/unittests/sessions/test_session_service.py

@@ -1370,6 +1370,137 @@ async def test_prepare_tables_serializes_schema_detection_and_creation():
     await service.close()
 
 
+@pytest.mark.asyncio
+async def test_get_or_create_state_returns_existing_row():
+  """_get_or_create_state returns an existing row without inserting."""
+  service = DatabaseSessionService('sqlite+aiosqlite:///:memory:')
+  try:
+    await service._prepare_tables()
+    schema = service._get_schema_classes()
+
+    # Pre-create the app_state row.
+    async with service.database_session_factory() as sql_session:
+      sql_session.add(schema.StorageAppState(app_name='app1', state={'k': 'v'}))
+      await sql_session.commit()
+
+    # _get_or_create_state should find and return it.
+    async with service.database_session_factory() as sql_session:
+      row = await database_session_service._get_or_create_state(
+          sql_session=sql_session,
+          state_model=schema.StorageAppState,
+          primary_key='app1',
+          defaults={'app_name': 'app1', 'state': {}},
+      )
+      assert row.app_name == 'app1'
+      assert row.state == {'k': 'v'}
+  finally:
+    await service.close()
+
+
+@pytest.mark.asyncio
+async def test_get_or_create_state_creates_new_row():
+  """_get_or_create_state creates a row when none exists."""
+  service = DatabaseSessionService('sqlite+aiosqlite:///:memory:')
+  try:
+    await service._prepare_tables()
+    schema = service._get_schema_classes()
+
+    async with service.database_session_factory() as sql_session:
+      row = await database_session_service._get_or_create_state(
+          sql_session=sql_session,
+          state_model=schema.StorageAppState,
+          primary_key='new_app',
+          defaults={'app_name': 'new_app', 'state': {}},
+      )
+      await sql_session.commit()
+      assert row.app_name == 'new_app'
+      assert row.state == {}
+
+    # Verify the row was actually persisted.
+    async with service.database_session_factory() as sql_session:
+      persisted = await sql_session.get(schema.StorageAppState, 'new_app')
+      assert persisted is not None
+  finally:
+    await service.close()
+
+
+@pytest.mark.asyncio
+async def test_get_or_create_state_handles_race_condition():
+  """_get_or_create_state recovers when a concurrent INSERT wins the race.
+
+  Simulates the race from https://github.com/google/adk-python/issues/4954:
+  the initial SELECT returns None (another caller hasn't committed yet), but
+  by the time we INSERT, the other caller has committed — so the INSERT fails
+  with IntegrityError and we fall back to re-fetching.
+  """
+  service = DatabaseSessionService('sqlite+aiosqlite:///:memory:')
+  try:
+    await service._prepare_tables()
+    schema = service._get_schema_classes()
+
+    # Pre-create the row to guarantee the INSERT will fail.
+    async with service.database_session_factory() as sql_session:
+      sql_session.add(schema.StorageAppState(app_name='race_app', state={}))
+      await sql_session.commit()
+
+    # Patch session.get to return None on the first call (simulating the
+    # race window), then fall through to the real implementation.
+    async with service.database_session_factory() as sql_session:
+      original_get = sql_session.get
+      call_count = 0
+
+      async def patched_get(*args, **kwargs):
+        nonlocal call_count
+        call_count += 1
+        if call_count == 1:
+          return None  # Simulate: row not yet visible
+        return await original_get(*args, **kwargs)
+
+      sql_session.get = patched_get
+
+      row = await database_session_service._get_or_create_state(
+          sql_session=sql_session,
+          state_model=schema.StorageAppState,
+          primary_key='race_app',
+          defaults={'app_name': 'race_app', 'state': {}},
+      )
+      assert row.app_name == 'race_app'
+      # The function should have called get twice: once before the INSERT
+      # (patched to return None) and once after the IntegrityError.
+      assert call_count == 2
+  finally:
+    await service.close()
+
+
+@pytest.mark.asyncio
+async def test_create_session_sequential_same_app_name():
+  """Sequential create_session calls for the same app_name work correctly.
+
+  The second call reuses the existing app_states row.
+  """
+  service = DatabaseSessionService('sqlite+aiosqlite:///:memory:')
+  try:
+    s1 = await service.create_session(
+        app_name='shared', user_id='u1', session_id='s1'
+    )
+    s2 = await service.create_session(
+        app_name='shared', user_id='u2', session_id='s2'
+    )
+    assert s1.app_name == 'shared'
+    assert s2.app_name == 'shared'
+
+    got1 = await service.get_session(
+        app_name='shared', user_id='u1', session_id='s1'
+    )
+    got2 = await service.get_session(
+        app_name='shared', user_id='u2', session_id='s2'
+    )
+    assert got1 is not None
+    assert got2 is not None
+  finally:
+    await service.close()
+
+
 @pytest.mark.asyncio
 async def test_prepare_tables_idempotent_after_creation():
   """Calling _prepare_tables multiple times is safe and idempotent.