fix: make AdvancedSQLiteSession.pop_item branch-aware

src/agents/extensions/memory/advanced_sqlite_session.py

@@ -266,6 +266,65 @@ 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.
+
+        The base SQLiteSession.pop_item deletes the newest row for the whole
+        session regardless of branch, so on a non-main branch it removes an item
+        from another branch (usually main) and corrupts it. Scope the pop to the
+        active branch instead: drop the branch's newest message_structure row,
+        and delete the underlying message only when no other branch still
+        references it (branches share message rows, and foreign keys are not
+        enforced so ON DELETE CASCADE does not run).
+
+        Returns:
+            The most recent item on the current branch, or None if it is empty.
+        """
+
+        def _pop_item_sync() -> TResponseInputItem | None:
+            with self._locked_connection() as conn:
+                while True:
+                    cursor = conn.execute(
+                        f"""
+                        SELECT s.rowid, m.id, m.message_data
+                        FROM {self.messages_table} m
+                        JOIN message_structure s ON m.id = s.message_id
+                        WHERE m.session_id = ? AND s.branch_id = ?
+                        ORDER BY s.sequence_number DESC
+                        LIMIT 1
+                        """,
+                        (self.session_id, self._current_branch_id),
+                    )
+                    row = cursor.fetchone()
+                    if row is None:
+                        return None
+
+                    structure_rowid, message_id, message_data = row
+                    conn.execute(
+                        "DELETE FROM message_structure WHERE rowid = ?",
+                        (structure_rowid,),
+                    )
+                    # Branches share message rows, so only delete the underlying
+                    # message when no other branch still references it.
+                    still_referenced = conn.execute(
+                        "SELECT 1 FROM message_structure WHERE message_id = ? LIMIT 1",
+                        (message_id,),
+                    ).fetchone()
+                    if still_referenced is None:
+                        conn.execute(
+                            f"DELETE FROM {self.messages_table} WHERE id = ?",
+                            (message_id,),
+                        )
+                    conn.commit()
+
+                    try:
+                        return cast(TResponseInputItem, json.loads(message_data))
+                    except (json.JSONDecodeError, TypeError):
+                        # Drop corrupted JSON entries and keep looking on this branch.
+                        continue
+
+        return await asyncio.to_thread(_pop_item_sync)
+
     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

@@ -615,6 +615,65 @@ async def test_branching_functionality(agent: Agent):
     session.close()
 
 
+async def test_pop_item_is_branch_scoped():
+    """pop_item must only affect the current branch and never corrupt main."""
+    session = AdvancedSQLiteSession(session_id="pop_item_branch_test", create_tables=True)
+
+    await session.add_items(
+        [
+            {"role": "user", "content": "First question"},
+            {"role": "assistant", "content": "First answer"},
+        ]
+    )
+    await session.add_items(
+        [
+            {"role": "user", "content": "Second question"},
+            {"role": "assistant", "content": "Second answer"},
+        ]
+    )
+
+    # Branch from turn 2 -> the current branch holds only the first turn.
+    await session.create_branch_from_turn(2, "test_branch")
+    assert session._current_branch_id == "test_branch"
+    assert [i.get("content") for i in await session.get_items()] == [
+        "First question",
+        "First answer",
+    ]
+
+    # Popping on the branch returns the branch tail, not main's newest item.
+    popped = await session.pop_item()
+    assert popped is not None
+    assert popped.get("content") == "First answer"
+    assert [i.get("content") for i in await session.get_items()] == ["First question"]
+
+    # main must be left untouched.
+    main_items = await session.get_items(branch_id="main")
+    assert [i.get("content") for i in main_items] == [
+        "First question",
+        "First answer",
+        "Second question",
+        "Second answer",
+    ]
+
+    # Popping the current (main) branch still works normally.
+    await session.switch_to_branch("main")
+    popped_main = await session.pop_item()
+    assert popped_main is not None
+    assert popped_main.get("content") == "Second answer"
+    assert [i.get("content") for i in await session.get_items()] == [
+        "First question",
+        "First answer",
+        "Second question",
+    ]
+
+    session.close()
+
+    # An empty branch pops to None.
+    empty = AdvancedSQLiteSession(session_id="empty_pop_test", create_tables=True)
+    assert await empty.pop_item() is None
+    empty.close()
+
+
 async def test_delete_branch_removes_branch_only_messages():
     """Deleting a branch should not leave unreferenced branch-only messages behind."""
     session_id = "branch_delete_cleanup_test"