Breaking: stop the OpenAI Agents streaming model activity when a parallel input guardrail trips

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.

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",

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)

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

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