Add support tool approval and deferred tools

Author: igalshilman

pydantic_ai_slim/pydantic_ai/durable_exec/restate/_toolset.py

@@ -2,12 +2,13 @@
 
 from collections.abc import Callable
 from dataclasses import dataclass
-from typing import Any, Generic
+from typing import Any, Literal
 
-from restate import Context, RunOptions
+from restate import Context, RunOptions, TerminalError
 
 from pydantic_ai import ToolDefinition
 from pydantic_ai._run_context import AgentDepsT
+from pydantic_ai.exceptions import ApprovalRequired, CallDeferred, ModelRetry, UserError
 from pydantic_ai.mcp import MCPServer, ToolResult
 from pydantic_ai.tools import RunContext
 from pydantic_ai.toolsets.abstract import AbstractToolset, ToolsetTool
@@ -20,40 +21,71 @@
 class RestateContextRunResult:
     """A simple wrapper for tool results to be used with Restate's run_typed."""
 
+    kind: Literal['output', 'call_deferred', 'approval_required']
     output: Any
 
 
+CONTEXT_RUN_SERDE = PydanticTypeAdapter(RestateContextRunResult)
+
+
 @dataclass
 class RestateMCPGetToolsContextRunResult:
     """A simple wrapper for tool results to be used with Restate's run_typed."""
 
     output: dict[str, ToolDefinition]
 
 
+MCP_GET_TOOLS_SERDE = PydanticTypeAdapter(RestateMCPGetToolsContextRunResult)
+
+
 @dataclass
-class RestateMCPToolRunResult(Generic[AgentDepsT]):
+class RestateMCPToolRunResult:
     """A simple wrapper for tool results to be used with Restate's run_typed."""
 
     output: ToolResult
 
 
+MCP_RUN_SERDE = PydanticTypeAdapter(RestateMCPToolRunResult)
+
+
 class RestateContextRunToolSet(WrapperToolset[AgentDepsT]):
     """A toolset that automatically wraps tool calls with restate's `ctx.run_typed()`."""
 
     def __init__(self, wrapped: AbstractToolset[AgentDepsT], context: Context):
         super().__init__(wrapped)
         self._context = context
-        self.options = RunOptions[RestateContextRunResult](serde=PydanticTypeAdapter(RestateContextRunResult))
+        self.options = RunOptions[RestateContextRunResult](serde=CONTEXT_RUN_SERDE)
 
     async def call_tool(
         self, name: str, tool_args: dict[str, Any], ctx: RunContext[AgentDepsT], tool: ToolsetTool[AgentDepsT]
     ) -> Any:
         async def action() -> RestateContextRunResult:
-            output = await self.wrapped.call_tool(name, tool_args, ctx, tool)
-            return RestateContextRunResult(output=output)
+            try:
+                # A tool may raise ModelRetry, CallDeferred, ApprovalRequired, or UserError
+                # to signal special conditions to the caller.
+                # Since, restate ctx.run() will retry this exception we need to convert these exceptions
+                # to a return value and handle them outside of the ctx.run().
+                output = await self.wrapped.call_tool(name, tool_args, ctx, tool)
+                return RestateContextRunResult(kind='output', output=output)
+            except ModelRetry:
+                # we let restate to retry this
+                raise
+            except CallDeferred:
+                return RestateContextRunResult(kind='call_deferred', output=None)
+            except ApprovalRequired:
+                return RestateContextRunResult(kind='approval_required', output=None)
+            except UserError as e:
+                raise TerminalError(str(e)) from e
 
         res = await self._context.run_typed(f'Calling {name}', action, self.options)
-        return res.output
+
+        if res.kind == 'call_deferred':
+            raise CallDeferred()
+        elif res.kind == 'approval_required':
+            raise ApprovalRequired()
+        else:
+            assert res.kind == 'output'
+            return res.output
 
     def visit_and_replace(
         self, visitor: Callable[[AbstractToolset[AgentDepsT]], AbstractToolset[AgentDepsT]]
@@ -68,8 +100,6 @@ def __init__(self, wrapped: MCPServer, context: Context):
         super().__init__(wrapped)
         self._wrapped = wrapped
         self._context = context
-        self._mcp_tool_run_serde = PydanticTypeAdapter(RestateMCPToolRunResult[AgentDepsT])
-        self._mcp_get_tools_serde = PydanticTypeAdapter(RestateMCPGetToolsContextRunResult)
 
     def visit_and_replace(
         self, visitor: Callable[[AbstractToolset[AgentDepsT]], AbstractToolset[AgentDepsT]]
@@ -79,11 +109,12 @@ def visit_and_replace(
     async def get_tools(self, ctx: RunContext[AgentDepsT]) -> dict[str, ToolsetTool[AgentDepsT]]:
         async def get_tools_in_context() -> RestateMCPGetToolsContextRunResult:
             res = await self._wrapped.get_tools(ctx)
-            # ToolsetTool is not serializable as it holds a SchemaValidator (which is also the same for every MCP tool so unnecessary to pass along the wire every time),
+            # ToolsetTool is not serializable as it holds a SchemaValidator
+            # (which is also the same for every MCP tool so unnecessary to pass along the wire every time),
             # so we just return the ToolDefinitions and wrap them in ToolsetTool outside of the activity.
             return RestateMCPGetToolsContextRunResult(output={name: tool.tool_def for name, tool in res.items()})
 
-        options = RunOptions(serde=self._mcp_get_tools_serde)
+        options = RunOptions(serde=MCP_GET_TOOLS_SERDE)
 
         tool_defs = await self._context.run_typed('get mcp tools', get_tools_in_context, options)
 
@@ -100,11 +131,11 @@ async def call_tool(
         ctx: RunContext[AgentDepsT],
         tool: ToolsetTool[AgentDepsT],
     ) -> ToolResult:
-        async def call_tool_in_context() -> RestateMCPToolRunResult[AgentDepsT]:
+        async def call_tool_in_context() -> RestateMCPToolRunResult:
             res = await self._wrapped.call_tool(name, tool_args, ctx, tool)
             return RestateMCPToolRunResult(output=res)
 
-        options = RunOptions(serde=self._mcp_tool_run_serde)
+        options = RunOptions(serde=MCP_RUN_SERDE)
         res = await self._context.run_typed(f'Calling mcp tool {name}', call_tool_in_context, options)
 
         return res.output