fix: AdvancedSQLiteSession clear_session and pop_item metadata leaks

src/agents/extensions/memory/advanced_sqlite_session.py

@@ -48,8 +48,28 @@ def __init__(
         if create_tables:
             self._init_structure_tables()
         self._current_branch_id = "main"
+        # Bumped (under the connection lock) whenever clear_session() wipes the
+        # session. Branch-pointer and usage writes capture the generation before
+        # their DB work and skip their mutation if a clear has committed since,
+        # so a stale switch/create/store_run_usage cannot resurrect a branch or
+        # a turn that clear already removed.
+        self._generation = 0
         self._logger = logger or logging.getLogger(__name__)
 
+    def _commit_branch_pointer(self, branch_id: str, generation: int) -> bool:
+        """Set the current-branch pointer unless a clear has committed meanwhile.
+
+        Acquires the connection lock so the generation check and the assignment
+        are atomic with clear_session's reset. Returns True if the pointer was
+        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:
+                return False
+            self._current_branch_id = branch_id
+            return True
+
     def _init_structure_tables(self):
         """Add structure and usage tracking tables.
 
@@ -252,6 +272,132 @@ def _get_items_sync():
 
         return await asyncio.to_thread(_get_items_sync)
 
+    async def pop_item(self) -> TResponseInputItem | None:
+        """Remove and return the most recent item from the current branch.
+
+        Overrides the base implementation so the popped message's
+        `message_structure` row is removed in the same transaction and only the
+        current branch is affected. The underlying message row is deleted only
+        when no other branch still references it, mirroring `delete_branch`. When
+        popping empties a turn on the current branch, its `turn_usage` row is
+        removed as well so usage analytics do not report a turn that no longer
+        exists.
+        """
+
+        # Snapshot the current branch at call time so a concurrent
+        # 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
+
+        def _pop_item_sync():
+            with self._locked_connection() as conn:
+                while True:
+                    with closing(conn.cursor()) as cursor:
+                        # Find the most recent item on the snapshotted branch.
+                        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
+                            """,
+                            (self.session_id, branch_id),
+                        )
+                        row = cursor.fetchone()
+                        if row is None:
+                            return None
+
+                        structure_id, message_id, user_turn_number = row
+
+                        # Read the message payload before removing anything.
+                        cursor.execute(
+                            f"SELECT message_data FROM {self.messages_table} WHERE id = ?",
+                            (message_id,),
+                        )
+                        message_row = cursor.fetchone()
+
+                        # Remove the structure row for this branch, then drop the
+                        # underlying message only if no other branch references it.
+                        cursor.execute(
+                            "DELETE FROM message_structure WHERE id = ?",
+                            (structure_id,),
+                        )
+                        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 that turn.
+                        if user_turn_number is not None:
+                            cursor.execute(
+                                """
+                                SELECT COUNT(*) FROM message_structure
+                                WHERE session_id = ? AND branch_id = ? AND user_turn_number = ?
+                                """,
+                                (self.session_id, branch_id, user_turn_number),
+                            )
+                            if cursor.fetchone()[0] == 0:
+                                cursor.execute(
+                                    """
+                                    DELETE FROM turn_usage
+                                    WHERE session_id = ? AND branch_id = ? AND user_turn_number = ?
+                                    """,
+                                    (self.session_id, branch_id, user_turn_number),
+                                )
+
+                        conn.commit()
+
+                        if message_row is None:
+                            # Structure row pointed at a missing message; keep looking.
+                            continue
+
+                        try:
+                            return json.loads(message_row[0])
+                        except (json.JSONDecodeError, TypeError):
+                            # Drop corrupted JSON entries and keep looking for a valid item.
+                            continue
+
+        return await asyncio.to_thread(_pop_item_sync)
+
+    async def clear_session(self) -> None:
+        """Clear all items for this session.
+
+        Overrides the base implementation so the `message_structure` and
+        `turn_usage` metadata tables are cleared in the same transaction. Those
+        rows declare an `ON DELETE CASCADE` foreign key, but SQLite does not
+        enforce foreign keys unless `PRAGMA foreign_keys=ON` is set, so they must
+        be deleted explicitly to avoid leaking stale structure and usage data.
+        """
+
+        def _clear_session_sync():
+            with self._locked_connection() as conn:
+                conn.execute(
+                    f"DELETE FROM {self.messages_table} WHERE session_id = ?",
+                    (self.session_id,),
+                )
+                conn.execute(
+                    f"DELETE FROM {self.sessions_table} WHERE session_id = ?",
+                    (self.session_id,),
+                )
+                conn.execute(
+                    "DELETE FROM message_structure WHERE session_id = ?",
+                    (self.session_id,),
+                )
+                conn.execute(
+                    "DELETE FROM turn_usage WHERE session_id = ?",
+                    (self.session_id,),
+                )
+                conn.commit()
+                # All branches were removed, so reset the in-memory pointer to
+                # 'main' while still holding the lock. Doing this inside the
+                # locked operation keeps the reset atomic with the clear, so no
+                # other locked operation observes the session as cleared while
+                # the pointer still references a deleted branch. Bumping the
+                # generation invalidates any in-flight switch/create/usage write
+                # that captured the pre-clear generation.
+                self._generation += 1
+                self._current_branch_id = "main"
+
+        await asyncio.to_thread(_clear_session_sync)
+
     async def store_run_usage(self, result: RunResult) -> None:
         """Store usage data for the current conversation turn.
 
@@ -263,10 +409,15 @@ async def store_run_usage(self, result: RunResult) -> None:
         """
         try:
             if result.context_wrapper.usage is not None:
+                # Capture the generation before reading the turn so a clear that
+                # commits before the usage write is detected and skipped.
+                generation = self._generation
                 # Get the current turn number for this branch
                 current_turn = self._get_current_turn_number()
                 # Only update turn-level usage - session usage is aggregated on demand
-                await self._update_turn_usage_internal(current_turn, result.context_wrapper.usage)
+                await self._update_turn_usage_internal(
+                    current_turn, result.context_wrapper.usage, generation
+                )
         except Exception as e:
             self._logger.error(f"Failed to store usage for session {self.session_id}: {e}")
 
@@ -594,6 +745,11 @@ async def create_branch_from_turn(
         """
         import time
 
+        # Capture the generation before any DB work so a clear that commits
+        # while this branch is being created cannot be overwritten by the
+        # pointer update below.
+        generation = self._generation
+
         # Validate the turn exists and contains a user message
         def _validate_turn():
             """Synchronous helper to validate turn exists and contains user message."""
@@ -634,9 +790,11 @@ def _validate_turn():
         # Copy messages before the branch point to the new branch
         await self._copy_messages_to_new_branch(branch_name, turn_number)
 
-        # Switch to new branch
+        # 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.
         old_branch = self._current_branch_id
-        self._current_branch_id = branch_name
+        await asyncio.to_thread(self._commit_branch_pointer, branch_name, generation)
 
         self._logger.debug(
             f"Created branch '{branch_name}' from turn {turn_number} ('{turn_content}') in '{old_branch}'"  # noqa: E501
@@ -676,6 +834,10 @@ 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."""
@@ -696,8 +858,11 @@ def _validate_branch():
         await asyncio.to_thread(_validate_branch)
 
         old_branch = self._current_branch_id
-        self._current_branch_id = branch_id
-        self._logger.info(f"Switched from branch '{old_branch}' to '{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)
+        if switched:
+            self._logger.info(f"Switched from branch '{old_branch}' to '{branch_id}'")
 
     async def delete_branch(self, branch_id: str, force: bool = False) -> None:
         """Delete a branch and all its associated data.
@@ -1286,17 +1451,40 @@ def _get_turn_usage_sync():
 
         return cast(list[dict[str, Any]] | dict[str, Any], result)
 
-    async def _update_turn_usage_internal(self, user_turn_number: int, usage_data: Usage) -> None:
+    async def _update_turn_usage_internal(
+        self, user_turn_number: int, usage_data: Usage, generation: int | None = None
+    ) -> None:
         """Internal method to update usage for a specific turn with full JSON details.
 
         Args:
             user_turn_number: The turn number to update usage for.
             usage_data: The usage data to store.
+            generation: The generation captured before the turn was read. When
+                provided, the write is skipped if a clear_session has committed
+                since (generation mismatch) or if the turn no longer exists on
+                the current branch (e.g. it was removed by pop_item), so stale
+                usage is never recorded for a turn that no longer exists.
         """
 
         def _update_sync():
             """Synchronous helper to update turn usage data."""
             with self._locked_connection() as conn:
+                if generation is not None:
+                    if self._generation != generation:
+                        # A clear_session committed after the turn was read.
+                        return
+                    with closing(conn.cursor()) as guard_cursor:
+                        guard_cursor.execute(
+                            """
+                            SELECT COUNT(*) FROM message_structure
+                            WHERE session_id = ? AND branch_id = ? AND user_turn_number = ?
+                            """,
+                            (self.session_id, self._current_branch_id, user_turn_number),
+                        )
+                        if guard_cursor.fetchone()[0] == 0:
+                            # The turn was removed (e.g. by pop_item) after it was
+                            # read; do not resurrect usage for a nonexistent turn.
+                            return
                 # Serialize token details as JSON
                 input_details_json = None
                 output_details_json = None