From 223fab9be07df16c3f837ebc3bcf97a69d9ec703 Mon Sep 17 00:00:00 2001
From: Aditya Jethani <aditya.ce21@sot.pdpu.ac.in>
Date: Wed, 8 Jul 2026 13:25:46 +0530
Subject: [PATCH 01/12] 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 <aditya.ce21@sot.pdpu.ac.in>
Date: Wed, 8 Jul 2026 13:42:41 +0530
Subject: [PATCH 02/12] 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 <aditya.ce21@sot.pdpu.ac.in>
Date: Wed, 8 Jul 2026 13:56:44 +0530
Subject: [PATCH 03/12] 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 <aditya.ce21@sot.pdpu.ac.in>
Date: Wed, 8 Jul 2026 13:56:54 +0530
Subject: [PATCH 04/12] 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 <aditya.ce21@sot.pdpu.ac.in>
Date: Wed, 8 Jul 2026 13:57:09 +0530
Subject: [PATCH 05/12] 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()

From 154535792e19604988acb4f0a03b94e59f0772b1 Mon Sep 17 00:00:00 2001
From: Aditya Jethani <aditya.ce21@sot.pdpu.ac.in>
Date: Thu, 9 Jul 2026 10:24:57 +0530
Subject: [PATCH 06/12] Guard the branch pointer against stale switch/create
 after clear_session

Add a generation counter bumped under the lock in clear_session, and route
switch_to_branch/create_branch_from_turn pointer updates through
_commit_branch_pointer, which captures the generation before its DB work and
sets _current_branch_id only if no clear_session has committed since. This
stops a switch or create that completes after a clear from resurrecting a
branch clear already deleted; the clear's reset to 'main' wins.
---
 .../memory/advanced_sqlite_session.py         | 47 +++++++++++++++++--
 1 file changed, 42 insertions(+), 5 deletions(-)

diff --git a/src/agents/extensions/memory/advanced_sqlite_session.py b/src/agents/extensions/memory/advanced_sqlite_session.py
index 4558920628..69ab05ba4f 100644
--- a/src/agents/extensions/memory/advanced_sqlite_session.py
+++ b/src/agents/extensions/memory/advanced_sqlite_session.py
@@ -48,8 +48,28 @@ def __init__(
         if create_tables:
             self._init_structure_tables()
         self._current_branch_id = "main"
+        # Bumped (under the connection lock) whenever clear_session() wipes the
+        # session. Branch-pointer and usage writes capture the generation before
+        # their DB work and skip their mutation if a clear has committed since,
+        # so a stale switch/create/store_run_usage cannot resurrect a branch or
+        # a turn that clear already removed.
+        self._generation = 0
         self._logger = logger or logging.getLogger(__name__)
 
+    def _commit_branch_pointer(self, branch_id: str, generation: int) -> bool:
+        """Set the current-branch pointer unless a clear has committed meanwhile.
+
+        Acquires the connection lock so the generation check and the assignment
+        are atomic with clear_session's reset. Returns True if the pointer was
+        updated, False if a clear_session committed after ``generation`` was
+        captured (in which case its reset to 'main' wins).
+        """
+        with self._lock:
+            if self._generation != generation:
+                return False
+            self._current_branch_id = branch_id
+            return True
+
     def _init_structure_tables(self):
         """Add structure and usage tracking tables.
 
@@ -370,7 +390,10 @@ def _clear_session_sync():
                 # '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.
+                # the pointer still references a deleted branch. Bumping the
+                # generation invalidates any in-flight switch/create/usage write
+                # that captured the pre-clear generation.
+                self._generation += 1
                 self._current_branch_id = "main"
 
         await asyncio.to_thread(_clear_session_sync)
@@ -717,6 +740,11 @@ async def create_branch_from_turn(
         """
         import time
 
+        # Capture the generation before any DB work so a clear that commits
+        # while this branch is being created cannot be overwritten by the
+        # pointer update below.
+        generation = self._generation
+
         # Validate the turn exists and contains a user message
         def _validate_turn():
             """Synchronous helper to validate turn exists and contains user message."""
@@ -757,9 +785,11 @@ def _validate_turn():
         # Copy messages before the branch point to the new branch
         await self._copy_messages_to_new_branch(branch_name, turn_number)
 
-        # Switch to new branch
+        # Switch to new branch under the lock; skipped if a clear_session has
+        # committed since `generation` was captured (its reset to 'main' wins),
+        # so we never point at a branch that clear removed.
         old_branch = self._current_branch_id
-        self._current_branch_id = branch_name
+        await asyncio.to_thread(self._commit_branch_pointer, branch_name, generation)
 
         self._logger.debug(
             f"Created branch '{branch_name}' from turn {turn_number} ('{turn_content}') in '{old_branch}'"  # noqa: E501
@@ -799,6 +829,10 @@ async def switch_to_branch(self, branch_id: str) -> None:
             ValueError: If the branch doesn't exist.
         """
 
+        # Capture the generation before validating so a clear that commits
+        # between validation and the pointer update is detected and skipped.
+        generation = self._generation
+
         # Validate branch exists
         def _validate_branch():
             """Synchronous helper to validate branch exists."""
@@ -819,8 +853,11 @@ def _validate_branch():
         await asyncio.to_thread(_validate_branch)
 
         old_branch = self._current_branch_id
-        self._current_branch_id = branch_id
-        self._logger.info(f"Switched from branch '{old_branch}' to '{branch_id}'")
+        # Update the pointer under the lock; a no-op if a clear_session has
+        # committed since `generation` was captured (its reset to 'main' wins).
+        switched = await asyncio.to_thread(self._commit_branch_pointer, branch_id, generation)
+        if switched:
+            self._logger.info(f"Switched from branch '{old_branch}' to '{branch_id}'")
 
     async def delete_branch(self, branch_id: str, force: bool = False) -> None:
         """Delete a branch and all its associated data.

From 9c5e44cfb2f104d66b199cf88c9865e891f13a5f Mon Sep 17 00:00:00 2001
From: Aditya Jethani <aditya.ce21@sot.pdpu.ac.in>
Date: Thu, 9 Jul 2026 10:25:08 +0530
Subject: [PATCH 07/12] Skip stale turn-usage writes after
 pop_item/clear_session

store_run_usage now captures the generation before reading the current turn
and passes it to _update_turn_usage_internal, which skips the write if a
clear_session has committed since (generation mismatch) or if the turn no
longer has any message_structure rows on the current branch (e.g. it was
removed by pop_item). This prevents reinserting usage for a turn that no
longer exists.
---
 .../memory/advanced_sqlite_session.py         | 32 +++++++++++++++++--
 1 file changed, 30 insertions(+), 2 deletions(-)

diff --git a/src/agents/extensions/memory/advanced_sqlite_session.py b/src/agents/extensions/memory/advanced_sqlite_session.py
index 69ab05ba4f..c27208d689 100644
--- a/src/agents/extensions/memory/advanced_sqlite_session.py
+++ b/src/agents/extensions/memory/advanced_sqlite_session.py
@@ -409,10 +409,15 @@ async def store_run_usage(self, result: RunResult) -> None:
         """
         try:
             if result.context_wrapper.usage is not None:
+                # Capture the generation before reading the turn so a clear that
+                # commits before the usage write is detected and skipped.
+                generation = self._generation
                 # Get the current turn number for this branch
                 current_turn = self._get_current_turn_number()
                 # Only update turn-level usage - session usage is aggregated on demand
-                await self._update_turn_usage_internal(current_turn, result.context_wrapper.usage)
+                await self._update_turn_usage_internal(
+                    current_turn, result.context_wrapper.usage, generation
+                )
         except Exception as e:
             self._logger.error(f"Failed to store usage for session {self.session_id}: {e}")
 
@@ -1446,17 +1451,40 @@ def _get_turn_usage_sync():
 
         return cast(list[dict[str, Any]] | dict[str, Any], result)
 
-    async def _update_turn_usage_internal(self, user_turn_number: int, usage_data: Usage) -> None:
+    async def _update_turn_usage_internal(
+        self, user_turn_number: int, usage_data: Usage, generation: int | None = None
+    ) -> None:
         """Internal method to update usage for a specific turn with full JSON details.
 
         Args:
             user_turn_number: The turn number to update usage for.
             usage_data: The usage data to store.
+            generation: The generation captured before the turn was read. When
+                provided, the write is skipped if a clear_session has committed
+                since (generation mismatch) or if the turn no longer exists on
+                the current branch (e.g. it was removed by pop_item), so stale
+                usage is never recorded for a turn that no longer exists.
         """
 
         def _update_sync():
             """Synchronous helper to update turn usage data."""
             with self._locked_connection() as conn:
+                if generation is not None:
+                    if self._generation != generation:
+                        # A clear_session committed after the turn was read.
+                        return
+                    with closing(conn.cursor()) as guard_cursor:
+                        guard_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 guard_cursor.fetchone()[0] == 0:
+                            # The turn was removed (e.g. by pop_item) after it was
+                            # read; do not resurrect usage for a nonexistent turn.
+                            return
                 # Serialize token details as JSON
                 input_details_json = None
                 output_details_json = None

From a9d6a30cf95ee5cbdf00f6b890a62d4310fe285f Mon Sep 17 00:00:00 2001
From: Aditya Jethani <aditya.ce21@sot.pdpu.ac.in>
Date: Thu, 9 Jul 2026 10:25:21 +0530
Subject: [PATCH 08/12] test: barrier-based interleaving tests for
 clear/switch/pop ordering

Replace the sleep(0) branch-snapshot test with a deterministic barrier that
holds the pop worker after its snapshot while a switch completes. Add
interleaving regression tests (via a gated asyncio.to_thread) proving a stale
switch or create after clear_session leaves the pointer on 'main', and a
store_run_usage racing with pop_item does not reinsert usage for the removed
turn.
---
 .../memory/test_advanced_sqlite_session.py    | 148 +++++++++++++++++-
 1 file changed, 142 insertions(+), 6 deletions(-)

diff --git a/tests/extensions/memory/test_advanced_sqlite_session.py b/tests/extensions/memory/test_advanced_sqlite_session.py
index cca2821b30..82901b827a 100644
--- a/tests/extensions/memory/test_advanced_sqlite_session.py
+++ b/tests/extensions/memory/test_advanced_sqlite_session.py
@@ -1,10 +1,13 @@
 """Tests for AdvancedSQLiteSession functionality."""
 
 import asyncio
+import contextlib
 import json
 import tempfile
+import threading
 from pathlib import Path
 from typing import Any, cast
+from unittest.mock import patch
 
 import pytest
 
@@ -1990,9 +1993,41 @@ def message_count() -> int:
         session.close()
 
 
+@contextlib.contextmanager
+def _gate_worker(target_name: str):
+    """Deterministically pause a session worker to control interleaving.
+
+    Patches ``asyncio.to_thread`` in the session module so the first dispatch of
+    a worker whose ``__name__`` equals ``target_name`` signals ``started`` and
+    blocks on ``release`` before running. The pause happens before the worker
+    acquires the connection lock, so other operations can run to completion
+    while it is held. Yields ``(started, release)`` threading events.
+    """
+    started = threading.Event()
+    release = threading.Event()
+    real_to_thread = asyncio.to_thread
+    state = {"gated": False}
+
+    async def gated(func, /, *args, **kwargs):
+        if not state["gated"] and getattr(func, "__name__", "") == target_name:
+            state["gated"] = True
+            started.set()
+            await real_to_thread(release.wait)
+        return await real_to_thread(func, *args, **kwargs)
+
+    with patch(
+        "agents.extensions.memory.advanced_sqlite_session.asyncio.to_thread",
+        gated,
+    ):
+        yield started, real_to_thread, release
+
+
 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.
+
+    Uses a barrier (not sleep) to prove the ordering: the pop worker is held after
+    its branch snapshot is taken while a full switch_to_branch("main") completes.
     """
     session = AdvancedSQLiteSession(session_id="pop_snapshot_test", create_tables=True)
 
@@ -2009,12 +2044,15 @@ async def test_pop_item_uses_branch_snapshot_when_branch_switches_concurrently()
         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
+        with _gate_worker("_pop_item_sync") as (started, real_to_thread, release):
+            # pop_item snapshots _current_branch_id ("branch_a") synchronously,
+            # then dispatches its worker, which parks at the barrier.
+            task = asyncio.ensure_future(session.pop_item())
+            await real_to_thread(started.wait)
+            # Switch to main completes fully while the pop worker is parked.
+            await session.switch_to_branch("main")
+            release.set()
+            popped = await task
 
         # The pop targeted branch_a (its state at call time), not main.
         assert popped == {"role": "user", "content": "branch-only"}
@@ -2023,6 +2061,104 @@ async def test_pop_item_uses_branch_snapshot_when_branch_switches_concurrently()
         session.close()
 
 
+async def test_stale_switch_after_clear_does_not_repoint_to_deleted_branch():
+    """A switch_to_branch that commits its pointer after clear_session must not
+    resurrect the deleted branch; the generation guard makes it a no-op.
+    """
+    session = AdvancedSQLiteSession(session_id="stale_switch_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("main")
+        assert session._current_branch_id == "main"
+
+        with _gate_worker("_commit_branch_pointer") as (started, real_to_thread, release):
+            # switch validates branch_a and captures the generation, then parks
+            # right before committing the pointer.
+            task = asyncio.ensure_future(session.switch_to_branch("branch_a"))
+            await real_to_thread(started.wait)
+            # A full clear commits: it bumps the generation and resets to main.
+            await session.clear_session()
+            release.set()
+            await task
+
+        # The stale switch saw a newer generation and left the pointer on main.
+        assert session._current_branch_id == "main"
+        assert await session.get_items() == []
+    finally:
+        session.close()
+
+
+async def test_stale_create_branch_after_clear_does_not_repoint():
+    """A create_branch_from_turn that commits its pointer after clear_session must
+    not point at the branch clear removed.
+    """
+    session = AdvancedSQLiteSession(session_id="stale_create_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"},
+            ]
+        )
+
+        with _gate_worker("_commit_branch_pointer") as (started, real_to_thread, release):
+            task = asyncio.ensure_future(session.create_branch_from_turn(2, "branch_b"))
+            await real_to_thread(started.wait)
+            await session.clear_session()
+            release.set()
+            await task
+
+        # clear won: the pointer stays on main, not the wiped branch_b.
+        assert session._current_branch_id == "main"
+        assert await session.get_items() == []
+    finally:
+        session.close()
+
+
+async def test_stale_store_run_usage_skipped_when_turn_removed_by_pop(usage_data: Usage):
+    """A store_run_usage that reads a turn and then races with pop_item removing
+    that turn must not reinsert usage for the now-nonexistent turn.
+    """
+    session = AdvancedSQLiteSession(session_id="stale_usage_test", create_tables=True)
+
+    try:
+        await session.add_items(
+            [
+                {"role": "user", "content": "u1"},
+                {"role": "assistant", "content": "a1"},
+            ]
+        )
+        result = create_mock_run_result(usage_data)
+
+        with _gate_worker("_update_sync") as (started, real_to_thread, release):
+            # store_run_usage reads current_turn (1) and captures the generation,
+            # then parks before writing turn_usage.
+            task = asyncio.ensure_future(session.store_run_usage(result))
+            await real_to_thread(started.wait)
+            # Pop both items of turn 1 so the turn no longer exists.
+            await session.pop_item()
+            await session.pop_item()
+            release.set()
+            await task
+
+        # The stale usage write was skipped: no row for the removed turn.
+        assert _count_rows(session, "turn_usage") == 0
+    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.

From 3c3849d0ccf3532c267ab23de179bd5e4ee63125 Mon Sep 17 00:00:00 2001
From: Aditya Jethani <aditya.ce21@sot.pdpu.ac.in>
Date: Thu, 9 Jul 2026 10:46:40 +0530
Subject: [PATCH 09/12] Invalidate stale usage writes with a turn-usage version
 counter

The existence-only usage guard was vulnerable to turn-id reuse (ABA): if a
store_run_usage captured turn N, then pop_item removed that turn and a new
turn reused id N before the write committed, the existence check passed and
the old run's usage was recorded against the new turn.

Add a _turn_usage_version counter bumped under the lock whenever a turn is
removed (clear_session, delete_branch, or a pop_item that empties a turn).
store_run_usage captures it before reading the turn and the write is skipped
if it changed, so usage is never recorded against a removed turn even when
its numeric id is reused. The turn-existence check remains as a backstop.
---
 .../memory/advanced_sqlite_session.py         | 65 +++++++++++++------
 1 file changed, 44 insertions(+), 21 deletions(-)

diff --git a/src/agents/extensions/memory/advanced_sqlite_session.py b/src/agents/extensions/memory/advanced_sqlite_session.py
index c27208d689..a358f0cd48 100644
--- a/src/agents/extensions/memory/advanced_sqlite_session.py
+++ b/src/agents/extensions/memory/advanced_sqlite_session.py
@@ -49,11 +49,17 @@ def __init__(
             self._init_structure_tables()
         self._current_branch_id = "main"
         # Bumped (under the connection lock) whenever clear_session() wipes the
-        # session. Branch-pointer and usage writes capture the generation before
-        # their DB work and skip their mutation if a clear has committed since,
-        # so a stale switch/create/store_run_usage cannot resurrect a branch or
-        # a turn that clear already removed.
+        # session. switch_to_branch / create_branch_from_turn capture the
+        # generation before their DB work and only update the branch pointer if
+        # no clear has committed since, so a stale switch/create cannot resurrect
+        # a branch that clear already removed.
         self._generation = 0
+        # Bumped (under the connection lock) whenever a turn is removed
+        # (clear_session, delete_branch, or a pop_item that empties a turn).
+        # store_run_usage captures this before reading the turn and the write is
+        # skipped if it changed, so usage is never recorded against a turn that
+        # was removed — even if a new turn later reuses the same numeric id.
+        self._turn_usage_version = 0
         self._logger = logger or logging.getLogger(__name__)
 
     def _commit_branch_pointer(self, branch_id: str, generation: int) -> bool:
@@ -342,6 +348,10 @@ def _pop_item_sync():
                                     """,
                                     (self.session_id, branch_id, user_turn_number),
                                 )
+                                # The turn is gone; invalidate any in-flight usage
+                                # write for it, even if a new turn later reuses
+                                # this numeric id.
+                                self._turn_usage_version += 1
 
                         conn.commit()
 
@@ -391,9 +401,11 @@ def _clear_session_sync():
                 # 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. Bumping the
-                # generation invalidates any in-flight switch/create/usage write
-                # that captured the pre-clear generation.
+                # generation invalidates any in-flight switch/create that
+                # captured the pre-clear generation; bumping the turn-usage
+                # version invalidates any in-flight usage write.
                 self._generation += 1
+                self._turn_usage_version += 1
                 self._current_branch_id = "main"
 
         await asyncio.to_thread(_clear_session_sync)
@@ -409,14 +421,17 @@ async def store_run_usage(self, result: RunResult) -> None:
         """
         try:
             if result.context_wrapper.usage is not None:
-                # Capture the generation before reading the turn so a clear that
-                # commits before the usage write is detected and skipped.
-                generation = self._generation
+                # Capture the turn-usage version before reading the turn. If the
+                # turn is removed (by clear/pop/delete) before the write commits,
+                # the version changes and the write is skipped, so usage is never
+                # recorded against a turn that no longer exists — even when a new
+                # turn later reuses the same numeric id.
+                turn_usage_version = self._turn_usage_version
                 # Get the current turn number for this branch
                 current_turn = self._get_current_turn_number()
                 # Only update turn-level usage - session usage is aggregated on demand
                 await self._update_turn_usage_internal(
-                    current_turn, result.context_wrapper.usage, generation
+                    current_turn, result.context_wrapper.usage, turn_usage_version
                 )
         except Exception as e:
             self._logger.error(f"Failed to store usage for session {self.session_id}: {e}")
@@ -934,6 +949,10 @@ def _delete_sync():
 
                     orphaned_messages_deleted = self._cleanup_orphaned_messages_sync(conn)
 
+                    # Turns were removed; invalidate any in-flight usage write.
+                    if structure_deleted:
+                        self._turn_usage_version += 1
+
                     conn.commit()
 
                     return usage_deleted, structure_deleted, orphaned_messages_deleted
@@ -1452,26 +1471,31 @@ def _get_turn_usage_sync():
         return cast(list[dict[str, Any]] | dict[str, Any], result)
 
     async def _update_turn_usage_internal(
-        self, user_turn_number: int, usage_data: Usage, generation: int | None = None
+        self,
+        user_turn_number: int,
+        usage_data: Usage,
+        turn_usage_version: int | None = None,
     ) -> None:
         """Internal method to update usage for a specific turn with full JSON details.
 
         Args:
             user_turn_number: The turn number to update usage for.
             usage_data: The usage data to store.
-            generation: The generation captured before the turn was read. When
-                provided, the write is skipped if a clear_session has committed
-                since (generation mismatch) or if the turn no longer exists on
-                the current branch (e.g. it was removed by pop_item), so stale
-                usage is never recorded for a turn that no longer exists.
+            turn_usage_version: The value of ``self._turn_usage_version`` captured
+                before the turn was read. When provided, the write is skipped if
+                the version has changed since — i.e. a clear_session, pop_item, or
+                delete_branch removed a turn — so usage is never recorded against a
+                turn that was removed, even when a new turn reuses the same numeric
+                id. The turn-existence check below is a defensive backstop.
         """
 
         def _update_sync():
             """Synchronous helper to update turn usage data."""
             with self._locked_connection() as conn:
-                if generation is not None:
-                    if self._generation != generation:
-                        # A clear_session committed after the turn was read.
+                if turn_usage_version is not None:
+                    if self._turn_usage_version != turn_usage_version:
+                        # A turn was removed after this turn number was read; the
+                        # captured turn is stale (its id may have been reused).
                         return
                     with closing(conn.cursor()) as guard_cursor:
                         guard_cursor.execute(
@@ -1482,8 +1506,7 @@ def _update_sync():
                             (self.session_id, self._current_branch_id, user_turn_number),
                         )
                         if guard_cursor.fetchone()[0] == 0:
-                            # The turn was removed (e.g. by pop_item) after it was
-                            # read; do not resurrect usage for a nonexistent turn.
+                            # The turn does not exist on the current branch.
                             return
                 # Serialize token details as JSON
                 input_details_json = None

From 38becd39370fc332875183d35f46a86e43a30f76 Mon Sep 17 00:00:00 2001
From: Aditya Jethani <aditya.ce21@sot.pdpu.ac.in>
Date: Thu, 9 Jul 2026 10:46:40 +0530
Subject: [PATCH 10/12] test: ABA regression for usage recorded against a
 reused turn number

Pop a turn while a store_run_usage is parked, add a new turn that reuses the
same numeric id, then let the write resume; it must be skipped. Fails with the
existence-only guard and passes with the turn-usage version counter.
---
 .../memory/test_advanced_sqlite_session.py    | 43 ++++++++++++++++++-
 1 file changed, 41 insertions(+), 2 deletions(-)

diff --git a/tests/extensions/memory/test_advanced_sqlite_session.py b/tests/extensions/memory/test_advanced_sqlite_session.py
index 82901b827a..01735f4e0c 100644
--- a/tests/extensions/memory/test_advanced_sqlite_session.py
+++ b/tests/extensions/memory/test_advanced_sqlite_session.py
@@ -2143,8 +2143,8 @@ async def test_stale_store_run_usage_skipped_when_turn_removed_by_pop(usage_data
         result = create_mock_run_result(usage_data)
 
         with _gate_worker("_update_sync") as (started, real_to_thread, release):
-            # store_run_usage reads current_turn (1) and captures the generation,
-            # then parks before writing turn_usage.
+            # store_run_usage reads current_turn (1) and captures the turn-usage
+            # version, then parks before writing turn_usage.
             task = asyncio.ensure_future(session.store_run_usage(result))
             await real_to_thread(started.wait)
             # Pop both items of turn 1 so the turn no longer exists.
@@ -2159,6 +2159,45 @@ async def test_stale_store_run_usage_skipped_when_turn_removed_by_pop(usage_data
         session.close()
 
 
+async def test_stale_store_run_usage_not_recorded_against_reused_turn_number(
+    usage_data: Usage,
+):
+    """A store_run_usage that read turn N must not record its usage when that turn
+    is popped and a *new* turn later reuses the same numeric id (the ABA case).
+
+    An existence-only guard would pass here because turn 1 exists again; the
+    turn-usage version counter invalidates the stale write.
+    """
+    session = AdvancedSQLiteSession(session_id="stale_usage_aba_test", create_tables=True)
+
+    try:
+        await session.add_items(
+            [
+                {"role": "user", "content": "u1"},
+                {"role": "assistant", "content": "a1"},
+            ]
+        )
+        result = create_mock_run_result(usage_data)
+
+        with _gate_worker("_update_sync") as (started, real_to_thread, release):
+            # Reads current_turn (1), captures the version, parks before writing.
+            task = asyncio.ensure_future(session.store_run_usage(result))
+            await real_to_thread(started.wait)
+            # Remove turn 1 entirely, then create a brand-new turn that reuses the
+            # numeric id 1.
+            await session.pop_item()
+            await session.pop_item()
+            await session.add_items([{"role": "user", "content": "fresh turn"}])
+            release.set()
+            await task
+
+        # The new turn 1 must not carry the previous run's usage.
+        assert _count_rows(session, "turn_usage") == 0
+        assert not await session.get_turn_usage(1)
+    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.

From d9f48a06dcfd5b5f5b4221b558c32eb0eb5156c9 Mon Sep 17 00:00:00 2001
From: Aditya Jethani <aditya.ce21@sot.pdpu.ac.in>
Date: Thu, 9 Jul 2026 11:21:55 +0530
Subject: [PATCH 11/12] Scope stale-usage invalidation to the captured turn via
 a row anchor

The session-wide _turn_usage_version was too coarse: an unrelated turn
removal (e.g. delete_branch on a non-current branch) bumped it and dropped a
legitimate usage write for a turn that still existed.

Replace it with a per-turn anchor: store_run_usage captures the id of the
turn's first message_structure row (ids are monotonic and never reused) via
_capture_current_turn, and the write is skipped only if that exact row no
longer exists for the captured branch/turn. This still catches turn-id reuse
(the anchor row is gone) but no longer over-invalidates on unrelated removals.
The write also targets the captured branch. The guard is backed by the
existing (session_id, branch_id, user_turn_number) index.
---
 .../memory/advanced_sqlite_session.py         | 111 +++++++++++-------
 1 file changed, 67 insertions(+), 44 deletions(-)

diff --git a/src/agents/extensions/memory/advanced_sqlite_session.py b/src/agents/extensions/memory/advanced_sqlite_session.py
index a358f0cd48..a4af87c52f 100644
--- a/src/agents/extensions/memory/advanced_sqlite_session.py
+++ b/src/agents/extensions/memory/advanced_sqlite_session.py
@@ -54,12 +54,6 @@ def __init__(
         # no clear has committed since, so a stale switch/create cannot resurrect
         # a branch that clear already removed.
         self._generation = 0
-        # Bumped (under the connection lock) whenever a turn is removed
-        # (clear_session, delete_branch, or a pop_item that empties a turn).
-        # store_run_usage captures this before reading the turn and the write is
-        # skipped if it changed, so usage is never recorded against a turn that
-        # was removed — even if a new turn later reuses the same numeric id.
-        self._turn_usage_version = 0
         self._logger = logger or logging.getLogger(__name__)
 
     def _commit_branch_pointer(self, branch_id: str, generation: int) -> bool:
@@ -348,10 +342,6 @@ def _pop_item_sync():
                                     """,
                                     (self.session_id, branch_id, user_turn_number),
                                 )
-                                # The turn is gone; invalidate any in-flight usage
-                                # write for it, even if a new turn later reuses
-                                # this numeric id.
-                                self._turn_usage_version += 1
 
                         conn.commit()
 
@@ -402,10 +392,8 @@ def _clear_session_sync():
                 # other locked operation observes the session as cleared while
                 # the pointer still references a deleted branch. Bumping the
                 # generation invalidates any in-flight switch/create that
-                # captured the pre-clear generation; bumping the turn-usage
-                # version invalidates any in-flight usage write.
+                # captured the pre-clear generation.
                 self._generation += 1
-                self._turn_usage_version += 1
                 self._current_branch_id = "main"
 
         await asyncio.to_thread(_clear_session_sync)
@@ -421,21 +409,56 @@ async def store_run_usage(self, result: RunResult) -> None:
         """
         try:
             if result.context_wrapper.usage is not None:
-                # Capture the turn-usage version before reading the turn. If the
-                # turn is removed (by clear/pop/delete) before the write commits,
-                # the version changes and the write is skipped, so usage is never
-                # recorded against a turn that no longer exists — even when a new
-                # turn later reuses the same numeric id.
-                turn_usage_version = self._turn_usage_version
-                # Get the current turn number for this branch
-                current_turn = self._get_current_turn_number()
+                # Capture the current turn together with an anchor that pins the
+                # exact turn incarnation: the id of its first message_structure
+                # row (ids are monotonic and never reused). If that turn is
+                # removed before the write commits — even if a new turn later
+                # reuses the same numeric id — the anchor row is gone and the
+                # write is skipped. The anchor is scoped to this branch/turn, so
+                # unrelated removals (e.g. delete_branch on another branch) do
+                # not drop this write.
+                current_turn, branch_id, turn_anchor = self._capture_current_turn()
                 # Only update turn-level usage - session usage is aggregated on demand
                 await self._update_turn_usage_internal(
-                    current_turn, result.context_wrapper.usage, turn_usage_version
+                    current_turn,
+                    result.context_wrapper.usage,
+                    branch_id=branch_id,
+                    turn_anchor=turn_anchor,
                 )
         except Exception as e:
             self._logger.error(f"Failed to store usage for session {self.session_id}: {e}")
 
+    def _capture_current_turn(self) -> tuple[int, str, int | None]:
+        """Return (current_turn, branch_id, turn_anchor) in one locked read.
+
+        ``turn_anchor`` is the smallest ``message_structure.id`` of the current
+        turn on the current branch (``None`` if the turn has no rows). Because
+        ids are monotonic and never reused, it uniquely identifies this turn
+        incarnation, so a later pop+recreate that reuses the numeric turn id
+        yields a different anchor.
+        """
+        with self._locked_connection() as conn:
+            with closing(conn.cursor()) as cursor:
+                branch_id = self._current_branch_id
+                cursor.execute(
+                    """
+                    SELECT COALESCE(MAX(user_turn_number), 0)
+                    FROM message_structure
+                    WHERE session_id = ? AND branch_id = ?
+                    """,
+                    (self.session_id, branch_id),
+                )
+                current_turn = cursor.fetchone()[0]
+                cursor.execute(
+                    """
+                    SELECT MIN(id) FROM message_structure
+                    WHERE session_id = ? AND branch_id = ? AND user_turn_number = ?
+                    """,
+                    (self.session_id, branch_id, current_turn),
+                )
+                turn_anchor = cursor.fetchone()[0]
+                return current_turn, branch_id, turn_anchor
+
     def _get_next_turn_number(self, branch_id: str) -> int:
         """Get the next turn number for a specific branch.
 
@@ -949,10 +972,6 @@ def _delete_sync():
 
                     orphaned_messages_deleted = self._cleanup_orphaned_messages_sync(conn)
 
-                    # Turns were removed; invalidate any in-flight usage write.
-                    if structure_deleted:
-                        self._turn_usage_version += 1
-
                     conn.commit()
 
                     return usage_deleted, structure_deleted, orphaned_messages_deleted
@@ -1474,39 +1493,43 @@ async def _update_turn_usage_internal(
         self,
         user_turn_number: int,
         usage_data: Usage,
-        turn_usage_version: int | None = None,
+        branch_id: str | None = None,
+        turn_anchor: int | None = None,
     ) -> None:
         """Internal method to update usage for a specific turn with full JSON details.
 
         Args:
             user_turn_number: The turn number to update usage for.
             usage_data: The usage data to store.
-            turn_usage_version: The value of ``self._turn_usage_version`` captured
-                before the turn was read. When provided, the write is skipped if
-                the version has changed since — i.e. a clear_session, pop_item, or
-                delete_branch removed a turn — so usage is never recorded against a
-                turn that was removed, even when a new turn reuses the same numeric
-                id. The turn-existence check below is a defensive backstop.
+            branch_id: The branch the turn was read from. Defaults to the current
+                branch when not provided.
+            turn_anchor: The id of the turn's first ``message_structure`` row,
+                captured when the turn was read. When provided, the write is
+                skipped unless that exact row still exists for the given
+                branch/turn, so usage is never recorded against a turn that was
+                removed — even if a new turn reused the same numeric id. Because
+                the check is scoped to this branch/turn, unrelated removals (e.g.
+                delete_branch on another branch) do not drop this write.
         """
 
+        target_branch = branch_id if branch_id is not None else self._current_branch_id
+
         def _update_sync():
             """Synchronous helper to update turn usage data."""
             with self._locked_connection() as conn:
-                if turn_usage_version is not None:
-                    if self._turn_usage_version != turn_usage_version:
-                        # A turn was removed after this turn number was read; the
-                        # captured turn is stale (its id may have been reused).
-                        return
+                if turn_anchor is not None:
                     with closing(conn.cursor()) as guard_cursor:
                         guard_cursor.execute(
                             """
-                            SELECT COUNT(*) FROM message_structure
-                            WHERE session_id = ? AND branch_id = ? AND user_turn_number = ?
+                            SELECT 1 FROM message_structure
+                            WHERE session_id = ? AND branch_id = ?
+                            AND user_turn_number = ? AND id = ?
                             """,
-                            (self.session_id, self._current_branch_id, user_turn_number),
+                            (self.session_id, target_branch, user_turn_number, turn_anchor),
                         )
-                        if guard_cursor.fetchone()[0] == 0:
-                            # The turn does not exist on the current branch.
+                        if guard_cursor.fetchone() is None:
+                            # The exact turn incarnation is gone (removed, or its
+                            # numeric id reused by a new turn); skip the stale write.
                             return
                 # Serialize token details as JSON
                 input_details_json = None
@@ -1539,7 +1562,7 @@ def _update_sync():
                     """,  # noqa: E501
                         (
                             self.session_id,
-                            self._current_branch_id,
+                            target_branch,
                             user_turn_number,
                             usage_data.requests or 0,
                             usage_data.input_tokens or 0,

From 7353a4aea21d4511c5ea7fb610bd7830b5ba526d Mon Sep 17 00:00:00 2001
From: Aditya Jethani <aditya.ce21@sot.pdpu.ac.in>
Date: Thu, 9 Jul 2026 11:21:55 +0530
Subject: [PATCH 12/12] test: usage write survives unrelated branch deletion
 (invalidation scoping)

Add a regression test asserting a parked store_run_usage still records its
turn's usage when an unrelated delete_branch runs meanwhile; it fails with the
session-wide counter and passes with the per-turn anchor. Refresh comments to
reference the anchor.
---
 .../memory/test_advanced_sqlite_session.py    | 42 ++++++++++++++++++-
 1 file changed, 41 insertions(+), 1 deletion(-)

diff --git a/tests/extensions/memory/test_advanced_sqlite_session.py b/tests/extensions/memory/test_advanced_sqlite_session.py
index 01735f4e0c..034cca0d35 100644
--- a/tests/extensions/memory/test_advanced_sqlite_session.py
+++ b/tests/extensions/memory/test_advanced_sqlite_session.py
@@ -2180,7 +2180,7 @@ async def test_stale_store_run_usage_not_recorded_against_reused_turn_number(
         result = create_mock_run_result(usage_data)
 
         with _gate_worker("_update_sync") as (started, real_to_thread, release):
-            # Reads current_turn (1), captures the version, parks before writing.
+            # Reads current_turn (1), captures the turn anchor, parks before write.
             task = asyncio.ensure_future(session.store_run_usage(result))
             await real_to_thread(started.wait)
             # Remove turn 1 entirely, then create a brand-new turn that reuses the
@@ -2198,6 +2198,46 @@ async def test_stale_store_run_usage_not_recorded_against_reused_turn_number(
         session.close()
 
 
+async def test_store_run_usage_survives_unrelated_branch_deletion(usage_data: Usage):
+    """A store_run_usage in flight must not be dropped when an unrelated turn is
+    removed (e.g. delete_branch on a non-current branch). The invalidation is
+    scoped to the captured branch/turn, so the write still lands.
+    """
+    session = AdvancedSQLiteSession(session_id="usage_scope_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"},
+            ]
+        )
+        # A separate branch that shares turn 1's messages; deleting it must not
+        # affect usage captured for the current (main) branch.
+        await session.create_branch_from_turn(2, "side_branch")
+        await session.switch_to_branch("main")
+        result = create_mock_run_result(usage_data)
+
+        with _gate_worker("_update_sync") as (started, real_to_thread, release):
+            # Captures main/turn 2 and its anchor, then parks before the write.
+            task = asyncio.ensure_future(session.store_run_usage(result))
+            await real_to_thread(started.wait)
+            # Delete an unrelated branch while the usage write is parked.
+            await session.delete_branch("side_branch")
+            release.set()
+            await task
+
+        # The write landed: main's turn 2 usage is recorded despite the deletion.
+        assert _count_rows(session, "turn_usage") == 1
+        turn_2_usage = await session.get_turn_usage(2)
+        assert isinstance(turn_2_usage, dict)
+        assert turn_2_usage["total_tokens"] == usage_data.total_tokens
+    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.