diff --git a/python/openai_agents/README.md b/python/openai_agents/README.md index 5fbb5bf5..44b42260 100644 --- a/python/openai_agents/README.md +++ b/python/openai_agents/README.md @@ -10,7 +10,7 @@ We welcome questions and feedback in the [#python-sdk](https://temporalio.slack. uv add temporalio-openai-agents ``` -This integration requires Temporal 1.34 or later. The SDK extends its package +This integration requires Temporal 1.35 or later. The SDK extends its package search path to support split-directory installations, including editable installs, separate user and system sites, layered deployments, and `--target` installs. @@ -818,6 +818,29 @@ retries the same way as any other workflow_streams publisher — see the conventional `RETRY` event pattern for surfacing the transition to consumers. +Input guardrails run in parallel with the model call by default. When one +trips while the model is streaming, the plugin cancels the streaming +activity, and the run raises `InputGuardrailTripwireTriggered` as +`stream_events()` does. The cancel reaches the activity through its +heartbeat, so set `heartbeat_timeout`: the activity then stops and +publishes an `error` event with the code `cancelled` to the topic. +Subscribers can still get the events the model sent before that, as a +native `stream_events()` consumer can. If the model has already finished +but the activity is still publishing its events when the cancel arrives, +the `error` event comes after `response.completed`. If the activity has +already returned, it publishes no `error` event. When the guardrail +has already tripped, the plugin does not start the model call. + +The workflow must still be open when the activity publishes that +`error` event. With the default `cancellation_type` (`TRY_CANCEL`), the +run raises without waiting for the activity to stop, so a workflow that +completes right after the trip can close first, and subscribers then see +the stream end without the event. Set +`cancellation_type=ActivityCancellationType.WAIT_CANCELLATION_COMPLETED` +to make the run raise only after the activity has published it. Other +kinds of cancel, such as a Worker shutdown or a heartbeat timeout, +publish no `error` event, because the activity's retry can still answer. + ## Feature Support This integration is presently subject to certain limitations. diff --git a/python/openai_agents/pyproject.toml b/python/openai_agents/pyproject.toml index b558ea88..e6dfc5b3 100644 --- a/python/openai_agents/pyproject.toml +++ b/python/openai_agents/pyproject.toml @@ -19,7 +19,7 @@ classifiers = [ dependencies = [ # opentelemetry and pydantic extras: the plugin imports temporalio.contrib.opentelemetry and # temporalio.contrib.pydantic at module level. - "temporalio[opentelemetry,pydantic]>=1.34.0,<2", + "temporalio[opentelemetry,pydantic]>=1.35.0,<2", "openai-agents>=0.22,<0.23", "openai>=3,<4", # imports openai._models.construct_type and openai.types.responses "mcp>=1.24,<3", diff --git a/python/openai_agents/src/temporalio/openai_agents/_heartbeat_decorator.py b/python/openai_agents/src/temporalio/openai_agents/_heartbeat_decorator.py index ad926718..2b3b6ae3 100644 --- a/python/openai_agents/src/temporalio/openai_agents/_heartbeat_decorator.py +++ b/python/openai_agents/src/temporalio/openai_agents/_heartbeat_decorator.py @@ -1,5 +1,6 @@ import asyncio -from collections.abc import Awaitable, Callable +from collections.abc import AsyncGenerator, Awaitable, Callable +from contextlib import asynccontextmanager from functools import wraps from typing import Any, TypeVar, cast, overload @@ -29,29 +30,41 @@ def auto_heartbeater( def decorate(fn: F) -> F: @wraps(fn) async def wrapper(*args: Any, **kwargs: Any) -> Any: - delay = interval - if delay is None: - heartbeat_timeout = activity.info().heartbeat_timeout - if heartbeat_timeout: - delay = heartbeat_timeout.total_seconds() / 2 - heartbeat_task = None - if delay: - heartbeat_task = asyncio.create_task(_heartbeat_every(delay)) - try: + async with heartbeating(interval): return await fn(*args, **kwargs) - finally: - if heartbeat_task: - heartbeat_task.cancel() - try: - await heartbeat_task - except asyncio.CancelledError: - pass return cast(F, wrapper) return decorate(fn) if fn is not None else decorate +@asynccontextmanager +async def heartbeating(interval: float | None = None) -> AsyncGenerator[None]: + """Heartbeat in the background while the block runs, as ``auto_heartbeater`` does. + + Stopping the heartbeats on exit awaits the heartbeat task and swallows a cancel + that arrives meanwhile. An activity that must act on every requested cancel checks + ``activity.cancellation_details()`` after the block. + """ + delay = interval + if delay is None: + heartbeat_timeout = activity.info().heartbeat_timeout + if heartbeat_timeout: + delay = heartbeat_timeout.total_seconds() / 2 + heartbeat_task = None + if delay: + heartbeat_task = asyncio.create_task(_heartbeat_every(delay)) + try: + yield + finally: + if heartbeat_task: + heartbeat_task.cancel() + try: + await heartbeat_task + except asyncio.CancelledError: + pass + + async def _heartbeat_every(delay: float) -> None: while True: await asyncio.sleep(delay) diff --git a/python/openai_agents/src/temporalio/openai_agents/_invoke_model_activity.py b/python/openai_agents/src/temporalio/openai_agents/_invoke_model_activity.py index 727d14e0..57e83b71 100644 --- a/python/openai_agents/src/temporalio/openai_agents/_invoke_model_activity.py +++ b/python/openai_agents/src/temporalio/openai_agents/_invoke_model_activity.py @@ -3,6 +3,7 @@ Implements mapping of OpenAI datastructures to Pydantic friendly types. """ +import asyncio import enum from collections.abc import Collection from dataclasses import dataclass @@ -44,14 +45,17 @@ APIStatusError, AsyncOpenAI, ) -from openai.types.responses import CustomToolParam +from openai.types.responses import CustomToolParam, ResponseErrorEvent from openai.types.responses.tool_param import Mcp from typing_extensions import Required, TypedDict from temporalio import activity from temporalio.contrib.workflow_streams import WorkflowStreamClient from temporalio.exceptions import ApplicationError -from temporalio.openai_agents._heartbeat_decorator import auto_heartbeater +from temporalio.openai_agents._heartbeat_decorator import ( + auto_heartbeater, + heartbeating, +) from temporalio.openai_agents._temporal_worker_env_ref import ( AllowAllWorkerEnvVars, _WorkerEnvRefResolver, @@ -400,6 +404,12 @@ def _raise_for_openai_status(e: APIStatusError) -> NoReturn: ) from e +def _cancel_requested() -> bool: + """Whether the workflow asked to cancel the running activity.""" + details = activity.cancellation_details() + return details is not None and details.cancel_requested + + def _default_openai_client() -> AsyncOpenAI: """Return the framework's default client as it is, or a client on its default key. @@ -453,7 +463,6 @@ async def invoke_model_activity(self, input: ActivityModelInput) -> ModelRespons _raise_for_openai_status(e) @activity.defn - @auto_heartbeater async def invoke_model_activity_streaming( self, input: StreamingActivityModelInput ) -> list[TResponseStreamEvent]: @@ -473,7 +482,7 @@ async def invoke_model_activity_streaming( ``streaming_topic`` so external consumers (UIs, tracing, etc.) can observe events as they arrive. - Heartbeats run on a background task via ``auto_heartbeater`` so + Heartbeats run on a background task via ``heartbeating`` so long initial-token latency or long pauses between chunks do not trip ``heartbeat_timeout``. """ @@ -496,26 +505,52 @@ async def invoke_model_activity_streaming( # typed decode can pass result_type=TResponseStreamEvent on # their own subscribe call. events_topic = stream.topic(topic) - async with stream: - try: - async for event in model.stream_response( - system_instructions=input.get("system_instructions"), - input=input["input"], - model_settings=model_settings, - tools=tools, - output_schema=input.get("output_schema"), - handoffs=handoffs, - tracing=ModelTracing(input["tracing"]), - previous_response_id=input.get("previous_response_id"), - conversation_id=input.get("conversation_id"), - prompt=input.get("prompt"), - ): - # OpenAI models set defer_build=True, so an event's pydantic - # schema may still be an unbuilt placeholder. - type(event).model_rebuild() - events.append(event) - events_topic.publish(event) - except APIStatusError as e: - _raise_for_openai_status(e) + try: + async with heartbeating(): + async with stream: + try: + async for event in model.stream_response( + system_instructions=input.get("system_instructions"), + input=input["input"], + model_settings=model_settings, + tools=tools, + output_schema=input.get("output_schema"), + handoffs=handoffs, + tracing=ModelTracing(input["tracing"]), + previous_response_id=input.get("previous_response_id"), + conversation_id=input.get("conversation_id"), + prompt=input.get("prompt"), + ): + # OpenAI models set defer_build=True, so an event's pydantic + # schema may still be an unbuilt placeholder. + type(event).model_rebuild() + events.append(event) + events_topic.publish(event) + except APIStatusError as e: + _raise_for_openai_status(e) + # Leaving the context managers swallows a cancel that arrives while they + # stop the stream's background flusher or the heartbeats, so raise that + # cancel here. + if _cancel_requested(): + raise asyncio.CancelledError() + except asyncio.CancelledError: + # When the workflow cancelled the call, for example because a parallel + # input guardrail tripped, tell subscribers the response ends here. The + # cancel can arrive while the model streams or while leaving the context + # managers flushes the events, so publish after both, behind the events + # still buffered. Other cancels, such as a Worker shutdown or a heartbeat + # timeout, can be retried, so publish nothing and let the retry answer. + if _cancel_requested(): + async with heartbeating(): + events_topic.publish( + ResponseErrorEvent( + type="error", + code="cancelled", + message="The model call was cancelled.", + sequence_number=len(events), + ) + ) + await stream.flush() + raise return events diff --git a/python/openai_agents/src/temporalio/openai_agents/_openai_runner.py b/python/openai_agents/src/temporalio/openai_agents/_openai_runner.py index 9b034e81..4876a11f 100644 --- a/python/openai_agents/src/temporalio/openai_agents/_openai_runner.py +++ b/python/openai_agents/src/temporalio/openai_agents/_openai_runner.py @@ -40,7 +40,11 @@ capture_scope, install_model_provider, ) -from temporalio.openai_agents._temporal_model_stub import _TemporalModelStub +from temporalio.openai_agents._temporal_model_stub import ( + _current_streamed_run, + _StreamedRun, + _TemporalModelStub, +) from temporalio.openai_agents._trace_interceptor import _refuse_tracing_key from temporalio.openai_agents.sandbox._temporal_sandbox_client import ( TemporalSandboxClient, @@ -474,12 +478,22 @@ def run_streamed( # The streaming loop runs in a task created inside run_streamed, which copies # this context, so the scope only needs to cover the call itself. - with capture_scope(run_config.model_settings), _inline_sync_function_tools(): - streamed_result = self._runner.run_streamed( - starting_agent=converted_agent, - input=input, - **kwargs, - ) + streamed_run = _StreamedRun() + token = _current_streamed_run.set(streamed_run) + try: + with ( + capture_scope(run_config.model_settings), + _inline_sync_function_tools(), + ): + streamed_result = self._runner.run_streamed( + starting_agent=converted_agent, + input=input, + **kwargs, + ) + finally: + _current_streamed_run.reset(token) + # The model stub reads it to watch the run's parallel input guardrails. + streamed_run.result = streamed_result # Mirror the AgentsException -> AgentsWorkflowError rewrap done # in run() above. The streaming runner attaches the actual run diff --git a/python/openai_agents/src/temporalio/openai_agents/_temporal_model_stub.py b/python/openai_agents/src/temporalio/openai_agents/_temporal_model_stub.py index 56a36614..c3428b0b 100644 --- a/python/openai_agents/src/temporalio/openai_agents/_temporal_model_stub.py +++ b/python/openai_agents/src/temporalio/openai_agents/_temporal_model_stub.py @@ -1,6 +1,8 @@ from __future__ import annotations +import asyncio from collections.abc import AsyncIterator, Callable +from contextvars import ContextVar from typing import Any from agents import ( @@ -13,10 +15,13 @@ Handoff, HostedMCPTool, ImageGenerationTool, + InputGuardrailResult, + InputGuardrailTripwireTriggered, Model, ModelResponse, ModelSettings, ModelTracing, + RunResultStreaming, Tool, TResponseInputItem, WebSearchTool, @@ -32,6 +37,7 @@ from openai.types.responses.response_prompt_param import ResponsePromptParam from temporalio import workflow +from temporalio.exceptions import ActivityError from temporalio.openai_agents._errors import AgentsWorkflowError from temporalio.openai_agents._invoke_model_activity import ( ActivityModelInput, @@ -50,6 +56,28 @@ from temporalio.openai_agents._model_parameters import ModelActivityParameters +class _StreamedRun: + """Holds the ``RunResultStreaming`` of the streamed run in progress.""" + + result: RunResultStreaming | None = None + + +# Set by the runner around ``Runner.run_streamed``; the run loop task copies it. +_current_streamed_run: ContextVar[_StreamedRun | None] = ContextVar( + "temporal_openai_agents_streamed_run", default=None +) + + +def _parallel_input_guardrails() -> tuple[RunResultStreaming, asyncio.Task[Any]] | None: + """Return the streamed run in progress and its parallel input guardrails task, if any.""" + streamed_run = _current_streamed_run.get() + result = streamed_run.result if streamed_run is not None else None + task = result._input_guardrails_task if result is not None else None + if result is None or task is None: + return None + return result, task + + class _TemporalModelStub(Model): # type:ignore[reportUnusedClass] """A stub that allows invoking models as Temporal activities.""" @@ -318,7 +346,19 @@ async def stream_response( "streaming_batch_interval": self.model_params.streaming_batch_interval, } - events = await workflow.execute_activity_method( + # The activity publishes each event to the topic as the model sends it, but the + # framework only checks the parallel input guardrails against the events it yields + # after the activity returns. So once one of them trips, don't start the activity, + # or cancel it: it then publishes an error event and stops. The framework records + # the trip before it waits for the other guardrails to clean up, so act on the + # record without waiting for the guardrails task to finish. + guardrails = _parallel_input_guardrails() + if guardrails is not None: + triggered = guardrails[0]._triggered_input_guardrail_result + if triggered is not None: + raise InputGuardrailTripwireTriggered(triggered) + + handle = workflow.start_activity_method( ModelActivity.invoke_model_activity_streaming, streaming_input, summary=summary, @@ -332,5 +372,33 @@ async def stream_response( versioning_intent=self.model_params.versioning_intent, priority=self.model_params.priority, ) + + cancelled_for: list[InputGuardrailResult] = [] + watcher: asyncio.Task[None] | None = None + if guardrails is not None and not guardrails[1].done(): + result, task = guardrails + + async def cancel_on_tripwire() -> None: + await workflow.wait_condition( + lambda: ( + result._triggered_input_guardrail_result is not None + or task.done() + ) + ) + triggered = result._triggered_input_guardrail_result + if triggered is not None and not handle.done(): + cancelled_for.append(triggered) + handle.cancel() + + watcher = asyncio.create_task(cancel_on_tripwire()) + try: + events = await handle + except ActivityError: + if cancelled_for: + raise InputGuardrailTripwireTriggered(cancelled_for[0]) from None + raise + finally: + if watcher is not None: + watcher.cancel() for event in events: yield event diff --git a/python/openai_agents/tests/test_openai_streaming.py b/python/openai_agents/tests/test_openai_streaming.py index 8996a338..0369c16d 100644 --- a/python/openai_agents/tests/test_openai_streaming.py +++ b/python/openai_agents/tests/test_openai_streaming.py @@ -10,31 +10,39 @@ import logging import uuid from collections.abc import AsyncIterator +from dataclasses import dataclass from datetime import timedelta from typing import Any import pytest from agents import ( Agent, + AgentHookContext, + AgentHooks, AgentOutputSchemaBase, + GuardrailFunctionOutput, Handoff, + InputGuardrailTripwireTriggered, Model, ModelResponse, ModelSettings, ModelTracing, RunConfig, + RunContextWrapper, Runner, Tool, ToolCallOutputItem, TResponseInputItem, Usage, function_tool, + input_guardrail, ) from agents.items import TResponseStreamEvent from agents.models import get_default_model_settings from openai.types.responses import ( Response, ResponseCompletedEvent, + ResponseErrorEvent, ResponseFunctionToolCall, ResponseOutputMessage, ResponseOutputText, @@ -49,12 +57,13 @@ from openai.types.shared import Reasoning from openai.types.shared.response_format_text import ResponseFormatText -from temporalio import workflow -from temporalio.client import Client, WorkflowFailureError +from temporalio import activity, workflow +from temporalio.client import Client, WorkflowFailureError, WorkflowHandle from temporalio.contrib.workflow_streams import WorkflowStream, WorkflowStreamClient -from temporalio.openai_agents import ModelActivityParameters +from temporalio.openai_agents import ModelActivityParameters, _heartbeat_decorator from temporalio.openai_agents.testing import AgentEnvironment -from tests.helpers import new_worker +from temporalio.workflow import ActivityCancellationType +from tests.helpers import assert_eventually, new_worker logger = logging.getLogger(__name__) @@ -528,3 +537,454 @@ async def test_streamed_run_runs_sync_function_tools(client: Client): ) # The sync tool ran in the streamed run, so the model got its result. assert outputs == [5] + + +class HoldingStreamingModel(StreamingTestModel): + """Streaming test model that sends its first delta, then, on the activity's first + attempt, waits up to ``hold`` seconds (10 by default) before the rest. + + Records the attempts and whether its stream was cancelled during the wait. + """ + + def __init__(self, hold: float = 10) -> None: + self.hold = hold + self.cancelled = False + self.attempts: list[int] = [] + + async def stream_response( + self, + system_instructions: str | None, + input: str | list[TResponseInputItem], + model_settings: ModelSettings, + tools: list[Tool], + output_schema: AgentOutputSchemaBase | None, + handoffs: list[Handoff], + tracing: ModelTracing, + **kwargs: Any, + ) -> AsyncIterator[TResponseStreamEvent]: + events = super().stream_response( + system_instructions, + input, + model_settings, + tools, + output_schema, + handoffs, + tracing, + **kwargs, + ) + attempt = activity.info().attempt + self.attempts.append(attempt) + yield await anext(events) + if attempt == 1: + try: + await asyncio.sleep(self.hold) + except asyncio.CancelledError: + self.cancelled = True + raise + async for event in events: + yield event + + +@input_guardrail +async def refuse_after_a_second( + _context: RunContextWrapper[None], + _agent: Agent, + _input: str | list[TResponseInputItem], +) -> GuardrailFunctionOutput: + """Runs in parallel with the model call (the default) and trips after one second.""" + await workflow.sleep(1) + return GuardrailFunctionOutput(output_info=None, tripwire_triggered=True) + + +@input_guardrail +async def clean_up_slowly( + _context: RunContextWrapper[None], + _agent: Agent, + _input: str | list[TResponseInputItem], +) -> GuardrailFunctionOutput: + """A parallel guardrail that takes 12 s to clean up when the framework cancels it.""" + try: + await workflow.sleep(60) + except asyncio.CancelledError: + await workflow.sleep(12) + raise + return GuardrailFunctionOutput(output_info=None, tripwire_triggered=False) + + +class DelayModelCall(AgentHooks[None]): + """Delays the model call by two seconds, so the guardrail trips before it.""" + + async def on_start(self, context: AgentHookContext[None], agent: Agent) -> None: + await workflow.sleep(2) + + +@dataclass +class GuardrailRun: + trip: bool = True + sibling: bool = False + delay_model: bool = False + # Stay open until the subscriber has read the activity's last event. + hold_open: bool = True + + +@workflow.defn +class StreamingInputGuardrailWorkflow: + @workflow.init + def __init__(self, run: GuardrailRun) -> None: + self.stream = WorkflowStream() + self.subscriber_done = False + + @workflow.signal + def finish(self) -> None: + self.subscriber_done = True + + @workflow.run + async def run(self, run: GuardrailRun) -> str: + guardrails = [refuse_after_a_second] if run.trip else [] + if run.sibling: + guardrails.append(clean_up_slowly) + agent = Agent[None]( + name="Assistant", + instructions="You are a test agent.", + input_guardrails=guardrails, + hooks=DelayModelCall() if run.delay_model else None, + ) + result = Runner.run_streamed(starting_agent=agent, input="Hello") + try: + async for _ in result.stream_events(): + pass + except InputGuardrailTripwireTriggered: + output = "refused" + else: + output = str(result.final_output) + if run.hold_open: + await workflow.wait_condition(lambda: self.subscriber_done) + return output + + +async def collect_until_the_end( + client: Client, + workflow_id: str, + published: list[TResponseStreamEvent] | None = None, +) -> list[TResponseStreamEvent]: + """Read the topic into ``published`` until an ``error`` or ``response.completed`` event.""" + stream = WorkflowStreamClient.create(client, workflow_id) + if published is None: + published = [] + async for item in stream.subscribe( + ["events"], + from_offset=0, + result_type=TResponseStreamEvent, # type: ignore[arg-type,call-overload] + poll_cooldown=timedelta(milliseconds=50), + ): + published.append(item.data) + if item.data.type in ("error", "response.completed"): + break + return published + + +async def streaming_activity_count(handle: WorkflowHandle[Any, Any]) -> int: + count = 0 + async for event in handle.fetch_history_events(): + if ( + event.HasField("activity_task_scheduled_event_attributes") + and event.activity_task_scheduled_event_attributes.activity_type.name + == "invoke_model_activity_streaming" + ): + count += 1 + return count + + +@pytest.mark.asyncio +@pytest.mark.parametrize("sibling", [False, True], ids=["alone", "sibling-cleaning-up"]) +async def test_streaming_input_guardrail_trip_stops_the_topic( + client: Client, sibling: bool +): + """A parallel input guardrail that trips while the model is streaming cancels the + streaming activity: topic subscribers get an error event, not the rest of the answer + and ``response.completed``. The cancel does not wait for the other parallel + guardrails to finish cleaning up.""" + model = HoldingStreamingModel() + async with AgentEnvironment( + model=model, + model_params=ModelActivityParameters( + start_to_close_timeout=timedelta(seconds=30), + heartbeat_timeout=timedelta(seconds=3), + streaming_topic="events", + ), + ) as env: + client = env.applied_on_client(client) + workflow_id = f"openai-streaming-guardrail-{uuid.uuid4()}" + + async with new_worker( + client, StreamingInputGuardrailWorkflow, max_cached_workflows=0 + ) as worker: + handle = await client.start_workflow( + StreamingInputGuardrailWorkflow.run, + GuardrailRun(sibling=sibling), + id=workflow_id, + task_queue=worker.task_queue, + execution_timeout=timedelta(seconds=50), + ) + published = await asyncio.wait_for( + collect_until_the_end(client, workflow_id), timeout=20.0 + ) + await handle.signal(StreamingInputGuardrailWorkflow.finish) + result = await handle.result() + + assert result == "refused" + published_types = [e.type for e in published] + assert "response.completed" not in published_types, published_types + last = published[-1] + assert isinstance(last, ResponseErrorEvent) + assert last.code == "cancelled" + assert model.cancelled + assert model.attempts == [1] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("sibling", [False, True], ids=["alone", "sibling-cleaning-up"]) +async def test_streaming_input_guardrail_known_trip_skips_the_model( + client: Client, sibling: bool +): + """When a parallel input guardrail has already tripped as the model is called, the + streaming activity is not scheduled, even while other guardrails are still cleaning + up.""" + model = HoldingStreamingModel() + async with AgentEnvironment( + model=model, + model_params=ModelActivityParameters( + start_to_close_timeout=timedelta(seconds=30), + heartbeat_timeout=timedelta(seconds=3), + streaming_topic="events", + ), + ) as env: + client = env.applied_on_client(client) + async with new_worker( + client, StreamingInputGuardrailWorkflow, max_cached_workflows=0 + ) as worker: + handle = await client.start_workflow( + StreamingInputGuardrailWorkflow.run, + GuardrailRun(sibling=sibling, delay_model=True, hold_open=False), + id=f"openai-streaming-known-trip-{uuid.uuid4()}", + task_queue=worker.task_queue, + execution_timeout=timedelta(seconds=50), + ) + result = await handle.result() + scheduled = await streaming_activity_count(handle) + + assert result == "refused" + assert scheduled == 0 + assert model.attempts == [] + + +def delay_stream_publishing(monkeypatch: pytest.MonkeyPatch, seconds: float) -> None: + """Hold every stream publish signal until ``seconds`` after the first one is sent.""" + signal = WorkflowHandle.signal + deadline: float | None = None + + async def delayed_signal( + self: WorkflowHandle[Any, Any], name: Any, *args: Any, **kwargs: Any + ) -> None: + nonlocal deadline + if name == "__temporal_workflow_stream_publish": + loop = asyncio.get_running_loop() + if deadline is None: + deadline = loop.time() + seconds + await asyncio.sleep(max(0.0, deadline - loop.time())) + await signal(self, name, *args, **kwargs) + + monkeypatch.setattr(WorkflowHandle, "signal", delayed_signal) + + +def clean_up_first_publish_slowly( + monkeypatch: pytest.MonkeyPatch, seconds: float +) -> None: + """Hold the first stream publish signal until it is cancelled, then take ``seconds`` + to clean up before it raises.""" + signal = WorkflowHandle.signal + first = True + + async def slow_cleanup_signal( + self: WorkflowHandle[Any, Any], name: Any, *args: Any, **kwargs: Any + ) -> None: + nonlocal first + if name == "__temporal_workflow_stream_publish" and first: + first = False + try: + await asyncio.sleep(60) + except asyncio.CancelledError: + await asyncio.sleep(seconds) + raise + await signal(self, name, *args, **kwargs) + + monkeypatch.setattr(WorkflowHandle, "signal", slow_cleanup_signal) + + +def stop_heartbeats_slowly(monkeypatch: pytest.MonkeyPatch, seconds: float) -> None: + """Make stopping the activity's heartbeats take ``seconds``, still heartbeating + meanwhile so that a cancel reaches the activity.""" + heartbeat_every = _heartbeat_decorator._heartbeat_every + + async def slow_to_stop(delay: float) -> None: + try: + await heartbeat_every(delay) + except asyncio.CancelledError: + loop = asyncio.get_running_loop() + deadline = loop.time() + seconds + while loop.time() < deadline: + activity.heartbeat() + await asyncio.sleep(0.5) + raise + + monkeypatch.setattr(_heartbeat_decorator, "_heartbeat_every", slow_to_stop) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("max_cached_workflows", [1000, 0], ids=["cached", "uncached"]) +@pytest.mark.parametrize( + "cancelled_while", + ["streaming", "publishing", "stopping-the-flusher", "stopping-the-heartbeats"], + ids=["streaming", "publishing", "stopping-the-flusher", "stopping-the-heartbeats"], +) +async def test_streaming_input_guardrail_trip_waits_for_the_error_event( + client: Client, + monkeypatch: pytest.MonkeyPatch, + max_cached_workflows: int, + cancelled_while: str, +): + """With ``WAIT_CANCELLATION_COMPLETED``, a workflow that returns right after the trip + still has the error event at the end of its stream, and the Worker shuts down after + the cancel, with or without a Workflow cache. That holds whether the cancel reaches + the activity while the model streams or after the model finished, while the activity + still publishes its events, stops the stream's background flusher or stops its + heartbeats.""" + model: StreamingTestModel + if cancelled_while == "streaming": + model = HoldingStreamingModel() + elif cancelled_while == "publishing": + # The model finishes at once, but the topic gets nothing for 8 s, so the cancel + # arrives while the activity is still publishing. + model = StreamingTestModel() + delay_stream_publishing(monkeypatch, 8) + elif cancelled_while == "stopping-the-flusher": + # The background flusher's first publish hangs. The model finishes after + # 0.3 s, and stopping the flusher takes 8 s, so the cancel arrives while the + # stream's context manager waits for the flusher. + model = HoldingStreamingModel(hold=0.3) + clean_up_first_publish_slowly(monkeypatch, 8) + else: + # The model finishes after 0.3 s and its events are published, but stopping + # the heartbeats takes 8 s, so the cancel arrives while the activity waits for + # its heartbeat task. + model = HoldingStreamingModel(hold=0.3) + stop_heartbeats_slowly(monkeypatch, 8) + async with AgentEnvironment( + model=model, + model_params=ModelActivityParameters( + start_to_close_timeout=timedelta(seconds=30), + heartbeat_timeout=timedelta(seconds=3), + streaming_topic="events", + cancellation_type=ActivityCancellationType.WAIT_CANCELLATION_COMPLETED, + ), + ) as env: + client = env.applied_on_client(client) + workflow_id = f"openai-streaming-guardrail-wait-{uuid.uuid4()}" + async with new_worker( + client, + StreamingInputGuardrailWorkflow, + max_cached_workflows=max_cached_workflows, + ) as worker: + result = await client.execute_workflow( + StreamingInputGuardrailWorkflow.run, + GuardrailRun(hold_open=False), + id=workflow_id, + task_queue=worker.task_queue, + execution_timeout=timedelta(seconds=30), + ) + stream = WorkflowStreamClient.create(client, workflow_id) + offset = await stream.get_offset() + last: TResponseStreamEvent | None = None + async for item in stream.subscribe( + ["events"], + from_offset=offset - 1, + result_type=TResponseStreamEvent, # type: ignore[arg-type,call-overload] + ): + last = item.data + break + + assert result == "refused" + assert isinstance(last, ResponseErrorEvent) + assert last.code == "cancelled" + if cancelled_while == "streaming": + # The model held after its first delta and was cancelled, so the stream holds + # that delta and the error event. + assert isinstance(model, HoldingStreamingModel) + assert model.cancelled + assert offset == 2 + else: + # Two deltas, response.completed, and the error event. + if isinstance(model, HoldingStreamingModel): + assert not model.cancelled + assert offset == 4 + + +@pytest.mark.asyncio +async def test_streaming_worker_shutdown_publishes_no_error_event(client: Client): + """A Worker shutdown cancels the streaming activity, but its retry can still answer, + so subscribers get no error event before the retry's events.""" + model = HoldingStreamingModel() + async with AgentEnvironment( + model=model, + model_params=ModelActivityParameters( + start_to_close_timeout=timedelta(seconds=30), + heartbeat_timeout=timedelta(seconds=3), + streaming_topic="events", + ), + ) as env: + client = env.applied_on_client(client) + workflow_id = f"openai-streaming-shutdown-{uuid.uuid4()}" + task_queue = str(uuid.uuid4()) + + # Only one Worker per task queue can be registered in a process at a time. + async with new_worker( + client, + StreamingInputGuardrailWorkflow, + task_queue=task_queue, + max_cached_workflows=0, + ): + handle = await client.start_workflow( + StreamingInputGuardrailWorkflow.run, + GuardrailRun(trip=False), + id=workflow_id, + task_queue=task_queue, + execution_timeout=timedelta(seconds=50), + ) + published: list[TResponseStreamEvent] = [] + collect_task = asyncio.create_task( + collect_until_the_end(client, workflow_id, published) + ) + + async def first_delta_read() -> None: + assert published + + # Shut the Worker down once the subscriber has the first attempt's delta. + await assert_eventually(first_delta_read, timeout=timedelta(seconds=20)) + async with new_worker( + client, + StreamingInputGuardrailWorkflow, + task_queue=task_queue, + max_cached_workflows=0, + ): + await asyncio.wait_for(collect_task, timeout=30.0) + await handle.signal(StreamingInputGuardrailWorkflow.finish) + result = await handle.result() + + assert result == "Hello world!" + assert model.attempts == [1, 2] + assert [e.type for e in published] == [ + "response.output_text.delta", + "response.output_text.delta", + "response.output_text.delta", + "response.completed", + ] diff --git a/python/openai_agents/uv.lock b/python/openai_agents/uv.lock index 636c9aee..01066cff 100644 --- a/python/openai_agents/uv.lock +++ b/python/openai_agents/uv.lock @@ -2942,7 +2942,7 @@ wheels = [ [[package]] name = "temporalio" -version = "1.34.0" +version = "1.35.0" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "nexus-rpc" }, @@ -2951,15 +2951,15 @@ dependencies = [ { name = "types-protobuf" }, { name = "typing-extensions" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/13/85/9ece6f400552ed21951c8c5c534056e7d015cb065e3384470ec23ac24f28/temporalio-1.34.0.tar.gz", hash = "sha256:6453cb20e18df485e16578b22c82a9c4bcb1cf7eedd94147dfd373551d80f5b6", size = 3163890, upload-time = "2026-09-30T20:22:56.056Z" } +sdist = { url = "https://files.pythonhosted.org/packages/69/38/03641a02eb18e2b16f58c1c136deafdffeb30828dcd83ebd8ac41eb39fd7/temporalio-1.35.0.tar.gz", hash = "sha256:c7a14f29909f8d847f78de49305af89c0ec5aea39272cd4050124386ac3ca295", size = 2879693, upload-time = "2026-10-28T20:13:38.726Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/3c/71/e55380e7820819357afada375d4e5c33d141bb94a08fda35aa0551b4369b/temporalio-1.34.0-cp310-abi3-macosx_10_12_x86_64.whl", hash = "sha256:87118447ad13e1062b79bfc8b44b1695e8328ac9c6be1776e426e87b501b135a", size = 14064102, upload-time = "2026-09-30T20:22:35.601Z" }, - { url = "https://files.pythonhosted.org/packages/db/78/356c8e2f0ec1678b757d4c5e8a1b0b1c65b15b8afb2335ad2944c158ea10/temporalio-1.34.0-cp310-abi3-macosx_11_0_arm64.whl", hash = "sha256:023fff9cd9dd21860061e003880dcf95250700f5afab96c030c9cb258504dab8", size = 13763219, upload-time = "2026-09-30T20:22:38.774Z" }, - { url = "https://files.pythonhosted.org/packages/e9/00/4a4b4e018c4c12b4691211aca1d97db8707ed4c9e31f40dcf09784c9913d/temporalio-1.34.0-cp310-abi3-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:cdb6e2fb525ea5635afa4a3f1ae4c992c16dbe29dc92939a85153baeb69d95a2", size = 14121901, upload-time = "2026-09-30T20:22:41.428Z" }, - { url = "https://files.pythonhosted.org/packages/c1/20/b032d4a0df51d466fd929d8205e1288e83d97a890f021479f43af9bd6f9b/temporalio-1.34.0-cp310-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:540761f738bdfe5cb5bd7240b659e116a0b5094b94282aef09b8d8c2d66e9c52", size = 14457183, upload-time = "2026-09-30T20:22:44.215Z" }, - { url = "https://files.pythonhosted.org/packages/b6/0d/b08ffeca93bbc29200a5bd85a45cca3c6be0a123de605ecd1798b617df4f/temporalio-1.34.0-cp310-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:87647f87f42ecd45efb675642e2ec8ca34aa42359f5968f31f5391f1db5ebcbe", size = 14179836, upload-time = "2026-09-30T20:22:47.474Z" }, - { url = "https://files.pythonhosted.org/packages/81/55/acde4d1b7c9f23434e263752ef0a3ef58adeb3b95ef8472607414798dfaa/temporalio-1.34.0-cp310-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:71caa4b9061628b22a457c87c3b16da40f7ac6ac1a26ece9aa1ee6b4934404d7", size = 14570248, upload-time = "2026-09-30T20:22:50.569Z" }, - { url = "https://files.pythonhosted.org/packages/cc/c5/6bdc3b02ec483093a71e9f887e8a50ae20fcc79e0578ed7435a81d4271b1/temporalio-1.34.0-cp310-abi3-win_amd64.whl", hash = "sha256:03bd86561188c18d88425178bc690fe0791104b78566dcebd98f59cbdbce0952", size = 15415684, upload-time = "2026-09-30T20:22:53.342Z" }, + { url = "https://files.pythonhosted.org/packages/9e/d9/8bda464beda3adb897d8d9f0e5fa26b28f2bf63c831c6800cbd73f7728cb/temporalio-1.35.0-cp310-abi3-macosx_10_12_x86_64.whl", hash = "sha256:ac3ab920cc05e24a1d420688a49aabed7b2be8f25dd275200dfa839be30980fd", size = 13989719, upload-time = "2026-10-28T20:13:16.897Z" }, + { url = "https://files.pythonhosted.org/packages/26/c1/94861d72cbe5e1778fcd4685db63231af01d40ad99f0375b4df6d4e60d06/temporalio-1.35.0-cp310-abi3-macosx_11_0_arm64.whl", hash = "sha256:51bae017d484450084daa19afafad12fa0af1a00af351ba304104d7b97b44f9a", size = 13552491, upload-time = "2026-10-28T20:13:19.725Z" }, + { url = "https://files.pythonhosted.org/packages/3c/ed/a776afec1033f749ec8f4c0c5264167aedc56a108a5adebfa4084af8aa82/temporalio-1.35.0-cp310-abi3-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:b93845a2200d23b8edc9a8b0bf78b58bdb5e78721602b3f64072867e07e28504", size = 13927445, upload-time = "2026-10-28T20:13:21.984Z" }, + { url = "https://files.pythonhosted.org/packages/ca/0b/bf0ed1f1b623e1ffd5c21c9e65a331a5c1c4f77f3829dc99ceae99b69f39/temporalio-1.35.0-cp310-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:45b5438f41ab154eda4eea40069df2fa8779cb8f5c27f266b917a700bfba8bb8", size = 14386743, upload-time = "2026-10-28T20:13:28.201Z" }, + { url = "https://files.pythonhosted.org/packages/76/d6/373d5ffe28263fc56ea0013bdd45715b4ff84f97d1052519b7c0d9ecd437/temporalio-1.35.0-cp310-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:8a5ce8b1e9c9e2df5fb9f17a345c0ef0a132c787cef8cf5a6593b5d707648253", size = 13988280, upload-time = "2026-10-28T20:13:31.667Z" }, + { url = "https://files.pythonhosted.org/packages/9c/41/47ad6f26002f84d8293d99098c414f75e9e8ff7955d034b7121bd211649b/temporalio-1.35.0-cp310-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:96310ae91b098c9fe2194f0892cef25381e8f870e48d5c069230312bbd1c6262", size = 14494297, upload-time = "2026-10-28T20:13:33.988Z" }, + { url = "https://files.pythonhosted.org/packages/db/33/30942b2d568df2c50089e381ee63b2d141f2ebfb7b4787bac3540f5cf617/temporalio-1.35.0-cp310-abi3-win_amd64.whl", hash = "sha256:6f434c0d3b6d049f4473f9f5d3e80b5e80c986440fd0490f2ee70cc631cbcddd", size = 15317073, upload-time = "2026-10-28T20:13:36.195Z" }, ] [package.optional-dependencies] @@ -3045,7 +3045,7 @@ requires-dist = [ { name = "openinference-instrumentation-openai-agents", marker = "extra == 'openinference'", specifier = ">=2,<3" }, { name = "opentelemetry-sdk", marker = "extra == 'openinference'", specifier = ">=1.26,<2" }, { name = "pydantic", specifier = ">=2,<3" }, - { name = "temporalio", extras = ["opentelemetry", "pydantic"], specifier = ">=1.34.0,<2" }, + { name = "temporalio", extras = ["opentelemetry", "pydantic"], specifier = ">=1.35.0,<2" }, { name = "temporalio-mcp", marker = "extra == 'mcp'", specifier = ">=0.2.0,<0.3" }, { name = "typing-extensions", specifier = ">=4.2,<5" }, ]