From 223fab9be07df16c3f837ebc3bcf97a69d9ec703 Mon Sep 17 00:00:00 2001 From: Aditya Jethani Date: Wed, 8 Jul 2026 13:25:46 +0530 Subject: [PATCH 1/5] Fix AdvancedSQLiteSession clear_session and pop_item metadata leaks AdvancedSQLiteSession maintains the auxiliary message_structure and turn_usage tables in its add_items override, but inherited clear_session and pop_item unchanged from SQLiteSession. Those base methods only touch the messages and sessions tables. The aux tables declare ON DELETE CASCADE foreign keys, but SQLite does not enforce foreign keys unless PRAGMA foreign_keys=ON is set, and the SDK never enables it. As a result: - clear_session left every message_structure and turn_usage row behind, so list_branches/get_session_usage/get_turn_usage reported stale data and sequence/turn numbering stayed permanently offset for items added after clearing. - pop_item deleted the globally highest-id message ignoring the current branch, and never removed the corresponding message_structure row, orphaning it and corrupting later MAX(sequence_number) numbering. Override both methods to keep the metadata tables consistent, mirroring the existing delete_branch cleanup. pop_item now pops the most recent item of the current branch and reuses _cleanup_orphaned_messages_sync so the shared message row is removed only when no other branch references it. Adds regression tests that fail without the fix. --- .../memory/advanced_sqlite_session.py | 91 +++++++++++ .../memory/test_advanced_sqlite_session.py | 142 ++++++++++++++++++ 2 files changed, 233 insertions(+) diff --git a/src/agents/extensions/memory/advanced_sqlite_session.py b/src/agents/extensions/memory/advanced_sqlite_session.py index 704c85058e..dadc49fa34 100644 --- a/src/agents/extensions/memory/advanced_sqlite_session.py +++ b/src/agents/extensions/memory/advanced_sqlite_session.py @@ -252,6 +252,97 @@ 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. + + Overrides the base implementation so the popped message's + `message_structure` row is removed in the same transaction and only the + current branch is affected. The underlying message row is deleted only + when no other branch still references it, mirroring `delete_branch`. + """ + + def _pop_item_sync(): + with self._locked_connection() as conn: + while True: + with closing(conn.cursor()) as cursor: + # Find the most recent item on the current branch. + cursor.execute( + """ + SELECT id, message_id FROM message_structure + WHERE session_id = ? AND branch_id = ? + ORDER BY sequence_number DESC + LIMIT 1 + """, + (self.session_id, self._current_branch_id), + ) + row = cursor.fetchone() + if row is None: + return None + + structure_id, message_id = row + + # Read the message payload before removing anything. + cursor.execute( + f"SELECT message_data FROM {self.messages_table} WHERE id = ?", + (message_id,), + ) + message_row = cursor.fetchone() + + # Remove the structure row for this branch, then drop the + # underlying message only if no other branch references it. + cursor.execute( + "DELETE FROM message_structure WHERE id = ?", + (structure_id,), + ) + self._cleanup_orphaned_messages_sync(conn) + conn.commit() + + if message_row is None: + # Structure row pointed at a missing message; keep looking. + continue + + try: + return json.loads(message_row[0]) + except (json.JSONDecodeError, TypeError): + # Drop corrupted JSON entries and keep looking for a valid item. + continue + + return await asyncio.to_thread(_pop_item_sync) + + async def clear_session(self) -> None: + """Clear all items for this session. + + Overrides the base implementation so the `message_structure` and + `turn_usage` metadata tables are cleared in the same transaction. Those + rows declare an `ON DELETE CASCADE` foreign key, but SQLite does not + enforce foreign keys unless `PRAGMA foreign_keys=ON` is set, so they must + be deleted explicitly to avoid leaking stale structure and usage data. + """ + + def _clear_session_sync(): + with self._locked_connection() as conn: + conn.execute( + f"DELETE FROM {self.messages_table} WHERE session_id = ?", + (self.session_id,), + ) + conn.execute( + f"DELETE FROM {self.sessions_table} WHERE session_id = ?", + (self.session_id,), + ) + conn.execute( + "DELETE FROM message_structure WHERE session_id = ?", + (self.session_id,), + ) + conn.execute( + "DELETE FROM turn_usage WHERE session_id = ?", + (self.session_id,), + ) + conn.commit() + + await asyncio.to_thread(_clear_session_sync) + # All branches were removed, so reset the in-memory pointer to 'main'. + self._current_branch_id = "main" + async def store_run_usage(self, result: RunResult) -> None: """Store usage data for the current conversation turn. diff --git a/tests/extensions/memory/test_advanced_sqlite_session.py b/tests/extensions/memory/test_advanced_sqlite_session.py index 9c543adc20..ff4e654cfb 100644 --- a/tests/extensions/memory/test_advanced_sqlite_session.py +++ b/tests/extensions/memory/test_advanced_sqlite_session.py @@ -1757,3 +1757,145 @@ async def test_output_tokens_details_persisted_when_input_details_missing(): assert turn_usage["output_tokens_details"] == {"reasoning_tokens": 42} assert turn_usage["input_tokens_details"] is None session.close() + + +def _count_rows(session: AdvancedSQLiteSession, table: str) -> int: + """Helper: count rows for the session in one of the metadata tables.""" + with session._locked_connection() as conn: + row = conn.execute( + f"SELECT COUNT(*) FROM {table} WHERE session_id = ?", + (session.session_id,), + ).fetchone() + return cast(int, row[0]) + + +async def test_clear_session_removes_structure_and_usage_metadata(usage_data: Usage): + """Regression: clear_session must also clear message_structure and turn_usage. + + Those tables declare an ON DELETE CASCADE foreign key, but SQLite does not + enforce foreign keys by default, so the inherited base clear_session left the + rows behind. That leaked stale structure/usage data and permanently offset + sequence and turn numbering for items added after clearing. + """ + session = AdvancedSQLiteSession(session_id="clear_metadata_test", create_tables=True) + + await session.add_items( + [ + {"role": "user", "content": "First question"}, + {"role": "assistant", "content": "First answer"}, + ] + ) + await session.store_run_usage(create_mock_run_result(usage_data)) + + assert _count_rows(session, "message_structure") > 0 + assert _count_rows(session, "turn_usage") > 0 + + await session.clear_session() + + assert await session.get_items() == [] + assert _count_rows(session, "message_structure") == 0 + assert _count_rows(session, "turn_usage") == 0 + + # Numbering must reset: the next item starts a fresh sequence and turn. + await session.add_items([{"role": "user", "content": "Fresh start"}]) + with session._locked_connection() as conn: + rows = conn.execute( + """ + SELECT sequence_number, user_turn_number + FROM message_structure + WHERE session_id = ? + """, + (session.session_id,), + ).fetchall() + assert rows == [(1, 1)] + + session.close() + + +async def test_pop_item_removes_its_structure_row(): + """Regression: pop_item must delete the popped message's structure row. + + The inherited base pop_item removed only the message row, leaving an orphaned + message_structure row that corrupted later MAX(sequence_number)/turn numbering. + """ + session = AdvancedSQLiteSession(session_id="pop_structure_test", create_tables=True) + + await session.add_items( + [ + {"role": "user", "content": "Question"}, + {"role": "assistant", "content": "Answer"}, + ] + ) + + popped = await session.pop_item() + assert popped == {"role": "assistant", "content": "Answer"} + + with session._locked_connection() as conn: + message_ids = { + row[0] + for row in conn.execute( + f"SELECT id FROM {session.messages_table} WHERE session_id = ?", + (session.session_id,), + ).fetchall() + } + structure_message_ids = { + row[0] + for row in conn.execute( + "SELECT message_id FROM message_structure WHERE session_id = ?", + (session.session_id,), + ).fetchall() + } + + # No structure row may reference a message that no longer exists. + assert structure_message_ids <= message_ids + assert await session.get_items() == [{"role": "user", "content": "Question"}] + + session.close() + + +async def test_pop_item_respects_current_branch_and_keeps_shared_messages(): + """Regression: pop_item must pop from the current branch and preserve messages + still referenced by another branch (branches share the underlying message rows). + """ + session = AdvancedSQLiteSession(session_id="pop_branch_test", create_tables=True) + + main_items: list[TResponseInputItem] = [ + {"role": "user", "content": "Main first question"}, + {"role": "assistant", "content": "Main first answer"}, + {"role": "user", "content": "Main second question"}, + {"role": "assistant", "content": "Main second answer"}, + ] + + try: + await session.add_items(main_items) + # Branch from turn 2 copies turn 1's shared messages into the new branch. + await session.create_branch_from_turn(2, "branch_a") + await session.switch_to_branch("branch_a") + await session.add_items([{"role": "user", "content": "Branch-only question"}]) + + # Popping on branch_a removes only its own newest item. + popped = await session.pop_item() + assert popped == {"role": "user", "content": "Branch-only question"} + + # The main branch, which shares turn 1's messages, is untouched. + assert await session.get_items(branch_id="main") == main_items + + # No orphaned structure rows anywhere in the session. + with session._locked_connection() as conn: + message_ids = { + row[0] + for row in conn.execute( + f"SELECT id FROM {session.messages_table} WHERE session_id = ?", + (session.session_id,), + ).fetchall() + } + structure_message_ids = { + row[0] + for row in conn.execute( + "SELECT message_id FROM message_structure WHERE session_id = ?", + (session.session_id,), + ).fetchall() + } + assert structure_message_ids <= message_ids + finally: + session.close() From 2c9725026b1bd78c285f723bce477fa7a30b5e07 Mon Sep 17 00:00:00 2001 From: Aditya Jethani Date: Wed, 8 Jul 2026 13:42:41 +0530 Subject: [PATCH 2/5] Clean up stale turn_usage in pop_item when a turn is emptied Address automated review feedback: when pop_item removes the last message_structure row of a turn on the current branch, delete that turn's turn_usage row as well so get_turn_usage/get_session_usage do not report a turn that no longer exists. A turn that still has items keeps its usage. Mirrors the turn_usage cleanup already done by delete_branch and clear_session. Adds a regression test. --- .../memory/advanced_sqlite_session.py | 29 +++++++++++++++++-- .../memory/test_advanced_sqlite_session.py | 28 ++++++++++++++++++ 2 files changed, 54 insertions(+), 3 deletions(-) diff --git a/src/agents/extensions/memory/advanced_sqlite_session.py b/src/agents/extensions/memory/advanced_sqlite_session.py index dadc49fa34..9d56b6e81a 100644 --- a/src/agents/extensions/memory/advanced_sqlite_session.py +++ b/src/agents/extensions/memory/advanced_sqlite_session.py @@ -258,7 +258,10 @@ async def pop_item(self) -> TResponseInputItem | None: Overrides the base implementation so the popped message's `message_structure` row is removed in the same transaction and only the current branch is affected. The underlying message row is deleted only - when no other branch still references it, mirroring `delete_branch`. + when no other branch still references it, mirroring `delete_branch`. When + popping empties a turn on the current branch, its `turn_usage` row is + removed as well so usage analytics do not report a turn that no longer + exists. """ def _pop_item_sync(): @@ -268,7 +271,7 @@ def _pop_item_sync(): # Find the most recent item on the current branch. cursor.execute( """ - SELECT id, message_id FROM message_structure + SELECT id, message_id, user_turn_number FROM message_structure WHERE session_id = ? AND branch_id = ? ORDER BY sequence_number DESC LIMIT 1 @@ -279,7 +282,7 @@ def _pop_item_sync(): if row is None: return None - structure_id, message_id = row + structure_id, message_id, user_turn_number = row # Read the message payload before removing anything. cursor.execute( @@ -295,6 +298,26 @@ def _pop_item_sync(): (structure_id,), ) self._cleanup_orphaned_messages_sync(conn) + + # If this was the last item of the turn on this branch, + # drop the now-stale turn_usage row for that turn. + if user_turn_number is not None: + cursor.execute( + """ + SELECT COUNT(*) FROM message_structure + WHERE session_id = ? AND branch_id = ? AND user_turn_number = ? + """, + (self.session_id, self._current_branch_id, user_turn_number), + ) + if cursor.fetchone()[0] == 0: + cursor.execute( + """ + DELETE FROM turn_usage + WHERE session_id = ? AND branch_id = ? AND user_turn_number = ? + """, + (self.session_id, self._current_branch_id, user_turn_number), + ) + conn.commit() if message_row is None: diff --git a/tests/extensions/memory/test_advanced_sqlite_session.py b/tests/extensions/memory/test_advanced_sqlite_session.py index ff4e654cfb..48ffc66a83 100644 --- a/tests/extensions/memory/test_advanced_sqlite_session.py +++ b/tests/extensions/memory/test_advanced_sqlite_session.py @@ -1853,6 +1853,34 @@ async def test_pop_item_removes_its_structure_row(): session.close() +async def test_pop_item_removes_turn_usage_only_when_turn_emptied(usage_data: Usage): + """Regression: pop_item must drop a turn's turn_usage row once the turn has no + remaining items on the current branch, while keeping it for a partial pop. + """ + session = AdvancedSQLiteSession(session_id="pop_turn_usage_test", create_tables=True) + + # One turn with two items, plus stored usage for that turn. + await session.add_items( + [ + {"role": "user", "content": "Question"}, + {"role": "assistant", "content": "Answer"}, + ] + ) + await session.store_run_usage(create_mock_run_result(usage_data)) + assert _count_rows(session, "turn_usage") == 1 + + # Popping only the assistant item leaves the turn non-empty: usage is kept. + await session.pop_item() + assert _count_rows(session, "turn_usage") == 1 + + # Popping the last item of the turn removes the now-stale usage row. + await session.pop_item() + assert _count_rows(session, "turn_usage") == 0 + assert not await session.get_turn_usage(1) + + session.close() + + async def test_pop_item_respects_current_branch_and_keeps_shared_messages(): """Regression: pop_item must pop from the current branch and preserve messages still referenced by another branch (branches share the underlying message rows). From 50d2b2f9ecb1c894061826ba1262ceb02fb54322 Mon Sep 17 00:00:00 2001 From: Aditya Jethani Date: Wed, 8 Jul 2026 13:56:44 +0530 Subject: [PATCH 3/5] pop_item: snapshot current branch before dispatching to the worker Read _current_branch_id once at call time and use that snapshot inside the worker thread, so a concurrent switch_to_branch() cannot redirect an in-flight pop to a different branch. --- .../extensions/memory/advanced_sqlite_session.py | 13 +++++++++---- 1 file changed, 9 insertions(+), 4 deletions(-) diff --git a/src/agents/extensions/memory/advanced_sqlite_session.py b/src/agents/extensions/memory/advanced_sqlite_session.py index 9d56b6e81a..0e81ee80cb 100644 --- a/src/agents/extensions/memory/advanced_sqlite_session.py +++ b/src/agents/extensions/memory/advanced_sqlite_session.py @@ -264,11 +264,16 @@ async def pop_item(self) -> TResponseInputItem | None: exists. """ + # Snapshot the current branch at call time so a concurrent + # switch_to_branch() cannot redirect this pop to a different branch once + # it has been dispatched to the worker thread. + branch_id = self._current_branch_id + def _pop_item_sync(): with self._locked_connection() as conn: while True: with closing(conn.cursor()) as cursor: - # Find the most recent item on the current branch. + # Find the most recent item on the snapshotted branch. cursor.execute( """ SELECT id, message_id, user_turn_number FROM message_structure @@ -276,7 +281,7 @@ def _pop_item_sync(): ORDER BY sequence_number DESC LIMIT 1 """, - (self.session_id, self._current_branch_id), + (self.session_id, branch_id), ) row = cursor.fetchone() if row is None: @@ -307,7 +312,7 @@ def _pop_item_sync(): SELECT COUNT(*) FROM message_structure WHERE session_id = ? AND branch_id = ? AND user_turn_number = ? """, - (self.session_id, self._current_branch_id, user_turn_number), + (self.session_id, branch_id, user_turn_number), ) if cursor.fetchone()[0] == 0: cursor.execute( @@ -315,7 +320,7 @@ def _pop_item_sync(): DELETE FROM turn_usage WHERE session_id = ? AND branch_id = ? AND user_turn_number = ? """, - (self.session_id, self._current_branch_id, user_turn_number), + (self.session_id, branch_id, user_turn_number), ) conn.commit() From f44364f22e2314fdd69dc71ff9d77e7edd3dd0c3 Mon Sep 17 00:00:00 2001 From: Aditya Jethani Date: Wed, 8 Jul 2026 13:56:54 +0530 Subject: [PATCH 4/5] clear_session: reset the branch pointer inside the locked operation Move the reset of _current_branch_id back to 'main' inside the locked clear so it is atomic with the row deletions and commit. This prevents any other locked operation from observing the session as cleared while the in-memory pointer still references a now-deleted branch. --- src/agents/extensions/memory/advanced_sqlite_session.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/src/agents/extensions/memory/advanced_sqlite_session.py b/src/agents/extensions/memory/advanced_sqlite_session.py index 0e81ee80cb..4558920628 100644 --- a/src/agents/extensions/memory/advanced_sqlite_session.py +++ b/src/agents/extensions/memory/advanced_sqlite_session.py @@ -366,10 +366,14 @@ def _clear_session_sync(): (self.session_id,), ) conn.commit() + # All branches were removed, so reset the in-memory pointer to + # 'main' while still holding the lock. Doing this inside the + # locked operation keeps the reset atomic with the clear, so no + # other locked operation observes the session as cleared while + # the pointer still references a deleted branch. + self._current_branch_id = "main" await asyncio.to_thread(_clear_session_sync) - # All branches were removed, so reset the in-memory pointer to 'main'. - self._current_branch_id = "main" async def store_run_usage(self, result: RunResult) -> None: """Store usage data for the current conversation turn. From a1f9a5a4aa94677c022943b51aaa2c4f6c12b323 Mon Sep 17 00:00:00 2001 From: Aditya Jethani Date: Wed, 8 Jul 2026 13:57:09 +0530 Subject: [PATCH 5/5] test: cover branch snapshot, locked clear reset, and shared copied-message pop Add regression tests for the review-requested behaviors: - popping a message copied (shared) into a branch keeps it while another branch still references it, and removes it once unreferenced - pop_item uses the call-time branch snapshot when a switch interleaves - clear_session resets the current branch pointer to 'main' --- .../memory/test_advanced_sqlite_session.py | 121 ++++++++++++++++++ 1 file changed, 121 insertions(+) diff --git a/tests/extensions/memory/test_advanced_sqlite_session.py b/tests/extensions/memory/test_advanced_sqlite_session.py index 48ffc66a83..cca2821b30 100644 --- a/tests/extensions/memory/test_advanced_sqlite_session.py +++ b/tests/extensions/memory/test_advanced_sqlite_session.py @@ -1927,3 +1927,124 @@ async def test_pop_item_respects_current_branch_and_keeps_shared_messages(): assert structure_message_ids <= message_ids finally: session.close() + + +async def test_pop_item_deletes_shared_copied_message_only_when_unreferenced(): + """Regression: popping a message that was copied into a branch (branches share + the underlying message row) must keep the message while another branch still + references it, and only remove it once no branch references it anymore. + """ + session = AdvancedSQLiteSession(session_id="pop_shared_copy_test", create_tables=True) + + main_items: list[TResponseInputItem] = [ + {"role": "user", "content": "u1"}, + {"role": "assistant", "content": "a1"}, + {"role": "user", "content": "u2"}, + {"role": "assistant", "content": "a2"}, + ] + + def message_count() -> int: + with session._locked_connection() as conn: + row = conn.execute( + f"SELECT COUNT(*) FROM {session.messages_table} WHERE session_id = ?", + (session.session_id,), + ).fetchone() + return cast(int, row[0]) + + try: + await session.add_items(main_items) + # Branch from turn 2 copies turn 1 (u1, a1) into branch_a as shared rows. + await session.create_branch_from_turn(2, "branch_a") + await session.switch_to_branch("branch_a") + await session.add_items([{"role": "user", "content": "branch-only"}]) + + assert message_count() == 5 # u1, a1, u2, a2, branch-only + + # Pop the branch-only item (not shared): its message row is removed. + assert await session.pop_item() == {"role": "user", "content": "branch-only"} + assert message_count() == 4 + + # Pop the copied, shared a1 and u1 off branch_a. They remain in the + # messages table because the main branch still references them. + assert await session.pop_item() == {"role": "assistant", "content": "a1"} + assert await session.pop_item() == {"role": "user", "content": "u1"} + assert message_count() == 4 + assert await session.get_items(branch_id="main") == main_items + assert await session.get_items(branch_id="branch_a") == [] + + # Now drain main: once no branch references u1/a1, the rows are removed. + await session.switch_to_branch("main") + for _ in range(len(main_items)): + await session.pop_item() + assert message_count() == 0 + assert await session.get_items() == [] + + # No orphaned structure rows at any point. + with session._locked_connection() as conn: + leftover = conn.execute( + "SELECT COUNT(*) FROM message_structure WHERE session_id = ?", + (session.session_id,), + ).fetchone()[0] + assert leftover == 0 + finally: + session.close() + + +async def test_pop_item_uses_branch_snapshot_when_branch_switches_concurrently(): + """Regression: pop_item snapshots the current branch at call time, so a branch + switch that interleaves after dispatch cannot redirect the pop to another branch. + """ + session = AdvancedSQLiteSession(session_id="pop_snapshot_test", create_tables=True) + + main_items: list[TResponseInputItem] = [ + {"role": "user", "content": "u1"}, + {"role": "assistant", "content": "a1"}, + {"role": "user", "content": "u2"}, + {"role": "assistant", "content": "a2"}, + ] + + try: + await session.add_items(main_items) + await session.create_branch_from_turn(2, "branch_a") + await session.switch_to_branch("branch_a") + await session.add_items([{"role": "user", "content": "branch-only"}]) + + # Start the pop while on branch_a, let it take its branch snapshot, then + # switch to main before it completes. + task = asyncio.ensure_future(session.pop_item()) + await asyncio.sleep(0) # let pop_item run up to its first await (snapshot taken) + await session.switch_to_branch("main") + popped = await task + + # The pop targeted branch_a (its state at call time), not main. + assert popped == {"role": "user", "content": "branch-only"} + assert await session.get_items(branch_id="main") == main_items + finally: + session.close() + + +async def test_clear_session_resets_current_branch_to_main(): + """Regression: clear_session must reset the in-memory branch pointer to 'main' + (inside the locked operation) since every branch was removed. + """ + session = AdvancedSQLiteSession(session_id="clear_branch_reset_test", create_tables=True) + + try: + await session.add_items( + [ + {"role": "user", "content": "u1"}, + {"role": "assistant", "content": "a1"}, + {"role": "user", "content": "u2"}, + {"role": "assistant", "content": "a2"}, + ] + ) + await session.create_branch_from_turn(2, "branch_a") + await session.switch_to_branch("branch_a") + assert session._current_branch_id == "branch_a" + + await session.clear_session() + + assert session._current_branch_id == "main" + assert await session.get_items() == [] + finally: + session.close()