diff --git a/crates/libsy/src/algorithms/rand.rs b/crates/libsy/src/algorithms/rand.rs index a8017cad7..5a0c3243c 100644 --- a/crates/libsy/src/algorithms/rand.rs +++ b/crates/libsy/src/algorithms/rand.rs @@ -179,7 +179,6 @@ mod tests { use switchyard_protocol::{Metadata, completion_text, text_request}; - use crate::DriverError; use crate::algorithms::util::affinity::AffinityRouter; use crate::core::algorithm::LlmTarget; use crate::core::testing::{echo, test_drive}; @@ -445,10 +444,8 @@ mod tests { let concrete = decision .as_any() .downcast_ref::() - .ok_or_else(|| { - LibsyError::from(DriverError::TypeMismatch { - expected: "RandomDecision", - }) + .ok_or_else(|| LibsyError::AlgorithmError { + message: "decision did not downcast to RandomDecision".to_string(), })?; assert_eq!(concrete.selected_model, "only/model"); Ok(()) diff --git a/crates/libsy/src/algorithms/util/llm_judge.rs b/crates/libsy/src/algorithms/util/llm_judge.rs index 2f281187e..dad60c31d 100644 --- a/crates/libsy/src/algorithms/util/llm_judge.rs +++ b/crates/libsy/src/algorithms/util/llm_judge.rs @@ -481,11 +481,10 @@ mod tests { }) } - /// Serves the single offloaded judge call with `reply`. The stream is taken first - /// because the driver refuses to publish a step until a consumer exists. + /// Serves the single offloaded judge call with `reply` through a standalone step receiver. async fn score_served_with(reply: Result) -> Result { - let driver = Driver::new(); - let mut steps = Box::pin(driver.stream()); + let (driver, step_rx) = Driver::new(); + let mut steps = tokio_stream::wrappers::ReceiverStream::new(step_rx); let classifier = classifier(); let mut state = State::default(); let mut request = request(); diff --git a/crates/libsy/src/core.rs b/crates/libsy/src/core.rs index e922faf13..b22ccc2ee 100644 --- a/crates/libsy/src/core.rs +++ b/crates/libsy/src/core.rs @@ -1,11 +1,10 @@ // SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. // SPDX-License-Identifier: Apache-2.0 -//! Core orchestration: the [`Algorithm`](algorithm::Algorithm) trait and its -//! [`Driver`](algorithm::Driver), built on the type-erased promise-over-a-stream -//! pump in [`driver`]. Algorithm implementations live in [`crate::algorithms`]. +//! Core orchestration: the [`Algorithm`](algorithm::Algorithm) trait and its typed +//! step-stream [`Driver`](algorithm::Driver). Algorithm implementations live in +//! [`crate::algorithms`]. -mod driver; #[cfg(test)] pub(crate) mod testing; diff --git a/crates/libsy/src/core/algorithm.rs b/crates/libsy/src/core/algorithm.rs index 463ebd647..f888aad54 100644 --- a/crates/libsy/src/core/algorithm.rs +++ b/crates/libsy/src/core/algorithm.rs @@ -16,6 +16,8 @@ use std::{ use async_trait::async_trait; use futures::{Stream, StreamExt}; use parking_lot::Mutex; +use tokio::sync::{mpsc, oneshot}; +use tokio_stream::wrappers::ReceiverStream; use tracing::Instrument; /// The request/response protocol types come from [`switchyard_protocol`]. @@ -29,7 +31,6 @@ use switchyard_protocol::{ Context, Decision, LlmClientError, Request, Response, RoutingFallbackReason, Signals, }; -use super::driver::{DriverRequest, DriverStep, TypeErasedDriver}; use crate::{DriverError, LibsyError, Result, observability}; /// A boxed, `Send` stream of [`Step`]s — the output of @@ -58,30 +59,17 @@ pub struct RoutedRequest { /// The host-facing half of an offloaded model call, surfaced inside [`Step::CallLlm`]. /// -/// Wraps a `DriverRequest` whose payload is a [`RoutedRequest`]. The host reads the -/// routed request ([`get_routed`](Self::get_routed)) and the decision behind it -/// ([`get_decision`](Self::get_decision)), performs (or delegates) the model call, and -/// fulfills it with [`respond`](Self::respond) — unblocking the algorithm's -/// [`Driver::call_llm`] on the other side. `switchyard-llm-client`'s `run` is the -/// ready-made consumer that does this for you. +/// The host reads the routed request ([`get_routed`](Self::get_routed)) and the decision +/// behind it ([`get_decision`](Self::get_decision)), performs (or delegates) the model +/// call, and fulfills it with [`respond`](Self::respond) — unblocking the algorithm's +/// [`Driver::call_llm`] on the other side. `switchyard-llm-client`'s `run` is the ready-made +/// consumer that does this for you. pub struct CallLlmRequest { - inner: DriverRequest, routed: RoutedRequest, + reply: oneshot::Sender>, } impl CallLlmRequest { - /// Wrap a driver request whose payload is a [`RoutedRequest`]. Caches an owned copy - /// so the accessors are plain field reads. - fn new(inner: DriverRequest) -> Self { - // The payload is always a `RoutedRequest` (set by `Driver::call_llm`); a - // mismatch would be a libsy bug, not a runtime condition. - let routed = match inner.request::() { - Ok(routed) => routed.clone(), - Err(_) => unreachable!("CallLlmRequest payload is always a RoutedRequest"), - }; - Self { inner, routed } - } - /// The routed request the host should serve; its `decision.selected_model()` names /// the model to hit. pub fn get_routed(&self) -> &RoutedRequest { @@ -102,7 +90,9 @@ impl CallLlmRequest { /// propagate a failed model call back to the algorithm. Consumes the promise: it /// can only be fulfilled once. pub fn respond(self, result: Result) -> Result<()> { - self.inner.respond::(result) + self.reply + .send(result) + .map_err(|_| DriverError::ResponseDropped.into()) } } @@ -114,16 +104,19 @@ impl CallLlmRequest { /// the algorithm one step at a time. #[derive(Clone)] pub struct Driver { - driver: TypeErasedDriver, + step_tx: mpsc::Sender>, } impl Driver { /// Build an empty driver with its step channel ready. Created per call by - /// [`run_stream`](Algorithm::run_stream). - pub(crate) fn new() -> Self { - Self { - driver: TypeErasedDriver::new(), - } + /// [`run_stream`](Algorithm::run_stream). Also returns the Step receiver. + pub(crate) fn new() -> (Self, mpsc::Receiver>) { + // Capacity one keeps the algorithm paced by the stream consumer. It limits queued steps, + // not model calls already pulled from the stream, which can still run at the same time. + // A larger buffer would use more memory and let the algorithm run farther ahead with + // little benefit because reading a step is cheap compared with serving a model call. + let (step_tx, step_rx) = mpsc::channel(1); + (Self { step_tx }, step_rx) } /// Offload a model call: publish `routed` as a [`Step::CallLlm`] and await the @@ -156,10 +149,18 @@ impl Driver { let tier = routed.decision.routing_tier().map(str::to_string); let is_routed = routed.decision.is_routed_call(); let started = Instant::now(); - let result = self - .driver - .fulfill_request::(routed.ctx.clone(), routed) - .await; + let (reply, response) = oneshot::channel::>(); + let call = CallLlmRequest { routed, reply }; + let result = async { + self.step_tx + .send(Ok(Step::CallLlm(Box::new(call)))) + .await + .map_err(|_| DriverError::StreamClosed)?; + response + .await + .map_err(|_| LibsyError::from(DriverError::ResponseDropped))? + } + .await; let elapsed = started.elapsed(); observability::record_llm_call( &algorithm, @@ -177,7 +178,10 @@ impl Driver { /// Each successfully published decision is counted and logged with its /// reasoning; a decision the stream never accepted is not recorded. pub async fn info(&self, ctx: Context, decision: Arc) -> Result<()> { - self.driver.info(ctx.clone(), decision.clone()).await?; + self.step_tx + .send(Ok(Step::Decision(decision.clone()))) + .await + .map_err(|_| DriverError::StreamClosed)?; observability::record_decision(&ctx, decision.as_ref()); Ok(()) } @@ -185,48 +189,16 @@ impl Driver { /// Emit the terminal step: [`Step::ReturnToAgent`] on `Ok`, or an `Err` stream /// item on failure. Internal: called once by [`run_stream`](Algorithm::run_stream) /// when the algorithm finishes. - pub(crate) async fn finish(&self, ctx: Context, result: Result) -> Result<()> { - match result { - Ok(response) => self.driver.done(ctx, response).await, - Err(err) => self.driver.fail(ctx, err).await, - } - } - - /// Transform the raw driver stream into a stream of [`Step`]s. Internal: the - /// consumer stream is taken (once) by [`run_stream`](Algorithm::run_stream). A - /// payload that does not match the expected type for its step becomes an `Err` item. - pub(crate) fn stream(&self) -> impl Stream> + use<> { - self.driver.stream().map(|item| match item? { - DriverStep::Request(req) => Ok(Step::CallLlm(Box::new(CallLlmRequest::new(req)))), - DriverStep::Info(payload) => payload - .downcast::>() - .map(|decision| Step::Decision(*decision)) - .map_err(|_| { - DriverError::TypeMismatch { - expected: "Arc", - } - .into() - }), - DriverStep::Done(payload) => payload - .downcast::() - .map(Step::ReturnToAgent) - .map_err(|_| { - DriverError::TypeMismatch { - expected: "Response", - } - .into() - }), - }) - } -} - -impl Default for Driver { - fn default() -> Self { - Self::new() + pub(crate) async fn finish(&self, result: Result) -> Result<()> { + let step = result.map(|response| Step::ReturnToAgent(Box::new(response))); + self.step_tx + .send(step) + .await + .map_err(|_| DriverError::StreamClosed.into()) } } -/// One item in the stream returned by `Driver::stream` / [`Algorithm::run_stream`]. +/// One item in the stream returned by [`Algorithm::run_stream`]. pub enum Step { /// The algorithm needs this model call performed. The host serves it and fulfills /// it with [`CallLlmRequest::respond`]. Boxed: it is by far the largest variant. @@ -609,10 +581,10 @@ pub trait Algorithm: Send + Sync + 'static { observability::ALGORITHM_KEY.to_string(), self.name().to_string(), ); - let driver = Driver::new(); + let (driver, step_rx) = Driver::new(); let task_driver = driver.clone(); let task_ctx = ctx.clone(); - let stream = task_driver.stream(); + let stream = ReceiverStream::new(step_rx); // One `libsy.run` span covers the whole algorithm task; the driver's // `libsy.llm_call` spans and decision logs nest inside it via `tracing`'s // contextual parenting. @@ -631,14 +603,13 @@ pub trait Algorithm: Send + Sync + 'static { let abort_guard = AbortOnDrop(handle.abort_handle()); let finish_driver = driver.clone(); - let finish_ctx = ctx; let tail: StepStream = Box::pin( futures::stream::once(async move { let result = match handle.await { Ok(response) => response, Err(source) => Err(LibsyError::AlgorithmTask { source }), }; - finish_driver.finish(finish_ctx, result).await + finish_driver.finish(result).await }) .filter_map(|finish_result| async move { finish_result.err().map(Err) }), ); @@ -801,6 +772,98 @@ mod tests { LlmTargetSet::new(targets) } + fn routed(model: &str) -> RoutedRequest { + RoutedRequest { + request: request(), + decision: Arc::new(TestDecision { + model: model.to_string(), + }), + ctx: Context::default(), + } + } + + #[tokio::test] + async fn typed_driver_preserves_call_and_stream_boundaries() -> Result<()> { + tokio::time::timeout(std::time::Duration::from_secs(1), async { + // Distinct oneshots keep reverse-order replies paired with their producers, and a + // retained call remains pending until the host responds. + let (driver, mut step_rx) = Driver::new(); + let first_driver = driver.clone(); + let mut first = + tokio::spawn(async move { first_driver.call_llm(routed("first")).await }); + let second = tokio::spawn(async move { driver.call_llm(routed("second")).await }); + + let mut calls = HashMap::new(); + for _ in 0..2 { + let step = step_rx.recv().await.ok_or(DriverError::StreamClosed)??; + let Step::CallLlm(call) = step else { + return Err(test_error("expected a CallLlm step")); + }; + calls.insert(call.get_decision().selected_model().to_string(), call); + } + assert!( + tokio::time::timeout(std::time::Duration::from_millis(20), &mut first) + .await + .is_err(), + "call completed before the host responded" + ); + calls + .remove("second") + .ok_or_else(|| test_error("missing second call"))? + .respond(Ok(reply("second response")))?; + calls + .remove("first") + .ok_or_else(|| test_error("missing first call"))? + .respond(Ok(reply("first response")))?; + + let first_response = first + .await + .map_err(|source| LibsyError::AlgorithmTask { source })??; + let second_response = second + .await + .map_err(|source| LibsyError::AlgorithmTask { source })??; + assert_eq!( + first_response.llm_response.as_agg().map(completion_text), + Some("first response".to_string()) + ); + assert_eq!( + second_response.llm_response.as_agg().map(completion_text), + Some("second response".to_string()) + ); + + // Dropping the host-facing promise closes only that call's reply channel. + let (driver, mut step_rx) = Driver::new(); + let producer = tokio::spawn(async move { driver.call_llm(routed("dropped")).await }); + let step = step_rx.recv().await.ok_or(DriverError::StreamClosed)??; + let Step::CallLlm(call) = step else { + return Err(test_error("expected a CallLlm step")); + }; + drop(call); + let result = producer + .await + .map_err(|source| LibsyError::AlgorithmTask { source })?; + assert!(matches!( + result, + Err(LibsyError::Driver(DriverError::ResponseDropped)) + )); + + // A standalone driver reports the typed step receiver disappearing at its next send. + let (driver, step_rx) = Driver::new(); + drop(step_rx); + let decision: Arc = Arc::new(TestDecision { + model: "closed".to_string(), + }); + let result = driver.info(Context::default(), decision).await; + assert!(matches!( + result, + Err(LibsyError::Driver(DriverError::StreamClosed)) + )); + Ok(()) + }) + .await + .map_err(|error| LibsyError::external("waiting for typed driver boundaries", error))? + } + #[test] fn target_lookup_returns_the_missing_target() { let error = target_set(&[]).get_target("missing").err(); diff --git a/crates/libsy/src/core/driver.rs b/crates/libsy/src/core/driver.rs deleted file mode 100644 index 62f5e9763..000000000 --- a/crates/libsy/src/core/driver.rs +++ /dev/null @@ -1,599 +0,0 @@ -// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -// SPDX-License-Identifier: Apache-2.0 - -//! # driver — a type-erased promise-over-a-stream request pump -//! -//! [`TypeErasedDriver`] is the generic offload primitive; its sibling -//! [`algorithm`](super::algorithm) builds the libsy-typed `Driver` on top of it. This -//! module has no dependency on the rest of the crate — the coupling is one-directional -//! (`core/algorithm.rs` → `core/driver.rs`). -//! -//! A [`TypeErasedDriver`] lets a *producer* (e.g. a routing algorithm) fulfill -//! arbitrary requests by publishing promises onto a stream that a single *consumer* -//! drains. It is the type-erased generalization of the crate's `run_stream` offload: -//! instead of one fixed request/response shape, a producer calls -//! [`fulfill_request`](TypeErasedDriver::fulfill_request) -//! with *any* `REQ` and awaits *any* `RES`. -//! -//! - [`fulfill_request`](TypeErasedDriver::fulfill_request) enqueues a [`DriverStep::Request`] -//! carrying a [`DriverRequest`], then awaits the consumer's response. -//! - [`info`](TypeErasedDriver::info) pushes a fire-and-forget [`DriverStep::Info`] — no promise -//! to await. -//! - [`done`](TypeErasedDriver::done) emits the terminal [`DriverStep::Done`] with a final payload. -//! - [`stream`](TypeErasedDriver::stream) hands the single consumer the [`Stream`] of steps; for -//! each [`DriverStep::Request`] the consumer downcasts the request, computes a -//! response, and writes it back with [`DriverRequest::respond`]. -//! -//! Payloads are erased to `Box`, so one `TypeErasedDriver` serves any request -//! type; the consumer downcasts to the concrete type it expects. `TypeErasedDriver` is `Clone` -//! (many producer tasks may call it concurrently — multi-producer), while the stream -//! is single-consumer. Each request rides its own `oneshot`, so concurrent -//! `fulfill_request` calls never cross responses. -//! -//! ## Pacing (bounded step channel) -//! -//! The step channel has capacity 1, so a producer cannot publish its next step until -//! the consumer has pulled the previous one. The consumer therefore *paces* the -//! algorithm: it advances one step for each `.next().await`. Every producer method is -//! `async` because publishing a step awaits channel capacity. -//! -//! ## Termination -//! -//! There is no explicit stop method — the consumer terminates by **dropping the -//! stream** (and any [`DriverRequest`] it is holding). The producer's next publish -//! (`fulfill_request`/`info`/`done`/`fail`) then resolves to `Err`, and a producer -//! awaiting a response sees `Err` once the promise it handed out is dropped. Either -//! way the algorithm unwinds cooperatively at its next driver interaction. Because the -//! producer runs on a task the driver does not own, hard cancellation (e.g. mid-compute -//! that never touches the driver) is the caller's concern — abort the producer task. - -use std::{any::Any, sync::Arc}; - -use crate::{DriverError, LibsyError, Result}; -use parking_lot::Mutex; - -use futures::{Stream, StreamExt}; -use switchyard_protocol::Context; -use tokio::sync::{mpsc, oneshot}; -use tokio::time::{Duration, timeout}; -use tokio_stream::wrappers::ReceiverStream; - -type BoxAny = Box; -type StepResult = Result; - -// TODO make request timeout configurable -const FULFILL_REQUEST_TIMEOUT: Duration = Duration::from_mins(10); - -/// One item on the stream returned by [`TypeErasedDriver::stream`]. -pub enum DriverStep { - /// A request awaiting a response. The consumer downcasts it and fulfills the - /// paired promise with [`DriverRequest::respond`]. - Request(DriverRequest), - /// A fire-and-forget payload from a producer; no response is expected. - Info(BoxAny), - /// A producer's terminal result. The consumer treats it as the last meaningful - /// step (the stream itself closes when every [`TypeErasedDriver`] clone drops). - Done(BoxAny), -} - -/// The consumer-facing half of one [`TypeErasedDriver::fulfill_request`] call. -/// -/// Yielded inside [`DriverStep::Request`]. The consumer reads the request via -/// [`request`](Self::request), does whatever work it names, and fulfills the promise -/// with [`respond`](Self::respond) — unblocking the producer's `fulfill_request`. -pub struct DriverRequest { - request: BoxAny, - // Fulfilled exactly once — `respond` consumes `self` to send, so no `Option`. - tx: oneshot::Sender>, -} - -impl DriverRequest { - /// Borrow the request payload as `REQ`. Errors if the producer enqueued a - /// different type than the consumer expected. - pub fn request(&self) -> Result<&REQ> { - self.request.downcast_ref::().ok_or_else(|| { - DriverError::TypeMismatch { - expected: std::any::type_name::(), - } - .into() - }) - } - - /// Fulfill the promise with a typed response, or an `Err` to propagate a failure - /// back to the producer. Consumes `self`: a promise is fulfilled exactly once. - pub fn respond(self, res: Result) -> Result<()> { - // Erase the response so the single stream item type can carry any RES; the - // producer downcasts it back in `fulfill_request`. - let boxed: Result = res.map(|r| Box::new(r) as BoxAny); - self.tx - .send(boxed) - .map_err(|_| DriverError::ResponseDropped.into()) - } -} - -/// Internal shared state: the step channel plus the single, take-once receiver. -struct DriverInner { - // Multi-producer: cloned into every `TypeErasedDriver`, so many tasks can enqueue steps. - // Capacity 1: a producer blocks publishing its next step until the consumer pulls - // the previous one, so the consumer paces the algorithm. - step_tx: mpsc::Sender, - // Single-consumer: taken out (once) by `stream`. `None` after the first take. - step_rx: Mutex>>, -} - -/// A promise-over-a-stream request pump. See the [module docs](self) for the model. -/// -/// Cheap to clone (shares one `Arc`); clone it to hand a producer handle to another -/// task. The consumer calls [`stream`](Self::stream) exactly once to drain steps. -#[derive(Clone)] -pub struct TypeErasedDriver { - inner: Arc, -} - -impl TypeErasedDriver { - /// Build an empty driver with its step channel ready. Take the consumer stream - /// with [`stream`](Self::stream); enqueue work with the other methods. - pub fn new() -> Self { - let (step_tx, step_rx) = mpsc::channel(1); - TypeErasedDriver { - inner: Arc::new(DriverInner { - step_tx, - step_rx: Mutex::new(Some(step_rx)), - }), - } - } - - fn ensure_started(&self) -> Result<()> { - let guard = self.inner.step_rx.lock(); - if guard.is_none() { - Ok(()) - } else { - Err(DriverError::NotStarted.into()) - } - } - - /// Enqueue `req` as a [`DriverStep::Request`], await the consumer's response, and - /// downcast it to `RES`. Errors if the stream is closed, the promise is dropped - /// unfulfilled, the consumer responded with `Err`, or the response was not a `RES`. - pub async fn fulfill_request(&self, _ctx: Context, req: REQ) -> Result - where - REQ: Any + Send + 'static, - RES: Any + Send + 'static, - { - self.ensure_started()?; - - let (tx, rx) = oneshot::channel::>(); - let promise = DriverRequest { - request: Box::new(req), - tx, - }; - self.inner - .step_tx - .send(Ok(DriverStep::Request(promise))) - .await - .map_err(|_| DriverError::StreamClosed)?; - - // Outer error: the promise was dropped without a response. Inner error: the - // consumer fulfilled it with an explicit `Err` — propagate it as-is. - let response = timeout(FULFILL_REQUEST_TIMEOUT, rx) - .await - .map_err(|_| DriverError::ResponseTimedOut { - timeout: FULFILL_REQUEST_TIMEOUT, - })? - .map_err(|_| DriverError::ResponseDropped)??; - response.downcast::().map(|boxed| *boxed).map_err(|_| { - DriverError::TypeMismatch { - expected: std::any::type_name::(), - } - .into() - }) - } - - /// Push a fire-and-forget [`DriverStep::Info`] payload; there is no promise to - /// await for a response. Awaits channel capacity (the consumer pacing the stream) - /// and errors only if the stream is closed. - pub async fn info(&self, _ctx: Context, info: INFO) -> Result<()> - where - INFO: Any + Send + 'static, - { - self.ensure_started()?; - - self.inner - .step_tx - .send(Ok(DriverStep::Info(Box::new(info)))) - .await - .map_err(|_| DriverError::StreamClosed.into()) - } - - /// Emit the terminal [`DriverStep::Done`] with a final payload. Does not close the - /// stream (that happens when every `TypeErasedDriver` clone drops); the consumer treats it - /// as the last meaningful step. Awaits channel capacity and errors only if the - /// stream is closed. - pub async fn done(&self, _ctx: Context, payload: T) -> Result<()> - where - T: Any + Send + 'static, - { - self.ensure_started()?; - - self.inner - .step_tx - .send(Ok(DriverStep::Done(Box::new(payload)))) - .await - .map_err(|_| DriverError::StreamClosed.into()) - } - - /// Terminate the stream with an error item — the producer-side way to surface a - /// failure to the consumer (mirrors how the crate's `run_stream` yields an `Err` - /// step). Awaits channel capacity and errors only if the stream is - /// already closed. - pub async fn fail(&self, _ctx: Context, err: LibsyError) -> Result<()> { - self.ensure_started()?; - - self.inner - .step_tx - .send(Err(err)) - .await - .map_err(|_| DriverError::StreamClosed.into()) - } - - /// Take the single consumer stream of [`DriverStep`]s. Callable once: a second - /// call yields a one-item stream carrying an `Err`, since the receiver is gone. - pub fn stream(&self) -> impl Stream> + use<> { - let receiver = self - .inner - .step_rx - .lock() - .take() - .ok_or_else(|| LibsyError::from(DriverError::StreamAlreadyTaken)); - match receiver { - Ok(rx) => ReceiverStream::new(rx).left_stream(), - Err(error) => futures::stream::once(async move { Err(error) }).right_stream(), - } - } -} - -impl Default for TypeErasedDriver { - fn default() -> Self { - Self::new() - } -} - -#[cfg(test)] -mod tests { - use super::*; - use futures::StreamExt; - - #[derive(Debug, thiserror::Error)] - #[error("{0}")] - struct TestError(&'static str); - - fn test_error(message: &'static str) -> LibsyError { - LibsyError::external("test", TestError(message)) - } - - #[tokio::test] - async fn fulfill_request_round_trips_typed_values() -> Result<()> { - let driver = TypeErasedDriver::new(); - let stream = driver.stream(); - - // Producer asks for a u32 -> String on its own task. - let producer = driver.clone(); - let handle = tokio::spawn(async move { - producer - .fulfill_request::(Context::default(), 7u32) - .await - }); - - tokio::pin!(stream); - match stream.next().await.ok_or(DriverError::StreamClosed)?? { - DriverStep::Request(promise) => { - let req = *promise.request::()?; - assert_eq!(req, 7); - promise.respond::(Ok(format!("got {req}")))?; - } - _ => return Err(test_error("expected a Request step")), - } - - assert_eq!(handle.await??, "got 7"); - Ok(()) - } - - #[tokio::test] - async fn info_pushes_a_typed_payload() -> Result<()> { - let driver = TypeErasedDriver::new(); - let stream = driver.stream(); - driver.info(Context::default(), 42u64).await?; - - tokio::pin!(stream); - match stream.next().await.ok_or(DriverError::StreamClosed)?? { - DriverStep::Info(payload) => { - let value = payload - .downcast::() - .map_err(|_| LibsyError::from(DriverError::TypeMismatch { expected: "u64" }))?; - assert_eq!(*value, 42); - } - _ => return Err(test_error("expected an Info step")), - } - Ok(()) - } - - #[tokio::test] - async fn done_emits_the_terminal_payload() -> Result<()> { - let driver = TypeErasedDriver::new(); - let stream = driver.stream(); - driver - .done(Context::default(), "finished".to_string()) - .await?; - - tokio::pin!(stream); - match stream.next().await.ok_or(DriverError::StreamClosed)?? { - DriverStep::Done(payload) => { - let value = payload.downcast::().map_err(|_| { - LibsyError::from(DriverError::TypeMismatch { expected: "String" }) - })?; - assert_eq!(*value, "finished"); - } - _ => return Err(test_error("expected a Done step")), - } - Ok(()) - } - - #[tokio::test] - async fn respond_error_propagates_to_the_producer() -> Result<()> { - let driver = TypeErasedDriver::new(); - let stream = driver.stream(); - - let producer = driver.clone(); - let handle = tokio::spawn(async move { - producer - .fulfill_request::(Context::default(), 1u32) - .await - }); - - tokio::pin!(stream); - match stream.next().await.ok_or(DriverError::StreamClosed)?? { - DriverStep::Request(promise) => { - promise.respond::(Err(test_error("upstream failed")))?; - } - _ => return Err(test_error("expected a Request step")), - } - - match handle.await? { - Ok(_) => Err(test_error("expected the error to propagate")), - Err(err) => { - assert!(err.to_string().contains("upstream failed")); - Ok(()) - } - } - } - - #[tokio::test] - async fn response_type_mismatch_errors() -> Result<()> { - let driver = TypeErasedDriver::new(); - let stream = driver.stream(); - - // Producer expects a String back. - let producer = driver.clone(); - let handle = tokio::spawn(async move { - producer - .fulfill_request::(Context::default(), 1u32) - .await - }); - - tokio::pin!(stream); - match stream.next().await.ok_or(DriverError::StreamClosed)?? { - DriverStep::Request(promise) => { - // But the consumer responds with a u32. - promise.respond::(Ok(99u32))?; - } - _ => return Err(test_error("expected a Request step")), - } - - match handle.await? { - Ok(_) => Err(test_error("expected a response type mismatch")), - Err(err) => { - assert!(matches!( - err, - LibsyError::Driver(DriverError::TypeMismatch { expected }) - if expected == std::any::type_name::() - )); - Ok(()) - } - } - } - - #[tokio::test] - async fn request_downcast_to_wrong_type_errors() -> Result<()> { - let driver = TypeErasedDriver::new(); - let stream = driver.stream(); - - let producer = driver.clone(); - let handle = tokio::spawn(async move { - producer - .fulfill_request::(Context::default(), 5u32) - .await - }); - - tokio::pin!(stream); - match stream.next().await.ok_or(DriverError::StreamClosed)?? { - DriverStep::Request(promise) => { - assert!(promise.request::().is_err()); - // Unblock the producer so its task can finish. - promise.respond::(Ok(5u32))?; - } - _ => return Err(test_error("expected a Request step")), - } - - assert_eq!(handle.await??, 5); - Ok(()) - } - - #[tokio::test] - async fn closed_stream_errors_on_send() -> Result<()> { - let driver = TypeErasedDriver::new(); - // Drop the consumer stream (and its receiver) before producing anything. - drop(driver.stream()); - - assert!( - driver - .fulfill_request::(Context::default(), 1u32) - .await - .is_err() - ); - assert!(driver.info(Context::default(), 1u32).await.is_err()); - assert!(driver.done(Context::default(), 1u32).await.is_err()); - Ok(()) - } - - #[tokio::test] - async fn promise_dropped_without_response_errors() -> Result<()> { - let driver = TypeErasedDriver::new(); - let stream = driver.stream(); - - let producer = driver.clone(); - let handle = tokio::spawn(async move { - producer - .fulfill_request::(Context::default(), 1u32) - .await - }); - - tokio::pin!(stream); - match stream.next().await.ok_or(DriverError::StreamClosed)?? { - // Drop the promise without responding. - DriverStep::Request(_promise) => {} - _ => return Err(test_error("expected a Request step")), - } - - assert!(handle.await?.is_err()); - Ok(()) - } - - #[tokio::test(flavor = "multi_thread", worker_threads = 4)] - async fn concurrent_producers_do_not_cross_responses() -> Result<()> { - const N: usize = 8; - let driver = TypeErasedDriver::new(); - let stream = driver.stream(); - - // N producers each fulfill their own request concurrently. - let mut handles = Vec::new(); - for i in 0..N { - let producer = driver.clone(); - handles.push(( - i, - tokio::spawn(async move { - producer - .fulfill_request::(Context::default(), i) - .await - }), - )); - } - - // Single consumer responds `req * 10` to each request. - tokio::pin!(stream); - let mut served = 0; - while served < N { - match stream.next().await.ok_or(DriverError::StreamClosed)?? { - DriverStep::Request(promise) => { - let req = *promise.request::()?; - promise.respond::(Ok(req * 10))?; - served += 1; - } - _ => return Err(test_error("expected a Request step")), - } - } - - // Each producer must see exactly its own response, not another's. - for (i, handle) in handles { - assert_eq!(handle.await??, i * 10); - } - Ok(()) - } - - #[tokio::test] - async fn stream_taken_twice_yields_an_error_item() -> Result<()> { - let driver = TypeErasedDriver::new(); - let _first = driver.stream(); - let second = driver.stream(); - - tokio::pin!(second); - match second.next().await.ok_or(DriverError::StreamClosed)? { - Err(err) => { - assert!(err.to_string().contains("already taken")); - Ok(()) - } - Ok(_) => Err(test_error("expected an error item")), - } - } - - #[tokio::test] - async fn fail_surfaces_an_error_item_on_the_stream() -> Result<()> { - let driver = TypeErasedDriver::new(); - let stream = driver.stream(); - driver - .fail(Context::default(), test_error("kaboom")) - .await?; - - tokio::pin!(stream); - match stream.next().await.ok_or(DriverError::StreamClosed)? { - Err(err) => { - assert!(err.to_string().contains("kaboom")); - Ok(()) - } - Ok(_) => Err(test_error("expected an error item")), - } - } - - #[tokio::test] - async fn dropping_the_stream_terminates_the_producer() -> Result<()> { - let driver = TypeErasedDriver::new(); - // Box::pin so the stream is owned here and `drop` actually drops the receiver. - // (`tokio::pin!` would rebind to a `Pin<&mut _>`, making `drop` a no-op.) - let mut stream = Box::pin(driver.stream()); - - // Producer publishes paced steps until the consumer goes away, then reports - // how many it managed to send. - let producer = driver.clone(); - let handle = tokio::spawn(async move { - let mut sent = 0usize; - while producer.info(Context::default(), sent).await.is_ok() { - sent += 1; - } - sent - }); - - // Pace two steps, then terminate by dropping the stream. - for _ in 0..2 { - match stream.next().await.ok_or(DriverError::StreamClosed)?? { - DriverStep::Info(_) => {} - _ => return Err(test_error("expected an Info step")), - } - } - drop(stream); - - // With the consumer gone, the producer's next publish errors and it stops. - let sent = handle.await?; - assert!(sent >= 2, "producer should have published the paced steps"); - Ok(()) - } - - #[test] - fn ensure_started_requires_the_stream_be_taken_first() -> Result<()> { - let driver = TypeErasedDriver::new(); - - // Before the consumer takes the stream, the receiver is still present, so producing - // is refused — with a message that points at the fix. - match driver.ensure_started() { - Ok(()) => { - return Err(test_error( - "expected ensure_started to reject an untaken stream", - )); - } - Err(err) => assert!(matches!(err, LibsyError::Driver(DriverError::NotStarted))), - } - - // Taking the stream claims the receiver (`Some` -> `None`), so the driver is started. - let _stream = driver.stream(); - assert!(driver.ensure_started().is_ok()); - Ok(()) - } -} diff --git a/crates/libsy/src/error.rs b/crates/libsy/src/error.rs index 648245d82..1e21e950f 100644 --- a/crates/libsy/src/error.rs +++ b/crates/libsy/src/error.rs @@ -3,7 +3,7 @@ //! Typed failures surfaced by libsy's orchestration APIs. -use std::{error::Error as StdError, time::Duration}; +use std::error::Error as StdError; use switchyard_protocol::LlmClientError; use thiserror::Error; @@ -32,7 +32,7 @@ pub enum LibsyError { message: String, }, - /// The type-erased offload driver could not complete an operation. + /// The step-stream driver could not complete an operation. #[error(transparent)] Driver(#[from] DriverError), @@ -94,38 +94,16 @@ impl LibsyError { } } -/// Failures in the type-erased promise-over-stream driver. +/// Failures in the step-stream driver. #[derive(Debug, Error, PartialEq, Eq)] pub enum DriverError { - /// A producer operation was attempted before taking the consumer stream. - #[error("driver stream must be taken before calling producer methods")] - NotStarted, - /// The consumer side of the step channel was dropped. #[error("driver stream is closed")] StreamClosed, - /// The single-consumer stream had already been taken. - #[error("driver stream was already taken")] - StreamAlreadyTaken, - /// One side of a response promise was dropped before delivery. #[error("driver response promise was dropped")] ResponseDropped, - - /// A consumer did not fulfill a request before its deadline. - #[error("driver response timed out after {timeout:?}")] - ResponseTimedOut { - /// Maximum time allowed for request fulfillment. - timeout: Duration, - }, - - /// A type-erased payload did not contain the expected concrete type. - #[error("driver payload type mismatch: expected {expected}")] - TypeMismatch { - /// Human-readable expected payload type or role. - expected: &'static str, - }, } #[cfg(test)]