Support named toolset pre-registration

qyinm · qyinm · commit db55bb0c2e2b · 2025-12-20T21:56:30.000+09:00
diff --git a/docs/durable_execution/temporal.md b/docs/durable_execution/temporal.md
@@ -276,7 +276,11 @@ To use toolsets at runtime with `TemporalAgent`, you need to:
 
 This pattern allows multiple agents to share the same set of tools without duplicating activity registrations, enabling dynamic agent creation while maintaining proper tool registration.
 
-Here's an example showing how to register and use shared toolsets across multiple agents:
+Alternatively, you can pre-register toolsets with the `TemporalAgent` constructor and reference them by name at runtime. This is similar to how models are handled.
+
+### Using Toolset Instances
+
+Here's an example showing how to register and use shared toolsets across multiple agents using toolset instances:
 
 ```python {title="shared_toolset_temporal.py" test="skip"}
 from datetime import timedelta
@@ -381,6 +385,54 @@ async def main():
 5. Pass the wrapped toolset at runtime to any agent that needs it.
 6. Register the shared toolset's activities once with the worker. Agent activities are automatically registered via `__pydantic_ai_agents__`.
 
+6. Register the shared toolset's activities once with the worker. Agent activities are automatically registered via `__pydantic_ai_agents__`.
+
+### Using Named Toolsets
+
+You can also pre-register toolsets with names and reference them by name at runtime:
+
+```python {title="named_toolset_temporal.py" test="skip"}
+from temporalio import workflow
+from pydantic_ai import Agent, FunctionToolset
+from pydantic_ai.durable_exec.temporal import TemporalAgent, TemporalFunctionToolset
+
+# Define tools and toolset
+def magic_trick(input: str) -> str:
+    return f"Magic: {input}"
+
+magic_toolset = FunctionToolset(tools=[magic_trick], id='magic')
+
+# Wrap toolset
+wrapped_magic_toolset = TemporalFunctionToolset(
+    magic_toolset,
+    activity_name_prefix='magic',
+    deps_type=type(None),
+)
+
+# Create agent with pre-registered toolset
+agent = Agent('openai:gpt-5', name='magic_agent')
+temporal_agent = TemporalAgent(
+    agent,
+    toolsets={'magic_tools': wrapped_magic_toolset},  # (1)!
+)
+
+@workflow.defn
+class MagicWorkflow:
+    __pydantic_ai_agents__ = [temporal_agent]
+
+    @workflow.run
+    async def run(self, input: str) -> str:
+        # Reference toolset by name
+        result = await temporal_agent.run(
+            input,
+            toolsets=['magic_tools'],  # (2)!
+        )
+        return result.output
+```
+
+1. Pass a dictionary of toolsets to `TemporalAgent` to pre-register them. The keys are the names used to reference the toolsets at runtime.
+2. Pass the listing of toolset names to `run(toolsets=[...])` to use the pre-registered toolsets.
+
 You can also wrap toolsets at agent creation time by passing them to the wrapped agent's constructor, which will automatically temporalize them. The runtime pattern shown above is useful when you want to share toolsets across multiple agents or select toolsets dynamically based on workflow parameters.
 
 ## Activity Configuration
diff --git a/pydantic_ai_slim/pydantic_ai/durable_exec/temporal/_agent.py b/pydantic_ai_slim/pydantic_ai/durable_exec/temporal/_agent.py
@@ -117,6 +117,7 @@ def __init__(
         wrapped: AbstractAgent[AgentDepsT, OutputDataT],
         *,
         name: str | None = None,
+        toolsets: Sequence[AbstractToolset[AgentDepsT]] | Mapping[str, AbstractToolset[AgentDepsT]] | None = None,
         models: Mapping[str, Model] | None = None,
         provider_factory: TemporalProviderFactory | None = None,
         event_stream_handler: EventStreamHandler[AgentDepsT] | None = None,
@@ -144,6 +145,11 @@ def __init__(
         Args:
             wrapped: The agent to wrap.
             name: Optional unique agent name to use in the Temporal activities' names. If not provided, the agent's `name` will be used.
+            toolsets:
+                Optional additional toolsets to register with the agent, or a mapping of toolset names to toolset instances.
+                Toolsets passed here will be temporalized and their activities registered alongside the wrapped agent's existing toolsets.
+                If a mapping is provided, toolsets can be referenced by name in `run(toolsets=['name'])`.
+
             models:
                 Optional mapping of model instances to register with the agent.
                 Keys define the names that can be referenced at runtime and the values are `Model` instances.
@@ -191,6 +197,7 @@ def __init__(
         ]
         activity_config['retry_policy'] = retry_policy
         self.activity_config = activity_config
+        self._named_toolsets: Mapping[str, AbstractToolset[AgentDepsT]] | None = None
 
         model_activity_config = model_activity_config or {}
         toolset_activity_config = toolset_activity_config or {}
@@ -235,6 +242,18 @@ async def streamed_response():
         activities.extend(temporal_model.temporal_activities)
         self._temporal_model = temporal_model
 
+        if toolsets:
+            if isinstance(toolsets, Mapping):
+                # Flatten the mapping for temporalization, but keep track of names
+                additional_toolsets = list(toolsets.values())
+                self._named_toolsets = toolsets
+            else:
+                additional_toolsets = list(toolsets)
+                self._named_toolsets = {}
+        else:
+            additional_toolsets = []
+            self._named_toolsets = {}
+
         def temporalize_toolset(toolset: AbstractToolset[AgentDepsT]) -> AbstractToolset[AgentDepsT]:
             id = toolset.id
             if id is None:
@@ -254,9 +273,35 @@ def temporalize_toolset(toolset: AbstractToolset[AgentDepsT]) -> AbstractToolset
                 activities.extend(toolset.temporal_activities)
             return toolset
 
-        temporal_toolsets = [toolset.visit_and_replace(temporalize_toolset) for toolset in wrapped.toolsets]
+        all_toolsets = [*wrapped.toolsets, *additional_toolsets]
+        temporal_toolsets = [toolset.visit_and_replace(temporalize_toolset) for toolset in all_toolsets]
+
+        # If we had named toolsets, we need to map the names to the temporalized versions
+        # We know that visit_and_replace returns a new instance (or the same one if no replacement needed)
+        # matching the structure of the input.
+        # However, since we flattened everything into `all_toolsets` and then mapped it to `temporal_toolsets`,
+        # we can reconstruct the named mapping by index.
+        # But wait, `wrapped.toolsets` are first.
+        if self._named_toolsets:
+            # The additional toolsets are at the end of `temporal_toolsets`
+            num_wrapped_toolsets = len(wrapped.toolsets)
+            # The temporalized additional toolsets
+            temporal_additional_toolsets = temporal_toolsets[num_wrapped_toolsets:]
+
+            # create a new mapping pointing to the temporalized versions
+            new_named_toolsets: dict[str, AbstractToolset[AgentDepsT]] = {}
+            for name, temporal_toolset in zip(self._named_toolsets, temporal_additional_toolsets):
+                new_named_toolsets[name] = temporal_toolset
+            self._named_toolsets = new_named_toolsets
+
+            # If toolsets were passed as a mapping, they are not added to the active toolsets by default
+            if isinstance(toolsets, Mapping):
+                self._toolsets = temporal_toolsets[:num_wrapped_toolsets]
+            else:
+                self._toolsets = temporal_toolsets
+        else:
+            self._toolsets = temporal_toolsets
 
-        self._toolsets = temporal_toolsets
         self._temporal_activities = activities
 
         self._temporal_overrides_active: ContextVar[bool] = ContextVar('_temporal_overrides_active', default=False)
@@ -314,16 +359,21 @@ def temporal_activities(self) -> list[Callable[..., Any]]:
     @contextmanager
     def _temporal_overrides(
         self,
-        toolsets: Sequence[AbstractToolset[AgentDepsT]] | None = None,
+        toolsets: Sequence[AbstractToolset[AgentDepsT] | str] | None = None,
         model: models.Model | models.KnownModelName | str | None = None,
         force: bool = False,
     ) -> Iterator[None]:
         in_workflow = workflow.in_workflow()
 
         if toolsets:
-            if in_workflow:
-                _validate_temporal_toolsets(toolsets)
-            overridden_toolsets = [*self._toolsets, *toolsets]
+            if workflow.in_workflow():
+                # If toolsets are provided as strings, we can't validate them directly here as they are resolved later.
+                # We only validate if they are already AbstractToolset instances.
+                _validate_temporal_toolsets([t for t in toolsets if not isinstance(t, str)])
+
+            resolved_toolsets = self._resolve_toolsets(toolsets)
+            assert resolved_toolsets is not None
+            overridden_toolsets = [*self._toolsets, *resolved_toolsets]
         else:
             overridden_toolsets = list(self._toolsets)
 
@@ -352,6 +402,26 @@ def _temporal_overrides(
             finally:
                 self._temporal_overrides_active.reset(temporal_active_token)
 
+    def _resolve_toolsets(
+        self, toolsets: Sequence[AbstractToolset[AgentDepsT] | str] | None
+    ) -> Sequence[AbstractToolset[AgentDepsT]] | None:
+        if toolsets is None:
+            return None
+
+        resolved_toolsets: list[AbstractToolset[AgentDepsT]] = []
+        for t in toolsets:
+            if isinstance(t, str):
+                if self._named_toolsets is None:
+                    raise UserError(f"Unknown toolset name: '{t}'. No named toolsets registered.")
+                if t not in self._named_toolsets:
+                    raise UserError(
+                        f"Unknown toolset name: '{t}'. Available toolsets: {list(self._named_toolsets.keys())}"
+                    )
+                resolved_toolsets.append(self._named_toolsets[t])
+            else:
+                resolved_toolsets.append(t)
+        return resolved_toolsets
+
     @overload
     async def run(
         self,
@@ -367,7 +437,7 @@ async def run(
         usage_limits: _usage.UsageLimits | None = None,
         usage: _usage.RunUsage | None = None,
         infer_name: bool = True,
-        toolsets: Sequence[AbstractToolset[AgentDepsT]] | None = None,
+        toolsets: Sequence[AbstractToolset[AgentDepsT] | str] | None = None,
         builtin_tools: Sequence[AbstractBuiltinTool | BuiltinToolFunc[AgentDepsT]] | None = None,
         event_stream_handler: EventStreamHandler[AgentDepsT] | None = None,
     ) -> AgentRunResult[OutputDataT]: ...
@@ -387,7 +457,7 @@ async def run(
         usage_limits: _usage.UsageLimits | None = None,
         usage: _usage.RunUsage | None = None,
         infer_name: bool = True,
-        toolsets: Sequence[AbstractToolset[AgentDepsT]] | None = None,
+        toolsets: Sequence[AbstractToolset[AgentDepsT] | str] | None = None,
         builtin_tools: Sequence[AbstractBuiltinTool | BuiltinToolFunc[AgentDepsT]] | None = None,
         event_stream_handler: EventStreamHandler[AgentDepsT] | None = None,
     ) -> AgentRunResult[RunOutputDataT]: ...
@@ -406,7 +476,7 @@ async def run(
         usage_limits: _usage.UsageLimits | None = None,
         usage: _usage.RunUsage | None = None,
         infer_name: bool = True,
-        toolsets: Sequence[AbstractToolset[AgentDepsT]] | None = None,
+        toolsets: Sequence[AbstractToolset[AgentDepsT] | str] | None = None,
         builtin_tools: Sequence[AbstractBuiltinTool | BuiltinToolFunc[AgentDepsT]] | None = None,
         event_stream_handler: EventStreamHandler[AgentDepsT] | None = None,
         **_deprecated_kwargs: Never,
@@ -492,7 +562,7 @@ def run_sync(
         usage_limits: _usage.UsageLimits | None = None,
         usage: _usage.RunUsage | None = None,
         infer_name: bool = True,
-        toolsets: Sequence[AbstractToolset[AgentDepsT]] | None = None,
+        toolsets: Sequence[AbstractToolset[AgentDepsT] | str] | None = None,
         builtin_tools: Sequence[AbstractBuiltinTool | BuiltinToolFunc[AgentDepsT]] | None = None,
         event_stream_handler: EventStreamHandler[AgentDepsT] | None = None,
     ) -> AgentRunResult[OutputDataT]: ...
@@ -512,7 +582,7 @@ def run_sync(
         usage_limits: _usage.UsageLimits | None = None,
         usage: _usage.RunUsage | None = None,
         infer_name: bool = True,
-        toolsets: Sequence[AbstractToolset[AgentDepsT]] | None = None,
+        toolsets: Sequence[AbstractToolset[AgentDepsT] | str] | None = None,
         builtin_tools: Sequence[AbstractBuiltinTool | BuiltinToolFunc[AgentDepsT]] | None = None,
         event_stream_handler: EventStreamHandler[AgentDepsT] | None = None,
     ) -> AgentRunResult[RunOutputDataT]: ...
@@ -531,7 +601,7 @@ def run_sync(
         usage_limits: _usage.UsageLimits | None = None,
         usage: _usage.RunUsage | None = None,
         infer_name: bool = True,
-        toolsets: Sequence[AbstractToolset[AgentDepsT]] | None = None,
+        toolsets: Sequence[AbstractToolset[AgentDepsT] | str] | None = None,
         builtin_tools: Sequence[AbstractBuiltinTool | BuiltinToolFunc[AgentDepsT]] | None = None,
         event_stream_handler: EventStreamHandler[AgentDepsT] | None = None,
         **_deprecated_kwargs: Never,
@@ -589,7 +659,7 @@ def run_sync(
             usage_limits=usage_limits,
             usage=usage,
             infer_name=infer_name,
-            toolsets=toolsets,
+            toolsets=self._resolve_toolsets(toolsets),
             builtin_tools=builtin_tools,
             event_stream_handler=event_stream_handler,
             **_deprecated_kwargs,
@@ -610,7 +680,7 @@ def run_stream(
         usage_limits: _usage.UsageLimits | None = None,
         usage: _usage.RunUsage | None = None,
         infer_name: bool = True,
-        toolsets: Sequence[AbstractToolset[AgentDepsT]] | None = None,
+        toolsets: Sequence[AbstractToolset[AgentDepsT] | str] | None = None,
         builtin_tools: Sequence[AbstractBuiltinTool | BuiltinToolFunc[AgentDepsT]] | None = None,
         event_stream_handler: EventStreamHandler[AgentDepsT] | None = None,
     ) -> AbstractAsyncContextManager[StreamedRunResult[AgentDepsT, OutputDataT]]: ...
@@ -630,7 +700,7 @@ def run_stream(
         usage_limits: _usage.UsageLimits | None = None,
         usage: _usage.RunUsage | None = None,
         infer_name: bool = True,
-        toolsets: Sequence[AbstractToolset[AgentDepsT]] | None = None,
+        toolsets: Sequence[AbstractToolset[AgentDepsT] | str] | None = None,
         builtin_tools: Sequence[AbstractBuiltinTool | BuiltinToolFunc[AgentDepsT]] | None = None,
         event_stream_handler: EventStreamHandler[AgentDepsT] | None = None,
     ) -> AbstractAsyncContextManager[StreamedRunResult[AgentDepsT, RunOutputDataT]]: ...
@@ -650,7 +720,7 @@ async def run_stream(
         usage_limits: _usage.UsageLimits | None = None,
         usage: _usage.RunUsage | None = None,
         infer_name: bool = True,
-        toolsets: Sequence[AbstractToolset[AgentDepsT]] | None = None,
+        toolsets: Sequence[AbstractToolset[AgentDepsT] | str] | None = None,
         builtin_tools: Sequence[AbstractBuiltinTool | BuiltinToolFunc[AgentDepsT]] | None = None,
         event_stream_handler: EventStreamHandler[AgentDepsT] | None = None,
         **_deprecated_kwargs: Never,
@@ -707,7 +777,7 @@ async def main():
             usage_limits=usage_limits,
             usage=usage,
             infer_name=infer_name,
-            toolsets=toolsets,
+            toolsets=self._resolve_toolsets(toolsets),
             event_stream_handler=event_stream_handler,
             builtin_tools=builtin_tools,
             **_deprecated_kwargs,
@@ -729,8 +799,9 @@ def run_stream_events(
         usage_limits: _usage.UsageLimits | None = None,
         usage: _usage.RunUsage | None = None,
         infer_name: bool = True,
-        toolsets: Sequence[AbstractToolset[AgentDepsT]] | None = None,
+        toolsets: Sequence[AbstractToolset[AgentDepsT] | str] | None = None,
         builtin_tools: Sequence[AbstractBuiltinTool | BuiltinToolFunc[AgentDepsT]] | None = None,
+        event_stream_handler: EventStreamHandler[AgentDepsT] | None = None,
     ) -> AsyncIterator[_messages.AgentStreamEvent | AgentRunResultEvent[OutputDataT]]: ...
 
     @overload
@@ -748,8 +819,9 @@ def run_stream_events(
         usage_limits: _usage.UsageLimits | None = None,
         usage: _usage.RunUsage | None = None,
         infer_name: bool = True,
-        toolsets: Sequence[AbstractToolset[AgentDepsT]] | None = None,
+        toolsets: Sequence[AbstractToolset[AgentDepsT] | str] | None = None,
         builtin_tools: Sequence[AbstractBuiltinTool | BuiltinToolFunc[AgentDepsT]] | None = None,
+        event_stream_handler: EventStreamHandler[AgentDepsT] | None = None,
     ) -> AsyncIterator[_messages.AgentStreamEvent | AgentRunResultEvent[RunOutputDataT]]: ...
 
     def run_stream_events(
@@ -766,8 +838,9 @@ def run_stream_events(
         usage_limits: _usage.UsageLimits | None = None,
         usage: _usage.RunUsage | None = None,
         infer_name: bool = True,
-        toolsets: Sequence[AbstractToolset[AgentDepsT]] | None = None,
+        toolsets: Sequence[AbstractToolset[AgentDepsT] | str] | None = None,
         builtin_tools: Sequence[AbstractBuiltinTool | BuiltinToolFunc[AgentDepsT]] | None = None,
+        event_stream_handler: EventStreamHandler[AgentDepsT] | None = None,
     ) -> AsyncIterator[_messages.AgentStreamEvent | AgentRunResultEvent[Any]]:
         """Run the agent with a user prompt in async mode and stream events from the run.
 
@@ -818,6 +891,7 @@ async def main():
             infer_name: Whether to try to infer the agent name from the call frame if it's not set.
             toolsets: Optional additional toolsets for this run.
             builtin_tools: Optional additional builtin tools for this run.
+            event_stream_handler: Optional event stream handler to use for this run. It will receive all the events up until the final result is found, which you can then read or stream from inside the context manager.
 
         Returns:
             An async iterable of stream events `AgentStreamEvent` and finally a `AgentRunResultEvent` with the final
@@ -841,7 +915,7 @@ async def main():
             usage_limits=usage_limits,
             usage=usage,
             infer_name=infer_name,
-            toolsets=toolsets,
+            toolsets=self._resolve_toolsets(toolsets),
             builtin_tools=builtin_tools,
         )
 
@@ -979,6 +1053,7 @@ async def main():
             infer_name: Whether to try to infer the agent name from the call frame if it's not set.
             toolsets: Optional additional toolsets for this run.
             builtin_tools: Optional additional builtin tools for this run.
+            event_stream_handler: Optional event stream handler to use for this run.
 
         Returns:
             The result of the run.
diff --git a/tests/test_temporal.py b/tests/test_temporal.py
