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

src/agents/extensions/memory/sqlalchemy_session.py

@@ -25,7 +25,9 @@
 
 import asyncio
 import json
-from typing import Any
+import threading
+from typing import Any, ClassVar
+from weakref import WeakKeyDictionary
 
 from sqlalchemy import (
     TIMESTAMP,
@@ -43,6 +45,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 +56,24 @@
 class SQLAlchemySession(SessionABC):
     """SQLAlchemy implementation of :pyclass:`agents.memory.session.Session`."""
 
+    _engine_init_locks: ClassVar[WeakKeyDictionary[AsyncEngine, asyncio.Lock]] = (
+        WeakKeyDictionary()
+    )
+    _engine_init_locks_guard: ClassVar[threading.Lock] = threading.Lock()
     _metadata: MetaData
     _sessions: Table
     _messages: Table
     session_settings: SessionSettings | None = None
 
+    @classmethod
+    def _get_engine_init_lock(cls, engine: AsyncEngine) -> asyncio.Lock:
+        with cls._engine_init_locks_guard:
+            lock = cls._engine_init_locks.get(engine)
+            if lock is None:
+                lock = asyncio.Lock()
+                cls._engine_init_locks[engine] = lock
+            return lock
+
     def __init__(
         self,
         session_id: str,
@@ -86,6 +102,7 @@ def __init__(
         self.session_settings = session_settings or SessionSettings()
         self._engine = engine
         self._lock = asyncio.Lock()
+        self._init_lock = self._get_engine_init_lock(engine)
 
         self._metadata = MetaData()
         self._sessions = Table(
@@ -182,7 +199,13 @@ 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
+
+        async with self._init_lock:
+            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
@@ -248,40 +271,45 @@ async def add_items(self, items: list[TResponseInputItem]) -> None:
         if not items:
             return
 
-        await self._ensure_tables()
-        payload = [
-            {
-                "session_id": self.session_id,
-                "message_data": await self._serialize_item(item),
-            }
-            for item in items
-        ]
-
-        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
-                existing = await sess.execute(
-                    select(self._sessions.c.session_id).where(
-                        self._sessions.c.session_id == self.session_id
+        async with self._lock:
+            await self._ensure_tables()
+            payload = [
+                {
+                    "session_id": self.session_id,
+                    "message_data": await self._serialize_item(item),
+                }
+                for item in items
+            ]
+
+            async with self._session_factory() as sess:
+                async with sess.begin():
+                    # 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
+                    if not existing.scalar_one_or_none():
+                        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)
+
+                    # Touch updated_at column
                     await sess.execute(
-                        insert(self._sessions).values({"session_id": self.session_id})
+                        update(self._sessions)
+                        .where(self._sessions.c.session_id == self.session_id)
+                        .values(updated_at=sql_text("CURRENT_TIMESTAMP"))
                     )
 
-                # Insert messages in bulk
-                await sess.execute(insert(self._messages), payload)
-
-                # Touch updated_at column
-                await sess.execute(
-                    update(self._sessions)
-                    .where(self._sessions.c.session_id == self.session_id)
-                    .values(updated_at=sql_text("CURRENT_TIMESTAMP"))
-                )
-
     async def pop_item(self) -> TResponseInputItem | None:
         """Remove and return the most recent item from the session.
 

tests/extensions/memory/test_sqlalchemy_session.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import asyncio
 import json
 from collections.abc import Iterable, Sequence
 from contextlib import asynccontextmanager
@@ -203,6 +204,100 @@ 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_get_items_same_timestamp_consistent_order():
     """Test that items with identical timestamps keep insertion order."""
     session_id = "same_timestamp_test"