Let Agent be run in a Temporal workflow by moving model requests, tool calls, and MCP to Temporal activities (#2225)

Author: DouweM

docs/api/agent.md

@@ -12,3 +12,4 @@
             - RunOutputDataT
             - capture_run_messages
             - InstrumentationSettings
+            - EventStreamHandler

docs/api/durable_exec.md

@@ -0,0 +1,3 @@
+# `pydantic_ai.durable_exec`
+
+::: pydantic_ai.durable_exec.temporal

docs/changelog.md

@@ -12,7 +12,7 @@ Pydantic AI is still pre-version 1, so breaking changes will occur, however:
 !!! note
     Here's a filtered list of the breaking changes for each version to help you upgrade Pydantic AI.
 
-### v0.7.0 (2025-08-08)
+### v0.7.0 (2025-08-12)
 
 See [#2458](https://github.com/pydantic/pydantic-ai/pull/2458) - `pydantic_ai.models.StreamedResponse` now yields a `FinalResultEvent` along with the existing `PartStartEvent` and `PartDeltaEvent`. If you're using `pydantic_ai.direct.model_request_stream` or `pydantic_ai.direct.model_request_stream_sync`, you may need to update your code to account for this.
 

docs/temporal.md

@@ -44,6 +44,7 @@ nav:
       - builtin-tools.md
       - common-tools.md
       - retries.md
+      - temporal.md
       - MCP:
           - mcp/index.md
           - mcp/client.md
@@ -75,6 +76,7 @@ nav:
       - api/toolsets.md
       - api/builtin_tools.md
       - api/common_tools.md
+      - api/durable_exec.md
       - api/output.md
       - api/result.md
       - api/messages.md

mkdocs.yml

@@ -82,6 +82,7 @@
     'InstrumentationSettings',
     'WrapperAgent',
     'AbstractAgent',
+    'EventStreamHandler',
 )
 
 
@@ -401,6 +402,11 @@ def name(self, value: str | None) -> None:
         """Set the name of the agent, used for logging."""
         self._name = value
 
+    @property
+    def deps_type(self) -> type:
+        """The type of dependencies used by the agent."""
+        return self._deps_type
+
     @property
     def output_type(self) -> OutputSpec[OutputDataT]:
         """The type of data output by agent runs, used to validate the data returned by the model, defaults to `str`."""
@@ -593,12 +599,7 @@ async def main():
             run_step=state.run_step,
         )
 
-        toolset = self._get_toolset(additional=toolsets)
-
-        if output_toolset is not None:
-            if self._prepare_output_tools:
-                output_toolset = PreparedToolset(output_toolset, self._prepare_output_tools)
-            toolset = CombinedToolset([output_toolset, toolset])
+        toolset = self._get_toolset(output_toolset=output_toolset, additional_toolsets=toolsets)
 
         async with toolset:
             # This will raise errors for any name conflicts
@@ -1240,48 +1241,64 @@ def _get_deps(self: Agent[T, OutputDataT], deps: T) -> T:
             return deps
 
     def _get_toolset(
-        self, additional: Sequence[AbstractToolset[AgentDepsT]] | None = None
+        self,
+        output_toolset: AbstractToolset[AgentDepsT] | None | _utils.Unset = _utils.UNSET,
+        additional_toolsets: Sequence[AbstractToolset[AgentDepsT]] | None = None,
     ) -> AbstractToolset[AgentDepsT]:
-        """Get the combined toolset containing function tools registered directly to the agent and user-provided toolsets including MCP servers.
+        """Get the complete toolset.
 
         Args:
-            additional: Additional toolsets to add.
+            output_toolset: The output toolset to use instead of the one built at agent construction time.
+            additional_toolsets: Additional toolsets to add, unless toolsets have been overridden.
         """
-        if some_tools := self._override_tools.get():
-            function_toolset = _AgentFunctionToolset(some_tools.value, max_retries=self._max_tool_retries)
-        else:
-            function_toolset = self._function_toolset
+        toolsets = self.toolsets
+        # Don't add additional toolsets if the toolsets have been overridden
+        if additional_toolsets and self._override_toolsets.get() is None:
+            toolsets = [*toolsets, *additional_toolsets]
 
-        if some_user_toolsets := self._override_toolsets.get():
-            user_toolsets = some_user_toolsets.value
-        else:
-            # Copy the dynamic toolsets to ensure each run has its own instances
-            dynamic_toolsets = [dataclasses.replace(toolset) for toolset in self._dynamic_toolsets]
-            user_toolsets = [*self._user_toolsets, *dynamic_toolsets, *(additional or [])]
+        toolset = CombinedToolset(toolsets)
 
-        if user_toolsets:
-            toolset = CombinedToolset([function_toolset, *user_toolsets])
-        else:
-            toolset = function_toolset
+        # Copy the dynamic toolsets to ensure each run has its own instances
+        def copy_dynamic_toolsets(toolset: AbstractToolset[AgentDepsT]) -> AbstractToolset[AgentDepsT]:
+            if isinstance(toolset, DynamicToolset):
+                return dataclasses.replace(toolset)
+            else:
+                return toolset
+
+        toolset = toolset.visit_and_replace(copy_dynamic_toolsets)
 
         if self._prepare_tools:
             toolset = PreparedToolset(toolset, self._prepare_tools)
 
+        output_toolset = output_toolset if _utils.is_set(output_toolset) else self._output_toolset
+        if output_toolset is not None:
+            if self._prepare_output_tools:
+                output_toolset = PreparedToolset(output_toolset, self._prepare_output_tools)
+            toolset = CombinedToolset([output_toolset, toolset])
+
         return toolset
 
     @property
     def toolsets(self) -> Sequence[AbstractToolset[AgentDepsT]]:
         """All toolsets registered on the agent, including a function toolset holding tools that were registered on the agent directly.
 
-        If a `prepare_tools` function was configured on the agent, this will contain just a `PreparedToolset` wrapping the original toolsets.
-
         Output tools are not included.
         """
-        toolset = self._get_toolset()
-        if isinstance(toolset, CombinedToolset):
-            return toolset.toolsets
+        toolsets: list[AbstractToolset[AgentDepsT]] = []
+
+        if some_tools := self._override_tools.get():
+            function_toolset = _AgentFunctionToolset(some_tools.value, max_retries=self._max_tool_retries)
+        else:
+            function_toolset = self._function_toolset
+        toolsets.append(function_toolset)
+
+        if some_user_toolsets := self._override_toolsets.get():
+            user_toolsets = some_user_toolsets.value
         else:
-            return [toolset]
+            user_toolsets = [*self._user_toolsets, *self._dynamic_toolsets]
+        toolsets.extend(user_toolsets)
+
+        return toolsets
 
     def _prepare_output_schema(
         self, output_type: OutputSpec[RunOutputDataT] | None, model_profile: ModelProfile
@@ -1369,7 +1386,7 @@ async def run_mcp_servers(
 class _AgentFunctionToolset(FunctionToolset[AgentDepsT]):
     @property
     def id(self) -> str:
-        return '<agent>'  # pragma: no cover
+        return '<agent>'
 
     @property
     def label(self) -> str:

pydantic_ai_slim/pydantic_ai/agent/__init__.py

@@ -93,6 +93,12 @@ def name(self, value: str | None) -> None:
         """Set the name of the agent, used for logging."""
         raise NotImplementedError
 
+    @property
+    @abstractmethod
+    def deps_type(self) -> type:
+        """The type of dependencies used by the agent."""
+        raise NotImplementedError
+
     @property
     @abstractmethod
     def output_type(self) -> OutputSpec[OutputDataT]:

pydantic_ai_slim/pydantic_ai/agent/abstract.py

@@ -43,6 +43,10 @@ def name(self) -> str | None:
     def name(self, value: str | None) -> None:
         self.wrapped.name = value
 
+    @property
+    def deps_type(self) -> type:
+        return self.wrapped.deps_type
+
     @property
     def output_type(self) -> OutputSpec[OutputDataT]:
         return self.wrapped.output_type
@@ -196,8 +200,8 @@ async def main():
             usage=usage,
             infer_name=infer_name,
             toolsets=toolsets,
-        ) as result:
-            yield result
+        ) as run:
+            yield run
 
     @contextmanager
     def override(

pydantic_ai_slim/pydantic_ai/agent/wrapper.py

@@ -0,0 +1,83 @@
+from __future__ import annotations
+
+import warnings
+from collections.abc import Sequence
+from dataclasses import replace
+from typing import Any, Callable
+
+from pydantic.errors import PydanticUserError
+from temporalio.client import ClientConfig, Plugin as ClientPlugin
+from temporalio.contrib.pydantic import PydanticPayloadConverter, pydantic_data_converter
+from temporalio.converter import DefaultPayloadConverter
+from temporalio.worker import Plugin as WorkerPlugin, WorkerConfig
+from temporalio.worker.workflow_sandbox import SandboxedWorkflowRunner
+
+from ...exceptions import UserError
+from ._agent import TemporalAgent
+from ._logfire import LogfirePlugin
+from ._run_context import TemporalRunContext
+from ._toolset import TemporalWrapperToolset
+
+__all__ = [
+    'TemporalAgent',
+    'PydanticAIPlugin',
+    'LogfirePlugin',
+    'AgentPlugin',
+    'TemporalRunContext',
+    'TemporalWrapperToolset',
+]
+
+
+class PydanticAIPlugin(ClientPlugin, WorkerPlugin):
+    """Temporal client and worker plugin for Pydantic AI."""
+
+    def configure_client(self, config: ClientConfig) -> ClientConfig:
+        if (data_converter := config.get('data_converter')) and data_converter.payload_converter_class not in (
+            DefaultPayloadConverter,
+            PydanticPayloadConverter,
+        ):
+            warnings.warn(  # pragma: no cover
+                'A non-default Temporal data converter was used which has been replaced with the Pydantic data converter.'
+            )
+
+        config['data_converter'] = pydantic_data_converter
+        return super().configure_client(config)
+
+    def configure_worker(self, config: WorkerConfig) -> WorkerConfig:
+        runner = config.get('workflow_runner')  # pyright: ignore[reportUnknownMemberType]
+        if isinstance(runner, SandboxedWorkflowRunner):  # pragma: no branch
+            config['workflow_runner'] = replace(
+                runner,
+                restrictions=runner.restrictions.with_passthrough_modules(
+                    'pydantic_ai',
+                    'logfire',
+                    'rich',
+                    'httpx',
+                    # Imported inside `logfire._internal.json_encoder` when running `logfire.info` inside an activity with attributes to serialize
+                    'attrs',
+                    # Imported inside `logfire._internal.json_schema` when running `logfire.info` inside an activity with attributes to serialize
+                    'numpy',
+                    'pandas',
+                ),
+            )
+
+        config['workflow_failure_exception_types'] = [
+            *config.get('workflow_failure_exception_types', []),  # pyright: ignore[reportUnknownMemberType]
+            UserError,
+            PydanticUserError,
+        ]
+
+        return super().configure_worker(config)
+
+
+class AgentPlugin(WorkerPlugin):
+    """Temporal worker plugin for a specific Pydantic AI agent."""
+
+    def __init__(self, agent: TemporalAgent[Any, Any]):
+        self.agent = agent
+
+    def configure_worker(self, config: WorkerConfig) -> WorkerConfig:
+        activities: Sequence[Callable[..., Any]] = config.get('activities', [])  # pyright: ignore[reportUnknownMemberType]
+        # Activities are checked for name conflicts by Temporal.
+        config['activities'] = [*activities, *self.agent.temporal_activities]
+        return super().configure_worker(config)