diff --git a/src/agents/extensions/memory/advanced_sqlite_session.py b/src/agents/extensions/memory/advanced_sqlite_session.py index 83c289bdf8..b099ae3814 100644 --- a/src/agents/extensions/memory/advanced_sqlite_session.py +++ b/src/agents/extensions/memory/advanced_sqlite_session.py @@ -788,12 +788,18 @@ def _delete_sync(): conn.commit() - return usage_deleted, structure_deleted + # Clean up messages that are no longer referenced by any branch. + orphaned_deleted = self._cleanup_orphaned_messages_sync(conn) + if orphaned_deleted: + conn.commit() - usage_deleted, structure_deleted = await asyncio.to_thread(_delete_sync) + return usage_deleted, structure_deleted, orphaned_deleted + + usage_deleted, structure_deleted, orphaned_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} structure entries, " + f"{usage_deleted} usage entries, {orphaned_deleted} orphaned messages removed" ) async def list_branches(self) -> list[dict[str, Any]]: diff --git a/tests/extensions/memory/test_advanced_sqlite_session.py b/tests/extensions/memory/test_advanced_sqlite_session.py index ad4b5c4d86..a09e069281 100644 --- a/tests/extensions/memory/test_advanced_sqlite_session.py +++ b/tests/extensions/memory/test_advanced_sqlite_session.py @@ -1396,6 +1396,44 @@ async def add_batch(worker_id: int) -> list[str]: session.close() +async def test_delete_branch_cleans_up_orphaned_messages(): + """Verify that delete_branch removes messages only referenced by the deleted branch. + + Regression test for #3346: delete_branch() previously left branch-only + messages in the messages table after removing their structure metadata. + """ + session = AdvancedSQLiteSession(session_id="orphan_cleanup", create_tables=True) + + # Add items to main branch. + await session.add_items([{"role": "user", "content": "Shared message"}]) + + # Create a branch from turn 1 and add a branch-only message. + await session.create_branch_from_turn(1, "feature_branch") + await session.add_items([{"role": "user", "content": "Branch-only message"}]) + + # Record total message count before deletion. + with session._locked_connection() as conn: + total_before = conn.execute( + f"SELECT COUNT(*) FROM {session.messages_table} WHERE session_id = ?", + (session.session_id,), + ).fetchone()[0] + + assert total_before == 2 # shared + branch-only + + # Delete the feature branch. + await session.delete_branch("feature_branch", force=True) + + # The branch-only message should now be removed from messages table. + with session._locked_connection() as conn: + total_after = conn.execute( + f"SELECT COUNT(*) FROM {session.messages_table} WHERE session_id = ?", + (session.session_id,), + ).fetchone()[0] + + assert total_after == 1 # only the shared message remains + session.close() + + async def test_output_tokens_details_persisted_when_input_details_missing(): """Regression: output_tokens_details must persist even if input_tokens_details is None.