fix: AdvancedSQLiteSession clear_session and pop_item metadata leaks

src/agents/extensions/memory/advanced_sqlite_session.py

@@ -62,8 +62,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. 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.
+        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.
 
@@ -266,6 +286,154 @@ 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()
+
+                        try:
+                            # 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 it.
+                            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()
+                        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
+
+                        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:
+                try:
+                    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
+                # 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 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.
 
@@ -277,13 +445,56 @@ async def store_run_usage(self, result: RunResult) -> None:
         """
         try:
             if result.context_wrapper.usage is not None:
-                # Get the current turn number for this branch
-                current_turn = self._get_current_turn_number()
+                # Capture the current turn together with an anchor that pins the
+                # exact turn incarnation: the id of its first message_structure
+                # row (ids are monotonic and never reused). If that turn is
+                # removed before the write commits — even if a new turn later
+                # reuses the same numeric id — the anchor row is gone and the
+                # write is skipped. The anchor is scoped to this branch/turn, so
+                # unrelated removals (e.g. delete_branch on another branch) do
+                # not drop this write.
+                current_turn, branch_id, turn_anchor = self._capture_current_turn()
                 # 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,
+                    branch_id=branch_id,
+                    turn_anchor=turn_anchor,
+                )
         except Exception as e:
             self._logger.error("Failed to store usage for session %s: %s", self.session_id, e)
 
+    def _capture_current_turn(self) -> tuple[int, str, int | None]:
+        """Return (current_turn, branch_id, turn_anchor) in one locked read.
+
+        ``turn_anchor`` is the smallest ``message_structure.id`` of the current
+        turn on the current branch (``None`` if the turn has no rows). Because
+        ids are monotonic and never reused, it uniquely identifies this turn
+        incarnation, so a later pop+recreate that reuses the numeric turn id
+        yields a different anchor.
+        """
+        with self._locked_connection() as conn:
+            with closing(conn.cursor()) as cursor:
+                branch_id = self._current_branch_id
+                cursor.execute(
+                    """
+                    SELECT COALESCE(MAX(user_turn_number), 0)
+                    FROM message_structure
+                    WHERE session_id = ? AND branch_id = ?
+                    """,
+                    (self.session_id, branch_id),
+                )
+                current_turn = cursor.fetchone()[0]
+                cursor.execute(
+                    """
+                    SELECT MIN(id) FROM message_structure
+                    WHERE session_id = ? AND branch_id = ? AND user_turn_number = ?
+                    """,
+                    (self.session_id, branch_id, current_turn),
+                )
+                turn_anchor = cursor.fetchone()[0]
+                return current_turn, branch_id, turn_anchor
+
     def _get_next_turn_number(self, branch_id: str) -> int:
         """Get the next turn number for a specific branch.
 
@@ -608,6 +819,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."""
@@ -648,9 +864,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(
             "Created branch '%s' from turn %s ('%s') in '%s'",
@@ -694,6 +912,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."""
@@ -714,8 +936,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("Switched from branch '%s' to '%s'", old_branch, 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("Switched from branch '%s' to '%s'", old_branch, branch_id)
 
     async def delete_branch(self, branch_id: str, force: bool = False) -> None:
         """Delete a branch and all its associated data.
@@ -1305,17 +1530,48 @@ 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,
+        branch_id: str | None = None,
+        turn_anchor: 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.
+            branch_id: The branch the turn was read from. Defaults to the current
+                branch when not provided.
+            turn_anchor: The id of the turn's first ``message_structure`` row,
+                captured when the turn was read. When provided, the write is
+                skipped unless that exact row still exists for the given
+                branch/turn, so usage is never recorded against a turn that was
+                removed — even if a new turn reused the same numeric id. Because
+                the check is scoped to this branch/turn, unrelated removals (e.g.
+                delete_branch on another branch) do not drop this write.
         """
 
+        target_branch = branch_id if branch_id is not None else self._current_branch_id
+
         def _update_sync():
             """Synchronous helper to update turn usage data."""
             with self._locked_connection() as conn:
+                if turn_anchor is not None:
+                    with closing(conn.cursor()) as guard_cursor:
+                        guard_cursor.execute(
+                            """
+                            SELECT 1 FROM message_structure
+                            WHERE session_id = ? AND branch_id = ?
+                            AND user_turn_number = ? AND id = ?
+                            """,
+                            (self.session_id, target_branch, user_turn_number, turn_anchor),
+                        )
+                        if guard_cursor.fetchone() is None:
+                            # The exact turn incarnation is gone (removed, or its
+                            # numeric id reused by a new turn); skip the stale write.
+                            return
                 # Serialize token details as JSON
                 input_details_json = None
                 output_details_json = None
@@ -1347,7 +1603,7 @@ def _update_sync():
                     """,  # noqa: E501
                         (
                             self.session_id,
-                            self._current_branch_id,
+                            target_branch,
                             user_turn_number,
                             usage_data.requests or 0,
                             usage_data.input_tokens or 0,