pydantic ext: request_stream() never initializes the turnstile — KeyError('call_...') on first tool call of a streamed run

[thuantruong-kite]

Summary

In a streamed run (RestateAgent with an event_stream_handler), the first tool call fails with:

KeyError: 'call_XXXXXXXXXXXX'

raised from Turnstile.wait_for via RestateContextRunToolSet.call_tool.

Non-streamed runs (no event_stream_handler) work fine.

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

Tool executions are gated by a Turnstile built from the model response's tool-call ids so parallel tool calls run in deterministic journal order.

The non-streaming path arms it — restate/ext/pydantic/_model.py, RestateModelWrapper.request:

res = await context.run_typed("Model call", self.wrapped.request, self._options, *args, **kwargs)
ids = [c.tool_call_id for c in res.tool_calls]
current_state().turnstile = Turnstile(ids)      # <- line 82

The streaming path (RestateModelWrapper.request_stream) never does this, so the state keeps the default Turnstile([]) from State.__init__. Then in _toolset.py::RestateContextRunToolSet.call_tool:

await turnstile.wait_for(id)   # Turnstile.wait_for: self.events[id] -> KeyError

Reproduction

Same setup as the companion issue (agent with one tool, RestateAgent(agent, event_stream_handler=..., auto_wrap_tools=True)), with the "run event" coroutine bug patched or avoided. Any prompt that triggers a tool call raises:

File ".../restate/ext/pydantic/_toolset.py", line 92, in call_tool
    await turnstile.wait_for(id)
File ".../restate/ext/turnstile.py", line 25, in wait_for
    event = self.events[id]
KeyError: 'call_HtCkhcyMnGofOXHadHWA9AXS'

Suggested fix

Mirror the non-streaming path after the journaled stream step resolves (deterministic on replay since the response is recorded):

response = await context.run_typed("Model stream call", request_stream_run, self._options)
ids = [c.tool_call_id for c in response.tool_calls]
current_state().turnstile = Turnstile(ids)
yield RestateStreamedResponse(model_request_parameters, response)

We are running this exact patch as a local subclass workaround and streamed multi-tool runs (3 tool calls in one turn) complete correctly with proper Calling <tool> journal entries.

Related

• Companion issue: pydantic ext: auto_wrap_tools + async event_stream_handler crashes with 'Object of type coroutine is not JSON serializable' on first tool call #222 (the two bugs stack on the same code path).
• Fix race condition langchain turnstile #198 fixed a similar turnstile race in the langchain integration.

[prescod]

Another repro demo:

streaming_turnstile_repro.py