pydantic ext: auto_wrap_tools + async event_stream_handler crashes with 'Object of type coroutine is not JSON serializable' on first tool call #222

Summary

With RestateAgent(..., event_stream_handler=<async fn>, auto_wrap_tools=True), the first tool call of a run fails the invocation with:

TypeError: Object of type coroutine is not JSON serializable

Text-only runs work fine — the crash only happens once the model calls a tool, because that's when wrapped_event_stream_handler journals tool-side events.

Environment

• restate-sdk 1.0.3 (bug also present on current main)
• pydantic-ai-slim 2.4.0
• Python 3.12, Restate server 1.7.2

Root cause

restate/ext/pydantic/_agent.py (line 171 on main), in wrapped_event_stream_handler:

await context.run_typed("run event", lambda: fn(ctx, single_event()))

fn is the user's event_stream_handler, which is an async function. The lambda wrapping it is a plain sync callable, so in server_context.create_run_coroutine:

if inspect.iscoroutinefunction(action):   # False for the lambda
    action_result = await action()
else:
    ...run in executor...
    action_result = await action_result_future   # <- the un-awaited coroutine object

buffer = serde.serialize(action_result)   # TypeError: coroutine is not JSON serializable

The handler coroutine is never awaited (so the user's handler never runs for tool events), and Restate then tries to JSON-serialize the coroutine object itself.

Reproduction

import restate
from pydantic_ai import Agent
from restate.ext.pydantic import RestateAgent

agent = Agent("openai:gpt-4o")

@agent.tool_plain
def get_weather(city: str) -> str:
    return "sunny"

svc = restate.Service("Repro")

@svc.handler()
async def run(ctx: restate.Context, prompt: str) -> str:
    async def on_events(run_ctx, events):
        async for _ in events:
            pass

    ragent = RestateAgent(agent, event_stream_handler=on_events, auto_wrap_tools=True)
    result = await ragent.run(prompt)
    return result.output

app = restate.app([svc])

Invoke with a prompt that triggers the tool, e.g. "What's the weather in Hanoi? Use the tool." → the run event journal entry fails with the TypeError above and the invocation retries forever.

Traceback

File ".../restate/server_context.py", line 857, in create_run_coroutine
    buffer = serde.serialize(action_result)
File ".../restate/serde.py", line 317, in serialize
    return json.dumps(obj).encode("utf-8")
...
TypeError: Object of type coroutine is not JSON serializable

Suggested fix

Use an async closure so create_run_coroutine awaits it:

async for event in stream:

    async def single_event():
        yield event

    async def deliver() -> None:
        await fn(ctx, single_event())

    await context.run_typed("run event", deliver)

Related

• Found together with a second streaming+tools bug (turnstile never initialized in request_stream), reported separately.
• fix(pydantic): use current_context() in RestateModelWrapper.request_stream (AttributeError on streaming / event_stream_handler) #201 fixed an earlier crash on this same streaming path — this combination (streaming handler + auto_wrap_tools + an actual tool call) seems to not be covered by tests yet.