From 92b3fb008e011ebf0588e98c35ae6ef3773f9ffc Mon Sep 17 00:00:00 2001 From: Kazuhiro Sera Date: Mon, 24 Aug 2026 09:40:29 +0900 Subject: [PATCH] fix: preserve serialized approval resume ownership --- src/agents/run_state.py | 178 ++++++++++++- tests/test_agent_runner_streamed.py | 389 +++++++++++++++++++++++++++- 2 files changed, 557 insertions(+), 10 deletions(-) diff --git a/src/agents/run_state.py b/src/agents/run_state.py index b3a0af49b4..e00a73bdce 100644 --- a/src/agents/run_state.py +++ b/src/agents/run_state.py @@ -182,6 +182,7 @@ def _default_run_state_validation_error( CURRENT_SCHEMA_VERSION = "1.17" _PROGRAMMATIC_TOOL_CALLING_MIN_SCHEMA_VERSION = "1.13" _HOSTED_MCP_APPROVALS_MIN_SCHEMA_VERSION = "1.14" +_CURRENT_RESPONSE_OWNERSHIP_MIN_SCHEMA_VERSION = "1.17" # Keep this mapping in chronological order. Every schema bump must add a one-line summary here. SCHEMA_VERSION_SUMMARIES: dict[str, str] = { "1.0": "Initial RunState snapshot format for HITL pause/resume flows.", @@ -213,7 +214,10 @@ def _default_run_state_validation_error( "Persists Docker network-isolation state and lets an exact call approval decision " "override a sticky decision for the same tool." ), - "1.17": "Persists Docker container labels across sandbox resume and replacement.", + "1.17": ( + "Persists Docker container labels and current-response generated-item ownership across " + "resume flows." + ), } SUPPORTED_SCHEMA_VERSIONS = frozenset(SCHEMA_VERSION_SUMMARIES) @@ -1457,6 +1461,46 @@ def _take_unused(candidates: deque[int] | None) -> int | None: indexes.append(session_index) return indexes + def _current_response_generated_item_ownership( + self, + generated_items: Sequence[RunItem], + ) -> dict[str, Any] | None: + """Record the response range and approval occurrences from live item identities.""" + from .run_internal.run_steps import NextStepInterruption + + if self._last_processed_response is None: + return None + if not isinstance(self._current_step, NextStepInterruption): + return None + + processed_items = self._last_processed_response.new_items + interruptions = self._current_step.interruptions + if not processed_items or not interruptions or len(processed_items) > len(generated_items): + return None + + candidate_starts = [ + start + for start in range(len(generated_items) - len(processed_items) + 1) + if all( + generated_items[start + offset] is item + for offset, item in enumerate(processed_items) + ) + ] + if len(candidate_starts) != 1: + return None + + start = candidate_starts[0] + indexes_by_identity: dict[int, list[int]] = {} + for index in range(start + len(processed_items), len(generated_items)): + indexes_by_identity.setdefault(id(generated_items[index]), []).append(index) + interruption_indexes: list[int] = [] + for item in interruptions: + indexes = indexes_by_identity.pop(id(item), []) + if len(indexes) != 1: + return None + interruption_indexes.append(indexes[0]) + return {"start": start, "end": len(generated_items), "interruptions": interruption_indexes} + def _serialize_context_payload( self, *, @@ -1805,6 +1849,14 @@ def to_json( "generated_session_item_indexes": self._generated_session_item_indexes(generated_items), } + current_response_generated_item_ownership = self._current_response_generated_item_ownership( + generated_items + ) + if current_response_generated_item_ownership is not None: + result["current_response_generated_item_ownership"] = ( + current_response_generated_item_ownership + ) + result["generated_items"] = [ self._serialize_item(item, agent_identity_keys_by_id=agent_identity_keys_by_id) for item in generated_items @@ -4250,6 +4302,24 @@ async def _build_run_state_from_json( current_step_data.get("data", {}).get("llm_end_hooks_started", True) ), ) + _restore_current_response_item_identities( + state, + serialized_generated_items=serialized_generated_items, + generated_source_indexes=generated_source_indexes, + last_processed_response_data=last_processed_response_data, + current_step_data=current_step_data, + current_response_generated_item_ownership=( + state_json.get("current_response_generated_item_ownership") + if (schema_major, schema_minor) + >= tuple( + int(part) + for part in _CURRENT_RESPONSE_OWNERSHIP_MIN_SCHEMA_VERSION.split( + ".", maxsplit=1 + ) + ) + else None + ), + ) if state._current_step.response_accepted: state._clear_generated_items_last_processed_marker() for approval_item in state._current_step.interruptions: @@ -5396,6 +5466,112 @@ def _deserialize_items_with_source_indexes( return items, source_indexes +def _restore_current_response_item_identities( + state: RunState[Any], + *, + serialized_generated_items: Any, + generated_source_indexes: Sequence[int], + last_processed_response_data: Any, + current_step_data: Mapping[str, Any], + current_response_generated_item_ownership: Any, +) -> None: + """Relink one response from explicit generated-item ownership after deserialization.""" + from .run_internal.run_steps import NextStepInterruption + + processed_response = state._last_processed_response + if processed_response is None: + return + current_step = state._current_step + if not isinstance(current_step, NextStepInterruption): + return + if not isinstance(serialized_generated_items, list): + return + if not isinstance(last_processed_response_data, Mapping): + return + + serialized_processed_items = last_processed_response_data.get("new_items") + if not isinstance(serialized_processed_items, list) or not serialized_processed_items: + return + if len(processed_response.new_items) != len(serialized_processed_items): + return + current_step_payload = current_step_data.get("data") + if not isinstance(current_step_payload, Mapping): + return + serialized_interruptions = current_step_payload.get("interruptions") + if not isinstance(serialized_interruptions, list) or not serialized_interruptions: + return + if len(current_step.interruptions) != len(serialized_interruptions): + return + ownership = current_response_generated_item_ownership + if not isinstance(ownership, Mapping): + return + source_start = ownership.get("start") + source_end = ownership.get("end") + interruption_indexes = ownership.get("interruptions") + if type(source_start) is not int or type(source_end) is not int: + return + if source_start < 0 or source_end != len(serialized_generated_items): + return + processed_end = source_start + len(serialized_processed_items) + # Handoff filters can clear prior items without resetting the model turn count. + if processed_end > source_end: + return + if not isinstance(interruption_indexes, list): + return + if len(interruption_indexes) != len(serialized_interruptions) or any( + type(index) is not int or index < processed_end or index >= source_end + for index in interruption_indexes + ): + return + if len(set(interruption_indexes)) != len(interruption_indexes): + return + source_indexes = [*range(source_start, processed_end), *interruption_indexes] + serialized_current_response_items = [*serialized_processed_items, *serialized_interruptions] + if any( + serialized_generated_items[source_index] != expected_item + for source_index, expected_item in zip( + source_indexes, + serialized_current_response_items, + strict=True, + ) + ): + return + + restored_indexes_by_source: dict[int, list[int]] = {} + for restored_index, source_index in enumerate(generated_source_indexes): + restored_indexes_by_source.setdefault(source_index, []).append(restored_index) + + restored_current_response_items: list[RunItem] = [] + for source_index in range(source_start, source_end): + restored_indexes = restored_indexes_by_source.get(source_index) + if restored_indexes is None or len(restored_indexes) != 1: + return + restored_current_response_items.append(state._generated_items[restored_indexes[0]]) + + processed_item_count = len(serialized_processed_items) + restored_processed_items = restored_current_response_items[:processed_item_count] + restored_interruptions = [ + restored_current_response_items[index - source_start] for index in interruption_indexes + ] + if not all(isinstance(item, ToolApprovalItem) for item in restored_interruptions): + return + + # The complete current response must be the same terminal suffix in both histories. + session_start = len(state._session_items) - len(restored_current_response_items) + if session_start < 0 or any( + generated_item is not session_item + for generated_item, session_item in zip( + restored_current_response_items, + state._session_items[session_start:], + strict=True, + ) + ): + return + + processed_response.new_items = restored_processed_items + current_step.interruptions = cast(list[ToolApprovalItem], restored_interruptions) + + def _clone_original_input(original_input: str | list[Any]) -> str | list[Any]: """Return a deep copy of the original input so later mutations don't leak into saved state.""" if isinstance(original_input, str): diff --git a/tests/test_agent_runner_streamed.py b/tests/test_agent_runner_streamed.py index a38424a8b3..9664d65513 100644 --- a/tests/test_agent_runner_streamed.py +++ b/tests/test_agent_runner_streamed.py @@ -2572,16 +2572,363 @@ async def run_once(input_value: Any) -> Any: @pytest.mark.parametrize("mode", ["non_streamed", "streamed"]) +@pytest.mark.parametrize("attach_session", [False, True], ids=["without-session", "with-session"]) +@pytest.mark.parametrize("mixed_tool_position", [None, "before", "after"]) @pytest.mark.asyncio -async def test_ambiguous_serialized_approval_state_fails_before_tool_execution( +async def test_serialized_later_turn_approval_with_output_guardrail_resumes( mode: str, + attach_session: bool, + mixed_tool_position: str | None, ) -> None: - tool_calls = 0 + tool_calls = {"normal": 0, "approval": 0} + + @function_tool(name_override="normal_tool") + def normal_tool() -> str: + tool_calls["normal"] += 1 + return "normal-result" @function_tool(name_override="approval_tool", needs_approval=True) + def approval_tool() -> str: + tool_calls["approval"] += 1 + return "approved-result" + + def output_guardrail( + _context: RunContextWrapper[Any], + _agent: Agent[Any], + _output: Any, + ) -> GuardrailFunctionOutput: + return GuardrailFunctionOutput(output_info=None, tripwire_triggered=False) + + approval_response = [get_function_tool_call("approval_tool", "{}", call_id="call-approved")] + if mixed_tool_position is not None: + extra_call = get_function_tool_call("normal_tool", "{}", call_id="call-extra") + approval_response.insert(0 if mixed_tool_position == "before" else 1, extra_call) + expected_normal_calls = 1 if mixed_tool_position is None else 2 + model = ScriptedModel( + [ + [get_function_tool_call("normal_tool", "{}", call_id="call-normal")], + approval_response, + [get_text_message("done")], + ] + ) + agent = Agent( + name="test", + model=model, + tools=[normal_tool, approval_tool], + output_guardrails=[OutputGuardrail(guardrail_function=output_guardrail)], + ) + session = SimpleListSession() if attach_session else None + + async def run_once(input_value: Any) -> Any: + if mode == "non_streamed": + return await Runner.run(agent, input_value, session=session) + result = Runner.run_streamed(agent, input_value, session=session) + await consume_stream(result) + return result + + first = await run_once("Use normal_tool, then approval_tool") + state = first.to_state() + assert state._current_turn == 2 + assert [item.tool_name for item in first.interruptions] == ["approval_tool"] + assert tool_calls == {"normal": expected_normal_calls, "approval": 0} + + serialized = state.to_json() + assert serialized["current_response_generated_item_ownership"] == { + "start": 2, + "end": 4 if mixed_tool_position is None else 6, + "interruptions": [ + 3 if mixed_tool_position is None else (5 if mixed_tool_position == "before" else 4) + ], + } + + released_payload = json.loads(json.dumps(serialized)) + released_payload["$schemaVersion"] = "1.16" + released = await RunState.from_string(agent, json.dumps(released_payload)) + released.approve(released.get_interruptions()[0]) + with pytest.raises(UserError, match="current response boundary cannot be proven"): + await run_once(released) + assert tool_calls == {"normal": expected_normal_calls, "approval": 0} + + malformed_payload = json.loads(json.dumps(serialized)) + malformed_payload["current_response_generated_item_ownership"]["interruptions"][0] = True + malformed = await RunState.from_string(agent, json.dumps(malformed_payload)) + malformed.approve(malformed.get_interruptions()[0]) + with pytest.raises(UserError, match="current response boundary cannot be proven"): + await run_once(malformed) + assert tool_calls == {"normal": expected_normal_calls, "approval": 0} + + restored_once = await RunState.from_string(agent, json.dumps(serialized)) + reserialized = restored_once.to_json() + assert ( + reserialized["current_response_generated_item_ownership"] + == serialized["current_response_generated_item_ownership"] + ) + + restored = await RunState.from_string(agent, json.dumps(reserialized)) + assert restored._last_processed_response is not None + assert any( + restored._last_processed_response.new_items[0] is item for item in restored._generated_items + ) + restored.approve(restored.get_interruptions()[0]) + + resumed = await run_once(restored) + + assert resumed.final_output == "done" + assert tool_calls == {"normal": expected_normal_calls, "approval": 1} + if session is not None: + saved_items = await session.get_items() + saved_tool_items = [ + item + for item in saved_items + if isinstance(item, dict) + and item.get("type") in {"function_call", "function_call_output"} + ] + expected_tool_items = [ + ("function_call", "call-normal"), + ("function_call_output", "call-normal"), + ("function_call", "call-approved"), + ("function_call_output", "call-approved"), + ] + if mixed_tool_position is not None: + expected_tool_items.insert( + 2 if mixed_tool_position == "before" else 3, ("function_call", "call-extra") + ) + expected_tool_items.insert(4, ("function_call_output", "call-extra")) + assert [(item.get("type"), item.get("call_id")) for item in saved_tool_items] == ( + expected_tool_items + ) + + +@pytest.mark.parametrize("mode", ["non_streamed", "streamed"]) +@pytest.mark.parametrize( + "session_ownership", + [ + "valid", + "missing-prefix", + "missing", + "invalid", + "missing-current", + "prefix-anchor", + "nonterminal-session", + ], +) +@pytest.mark.asyncio +async def test_serialized_mixed_approval_guardrail_preserves_only_accepted_outputs( + mode: str, + session_ownership: str, +) -> None: + tool_calls = {"normal": 0, "approval": 0} + + @function_tool + def normal_tool() -> str: + tool_calls["normal"] += 1 + return "accepted-result" if tool_calls["normal"] == 1 else "sibling-secret" + + @function_tool(needs_approval=True) + def approval_tool() -> str: + tool_calls["approval"] += 1 + return "approved-secret" + + model = ScriptedModel( + [ + [get_function_tool_call("normal_tool", "{}", call_id="call-normal")], + [ + get_function_tool_call("normal_tool", "{}", call_id="call-extra"), + get_function_tool_call("approval_tool", "{}", call_id="call-approved"), + ], + ] + ) + agent = Agent( + name="test", + model=model, + tools=[normal_tool, approval_tool], + tool_use_behavior={"stop_at_tool_names": ["approval_tool"]}, + output_guardrails=[ + OutputGuardrail( + guardrail_function=lambda *_: GuardrailFunctionOutput( + output_info=None, tripwire_triggered=True + ) + ) + ], + ) + session = SimpleListSession() + + async def run_once(input_value: Any) -> Any: + if mode == "non_streamed": + return await Runner.run(agent, input_value, session=session) + result = Runner.run_streamed(agent, input_value, session=session) + await consume_stream(result) + return result + + first = await run_once("Use the tools") + payload = first.to_state().to_json() + if session_ownership == "missing": + payload.pop("generated_session_item_indexes") + elif session_ownership == "missing-prefix": + payload["generated_session_item_indexes"][0] = None + elif session_ownership == "invalid": + payload["generated_session_item_indexes"][0] = True + elif session_ownership == "missing-current": + current_start = payload["current_response_generated_item_ownership"]["start"] + payload["generated_session_item_indexes"][current_start] = None + elif session_ownership == "prefix-anchor": + earlier_copies = json.loads(json.dumps(payload["last_processed_response"]["new_items"])) + payload["session_items"][:0] = earlier_copies + payload["generated_session_item_indexes"] = [ + index + len(earlier_copies) if index is not None else None + for index in payload["generated_session_item_indexes"] + ] + current_start = payload["current_response_generated_item_ownership"]["start"] + for offset in range(len(earlier_copies)): + payload["generated_session_item_indexes"][current_start + offset] = offset + elif session_ownership == "nonterminal-session": + payload["session_items"].append(json.loads(json.dumps(payload["session_items"][0]))) + restored = await RunState.from_string(agent, json.dumps(payload)) + restored = await RunState.from_string(agent, restored.to_string()) + restored.approve(restored.get_interruptions()[0]) + if session_ownership not in {"valid", "missing-prefix"}: + saved_before = json.loads(json.dumps(await session.get_items())) + state_session_before = restored.to_json()["session_items"] + with pytest.raises(UserError, match="current response boundary cannot be proven"): + await run_once(restored) + assert tool_calls == {"normal": 2, "approval": 0} + assert await session.get_items() == saved_before + assert restored.to_json()["session_items"] == state_session_before + assert "accepted-result" in json.dumps(state_session_before) + return + with pytest.raises(OutputGuardrailTripwireTriggered): + await run_once(restored) + + assert tool_calls == {"normal": 2, "approval": 1} + outputs = [ + (item["call_id"], item["output"]) + for item in await session.get_items() + if isinstance(item, dict) and item.get("type") == "function_call_output" + ] + assert outputs == [ + ("call-normal", "accepted-result"), + ("call-extra", run_loop._OUTPUT_GUARDRAIL_BLOCKED_TOOL_OUTPUT), + ("call-approved", run_loop._OUTPUT_GUARDRAIL_BLOCKED_TOOL_OUTPUT), + ] + serialized = restored.to_string() + assert "accepted-result" in serialized + assert "sibling-secret" not in serialized + assert "approved-secret" not in serialized + + +@pytest.mark.parametrize("mode", ["non_streamed", "streamed"]) +@pytest.mark.parametrize("attach_session", [False, True], ids=["without-session", "with-session"]) +@pytest.mark.parametrize("tripwire", [False, True], ids=["accepted", "blocked"]) +@pytest.mark.asyncio +async def test_serialized_filtered_handoff_approval_with_empty_prefix_resumes( + mode: str, + attach_session: bool, + tripwire: bool, +) -> None: + tool_calls = 0 + + @function_tool(needs_approval=True) def approval_tool() -> str: nonlocal tool_calls tool_calls += 1 + return "approved-secret" + + def clear_generated_history(data: HandoffInputData) -> HandoffInputData: + return HandoffInputData( + input_history=data.input_history, + pre_handoff_items=(), + new_items=(), + run_context=data.run_context, + ) + + target = Agent( + name="target", + model=ScriptedModel( + [[get_function_tool_call("approval_tool", "{}", call_id="call-approved")]] + ), + tools=[approval_tool], + tool_use_behavior="stop_on_first_tool", + output_guardrails=[ + OutputGuardrail( + guardrail_function=lambda *_: GuardrailFunctionOutput( + output_info=None, tripwire_triggered=tripwire + ) + ) + ], + ) + starting = Agent( + name="starting", + model=ScriptedModel([[get_handoff_tool_call(target)]]), + handoffs=[handoff(target, input_filter=clear_generated_history)], + ) + session = SimpleListSession() if attach_session else None + + async def run_once(input_value: Any) -> Any: + if mode == "non_streamed": + return await Runner.run(starting, input_value, session=session) + result = Runner.run_streamed(starting, input_value, session=session) + await consume_stream(result) + return result + + first = await run_once("Transfer and run approval_tool") + state = first.to_state() + assert state._current_turn == 2 + assert tool_calls == 0 + ownership = {"start": 0, "end": 2, "interruptions": [1]} + assert state.to_json()["current_response_generated_item_ownership"] == ownership + restored = await RunState.from_string(starting, state.to_string()) + assert restored.to_json()["current_response_generated_item_ownership"] == ownership + restored = await RunState.from_string(starting, restored.to_string()) + restored.approve(restored.get_interruptions()[0]) + + if tripwire: + with pytest.raises(OutputGuardrailTripwireTriggered): + await run_once(restored) + assert "approved-secret" not in restored.to_string() + else: + resumed = await run_once(restored) + assert resumed.final_output == "approved-secret" + assert tool_calls == 1 + if session is not None: + saved_outputs = [ + item + for item in await session.get_items() + if isinstance(item, dict) and item.get("type") == "function_call_output" + ] + assert [(item["call_id"], item["output"]) for item in saved_outputs] == [ + ( + "call-approved", + run_loop._OUTPUT_GUARDRAIL_BLOCKED_TOOL_OUTPUT if tripwire else "approved-secret", + ) + ] + + +@pytest.mark.parametrize("mode", ["non_streamed", "streamed"]) +@pytest.mark.parametrize( + "corruption", + [ + "missing-prefix", + "nonterminal-indexes", + "invalid-interruption-index", + "processed-mismatch", + "interruption-mismatch", + ], +) +@pytest.mark.asyncio +async def test_ambiguous_serialized_approval_state_fails_before_tool_execution( + mode: str, + corruption: str, +) -> None: + tool_calls = {"normal": 0, "approval": 0} + + @function_tool(name_override="normal_tool") + def normal_tool() -> str: + tool_calls["normal"] += 1 + return "normal-result" + + @function_tool(name_override="approval_tool", needs_approval=True) + def approval_tool() -> str: + tool_calls["approval"] += 1 return "secret-result" def output_guardrail( @@ -2592,19 +2939,43 @@ def output_guardrail( return GuardrailFunctionOutput(output_info=None, tripwire_triggered=False) model = ScriptedModel( - [[get_function_tool_call("approval_tool", "{}", call_id="call-approved")]] + [ + [get_function_tool_call("normal_tool", "{}", call_id="call-normal")], + [ + get_function_tool_call("normal_tool", "{}", call_id="call-extra"), + get_function_tool_call("approval_tool", "{}", call_id="call-approved"), + ], + ] ) agent = Agent( name="test", model=model, - tools=[approval_tool], + tools=[normal_tool, approval_tool], output_guardrails=[OutputGuardrail(guardrail_function=output_guardrail)], ) - first = await Runner.run(agent, "Use approval_tool") + first = await Runner.run(agent, "Use normal_tool, then approval_tool") state = first.to_state() - state._current_turn = 2 - state._current_turn_persisted_item_count = 1 - restored = await RunState.from_json(agent, state.to_json()) + assert state._current_turn == 2 + assert tool_calls == {"normal": 2, "approval": 0} + serialized = state.to_json() + ownership = serialized["current_response_generated_item_ownership"] + assert ownership == {"start": 2, "end": 6, "interruptions": [5]} + + if corruption == "missing-prefix": + serialized["generated_items"] = serialized["generated_items"][2:] + elif corruption == "nonterminal-indexes": + trailing_item = json.loads(json.dumps(serialized["generated_items"][0])) + trailing_item["raw_item"]["call_id"] = "call-trailing" + serialized["generated_items"].append(trailing_item) + elif corruption == "invalid-interruption-index": + ownership["interruptions"] = [3] + elif corruption == "processed-mismatch": + serialized["generated_items"][ownership["start"]]["raw_item"]["call_id"] = "call-mismatch" + elif corruption == "interruption-mismatch": + serialized["generated_items"][ownership["interruptions"][0]]["raw_item"]["call_id"] = ( + "call-mismatch" + ) + restored = await RunState.from_json(agent, serialized) restored.approve(restored.get_interruptions()[0]) with pytest.raises(UserError, match="current response boundary cannot be proven"): @@ -2614,7 +2985,7 @@ def output_guardrail( result = Runner.run_streamed(agent, restored, session=None) await consume_stream(result) - assert tool_calls == 0 + assert tool_calls == {"normal": 2, "approval": 0} @pytest.mark.parametrize("mode", ["non_streamed", "streamed"])