Skip to content

Commit 26345c1

Browse files
jbeckwith-oaijbeckwith-oai
authored
fix: batch OpenAI conversation session writes within the API limit (#5334)
* fix: batch OpenAI conversation session writes within the API limit * fix: serialize conversation session mutations across batches * fix: serialize conversation reads with batched writes * fix: reject ambiguous partial session writes on resume * fix: preserve atomic session recovery and use complete history counts * fix: settle cancelled conversation batches before releasing ownership * test: inspect cancellation before Python 3.10 task propagation * fix: disable automatic retries for chunked conversation writes * fix: deliver cancellation before sending subsequent chunks
1 parent f007b3f commit 26345c1

3 files changed

Lines changed: 737 additions & 24 deletions

File tree

‎src/agents/memory/openai_conversations_session.py‎

Lines changed: 57 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,8 @@
1111
from .session import SessionABC, _await_mutation
1212
from .session_settings import SessionSettings, coerce_session_settings, resolve_session_limit
1313

14+
_MAX_ITEMS_PER_REQUEST = 20
15+
1416

1517
async def start_openai_conversations_session(openai_client: AsyncOpenAI | None = None) -> str:
1618
_maybe_openai_client = openai_client
@@ -36,6 +38,7 @@ def __init__(
3638
):
3739
self._session_id: str | None = conversation_id
3840
self._session_id_lock = asyncio.Lock()
41+
self._mutation_lock = asyncio.Lock()
3942
self.session_settings = (
4043
coerce_session_settings(session_settings)
4144
if session_settings is not None
@@ -83,6 +86,11 @@ async def _clear_session_id(self) -> None:
8386
self._session_id = None
8487

8588
async def get_items(self, limit: int | None = None) -> list[TResponseInputItem]:
89+
async with self._mutation_lock:
90+
return await self._get_items(limit)
91+
92+
async def _get_items(self, limit: int | None = None) -> list[TResponseInputItem]:
93+
"""Read history while the caller holds the instance's mutation lock."""
8694
session_id = await self._get_session_id()
8795

8896
session_limit = resolve_session_limit(limit, self.session_settings)
@@ -112,28 +120,62 @@ async def get_items(self, limit: int | None = None) -> list[TResponseInputItem]:
112120
return all_items # type: ignore
113121

114122
async def add_items(self, items: list[TResponseInputItem]) -> None:
123+
"""Append items in order, in requests of at most 20 items each.
124+
125+
Reads and mutations through this session instance are serialized. Separate instances
126+
or external writers require application-level coordination.
127+
128+
Writes spanning multiple requests are not atomic. If a request fails or
129+
is cancelled, earlier batches remain saved and later batches are not
130+
sent. For multi-request writes, cancellation waits for the current request
131+
to complete or fail before releasing the instance lock and propagating
132+
cancellation. Single-request writes retain their existing cancellation behavior.
133+
Automatic client retries are disabled for multi-request writes, even if
134+
the supplied client enables retries, to avoid duplicating a saved chunk
135+
after a lost response. Single-request writes retain the client's retry policy.
136+
The original exception propagates. Before retrying, callers must
137+
reconcile the remote history; retrying the entire list can duplicate
138+
items that were already saved, including an unacknowledged request.
139+
"""
115140
if not items:
116141
return
117142

118-
session_id = await self._get_session_id()
119-
await self._openai_client.conversations.items.create(
120-
conversation_id=session_id,
121-
items=items,
122-
)
143+
async with self._mutation_lock:
144+
session_id = await self._get_session_id()
145+
multiple_requests = len(items) > _MAX_ITEMS_PER_REQUEST
146+
client = (
147+
self._openai_client.with_options(max_retries=0)
148+
if multiple_requests
149+
else self._openai_client
150+
)
151+
# The Conversations items-create endpoint accepts up to 20 items per request.
152+
for offset in range(0, len(items), _MAX_ITEMS_PER_REQUEST):
153+
if offset:
154+
# Deliver queued cancellation before scheduling another mutation.
155+
await asyncio.sleep(0)
156+
request = client.conversations.items.create(
157+
conversation_id=session_id,
158+
items=items[offset : offset + _MAX_ITEMS_PER_REQUEST],
159+
)
160+
if multiple_requests:
161+
await _await_mutation(request)
162+
else:
163+
await request
123164

124165
async def pop_item(self) -> TResponseInputItem | None:
125-
session_id = await self._get_session_id()
126-
items = await self.get_items(limit=1)
127-
if not items:
128-
return None
129-
item_id: str = str(items[0]["id"]) # type: ignore [typeddict-item]
130-
await self._openai_client.conversations.items.delete(
131-
conversation_id=session_id, item_id=item_id
132-
)
133-
return items[0]
166+
async with self._mutation_lock:
167+
session_id = await self._get_session_id()
168+
items = await self._get_items(limit=1)
169+
if not items:
170+
return None
171+
item_id: str = str(items[0]["id"]) # type: ignore [typeddict-item]
172+
await self._openai_client.conversations.items.delete(
173+
conversation_id=session_id, item_id=item_id
174+
)
175+
return items[0]
134176

135177
async def clear_session(self) -> None:
136-
async with self._session_id_lock:
178+
async with self._mutation_lock, self._session_id_lock:
137179
if self._session_id is None:
138180
return
139181

‎src/agents/run_internal/session_persistence.py‎

Lines changed: 29 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -39,7 +39,10 @@
3939
SessionSettings,
4040
is_openai_responses_compaction_aware_session,
4141
)
42-
from ..memory.openai_conversations_session import OpenAIConversationsSession
42+
from ..memory.openai_conversations_session import (
43+
_MAX_ITEMS_PER_REQUEST,
44+
OpenAIConversationsSession,
45+
)
4346
from ..memory.session import _call_session_method, _get_session_wrapper
4447
from ..models.fake_id import FAKE_RESPONSES_ID
4548
from ..run_context import RunContextWrapper
@@ -1496,7 +1499,8 @@ def digests(items: Sequence[TResponseInputItem]) -> list[str]:
14961499
pending["before"] = digests(tail)
14971500
append = True
14981501
else:
1499-
expected = (before or []) + digests(pending["items"])
1502+
batch = digests(pending["items"])
1503+
expected = (before or []) + batch
15001504
observed_generation: int | None = None
15011505
get_with_generation = getattr(session, "_get_items_with_generation", None)
15021506
if wrapper is not None and callable(get_with_generation):
@@ -1509,7 +1513,29 @@ def digests(items: Sequence[TResponseInputItem]) -> list[str]:
15091513
observed = digests(tail)
15101514
committed = observed == expected
15111515
unchanged = observed[-len(before) :] == before if before else not observed
1512-
if (acknowledged and not committed) or (not acknowledged and committed == unchanged):
1516+
# Only multi-request Conversations writes admit partially committed chunks.
1517+
# Other backends and single-request writes retain their atomic recovery contract.
1518+
partial = False
1519+
if (
1520+
isinstance(session, OpenAIConversationsSession)
1521+
and not acknowledged
1522+
and len(batch) > _MAX_ITEMS_PER_REQUEST
1523+
):
1524+
before_count = len(before or [])
1525+
complete_before = before_count < len(batch) + 1
1526+
# A short snapshot captured all prior history: count must agree with
1527+
# zero progress as well as each candidate chunk prefix. A full-size
1528+
# snapshot may omit older items, so repeated content remains ambiguous.
1529+
if complete_before:
1530+
unchanged = unchanged and len(observed) == before_count
1531+
partial = any(
1532+
(not complete_before or len(observed) == before_count + count)
1533+
and observed[-(before_count + count) :] == (before or []) + batch[:count]
1534+
for count in range(_MAX_ITEMS_PER_REQUEST, len(batch), _MAX_ITEMS_PER_REQUEST)
1535+
)
1536+
if (acknowledged and not committed) or (
1537+
not acknowledged and (partial or committed == unchanged)
1538+
):
15131539
raise UserError(
15141540
"Cannot reconcile the pending Session write: history changed or is "
15151541
"ambiguous. Repair the original Session before resuming; do not rerun "

0 commit comments

Comments
 (0)