From 14c9758cc7a4ff2d4dc4c593822f8ceab0ef70ac Mon Sep 17 00:00:00 2001 From: Quoc Thuan Truong Date: Mon, 27 Jul 2026 15:10:25 +0700 Subject: [PATCH] fix(pydantic): make streamed runs work with auto_wrap_tools tool calls MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Two bugs that stack on the first tool call of a streamed run (RestateAgent with an event_stream_handler): 1. wrapped_event_stream_handler journaled each tool-side event via run_typed("run event", lambda: fn(ctx, single_event())). The sync lambda returns the async handler's coroutine un-awaited, so the handler never runs and the journal entry fails with "TypeError: Object of type coroutine is not JSON serializable". Use an async action so ctx.run() awaits it. 2. request_stream() never armed the turnstile that orders journaled tool executions — only the non-streaming request() did — so the first tool call of a streamed run raised KeyError('call_...') in Turnstile.wait_for. Arm it from the recorded response, mirroring request(). Fixes #222 Fixes #223 --- python/restate/ext/pydantic/_agent.py | 8 +++++++- python/restate/ext/pydantic/_model.py | 5 +++++ 2 files changed, 12 insertions(+), 1 deletion(-) diff --git a/python/restate/ext/pydantic/_agent.py b/python/restate/ext/pydantic/_agent.py index cc7888c..754bd4e 100644 --- a/python/restate/ext/pydantic/_agent.py +++ b/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]]: diff --git a/python/restate/ext/pydantic/_model.py b/python/restate/ext/pydantic/_model.py index fd162da..b61ff5c 100644 --- a/python/restate/ext/pydantic/_model.py +++ b/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