fix: cancel model task when parallel input guardrail trips (#2416)

Author: seratch

src/agents/run.py

@@ -52,6 +52,7 @@
     resolve_resumed_context,
     resolve_trace_settings,
     save_turn_items_if_needed,
+    should_cancel_parallel_model_task_on_input_guardrail_trip,
     update_run_state_for_interruption,
     validate_session_conversation_settings,
 )
@@ -1007,6 +1008,10 @@ async def run(
                                     model_task,
                                 )
                             except InputGuardrailTripwireTriggered:
+                                if should_cancel_parallel_model_task_on_input_guardrail_trip():
+                                    if not model_task.done():
+                                        model_task.cancel()
+                                    await asyncio.gather(model_task, return_exceptions=True)
                                 session_input_items_for_persistence = (
                                     await persist_session_items_for_guardrail_trip(
                                         session,

src/agents/run_internal/agent_runner_helpers.py

@@ -43,9 +43,30 @@
     "resolve_processed_response",
     "resolve_resumed_context",
     "save_turn_items_if_needed",
+    "should_cancel_parallel_model_task_on_input_guardrail_trip",
     "update_run_state_for_interruption",
 ]
 
+_PARALLEL_INPUT_GUARDRAIL_CANCEL_PATCH_ID = (
+    "openai_agents.cancel_parallel_model_task_on_input_guardrail_trip.v1"
+)
+
+
+def should_cancel_parallel_model_task_on_input_guardrail_trip() -> bool:
+    """Return whether an in-flight model task should be cancelled on guardrail trip."""
+    try:
+        from temporalio import workflow as temporal_workflow  # type: ignore[import-not-found]
+    except Exception:
+        return True
+
+    try:
+        if not temporal_workflow.in_workflow():
+            return True
+        # Preserve replay compatibility for histories created before cancellation.
+        return bool(temporal_workflow.patched(_PARALLEL_INPUT_GUARDRAIL_CANCEL_PATCH_ID))
+    except Exception:
+        return True
+
 
 def apply_resumed_conversation_settings(
     *,

tests/test_guardrails.py

@@ -509,6 +509,104 @@ async def slow_parallel_check(
     assert model.first_turn_args is not None, "Model should have been called in parallel mode"
 
 
+@pytest.mark.asyncio
+async def test_parallel_guardrail_trip_cancels_model_task():
+    model_started = asyncio.Event()
+    model_cancelled = asyncio.Event()
+    model_finished = asyncio.Event()
+
+    @input_guardrail(run_in_parallel=True)
+    async def tripwire_after_model_starts(
+        ctx: RunContextWrapper[Any], agent: Agent[Any], input: str | list[TResponseInputItem]
+    ) -> GuardrailFunctionOutput:
+        await asyncio.wait_for(model_started.wait(), timeout=1)
+        return GuardrailFunctionOutput(
+            output_info="parallel_tripwire",
+            tripwire_triggered=True,
+        )
+
+    model = FakeModel()
+    original_get_response = model.get_response
+
+    async def slow_get_response(*args, **kwargs):
+        model_started.set()
+        try:
+            await asyncio.sleep(0.2)
+            return await original_get_response(*args, **kwargs)
+        except asyncio.CancelledError:
+            model_cancelled.set()
+            raise
+        finally:
+            model_finished.set()
+
+    agent = Agent(
+        name="parallel_tripwire_agent",
+        instructions="Reply with 'hello'",
+        input_guardrails=[tripwire_after_model_starts],
+        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(InputGuardrailTripwireTriggered):
+            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_trip_compat_mode_does_not_cancel_model_task():
+    model_started = asyncio.Event()
+    model_cancelled = asyncio.Event()
+    model_finished = asyncio.Event()
+
+    @input_guardrail(run_in_parallel=True)
+    async def tripwire_after_model_starts(
+        ctx: RunContextWrapper[Any], agent: Agent[Any], input: str | list[TResponseInputItem]
+    ) -> GuardrailFunctionOutput:
+        await asyncio.wait_for(model_started.wait(), timeout=1)
+        return GuardrailFunctionOutput(
+            output_info="parallel_tripwire",
+            tripwire_triggered=True,
+        )
+
+    model = FakeModel()
+    original_get_response = model.get_response
+
+    async def slow_get_response(*args, **kwargs):
+        model_started.set()
+        try:
+            await asyncio.sleep(0.2)
+            return await original_get_response(*args, **kwargs)
+        except asyncio.CancelledError:
+            model_cancelled.set()
+            raise
+        finally:
+            model_finished.set()
+
+    agent = Agent(
+        name="parallel_tripwire_agent",
+        instructions="Reply with 'hello'",
+        input_guardrails=[tripwire_after_model_starts],
+        model=model,
+    )
+    model.set_next_output([get_text_message("should_finish_without_cancel")])
+
+    with patch.object(model, "get_response", side_effect=slow_get_response):
+        with patch(
+            "agents.run.should_cancel_parallel_model_task_on_input_guardrail_trip",
+            return_value=False,
+        ):
+            with pytest.raises(InputGuardrailTripwireTriggered):
+                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 False
+
+
 @pytest.mark.asyncio
 async def test_parallel_guardrail_may_not_prevent_tool_execution_streaming():
     tool_was_executed = False