fix(pydantic): make streamed runs work with auto_wrap_tools tool calls

python/restate/ext/pydantic/_agent.py

@@ -168,7 +168,13 @@ async def wrapped_event_stream_handler(
             async def single_event():
                 yield event
 
-            await context.run_typed("run event", lambda: fn(ctx, single_event()))
+            async def run_event() -> None:
+                # An async action so ctx.run() awaits the handler; a sync lambda
+                # returning the coroutine would never run it and the un-awaited
+                # coroutine fails JSON serialization of the journal entry.
+                await fn(ctx, single_event())
+
+            await context.run_typed("run event", run_event)
 
     @property
     def toolsets(self) -> Sequence[AbstractToolset[AgentDepsT]]:

python/restate/ext/pydantic/_model.py

@@ -120,6 +120,11 @@ async def request_stream_run():
             )
         try:
             response = await context.run_typed("Model stream call", request_stream_run, self._options)
+            # Arm the turnstile that orders journaled tool executions, exactly
+            # like the non-streaming request() does — without this the first
+            # tool call of a streamed run fails with KeyError in wait_for().
+            ids = [c.tool_call_id for c in response.tool_calls]
+            current_state().turnstile = Turnstile(ids)
             yield RestateStreamedResponse(model_request_parameters, response)
         except SdkInternalBaseException as e:
             raise Exception("Internal error during model stream call") from e