fix: AdvancedSQLiteSession clear_session and pop_item metadata leaks#3755
fix: AdvancedSQLiteSession clear_session and pop_item metadata leaks#3755okaditya84 wants to merge 13 commits into
Conversation
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.
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 223fab9be0
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
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.
|
@seratch would you be able to take a look when you have a moment? 🙏 This fixes the two remaining |
seratch
left a comment
seratch
left a comment
There was a problem hiding this comment.
Thanks for the contribution. The underlying clear_session() and branch-aware pop_item() bugs are valid, and overriding these methods in AdvancedSQLiteSession is the right implementation layer.
Before we can merge this, please make three focused changes: remove the matching turn_usage row when popping the final structure row for a branch turn; snapshot the branch before dispatching pop_item() to the worker and reset the branch to main inside the locked clear operation after commit; and extend the tests to actually pop a shared copied message and cover the branch-state interleaving or cancellation case.
These changes are needed to preserve metadata consistency and prevent stale operations from affecting surviving branch state. Once covered by focused regression tests, this should be ready for another review.
seratch
changed the title
Thanks for the review. I will do the respective changes and push right away. |
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.
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.
…ssage 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'
|
Thanks for the detailed review @seratch! All three changes are done and pushed as focused commits:
Verified locally: |
seratch
left a comment
seratch
left a comment
There was a problem hiding this comment.
Thanks for the update. The original findings around turn-usage cleanup, branch snapshotting, and shared-message ownership are now addressed.
One concurrency invariant still needs to be fixed. An earlier switch_to_branch() or create_branch_from_turn() can complete after clear_session() and overwrite _current_branch_id with a branch that clear already deleted. Likewise, store_run_usage() can read a turn, race with pop/clear, and then reinsert usage for a nonexistent turn.
Please enforce these orderings: once clear commits, stale switch/create completions must not change the branch pointer; once pop/clear removes a turn, stale usage writes must be skipped. Add controlled event/future interleaving tests for these cases, and replace the sleep(0) test with a barrier that proves the intended worker ordering. The rest of the patch and CI look good.
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.
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.
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.
|
Thanks for the thorough review @seratch — both concurrency invariants are now enforced and covered by deterministic interleaving tests. 1. Branch pointer can't be resurrected after clear ( 2. Stale usage writes are skipped after pop/clear ( 3. Interleaving tests + barrier (
Each fails when the corresponding guard is neutered and passes with it. Verified locally: |
|
Done @seratch — both invariants enforced:
format/lint/typecheck clean; full suite 4910 passed. |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 39b711a5bd
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| """, | ||
| (self.session_id, self._current_branch_id, user_turn_number), | ||
| ) | ||
| if guard_cursor.fetchone()[0] == 0: |
There was a problem hiding this comment.
If a store_run_usage() call has captured user_turn_number == 1, then pop_item() removes that whole turn and another add_items() creates a new turn before the paused usage write resumes, this existence-only guard passes because turn 1 exists again. The subsequent insert records the old run's usage against the new turn, corrupting usage analytics; popping a turn needs to invalidate pending usage writes more strongly than checking whether the same numeric turn currently exists.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Good catch on the ABA race — fixed in 3c3849d. The existence-only guard is replaced by 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 _update_turn_usage_internal skips the write if it changed, so a paused usage write is invalidated when its turn is popped — even if a new turn later reuses the same numeric id. (The turn-existence check stays as a backstop.) Added test_stale_store_run_usage_not_recorded_against_reused_turn_number, which reproduces exactly this pop-then-reuse sequence; it fails with the old existence-only guard and passes with the version counter. Full suite: 4911 passed.
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.
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.
|
Addressed the Codex ABA finding ( |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 38becd3937
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| 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 |
There was a problem hiding this comment.
When store_run_usage() is in flight and an unrelated operation removes any other turn, such as delete_branch() deleting a non-current branch, the session-wide _turn_usage_version changes and this check drops the usage write even though the captured turn on the current branch still exists. That loses usage analytics for completed runs during normal branch cleanup races; the invalidation needs to be scoped to the captured branch/turn rather than every turn removal in the session.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Good point — the session-wide counter was over-invalidating. Fixed in d9f48a06 by scoping the guard to the captured turn: store_run_usage now captures a per-turn anchor (the id of the turn's first message_structure row — ids are monotonic and never reused), and the write is skipped only if that exact row no longer exists for the captured branch/turn. So an unrelated delete_branch on a non-current branch no longer drops the write, while turn-id reuse is still caught (the anchor row is gone). The write also targets the captured branch, and the check rides the existing (session_id, branch_id, user_turn_number) index.
Added test_store_run_usage_survives_unrelated_branch_deletion (7353a4ae) — it fails with the old session-wide counter and passes with the anchor. Full suite: 4912 passed / 6 skipped, format/lint/typecheck clean.
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.
…oping) 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.
Summary
AdvancedSQLiteSessionmaintains the auxiliarymessage_structureandturn_usagetables in itsadd_itemsoverride, but it inheritedclear_sessionandpop_itemunchanged fromSQLiteSession, which only touch themessagesandsessionstables.Those aux tables declare
ON DELETE CASCADEforeign keys, but SQLite does not enforce foreign keys unlessPRAGMA foreign_keys=ONis set, and the SDK never enables it (only WAL is configured). So the cascade never fires and the rows leak. Concretely:clear_session()left everymessage_structureandturn_usagerow behind.list_branches,get_session_usage, andget_turn_usagekept reporting stale data, and because_insert_structure_metadataseeds fromMAX(sequence_number)/MAX(user_turn_number), sequence and turn numbering stayed permanently offset for any items added after clearing (e.g. the next item gotsequence_number=3, user_turn_number=2instead of1, 1).pop_item()deleted the globally highest-idmessage for the session — ignoring the current branch — and never removed the matchingmessage_structurerow, orphaning it and corrupting subsequent numbering. On a branched session it could also delete a message belonging to a different branch.This is the same class of bug already fixed for the sibling methods
delete_branch(#3347) andadd_itemsatomicity (#3349);clear_sessionandpop_itemwere the remaining unpatched siblings.Fix: override both methods so the metadata tables stay consistent, mirroring the existing
delete_branchcleanup.pop_itemnow pops the most recent item of the current branch and reuses_cleanup_orphaned_messages_sync, so the underlying (branch-shared) message row is removed only when no other branch still references it.clear_sessionclears both aux tables in the same transaction and resets the in-memory branch pointer tomain.Test plan
Added three regression tests to
tests/extensions/memory/test_advanced_sqlite_session.py:test_clear_session_removes_structure_and_usage_metadata— afterclear_session, both aux tables are empty and numbering resets to(1, 1).test_pop_item_removes_its_structure_row— nomessage_structurerow references a deleted message after a pop.test_pop_item_respects_current_branch_and_keeps_shared_messages— popping on a branch removes only that branch's item and leaves messages still referenced by another branch intact.All three fail on
mainand pass with this change. Verified locally:make format/make lint— cleanmake typecheck— 0 errors (pyright + mypy)make tests— full suite: 4903 passed, 6 skippedIssue number
N/A (bug fix; no existing issue — happy to open one if preferred).
Checks
.agents/skills/code-change-verification/scripts/run.sh/reviewbefore submitting this PR