From 26020758244552a9e120074bd87c9c17a051d77a Mon Sep 17 00:00:00 2001 From: Aryan Putta Date: Tue, 4 Aug 2026 19:43:32 -0400 Subject: [PATCH] fix(memory): reject an existing branch id in create_branch_from_turn create_branch_from_turn() passed branch_name straight to _copy_messages_to_new_branch(), which unconditionally inserts message_structure rows tagged with that id. Passing the id of a branch that already exists therefore merged the copied turns into it and then switched to it, with no error and no way to tell the two histories apart afterwards. Check the id first and raise ValueError, which the docstring already documents as the failure mode for this method. A generated name is unaffected, since it embeds a timestamp. Closes #4150 --- .../memory/advanced_sqlite_session.py | 24 +++++++++++++++++ .../memory/test_advanced_sqlite_session.py | 26 +++++++++++++++++++ 2 files changed, 50 insertions(+) diff --git a/src/agents/extensions/memory/advanced_sqlite_session.py b/src/agents/extensions/memory/advanced_sqlite_session.py index cbf8f510f1..3484e1e800 100644 --- a/src/agents/extensions/memory/advanced_sqlite_session.py +++ b/src/agents/extensions/memory/advanced_sqlite_session.py @@ -863,6 +863,11 @@ def _validate_turn(): if branch_name is None: timestamp = int(time.time()) branch_name = f"branch_from_turn_{turn_number}_{timestamp}" + elif await self._branch_exists(branch_name): + raise ValueError( + f"Branch '{branch_name}' already exists in session '{self.session_id}'. " + "Pass a different branch_name, or omit it to generate one." + ) # Copy messages before the branch point to the new branch await self._copy_messages_to_new_branch(branch_name, turn_number) @@ -1087,6 +1092,25 @@ def _list_branches_sync(): return await asyncio.to_thread(_list_branches_sync) + async def _branch_exists(self, branch_id: str) -> bool: + """Whether any message in this session is already tagged with ``branch_id``.""" + + def _branch_exists_sync() -> bool: + with self._locked_connection() as conn: + with closing(conn.cursor()) as cursor: + cursor.execute( + """ + SELECT 1 + FROM message_structure + WHERE session_id = ? AND branch_id = ? + LIMIT 1 + """, + (self.session_id, branch_id), + ) + return cursor.fetchone() is not None + + return await asyncio.to_thread(_branch_exists_sync) + async def _copy_messages_to_new_branch(self, new_branch_id: str, from_turn_number: int) -> None: """Copy messages before the branch point to the new branch. diff --git a/tests/extensions/memory/test_advanced_sqlite_session.py b/tests/extensions/memory/test_advanced_sqlite_session.py index ae1606f249..5264fed783 100644 --- a/tests/extensions/memory/test_advanced_sqlite_session.py +++ b/tests/extensions/memory/test_advanced_sqlite_session.py @@ -182,6 +182,32 @@ async def test_create_branch_logging_respects_model_data_policy(monkeypatch, red session.close() +async def test_create_branch_rejects_existing_branch_id(): + session = AdvancedSQLiteSession(session_id="s4150", create_tables=True) + try: + await session.add_items( + [ + {"role": "user", "content": "one"}, + {"role": "assistant", "content": "a"}, + {"role": "user", "content": "two"}, + {"role": "assistant", "content": "b"}, + ] + ) + await session.create_branch_from_turn(2, "dup") + # Return to a branch where turn 2 still resolves, so the only thing + # left for the second call to object to is the duplicate branch id. + await session.switch_to_branch("main") + before = {b["branch_id"]: b["message_count"] for b in await session.list_branches()} + + with pytest.raises(ValueError, match="already exists"): + await session.create_branch_from_turn(2, "dup") + + after = {b["branch_id"]: b["message_count"] for b in await session.list_branches()} + assert after == before, "a rejected branch creation must not modify any branch" + finally: + session.close() + + async def test_advanced_session_respects_custom_table_names(): """AdvancedSQLiteSession should consistently use configured table names.""" session = AdvancedSQLiteSession(