Moving mcp classes into plugin

Author: tconley1428

temporalio/contrib/openai_agents/__init__.py

@@ -8,6 +8,7 @@
     Use with caution in production environments.
 """
 
+from temporalio.contrib.openai_agents._mcp import TemporalMCPServerWorkflowShim
 from temporalio.contrib.openai_agents._model_parameters import ModelActivityParameters
 from temporalio.contrib.openai_agents._temporal_openai_agents import (
     OpenAIAgentsPlugin,
@@ -24,6 +25,7 @@
     "OpenAIAgentsPlugin",
     "ModelActivityParameters",
     "workflow",
+    "TemporalMCPServerWorkflowShim",
     "TestModel",
     "TestModelProvider",
 ]

temporalio/contrib/openai_agents/_mcp.py

@@ -0,0 +1,115 @@
+from datetime import timedelta
+from typing import Any, Callable, Optional, Sequence
+
+from agents import AgentBase, RunContextWrapper
+from agents.mcp import MCPServer
+from mcp import GetPromptResult, ListPromptsResult
+from mcp import Tool as MCPTool
+from mcp.types import CallToolResult
+
+from temporalio import activity, workflow
+
+
+class TemporalMCPServerWorkflowShim(MCPServer):
+    def __init__(self, name: str):
+        self.server_name = name
+        super().__init__()
+
+    @property
+    def name(self) -> str:
+        return self.server_name
+
+    async def connect(self) -> None:
+        raise ValueError("Cannot connect to a server shim")
+
+    async def cleanup(self) -> None:
+        raise ValueError("Cannot clean up a server shim")
+
+    async def list_tools(
+        self,
+        run_context: Optional[RunContextWrapper[Any]] = None,
+        agent: Optional[AgentBase] = None,
+    ) -> list[MCPTool]:
+        workflow.logger.info("Listing tools")
+        tools: list[MCPTool] = await workflow.execute_local_activity(
+            self.name + "-list-tools",
+            start_to_close_timeout=timedelta(seconds=30),
+            result_type=list[MCPTool],
+        )
+        print(tools[0])
+        print("Tool type:", type(tools[0]))
+        # print(type(MCPTool(**tools[0])))
+        return tools
+
+    async def call_tool(
+        self, tool_name: str, arguments: Optional[dict[str, Any]]
+    ) -> CallToolResult:
+        return await workflow.execute_local_activity(
+            self.name + "-call-tool",
+            args=[tool_name, arguments],
+            start_to_close_timeout=timedelta(seconds=30),
+            result_type=CallToolResult,
+        )
+
+    async def list_prompts(self) -> ListPromptsResult:
+        raise NotImplementedError()
+
+    async def get_prompt(
+        self, name: str, arguments: Optional[dict[str, Any]] = None
+    ) -> GetPromptResult:
+        raise NotImplementedError()
+
+
+class TemporalMCPServer(TemporalMCPServerWorkflowShim):
+    def __init__(self, server: MCPServer):
+        self.server = server
+        super().__init__(server.name)
+
+    @property
+    def name(self) -> str:
+        return self.server.name
+
+    async def connect(self) -> None:
+        await self.server.connect()
+
+    async def cleanup(self) -> None:
+        await self.server.cleanup()
+
+    async def list_tools(
+        self,
+        run_context: Optional[RunContextWrapper[Any]] = None,
+        agent: Optional[AgentBase] = None,
+    ) -> list[MCPTool]:
+        if not workflow.in_workflow():
+            return await self.server.list_tools(run_context, agent)
+
+        return await super().list_tools(run_context, agent)
+
+    async def call_tool(
+        self, tool_name: str, arguments: Optional[dict[str, Any]]
+    ) -> CallToolResult:
+        if not workflow.in_workflow():
+            return await self.server.call_tool(tool_name, arguments)
+
+        return await super().call_tool(tool_name, arguments)
+
+    async def __aenter__(self):
+        await self.connect()
+        return self
+
+    async def __aexit__(self, exc_type, exc_value, traceback):
+        await self.cleanup()
+
+    def get_activities(self) -> Sequence[Callable]:
+        @activity.defn(name=self.name + "-list-tools")
+        async def list_tools() -> list[MCPTool]:
+            activity.logger.info("Listing tools in activity")
+            return await self.server.list_tools()
+
+        @activity.defn(name=self.name + "-call-tool")
+        async def call_tool(
+            tool_name: str, arguments: Optional[dict[str, Any]]
+        ) -> CallToolResult:
+            return await self.server.call_tool(tool_name, arguments)
+
+        return list_tools, call_tool

temporalio/contrib/openai_agents/_temporal_openai_agents.py

@@ -1,8 +1,9 @@
 """Initialize Temporal OpenAI Agents overrides."""
 
-from contextlib import contextmanager
+import dataclasses
+from contextlib import AsyncExitStack, contextmanager
 from datetime import timedelta
-from typing import AsyncIterator, Callable, Optional, Union
+from typing import AsyncIterator, Callable, Optional, Sequence, Union
 
 from agents import (
     AgentOutputSchemaBase,
@@ -17,6 +18,7 @@
     set_trace_provider,
 )
 from agents.items import TResponseStreamEvent
+from agents.mcp import MCPServer
 from agents.run import get_default_agent_runner, set_default_agent_runner
 from agents.tracing import get_trace_provider
 from agents.tracing.provider import DefaultTraceProvider
@@ -26,6 +28,7 @@
 import temporalio.worker
 from temporalio.client import ClientConfig
 from temporalio.contrib.openai_agents._invoke_model_activity import ModelActivity
+from temporalio.contrib.openai_agents._mcp import TemporalMCPServer
 from temporalio.contrib.openai_agents._model_parameters import ModelActivityParameters
 from temporalio.contrib.openai_agents._openai_runner import TemporalOpenAIRunner
 from temporalio.contrib.openai_agents._temporal_trace_provider import (
@@ -42,6 +45,7 @@
     DataConverter,
 )
 from temporalio.worker import Worker, WorkerConfig
+from temporalio.worker.workflow_sandbox import SandboxedWorkflowRunner
 
 
 @contextmanager
@@ -203,6 +207,7 @@ def __init__(
         self,
         model_params: Optional[ModelActivityParameters] = None,
         model_provider: Optional[ModelProvider] = None,
+        mcp_servers: Sequence[MCPServer] = (),
     ) -> None:
         """Initialize the OpenAI agents plugin.
 
@@ -231,6 +236,13 @@ def __init__(
         self._model_params = model_params
         self._model_provider = model_provider
 
+        self._mcp_servers = [
+            server
+            if isinstance(server, TemporalMCPServer)
+            else TemporalMCPServer(server)
+            for server in mcp_servers
+        ]
+
     def configure_client(self, config: ClientConfig) -> ClientConfig:
         """Configure the Temporal client for OpenAI agents integration.
 
@@ -265,9 +277,18 @@ def configure_worker(self, config: WorkerConfig) -> WorkerConfig:
         config["interceptors"] = list(config.get("interceptors") or []) + [
             OpenAIAgentsTracingInterceptor()
         ]
-        config["activities"] = list(config.get("activities") or []) + [
-            ModelActivity(self._model_provider).invoke_model_activity
-        ]
+        new_activities = [ModelActivity(self._model_provider).invoke_model_activity]
+        for mcp_server in self._mcp_servers:
+            new_activities.extend(mcp_server.get_activities())
+        config["activities"] = list(config.get("activities") or []) + new_activities
+
+        runner = config.get("workflow_runner")
+        if isinstance(runner, SandboxedWorkflowRunner):
+            config["workflow_runner"] = dataclasses.replace(
+                runner,
+                restrictions=runner.restrictions.with_passthrough_modules("mcp"),
+            )
+
         return super().configure_worker(config)
 
     async def run_worker(self, worker: Worker) -> None:
@@ -281,4 +302,7 @@ async def run_worker(self, worker: Worker) -> None:
             worker: The worker instance to run.
         """
         with set_open_ai_agent_temporal_overrides(self._model_params):
-            await super().run_worker(worker)
+            async with AsyncExitStack() as stack:
+                for mcp_server in self._mcp_servers:
+                    await stack.enter_async_context(mcp_server)
+                await super().run_worker(worker)