refactor(rust-bridge): generator-style execution protocol with stream step names (#46282)

* refactor(rust-bridge): drive executions with generator-style send, throw and close

Co-Authored-By: Yujong Lee <yujong@berri.ai>

* refactor(rust-bridge): name stream steps StreamOpen and StreamChunk and reject out-of-phase steps

Co-Authored-By: Yujong Lee <yujong@berri.ai>

* refactor(rust-bridge): drive with ExitStack and require the stream factory

Co-Authored-By: Yujong Lee <yujong@berri.ai>

* test(host-python): give the hook_chain lifecycle binding the stream factory

Co-Authored-By: Yujong Lee <yujong@berri.ai>

---------

Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

litellm-rust/crates/host-python/src/caller.rs

@@ -32,12 +32,12 @@ impl PythonCaller {
 
     /// Hands the caller a stream head and waits for demand or closure.
     pub(crate) async fn open_stream(&self, head: Py<PyAny>) -> ControlFlow<()> {
-        self.stream_boundary(Suspension::Open(head)).await
+        self.stream_boundary(Suspension::StreamOpen(head)).await
     }
 
     /// Delivers a chunk and waits for the caller to request another or close.
     pub(crate) async fn yield_chunk(&self, chunk: Py<PyAny>) -> ControlFlow<()> {
-        self.stream_boundary(Suspension::Yield(chunk)).await
+        self.stream_boundary(Suspension::StreamChunk(chunk)).await
     }
 
     /// Maps stream demand to continuation and any error to detachment.

litellm-rust/crates/host-python/src/driver.rs

@@ -645,12 +645,12 @@ where
     }
     match driver.resume(None)? {
         ExecutionStep::Return(value) => Ok(value),
-        ExecutionStep::Suspended(Suspension::Open(head)) => {
+        ExecutionStep::Suspended(Suspension::StreamOpen(head)) => {
             Execution::suspended(driver, options.lifecycle)
                 .into_sync_stream(py, head)
                 .map(Bound::unbind)
         }
-        ExecutionStep::Suspended(Suspension::Await(_) | Suspension::Yield(_)) => {
+        ExecutionStep::Suspended(Suspension::Await(_) | Suspension::StreamChunk(_)) => {
             Err(Error::SyncCallSuspended.into())
         }
     }

litellm-rust/crates/host-python/src/error.rs

@@ -12,6 +12,8 @@ pub enum Error {
     ExecutionRunning,
     #[error("execution is closed")]
     ExecutionClosed,
+    #[error("can't send non-None value to a just-started execution")]
+    SendBeforeStart,
     #[error("execution requires start before resume and can only start once")]
     ExecutionOutOfOrder,
     #[error("execution was closed while running")]
@@ -31,7 +33,9 @@ pub enum Error {
 impl From<Error> for PyErr {
     fn from(error: Error) -> Self {
         match error {
-            Error::FileContentType { .. } => PyTypeError::new_err(error.to_string()),
+            Error::SendBeforeStart | Error::FileContentType { .. } => {
+                PyTypeError::new_err(error.to_string())
+            }
             Error::MissingState
             | Error::ExecutionRunning
             | Error::ExecutionClosed

litellm-rust/crates/host-python/src/handle.rs

@@ -13,8 +13,13 @@ pub enum Suspension {
     Await(Py<PyAny>),
     /// The call streams: the caller gets a stream over this execution carrying this head,
     /// and the execution stays suspended until the stream asks for a chunk.
-    Open(Py<PyAny>),
-    Yield(Py<PyAny>),
+    StreamOpen(Py<PyAny>),
+    StreamChunk(Py<PyAny>),
+}
+
+enum Resume {
+    Send(Py<PyAny>),
+    Throw(PyErr),
 }
 
 pub enum ExecutionStep {
@@ -72,37 +77,45 @@ impl Execution {
         }
     }
 
-    fn advance(
-        slf: &Bound<'_, Self>,
-        py: Python<'_>,
-        result: Option<PyResult<Py<PyAny>>>,
-    ) -> PyResult<Py<PyAny>> {
-        let mut body = {
+    fn advance(slf: &Bound<'_, Self>, py: Python<'_>, input: Resume) -> PyResult<Py<PyAny>> {
+        let (mut body, started) = {
             let mut execution = slf.borrow_mut();
-            match (&execution.state, result.is_some()) {
-                (ExecutionState::Created(_), false) | (ExecutionState::Suspended(_), true) => {}
+            let started = match (&execution.state, &input) {
+                (ExecutionState::Created(_), Resume::Send(value)) if value.is_none(py) => false,
+                (ExecutionState::Created(_), Resume::Send(_)) => {
+                    return Err(Error::SendBeforeStart.into());
+                }
+                (ExecutionState::Suspended(_), _) => true,
                 (ExecutionState::Running, _) => {
                     return Err(Error::ExecutionRunning.into());
                 }
                 (ExecutionState::Closed, _) => {
                     return Err(Error::ExecutionClosed.into());
                 }
-                _ => {
+                (ExecutionState::Created(_), Resume::Throw(_)) => {
                     return Err(Error::ExecutionOutOfOrder.into());
                 }
-            }
+            };
             match std::mem::replace(&mut execution.state, ExecutionState::Running) {
-                ExecutionState::Created(body) | ExecutionState::Suspended(body) => body,
+                ExecutionState::Created(body) | ExecutionState::Suspended(body) => (body, started),
                 _ => unreachable!(),
             }
         };
+        let result = match input {
+            Resume::Send(value) => started.then_some(Ok(value)),
+            Resume::Throw(error) => Some(Err(error)),
+        };
         let lifecycle = slf.borrow().lifecycle;
         let outcome = catch_unwind(AssertUnwindSafe(|| {
             let step = body.resume(result)?;
             let (tag, value, suspended) = match step {
                 ExecutionStep::Suspended(Suspension::Await(value)) => ("Await", value, true),
-                ExecutionStep::Suspended(Suspension::Open(head)) => ("Open", head, true),
-                ExecutionStep::Suspended(Suspension::Yield(value)) => ("Yield", value, true),
+                ExecutionStep::Suspended(Suspension::StreamOpen(head)) => {
+                    ("StreamOpen", head, true)
+                }
+                ExecutionStep::Suspended(Suspension::StreamChunk(chunk)) => {
+                    ("StreamChunk", chunk, true)
+                }
                 ExecutionStep::Return(value) => ("Complete", value, false),
             };
             let step = lifecycle(py)?.getattr(tag)?.call1((value,))?.unbind();
@@ -132,24 +145,16 @@ impl Execution {
 
 #[pymethods]
 impl Execution {
-    fn start(slf: &Bound<'_, Self>, py: Python<'_>) -> PyResult<Py<PyAny>> {
-        Self::advance(slf, py, None)
+    fn send(slf: &Bound<'_, Self>, py: Python<'_>, value: Py<PyAny>) -> PyResult<Py<PyAny>> {
+        Self::advance(slf, py, Resume::Send(value))
     }
 
-    fn resume_value(
-        slf: &Bound<'_, Self>,
-        py: Python<'_>,
-        value: Py<PyAny>,
-    ) -> PyResult<Py<PyAny>> {
-        Self::advance(slf, py, Some(Ok(value)))
-    }
-
-    fn resume_error(
+    fn throw(
         slf: &Bound<'_, Self>,
         py: Python<'_>,
         error: Bound<'_, PyBaseException>,
     ) -> PyResult<Py<PyAny>> {
-        Self::advance(slf, py, Some(Err(PyErr::from_value(error.into_any()))))
+        Self::advance(slf, py, Resume::Throw(PyErr::from_value(error.into_any())))
     }
 
     fn close(slf: &Bound<'_, Self>) {

litellm-rust/crates/host-python/tests/hook_chain.rs

@@ -519,9 +519,11 @@ impl litellm_host_python::ExecutionBody for Preparing {
 }
 
 fn lifecycle(py: Python<'_>) -> PyResult<Bound<'_, PyModule>> {
-    let source =
-        std::ffi::CString::new(include_str!("../../../../litellm/rust_bridge/lifecycle.py"))
-            .unwrap();
+    let source = std::ffi::CString::new(concat!(
+        include_str!("../../../../litellm/rust_bridge/lifecycle.py"),
+        "\n_drive = drive\n\n\nasync def drive(execution):\n    return await _drive(execution, Stream)\n",
+    ))
+    .unwrap();
     PyModule::from_code(py, &source, c"lifecycle.py", c"hook_chain_test_lifecycle")
 }
 

litellm-rust/crates/host-python/tests/lifecycle.py

@@ -36,11 +36,17 @@ async def exercise():
     releaser = asyncio.create_task(release())
     execution = await_execution(CustomAwaitable())
     try:
-        execution.resume_value(None)
+        execution.throw(ValueError("early"))
     except RuntimeError:
         pass
     else:
         raise AssertionError("resumed an unstarted execution")
+    try:
+        execution.send(1)
+    except TypeError:
+        pass
+    else:
+        raise AssertionError("sent a value into an unstarted execution")
     wrapped = drive(execution)
     try:
         wrapped.send(1)
@@ -129,7 +135,7 @@ async def exercise():
 
     def reenter():
         try:
-            active.start()
+            active.send(None)
         except RuntimeError as error:
             observed.append(str(error))
         return result
@@ -151,7 +157,7 @@ async def exercise():
         execution = calling_execution(callback)
         callback.owner = execution
         if started:
-            assert execution.start().value is result
+            assert execution.send(None).value is result
         return weakref.ref(callback)
 
     for started in (False, True):

litellm-rust/crates/host-python/tests/support/mod.rs

@@ -50,8 +50,9 @@ pub fn call_options(asynchronous: bool) -> CallOptions {
 }
 
 pub fn install_lifecycle_module(py: Python<'_>) -> Bound<'_, PyModule> {
-    let source = std::ffi::CString::new(include_str!(
-        "../../../../../litellm/rust_bridge/lifecycle.py"
+    let source = std::ffi::CString::new(concat!(
+        include_str!("../../../../../litellm/rust_bridge/lifecycle.py"),
+        "\n_drive = drive\n\n\nasync def drive(execution):\n    return await _drive(execution, Stream)\n",
     ))
     .unwrap();
     PyModule::from_code(

litellm-rust/crates/python-bridge/AGENTS.md

@@ -49,7 +49,7 @@ Each step needs focused regression tests in the owning crate and Python integrat
 - Conversion errors and every failure after the call starts are terminal
   - Disabled/unavailable native execution may select legacy once. Callback exceptions never authorize fallback or replay
 - Use one ordinary inline `async def` driver in `litellm/rust_bridge/lifecycle.py`, with the native handle and call driver in `litellm-host-python`
-  - Contract: `start`, `resume_value`, `resume_error`, idempotent `close`. Explicitly tagged `Await`/`Complete` preserve awaitable final values
+  - Contract: generator-style `send` (`send(None)` starts), `throw`, idempotent `close`. Explicitly tagged `Await`/`Complete` preserve awaitable final values
   - Validate Created/Running/Suspended/Closed protocol states. The machine yields ops, the driver emits one terminal event, the adapter chooses dispatch policy
   - Defer effectful setup/context reads/timestamps until start. Unstarted-handle destruction releases inputs independently of Python `finally`
   - Catch only the selected await's errors. Start/resume errors propagate, `GeneratorExit` closes without further awaits

litellm/rust_bridge/AGENTS.md

@@ -4,7 +4,7 @@ This package owns native rollout and fallback selection, Python public API compa
 
 `host/` holds the Python modules Rust imports by name to read Python-owned state (provider resolution, model capabilities, settings, caches, diagnostics, legacy callbacks, the secret manager, transport and timeout conversions). Its AGENTS.md states that contract
 
-`lifecycle.py` owns generic inline execution and stream iteration. `streams.py` supplies the product binding and public stream wrappers through the bridge rather than let the generic Rust host import this package by name. Keep one driver implementing `start`, `resume_value`, `resume_error` and idempotent `close`. Do not create a second implementation
+`lifecycle.py` owns generic inline execution and stream iteration. `streams.py` supplies the product binding and public stream wrappers through the bridge rather than let the generic Rust host import this package by name. Keep one driver implementing generator-style `send` (`send(None)` starts), `throw` and idempotent `close`. Do not create a second implementation
 
 Generic execution steps and inline suspension handling must not depend on LiteLLM response metadata. Public stream construction, `_hidden_params` and header compatibility remain product responsibilities. Keep generic stream heads opaque and preserve public header behavior in the product wrappers
 

litellm/rust_bridge/lifecycle.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 from collections.abc import AsyncIterator, Awaitable, Callable, Iterator
+from contextlib import ExitStack
 from dataclasses import dataclass
 from typing import Final, Protocol
 
@@ -16,29 +17,31 @@ class Complete:
 
 
 @dataclass(frozen=True, slots=True)
-class Open:
-    value: object
+class StreamOpen:
+    head: object
 
 
 @dataclass(frozen=True, slots=True)
-class Yield:
-    value: object
+class StreamChunk:
+    chunk: object
 
 
-Settled = Complete | Open | Yield
+StreamStep = StreamOpen | StreamChunk
+Settled = Complete | StreamStep
 Step = Await | Settled
 
 
 class Execution(Protocol):
-    def start(self) -> Step: ...
+    def send(self, value: object) -> Step: ...
 
-    def resume_value(self, value: object) -> Step: ...
-
-    def resume_error(self, error: BaseException) -> Step: ...
+    def throw(self, error: BaseException) -> Step: ...
 
     def close(self) -> None: ...
 
 
+_REOPENED: Final = "execution opened a stream twice"
+
+
 class StreamClosed(Exception):
     """Tells a streaming execution that its caller stopped reading."""
 
@@ -50,9 +53,9 @@ async def _settle(execution: Execution, step: Step) -> Settled:
         except GeneratorExit:
             raise
         except BaseException as error:
-            step = execution.resume_error(error)  # rebind-ok: advance the execution protocol
+            step = execution.throw(error)  # rebind-ok: advance the execution protocol
         else:
-            step = execution.resume_value(value)  # rebind-ok: advance the execution protocol
+            step = execution.send(value)  # rebind-ok: advance the execution protocol
     return step
 
 
@@ -62,19 +65,18 @@ def _settled(step: Step) -> Settled:
     return step
 
 
-async def drive(execution: Execution, stream_factory: Callable[[Execution, object], object] | None = None) -> object:
-    handed_off = False  # rebind-ok: set once the execution belongs to the returned stream
-    try:
-        step: Final = await _settle(execution, execution.start())
-        if isinstance(step, Open):
-            factory: Final = Stream if stream_factory is None else stream_factory
-            stream: Final = factory(execution, step.value)
-            handed_off = True
-            return stream
-        return step.value
-    finally:
-        if not handed_off:
-            execution.close()
+async def drive(execution: Execution, open_stream: Callable[[Execution, object], object]) -> object:
+    with ExitStack() as cleanup:
+        cleanup.callback(execution.close)
+        match await _settle(execution, execution.send(None)):
+            case Complete(value):
+                return value
+            case StreamOpen(head):
+                stream: Final = open_stream(execution, head)
+                cleanup.pop_all()
+                return stream
+            case StreamChunk():
+                raise RuntimeError("execution yielded a chunk before opening a stream")
 
 
 class Stream(AsyncIterator[object]):
@@ -92,20 +94,25 @@ class Stream(AsyncIterator[object]):
         if self._done:
             raise StopAsyncIteration
         try:
-            step: Final = await _settle(self._execution, self._execution.resume_value(None))
+            step: Final = await _settle(self._execution, self._execution.send(None))
         except BaseException:
             self._finish()
             raise
-        if isinstance(step, Yield):
-            return step.value
-        self._finish()
-        raise StopAsyncIteration
+        match step:
+            case StreamChunk(chunk):
+                return chunk
+            case Complete():
+                self._finish()
+                raise StopAsyncIteration
+            case StreamOpen():
+                self._finish()
+                raise RuntimeError(_REOPENED)
 
     async def aclose(self) -> None:
         if self._done:
             return
         try:
-            await _settle(self._execution, self._execution.resume_error(StreamClosed()))
+            await _settle(self._execution, self._execution.throw(StreamClosed()))
         finally:
             self._finish()
 
@@ -129,20 +136,25 @@ class SyncStream(Iterator[object]):
         if self._done:
             raise StopIteration
         try:
-            step: Final = _settled(self._execution.resume_value(None))
+            step: Final = _settled(self._execution.send(None))
         except BaseException:
             self._finish()
             raise
-        if isinstance(step, Yield):
-            return step.value
-        self._finish()
-        raise StopIteration
+        match step:
+            case StreamChunk(chunk):
+                return chunk
+            case Complete():
+                self._finish()
+                raise StopIteration
+            case StreamOpen():
+                self._finish()
+                raise RuntimeError(_REOPENED)
 
     def close(self) -> None:
         if self._done:
             return
         try:
-            _settled(self._execution.resume_error(StreamClosed()))
+            _settled(self._execution.throw(StreamClosed()))
         finally:
             self._finish()
 

litellm/rust_bridge/streams.py

@@ -10,8 +10,8 @@ from litellm.rust_bridge.lifecycle import Execution
 
 Await: Final = lifecycle.Await
 Complete: Final = lifecycle.Complete
-Open: Final = lifecycle.Open
-Yield: Final = lifecycle.Yield
+StreamOpen: Final = lifecycle.StreamOpen
+StreamChunk: Final = lifecycle.StreamChunk
 
 _HEADERS: Final = TypeAdapter(Mapping[str, object])
 

tests/unit/rust_bridge/test_lifecycle.py

@@ -6,7 +6,17 @@ from typing import Final
 
 import pytest
 
-from litellm.rust_bridge.lifecycle import Await, Complete, Execution, Open, Step, drive
+from litellm.rust_bridge.lifecycle import (
+    Await,
+    Complete,
+    Execution,
+    Step,
+    Stream,
+    StreamChunk,
+    StreamOpen,
+    SyncStream,
+    drive,
+)
 
 
 class ScriptedExecution:
@@ -17,14 +27,11 @@ class ScriptedExecution:
         self.resumed: list[tuple[str, object]] = []
         self.closed = False
 
-    def start(self) -> Step:
-        return self._steps.pop(0)
-
-    def resume_value(self, value: object) -> Step:
+    def send(self, value: object) -> Step:
         self.resumed.append(("value", value))
         return self._steps.pop(0)
 
-    def resume_error(self, error: BaseException) -> Step:
+    def throw(self, error: BaseException) -> Step:
         self.resumed.append(("error", type(error)))
         return self._steps.pop(0)
 
@@ -43,9 +50,9 @@ async def failing() -> object:
 def test_drive_resumes_each_await_with_its_result_or_error_and_returns_the_completed_value() -> None:
     execution: Final = ScriptedExecution([Await(ready(1)), Await(failing()), Complete("done")])
 
-    assert asyncio.run(drive(execution)) == "done"
+    assert asyncio.run(drive(execution, Stream)) == "done"
 
-    assert execution.resumed == [("value", 1), ("error", ValueError)]
+    assert execution.resumed == [("value", None), ("value", 1), ("error", ValueError)]
     assert execution.closed
 
 
@@ -53,7 +60,7 @@ def test_drive_resumes_each_await_with_its_result_or_error_and_returns_the_compl
 def test_stream_handoff_preserves_head_identity_and_closes_on_construction_failure(factory_fails: bool) -> None:
     head: Final = object()
     stream: Final = object()
-    execution: Final = ScriptedExecution([Open(head)])
+    execution: Final = ScriptedExecution([StreamOpen(head)])
     failure: Final = ValueError("stream construction failed")
 
     def construct(owner: Execution, received: object) -> object:
@@ -72,3 +79,30 @@ def test_stream_handoff_preserves_head_identity_and_closes_on_construction_failu
         assert asyncio.run(drive(execution, construct)) is stream
         assert not execution.closed
         execution.close()
+
+
+def test_drive_rejects_a_chunk_before_the_stream_opens_and_closes_the_execution() -> None:
+    execution: Final = ScriptedExecution([StreamChunk("chunk")])
+
+    with pytest.raises(RuntimeError, match="chunk before opening a stream"):
+        asyncio.run(drive(execution, Stream))
+
+    assert execution.closed
+
+
+@pytest.mark.parametrize("asynchronous", (False, True))
+def test_stream_read_rejects_a_second_open_and_closes_the_execution(asynchronous: bool) -> None:
+    execution: Final = ScriptedExecution([StreamChunk("first"), StreamOpen("head")])
+
+    async def read() -> tuple[object, ...]:
+        stream: Final = Stream(execution)
+        return (await anext(stream), await anext(stream))
+
+    def read_sync() -> tuple[object, ...]:
+        stream: Final = SyncStream(execution)
+        return (next(stream), next(stream))
+
+    with pytest.raises(RuntimeError, match="opened a stream twice"):
+        asyncio.run(read()) if asynchronous else read_sync()
+
+    assert execution.closed

tests/unit/rust_bridge/test_runtime.py

@@ -12,7 +12,7 @@ from litellm.router_utils.add_retry_fallback_headers import get_hidden_params_di
 from litellm.rust_bridge import bindings, configuration, runtime
 from litellm.rust_bridge.catalog import Route, RouteContext, RouteRule
 from litellm.rust_bridge.configuration import Rollout
-from litellm.rust_bridge.lifecycle import Complete, Open, Yield
+from litellm.rust_bridge.lifecycle import Complete, StreamChunk
 from litellm.rust_bridge.streams import Stream, SyncStream
 
 
@@ -276,16 +276,13 @@ async def test_native_response_marker_reaches_caller_with_existing_metadata(
 
 class ScriptedStreamExecution:
     def __init__(self, chunks: tuple[bytes, ...]) -> None:
-        self._steps: Final = iter((*(Yield(chunk) for chunk in chunks), Complete(None)))
+        self._steps: Final = iter((*(StreamChunk(chunk) for chunk in chunks), Complete(None)))
         self.closed = False
 
-    def start(self) -> Open:
-        return Open(None)
-
-    def resume_value(self, value: object) -> Yield | Complete:
+    def send(self, value: object) -> StreamChunk | Complete:
         return next(self._steps)
 
-    def resume_error(self, error: BaseException) -> Complete:
+    def throw(self, error: BaseException) -> Complete:
         return Complete(None)
 
     def close(self) -> None: