From 36d1d9fd170b80d8c3c0db1393c9822c11091340 Mon Sep 17 00:00:00 2001 From: Kazuhiro Sera Date: Tue, 18 Aug 2026 16:36:20 +0900 Subject: [PATCH 01/12] fix: redact blocked tool outputs from replay state --- src/agents/run.py | 74 +++++++++++----- src/agents/run_internal/run_loop.py | 133 +++++++++++++++------------- tests/test_agent_runner.py | 66 +++++++++++++- tests/test_agent_runner_streamed.py | 113 +++++++++++++++++++++-- 4 files changed, 295 insertions(+), 91 deletions(-) diff --git a/src/agents/run.py b/src/agents/run.py index b3fa3f132d..e9f85f6ed1 100644 --- a/src/agents/run.py +++ b/src/agents/run.py @@ -1086,6 +1086,7 @@ def _mark_response_hooks_started() -> None: session_persistence_enabled and turn_session_items and run_state is not None + and not isinstance(turn_result.next_step, NextStepFinalOutput) ): run_state._current_turn_persisted_item_count = ( await save_resumed_turn_items( @@ -1168,13 +1169,57 @@ def _mark_response_hooks_started() -> None: ) if isinstance(turn_result.next_step, NextStepFinalOutput): - await run_output_guardrails( - current_agent.output_guardrails - + (run_config.output_guardrails or []), - current_agent, - turn_result.next_step.output, - context_wrapper, - output_guardrail_results, + try: + await run_output_guardrails( + current_agent.output_guardrails + + (run_config.output_guardrails or []), + current_agent, + turn_result.next_step.output, + context_wrapper, + output_guardrail_results, + ) + except OutputGuardrailTripwireTriggered: + await save_final_turn_items_after_guardrails( + session=session, + run_state=run_state, + session_persistence_enabled=session_persistence_enabled, + input_guardrail_results=( + _attempt_input_guardrail_results() + ), + items=_retained_items_for_blocked_output( + turn_session_items + ), + response_id=turn_result.model_response.response_id, + store=store_setting, + wrapper=context_wrapper, + ) + raise + except (Exception, asyncio.CancelledError): + # An ordinary guardrail failure leaves the verdict unknown, so + # preserve the completed turn exactly as fresh execution does. + await save_final_turn_items_after_guardrails( + session=session, + run_state=run_state, + session_persistence_enabled=session_persistence_enabled, + input_guardrail_results=( + _attempt_input_guardrail_results() + ), + items=turn_session_items, + response_id=turn_result.model_response.response_id, + store=store_setting, + wrapper=context_wrapper, + ) + raise + + await save_final_turn_items_after_guardrails( + session=session, + run_state=run_state, + session_persistence_enabled=session_persistence_enabled, + input_guardrail_results=_attempt_input_guardrail_results(), + items=turn_session_items, + response_id=turn_result.model_response.response_id, + store=store_setting, + wrapper=context_wrapper, ) current_step = getattr(run_state, "_current_step", None) approvals_from_state = approvals_from_step(current_step) @@ -1202,21 +1247,6 @@ def _mark_response_hooks_started() -> None: ) != list(session_items) if run_state is not None: result._trace_state = run_state._trace_state - if session_persistence_enabled: - input_items_for_save_1: list[TResponseInputItem] = ( - session_input_items_for_persistence - if session_input_items_for_persistence is not None - else [] - ) - await save_result_to_session( - session, - input_items_for_save_1, - session_items_for_turn(turn_result), - run_state, - response_id=turn_result.model_response.response_id, - store=store_setting, - wrapper=context_wrapper, - ) result._original_input = copy_input_items(original_input) run_state._current_step = None return _finalize_result(result) diff --git a/src/agents/run_internal/run_loop.py b/src/agents/run_internal/run_loop.py index 90a1320a84..896b669612 100644 --- a/src/agents/run_internal/run_loop.py +++ b/src/agents/run_internal/run_loop.py @@ -53,6 +53,7 @@ ModelResponse, RunItem, ToolApprovalItem, + ToolCallOutputItem, TResponseInputItem, ) from ..lifecycle import RunHooks @@ -459,6 +460,26 @@ async def _run_output_guardrails_for_stream( _SIDE_EFFECT_ITEM_TYPES = frozenset({"tool_call_item", "tool_call_output_item"}) +_OUTPUT_GUARDRAIL_BLOCKED_TOOL_OUTPUT = "Output withheld by an output guardrail." + + +def _sanitize_retained_function_tool_output(item: RunItem) -> None: + """Replace a blocked function-tool result before it enters replayable state.""" + if not isinstance(item, ToolCallOutputItem) or not isinstance(item.raw_item, dict): + return + if item.raw_item.get("type") != "function_call_output": + return + + item.raw_item = cast( + dict[str, Any], + { + **item.raw_item, + "output": _OUTPUT_GUARDRAIL_BLOCKED_TOOL_OUTPUT, + }, + ) + item.output = _OUTPUT_GUARDRAIL_BLOCKED_TOOL_OUTPUT + # Custom data is SDK-only and may duplicate the rejected tool result. + item.custom_data = None def _reasoning_indexes_tied_to_retained_items( @@ -512,6 +533,8 @@ def _retained_items_for_blocked_output(items: list[RunItem]) -> list[RunItem]: # Reasoning items are not side effects themselves, but a reasoning model requires the reasoning # item tied to a function call to accompany it in the next request. retained_indexes |= _reasoning_indexes_tied_to_retained_items(items, retained_indexes) + for index in retained_indexes: + _sanitize_retained_function_tool_output(items[index]) # Indexed rather than filtered by type so the retained items keep the model's own order. return [item for index, item in enumerate(items) if index in retained_indexes] @@ -527,15 +550,9 @@ async def _finalize_streamed_final_output( items: list[RunItem], response_id: str | None, store_setting: bool | None, - persist_before_output_guardrails: bool, on_persisted_after_guardrails: Callable[[bool], None] | None = None, ) -> None: redacted_persistence_error: BaseException | None = None - if persist_before_output_guardrails: - # A resumed approval has already committed the tool side effect, so keep its call/output - # pair even when an agent output guardrail blocks delivery of the final result. - await save_items(items, response_id, store_setting) - try: output_guardrail_results = await _run_output_guardrails_for_stream( agent=agent, @@ -549,10 +566,9 @@ async def _finalize_streamed_final_output( # has to see that side effect rather than re-issue it. This turn reaches here with tool # items when `tool_use_behavior="stop_on_first_tool"` (or `stop_at_tool_names`, or a custom # callable) turned a tool result straight into the final output. - if not persist_before_output_guardrails: - retained_items = _retained_items_for_blocked_output(items) - if retained_items: - await save_items(retained_items, response_id, store_setting) + retained_items = _retained_items_for_blocked_output(items) + if retained_items: + await save_items(retained_items, response_id, store_setting) raise except Exception as guardrail_error: # Only a tripwire means the output was judged undeliverable. A guardrail error leaves the @@ -564,42 +580,41 @@ async def _finalize_streamed_final_output( guardrail_error_is_redacted = _is_error_data_redacted(guardrail_error) if guardrail_error_is_redacted: _detach_data_redacted_error_traceback(guardrail_error) - if not persist_before_output_guardrails: - try: - await save_items(items, response_id, store_setting) - except BaseException as persistence_error: - if guardrail_error_is_redacted: - safe_persistence_error = _safe_redacted_persistence_error(persistence_error) - if ( - isinstance(safe_persistence_error, asyncio.CancelledError) - and streamed_result._cancel_mode != "immediate" - ): - # A cancelled session write is distinct from the caller requesting - # immediate cancellation. Retain a safe cancellation for `stream_events()` - # without completing the run-loop task with the payload-bearing backend - # exception. - streamed_result._stored_exception = safe_persistence_error - streamed_result.is_complete = True - streamed_result._event_queue.put_nowait(QueueCompleteSentinel()) - return - if isinstance(safe_persistence_error, asyncio.CancelledError): - # Public immediate cancellation already owns stream completion and must - # not surface a recovery failure. - return - redacted_persistence_error = safe_persistence_error + try: + await save_items(items, response_id, store_setting) + except BaseException as persistence_error: + if guardrail_error_is_redacted: + safe_persistence_error = _safe_redacted_persistence_error(persistence_error) if ( - isinstance(persistence_error, asyncio.CancelledError) + isinstance(safe_persistence_error, asyncio.CancelledError) and streamed_result._cancel_mode != "immediate" ): - # A cancelled session write is distinct from the caller requesting immediate - # cancellation. The run-loop task itself becomes cancelled, so retain the - # backend cancellation for `stream_events()` to surface. - streamed_result._stored_exception = persistence_error - if redacted_persistence_error is None: - raise - else: - if on_persisted_after_guardrails is not None: - on_persisted_after_guardrails(False) + # A cancelled session write is distinct from the caller requesting + # immediate cancellation. Retain a safe cancellation for `stream_events()` + # without completing the run-loop task with the payload-bearing backend + # exception. + streamed_result._stored_exception = safe_persistence_error + streamed_result.is_complete = True + streamed_result._event_queue.put_nowait(QueueCompleteSentinel()) + return + if isinstance(safe_persistence_error, asyncio.CancelledError): + # Public immediate cancellation already owns stream completion and must + # not surface a recovery failure. + return + redacted_persistence_error = safe_persistence_error + if ( + isinstance(persistence_error, asyncio.CancelledError) + and streamed_result._cancel_mode != "immediate" + ): + # A cancelled session write is distinct from the caller requesting immediate + # cancellation. The run-loop task itself becomes cancelled, so retain the + # backend cancellation for `stream_events()` to surface. + streamed_result._stored_exception = persistence_error + if redacted_persistence_error is None: + raise + else: + if on_persisted_after_guardrails is not None: + on_persisted_after_guardrails(False) if redacted_persistence_error is None: raise @@ -608,22 +623,21 @@ async def _finalize_streamed_final_output( streamed_result.output_guardrail_results.extend(output_guardrail_results) - if not persist_before_output_guardrails: - # Saved as one ordered batch so the session mirrors the model response. Doing it in two - # halves would both reorder the turn and, because the first save advances the turn's - # persisted-item count, make the second one a no-op. - if on_persisted_after_guardrails is None: + # Saved as one ordered batch so the session mirrors the model response. Doing it in two + # halves would both reorder the turn and, because the first save advances the turn's + # persisted-item count, make the second one a no-op. + if on_persisted_after_guardrails is None: + await save_items(items, response_id, store_setting) + else: + try: await save_items(items, response_id, store_setting) - else: - try: - await save_items(items, response_id, store_setting) - except asyncio.CancelledError as persistence_error: - if streamed_result._cancel_mode == "immediate": - raise - streamed_result._stored_exception = persistence_error - streamed_result.is_complete = True - streamed_result._event_queue.put_nowait(QueueCompleteSentinel()) - return + except asyncio.CancelledError as persistence_error: + if streamed_result._cancel_mode == "immediate": + raise + streamed_result._stored_exception = persistence_error + streamed_result.is_complete = True + streamed_result._event_queue.put_nowait(QueueCompleteSentinel()) + return streamed_result.final_output = output if on_persisted_after_guardrails is not None: @@ -1348,7 +1362,6 @@ async def _save_max_turns_items( items=list(turn_session_items), response_id=turn_result.model_response.response_id, store_setting=store_setting, - persist_before_output_guardrails=True, ) run_state._current_step = None break @@ -1538,7 +1551,6 @@ def _record_max_turns_handler_output( items=[synthesized_item] if include_in_history else [], response_id=None, store_setting=store_setting, - persist_before_output_guardrails=False, on_persisted_after_guardrails=_record_max_turns_handler_output, ) streamed_result._max_turns_handled = True @@ -1740,7 +1752,6 @@ def _record_max_turns_handler_output( items=turn_session_items, response_id=turn_result.model_response.response_id, store_setting=store_setting, - persist_before_output_guardrails=False, ) if run_state is not None: run_state._current_step = None diff --git a/tests/test_agent_runner.py b/tests/test_agent_runner.py index 8f33883012..59dc967c15 100644 --- a/tests/test_agent_runner.py +++ b/tests/test_agent_runner.py @@ -114,6 +114,8 @@ from .utils.hitl import make_context_wrapper, make_model_and_agent, make_shell_call from .utils.simple_session import CountingSession, IdStrippingSession, SimpleListSession +_BLOCKED_TOOL_OUTPUT = "Output withheld by an output guardrail." + class _DummyRunItem: def __init__(self, payload: dict[str, Any], item_type: str = "tool_call_output_item"): @@ -4583,9 +4585,11 @@ def foo(a: str) -> str: @pytest.mark.parametrize("tripwire_triggered", [False, True]) @pytest.mark.asyncio -async def test_resumed_final_tool_persists_call_and_output_after_output_guardrail( +async def test_resumed_final_tool_sanitizes_output_after_output_guardrail_tripwire( tripwire_triggered: bool, ) -> None: + tool_calls = 0 + def guardrail_function( _context: RunContextWrapper[Any], _agent: Agent[Any], _agent_output: Any ) -> GuardrailFunctionOutput: @@ -4596,7 +4600,9 @@ def guardrail_function( @function_tool(name_override="commit_tool") def commit_tool() -> str: - return "committed-result" + nonlocal tool_calls + tool_calls += 1 + return f"committed-result-{tool_calls}" session = SimpleListSession() model = ScriptedModel() @@ -4624,7 +4630,7 @@ def commit_tool() -> str: await Runner.run(agent, state, session=session) else: result = await Runner.run(agent, state, session=session) - assert result.final_output == "committed-result" + assert result.final_output == "committed-result-2" assert state._current_turn_persisted_item_count == 4 items = await session.get_items() @@ -4641,6 +4647,60 @@ def commit_tool() -> str: ("function_call", "call-second"), ("function_call_output", "call-second"), ] + second_output = cast(dict[str, Any], items[-1]).get("output") + assert second_output == (_BLOCKED_TOOL_OUTPUT if tripwire_triggered else "committed-result-2") + + serialized_state = json.dumps(state.to_json()) + if tripwire_triggered: + assert "committed-result-2" not in serialized_state + + +@pytest.mark.parametrize("behavior_kind", ["stop_at_tools", "custom"]) +@pytest.mark.asyncio +async def test_blocked_tool_output_is_sanitized_for_terminal_tool_behaviors( + behavior_kind: str, +) -> None: + @function_tool(name_override="commit_tool") + def commit_tool() -> str: + return "sensitive-result" + + def custom_behavior( + _context: RunContextWrapper[Any], results: list[FunctionToolResult] + ) -> ToolsToFinalOutputResult: + return ToolsToFinalOutputResult(is_final_output=True, final_output=results[0].output) + + def output_guardrail( + _context: RunContextWrapper[Any], _agent: Agent[Any], _output: Any + ) -> GuardrailFunctionOutput: + return GuardrailFunctionOutput(output_info=None, tripwire_triggered=True) + + tool_use_behavior: Any = ( + {"stop_at_tool_names": ["commit_tool"]} + if behavior_kind == "stop_at_tools" + else custom_behavior + ) + model = ScriptedModel([[get_function_tool_call("commit_tool", "{}", call_id="call-committed")]]) + session = SimpleListSession() + agent = Agent( + name="test", + model=model, + tools=[commit_tool], + tool_use_behavior=tool_use_behavior, + output_guardrails=[OutputGuardrail(guardrail_function=output_guardrail)], + ) + + with pytest.raises(OutputGuardrailTripwireTriggered): + await Runner.run(agent, "Use commit_tool", session=session) + + items = await session.get_items() + tool_output = next( + item + for item in items + if isinstance(item, dict) and item.get("type") == "function_call_output" + ) + assert tool_output.get("call_id") == "call-committed" + assert tool_output.get("output") == _BLOCKED_TOOL_OUTPUT + assert "sensitive-result" not in json.dumps(items) @pytest.mark.asyncio diff --git a/tests/test_agent_runner_streamed.py b/tests/test_agent_runner_streamed.py index d962a1b393..3b708bad9f 100644 --- a/tests/test_agent_runner_streamed.py +++ b/tests/test_agent_runner_streamed.py @@ -75,6 +75,8 @@ ) from .utils.simple_session import CountingSession, SimpleListSession +_BLOCKED_TOOL_OUTPUT = "Output withheld by an output guardrail." + def _conversation_locked_error() -> BadRequestError: request = httpx.Request("POST", "https://example.com") @@ -2239,13 +2241,20 @@ async def test_tool() -> str: @pytest.mark.parametrize("mode", ["non_streamed", "streamed"]) @pytest.mark.parametrize("tripwire", [False, True], ids=["passes", "trips"]) @pytest.mark.asyncio -async def test_resumed_approved_tool_final_persists_call_output_before_output_guardrails( +async def test_resumed_approved_tool_final_persists_output_after_output_guardrail_verdict( mode: str, tripwire: bool, ) -> None: guardrail_state = {"tripwire": tripwire} - @function_tool(name_override="approval_tool", needs_approval=True) + def extract_custom_data(_context: Any) -> dict[str, str]: + return {"duplicate": "approved-result"} + + @function_tool( + name_override="approval_tool", + needs_approval=True, + custom_data_extractor=extract_custom_data, + ) def approval_tool() -> str: return "approved-result" @@ -2303,7 +2312,15 @@ async def run_once(input_value: Any) -> Any: ("function_call", "call-approved"), ("function_call_output", "call-approved"), ] - assert saved_tool_items[1].get("output") == "approved-result" + expected_output = _BLOCKED_TOOL_OUTPUT if tripwire else "approved-result" + assert saved_tool_items[1].get("output") == expected_output + + serialized_state = json.dumps(state.to_json()) + if tripwire: + assert "approved-result" not in serialized_state + assert _BLOCKED_TOOL_OUTPUT in serialized_state + else: + assert "approved-result" in serialized_state if tripwire: guardrail_state["tripwire"] = False @@ -2323,7 +2340,7 @@ async def run_once(input_value: Any) -> Any: ("function_call", "call-approved"), ("function_call_output", "call-approved"), ] - assert replayed_tool_items[1].get("output") == "approved-result" + assert replayed_tool_items[1].get("output") == _BLOCKED_TOOL_OUTPUT @pytest.mark.parametrize("mode", ["non_streamed", "streamed"]) @@ -2338,7 +2355,7 @@ async def test_stop_on_first_tool_final_persists_committed_tool_items_on_tripwir @function_tool(name_override="commit_tool") def commit_tool() -> str: calls.append("ran") - return "committed-result" + return "sensitive-result" def output_guardrail( _context: RunContextWrapper[Any], @@ -2368,6 +2385,7 @@ def output_guardrail( assert calls == ["ran"], "the tool never ran, so the test proves nothing" saved_items = await session.get_items() + assert "sensitive-result" not in json.dumps(saved_items) saved = [ (item.get("type") or item.get("role"), item.get("call_id")) for item in saved_items @@ -2378,6 +2396,7 @@ def output_guardrail( ("function_call", "call-committed"), ("function_call_output", "call-committed"), ] + assert cast(dict[str, Any], saved_items[-1]).get("output") == _BLOCKED_TOOL_OUTPUT # The next run must see the completed call instead of re-issuing the same side effect. agent.output_guardrails = [] @@ -2401,6 +2420,90 @@ def output_guardrail( ("function_call", "call-committed"), ("function_call_output", "call-committed"), ] + replayed_output = next( + item.get("output") + for item in model_input + if isinstance(item, dict) and item.get("type") == "function_call_output" + ) + assert replayed_output == _BLOCKED_TOOL_OUTPUT + assert "sensitive-result" not in json.dumps(model_input) + + +@pytest.mark.parametrize("mode", ["non_streamed", "streamed"]) +@pytest.mark.asyncio +async def test_blocked_tool_output_cannot_be_forwarded_by_a_later_tool(mode: str) -> None: + secret_calls = 0 + forwarded_values: list[str] = [] + guardrail_outputs: list[Any] = [] + + @function_tool(name_override="secret_tool") + def secret_tool() -> str: + nonlocal secret_calls + secret_calls += 1 + return "private-value" + + @function_tool(name_override="record_replay") + def record_replay(value: str) -> str: + forwarded_values.append(value) + return "recorded" + + def output_guardrail( + _context: RunContextWrapper[Any], + _agent: Agent[Any], + output: Any, + ) -> GuardrailFunctionOutput: + guardrail_outputs.append(output) + return GuardrailFunctionOutput( + output_info=None, + tripwire_triggered=output == "private-value", + ) + + def forward_replayed_output(call: Any) -> list[Any]: + assert isinstance(call.input, list) + replayed_output = next( + item.get("output") + for item in call.input + if isinstance(item, dict) and item.get("type") == "function_call_output" + ) + return [ + get_function_tool_call( + "record_replay", + json.dumps({"value": replayed_output}), + call_id="call-record", + ) + ] + + model = ScriptedModel( + [ + [get_function_tool_call("secret_tool", "{}", call_id="call-secret")], + ModelStep.respond(forward_replayed_output), + ] + ) + session = SimpleListSession() + agent = Agent( + name="test", + model=model, + tools=[secret_tool, record_replay], + tool_use_behavior="stop_on_first_tool", + output_guardrails=[OutputGuardrail(guardrail_function=output_guardrail)], + ) + + async def run_once(input_value: str) -> 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 + + with pytest.raises(OutputGuardrailTripwireTriggered): + await run_once("Read the secret") + + followup = await run_once("Forward the previous result") + assert followup.final_output == "recorded" + assert secret_calls == 1 + assert forwarded_values == [_BLOCKED_TOOL_OUTPUT] + assert guardrail_outputs == ["private-value", "recorded"] + assert "private-value" not in json.dumps(await session.get_items()) @pytest.mark.parametrize("mode", ["non_streamed", "streamed"]) From 9c67b9397d7b71199da7b0e0214cf1e00723cd78 Mon Sep 17 00:00:00 2001 From: Kazuhiro Sera Date: Tue, 18 Aug 2026 19:19:10 +0900 Subject: [PATCH 02/12] fix: cover retained tool output variants --- src/agents/run.py | 99 +++- src/agents/run_internal/run_loop.py | 611 ++++++++++++++++++--- src/agents/run_internal/run_steps.py | 71 +++ src/agents/run_internal/tool_actions.py | 30 +- src/agents/run_internal/tool_execution.py | 123 ++++- tests/test_agent_runner.py | 631 +++++++++++++++++++++- tests/test_agent_runner_streamed.py | 476 +++++++++++++++- 7 files changed, 1910 insertions(+), 131 deletions(-) diff --git a/src/agents/run.py b/src/agents/run.py index e9f85f6ed1..8591b7dd59 100644 --- a/src/agents/run.py +++ b/src/agents/run.py @@ -3,6 +3,7 @@ import asyncio import contextlib import warnings +from functools import partial from typing import TYPE_CHECKING, Any, cast from typing_extensions import Unpack @@ -34,6 +35,7 @@ ItemHelpers, ModelResponse, RunItem, + ToolCallOutputItem, TResponseInputItem, ) from .lifecycle import RunHooks @@ -94,7 +96,12 @@ from .run_internal.prompt_cache_key import PromptCacheKeyResolver from .run_internal.run_grouping import resolve_run_grouping_id from .run_internal.run_loop import ( + _finish_blocked_output_tool_spans, + _redact_blocked_output_state_step, _retained_items_for_blocked_output, + _run_with_deferred_tool_spans, + _sanitize_blocked_output_guardrail_results, + _sanitize_blocked_tool_output_guardrail_results, cleanup_models_after_run, finalize_max_turns_handler_output, get_all_tools, @@ -127,6 +134,7 @@ session_items_for_turn, update_run_state_after_resume, ) +from .run_internal.tool_execution import finish_deferred_tool_spans from .run_internal.tool_use_tracker import ( AgentToolUseTracker, hydrate_tool_use_tracker, @@ -1040,18 +1048,25 @@ def _mark_response_hooks_started() -> None: ) raise UserError("No processed response found in previous state") - turn_result = await resolve_interrupted_turn( - bindings=current_bindings, - original_input=original_input, - original_pre_step_items=generated_items, - new_response=run_state._model_responses[-1], - processed_response=run_state._last_processed_response, - hooks=hooks, - context_wrapper=context_wrapper, + turn_result = await _run_with_deferred_tool_spans( + agent=current_agent, run_config=run_config, - server_manages_conversation=server_conversation_tracker is not None, - run_state=run_state, - error_handlers=error_handlers, + run_turn=partial( + resolve_interrupted_turn, + bindings=current_bindings, + original_input=original_input, + original_pre_step_items=generated_items, + new_response=run_state._model_responses[-1], + processed_response=run_state._last_processed_response, + hooks=hooks, + context_wrapper=context_wrapper, + run_config=run_config, + server_manages_conversation=( + server_conversation_tracker is not None + ), + run_state=run_state, + error_handlers=error_handlers, + ), ) if run_state._last_processed_response is not None: @@ -1169,6 +1184,7 @@ def _mark_response_hooks_started() -> None: ) if isinstance(turn_result.next_step, NextStepFinalOutput): + output_guardrail_result_start = len(output_guardrail_results) try: await run_output_guardrails( current_agent.output_guardrails @@ -1178,7 +1194,30 @@ def _mark_response_hooks_started() -> None: context_wrapper, output_guardrail_results, ) - except OutputGuardrailTripwireTriggered: + except OutputGuardrailTripwireTriggered as exc: + _finish_blocked_output_tool_spans( + turn_result.deferred_tool_spans + ) + has_tool_output = any( + isinstance(item, ToolCallOutputItem) + for item in turn_session_items + ) + if has_tool_output: + _sanitize_blocked_output_guardrail_results( + output_guardrail_results[ + output_guardrail_result_start: + ], + exc, + ) + _sanitize_blocked_tool_output_guardrail_results( + turn_result.tool_output_guardrail_results + ) + _redact_blocked_output_state_step(run_state) + retained_items = _retained_items_for_blocked_output( + turn_session_items, + turn_result.model_response, + turn_result.pre_step_items, + ) await save_final_turn_items_after_guardrails( session=session, run_state=run_state, @@ -1186,15 +1225,16 @@ def _mark_response_hooks_started() -> None: input_guardrail_results=( _attempt_input_guardrail_results() ), - items=_retained_items_for_blocked_output( - turn_session_items - ), + items=retained_items, response_id=turn_result.model_response.response_id, store=store_setting, wrapper=context_wrapper, ) + if has_tool_output: + run_state._current_step = None raise except (Exception, asyncio.CancelledError): + finish_deferred_tool_spans(turn_result.deferred_tool_spans) # An ordinary guardrail failure leaves the verdict unknown, so # preserve the completed turn exactly as fresh execution does. await save_final_turn_items_after_guardrails( @@ -1211,6 +1251,7 @@ def _mark_response_hooks_started() -> None: ) raise + finish_deferred_tool_spans(turn_result.deferred_tool_spans) await save_final_turn_items_after_guardrails( session=session, run_state=run_state, @@ -1695,6 +1736,7 @@ async def _save_max_turns_handler_output( try: if isinstance(turn_result.next_step, NextStepFinalOutput): + output_guardrail_result_start = len(output_guardrail_results) try: await run_output_guardrails( current_agent.output_guardrails @@ -1704,19 +1746,41 @@ async def _save_max_turns_handler_output( context_wrapper, output_guardrail_results, ) - except OutputGuardrailTripwireTriggered: + except OutputGuardrailTripwireTriggered as exc: + _finish_blocked_output_tool_spans(turn_result.deferred_tool_spans) + has_tool_output = any( + isinstance(item, ToolCallOutputItem) + for item in items_to_save_turn + ) + if has_tool_output: + _sanitize_blocked_output_guardrail_results( + output_guardrail_results[output_guardrail_result_start:], + exc, + ) + _sanitize_blocked_tool_output_guardrail_results( + turn_result.tool_output_guardrail_results + ) + _redact_blocked_output_state_step(run_state) + retained_items = _retained_items_for_blocked_output( + items_to_save_turn, + turn_result.model_response, + turn_result.pre_step_items, + ) await save_final_turn_items_after_guardrails( session=session, run_state=run_state, session_persistence_enabled=session_persistence_enabled, input_guardrail_results=_attempt_input_guardrail_results(), - items=_retained_items_for_blocked_output(items_to_save_turn), + items=retained_items, response_id=turn_result.model_response.response_id, store=store_setting, wrapper=context_wrapper, ) + if has_tool_output and run_state is not None: + run_state._current_step = None raise except (Exception, asyncio.CancelledError): + finish_deferred_tool_spans(turn_result.deferred_tool_spans) # Preserve the released non-stream behavior for guardrail errors # and cancellation: the completed final turn remains replayable. await save_final_turn_items_after_guardrails( @@ -1731,6 +1795,7 @@ async def _save_max_turns_handler_output( ) raise + finish_deferred_tool_spans(turn_result.deferred_tool_spans) await save_final_turn_items_after_guardrails( session=session, run_state=run_state, diff --git a/src/agents/run_internal/run_loop.py b/src/agents/run_internal/run_loop.py index 896b669612..0a27ed8ad1 100644 --- a/src/agents/run_internal/run_loop.py +++ b/src/agents/run_internal/run_loop.py @@ -7,7 +7,7 @@ import asyncio import dataclasses as _dc -from collections.abc import Awaitable, Callable +from collections.abc import Awaitable, Callable, Mapping, Sequence from contextlib import aclosing from functools import partial from typing import Any, TypeVar, cast @@ -74,7 +74,7 @@ from ..run_config import ReasoningItemIdPolicy, RunConfig from ..run_context import AgentHookContext, RunContextWrapper, TContext from ..run_error_handlers import RunErrorHandlers -from ..run_state import RunState +from ..run_state import RunState, _deserialize_tool_call_output_raw_item from ..sandbox.runtime import SandboxRuntime from ..stream_events import ( AgentUpdatedStreamEvent, @@ -85,6 +85,7 @@ Tool, dispose_resolved_computers, ) +from ..tool_guardrails import ToolOutputGuardrailResult from ..tracing import Span, SpanError, agent_span, get_current_trace, task_span, turn_span from ..tracing.config import include_task_and_turn_spans from ..tracing.model_tracing import get_model_tracing_impl @@ -140,6 +141,7 @@ from .oai_conversation import OpenAIServerConversationTracker from .prompt_cache_key import PromptCacheKeyResolver, model_settings_with_prompt_cache_key from .run_steps import ( + DeferredToolSpan, NextStepFinalOutput, NextStepHandoff, NextStepInterruption, @@ -173,12 +175,15 @@ from .tool_actions import ApplyPatchAction, ComputerAction, LocalShellAction, ShellAction from .tool_execution import ( coerce_shell_call, + collect_deferred_tool_spans, execute_apply_patch_calls, execute_computer_actions, execute_function_tool_calls, execute_local_shell_calls, execute_shell_calls, extract_tool_call_id, + finish_deferred_tool_spans, + get_mapping_or_attr, initialize_computer_tools, maybe_reset_tool_choice, normalize_shell_output, @@ -200,6 +205,7 @@ validate_run_hooks, ) from .turn_resolution import ( + _collect_program_parent_state, check_for_final_output_from_tools, execute_final_output, execute_handoffs, @@ -451,35 +457,401 @@ async def _run_output_guardrails_for_stream( # Publish at a single boundary so no failure path can omit results that already # finished. A guardrail raising a non-tripwire error reports the same completed # results a tripwire does. + if not isinstance(exc, OutputGuardrailTripwireTriggered): + log_model_action_error(logger, "Unexpected error in output guardrails", exc) streamed_result.output_guardrail_results = ( streamed_result.output_guardrail_results + completed_results ) - if not isinstance(exc, OutputGuardrailTripwireTriggered): - log_model_action_error(logger, "Unexpected error in output guardrails", exc) raise _SIDE_EFFECT_ITEM_TYPES = frozenset({"tool_call_item", "tool_call_output_item"}) _OUTPUT_GUARDRAIL_BLOCKED_TOOL_OUTPUT = "Output withheld by an output guardrail." +_OUTPUT_GUARDRAIL_BLOCKED_TOOL_CALL_ID = "blocked-tool-output" +_OUTPUT_GUARDRAIL_BLOCKED_COMPUTER_SCREENSHOT = ( + "data:image/png;base64," + "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mNk+A8AAQUBAScY42YAAAAASUVORK5CYII=" +) +_RESPONSE_OUTPUT_STATUSES = frozenset({"in_progress", "completed", "incomplete"}) +_HOSTED_TOOL_OUTPUT_STATUSES = frozenset({"completed", "failed"}) +_SHELL_TOOL_OUTPUT_STATUSES = _RESPONSE_OUTPUT_STATUSES | _HOSTED_TOOL_OUTPUT_STATUSES +_PROGRAM_OUTPUT_STATUSES = frozenset({"completed", "incomplete"}) +_KNOWN_TOOL_OUTPUT_TYPES = frozenset( + { + "function_call_output", + "custom_tool_call_output", + "local_shell_call_output", + "apply_patch_call_output", + "shell_call_output", + "computer_call_output", + "program_output", + } +) + + +async def _run_with_deferred_tool_spans( + *, + agent: Agent[TContext], + run_config: RunConfig, + run_turn: Callable[[], Awaitable[SingleStepResult]], +) -> SingleStepResult: + """Delay terminal tool span publication until output guardrails select the payload.""" + should_defer = ( + run_config.trace_include_sensitive_data + and not run_config.tracing_disabled + and agent.tool_use_behavior != "run_llm_again" + and bool(agent.output_guardrails or run_config.output_guardrails) + ) + with collect_deferred_tool_spans(should_defer) as deferred_spans: + try: + result = await run_turn() + except BaseException: + finish_deferred_tool_spans(deferred_spans) + raise + + if isinstance(result.next_step, NextStepFinalOutput): + result.deferred_tool_spans = deferred_spans + else: + finish_deferred_tool_spans(deferred_spans) + return result + + +def _tool_output_payload(raw_item: Any) -> dict[str, Any]: + """Convert a raw tool output into a mutable mapping.""" + if isinstance(raw_item, dict): + return dict(raw_item) + model_dump = getattr(raw_item, "model_dump", None) + if callable(model_dump): + return cast(dict[str, Any], model_dump(exclude_unset=True)) + raise AgentsException(f"Unexpected raw tool output type: {type(raw_item)}") + + +def _tool_output_identity(raw_item: Any) -> tuple[str, str] | None: + """Return the raw output type and provider identity used to join replay copies.""" + payload = _tool_output_payload(raw_item) + output_type = payload.get("type") + call_id = payload.get("call_id") or payload.get("id") + if not isinstance(output_type, str) or not isinstance(call_id, str): + return None + return output_type, call_id + + +def _required_tool_output_string( + raw_payload: Mapping[str, Any], + field: str, + output_type: str, +) -> str: + value = raw_payload.get(field) + if not isinstance(value, str) or not value: + raise AgentsException(f"Cannot sanitize {output_type} without a non-empty string {field}.") + return value + + +def _copy_safe_status( + sanitized: dict[str, Any], + raw_payload: Mapping[str, Any], + output_type: str, + allowed: frozenset[str], + *, + required: bool = False, +) -> None: + status = raw_payload.get("status") + if status is None and not required: + return + if not isinstance(status, str) or status not in allowed: + raise AgentsException(f"Cannot sanitize {output_type} with an invalid status.") + sanitized["status"] = status -def _sanitize_retained_function_tool_output(item: RunItem) -> None: - """Replace a blocked function-tool result before it enters replayable state.""" - if not isinstance(item, ToolCallOutputItem) or not isinstance(item.raw_item, dict): +def _copy_safe_caller( + sanitized: dict[str, Any], + raw_payload: Mapping[str, Any], + output_type: str, +) -> None: + caller = raw_payload.get("caller") + if caller is None: + return + if not isinstance(caller, Mapping): + raise AgentsException(f"Cannot sanitize {output_type} with an invalid caller.") + if caller.get("type") == "direct": + sanitized["caller"] = {"type": "direct"} return - if item.raw_item.get("type") != "function_call_output": + caller_id = caller.get("caller_id") + if caller.get("type") == "program" and isinstance(caller_id, str) and caller_id: + sanitized["caller"] = {"type": "program", "caller_id": caller_id} return + raise AgentsException(f"Cannot sanitize {output_type} with an invalid caller.") + + +def _copy_safe_safety_checks( + sanitized: dict[str, Any], + raw_payload: Mapping[str, Any], +) -> None: + checks = raw_payload.get("acknowledged_safety_checks") + if checks is None: + return + if not isinstance(checks, Sequence) or isinstance(checks, str | bytes): + raise AgentsException( + "Cannot sanitize computer_call_output with invalid acknowledged safety checks." + ) + safe_checks: list[dict[str, str]] = [] + for check in checks: + if not isinstance(check, Mapping): + raise AgentsException( + "Cannot sanitize computer_call_output with invalid acknowledged safety checks." + ) + check_id = check.get("id") + if not isinstance(check_id, str) or not check_id: + raise AgentsException( + "Cannot sanitize computer_call_output with invalid acknowledged safety checks." + ) + safe_checks.append({"id": check_id}) + sanitized["acknowledged_safety_checks"] = safe_checks + + +def _blocked_tool_output_payload(raw_item: Any) -> dict[str, Any]: + """Build a data-free replay payload from validated protocol fields.""" + raw_payload = _tool_output_payload(raw_item) + output_type = raw_payload.get("type") + if not isinstance(output_type, str) or not output_type: + raise AgentsException("Cannot sanitize a tool output without a non-empty string type.") + + sanitized_raw_item: dict[str, Any] = {"type": output_type} + if output_type not in _KNOWN_TOOL_OUTPUT_TYPES: + call_id = raw_payload.get("call_id") + item_id = raw_payload.get("id") + if isinstance(call_id, str) and call_id: + sanitized_raw_item["call_id"] = call_id + elif isinstance(item_id, str) and item_id: + sanitized_raw_item["id"] = item_id + else: + raise AgentsException( + f"Cannot sanitize {output_type} without a non-empty string call_id or id." + ) + sanitized_raw_item["output"] = _OUTPUT_GUARDRAIL_BLOCKED_TOOL_OUTPUT + return sanitized_raw_item + if output_type == "program_output": + sanitized_raw_item["id"] = _required_tool_output_string(raw_payload, "id", output_type) + sanitized_raw_item["call_id"] = _required_tool_output_string( + raw_payload, "call_id", output_type + ) + _copy_safe_status( + sanitized_raw_item, + raw_payload, + output_type, + _PROGRAM_OUTPUT_STATUSES, + required=True, + ) + sanitized_raw_item["result"] = _OUTPUT_GUARDRAIL_BLOCKED_TOOL_OUTPUT + else: + sanitized_raw_item["call_id"] = _required_tool_output_string( + raw_payload, "call_id", output_type + ) + + item_id = raw_payload.get("id") + if item_id is not None: + if not isinstance(item_id, str) or not item_id: + raise AgentsException(f"Cannot sanitize {output_type} with an invalid id.") + sanitized_raw_item["id"] = item_id + + if output_type in { + "function_call_output", + "custom_tool_call_output", + "shell_call_output", + "apply_patch_call_output", + }: + _copy_safe_caller(sanitized_raw_item, raw_payload, output_type) + + if output_type in {"function_call_output", "computer_call_output"}: + _copy_safe_status( + sanitized_raw_item, + raw_payload, + output_type, + _RESPONSE_OUTPUT_STATUSES, + ) + elif output_type == "shell_call_output": + _copy_safe_status( + sanitized_raw_item, + raw_payload, + output_type, + _SHELL_TOOL_OUTPUT_STATUSES, + ) + elif output_type == "apply_patch_call_output": + _copy_safe_status( + sanitized_raw_item, + raw_payload, + output_type, + _HOSTED_TOOL_OUTPUT_STATUSES, + ) + + if output_type == "computer_call_output": + _copy_safe_safety_checks(sanitized_raw_item, raw_payload) + + if output_type == "computer_call_output": + sanitized_raw_item["output"] = { + "type": "computer_screenshot", + "image_url": _OUTPUT_GUARDRAIL_BLOCKED_COMPUTER_SCREENSHOT, + } + elif output_type == "shell_call_output": + sanitized_raw_item["output"] = [ + { + "stdout": "", + "stderr": _OUTPUT_GUARDRAIL_BLOCKED_TOOL_OUTPUT, + "outcome": {"type": "exit", "exit_code": 1}, + } + ] + elif output_type != "program_output": + sanitized_raw_item["output"] = _OUTPUT_GUARDRAIL_BLOCKED_TOOL_OUTPUT + + try: + restored = _deserialize_tool_call_output_raw_item(sanitized_raw_item) + except Exception: + raise AgentsException(f"Sanitized {output_type} is not valid for durable replay.") from None + if restored is None: + raise AgentsException(f"Sanitized {output_type} is not valid for durable replay.") + return sanitized_raw_item - item.raw_item = cast( - dict[str, Any], - { - **item.raw_item, - "output": _OUTPUT_GUARDRAIL_BLOCKED_TOOL_OUTPUT, - }, + +def _sanitize_blocked_output_guardrail_results( + results: Sequence[OutputGuardrailResult], + tripwire: OutputGuardrailTripwireTriggered, +) -> None: + """Remove blocked output aliases from completed guardrail results and the exception.""" + seen: set[int] = set() + for result in (*results, tripwire.guardrail_result): + if id(result) in seen: + continue + seen.add(id(result)) + result.agent_output = _OUTPUT_GUARDRAIL_BLOCKED_TOOL_OUTPUT + result.output.output_info = None + _mark_error_data_redacted(tripwire) + _detach_data_redacted_error_traceback(tripwire) + + +def _sanitize_blocked_tool_output_guardrail_results( + results: Sequence[ToolOutputGuardrailResult], +) -> None: + """Remove blocked tool-output aliases from the current turn's guardrail results.""" + for result in results: + result.output.output_info = None + if result.output.behavior["type"] == "reject_content": + result.output.behavior = { + "type": "reject_content", + "message": _OUTPUT_GUARDRAIL_BLOCKED_TOOL_OUTPUT, + } + + +def _redact_blocked_output_state_step(run_state: RunState[Any] | None) -> None: + """Remove a rejected final output from the live resumable step before propagation.""" + if run_state is not None and isinstance(run_state._current_step, NextStepFinalOutput): + run_state._current_step.output = _OUTPUT_GUARDRAIL_BLOCKED_TOOL_OUTPUT + + +def _finish_blocked_output_tool_spans(spans: list[DeferredToolSpan]) -> None: + """Publish terminal tool spans with the same placeholder used by replay state.""" + finish_deferred_tool_spans( + spans, + output_override=_OUTPUT_GUARDRAIL_BLOCKED_TOOL_OUTPUT, ) - item.output = _OUTPUT_GUARDRAIL_BLOCKED_TOOL_OUTPUT - # Custom data is SDK-only and may duplicate the rejected tool result. - item.custom_data = None + + +def _sanitize_retained_tool_outputs( + items: list[RunItem], + model_response: ModelResponse | None, +) -> None: + """Sanitize retained run items and matching archived raw-response outputs.""" + retained_identities: set[tuple[str, str]] = set() + retained_raw_item_ids: set[int] = set() + for item in items: + if not isinstance(item, ToolCallOutputItem): + continue + + original_raw_item = item.raw_item + identity = _tool_output_identity(original_raw_item) + if identity is not None: + retained_identities.add(identity) + retained_raw_item_ids.add(id(original_raw_item)) + + item.raw_item = cast(Any, _blocked_tool_output_payload(original_raw_item)) + item.output = _OUTPUT_GUARDRAIL_BLOCKED_TOOL_OUTPUT + # Custom data is SDK-only and may duplicate the blocked tool result. + item.custom_data = None + + if model_response is None: + return + + sanitized_response_output = [] + for raw_output in model_response.output: + identity = _tool_output_identity(raw_output) + if id(raw_output) in retained_raw_item_ids or ( + identity is not None and identity in retained_identities + ): + sanitized_response_output.append(cast(Any, _blocked_tool_output_payload(raw_output))) + else: + sanitized_response_output.append(raw_output) + model_response.output = sanitized_response_output + + +def _validate_retained_program_relationships( + items: list[RunItem], + preceding_items: Sequence[RunItem], +) -> None: + """Reject retained program relationships that cannot be replayed in order.""" + preceding_items = list(preceding_items) + for item in items: + raw_item = getattr(item, "raw_item", item) + output_type = get_mapping_or_attr(raw_item, "type") + program_call_ids, completed_program_call_ids = _collect_program_parent_state( + preceding_items + ) + + caller = get_mapping_or_attr(raw_item, "caller") + if get_mapping_or_attr(caller, "type") == "program": + caller_id = get_mapping_or_attr(caller, "caller_id") + if ( + not isinstance(caller_id, str) + or caller_id not in program_call_ids + or caller_id in completed_program_call_ids + ): + raise AgentsException( + f"Cannot sanitize {output_type} with an invalid program caller." + ) + + if output_type == "program_output": + call_id = get_mapping_or_attr(raw_item, "call_id") + if ( + not isinstance(call_id, str) + or call_id not in program_call_ids + or call_id in completed_program_call_ids + ): + raise AgentsException( + "Cannot sanitize program_output without an active retained program parent." + ) + + preceding_items.append(item) + + +def _scrub_failed_blocked_output_aliases( + items: list[RunItem], + model_response: ModelResponse | None, +) -> None: + """Remove output data after replay-payload validation fails.""" + for item in items: + if not isinstance(item, ToolCallOutputItem): + continue + item.raw_item = cast( + Any, + { + "type": "function_call_output", + "call_id": _OUTPUT_GUARDRAIL_BLOCKED_TOOL_CALL_ID, + "output": _OUTPUT_GUARDRAIL_BLOCKED_TOOL_OUTPUT, + }, + ) + item.output = _OUTPUT_GUARDRAIL_BLOCKED_TOOL_OUTPUT + item.custom_data = None + + if model_response is not None: + model_response.output = [] def _reasoning_indexes_tied_to_retained_items( @@ -513,7 +885,11 @@ def _reasoning_indexes_tied_to_retained_items( return tied -def _retained_items_for_blocked_output(items: list[RunItem]) -> list[RunItem]: +def _retained_items_for_blocked_output( + items: list[RunItem], + model_response: ModelResponse | None = None, + preceding_items: Sequence[RunItem] = (), +) -> list[RunItem]: """Pick out the items of a final turn to keep when its output is not deliverable. A tool that already ran has to stay in the session, together with the context needed to replay @@ -525,18 +901,32 @@ def _retained_items_for_blocked_output(items: list[RunItem]) -> list[RunItem]: that goes unclassified is a bug, so the safer default is the one that surfaces as a missing item rather than as a rejected message quietly reaching the session. """ - retained_indexes = { - index for index, item in enumerate(items) if item.type in _SIDE_EFFECT_ITEM_TYPES - } - if not retained_indexes: - return [] - # Reasoning items are not side effects themselves, but a reasoning model requires the reasoning - # item tied to a function call to accompany it in the next request. - retained_indexes |= _reasoning_indexes_tied_to_retained_items(items, retained_indexes) - for index in retained_indexes: - _sanitize_retained_function_tool_output(items[index]) - # Indexed rather than filtered by type so the retained items keep the model's own order. - return [item for index, item in enumerate(items) if index in retained_indexes] + redacted_error: AgentsException | None = None + try: + retained_indexes = { + index for index, item in enumerate(items) if item.type in _SIDE_EFFECT_ITEM_TYPES + } + if not retained_indexes: + return [] + # Reasoning items are not side effects themselves, but a reasoning model requires the + # reasoning item tied to a function call to accompany it in the next request. + retained_indexes |= _reasoning_indexes_tied_to_retained_items(items, retained_indexes) + retained_items = [items[index] for index in sorted(retained_indexes)] + _validate_retained_program_relationships(retained_items, preceding_items) + _sanitize_retained_tool_outputs(retained_items, model_response) + # Indexed rather than filtered by type so the retained items keep the model's own order. + return [item for index, item in enumerate(items) if index in retained_indexes] + except AgentsException as error: + _scrub_failed_blocked_output_aliases(items, model_response) + _mark_error_data_redacted(error) + _detach_data_redacted_error_traceback(error) + redacted_error = AgentsException("Cannot sanitize a blocked tool output for replay.") + _mark_error_data_redacted(redacted_error) + + items = [] + model_response = None + assert redacted_error is not None + raise redacted_error from None async def _finalize_streamed_final_output( @@ -548,11 +938,16 @@ async def _finalize_streamed_final_output( context_wrapper: RunContextWrapper[TContext], save_items: Callable[[list[RunItem], str | None, bool | None], Awaitable[None]], items: list[RunItem], + model_response: ModelResponse | None, + deferred_tool_spans: list[DeferredToolSpan], + preceding_items: Sequence[RunItem], + tool_output_guardrail_results: Sequence[ToolOutputGuardrailResult], response_id: str | None, store_setting: bool | None, on_persisted_after_guardrails: Callable[[bool], None] | None = None, ) -> None: redacted_persistence_error: BaseException | None = None + output_guardrail_result_start = len(streamed_result.output_guardrail_results) try: output_guardrail_results = await _run_output_guardrails_for_stream( agent=agent, @@ -561,16 +956,46 @@ async def _finalize_streamed_final_output( context_wrapper=context_wrapper, streamed_result=streamed_result, ) - except OutputGuardrailTripwireTriggered: + except OutputGuardrailTripwireTriggered as exc: # The blocked output itself is not persisted, but a tool that already ran is: the next run # has to see that side effect rather than re-issue it. This turn reaches here with tool # items when `tool_use_behavior="stop_on_first_tool"` (or `stop_at_tool_names`, or a custom # callable) turned a tool result straight into the final output. - retained_items = _retained_items_for_blocked_output(items) + has_tool_output = any(isinstance(item, ToolCallOutputItem) for item in items) + if has_tool_output: + _sanitize_blocked_output_guardrail_results( + streamed_result.output_guardrail_results[output_guardrail_result_start:], + exc, + ) + _sanitize_blocked_tool_output_guardrail_results(tool_output_guardrail_results) + _finish_blocked_output_tool_spans(deferred_tool_spans) + if has_tool_output: + _redact_blocked_output_state_step(streamed_result._state) + retained_items = _retained_items_for_blocked_output( + items, + model_response, + preceding_items, + ) if retained_items: - await save_items(retained_items, response_id, store_setting) + try: + await save_items(retained_items, response_id, store_setting) + except asyncio.CancelledError as persistence_error: + if streamed_result._cancel_mode == "immediate": + raise + streamed_result._stored_exception = _safe_redacted_persistence_error( + persistence_error + ) + streamed_result.is_complete = True + streamed_result._event_queue.put_nowait(QueueCompleteSentinel()) + return + if has_tool_output and streamed_result._state is not None: + streamed_result._state._current_step = None + raise + except asyncio.CancelledError: + finish_deferred_tool_spans(deferred_tool_spans) raise except Exception as guardrail_error: + finish_deferred_tool_spans(deferred_tool_spans) # Only a tripwire means the output was judged undeliverable. A guardrail error leaves the # verdict unknown, so the completed final turn is persisted whole and remains replayable. # `asyncio.CancelledError` is deliberately not caught here: `cancel()` in its default @@ -621,6 +1046,7 @@ async def _finalize_streamed_final_output( if redacted_persistence_error is not None: raise redacted_persistence_error from None + finish_deferred_tool_spans(deferred_tool_spans) streamed_result.output_guardrail_results.extend(output_guardrail_results) # Saved as one ordered batch so the session mirrors the model response. Doing it in two @@ -1242,18 +1668,23 @@ async def _save_max_turns_items( last_model_response = run_state._model_responses[-1] - turn_result = await resolve_interrupted_turn( - bindings=current_bindings, - original_input=run_state._original_input, - original_pre_step_items=run_state._generated_items, - new_response=last_model_response, - processed_response=run_state._last_processed_response, - hooks=hooks, - context_wrapper=context_wrapper, + turn_result = await _run_with_deferred_tool_spans( + agent=current_agent, run_config=run_config, - server_manages_conversation=server_conversation_tracker is not None, - run_state=run_state, - error_handlers=error_handlers, + run_turn=partial( + resolve_interrupted_turn, + bindings=current_bindings, + original_input=run_state._original_input, + original_pre_step_items=run_state._generated_items, + new_response=last_model_response, + processed_response=run_state._last_processed_response, + hooks=hooks, + context_wrapper=context_wrapper, + run_config=run_config, + server_manages_conversation=server_conversation_tracker is not None, + run_state=run_state, + error_handlers=error_handlers, + ), ) tool_use_tracker.record_processed_response( @@ -1360,9 +1791,17 @@ async def _save_max_turns_items( context_wrapper=context_wrapper, save_items=_save_resumed_items, items=list(turn_session_items), + model_response=turn_result.model_response, + deferred_tool_spans=turn_result.deferred_tool_spans, + preceding_items=turn_result.pre_step_items, + tool_output_guardrail_results=( + turn_result.tool_output_guardrail_results + ), response_id=turn_result.model_response.response_id, store_setting=store_setting, ) + if streamed_result._stored_exception is not None: + break run_state._current_step = None break @@ -1549,10 +1988,16 @@ def _record_max_turns_handler_output( context_wrapper=context_wrapper, save_items=_save_max_turns_items, items=[synthesized_item] if include_in_history else [], + model_response=None, + deferred_tool_spans=[], + preceding_items=[], + tool_output_guardrail_results=[], response_id=None, store_setting=store_setting, on_persisted_after_guardrails=_record_max_turns_handler_output, ) + if streamed_result._stored_exception is not None: + break streamed_result._max_turns_handled = True streamed_result.current_turn = max_turns if run_state is not None and not is_resumed_state: @@ -1750,9 +2195,15 @@ def _record_max_turns_handler_output( context_wrapper=context_wrapper, save_items=_save_stream_items_with_count, items=turn_session_items, + model_response=turn_result.model_response, + deferred_tool_spans=turn_result.deferred_tool_spans, + preceding_items=turn_result.pre_step_items, + tool_output_guardrail_results=turn_result.tool_output_guardrail_results, response_id=turn_result.model_response.response_id, store_setting=store_setting, ) + if streamed_result._stored_exception is not None: + break if run_state is not None: run_state._current_step = None break @@ -2197,23 +2648,28 @@ async def after_invocation_validation( async def check_input_guardrails_before_side_effects() -> None: await raise_if_input_guardrail_tripwire_known() - single_step_result = await get_single_step_result_from_response( - bindings=bindings, - original_input=streamed_result.input, - pre_step_items=streamed_result._model_input_items, - new_response=final_response, - output_schema=output_schema, - all_tools=all_tools, - handoffs=handoffs, - hooks=hooks, - context_wrapper=context_wrapper, + single_step_result = await _run_with_deferred_tool_spans( + agent=public_agent, run_config=run_config, - error_handlers=error_handlers, - tool_use_tracker=tool_use_tracker, - server_manages_conversation=server_conversation_tracker is not None, - after_invocation_validation=after_invocation_validation, - before_side_effects=check_input_guardrails_before_side_effects, - run_state=run_state, + run_turn=partial( + get_single_step_result_from_response, + bindings=bindings, + original_input=streamed_result.input, + pre_step_items=streamed_result._model_input_items, + new_response=final_response, + output_schema=output_schema, + all_tools=all_tools, + handoffs=handoffs, + hooks=hooks, + context_wrapper=context_wrapper, + run_config=run_config, + error_handlers=error_handlers, + tool_use_tracker=tool_use_tracker, + server_manages_conversation=server_conversation_tracker is not None, + after_invocation_validation=after_invocation_validation, + before_side_effects=check_input_guardrails_before_side_effects, + run_state=run_state, + ), ) items_to_filter = session_items_for_turn(single_step_result) @@ -2341,22 +2797,27 @@ async def after_invocation_validation( ) return response_accepted - return await get_single_step_result_from_response( - bindings=bindings, - original_input=original_input, - pre_step_items=generated_items, - new_response=new_response, - output_schema=output_schema, - all_tools=all_tools, - handoffs=handoffs, - hooks=hooks, - context_wrapper=context_wrapper, + return await _run_with_deferred_tool_spans( + agent=public_agent, run_config=run_config, - error_handlers=error_handlers, - tool_use_tracker=tool_use_tracker, - server_manages_conversation=server_conversation_tracker is not None, - after_invocation_validation=after_invocation_validation, - run_state=run_state, + run_turn=partial( + get_single_step_result_from_response, + bindings=bindings, + original_input=original_input, + pre_step_items=generated_items, + new_response=new_response, + output_schema=output_schema, + all_tools=all_tools, + handoffs=handoffs, + hooks=hooks, + context_wrapper=context_wrapper, + run_config=run_config, + error_handlers=error_handlers, + tool_use_tracker=tool_use_tracker, + server_manages_conversation=server_conversation_tracker is not None, + after_invocation_validation=after_invocation_validation, + run_state=run_state, + ), ) diff --git a/src/agents/run_internal/run_steps.py b/src/agents/run_internal/run_steps.py index f692e1ca4e..3a67bac5ee 100644 --- a/src/agents/run_internal/run_steps.py +++ b/src/agents/run_internal/run_steps.py @@ -26,6 +26,7 @@ ShellTool, ) from ..tool_guardrails import ToolInputGuardrailResult, ToolOutputGuardrailResult +from ..tracing import Span, SpanError from .items import NestedHistoryOwnedItem __all__ = [ @@ -181,6 +182,73 @@ class NextStepInterruption: """Whether response-end hooks started before the interruption was persisted.""" +@dataclass +class DeferredToolSpan(Span[Any]): + """A tool span whose sensitive output is held until the terminal verdict.""" + + span: Span[Any] + output: Any = None + has_output: bool = False + deferred_error: SpanError | None = None + has_error: bool = False + + @property + def trace_id(self) -> str: + return self.span.trace_id + + @property + def span_id(self) -> str: + return self.span.span_id + + @property + def span_data(self) -> Any: + return self.span.span_data + + @property + def parent_id(self) -> str | None: + return self.span.parent_id + + def start(self, mark_as_current: bool = False) -> None: + self.span.start(mark_as_current=mark_as_current) + + def finish(self, reset_current: bool = False) -> None: + self.span.finish(reset_current=reset_current) + + def __enter__(self) -> DeferredToolSpan: + self.start(mark_as_current=True) + return self + + def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None: + self.finish(reset_current=True) + + def set_error(self, error: SpanError) -> None: + self.deferred_error = error + self.has_error = True + + @property + def error(self) -> SpanError | None: + return self.deferred_error if self.has_error else self.span.error + + def export(self) -> dict[str, Any] | None: + return self.span.export() + + @property + def started_at(self) -> str | None: + return self.span.started_at + + @property + def ended_at(self) -> str | None: + return self.span.ended_at + + @property + def tracing_api_key(self) -> str | None: + return self.span.tracing_api_key + + @property + def trace_metadata(self) -> dict[str, Any] | None: + return self.span.trace_metadata + + @dataclass class SingleStepResult: original_input: str | list[TResponseInputItem] @@ -222,6 +290,9 @@ class SingleStepResult: processed_response: ProcessedResponse | None = None """The processed model response. This is needed for resuming from interruptions.""" + deferred_tool_spans: list[DeferredToolSpan] = dataclasses.field(default_factory=list) + """Tool spans waiting for the terminal output-guardrail verdict before publication.""" + @property def generated_items(self) -> list[RunItem]: """Items generated during the agent run (i.e. everything generated after diff --git a/src/agents/run_internal/tool_actions.py b/src/agents/run_internal/tool_actions.py index 6bffae29cd..95a28af452 100644 --- a/src/agents/run_internal/tool_actions.py +++ b/src/agents/run_internal/tool_actions.py @@ -58,6 +58,8 @@ resolve_approval_rejection_message, resolve_approval_status, serialize_shell_output, + set_tool_span_error, + set_tool_span_output, truncate_shell_outputs, with_tool_function_span, ) @@ -150,14 +152,15 @@ async def _run_action(span: Any | None) -> RunItem: error_message=error_text, ) if span is not None: - span.set_error( + set_tool_span_error( + span, SpanError( message="Error running tool", data={ "tool_name": trace_tool_name, "error": trace_error, }, - ) + ), ) log_tool_action_error("Failed to execute computer action", exc) raise @@ -201,7 +204,7 @@ async def _run_action(span: Any | None) -> RunItem: ) if span is not None and config.trace_include_sensitive_data: - span.span_data.output = image_url + set_tool_span_output(span, image_url) return output_item @@ -590,14 +593,15 @@ async def _run_call(span: Any | None) -> RunItem: error_message=output_text, ) if span is not None: - span.set_error( + set_tool_span_error( + span, SpanError( message="Error running tool", data={ "tool_name": shell_tool.name, "error": trace_error, }, - ) + ), ) if requested_max_output_length is not None: max_output_length = requested_max_output_length @@ -651,7 +655,7 @@ async def _run_call(span: Any | None) -> RunItem: ) if span is not None and config.trace_include_sensitive_data: - span.span_data.output = output_text + set_tool_span_output(span, output_text) return output_item @@ -777,14 +781,15 @@ async def _run_call(span: Any | None) -> RunItem: error_message=output_text, ) if span is not None: - span.set_error( + set_tool_span_error( + span, SpanError( message="Error running tool", data={ "tool_name": custom_tool.name, "error": trace_error, }, - ) + ), ) log_tool_action_error("Custom tool failed", exc) @@ -823,7 +828,7 @@ async def _run_call(span: Any | None) -> RunItem: ) if span is not None and config.trace_include_sensitive_data: - span.span_data.output = output_text + set_tool_span_output(span, output_text) return output_item return await with_tool_function_span( @@ -1009,14 +1014,15 @@ async def _run_call(span: Any | None) -> RunItem: error_message=output_text, ) if span is not None: - span.set_error( + set_tool_span_error( + span, SpanError( message="Error running tool", data={ "tool_name": apply_patch_tool.name, "error": trace_error, }, - ) + ), ) log_tool_action_error("Apply patch editor failed", exc) @@ -1060,7 +1066,7 @@ async def _run_call(span: Any | None) -> RunItem: ) if span is not None and config.trace_include_sensitive_data: - span.span_data.output = output_text + set_tool_span_output(span, output_text) return output_item diff --git a/src/agents/run_internal/tool_execution.py b/src/agents/run_internal/tool_execution.py index af257b165f..e32cac9efd 100644 --- a/src/agents/run_internal/tool_execution.py +++ b/src/agents/run_internal/tool_execution.py @@ -6,12 +6,15 @@ from __future__ import annotations import asyncio +import contextvars import copy import dataclasses import functools import inspect import json -from collections.abc import Awaitable, Callable, Mapping, Sequence +from collections.abc import Awaitable, Callable, Iterator, Mapping, Sequence +from contextlib import contextmanager +from contextvars import Token from typing import TYPE_CHECKING, Any, Literal, TypeVar, cast from openai.types.responses import ResponseFunctionToolCall @@ -93,6 +96,7 @@ ToolOutputGuardrailResult, ) from ..tracing import Span, SpanError, function_span, get_current_trace +from ..tracing.scope import Scope from ..util import _coro, _error_tracing from ..util._approvals import evaluate_needs_approval_setting, parse_function_tool_arguments from ..util._asyncio_tasks import gather_with_cancel @@ -109,7 +113,7 @@ function_rejection_item, function_tool_error_output, ) -from .run_steps import ToolRunFunction +from .run_steps import DeferredToolSpan, ToolRunFunction from .tool_use_tracker import AgentToolUseTracker if TYPE_CHECKING: @@ -123,6 +127,106 @@ ToolRunShellCall, ) + +_deferred_tool_spans: contextvars.ContextVar[list[DeferredToolSpan] | None] = ( + contextvars.ContextVar( + "deferred_tool_spans", + default=None, + ) +) + + +@contextmanager +def collect_deferred_tool_spans(enabled: bool) -> Iterator[list[DeferredToolSpan]]: + """Collect tool spans without publishing their output until the turn verdict is known.""" + spans: list[DeferredToolSpan] = [] + token: Token[list[DeferredToolSpan] | None] = _deferred_tool_spans.set( + spans if enabled else None + ) + try: + yield spans + finally: + _deferred_tool_spans.reset(token) + + +def finish_deferred_tool_spans( + spans: list[DeferredToolSpan], + *, + output_override: Any | None = None, +) -> None: + """Publish a deferred tool-span batch exactly once, optionally replacing its output.""" + pending_spans = list(spans) + spans.clear() + for deferred_span in pending_spans: + span = deferred_span.span + if span.ended_at is not None: + continue + if output_override is not None: + cast(Any, span.span_data).output = output_override + elif deferred_span.has_output: + cast(Any, span.span_data).output = deferred_span.output + if output_override is None and deferred_span.has_error: + assert deferred_span.deferred_error is not None + span.set_error(deferred_span.deferred_error) + deferred_span.output = None + deferred_span.has_output = False + deferred_span.deferred_error = None + deferred_span.has_error = False + span.finish() + + +def set_tool_span_output(span: Span[Any], output: Any) -> None: + """Store sensitive tool output outside a processor-visible span until finalization.""" + deferred_spans = _deferred_tool_spans.get() + if deferred_spans is not None: + for deferred_span in reversed(deferred_spans): + if deferred_span.span is span: + deferred_span.output = output + deferred_span.has_output = True + return + cast(Any, span.span_data).output = output + + +def set_tool_span_error(span: Span[Any], error: SpanError) -> None: + """Store a tool error outside a processor-visible span until finalization.""" + deferred_spans = _deferred_tool_spans.get() + if deferred_spans is not None: + for deferred_span in reversed(deferred_spans): + if deferred_span is span or deferred_span.span is span: + deferred_span.set_error(error) + return + span.set_error(error) + + +@contextmanager +def _tool_function_span(tool_name: str) -> Iterator[Span[Any]]: + """Keep current-span ownership in the tool task while allowing deferred publication.""" + span = function_span(tool_name) + span.start() + deferred_spans = _deferred_tool_spans.get() + if deferred_spans is None: + deferred_span = None + else: + deferred_span = DeferredToolSpan(span=span) + deferred_spans.append(deferred_span) + token = Scope.set_current_span(deferred_span or span) + try: + yield span + except BaseException: + Scope.reset_current_span(token) + if deferred_span is None: + span.finish() + else: + assert deferred_spans is not None + deferred_spans.remove(deferred_span) + finish_deferred_tool_spans([deferred_span]) + raise + else: + Scope.reset_current_span(token) + if deferred_span is None: + span.finish() + + __all__ = [ "maybe_reset_tool_choice", "initialize_computer_tools", @@ -146,6 +250,8 @@ "format_shell_error", "get_trace_tool_error", "with_tool_function_span", + "set_tool_span_output", + "set_tool_span_error", "build_litellm_json_tool_call", "collect_manual_mcp_approvals", "index_approval_items_by_call_id", @@ -1131,7 +1237,7 @@ async def with_tool_function_span( direct_result: object = result return cast(TToolSpanResult, direct_result) - with function_span(tool_name) as span: + with _tool_function_span(tool_name) as span: result = fn(span) if inspect.isawaitable(result): return await result @@ -1802,7 +1908,7 @@ async def _run_single_tool( or get_function_tool_trace_name(func_tool) or func_tool.name ) - with function_span(trace_tool_name) as span_fn: + with _tool_function_span(trace_tool_name) as span_fn: tool_context_namespace = get_tool_call_namespace(raw_tool_call) if tool_context_namespace is None: tool_context_namespace = get_tool_call_namespace(tool_call) @@ -1852,7 +1958,7 @@ async def _run_single_tool( raise UserError(f"Error running tool {func_tool.name}: {e}") from e if self.config.trace_include_sensitive_data: - span_fn.span_data.output = result + set_tool_span_output(span_fn, result) return result async def _maybe_execute_tool_approval( @@ -1964,7 +2070,8 @@ async def _maybe_execute_tool_approval( tool_namespace=tool_namespace, tool_lookup_key=tool_lookup_key, ) - span_fn.set_error( + set_tool_span_error( + span_fn, SpanError( message=rejection_message, data={ @@ -1973,9 +2080,9 @@ async def _maybe_execute_tool_approval( f"Tool execution for {tool_call.call_id} was manually rejected by user." ), }, - ) + ), ) - span_fn.span_data.output = rejection_message + set_tool_span_output(span_fn, rejection_message) return FunctionToolResult( tool=func_tool, output=rejection_message, diff --git a/tests/test_agent_runner.py b/tests/test_agent_runner.py index 59dc967c15..7e8cc0d26f 100644 --- a/tests/test_agent_runner.py +++ b/tests/test_agent_runner.py @@ -13,8 +13,8 @@ import httpx import pytest from openai import APIConnectionError, BadRequestError, NotFoundError -from openai.types.responses import ResponseFunctionToolCall -from openai.types.responses.response_output_item import McpApprovalRequest +from openai.types.responses import ResponseFunctionShellToolCallOutput, ResponseFunctionToolCall +from openai.types.responses.response_output_item import McpApprovalRequest, ProgramOutput from openai.types.responses.response_output_text import AnnotationFileCitation, ResponseOutputText from openai.types.responses.response_reasoning_item import ResponseReasoningItem, Summary from openai.types.responses.tool_param import Mcp @@ -24,6 +24,7 @@ from agents import ( Agent, AgentOutputSchema, + AgentsException, GuardrailFunctionOutput, Handoff, HandoffInputData, @@ -71,6 +72,7 @@ from agents.models.fake_id import FAKE_RESPONSES_ID from agents.run import AgentRunner, get_default_agent_runner, set_default_agent_runner from agents.run_config import _default_trace_include_sensitive_data +from agents.run_internal import run_loop from agents.run_internal.agent_bindings import bind_public_agent from agents.run_internal.agent_runner_helpers import build_resumed_stream_debug_extra from agents.run_internal.items import ( @@ -110,6 +112,7 @@ get_text_input_item, get_text_message, ) +from .testing_processor import SPAN_PROCESSOR_TESTING, fetch_ordered_spans from .utils.factories import make_run_state from .utils.hitl import make_context_wrapper, make_model_and_agent, make_shell_call from .utils.simple_session import CountingSession, IdStrippingSession, SimpleListSession @@ -117,6 +120,34 @@ _BLOCKED_TOOL_OUTPUT = "Output withheld by an output guardrail." +def _sdk_exception_traceback_repr_locations( + error: BaseException, expected: str +) -> list[tuple[str, str]]: + pending = [error] + seen: set[int] = set() + locations: list[tuple[str, str]] = [] + while pending: + current = pending.pop() + if id(current) in seen: + continue + seen.add(id(current)) + + traceback = current.__traceback__ + while traceback is not None: + if "/src/agents/" in traceback.tb_frame.f_code.co_filename: + locations.extend( + (traceback.tb_frame.f_code.co_name, name) + for name, value in traceback.tb_frame.f_locals.items() + if expected in repr(value) + ) + traceback = traceback.tb_next + + for linked in (current.__cause__, current.__context__): + if linked is not None: + pending.append(linked) + return locations + + class _DummyRunItem: def __init__(self, payload: dict[str, Any], item_type: str = "tool_call_output_item"): self._payload = payload @@ -126,6 +157,432 @@ def to_input_item(self) -> dict[str, Any]: return self._payload +@pytest.mark.parametrize( + "raw_item", + [ + { + "type": "function_call_output", + "call_id": "call-function", + "output": "secret", + "caller": { + "type": "program", + "caller_id": "program-1", + "duplicate": "secret", + }, + }, + { + "type": "custom_tool_call_output", + "call_id": "call-custom", + "output": "secret", + }, + { + "type": "local_shell_call_output", + "call_id": "call-local-shell", + "output": "secret", + }, + { + "type": "apply_patch_call_output", + "call_id": "call-apply-patch", + "output": "secret", + "status": "completed", + }, + { + "type": "shell_call_output", + "call_id": "call-shell", + "output": [ + { + "stdout": "secret", + "stderr": "", + "outcome": {"type": "exit", "exit_code": 0}, + } + ], + "status": "completed", + "shell_output": [{"stdout": "secret"}], + "provider_data": {"raw_output": "secret"}, + "caller": { + "type": "program", + "caller_id": "program-1", + "duplicate": "secret", + }, + }, + ResponseFunctionShellToolCallOutput( + id="shell-output-in-progress", + call_id="call-shell-in-progress", + output=[ + { + "stdout": "secret", + "stderr": "", + "outcome": {"type": "exit", "exit_code": 0}, + } + ], + status="in_progress", + type="shell_call_output", + ), + ResponseFunctionShellToolCallOutput( + id="shell-output-incomplete", + call_id="call-shell-incomplete", + output=[ + { + "stdout": "secret", + "stderr": "", + "outcome": {"type": "exit", "exit_code": 0}, + } + ], + status="incomplete", + type="shell_call_output", + ), + { + "type": "computer_call_output", + "call_id": "call-computer", + "output": { + "type": "computer_screenshot", + "image_url": "data:image/png;base64,c2VjcmV0", + }, + "acknowledged_safety_checks": [ + { + "id": "check-1", + "code": "secret", + "message": "secret", + "duplicate": "secret", + } + ], + }, + { + "type": "program_output", + "id": "program-output", + "call_id": "call-program", + "result": "secret", + "status": "completed", + }, + ], +) +@pytest.mark.asyncio +async def test_blocked_retained_tool_output_uses_replay_valid_payload( + raw_item: Any, +) -> None: + raw_payload = ( + raw_item if isinstance(raw_item, dict) else raw_item.model_dump(exclude_unset=True) + ) + agent = Agent(name="test") + item = ToolCallOutputItem( + agent=agent, + raw_item=cast(Any, raw_item), + output="secret", + custom_data={"duplicate": "secret"}, + ) + retained_input: list[RunItem] = [] + if raw_payload.get("caller") is not None or raw_payload["type"] == "program_output": + retained_input.append( + ToolCallItem( + agent=agent, + raw_item=cast( + Any, + { + "type": "program", + "call_id": ( + raw_payload["call_id"] + if raw_payload["type"] == "program_output" + else raw_payload["caller"]["caller_id"] + ), + "code": "return await tools.lookup({});", + }, + ), + ) + ) + retained_input.append(item) + + retained = run_loop._retained_items_for_blocked_output(retained_input) + + assert retained == retained_input + assert item.output == _BLOCKED_TOOL_OUTPUT + assert item.custom_data is None + payload = cast(dict[str, Any], item.to_input_item()) + assert payload["type"] == raw_payload["type"] + assert "secret" not in json.dumps(payload) + if "caller" in raw_payload: + assert payload["caller"] == {"type": "program", "caller_id": "program-1"} + sanitized_raw_item = cast(dict[str, Any], item.raw_item) + if raw_payload["type"] == "shell_call_output" and "status" in raw_payload: + assert sanitized_raw_item["status"] == raw_payload["status"] + + if payload["type"] == "computer_call_output": + assert payload["output"]["type"] == "computer_screenshot" + assert payload["output"]["image_url"].startswith("data:image/png;base64,") + elif payload["type"] == "shell_call_output": + assert payload["output"] == [ + { + "stdout": "", + "stderr": _BLOCKED_TOOL_OUTPUT, + "outcome": {"type": "exit", "exit_code": 1}, + } + ] + elif payload["type"] == "program_output": + assert payload["result"] == _BLOCKED_TOOL_OUTPUT + else: + assert payload["output"] == _BLOCKED_TOOL_OUTPUT + + state = make_run_state(agent) + state._generated_items = retained_input + state._session_items = retained_input + state_json = state.to_json() + assert "secret" not in json.dumps(state_json) + restored = await RunState.from_json(agent, state_json) + restored_item = cast(ToolCallOutputItem, restored._generated_items[-1]) + restored_payload = cast(dict[str, Any], restored_item.to_input_item()) + assert restored_payload == payload + + +@pytest.mark.parametrize( + "raw_item", + [ + { + "type": "local_shell_call_output", + "id": "unsupported-provider-id", + "output": "secret", + }, + { + "type": "function_call_output", + "call_id": "call-function", + "output": "secret", + "caller": {"type": "program", "duplicate": "secret"}, + }, + { + "type": "computer_call_output", + "call_id": "call-computer", + "output": { + "type": "computer_screenshot", + "image_url": "data:image/png;base64,c2VjcmV0", + }, + "acknowledged_safety_checks": [{"id": ""}], + }, + { + "type": "function_call_output", + "call_id": "call-function", + "output": "secret", + "status": "failed", + }, + ], +) +def test_blocked_retained_tool_output_rejects_malformed_replay_identity( + raw_item: dict[str, Any], +) -> None: + item = ToolCallOutputItem( + agent=Agent(name="test"), + raw_item=cast(Any, raw_item), + output="secret", + ) + model_response = ModelResponse( + output=[get_text_message("archived-secret")], + usage=Usage(), + response_id="response-malformed", + ) + + with pytest.raises( + AgentsException, match="Cannot sanitize a blocked tool output for replay" + ) as exc_info: + run_loop._retained_items_for_blocked_output([item], model_response) + + assert "secret" not in str(exc_info.value) + assert exc_info.value.run_data is None + assert exc_info.value.__cause__ is None + assert exc_info.value.__context__ is None + assert _sdk_exception_traceback_repr_locations(exc_info.value, "secret") == [] + assert item.output == _BLOCKED_TOOL_OUTPUT + assert item.custom_data is None + assert item.to_input_item() == { + "type": "function_call_output", + "call_id": "blocked-tool-output", + "output": _BLOCKED_TOOL_OUTPUT, + } + assert model_response.output == [] + + +@pytest.mark.parametrize( + "raw_item", + [ + { + "type": "function_call_output", + "call_id": "call-orphan", + "output": "secret", + "caller": {"type": "program", "caller_id": "program-orphan"}, + }, + { + "type": "program_output", + "id": "program-output-orphan", + "call_id": "program-orphan", + "result": "secret", + "status": "completed", + }, + ], +) +def test_blocked_retained_tool_output_rejects_orphan_program_relationship( + raw_item: dict[str, Any], +) -> None: + item = ToolCallOutputItem( + agent=Agent(name="test"), + raw_item=cast(Any, raw_item), + output="secret", + custom_data={"duplicate": "secret"}, + ) + + with pytest.raises( + AgentsException, match="Cannot sanitize a blocked tool output for replay" + ) as exc_info: + run_loop._retained_items_for_blocked_output([item]) + + assert "secret" not in str(exc_info.value) + assert exc_info.value.__cause__ is None + assert exc_info.value.__context__ is None + assert _sdk_exception_traceback_repr_locations(exc_info.value, "secret") == [] + assert item.output == _BLOCKED_TOOL_OUTPUT + assert item.custom_data is None + + +def test_blocked_retained_tool_output_accepts_program_parent_from_prior_history() -> None: + agent = Agent(name="test") + program_item = ToolCallItem( + agent=agent, + raw_item=cast( + Any, + { + "type": "program", + "call_id": "program-prior", + "code": "return await tools.lookup({});", + }, + ), + ) + child_call = ToolCallItem( + agent=agent, + raw_item=cast( + Any, + { + "type": "function_call", + "call_id": "call-child", + "name": "lookup", + "arguments": "{}", + "caller": {"type": "program", "caller_id": "program-prior"}, + }, + ), + ) + child_output = ToolCallOutputItem( + agent=agent, + raw_item=cast( + Any, + { + "type": "function_call_output", + "call_id": "call-child", + "output": "program-secret", + "caller": {"type": "program", "caller_id": "program-prior"}, + }, + ), + output="program-secret", + ) + + retained = run_loop._retained_items_for_blocked_output( + [child_call, child_output], + preceding_items=[program_item], + ) + + assert retained == [child_call, child_output] + assert child_output.output == _BLOCKED_TOOL_OUTPUT + assert program_item not in retained + + +@pytest.mark.asyncio +async def test_malformed_blocked_tool_output_failure_is_data_free_at_runner_boundary( + monkeypatch: pytest.MonkeyPatch, +) -> None: + original_sanitizer = run_loop._blocked_tool_output_payload + + def reject_with_malformed_caller(_raw_item: Any) -> dict[str, Any]: + return original_sanitizer( + { + "type": "function_call_output", + "call_id": "call-malformed", + "output": "runner-boundary-secret", + "caller": {"type": "program", "duplicate": "runner-boundary-secret"}, + } + ) + + monkeypatch.setattr(run_loop, "_blocked_tool_output_payload", reject_with_malformed_caller) + + @function_tool(name_override="terminal_tool") + def terminal_tool() -> str: + return "runner-boundary-secret" + + def output_guardrail( + _context: RunContextWrapper[Any], + _agent: Agent[Any], + output: Any, + ) -> GuardrailFunctionOutput: + return GuardrailFunctionOutput(output_info=output, tripwire_triggered=True) + + model = ScriptedModel( + [[get_function_tool_call("terminal_tool", "{}", call_id="call-terminal")]] + ) + agent = Agent( + name="test", + model=model, + tools=[terminal_tool], + tool_use_behavior="stop_on_first_tool", + output_guardrails=[OutputGuardrail(guardrail_function=output_guardrail)], + ) + + with pytest.raises( + AgentsException, match="Cannot sanitize a blocked tool output for replay" + ) as exc_info: + await Runner.run(agent, "Run terminal_tool") + + assert exc_info.value.run_data is None + assert exc_info.value.__cause__ is None + assert exc_info.value.__context__ is None + assert _sdk_exception_traceback_repr_locations(exc_info.value, "runner-boundary-secret") == [] + + +@pytest.mark.parametrize("shared_raw_item", [False, True]) +def test_blocked_retained_tool_output_sanitizes_archived_model_response( + shared_raw_item: bool, +) -> None: + agent = Agent(name="test") + raw_output = ProgramOutput( + id="program-output", + call_id="call-program", + result="secret-program-result", + status="completed", + type="program_output", + ) + item = ToolCallOutputItem( + agent=agent, + raw_item=raw_output if shared_raw_item else cast(Any, raw_output.model_dump()), + output=raw_output.result, + ) + program_item = ToolCallItem( + agent=agent, + raw_item=cast( + Any, + { + "type": "program", + "call_id": "call-program", + "code": "return await tools.lookup({});", + }, + ), + ) + model_response = ModelResponse(output=[raw_output], usage=Usage(), response_id="response") + state = make_run_state(agent) + state._model_responses = [model_response] + state._generated_items = [program_item, item] + state._session_items = [program_item, item] + + retained = run_loop._retained_items_for_blocked_output([program_item, item], model_response) + + assert retained == [program_item, item] + assert "secret-program-result" not in json.dumps(state.to_json()) + archived_output = cast(dict[str, Any], model_response.output[0]) + assert archived_output["type"] == "program_output" + assert archived_output["call_id"] == "call-program" + assert archived_output["result"] == _BLOCKED_TOOL_OUTPUT + + async def run_execute_approved_tools( agent: Agent[Any], approval_item: ToolApprovalItem, @@ -4434,7 +4891,7 @@ def guardrail_function( context: RunContextWrapper[Any], agent: Agent[Any], agent_output: Any ) -> GuardrailFunctionOutput: return GuardrailFunctionOutput( - output_info=None, + output_info="safe-info", tripwire_triggered=True, ) @@ -4446,9 +4903,12 @@ def guardrail_function( ) model.enqueue([get_text_message("user_message")]) - with pytest.raises(OutputGuardrailTripwireTriggered): + with pytest.raises(OutputGuardrailTripwireTriggered) as exc_info: await Runner.run(agent, input="user_message") + assert exc_info.value.guardrail_result.agent_output == "user_message" + assert exc_info.value.guardrail_result.output.output_info == "safe-info" + def test_output_guardrail_tripwire_does_not_save_assistant_message_to_session_sync() -> None: def guardrail_function( @@ -4591,10 +5051,10 @@ async def test_resumed_final_tool_sanitizes_output_after_output_guardrail_tripwi tool_calls = 0 def guardrail_function( - _context: RunContextWrapper[Any], _agent: Agent[Any], _agent_output: Any + _context: RunContextWrapper[Any], _agent: Agent[Any], agent_output: Any ) -> GuardrailFunctionOutput: return GuardrailFunctionOutput( - output_info=None, + output_info=agent_output, tripwire_triggered=tripwire_triggered, ) @@ -4626,8 +5086,10 @@ def commit_tool() -> str: model.enqueue([get_function_tool_call("commit_tool", "{}", call_id="call-second")]) if tripwire_triggered: - with pytest.raises(OutputGuardrailTripwireTriggered): + with pytest.raises(OutputGuardrailTripwireTriggered) as exc_info: await Runner.run(agent, state, session=session) + assert exc_info.value.guardrail_result.agent_output == _BLOCKED_TOOL_OUTPUT + assert exc_info.value.guardrail_result.output.output_info is None else: result = await Runner.run(agent, state, session=session) assert result.final_output == "committed-result-2" @@ -4655,6 +5117,161 @@ def commit_tool() -> str: assert "committed-result-2" not in serialized_state +@pytest.mark.parametrize("mode", ["non_streamed", "streamed"]) +@pytest.mark.parametrize("tripwire_triggered", [False, True]) +@pytest.mark.asyncio +async def test_terminal_tool_span_waits_for_output_guardrail_verdict( + mode: str, + tripwire_triggered: bool, + monkeypatch: pytest.MonkeyPatch, +) -> None: + started_function_spans: list[Any] = [] + original_on_span_start = SPAN_PROCESSOR_TESTING.on_span_start + + def capture_span_start(span: Any) -> None: + original_on_span_start(span) + if span.span_data.type == "function": + started_function_spans.append(span) + + monkeypatch.setattr(SPAN_PROCESSOR_TESTING, "on_span_start", capture_span_start) + + @function_tool(name_override="trace_tool") + def trace_tool() -> str: + return "trace-secret" + + def output_guardrail( + _context: RunContextWrapper[Any], + _agent: Agent[Any], + output: Any, + ) -> GuardrailFunctionOutput: + assert len(started_function_spans) == 1 + assert cast(Any, started_function_spans[0].span_data).output is None + assert not [span for span in fetch_ordered_spans() if span.span_data.type == "function"] + return GuardrailFunctionOutput( + output_info=output, + tripwire_triggered=tripwire_triggered, + ) + + model_output = [get_function_tool_call("trace_tool", "{}", call_id="call-trace")] + model = ScriptedModel( + [model_output if mode == "non_streamed" else get_exact_output_stream_step(model_output)] + ) + agent = Agent( + name="test", + model=model, + tools=[trace_tool], + tool_use_behavior="stop_on_first_tool", + output_guardrails=[OutputGuardrail(guardrail_function=output_guardrail)], + ) + run_config = RunConfig(trace_include_sensitive_data=True) + + if tripwire_triggered: + with pytest.raises(OutputGuardrailTripwireTriggered): + if mode == "non_streamed": + await Runner.run(agent, "Use trace_tool", run_config=run_config) + else: + result = Runner.run_streamed(agent, "Use trace_tool", run_config=run_config) + async for _ in result.stream_events(): + pass + elif mode == "non_streamed": + await Runner.run(agent, "Use trace_tool", run_config=run_config) + else: + result = Runner.run_streamed(agent, "Use trace_tool", run_config=run_config) + async for _ in result.stream_events(): + pass + + function_spans = [span for span in fetch_ordered_spans() if span.span_data.type == "function"] + assert len(function_spans) == 1 + expected_output = _BLOCKED_TOOL_OUTPUT if tripwire_triggered else "trace-secret" + assert cast(Any, function_spans[0].span_data).output == expected_output + assert cast(Any, started_function_spans[0].span_data).output == expected_output + if tripwire_triggered: + assert "trace-secret" not in json.dumps(function_spans[0].export()) + + +@pytest.mark.parametrize("mode", ["non_streamed", "streamed"]) +@pytest.mark.parametrize("tripwire_triggered", [False, True]) +@pytest.mark.asyncio +async def test_blocked_terminal_tool_span_discards_deferred_error( + mode: str, + tripwire_triggered: bool, + monkeypatch: pytest.MonkeyPatch, +) -> None: + started_function_spans: list[Any] = [] + original_on_span_start = SPAN_PROCESSOR_TESTING.on_span_start + + def capture_span_start(span: Any) -> None: + original_on_span_start(span) + if span.span_data.type == "function": + started_function_spans.append(span) + + monkeypatch.setattr(SPAN_PROCESSOR_TESTING, "on_span_start", capture_span_start) + + def expose_error(_context: RunContextWrapper[Any], error: Exception) -> str: + return str(error) + + @function_tool(name_override="failing_tool", failure_error_function=expose_error) + def failing_tool() -> str: + raise ValueError("span-error-secret") + + def output_guardrail( + _context: RunContextWrapper[Any], + _agent: Agent[Any], + output: Any, + ) -> GuardrailFunctionOutput: + assert output == "span-error-secret" + assert len(started_function_spans) == 1 + assert started_function_spans[0].error is None + return GuardrailFunctionOutput( + output_info=output, + tripwire_triggered=tripwire_triggered, + ) + + model_output = [get_function_tool_call("failing_tool", "{}", call_id="call-failing")] + model = ScriptedModel( + [model_output if mode == "non_streamed" else get_exact_output_stream_step(model_output)] + ) + agent = Agent( + name="test", + model=model, + tools=[failing_tool], + tool_use_behavior="stop_on_first_tool", + output_guardrails=[OutputGuardrail(guardrail_function=output_guardrail)], + ) + + async def run_once() -> None: + if mode == "non_streamed": + await Runner.run( + agent, + "Use failing_tool", + run_config=RunConfig(trace_include_sensitive_data=True), + ) + else: + result = Runner.run_streamed( + agent, + "Use failing_tool", + run_config=RunConfig(trace_include_sensitive_data=True), + ) + async for _ in result.stream_events(): + pass + + if tripwire_triggered: + with pytest.raises(OutputGuardrailTripwireTriggered): + await run_once() + else: + await run_once() + + function_spans = [span for span in fetch_ordered_spans() if span.span_data.type == "function"] + assert len(function_spans) == 1 + if tripwire_triggered: + assert function_spans[0].error is None + assert cast(Any, function_spans[0].span_data).output == _BLOCKED_TOOL_OUTPUT + assert "span-error-secret" not in json.dumps(function_spans[0].export()) + else: + assert function_spans[0].error is not None + assert "span-error-secret" in json.dumps(function_spans[0].export()) + + @pytest.mark.parametrize("behavior_kind", ["stop_at_tools", "custom"]) @pytest.mark.asyncio async def test_blocked_tool_output_is_sanitized_for_terminal_tool_behaviors( diff --git a/tests/test_agent_runner_streamed.py b/tests/test_agent_runner_streamed.py index 3b708bad9f..b2d20d7b1b 100644 --- a/tests/test_agent_runner_streamed.py +++ b/tests/test_agent_runner_streamed.py @@ -12,8 +12,10 @@ from openai import APIConnectionError, BadRequestError, NotFoundError from openai.types.responses import ( ResponseCompletedEvent, + ResponseCustomToolCall, ResponseErrorEvent, ResponseFailedEvent, + ResponseFunctionShellToolCallOutput, ResponseFunctionToolCall, ResponseIncompleteEvent, ) @@ -52,12 +54,13 @@ from agents.run import RunConfig from agents.run_internal import run_loop from agents.run_internal.run_loop import QueueCompleteSentinel +from agents.run_internal.run_steps import NextStepRunAgain from agents.stream_events import AgentUpdatedStreamEvent, RawResponsesStreamEvent, StreamEvent from agents.testing import ModelStep, ScriptedModel -from agents.tool import FunctionTool +from agents.tool import CustomTool, FunctionTool from agents.tool_guardrails import tool_input_guardrail, tool_output_guardrail from agents.usage import Usage, _attach_raw_usage_snapshot -from tests.model_test_helpers import get_response_obj +from tests.model_test_helpers import get_exact_output_stream_step, get_response_obj from .test_responses import ( get_final_output_message, @@ -78,6 +81,34 @@ _BLOCKED_TOOL_OUTPUT = "Output withheld by an output guardrail." +def _sdk_exception_traceback_string_locations( + error: BaseException, expected: str +) -> list[tuple[str, str]]: + pending = [error] + seen: set[int] = set() + locations: list[tuple[str, str]] = [] + while pending: + current = pending.pop() + if id(current) in seen: + continue + seen.add(id(current)) + + traceback = current.__traceback__ + while traceback is not None: + if "/src/agents/" in traceback.tb_frame.f_code.co_filename: + locations.extend( + (traceback.tb_frame.f_code.co_name, name) + for name, value in traceback.tb_frame.f_locals.items() + if isinstance(value, str) and value == expected + ) + traceback = traceback.tb_next + + for linked in (current.__cause__, current.__context__): + if linked is not None: + pending.append(linked) + return locations + + def _conversation_locked_error() -> BadRequestError: request = httpx.Request("POST", "https://example.com") response = httpx.Response( @@ -1872,7 +1903,7 @@ def guardrail_function( context: RunContextWrapper[Any], agent: Agent[Any], agent_output: Any ) -> GuardrailFunctionOutput: return GuardrailFunctionOutput( - output_info=None, + output_info="safe-info", tripwire_triggered=True, ) @@ -1884,11 +1915,16 @@ def guardrail_function( model=model, ) - with pytest.raises(OutputGuardrailTripwireTriggered): + with pytest.raises(OutputGuardrailTripwireTriggered) as exc_info: result = Runner.run_streamed(agent, input="user_message") async for _ in result.stream_events(): pass + assert exc_info.value.guardrail_result.agent_output == "first_test" + assert exc_info.value.guardrail_result.output.output_info == "safe-info" + assert result.output_guardrail_results[0].agent_output == "first_test" + assert result.output_guardrail_results[0].output.output_info == "safe-info" + @pytest.mark.asyncio async def test_output_guardrail_tripwire_raises_from_run_loop_task_before_stream_consumption(): @@ -2261,10 +2297,10 @@ def approval_tool() -> str: def output_guardrail( _context: RunContextWrapper[Any], _agent: Agent[Any], - _output: Any, + output: Any, ) -> GuardrailFunctionOutput: return GuardrailFunctionOutput( - output_info=None, + output_info=output, tripwire_triggered=guardrail_state["tripwire"], ) @@ -2292,8 +2328,10 @@ async def run_once(input_value: Any) -> Any: state.approve(first.interruptions[0]) if tripwire: - with pytest.raises(OutputGuardrailTripwireTriggered): + with pytest.raises(OutputGuardrailTripwireTriggered) as exc_info: await run_once(state) + assert exc_info.value.guardrail_result.agent_output == _BLOCKED_TOOL_OUTPUT + assert exc_info.value.guardrail_result.output.output_info is None else: resumed = await run_once(state) assert resumed.final_output == "approved-result" @@ -2319,6 +2357,7 @@ async def run_once(input_value: Any) -> Any: if tripwire: assert "approved-result" not in serialized_state assert _BLOCKED_TOOL_OUTPUT in serialized_state + assert state._current_step is None else: assert "approved-result" in serialized_state @@ -2343,6 +2382,155 @@ async def run_once(input_value: Any) -> Any: assert replayed_tool_items[1].get("output") == _BLOCKED_TOOL_OUTPUT +@pytest.mark.parametrize("mode", ["non_streamed", "streamed"]) +@pytest.mark.asyncio +async def test_resumed_blocked_tool_redacts_live_state_before_failed_session_save( + mode: str, +) -> None: + class FailingResumedTurnSession(SimpleListSession): + fail_writes = False + + async def add_items(self, items: list[TResponseInputItem]) -> None: + if self.fail_writes: + raise LookupError("session save failed") + await super().add_items(items) + + @function_tool(name_override="approval_tool", needs_approval=True) + def approval_tool() -> str: + return "state-secret" + + def output_guardrail( + _context: RunContextWrapper[Any], + _agent: Agent[Any], + output: Any, + ) -> GuardrailFunctionOutput: + return GuardrailFunctionOutput(output_info=output, tripwire_triggered=True) + + model = ScriptedModel( + [[get_function_tool_call("approval_tool", "{}", call_id="call-approved")]] + ) + agent = Agent( + name="test", + model=model, + tools=[approval_tool], + tool_use_behavior="stop_on_first_tool", + output_guardrails=[OutputGuardrail(guardrail_function=output_guardrail)], + ) + session = FailingResumedTurnSession() + first = await Runner.run(agent, "Use approval_tool", session=session) + state = first.to_state() + state.approve(first.interruptions[0]) + session.fail_writes = True + + with pytest.raises(LookupError, match="session save failed"): + if mode == "non_streamed": + await Runner.run(agent, state, session=session) + else: + result = Runner.run_streamed(agent, state, session=session) + await consume_stream(result) + + assert state._current_step is not None + assert getattr(state._current_step, "output", None) == _BLOCKED_TOOL_OUTPUT + assert "state-secret" not in json.dumps(state.to_json()) + + +@pytest.mark.asyncio +async def test_resumed_blocked_tool_session_save_cancellation_remains_observable() -> None: + class CancellingResumedTurnSession(SimpleListSession): + cancel_writes = False + + async def add_items(self, items: list[TResponseInputItem]) -> None: + if self.cancel_writes: + raise asyncio.CancelledError("session-secret") + await super().add_items(items) + + @function_tool(name_override="approval_tool", needs_approval=True) + def approval_tool() -> str: + return "state-secret" + + def output_guardrail( + _context: RunContextWrapper[Any], + _agent: Agent[Any], + output: Any, + ) -> GuardrailFunctionOutput: + return GuardrailFunctionOutput(output_info=output, tripwire_triggered=True) + + model = ScriptedModel( + [[get_function_tool_call("approval_tool", "{}", call_id="call-approved")]] + ) + agent = Agent( + name="test", + model=model, + tools=[approval_tool], + tool_use_behavior="stop_on_first_tool", + output_guardrails=[OutputGuardrail(guardrail_function=output_guardrail)], + ) + session = CancellingResumedTurnSession() + first = await Runner.run(agent, "Use approval_tool", session=session) + state = first.to_state() + state.approve(first.interruptions[0]) + session.cancel_writes = True + + result = Runner.run_streamed(agent, state, session=session) + with pytest.raises(asyncio.CancelledError) as exc_info: + await consume_stream(result) + + assert result._cancel_mode == "none" + assert result._stored_exception is exc_info.value + assert "session-secret" not in str(exc_info.value) + assert state._current_step is not None + assert getattr(state._current_step, "output", None) == _BLOCKED_TOOL_OUTPUT + serialized_state = json.dumps(state.to_json()) + assert "state-secret" not in serialized_state + assert "session-secret" not in serialized_state + + +@pytest.mark.parametrize("mode", ["non_streamed", "streamed"]) +@pytest.mark.asyncio +async def test_resumed_non_tool_tripwire_preserves_live_final_step(mode: str) -> None: + @function_tool(name_override="approval_tool", needs_approval=True) + def approval_tool() -> str: + return "approved" + + def output_guardrail( + _context: RunContextWrapper[Any], + _agent: Agent[Any], + output: Any, + ) -> GuardrailFunctionOutput: + return GuardrailFunctionOutput(output_info=output, tripwire_triggered=True) + + text_output = [get_text_message("blocked-text")] + second_step: Any = ( + text_output if mode == "non_streamed" else get_exact_output_stream_step(text_output) + ) + model = ScriptedModel( + [ + [get_function_tool_call("approval_tool", "{}", call_id="call-approved")], + second_step, + ] + ) + agent = Agent( + name="test", + model=model, + tools=[approval_tool], + output_guardrails=[OutputGuardrail(guardrail_function=output_guardrail)], + ) + first = await Runner.run(agent, "Use approval_tool") + state = first.to_state() + state.approve(first.interruptions[0]) + + with pytest.raises(OutputGuardrailTripwireTriggered) as exc_info: + if mode == "non_streamed": + await Runner.run(agent, state) + else: + result = Runner.run_streamed(agent, state) + await consume_stream(result) + + assert isinstance(state._current_step, NextStepRunAgain) + assert exc_info.value.guardrail_result.agent_output == "blocked-text" + assert exc_info.value.guardrail_result.output.output_info == "blocked-text" + + @pytest.mark.parametrize("mode", ["non_streamed", "streamed"]) @pytest.mark.asyncio async def test_stop_on_first_tool_final_persists_committed_tool_items_on_tripwire( @@ -2360,9 +2548,9 @@ def commit_tool() -> str: def output_guardrail( _context: RunContextWrapper[Any], _agent: Agent[Any], - _output: Any, + output: Any, ) -> GuardrailFunctionOutput: - return GuardrailFunctionOutput(output_info=None, tripwire_triggered=True) + return GuardrailFunctionOutput(output_info=output, tripwire_triggered=True) model = ScriptedModel() model.enqueue([get_function_tool_call("commit_tool", "{}", call_id="call-committed")]) @@ -2375,14 +2563,26 @@ def output_guardrail( ) session = SimpleListSession() - with pytest.raises(OutputGuardrailTripwireTriggered): + streamed_result: Any = None + with pytest.raises(OutputGuardrailTripwireTriggered) as exc_info: if mode == "non_streamed": await Runner.run(agent, "Use commit_tool", session=session) else: - result = Runner.run_streamed(agent, "Use commit_tool", session=session) - await consume_stream(result) + streamed_result = Runner.run_streamed(agent, "Use commit_tool", session=session) + await consume_stream(streamed_result) assert calls == ["ran"], "the tool never ran, so the test proves nothing" + assert exc_info.value.guardrail_result.agent_output == _BLOCKED_TOOL_OUTPUT + assert exc_info.value.guardrail_result.output.output_info is None + assert exc_info.value.__cause__ is None + assert exc_info.value.__context__ is None + assert _sdk_exception_traceback_string_locations(exc_info.value, "sensitive-result") == [] + if streamed_result is not None: + assert all( + result.agent_output == _BLOCKED_TOOL_OUTPUT + for result in streamed_result.output_guardrail_results + ) + assert "sensitive-result" not in json.dumps(streamed_result.to_state().to_json()) saved_items = await session.get_items() assert "sensitive-result" not in json.dumps(saved_items) @@ -2429,6 +2629,216 @@ def output_guardrail( assert "sensitive-result" not in json.dumps(model_input) +@pytest.mark.parametrize("mode", ["non_streamed", "streamed"]) +@pytest.mark.asyncio +async def test_blocked_terminal_turn_sanitizes_concurrent_custom_tool_output(mode: str) -> None: + function_calls = 0 + custom_calls = 0 + + @function_tool(name_override="terminal_tool") + def terminal_tool() -> str: + nonlocal function_calls + function_calls += 1 + return "function-secret" + + def run_custom_tool(_context: Any, _input: str) -> str: + nonlocal custom_calls + custom_calls += 1 + return "custom-secret" + + custom_tool = CustomTool( + name="custom_side_effect", + description="Return a custom result.", + on_invoke_tool=run_custom_tool, + format={"type": "text"}, + ) + + def output_guardrail( + _context: RunContextWrapper[Any], + _agent: Agent[Any], + _output: Any, + ) -> GuardrailFunctionOutput: + return GuardrailFunctionOutput(output_info=None, tripwire_triggered=True) + + model_output = [ + get_function_tool_call("terminal_tool", "{}", call_id="call-function"), + ResponseCustomToolCall( + type="custom_tool_call", + name="custom_side_effect", + call_id="call-custom", + input="custom input", + ), + ] + model = ScriptedModel( + [model_output if mode == "non_streamed" else get_exact_output_stream_step(model_output)] + ) + session = SimpleListSession() + agent = Agent( + name="test", + model=model, + tools=[terminal_tool, custom_tool], + tool_use_behavior="stop_on_first_tool", + output_guardrails=[OutputGuardrail(guardrail_function=output_guardrail)], + ) + + with pytest.raises(OutputGuardrailTripwireTriggered): + if mode == "non_streamed": + await Runner.run(agent, "Run both tools", session=session) + else: + result = Runner.run_streamed(agent, "Run both tools", session=session) + await consume_stream(result) + + assert function_calls == 1 + assert custom_calls == 1 + saved_items = await session.get_items() + serialized_items = json.dumps(saved_items) + assert "function-secret" not in serialized_items + assert "custom-secret" not in serialized_items + + saved_outputs = { + cast(dict[str, Any], item).get("type"): cast(dict[str, Any], item).get("output") + for item in saved_items + if isinstance(item, dict) + and item.get("type") in {"function_call_output", "custom_tool_call_output"} + } + assert saved_outputs == { + "function_call_output": _BLOCKED_TOOL_OUTPUT, + "custom_tool_call_output": _BLOCKED_TOOL_OUTPUT, + } + + +@pytest.mark.parametrize("mode", ["non_streamed", "streamed"]) +@pytest.mark.parametrize("behavior", ["allow", "reject_content"]) +@pytest.mark.asyncio +async def test_blocked_terminal_tool_sanitizes_tool_output_guardrail_aliases( + mode: str, + behavior: str, +) -> None: + guardrail_outputs: list[ToolGuardrailFunctionOutput] = [] + + @tool_output_guardrail + def retain_tool_output(data: ToolOutputGuardrailData) -> ToolGuardrailFunctionOutput: + if behavior == "allow": + output = ToolGuardrailFunctionOutput.allow(output_info=data.output) + else: + output = ToolGuardrailFunctionOutput.reject_content( + message=data.output, + output_info=data.output, + ) + guardrail_outputs.append(output) + return output + + @function_tool( + name_override="terminal_tool", + tool_output_guardrails=[retain_tool_output], + ) + def terminal_tool() -> str: + return "tool-guardrail-secret" + + def output_guardrail( + _context: RunContextWrapper[Any], + _agent: Agent[Any], + output: Any, + ) -> GuardrailFunctionOutput: + return GuardrailFunctionOutput(output_info=output, tripwire_triggered=True) + + model_output = [get_function_tool_call("terminal_tool", "{}", call_id="call-terminal")] + model = ScriptedModel( + [model_output if mode == "non_streamed" else get_exact_output_stream_step(model_output)] + ) + agent = Agent( + name="test", + model=model, + tools=[terminal_tool], + tool_use_behavior="stop_on_first_tool", + output_guardrails=[OutputGuardrail(guardrail_function=output_guardrail)], + ) + streamed_result: Any = None + + with pytest.raises(OutputGuardrailTripwireTriggered): + if mode == "non_streamed": + await Runner.run(agent, "Use terminal_tool") + else: + streamed_result = Runner.run_streamed(agent, "Use terminal_tool") + await consume_stream(streamed_result) + + assert len(guardrail_outputs) == 1 + assert guardrail_outputs[0].output_info is None + if behavior == "allow": + assert guardrail_outputs[0].behavior == {"type": "allow"} + else: + assert guardrail_outputs[0].behavior == { + "type": "reject_content", + "message": _BLOCKED_TOOL_OUTPUT, + } + if streamed_result is not None: + assert "tool-guardrail-secret" not in json.dumps(streamed_result.to_state().to_json()) + + +@pytest.mark.parametrize("status", ["in_progress", "incomplete"]) +@pytest.mark.asyncio +async def test_blocked_streamed_terminal_turn_accepts_provider_shell_status( + status: str, +) -> None: + @function_tool(name_override="terminal_tool") + def terminal_tool() -> str: + return "terminal-secret" + + def output_guardrail( + _context: RunContextWrapper[Any], + _agent: Agent[Any], + output: Any, + ) -> GuardrailFunctionOutput: + return GuardrailFunctionOutput(output_info=output, tripwire_triggered=True) + + shell_output = ResponseFunctionShellToolCallOutput( + id=f"shell-output-{status}", + call_id="call-shell", + output=[ + { + "stdout": "shell-secret", + "stderr": "", + "outcome": {"type": "exit", "exit_code": 0}, + } + ], + status=cast(Any, status), + type="shell_call_output", + ) + model_output = [ + get_function_tool_call("terminal_tool", "{}", call_id="call-terminal"), + shell_output, + ] + model = ScriptedModel([get_exact_output_stream_step(model_output)]) + session = SimpleListSession() + agent = Agent( + name="test", + model=model, + tools=[terminal_tool], + tool_use_behavior="stop_on_first_tool", + output_guardrails=[OutputGuardrail(guardrail_function=output_guardrail)], + ) + + result = Runner.run_streamed(agent, "Run terminal_tool", session=session) + with pytest.raises(OutputGuardrailTripwireTriggered) as exc_info: + await consume_stream(result) + + assert exc_info.value.guardrail_result.agent_output == _BLOCKED_TOOL_OUTPUT + assert exc_info.value.guardrail_result.output.output_info is None + assert all( + guardrail_result.agent_output == _BLOCKED_TOOL_OUTPUT + and guardrail_result.output.output_info is None + for guardrail_result in result.output_guardrail_results + ) + serialized_state = json.dumps(result.to_state().to_json()) + serialized_session = json.dumps(await session.get_items()) + assert "terminal-secret" not in serialized_state + assert "shell-secret" not in serialized_state + assert "terminal-secret" not in serialized_session + assert "shell-secret" not in serialized_session + assert _BLOCKED_TOOL_OUTPUT in serialized_state + assert _BLOCKED_TOOL_OUTPUT in serialized_session + + @pytest.mark.parametrize("mode", ["non_streamed", "streamed"]) @pytest.mark.asyncio async def test_blocked_tool_output_cannot_be_forwarded_by_a_later_tool(mode: str) -> None: @@ -2743,6 +3153,48 @@ def output_guardrail( assert await session.get_items() == [{"content": "Hello", "role": "user"}] +@pytest.mark.asyncio +async def test_blocked_tool_output_redacts_live_state_before_failed_session_save() -> None: + guardrail_tripped = False + + class FailingBlockedTurnSession(SimpleListSession): + async def add_items(self, items: list[TResponseInputItem]) -> None: + if guardrail_tripped: + raise LookupError("session save failed") + await super().add_items(items) + + @function_tool(name_override="commit_tool") + def commit_tool() -> str: + return "state-secret" + + def output_guardrail( + _context: RunContextWrapper[Any], + _agent: Agent[Any], + output: Any, + ) -> GuardrailFunctionOutput: + nonlocal guardrail_tripped + guardrail_tripped = True + return GuardrailFunctionOutput(output_info=output, tripwire_triggered=True) + + model = ScriptedModel([[get_function_tool_call("commit_tool", "{}", call_id="call-committed")]]) + agent = Agent( + name="test", + model=model, + tools=[commit_tool], + tool_use_behavior="stop_on_first_tool", + output_guardrails=[OutputGuardrail(guardrail_function=output_guardrail)], + ) + session = FailingBlockedTurnSession() + result = Runner.run_streamed(agent, "Use commit_tool", session=session) + + with pytest.raises(LookupError, match="session save failed"): + await consume_stream(result) + + assert result._state is not None + assert result._state._current_step is None + assert "state-secret" not in json.dumps(result.to_state().to_json()) + + @pytest.mark.asyncio async def test_streamed_session_save_cancellation_is_not_a_public_immediate_cancel() -> None: guardrail_failed = False From a37fc36900ebf6ee7c6962baffdc770f58fc2ee0 Mon Sep 17 00:00:00 2001 From: Kazuhiro Sera Date: Wed, 19 Aug 2026 04:11:04 +0900 Subject: [PATCH 03/12] Prevent replay of blocked terminal tool output --- src/agents/run.py | 254 ++-- .../run_internal/agent_runner_helpers.py | 21 + src/agents/run_internal/blocked_output.py | 106 ++ src/agents/run_internal/run_loop.py | 1343 ++++++++++------- tests/test_agent_runner.py | 1052 ++++++------- tests/test_agent_runner_streamed.py | 1010 ++++--------- tests/test_error_logging_redaction.py | 172 +-- tests/test_max_turns.py | 116 +- 8 files changed, 1910 insertions(+), 2164 deletions(-) create mode 100644 src/agents/run_internal/blocked_output.py diff --git a/src/agents/run.py b/src/agents/run.py index 8591b7dd59..c955c0484a 100644 --- a/src/agents/run.py +++ b/src/agents/run.py @@ -3,7 +3,6 @@ import asyncio import contextlib import warnings -from functools import partial from typing import TYPE_CHECKING, Any, cast from typing_extensions import Unpack @@ -35,7 +34,6 @@ ItemHelpers, ModelResponse, RunItem, - ToolCallOutputItem, TResponseInputItem, ) from .lifecycle import RunHooks @@ -79,6 +77,7 @@ snapshot_usage, update_run_state_for_interruption, usage_delta, + validate_output_guardrails_with_server_managed_conversation, validate_session_conversation_settings, ) from .run_internal.approvals import approvals_from_step @@ -96,12 +95,14 @@ from .run_internal.prompt_cache_key import PromptCacheKeyResolver from .run_internal.run_grouping import resolve_run_grouping_id from .run_internal.run_loop import ( - _finish_blocked_output_tool_spans, - _redact_blocked_output_state_step, - _retained_items_for_blocked_output, - _run_with_deferred_tool_spans, + _BlockedOutputOwnerStarts, + _current_response_boundary, + _final_turn_items_for_persistence, + _is_terminal_tool_output_response, + _retained_items_for_blocked_response, _sanitize_blocked_output_guardrail_results, - _sanitize_blocked_tool_output_guardrail_results, + _should_defer_interrupted_session_items, + _validate_resumed_session_output_guardrail_safety, cleanup_models_after_run, finalize_max_turns_handler_output, get_all_tools, @@ -134,7 +135,6 @@ session_items_for_turn, update_run_state_after_resume, ) -from .run_internal.tool_execution import finish_deferred_tool_spans from .run_internal.tool_use_tracker import ( AgentToolUseTracker, hydrate_tool_use_tracker, @@ -900,6 +900,12 @@ def _mark_response_hooks_started() -> None: current_agent = run_state._current_agent else: current_agent = starting_agent + _validate_resumed_session_output_guardrail_safety( + agent=current_agent, + run_config=run_config, + session=session, + run_state=run_state if is_resumed_state else None, + ) sandbox_runtime.assert_agent_supported(current_agent) should_run_agent_start_hooks = True store_setting = current_agent.model_settings.resolve( @@ -945,6 +951,13 @@ def _mark_response_hooks_started() -> None: try: while True: + validate_output_guardrails_with_server_managed_conversation( + current_agent, + run_config, + conversation_id=conversation_id, + previous_response_id=previous_response_id, + auto_previous_response_id=auto_previous_response_id, + ) if TYPE_CHECKING: # Keep loop-carried types explicit to bound Pyright's flow analysis. original_input = cast( # type: ignore[redundant-cast] @@ -1048,25 +1061,36 @@ def _mark_response_hooks_started() -> None: ) raise UserError("No processed response found in previous state") - turn_result = await _run_with_deferred_tool_spans( - agent=current_agent, + resumed_response_boundary = _current_response_boundary( + (), + run_state._last_processed_response, + run_state, + ) + blocked_output_owner_starts = _BlockedOutputOwnerStarts( + run_state_generated_items=( + resumed_response_boundary.generated_start + ), + run_state_session_items=resumed_response_boundary.session_start, + run_state_model_responses=len(run_state._model_responses) - 1, + run_state_tool_output_guardrail_results=len( + run_state._tool_output_guardrail_results + ), + ) + + turn_result = await resolve_interrupted_turn( + bindings=current_bindings, + original_input=original_input, + original_pre_step_items=generated_items, + new_response=run_state._model_responses[-1], + processed_response=run_state._last_processed_response, + hooks=hooks, + context_wrapper=context_wrapper, run_config=run_config, - run_turn=partial( - resolve_interrupted_turn, - bindings=current_bindings, - original_input=original_input, - original_pre_step_items=generated_items, - new_response=run_state._model_responses[-1], - processed_response=run_state._last_processed_response, - hooks=hooks, - context_wrapper=context_wrapper, - run_config=run_config, - server_manages_conversation=( - server_conversation_tracker is not None - ), - run_state=run_state, - error_handlers=error_handlers, + server_manages_conversation=( + server_conversation_tracker is not None ), + run_state=run_state, + error_handlers=error_handlers, ) if run_state._last_processed_response is not None: @@ -1102,6 +1126,13 @@ def _mark_response_hooks_started() -> None: and turn_session_items and run_state is not None and not isinstance(turn_result.next_step, NextStepFinalOutput) + and not ( + isinstance(turn_result.next_step, NextStepInterruption) + and _should_defer_interrupted_session_items( + current_agent, + run_config, + ) + ) ): run_state._current_turn_persisted_item_count = ( await save_resumed_turn_items( @@ -1184,6 +1215,11 @@ def _mark_response_hooks_started() -> None: ) if isinstance(turn_result.next_step, NextStepFinalOutput): + current_processed_response = ( + turn_result.processed_response + if turn_result.processed_response is not None + else run_state._last_processed_response + ) output_guardrail_result_start = len(output_guardrail_results) try: await run_output_guardrails( @@ -1195,28 +1231,25 @@ def _mark_response_hooks_started() -> None: output_guardrail_results, ) except OutputGuardrailTripwireTriggered as exc: - _finish_blocked_output_tool_spans( - turn_result.deferred_tool_spans + if not _is_terminal_tool_output_response( + turn_session_items, + current_processed_response, + run_state, + ): + raise + sanitized_results = _sanitize_blocked_output_guardrail_results( + output_guardrail_results[output_guardrail_result_start:], + exc, ) - has_tool_output = any( - isinstance(item, ToolCallOutputItem) - for item in turn_session_items + output_guardrail_results[output_guardrail_result_start:] = ( + sanitized_results ) - if has_tool_output: - _sanitize_blocked_output_guardrail_results( - output_guardrail_results[ - output_guardrail_result_start: - ], - exc, - ) - _sanitize_blocked_tool_output_guardrail_results( - turn_result.tool_output_guardrail_results - ) - _redact_blocked_output_state_step(run_state) - retained_items = _retained_items_for_blocked_output( + retained_items = _retained_items_for_blocked_response( turn_session_items, turn_result.model_response, - turn_result.pre_step_items, + run_state, + current_processed_response, + owner_starts=blocked_output_owner_starts, ) await save_final_turn_items_after_guardrails( session=session, @@ -1230,34 +1263,24 @@ def _mark_response_hooks_started() -> None: store=store_setting, wrapper=context_wrapper, ) - if has_tool_output: - run_state._current_step = None raise except (Exception, asyncio.CancelledError): - finish_deferred_tool_spans(turn_result.deferred_tool_spans) - # An ordinary guardrail failure leaves the verdict unknown, so - # preserve the completed turn exactly as fresh execution does. - await save_final_turn_items_after_guardrails( - session=session, - run_state=run_state, - session_persistence_enabled=session_persistence_enabled, - input_guardrail_results=( - _attempt_input_guardrail_results() - ), - items=turn_session_items, - response_id=turn_result.model_response.response_id, - store=store_setting, - wrapper=context_wrapper, - ) + # Without a verdict, do not persist any part of this response. raise - finish_deferred_tool_spans(turn_result.deferred_tool_spans) + final_turn_items = _final_turn_items_for_persistence( + turn_session_items, + current_processed_response, + run_state, + current_agent, + run_config, + ) await save_final_turn_items_after_guardrails( session=session, run_state=run_state, session_persistence_enabled=session_persistence_enabled, input_guardrail_results=_attempt_input_guardrail_results(), - items=turn_session_items, + items=final_turn_items, response_id=turn_result.model_response.response_id, store=store_setting, wrapper=context_wrapper, @@ -1431,8 +1454,6 @@ async def _save_max_turns_handler_output( output=handler_result.final_output, context_wrapper=context_wrapper, output_guardrail_results=output_guardrail_results, - save_items_after_guardrails=_save_max_turns_handler_output, - include_in_history=include_in_history, ) if include_in_history and not handler_output_recorded: await _save_max_turns_handler_output([synthesized_item]) @@ -1464,7 +1485,9 @@ async def _save_max_turns_handler_output( result._original_input = copy_input_items(original_input) return _finalize_result(result) - if run_state is not None and not resuming_turn: + if run_state is not None and ( + not resuming_turn or isinstance(run_state._current_step, NextStepRunAgain) + ): run_state._current_turn_persisted_item_count = 0 logger.debug("Running agent %s (turn %s)", current_agent.name, current_turn) @@ -1477,6 +1500,23 @@ async def _save_max_turns_handler_output( except Exception: last_saved_input_snapshot_for_rewind = None + blocked_output_owner_starts = _BlockedOutputOwnerStarts( + run_state_generated_items=( + len(run_state._generated_items) if run_state is not None else None + ), + run_state_session_items=( + len(run_state._session_items) if run_state is not None else None + ), + run_state_model_responses=( + len(run_state._model_responses) if run_state is not None else None + ), + run_state_tool_output_guardrail_results=( + len(run_state._tool_output_guardrail_results) + if run_state is not None + else None + ), + ) + items_for_model = ( pending_server_items if server_conversation_tracker is not None and pending_server_items @@ -1747,24 +1787,25 @@ async def _save_max_turns_handler_output( output_guardrail_results, ) except OutputGuardrailTripwireTriggered as exc: - _finish_blocked_output_tool_spans(turn_result.deferred_tool_spans) - has_tool_output = any( - isinstance(item, ToolCallOutputItem) - for item in items_to_save_turn + if not _is_terminal_tool_output_response( + turn_session_items, + turn_result.processed_response, + run_state, + ): + raise + sanitized_results = _sanitize_blocked_output_guardrail_results( + output_guardrail_results[output_guardrail_result_start:], + exc, ) - if has_tool_output: - _sanitize_blocked_output_guardrail_results( - output_guardrail_results[output_guardrail_result_start:], - exc, - ) - _sanitize_blocked_tool_output_guardrail_results( - turn_result.tool_output_guardrail_results - ) - _redact_blocked_output_state_step(run_state) - retained_items = _retained_items_for_blocked_output( - items_to_save_turn, + output_guardrail_results[output_guardrail_result_start:] = ( + sanitized_results + ) + retained_items = _retained_items_for_blocked_response( + turn_session_items, turn_result.model_response, - turn_result.pre_step_items, + run_state, + turn_result.processed_response, + owner_starts=blocked_output_owner_starts, ) await save_final_turn_items_after_guardrails( session=session, @@ -1776,32 +1817,24 @@ async def _save_max_turns_handler_output( store=store_setting, wrapper=context_wrapper, ) - if has_tool_output and run_state is not None: - run_state._current_step = None raise except (Exception, asyncio.CancelledError): - finish_deferred_tool_spans(turn_result.deferred_tool_spans) - # Preserve the released non-stream behavior for guardrail errors - # and cancellation: the completed final turn remains replayable. - await save_final_turn_items_after_guardrails( - session=session, - run_state=run_state, - session_persistence_enabled=session_persistence_enabled, - input_guardrail_results=_attempt_input_guardrail_results(), - items=items_to_save_turn, - response_id=turn_result.model_response.response_id, - store=store_setting, - wrapper=context_wrapper, - ) + # Without a verdict, do not persist any part of this response. raise - finish_deferred_tool_spans(turn_result.deferred_tool_spans) + final_turn_items = _final_turn_items_for_persistence( + turn_session_items, + turn_result.processed_response, + run_state, + current_agent, + run_config, + ) await save_final_turn_items_after_guardrails( session=session, run_state=run_state, session_persistence_enabled=session_persistence_enabled, input_guardrail_results=_attempt_input_guardrail_results(), - items=items_to_save_turn, + items=final_turn_items, response_id=turn_result.model_response.response_id, store=store_setting, wrapper=context_wrapper, @@ -1840,7 +1873,12 @@ async def _save_max_turns_handler_output( run_state._current_step = None return _finalize_result(result) elif isinstance(turn_result.next_step, NextStepInterruption): - if session_persistence_enabled: + if session_persistence_enabled and not ( + _should_defer_interrupted_session_items( + current_agent, + run_config, + ) + ): if not input_guardrails_triggered( _attempt_input_guardrail_results() ): @@ -2236,6 +2274,19 @@ def run_streamed( if run_state is not None: run_state._reasoning_item_id_policy = resolved_reasoning_item_id_policy + schema_agent = ( + run_state._current_agent + if run_state is not None and run_state._current_agent is not None + else starting_agent + ) + validate_output_guardrails_with_server_managed_conversation( + schema_agent, + run_config, + conversation_id=conversation_id, + previous_response_id=previous_response_id, + auto_previous_response_id=auto_previous_response_id, + ) + ( trace_workflow_name, trace_id, @@ -2271,11 +2322,6 @@ def run_streamed( run_state=run_state, ) - schema_agent = ( - run_state._current_agent - if run_state is not None and run_state._current_agent is not None - else starting_agent - ) sandbox_runtime.assert_agent_supported(schema_agent) output_schema = get_output_schema(schema_agent) diff --git a/src/agents/run_internal/agent_runner_helpers.py b/src/agents/run_internal/agent_runner_helpers.py index 7d5b73ad5a..6c564230d3 100644 --- a/src/agents/run_internal/agent_runner_helpers.py +++ b/src/agents/run_internal/agent_runner_helpers.py @@ -55,6 +55,7 @@ "finalize_conversation_tracking", "get_unsent_tool_call_ids_for_interrupted_state", "input_guardrails_triggered", + "validate_output_guardrails_with_server_managed_conversation", "validate_session_conversation_settings", "resolve_trace_settings", "resolve_processed_response", @@ -257,6 +258,26 @@ def validate_session_conversation_settings( ) +def validate_output_guardrails_with_server_managed_conversation( + agent: Agent[Any], + run_config: RunConfig, + *, + conversation_id: str | None, + previous_response_id: str | None, + auto_previous_response_id: bool, +) -> None: + """Reject an output-guardrail run whose rejected history cannot be locally replaced.""" + if conversation_id is None and previous_response_id is None and not auto_previous_response_id: + return + if not agent.output_guardrails and not run_config.output_guardrails: + return + raise UserError( + "Output guardrails cannot be combined with conversation_id, previous_response_id, " + "or auto_previous_response_id because rejected output cannot be removed from " + "server-managed conversation history." + ) + + def resolve_trace_settings( *, run_state: RunState[TContext] | None, diff --git a/src/agents/run_internal/blocked_output.py b/src/agents/run_internal/blocked_output.py new file mode 100644 index 0000000000..084935f5dc --- /dev/null +++ b/src/agents/run_internal/blocked_output.py @@ -0,0 +1,106 @@ +"""Canonical data-free function-tool payloads rejected by an output guardrail.""" + +from __future__ import annotations + +from typing import Any + +from openai.types.responses import ResponseFunctionToolCall + +from ..exceptions import AgentsException + +OUTPUT_GUARDRAIL_BLOCKED_TOOL_OUTPUT = "Output withheld by an output guardrail." + +_RESPONSE_OUTPUT_STATUSES = frozenset({"in_progress", "completed", "incomplete"}) + + +def _payload_field(raw_item: Any, field: str) -> Any: + """Read an allowlisted field without copying extras or invoking instance hooks.""" + if type(raw_item) is dict: + return dict.get(raw_item, field) + if type(raw_item) is ResponseFunctionToolCall: + values = object.__getattribute__(raw_item, "__dict__") + return dict.get(values, field) + raise AgentsException("Cannot sanitize an unsupported tool item variant.") + + +def _required_string(raw_item: Any, field: str) -> str: + value = _payload_field(raw_item, field) + if type(value) is not str or not value: + raise AgentsException(f"Cannot sanitize a function tool item without {field}.") + return value + + +def _copy_optional_string( + sanitized: dict[str, Any], + raw_item: Any, + field: str, +) -> None: + value = _payload_field(raw_item, field) + if value is None: + return + if type(value) is not str or not value: + raise AgentsException(f"Cannot sanitize a function tool item with an invalid {field}.") + sanitized[field] = value + + +def _copy_optional_status(sanitized: dict[str, Any], raw_item: Any) -> None: + status = _payload_field(raw_item, "status") + if status is None: + return + if type(status) is not str or status not in _RESPONSE_OUTPUT_STATUSES: + raise AgentsException("Cannot sanitize a function tool item with an invalid status.") + sanitized["status"] = status + + +def _copy_optional_direct_caller(sanitized: dict[str, Any], raw_item: Any) -> None: + caller = _payload_field(raw_item, "caller") + if caller is None: + return + if type(caller) is dict and dict.get(caller, "type") == "direct": + sanitized["caller"] = {"type": "direct"} + return + raise AgentsException("Cannot sanitize a function tool item with a non-direct caller.") + + +def blocked_function_call_payload(raw_item: Any) -> dict[str, Any]: + """Build a provider-valid function call from explicitly allowlisted fields.""" + if _payload_field(raw_item, "type") != "function_call": + raise AgentsException("Cannot sanitize an unsupported tool call variant.") + sanitized: dict[str, Any] = { + "type": "function_call", + "name": _required_string(raw_item, "name"), + "arguments": _required_string(raw_item, "arguments"), + "call_id": _required_string(raw_item, "call_id"), + } + _copy_optional_string(sanitized, raw_item, "id") + _copy_optional_string(sanitized, raw_item, "namespace") + _copy_optional_status(sanitized, raw_item) + _copy_optional_direct_caller(sanitized, raw_item) + try: + validated = ResponseFunctionToolCall(**sanitized) + except Exception: + raise AgentsException("Sanitized function_call is not valid for replay.") from None + return validated.model_dump(exclude_unset=True) + + +def blocked_function_output_payload(raw_item: Any) -> dict[str, Any]: + """Build a replay-valid function output from explicitly allowlisted fields.""" + if _payload_field(raw_item, "type") != "function_call_output": + raise AgentsException("Cannot sanitize an unsupported tool output variant.") + sanitized: dict[str, Any] = { + "type": "function_call_output", + "call_id": _required_string(raw_item, "call_id"), + "output": OUTPUT_GUARDRAIL_BLOCKED_TOOL_OUTPUT, + } + _copy_optional_string(sanitized, raw_item, "id") + _copy_optional_status(sanitized, raw_item) + _copy_optional_direct_caller(sanitized, raw_item) + try: + from ..run_state import _deserialize_tool_call_output_raw_item + + restored = _deserialize_tool_call_output_raw_item(sanitized) + except Exception: + raise AgentsException("Sanitized function_call_output is not valid for replay.") from None + if restored is None: + raise AgentsException("Sanitized function_call_output is not valid for replay.") + return sanitized diff --git a/src/agents/run_internal/run_loop.py b/src/agents/run_internal/run_loop.py index 0a27ed8ad1..e805119d9f 100644 --- a/src/agents/run_internal/run_loop.py +++ b/src/agents/run_internal/run_loop.py @@ -7,7 +7,7 @@ import asyncio import dataclasses as _dc -from collections.abc import Awaitable, Callable, Mapping, Sequence +from collections.abc import Awaitable, Callable, Sequence from contextlib import aclosing from functools import partial from typing import Any, TypeVar, cast @@ -45,7 +45,7 @@ _mark_error_data_redacted, _prepare_data_redacted_error, ) -from ..guardrail import OutputGuardrailResult +from ..guardrail import GuardrailFunctionOutput, OutputGuardrailResult from ..handoffs import Handoff from ..items import ( InputItem, @@ -53,6 +53,7 @@ ModelResponse, RunItem, ToolApprovalItem, + ToolCallItem, ToolCallOutputItem, TResponseInputItem, ) @@ -74,7 +75,7 @@ from ..run_config import ReasoningItemIdPolicy, RunConfig from ..run_context import AgentHookContext, RunContextWrapper, TContext from ..run_error_handlers import RunErrorHandlers -from ..run_state import RunState, _deserialize_tool_call_output_raw_item +from ..run_state import RunState from ..sandbox.runtime import SandboxRuntime from ..stream_events import ( AgentUpdatedStreamEvent, @@ -85,7 +86,7 @@ Tool, dispose_resolved_computers, ) -from ..tool_guardrails import ToolOutputGuardrailResult +from ..tool_guardrails import ToolGuardrailFunctionOutput, ToolOutputGuardrailResult from ..tracing import Span, SpanError, agent_span, get_current_trace, task_span, turn_span from ..tracing.config import include_task_and_turn_spans from ..tracing.model_tracing import get_model_tracing_impl @@ -105,8 +106,14 @@ get_unsent_tool_call_ids_for_interrupted_state, snapshot_usage, usage_delta, + validate_output_guardrails_with_server_managed_conversation, ) from .approvals import approvals_from_step +from .blocked_output import ( + OUTPUT_GUARDRAIL_BLOCKED_TOOL_OUTPUT as _OUTPUT_GUARDRAIL_BLOCKED_TOOL_OUTPUT, + blocked_function_call_payload as _blocked_function_call_payload, + blocked_function_output_payload as _blocked_function_output_payload, +) from .error_handlers import ( attach_generic_agent_error, build_run_error_data, @@ -141,7 +148,6 @@ from .oai_conversation import OpenAIServerConversationTracker from .prompt_cache_key import PromptCacheKeyResolver, model_settings_with_prompt_cache_key from .run_steps import ( - DeferredToolSpan, NextStepFinalOutput, NextStepHandoff, NextStepInterruption, @@ -175,15 +181,12 @@ from .tool_actions import ApplyPatchAction, ComputerAction, LocalShellAction, ShellAction from .tool_execution import ( coerce_shell_call, - collect_deferred_tool_spans, execute_apply_patch_calls, execute_computer_actions, execute_function_tool_calls, execute_local_shell_calls, execute_shell_calls, extract_tool_call_id, - finish_deferred_tool_spans, - get_mapping_or_attr, initialize_computer_tools, maybe_reset_tool_choice, normalize_shell_output, @@ -205,7 +208,6 @@ validate_run_hooks, ) from .turn_resolution import ( - _collect_program_parent_state, check_for_final_output_from_tools, execute_final_output, execute_handoffs, @@ -466,467 +468,635 @@ async def _run_output_guardrails_for_stream( _SIDE_EFFECT_ITEM_TYPES = frozenset({"tool_call_item", "tool_call_output_item"}) -_OUTPUT_GUARDRAIL_BLOCKED_TOOL_OUTPUT = "Output withheld by an output guardrail." -_OUTPUT_GUARDRAIL_BLOCKED_TOOL_CALL_ID = "blocked-tool-output" -_OUTPUT_GUARDRAIL_BLOCKED_COMPUTER_SCREENSHOT = ( - "data:image/png;base64," - "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mNk+A8AAQUBAScY42YAAAAASUVORK5CYII=" -) -_RESPONSE_OUTPUT_STATUSES = frozenset({"in_progress", "completed", "incomplete"}) -_HOSTED_TOOL_OUTPUT_STATUSES = frozenset({"completed", "failed"}) -_SHELL_TOOL_OUTPUT_STATUSES = _RESPONSE_OUTPUT_STATUSES | _HOSTED_TOOL_OUTPUT_STATUSES -_PROGRAM_OUTPUT_STATUSES = frozenset({"completed", "incomplete"}) -_KNOWN_TOOL_OUTPUT_TYPES = frozenset( - { - "function_call_output", - "custom_tool_call_output", - "local_shell_call_output", - "apply_patch_call_output", - "shell_call_output", - "computer_call_output", - "program_output", - } -) -async def _run_with_deferred_tool_spans( - *, - agent: Agent[TContext], - run_config: RunConfig, - run_turn: Callable[[], Awaitable[SingleStepResult]], -) -> SingleStepResult: - """Delay terminal tool span publication until output guardrails select the payload.""" - should_defer = ( - run_config.trace_include_sensitive_data - and not run_config.tracing_disabled - and agent.tool_use_behavior != "run_llm_again" - and bool(agent.output_guardrails or run_config.output_guardrails) - ) - with collect_deferred_tool_spans(should_defer) as deferred_spans: - try: - result = await run_turn() - except BaseException: - finish_deferred_tool_spans(deferred_spans) - raise +def _sanitize_blocked_output_guardrail_results( + results: Sequence[OutputGuardrailResult], + tripwire: OutputGuardrailTripwireTriggered, +) -> list[OutputGuardrailResult]: + """Build data-free guardrail results and detach the tripwire from raw output.""" + sanitized_by_id: dict[int, OutputGuardrailResult] = {} + + def sanitize(result: OutputGuardrailResult) -> OutputGuardrailResult: + existing = sanitized_by_id.get(id(result)) + if existing is not None: + return existing + sanitized = OutputGuardrailResult( + guardrail=result.guardrail, + agent_output=_OUTPUT_GUARDRAIL_BLOCKED_TOOL_OUTPUT, + agent=result.agent, + output=GuardrailFunctionOutput( + output_info=None, + tripwire_triggered=result.output.tripwire_triggered, + ), + ) + sanitized_by_id[id(result)] = sanitized + return sanitized - if isinstance(result.next_step, NextStepFinalOutput): - result.deferred_tool_spans = deferred_spans - else: - finish_deferred_tool_spans(deferred_spans) - return result - - -def _tool_output_payload(raw_item: Any) -> dict[str, Any]: - """Convert a raw tool output into a mutable mapping.""" - if isinstance(raw_item, dict): - return dict(raw_item) - model_dump = getattr(raw_item, "model_dump", None) - if callable(model_dump): - return cast(dict[str, Any], model_dump(exclude_unset=True)) - raise AgentsException(f"Unexpected raw tool output type: {type(raw_item)}") - - -def _tool_output_identity(raw_item: Any) -> tuple[str, str] | None: - """Return the raw output type and provider identity used to join replay copies.""" - payload = _tool_output_payload(raw_item) - output_type = payload.get("type") - call_id = payload.get("call_id") or payload.get("id") - if not isinstance(output_type, str) or not isinstance(call_id, str): - return None - return output_type, call_id + sanitized_results = [sanitize(result) for result in results] + object.__setattr__(tripwire, "guardrail_result", sanitize(tripwire.guardrail_result)) + _mark_error_data_redacted(tripwire) + _detach_data_redacted_error_traceback(tripwire) + return sanitized_results -def _required_tool_output_string( - raw_payload: Mapping[str, Any], - field: str, - output_type: str, -) -> str: - value = raw_payload.get(field) - if not isinstance(value, str) or not value: - raise AgentsException(f"Cannot sanitize {output_type} without a non-empty string {field}.") - return value +@_dc.dataclass(frozen=True) +class _CurrentResponseBoundary: + """A current-response suffix proven only by lifecycle position or object identity.""" + items: tuple[RunItem, ...] + processed_items: tuple[RunItem, ...] + generated_start: int | None + session_start: int | None + proven: bool -def _copy_safe_status( - sanitized: dict[str, Any], - raw_payload: Mapping[str, Any], - output_type: str, - allowed: frozenset[str], - *, - required: bool = False, -) -> None: - status = raw_payload.get("status") - if status is None and not required: - return - if not isinstance(status, str) or status not in allowed: - raise AgentsException(f"Cannot sanitize {output_type} with an invalid status.") - sanitized["status"] = status +@_dc.dataclass(frozen=True) +class _BlockedOutputSnapshot: + """Prepared data-free replacements for one complete current response.""" -def _copy_safe_caller( - sanitized: dict[str, Any], - raw_payload: Mapping[str, Any], - output_type: str, -) -> None: - caller = raw_payload.get("caller") - if caller is None: - return - if not isinstance(caller, Mapping): - raise AgentsException(f"Cannot sanitize {output_type} with an invalid caller.") - if caller.get("type") == "direct": - sanitized["caller"] = {"type": "direct"} - return - caller_id = caller.get("caller_id") - if caller.get("type") == "program" and isinstance(caller_id, str) and caller_id: - sanitized["caller"] = {"type": "program", "caller_id": caller_id} - return - raise AgentsException(f"Cannot sanitize {output_type} with an invalid caller.") + items: tuple[RunItem, ...] + processed_items: tuple[RunItem, ...] + model_response: ModelResponse | None + + +@_dc.dataclass(frozen=True) +class _BlockedOutputOwnerPlan: + """Prebuilt trusted-owner assignments for application or emergency cleanup.""" + + assignments: tuple[tuple[Any, str, Any], ...] + + +@_dc.dataclass(frozen=True) +class _BlockedOutputOwnerStarts: + """Owner-specific current-response starts captured at trusted lifecycle boundaries.""" + + run_state_generated_items: int | None = None + run_state_session_items: int | None = None + run_state_model_responses: int | None = None + run_state_tool_output_guardrail_results: int | None = None + streamed_new_items: int | None = None + streamed_model_input_items: int | None = None + streamed_raw_responses: int | None = None + streamed_tool_output_guardrail_results: int | None = None + + +@_dc.dataclass(frozen=True) +class _BlockedOutputOwnerPrefixes: + """Accepted owner prefixes allocated before any blocked-output replacement begins.""" + run_state_generated_items: list[RunItem] + run_state_session_items: list[RunItem] + run_state_model_responses: list[ModelResponse] + run_state_tool_output_guardrail_results: list[ToolOutputGuardrailResult] + streamed_new_items: list[RunItem] + streamed_model_input_items: list[RunItem] + streamed_raw_responses: list[ModelResponse] + streamed_tool_output_guardrail_results: list[ToolOutputGuardrailResult] -def _copy_safe_safety_checks( - sanitized: dict[str, Any], - raw_payload: Mapping[str, Any], + +_OwnerItemT = TypeVar("_OwnerItemT") + + +def _has_output_guardrails(agent: Agent[Any], run_config: RunConfig) -> bool: + return bool(agent.output_guardrails or run_config.output_guardrails) + + +def _should_defer_interrupted_session_items( + agent: Agent[Any], + run_config: RunConfig, +) -> bool: + """Keep pre-verdict approval state in RunState instead of durable Session history.""" + return _has_output_guardrails(agent, run_config) + + +def _validate_resumed_session_output_guardrail_safety( + *, + agent: Agent[Any], + run_config: RunConfig, + session: Session | None, + run_state: RunState[Any] | None, ) -> None: - checks = raw_payload.get("acknowledged_safety_checks") - if checks is None: + """Reject approval resumes whose current-response boundary is not structurally provable.""" + del session + if run_state is None or not _has_output_guardrails(agent, run_config): return - if not isinstance(checks, Sequence) or isinstance(checks, str | bytes): - raise AgentsException( - "Cannot sanitize computer_call_output with invalid acknowledged safety checks." - ) - safe_checks: list[dict[str, str]] = [] - for check in checks: - if not isinstance(check, Mapping): - raise AgentsException( - "Cannot sanitize computer_call_output with invalid acknowledged safety checks." - ) - check_id = check.get("id") - if not isinstance(check_id, str) or not check_id: - raise AgentsException( - "Cannot sanitize computer_call_output with invalid acknowledged safety checks." - ) - safe_checks.append({"id": check_id}) - sanitized["acknowledged_safety_checks"] = safe_checks - - -def _blocked_tool_output_payload(raw_item: Any) -> dict[str, Any]: - """Build a data-free replay payload from validated protocol fields.""" - raw_payload = _tool_output_payload(raw_item) - output_type = raw_payload.get("type") - if not isinstance(output_type, str) or not output_type: - raise AgentsException("Cannot sanitize a tool output without a non-empty string type.") - - sanitized_raw_item: dict[str, Any] = {"type": output_type} - if output_type not in _KNOWN_TOOL_OUTPUT_TYPES: - call_id = raw_payload.get("call_id") - item_id = raw_payload.get("id") - if isinstance(call_id, str) and call_id: - sanitized_raw_item["call_id"] = call_id - elif isinstance(item_id, str) and item_id: - sanitized_raw_item["id"] = item_id - else: - raise AgentsException( - f"Cannot sanitize {output_type} without a non-empty string call_id or id." - ) - sanitized_raw_item["output"] = _OUTPUT_GUARDRAIL_BLOCKED_TOOL_OUTPUT - return sanitized_raw_item - if output_type == "program_output": - sanitized_raw_item["id"] = _required_tool_output_string(raw_payload, "id", output_type) - sanitized_raw_item["call_id"] = _required_tool_output_string( - raw_payload, "call_id", output_type - ) - _copy_safe_status( - sanitized_raw_item, - raw_payload, - output_type, - _PROGRAM_OUTPUT_STATUSES, - required=True, - ) - sanitized_raw_item["result"] = _OUTPUT_GUARDRAIL_BLOCKED_TOOL_OUTPUT - else: - sanitized_raw_item["call_id"] = _required_tool_output_string( - raw_payload, "call_id", output_type + if not isinstance(run_state._current_step, NextStepInterruption): + return + if run_state._current_turn_persisted_item_count > 0: + raise UserError( + "Cannot resume an approval checkpoint with output guardrails after current-turn " + "items were persisted. Start a new run from safe input." ) + boundary = _current_response_boundary( + (), + run_state._last_processed_response, + run_state, + ) + if boundary.proven: + return + raise UserError( + "Cannot resume a serialized approval checkpoint with output guardrails because the " + "current response boundary cannot be proven. Start a new run from safe input." + ) - item_id = raw_payload.get("id") - if item_id is not None: - if not isinstance(item_id, str) or not item_id: - raise AgentsException(f"Cannot sanitize {output_type} with an invalid id.") - sanitized_raw_item["id"] = item_id - - if output_type in { - "function_call_output", - "custom_tool_call_output", - "shell_call_output", - "apply_patch_call_output", - }: - _copy_safe_caller(sanitized_raw_item, raw_payload, output_type) - - if output_type in {"function_call_output", "computer_call_output"}: - _copy_safe_status( - sanitized_raw_item, - raw_payload, - output_type, - _RESPONSE_OUTPUT_STATUSES, - ) - elif output_type == "shell_call_output": - _copy_safe_status( - sanitized_raw_item, - raw_payload, - output_type, - _SHELL_TOOL_OUTPUT_STATUSES, - ) - elif output_type == "apply_patch_call_output": - _copy_safe_status( - sanitized_raw_item, - raw_payload, - output_type, - _HOSTED_TOOL_OUTPUT_STATUSES, - ) - if output_type == "computer_call_output": - _copy_safe_safety_checks(sanitized_raw_item, raw_payload) - - if output_type == "computer_call_output": - sanitized_raw_item["output"] = { - "type": "computer_screenshot", - "image_url": _OUTPUT_GUARDRAIL_BLOCKED_COMPUTER_SCREENSHOT, - } - elif output_type == "shell_call_output": - sanitized_raw_item["output"] = [ - { - "stdout": "", - "stderr": _OUTPUT_GUARDRAIL_BLOCKED_TOOL_OUTPUT, - "outcome": {"type": "exit", "exit_code": 1}, - } - ] - elif output_type != "program_output": - sanitized_raw_item["output"] = _OUTPUT_GUARDRAIL_BLOCKED_TOOL_OUTPUT +def _identity_sequence_start( + container: Sequence[RunItem], + sequence: Sequence[RunItem], +) -> int | None: + if not sequence or len(sequence) > len(container): + return None + for start in range(len(container) - len(sequence) + 1): + if all(container[start + offset] is item for offset, item in enumerate(sequence)): + return start + return None - try: - restored = _deserialize_tool_call_output_raw_item(sanitized_raw_item) - except Exception: - raise AgentsException(f"Sanitized {output_type} is not valid for durable replay.") from None - if restored is None: - raise AgentsException(f"Sanitized {output_type} is not valid for durable replay.") - return sanitized_raw_item +def _current_response_boundary( + new_items: Sequence[RunItem], + processed_response: ProcessedResponse | None, + run_state: RunState[Any] | None, +) -> _CurrentResponseBoundary: + """Collect one response using only SDK lifecycle position and exact object identity.""" + processed_items = tuple(processed_response.new_items) if processed_response is not None else () + supplied_items = tuple(new_items) + supplied_start = _identity_sequence_start(supplied_items, processed_items) + response_items = ( + supplied_items[supplied_start:] if supplied_start is not None else supplied_items + ) + generated_start = None + session_start = None + proven = run_state is None or not processed_items or supplied_start is not None + suffixes: list[RunItem] = [] + if run_state is not None: + anchor_items = processed_items or response_items + if anchor_items: + generated_start = _identity_sequence_start(run_state._generated_items, anchor_items) + session_start = _identity_sequence_start(run_state._session_items, anchor_items) + if generated_start is not None: + suffixes.extend(run_state._generated_items[generated_start:]) + proven = True + if session_start is not None: + suffixes.extend(run_state._session_items[session_start:]) + proven = True + if ( + not supplied_items + and generated_start is None + and session_start is None + and run_state._current_turn == 1 + ): + generated_start = 0 + session_start = 0 + suffixes.extend(run_state._generated_items) + suffixes.extend(run_state._session_items) + proven = True -def _sanitize_blocked_output_guardrail_results( - results: Sequence[OutputGuardrailResult], - tripwire: OutputGuardrailTripwireTriggered, -) -> None: - """Remove blocked output aliases from completed guardrail results and the exception.""" + current_items: list[RunItem] = [] seen: set[int] = set() - for result in (*results, tripwire.guardrail_result): - if id(result) in seen: + for item in (*processed_items, *suffixes, *response_items): + if id(item) in seen: continue - seen.add(id(result)) - result.agent_output = _OUTPUT_GUARDRAIL_BLOCKED_TOOL_OUTPUT - result.output.output_info = None - _mark_error_data_redacted(tripwire) - _detach_data_redacted_error_traceback(tripwire) + seen.add(id(item)) + current_items.append(item) + return _CurrentResponseBoundary( + items=tuple(current_items), + processed_items=processed_items, + generated_start=generated_start, + session_start=session_start, + proven=proven, + ) -def _sanitize_blocked_tool_output_guardrail_results( - results: Sequence[ToolOutputGuardrailResult], -) -> None: - """Remove blocked tool-output aliases from the current turn's guardrail results.""" - for result in results: - result.output.output_info = None - if result.output.behavior["type"] == "reject_content": - result.output.behavior = { - "type": "reject_content", - "message": _OUTPUT_GUARDRAIL_BLOCKED_TOOL_OUTPUT, - } - - -def _redact_blocked_output_state_step(run_state: RunState[Any] | None) -> None: - """Remove a rejected final output from the live resumable step before propagation.""" - if run_state is not None and isinstance(run_state._current_step, NextStepFinalOutput): - run_state._current_step.output = _OUTPUT_GUARDRAIL_BLOCKED_TOOL_OUTPUT - - -def _finish_blocked_output_tool_spans(spans: list[DeferredToolSpan]) -> None: - """Publish terminal tool spans with the same placeholder used by replay state.""" - finish_deferred_tool_spans( - spans, - output_override=_OUTPUT_GUARDRAIL_BLOCKED_TOOL_OUTPUT, - ) +def _current_response_items_for_persistence( + items: Sequence[RunItem], + processed_response: ProcessedResponse | None, + run_state: RunState[Any] | None = None, +) -> list[RunItem]: + """Return the complete current response or fail before using an ambiguous boundary.""" + boundary = _current_response_boundary(items, processed_response, run_state) + if not boundary.proven: + raise UserError( + "Cannot persist an ambiguous resumed response with output guardrails. " + "Start a new run from safe input." + ) + return list(boundary.items) -def _sanitize_retained_tool_outputs( - items: list[RunItem], - model_response: ModelResponse | None, -) -> None: - """Sanitize retained run items and matching archived raw-response outputs.""" - retained_identities: set[tuple[str, str]] = set() - retained_raw_item_ids: set[int] = set() - for item in items: - if not isinstance(item, ToolCallOutputItem): - continue +def _final_turn_items_for_persistence( + items: Sequence[RunItem], + processed_response: ProcessedResponse | None, + run_state: RunState[Any] | None, + agent: Agent[Any], + run_config: RunConfig, +) -> list[RunItem]: + """Use released resumed suffix persistence unless output guardrails defer the response.""" + if not _has_output_guardrails(agent, run_config): + return list(items) + return _current_response_items_for_persistence(items, processed_response, run_state) - original_raw_item = item.raw_item - identity = _tool_output_identity(original_raw_item) - if identity is not None: - retained_identities.add(identity) - retained_raw_item_ids.add(id(original_raw_item)) - item.raw_item = cast(Any, _blocked_tool_output_payload(original_raw_item)) - item.output = _OUTPUT_GUARDRAIL_BLOCKED_TOOL_OUTPUT - # Custom data is SDK-only and may duplicate the blocked tool result. - item.custom_data = None +def _is_terminal_tool_output_response( + items: Sequence[RunItem], + processed_response: ProcessedResponse | None, + run_state: RunState[Any] | None = None, +) -> bool: + """Return whether the structurally owned current response produced a tool final output.""" + boundary = _current_response_boundary(items, processed_response, run_state) + return boundary.proven and any(isinstance(item, ToolCallOutputItem) for item in boundary.items) - if model_response is None: - return - sanitized_response_output = [] - for raw_output in model_response.output: - identity = _tool_output_identity(raw_output) - if id(raw_output) in retained_raw_item_ids or ( - identity is not None and identity in retained_identities - ): - sanitized_response_output.append(cast(Any, _blocked_tool_output_payload(raw_output))) +def _prepare_blocked_output_snapshot( + boundary: _CurrentResponseBoundary, + model_response: ModelResponse | None, +) -> _BlockedOutputSnapshot: + """Build an allowlist-only function call/output snapshot before changing live state.""" + current_items = list(boundary.items) + if any(item.type == "reasoning_item" for item in current_items): + raise AgentsException("Cannot sanitize a response containing reasoning items.") + retained_indexes = { + index for index, item in enumerate(current_items) if item.type in _SIDE_EFFECT_ITEM_TYPES + } + replacements: dict[int, RunItem] = {} + calls_by_id: dict[str, int] = {} + outputs_by_id: dict[str, int] = {} + for index in sorted(retained_indexes): + item = current_items[index] + if isinstance(item, ToolCallItem): + payload = _blocked_function_call_payload(item.raw_item) + call_id = cast(str, payload["call_id"]) + if call_id in calls_by_id: + raise AgentsException("Cannot sanitize duplicate function calls.") + calls_by_id[call_id] = index + replacements[index] = ToolCallItem( + agent=item.agent, + raw_item=cast(Any, payload), + description=item.description, + title=item.title, + tool_origin=item.tool_origin, + _resolved_tool_name=item._resolved_tool_name, + ) + elif isinstance(item, ToolCallOutputItem): + payload = _blocked_function_output_payload(item.raw_item) + call_id = cast(str, payload["call_id"]) + if call_id in outputs_by_id: + raise AgentsException("Cannot sanitize duplicate function outputs.") + outputs_by_id[call_id] = index + replacements[index] = ToolCallOutputItem( + agent=item.agent, + raw_item=cast(Any, payload), + output=_OUTPUT_GUARDRAIL_BLOCKED_TOOL_OUTPUT, + tool_origin=item.tool_origin, + custom_data=None, + ) else: - sanitized_response_output.append(raw_output) - model_response.output = sanitized_response_output + raise AgentsException("Cannot sanitize an unsupported side-effect item.") + + if not outputs_by_id or set(outputs_by_id) - set(calls_by_id): + raise AgentsException("Cannot sanitize an incomplete function call/output batch.") + retained_indexes = { + index + for call_id in outputs_by_id + for index in (calls_by_id[call_id], outputs_by_id[call_id]) + } + retained_items = tuple(replacements[index] for index in sorted(retained_indexes)) + processed_indexes = {id(item): index for index, item in enumerate(current_items)} + retained_processed_items = tuple( + replacements.get(processed_indexes[id(item)], item) + for item in boundary.processed_items + if processed_indexes.get(id(item)) in retained_indexes + ) + sanitized_response = None + if model_response is not None: + sanitized_response = ModelResponse( + output=cast(Any, [item.raw_item for item in retained_processed_items]), + usage=model_response.usage, + response_id=model_response.response_id, + request_id=model_response.request_id, + raw_usage=model_response.raw_usage, + ) + return _BlockedOutputSnapshot( + items=retained_items, + processed_items=retained_processed_items, + model_response=sanitized_response, + ) -def _validate_retained_program_relationships( - items: list[RunItem], - preceding_items: Sequence[RunItem], -) -> None: - """Reject retained program relationships that cannot be replayed in order.""" - preceding_items = list(preceding_items) - for item in items: - raw_item = getattr(item, "raw_item", item) - output_type = get_mapping_or_attr(raw_item, "type") - program_call_ids, completed_program_call_ids = _collect_program_parent_state( - preceding_items +def _blocked_output_owner_prefix(items: list[_OwnerItemT], start: int | None) -> list[_OwnerItemT]: + """Copy a structurally captured prefix without consulting item values or identities.""" + if start is None or start < 0 or start > len(items): + return [] + return list.__getitem__(items, slice(0, start)) + + +def _prepare_blocked_output_owner_prefixes( + run_state: RunState[Any] | None, + streamed_result: RunResultStreaming | None, + owner_starts: _BlockedOutputOwnerStarts, +) -> _BlockedOutputOwnerPrefixes: + """Allocate every accepted owner prefix before snapshot application begins.""" + return _BlockedOutputOwnerPrefixes( + run_state_generated_items=( + _blocked_output_owner_prefix( + run_state._generated_items, + owner_starts.run_state_generated_items, + ) + if run_state is not None + else [] + ), + run_state_session_items=( + _blocked_output_owner_prefix( + run_state._session_items, + owner_starts.run_state_session_items, + ) + if run_state is not None + else [] + ), + run_state_model_responses=( + _blocked_output_owner_prefix( + run_state._model_responses, + owner_starts.run_state_model_responses, + ) + if run_state is not None + else [] + ), + run_state_tool_output_guardrail_results=( + _blocked_output_owner_prefix( + run_state._tool_output_guardrail_results, + owner_starts.run_state_tool_output_guardrail_results, + ) + if run_state is not None + else [] + ), + streamed_new_items=( + _blocked_output_owner_prefix( + streamed_result.new_items, + owner_starts.streamed_new_items, + ) + if streamed_result is not None + else [] + ), + streamed_model_input_items=( + _blocked_output_owner_prefix( + streamed_result._model_input_items, + owner_starts.streamed_model_input_items, + ) + if streamed_result is not None + else [] + ), + streamed_raw_responses=( + _blocked_output_owner_prefix( + streamed_result.raw_responses, + owner_starts.streamed_raw_responses, + ) + if streamed_result is not None + else [] + ), + streamed_tool_output_guardrail_results=( + _blocked_output_owner_prefix( + streamed_result.tool_output_guardrail_results, + owner_starts.streamed_tool_output_guardrail_results, + ) + if streamed_result is not None + else [] + ), + ) + + +def _prepare_blocked_output_cleanup_plan( + run_state: RunState[Any] | None, + streamed_result: RunResultStreaming | None, + prefixes: _BlockedOutputOwnerPrefixes, +) -> _BlockedOutputOwnerPlan: + """Prepare accepted-prefix cleanup containers before snapshot application begins.""" + assignments: list[tuple[Any, str, Any]] = [] + if run_state is not None: + assignments.extend( + [ + (run_state, "_generated_items", prefixes.run_state_generated_items), + (run_state, "_session_items", prefixes.run_state_session_items), + (run_state, "_model_responses", prefixes.run_state_model_responses), + (run_state, "_last_processed_response", None), + (run_state, "_current_step", None), + (run_state, "_generated_items_last_processed_marker", None), + ( + run_state, + "_tool_output_guardrail_results", + prefixes.run_state_tool_output_guardrail_results, + ), + ] + ) + if streamed_result is not None: + assignments.extend( + [ + (streamed_result, "new_items", prefixes.streamed_new_items), + (streamed_result, "raw_responses", prefixes.streamed_raw_responses), + ( + streamed_result, + "_model_input_items", + prefixes.streamed_model_input_items, + ), + (streamed_result, "_last_processed_response", None), + ( + streamed_result, + "tool_output_guardrail_results", + prefixes.streamed_tool_output_guardrail_results, + ), + ] ) + return _BlockedOutputOwnerPlan(assignments=tuple(assignments)) - caller = get_mapping_or_attr(raw_item, "caller") - if get_mapping_or_attr(caller, "type") == "program": - caller_id = get_mapping_or_attr(caller, "caller_id") - if ( - not isinstance(caller_id, str) - or caller_id not in program_call_ids - or caller_id in completed_program_call_ids - ): - raise AgentsException( - f"Cannot sanitize {output_type} with an invalid program caller." - ) - if output_type == "program_output": - call_id = get_mapping_or_attr(raw_item, "call_id") - if ( - not isinstance(call_id, str) - or call_id not in program_call_ids - or call_id in completed_program_call_ids - ): - raise AgentsException( - "Cannot sanitize program_output without an active retained program parent." - ) +def _sever_blocked_output_replay_graph(cleanup_plan: _BlockedOutputOwnerPlan) -> None: + """Best-effort leaf cleanup using only containers allocated before application.""" + for owner, field, value in cleanup_plan.assignments: + try: + object.__setattr__(owner, field, value) + except BaseException: + continue - preceding_items.append(item) + +def _data_free_tool_output_guardrail_results( + results: Sequence[ToolOutputGuardrailResult], +) -> tuple[ToolOutputGuardrailResult, ...]: + """Rebuild current-turn tool guardrail results without retaining caller output data.""" + replacements: list[ToolOutputGuardrailResult] = [] + try: + for result in results: + if not isinstance(result, ToolOutputGuardrailResult): + return () + replacements.append( + ToolOutputGuardrailResult( + guardrail=object.__getattribute__(result, "guardrail"), + output=ToolGuardrailFunctionOutput( + output_info=_OUTPUT_GUARDRAIL_BLOCKED_TOOL_OUTPUT, + behavior={"type": "allow"}, + ), + ) + ) + except Exception: + return () + return tuple(replacements) -def _scrub_failed_blocked_output_aliases( - items: list[RunItem], +def _prepare_blocked_output_owner_plan( + boundary: _CurrentResponseBoundary, + snapshot: _BlockedOutputSnapshot | None, model_response: ModelResponse | None, -) -> None: - """Remove output data after replay-payload validation fails.""" - for item in items: - if not isinstance(item, ToolCallOutputItem): - continue - item.raw_item = cast( - Any, - { - "type": "function_call_output", - "call_id": _OUTPUT_GUARDRAIL_BLOCKED_TOOL_CALL_ID, - "output": _OUTPUT_GUARDRAIL_BLOCKED_TOOL_OUTPUT, - }, + run_state: RunState[Any] | None, + streamed_result: RunResultStreaming | None, + prefixes: _BlockedOutputOwnerPrefixes, + cleanup_plan: _BlockedOutputOwnerPlan, +) -> _BlockedOutputOwnerPlan: + """Build every owner replacement before applying any of them.""" + safe_items = list(snapshot.items) if snapshot is not None else [] + safe_response = snapshot.model_response if snapshot is not None else None + assignments: list[tuple[Any, str, Any]] = [] + safe_tool_output_guardrail_results: tuple[ToolOutputGuardrailResult, ...] = () + if streamed_result is not None: + public_results = streamed_result.tool_output_guardrail_results + current_results = list.__getitem__( + public_results, + slice(len(prefixes.streamed_tool_output_guardrail_results), None), + ) + safe_tool_output_guardrail_results = _data_free_tool_output_guardrail_results( + current_results ) - item.output = _OUTPUT_GUARDRAIL_BLOCKED_TOOL_OUTPUT - item.custom_data = None + public_safe_results = [ + *prefixes.streamed_tool_output_guardrail_results, + *safe_tool_output_guardrail_results, + ] + else: + public_safe_results = [] - if model_response is not None: - model_response.output = [] + if run_state is not None: + if boundary.proven: + responses = [ + *prefixes.run_state_model_responses, + *( + [safe_response] + if model_response is not None and safe_response is not None + else [] + ), + ] + if streamed_result is not None: + run_state_safe_results = [ + *prefixes.run_state_tool_output_guardrail_results, + *safe_tool_output_guardrail_results, + ] + else: + run_state_safe_results = list(prefixes.run_state_tool_output_guardrail_results) + assignments.extend( + [ + ( + run_state, + "_generated_items", + [*prefixes.run_state_generated_items, *safe_items], + ), + ( + run_state, + "_session_items", + [*prefixes.run_state_session_items, *safe_items], + ), + (run_state, "_model_responses", responses), + (run_state, "_last_processed_response", None), + (run_state, "_current_step", None), + (run_state, "_generated_items_last_processed_marker", None), + (run_state, "_tool_output_guardrail_results", run_state_safe_results), + ] + ) + else: + return cleanup_plan + + if streamed_result is not None: + responses = [ + *prefixes.streamed_raw_responses, + *([safe_response] if model_response is not None and safe_response is not None else []), + ] + assignments.extend( + [ + ( + streamed_result, + "new_items", + [*prefixes.streamed_new_items, *safe_items], + ), + ( + streamed_result, + "_model_input_items", + [*prefixes.streamed_model_input_items, *safe_items], + ), + (streamed_result, "raw_responses", responses), + (streamed_result, "_last_processed_response", None), + ( + streamed_result, + "tool_output_guardrail_results", + public_safe_results, + ), + ] + ) + return _BlockedOutputOwnerPlan(assignments=tuple(assignments)) -def _reasoning_indexes_tied_to_retained_items( - items: list[RunItem], - retained_indexes: set[int], -) -> set[int]: - """Indexes of the reasoning items whose tied item is being retained. - - Applies the same association rule as - ``agents.run_internal.items._drop_reasoning_items_preceding_dropped_calls``: a reasoning item - is tied to the next *non-reasoning* model-emitted item. Keeping a group whose following item is - dropped would leave a dangling reasoning item, which the Responses API rejects on the next - request (``reasoning was provided without its required following item``); dropping a group - whose following item is retained would strip the context that call needs to be replayed. - - A trailing reasoning group - one with no following non-reasoning item at all - is not tied to - anything retained, so it is dropped. Note this is stricter than the reference, which keeps such - a group because the item it belongs to may still arrive later in a longer history; here the - turn is complete, so there is nothing left to tie it to. - """ - tied: set[int] = set() - for index in range(len(items) - 1, -1, -1): - if items[index].type != "reasoning_item": - continue - for next_index in range(index + 1, len(items)): - if items[next_index].type == "reasoning_item": - continue - if next_index in retained_indexes: - tied.add(index) - break - return tied +def _apply_blocked_output_owner_plan(plan: _BlockedOutputOwnerPlan) -> None: + """Apply only values that were fully constructed before the first owner swap.""" + for owner, field, value in plan.assignments: + object.__setattr__(owner, field, value) def _retained_items_for_blocked_output( items: list[RunItem], model_response: ModelResponse | None = None, - preceding_items: Sequence[RunItem] = (), ) -> list[RunItem]: - """Pick out the items of a final turn to keep when its output is not deliverable. + """Return trusted retained items without consulting earlier provider identities.""" + return _retained_items_for_blocked_response( + items, + model_response, + ) - A tool that already ran has to stay in the session, together with the context needed to replay - its call. Everything else - the assistant message the guardrail rejected above all - is dropped, - including the reasoning that belongs to the rejected message rather than to a retained call. - ``_SIDE_EFFECT_ITEM_TYPES`` is enumerated rather than derived, so an item type added later is - *discarded* here by default and has to be classified deliberately. A record of a side effect - that goes unclassified is a bug, so the safer default is the one that surfaces as a missing item - rather than as a rejected message quietly reaching the session. - """ - redacted_error: AgentsException | None = None +def _retained_items_for_blocked_response( + items: list[RunItem], + model_response: ModelResponse | None, + run_state: RunState[Any] | None = None, + processed_response: ProcessedResponse | None = None, + streamed_result: RunResultStreaming | None = None, + owner_starts: _BlockedOutputOwnerStarts | None = None, +) -> list[RunItem]: + """Return a complete data-free response or discard the entire unsupported suffix.""" + boundary = _current_response_boundary(items, processed_response, run_state) + prefixes = _prepare_blocked_output_owner_prefixes( + run_state, + streamed_result, + owner_starts if owner_starts is not None else _BlockedOutputOwnerStarts(), + ) + cleanup_plan = _prepare_blocked_output_cleanup_plan(run_state, streamed_result, prefixes) + snapshot: _BlockedOutputSnapshot | None = None + try: + if boundary.proven: + snapshot = _prepare_blocked_output_snapshot(boundary, model_response) + except Exception: + snapshot = None + except BaseException: + _sever_blocked_output_replay_graph(cleanup_plan) + raise try: - retained_indexes = { - index for index, item in enumerate(items) if item.type in _SIDE_EFFECT_ITEM_TYPES - } - if not retained_indexes: - return [] - # Reasoning items are not side effects themselves, but a reasoning model requires the - # reasoning item tied to a function call to accompany it in the next request. - retained_indexes |= _reasoning_indexes_tied_to_retained_items(items, retained_indexes) - retained_items = [items[index] for index in sorted(retained_indexes)] - _validate_retained_program_relationships(retained_items, preceding_items) - _sanitize_retained_tool_outputs(retained_items, model_response) - # Indexed rather than filtered by type so the retained items keep the model's own order. - return [item for index, item in enumerate(items) if index in retained_indexes] - except AgentsException as error: - _scrub_failed_blocked_output_aliases(items, model_response) - _mark_error_data_redacted(error) - _detach_data_redacted_error_traceback(error) - redacted_error = AgentsException("Cannot sanitize a blocked tool output for replay.") - _mark_error_data_redacted(redacted_error) - - items = [] - model_response = None - assert redacted_error is not None - raise redacted_error from None + owner_plan = _prepare_blocked_output_owner_plan( + boundary, + snapshot, + model_response, + run_state, + streamed_result, + prefixes, + cleanup_plan, + ) + _apply_blocked_output_owner_plan(owner_plan) + except Exception as error: + _sever_blocked_output_replay_graph(cleanup_plan) + raise _prepare_data_redacted_error(error) from None + except BaseException: + _sever_blocked_output_replay_graph(cleanup_plan) + raise + return list(snapshot.items) if snapshot is not None else [] async def _finalize_streamed_final_output( @@ -939,14 +1109,12 @@ async def _finalize_streamed_final_output( save_items: Callable[[list[RunItem], str | None, bool | None], Awaitable[None]], items: list[RunItem], model_response: ModelResponse | None, - deferred_tool_spans: list[DeferredToolSpan], - preceding_items: Sequence[RunItem], - tool_output_guardrail_results: Sequence[ToolOutputGuardrailResult], + processed_response: ProcessedResponse | None, + owner_starts: _BlockedOutputOwnerStarts, response_id: str | None, store_setting: bool | None, on_persisted_after_guardrails: Callable[[bool], None] | None = None, ) -> None: - redacted_persistence_error: BaseException | None = None output_guardrail_result_start = len(streamed_result.output_guardrail_results) try: output_guardrail_results = await _run_output_guardrails_for_stream( @@ -957,24 +1125,31 @@ async def _finalize_streamed_final_output( streamed_result=streamed_result, ) except OutputGuardrailTripwireTriggered as exc: + if not _is_terminal_tool_output_response( + items, + processed_response, + streamed_result._state, + ): + raise # The blocked output itself is not persisted, but a tool that already ran is: the next run # has to see that side effect rather than re-issue it. This turn reaches here with tool # items when `tool_use_behavior="stop_on_first_tool"` (or `stop_at_tool_names`, or a custom # callable) turned a tool result straight into the final output. - has_tool_output = any(isinstance(item, ToolCallOutputItem) for item in items) - if has_tool_output: - _sanitize_blocked_output_guardrail_results( - streamed_result.output_guardrail_results[output_guardrail_result_start:], - exc, - ) - _sanitize_blocked_tool_output_guardrail_results(tool_output_guardrail_results) - _finish_blocked_output_tool_spans(deferred_tool_spans) - if has_tool_output: - _redact_blocked_output_state_step(streamed_result._state) - retained_items = _retained_items_for_blocked_output( + sanitized_results = _sanitize_blocked_output_guardrail_results( + streamed_result.output_guardrail_results[output_guardrail_result_start:], + exc, + ) + streamed_result.output_guardrail_results = [ + *streamed_result.output_guardrail_results[:output_guardrail_result_start], + *sanitized_results, + ] + retained_items = _retained_items_for_blocked_response( items, model_response, - preceding_items, + streamed_result._state, + processed_response, + streamed_result, + owner_starts, ) if retained_items: try: @@ -988,75 +1163,28 @@ async def _finalize_streamed_final_output( streamed_result.is_complete = True streamed_result._event_queue.put_nowait(QueueCompleteSentinel()) return - if has_tool_output and streamed_result._state is not None: - streamed_result._state._current_step = None raise - except asyncio.CancelledError: - finish_deferred_tool_spans(deferred_tool_spans) + except (Exception, asyncio.CancelledError): + # Without a verdict, the SDK does not persist any part of the terminal response. raise - except Exception as guardrail_error: - finish_deferred_tool_spans(deferred_tool_spans) - # Only a tripwire means the output was judged undeliverable. A guardrail error leaves the - # verdict unknown, so the completed final turn is persisted whole and remains replayable. - # `asyncio.CancelledError` is deliberately not caught here: `cancel()` in its default - # immediate mode has to stay prompt, and awaiting a session write would block - # `stream_events()` on an arbitrary backend. `after_turn` is the mode that finishes the - # turn and saves. - guardrail_error_is_redacted = _is_error_data_redacted(guardrail_error) - if guardrail_error_is_redacted: - _detach_data_redacted_error_traceback(guardrail_error) - try: - await save_items(items, response_id, store_setting) - except BaseException as persistence_error: - if guardrail_error_is_redacted: - safe_persistence_error = _safe_redacted_persistence_error(persistence_error) - if ( - isinstance(safe_persistence_error, asyncio.CancelledError) - and streamed_result._cancel_mode != "immediate" - ): - # A cancelled session write is distinct from the caller requesting - # immediate cancellation. Retain a safe cancellation for `stream_events()` - # without completing the run-loop task with the payload-bearing backend - # exception. - streamed_result._stored_exception = safe_persistence_error - streamed_result.is_complete = True - streamed_result._event_queue.put_nowait(QueueCompleteSentinel()) - return - if isinstance(safe_persistence_error, asyncio.CancelledError): - # Public immediate cancellation already owns stream completion and must - # not surface a recovery failure. - return - redacted_persistence_error = safe_persistence_error - if ( - isinstance(persistence_error, asyncio.CancelledError) - and streamed_result._cancel_mode != "immediate" - ): - # A cancelled session write is distinct from the caller requesting immediate - # cancellation. The run-loop task itself becomes cancelled, so retain the - # backend cancellation for `stream_events()` to surface. - streamed_result._stored_exception = persistence_error - if redacted_persistence_error is None: - raise - else: - if on_persisted_after_guardrails is not None: - on_persisted_after_guardrails(False) - if redacted_persistence_error is None: - raise - - if redacted_persistence_error is not None: - raise redacted_persistence_error from None - finish_deferred_tool_spans(deferred_tool_spans) streamed_result.output_guardrail_results.extend(output_guardrail_results) + final_turn_items = _final_turn_items_for_persistence( + items, + processed_response, + streamed_result._state, + agent, + run_config, + ) # Saved as one ordered batch so the session mirrors the model response. Doing it in two # halves would both reorder the turn and, because the first save advances the turn's # persisted-item count, make the second one a no-op. if on_persisted_after_guardrails is None: - await save_items(items, response_id, store_setting) + await save_items(final_turn_items, response_id, store_setting) else: try: - await save_items(items, response_id, store_setting) + await save_items(final_turn_items, response_id, store_setting) except asyncio.CancelledError as persistence_error: if streamed_result._cancel_mode == "immediate": raise @@ -1173,8 +1301,6 @@ async def finalize_max_turns_handler_output( output: Any, context_wrapper: RunContextWrapper[TContext], output_guardrail_results: list[OutputGuardrailResult], - save_items_after_guardrails: Callable[[list[RunItem]], Awaitable[None]], - include_in_history: bool, ) -> tuple[Any, RunItem]: """Validate and finalize one synthesized max-turn handler output.""" validated_output = validate_handler_final_output(agent, output) @@ -1183,7 +1309,6 @@ async def finalize_max_turns_handler_output( await run_final_output_hooks(agent, hooks, context_wrapper, validated_output) - redacted_persistence_error: BaseException | None = None try: await run_output_guardrails( agent.output_guardrails + (run_config.output_guardrails or []), @@ -1192,23 +1317,10 @@ async def finalize_max_turns_handler_output( context_wrapper, output_guardrail_results, ) - except OutputGuardrailTripwireTriggered: - raise except Exception as guardrail_error: - guardrail_error_is_redacted = _is_error_data_redacted(guardrail_error) - if guardrail_error_is_redacted: + if _is_error_data_redacted(guardrail_error): _detach_data_redacted_error_traceback(guardrail_error) - try: - await save_items_after_guardrails([synthesized_item] if include_in_history else []) - except BaseException as persistence_error: - if not guardrail_error_is_redacted: - raise - redacted_persistence_error = _safe_redacted_persistence_error(persistence_error) - if redacted_persistence_error is None: - raise - - if redacted_persistence_error is not None: - raise redacted_persistence_error from None + raise return validated_output, synthesized_item @@ -1410,6 +1522,12 @@ def _sync_conversation_tracking_from_tracker() -> None: current_agent = run_state._current_agent else: current_agent = starting_agent + _validate_resumed_session_output_guardrail_safety( + agent=current_agent, + run_config=run_config, + session=session, + run_state=run_state if is_resumed_state else None, + ) if run_state is not None: current_turn = run_state._current_turn else: @@ -1586,6 +1704,13 @@ async def _save_max_turns_items( try: while True: + validate_output_guardrails_with_server_managed_conversation( + current_agent, + run_config, + conversation_id=conversation_id, + previous_response_id=previous_response_id, + auto_previous_response_id=auto_previous_response_id, + ) all_input_guardrails = ( starting_agent.input_guardrails + (run_config.input_guardrails or []) if current_turn == 0 and not is_resumed_state @@ -1667,24 +1792,38 @@ async def _save_max_turns_items( raise UserError("No processed response found in previous state") last_model_response = run_state._model_responses[-1] + resumed_response_boundary = _current_response_boundary( + (), + run_state._last_processed_response, + run_state, + ) + blocked_output_owner_starts = _BlockedOutputOwnerStarts( + run_state_generated_items=resumed_response_boundary.generated_start, + run_state_session_items=resumed_response_boundary.session_start, + run_state_model_responses=len(run_state._model_responses) - 1, + run_state_tool_output_guardrail_results=len( + run_state._tool_output_guardrail_results + ), + streamed_new_items=resumed_response_boundary.session_start, + streamed_model_input_items=resumed_response_boundary.generated_start, + streamed_raw_responses=len(streamed_result.raw_responses) - 1, + streamed_tool_output_guardrail_results=len( + streamed_result.tool_output_guardrail_results + ), + ) - turn_result = await _run_with_deferred_tool_spans( - agent=current_agent, + turn_result = await resolve_interrupted_turn( + bindings=current_bindings, + original_input=run_state._original_input, + original_pre_step_items=run_state._generated_items, + new_response=last_model_response, + processed_response=run_state._last_processed_response, + hooks=hooks, + context_wrapper=context_wrapper, run_config=run_config, - run_turn=partial( - resolve_interrupted_turn, - bindings=current_bindings, - original_input=run_state._original_input, - original_pre_step_items=run_state._generated_items, - new_response=last_model_response, - processed_response=run_state._last_processed_response, - hooks=hooks, - context_wrapper=context_wrapper, - run_config=run_config, - server_manages_conversation=server_conversation_tracker is not None, - run_state=run_state, - error_handlers=error_handlers, - ), + server_manages_conversation=server_conversation_tracker is not None, + run_state=run_state, + error_handlers=error_handlers, ) tool_use_tracker.record_processed_response( @@ -1750,7 +1889,14 @@ async def _save_max_turns_items( await _finalize_streamed_interruption( streamed_result=streamed_result, save_items=_save_resumed_items, - items=list(turn_session_items), + items=( + [] + if _should_defer_interrupted_session_items( + current_agent, + run_config, + ) + else list(turn_session_items) + ), response_id=turn_result.model_response.response_id, store_setting=store_setting, interruptions=approvals_from_step(turn_result.next_step), @@ -1792,11 +1938,12 @@ async def _save_max_turns_items( save_items=_save_resumed_items, items=list(turn_session_items), model_response=turn_result.model_response, - deferred_tool_spans=turn_result.deferred_tool_spans, - preceding_items=turn_result.pre_step_items, - tool_output_guardrail_results=( - turn_result.tool_output_guardrail_results + processed_response=( + turn_result.processed_response + if turn_result.processed_response is not None + else run_state._last_processed_response ), + owner_starts=blocked_output_owner_starts, response_id=turn_result.model_response.response_id, store_setting=store_setting, ) @@ -1989,9 +2136,29 @@ def _record_max_turns_handler_output( save_items=_save_max_turns_items, items=[synthesized_item] if include_in_history else [], model_response=None, - deferred_tool_spans=[], - preceding_items=[], - tool_output_guardrail_results=[], + processed_response=None, + owner_starts=_BlockedOutputOwnerStarts( + run_state_generated_items=( + len(run_state._generated_items) if run_state is not None else None + ), + run_state_session_items=( + len(run_state._session_items) if run_state is not None else None + ), + run_state_model_responses=( + len(run_state._model_responses) if run_state is not None else None + ), + run_state_tool_output_guardrail_results=( + len(run_state._tool_output_guardrail_results) + if run_state is not None + else None + ), + streamed_new_items=len(streamed_result.new_items), + streamed_model_input_items=len(streamed_result._model_input_items), + streamed_raw_responses=len(streamed_result.raw_responses), + streamed_tool_output_guardrail_results=len( + streamed_result.tool_output_guardrail_results + ), + ), response_id=None, store_setting=store_setting, on_persisted_after_guardrails=_record_max_turns_handler_output, @@ -2045,6 +2212,28 @@ def _record_max_turns_handler_output( ) ) try: + blocked_output_owner_starts = _BlockedOutputOwnerStarts( + run_state_generated_items=( + len(run_state._generated_items) if run_state is not None else None + ), + run_state_session_items=( + len(run_state._session_items) if run_state is not None else None + ), + run_state_model_responses=( + len(run_state._model_responses) if run_state is not None else None + ), + run_state_tool_output_guardrail_results=( + len(run_state._tool_output_guardrail_results) + if run_state is not None + else None + ), + streamed_new_items=len(streamed_result.new_items), + streamed_model_input_items=len(streamed_result._model_input_items), + streamed_raw_responses=len(streamed_result.raw_responses), + streamed_tool_output_guardrail_results=len( + streamed_result.tool_output_guardrail_results + ), + ) logger.debug( "Starting turn %s, current_agent=%s", current_turn, @@ -2196,9 +2385,8 @@ def _record_max_turns_handler_output( save_items=_save_stream_items_with_count, items=turn_session_items, model_response=turn_result.model_response, - deferred_tool_spans=turn_result.deferred_tool_spans, - preceding_items=turn_result.pre_step_items, - tool_output_guardrail_results=turn_result.tool_output_guardrail_results, + processed_response=turn_result.processed_response, + owner_starts=blocked_output_owner_starts, response_id=turn_result.model_response.response_id, store_setting=store_setting, ) @@ -2225,7 +2413,14 @@ def _record_max_turns_handler_output( await _finalize_streamed_interruption( streamed_result=streamed_result, save_items=_save_stream_items_with_count, - items=turn_session_items, + items=( + [] + if _should_defer_interrupted_session_items( + current_agent, + run_config, + ) + else turn_session_items + ), response_id=turn_result.model_response.response_id, store_setting=store_setting, interruptions=approvals_from_step(turn_result.next_step), @@ -2648,28 +2843,23 @@ async def after_invocation_validation( async def check_input_guardrails_before_side_effects() -> None: await raise_if_input_guardrail_tripwire_known() - single_step_result = await _run_with_deferred_tool_spans( - agent=public_agent, + single_step_result = await get_single_step_result_from_response( + bindings=bindings, + original_input=streamed_result.input, + pre_step_items=streamed_result._model_input_items, + new_response=final_response, + output_schema=output_schema, + all_tools=all_tools, + handoffs=handoffs, + hooks=hooks, + context_wrapper=context_wrapper, run_config=run_config, - run_turn=partial( - get_single_step_result_from_response, - bindings=bindings, - original_input=streamed_result.input, - pre_step_items=streamed_result._model_input_items, - new_response=final_response, - output_schema=output_schema, - all_tools=all_tools, - handoffs=handoffs, - hooks=hooks, - context_wrapper=context_wrapper, - run_config=run_config, - error_handlers=error_handlers, - tool_use_tracker=tool_use_tracker, - server_manages_conversation=server_conversation_tracker is not None, - after_invocation_validation=after_invocation_validation, - before_side_effects=check_input_guardrails_before_side_effects, - run_state=run_state, - ), + error_handlers=error_handlers, + tool_use_tracker=tool_use_tracker, + server_manages_conversation=server_conversation_tracker is not None, + after_invocation_validation=after_invocation_validation, + before_side_effects=check_input_guardrails_before_side_effects, + run_state=run_state, ) items_to_filter = session_items_for_turn(single_step_result) @@ -2797,27 +2987,22 @@ async def after_invocation_validation( ) return response_accepted - return await _run_with_deferred_tool_spans( - agent=public_agent, + return await get_single_step_result_from_response( + bindings=bindings, + original_input=original_input, + pre_step_items=generated_items, + new_response=new_response, + output_schema=output_schema, + all_tools=all_tools, + handoffs=handoffs, + hooks=hooks, + context_wrapper=context_wrapper, run_config=run_config, - run_turn=partial( - get_single_step_result_from_response, - bindings=bindings, - original_input=original_input, - pre_step_items=generated_items, - new_response=new_response, - output_schema=output_schema, - all_tools=all_tools, - handoffs=handoffs, - hooks=hooks, - context_wrapper=context_wrapper, - run_config=run_config, - error_handlers=error_handlers, - tool_use_tracker=tool_use_tracker, - server_manages_conversation=server_conversation_tracker is not None, - after_invocation_validation=after_invocation_validation, - run_state=run_state, - ), + error_handlers=error_handlers, + tool_use_tracker=tool_use_tracker, + server_manages_conversation=server_conversation_tracker is not None, + after_invocation_validation=after_invocation_validation, + run_state=run_state, ) diff --git a/tests/test_agent_runner.py b/tests/test_agent_runner.py index 7e8cc0d26f..1e5569ebeb 100644 --- a/tests/test_agent_runner.py +++ b/tests/test_agent_runner.py @@ -13,8 +13,8 @@ import httpx import pytest from openai import APIConnectionError, BadRequestError, NotFoundError -from openai.types.responses import ResponseFunctionShellToolCallOutput, ResponseFunctionToolCall -from openai.types.responses.response_output_item import McpApprovalRequest, ProgramOutput +from openai.types.responses import ResponseFunctionToolCall +from openai.types.responses.response_output_item import McpApprovalRequest from openai.types.responses.response_output_text import AnnotationFileCitation, ResponseOutputText from openai.types.responses.response_reasoning_item import ResponseReasoningItem, Summary from openai.types.responses.tool_param import Mcp @@ -24,7 +24,6 @@ from agents import ( Agent, AgentOutputSchema, - AgentsException, GuardrailFunctionOutput, Handoff, HandoffInputData, @@ -70,6 +69,7 @@ from agents.lifecycle import RunHooks from agents.memory import SessionSettings from agents.models.fake_id import FAKE_RESPONSES_ID +from agents.result import RunResultStreaming from agents.run import AgentRunner, get_default_agent_runner, set_default_agent_runner from agents.run_config import _default_trace_include_sensitive_data from agents.run_internal import run_loop @@ -112,41 +112,10 @@ get_text_input_item, get_text_message, ) -from .testing_processor import SPAN_PROCESSOR_TESTING, fetch_ordered_spans from .utils.factories import make_run_state from .utils.hitl import make_context_wrapper, make_model_and_agent, make_shell_call from .utils.simple_session import CountingSession, IdStrippingSession, SimpleListSession -_BLOCKED_TOOL_OUTPUT = "Output withheld by an output guardrail." - - -def _sdk_exception_traceback_repr_locations( - error: BaseException, expected: str -) -> list[tuple[str, str]]: - pending = [error] - seen: set[int] = set() - locations: list[tuple[str, str]] = [] - while pending: - current = pending.pop() - if id(current) in seen: - continue - seen.add(id(current)) - - traceback = current.__traceback__ - while traceback is not None: - if "/src/agents/" in traceback.tb_frame.f_code.co_filename: - locations.extend( - (traceback.tb_frame.f_code.co_name, name) - for name, value in traceback.tb_frame.f_locals.items() - if expected in repr(value) - ) - traceback = traceback.tb_next - - for linked in (current.__cause__, current.__context__): - if linked is not None: - pending.append(linked) - return locations - class _DummyRunItem: def __init__(self, payload: dict[str, Any], item_type: str = "tool_call_output_item"): @@ -157,430 +126,464 @@ def to_input_item(self) -> dict[str, Any]: return self._payload -@pytest.mark.parametrize( - "raw_item", - [ - { +def test_blocked_function_batch_is_rebuilt_from_allowlisted_fields() -> None: + agent = Agent(name="test") + call = ToolCallItem( + agent=agent, + raw_item={ + "type": "function_call", + "name": "commit_tool", + "arguments": "{}", + "call_id": "call-commit", + "provider_data": {"secret": "call-secret"}, + }, + ) + output = ToolCallOutputItem( + agent=agent, + raw_item={ "type": "function_call_output", - "call_id": "call-function", - "output": "secret", - "caller": { - "type": "program", - "caller_id": "program-1", - "duplicate": "secret", - }, + "call_id": "call-commit", + "output": "raw-secret", + "provider_data": {"secret": "output-secret"}, }, - { + output="sdk-secret", + custom_data={"secret": "custom-secret"}, + ) + + retained = run_loop._retained_items_for_blocked_output([call, output]) + + assert [item.type for item in retained] == ["tool_call_item", "tool_call_output_item"] + retained_call = cast(ToolCallItem, retained[0]) + retained_output = cast(ToolCallOutputItem, retained[1]) + assert "provider_data" not in cast(dict[str, Any], retained_call.raw_item) + assert cast(dict[str, Any], retained_output.raw_item) == { + "type": "function_call_output", + "call_id": "call-commit", + "output": run_loop._OUTPUT_GUARDRAIL_BLOCKED_TOOL_OUTPUT, + } + assert retained_output.output == run_loop._OUTPUT_GUARDRAIL_BLOCKED_TOOL_OUTPUT + assert retained_output.custom_data is None + + +def test_blocked_unknown_tool_variant_discards_the_complete_response() -> None: + agent = Agent(name="test") + call = ToolCallItem( + agent=agent, + raw_item={"type": "custom_tool_call", "call_id": "call-custom", "secret": "call"}, + ) + output = ToolCallOutputItem( + agent=agent, + raw_item={ "type": "custom_tool_call_output", "call_id": "call-custom", - "output": "secret", - }, - { - "type": "local_shell_call_output", - "call_id": "call-local-shell", - "output": "secret", + "output": "raw-secret", }, - { - "type": "apply_patch_call_output", - "call_id": "call-apply-patch", - "output": "secret", - "status": "completed", - }, - { - "type": "shell_call_output", - "call_id": "call-shell", - "output": [ - { - "stdout": "secret", - "stderr": "", - "outcome": {"type": "exit", "exit_code": 0}, - } - ], - "status": "completed", - "shell_output": [{"stdout": "secret"}], - "provider_data": {"raw_output": "secret"}, - "caller": { - "type": "program", - "caller_id": "program-1", - "duplicate": "secret", - }, - }, - ResponseFunctionShellToolCallOutput( - id="shell-output-in-progress", - call_id="call-shell-in-progress", - output=[ - { - "stdout": "secret", - "stderr": "", - "outcome": {"type": "exit", "exit_code": 0}, - } - ], - status="in_progress", - type="shell_call_output", - ), - ResponseFunctionShellToolCallOutput( - id="shell-output-incomplete", - call_id="call-shell-incomplete", - output=[ - { - "stdout": "secret", - "stderr": "", - "outcome": {"type": "exit", "exit_code": 0}, - } - ], - status="incomplete", - type="shell_call_output", + output="sdk-secret", + custom_data={"secret": "custom-secret"}, + ) + + assert run_loop._retained_items_for_blocked_output([call, output]) == [] + + +def test_blocked_reasoning_item_discards_the_complete_response() -> None: + agent = Agent(name="test") + reasoning = ReasoningItem( + agent=agent, + raw_item=ResponseReasoningItem( + id="reasoning-id", + type="reasoning", + summary=[], + encrypted_content="reasoning-secret", ), - { - "type": "computer_call_output", - "call_id": "call-computer", - "output": { - "type": "computer_screenshot", - "image_url": "data:image/png;base64,c2VjcmV0", - }, - "acknowledged_safety_checks": [ - { - "id": "check-1", - "code": "secret", - "message": "secret", - "duplicate": "secret", - } - ], - }, - { - "type": "program_output", - "id": "program-output", - "call_id": "call-program", - "result": "secret", - "status": "completed", + ) + call = ToolCallItem( + agent=agent, + raw_item={ + "type": "function_call", + "name": "commit_tool", + "arguments": "{}", + "call_id": "call-commit", }, - ], -) -@pytest.mark.asyncio -async def test_blocked_retained_tool_output_uses_replay_valid_payload( - raw_item: Any, -) -> None: - raw_payload = ( - raw_item if isinstance(raw_item, dict) else raw_item.model_dump(exclude_unset=True) ) - agent = Agent(name="test") - item = ToolCallOutputItem( + output = ToolCallOutputItem( agent=agent, - raw_item=cast(Any, raw_item), - output="secret", - custom_data={"duplicate": "secret"}, - ) - retained_input: list[RunItem] = [] - if raw_payload.get("caller") is not None or raw_payload["type"] == "program_output": - retained_input.append( - ToolCallItem( - agent=agent, - raw_item=cast( - Any, - { - "type": "program", - "call_id": ( - raw_payload["call_id"] - if raw_payload["type"] == "program_output" - else raw_payload["caller"]["caller_id"] - ), - "code": "return await tools.lookup({});", - }, - ), - ) - ) - retained_input.append(item) - - retained = run_loop._retained_items_for_blocked_output(retained_input) - - assert retained == retained_input - assert item.output == _BLOCKED_TOOL_OUTPUT - assert item.custom_data is None - payload = cast(dict[str, Any], item.to_input_item()) - assert payload["type"] == raw_payload["type"] - assert "secret" not in json.dumps(payload) - if "caller" in raw_payload: - assert payload["caller"] == {"type": "program", "caller_id": "program-1"} - sanitized_raw_item = cast(dict[str, Any], item.raw_item) - if raw_payload["type"] == "shell_call_output" and "status" in raw_payload: - assert sanitized_raw_item["status"] == raw_payload["status"] - - if payload["type"] == "computer_call_output": - assert payload["output"]["type"] == "computer_screenshot" - assert payload["output"]["image_url"].startswith("data:image/png;base64,") - elif payload["type"] == "shell_call_output": - assert payload["output"] == [ - { - "stdout": "", - "stderr": _BLOCKED_TOOL_OUTPUT, - "outcome": {"type": "exit", "exit_code": 1}, - } - ] - elif payload["type"] == "program_output": - assert payload["result"] == _BLOCKED_TOOL_OUTPUT - else: - assert payload["output"] == _BLOCKED_TOOL_OUTPUT + raw_item={ + "type": "function_call_output", + "call_id": "call-commit", + "output": "raw-secret", + }, + output="sdk-secret", + ) - state = make_run_state(agent) - state._generated_items = retained_input - state._session_items = retained_input - state_json = state.to_json() - assert "secret" not in json.dumps(state_json) - restored = await RunState.from_json(agent, state_json) - restored_item = cast(ToolCallOutputItem, restored._generated_items[-1]) - restored_payload = cast(dict[str, Any], restored_item.to_input_item()) - assert restored_payload == payload + assert run_loop._retained_items_for_blocked_output([reasoning, call, output]) == [] -@pytest.mark.parametrize( - "raw_item", - [ - { - "type": "local_shell_call_output", - "id": "unsupported-provider-id", - "output": "secret", +def test_blocked_snapshot_preserves_accepted_prefix_with_reused_provider_id() -> None: + agent = Agent(name="test") + prior_call = ToolCallItem( + agent=agent, + raw_item={ + "type": "function_call", + "name": "prior_tool", + "arguments": "{}", + "call_id": "reused-call-id", }, - { + ) + prior_output = ToolCallOutputItem( + agent=agent, + raw_item={ "type": "function_call_output", - "call_id": "call-function", - "output": "secret", - "caller": {"type": "program", "duplicate": "secret"}, + "call_id": "reused-call-id", + "output": "accepted-prior-output", }, - { - "type": "computer_call_output", - "call_id": "call-computer", - "output": { - "type": "computer_screenshot", - "image_url": "data:image/png;base64,c2VjcmV0", - }, - "acknowledged_safety_checks": [{"id": ""}], + output="accepted-prior-output", + ) + current_call = ToolCallItem( + agent=agent, + raw_item={ + "type": "function_call", + "name": "current_tool", + "arguments": "{}", + "call_id": "reused-call-id", }, - { + ) + current_output = ToolCallOutputItem( + agent=agent, + raw_item={ "type": "function_call_output", - "call_id": "call-function", - "output": "secret", - "status": "failed", + "call_id": "reused-call-id", + "output": "rejected-current-output", }, - ], -) -def test_blocked_retained_tool_output_rejects_malformed_replay_identity( - raw_item: dict[str, Any], -) -> None: - item = ToolCallOutputItem( - agent=Agent(name="test"), - raw_item=cast(Any, raw_item), - output="secret", + output="rejected-current-output", ) - model_response = ModelResponse( - output=[get_text_message("archived-secret")], + prior_response = ModelResponse( + output=[cast(Any, prior_call.raw_item)], usage=Usage(), - response_id="response-malformed", + response_id="prior-response", + ) + current_response = ModelResponse( + output=[cast(Any, current_call.raw_item)], + usage=Usage(), + response_id="current-response", + ) + state = make_run_state( + agent, + context=make_context_wrapper(), + original_input="test", + max_turns=2, + ) + state._generated_items = [prior_call, prior_output, current_call, current_output] + state._session_items = [prior_call, prior_output, current_call, current_output] + state._model_responses = [prior_response, current_response] + + retained = run_loop._retained_items_for_blocked_response( + [current_call, current_output], + current_response, + run_state=state, + owner_starts=run_loop._BlockedOutputOwnerStarts( + run_state_generated_items=2, + run_state_session_items=2, + run_state_model_responses=1, + run_state_tool_output_guardrail_results=0, + ), ) - with pytest.raises( - AgentsException, match="Cannot sanitize a blocked tool output for replay" - ) as exc_info: - run_loop._retained_items_for_blocked_output([item], model_response) - - assert "secret" not in str(exc_info.value) - assert exc_info.value.run_data is None - assert exc_info.value.__cause__ is None - assert exc_info.value.__context__ is None - assert _sdk_exception_traceback_repr_locations(exc_info.value, "secret") == [] - assert item.output == _BLOCKED_TOOL_OUTPUT - assert item.custom_data is None - assert item.to_input_item() == { - "type": "function_call_output", - "call_id": "blocked-tool-output", - "output": _BLOCKED_TOOL_OUTPUT, - } - assert model_response.output == [] + assert state._generated_items[:2] == [prior_call, prior_output] + assert state._generated_items[0] is prior_call + assert state._generated_items[1] is prior_output + assert state._session_items[:2] == [prior_call, prior_output] + assert state._model_responses[0] is prior_response + assert retained == state._generated_items[2:] + assert cast(ToolCallOutputItem, retained[1]).output == ( + run_loop._OUTPUT_GUARDRAIL_BLOCKED_TOOL_OUTPUT + ) + assert prior_output.output == "accepted-prior-output" -@pytest.mark.parametrize( - "raw_item", - [ - { - "type": "function_call_output", - "call_id": "call-orphan", - "output": "secret", - "caller": {"type": "program", "caller_id": "program-orphan"}, +def test_blocked_snapshot_cancellation_severs_replay_graph_and_propagates( + monkeypatch: pytest.MonkeyPatch, +) -> None: + agent = Agent(name="test") + call = ToolCallItem( + agent=agent, + raw_item={ + "type": "function_call", + "name": "commit_tool", + "arguments": "{}", + "call_id": "call-commit", }, - { - "type": "program_output", - "id": "program-output-orphan", - "call_id": "program-orphan", - "result": "secret", - "status": "completed", + ) + output = ToolCallOutputItem( + agent=agent, + raw_item={ + "type": "function_call_output", + "call_id": "call-commit", + "output": "raw-secret", }, - ], -) -def test_blocked_retained_tool_output_rejects_orphan_program_relationship( - raw_item: dict[str, Any], -) -> None: - item = ToolCallOutputItem( - agent=Agent(name="test"), - raw_item=cast(Any, raw_item), - output="secret", - custom_data={"duplicate": "secret"}, + output="sdk-secret", + custom_data={"secret": "custom-secret"}, ) + response = ModelResponse( + output=[cast(Any, call.raw_item)], + usage=Usage(), + response_id="response-id", + ) + state = make_run_state( + agent, + context=make_context_wrapper(), + original_input="test", + max_turns=1, + ) + state._generated_items = [call, output] + state._session_items = [call, output] + state._model_responses = [response] + cancellation = asyncio.CancelledError("original cancellation") - with pytest.raises( - AgentsException, match="Cannot sanitize a blocked tool output for replay" - ) as exc_info: - run_loop._retained_items_for_blocked_output([item]) + def cancel_preparation(_raw_item: Any) -> dict[str, Any]: + raise cancellation - assert "secret" not in str(exc_info.value) - assert exc_info.value.__cause__ is None - assert exc_info.value.__context__ is None - assert _sdk_exception_traceback_repr_locations(exc_info.value, "secret") == [] - assert item.output == _BLOCKED_TOOL_OUTPUT - assert item.custom_data is None + monkeypatch.setattr(run_loop, "_blocked_function_output_payload", cancel_preparation) + with pytest.raises(asyncio.CancelledError) as exc_info: + run_loop._retained_items_for_blocked_response( + [call, output], + response, + run_state=state, + ) -def test_blocked_retained_tool_output_accepts_program_parent_from_prior_history() -> None: + assert exc_info.value is cancellation + assert state._generated_items == [] + assert state._session_items == [] + assert state._model_responses == [] + + +@pytest.mark.parametrize("fail_after_first_swap", [False, True]) +def test_blocked_snapshot_application_baseexception_severs_every_owner( + monkeypatch: pytest.MonkeyPatch, + fail_after_first_swap: bool, +) -> None: agent = Agent(name="test") - program_item = ToolCallItem( + prior_call = ToolCallItem( agent=agent, - raw_item=cast( - Any, - { - "type": "program", - "call_id": "program-prior", - "code": "return await tools.lookup({});", - }, - ), + raw_item={ + "type": "function_call", + "name": "prior_tool", + "arguments": "{}", + "call_id": "prior-call", + }, ) - child_call = ToolCallItem( + prior_output = ToolCallOutputItem( agent=agent, - raw_item=cast( - Any, - { - "type": "function_call", - "call_id": "call-child", - "name": "lookup", - "arguments": "{}", - "caller": {"type": "program", "caller_id": "program-prior"}, - }, - ), + raw_item={ + "type": "function_call_output", + "call_id": "prior-call", + "output": "accepted-prior-output", + }, + output="accepted-prior-output", ) - child_output = ToolCallOutputItem( + call = ToolCallItem( agent=agent, - raw_item=cast( - Any, - { - "type": "function_call_output", - "call_id": "call-child", - "output": "program-secret", - "caller": {"type": "program", "caller_id": "program-prior"}, - }, - ), - output="program-secret", + raw_item={ + "type": "function_call", + "name": "commit_tool", + "arguments": "{}", + "call_id": "call-commit", + }, ) - - retained = run_loop._retained_items_for_blocked_output( - [child_call, child_output], - preceding_items=[program_item], + output = ToolCallOutputItem( + agent=agent, + raw_item={ + "type": "function_call_output", + "call_id": "call-commit", + "output": "raw-secret", + }, + output="sdk-secret", + ) + prior_response = ModelResponse( + output=[cast(Any, prior_call.raw_item)], + usage=Usage(), + response_id="prior-response", + ) + response = ModelResponse( + output=[cast(Any, call.raw_item)], + usage=Usage(), + response_id="response-id", + ) + state = make_run_state( + agent, + context=make_context_wrapper(), + original_input="test", + max_turns=1, ) + prior_guardrail_result = cast(Any, object()) + current_guardrail_result = cast(Any, object()) + state._generated_items = [prior_call, prior_output, call, output] + state._session_items = [prior_call, prior_output, call, output] + state._model_responses = [prior_response, response] + state._tool_output_guardrail_results = [prior_guardrail_result, current_guardrail_result] + streamed_result = RunResultStreaming( + input="test", + new_items=[prior_call, prior_output, call, output], + raw_responses=[prior_response, response], + final_output=None, + input_guardrail_results=[], + output_guardrail_results=[], + tool_input_guardrail_results=[], + tool_output_guardrail_results=[prior_guardrail_result, current_guardrail_result], + context_wrapper=make_context_wrapper(), + current_agent=agent, + current_turn=2, + max_turns=2, + _current_agent_output_schema=None, + trace=None, + ) + streamed_result._model_input_items = [prior_call, prior_output, call, output] + streamed_result._state = state + application_error = KeyboardInterrupt("application failed") + + def fail_application(plan: Any) -> None: + if fail_after_first_swap: + owner, field, value = plan.assignments[0] + object.__setattr__(owner, field, value) + raise application_error + + monkeypatch.setattr(run_loop, "_apply_blocked_output_owner_plan", fail_application) + + with pytest.raises(KeyboardInterrupt) as exc_info: + run_loop._retained_items_for_blocked_response( + [call, output], + response, + run_state=state, + streamed_result=streamed_result, + owner_starts=run_loop._BlockedOutputOwnerStarts( + run_state_generated_items=2, + run_state_session_items=2, + run_state_model_responses=1, + run_state_tool_output_guardrail_results=1, + streamed_new_items=2, + streamed_model_input_items=2, + streamed_raw_responses=1, + streamed_tool_output_guardrail_results=1, + ), + ) - assert retained == [child_call, child_output] - assert child_output.output == _BLOCKED_TOOL_OUTPUT - assert program_item not in retained + assert exc_info.value is application_error + assert state._generated_items == [prior_call, prior_output] + assert state._session_items == [prior_call, prior_output] + assert state._model_responses == [prior_response] + assert state._tool_output_guardrail_results == [prior_guardrail_result] + assert streamed_result.new_items == [prior_call, prior_output] + assert streamed_result._model_input_items == [prior_call, prior_output] + assert streamed_result.raw_responses == [prior_response] + assert streamed_result.tool_output_guardrail_results == [prior_guardrail_result] -@pytest.mark.asyncio -async def test_malformed_blocked_tool_output_failure_is_data_free_at_runner_boundary( +def test_blocked_snapshot_application_exception_becomes_fixed_error( monkeypatch: pytest.MonkeyPatch, ) -> None: - original_sanitizer = run_loop._blocked_tool_output_payload + agent = Agent(name="test") + call = ToolCallItem( + agent=agent, + raw_item={ + "type": "function_call", + "name": "commit_tool", + "arguments": "{}", + "call_id": "call-commit", + }, + ) + output = ToolCallOutputItem( + agent=agent, + raw_item={ + "type": "function_call_output", + "call_id": "call-commit", + "output": "raw-secret", + }, + output="sdk-secret", + ) + state = make_run_state( + agent, + context=make_context_wrapper(), + original_input="test", + max_turns=1, + ) + state._generated_items = [call, output] + state._session_items = [call, output] - def reject_with_malformed_caller(_raw_item: Any) -> dict[str, Any]: - return original_sanitizer( - { - "type": "function_call_output", - "call_id": "call-malformed", - "output": "runner-boundary-secret", - "caller": {"type": "program", "duplicate": "runner-boundary-secret"}, - } + def fail_application(_plan: Any) -> None: + raise ValueError("application-secret") + + monkeypatch.setattr(run_loop, "_apply_blocked_output_owner_plan", fail_application) + + with pytest.raises(RuntimeError) as exc_info: + run_loop._retained_items_for_blocked_response( + [call, output], + None, + run_state=state, ) - monkeypatch.setattr(run_loop, "_blocked_tool_output_payload", reject_with_malformed_caller) + assert "application-secret" not in str(exc_info.value) + assert state._generated_items == [] + assert state._session_items == [] + + +@pytest.mark.asyncio +async def test_non_streamed_trip_preserves_prior_run_state_side_effect() -> None: + side_effects: list[str] = [] + + @function_tool(name_override="accepted_tool") + def accepted_tool() -> str: + side_effects.append("accepted") + return "accepted-output" @function_tool(name_override="terminal_tool") def terminal_tool() -> str: - return "runner-boundary-secret" - - def output_guardrail( - _context: RunContextWrapper[Any], - _agent: Agent[Any], - output: Any, - ) -> GuardrailFunctionOutput: - return GuardrailFunctionOutput(output_info=output, tripwire_triggered=True) + side_effects.append("terminal") + return "rejected-output" model = ScriptedModel( - [[get_function_tool_call("terminal_tool", "{}", call_id="call-terminal")]] - ) - agent = Agent( - name="test", - model=model, - tools=[terminal_tool], - tool_use_behavior="stop_on_first_tool", - output_guardrails=[OutputGuardrail(guardrail_function=output_guardrail)], + steps=[ + [get_function_tool_call("accepted_tool", "{}", call_id="accepted-call")], + [get_text_message("accepted-final")], + ] ) + agent = Agent(name="test", model=model, tools=[accepted_tool, terminal_tool]) + first = await Runner.run(agent, "run accepted tool", max_turns=5) + state = first.to_state() + prior_generated = list(state._generated_items) + prior_session = list(state._session_items) + prior_responses = list(state._model_responses) - with pytest.raises( - AgentsException, match="Cannot sanitize a blocked tool output for replay" - ) as exc_info: - await Runner.run(agent, "Run terminal_tool") - - assert exc_info.value.run_data is None - assert exc_info.value.__cause__ is None - assert exc_info.value.__context__ is None - assert _sdk_exception_traceback_repr_locations(exc_info.value, "runner-boundary-secret") == [] - + def reject_output( + _context: RunContextWrapper[Any], + _agent: Agent[Any], + _output: Any, + ) -> GuardrailFunctionOutput: + return GuardrailFunctionOutput(output_info=None, tripwire_triggered=True) -@pytest.mark.parametrize("shared_raw_item", [False, True]) -def test_blocked_retained_tool_output_sanitizes_archived_model_response( - shared_raw_item: bool, -) -> None: - agent = Agent(name="test") - raw_output = ProgramOutput( - id="program-output", - call_id="call-program", - result="secret-program-result", - status="completed", - type="program_output", - ) - item = ToolCallOutputItem( - agent=agent, - raw_item=raw_output if shared_raw_item else cast(Any, raw_output.model_dump()), - output=raw_output.result, - ) - program_item = ToolCallItem( - agent=agent, - raw_item=cast( - Any, - { - "type": "program", - "call_id": "call-program", - "code": "return await tools.lookup({});", - }, - ), + agent.tool_use_behavior = {"stop_at_tool_names": ["terminal_tool"]} + agent.output_guardrails = [OutputGuardrail(guardrail_function=reject_output)] + model.enqueue( + [ + ResponseReasoningItem( + id="reasoning-current", + type="reasoning", + summary=[Summary(text="calling terminal tool", type="summary_text")], + ), + get_function_tool_call("terminal_tool", "{}", call_id="current-call"), + ] ) - model_response = ModelResponse(output=[raw_output], usage=Usage(), response_id="response") - state = make_run_state(agent) - state._model_responses = [model_response] - state._generated_items = [program_item, item] - state._session_items = [program_item, item] - retained = run_loop._retained_items_for_blocked_output([program_item, item], model_response) + with pytest.raises(OutputGuardrailTripwireTriggered): + await Runner.run(agent, state) - assert retained == [program_item, item] - assert "secret-program-result" not in json.dumps(state.to_json()) - archived_output = cast(dict[str, Any], model_response.output[0]) - assert archived_output["type"] == "program_output" - assert archived_output["call_id"] == "call-program" - assert archived_output["result"] == _BLOCKED_TOOL_OUTPUT + assert side_effects == ["accepted", "terminal"] + assert state._generated_items == prior_generated + assert state._session_items == prior_session + assert state._model_responses == prior_responses + serialized_state = json.dumps(state.to_json()) + assert "accepted-output" in serialized_state + assert "rejected-output" not in serialized_state + assert "reasoning-current" not in serialized_state async def run_execute_approved_tools( @@ -4891,7 +4894,7 @@ def guardrail_function( context: RunContextWrapper[Any], agent: Agent[Any], agent_output: Any ) -> GuardrailFunctionOutput: return GuardrailFunctionOutput( - output_info="safe-info", + output_info=None, tripwire_triggered=True, ) @@ -4903,12 +4906,9 @@ def guardrail_function( ) model.enqueue([get_text_message("user_message")]) - with pytest.raises(OutputGuardrailTripwireTriggered) as exc_info: + with pytest.raises(OutputGuardrailTripwireTriggered): await Runner.run(agent, input="user_message") - assert exc_info.value.guardrail_result.agent_output == "user_message" - assert exc_info.value.guardrail_result.output.output_info == "safe-info" - def test_output_guardrail_tripwire_does_not_save_assistant_message_to_session_sync() -> None: def guardrail_function( @@ -4936,7 +4936,7 @@ def guardrail_function( @pytest.mark.asyncio -async def test_output_guardrail_error_preserves_final_output_in_session() -> None: +async def test_output_guardrail_error_does_not_persist_unverdictable_output() -> None: def guardrail_function( _context: RunContextWrapper[Any], _agent: Agent[Any], _agent_output: Any ) -> GuardrailFunctionOutput: @@ -4944,7 +4944,7 @@ def guardrail_function( session = SimpleListSession() model = ScriptedModel() - model.enqueue([get_text_message("preserved_on_guardrail_error")]) + model.enqueue([get_text_message("not_persisted_on_guardrail_error")]) agent = Agent( name="test", model=model, @@ -4958,11 +4958,11 @@ def guardrail_function( assert [ cast(dict[str, Any], item).get("type") or cast(dict[str, Any], item).get("role") for item in items - ] == ["user", "message"] + ] == ["user"] @pytest.mark.asyncio -async def test_output_guardrail_cancellation_preserves_final_output_in_session() -> None: +async def test_output_guardrail_cancellation_does_not_start_a_final_session_write() -> None: guardrail_started = asyncio.Event() async def guardrail_function( @@ -4974,7 +4974,7 @@ async def guardrail_function( session = SimpleListSession() model = ScriptedModel() - model.enqueue([get_text_message("preserved_on_guardrail_cancellation")]) + model.enqueue([get_text_message("not_persisted_on_guardrail_cancellation")]) agent = Agent( name="test", model=model, @@ -4992,7 +4992,7 @@ async def guardrail_function( assert [ cast(dict[str, Any], item).get("type") or cast(dict[str, Any], item).get("role") for item in items - ] == ["user", "message"] + ] == ["user"] @pytest.mark.asyncio @@ -5045,24 +5045,20 @@ def foo(a: str) -> str: @pytest.mark.parametrize("tripwire_triggered", [False, True]) @pytest.mark.asyncio -async def test_resumed_final_tool_sanitizes_output_after_output_guardrail_tripwire( +async def test_resumed_final_tool_persists_call_and_output_after_output_guardrail( tripwire_triggered: bool, ) -> None: - tool_calls = 0 - def guardrail_function( - _context: RunContextWrapper[Any], _agent: Agent[Any], agent_output: Any + _context: RunContextWrapper[Any], _agent: Agent[Any], _agent_output: Any ) -> GuardrailFunctionOutput: return GuardrailFunctionOutput( - output_info=agent_output, + output_info=None, tripwire_triggered=tripwire_triggered, ) @function_tool(name_override="commit_tool") def commit_tool() -> str: - nonlocal tool_calls - tool_calls += 1 - return f"committed-result-{tool_calls}" + return "committed-result" session = SimpleListSession() model = ScriptedModel() @@ -5086,15 +5082,13 @@ def commit_tool() -> str: model.enqueue([get_function_tool_call("commit_tool", "{}", call_id="call-second")]) if tripwire_triggered: - with pytest.raises(OutputGuardrailTripwireTriggered) as exc_info: + with pytest.raises(OutputGuardrailTripwireTriggered): await Runner.run(agent, state, session=session) - assert exc_info.value.guardrail_result.agent_output == _BLOCKED_TOOL_OUTPUT - assert exc_info.value.guardrail_result.output.output_info is None else: result = await Runner.run(agent, state, session=session) - assert result.final_output == "committed-result-2" + assert result.final_output == "committed-result" - assert state._current_turn_persisted_item_count == 4 + assert state._current_turn_persisted_item_count == 2 items = await session.get_items() assert [ ( @@ -5109,215 +5103,11 @@ def commit_tool() -> str: ("function_call", "call-second"), ("function_call_output", "call-second"), ] - second_output = cast(dict[str, Any], items[-1]).get("output") - assert second_output == (_BLOCKED_TOOL_OUTPUT if tripwire_triggered else "committed-result-2") - - serialized_state = json.dumps(state.to_json()) - if tripwire_triggered: - assert "committed-result-2" not in serialized_state - - -@pytest.mark.parametrize("mode", ["non_streamed", "streamed"]) -@pytest.mark.parametrize("tripwire_triggered", [False, True]) -@pytest.mark.asyncio -async def test_terminal_tool_span_waits_for_output_guardrail_verdict( - mode: str, - tripwire_triggered: bool, - monkeypatch: pytest.MonkeyPatch, -) -> None: - started_function_spans: list[Any] = [] - original_on_span_start = SPAN_PROCESSOR_TESTING.on_span_start - - def capture_span_start(span: Any) -> None: - original_on_span_start(span) - if span.span_data.type == "function": - started_function_spans.append(span) - - monkeypatch.setattr(SPAN_PROCESSOR_TESTING, "on_span_start", capture_span_start) - - @function_tool(name_override="trace_tool") - def trace_tool() -> str: - return "trace-secret" - - def output_guardrail( - _context: RunContextWrapper[Any], - _agent: Agent[Any], - output: Any, - ) -> GuardrailFunctionOutput: - assert len(started_function_spans) == 1 - assert cast(Any, started_function_spans[0].span_data).output is None - assert not [span for span in fetch_ordered_spans() if span.span_data.type == "function"] - return GuardrailFunctionOutput( - output_info=output, - tripwire_triggered=tripwire_triggered, - ) - - model_output = [get_function_tool_call("trace_tool", "{}", call_id="call-trace")] - model = ScriptedModel( - [model_output if mode == "non_streamed" else get_exact_output_stream_step(model_output)] - ) - agent = Agent( - name="test", - model=model, - tools=[trace_tool], - tool_use_behavior="stop_on_first_tool", - output_guardrails=[OutputGuardrail(guardrail_function=output_guardrail)], + assert cast(dict[str, Any], items[-1]).get("output") == ( + run_loop._OUTPUT_GUARDRAIL_BLOCKED_TOOL_OUTPUT if tripwire_triggered else "committed-result" ) - run_config = RunConfig(trace_include_sensitive_data=True) - if tripwire_triggered: - with pytest.raises(OutputGuardrailTripwireTriggered): - if mode == "non_streamed": - await Runner.run(agent, "Use trace_tool", run_config=run_config) - else: - result = Runner.run_streamed(agent, "Use trace_tool", run_config=run_config) - async for _ in result.stream_events(): - pass - elif mode == "non_streamed": - await Runner.run(agent, "Use trace_tool", run_config=run_config) - else: - result = Runner.run_streamed(agent, "Use trace_tool", run_config=run_config) - async for _ in result.stream_events(): - pass - - function_spans = [span for span in fetch_ordered_spans() if span.span_data.type == "function"] - assert len(function_spans) == 1 - expected_output = _BLOCKED_TOOL_OUTPUT if tripwire_triggered else "trace-secret" - assert cast(Any, function_spans[0].span_data).output == expected_output - assert cast(Any, started_function_spans[0].span_data).output == expected_output - if tripwire_triggered: - assert "trace-secret" not in json.dumps(function_spans[0].export()) - - -@pytest.mark.parametrize("mode", ["non_streamed", "streamed"]) -@pytest.mark.parametrize("tripwire_triggered", [False, True]) -@pytest.mark.asyncio -async def test_blocked_terminal_tool_span_discards_deferred_error( - mode: str, - tripwire_triggered: bool, - monkeypatch: pytest.MonkeyPatch, -) -> None: - started_function_spans: list[Any] = [] - original_on_span_start = SPAN_PROCESSOR_TESTING.on_span_start - - def capture_span_start(span: Any) -> None: - original_on_span_start(span) - if span.span_data.type == "function": - started_function_spans.append(span) - - monkeypatch.setattr(SPAN_PROCESSOR_TESTING, "on_span_start", capture_span_start) - - def expose_error(_context: RunContextWrapper[Any], error: Exception) -> str: - return str(error) - - @function_tool(name_override="failing_tool", failure_error_function=expose_error) - def failing_tool() -> str: - raise ValueError("span-error-secret") - - def output_guardrail( - _context: RunContextWrapper[Any], - _agent: Agent[Any], - output: Any, - ) -> GuardrailFunctionOutput: - assert output == "span-error-secret" - assert len(started_function_spans) == 1 - assert started_function_spans[0].error is None - return GuardrailFunctionOutput( - output_info=output, - tripwire_triggered=tripwire_triggered, - ) - - model_output = [get_function_tool_call("failing_tool", "{}", call_id="call-failing")] - model = ScriptedModel( - [model_output if mode == "non_streamed" else get_exact_output_stream_step(model_output)] - ) - agent = Agent( - name="test", - model=model, - tools=[failing_tool], - tool_use_behavior="stop_on_first_tool", - output_guardrails=[OutputGuardrail(guardrail_function=output_guardrail)], - ) - - async def run_once() -> None: - if mode == "non_streamed": - await Runner.run( - agent, - "Use failing_tool", - run_config=RunConfig(trace_include_sensitive_data=True), - ) - else: - result = Runner.run_streamed( - agent, - "Use failing_tool", - run_config=RunConfig(trace_include_sensitive_data=True), - ) - async for _ in result.stream_events(): - pass - - if tripwire_triggered: - with pytest.raises(OutputGuardrailTripwireTriggered): - await run_once() - else: - await run_once() - - function_spans = [span for span in fetch_ordered_spans() if span.span_data.type == "function"] - assert len(function_spans) == 1 - if tripwire_triggered: - assert function_spans[0].error is None - assert cast(Any, function_spans[0].span_data).output == _BLOCKED_TOOL_OUTPUT - assert "span-error-secret" not in json.dumps(function_spans[0].export()) - else: - assert function_spans[0].error is not None - assert "span-error-secret" in json.dumps(function_spans[0].export()) - - -@pytest.mark.parametrize("behavior_kind", ["stop_at_tools", "custom"]) -@pytest.mark.asyncio -async def test_blocked_tool_output_is_sanitized_for_terminal_tool_behaviors( - behavior_kind: str, -) -> None: - @function_tool(name_override="commit_tool") - def commit_tool() -> str: - return "sensitive-result" - - def custom_behavior( - _context: RunContextWrapper[Any], results: list[FunctionToolResult] - ) -> ToolsToFinalOutputResult: - return ToolsToFinalOutputResult(is_final_output=True, final_output=results[0].output) - - def output_guardrail( - _context: RunContextWrapper[Any], _agent: Agent[Any], _output: Any - ) -> GuardrailFunctionOutput: - return GuardrailFunctionOutput(output_info=None, tripwire_triggered=True) - - tool_use_behavior: Any = ( - {"stop_at_tool_names": ["commit_tool"]} - if behavior_kind == "stop_at_tools" - else custom_behavior - ) - model = ScriptedModel([[get_function_tool_call("commit_tool", "{}", call_id="call-committed")]]) - session = SimpleListSession() - agent = Agent( - name="test", - model=model, - tools=[commit_tool], - tool_use_behavior=tool_use_behavior, - output_guardrails=[OutputGuardrail(guardrail_function=output_guardrail)], - ) - - with pytest.raises(OutputGuardrailTripwireTriggered): - await Runner.run(agent, "Use commit_tool", session=session) - - items = await session.get_items() - tool_output = next( - item - for item in items - if isinstance(item, dict) and item.get("type") == "function_call_output" - ) - assert tool_output.get("call_id") == "call-committed" - assert tool_output.get("output") == _BLOCKED_TOOL_OUTPUT - assert "sensitive-result" not in json.dumps(items) + assert "committed-result" not in json.dumps(items[-2:]) @pytest.mark.asyncio diff --git a/tests/test_agent_runner_streamed.py b/tests/test_agent_runner_streamed.py index b2d20d7b1b..e712af1a53 100644 --- a/tests/test_agent_runner_streamed.py +++ b/tests/test_agent_runner_streamed.py @@ -12,10 +12,8 @@ from openai import APIConnectionError, BadRequestError, NotFoundError from openai.types.responses import ( ResponseCompletedEvent, - ResponseCustomToolCall, ResponseErrorEvent, ResponseFailedEvent, - ResponseFunctionShellToolCallOutput, ResponseFunctionToolCall, ResponseIncompleteEvent, ) @@ -54,13 +52,17 @@ from agents.run import RunConfig from agents.run_internal import run_loop from agents.run_internal.run_loop import QueueCompleteSentinel -from agents.run_internal.run_steps import NextStepRunAgain +from agents.run_state import RunState from agents.stream_events import AgentUpdatedStreamEvent, RawResponsesStreamEvent, StreamEvent from agents.testing import ModelStep, ScriptedModel -from agents.tool import CustomTool, FunctionTool -from agents.tool_guardrails import tool_input_guardrail, tool_output_guardrail +from agents.tool import FunctionTool +from agents.tool_guardrails import ( + ToolOutputGuardrailResult, + tool_input_guardrail, + tool_output_guardrail, +) from agents.usage import Usage, _attach_raw_usage_snapshot -from tests.model_test_helpers import get_exact_output_stream_step, get_response_obj +from tests.model_test_helpers import get_response_obj from .test_responses import ( get_final_output_message, @@ -78,36 +80,6 @@ ) from .utils.simple_session import CountingSession, SimpleListSession -_BLOCKED_TOOL_OUTPUT = "Output withheld by an output guardrail." - - -def _sdk_exception_traceback_string_locations( - error: BaseException, expected: str -) -> list[tuple[str, str]]: - pending = [error] - seen: set[int] = set() - locations: list[tuple[str, str]] = [] - while pending: - current = pending.pop() - if id(current) in seen: - continue - seen.add(id(current)) - - traceback = current.__traceback__ - while traceback is not None: - if "/src/agents/" in traceback.tb_frame.f_code.co_filename: - locations.extend( - (traceback.tb_frame.f_code.co_name, name) - for name, value in traceback.tb_frame.f_locals.items() - if isinstance(value, str) and value == expected - ) - traceback = traceback.tb_next - - for linked in (current.__cause__, current.__context__): - if linked is not None: - pending.append(linked) - return locations - def _conversation_locked_error() -> BadRequestError: request = httpx.Request("POST", "https://example.com") @@ -1903,7 +1875,7 @@ def guardrail_function( context: RunContextWrapper[Any], agent: Agent[Any], agent_output: Any ) -> GuardrailFunctionOutput: return GuardrailFunctionOutput( - output_info="safe-info", + output_info=None, tripwire_triggered=True, ) @@ -1915,16 +1887,11 @@ def guardrail_function( model=model, ) - with pytest.raises(OutputGuardrailTripwireTriggered) as exc_info: + with pytest.raises(OutputGuardrailTripwireTriggered): result = Runner.run_streamed(agent, input="user_message") async for _ in result.stream_events(): pass - assert exc_info.value.guardrail_result.agent_output == "first_test" - assert exc_info.value.guardrail_result.output.output_info == "safe-info" - assert result.output_guardrail_results[0].agent_output == "first_test" - assert result.output_guardrail_results[0].output.output_info == "safe-info" - @pytest.mark.asyncio async def test_output_guardrail_tripwire_raises_from_run_loop_task_before_stream_consumption(): @@ -2277,30 +2244,23 @@ async def test_tool() -> str: @pytest.mark.parametrize("mode", ["non_streamed", "streamed"]) @pytest.mark.parametrize("tripwire", [False, True], ids=["passes", "trips"]) @pytest.mark.asyncio -async def test_resumed_approved_tool_final_persists_output_after_output_guardrail_verdict( +async def test_resumed_approved_tool_final_persists_complete_post_verdict_batch( mode: str, tripwire: bool, ) -> None: guardrail_state = {"tripwire": tripwire} - def extract_custom_data(_context: Any) -> dict[str, str]: - return {"duplicate": "approved-result"} - - @function_tool( - name_override="approval_tool", - needs_approval=True, - custom_data_extractor=extract_custom_data, - ) + @function_tool(name_override="approval_tool", needs_approval=True) def approval_tool() -> str: return "approved-result" def output_guardrail( _context: RunContextWrapper[Any], _agent: Agent[Any], - output: Any, + _output: Any, ) -> GuardrailFunctionOutput: return GuardrailFunctionOutput( - output_info=output, + output_info=None, tripwire_triggered=guardrail_state["tripwire"], ) @@ -2328,10 +2288,8 @@ async def run_once(input_value: Any) -> Any: state.approve(first.interruptions[0]) if tripwire: - with pytest.raises(OutputGuardrailTripwireTriggered) as exc_info: + with pytest.raises(OutputGuardrailTripwireTriggered): await run_once(state) - assert exc_info.value.guardrail_result.agent_output == _BLOCKED_TOOL_OUTPUT - assert exc_info.value.guardrail_result.output.output_info is None else: resumed = await run_once(state) assert resumed.final_output == "approved-result" @@ -2350,17 +2308,11 @@ async def run_once(input_value: Any) -> Any: ("function_call", "call-approved"), ("function_call_output", "call-approved"), ] - expected_output = _BLOCKED_TOOL_OUTPUT if tripwire else "approved-result" + expected_output = ( + run_loop._OUTPUT_GUARDRAIL_BLOCKED_TOOL_OUTPUT if tripwire else "approved-result" + ) assert saved_tool_items[1].get("output") == expected_output - serialized_state = json.dumps(state.to_json()) - if tripwire: - assert "approved-result" not in serialized_state - assert _BLOCKED_TOOL_OUTPUT in serialized_state - assert state._current_step is None - else: - assert "approved-result" in serialized_state - if tripwire: guardrail_state["tripwire"] = False model.enqueue([get_text_message("done")]) @@ -2379,32 +2331,31 @@ async def run_once(input_value: Any) -> Any: ("function_call", "call-approved"), ("function_call_output", "call-approved"), ] - assert replayed_tool_items[1].get("output") == _BLOCKED_TOOL_OUTPUT + assert replayed_tool_items[1].get("output") == ( + run_loop._OUTPUT_GUARDRAIL_BLOCKED_TOOL_OUTPUT + ) + assert "approved-result" not in json.dumps(model_input) @pytest.mark.parametrize("mode", ["non_streamed", "streamed"]) @pytest.mark.asyncio -async def test_resumed_blocked_tool_redacts_live_state_before_failed_session_save( +async def test_ambiguous_serialized_approval_state_fails_before_tool_execution( mode: str, ) -> None: - class FailingResumedTurnSession(SimpleListSession): - fail_writes = False - - async def add_items(self, items: list[TResponseInputItem]) -> None: - if self.fail_writes: - raise LookupError("session save failed") - await super().add_items(items) + tool_calls = 0 @function_tool(name_override="approval_tool", needs_approval=True) def approval_tool() -> str: - return "state-secret" + nonlocal tool_calls + tool_calls += 1 + return "secret-result" def output_guardrail( _context: RunContextWrapper[Any], _agent: Agent[Any], - output: Any, + _output: Any, ) -> GuardrailFunctionOutput: - return GuardrailFunctionOutput(output_info=output, tripwire_triggered=True) + return GuardrailFunctionOutput(output_info=None, tripwire_triggered=False) model = ScriptedModel( [[get_function_tool_call("approval_tool", "{}", call_id="call-approved")]] @@ -2413,128 +2364,56 @@ def output_guardrail( name="test", model=model, tools=[approval_tool], - tool_use_behavior="stop_on_first_tool", output_guardrails=[OutputGuardrail(guardrail_function=output_guardrail)], ) - session = FailingResumedTurnSession() - first = await Runner.run(agent, "Use approval_tool", session=session) + first = await Runner.run(agent, "Use approval_tool") state = first.to_state() - state.approve(first.interruptions[0]) - session.fail_writes = True + state._current_turn = 2 + restored = await RunState.from_json(agent, state.to_json()) + restored.approve(restored.get_interruptions()[0]) - with pytest.raises(LookupError, match="session save failed"): + with pytest.raises(UserError, match="current response boundary cannot be proven"): if mode == "non_streamed": - await Runner.run(agent, state, session=session) + await Runner.run(agent, restored, session=None) else: - result = Runner.run_streamed(agent, state, session=session) + result = Runner.run_streamed(agent, restored, session=None) await consume_stream(result) - assert state._current_step is not None - assert getattr(state._current_step, "output", None) == _BLOCKED_TOOL_OUTPUT - assert "state-secret" not in json.dumps(state.to_json()) - - -@pytest.mark.asyncio -async def test_resumed_blocked_tool_session_save_cancellation_remains_observable() -> None: - class CancellingResumedTurnSession(SimpleListSession): - cancel_writes = False - - async def add_items(self, items: list[TResponseInputItem]) -> None: - if self.cancel_writes: - raise asyncio.CancelledError("session-secret") - await super().add_items(items) - - @function_tool(name_override="approval_tool", needs_approval=True) - def approval_tool() -> str: - return "state-secret" - - def output_guardrail( - _context: RunContextWrapper[Any], - _agent: Agent[Any], - output: Any, - ) -> GuardrailFunctionOutput: - return GuardrailFunctionOutput(output_info=output, tripwire_triggered=True) - - model = ScriptedModel( - [[get_function_tool_call("approval_tool", "{}", call_id="call-approved")]] - ) - agent = Agent( - name="test", - model=model, - tools=[approval_tool], - tool_use_behavior="stop_on_first_tool", - output_guardrails=[OutputGuardrail(guardrail_function=output_guardrail)], - ) - session = CancellingResumedTurnSession() - first = await Runner.run(agent, "Use approval_tool", session=session) - state = first.to_state() - state.approve(first.interruptions[0]) - session.cancel_writes = True - - result = Runner.run_streamed(agent, state, session=session) - with pytest.raises(asyncio.CancelledError) as exc_info: - await consume_stream(result) - - assert result._cancel_mode == "none" - assert result._stored_exception is exc_info.value - assert "session-secret" not in str(exc_info.value) - assert state._current_step is not None - assert getattr(state._current_step, "output", None) == _BLOCKED_TOOL_OUTPUT - serialized_state = json.dumps(state.to_json()) - assert "state-secret" not in serialized_state - assert "session-secret" not in serialized_state + assert tool_calls == 0 @pytest.mark.parametrize("mode", ["non_streamed", "streamed"]) @pytest.mark.asyncio -async def test_resumed_non_tool_tripwire_preserves_live_final_step(mode: str) -> None: - @function_tool(name_override="approval_tool", needs_approval=True) - def approval_tool() -> str: - return "approved" - - def output_guardrail( - _context: RunContextWrapper[Any], - _agent: Agent[Any], - output: Any, - ) -> GuardrailFunctionOutput: - return GuardrailFunctionOutput(output_info=output, tripwire_triggered=True) - - text_output = [get_text_message("blocked-text")] - second_step: Any = ( - text_output if mode == "non_streamed" else get_exact_output_stream_step(text_output) - ) - model = ScriptedModel( - [ - [get_function_tool_call("approval_tool", "{}", call_id="call-approved")], - second_step, - ] - ) +async def test_output_guardrails_fail_closed_with_server_managed_history(mode: str) -> None: + model = ScriptedModel([[get_text_message("unreachable")]]) agent = Agent( name="test", model=model, - tools=[approval_tool], - output_guardrails=[OutputGuardrail(guardrail_function=output_guardrail)], + output_guardrails=[ + OutputGuardrail( + guardrail_function=lambda _context, _agent, _output: GuardrailFunctionOutput( + output_info=None, + tripwire_triggered=False, + ) + ) + ], ) - first = await Runner.run(agent, "Use approval_tool") - state = first.to_state() - state.approve(first.interruptions[0]) - with pytest.raises(OutputGuardrailTripwireTriggered) as exc_info: + with pytest.raises(UserError, match="server-managed conversation history"): if mode == "non_streamed": - await Runner.run(agent, state) + await Runner.run(agent, "hello", previous_response_id="response-id") else: - result = Runner.run_streamed(agent, state) - await consume_stream(result) + Runner.run_streamed(agent, "hello", previous_response_id="response-id") - assert isinstance(state._current_step, NextStepRunAgain) - assert exc_info.value.guardrail_result.agent_output == "blocked-text" - assert exc_info.value.guardrail_result.output.output_info == "blocked-text" + assert not model.calls @pytest.mark.parametrize("mode", ["non_streamed", "streamed"]) +@pytest.mark.parametrize("session_kind", ["simple", "openai_conversations"]) @pytest.mark.asyncio async def test_stop_on_first_tool_final_persists_committed_tool_items_on_tripwire( mode: str, + session_kind: str, ) -> None: """A blocked final output must not discard the session record of a tool that already ran.""" @@ -2543,14 +2422,14 @@ async def test_stop_on_first_tool_final_persists_committed_tool_items_on_tripwir @function_tool(name_override="commit_tool") def commit_tool() -> str: calls.append("ran") - return "sensitive-result" + return "committed-result" def output_guardrail( _context: RunContextWrapper[Any], _agent: Agent[Any], - output: Any, + _output: Any, ) -> GuardrailFunctionOutput: - return GuardrailFunctionOutput(output_info=output, tripwire_triggered=True) + return GuardrailFunctionOutput(output_info=None, tripwire_triggered=True) model = ScriptedModel() model.enqueue([get_function_tool_call("commit_tool", "{}", call_id="call-committed")]) @@ -2561,31 +2440,43 @@ def output_guardrail( tool_use_behavior="stop_on_first_tool", output_guardrails=[OutputGuardrail(guardrail_function=output_guardrail)], ) - session = SimpleListSession() - streamed_result: Any = None - with pytest.raises(OutputGuardrailTripwireTriggered) as exc_info: + class DummyOpenAIConversationsSession(OpenAIConversationsSession): + def __init__(self) -> None: + self.history: list[TResponseInputItem] = [] + + async def get_items(self, limit: int | None = None) -> list[TResponseInputItem]: + return list(self.history if limit is None else self.history[-limit:]) + + async def add_items(self, items: list[TResponseInputItem]) -> None: + self.history.extend(items) + + async def pop_item(self) -> TResponseInputItem | None: + return self.history.pop() if self.history else None + + async def clear_session(self) -> None: + self.history.clear() + + session = SimpleListSession() if session_kind == "simple" else DummyOpenAIConversationsSession() + run_config = RunConfig( + session_input_callback=lambda history, new_input: [*reversed(history), *new_input] + ) + + with pytest.raises(OutputGuardrailTripwireTriggered): if mode == "non_streamed": - await Runner.run(agent, "Use commit_tool", session=session) + await Runner.run(agent, "Use commit_tool", session=session, run_config=run_config) else: - streamed_result = Runner.run_streamed(agent, "Use commit_tool", session=session) - await consume_stream(streamed_result) + result = Runner.run_streamed( + agent, + "Use commit_tool", + session=session, + run_config=run_config, + ) + await consume_stream(result) assert calls == ["ran"], "the tool never ran, so the test proves nothing" - assert exc_info.value.guardrail_result.agent_output == _BLOCKED_TOOL_OUTPUT - assert exc_info.value.guardrail_result.output.output_info is None - assert exc_info.value.__cause__ is None - assert exc_info.value.__context__ is None - assert _sdk_exception_traceback_string_locations(exc_info.value, "sensitive-result") == [] - if streamed_result is not None: - assert all( - result.agent_output == _BLOCKED_TOOL_OUTPUT - for result in streamed_result.output_guardrail_results - ) - assert "sensitive-result" not in json.dumps(streamed_result.to_state().to_json()) saved_items = await session.get_items() - assert "sensitive-result" not in json.dumps(saved_items) saved = [ (item.get("type") or item.get("role"), item.get("call_id")) for item in saved_items @@ -2596,15 +2487,28 @@ def output_guardrail( ("function_call", "call-committed"), ("function_call_output", "call-committed"), ] - assert cast(dict[str, Any], saved_items[-1]).get("output") == _BLOCKED_TOOL_OUTPUT + assert cast(dict[str, Any], saved_items[-1]).get("output") == ( + run_loop._OUTPUT_GUARDRAIL_BLOCKED_TOOL_OUTPUT + ) + assert "committed-result" not in json.dumps(saved_items) # The next run must see the completed call instead of re-issuing the same side effect. agent.output_guardrails = [] model.enqueue([get_text_message("done")]) if mode == "non_streamed": - followup: Any = await Runner.run(agent, "Continue", session=session) + followup: Any = await Runner.run( + agent, + "Continue", + session=session, + run_config=run_config, + ) else: - followup = Runner.run_streamed(agent, "Continue", session=session) + followup = Runner.run_streamed( + agent, + "Continue", + session=session, + run_config=run_config, + ) await consume_stream(followup) assert followup.final_output == "done" assert calls == ["ran"] @@ -2616,304 +2520,17 @@ def output_guardrail( for item in model_input if isinstance(item, dict) and item.get("type") in {"function_call", "function_call_output"} ] - assert replayed == [ + assert set(replayed) == { ("function_call", "call-committed"), ("function_call_output", "call-committed"), - ] + } replayed_output = next( item.get("output") for item in model_input if isinstance(item, dict) and item.get("type") == "function_call_output" ) - assert replayed_output == _BLOCKED_TOOL_OUTPUT - assert "sensitive-result" not in json.dumps(model_input) - - -@pytest.mark.parametrize("mode", ["non_streamed", "streamed"]) -@pytest.mark.asyncio -async def test_blocked_terminal_turn_sanitizes_concurrent_custom_tool_output(mode: str) -> None: - function_calls = 0 - custom_calls = 0 - - @function_tool(name_override="terminal_tool") - def terminal_tool() -> str: - nonlocal function_calls - function_calls += 1 - return "function-secret" - - def run_custom_tool(_context: Any, _input: str) -> str: - nonlocal custom_calls - custom_calls += 1 - return "custom-secret" - - custom_tool = CustomTool( - name="custom_side_effect", - description="Return a custom result.", - on_invoke_tool=run_custom_tool, - format={"type": "text"}, - ) - - def output_guardrail( - _context: RunContextWrapper[Any], - _agent: Agent[Any], - _output: Any, - ) -> GuardrailFunctionOutput: - return GuardrailFunctionOutput(output_info=None, tripwire_triggered=True) - - model_output = [ - get_function_tool_call("terminal_tool", "{}", call_id="call-function"), - ResponseCustomToolCall( - type="custom_tool_call", - name="custom_side_effect", - call_id="call-custom", - input="custom input", - ), - ] - model = ScriptedModel( - [model_output if mode == "non_streamed" else get_exact_output_stream_step(model_output)] - ) - session = SimpleListSession() - agent = Agent( - name="test", - model=model, - tools=[terminal_tool, custom_tool], - tool_use_behavior="stop_on_first_tool", - output_guardrails=[OutputGuardrail(guardrail_function=output_guardrail)], - ) - - with pytest.raises(OutputGuardrailTripwireTriggered): - if mode == "non_streamed": - await Runner.run(agent, "Run both tools", session=session) - else: - result = Runner.run_streamed(agent, "Run both tools", session=session) - await consume_stream(result) - - assert function_calls == 1 - assert custom_calls == 1 - saved_items = await session.get_items() - serialized_items = json.dumps(saved_items) - assert "function-secret" not in serialized_items - assert "custom-secret" not in serialized_items - - saved_outputs = { - cast(dict[str, Any], item).get("type"): cast(dict[str, Any], item).get("output") - for item in saved_items - if isinstance(item, dict) - and item.get("type") in {"function_call_output", "custom_tool_call_output"} - } - assert saved_outputs == { - "function_call_output": _BLOCKED_TOOL_OUTPUT, - "custom_tool_call_output": _BLOCKED_TOOL_OUTPUT, - } - - -@pytest.mark.parametrize("mode", ["non_streamed", "streamed"]) -@pytest.mark.parametrize("behavior", ["allow", "reject_content"]) -@pytest.mark.asyncio -async def test_blocked_terminal_tool_sanitizes_tool_output_guardrail_aliases( - mode: str, - behavior: str, -) -> None: - guardrail_outputs: list[ToolGuardrailFunctionOutput] = [] - - @tool_output_guardrail - def retain_tool_output(data: ToolOutputGuardrailData) -> ToolGuardrailFunctionOutput: - if behavior == "allow": - output = ToolGuardrailFunctionOutput.allow(output_info=data.output) - else: - output = ToolGuardrailFunctionOutput.reject_content( - message=data.output, - output_info=data.output, - ) - guardrail_outputs.append(output) - return output - - @function_tool( - name_override="terminal_tool", - tool_output_guardrails=[retain_tool_output], - ) - def terminal_tool() -> str: - return "tool-guardrail-secret" - - def output_guardrail( - _context: RunContextWrapper[Any], - _agent: Agent[Any], - output: Any, - ) -> GuardrailFunctionOutput: - return GuardrailFunctionOutput(output_info=output, tripwire_triggered=True) - - model_output = [get_function_tool_call("terminal_tool", "{}", call_id="call-terminal")] - model = ScriptedModel( - [model_output if mode == "non_streamed" else get_exact_output_stream_step(model_output)] - ) - agent = Agent( - name="test", - model=model, - tools=[terminal_tool], - tool_use_behavior="stop_on_first_tool", - output_guardrails=[OutputGuardrail(guardrail_function=output_guardrail)], - ) - streamed_result: Any = None - - with pytest.raises(OutputGuardrailTripwireTriggered): - if mode == "non_streamed": - await Runner.run(agent, "Use terminal_tool") - else: - streamed_result = Runner.run_streamed(agent, "Use terminal_tool") - await consume_stream(streamed_result) - - assert len(guardrail_outputs) == 1 - assert guardrail_outputs[0].output_info is None - if behavior == "allow": - assert guardrail_outputs[0].behavior == {"type": "allow"} - else: - assert guardrail_outputs[0].behavior == { - "type": "reject_content", - "message": _BLOCKED_TOOL_OUTPUT, - } - if streamed_result is not None: - assert "tool-guardrail-secret" not in json.dumps(streamed_result.to_state().to_json()) - - -@pytest.mark.parametrize("status", ["in_progress", "incomplete"]) -@pytest.mark.asyncio -async def test_blocked_streamed_terminal_turn_accepts_provider_shell_status( - status: str, -) -> None: - @function_tool(name_override="terminal_tool") - def terminal_tool() -> str: - return "terminal-secret" - - def output_guardrail( - _context: RunContextWrapper[Any], - _agent: Agent[Any], - output: Any, - ) -> GuardrailFunctionOutput: - return GuardrailFunctionOutput(output_info=output, tripwire_triggered=True) - - shell_output = ResponseFunctionShellToolCallOutput( - id=f"shell-output-{status}", - call_id="call-shell", - output=[ - { - "stdout": "shell-secret", - "stderr": "", - "outcome": {"type": "exit", "exit_code": 0}, - } - ], - status=cast(Any, status), - type="shell_call_output", - ) - model_output = [ - get_function_tool_call("terminal_tool", "{}", call_id="call-terminal"), - shell_output, - ] - model = ScriptedModel([get_exact_output_stream_step(model_output)]) - session = SimpleListSession() - agent = Agent( - name="test", - model=model, - tools=[terminal_tool], - tool_use_behavior="stop_on_first_tool", - output_guardrails=[OutputGuardrail(guardrail_function=output_guardrail)], - ) - - result = Runner.run_streamed(agent, "Run terminal_tool", session=session) - with pytest.raises(OutputGuardrailTripwireTriggered) as exc_info: - await consume_stream(result) - - assert exc_info.value.guardrail_result.agent_output == _BLOCKED_TOOL_OUTPUT - assert exc_info.value.guardrail_result.output.output_info is None - assert all( - guardrail_result.agent_output == _BLOCKED_TOOL_OUTPUT - and guardrail_result.output.output_info is None - for guardrail_result in result.output_guardrail_results - ) - serialized_state = json.dumps(result.to_state().to_json()) - serialized_session = json.dumps(await session.get_items()) - assert "terminal-secret" not in serialized_state - assert "shell-secret" not in serialized_state - assert "terminal-secret" not in serialized_session - assert "shell-secret" not in serialized_session - assert _BLOCKED_TOOL_OUTPUT in serialized_state - assert _BLOCKED_TOOL_OUTPUT in serialized_session - - -@pytest.mark.parametrize("mode", ["non_streamed", "streamed"]) -@pytest.mark.asyncio -async def test_blocked_tool_output_cannot_be_forwarded_by_a_later_tool(mode: str) -> None: - secret_calls = 0 - forwarded_values: list[str] = [] - guardrail_outputs: list[Any] = [] - - @function_tool(name_override="secret_tool") - def secret_tool() -> str: - nonlocal secret_calls - secret_calls += 1 - return "private-value" - - @function_tool(name_override="record_replay") - def record_replay(value: str) -> str: - forwarded_values.append(value) - return "recorded" - - def output_guardrail( - _context: RunContextWrapper[Any], - _agent: Agent[Any], - output: Any, - ) -> GuardrailFunctionOutput: - guardrail_outputs.append(output) - return GuardrailFunctionOutput( - output_info=None, - tripwire_triggered=output == "private-value", - ) - - def forward_replayed_output(call: Any) -> list[Any]: - assert isinstance(call.input, list) - replayed_output = next( - item.get("output") - for item in call.input - if isinstance(item, dict) and item.get("type") == "function_call_output" - ) - return [ - get_function_tool_call( - "record_replay", - json.dumps({"value": replayed_output}), - call_id="call-record", - ) - ] - - model = ScriptedModel( - [ - [get_function_tool_call("secret_tool", "{}", call_id="call-secret")], - ModelStep.respond(forward_replayed_output), - ] - ) - session = SimpleListSession() - agent = Agent( - name="test", - model=model, - tools=[secret_tool, record_replay], - tool_use_behavior="stop_on_first_tool", - output_guardrails=[OutputGuardrail(guardrail_function=output_guardrail)], - ) - - async def run_once(input_value: str) -> 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 - - with pytest.raises(OutputGuardrailTripwireTriggered): - await run_once("Read the secret") - - followup = await run_once("Forward the previous result") - assert followup.final_output == "recorded" - assert secret_calls == 1 - assert forwarded_values == [_BLOCKED_TOOL_OUTPUT] - assert guardrail_outputs == ["private-value", "recorded"] - assert "private-value" not in json.dumps(await session.get_items()) + assert replayed_output == run_loop._OUTPUT_GUARDRAIL_BLOCKED_TOOL_OUTPUT + assert "committed-result" not in json.dumps(model_input) @pytest.mark.parametrize("mode", ["non_streamed", "streamed"]) @@ -2963,7 +2580,10 @@ def output_guardrail( _agent: Agent[Any], _output: Any, ) -> GuardrailFunctionOutput: - return GuardrailFunctionOutput(output_info=None, tripwire_triggered=True) + return GuardrailFunctionOutput( + output_info={"reason": "message rejected"}, + tripwire_triggered=True, + ) model = ScriptedModel() model.extend( @@ -2980,13 +2600,16 @@ def output_guardrail( ) session = SimpleListSession() - with pytest.raises(OutputGuardrailTripwireTriggered): + with pytest.raises(OutputGuardrailTripwireTriggered) as exc_info: if mode == "non_streamed": await Runner.run(agent, "Use commit_tool", session=session) else: result = Runner.run_streamed(agent, "Use commit_tool", session=session) await consume_stream(result) + assert exc_info.value.guardrail_result.agent_output == "should_not_be_saved" + assert exc_info.value.guardrail_result.output.output_info == {"reason": "message rejected"} + saved_items = await session.get_items() saved = [item.get("type") or item.get("role") for item in saved_items if isinstance(item, dict)] assert saved == ["user", "function_call", "function_call_output"] @@ -3059,15 +2682,10 @@ async def run_once() -> Any: @pytest.mark.parametrize("mode", ["non_streamed", "streamed"]) @pytest.mark.asyncio -async def test_failing_output_guardrail_keeps_the_whole_final_turn( +async def test_failing_output_guardrail_does_not_persist_the_unverdictable_turn( mode: str, ) -> None: - """A guardrail *error* is not a tripwire: the completed final turn stays replayable. - - Only a tripwire means the output was judged undeliverable. An ordinary guardrail exception - leaves the verdict unknown, so the turn must be persisted whole, exactly as the non-streamed - path does. - """ + """A guardrail error leaves no verdict, so its response is not persisted.""" @function_tool(name_override="commit_tool") def commit_tool() -> str: @@ -3108,235 +2726,17 @@ async def run_once() -> None: saved_items = await session.get_items() saved = [item.get("type") or item.get("role") for item in saved_items if isinstance(item, dict)] - assert saved == ["user", "message", "function_call", "function_call_output"] - - -@pytest.mark.asyncio -async def test_streamed_session_save_error_takes_precedence_over_output_guardrail_error() -> None: - guardrail_failed = False - final_turn_save_attempted = False - - class FailingFinalTurnSession(SimpleListSession): - async def add_items(self, items: list[TResponseInputItem]) -> None: - nonlocal final_turn_save_attempted - if guardrail_failed: - final_turn_save_attempted = True - raise LookupError("session save failed") - await super().add_items(items) - - def output_guardrail( - _context: RunContextWrapper[Any], - _agent: Agent[Any], - _output: Any, - ) -> GuardrailFunctionOutput: - nonlocal guardrail_failed - guardrail_failed = True - raise RuntimeError("guardrail failed") - - model = ScriptedModel() - model.enqueue([get_text_message("assistant-preamble")]) - agent = Agent( - name="test", - model=model, - output_guardrails=[OutputGuardrail(guardrail_function=output_guardrail)], - ) - session = FailingFinalTurnSession() - result = Runner.run_streamed(agent, "Hello", session=session) - - with pytest.raises(LookupError, match="session save failed") as exc_info: - await consume_stream(result) - - assert final_turn_save_attempted is True - assert isinstance(exc_info.value.__context__, RuntimeError) - assert str(exc_info.value.__context__) == "guardrail failed" - assert result.run_loop_exception is exc_info.value - assert await session.get_items() == [{"content": "Hello", "role": "user"}] - - -@pytest.mark.asyncio -async def test_blocked_tool_output_redacts_live_state_before_failed_session_save() -> None: - guardrail_tripped = False - - class FailingBlockedTurnSession(SimpleListSession): - async def add_items(self, items: list[TResponseInputItem]) -> None: - if guardrail_tripped: - raise LookupError("session save failed") - await super().add_items(items) - - @function_tool(name_override="commit_tool") - def commit_tool() -> str: - return "state-secret" - - def output_guardrail( - _context: RunContextWrapper[Any], - _agent: Agent[Any], - output: Any, - ) -> GuardrailFunctionOutput: - nonlocal guardrail_tripped - guardrail_tripped = True - return GuardrailFunctionOutput(output_info=output, tripwire_triggered=True) - - model = ScriptedModel([[get_function_tool_call("commit_tool", "{}", call_id="call-committed")]]) - agent = Agent( - name="test", - model=model, - tools=[commit_tool], - tool_use_behavior="stop_on_first_tool", - output_guardrails=[OutputGuardrail(guardrail_function=output_guardrail)], - ) - session = FailingBlockedTurnSession() - result = Runner.run_streamed(agent, "Use commit_tool", session=session) - - with pytest.raises(LookupError, match="session save failed"): - await consume_stream(result) - - assert result._state is not None - assert result._state._current_step is None - assert "state-secret" not in json.dumps(result.to_state().to_json()) - - -@pytest.mark.asyncio -async def test_streamed_session_save_cancellation_is_not_a_public_immediate_cancel() -> None: - guardrail_failed = False - - class CancellingFinalTurnSession(SimpleListSession): - async def add_items(self, items: list[TResponseInputItem]) -> None: - if guardrail_failed: - raise asyncio.CancelledError("session save cancelled") - await super().add_items(items) - - def output_guardrail( - _context: RunContextWrapper[Any], - _agent: Agent[Any], - _output: Any, - ) -> GuardrailFunctionOutput: - nonlocal guardrail_failed - guardrail_failed = True - raise RuntimeError("guardrail failed") - - model = ScriptedModel(steps=[[get_text_message("assistant-preamble")]]) - agent = Agent( - name="test", - model=model, - output_guardrails=[OutputGuardrail(guardrail_function=output_guardrail)], - ) - session = CancellingFinalTurnSession() - result = Runner.run_streamed(agent, "Hello", session=session) - - with pytest.raises(asyncio.CancelledError, match="session save cancelled") as exc_info: - await consume_stream(result) - - assert result._cancel_mode == "none" - assert result._stored_exception is exc_info.value - assert isinstance(exc_info.value.__context__, RuntimeError) - assert str(exc_info.value.__context__) == "guardrail failed" - assert await session.get_items() == [{"content": "Hello", "role": "user"}] - - -@pytest.mark.asyncio -async def test_streamed_session_save_direct_base_exception_is_terminal() -> None: - guardrail_failed = False - - class DirectAbort(BaseException): - pass - - class AbortingFinalTurnSession(SimpleListSession): - async def add_items(self, items: list[TResponseInputItem]) -> None: - if guardrail_failed: - raise DirectAbort("session save aborted") - await super().add_items(items) - - def output_guardrail( - _context: RunContextWrapper[Any], - _agent: Agent[Any], - _output: Any, - ) -> GuardrailFunctionOutput: - nonlocal guardrail_failed - guardrail_failed = True - raise RuntimeError("guardrail failed") - - agent = Agent( - name="test", - model=ScriptedModel(steps=[[get_text_message("assistant-preamble")]]), - output_guardrails=[OutputGuardrail(guardrail_function=output_guardrail)], - ) - session = AbortingFinalTurnSession() - result = Runner.run_streamed(agent, "Hello", session=session) - - with pytest.raises(DirectAbort, match="session save aborted") as exc_info: - await consume_stream(result) - - assert result._stored_exception is exc_info.value - assert result.run_loop_exception is exc_info.value - assert await session.get_items() == [{"content": "Hello", "role": "user"}] - - -@pytest.mark.asyncio -async def test_public_immediate_cancel_during_guardrail_recovery_save_stays_prompt() -> None: - guardrail_failed = False - save_started = asyncio.Event() - save_cancelled = asyncio.Event() - never_set = asyncio.Event() - - class BlockingFinalTurnSession(SimpleListSession): - async def add_items(self, items: list[TResponseInputItem]) -> None: - if guardrail_failed: - save_started.set() - try: - await never_set.wait() - finally: - save_cancelled.set() - return - await super().add_items(items) - - def output_guardrail( - _context: RunContextWrapper[Any], - _agent: Agent[Any], - _output: Any, - ) -> GuardrailFunctionOutput: - nonlocal guardrail_failed - guardrail_failed = True - raise RuntimeError("guardrail failed") - - model = ScriptedModel(steps=[[get_text_message("assistant-preamble")]]) - agent = Agent( - name="test", - model=model, - output_guardrails=[OutputGuardrail(guardrail_function=output_guardrail)], - ) - session = BlockingFinalTurnSession() - result = Runner.run_streamed(agent, "Hello", session=session) - drain_task = asyncio.create_task(consume_stream(result)) - - try: - await asyncio.wait_for(save_started.wait(), timeout=1) - result.cancel() - await asyncio.wait_for(drain_task, timeout=1) - finally: - if not drain_task.done(): - result.cancel() - drain_task.cancel() - with contextlib.suppress(asyncio.CancelledError): - await drain_task - - assert save_cancelled.is_set() - assert result._stored_exception is None - assert await session.get_items() == [{"content": "Hello", "role": "user"}] + assert saved == ["user"] @pytest.mark.parametrize("mode", ["non_streamed", "streamed"]) @pytest.mark.parametrize("tripwire", [False, True], ids=["passes", "trips"]) @pytest.mark.asyncio -async def test_blocked_tool_final_keeps_reasoning_context_with_the_committed_call( +async def test_blocked_tool_final_discards_reasoning_response_suffix_on_trip( mode: str, tripwire: bool, ) -> None: - """A retained tool call keeps the reasoning item it belongs to, in order. - - A reasoning model requires the reasoning item that preceded a function call to accompany that - call in the next request, so persisting the call/output pair without it leaves an unreplayable - turn. Asserted on both the session contents and the next run's model input. - """ + """A reasoning-bearing response is preserved on pass and discarded completely on trip.""" @function_tool(name_override="commit_tool") def commit_tool() -> str: @@ -3384,7 +2784,10 @@ async def run_once(input_value: Any) -> Any: saved_items = await session.get_items() saved = [item.get("type") or item.get("role") for item in saved_items if isinstance(item, dict)] - assert saved == ["user", "reasoning", "function_call", "function_call_output"] + expected_saved = ( + ["user"] if tripwire else ["user", "reasoning", "function_call", "function_call_output"] + ) + assert saved == expected_saved # The reasoning/call/output group has to reach the next request in that order. agent.output_guardrails = [] @@ -3400,22 +2803,16 @@ async def run_once(input_value: Any) -> Any: if isinstance(item, dict) and item.get("type") in {"reasoning", "function_call", "function_call_output"} ] - assert replayed == ["reasoning", "function_call", "function_call_output"] + expected_replayed = [] if tripwire else ["reasoning", "function_call", "function_call_output"] + assert replayed == expected_replayed @pytest.mark.parametrize("mode", ["non_streamed", "streamed"]) @pytest.mark.asyncio -async def test_blocked_tool_final_drops_reasoning_tied_to_the_rejected_message( +async def test_blocked_tool_final_discards_suffix_with_multiple_reasoning_groups( mode: str, ) -> None: - """Only the reasoning tied to a retained call survives; the message's reasoning goes with it. - - The turn is `reasoning_for_message -> message -> reasoning_for_call -> function_call`. A - reasoning item belongs to the next non-reasoning item, so retaining every reasoning item - whenever the turn happens to contain a tool call would leave the rejected message's reasoning - dangling in the next request. - - """ + """Any reasoning item makes the complete rejected current-response suffix unsupported.""" @function_tool(name_override="commit_tool") def commit_tool() -> str: @@ -3466,17 +2863,14 @@ async def run_once(input_value: Any) -> Any: saved_items = await session.get_items() saved = [item.get("type") or item.get("role") for item in saved_items if isinstance(item, dict)] - assert saved == ["user", "reasoning", "function_call", "function_call_output"] + assert saved == ["user"] saved_reasoning_ids = [ item.get("id") for item in saved_items if isinstance(item, dict) and item.get("id") ] - assert "rs_committed" in saved_reasoning_ids - assert "rs_rejected" not in saved_reasoning_ids, ( - "reasoning tied to the rejected message must not be persisted" - ) + assert saved_reasoning_ids == [] - # ...and the surviving group still replays in order, with no dangling reasoning item. + # The unsupported response contributes nothing to the next model request. agent.output_guardrails = [] model.enqueue([get_text_message("done")]) followup = await run_once("Continue") @@ -3490,7 +2884,94 @@ async def run_once(input_value: Any) -> Any: if isinstance(item, dict) and item.get("type") in {"reasoning", "message", "function_call", "function_call_output"} ] - assert replayed == ["reasoning", "function_call", "function_call_output"] + assert replayed == [] + + +@pytest.mark.parametrize("reasoning_suffix", [False, True], ids=["canonical", "reasoning"]) +@pytest.mark.asyncio +async def test_streamed_trip_preserves_accepted_tool_prefix( + reasoning_suffix: bool, +) -> None: + """Only the rejected current response is replaced or dropped from replay owners.""" + side_effects: list[str] = [] + + @function_tool(name_override="accepted_tool") + def accepted_tool() -> str: + side_effects.append("accepted") + return "accepted-output" + + @function_tool(name_override="terminal_tool") + def terminal_tool() -> str: + side_effects.append("terminal") + return "rejected-output" + + def reject_output( + _context: RunContextWrapper[Any], + _agent: Agent[Any], + _output: Any, + ) -> GuardrailFunctionOutput: + return GuardrailFunctionOutput(output_info=None, tripwire_triggered=True) + + terminal_response: list[Any] = [] + if reasoning_suffix: + terminal_response.append( + ResponseReasoningItem( + id="reasoning-current", + summary=[Summary(text="calling terminal tool", type="summary_text")], + type="reasoning", + ) + ) + terminal_response.append(get_function_tool_call("terminal_tool", "{}", call_id="current-call")) + model = ScriptedModel( + steps=[ + [get_function_tool_call("accepted_tool", "{}", call_id="accepted-call")], + terminal_response, + ] + ) + agent = Agent( + name="test", + model=model, + tools=[accepted_tool, terminal_tool], + tool_use_behavior={"stop_at_tool_names": ["terminal_tool"]}, + output_guardrails=[OutputGuardrail(guardrail_function=reject_output)], + ) + + result = Runner.run_streamed(agent, "run both tools") + with pytest.raises(OutputGuardrailTripwireTriggered): + await consume_stream(result) + + assert side_effects == ["accepted", "terminal"] + + def call_ids(items: list[RunItem]) -> list[str]: + return [ + call_id + for item in items + if ( + call_id := ( + item.raw_item.get("call_id") + if isinstance(item.raw_item, dict) + else getattr(item.raw_item, "call_id", None) + ) + ) + is not None + ] + + expected_call_ids = ["accepted-call", "accepted-call"] + if not reasoning_suffix: + expected_call_ids.extend(["current-call", "current-call"]) + assert call_ids(result.new_items) == expected_call_ids + assert call_ids(result._model_input_items) == expected_call_ids + + state = result.to_state() + assert call_ids(state._generated_items) == expected_call_ids + assert call_ids(state._session_items) == expected_call_ids + serialized_state = json.dumps(state.to_json()) + assert "accepted-output" in serialized_state + assert "rejected-output" not in serialized_state + if reasoning_suffix: + assert "reasoning-current" not in serialized_state + else: + assert run_loop._OUTPUT_GUARDRAIL_BLOCKED_TOOL_OUTPUT in serialized_state @pytest.mark.asyncio @@ -3818,6 +3299,61 @@ async def test_streamed_run_reports_tool_guardrail_results(): assert result.tool_output_guardrail_results[0].output.output_info == "output-checked" +@pytest.mark.asyncio +async def test_streamed_trip_replaces_current_tool_output_guardrail_results() -> None: + """A copied terminal tool result is replaced in public and RunState guardrail results.""" + original_outputs: list[ToolGuardrailFunctionOutput] = [] + + @tool_output_guardrail + def retain_output(data: ToolOutputGuardrailData) -> ToolGuardrailFunctionOutput: + output = ToolGuardrailFunctionOutput.allow(output_info=data.output) + original_outputs.append(output) + return output + + @function_tool(name_override="secret_tool", tool_output_guardrails=[retain_output]) + def secret_tool() -> str: + return "blocked-secret" + + def reject_output( + _context: RunContextWrapper[Any], + _agent: Agent[Any], + _output: Any, + ) -> GuardrailFunctionOutput: + return GuardrailFunctionOutput(output_info=None, tripwire_triggered=True) + + model = ScriptedModel() + model.enqueue([get_function_tool_call("secret_tool", "{}", call_id="call-secret")]) + agent = Agent( + name="test", + model=model, + tools=[secret_tool], + tool_use_behavior="stop_on_first_tool", + output_guardrails=[OutputGuardrail(guardrail_function=reject_output)], + ) + + result = Runner.run_streamed(agent, "run") + prior_output = ToolGuardrailFunctionOutput.allow(output_info="prior-safe") + prior_result = ToolOutputGuardrailResult(guardrail=retain_output, output=prior_output) + result.tool_output_guardrail_results.append(prior_result) + with pytest.raises(OutputGuardrailTripwireTriggered): + await consume_stream(result) + + assert len(original_outputs) == 1 + assert len(result.tool_output_guardrail_results) == 2 + assert result.tool_output_guardrail_results[0] is prior_result + assert result.tool_output_guardrail_results[0].output is prior_output + public_output = result.tool_output_guardrail_results[1].output + assert public_output is not original_outputs[0] + assert public_output.output_info == run_loop._OUTPUT_GUARDRAIL_BLOCKED_TOOL_OUTPUT + assert result._state is not None + # The caller-added public result was never owned by RunState, so only the current + # data-free result is added to that owner. + assert len(result._state._tool_output_guardrail_results) == 1 + state_output = result._state._tool_output_guardrail_results[0].output + assert state_output is public_output + assert state_output.output_info == run_loop._OUTPUT_GUARDRAIL_BLOCKED_TOOL_OUTPUT + + @pytest.mark.asyncio async def test_streamed_tool_guardrail_results_match_non_streamed(): """The same run reports the same tool guardrail results in both execution modes.""" diff --git a/tests/test_error_logging_redaction.py b/tests/test_error_logging_redaction.py index 6c26da31f0..ff8722aac3 100644 --- a/tests/test_error_logging_redaction.py +++ b/tests/test_error_logging_redaction.py @@ -1532,28 +1532,24 @@ def output_guardrail( @pytest.mark.parametrize("redacted", [False, True]) -@pytest.mark.parametrize("persistence_failure", ["error", "cancelled"]) @pytest.mark.asyncio -async def test_streamed_session_error_after_output_guardrail_respects_redaction( +async def test_streamed_output_guardrail_error_skips_final_session_write( monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture, redacted: bool, - persistence_failure: Literal["error", "cancelled"], ) -> None: monkeypatch.setattr(_debug, "DONT_LOG_MODEL_DATA", redacted) monkeypatch.setattr(_debug, "DONT_LOG_TOOL_DATA", False) guardrail_failed = False + post_guardrail_save_calls = 0 payload = f'{{"answer": "{_MODEL_OUTPUT_SECRET}"}}' class FailingFinalTurnSession(SimpleListSession): async def add_items(self, items: list[Any]) -> None: + nonlocal post_guardrail_save_calls if guardrail_failed: - cause = RuntimeError(f"session save cause: {_MODEL_OUTPUT_SECRET}") - if persistence_failure == "cancelled": - raise asyncio.CancelledError( - f"session save cancelled: {_MODEL_OUTPUT_SECRET}" - ) from cause - raise LookupError(f"session save failed: {_MODEL_OUTPUT_SECRET}") from cause + post_guardrail_save_calls += 1 + raise AssertionError("the final Session write must not start without a verdict") await super().add_items(items) def output_guardrail( @@ -1574,86 +1570,56 @@ def output_guardrail( output_guardrails=[OutputGuardrail(guardrail_function=output_guardrail)], ) result = Runner.run_streamed(agent, "go", session=FailingFinalTurnSession()) - run_loop_callback_errors: list[BaseException] = [] - run_loop_done = asyncio.Event() - if redacted and persistence_failure == "cancelled": - assert result.run_loop_task is not None - - def inspect_run_loop_task(task: asyncio.Task[Any]) -> None: - try: - task.result() - except BaseException as error: - run_loop_callback_errors.append(error) - finally: - run_loop_done.set() - - result.run_loop_task.add_done_callback(inspect_run_loop_task) - expected_error_type = ( - asyncio.CancelledError - if persistence_failure == "cancelled" - else UserError - if redacted - else LookupError - ) - expected_message = ( - "Error details are redacted." - if redacted - else "session save cancelled" - if persistence_failure == "cancelled" - else "session save failed" - ) - - with pytest.raises(expected_error_type, match=expected_message) as exc_info: + with pytest.raises(ModelBehaviorError) as exc_info: async for _ in result.stream_events(): pass error = exc_info.value - guardrail_error = error.__context__ + assert post_guardrail_save_calls == 0 + assert result.run_loop_exception is error + assert result._stored_exception is error if redacted: - assert guardrail_error is None + assert error.run_data is None assert error.__cause__ is None + assert error.__context__ is None assert _MODEL_OUTPUT_SECRET not in str(error) - _assert_secret_absent_from_agents_traceback(error, _MODEL_OUTPUT_SECRET) + _assert_secret_absent_from_agents_traceback( + error, + _MODEL_OUTPUT_SECRET, + require_agents_frames=False, + ) for record in caplog.records: assert _MODEL_OUTPUT_SECRET not in repr(record.__dict__) assert _MODEL_OUTPUT_SECRET not in logging.Formatter().format(record) assert record.exc_info is None else: - assert isinstance(guardrail_error, ModelBehaviorError) - assert error.__cause__ is not None - assert _MODEL_OUTPUT_SECRET in str(error.__cause__) - assert _MODEL_OUTPUT_SECRET in str(guardrail_error) + assert _MODEL_OUTPUT_SECRET in str(error) + assert isinstance(error.__cause__, ValidationError) assert any( _MODEL_OUTPUT_SECRET in repr(frame) for frame in _agents_traceback_frame_locals(error) ) - - if persistence_failure == "cancelled": - assert result._stored_exception is error - assert result.run_loop_exception is None - if redacted: - await asyncio.wait_for(run_loop_done.wait(), timeout=1) - assert run_loop_callback_errors == [] - else: - assert result.run_loop_exception is error - if redacted: - assert error.__traceback__ is None + if redacted: + assert error.__traceback__ is None @pytest.mark.asyncio -async def test_streamed_session_hostile_error_after_redacted_output_guardrail_is_replaced( +async def test_streamed_redacted_output_guardrail_does_not_invoke_hostile_session( monkeypatch: pytest.MonkeyPatch, ) -> None: persistence_secret = "HOSTILE_SESSION_FAILURE_SECRET" monkeypatch.setattr(_debug, "DONT_LOG_MODEL_DATA", True) monkeypatch.setattr(_debug, "DONT_LOG_TOOL_DATA", False) guardrail_failed = False + post_guardrail_save_calls = 0 payload = f'{{"answer": "{_MODEL_OUTPUT_SECRET}"}}' class FailingFinalTurnSession(SimpleListSession): async def add_items(self, items: list[Any]) -> None: + nonlocal post_guardrail_save_calls if guardrail_failed: + post_guardrail_save_calls += 1 raise _HostileAttributeWriteException(persistence_secret) await super().add_items(items) @@ -1674,11 +1640,13 @@ def output_guardrail( ) result = Runner.run_streamed(agent, "go", session=FailingFinalTurnSession()) - with pytest.raises(UserError, match="Error details are redacted.") as exc_info: + with pytest.raises(ModelBehaviorError) as exc_info: async for _ in result.stream_events(): pass error = exc_info.value + assert post_guardrail_save_calls == 0 + assert error.run_data is None assert error.__cause__ is None assert error.__context__ is None assert persistence_secret not in str(error) @@ -1846,25 +1814,26 @@ def inspect_public_task(task: asyncio.Task[Any]) -> None: @pytest.mark.parametrize("streamed", [False, True]) -@pytest.mark.parametrize("failure_kind", ["exception", "cancelled", "direct_base", "group"]) @pytest.mark.asyncio -async def test_max_turns_recovery_session_failure_preserves_complete_redaction_boundary( +async def test_max_turns_guardrail_error_skips_final_session_write( monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture, streamed: bool, - failure_kind: Literal["exception", "cancelled", "direct_base", "group"], ) -> None: persistence_secret = "MAX_TURNS_SESSION_FAILURE_SECRET" fallback_secret = "MAX_TURNS_HANDLER_OUTPUT_SECRET" payload = f'{{"answer": "{_MODEL_OUTPUT_SECRET}"}}' guardrail_failed = False + post_guardrail_save_calls = 0 monkeypatch.setattr(_debug, "DONT_LOG_MODEL_DATA", True) monkeypatch.setattr(_debug, "DONT_LOG_TOOL_DATA", False) class FailingMaxTurnsSession(SimpleListSession): async def add_items(self, items: list[Any]) -> None: + nonlocal post_guardrail_save_calls if guardrail_failed: - raise _persistence_failure(failure_kind, persistence_secret) + post_guardrail_save_calls += 1 + raise AssertionError("the final Session write must not start without a verdict") await super().add_items(items) def output_guardrail( @@ -1908,23 +1877,12 @@ def output_guardrail( except BaseException as error: captured_error = error else: # pragma: no cover - raise AssertionError("the session failure must propagate") + raise AssertionError("the guardrail failure must propagate") - assert captured_error is not None + assert isinstance(captured_error, ModelBehaviorError) + assert post_guardrail_save_calls == 0 error = captured_error - if failure_kind == "exception": - assert isinstance(error, UserError) - elif failure_kind == "cancelled": - assert isinstance(error, asyncio.CancelledError) - elif failure_kind == "direct_base": - assert type(error) is BaseException - else: - assert isinstance(error, BaseExceptionGroup) - assert not isinstance(error, Exception) - assert {type(child) for child in error.exceptions} == { - UserError, - asyncio.CancelledError, - } + assert error.run_data is None error_graph = _exception_graph(error) assert error_graph @@ -2064,7 +2022,7 @@ def test_agent_runner_run_sync_detaches_all_marked_recovery_failures( @pytest.mark.parametrize("streamed", [False, True]) @pytest.mark.asyncio -async def test_max_turns_recovery_deep_exception_group_preserves_redaction_boundary( +async def test_max_turns_guardrail_error_does_not_build_session_exception_group( monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture, streamed: bool, @@ -2072,12 +2030,15 @@ async def test_max_turns_recovery_deep_exception_group_preserves_redaction_bound persistence_secret = "DEEP_MAX_TURNS_SESSION_FAILURE_SECRET" payload = f'{{"answer": "{_MODEL_OUTPUT_SECRET}"}}' guardrail_failed = False + post_guardrail_save_calls = 0 monkeypatch.setattr(_debug, "DONT_LOG_MODEL_DATA", True) monkeypatch.setattr(_debug, "DONT_LOG_TOOL_DATA", False) class DeepGroupSession(SimpleListSession): async def add_items(self, items: list[Any]) -> None: + nonlocal post_guardrail_save_calls if guardrail_failed: + post_guardrail_save_calls += 1 error: BaseException = asyncio.CancelledError(persistence_secret) for _ in range(1200): error = BaseExceptionGroup(persistence_secret, [error]) @@ -2125,23 +2086,14 @@ def output_guardrail( except BaseException as error: captured_error = error else: # pragma: no cover - raise AssertionError("the deep exception group must propagate") + raise AssertionError("the guardrail failure must propagate") - assert isinstance(captured_error, BaseExceptionGroup) - error_graph = _exception_graph(captured_error) - assert len(error_graph) == 1201 - for current in error_graph: - assert current.__cause__ is None - assert current.__context__ is None - if isinstance(current, BaseExceptionGroup): - assert current.message == "Error details are redacted." - else: - assert isinstance(current, asyncio.CancelledError) - assert persistence_secret not in str(current) - assert persistence_secret not in repr(current) - for frame_locals in _agents_traceback_frame_locals(current): - _assert_secret_absent_from_value_graph(frame_locals, persistence_secret) - _assert_secret_absent_from_value_graph(frame_locals, _MODEL_OUTPUT_SECRET) + assert isinstance(captured_error, ModelBehaviorError) + assert captured_error.run_data is None + assert post_guardrail_save_calls == 0 + assert captured_error.__cause__ is None + assert captured_error.__context__ is None + _assert_secret_absent_from_agents_traceback(captured_error, _MODEL_OUTPUT_SECRET) for record in caplog.records: assert persistence_secret not in logging.Formatter().format(record) @@ -2152,24 +2104,25 @@ def output_guardrail( @pytest.mark.parametrize("redacted", [False, True]) -@pytest.mark.parametrize("failure_kind", ["direct_base", "group"]) @pytest.mark.asyncio -async def test_max_turns_run_loop_exception_follows_redaction_policy_for_base_exceptions( +async def test_max_turns_run_loop_guardrail_error_skips_final_session_write( monkeypatch: pytest.MonkeyPatch, redacted: bool, - failure_kind: Literal["direct_base", "group"], ) -> None: persistence_secret = "RUN_LOOP_EXCEPTION_PERSISTENCE_SECRET" fallback_secret = "RUN_LOOP_EXCEPTION_FALLBACK_SECRET" payload = f'{{"answer": "{_MODEL_OUTPUT_SECRET}"}}' guardrail_failed = False + post_guardrail_save_calls = 0 monkeypatch.setattr(_debug, "DONT_LOG_MODEL_DATA", redacted) monkeypatch.setattr(_debug, "DONT_LOG_TOOL_DATA", False) class FailingMaxTurnsSession(SimpleListSession): async def add_items(self, items: list[Any]) -> None: + nonlocal post_guardrail_save_calls if guardrail_failed: - raise _persistence_failure(failure_kind, persistence_secret) + post_guardrail_save_calls += 1 + raise _DirectBaseException(persistence_secret) await super().add_items(items) def output_guardrail( @@ -2209,7 +2162,8 @@ def inspect_run_loop_task(task: asyncio.Task[Any]) -> None: await asyncio.wait_for(run_loop_done.wait(), timeout=1) error = result.run_loop_exception - assert error is not None + assert isinstance(error, ModelBehaviorError) + assert post_guardrail_save_calls == 0 frame_locals = _agents_traceback_frame_locals(error) if redacted: for traceback_locals in callback_frame_locals + frame_locals: @@ -2229,7 +2183,7 @@ def inspect_run_loop_task(task: asyncio.Task[Any]) -> None: assert any(fallback_secret in repr(frame) for frame in callback_frame_locals) assert frame_locals assert any(fallback_secret in repr(frame) for frame in frame_locals) - assert persistence_secret in str(error) + assert _MODEL_OUTPUT_SECRET in str(error) try: async for _ in result.stream_events(): @@ -2237,7 +2191,7 @@ def inspect_run_loop_task(task: asyncio.Task[Any]) -> None: except BaseException as streamed_error: assert streamed_error is error else: # pragma: no cover - raise AssertionError("the session failure must propagate through the stream") + raise AssertionError("the guardrail failure must propagate through the stream") @pytest.mark.parametrize( @@ -2371,18 +2325,21 @@ def test_safe_redacted_persistence_error_preserves_hybrid_cancellation() -> None @pytest.mark.parametrize("streamed", [False, True]) @pytest.mark.asyncio -async def test_max_turns_recovery_session_failure_preserves_diagnostic_context( +async def test_max_turns_guardrail_error_preserves_diagnostic_without_session_write( monkeypatch: pytest.MonkeyPatch, streamed: bool, ) -> None: persistence_secret = "DIAGNOSTIC_MAX_TURNS_SESSION_SECRET" guardrail_secret = "DIAGNOSTIC_MAX_TURNS_GUARDRAIL_SECRET" guardrail_failed = False + post_guardrail_save_calls = 0 monkeypatch.setattr(_debug, "DONT_LOG_MODEL_DATA", False) class FailingMaxTurnsSession(SimpleListSession): async def add_items(self, items: list[Any]) -> None: + nonlocal post_guardrail_save_calls if guardrail_failed: + post_guardrail_save_calls += 1 raise LookupError(persistence_secret) await super().add_items(items) @@ -2410,11 +2367,11 @@ def output_guardrail( session=session, error_handlers={"max_turns": lambda data: "fallback"}, ) - with pytest.raises(LookupError, match=persistence_secret) as exc_info: + with pytest.raises(RuntimeError, match=guardrail_secret) as exc_info: async for _ in result.stream_events(): pass else: - with pytest.raises(LookupError, match=persistence_secret) as exc_info: + with pytest.raises(RuntimeError, match=guardrail_secret) as exc_info: await Runner.run( agent, "go", @@ -2423,9 +2380,8 @@ def output_guardrail( error_handlers={"max_turns": lambda data: "fallback"}, ) - guardrail_error = exc_info.value.__context__ - assert isinstance(guardrail_error, RuntimeError) - assert guardrail_secret in str(guardrail_error) + assert post_guardrail_save_calls == 0 + assert persistence_secret not in str(exc_info.value) @pytest.mark.asyncio diff --git a/tests/test_max_turns.py b/tests/test_max_turns.py index d002609523..dedb195674 100644 --- a/tests/test_max_turns.py +++ b/tests/test_max_turns.py @@ -662,8 +662,10 @@ async def run_once() -> Any: with pytest.raises(RuntimeError, match="guardrail failed"): await run_once() elif outcome == "tripwire": - with pytest.raises(OutputGuardrailTripwireTriggered): + with pytest.raises(OutputGuardrailTripwireTriggered) as exc_info: await run_once() + assert exc_info.value.guardrail_result.agent_output == "fallback answer" + assert exc_info.value.guardrail_result.output.output_info == "tripwire" else: result = await run_once() assert result.final_output == "fallback answer" @@ -672,10 +674,10 @@ async def run_once() -> Any: saved_items = await session.get_items() saved_types = [str(item.get("type", item.get("role"))) for item in saved_items] - if outcome == "tripwire": - assert saved_types == ["user"] - else: + if outcome == "pass": assert saved_types == ["user", "message"] + else: + assert saved_types == ["user"] fallback_events = [ event @@ -688,7 +690,7 @@ async def run_once() -> Any: if streamed: assert streamed_result is not None - expected_history_count = 0 if outcome == "tripwire" else 1 + expected_history_count = 1 if outcome == "pass" else 0 assert ( len([item for item in streamed_result.new_items if isinstance(item, MessageOutputItem)]) == expected_history_count @@ -700,6 +702,64 @@ async def run_once() -> Any: ) +@pytest.mark.asyncio +async def test_streamed_max_turns_trip_preserves_completed_tool_prefix() -> None: + def reject_output( + _context: RunContextWrapper[Any], + _agent: Agent[Any], + _output: Any, + ) -> GuardrailFunctionOutput: + return GuardrailFunctionOutput(output_info=None, tripwire_triggered=True) + + model = ScriptedModel( + steps=[[get_function_tool_call("some_function", "{}", call_id="accepted-call")]] + ) + agent = Agent( + name="test", + model=model, + tools=[get_function_tool("some_function", "accepted-output")], + output_guardrails=[OutputGuardrail(guardrail_function=reject_output)], + ) + session = SimpleListSession() + result = Runner.run_streamed( + agent, + "run the tool", + max_turns=1, + session=session, + error_handlers={"max_turns": lambda data: "rejected fallback"}, + ) + + with pytest.raises(OutputGuardrailTripwireTriggered): + async for _ in result.stream_events(): + pass + + def call_ids(items: list[Any]) -> list[str]: + return [ + call_id + for item in items + if ( + call_id := ( + item.raw_item.get("call_id") + if isinstance(item.raw_item, dict) + else getattr(item.raw_item, "call_id", None) + ) + ) + is not None + ] + + assert call_ids(result.new_items) == ["accepted-call", "accepted-call"] + assert call_ids(result._model_input_items) == ["accepted-call", "accepted-call"] + state = result.to_state() + assert call_ids(state._generated_items) == ["accepted-call", "accepted-call"] + assert call_ids(state._session_items) == ["accepted-call", "accepted-call"] + serialized_state = json.dumps(state.to_json()) + assert "accepted-output" in serialized_state + assert "rejected fallback" in serialized_state + + saved_types = [item.get("type", item.get("role")) for item in await session.get_items()] + assert saved_types == ["user", "function_call", "function_call_output"] + + @pytest.mark.asyncio async def test_streamed_max_turns_handler_validation_failure_persists_input() -> None: agent = Agent(name="test", model=ScriptedModel(), output_type=Foo) @@ -995,6 +1055,52 @@ def output_guardrail( ] +@pytest.mark.parametrize("streamed", [False, True]) +@pytest.mark.asyncio +async def test_resumed_max_turns_trip_preserves_current_guardrail_result( + streamed: bool, +) -> None: + def output_guardrail( + _context: RunContextWrapper[Any], + _agent: Agent[Any], + output: Any, + ) -> GuardrailFunctionOutput: + return GuardrailFunctionOutput( + output_info=output, + tripwire_triggered=output == "fallback answer", + ) + + agent = Agent( + name="test", + model=ScriptedModel(steps=[[get_text_message("first response")]]), + output_guardrails=[OutputGuardrail(guardrail_function=output_guardrail)], + ) + first = await Runner.run(agent, "first input", max_turns=1) + state = first.to_state() + prior_result = state._output_guardrail_results[0] + + with pytest.raises(OutputGuardrailTripwireTriggered) as exc_info: + if streamed: + result = Runner.run_streamed( + agent, + state, + error_handlers={"max_turns": lambda data: "fallback answer"}, + ) + async for _ in result.stream_events(): + pass + else: + await Runner.run( + agent, + state, + error_handlers={"max_turns": lambda data: "fallback answer"}, + ) + + assert state._output_guardrail_results == [prior_result] + assert prior_result.output.output_info == "first response" + assert exc_info.value.guardrail_result.agent_output == "fallback answer" + assert exc_info.value.guardrail_result.output.output_info == "fallback answer" + + @pytest.mark.parametrize("streamed", [False, True]) @pytest.mark.asyncio async def test_resumed_max_turns_handler_preserves_checkpoint_after_continuation( From 72834a8df999aef2c640d688c3a648bd567f951c Mon Sep 17 00:00:00 2001 From: Kazuhiro Sera Date: Wed, 19 Aug 2026 04:11:46 +0900 Subject: [PATCH 04/12] Remove deferred tool span machinery --- src/agents/run_internal/run_steps.py | 71 ------------- src/agents/run_internal/tool_actions.py | 30 +++--- src/agents/run_internal/tool_execution.py | 123 ++-------------------- 3 files changed, 20 insertions(+), 204 deletions(-) diff --git a/src/agents/run_internal/run_steps.py b/src/agents/run_internal/run_steps.py index 3a67bac5ee..f692e1ca4e 100644 --- a/src/agents/run_internal/run_steps.py +++ b/src/agents/run_internal/run_steps.py @@ -26,7 +26,6 @@ ShellTool, ) from ..tool_guardrails import ToolInputGuardrailResult, ToolOutputGuardrailResult -from ..tracing import Span, SpanError from .items import NestedHistoryOwnedItem __all__ = [ @@ -182,73 +181,6 @@ class NextStepInterruption: """Whether response-end hooks started before the interruption was persisted.""" -@dataclass -class DeferredToolSpan(Span[Any]): - """A tool span whose sensitive output is held until the terminal verdict.""" - - span: Span[Any] - output: Any = None - has_output: bool = False - deferred_error: SpanError | None = None - has_error: bool = False - - @property - def trace_id(self) -> str: - return self.span.trace_id - - @property - def span_id(self) -> str: - return self.span.span_id - - @property - def span_data(self) -> Any: - return self.span.span_data - - @property - def parent_id(self) -> str | None: - return self.span.parent_id - - def start(self, mark_as_current: bool = False) -> None: - self.span.start(mark_as_current=mark_as_current) - - def finish(self, reset_current: bool = False) -> None: - self.span.finish(reset_current=reset_current) - - def __enter__(self) -> DeferredToolSpan: - self.start(mark_as_current=True) - return self - - def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None: - self.finish(reset_current=True) - - def set_error(self, error: SpanError) -> None: - self.deferred_error = error - self.has_error = True - - @property - def error(self) -> SpanError | None: - return self.deferred_error if self.has_error else self.span.error - - def export(self) -> dict[str, Any] | None: - return self.span.export() - - @property - def started_at(self) -> str | None: - return self.span.started_at - - @property - def ended_at(self) -> str | None: - return self.span.ended_at - - @property - def tracing_api_key(self) -> str | None: - return self.span.tracing_api_key - - @property - def trace_metadata(self) -> dict[str, Any] | None: - return self.span.trace_metadata - - @dataclass class SingleStepResult: original_input: str | list[TResponseInputItem] @@ -290,9 +222,6 @@ class SingleStepResult: processed_response: ProcessedResponse | None = None """The processed model response. This is needed for resuming from interruptions.""" - deferred_tool_spans: list[DeferredToolSpan] = dataclasses.field(default_factory=list) - """Tool spans waiting for the terminal output-guardrail verdict before publication.""" - @property def generated_items(self) -> list[RunItem]: """Items generated during the agent run (i.e. everything generated after diff --git a/src/agents/run_internal/tool_actions.py b/src/agents/run_internal/tool_actions.py index 95a28af452..6bffae29cd 100644 --- a/src/agents/run_internal/tool_actions.py +++ b/src/agents/run_internal/tool_actions.py @@ -58,8 +58,6 @@ resolve_approval_rejection_message, resolve_approval_status, serialize_shell_output, - set_tool_span_error, - set_tool_span_output, truncate_shell_outputs, with_tool_function_span, ) @@ -152,15 +150,14 @@ async def _run_action(span: Any | None) -> RunItem: error_message=error_text, ) if span is not None: - set_tool_span_error( - span, + span.set_error( SpanError( message="Error running tool", data={ "tool_name": trace_tool_name, "error": trace_error, }, - ), + ) ) log_tool_action_error("Failed to execute computer action", exc) raise @@ -204,7 +201,7 @@ async def _run_action(span: Any | None) -> RunItem: ) if span is not None and config.trace_include_sensitive_data: - set_tool_span_output(span, image_url) + span.span_data.output = image_url return output_item @@ -593,15 +590,14 @@ async def _run_call(span: Any | None) -> RunItem: error_message=output_text, ) if span is not None: - set_tool_span_error( - span, + span.set_error( SpanError( message="Error running tool", data={ "tool_name": shell_tool.name, "error": trace_error, }, - ), + ) ) if requested_max_output_length is not None: max_output_length = requested_max_output_length @@ -655,7 +651,7 @@ async def _run_call(span: Any | None) -> RunItem: ) if span is not None and config.trace_include_sensitive_data: - set_tool_span_output(span, output_text) + span.span_data.output = output_text return output_item @@ -781,15 +777,14 @@ async def _run_call(span: Any | None) -> RunItem: error_message=output_text, ) if span is not None: - set_tool_span_error( - span, + span.set_error( SpanError( message="Error running tool", data={ "tool_name": custom_tool.name, "error": trace_error, }, - ), + ) ) log_tool_action_error("Custom tool failed", exc) @@ -828,7 +823,7 @@ async def _run_call(span: Any | None) -> RunItem: ) if span is not None and config.trace_include_sensitive_data: - set_tool_span_output(span, output_text) + span.span_data.output = output_text return output_item return await with_tool_function_span( @@ -1014,15 +1009,14 @@ async def _run_call(span: Any | None) -> RunItem: error_message=output_text, ) if span is not None: - set_tool_span_error( - span, + span.set_error( SpanError( message="Error running tool", data={ "tool_name": apply_patch_tool.name, "error": trace_error, }, - ), + ) ) log_tool_action_error("Apply patch editor failed", exc) @@ -1066,7 +1060,7 @@ async def _run_call(span: Any | None) -> RunItem: ) if span is not None and config.trace_include_sensitive_data: - set_tool_span_output(span, output_text) + span.span_data.output = output_text return output_item diff --git a/src/agents/run_internal/tool_execution.py b/src/agents/run_internal/tool_execution.py index e32cac9efd..af257b165f 100644 --- a/src/agents/run_internal/tool_execution.py +++ b/src/agents/run_internal/tool_execution.py @@ -6,15 +6,12 @@ from __future__ import annotations import asyncio -import contextvars import copy import dataclasses import functools import inspect import json -from collections.abc import Awaitable, Callable, Iterator, Mapping, Sequence -from contextlib import contextmanager -from contextvars import Token +from collections.abc import Awaitable, Callable, Mapping, Sequence from typing import TYPE_CHECKING, Any, Literal, TypeVar, cast from openai.types.responses import ResponseFunctionToolCall @@ -96,7 +93,6 @@ ToolOutputGuardrailResult, ) from ..tracing import Span, SpanError, function_span, get_current_trace -from ..tracing.scope import Scope from ..util import _coro, _error_tracing from ..util._approvals import evaluate_needs_approval_setting, parse_function_tool_arguments from ..util._asyncio_tasks import gather_with_cancel @@ -113,7 +109,7 @@ function_rejection_item, function_tool_error_output, ) -from .run_steps import DeferredToolSpan, ToolRunFunction +from .run_steps import ToolRunFunction from .tool_use_tracker import AgentToolUseTracker if TYPE_CHECKING: @@ -127,106 +123,6 @@ ToolRunShellCall, ) - -_deferred_tool_spans: contextvars.ContextVar[list[DeferredToolSpan] | None] = ( - contextvars.ContextVar( - "deferred_tool_spans", - default=None, - ) -) - - -@contextmanager -def collect_deferred_tool_spans(enabled: bool) -> Iterator[list[DeferredToolSpan]]: - """Collect tool spans without publishing their output until the turn verdict is known.""" - spans: list[DeferredToolSpan] = [] - token: Token[list[DeferredToolSpan] | None] = _deferred_tool_spans.set( - spans if enabled else None - ) - try: - yield spans - finally: - _deferred_tool_spans.reset(token) - - -def finish_deferred_tool_spans( - spans: list[DeferredToolSpan], - *, - output_override: Any | None = None, -) -> None: - """Publish a deferred tool-span batch exactly once, optionally replacing its output.""" - pending_spans = list(spans) - spans.clear() - for deferred_span in pending_spans: - span = deferred_span.span - if span.ended_at is not None: - continue - if output_override is not None: - cast(Any, span.span_data).output = output_override - elif deferred_span.has_output: - cast(Any, span.span_data).output = deferred_span.output - if output_override is None and deferred_span.has_error: - assert deferred_span.deferred_error is not None - span.set_error(deferred_span.deferred_error) - deferred_span.output = None - deferred_span.has_output = False - deferred_span.deferred_error = None - deferred_span.has_error = False - span.finish() - - -def set_tool_span_output(span: Span[Any], output: Any) -> None: - """Store sensitive tool output outside a processor-visible span until finalization.""" - deferred_spans = _deferred_tool_spans.get() - if deferred_spans is not None: - for deferred_span in reversed(deferred_spans): - if deferred_span.span is span: - deferred_span.output = output - deferred_span.has_output = True - return - cast(Any, span.span_data).output = output - - -def set_tool_span_error(span: Span[Any], error: SpanError) -> None: - """Store a tool error outside a processor-visible span until finalization.""" - deferred_spans = _deferred_tool_spans.get() - if deferred_spans is not None: - for deferred_span in reversed(deferred_spans): - if deferred_span is span or deferred_span.span is span: - deferred_span.set_error(error) - return - span.set_error(error) - - -@contextmanager -def _tool_function_span(tool_name: str) -> Iterator[Span[Any]]: - """Keep current-span ownership in the tool task while allowing deferred publication.""" - span = function_span(tool_name) - span.start() - deferred_spans = _deferred_tool_spans.get() - if deferred_spans is None: - deferred_span = None - else: - deferred_span = DeferredToolSpan(span=span) - deferred_spans.append(deferred_span) - token = Scope.set_current_span(deferred_span or span) - try: - yield span - except BaseException: - Scope.reset_current_span(token) - if deferred_span is None: - span.finish() - else: - assert deferred_spans is not None - deferred_spans.remove(deferred_span) - finish_deferred_tool_spans([deferred_span]) - raise - else: - Scope.reset_current_span(token) - if deferred_span is None: - span.finish() - - __all__ = [ "maybe_reset_tool_choice", "initialize_computer_tools", @@ -250,8 +146,6 @@ def _tool_function_span(tool_name: str) -> Iterator[Span[Any]]: "format_shell_error", "get_trace_tool_error", "with_tool_function_span", - "set_tool_span_output", - "set_tool_span_error", "build_litellm_json_tool_call", "collect_manual_mcp_approvals", "index_approval_items_by_call_id", @@ -1237,7 +1131,7 @@ async def with_tool_function_span( direct_result: object = result return cast(TToolSpanResult, direct_result) - with _tool_function_span(tool_name) as span: + with function_span(tool_name) as span: result = fn(span) if inspect.isawaitable(result): return await result @@ -1908,7 +1802,7 @@ async def _run_single_tool( or get_function_tool_trace_name(func_tool) or func_tool.name ) - with _tool_function_span(trace_tool_name) as span_fn: + with function_span(trace_tool_name) as span_fn: tool_context_namespace = get_tool_call_namespace(raw_tool_call) if tool_context_namespace is None: tool_context_namespace = get_tool_call_namespace(tool_call) @@ -1958,7 +1852,7 @@ async def _run_single_tool( raise UserError(f"Error running tool {func_tool.name}: {e}") from e if self.config.trace_include_sensitive_data: - set_tool_span_output(span_fn, result) + span_fn.span_data.output = result return result async def _maybe_execute_tool_approval( @@ -2070,8 +1964,7 @@ async def _maybe_execute_tool_approval( tool_namespace=tool_namespace, tool_lookup_key=tool_lookup_key, ) - set_tool_span_error( - span_fn, + span_fn.set_error( SpanError( message=rejection_message, data={ @@ -2080,9 +1973,9 @@ async def _maybe_execute_tool_approval( f"Tool execution for {tool_call.call_id} was manually rejected by user." ), }, - ), + ) ) - set_tool_span_output(span_fn, rejection_message) + span_fn.span_data.output = rejection_message return FunctionToolResult( tool=func_tool, output=rejection_message, From 4a3e4a58b67f85def8b800b60f6a337063db906f Mon Sep 17 00:00:00 2001 From: Kazuhiro Sera Date: Wed, 19 Aug 2026 07:58:49 +0900 Subject: [PATCH 05/12] fix: harden blocked output snapshots --- src/agents/run.py | 74 +++++-- src/agents/run_internal/blocked_output.py | 35 ++- src/agents/run_internal/run_loop.py | 33 ++- tests/test_agent_runner.py | 258 +++++++++++++++++++++- tests/test_error_logging_redaction.py | 63 ++++++ 5 files changed, 424 insertions(+), 39 deletions(-) diff --git a/src/agents/run.py b/src/agents/run.py index c955c0484a..6534f7ca95 100644 --- a/src/agents/run.py +++ b/src/agents/run.py @@ -95,11 +95,13 @@ from .run_internal.prompt_cache_key import PromptCacheKeyResolver from .run_internal.run_grouping import resolve_run_grouping_id from .run_internal.run_loop import ( + _blocked_output_failure_items, _BlockedOutputOwnerStarts, _current_response_boundary, _final_turn_items_for_persistence, _is_terminal_tool_output_response, _retained_items_for_blocked_response, + _safe_redacted_persistence_error, _sanitize_blocked_output_guardrail_results, _should_defer_interrupted_session_items, _validate_resumed_session_output_guardrail_safety, @@ -1067,6 +1069,7 @@ def _mark_response_hooks_started() -> None: run_state, ) blocked_output_owner_starts = _BlockedOutputOwnerStarts( + nonstreamed_session_items=(resumed_response_boundary.session_start), run_state_generated_items=( resumed_response_boundary.generated_start ), @@ -1244,6 +1247,11 @@ def _mark_response_hooks_started() -> None: output_guardrail_results[output_guardrail_result_start:] = ( sanitized_results ) + session_items = _blocked_output_failure_items( + session_items, + (), + blocked_output_owner_starts, + ) retained_items = _retained_items_for_blocked_response( turn_session_items, turn_result.model_response, @@ -1251,18 +1259,26 @@ def _mark_response_hooks_started() -> None: current_processed_response, owner_starts=blocked_output_owner_starts, ) - await save_final_turn_items_after_guardrails( - session=session, - run_state=run_state, - session_persistence_enabled=session_persistence_enabled, - input_guardrail_results=( - _attempt_input_guardrail_results() - ), - items=retained_items, - response_id=turn_result.model_response.response_id, - store=store_setting, - wrapper=context_wrapper, - ) + list.extend(session_items, retained_items) + try: + await save_final_turn_items_after_guardrails( + session=session, + run_state=run_state, + session_persistence_enabled=( + session_persistence_enabled + ), + input_guardrail_results=( + _attempt_input_guardrail_results() + ), + items=retained_items, + response_id=turn_result.model_response.response_id, + store=store_setting, + wrapper=context_wrapper, + ) + except BaseException as persistence_error: + raise _safe_redacted_persistence_error( + persistence_error + ) from None raise except (Exception, asyncio.CancelledError): # Without a verdict, do not persist any part of this response. @@ -1501,6 +1517,7 @@ async def _save_max_turns_handler_output( last_saved_input_snapshot_for_rewind = None blocked_output_owner_starts = _BlockedOutputOwnerStarts( + nonstreamed_session_items=len(session_items), run_state_generated_items=( len(run_state._generated_items) if run_state is not None else None ), @@ -1800,6 +1817,11 @@ async def _save_max_turns_handler_output( output_guardrail_results[output_guardrail_result_start:] = ( sanitized_results ) + session_items = _blocked_output_failure_items( + session_items, + (), + blocked_output_owner_starts, + ) retained_items = _retained_items_for_blocked_response( turn_session_items, turn_result.model_response, @@ -1807,16 +1829,24 @@ async def _save_max_turns_handler_output( turn_result.processed_response, owner_starts=blocked_output_owner_starts, ) - await save_final_turn_items_after_guardrails( - session=session, - run_state=run_state, - session_persistence_enabled=session_persistence_enabled, - input_guardrail_results=_attempt_input_guardrail_results(), - items=retained_items, - response_id=turn_result.model_response.response_id, - store=store_setting, - wrapper=context_wrapper, - ) + list.extend(session_items, retained_items) + try: + await save_final_turn_items_after_guardrails( + session=session, + run_state=run_state, + session_persistence_enabled=(session_persistence_enabled), + input_guardrail_results=( + _attempt_input_guardrail_results() + ), + items=retained_items, + response_id=turn_result.model_response.response_id, + store=store_setting, + wrapper=context_wrapper, + ) + except BaseException as persistence_error: + raise _safe_redacted_persistence_error( + persistence_error + ) from None raise except (Exception, asyncio.CancelledError): # Without a verdict, do not persist any part of this response. diff --git a/src/agents/run_internal/blocked_output.py b/src/agents/run_internal/blocked_output.py index 084935f5dc..e68705728f 100644 --- a/src/agents/run_internal/blocked_output.py +++ b/src/agents/run_internal/blocked_output.py @@ -5,6 +5,7 @@ from typing import Any from openai.types.responses import ResponseFunctionToolCall +from openai.types.responses.response_function_tool_call import CallerDirect from ..exceptions import AgentsException @@ -13,14 +14,25 @@ _RESPONSE_OUTPUT_STATUSES = frozenset({"in_progress", "completed", "incomplete"}) +def _exact_dict_field(values: dict[Any, Any], field: str) -> Any: + """Read one exact string key without invoking stored-key equality hooks.""" + for key, value in dict.items(values): + if type(key) is str and str.__eq__(key, field) is True: + return value + return None + + def _payload_field(raw_item: Any, field: str) -> Any: """Read an allowlisted field without copying extras or invoking instance hooks.""" if type(raw_item) is dict: - return dict.get(raw_item, field) - if type(raw_item) is ResponseFunctionToolCall: + values = raw_item + elif type(raw_item) is ResponseFunctionToolCall: values = object.__getattribute__(raw_item, "__dict__") - return dict.get(values, field) - raise AgentsException("Cannot sanitize an unsupported tool item variant.") + else: + raise AgentsException("Cannot sanitize an unsupported tool item variant.") + if type(values) is not dict: + raise AgentsException("Cannot sanitize an unsupported tool item representation.") + return _exact_dict_field(values, field) def _required_string(raw_item: Any, field: str) -> str: @@ -56,7 +68,14 @@ def _copy_optional_direct_caller(sanitized: dict[str, Any], raw_item: Any) -> No caller = _payload_field(raw_item, "caller") if caller is None: return - if type(caller) is dict and dict.get(caller, "type") == "direct": + if type(caller) is CallerDirect: + values = object.__getattribute__(caller, "__dict__") + caller_type = _exact_dict_field(values, "type") if type(values) is dict else None + elif type(caller) is dict: + caller_type = _exact_dict_field(caller, "type") + else: + caller_type = None + if type(caller_type) is str and str.__eq__(caller_type, "direct") is True: sanitized["caller"] = {"type": "direct"} return raise AgentsException("Cannot sanitize a function tool item with a non-direct caller.") @@ -64,7 +83,8 @@ def _copy_optional_direct_caller(sanitized: dict[str, Any], raw_item: Any) -> No def blocked_function_call_payload(raw_item: Any) -> dict[str, Any]: """Build a provider-valid function call from explicitly allowlisted fields.""" - if _payload_field(raw_item, "type") != "function_call": + item_type = _payload_field(raw_item, "type") + if type(item_type) is not str or str.__eq__(item_type, "function_call") is not True: raise AgentsException("Cannot sanitize an unsupported tool call variant.") sanitized: dict[str, Any] = { "type": "function_call", @@ -85,7 +105,8 @@ def blocked_function_call_payload(raw_item: Any) -> dict[str, Any]: def blocked_function_output_payload(raw_item: Any) -> dict[str, Any]: """Build a replay-valid function output from explicitly allowlisted fields.""" - if _payload_field(raw_item, "type") != "function_call_output": + item_type = _payload_field(raw_item, "type") + if type(item_type) is not str or str.__eq__(item_type, "function_call_output") is not True: raise AgentsException("Cannot sanitize an unsupported tool output variant.") sanitized: dict[str, Any] = { "type": "function_call_output", diff --git a/src/agents/run_internal/run_loop.py b/src/agents/run_internal/run_loop.py index e805119d9f..ab014062cb 100644 --- a/src/agents/run_internal/run_loop.py +++ b/src/agents/run_internal/run_loop.py @@ -531,6 +531,7 @@ class _BlockedOutputOwnerPlan: class _BlockedOutputOwnerStarts: """Owner-specific current-response starts captured at trusted lifecycle boundaries.""" + nonstreamed_session_items: int | None = None run_state_generated_items: int | None = None run_state_session_items: int | None = None run_state_model_responses: int | None = None @@ -789,6 +790,18 @@ def _blocked_output_owner_prefix(items: list[_OwnerItemT], start: int | None) -> return list.__getitem__(items, slice(0, start)) +def _blocked_output_failure_items( + items: list[RunItem], + retained_items: Sequence[RunItem], + owner_starts: _BlockedOutputOwnerStarts, +) -> list[RunItem]: + """Build the non-streamed accepted prefix plus the data-free current response.""" + return [ + *_blocked_output_owner_prefix(items, owner_starts.nonstreamed_session_items), + *retained_items, + ] + + def _prepare_blocked_output_owner_prefixes( run_state: RunState[Any] | None, streamed_result: RunResultStreaming | None, @@ -1154,15 +1167,17 @@ async def _finalize_streamed_final_output( if retained_items: try: await save_items(retained_items, response_id, store_setting) - except asyncio.CancelledError as persistence_error: - if streamed_result._cancel_mode == "immediate": - raise - streamed_result._stored_exception = _safe_redacted_persistence_error( - persistence_error - ) - streamed_result.is_complete = True - streamed_result._event_queue.put_nowait(QueueCompleteSentinel()) - return + except BaseException as persistence_error: + safe_error = _safe_redacted_persistence_error(persistence_error) + if ( + isinstance(persistence_error, asyncio.CancelledError) + and streamed_result._cancel_mode != "immediate" + ): + streamed_result._stored_exception = safe_error + streamed_result.is_complete = True + streamed_result._event_queue.put_nowait(QueueCompleteSentinel()) + return + raise safe_error from None raise except (Exception, asyncio.CancelledError): # Without a verdict, the SDK does not persist any part of the terminal response. diff --git a/tests/test_agent_runner.py b/tests/test_agent_runner.py index 1e5569ebeb..a4dbc87609 100644 --- a/tests/test_agent_runner.py +++ b/tests/test_agent_runner.py @@ -14,10 +14,12 @@ import pytest from openai import APIConnectionError, BadRequestError, NotFoundError from openai.types.responses import ResponseFunctionToolCall +from openai.types.responses.response_function_tool_call import CallerDirect, CallerProgram from openai.types.responses.response_output_item import McpApprovalRequest from openai.types.responses.response_output_text import AnnotationFileCitation, ResponseOutputText from openai.types.responses.response_reasoning_item import ResponseReasoningItem, Summary from openai.types.responses.tool_param import Mcp +from pydantic import BaseModel from typing_extensions import TypedDict import agents._debug as _debug @@ -165,6 +167,156 @@ def test_blocked_function_batch_is_rebuilt_from_allowlisted_fields() -> None: assert retained_output.custom_data is None +def test_blocked_function_batch_accepts_exact_typed_direct_caller() -> None: + agent = Agent(name="test") + call = ToolCallItem( + agent=agent, + raw_item=ResponseFunctionToolCall( + type="function_call", + name="commit_tool", + arguments="{}", + call_id="call-commit", + caller=CallerDirect(type="direct"), + ), + ) + output = ToolCallOutputItem( + agent=agent, + raw_item={ + "type": "function_call_output", + "call_id": "call-commit", + "output": "raw-secret", + }, + output="sdk-secret", + ) + + retained = run_loop._retained_items_for_blocked_output([call, output]) + + assert len(retained) == 2 + retained_call = cast(ToolCallItem, retained[0]) + assert cast(dict[str, Any], retained_call.raw_item)["caller"] == {"type": "direct"} + + +def test_blocked_function_batch_ignores_hash_collision_key_hooks() -> None: + equality_calls: list[Any] = [] + + class HashCollisionKey: + def __init__(self, field: str) -> None: + self.field = field + + def __hash__(self) -> int: + return hash(self.field) + + def __eq__(self, other: object) -> bool: + equality_calls.append(other) + return False + + caller: dict[Any, Any] = { + HashCollisionKey("type"): "caller-secret", + "type": "direct", + } + raw_call: dict[Any, Any] = { + HashCollisionKey("type"): "type-secret", + HashCollisionKey("name"): "name-secret", + HashCollisionKey("arguments"): "arguments-secret", + HashCollisionKey("call_id"): "call-id-secret", + HashCollisionKey("id"): "id-secret", + HashCollisionKey("namespace"): "namespace-secret", + HashCollisionKey("status"): "status-secret", + HashCollisionKey("caller"): "caller-secret", + "type": "function_call", + "name": "commit_tool", + "arguments": "{}", + "call_id": "call-commit", + "caller": caller, + } + equality_calls.clear() + agent = Agent(name="test") + call = ToolCallItem(agent=agent, raw_item=cast(Any, raw_call)) + output = ToolCallOutputItem( + agent=agent, + raw_item={ + "type": "function_call_output", + "call_id": "call-commit", + "output": "raw-secret", + }, + output="sdk-secret", + ) + + retained = run_loop._retained_items_for_blocked_output([call, output]) + + assert len(retained) == 2 + assert equality_calls == [] + retained_call = cast(ToolCallItem, retained[0]) + assert cast(dict[str, Any], retained_call.raw_item)["caller"] == {"type": "direct"} + + +@pytest.mark.parametrize("discriminator", ["call", "output", "caller"]) +def test_blocked_function_batch_rejects_equality_impostor_discriminators_without_hooks( + discriminator: str, +) -> None: + equality_calls: list[Any] = [] + + class EqualityImpostor: + def __eq__(self, other: object) -> bool: + equality_calls.append(other) + return True + + raw_call: dict[str, Any] = { + "type": EqualityImpostor() if discriminator == "call" else "function_call", + "name": "commit_tool", + "arguments": "{}", + "call_id": "call-commit", + } + if discriminator == "caller": + raw_call["caller"] = {"type": EqualityImpostor()} + agent = Agent(name="test") + call = ToolCallItem(agent=agent, raw_item=cast(Any, raw_call)) + output = ToolCallOutputItem( + agent=agent, + raw_item={ + "type": (EqualityImpostor() if discriminator == "output" else "function_call_output"), + "call_id": "call-commit", + "output": "raw-secret", + }, + output="sdk-secret", + ) + + assert run_loop._retained_items_for_blocked_output([call, output]) == [] + assert equality_calls == [] + + +def test_blocked_function_batch_rejects_non_direct_typed_callers() -> None: + class GenericCaller(BaseModel): + type: str + + agent = Agent(name="test") + output = ToolCallOutputItem( + agent=agent, + raw_item={ + "type": "function_call_output", + "call_id": "call-commit", + "output": "raw-secret", + }, + output="sdk-secret", + ) + for caller in ( + CallerProgram(type="program", caller_id="program-call"), + GenericCaller(type="direct"), + ): + call = ToolCallItem( + agent=agent, + raw_item={ + "type": "function_call", + "name": "commit_tool", + "arguments": "{}", + "call_id": "call-commit", + "caller": caller, + }, + ) + + assert run_loop._retained_items_for_blocked_output([call, output]) == [] + + def test_blocked_unknown_tool_variant_discards_the_complete_response() -> None: agent = Agent(name="test") call = ToolCallItem( @@ -527,8 +679,11 @@ def fail_application(_plan: Any) -> None: @pytest.mark.asyncio -async def test_non_streamed_trip_preserves_prior_run_state_side_effect() -> None: +async def test_non_streamed_trip_preserves_prior_run_state_side_effect( + monkeypatch: pytest.MonkeyPatch, +) -> None: side_effects: list[str] = [] + memory_items: list[RunItem] | None = None @function_tool(name_override="accepted_tool") def accepted_tool() -> str: @@ -562,6 +717,24 @@ def reject_output( agent.tool_use_behavior = {"stop_at_tool_names": ["terminal_tool"]} agent.output_guardrails = [OutputGuardrail(guardrail_function=reject_output)] + + async def capture_memory_payload( + _runtime: Any, + *, + input: Any, + new_items: list[Any], + final_output: object, + interruptions: list[Any], + terminal_metadata: Any, + ) -> None: + del input, final_output, interruptions, terminal_metadata + nonlocal memory_items + memory_items = cast(list[RunItem], new_items) + + monkeypatch.setattr( + "agents.run.SandboxRuntime.enqueue_memory_payload", + capture_memory_payload, + ) model.enqueue( [ ResponseReasoningItem( @@ -584,6 +757,89 @@ def reject_output( assert "accepted-output" in serialized_state assert "rejected-output" not in serialized_state assert "reasoning-current" not in serialized_state + assert memory_items is not None + serialized_memory_items = json.dumps([item.to_input_item() for item in memory_items]) + assert "accepted-output" in serialized_memory_items + assert "rejected-output" not in serialized_memory_items + assert "reasoning-current" not in serialized_memory_items + + +@pytest.mark.asyncio +async def test_non_streamed_trip_uses_safe_items_for_sandbox_memory_after_session_failure( + monkeypatch: pytest.MonkeyPatch, +) -> None: + accepted_output = "accepted-tool-output" + tool_output_secret = "sandbox-memory-tool-output-secret" + persistence_secret = "sandbox-memory-session-failure-secret" + memory_items: list[RunItem] | None = None + + @function_tool(name_override="accepted_tool") + def accepted_tool() -> str: + return accepted_output + + @function_tool(name_override="terminal_tool") + def terminal_tool() -> str: + return tool_output_secret + + def reject_output( + _context: RunContextWrapper[Any], + _agent: Agent[Any], + _output: Any, + ) -> GuardrailFunctionOutput: + return GuardrailFunctionOutput(output_info=None, tripwire_triggered=True) + + class FailingBlockedSession(SimpleListSession): + async def add_items(self, items: list[Any]) -> None: + if any( + type(item) is dict + and item.get("type") == "function_call_output" + and item.get("call_id") == "terminal-call" + for item in items + ): + raise LookupError(persistence_secret) + await super().add_items(items) + + async def capture_memory_payload( + _runtime: Any, + *, + input: Any, + new_items: list[Any], + final_output: object, + interruptions: list[Any], + terminal_metadata: Any, + ) -> None: + del input, final_output, interruptions, terminal_metadata + nonlocal memory_items + memory_items = cast(list[RunItem], new_items) + + monkeypatch.setattr( + "agents.run.SandboxRuntime.enqueue_memory_payload", + capture_memory_payload, + ) + agent = Agent( + name="test", + model=ScriptedModel( + steps=[ + [get_function_tool_call("accepted_tool", "{}", call_id="accepted-call")], + [get_function_tool_call("terminal_tool", "{}", call_id="terminal-call")], + ] + ), + tools=[accepted_tool, terminal_tool], + tool_use_behavior={"stop_at_tool_names": ["terminal_tool"]}, + output_guardrails=[OutputGuardrail(guardrail_function=reject_output)], + ) + + with pytest.raises(UserError, match="Error details are redacted") as exc_info: + await Runner.run(agent, "run terminal tool", session=FailingBlockedSession()) + + assert exc_info.value.run_data is None + assert persistence_secret not in str(exc_info.value) + assert memory_items is not None + serialized_memory_items = json.dumps([item.to_input_item() for item in memory_items]) + assert accepted_output in serialized_memory_items + assert tool_output_secret not in serialized_memory_items + assert persistence_secret not in serialized_memory_items + assert run_loop._OUTPUT_GUARDRAIL_BLOCKED_TOOL_OUTPUT in serialized_memory_items async def run_execute_approved_tools( diff --git a/tests/test_error_logging_redaction.py b/tests/test_error_logging_redaction.py index ff8722aac3..d4eeccfa10 100644 --- a/tests/test_error_logging_redaction.py +++ b/tests/test_error_logging_redaction.py @@ -1654,6 +1654,69 @@ def output_guardrail( _assert_secret_absent_from_agents_traceback(error, _MODEL_OUTPUT_SECRET) +@pytest.mark.parametrize("streamed", [False, True]) +@pytest.mark.asyncio +async def test_blocked_terminal_tool_session_failure_is_data_redacted( + streamed: bool, +) -> None: + tool_output_secret = "BLOCKED_TERMINAL_TOOL_OUTPUT_SECRET" + persistence_secret = "BLOCKED_TERMINAL_SESSION_FAILURE_SECRET" + + @function_tool(name_override="terminal_tool") + def terminal_tool() -> str: + return tool_output_secret + + def reject_output( + _context: RunContextWrapper[Any], + _agent: Agent[Any], + _output: Any, + ) -> GuardrailFunctionOutput: + return GuardrailFunctionOutput(output_info=None, tripwire_triggered=True) + + class FailingBlockedSession(SimpleListSession): + async def add_items(self, items: list[Any]) -> None: + if any( + type(item) is dict and item.get("type") == "function_call_output" for item in items + ): + error = LookupError(f"session save failed: {persistence_secret}") + error.run_data = items # type: ignore[attr-defined] + raise error + await super().add_items(items) + + agent = Agent( + name="test", + model=ScriptedModel( + steps=[[get_function_tool_call("terminal_tool", "{}", call_id="terminal-call")]] + ), + tools=[terminal_tool], + tool_use_behavior="stop_on_first_tool", + output_guardrails=[OutputGuardrail(guardrail_function=reject_output)], + ) + session = FailingBlockedSession() + + if streamed: + result = Runner.run_streamed(agent, "run terminal tool", session=session) + with pytest.raises(UserError) as exc_info: + async for _ in result.stream_events(): + pass + else: + with pytest.raises(UserError) as exc_info: + await Runner.run(agent, "run terminal tool", session=session) + + error = exc_info.value + assert str(error) == "Error details are redacted." + assert error.run_data is None + assert error.__cause__ is None + assert error.__context__ is None + for secret in (tool_output_secret, persistence_secret): + assert secret not in repr(error) + _assert_secret_absent_from_agents_traceback( + error, + secret, + require_agents_frames=False, + ) + + def _persistence_failure( kind: Literal["exception", "cancelled", "direct_base", "exception_group", "group"], secret: str, From 77648c97c74acd9c5afd69ddeb2bcdb577fee1e8 Mon Sep 17 00:00:00 2001 From: Kazuhiro Sera Date: Wed, 19 Aug 2026 09:12:05 +0900 Subject: [PATCH 06/12] fix: retain zero-argument tool calls after guardrail trips --- src/agents/run_internal/blocked_output.py | 5 ++++- tests/test_agent_runner.py | 6 ++++-- tests/test_agent_runner_streamed.py | 4 +++- 3 files changed, 11 insertions(+), 4 deletions(-) diff --git a/src/agents/run_internal/blocked_output.py b/src/agents/run_internal/blocked_output.py index e68705728f..a7f2f31928 100644 --- a/src/agents/run_internal/blocked_output.py +++ b/src/agents/run_internal/blocked_output.py @@ -86,10 +86,13 @@ def blocked_function_call_payload(raw_item: Any) -> dict[str, Any]: item_type = _payload_field(raw_item, "type") if type(item_type) is not str or str.__eq__(item_type, "function_call") is not True: raise AgentsException("Cannot sanitize an unsupported tool call variant.") + arguments = _payload_field(raw_item, "arguments") + if type(arguments) is not str: + raise AgentsException("Cannot sanitize a function tool item without arguments.") sanitized: dict[str, Any] = { "type": "function_call", "name": _required_string(raw_item, "name"), - "arguments": _required_string(raw_item, "arguments"), + "arguments": arguments, "call_id": _required_string(raw_item, "call_id"), } _copy_optional_string(sanitized, raw_item, "id") diff --git a/tests/test_agent_runner.py b/tests/test_agent_runner.py index a4dbc87609..9f0e3998f4 100644 --- a/tests/test_agent_runner.py +++ b/tests/test_agent_runner.py @@ -128,14 +128,15 @@ def to_input_item(self) -> dict[str, Any]: return self._payload -def test_blocked_function_batch_is_rebuilt_from_allowlisted_fields() -> None: +@pytest.mark.parametrize("arguments", ["{}", ""], ids=["json-object", "empty"]) +def test_blocked_function_batch_is_rebuilt_from_allowlisted_fields(arguments: str) -> None: agent = Agent(name="test") call = ToolCallItem( agent=agent, raw_item={ "type": "function_call", "name": "commit_tool", - "arguments": "{}", + "arguments": arguments, "call_id": "call-commit", "provider_data": {"secret": "call-secret"}, }, @@ -158,6 +159,7 @@ def test_blocked_function_batch_is_rebuilt_from_allowlisted_fields() -> None: retained_call = cast(ToolCallItem, retained[0]) retained_output = cast(ToolCallOutputItem, retained[1]) assert "provider_data" not in cast(dict[str, Any], retained_call.raw_item) + assert cast(dict[str, Any], retained_call.raw_item)["arguments"] == arguments assert cast(dict[str, Any], retained_output.raw_item) == { "type": "function_call_output", "call_id": "call-commit", diff --git a/tests/test_agent_runner_streamed.py b/tests/test_agent_runner_streamed.py index e712af1a53..4bff39ee2e 100644 --- a/tests/test_agent_runner_streamed.py +++ b/tests/test_agent_runner_streamed.py @@ -2410,10 +2410,12 @@ async def test_output_guardrails_fail_closed_with_server_managed_history(mode: s @pytest.mark.parametrize("mode", ["non_streamed", "streamed"]) @pytest.mark.parametrize("session_kind", ["simple", "openai_conversations"]) +@pytest.mark.parametrize("arguments", ["{}", ""], ids=["json-object", "empty"]) @pytest.mark.asyncio async def test_stop_on_first_tool_final_persists_committed_tool_items_on_tripwire( mode: str, session_kind: str, + arguments: str, ) -> None: """A blocked final output must not discard the session record of a tool that already ran.""" @@ -2432,7 +2434,7 @@ def output_guardrail( return GuardrailFunctionOutput(output_info=None, tripwire_triggered=True) model = ScriptedModel() - model.enqueue([get_function_tool_call("commit_tool", "{}", call_id="call-committed")]) + model.enqueue([get_function_tool_call("commit_tool", arguments, call_id="call-committed")]) agent = Agent( name="test", model=model, From 2c00703ebd7e8cfc4986c950693e2f2b38946d3d Mon Sep 17 00:00:00 2001 From: Kazuhiro Sera Date: Wed, 19 Aug 2026 09:36:09 +0900 Subject: [PATCH 07/12] fix: preserve released guardrail failure persistence --- src/agents/run.py | 38 +++++- src/agents/run_internal/run_loop.py | 65 +++++++++- tests/test_agent_runner.py | 20 +-- tests/test_agent_runner_streamed.py | 171 +++++++++++++++++++++++++ tests/test_error_logging_redaction.py | 172 ++++++++++++++++---------- tests/test_max_turns.py | 4 +- 6 files changed, 390 insertions(+), 80 deletions(-) diff --git a/src/agents/run.py b/src/agents/run.py index 6534f7ca95..71f29551a4 100644 --- a/src/agents/run.py +++ b/src/agents/run.py @@ -1281,7 +1281,23 @@ def _mark_response_hooks_started() -> None: ) from None raise except (Exception, asyncio.CancelledError): - # Without a verdict, do not persist any part of this response. + if not _is_terminal_tool_output_response( + turn_session_items, + current_processed_response, + run_state, + ): + await save_final_turn_items_after_guardrails( + session=session, + run_state=run_state, + session_persistence_enabled=session_persistence_enabled, + input_guardrail_results=( + _attempt_input_guardrail_results() + ), + items=turn_session_items, + response_id=turn_result.model_response.response_id, + store=store_setting, + wrapper=context_wrapper, + ) raise final_turn_items = _final_turn_items_for_persistence( @@ -1470,6 +1486,8 @@ async def _save_max_turns_handler_output( output=handler_result.final_output, context_wrapper=context_wrapper, output_guardrail_results=output_guardrail_results, + save_items_after_guardrails=_save_max_turns_handler_output, + include_in_history=include_in_history, ) if include_in_history and not handler_output_recorded: await _save_max_turns_handler_output([synthesized_item]) @@ -1849,7 +1867,23 @@ async def _save_max_turns_handler_output( ) from None raise except (Exception, asyncio.CancelledError): - # Without a verdict, do not persist any part of this response. + if not _is_terminal_tool_output_response( + turn_session_items, + turn_result.processed_response, + run_state, + ): + await save_final_turn_items_after_guardrails( + session=session, + run_state=run_state, + session_persistence_enabled=session_persistence_enabled, + input_guardrail_results=( + _attempt_input_guardrail_results() + ), + items=items_to_save_turn, + response_id=turn_result.model_response.response_id, + store=store_setting, + wrapper=context_wrapper, + ) raise final_turn_items = _final_turn_items_for_persistence( diff --git a/src/agents/run_internal/run_loop.py b/src/agents/run_internal/run_loop.py index ab014062cb..f6ea75ea7d 100644 --- a/src/agents/run_internal/run_loop.py +++ b/src/agents/run_internal/run_loop.py @@ -1129,6 +1129,7 @@ async def _finalize_streamed_final_output( on_persisted_after_guardrails: Callable[[bool], None] | None = None, ) -> None: output_guardrail_result_start = len(streamed_result.output_guardrail_results) + redacted_persistence_error: BaseException | None = None try: output_guardrail_results = await _run_output_guardrails_for_stream( agent=agent, @@ -1179,9 +1180,47 @@ async def _finalize_streamed_final_output( return raise safe_error from None raise - except (Exception, asyncio.CancelledError): - # Without a verdict, the SDK does not persist any part of the terminal response. - raise + except Exception as guardrail_error: + if _is_terminal_tool_output_response( + items, + processed_response, + streamed_result._state, + ): + raise + guardrail_error_is_redacted = _is_error_data_redacted(guardrail_error) + if guardrail_error_is_redacted: + _detach_data_redacted_error_traceback(guardrail_error) + try: + await save_items(items, response_id, store_setting) + except BaseException as persistence_error: + if guardrail_error_is_redacted: + safe_persistence_error = _safe_redacted_persistence_error(persistence_error) + if ( + isinstance(safe_persistence_error, asyncio.CancelledError) + and streamed_result._cancel_mode != "immediate" + ): + streamed_result._stored_exception = safe_persistence_error + streamed_result.is_complete = True + streamed_result._event_queue.put_nowait(QueueCompleteSentinel()) + return + if isinstance(safe_persistence_error, asyncio.CancelledError): + return + redacted_persistence_error = safe_persistence_error + if ( + isinstance(persistence_error, asyncio.CancelledError) + and streamed_result._cancel_mode != "immediate" + ): + streamed_result._stored_exception = persistence_error + if redacted_persistence_error is None: + raise + else: + if on_persisted_after_guardrails is not None: + on_persisted_after_guardrails(False) + if redacted_persistence_error is None: + raise + + if redacted_persistence_error is not None: + raise redacted_persistence_error from None streamed_result.output_guardrail_results.extend(output_guardrail_results) final_turn_items = _final_turn_items_for_persistence( @@ -1316,6 +1355,8 @@ async def finalize_max_turns_handler_output( output: Any, context_wrapper: RunContextWrapper[TContext], output_guardrail_results: list[OutputGuardrailResult], + save_items_after_guardrails: Callable[[list[RunItem]], Awaitable[None]], + include_in_history: bool, ) -> tuple[Any, RunItem]: """Validate and finalize one synthesized max-turn handler output.""" validated_output = validate_handler_final_output(agent, output) @@ -1324,6 +1365,7 @@ async def finalize_max_turns_handler_output( await run_final_output_hooks(agent, hooks, context_wrapper, validated_output) + redacted_persistence_error: BaseException | None = None try: await run_output_guardrails( agent.output_guardrails + (run_config.output_guardrails or []), @@ -1332,10 +1374,23 @@ async def finalize_max_turns_handler_output( context_wrapper, output_guardrail_results, ) + except OutputGuardrailTripwireTriggered: + raise except Exception as guardrail_error: - if _is_error_data_redacted(guardrail_error): + guardrail_error_is_redacted = _is_error_data_redacted(guardrail_error) + if guardrail_error_is_redacted: _detach_data_redacted_error_traceback(guardrail_error) - raise + try: + await save_items_after_guardrails([synthesized_item] if include_in_history else []) + except BaseException as persistence_error: + if not guardrail_error_is_redacted: + raise + redacted_persistence_error = _safe_redacted_persistence_error(persistence_error) + if redacted_persistence_error is None: + raise + + if redacted_persistence_error is not None: + raise redacted_persistence_error from None return validated_output, synthesized_item diff --git a/tests/test_agent_runner.py b/tests/test_agent_runner.py index 9f0e3998f4..c56a45f86c 100644 --- a/tests/test_agent_runner.py +++ b/tests/test_agent_runner.py @@ -5193,8 +5193,9 @@ def guardrail_function( ] == ["user"] +@pytest.mark.parametrize("streamed", [False, True]) @pytest.mark.asyncio -async def test_output_guardrail_error_does_not_persist_unverdictable_output() -> None: +async def test_output_guardrail_error_preserves_final_output_in_session(streamed: bool) -> None: def guardrail_function( _context: RunContextWrapper[Any], _agent: Agent[Any], _agent_output: Any ) -> GuardrailFunctionOutput: @@ -5202,7 +5203,7 @@ def guardrail_function( session = SimpleListSession() model = ScriptedModel() - model.enqueue([get_text_message("not_persisted_on_guardrail_error")]) + model.enqueue([get_text_message("preserved_on_guardrail_error")]) agent = Agent( name="test", model=model, @@ -5210,17 +5211,22 @@ def guardrail_function( ) with pytest.raises(RuntimeError, match="guardrail failed"): - await Runner.run(agent, input="user_message", session=session) + if streamed: + result = Runner.run_streamed(agent, input="user_message", session=session) + async for _ in result.stream_events(): + pass + else: + await Runner.run(agent, input="user_message", session=session) items = await session.get_items() assert [ cast(dict[str, Any], item).get("type") or cast(dict[str, Any], item).get("role") for item in items - ] == ["user"] + ] == ["user", "message"] @pytest.mark.asyncio -async def test_output_guardrail_cancellation_does_not_start_a_final_session_write() -> None: +async def test_output_guardrail_cancellation_preserves_final_output_in_session() -> None: guardrail_started = asyncio.Event() async def guardrail_function( @@ -5232,7 +5238,7 @@ async def guardrail_function( session = SimpleListSession() model = ScriptedModel() - model.enqueue([get_text_message("not_persisted_on_guardrail_cancellation")]) + model.enqueue([get_text_message("preserved_on_guardrail_cancellation")]) agent = Agent( name="test", model=model, @@ -5250,7 +5256,7 @@ async def guardrail_function( assert [ cast(dict[str, Any], item).get("type") or cast(dict[str, Any], item).get("role") for item in items - ] == ["user"] + ] == ["user", "message"] @pytest.mark.asyncio diff --git a/tests/test_agent_runner_streamed.py b/tests/test_agent_runner_streamed.py index 4bff39ee2e..cd0f1a495f 100644 --- a/tests/test_agent_runner_streamed.py +++ b/tests/test_agent_runner_streamed.py @@ -2731,6 +2731,177 @@ async def run_once() -> None: assert saved == ["user"] +@pytest.mark.asyncio +async def test_streamed_session_save_error_takes_precedence_over_output_guardrail_error() -> None: + guardrail_failed = False + final_turn_save_attempted = False + + class FailingFinalTurnSession(SimpleListSession): + async def add_items(self, items: list[TResponseInputItem]) -> None: + nonlocal final_turn_save_attempted + if guardrail_failed: + final_turn_save_attempted = True + raise LookupError("session save failed") + await super().add_items(items) + + def output_guardrail( + _context: RunContextWrapper[Any], + _agent: Agent[Any], + _output: Any, + ) -> GuardrailFunctionOutput: + nonlocal guardrail_failed + guardrail_failed = True + raise RuntimeError("guardrail failed") + + model = ScriptedModel() + model.enqueue([get_text_message("assistant-preamble")]) + agent = Agent( + name="test", + model=model, + output_guardrails=[OutputGuardrail(guardrail_function=output_guardrail)], + ) + session = FailingFinalTurnSession() + result = Runner.run_streamed(agent, "Hello", session=session) + + with pytest.raises(LookupError, match="session save failed") as exc_info: + await consume_stream(result) + + assert final_turn_save_attempted is True + assert isinstance(exc_info.value.__context__, RuntimeError) + assert str(exc_info.value.__context__) == "guardrail failed" + assert result.run_loop_exception is exc_info.value + assert await session.get_items() == [{"content": "Hello", "role": "user"}] + + +@pytest.mark.asyncio +async def test_streamed_session_save_cancellation_is_not_a_public_immediate_cancel() -> None: + guardrail_failed = False + + class CancellingFinalTurnSession(SimpleListSession): + async def add_items(self, items: list[TResponseInputItem]) -> None: + if guardrail_failed: + raise asyncio.CancelledError("session save cancelled") + await super().add_items(items) + + def output_guardrail( + _context: RunContextWrapper[Any], + _agent: Agent[Any], + _output: Any, + ) -> GuardrailFunctionOutput: + nonlocal guardrail_failed + guardrail_failed = True + raise RuntimeError("guardrail failed") + + model = ScriptedModel(steps=[[get_text_message("assistant-preamble")]]) + agent = Agent( + name="test", + model=model, + output_guardrails=[OutputGuardrail(guardrail_function=output_guardrail)], + ) + session = CancellingFinalTurnSession() + result = Runner.run_streamed(agent, "Hello", session=session) + + with pytest.raises(asyncio.CancelledError, match="session save cancelled") as exc_info: + await consume_stream(result) + + assert result._cancel_mode == "none" + assert result._stored_exception is exc_info.value + assert isinstance(exc_info.value.__context__, RuntimeError) + assert str(exc_info.value.__context__) == "guardrail failed" + assert await session.get_items() == [{"content": "Hello", "role": "user"}] + + +@pytest.mark.asyncio +async def test_streamed_session_save_direct_base_exception_is_terminal() -> None: + guardrail_failed = False + + class DirectAbort(BaseException): + pass + + class AbortingFinalTurnSession(SimpleListSession): + async def add_items(self, items: list[TResponseInputItem]) -> None: + if guardrail_failed: + raise DirectAbort("session save aborted") + await super().add_items(items) + + def output_guardrail( + _context: RunContextWrapper[Any], + _agent: Agent[Any], + _output: Any, + ) -> GuardrailFunctionOutput: + nonlocal guardrail_failed + guardrail_failed = True + raise RuntimeError("guardrail failed") + + agent = Agent( + name="test", + model=ScriptedModel(steps=[[get_text_message("assistant-preamble")]]), + output_guardrails=[OutputGuardrail(guardrail_function=output_guardrail)], + ) + session = AbortingFinalTurnSession() + result = Runner.run_streamed(agent, "Hello", session=session) + + with pytest.raises(DirectAbort, match="session save aborted") as exc_info: + await consume_stream(result) + + assert result._stored_exception is exc_info.value + assert result.run_loop_exception is exc_info.value + assert await session.get_items() == [{"content": "Hello", "role": "user"}] + + +@pytest.mark.asyncio +async def test_public_immediate_cancel_during_guardrail_recovery_save_stays_prompt() -> None: + guardrail_failed = False + save_started = asyncio.Event() + save_cancelled = asyncio.Event() + never_set = asyncio.Event() + + class BlockingFinalTurnSession(SimpleListSession): + async def add_items(self, items: list[TResponseInputItem]) -> None: + if guardrail_failed: + save_started.set() + try: + await never_set.wait() + finally: + save_cancelled.set() + return + await super().add_items(items) + + def output_guardrail( + _context: RunContextWrapper[Any], + _agent: Agent[Any], + _output: Any, + ) -> GuardrailFunctionOutput: + nonlocal guardrail_failed + guardrail_failed = True + raise RuntimeError("guardrail failed") + + model = ScriptedModel(steps=[[get_text_message("assistant-preamble")]]) + agent = Agent( + name="test", + model=model, + output_guardrails=[OutputGuardrail(guardrail_function=output_guardrail)], + ) + session = BlockingFinalTurnSession() + result = Runner.run_streamed(agent, "Hello", session=session) + drain_task = asyncio.create_task(consume_stream(result)) + + try: + await asyncio.wait_for(save_started.wait(), timeout=1) + result.cancel() + await asyncio.wait_for(drain_task, timeout=1) + finally: + if not drain_task.done(): + result.cancel() + drain_task.cancel() + with contextlib.suppress(asyncio.CancelledError): + await drain_task + + assert save_cancelled.is_set() + assert result._stored_exception is None + assert await session.get_items() == [{"content": "Hello", "role": "user"}] + + @pytest.mark.parametrize("mode", ["non_streamed", "streamed"]) @pytest.mark.parametrize("tripwire", [False, True], ids=["passes", "trips"]) @pytest.mark.asyncio diff --git a/tests/test_error_logging_redaction.py b/tests/test_error_logging_redaction.py index d4eeccfa10..1de64711df 100644 --- a/tests/test_error_logging_redaction.py +++ b/tests/test_error_logging_redaction.py @@ -1532,24 +1532,28 @@ def output_guardrail( @pytest.mark.parametrize("redacted", [False, True]) +@pytest.mark.parametrize("persistence_failure", ["error", "cancelled"]) @pytest.mark.asyncio -async def test_streamed_output_guardrail_error_skips_final_session_write( +async def test_streamed_session_error_after_output_guardrail_respects_redaction( monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture, redacted: bool, + persistence_failure: Literal["error", "cancelled"], ) -> None: monkeypatch.setattr(_debug, "DONT_LOG_MODEL_DATA", redacted) monkeypatch.setattr(_debug, "DONT_LOG_TOOL_DATA", False) guardrail_failed = False - post_guardrail_save_calls = 0 payload = f'{{"answer": "{_MODEL_OUTPUT_SECRET}"}}' class FailingFinalTurnSession(SimpleListSession): async def add_items(self, items: list[Any]) -> None: - nonlocal post_guardrail_save_calls if guardrail_failed: - post_guardrail_save_calls += 1 - raise AssertionError("the final Session write must not start without a verdict") + cause = RuntimeError(f"session save cause: {_MODEL_OUTPUT_SECRET}") + if persistence_failure == "cancelled": + raise asyncio.CancelledError( + f"session save cancelled: {_MODEL_OUTPUT_SECRET}" + ) from cause + raise LookupError(f"session save failed: {_MODEL_OUTPUT_SECRET}") from cause await super().add_items(items) def output_guardrail( @@ -1570,56 +1574,86 @@ def output_guardrail( output_guardrails=[OutputGuardrail(guardrail_function=output_guardrail)], ) result = Runner.run_streamed(agent, "go", session=FailingFinalTurnSession()) + run_loop_callback_errors: list[BaseException] = [] + run_loop_done = asyncio.Event() - with pytest.raises(ModelBehaviorError) as exc_info: + if redacted and persistence_failure == "cancelled": + assert result.run_loop_task is not None + + def inspect_run_loop_task(task: asyncio.Task[Any]) -> None: + try: + task.result() + except BaseException as error: + run_loop_callback_errors.append(error) + finally: + run_loop_done.set() + + result.run_loop_task.add_done_callback(inspect_run_loop_task) + expected_error_type = ( + asyncio.CancelledError + if persistence_failure == "cancelled" + else UserError + if redacted + else LookupError + ) + expected_message = ( + "Error details are redacted." + if redacted + else "session save cancelled" + if persistence_failure == "cancelled" + else "session save failed" + ) + + with pytest.raises(expected_error_type, match=expected_message) as exc_info: async for _ in result.stream_events(): pass error = exc_info.value - assert post_guardrail_save_calls == 0 - assert result.run_loop_exception is error - assert result._stored_exception is error + guardrail_error = error.__context__ if redacted: - assert error.run_data is None + assert guardrail_error is None assert error.__cause__ is None - assert error.__context__ is None assert _MODEL_OUTPUT_SECRET not in str(error) - _assert_secret_absent_from_agents_traceback( - error, - _MODEL_OUTPUT_SECRET, - require_agents_frames=False, - ) + _assert_secret_absent_from_agents_traceback(error, _MODEL_OUTPUT_SECRET) for record in caplog.records: assert _MODEL_OUTPUT_SECRET not in repr(record.__dict__) assert _MODEL_OUTPUT_SECRET not in logging.Formatter().format(record) assert record.exc_info is None else: - assert _MODEL_OUTPUT_SECRET in str(error) - assert isinstance(error.__cause__, ValidationError) + assert isinstance(guardrail_error, ModelBehaviorError) + assert error.__cause__ is not None + assert _MODEL_OUTPUT_SECRET in str(error.__cause__) + assert _MODEL_OUTPUT_SECRET in str(guardrail_error) assert any( _MODEL_OUTPUT_SECRET in repr(frame) for frame in _agents_traceback_frame_locals(error) ) - if redacted: - assert error.__traceback__ is None + + if persistence_failure == "cancelled": + assert result._stored_exception is error + assert result.run_loop_exception is None + if redacted: + await asyncio.wait_for(run_loop_done.wait(), timeout=1) + assert run_loop_callback_errors == [] + else: + assert result.run_loop_exception is error + if redacted: + assert error.__traceback__ is None @pytest.mark.asyncio -async def test_streamed_redacted_output_guardrail_does_not_invoke_hostile_session( +async def test_streamed_session_hostile_error_after_redacted_output_guardrail_is_replaced( monkeypatch: pytest.MonkeyPatch, ) -> None: persistence_secret = "HOSTILE_SESSION_FAILURE_SECRET" monkeypatch.setattr(_debug, "DONT_LOG_MODEL_DATA", True) monkeypatch.setattr(_debug, "DONT_LOG_TOOL_DATA", False) guardrail_failed = False - post_guardrail_save_calls = 0 payload = f'{{"answer": "{_MODEL_OUTPUT_SECRET}"}}' class FailingFinalTurnSession(SimpleListSession): async def add_items(self, items: list[Any]) -> None: - nonlocal post_guardrail_save_calls if guardrail_failed: - post_guardrail_save_calls += 1 raise _HostileAttributeWriteException(persistence_secret) await super().add_items(items) @@ -1640,13 +1674,11 @@ def output_guardrail( ) result = Runner.run_streamed(agent, "go", session=FailingFinalTurnSession()) - with pytest.raises(ModelBehaviorError) as exc_info: + with pytest.raises(UserError, match="Error details are redacted.") as exc_info: async for _ in result.stream_events(): pass error = exc_info.value - assert post_guardrail_save_calls == 0 - assert error.run_data is None assert error.__cause__ is None assert error.__context__ is None assert persistence_secret not in str(error) @@ -1877,26 +1909,25 @@ def inspect_public_task(task: asyncio.Task[Any]) -> None: @pytest.mark.parametrize("streamed", [False, True]) +@pytest.mark.parametrize("failure_kind", ["exception", "cancelled", "direct_base", "group"]) @pytest.mark.asyncio -async def test_max_turns_guardrail_error_skips_final_session_write( +async def test_max_turns_recovery_session_failure_preserves_complete_redaction_boundary( monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture, streamed: bool, + failure_kind: Literal["exception", "cancelled", "direct_base", "group"], ) -> None: persistence_secret = "MAX_TURNS_SESSION_FAILURE_SECRET" fallback_secret = "MAX_TURNS_HANDLER_OUTPUT_SECRET" payload = f'{{"answer": "{_MODEL_OUTPUT_SECRET}"}}' guardrail_failed = False - post_guardrail_save_calls = 0 monkeypatch.setattr(_debug, "DONT_LOG_MODEL_DATA", True) monkeypatch.setattr(_debug, "DONT_LOG_TOOL_DATA", False) class FailingMaxTurnsSession(SimpleListSession): async def add_items(self, items: list[Any]) -> None: - nonlocal post_guardrail_save_calls if guardrail_failed: - post_guardrail_save_calls += 1 - raise AssertionError("the final Session write must not start without a verdict") + raise _persistence_failure(failure_kind, persistence_secret) await super().add_items(items) def output_guardrail( @@ -1940,12 +1971,23 @@ def output_guardrail( except BaseException as error: captured_error = error else: # pragma: no cover - raise AssertionError("the guardrail failure must propagate") + raise AssertionError("the session failure must propagate") - assert isinstance(captured_error, ModelBehaviorError) - assert post_guardrail_save_calls == 0 + assert captured_error is not None error = captured_error - assert error.run_data is None + if failure_kind == "exception": + assert isinstance(error, UserError) + elif failure_kind == "cancelled": + assert isinstance(error, asyncio.CancelledError) + elif failure_kind == "direct_base": + assert type(error) is BaseException + else: + assert isinstance(error, BaseExceptionGroup) + assert not isinstance(error, Exception) + assert {type(child) for child in error.exceptions} == { + UserError, + asyncio.CancelledError, + } error_graph = _exception_graph(error) assert error_graph @@ -2085,7 +2127,7 @@ def test_agent_runner_run_sync_detaches_all_marked_recovery_failures( @pytest.mark.parametrize("streamed", [False, True]) @pytest.mark.asyncio -async def test_max_turns_guardrail_error_does_not_build_session_exception_group( +async def test_max_turns_recovery_deep_exception_group_preserves_redaction_boundary( monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture, streamed: bool, @@ -2093,15 +2135,12 @@ async def test_max_turns_guardrail_error_does_not_build_session_exception_group( persistence_secret = "DEEP_MAX_TURNS_SESSION_FAILURE_SECRET" payload = f'{{"answer": "{_MODEL_OUTPUT_SECRET}"}}' guardrail_failed = False - post_guardrail_save_calls = 0 monkeypatch.setattr(_debug, "DONT_LOG_MODEL_DATA", True) monkeypatch.setattr(_debug, "DONT_LOG_TOOL_DATA", False) class DeepGroupSession(SimpleListSession): async def add_items(self, items: list[Any]) -> None: - nonlocal post_guardrail_save_calls if guardrail_failed: - post_guardrail_save_calls += 1 error: BaseException = asyncio.CancelledError(persistence_secret) for _ in range(1200): error = BaseExceptionGroup(persistence_secret, [error]) @@ -2149,14 +2188,23 @@ def output_guardrail( except BaseException as error: captured_error = error else: # pragma: no cover - raise AssertionError("the guardrail failure must propagate") + raise AssertionError("the deep exception group must propagate") - assert isinstance(captured_error, ModelBehaviorError) - assert captured_error.run_data is None - assert post_guardrail_save_calls == 0 - assert captured_error.__cause__ is None - assert captured_error.__context__ is None - _assert_secret_absent_from_agents_traceback(captured_error, _MODEL_OUTPUT_SECRET) + assert isinstance(captured_error, BaseExceptionGroup) + error_graph = _exception_graph(captured_error) + assert len(error_graph) == 1201 + for current in error_graph: + assert current.__cause__ is None + assert current.__context__ is None + if isinstance(current, BaseExceptionGroup): + assert current.message == "Error details are redacted." + else: + assert isinstance(current, asyncio.CancelledError) + assert persistence_secret not in str(current) + assert persistence_secret not in repr(current) + for frame_locals in _agents_traceback_frame_locals(current): + _assert_secret_absent_from_value_graph(frame_locals, persistence_secret) + _assert_secret_absent_from_value_graph(frame_locals, _MODEL_OUTPUT_SECRET) for record in caplog.records: assert persistence_secret not in logging.Formatter().format(record) @@ -2167,25 +2215,24 @@ def output_guardrail( @pytest.mark.parametrize("redacted", [False, True]) +@pytest.mark.parametrize("failure_kind", ["direct_base", "group"]) @pytest.mark.asyncio -async def test_max_turns_run_loop_guardrail_error_skips_final_session_write( +async def test_max_turns_run_loop_exception_follows_redaction_policy_for_base_exceptions( monkeypatch: pytest.MonkeyPatch, redacted: bool, + failure_kind: Literal["direct_base", "group"], ) -> None: persistence_secret = "RUN_LOOP_EXCEPTION_PERSISTENCE_SECRET" fallback_secret = "RUN_LOOP_EXCEPTION_FALLBACK_SECRET" payload = f'{{"answer": "{_MODEL_OUTPUT_SECRET}"}}' guardrail_failed = False - post_guardrail_save_calls = 0 monkeypatch.setattr(_debug, "DONT_LOG_MODEL_DATA", redacted) monkeypatch.setattr(_debug, "DONT_LOG_TOOL_DATA", False) class FailingMaxTurnsSession(SimpleListSession): async def add_items(self, items: list[Any]) -> None: - nonlocal post_guardrail_save_calls if guardrail_failed: - post_guardrail_save_calls += 1 - raise _DirectBaseException(persistence_secret) + raise _persistence_failure(failure_kind, persistence_secret) await super().add_items(items) def output_guardrail( @@ -2225,8 +2272,7 @@ def inspect_run_loop_task(task: asyncio.Task[Any]) -> None: await asyncio.wait_for(run_loop_done.wait(), timeout=1) error = result.run_loop_exception - assert isinstance(error, ModelBehaviorError) - assert post_guardrail_save_calls == 0 + assert error is not None frame_locals = _agents_traceback_frame_locals(error) if redacted: for traceback_locals in callback_frame_locals + frame_locals: @@ -2246,7 +2292,7 @@ def inspect_run_loop_task(task: asyncio.Task[Any]) -> None: assert any(fallback_secret in repr(frame) for frame in callback_frame_locals) assert frame_locals assert any(fallback_secret in repr(frame) for frame in frame_locals) - assert _MODEL_OUTPUT_SECRET in str(error) + assert persistence_secret in str(error) try: async for _ in result.stream_events(): @@ -2254,7 +2300,7 @@ def inspect_run_loop_task(task: asyncio.Task[Any]) -> None: except BaseException as streamed_error: assert streamed_error is error else: # pragma: no cover - raise AssertionError("the guardrail failure must propagate through the stream") + raise AssertionError("the session failure must propagate through the stream") @pytest.mark.parametrize( @@ -2388,21 +2434,18 @@ def test_safe_redacted_persistence_error_preserves_hybrid_cancellation() -> None @pytest.mark.parametrize("streamed", [False, True]) @pytest.mark.asyncio -async def test_max_turns_guardrail_error_preserves_diagnostic_without_session_write( +async def test_max_turns_recovery_session_failure_preserves_diagnostic_context( monkeypatch: pytest.MonkeyPatch, streamed: bool, ) -> None: persistence_secret = "DIAGNOSTIC_MAX_TURNS_SESSION_SECRET" guardrail_secret = "DIAGNOSTIC_MAX_TURNS_GUARDRAIL_SECRET" guardrail_failed = False - post_guardrail_save_calls = 0 monkeypatch.setattr(_debug, "DONT_LOG_MODEL_DATA", False) class FailingMaxTurnsSession(SimpleListSession): async def add_items(self, items: list[Any]) -> None: - nonlocal post_guardrail_save_calls if guardrail_failed: - post_guardrail_save_calls += 1 raise LookupError(persistence_secret) await super().add_items(items) @@ -2430,11 +2473,11 @@ def output_guardrail( session=session, error_handlers={"max_turns": lambda data: "fallback"}, ) - with pytest.raises(RuntimeError, match=guardrail_secret) as exc_info: + with pytest.raises(LookupError, match=persistence_secret) as exc_info: async for _ in result.stream_events(): pass else: - with pytest.raises(RuntimeError, match=guardrail_secret) as exc_info: + with pytest.raises(LookupError, match=persistence_secret) as exc_info: await Runner.run( agent, "go", @@ -2443,8 +2486,9 @@ def output_guardrail( error_handlers={"max_turns": lambda data: "fallback"}, ) - assert post_guardrail_save_calls == 0 - assert persistence_secret not in str(exc_info.value) + guardrail_error = exc_info.value.__context__ + assert isinstance(guardrail_error, RuntimeError) + assert guardrail_secret in str(guardrail_error) @pytest.mark.asyncio diff --git a/tests/test_max_turns.py b/tests/test_max_turns.py index dedb195674..02817e57cc 100644 --- a/tests/test_max_turns.py +++ b/tests/test_max_turns.py @@ -674,7 +674,7 @@ async def run_once() -> Any: saved_items = await session.get_items() saved_types = [str(item.get("type", item.get("role"))) for item in saved_items] - if outcome == "pass": + if outcome in {"pass", "error"}: assert saved_types == ["user", "message"] else: assert saved_types == ["user"] @@ -690,7 +690,7 @@ async def run_once() -> Any: if streamed: assert streamed_result is not None - expected_history_count = 1 if outcome == "pass" else 0 + expected_history_count = 1 if outcome in {"pass", "error"} else 0 assert ( len([item for item in streamed_result.new_items if isinstance(item, MessageOutputItem)]) == expected_history_count From 93e2668312ec21162a7705bc69a0443afd541731 Mon Sep 17 00:00:00 2001 From: Kazuhiro Sera Date: Wed, 19 Aug 2026 11:51:26 +0900 Subject: [PATCH 08/12] refactor: consolidate blocked output handling --- src/agents/run.py | 20 +- src/agents/run_internal/blocked_output.py | 656 +++++++++++++++++++++- src/agents/run_internal/run_loop.py | 654 +-------------------- tests/test_agent_runner.py | 8 +- 4 files changed, 681 insertions(+), 657 deletions(-) diff --git a/src/agents/run.py b/src/agents/run.py index 71f29551a4..cd880716a3 100644 --- a/src/agents/run.py +++ b/src/agents/run.py @@ -81,6 +81,17 @@ validate_session_conversation_settings, ) from .run_internal.approvals import approvals_from_step +from .run_internal.blocked_output import ( + _blocked_output_failure_items, + _BlockedOutputOwnerStarts, + _current_response_boundary, + _final_turn_items_for_persistence, + _is_terminal_tool_output_response, + _retained_items_for_blocked_response, + _sanitize_blocked_output_guardrail_results, + _should_defer_interrupted_session_items, + _validate_resumed_session_output_guardrail_safety, +) from .run_internal.error_handlers import ( attach_generic_agent_error, build_run_error_data, @@ -95,16 +106,7 @@ from .run_internal.prompt_cache_key import PromptCacheKeyResolver from .run_internal.run_grouping import resolve_run_grouping_id from .run_internal.run_loop import ( - _blocked_output_failure_items, - _BlockedOutputOwnerStarts, - _current_response_boundary, - _final_turn_items_for_persistence, - _is_terminal_tool_output_response, - _retained_items_for_blocked_response, _safe_redacted_persistence_error, - _sanitize_blocked_output_guardrail_results, - _should_defer_interrupted_session_items, - _validate_resumed_session_output_guardrail_safety, cleanup_models_after_run, finalize_max_turns_handler_output, get_all_tools, diff --git a/src/agents/run_internal/blocked_output.py b/src/agents/run_internal/blocked_output.py index a7f2f31928..5c3caf3f07 100644 --- a/src/agents/run_internal/blocked_output.py +++ b/src/agents/run_internal/blocked_output.py @@ -2,12 +2,30 @@ from __future__ import annotations -from typing import Any +import dataclasses as _dc +from collections.abc import Sequence +from typing import Any, TypeVar, cast from openai.types.responses import ResponseFunctionToolCall from openai.types.responses.response_function_tool_call import CallerDirect -from ..exceptions import AgentsException +from ..agent import Agent +from ..exceptions import ( + AgentsException, + OutputGuardrailTripwireTriggered, + UserError, + _detach_data_redacted_error_traceback, + _mark_error_data_redacted, + _prepare_data_redacted_error, +) +from ..guardrail import GuardrailFunctionOutput, OutputGuardrailResult +from ..items import ModelResponse, RunItem, ToolCallItem, ToolCallOutputItem +from ..memory import Session +from ..result import RunResultStreaming +from ..run_config import RunConfig +from ..run_state import RunState +from ..tool_guardrails import ToolGuardrailFunctionOutput, ToolOutputGuardrailResult +from .run_steps import NextStepInterruption, ProcessedResponse OUTPUT_GUARDRAIL_BLOCKED_TOOL_OUTPUT = "Output withheld by an output guardrail." @@ -128,3 +146,637 @@ def blocked_function_output_payload(raw_item: Any) -> dict[str, Any]: if restored is None: raise AgentsException("Sanitized function_call_output is not valid for replay.") return sanitized + + +_SIDE_EFFECT_ITEM_TYPES = frozenset({"tool_call_item", "tool_call_output_item"}) + + +def _sanitize_blocked_output_guardrail_results( + results: Sequence[OutputGuardrailResult], + tripwire: OutputGuardrailTripwireTriggered, +) -> list[OutputGuardrailResult]: + """Build data-free guardrail results and detach the tripwire from raw output.""" + sanitized_by_id: dict[int, OutputGuardrailResult] = {} + + def sanitize(result: OutputGuardrailResult) -> OutputGuardrailResult: + existing = sanitized_by_id.get(id(result)) + if existing is not None: + return existing + sanitized = OutputGuardrailResult( + guardrail=result.guardrail, + agent_output=OUTPUT_GUARDRAIL_BLOCKED_TOOL_OUTPUT, + agent=result.agent, + output=GuardrailFunctionOutput( + output_info=None, + tripwire_triggered=result.output.tripwire_triggered, + ), + ) + sanitized_by_id[id(result)] = sanitized + return sanitized + + sanitized_results = [sanitize(result) for result in results] + object.__setattr__(tripwire, "guardrail_result", sanitize(tripwire.guardrail_result)) + _mark_error_data_redacted(tripwire) + _detach_data_redacted_error_traceback(tripwire) + return sanitized_results + + +@_dc.dataclass(frozen=True) +class _CurrentResponseBoundary: + """A current-response suffix proven only by lifecycle position or object identity.""" + + items: tuple[RunItem, ...] + processed_items: tuple[RunItem, ...] + generated_start: int | None + session_start: int | None + proven: bool + + +@_dc.dataclass(frozen=True) +class _BlockedOutputSnapshot: + """Prepared data-free replacements for one complete current response.""" + + items: tuple[RunItem, ...] + processed_items: tuple[RunItem, ...] + model_response: ModelResponse | None + + +@_dc.dataclass(frozen=True) +class _BlockedOutputOwnerPlan: + """Prebuilt trusted-owner assignments for application or emergency cleanup.""" + + assignments: tuple[tuple[Any, str, Any], ...] + + +@_dc.dataclass(frozen=True) +class _BlockedOutputOwnerStarts: + """Owner-specific current-response starts captured at trusted lifecycle boundaries.""" + + nonstreamed_session_items: int | None = None + run_state_generated_items: int | None = None + run_state_session_items: int | None = None + run_state_model_responses: int | None = None + run_state_tool_output_guardrail_results: int | None = None + streamed_new_items: int | None = None + streamed_model_input_items: int | None = None + streamed_raw_responses: int | None = None + streamed_tool_output_guardrail_results: int | None = None + + +@_dc.dataclass(frozen=True) +class _BlockedOutputOwnerPrefixes: + """Accepted owner prefixes allocated before any blocked-output replacement begins.""" + + run_state_generated_items: list[RunItem] + run_state_session_items: list[RunItem] + run_state_model_responses: list[ModelResponse] + run_state_tool_output_guardrail_results: list[ToolOutputGuardrailResult] + streamed_new_items: list[RunItem] + streamed_model_input_items: list[RunItem] + streamed_raw_responses: list[ModelResponse] + streamed_tool_output_guardrail_results: list[ToolOutputGuardrailResult] + + +_OwnerItemT = TypeVar("_OwnerItemT") + + +def _has_output_guardrails(agent: Agent[Any], run_config: RunConfig) -> bool: + return bool(agent.output_guardrails or run_config.output_guardrails) + + +def _should_defer_interrupted_session_items( + agent: Agent[Any], + run_config: RunConfig, +) -> bool: + """Keep pre-verdict approval state in RunState instead of durable Session history.""" + return _has_output_guardrails(agent, run_config) + + +def _validate_resumed_session_output_guardrail_safety( + *, + agent: Agent[Any], + run_config: RunConfig, + session: Session | None, + run_state: RunState[Any] | None, +) -> None: + """Reject approval resumes whose current-response boundary is not structurally provable.""" + del session + if run_state is None or not _has_output_guardrails(agent, run_config): + return + if not isinstance(run_state._current_step, NextStepInterruption): + return + if run_state._current_turn_persisted_item_count > 0: + raise UserError( + "Cannot resume an approval checkpoint with output guardrails after current-turn " + "items were persisted. Start a new run from safe input." + ) + boundary = _current_response_boundary( + (), + run_state._last_processed_response, + run_state, + ) + if boundary.proven: + return + raise UserError( + "Cannot resume a serialized approval checkpoint with output guardrails because the " + "current response boundary cannot be proven. Start a new run from safe input." + ) + + +def _identity_sequence_start( + container: Sequence[RunItem], + sequence: Sequence[RunItem], +) -> int | None: + if not sequence or len(sequence) > len(container): + return None + for start in range(len(container) - len(sequence) + 1): + if all(container[start + offset] is item for offset, item in enumerate(sequence)): + return start + return None + + +def _current_response_boundary( + new_items: Sequence[RunItem], + processed_response: ProcessedResponse | None, + run_state: RunState[Any] | None, +) -> _CurrentResponseBoundary: + """Collect one response using only SDK lifecycle position and exact object identity.""" + processed_items = tuple(processed_response.new_items) if processed_response is not None else () + supplied_items = tuple(new_items) + supplied_start = _identity_sequence_start(supplied_items, processed_items) + response_items = ( + supplied_items[supplied_start:] if supplied_start is not None else supplied_items + ) + generated_start = None + session_start = None + proven = run_state is None or not processed_items or supplied_start is not None + suffixes: list[RunItem] = [] + if run_state is not None: + anchor_items = processed_items or response_items + if anchor_items: + generated_start = _identity_sequence_start(run_state._generated_items, anchor_items) + session_start = _identity_sequence_start(run_state._session_items, anchor_items) + if generated_start is not None: + suffixes.extend(run_state._generated_items[generated_start:]) + proven = True + if session_start is not None: + suffixes.extend(run_state._session_items[session_start:]) + proven = True + if ( + not supplied_items + and generated_start is None + and session_start is None + and run_state._current_turn == 1 + ): + generated_start = 0 + session_start = 0 + suffixes.extend(run_state._generated_items) + suffixes.extend(run_state._session_items) + proven = True + + current_items: list[RunItem] = [] + seen: set[int] = set() + for item in (*processed_items, *suffixes, *response_items): + if id(item) in seen: + continue + seen.add(id(item)) + current_items.append(item) + return _CurrentResponseBoundary( + items=tuple(current_items), + processed_items=processed_items, + generated_start=generated_start, + session_start=session_start, + proven=proven, + ) + + +def _current_response_items_for_persistence( + items: Sequence[RunItem], + processed_response: ProcessedResponse | None, + run_state: RunState[Any] | None = None, +) -> list[RunItem]: + """Return the complete current response or fail before using an ambiguous boundary.""" + boundary = _current_response_boundary(items, processed_response, run_state) + if not boundary.proven: + raise UserError( + "Cannot persist an ambiguous resumed response with output guardrails. " + "Start a new run from safe input." + ) + return list(boundary.items) + + +def _final_turn_items_for_persistence( + items: Sequence[RunItem], + processed_response: ProcessedResponse | None, + run_state: RunState[Any] | None, + agent: Agent[Any], + run_config: RunConfig, +) -> list[RunItem]: + """Use released resumed suffix persistence unless output guardrails defer the response.""" + if not _has_output_guardrails(agent, run_config): + return list(items) + return _current_response_items_for_persistence(items, processed_response, run_state) + + +def _is_terminal_tool_output_response( + items: Sequence[RunItem], + processed_response: ProcessedResponse | None, + run_state: RunState[Any] | None = None, +) -> bool: + """Return whether the structurally owned current response produced a tool final output.""" + boundary = _current_response_boundary(items, processed_response, run_state) + return boundary.proven and any(isinstance(item, ToolCallOutputItem) for item in boundary.items) + + +def _prepare_blocked_output_snapshot( + boundary: _CurrentResponseBoundary, + model_response: ModelResponse | None, +) -> _BlockedOutputSnapshot: + """Build an allowlist-only function call/output snapshot before changing live state.""" + current_items = list(boundary.items) + if any(item.type == "reasoning_item" for item in current_items): + raise AgentsException("Cannot sanitize a response containing reasoning items.") + retained_indexes = { + index for index, item in enumerate(current_items) if item.type in _SIDE_EFFECT_ITEM_TYPES + } + replacements: dict[int, RunItem] = {} + calls_by_id: dict[str, int] = {} + outputs_by_id: dict[str, int] = {} + for index in sorted(retained_indexes): + item = current_items[index] + if isinstance(item, ToolCallItem): + payload = blocked_function_call_payload(item.raw_item) + call_id = cast(str, payload["call_id"]) + if call_id in calls_by_id: + raise AgentsException("Cannot sanitize duplicate function calls.") + calls_by_id[call_id] = index + replacements[index] = ToolCallItem( + agent=item.agent, + raw_item=cast(Any, payload), + description=item.description, + title=item.title, + tool_origin=item.tool_origin, + _resolved_tool_name=item._resolved_tool_name, + ) + elif isinstance(item, ToolCallOutputItem): + payload = blocked_function_output_payload(item.raw_item) + call_id = cast(str, payload["call_id"]) + if call_id in outputs_by_id: + raise AgentsException("Cannot sanitize duplicate function outputs.") + outputs_by_id[call_id] = index + replacements[index] = ToolCallOutputItem( + agent=item.agent, + raw_item=cast(Any, payload), + output=OUTPUT_GUARDRAIL_BLOCKED_TOOL_OUTPUT, + tool_origin=item.tool_origin, + custom_data=None, + ) + else: + raise AgentsException("Cannot sanitize an unsupported side-effect item.") + + if not outputs_by_id or set(outputs_by_id) - set(calls_by_id): + raise AgentsException("Cannot sanitize an incomplete function call/output batch.") + retained_indexes = { + index + for call_id in outputs_by_id + for index in (calls_by_id[call_id], outputs_by_id[call_id]) + } + retained_items = tuple(replacements[index] for index in sorted(retained_indexes)) + processed_indexes = {id(item): index for index, item in enumerate(current_items)} + retained_processed_items = tuple( + replacements.get(processed_indexes[id(item)], item) + for item in boundary.processed_items + if processed_indexes.get(id(item)) in retained_indexes + ) + sanitized_response = None + if model_response is not None: + sanitized_response = ModelResponse( + output=cast(Any, [item.raw_item for item in retained_processed_items]), + usage=model_response.usage, + response_id=model_response.response_id, + request_id=model_response.request_id, + raw_usage=model_response.raw_usage, + ) + return _BlockedOutputSnapshot( + items=retained_items, + processed_items=retained_processed_items, + model_response=sanitized_response, + ) + + +def _blocked_output_owner_prefix(items: list[_OwnerItemT], start: int | None) -> list[_OwnerItemT]: + """Copy a structurally captured prefix without consulting item values or identities.""" + if start is None or start < 0 or start > len(items): + return [] + return list.__getitem__(items, slice(0, start)) + + +def _blocked_output_failure_items( + items: list[RunItem], + retained_items: Sequence[RunItem], + owner_starts: _BlockedOutputOwnerStarts, +) -> list[RunItem]: + """Build the non-streamed accepted prefix plus the data-free current response.""" + return [ + *_blocked_output_owner_prefix(items, owner_starts.nonstreamed_session_items), + *retained_items, + ] + + +def _prepare_blocked_output_owner_prefixes( + run_state: RunState[Any] | None, + streamed_result: RunResultStreaming | None, + owner_starts: _BlockedOutputOwnerStarts, +) -> _BlockedOutputOwnerPrefixes: + """Allocate every accepted owner prefix before snapshot application begins.""" + return _BlockedOutputOwnerPrefixes( + run_state_generated_items=( + _blocked_output_owner_prefix( + run_state._generated_items, + owner_starts.run_state_generated_items, + ) + if run_state is not None + else [] + ), + run_state_session_items=( + _blocked_output_owner_prefix( + run_state._session_items, + owner_starts.run_state_session_items, + ) + if run_state is not None + else [] + ), + run_state_model_responses=( + _blocked_output_owner_prefix( + run_state._model_responses, + owner_starts.run_state_model_responses, + ) + if run_state is not None + else [] + ), + run_state_tool_output_guardrail_results=( + _blocked_output_owner_prefix( + run_state._tool_output_guardrail_results, + owner_starts.run_state_tool_output_guardrail_results, + ) + if run_state is not None + else [] + ), + streamed_new_items=( + _blocked_output_owner_prefix( + streamed_result.new_items, + owner_starts.streamed_new_items, + ) + if streamed_result is not None + else [] + ), + streamed_model_input_items=( + _blocked_output_owner_prefix( + streamed_result._model_input_items, + owner_starts.streamed_model_input_items, + ) + if streamed_result is not None + else [] + ), + streamed_raw_responses=( + _blocked_output_owner_prefix( + streamed_result.raw_responses, + owner_starts.streamed_raw_responses, + ) + if streamed_result is not None + else [] + ), + streamed_tool_output_guardrail_results=( + _blocked_output_owner_prefix( + streamed_result.tool_output_guardrail_results, + owner_starts.streamed_tool_output_guardrail_results, + ) + if streamed_result is not None + else [] + ), + ) + + +def _prepare_blocked_output_cleanup_plan( + run_state: RunState[Any] | None, + streamed_result: RunResultStreaming | None, + prefixes: _BlockedOutputOwnerPrefixes, +) -> _BlockedOutputOwnerPlan: + """Prepare accepted-prefix cleanup containers before snapshot application begins.""" + assignments: list[tuple[Any, str, Any]] = [] + if run_state is not None: + assignments.extend( + [ + (run_state, "_generated_items", prefixes.run_state_generated_items), + (run_state, "_session_items", prefixes.run_state_session_items), + (run_state, "_model_responses", prefixes.run_state_model_responses), + (run_state, "_last_processed_response", None), + (run_state, "_current_step", None), + (run_state, "_generated_items_last_processed_marker", None), + ( + run_state, + "_tool_output_guardrail_results", + prefixes.run_state_tool_output_guardrail_results, + ), + ] + ) + if streamed_result is not None: + assignments.extend( + [ + (streamed_result, "new_items", prefixes.streamed_new_items), + (streamed_result, "raw_responses", prefixes.streamed_raw_responses), + ( + streamed_result, + "_model_input_items", + prefixes.streamed_model_input_items, + ), + (streamed_result, "_last_processed_response", None), + ( + streamed_result, + "tool_output_guardrail_results", + prefixes.streamed_tool_output_guardrail_results, + ), + ] + ) + return _BlockedOutputOwnerPlan(assignments=tuple(assignments)) + + +def _sever_blocked_output_replay_graph(cleanup_plan: _BlockedOutputOwnerPlan) -> None: + """Best-effort leaf cleanup using only containers allocated before application.""" + for owner, field, value in cleanup_plan.assignments: + try: + object.__setattr__(owner, field, value) + except BaseException: + continue + + +def _data_free_tool_output_guardrail_results( + results: Sequence[ToolOutputGuardrailResult], +) -> tuple[ToolOutputGuardrailResult, ...]: + """Rebuild current-turn tool guardrail results without retaining caller output data.""" + replacements: list[ToolOutputGuardrailResult] = [] + try: + for result in results: + if not isinstance(result, ToolOutputGuardrailResult): + return () + replacements.append( + ToolOutputGuardrailResult( + guardrail=object.__getattribute__(result, "guardrail"), + output=ToolGuardrailFunctionOutput( + output_info=OUTPUT_GUARDRAIL_BLOCKED_TOOL_OUTPUT, + behavior={"type": "allow"}, + ), + ) + ) + except Exception: + return () + return tuple(replacements) + + +def _prepare_blocked_output_owner_plan( + boundary: _CurrentResponseBoundary, + snapshot: _BlockedOutputSnapshot | None, + model_response: ModelResponse | None, + run_state: RunState[Any] | None, + streamed_result: RunResultStreaming | None, + prefixes: _BlockedOutputOwnerPrefixes, + cleanup_plan: _BlockedOutputOwnerPlan, +) -> _BlockedOutputOwnerPlan: + """Build every owner replacement before applying any of them.""" + safe_items = list(snapshot.items) if snapshot is not None else [] + safe_response = snapshot.model_response if snapshot is not None else None + assignments: list[tuple[Any, str, Any]] = [] + safe_tool_output_guardrail_results: tuple[ToolOutputGuardrailResult, ...] = () + if streamed_result is not None: + public_results = streamed_result.tool_output_guardrail_results + current_results = list.__getitem__( + public_results, + slice(len(prefixes.streamed_tool_output_guardrail_results), None), + ) + safe_tool_output_guardrail_results = _data_free_tool_output_guardrail_results( + current_results + ) + public_safe_results = [ + *prefixes.streamed_tool_output_guardrail_results, + *safe_tool_output_guardrail_results, + ] + else: + public_safe_results = [] + + if run_state is not None: + if boundary.proven: + responses = [ + *prefixes.run_state_model_responses, + *( + [safe_response] + if model_response is not None and safe_response is not None + else [] + ), + ] + if streamed_result is not None: + run_state_safe_results = [ + *prefixes.run_state_tool_output_guardrail_results, + *safe_tool_output_guardrail_results, + ] + else: + run_state_safe_results = list(prefixes.run_state_tool_output_guardrail_results) + assignments.extend( + [ + ( + run_state, + "_generated_items", + [*prefixes.run_state_generated_items, *safe_items], + ), + ( + run_state, + "_session_items", + [*prefixes.run_state_session_items, *safe_items], + ), + (run_state, "_model_responses", responses), + (run_state, "_last_processed_response", None), + (run_state, "_current_step", None), + (run_state, "_generated_items_last_processed_marker", None), + (run_state, "_tool_output_guardrail_results", run_state_safe_results), + ] + ) + else: + return cleanup_plan + + if streamed_result is not None: + responses = [ + *prefixes.streamed_raw_responses, + *([safe_response] if model_response is not None and safe_response is not None else []), + ] + assignments.extend( + [ + ( + streamed_result, + "new_items", + [*prefixes.streamed_new_items, *safe_items], + ), + ( + streamed_result, + "_model_input_items", + [*prefixes.streamed_model_input_items, *safe_items], + ), + (streamed_result, "raw_responses", responses), + (streamed_result, "_last_processed_response", None), + ( + streamed_result, + "tool_output_guardrail_results", + public_safe_results, + ), + ] + ) + return _BlockedOutputOwnerPlan(assignments=tuple(assignments)) + + +def _apply_blocked_output_owner_plan(plan: _BlockedOutputOwnerPlan) -> None: + """Apply only values that were fully constructed before the first owner swap.""" + for owner, field, value in plan.assignments: + object.__setattr__(owner, field, value) + + +def _retained_items_for_blocked_response( + items: list[RunItem], + model_response: ModelResponse | None, + run_state: RunState[Any] | None = None, + processed_response: ProcessedResponse | None = None, + streamed_result: RunResultStreaming | None = None, + owner_starts: _BlockedOutputOwnerStarts | None = None, +) -> list[RunItem]: + """Return a complete data-free response or discard the entire unsupported suffix.""" + boundary = _current_response_boundary(items, processed_response, run_state) + prefixes = _prepare_blocked_output_owner_prefixes( + run_state, + streamed_result, + owner_starts if owner_starts is not None else _BlockedOutputOwnerStarts(), + ) + cleanup_plan = _prepare_blocked_output_cleanup_plan(run_state, streamed_result, prefixes) + snapshot: _BlockedOutputSnapshot | None = None + try: + if boundary.proven: + snapshot = _prepare_blocked_output_snapshot(boundary, model_response) + except Exception: + snapshot = None + except BaseException: + _sever_blocked_output_replay_graph(cleanup_plan) + raise + try: + owner_plan = _prepare_blocked_output_owner_plan( + boundary, + snapshot, + model_response, + run_state, + streamed_result, + prefixes, + cleanup_plan, + ) + _apply_blocked_output_owner_plan(owner_plan) + except Exception as error: + _sever_blocked_output_replay_graph(cleanup_plan) + raise _prepare_data_redacted_error(error) from None + except BaseException: + _sever_blocked_output_replay_graph(cleanup_plan) + raise + return list(snapshot.items) if snapshot is not None else [] diff --git a/src/agents/run_internal/run_loop.py b/src/agents/run_internal/run_loop.py index f6ea75ea7d..e7c41c175a 100644 --- a/src/agents/run_internal/run_loop.py +++ b/src/agents/run_internal/run_loop.py @@ -7,7 +7,7 @@ import asyncio import dataclasses as _dc -from collections.abc import Awaitable, Callable, Sequence +from collections.abc import Awaitable, Callable from contextlib import aclosing from functools import partial from typing import Any, TypeVar, cast @@ -45,7 +45,7 @@ _mark_error_data_redacted, _prepare_data_redacted_error, ) -from ..guardrail import GuardrailFunctionOutput, OutputGuardrailResult +from ..guardrail import OutputGuardrailResult from ..handoffs import Handoff from ..items import ( InputItem, @@ -53,8 +53,6 @@ ModelResponse, RunItem, ToolApprovalItem, - ToolCallItem, - ToolCallOutputItem, TResponseInputItem, ) from ..lifecycle import RunHooks @@ -86,7 +84,6 @@ Tool, dispose_resolved_computers, ) -from ..tool_guardrails import ToolGuardrailFunctionOutput, ToolOutputGuardrailResult from ..tracing import Span, SpanError, agent_span, get_current_trace, task_span, turn_span from ..tracing.config import include_task_and_turn_spans from ..tracing.model_tracing import get_model_tracing_impl @@ -110,9 +107,15 @@ ) from .approvals import approvals_from_step from .blocked_output import ( - OUTPUT_GUARDRAIL_BLOCKED_TOOL_OUTPUT as _OUTPUT_GUARDRAIL_BLOCKED_TOOL_OUTPUT, - blocked_function_call_payload as _blocked_function_call_payload, - blocked_function_output_payload as _blocked_function_output_payload, + OUTPUT_GUARDRAIL_BLOCKED_TOOL_OUTPUT, + _BlockedOutputOwnerStarts, + _current_response_boundary, + _final_turn_items_for_persistence, + _is_terminal_tool_output_response, + _retained_items_for_blocked_response, + _sanitize_blocked_output_guardrail_results, + _should_defer_interrupted_session_items, + _validate_resumed_session_output_guardrail_safety, ) from .error_handlers import ( attach_generic_agent_error, @@ -282,6 +285,7 @@ "input_guardrail_tripwire_triggered_for_stream", ] +_OUTPUT_GUARDRAIL_BLOCKED_TOOL_OUTPUT = OUTPUT_GUARDRAIL_BLOCKED_TOOL_OUTPUT _STREAM_EVENT_ITEM_OCCURRENCE_KEY = "_agents_stream_event_item_occurrence_key" @@ -467,595 +471,6 @@ async def _run_output_guardrails_for_stream( raise -_SIDE_EFFECT_ITEM_TYPES = frozenset({"tool_call_item", "tool_call_output_item"}) - - -def _sanitize_blocked_output_guardrail_results( - results: Sequence[OutputGuardrailResult], - tripwire: OutputGuardrailTripwireTriggered, -) -> list[OutputGuardrailResult]: - """Build data-free guardrail results and detach the tripwire from raw output.""" - sanitized_by_id: dict[int, OutputGuardrailResult] = {} - - def sanitize(result: OutputGuardrailResult) -> OutputGuardrailResult: - existing = sanitized_by_id.get(id(result)) - if existing is not None: - return existing - sanitized = OutputGuardrailResult( - guardrail=result.guardrail, - agent_output=_OUTPUT_GUARDRAIL_BLOCKED_TOOL_OUTPUT, - agent=result.agent, - output=GuardrailFunctionOutput( - output_info=None, - tripwire_triggered=result.output.tripwire_triggered, - ), - ) - sanitized_by_id[id(result)] = sanitized - return sanitized - - sanitized_results = [sanitize(result) for result in results] - object.__setattr__(tripwire, "guardrail_result", sanitize(tripwire.guardrail_result)) - _mark_error_data_redacted(tripwire) - _detach_data_redacted_error_traceback(tripwire) - return sanitized_results - - -@_dc.dataclass(frozen=True) -class _CurrentResponseBoundary: - """A current-response suffix proven only by lifecycle position or object identity.""" - - items: tuple[RunItem, ...] - processed_items: tuple[RunItem, ...] - generated_start: int | None - session_start: int | None - proven: bool - - -@_dc.dataclass(frozen=True) -class _BlockedOutputSnapshot: - """Prepared data-free replacements for one complete current response.""" - - items: tuple[RunItem, ...] - processed_items: tuple[RunItem, ...] - model_response: ModelResponse | None - - -@_dc.dataclass(frozen=True) -class _BlockedOutputOwnerPlan: - """Prebuilt trusted-owner assignments for application or emergency cleanup.""" - - assignments: tuple[tuple[Any, str, Any], ...] - - -@_dc.dataclass(frozen=True) -class _BlockedOutputOwnerStarts: - """Owner-specific current-response starts captured at trusted lifecycle boundaries.""" - - nonstreamed_session_items: int | None = None - run_state_generated_items: int | None = None - run_state_session_items: int | None = None - run_state_model_responses: int | None = None - run_state_tool_output_guardrail_results: int | None = None - streamed_new_items: int | None = None - streamed_model_input_items: int | None = None - streamed_raw_responses: int | None = None - streamed_tool_output_guardrail_results: int | None = None - - -@_dc.dataclass(frozen=True) -class _BlockedOutputOwnerPrefixes: - """Accepted owner prefixes allocated before any blocked-output replacement begins.""" - - run_state_generated_items: list[RunItem] - run_state_session_items: list[RunItem] - run_state_model_responses: list[ModelResponse] - run_state_tool_output_guardrail_results: list[ToolOutputGuardrailResult] - streamed_new_items: list[RunItem] - streamed_model_input_items: list[RunItem] - streamed_raw_responses: list[ModelResponse] - streamed_tool_output_guardrail_results: list[ToolOutputGuardrailResult] - - -_OwnerItemT = TypeVar("_OwnerItemT") - - -def _has_output_guardrails(agent: Agent[Any], run_config: RunConfig) -> bool: - return bool(agent.output_guardrails or run_config.output_guardrails) - - -def _should_defer_interrupted_session_items( - agent: Agent[Any], - run_config: RunConfig, -) -> bool: - """Keep pre-verdict approval state in RunState instead of durable Session history.""" - return _has_output_guardrails(agent, run_config) - - -def _validate_resumed_session_output_guardrail_safety( - *, - agent: Agent[Any], - run_config: RunConfig, - session: Session | None, - run_state: RunState[Any] | None, -) -> None: - """Reject approval resumes whose current-response boundary is not structurally provable.""" - del session - if run_state is None or not _has_output_guardrails(agent, run_config): - return - if not isinstance(run_state._current_step, NextStepInterruption): - return - if run_state._current_turn_persisted_item_count > 0: - raise UserError( - "Cannot resume an approval checkpoint with output guardrails after current-turn " - "items were persisted. Start a new run from safe input." - ) - boundary = _current_response_boundary( - (), - run_state._last_processed_response, - run_state, - ) - if boundary.proven: - return - raise UserError( - "Cannot resume a serialized approval checkpoint with output guardrails because the " - "current response boundary cannot be proven. Start a new run from safe input." - ) - - -def _identity_sequence_start( - container: Sequence[RunItem], - sequence: Sequence[RunItem], -) -> int | None: - if not sequence or len(sequence) > len(container): - return None - for start in range(len(container) - len(sequence) + 1): - if all(container[start + offset] is item for offset, item in enumerate(sequence)): - return start - return None - - -def _current_response_boundary( - new_items: Sequence[RunItem], - processed_response: ProcessedResponse | None, - run_state: RunState[Any] | None, -) -> _CurrentResponseBoundary: - """Collect one response using only SDK lifecycle position and exact object identity.""" - processed_items = tuple(processed_response.new_items) if processed_response is not None else () - supplied_items = tuple(new_items) - supplied_start = _identity_sequence_start(supplied_items, processed_items) - response_items = ( - supplied_items[supplied_start:] if supplied_start is not None else supplied_items - ) - generated_start = None - session_start = None - proven = run_state is None or not processed_items or supplied_start is not None - suffixes: list[RunItem] = [] - if run_state is not None: - anchor_items = processed_items or response_items - if anchor_items: - generated_start = _identity_sequence_start(run_state._generated_items, anchor_items) - session_start = _identity_sequence_start(run_state._session_items, anchor_items) - if generated_start is not None: - suffixes.extend(run_state._generated_items[generated_start:]) - proven = True - if session_start is not None: - suffixes.extend(run_state._session_items[session_start:]) - proven = True - if ( - not supplied_items - and generated_start is None - and session_start is None - and run_state._current_turn == 1 - ): - generated_start = 0 - session_start = 0 - suffixes.extend(run_state._generated_items) - suffixes.extend(run_state._session_items) - proven = True - - current_items: list[RunItem] = [] - seen: set[int] = set() - for item in (*processed_items, *suffixes, *response_items): - if id(item) in seen: - continue - seen.add(id(item)) - current_items.append(item) - return _CurrentResponseBoundary( - items=tuple(current_items), - processed_items=processed_items, - generated_start=generated_start, - session_start=session_start, - proven=proven, - ) - - -def _current_response_items_for_persistence( - items: Sequence[RunItem], - processed_response: ProcessedResponse | None, - run_state: RunState[Any] | None = None, -) -> list[RunItem]: - """Return the complete current response or fail before using an ambiguous boundary.""" - boundary = _current_response_boundary(items, processed_response, run_state) - if not boundary.proven: - raise UserError( - "Cannot persist an ambiguous resumed response with output guardrails. " - "Start a new run from safe input." - ) - return list(boundary.items) - - -def _final_turn_items_for_persistence( - items: Sequence[RunItem], - processed_response: ProcessedResponse | None, - run_state: RunState[Any] | None, - agent: Agent[Any], - run_config: RunConfig, -) -> list[RunItem]: - """Use released resumed suffix persistence unless output guardrails defer the response.""" - if not _has_output_guardrails(agent, run_config): - return list(items) - return _current_response_items_for_persistence(items, processed_response, run_state) - - -def _is_terminal_tool_output_response( - items: Sequence[RunItem], - processed_response: ProcessedResponse | None, - run_state: RunState[Any] | None = None, -) -> bool: - """Return whether the structurally owned current response produced a tool final output.""" - boundary = _current_response_boundary(items, processed_response, run_state) - return boundary.proven and any(isinstance(item, ToolCallOutputItem) for item in boundary.items) - - -def _prepare_blocked_output_snapshot( - boundary: _CurrentResponseBoundary, - model_response: ModelResponse | None, -) -> _BlockedOutputSnapshot: - """Build an allowlist-only function call/output snapshot before changing live state.""" - current_items = list(boundary.items) - if any(item.type == "reasoning_item" for item in current_items): - raise AgentsException("Cannot sanitize a response containing reasoning items.") - retained_indexes = { - index for index, item in enumerate(current_items) if item.type in _SIDE_EFFECT_ITEM_TYPES - } - replacements: dict[int, RunItem] = {} - calls_by_id: dict[str, int] = {} - outputs_by_id: dict[str, int] = {} - for index in sorted(retained_indexes): - item = current_items[index] - if isinstance(item, ToolCallItem): - payload = _blocked_function_call_payload(item.raw_item) - call_id = cast(str, payload["call_id"]) - if call_id in calls_by_id: - raise AgentsException("Cannot sanitize duplicate function calls.") - calls_by_id[call_id] = index - replacements[index] = ToolCallItem( - agent=item.agent, - raw_item=cast(Any, payload), - description=item.description, - title=item.title, - tool_origin=item.tool_origin, - _resolved_tool_name=item._resolved_tool_name, - ) - elif isinstance(item, ToolCallOutputItem): - payload = _blocked_function_output_payload(item.raw_item) - call_id = cast(str, payload["call_id"]) - if call_id in outputs_by_id: - raise AgentsException("Cannot sanitize duplicate function outputs.") - outputs_by_id[call_id] = index - replacements[index] = ToolCallOutputItem( - agent=item.agent, - raw_item=cast(Any, payload), - output=_OUTPUT_GUARDRAIL_BLOCKED_TOOL_OUTPUT, - tool_origin=item.tool_origin, - custom_data=None, - ) - else: - raise AgentsException("Cannot sanitize an unsupported side-effect item.") - - if not outputs_by_id or set(outputs_by_id) - set(calls_by_id): - raise AgentsException("Cannot sanitize an incomplete function call/output batch.") - retained_indexes = { - index - for call_id in outputs_by_id - for index in (calls_by_id[call_id], outputs_by_id[call_id]) - } - retained_items = tuple(replacements[index] for index in sorted(retained_indexes)) - processed_indexes = {id(item): index for index, item in enumerate(current_items)} - retained_processed_items = tuple( - replacements.get(processed_indexes[id(item)], item) - for item in boundary.processed_items - if processed_indexes.get(id(item)) in retained_indexes - ) - sanitized_response = None - if model_response is not None: - sanitized_response = ModelResponse( - output=cast(Any, [item.raw_item for item in retained_processed_items]), - usage=model_response.usage, - response_id=model_response.response_id, - request_id=model_response.request_id, - raw_usage=model_response.raw_usage, - ) - return _BlockedOutputSnapshot( - items=retained_items, - processed_items=retained_processed_items, - model_response=sanitized_response, - ) - - -def _blocked_output_owner_prefix(items: list[_OwnerItemT], start: int | None) -> list[_OwnerItemT]: - """Copy a structurally captured prefix without consulting item values or identities.""" - if start is None or start < 0 or start > len(items): - return [] - return list.__getitem__(items, slice(0, start)) - - -def _blocked_output_failure_items( - items: list[RunItem], - retained_items: Sequence[RunItem], - owner_starts: _BlockedOutputOwnerStarts, -) -> list[RunItem]: - """Build the non-streamed accepted prefix plus the data-free current response.""" - return [ - *_blocked_output_owner_prefix(items, owner_starts.nonstreamed_session_items), - *retained_items, - ] - - -def _prepare_blocked_output_owner_prefixes( - run_state: RunState[Any] | None, - streamed_result: RunResultStreaming | None, - owner_starts: _BlockedOutputOwnerStarts, -) -> _BlockedOutputOwnerPrefixes: - """Allocate every accepted owner prefix before snapshot application begins.""" - return _BlockedOutputOwnerPrefixes( - run_state_generated_items=( - _blocked_output_owner_prefix( - run_state._generated_items, - owner_starts.run_state_generated_items, - ) - if run_state is not None - else [] - ), - run_state_session_items=( - _blocked_output_owner_prefix( - run_state._session_items, - owner_starts.run_state_session_items, - ) - if run_state is not None - else [] - ), - run_state_model_responses=( - _blocked_output_owner_prefix( - run_state._model_responses, - owner_starts.run_state_model_responses, - ) - if run_state is not None - else [] - ), - run_state_tool_output_guardrail_results=( - _blocked_output_owner_prefix( - run_state._tool_output_guardrail_results, - owner_starts.run_state_tool_output_guardrail_results, - ) - if run_state is not None - else [] - ), - streamed_new_items=( - _blocked_output_owner_prefix( - streamed_result.new_items, - owner_starts.streamed_new_items, - ) - if streamed_result is not None - else [] - ), - streamed_model_input_items=( - _blocked_output_owner_prefix( - streamed_result._model_input_items, - owner_starts.streamed_model_input_items, - ) - if streamed_result is not None - else [] - ), - streamed_raw_responses=( - _blocked_output_owner_prefix( - streamed_result.raw_responses, - owner_starts.streamed_raw_responses, - ) - if streamed_result is not None - else [] - ), - streamed_tool_output_guardrail_results=( - _blocked_output_owner_prefix( - streamed_result.tool_output_guardrail_results, - owner_starts.streamed_tool_output_guardrail_results, - ) - if streamed_result is not None - else [] - ), - ) - - -def _prepare_blocked_output_cleanup_plan( - run_state: RunState[Any] | None, - streamed_result: RunResultStreaming | None, - prefixes: _BlockedOutputOwnerPrefixes, -) -> _BlockedOutputOwnerPlan: - """Prepare accepted-prefix cleanup containers before snapshot application begins.""" - assignments: list[tuple[Any, str, Any]] = [] - if run_state is not None: - assignments.extend( - [ - (run_state, "_generated_items", prefixes.run_state_generated_items), - (run_state, "_session_items", prefixes.run_state_session_items), - (run_state, "_model_responses", prefixes.run_state_model_responses), - (run_state, "_last_processed_response", None), - (run_state, "_current_step", None), - (run_state, "_generated_items_last_processed_marker", None), - ( - run_state, - "_tool_output_guardrail_results", - prefixes.run_state_tool_output_guardrail_results, - ), - ] - ) - if streamed_result is not None: - assignments.extend( - [ - (streamed_result, "new_items", prefixes.streamed_new_items), - (streamed_result, "raw_responses", prefixes.streamed_raw_responses), - ( - streamed_result, - "_model_input_items", - prefixes.streamed_model_input_items, - ), - (streamed_result, "_last_processed_response", None), - ( - streamed_result, - "tool_output_guardrail_results", - prefixes.streamed_tool_output_guardrail_results, - ), - ] - ) - return _BlockedOutputOwnerPlan(assignments=tuple(assignments)) - - -def _sever_blocked_output_replay_graph(cleanup_plan: _BlockedOutputOwnerPlan) -> None: - """Best-effort leaf cleanup using only containers allocated before application.""" - for owner, field, value in cleanup_plan.assignments: - try: - object.__setattr__(owner, field, value) - except BaseException: - continue - - -def _data_free_tool_output_guardrail_results( - results: Sequence[ToolOutputGuardrailResult], -) -> tuple[ToolOutputGuardrailResult, ...]: - """Rebuild current-turn tool guardrail results without retaining caller output data.""" - replacements: list[ToolOutputGuardrailResult] = [] - try: - for result in results: - if not isinstance(result, ToolOutputGuardrailResult): - return () - replacements.append( - ToolOutputGuardrailResult( - guardrail=object.__getattribute__(result, "guardrail"), - output=ToolGuardrailFunctionOutput( - output_info=_OUTPUT_GUARDRAIL_BLOCKED_TOOL_OUTPUT, - behavior={"type": "allow"}, - ), - ) - ) - except Exception: - return () - return tuple(replacements) - - -def _prepare_blocked_output_owner_plan( - boundary: _CurrentResponseBoundary, - snapshot: _BlockedOutputSnapshot | None, - model_response: ModelResponse | None, - run_state: RunState[Any] | None, - streamed_result: RunResultStreaming | None, - prefixes: _BlockedOutputOwnerPrefixes, - cleanup_plan: _BlockedOutputOwnerPlan, -) -> _BlockedOutputOwnerPlan: - """Build every owner replacement before applying any of them.""" - safe_items = list(snapshot.items) if snapshot is not None else [] - safe_response = snapshot.model_response if snapshot is not None else None - assignments: list[tuple[Any, str, Any]] = [] - safe_tool_output_guardrail_results: tuple[ToolOutputGuardrailResult, ...] = () - if streamed_result is not None: - public_results = streamed_result.tool_output_guardrail_results - current_results = list.__getitem__( - public_results, - slice(len(prefixes.streamed_tool_output_guardrail_results), None), - ) - safe_tool_output_guardrail_results = _data_free_tool_output_guardrail_results( - current_results - ) - public_safe_results = [ - *prefixes.streamed_tool_output_guardrail_results, - *safe_tool_output_guardrail_results, - ] - else: - public_safe_results = [] - - if run_state is not None: - if boundary.proven: - responses = [ - *prefixes.run_state_model_responses, - *( - [safe_response] - if model_response is not None and safe_response is not None - else [] - ), - ] - if streamed_result is not None: - run_state_safe_results = [ - *prefixes.run_state_tool_output_guardrail_results, - *safe_tool_output_guardrail_results, - ] - else: - run_state_safe_results = list(prefixes.run_state_tool_output_guardrail_results) - assignments.extend( - [ - ( - run_state, - "_generated_items", - [*prefixes.run_state_generated_items, *safe_items], - ), - ( - run_state, - "_session_items", - [*prefixes.run_state_session_items, *safe_items], - ), - (run_state, "_model_responses", responses), - (run_state, "_last_processed_response", None), - (run_state, "_current_step", None), - (run_state, "_generated_items_last_processed_marker", None), - (run_state, "_tool_output_guardrail_results", run_state_safe_results), - ] - ) - else: - return cleanup_plan - - if streamed_result is not None: - responses = [ - *prefixes.streamed_raw_responses, - *([safe_response] if model_response is not None and safe_response is not None else []), - ] - assignments.extend( - [ - ( - streamed_result, - "new_items", - [*prefixes.streamed_new_items, *safe_items], - ), - ( - streamed_result, - "_model_input_items", - [*prefixes.streamed_model_input_items, *safe_items], - ), - (streamed_result, "raw_responses", responses), - (streamed_result, "_last_processed_response", None), - ( - streamed_result, - "tool_output_guardrail_results", - public_safe_results, - ), - ] - ) - return _BlockedOutputOwnerPlan(assignments=tuple(assignments)) - - -def _apply_blocked_output_owner_plan(plan: _BlockedOutputOwnerPlan) -> None: - """Apply only values that were fully constructed before the first owner swap.""" - for owner, field, value in plan.assignments: - object.__setattr__(owner, field, value) - - def _retained_items_for_blocked_output( items: list[RunItem], model_response: ModelResponse | None = None, @@ -1067,51 +482,6 @@ def _retained_items_for_blocked_output( ) -def _retained_items_for_blocked_response( - items: list[RunItem], - model_response: ModelResponse | None, - run_state: RunState[Any] | None = None, - processed_response: ProcessedResponse | None = None, - streamed_result: RunResultStreaming | None = None, - owner_starts: _BlockedOutputOwnerStarts | None = None, -) -> list[RunItem]: - """Return a complete data-free response or discard the entire unsupported suffix.""" - boundary = _current_response_boundary(items, processed_response, run_state) - prefixes = _prepare_blocked_output_owner_prefixes( - run_state, - streamed_result, - owner_starts if owner_starts is not None else _BlockedOutputOwnerStarts(), - ) - cleanup_plan = _prepare_blocked_output_cleanup_plan(run_state, streamed_result, prefixes) - snapshot: _BlockedOutputSnapshot | None = None - try: - if boundary.proven: - snapshot = _prepare_blocked_output_snapshot(boundary, model_response) - except Exception: - snapshot = None - except BaseException: - _sever_blocked_output_replay_graph(cleanup_plan) - raise - try: - owner_plan = _prepare_blocked_output_owner_plan( - boundary, - snapshot, - model_response, - run_state, - streamed_result, - prefixes, - cleanup_plan, - ) - _apply_blocked_output_owner_plan(owner_plan) - except Exception as error: - _sever_blocked_output_replay_graph(cleanup_plan) - raise _prepare_data_redacted_error(error) from None - except BaseException: - _sever_blocked_output_replay_graph(cleanup_plan) - raise - return list(snapshot.items) if snapshot is not None else [] - - async def _finalize_streamed_final_output( *, streamed_result: RunResultStreaming, diff --git a/tests/test_agent_runner.py b/tests/test_agent_runner.py index c56a45f86c..f7eb002c88 100644 --- a/tests/test_agent_runner.py +++ b/tests/test_agent_runner.py @@ -74,7 +74,7 @@ from agents.result import RunResultStreaming from agents.run import AgentRunner, get_default_agent_runner, set_default_agent_runner from agents.run_config import _default_trace_include_sensitive_data -from agents.run_internal import run_loop +from agents.run_internal import blocked_output, run_loop from agents.run_internal.agent_bindings import bind_public_agent from agents.run_internal.agent_runner_helpers import build_resumed_stream_debug_extra from agents.run_internal.items import ( @@ -496,7 +496,7 @@ def test_blocked_snapshot_cancellation_severs_replay_graph_and_propagates( def cancel_preparation(_raw_item: Any) -> dict[str, Any]: raise cancellation - monkeypatch.setattr(run_loop, "_blocked_function_output_payload", cancel_preparation) + monkeypatch.setattr(blocked_output, "blocked_function_output_payload", cancel_preparation) with pytest.raises(asyncio.CancelledError) as exc_info: run_loop._retained_items_for_blocked_response( @@ -601,7 +601,7 @@ def fail_application(plan: Any) -> None: object.__setattr__(owner, field, value) raise application_error - monkeypatch.setattr(run_loop, "_apply_blocked_output_owner_plan", fail_application) + monkeypatch.setattr(blocked_output, "_apply_blocked_output_owner_plan", fail_application) with pytest.raises(KeyboardInterrupt) as exc_info: run_loop._retained_items_for_blocked_response( @@ -666,7 +666,7 @@ def test_blocked_snapshot_application_exception_becomes_fixed_error( def fail_application(_plan: Any) -> None: raise ValueError("application-secret") - monkeypatch.setattr(run_loop, "_apply_blocked_output_owner_plan", fail_application) + monkeypatch.setattr(blocked_output, "_apply_blocked_output_owner_plan", fail_application) with pytest.raises(RuntimeError) as exc_info: run_loop._retained_items_for_blocked_response( From 669c461511386860e94439f6fefeda00bbf37430 Mon Sep 17 00:00:00 2001 From: Kazuhiro Sera Date: Wed, 19 Aug 2026 12:28:06 +0900 Subject: [PATCH 09/12] fix: preserve redacted tool guardrail verdicts --- src/agents/run_internal/blocked_output.py | 27 +++++++++++++++++++---- tests/test_agent_runner_streamed.py | 27 +++++++++++++++++++++-- 2 files changed, 48 insertions(+), 6 deletions(-) diff --git a/src/agents/run_internal/blocked_output.py b/src/agents/run_internal/blocked_output.py index 5c3caf3f07..ee5b8d76e8 100644 --- a/src/agents/run_internal/blocked_output.py +++ b/src/agents/run_internal/blocked_output.py @@ -619,13 +619,32 @@ def _data_free_tool_output_guardrail_results( for result in results: if not isinstance(result, ToolOutputGuardrailResult): return () + original_output = object.__getattribute__(result, "output") + behavior = object.__getattribute__(original_output, "behavior") + if type(behavior) is not dict: + return () + behavior_type = _exact_dict_field(behavior, "type") + if type(behavior_type) is not str: + return () + if str.__eq__(behavior_type, "allow") is True: + sanitized_output = ToolGuardrailFunctionOutput.allow( + output_info=OUTPUT_GUARDRAIL_BLOCKED_TOOL_OUTPUT, + ) + elif str.__eq__(behavior_type, "reject_content") is True: + sanitized_output = ToolGuardrailFunctionOutput.reject_content( + message=OUTPUT_GUARDRAIL_BLOCKED_TOOL_OUTPUT, + output_info=OUTPUT_GUARDRAIL_BLOCKED_TOOL_OUTPUT, + ) + elif str.__eq__(behavior_type, "raise_exception") is True: + sanitized_output = ToolGuardrailFunctionOutput.raise_exception( + output_info=OUTPUT_GUARDRAIL_BLOCKED_TOOL_OUTPUT, + ) + else: + return () replacements.append( ToolOutputGuardrailResult( guardrail=object.__getattribute__(result, "guardrail"), - output=ToolGuardrailFunctionOutput( - output_info=OUTPUT_GUARDRAIL_BLOCKED_TOOL_OUTPUT, - behavior={"type": "allow"}, - ), + output=sanitized_output, ) ) except Exception: diff --git a/tests/test_agent_runner_streamed.py b/tests/test_agent_runner_streamed.py index cd0f1a495f..ddecd5e3c0 100644 --- a/tests/test_agent_runner_streamed.py +++ b/tests/test_agent_runner_streamed.py @@ -3472,14 +3472,23 @@ async def test_streamed_run_reports_tool_guardrail_results(): assert result.tool_output_guardrail_results[0].output.output_info == "output-checked" +@pytest.mark.parametrize("tool_guardrail_behavior", ["allow", "reject_content"]) @pytest.mark.asyncio -async def test_streamed_trip_replaces_current_tool_output_guardrail_results() -> None: +async def test_streamed_trip_replaces_current_tool_output_guardrail_results( + tool_guardrail_behavior: str, +) -> None: """A copied terminal tool result is replaced in public and RunState guardrail results.""" original_outputs: list[ToolGuardrailFunctionOutput] = [] @tool_output_guardrail def retain_output(data: ToolOutputGuardrailData) -> ToolGuardrailFunctionOutput: - output = ToolGuardrailFunctionOutput.allow(output_info=data.output) + if tool_guardrail_behavior == "reject_content": + output = ToolGuardrailFunctionOutput.reject_content( + message=f"Rejected sensitive tool output: {data.output}", + output_info=data.output, + ) + else: + output = ToolGuardrailFunctionOutput.allow(output_info=data.output) original_outputs.append(output) return output @@ -3518,6 +3527,11 @@ def reject_output( public_output = result.tool_output_guardrail_results[1].output assert public_output is not original_outputs[0] assert public_output.output_info == run_loop._OUTPUT_GUARDRAIL_BLOCKED_TOOL_OUTPUT + assert public_output.behavior["type"] == tool_guardrail_behavior + if public_output.behavior["type"] == "reject_content": + assert public_output.behavior["message"] == run_loop._OUTPUT_GUARDRAIL_BLOCKED_TOOL_OUTPUT + assert original_outputs[0].behavior["type"] == "reject_content" + assert "blocked-secret" in original_outputs[0].behavior["message"] assert result._state is not None # The caller-added public result was never owned by RunState, so only the current # data-free result is added to that owner. @@ -3525,6 +3539,15 @@ def reject_output( state_output = result._state._tool_output_guardrail_results[0].output assert state_output is public_output assert state_output.output_info == run_loop._OUTPUT_GUARDRAIL_BLOCKED_TOOL_OUTPUT + assert state_output.behavior["type"] == tool_guardrail_behavior + serialized_state = result.to_state().to_json() + serialized_results = serialized_state["tool_output_guardrail_results"] + assert serialized_results[0]["output"]["behavior"]["type"] == "allow" + serialized_behavior = serialized_results[-1]["output"]["behavior"] + assert serialized_behavior["type"] == tool_guardrail_behavior + if tool_guardrail_behavior == "reject_content": + assert serialized_behavior["message"] == run_loop._OUTPUT_GUARDRAIL_BLOCKED_TOOL_OUTPUT + assert "blocked-secret" not in json.dumps(serialized_state) @pytest.mark.asyncio From cfaffef005e459e14961ba893df8f3e68e9935e2 Mon Sep 17 00:00:00 2001 From: Kazuhiro Sera Date: Wed, 19 Aug 2026 14:07:12 +0900 Subject: [PATCH 10/12] fix: preserve accepted history after resumed guardrail trips --- src/agents/run.py | 13 +++ src/agents/run_internal/blocked_output.py | 25 ++++- src/agents/run_internal/run_loop.py | 40 +++++++- tests/test_agent_runner.py | 116 ++++++++++++++++++++++ 4 files changed, 191 insertions(+), 3 deletions(-) diff --git a/src/agents/run.py b/src/agents/run.py index cd880716a3..6e3f517dca 100644 --- a/src/agents/run.py +++ b/src/agents/run.py @@ -86,10 +86,12 @@ _BlockedOutputOwnerStarts, _current_response_boundary, _final_turn_items_for_persistence, + _has_output_guardrails, _is_terminal_tool_output_response, _retained_items_for_blocked_response, _sanitize_blocked_output_guardrail_results, _should_defer_interrupted_session_items, + _synchronize_accepted_run_state, _validate_resumed_session_output_guardrail_safety, ) from .run_internal.error_handlers import ( @@ -1536,6 +1538,17 @@ async def _save_max_turns_handler_output( except Exception: last_saved_input_snapshot_for_rewind = None + if run_state is not None and _has_output_guardrails(current_agent, run_config): + _synchronize_accepted_run_state( + run_state, + generated_items=generated_items, + session_items=session_items, + model_responses=model_responses, + tool_input_guardrail_results=tool_input_guardrail_results, + tool_output_guardrail_results=tool_output_guardrail_results, + current_turn=current_turn, + ) + blocked_output_owner_starts = _BlockedOutputOwnerStarts( nonstreamed_session_items=len(session_items), run_state_generated_items=( diff --git a/src/agents/run_internal/blocked_output.py b/src/agents/run_internal/blocked_output.py index ee5b8d76e8..89a6420441 100644 --- a/src/agents/run_internal/blocked_output.py +++ b/src/agents/run_internal/blocked_output.py @@ -24,7 +24,11 @@ from ..result import RunResultStreaming from ..run_config import RunConfig from ..run_state import RunState -from ..tool_guardrails import ToolGuardrailFunctionOutput, ToolOutputGuardrailResult +from ..tool_guardrails import ( + ToolGuardrailFunctionOutput, + ToolInputGuardrailResult, + ToolOutputGuardrailResult, +) from .run_steps import NextStepInterruption, ProcessedResponse OUTPUT_GUARDRAIL_BLOCKED_TOOL_OUTPUT = "Output withheld by an output guardrail." @@ -244,6 +248,25 @@ def _has_output_guardrails(agent: Agent[Any], run_config: RunConfig) -> bool: return bool(agent.output_guardrails or run_config.output_guardrails) +def _synchronize_accepted_run_state( + run_state: RunState[Any], + *, + generated_items: Sequence[RunItem], + session_items: Sequence[RunItem], + model_responses: Sequence[ModelResponse], + tool_input_guardrail_results: Sequence[ToolInputGuardrailResult], + tool_output_guardrail_results: Sequence[ToolOutputGuardrailResult], + current_turn: int, +) -> None: + """Capture accepted run history before a guardrail-owned model response begins.""" + run_state._generated_items = list(generated_items) + run_state._session_items = list(session_items) + run_state._model_responses = list(model_responses) + run_state._tool_input_guardrail_results = list(tool_input_guardrail_results) + run_state._tool_output_guardrail_results = list(tool_output_guardrail_results) + run_state._current_turn = current_turn + + def _should_defer_interrupted_session_items( agent: Agent[Any], run_config: RunConfig, diff --git a/src/agents/run_internal/run_loop.py b/src/agents/run_internal/run_loop.py index e7c41c175a..6b94fae5e5 100644 --- a/src/agents/run_internal/run_loop.py +++ b/src/agents/run_internal/run_loop.py @@ -84,6 +84,7 @@ Tool, dispose_resolved_computers, ) +from ..tool_guardrails import ToolInputGuardrailResult, ToolOutputGuardrailResult from ..tracing import Span, SpanError, agent_span, get_current_trace, task_span, turn_span from ..tracing.config import include_task_and_turn_spans from ..tracing.model_tracing import get_model_tracing_impl @@ -111,10 +112,12 @@ _BlockedOutputOwnerStarts, _current_response_boundary, _final_turn_items_for_persistence, + _has_output_guardrails, _is_terminal_tool_output_response, _retained_items_for_blocked_response, _sanitize_blocked_output_guardrail_results, _should_defer_interrupted_session_items, + _synchronize_accepted_run_state, _validate_resumed_session_output_guardrail_safety, ) from .error_handlers import ( @@ -800,6 +803,9 @@ async def _persist_stream_input_if_needed( def _accumulate_tool_guardrail_results( streamed_result: RunResultStreaming, turn_result: SingleStepResult, + *, + accepted_input_results: list[ToolInputGuardrailResult], + accepted_output_results: list[ToolOutputGuardrailResult], ) -> None: """Carry a turn's tool guardrail results onto the streamed result. @@ -812,6 +818,9 @@ def _accumulate_tool_guardrail_results( streamed_result.tool_output_guardrail_results = ( streamed_result.tool_output_guardrail_results + turn_result.tool_output_guardrail_results ) + if isinstance(turn_result.next_step, NextStepRunAgain | NextStepHandoff): + accepted_input_results.extend(turn_result.tool_input_guardrail_results) + accepted_output_results.extend(turn_result.tool_output_guardrail_results) async def _finalize_streamed_interruption( @@ -972,6 +981,12 @@ def _sync_conversation_tracking_from_tracker() -> None: current_turn = run_state._current_turn else: current_turn = 0 + accepted_tool_input_guardrail_results = ( + list(run_state._tool_input_guardrail_results) if run_state is not None else [] + ) + accepted_tool_output_guardrail_results = ( + list(run_state._tool_output_guardrail_results) if run_state is not None else [] + ) should_run_agent_start_hooks = True tool_use_tracker = AgentToolUseTracker() if run_state is not None: @@ -1323,7 +1338,12 @@ async def _save_max_turns_items( # but skips a resumed turn that loops back to the model, so a guardrail that # re-runs for the same tool call on resume is not counted twice. if not isinstance(turn_result.next_step, NextStepRunAgain): - _accumulate_tool_guardrail_results(streamed_result, turn_result) + _accumulate_tool_guardrail_results( + streamed_result, + turn_result, + accepted_input_results=accepted_tool_input_guardrail_results, + accepted_output_results=accepted_tool_output_guardrail_results, + ) if isinstance(turn_result.next_step, NextStepInterruption): await _finalize_streamed_interruption( @@ -1652,6 +1672,17 @@ def _record_max_turns_handler_output( ) ) try: + if run_state is not None and _has_output_guardrails(current_agent, run_config): + _synchronize_accepted_run_state( + run_state, + generated_items=streamed_result._model_input_items, + session_items=streamed_result.new_items, + model_responses=streamed_result.raw_responses, + tool_input_guardrail_results=accepted_tool_input_guardrail_results, + tool_output_guardrail_results=accepted_tool_output_guardrail_results, + current_turn=current_turn, + ) + blocked_output_owner_starts = _BlockedOutputOwnerStarts( run_state_generated_items=( len(run_state._generated_items) if run_state is not None else None @@ -1746,7 +1777,12 @@ def _record_max_turns_handler_output( streamed_result.raw_responses = streamed_result.raw_responses + [ turn_result.model_response ] - _accumulate_tool_guardrail_results(streamed_result, turn_result) + _accumulate_tool_guardrail_results( + streamed_result, + turn_result, + accepted_input_results=accepted_tool_input_guardrail_results, + accepted_output_results=accepted_tool_output_guardrail_results, + ) input_before_turn_rewrite = streamed_result.input streamed_result.input = turn_result.original_input if isinstance(turn_result.next_step, NextStepHandoff): diff --git a/tests/test_agent_runner.py b/tests/test_agent_runner.py index f7eb002c88..f822fefea0 100644 --- a/tests/test_agent_runner.py +++ b/tests/test_agent_runner.py @@ -31,6 +31,7 @@ HandoffInputData, InputGuardrail, InputGuardrailTripwireTriggered, + MaxTurnsExceeded, ModelBehaviorError, ModelRetryAdvice, ModelRetrySettings, @@ -48,12 +49,14 @@ ToolGuardrailFunctionOutput, ToolInputGuardrailData, ToolNameCollisionPolicy, + ToolOutputGuardrailData, ToolTimeoutError, UserError, handoff, retry_policies, tool_input_guardrail, tool_namespace, + tool_output_guardrail, ) from agents._tool_identity import resolve_tool_name_collisions from agents.agent import ToolsToFinalOutputResult @@ -766,6 +769,119 @@ async def capture_memory_payload( assert "reasoning-current" not in serialized_memory_items +@pytest.mark.parametrize("streamed", [False, True], ids=["non-streamed", "streamed"]) +@pytest.mark.parametrize("handoff_turn", [False, True], ids=["run-again", "handoff"]) +@pytest.mark.asyncio +async def test_resumed_trip_preserves_accepted_turns_and_turn_budget( + streamed: bool, + handoff_turn: bool, +) -> None: + side_effects: list[str] = [] + + @tool_input_guardrail + def record_accepted_input(_data: ToolInputGuardrailData) -> ToolGuardrailFunctionOutput: + return ToolGuardrailFunctionOutput.allow(output_info="accepted-input-audit") + + @tool_output_guardrail + def record_accepted_output(_data: ToolOutputGuardrailData) -> ToolGuardrailFunctionOutput: + return ToolGuardrailFunctionOutput.allow(output_info="accepted-output-audit") + + @function_tool(name_override="approval_tool", needs_approval=True) + def approval_tool() -> str: + side_effects.append("approved") + return "approved-output" + + @function_tool( + name_override="accepted_tool", + tool_input_guardrails=[record_accepted_input], + tool_output_guardrails=[record_accepted_output], + ) + def accepted_tool() -> str: + side_effects.append("accepted") + return "accepted-output" + + @function_tool(name_override="terminal_tool") + def terminal_tool() -> str: + side_effects.append("terminal") + return "rejected-secret" + + def reject_output( + _context: RunContextWrapper[Any], + _agent: Agent[Any], + _output: Any, + ) -> GuardrailFunctionOutput: + return GuardrailFunctionOutput(output_info=None, tripwire_triggered=True) + + model = ScriptedModel() + target = Agent( + name="target", + model=model, + tools=[terminal_tool], + tool_use_behavior={"stop_at_tool_names": ["terminal_tool"]}, + output_guardrails=[OutputGuardrail(guardrail_function=reject_output)], + ) + agent = Agent( + name="source", + model=model, + tools=[approval_tool, accepted_tool, terminal_tool], + handoffs=[target] if handoff_turn else [], + tool_use_behavior={"stop_at_tool_names": ["terminal_tool"]}, + output_guardrails=( + [] if handoff_turn else [OutputGuardrail(guardrail_function=reject_output)] + ), + ) + accepted_response = [get_function_tool_call("accepted_tool", "{}", call_id="accepted-call")] + if handoff_turn: + accepted_response.append(get_handoff_tool_call(target)) + model.extend( + [ + [get_function_tool_call("approval_tool", "{}", call_id="approved-call")], + accepted_response, + [get_function_tool_call("terminal_tool", "{}", call_id="terminal-call")], + ] + ) + + interrupted = await Runner.run(agent, "run approved tools", max_turns=3) + state = interrupted.to_state() + state.approve(interrupted.interruptions[0]) + + with pytest.raises(OutputGuardrailTripwireTriggered): + if streamed: + result = Runner.run_streamed(agent, state) + async for _ in result.stream_events(): + pass + else: + await Runner.run(agent, state) + + assert side_effects == ["approved", "accepted", "terminal"] + assert state._current_turn == 3 + assert len(state._model_responses) == 3 + assert [result.output.output_info for result in state._tool_input_guardrail_results] == [ + "accepted-input-audit" + ] + assert [result.output.output_info for result in state._tool_output_guardrail_results] == [ + "accepted-output-audit" + ] + for items in (state._generated_items, state._session_items): + outputs = [item for item in items if isinstance(item, ToolCallOutputItem)] + assert [item.output for item in outputs] == [ + "approved-output", + "accepted-output", + run_loop._OUTPUT_GUARDRAIL_BLOCKED_TOOL_OUTPUT, + ] + + serialized_state = json.dumps(state.to_json()) + assert "approved-output" in serialized_state + assert "accepted-output" in serialized_state + assert "accepted-input-audit" in serialized_state + assert "accepted-output-audit" in serialized_state + assert "rejected-secret" not in serialized_state + + with pytest.raises(MaxTurnsExceeded): + await Runner.run(agent, state) + assert side_effects == ["approved", "accepted", "terminal"] + + @pytest.mark.asyncio async def test_non_streamed_trip_uses_safe_items_for_sandbox_memory_after_session_failure( monkeypatch: pytest.MonkeyPatch, From 6aec076e74e7273b5bfbb6c604f5a07eda4a98a8 Mon Sep 17 00:00:00 2001 From: Kazuhiro Sera Date: Wed, 19 Aug 2026 14:43:25 +0900 Subject: [PATCH 11/12] fix: restore released guardrail lifecycle compatibility --- src/agents/run.py | 42 +++- .../run_internal/agent_runner_helpers.py | 5 + src/agents/run_internal/blocked_output.py | 71 +++---- src/agents/run_internal/run_loop.py | 19 +- tests/test_agent_runner_streamed.py | 183 ++++++++++++++++-- 5 files changed, 262 insertions(+), 58 deletions(-) diff --git a/src/agents/run.py b/src/agents/run.py index 6e3f517dca..9d6621e66c 100644 --- a/src/agents/run.py +++ b/src/agents/run.py @@ -1222,6 +1222,12 @@ def _mark_response_hooks_started() -> None: ) if isinstance(turn_result.next_step, NextStepFinalOutput): + if run_state is not None and _has_output_guardrails( + current_agent, run_config + ): + run_state._tool_output_guardrail_results = list( + tool_output_guardrail_results + ) current_processed_response = ( turn_result.processed_response if turn_result.processed_response is not None @@ -1284,12 +1290,21 @@ def _mark_response_hooks_started() -> None: persistence_error ) from None raise - except (Exception, asyncio.CancelledError): - if not _is_terminal_tool_output_response( + except (Exception, asyncio.CancelledError) as guardrail_error: + if not isinstance( + guardrail_error, asyncio.CancelledError + ) or not _is_terminal_tool_output_response( turn_session_items, current_processed_response, run_state, ): + final_turn_items = _final_turn_items_for_persistence( + turn_session_items, + current_processed_response, + run_state, + current_agent, + run_config, + ) await save_final_turn_items_after_guardrails( session=session, run_state=run_state, @@ -1297,7 +1312,7 @@ def _mark_response_hooks_started() -> None: input_guardrail_results=( _attempt_input_guardrail_results() ), - items=turn_session_items, + items=final_turn_items, response_id=turn_result.model_response.response_id, store=store_setting, wrapper=context_wrapper, @@ -1826,6 +1841,12 @@ async def _save_max_turns_handler_output( try: if isinstance(turn_result.next_step, NextStepFinalOutput): + if run_state is not None and _has_output_guardrails( + current_agent, run_config + ): + run_state._tool_output_guardrail_results = list( + tool_output_guardrail_results + ) output_guardrail_result_start = len(output_guardrail_results) try: await run_output_guardrails( @@ -1881,12 +1902,21 @@ async def _save_max_turns_handler_output( persistence_error ) from None raise - except (Exception, asyncio.CancelledError): - if not _is_terminal_tool_output_response( + except (Exception, asyncio.CancelledError) as guardrail_error: + if not isinstance( + guardrail_error, asyncio.CancelledError + ) or not _is_terminal_tool_output_response( turn_session_items, turn_result.processed_response, run_state, ): + final_turn_items = _final_turn_items_for_persistence( + turn_session_items, + turn_result.processed_response, + run_state, + current_agent, + run_config, + ) await save_final_turn_items_after_guardrails( session=session, run_state=run_state, @@ -1894,7 +1924,7 @@ async def _save_max_turns_handler_output( input_guardrail_results=( _attempt_input_guardrail_results() ), - items=items_to_save_turn, + items=final_turn_items, response_id=turn_result.model_response.response_id, store=store_setting, wrapper=context_wrapper, diff --git a/src/agents/run_internal/agent_runner_helpers.py b/src/agents/run_internal/agent_runner_helpers.py index 6c564230d3..be1d976724 100644 --- a/src/agents/run_internal/agent_runner_helpers.py +++ b/src/agents/run_internal/agent_runner_helpers.py @@ -14,6 +14,7 @@ from ..items import ModelResponse, RunItem, ToolApprovalItem, TResponseInputItem from ..memory import Session from ..models.openai_agent_registration import add_openai_harness_id_to_metadata +from ..models.openai_chatcompletions import OpenAIChatCompletionsModel from ..result import RunResult from ..run_config import ReasoningItemIdPolicy, RunConfig from ..run_context import RunContextWrapper, TContext @@ -42,6 +43,7 @@ ) from .session_persistence import save_result_to_session, save_resumed_turn_items from .tool_use_tracker import AgentToolUseTracker, serialize_tool_use_tracker +from .turn_preparation import get_model __all__ = [ "apply_resumed_conversation_settings", @@ -271,6 +273,9 @@ def validate_output_guardrails_with_server_managed_conversation( return if not agent.output_guardrails and not run_config.output_guardrails: return + if isinstance(get_model(agent, run_config), OpenAIChatCompletionsModel): + # Chat Completions owns its released warn-and-ignore or strict rejection behavior. + return raise UserError( "Output guardrails cannot be combined with conversation_id, previous_response_id, " "or auto_previous_response_id because rejected output cannot be removed from " diff --git a/src/agents/run_internal/blocked_output.py b/src/agents/run_internal/blocked_output.py index 89a6420441..7011965843 100644 --- a/src/agents/run_internal/blocked_output.py +++ b/src/agents/run_internal/blocked_output.py @@ -283,27 +283,28 @@ def _validate_resumed_session_output_guardrail_safety( run_state: RunState[Any] | None, ) -> None: """Reject approval resumes whose current-response boundary is not structurally provable.""" - del session if run_state is None or not _has_output_guardrails(agent, run_config): return if not isinstance(run_state._current_step, NextStepInterruption): return - if run_state._current_turn_persisted_item_count > 0: - raise UserError( - "Cannot resume an approval checkpoint with output guardrails after current-turn " - "items were persisted. Start a new run from safe input." - ) boundary = _current_response_boundary( (), run_state._last_processed_response, run_state, ) - if boundary.proven: - return - raise UserError( - "Cannot resume a serialized approval checkpoint with output guardrails because the " - "current response boundary cannot be proven. Start a new run from safe input." - ) + if not boundary.proven: + raise UserError( + "Cannot resume a serialized approval checkpoint with output guardrails because the " + "current response boundary cannot be proven. Start a new run from safe input." + ) + if run_state._current_turn_persisted_item_count > 0: + if session is not None: + raise UserError( + "Cannot resume an approval checkpoint with output guardrails after current-turn " + "items were persisted. Start a new run from safe input." + ) + # A detached Session cannot contribute its old persisted prefix to this run. + run_state._current_turn_persisted_item_count = 0 def _identity_sequence_start( @@ -345,17 +346,19 @@ def _current_response_boundary( if session_start is not None: suffixes.extend(run_state._session_items[session_start:]) proven = True - if ( - not supplied_items - and generated_start is None - and session_start is None - and run_state._current_turn == 1 - ): - generated_start = 0 - session_start = 0 - suffixes.extend(run_state._generated_items) - suffixes.extend(run_state._session_items) - proven = True + if generated_start is None and session_start is None and run_state._current_turn == 1: + current_response_prefix = tuple(run_state._generated_items[: len(processed_items)]) + if len(current_response_prefix) == len(processed_items) and all( + type(actual) is type(expected) + for actual, expected in zip(current_response_prefix, processed_items, strict=False) + ): + # Serialization rebuilds item identities, but turn one has no accepted prefix. + processed_items = current_response_prefix + generated_start = 0 + session_start = 0 + suffixes.extend(run_state._generated_items) + suffixes.extend(run_state._session_items) + proven = True current_items: list[RunItem] = [] seen: set[int] = set() @@ -688,16 +691,21 @@ def _prepare_blocked_output_owner_plan( safe_items = list(snapshot.items) if snapshot is not None else [] safe_response = snapshot.model_response if snapshot is not None else None assignments: list[tuple[Any, str, Any]] = [] - safe_tool_output_guardrail_results: tuple[ToolOutputGuardrailResult, ...] = () if streamed_result is not None: public_results = streamed_result.tool_output_guardrail_results current_results = list.__getitem__( public_results, slice(len(prefixes.streamed_tool_output_guardrail_results), None), ) - safe_tool_output_guardrail_results = _data_free_tool_output_guardrail_results( - current_results + elif run_state is not None: + current_results = list.__getitem__( + run_state._tool_output_guardrail_results, + slice(len(prefixes.run_state_tool_output_guardrail_results), None), ) + else: + current_results = [] + safe_tool_output_guardrail_results = _data_free_tool_output_guardrail_results(current_results) + if streamed_result is not None: public_safe_results = [ *prefixes.streamed_tool_output_guardrail_results, *safe_tool_output_guardrail_results, @@ -715,13 +723,10 @@ def _prepare_blocked_output_owner_plan( else [] ), ] - if streamed_result is not None: - run_state_safe_results = [ - *prefixes.run_state_tool_output_guardrail_results, - *safe_tool_output_guardrail_results, - ] - else: - run_state_safe_results = list(prefixes.run_state_tool_output_guardrail_results) + run_state_safe_results = [ + *prefixes.run_state_tool_output_guardrail_results, + *safe_tool_output_guardrail_results, + ] assignments.extend( [ ( diff --git a/src/agents/run_internal/run_loop.py b/src/agents/run_internal/run_loop.py index 6b94fae5e5..6125e56d8f 100644 --- a/src/agents/run_internal/run_loop.py +++ b/src/agents/run_internal/run_loop.py @@ -554,17 +554,18 @@ async def _finalize_streamed_final_output( raise safe_error from None raise except Exception as guardrail_error: - if _is_terminal_tool_output_response( - items, - processed_response, - streamed_result._state, - ): - raise guardrail_error_is_redacted = _is_error_data_redacted(guardrail_error) if guardrail_error_is_redacted: _detach_data_redacted_error_traceback(guardrail_error) try: - await save_items(items, response_id, store_setting) + final_turn_items = _final_turn_items_for_persistence( + items, + processed_response, + streamed_result._state, + agent, + run_config, + ) + await save_items(final_turn_items, response_id, store_setting) except BaseException as persistence_error: if guardrail_error_is_redacted: safe_persistence_error = _safe_redacted_persistence_error(persistence_error) @@ -977,6 +978,10 @@ def _sync_conversation_tracking_from_tracker() -> None: session=session, run_state=run_state if is_resumed_state else None, ) + if run_state is not None and session is None: + streamed_result._current_turn_persisted_item_count = ( + run_state._current_turn_persisted_item_count + ) if run_state is not None: current_turn = run_state._current_turn else: diff --git a/tests/test_agent_runner_streamed.py b/tests/test_agent_runner_streamed.py index ddecd5e3c0..44246f25d8 100644 --- a/tests/test_agent_runner_streamed.py +++ b/tests/test_agent_runner_streamed.py @@ -32,6 +32,7 @@ ModelBehaviorError, ModelRetrySettings, ModelSettings, + OpenAIChatCompletionsModel, OpenAIResponsesWSModel, OutputGuardrail, OutputGuardrailTripwireTriggered, @@ -2242,15 +2243,23 @@ async def test_tool() -> str: @pytest.mark.parametrize("mode", ["non_streamed", "streamed"]) -@pytest.mark.parametrize("tripwire", [False, True], ids=["passes", "trips"]) +@pytest.mark.parametrize("outcome", ["passes", "trips", "error"]) @pytest.mark.asyncio async def test_resumed_approved_tool_final_persists_complete_post_verdict_batch( mode: str, - tripwire: bool, + outcome: str, ) -> None: - guardrail_state = {"tripwire": tripwire} + guardrail_state = {"outcome": outcome} - @function_tool(name_override="approval_tool", needs_approval=True) + @tool_output_guardrail + def record_tool_output(data: ToolOutputGuardrailData) -> ToolGuardrailFunctionOutput: + return ToolGuardrailFunctionOutput.allow(output_info=data.output) + + @function_tool( + name_override="approval_tool", + needs_approval=True, + tool_output_guardrails=[record_tool_output], + ) def approval_tool() -> str: return "approved-result" @@ -2259,9 +2268,11 @@ def output_guardrail( _agent: Agent[Any], _output: Any, ) -> GuardrailFunctionOutput: + if guardrail_state["outcome"] == "error": + raise RuntimeError("guardrail failed") return GuardrailFunctionOutput( output_info=None, - tripwire_triggered=guardrail_state["tripwire"], + tripwire_triggered=guardrail_state["outcome"] == "trips", ) model = ScriptedModel() @@ -2287,9 +2298,15 @@ async def run_once(input_value: Any) -> Any: state = first.to_state() state.approve(first.interruptions[0]) - if tripwire: + if outcome == "trips": with pytest.raises(OutputGuardrailTripwireTriggered): await run_once(state) + assert [result.output.output_info for result in state._tool_output_guardrail_results] == [ + run_loop._OUTPUT_GUARDRAIL_BLOCKED_TOOL_OUTPUT + ] + elif outcome == "error": + with pytest.raises(RuntimeError, match="guardrail failed"): + await run_once(state) else: resumed = await run_once(state) assert resumed.final_output == "approved-result" @@ -2309,12 +2326,12 @@ async def run_once(input_value: Any) -> Any: ("function_call_output", "call-approved"), ] expected_output = ( - run_loop._OUTPUT_GUARDRAIL_BLOCKED_TOOL_OUTPUT if tripwire else "approved-result" + run_loop._OUTPUT_GUARDRAIL_BLOCKED_TOOL_OUTPUT if outcome == "trips" else "approved-result" ) assert saved_tool_items[1].get("output") == expected_output - if tripwire: - guardrail_state["tripwire"] = False + if outcome == "trips": + guardrail_state["outcome"] = "passes" model.enqueue([get_text_message("done")]) next_result = await run_once("Continue") assert next_result.final_output == "done" @@ -2369,6 +2386,7 @@ def output_guardrail( first = await Runner.run(agent, "Use 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()) restored.approve(restored.get_interruptions()[0]) @@ -2382,6 +2400,69 @@ def output_guardrail( assert tool_calls == 0 +@pytest.mark.parametrize("mode", ["non_streamed", "streamed"]) +@pytest.mark.parametrize("serialized", [False, True], ids=["live", "serialized"]) +@pytest.mark.parametrize("attach_session", [False, True], ids=["without-session", "with-session"]) +@pytest.mark.asyncio +async def test_legacy_approval_checkpoint_uses_current_session_ownership( + mode: str, + serialized: bool, + attach_session: bool, +) -> None: + side_effects: list[str] = [] + + @function_tool(name_override="approval_tool", needs_approval=True) + def approval_tool() -> str: + side_effects.append("executed") + return "approved-result" + + model = ScriptedModel( + [[get_function_tool_call("approval_tool", "{}", call_id="call-approved")]] + ) + agent = Agent( + name="test", + model=model, + tools=[approval_tool], + tool_use_behavior="stop_on_first_tool", + ) + legacy_session = SimpleListSession() + first = await Runner.run(agent, "Use approval_tool", session=legacy_session) + state = first.to_state() + assert state._current_turn_persisted_item_count > 0 + if serialized: + state = await RunState.from_json(agent, state.to_json()) + state.approve(state.get_interruptions()[0]) + agent.output_guardrails = [ + OutputGuardrail( + guardrail_function=lambda _context, _agent, _output: GuardrailFunctionOutput( + output_info=None, + tripwire_triggered=False, + ) + ) + ] + session = legacy_session if attach_session else None + + if attach_session: + with pytest.raises(UserError, match="after current-turn items were persisted"): + if mode == "non_streamed": + await Runner.run(agent, state, session=session) + else: + result = Runner.run_streamed(agent, state, session=session) + await consume_stream(result) + assert side_effects == [] + return + + if mode == "non_streamed": + result = await Runner.run(agent, state, session=None) + else: + result = Runner.run_streamed(agent, state, session=None) + await consume_stream(result) + + assert result.final_output == "approved-result" + assert state._current_turn_persisted_item_count == 0 + assert side_effects == ["executed"] + + @pytest.mark.parametrize("mode", ["non_streamed", "streamed"]) @pytest.mark.asyncio async def test_output_guardrails_fail_closed_with_server_managed_history(mode: str) -> None: @@ -2408,6 +2489,84 @@ async def test_output_guardrails_fail_closed_with_server_managed_history(mode: s assert not model.calls +@pytest.mark.parametrize("mode", ["non_streamed", "streamed"]) +@pytest.mark.parametrize("strict", [False, True], ids=["default", "strict"]) +@pytest.mark.parametrize("use_run_config_model", [False, True], ids=["agent-model", "run-model"]) +@pytest.mark.asyncio +async def test_chat_completions_output_guardrails_use_adapter_conversation_policy( + mode: str, + strict: bool, + use_run_config_model: bool, + monkeypatch: pytest.MonkeyPatch, + caplog: pytest.LogCaptureFixture, +) -> None: + scripted_model = ScriptedModel([[get_text_message("accepted-output")]]) + chat_model = OpenAIChatCompletionsModel( + model="test", + openai_client=cast(Any, object()), + strict_feature_validation=strict, + ) + + async def get_response(*args: Any, **kwargs: Any) -> Any: + chat_model._handle_unsupported_server_managed_conversation_state( + previous_response_id=kwargs.get("previous_response_id"), + conversation_id=kwargs.get("conversation_id"), + ) + return await scripted_model.get_response(*args, **kwargs) + + async def stream_response(*args: Any, **kwargs: Any) -> AsyncIterator[Any]: + chat_model._handle_unsupported_server_managed_conversation_state( + previous_response_id=kwargs.get("previous_response_id"), + conversation_id=kwargs.get("conversation_id"), + ) + async for event in scripted_model.stream_response(*args, **kwargs): + yield event + + monkeypatch.setattr(chat_model, "get_response", get_response) + monkeypatch.setattr(chat_model, "stream_response", stream_response) + agent = Agent( + name="test", + model=ScriptedModel() if use_run_config_model else chat_model, + output_guardrails=[ + OutputGuardrail( + guardrail_function=lambda _context, _agent, _output: GuardrailFunctionOutput( + output_info=None, + tripwire_triggered=False, + ) + ) + ], + ) + run_config = RunConfig(model=chat_model) if use_run_config_model else None + caplog.set_level(logging.WARNING, logger="openai.agents") + + async def run_once() -> Any: + if mode == "non_streamed": + return await Runner.run( + agent, + "hello", + previous_response_id="response-id", + run_config=run_config, + ) + result = Runner.run_streamed( + agent, + "hello", + previous_response_id="response-id", + run_config=run_config, + ) + await consume_stream(result) + return result + + if strict: + with pytest.raises(UserError, match="OpenAIChatCompletionsModel does not support"): + await run_once() + assert not scripted_model.calls + return + + assert (await run_once()).final_output == "accepted-output" + assert "Ignoring unsupported server-managed conversation state" in caplog.text + assert len(scripted_model.calls) == 1 + + @pytest.mark.parametrize("mode", ["non_streamed", "streamed"]) @pytest.mark.parametrize("session_kind", ["simple", "openai_conversations"]) @pytest.mark.parametrize("arguments", ["{}", ""], ids=["json-object", "empty"]) @@ -2684,10 +2843,10 @@ async def run_once() -> Any: @pytest.mark.parametrize("mode", ["non_streamed", "streamed"]) @pytest.mark.asyncio -async def test_failing_output_guardrail_does_not_persist_the_unverdictable_turn( +async def test_failing_output_guardrail_keeps_the_whole_final_turn( mode: str, ) -> None: - """A guardrail error leaves no verdict, so its response is not persisted.""" + """A guardrail error leaves no rejection, so the completed turn remains replayable.""" @function_tool(name_override="commit_tool") def commit_tool() -> str: @@ -2728,7 +2887,7 @@ async def run_once() -> None: saved_items = await session.get_items() saved = [item.get("type") or item.get("role") for item in saved_items if isinstance(item, dict)] - assert saved == ["user"] + assert saved == ["user", "message", "function_call", "function_call_output"] @pytest.mark.asyncio From f0a1a46d799e5ab7014b4ae9b80bc6f91e4c8645 Mon Sep 17 00:00:00 2001 From: Kazuhiro Sera Date: Wed, 19 Aug 2026 15:00:25 +0900 Subject: [PATCH 12/12] fix: persist nonterminal approval sibling outputs --- src/agents/run_internal/blocked_output.py | 8 +- tests/test_agent_runner_streamed.py | 210 ++++++++++++++++++++++ 2 files changed, 215 insertions(+), 3 deletions(-) diff --git a/src/agents/run_internal/blocked_output.py b/src/agents/run_internal/blocked_output.py index 7011965843..bb1da4ffaa 100644 --- a/src/agents/run_internal/blocked_output.py +++ b/src/agents/run_internal/blocked_output.py @@ -271,8 +271,8 @@ def _should_defer_interrupted_session_items( agent: Agent[Any], run_config: RunConfig, ) -> bool: - """Keep pre-verdict approval state in RunState instead of durable Session history.""" - return _has_output_guardrails(agent, run_config) + """Defer only approval state that could still become guarded terminal tool output.""" + return _has_output_guardrails(agent, run_config) and agent.tool_use_behavior != "run_llm_again" def _validate_resumed_session_output_guardrail_safety( @@ -297,7 +297,9 @@ def _validate_resumed_session_output_guardrail_safety( "Cannot resume a serialized approval checkpoint with output guardrails because the " "current response boundary cannot be proven. Start a new run from safe input." ) - if run_state._current_turn_persisted_item_count > 0: + if run_state._current_turn_persisted_item_count > 0 and ( + _should_defer_interrupted_session_items(agent, run_config) + ): if session is not None: raise UserError( "Cannot resume an approval checkpoint with output guardrails after current-turn " diff --git a/tests/test_agent_runner_streamed.py b/tests/test_agent_runner_streamed.py index 44246f25d8..e5c2a03fd8 100644 --- a/tests/test_agent_runner_streamed.py +++ b/tests/test_agent_runner_streamed.py @@ -42,6 +42,7 @@ ToolGuardrailFunctionOutput, ToolInputGuardrailData, ToolOutputGuardrailData, + ToolsToFinalOutputResult, UserError, function_tool, handoff, @@ -2242,6 +2243,215 @@ async def test_tool() -> str: assert output_count == 1 +@pytest.mark.parametrize("mode", ["non_streamed", "streamed"]) +@pytest.mark.asyncio +async def test_run_llm_again_approval_persists_completed_sibling(mode: str) -> None: + side_effects: list[str] = [] + + @function_tool(name_override="approval_tool", needs_approval=True) + def approval_tool() -> str: + side_effects.append("approved") + return "approved-output" + + @function_tool(name_override="sibling_tool") + def sibling_tool() -> str: + side_effects.append("sibling") + return "sibling-output" + + model = ScriptedModel( + [ + [ + get_function_tool_call("approval_tool", "{}", call_id="call-approved"), + get_function_tool_call("sibling_tool", "{}", call_id="call-sibling"), + ], + [get_text_message("done")], + ] + ) + agent = Agent( + name="test", + model=model, + tools=[approval_tool, sibling_tool], + output_guardrails=[ + OutputGuardrail( + guardrail_function=lambda _context, _agent, _output: GuardrailFunctionOutput( + output_info=None, + tripwire_triggered=False, + ) + ) + ], + ) + 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 both tools") + assert len(first.interruptions) == 1 + assert side_effects == ["sibling"] + + saved_before_resume = await session.get_items() + saved_sibling_items = [ + item + for item in saved_before_resume + if isinstance(item, dict) and item.get("call_id") == "call-sibling" + ] + assert [item.get("type") for item in saved_sibling_items] == [ + "function_call", + "function_call_output", + ] + assert saved_sibling_items[1].get("output") == "sibling-output" + + state = first.to_state() + assert state._current_turn_persisted_item_count > 0 + state.approve(first.interruptions[0]) + resumed = await run_once(state) + + assert resumed.final_output == "done" + assert side_effects == ["sibling", "approved"] + saved_after_resume = await session.get_items() + for call_id in ("call-sibling", "call-approved"): + assert [ + item.get("type") + for item in saved_after_resume + if isinstance(item, dict) and item.get("call_id") == call_id + ] == ["function_call", "function_call_output"] + + +@pytest.mark.parametrize("mode", ["non_streamed", "streamed"]) +@pytest.mark.parametrize("terminal_behavior", ["first", "named", "custom"]) +@pytest.mark.asyncio +async def test_terminal_behaviors_defer_completed_approval_siblings( + mode: str, + terminal_behavior: str, +) -> None: + @function_tool(name_override="approval_tool", needs_approval=True) + def approval_tool() -> str: + return "approved-output" + + @function_tool(name_override="sibling_tool") + def sibling_tool() -> str: + return "sibling-output" + + model = ScriptedModel( + [ + [ + get_function_tool_call("approval_tool", "{}", call_id="call-approved"), + get_function_tool_call("sibling_tool", "{}", call_id="call-sibling"), + ] + ] + ) + agent = Agent( + name="test", + model=model, + tools=[approval_tool, sibling_tool], + output_guardrails=[ + OutputGuardrail( + guardrail_function=lambda _context, _agent, _output: GuardrailFunctionOutput( + output_info=None, + tripwire_triggered=False, + ) + ) + ], + ) + if terminal_behavior == "first": + agent.tool_use_behavior = "stop_on_first_tool" + elif terminal_behavior == "named": + agent.tool_use_behavior = {"stop_at_tool_names": ["approval_tool"]} + else: + agent.tool_use_behavior = lambda _context, results: ToolsToFinalOutputResult( + is_final_output=True, + final_output=results[0].output, + ) + + session = SimpleListSession() + if mode == "non_streamed": + result = await Runner.run(agent, "Use both tools", session=session) + else: + result = Runner.run_streamed(agent, "Use both tools", session=session) + await consume_stream(result) + + assert len(result.interruptions) == 1 + assert result.to_state()._current_turn_persisted_item_count == 0 + assert "sibling-output" not in json.dumps(await session.get_items()) + + +@pytest.mark.parametrize("mode", ["non_streamed", "streamed"]) +@pytest.mark.parametrize("terminal_behavior", ["first", "named", "custom"]) +@pytest.mark.asyncio +async def test_persisted_run_llm_again_checkpoint_rejects_terminal_behavior_change( + mode: str, + terminal_behavior: str, +) -> None: + side_effects: list[str] = [] + + @function_tool(name_override="approval_tool", needs_approval=True) + def approval_tool() -> str: + side_effects.append("approved") + return "approved-output" + + @function_tool(name_override="sibling_tool") + def sibling_tool() -> str: + side_effects.append("sibling") + return "sibling-output" + + model = ScriptedModel( + [ + [ + get_function_tool_call("approval_tool", "{}", call_id="call-approved"), + get_function_tool_call("sibling_tool", "{}", call_id="call-sibling"), + ] + ] + ) + agent = Agent( + name="test", + model=model, + tools=[approval_tool, sibling_tool], + output_guardrails=[ + OutputGuardrail( + guardrail_function=lambda _context, _agent, _output: GuardrailFunctionOutput( + output_info=None, + tripwire_triggered=False, + ) + ) + ], + ) + session = SimpleListSession() + + if mode == "non_streamed": + first = await Runner.run(agent, "Use both tools", session=session) + else: + first = Runner.run_streamed(agent, "Use both tools", session=session) + await consume_stream(first) + + assert side_effects == ["sibling"] + state = first.to_state() + assert state._current_turn_persisted_item_count > 0 + state.approve(first.interruptions[0]) + + if terminal_behavior == "first": + agent.tool_use_behavior = "stop_on_first_tool" + elif terminal_behavior == "named": + agent.tool_use_behavior = {"stop_at_tool_names": ["approval_tool"]} + else: + agent.tool_use_behavior = lambda _context, results: ToolsToFinalOutputResult( + is_final_output=True, + final_output=results[0].output, + ) + + with pytest.raises(UserError, match="after current-turn items were persisted"): + if mode == "non_streamed": + await Runner.run(agent, state, session=session) + else: + result = Runner.run_streamed(agent, state, session=session) + await consume_stream(result) + + assert side_effects == ["sibling"] + + @pytest.mark.parametrize("mode", ["non_streamed", "streamed"]) @pytest.mark.parametrize("outcome", ["passes", "trips", "error"]) @pytest.mark.asyncio