diff --git a/src/agents/memory/openai_responses_compaction_session.py b/src/agents/memory/openai_responses_compaction_session.py index f09c3a6edd..fed2d89ff4 100644 --- a/src/agents/memory/openai_responses_compaction_session.py +++ b/src/agents/memory/openai_responses_compaction_session.py @@ -142,6 +142,10 @@ def __init__( # Serialize wrapper mutations against compaction snapshot/replace/restore so a # cancellation rollback cannot rewrite past a newer concurrent write. self._mutation_lock = asyncio.Lock() + # Runner persistence can carry this wrapper-local generation across the + # append-to-compaction gap. A later wrapper mutation revokes that one + # pending automatic replacement without inferring ownership from history. + self._mutation_generation = 0 @property def client(self) -> AsyncOpenAI: @@ -178,6 +182,35 @@ async def run_compaction( When a run context is provided, the billed compaction request contributes to that run's usage totals. """ + # Keep one wrapper mutation boundary from the snapshot through replacement. + # A concurrent add, pop, or clear waits here and then runs against the + # compacted state instead of being overwritten by a stale replacement. + async with self._mutation_lock: + has_expected_generation = wrapper is not None and hasattr( + wrapper, "_session_compaction_generation" + ) + expected_generation = ( + getattr(wrapper, "_session_compaction_generation", None) + if has_expected_generation + else None + ) + if has_expected_generation and ( + not isinstance(expected_generation, int) + or expected_generation != self._mutation_generation + ): + logger.warning( + "Skipped compaction because Session history changed after this " + "run appended its items." + ) + return + await self._run_compaction_locked(args, wrapper=wrapper) + + async def _run_compaction_locked( + self, + args: OpenAIResponsesCompactionArgs | None, + *, + wrapper: RunContextWrapper[Any] | None, + ) -> None: if args and args.get("response_id"): self._response_id = args["response_id"] requested_mode = args.get("compaction_mode") if args else None @@ -245,14 +278,18 @@ async def run_compaction( _normalize_compaction_output_items(compacted.output or []) ) - async with self._mutation_lock: - previous_items = await self._get_all_underlying_session_items() + previous_items = await self._get_all_underlying_session_items() + try: await self._replace_underlying_session_items( output_items=output_items, previous_items=previous_items, ) - self._compaction_candidate_items = select_compaction_candidate_items(output_items) - self._session_items = output_items + except (Exception, asyncio.CancelledError): + self._mutation_generation += 1 + raise + self._mutation_generation += 1 + self._compaction_candidate_items = select_compaction_candidate_items(output_items) + self._session_items = output_items logger.debug( "compact: done for %s (mode=%s, output=%s, candidates=%s)", @@ -265,6 +302,14 @@ async def run_compaction( async def get_items(self, limit: int | None = None) -> list[TResponseInputItem]: return await self.underlying_session.get_items(limit) + async def _get_items_with_generation( + self, limit: int | None = None + ) -> tuple[list[TResponseInputItem], int]: + """Read one Runner snapshot with its exact wrapper generation.""" + async with self._mutation_lock: + items = await self.underlying_session.get_items(limit) + return items, self._mutation_generation + async def _get_all_underlying_session_items(self) -> list[TResponseInputItem]: return await self.underlying_session.get_items(limit=_ALL_SESSION_ITEMS_LIMIT) @@ -410,37 +455,69 @@ def _clear_deferred_compaction(self) -> None: self._deferred_response_id = None async def add_items(self, items: list[TResponseInputItem]) -> None: + async with self._mutation_lock: + await self._add_items_locked(items) + + async def _add_items_with_generation( + self, + items: list[TResponseInputItem], + *, + expected_generation: int | None, + ) -> int | None: + """Append one Runner batch and retain ownership only when its read stayed current.""" + async with self._mutation_lock: + owns_generation = expected_generation == self._mutation_generation + await self._add_items_locked(items) + return self._mutation_generation if owns_generation else None + + async def _add_items_locked(self, items: list[TResponseInputItem]) -> None: + try: + await self.underlying_session.add_items(items) + except (Exception, asyncio.CancelledError): + # The backend may have committed before acknowledgement failed. Re-read its + # authoritative history before compaction instead of retaining a stale cache. + self._compaction_candidate_items = None + self._session_items = None + self._mutation_generation += 1 + raise + self._mutation_generation += 1 + if self._compaction_candidate_items is not None: + new_items = _normalize_compaction_session_items(items) + new_candidates = select_compaction_candidate_items(new_items) + if new_candidates: + self._compaction_candidate_items.extend(new_candidates) + if self._session_items is not None: + self._session_items.extend(_normalize_compaction_session_items(items)) + + async def pop_item(self) -> TResponseInputItem | None: async with self._mutation_lock: try: - await self.underlying_session.add_items(items) + popped = await self.underlying_session.pop_item() except (Exception, asyncio.CancelledError): - # The backend may have committed before acknowledgement failed. Re-read its - # authoritative history before compaction instead of retaining a stale cache. self._compaction_candidate_items = None self._session_items = None + self._mutation_generation += 1 raise - if self._compaction_candidate_items is not None: - new_items = _normalize_compaction_session_items(items) - new_candidates = select_compaction_candidate_items(new_items) - if new_candidates: - self._compaction_candidate_items.extend(new_candidates) - if self._session_items is not None: - self._session_items.extend(_normalize_compaction_session_items(items)) - - async def pop_item(self) -> TResponseInputItem | None: - async with self._mutation_lock: - popped = await self.underlying_session.pop_item() if popped: self._compaction_candidate_items = None self._session_items = None + self._mutation_generation += 1 return popped async def clear_session(self) -> None: async with self._mutation_lock: - await self.underlying_session.clear_session() + try: + await self.underlying_session.clear_session() + except (Exception, asyncio.CancelledError): + self._compaction_candidate_items = None + self._session_items = None + self._deferred_response_id = None + self._mutation_generation += 1 + raise self._compaction_candidate_items = [] self._session_items = [] self._deferred_response_id = None + self._mutation_generation += 1 async def _ensure_compaction_candidates( self, diff --git a/src/agents/run_internal/session_persistence.py b/src/agents/run_internal/session_persistence.py index 2e53667e7e..de5d3e6b3c 100644 --- a/src/agents/run_internal/session_persistence.py +++ b/src/agents/run_internal/session_persistence.py @@ -199,8 +199,17 @@ async def _session_get_items( limit: int | None | object = _SESSION_LIMIT_UNSET, *, wrapper: RunContextWrapper[Any] | None = None, + capture_compaction_generation: bool = False, ) -> list[TResponseInputItem]: """Read session items while preserving the legacy method call shape.""" + get_with_generation = getattr(session, "_get_items_with_generation", None) + if capture_compaction_generation and wrapper is not None and callable(get_with_generation): + if limit is _SESSION_LIMIT_UNSET: + result, generation = await _call_session_method(get_with_generation) + else: + result, generation = await _call_session_method(get_with_generation, limit=limit) + wrapper._session_compaction_generation = generation # type: ignore[attr-defined] + return cast(list[TResponseInputItem], result) wrapper = _get_session_wrapper(session, wrapper) if limit is _SESSION_LIMIT_UNSET: result = await _call_session_method(session.get_items, wrapper=wrapper) @@ -216,6 +225,16 @@ async def _session_add_items( wrapper: RunContextWrapper[Any] | None = None, ) -> None: """Append session items while preserving the legacy method call shape.""" + add_with_generation = getattr(session, "_add_items_with_generation", None) + if wrapper is not None and callable(add_with_generation): + expected_generation = getattr(wrapper, "_session_compaction_generation", None) + generation = await _call_session_method( + add_with_generation, + items, + expected_generation=expected_generation, + ) + wrapper._session_compaction_generation = generation # type: ignore[attr-defined] + return wrapper = _get_session_wrapper(session, wrapper) await _call_session_method(session.add_items, items, wrapper=wrapper) @@ -360,9 +379,14 @@ async def prepare_input_with_session( session, limit=resolved_settings.limit, wrapper=wrapper, + capture_compaction_generation=True, ) else: - history = await _session_get_items(session, wrapper=wrapper) + history = await _session_get_items( + session, + wrapper=wrapper, + capture_compaction_generation=True, + ) is_openai_conversation_session = isinstance(session, OpenAIConversationsSession) converted_history = [ strip_internal_input_item_metadata(ensure_input_item_format(item)) for item in history @@ -674,9 +698,13 @@ async def save_result_to_session( resumed_write_state._current_turn_persisted_item_count + saved_run_items_count ), } - await resume_pending_session_write(resumed_write_state, session, wrapper=wrapper) + await resume_pending_session_write( + resumed_write_state, + session, + wrapper=compaction_wrapper, + ) else: - await _session_add_items(session, items_to_save, wrapper=wrapper) + await _session_add_items(session, items_to_save, wrapper=compaction_wrapper) if run_state is not None: run_state._current_turn_persisted_item_count = already_persisted + saved_run_items_count @@ -801,7 +829,15 @@ def digests(items: Sequence[TResponseInputItem]) -> list[str]: append = True else: expected = before + digests(pending["items"]) - tail = await _session_get_items(session, limit=len(expected), wrapper=wrapper) + committed_generation: int | None = None + get_with_generation = getattr(session, "_get_items_with_generation", None) + if wrapper is not None and callable(get_with_generation): + tail, committed_generation = await _call_session_method( + get_with_generation, + limit=len(expected), + ) + else: + tail = await _session_get_items(session, limit=len(expected), wrapper=wrapper) observed = digests(tail) committed = observed == expected unchanged = observed[-len(before) :] == before if before else not observed @@ -811,6 +847,8 @@ def digests(items: Sequence[TResponseInputItem]) -> list[str]: "Repair the original Session before resuming; do not rerun the completed tool." ) append = unchanged + if committed and committed_generation is not None and wrapper is not None: + wrapper._session_compaction_generation = committed_generation # type: ignore[attr-defined] if append: # Backends may retain or transform their input; the durable checkpoint stays detached. await _session_add_items(session, copy.deepcopy(pending["items"]), wrapper=wrapper) diff --git a/tests/memory/test_openai_responses_compaction_session.py b/tests/memory/test_openai_responses_compaction_session.py index 5519228ea6..7f8895a31c 100644 --- a/tests/memory/test_openai_responses_compaction_session.py +++ b/tests/memory/test_openai_responses_compaction_session.py @@ -16,7 +16,7 @@ import agents._debug as _debug from agents import Agent, Runner -from agents.items import TResponseInputItem +from agents.items import MessageOutputItem, TResponseInputItem from agents.memory import ( OpenAIResponsesCompactionSession, Session, @@ -30,11 +30,17 @@ is_openai_model_name, select_compaction_candidate_items, ) +from agents.run_context import RunContextWrapper from agents.run_internal.items import ( TOOL_CALL_SESSION_DESCRIPTION_KEY, TOOL_CALL_SESSION_TITLE_KEY, ) -from agents.testing import ScriptedModel +from agents.run_internal.session_persistence import ( + prepare_input_with_session, + save_result_to_session, +) +from agents.run_state import RunState +from agents.testing import ModelStep, ScriptedModel from tests.test_responses import get_function_tool, get_function_tool_call, get_text_message from tests.utils.simple_session import SimpleListSession @@ -1702,6 +1708,313 @@ def test_strips_multiple_assistant_ids(self) -> None: assert "id" not in item +class TestCompactionMutationSerialization: + @pytest.mark.asyncio + async def test_resumed_save_carries_generation_into_compaction(self) -> None: + """A resumed append uses the same ownership handoff as a fresh save.""" + underlying = SimpleListSession() + client = MagicMock() + client.responses.compact = AsyncMock(return_value=SimpleNamespace(output=[])) + resumed_persisted = asyncio.Event() + release_resumed = asyncio.Event() + add_calls = 0 + + class PausingCompactionSession(OpenAIResponsesCompactionSession): + async def _add_items_with_generation( + self, + items: list[TResponseInputItem], + *, + expected_generation: int | None, + ) -> int | None: + nonlocal add_calls + generation = await super()._add_items_with_generation( + items, + expected_generation=expected_generation, + ) + add_calls += 1 + if add_calls == 1: + resumed_persisted.set() + await release_resumed.wait() + return generation + + session = PausingCompactionSession( + session_id="resumed-interleaved-persist", + underlying_session=underlying, + client=client, + compaction_mode="previous_response_id", + should_trigger_compaction=lambda context: context["response_id"] == "resp-a", + ) + agent = Agent(name="worker-a") + wrapper = RunContextWrapper(context=None) + state: RunState[Any] = RunState( + context=wrapper, + original_input=[], + starting_agent=agent, + ) + resumed_item = MessageOutputItem(agent=agent, raw_item=get_text_message("run-A")) + await prepare_input_with_session([], session, None, wrapper=wrapper) + + resumed_save = asyncio.create_task( + save_result_to_session( + session, + [], + [resumed_item], + state, + response_id="resp-a", + wrapper=wrapper, + resumed_write_state=state, + ) + ) + await asyncio.wait_for(resumed_persisted.wait(), timeout=1) + await asyncio.wait_for( + Runner.run( + Agent( + name="worker-b", + model=ScriptedModel( + steps=[ + ModelStep( + output=[get_text_message("run-B")], + response_id="resp-b", + ) + ] + ), + ), + "input-B", + session=session, + ), + timeout=1, + ) + release_resumed.set() + await resumed_save + + client.responses.compact.assert_not_awaited() + stored = str(await underlying.get_items()) + assert "run-A" in stored + assert "run-B" in stored + + @pytest.mark.asyncio + async def test_runner_skips_compaction_after_interleaved_model_wait(self) -> None: + """A later run that completes during A's model wait revokes A's replacement.""" + underlying = SimpleListSession() + client = MagicMock() + client.responses.compact = AsyncMock(return_value=SimpleNamespace(output=[])) + model_entered = asyncio.Event() + release_model = asyncio.Event() + + async def respond_after_b(_: Any) -> ModelStep: + model_entered.set() + await release_model.wait() + return ModelStep(output=[get_text_message("run-A")], response_id="resp-a") + + session = OpenAIResponsesCompactionSession( + session_id="interleaved-model-wait", + underlying_session=underlying, + client=client, + compaction_mode="previous_response_id", + should_trigger_compaction=lambda context: context["response_id"] == "resp-a", + ) + run_a = asyncio.create_task( + Runner.run( + Agent( + name="worker-a", + model=ScriptedModel(steps=[ModelStep.respond(respond_after_b)]), + ), + "input-A", + session=session, + ) + ) + await asyncio.wait_for(model_entered.wait(), timeout=1) + await asyncio.wait_for( + Runner.run( + Agent( + name="worker-b", + model=ScriptedModel( + steps=[ + ModelStep( + output=[get_text_message("run-B")], + response_id="resp-b", + ) + ] + ), + ), + "input-B", + session=session, + ), + timeout=1, + ) + release_model.set() + await run_a + + client.responses.compact.assert_not_awaited() + stored = str(await underlying.get_items()) + assert "run-A" in stored + assert "run-B" in stored + + @pytest.mark.asyncio + async def test_runner_skips_compaction_after_interleaved_persist(self) -> None: + """A later Runner append between save and compact revokes replacement.""" + underlying = SimpleListSession() + client = MagicMock() + client.responses.compact = AsyncMock(return_value=SimpleNamespace(output=[])) + run_a_persisted = asyncio.Event() + release_run_a = asyncio.Event() + add_calls = 0 + + class PausingCompactionSession(OpenAIResponsesCompactionSession): + async def _add_items_with_generation( + self, + items: list[TResponseInputItem], + *, + expected_generation: int | None, + ) -> int | None: + nonlocal add_calls + generation = await super()._add_items_with_generation( + items, + expected_generation=expected_generation, + ) + add_calls += 1 + if add_calls == 2: + run_a_persisted.set() + await release_run_a.wait() + return generation + + session = PausingCompactionSession( + session_id="interleaved-persist", + underlying_session=underlying, + client=client, + compaction_mode="previous_response_id", + should_trigger_compaction=lambda context: context["response_id"] == "resp-a", + ) + + run_a = asyncio.create_task( + Runner.run( + Agent( + name="worker-a", + model=ScriptedModel( + steps=[ + ModelStep( + output=[get_text_message("run-A")], + response_id="resp-a", + ) + ] + ), + ), + "input-A", + session=session, + ) + ) + await asyncio.wait_for(run_a_persisted.wait(), timeout=1) + await asyncio.wait_for( + Runner.run( + Agent( + name="worker-b", + model=ScriptedModel( + steps=[ + ModelStep( + output=[get_text_message("run-B")], + response_id="resp-b", + ) + ] + ), + ), + "input-B", + session=session, + ), + timeout=1, + ) + release_run_a.set() + await run_a + + client.responses.compact.assert_not_awaited() + stored = str(await underlying.get_items()) + assert "run-A" in stored + assert "run-B" in stored + + @pytest.mark.asyncio + async def test_add_waits_for_in_flight_compaction_and_survives(self) -> None: + old_item = cast( + TResponseInputItem, + {"type": "message", "role": "assistant", "content": "old"}, + ) + concurrent_item = cast( + TResponseInputItem, + {"type": "message", "role": "user", "content": "concurrent"}, + ) + compacted_item = cast( + TResponseInputItem, + {"type": "compaction", "summary": "compacted"}, + ) + underlying = SimpleListSession(history=[old_item]) + compact_entered = asyncio.Event() + release_compact = asyncio.Event() + + async def compact(**_: Any) -> SimpleNamespace: + compact_entered.set() + await release_compact.wait() + return SimpleNamespace(output=[compacted_item]) + + client = MagicMock() + client.responses.compact = AsyncMock(side_effect=compact) + session = OpenAIResponsesCompactionSession( + session_id="serialized-add", + underlying_session=underlying, + client=client, + compaction_mode="input", + ) + + compaction_task = asyncio.create_task(session.run_compaction({"force": True})) + await compact_entered.wait() + add_task = asyncio.create_task(session.add_items([concurrent_item])) + await asyncio.sleep(0) + assert not add_task.done() + + release_compact.set() + await compaction_task + await add_task + + assert await underlying.get_items() == [compacted_item, concurrent_item] + + @pytest.mark.asyncio + async def test_clear_waits_for_in_flight_compaction_and_stays_empty(self) -> None: + old_item = cast( + TResponseInputItem, + {"type": "message", "role": "assistant", "content": "old"}, + ) + compacted_item = cast( + TResponseInputItem, + {"type": "compaction", "summary": "compacted"}, + ) + underlying = SimpleListSession(history=[old_item]) + compact_entered = asyncio.Event() + release_compact = asyncio.Event() + + async def compact(**_: Any) -> SimpleNamespace: + compact_entered.set() + await release_compact.wait() + return SimpleNamespace(output=[compacted_item]) + + client = MagicMock() + client.responses.compact = AsyncMock(side_effect=compact) + session = OpenAIResponsesCompactionSession( + session_id="serialized-clear", + underlying_session=underlying, + client=client, + compaction_mode="input", + ) + + compaction_task = asyncio.create_task(session.run_compaction({"force": True})) + await compact_entered.wait() + clear_task = asyncio.create_task(session.clear_session()) + await asyncio.sleep(0) + assert not clear_task.done() + + release_compact.set() + await compaction_task + await clear_task + + assert await underlying.get_items() == [] + + class TestCompactionStripsOrphanedIds: """Regression test for #2727: gpt-5.4 compact retains assistant msg IDs after stripping reasoning items, causing 400 errors on the next responses.create call."""