From 787f9d6cc280b476f0367a788077ca7e400d59c1 Mon Sep 17 00:00:00 2001 From: Kazuhiro Sera Date: Tue, 25 Aug 2026 07:03:32 +0900 Subject: [PATCH] fix: recover failed resumed Session writes before model calls --- .../openai_responses_compaction_session.py | 9 +- src/agents/result.py | 1 + src/agents/run.py | 3 + src/agents/run_internal/run_loop.py | 8 + .../run_internal/session_persistence.py | 87 +++- src/agents/run_state.py | 56 ++- tests/test_agent_runner_streamed.py | 2 + tests/test_run_impl_resume_paths.py | 421 +++++++++++++++++- 8 files changed, 582 insertions(+), 5 deletions(-) diff --git a/src/agents/memory/openai_responses_compaction_session.py b/src/agents/memory/openai_responses_compaction_session.py index 11cc8bb682..f09c3a6edd 100644 --- a/src/agents/memory/openai_responses_compaction_session.py +++ b/src/agents/memory/openai_responses_compaction_session.py @@ -411,7 +411,14 @@ def _clear_deferred_compaction(self) -> None: async def add_items(self, items: list[TResponseInputItem]) -> None: async with self._mutation_lock: - await self.underlying_session.add_items(items) + 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 + raise if self._compaction_candidate_items is not None: new_items = _normalize_compaction_session_items(items) new_candidates = select_compaction_candidate_items(new_items) diff --git a/src/agents/result.py b/src/agents/result.py index f88819df54..0ceb0d7187 100644 --- a/src/agents/result.py +++ b/src/agents/result.py @@ -148,6 +148,7 @@ def _populate_state_from_result( if isinstance(source_state, RunState): state._generated_prompt_cache_key = source_state._generated_prompt_cache_key state._pending_input = copy.deepcopy(source_state._pending_input) + state._pending_session_write = copy.deepcopy(source_state._pending_session_write) state._current_step = source_state._current_step else: state._generated_prompt_cache_key = getattr(result, "_generated_prompt_cache_key", None) diff --git a/src/agents/run.py b/src/agents/run.py index a782f80cce..629b5ff23f 100644 --- a/src/agents/run.py +++ b/src/agents/run.py @@ -139,6 +139,7 @@ persist_session_items_for_guardrail_trip, prepare_input_with_session, reconcile_nested_history_owned_session_item_refs, + resume_pending_session_write, resumed_turn_items, save_result_to_session, save_resumed_turn_items, @@ -634,6 +635,7 @@ async def _run_impl( ) context = context_wrapper.context + await resume_pending_session_write(run_state, session, wrapper=context_wrapper) max_turns = run_state._max_turns else: raw_input = cast(str | list[TResponseInputItem], input) @@ -1149,6 +1151,7 @@ def _mark_response_hooks_started() -> None: ): run_state._current_turn_persisted_item_count = ( await save_resumed_turn_items( + run_state=run_state, session=session, items=turn_session_items, persisted_count=( diff --git a/src/agents/run_internal/run_loop.py b/src/agents/run_internal/run_loop.py index ba0c02f351..7c22fee317 100644 --- a/src/agents/run_internal/run_loop.py +++ b/src/agents/run_internal/run_loop.py @@ -177,6 +177,7 @@ persist_session_items_for_guardrail_trip, prepare_input_with_session, reconcile_nested_history_owned_session_item_refs, + resume_pending_session_write, resumed_turn_items, rewind_session_items, save_result_to_session, @@ -392,6 +393,7 @@ async def _save_resumed_stream_items( ): return streamed_result._current_turn_persisted_item_count = await save_resumed_turn_items( + run_state=run_state, session=session, items=items, persisted_count=streamed_result._current_turn_persisted_item_count, @@ -920,6 +922,12 @@ async def start_streaming( run_state._reasoning_item_id_policy = resolved_reasoning_item_id_policy streamed_result._reasoning_item_id_policy = resolved_reasoning_item_id_policy + if is_resumed_state and run_state is not None: + await resume_pending_session_write(run_state, session, wrapper=context_wrapper) + streamed_result._current_turn_persisted_item_count = ( + run_state._current_turn_persisted_item_count + ) + if ( conversation_id is not None or previous_response_id is not None diff --git a/src/agents/run_internal/session_persistence.py b/src/agents/run_internal/session_persistence.py index 8ebba55802..809f1648c3 100644 --- a/src/agents/run_internal/session_persistence.py +++ b/src/agents/run_internal/session_persistence.py @@ -7,6 +7,7 @@ import asyncio import copy +import hashlib import inspect import json from collections import deque @@ -60,7 +61,7 @@ strip_internal_input_item_metadata, ) from .oai_conversation import OpenAIServerConversationTracker -from .run_steps import NextStepInterruption, ProcessedResponse, SingleStepResult +from .run_steps import NextStepInterruption, NextStepRunAgain, ProcessedResponse, SingleStepResult __all__ = [ "admit_pending_input", @@ -73,6 +74,7 @@ "resumed_turn_items", "save_result_to_session", "save_resumed_turn_items", + "resume_pending_session_write", "update_run_state_after_resume", "rewind_session_items", "wait_for_session_cleanup", @@ -552,6 +554,7 @@ async def save_result_to_session( reasoning_item_id_policy: ReasoningItemIdPolicy | None = None, store: bool | None = None, wrapper: RunContextWrapper[Any] | None = None, + resumed_write_state: RunState | None = None, ) -> int: """ Persist a turn to the session store, keeping track of what was already saved so retries @@ -648,7 +651,20 @@ async def save_result_to_session( run_state._current_turn_persisted_item_count = already_persisted + saved_run_items_count return saved_run_items_count - await _session_add_items(session, items_to_save, wrapper=wrapper) + if resumed_write_state is not None: + if resumed_write_state._pending_session_write is not None: + raise UserError("Resolve the pending Session write before saving another batch") + resumed_write_state._pending_session_write = { + "session_id": session.session_id, + "items": copy.deepcopy(items_to_save), + "before": None, + "persisted_count": ( + resumed_write_state._current_turn_persisted_item_count + saved_run_items_count + ), + } + await resume_pending_session_write(resumed_write_state, session, wrapper=wrapper) + else: + await _session_add_items(session, items_to_save, wrapper=wrapper) if run_state is not None: run_state._current_turn_persisted_item_count = already_persisted + saved_run_items_count @@ -707,6 +723,7 @@ async def save_resumed_turn_items( reasoning_item_id_policy: ReasoningItemIdPolicy | None = None, store: bool | None = None, wrapper: RunContextWrapper[Any] | None = None, + run_state: RunState | None = None, ) -> int: """Persist resumed turn items and return the updated persisted count.""" if session is None or not items: @@ -720,10 +737,76 @@ async def save_resumed_turn_items( reasoning_item_id_policy=reasoning_item_id_policy, store=store, wrapper=wrapper, + resumed_write_state=( + run_state + if run_state is not None and isinstance(run_state._current_step, NextStepRunAgain) + else None + ), ) return persisted_count + saved_count +async def resume_pending_session_write( + run_state: RunState, + session: Session | None, + *, + wrapper: RunContextWrapper[Any] | None = None, +) -> None: + """Settle a resumed output batch before allowing further model work. + + The application must supply the original backend and serialize access to its history, + including independently restored RunState copies. Session has no distributed compare-and-swap + or backend identity contract. A changed tail is not repaired or searched for similar items. + """ + pending = run_state._pending_session_write + if pending is None: + return + if run_state._session_write_in_progress: + raise UserError("The pending Session write is already in progress for this RunState") + if session is None or session.session_id != pending["session_id"]: + raise UserError("Resume the pending Session write with the original Session and session ID") + + def digests(items: Sequence[TResponseInputItem]) -> list[str]: + return [ + hashlib.sha256( + _fingerprint_or_repr( + item, ignore_ids_for_matching=_ignore_ids_for_matching(session) + ).encode("utf-8") + ).hexdigest() + for item in items + ] + + run_state._session_write_in_progress = True + try: + before = pending["before"] + if before is None: + # No append has started. Retain the batch even if this first read fails. + tail = await _session_get_items( + session, limit=len(pending["items"]) + 1, wrapper=wrapper + ) + pending["before"] = digests(tail) + append = True + else: + expected = before + digests(pending["items"]) + 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 + if committed == unchanged: + raise UserError( + "Cannot reconcile the pending Session write: history changed or is ambiguous. " + "Repair the original Session before resuming; do not rerun the completed tool." + ) + append = unchanged + 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) + run_state._current_turn_persisted_item_count = pending["persisted_count"] + run_state._pending_session_write = None + finally: + run_state._session_write_in_progress = False + + async def rewind_session_items( session: Session | None, items: Sequence[TResponseInputItem], diff --git a/src/agents/run_state.py b/src/agents/run_state.py index e00a73bdce..b196bbaf85 100644 --- a/src/agents/run_state.py +++ b/src/agents/run_state.py @@ -165,6 +165,15 @@ ] +class _PendingSessionWrite(TypedDict): + """One canonical resumed-output append awaiting acknowledgement.""" + + session_id: str + items: list[TResponseInputItem] + before: list[str] | None + persisted_count: int + + def _default_run_state_validation_error( message: str, error_type: RunStateValidationErrorType, @@ -216,7 +225,7 @@ def _default_run_state_validation_error( ), "1.17": ( "Persists Docker container labels and current-response generated-item ownership across " - "resume flows." + "resume flows, including pending resumed Session writes." ), } SUPPORTED_SCHEMA_VERSIONS = frozenset(SCHEMA_VERSION_SUMMARIES) @@ -757,6 +766,13 @@ class RunState(Generic[TContext, TAgent]): enough information to continue an interrupted run, including model responses, generated items, approval state, and optional server-managed conversation identifiers. + A failed Session append after resumed tool work that continues to another model call remains + pending across serialization. + Resume with the original Session backend and session ID, with exclusive access to that history. + Runner reconciles the exact pending batch before the next model call without rerunning the tool. + Changed or ambiguous history requires application repair. Independently restored snapshots must + not be resumed concurrently against the same Session. + Context serialization is intentionally conservative: - Mapping contexts round-trip directly. @@ -854,6 +870,12 @@ class RunState(Generic[TContext, TAgent]): _schema_version: str = field(default=CURRENT_SCHEMA_VERSION, repr=False) """Schema version the snapshot was loaded from for schema-gated resume compatibility.""" + _pending_session_write: _PendingSessionWrite | None = field(default=None, repr=False) + """Canonical Session append that must settle before another model call.""" + + _session_write_in_progress: bool = field(default=False, repr=False) + """Live ownership guard; independent serialized copies require caller serialization.""" + def __init__( self, context: RunContextWrapper[TContext], @@ -894,6 +916,8 @@ def __init__( self._trace_state = None self._sandbox = None self._schema_version = CURRENT_SCHEMA_VERSION + self._pending_session_write = None + self._session_write_in_progress = False from .agent_tool_state import get_agent_tool_state_scope self._agent_tool_state_scope_id = get_agent_tool_state_scope(context) @@ -901,6 +925,8 @@ def __init__( def _copy_for_result_checkpoint(self) -> RunState[TContext, TAgent]: """Copy SDK-owned decision state when nesting this checkpoint in a result snapshot.""" copied = copy.copy(self) + copied._pending_session_write = copy.deepcopy(self._pending_session_write) + copied._session_write_in_progress = False if self._context is None: return copied copied._context = self._context._copy_for_run_state() @@ -1879,6 +1905,8 @@ def to_json( else None ) result["current_turn_persisted_item_count"] = self._current_turn_persisted_item_count + if self._pending_session_write is not None: + result["pending_session_write"] = copy.deepcopy(self._pending_session_write) result["trace"] = self._serialize_trace_data( include_tracing_api_key=include_tracing_api_key ) @@ -4328,6 +4356,31 @@ async def _build_run_state_from_json( state._current_turn_persisted_item_count = state_json.get( "current_turn_persisted_item_count", 0 ) + pending_write = state_json.get("pending_session_write") + if pending_write is not None: + from .run_internal.run_steps import NextStepRunAgain + + if ( + (schema_major, schema_minor) < (1, 17) + or not isinstance(state._current_step, NextStepRunAgain) + or not isinstance(pending_write, dict) + or set(pending_write) != {"session_id", "items", "before", "persisted_count"} + or not isinstance(pending_write.get("session_id"), str) + or not isinstance(pending_write.get("items"), list) + or not pending_write["items"] + or not all(isinstance(item, dict) for item in pending_write["items"]) + or ( + pending_write.get("before") is not None + and ( + not isinstance(pending_write["before"], list) + or not all(isinstance(item, str) for item in pending_write["before"]) + ) + ) + or type(pending_write.get("persisted_count")) is not int + or pending_write["persisted_count"] < 0 + ): + raise validation_error_factory("Run state pending Session write is invalid", UserError) + state._pending_session_write = copy.deepcopy(cast(_PendingSessionWrite, pending_write)) serialized_policy = state_json.get("reasoning_item_id_policy") if serialized_policy in {"preserve", "omit"}: state._reasoning_item_id_policy = cast(Literal["preserve", "omit"], serialized_policy) @@ -5591,6 +5644,7 @@ def _clone_original_input(original_input: str | list[Any]) -> str | list[Any]: ), "Run state agent not found in agent map", "Run state pending_input must be a list", + "Run state pending Session write is invalid", "Run state references an agent identity that is not present in the restored graph", ( "RunState context was serialized from a custom type; provide context_deserializer " diff --git a/tests/test_agent_runner_streamed.py b/tests/test_agent_runner_streamed.py index 9664d65513..2923286924 100644 --- a/tests/test_agent_runner_streamed.py +++ b/tests/test_agent_runner_streamed.py @@ -4317,6 +4317,7 @@ async def save_wrapper( reasoning_item_id_policy: str | None = None, store: bool | None = None, wrapper: RunContextWrapper[Any] | None = None, + run_state: RunState | None = None, ) -> int: observed_counts.append(persisted_count) result = await real_save_resumed( @@ -4327,6 +4328,7 @@ async def save_wrapper( reasoning_item_id_policy=reasoning_item_id_policy, store=store, wrapper=wrapper, + run_state=run_state, ) return int(result) diff --git a/tests/test_run_impl_resume_paths.py b/tests/test_run_impl_resume_paths.py index c518c95ee0..9a4d88f061 100644 --- a/tests/test_run_impl_resume_paths.py +++ b/tests/test_run_impl_resume_paths.py @@ -1,6 +1,9 @@ import asyncio +import copy import json -from typing import Any, cast +from pathlib import Path +from types import SimpleNamespace +from typing import Any, Literal, cast import pytest from openai.types.responses import ResponseFunctionToolCall, ResponseOutputMessage @@ -9,14 +12,18 @@ from agents import Agent, Runner, function_tool from agents.agent import ToolsToFinalOutputResult from agents.agent_output import AgentOutputSchema +from agents.decorators import tool +from agents.exceptions import UserError from agents.items import ( MessageOutputItem, ModelResponse, ToolApprovalItem, ToolCallItem, ToolCallOutputItem, + TResponseInputItem, ) from agents.lifecycle import RunHooks +from agents.memory import OpenAIResponsesCompactionSession, Session, SQLiteSession from agents.run import RunConfig from agents.run_context import RunContextWrapper from agents.run_internal import run_loop, turn_resolution @@ -42,6 +49,418 @@ from tests.utils.simple_session import SimpleListSession +class _FailingResumeSession(SimpleListSession): + """Control append acknowledgement at the public Session boundary.""" + + def __init__(self) -> None: + super().__init__() + self.failure: str | None = None + self.error = RuntimeError("session append failed") + self.block_next_add = False + self.add_started = asyncio.Event() + self.release_add = asyncio.Event() + + async def add_items(self, items: list[TResponseInputItem]) -> None: + failure, self.failure = self.failure, None + if failure == "before": + raise self.error + if self.block_next_add: + self.block_next_add = False + self.add_started.set() + await self.release_add.wait() + if failure == "partial": + await super().add_items(items[:1]) + raise self.error + await super().add_items(items) + if failure == "after": + raise self.error + + +class _LostAckSQLiteSession(SQLiteSession): + fail_after_commit = False + error = RuntimeError("session append failed") + + async def add_items(self, items: list[TResponseInputItem]) -> None: + await super().add_items(items) + if self.fail_after_commit: + self.fail_after_commit = False + raise self.error + + +async def _run_session_resume( + agent: Agent[Any], value: str | RunState[Any], session: Session | None, streamed: bool +): + config = RunConfig(tracing_disabled=True) + if not streamed: + return await Runner.run(agent, value, session=session, run_config=config) + result = Runner.run_streamed(agent, value, session=session, run_config=config) + async for _ in result.stream_events(): + pass + return result + + +async def _approved_session_state(streamed: bool, session: Session | None = None): + effects: list[int] = [] + + @tool(needs_approval=True) + async def charge(amount: int) -> str: + effects.append(amount) + return "receipt-7" + + model = ScriptedModel( + [ + [get_function_tool_call("charge", '{"amount":7}', call_id="charge-1")], + [get_text_message("done")], + [get_text_message("fresh")], + ] + ) + agent = Agent(name="payment", model=model, tools=[charge]) + session = session if session is not None else _FailingResumeSession() + paused = await _run_session_resume(agent, "charge 7", session, streamed) + state = paused.to_state() + state.approve(state.get_interruptions()[0]) + return agent, model, session, state, effects + + +def _charge_pair(items: list[TResponseInputItem]) -> list[str]: + return [ + str(item.get("type")) + for item in items + if isinstance(item, dict) and item.get("call_id") == "charge-1" + ] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "failing_streamed,retry_streamed", [(False, False), (False, True), (True, False), (True, True)] +) +@pytest.mark.parametrize("round_trip", [False, True], ids=["live", "json"]) +@pytest.mark.parametrize("failure", ["before", "after"], ids=["atomic-failure", "lost-ack"]) +async def test_resumed_session_append_is_recovered_before_next_model( + failing_streamed: bool, retry_streamed: bool, round_trip: bool, failure: str +) -> None: + agent, model, session, state, effects = await _approved_session_state(failing_streamed) + session.failure = failure + with pytest.raises(RuntimeError) as error: + await _run_session_resume(agent, state, session, failing_streamed) + assert error.value is session.error + assert effects == [7] + assert len(model.calls) == 1 + if round_trip: + state = await RunState.from_json(agent, state.to_json()) + + result = await _run_session_resume(agent, state, session, retry_streamed) + assert result.final_output == "done" + assert effects == [7] + expected_pair = ["function_call", "function_call_output"] + assert _charge_pair(await session.get_items()) == expected_pair + assert _charge_pair(result.to_input_list()) == expected_pair + await _run_session_resume(agent, "What was the receipt?", session, retry_streamed) + assert _charge_pair(model.calls[-1].input) == expected_pair + assert "pending_session_write" not in result.to_state().to_json() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("retry_streamed", [False, True]) +@pytest.mark.parametrize("mismatch", ["missing", "different-id", "changed-tail"]) +async def test_resumed_session_append_rejects_ambiguous_recovery( + retry_streamed: bool, mismatch: str +) -> None: + agent, model, session, state, effects = await _approved_session_state(False) + session.failure = "before" + with pytest.raises(RuntimeError, match="session append failed"): + await _run_session_resume(agent, state, session, False) + state = await RunState.from_json(agent, state.to_json()) + supplied_session: Session | None = session + if mismatch == "missing": + supplied_session = None + elif mismatch == "different-id": + supplied_session = SimpleListSession("other", await session.get_items()) + else: + await session.add_items([{"role": "user", "content": "another writer"}]) + before = await session.get_items() + with pytest.raises(UserError, match="pending Session write"): + await _run_session_resume(agent, state, supplied_session, retry_streamed) + assert len(model.calls) == 1 + assert effects == [7] + assert await session.get_items() == before + + +@pytest.mark.asyncio +async def test_resumed_session_append_survives_repeated_failure_and_late_input() -> None: + agent, model, session, state, effects = await _approved_session_state(False) + for _ in range(2): + session.failure = "before" + with pytest.raises(RuntimeError, match="session append failed"): + await _run_session_resume(agent, state, session, False) + state = await RunState.from_json(agent, state.to_json()) + assert len(model.calls) == 1 + assert effects == [7] + state.add_input("What was the receipt?") + result = await _run_session_resume(agent, state, session, True) + assert result.final_output == "done" + stored = await session.get_items() + output_index = next( + i for i, item in enumerate(stored) if item.get("type") == "function_call_output" + ) + late_index = next( + i for i, item in enumerate(stored) if item.get("content") == "What was the receipt?" + ) + assert output_index < late_index + assert effects == [7] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("streamed", [False, True]) +@pytest.mark.parametrize("round_trip", [False, True], ids=["live", "json"]) +async def test_resumed_committed_append_refreshes_compaction_input( + streamed: bool, round_trip: bool, tmp_path: Path +) -> None: + backend = _LostAckSQLiteSession("compaction-recovery", tmp_path / "history.db") + compaction_inputs: list[list[TResponseInputItem]] = [] + compact_enabled = False + + async def compact(**kwargs: Any) -> SimpleNamespace: + items = copy.deepcopy(kwargs["input"]) + compaction_inputs.append(items) + return SimpleNamespace(output=items, usage=None) + + session = OpenAIResponsesCompactionSession( + backend.session_id, + underlying_session=backend, + client=cast(Any, SimpleNamespace(responses=SimpleNamespace(compact=compact))), + compaction_mode="input", + should_trigger_compaction=lambda _: compact_enabled, + ) + try: + agent, model, _, state, effects = await _approved_session_state(streamed, session) + # A normal declined compaction initializes the retained wrapper's history cache. + await session.run_compaction() + assert compaction_inputs == [] + backend.fail_after_commit = True + with pytest.raises(RuntimeError) as error: + await _run_session_resume(agent, state, session, streamed) + assert error.value is backend.error + expected_pair = ["function_call", "function_call_output"] + assert _charge_pair(await backend.get_items(limit=100)) == expected_pair + if round_trip: + state = await RunState.from_json(agent, state.to_json()) + + compact_enabled = True + result = await _run_session_resume(agent, state, session, streamed) + assert result.final_output == "done" + assert effects == [7] + assert len(model.calls) == 2 + assert len(compaction_inputs) == 1 + assert _charge_pair(compaction_inputs[0]) == expected_pair + assert _charge_pair(await backend.get_items(limit=100)) == expected_pair + assert _charge_pair(result.to_input_list()) == expected_pair + assert "pending_session_write" not in result.to_state().to_json() + finally: + backend.close() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("mode", ["input", "auto"]) +async def test_compaction_reload_preserves_session_retrieval_window( + mode: Literal["input", "auto"], tmp_path: Path +) -> None: + backend = _LostAckSQLiteSession( + "bounded-compaction", tmp_path / "history.db", session_settings={"limit": 1} + ) + compaction_inputs: list[list[TResponseInputItem]] = [] + + async def compact(**kwargs: Any) -> SimpleNamespace: + assert "previous_response_id" not in kwargs + items = copy.deepcopy(kwargs["input"]) + compaction_inputs.append(items) + return SimpleNamespace(output=items, usage=None) + + session = OpenAIResponsesCompactionSession( + backend.session_id, + underlying_session=backend, + client=cast(Any, SimpleNamespace(responses=SimpleNamespace(compact=compact))), + compaction_mode=mode, + ) + old_items: list[TResponseInputItem] = [ + {"role": "assistant", "content": f"old message {index}"} for index in range(12) + ] + recovered_item: TResponseInputItem = {"role": "assistant", "content": "committed reply"} + try: + await backend.add_items(old_items) + # The configured window has one candidate, so the default threshold is not met. + await session.run_compaction({"response_id": "unstored-response", "store": False}) + assert compaction_inputs == [] + assert await backend.get_items(limit=100) == old_items + + backend.fail_after_commit = True + with pytest.raises(RuntimeError) as error: + await session.add_items([recovered_item]) + assert error.value is backend.error + assert await backend.get_items(limit=100) == [*old_items, recovered_item] + + await session.run_compaction({"force": True, "store": False}) + assert compaction_inputs == [[recovered_item]] + assert await backend.get_items(limit=100) == [recovered_item] + finally: + backend.close() + + +@pytest.mark.asyncio +async def test_cancelled_compaction_append_preserves_committed_and_surviving_writes() -> None: + appended = asyncio.Event() + wait_for_ack = asyncio.Event() + + class DelayedAckSession(SimpleListSession): + delay_next_ack = True + + async def add_items(self, items: list[TResponseInputItem]) -> None: + await super().add_items(items) + if self.delay_next_ack: + self.delay_next_ack = False + appended.set() + await wait_for_ack.wait() + + backend = DelayedAckSession() + compaction_inputs: list[list[TResponseInputItem]] = [] + + async def compact(**kwargs: Any) -> SimpleNamespace: + items = copy.deepcopy(kwargs["input"]) + compaction_inputs.append(items) + return SimpleNamespace(output=items, usage=None) + + session = OpenAIResponsesCompactionSession( + backend.session_id, + underlying_session=backend, + client=cast(Any, SimpleNamespace(responses=SimpleNamespace(compact=compact))), + compaction_mode="input", + should_trigger_compaction=lambda _: False, + ) + await session.run_compaction() + first_item: TResponseInputItem = {"role": "user", "content": "committed before cancellation"} + newer_item: TResponseInputItem = {"role": "user", "content": "surviving writer"} + first = asyncio.create_task(session.add_items([first_item])) + newer: asyncio.Task[None] | None = None + newer_started = asyncio.Event() + + async def write_newer() -> None: + newer_started.set() + await session.add_items([newer_item]) + + try: + await asyncio.wait_for(appended.wait(), timeout=5) + newer = asyncio.create_task(write_newer()) + await asyncio.wait_for(newer_started.wait(), timeout=5) + assert not newer.done() + first.cancel() + with pytest.raises(asyncio.CancelledError): + await first + await asyncio.wait_for(newer, timeout=5) + await session.run_compaction({"force": True}) + assert compaction_inputs == [[first_item, newer_item]] + assert await backend.get_items() == [first_item, newer_item] + finally: + wait_for_ack.set() + tasks = [first, *([newer] if newer is not None else [])] + for task in tasks: + if not task.done(): + task.cancel() + await asyncio.gather(*tasks, return_exceptions=True) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("streamed", [False, True]) +async def test_resumed_session_append_cancellation_retains_recoverable_state( + streamed: bool, +) -> None: + agent, model, session, state, effects = await _approved_session_state(streamed) + session.block_next_add = True + attempt = asyncio.create_task(_run_session_resume(agent, state, session, streamed)) + try: + await asyncio.wait_for(session.add_started.wait(), timeout=5) + with pytest.raises(UserError, match="pending Session write is already in progress"): + await _run_session_resume(agent, state, session, not streamed) + assert len(model.calls) == 1 + attempt.cancel() + with pytest.raises(asyncio.CancelledError): + await attempt + finally: + session.release_add.set() + if not attempt.done(): + attempt.cancel() + await asyncio.gather(attempt, return_exceptions=True) + + restored = await RunState.from_json(agent, state.to_json()) + result = await _run_session_resume(agent, restored, session, not streamed) + assert result.final_output == "done" + assert effects == [7] + assert _charge_pair(await session.get_items()) == ["function_call", "function_call_output"] + + +@pytest.mark.asyncio +async def test_failed_streamed_result_checkpoint_retains_detached_pending_write() -> None: + agent, model, session, state, effects = await _approved_session_state(True) + session.failure = "before" + result = Runner.run_streamed(agent, state, session=session) + with pytest.raises(RuntimeError, match="session append failed"): + async for _ in result.stream_events(): + pass + snapshot = result.to_state() + payload = snapshot.to_json() + payload["pending_session_write"]["items"][0]["output"] = "changed snapshot" + assert state.to_json()["pending_session_write"]["items"][0]["output"] == "receipt-7" + assert snapshot.to_json()["pending_session_write"]["items"][0]["output"] == "receipt-7" + await _run_session_resume(agent, snapshot, session, False) + assert effects == [7] + assert len(model.calls) == 2 + assert _charge_pair(await session.get_items()) == ["function_call", "function_call_output"] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("invalid", ["old-schema", "batch-shape"]) +async def test_pending_session_write_rejects_invalid_serialized_checkpoint(invalid: str) -> None: + agent, _, session, state, _ = await _approved_session_state(False) + session.failure = "before" + with pytest.raises(RuntimeError): + await _run_session_resume(agent, state, session, False) + payload = state.to_json() + if invalid == "old-schema": + payload["$schemaVersion"] = "1.16" + else: + payload["pending_session_write"]["items"] = "not an item batch" + with pytest.raises(UserError, match="pending Session write is invalid"): + await RunState.from_json(agent, payload) + + +@pytest.mark.asyncio +async def test_resumed_session_append_partial_commit_fails_closed() -> None: + agent, model, session, state, effects = await _approved_session_state(False) + # Two approved calls produce one resumed batch, allowing an actual partial append. + second_call = get_function_tool_call("charge", '{"amount":7}', call_id="charge-2") + model = ScriptedModel( + [ + [get_function_tool_call("charge", '{"amount":7}', call_id="charge-1"), second_call], + [get_text_message("done")], + ] + ) + agent.model = model + session = _FailingResumeSession() + paused = await _run_session_resume(agent, "charge twice", session, False) + state = paused.to_state() + for interruption in state.get_interruptions(): + state.approve(interruption) + session.failure = "partial" + with pytest.raises(RuntimeError, match="session append failed"): + await _run_session_resume(agent, state, session, False) + before = await session.get_items() + restored = await RunState.from_json(agent, state.to_json()) + with pytest.raises(UserError, match="history changed or is ambiguous"): + await _run_session_resume(agent, restored, session, True) + assert effects == [7, 7] + assert len(model.calls) == 1 + assert await session.get_items() == before + + @pytest.mark.asyncio async def test_resolve_interrupted_turn_final_output_short_circuit(monkeypatch) -> None: agent: Agent[dict[str, str]] = make_agent(model=ScriptedModel())