fix(memory): make SQLAlchemySession.pop_item atomic - #4206
Conversation
pop_item selects the newest id, reads its payload, then deletes it in three separate statements. Concurrent callers can select the same id before either delete lands, so both return that item while the rows neither of them claimed stay in the store. With forty concurrent pops of a forty item session, twenty six distinct items were returned -- one of them five times -- and fourteen rows were still present afterwards, even though every call reported success. The delete is the only atomic claim on a row, so treat deleting zero rows as losing the race and continue to the next item instead of returning one this caller did not remove. Checking rowcount keeps the dialect-agnostic fallback rather than requiring DELETE ... RETURNING support. The SQLite backends are unaffected: they claim the row with a single DELETE ... RETURNING.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 84606a26ac
ℹ️ 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 deleted.rowcount == 0: | ||
| continue |
There was a problem hiding this comment.
Restart the pop after losing the delete race
When two concurrent pop_item() calls run under MySQL/InnoDB's default REPEATABLE READ isolation, the loser can enter a tight loop here: the transaction snapshot keeps selecting the same tail row that another transaction already deleted, DELETE returns rowcount == 0, and this continue repeats without ever seeing the next item. Since SQLAlchemySession supports mysql+aiomysql, retry the pop in a new transaction or take a locking/current read before looping.
AGENTS.md reference: AGENTS.md:L134-L134
Useful? React with 👍 / 👎.
| ) | ||
| ) | ||
|
|
||
| popped = await asyncio.gather(*[session.pop_item() for _ in range(item_count)]) |
There was a problem hiding this comment.
Make the concurrent pop test deterministic
With the SQLite URL used in this test, running 40 real writers concurrently can raise OperationalError: database is locked instead of producing a stable list of popped items: pop_item() starts with reads and then upgrades to a DELETE, and unlike add_items() it is not wrapped in the SQLite write retry helper. This makes the new regression test flaky/failing on normal SQLite scheduling; use a deterministic interleaving/fake result or add the retry path before relying on this gather.
AGENTS.md reference: AGENTS.md:L169-L169
Useful? React with 👍 / 👎.
|
Thanks for sharing this patch. We'll close this PR in favor of #4212, which covers all similar patterns across the SDK. |
Summary
SQLAlchemySession.pop_item()picks the newest row id, reads its payload, then deletes it — three separate statements. Concurrent callers can complete the select for the same id before either delete lands, so both return that item. The second delete removes nothing, but its result is discarded, so the caller cannot tell. The rows no caller claimed are then never popped.pop_itemis a destructive read, so this is not a benign duplicate: an application draining a session, retrying a turn, or rewinding history processes the same item twice while other items are silently left behind.Forty concurrent pops of a forty-item session, on
main:The delete is the only statement that atomically claims a row, so this treats deleting zero rows as losing the race and continues to the next item. Checking
rowcountkeeps the existing dialect-agnostic fallback rather than requiringDELETE ... RETURNINGsupport, which is what the current code comment ("Fallback for all dialects") is deliberately written around.The SQLite backends are unaffected — they already claim the row with a single
DELETE ... RETURNING. I verified that by running the same probe against all three backends; only the SQLAlchemy one violates the contract.Test plan
tests/extensions/memory/test_sqlalchemy_session.py::test_concurrent_pop_item_hands_each_item_to_one_callerruns 40 concurrentpop_item()calls against a 40-item session and asserts every item is returned exactly once and the store ends empty.Fails on
main, passes with the fix:Verification from the repository root:
make formatmake lintmake mypymain, none in the touched filesmake pyrightmain(src/agents/sandbox/util/tar_utils.py:161)uv run pytest tests/extensions/memory/test_sqlalchemy_session.pymake testsThe full-suite run was done on Windows, where some sandbox symlink and tracing/realtime timing tests fail independently of this change. I diffed the failing set against a clean
maincheckout in the same environment: the two sets are identical (58 vs 58, no differences either way).Issue number
Closes #4205
Checks
.agents/skills/code-change-verification/scripts/run.sh/reviewbefore submitting this PRThe verification script is a bash script that shells out to
make; I ran the underlying steps individually instead, with the results above.@seratch — this came out of running the same concurrency contract against every session backend rather than reading one in isolation, which is why I am fairly confident the scope is exactly this one method.
One thing worth your call:
rowcountis the minimal fix and keeps the fallback working on every dialect, but it does leavepop_itemas three round trips per attempt. If you would rather have the atomic form where the dialect supports it,delete(...).returning(message_data)would collapse the read and the claim into one statement on PostgreSQL, SQLite 3.35+ and MariaDB, with thisrowcountpath kept for everything else. I did not do that here because it adds a dialect branch to a function that was deliberately written without one.