From 84606a26acf97b465aa4a143d8ceb2674a34d466 Mon Sep 17 00:00:00 2001 From: abhay-codes07 Date: Wed, 5 Aug 2026 16:30:19 +0530 Subject: [PATCH] fix(memory): make SQLAlchemySession.pop_item atomic 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. --- .../extensions/memory/sqlalchemy_session.py | 11 +++++- .../memory/test_sqlalchemy_session.py | 37 +++++++++++++++++++ 2 files changed, 47 insertions(+), 1 deletion(-) diff --git a/src/agents/extensions/memory/sqlalchemy_session.py b/src/agents/extensions/memory/sqlalchemy_session.py index 977c25cfa4..1bc55da4a9 100644 --- a/src/agents/extensions/memory/sqlalchemy_session.py +++ b/src/agents/extensions/memory/sqlalchemy_session.py @@ -435,7 +435,16 @@ async def pop_item(self) -> TResponseInputItem | None: select(self._messages.c.message_data).where(self._messages.c.id == row_id) ) row = res_data.scalar_one_or_none() - await sess.execute(delete(self._messages).where(self._messages.c.id == row_id)) + deleted = await sess.execute( + delete(self._messages).where(self._messages.c.id == row_id) + ) + + # The DELETE is the only atomic claim on this row: the select above can + # return the same id to concurrent poppers, so treat deleting zero rows as + # losing the race and look for the next item instead of handing the same + # item to two callers. + if deleted.rowcount == 0: + continue if row is None: continue diff --git a/tests/extensions/memory/test_sqlalchemy_session.py b/tests/extensions/memory/test_sqlalchemy_session.py index 25f3001a2e..bde700c137 100644 --- a/tests/extensions/memory/test_sqlalchemy_session.py +++ b/tests/extensions/memory/test_sqlalchemy_session.py @@ -6,6 +6,7 @@ from collections.abc import Iterable, Sequence from contextlib import asynccontextmanager from datetime import datetime, timedelta +from pathlib import Path from typing import Any, cast import pytest @@ -990,3 +991,39 @@ async def test_runner_with_session_settings_override(agent: Agent): history_items = [item for item in last_input if item.get("content") != "New question"] # Should have 2 history items (last two from the 10 we added) assert len(history_items) == 2 + + +@pytest.mark.asyncio +async def test_concurrent_pop_item_hands_each_item_to_one_caller( + agent: Agent, tmp_path: Path +) -> None: + """pop_item removes an item, so concurrent callers must never receive the same one. + + pop_item selects the newest id, reads its payload, then deletes it in three separate + statements. Concurrent callers can therefore select the same id before either delete + lands, and both return that item while the rows they never claimed stay in the store. + """ + item_count = 40 + session_id = "concurrent_pop" + session = SQLAlchemySession.from_url( + session_id, + url=f"sqlite+aiosqlite:///{(tmp_path / 'concurrent_pop.db').as_posix()}", + create_tables=True, + ) + + await session.add_items( + cast( + list[TResponseInputItem], + [{"role": "user", "content": f"m{index}"} for index in range(item_count)], + ) + ) + + popped = await asyncio.gather(*[session.pop_item() for _ in range(item_count)]) + contents = [item.get("content") for item in popped if item is not None] + + assert len(contents) == item_count + # Each item belongs to exactly one caller. + assert len(set(contents)) == item_count + assert sorted(cast(list[str], contents)) == sorted(f"m{index}" for index in range(item_count)) + # And nothing is left behind by a pop that returned an item it did not delete. + assert await session.get_items() == []