fix(memory): make SQLAlchemySession first writes race-safe (#2725)

Author: felmonon

src/agents/extensions/memory/sqlalchemy_session.py

@@ -25,7 +25,8 @@
 
 import asyncio
 import json
-from typing import Any
+import threading
+from typing import Any, ClassVar
 
 from sqlalchemy import (
     TIMESTAMP,
@@ -43,6 +44,7 @@
     text as sql_text,
     update,
 )
+from sqlalchemy.exc import IntegrityError
 from sqlalchemy.ext.asyncio import AsyncEngine, async_sessionmaker, create_async_engine
 
 from ...items import TResponseInputItem
@@ -53,11 +55,29 @@
 class SQLAlchemySession(SessionABC):
     """SQLAlchemy implementation of :pyclass:`agents.memory.session.Session`."""
 
+    _table_init_locks: ClassVar[dict[tuple[str, str, str], threading.Lock]] = {}
+    _table_init_locks_guard: ClassVar[threading.Lock] = threading.Lock()
     _metadata: MetaData
     _sessions: Table
     _messages: Table
     session_settings: SessionSettings | None = None
 
+    @classmethod
+    def _get_table_init_lock(
+        cls, engine: AsyncEngine, sessions_table: str, messages_table: str
+    ) -> threading.Lock:
+        lock_key = (
+            engine.url.render_as_string(hide_password=True),
+            sessions_table,
+            messages_table,
+        )
+        with cls._table_init_locks_guard:
+            lock = cls._table_init_locks.get(lock_key)
+            if lock is None:
+                lock = threading.Lock()
+                cls._table_init_locks[lock_key] = lock
+            return lock
+
     def __init__(
         self,
         session_id: str,
@@ -85,7 +105,11 @@ def __init__(
         self.session_id = session_id
         self.session_settings = session_settings or SessionSettings()
         self._engine = engine
-        self._lock = asyncio.Lock()
+        self._init_lock = (
+            self._get_table_init_lock(engine, sessions_table, messages_table)
+            if create_tables
+            else None
+        )
 
         self._metadata = MetaData()
         self._sessions = Table(
@@ -182,10 +206,23 @@ async def _deserialize_item(self, item: str) -> TResponseInputItem:
     # ------------------------------------------------------------------
     async def _ensure_tables(self) -> None:
         """Ensure tables are created before any database operations."""
-        if self._create_tables:
+        if not self._create_tables:
+            return
+
+        assert self._init_lock is not None
+        while not self._init_lock.acquire(blocking=False):
+            # Poll without handing lock acquisition to a background thread so
+            # cancellation cannot strand the shared init lock in the acquired state.
+            await asyncio.sleep(0.01)
+        try:
+            if not self._create_tables:
+                return
+
             async with self._engine.begin() as conn:
                 await conn.run_sync(self._metadata.create_all)
             self._create_tables = False  # Only create once
+        finally:
+            self._init_lock.release()
 
     async def get_items(self, limit: int | None = None) -> list[TResponseInputItem]:
         """Retrieve the conversation history for this session.
@@ -259,18 +296,22 @@ async def add_items(self, items: list[TResponseInputItem]) -> None:
 
         async with self._session_factory() as sess:
             async with sess.begin():
-                # Ensure the parent session row exists - use merge for cross-DB compatibility
-                # Check if session exists
+                # Avoid check-then-insert races on the first write while keeping
+                # the common path free of avoidable integrity exceptions.
                 existing = await sess.execute(
                     select(self._sessions.c.session_id).where(
                         self._sessions.c.session_id == self.session_id
                     )
                 )
                 if not existing.scalar_one_or_none():
-                    # Session doesn't exist, create it
-                    await sess.execute(
-                        insert(self._sessions).values({"session_id": self.session_id})
-                    )
+                    try:
+                        async with sess.begin_nested():
+                            await sess.execute(
+                                insert(self._sessions).values({"session_id": self.session_id})
+                            )
+                    except IntegrityError:
+                        # Another concurrent writer created the parent row first.
+                        pass
 
                 # Insert messages in bulk
                 await sess.execute(insert(self._messages), payload)

tests/extensions/memory/test_sqlalchemy_session.py

@@ -1,6 +1,8 @@
 from __future__ import annotations
 
+import asyncio
 import json
+import threading
 from collections.abc import Iterable, Sequence
 from contextlib import asynccontextmanager
 from datetime import datetime, timedelta
@@ -203,6 +205,275 @@ async def test_add_empty_items_list():
     assert len(items_after_add) == 0
 
 
+async def test_add_items_concurrent_first_access_with_create_tables(tmp_path):
+    """Concurrent first writes should not race table creation or drop items."""
+    db_url = f"sqlite+aiosqlite:///{tmp_path / 'concurrent_first_access.db'}"
+    session = SQLAlchemySession.from_url(
+        "concurrent_first_access",
+        url=db_url,
+        create_tables=True,
+    )
+    submitted = [f"msg-{i}" for i in range(25)]
+
+    async def worker(content: str) -> None:
+        await session.add_items([{"role": "user", "content": content}])
+
+    results = await asyncio.gather(
+        *(worker(content) for content in submitted),
+        return_exceptions=True,
+    )
+
+    assert [result for result in results if isinstance(result, Exception)] == []
+
+    stored = await session.get_items()
+    assert len(stored) == len(submitted)
+    stored_contents: list[str] = []
+    for item in stored:
+        content = item.get("content")
+        assert isinstance(content, str)
+        stored_contents.append(content)
+    assert sorted(stored_contents) == sorted(submitted)
+
+
+async def test_add_items_concurrent_first_write_after_tables_exist(tmp_path):
+    """Concurrent first writes should not race parent session creation."""
+    db_url = f"sqlite+aiosqlite:///{tmp_path / 'concurrent_first_write.db'}"
+    setup_session = SQLAlchemySession.from_url(
+        "concurrent_first_write",
+        url=db_url,
+        create_tables=True,
+    )
+    await setup_session.get_items()
+
+    session = SQLAlchemySession.from_url(
+        "concurrent_first_write",
+        url=db_url,
+        create_tables=False,
+    )
+    submitted = [f"msg-{i}" for i in range(25)]
+
+    async def worker(content: str) -> None:
+        await session.add_items([{"role": "user", "content": content}])
+
+    results = await asyncio.gather(
+        *(worker(content) for content in submitted),
+        return_exceptions=True,
+    )
+
+    assert [result for result in results if isinstance(result, Exception)] == []
+
+    stored = await session.get_items()
+    assert len(stored) == len(submitted)
+    stored_contents: list[str] = []
+    for item in stored:
+        content = item.get("content")
+        assert isinstance(content, str)
+        stored_contents.append(content)
+    assert sorted(stored_contents) == sorted(submitted)
+
+
+async def test_add_items_concurrent_first_access_across_sessions_with_shared_engine(tmp_path):
+    """Concurrent first writes should not race table creation across session instances."""
+    db_url = f"sqlite+aiosqlite:///{tmp_path / 'concurrent_shared_engine.db'}"
+    engine = create_async_engine(db_url)
+    try:
+        session_a = SQLAlchemySession("shared_engine_a", engine=engine, create_tables=True)
+        session_b = SQLAlchemySession("shared_engine_b", engine=engine, create_tables=True)
+
+        results = await asyncio.gather(
+            session_a.add_items([{"role": "user", "content": "one"}]),
+            session_b.add_items([{"role": "user", "content": "two"}]),
+            return_exceptions=True,
+        )
+
+        assert [result for result in results if isinstance(result, Exception)] == []
+
+        stored_a = await session_a.get_items()
+        assert len(stored_a) == 1
+        assert stored_a[0].get("content") == "one"
+
+        stored_b = await session_b.get_items()
+        assert len(stored_b) == 1
+        assert stored_b[0].get("content") == "two"
+    finally:
+        await engine.dispose()
+
+
+async def test_add_items_concurrent_first_access_across_from_url_sessions(tmp_path):
+    """Concurrent first writes should not race table creation across from_url sessions."""
+    db_url = f"sqlite+aiosqlite:///{tmp_path / 'concurrent_from_url.db'}"
+    session_a = SQLAlchemySession.from_url("from_url_a", url=db_url, create_tables=True)
+    session_b = SQLAlchemySession.from_url("from_url_b", url=db_url, create_tables=True)
+    try:
+        results = await asyncio.gather(
+            session_a.add_items([{"role": "user", "content": "one"}]),
+            session_b.add_items([{"role": "user", "content": "two"}]),
+            return_exceptions=True,
+        )
+
+        assert [result for result in results if isinstance(result, Exception)] == []
+
+        stored_a = await session_a.get_items()
+        assert len(stored_a) == 1
+        assert stored_a[0].get("content") == "one"
+
+        stored_b = await session_b.get_items()
+        assert len(stored_b) == 1
+        assert stored_b[0].get("content") == "two"
+    finally:
+        await session_a.engine.dispose()
+        await session_b.engine.dispose()
+
+
+async def test_add_items_concurrent_first_access_across_from_url_sessions_cross_loop(tmp_path):
+    """Concurrent first writes should not race or hang across event loops."""
+    db_url = f"sqlite+aiosqlite:///{tmp_path / 'concurrent_from_url_cross_loop.db'}"
+    barrier = threading.Barrier(2)
+    results: list[tuple[str, str, Any]] = []
+    results_lock = threading.Lock()
+
+    def worker(session_id: str, content: str) -> None:
+        async def run() -> tuple[str, Any]:
+            session = SQLAlchemySession.from_url(session_id, url=db_url, create_tables=True)
+            barrier.wait()
+            try:
+                await asyncio.wait_for(
+                    session.add_items([{"role": "user", "content": content}]),
+                    timeout=5,
+                )
+                stored = await session.get_items()
+                return ("ok", stored)
+            finally:
+                await session.engine.dispose()
+
+        try:
+            status, payload = asyncio.run(run())
+        except Exception as exc:
+            status, payload = type(exc).__name__, str(exc)
+
+        with results_lock:
+            results.append((session_id, status, payload))
+
+    threads = [
+        threading.Thread(target=worker, args=("from_url_cross_loop_a", "one")),
+        threading.Thread(target=worker, args=("from_url_cross_loop_b", "two")),
+    ]
+    for thread in threads:
+        thread.start()
+    for thread in threads:
+        await asyncio.to_thread(thread.join)
+
+    assert len(results) == 2
+    assert [status for _, status, _ in results] == ["ok", "ok"]
+
+    stored_by_session = {
+        session_id: cast(list[TResponseInputItem], payload) for session_id, _, payload in results
+    }
+    assert stored_by_session["from_url_cross_loop_a"][0].get("content") == "one"
+    assert stored_by_session["from_url_cross_loop_b"][0].get("content") == "two"
+
+
+async def test_add_items_concurrent_first_access_with_shared_session_cross_loop(tmp_path):
+    """A shared session instance should not hang when used from two event loops."""
+    db_url = f"sqlite+aiosqlite:///{tmp_path / 'shared_session_cross_loop.db'}"
+    session = SQLAlchemySession.from_url(
+        "shared_session_cross_loop",
+        url=db_url,
+        create_tables=True,
+    )
+    barrier = threading.Barrier(2)
+    results: list[tuple[str, str]] = []
+    results_lock = threading.Lock()
+
+    def worker(content: str) -> None:
+        async def run() -> None:
+            barrier.wait()
+            await asyncio.wait_for(
+                session.add_items([{"role": "user", "content": content}]),
+                timeout=5,
+            )
+
+        try:
+            asyncio.run(run())
+            status = "ok"
+        except Exception as exc:
+            status = type(exc).__name__
+
+        with results_lock:
+            results.append((content, status))
+
+    threads = [
+        threading.Thread(target=worker, args=("one",)),
+        threading.Thread(target=worker, args=("two",)),
+    ]
+    try:
+        for thread in threads:
+            thread.start()
+        for thread in threads:
+            await asyncio.to_thread(thread.join)
+
+        assert sorted(results) == [("one", "ok"), ("two", "ok")]
+
+        stored = await session.get_items()
+        stored_contents: list[str] = []
+        for item in stored:
+            content = item.get("content")
+            assert isinstance(content, str)
+            stored_contents.append(content)
+        assert sorted(stored_contents) == ["one", "two"]
+    finally:
+        await session.engine.dispose()
+
+
+async def test_add_items_cancelled_waiter_does_not_strand_table_init_lock(tmp_path):
+    """Cancelling a waiting initializer must not leave the shared init lock acquired."""
+    db_url = f"sqlite+aiosqlite:///{tmp_path / 'cancelled_table_init_waiter.db'}"
+    holder = SQLAlchemySession.from_url("holder", url=db_url, create_tables=True)
+    waiter = SQLAlchemySession.from_url("waiter", url=db_url, create_tables=True)
+    follower = SQLAlchemySession.from_url("follower", url=db_url, create_tables=True)
+
+    assert holder._init_lock is waiter._init_lock
+    assert waiter._init_lock is follower._init_lock
+    assert holder._init_lock is not None
+
+    acquired = holder._init_lock.acquire(blocking=False)
+    assert acquired
+
+    try:
+        blocked = asyncio.create_task(waiter.add_items([{"role": "user", "content": "waiter"}]))
+        await asyncio.sleep(0.05)
+        blocked.cancel()
+        with pytest.raises(asyncio.CancelledError):
+            await blocked
+    finally:
+        holder._init_lock.release()
+
+    try:
+        await asyncio.wait_for(
+            follower.add_items([{"role": "user", "content": "follower"}]),
+            timeout=2,
+        )
+        stored = await follower.get_items()
+        assert len(stored) == 1
+        assert stored[0].get("content") == "follower"
+    finally:
+        await holder.engine.dispose()
+        await waiter.engine.dispose()
+        await follower.engine.dispose()
+
+
+async def test_create_tables_false_does_not_allocate_shared_init_lock(tmp_path):
+    """Sessions that skip auto-create should not populate the shared lock map."""
+    db_url = f"sqlite+aiosqlite:///{tmp_path / 'no_create_tables_lock.db'}"
+    before = len(SQLAlchemySession._table_init_locks)
+    session = SQLAlchemySession.from_url("no_create_tables_lock", url=db_url, create_tables=False)
+    try:
+        assert session._init_lock is None
+        assert len(SQLAlchemySession._table_init_locks) == before
+    finally:
+        await session.engine.dispose()
+
+
 async def test_get_items_same_timestamp_consistent_order():
     """Test that items with identical timestamps keep insertion order."""
     session_id = "same_timestamp_test"