feat: add run_in_parallel parameter to input guardrails (#1986)

Author: RVV-karma

src/agents/guardrail.py

@@ -70,7 +70,7 @@ class OutputGuardrailResult:
 
 @dataclass
 class InputGuardrail(Generic[TContext]):
-    """Input guardrails are checks that run in parallel to the agent's execution.
+    """Input guardrails are checks that run either in parallel with the agent or before it starts.
     They can be used to do things like:
     - Check if input messages are off-topic
     - Take over control of the agent's execution if an unexpected input is detected
@@ -97,6 +97,11 @@ class InputGuardrail(Generic[TContext]):
     function's name.
     """
 
+    run_in_parallel: bool = True
+    """Whether the guardrail runs concurrently with the agent (True, default) or before
+    the agent starts (False).
+    """
+
     def get_name(self) -> str:
         if self.name:
             return self.name
@@ -209,6 +214,7 @@ def input_guardrail(
 def input_guardrail(
     *,
     name: str | None = None,
+    run_in_parallel: bool = True,
 ) -> Callable[
     [_InputGuardrailFuncSync[TContext_co] | _InputGuardrailFuncAsync[TContext_co]],
     InputGuardrail[TContext_co],
@@ -221,6 +227,7 @@ def input_guardrail(
     | None = None,
     *,
     name: str | None = None,
+    run_in_parallel: bool = True,
 ) -> (
     InputGuardrail[TContext_co]
     | Callable[
@@ -235,8 +242,14 @@ def input_guardrail(
         @input_guardrail
         def my_sync_guardrail(...): ...
 
-        @input_guardrail(name="guardrail_name")
+        @input_guardrail(name="guardrail_name", run_in_parallel=False)
         async def my_async_guardrail(...): ...
+
+    Args:
+        func: The guardrail function to wrap.
+        name: Optional name for the guardrail. If not provided, uses the function's name.
+        run_in_parallel: Whether to run the guardrail concurrently with the agent (True, default)
+            or before the agent starts (False).
     """
 
     def decorator(
@@ -246,6 +259,7 @@ def decorator(
             guardrail_function=f,
             # If not set, guardrail name uses the function’s name by default.
             name=name if name else f.__name__,
+            run_in_parallel=run_in_parallel,
         )
 
     if func is not None:

src/agents/run.py

@@ -614,11 +614,31 @@ async def run(
                     )
 
                     if current_turn == 1:
+                        # Separate guardrails based on execution mode.
+                        all_input_guardrails = starting_agent.input_guardrails + (
+                            run_config.input_guardrails or []
+                        )
+                        sequential_guardrails = [
+                            g for g in all_input_guardrails if not g.run_in_parallel
+                        ]
+                        parallel_guardrails = [g for g in all_input_guardrails if g.run_in_parallel]
+
+                        # Run blocking guardrails first, before agent starts.
+                        # (will raise exception if tripwire triggered).
+                        sequential_results = []
+                        if sequential_guardrails:
+                            sequential_results = await self._run_input_guardrails(
+                                starting_agent,
+                                sequential_guardrails,
+                                _copy_str_or_list(prepared_input),
+                                context_wrapper,
+                            )
+
+                        # Run parallel guardrails + agent together.
                         input_guardrail_results, turn_result = await asyncio.gather(
                             self._run_input_guardrails(
                                 starting_agent,
-                                starting_agent.input_guardrails
-                                + (run_config.input_guardrails or []),
+                                parallel_guardrails,
                                 _copy_str_or_list(prepared_input),
                                 context_wrapper,
                             ),
@@ -635,6 +655,9 @@ async def run(
                                 server_conversation_tracker=server_conversation_tracker,
                             ),
                         )
+
+                        # Combine sequential and parallel results.
+                        input_guardrail_results = sequential_results + input_guardrail_results
                     else:
                         turn_result = await self._run_single_turn(
                             agent=current_agent,
@@ -954,6 +977,11 @@ async def _run_input_guardrails_with_queue(
             for done in asyncio.as_completed(guardrail_tasks):
                 result = await done
                 if result.output.tripwire_triggered:
+                    # Cancel all remaining guardrail tasks if a tripwire is triggered.
+                    for t in guardrail_tasks:
+                        t.cancel()
+                    # Wait for cancellations to propagate by awaiting the cancelled tasks.
+                    await asyncio.gather(*guardrail_tasks, return_exceptions=True)
                     _error_tracing.attach_error_to_span(
                         parent_span,
                         SpanError(
@@ -964,14 +992,19 @@ async def _run_input_guardrails_with_queue(
                             },
                         ),
                     )
+                    queue.put_nowait(result)
+                    guardrail_results.append(result)
+                    break
                 queue.put_nowait(result)
                 guardrail_results.append(result)
         except Exception:
             for t in guardrail_tasks:
                 t.cancel()
             raise
 
-        streamed_result.input_guardrail_results = guardrail_results
+        streamed_result.input_guardrail_results = (
+            streamed_result.input_guardrail_results + guardrail_results
+        )
 
     @classmethod
     async def _start_streaming(
@@ -1063,11 +1096,36 @@ async def _start_streaming(
                     break
 
                 if current_turn == 1:
-                    # Run the input guardrails in the background and put the results on the queue
+                    # Separate guardrails based on execution mode.
+                    all_input_guardrails = starting_agent.input_guardrails + (
+                        run_config.input_guardrails or []
+                    )
+                    sequential_guardrails = [
+                        g for g in all_input_guardrails if not g.run_in_parallel
+                    ]
+                    parallel_guardrails = [g for g in all_input_guardrails if g.run_in_parallel]
+
+                    # Run sequential guardrails first.
+                    if sequential_guardrails:
+                        await cls._run_input_guardrails_with_queue(
+                            starting_agent,
+                            sequential_guardrails,
+                            ItemHelpers.input_to_new_input_list(prepared_input),
+                            context_wrapper,
+                            streamed_result,
+                            current_span,
+                        )
+                        # Check if any blocking guardrail triggered and raise before starting agent.
+                        for result in streamed_result.input_guardrail_results:
+                            if result.output.tripwire_triggered:
+                                streamed_result._event_queue.put_nowait(QueueCompleteSentinel())
+                                raise InputGuardrailTripwireTriggered(result)
+
+                    # Run parallel guardrails in background.
                     streamed_result._input_guardrails_task = asyncio.create_task(
                         cls._run_input_guardrails_with_queue(
                             starting_agent,
-                            starting_agent.input_guardrails + (run_config.input_guardrails or []),
+                            parallel_guardrails,
                             ItemHelpers.input_to_new_input_list(prepared_input),
                             context_wrapper,
                             streamed_result,
@@ -1632,6 +1690,8 @@ async def _run_input_guardrails(
                 # Cancel all guardrail tasks if a tripwire is triggered.
                 for t in guardrail_tasks:
                     t.cancel()
+                # Wait for cancellations to propagate by awaiting the cancelled tasks.
+                await asyncio.gather(*guardrail_tasks, return_exceptions=True)
                 _error_tracing.attach_error_to_current_span(
                     SpanError(
                         message="Guardrail tripwire triggered",