Skip to content

fix(advanced-sqlite-session): cleanup branch-only messages on delete_branch#3531

Open
LeSingh1 wants to merge 1 commit into
openai:mainfrom
LeSingh1:fix/delete-branch-orphaned-messages-3346
Open

fix(advanced-sqlite-session): cleanup branch-only messages on delete_branch#3531
LeSingh1 wants to merge 1 commit into
openai:mainfrom
LeSingh1:fix/delete-branch-orphaned-messages-3346

Conversation

@LeSingh1
Copy link
Copy Markdown
Contributor

Fixes #3346.

AdvancedSQLiteSession.delete_branch deletes the matching turn_usage and message_structure rows but never touched agent_messages. When messages were only referenced by the deleted branch, their agent_messages rows survived the deletion. After the delete:

  • branch queries via get_items() correctly omit them (no message_structure row), and
  • the rows are still sitting in the base table.

For branch-heavy sessions that meant invisible orphans accumulating.

Repro from the issue:

import asyncio
from agents.extensions.memory import AdvancedSQLiteSession

async def main():
    session = AdvancedSQLiteSession(session_id="repro", create_tables=True)
    await session.add_items([
        {"role": "user", "content": "main question"},
        {"role": "assistant", "content": "main answer"},
    ])
    await session.create_branch_from_turn(1, "branch_only")
    await session.add_items([
        {"role": "user", "content": "branch-only question"},
        {"role": "assistant", "content": "branch-only answer"},
    ])
    await session.delete_branch("branch_only", force=True)

    with session._locked_connection() as conn:
        rows = conn.execute(
            "SELECT id, message_data FROM agent_messages WHERE session_id = ? ORDER BY id",
            (session.session_id,),
        ).fetchall()
    print(rows)
    session.close()

asyncio.run(main())

Before: four rows (the branch-only pair is still there). After this PR: two rows (only main).

The existing _cleanup_orphaned_messages_sync helper already does exactly the right thing — it scopes to self.session_id and removes only rows with no message_structure reference, so it preserves messages that other branches still share via _copy_messages_to_new_branch. The fix calls it inside delete_branch's existing transaction, after the message_structure delete, so the cleanup commits with the rest of the deletion.

Tests:

  • test_delete_branch_cleans_branch_only_messages reproduces the issue and asserts only the main messages remain after the branch is deleted.
  • test_delete_branch_preserves_messages_shared_with_main forks a branch at turn 2, deletes it, and confirms turn 1's messages stay because main still references them.

Verification: uv run pytest tests/extensions/memory/test_advanced_sqlite_session.py → 38 passed. make format / make lint / make typecheck clean on the touched files.

…branch (fixes openai#3346)

AdvancedSQLiteSession.delete_branch removed turn_usage and
message_structure rows for the deleted branch but left the underlying
agent_messages rows in the table when those messages were only
referenced by the deleted branch. After deletion those rows were
invisible to advanced branch queries such as get_items() yet still
took up space, so branch-heavy sessions accumulated orphaned rows.

The existing _cleanup_orphaned_messages_sync helper already scopes to
self.session_id and only removes messages with no message_structure
reference, so it preserves messages that the main branch (or any
other branch) still shares via _copy_messages_to_new_branch. Reuse it
inside delete_branch's transaction, after the message_structure delete,
so the cleanup commits with the rest of the deletion.

Adds two regression tests: one for the branch-only cleanup case from
the issue, and one that verifies messages still referenced by main are
preserved when a forked branch is deleted.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

AdvancedSQLiteSession.delete_branch() leaves branch-only messages in the base table

1 participant