Use interceptors instead of forking open ai agents plugin

Author: danielmillerp

src/agentex/lib/core/clients/temporal/utils.py

@@ -6,6 +6,7 @@
 from temporalio.runtime import Runtime, TelemetryConfig, OpenTelemetryConfig
 from temporalio.contrib.pydantic import pydantic_data_converter
 from temporalio.contrib.openai_agents import OpenAIAgentsPlugin
+from temporalio.worker import Interceptor
 
 # class DateTimeJSONEncoder(AdvancedJSONEncoder):
 #     def default(self, o: Any) -> Any:
@@ -61,6 +62,24 @@ def validate_client_plugins(plugins: list[Any]) -> None:
             )
 
 
+def validate_worker_interceptors(interceptors: list[Any]) -> None:
+    """
+    Validate that all items in the interceptors list are valid Temporal worker interceptors.
+
+    Args:
+        interceptors: List of interceptors to validate
+
+    Raises:
+        TypeError: If any interceptor is not a valid Interceptor instance
+    """
+    for i, interceptor in enumerate(interceptors):
+        if not isinstance(interceptor, Interceptor):
+            raise TypeError(
+                f"Interceptor at index {i} must be an instance of temporalio.worker.Interceptor, "
+                f"got {type(interceptor).__name__}"
+            )
+
+
 async def get_temporal_client(temporal_address: str, metrics_url: str | None = None, plugins: list[Any] = []) -> Client:
     """
     Create a Temporal client with plugin integration.

src/agentex/lib/core/temporal/plugins/__init__.py

@@ -0,0 +1,55 @@
+"""OpenAI Agents SDK Temporal Plugin with Streaming Support.
+
+This module provides streaming capabilities for the OpenAI Agents SDK in Temporal
+using interceptors to thread task_id through workflows to activities.
+
+The streaming implementation works by:
+1. Using Temporal interceptors to thread task_id through the execution
+2. Streaming LLM responses to Redis in real-time from activities
+3. Returning complete responses to maintain Temporal determinism
+
+Example:
+    >>> from agentex.lib.core.temporal.plugins.openai_agents import (
+    ...     StreamingModelProvider,
+    ...     StreamingInterceptor,
+    ... )
+    >>> from temporalio.contrib.openai_agents import OpenAIAgentsPlugin, ModelActivityParameters
+    >>> from datetime import timedelta
+    >>>
+    >>> # Create streaming model provider
+    >>> model_provider = StreamingModelProvider()
+    >>>
+    >>> # Create STANDARD plugin with streaming model provider
+    >>> plugin = OpenAIAgentsPlugin(
+    ...     model_params=ModelActivityParameters(
+    ...         start_to_close_timeout=timedelta(seconds=120),
+    ...     ),
+    ...     model_provider=model_provider,
+    ... )
+    >>>
+    >>> # Register interceptor with worker
+    >>> interceptor = StreamingInterceptor()
+    >>> # Add interceptor to worker configuration
+"""
+
+from agentex.lib.core.temporal.plugins.openai_agents import (
+    StreamingModel,
+    StreamingModelProvider,
+    StreamingInterceptor,
+    streaming_task_id,
+    streaming_trace_id,
+    streaming_parent_span_id,
+    TemporalStreamingHooks,
+    stream_lifecycle_content,
+)
+
+__all__ = [
+    "StreamingModel",
+    "StreamingModelProvider",
+    "StreamingInterceptor",
+    "streaming_task_id",
+    "streaming_trace_id",
+    "streaming_parent_span_id",
+    "TemporalStreamingHooks",
+    "stream_lifecycle_content",
+]

src/agentex/lib/core/temporal/plugins/openai_agents/README.md

@@ -1,62 +1,65 @@
 """OpenAI Agents SDK Temporal Plugin with Streaming Support.
 
-This module provides a custom Temporal plugin for the OpenAI Agents SDK that adds
-real-time streaming capabilities while maintaining Temporal's durability guarantees.
+This module provides streaming capabilities for the OpenAI Agents SDK in Temporal
+using interceptors to thread task_id through workflows to activities.
 
 The streaming implementation works by:
-1. Threading a task_id through the workflow execution
+1. Using Temporal interceptors to thread task_id through the execution
 2. Streaming LLM responses to Redis in real-time from activities
 3. Streaming lifecycle events (tool calls, handoffs) via hooks and activities
 4. Returning complete responses to maintain Temporal determinism
 
 Example - Complete Setup:
     >>> from agentex.lib.core.temporal.plugins.openai_agents import (
-    ...     CustomStreamingOpenAIAgentsPlugin,
     ...     StreamingModelProvider,
     ...     TemporalStreamingHooks,
+    ...     StreamingInterceptor,
     ... )
-    >>> from temporalio.contrib.openai_agents import ModelActivityParameters
+    >>> from temporalio.contrib.openai_agents import OpenAIAgentsPlugin, ModelActivityParameters
     >>> from datetime import timedelta
     >>> from agents import Agent, Runner
     >>>
     >>> # 1. Create streaming model provider
     >>> model_provider = StreamingModelProvider()
     >>>
-    >>> # 2. Create plugin with streaming support
-    >>> plugin = CustomStreamingOpenAIAgentsPlugin(
+    >>> # 2. Create STANDARD plugin with streaming model provider
+    >>> plugin = OpenAIAgentsPlugin(
     ...     model_params=ModelActivityParameters(
     ...         start_to_close_timeout=timedelta(seconds=120),
     ...     ),
     ...     model_provider=model_provider,
     ... )
     >>>
-    >>> # 3. In workflow, create hooks for streaming lifecycle events
-    >>> # The hooks automatically use the built-in streaming activity
+    >>> # 3. Register interceptor with worker
+    >>> interceptor = StreamingInterceptor()
+    >>> # Add interceptor to worker configuration
+    >>>
+    >>> # 4. In workflow, store task_id in instance variable
+    >>> self._task_id = params.task.id
+    >>>
+    >>> # 5. Create hooks for streaming lifecycle events
     >>> hooks = TemporalStreamingHooks(task_id="your-task-id")
     >>>
-    >>> # 4. Run agent with streaming context and hooks
-    >>> context = {"task_id": "your-task-id"}
-    >>> result = await Runner.run(agent, input, context=context, hooks=hooks)
+    >>> # 6. Run agent - interceptor handles task_id threading automatically
+    >>> result = await Runner.run(agent, input, hooks=hooks)
 
 This gives you:
-- Real-time streaming of LLM responses (via StreamingModel)
+- Real-time streaming of LLM responses (via StreamingModel + interceptors)
 - Real-time streaming of tool calls (via TemporalStreamingHooks)
 - Real-time streaming of agent handoffs (via TemporalStreamingHooks)
 - Full Temporal durability and observability
+- No forked plugin required - uses standard OpenAIAgentsPlugin
 """
 
 from agentex.lib.core.temporal.plugins.openai_agents.streaming_model import (
     StreamingModel,
     StreamingModelProvider,
 )
-from agentex.lib.core.temporal.plugins.openai_agents.streaming_plugin import (
-    CustomStreamingOpenAIAgentsPlugin,
-    StreamingModelActivity,
-    StreamingActivityModelInput,
-)
-from agentex.lib.core.temporal.plugins.openai_agents.streaming_runner import (
-    StreamingTemporalRunner,
-    StreamingTemporalModelStub,
+from agentex.lib.core.temporal.plugins.openai_agents.streaming_interceptor import (
+    StreamingInterceptor,
+    streaming_task_id,
+    streaming_trace_id,
+    streaming_parent_span_id,
 )
 from agentex.lib.core.temporal.plugins.openai_agents.hooks import (
     TemporalStreamingHooks,
@@ -66,13 +69,12 @@
 )
 
 __all__ = [
-    "CustomStreamingOpenAIAgentsPlugin",
     "StreamingModel",
     "StreamingModelProvider",
-    "StreamingModelActivity",
-    "StreamingActivityModelInput",
-    "StreamingTemporalRunner",
-    "StreamingTemporalModelStub",
+    "StreamingInterceptor",
+    "streaming_task_id",
+    "streaming_trace_id",
+    "streaming_parent_span_id",
     "TemporalStreamingHooks",
     "stream_lifecycle_content",
 ]

src/agentex/lib/core/temporal/plugins/openai_agents/__init__.py

@@ -0,0 +1,152 @@
+"""
+Simple Temporal interceptors for threading task_id to enable streaming.
+
+This module provides minimal interceptors to pass task_id from workflows
+to activities via headers, making it available to the StreamingModel.
+"""
+
+from contextvars import ContextVar
+from typing import Optional, Any, Type
+import logging
+
+from temporalio import workflow
+from temporalio.worker import (
+    Interceptor,
+    WorkflowInboundInterceptor,
+    WorkflowOutboundInterceptor,
+    ActivityInboundInterceptor,
+    ExecuteWorkflowInput,
+    StartActivityInput,
+    ExecuteActivityInput,
+)
+from temporalio.converter import default
+
+# Set up logging
+logger = logging.getLogger("streaming.interceptor")
+
+# Global context variable that StreamingModel will read
+# This is thread-safe and works across async boundaries
+streaming_task_id: ContextVar[Optional[str]] = ContextVar('streaming_task_id', default=None)
+streaming_trace_id: ContextVar[Optional[str]] = ContextVar('streaming_trace_id', default=None)
+streaming_parent_span_id: ContextVar[Optional[str]] = ContextVar('streaming_parent_span_id', default=None)
+# Header key for passing task_id
+TASK_ID_HEADER = "streaming-task-id"
+TRACE_ID_HEADER = "trace-id"
+PARENT_SPAN_ID_HEADER = "parent-span-id"
+
+class StreamingInterceptor(Interceptor):
+    """Main interceptor that enables task_id threading."""
+
+    def __init__(self):
+        self._payload_converter = default().payload_converter
+        logger.info("[StreamingInterceptor] Initialized")
+
+    def intercept_activity(self, next: ActivityInboundInterceptor) -> ActivityInboundInterceptor:
+        """Create activity interceptor to read task_id from headers."""
+        return StreamingActivityInboundInterceptor(next, self._payload_converter)
+
+    def workflow_interceptor_class(self, input: Any) -> Optional[Type[WorkflowInboundInterceptor]]:
+        """Return workflow interceptor class."""
+        return StreamingWorkflowInboundInterceptor
+
+
+class StreamingWorkflowInboundInterceptor(WorkflowInboundInterceptor):
+    """Workflow interceptor that creates the outbound interceptor."""
+
+    def __init__(self, next: WorkflowInboundInterceptor):
+        super().__init__(next)
+        self._payload_converter = default().payload_converter
+
+    async def execute_workflow(self, input: ExecuteWorkflowInput) -> Any:
+        """Execute workflow - just pass through."""
+        return await self.next.execute_workflow(input)
+
+    def init(self, outbound: WorkflowOutboundInterceptor) -> None:
+        """Initialize with our custom outbound interceptor."""
+        self.next.init(StreamingWorkflowOutboundInterceptor(
+            outbound, self._payload_converter
+        ))
+
+
+class StreamingWorkflowOutboundInterceptor(WorkflowOutboundInterceptor):
+    """Outbound interceptor that adds task_id to activity headers."""
+
+    def __init__(self, next, payload_converter):
+        super().__init__(next)
+        self._payload_converter = payload_converter
+
+    def start_activity(self, input: StartActivityInput) -> workflow.ActivityHandle:
+        """Add task_id, trace_id, and parent_span_id to headers when starting model activities."""
+
+        # Only add headers for invoke_model_activity calls
+        activity_name = str(input.activity) if hasattr(input, 'activity') else ""
+
+        if "invoke_model_activity" in activity_name or "invoke-model-activity" in activity_name:
+            # Get task_id, trace_id, and parent_span_id from workflow instance instead of inbound interceptor
+            try:
+                workflow_instance = workflow.instance()
+                task_id = getattr(workflow_instance, '_task_id', None)
+                trace_id = getattr(workflow_instance, '_trace_id', None)
+                parent_span_id = getattr(workflow_instance, '_parent_span_id', None)
+
+                if task_id and trace_id and parent_span_id:
+                    # Initialize headers if needed
+                    if not input.headers:
+                        input.headers = {}
+
+                    # Add task_id to headers
+                    input.headers[TASK_ID_HEADER] = self._payload_converter.to_payload(task_id)
+                    input.headers[TRACE_ID_HEADER] = self._payload_converter.to_payload(trace_id)
+                    input.headers[PARENT_SPAN_ID_HEADER] = self._payload_converter.to_payload(parent_span_id)
+                    logger.debug(f"[OutboundInterceptor] Added task_id, trace_id, and parent_span_id to activity headers: {task_id}, {trace_id}, {parent_span_id}")
+                else:
+                    logger.warning("[OutboundInterceptor] No _task_id, _trace_id, or _parent_span_id found in workflow instance")
+            except Exception as e:
+                logger.error(f"[OutboundInterceptor] Failed to get task_id, trace_id, or parent_span_id from workflow instance: {e}")
+
+        return self.next.start_activity(input)
+
+
+class StreamingActivityInboundInterceptor(ActivityInboundInterceptor):
+    """Activity interceptor that extracts task_id, trace_id, and parent_span_id from headers and sets context variables."""
+
+    def __init__(self, next, payload_converter):
+        super().__init__(next)
+        self._payload_converter = payload_converter
+
+    async def execute_activity(self, input: ExecuteActivityInput) -> Any:
+        """Extract task_id, trace_id, and parent_span_id from headers and set context variables."""
+
+        # Extract task_id from headers if present
+        if input.headers and TASK_ID_HEADER in input.headers:
+            task_id_value = self._payload_converter.from_payload(
+                input.headers[TASK_ID_HEADER], str
+            )
+            trace_id_value = self._payload_converter.from_payload(
+                input.headers[TRACE_ID_HEADER], str
+            )
+            parent_span_id_value = self._payload_converter.from_payload(
+                input.headers[PARENT_SPAN_ID_HEADER], str
+            )
+
+            # P THIS IS THE KEY PART - Set the context variable!
+            # This makes task_id available to StreamingModel.get_response()
+            streaming_task_id.set(task_id_value)
+            streaming_trace_id.set(trace_id_value)
+            streaming_parent_span_id.set(parent_span_id_value)
+            logger.info(f"[ActivityInterceptor] Set task_id, trace_id, and parent_span_id in context: {task_id_value}, {trace_id_value}, {parent_span_id_value}")
+        else:
+            logger.debug("[ActivityInterceptor] No task_id, trace_id, or parent_span_id in headers")
+
+        try:
+            # Execute the activity
+            # The StreamingModel can now read streaming_task_id.get()
+            result = await self.next.execute_activity(input)
+            return result
+        finally:
+            # Clean up context after activity
+            streaming_task_id.set(None)
+            streaming_trace_id.set(None)
+            streaming_parent_span_id.set(None)
+            logger.debug("[ActivityInterceptor] Cleared task_id, trace_id, and parent_span_id from context")
+

src/agentex/lib/core/temporal/plugins/openai_agents/streaming_interceptor.py

@@ -59,7 +59,9 @@
 from agentex.types.task_message_delta import ReasoningSummaryDelta, TextDelta, ReasoningContentDelta
 from agentex.types.task_message_content import TextContent, ReasoningContent
 from agentex.types.task_message_update import StreamTaskMessageDelta, StreamTaskMessageFull
-
+from agentex.lib.core.temporal.plugins.openai_agents.streaming_interceptor import streaming_task_id
+from agentex.lib.core.temporal.plugins.openai_agents.streaming_interceptor import streaming_trace_id
+from agentex.lib.core.temporal.plugins.openai_agents.streaming_interceptor import streaming_parent_span_id
 # Create logger for this module
 logger = logging.getLogger("agentex.temporal.streaming")
 
@@ -383,20 +385,15 @@ async def get_response(
         This method is used by Temporal activities and needs to return a complete
         response, but we stream the response to Redis while generating it.
         """
-        # Get optional parameters from kwargs
-        trace_id = kwargs.get('trace_id')
-        parent_span_id = kwargs.get('parent_span_id')
-
-        # Use the class's tracer instance
-        tracer = self.tracer
+        
+        task_id = streaming_task_id.get()
+        trace_id = streaming_trace_id.get()
+        parent_span_id = streaming_parent_span_id.get()
 
-        # Use task_id as trace_id if not provided
-        task_id = kwargs.get('task_id')
-        if trace_id is None:
-            trace_id = task_id
-            logger.info(f"[StreamingModel] Using task_id as trace_id: {trace_id}")
+        if not task_id or not trace_id or not parent_span_id:
+            raise ValueError("task_id, trace_id, and parent_span_id are required for streaming with Responses API")
 
-        trace = tracer.trace(trace_id)
+        trace = self.tracer.trace(trace_id)
 
         async with trace.span(
             parent_id=parent_span_id,