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

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

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/handle.rs

@@ -13,8 +13,8 @@ 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 {
@@ -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();

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:

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:
@@ -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

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: