fix: AdvancedSQLiteSession clear_session and pop_item metadata leaks#3755
fix: AdvancedSQLiteSession clear_session and pop_item metadata leaks#3755okaditya84 wants to merge 5 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.
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".
| cursor.execute( | ||
| "DELETE FROM message_structure WHERE id = ?", | ||
| (structure_id,), | ||
| ) | ||
| self._cleanup_orphaned_messages_sync(conn) |
There was a problem hiding this comment.
Remove stale usage when popping a whole turn
When callers have stored usage for the latest branch turn and then pop_item() removes the last message_structure row for that turn, this path only deletes the structure row and orphaned message rows. The matching turn_usage row remains, so get_turn_usage() and get_session_usage() continue to report usage for a turn that no longer exists in the current branch; clear_session() and delete_branch() already clean this table, so this pop path should delete the turn's usage once no structure rows reference it.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Good call — fixed in 2c97250. pop_item now deletes the turn's turn_usage row once no message_structure rows reference that (branch, user_turn_number) anymore, so a fully-popped turn no longer shows up in get_turn_usage/get_session_usage, while a partially-popped turn keeps its usage. This matches the cleanup delete_branch and clear_session already do. Added a regression test covering both the partial and full-pop cases.
There was a problem hiding this comment.
Pull request overview
This pull request fixes metadata leaks and branch-incorrect behavior in AdvancedSQLiteSession by overriding inherited SQLiteSession methods that previously only modified the base messages/sessions tables. The change keeps message_structure and turn_usage consistent with item deletion/clearing, preventing stale analytics and sequence/turn-number drift.
Changes:
- Override
AdvancedSQLiteSession.clear_session()to deletemessage_structureandturn_usagerows explicitly (not relying on SQLite FK cascades). - Override
AdvancedSQLiteSession.pop_item()to pop the latest item from the current branch and delete the correspondingmessage_structurerow, then clean up orphaned message rows safely. - Add regression tests ensuring metadata tables are cleared on
clear_session()and remain non-orphaned / branch-correct afterpop_item().
Reviewed changes
Copilot reviewed 2 out of 2 changed files in this pull request and generated no comments.
| File | Description |
|---|---|
src/agents/extensions/memory/advanced_sqlite_session.py |
Overrides pop_item/clear_session to keep message_structure/turn_usage consistent and branch-correct without relying on SQLite FK enforcement. |
tests/extensions/memory/test_advanced_sqlite_session.py |
Adds regression coverage for metadata cleanup, correct branch popping, and preventing orphaned message_structure rows. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
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
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.
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: |
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