From 5bbfd27074048ba5eeb898497100a40c7bf2e9a1 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Tue, 13 Oct 2026 20:12:40 +0000 Subject: [PATCH] refactor(rust-bridge): name stream steps StreamOpen and StreamChunk and reject out-of-phase steps Co-Authored-By: Yujong Lee --- litellm-rust/crates/host-python/src/caller.rs | 4 +- litellm-rust/crates/host-python/src/driver.rs | 4 +- litellm-rust/crates/host-python/src/handle.rs | 12 ++-- litellm/rust_bridge/lifecycle.py | 57 ++++++++++++------- litellm/rust_bridge/streams.py | 4 +- tests/unit/rust_bridge/test_lifecycle.py | 41 ++++++++++++- tests/unit/rust_bridge/test_runtime.py | 6 +- 7 files changed, 93 insertions(+), 35 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/handle.rs b/litellm-rust/crates/host-python/src/handle.rs index 95ffc78985e..39a36b72a20 100644 --- a/litellm-rust/crates/host-python/src/handle.rs +++ b/litellm-rust/crates/host-python/src/handle.rs @@ -13,8 +13,8 @@ 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 { @@ -110,8 +110,12 @@ impl Execution { 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(); diff --git a/litellm/rust_bridge/lifecycle.py b/litellm/rust_bridge/lifecycle.py index 91af67fa7dd..9b19d976c12 100644 --- a/litellm/rust_bridge/lifecycle.py +++ b/litellm/rust_bridge/lifecycle.py @@ -16,16 +16,17 @@ 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 @@ -37,6 +38,9 @@ class Execution(Protocol): def close(self) -> None: ... +_REOPENED: Final = "execution opened a stream twice" + + class StreamClosed(Exception): """Tells a streaming execution that its caller stopped reading.""" @@ -63,13 +67,16 @@ def _settled(step: Step) -> Settled: 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.send(None)) - 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 + match await _settle(execution, execution.send(None)): + case Complete(value): + return value + case StreamOpen(head): + factory: Final = Stream if stream_factory is None else stream_factory + stream: Final = factory(execution, head) + handed_off = True + return stream + case StreamChunk(): + raise RuntimeError("execution yielded a chunk before opening a stream") finally: if not handed_off: execution.close() @@ -94,10 +101,15 @@ class Stream(AsyncIterator[object]): 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: @@ -131,10 +143,15 @@ class SyncStream(Iterator[object]): 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: 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 a2dcad7de0b..e4e01a2b849 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: @@ -50,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: @@ -69,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)) + + 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 84a8a50c4ab..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,10 +276,10 @@ 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 send(self, value: object) -> Yield | Complete: + def send(self, value: object) -> StreamChunk | Complete: return next(self._steps) def throw(self, error: BaseException) -> Complete: