fix: AdvancedSQLiteSession clear_session and pop_item metadata leaks

src/agents/extensions/memory/advanced_sqlite_session.py

@@ -252,6 +252,120 @@ 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.
+        """
+
+        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 current 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, self._current_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, self._current_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, self._current_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()
+
+        await asyncio.to_thread(_clear_session_sync)
+        # All branches were removed, so reset the in-memory pointer to 'main'.
+        self._current_branch_id = "main"
+
     async def store_run_usage(self, result: RunResult) -> None:
         """Store usage data for the current conversation turn.
 

tests/extensions/memory/test_advanced_sqlite_session.py

@@ -1757,3 +1757,173 @@ async def test_output_tokens_details_persisted_when_input_details_missing():
     assert turn_usage["output_tokens_details"] == {"reasoning_tokens": 42}
     assert turn_usage["input_tokens_details"] is None
     session.close()
+
+
+def _count_rows(session: AdvancedSQLiteSession, table: str) -> int:
+    """Helper: count rows for the session in one of the metadata tables."""
+    with session._locked_connection() as conn:
+        row = conn.execute(
+            f"SELECT COUNT(*) FROM {table} WHERE session_id = ?",
+            (session.session_id,),
+        ).fetchone()
+        return cast(int, row[0])
+
+
+async def test_clear_session_removes_structure_and_usage_metadata(usage_data: Usage):
+    """Regression: clear_session must also clear message_structure and turn_usage.
+
+    Those tables declare an ON DELETE CASCADE foreign key, but SQLite does not
+    enforce foreign keys by default, so the inherited base clear_session left the
+    rows behind. That leaked stale structure/usage data and permanently offset
+    sequence and turn numbering for items added after clearing.
+    """
+    session = AdvancedSQLiteSession(session_id="clear_metadata_test", create_tables=True)
+
+    await session.add_items(
+        [
+            {"role": "user", "content": "First question"},
+            {"role": "assistant", "content": "First answer"},
+        ]
+    )
+    await session.store_run_usage(create_mock_run_result(usage_data))
+
+    assert _count_rows(session, "message_structure") > 0
+    assert _count_rows(session, "turn_usage") > 0
+
+    await session.clear_session()
+
+    assert await session.get_items() == []
+    assert _count_rows(session, "message_structure") == 0
+    assert _count_rows(session, "turn_usage") == 0
+
+    # Numbering must reset: the next item starts a fresh sequence and turn.
+    await session.add_items([{"role": "user", "content": "Fresh start"}])
+    with session._locked_connection() as conn:
+        rows = conn.execute(
+            """
+            SELECT sequence_number, user_turn_number
+            FROM message_structure
+            WHERE session_id = ?
+            """,
+            (session.session_id,),
+        ).fetchall()
+    assert rows == [(1, 1)]
+
+    session.close()
+
+
+async def test_pop_item_removes_its_structure_row():
+    """Regression: pop_item must delete the popped message's structure row.
+
+    The inherited base pop_item removed only the message row, leaving an orphaned
+    message_structure row that corrupted later MAX(sequence_number)/turn numbering.
+    """
+    session = AdvancedSQLiteSession(session_id="pop_structure_test", create_tables=True)
+
+    await session.add_items(
+        [
+            {"role": "user", "content": "Question"},
+            {"role": "assistant", "content": "Answer"},
+        ]
+    )
+
+    popped = await session.pop_item()
+    assert popped == {"role": "assistant", "content": "Answer"}
+
+    with session._locked_connection() as conn:
+        message_ids = {
+            row[0]
+            for row in conn.execute(
+                f"SELECT id FROM {session.messages_table} WHERE session_id = ?",
+                (session.session_id,),
+            ).fetchall()
+        }
+        structure_message_ids = {
+            row[0]
+            for row in conn.execute(
+                "SELECT message_id FROM message_structure WHERE session_id = ?",
+                (session.session_id,),
+            ).fetchall()
+        }
+
+    # No structure row may reference a message that no longer exists.
+    assert structure_message_ids <= message_ids
+    assert await session.get_items() == [{"role": "user", "content": "Question"}]
+
+    session.close()
+
+
+async def test_pop_item_removes_turn_usage_only_when_turn_emptied(usage_data: Usage):
+    """Regression: pop_item must drop a turn's turn_usage row once the turn has no
+    remaining items on the current branch, while keeping it for a partial pop.
+    """
+    session = AdvancedSQLiteSession(session_id="pop_turn_usage_test", create_tables=True)
+
+    # One turn with two items, plus stored usage for that turn.
+    await session.add_items(
+        [
+            {"role": "user", "content": "Question"},
+            {"role": "assistant", "content": "Answer"},
+        ]
+    )
+    await session.store_run_usage(create_mock_run_result(usage_data))
+    assert _count_rows(session, "turn_usage") == 1
+
+    # Popping only the assistant item leaves the turn non-empty: usage is kept.
+    await session.pop_item()
+    assert _count_rows(session, "turn_usage") == 1
+
+    # Popping the last item of the turn removes the now-stale usage row.
+    await session.pop_item()
+    assert _count_rows(session, "turn_usage") == 0
+    assert not await session.get_turn_usage(1)
+
+    session.close()
+
+
+async def test_pop_item_respects_current_branch_and_keeps_shared_messages():
+    """Regression: pop_item must pop from the current branch and preserve messages
+    still referenced by another branch (branches share the underlying message rows).
+    """
+    session = AdvancedSQLiteSession(session_id="pop_branch_test", create_tables=True)
+
+    main_items: list[TResponseInputItem] = [
+        {"role": "user", "content": "Main first question"},
+        {"role": "assistant", "content": "Main first answer"},
+        {"role": "user", "content": "Main second question"},
+        {"role": "assistant", "content": "Main second answer"},
+    ]
+
+    try:
+        await session.add_items(main_items)
+        # Branch from turn 2 copies turn 1's shared messages into the new branch.
+        await session.create_branch_from_turn(2, "branch_a")
+        await session.switch_to_branch("branch_a")
+        await session.add_items([{"role": "user", "content": "Branch-only question"}])
+
+        # Popping on branch_a removes only its own newest item.
+        popped = await session.pop_item()
+        assert popped == {"role": "user", "content": "Branch-only question"}
+
+        # The main branch, which shares turn 1's messages, is untouched.
+        assert await session.get_items(branch_id="main") == main_items
+
+        # No orphaned structure rows anywhere in the session.
+        with session._locked_connection() as conn:
+            message_ids = {
+                row[0]
+                for row in conn.execute(
+                    f"SELECT id FROM {session.messages_table} WHERE session_id = ?",
+                    (session.session_id,),
+                ).fetchall()
+            }
+            structure_message_ids = {
+                row[0]
+                for row in conn.execute(
+                    "SELECT message_id FROM message_structure WHERE session_id = ?",
+                    (session.session_id,),
+                ).fetchall()
+            }
+        assert structure_message_ids <= message_ids
+    finally:
+        session.close()