test(host-python): move driver tests out of the inline module (#46325)

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

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

@@ -0,0 +1,866 @@
+mod support;
+
+use std::num::NonZeroUsize;
+use std::sync::Arc;
+use std::sync::atomic::{AtomicBool, Ordering};
+
+use litellm_host::{
+    hooks::CallHooks,
+    lifecycle::{CallEvent, ExecutionEvent, FailureOrigin},
+    machine::CallMachine,
+    observation::observation_channel,
+    protocol::Reply,
+};
+use litellm_host_python::{
+    CallOptions, HookChain, InvokeError, PythonBinding, PythonHostCalls, PythonOwned,
+    PythonRuntime, missing_state, run_call,
+};
+use pyo3::{
+    exceptions::{
+        PyBaseException, PyException, PyLookupError, PyRuntimeError, PyStopIteration, PyTypeError,
+        PyValueError, asyncio::CancelledError,
+    },
+    gc::{PyTraverseError, PyVisit},
+    prelude::*,
+    types::PyDict,
+};
+use rstest::rstest;
+use support::{
+    Classified, Error, HookScript, Interpreter, Log, MachineFactory, OpScript, Synthetic, binding,
+    call_options, failing_machine, hooks, interpreter, lifecycle_binding, open_stream, read_all,
+    run_composed, run_scripted, settled, success_machine,
+};
+
+fn observed(capacity: usize) -> (CallOptions, impl FnMut() -> Vec<CallEvent>) {
+    let (sender, mut receiver) = observation_channel(NonZeroUsize::new(capacity).unwrap());
+    let options = CallOptions {
+        asynchronous: false,
+        lifecycle: lifecycle_binding,
+        observers: Some(sender),
+    };
+    (options, move || {
+        std::iter::from_fn(|| receiver.try_recv().ok()).collect()
+    })
+}
+
+#[rstest]
+#[case::success(OpScript::AwaitAnswer)]
+#[case::failure(OpScript::AwaitFailure)]
+fn asynchronous_host_operations_resume_the_same_machine(
+    #[from(interpreter)] _python: Interpreter,
+    #[case] op: OpScript,
+) {
+    Python::attach(|py| {
+        let machine: MachineFactory = Box::new(|_| {
+            CallMachine::<Synthetic>::new(None, |host| {
+                Box::pin(async move { host.services.call(|reply| ("read", reply)).await })
+            })
+        });
+        let (result, log) = run_scripted(py, machine, op, HookScript::Plain, true);
+        match op {
+            OpScript::AwaitAnswer => {
+                assert_eq!(result.unwrap().extract::<String>(py).unwrap(), "awaited")
+            }
+            OpScript::AwaitFailure => {
+                assert!(result.unwrap_err().is_instance_of::<PyLookupError>(py));
+                assert_eq!(
+                    log.iter()
+                        .filter(|entry| entry.starts_with("failed:"))
+                        .count(),
+                    1
+                );
+                assert!(!log.iter().any(|entry| entry.starts_with("succeeded:")));
+            }
+            _ => unreachable!(),
+        }
+    });
+}
+
+#[rstest]
+fn composed_hooks_share_prepared_arguments_and_one_execution(
+    #[from(interpreter)] _python: Interpreter,
+    #[values(false, true)] asynchronous: bool,
+) {
+    Python::attach(|py| {
+        let log = Log::default();
+        let chain = HookChain::new()
+            .with(hooks(&log, HookScript::ObserveArguments))
+            .with(hooks(&log, HookScript::RewriteArguments))
+            .with(hooks(&log, HookScript::ReplaceResponse));
+        let result = run_call(
+            py,
+            move |_, _, request| Ok(success_machine()(request)),
+            binding(&log, OpScript::Answer),
+            chain,
+            PyDict::new(py).unbind(),
+            call_options(asynchronous),
+        );
+        if asynchronous {
+            assert!(log.entries().is_empty());
+        }
+        let response = settled(py, result, asynchronous).unwrap();
+        assert_eq!(response.extract::<String>(py).unwrap(), "replaced");
+        let entries = log.entries();
+        assert!(entries.iter().any(|entry| entry == "adopted:hook"));
+        assert_eq!(
+            entries.iter().filter(|entry| *entry == "project").count(),
+            1
+        );
+        assert_eq!(
+            entries.iter().filter(|entry| *entry == "op:sign").count(),
+            1
+        );
+        assert_eq!(
+            entries
+                .iter()
+                .filter(|entry| *entry == "succeeded:replaced")
+                .count(),
+            3
+        );
+        assert!(!entries.iter().any(|entry| entry.starts_with("failed:")));
+    });
+}
+
+#[rstest]
+#[case::draining(4, false)]
+#[case::full(1, false)]
+#[case::closed(4, true)]
+fn observation_delivery_preserves_callback_order_and_response(
+    #[from(interpreter)] _python: Interpreter,
+    #[case] capacity: usize,
+    #[case] closed: bool,
+    #[values(false, true)] asynchronous: bool,
+) {
+    Python::attach(|py| {
+        let (expected, expected_log) = run_scripted(
+            py,
+            success_machine(),
+            OpScript::Answer,
+            HookScript::ReplaceResponse,
+            asynchronous,
+        );
+        let (sender, mut receiver) = observation_channel(NonZeroUsize::new(capacity).unwrap());
+        if closed {
+            receiver.close();
+        }
+        let log = Log::default();
+        let (actual, actual_log) = run_composed(
+            py,
+            success_machine(),
+            binding(&log, OpScript::Answer),
+            hooks(&log, HookScript::ReplaceResponse),
+            std::convert::identity,
+            CallOptions {
+                asynchronous,
+                lifecycle: lifecycle_binding,
+                observers: Some(sender.clone()),
+            },
+        );
+        assert_eq!(
+            actual.unwrap().extract::<String>(py).unwrap(),
+            expected.unwrap().extract::<String>(py).unwrap()
+        );
+        assert_eq!(actual_log, expected_log);
+        let events: Vec<_> = std::iter::from_fn(|| receiver.try_recv().ok()).collect();
+        assert_eq!(events.len() as u64 + sender.dropped_events(), 3);
+        if !closed && capacity >= 3 {
+            let [
+                CallEvent::Started { start_time },
+                CallEvent::Execution(ExecutionEvent::ProviderResponseReceived { raw }),
+                CallEvent::Succeeded { timing, .. },
+            ] = events.as_slice()
+            else {
+                panic!("expected start, provider response and success: {events:?}");
+            };
+            assert_eq!(raw.body, "raw");
+            assert_eq!(*start_time, timing.start_time);
+            assert!(timing.end_time >= timing.start_time);
+        }
+    });
+}
+
+#[rstest]
+#[case::preparation(HookScript::FailBegin, false, 1)]
+#[case::transformation(HookScript::FailAfterSuccess, false, 1)]
+#[case::terminal_cancellation(HookScript::CancelTerminal, true, 0)]
+#[case::terminal_failure(HookScript::FailTerminal, false, 0)]
+fn a_terminal_hook_failure_is_reported_once_and_never_replays_callbacks(
+    #[from(interpreter)] _python: Interpreter,
+    #[case] script: HookScript,
+    #[case] cancelled: bool,
+    #[case] failure_callbacks: usize,
+    #[values(false, true)] asynchronous: bool,
+) {
+    Python::attach(|py| {
+        let (options, mut events) = observed(4);
+        let log = Log::default();
+        let (result, log) = run_composed(
+            py,
+            success_machine(),
+            binding(&log, OpScript::Answer),
+            hooks(&log, script),
+            std::convert::identity,
+            CallOptions {
+                asynchronous,
+                ..options
+            },
+        );
+        let error = result.unwrap_err();
+        assert_eq!(error.is_instance_of::<CancelledError>(py), cancelled);
+        let events = events();
+        assert!(matches!(events.first(), Some(CallEvent::Started { .. })));
+        assert_eq!(
+            matches!(events.last(), Some(CallEvent::Cancelled { .. })),
+            cancelled
+        );
+        assert_eq!(
+            matches!(
+                events.last(),
+                Some(CallEvent::Failed {
+                    origin: FailureOrigin::Host,
+                    ..
+                })
+            ),
+            !cancelled
+        );
+        let terminal = events.iter().filter(|event| {
+            matches!(
+                event,
+                CallEvent::Failed { .. }
+                    | CallEvent::Succeeded { .. }
+                    | CallEvent::Cancelled { .. }
+            )
+        });
+        assert_eq!(terminal.count(), 1);
+        assert_eq!(
+            log.iter()
+                .filter(|entry| entry.starts_with("failed:"))
+                .count(),
+            failure_callbacks
+        );
+        assert!(
+            !log.iter()
+                .any(|entry| entry.starts_with("succeeded:") || entry == "cancelled")
+        );
+    });
+}
+
+#[rstest]
+#[case::preparation(HookScript::FailBegin, OpScript::Answer)]
+#[case::native_decode(HookScript::Plain, OpScript::RejectRequestNatively)]
+#[case::python_decode(HookScript::Plain, OpScript::RaiseRequestPython)]
+fn startup_failures_release_the_factory_without_constructing_a_machine(
+    #[from(interpreter)] _python: Interpreter,
+    #[case] script: HookScript,
+    #[case] op: OpScript,
+    #[values(false, true)] asynchronous: bool,
+) {
+    struct Release(Arc<AtomicBool>);
+    impl Drop for Release {
+        fn drop(&mut self) {
+            self.0.store(true, Ordering::SeqCst);
+        }
+    }
+
+    Python::attach(|py| {
+        let constructed = Arc::new(AtomicBool::new(false));
+        let did_construct = constructed.clone();
+        let released = Arc::new(AtomicBool::new(false));
+        let release = Release(released.clone());
+        let machine: MachineFactory = Box::new(move |request| {
+            let _release = release;
+            did_construct.store(true, Ordering::SeqCst);
+            success_machine()(request)
+        });
+        let (result, log) = run_scripted(py, machine, op, script, asynchronous);
+        assert!(result.is_err());
+        assert!(!constructed.load(Ordering::SeqCst));
+        assert!(released.load(Ordering::SeqCst));
+        assert_eq!(
+            log.iter()
+                .filter(|event| event.starts_with("failed:"))
+                .count(),
+            1
+        );
+        assert_eq!(
+            log.iter().filter(|event| *event == "adapter.close").count(),
+            1
+        );
+        assert_eq!(log.iter().filter(|event| *event == "host.close").count(), 1);
+    });
+}
+
+#[rstest]
+fn success_runs_every_step_in_order_and_returns_the_public_response(
+    #[from(interpreter)] _python: Interpreter,
+    #[values(false, true)] asynchronous: bool,
+) {
+    Python::attach(|py| {
+        let (result, log) = run_scripted(
+            py,
+            success_machine(),
+            OpScript::Answer,
+            HookScript::Plain,
+            asynchronous,
+        );
+        assert_eq!(
+            result.unwrap().extract::<String>(py).unwrap(),
+            "project:1|sign|rewritten"
+        );
+        assert_eq!(
+            log,
+            [
+                "started",
+                "begin",
+                "project",
+                "op:sign",
+                "before_provider_request",
+                "response:raw",
+                "complete",
+                "after_success",
+                "succeeded:project:1|sign|rewritten",
+                "adapter.close",
+                "host.close",
+            ]
+        );
+    });
+}
+
+#[rstest]
+fn explicitly_closing_a_stream_dispatches_success_for_delivered_chunks(
+    #[from(interpreter)] _python: Interpreter,
+    #[values(false, true)] asynchronous: bool,
+) {
+    Python::attach(|py| {
+        let log = Log::default();
+        let (options, mut events) = observed(3);
+        let stream = open_stream(
+            py,
+            hooks(&log, HookScript::Plain),
+            CallOptions {
+                asynchronous,
+                ..options
+            },
+        );
+        if asynchronous {
+            for method in ["__anext__", "aclose", "aclose"] {
+                let stop = stream
+                    .call_method0(method)
+                    .unwrap()
+                    .call_method1("send", (py.None(),))
+                    .unwrap_err();
+                assert!(stop.is_instance_of::<PyStopIteration>(py));
+            }
+        } else {
+            assert_eq!(
+                stream
+                    .call_method0("__next__")
+                    .unwrap()
+                    .extract::<String>()
+                    .unwrap(),
+                "first"
+            );
+            stream.call_method0("close").unwrap();
+            stream.call_method0("close").unwrap();
+        }
+        assert!(matches!(
+            events().as_slice(),
+            [CallEvent::Started { .. }, CallEvent::Succeeded { .. }]
+        ));
+        assert_eq!(
+            log.entries(),
+            [
+                "started",
+                "begin",
+                "opened",
+                "delivered",
+                "succeeded:None",
+                "adapter.close"
+            ]
+        );
+    });
+}
+
+#[rstest]
+fn a_stream_carries_its_head_before_the_first_chunk(
+    #[from(interpreter)] _python: Interpreter,
+    #[values(false, true)] asynchronous: bool,
+) {
+    Python::attach(|py| {
+        let log = Log::default();
+        let stream = open_stream(
+            py,
+            hooks(&log, HookScript::Plain),
+            call_options(asynchronous),
+        );
+        let hidden: std::collections::HashMap<String, std::collections::HashMap<String, String>> =
+            stream.getattr("head").unwrap().extract().unwrap();
+        assert_eq!(
+            hidden["additional_headers"],
+            std::collections::HashMap::from([("request-id".to_string(), "req_1".to_string())])
+        );
+        assert_eq!(log.entries(), ["started", "begin", "opened"]);
+        assert_eq!(read_all(py, &stream, asynchronous), ["first", "second"]);
+    });
+}
+
+#[rstest]
+fn a_native_failure_is_classified_once_and_reported_classified(
+    #[from(interpreter)] _python: Interpreter,
+    #[values(false, true)] asynchronous: bool,
+) {
+    Python::attach(|py| {
+        let (result, log) = run_scripted(
+            py,
+            failing_machine(),
+            OpScript::Answer,
+            HookScript::Plain,
+            asynchronous,
+        );
+        let error = result.unwrap_err();
+        assert!(error.is_instance_of::<PyValueError>(py));
+        assert_eq!(error.value(py).to_string(), "classified: provider exploded");
+        assert_eq!(
+            log,
+            [
+                "started",
+                "begin",
+                "project",
+                "classify:provider exploded",
+                "failed:Call:classified: provider exploded",
+                "adapter.close",
+                "host.close",
+            ]
+        );
+    });
+}
+
+#[rstest]
+fn a_native_rejection_from_a_host_operation_is_classified_once(
+    #[from(interpreter)] _python: Interpreter,
+) {
+    Python::attach(|py| {
+        let (result, log) = run_scripted(
+            py,
+            success_machine(),
+            OpScript::RejectNatively,
+            HookScript::Plain,
+            false,
+        );
+        assert_eq!(
+            result.unwrap_err().value(py).to_string(),
+            "classified: op rejected"
+        );
+        assert_eq!(
+            log,
+            [
+                "started",
+                "begin",
+                "project",
+                "op:sign",
+                "classify:op rejected",
+                "failed:Call:classified: op rejected",
+                "adapter.close",
+                "host.close",
+            ]
+        );
+    });
+}
+
+#[rstest]
+fn a_python_exception_from_a_host_operation_is_reported_as_raised(
+    #[from(interpreter)] _python: Interpreter,
+) {
+    Python::attach(|py| {
+        let (result, log) = run_scripted(
+            py,
+            success_machine(),
+            OpScript::RaisePython,
+            HookScript::Plain,
+            false,
+        );
+        let error = result.unwrap_err();
+        assert!(error.is_instance_of::<PyValueError>(py));
+        assert_eq!(error.value(py).to_string(), "op failed");
+        assert_eq!(
+            log,
+            [
+                "started",
+                "begin",
+                "project",
+                "op:sign",
+                "failed:Call:op failed",
+                "adapter.close",
+                "host.close",
+            ]
+        );
+    });
+}
+
+#[rstest]
+fn a_failing_classifier_surfaces_with_the_native_error_as_context(
+    #[from(interpreter)] _python: Interpreter,
+) {
+    Python::attach(|py| {
+        let log_handle = Log::default();
+        let (result, log) = run_composed(
+            py,
+            failing_machine(),
+            binding(&log_handle, OpScript::Answer).with_failing_classifier(),
+            hooks(&log_handle, HookScript::Plain),
+            std::convert::identity,
+            call_options(false),
+        );
+        let error = result.unwrap_err();
+        assert!(error.is_instance_of::<PyTypeError>(py));
+        assert_eq!(error.value(py).to_string(), "classifier failed");
+        let context = error.context(py).unwrap();
+        assert!(context.is_instance_of::<PyRuntimeError>(py));
+        assert_eq!(context.value(py).to_string(), "provider exploded");
+        assert_eq!(
+            log,
+            [
+                "started",
+                "begin",
+                "project",
+                "classify:provider exploded",
+                "failed:Call:classifier failed",
+                "adapter.close",
+                "host.close",
+            ]
+        );
+    });
+}
+
+#[rstest]
+fn begin_failures_are_host_failures_without_provider_mapping(
+    #[from(interpreter)] _python: Interpreter,
+) {
+    Python::attach(|py| {
+        let (result, log) = run_scripted(
+            py,
+            success_machine(),
+            OpScript::Answer,
+            HookScript::FailBegin,
+            false,
+        );
+        assert_eq!(result.unwrap_err().value(py).to_string(), "begin failed");
+        assert_eq!(
+            log,
+            [
+                "started",
+                "begin",
+                "failed:Host:begin failed",
+                "adapter.close",
+                "host.close"
+            ]
+        );
+    });
+}
+
+enum ArgumentPolicy {
+    Inherit,
+    Reject(Py<PyBaseException>),
+}
+
+impl CallHooks<PythonRuntime> for ArgumentPolicy {
+    fn arguments_prepared(&mut self, py: Python<'_>, arguments: &Py<PyDict>) -> PyResult<()> {
+        match self {
+            Self::Inherit => arguments.bind(py).set_item("api_key", "inherited"),
+            Self::Reject(error) => Err(PyErr::from_value(error.bind(py).clone().into_any())),
+        }
+    }
+}
+
+impl PythonOwned for ArgumentPolicy {
+    fn close(&mut self, _: Python<'_>) {}
+
+    fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> {
+        match self {
+            Self::Inherit => Ok(()),
+            Self::Reject(error) => visit.call(error),
+        }
+    }
+}
+
+#[rstest]
+fn resource_setup_uses_prepared_arguments_and_failures_are_terminal(
+    #[from(interpreter)] _python: Interpreter,
+    #[values(false, true)] asynchronous: bool,
+    #[values(false, true)] fail_setup: bool,
+) {
+    Python::attach(|py| {
+        let log = Log::default();
+        let setup_log = log.clone();
+        let error = PyValueError::new_err("setup failed").into_value(py);
+        let setup_error = error.clone_ref(py);
+        let arguments = PyDict::new(py);
+        arguments.set_item("api_key", "original").unwrap();
+        let result = run_call(
+            py,
+            move |py, prepared, request| {
+                let source: String = prepared.get_item("prepared")?.unwrap().extract()?;
+                let key: String = prepared.get_item("api_key")?.unwrap().extract()?;
+                setup_log.push(format!("setup:{source}:{key}"));
+                if fail_setup {
+                    return Err(PyErr::from_value(setup_error.into_bound(py).into_any()));
+                }
+                Ok(success_machine()(request))
+            },
+            binding(&log, OpScript::Answer),
+            HookChain::new()
+                .with(hooks(&log, HookScript::RewriteArguments))
+                .with(ArgumentPolicy::Inherit),
+            arguments.clone().unbind(),
+            call_options(asynchronous),
+        );
+        if asynchronous {
+            assert!(log.entries().is_empty());
+        }
+        let settled = settled(py, result, asynchronous);
+        assert_eq!(
+            arguments
+                .get_item("api_key")
+                .unwrap()
+                .unwrap()
+                .extract::<String>()
+                .unwrap(),
+            "original"
+        );
+        assert_eq!(
+            &log.entries()[..4],
+            ["started", "begin", "project", "setup:hook:inherited"]
+        );
+        if fail_setup {
+            assert!(settled.unwrap_err().value(py).is(error.bind(py)));
+            assert_eq!(
+                log.entries(),
+                [
+                    "started",
+                    "begin",
+                    "project",
+                    "setup:hook:inherited",
+                    "failed:Host:setup failed",
+                    "adapter.close",
+                    "host.close",
+                ]
+            );
+        } else {
+            assert!(settled.is_ok());
+            assert!(
+                log.entries()
+                    .iter()
+                    .any(|entry| entry.starts_with("succeeded:"))
+            );
+        }
+    });
+}
+
+#[rstest]
+fn closing_an_unstarted_call_never_initializes_resources(
+    #[from(interpreter)] _python: Interpreter,
+) {
+    Python::attach(|py| {
+        let log = Log::default();
+        let setup_log = log.clone();
+        let pending = run_call(
+            py,
+            move |_, _, request| {
+                setup_log.push("setup");
+                Ok(success_machine()(request))
+            },
+            binding(&log, OpScript::Answer),
+            hooks(&log, HookScript::Plain),
+            PyDict::new(py).unbind(),
+            call_options(true),
+        )
+        .unwrap();
+        assert!(log.entries().is_empty());
+        pending.call_method0(py, "close").unwrap();
+        drop(pending);
+        assert_eq!(log.entries(), ["adapter.close", "host.close"]);
+    });
+}
+
+#[rstest]
+fn an_argument_policy_rejection_is_the_callers_error_and_the_machine_never_starts(
+    #[from(interpreter)] _python: Interpreter,
+    #[values(false, true)] asynchronous: bool,
+) {
+    Python::attach(|py| {
+        let raised = PyValueError::new_err("over budget").into_value(py);
+        let log = Log::default();
+        let (result, log) = run_composed(
+            py,
+            success_machine(),
+            binding(&log, OpScript::Answer),
+            hooks(&log, HookScript::Plain),
+            |hooks| {
+                HookChain::new()
+                    .with(hooks)
+                    .with(ArgumentPolicy::Reject(raised.clone_ref(py)))
+            },
+            call_options(asynchronous),
+        );
+        assert!(result.unwrap_err().value(py).is(raised.bind(py)));
+        assert_eq!(
+            log,
+            [
+                "started",
+                "begin",
+                "failed:Host:over budget",
+                "adapter.close",
+                "host.close"
+            ]
+        );
+    });
+}
+
+#[rstest]
+fn the_host_projects_from_the_keyword_view_the_argument_policy_rewrote(
+    #[from(interpreter)] _python: Interpreter,
+    #[values(false, true)] asynchronous: bool,
+) {
+    Python::attach(|py| {
+        let log = Log::default();
+        let (result, _) = run_composed(
+            py,
+            success_machine(),
+            binding(&log, OpScript::Answer),
+            hooks(&log, HookScript::Plain),
+            |hooks| HookChain::new().with(hooks).with(ArgumentPolicy::Inherit),
+            call_options(asynchronous),
+        );
+        assert_eq!(
+            result.unwrap().extract::<String>(py).unwrap(),
+            "project:2|sign|rewritten"
+        );
+    });
+}
+
+#[rstest]
+fn the_adapters_finalized_response_is_what_the_call_returns_and_reports(
+    #[from(interpreter)] _python: Interpreter,
+    #[values(false, true)] asynchronous: bool,
+) {
+    Python::attach(|py| {
+        let (result, log) = run_scripted(
+            py,
+            success_machine(),
+            OpScript::Answer,
+            HookScript::ReplaceResponse,
+            asynchronous,
+        );
+        assert_eq!(result.unwrap().extract::<String>(py).unwrap(), "replaced");
+        assert!(log.contains(&"succeeded:replaced".to_string()));
+        assert!(!log.contains(&"succeeded:project:1|rewritten".to_string()));
+    });
+}
+
+#[rstest]
+fn a_failure_while_finalizing_fails_the_call_instead_of_succeeding(
+    #[from(interpreter)] _python: Interpreter,
+    #[values(false, true)] asynchronous: bool,
+) {
+    Python::attach(|py| {
+        let (result, log) = run_scripted(
+            py,
+            success_machine(),
+            OpScript::Answer,
+            HookScript::FailAfterSuccess,
+            asynchronous,
+        );
+        assert_eq!(
+            result.unwrap_err().value(py).to_string(),
+            "after_success failed"
+        );
+        assert_eq!(
+            &log[log.len() - 4..],
+            [
+                "after_success",
+                "failed:Host:after_success failed",
+                "adapter.close",
+                "host.close"
+            ]
+        );
+        assert!(!log.iter().any(|entry| entry.starts_with("succeeded")));
+    });
+}
+
+#[rstest]
+fn cancellation_ends_the_call_without_terminal_dispatch(#[from(interpreter)] _python: Interpreter) {
+    Python::attach(|py| {
+        struct Cancelling(Log);
+        impl PythonBinding for Cancelling {
+            type Protocol = Synthetic;
+            type Failure = Classified;
+            fn decode_request(
+                &mut self,
+                _: Python<'_>,
+                _: &Bound<'_, PyDict>,
+            ) -> Result<String, InvokeError<Error>> {
+                self.0.push("project");
+                Err(pyo3::exceptions::asyncio::CancelledError::new_err(()).into())
+            }
+
+            fn encode_stream_head(
+                &mut self,
+                _: Python<'_>,
+                head: std::convert::Infallible,
+            ) -> PyResult<Py<PyAny>> {
+                match head {}
+            }
+            fn encode_chunk(
+                &mut self,
+                _: Python<'_>,
+                chunk: std::convert::Infallible,
+            ) -> PyResult<Py<PyAny>> {
+                match chunk {}
+            }
+            fn encode_response(&mut self, _: Python<'_>, _: String) -> PyResult<Py<PyAny>> {
+                Err(missing_state())
+            }
+            fn map_error(&self, _: Python<'_>, error: Error) -> PyResult<Classified> {
+                self.0.push("classify");
+                Ok(Classified(error.0))
+            }
+            fn host_error(error: &PyErr) -> Error {
+                Error(error.to_string())
+            }
+        }
+
+        impl PythonHostCalls<Synthetic> for Cancelling {
+            fn handle_host_call(
+                &mut self,
+                _: Python<'_>,
+                _: (&'static str, Reply<String>),
+            ) -> Result<(), InvokeError<Error>> {
+                Err(missing_state().into())
+            }
+        }
+
+        impl PythonOwned for Cancelling {
+            fn close(&mut self, _: Python<'_>) {}
+            fn traverse(&self, _: &PyVisit<'_>) -> Result<(), PyTraverseError> {
+                Ok(())
+            }
+        }
+        let log = Log::default();
+        let host = Cancelling(log.clone());
+        let adapter = hooks(&log, HookScript::Plain);
+        let error = run_call(
+            py,
+            |_, _, request| Ok(success_machine()(request)),
+            host,
+            adapter,
+            PyDict::new(py).unbind(),
+            call_options(false),
+        )
+        .unwrap_err();
+        assert!(!error.is_instance_of::<PyException>(py));
+        assert_eq!(
+            log.entries(),
+            ["started", "begin", "project", "adapter.close"]
+        );
+    });
+}

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

@@ -0,0 +1,201 @@
+mod support;
+
+use litellm_host_python::{Execution, ExecutionBody, ExecutionStep};
+use pyo3::{
+    exceptions::PyBaseException,
+    gc::{PyTraverseError, PyVisit},
+    prelude::*,
+    types::PyDict,
+};
+use rstest::rstest;
+use support::{Interpreter, install_lifecycle_module, interpreter, lifecycle_binding};
+
+#[rstest]
+fn python_driver_preserves_inline_await_and_native_ownership(
+    #[from(interpreter)] _python: Interpreter,
+) {
+    Python::attach(|py| {
+        py.import("asyncio").unwrap();
+        let module = install_lifecycle_module(py);
+        let locals = PyDict::new(py);
+        locals
+            .set_item("drive", module.getattr("drive").unwrap())
+            .unwrap();
+        locals
+            .set_item(
+                "await_execution",
+                wrap_pyfunction!(await_execution, py).unwrap(),
+            )
+            .unwrap();
+        locals
+            .set_item(
+                "calling_execution",
+                wrap_pyfunction!(calling_execution, py).unwrap(),
+            )
+            .unwrap();
+        let probe = std::ffi::CString::new(include_str!("lifecycle.py")).unwrap();
+        py.run(&probe, Some(&locals), Some(&locals)).unwrap();
+    });
+}
+struct RetainingHost {
+    retained: Option<Py<PyAny>>,
+}
+
+impl ExecutionBody for RetainingHost {
+    fn resume(&mut self, _: Option<PyResult<Py<PyAny>>>) -> PyResult<ExecutionStep> {
+        Python::attach(|py| Ok(ExecutionStep::Return(py.None())))
+    }
+
+    fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> {
+        visit.call(&self.retained)
+    }
+}
+
+#[pyfunction]
+fn retaining_coroutine(py: Python<'_>, retained: Py<PyAny>) -> PyResult<Py<Execution>> {
+    Py::new(
+        py,
+        Execution::new(
+            RetainingHost {
+                retained: Some(retained),
+            },
+            lifecycle_binding,
+        ),
+    )
+}
+
+struct AwaitBody(Option<Py<PyAny>>);
+
+impl ExecutionBody for AwaitBody {
+    fn resume(&mut self, result: Option<PyResult<Py<PyAny>>>) -> PyResult<ExecutionStep> {
+        match self.0.take() {
+            Some(awaitable) => Ok(ExecutionStep::Await(awaitable)),
+            None => result
+                .expect("selected await completed")
+                .map(ExecutionStep::Return),
+        }
+    }
+
+    fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> {
+        visit.call(&self.0)
+    }
+}
+
+#[pyfunction]
+fn await_execution(awaitable: Py<PyAny>) -> Execution {
+    Execution::new(AwaitBody(Some(awaitable)), lifecycle_binding)
+}
+
+struct CallingBody(Py<PyAny>);
+
+impl ExecutionBody for CallingBody {
+    fn resume(&mut self, _: Option<PyResult<Py<PyAny>>>) -> PyResult<ExecutionStep> {
+        Python::attach(|py| self.0.call0(py).map(ExecutionStep::Return))
+    }
+
+    fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> {
+        visit.call(&self.0)
+    }
+}
+
+#[pyfunction]
+fn calling_execution(callback: Py<PyAny>) -> Execution {
+    Execution::new(CallingBody(callback), lifecycle_binding)
+}
+
+struct ErrorBody(Option<Py<PyBaseException>>);
+
+impl ExecutionBody for ErrorBody {
+    fn resume(&mut self, _: Option<PyResult<Py<PyAny>>>) -> PyResult<ExecutionStep> {
+        Python::attach(|py| {
+            Err(PyErr::from_value(
+                self.0.take().unwrap().into_bound(py).into_any(),
+            ))
+        })
+    }
+
+    fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> {
+        visit.call(&self.0)
+    }
+}
+
+#[pyfunction]
+fn error_execution(error: Bound<'_, PyBaseException>) -> Execution {
+    Execution::new(ErrorBody(Some(error.unbind())), lifecycle_binding)
+}
+
+#[rstest]
+fn retained_exception_frames_are_collectable(#[from(interpreter)] _python: Interpreter) {
+    Python::attach(|py| {
+        let locals = PyDict::new(py);
+        locals
+            .set_item(
+                "error_execution",
+                wrap_pyfunction!(error_execution, py).unwrap(),
+            )
+            .unwrap();
+        py.run(
+            pyo3::ffi::c_str!(
+                r#"
+import gc
+import weakref
+
+class Retained:
+    pass
+
+def cycle():
+    retained = Retained()
+    try:
+        raise ValueError('retained traceback')
+    except ValueError as error:
+        retained.owner = error_execution(error)
+    return weakref.ref(retained)
+
+reference = cycle()
+gc.collect()
+assert reference() is None
+"#
+            ),
+            Some(&locals),
+            Some(&locals),
+        )
+        .unwrap();
+    });
+}
+
+#[rstest]
+fn cycles_through_a_retained_execution_are_collectable(#[from(interpreter)] _python: Interpreter) {
+    Python::attach(|py| {
+        let locals = PyDict::new(py);
+        locals
+            .set_item(
+                "retaining_coroutine",
+                wrap_pyfunction!(retaining_coroutine, py).unwrap(),
+            )
+            .unwrap();
+        py.run(
+            pyo3::ffi::c_str!(
+                r#"
+import gc
+import weakref
+
+class Retained:
+    pass
+
+def cycle():
+    retained = Retained()
+    coroutine = retaining_coroutine(retained)
+    retained.coroutine = coroutine
+    return weakref.ref(retained)
+
+retained_ref = cycle()
+gc.collect()
+assert retained_ref() is None
+"#
+            ),
+            Some(&locals),
+            Some(&locals),
+        )
+        .unwrap();
+    });
+}

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

@@ -0,0 +1,656 @@
+#![allow(dead_code)] // each test binary compiles this module on its own and uses a different subset
+
+use std::sync::{Arc, Mutex, MutexGuard, PoisonError};
+
+use litellm_host::{
+    call::{CallOutput, HostedMachine, hosted_call},
+    hooks::CallHooks,
+    interceptors::{Interceptors, RawResponse, RequestContext, WireRequest},
+    lifecycle::{ExecutionEvent, Timing},
+    machine::{CallMachine, MachineFault},
+    protocol::{Protocol, Reply},
+};
+use litellm_host_python::{
+    CallOptions, HookStep, InvokeError, PythonBinding, PythonCallEvent, PythonCallHooks,
+    PythonHostCalls, PythonOwned, PythonRuntime, run_call,
+};
+use pyo3::{
+    exceptions::{PyTypeError, PyValueError, asyncio::CancelledError},
+    gc::{PyTraverseError, PyVisit},
+    prelude::*,
+    types::{PyDict, PyString},
+};
+use rstest::fixture;
+
+static PYTHON_GLOBALS: Mutex<()> = Mutex::new(());
+
+/// One initialized interpreter with the real lifecycle protocol installed, held for the
+/// length of a test so tests sharing the interpreter run one at a time.
+pub struct Interpreter(MutexGuard<'static, ()>);
+
+#[fixture]
+pub fn interpreter() -> Interpreter {
+    let guard = PYTHON_GLOBALS
+        .lock()
+        .unwrap_or_else(PoisonError::into_inner);
+    Python::initialize();
+    Python::attach(|py| {
+        py.import("asyncio").unwrap();
+        install_lifecycle_module(py);
+    });
+    Interpreter(guard)
+}
+
+pub fn lifecycle_binding(py: Python<'_>) -> PyResult<Bound<'_, PyModule>> {
+    py.import("host_test_lifecycle")
+}
+
+pub fn call_options(asynchronous: bool) -> CallOptions {
+    CallOptions::new(asynchronous, lifecycle_binding)
+}
+
+pub fn install_lifecycle_module(py: Python<'_>) -> Bound<'_, PyModule> {
+    let source = std::ffi::CString::new(include_str!(
+        "../../../../../litellm/rust_bridge/lifecycle.py"
+    ))
+    .unwrap();
+    PyModule::from_code(
+        py,
+        &source,
+        pyo3::ffi::c_str!("lifecycle.py"),
+        pyo3::ffi::c_str!("host_test_lifecycle"),
+    )
+    .unwrap()
+}
+
+#[derive(Clone, Debug, PartialEq, Eq)]
+pub struct Error(pub String);
+
+impl std::fmt::Display for Error {
+    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        formatter.write_str(&self.0)
+    }
+}
+
+impl From<MachineFault> for Error {
+    fn from(fault: MachineFault) -> Self {
+        Self(format!("{fault:?}"))
+    }
+}
+
+pub struct Synthetic;
+
+impl Protocol for Synthetic {
+    type Response = String;
+    type Error = Error;
+    type Request = String;
+    type HostCall = (&'static str, Reply<String>);
+    type Chunk = std::convert::Infallible;
+    type StreamHead = std::convert::Infallible;
+}
+
+pub fn wire() -> WireRequest {
+    WireRequest {
+        url: "https://example.invalid".into(),
+        headers: Vec::new(),
+        body: serde_json::json!({}),
+    }
+}
+
+pub fn context() -> RequestContext {
+    RequestContext {
+        model: "model".into(),
+        custom_llm_provider: "provider".into(),
+        optional_params: serde_json::json!({}),
+        secret_fields: Vec::new(),
+        api_key: None,
+    }
+}
+
+#[derive(Clone, Default)]
+pub struct Log(Arc<Mutex<Vec<String>>>);
+
+impl Log {
+    pub fn push(&self, entry: impl Into<String>) {
+        self.0.lock().unwrap().push(entry.into());
+    }
+
+    pub fn entries(&self) -> Vec<String> {
+        self.0.lock().unwrap().clone()
+    }
+}
+
+#[derive(Clone, Copy, Debug)]
+pub enum OpScript {
+    Answer,
+    RaisePython,
+    RejectNatively,
+    RejectRequestNatively,
+    RaiseRequestPython,
+    CancelRequest,
+    AwaitAnswer,
+    AwaitFailure,
+}
+
+/// The fake route's public exception, kept as a value so a test sees what `classify`
+/// produced before the driver raises it.
+#[derive(Debug, PartialEq, Eq)]
+pub struct Classified(pub String);
+
+impl From<Classified> for PyErr {
+    fn from(classified: Classified) -> Self {
+        PyValueError::new_err(format!("classified: {}", classified.0))
+    }
+}
+
+pub struct SyntheticBinding {
+    log: Log,
+    op: OpScript,
+    classifier_fails: bool,
+    pending_reply: Option<Reply<String>>,
+}
+
+pub fn binding(log: &Log, op: OpScript) -> SyntheticBinding {
+    SyntheticBinding {
+        log: log.clone(),
+        op,
+        classifier_fails: false,
+        pending_reply: None,
+    }
+}
+
+impl SyntheticBinding {
+    pub fn with_failing_classifier(self) -> Self {
+        Self {
+            classifier_fails: true,
+            ..self
+        }
+    }
+
+    fn answer(&self, value: impl FnOnce() -> String) -> Result<String, InvokeError<Error>> {
+        match self.op {
+            OpScript::Answer | OpScript::AwaitAnswer | OpScript::AwaitFailure => Ok(value()),
+            OpScript::RaisePython | OpScript::RaiseRequestPython => {
+                Err(PyValueError::new_err("op failed").into())
+            }
+            OpScript::RejectNatively | OpScript::RejectRequestNatively => {
+                Err(InvokeError::Native(Error("op rejected".into())))
+            }
+            OpScript::CancelRequest => Err(CancelledError::new_err(()).into()),
+        }
+    }
+}
+
+impl PythonBinding for SyntheticBinding {
+    type Protocol = Synthetic;
+    type Failure = Classified;
+
+    fn decode_request(
+        &mut self,
+        _: Python<'_>,
+        arguments: &Bound<'_, PyDict>,
+    ) -> Result<String, InvokeError<Error>> {
+        self.log.push("project");
+        match self.op {
+            OpScript::RejectRequestNatively
+            | OpScript::RaiseRequestPython
+            | OpScript::CancelRequest => self.answer(String::new),
+            _ => Ok(format!("project:{}", arguments.len())),
+        }
+    }
+
+    fn encode_stream_head(
+        &mut self,
+        _: Python<'_>,
+        head: std::convert::Infallible,
+    ) -> PyResult<Py<PyAny>> {
+        match head {}
+    }
+
+    fn encode_chunk(
+        &mut self,
+        _: Python<'_>,
+        chunk: std::convert::Infallible,
+    ) -> PyResult<Py<PyAny>> {
+        match chunk {}
+    }
+
+    fn encode_response(&mut self, py: Python<'_>, response: String) -> PyResult<Py<PyAny>> {
+        self.log.push("complete");
+        Ok(PyString::new(py, &response).into_any().unbind())
+    }
+
+    fn map_error(&self, _: Python<'_>, error: Error) -> PyResult<Classified> {
+        self.log.push(format!("classify:{error}"));
+        if self.classifier_fails {
+            return Err(PyTypeError::new_err("classifier failed"));
+        }
+        Ok(Classified(error.0))
+    }
+
+    fn host_error(error: &PyErr) -> Error {
+        Error(error.to_string())
+    }
+}
+
+impl PythonHostCalls<Synthetic> for SyntheticBinding {
+    fn handle_host_call(
+        &mut self,
+        _: Python<'_>,
+        (op, reply): (&'static str, Reply<String>),
+    ) -> Result<(), InvokeError<Error>> {
+        self.log.push(format!("op:{op}"));
+        self.answer(|| op.to_string())
+            .map(|answer| reply.send(answer))
+    }
+
+    fn begin_host_call(
+        &mut self,
+        py: Python<'_>,
+        (op, reply): (&'static str, Reply<String>),
+    ) -> Result<Option<Py<PyAny>>, InvokeError<Error>> {
+        if !matches!(self.op, OpScript::AwaitAnswer | OpScript::AwaitFailure) {
+            return self.handle_host_call(py, (op, reply)).map(|()| None);
+        }
+        self.pending_reply = Some(reply);
+        let module = PyModule::from_code(
+            py,
+            pyo3::ffi::c_str!(
+                "async def answer(fail):\n    if fail:\n        raise LookupError('async host failed')\n    return 'awaited'\n"
+            ),
+            pyo3::ffi::c_str!("host_op.py"),
+            pyo3::ffi::c_str!("host_op"),
+        )?;
+        Ok(Some(
+            module
+                .getattr("answer")?
+                .call1((matches!(self.op, OpScript::AwaitFailure),))?
+                .unbind(),
+        ))
+    }
+
+    fn resume_host_call(
+        &mut self,
+        py: Python<'_>,
+        result: PyResult<Py<PyAny>>,
+    ) -> Result<Option<Py<PyAny>>, InvokeError<Error>> {
+        let answer = result?.extract::<String>(py)?;
+        self.pending_reply.take().unwrap().send(answer);
+        Ok(None)
+    }
+}
+
+impl PythonOwned for SyntheticBinding {
+    fn close(&mut self, _: Python<'_>) {
+        self.log.push("host.close");
+    }
+
+    fn traverse(&self, _: &PyVisit<'_>) -> Result<(), PyTraverseError> {
+        Ok(())
+    }
+}
+
+#[derive(Clone, Copy, Debug)]
+pub enum HookScript {
+    Plain,
+    RewriteArguments,
+    ObserveArguments,
+    FailBegin,
+    ReplaceResponse,
+    FailAfterSuccess,
+    CancelTerminal,
+    FailTerminal,
+}
+
+pub struct SyntheticHooks {
+    log: Log,
+    script: HookScript,
+}
+
+pub fn hooks(log: &Log, script: HookScript) -> SyntheticHooks {
+    SyntheticHooks {
+        log: log.clone(),
+        script,
+    }
+}
+
+impl CallHooks<PythonRuntime> for SyntheticHooks {
+    fn prepare_arguments(
+        &mut self,
+        _: Python<'_>,
+        arguments: Py<PyDict>,
+        _: f64,
+    ) -> PyResult<HookStep<Self, Py<PyDict>>> {
+        self.log.push("begin");
+        if matches!(self.script, HookScript::FailBegin) {
+            return Err(PyValueError::new_err("begin failed"));
+        }
+        if matches!(self.script, HookScript::RewriteArguments) {
+            let prepared = Python::attach(|py| -> PyResult<Py<PyDict>> {
+                let prepared = arguments.bind(py).copy()?;
+                prepared.set_item("prepared", "hook")?;
+                prepared.set_item("api_key", "hook-key")?;
+                Ok(prepared.unbind())
+            })?;
+            return Ok(HookStep::Ready(prepared));
+        }
+        Ok(HookStep::Ready(arguments))
+    }
+
+    fn arguments_prepared(&mut self, py: Python<'_>, arguments: &Py<PyDict>) -> PyResult<()> {
+        if matches!(self.script, HookScript::ObserveArguments) {
+            let value: String = arguments
+                .bind(py)
+                .get_item("prepared")?
+                .unwrap()
+                .extract()?;
+            self.log.push(format!("adopted:{value}"));
+        }
+        Ok(())
+    }
+
+    fn before_provider_request(
+        &mut self,
+        _: Python<'_>,
+        wire: Box<WireRequest>,
+        _: &RequestContext,
+    ) -> PyResult<HookStep<Self, Box<WireRequest>>> {
+        self.log.push("before_provider_request");
+        Ok(HookStep::Ready(Box::new(WireRequest {
+            url: "rewritten".into(),
+            ..*wire
+        })))
+    }
+
+    fn transform_response(
+        &mut self,
+        py: Python<'_>,
+        response: Py<PyAny>,
+        _: Timing,
+    ) -> PyResult<HookStep<Self, Py<PyAny>>> {
+        self.log.push("after_success");
+        match self.script {
+            HookScript::ReplaceResponse => Ok(HookStep::Ready(
+                "replaced".into_pyobject(py)?.into_any().unbind(),
+            )),
+            HookScript::FailAfterSuccess => Err(PyValueError::new_err("after_success failed")),
+            _ => Ok(HookStep::Ready(response)),
+        }
+    }
+
+    fn on_event(
+        &mut self,
+        py: Python<'_>,
+        event: PythonCallEvent<'_>,
+    ) -> PyResult<HookStep<Self, ()>> {
+        let terminal = matches!(
+            event,
+            PythonCallEvent::Succeeded { .. } | PythonCallEvent::Failed { .. }
+        );
+        if terminal && matches!(self.script, HookScript::CancelTerminal) {
+            return Err(CancelledError::new_err("callback cancelled"));
+        }
+        if terminal && matches!(self.script, HookScript::FailTerminal) {
+            return Err(PyValueError::new_err("callback failed"));
+        }
+        self.log.push(match event {
+            PythonCallEvent::Execution(ExecutionEvent::ResultReady { .. }) => "cache_hit".into(),
+            PythonCallEvent::Started { .. } => "started".into(),
+            PythonCallEvent::Cancelled { .. } => "cancelled".into(),
+            PythonCallEvent::Execution(ExecutionEvent::ProviderResponseReceived { raw }) => {
+                format!("response:{}", raw.body)
+            }
+            PythonCallEvent::Succeeded { response, .. } => {
+                format!("succeeded:{}", response.bind(py))
+            }
+            PythonCallEvent::Failed { origin, error, .. } => {
+                format!("failed:{origin:?}:{}", error.value(py))
+            }
+        });
+        Ok(HookStep::Ready(()))
+    }
+
+    fn on_stream_open(&mut self, _: Python<'_>, _: &Py<PyAny>) -> PyResult<()> {
+        self.log.push("opened");
+        Ok(())
+    }
+
+    fn on_stream_chunk(&mut self, _: Python<'_>, _: &Py<PyAny>) -> PyResult<()> {
+        self.log.push("delivered");
+        Ok(())
+    }
+}
+
+impl PythonOwned for SyntheticHooks {
+    fn close(&mut self, _: Python<'_>) {
+        self.log.push("adapter.close");
+    }
+
+    fn traverse(&self, _: &PyVisit<'_>) -> Result<(), PyTraverseError> {
+        Ok(())
+    }
+}
+
+pub type MachineFactory = Box<dyn FnOnce(String) -> CallMachine<Synthetic> + Send + Sync>;
+
+/// Answers to projection, to the route op and to `before_provider_request` all reach the
+/// response, so a driver that misroutes a reply changes what the call returns.
+pub fn success_machine() -> MachineFactory {
+    Box::new(move |projected| {
+        CallMachine::new(None, move |host| {
+            Box::pin(async move {
+                let signed = host.services.call(|reply| ("sign", reply)).await?;
+                let wire = host
+                    .interceptors
+                    .before_provider_request(wire(), context())
+                    .await?;
+                host.interceptors
+                    .after_provider_response(RawResponse { body: "raw".into() })
+                    .await?;
+                Ok(format!("{projected}|{signed}|{}", wire.url))
+            })
+        })
+    })
+}
+
+pub fn failing_machine() -> MachineFactory {
+    Box::new(move |_| {
+        CallMachine::new(None, |_| {
+            Box::pin(async move { Err(Error("provider exploded".into())) })
+        })
+    })
+}
+
+/// Drives an asynchronous call's coroutine to its settled value from a sync test. A
+/// synchronous call is already settled.
+pub fn settled(
+    py: Python<'_>,
+    result: PyResult<Py<PyAny>>,
+    asynchronous: bool,
+) -> PyResult<Py<PyAny>> {
+    if !asynchronous {
+        return result;
+    }
+    let coroutine = result?;
+    let completed = coroutine
+        .call_method1(py, "send", (py.None(),))
+        .unwrap_err();
+    if !completed.is_instance_of::<pyo3::exceptions::PyStopIteration>(py) {
+        return Err(completed);
+    }
+    completed.value(py).getattr("value").map(Bound::unbind)
+}
+
+pub fn run_scripted(
+    py: Python<'_>,
+    machine: MachineFactory,
+    op: OpScript,
+    script: HookScript,
+    asynchronous: bool,
+) -> (PyResult<Py<PyAny>>, Vec<String>) {
+    let log = Log::default();
+    run_composed(
+        py,
+        machine,
+        binding(&log, op),
+        hooks(&log, script),
+        std::convert::identity,
+        call_options(asynchronous),
+    )
+}
+
+pub fn run_composed<L: PythonCallHooks + 'static>(
+    py: Python<'_>,
+    machine: MachineFactory,
+    binding: SyntheticBinding,
+    hooks: SyntheticHooks,
+    compose: impl FnOnce(SyntheticHooks) -> L,
+    options: CallOptions,
+) -> (PyResult<Py<PyAny>>, Vec<String>) {
+    let asynchronous = options.asynchronous;
+    let log = binding.log.clone();
+    let arguments = PyDict::new(py);
+    arguments.set_item("model", "m").unwrap();
+    let result = run_call(
+        py,
+        move |_, _, request| Ok(machine(request)),
+        binding,
+        compose(hooks),
+        arguments.unbind(),
+        options,
+    );
+    (settled(py, result, asynchronous), log.entries())
+}
+
+pub struct Streaming;
+
+impl Protocol for Streaming {
+    type Response = ();
+    type Error = Error;
+    type Request = ();
+    type HostCall = std::convert::Infallible;
+    type Chunk = &'static str;
+    type StreamHead = Vec<(&'static str, &'static str)>;
+}
+
+pub struct StreamingBinding;
+
+impl PythonBinding for StreamingBinding {
+    type Protocol = Streaming;
+    type Failure = Classified;
+
+    fn decode_request(
+        &mut self,
+        _: Python<'_>,
+        _: &Bound<'_, PyDict>,
+    ) -> Result<(), InvokeError<Error>> {
+        Ok(())
+    }
+
+    fn encode_stream_head(
+        &mut self,
+        py: Python<'_>,
+        head: Vec<(&'static str, &'static str)>,
+    ) -> PyResult<Py<PyAny>> {
+        let headers = PyDict::new(py);
+        for (name, value) in head {
+            headers.set_item(name, value)?;
+        }
+        let hidden = PyDict::new(py);
+        hidden.set_item("additional_headers", headers)?;
+        Ok(hidden.into_any().unbind())
+    }
+
+    fn encode_chunk(&mut self, py: Python<'_>, chunk: &'static str) -> PyResult<Py<PyAny>> {
+        Ok(PyString::new(py, chunk).into_any().unbind())
+    }
+
+    fn encode_response(&mut self, py: Python<'_>, (): ()) -> PyResult<Py<PyAny>> {
+        Ok(py.None())
+    }
+
+    fn map_error(&self, _: Python<'_>, error: Error) -> PyResult<Classified> {
+        Ok(Classified(error.0))
+    }
+
+    fn host_error(error: &PyErr) -> Error {
+        Error(error.to_string())
+    }
+}
+
+impl PythonHostCalls<Streaming> for StreamingBinding {
+    fn handle_host_call(
+        &mut self,
+        _: Python<'_>,
+        op: std::convert::Infallible,
+    ) -> Result<(), InvokeError<Error>> {
+        match op {}
+    }
+}
+
+impl PythonOwned for StreamingBinding {
+    fn close(&mut self, _: Python<'_>) {}
+
+    fn traverse(&self, _: &PyVisit<'_>) -> Result<(), PyTraverseError> {
+        Ok(())
+    }
+}
+
+pub fn streaming_machine() -> impl FnOnce(()) -> HostedMachine<Streaming> + Send + Sync {
+    move |()| {
+        hosted_call((), None, |(), _, _, _observations| async {
+            Ok(CallOutput::Stream {
+                head: vec![("request-id", "req_1")],
+                chunks: Box::pin(futures_util::stream::iter([Ok("first"), Ok("second")])),
+            })
+        })
+    }
+}
+
+/// Starts a streaming call and returns the caller's stream object.
+pub fn open_stream(
+    py: Python<'_>,
+    hooks: SyntheticHooks,
+    options: CallOptions,
+) -> Bound<'_, PyAny> {
+    let asynchronous = options.asynchronous;
+    let handed = run_call(
+        py,
+        |_, _, request| Ok(streaming_machine()(request)),
+        StreamingBinding,
+        hooks,
+        PyDict::new(py).unbind(),
+        options,
+    )
+    .unwrap();
+    if asynchronous {
+        let stop = handed.call_method1(py, "send", (py.None(),)).unwrap_err();
+        stop.value(py).getattr("value").unwrap()
+    } else {
+        handed.into_bound(py)
+    }
+}
+
+/// Drives a `Stream` (async) or `SyncStream` to completion from a sync test.
+pub fn read_all(py: Python<'_>, stream: &Bound<'_, PyAny>, asynchronous: bool) -> Vec<String> {
+    if !asynchronous {
+        return stream
+            .try_iter()
+            .unwrap()
+            .map(|chunk| chunk.unwrap().extract().unwrap())
+            .collect();
+    }
+    std::iter::from_fn(|| {
+        let stop = stream
+            .call_method0("__anext__")
+            .unwrap()
+            .call_method1("send", (py.None(),))
+            .unwrap_err();
+        if stop.is_instance_of::<pyo3::exceptions::PyStopAsyncIteration>(py) {
+            return None;
+        }
+        assert!(stop.is_instance_of::<pyo3::exceptions::PyStopIteration>(py));
+        Some(stop.value(py).getattr("value").unwrap().extract().unwrap())
+    })
+    .collect()
+}