From 1518f669c7866d97d67afd85e5fc7ddc6cd17bd3 Mon Sep 17 00:00:00 2001 From: dnandakumar-nv Date: Thu, 10 Sep 2026 11:48:24 -0400 Subject: [PATCH] feat(python): add native host driver Signed-off-by: dnandakumar-nv --- crates/switchyard-py/src/libsy_bindings.rs | 289 +++++++++++++++++---- switchyard/libsy/__init__.py | 2 + switchyard_rust/libsy.py | 189 +++++++++++++- tests/test_libsy_minimal_bindings.py | 149 +++++++++++ 4 files changed, 572 insertions(+), 57 deletions(-) diff --git a/crates/switchyard-py/src/libsy_bindings.rs b/crates/switchyard-py/src/libsy_bindings.rs index 79adca287..d3f7b75f4 100644 --- a/crates/switchyard-py/src/libsy_bindings.rs +++ b/crates/switchyard-py/src/libsy_bindings.rs @@ -10,6 +10,7 @@ use futures::StreamExt; use http::header::{HeaderName, HeaderValue}; use pyo3::exceptions::{PyBaseException, PyStopAsyncIteration, PyTypeError, PyValueError}; use pyo3::prelude::*; +use pyo3::types::PyList; use serde_json::Value; use switchyard_libsy::{ Algorithm, CallModel, ClassifierContractConfig, ClassifierResponseFormat, ClassifyTrigger, @@ -22,11 +23,25 @@ use switchyard_protocol::{ LlmClientError, LlmResponse, LlmResponseStream, LlmResponseStreamEvent, Metadata, ModelId, Request, Response, }; -use tokio::sync::Mutex; +use tokio::sync::{Mutex, watch}; use crate::errors::{ContextWindowExceededError, py_libsy_error}; use crate::py_serde::{from_python, to_python}; +const RESPONSE_SOURCE_ID: &str = "switchyard.response_source_id"; + +fn request_from_python( + request: &Bound<'_, PyAny>, + headers: Option>, +) -> PyResult { + let headers = headers.as_ref().map(header_map_from_python).transpose()?; + Ok(Request { + llm_request: from_python(request)?, + raw_request: None, + metadata: headers.map(|headers| Metadata::from_headers(&headers)), + }) +} + /// The Python API keeps its `session_affinity` flag, which selects the per-session trigger. fn classify_trigger(session_affinity: bool) -> ClassifyTrigger { if session_affinity { @@ -357,11 +372,17 @@ enum PyLlmResponse { } impl PyLlmResponse { - fn to_core(&self, py: Python<'_>, model: ModelId) -> PyResult { + fn to_core( + &self, + py: Python<'_>, + model: ModelId, + streams: &Py, + ) -> PyResult { match self { Self::Agg { response } => from_python(response.bind(py)).map(LlmResponse::Agg), Self::Stream { stream } => { - python_response_stream(py, stream.clone_ref(py), model).map(LlmResponse::Stream) + python_response_stream(py, stream.clone_ref(py), model, streams) + .map(LlmResponse::Stream) } } } @@ -384,12 +405,29 @@ fn python_client_error(py: Python<'_>, error: PyErr, model: &ModelId) -> LlmClie } } +struct PythonStreamGuard(Py); + +impl Drop for PythonStreamGuard { + fn drop(&mut self) { + Python::attach(|py| { + let _ = self.0.bind(py).call_method0("_release"); + }); + } +} + fn python_response_stream( py: Python<'_>, stream: Py, model: ModelId, + streams: &Py, ) -> PyResult { - let iterator = stream.bind(py).call_method0("__aiter__")?.unbind(); + let iterator = py + .import("switchyard_rust.libsy")? + .getattr("_InputStream")? + .call1((stream,))? + .unbind(); + streams.bind(py).append(iterator.bind(py))?; + let iterator = PythonStreamGuard(iterator); // Rust polls on Tokio, so retain the Python task's event loop and context for every item. let locals = pyo3_async_runtimes::tokio::get_current_locals(py)?; let stream = futures::stream::unfold(Some((iterator, locals, model)), |state| async move { @@ -397,7 +435,7 @@ fn python_response_stream( let next = Python::attach(|py| { pyo3_async_runtimes::into_future_with_locals( &locals, - iterator.bind(py).call_method0("__anext__")?, + iterator.0.bind(py).call_method0("__anext__")?, ) }); match next { @@ -430,16 +468,18 @@ struct PyModelCall { algorithm: String, request: Py, models: Vec, + streams: Py, } impl PyModelCall { - fn new(py: Python<'_>, call: CallModel) -> PyResult { + fn new(py: Python<'_>, call: CallModel, streams: Py) -> PyResult { let request = to_python(py, &call.request.llm_request)?; Ok(Self { algorithm: call.algorithm.clone(), models: call.models.iter().map(ToString::to_string).collect(), inner: Some(call), request, + streams, }) } @@ -471,7 +511,23 @@ impl PyModelCall { } /// Fulfill this call with a normalized aggregate or streamed response. - fn respond(&mut self, py: Python<'_>, response: PyRef<'_, PyLlmResponse>) -> PyResult<()> { + /// + /// `served_model` names the Switchyard model that answered. `source_id` is a + /// nonempty receipt in the host's per-run records. Omitted identifiers stay + /// unknown. These fields do not authorize fallbacks or response reuse. + #[pyo3(signature = (response, *, served_model=None, source_id=None))] + fn respond( + &mut self, + py: Python<'_>, + response: PyRef<'_, PyLlmResponse>, + served_model: Option, + source_id: Option, + ) -> PyResult<()> { + for (name, value) in [("served_model", &served_model), ("source_id", &source_id)] { + if value.as_ref().is_some_and(String::is_empty) { + return Err(PyValueError::new_err(format!("{name} must be nonempty"))); + } + } let model = self .inner .as_ref() @@ -482,16 +538,34 @@ impl PyModelCall { .as_ref() .map(ModelId::new) .ok_or_else(|| py_libsy_error("model call request is missing its selected model"))?; - let llm_response = response.to_core(py, model)?; + let llm_response = response.to_core( + py, + served_model.as_ref().map(ModelId::new).unwrap_or(model), + &self.streams, + )?; let call = self.take()?; - let metadata = call.request.metadata.clone(); + let mut metadata = call.request.metadata.clone().unwrap_or_default(); + metadata.served_model = served_model.map(ModelId::new); + if let Some(extra) = metadata.extra_metadata.as_mut() { + extra.remove(RESPONSE_SOURCE_ID); + } + if let Some(source_id) = source_id { + metadata + .extra_metadata + .get_or_insert_default() + .insert(RESPONSE_SOURCE_ID.to_string(), source_id); + } call.respond(Ok(Response { llm_response, - metadata, + metadata: Some(metadata), })) .map_err(py_libsy_error) } + fn _is_completed(&self) -> bool { + self.inner.is_none() + } + /// Fulfill this call with a Python client failure. fn fail(&mut self, error: &Bound<'_, PyAny>) -> PyResult<()> { if !error.is_instance_of::() { @@ -551,6 +625,12 @@ struct PyRoutingOutcome { selected_model_ids: Vec, request: Py, response: Option>, + /// The Switchyard model that served the response, if explicitly supplied. + #[pyo3(get)] + served_model: Option, + /// The host's opaque response receipt, if explicitly supplied. + #[pyo3(get)] + source_id: Option, #[pyo3(get)] metadata: Option>, } @@ -581,7 +661,9 @@ impl PyRoutingOutcome { /// Async Python iterator over one normalized Rust response stream. #[pyclass(name = "_LlmResponseStream", module = "switchyard.libsy", frozen)] struct PyLlmResponseStream { - inner: Arc>, + inner: Arc>>, + closed: watch::Sender, + streams: Py, } #[pymethods] @@ -592,17 +674,54 @@ impl PyLlmResponseStream { fn __anext__<'py>(&self, py: Python<'py>) -> PyResult> { let stream = Arc::clone(&self.inner); + let mut closed = self.closed.subscribe(); pyo3_async_runtimes::tokio::future_into_py(py, async move { - match stream.lock().await.next().await { + let next = async { + let mut stream = stream.lock().await; + match stream.as_mut() { + Some(stream) => stream.next().await, + None => None, + } + }; + let event = tokio::select! { + biased; + _ = closed.wait_for(|closed| *closed) => None, + event = next => event, + }; + match event { Some(Ok(event)) => Python::attach(|py| to_python(py, &event)), Some(Err(error)) => Err(py_libsy_error(error)), None => Err(PyStopAsyncIteration::new_err(())), } }) } + + fn aclose<'py>(&self, py: Python<'py>) -> PyResult> { + self.closed.send_replace(true); + let stream = Arc::clone(&self.inner); + let streams = self.streams.clone_ref(py); + let locals = pyo3_async_runtimes::tokio::get_current_locals(py)?; + pyo3_async_runtimes::tokio::future_into_py(py, async move { + stream.lock().await.take(); + Python::attach(|py| { + pyo3_async_runtimes::into_future_with_locals( + &locals, + py.import("switchyard_rust.libsy")? + .getattr("_close_streams")? + .call1((streams, false))?, + ) + })? + .await?; + Ok(()) + }) + } } -fn response_to_python(py: Python<'_>, response: LlmResponse) -> PyResult> { +fn response_to_python( + py: Python<'_>, + response: LlmResponse, + streams: Py, +) -> PyResult> { let response = match response { LlmResponse::Agg(response) => PyLlmResponse::Agg { response: to_python(py, &response)?, @@ -611,7 +730,9 @@ fn response_to_python(py: Python<'_>, response: LlmResponse) -> PyResult>, + streams: Py, } #[pymethods] @@ -645,10 +767,11 @@ impl PyRunStream { fn __anext__<'py>(&self, py: Python<'py>) -> PyResult> { let stream = Arc::clone(&self.inner); + let streams = self.streams.clone_ref(py); pyo3_async_runtimes::tokio::future_into_py(py, async move { let step = stream.lock().await.next().await; match step { - Some(Ok(step)) => step_to_python(step), + Some(Ok(step)) => step_to_python(step, streams), Some(Err(error)) => Err(py_libsy_error(error)), None => Err(PyStopAsyncIteration::new_err(())), } @@ -676,18 +799,15 @@ impl PyAlgorithm { request: &Bound<'_, PyAny>, headers: Option>, ) -> PyResult { - let headers = headers.as_ref().map(header_map_from_python).transpose()?; - let request = Request { - llm_request: from_python(request)?, - raw_request: None, - metadata: headers.map(|headers| Metadata::from_headers(&headers)), - }; + let streams = PyList::empty(request.py()).unbind(); + let request = request_from_python(request, headers)?; let stream = { let _guard = pyo3_async_runtimes::tokio::get_runtime().enter(); Arc::clone(&self.inner).run_stream(request) }; Ok(PyRunStream { inner: Arc::new(Mutex::new(stream)), + streams, }) } @@ -696,44 +816,112 @@ impl PyAlgorithm { } } -fn step_to_python(step: RustStep) -> PyResult { +fn step_to_python(step: RustStep, streams: Py) -> PyResult { match step { RustStep::CallModel(call) => Python::attach(|py| { Ok(PyStep::CallModel { - call: Py::new(py, PyModelCall::new(py, *call)?)?, + call: Py::new(py, PyModelCall::new(py, *call, streams)?)?, }) }), - RustStep::Done(outcome) => { - let RoutingOutcome { - selected_model_ids, - request, - response, - metadata, - } = *outcome; - Python::attach(|py| { - Ok(PyStep::Done { - outcome: Py::new( - py, - PyRoutingOutcome { - metadata: metadata - .map(|inner| Py::new(py, PyOutcomeMetadata { inner })) - .transpose()?, - selected_model_ids: selected_model_ids - .iter() - .map(ToString::to_string) - .collect(), - request: to_python(py, &request.llm_request)?, - response: response - .map(|response| response_to_python(py, response.llm_response)) - .transpose()?, - }, - )?, - }) + RustStep::Done(outcome) => Python::attach(|py| { + Ok(PyStep::Done { + outcome: outcome_to_python(py, *outcome, streams)?, }) - } + }), } } +fn outcome_to_python( + py: Python<'_>, + outcome: RoutingOutcome, + streams: Py, +) -> PyResult> { + let source = outcome + .response + .as_ref() + .and_then(|response| response.metadata.as_ref()); + let served_model = source + .and_then(|metadata| metadata.served_model.as_ref()) + .map(ToString::to_string); + let source_id = source + .and_then(|metadata| metadata.extra_metadata.as_ref()) + .and_then(|extra| extra.get(RESPONSE_SOURCE_ID)) + .cloned(); + Py::new( + py, + PyRoutingOutcome { + served_model, + source_id, + selected_model_ids: outcome + .selected_model_ids + .iter() + .map(ToString::to_string) + .collect(), + request: to_python(py, &outcome.request.llm_request)?, + response: outcome + .response + .map(|response| response_to_python(py, response.llm_response, streams)) + .transpose()?, + metadata: outcome + .metadata + .map(|inner| Py::new(py, PyOutcomeMetadata { inner })) + .transpose()?, + }, + ) +} + +/// The Python wrapper owns cancellation and joins callbacks after this future stops. +#[pyfunction] +#[pyo3(signature = (algorithm, request, serve, stop, streams, headers=None))] +fn _drive<'py>( + py: Python<'py>, + algorithm: PyRef<'_, PyAlgorithm>, + request: &Bound<'_, PyAny>, + serve: Py, + stop: &Bound<'_, PyAny>, + streams: Py, + headers: Option>, +) -> PyResult> { + let algorithm = Arc::clone(&algorithm.inner); + let request = request_from_python(request, headers)?; + let locals = pyo3_async_runtimes::tokio::get_current_locals(py)?; + let stop = pyo3_async_runtimes::into_future_with_locals(&locals, stop.call_method0("wait")?)?; + pyo3_async_runtimes::tokio::future_into_py(py, async move { + let run = switchyard_libsy::drive(algorithm, request, |call| { + let future = Python::attach(|py| { + let call = Py::new(py, PyModelCall::new(py, call, streams.clone_ref(py))?)?; + pyo3_async_runtimes::into_future_with_locals( + &locals, + serve.bind(py).call1((call,))?, + ) + }); + async move { + match future { + Ok(future) => future.await.map(|_| ()), + Err(error) => Err(error), + } + .map_err(|error| RustLibsyError::external("Python host callback", error)) + } + }); + let outcome = tokio::select! { + _ = stop => return Ok(None), + outcome = run => outcome, + }; + Python::attach(|py| { + let outcome = outcome.map_err(|error| { + let exception = py_libsy_error(&error); + if let RustLibsyError::External { source, .. } = error + && let Ok(cause) = source.downcast::() + { + exception.set_cause(py, Some(*cause)); + } + exception + })?; + outcome_to_python(py, outcome, streams).map(Some) + }) + }) +} + /// Construct the no-op reference algorithm. #[pyfunction(name = "noop")] fn noop_algorithm() -> PyAlgorithm { @@ -890,6 +1078,7 @@ pub(crate) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> { libsy_module.add_class::()?; libsy_module.add_class::()?; libsy_module.add_class::()?; + libsy_module.add_function(wrap_pyfunction!(_drive, &libsy_module)?)?; libsy_module.add_function(wrap_pyfunction!(noop_algorithm, &libsy_module)?)?; libsy_module.add_function(wrap_pyfunction!(random_algorithm, &libsy_module)?)?; libsy_module.add_function(wrap_pyfunction!(llm_classifier_algorithm, &libsy_module)?)?; diff --git a/switchyard/libsy/__init__.py b/switchyard/libsy/__init__.py index 88a5aa372..fd2cddd22 100644 --- a/switchyard/libsy/__init__.py +++ b/switchyard/libsy/__init__.py @@ -17,6 +17,7 @@ RoutingOutcome, Step, TaskClassifierConfig, + drive, ) from . import algorithms as algorithms @@ -36,4 +37,5 @@ "Step", "TaskClassifierConfig", "algorithms", + "drive", ] diff --git a/switchyard_rust/libsy.py b/switchyard_rust/libsy.py index 9756b1c6b..7dceae0ab 100644 --- a/switchyard_rust/libsy.py +++ b/switchyard_rust/libsy.py @@ -5,7 +5,9 @@ from __future__ import annotations -from collections.abc import Mapping +import asyncio +from collections.abc import Awaitable, Callable, Mapping +from contextvars import copy_context from typing import TYPE_CHECKING, Any from switchyard_rust._native import load_native @@ -25,6 +27,7 @@ "RoutingOutcome", "Step", "TaskClassifierConfig", + "drive", "llm_classifier", "llm_task_classifier", "noop", @@ -35,12 +38,22 @@ if TYPE_CHECKING: from collections.abc import AsyncIterator, Sequence - from typing import ClassVar, Literal, final + from typing import ClassVar, Generic, Literal, TypeVar, final + + _Stream = TypeVar("_Stream", bound=AsyncIterator[Mapping[str, object]]) class LibsyError(RuntimeError): ... class ContextWindowExceededError(RuntimeError): ... + @final + class _LlmResponseStream(AsyncIterator[dict[str, object]]): + """Native outcome stream. Host-supplied iterators need not support close.""" + + def __aiter__(self) -> _LlmResponseStream: ... + async def __anext__(self) -> dict[str, object]: ... + async def aclose(self) -> None: ... + class LlmResponse: """A normalized aggregate response or live normalized event stream.""" @@ -52,11 +65,11 @@ class Agg: def __init__(self, response: Mapping[str, object]) -> None: ... @final - class Stream: + class Stream(Generic[_Stream]): __match_args__: ClassVar[tuple[Literal["stream"]]] = ("stream",) - stream: AsyncIterator[dict[str, object]] + stream: _Stream - def __init__(self, stream: AsyncIterator[Mapping[str, object]]) -> None: ... + def __init__(self, stream: _Stream) -> None: ... @final class CustomClassifierConfig: @@ -108,10 +121,18 @@ def request(self) -> dict[str, object]: ... @property def models(self) -> list[str]: ... - def respond(self, response: LlmResponse.Agg | LlmResponse.Stream) -> None: ... + def respond( + self, + response: LlmResponse.Agg | LlmResponse.Stream[Any], + *, + served_model: str | None = None, + source_id: str | None = None, + ) -> None: ... def fail(self, error: BaseException) -> None: ... + def _is_completed(self) -> bool: ... + @final class OutcomeMetadata: """Read-only outcome identity and optional algorithm evidence.""" @@ -127,6 +148,12 @@ def evidence(self) -> Any | None: ... @final class RoutingOutcome: + @property + def served_model(self) -> str | None: ... + + @property + def source_id(self) -> str | None: ... + @property def metadata(self) -> OutcomeMetadata | None: ... @@ -137,7 +164,7 @@ def selected_model_ids(self) -> list[str]: ... def request(self) -> dict[str, object]: ... @property - def response(self) -> LlmResponse.Agg | LlmResponse.Stream | None: ... + def response(self) -> LlmResponse.Agg | LlmResponse.Stream[_LlmResponseStream] | None: ... class Step: @final @@ -265,6 +292,154 @@ def stage_router( ) -> Algorithm: ... +class _InputStream: + """Keep Python reads alive only while Rust owns their stream.""" + + def __init__(self, stream: Any) -> None: + self.iterator = stream.__aiter__() + self.loop = asyncio.get_running_loop() + self.context = copy_context() + self.read: asyncio.Future[Any] | None = None + self.close_task: asyncio.Task[None] | None = None + self.released = False + + async def __anext__(self) -> Any: + if self.close_task is not None: + raise StopAsyncIteration + self.read = asyncio.ensure_future(self.iterator.__anext__()) + try: + return await self.read + finally: + self.read = None + + def _release(self) -> None: + self.released = True + if not self.loop.is_closed(): + self.loop.call_soon_threadsafe(self._start_close) + + def _start_close(self) -> asyncio.Task[None]: + if self.close_task is None: + self.close_task = self.context.run(self.loop.create_task, self._close()) + return self.close_task + + async def _close(self) -> None: + iterator, self.iterator = self.iterator, None + if self.read is not None: + self.read.cancel() + await asyncio.gather(self.read, return_exceptions=True) + close = getattr(iterator, "aclose", None) + if close is not None: + await close() + + +async def _close_streams(streams: list[_InputStream], close_all: bool) -> None: + results = await asyncio.gather( + *(stream._start_close() for stream in streams if close_all or stream.released), + return_exceptions=True, + ) + for result in results: + if isinstance(result, BaseException): + raise result + + +async def drive( + algorithm: Algorithm, + request: Mapping[str, object], + serve: Callable[[ModelCall], Awaitable[None]], + *, + headers: Mapping[str, str] | None = None, +) -> RoutingOutcome: + """Run the native driver with a host callback for each model call. + + ``serve`` must finish each call with ``respond`` or ``fail`` and return None. + Record host receipts before completing the call. Complete it as the last + action apart from resource cleanup. Callbacks must honor cancellation and + release their resources in finally blocks. Retries and accounting belong + to the host. This function does not make a final-answer call for route-only + outcomes. + + A returned native response stream belongs to the caller. Consume it or + await its ``aclose()``, including when a client disconnects. Host-supplied + input iterators do not need an ``aclose`` method. + """ + native: Any = load_native().libsy + tasks: dict[asyncio.Task[Any], ModelCall] = {} + cancelled: set[asyncio.Task[Any]] = set() + streams: list[_InputStream] = [] + failure: BaseException | None = None + stop = asyncio.Event() + + async def serve_owned(call: ModelCall) -> None: + nonlocal failure + if stop.is_set(): + return + task = asyncio.current_task() + assert task is not None + tasks[task] = call + try: + await serve(call) + if not call._is_completed(): + raise native.LibsyError("serve returned without completing its model call") + except BaseException as error: + if failure is None and not ( + task in cancelled and isinstance(error, asyncio.CancelledError) + ): + failure = error + raise + + run = native._drive( + algorithm, + request, + serve_owned, + stop, + streams, + dict(headers) if headers is not None else None, + ) + outcome: RoutingOutcome | None = None + error: BaseException | None = None + try: + outcome = await asyncio.shield(run) + except BaseException as caught: + error = caught + stop.set() + + async def cleanup() -> None: + await asyncio.gather(run, return_exceptions=True) + for task, call in tasks.items(): + if not task.done() and not call._is_completed(): + cancelled.add(task) + task.cancel() + await asyncio.gather(*tasks, return_exceptions=True) + await _close_streams(streams, error is not None or failure is not None) + + # Repeated caller cancellation must not interrupt provider finally blocks. + joined = asyncio.create_task(cleanup()) + closed_all = False + while True: + try: + await asyncio.shield(joined) + except asyncio.CancelledError as caught: + error = caught + if not joined.done(): + continue + except BaseException as caught: + if error is None: + error = caught + if (error is not None or failure is not None) and not closed_all: + joined = asyncio.create_task(_close_streams(streams, True)) + closed_all = True + else: + break + if isinstance(error, asyncio.CancelledError): + raise error + if failure is not None: + raise native.LibsyError(f"Python host callback failed: {failure}") from failure + if error is not None: + raise error + assert outcome is not None + return outcome + + def __getattr__(name: str) -> object: if name in _EXPORTS: native: Any = load_native() diff --git a/tests/test_libsy_minimal_bindings.py b/tests/test_libsy_minimal_bindings.py index b693d4d87..1e709c41d 100644 --- a/tests/test_libsy_minimal_bindings.py +++ b/tests/test_libsy_minimal_bindings.py @@ -3,6 +3,7 @@ """Tests for the dictionary-based libsy Python API.""" +import asyncio from collections.abc import AsyncIterator from typing import Any from uuid import UUID @@ -13,13 +14,17 @@ Algorithm, ContextWindowExceededError, CustomClassifierConfig, + EscalationClassifierConfig, + LibsyError, LlmClassifierConfig, LlmResponse, + ModelCall, OutcomeMetadata, RoutingOutcome, Step, TaskClassifierConfig, algorithms, + drive, ) @@ -436,3 +441,147 @@ async def call(self, request: dict[str, Any]) -> dict[str, Any]: assert selected_model == "fast" assert response["model"] == "strong" + + +def escalation() -> Algorithm: + return algorithms.llm_classifier( + LlmClassifierConfig.escalation( + "judge", "efficient", "capable", config=EscalationClassifierConfig(confirmations=1) + ) + ) + + +def answer(text: str) -> Any: + return LlmResponse.Agg( + { + "model": "provider-model", + "outputs": [{"role": "assistant", "content": [{"type": "text", "text": text}]}], + } + ) + + +@pytest.mark.parametrize( + "picker,expected", [("efficient_first", "efficient"), ("capable_first", "capable")] +) +async def test_drive_stage_routes_without_calling_host(picker: str, expected: str) -> None: + async def serve(call: ModelCall) -> None: + pytest.fail("Stage must not call the host") + + outcome = await drive( + algorithms.stage_router("capable", "efficient", picker=picker, confidence_threshold=0.5), + request_body(), + serve, + ) + assert outcome.selected_model_ids[0] == expected + assert outcome.response is outcome.served_model is outcome.source_id is None + + +@pytest.mark.parametrize("reject", [False, True]) +async def test_drive_escalation_response_source(reject: bool) -> None: + calls = [] + + async def serve(call: ModelCall) -> None: + model = call.models[0] + calls.append(model) + verdict = ( + '{"escalate":true,"reason":"stuck"}' + if reject + else '{"escalate":false,"reason":"progressing"}' + ) + call.respond( + answer(verdict if model == "judge" else "accepted answer"), + served_model=model + "-deployment", + source_id=model + "-receipt", + ) + + outcome = await drive(escalation(), request_body(), serve) + assert calls == ["efficient", "judge"] + if reject: + assert outcome.selected_model_ids == ["capable", "efficient"] + assert outcome.response is outcome.served_model is outcome.source_id is None + else: + assert outcome.selected_model_ids == ["efficient"] + assert outcome.served_model == "efficient-deployment" + assert outcome.source_id == "efficient-receipt" + assert outcome.response.response["model"] == "provider-model" + assert outcome.response.response["outputs"][0]["content"][0]["text"] == "accepted answer" + + +async def test_drive_callback_failure_and_retained_incomplete_call() -> None: + retained = [] + original = ValueError("host failed") + + async def failing(call: ModelCall) -> None: + raise original + + with pytest.raises(LibsyError) as caught: + await drive(escalation(), request_body(), failing) + assert caught.value.__cause__ is original + + async def incomplete(call: ModelCall) -> None: + retained.append(call) + + with pytest.raises(LibsyError, match="without completing"): + await asyncio.wait_for(drive(escalation(), request_body(), incomplete), 2) + assert len(retained) == 1 + + +@pytest.mark.parametrize("streaming", [False, True]) +async def test_drive_cancellation_joins_owned_cleanup(streaming: bool) -> None: + started, cleaning, finish = asyncio.Event(), asyncio.Event(), asyncio.Event() + cleaned = [] + + async def pending() -> None: + try: + started.set() + await asyncio.Future() + finally: + cleaning.set() + await finish.wait() + cleaned.append(True) + + async def events() -> AsyncIterator[dict[str, object]]: + await pending() + yield {} + + async def serve(call: ModelCall) -> None: + if streaming: + call.respond(LlmResponse.Stream(events())) + else: + await pending() + + task = asyncio.create_task(drive(escalation(), request_body(), serve)) + await asyncio.wait_for(started.wait(), 2) + task.cancel() + await asyncio.wait_for(cleaning.wait(), 2) + task.cancel() + await asyncio.sleep(0) + finish.set() + with pytest.raises(asyncio.CancelledError): + await task + assert cleaned == [True] + + +async def test_drive_accepted_stream_survives_callback_cleanup() -> None: + cleaned = [] + + async def serve(call: ModelCall) -> None: + try: + text = ( + "accepted answer" + if call.models == ["efficient"] + else '{"escalate":false,"reason":"progressing"}' + ) + call.respond(answer(text), source_id=call.models[0]) + finally: + await asyncio.sleep(0.01) + cleaned.append(call.models[0]) + + outcome = await drive(escalation(), {**request_body(), "stream": True}, serve) + assert sorted(cleaned) == ["efficient", "judge"] + assert outcome.source_id == "efficient" + stream = outcome.response.stream + assert await anext(stream) + await stream.aclose() + await stream.aclose() + assert [event async for event in stream] == []