fix: make session mutations atomic

docs/sessions/index.md

@@ -450,7 +450,7 @@ Notes:
 
 -   `from_uri(...)` creates and owns the `AsyncMongoClient` and closes it on `session.close()`. An owned-client session is terminal after `close()`, and subsequent session operations raise `RuntimeError`. If your application already manages a client, construct `MongoDBSession(...)` directly with `client=...`; in that case `session.close()` is a no-op, and lifecycle plus session usability stay with the caller.
 -   Connect to [MongoDB Atlas](https://www.mongodb.com/products/platform) by passing an `mongodb+srv://user:password@cluster.example.mongodb.net` URI to `from_uri(...)` with no other changes.
--   Two collections are used and both names are configurable via `sessions_collection=` (default `agent_sessions`) and `messages_collection=` (default `agent_messages`). Indexes are created automatically on first use. Each message document carries a monotonically increasing `seq` counter that preserves ordering across concurrent writers and processes.
+-   Two collections are used and both names are configurable via `sessions_collection=` (default `agent_sessions`) and `messages_collection=` (default `agent_messages`). Indexes are created automatically on first use. Each non-empty `add_items()` call writes one logical-batch document whose monotonically increasing `seq` orders the batch by its final item; legacy per-item message documents remain readable. A logical batch must fit within MongoDB's single-document size limit; an oversized batch fails atomically without storing a partial batch.
 -   Use `await session.ping()` to verify connectivity before your first run.
 
 ### Advanced SQLite sessions

src/agents/extensions/memory/async_sqlite_session.py

@@ -2,10 +2,10 @@
 
 import asyncio
 import json
-from collections.abc import AsyncIterator
+from collections.abc import AsyncIterator, Awaitable
 from contextlib import asynccontextmanager
 from pathlib import Path
-from typing import Any, cast
+from typing import Any, TypeVar, cast
 
 import aiosqlite
 
@@ -16,6 +16,20 @@
     coerce_session_settings,
     resolve_session_limit,
 )
+from ...memory.sqlite_session import _await_mutation
+
+_T = TypeVar("_T")
+
+
+async def _await_cleanup(awaitable: Awaitable[_T]) -> _T:
+    """Keep ownership of cleanup until it finishes despite repeated cancellation."""
+    task = asyncio.ensure_future(awaitable)
+    while True:
+        try:
+            return await asyncio.shield(task)
+        except asyncio.CancelledError:
+            if task.done():
+                return task.result()
 
 
 class AsyncSQLiteSession(SessionABC):
@@ -57,6 +71,7 @@ def __init__(
         self.sessions_table = sessions_table
         self.messages_table = messages_table
         self._connection: aiosqlite.Connection | None = None
+        self._quarantined_connections: set[aiosqlite.Connection] = set()
         self._lock = asyncio.Lock()
         self._init_lock = asyncio.Lock()
         self._closed = False
@@ -102,9 +117,24 @@ async def _get_connection(self) -> aiosqlite.Connection:
 
         async with self._init_lock:
             if self._connection is None:
-                self._connection = await aiosqlite.connect(str(self.db_path))
-                await self._connection.execute("PRAGMA journal_mode=WAL")
-                await self._init_db_for_connection(self._connection)
+                connect_task = asyncio.ensure_future(aiosqlite.connect(str(self.db_path)))
+                try:
+                    connection = await asyncio.shield(connect_task)
+                except BaseException:
+                    try:
+                        connection = await _await_cleanup(connect_task)
+                    except BaseException:
+                        pass
+                    else:
+                        await self._close_owned_connection(connection)
+                    raise
+                try:
+                    await connection.execute("PRAGMA journal_mode=WAL")
+                    await self._init_db_for_connection(connection)
+                except BaseException:
+                    await self._close_owned_connection(connection)
+                    raise
+                self._connection = connection
 
         return self._connection
 
@@ -121,6 +151,38 @@ async def _locked_connection(self) -> AsyncIterator[aiosqlite.Connection]:
             conn = await self._get_connection()
             yield conn
 
+    @asynccontextmanager
+    async def _write_connection(self) -> AsyncIterator[aiosqlite.Connection]:
+        """Provide a connection that cannot retain a failed write transaction."""
+        async with self._locked_connection() as conn:
+            try:
+                yield conn
+            except BaseException:
+                try:
+                    await _await_cleanup(conn.rollback())
+                except BaseException:
+                    await self._invalidate_connection(conn)
+                raise
+
+    async def _invalidate_connection(self, conn: aiosqlite.Connection) -> None:
+        """Close and evict a connection that could not roll back safely."""
+        close_error = await self._close_owned_connection(conn)
+        if self._connection is conn:
+            self._connection = None
+        if str(self.db_path) == ":memory:" or close_error is not None:
+            self._closed = True
+
+    async def _close_owned_connection(self, conn: aiosqlite.Connection) -> BaseException | None:
+        """Close an owned connection or retain it for a later cleanup retry."""
+        try:
+            await _await_cleanup(conn.close())
+        except BaseException as exc:
+            self._quarantined_connections.add(conn)
+            self._closed = True
+            return exc
+        self._quarantined_connections.discard(conn)
+        return None
+
     async def get_items(self, limit: int | None = None) -> list[TResponseInputItem]:
         """Retrieve the conversation history for this session.
 
@@ -206,7 +268,7 @@ async def add_items(self, items: list[TResponseInputItem]) -> None:
         if not items:
             return
 
-        async with self._locked_connection() as conn:
+        async with self._write_connection() as conn:
             await conn.execute(
                 f"""
                 INSERT OR IGNORE INTO {self.sessions_table} (session_id) VALUES (?)
@@ -231,15 +293,16 @@ async def add_items(self, items: list[TResponseInputItem]) -> None:
                 (self.session_id,),
             )
 
-            await conn.commit()
+            await _await_mutation(conn.commit())
 
     async def pop_item(self) -> TResponseInputItem | None:
         """Remove and return the most recent item from the session.
 
         Returns:
             The most recent item if it exists, None if the session is empty
         """
-        async with self._locked_connection() as conn:
+
+        async with self._write_connection() as conn:
             cursor = await conn.execute(
                 f"""
                 DELETE FROM {self.messages_table}
@@ -256,7 +319,7 @@ async def pop_item(self) -> TResponseInputItem | None:
 
             result = await cursor.fetchone()
             await cursor.close()
-            await conn.commit()
+            await _await_mutation(conn.commit())
 
             while result:
                 message_data = result[0]
@@ -278,13 +341,14 @@ async def pop_item(self) -> TResponseInputItem | None:
                     )
                     result = await cursor.fetchone()
                     await cursor.close()
-                    await conn.commit()
+                    await _await_mutation(conn.commit())
 
         return None
 
     async def clear_session(self) -> None:
         """Clear all items for this session."""
-        async with self._locked_connection() as conn:
+
+        async with self._write_connection() as conn:
             await conn.execute(
                 f"DELETE FROM {self.messages_table} WHERE session_id = ?",
                 (self.session_id,),
@@ -293,18 +357,39 @@ async def clear_session(self) -> None:
                 f"DELETE FROM {self.sessions_table} WHERE session_id = ?",
                 (self.session_id,),
             )
-            await conn.commit()
+            await _await_mutation(conn.commit())
 
     async def close(self) -> None:
         """Close the database connection.
 
         The session becomes terminal from the first close attempt: subsequent
         operations raise RuntimeError rather than reopening the database. Repeated
-        and concurrent calls are safe no-ops.
+        and concurrent calls are safe. A repeated call retries any owned
+        connection whose previous close did not complete.
         """
         async with self._lock:
             self._closed = True
-            if self._connection is None:
-                return
-            await self._connection.close()
-            self._connection = None
+            connections = set(self._quarantined_connections)
+            if self._connection is not None:
+                connections.add(self._connection)
+
+            first_error: BaseException | None = None
+            cancellation: asyncio.CancelledError | None = None
+            for connection in connections:
+                close_task = asyncio.create_task(self._close_owned_connection(connection))
+                try:
+                    close_error = await asyncio.shield(close_task)
+                except asyncio.CancelledError as exc:
+                    if cancellation is None:
+                        cancellation = exc
+                    close_error = await _await_cleanup(close_task)
+                if close_error is None:
+                    if self._connection is connection:
+                        self._connection = None
+                elif first_error is None:
+                    first_error = close_error
+
+            if cancellation is not None:
+                raise cancellation
+            if first_error is not None:
+                raise first_error

src/agents/extensions/memory/dapr_session.py

@@ -334,6 +334,23 @@ async def add_items(self, items: list[TResponseInputItem]) -> None:
         async with self._lock:
             self._check_not_closed()
             serialized_items: list[str] = [await self._serialize_item(item) for item in items]
+
+            # Persist ancillary metadata before the authoritative history so a metadata
+            # failure cannot make a committed batch look safe to retry.
+            now = str(int(time.time()))
+            metadata = {
+                "session_id": self.session_id,
+                "created_at": now,
+                "updated_at": now,
+            }
+            await self._dapr_client.save_state(
+                store_name=self._state_store_name,
+                key=self._metadata_key,
+                value=json.dumps(metadata),
+                state_metadata=self._get_metadata(),
+                options=self._get_state_options(),
+            )
+
             attempt = 0
             while True:
                 attempt += 1
@@ -362,20 +379,6 @@ async def add_items(self, items: list[TResponseInputItem]) -> None:
                         continue
                     raise
 
-            # Update metadata
-            metadata = {
-                "session_id": self.session_id,
-                "created_at": str(int(time.time())),
-                "updated_at": str(int(time.time())),
-            }
-            await self._dapr_client.save_state(
-                store_name=self._state_store_name,
-                key=self._metadata_key,
-                value=json.dumps(metadata),
-                state_metadata=self._get_metadata(),
-                options=self._get_state_options(),
-            )
-
     async def pop_item(self) -> TResponseInputItem | None:
         """Remove and return the most recent item from the session.