wehub-resource-sync
67 lines
2.4 KiB
Python
67 lines
2.4 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
from typing import Any
|
|
|
|
import pytest
|
|
|
|
from agents import (
|
|
Agent,
|
|
GuardrailFunctionOutput,
|
|
OutputGuardrail,
|
|
OutputGuardrailTripwireTriggered,
|
|
RunContextWrapper,
|
|
)
|
|
from agents.run_internal.guardrails import run_output_guardrails
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_run_output_guardrails_awaits_cancelled_tasks():
|
|
"""When one output guardrail trips, sibling guardrails must be awaited after cancel.
|
|
|
|
Regression test: ``run_output_guardrails`` previously cancelled sibling tasks but
|
|
did not await them, which can leak pending tasks and emit
|
|
``Task was destroyed but it is pending!`` warnings. The input-guardrail variants
|
|
already await on cancel; the output variant should match.
|
|
"""
|
|
|
|
slow_started = asyncio.Event()
|
|
cancelled_observed = asyncio.Event()
|
|
|
|
async def slow_then_observe_cancel(
|
|
context: RunContextWrapper[Any], agent: Agent[Any], agent_output: Any
|
|
) -> GuardrailFunctionOutput:
|
|
slow_started.set()
|
|
try:
|
|
await asyncio.sleep(5)
|
|
except asyncio.CancelledError:
|
|
cancelled_observed.set()
|
|
raise
|
|
return GuardrailFunctionOutput(output_info=None, tripwire_triggered=False)
|
|
|
|
async def fast_tripwire(
|
|
context: RunContextWrapper[Any], agent: Agent[Any], agent_output: Any
|
|
) -> GuardrailFunctionOutput:
|
|
# Wait until the slow guardrail is actually parked on its sleep so the
|
|
# subsequent cancel hits the installed CancelledError handler. Without
|
|
# this, the slow task could be cancelled before reaching the try block
|
|
# and ``cancelled_observed`` would stay unset even with the production
|
|
# fix in place.
|
|
await slow_started.wait()
|
|
return GuardrailFunctionOutput(output_info=None, tripwire_triggered=True)
|
|
|
|
guardrails = [
|
|
OutputGuardrail(guardrail_function=slow_then_observe_cancel),
|
|
OutputGuardrail(guardrail_function=fast_tripwire),
|
|
]
|
|
agent: Agent[Any] = Agent(name="test")
|
|
context: RunContextWrapper[Any] = RunContextWrapper(context=None)
|
|
|
|
with pytest.raises(OutputGuardrailTripwireTriggered):
|
|
await run_output_guardrails(guardrails, agent, "agent output", context)
|
|
|
|
# The slow guardrail must have observed cancellation and finished before
|
|
# ``run_output_guardrails`` returned. If it had not been awaited, this event
|
|
# would still be unset at this point.
|
|
assert cancelled_observed.is_set()
|