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..c67e3f8a6a 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,13 +68,18 @@ 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 - # 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__) @@ -85,19 +91,31 @@ 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._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 ( @@ -139,6 +157,7 @@ def _init_structure_tables(self): """) self._ensure_branch_reservations_table(conn) + self._ensure_session_clear_generations_table(conn) # Indexes conn.execute(""" @@ -178,20 +197,18 @@ 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: + 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) + 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, @@ -209,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: @@ -225,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: @@ -236,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()) @@ -255,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))) @@ -276,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()))) @@ -298,79 +313,72 @@ 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._locked_connection() as conn: + 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: - # 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), + (self.session_id, resolved_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, resolved_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), + (self.session_id, resolved_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 +390,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 +406,43 @@ 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) + self._ensure_session_clear_generations_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.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 # locked operation keeps the reset atomic with the clear, so no @@ -436,10 +450,10 @@ 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 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. @@ -490,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) @@ -564,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( """ @@ -571,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 @@ -591,12 +606,12 @@ 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() 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, @@ -709,15 +724,12 @@ 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) + 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: @@ -839,39 +851,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 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 @@ -908,14 +924,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( """ @@ -928,13 +941,27 @@ 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 # 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) @@ -971,57 +998,53 @@ 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 + usage_deleted, structure_deleted, orphaned_messages_deleted = await _await_mutation( + asyncio.to_thread(_delete_sync) ) self._logger.info( @@ -1047,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( """ @@ -1071,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, } ) @@ -1113,6 +1137,64 @@ 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, + *, + initialize: bool = True, + ) -> None: + """Reset a stale branch pointer after another session instance clears history.""" + 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() + 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: @@ -1148,127 +1230,124 @@ 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._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) + 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""" + 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, source_branch_id, generation return await asyncio.to_thread(_copy_sync) @@ -1286,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""" @@ -1305,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 = [] @@ -1341,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""" @@ -1361,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 = [] @@ -1396,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( """ @@ -1410,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]]] = {} @@ -1432,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( """ @@ -1472,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() @@ -1554,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, @@ -1569,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: @@ -1608,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 @@ -1671,7 +1747,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( @@ -1733,4 +1809,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 06d0cc1755..215a668902 100644 --- a/src/agents/extensions/memory/async_sqlite_session.py +++ b/src/agents/extensions/memory/async_sqlite_session.py @@ -16,6 +16,7 @@ coerce_session_settings, resolve_session_limit, ) +from ...memory.sqlite_session import _await_mutation class AsyncSQLiteSession(SessionABC): @@ -57,6 +58,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 +104,46 @@ 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 as acquisition_error: + connection = None + cleanup_cancellation: asyncio.CancelledError | None = None + try: + 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 + 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 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 return self._connection @@ -121,6 +160,71 @@ 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 as operation_error: + rollback_task = asyncio.create_task(conn.rollback()) + rollback_error: BaseException | None = None + rollback_cancellation: asyncio.CancelledError | None = None + try: + 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) -> 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_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 + 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. @@ -206,7 +310,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 +335,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 +343,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 +361,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 +383,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 +399,42 @@ 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 + 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 + 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/mongodb_session.py b/src/agents/extensions/memory/mongodb_session.py index b2ba601ab0..3887409ca1 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 @@ -50,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", @@ -65,6 +67,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) @@ -73,19 +76,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 +131,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). @@ -245,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() @@ -263,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 # ------------------------------------------------------------------ @@ -287,16 +309,23 @@ 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] = [] 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: @@ -319,79 +348,211 @@ 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}, + "$setOnInsert": { + "session_id": self.session_id, + "created_at": now, + "_generation": 0, + }, "$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 + ) + generation = result.get("_generation", 0) if result else 0 + if not isinstance(generation, int): + generation = 0 - 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, + "generation": generation, + "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 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: 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() + 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, + **self._generation_query(generation), + "message_data": [], + } + ) + except asyncio.CancelledError: + raise + except Exception: + pass while True: - doc = await self._messages.find_one_and_delete( - {"session_id": self.session_id}, + doc = await self._messages.find_one_and_update( + { + "session_id": self.session_id, + **self._generation_query(generation), + "message_data": {"$ne": []}, + }, + [ + { + "$set": { + "message_data": { + "$cond": [ + {"$isArray": "$message_data"}, + { + "$slice": [ + "$message_data", + { + "$subtract": [ + {"$size": "$message_data"}, + 1, + ] + }, + ] + }, + [], + ] + } + } + } + ], sort=[("seq", -1)], + 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): + if not raw: + continue + claimed_raw = raw[-1] + exhausted = len(raw) == 1 + else: + 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(doc["message_data"]) - except (json.JSONDecodeError, KeyError, TypeError): + 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 self._messages.delete_many({"session_id": self.session_id}) - await self._sessions.delete_one({"session_id": self.session_id}) + await _await_mutation(self._clear_session()) + + async def _clear_session(self) -> None: + """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 de9efbf6c8..953b1cf683 100644 --- a/src/agents/extensions/memory/redis_session.py +++ b/src/agents/extensions/memory/redis_session.py @@ -24,13 +24,17 @@ 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 + import redis.asyncio.connection as redis_connection + from redis.asyncio import BlockingConnectionPool, Redis + from redis.event import AsyncAfterConnectionReleasedEvent + from redis.exceptions import ConnectionError as RedisConnectionError, ResponseError, WatchError except ImportError as e: raise_optional_dependency_error( "RedisSession", @@ -46,6 +50,308 @@ coerce_session_settings, resolve_session_limit, ) +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 + + +class _PipelineConnectionPool: + """Track one pipeline's connection release without changing the shared pool.""" + + 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 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: + pipe.connection = None + return None, True, None + + reset_error: BaseException | None = None + if not connection_pool.release_started and not discard_connection: + try: + await pipe.reset() + except BaseException as exc: + reset_error = exc + + if connection_pool.release_completed: + pipe.connection = None + return reset_error, True, None + + connection = connection_pool.connection or getattr(pipe, "connection", None) + if connection is None: + 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: + await connection_pool.record_discard() + except BaseException as exc: + metrics_error = exc + 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 + 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 + + +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 not attempt.done(): + try: + await asyncio.wait({attempt}) + except asyncio.CancelledError as exc: + if cancellation is None: + cancellation = exc + if not completion_owned.is_set() and not attempt_cancelled: + attempt.cancel() + attempt_cancelled = True + + try: + outcome = attempt.result() + except BaseException: + if cancellation is not None: + raise cancellation from None + raise + return outcome, cancellation class RedisSession(SessionABC): @@ -70,7 +376,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(). """ @@ -87,6 +394,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}" @@ -143,14 +451,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 +460,162 @@ 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_attempt( + self, + pipe: Any, + keys: tuple[str, str, str], + serialized_items: list[str], + completion_owned: asyncio.Event, + ) -> _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 + + 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 + + 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) + 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) + 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) + + 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 + except WatchError as exc: + 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, + discard_connection=discard_connection, + ) + if detached_connection is not None: + self._detached_connections.add(detached_connection) + pipe.connection_pool = raw_connection_pool + + return _PipelineAttemptOutcome( + committed=committed, + retryable_watch_conflict=retryable_watch_conflict, + 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 + if cancellation is not None: + raise cancellation + if outcome.cleanup_error is not None: + raise outcome.cleanup_error + if outcome.retryable_watch_conflict: + 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. @@ -224,32 +680,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. @@ -259,34 +695,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. @@ -303,12 +746,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/src/agents/extensions/memory/sqlalchemy_session.py b/src/agents/extensions/memory/sqlalchemy_session.py index 8751cca68b..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() - async with self._session_factory() as sess: - async with sess.begin(): - while True: - # Fallback for all dialects - get ID first, then delete - subq = ( - select(self._messages.c.id) + while True: + retry_claim = False + async with self._session_factory() as sess: + async with sess.begin(): + 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,27 +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() - 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 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 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/src/agents/memory/sqlite_session.py b/src/agents/memory/sqlite_session.py index 61bf4e563b..fc9f3fdb8f 100644 --- a/src/agents/memory/sqlite_session.py +++ b/src/agents/memory/sqlite_session.py @@ -4,15 +4,39 @@ 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 not task.done(): + try: + await asyncio.wait({task}) + except asyncio.CancelledError as exc: + if cancellation is None: + cancellation = exc + + try: + result = task.result() + except BaseException: + if cancellation is not None: + raise cancellation from None + raise + if cancellation is not None: + raise cancellation from None + return result + class SQLiteSession(SessionABC): """SQLite-based implementation of session storage. @@ -57,6 +81,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 +150,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 +352,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 +366,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 +410,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 +427,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..389d174e68 100644 --- a/tests/extensions/memory/test_advanced_sqlite_session.py +++ b/tests/extensions/memory/test_advanced_sqlite_session.py @@ -5,6 +5,8 @@ import json import logging import multiprocessing +import sqlite3 +import sys import tempfile import threading import time @@ -31,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.""" @@ -373,6 +381,297 @@ 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_propagates_after_known_mutation_outcome( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + operation: str, +): + """Cancellation after commit must propagate without inviting a 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() + with pytest.raises(asyncio.CancelledError): + await mutation + finally: + conn.allow_return.set() + if not mutation.done(): + mutation.cancel() + await asyncio.gather(mutation, return_exceptions=True) + + if operation == "add": + assert await session.get_items() == [item] + elif operation == "pop": + assert await session.get_items() == [] + else: + assert await session.get_items() == [] + assert mutation.cancelled() + 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_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 == "cancelled_branch" + 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" @@ -751,6 +1050,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 +1998,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 +3079,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. @@ -2957,6 +3365,305 @@ 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() + + +@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, +): + """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 b45cbdf4e7..5ade0e2cc5 100644 --- a/tests/extensions/memory/test_async_sqlite_session.py +++ b/tests/extensions/memory/test_async_sqlite_session.py @@ -2,7 +2,10 @@ from __future__ import annotations +import asyncio import json +import sqlite3 +import sys import tempfile from collections.abc import Sequence from datetime import datetime @@ -22,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.""" @@ -495,3 +504,622 @@ 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: 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() + + 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 close_task is not None and not close_task.done(): + close_task.cancel() + await asyncio.gather(close_task, return_exceptions=True) + 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"]) +async def test_post_commit_cancellation_propagates_after_known_mutation_outcome( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + operation: str, +): + """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"} + try: + 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() + with pytest.raises(asyncio.CancelledError): + await mutation + finally: + allow_return.set() + if not mutation.done(): + mutation.cancel() + await asyncio.gather(mutation, return_exceptions=True) + + if operation == "add": + assert await session.get_items() == [item] + elif operation == "pop": + assert await session.get_items() == [] + else: + assert await session.get_items() == [] + assert mutation.cancelled() + finally: + 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()}) + try: + 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"] + finally: + 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) + try: + 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) + finally: + 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) + try: + 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) + finally: + 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") + try: + with pytest.raises(RuntimeError, match="initialization failed"): + await session.get_items() + + 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( + 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) + try: + 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() == [] + finally: + 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_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() == [] + 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, +): + """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) + try: + 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"] + finally: + 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) + 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_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) + 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( + 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") + 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 + + monkeypatch.setattr(conn, "close", session.real_close) + await session.close() + + 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( + 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 + 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) + init_started.set() + await asyncio.Event().wait() + + session = CancelledInitSession("cancelled_init") + task = asyncio.create_task(session.get_items()) + try: + 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 + 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_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_cancel_message(exc_info.value, "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, +): + """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: + 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 + 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 3bd8f7c034..da8b6214f4 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 @@ -43,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})" @@ -87,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)] @@ -117,22 +128,44 @@ 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], + update: dict[str, Any] | list[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) + 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(): + 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: + assert isinstance(update, dict) new_doc: dict[str, Any] = {"_id": FakeObjectId()} new_doc.update(update.get("$setOnInsert", {})) new_doc.update(update.get("$set", {})) @@ -169,7 +202,25 @@ 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(): + 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 + return True class FakeAsyncDatabase: @@ -202,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.""" @@ -237,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() @@ -334,9 +394,149 @@ 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() == [] + 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() == [] @@ -974,3 +1174,311 @@ 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)) + + 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] + + 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) + + tasks = [asyncio.create_task(session.pop_item()) for _ in items] + try: + popped = await asyncio.wait_for(asyncio.gather(*tasks), timeout=1) + finally: + 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 == { + "first", + "second", + "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: + """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..661124ddcb 100644 --- a/tests/extensions/memory/test_redis_session.py +++ b/tests/extensions/memory/test_redis_session.py @@ -1,5 +1,10 @@ from __future__ import annotations +import asyncio +import json +import sys +import time +from collections.abc import Awaitable from typing import Any, cast import pytest @@ -14,6 +19,21 @@ # 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 + + +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 @@ -125,6 +145,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() @@ -1267,6 +1341,1707 @@ 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) + + def __setattr__(self, name: str, value: Any) -> None: + setattr(delegate, name, value) + + 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") + + 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 __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) + + async def time(self) -> tuple[int, int]: + return next(server_times) + + async def execute(self) -> Any: + if attempt == 1: + await client.hset( # type: ignore[misc] + session._session_key, "concurrent_write", "1" + ) + 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_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: + 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: + """Cancellation before an immediate command must discard the watched connection.""" + 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 + pool = client.connection_pool + time_started = asyncio.Event() + dirty_connection: Any = None + + 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") + + 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("first-pre-exec-cancel") + 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, "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"], +) +async def test_post_commit_cancellation_propagates_after_cleanup( + monkeypatch: pytest.MonkeyPatch, + cleanup_mode: str, +) -> None: + """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=f"post_commit_cancellation_{cleanup_mode}", + 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() + if cleanup_mode == "failing_internal_reset" and reset_calls == 1: + raise RuntimeError("post-commit reset failed") + + 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("first-post-commit-cancel") + await asyncio.sleep(0) + add_task.cancel("second-during-cleanup") + await asyncio.sleep(0) + allow_internal_reset.set() + with pytest.raises(asyncio.CancelledError) as exc_info: + 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_cancel_message(exc_info.value, "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_reset_failure_detaches_without_release_listener( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """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) + 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) + 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] + + +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_close_failure_is_quarantined( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """A failed dirty-connection close must remain owned for close retry.""" + 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() + 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") + + 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"}])) + await time_started.wait() + add_task.cancel("caller-cancel") + + with pytest.raises(asyncio.CancelledError) as exc_info: + 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_skips_pipeline_reset( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """A protocol-dirty cancellation must not expose the connection through reset.""" + 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) + + async def controlled_time() -> tuple[int, int]: + time_started.set() + await asyncio.Event().wait() + raise AssertionError("unreachable") + + async def fail_if_reset() -> None: + raise AssertionError("dirty connection reached pipeline reset") + + monkeypatch.setattr(delegate, "time", controlled_time) + monkeypatch.setattr(delegate, "reset", fail_if_reset) + 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_cancel_message(exc_info.value, "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_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_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_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_transparent_overridden_pool_release_is_supported( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """A custom release that preserves pool ownership semantics remains supported.""" + 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 transparent_release(connection: Any) -> None: + nonlocal release_calls + release_calls += 1 + await real_release(connection) + + 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_release_failure_after_reborrow_preserves_new_borrower( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """A release-then-reborrow failure must not reclaim the new borrower's 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="ambiguous_pool_release", + redis_client=cast("Redis", client), + 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 + 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) + 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 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 borrowed_connection + await real_release(replacement_connection) + assert await session.get_items() == [item] + + +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) + 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 _release_after_detaching_pipeline_connection(delegate) + + 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) + + 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 _release_after_detaching_pipeline_connection(delegate) + + 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_cancel_message(exc_info.value, "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) + + 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 _release_after_detaching_pipeline_connection(delegate) + + 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( + 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..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,6 +238,291 @@ async def test_pop_from_empty_session(): assert popped is None +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, row: Any = None, rowcount: int = -1) -> None: + self._row = row + self.rowcount = rowcount + + def one_or_none(self) -> Any: + return self._row + + 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: + nonlocal delete_executed + if isinstance(statement, Select): + 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: + return FakeSession() + + 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", False) + + try: + popped = await session.pop_item() + finally: + await session.engine.dispose() + + 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_retries_returning_claim_lost_to_concurrent_delete( + monkeypatch: pytest.MonkeyPatch, +) -> 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, value: Any = None) -> None: + self._value = value + + def scalar_one_or_none(self) -> Any: + return self._value + + class FakeTransaction: + async def __aenter__(self) -> None: + return None + + 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 + + async def __aexit__(self, *args: Any) -> None: + return None + + def begin(self) -> FakeTransaction: + return FakeTransaction() + + async def execute(self, statement: Any) -> FakeResult: + if self._attempt == 1: + if isinstance(statement, Select): + return FakeResult(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: + return FakeSession() + + 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) + + 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: + 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_returned is True + + 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..8c8bf5ea0d 100644 --- a/tests/memory/test_session.py +++ b/tests/memory/test_session.py @@ -3,16 +3,51 @@ import asyncio import sqlite3 import tempfile +import threading from pathlib import Path -from typing import cast +from typing import Any, cast 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.""" @@ -967,3 +1002,213 @@ 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_propagates_after_known_outcome( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + operation: str, +): + """Cancellation after a worker commit must propagate without inviting a 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() + with pytest.raises(asyncio.CancelledError): + await mutation + finally: + conn.allow_return.set() + if not mutation.done(): + mutation.cancel() + await asyncio.gather(mutation, return_exceptions=True) + + if operation == "add": + assert await session.get_items() == [item] + elif operation == "pop": + assert await session.get_items() == [] + else: + assert await session.get_items() == [] + assert mutation.cancelled() + session.close()