fix(pydantic): use current_context() in RestateModelWrapper.request_stream (AttributeError on streaming / event_stream_handler)

[selimacerbas]

Summary

RestateModelWrapper.request_stream references self._context, which is never assigned anywhere in restate.ext.pydantic. The sibling request() method — and every other method in _model.py / _toolset.py / _agent.py — uses the current_context() extension point added in #148. request_stream was missed in that migration, so any agent run that uses event_stream_handler (the only way restate.ext.pydantic surfaces per-turn model events) fails the moment the streaming model path executes.

self._context → WrapperModel.__getattr__ → delegates to the wrapped model → AttributeError.

Reproduction

Build a RestateAgent(agent, event_stream_handler=handler) and agent.run(...) it inside a Restate handler:

File ".../restate/ext/pydantic/_model.py", line 116, in request_stream
    response = await self._context.run_typed("Model stream call", request_stream_run, self._options)
File ".../pydantic_ai/models/wrapper.py", line 161, in __getattr__
    return getattr(self.wrapped, item)
AttributeError: 'OpenAIChatModel' object has no attribute '_context'

(event_stream_handler is required for streaming because pydantic-ai routes through the streaming model path — _agent_graph.py _streaming_handler → req_ctx.model.request_stream(...) — only when a handler is set.)

Fix

Use current_context() in request_stream, mirroring request() (including its None guard). Both current_context and UserError are already imported in this module.

        context = current_context()
        if context is None:
            raise UserError(
                "A model cannot be used without a Restate context. Make sure to run it within an agent or a run context."
            )
        try:
            response = await context.run_typed("Model stream call", request_stream_run, self._options)
            ...

Test

Not included here — happy to add one in your preferred harness. A targeted test would run a RestateAgent with an event_stream_handler (e.g. a TestModel / FunctionModel) under a Restate context and assert the handler receives the stream and the run completes; it raises AttributeError before this change. Point me at the right module and I'll add it.

Notes

• One-line behavioral change (plus the None guard for parity with request()); no public API change.
• Discovered while building a durable agent runtime on restate-sdk[pydantic]==0.18.1; the bug is also present on main.

[github-actions]

All contributors have signed the CLA ✍️ ✅
_{Posted by the CLA Assistant Lite bot.}

[selimacerbas]

I have read the CLA Document and I hereby sign the CLA

[selimacerbas]

recheck

[igalshilman]

Thanks, LGTM!