From a0c61eea15cdb3c3da2e90b7809c9a1b85b4498c Mon Sep 17 00:00:00 2001 From: Kazuhiro Sera Date: Wed, 5 Aug 2026 21:44:07 +0900 Subject: [PATCH 1/7] fix: make session mutations atomic --- docs/sessions/index.md | 2 +- .../memory/advanced_sqlite_session.py | 436 ++++++++---------- .../extensions/memory/async_sqlite_session.py | 119 ++++- src/agents/extensions/memory/dapr_session.py | 31 +- .../extensions/memory/mongodb_session.py | 100 ++-- src/agents/extensions/memory/redis_session.py | 165 +++++-- .../extensions/memory/sqlalchemy_session.py | 14 +- src/agents/memory/sqlite_session.py | 131 ++++-- .../memory/test_advanced_sqlite_session.py | 297 ++++++++++++ .../memory/test_async_sqlite_session.py | 431 +++++++++++++++++ tests/extensions/memory/test_dapr_session.py | 48 ++ .../extensions/memory/test_mongodb_session.py | 169 ++++++- tests/extensions/memory/test_redis_session.py | 374 +++++++++++++++ .../memory/test_sqlalchemy_session.py | 85 ++++ tests/memory/test_session.py | 214 ++++++++- 15 files changed, 2245 insertions(+), 371 deletions(-) diff --git a/docs/sessions/index.md b/docs/sessions/index.md index 95f66172d0..2c0cd5f2c1 100644 --- a/docs/sessions/index.md +++ b/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 diff --git a/src/agents/extensions/memory/advanced_sqlite_session.py b/src/agents/extensions/memory/advanced_sqlite_session.py index de88927bc2..67d86f1cc6 100644 --- a/src/agents/extensions/memory/advanced_sqlite_session.py +++ b/src/agents/extensions/memory/advanced_sqlite_session.py @@ -22,6 +22,7 @@ ) from ...memory import SQLiteSession from ...memory.session_settings import SessionSettings, resolve_session_limit +from ...memory.sqlite_session import _await_mutation def _content_preview(content: Any, max_length: int | None = None) -> str: @@ -67,7 +68,14 @@ def __init__( **kwargs, ) if create_tables: - self._init_structure_tables() + try: + self._init_structure_tables() + except BaseException: + try: + self.close() + except BaseException: + pass + raise self._current_branch_id = "main" # Bumped (under the connection lock) whenever clear_session() wipes the # session. switch_to_branch / create_branch_from_turn capture the @@ -97,7 +105,7 @@ def _init_structure_tables(self): Creates the message_structure, branch_reservations, and turn_usage tables with appropriate indexes for conversation branching and usage analytics. """ - with self._locked_connection() as conn: + with self._write_connection() as conn: # Message structure with branch support conn.execute(f""" CREATE TABLE IF NOT EXISTS message_structure ( @@ -178,20 +186,17 @@ async def add_items(self, items: list[TResponseInputItem]) -> None: def _add_items_sync(): """Synchronous helper to add items and structure metadata together.""" - with self._locked_connection() as conn: - try: - # Keep both writes in one transaction so metadata failures do not leave orphans. - self._insert_items(conn, items) - self._insert_structure_metadata(conn, items) - conn.commit() - except Exception as exc: - conn.rollback() - log_model_and_tool_action_error( - self._logger, "Failed to add session items", exc - ) - raise + with self._write_connection() as conn: + # Keep both writes in one transaction so metadata failures do not leave orphans. + self._insert_items(conn, items) + self._insert_structure_metadata(conn, items) + conn.commit() - await asyncio.to_thread(_add_items_sync) + try: + await _await_mutation(asyncio.to_thread(_add_items_sync)) + except Exception as exc: + log_model_and_tool_action_error(self._logger, "Failed to add session items", exc) + raise async def get_items( self, @@ -300,77 +305,65 @@ async def pop_item(self) -> TResponseInputItem | None: branch_id = self._current_branch_id def _pop_item_sync(): - with self._locked_connection() as conn: + with self._write_connection() as conn: while True: with closing(conn.cursor()) as cursor: - # Find the most recent item on the snapshotted branch. + # Preserve every legacy branch ID before a pop can remove its + # final message_structure row. This stays inside the existing + # rollback boundary for the mutation. + self._ensure_branch_reservations_table(conn) + + # Atomically claim the newest structure row across processes. cursor.execute( """ - SELECT id, message_id, user_turn_number FROM message_structure - WHERE session_id = ? AND branch_id = ? - ORDER BY sequence_number DESC - LIMIT 1 + DELETE FROM message_structure + WHERE id = ( + SELECT id FROM message_structure + WHERE session_id = ? AND branch_id = ? + ORDER BY sequence_number DESC + LIMIT 1 + ) + RETURNING message_id, user_turn_number """, (self.session_id, branch_id), ) - row = cursor.fetchone() - if row is None: + claimed_row = cursor.fetchone() + if claimed_row is None: + conn.commit() return None - structure_id, message_id, user_turn_number = row - - # Read the message payload before removing anything. + message_id, user_turn_number = claimed_row cursor.execute( f"SELECT message_data FROM {self.messages_table} WHERE id = ?", (message_id,), ) message_row = cursor.fetchone() - try: - # Preserve every legacy branch ID before a pop can remove its - # final message_structure row. This stays inside the existing - # rollback boundary for the mutation. - self._ensure_branch_reservations_table(conn) - - # Remove the structure row for this branch, then drop - # the underlying message only if no other branch - # references it. + # Drop the underlying message only if no other branch references it. + self._cleanup_orphaned_messages_sync(conn) + + # If this was the last item of the turn on this + # branch, drop the now-stale turn_usage row for it. + if user_turn_number is not None: cursor.execute( - "DELETE FROM message_structure WHERE id = ?", - (structure_id,), + """ + SELECT COUNT(*) FROM message_structure + WHERE session_id = ? AND branch_id = ? + AND user_turn_number = ? + """, + (self.session_id, branch_id, user_turn_number), ) - self._cleanup_orphaned_messages_sync(conn) - - # If this was the last item of the turn on this - # branch, drop the now-stale turn_usage row for it. - if user_turn_number is not None: + if cursor.fetchone()[0] == 0: cursor.execute( """ - SELECT COUNT(*) FROM message_structure + DELETE FROM turn_usage WHERE session_id = ? AND branch_id = ? AND user_turn_number = ? """, (self.session_id, branch_id, user_turn_number), ) - if cursor.fetchone()[0] == 0: - cursor.execute( - """ - DELETE FROM turn_usage - WHERE session_id = ? AND branch_id = ? - AND user_turn_number = ? - """, - (self.session_id, branch_id, user_turn_number), - ) - conn.commit() - except Exception: - # _locked_connection() does not manage transactions; - # roll back explicitly so a failure partway through - # this delete sequence never leaves a partial - # mutation or an open transaction for a later - # operation on this connection to inherit. - conn.rollback() - raise + conn.commit() if message_row is None: # Structure row pointed at a missing message; keep looking. @@ -382,7 +375,7 @@ def _pop_item_sync(): # Drop corrupted JSON entries and keep looking for a valid item. continue - return await asyncio.to_thread(_pop_item_sync) + return await _await_mutation(asyncio.to_thread(_pop_item_sync)) async def clear_session(self) -> None: """Clear all items for this session. @@ -398,37 +391,27 @@ async def clear_session(self) -> None: """ def _clear_session_sync(): - with self._locked_connection() as conn: - try: - # Backfill legacy branch IDs before clearing their only durable - # identity evidence. - self._ensure_branch_reservations_table(conn) - conn.execute( - f"DELETE FROM {self.messages_table} WHERE session_id = ?", - (self.session_id,), - ) - conn.execute( - f"DELETE FROM {self.sessions_table} WHERE session_id = ?", - (self.session_id,), - ) - conn.execute( - "DELETE FROM message_structure WHERE session_id = ?", - (self.session_id,), - ) - conn.execute( - "DELETE FROM turn_usage WHERE session_id = ?", - (self.session_id,), - ) - conn.commit() - except Exception: - # _locked_connection() does not manage transactions; roll - # back explicitly so a failure partway through this delete - # sequence never leaves a partial mutation or an open - # transaction for a later operation on this connection to - # inherit. The in-memory branch state below is only updated - # after a successful commit, so it stays consistent with it. - conn.rollback() - raise + with self._write_connection() as conn: + # Backfill legacy branch IDs before clearing their only durable + # identity evidence. + self._ensure_branch_reservations_table(conn) + conn.execute( + f"DELETE FROM {self.messages_table} WHERE session_id = ?", + (self.session_id,), + ) + conn.execute( + f"DELETE FROM {self.sessions_table} WHERE session_id = ?", + (self.session_id,), + ) + conn.execute( + "DELETE FROM message_structure WHERE session_id = ?", + (self.session_id,), + ) + conn.execute( + "DELETE FROM turn_usage WHERE session_id = ?", + (self.session_id,), + ) + conn.commit() # All branches were removed, so reset the in-memory pointer to # 'main' while still holding the lock. Doing this inside the # locked operation keeps the reset atomic with the clear, so no @@ -439,7 +422,7 @@ def _clear_session_sync(): self._generation += 1 self._current_branch_id = "main" - await asyncio.to_thread(_clear_session_sync) + await _await_mutation(asyncio.to_thread(_clear_session_sync)) async def store_run_usage(self, result: RunResult) -> None: """Store usage data for the current conversation turn. @@ -591,7 +574,7 @@ async def _add_structure_metadata(self, items: list[TResponseInputItem]) -> None def _add_structure_sync(): """Synchronous helper to add structure metadata to database.""" - with self._locked_connection() as conn: + with self._write_connection() as conn: self._insert_structure_metadata(conn, items) conn.commit() @@ -709,12 +692,9 @@ async def _cleanup_orphaned_messages(self) -> int: def _cleanup_sync(): """Synchronous helper to cleanup orphaned messages.""" - with self._locked_connection() as conn: + with self._write_connection() as conn: deleted_count = self._cleanup_orphaned_messages_sync(conn) - if deleted_count: - conn.commit() - else: - conn.rollback() + conn.commit() return deleted_count return await asyncio.to_thread(_cleanup_sync) @@ -971,54 +951,50 @@ async def delete_branch(self, branch_id: str, force: bool = False) -> None: def _delete_sync(): """Synchronous helper to delete branch and associated data.""" - with self._locked_connection() as conn: - try: - # Backfill legacy branch IDs before deleting their message structure. - self._ensure_branch_reservations_table(conn) - with closing(conn.cursor()) as cursor: - # First verify the branch exists - cursor.execute( - """ - SELECT COUNT(*) FROM message_structure - WHERE session_id = ? AND branch_id = ? - """, - (self.session_id, branch_id), - ) + with self._write_connection() as conn: + # Backfill legacy branch IDs before deleting their message structure. + self._ensure_branch_reservations_table(conn) + with closing(conn.cursor()) as cursor: + # First verify the branch exists + cursor.execute( + """ + SELECT COUNT(*) FROM message_structure + WHERE session_id = ? AND branch_id = ? + """, + (self.session_id, branch_id), + ) - count = cursor.fetchone()[0] - if count == 0: - raise ValueError(f"Branch '{branch_id}' does not exist") + count = cursor.fetchone()[0] + if count == 0: + raise ValueError(f"Branch '{branch_id}' does not exist") - # Delete from turn_usage first (foreign key constraint) - cursor.execute( - """ - DELETE FROM turn_usage - WHERE session_id = ? AND branch_id = ? - """, - (self.session_id, branch_id), - ) + # Delete from turn_usage first (foreign key constraint) + cursor.execute( + """ + DELETE FROM turn_usage + WHERE session_id = ? AND branch_id = ? + """, + (self.session_id, branch_id), + ) - usage_deleted = cursor.rowcount + usage_deleted = cursor.rowcount - # Delete from message_structure - cursor.execute( - """ - DELETE FROM message_structure - WHERE session_id = ? AND branch_id = ? - """, - (self.session_id, branch_id), - ) + # Delete from message_structure + cursor.execute( + """ + DELETE FROM message_structure + WHERE session_id = ? AND branch_id = ? + """, + (self.session_id, branch_id), + ) - structure_deleted = cursor.rowcount + structure_deleted = cursor.rowcount - orphaned_messages_deleted = self._cleanup_orphaned_messages_sync(conn) + orphaned_messages_deleted = self._cleanup_orphaned_messages_sync(conn) - conn.commit() + conn.commit() - return usage_deleted, structure_deleted, orphaned_messages_deleted - except Exception: - conn.rollback() - raise + return usage_deleted, structure_deleted, orphaned_messages_deleted usage_deleted, structure_deleted, orphaned_messages_deleted = await asyncio.to_thread( _delete_sync @@ -1166,109 +1142,105 @@ async def _copy_messages_to_new_branch( def _copy_sync() -> tuple[str, Any]: """Synchronous helper to copy messages to new branch.""" - with self._locked_connection() as conn: - try: - # Acquire SQLite's write reservation before checking the branch ID so - # sessions in other processes cannot pass the same check concurrently. - conn.execute("BEGIN IMMEDIATE") - self._ensure_branch_reservations_table(conn) - with closing(conn.cursor()) as cursor: - cursor.execute( - f""" - SELECT am.message_data - FROM message_structure ms - JOIN {self.messages_table} am ON ms.message_id = am.id - WHERE ms.session_id = ? AND ms.branch_id = ? - AND ms.branch_turn_number = ? AND ms.message_type = 'user' - """, - (self.session_id, source_branch_id, from_turn_number), + with self._write_connection() as conn: + # Acquire SQLite's write reservation before checking the branch ID so + # sessions in other processes cannot pass the same check concurrently. + conn.execute("BEGIN IMMEDIATE") + self._ensure_branch_reservations_table(conn) + with closing(conn.cursor()) as cursor: + cursor.execute( + f""" + SELECT am.message_data + FROM message_structure ms + JOIN {self.messages_table} am ON ms.message_id = am.id + WHERE ms.session_id = ? AND ms.branch_id = ? + AND ms.branch_turn_number = ? AND ms.message_type = 'user' + """, + (self.session_id, source_branch_id, from_turn_number), + ) + result = cursor.fetchone() + if result is None: + raise ValueError( + f"Turn {from_turn_number} does not contain a user message " + f"in branch '{source_branch_id}'" ) - result = cursor.fetchone() - if result is None: - raise ValueError( - f"Turn {from_turn_number} does not contain a user message " - f"in branch '{source_branch_id}'" - ) - try: - content = json.loads(result[0]).get("content", "") - turn_content = content[:50] + "..." if len(content) > 50 else content - except Exception: - turn_content = "Unable to parse content" + try: + content = json.loads(result[0]).get("content", "") + turn_content = content[:50] + "..." if len(content) > 50 else content + except Exception: + turn_content = "Unable to parse content" + + branch_id = self._reserve_branch_id(cursor, new_branch_id, from_turn_number) + + # Get all messages before the branch point + cursor.execute( + """ + SELECT + ms.message_id, + ms.message_type, + ms.sequence_number, + ms.user_turn_number, + ms.branch_turn_number, + ms.tool_name + FROM message_structure ms + WHERE ms.session_id = ? AND ms.branch_id = ? + AND ms.branch_turn_number < ? + ORDER BY ms.sequence_number + """, + (self.session_id, source_branch_id, from_turn_number), + ) - branch_id = self._reserve_branch_id(cursor, new_branch_id, from_turn_number) + messages_to_copy = cursor.fetchall() - # Get all messages before the branch point + if messages_to_copy: + # Get the max sequence number for the new inserts cursor.execute( """ - SELECT - ms.message_id, - ms.message_type, - ms.sequence_number, - ms.user_turn_number, - ms.branch_turn_number, - ms.tool_name - FROM message_structure ms - WHERE ms.session_id = ? AND ms.branch_id = ? - AND ms.branch_turn_number < ? - ORDER BY ms.sequence_number + SELECT COALESCE(MAX(sequence_number), 0) + FROM message_structure + WHERE session_id = ? """, - (self.session_id, source_branch_id, from_turn_number), + (self.session_id,), ) - messages_to_copy = cursor.fetchall() - - if messages_to_copy: - # Get the max sequence number for the new inserts - cursor.execute( - """ - SELECT COALESCE(MAX(sequence_number), 0) - FROM message_structure - WHERE session_id = ? - """, - (self.session_id,), - ) - - seq_start = cursor.fetchone()[0] - - # Insert copied messages with new branch_id - new_structure_data = [] - for i, ( - msg_id, - msg_type, - _, - user_turn, - branch_turn, - tool_name, - ) in enumerate(messages_to_copy): - new_structure_data.append( - ( - self.session_id, - msg_id, # Same message_id (sharing the actual message data) - branch_id, - msg_type, - seq_start + i + 1, # New sequence number - user_turn, # Keep same global turn number - branch_turn, # Keep same branch turn number - tool_name, - ) + seq_start = cursor.fetchone()[0] + + # Insert copied messages with new branch_id + new_structure_data = [] + for i, ( + msg_id, + msg_type, + _, + user_turn, + branch_turn, + tool_name, + ) in enumerate(messages_to_copy): + new_structure_data.append( + ( + self.session_id, + msg_id, # Same message_id (sharing the actual message data) + branch_id, + msg_type, + seq_start + i + 1, # New sequence number + user_turn, # Keep same global turn number + branch_turn, # Keep same branch turn number + tool_name, ) - - cursor.executemany( - """ - INSERT INTO message_structure - (session_id, message_id, branch_id, message_type, sequence_number, - user_turn_number, branch_turn_number, tool_name) - VALUES (?, ?, ?, ?, ?, ?, ?, ?) - """, - new_structure_data, ) - conn.commit() - return branch_id, turn_content - except Exception: - conn.rollback() - raise + cursor.executemany( + """ + INSERT INTO message_structure + (session_id, message_id, branch_id, message_type, sequence_number, + user_turn_number, branch_turn_number, tool_name) + VALUES (?, ?, ?, ?, ?, ?, ?, ?) + """, + new_structure_data, + ) + + conn.commit() + return branch_id, turn_content return await asyncio.to_thread(_copy_sync) @@ -1671,7 +1643,7 @@ async def _update_turn_usage_internal( def _update_sync(): """Synchronous helper to update turn usage data.""" - with self._locked_connection() as conn: + with self._write_connection() as conn: if turn_anchor is not None: with closing(conn.cursor()) as guard_cursor: guard_cursor.execute( diff --git a/src/agents/extensions/memory/async_sqlite_session.py b/src/agents/extensions/memory/async_sqlite_session.py index 06d0cc1755..d9b72f4193 100644 --- a/src/agents/extensions/memory/async_sqlite_session.py +++ b/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,7 +293,7 @@ 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. @@ -239,7 +301,8 @@ async def pop_item(self) -> TResponseInputItem | None: 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 diff --git a/src/agents/extensions/memory/dapr_session.py b/src/agents/extensions/memory/dapr_session.py index 20b300ce3c..ca72174368 100644 --- a/src/agents/extensions/memory/dapr_session.py +++ b/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. diff --git a/src/agents/extensions/memory/mongodb_session.py b/src/agents/extensions/memory/mongodb_session.py index b2ba601ab0..b10820267c 100644 --- a/src/agents/extensions/memory/mongodb_session.py +++ b/src/agents/extensions/memory/mongodb_session.py @@ -73,19 +73,22 @@ class MongoDBSession(SessionABC): """MongoDB implementation of [`Session`][agents.memory.session.Session]. - Conversation items are stored as individual documents in a ``messages`` - collection. A lightweight ``sessions`` collection tracks metadata - (creation time, last-updated time) for each session. + Conversation items are stored as logical-batch documents in a ``messages`` + collection. Legacy per-item documents remain readable. A lightweight + ``sessions`` collection tracks metadata (creation time, last-updated time) + for each session. Each logical batch must fit within MongoDB's single-document + size limit; an oversized batch fails atomically without storing a partial batch. Indexes are created once per ``(client, database, sessions_collection, messages_collection)`` combination on the first call to any of the session protocol methods. Subsequent calls skip the setup entirely. - Each message document carries a ``seq`` field — an integer assigned by - atomically incrementing a counter on the session metadata document. This - guarantees a strictly monotonic insertion order that is safe across - multiple writers and processes, unlike sorting by ``_id`` / ObjectId which - is only second-level accurate and non-monotonic across machines. + Each message document carries a ``seq`` field for the final item in that + document. Sequence ranges are assigned by atomically incrementing a counter + on the session metadata document. This guarantees a strictly monotonic + insertion order that is safe across multiple writers and processes, unlike + sorting by ``_id`` / ObjectId which is only second-level accurate and + non-monotonic across machines. """ # Class-level registry so index creation runs only once per unique @@ -125,8 +128,9 @@ def __init__( Defaults to ``"agents"``. sessions_collection: Name of the collection that stores session metadata. Defaults to ``"agent_sessions"``. - messages_collection: Name of the collection that stores individual - conversation items. Defaults to ``"agent_messages"``. + messages_collection: Name of the collection that stores logical + conversation batches and legacy per-item records. Defaults to + ``"agent_messages"``. session_settings: Optional session configuration. When ``None`` a default [`SessionSettings`][agents.memory.session_settings.SessionSettings] is used (no item limit). @@ -292,11 +296,14 @@ async def get_items(self, limit: int | None = None) -> list[TResponseInputItem]: async def _decode_docs(docs: list[Any]) -> list[TResponseInputItem]: items: list[TResponseInputItem] = [] for doc in docs: - try: - items.append(await self._deserialize_item(doc["message_data"])) - except (json.JSONDecodeError, KeyError, TypeError): - # Skip corrupted or malformed documents (including non-string BSON values). - continue + raw = doc.get("message_data") + raw_items = raw if isinstance(raw, list) else [raw] + for raw_item in raw_items: + try: + items.append(await self._deserialize_item(raw_item)) + except (json.JSONDecodeError, TypeError): + # Skip corrupted or malformed entries, including legacy non-string values. + continue return items if session_limit is None: @@ -333,6 +340,7 @@ async def add_items(self, items: list[TResponseInputItem]) -> None: await self._ensure_indexes() + serialized_items = [await self._serialize_item(item) for item in items] now = datetime.now(timezone.utc) # Atomically reserve a block of sequence numbers for this batch. @@ -350,16 +358,15 @@ async def add_items(self, items: list[TResponseInputItem]) -> None: ) next_seq: int = (result["_seq"] if result else len(items)) - len(items) - payload = [ + # One document is the commit boundary for the logical batch. This keeps + # standalone MongoDB deployments failure-atomic without requiring transactions. + await self._messages.insert_one( { "session_id": self.session_id, - "seq": next_seq + i, - "message_data": await self._serialize_item(item), + "seq": next_seq + len(serialized_items) - 1, + "message_data": serialized_items, } - for i, item in enumerate(items) - ] - - await self._messages.insert_many(payload, ordered=True) + ) async def pop_item(self) -> TResponseInputItem | None: """Remove and return the most recent item from the session. @@ -367,23 +374,52 @@ async def pop_item(self) -> TResponseInputItem | None: Returns: The most recent item if it exists, ``None`` if the session is empty. - Corrupt documents (invalid JSON, missing/non-string ``message_data``) - are silently discarded and the next-most-recent item is returned. This - matches :meth:`get_items`, which also skips corrupt documents, so a - single bad row cannot make a non-empty session look empty to callers. + New list-valued logical batches and legacy string-valued items are both + supported. Malformed entries (invalid JSON, missing ``message_data``, or + other value types) are silently discarded and the next-most-recent item + is returned. This matches :meth:`get_items`, which also skips malformed + entries, so one bad record cannot make a non-empty session look empty. """ await self._ensure_indexes() while True: - doc = await self._messages.find_one_and_delete( - {"session_id": self.session_id}, - sort=[("seq", -1)], + docs = await ( + self._messages.find( + { + "session_id": self.session_id, + "message_data": {"$ne": []}, + } + ) + .sort("seq", -1) + .limit(1) + .to_list() ) - if doc is None: + if not docs: return None + candidate = docs[0] + raw = candidate.get("message_data") + + if isinstance(raw, list): + doc = await self._messages.find_one_and_update( + {"_id": candidate["_id"], "message_data": {"$ne": []}}, + {"$pop": {"message_data": 1}}, + return_document=False, + ) + if doc is None: + continue + claimed_batch = doc.get("message_data") + if not isinstance(claimed_batch, list) or not claimed_batch: + continue + raw = claimed_batch[-1] + else: + doc = await self._messages.find_one_and_delete({"_id": candidate["_id"]}) + if doc is None: + continue + raw = doc.get("message_data") + try: - return await self._deserialize_item(doc["message_data"]) - except (json.JSONDecodeError, KeyError, TypeError): + return await self._deserialize_item(raw) + except (json.JSONDecodeError, TypeError): # Corrupt — drop it and try the next-most-recent document. continue diff --git a/src/agents/extensions/memory/redis_session.py b/src/agents/extensions/memory/redis_session.py index de9efbf6c8..cbf5874c39 100644 --- a/src/agents/extensions/memory/redis_session.py +++ b/src/agents/extensions/memory/redis_session.py @@ -31,6 +31,7 @@ try: import redis.asyncio as redis from redis.asyncio import Redis + from redis.exceptions import ResponseError, WatchError except ImportError as e: raise_optional_dependency_error( "RedisSession", @@ -48,6 +49,62 @@ ) +async def _reset_pipeline(pipe: Any) -> asyncio.CancelledError | None: + """Finish releasing a watched pipeline despite repeated caller cancellation.""" + reset_task = asyncio.create_task(pipe.reset()) + cancellation: asyncio.CancelledError | None = None + while True: + try: + await asyncio.shield(reset_task) + except asyncio.CancelledError as exc: + if cancellation is None: + cancellation = exc + if reset_task.done(): + reset_task.result() + return cancellation + else: + return cancellation + + +async def _execute_pipeline( + pipe: Any, +) -> tuple[BaseException | None, asyncio.CancelledError | None, bool]: + """Wait for transaction outcome and driver cleanup despite caller cancellation.""" + execute_transaction = pipe._execute_transaction + transaction_succeeded = False + + async def execute_transaction_and_mark(*args: Any, **kwargs: Any) -> Any: + nonlocal transaction_succeeded + result = await execute_transaction(*args, **kwargs) + transaction_succeeded = True + return result + + pipe._execute_transaction = execute_transaction_and_mark + execute_task = asyncio.create_task(pipe.execute()) + cancellation: asyncio.CancelledError | None = None + try: + while True: + try: + await asyncio.shield(execute_task) + except asyncio.CancelledError as exc: + if cancellation is None: + cancellation = exc + if not execute_task.done(): + continue + except BaseException as exc: + return exc, cancellation, transaction_succeeded + else: + return None, cancellation, transaction_succeeded + + try: + execute_task.result() + except BaseException as exc: + return exc, cancellation, transaction_succeeded + return None, cancellation, transaction_succeeded + finally: + pipe._execute_transaction = execute_transaction + + class RedisSession(SessionABC): """Redis implementation of [`Session`][agents.memory.session.Session].""" @@ -70,7 +127,8 @@ def __init__( key_prefix (str, optional): Prefix for Redis keys to avoid collisions. Defaults to "agents:session". ttl (int | None, optional): Time-to-live in seconds for session data. - If None, data persists indefinitely. Defaults to None. + If None, data persists indefinitely. Values outside Redis's supported expiration + range raise ValueError when adding items. Defaults to None. session_settings (SessionSettings | None): Session configuration settings including default limit for retrieving items. If None, uses default SessionSettings(). """ @@ -143,14 +201,6 @@ async def _get_next_id(self) -> int: result = await self._redis.incr(self._counter_key) return int(result) - async def _set_ttl_if_configured(self, *keys: str) -> None: - """Set TTL on keys if configured.""" - if self._ttl is not None: - pipe = self._redis.pipeline() - for key in keys: - pipe.expire(key, self._ttl) - await pipe.execute() - # ------------------------------------------------------------------ # Session protocol implementation # ------------------------------------------------------------------ @@ -160,6 +210,81 @@ def _check_not_closed(self) -> None: if self._closed: raise RuntimeError("RedisSession is closed") + @staticmethod + def _key_type_name(key_type: Any) -> str: + """Normalize Redis TYPE responses from bytes and decoded clients.""" + if isinstance(key_type, bytes): + return key_type.decode("utf-8") + return str(key_type) + + async def _write_items( + self, + serialized_items: list[str], + ) -> None: + """Validate key types and atomically write one batch with optimistic locking.""" + keys = (self._session_key, self._messages_key, self._counter_key) + while True: + pipe = self._redis.pipeline() + committed = False + try: + await pipe.watch(*keys) + session_key_type = self._key_type_name(await pipe.type(self._session_key)) + messages_key_type = self._key_type_name(await pipe.type(self._messages_key)) + if session_key_type not in ("none", "hash"): + raise ResponseError("WRONGTYPE session metadata key must contain a hash") + if messages_key_type not in ("none", "list"): + raise ResponseError("WRONGTYPE session messages key must contain a list") + + if self._ttl is None: + now = str(int(time.time())) + expiration_time_ms = None + else: + server_seconds, server_microseconds = await pipe.time() + now = str(int(server_seconds)) + expiration_time_ms = ( + int(server_seconds) * 1000 + + int(server_microseconds) // 1000 + + self._ttl * 1000 + ) + min_int64 = -(2**63) + max_int64 = 2**63 - 1 + if not min_int64 <= expiration_time_ms <= max_int64: + raise ValueError("ttl is outside Redis's supported expiration range") + + pipe.multi() + pipe.hset(self._session_key, "session_id", self.session_id) + pipe.hsetnx(self._session_key, "created_at", now) + pipe.rpush(self._messages_key, *serialized_items) + pipe.hset(self._session_key, "updated_at", now) + if expiration_time_ms is not None: + for key in keys: + pipe.pexpireat(key, expiration_time_ms) + ( + execute_error, + execute_cancellation, + transaction_succeeded, + ) = await _execute_pipeline(pipe) + if transaction_succeeded: + committed = True + return + if isinstance(execute_error, WatchError): + if execute_cancellation is not None: + raise execute_cancellation + continue + if execute_error is not None: + raise execute_error + committed = True + return + finally: + try: + reset_cancellation = await _reset_pipeline(pipe) + except BaseException: + if not committed: + raise + else: + if reset_cancellation is not None and not committed: + raise reset_cancellation + async def get_items(self, limit: int | None = None) -> list[TResponseInputItem]: """Retrieve the conversation history for this session. @@ -224,32 +349,12 @@ async def add_items(self, items: list[TResponseInputItem]) -> None: async with self._lock: self._check_not_closed() - pipe = self._redis.pipeline() - now = str(int(time.time())) - - # Set session metadata, preserving created_at across subsequent writes. - pipe.hset(self._session_key, "session_id", self.session_id) - pipe.hsetnx(self._session_key, "created_at", now) - - # Add all items to the messages list serialized_items = [] for item in items: serialized = await self._serialize_item(item) serialized_items.append(serialized) - if serialized_items: - pipe.rpush(self._messages_key, *serialized_items) - - # Update the session timestamp - pipe.hset(self._session_key, "updated_at", now) - - # Execute all commands - await pipe.execute() - - # Set TTL if configured - await self._set_ttl_if_configured( - self._session_key, self._messages_key, self._counter_key - ) + await self._write_items(serialized_items) async def pop_item(self) -> TResponseInputItem | None: """Remove and return the most recent item from the session. diff --git a/src/agents/extensions/memory/sqlalchemy_session.py b/src/agents/extensions/memory/sqlalchemy_session.py index 8751cca68b..713b5986a2 100644 --- a/src/agents/extensions/memory/sqlalchemy_session.py +++ b/src/agents/extensions/memory/sqlalchemy_session.py @@ -421,9 +421,9 @@ async def pop_item(self) -> TResponseInputItem | None: The most recent item if it exists, None if the session is empty """ await self._ensure_tables() - async with self._session_factory() as sess: - async with sess.begin(): - while True: + while True: + async with self._session_factory() as sess: + async with sess.begin(): # Fallback for all dialects - get ID first, then delete subq = ( select(self._messages.c.id) @@ -443,7 +443,13 @@ async def pop_item(self) -> TResponseInputItem | None: select(self._messages.c.message_data).where(self._messages.c.id == row_id) ) row = res_data.scalar_one_or_none() - await sess.execute(delete(self._messages).where(self._messages.c.id == row_id)) + delete_result = await sess.execute( + delete(self._messages).where(self._messages.c.id == row_id) + ) + if delete_result.rowcount == 0: + # A concurrent caller claimed this row. Start a new transaction so + # dialects with repeatable-read snapshots can observe the next row. + continue if row is None: continue diff --git a/src/agents/memory/sqlite_session.py b/src/agents/memory/sqlite_session.py index 61bf4e563b..a84ac59452 100644 --- a/src/agents/memory/sqlite_session.py +++ b/src/agents/memory/sqlite_session.py @@ -4,15 +4,41 @@ import json import sqlite3 import threading -from collections.abc import Iterator +from collections.abc import Awaitable, Iterator from contextlib import contextmanager from pathlib import Path -from typing import Any, ClassVar +from typing import Any, ClassVar, TypeVar from ..items import TResponseInputItem from .session import SessionABC from .session_settings import SessionSettings, coerce_session_settings, resolve_session_limit +_T = TypeVar("_T") + + +async def _await_mutation(awaitable: Awaitable[_T]) -> _T: + """Wait for a mutation outcome despite repeated caller cancellation.""" + task = asyncio.ensure_future(awaitable) + cancellation: asyncio.CancelledError | None = None + while True: + try: + return await asyncio.shield(task) + except asyncio.CancelledError as exc: + if cancellation is None: + cancellation = exc + if not task.done(): + continue + except BaseException: + if cancellation is not None: + raise cancellation from None + raise + + try: + return task.result() + except BaseException: + assert cancellation is not None + raise cancellation from None + class SQLiteSession(SessionABC): """SQLite-based implementation of session storage. @@ -57,6 +83,7 @@ def __init__( self.messages_table = messages_table self._local = threading.local() self._connections: set[sqlite3.Connection] = set() + self._quarantined_connections: set[sqlite3.Connection] = set() self._connections_lock = threading.Lock() self._closed = False @@ -125,6 +152,39 @@ def _check_not_closed(self) -> None: if self._closed: raise RuntimeError("SQLiteSession is closed") + @contextmanager + def _write_connection(self) -> Iterator[sqlite3.Connection]: + """Provide a connection that cannot retain a failed write transaction.""" + with self._locked_connection() as conn: + try: + yield conn + except BaseException: + try: + conn.rollback() + except BaseException: + self._invalidate_connection(conn) + raise + + def _invalidate_connection(self, conn: sqlite3.Connection) -> None: + """Close and evict a connection that could not roll back safely.""" + try: + conn.close() + except BaseException: + close_failed = True + else: + close_failed = False + + with self._connections_lock: + self._connections.discard(conn) + if close_failed: + self._quarantined_connections.add(conn) + else: + self._quarantined_connections.discard(conn) + if getattr(self._local, "connection", None) is conn: + del self._local.connection + if self._is_memory_db or close_failed: + self._closed = True + def _get_connection(self) -> sqlite3.Connection: """Get a database connection.""" self._check_not_closed() @@ -294,20 +354,11 @@ async def add_items(self, items: list[TResponseInputItem]) -> None: return def _add_items_sync(): - with self._locked_connection() as conn: - try: - self._insert_items(conn, items) - conn.commit() - except Exception: - # _locked_connection() does not manage transactions; roll back - # explicitly so a failure partway through the insert never leaves a - # partial mutation or an open transaction on this cached connection. - # An open write transaction would hold the SQLite write lock for the - # lifetime of the connection and block every later writer. - conn.rollback() - raise + with self._write_connection() as conn: + self._insert_items(conn, items) + conn.commit() - await asyncio.to_thread(_add_items_sync) + await _await_mutation(asyncio.to_thread(_add_items_sync)) async def pop_item(self) -> TResponseInputItem | None: """Remove and return the most recent item from the session. @@ -317,7 +368,7 @@ async def pop_item(self) -> TResponseInputItem | None: """ def _pop_item_sync(): - with self._locked_connection() as conn: + with self._write_connection() as conn: # Use DELETE with RETURNING to atomically delete and return the most recent item cursor = conn.execute( f""" @@ -361,13 +412,13 @@ def _pop_item_sync(): return None - return await asyncio.to_thread(_pop_item_sync) + return await _await_mutation(asyncio.to_thread(_pop_item_sync)) async def clear_session(self) -> None: """Clear all items for this session.""" def _clear_session_sync(): - with self._locked_connection() as conn: + with self._write_connection() as conn: conn.execute( f"DELETE FROM {self.messages_table} WHERE session_id = ?", (self.session_id,), @@ -378,24 +429,44 @@ def _clear_session_sync(): ) conn.commit() - await asyncio.to_thread(_clear_session_sync) + await _await_mutation(asyncio.to_thread(_clear_session_sync)) def close(self) -> None: """Close the database connection.""" with self._lock: - if self._closed: - return - self._closed = True + with self._connections_lock: + connections = self._connections | self._quarantined_connections if self._is_memory_db: if hasattr(self, "_shared_connection"): - self._shared_connection.close() - else: + connections.add(self._shared_connection) + + first_error: BaseException | None = None + for connection in connections: + try: + connection.close() + except BaseException as exc: + if first_error is None: + first_error = exc + with self._connections_lock: + self._connections.discard(connection) + self._quarantined_connections.add(connection) + else: + with self._connections_lock: + self._connections.discard(connection) + self._quarantined_connections.discard(connection) + + if getattr(self._local, "connection", None) in connections: + del self._local.connection + + with self._connections_lock: + has_unclosed_connections = bool(self._quarantined_connections) + if not has_unclosed_connections and self._lock_path is not None: with self._connections_lock: - connections = list(self._connections) self._connections.clear() - for connection in connections: - connection.close() - if self._lock_path is not None and not self._lock_released: - self._release_file_lock(self._lock_path) - self._lock_released = True + if not self._lock_released: + self._release_file_lock(self._lock_path) + self._lock_released = True + + if first_error is not None: + raise first_error diff --git a/tests/extensions/memory/test_advanced_sqlite_session.py b/tests/extensions/memory/test_advanced_sqlite_session.py index ef1e72e16a..6eea7f805a 100644 --- a/tests/extensions/memory/test_advanced_sqlite_session.py +++ b/tests/extensions/memory/test_advanced_sqlite_session.py @@ -5,6 +5,7 @@ import json import logging import multiprocessing +import sqlite3 import tempfile import threading import time @@ -373,6 +374,193 @@ async def test_add_items_rolls_back_partial_structure_metadata_write(): session.close() +async def test_add_items_rollback_failure_invalidates_connection( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +): + """Advanced add failures must use the base rollback-failure invalidation path.""" + + class FailingRollbackConnection(sqlite3.Connection): + def rollback(self) -> None: + raise RuntimeError("rollback failed") + + db_path = tmp_path / "advanced_rollback_failure.db" + session = AdvancedSQLiteSession( + session_id="advanced_rollback_failure", + db_path=db_path, + create_tables=True, + ) + conn = sqlite3.connect( + str(db_path), + check_same_thread=False, + factory=FailingRollbackConnection, + ) + with session._connections_lock: + session._connections.add(conn) + real_get_connection = session._get_connection + real_insert_structure_metadata = session._insert_structure_metadata + monkeypatch.setattr(session, "_get_connection", lambda: conn) + + def fail_structure_metadata(*_args: Any) -> None: + raise RuntimeError("structure metadata failed") + + monkeypatch.setattr(session, "_insert_structure_metadata", fail_structure_metadata) + + with pytest.raises(RuntimeError, match="structure metadata failed"): + await session.add_items([{"role": "user", "content": "not saved"}]) + + assert conn not in session._connections + probe = sqlite3.connect(str(db_path), timeout=0) + try: + probe.execute("CREATE TABLE IF NOT EXISTS probe_lock (x INTEGER)") + probe.commit() + finally: + probe.close() + + monkeypatch.setattr(session, "_get_connection", real_get_connection) + monkeypatch.setattr(session, "_insert_structure_metadata", real_insert_structure_metadata) + await session.add_items([{"role": "user", "content": "after failure"}]) + assert await session.get_items() == [{"role": "user", "content": "after failure"}] + session.close() + + +async def test_structure_initialization_failure_invalidates_connection( + tmp_path: Path, +): + """Initialization must release its write lock even when rollback also fails.""" + + class FailingRollbackConnection(sqlite3.Connection): + def rollback(self) -> None: + raise RuntimeError("rollback failed") + + captured_connections: list[sqlite3.Connection] = [] + + class FailingRollbackInitSession(AdvancedSQLiteSession): + def _get_connection(self) -> sqlite3.Connection: + if not hasattr(self, "_test_connection"): + connection = sqlite3.connect( + str(self.db_path), + check_same_thread=False, + factory=FailingRollbackConnection, + ) + self._test_connection = connection + captured_connections.append(connection) + with self._connections_lock: + self._connections.add(connection) + return self._test_connection + + db_path = tmp_path / "advanced_init_failure.db" + setup = AdvancedSQLiteSession( + session_id="advanced_init_failure", + db_path=db_path, + create_tables=True, + ) + try: + await setup.add_items([{"role": "user", "content": "existing"}]) + finally: + setup.close() + + conflict = sqlite3.connect(str(db_path)) + try: + conflict.execute("DROP TABLE branch_reservations") + conflict.execute("DROP INDEX idx_structure_session_seq") + conflict.execute("CREATE TABLE idx_structure_session_seq (value INTEGER)") + conflict.commit() + finally: + conflict.close() + + with pytest.raises(sqlite3.OperationalError, match="already a table"): + FailingRollbackInitSession( + session_id="advanced_init_failure", + db_path=db_path, + create_tables=True, + ) + + assert len(captured_connections) == 1 + with pytest.raises(sqlite3.ProgrammingError): + captured_connections[0].execute("SELECT 1") + + probe = sqlite3.connect(str(db_path), timeout=0) + try: + probe.execute("CREATE TABLE IF NOT EXISTS probe_lock (x INTEGER)") + probe.commit() + finally: + probe.close() + + +@pytest.mark.parametrize("operation", ["add", "pop", "clear"]) +async def test_post_commit_cancellation_returns_known_mutation_outcome( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + operation: str, +): + """Cancellation after commit must not invite an AdvancedSQLite mutation retry.""" + + class PausingCommitConnection(sqlite3.Connection): + pause_commit = False + commit_finished = threading.Event() + allow_return = threading.Event() + + def commit(self) -> None: + super().commit() + if self.pause_commit: + self.pause_commit = False + self.commit_finished.set() + assert self.allow_return.wait(timeout=10) + + db_path = tmp_path / f"advanced_post_commit_{operation}.db" + session = AdvancedSQLiteSession( + session_id=f"advanced_post_commit_{operation}", + db_path=db_path, + create_tables=True, + ) + item: TResponseInputItem = {"role": "user", "content": "once"} + if operation != "add": + await session.add_items([item]) + + conn = sqlite3.connect( + str(db_path), + check_same_thread=False, + factory=PausingCommitConnection, + ) + with session._connections_lock: + session._connections.add(conn) + monkeypatch.setattr(session, "_get_connection", lambda: conn) + conn.pause_commit = True + + if operation == "add": + mutation: asyncio.Task[Any] = asyncio.create_task(session.add_items([item])) + elif operation == "pop": + mutation = asyncio.create_task(session.pop_item()) + else: + mutation = asyncio.create_task(session.clear_session()) + + try: + assert await asyncio.to_thread(conn.commit_finished.wait, 10) + mutation.cancel() + await asyncio.sleep(0) + mutation.cancel() + await asyncio.sleep(0) + conn.allow_return.set() + result = await mutation + finally: + conn.allow_return.set() + if not mutation.done(): + mutation.cancel() + await asyncio.gather(mutation, return_exceptions=True) + + if operation == "add": + assert result is None + assert await session.get_items() == [item] + elif operation == "pop": + assert result == item + assert await session.get_items() == [] + else: + assert result is None + assert await session.get_items() == [] + session.close() + + async def test_message_structure_tracking(agent: Agent): """Test that message structure is properly tracked.""" session_id = "structure_test" @@ -751,6 +939,30 @@ def _reserve_branch_id( session.close() +def _pop_item_in_process( + db_path: str, + session_id: str, + ready: Any, + start: Any, + results: Any, +) -> None: + """Pop one AdvancedSQLite item in a separately synchronized process.""" + session = AdvancedSQLiteSession( + session_id=session_id, + db_path=db_path, + create_tables=False, + ) + try: + ready.set() + if not start.wait(timeout=10): + raise TimeoutError("Timed out waiting to start pop") + results.put(("ok", asyncio.run(session.pop_item()))) + except Exception as exc: + results.put(("error", type(exc).__name__, str(exc))) + finally: + session.close() + + @pytest.mark.parametrize("branch_id", ["main", "existing_branch"]) async def test_create_branch_rejects_populated_branch_id(branch_id: str): """Creating a branch must not append history to a populated branch.""" @@ -1675,6 +1887,46 @@ async def test_usage_tracking_storage(agent: Agent, usage_data: Usage): session.close() +async def test_failed_usage_write_rolls_back_cached_connection(usage_data: Usage): + """A swallowed usage-write failure must not strand a transaction or SQLite lock.""" + with tempfile.TemporaryDirectory() as temp_dir: + db_path = Path(temp_dir) / "usage_rollback.db" + session = AdvancedSQLiteSession( + session_id="usage_rollback", + db_path=db_path, + create_tables=True, + ) + await session.add_items([{"role": "user", "content": "turn"}]) + + helper = session._get_connection() + helper.execute( + """ + CREATE TRIGGER fail_turn_usage + BEFORE INSERT ON turn_usage + BEGIN + SELECT RAISE(ABORT, 'usage write failed'); + END + """ + ) + helper.commit() + + await session.store_run_usage(create_mock_run_result(usage_data)) + + assert all(not conn.in_transaction for conn in session._connections) + probe = sqlite3.connect(str(db_path), timeout=0) + try: + probe.execute("CREATE TABLE usage_lock_probe (x INTEGER)") + probe.commit() + finally: + probe.close() + + helper.execute("DROP TRIGGER fail_turn_usage") + helper.commit() + await session.store_run_usage(create_mock_run_result(usage_data)) + assert await session.get_turn_usage(1) + session.close() + + async def test_runner_integration_with_usage_tracking(agent: Agent): """Test integration with Runner and automatic usage tracking pattern.""" session_id = "integration_test" @@ -2716,6 +2968,51 @@ async def test_pop_item_uses_branch_snapshot_when_branch_switches_concurrently() session.close() +async def test_pop_item_claim_is_unique_across_processes(tmp_path: Path): + """Two processes must not return the same destructively read item.""" + db_path = tmp_path / "advanced_pop_processes.db" + session_id = "advanced_pop_processes" + item: TResponseInputItem = {"role": "user", "content": "only"} + setup = AdvancedSQLiteSession(session_id=session_id, db_path=db_path, create_tables=True) + await setup.add_items([item]) + setup.close() + + context = multiprocessing.get_context("spawn") + start = context.Event() + results = context.Queue() + ready_events = [context.Event(), context.Event()] + processes = [ + context.Process( + target=_pop_item_in_process, + args=(str(db_path), session_id, ready, start, results), + ) + for ready in ready_events + ] + + try: + for process in processes: + process.start() + for ready in ready_events: + assert ready.wait(timeout=10) + start.set() + for process in processes: + process.join(timeout=10) + assert process.exitcode == 0 + + outcomes = [results.get(timeout=5), results.get(timeout=5)] + assert all(outcome[0] == "ok" for outcome in outcomes) + popped_items = [outcome[1] for outcome in outcomes] + assert popped_items.count(item) == 1 + assert popped_items.count(None) == 1 + finally: + start.set() + for process in processes: + if process.is_alive(): + process.terminate() + process.join(timeout=5) + results.close() + + async def test_stale_switch_after_clear_does_not_repoint_to_deleted_branch(): """A switch_to_branch that commits its pointer after clear_session must not resurrect the deleted branch; the generation guard makes it a no-op. diff --git a/tests/extensions/memory/test_async_sqlite_session.py b/tests/extensions/memory/test_async_sqlite_session.py index b45cbdf4e7..6d843ab8e4 100644 --- a/tests/extensions/memory/test_async_sqlite_session.py +++ b/tests/extensions/memory/test_async_sqlite_session.py @@ -2,7 +2,9 @@ from __future__ import annotations +import asyncio import json +import sqlite3 import tempfile from collections.abc import Sequence from datetime import datetime @@ -495,3 +497,432 @@ async def test_async_sqlite_session_close_is_idempotent(): with pytest.raises(RuntimeError, match="AsyncSQLiteSession is closed"): await session.get_items() + + +async def test_cancelled_close_finishes_cleanup_and_propagates_cancellation( + monkeypatch: pytest.MonkeyPatch, +): + """Repeated cancellation must propagate after the owned connection closes.""" + close_started = asyncio.Event() + allow_close = asyncio.Event() + session = AsyncSQLiteSession("cancelled_close") + conn = await session._get_connection() + real_close = conn.close + + async def controlled_close() -> None: + close_started.set() + await allow_close.wait() + await real_close() + + monkeypatch.setattr(conn, "close", controlled_close) + close_task = asyncio.create_task(session.close()) + try: + await close_started.wait() + close_task.cancel() + await asyncio.sleep(0) + close_task.cancel() + await asyncio.sleep(0) + allow_close.set() + with pytest.raises(asyncio.CancelledError): + await close_task + finally: + allow_close.set() + if not close_task.done(): + close_task.cancel() + await asyncio.gather(close_task, return_exceptions=True) + + assert session._closed is True + assert session._connection is None + assert session._quarantined_connections == set() + assert conn._running is False + + +@pytest.mark.parametrize("operation", ["add", "pop", "clear"]) +async def test_post_commit_cancellation_returns_known_mutation_outcome( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + operation: str, +): + """Cancellation after async commit must not invite a duplicate mutation retry.""" + db_path = tmp_path / f"async_post_commit_{operation}.db" + session = AsyncSQLiteSession(f"async_post_commit_{operation}", db_path) + item: TResponseInputItem = {"role": "user", "content": "once"} + if operation != "add": + await session.add_items([item]) + + conn = await session._get_connection() + real_commit = conn.commit + commit_finished = asyncio.Event() + allow_return = asyncio.Event() + pause_commit = True + + async def controlled_commit() -> None: + nonlocal pause_commit + await real_commit() + if pause_commit: + pause_commit = False + commit_finished.set() + await allow_return.wait() + + monkeypatch.setattr(conn, "commit", controlled_commit) + if operation == "add": + mutation: asyncio.Task[Any] = asyncio.create_task(session.add_items([item])) + elif operation == "pop": + mutation = asyncio.create_task(session.pop_item()) + else: + mutation = asyncio.create_task(session.clear_session()) + + try: + await commit_finished.wait() + mutation.cancel() + await asyncio.sleep(0) + mutation.cancel() + await asyncio.sleep(0) + allow_return.set() + result = await mutation + finally: + allow_return.set() + if not mutation.done(): + mutation.cancel() + await asyncio.gather(mutation, return_exceptions=True) + + if operation == "add": + assert result is None + assert await session.get_items() == [item] + elif operation == "pop": + assert result == item + assert await session.get_items() == [] + else: + assert result is None + assert await session.get_items() == [] + await session.close() + + +def _drop_sqlite_table(db_path: Path, table: str) -> None: + """Drop a table from an independent connection to make a later statement fail.""" + helper = sqlite3.connect(str(db_path)) + try: + helper.execute(f"DROP TABLE {table}") + helper.commit() + finally: + helper.close() + + +def _sqlite_write_lock_is_free(db_path: Path) -> bool: + """Return whether an independent writer can take the SQLite write lock.""" + probe = sqlite3.connect(str(db_path), timeout=0) + try: + probe.execute("CREATE TABLE IF NOT EXISTS probe_lock (x INTEGER)") + probe.commit() + return True + except sqlite3.OperationalError: + return False + finally: + probe.close() + + +async def test_failed_add_items_rolls_back_and_reuses_connection(): + """A failed add must roll back its partial write and leave the session reusable.""" + with tempfile.TemporaryDirectory() as temp_dir: + db_path = Path(temp_dir) / "add_rollback.db" + session = AsyncSQLiteSession("add_rollback", db_path) + unserializable = cast(TResponseInputItem, {"role": "user", "content": object()}) + + with pytest.raises(TypeError): + await session.add_items([unserializable]) + + conn = await session._get_connection() + assert conn.in_transaction is False + assert _sqlite_write_lock_is_free(db_path) + + await session.add_items([{"role": "user", "content": "after failure"}]) + assert [item.get("content") for item in await session.get_items()] == ["after failure"] + await session.close() + + +async def test_failed_clear_session_rolls_back(): + """A failed clear must restore earlier statements and release the shared write lock.""" + with tempfile.TemporaryDirectory() as temp_dir: + db_path = Path(temp_dir) / "clear_rollback.db" + session = AsyncSQLiteSession("clear_rollback", db_path) + await session.add_items([{"role": "user", "content": "kept"}]) + + _drop_sqlite_table(db_path, "agent_sessions") + + with pytest.raises(sqlite3.OperationalError): + await session.clear_session() + + conn = await session._get_connection() + assert conn.in_transaction is False + assert _sqlite_write_lock_is_free(db_path) + await session.close() + + +async def test_failed_pop_item_releases_write_lock(): + """A failed pop must not leave a write transaction on the shared connection.""" + with tempfile.TemporaryDirectory() as temp_dir: + db_path = Path(temp_dir) / "pop_rollback.db" + session = AsyncSQLiteSession("pop_rollback", db_path) + await session.add_items([{"role": "user", "content": "kept"}]) + + _drop_sqlite_table(db_path, "agent_messages") + + with pytest.raises(sqlite3.OperationalError): + await session.pop_item() + + conn = await session._get_connection() + assert conn.in_transaction is False + assert _sqlite_write_lock_is_free(db_path) + await session.close() + + +async def test_failed_initialization_closes_candidate_connection(): + """A failed initialization must not retain a half-initialized connection.""" + + class FailingInitSession(AsyncSQLiteSession): + captured_connection: Any = None + + async def _init_db_for_connection(self, conn: Any) -> None: + self.captured_connection = conn + raise RuntimeError("initialization failed") + + session = FailingInitSession("failed_init") + with pytest.raises(RuntimeError, match="initialization failed"): + await session.get_items() + + assert session._connection is None + assert session.captured_connection._running is False + + +async def test_cancelled_add_items_rolls_back_write_transaction( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +): + """Cancellation after the first write must release the transaction and database lock.""" + db_path = tmp_path / "cancelled_add.db" + session = AsyncSQLiteSession("cancelled_add", db_path) + await session.get_items() + conn = await session._get_connection() + real_execute = conn.execute + real_rollback = conn.rollback + first_write_started = asyncio.Event() + rollback_started = asyncio.Event() + allow_rollback = asyncio.Event() + + async def pause_after_first_write(*args: Any, **kwargs: Any) -> Any: + cursor = await real_execute(*args, **kwargs) + first_write_started.set() + await asyncio.Event().wait() + return cursor + + async def controlled_rollback() -> None: + rollback_started.set() + await allow_rollback.wait() + await real_rollback() + + monkeypatch.setattr(conn, "execute", pause_after_first_write) + monkeypatch.setattr(conn, "rollback", controlled_rollback) + task = asyncio.create_task(session.add_items([{"role": "user", "content": "cancelled"}])) + try: + await first_write_started.wait() + task.cancel() + await rollback_started.wait() + task.cancel() + allow_rollback.set() + with pytest.raises(asyncio.CancelledError): + await task + finally: + allow_rollback.set() + monkeypatch.setattr(conn, "execute", real_execute) + monkeypatch.setattr(conn, "rollback", real_rollback) + if not task.done(): + task.cancel() + await asyncio.gather(task, return_exceptions=True) + + assert conn.in_transaction is False + assert _sqlite_write_lock_is_free(db_path) + assert await session.get_items() == [] + await session.close() + + +async def test_rollback_failure_closes_and_evicts_connection( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +): + """A connection that cannot roll back must not remain cached or retain its write lock.""" + db_path = tmp_path / "rollback_failure.db" + session = AsyncSQLiteSession("rollback_failure", db_path) + await session.get_items() + conn = await session._get_connection() + + async def fail_rollback() -> None: + raise RuntimeError("rollback failed") + + monkeypatch.setattr(conn, "rollback", fail_rollback) + unserializable = cast(TResponseInputItem, {"role": "user", "content": object()}) + + with pytest.raises(TypeError): + await session.add_items([unserializable]) + + assert session._connection is None + assert session._closed is False + assert _sqlite_write_lock_is_free(db_path) + + await session.add_items([{"role": "user", "content": "after failure"}]) + assert [item.get("content") for item in await session.get_items()] == ["after failure"] + await session.close() + + +async def test_close_retries_connection_quarantined_after_rollback_failure( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +): + """A failed invalidation close must remain owned until a later close succeeds.""" + db_path = tmp_path / "close_retry.db" + session = AsyncSQLiteSession("close_retry", db_path) + await session.get_items() + conn = await session._get_connection() + real_close = conn.close + + async def fail_rollback() -> None: + raise RuntimeError("rollback failed") + + async def fail_close() -> None: + raise RuntimeError("close failed") + + monkeypatch.setattr(conn, "rollback", fail_rollback) + monkeypatch.setattr(conn, "close", fail_close) + unserializable = cast(TResponseInputItem, {"role": "user", "content": object()}) + + with pytest.raises(TypeError): + await session.add_items([unserializable]) + + assert session._closed is True + assert session._connection is None + assert conn in session._quarantined_connections + assert conn._running is True + assert _sqlite_write_lock_is_free(db_path) is False + + monkeypatch.setattr(conn, "close", real_close) + await session.close() + + assert session._quarantined_connections == set() + assert conn._running is False + assert _sqlite_write_lock_is_free(db_path) + + +async def test_close_retries_quarantined_failed_initialization_candidate( + monkeypatch: pytest.MonkeyPatch, +): + """A failed initialization candidate close must remain owned for close retry.""" + + class FailingInitSession(AsyncSQLiteSession): + captured_connection: Any = None + real_close: Any = None + + async def _init_db_for_connection(self, conn: Any) -> None: + self.captured_connection = conn + self.real_close = conn.close + + async def fail_close() -> None: + raise RuntimeError("close failed") + + monkeypatch.setattr(conn, "close", fail_close) + raise RuntimeError("initialization failed") + + session = FailingInitSession("failed_init_close_retry") + with pytest.raises(RuntimeError, match="initialization failed"): + await session.get_items() + + conn = session.captured_connection + assert session._closed is True + assert conn in session._quarantined_connections + assert conn._running is True + + monkeypatch.setattr(conn, "close", session.real_close) + await session.close() + + assert session._quarantined_connections == set() + assert conn._running is False + + +async def test_cancelled_initialization_finishes_candidate_close( + monkeypatch: pytest.MonkeyPatch, +): + """Repeated cancellation must not release initialization ownership before close finishes.""" + init_started = asyncio.Event() + close_started = asyncio.Event() + allow_close = asyncio.Event() + + class CancelledInitSession(AsyncSQLiteSession): + captured_connection: Any = None + + async def _init_db_for_connection(self, conn: Any) -> None: + self.captured_connection = conn + real_close = conn.close + + async def controlled_close() -> None: + close_started.set() + await allow_close.wait() + await real_close() + + monkeypatch.setattr(conn, "close", controlled_close) + init_started.set() + await asyncio.Event().wait() + + session = CancelledInitSession("cancelled_init") + task = asyncio.create_task(session.get_items()) + try: + await init_started.wait() + task.cancel() + await close_started.wait() + task.cancel() + allow_close.set() + with pytest.raises(asyncio.CancelledError): + await task + finally: + allow_close.set() + if not task.done(): + task.cancel() + await asyncio.gather(task, return_exceptions=True) + + assert session._connection is None + assert session.captured_connection._running is False + + +async def test_cancelled_connect_closes_eventually_acquired_connection( + monkeypatch: pytest.MonkeyPatch, +): + """Cancellation during connect must wait for and close the eventual connection.""" + import aiosqlite + + real_connect = aiosqlite.connect + connect_started = asyncio.Event() + allow_connect = asyncio.Event() + created_connections: list[Any] = [] + + async def controlled_connect(database: str) -> Any: + connect_started.set() + await allow_connect.wait() + conn = await real_connect(database) + created_connections.append(conn) + return conn + + monkeypatch.setattr(aiosqlite, "connect", controlled_connect) + session = AsyncSQLiteSession("cancelled_connect") + task = asyncio.create_task(session.get_items()) + try: + await connect_started.wait() + task.cancel() + allow_connect.set() + with pytest.raises(asyncio.CancelledError): + await task + finally: + allow_connect.set() + if not task.done(): + task.cancel() + await asyncio.gather(task, return_exceptions=True) + + assert len(created_connections) == 1 + assert created_connections[0]._running is False + assert session._connection is None diff --git a/tests/extensions/memory/test_dapr_session.py b/tests/extensions/memory/test_dapr_session.py index 702cb7388c..7cdde0f124 100644 --- a/tests/extensions/memory/test_dapr_session.py +++ b/tests/extensions/memory/test_dapr_session.py @@ -157,6 +157,35 @@ async def save_state( ) +class FailingMetadataFakeDaprClient(FakeDaprClient): + """Fake client that rejects the next metadata write.""" + + def __init__(self) -> None: + super().__init__() + self.fail_metadata_save = True + + async def save_state( + self, + store_name: str, + key: str, + value: str | bytes, + state_metadata: dict[str, str] | None = None, + options: Any = None, + etag: str | None = None, + ) -> None: + if key.endswith(":metadata") and self.fail_metadata_save: + self.fail_metadata_save = False + raise RuntimeError("metadata write failed") + await super().save_state( + store_name=store_name, + key=key, + value=value, + state_metadata=state_metadata, + options=options, + etag=etag, + ) + + @pytest.fixture def conflict_dapr_client() -> ConflictFakeDaprClient: """Fixture for fake client that forces concurrency conflicts.""" @@ -1327,3 +1356,22 @@ def entry_signalling_resolve(*args: Any, **kwargs: Any) -> Any: task.cancel() with suppress(asyncio.CancelledError, RuntimeError): await task + + +async def test_metadata_failure_does_not_commit_messages() -> None: + """An ancillary metadata failure must happen before conversation history is committed.""" + client = FailingMetadataFakeDaprClient() + session = DaprSession( + session_id="metadata_failure", + state_store_name="statestore", + dapr_client=client, # type: ignore[arg-type] + ) + item: TResponseInputItem = {"role": "user", "content": "once"} + + with pytest.raises(RuntimeError, match="metadata write failed"): + await session.add_items([item]) + + assert await session.get_items() == [] + + await session.add_items([item]) + assert await session.get_items() == [item] diff --git a/tests/extensions/memory/test_mongodb_session.py b/tests/extensions/memory/test_mongodb_session.py index 3bd8f7c034..3496b7fde6 100644 --- a/tests/extensions/memory/test_mongodb_session.py +++ b/tests/extensions/memory/test_mongodb_session.py @@ -9,11 +9,13 @@ from __future__ import annotations import asyncio +import copy +import json import sys import types from collections import defaultdict from datetime import datetime, timezone -from typing import Any +from typing import Any, cast from unittest.mock import AsyncMock, patch import pytest @@ -117,21 +119,38 @@ async def insert_many( doc["_id"] = FakeObjectId() self._docs[id(doc["_id"])] = dict(doc) + async def insert_one(self, document: dict[str, Any]) -> Any: + if "_id" not in document: + document["_id"] = FakeObjectId() + stored = dict(document) + if isinstance(stored.get("message_data"), list): + stored["message_data"] = list(stored["message_data"]) + self._docs[id(document["_id"])] = stored + async def find_one_and_update( self, query: dict[str, Any], update: dict[str, Any], upsert: bool = False, return_document: bool = False, + sort: list[tuple[str, int]] | None = None, ) -> dict[str, Any] | None: - for doc in self._docs.values(): - if self._matches(doc, query): - # Apply $inc fields. - for field, delta in update.get("$inc", {}).items(): - doc[field] = doc.get(field, 0) + delta - for field, value in update.get("$set", {}).items(): - doc[field] = value - return dict(doc) if return_document else None + matches = [doc for doc in self._docs.values() if self._matches(doc, query)] + if sort: + for field, direction in reversed(sort): + matches.sort(key=lambda doc: doc.get(field, 0), reverse=(direction == -1)) + if matches: + doc = matches[0] + before = copy.deepcopy(doc) + for field, delta in update.get("$inc", {}).items(): + doc[field] = doc.get(field, 0) + delta + for field, value in update.get("$set", {}).items(): + doc[field] = value + for field, direction in update.get("$pop", {}).items(): + values = doc.get(field) + if isinstance(values, list) and values: + values.pop(-1 if direction == 1 else 0) + return copy.deepcopy(doc) if return_document else before if upsert: new_doc: dict[str, Any] = {"_id": FakeObjectId()} new_doc.update(update.get("$setOnInsert", {})) @@ -169,7 +188,15 @@ async def delete_one(self, query: dict[str, Any]) -> None: @staticmethod def _matches(doc: dict[str, Any], query: dict[str, Any]) -> bool: - return all(doc.get(k) == v for k, v in query.items()) + for key, expected in query.items(): + actual = doc.get(key) + if isinstance(expected, dict): + if "$ne" in expected and actual == expected["$ne"]: + return False + continue + if actual != expected: + return False + return True class FakeAsyncDatabase: @@ -974,3 +1001,125 @@ def _fake_client(uri: str, **kwargs: Any) -> FakeAsyncMongoClient: # The caller's value must be preserved — setdefault must not overwrite it. assert captured_kwargs["driver"] is custom_info + + +async def test_add_items_serializes_before_reserving_sequence_numbers() -> None: + """Serialization failure must not advance the durable sequence counter.""" + + class FailingSerializationSession(MongoDBSession): + async def _serialize_item(self, item: TResponseInputItem) -> str: + if item.get("content") == "fail": + raise TypeError("serialization failed") + return await super()._serialize_item(item) + + MongoDBSession._init_state.clear() + client = FakeAsyncMongoClient() + session = FailingSerializationSession( + "serialize-first", + client=client, # type: ignore[arg-type] + database="agents_test", + ) + + with pytest.raises(TypeError, match="serialization failed"): + await session.add_items( + [ + {"role": "user", "content": "valid"}, + {"role": "assistant", "content": "fail"}, + ] + ) + + assert not session._sessions._docs + assert await session.get_items() == [] + + +async def test_add_items_is_invisible_until_single_document_commit( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """A failed batch must remain invisible while its single-document insert is pending.""" + MongoDBSession._init_state.clear() + client = FakeAsyncMongoClient() + session = MongoDBSession( + "partial-batch", + client=client, # type: ignore[arg-type] + database="agents_test", + ) + seed: TResponseInputItem = {"role": "user", "content": "seed"} + await session.add_items([seed]) + + real_insert_one = session._messages.insert_one + insert_started = asyncio.Event() + allow_failure = asyncio.Event() + + async def pause_then_fail(document: dict[str, Any]) -> None: + insert_started.set() + await allow_failure.wait() + raise RuntimeError("insert failed") + + monkeypatch.setattr(session._messages, "insert_one", pause_then_fail) + batch: list[TResponseInputItem] = [ + {"role": "assistant", "content": "first"}, + {"role": "user", "content": "second"}, + ] + add_task = asyncio.create_task(session.add_items(batch)) + + await insert_started.wait() + assert await session.get_items() == [seed] + allow_failure.set() + with pytest.raises(RuntimeError, match="insert failed"): + await add_task + + assert await session.get_items() == [seed] + + monkeypatch.setattr(session._messages, "insert_one", real_insert_one) + await session.add_items(batch) + assert await session.get_items() == [seed, *batch] + + +async def test_concurrent_pop_item_claims_distinct_items_from_batch() -> None: + """Atomic array pops must return each item in a logical batch at most once.""" + MongoDBSession._init_state.clear() + client = FakeAsyncMongoClient() + session = MongoDBSession( + "concurrent-pop", + client=client, # type: ignore[arg-type] + database="agents_test", + ) + items: list[TResponseInputItem] = [ + {"role": "user", "content": "first"}, + {"role": "assistant", "content": "second"}, + {"role": "user", "content": "third"}, + ] + await session.add_items(items) + + popped = await asyncio.gather(*(session.pop_item() for _ in items)) + + contents = {cast(dict[str, Any], item).get("content") for item in popped if item is not None} + assert contents == { + "first", + "second", + "third", + } + assert await session.get_items() == [] + + +async def test_reads_legacy_item_documents_with_new_batch_documents() -> None: + """New readers must preserve histories written by released per-item storage.""" + MongoDBSession._init_state.clear() + client = FakeAsyncMongoClient() + session = MongoDBSession( + "legacy-read", + client=client, # type: ignore[arg-type] + database="agents_test", + ) + await session.add_items([{"role": "assistant", "content": "new"}]) + legacy_doc = { + "_id": FakeObjectId(), + "session_id": session.session_id, + "seq": -1, + "message_data": json.dumps({"role": "user", "content": "legacy"}), + } + session._messages._docs[id(legacy_doc["_id"])] = legacy_doc + + assert [item.get("content") for item in await session.get_items()] == ["legacy", "new"] + assert (await session.pop_item() or {}).get("content") == "new" + assert (await session.pop_item() or {}).get("content") == "legacy" diff --git a/tests/extensions/memory/test_redis_session.py b/tests/extensions/memory/test_redis_session.py index e906387c0c..901019c145 100644 --- a/tests/extensions/memory/test_redis_session.py +++ b/tests/extensions/memory/test_redis_session.py @@ -1,5 +1,8 @@ from __future__ import annotations +import asyncio +import json +import time from typing import Any, cast import pytest @@ -1267,6 +1270,377 @@ async def test_redis_session_close_is_noop_for_injected_client(): await session.clear_session() +async def test_add_items_applies_ttl_in_the_write_transaction( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """TTL setup must not create a post-commit failure window for add_items.""" + if not USE_FAKE_REDIS: + pytest.skip("This test requires fakeredis for pipeline instrumentation") + + client = fakeredis.aioredis.FakeRedis() + session = RedisSession( + session_id="atomic_ttl", + redis_client=cast("Redis", client), + key_prefix="test:", + ttl=60, + ) + real_pipeline = client.pipeline + execute_calls = 0 + + def pipeline(*args: Any, **kwargs: Any) -> Any: + delegate = real_pipeline(*args, **kwargs) + + class PipelineProxy: + def __getattr__(self, name: str) -> Any: + return getattr(delegate, name) + + async def execute(self) -> Any: + nonlocal execute_calls + execute_calls += 1 + result = await delegate.execute() + if execute_calls == 2: + raise RuntimeError("post-commit TTL failure") + return result + + return PipelineProxy() + + monkeypatch.setattr(client, "pipeline", pipeline) + item: TResponseInputItem = {"role": "user", "content": "once"} + + await session.add_items([item]) + + assert execute_calls == 1 + assert await session.get_items() == [item] + + +async def test_add_items_rejects_unrepresentable_ttl_before_writing() -> None: + """An invalid Redis TTL must not turn a failed write retry into duplicates.""" + if not USE_FAKE_REDIS: + pytest.skip("This test requires fakeredis") + + client = fakeredis.aioredis.FakeRedis() + session = RedisSession( + session_id="invalid_ttl", + redis_client=cast("Redis", client), + key_prefix="test:", + ttl=2**63, + ) + item: TResponseInputItem = {"role": "user", "content": "never committed"} + + for _ in range(2): + with pytest.raises(ValueError, match="outside Redis's supported expiration range"): + await session.add_items([item]) + + assert await session.get_items() == [] + + +async def test_add_items_rejects_wrong_metadata_key_type_before_writing() -> None: + """A metadata type error must not commit history that a retry duplicates.""" + if not USE_FAKE_REDIS: + pytest.skip("This test requires fakeredis") + + from redis.exceptions import ResponseError + + client = fakeredis.aioredis.FakeRedis() + session = RedisSession( + session_id="wrong_metadata_type", + redis_client=cast("Redis", client), + key_prefix="test:", + ) + await client.set(session._session_key, "not a hash") + item: TResponseInputItem = {"role": "user", "content": "never committed"} + + for _ in range(2): + with pytest.raises(ResponseError, match="metadata key must contain a hash"): + await session.add_items([item]) + + assert await session.get_items() == [] + + +async def test_add_items_uses_server_absolute_expiration_after_serialization( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """TTL validation must not become stale between serialization and Redis execution.""" + if not USE_FAKE_REDIS: + pytest.skip("This test requires fakeredis") + + client = fakeredis.aioredis.FakeRedis() + server_seconds = 1_750_000_000 + server_microseconds = 500_000 + server_time_ms = server_seconds * 1000 + server_microseconds // 1000 + ttl = (2**63 - 1 - server_time_ms) // 1000 + expected_expiration_ms = server_time_ms + ttl * 1000 + order: list[str] = [] + real_pipeline = client.pipeline + + session = RedisSession( + session_id="absolute_ttl", + redis_client=cast("Redis", client), + key_prefix="test:", + ttl=ttl, + ) + + async def serialize(item: TResponseInputItem) -> str: + order.append("serialize") + return json.dumps(item) + + def pipeline(*args: Any, **kwargs: Any) -> Any: + delegate = real_pipeline(*args, **kwargs) + + class PipelineProxy: + def __getattr__(self, name: str) -> Any: + return getattr(delegate, name) + + async def time(self) -> tuple[int, int]: + order.append("time") + return server_seconds, server_microseconds + + return PipelineProxy() + + monkeypatch.setattr(session, "_serialize_item", serialize) + monkeypatch.setattr(client, "pipeline", pipeline) + item: TResponseInputItem = {"role": "user", "content": "once"} + + await session.add_items([item]) + + assert order == ["serialize", "time"] + assert -(2**63) <= expected_expiration_ms <= 2**63 - 1 + assert await session.get_items() == [item] + + +async def test_add_items_refreshes_server_timestamps_after_watch_retry( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """A watched retry must derive metadata and expiration from its own attempt.""" + if not USE_FAKE_REDIS: + pytest.skip("This test requires fakeredis") + + from redis.exceptions import WatchError + + client = fakeredis.aioredis.FakeRedis() + ttl = 60 + first_seconds = int(time.time()) + second_seconds = first_seconds + ttl * 2 + server_times = iter([(first_seconds, 0), (second_seconds, 0)]) + expirations: list[int] = [] + pipeline_attempt = 0 + real_pipeline = client.pipeline + + session = RedisSession( + session_id="watch_retry_timestamps", + redis_client=cast("Redis", client), + key_prefix="test:", + ttl=ttl, + ) + + def pipeline(*args: Any, **kwargs: Any) -> Any: + nonlocal pipeline_attempt + pipeline_attempt += 1 + attempt = pipeline_attempt + delegate = real_pipeline(*args, **kwargs) + + class PipelineProxy: + def __getattr__(self, name: str) -> Any: + return getattr(delegate, name) + + def pexpireat(self, key: str, expiration_time_ms: int) -> Any: + expirations.append(expiration_time_ms) + return delegate.pexpireat(key, expiration_time_ms) + + async def time(self) -> tuple[int, int]: + return next(server_times) + + async def execute(self) -> Any: + if attempt == 1: + raise WatchError("retry") + return await delegate.execute() + + return PipelineProxy() + + monkeypatch.setattr(client, "pipeline", pipeline) + item: TResponseInputItem = {"role": "user", "content": "once"} + + await session.add_items([item]) + + first_expiration = (first_seconds + ttl) * 1000 + second_expiration = (second_seconds + ttl) * 1000 + assert expirations == [first_expiration] * 3 + [second_expiration] * 3 + metadata = await client.hgetall(session._session_key) # type: ignore[misc] + updated_at = metadata.get(b"updated_at") or metadata.get("updated_at") + assert updated_at in (str(second_seconds), str(second_seconds).encode()) + assert await session.get_items() == [item] + + +async def test_add_items_with_ttl_supports_single_connection_pool() -> None: + """WATCH and server time must share one caller-managed Redis connection.""" + if not USE_FAKE_REDIS: + pytest.skip("This test requires fakeredis") + + client = fakeredis.aioredis.FakeRedis(max_connections=1) + session = RedisSession( + session_id="single_connection_ttl", + redis_client=cast("Redis", client), + key_prefix="test:", + ttl=60, + ) + item: TResponseInputItem = {"role": "user", "content": "once"} + + await session.add_items([item]) + + assert await session.get_items() == [item] + + +async def test_cancelled_watch_releases_single_connection_pool( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Repeated cancellation must not interrupt watched-pipeline cleanup.""" + if not USE_FAKE_REDIS: + pytest.skip("This test requires fakeredis") + + client = fakeredis.aioredis.FakeRedis(max_connections=1) + session = RedisSession( + session_id="cancelled_watch_cleanup", + redis_client=cast("Redis", client), + key_prefix="test:", + ttl=60, + ) + real_pipeline = client.pipeline + time_started = asyncio.Event() + reset_started = asyncio.Event() + allow_reset = asyncio.Event() + + def pipeline(*args: Any, **kwargs: Any) -> Any: + delegate = real_pipeline(*args, **kwargs) + + class PipelineProxy: + def __getattr__(self, name: str) -> Any: + return getattr(delegate, name) + + async def time(self) -> tuple[int, int]: + time_started.set() + await asyncio.Event().wait() + raise AssertionError("unreachable") + + async def reset(self) -> None: + reset_started.set() + await allow_reset.wait() + await delegate.reset() + + return PipelineProxy() + + monkeypatch.setattr(client, "pipeline", pipeline) + add_task = asyncio.create_task(session.add_items([{"role": "user", "content": "cancelled"}])) + try: + await time_started.wait() + add_task.cancel() + await reset_started.wait() + add_task.cancel() + await asyncio.sleep(0) + allow_reset.set() + with pytest.raises(asyncio.CancelledError): + await add_task + finally: + allow_reset.set() + if not add_task.done(): + add_task.cancel() + await asyncio.gather(add_task, return_exceptions=True) + + assert await client.ping() is True # type: ignore[misc] + + +async def test_post_commit_cancellation_does_not_report_retryable_failure( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Cancellation during driver cleanup must not hide a known successful commit.""" + if not USE_FAKE_REDIS: + pytest.skip("This test requires fakeredis") + + client = fakeredis.aioredis.FakeRedis(max_connections=1) + session = RedisSession( + session_id="post_commit_cancellation", + redis_client=cast("Redis", client), + key_prefix="test:", + ) + real_pipeline = client.pipeline + internal_reset_started = asyncio.Event() + allow_internal_reset = asyncio.Event() + + def pipeline(*args: Any, **kwargs: Any) -> Any: + delegate = real_pipeline(*args, **kwargs) + real_reset = delegate.reset + reset_calls = 0 + + async def controlled_reset() -> None: + nonlocal reset_calls + reset_calls += 1 + if reset_calls == 1: + internal_reset_started.set() + await allow_internal_reset.wait() + await real_reset() + + monkeypatch.setattr(delegate, "reset", controlled_reset) + return delegate + + monkeypatch.setattr(client, "pipeline", pipeline) + item: TResponseInputItem = {"role": "user", "content": "committed"} + add_task = asyncio.create_task(session.add_items([item])) + try: + await internal_reset_started.wait() + add_task.cancel() + await asyncio.sleep(0) + add_task.cancel() + await asyncio.sleep(0) + allow_internal_reset.set() + await add_task + finally: + allow_internal_reset.set() + if not add_task.done(): + add_task.cancel() + await asyncio.gather(add_task, return_exceptions=True) + + assert await session.get_items() == [item] + assert await client.ping() is True # type: ignore[misc] + + +async def test_post_commit_internal_reset_failure_does_not_report_write_failure( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """A cleanup failure after a successful EXEC must not invite a duplicate retry.""" + if not USE_FAKE_REDIS: + pytest.skip("This test requires fakeredis") + + client = fakeredis.aioredis.FakeRedis(max_connections=1) + session = RedisSession( + session_id="post_commit_reset_failure", + redis_client=cast("Redis", client), + key_prefix="test:", + ) + real_pipeline = client.pipeline + + def pipeline(*args: Any, **kwargs: Any) -> Any: + delegate = real_pipeline(*args, **kwargs) + real_reset = delegate.reset + reset_calls = 0 + + async def fail_first_reset_after_release() -> None: + nonlocal reset_calls + reset_calls += 1 + await real_reset() + if reset_calls == 1: + raise RuntimeError("post-commit reset failed") + + monkeypatch.setattr(delegate, "reset", fail_first_reset_after_release) + return delegate + + monkeypatch.setattr(client, "pipeline", pipeline) + item: TResponseInputItem = {"role": "user", "content": "committed"} + + await session.add_items([item]) + + assert await session.get_items() == [item] + assert await client.ping() is True # type: ignore[misc] + + async def test_redis_session_operation_waiting_behind_close_raises(): """An operation queued behind close() must fail rather than run after shutdown completes.""" if not USE_FAKE_REDIS: diff --git a/tests/extensions/memory/test_sqlalchemy_session.py b/tests/extensions/memory/test_sqlalchemy_session.py index b1984f4e4a..2829cfeedd 100644 --- a/tests/extensions/memory/test_sqlalchemy_session.py +++ b/tests/extensions/memory/test_sqlalchemy_session.py @@ -237,6 +237,91 @@ async def test_pop_from_empty_session(): assert popped is None +async def test_concurrent_pop_item_returns_each_row_once(monkeypatch: pytest.MonkeyPatch): + """A caller that loses the DELETE race must retry instead of returning the same row.""" + session = SQLAlchemySession.from_url("concurrent_pop", url=DB_URL, create_tables=False) + rows = { + 1: json.dumps({"role": "user", "content": "first"}), + 2: json.dumps({"role": "assistant", "content": "second"}), + } + first_reads = 0 + both_read_newest = asyncio.Event() + first_payload_reads = 0 + both_read_newest_payload = asyncio.Event() + + class FakeResult: + def __init__(self, scalar: Any = None, rowcount: int = -1) -> None: + self._scalar = scalar + self.rowcount = rowcount + + def scalar_one_or_none(self) -> Any: + return self._scalar + + class FakeTransaction: + async def __aenter__(self) -> None: + return None + + async def __aexit__(self, *args: Any) -> None: + return None + + class FakeSession: + selected_row_id: int | None = None + selected_row: str | None = None + + async def __aenter__(self) -> FakeSession: + return self + + async def __aexit__(self, *args: Any) -> None: + return None + + def begin(self) -> FakeTransaction: + return FakeTransaction() + + async def execute(self, statement: Any) -> FakeResult: + nonlocal first_payload_reads, first_reads + if isinstance(statement, Select): + selected_columns = list(statement.selected_columns.keys()) + if selected_columns == ["id"]: + self.selected_row_id = max(rows, default=None) + if first_reads < 2: + first_reads += 1 + if first_reads == 2: + both_read_newest.set() + await both_read_newest.wait() + return FakeResult(self.selected_row_id) + assert selected_columns == ["message_data"] + assert self.selected_row_id is not None + self.selected_row = rows.get(self.selected_row_id) + if first_payload_reads < 2: + first_payload_reads += 1 + if first_payload_reads == 2: + both_read_newest_payload.set() + await both_read_newest_payload.wait() + return FakeResult(self.selected_row) + + assert self.selected_row_id is not None + if self.selected_row_id in rows: + del rows[self.selected_row_id] + return FakeResult(rowcount=1) + return FakeResult(rowcount=0) + + class FakeSessionFactory: + def __call__(self) -> FakeSession: + return FakeSession() + + async def tables_ready() -> None: + return None + + monkeypatch.setattr(session, "_ensure_tables", tables_ready) + monkeypatch.setattr(session, "_session_factory", FakeSessionFactory()) + + popped = await asyncio.gather(session.pop_item(), session.pop_item()) + + contents = {cast(dict[str, Any], item)["content"] for item in popped if item is not None} + assert contents == {"first", "second"} + assert rows == {} + + async def test_pop_item_skips_corrupt_most_recent(): """pop_item skips corrupt newest rows and returns the next valid item.""" session = SQLAlchemySession.from_url("pop_corrupt", url=DB_URL, create_tables=True) diff --git a/tests/memory/test_session.py b/tests/memory/test_session.py index a2df310df1..60fbc41c10 100644 --- a/tests/memory/test_session.py +++ b/tests/memory/test_session.py @@ -3,8 +3,9 @@ import asyncio import sqlite3 import tempfile +import threading from pathlib import Path -from typing import cast +from typing import Any, cast import pytest @@ -967,3 +968,214 @@ async def test_runner_with_session_settings_override(): assert len(history_items) == 2 session.close() + + +def _drop_sqlite_table(db_path: Path, table: str) -> None: + """Drop a table from an independent connection to make a later statement fail.""" + helper = sqlite3.connect(str(db_path)) + try: + helper.execute(f"DROP TABLE {table}") + helper.commit() + finally: + helper.close() + + +def _sqlite_write_lock_is_free(db_path: Path) -> bool: + """Return whether an independent writer can take the SQLite write lock.""" + probe = sqlite3.connect(str(db_path), timeout=0) + try: + probe.execute("CREATE TABLE IF NOT EXISTS probe_lock (x INTEGER)") + probe.commit() + return True + except sqlite3.OperationalError: + return False + finally: + probe.close() + + +@pytest.mark.asyncio +async def test_sqlite_session_failed_clear_session_rolls_back(): + """A failed clear must restore earlier statements and release the cached write lock.""" + with tempfile.TemporaryDirectory() as temp_dir: + db_path = Path(temp_dir) / "clear_rollback.db" + session = SQLiteSession("clear_rollback", db_path) + await session.add_items([{"role": "user", "content": "kept"}]) + + _drop_sqlite_table(db_path, "agent_sessions") + + with pytest.raises(sqlite3.OperationalError): + await session.clear_session() + + assert all(not conn.in_transaction for conn in session._connections) + assert _sqlite_write_lock_is_free(db_path) + session.close() + + +@pytest.mark.asyncio +async def test_sqlite_session_failed_pop_item_releases_write_lock(): + """A failed pop must not leave a write transaction on the cached connection.""" + with tempfile.TemporaryDirectory() as temp_dir: + db_path = Path(temp_dir) / "pop_rollback.db" + session = SQLiteSession("pop_rollback", db_path) + await session.add_items([{"role": "user", "content": "kept"}]) + + _drop_sqlite_table(db_path, "agent_messages") + + with pytest.raises(sqlite3.OperationalError): + await session.pop_item() + + assert all(not conn.in_transaction for conn in session._connections) + assert _sqlite_write_lock_is_free(db_path) + session.close() + + +@pytest.mark.asyncio +async def test_sqlite_session_rollback_failure_evicts_connection( + monkeypatch: pytest.MonkeyPatch, +): + """A file connection that cannot roll back must be closed and replaced.""" + + class FailingRollbackConnection(sqlite3.Connection): + def rollback(self) -> None: + raise RuntimeError("rollback failed") + + with tempfile.TemporaryDirectory() as temp_dir: + db_path = Path(temp_dir) / "rollback_failure.db" + session = SQLiteSession("rollback_failure", db_path) + conn = sqlite3.connect( + str(db_path), + check_same_thread=False, + factory=FailingRollbackConnection, + ) + with session._connections_lock: + session._connections.add(conn) + real_get_connection = session._get_connection + monkeypatch.setattr(session, "_get_connection", lambda: conn) + unserializable = cast(TResponseInputItem, {"role": "user", "content": object()}) + + with pytest.raises(TypeError): + await session.add_items([unserializable]) + + assert conn not in session._connections + assert _sqlite_write_lock_is_free(db_path) + + monkeypatch.setattr(session, "_get_connection", real_get_connection) + await session.add_items([{"role": "user", "content": "after failure"}]) + assert [item.get("content") for item in await session.get_items()] == ["after failure"] + session.close() + + +@pytest.mark.asyncio +async def test_sqlite_session_close_retries_quarantined_connection( + monkeypatch: pytest.MonkeyPatch, +): + """A failed invalidation close must remain owned until a later close succeeds.""" + + class FailingRollbackAndCloseConnection(sqlite3.Connection): + fail_close = True + + def rollback(self) -> None: + raise RuntimeError("rollback failed") + + def close(self) -> None: + if self.fail_close: + raise RuntimeError("close failed") + super().close() + + with tempfile.TemporaryDirectory() as temp_dir: + db_path = Path(temp_dir) / "close_retry.db" + session = SQLiteSession("close_retry", db_path) + conn = sqlite3.connect( + str(db_path), + check_same_thread=False, + factory=FailingRollbackAndCloseConnection, + ) + with session._connections_lock: + session._connections.add(conn) + monkeypatch.setattr(session, "_get_connection", lambda: conn) + unserializable = cast(TResponseInputItem, {"role": "user", "content": object()}) + + with pytest.raises(TypeError): + await session.add_items([unserializable]) + + assert session._closed is True + assert conn in session._quarantined_connections + assert _sqlite_write_lock_is_free(db_path) is False + + conn.fail_close = False + session.close() + + assert session._quarantined_connections == set() + assert _sqlite_write_lock_is_free(db_path) + with pytest.raises(sqlite3.ProgrammingError): + conn.execute("SELECT 1") + + +@pytest.mark.asyncio +@pytest.mark.parametrize("operation", ["add", "pop", "clear"]) +async def test_sqlite_session_post_commit_cancellation_returns_known_outcome( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + operation: str, +): + """Cancellation after a worker commit must not invite a duplicate mutation retry.""" + + class PausingCommitConnection(sqlite3.Connection): + pause_commit = False + commit_finished = threading.Event() + allow_return = threading.Event() + + def commit(self) -> None: + super().commit() + if self.pause_commit: + self.pause_commit = False + self.commit_finished.set() + assert self.allow_return.wait(timeout=10) + + db_path = tmp_path / f"post_commit_{operation}.db" + session = SQLiteSession(f"post_commit_{operation}", db_path) + item: TResponseInputItem = {"role": "user", "content": "once"} + if operation != "add": + await session.add_items([item]) + + conn = sqlite3.connect( + str(db_path), + check_same_thread=False, + factory=PausingCommitConnection, + ) + with session._connections_lock: + session._connections.add(conn) + monkeypatch.setattr(session, "_get_connection", lambda: conn) + conn.pause_commit = True + + if operation == "add": + mutation: asyncio.Task[Any] = asyncio.create_task(session.add_items([item])) + elif operation == "pop": + mutation = asyncio.create_task(session.pop_item()) + else: + mutation = asyncio.create_task(session.clear_session()) + + try: + assert await asyncio.to_thread(conn.commit_finished.wait, 10) + mutation.cancel() + await asyncio.sleep(0) + mutation.cancel() + await asyncio.sleep(0) + conn.allow_return.set() + result = await mutation + finally: + conn.allow_return.set() + if not mutation.done(): + mutation.cancel() + await asyncio.gather(mutation, return_exceptions=True) + + if operation == "add": + assert result is None + assert await session.get_items() == [item] + elif operation == "pop": + assert result == item + assert await session.get_items() == [] + else: + assert result is None + assert await session.get_items() == [] + session.close() From 068698f192832240f3999144e01732e81640772c Mon Sep 17 00:00:00 2001 From: Kazuhiro Sera Date: Wed, 5 Aug 2026 22:20:47 +0900 Subject: [PATCH 2/7] fix: harden session mutation ownership --- .../extensions/memory/sqlalchemy_session.py | 8 + .../memory/test_async_sqlite_session.py | 504 ++++++++++-------- .../extensions/memory/test_mongodb_session.py | 48 +- .../memory/test_sqlalchemy_session.py | 91 +++- 4 files changed, 425 insertions(+), 226 deletions(-) diff --git a/src/agents/extensions/memory/sqlalchemy_session.py b/src/agents/extensions/memory/sqlalchemy_session.py index 713b5986a2..40b2ca5e1d 100644 --- a/src/agents/extensions/memory/sqlalchemy_session.py +++ b/src/agents/extensions/memory/sqlalchemy_session.py @@ -450,6 +450,14 @@ async def pop_item(self) -> TResponseInputItem | None: # A concurrent caller claimed this row. Start a new transaction so # dialects with repeatable-read snapshots can observe the next row. continue + if delete_result.rowcount != 1: + # Unknown or unexpected row counts cannot prove that this caller owns + # the destructive read. Raising inside the transaction rolls the DELETE + # back instead of risking the same payload being returned twice. + raise RuntimeError( + "SQLAlchemySession.pop_item() could not confirm that exactly one " + "item was deleted" + ) if row is None: continue diff --git a/tests/extensions/memory/test_async_sqlite_session.py b/tests/extensions/memory/test_async_sqlite_session.py index 6d843ab8e4..5b353ce1f2 100644 --- a/tests/extensions/memory/test_async_sqlite_session.py +++ b/tests/extensions/memory/test_async_sqlite_session.py @@ -506,35 +506,51 @@ async def test_cancelled_close_finishes_cleanup_and_propagates_cancellation( close_started = asyncio.Event() allow_close = asyncio.Event() session = AsyncSQLiteSession("cancelled_close") - conn = await session._get_connection() - real_close = conn.close + conn: Any = None + real_close: Any = None + close_task: asyncio.Task[None] | None = None + try: + conn = await session._get_connection() + real_close = conn.close - async def controlled_close() -> None: - close_started.set() - await allow_close.wait() - await real_close() + async def controlled_close() -> None: + close_started.set() + await allow_close.wait() + await real_close() - monkeypatch.setattr(conn, "close", controlled_close) - close_task = asyncio.create_task(session.close()) - try: - await close_started.wait() - close_task.cancel() - await asyncio.sleep(0) - close_task.cancel() - await asyncio.sleep(0) - allow_close.set() - with pytest.raises(asyncio.CancelledError): - await close_task + monkeypatch.setattr(conn, "close", controlled_close) + close_task = asyncio.create_task(session.close()) + try: + await close_started.wait() + close_task.cancel() + await asyncio.sleep(0) + close_task.cancel() + await asyncio.sleep(0) + allow_close.set() + with pytest.raises(asyncio.CancelledError): + await close_task + finally: + allow_close.set() + if not close_task.done(): + close_task.cancel() + await asyncio.gather(close_task, return_exceptions=True) + + assert session._closed is True + assert session._connection is None + assert session._quarantined_connections == set() + assert conn._running is False finally: allow_close.set() - if not close_task.done(): + if close_task is not None and not close_task.done(): close_task.cancel() await asyncio.gather(close_task, return_exceptions=True) - - assert session._closed is True - assert session._connection is None - assert session._quarantined_connections == set() - assert conn._running is False + if conn is not None and real_close is not None: + monkeypatch.setattr(conn, "close", real_close) + try: + await session.close() + finally: + if conn is not None and real_close is not None and conn._running: + await real_close() @pytest.mark.parametrize("operation", ["add", "pop", "clear"]) @@ -547,55 +563,57 @@ async def test_post_commit_cancellation_returns_known_mutation_outcome( db_path = tmp_path / f"async_post_commit_{operation}.db" session = AsyncSQLiteSession(f"async_post_commit_{operation}", db_path) item: TResponseInputItem = {"role": "user", "content": "once"} - if operation != "add": - await session.add_items([item]) - - conn = await session._get_connection() - real_commit = conn.commit - commit_finished = asyncio.Event() - allow_return = asyncio.Event() - pause_commit = True - - async def controlled_commit() -> None: - nonlocal pause_commit - await real_commit() - if pause_commit: - pause_commit = False - commit_finished.set() - await allow_return.wait() - - monkeypatch.setattr(conn, "commit", controlled_commit) - if operation == "add": - mutation: asyncio.Task[Any] = asyncio.create_task(session.add_items([item])) - elif operation == "pop": - mutation = asyncio.create_task(session.pop_item()) - else: - mutation = asyncio.create_task(session.clear_session()) - try: - await commit_finished.wait() - mutation.cancel() - await asyncio.sleep(0) - mutation.cancel() - await asyncio.sleep(0) - allow_return.set() - result = await mutation - finally: - allow_return.set() - if not mutation.done(): - mutation.cancel() - await asyncio.gather(mutation, return_exceptions=True) + if operation != "add": + await session.add_items([item]) - if operation == "add": - assert result is None - assert await session.get_items() == [item] - elif operation == "pop": - assert result == item - assert await session.get_items() == [] - else: - assert result is None - assert await session.get_items() == [] - await session.close() + conn = await session._get_connection() + real_commit = conn.commit + commit_finished = asyncio.Event() + allow_return = asyncio.Event() + pause_commit = True + + async def controlled_commit() -> None: + nonlocal pause_commit + await real_commit() + if pause_commit: + pause_commit = False + commit_finished.set() + await allow_return.wait() + + monkeypatch.setattr(conn, "commit", controlled_commit) + if operation == "add": + mutation: asyncio.Task[Any] = asyncio.create_task(session.add_items([item])) + elif operation == "pop": + mutation = asyncio.create_task(session.pop_item()) + else: + mutation = asyncio.create_task(session.clear_session()) + + try: + await commit_finished.wait() + mutation.cancel() + await asyncio.sleep(0) + mutation.cancel() + await asyncio.sleep(0) + allow_return.set() + result = await mutation + finally: + allow_return.set() + if not mutation.done(): + mutation.cancel() + await asyncio.gather(mutation, return_exceptions=True) + + if operation == "add": + assert result is None + assert await session.get_items() == [item] + elif operation == "pop": + assert result == item + assert await session.get_items() == [] + else: + assert result is None + assert await session.get_items() == [] + finally: + await session.close() def _drop_sqlite_table(db_path: Path, table: str) -> None: @@ -627,17 +645,18 @@ async def test_failed_add_items_rolls_back_and_reuses_connection(): db_path = Path(temp_dir) / "add_rollback.db" session = AsyncSQLiteSession("add_rollback", db_path) unserializable = cast(TResponseInputItem, {"role": "user", "content": object()}) + try: + with pytest.raises(TypeError): + await session.add_items([unserializable]) - with pytest.raises(TypeError): - await session.add_items([unserializable]) - - conn = await session._get_connection() - assert conn.in_transaction is False - assert _sqlite_write_lock_is_free(db_path) + conn = await session._get_connection() + assert conn.in_transaction is False + assert _sqlite_write_lock_is_free(db_path) - await session.add_items([{"role": "user", "content": "after failure"}]) - assert [item.get("content") for item in await session.get_items()] == ["after failure"] - await session.close() + await session.add_items([{"role": "user", "content": "after failure"}]) + assert [item.get("content") for item in await session.get_items()] == ["after failure"] + finally: + await session.close() async def test_failed_clear_session_rolls_back(): @@ -645,17 +664,19 @@ async def test_failed_clear_session_rolls_back(): with tempfile.TemporaryDirectory() as temp_dir: db_path = Path(temp_dir) / "clear_rollback.db" session = AsyncSQLiteSession("clear_rollback", db_path) - await session.add_items([{"role": "user", "content": "kept"}]) + try: + await session.add_items([{"role": "user", "content": "kept"}]) - _drop_sqlite_table(db_path, "agent_sessions") + _drop_sqlite_table(db_path, "agent_sessions") - with pytest.raises(sqlite3.OperationalError): - await session.clear_session() + with pytest.raises(sqlite3.OperationalError): + await session.clear_session() - conn = await session._get_connection() - assert conn.in_transaction is False - assert _sqlite_write_lock_is_free(db_path) - await session.close() + conn = await session._get_connection() + assert conn.in_transaction is False + assert _sqlite_write_lock_is_free(db_path) + finally: + await session.close() async def test_failed_pop_item_releases_write_lock(): @@ -663,17 +684,19 @@ async def test_failed_pop_item_releases_write_lock(): with tempfile.TemporaryDirectory() as temp_dir: db_path = Path(temp_dir) / "pop_rollback.db" session = AsyncSQLiteSession("pop_rollback", db_path) - await session.add_items([{"role": "user", "content": "kept"}]) + try: + await session.add_items([{"role": "user", "content": "kept"}]) - _drop_sqlite_table(db_path, "agent_messages") + _drop_sqlite_table(db_path, "agent_messages") - with pytest.raises(sqlite3.OperationalError): - await session.pop_item() + with pytest.raises(sqlite3.OperationalError): + await session.pop_item() - conn = await session._get_connection() - assert conn.in_transaction is False - assert _sqlite_write_lock_is_free(db_path) - await session.close() + conn = await session._get_connection() + assert conn.in_transaction is False + assert _sqlite_write_lock_is_free(db_path) + finally: + await session.close() async def test_failed_initialization_closes_candidate_connection(): @@ -687,11 +710,18 @@ async def _init_db_for_connection(self, conn: Any) -> None: raise RuntimeError("initialization failed") session = FailingInitSession("failed_init") - with pytest.raises(RuntimeError, match="initialization failed"): - await session.get_items() + try: + with pytest.raises(RuntimeError, match="initialization failed"): + await session.get_items() - assert session._connection is None - assert session.captured_connection._running is False + assert session._connection is None + assert session.captured_connection._running is False + finally: + try: + await session.close() + finally: + if session.captured_connection is not None and session.captured_connection._running: + await session.captured_connection.close() async def test_cancelled_add_items_rolls_back_write_transaction( @@ -701,48 +731,50 @@ async def test_cancelled_add_items_rolls_back_write_transaction( """Cancellation after the first write must release the transaction and database lock.""" db_path = tmp_path / "cancelled_add.db" session = AsyncSQLiteSession("cancelled_add", db_path) - await session.get_items() - conn = await session._get_connection() - real_execute = conn.execute - real_rollback = conn.rollback - first_write_started = asyncio.Event() - rollback_started = asyncio.Event() - allow_rollback = asyncio.Event() - - async def pause_after_first_write(*args: Any, **kwargs: Any) -> Any: - cursor = await real_execute(*args, **kwargs) - first_write_started.set() - await asyncio.Event().wait() - return cursor - - async def controlled_rollback() -> None: - rollback_started.set() - await allow_rollback.wait() - await real_rollback() - - monkeypatch.setattr(conn, "execute", pause_after_first_write) - monkeypatch.setattr(conn, "rollback", controlled_rollback) - task = asyncio.create_task(session.add_items([{"role": "user", "content": "cancelled"}])) try: - await first_write_started.wait() - task.cancel() - await rollback_started.wait() - task.cancel() - allow_rollback.set() - with pytest.raises(asyncio.CancelledError): - await task - finally: - allow_rollback.set() - monkeypatch.setattr(conn, "execute", real_execute) - monkeypatch.setattr(conn, "rollback", real_rollback) - if not task.done(): + await session.get_items() + conn = await session._get_connection() + real_execute = conn.execute + real_rollback = conn.rollback + first_write_started = asyncio.Event() + rollback_started = asyncio.Event() + allow_rollback = asyncio.Event() + + async def pause_after_first_write(*args: Any, **kwargs: Any) -> Any: + cursor = await real_execute(*args, **kwargs) + first_write_started.set() + await asyncio.Event().wait() + return cursor + + async def controlled_rollback() -> None: + rollback_started.set() + await allow_rollback.wait() + await real_rollback() + + monkeypatch.setattr(conn, "execute", pause_after_first_write) + monkeypatch.setattr(conn, "rollback", controlled_rollback) + task = asyncio.create_task(session.add_items([{"role": "user", "content": "cancelled"}])) + try: + await first_write_started.wait() + task.cancel() + await rollback_started.wait() task.cancel() - await asyncio.gather(task, return_exceptions=True) + allow_rollback.set() + with pytest.raises(asyncio.CancelledError): + await task + finally: + allow_rollback.set() + monkeypatch.setattr(conn, "execute", real_execute) + monkeypatch.setattr(conn, "rollback", real_rollback) + if not task.done(): + task.cancel() + await asyncio.gather(task, return_exceptions=True) - assert conn.in_transaction is False - assert _sqlite_write_lock_is_free(db_path) - assert await session.get_items() == [] - await session.close() + assert conn.in_transaction is False + assert _sqlite_write_lock_is_free(db_path) + assert await session.get_items() == [] + finally: + await session.close() async def test_rollback_failure_closes_and_evicts_connection( @@ -752,25 +784,27 @@ async def test_rollback_failure_closes_and_evicts_connection( """A connection that cannot roll back must not remain cached or retain its write lock.""" db_path = tmp_path / "rollback_failure.db" session = AsyncSQLiteSession("rollback_failure", db_path) - await session.get_items() - conn = await session._get_connection() + try: + await session.get_items() + conn = await session._get_connection() - async def fail_rollback() -> None: - raise RuntimeError("rollback failed") + async def fail_rollback() -> None: + raise RuntimeError("rollback failed") - monkeypatch.setattr(conn, "rollback", fail_rollback) - unserializable = cast(TResponseInputItem, {"role": "user", "content": object()}) + monkeypatch.setattr(conn, "rollback", fail_rollback) + unserializable = cast(TResponseInputItem, {"role": "user", "content": object()}) - with pytest.raises(TypeError): - await session.add_items([unserializable]) + with pytest.raises(TypeError): + await session.add_items([unserializable]) - assert session._connection is None - assert session._closed is False - assert _sqlite_write_lock_is_free(db_path) + assert session._connection is None + assert session._closed is False + assert _sqlite_write_lock_is_free(db_path) - await session.add_items([{"role": "user", "content": "after failure"}]) - assert [item.get("content") for item in await session.get_items()] == ["after failure"] - await session.close() + await session.add_items([{"role": "user", "content": "after failure"}]) + assert [item.get("content") for item in await session.get_items()] == ["after failure"] + finally: + await session.close() async def test_close_retries_connection_quarantined_after_rollback_failure( @@ -780,35 +814,46 @@ async def test_close_retries_connection_quarantined_after_rollback_failure( """A failed invalidation close must remain owned until a later close succeeds.""" db_path = tmp_path / "close_retry.db" session = AsyncSQLiteSession("close_retry", db_path) - await session.get_items() - conn = await session._get_connection() - real_close = conn.close + conn: Any = None + real_close: Any = None + try: + await session.get_items() + conn = await session._get_connection() + real_close = conn.close - async def fail_rollback() -> None: - raise RuntimeError("rollback failed") + async def fail_rollback() -> None: + raise RuntimeError("rollback failed") - async def fail_close() -> None: - raise RuntimeError("close failed") + async def fail_close() -> None: + raise RuntimeError("close failed") - monkeypatch.setattr(conn, "rollback", fail_rollback) - monkeypatch.setattr(conn, "close", fail_close) - unserializable = cast(TResponseInputItem, {"role": "user", "content": object()}) + monkeypatch.setattr(conn, "rollback", fail_rollback) + monkeypatch.setattr(conn, "close", fail_close) + unserializable = cast(TResponseInputItem, {"role": "user", "content": object()}) - with pytest.raises(TypeError): - await session.add_items([unserializable]) + with pytest.raises(TypeError): + await session.add_items([unserializable]) - assert session._closed is True - assert session._connection is None - assert conn in session._quarantined_connections - assert conn._running is True - assert _sqlite_write_lock_is_free(db_path) is False + assert session._closed is True + assert session._connection is None + assert conn in session._quarantined_connections + assert conn._running is True + assert _sqlite_write_lock_is_free(db_path) is False - monkeypatch.setattr(conn, "close", real_close) - await session.close() + monkeypatch.setattr(conn, "close", real_close) + await session.close() - assert session._quarantined_connections == set() - assert conn._running is False - assert _sqlite_write_lock_is_free(db_path) + assert session._quarantined_connections == set() + assert conn._running is False + assert _sqlite_write_lock_is_free(db_path) + finally: + if conn is not None and real_close is not None: + monkeypatch.setattr(conn, "close", real_close) + try: + await session.close() + finally: + if conn is not None and real_close is not None and conn._running: + await real_close() async def test_close_retries_quarantined_failed_initialization_candidate( @@ -831,19 +876,32 @@ async def fail_close() -> None: raise RuntimeError("initialization failed") session = FailingInitSession("failed_init_close_retry") - with pytest.raises(RuntimeError, match="initialization failed"): - await session.get_items() + try: + with pytest.raises(RuntimeError, match="initialization failed"): + await session.get_items() - conn = session.captured_connection - assert session._closed is True - assert conn in session._quarantined_connections - assert conn._running is True + conn = session.captured_connection + assert session._closed is True + assert conn in session._quarantined_connections + assert conn._running is True - monkeypatch.setattr(conn, "close", session.real_close) - await session.close() + monkeypatch.setattr(conn, "close", session.real_close) + await session.close() - assert session._quarantined_connections == set() - assert conn._running is False + assert session._quarantined_connections == set() + assert conn._running is False + finally: + if session.captured_connection is not None and session.real_close is not None: + monkeypatch.setattr(session.captured_connection, "close", session.real_close) + try: + await session.close() + finally: + if ( + session.captured_connection is not None + and session.real_close is not None + and session.captured_connection._running + ): + await session.real_close() async def test_cancelled_initialization_finishes_candidate_close( @@ -856,15 +914,16 @@ async def test_cancelled_initialization_finishes_candidate_close( class CancelledInitSession(AsyncSQLiteSession): captured_connection: Any = None + real_close: Any = None async def _init_db_for_connection(self, conn: Any) -> None: self.captured_connection = conn - real_close = conn.close + self.real_close = conn.close async def controlled_close() -> None: close_started.set() await allow_close.wait() - await real_close() + await self.real_close() monkeypatch.setattr(conn, "close", controlled_close) init_started.set() @@ -873,21 +932,34 @@ async def controlled_close() -> None: session = CancelledInitSession("cancelled_init") task = asyncio.create_task(session.get_items()) try: - await init_started.wait() - task.cancel() - await close_started.wait() - task.cancel() - allow_close.set() - with pytest.raises(asyncio.CancelledError): - await task - finally: - allow_close.set() - if not task.done(): + try: + await init_started.wait() task.cancel() - await asyncio.gather(task, return_exceptions=True) - - assert session._connection is None - assert session.captured_connection._running is False + await close_started.wait() + task.cancel() + allow_close.set() + with pytest.raises(asyncio.CancelledError): + await task + finally: + allow_close.set() + if not task.done(): + task.cancel() + await asyncio.gather(task, return_exceptions=True) + + assert session._connection is None + assert session.captured_connection._running is False + finally: + if session.captured_connection is not None and session.real_close is not None: + monkeypatch.setattr(session.captured_connection, "close", session.real_close) + try: + await session.close() + finally: + if ( + session.captured_connection is not None + and session.real_close is not None + and session.captured_connection._running + ): + await session.real_close() async def test_cancelled_connect_closes_eventually_acquired_connection( @@ -912,17 +984,23 @@ async def controlled_connect(database: str) -> Any: session = AsyncSQLiteSession("cancelled_connect") task = asyncio.create_task(session.get_items()) try: - await connect_started.wait() - task.cancel() - allow_connect.set() - with pytest.raises(asyncio.CancelledError): - await task - finally: - allow_connect.set() - if not task.done(): + try: + await connect_started.wait() task.cancel() - await asyncio.gather(task, return_exceptions=True) - - assert len(created_connections) == 1 - assert created_connections[0]._running is False - assert session._connection is None + allow_connect.set() + with pytest.raises(asyncio.CancelledError): + await task + finally: + allow_connect.set() + if not task.done(): + task.cancel() + await asyncio.gather(task, return_exceptions=True) + + assert len(created_connections) == 1 + assert created_connections[0]._running is False + assert session._connection is None + finally: + await session.close() + for conn in created_connections: + if conn._running: + await conn.close() diff --git a/tests/extensions/memory/test_mongodb_session.py b/tests/extensions/memory/test_mongodb_session.py index 3496b7fde6..31cd2c496f 100644 --- a/tests/extensions/memory/test_mongodb_session.py +++ b/tests/extensions/memory/test_mongodb_session.py @@ -1062,11 +1062,17 @@ async def pause_then_fail(document: dict[str, Any]) -> None: ] add_task = asyncio.create_task(session.add_items(batch)) - await insert_started.wait() - assert await session.get_items() == [seed] - allow_failure.set() - with pytest.raises(RuntimeError, match="insert failed"): - await add_task + try: + await asyncio.wait_for(insert_started.wait(), timeout=1) + assert await session.get_items() == [seed] + allow_failure.set() + with pytest.raises(RuntimeError, match="insert failed"): + await asyncio.wait_for(add_task, timeout=1) + finally: + allow_failure.set() + if not add_task.done(): + add_task.cancel() + await asyncio.gather(add_task, return_exceptions=True) assert await session.get_items() == [seed] @@ -1075,7 +1081,9 @@ async def pause_then_fail(document: dict[str, Any]) -> None: assert await session.get_items() == [seed, *batch] -async def test_concurrent_pop_item_claims_distinct_items_from_batch() -> None: +async def test_concurrent_pop_item_claims_distinct_items_from_batch( + monkeypatch: pytest.MonkeyPatch, +) -> None: """Atomic array pops must return each item in a logical batch at most once.""" MongoDBSession._init_state.clear() client = FakeAsyncMongoClient() @@ -1091,7 +1099,33 @@ async def test_concurrent_pop_item_claims_distinct_items_from_batch() -> None: ] await session.add_items(items) - popped = await asyncio.gather(*(session.pop_item() for _ in items)) + real_to_list = FakeCursor.to_list + initial_reads = 0 + all_initial_reads_ready = asyncio.Event() + + async def pause_initial_candidate_reads(cursor: FakeCursor) -> list[dict[str, Any]]: + nonlocal initial_reads + docs = [ + {**doc, "message_data": copy.deepcopy(doc.get("message_data"))} + for doc in await real_to_list(cursor) + ] + if initial_reads < len(items): + initial_reads += 1 + if initial_reads == len(items): + all_initial_reads_ready.set() + await all_initial_reads_ready.wait() + return docs + + monkeypatch.setattr(FakeCursor, "to_list", pause_initial_candidate_reads) + tasks = [asyncio.create_task(session.pop_item()) for _ in items] + try: + popped = await asyncio.wait_for(asyncio.gather(*tasks), timeout=1) + finally: + all_initial_reads_ready.set() + for task in tasks: + if not task.done(): + task.cancel() + await asyncio.gather(*tasks, return_exceptions=True) contents = {cast(dict[str, Any], item).get("content") for item in popped if item is not None} assert contents == { diff --git a/tests/extensions/memory/test_sqlalchemy_session.py b/tests/extensions/memory/test_sqlalchemy_session.py index 2829cfeedd..2a1fca969a 100644 --- a/tests/extensions/memory/test_sqlalchemy_session.py +++ b/tests/extensions/memory/test_sqlalchemy_session.py @@ -265,8 +265,12 @@ async def __aexit__(self, *args: Any) -> None: return None class FakeSession: - selected_row_id: int | None = None - selected_row: str | None = None + def __init__(self) -> None: + # Model a repeatable-read transaction: selects keep reading this snapshot even + # after another transaction deletes from the durable rows mapping. + self.snapshot = dict(rows) + self.selected_row_id: int | None = None + self.selected_row: str | None = None async def __aenter__(self) -> FakeSession: return self @@ -282,7 +286,7 @@ async def execute(self, statement: Any) -> FakeResult: if isinstance(statement, Select): selected_columns = list(statement.selected_columns.keys()) if selected_columns == ["id"]: - self.selected_row_id = max(rows, default=None) + self.selected_row_id = max(self.snapshot, default=None) if first_reads < 2: first_reads += 1 if first_reads == 2: @@ -291,7 +295,7 @@ async def execute(self, statement: Any) -> FakeResult: return FakeResult(self.selected_row_id) assert selected_columns == ["message_data"] assert self.selected_row_id is not None - self.selected_row = rows.get(self.selected_row_id) + self.selected_row = self.snapshot.get(self.selected_row_id) if first_payload_reads < 2: first_payload_reads += 1 if first_payload_reads == 2: @@ -312,16 +316,91 @@ def __call__(self) -> FakeSession: async def tables_ready() -> None: return None + session_factory = FakeSessionFactory() monkeypatch.setattr(session, "_ensure_tables", tables_ready) - monkeypatch.setattr(session, "_session_factory", FakeSessionFactory()) + monkeypatch.setattr(session, "_session_factory", session_factory) - popped = await asyncio.gather(session.pop_item(), session.pop_item()) + tasks = [asyncio.create_task(session.pop_item()) for _ in range(2)] + try: + popped = await asyncio.wait_for(asyncio.gather(*tasks), timeout=1) + finally: + both_read_newest.set() + both_read_newest_payload.set() + for task in tasks: + if not task.done(): + task.cancel() + await asyncio.gather(*tasks, return_exceptions=True) + await session.engine.dispose() contents = {cast(dict[str, Any], item)["content"] for item in popped if item is not None} assert contents == {"first", "second"} assert rows == {} +async def test_pop_item_rejects_unknown_delete_rowcount( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """An unknown DELETE row count cannot authorize a destructive read.""" + session = SQLAlchemySession.from_url("unknown_rowcount", url=DB_URL, create_tables=False) + transaction_exit_errors: list[type[BaseException] | None] = [] + + class FakeResult: + def __init__(self, scalar: Any = None, rowcount: int = -1) -> None: + self._scalar = scalar + self.rowcount = rowcount + + def scalar_one_or_none(self) -> Any: + return self._scalar + + class FakeTransaction: + async def __aenter__(self) -> None: + return None + + async def __aexit__( + self, + exc_type: type[BaseException] | None, + exc: BaseException | None, + traceback: Any, + ) -> None: + transaction_exit_errors.append(exc_type) + + class FakeSession: + async def __aenter__(self) -> FakeSession: + return self + + async def __aexit__(self, *args: Any) -> None: + return None + + def begin(self) -> FakeTransaction: + return FakeTransaction() + + async def execute(self, statement: Any) -> FakeResult: + if isinstance(statement, Select): + selected_columns = list(statement.selected_columns.keys()) + if selected_columns == ["id"]: + return FakeResult(1) + return FakeResult(json.dumps({"role": "user", "content": "unclaimed"})) + return FakeResult(rowcount=-1) + + class FakeSessionFactory: + def __call__(self) -> FakeSession: + return FakeSession() + + async def tables_ready() -> None: + return None + + monkeypatch.setattr(session, "_ensure_tables", tables_ready) + monkeypatch.setattr(session, "_session_factory", FakeSessionFactory()) + + try: + with pytest.raises(RuntimeError, match="confirm that exactly one item was deleted"): + await session.pop_item() + finally: + await session.engine.dispose() + + assert transaction_exit_errors == [RuntimeError] + + async def test_pop_item_skips_corrupt_most_recent(): """pop_item skips corrupt newest rows and returns the next valid item.""" session = SQLAlchemySession.from_url("pop_corrupt", url=DB_URL, create_tables=True) From 125de8ed6a9e9f5857be53a6ffb84bada73ecba0 Mon Sep 17 00:00:00 2001 From: Kazuhiro Sera Date: Wed, 5 Aug 2026 22:56:46 +0900 Subject: [PATCH 3/7] fix: preserve cancellation after session mutations --- src/agents/memory/sqlite_session.py | 10 ++++++++-- .../extensions/memory/test_advanced_sqlite_session.py | 11 +++++------ tests/extensions/memory/test_async_sqlite_session.py | 11 +++++------ tests/memory/test_session.py | 11 +++++------ 4 files changed, 23 insertions(+), 20 deletions(-) diff --git a/src/agents/memory/sqlite_session.py b/src/agents/memory/sqlite_session.py index a84ac59452..871f5f04ee 100644 --- a/src/agents/memory/sqlite_session.py +++ b/src/agents/memory/sqlite_session.py @@ -22,7 +22,7 @@ async def _await_mutation(awaitable: Awaitable[_T]) -> _T: cancellation: asyncio.CancelledError | None = None while True: try: - return await asyncio.shield(task) + result = await asyncio.shield(task) except asyncio.CancelledError as exc: if cancellation is None: cancellation = exc @@ -32,12 +32,18 @@ async def _await_mutation(awaitable: Awaitable[_T]) -> _T: if cancellation is not None: raise cancellation from None raise + else: + if cancellation is not None: + raise cancellation from None + return result try: - return task.result() + result = task.result() except BaseException: assert cancellation is not None raise cancellation from None + assert cancellation is not None + raise cancellation from None class SQLiteSession(SessionABC): diff --git a/tests/extensions/memory/test_advanced_sqlite_session.py b/tests/extensions/memory/test_advanced_sqlite_session.py index 6eea7f805a..beef62c9bb 100644 --- a/tests/extensions/memory/test_advanced_sqlite_session.py +++ b/tests/extensions/memory/test_advanced_sqlite_session.py @@ -489,12 +489,12 @@ def _get_connection(self) -> sqlite3.Connection: @pytest.mark.parametrize("operation", ["add", "pop", "clear"]) -async def test_post_commit_cancellation_returns_known_mutation_outcome( +async def test_post_commit_cancellation_propagates_after_known_mutation_outcome( tmp_path: Path, monkeypatch: pytest.MonkeyPatch, operation: str, ): - """Cancellation after commit must not invite an AdvancedSQLite mutation retry.""" + """Cancellation after commit must propagate without inviting a mutation retry.""" class PausingCommitConnection(sqlite3.Connection): pause_commit = False @@ -542,7 +542,8 @@ def commit(self) -> None: mutation.cancel() await asyncio.sleep(0) conn.allow_return.set() - result = await mutation + with pytest.raises(asyncio.CancelledError): + await mutation finally: conn.allow_return.set() if not mutation.done(): @@ -550,14 +551,12 @@ def commit(self) -> None: await asyncio.gather(mutation, return_exceptions=True) if operation == "add": - assert result is None assert await session.get_items() == [item] elif operation == "pop": - assert result == item assert await session.get_items() == [] else: - assert result is None assert await session.get_items() == [] + assert mutation.cancelled() session.close() diff --git a/tests/extensions/memory/test_async_sqlite_session.py b/tests/extensions/memory/test_async_sqlite_session.py index 5b353ce1f2..a4a6e8d9b8 100644 --- a/tests/extensions/memory/test_async_sqlite_session.py +++ b/tests/extensions/memory/test_async_sqlite_session.py @@ -554,12 +554,12 @@ async def controlled_close() -> None: @pytest.mark.parametrize("operation", ["add", "pop", "clear"]) -async def test_post_commit_cancellation_returns_known_mutation_outcome( +async def test_post_commit_cancellation_propagates_after_known_mutation_outcome( tmp_path: Path, monkeypatch: pytest.MonkeyPatch, operation: str, ): - """Cancellation after async commit must not invite a duplicate mutation retry.""" + """Cancellation after async commit must propagate without inviting a retry.""" db_path = tmp_path / f"async_post_commit_{operation}.db" session = AsyncSQLiteSession(f"async_post_commit_{operation}", db_path) item: TResponseInputItem = {"role": "user", "content": "once"} @@ -596,7 +596,8 @@ async def controlled_commit() -> None: mutation.cancel() await asyncio.sleep(0) allow_return.set() - result = await mutation + with pytest.raises(asyncio.CancelledError): + await mutation finally: allow_return.set() if not mutation.done(): @@ -604,14 +605,12 @@ async def controlled_commit() -> None: await asyncio.gather(mutation, return_exceptions=True) if operation == "add": - assert result is None assert await session.get_items() == [item] elif operation == "pop": - assert result == item assert await session.get_items() == [] else: - assert result is None assert await session.get_items() == [] + assert mutation.cancelled() finally: await session.close() diff --git a/tests/memory/test_session.py b/tests/memory/test_session.py index 60fbc41c10..efc9fe548f 100644 --- a/tests/memory/test_session.py +++ b/tests/memory/test_session.py @@ -1113,12 +1113,12 @@ def close(self) -> None: @pytest.mark.asyncio @pytest.mark.parametrize("operation", ["add", "pop", "clear"]) -async def test_sqlite_session_post_commit_cancellation_returns_known_outcome( +async def test_sqlite_session_post_commit_cancellation_propagates_after_known_outcome( tmp_path: Path, monkeypatch: pytest.MonkeyPatch, operation: str, ): - """Cancellation after a worker commit must not invite a duplicate mutation retry.""" + """Cancellation after a worker commit must propagate without inviting a retry.""" class PausingCommitConnection(sqlite3.Connection): pause_commit = False @@ -1162,7 +1162,8 @@ def commit(self) -> None: mutation.cancel() await asyncio.sleep(0) conn.allow_return.set() - result = await mutation + with pytest.raises(asyncio.CancelledError): + await mutation finally: conn.allow_return.set() if not mutation.done(): @@ -1170,12 +1171,10 @@ def commit(self) -> None: await asyncio.gather(mutation, return_exceptions=True) if operation == "add": - assert result is None assert await session.get_items() == [item] elif operation == "pop": - assert result == item assert await session.get_items() == [] else: - assert result is None assert await session.get_items() == [] + assert mutation.cancelled() session.close() From 4b91149a1f125658222bbeefd93389b0f4dcbc14 Mon Sep 17 00:00:00 2001 From: Kazuhiro Sera Date: Thu, 6 Aug 2026 00:45:59 +0900 Subject: [PATCH 4/7] fix: complete session mutation cleanup --- .../memory/advanced_sqlite_session.py | 20 +- .../extensions/memory/async_sqlite_session.py | 105 ++- src/agents/extensions/memory/redis_session.py | 311 +++++-- .../memory/test_advanced_sqlite_session.py | 105 +++ .../memory/test_async_sqlite_session.py | 113 +++ tests/extensions/memory/test_redis_session.py | 770 +++++++++++++++++- 6 files changed, 1300 insertions(+), 124 deletions(-) diff --git a/src/agents/extensions/memory/advanced_sqlite_session.py b/src/agents/extensions/memory/advanced_sqlite_session.py index 67d86f1cc6..205174263d 100644 --- a/src/agents/extensions/memory/advanced_sqlite_session.py +++ b/src/agents/extensions/memory/advanced_sqlite_session.py @@ -579,7 +579,7 @@ def _add_structure_sync(): conn.commit() try: - await asyncio.to_thread(_add_structure_sync) + await _await_mutation(asyncio.to_thread(_add_structure_sync)) except Exception as exc: log_model_and_tool_action_error( self._logger, @@ -697,7 +697,7 @@ def _cleanup_sync(): conn.commit() return deleted_count - return await asyncio.to_thread(_cleanup_sync) + return await _await_mutation(asyncio.to_thread(_cleanup_sync)) def _cleanup_orphaned_messages_sync(self, conn: sqlite3.Connection) -> int: with closing(conn.cursor()) as cursor: @@ -834,7 +834,9 @@ async def create_branch_from_turn( # Switch to new branch under the lock; skipped if a clear_session has # committed since `generation` was captured (its reset to 'main' wins), # so we never point at a branch that clear removed. - await asyncio.to_thread(self._commit_branch_pointer, branch_name, generation) + await _await_mutation( + asyncio.to_thread(self._commit_branch_pointer, branch_name, generation) + ) if _debug.DONT_LOG_MODEL_DATA: self._logger.debug( @@ -914,7 +916,9 @@ def _validate_branch(): old_branch = self._current_branch_id # Update the pointer under the lock; a no-op if a clear_session has # committed since `generation` was captured (its reset to 'main' wins). - switched = await asyncio.to_thread(self._commit_branch_pointer, branch_id, generation) + switched = await _await_mutation( + asyncio.to_thread(self._commit_branch_pointer, branch_id, generation) + ) if switched: self._logger.info("Switched from branch '%s' to '%s'", old_branch, branch_id) @@ -996,8 +1000,8 @@ def _delete_sync(): return usage_deleted, structure_deleted, orphaned_messages_deleted - usage_deleted, structure_deleted, orphaned_messages_deleted = await asyncio.to_thread( - _delete_sync + usage_deleted, structure_deleted, orphaned_messages_deleted = await _await_mutation( + asyncio.to_thread(_delete_sync) ) self._logger.info( @@ -1242,7 +1246,7 @@ def _copy_sync() -> tuple[str, Any]: conn.commit() return branch_id, turn_content - return await asyncio.to_thread(_copy_sync) + return await _await_mutation(asyncio.to_thread(_copy_sync)) async def get_conversation_turns(self, branch_id: str | None = None) -> list[dict[str, Any]]: """Get user turns with content for easy browsing and branching decisions. @@ -1705,4 +1709,4 @@ def _update_sync(): ) conn.commit() - await asyncio.to_thread(_update_sync) + await _await_mutation(asyncio.to_thread(_update_sync)) diff --git a/src/agents/extensions/memory/async_sqlite_session.py b/src/agents/extensions/memory/async_sqlite_session.py index d9b72f4193..215a668902 100644 --- a/src/agents/extensions/memory/async_sqlite_session.py +++ b/src/agents/extensions/memory/async_sqlite_session.py @@ -2,10 +2,10 @@ import asyncio import json -from collections.abc import AsyncIterator, Awaitable +from collections.abc import AsyncIterator from contextlib import asynccontextmanager from pathlib import Path -from typing import Any, TypeVar, cast +from typing import Any, cast import aiosqlite @@ -18,19 +18,6 @@ ) 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): """Async SQLite-based implementation of session storage. @@ -120,19 +107,41 @@ async def _get_connection(self) -> aiosqlite.Connection: connect_task = asyncio.ensure_future(aiosqlite.connect(str(self.db_path))) try: connection = await asyncio.shield(connect_task) - except BaseException: + except BaseException as acquisition_error: + connection = None + cleanup_cancellation: asyncio.CancelledError | None = None try: - connection = await _await_cleanup(connect_task) + connection = await _await_mutation(connect_task) + except asyncio.CancelledError as exc: + cleanup_cancellation = exc + try: + connection = connect_task.result() + except BaseException: + pass except BaseException: pass - else: + close_error = ( await self._close_owned_connection(connection) + if connection is not None + else None + ) + if isinstance(acquisition_error, asyncio.CancelledError): + raise + if cleanup_cancellation is not None: + raise cleanup_cancellation from None + if isinstance(close_error, asyncio.CancelledError): + raise close_error from None raise + assert connection is not None try: await connection.execute("PRAGMA journal_mode=WAL") await self._init_db_for_connection(connection) - except BaseException: - await self._close_owned_connection(connection) + except BaseException as initialization_error: + close_error = await self._close_owned_connection(connection) + if isinstance(initialization_error, asyncio.CancelledError): + raise + if isinstance(close_error, asyncio.CancelledError): + raise close_error from None raise self._connection = connection @@ -157,31 +166,64 @@ async def _write_connection(self) -> AsyncIterator[aiosqlite.Connection]: async with self._locked_connection() as conn: try: yield conn - except BaseException: + except BaseException as operation_error: + rollback_task = asyncio.create_task(conn.rollback()) + rollback_error: BaseException | None = None + rollback_cancellation: asyncio.CancelledError | None = None try: - await _await_cleanup(conn.rollback()) - except BaseException: - await self._invalidate_connection(conn) + await _await_mutation(rollback_task) + except asyncio.CancelledError as exc: + rollback_cancellation = exc + try: + rollback_task.result() + except BaseException as outcome_error: + rollback_error = outcome_error + except BaseException as exc: + rollback_error = exc + + invalidation_error = None + if rollback_error is not None: + invalidation_error = await self._invalidate_connection(conn) + + if isinstance(operation_error, asyncio.CancelledError): + raise + if rollback_cancellation is not None: + raise rollback_cancellation from None + if isinstance(invalidation_error, asyncio.CancelledError): + raise invalidation_error from None raise - async def _invalidate_connection(self, conn: aiosqlite.Connection) -> None: + async def _invalidate_connection(self, conn: aiosqlite.Connection) -> BaseException | 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 + return close_error async def _close_owned_connection(self, conn: aiosqlite.Connection) -> BaseException | None: """Close an owned connection or retain it for a later cleanup retry.""" + close_task = asyncio.create_task(conn.close()) + cancellation: asyncio.CancelledError | None = None + close_error: BaseException | None = None try: - await _await_cleanup(conn.close()) + await _await_mutation(close_task) + except asyncio.CancelledError as exc: + cancellation = exc + try: + close_task.result() + except BaseException as outcome_error: + close_error = outcome_error except BaseException as exc: + close_error = exc + + if close_error is not None: self._quarantined_connections.add(conn) self._closed = True - return exc - self._quarantined_connections.discard(conn) - return None + else: + self._quarantined_connections.discard(conn) + return cancellation or close_error async def get_items(self, limit: int | None = None) -> list[TResponseInputItem]: """Retrieve the conversation history for this session. @@ -382,7 +424,10 @@ async def close(self) -> None: except asyncio.CancelledError as exc: if cancellation is None: cancellation = exc - close_error = await _await_cleanup(close_task) + try: + close_error = await _await_mutation(close_task) + except asyncio.CancelledError: + close_error = close_task.result() if close_error is None: if self._connection is connection: self._connection = None diff --git a/src/agents/extensions/memory/redis_session.py b/src/agents/extensions/memory/redis_session.py index cbf5874c39..2fe5647184 100644 --- a/src/agents/extensions/memory/redis_session.py +++ b/src/agents/extensions/memory/redis_session.py @@ -24,13 +24,14 @@ import asyncio import json import time +from dataclasses import dataclass from typing import Any from ._optional_imports import raise_optional_dependency_error try: import redis.asyncio as redis - from redis.asyncio import Redis + from redis.asyncio import BlockingConnectionPool, ConnectionPool, Redis from redis.exceptions import ResponseError, WatchError except ImportError as e: raise_optional_dependency_error( @@ -49,60 +50,146 @@ ) -async def _reset_pipeline(pipe: Any) -> asyncio.CancelledError | None: - """Finish releasing a watched pipeline despite repeated caller cancellation.""" - reset_task = asyncio.create_task(pipe.reset()) - cancellation: asyncio.CancelledError | None = None - while True: +@dataclass +class _PipelineAttemptOutcome: + committed: bool + operation_error: BaseException | None + cleanup_error: BaseException | None + settled: bool + + +class _PipelineConnectionPool: + """Track one pipeline's connection release without changing the shared pool.""" + + def __init__(self, pool: Any): + self._pool = pool + self.release_started = False + self.release_completed = False + release_implementation = getattr(pool.release, "__func__", None) + if release_implementation not in { + ConnectionPool.release, + BlockingConnectionPool.release, + }: + raise RuntimeError( + "RedisSession requires the standard redis-py connection pool release method" + ) + self._uses_standard_release = release_implementation is ConnectionPool.release + + def __getattr__(self, name: str) -> Any: + return getattr(self._pool, name) + + async def notify_capacity_available(self) -> None: + if isinstance(self._pool, BlockingConnectionPool): + async with self._pool._condition: + self._pool._condition.notify() + + async def release(self, connection: Any) -> None: + self.release_started = True + was_in_use = connection in self._pool._in_use_connections try: - await asyncio.shield(reset_task) - except asyncio.CancelledError as exc: - if cancellation is None: - cancellation = exc - if reset_task.done(): - reset_task.result() - return cancellation + await self._pool.release(connection) + except BaseException: + if was_in_use and connection in self._pool._available_connections: + self.release_completed = True + await self.notify_capacity_available() + elif ( + was_in_use + and self._uses_standard_release + and connection in self._pool._in_use_connections + ): + self.release_completed = True + raise else: - return cancellation + self.release_completed = True -async def _execute_pipeline( +async def _finish_pipeline( pipe: Any, -) -> tuple[BaseException | None, asyncio.CancelledError | None, bool]: - """Wait for transaction outcome and driver cleanup despite caller cancellation.""" - execute_transaction = pipe._execute_transaction - transaction_succeeded = False - - async def execute_transaction_and_mark(*args: Any, **kwargs: Any) -> Any: - nonlocal transaction_succeeded - result = await execute_transaction(*args, **kwargs) - transaction_succeeded = True - return result - - pipe._execute_transaction = execute_transaction_and_mark - execute_task = asyncio.create_task(pipe.execute()) - cancellation: asyncio.CancelledError | None = None + connection_pool: _PipelineConnectionPool, +) -> tuple[BaseException | None, bool, Any | None]: + """Reset a pipeline and prove whether its connection left pipeline ownership.""" + if connection_pool.release_completed: + pipe.connection = None + return None, True, None + + reset_error: BaseException | None = None try: - while True: - try: - await asyncio.shield(execute_task) - except asyncio.CancelledError as exc: - if cancellation is None: - cancellation = exc - if not execute_task.done(): - continue - except BaseException as exc: - return exc, cancellation, transaction_succeeded - else: - return None, cancellation, transaction_succeeded + await pipe.reset() + except BaseException as exc: + reset_error = exc + + if connection_pool.release_completed: + pipe.connection = None + return reset_error, True, None + + connection = getattr(pipe, "connection", None) + if connection is None: + return reset_error, True, None + if ( + connection_pool.release_started + and connection not in connection_pool._pool._in_use_connections + ): + try: + connection._close() + except BaseException as close_error: + pipe.connection = None + return close_error, True, connection + await connection_pool.notify_capacity_available() + pipe.connection = None + if reset_error is None: + reset_error = RuntimeError("Redis pipeline release did not complete") + return reset_error, True, None + + try: + connection.mark_for_reconnect() + except BaseException as reconnect_error: + return reconnect_error, False, None + + try: + await connection_pool.release(connection) + except BaseException as release_error: + if ( + not connection_pool.release_completed + and connection not in connection_pool._pool._in_use_connections + ): try: - execute_task.result() - except BaseException as exc: - return exc, cancellation, transaction_succeeded - return None, cancellation, transaction_succeeded - finally: - pipe._execute_transaction = execute_transaction + connection._close() + except BaseException as close_error: + pipe.connection = None + return close_error, True, connection + await connection_pool.notify_capacity_available() + connection_pool.release_completed = True + pipe.connection = None + return release_error, connection_pool.release_completed, None + + pipe.connection = None + return reset_error, True, None + + +async def _await_pipeline_attempt( + attempt: asyncio.Task[_PipelineAttemptOutcome], + completion_owned: asyncio.Event, +) -> tuple[_PipelineAttemptOutcome, asyncio.CancelledError | None]: + """Wait for an attempt while preserving only caller-originated cancellation.""" + cancellation: asyncio.CancelledError | None = None + attempt_cancelled = False + + while True: + try: + outcome = await asyncio.shield(attempt) + except asyncio.CancelledError as exc: + if attempt.done() and attempt.cancelled(): + if cancellation is not None: + raise cancellation from None + raise + if cancellation is None: + cancellation = exc + if not completion_owned.is_set() and not attempt_cancelled: + attempt.cancel() + attempt_cancelled = True + else: + return outcome, cancellation class RedisSession(SessionABC): @@ -145,6 +232,7 @@ def __init__( self._owns_client = False # Track if we own the Redis client self._closed = False self._client_released = False + self._detached_connections: set[Any] = set() # Redis key patterns self._session_key = f"{self._key_prefix}:{self.session_id}" @@ -217,15 +305,37 @@ def _key_type_name(key_type: Any) -> str: return key_type.decode("utf-8") return str(key_type) - async def _write_items( + async def _write_items_attempt( self, + pipe: Any, + keys: tuple[str, str, str], serialized_items: list[str], - ) -> None: - """Validate key types and atomically write one batch with optimistic locking.""" - keys = (self._session_key, self._messages_key, self._counter_key) - while True: - pipe = self._redis.pipeline() - committed = False + completion_owned: asyncio.Event, + ) -> _PipelineAttemptOutcome: + """Run one watched write attempt and finish its pipeline before returning.""" + committed = False + operation_error: BaseException | None = None + batch_response_index: int | None = None + raise_first_error = pipe.raise_first_error + raw_connection_pool = pipe.connection_pool + connection_pool = _PipelineConnectionPool(raw_connection_pool) + pipe.connection_pool = connection_pool + + def raise_first_error_and_mark(*args: Any, **kwargs: Any) -> Any: + nonlocal committed + response = args[1] if len(args) > 1 else kwargs.get("response") + if ( + batch_response_index is not None + and isinstance(response, list) + and batch_response_index < len(response) + and not isinstance(response[batch_response_index], BaseException) + ): + committed = True + result = raise_first_error(*args, **kwargs) + committed = True + return result + + try: try: await pipe.watch(*keys) session_key_type = self._key_type_name(await pipe.type(self._session_key)) @@ -254,36 +364,69 @@ async def _write_items( pipe.multi() pipe.hset(self._session_key, "session_id", self.session_id) pipe.hsetnx(self._session_key, "created_at", now) + batch_response_index = len(pipe.command_stack) pipe.rpush(self._messages_key, *serialized_items) pipe.hset(self._session_key, "updated_at", now) if expiration_time_ms is not None: for key in keys: pipe.pexpireat(key, expiration_time_ms) - ( - execute_error, - execute_cancellation, - transaction_succeeded, - ) = await _execute_pipeline(pipe) - if transaction_succeeded: - committed = True - return - if isinstance(execute_error, WatchError): - if execute_cancellation is not None: - raise execute_cancellation - continue - if execute_error is not None: - raise execute_error + + pipe.raise_first_error = raise_first_error_and_mark + completion_owned.set() + await pipe.execute() committed = True + except WatchError as exc: + operation_error = exc + except BaseException as exc: + operation_error = exc + finally: + completion_owned.set() + pipe.raise_first_error = raise_first_error + cleanup_error, settled, detached_connection = await _finish_pipeline( + pipe, connection_pool + ) + if detached_connection is not None: + self._detached_connections.add(detached_connection) + pipe.connection_pool = raw_connection_pool + + return _PipelineAttemptOutcome( + committed=committed, + operation_error=operation_error, + cleanup_error=cleanup_error, + settled=settled, + ) + + async def _write_items( + self, + serialized_items: list[str], + ) -> None: + """Validate key types and atomically write one batch with optimistic locking.""" + keys = (self._session_key, self._messages_key, self._counter_key) + while True: + pipe = self._redis.pipeline() + completion_owned = asyncio.Event() + attempt = asyncio.create_task( + self._write_items_attempt(pipe, keys, serialized_items, completion_owned) + ) + outcome, cancellation = await _await_pipeline_attempt(attempt, completion_owned) + + if not outcome.settled: + if outcome.cleanup_error is not None: + raise outcome.cleanup_error + raise RuntimeError("Redis pipeline cleanup did not settle its connection") + if outcome.committed: + if cancellation is not None: + raise cancellation return - finally: - try: - reset_cancellation = await _reset_pipeline(pipe) - except BaseException: - if not committed: - raise - else: - if reset_cancellation is not None and not committed: - raise reset_cancellation + if cancellation is not None: + raise cancellation + if outcome.cleanup_error is not None: + raise outcome.cleanup_error + if isinstance(outcome.operation_error, WatchError): + continue + if outcome.operation_error is not None: + raise outcome.operation_error + return async def get_items(self, limit: int | None = None) -> list[TResponseInputItem]: """Retrieve the conversation history for this session. @@ -408,12 +551,26 @@ async def close(self) -> None: concurrent calls are safe no-ops. """ async with self._lock: + detached_error: BaseException | None = None + for connection in tuple(self._detached_connections): + try: + connection._close() + except BaseException as exc: + if detached_error is None: + detached_error = exc + else: + self._detached_connections.discard(connection) + if not self._owns_client: + if detached_error is not None: + raise detached_error return self._closed = True if not self._client_released: await self._redis.aclose() self._client_released = True + if detached_error is not None: + raise detached_error async def ping(self) -> bool: """Test Redis connectivity. diff --git a/tests/extensions/memory/test_advanced_sqlite_session.py b/tests/extensions/memory/test_advanced_sqlite_session.py index beef62c9bb..548052c8be 100644 --- a/tests/extensions/memory/test_advanced_sqlite_session.py +++ b/tests/extensions/memory/test_advanced_sqlite_session.py @@ -560,6 +560,111 @@ def commit(self) -> None: session.close() +@pytest.mark.parametrize("operation", ["create_branch", "delete_branch", "cleanup", "usage"]) +async def test_auxiliary_mutation_cancellation_waits_for_commit( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + usage_data: Usage, + operation: str, +): + """Branch and ancillary mutations must settle before cancellation propagates.""" + + class PausingCommitConnection(sqlite3.Connection): + pause_commit = False + commit_finished = threading.Event() + allow_return = threading.Event() + + def commit(self) -> None: + super().commit() + if self.pause_commit: + self.pause_commit = False + self.commit_finished.set() + assert self.allow_return.wait(timeout=10) + + session = AdvancedSQLiteSession( + session_id=f"advanced_auxiliary_cancel_{operation}", + db_path=tmp_path / f"advanced_auxiliary_cancel_{operation}.db", + create_tables=True, + ) + items: list[TResponseInputItem] = [ + {"role": "user", "content": "u1"}, + {"role": "assistant", "content": "a1"}, + {"role": "user", "content": "u2"}, + {"role": "assistant", "content": "a2"}, + ] + mutation: asyncio.Task[Any] | None = None + + try: + if operation in {"create_branch", "delete_branch"}: + await session.add_items(items) + if operation == "delete_branch": + await session.create_branch_from_turn(2, "cancelled_branch") + await session.switch_to_branch("main") + elif operation == "cleanup": + with session._write_connection() as setup_connection: + session._insert_items( + setup_connection, + [{"role": "user", "content": "orphan"}], + ) + setup_connection.commit() + elif operation == "usage": + await session.add_items([{"role": "user", "content": "usage turn"}]) + + connection = sqlite3.connect( + str(session.db_path), + check_same_thread=False, + factory=PausingCommitConnection, + ) + with session._connections_lock: + session._connections.add(connection) + monkeypatch.setattr(session, "_get_connection", lambda: connection) + connection.pause_commit = True + + if operation == "create_branch": + mutation = asyncio.create_task(session.create_branch_from_turn(2, "cancelled_branch")) + elif operation == "delete_branch": + mutation = asyncio.create_task(session.delete_branch("cancelled_branch")) + elif operation == "cleanup": + mutation = asyncio.create_task(session._cleanup_orphaned_messages()) + else: + mutation = asyncio.create_task( + session.store_run_usage(create_mock_run_result(usage_data)) + ) + + assert await asyncio.to_thread(connection.commit_finished.wait, 10) + mutation.cancel("first-caller-cancel") + await asyncio.sleep(0) + mutation.cancel("second-caller-cancel") + await asyncio.sleep(0) + connection.allow_return.set() + + with pytest.raises(asyncio.CancelledError) as exc_info: + await mutation + + assert exc_info.value.args == ("first-caller-cancel",) + assert mutation.cancelled() + + if operation == "create_branch": + branches = await session.list_branches() + assert {branch["branch_id"] for branch in branches} == {"main", "cancelled_branch"} + assert session._current_branch_id == "main" + elif operation == "delete_branch": + branches = await session.list_branches() + assert {branch["branch_id"] for branch in branches} == {"main"} + elif operation == "cleanup": + assert _count_rows(session, session.messages_table) == 0 + else: + turn_usage = await session.get_turn_usage(1) + assert isinstance(turn_usage, dict) + assert turn_usage["total_tokens"] == usage_data.total_tokens + finally: + PausingCommitConnection.allow_return.set() + if mutation is not None and not mutation.done(): + mutation.cancel() + await asyncio.gather(mutation, return_exceptions=True) + session.close() + + async def test_message_structure_tracking(agent: Agent): """Test that message structure is properly tracked.""" session_id = "structure_test" diff --git a/tests/extensions/memory/test_async_sqlite_session.py b/tests/extensions/memory/test_async_sqlite_session.py index a4a6e8d9b8..abaffcc8f2 100644 --- a/tests/extensions/memory/test_async_sqlite_session.py +++ b/tests/extensions/memory/test_async_sqlite_session.py @@ -776,6 +776,61 @@ async def controlled_rollback() -> None: await session.close() +async def test_operation_failure_then_cancellation_during_rollback( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +): + """Cancellation during rollback must supersede an earlier operation failure.""" + db_path = tmp_path / "failure_then_cancelled_rollback.db" + session = AsyncSQLiteSession("failure_then_cancelled_rollback", db_path) + task: asyncio.Task[None] | None = None + try: + await session.get_items() + conn = await session._get_connection() + real_execute = conn.execute + real_rollback = conn.rollback + rollback_started = asyncio.Event() + allow_rollback = asyncio.Event() + + async def fail_after_write(*args: Any, **kwargs: Any) -> Any: + await real_execute(*args, **kwargs) + raise RuntimeError("operation failed") + + async def controlled_rollback() -> None: + rollback_started.set() + await allow_rollback.wait() + await real_rollback() + + monkeypatch.setattr(conn, "execute", fail_after_write) + monkeypatch.setattr(conn, "rollback", controlled_rollback) + task = asyncio.create_task(session.add_items([{"role": "user", "content": "cancelled"}])) + try: + await rollback_started.wait() + task.cancel("first-caller-cancel") + await asyncio.sleep(0) + task.cancel("second-caller-cancel") + allow_rollback.set() + with pytest.raises(asyncio.CancelledError) as exc_info: + await task + finally: + allow_rollback.set() + monkeypatch.setattr(conn, "execute", real_execute) + monkeypatch.setattr(conn, "rollback", real_rollback) + if task is not None and not task.done(): + task.cancel() + await asyncio.gather(task, return_exceptions=True) + + assert exc_info.value.args == ("first-caller-cancel",) + assert conn.in_transaction is False + assert _sqlite_write_lock_is_free(db_path) + assert await session.get_items() == [] + finally: + if task is not None and not task.done(): + task.cancel() + await asyncio.gather(task, return_exceptions=True) + await session.close() + + async def test_rollback_failure_closes_and_evicts_connection( tmp_path: Path, monkeypatch: pytest.MonkeyPatch, @@ -961,6 +1016,64 @@ async def controlled_close() -> None: await session.real_close() +async def test_initialization_failure_then_cancellation_during_candidate_close( + monkeypatch: pytest.MonkeyPatch, +): + """Cancellation during candidate close must supersede an initialization failure.""" + close_started = asyncio.Event() + allow_close = asyncio.Event() + + class FailingInitSession(AsyncSQLiteSession): + captured_connection: Any = None + real_close: Any = None + + async def _init_db_for_connection(self, conn: Any) -> None: + self.captured_connection = conn + self.real_close = conn.close + + async def controlled_close() -> None: + close_started.set() + await allow_close.wait() + await self.real_close() + + monkeypatch.setattr(conn, "close", controlled_close) + raise RuntimeError("initialization failed") + + session = FailingInitSession("failed_init_then_cancelled_close") + task = asyncio.create_task(session.get_items()) + try: + try: + await close_started.wait() + task.cancel("first-caller-cancel") + await asyncio.sleep(0) + task.cancel("second-caller-cancel") + allow_close.set() + with pytest.raises(asyncio.CancelledError) as exc_info: + await task + finally: + allow_close.set() + if not task.done(): + task.cancel() + await asyncio.gather(task, return_exceptions=True) + + assert exc_info.value.args == ("first-caller-cancel",) + assert session._connection is None + assert session.captured_connection._running is False + finally: + allow_close.set() + if session.captured_connection is not None and session.real_close is not None: + monkeypatch.setattr(session.captured_connection, "close", session.real_close) + try: + await session.close() + finally: + if ( + session.captured_connection is not None + and session.real_close is not None + and session.captured_connection._running + ): + await session.real_close() + + async def test_cancelled_connect_closes_eventually_acquired_connection( monkeypatch: pytest.MonkeyPatch, ): diff --git a/tests/extensions/memory/test_redis_session.py b/tests/extensions/memory/test_redis_session.py index 901019c145..6fef22e38c 100644 --- a/tests/extensions/memory/test_redis_session.py +++ b/tests/extensions/memory/test_redis_session.py @@ -1532,12 +1532,12 @@ async def reset(self) -> None: add_task = asyncio.create_task(session.add_items([{"role": "user", "content": "cancelled"}])) try: await time_started.wait() - add_task.cancel() + add_task.cancel("first-pre-exec-cancel") await reset_started.wait() - add_task.cancel() + add_task.cancel("second-during-reset") await asyncio.sleep(0) allow_reset.set() - with pytest.raises(asyncio.CancelledError): + with pytest.raises(asyncio.CancelledError) as exc_info: await add_task finally: allow_reset.set() @@ -1545,19 +1545,26 @@ async def reset(self) -> None: add_task.cancel() await asyncio.gather(add_task, return_exceptions=True) + assert exc_info.value.args == ("first-pre-exec-cancel",) + assert await session.get_items() == [] assert await client.ping() is True # type: ignore[misc] -async def test_post_commit_cancellation_does_not_report_retryable_failure( +@pytest.mark.parametrize( + "cleanup_mode", + ["internal_reset", "failing_internal_reset"], +) +async def test_post_commit_cancellation_propagates_after_cleanup( monkeypatch: pytest.MonkeyPatch, + cleanup_mode: str, ) -> None: - """Cancellation during driver cleanup must not hide a known successful commit.""" + """Cancellation after EXEC must propagate once driver cleanup settles.""" if not USE_FAKE_REDIS: pytest.skip("This test requires fakeredis") client = fakeredis.aioredis.FakeRedis(max_connections=1) session = RedisSession( - session_id="post_commit_cancellation", + session_id=f"post_commit_cancellation_{cleanup_mode}", redis_client=cast("Redis", client), key_prefix="test:", ) @@ -1577,6 +1584,8 @@ async def controlled_reset() -> None: internal_reset_started.set() await allow_internal_reset.wait() await real_reset() + if cleanup_mode == "failing_internal_reset" and reset_calls == 1: + raise RuntimeError("post-commit reset failed") monkeypatch.setattr(delegate, "reset", controlled_reset) return delegate @@ -1586,12 +1595,13 @@ async def controlled_reset() -> None: add_task = asyncio.create_task(session.add_items([item])) try: await internal_reset_started.wait() - add_task.cancel() + add_task.cancel("first-post-commit-cancel") await asyncio.sleep(0) - add_task.cancel() + add_task.cancel("second-during-cleanup") await asyncio.sleep(0) allow_internal_reset.set() - await add_task + with pytest.raises(asyncio.CancelledError) as exc_info: + await add_task finally: allow_internal_reset.set() if not add_task.done(): @@ -1599,7 +1609,749 @@ async def controlled_reset() -> None: await asyncio.gather(add_task, return_exceptions=True) assert await session.get_items() == [item] + assert exc_info.value.args == ("first-post-commit-cancel",) + assert add_task.cancelled() + assert await client.ping() is True # type: ignore[misc] + + +async def test_post_commit_response_callback_failure_does_not_retry( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """A response callback failure after EXEC must not duplicate the committed batch.""" + if not USE_FAKE_REDIS: + pytest.skip("This test requires fakeredis") + + client = fakeredis.aioredis.FakeRedis(max_connections=1) + session = RedisSession( + session_id="post_commit_callback_failure", + redis_client=cast("Redis", client), + key_prefix="test:", + ) + + def fail_rpush_callback(response: Any, **kwargs: Any) -> Any: + raise RuntimeError("response callback failed") + + client.set_response_callback("RPUSH", fail_rpush_callback) + item: TResponseInputItem = {"role": "user", "content": "committed"} + + await session.add_items([item]) + + monkeypatch.delitem(client.response_callbacks, "RPUSH") + assert await session.get_items() == [item] + assert await client.ping() is True # type: ignore[misc] + + +async def test_successful_batch_response_with_sibling_error_does_not_retry( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """A successful RPUSH remains committed when a sibling EXEC response fails.""" + if not USE_FAKE_REDIS: + pytest.skip("This test requires fakeredis") + + client = fakeredis.aioredis.FakeRedis(max_connections=1) + session = RedisSession( + session_id="successful_batch_with_sibling_error", + redis_client=cast("Redis", client), + key_prefix="test:", + ) + real_pipeline = client.pipeline + + def pipeline(*args: Any, **kwargs: Any) -> Any: + delegate = real_pipeline(*args, **kwargs) + real_hset = delegate.hset + hset_calls = 0 + + def replace_updated_at_with_wrong_type(*args: Any, **kwargs: Any) -> Any: + nonlocal hset_calls + hset_calls += 1 + if hset_calls == 2: + return delegate.incr(session._session_key) + return real_hset(*args, **kwargs) + + monkeypatch.setattr(delegate, "hset", replace_updated_at_with_wrong_type) + return delegate + + monkeypatch.setattr(client, "pipeline", pipeline) + item: TResponseInputItem = {"role": "user", "content": "committed"} + + await session.add_items([item]) + + monkeypatch.setattr(client, "pipeline", real_pipeline) + assert await session.get_items() == [item] + assert await client.ping() is True # type: ignore[misc] + + +async def test_post_commit_reset_self_cancellation_does_not_cancel_caller( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """A reset child cancellation after EXEC must not impersonate caller cancellation.""" + if not USE_FAKE_REDIS: + pytest.skip("This test requires fakeredis") + + client = fakeredis.aioredis.FakeRedis(max_connections=1) + session = RedisSession( + session_id="post_commit_reset_self_cancel", + redis_client=cast("Redis", client), + key_prefix="test:", + ) + real_pipeline = client.pipeline + + def pipeline(*args: Any, **kwargs: Any) -> Any: + delegate = real_pipeline(*args, **kwargs) + real_reset = delegate.reset + reset_calls = 0 + + async def controlled_reset() -> None: + nonlocal reset_calls + reset_calls += 1 + if reset_calls == 1: + raise asyncio.CancelledError("reset self-cancelled") + await real_reset() + + monkeypatch.setattr(delegate, "reset", controlled_reset) + return delegate + + monkeypatch.setattr(client, "pipeline", pipeline) + item: TResponseInputItem = {"role": "user", "content": "committed"} + + await session.add_items([item]) + + assert await session.get_items() == [item] + assert await client.ping() is True # type: ignore[misc] + + +async def test_post_commit_repeated_reset_failure_force_releases_connection( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Repeated reset failures after EXEC must not leak the pipeline connection.""" + if not USE_FAKE_REDIS: + pytest.skip("This test requires fakeredis") + + client = fakeredis.aioredis.FakeRedis(max_connections=1) + session = RedisSession( + session_id="post_commit_repeated_reset_failure", + redis_client=cast("Redis", client), + key_prefix="test:", + ) + real_pipeline = client.pipeline + + def pipeline(*args: Any, **kwargs: Any) -> Any: + delegate = real_pipeline(*args, **kwargs) + real_reset = delegate.reset + reset_calls = 0 + + async def fail_two_resets_before_release() -> None: + nonlocal reset_calls + reset_calls += 1 + if reset_calls <= 2: + raise RuntimeError("reset failed before release") + await real_reset() + + monkeypatch.setattr(delegate, "reset", fail_two_resets_before_release) + return delegate + + monkeypatch.setattr(client, "pipeline", pipeline) + item: TResponseInputItem = {"role": "user", "content": "committed"} + + await session.add_items([item]) + + assert await session.get_items() == [item] + assert await client.ping() is True # type: ignore[misc] + + +async def test_post_commit_release_listener_failure_does_not_retry( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """A listener failure after pool release must not make a committed batch retryable.""" + if not USE_FAKE_REDIS: + pytest.skip("This test requires fakeredis") + + from redis.event import AsyncAfterConnectionReleasedEvent + + client = fakeredis.aioredis.FakeRedis(max_connections=1) + session = RedisSession( + session_id="post_commit_release_listener_failure", + redis_client=cast("Redis", client), + key_prefix="test:", + ) + dispatcher = client.connection_pool._event_dispatcher + assert dispatcher is not None + listeners = dispatcher._event_listeners_mapping[AsyncAfterConnectionReleasedEvent] + + class FailingReleaseListener: + async def listen(self, event: Any) -> None: + raise RuntimeError("release listener failed after pool return") + + monkeypatch.setitem( + dispatcher._event_listeners_mapping, + AsyncAfterConnectionReleasedEvent, + [*listeners, FailingReleaseListener()], + ) + item: TResponseInputItem = {"role": "user", "content": "committed"} + + await session.add_items([item]) + + monkeypatch.setitem( + dispatcher._event_listeners_mapping, + AsyncAfterConnectionReleasedEvent, + listeners, + ) + assert await session.get_items() == [item] + assert await client.ping() is True # type: ignore[misc] + + +async def test_pre_exec_disconnect_failure_does_not_reuse_watched_connection( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """A failed disconnect must remove the watched connection from pool reuse.""" + if not USE_FAKE_REDIS: + pytest.skip("This test requires fakeredis") + + client = fakeredis.aioredis.FakeRedis(max_connections=1) + session = RedisSession( + session_id="pre_exec_disconnect_failure", + redis_client=cast("Redis", client), + key_prefix="test:", + ttl=60, + ) + real_pipeline = client.pipeline + time_started = asyncio.Event() + + def pipeline(*args: Any, **kwargs: Any) -> Any: + delegate = real_pipeline(*args, **kwargs) + + async def controlled_time() -> tuple[int, int]: + time_started.set() + await asyncio.Event().wait() + raise AssertionError("unreachable") + + async def fail_reset_before_release() -> None: + connection = delegate.connection + assert connection is not None + + async def fail_disconnect(nowait: bool = False) -> None: + raise RuntimeError("disconnect failed") + + monkeypatch.setattr(connection, "disconnect", fail_disconnect) + raise RuntimeError("reset failed before release") + + monkeypatch.setattr(delegate, "time", controlled_time) + monkeypatch.setattr(delegate, "reset", fail_reset_before_release) + return delegate + + monkeypatch.setattr(client, "pipeline", pipeline) + add_task = asyncio.create_task(session.add_items([{"role": "user", "content": "cancelled"}])) + await time_started.wait() + add_task.cancel("caller-cancel") + + with pytest.raises(asyncio.CancelledError) as exc_info: + await add_task + + assert exc_info.value.args == ("caller-cancel",) + monkeypatch.setattr(client, "pipeline", real_pipeline) assert await client.ping() is True # type: ignore[misc] + assert await session.get_items() == [] + + +async def test_pre_exec_cancellation_wins_after_released_reset_failure( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Caller cancellation must win once a failing reset has released its connection.""" + if not USE_FAKE_REDIS: + pytest.skip("This test requires fakeredis") + + client = fakeredis.aioredis.FakeRedis(max_connections=1) + session = RedisSession( + session_id="pre_exec_cancel_released_reset_failure", + redis_client=cast("Redis", client), + key_prefix="test:", + ttl=60, + ) + real_pipeline = client.pipeline + time_started = asyncio.Event() + + def pipeline(*args: Any, **kwargs: Any) -> Any: + delegate = real_pipeline(*args, **kwargs) + real_reset = delegate.reset + + async def controlled_time() -> tuple[int, int]: + time_started.set() + await asyncio.Event().wait() + raise AssertionError("unreachable") + + async def fail_reset_after_release() -> None: + await real_reset() + raise RuntimeError("post-release reset callback failed") + + monkeypatch.setattr(delegate, "time", controlled_time) + monkeypatch.setattr(delegate, "reset", fail_reset_after_release) + return delegate + + monkeypatch.setattr(client, "pipeline", pipeline) + add_task = asyncio.create_task(session.add_items([{"role": "user", "content": "cancelled"}])) + await time_started.wait() + add_task.cancel("first-caller-cancel") + + with pytest.raises(asyncio.CancelledError) as exc_info: + await add_task + + assert exc_info.value.args == ("first-caller-cancel",) + assert await client.ping() is True # type: ignore[misc] + assert await session.get_items() == [] + + +async def test_blocking_pool_cancellation_completion_owns_release( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Caller cancellation must not interrupt blocking-pool cleanup after it starts.""" + if not USE_FAKE_REDIS: + pytest.skip("This test requires fakeredis") + + from redis.asyncio import BlockingConnectionPool + + seed_client = fakeredis.aioredis.FakeRedis() + source_pool = seed_client.connection_pool + pool = BlockingConnectionPool( + max_connections=1, + timeout=1, + connection_class=source_pool.connection_class, + **source_pool.connection_kwargs, + ) + client = fakeredis.aioredis.FakeRedis(connection_pool=pool) + session = RedisSession( + session_id="blocking_pool_cancelled_release", + redis_client=cast("Redis", client), + key_prefix="test:", + ttl=60, + ) + real_pipeline = client.pipeline + time_started = asyncio.Event() + fail_time = asyncio.Event() + reset_started = asyncio.Event() + + def pipeline(*args: Any, **kwargs: Any) -> Any: + delegate = real_pipeline(*args, **kwargs) + real_reset = delegate.reset + + async def controlled_time() -> tuple[int, int]: + time_started.set() + await fail_time.wait() + raise RuntimeError("pre-execute failure") + + async def observed_reset() -> None: + reset_started.set() + await real_reset() + + monkeypatch.setattr(delegate, "time", controlled_time) + monkeypatch.setattr(delegate, "reset", observed_reset) + return delegate + + monkeypatch.setattr(client, "pipeline", pipeline) + add_task = asyncio.create_task(session.add_items([{"role": "user", "content": "failed"}])) + await time_started.wait() + await pool._condition.acquire() + try: + fail_time.set() + await reset_started.wait() + add_task.cancel("caller-cancel") + await asyncio.sleep(0) + finally: + pool._condition.release() + + with pytest.raises(asyncio.CancelledError) as exc_info: + await add_task + + assert exc_info.value.args == ("caller-cancel",) + assert not pool._in_use_connections + monkeypatch.setattr(client, "pipeline", real_pipeline) + assert await client.ping() is True # type: ignore[misc] + assert await session.get_items() == [] + + +async def test_ordinary_pool_cancellation_completion_owns_release( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Caller cancellation must not interrupt ordinary-pool cleanup after it starts.""" + if not USE_FAKE_REDIS: + pytest.skip("This test requires fakeredis") + + client = fakeredis.aioredis.FakeRedis(max_connections=1) + pool = client.connection_pool + session = RedisSession( + session_id="ordinary_pool_cancelled_release", + redis_client=cast("Redis", client), + key_prefix="test:", + ttl=60, + ) + real_pipeline = client.pipeline + time_started = asyncio.Event() + fail_time = asyncio.Event() + reset_started = asyncio.Event() + allow_reset = asyncio.Event() + + def pipeline(*args: Any, **kwargs: Any) -> Any: + delegate = real_pipeline(*args, **kwargs) + real_reset = delegate.reset + + async def controlled_time() -> tuple[int, int]: + time_started.set() + await fail_time.wait() + raise RuntimeError("pre-execute failure") + + async def controlled_reset() -> None: + reset_started.set() + await allow_reset.wait() + await real_reset() + + monkeypatch.setattr(delegate, "time", controlled_time) + monkeypatch.setattr(delegate, "reset", controlled_reset) + return delegate + + monkeypatch.setattr(client, "pipeline", pipeline) + add_task = asyncio.create_task(session.add_items([{"role": "user", "content": "failed"}])) + try: + await time_started.wait() + fail_time.set() + await reset_started.wait() + add_task.cancel("caller-cancel") + await asyncio.sleep(0) + allow_reset.set() + with pytest.raises(asyncio.CancelledError) as exc_info: + await add_task + finally: + allow_reset.set() + if not add_task.done(): + add_task.cancel() + await asyncio.gather(add_task, return_exceptions=True) + + assert exc_info.value.args == ("caller-cancel",) + assert not pool._in_use_connections + monkeypatch.setattr(client, "pipeline", real_pipeline) + assert await client.ping() is True # type: ignore[misc] + assert await session.get_items() == [] + + +async def test_overridden_pool_release_is_rejected_before_mutation( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Ambiguous pool ownership semantics must be rejected before acquiring a connection.""" + if not USE_FAKE_REDIS: + pytest.skip("This test requires fakeredis") + + client = fakeredis.aioredis.FakeRedis(max_connections=1) + pool = client.connection_pool + session = RedisSession( + session_id="overridden_pool_release", + redis_client=cast("Redis", client), + key_prefix="test:", + ) + real_release = pool.release + release_calls = 0 + + async def transfer_reborrow_then_fail(connection: Any) -> None: + nonlocal release_calls + release_calls += 1 + await real_release(connection) + borrowed = await pool.get_connection() + assert borrowed is connection + raise RuntimeError("release failed after reborrow") + + monkeypatch.setattr(pool, "release", transfer_reborrow_then_fail) + item: TResponseInputItem = {"role": "user", "content": "not committed"} + + with pytest.raises(RuntimeError, match="standard redis-py connection pool release method"): + await session.add_items([item]) + + assert release_calls == 0 + assert not pool._in_use_connections + monkeypatch.setattr(pool, "release", real_release) + assert await session.get_items() == [] + + +async def test_post_commit_disconnect_failure_is_not_retryable( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """A detached failed-disconnect connection must not make a committed batch retryable.""" + if not USE_FAKE_REDIS: + pytest.skip("This test requires fakeredis") + + client = fakeredis.aioredis.FakeRedis(max_connections=1) + session = RedisSession( + session_id="post_commit_disconnect_failure", + redis_client=cast("Redis", client), + key_prefix="test:", + ) + real_pipeline = client.pipeline + + def pipeline(*args: Any, **kwargs: Any) -> Any: + delegate = real_pipeline(*args, **kwargs) + real_reset = delegate.reset + reset_calls = 0 + + async def fail_disconnect_after_commit() -> None: + nonlocal reset_calls + reset_calls += 1 + if reset_calls == 1: + connection = delegate.connection + assert connection is not None + connection.mark_for_reconnect() + + async def fail_disconnect(nowait: bool = False) -> None: + raise RuntimeError("disconnect failed") + + monkeypatch.setattr(connection, "disconnect", fail_disconnect) + await real_reset() + + monkeypatch.setattr(delegate, "reset", fail_disconnect_after_commit) + return delegate + + monkeypatch.setattr(client, "pipeline", pipeline) + item: TResponseInputItem = {"role": "user", "content": "committed"} + + await session.add_items([item]) + + monkeypatch.setattr(client, "pipeline", real_pipeline) + assert await session.get_items() == [item] + assert await client.ping() is True # type: ignore[misc] + + +@pytest.mark.parametrize("cancelled", [False, True]) +async def test_post_commit_detached_close_failure_is_quarantined( + monkeypatch: pytest.MonkeyPatch, + cancelled: bool, +) -> None: + """A non-reusable connection close failure must not make a commit retryable.""" + if not USE_FAKE_REDIS: + pytest.skip("This test requires fakeredis") + + client = fakeredis.aioredis.FakeRedis(max_connections=1) + pool = client.connection_pool + session = RedisSession( + session_id=f"post_commit_detached_close_failure_{cancelled}", + redis_client=cast("Redis", client), + key_prefix="test:", + ) + real_pipeline = client.pipeline + reset_started = asyncio.Event() + allow_reset = asyncio.Event() + detached_connection: Any = None + real_close: Any = None + + def pipeline(*args: Any, **kwargs: Any) -> Any: + delegate = real_pipeline(*args, **kwargs) + real_reset = delegate.reset + + async def fail_disconnect_and_close_after_commit() -> None: + nonlocal detached_connection, real_close + reset_started.set() + await allow_reset.wait() + if detached_connection is None: + detached_connection = delegate.connection + assert detached_connection is not None + real_close = detached_connection._close + detached_connection.mark_for_reconnect() + + async def fail_disconnect(nowait: bool = False) -> None: + raise RuntimeError("disconnect failed") + + def fail_close() -> None: + raise RuntimeError("close failed") + + monkeypatch.setattr(detached_connection, "disconnect", fail_disconnect) + monkeypatch.setattr(detached_connection, "_close", fail_close) + await real_reset() + + monkeypatch.setattr(delegate, "reset", fail_disconnect_and_close_after_commit) + return delegate + + monkeypatch.setattr(client, "pipeline", pipeline) + item: TResponseInputItem = {"role": "user", "content": "committed"} + add_task = asyncio.create_task(session.add_items([item])) + try: + await reset_started.wait() + if cancelled: + add_task.cancel("caller-cancel") + await asyncio.sleep(0) + allow_reset.set() + if cancelled: + with pytest.raises(asyncio.CancelledError) as exc_info: + await add_task + assert exc_info.value.args == ("caller-cancel",) + else: + await add_task + + assert not pool._in_use_connections + assert session._detached_connections == {detached_connection} + monkeypatch.setattr(detached_connection, "_close", real_close) + await session.close() + assert session._detached_connections == set() + assert await session.get_items() == [item] + assert await client.ping() is True # type: ignore[misc] + finally: + allow_reset.set() + if not add_task.done(): + add_task.cancel() + await asyncio.gather(add_task, return_exceptions=True) + if detached_connection is not None and real_close is not None: + monkeypatch.setattr(detached_connection, "_close", real_close) + await session.close() + + +async def test_blocking_pool_listener_failure_notifies_waiter( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """A completed blocking-pool release must notify a waiter despite listener failure.""" + if not USE_FAKE_REDIS: + pytest.skip("This test requires fakeredis") + + from redis.asyncio import BlockingConnectionPool + from redis.event import AsyncAfterConnectionReleasedEvent + + seed_client = fakeredis.aioredis.FakeRedis() + source_pool = seed_client.connection_pool + pool = BlockingConnectionPool( + max_connections=1, + timeout=1, + connection_class=source_pool.connection_class, + **source_pool.connection_kwargs, + ) + client = fakeredis.aioredis.FakeRedis(connection_pool=pool) + session = RedisSession( + session_id="blocking_pool_listener_failure", + redis_client=cast("Redis", client), + key_prefix="test:", + ) + dispatcher = pool._event_dispatcher + assert dispatcher is not None + listeners = dispatcher._event_listeners_mapping[AsyncAfterConnectionReleasedEvent] + real_pipeline = client.pipeline + reset_started = asyncio.Event() + allow_reset = asyncio.Event() + + class FailFirstReleaseListener: + def __init__(self) -> None: + self.calls = 0 + + async def listen(self, event: Any) -> None: + self.calls += 1 + if self.calls == 1: + raise RuntimeError("release listener failed") + + def pipeline(*args: Any, **kwargs: Any) -> Any: + delegate = real_pipeline(*args, **kwargs) + real_reset = delegate.reset + + async def controlled_reset() -> None: + reset_started.set() + await allow_reset.wait() + await real_reset() + + monkeypatch.setattr(delegate, "reset", controlled_reset) + return delegate + + monkeypatch.setitem( + dispatcher._event_listeners_mapping, + AsyncAfterConnectionReleasedEvent, + [*listeners, FailFirstReleaseListener()], + ) + monkeypatch.setattr(client, "pipeline", pipeline) + item: TResponseInputItem = {"role": "user", "content": "committed"} + add_task = asyncio.create_task(session.add_items([item])) + ping_task: asyncio.Task[Any] | None = None + + async def ping() -> Any: + return await client.ping() # type: ignore[misc] + + try: + await reset_started.wait() + ping_task = asyncio.create_task(ping()) + await asyncio.sleep(0) + allow_reset.set() + await add_task + assert await ping_task is True + finally: + allow_reset.set() + pending = [task for task in (add_task, ping_task) if task is not None and not task.done()] + for task in pending: + task.cancel() + if pending: + await asyncio.gather(*pending, return_exceptions=True) + + monkeypatch.setattr(client, "pipeline", real_pipeline) + assert await session.get_items() == [item] + + +async def test_blocking_pool_detached_disconnect_notifies_waiter( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Closing a detached blocking-pool connection must notify an existing waiter.""" + if not USE_FAKE_REDIS: + pytest.skip("This test requires fakeredis") + + from redis.asyncio import BlockingConnectionPool + + seed_client = fakeredis.aioredis.FakeRedis() + source_pool = seed_client.connection_pool + pool = BlockingConnectionPool( + max_connections=1, + timeout=1, + connection_class=source_pool.connection_class, + **source_pool.connection_kwargs, + ) + client = fakeredis.aioredis.FakeRedis(connection_pool=pool) + session = RedisSession( + session_id="blocking_pool_detached_disconnect", + redis_client=cast("Redis", client), + key_prefix="test:", + ) + real_pipeline = client.pipeline + reset_started = asyncio.Event() + allow_reset = asyncio.Event() + + def pipeline(*args: Any, **kwargs: Any) -> Any: + delegate = real_pipeline(*args, **kwargs) + real_reset = delegate.reset + + async def fail_disconnect_after_commit() -> None: + reset_started.set() + await allow_reset.wait() + connection = delegate.connection + assert connection is not None + connection.mark_for_reconnect() + + async def fail_disconnect(nowait: bool = False) -> None: + raise RuntimeError("disconnect failed") + + monkeypatch.setattr(connection, "disconnect", fail_disconnect) + await real_reset() + + monkeypatch.setattr(delegate, "reset", fail_disconnect_after_commit) + return delegate + + async def ping() -> Any: + return await client.ping() # type: ignore[misc] + + monkeypatch.setattr(client, "pipeline", pipeline) + item: TResponseInputItem = {"role": "user", "content": "committed"} + add_task = asyncio.create_task(session.add_items([item])) + ping_task: asyncio.Task[Any] | None = None + try: + await reset_started.wait() + ping_task = asyncio.create_task(ping()) + await asyncio.sleep(0) + allow_reset.set() + await add_task + assert await ping_task is True + finally: + allow_reset.set() + pending = [task for task in (add_task, ping_task) if task is not None and not task.done()] + for task in pending: + task.cancel() + if pending: + await asyncio.gather(*pending, return_exceptions=True) + + monkeypatch.setattr(client, "pipeline", real_pipeline) + assert await session.get_items() == [item] async def test_post_commit_internal_reset_failure_does_not_report_write_failure( From 4f2c1cd764ed6006252155616c440bfa449653ce Mon Sep 17 00:00:00 2001 From: Kazuhiro Sera Date: Thu, 6 Aug 2026 09:19:05 +0900 Subject: [PATCH 5/7] fix: complete session mutation atomicity --- .../memory/advanced_sqlite_session.py | 171 +++++++--- src/agents/extensions/memory/dapr_session.py | 156 +++++----- .../extensions/memory/mongodb_session.py | 130 +++++--- src/agents/extensions/memory/redis_session.py | 104 ++++--- .../extensions/memory/sqlalchemy_session.py | 113 ++++--- .../memory/test_advanced_sqlite_session.py | 203 +++++++++++- .../memory/test_async_sqlite_session.py | 11 +- tests/extensions/memory/test_dapr_session.py | 62 ++++ .../extensions/memory/test_mongodb_session.py | 220 +++++++++++-- tests/extensions/memory/test_redis_session.py | 138 ++++++-- .../memory/test_sqlalchemy_session.py | 294 +++++++++++++----- 11 files changed, 1213 insertions(+), 389 deletions(-) diff --git a/src/agents/extensions/memory/advanced_sqlite_session.py b/src/agents/extensions/memory/advanced_sqlite_session.py index 205174263d..f9946e0fc6 100644 --- a/src/agents/extensions/memory/advanced_sqlite_session.py +++ b/src/agents/extensions/memory/advanced_sqlite_session.py @@ -77,11 +77,9 @@ def __init__( pass raise self._current_branch_id = "main" - # Bumped (under the connection lock) whenever clear_session() wipes the - # session. switch_to_branch / create_branch_from_turn capture the - # generation before their DB work and only update the branch pointer if - # no clear has committed since, so a stale switch/create cannot resurrect - # a branch that clear already removed. + # Synchronized with the durable session_clear_generations row whenever a + # branch pointer is established or a write begins. A mismatch means + # another instance cleared the session, so the local pointer resets to main. self._generation = 0 self._logger = logger or logging.getLogger(__name__) @@ -93,17 +91,29 @@ def _commit_branch_pointer(self, branch_id: str, generation: int) -> bool: updated, False if a clear_session committed after ``generation`` was captured (in which case its reset to 'main' wins). """ - with self._lock: - if self._generation != generation: + with self._locked_connection() as conn: + row = conn.execute( + """ + SELECT generation FROM session_clear_generations + WHERE session_id = ? + """, + (self.session_id,), + ).fetchone() + durable_generation = row[0] if row is not None else 0 + if durable_generation != generation: + self._generation = durable_generation + self._current_branch_id = "main" return False + self._generation = durable_generation self._current_branch_id = branch_id return True def _init_structure_tables(self): """Add structure and usage tracking tables. - Creates the message_structure, branch_reservations, and turn_usage tables - with appropriate indexes for conversation branching and usage analytics. + Creates the message_structure, branch_reservations, session_clear_generations, + and turn_usage tables with appropriate indexes for conversation branching + and usage analytics. """ with self._write_connection() as conn: # Message structure with branch support @@ -147,6 +157,7 @@ def _init_structure_tables(self): """) self._ensure_branch_reservations_table(conn) + self._ensure_session_clear_generations_table(conn) # Indexes conn.execute(""" @@ -187,6 +198,7 @@ async def add_items(self, items: list[TResponseInputItem]) -> None: def _add_items_sync(): """Synchronous helper to add items and structure metadata together.""" with self._write_connection() as conn: + self._refresh_branch_after_external_clear(conn) # Keep both writes in one transaction so metadata failures do not leave orphans. self._insert_items(conn, items) self._insert_structure_metadata(conn, items) @@ -303,9 +315,14 @@ async def pop_item(self) -> TResponseInputItem | None: # switch_to_branch() cannot redirect this pop to a different branch once # it has been dispatched to the worker thread. branch_id = self._current_branch_id + generation = self._generation def _pop_item_sync(): with self._write_connection() as conn: + self._refresh_branch_after_external_clear(conn) + resolved_branch_id = ( + self._current_branch_id if self._generation != generation else branch_id + ) while True: with closing(conn.cursor()) as cursor: # Preserve every legacy branch ID before a pop can remove its @@ -325,7 +342,7 @@ def _pop_item_sync(): ) RETURNING message_id, user_turn_number """, - (self.session_id, branch_id), + (self.session_id, resolved_branch_id), ) claimed_row = cursor.fetchone() if claimed_row is None: @@ -351,7 +368,7 @@ def _pop_item_sync(): WHERE session_id = ? AND branch_id = ? AND user_turn_number = ? """, - (self.session_id, branch_id, user_turn_number), + (self.session_id, resolved_branch_id, user_turn_number), ) if cursor.fetchone()[0] == 0: cursor.execute( @@ -360,7 +377,7 @@ def _pop_item_sync(): WHERE session_id = ? AND branch_id = ? AND user_turn_number = ? """, - (self.session_id, branch_id, user_turn_number), + (self.session_id, resolved_branch_id, user_turn_number), ) conn.commit() @@ -395,6 +412,7 @@ def _clear_session_sync(): # Backfill legacy branch IDs before clearing their only durable # identity evidence. self._ensure_branch_reservations_table(conn) + self._ensure_session_clear_generations_table(conn) conn.execute( f"DELETE FROM {self.messages_table} WHERE session_id = ?", (self.session_id,), @@ -411,6 +429,21 @@ def _clear_session_sync(): "DELETE FROM turn_usage WHERE session_id = ?", (self.session_id,), ) + conn.execute( + """ + UPDATE session_clear_generations + SET generation = generation + 1 + WHERE session_id = ? + """, + (self.session_id,), + ) + generation = conn.execute( + """ + SELECT generation FROM session_clear_generations + WHERE session_id = ? + """, + (self.session_id,), + ).fetchone()[0] conn.commit() # All branches were removed, so reset the in-memory pointer to # 'main' while still holding the lock. Doing this inside the @@ -419,7 +452,7 @@ def _clear_session_sync(): # the pointer still references a deleted branch. Bumping the # generation invalidates any in-flight switch/create that # captured the pre-clear generation. - self._generation += 1 + self._generation = generation self._current_branch_id = "main" await _await_mutation(asyncio.to_thread(_clear_session_sync)) @@ -819,41 +852,43 @@ async def create_branch_from_turn( ValueError: If turn doesn't exist, doesn't contain a user message, or `branch_name` has already been used in this session """ - # Snapshot the source branch and clear generation together. The source turn is - # revalidated inside the reservation transaction below. - with self._lock: - generation = self._generation - source_branch_id = self._current_branch_id - - # Resolve the target branch ID under the same transaction that performs the copy - # so concurrent creators cannot reserve the same branch. - branch_name, turn_content = await self._copy_messages_to_new_branch( - branch_name, turn_number, source_branch_id - ) - # Switch to new branch under the lock; skipped if a clear_session has - # committed since `generation` was captured (its reset to 'main' wins), - # so we never point at a branch that clear removed. - await _await_mutation( - asyncio.to_thread(self._commit_branch_pointer, branch_name, generation) + async def _create_and_switch() -> tuple[str, Any, str]: + # Copying the branch is the first durable side effect. Keep the + # generation-guarded pointer update in the same completion-owned task. + ( + resolved_name, + turn_content, + source_branch_id, + generation, + ) = await self._copy_messages_to_new_branch(branch_name, turn_number) + await asyncio.to_thread( + self._commit_branch_pointer, + resolved_name, + generation, + ) + return resolved_name, turn_content, source_branch_id + + resolved_branch_name, turn_content, source_branch_id = await _await_mutation( + _create_and_switch() ) if _debug.DONT_LOG_MODEL_DATA: self._logger.debug( "Created branch '%s' from turn %s in '%s'", - branch_name, + resolved_branch_name, turn_number, source_branch_id, ) else: self._logger.debug( "Created branch '%s' from turn %s ('%s') in '%s'", - branch_name, + resolved_branch_name, turn_number, turn_content, source_branch_id, ) - return branch_name + return resolved_branch_name async def create_branch_from_content( self, search_term: str, branch_name: str | None = None @@ -890,14 +925,11 @@ async def switch_to_branch(self, branch_id: str) -> None: ValueError: If the branch doesn't exist. """ - # Capture the generation before validating so a clear that commits - # between validation and the pointer update is detected and skipped. - generation = self._generation - # Validate branch exists - def _validate_branch(): - """Synchronous helper to validate branch exists.""" - with self._locked_connection() as conn: + def _validate_branch() -> int: + """Validate the branch and return its current durable clear generation.""" + with self._write_connection() as conn: + self._ensure_session_clear_generations_table(conn) with closing(conn.cursor()) as cursor: cursor.execute( """ @@ -910,8 +942,20 @@ def _validate_branch(): count = cursor.fetchone()[0] if count == 0: raise ValueError(f"Branch '{branch_id}' does not exist") + generation = cast( + int, + cursor.execute( + """ + SELECT generation FROM session_clear_generations + WHERE session_id = ? + """, + (self.session_id,), + ).fetchone()[0], + ) + conn.commit() + return generation - await asyncio.to_thread(_validate_branch) + generation = await _await_mutation(asyncio.to_thread(_validate_branch)) old_branch = self._current_branch_id # Update the pointer under the lock; a no-op if a clear_session has @@ -1093,6 +1137,36 @@ def _ensure_branch_reservations_table(self, conn: sqlite3.Connection) -> None: (self.session_id,), ) + def _ensure_session_clear_generations_table(self, conn: sqlite3.Connection) -> None: + """Create and initialize the durable clear generation for this session.""" + conn.execute(""" + CREATE TABLE IF NOT EXISTS session_clear_generations ( + session_id TEXT PRIMARY KEY, + generation INTEGER NOT NULL DEFAULT 0 + ) + """) + conn.execute( + """ + INSERT OR IGNORE INTO session_clear_generations (session_id, generation) + VALUES (?, 0) + """, + (self.session_id,), + ) + + def _refresh_branch_after_external_clear(self, conn: sqlite3.Connection) -> None: + """Reset a stale branch pointer after another session instance clears history.""" + self._ensure_session_clear_generations_table(conn) + generation = conn.execute( + """ + SELECT generation FROM session_clear_generations + WHERE session_id = ? + """, + (self.session_id,), + ).fetchone()[0] + if generation != self._generation: + self._generation = generation + self._current_branch_id = "main" + def _reserve_branch_id( self, cursor: sqlite3.Cursor, new_branch_id: str | None, from_turn_number: int ) -> str: @@ -1128,29 +1202,30 @@ def _reserve_branch_id( branch_id = f"{base_branch_id}_{suffix}" async def _copy_messages_to_new_branch( - self, new_branch_id: str | None, from_turn_number: int, source_branch_id: str - ) -> tuple[str, Any]: + self, new_branch_id: str | None, from_turn_number: int + ) -> tuple[str, Any, str, int]: """Copy messages before the branch point to the new branch. Args: new_branch_id: The ID of the new branch, or None to generate an unused ID. from_turn_number: The turn number to copy messages up to (exclusive). - source_branch_id: The branch to copy messages from. - Returns: - The resolved branch ID and a preview of the source turn content. + The resolved branch ID, source preview, source branch, and clear generation. Raises: ValueError: If `new_branch_id` has already been used in this session. """ - def _copy_sync() -> tuple[str, Any]: + def _copy_sync() -> tuple[str, Any, str, int]: """Synchronous helper to copy messages to new branch.""" with self._write_connection() as conn: # Acquire SQLite's write reservation before checking the branch ID so # sessions in other processes cannot pass the same check concurrently. conn.execute("BEGIN IMMEDIATE") self._ensure_branch_reservations_table(conn) + self._refresh_branch_after_external_clear(conn) + source_branch_id = self._current_branch_id + generation = self._generation with closing(conn.cursor()) as cursor: cursor.execute( f""" @@ -1244,9 +1319,9 @@ def _copy_sync() -> tuple[str, Any]: ) conn.commit() - return branch_id, turn_content + return branch_id, turn_content, source_branch_id, generation - return await _await_mutation(asyncio.to_thread(_copy_sync)) + return await asyncio.to_thread(_copy_sync) async def get_conversation_turns(self, branch_id: str | None = None) -> list[dict[str, Any]]: """Get user turns with content for easy browsing and branching decisions. diff --git a/src/agents/extensions/memory/dapr_session.py b/src/agents/extensions/memory/dapr_session.py index ca72174368..d13c042664 100644 --- a/src/agents/extensions/memory/dapr_session.py +++ b/src/agents/extensions/memory/dapr_session.py @@ -50,6 +50,7 @@ coerce_session_settings, resolve_session_limit, ) +from ...memory.sqlite_session import _await_mutation # Type alias for consistency levels ConsistencyLevel = Literal["eventual", "strong"] @@ -334,23 +335,67 @@ 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] + await _await_mutation(self._add_items_locked(serialized_items)) + + async def _add_items_locked(self, serialized_items: list[str]) -> None: + """Persist one batch while the caller retains the session lock.""" + # 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(), + ) - # 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( + attempt = 0 + while True: + attempt += 1 + response = await self._dapr_client.get_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(), + key=self._messages_key, + state_metadata=self._get_read_metadata(), ) + existing_messages = self._decode_messages_for_update(response.data) + updated_messages = existing_messages + serialized_items + messages_json = json.dumps(updated_messages, separators=(",", ":")) + etag = response.etag + try: + await self._dapr_client.save_state( + store_name=self._state_store_name, + key=self._messages_key, + value=messages_json, + etag=etag, + state_metadata=self._get_metadata(), + options=self._get_state_options(concurrency=Concurrency.first_write), + ) + return + except Exception as error: + should_retry = await self._handle_concurrency_conflict(error, attempt) + if should_retry: + continue + raise + 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._lock: + self._check_not_closed() + return await _await_mutation(self._pop_item_locked()) + + async def _pop_item_locked(self) -> TResponseInputItem | None: + """Claim one item while the caller retains the session lock.""" + while True: attempt = 0 while True: attempt += 1 @@ -359,10 +404,12 @@ async def add_items(self, items: list[TResponseInputItem]) -> None: key=self._messages_key, state_metadata=self._get_read_metadata(), ) - existing_messages = self._decode_messages_for_update(response.data) - updated_messages = existing_messages + serialized_items - messages_json = json.dumps(updated_messages, separators=(",", ":")) - etag = response.etag + messages = self._decode_messages(response.data) + if not messages: + return None + last_item = messages.pop() + messages_json = json.dumps(messages, separators=(",", ":")) + etag = getattr(response, "etag", None) or None try: await self._dapr_client.save_state( store_name=self._state_store_name, @@ -378,68 +425,29 @@ async def add_items(self, items: list[TResponseInputItem]) -> None: if should_retry: continue raise - - 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._lock: - self._check_not_closed() - while True: - attempt = 0 - while True: - attempt += 1 - response = await self._dapr_client.get_state( - store_name=self._state_store_name, - key=self._messages_key, - state_metadata=self._get_read_metadata(), - ) - messages = self._decode_messages(response.data) - if not messages: - return None - last_item = messages.pop() - messages_json = json.dumps(messages, separators=(",", ":")) - etag = getattr(response, "etag", None) or None - try: - await self._dapr_client.save_state( - store_name=self._state_store_name, - key=self._messages_key, - value=messages_json, - etag=etag, - state_metadata=self._get_metadata(), - options=self._get_state_options(concurrency=Concurrency.first_write), - ) - break - except Exception as error: - should_retry = await self._handle_concurrency_conflict(error, attempt) - if should_retry: - continue - raise - try: - if isinstance(last_item, str): - return await self._deserialize_item(last_item) - return last_item # type: ignore[no-any-return] - except (json.JSONDecodeError, TypeError): - continue + try: + if isinstance(last_item, str): + return await self._deserialize_item(last_item) + return last_item # type: ignore[no-any-return] + except (json.JSONDecodeError, TypeError): + continue async def clear_session(self) -> None: """Clear all items for this session.""" async with self._lock: self._check_not_closed() - # Delete messages and metadata keys - await self._dapr_client.delete_state( - store_name=self._state_store_name, - key=self._messages_key, - options=self._get_state_options(), - ) - - await self._dapr_client.delete_state( - store_name=self._state_store_name, - key=self._metadata_key, - options=self._get_state_options(), - ) + await _await_mutation(self._clear_session_locked()) + + async def _clear_session_locked(self) -> None: + """Delete authoritative history while the caller retains the session lock.""" + # The messages key is the authoritative history boundary. Keep the + # ancillary metadata key so clear_session remains one state-store + # mutation instead of partially failing after history is cleared. + await self._dapr_client.delete_state( + store_name=self._state_store_name, + key=self._messages_key, + options=self._get_state_options(), + ) async def close(self) -> None: """Close the Dapr client connection. diff --git a/src/agents/extensions/memory/mongodb_session.py b/src/agents/extensions/memory/mongodb_session.py index b10820267c..c417cec937 100644 --- a/src/agents/extensions/memory/mongodb_session.py +++ b/src/agents/extensions/memory/mongodb_session.py @@ -31,6 +31,7 @@ from __future__ import annotations +import asyncio import json import threading import weakref @@ -65,6 +66,7 @@ coerce_session_settings, resolve_session_limit, ) +from ...memory.sqlite_session import _await_mutation # Identifies this library in the MongoDB handshake for server-side telemetry. _DRIVER_INFO = DriverInfo(name="openai-agents", version=_VERSION) @@ -326,37 +328,34 @@ async def _decode_docs(docs: list[Any]) -> list[TResponseInputItem]: window *= 2 async def add_items(self, items: list[TResponseInputItem]) -> None: - """Add new items to the conversation history. - - Args: - items: List of input items to append to the session. - """ - # Checked before the empty-list fast path, which would otherwise return - # successfully on a closed session. + """Add new items and wait until the batch outcome is known.""" self._check_not_closed() - if not items: return - await self._ensure_indexes() - serialized_items = [await self._serialize_item(item) for item in items] + await _await_mutation(self._add_items(serialized_items)) + + async def _add_items(self, serialized_items: list[str]) -> None: + """Store one pre-serialized logical batch.""" now = datetime.now(timezone.utc) # Atomically reserve a block of sequence numbers for this batch. - # $inc returns the new value, so subtract len(items) to get the first - # number in the block. + # $inc returns the new value, so subtract the batch size to get the + # first number in the block. result = await self._sessions.find_one_and_update( {"session_id": self.session_id}, { "$setOnInsert": {"session_id": self.session_id, "created_at": now}, "$set": {"updated_at": now}, - "$inc": {"_seq": len(items)}, + "$inc": {"_seq": len(serialized_items)}, }, upsert=True, return_document=True, ) - next_seq: int = (result["_seq"] if result else len(items)) - len(items) + next_seq: int = (result["_seq"] if result else len(serialized_items)) - len( + serialized_items + ) # One document is the commit boundary for the logical batch. This keeps # standalone MongoDB deployments failure-atomic without requiring transactions. @@ -369,6 +368,11 @@ async def add_items(self, items: list[TResponseInputItem]) -> None: ) async def pop_item(self) -> TResponseInputItem | None: + """Remove the most recent item after the destructive claim settles.""" + await self._ensure_indexes() + return await _await_mutation(self._pop_item()) + + async def _pop_item(self) -> TResponseInputItem | None: """Remove and return the most recent item from the session. Returns: @@ -380,54 +384,94 @@ async def pop_item(self) -> TResponseInputItem | None: is returned. This matches :meth:`get_items`, which also skips malformed entries, so one bad record cannot make a non-empty session look empty. """ - await self._ensure_indexes() + # Retry cleanup left by a prior post-claim failure. Empty markers are + # never model-visible or claimable, so cleanup failure must not block a + # later valid tail claim. + try: + await self._messages.delete_many({"session_id": self.session_id, "message_data": []}) + except asyncio.CancelledError: + raise + except Exception: + pass while True: - docs = await ( - self._messages.find( + doc = await self._messages.find_one_and_update( + { + "session_id": self.session_id, + "message_data": {"$ne": []}, + }, + [ { - "session_id": self.session_id, - "message_data": {"$ne": []}, + "$set": { + "message_data": { + "$cond": [ + {"$isArray": "$message_data"}, + { + "$slice": [ + "$message_data", + { + "$subtract": [ + {"$size": "$message_data"}, + 1, + ] + }, + ] + }, + [], + ] + } + } } - ) - .sort("seq", -1) - .limit(1) - .to_list() + ], + sort=[("seq", -1)], + return_document=False, ) - if not docs: + if doc is None: return None - candidate = docs[0] - raw = candidate.get("message_data") + raw = doc.get("message_data") if isinstance(raw, list): - doc = await self._messages.find_one_and_update( - {"_id": candidate["_id"], "message_data": {"$ne": []}}, - {"$pop": {"message_data": 1}}, - return_document=False, - ) - if doc is None: + if not raw: continue - claimed_batch = doc.get("message_data") - if not isinstance(claimed_batch, list) or not claimed_batch: - continue - raw = claimed_batch[-1] + claimed_raw = raw[-1] + exhausted = len(raw) == 1 else: - doc = await self._messages.find_one_and_delete({"_id": candidate["_id"]}) - if doc is None: - continue - raw = doc.get("message_data") + claimed_raw = raw + exhausted = True + + if exhausted: + # The atomic claim above leaves an empty marker so another pop cannot + # claim the same item. Remove that marker before returning to avoid + # accumulating exhausted logical-batch and legacy documents. + try: + await self._messages.delete_one({"_id": doc["_id"], "message_data": []}) + except asyncio.CancelledError: + raise + except Exception: + # The item is already claimed. Do not turn a known destructive + # outcome into a retry-visible failure; the next pop retries the + # best-effort empty-marker sweep above. + pass try: - return await self._deserialize_item(raw) + return await self._deserialize_item(claimed_raw) except (json.JSONDecodeError, TypeError): # Corrupt — drop it and try the next-most-recent document. continue async def clear_session(self) -> None: - """Clear all items for this session.""" + """Clear history after the authoritative delete settles.""" await self._ensure_indexes() + await _await_mutation(self._clear_session()) + + async def _clear_session(self) -> None: + """Clear all items for this session.""" await self._messages.delete_many({"session_id": self.session_id}) - await self._sessions.delete_one({"session_id": self.session_id}) + # Keep the ancillary metadata document and its sequence counter. The + # messages collection is the authoritative history boundary, and a + # second delete would make clear_session partially fail after history + # was already cleared. Retaining the counter also prevents a concurrent + # add from losing its ordering authority across this clear. # ------------------------------------------------------------------ # Lifecycle helpers diff --git a/src/agents/extensions/memory/redis_session.py b/src/agents/extensions/memory/redis_session.py index 2fe5647184..694d6888fc 100644 --- a/src/agents/extensions/memory/redis_session.py +++ b/src/agents/extensions/memory/redis_session.py @@ -31,7 +31,7 @@ try: import redis.asyncio as redis - from redis.asyncio import BlockingConnectionPool, ConnectionPool, Redis + from redis.asyncio import BlockingConnectionPool, Redis from redis.exceptions import ResponseError, WatchError except ImportError as e: raise_optional_dependency_error( @@ -48,6 +48,7 @@ coerce_session_settings, resolve_session_limit, ) +from ...memory.sqlite_session import _await_mutation @dataclass @@ -65,15 +66,6 @@ def __init__(self, pool: Any): self._pool = pool self.release_started = False self.release_completed = False - release_implementation = getattr(pool.release, "__func__", None) - if release_implementation not in { - ConnectionPool.release, - BlockingConnectionPool.release, - }: - raise RuntimeError( - "RedisSession requires the standard redis-py connection pool release method" - ) - self._uses_standard_release = release_implementation is ConnectionPool.release def __getattr__(self, name: str) -> Any: return getattr(self._pool, name) @@ -92,12 +84,6 @@ async def release(self, connection: Any) -> None: if was_in_use and connection in self._pool._available_connections: self.release_completed = True await self.notify_capacity_available() - elif ( - was_in_use - and self._uses_standard_release - and connection in self._pool._in_use_connections - ): - self.release_completed = True raise else: self.release_completed = True @@ -113,10 +99,11 @@ async def _finish_pipeline( return None, True, None reset_error: BaseException | None = None - try: - await pipe.reset() - except BaseException as exc: - reset_error = exc + if not connection_pool.release_started: + try: + await pipe.reset() + except BaseException as exc: + reset_error = exc if connection_pool.release_completed: pipe.connection = None @@ -126,19 +113,22 @@ async def _finish_pipeline( if connection is None: return reset_error, True, None - if ( - connection_pool.release_started - and connection not in connection_pool._pool._in_use_connections - ): + if connection_pool.release_started: + # The first release may have transferred and reborrowed the connection + # before failing. Its ownership is ambiguous, so never release it again. + connection_pool._pool._in_use_connections.discard(connection) + while connection in connection_pool._pool._available_connections: + connection_pool._pool._available_connections.remove(connection) try: connection._close() except BaseException as close_error: pipe.connection = None + connection_pool.release_completed = True + await connection_pool.notify_capacity_available() return close_error, True, connection await connection_pool.notify_capacity_available() pipe.connection = None - if reset_error is None: - reset_error = RuntimeError("Redis pipeline release did not complete") + connection_pool.release_completed = True return reset_error, True, None try: @@ -149,14 +139,19 @@ async def _finish_pipeline( try: await connection_pool.release(connection) except BaseException as release_error: - if ( - not connection_pool.release_completed - and connection not in connection_pool._pool._in_use_connections - ): + if not connection_pool.release_completed: + # A custom release implementation may transfer and reborrow the + # connection before failing. Remove ambiguous pool membership and + # quarantine the connection instead of leaving it checked out. + connection_pool._pool._in_use_connections.discard(connection) + while connection in connection_pool._pool._available_connections: + connection_pool._pool._available_connections.remove(connection) try: connection._close() except BaseException as close_error: pipe.connection = None + connection_pool.release_completed = True + await connection_pool.notify_capacity_available() return close_error, True, connection await connection_pool.notify_capacity_available() connection_pool.release_completed = True @@ -507,34 +502,41 @@ async def pop_item(self) -> TResponseInputItem | None: """ async with self._lock: self._check_not_closed() - while True: - # Use RPOP to atomically remove and return the rightmost (most recent) item - raw_msg = await self._redis.rpop(self._messages_key) # type: ignore[misc] # Redis library returns Union[Awaitable[T], T] in async context + return await _await_mutation(self._pop_item_locked()) - if raw_msg is None: - return None + async def _pop_item_locked(self) -> TResponseInputItem | None: + """Claim one item while the caller retains the session lock.""" + while True: + # Use RPOP to atomically remove and return the rightmost (most recent) item + raw_msg = await self._redis.rpop(self._messages_key) # type: ignore[misc] # Redis library returns Union[Awaitable[T], T] in async context - try: - # Handle both bytes (default) and str (decode_responses=True) Redis clients - if isinstance(raw_msg, bytes): - msg_str = raw_msg.decode("utf-8") - else: - msg_str = raw_msg # Already a string - return await self._deserialize_item(msg_str) - except (json.JSONDecodeError, UnicodeDecodeError): - # Drop corrupted messages and keep looking for a valid item. - continue + if raw_msg is None: + return None + + try: + # Handle both bytes (default) and str (decode_responses=True) Redis clients + if isinstance(raw_msg, bytes): + msg_str = raw_msg.decode("utf-8") + else: + msg_str = raw_msg # Already a string + return await self._deserialize_item(msg_str) + except (json.JSONDecodeError, UnicodeDecodeError): + # Drop corrupted messages and keep looking for a valid item. + continue async def clear_session(self) -> None: """Clear all items for this session.""" async with self._lock: self._check_not_closed() - # Delete all keys associated with this session - await self._redis.delete( - self._session_key, - self._messages_key, - self._counter_key, - ) + await _await_mutation(self._clear_session_locked()) + + async def _clear_session_locked(self) -> None: + """Delete all session keys while the caller retains the session lock.""" + await self._redis.delete( + self._session_key, + self._messages_key, + self._counter_key, + ) async def close(self) -> None: """Close the Redis connection. diff --git a/src/agents/extensions/memory/sqlalchemy_session.py b/src/agents/extensions/memory/sqlalchemy_session.py index 40b2ca5e1d..81f7dcdae8 100644 --- a/src/agents/extensions/memory/sqlalchemy_session.py +++ b/src/agents/extensions/memory/sqlalchemy_session.py @@ -27,7 +27,8 @@ import json import threading import weakref -from typing import Any, ClassVar +from collections.abc import Awaitable, Callable +from typing import Any, ClassVar, TypeVar from sqlalchemy import ( TIMESTAMP, @@ -57,6 +58,9 @@ coerce_session_settings, resolve_session_limit, ) +from ...memory.sqlite_session import _await_mutation + +_T = TypeVar("_T") class SQLAlchemySession(SessionABC): @@ -122,23 +126,22 @@ def _configure_sqlite_connection(dbapi_connection: Any, _: Any) -> None: def _is_sqlite_lock_error(exc: OperationalError) -> bool: return "database is locked" in str(exc).lower() - async def _run_sqlite_write_with_retry(self, operation: Any) -> None: + async def _run_sqlite_write_with_retry(self, operation: Callable[[], Awaitable[_T]]) -> _T: """Retry transient SQLite write lock failures with bounded backoff.""" if self._engine.dialect.name != "sqlite": - await operation() - return + return await operation() for attempt, delay in enumerate((0.0, *self._SQLITE_LOCK_RETRY_DELAYS)): if delay: await asyncio.sleep(delay) try: - await operation() - return + return await operation() except OperationalError as exc: if not self._is_sqlite_lock_error(exc): raise if attempt == len(self._SQLITE_LOCK_RETRY_DELAYS): raise + raise AssertionError("SQLite write retry loop exited unexpectedly") def __init__( self, @@ -412,21 +415,32 @@ async def _write_items() -> None: .values(updated_at=sql_text("CURRENT_TIMESTAMP")) ) - await self._run_sqlite_write_with_retry(_write_items) + await _await_mutation(self._run_sqlite_write_with_retry(_write_items)) async def pop_item(self) -> TResponseInputItem | None: + """Remove the most recent item after its transaction settles.""" + await self._ensure_tables() + return await _await_mutation(self._run_sqlite_write_with_retry(self._pop_item)) + + 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 """ - await self._ensure_tables() while True: + retry_claim = False async with self._session_factory() as sess: async with sess.begin(): - # Fallback for all dialects - get ID first, then delete - subq = ( - select(self._messages.c.id) + if ( + self._engine.dialect.name == "sqlite" + and not self._engine.dialect.delete_returning + ): + # SQLite ignores SELECT ... FOR UPDATE. Reserve the single + # writer before selecting so the fallback claim remains unique. + await sess.execute(sql_text("BEGIN IMMEDIATE")) + tail = ( + select(self._messages.c.id, self._messages.c.message_data) .where(self._messages.c.session_id == self.session_id) .order_by( self._messages.c.created_at.desc(), @@ -434,41 +448,58 @@ async def pop_item(self) -> TResponseInputItem | None: ) .limit(1) ) - res = await sess.execute(subq) - row_id = res.scalar_one_or_none() - if row_id is None: - return None - # Fetch data before deleting - res_data = await sess.execute( - select(self._messages.c.message_data).where(self._messages.c.id == row_id) - ) - row = res_data.scalar_one_or_none() - delete_result = await sess.execute( - delete(self._messages).where(self._messages.c.id == row_id) - ) - if delete_result.rowcount == 0: - # A concurrent caller claimed this row. Start a new transaction so - # dialects with repeatable-read snapshots can observe the next row. - continue - if delete_result.rowcount != 1: - # Unknown or unexpected row counts cannot prove that this caller owns - # the destructive read. Raising inside the transaction rolls the DELETE - # back instead of risking the same payload being returned twice. - raise RuntimeError( - "SQLAlchemySession.pop_item() could not confirm that exactly one " - "item was deleted" + + if self._engine.dialect.delete_returning: + # DELETE ... RETURNING is the claim: only the transaction that + # removes the current tail receives its payload. This avoids relying + # on DBAPI rowcount, which some dialects report as unknown. + result = await sess.execute( + delete(self._messages) + .where( + self._messages.c.id + == tail.with_only_columns(self._messages.c.id).scalar_subquery() + ) + .returning(self._messages.c.message_data) + ) + row = result.scalar_one_or_none() + if row is None: + # A concurrent DELETE can win the same tail between the + # subquery read and this claim. Distinguish that race from + # an empty session before retrying with a fresh transaction. + remaining = await sess.execute( + tail.with_only_columns(self._messages.c.id) + ) + if remaining.scalar_one_or_none() is None: + return None + retry_claim = True + else: + # Dialects without DELETE ... RETURNING claim the row with a + # transaction-scoped lock before deleting it. The lock, rather than + # rowcount, establishes ownership of the returned payload. + result = await sess.execute(tail.with_for_update()) + claimed = result.one_or_none() + if claimed is None: + return None + row_id, row = claimed + await sess.execute( + delete(self._messages).where(self._messages.c.id == row_id) ) - if row is None: - continue - try: - return await self._deserialize_item(row) - except (json.JSONDecodeError, TypeError): - continue + if retry_claim: + continue + assert row is not None + try: + return await self._deserialize_item(row) + except (json.JSONDecodeError, TypeError): + continue async def clear_session(self) -> None: - """Clear all items for this session.""" + """Clear history after its transaction settles.""" await self._ensure_tables() + await _await_mutation(self._clear_session()) + + async def _clear_session(self) -> None: + """Clear all items for this session.""" async with self._session_factory() as sess: async with sess.begin(): await sess.execute( diff --git a/tests/extensions/memory/test_advanced_sqlite_session.py b/tests/extensions/memory/test_advanced_sqlite_session.py index 548052c8be..cd4397423a 100644 --- a/tests/extensions/memory/test_advanced_sqlite_session.py +++ b/tests/extensions/memory/test_advanced_sqlite_session.py @@ -6,6 +6,7 @@ import logging import multiprocessing import sqlite3 +import sys import tempfile import threading import time @@ -32,6 +33,12 @@ pytestmark = pytest.mark.asyncio +def _assert_cancel_message(exc: asyncio.CancelledError, expected: str) -> None: + """Account for Python 3.10 dropping Task cancellation messages when re-awaited.""" + expected_args = (expected,) if sys.version_info >= (3, 11) else () + assert exc.args == expected_args + + @function_tool async def test_tool(query: str) -> str: """A test tool for testing tool call tracking.""" @@ -641,13 +648,13 @@ def commit(self) -> None: with pytest.raises(asyncio.CancelledError) as exc_info: await mutation - assert exc_info.value.args == ("first-caller-cancel",) + _assert_cancel_message(exc_info.value, "first-caller-cancel") assert mutation.cancelled() if operation == "create_branch": branches = await session.list_branches() assert {branch["branch_id"] for branch in branches} == {"main", "cancelled_branch"} - assert session._current_branch_id == "main" + assert session._current_branch_id == "cancelled_branch" elif operation == "delete_branch": branches = await session.list_branches() assert {branch["branch_id"] for branch in branches} == {"main"} @@ -3358,6 +3365,198 @@ async def test_clear_session_resets_current_branch_to_main(): session.close() +async def test_external_clear_resets_stale_branch_before_next_write(tmp_path: Path): + """A second instance's clear must prevent stale branch resurrection.""" + db_path = tmp_path / "external_clear_generation.db" + stale = AdvancedSQLiteSession( + session_id="external_clear_generation", + db_path=db_path, + create_tables=True, + ) + clearer = AdvancedSQLiteSession( + session_id="external_clear_generation", + db_path=db_path, + ) + + try: + await stale.add_items( + [ + {"role": "user", "content": "u1"}, + {"role": "assistant", "content": "a1"}, + {"role": "user", "content": "u2"}, + ] + ) + await stale.create_branch_from_turn(2, "stale") + assert stale._current_branch_id == "stale" + + await clearer.clear_session() + await stale.add_items([{"role": "user", "content": "after clear"}]) + + assert stale._current_branch_id == "main" + assert [item.get("content") for item in await stale.get_items()] == ["after clear"] + assert await stale.get_items(branch_id="stale") == [] + assert {branch["branch_id"] for branch in await stale.list_branches()} == {"main"} + finally: + stale.close() + clearer.close() + + +async def test_external_clear_resets_stale_branch_before_pop(tmp_path: Path): + """A stale instance must pop the current main tail after an external clear.""" + db_path = tmp_path / "external_clear_pop_generation.db" + stale = AdvancedSQLiteSession( + session_id="external_clear_pop_generation", + db_path=db_path, + create_tables=True, + ) + clearer = AdvancedSQLiteSession( + session_id="external_clear_pop_generation", + db_path=db_path, + ) + + try: + await stale.add_items( + [ + {"role": "user", "content": "u1"}, + {"role": "assistant", "content": "a1"}, + {"role": "user", "content": "u2"}, + ] + ) + await stale.create_branch_from_turn(2, "stale") + assert stale._current_branch_id == "stale" + + await clearer.clear_session() + item: TResponseInputItem = {"role": "user", "content": "after clear"} + await clearer.add_items([item]) + + assert await stale.pop_item() == item + assert stale._current_branch_id == "main" + assert await clearer.get_items() == [] + finally: + stale.close() + clearer.close() + + +async def test_switch_validation_cancellation_waits_for_generation_commit( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +): + """Legacy generation initialization must settle before cancellation propagates.""" + + class PausingCommitConnection(sqlite3.Connection): + pause_commit = False + commit_finished = threading.Event() + allow_return = threading.Event() + + def commit(self) -> None: + super().commit() + if self.pause_commit: + self.pause_commit = False + self.commit_finished.set() + assert self.allow_return.wait(timeout=10) + + db_path = tmp_path / "switch_validation_cancellation.db" + session = AdvancedSQLiteSession( + session_id="switch_validation_cancellation", + db_path=db_path, + create_tables=True, + ) + mutation: asyncio.Task[Any] | None = None + + try: + await session.add_items( + [ + {"role": "user", "content": "u1"}, + {"role": "assistant", "content": "a1"}, + {"role": "user", "content": "u2"}, + ] + ) + await session.create_branch_from_turn(2, "target") + await session.switch_to_branch("main") + with session._write_connection() as setup_connection: + setup_connection.execute("DROP TABLE session_clear_generations") + setup_connection.commit() + + connection = sqlite3.connect( + str(db_path), + check_same_thread=False, + factory=PausingCommitConnection, + ) + with session._connections_lock: + session._connections.add(connection) + monkeypatch.setattr(session, "_get_connection", lambda: connection) + connection.pause_commit = True + + mutation = asyncio.create_task(session.switch_to_branch("target")) + assert await asyncio.to_thread(connection.commit_finished.wait, 10) + mutation.cancel("first-caller-cancel") + await asyncio.sleep(0) + mutation.cancel("second-caller-cancel") + await asyncio.sleep(0) + assert mutation.done() is False + connection.allow_return.set() + + with pytest.raises(asyncio.CancelledError) as exc_info: + await mutation + + _assert_cancel_message(exc_info.value, "first-caller-cancel") + assert session._current_branch_id == "main" + row = connection.execute( + "SELECT generation FROM session_clear_generations WHERE session_id = ?", + (session.session_id,), + ).fetchone() + assert row == (0,) + finally: + PausingCommitConnection.allow_return.set() + if mutation is not None and not mutation.done(): + mutation.cancel() + await asyncio.gather(mutation, return_exceptions=True) + session.close() + + +async def test_post_clear_switch_synchronizes_generation_before_next_write(tmp_path: Path): + """A new instance may select and write to a branch created after an earlier clear.""" + db_path = tmp_path / "post_clear_branch_switch.db" + owner = AdvancedSQLiteSession( + session_id="post_clear_branch_switch", + db_path=db_path, + create_tables=True, + ) + other = AdvancedSQLiteSession( + session_id="post_clear_branch_switch", + db_path=db_path, + ) + + try: + await owner.clear_session() + await owner.add_items( + [ + {"role": "user", "content": "u1"}, + {"role": "assistant", "content": "a1"}, + {"role": "user", "content": "u2"}, + ] + ) + await owner.create_branch_from_turn(2, "fresh") + + await other.switch_to_branch("fresh") + await other.add_items([{"role": "assistant", "content": "on fresh"}]) + + assert other._current_branch_id == "fresh" + assert [item.get("content") for item in await other.get_items()] == [ + "u1", + "a1", + "on fresh", + ] + assert [item.get("content") for item in await other.get_items(branch_id="main")] == [ + "u1", + "a1", + "u2", + ] + finally: + owner.close() + other.close() + + async def test_pop_item_rolls_back_on_failure_after_earlier_delete(): """Regression: a failure partway through pop_item's delete sequence must roll back so no partial mutation or open transaction survives. diff --git a/tests/extensions/memory/test_async_sqlite_session.py b/tests/extensions/memory/test_async_sqlite_session.py index abaffcc8f2..5ade0e2cc5 100644 --- a/tests/extensions/memory/test_async_sqlite_session.py +++ b/tests/extensions/memory/test_async_sqlite_session.py @@ -5,6 +5,7 @@ import asyncio import json import sqlite3 +import sys import tempfile from collections.abc import Sequence from datetime import datetime @@ -24,6 +25,12 @@ pytestmark = pytest.mark.asyncio +def _assert_cancel_message(exc: asyncio.CancelledError, expected: str) -> None: + """Account for Python 3.10 dropping Task cancellation messages when re-awaited.""" + expected_args = (expected,) if sys.version_info >= (3, 11) else () + assert exc.args == expected_args + + @pytest.fixture def agent() -> Agent: """Fixture for a basic agent with a fake model.""" @@ -820,7 +827,7 @@ async def controlled_rollback() -> None: task.cancel() await asyncio.gather(task, return_exceptions=True) - assert exc_info.value.args == ("first-caller-cancel",) + _assert_cancel_message(exc_info.value, "first-caller-cancel") assert conn.in_transaction is False assert _sqlite_write_lock_is_free(db_path) assert await session.get_items() == [] @@ -1056,7 +1063,7 @@ async def controlled_close() -> None: task.cancel() await asyncio.gather(task, return_exceptions=True) - assert exc_info.value.args == ("first-caller-cancel",) + _assert_cancel_message(exc_info.value, "first-caller-cancel") assert session._connection is None assert session.captured_connection._running is False finally: diff --git a/tests/extensions/memory/test_dapr_session.py b/tests/extensions/memory/test_dapr_session.py index 7cdde0f124..c1a765fc10 100644 --- a/tests/extensions/memory/test_dapr_session.py +++ b/tests/extensions/memory/test_dapr_session.py @@ -1,5 +1,6 @@ from __future__ import annotations +import asyncio import json from typing import Any from unittest.mock import Mock @@ -250,11 +251,72 @@ async def test_dapr_session_direct_ops(fake_dapr_client: FakeDaprClient): await session.clear_session() retrieved_after_clear = await session.get_items() assert len(retrieved_after_clear) == 0 + assert session._metadata_key in fake_dapr_client._state finally: await session.close() +@pytest.mark.parametrize("operation", ["add", "pop", "clear"]) +async def test_mutation_cancellation_waits_for_authoritative_outcome( + fake_dapr_client: FakeDaprClient, + monkeypatch: pytest.MonkeyPatch, + operation: str, +) -> None: + """Cancellation must wait after the state store applies a history mutation.""" + session = await _create_test_session(fake_dapr_client) + item: TResponseInputItem = {"role": "user", "content": "once"} + if operation != "add": + await session.add_items([item]) + + mutation_applied = asyncio.Event() + allow_return = asyncio.Event() + + if operation in {"add", "pop"}: + original_save = fake_dapr_client.save_state + + async def controlled_save(*args: Any, **kwargs: Any) -> None: + await original_save(*args, **kwargs) + if kwargs.get("key") == session._messages_key: + mutation_applied.set() + await allow_return.wait() + + monkeypatch.setattr(fake_dapr_client, "save_state", controlled_save) + if operation == "add": + task: asyncio.Task[Any] = asyncio.create_task(session.add_items([item])) + else: + task = asyncio.create_task(session.pop_item()) + else: + original_delete = fake_dapr_client.delete_state + + async def controlled_delete(*args: Any, **kwargs: Any) -> None: + await original_delete(*args, **kwargs) + mutation_applied.set() + await allow_return.wait() + + monkeypatch.setattr(fake_dapr_client, "delete_state", controlled_delete) + task = asyncio.create_task(session.clear_session()) + + try: + await mutation_applied.wait() + task.cancel() + await asyncio.sleep(0) + task.cancel() + await asyncio.sleep(0) + assert task.done() is False + allow_return.set() + with pytest.raises(asyncio.CancelledError): + await task + finally: + allow_return.set() + if not task.done(): + task.cancel() + await asyncio.gather(task, return_exceptions=True) + + expected = [item] if operation == "add" else [] + assert await session.get_items() == expected + + async def test_runner_integration(agent: Agent, fake_dapr_client: FakeDaprClient): """Test that DaprSession works correctly with the agent Runner.""" session = await _create_test_session(fake_dapr_client) diff --git a/tests/extensions/memory/test_mongodb_session.py b/tests/extensions/memory/test_mongodb_session.py index 31cd2c496f..acd9511343 100644 --- a/tests/extensions/memory/test_mongodb_session.py +++ b/tests/extensions/memory/test_mongodb_session.py @@ -45,6 +45,12 @@ def __init__(self) -> None: def __lt__(self, other: FakeObjectId) -> bool: return self._value < other._value + def __eq__(self, other: object) -> bool: + return isinstance(other, FakeObjectId) and self._value == other._value + + def __hash__(self) -> int: + return hash(self._value) + def __repr__(self) -> str: return f"FakeObjectId({self._value})" @@ -130,7 +136,7 @@ async def insert_one(self, document: dict[str, Any]) -> Any: async def find_one_and_update( self, query: dict[str, Any], - update: dict[str, Any], + update: dict[str, Any] | list[dict[str, Any]], upsert: bool = False, return_document: bool = False, sort: list[tuple[str, int]] | None = None, @@ -142,6 +148,10 @@ async def find_one_and_update( if matches: doc = matches[0] before = copy.deepcopy(doc) + if isinstance(update, list): + raw = doc.get("message_data") + doc["message_data"] = raw[:-1] if isinstance(raw, list) else [] + return copy.deepcopy(doc) if return_document else before for field, delta in update.get("$inc", {}).items(): doc[field] = doc.get(field, 0) + delta for field, value in update.get("$set", {}).items(): @@ -152,6 +162,7 @@ async def find_one_and_update( values.pop(-1 if direction == 1 else 0) return copy.deepcopy(doc) if return_document else before if upsert: + assert isinstance(update, dict) new_doc: dict[str, Any] = {"_id": FakeObjectId()} new_doc.update(update.get("$setOnInsert", {})) new_doc.update(update.get("$set", {})) @@ -361,10 +372,19 @@ async def test_pop_item_empty_session(session: MongoDBSession) -> None: async def test_clear_session(session: MongoDBSession) -> None: - """clear_session must remove all items and session metadata.""" + """clear_session removes history while preserving monotonic ordering metadata.""" await session.add_items([{"role": "user", "content": "x"}]) + metadata: dict[str, Any] = next(iter(session._sessions._docs.values())) + sequence_before_clear = metadata["_seq"] + await session.clear_session() + assert await session.get_items() == [] + assert next(iter(session._sessions._docs.values()))["_seq"] == sequence_before_clear + + await session.add_items([{"role": "user", "content": "after clear"}]) + batch = next(iter(session._messages._docs.values())) + assert batch["seq"] == sequence_before_clear async def test_multiple_add_calls_accumulate(session: MongoDBSession) -> None: @@ -1081,9 +1101,7 @@ async def pause_then_fail(document: dict[str, Any]) -> None: assert await session.get_items() == [seed, *batch] -async def test_concurrent_pop_item_claims_distinct_items_from_batch( - monkeypatch: pytest.MonkeyPatch, -) -> None: +async def test_concurrent_pop_item_claims_distinct_items_from_batch() -> None: """Atomic array pops must return each item in a logical batch at most once.""" MongoDBSession._init_state.clear() client = FakeAsyncMongoClient() @@ -1099,29 +1117,10 @@ async def test_concurrent_pop_item_claims_distinct_items_from_batch( ] await session.add_items(items) - real_to_list = FakeCursor.to_list - initial_reads = 0 - all_initial_reads_ready = asyncio.Event() - - async def pause_initial_candidate_reads(cursor: FakeCursor) -> list[dict[str, Any]]: - nonlocal initial_reads - docs = [ - {**doc, "message_data": copy.deepcopy(doc.get("message_data"))} - for doc in await real_to_list(cursor) - ] - if initial_reads < len(items): - initial_reads += 1 - if initial_reads == len(items): - all_initial_reads_ready.set() - await all_initial_reads_ready.wait() - return docs - - monkeypatch.setattr(FakeCursor, "to_list", pause_initial_candidate_reads) tasks = [asyncio.create_task(session.pop_item()) for _ in items] try: popped = await asyncio.wait_for(asyncio.gather(*tasks), timeout=1) finally: - all_initial_reads_ready.set() for task in tasks: if not task.done(): task.cancel() @@ -1134,6 +1133,179 @@ async def pause_initial_candidate_reads(cursor: FakeCursor) -> list[dict[str, An "third", } assert await session.get_items() == [] + assert session._messages._docs == {} + + +async def test_pop_item_atomically_selects_newest_concurrent_batch( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """The claim operation must select a batch appended before its linearization point.""" + MongoDBSession._init_state.clear() + client = FakeAsyncMongoClient() + session = MongoDBSession( + "concurrent-add-pop", + client=client, # type: ignore[arg-type] + database="agents_test", + ) + await session.add_items([{"role": "user", "content": "older"}]) + + real_find_one_and_update = session._messages.find_one_and_update + appended = False + + async def append_before_claim(*args: Any, **kwargs: Any) -> Any: + nonlocal appended + if not appended: + appended = True + await session.add_items([{"role": "assistant", "content": "newer"}]) + return await real_find_one_and_update(*args, **kwargs) + + monkeypatch.setattr(session._messages, "find_one_and_update", append_before_claim) + + popped = await session.pop_item() + + assert popped is not None + assert popped.get("content") == "newer" + assert [item.get("content") for item in await session.get_items()] == ["older"] + + +async def test_pop_item_deletes_exhausted_batch_document(session: MongoDBSession) -> None: + """Popping a one-item batch must not leave an empty batch document behind.""" + await session.add_items([{"role": "user", "content": "only"}]) + + popped = await session.pop_item() + assert popped is not None + assert popped.get("content") == "only" + assert session._messages._docs == {} + + +async def test_pop_item_cancellation_waits_for_exhausted_batch_cleanup( + session: MongoDBSession, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Cancellation after a claim must wait until its empty marker is deleted.""" + await session.add_items([{"role": "user", "content": "claimed"}]) + real_delete_one = session._messages.delete_one + delete_started = asyncio.Event() + allow_delete = asyncio.Event() + + async def controlled_delete(query: dict[str, Any]) -> None: + delete_started.set() + await allow_delete.wait() + await real_delete_one(query) + + monkeypatch.setattr(session._messages, "delete_one", controlled_delete) + task = asyncio.create_task(session.pop_item()) + try: + await delete_started.wait() + task.cancel("first-caller-cancel") + await asyncio.sleep(0) + task.cancel("second-caller-cancel") + await asyncio.sleep(0) + assert task.done() is False + allow_delete.set() + with pytest.raises(asyncio.CancelledError): + await task + finally: + allow_delete.set() + if not task.done(): + task.cancel() + await asyncio.gather(task, return_exceptions=True) + + assert session._messages._docs == {} + + +@pytest.mark.parametrize("operation", ["add", "pop", "clear"]) +async def test_mutation_cancellation_waits_for_authoritative_outcome( + session: MongoDBSession, + monkeypatch: pytest.MonkeyPatch, + operation: str, +) -> None: + """Cancellation must wait after the server has applied a history mutation.""" + item: TResponseInputItem = {"role": "user", "content": "once"} + if operation != "add": + await session.add_items([item]) + + mutation_applied = asyncio.Event() + allow_return = asyncio.Event() + task: asyncio.Task[Any] + + if operation == "add": + original_insert = session._messages.insert_one + + async def controlled_insert(document: dict[str, Any]) -> Any: + result = await original_insert(document) + mutation_applied.set() + await allow_return.wait() + return result + + monkeypatch.setattr(session._messages, "insert_one", controlled_insert) + task = asyncio.create_task(session.add_items([item])) + elif operation == "pop": + original_claim = session._messages.find_one_and_update + + async def controlled_claim(*args: Any, **kwargs: Any) -> Any: + result = await original_claim(*args, **kwargs) + mutation_applied.set() + await allow_return.wait() + return result + + monkeypatch.setattr(session._messages, "find_one_and_update", controlled_claim) + task = asyncio.create_task(session.pop_item()) + else: + original_clear = session._messages.delete_many + + async def controlled_clear(*args: Any, **kwargs: Any) -> Any: + result = await original_clear(*args, **kwargs) + mutation_applied.set() + await allow_return.wait() + return result + + monkeypatch.setattr(session._messages, "delete_many", controlled_clear) + task = asyncio.create_task(session.clear_session()) + + try: + await mutation_applied.wait() + task.cancel() + await asyncio.sleep(0) + task.cancel() + await asyncio.sleep(0) + assert task.done() is False + allow_return.set() + with pytest.raises(asyncio.CancelledError): + await task + finally: + allow_return.set() + if not task.done(): + task.cancel() + await asyncio.gather(task, return_exceptions=True) + + expected = [item] if operation == "add" else [] + assert await session.get_items() == expected + + +async def test_pop_item_cleanup_failure_does_not_hide_known_claim( + session: MongoDBSession, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """A failed empty-marker delete must not make a claimed item retry-visible.""" + await session.add_items([{"role": "user", "content": "claimed"}]) + real_delete_one = session._messages.delete_one + + async def fail_delete_one(query: dict[str, Any]) -> None: + raise RuntimeError("cleanup failed") + + monkeypatch.setattr(session._messages, "delete_one", fail_delete_one) + popped = await session.pop_item() + + assert popped is not None + assert popped.get("content") == "claimed" + assert len(session._messages._docs) == 1 + remaining_doc: dict[str, Any] = next(iter(session._messages._docs.values())) + assert remaining_doc["message_data"] == [] + + monkeypatch.setattr(session._messages, "delete_one", real_delete_one) + assert await session.pop_item() is None + assert session._messages._docs == {} async def test_reads_legacy_item_documents_with_new_batch_documents() -> None: diff --git a/tests/extensions/memory/test_redis_session.py b/tests/extensions/memory/test_redis_session.py index 6fef22e38c..1a0a689853 100644 --- a/tests/extensions/memory/test_redis_session.py +++ b/tests/extensions/memory/test_redis_session.py @@ -2,7 +2,9 @@ import asyncio import json +import sys import time +from collections.abc import Awaitable from typing import Any, cast import pytest @@ -17,6 +19,13 @@ # Keep the fallback-to-real-Redis path isolated from xdist workers. pytestmark = [pytest.mark.asyncio, pytest.mark.serial] + +def _assert_cancel_message(exc: asyncio.CancelledError, expected: str) -> None: + """Account for Python 3.10 dropping Task cancellation messages when re-awaited.""" + expected_args = (expected,) if sys.version_info >= (3, 11) else () + assert exc.args == expected_args + + # Try to use fakeredis for in-memory testing, fall back to real Redis if not available try: import fakeredis.aioredis @@ -128,6 +137,60 @@ async def test_redis_session_direct_ops(): await session.close() +@pytest.mark.parametrize("operation", ["pop", "clear"]) +async def test_mutation_cancellation_waits_for_authoritative_outcome( + monkeypatch: pytest.MonkeyPatch, + operation: str, +) -> None: + """Cancellation must wait after Redis applies a destructive mutation.""" + session = await _create_test_session() + item: TResponseInputItem = {"role": "user", "content": "once"} + await session.add_items([item]) + mutation_applied = asyncio.Event() + allow_return = asyncio.Event() + + if operation == "pop": + original_rpop = session._redis.rpop + + async def controlled_rpop(*args: Any, **kwargs: Any) -> Any: + result = await cast(Awaitable[Any], original_rpop(*args, **kwargs)) + mutation_applied.set() + await allow_return.wait() + return result + + monkeypatch.setattr(session._redis, "rpop", controlled_rpop) + task: asyncio.Task[Any] = asyncio.create_task(session.pop_item()) + else: + original_delete = session._redis.delete + + async def controlled_delete(*args: Any, **kwargs: Any) -> Any: + result = await original_delete(*args, **kwargs) + mutation_applied.set() + await allow_return.wait() + return result + + monkeypatch.setattr(session._redis, "delete", controlled_delete) + task = asyncio.create_task(session.clear_session()) + + try: + await mutation_applied.wait() + task.cancel() + await asyncio.sleep(0) + task.cancel() + await asyncio.sleep(0) + assert task.done() is False + allow_return.set() + with pytest.raises(asyncio.CancelledError): + await task + finally: + allow_return.set() + if not task.done(): + task.cancel() + await asyncio.gather(task, return_exceptions=True) + + assert await session.get_items() == [] + + async def test_runner_integration(agent: Agent): """Test that RedisSession works correctly with the agent Runner.""" session = await _create_test_session() @@ -1545,7 +1608,7 @@ async def reset(self) -> None: add_task.cancel() await asyncio.gather(add_task, return_exceptions=True) - assert exc_info.value.args == ("first-pre-exec-cancel",) + _assert_cancel_message(exc_info.value, "first-pre-exec-cancel") assert await session.get_items() == [] assert await client.ping() is True # type: ignore[misc] @@ -1609,7 +1672,7 @@ async def controlled_reset() -> None: await asyncio.gather(add_task, return_exceptions=True) assert await session.get_items() == [item] - assert exc_info.value.args == ("first-post-commit-cancel",) + _assert_cancel_message(exc_info.value, "first-post-commit-cancel") assert add_task.cancelled() assert await client.ping() is True # type: ignore[misc] @@ -1847,7 +1910,7 @@ async def fail_disconnect(nowait: bool = False) -> None: with pytest.raises(asyncio.CancelledError) as exc_info: await add_task - assert exc_info.value.args == ("caller-cancel",) + _assert_cancel_message(exc_info.value, "caller-cancel") monkeypatch.setattr(client, "pipeline", real_pipeline) assert await client.ping() is True # type: ignore[misc] assert await session.get_items() == [] @@ -1895,7 +1958,7 @@ async def fail_reset_after_release() -> None: with pytest.raises(asyncio.CancelledError) as exc_info: await add_task - assert exc_info.value.args == ("first-caller-cancel",) + _assert_cancel_message(exc_info.value, "first-caller-cancel") assert await client.ping() is True # type: ignore[misc] assert await session.get_items() == [] @@ -1961,7 +2024,7 @@ async def observed_reset() -> None: with pytest.raises(asyncio.CancelledError) as exc_info: await add_task - assert exc_info.value.args == ("caller-cancel",) + _assert_cancel_message(exc_info.value, "caller-cancel") assert not pool._in_use_connections monkeypatch.setattr(client, "pipeline", real_pipeline) assert await client.ping() is True # type: ignore[misc] @@ -2024,17 +2087,17 @@ async def controlled_reset() -> None: add_task.cancel() await asyncio.gather(add_task, return_exceptions=True) - assert exc_info.value.args == ("caller-cancel",) + _assert_cancel_message(exc_info.value, "caller-cancel") assert not pool._in_use_connections monkeypatch.setattr(client, "pipeline", real_pipeline) assert await client.ping() is True # type: ignore[misc] assert await session.get_items() == [] -async def test_overridden_pool_release_is_rejected_before_mutation( +async def test_transparent_overridden_pool_release_is_supported( monkeypatch: pytest.MonkeyPatch, ) -> None: - """Ambiguous pool ownership semantics must be rejected before acquiring a connection.""" + """A custom release that preserves pool ownership semantics remains supported.""" if not USE_FAKE_REDIS: pytest.skip("This test requires fakeredis") @@ -2048,24 +2111,63 @@ async def test_overridden_pool_release_is_rejected_before_mutation( real_release = pool.release release_calls = 0 - async def transfer_reborrow_then_fail(connection: Any) -> None: + async def transparent_release(connection: Any) -> None: nonlocal release_calls release_calls += 1 await real_release(connection) - borrowed = await pool.get_connection() - assert borrowed is connection - raise RuntimeError("release failed after reborrow") + + monkeypatch.setattr(pool, "release", transparent_release) + item: TResponseInputItem = {"role": "user", "content": "committed"} + + await session.add_items([item]) + + assert release_calls > 0 + assert not pool._in_use_connections + monkeypatch.setattr(pool, "release", real_release) + assert await session.get_items() == [item] + + +async def test_ambiguous_overridden_pool_release_is_detached_after_commit( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """A release-then-reborrow failure must not expose a committed batch as retryable.""" + if not USE_FAKE_REDIS: + pytest.skip("This test requires fakeredis") + + client = fakeredis.aioredis.FakeRedis(max_connections=1) + pool = client.connection_pool + session = RedisSession( + session_id="ambiguous_pool_release", + redis_client=cast("Redis", client), + key_prefix="test:", + ) + real_release = pool.release + release_calls = 0 + borrowed_connection: Any = None + + async def transfer_reborrow_then_fail(connection: Any) -> None: + nonlocal borrowed_connection, release_calls + release_calls += 1 + await real_release(connection) + if release_calls == 1: + borrowed_connection = await pool.get_connection() + assert borrowed_connection is connection + raise RuntimeError("release failed after reborrow") monkeypatch.setattr(pool, "release", transfer_reborrow_then_fail) - item: TResponseInputItem = {"role": "user", "content": "not committed"} + item: TResponseInputItem = {"role": "user", "content": "committed"} - with pytest.raises(RuntimeError, match="standard redis-py connection pool release method"): - await session.add_items([item]) + await session.add_items([item]) - assert release_calls == 0 + assert release_calls == 1 + assert borrowed_connection is not None assert not pool._in_use_connections + assert borrowed_connection not in pool._available_connections monkeypatch.setattr(pool, "release", real_release) - assert await session.get_items() == [] + replacement_connection = await pool.get_connection() + assert replacement_connection is not borrowed_connection + await real_release(replacement_connection) + assert await session.get_items() == [item] async def test_post_commit_disconnect_failure_is_not_retryable( @@ -2176,7 +2278,7 @@ def fail_close() -> None: if cancelled: with pytest.raises(asyncio.CancelledError) as exc_info: await add_task - assert exc_info.value.args == ("caller-cancel",) + _assert_cancel_message(exc_info.value, "caller-cancel") else: await add_task diff --git a/tests/extensions/memory/test_sqlalchemy_session.py b/tests/extensions/memory/test_sqlalchemy_session.py index 2a1fca969a..7019b000c4 100644 --- a/tests/extensions/memory/test_sqlalchemy_session.py +++ b/tests/extensions/memory/test_sqlalchemy_session.py @@ -7,6 +7,7 @@ from collections.abc import Iterable, Sequence from contextlib import asynccontextmanager from datetime import datetime, timedelta +from pathlib import Path from typing import Any, cast import pytest @@ -16,7 +17,7 @@ ResponseReasoningItemParam, Summary, ) -from sqlalchemy import insert, select, text, update +from sqlalchemy import event, insert, select, text, update from sqlalchemy.ext.asyncio import AsyncEngine, create_async_engine from sqlalchemy.sql import Select @@ -237,41 +238,101 @@ async def test_pop_from_empty_session(): assert popped is None -async def test_concurrent_pop_item_returns_each_row_once(monkeypatch: pytest.MonkeyPatch): - """A caller that loses the DELETE race must retry instead of returning the same row.""" - session = SQLAlchemySession.from_url("concurrent_pop", url=DB_URL, create_tables=False) - rows = { - 1: json.dumps({"role": "user", "content": "first"}), - 2: json.dumps({"role": "assistant", "content": "second"}), - } - first_reads = 0 - both_read_newest = asyncio.Event() - first_payload_reads = 0 - both_read_newest_payload = asyncio.Event() +async def test_concurrent_pop_item_returns_each_row_once(tmp_path): + """Concurrent atomic DELETE claims must return each stored row at most once.""" + engine = create_async_engine(f"sqlite+aiosqlite:///{tmp_path / 'concurrent_pop.db'}") + writer = SQLAlchemySession("concurrent_pop", engine=engine, create_tables=True) + other = SQLAlchemySession("concurrent_pop", engine=engine) + await writer.add_items( + [ + {"role": "user", "content": "first"}, + {"role": "assistant", "content": "second"}, + ] + ) + + tasks = [asyncio.create_task(session.pop_item()) for session in (writer, other)] + try: + popped = await asyncio.wait_for(asyncio.gather(*tasks), timeout=2) + finally: + for task in tasks: + if not task.done(): + task.cancel() + await asyncio.gather(*tasks, return_exceptions=True) + await engine.dispose() + + contents = {cast(dict[str, Any], item)["content"] for item in popped if item is not None} + assert contents == {"first", "second"} + + +async def test_sqlite_fallback_reserves_writer_before_tail_claim( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """SQLite without DELETE RETURNING must serialize the select-delete fallback.""" + engine = create_async_engine(f"sqlite+aiosqlite:///{tmp_path / 'fallback_pop.db'}") + writer = SQLAlchemySession("fallback_pop", engine=engine, create_tables=True) + other = SQLAlchemySession("fallback_pop", engine=engine) + statements: list[str] = [] + + def record_statement( + conn: Any, + cursor: Any, + statement: str, + parameters: Any, + context: Any, + executemany: bool, + ) -> None: + statements.append(statement) + + event.listen(engine.sync_engine, "before_cursor_execute", record_statement) + monkeypatch.setattr(engine.dialect, "delete_returning", False) + await writer.add_items([{"role": "user", "content": "only"}]) + + tasks = [asyncio.create_task(session.pop_item()) for session in (writer, other)] + try: + popped = await asyncio.wait_for(asyncio.gather(*tasks), timeout=2) + finally: + for task in tasks: + if not task.done(): + task.cancel() + await asyncio.gather(*tasks, return_exceptions=True) + event.remove(engine.sync_engine, "before_cursor_execute", record_statement) + await engine.dispose() + + assert sum(item is not None for item in popped) == 1 + assert [item.get("content") for item in popped if item is not None] == ["only"] + assert any(statement.strip().upper() == "BEGIN IMMEDIATE" for statement in statements) + + +async def test_pop_item_supports_unknown_delete_rowcount( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """A locked fallback claim must not depend on the DBAPI DELETE row count.""" + session = SQLAlchemySession.from_url("unknown_rowcount", url=DB_URL, create_tables=False) + transaction_exit_errors: list[type[BaseException] | None] = [] + delete_executed = False class FakeResult: - def __init__(self, scalar: Any = None, rowcount: int = -1) -> None: - self._scalar = scalar + def __init__(self, row: Any = None, rowcount: int = -1) -> None: + self._row = row self.rowcount = rowcount - def scalar_one_or_none(self) -> Any: - return self._scalar + def one_or_none(self) -> Any: + return self._row class FakeTransaction: async def __aenter__(self) -> None: return None - async def __aexit__(self, *args: Any) -> None: - return None + async def __aexit__( + self, + exc_type: type[BaseException] | None, + exc: BaseException | None, + traceback: Any, + ) -> None: + transaction_exit_errors.append(exc_type) class FakeSession: - def __init__(self) -> None: - # Model a repeatable-read transaction: selects keep reading this snapshot even - # after another transaction deletes from the durable rows mapping. - self.snapshot = dict(rows) - self.selected_row_id: int | None = None - self.selected_row: str | None = None - async def __aenter__(self) -> FakeSession: return self @@ -282,32 +343,12 @@ def begin(self) -> FakeTransaction: return FakeTransaction() async def execute(self, statement: Any) -> FakeResult: - nonlocal first_payload_reads, first_reads + nonlocal delete_executed if isinstance(statement, Select): - selected_columns = list(statement.selected_columns.keys()) - if selected_columns == ["id"]: - self.selected_row_id = max(self.snapshot, default=None) - if first_reads < 2: - first_reads += 1 - if first_reads == 2: - both_read_newest.set() - await both_read_newest.wait() - return FakeResult(self.selected_row_id) - assert selected_columns == ["message_data"] - assert self.selected_row_id is not None - self.selected_row = self.snapshot.get(self.selected_row_id) - if first_payload_reads < 2: - first_payload_reads += 1 - if first_payload_reads == 2: - both_read_newest_payload.set() - await both_read_newest_payload.wait() - return FakeResult(self.selected_row) - - assert self.selected_row_id is not None - if self.selected_row_id in rows: - del rows[self.selected_row_id] - return FakeResult(rowcount=1) - return FakeResult(rowcount=0) + assert statement._for_update_arg is not None + return FakeResult((1, json.dumps({"role": "user", "content": "claimed"}))) + delete_executed = True + return FakeResult(rowcount=-1) class FakeSessionFactory: def __call__(self) -> FakeSession: @@ -316,55 +357,46 @@ def __call__(self) -> FakeSession: async def tables_ready() -> None: return None - session_factory = FakeSessionFactory() monkeypatch.setattr(session, "_ensure_tables", tables_ready) - monkeypatch.setattr(session, "_session_factory", session_factory) + monkeypatch.setattr(session, "_session_factory", FakeSessionFactory()) + monkeypatch.setattr(session.engine.dialect, "delete_returning", False) - tasks = [asyncio.create_task(session.pop_item()) for _ in range(2)] try: - popped = await asyncio.wait_for(asyncio.gather(*tasks), timeout=1) + popped = await session.pop_item() finally: - both_read_newest.set() - both_read_newest_payload.set() - for task in tasks: - if not task.done(): - task.cancel() - await asyncio.gather(*tasks, return_exceptions=True) await session.engine.dispose() - contents = {cast(dict[str, Any], item)["content"] for item in popped if item is not None} - assert contents == {"first", "second"} - assert rows == {} + assert popped is not None + assert popped.get("content") == "claimed" + assert delete_executed is True + assert transaction_exit_errors == [None] -async def test_pop_item_rejects_unknown_delete_rowcount( +async def test_pop_item_retries_returning_claim_lost_to_concurrent_delete( monkeypatch: pytest.MonkeyPatch, ) -> None: - """An unknown DELETE row count cannot authorize a destructive read.""" - session = SQLAlchemySession.from_url("unknown_rowcount", url=DB_URL, create_tables=False) - transaction_exit_errors: list[type[BaseException] | None] = [] + """A lost DELETE RETURNING race must retry if older rows remain.""" + session = SQLAlchemySession.from_url("returning_retry", url=DB_URL, create_tables=False) + session_count = 0 class FakeResult: - def __init__(self, scalar: Any = None, rowcount: int = -1) -> None: - self._scalar = scalar - self.rowcount = rowcount + def __init__(self, value: Any = None) -> None: + self._value = value def scalar_one_or_none(self) -> Any: - return self._scalar + return self._value class FakeTransaction: async def __aenter__(self) -> None: return None - async def __aexit__( - self, - exc_type: type[BaseException] | None, - exc: BaseException | None, - traceback: Any, - ) -> None: - transaction_exit_errors.append(exc_type) + async def __aexit__(self, *args: Any) -> None: + return None class FakeSession: + def __init__(self, attempt: int) -> None: + self._attempt = attempt + async def __aenter__(self) -> FakeSession: return self @@ -375,12 +407,81 @@ def begin(self) -> FakeTransaction: return FakeTransaction() async def execute(self, statement: Any) -> FakeResult: - if isinstance(statement, Select): - selected_columns = list(statement.selected_columns.keys()) - if selected_columns == ["id"]: + if self._attempt == 1: + if isinstance(statement, Select): return FakeResult(1) - return FakeResult(json.dumps({"role": "user", "content": "unclaimed"})) - return FakeResult(rowcount=-1) + return FakeResult() + assert not isinstance(statement, Select) + return FakeResult(json.dumps({"role": "user", "content": "older"})) + + class FakeSessionFactory: + def __call__(self) -> FakeSession: + nonlocal session_count + session_count += 1 + return FakeSession(session_count) + + async def tables_ready() -> None: + return None + + monkeypatch.setattr(session, "_ensure_tables", tables_ready) + monkeypatch.setattr(session, "_session_factory", FakeSessionFactory()) + monkeypatch.setattr(session.engine.dialect, "delete_returning", True) + + try: + popped = await session.pop_item() + finally: + await session.engine.dispose() + + assert popped is not None + assert popped.get("content") == "older" + assert session_count == 2 + + +@pytest.mark.parametrize("operation", ["add", "pop", "clear"]) +async def test_mutation_cancellation_waits_for_transaction_exit( + monkeypatch: pytest.MonkeyPatch, + operation: str, +) -> None: + """Cancellation must wait until the transaction context finishes settling.""" + session = SQLAlchemySession.from_url( + f"transaction_cancellation_{operation}", + url=DB_URL, + create_tables=False, + ) + transaction_applied = asyncio.Event() + allow_return = asyncio.Event() + transaction_returned = False + + class FakeResult: + def scalar_one_or_none(self) -> Any: + if operation == "add": + return 1 + if operation == "pop": + return json.dumps({"role": "user", "content": "claimed"}) + return None + + class FakeTransaction: + async def __aenter__(self) -> None: + return None + + async def __aexit__(self, *args: Any) -> None: + nonlocal transaction_returned + transaction_applied.set() + await allow_return.wait() + transaction_returned = True + + class FakeSession: + async def __aenter__(self) -> FakeSession: + return self + + async def __aexit__(self, *args: Any) -> None: + return None + + def begin(self) -> FakeTransaction: + return FakeTransaction() + + async def execute(self, statement: Any) -> FakeResult: + return FakeResult() class FakeSessionFactory: def __call__(self) -> FakeSession: @@ -391,14 +492,35 @@ async def tables_ready() -> None: monkeypatch.setattr(session, "_ensure_tables", tables_ready) monkeypatch.setattr(session, "_session_factory", FakeSessionFactory()) + monkeypatch.setattr(session.engine.dialect, "delete_returning", True) + + if operation == "add": + task: asyncio.Task[Any] = asyncio.create_task( + session.add_items([{"role": "user", "content": "once"}]) + ) + elif operation == "pop": + task = asyncio.create_task(session.pop_item()) + else: + task = asyncio.create_task(session.clear_session()) try: - with pytest.raises(RuntimeError, match="confirm that exactly one item was deleted"): - await session.pop_item() + await transaction_applied.wait() + task.cancel() + await asyncio.sleep(0) + task.cancel() + await asyncio.sleep(0) + assert task.done() is False + allow_return.set() + with pytest.raises(asyncio.CancelledError): + await task finally: + allow_return.set() + if not task.done(): + task.cancel() + await asyncio.gather(task, return_exceptions=True) await session.engine.dispose() - assert transaction_exit_errors == [RuntimeError] + assert transaction_returned is True async def test_pop_item_skips_corrupt_most_recent(): From e3f7014b38e19b12abb5b1d3e962d60e2cf701e7 Mon Sep 17 00:00:00 2001 From: Kazuhiro Sera Date: Thu, 6 Aug 2026 11:52:04 +0900 Subject: [PATCH 6/7] fix: finalize session mutation atomicity --- .../memory/advanced_sqlite_session.py | 87 ++- src/agents/extensions/memory/dapr_session.py | 161 ++--- .../extensions/memory/mongodb_session.py | 107 ++- src/agents/extensions/memory/redis_session.py | 307 +++++++-- src/agents/memory/sqlite_session.py | 26 +- .../memory/test_advanced_sqlite_session.py | 107 +++ tests/extensions/memory/test_dapr_session.py | 110 --- .../extensions/memory/test_mongodb_session.py | 155 ++++- tests/extensions/memory/test_redis_session.py | 643 ++++++++++++++++-- tests/memory/test_session.py | 34 + 10 files changed, 1374 insertions(+), 363 deletions(-) diff --git a/src/agents/extensions/memory/advanced_sqlite_session.py b/src/agents/extensions/memory/advanced_sqlite_session.py index f9946e0fc6..c67e3f8a6a 100644 --- a/src/agents/extensions/memory/advanced_sqlite_session.py +++ b/src/agents/extensions/memory/advanced_sqlite_session.py @@ -226,9 +226,6 @@ async def get_items( """ session_limit = resolve_session_limit(limit, self.session_settings) - if branch_id is None: - branch_id = self._current_branch_id - def _decode_rows(rows: list[Any]) -> list[TResponseInputItem]: items: list[TResponseInputItem] = [] for (message_data,) in rows: @@ -242,6 +239,7 @@ def _decode_rows(rows: list[Any]) -> list[TResponseInputItem]: def _get_items_sync(): """Synchronous helper to get items for a specific branch.""" with self._locked_connection() as conn: + resolved_branch_id = self._resolve_read_branch(conn, branch_id) with closing(conn.cursor()) as cursor: # Get message IDs in correct order for this branch if session_limit is None: @@ -253,7 +251,7 @@ def _get_items_sync(): WHERE m.session_id = ? AND s.branch_id = ? ORDER BY s.sequence_number ASC """, - (self.session_id, branch_id), + (self.session_id, resolved_branch_id), ) return _decode_rows(cursor.fetchall()) @@ -272,7 +270,7 @@ def _get_items_sync(): ORDER BY s.sequence_number DESC LIMIT ? """, - (self.session_id, branch_id, window), + (self.session_id, resolved_branch_id, window), ) rows = cursor.fetchall() items = _decode_rows(list(reversed(rows))) @@ -293,7 +291,7 @@ def _get_items_sync(): ORDER BY s.sequence_number DESC LIMIT ? """, - (self.session_id, branch_id, session_limit), + (self.session_id, resolved_branch_id, session_limit), ) return _decode_rows(list(reversed(cursor.fetchall()))) @@ -506,8 +504,8 @@ def _capture_current_turn(self) -> tuple[int, str, int | None]: yields a different anchor. """ with self._locked_connection() as conn: + branch_id = self._resolve_read_branch(conn, None) with closing(conn.cursor()) as cursor: - branch_id = self._current_branch_id cursor.execute( """ SELECT COALESCE(MAX(user_turn_number), 0) @@ -580,6 +578,7 @@ def _get_current_turn_number(self) -> int: The current turn number for the active branch. """ with self._locked_connection() as conn: + branch_id = self._resolve_read_branch(conn, None) with closing(conn.cursor()) as cursor: cursor.execute( """ @@ -587,7 +586,7 @@ def _get_current_turn_number(self) -> int: FROM message_structure WHERE session_id = ? AND branch_id = ? """, - (self.session_id, self._current_branch_id), + (self.session_id, branch_id), ) result = cursor.fetchone() return result[0] if result else 0 @@ -1071,6 +1070,7 @@ async def list_branches(self) -> list[dict[str, Any]]: def _list_branches_sync(): """Synchronous helper to list all branches.""" with self._locked_connection() as conn: + current_branch_id = self._resolve_read_branch(conn, None) with closing(conn.cursor()) as cursor: cursor.execute( """ @@ -1095,7 +1095,7 @@ def _list_branches_sync(): "branch_id": branch_id, "message_count": msg_count, "user_turns": user_turns, - "is_current": branch_id == self._current_branch_id, + "is_current": branch_id == current_branch_id, "created_at": created_at, } ) @@ -1153,20 +1153,48 @@ def _ensure_session_clear_generations_table(self, conn: sqlite3.Connection) -> N (self.session_id,), ) - def _refresh_branch_after_external_clear(self, conn: sqlite3.Connection) -> None: + def _refresh_branch_after_external_clear( + self, + conn: sqlite3.Connection, + *, + initialize: bool = True, + ) -> None: """Reset a stale branch pointer after another session instance clears history.""" - self._ensure_session_clear_generations_table(conn) - generation = conn.execute( + if initialize: + self._ensure_session_clear_generations_table(conn) + else: + table_exists = conn.execute( + """ + SELECT 1 FROM sqlite_master + WHERE type = 'table' AND name = 'session_clear_generations' + """ + ).fetchone() + if table_exists is None: + return + + row = conn.execute( """ SELECT generation FROM session_clear_generations WHERE session_id = ? """, (self.session_id,), - ).fetchone()[0] + ).fetchone() + generation = row[0] if row is not None else 0 if generation != self._generation: self._generation = generation self._current_branch_id = "main" + def _resolve_read_branch( + self, + conn: sqlite3.Connection, + branch_id: str | None, + ) -> str: + """Resolve an implicit branch after synchronizing an external clear.""" + if branch_id is not None: + return branch_id + self._refresh_branch_after_external_clear(conn, initialize=False) + return self._current_branch_id + def _reserve_branch_id( self, cursor: sqlite3.Cursor, new_branch_id: str | None, from_turn_number: int ) -> str: @@ -1337,12 +1365,11 @@ async def get_conversation_turns(self, branch_id: str | None = None) -> list[dic - 'timestamp': When the turn was created - 'can_branch': Always True (all user messages can branch) """ - if branch_id is None: - branch_id = self._current_branch_id def _get_turns_sync(): """Synchronous helper to get conversation turns.""" with self._locked_connection() as conn: + resolved_branch_id = self._resolve_read_branch(conn, branch_id) with closing(conn.cursor()) as cursor: cursor.execute( f""" @@ -1356,7 +1383,7 @@ def _get_turns_sync(): AND ms.message_type = 'user' ORDER BY ms.branch_turn_number """, - (self.session_id, branch_id), + (self.session_id, resolved_branch_id), ) turns = [] @@ -1392,12 +1419,11 @@ async def find_turns_by_content( Returns: List of matching turns with same format as get_conversation_turns(). """ - if branch_id is None: - branch_id = self._current_branch_id def _search_sync(): """Synchronous helper to search turns by content.""" with self._locked_connection() as conn: + resolved_branch_id = self._resolve_read_branch(conn, branch_id) with closing(conn.cursor()) as cursor: cursor.execute( f""" @@ -1412,7 +1438,7 @@ def _search_sync(): AND am.message_data LIKE ? ORDER BY ms.branch_turn_number """, - (self.session_id, branch_id, f"%{search_term}%"), + (self.session_id, resolved_branch_id, f"%{search_term}%"), ) matches = [] @@ -1447,12 +1473,11 @@ async def get_conversation_by_turns( Returns: Dictionary mapping turn numbers to lists of message metadata. """ - if branch_id is None: - branch_id = self._current_branch_id def _get_conversation_sync(): """Synchronous helper to get conversation by turns.""" with self._locked_connection() as conn: + resolved_branch_id = self._resolve_read_branch(conn, branch_id) with closing(conn.cursor()) as cursor: cursor.execute( """ @@ -1461,7 +1486,7 @@ def _get_conversation_sync(): WHERE session_id = ? AND branch_id = ? ORDER BY sequence_number """, - (self.session_id, branch_id), + (self.session_id, resolved_branch_id), ) turns: dict[int, list[dict[str, str | None]]] = {} @@ -1483,12 +1508,11 @@ async def get_tool_usage(self, branch_id: str | None = None) -> list[tuple[str, Returns: List of tuples containing (tool_name, usage_count, turn_number). """ - if branch_id is None: - branch_id = self._current_branch_id def _get_tool_usage_sync(): """Synchronous helper to get tool usage statistics.""" with self._locked_connection() as conn: + resolved_branch_id = self._resolve_read_branch(conn, branch_id) with closing(conn.cursor()) as cursor: cursor.execute( """ @@ -1523,9 +1547,9 @@ def _get_tool_usage_sync(): """, ( self.session_id, - branch_id, + resolved_branch_id, self.session_id, - branch_id, + resolved_branch_id, ), ) return cursor.fetchall() @@ -1605,12 +1629,10 @@ async def get_turn_usage( Dictionary with usage data for specific turn, or list of dictionaries for all turns. """ - if branch_id is None: - branch_id = self._current_branch_id - def _get_turn_usage_sync(): """Synchronous helper to get turn usage statistics.""" with self._locked_connection() as conn: + resolved_branch_id = self._resolve_read_branch(conn, branch_id) if user_turn_number is not None: query = """ SELECT requests, input_tokens, output_tokens, total_tokens, @@ -1620,7 +1642,10 @@ def _get_turn_usage_sync(): """ with closing(conn.cursor()) as cursor: - cursor.execute(query, (self.session_id, branch_id, user_turn_number)) + cursor.execute( + query, + (self.session_id, resolved_branch_id, user_turn_number), + ) row = cursor.fetchone() if row: @@ -1659,7 +1684,7 @@ def _get_turn_usage_sync(): """ with closing(conn.cursor()) as cursor: - cursor.execute(query, (self.session_id, branch_id)) + cursor.execute(query, (self.session_id, resolved_branch_id)) results = [] for row in cursor.fetchall(): # Parse JSON details if present diff --git a/src/agents/extensions/memory/dapr_session.py b/src/agents/extensions/memory/dapr_session.py index d13c042664..20b300ce3c 100644 --- a/src/agents/extensions/memory/dapr_session.py +++ b/src/agents/extensions/memory/dapr_session.py @@ -50,7 +50,6 @@ coerce_session_settings, resolve_session_limit, ) -from ...memory.sqlite_session import _await_mutation # Type alias for consistency levels ConsistencyLevel = Literal["eventual", "strong"] @@ -335,67 +334,6 @@ 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] - await _await_mutation(self._add_items_locked(serialized_items)) - - async def _add_items_locked(self, serialized_items: list[str]) -> None: - """Persist one batch while the caller retains the session lock.""" - # 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 - response = await self._dapr_client.get_state( - store_name=self._state_store_name, - key=self._messages_key, - state_metadata=self._get_read_metadata(), - ) - existing_messages = self._decode_messages_for_update(response.data) - updated_messages = existing_messages + serialized_items - messages_json = json.dumps(updated_messages, separators=(",", ":")) - etag = response.etag - try: - await self._dapr_client.save_state( - store_name=self._state_store_name, - key=self._messages_key, - value=messages_json, - etag=etag, - state_metadata=self._get_metadata(), - options=self._get_state_options(concurrency=Concurrency.first_write), - ) - return - except Exception as error: - should_retry = await self._handle_concurrency_conflict(error, attempt) - if should_retry: - continue - raise - - 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._lock: - self._check_not_closed() - return await _await_mutation(self._pop_item_locked()) - - async def _pop_item_locked(self) -> TResponseInputItem | None: - """Claim one item while the caller retains the session lock.""" - while True: attempt = 0 while True: attempt += 1 @@ -404,12 +342,10 @@ async def _pop_item_locked(self) -> TResponseInputItem | None: key=self._messages_key, state_metadata=self._get_read_metadata(), ) - messages = self._decode_messages(response.data) - if not messages: - return None - last_item = messages.pop() - messages_json = json.dumps(messages, separators=(",", ":")) - etag = getattr(response, "etag", None) or None + existing_messages = self._decode_messages_for_update(response.data) + updated_messages = existing_messages + serialized_items + messages_json = json.dumps(updated_messages, separators=(",", ":")) + etag = response.etag try: await self._dapr_client.save_state( store_name=self._state_store_name, @@ -425,29 +361,82 @@ async def _pop_item_locked(self) -> TResponseInputItem | None: if should_retry: continue raise - try: - if isinstance(last_item, str): - return await self._deserialize_item(last_item) - return last_item # type: ignore[no-any-return] - except (json.JSONDecodeError, TypeError): - continue + + # 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. + + Returns: + The most recent item if it exists, None if the session is empty + """ + async with self._lock: + self._check_not_closed() + while True: + attempt = 0 + while True: + attempt += 1 + response = await self._dapr_client.get_state( + store_name=self._state_store_name, + key=self._messages_key, + state_metadata=self._get_read_metadata(), + ) + messages = self._decode_messages(response.data) + if not messages: + return None + last_item = messages.pop() + messages_json = json.dumps(messages, separators=(",", ":")) + etag = getattr(response, "etag", None) or None + try: + await self._dapr_client.save_state( + store_name=self._state_store_name, + key=self._messages_key, + value=messages_json, + etag=etag, + state_metadata=self._get_metadata(), + options=self._get_state_options(concurrency=Concurrency.first_write), + ) + break + except Exception as error: + should_retry = await self._handle_concurrency_conflict(error, attempt) + if should_retry: + continue + raise + try: + if isinstance(last_item, str): + return await self._deserialize_item(last_item) + return last_item # type: ignore[no-any-return] + except (json.JSONDecodeError, TypeError): + continue async def clear_session(self) -> None: """Clear all items for this session.""" async with self._lock: self._check_not_closed() - await _await_mutation(self._clear_session_locked()) - - async def _clear_session_locked(self) -> None: - """Delete authoritative history while the caller retains the session lock.""" - # The messages key is the authoritative history boundary. Keep the - # ancillary metadata key so clear_session remains one state-store - # mutation instead of partially failing after history is cleared. - await self._dapr_client.delete_state( - store_name=self._state_store_name, - key=self._messages_key, - options=self._get_state_options(), - ) + # Delete messages and metadata keys + await self._dapr_client.delete_state( + store_name=self._state_store_name, + key=self._messages_key, + options=self._get_state_options(), + ) + + await self._dapr_client.delete_state( + store_name=self._state_store_name, + key=self._metadata_key, + options=self._get_state_options(), + ) async def close(self) -> None: """Close the Dapr client connection. diff --git a/src/agents/extensions/memory/mongodb_session.py b/src/agents/extensions/memory/mongodb_session.py index c417cec937..3887409ca1 100644 --- a/src/agents/extensions/memory/mongodb_session.py +++ b/src/agents/extensions/memory/mongodb_session.py @@ -51,6 +51,7 @@ from pymongo.asynchronous.collection import AsyncCollection from pymongo.asynchronous.mongo_client import AsyncMongoClient from pymongo.driver_info import DriverInfo + from pymongo.read_preferences import ReadPreference except ImportError as e: raise_optional_dependency_error( "MongoDBSession", @@ -251,9 +252,9 @@ async def _ensure_indexes(self) -> None: # sessions: unique index on session_id. await self._sessions.create_index("session_id", unique=True) - # messages: compound index for efficient per-session retrieval and - # sorting by the explicit seq counter. - await self._messages.create_index([("session_id", 1), ("seq", 1)]) + # messages: compound index for efficient active-generation retrieval + # and sorting by the explicit seq counter. + await self._messages.create_index([("session_id", 1), ("generation", 1), ("seq", 1)]) self._mark_init_done() @@ -269,6 +270,21 @@ async def _deserialize_item(self, raw: str) -> TResponseInputItem: """Deserialize a JSON string to an item. Can be overridden by subclasses.""" return json.loads(raw) # type: ignore[no-any-return] + async def _get_generation(self) -> int: + """Return the authoritative history generation for this session.""" + sessions = self._sessions.with_options(read_preference=ReadPreference.PRIMARY) + docs = await sessions.find({"session_id": self.session_id}).limit(1).to_list() + if not docs: + return 0 + generation = docs[0].get("_generation", 0) + return generation if isinstance(generation, int) else 0 + + def _generation_query(self, generation: int) -> dict[str, Any]: + """Match the active generation while retaining legacy generation-zero data.""" + if generation == 0: + return {"generation": {"$in": [0, None]}} + return {"generation": generation} + # ------------------------------------------------------------------ # Session protocol implementation # ------------------------------------------------------------------ @@ -293,7 +309,11 @@ async def get_items(self, limit: int | None = None) -> list[TResponseInputItem]: if session_limit is not None and session_limit <= 0: return [] - query = {"session_id": self.session_id} + generation = await self._get_generation() + query = { + "session_id": self.session_id, + **self._generation_query(generation), + } async def _decode_docs(docs: list[Any]) -> list[TResponseInputItem]: items: list[TResponseInputItem] = [] @@ -346,7 +366,11 @@ async def _add_items(self, serialized_items: list[str]) -> None: result = await self._sessions.find_one_and_update( {"session_id": self.session_id}, { - "$setOnInsert": {"session_id": self.session_id, "created_at": now}, + "$setOnInsert": { + "session_id": self.session_id, + "created_at": now, + "_generation": 0, + }, "$set": {"updated_at": now}, "$inc": {"_seq": len(serialized_items)}, }, @@ -356,6 +380,9 @@ async def _add_items(self, serialized_items: list[str]) -> None: next_seq: int = (result["_seq"] if result else len(serialized_items)) - len( serialized_items ) + generation = result.get("_generation", 0) if result else 0 + if not isinstance(generation, int): + generation = 0 # One document is the commit boundary for the logical batch. This keeps # standalone MongoDB deployments failure-atomic without requiring transactions. @@ -363,6 +390,7 @@ async def _add_items(self, serialized_items: list[str]) -> None: { "session_id": self.session_id, "seq": next_seq + len(serialized_items) - 1, + "generation": generation, "message_data": serialized_items, } ) @@ -384,11 +412,19 @@ async def _pop_item(self) -> TResponseInputItem | None: is returned. This matches :meth:`get_items`, which also skips malformed entries, so one bad record cannot make a non-empty session look empty. """ + generation = await self._get_generation() + # Retry cleanup left by a prior post-claim failure. Empty markers are # never model-visible or claimable, so cleanup failure must not block a # later valid tail claim. try: - await self._messages.delete_many({"session_id": self.session_id, "message_data": []}) + await self._messages.delete_many( + { + "session_id": self.session_id, + **self._generation_query(generation), + "message_data": [], + } + ) except asyncio.CancelledError: raise except Exception: @@ -398,6 +434,7 @@ async def _pop_item(self) -> TResponseInputItem | None: doc = await self._messages.find_one_and_update( { "session_id": self.session_id, + **self._generation_query(generation), "message_data": {"$ne": []}, }, [ @@ -427,7 +464,22 @@ async def _pop_item(self) -> TResponseInputItem | None: return_document=False, ) if doc is None: + current_generation = await self._get_generation() + if current_generation != generation: + generation = current_generation + continue return None + + current_generation = await self._get_generation() + if current_generation != generation: + try: + await self._messages.delete_one({"_id": doc["_id"]}) + except asyncio.CancelledError: + raise + except Exception: + pass + generation = current_generation + continue raw = doc.get("message_data") if isinstance(raw, list): @@ -465,13 +517,42 @@ async def clear_session(self) -> None: await _await_mutation(self._clear_session()) async def _clear_session(self) -> None: - """Clear all items for this session.""" - await self._messages.delete_many({"session_id": self.session_id}) - # Keep the ancillary metadata document and its sequence counter. The - # messages collection is the authoritative history boundary, and a - # second delete would make clear_session partially fail after history - # was already cleared. Retaining the counter also prevents a concurrent - # add from losing its ordering authority across this clear. + """Advance the authoritative generation and clean obsolete history.""" + now = datetime.now(timezone.utc) + result = await self._sessions.find_one_and_update( + {"session_id": self.session_id}, + { + "$setOnInsert": { + "session_id": self.session_id, + "created_at": now, + "_seq": 0, + }, + "$set": {"updated_at": now}, + "$inc": {"_generation": 1}, + }, + upsert=True, + return_document=True, + ) + generation = result.get("_generation", 1) if result else 1 + if not isinstance(generation, int): + generation = 1 + + # The metadata update above is the single-document clear boundary. + # Obsolete batches are no longer visible, so physical deletion is best effort. + try: + await self._messages.delete_many( + { + "session_id": self.session_id, + "$or": [ + {"generation": {"$lt": generation}}, + {"generation": {"$exists": False}}, + ], + } + ) + except asyncio.CancelledError: + raise + except Exception: + pass # ------------------------------------------------------------------ # Lifecycle helpers diff --git a/src/agents/extensions/memory/redis_session.py b/src/agents/extensions/memory/redis_session.py index 694d6888fc..953b1cf683 100644 --- a/src/agents/extensions/memory/redis_session.py +++ b/src/agents/extensions/memory/redis_session.py @@ -31,8 +31,10 @@ try: import redis.asyncio as redis + import redis.asyncio.connection as redis_connection from redis.asyncio import BlockingConnectionPool, Redis - from redis.exceptions import ResponseError, WatchError + from redis.event import AsyncAfterConnectionReleasedEvent + from redis.exceptions import ConnectionError as RedisConnectionError, ResponseError, WatchError except ImportError as e: raise_optional_dependency_error( "RedisSession", @@ -50,10 +52,13 @@ ) from ...memory.sqlite_session import _await_mutation +_redis_connection_api: Any = redis_connection + @dataclass class _PipelineAttemptOutcome: committed: bool + retryable_watch_conflict: bool operation_error: BaseException | None cleanup_error: BaseException | None settled: bool @@ -64,34 +69,214 @@ class _PipelineConnectionPool: def __init__(self, pool: Any): self._pool = pool + self.connection: Any | None = None self.release_started = False self.release_completed = False + self._checkout_recorded_used = False def __getattr__(self, name: str) -> Any: return getattr(self._pool, name) + async def get_connection(self, *args: Any, **kwargs: Any) -> Any: + """Retain an acquired identity before validating it for pipeline use.""" + del args, kwargs + + async def acquire() -> Any: + if isinstance(self._pool, BlockingConnectionPool): + start_time_acquired = time.monotonic() + has_timing_observability = all( + hasattr(_redis_connection_api, name) + for name in ( + "get_pool_name", + "record_connection_create_time", + "record_connection_wait_time", + ) + ) + try: + async with self._pool._condition: + await asyncio.wait_for( + self._pool._condition.wait_for(self._pool.can_get_connection), + timeout=self._pool.timeout, + ) + maybe_pool_lock = getattr(self._pool, "_maybe_pool_lock", None) + if has_timing_observability: + connections_before = len(self._pool._available_connections) + len( + self._pool._in_use_connections + ) + start_time_created = time.monotonic() + if maybe_pool_lock is None: + connection = self._pool.get_available_connection() + self.connection = connection + else: + async with maybe_pool_lock(): + connection = self._pool.get_available_connection() + self.connection = connection + if has_timing_observability: + connections_after = len(self._pool._available_connections) + len( + self._pool._in_use_connections + ) + is_created = connections_after > connections_before + except asyncio.TimeoutError as exc: + raise RedisConnectionError("No connection available.") from exc + await self._pool.ensure_connection(connection) + if has_timing_observability: + if is_created: + await _redis_connection_api.record_connection_create_time( + connection_pool=self._pool, + duration_seconds=time.monotonic() - start_time_created, + ) + await _redis_connection_api.record_connection_wait_time( + pool_name=_redis_connection_api.get_pool_name(self._pool), + duration_seconds=time.monotonic() - start_time_acquired, + ) + return connection + + has_observability = hasattr(_redis_connection_api, "record_connection_count") + async with self._pool._lock: + if has_observability: + connections_before = len(self._pool._available_connections) + len( + self._pool._in_use_connections + ) + start_time_created = time.monotonic() + connection = self._pool.get_available_connection() + self.connection = connection + if has_observability: + connections_after = len(self._pool._available_connections) + len( + self._pool._in_use_connections + ) + is_created = connections_after > connections_before + else: + await self._pool.ensure_connection(connection) + return connection + + pool_name = _redis_connection_api.get_pool_name(self._pool) + if is_created: + await _redis_connection_api.record_connection_count( + pool_name=pool_name, + connection_state=_redis_connection_api.ConnectionState.USED, + counter=1, + ) + else: + await _redis_connection_api.record_connection_count( + pool_name=pool_name, + connection_state=_redis_connection_api.ConnectionState.IDLE, + counter=-1, + ) + await _redis_connection_api.record_connection_count( + pool_name=pool_name, + connection_state=_redis_connection_api.ConnectionState.USED, + counter=1, + ) + self._checkout_recorded_used = True + await self._pool.ensure_connection(connection) + if is_created: + await _redis_connection_api.record_connection_create_time( + connection_pool=self._pool, + duration_seconds=time.monotonic() - start_time_created, + ) + return connection + + acquisition = asyncio.create_task(acquire()) + cancellation: asyncio.CancelledError | None = None + while not acquisition.done(): + try: + await asyncio.wait({acquisition}) + except asyncio.CancelledError as exc: + if cancellation is None: + cancellation = exc + if self.connection is None: + acquisition.cancel() + + try: + connection = acquisition.result() + except BaseException: + if cancellation is not None: + raise cancellation from None + raise + if cancellation is not None: + raise cancellation from None + return connection + + async def record_discard(self) -> None: + """Balance supported pool observability when a checked-out connection is removed.""" + if not self._checkout_recorded_used: + return + await _redis_connection_api.record_connection_count( + pool_name=_redis_connection_api.get_pool_name(self._pool), + connection_state=_redis_connection_api.ConnectionState.USED, + counter=-1, + ) + self._checkout_recorded_used = False + async def notify_capacity_available(self) -> None: if isinstance(self._pool, BlockingConnectionPool): async with self._pool._condition: self._pool._condition.notify() + def _install_release_transfer_listener(self, connection: Any) -> Any: + """Mark the exact point where redis-py exposes an identity for reuse.""" + owner = self + + class ReleaseTransferListener: + async def listen(self, event: Any) -> None: + if event.connection is connection: + owner.release_completed = True + owner._checkout_recorded_used = False + + listener = ReleaseTransferListener() + dispatcher = self._pool._event_dispatcher + with dispatcher._lock: + listeners = dispatcher._event_listeners_mapping.get( + AsyncAfterConnectionReleasedEvent, [] + ) + dispatcher._event_listeners_mapping[AsyncAfterConnectionReleasedEvent] = [ + listener, + *listeners, + ] + return listener + + def _remove_release_transfer_listener(self, listener: Any) -> None: + dispatcher = self._pool._event_dispatcher + with dispatcher._lock: + listeners = dispatcher._event_listeners_mapping.get( + AsyncAfterConnectionReleasedEvent, [] + ) + dispatcher._event_listeners_mapping[AsyncAfterConnectionReleasedEvent] = [ + current for current in listeners if current is not listener + ] + async def release(self, connection: Any) -> None: + self.connection = connection self.release_started = True was_in_use = connection in self._pool._in_use_connections + if was_in_use: + try: + if connection.should_reconnect(): + # Let Pipeline.reset() finish clearing its local state, then + # let _finish_pipeline() detach this retained identity without + # exposing it to shared-pool release listeners. + return + except BaseException: + pass + transfer_listener = self._install_release_transfer_listener(connection) try: await self._pool.release(connection) except BaseException: - if was_in_use and connection in self._pool._available_connections: - self.release_completed = True + if self.release_completed and connection in self._pool._available_connections: await self.notify_capacity_available() raise else: self.release_completed = True + self._checkout_recorded_used = False + finally: + self._remove_release_transfer_listener(transfer_listener) async def _finish_pipeline( pipe: Any, connection_pool: _PipelineConnectionPool, + *, + discard_connection: bool = False, ) -> tuple[BaseException | None, bool, Any | None]: """Reset a pipeline and prove whether its connection left pipeline ownership.""" if connection_pool.release_completed: @@ -99,7 +284,7 @@ async def _finish_pipeline( return None, True, None reset_error: BaseException | None = None - if not connection_pool.release_started: + if not connection_pool.release_started and not discard_connection: try: await pipe.reset() except BaseException as exc: @@ -109,56 +294,36 @@ async def _finish_pipeline( pipe.connection = None return reset_error, True, None - connection = getattr(pipe, "connection", None) + connection = connection_pool.connection or getattr(pipe, "connection", None) if connection is None: return reset_error, True, None - if connection_pool.release_started: - # The first release may have transferred and reborrowed the connection - # before failing. Its ownership is ambiguous, so never release it again. - connection_pool._pool._in_use_connections.discard(connection) - while connection in connection_pool._pool._available_connections: - connection_pool._pool._available_connections.remove(connection) - try: - connection._close() - except BaseException as close_error: - pipe.connection = None - connection_pool.release_completed = True - await connection_pool.notify_capacity_available() - return close_error, True, connection - await connection_pool.notify_capacity_available() - pipe.connection = None - connection_pool.release_completed = True - return reset_error, True, None - + # Any retained identity at this point has no proven pool transfer. Detach it + # directly instead of starting or repeating a shared-pool release whose + # listeners could reborrow the identity before reporting a failure. + connection_pool._pool._in_use_connections.discard(connection) + while connection in connection_pool._pool._available_connections: + connection_pool._pool._available_connections.remove(connection) + metrics_error: BaseException | None = None try: - connection.mark_for_reconnect() - except BaseException as reconnect_error: - return reconnect_error, False, None - + await connection_pool.record_discard() + except BaseException as exc: + metrics_error = exc try: - await connection_pool.release(connection) - except BaseException as release_error: - if not connection_pool.release_completed: - # A custom release implementation may transfer and reborrow the - # connection before failing. Remove ambiguous pool membership and - # quarantine the connection instead of leaving it checked out. - connection_pool._pool._in_use_connections.discard(connection) - while connection in connection_pool._pool._available_connections: - connection_pool._pool._available_connections.remove(connection) - try: - connection._close() - except BaseException as close_error: - pipe.connection = None - connection_pool.release_completed = True - await connection_pool.notify_capacity_available() - return close_error, True, connection - await connection_pool.notify_capacity_available() - connection_pool.release_completed = True + connection._close() + except BaseException as close_error: pipe.connection = None - return release_error, connection_pool.release_completed, None - + connection_pool.release_completed = True + await connection_pool.notify_capacity_available() + return close_error, True, connection + if metrics_error is not None: + pipe.connection = None + connection_pool.release_completed = True + await connection_pool.notify_capacity_available() + return metrics_error, True, None + await connection_pool.notify_capacity_available() pipe.connection = None + connection_pool.release_completed = True return reset_error, True, None @@ -170,21 +335,23 @@ async def _await_pipeline_attempt( cancellation: asyncio.CancelledError | None = None attempt_cancelled = False - while True: + while not attempt.done(): try: - outcome = await asyncio.shield(attempt) + await asyncio.wait({attempt}) except asyncio.CancelledError as exc: - if attempt.done() and attempt.cancelled(): - if cancellation is not None: - raise cancellation from None - raise if cancellation is None: cancellation = exc if not completion_owned.is_set() and not attempt_cancelled: attempt.cancel() attempt_cancelled = True - else: - return outcome, cancellation + + try: + outcome = attempt.result() + except BaseException: + if cancellation is not None: + raise cancellation from None + raise + return outcome, cancellation class RedisSession(SessionABC): @@ -309,9 +476,12 @@ async def _write_items_attempt( ) -> _PipelineAttemptOutcome: """Run one watched write attempt and finish its pipeline before returning.""" committed = False + retryable_watch_conflict = False operation_error: BaseException | None = None + discard_connection = False batch_response_index: int | None = None raise_first_error = pipe.raise_first_error + parse_response = pipe.parse_response raw_connection_pool = pipe.connection_pool connection_pool = _PipelineConnectionPool(raw_connection_pool) pipe.connection_pool = connection_pool @@ -330,6 +500,18 @@ def raise_first_error_and_mark(*args: Any, **kwargs: Any) -> Any: committed = True return result + parsed_transaction_responses = 0 + exec_response_position = 0 + + async def parse_response_and_classify(*args: Any, **kwargs: Any) -> Any: + nonlocal parsed_transaction_responses, retryable_watch_conflict + response_position = parsed_transaction_responses + parsed_transaction_responses += 1 + response = await parse_response(*args, **kwargs) + if response_position == exec_response_position and response is None: + retryable_watch_conflict = True + return response + try: try: await pipe.watch(*keys) @@ -367,6 +549,8 @@ def raise_first_error_and_mark(*args: Any, **kwargs: Any) -> Any: pipe.pexpireat(key, expiration_time_ms) pipe.raise_first_error = raise_first_error_and_mark + exec_response_position = len(pipe.command_stack) + 1 + pipe.parse_response = parse_response_and_classify completion_owned.set() await pipe.execute() committed = True @@ -374,11 +558,19 @@ def raise_first_error_and_mark(*args: Any, **kwargs: Any) -> Any: operation_error = exc except BaseException as exc: operation_error = exc + if isinstance(exc, asyncio.CancelledError) and not completion_owned.is_set(): + # An immediate WATCH command may have been sent without its + # response being consumed. Never return that connection to + # shared pool reuse or invoke release listeners with it. + discard_connection = True finally: completion_owned.set() pipe.raise_first_error = raise_first_error + pipe.parse_response = parse_response cleanup_error, settled, detached_connection = await _finish_pipeline( - pipe, connection_pool + pipe, + connection_pool, + discard_connection=discard_connection, ) if detached_connection is not None: self._detached_connections.add(detached_connection) @@ -386,6 +578,7 @@ def raise_first_error_and_mark(*args: Any, **kwargs: Any) -> Any: return _PipelineAttemptOutcome( committed=committed, + retryable_watch_conflict=retryable_watch_conflict, operation_error=operation_error, cleanup_error=cleanup_error, settled=settled, @@ -417,7 +610,7 @@ async def _write_items( raise cancellation if outcome.cleanup_error is not None: raise outcome.cleanup_error - if isinstance(outcome.operation_error, WatchError): + if outcome.retryable_watch_conflict: continue if outcome.operation_error is not None: raise outcome.operation_error diff --git a/src/agents/memory/sqlite_session.py b/src/agents/memory/sqlite_session.py index 871f5f04ee..fc9f3fdb8f 100644 --- a/src/agents/memory/sqlite_session.py +++ b/src/agents/memory/sqlite_session.py @@ -20,30 +20,22 @@ async def _await_mutation(awaitable: Awaitable[_T]) -> _T: """Wait for a mutation outcome despite repeated caller cancellation.""" task = asyncio.ensure_future(awaitable) cancellation: asyncio.CancelledError | None = None - while True: + while not task.done(): try: - result = await asyncio.shield(task) + await asyncio.wait({task}) except asyncio.CancelledError as exc: if cancellation is None: cancellation = exc - if not task.done(): - continue - except BaseException: - if cancellation is not None: - raise cancellation from None - raise - else: - if cancellation is not None: - raise cancellation from None - return result - try: - result = task.result() - except BaseException: - assert cancellation is not None + try: + result = task.result() + except BaseException: + if cancellation is not None: raise cancellation from None - assert cancellation is not None + raise + if cancellation is not None: raise cancellation from None + return result class SQLiteSession(SessionABC): diff --git a/tests/extensions/memory/test_advanced_sqlite_session.py b/tests/extensions/memory/test_advanced_sqlite_session.py index cd4397423a..389d174e68 100644 --- a/tests/extensions/memory/test_advanced_sqlite_session.py +++ b/tests/extensions/memory/test_advanced_sqlite_session.py @@ -3437,6 +3437,113 @@ async def test_external_clear_resets_stale_branch_before_pop(tmp_path: Path): clearer.close() +@pytest.mark.parametrize( + "read_path", + ["items", "turns", "search", "conversation", "tools", "usage", "branches"], +) +async def test_external_clear_resets_stale_branch_before_default_reads( + tmp_path: Path, + usage_data: Usage, + read_path: str, +): + """Default reads must recover from a stale branch pointer after an external clear.""" + db_path = tmp_path / f"external_clear_read_generation_{read_path}.db" + session_id = f"external_clear_read_generation_{read_path}" + stale = AdvancedSQLiteSession( + session_id=session_id, + db_path=db_path, + create_tables=True, + ) + clearer = AdvancedSQLiteSession(session_id=session_id, db_path=db_path) + + try: + await stale.add_items( + [ + {"role": "user", "content": "old question"}, + {"role": "assistant", "content": "old answer"}, + {"role": "user", "content": "old follow-up"}, + ] + ) + await stale.create_branch_from_turn(2, "stale") + assert stale._current_branch_id == "stale" + + await clearer.clear_session() + new_items: list[TResponseInputItem] = [ + {"role": "user", "content": "new main question"}, + { + "type": "function_call", + "name": "lookup", + "arguments": '{"query": "new"}', + "call_id": "lookup-new-main", + }, + {"role": "assistant", "content": "new main answer"}, + ] + await clearer.add_items(new_items) + await clearer.store_run_usage(create_mock_run_result(usage_data)) + + if read_path == "items": + assert await stale.get_items() == new_items + elif read_path == "turns": + assert [turn["full_content"] for turn in await stale.get_conversation_turns()] == [ + "new main question" + ] + elif read_path == "search": + assert [turn["full_content"] for turn in await stale.find_turns_by_content("new")] == [ + "new main question" + ] + elif read_path == "conversation": + assert set(await stale.get_conversation_by_turns()) == {1} + elif read_path == "tools": + assert await stale.get_tool_usage() == [("lookup", 1, 1)] + elif read_path == "usage": + assert await stale.get_turn_usage(1) == { + "requests": 1, + "input_tokens": 50, + "output_tokens": 30, + "total_tokens": 80, + "input_tokens_details": {"cache_write_tokens": 0, "cached_tokens": 10}, + "output_tokens_details": {"reasoning_tokens": 5}, + } + else: + assert [ + (branch["branch_id"], branch["is_current"]) + for branch in await stale.list_branches() + ] == [("main", True)] + + assert stale._current_branch_id == "main" + finally: + stale.close() + clearer.close() + + +async def test_default_read_does_not_initialize_clear_generation_table(tmp_path: Path): + """Reading a legacy database must not create the clear-generation table.""" + session = AdvancedSQLiteSession( + session_id="legacy_generation_read", + db_path=tmp_path / "legacy_generation_read.db", + create_tables=True, + ) + + try: + await session.add_items([{"role": "user", "content": "legacy history"}]) + with session._locked_connection() as conn: + conn.execute("DROP TABLE session_clear_generations") + conn.commit() + + assert await session.get_items() == [{"role": "user", "content": "legacy history"}] + + with session._locked_connection() as conn: + table_exists = conn.execute( + """ + SELECT 1 FROM sqlite_master + WHERE type = 'table' AND name = 'session_clear_generations' + """ + ).fetchone() + assert table_exists is None + finally: + session.close() + + async def test_switch_validation_cancellation_waits_for_generation_commit( tmp_path: Path, monkeypatch: pytest.MonkeyPatch, diff --git a/tests/extensions/memory/test_dapr_session.py b/tests/extensions/memory/test_dapr_session.py index c1a765fc10..702cb7388c 100644 --- a/tests/extensions/memory/test_dapr_session.py +++ b/tests/extensions/memory/test_dapr_session.py @@ -1,6 +1,5 @@ from __future__ import annotations -import asyncio import json from typing import Any from unittest.mock import Mock @@ -158,35 +157,6 @@ async def save_state( ) -class FailingMetadataFakeDaprClient(FakeDaprClient): - """Fake client that rejects the next metadata write.""" - - def __init__(self) -> None: - super().__init__() - self.fail_metadata_save = True - - async def save_state( - self, - store_name: str, - key: str, - value: str | bytes, - state_metadata: dict[str, str] | None = None, - options: Any = None, - etag: str | None = None, - ) -> None: - if key.endswith(":metadata") and self.fail_metadata_save: - self.fail_metadata_save = False - raise RuntimeError("metadata write failed") - await super().save_state( - store_name=store_name, - key=key, - value=value, - state_metadata=state_metadata, - options=options, - etag=etag, - ) - - @pytest.fixture def conflict_dapr_client() -> ConflictFakeDaprClient: """Fixture for fake client that forces concurrency conflicts.""" @@ -251,72 +221,11 @@ async def test_dapr_session_direct_ops(fake_dapr_client: FakeDaprClient): await session.clear_session() retrieved_after_clear = await session.get_items() assert len(retrieved_after_clear) == 0 - assert session._metadata_key in fake_dapr_client._state finally: await session.close() -@pytest.mark.parametrize("operation", ["add", "pop", "clear"]) -async def test_mutation_cancellation_waits_for_authoritative_outcome( - fake_dapr_client: FakeDaprClient, - monkeypatch: pytest.MonkeyPatch, - operation: str, -) -> None: - """Cancellation must wait after the state store applies a history mutation.""" - session = await _create_test_session(fake_dapr_client) - item: TResponseInputItem = {"role": "user", "content": "once"} - if operation != "add": - await session.add_items([item]) - - mutation_applied = asyncio.Event() - allow_return = asyncio.Event() - - if operation in {"add", "pop"}: - original_save = fake_dapr_client.save_state - - async def controlled_save(*args: Any, **kwargs: Any) -> None: - await original_save(*args, **kwargs) - if kwargs.get("key") == session._messages_key: - mutation_applied.set() - await allow_return.wait() - - monkeypatch.setattr(fake_dapr_client, "save_state", controlled_save) - if operation == "add": - task: asyncio.Task[Any] = asyncio.create_task(session.add_items([item])) - else: - task = asyncio.create_task(session.pop_item()) - else: - original_delete = fake_dapr_client.delete_state - - async def controlled_delete(*args: Any, **kwargs: Any) -> None: - await original_delete(*args, **kwargs) - mutation_applied.set() - await allow_return.wait() - - monkeypatch.setattr(fake_dapr_client, "delete_state", controlled_delete) - task = asyncio.create_task(session.clear_session()) - - try: - await mutation_applied.wait() - task.cancel() - await asyncio.sleep(0) - task.cancel() - await asyncio.sleep(0) - assert task.done() is False - allow_return.set() - with pytest.raises(asyncio.CancelledError): - await task - finally: - allow_return.set() - if not task.done(): - task.cancel() - await asyncio.gather(task, return_exceptions=True) - - expected = [item] if operation == "add" else [] - assert await session.get_items() == expected - - async def test_runner_integration(agent: Agent, fake_dapr_client: FakeDaprClient): """Test that DaprSession works correctly with the agent Runner.""" session = await _create_test_session(fake_dapr_client) @@ -1418,22 +1327,3 @@ def entry_signalling_resolve(*args: Any, **kwargs: Any) -> Any: task.cancel() with suppress(asyncio.CancelledError, RuntimeError): await task - - -async def test_metadata_failure_does_not_commit_messages() -> None: - """An ancillary metadata failure must happen before conversation history is committed.""" - client = FailingMetadataFakeDaprClient() - session = DaprSession( - session_id="metadata_failure", - state_store_name="statestore", - dapr_client=client, # type: ignore[arg-type] - ) - item: TResponseInputItem = {"role": "user", "content": "once"} - - with pytest.raises(RuntimeError, match="metadata write failed"): - await session.add_items([item]) - - assert await session.get_items() == [] - - await session.add_items([item]) - assert await session.get_items() == [item] diff --git a/tests/extensions/memory/test_mongodb_session.py b/tests/extensions/memory/test_mongodb_session.py index acd9511343..da8b6214f4 100644 --- a/tests/extensions/memory/test_mongodb_session.py +++ b/tests/extensions/memory/test_mongodb_session.py @@ -95,6 +95,9 @@ def __init__(self) -> None: async def create_index(self, keys: Any, **kwargs: Any) -> str: return "fake_index" + def with_options(self, **kwargs: Any) -> FakeAsyncCollection: + return self + def find(self, query: dict[str, Any] | None = None) -> FakeCursor: query = query or {} results = [doc for doc in self._docs.values() if self._matches(doc, query)] @@ -200,10 +203,20 @@ async def delete_one(self, query: dict[str, Any]) -> None: @staticmethod def _matches(doc: dict[str, Any], query: dict[str, Any]) -> bool: for key, expected in query.items(): + if key == "$or": + if not any(FakeAsyncCollection._matches(doc, branch) for branch in expected): + return False + continue actual = doc.get(key) if isinstance(expected, dict): if "$ne" in expected and actual == expected["$ne"]: return False + if "$in" in expected and actual not in expected["$in"]: + return False + if "$lt" in expected and (actual is None or actual >= expected["$lt"]): + return False + if "$exists" in expected and (key in doc) != expected["$exists"]: + return False continue if actual != expected: return False @@ -240,6 +253,12 @@ def __init__(self, name: str, version: str | None = None) -> None: self.version = version +class FakeReadPreference: + """Minimal stand-in for pymongo read preference constants.""" + + PRIMARY = object() + + class FakeAsyncMongoClient: """In-memory substitute for pymongo AsyncMongoClient.""" @@ -275,16 +294,19 @@ def _make_fake_pymongo_modules() -> None: collection_mod = types.ModuleType("pymongo.asynchronous.collection") client_mod = types.ModuleType("pymongo.asynchronous.mongo_client") driver_info_mod = types.ModuleType("pymongo.driver_info") + read_preferences_mod = types.ModuleType("pymongo.read_preferences") collection_mod.AsyncCollection = FakeAsyncCollection # type: ignore[attr-defined] client_mod.AsyncMongoClient = FakeAsyncMongoClient # type: ignore[attr-defined] driver_info_mod.DriverInfo = FakeDriverInfo # type: ignore[attr-defined] + read_preferences_mod.ReadPreference = FakeReadPreference # type: ignore[attr-defined] sys.modules["pymongo"] = pymongo_mod sys.modules["pymongo.asynchronous"] = async_pkg sys.modules["pymongo.asynchronous.collection"] = collection_mod sys.modules["pymongo.asynchronous.mongo_client"] = client_mod sys.modules["pymongo.driver_info"] = driver_info_mod + sys.modules["pymongo.read_preferences"] = read_preferences_mod _make_fake_pymongo_modules() @@ -380,11 +402,142 @@ async def test_clear_session(session: MongoDBSession) -> None: await session.clear_session() assert await session.get_items() == [] - assert next(iter(session._sessions._docs.values()))["_seq"] == sequence_before_clear + metadata = next(iter(session._sessions._docs.values())) + assert metadata["_seq"] == sequence_before_clear + assert metadata["_generation"] == 1 await session.add_items([{"role": "user", "content": "after clear"}]) batch = next(iter(session._messages._docs.values())) assert batch["seq"] == sequence_before_clear + assert batch["generation"] == 1 + + +async def test_clear_session_partial_cleanup_failure_is_logically_atomic( + session: MongoDBSession, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """A failed physical sweep must not expose a partially cleared generation.""" + await session.add_items([{"role": "user", "content": "first"}]) + await session.add_items([{"role": "assistant", "content": "second"}]) + messages = cast(FakeAsyncCollection, session._messages) + + async def delete_one_then_fail(query: dict[str, Any]) -> None: + matching = [ + key for key, doc in messages._docs.items() if FakeAsyncCollection._matches(doc, query) + ] + assert len(matching) == 2 + del messages._docs[matching[0]] + raise RuntimeError("partial cleanup failed") + + monkeypatch.setattr(session._messages, "delete_many", delete_one_then_fail) + + await session.clear_session() + + assert len(messages._docs) == 1 + assert await session.get_items() == [] + + +async def test_generation_reads_use_primary_after_failed_clear_cleanup( + session: MongoDBSession, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Lagging secondary metadata must not expose cleared history to reads or pops.""" + await session.add_items([{"role": "user", "content": "old"}]) + sessions = cast(FakeAsyncCollection, session._sessions) + messages = cast(FakeAsyncCollection, session._messages) + primary_find = sessions.find + + async def fail_cleanup(query: dict[str, Any]) -> None: + raise RuntimeError("cleanup failed") + + monkeypatch.setattr(messages, "delete_many", fail_cleanup) + await session.clear_session() + assert len(messages._docs) == 1 + + metadata = copy.deepcopy(next(iter(sessions._docs.values()))) + stale_metadata = {**metadata, "_generation": 0} + + def stale_secondary_find(query: dict[str, Any] | None = None) -> FakeCursor: + query = query or {} + docs = [stale_metadata] if FakeAsyncCollection._matches(stale_metadata, query) else [] + return FakeCursor(docs) + + class PrimaryCollectionView: + def find(self, query: dict[str, Any] | None = None) -> FakeCursor: + return primary_find(query) + + def with_options(**kwargs: Any) -> PrimaryCollectionView: + assert kwargs == {"read_preference": FakeReadPreference.PRIMARY} + return PrimaryCollectionView() + + monkeypatch.setattr(sessions, "find", stale_secondary_find) + monkeypatch.setattr(sessions, "with_options", with_options) + + assert await session.get_items() == [] + assert await session.pop_item() is None + assert len(messages._docs) == 1 + + +async def test_clear_session_hides_add_reserved_before_generation_advance( + session: MongoDBSession, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """A batch reserved before clear must remain in the cleared generation.""" + insert_started = asyncio.Event() + allow_insert = asyncio.Event() + original_insert = session._messages.insert_one + + async def controlled_insert(document: dict[str, Any]) -> Any: + insert_started.set() + await allow_insert.wait() + return await original_insert(document) + + monkeypatch.setattr(session._messages, "insert_one", controlled_insert) + add_task = asyncio.create_task(session.add_items([{"role": "user", "content": "old"}])) + try: + await insert_started.wait() + await session.clear_session() + allow_insert.set() + await add_task + finally: + allow_insert.set() + if not add_task.done(): + add_task.cancel() + await asyncio.gather(add_task, return_exceptions=True) + + assert await session.get_items() == [] + + +async def test_pop_rechecks_generation_after_concurrent_clear( + session: MongoDBSession, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """A claim from a cleared generation must not be returned to the caller.""" + await session.add_items([{"role": "user", "content": "old"}]) + claim_finished = asyncio.Event() + allow_claim_return = asyncio.Event() + original_claim = session._messages.find_one_and_update + + async def controlled_claim(*args: Any, **kwargs: Any) -> Any: + result = await original_claim(*args, **kwargs) + claim_finished.set() + await allow_claim_return.wait() + return result + + monkeypatch.setattr(session._messages, "find_one_and_update", controlled_claim) + pop_task = asyncio.create_task(session.pop_item()) + try: + await claim_finished.wait() + await session.clear_session() + allow_claim_return.set() + assert await pop_task is None + finally: + allow_claim_return.set() + if not pop_task.done(): + pop_task.cancel() + await asyncio.gather(pop_task, return_exceptions=True) + + assert await session.get_items() == [] async def test_multiple_add_calls_accumulate(session: MongoDBSession) -> None: diff --git a/tests/extensions/memory/test_redis_session.py b/tests/extensions/memory/test_redis_session.py index 1a0a689853..661124ddcb 100644 --- a/tests/extensions/memory/test_redis_session.py +++ b/tests/extensions/memory/test_redis_session.py @@ -26,6 +26,14 @@ def _assert_cancel_message(exc: asyncio.CancelledError, expected: str) -> None: assert exc.args == expected_args +async def _release_after_detaching_pipeline_connection(delegate: Any) -> None: + """Model Redis 8.1 clearing its pipeline reference before awaiting release.""" + connection = delegate.connection + delegate.connection = None + if connection is not None: + await delegate.connection_pool.release(connection) + + # Try to use fakeredis for in-memory testing, fall back to real Redis if not available try: import fakeredis.aioredis @@ -1357,6 +1365,9 @@ class PipelineProxy: def __getattr__(self, name: str) -> Any: return getattr(delegate, name) + def __setattr__(self, name: str, value: Any) -> None: + setattr(delegate, name, value) + async def execute(self) -> Any: nonlocal execute_calls execute_calls += 1 @@ -1478,8 +1489,6 @@ async def test_add_items_refreshes_server_timestamps_after_watch_retry( if not USE_FAKE_REDIS: pytest.skip("This test requires fakeredis") - from redis.exceptions import WatchError - client = fakeredis.aioredis.FakeRedis() ttl = 60 first_seconds = int(time.time()) @@ -1506,6 +1515,9 @@ class PipelineProxy: def __getattr__(self, name: str) -> Any: return getattr(delegate, name) + def __setattr__(self, name: str, value: Any) -> None: + setattr(delegate, name, value) + def pexpireat(self, key: str, expiration_time_ms: int) -> Any: expirations.append(expiration_time_ms) return delegate.pexpireat(key, expiration_time_ms) @@ -1515,7 +1527,9 @@ async def time(self) -> tuple[int, int]: async def execute(self) -> Any: if attempt == 1: - raise WatchError("retry") + await client.hset( # type: ignore[misc] + session._session_key, "concurrent_write", "1" + ) return await delegate.execute() return PipelineProxy() @@ -1534,6 +1548,195 @@ async def execute(self) -> Any: assert await session.get_items() == [item] +async def test_ambiguous_exec_watch_error_is_not_retried( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """A transport-derived WatchError after a possible commit must remain ambiguous.""" + if not USE_FAKE_REDIS: + pytest.skip("This test requires fakeredis") + + from redis.exceptions import WatchError + + client = fakeredis.aioredis.FakeRedis() + session = RedisSession( + session_id="ambiguous_exec_watch_error", + redis_client=cast("Redis", client), + key_prefix="test:", + ) + real_pipeline = client.pipeline + attempts = 0 + item: TResponseInputItem = {"role": "user", "content": "once"} + serialized = json.dumps(item, separators=(",", ":")) + + def pipeline(*args: Any, **kwargs: Any) -> Any: + nonlocal attempts + attempts += 1 + delegate = real_pipeline(*args, **kwargs) + + class PipelineProxy: + def __getattr__(self, name: str) -> Any: + return getattr(delegate, name) + + async def execute(self) -> Any: + await client.rpush(session._messages_key, serialized) # type: ignore[misc] + raise WatchError("A ConnectionError occurred while watching one or more keys") + + return PipelineProxy() + + monkeypatch.setattr(client, "pipeline", pipeline) + + with pytest.raises(WatchError, match="ConnectionError"): + await session.add_items([item]) + + assert attempts == 1 + assert await session.get_items() == [item] + + +async def test_redis81_standard_pool_checkout_records_native_observability( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """The retained checkout path must preserve Redis 8.1 standard-pool metrics.""" + if not USE_FAKE_REDIS: + pytest.skip("This test requires fakeredis") + + import agents.extensions.memory.redis_session as redis_session_module + + client = fakeredis.aioredis.FakeRedis(max_connections=1) + pool = client.connection_pool + session = RedisSession( + session_id="redis81_standard_observability", + redis_client=cast("Redis", client), + key_prefix="test:", + ) + counts: list[tuple[str, int]] = [] + create_times: list[float] = [] + real_release = pool.release + + class ConnectionState: + IDLE = "idle" + USED = "used" + + async def record_connection_count( + *, pool_name: str, connection_state: str, counter: int + ) -> None: + assert pool_name == "test-pool" + counts.append((connection_state, counter)) + + async def record_connection_create_time( + *, connection_pool: Any, duration_seconds: float + ) -> None: + assert connection_pool is pool + create_times.append(duration_seconds) + + async def release(connection: Any) -> None: + await real_release(connection) + await record_connection_count( + pool_name="test-pool", connection_state=ConnectionState.USED, counter=-1 + ) + await record_connection_count( + pool_name="test-pool", connection_state=ConnectionState.IDLE, counter=1 + ) + + connection_module = cast(Any, redis_session_module)._redis_connection_api + has_native_release_metrics = hasattr(connection_module, "record_connection_count") + monkeypatch.setattr(connection_module, "ConnectionState", ConnectionState, raising=False) + monkeypatch.setattr( + connection_module, "get_pool_name", lambda _pool: "test-pool", raising=False + ) + monkeypatch.setattr( + connection_module, "record_connection_count", record_connection_count, raising=False + ) + monkeypatch.setattr( + connection_module, + "record_connection_create_time", + record_connection_create_time, + raising=False, + ) + if not has_native_release_metrics: + monkeypatch.setattr(pool, "release", release) + + await session.add_items([{"role": "user", "content": "once"}]) + + assert sum(counter for state, counter in counts if state == ConnectionState.USED) == 0 + assert sum(counter for state, counter in counts if state == ConnectionState.IDLE) == 1 + assert len(create_times) == 1 + + +@pytest.mark.parametrize("has_maintenance_lock", [False, True]) +async def test_redis8_blocking_pool_checkout_preserves_native_timing( + monkeypatch: pytest.MonkeyPatch, + has_maintenance_lock: bool, +) -> None: + """The retained checkout path must preserve Redis 8.0 and 8.1 timing hooks.""" + if not USE_FAKE_REDIS: + pytest.skip("This test requires fakeredis") + + from contextlib import asynccontextmanager + + from redis.asyncio import BlockingConnectionPool + + import agents.extensions.memory.redis_session as redis_session_module + + seed_client = fakeredis.aioredis.FakeRedis() + source_pool = seed_client.connection_pool + pool = BlockingConnectionPool( + max_connections=1, + timeout=1, + connection_class=source_pool.connection_class, + **source_pool.connection_kwargs, + ) + client = fakeredis.aioredis.FakeRedis(connection_pool=pool) + session = RedisSession( + session_id="redis81_blocking_observability", + redis_client=cast("Redis", client), + key_prefix="test:", + ) + maintenance_entries = 0 + create_times: list[float] = [] + wait_times: list[float] = [] + + @asynccontextmanager + async def maybe_pool_lock() -> Any: + nonlocal maintenance_entries + maintenance_entries += 1 + yield + + async def record_connection_create_time( + *, connection_pool: Any, duration_seconds: float + ) -> None: + assert connection_pool is pool + create_times.append(duration_seconds) + + async def record_connection_wait_time(*, pool_name: str, duration_seconds: float) -> None: + assert pool_name == "test-pool" + wait_times.append(duration_seconds) + + connection_module = cast(Any, redis_session_module)._redis_connection_api + if has_maintenance_lock: + monkeypatch.setattr(pool, "_maybe_pool_lock", maybe_pool_lock, raising=False) + monkeypatch.setattr( + connection_module, "get_pool_name", lambda _pool: "test-pool", raising=False + ) + monkeypatch.setattr( + connection_module, + "record_connection_create_time", + record_connection_create_time, + raising=False, + ) + monkeypatch.setattr( + connection_module, + "record_connection_wait_time", + record_connection_wait_time, + raising=False, + ) + + await session.add_items([{"role": "user", "content": "once"}]) + + assert maintenance_entries == int(has_maintenance_lock) + assert len(create_times) == 1 + assert len(wait_times) == 1 + + async def test_add_items_with_ttl_supports_single_connection_pool() -> None: """WATCH and server time must share one caller-managed Redis connection.""" if not USE_FAKE_REDIS: @@ -1556,7 +1759,7 @@ async def test_add_items_with_ttl_supports_single_connection_pool() -> None: async def test_cancelled_watch_releases_single_connection_pool( monkeypatch: pytest.MonkeyPatch, ) -> None: - """Repeated cancellation must not interrupt watched-pipeline cleanup.""" + """Cancellation before an immediate command must discard the watched connection.""" if not USE_FAKE_REDIS: pytest.skip("This test requires fakeredis") @@ -1568,9 +1771,9 @@ async def test_cancelled_watch_releases_single_connection_pool( ttl=60, ) real_pipeline = client.pipeline + pool = client.connection_pool time_started = asyncio.Event() - reset_started = asyncio.Event() - allow_reset = asyncio.Event() + dirty_connection: Any = None def pipeline(*args: Any, **kwargs: Any) -> Any: delegate = real_pipeline(*args, **kwargs) @@ -1580,15 +1783,13 @@ def __getattr__(self, name: str) -> Any: return getattr(delegate, name) async def time(self) -> tuple[int, int]: + nonlocal dirty_connection + dirty_connection = delegate.connection + assert dirty_connection is not None time_started.set() await asyncio.Event().wait() raise AssertionError("unreachable") - async def reset(self) -> None: - reset_started.set() - await allow_reset.wait() - await delegate.reset() - return PipelineProxy() monkeypatch.setattr(client, "pipeline", pipeline) @@ -1596,23 +1797,212 @@ async def reset(self) -> None: try: await time_started.wait() add_task.cancel("first-pre-exec-cancel") - await reset_started.wait() - add_task.cancel("second-during-reset") - await asyncio.sleep(0) - allow_reset.set() with pytest.raises(asyncio.CancelledError) as exc_info: await add_task finally: - allow_reset.set() if not add_task.done(): add_task.cancel() await asyncio.gather(add_task, return_exceptions=True) _assert_cancel_message(exc_info.value, "first-pre-exec-cancel") + assert dirty_connection not in pool._in_use_connections + assert dirty_connection not in pool._available_connections assert await session.get_items() == [] assert await client.ping() is True # type: ignore[misc] +async def test_cancelled_in_flight_watch_command_reconnects_before_pool_reuse( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Cancellation with an unread reply must discard the dirty connection.""" + if not USE_FAKE_REDIS: + pytest.skip("This test requires fakeredis") + + client = fakeredis.aioredis.FakeRedis(max_connections=1) + session = RedisSession( + session_id="cancelled_in_flight_watch_command", + redis_client=cast("Redis", client), + key_prefix="test:", + ) + real_pipeline = client.pipeline + pool = client.connection_pool + real_release = pool.release + response_read_started = asyncio.Event() + dirty_connection: Any = None + close_calls = 0 + released_connections: list[Any] = [] + + async def track_release(connection: Any) -> None: + released_connections.append(connection) + await real_release(connection) + + def pipeline(*args: Any, **kwargs: Any) -> Any: + delegate = real_pipeline(*args, **kwargs) + real_type = delegate.type + + async def controlled_type(key: str) -> Any: + nonlocal close_calls, dirty_connection + connection = delegate.connection + assert connection is not None + dirty_connection = connection + real_read_response = connection.read_response + real_close = connection._close + read_calls = 0 + + async def block_first_read(*args: Any, **kwargs: Any) -> Any: + nonlocal read_calls + read_calls += 1 + if read_calls == 1: + response_read_started.set() + await asyncio.Event().wait() + raise AssertionError("unreachable") + return await real_read_response(*args, **kwargs) + + def track_close() -> None: + nonlocal close_calls + close_calls += 1 + real_close() + + monkeypatch.setattr(connection, "read_response", block_first_read) + monkeypatch.setattr(connection, "_close", track_close) + return await real_type(key) + + monkeypatch.setattr(delegate, "type", controlled_type) + return delegate + + monkeypatch.setattr(client, "pipeline", pipeline) + monkeypatch.setattr(pool, "release", track_release) + add_task = asyncio.create_task(session.add_items([{"role": "user", "content": "cancelled"}])) + try: + await response_read_started.wait() + add_task.cancel("cancel-during-response-read") + with pytest.raises(asyncio.CancelledError) as exc_info: + await add_task + finally: + if not add_task.done(): + add_task.cancel() + await asyncio.gather(add_task, return_exceptions=True) + + _assert_cancel_message(exc_info.value, "cancel-during-response-read") + assert dirty_connection is not None + assert close_calls == 1 + assert released_connections == [] + assert dirty_connection not in pool._in_use_connections + assert dirty_connection not in pool._available_connections + assert await client.ping() is True # type: ignore[misc] + assert len(released_connections) == 1 + assert released_connections[0] is not dirty_connection + assert await session.get_items() == [] + + +@pytest.mark.parametrize("blocking_pool", [False, True]) +@pytest.mark.parametrize("validation_fails", [False, True]) +async def test_cancelled_connection_acquisition_is_discarded_before_release( + monkeypatch: pytest.MonkeyPatch, + blocking_pool: bool, + validation_fails: bool, +) -> None: + """Cancellation during pool validation must retain and discard the acquired identity.""" + if not USE_FAKE_REDIS: + pytest.skip("This test requires fakeredis") + + from redis.asyncio import BlockingConnectionPool + from redis.event import AsyncAfterConnectionReleasedEvent + + pool: Any + if blocking_pool: + seed_client = fakeredis.aioredis.FakeRedis() + source_pool = seed_client.connection_pool + pool = BlockingConnectionPool( + max_connections=1, + timeout=1, + connection_class=source_pool.connection_class, + **source_pool.connection_kwargs, + ) + client = fakeredis.aioredis.FakeRedis(connection_pool=pool) + else: + client = fakeredis.aioredis.FakeRedis(max_connections=1) + pool = client.connection_pool + + session = RedisSession( + session_id=f"cancelled_connection_acquisition_{blocking_pool}", + redis_client=cast("Redis", client), + key_prefix="test:", + ) + dispatcher = pool._event_dispatcher + assert dispatcher is not None + listeners = dispatcher._event_listeners_mapping[AsyncAfterConnectionReleasedEvent] + released_connections: list[Any] = [] + ensure_started = asyncio.Event() + allow_ensure = asyncio.Event() + dirty_connection: Any = None + real_ensure_connection = pool.ensure_connection + loop = asyncio.get_running_loop() + previous_exception_handler = loop.get_exception_handler() + loop_errors: list[dict[str, Any]] = [] + + class RecordingReleaseListener: + async def listen(self, event: Any) -> None: + released_connections.append(event.connection) + + async def controlled_ensure_connection(connection: Any) -> None: + nonlocal dirty_connection + if dirty_connection is None: + dirty_connection = connection + ensure_started.set() + await allow_ensure.wait() + if validation_fails: + raise RuntimeError("connection validation failed") + await real_ensure_connection(connection) + + monkeypatch.setitem( + dispatcher._event_listeners_mapping, + AsyncAfterConnectionReleasedEvent, + [*listeners, RecordingReleaseListener()], + ) + monkeypatch.setattr(pool, "ensure_connection", controlled_ensure_connection) + loop.set_exception_handler(lambda _loop, context: loop_errors.append(context)) + add_task = asyncio.create_task(session.add_items([{"role": "user", "content": "cancelled"}])) + waiter: asyncio.Task[Any] | None = None + try: + await ensure_started.wait() + add_task.cancel("cancel-during-acquisition") + if blocking_pool: + + async def ping() -> Any: + return await client.ping() # type: ignore[misc] + + waiter = asyncio.create_task(ping()) + await asyncio.sleep(0) + allow_ensure.set() + + with pytest.raises(asyncio.CancelledError) as exc_info: + await add_task + _assert_cancel_message(exc_info.value, "cancel-during-acquisition") + + assert dirty_connection is not None + assert dirty_connection not in pool._in_use_connections + assert dirty_connection not in pool._available_connections + assert dirty_connection not in released_connections + if waiter is not None: + assert await waiter is True + else: + assert await client.ping() is True # type: ignore[misc] + assert released_connections + assert all(connection is not dirty_connection for connection in released_connections) + assert await session.get_items() == [] + await asyncio.sleep(0) + assert loop_errors == [] + finally: + loop.set_exception_handler(previous_exception_handler) + allow_ensure.set() + pending = [task for task in (add_task, waiter) if task is not None and not task.done()] + for task in pending: + task.cancel() + if pending: + await asyncio.gather(*pending, return_exceptions=True) + + @pytest.mark.parametrize( "cleanup_mode", ["internal_reset", "failing_internal_reset"], @@ -1783,20 +2173,35 @@ async def controlled_reset() -> None: assert await client.ping() is True # type: ignore[misc] -async def test_post_commit_repeated_reset_failure_force_releases_connection( +async def test_post_commit_reset_failure_detaches_without_release_listener( monkeypatch: pytest.MonkeyPatch, ) -> None: - """Repeated reset failures after EXEC must not leak the pipeline connection.""" + """A reset failure must detach rather than release through a reborrowing listener.""" if not USE_FAKE_REDIS: pytest.skip("This test requires fakeredis") + from redis.event import AsyncAfterConnectionReleasedEvent + client = fakeredis.aioredis.FakeRedis(max_connections=1) + pool = client.connection_pool session = RedisSession( session_id="post_commit_repeated_reset_failure", redis_client=cast("Redis", client), key_prefix="test:", ) real_pipeline = client.pipeline + dispatcher = pool._event_dispatcher + assert dispatcher is not None + listeners = dispatcher._event_listeners_mapping[AsyncAfterConnectionReleasedEvent] + listener_calls = 0 + borrowed_connection: Any = None + + class ReborrowingFailingReleaseListener: + async def listen(self, event: Any) -> None: + nonlocal borrowed_connection, listener_calls + listener_calls += 1 + borrowed_connection = await pool.get_connection() + raise RuntimeError("release listener failed after reborrow") def pipeline(*args: Any, **kwargs: Any) -> Any: delegate = real_pipeline(*args, **kwargs) @@ -1814,10 +2219,138 @@ async def fail_two_resets_before_release() -> None: return delegate monkeypatch.setattr(client, "pipeline", pipeline) + monkeypatch.setitem( + dispatcher._event_listeners_mapping, + AsyncAfterConnectionReleasedEvent, + [*listeners, ReborrowingFailingReleaseListener()], + ) item: TResponseInputItem = {"role": "user", "content": "committed"} await session.add_items([item]) + assert listener_calls == 0 + assert borrowed_connection is None + assert not pool._in_use_connections + monkeypatch.setitem( + dispatcher._event_listeners_mapping, + AsyncAfterConnectionReleasedEvent, + listeners, + ) + assert await session.get_items() == [item] + assert await client.ping() is True # type: ignore[misc] + + +async def test_reconnect_required_release_detaches_without_listener_reborrow( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """A reconnect-required identity must never enter shared-pool release.""" + if not USE_FAKE_REDIS: + pytest.skip("This test requires fakeredis") + + from redis.event import AsyncAfterConnectionReleasedEvent + + client = fakeredis.aioredis.FakeRedis(max_connections=1) + pool = client.connection_pool + session = RedisSession( + session_id="reconnect_required_release", + redis_client=cast("Redis", client), + key_prefix="test:", + ) + real_pipeline = client.pipeline + dispatcher = pool._event_dispatcher + assert dispatcher is not None + listeners = dispatcher._event_listeners_mapping[AsyncAfterConnectionReleasedEvent] + listener_calls = 0 + borrowed_connection: Any = None + + class ReborrowingFailingReleaseListener: + async def listen(self, event: Any) -> None: + nonlocal borrowed_connection, listener_calls + listener_calls += 1 + borrowed_connection = await pool.get_connection() + raise RuntimeError("release listener failed after reborrow") + + def pipeline(*args: Any, **kwargs: Any) -> Any: + delegate = real_pipeline(*args, **kwargs) + real_reset = delegate.reset + + async def mark_for_reconnect_before_reset() -> None: + connection = delegate.connection + assert connection is not None + connection.mark_for_reconnect() + await real_reset() + + monkeypatch.setattr(delegate, "reset", mark_for_reconnect_before_reset) + return delegate + + monkeypatch.setattr(client, "pipeline", pipeline) + monkeypatch.setitem( + dispatcher._event_listeners_mapping, + AsyncAfterConnectionReleasedEvent, + [*listeners, ReborrowingFailingReleaseListener()], + ) + item: TResponseInputItem = {"role": "user", "content": "committed"} + + await session.add_items([item]) + + assert listener_calls == 0 + assert borrowed_connection is None + assert not pool._in_use_connections + monkeypatch.setitem( + dispatcher._event_listeners_mapping, + AsyncAfterConnectionReleasedEvent, + listeners, + ) + monkeypatch.setattr(client, "pipeline", real_pipeline) + assert await session.get_items() == [item] + assert await client.ping() is True # type: ignore[misc] + + +async def test_reconnect_marked_inside_release_failure_is_detached( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """A reconnect race before native transfer must retain cleanup ownership.""" + if not USE_FAKE_REDIS: + pytest.skip("This test requires fakeredis") + + client = fakeredis.aioredis.FakeRedis(max_connections=1) + pool = client.connection_pool + session = RedisSession( + session_id="reconnect_during_release", + redis_client=cast("Redis", client), + key_prefix="test:", + ) + real_release = pool.release + real_pipeline = client.pipeline + close_calls = 0 + detached_connection: Any = None + + async def redis81_raced_release(connection: Any) -> None: + nonlocal close_calls, detached_connection + detached_connection = connection + real_close = connection._close + + def tracked_close() -> None: + nonlocal close_calls + close_calls += 1 + real_close() + + monkeypatch.setattr(connection, "_close", tracked_close) + connection.mark_for_reconnect() + pool._in_use_connections.remove(connection) + raise RuntimeError("disconnect failed before pool transfer") + + monkeypatch.setattr(pool, "release", redis81_raced_release) + item: TResponseInputItem = {"role": "user", "content": "committed"} + + await session.add_items([item]) + + assert detached_connection is not None + assert close_calls == 1 + assert detached_connection not in pool._in_use_connections + assert detached_connection not in pool._available_connections + monkeypatch.setattr(pool, "release", real_release) + monkeypatch.setattr(client, "pipeline", real_pipeline) assert await session.get_items() == [item] assert await client.ping() is True # type: ignore[misc] @@ -1863,10 +2396,10 @@ async def listen(self, event: Any) -> None: assert await client.ping() is True # type: ignore[misc] -async def test_pre_exec_disconnect_failure_does_not_reuse_watched_connection( +async def test_pre_exec_close_failure_is_quarantined( monkeypatch: pytest.MonkeyPatch, ) -> None: - """A failed disconnect must remove the watched connection from pool reuse.""" + """A failed dirty-connection close must remain owned for close retry.""" if not USE_FAKE_REDIS: pytest.skip("This test requires fakeredis") @@ -1879,27 +2412,27 @@ async def test_pre_exec_disconnect_failure_does_not_reuse_watched_connection( ) real_pipeline = client.pipeline time_started = asyncio.Event() + dirty_connection: Any = None + real_close: Any = None def pipeline(*args: Any, **kwargs: Any) -> Any: delegate = real_pipeline(*args, **kwargs) async def controlled_time() -> tuple[int, int]: + nonlocal dirty_connection, real_close + dirty_connection = delegate.connection + assert dirty_connection is not None + real_close = dirty_connection._close + + def fail_close() -> None: + raise RuntimeError("close failed") + + monkeypatch.setattr(dirty_connection, "_close", fail_close) time_started.set() await asyncio.Event().wait() raise AssertionError("unreachable") - async def fail_reset_before_release() -> None: - connection = delegate.connection - assert connection is not None - - async def fail_disconnect(nowait: bool = False) -> None: - raise RuntimeError("disconnect failed") - - monkeypatch.setattr(connection, "disconnect", fail_disconnect) - raise RuntimeError("reset failed before release") - monkeypatch.setattr(delegate, "time", controlled_time) - monkeypatch.setattr(delegate, "reset", fail_reset_before_release) return delegate monkeypatch.setattr(client, "pipeline", pipeline) @@ -1911,15 +2444,21 @@ async def fail_disconnect(nowait: bool = False) -> None: await add_task _assert_cancel_message(exc_info.value, "caller-cancel") + assert session._detached_connections == {dirty_connection} + assert dirty_connection not in client.connection_pool._in_use_connections + assert dirty_connection not in client.connection_pool._available_connections + monkeypatch.setattr(dirty_connection, "_close", real_close) + await session.close() + assert session._detached_connections == set() monkeypatch.setattr(client, "pipeline", real_pipeline) assert await client.ping() is True # type: ignore[misc] assert await session.get_items() == [] -async def test_pre_exec_cancellation_wins_after_released_reset_failure( +async def test_pre_exec_cancellation_skips_pipeline_reset( monkeypatch: pytest.MonkeyPatch, ) -> None: - """Caller cancellation must win once a failing reset has released its connection.""" + """A protocol-dirty cancellation must not expose the connection through reset.""" if not USE_FAKE_REDIS: pytest.skip("This test requires fakeredis") @@ -1935,19 +2474,17 @@ async def test_pre_exec_cancellation_wins_after_released_reset_failure( def pipeline(*args: Any, **kwargs: Any) -> Any: delegate = real_pipeline(*args, **kwargs) - real_reset = delegate.reset async def controlled_time() -> tuple[int, int]: time_started.set() await asyncio.Event().wait() raise AssertionError("unreachable") - async def fail_reset_after_release() -> None: - await real_reset() - raise RuntimeError("post-release reset callback failed") + async def fail_if_reset() -> None: + raise AssertionError("dirty connection reached pipeline reset") monkeypatch.setattr(delegate, "time", controlled_time) - monkeypatch.setattr(delegate, "reset", fail_reset_after_release) + monkeypatch.setattr(delegate, "reset", fail_if_reset) return delegate monkeypatch.setattr(client, "pipeline", pipeline) @@ -2127,10 +2664,10 @@ async def transparent_release(connection: Any) -> None: assert await session.get_items() == [item] -async def test_ambiguous_overridden_pool_release_is_detached_after_commit( +async def test_release_failure_after_reborrow_preserves_new_borrower( monkeypatch: pytest.MonkeyPatch, ) -> None: - """A release-then-reborrow failure must not expose a committed batch as retryable.""" + """A release-then-reborrow failure must not reclaim the new borrower's connection.""" if not USE_FAKE_REDIS: pytest.skip("This test requires fakeredis") @@ -2142,9 +2679,19 @@ async def test_ambiguous_overridden_pool_release_is_detached_after_commit( key_prefix="test:", ) real_release = pool.release + real_pipeline = client.pipeline release_calls = 0 borrowed_connection: Any = None + def pipeline(*args: Any, **kwargs: Any) -> Any: + delegate = real_pipeline(*args, **kwargs) + monkeypatch.setattr( + delegate, + "reset", + lambda: _release_after_detaching_pipeline_connection(delegate), + ) + return delegate + async def transfer_reborrow_then_fail(connection: Any) -> None: nonlocal borrowed_connection, release_calls release_calls += 1 @@ -2155,17 +2702,20 @@ async def transfer_reborrow_then_fail(connection: Any) -> None: raise RuntimeError("release failed after reborrow") monkeypatch.setattr(pool, "release", transfer_reborrow_then_fail) + monkeypatch.setattr(client, "pipeline", pipeline) item: TResponseInputItem = {"role": "user", "content": "committed"} await session.add_items([item]) assert release_calls == 1 assert borrowed_connection is not None - assert not pool._in_use_connections + assert borrowed_connection in pool._in_use_connections assert borrowed_connection not in pool._available_connections monkeypatch.setattr(pool, "release", real_release) + monkeypatch.setattr(client, "pipeline", real_pipeline) + await real_release(borrowed_connection) replacement_connection = await pool.get_connection() - assert replacement_connection is not borrowed_connection + assert replacement_connection is borrowed_connection await real_release(replacement_connection) assert await session.get_items() == [item] @@ -2187,7 +2737,6 @@ async def test_post_commit_disconnect_failure_is_not_retryable( def pipeline(*args: Any, **kwargs: Any) -> Any: delegate = real_pipeline(*args, **kwargs) - real_reset = delegate.reset reset_calls = 0 async def fail_disconnect_after_commit() -> None: @@ -2202,7 +2751,7 @@ async def fail_disconnect(nowait: bool = False) -> None: raise RuntimeError("disconnect failed") monkeypatch.setattr(connection, "disconnect", fail_disconnect) - await real_reset() + await _release_after_detaching_pipeline_connection(delegate) monkeypatch.setattr(delegate, "reset", fail_disconnect_after_commit) return delegate @@ -2241,7 +2790,6 @@ async def test_post_commit_detached_close_failure_is_quarantined( def pipeline(*args: Any, **kwargs: Any) -> Any: delegate = real_pipeline(*args, **kwargs) - real_reset = delegate.reset async def fail_disconnect_and_close_after_commit() -> None: nonlocal detached_connection, real_close @@ -2261,7 +2809,7 @@ def fail_close() -> None: monkeypatch.setattr(detached_connection, "disconnect", fail_disconnect) monkeypatch.setattr(detached_connection, "_close", fail_close) - await real_reset() + await _release_after_detaching_pipeline_connection(delegate) monkeypatch.setattr(delegate, "reset", fail_disconnect_and_close_after_commit) return delegate @@ -2412,7 +2960,6 @@ async def test_blocking_pool_detached_disconnect_notifies_waiter( def pipeline(*args: Any, **kwargs: Any) -> Any: delegate = real_pipeline(*args, **kwargs) - real_reset = delegate.reset async def fail_disconnect_after_commit() -> None: reset_started.set() @@ -2425,7 +2972,7 @@ async def fail_disconnect(nowait: bool = False) -> None: raise RuntimeError("disconnect failed") monkeypatch.setattr(connection, "disconnect", fail_disconnect) - await real_reset() + await _release_after_detaching_pipeline_connection(delegate) monkeypatch.setattr(delegate, "reset", fail_disconnect_after_commit) return delegate diff --git a/tests/memory/test_session.py b/tests/memory/test_session.py index efc9fe548f..8c8bf5ea0d 100644 --- a/tests/memory/test_session.py +++ b/tests/memory/test_session.py @@ -10,10 +10,44 @@ import pytest from agents import Agent, RunConfig, Runner, SessionSettings, SQLiteSession, TResponseInputItem +from agents.memory.sqlite_session import _await_mutation from tests.fake_model import FakeModel from tests.test_responses import get_text_message +@pytest.mark.asyncio +async def test_await_mutation_cancellation_hides_later_failure_without_loop_error() -> None: + """A failed mutation must not leak a false loop error after caller cancellation.""" + mutation_started = asyncio.Event() + allow_failure = asyncio.Event() + loop = asyncio.get_running_loop() + previous_exception_handler = loop.get_exception_handler() + loop_errors: list[dict[str, Any]] = [] + + async def mutation() -> None: + mutation_started.set() + await allow_failure.wait() + raise RuntimeError("mutation failed") + + loop.set_exception_handler(lambda _loop, context: loop_errors.append(context)) + task = asyncio.create_task(_await_mutation(mutation())) + try: + await mutation_started.wait() + task.cancel("caller-cancelled") + allow_failure.set() + + with pytest.raises(asyncio.CancelledError): + await task + await asyncio.sleep(0) + assert loop_errors == [] + finally: + loop.set_exception_handler(previous_exception_handler) + allow_failure.set() + if not task.done(): + task.cancel() + await asyncio.gather(task, return_exceptions=True) + + # Helper functions for parametrized testing of different Runner methods def _run_sync_wrapper(agent, input_data, **kwargs): """Wrapper for run_sync that properly sets up an event loop.""" From a9d4bc05f59f6374e51eb570e73f3567b0d05946 Mon Sep 17 00:00:00 2001 From: Kazuhiro Sera Date: Thu, 6 Aug 2026 12:18:10 +0900 Subject: [PATCH 7/7] test: avoid leaking Redis pipeline proxy --- tests/extensions/memory/test_redis_session.py | 21 ++++++++----------- 1 file changed, 9 insertions(+), 12 deletions(-) diff --git a/tests/extensions/memory/test_redis_session.py b/tests/extensions/memory/test_redis_session.py index 661124ddcb..8a6f35eee2 100644 --- a/tests/extensions/memory/test_redis_session.py +++ b/tests/extensions/memory/test_redis_session.py @@ -1778,19 +1778,16 @@ async def test_cancelled_watch_releases_single_connection_pool( def pipeline(*args: Any, **kwargs: Any) -> Any: delegate = real_pipeline(*args, **kwargs) - class PipelineProxy: - def __getattr__(self, name: str) -> Any: - return getattr(delegate, name) - - async def time(self) -> tuple[int, int]: - nonlocal dirty_connection - dirty_connection = delegate.connection - assert dirty_connection is not None - time_started.set() - await asyncio.Event().wait() - raise AssertionError("unreachable") + async def controlled_time() -> tuple[int, int]: + nonlocal dirty_connection + dirty_connection = delegate.connection + assert dirty_connection is not None + time_started.set() + await asyncio.Event().wait() + raise AssertionError("unreachable") - return PipelineProxy() + monkeypatch.setattr(delegate, "time", controlled_time) + return delegate monkeypatch.setattr(client, "pipeline", pipeline) add_task = asyncio.create_task(session.add_items([{"role": "user", "content": "cancelled"}]))