fix(sessions): recover fresh streamed handoffs after session append failures

src/agents/memory/openai_responses_compaction_session.py

@@ -221,7 +221,6 @@ async def run_compaction(
             )
             return
 
-        self._deferred_response_id = None
         logger.debug(
             "compact: start for %s using %s (mode=%s)",
             self._response_id,
@@ -254,6 +253,12 @@ async def run_compaction(
             self._compaction_candidate_items = select_compaction_candidate_items(output_items)
             self._session_items = output_items
 
+        # Clear the deferred marker only now that compaction has actually settled. Clearing it
+        # before the fallible API call/replacement above would let a failed forced compaction
+        # silently lose its "this must be forced" signal: a later retry recomputes `force` from
+        # this marker, so an early clear makes the retry decline work that was still owed.
+        self._deferred_response_id = None
+
         logger.debug(
             "compact: done for %s (mode=%s, output=%s, candidates=%s)",
             self._response_id,

src/agents/run_internal/run_loop.py

@@ -43,6 +43,7 @@
     _detach_data_redacted_error_traceback,
     _is_error_data_redacted,
     _mark_error_data_redacted,
+    _mark_error_to_drain_stream_events,
     _prepare_data_redacted_error,
 )
 from ..guardrail import OutputGuardrailResult
@@ -418,6 +419,7 @@ async def _save_stream_items(
     response_id: str | None,
     update_persisted_count: bool,
     store: bool | None = None,
+    resumed_write_state: RunState | None = None,
 ) -> None:
     if not await _should_persist_stream_items(
         session=session,
@@ -433,6 +435,7 @@ async def _save_stream_items(
         response_id=response_id,
         store=store,
         wrapper=streamed_result.context_wrapper,
+        resumed_write_state=resumed_write_state,
     )
     if update_persisted_count and streamed_result._state is not None:
         streamed_result._current_turn_persisted_item_count = (
@@ -1152,6 +1155,12 @@ async def _save_stream_items_without_count(
                 response_id=response_id,
                 update_persisted_count=False,
                 store=store_setting,
+                resumed_write_state=(
+                    run_state
+                    if run_state is not None
+                    and isinstance(run_state._current_step, NextStepRunAgain)
+                    else None
+                ),
             )
 
         async def _save_max_turns_items(
@@ -1879,23 +1888,34 @@ def _record_max_turns_handler_output(
                     server_conversation_tracker.track_server_items(turn_result.model_response)
 
                 if isinstance(turn_result.next_step, NextStepHandoff):
-                    await _save_stream_items_without_count(
-                        turn_session_items,
-                        turn_result.model_response.response_id,
-                        store_setting,
-                    )
+                    # Resolve any still-in-flight parallel input guardrail before committing the
+                    # handoff transition, so a tripwire or guardrail exception is surfaced instead
+                    # of the state (current_agent, run_state, published events) racing ahead of an
+                    # input guardrail that was still validating the original input.
+                    await input_guardrail_tripwire_triggered_for_stream(streamed_result)
                     current_agent = turn_result.next_step.new_agent
                     if run_state is not None:
                         run_state._current_agent = current_agent
                     _publish_streamed_result_agent(streamed_result, current_agent)
-                    current_span.finish(reset_current=True)
-                    current_span = None
-                    should_run_agent_start_hooks = True
+                    if streamed_result._state is not None:
+                        streamed_result._state._current_step = NextStepRunAgain()
+                    # Queue the agent-transition event before the fallible session append so
+                    # stream consumers observe the transition even if the append later raises.
                     streamed_result._event_queue.put_nowait(
                         AgentUpdatedStreamEvent(new_agent=current_agent)
                     )
-                    if streamed_result._state is not None:
-                        streamed_result._state._current_step = NextStepRunAgain()
+                    try:
+                        await _save_stream_items_without_count(
+                            turn_session_items,
+                            turn_result.model_response.response_id,
+                            store_setting,
+                        )
+                    except BaseException as session_persistence_error:
+                        _mark_error_to_drain_stream_events(session_persistence_error)
+                        raise
+                    current_span.finish(reset_current=True)
+                    current_span = None
+                    should_run_agent_start_hooks = True
 
                     if await _wait_for_streamed_turn_events_and_stop_if_cancelled(streamed_result):
                         break

src/agents/run_internal/session_persistence.py

@@ -556,6 +556,65 @@ def update_run_state_after_resume(
     run_state._current_step = next_step  # type: ignore[assignment]
 
 
+async def _apply_post_write_compaction(
+    session: Session,
+    *,
+    response_id: str | None,
+    store: bool | None,
+    has_local_tool_outputs: bool,
+    wrapper: RunContextWrapper[Any] | None = None,
+) -> None:
+    """Evaluate deferred/forced Responses compaction for a settled session append.
+
+    Shared by the immediate-write path in ``save_result_to_session`` and the checkpoint
+    replay path in ``resume_pending_session_write``, so a batch that only settles later
+    (via a separate resume) still gets the same compaction decision it would have gotten
+    had the original append succeeded inline. ``wrapper`` is the caller's raw (pre-gating)
+    context wrapper; it is used as-is for ``run_compaction`` and re-gated here for
+    ``_defer_compaction``, mirroring the two call sites this helper replaces.
+    """
+    if not response_id or not is_openai_responses_compaction_aware_session(session):
+        return
+
+    if has_local_tool_outputs:
+        defer_compaction = getattr(session, "_defer_compaction", None)
+        if callable(defer_compaction):
+            await _call_session_method(
+                defer_compaction,
+                response_id,
+                store=store,
+                wrapper=_get_session_wrapper(session, wrapper),
+            )
+        logger.debug(
+            "skip: deferring compaction for response %s due to local tool outputs",
+            response_id,
+        )
+        return
+
+    deferred_response_id = None
+    get_deferred = getattr(session, "_get_deferred_compaction_response_id", None)
+    if callable(get_deferred):
+        deferred_response_id = get_deferred()
+    force_compaction = deferred_response_id is not None
+    if force_compaction:
+        logger.debug(
+            "compact: forcing for response %s after deferred %s",
+            response_id,
+            deferred_response_id,
+        )
+    compaction_args: OpenAIResponsesCompactionArgs = {
+        "response_id": response_id,
+        "force": force_compaction,
+    }
+    if store is not None:
+        compaction_args["store"] = store
+    await _call_session_method(
+        session.run_compaction,
+        compaction_args,
+        wrapper=wrapper,
+    )
+
+
 async def save_result_to_session(
     session: Session | None,
     original_input: str | list[TResponseInputItem],
@@ -663,6 +722,10 @@ async def save_result_to_session(
             run_state._current_turn_persisted_item_count = already_persisted + saved_run_items_count
         return saved_run_items_count
 
+    has_local_tool_outputs = any(
+        isinstance(item, ToolCallOutputItem | HandoffOutputItem) for item in new_items
+    )
+
     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")
@@ -673,53 +736,31 @@ async def save_result_to_session(
             "persisted_count": (
                 resumed_write_state._current_turn_persisted_item_count + saved_run_items_count
             ),
+            "response_id": response_id,
+            "store": store,
+            "has_local_tool_outputs": has_local_tool_outputs,
         }
-        await resume_pending_session_write(resumed_write_state, session, wrapper=wrapper)
+        # resume_pending_session_write() applies post-write compaction itself once the
+        # checkpoint settles, whether that happens inline below or on a later, separate
+        # resume -- so it is not repeated after this call returns.
+        await resume_pending_session_write(
+            resumed_write_state,
+            session,
+            wrapper=wrapper,
+            compaction_wrapper=compaction_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
 
-    if response_id and is_openai_responses_compaction_aware_session(session):
-        has_local_tool_outputs = any(
-            isinstance(item, ToolCallOutputItem | HandoffOutputItem) for item in new_items
-        )
-        if has_local_tool_outputs:
-            defer_compaction = getattr(session, "_defer_compaction", None)
-            if callable(defer_compaction):
-                await _call_session_method(
-                    defer_compaction,
-                    response_id,
-                    store=store,
-                    wrapper=wrapper,
-                )
-            logger.debug(
-                "skip: deferring compaction for response %s due to local tool outputs",
-                response_id,
-            )
-            return saved_run_items_count
-
-        deferred_response_id = None
-        get_deferred = getattr(session, "_get_deferred_compaction_response_id", None)
-        if callable(get_deferred):
-            deferred_response_id = get_deferred()
-        force_compaction = deferred_response_id is not None
-        if force_compaction:
-            logger.debug(
-                "compact: forcing for response %s after deferred %s",
-                response_id,
-                deferred_response_id,
-            )
-        compaction_args: OpenAIResponsesCompactionArgs = {
-            "response_id": response_id,
-            "force": force_compaction,
-        }
-        if store is not None:
-            compaction_args["store"] = store
-        await _call_session_method(
-            session.run_compaction,
-            compaction_args,
+    if resumed_write_state is None:
+        await _apply_post_write_compaction(
+            session,
+            response_id=response_id,
+            store=store,
+            has_local_tool_outputs=has_local_tool_outputs,
             wrapper=compaction_wrapper,
         )
 
@@ -764,12 +805,18 @@ async def resume_pending_session_write(
     session: Session | None,
     *,
     wrapper: RunContextWrapper[Any] | None = None,
+    compaction_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.
+
+    ``compaction_wrapper`` defaults to ``wrapper`` when omitted; ``save_result_to_session``
+    passes its own raw (pre-gating) wrapper explicitly so a batch that settles here -- either
+    inline or on a later, separate resume -- gets the exact same post-write Responses
+    compaction decision ``save_result_to_session`` would otherwise have applied itself.
     """
     pending = run_state._pending_session_write
     if pending is None:
@@ -815,6 +862,17 @@ def digests(items: Sequence[TResponseInputItem]) -> list[str]:
             # 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"]
+        # Keep the checkpoint until compaction also settles: if _apply_post_write_compaction
+        # raises below, a later retry must still be able to redo just the compaction step
+        # instead of silently losing it. The append itself is retry-safe (the reconciliation
+        # above detects an already-committed batch and skips re-appending it).
+        await _apply_post_write_compaction(
+            session,
+            response_id=pending.get("response_id"),
+            store=pending.get("store"),
+            has_local_tool_outputs=pending.get("has_local_tool_outputs", False),
+            wrapper=compaction_wrapper if compaction_wrapper is not None else wrapper,
+        )
         run_state._pending_session_write = None
     finally:
         run_state._session_write_in_progress = False

src/agents/run_state.py

@@ -38,7 +38,7 @@
     ProgramOutput,
 )
 from pydantic import BaseModel, StringConstraints, TypeAdapter, ValidationError
-from typing_extensions import TypedDict, TypeVar
+from typing_extensions import NotRequired, TypedDict, TypeVar
 
 from ._tool_identity import (
     FunctionToolLookupKey,
@@ -172,6 +172,14 @@ class _PendingSessionWrite(TypedDict):
     items: list[TResponseInputItem]
     before: list[str] | None
     persisted_count: int
+    # Compaction inputs for the batch this checkpoint is settling, so a later, separate
+    # resume_pending_session_write() call (not the original save_result_to_session() call)
+    # can still apply the same post-write Responses compaction decision. Optional so a
+    # RunState serialized before these fields existed degrades to "skip compaction" on
+    # read instead of raising KeyError.
+    response_id: NotRequired[str | None]
+    store: NotRequired[bool | None]
+    has_local_tool_outputs: NotRequired[bool]
 
 
 def _default_run_state_validation_error(
@@ -4360,11 +4368,18 @@ async def _build_run_state_from_json(
     if pending_write is not None:
         from .run_internal.run_steps import NextStepInterruption, NextStepRunAgain
 
+        required_pending_write_keys = {"session_id", "items", "before", "persisted_count"}
+        # response_id/store/has_local_tool_outputs carry the compaction inputs needed to replay
+        # deferred/forced Responses compaction on a later, separate resume; they are optional so
+        # a RunState serialized before these fields existed (same, unreleased schema version)
+        # still round-trips.
+        optional_pending_write_keys = {"response_id", "store", "has_local_tool_outputs"}
         if (
             (schema_major, schema_minor) < (1, 17)
             or not isinstance(state._current_step, NextStepRunAgain | NextStepInterruption)
             or not isinstance(pending_write, dict)
-            or set(pending_write) != {"session_id", "items", "before", "persisted_count"}
+            or not required_pending_write_keys <= set(pending_write)
+            or not set(pending_write) <= required_pending_write_keys | optional_pending_write_keys
             or not isinstance(pending_write.get("session_id"), str)
             or not isinstance(pending_write.get("items"), list)
             or not pending_write["items"]
@@ -4378,6 +4393,20 @@ async def _build_run_state_from_json(
             )
             or type(pending_write.get("persisted_count")) is not int
             or pending_write["persisted_count"] < 0
+            or (
+                "response_id" in pending_write
+                and pending_write["response_id"] is not None
+                and not isinstance(pending_write["response_id"], str)
+            )
+            or (
+                "store" in pending_write
+                and pending_write["store"] is not None
+                and not isinstance(pending_write["store"], bool)
+            )
+            or (
+                "has_local_tool_outputs" in pending_write
+                and not isinstance(pending_write["has_local_tool_outputs"], bool)
+            )
         ):
             raise validation_error_factory("Run state pending Session write is invalid", UserError)
         state._pending_session_write = copy.deepcopy(cast(_PendingSessionWrite, pending_write))