diff --git a/src/agents/run.py b/src/agents/run.py index bb07bd2554..d0e3f9ee56 100644 --- a/src/agents/run.py +++ b/src/agents/run.py @@ -1266,14 +1266,17 @@ def _finalize_result(result: RunResult) -> RunResult: ) if parallel_guardrails: + guardrail_task = asyncio.create_task( + run_input_guardrails( + starting_agent, + parallel_guardrails, + copy_input_items(original_input), + context_wrapper, + ) + ) try: parallel_results, turn_result = await asyncio.gather( - run_input_guardrails( - starting_agent, - parallel_guardrails, - copy_input_items(original_input), - context_wrapper, - ), + guardrail_task, model_task, ) except InputGuardrailTripwireTriggered: @@ -1292,6 +1295,19 @@ def _finalize_result(result: RunResult) -> RunResult: ) ) raise + except BaseException: + # A non-tripwire failure (the model turn raising, or a + # guardrail raising a non-tripwire error) propagates from + # gather without cancelling the sibling task. Cancel and drain + # whichever side is still pending so it is not left running + # after the run has failed and its exception is not swallowed. + for pending_task in (guardrail_task, model_task): + if not pending_task.done(): + pending_task.cancel() + await asyncio.gather( + guardrail_task, model_task, return_exceptions=True + ) + raise else: turn_result = await model_task diff --git a/tests/test_guardrails.py b/tests/test_guardrails.py index 511c342ca5..08e57e67b1 100644 --- a/tests/test_guardrails.py +++ b/tests/test_guardrails.py @@ -627,6 +627,104 @@ async def slow_get_response(*args, **kwargs): assert model_cancelled.is_set() is False +@pytest.mark.asyncio +async def test_model_error_cancels_parallel_input_guardrail_task(): + """A non-tripwire model failure must cancel the still-running guardrail task. + + Without cancellation the guardrail task is orphaned and keeps running after + ``Runner.run`` has already raised. + """ + guardrail_started = asyncio.Event() + guardrail_cancelled = asyncio.Event() + guardrail_finished = asyncio.Event() + + @input_guardrail(run_in_parallel=True) + async def slow_parallel_check( + ctx: RunContextWrapper[Any], agent: Agent[Any], input: str | list[TResponseInputItem] + ) -> GuardrailFunctionOutput: + guardrail_started.set() + try: + await asyncio.sleep(LONG_DELAY) + guardrail_finished.set() + return GuardrailFunctionOutput( + output_info="parallel_ok", + tripwire_triggered=False, + ) + except asyncio.CancelledError: + guardrail_cancelled.set() + raise + + model = FakeModel() + + async def boom_get_response(*args, **kwargs): + # Only blow up once the guardrail is genuinely mid-flight. + await asyncio.wait_for(guardrail_started.wait(), timeout=1) + raise RuntimeError("model boom") + + agent = Agent( + name="model_error_agent", + input_guardrails=[slow_parallel_check], + model=model, + ) + + with patch.object(model, "get_response", side_effect=boom_get_response): + with pytest.raises(RuntimeError, match="model boom"): + await Runner.run(agent, "trigger guardrail") + + # By the time Runner.run returns, the guardrail task must already be + # cancelled rather than left running to completion in the background. + assert guardrail_started.is_set() is True + assert guardrail_cancelled.is_set() is True + assert guardrail_finished.is_set() is False + + +@pytest.mark.asyncio +async def test_parallel_guardrail_non_tripwire_error_not_swallowed(): + """A non-tripwire error raised inside a parallel guardrail must propagate. + + It should also cancel the in-flight model task rather than leave it running. + """ + model_started = asyncio.Event() + model_cancelled = asyncio.Event() + model_finished = asyncio.Event() + + @input_guardrail(run_in_parallel=True) + async def raising_parallel_check( + ctx: RunContextWrapper[Any], agent: Agent[Any], input: str | list[TResponseInputItem] + ) -> GuardrailFunctionOutput: + await asyncio.wait_for(model_started.wait(), timeout=1) + raise ValueError("guardrail boom") + + model = FakeModel() + original_get_response = model.get_response + + async def slow_get_response(*args, **kwargs): + model_started.set() + try: + await asyncio.sleep(LONG_DELAY) + return await original_get_response(*args, **kwargs) + except asyncio.CancelledError: + model_cancelled.set() + raise + finally: + model_finished.set() + + agent = Agent( + name="guardrail_error_agent", + input_guardrails=[raising_parallel_check], + model=model, + ) + model.set_next_output([get_text_message("should_not_finish")]) + + with patch.object(model, "get_response", side_effect=slow_get_response): + with pytest.raises(ValueError, match="guardrail boom"): + await Runner.run(agent, "trigger guardrail") + + await asyncio.wait_for(model_finished.wait(), timeout=1) + assert model_started.is_set() is True + assert model_cancelled.is_set() is True + + @pytest.mark.asyncio async def test_parallel_guardrail_may_not_prevent_tool_execution_streaming(): tool_was_executed = False