diff --git a/src/agents/extensions/memory/async_sqlite_session.py b/src/agents/extensions/memory/async_sqlite_session.py index 2eef596264..8bbf2d381e 100644 --- a/src/agents/extensions/memory/async_sqlite_session.py +++ b/src/agents/extensions/memory/async_sqlite_session.py @@ -211,12 +211,27 @@ async def pop_item(self) -> TResponseInputItem | None: await cursor.close() await conn.commit() - if result: - message_data = result[0] - try: - return cast(TResponseInputItem, json.loads(message_data)) - except json.JSONDecodeError: - return None + while result: + message_data = result[0] + try: + return cast(TResponseInputItem, json.loads(message_data)) + except (json.JSONDecodeError, TypeError): + cursor = await conn.execute( + f""" + DELETE FROM {self.messages_table} + WHERE id = ( + SELECT id FROM {self.messages_table} + WHERE session_id = ? + ORDER BY id DESC + LIMIT 1 + ) + RETURNING message_data + """, + (self.session_id,), + ) + result = await cursor.fetchone() + await cursor.close() + await conn.commit() return None diff --git a/src/agents/extensions/memory/dapr_session.py b/src/agents/extensions/memory/dapr_session.py index a1b50937f0..ebb47198b3 100644 --- a/src/agents/extensions/memory/dapr_session.py +++ b/src/agents/extensions/memory/dapr_session.py @@ -344,42 +344,42 @@ async def pop_item(self) -> TResponseInputItem | None: The most recent item if it exists, None if the session is empty """ async with self._lock: - attempt = 0 while True: - attempt += 1 - response = await self._dapr_client.get_state( - store_name=self._state_store_name, - key=self._messages_key, - state_metadata=self._get_read_metadata(), - ) - messages = self._decode_messages(response.data) - if not messages: - return None - last_item = messages.pop() - messages_json = json.dumps(messages, separators=(",", ":")) - etag = getattr(response, "etag", None) or None - etag = getattr(response, "etag", None) or None - try: - await self._dapr_client.save_state( + attempt = 0 + while True: + attempt += 1 + response = await self._dapr_client.get_state( store_name=self._state_store_name, key=self._messages_key, - value=messages_json, - etag=etag, - state_metadata=self._get_metadata(), - options=self._get_state_options(concurrency=Concurrency.first_write), + state_metadata=self._get_read_metadata(), ) - break - except Exception as error: - should_retry = await self._handle_concurrency_conflict(error, attempt) - if should_retry: - continue - raise - try: - if isinstance(last_item, str): - return await self._deserialize_item(last_item) - return last_item # type: ignore[no-any-return] - except (json.JSONDecodeError, TypeError): - return None + messages = self._decode_messages(response.data) + if not messages: + return None + last_item = messages.pop() + messages_json = json.dumps(messages, separators=(",", ":")) + etag = getattr(response, "etag", None) or None + try: + await self._dapr_client.save_state( + store_name=self._state_store_name, + key=self._messages_key, + value=messages_json, + etag=etag, + state_metadata=self._get_metadata(), + options=self._get_state_options(concurrency=Concurrency.first_write), + ) + break + except Exception as error: + should_retry = await self._handle_concurrency_conflict(error, attempt) + if should_retry: + continue + raise + try: + if isinstance(last_item, str): + return await self._deserialize_item(last_item) + return last_item # type: ignore[no-any-return] + except (json.JSONDecodeError, TypeError): + continue async def clear_session(self) -> None: """Clear all items for this session.""" diff --git a/src/agents/extensions/memory/redis_session.py b/src/agents/extensions/memory/redis_session.py index bfa9181ed3..2dd3d7165f 100644 --- a/src/agents/extensions/memory/redis_session.py +++ b/src/agents/extensions/memory/redis_session.py @@ -223,22 +223,23 @@ async def pop_item(self) -> TResponseInputItem | None: The most recent item if it exists, None if the session is empty """ async with self._lock: - # Use RPOP to atomically remove and return the rightmost (most recent) item - raw_msg = await self._redis.rpop(self._messages_key) # type: ignore[misc] # Redis library returns Union[Awaitable[T], T] in async context - - if raw_msg is None: - return None - - try: - # Handle both bytes (default) and str (decode_responses=True) Redis clients - if isinstance(raw_msg, bytes): - msg_str = raw_msg.decode("utf-8") - else: - msg_str = raw_msg # Already a string - return await self._deserialize_item(msg_str) - except (json.JSONDecodeError, UnicodeDecodeError): - # Return None for corrupted messages (already removed) - return None + while True: + # Use RPOP to atomically remove and return the rightmost (most recent) item + raw_msg = await self._redis.rpop(self._messages_key) # type: ignore[misc] # Redis library returns Union[Awaitable[T], T] in async context + + if raw_msg is None: + return None + + try: + # Handle both bytes (default) and str (decode_responses=True) Redis clients + if isinstance(raw_msg, bytes): + msg_str = raw_msg.decode("utf-8") + else: + msg_str = raw_msg # Already a string + return await self._deserialize_item(msg_str) + except (json.JSONDecodeError, UnicodeDecodeError): + # Drop corrupted messages and keep looking for a valid item. + continue async def clear_session(self) -> None: """Clear all items for this session.""" diff --git a/src/agents/extensions/memory/sqlalchemy_session.py b/src/agents/extensions/memory/sqlalchemy_session.py index d84f2c78fb..ff411b7ccc 100644 --- a/src/agents/extensions/memory/sqlalchemy_session.py +++ b/src/agents/extensions/memory/sqlalchemy_session.py @@ -385,33 +385,34 @@ async def pop_item(self) -> TResponseInputItem | None: await self._ensure_tables() async with self._session_factory() as sess: async with sess.begin(): - # Fallback for all dialects - get ID first, then delete - subq = ( - select(self._messages.c.id) - .where(self._messages.c.session_id == self.session_id) - .order_by( - self._messages.c.created_at.desc(), - self._messages.c.id.desc(), + while True: + # Fallback for all dialects - get ID first, then delete + subq = ( + select(self._messages.c.id) + .where(self._messages.c.session_id == self.session_id) + .order_by( + self._messages.c.created_at.desc(), + self._messages.c.id.desc(), + ) + .limit(1) ) - .limit(1) - ) - res = await sess.execute(subq) - row_id = res.scalar_one_or_none() - if row_id is None: - return None - # Fetch data before deleting - res_data = await sess.execute( - 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)) - - if row is None: - return None - try: - return await self._deserialize_item(row) - except json.JSONDecodeError: - return None + res = await sess.execute(subq) + row_id = res.scalar_one_or_none() + if row_id is None: + return None + # Fetch data before deleting + res_data = await sess.execute( + 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)) + + if row is None: + continue + try: + return await self._deserialize_item(row) + except (json.JSONDecodeError, TypeError): + continue async def clear_session(self) -> None: """Clear all items for this session.""" diff --git a/src/agents/memory/sqlite_session.py b/src/agents/memory/sqlite_session.py index a31347cdcd..3a69f9883a 100644 --- a/src/agents/memory/sqlite_session.py +++ b/src/agents/memory/sqlite_session.py @@ -246,7 +246,7 @@ def _get_items_sync(): try: item = json.loads(message_data) items.append(item) - except json.JSONDecodeError: + except (json.JSONDecodeError, TypeError): # Skip invalid JSON entries continue @@ -297,14 +297,28 @@ def _pop_item_sync(): result = cursor.fetchone() conn.commit() - if result: + while result: message_data = result[0] try: item = json.loads(message_data) return item - except json.JSONDecodeError: - # Return None for corrupted JSON entries (already deleted) - return None + except (json.JSONDecodeError, TypeError): + # Drop corrupted JSON entries and keep looking for a valid item. + cursor = conn.execute( + f""" + DELETE FROM {self.messages_table} + WHERE id = ( + SELECT id FROM {self.messages_table} + WHERE session_id = ? + ORDER BY id DESC + LIMIT 1 + ) + RETURNING message_data + """, + (self.session_id,), + ) + result = cursor.fetchone() + conn.commit() return None diff --git a/tests/extensions/memory/test_async_sqlite_session.py b/tests/extensions/memory/test_async_sqlite_session.py index 71a13b3b92..ef235c57f0 100644 --- a/tests/extensions/memory/test_async_sqlite_session.py +++ b/tests/extensions/memory/test_async_sqlite_session.py @@ -77,6 +77,47 @@ async def test_async_sqlite_session_pop_item(): await session.close() +async def test_async_sqlite_session_pop_item_skips_corrupt_most_recent(): + """pop_item skips corrupt newest rows and returns the next valid item.""" + with tempfile.TemporaryDirectory() as temp_dir: + db_path = Path(temp_dir) / "async_pop_corrupt.db" + session = AsyncSQLiteSession("async_pop_corrupt", db_path) + + valid_item: TResponseInputItem = {"role": "user", "content": "valid"} + await session.add_items([valid_item]) + + conn = await session._get_connection() + await conn.execute( + f"INSERT INTO {session.messages_table} (session_id, message_data) VALUES (?, ?)", + (session.session_id, "not valid json {{{"), + ) + await conn.commit() + + assert await session.pop_item() == valid_item + assert await session.get_items() == [] + + await session.close() + + +async def test_async_sqlite_session_pop_item_returns_none_after_dropping_only_corrupt_rows(): + """pop_item removes corrupt rows and returns None when no valid items remain.""" + with tempfile.TemporaryDirectory() as temp_dir: + db_path = Path(temp_dir) / "async_pop_only_corrupt.db" + session = AsyncSQLiteSession("async_pop_only_corrupt", db_path) + + conn = await session._get_connection() + await conn.execute( + f"INSERT INTO {session.messages_table} (session_id, message_data) VALUES (?, ?)", + (session.session_id, "not valid json {{{"), + ) + await conn.commit() + + assert await session.pop_item() is None + assert await session.get_items() == [] + + await session.close() + + async def test_async_sqlite_session_get_items_limit(): """Test AsyncSQLiteSession get_items limit handling.""" with tempfile.TemporaryDirectory() as temp_dir: diff --git a/tests/extensions/memory/test_dapr_session.py b/tests/extensions/memory/test_dapr_session.py index af0d78b84a..9766f35d40 100644 --- a/tests/extensions/memory/test_dapr_session.py +++ b/tests/extensions/memory/test_dapr_session.py @@ -396,6 +396,41 @@ async def test_pop_from_empty_session(fake_dapr_client: FakeDaprClient): await session.close() +async def test_pop_item_skips_corrupt_most_recent(fake_dapr_client: FakeDaprClient): + """pop_item skips corrupt newest entries and returns the next valid item.""" + session = await _create_test_session(fake_dapr_client, "pop_corrupt") + + try: + valid_item: TResponseInputItem = {"role": "user", "content": "valid"} + fake_dapr_client._state[session._messages_key] = json.dumps( + [await session._serialize_item(valid_item), "not valid json {{{"], + separators=(",", ":"), + ).encode("utf-8") + + assert await session.pop_item() == valid_item + assert await session.get_items() == [] + finally: + await session.close() + + +async def test_pop_item_returns_none_after_dropping_only_corrupt_entries( + fake_dapr_client: FakeDaprClient, +): + """pop_item removes corrupt entries and returns None when no valid items remain.""" + session = await _create_test_session(fake_dapr_client, "pop_only_corrupt") + + try: + fake_dapr_client._state[session._messages_key] = json.dumps( + ["not valid json {{{"], + separators=(",", ":"), + ).encode("utf-8") + + assert await session.pop_item() is None + assert await session.get_items() == [] + finally: + await session.close() + + async def test_add_empty_items_list(fake_dapr_client: FakeDaprClient): """Test that adding an empty list of items is a no-op.""" session = await _create_test_session(fake_dapr_client) diff --git a/tests/extensions/memory/test_redis_session.py b/tests/extensions/memory/test_redis_session.py index d10c490e74..b5011cdd4d 100644 --- a/tests/extensions/memory/test_redis_session.py +++ b/tests/extensions/memory/test_redis_session.py @@ -740,12 +740,11 @@ async def test_corrupted_data_handling(): assert items[0].get("content") == "valid message" assert items[1].get("content") == "valid after corruption" - # Test pop_item with corrupted data at the end + # Test pop_item with corrupted data at the end. await _safe_rpush(fake_redis, messages_key, "corrupted at end") - # The corrupted item should be handled gracefully - # Since it's at the end, pop_item will encounter it first and return None - # But first, let's pop the valid items to get to the corrupted one + # The corrupted item should be dropped and pop_item should keep looking + # for the next valid item. popped1 = await session.pop_item() assert popped1 is not None assert popped1.get("content") == "valid after corruption" @@ -754,8 +753,7 @@ async def test_corrupted_data_handling(): assert popped2 is not None assert popped2.get("content") == "valid message" - # Now we should hit the corrupted data - this should gracefully handle it - # by returning None (and removing the corrupted item) + # All corrupt items were removed while looking for valid messages. popped_corrupted = await session.pop_item() assert popped_corrupted is None diff --git a/tests/extensions/memory/test_sqlalchemy_session.py b/tests/extensions/memory/test_sqlalchemy_session.py index 3919ada9b6..fe30993699 100644 --- a/tests/extensions/memory/test_sqlalchemy_session.py +++ b/tests/extensions/memory/test_sqlalchemy_session.py @@ -15,7 +15,7 @@ ResponseReasoningItemParam, Summary, ) -from sqlalchemy import select, text, update +from sqlalchemy import insert, select, text, update from sqlalchemy.ext.asyncio import AsyncEngine, create_async_engine from sqlalchemy.sql import Select @@ -191,6 +191,43 @@ async def test_pop_from_empty_session(): assert popped is None +async def test_pop_item_skips_corrupt_most_recent(): + """pop_item skips corrupt newest rows and returns the next valid item.""" + session = SQLAlchemySession.from_url("pop_corrupt", url=DB_URL, create_tables=True) + + valid_item: TResponseInputItem = {"role": "user", "content": "valid"} + await session.add_items([valid_item]) + + await session._ensure_tables() + async with session._session_factory() as sess: + async with sess.begin(): + await sess.execute( + insert(session._messages).values( + {"session_id": session.session_id, "message_data": "not valid json {{{"} + ) + ) + + assert await session.pop_item() == valid_item + assert await session.get_items() == [] + + +async def test_pop_item_returns_none_after_dropping_only_corrupt_rows(): + """pop_item removes corrupt rows and returns None when no valid items remain.""" + session = SQLAlchemySession.from_url("pop_only_corrupt", url=DB_URL, create_tables=True) + + await session._ensure_tables() + async with session._session_factory() as sess: + async with sess.begin(): + await sess.execute( + insert(session._messages).values( + {"session_id": session.session_id, "message_data": "not valid json {{{"} + ) + ) + + assert await session.pop_item() is None + assert await session.get_items() == [] + + async def test_add_empty_items_list(): """Test that adding an empty list of items is a no-op.""" session_id = "add_empty_test" diff --git a/tests/memory/test_session.py b/tests/memory/test_session.py index 27b5c6fa7b..f9cc324d2e 100644 --- a/tests/memory/test_session.py +++ b/tests/memory/test_session.py @@ -334,6 +334,49 @@ async def test_session_memory_pop_different_sessions(): session_2.close() +@pytest.mark.asyncio +async def test_sqlite_session_pop_item_skips_corrupt_most_recent(): + """pop_item skips corrupt newest rows and returns the next valid item.""" + with tempfile.TemporaryDirectory() as temp_dir: + db_path = Path(temp_dir) / "test_pop_corrupt.db" + session = SQLiteSession("pop_corrupt", db_path) + + valid_item: TResponseInputItem = {"role": "user", "content": "valid"} + await session.add_items([valid_item]) + + with session._locked_connection() as conn: + conn.execute( + f"INSERT INTO {session.messages_table} (session_id, message_data) VALUES (?, ?)", + (session.session_id, "not valid json {{{"), + ) + conn.commit() + + assert await session.pop_item() == valid_item + assert await session.get_items() == [] + + session.close() + + +@pytest.mark.asyncio +async def test_sqlite_session_pop_item_returns_none_after_dropping_only_corrupt_rows(): + """pop_item removes corrupt rows and returns None when no valid items remain.""" + with tempfile.TemporaryDirectory() as temp_dir: + db_path = Path(temp_dir) / "test_pop_only_corrupt.db" + session = SQLiteSession("pop_only_corrupt", db_path) + + with session._locked_connection() as conn: + conn.execute( + f"INSERT INTO {session.messages_table} (session_id, message_data) VALUES (?, ?)", + (session.session_id, "not valid json {{{"), + ) + conn.commit() + + assert await session.pop_item() is None + assert await session.get_items() == [] + + session.close() + + @pytest.mark.asyncio async def test_sqlite_session_get_items_with_limit(): """Test SQLiteSession get_items with limit parameter."""