fix(sessions): clean up branch-only messages on delete

Author: Aphroq

src/agents/extensions/memory/advanced_sqlite_session.py

@@ -786,14 +786,19 @@ def _delete_sync():
 
                     structure_deleted = cursor.rowcount
 
+                    orphaned_messages_deleted = self._cleanup_orphaned_messages_sync(conn)
+
                     conn.commit()
 
-                    return usage_deleted, structure_deleted
+                    return usage_deleted, structure_deleted, orphaned_messages_deleted
 
-        usage_deleted, structure_deleted = await asyncio.to_thread(_delete_sync)
+        usage_deleted, structure_deleted, orphaned_messages_deleted = await asyncio.to_thread(
+            _delete_sync
+        )
 
         self._logger.info(
-            f"Deleted branch '{branch_id}': {structure_deleted} message entries, {usage_deleted} usage entries"  # noqa: E501
+            f"Deleted branch '{branch_id}': {structure_deleted} message entries, "
+            f"{usage_deleted} usage entries, {orphaned_messages_deleted} orphaned messages"
         )
 
     async def list_branches(self) -> list[dict[str, Any]]:

tests/extensions/memory/test_advanced_sqlite_session.py

@@ -442,6 +442,51 @@ async def test_branching_functionality(agent: Agent):
     session.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"
+    session = AdvancedSQLiteSession(session_id=session_id, create_tables=True)
+
+    main_items: list[TResponseInputItem] = [
+        {"role": "user", "content": "First question"},
+        {"role": "assistant", "content": "First answer"},
+        {"role": "user", "content": "Second question"},
+        {"role": "assistant", "content": "Second answer"},
+    ]
+    await session.add_items(main_items)
+
+    await session.create_branch_from_turn(2, "cleanup_branch")
+    branch_items: list[TResponseInputItem] = [
+        {"role": "user", "content": "Branch-only question"},
+        {"role": "assistant", "content": "Branch-only answer"},
+    ]
+    await session.add_items(branch_items)
+
+    await session.delete_branch("cleanup_branch", force=True)
+
+    with session._locked_connection() as conn:
+        rows = conn.execute(
+            f"""
+            SELECT message_data
+            FROM {session.messages_table}
+            WHERE session_id = ?
+            ORDER BY id
+            """,
+            (session.session_id,),
+        ).fetchall()
+
+    contents = [json.loads(message_data)["content"] for (message_data,) in rows]
+    assert contents == [
+        "First question",
+        "First answer",
+        "Second question",
+        "Second answer",
+    ]
+    assert await session.get_items(branch_id="main") == main_items
+
+    session.close()
+
+
 async def test_get_conversation_turns():
     """Test get_conversation_turns functionality."""
     session_id = "conversation_turns_test"