diff --git a/src/agents/extensions/memory/advanced_sqlite_session.py b/src/agents/extensions/memory/advanced_sqlite_session.py index 704c85058e..a4af87c52f 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. 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 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. @@ -252,6 +272,132 @@ 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`. 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. + """ + + # 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 snapshotted branch. + cursor.execute( + """ + SELECT id, message_id, user_turn_number FROM message_structure + WHERE session_id = ? AND branch_id = ? + ORDER BY sequence_number DESC + LIMIT 1 + """, + (self.session_id, branch_id), + ) + row = cursor.fetchone() + if row is None: + return None + + structure_id, message_id, user_turn_number = 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) + + # 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, 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, branch_id, user_turn_number), + ) + + 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() + # 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. Bumping the + # generation invalidates any in-flight switch/create that + # captured the pre-clear generation. + self._generation += 1 + self._current_branch_id = "main" + + await asyncio.to_thread(_clear_session_sync) + async def store_run_usage(self, result: RunResult) -> None: """Store usage data for the current conversation turn. @@ -263,13 +409,56 @@ async def store_run_usage(self, result: RunResult) -> None: """ try: if result.context_wrapper.usage is not None: - # 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) + await self._update_turn_usage_internal( + 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. @@ -594,6 +783,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.""" @@ -634,9 +828,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 @@ -676,6 +872,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.""" @@ -696,8 +896,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. @@ -1286,17 +1489,48 @@ 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, + 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. + 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_anchor is not None: + with closing(conn.cursor()) as guard_cursor: + guard_cursor.execute( + """ + SELECT 1 FROM message_structure + WHERE session_id = ? AND branch_id = ? + AND user_turn_number = ? AND id = ? + """, + (self.session_id, target_branch, user_turn_number, turn_anchor), + ) + 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 output_details_json = None @@ -1328,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, diff --git a/tests/extensions/memory/test_advanced_sqlite_session.py b/tests/extensions/memory/test_advanced_sqlite_session.py index 9c543adc20..034cca0d35 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 @@ -1757,3 +1760,506 @@ 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_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). + """ + 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() + + +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() + + +@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) + + 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"}]) + + 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"} + assert await session.get_items(branch_id="main") == main_items + finally: + 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 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. + 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_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 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 + # 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_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. + """ + 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()