fix(run): report input guardrail results when a tripwire aborts the run (#4071)
This commit is contained in:
Henry Su
+6
-8
@@ -797,16 +797,16 @@ class AgentRunner:
|
||||
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]
|
||||
sequential_results: list[InputGuardrailResult] = []
|
||||
if sandbox_runtime.enabled and sequential_guardrails:
|
||||
# Blocking first-turn guardrails must run before sandbox prep so a tripwire
|
||||
# can prevent session creation, startup, or live-session mutation.
|
||||
try:
|
||||
sequential_results = await run_input_guardrails(
|
||||
await run_input_guardrails(
|
||||
starting_agent,
|
||||
sequential_guardrails,
|
||||
copy_input_items(original_input),
|
||||
context_wrapper,
|
||||
input_guardrail_results,
|
||||
)
|
||||
except InputGuardrailTripwireTriggered:
|
||||
session_input_items_for_persistence = (
|
||||
@@ -1221,11 +1221,12 @@ class AgentRunner:
|
||||
if current_turn <= 1:
|
||||
try:
|
||||
if sequential_guardrails:
|
||||
sequential_results = await run_input_guardrails(
|
||||
await run_input_guardrails(
|
||||
starting_agent,
|
||||
sequential_guardrails,
|
||||
copy_input_items(original_input),
|
||||
context_wrapper,
|
||||
input_guardrail_results,
|
||||
)
|
||||
except InputGuardrailTripwireTriggered:
|
||||
session_input_items_for_persistence = (
|
||||
@@ -1240,7 +1241,6 @@ class AgentRunner:
|
||||
)
|
||||
raise
|
||||
|
||||
parallel_results: list[InputGuardrailResult] = []
|
||||
model_task = asyncio.create_task(
|
||||
run_single_turn(
|
||||
bindings=current_bindings,
|
||||
@@ -1272,10 +1272,11 @@ class AgentRunner:
|
||||
parallel_guardrails,
|
||||
copy_input_items(original_input),
|
||||
context_wrapper,
|
||||
input_guardrail_results,
|
||||
)
|
||||
)
|
||||
try:
|
||||
parallel_results, turn_result = await asyncio.gather(
|
||||
_, turn_result = await asyncio.gather(
|
||||
guardrail_task,
|
||||
model_task,
|
||||
)
|
||||
@@ -1310,9 +1311,6 @@ class AgentRunner:
|
||||
raise
|
||||
else:
|
||||
turn_result = await model_task
|
||||
|
||||
input_guardrail_results.extend(sequential_results)
|
||||
input_guardrail_results.extend(parallel_results)
|
||||
else:
|
||||
turn_result = await run_single_turn(
|
||||
bindings=current_bindings,
|
||||
|
||||
@@ -121,8 +121,14 @@ async def run_input_guardrails(
|
||||
guardrails: list[InputGuardrail[TContext]],
|
||||
input: str | list[TResponseInputItem],
|
||||
context: RunContextWrapper[TContext],
|
||||
results_sink: list[InputGuardrailResult] | None = None,
|
||||
) -> list[InputGuardrailResult]:
|
||||
"""Run input guardrails concurrently and raise on tripwires."""
|
||||
"""Run input guardrails concurrently and raise on tripwires.
|
||||
|
||||
Results are recorded into ``results_sink`` as each guardrail completes, including the
|
||||
tripping result, so callers can report them even when this function raises. The streamed
|
||||
path publishes the same results through `RunResultStreaming.input_guardrail_results`.
|
||||
"""
|
||||
if not guardrails:
|
||||
return []
|
||||
|
||||
@@ -133,10 +139,16 @@ async def run_input_guardrails(
|
||||
|
||||
guardrail_results: list[InputGuardrailResult] = []
|
||||
|
||||
def record(result: InputGuardrailResult) -> None:
|
||||
guardrail_results.append(result)
|
||||
if results_sink is not None:
|
||||
results_sink.append(result)
|
||||
|
||||
try:
|
||||
for done in asyncio.as_completed(guardrail_tasks):
|
||||
result = await done
|
||||
if result.output.tripwire_triggered:
|
||||
record(result)
|
||||
for t in guardrail_tasks:
|
||||
t.cancel()
|
||||
await asyncio.gather(*guardrail_tasks, return_exceptions=True)
|
||||
@@ -147,7 +159,7 @@ async def run_input_guardrails(
|
||||
)
|
||||
)
|
||||
raise InputGuardrailTripwireTriggered(result)
|
||||
guardrail_results.append(result)
|
||||
record(result)
|
||||
except BaseException:
|
||||
# On any error (including a guardrail raising or the caller being cancelled),
|
||||
# cancel and await siblings so they don't leak past this function's return.
|
||||
|
||||
@@ -1997,3 +1997,140 @@ async def test_output_guardrail_raise_cancels_siblings():
|
||||
|
||||
assert sibling_cancelled.is_set(), "Sibling task should have been cancelled"
|
||||
assert not sibling_completed.is_set(), "Sibling task should not have completed"
|
||||
|
||||
|
||||
def _ordered_input_guardrails(
|
||||
*, second_triggers: bool, second_raises: bool = False, run_in_parallel: bool = False
|
||||
) -> list[InputGuardrail[Any]]:
|
||||
"""Build two guardrails whose completion order is fixed by an explicit barrier."""
|
||||
first_done = asyncio.Event()
|
||||
|
||||
async def first_fn(
|
||||
context: RunContextWrapper[Any], agent: Agent[Any], input: str | list[TResponseInputItem]
|
||||
) -> GuardrailFunctionOutput:
|
||||
first_done.set()
|
||||
return GuardrailFunctionOutput(output_info="passes", tripwire_triggered=False)
|
||||
|
||||
async def second_fn(
|
||||
context: RunContextWrapper[Any], agent: Agent[Any], input: str | list[TResponseInputItem]
|
||||
) -> GuardrailFunctionOutput:
|
||||
await first_done.wait()
|
||||
if second_raises:
|
||||
raise RuntimeError("guardrail exploded")
|
||||
return GuardrailFunctionOutput(output_info="second", tripwire_triggered=second_triggers)
|
||||
|
||||
return [
|
||||
InputGuardrail(guardrail_function=first_fn, name="passes", run_in_parallel=run_in_parallel),
|
||||
InputGuardrail(
|
||||
guardrail_function=second_fn,
|
||||
name="raises" if second_raises else "trips",
|
||||
run_in_parallel=run_in_parallel,
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
def _tripwire_agent(model: FakeModel, *, run_in_parallel: bool) -> Agent[Any]:
|
||||
return Agent(
|
||||
name="guardrail_results_agent",
|
||||
model=model,
|
||||
input_guardrails=_ordered_input_guardrails(
|
||||
second_triggers=True, run_in_parallel=run_in_parallel
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def _result_names(results: list[Any]) -> list[str]:
|
||||
return [result.guardrail.get_name() for result in results]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("run_in_parallel", [False, True])
|
||||
async def test_input_guardrail_tripwire_reports_results(run_in_parallel: bool):
|
||||
"""Runner.run() reports every completed guardrail result on the raised tripwire."""
|
||||
model = FakeModel()
|
||||
model.set_next_output([get_text_message("hello")])
|
||||
|
||||
with pytest.raises(InputGuardrailTripwireTriggered) as exc_info:
|
||||
await Runner.run(_tripwire_agent(model, run_in_parallel=run_in_parallel), "test input")
|
||||
|
||||
run_data = exc_info.value.run_data
|
||||
assert run_data is not None
|
||||
assert _result_names(run_data.input_guardrail_results) == ["passes", "trips"]
|
||||
assert exc_info.value.guardrail_result.guardrail.get_name() == "trips"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("run_in_parallel", [False, True])
|
||||
async def test_input_guardrail_tripwire_reports_results_streamed(run_in_parallel: bool):
|
||||
"""The streamed path reports the same results, including on the streamed result object."""
|
||||
model = FakeModel()
|
||||
model.set_next_output([get_text_message("hello")])
|
||||
|
||||
result = Runner.run_streamed(
|
||||
_tripwire_agent(model, run_in_parallel=run_in_parallel), "test input"
|
||||
)
|
||||
with pytest.raises(InputGuardrailTripwireTriggered) as exc_info:
|
||||
async for _ in result.stream_events():
|
||||
pass
|
||||
|
||||
run_data = exc_info.value.run_data
|
||||
assert run_data is not None
|
||||
assert _result_names(run_data.input_guardrail_results) == ["passes", "trips"]
|
||||
assert _result_names(result.input_guardrail_results) == ["passes", "trips"]
|
||||
|
||||
|
||||
def test_input_guardrail_tripwire_reports_results_sync():
|
||||
"""Runner.run_sync() matches the async entry points."""
|
||||
model = FakeModel()
|
||||
model.set_next_output([get_text_message("hello")])
|
||||
|
||||
with pytest.raises(InputGuardrailTripwireTriggered) as exc_info:
|
||||
Runner.run_sync(_tripwire_agent(model, run_in_parallel=False), "test input")
|
||||
|
||||
run_data = exc_info.value.run_data
|
||||
assert run_data is not None
|
||||
assert _result_names(run_data.input_guardrail_results) == ["passes", "trips"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_input_guardrail_results_reported_on_success():
|
||||
"""Passing guardrails still land on the successful result exactly once."""
|
||||
model = FakeModel()
|
||||
model.set_next_output([get_text_message("hello")])
|
||||
agent = Agent(
|
||||
name="guardrail_results_agent",
|
||||
model=model,
|
||||
input_guardrails=[
|
||||
InputGuardrail(
|
||||
guardrail_function=get_sync_guardrail(triggers=False),
|
||||
name="blocking",
|
||||
run_in_parallel=False,
|
||||
),
|
||||
InputGuardrail(
|
||||
guardrail_function=get_sync_guardrail(triggers=False),
|
||||
name="parallel",
|
||||
run_in_parallel=True,
|
||||
),
|
||||
],
|
||||
)
|
||||
|
||||
result = await Runner.run(agent, "test input")
|
||||
|
||||
assert _result_names(result.input_guardrail_results) == ["blocking", "parallel"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_input_guardrail_exception_reports_completed_results():
|
||||
"""A guardrail raising a non-tripwire error still preserves earlier results."""
|
||||
|
||||
collected: list[Any] = []
|
||||
with pytest.raises(RuntimeError, match="guardrail exploded"):
|
||||
await run_input_guardrails(
|
||||
Agent(name="t"),
|
||||
_ordered_input_guardrails(second_triggers=False, second_raises=True),
|
||||
"test input",
|
||||
RunContextWrapper(context=None),
|
||||
collected,
|
||||
)
|
||||
|
||||
assert _result_names(collected) == ["passes"]
|
||||
|
||||
Reference in New Issue
Block a user