diff --git a/src/google/adk/sessions/database_session_service.py b/src/google/adk/sessions/database_session_service.py index a8716d0c9b..60bf0f995d 100644 --- a/src/google/adk/sessions/database_session_service.py +++ b/src/google/adk/sessions/database_session_service.py @@ -639,7 +639,18 @@ async def create_session( storage_app_state.state, storage_user_state.state, session_state ) # Call to_session before commit to avoid post-commit lazy-load. - await sql_session.flush() + try: + await sql_session.flush() + except IntegrityError: + # A concurrent caller won the race on this (app_name, user_id, + # session_id) primary key: the has_user_provided_id check above is + # not atomic with this insert, so two callers can both pass it and + # then race the same insert. Same failure mode _get_or_create_state + # guards against for app_state/user_state; surface the same clean + # error here instead of letting the raw IntegrityError propagate. + raise AlreadyExistsError( + f"Session with id {session_id} already exists." + ) session = storage_session.to_session( state=merged_state, is_sqlite=is_sqlite, is_postgresql=is_postgresql ) diff --git a/tests/unittests/sessions/test_session_service.py b/tests/unittests/sessions/test_session_service.py index a279701fab..b80e6a1d55 100644 --- a/tests/unittests/sessions/test_session_service.py +++ b/tests/unittests/sessions/test_session_service.py @@ -1190,6 +1190,66 @@ async def test_create_session_with_existing_id_raises_error(session_service): ) +@pytest.mark.asyncio +async def test_create_session_concurrent_same_id_raises_already_exists_error( + tmp_path, +): + """Two concurrent create_session() calls for the same caller-provided id. + + The has_user_provided_id existence check in create_session() is not atomic + with the insert that follows it, so both callers can pass the check and + then race the same INSERT. The loser must see a clean AlreadyExistsError + (mirroring the up-front check above and the _get_or_create_state + savepoint pattern for app_state/user_state), not a raw IntegrityError. + + Uses a file-backed sqlite db (not ':memory:') so the two concurrent + sessions get real, independent connections from the pool instead of + sharing the single StaticPool connection ':memory:' relies on to survive + across connections -- sharing one physical connection between the two + concurrent sessions here made the loser's rollback able to interleave + with the winner's commit on the same connection. + """ + db_path = tmp_path / 'race.db' + session_service = DatabaseSessionService(f'sqlite+aiosqlite:///{db_path}') + + async with session_service: + app_name = 'my_app' + user_id = 'user' + + # Pre-warm app_state/user_state with an unrelated session first, so the + # race below is purely on the StorageSession primary key and not + # confounded by the (separate) app_state/user_state creation race. + await session_service.create_session( + app_name=app_name, user_id=user_id, session_id='warmup-session' + ) + + for i in range(5): + session_id = f'race-session-{i}' + results = await asyncio.gather( + session_service.create_session( + app_name=app_name, user_id=user_id, session_id=session_id + ), + session_service.create_session( + app_name=app_name, user_id=user_id, session_id=session_id + ), + 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], AlreadyExistsError) + assert session_id in str(errors[0]) + + final_session = await session_service.get_session( + app_name=app_name, user_id=user_id, session_id=session_id + ) + assert final_session is not None + assert final_session.id == successes[0].id + + @pytest.mark.asyncio async def test_append_event_bytes(session_service): app_name = 'my_app'