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

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

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_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
 
@@ -64,22 +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:
+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):
-                factory: Final = Stream if stream_factory is None else stream_factory
-                stream: Final = factory(execution, head)
-                handed_off = True
+                stream: Final = open_stream(execution, head)
+                cleanup.pop_all()
                 return stream
             case StreamChunk():
                 raise RuntimeError("execution yielded a chunk before opening a stream")
-    finally:
-        if not handed_off:
-            execution.close()
 
 
 class Stream(AsyncIterator[object]):

tests/unit/rust_bridge/test_lifecycle.py

@@ -50,7 +50,7 @@ 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", None), ("value", 1), ("error", ValueError)]
     assert execution.closed
@@ -85,7 +85,7 @@ def test_drive_rejects_a_chunk_before_the_stream_opens_and_closes_the_execution(
     execution: Final = ScriptedExecution([StreamChunk("chunk")])
 
     with pytest.raises(RuntimeError, match="chunk before opening a stream"):
-        asyncio.run(drive(execution))
+        asyncio.run(drive(execution, Stream))
 
     assert execution.closed