diff --git a/src/agents/memory/openai_conversations_session.py b/src/agents/memory/openai_conversations_session.py index 8e0641067c..5aefc3dda6 100644 --- a/src/agents/memory/openai_conversations_session.py +++ b/src/agents/memory/openai_conversations_session.py @@ -1,6 +1,8 @@ from __future__ import annotations import asyncio +import contextlib +from collections.abc import Awaitable from typing import Any from openai import AsyncOpenAI @@ -11,6 +13,37 @@ from .session import SessionABC from .session_settings import SessionSettings, coerce_session_settings, resolve_session_limit +# Conversations items.create accepts at most 20 items per request. +# See https://developers.openai.com/api/reference/resources/conversations/subresources/items/. +_MAX_ITEMS_PER_CONVERSATION_CREATE = 20 + + +def _created_conversation_item_ids(result: object) -> list[str]: + data = getattr(result, "data", None) + if data is None: + return [] + ids: list[str] = [] + for item in data: + item_id = getattr(item, "id", None) + if isinstance(item_id, str) and item_id: + ids.append(item_id) + return ids + + +async def _await_despite_cancellation(awaitable: Awaitable[None]) -> None: + """Finish rollback even if the caller keeps cancelling the current task.""" + task = asyncio.ensure_future(awaitable) + try: + await asyncio.shield(task) + except asyncio.CancelledError: + while not task.done(): + try: + await asyncio.shield(task) + except asyncio.CancelledError: + continue + _ = task.exception() if not task.cancelled() else None + raise + async def start_openai_conversations_session(openai_client: AsyncOpenAI | None = None) -> str: _maybe_openai_client = openai_client @@ -36,6 +69,7 @@ def __init__( ): self._session_id: str | None = conversation_id self._session_id_lock = asyncio.Lock() + self._session_lock = asyncio.Lock() self.session_settings = ( coerce_session_settings(session_settings) if session_settings is not None @@ -83,6 +117,10 @@ async def _clear_session_id(self) -> None: self._session_id = None async def get_items(self, limit: int | None = None) -> list[TResponseInputItem]: + async with self._session_lock: + return await self._get_items_unlocked(limit) + + async def _get_items_unlocked(self, limit: int | None = None) -> list[TResponseInputItem]: session_id = await self._get_session_id() session_limit = resolve_session_limit(limit, self.session_settings) @@ -115,29 +153,48 @@ async def add_items(self, items: list[TResponseInputItem]) -> None: if not items: return - session_id = await self._get_session_id() - await self._openai_client.conversations.items.create( - conversation_id=session_id, - items=items, - ) + async with self._session_lock: + session_id = await self._get_session_id() + created_ids: list[str] = [] + try: + for offset in range(0, len(items), _MAX_ITEMS_PER_CONVERSATION_CREATE): + created = await self._openai_client.conversations.items.create( + conversation_id=session_id, + items=items[offset : offset + _MAX_ITEMS_PER_CONVERSATION_CREATE], + ) + created_ids.extend(_created_conversation_item_ids(created)) + except (Exception, asyncio.CancelledError): + await _await_despite_cancellation( + self._delete_created_items(session_id, created_ids) + ) + raise + + async def _delete_created_items(self, session_id: str, created_ids: list[str]) -> None: + for item_id in reversed(created_ids): + with contextlib.suppress(Exception): + await self._openai_client.conversations.items.delete( + conversation_id=session_id, item_id=item_id + ) async def pop_item(self) -> TResponseInputItem | None: - session_id = await self._get_session_id() - items = await self.get_items(limit=1) - if not items: - return None - item_id: str = str(items[0]["id"]) # type: ignore [typeddict-item] - await self._openai_client.conversations.items.delete( - conversation_id=session_id, item_id=item_id - ) - return items[0] + async with self._session_lock: + session_id = await self._get_session_id() + items = await self._get_items_unlocked(limit=1) + if not items: + return None + item_id: str = str(items[0]["id"]) # type: ignore [typeddict-item] + await self._openai_client.conversations.items.delete( + conversation_id=session_id, item_id=item_id + ) + return items[0] async def clear_session(self) -> None: - async with self._session_id_lock: - if self._session_id is None: - return - - await self._openai_client.conversations.delete( - conversation_id=self._session_id, - ) - self._session_id = None + async with self._session_lock: + async with self._session_id_lock: + if self._session_id is None: + return + + await self._openai_client.conversations.delete( + conversation_id=self._session_id, + ) + self._session_id = None diff --git a/tests/memory/test_openai_conversations_session.py b/tests/memory/test_openai_conversations_session.py index d0075418d9..7722a8f3fe 100644 --- a/tests/memory/test_openai_conversations_session.py +++ b/tests/memory/test_openai_conversations_session.py @@ -3,6 +3,7 @@ from __future__ import annotations import asyncio +from types import SimpleNamespace from typing import Any, cast from unittest.mock import AsyncMock, MagicMock, patch @@ -24,6 +25,56 @@ from tests.test_responses import get_text_message +class _InMemoryConversationItems: + """Store conversation items so tests can observe committed history.""" + + def __init__(self) -> None: + self.items: list[dict[str, Any]] = [] + self._next_id = 0 + self.create_calls = 0 + self.second_create_started = asyncio.Event() + self.release_second_create = asyncio.Event() + + async def create(self, *, conversation_id: str, items: list[Any]) -> SimpleNamespace: + self.create_calls += 1 + if self.create_calls == 2: + self.second_create_started.set() + await self.release_second_create.wait() + created: list[SimpleNamespace] = [] + for item in items: + self._next_id += 1 + stored = dict(item) + stored["id"] = f"item_{self._next_id}" + self.items.append(stored) + created.append(SimpleNamespace(id=stored["id"])) + return SimpleNamespace(data=created) + + async def delete(self, *, conversation_id: str, item_id: str) -> None: + self.items = [item for item in self.items if item["id"] != item_id] + + def list(self, **kwargs: Any) -> Any: + stored_items = list(self.items) + if kwargs.get("order") == "desc": + stored_items = list(reversed(stored_items)) + + async def _iter() -> Any: + for item in stored_items: + payload = dict(item) + yield SimpleNamespace( + model_dump=lambda exclude_unset=True, payload=payload: payload + ) + + return _iter() + + +def _bind_in_memory_conversation( + mock_openai_client: Any, store: _InMemoryConversationItems +) -> None: + mock_openai_client.conversations.items.create.side_effect = store.create + mock_openai_client.conversations.items.delete.side_effect = store.delete + mock_openai_client.conversations.items.list = MagicMock(side_effect=store.list) + + @pytest.fixture def mock_openai_client(): """Create a mock OpenAI client for testing.""" @@ -291,6 +342,176 @@ async def test_add_items_empty_keeps_existing_session_id(self, mock_openai_clien mock_openai_client.conversations.items.create.assert_not_called() assert session.session_id == "test_id" + @pytest.mark.asyncio + async def test_add_items_chunks_payloads_over_create_limit(self, mock_openai_client): + """The Conversations API accepts at most 20 items per create call.""" + session = OpenAIConversationsSession( + conversation_id="test_id", openai_client=mock_openai_client + ) + items: list[TResponseInputItem] = [ + {"role": "user", "content": f"message {index}"} for index in range(21) + ] + + await session.add_items(items) + + calls = mock_openai_client.conversations.items.create.call_args_list + assert len(calls) == 2 + assert calls[0].kwargs == {"conversation_id": "test_id", "items": items[:20]} + assert calls[1].kwargs == {"conversation_id": "test_id", "items": items[20:]} + + @pytest.mark.asyncio + async def test_add_items_does_not_chunk_payloads_at_create_limit(self, mock_openai_client): + session = OpenAIConversationsSession( + conversation_id="test_id", openai_client=mock_openai_client + ) + items: list[TResponseInputItem] = [ + {"role": "user", "content": f"message {index}"} for index in range(20) + ] + + await session.add_items(items) + + mock_openai_client.conversations.items.create.assert_called_once_with( + conversation_id="test_id", items=items + ) + + @pytest.mark.asyncio + async def test_add_items_rolls_back_created_chunks_when_a_later_create_fails( + self, mock_openai_client + ): + session = OpenAIConversationsSession( + conversation_id="test_id", openai_client=mock_openai_client + ) + items: list[TResponseInputItem] = [ + {"role": "user", "content": f"message {index}"} for index in range(21) + ] + created_ids = [f"item_{index}" for index in range(20)] + mock_openai_client.conversations.items.create.side_effect = [ + SimpleNamespace(data=[SimpleNamespace(id=item_id) for item_id in created_ids]), + RuntimeError("create failed"), + ] + + with pytest.raises(RuntimeError, match="create failed"): + await session.add_items(items) + + deleted_ids = [ + call.kwargs["item_id"] + for call in mock_openai_client.conversations.items.delete.call_args_list + ] + assert deleted_ids == list(reversed(created_ids)) + + @pytest.mark.asyncio + async def test_add_items_rolls_back_created_chunks_when_cancelled_between_creates( + self, mock_openai_client + ): + store = _InMemoryConversationItems() + seed = {"id": "seed", "role": "user", "content": "seed"} + store.items = [dict(seed)] + _bind_in_memory_conversation(mock_openai_client, store) + session = OpenAIConversationsSession( + conversation_id="test_id", openai_client=mock_openai_client + ) + items: list[TResponseInputItem] = [ + {"role": "user", "content": f"message {index}"} for index in range(21) + ] + + add_task = asyncio.create_task(session.add_items(items)) + await store.second_create_started.wait() + add_task.cancel() + + with pytest.raises(asyncio.CancelledError): + await add_task + + assert await session.get_items() == [seed] + + @pytest.mark.asyncio + async def test_cancelled_add_items_does_not_expose_partial_batch_to_overlapping_get_items( + self, mock_openai_client + ): + store = _InMemoryConversationItems() + seed = {"id": "seed", "role": "user", "content": "seed"} + store.items = [dict(seed)] + _bind_in_memory_conversation(mock_openai_client, store) + session = OpenAIConversationsSession( + conversation_id="test_id", openai_client=mock_openai_client + ) + items: list[TResponseInputItem] = [ + {"role": "user", "content": f"message {index}"} for index in range(21) + ] + + add_task = asyncio.create_task(session.add_items(items)) + await store.second_create_started.wait() + get_task = asyncio.create_task(session.get_items()) + await asyncio.sleep(0) + add_task.cancel() + + add_result, get_result = await asyncio.gather(add_task, get_task, return_exceptions=True) + + assert isinstance(add_result, asyncio.CancelledError) + assert get_result == [seed] + assert await session.get_items() == [seed] + + @pytest.mark.asyncio + async def test_cancelled_add_items_does_not_expose_partial_batch_to_overlapping_add_items( + self, mock_openai_client + ): + store = _InMemoryConversationItems() + seed = {"id": "seed", "role": "user", "content": "seed"} + store.items = [dict(seed)] + _bind_in_memory_conversation(mock_openai_client, store) + session = OpenAIConversationsSession( + conversation_id="test_id", openai_client=mock_openai_client + ) + items: list[TResponseInputItem] = [ + {"role": "user", "content": f"message {index}"} for index in range(21) + ] + surviving_item: TResponseInputItem = {"role": "user", "content": "survivor"} + + add_task = asyncio.create_task(session.add_items(items)) + await store.second_create_started.wait() + surviving_task = asyncio.create_task(session.add_items([surviving_item])) + await asyncio.sleep(0) + add_task.cancel() + + add_result, surviving_result = await asyncio.gather( + add_task, surviving_task, return_exceptions=True + ) + + assert isinstance(add_result, asyncio.CancelledError) + assert surviving_result is None + history = await session.get_items() + assert history[0] == seed + assert history[-1]["content"] == "survivor" + assert [item.get("content") for item in history if item.get("content") != "seed"] == [ + "survivor" + ] + + @pytest.mark.asyncio + async def test_cancelled_add_items_does_not_expose_partial_batch_to_overlapping_pop_item( + self, mock_openai_client + ): + store = _InMemoryConversationItems() + seed = {"id": "seed", "role": "user", "content": "seed"} + store.items = [dict(seed)] + _bind_in_memory_conversation(mock_openai_client, store) + session = OpenAIConversationsSession( + conversation_id="test_id", openai_client=mock_openai_client + ) + items: list[TResponseInputItem] = [ + {"role": "user", "content": f"message {index}"} for index in range(21) + ] + + add_task = asyncio.create_task(session.add_items(items)) + await store.second_create_started.wait() + pop_task = asyncio.create_task(session.pop_item()) + await asyncio.sleep(0) + add_task.cancel() + + add_result, popped = await asyncio.gather(add_task, pop_task, return_exceptions=True) + + assert isinstance(add_result, asyncio.CancelledError) + assert popped == seed + assert await session.get_items() == [] + @pytest.mark.asyncio async def test_pop_item_with_items(self, mock_openai_client): """Test popping item when items exist using method patching.""" @@ -301,7 +522,9 @@ async def test_pop_item_with_items(self, mock_openai_client): # Mock get_items to return one item latest_item = {"id": "item_123", "role": "assistant", "content": "Latest message"} - with patch.object(session, "get_items", return_value=[latest_item]): + with patch.object( + session, "_get_items_unlocked", new_callable=AsyncMock, return_value=[latest_item] + ): popped_item = await session.pop_item() assert popped_item == latest_item @@ -317,7 +540,7 @@ async def test_pop_item_empty_session(self, mock_openai_client): ) # Mock get_items to return empty list - with patch.object(session, "get_items", return_value=[]): + with patch.object(session, "_get_items_unlocked", new_callable=AsyncMock, return_value=[]): popped_item = await session.pop_item() assert popped_item is None @@ -609,7 +832,9 @@ async def test_invalid_item_id_in_pop_item(self, mock_openai_client): # Mock item without ID invalid_item = {"role": "assistant", "content": "No ID"} - with patch.object(session, "get_items", return_value=[invalid_item]): + with patch.object( + session, "_get_items_unlocked", new_callable=AsyncMock, return_value=[invalid_item] + ): # This should raise a KeyError because 'id' field is missing with pytest.raises(KeyError, match="'id'"): await session.pop_item()