From 92582ee738acc1d715b44a6ffca2419deed54a73 Mon Sep 17 00:00:00 2001 From: Jose Alvarez Date: Fri, 18 Sep 2026 15:32:07 +0200 Subject: [PATCH] Python: Synchronize workflow preemption test startup Wait for the inner stream to block before signalling shutdown or cancellation, independently of the retained preemption deadline. Drain pending test tasks and verify that the blocked call is cancelled before cleanup. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../foundry_hosting/tests/test_responses.py | 68 ++++++++++++------- 1 file changed, 44 insertions(+), 24 deletions(-) diff --git a/python/packages/foundry_hosting/tests/test_responses.py b/python/packages/foundry_hosting/tests/test_responses.py index d1ea6a0e06..0f0e2b2c05 100644 --- a/python/packages/foundry_hosting/tests/test_responses.py +++ b/python/packages/foundry_hosting/tests/test_responses.py @@ -16,6 +16,7 @@ import os import uuid from collections.abc import AsyncGenerator, AsyncIterator, Awaitable, Callable, Generator, Mapping, Sequence +from contextlib import asynccontextmanager from dataclasses import dataclass from importlib import import_module from pathlib import Path @@ -5597,7 +5598,8 @@ def __init__(self, name: str, texts: Sequence[str], *, gate: asyncio.Event | Non self._texts = list(texts) self._gate = gate self.run_count = 0 - self.started = asyncio.Event() # Set at the top of run(), before any gate wait. + self.started = asyncio.Event() + self.cancelled = asyncio.Event() def create_session(self, **kwargs: Any) -> AgentSession: del kwargs @@ -5638,14 +5640,18 @@ def run( del messages, session, kwargs assert stream is True, "The inner agent only runs in stream mode in Foundry Hosted Agents." self.run_count += 1 - self.started.set() texts = self._texts name = self.name gate = self._gate async def _aiter() -> AsyncIterator[AgentResponseUpdate]: + self.started.set() if gate is not None: - await gate.wait() # Simulates a stuck model/tool call for preemption tests. + try: + await gate.wait() # Simulates a stuck model/tool call for preemption tests. + except asyncio.CancelledError: + self.cancelled.set() + raise for text in texts: yield AgentResponseUpdate( contents=[Content.from_text(text=text)], @@ -5670,6 +5676,25 @@ async def start(messages: list[Message], ctx: WorkflowContext[AgentExecutorReque return WorkflowAgent(workflow=workflow, name="Multi Update Workflow Agent"), inner +@asynccontextmanager +async def _pending_workflow_event( + handler: AsyncGenerator[Any], started: asyncio.Event +) -> AsyncIterator[asyncio.Future[Any]]: + pending = asyncio.ensure_future(anext(handler)) + started_wait = asyncio.ensure_future(started.wait()) + try: + # Startup uses pytest's test timeout; only preemption has a short deadline. + await asyncio.wait([pending, started_wait], return_when=asyncio.FIRST_COMPLETED) + if pending.done(): + pytest.fail(f"Workflow returned before reaching the blocked call: {pending.result()!r}") + yield pending + finally: + started_wait.cancel() + pending.cancel() + await asyncio.gather(started_wait, pending, return_exceptions=True) + await handler.aclose() + + def _build_approval_workflow_agent( *, approval_request_id: str, @@ -5786,20 +5811,17 @@ async def test_cancellation_signal_preempts_stuck_workflow_call(self) -> None: await anext(handler) # response.created await anext(handler) # response.in_progress - # Pull the first workflow event in the background so we can wait for the inner agent's - # run() to actually start (proving it's genuinely stuck on `gate`) before signalling -- - # otherwise cancellation could preempt the pull before the workflow even reaches it. - pending = asyncio.ensure_future(anext(handler)) - await asyncio.wait_for(inner.started.wait(), timeout=1.0) - cancellation_signal.set() # Fires while the inner agent is stuck awaiting `gate`. + async with _pending_workflow_event(handler, inner.started) as pending: + cancellation_signal.set() # Fires while the inner agent is stuck awaiting `gate`. - async def _drain() -> list[Any]: - first = await pending - return [first, *[event async for event in handler]] + async def _drain() -> list[Any]: + first = await pending + return [first, *[event async for event in handler]] - # Bounded well below `gate` never being set: proves cancellation preempted the stuck - # call instead of only being observed after it (eventually) produced an update. - events = await asyncio.wait_for(_drain(), timeout=1.0) + # Bounded well below `gate` never being set: proves cancellation preempted the stuck + # call instead of only being observed after it (eventually) produced an update. + events = await asyncio.wait_for(_drain(), timeout=1.0) + assert inner.cancelled.is_set() types = [event.get("type") for event in events if isinstance(event, Mapping)] assert "response.output_text.delta" not in types @@ -5833,16 +5855,14 @@ async def test_shutdown_signal_preempts_stuck_workflow_call(self, tmp_path: Path await anext(handler) # response.created await anext(handler) # response.in_progress - # Pull the first workflow event in the background so we can wait for the inner agent's - # run() to actually start (proving it's genuinely stuck on `gate`) before signalling. - pending = asyncio.ensure_future(anext(handler)) - await asyncio.wait_for(inner.started.wait(), timeout=1.0) - context.shutdown.set() # Fires while the inner agent is stuck awaiting `gate`. + async with _pending_workflow_event(handler, inner.started) as pending: + context.shutdown.set() # Fires while the inner agent is stuck awaiting `gate`. - # Bounded well below `gate` never being set: proves shutdown preempted the stuck call - # instead of only being observed after it (eventually) produced an update. - with pytest.raises(ResponseExitForRecovery): - await asyncio.wait_for(pending, timeout=1.0) + # Bounded well below `gate` never being set: proves shutdown preempted the stuck call + # instead of only being observed after it (eventually) produced an update. + with pytest.raises(ResponseExitForRecovery): + await asyncio.wait_for(pending, timeout=1.0) + assert inner.cancelled.is_set() assert inner.run_count == 1