fix(sessions): dispose of the held batch when the run ends, not when its last step is chosen

Choosing a terminal step is not the same as ending the run. Validation, the
final-output hooks, the output guardrails and the final save all run after that
choice, any of them can raise, and a run that raises may still be retried or
reattached with the approved tool's call and output reachable only through the
held batch. Consuming the batch at the choice threw it away on every one of
those failures.

Four sites carried that ordering, and they are the whole class: the max-turn
handler finalization, the detached final output in both runners, and the
detached final output on the resumed streamed loop. Each now disposes of the
batch once finalization has completed, with a tripwire handled separately as
the decided blocked outcome it is. The five remaining disposal sites are
deliberate ones that follow an outcome already decided, and they are unchanged.

Also: a re-interruption no longer overwrites the storage setting the parked
response was produced under. Presence of the key decides, not its truthiness, so
an ordinary ``store=None`` park keeps its own setting and the settle resolves
that response's compaction mode from the right turn. ``response_id`` follows the
same rule for the same reason.

Tests pin the failure of each finalization stage in both runners, and the park
storage settings across None, False and True.

Author: dixso

src/agents/run.py

@@ -1973,12 +1973,6 @@ async def _save_max_turns_handler_output(
 
                     try:
                         if isinstance(turn_result.next_step, NextStepFinalOutput):
-                            if session is None and run_state is not None:
-                                # A detached completion has no Session to settle against
-                                # and the run ends here, so the batch is discarded
-                                # rather than left to invalidate the completed run's
-                                # checkpoint. Mirrors the resumed final exit.
-                                take_held_session_write(run_state)
                             if run_state is not None and _has_output_guardrails(
                                 current_agent, run_config
                             ):
@@ -2106,6 +2100,16 @@ async def _save_max_turns_handler_output(
                             if run_state is not None:
                                 run_state._terminal_unrecoverable = False
 
+                            if session is None and run_state is not None:
+                                # A detached completion has no Session to settle against
+                                # and the run ends here, so the batch is discarded
+                                # rather than left to invalidate the completed run's
+                                # checkpoint. Only here, though: the guardrails and the
+                                # final save above can raise, and a run that raises may
+                                # still be retried or reattached, with the executed
+                                # tool's call and output reachable only through it.
+                                take_held_session_write(run_state)
+
                             # Ensure starting_input is not None and not RunState
                             final_output_result_input: str | list[TResponseInputItem] = (
                                 normalized_starting_input

src/agents/run_internal/run_loop.py

@@ -784,12 +784,13 @@ async def finalize_max_turns_handler_output(
     A max-turn handler ends the run, so a held Session write still standing here has
     no later gate-legal exit to settle it: it is discarded, exactly as a detached
     completion discards it, so the finished run's checkpoint stays loadable and both
-    runners report the same terminal state.
+    runners report the same terminal state. The discard waits for the run to actually
+    end, which is either a completed finalization or a decided blocked outcome:
+    validation, the final-output hooks and the guardrails can all raise, and a run that
+    raises may still be retried or reattached, with the executed tool's call and output
+    reachable only through this batch.
     """
     validated_output = validate_handler_final_output(agent, output)
-    # Only past the validation does the handler actually end the run; discarding above
-    # it would throw the batch away on a rejection the streamed runner survives.
-    take_held_session_write(run_state)
     output_text = format_final_output_text(agent, validated_output)
     synthesized_item = create_message_output_item(agent, output_text)
 
@@ -805,6 +806,9 @@ async def finalize_max_turns_handler_output(
             output_guardrail_results,
         )
     except OutputGuardrailTripwireTriggered:
+        # A blocked outcome is decided and nothing of the withheld batch may reach the
+        # Session, exactly as every other tripwire path disposes of it.
+        take_held_session_write(run_state)
         raise
     except Exception as guardrail_error:
         guardrail_error_is_redacted = _is_error_data_redacted(guardrail_error)
@@ -821,6 +825,7 @@ async def finalize_max_turns_handler_output(
 
     if redacted_persistence_error is not None:
         raise redacted_persistence_error from None
+    take_held_session_write(run_state)
     return validated_output, synthesized_item
 
 
@@ -1552,11 +1557,6 @@ async def _save_max_turns_items(
                         continue
 
                     if isinstance(turn_result.next_step, NextStepFinalOutput):
-                        if session is None:
-                            # A detached final output has no Session to settle against
-                            # and the run ends here, so the batch is discarded rather
-                            # than left to invalidate the completed run's checkpoint.
-                            take_held_session_write(run_state)
                         await _finalize_streamed_final_output(
                             streamed_result=streamed_result,
                             agent=current_agent,
@@ -1577,6 +1577,16 @@ async def _save_max_turns_items(
                         )
                         if streamed_result._stored_exception is not None:
                             break
+                        if session is None:
+                            # A detached final output has no Session to settle against
+                            # and the run ends here, so the batch is discarded rather
+                            # than left to invalidate the completed run's checkpoint.
+                            # Only here, though: the finalization above runs the hooks,
+                            # the guardrails and the final save, any of which can raise,
+                            # and a run that raises may still be retried or reattached
+                            # with the executed tool's call and output reachable only
+                            # through this batch.
+                            take_held_session_write(run_state)
                         run_state._current_step = None
                         break
 
@@ -2056,12 +2066,6 @@ def _record_max_turns_handler_output(
                     if await _wait_for_streamed_turn_events_and_stop_if_cancelled(streamed_result):
                         break
                 elif isinstance(turn_result.next_step, NextStepFinalOutput):
-                    if session is None:
-                        # A detached completion has no Session to settle against and
-                        # the run ends here, so the batch is discarded rather than
-                        # left to invalidate the completed run's checkpoint. Mirrors
-                        # the resumed final exit.
-                        take_held_session_write(run_state)
                     await _finalize_streamed_final_output(
                         streamed_result=streamed_result,
                         agent=current_agent,
@@ -2078,6 +2082,15 @@ def _record_max_turns_handler_output(
                     )
                     if streamed_result._stored_exception is not None:
                         break
+                    if session is None:
+                        # A detached completion has no Session to settle against and
+                        # the run ends here, so the batch is discarded rather than
+                        # left to invalidate the completed run's checkpoint. Only here,
+                        # though: the finalization above runs the hooks, the guardrails
+                        # and the final save, any of which can raise, and a run that
+                        # raises may still be retried or reattached with the executed
+                        # tool's call and output reachable only through this batch.
+                        take_held_session_write(run_state)
                     if run_state is not None:
                         run_state._current_step = None
                     break

src/agents/run_internal/session_persistence.py

@@ -1124,11 +1124,15 @@ def defer_interrupted_session_write(
         "held": True,
         # The response the withheld batch belongs to, so the settle can run the same
         # compaction bookkeeping the ordinary persistence path runs for it. An extend
-        # keeps the original response: the batch is that response's write.
-        "response_id": (pending.get("response_id") if pending is not None else None) or response_id,
-        "store": (pending.get("store") if pending is not None else None)
-        if (pending is not None and pending.get("store") is not None)
-        else store,
+        # keeps the original response: the batch is that response's write, and the
+        # settle resolves its compaction mode from that response's own storage setting.
+        # Presence decides, not truthiness: a park under the ordinary ``store=None``
+        # records a real value, and letting a re-interruption's setting overwrite it
+        # would resolve the original response's compaction mode from the wrong turn.
+        "response_id": pending["response_id"]
+        if (pending is not None and "response_id" in pending)
+        else response_id,
+        "store": pending["store"] if (pending is not None and "store" in pending) else store,
     }
     run_state._pending_session_write = record
 

tests/test_deferred_interrupted_session_write.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import json
+from collections.abc import Callable
 from dataclasses import replace
 from typing import Any, Literal, cast
 
@@ -1347,6 +1348,145 @@ async def test_the_settled_count_matches_what_the_append_actually_wrote() -> Non
     assert count == len(await session.get_items())
 
 
+class _FinalOutputHookFailure(RunHooks[Any]):
+    """Fail the run at the final-output hook, after the terminal step is decided."""
+
+    async def on_agent_end(self, context: Any, agent: Any, output: Any) -> None:
+        raise RuntimeError("final output hook failed")
+
+
+@pytest.mark.asyncio
+async def test_a_failed_max_turns_finalization_keeps_the_held_record() -> None:
+    # The batch is disposed of when the run actually ends, not when the terminal step
+    # is chosen. Validation, the final-output hooks and the output guardrails all run
+    # after that choice and all can raise, and a run that raises may still be retried
+    # or reattached with the executed tool's call and output reachable only here.
+    from agents.run_internal.run_loop import finalize_max_turns_handler_output
+
+    session = SimpleListSession()
+    agent = _make_deferring_agent()
+    state = await _parked_and_approved(agent, session, streamed=False)
+    assert state._pending_session_write is not None
+
+    async def _no_save(items: list[Any]) -> None:
+        return None
+
+    with pytest.raises(RuntimeError):
+        await finalize_max_turns_handler_output(
+            agent=agent,
+            hooks=_FinalOutputHookFailure(),
+            run_config=RunConfig(tracing_disabled=True),
+            output="stopped at max turns",
+            context_wrapper=RunContextWrapper(context=None),
+            output_guardrail_results=[],
+            save_items_after_guardrails=_no_save,
+            include_in_history=False,
+            run_state=state,
+        )
+
+    assert state._pending_session_write is not None
+
+
+def _make_deferring_agent_with_a_turn_after_the_resume() -> Agent:
+    """A gated write whose resume runs one more model turn before finishing.
+
+    The extra turn moves the final output past the resumed boundary and onto the main
+    loop, which owns its own detached-completion disposal.
+    """
+    return Agent(
+        name="deferred repro (turn after resume)",
+        instructions="Always call write_thing.",
+        model=ScriptedModel(
+            [
+                ModelStep(output=[function_call("look_up", {"query": "x"}, call_id="call_LOOKUP")]),
+                ModelStep(
+                    output=[function_call("write_thing", {"query": "x"}, call_id="call_PARKED")]
+                ),
+                ModelStep(output=[function_call("look_up", {"query": "y"}, call_id="call_AFTER")]),
+                ModelStep(output=[assistant_message("done")]),
+            ]
+        ),
+        tools=[look_up, write_thing],
+        output_guardrails=[always_fine],
+        tool_use_behavior=_DEFERRING_BEHAVIOR,
+    )
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("streamed", [False, True])
+@pytest.mark.parametrize(
+    "make_agent",
+    [_make_deferring_agent, _make_deferring_agent_with_a_turn_after_the_resume],
+    ids=["final-on-the-resumed-turn", "final-on-a-later-turn"],
+)
+async def test_a_failed_detached_completion_keeps_the_held_record(
+    streamed: bool, make_agent: Callable[[], Agent]
+) -> None:
+    # A detached completion discards the batch because the run ends there, but only
+    # once it has ended: the guardrails and the final save run after the terminal step
+    # is chosen, and a failure there leaves a checkpoint whose reattach is the batch's
+    # only remaining way into the Session.
+    from agents import output_guardrail
+
+    @output_guardrail
+    async def _fails(ctx: Any, agent: Agent, output: Any) -> GuardrailFunctionOutput:
+        raise RuntimeError("output guardrail failed")
+
+    session = SimpleListSession()
+    agent = make_agent()
+    state = await _parked_and_approved(agent, session, streamed=streamed)
+    assert state._pending_session_write is not None
+    agent.output_guardrails = [*agent.output_guardrails, _fails]
+
+    with pytest.raises(RuntimeError):
+        await _run(agent, state, None, streamed=streamed)
+
+    assert state._pending_session_write is not None
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("parked_store", [None, False, True])
+async def test_a_re_park_keeps_the_storage_setting_the_response_was_produced_under(
+    parked_store: bool | None,
+) -> None:
+    # The batch belongs to the parked response, and the settle resolves that
+    # response's compaction mode from this value. Presence decides, not truthiness: a
+    # park under the ordinary ``store=None`` records a real setting, and a
+    # re-interruption under a different one must not overwrite it.
+    from agents.run_internal.session_persistence import defer_interrupted_session_write
+
+    class _Session:
+        session_id = "s1"
+
+    state = object.__new__(RunState)
+    state._pending_session_write = {
+        "session_id": "s1",
+        "items": [
+            {"type": "function_call", "call_id": "call_PARKED", "name": "t", "arguments": "{}"}
+        ],
+        "before": None,
+        "persisted_count": 1,
+        "held": True,
+        "response_id": "resp_parked",
+        "store": parked_store,
+    }
+    state._current_turn_persisted_item_count = 0
+    state._reasoning_item_id_policy = None
+
+    defer_interrupted_session_write(
+        state,
+        _Session(),  # type: ignore[arg-type]
+        run_items=[],
+        reasoning_item_id_policy=None,
+        response_id="resp_reinterrupted",
+        store=not parked_store,
+    )
+
+    assert state._pending_session_write is not None
+    assert state._pending_session_write["store"] is parked_store
+    assert state._pending_session_write["response_id"] == "resp_parked"
+
+
 class _CompactionRecordingSession(SimpleListSession):
     """Record the compaction bookkeeping a compaction-aware backend expects."""