From f50f14117fd879f16caa761cbe34081618508786 Mon Sep 17 00:00:00 2001 From: yujonglee Date: Wed, 14 Oct 2026 17:19:07 -0700 Subject: [PATCH] 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 * refactor(rust-bridge): name stream steps StreamOpen and StreamChunk and reject out-of-phase steps Co-Authored-By: Yujong Lee * refactor(rust-bridge): drive with ExitStack and require the stream factory Co-Authored-By: Yujong Lee * test(host-python): give the hook_chain lifecycle binding the stream factory Co-Authored-By: Yujong Lee --------- Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm-rust/crates/host-python/src/caller.rs | 4 +- litellm-rust/crates/host-python/src/driver.rs | 4 +- litellm-rust/crates/host-python/src/error.rs | 6 +- litellm-rust/crates/host-python/src/handle.rs | 59 +++++++------ .../crates/host-python/tests/hook_chain.rs | 8 +- .../crates/host-python/tests/lifecycle.py | 12 ++- .../crates/host-python/tests/support/mod.rs | 5 +- litellm-rust/crates/python-bridge/AGENTS.md | 2 +- litellm/rust_bridge/AGENTS.md | 2 +- litellm/rust_bridge/lifecycle.py | 84 +++++++++++-------- litellm/rust_bridge/streams.py | 4 +- tests/unit/rust_bridge/test_lifecycle.py | 52 ++++++++++-- tests/unit/rust_bridge/test_runtime.py | 11 +-- 13 files changed, 157 insertions(+), 96 deletions(-) diff --git a/litellm-rust/crates/host-python/src/caller.rs b/litellm-rust/crates/host-python/src/caller.rs index 3bae948c936..5e80b294b88 100644 --- a/litellm-rust/crates/host-python/src/caller.rs +++ b/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) -> 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) -> 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. diff --git a/litellm-rust/crates/host-python/src/driver.rs b/litellm-rust/crates/host-python/src/driver.rs index da38588088e..fed4cc1eeea 100644 --- a/litellm-rust/crates/host-python/src/driver.rs +++ b/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()) } } diff --git a/litellm-rust/crates/host-python/src/error.rs b/litellm-rust/crates/host-python/src/error.rs index bec2ff20924..b22959432c1 100644 --- a/litellm-rust/crates/host-python/src/error.rs +++ b/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 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 diff --git a/litellm-rust/crates/host-python/src/handle.rs b/litellm-rust/crates/host-python/src/handle.rs index a9f03dc9724..39a36b72a20 100644 --- a/litellm-rust/crates/host-python/src/handle.rs +++ b/litellm-rust/crates/host-python/src/handle.rs @@ -13,8 +13,13 @@ pub enum Suspension { Await(Py), /// 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), - Yield(Py), + StreamOpen(Py), + StreamChunk(Py), +} + +enum Resume { + Send(Py), + Throw(PyErr), } pub enum ExecutionStep { @@ -72,37 +77,45 @@ impl Execution { } } - fn advance( - slf: &Bound<'_, Self>, - py: Python<'_>, - result: Option>>, - ) -> PyResult> { - let mut body = { + fn advance(slf: &Bound<'_, Self>, py: Python<'_>, input: Resume) -> PyResult> { + 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> { - Self::advance(slf, py, None) + fn send(slf: &Bound<'_, Self>, py: Python<'_>, value: Py) -> PyResult> { + Self::advance(slf, py, Resume::Send(value)) } - fn resume_value( - slf: &Bound<'_, Self>, - py: Python<'_>, - value: Py, - ) -> PyResult> { - Self::advance(slf, py, Some(Ok(value))) - } - - fn resume_error( + fn throw( slf: &Bound<'_, Self>, py: Python<'_>, error: Bound<'_, PyBaseException>, ) -> PyResult> { - 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>) { diff --git a/litellm-rust/crates/host-python/tests/hook_chain.rs b/litellm-rust/crates/host-python/tests/hook_chain.rs index 1943cdfcb2f..faae56e6daf 100644 --- a/litellm-rust/crates/host-python/tests/hook_chain.rs +++ b/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> { - 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") } diff --git a/litellm-rust/crates/host-python/tests/lifecycle.py b/litellm-rust/crates/host-python/tests/lifecycle.py index fd6742102a4..b4f39adf753 100644 --- a/litellm-rust/crates/host-python/tests/lifecycle.py +++ b/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): diff --git a/litellm-rust/crates/host-python/tests/support/mod.rs b/litellm-rust/crates/host-python/tests/support/mod.rs index c2899c16878..be93fb822a8 100644 --- a/litellm-rust/crates/host-python/tests/support/mod.rs +++ b/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( diff --git a/litellm-rust/crates/python-bridge/AGENTS.md b/litellm-rust/crates/python-bridge/AGENTS.md index 254a493d1ea..28ea62468f9 100644 --- a/litellm-rust/crates/python-bridge/AGENTS.md +++ b/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 diff --git a/litellm/rust_bridge/AGENTS.md b/litellm/rust_bridge/AGENTS.md index c3fb11f0bfb..9845ba9f18e 100644 --- a/litellm/rust_bridge/AGENTS.md +++ b/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 diff --git a/litellm/rust_bridge/lifecycle.py b/litellm/rust_bridge/lifecycle.py index fe95e9ca6a2..15e889c868c 100644 --- a/litellm/rust_bridge/lifecycle.py +++ b/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() diff --git a/litellm/rust_bridge/streams.py b/litellm/rust_bridge/streams.py index 6ead17777d2..4f8f935d3da 100644 --- a/litellm/rust_bridge/streams.py +++ b/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]) diff --git a/tests/unit/rust_bridge/test_lifecycle.py b/tests/unit/rust_bridge/test_lifecycle.py index 021a4aad85a..41f71c3332f 100644 --- a/tests/unit/rust_bridge/test_lifecycle.py +++ b/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 diff --git a/tests/unit/rust_bridge/test_runtime.py b/tests/unit/rust_bridge/test_runtime.py index bc3c7b43d75..86cc9eda5cc 100644 --- a/tests/unit/rust_bridge/test_runtime.py +++ b/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: