Fix pydantic.ext ToolReturn/deferred metadata, turnstile handling, and durable iter()

python/restate/ext/pydantic/_agent.py

@@ -27,7 +27,7 @@
 from pydantic_ai.models import Model
 from pydantic_ai.output import OutputDataT, OutputSpec
 from pydantic_ai.result import StreamedRunResult
-from pydantic_ai.run import AgentRunResult
+from pydantic_ai.run import AgentRun, AgentRunResult
 from pydantic_ai.tools import DeferredToolResults, RunContext
 from pydantic_ai.toolsets.abstract import AbstractToolset
 from pydantic_ai.toolsets.function import FunctionToolset
@@ -308,6 +308,48 @@ async def main():
                 spec=spec,
             )
 
+    @overload
+    def iter(
+        self,
+        user_prompt: str | Sequence[UserContent] | None = None,
+        *,
+        output_type: None = None,
+        model: models.Model | models.KnownModelName | str | None = None,
+        **kwargs: Any,
+    ) -> AbstractAsyncContextManager[AgentRun[AgentDepsT, OutputDataT]]: ...
+
+    @overload
+    def iter(
+        self,
+        user_prompt: str | Sequence[UserContent] | None = None,
+        *,
+        output_type: OutputSpec[RunOutputDataT],
+        model: models.Model | models.KnownModelName | str | None = None,
+        **kwargs: Any,
+    ) -> AbstractAsyncContextManager[AgentRun[AgentDepsT, RunOutputDataT]]: ...
+
+    @asynccontextmanager
+    async def iter(
+        self,
+        user_prompt: str | Sequence[UserContent] | None = None,
+        *,
+        output_type: OutputSpec[RunOutputDataT] | None = None,
+        model: models.Model | models.KnownModelName | str | None = None,
+        **kwargs: Any,
+    ) -> AsyncIterator[AgentRun[AgentDepsT, Any]]:
+        """Iterate over the agent graph's nodes with the same durable model and toolsets as `run()`.
+
+        The remaining keyword arguments are forwarded unchanged so that this override accepts every
+        argument `AbstractAgent.run()` passes to `iter()` across supported pydantic-ai versions.
+        """
+        if model is not None:
+            raise TerminalError(
+                "An agent needs to have a `model` in order to be used with Restate, it cannot be set at agent run time."
+            )
+        with self._restate_overrides():
+            async with super().iter(user_prompt, output_type=output_type, model=model, **kwargs) as agent_run:
+                yield agent_run
+
     @overload
     def run_stream(
         self,

python/restate/ext/pydantic/_model.py

@@ -120,6 +120,7 @@ async def request_stream_run():
             )
         try:
             response = await context.run_typed("Model stream call", request_stream_run, self._options)
+            current_state().turnstile = Turnstile([c.tool_call_id for c in response.tool_calls])
             yield RestateStreamedResponse(model_request_parameters, response)
         except SdkInternalBaseException as e:
             raise Exception("Internal error during model stream call") from e

python/restate/ext/pydantic/_toolset.py

@@ -11,6 +11,7 @@
 from pydantic_ai import ToolDefinition
 from pydantic_ai._run_context import AgentDepsT
 from pydantic_ai.exceptions import ApprovalRequired, CallDeferred, ModelRetry, UserError
+from pydantic_ai.messages import ToolReturn
 from pydantic_ai.mcp import MCPToolset
 from pydantic_ai.tools import RunContext
 from pydantic_ai.toolsets.abstract import AbstractToolset, ToolsetTool
@@ -27,6 +28,9 @@ class RestateContextRunResult:
     kind: Literal["output", "call_deferred", "approval_required", "model_retry"]
     output: Any
     error: str | None = None
+    # `output: Any` round-trips a ToolReturn as a plain dict, so it is journaled in its own typed field.
+    tool_return: ToolReturn | None = None
+    metadata: dict[str, Any] | None = None
 
 
 CONTEXT_RUN_SERDE = PydanticTypeAdapter(RestateContextRunResult)
@@ -69,13 +73,15 @@ async def action() -> RestateContextRunResult:
                 # 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)
+                if isinstance(output, ToolReturn):
+                    return RestateContextRunResult(kind="output", output=None, tool_return=output)
                 return RestateContextRunResult(kind="output", output=output, error=None)
             except ModelRetry as e:
                 return RestateContextRunResult(kind="model_retry", output=None, error=e.message)
-            except CallDeferred:
-                return RestateContextRunResult(kind="call_deferred", output=None, error=None)
-            except ApprovalRequired:
-                return RestateContextRunResult(kind="approval_required", output=None, error=None)
+            except CallDeferred as e:
+                return RestateContextRunResult(kind="call_deferred", output=None, metadata=e.metadata)
+            except ApprovalRequired as e:
+                return RestateContextRunResult(kind="approval_required", output=None, metadata=e.metadata)
             except UserError as e:
                 raise TerminalError(str(e)) from e
 
@@ -93,16 +99,25 @@ async def action() -> RestateContextRunResult:
             res = await context.run_typed(f"Calling {name}", action, self.options)
 
             if res.kind == "call_deferred":
-                raise CallDeferred()
+                # Deferring, approval and retry concern only this call; parallel siblings still
+                # need to run so the agent can continue with their results.
+                turnstile.allow_next_after(id)
+                raise CallDeferred(metadata=res.metadata)
             elif res.kind == "approval_required":
-                raise ApprovalRequired()
+                turnstile.allow_next_after(id)
+                raise ApprovalRequired(metadata=res.metadata)
             elif res.kind == "model_retry":
                 assert res.error is not None
+                turnstile.allow_next_after(id)
                 raise ModelRetry(res.error)
             else:
                 assert res.kind == "output"
                 turnstile.allow_next_after(id)
+                if res.tool_return is not None:
+                    return res.tool_return
                 return res.output
+        except (CallDeferred, ApprovalRequired, ModelRetry):
+            raise
         except Exception as e:
             turnstile.cancel_all_after(id)
             raise e from None

python/restate/ext/turnstile.py

@@ -22,12 +22,21 @@ def __init__(self, ids: list[str]):
             event.set()
 
     async def wait_for(self, id: str) -> None:
-        event = self.events[id]
+        event = self.events.get(id)
+        if event is None:
+            # Untracked ids (resumed deferred calls, nested calls with synthesized ids) are not part of
+            # this turn's ordering, so they proceed unless the turn was already canceled.
+            if self.canceled:
+                raise SdkInternalException() from None
+            return
         await event.wait()
         if self.canceled:
             raise SdkInternalException() from None
 
     def cancel_all_after(self, id: str) -> None:
+        if id not in self.events:
+            # A failing untracked call must not cancel the tracked calls of this turn.
+            return
         self.canceled = True
         next_id = self.turns.get(id)
         while next_id is not None: