Skip to main content

switchyard_libsy/core/
algorithm.rs

1// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2// SPDX-License-Identifier: Apache-2.0
3
4//! The [`Algorithm`] trait and its [`Driver`] — the orchestration contract every
5//! algorithm implements, and the offload channel it uses to make model calls and
6//! publish [`Decision`]s.
7
8use std::{
9    collections::{HashMap, HashSet},
10    future::Future,
11    panic::AssertUnwindSafe,
12    pin::Pin,
13    sync::Arc,
14    time::Instant,
15};
16
17use async_trait::async_trait;
18use futures::{FutureExt, Stream, StreamExt};
19use parking_lot::Mutex;
20use tokio::sync::{mpsc, oneshot};
21use tokio_stream::wrappers::ReceiverStream;
22use tracing::Instrument;
23
24/// The request/response protocol types come from [`switchyard_protocol`].
25/// [`switchyard_protocol::LlmRequest`] is the normalized request;
26/// [`switchyard_protocol::AggLlmResponse`] is the buffered response;
27/// [`switchyard_protocol::LlmResponseChunk`] is normalized streaming content;
28/// [`switchyard_protocol::LlmResponseStreamEvent`] is its host/algorithm envelope; and
29/// [`switchyard_protocol::LlmResponse`] carries either a live
30/// [`switchyard_protocol::LlmResponseStream`] or the terminal aggregate.
31use switchyard_protocol::{Decision, LlmClientError, Request, Response, RoutingFallbackReason};
32
33use crate::{DriverError, LibsyError, Result, observability};
34
35/// A boxed, `Send` stream of [`Step`]s — the output of
36/// [`Algorithm::run_stream`]. Boxed so the trait method that produces it keeps
37/// `Arc<dyn Algorithm>` object-safe.
38pub type StepStream = Pin<Box<dyn Stream<Item = Result<Step>> + Send>>;
39
40/// An offloaded model call, surfaced inside [`Step::CallModel`].
41///
42/// The host reads the public fields, performs (or delegates) the model call, and fulfills it
43/// with [`respond`](Self::respond) — unblocking the algorithm's [`Driver::call_model`] on the
44/// other side. `switchyard-llm-client`'s `run` is the ready-made consumer that does this for
45/// you. A host that only wants the routing outcome can take the contents with
46/// [`into_parts`](Self::into_parts) and never respond; dropping the stream ends the run.
47///
48/// The selected model and inbound route name live in separate, unambiguous places: the model
49/// identifier is [`decision.selected_model_id()`](Decision::selected_model_id), while
50/// `request.llm_request.model` is the *inbound* name the agent asked for (libsy
51/// never overwrites it). A client maps `selected_model_id()` to the provider model
52/// id it hits.
53pub struct CallModel {
54    /// The name of the algorithm that produced this call, so a host instrumenting the
55    /// calls it serves can attribute its own spans to the algorithm behind them.
56    pub algorithm: String,
57    /// The request to serve; its `model` is the agent's original name NOT the selected model.
58    /// The caller making the request needs to change it to decision.selected_model_id() before
59    /// sending.
60    pub request: Request,
61    /// The routing decision behind this call; `selected_model_id()` identifies the model to use.
62    pub decision: Decision,
63    // How to send the response back to the algorithm
64    reply: oneshot::Sender<Result<Response>>,
65}
66
67impl CallModel {
68    /// Fulfill the promise with the caller's model-call result. Pass `Err(..)` to
69    /// propagate a failed model call back to the algorithm. Consumes the promise: it
70    /// can only be fulfilled once.
71    pub fn respond(self, result: Result<Response>) -> Result<()> {
72        self.reply
73            .send(result)
74            .map_err(|_| DriverError::ResponseDropped.into())
75    }
76
77    /// Take the call's contents without answering it, dropping the promise — the routing
78    /// outcome plus the request as the algorithm would have sent it, after any rewriting.
79    ///
80    /// Should only be called if `decision.is_answer_call` is true as that is the final call.
81    /// The algorithm's [`Driver::call_model`] will fail with [`DriverError::Abandoned`] and
82    /// the run ends there. Taking a call the algorithm does not depend on (a judge or
83    /// classifier call) may instead let it fail open and complete with degraded routing.
84    ///
85    /// An abandoned run is not recorded as a failed one. Dropping a [`CallModel`] without
86    /// calling this still yields [`DriverError::ResponseDropped`], which is.
87    pub fn into_parts(self) -> (Request, Decision) {
88        let Self {
89            request,
90            decision,
91            reply,
92            ..
93        } = self;
94        // Tell the algorithm the call was taken deliberately rather than lost, so its
95        // telemetry can tell an abandoned run from a failed one. The receiver is already
96        // gone if the algorithm stopped waiting, which is fine.
97        let _ = reply.send(Err(DriverError::Abandoned.into()));
98        (request, decision)
99    }
100}
101
102/// How an algorithm's [`route`](Algorithm::route) makes model calls.
103#[derive(Clone)]
104pub struct Driver {
105    step_tx: mpsc::Sender<Result<Step>>,
106    /// The owning algorithm's telemetry label, stamped onto every call and decision
107    /// this driver publishes.
108    algorithm: String,
109}
110
111impl Driver {
112    /// Build an empty driver with its step channel ready. Created per call by
113    /// [`run_stream`](Algorithm::run_stream). Also returns the Step receiver.
114    pub(crate) fn new(algorithm: &str) -> (Self, mpsc::Receiver<Result<Step>>) {
115        // Capacity one keeps the algorithm paced by the stream consumer. It limits queued steps,
116        // not model calls already pulled from the stream, which can still run at the same time.
117        // A larger buffer would use more memory and let the algorithm run farther ahead with
118        // little benefit because reading a step is cheap compared with serving a model call.
119        let (step_tx, step_rx) = mpsc::channel(1);
120        (
121            Self {
122                step_tx,
123                algorithm: algorithm.to_string(),
124            },
125            step_rx,
126        )
127    }
128
129    /// Offload a model call: publish it as a [`Step::CallModel`] and await the consumer's
130    /// [`Response`]. Errors if the stream is closed or the call failed.
131    /// The await is wrapped in a `libsy.llm_call` span measuring *fulfillment* as
132    /// the algorithm observes it (host queueing/serving included; a streamed
133    /// response resolves when its stream handle arrives); latency, outcome, and
134    /// token usage are recorded when it resolves. The provider call itself is the
135    /// host's, and is instrumented by whoever makes it.
136    #[tracing::instrument(
137        target = "libsy",
138        name = "libsy.llm_call",
139        skip_all,
140        fields(
141            algorithm = self.algorithm,
142            selected_model = decision.selected_model_id(),
143            openinference.span.kind = "CHAIN",
144            outcome = tracing::field::Empty,
145            error = tracing::field::Empty,
146            input_tokens = tracing::field::Empty,
147            output_tokens = tracing::field::Empty,
148            total_tokens = tracing::field::Empty,
149            reasoning_tokens = tracing::field::Empty,
150        )
151    )]
152    pub async fn call_model(&self, request: Request, decision: Decision) -> Result<Response> {
153        let selected_model_id = decision.selected_model_id().to_string();
154        let is_answer_call = decision.is_answer_call();
155        let started = Instant::now();
156        let (reply, response) = oneshot::channel::<Result<Response>>();
157        let call = CallModel {
158            algorithm: self.algorithm.clone(),
159            request,
160            decision,
161            reply,
162        };
163        let result = async {
164            self.step_tx
165                .send(Ok(Step::CallModel(Box::new(call))))
166                .await
167                .map_err(|_| DriverError::StreamClosed)?;
168            response
169                .await
170                .map_err(|_| LibsyError::from(DriverError::ResponseDropped))?
171        }
172        .await;
173        let elapsed = started.elapsed();
174        observability::record_llm_call(
175            &self.algorithm,
176            &selected_model_id,
177            is_answer_call,
178            elapsed,
179            &result,
180            &tracing::Span::current(),
181        );
182        result
183    }
184
185    /// Publish a routing [`Decision`] as a [`Step::Decision`] on the stream.
186    /// Each successfully published decision is counted and logged with its
187    /// reasoning; a decision the stream never accepted is not recorded.
188    pub async fn decide(&self, decision: Decision) -> Result<()> {
189        self.step_tx
190            .send(Ok(Step::Decision(decision.clone())))
191            .await
192            .map_err(|_| DriverError::StreamClosed)?;
193        observability::record_decision(&self.algorithm, &decision);
194        Ok(())
195    }
196
197    /// Emit the terminal step: [`Step::Done`] on `Ok`, or an `Err` stream
198    /// item on failure. Internal: called once by [`run_stream`](Algorithm::run_stream)
199    /// when the algorithm finishes.
200    pub(crate) async fn finish(&self, result: Result<Response>) -> Result<()> {
201        let step = result.map(|response| Step::Done(Box::new(response)));
202        self.step_tx
203            .send(step)
204            .await
205            .map_err(|_| DriverError::StreamClosed.into())
206    }
207}
208
209/// One item in the stream returned by [`Algorithm::run_stream`].
210pub enum Step {
211    /// The algorithm needs this model call performed. The host serves it and fulfills
212    /// it with [`CallModel::respond`]. Boxed: it is by far the largest variant.
213    CallModel(Box<CallModel>),
214    /// A routing decision the algorithm made, published via [`Driver::decide`] as it
215    /// happens (rather than collected into a trace returned at the end).
216    Decision(Decision),
217    /// The algorithm finished with its final response — the last step of a run.
218    Done(Box<Response>),
219}
220
221/// Drive [`Algorithm::run_stream`] to completion, handing each offloaded call to `serve`.
222///
223/// Returns the final [`Response`] and the trace of [`Decision`]s the algorithm published.
224/// `serve` owns the call: it performs it however the host likes and must fulfill the promise
225/// with [`CallModel::respond`]. A failed *model* call belongs in `respond` — the
226/// algorithm may route around it. Returning `Err` from `serve` aborts the whole run, so
227/// reserve it for infrastructure failures. Calls are served concurrently, so an algorithm
228/// that offloads several at once (hedging, fan-out) gets real parallelism.
229///
230/// libsy performs no I/O; this is only the mechanics of consuming its own step stream, kept
231/// here so every host does not reimplement the same loop. `switchyard-llm-client`'s `run`
232/// is this function plus an HTTP client.
233pub async fn drive<F, Fut>(
234    algorithm: Arc<dyn Algorithm>,
235    request: Request,
236    serve: F,
237) -> Result<(Vec<Decision>, Response)>
238where
239    F: Fn(CallModel) -> Fut,
240    Fut: Future<Output = Result<()>>,
241{
242    let stream = algorithm.run_stream(request);
243    tokio::pin!(stream);
244
245    let mut trace: Vec<Decision> = Vec::new();
246    let mut in_flight = futures::stream::FuturesUnordered::new();
247    let mut final_response: Option<Response> = None;
248
249    loop {
250        tokio::select! {
251            Some(result) = in_flight.next() => match result {
252                Ok(()) => {}, // CallModel completed successfully
253                Err(err) => return Err(err), // CallModel failed, propagate the error
254            },
255            step = stream.next() => {
256                match step {
257                    None => break, // stream has ended, no more steps
258                    Some(item) => match item? {
259                        Step::CallModel(call) => in_flight.push(serve(*call)),
260                        Step::Decision(decision) => trace.push(decision),
261                        Step::Done(response) => {
262                            final_response = Some(*response);
263                            break;
264                        }
265                    }
266                }
267            },
268        }
269    }
270    final_response
271        .map(|response| (trace, response))
272        .ok_or(LibsyError::MissingFinalResponse)
273}
274
275/// Recover the message from an algorithm's panic.
276fn panic_message(payload: &(dyn std::any::Any + Send)) -> String {
277    payload
278        .downcast_ref::<&'static str>()
279        .map(|message| (*message).to_string())
280        .or_else(|| payload.downcast_ref::<String>().cloned())
281        .unwrap_or_else(|| "unknown panic payload".to_string())
282}
283
284/// Abort guard
285struct AbortOnDrop(tokio::task::AbortHandle);
286
287impl Drop for AbortOnDrop {
288    fn drop(&mut self) {
289        self.0.abort();
290    }
291}
292
293/// A named routing target an algorithm routes by. Serving its calls is the stream
294/// consumer's concern: the selected identifier reaches the consumer as
295/// `decision.selected_model_id()` on the offloaded [`CallModel`].
296#[derive(Clone)]
297pub struct LlmTarget {
298    /// The routing label an algorithm selects this target by — a logical tier like
299    /// `"strong"`, or the model id when they coincide. Mapping it to a provider model
300    /// id is the consumer's concern, never the algorithm's.
301    pub semantic_name: String,
302}
303
304/// The set of targets an algorithm may route among. An algorithm is constructed
305/// with one and picks targets by position ([`targets`](Self::targets)) or by name
306/// ([`get_target`](Self::get_target)).
307#[derive(Clone)]
308pub struct LlmTargetSet {
309    targets: Vec<LlmTarget>,
310}
311
312impl LlmTargetSet {
313    /// Build a target set from a list of targets.
314    pub fn new(targets: Vec<LlmTarget>) -> Self {
315        Self { targets }
316    }
317
318    /// All targets in the set — e.g. for an algorithm to select among.
319    pub fn targets(&self) -> &[LlmTarget] {
320        &self.targets
321    }
322
323    /// Look up a target by name; errors if no target has that name.
324    pub fn get_target(&self, name: &str) -> Result<LlmTarget> {
325        self.targets
326            .iter()
327            .find(|t| t.semantic_name == name)
328            .cloned()
329            .ok_or_else(|| LibsyError::TargetNotFound {
330                target: name.to_string(),
331            })
332    }
333
334    /// The named target, or the first one not in `excluded` when it has been barred.
335    /// Errors if every target is excluded.
336    pub fn resolve_target(&self, name: &str, excluded: &HashSet<String>) -> Result<LlmTarget> {
337        let target = self.get_target(name)?;
338        if !excluded.contains(&target.semantic_name) {
339            return Ok(target);
340        }
341        self.targets
342            .iter()
343            .find(|t| !excluded.contains(&t.semantic_name))
344            .cloned()
345            .ok_or(LibsyError::AllTargetsExcluded)
346    }
347}
348
349/// Key for overflow history: a root request by its session, a child request by its session
350/// and agent. Keying a child finer than its session keeps one child's overflow from evicting
351/// a target for the parent or a sibling sharing the session.
352#[derive(Clone, Hash, PartialEq, Eq)]
353pub(crate) enum RoutingIdentity {
354    /// Root request, keyed by session ID.
355    Session(String),
356    /// Child request, keyed by session and agent IDs.
357    Subagent { session: String, agent: String },
358}
359
360impl RoutingIdentity {
361    /// Builds a root or child identity from non-empty request metadata.
362    ///
363    /// A child request missing either ID returns `None`, so it keeps no routing history
364    /// rather than sharing the parent's.
365    pub(crate) fn from_request(request: &Request) -> Option<Self> {
366        let metadata = request.metadata.as_ref()?;
367        let session = metadata.session_id.as_deref().filter(|id| !id.is_empty())?;
368        if metadata.is_subagent {
369            let agent = metadata.agent_id.as_deref().filter(|id| !id.is_empty())?;
370            Some(Self::Subagent {
371                session: session.to_string(),
372                agent: agent.to_string(),
373            })
374        } else {
375            Some(Self::Session(session.to_string()))
376        }
377    }
378
379    /// The session this identity belongs to; shared by a session's root and its children.
380    fn session(&self) -> &str {
381        match self {
382            Self::Session(session) | Self::Subagent { session, .. } => session,
383        }
384    }
385}
386
387/// Bounds process-local overflow history. Dropping a live entry costs one rediscovered
388/// overflow, so the victim choice does not need to be exact.
389const MAX_EVICTION_IDENTITIES: usize = 1_024;
390
391/// Per-identity record of the targets that overflowed their context window.
392///
393/// A conversation only grows, so a target that could not fit one turn will not fit a
394/// later one; remembering it lets the next turn skip a call certain to fail. Requests
395/// without a routing identity are not tracked — there is nothing to remember them by.
396#[derive(Default)]
397pub(crate) struct SessionEvictions {
398    by_identity: Mutex<HashMap<RoutingIdentity, HashSet<String>>>,
399}
400
401impl SessionEvictions {
402    /// Forgets overflow history for a completed session, including every child of it.
403    pub(crate) fn remove_session(&self, session: &str) {
404        self.by_identity
405            .lock()
406            .retain(|identity, _| identity.session() != session);
407    }
408
409    /// The targets `identity` has already overflowed; empty for an untracked request.
410    fn evicted_for(&self, identity: Option<&RoutingIdentity>) -> Vec<String> {
411        let Some(identity) = identity else {
412            return Vec::new();
413        };
414        self.by_identity
415            .lock()
416            .get(identity)
417            .map(|targets| targets.iter().cloned().collect())
418            .unwrap_or_default()
419    }
420
421    /// Remembers that `target` overflowed for `identity`, tracking at most
422    /// [`MAX_EVICTION_IDENTITIES`] identities.
423    fn record(&self, identity: Option<&RoutingIdentity>, target: &str) {
424        let Some(identity) = identity else { return };
425        let mut histories = self.by_identity.lock();
426        if histories.len() >= MAX_EVICTION_IDENTITIES
427            && !histories.contains_key(identity)
428            && let Some(oldest) = histories.keys().next().cloned()
429        {
430            histories.remove(&oldest);
431        }
432        histories
433            .entry(identity.clone())
434            .or_default()
435            .insert(target.to_string());
436    }
437}
438
439/// How many of `targets` this request is still allowed to reach.
440fn eligible_targets(targets: &LlmTargetSet, excluded: &HashSet<String>) -> usize {
441    targets
442        .targets()
443        .iter()
444        .filter(|t| !excluded.contains(&t.semantic_name))
445        .count()
446}
447
448/// Bars the targets `identity` has already overflowed from this request, so routing does
449/// not select one that is certain to fail again.
450pub(crate) fn exclude_evicted(
451    excluded: &mut HashSet<String>,
452    targets: &LlmTargetSet,
453    evictions: &SessionEvictions,
454    identity: Option<&RoutingIdentity>,
455) {
456    for target in evictions.evicted_for(identity) {
457        // Never seed the pool empty: a later turn may be small enough to serve, and the
458        // caller should get the upstream's answer rather than a routing error.
459        if eligible_targets(targets, excluded) <= 1 {
460            break;
461        }
462        excluded.insert(target);
463    }
464}
465
466/// Returns the failed target and routing fallback policy for a terminal client error.
467fn classify_fallback(error: &LibsyError) -> Option<(&str, RoutingFallbackReason)> {
468    let LibsyError::ClientCall { target, source } = error else {
469        return None;
470    };
471    let reason = match source {
472        LlmClientError::ContextWindowExceeded { .. } => RoutingFallbackReason::ContextWindow,
473        LlmClientError::Transport { .. } | LlmClientError::Timeout { .. } => {
474            RoutingFallbackReason::Unavailable
475        }
476        LlmClientError::UpstreamHttp { status, .. }
477            if matches!(*status, 403 | 408 | 429) || (500..=599).contains(status) =>
478        {
479            RoutingFallbackReason::Unavailable
480        }
481        _ => return None,
482    };
483    Some((target, reason))
484}
485
486/// Calls `target`, falling back to the next eligible target after a route-level failure,
487/// until a call succeeds or every target has been tried.
488///
489/// Routing is deliberately not re-run: the fallback replaces the target in place, so the
490/// caller's request-side work and retained state still see exactly one turn.
491/// `fallback_decision` builds the [`Decision`] published for a `from -> to` hop. Context
492/// overflows are recorded for `identity`; unavailable targets remain request-local.
493#[allow(clippy::too_many_arguments)]
494pub(crate) async fn call_model_with_fallback(
495    excluded: &mut HashSet<String>,
496    driver: &Driver,
497    targets: &LlmTargetSet,
498    mut target: LlmTarget,
499    mut decision: Decision,
500    request: Request,
501    identity: Option<&RoutingIdentity>,
502    evictions: &SessionEvictions,
503    target_unavailable: impl Fn(&Request, &str),
504    fallback_decision: impl Fn(&LlmTarget, &LlmTarget, RoutingFallbackReason) -> Decision,
505) -> Result<Response> {
506    loop {
507        let result = driver.call_model(request.clone(), decision.clone()).await;
508        let Err(error) = result else { return result };
509        let Some((failed, reason)) = classify_fallback(&error) else {
510            return Err(error);
511        };
512        // A target already excluded means the pool is spent; surface the client error
513        // so the caller still sees the concrete upstream failure.
514        if !excluded.insert(failed.to_string()) {
515            return Err(error);
516        }
517        match reason {
518            RoutingFallbackReason::ContextWindow => evictions.record(identity, failed),
519            RoutingFallbackReason::Unavailable => target_unavailable(&request, failed),
520        }
521        let Ok(next) = targets.resolve_target(&target.semantic_name, excluded) else {
522            return Err(error);
523        };
524        decision = fallback_decision(&target, &next, reason);
525        target = next;
526        driver.decide(decision.clone()).await?;
527    }
528}
529
530/// An optimization strategy. Implement [`route`](Self::route);
531/// callers drive it with [`run_stream`](Self::run_stream), serving each [`Step::CallModel`]
532/// it emits. `switchyard-llm-client`'s `run` is the ready-made consumer that does this
533/// over HTTP.
534///
535/// Methods take `self: Arc<Self>`: one algorithm (`Arc<dyn Algorithm>`) is shared across
536/// requests and run concurrently, so it owns its thread-safety and any shared state.
537///
538/// # Concurrency
539///
540/// A host may run the same algorithm concurrently for many requests. Implementations
541/// must synchronize their own mutable shared state. Each call to [`run_stream`](Self::run_stream)
542/// creates an independent [`Driver`], so model-call promises and emitted [`Step`]s cannot
543/// cross between runs.
544///
545/// # Observability
546///
547/// [`run_stream`](Self::run_stream) creates a `libsy.run` span, and each offloaded model
548/// call creates a `libsy.llm_call` span. Decisions and failures are emitted through
549/// `tracing`; metrics use the global OpenTelemetry meter provider. The provider call
550/// itself belongs to the host, and is instrumented by whoever makes it.
551#[async_trait]
552pub trait Algorithm: Send + Sync + 'static {
553    /// Stable, low-cardinality name identifying this algorithm — the
554    /// `algorithm` attribute on every span, metric, and log line the crate
555    /// emits for its runs.
556    fn name(&self) -> &str;
557
558    /// Run one request to completion: make model calls with [`Driver::call_model`],
559    /// publish [`Decision`]s with [`Driver::decide`], and return the final [`Response`].
560    /// The method an algorithm implements; [`run_stream`](Self::run_stream) drives it.
561    async fn route(self: Arc<Self>, driver: Driver, request: Request) -> Result<Response>;
562
563    /// Process a request to completion, returning a stream of [`Step`]s.
564    ///
565    /// The consumer must fulfill every [`Step::CallModel`] before the algorithm can
566    /// continue. Every run ends with exactly one terminal item — [`Step::Done`] on
567    /// success, an `Err` item on failure, including when the algorithm panics. Dropping
568    /// the stream aborts the spawned algorithm task.
569    ///
570    /// Every invocation owns a separate [`Driver`].
571    fn run_stream(self: Arc<Self>, request: Request) -> StepStream {
572        let (driver, step_rx) = Driver::new(self.name());
573        let span = observability::run_span(self.name(), &request);
574        let handle = tokio::spawn(
575            async move {
576                let algorithm = self.name().to_string();
577                // Catch a panicking algorithm so the run still publishes a terminal step.
578                let route = AssertUnwindSafe(self.route(driver.clone(), request)).catch_unwind();
579                let result = observability::observe_run(&algorithm, async move {
580                    route.await.unwrap_or_else(|payload| {
581                        Err(LibsyError::AlgorithmError {
582                            message: format!(
583                                "algorithm task panicked: {}",
584                                panic_message(payload.as_ref())
585                            ),
586                        })
587                    })
588                })
589                .await;
590
591                let _ = driver.finish(result).await;
592            }
593            .instrument(span),
594        );
595        // Dropping the stream aborts the algorithm task when its consumer goes away.
596        let abort_guard = AbortOnDrop(handle.abort_handle());
597        Box::pin(ReceiverStream::new(step_rx).map(move |step| {
598            // link abort guard to stream
599            let _keep_alive = &abort_guard;
600            step
601        }))
602    }
603}
604
605#[cfg(test)]
606mod tests {
607    use super::*;
608    use crate::core::testing::{Serve, ServeResult, echo, reply, test_drive};
609    use futures::StreamExt;
610    use switchyard_protocol::{
611        LlmResponse, LlmResponseChunk, completion_text, text_request, text_response,
612    };
613
614    #[derive(Debug, thiserror::Error)]
615    #[error("{0}")]
616    struct TestError(&'static str);
617
618    fn test_error(message: &'static str) -> LibsyError {
619        LibsyError::external("test", TestError(message))
620    }
621
622    fn classified_client_error(source: LlmClientError) -> Option<RoutingFallbackReason> {
623        classify_fallback(&LibsyError::client_call("target", source)).map(|(_, reason)| reason)
624    }
625
626    #[test]
627    fn route_fallback_only_accepts_context_and_unavailable_failures() {
628        assert_eq!(
629            classified_client_error(LlmClientError::ContextWindowExceeded {
630                model: "target".to_string(),
631                message: "too long".to_string(),
632            }),
633            Some(RoutingFallbackReason::ContextWindow)
634        );
635        for source in [
636            LlmClientError::Transport {
637                source: Box::new(std::io::Error::other("connection failed")),
638            },
639            LlmClientError::Timeout {
640                source: Box::new(std::io::Error::other("request timed out")),
641            },
642        ] {
643            assert_eq!(
644                classified_client_error(source),
645                Some(RoutingFallbackReason::Unavailable)
646            );
647        }
648        for (status, expected) in [
649            (400, None),
650            (401, None),
651            (403, Some(RoutingFallbackReason::Unavailable)),
652            (404, None),
653            (408, Some(RoutingFallbackReason::Unavailable)),
654            (409, None),
655            (429, Some(RoutingFallbackReason::Unavailable)),
656            (499, None),
657            (500, Some(RoutingFallbackReason::Unavailable)),
658            (599, Some(RoutingFallbackReason::Unavailable)),
659            (600, None),
660        ] {
661            assert_eq!(
662                classified_client_error(LlmClientError::UpstreamHttp {
663                    status,
664                    body: "failed".to_string(),
665                }),
666                expected
667            );
668        }
669        assert_eq!(
670            classified_client_error(LlmClientError::InvalidResponse {
671                source: Box::new(std::io::Error::other("invalid response")),
672            }),
673            None
674        );
675    }
676
677    /// Build a routed decision for orchestration tests.
678    fn test_decision(selected_model_id: String) -> Decision {
679        Decision::new(selected_model_id, None, true)
680    }
681
682    /// Trivial algo used only to exercise the orchestrator: calls the first target
683    /// and returns its response with a one-item trace.
684    struct TestAlgo {
685        target_set: LlmTargetSet,
686    }
687
688    #[async_trait]
689    impl Algorithm for TestAlgo {
690        fn name(&self) -> &str {
691            "test"
692        }
693
694        async fn route(self: Arc<Self>, driver: Driver, request: Request) -> Result<Response> {
695            let target = self
696                .target_set
697                .targets()
698                .first()
699                .ok_or(LibsyError::NoTargets)?
700                .clone();
701            let decision = test_decision(target.semantic_name.clone());
702            driver.decide(decision.clone()).await?;
703            driver.call_model(request, decision).await
704        }
705    }
706
707    /// Build a shared `TestAlgo` over the given target set.
708    fn orch(target_set: LlmTargetSet) -> Arc<dyn Algorithm> {
709        Arc::new(TestAlgo { target_set })
710    }
711
712    fn request() -> Request {
713        Request {
714            llm_request: text_request(Some("auto".to_string()), "hi".to_string()),
715            raw_request: None,
716            metadata: None,
717        }
718    }
719
720    fn target_set(names: &[&str]) -> LlmTargetSet {
721        let targets = names
722            .iter()
723            .map(|name| LlmTarget {
724                semantic_name: name.to_string(),
725            })
726            .collect();
727        LlmTargetSet::new(targets)
728    }
729
730    #[tokio::test]
731    async fn typed_driver_preserves_call_and_stream_boundaries() -> Result<()> {
732        tokio::time::timeout(std::time::Duration::from_secs(1), async {
733            // Distinct oneshots keep reverse-order replies paired with their producers, and a
734            // retained call remains pending until the host responds.
735            let (driver, mut step_rx) = Driver::new("test");
736            let first_driver = driver.clone();
737            let mut first = tokio::spawn(async move {
738                first_driver
739                    .call_model(request(), test_decision("first".to_string()))
740                    .await
741            });
742            let second = tokio::spawn(async move {
743                driver
744                    .call_model(request(), test_decision("second".to_string()))
745                    .await
746            });
747
748            let mut calls = HashMap::new();
749            for _ in 0..2 {
750                let step = step_rx.recv().await.ok_or(DriverError::StreamClosed)??;
751                let Step::CallModel(call) = step else {
752                    return Err(test_error("expected a CallModel step"));
753                };
754                calls.insert(call.decision.selected_model_id().to_string(), call);
755            }
756            assert!(
757                tokio::time::timeout(std::time::Duration::from_millis(20), &mut first)
758                    .await
759                    .is_err(),
760                "call completed before the host responded"
761            );
762            calls
763                .remove("second")
764                .ok_or_else(|| test_error("missing second call"))?
765                .respond(Ok(reply("second response")))?;
766            calls
767                .remove("first")
768                .ok_or_else(|| test_error("missing first call"))?
769                .respond(Ok(reply("first response")))?;
770
771            let first_response = first
772                .await
773                .map_err(|source| LibsyError::external("joining a test task", source))??;
774            let second_response = second
775                .await
776                .map_err(|source| LibsyError::external("joining a test task", source))??;
777            assert_eq!(
778                first_response.llm_response.as_agg().map(completion_text),
779                Some("first response".to_string())
780            );
781            assert_eq!(
782                second_response.llm_response.as_agg().map(completion_text),
783                Some("second response".to_string())
784            );
785
786            // Dropping the host-facing promise closes only that call's reply channel.
787            let (driver, mut step_rx) = Driver::new("test");
788            let producer = tokio::spawn(async move {
789                driver
790                    .call_model(request(), test_decision("dropped".to_string()))
791                    .await
792            });
793            let step = step_rx.recv().await.ok_or(DriverError::StreamClosed)??;
794            let Step::CallModel(call) = step else {
795                return Err(test_error("expected a CallModel step"));
796            };
797            drop(call);
798            let result = producer
799                .await
800                .map_err(|source| LibsyError::external("joining a test task", source))?;
801            assert!(matches!(
802                result,
803                Err(LibsyError::Driver(DriverError::ResponseDropped))
804            ));
805
806            // A standalone driver reports the typed step receiver disappearing at its next send.
807            let (driver, step_rx) = Driver::new("test");
808            drop(step_rx);
809            let decision = test_decision("closed".to_string());
810            let result = driver.decide(decision).await;
811            assert!(matches!(
812                result,
813                Err(LibsyError::Driver(DriverError::StreamClosed))
814            ));
815            Ok(())
816        })
817        .await
818        .map_err(|error| LibsyError::external("waiting for typed driver boundaries", error))?
819    }
820
821    /// A consumer that only wants the routing outcome takes the call apart instead of
822    /// answering it: it gets the request as the algorithm would have sent it, and the
823    /// algorithm learns the call was abandoned rather than lost — the distinction that
824    /// keeps a deliberate decision-only run out of the failure counters.
825    #[tokio::test]
826    async fn into_parts_yields_the_call_without_answering_it() -> Result<()> {
827        let (driver, mut step_rx) = Driver::new("test");
828        let decision = test_decision("answer/model".to_string());
829        let producer = tokio::spawn({
830            let decision = decision.clone();
831            async move { driver.call_model(request(), decision).await }
832        });
833
834        let step = step_rx.recv().await.ok_or(DriverError::StreamClosed)??;
835        let Step::CallModel(call) = step else {
836            return Err(test_error("expected a CallModel step"));
837        };
838        let (taken_request, taken_decision) = call.into_parts();
839        assert_eq!(taken_decision.selected_model_id(), "answer/model");
840        assert!(taken_decision.is_answer_call());
841        assert_eq!(taken_request.llm_request, request().llm_request);
842
843        let result = producer
844            .await
845            .map_err(|source| LibsyError::external("joining a test task", source))?;
846        assert!(matches!(
847            result,
848            Err(LibsyError::Driver(DriverError::Abandoned))
849        ));
850        Ok(())
851    }
852
853    /// The two ways a call goes unanswered are told apart: `into_parts` is the consumer's
854    /// choice, a bare drop is the promise being lost.
855    #[test]
856    fn abandoning_and_dropping_a_call_report_different_outcomes() {
857        let abandoned: Result<Response> = Err(DriverError::Abandoned.into());
858        let dropped: Result<Response> = Err(DriverError::ResponseDropped.into());
859        assert_eq!(observability::outcome_value(&abandoned), "abandoned");
860        assert_eq!(observability::outcome_value(&dropped), "error");
861        assert!(observability::is_abandoned(&abandoned));
862        assert!(!observability::is_abandoned(&dropped));
863    }
864
865    #[test]
866    fn target_lookup_returns_the_missing_target() {
867        let error = target_set(&[]).get_target("missing").err();
868        assert!(matches!(
869            error,
870            Some(LibsyError::TargetNotFound { target }) if target == "missing"
871        ));
872    }
873
874    /// Build a single-target algo, plus a `serve` that answers it as a token stream
875    /// replaying `chunks` in order (as `Ok` items).
876    fn streaming_orch(chunks: Vec<LlmResponseChunk>) -> (Arc<dyn Algorithm>, impl Serve) {
877        let algo = orch(target_set(&["stream/model"]));
878        let serve = move |_decision: Decision, _request: Request| {
879            let chunks = chunks.clone();
880            async move {
881                let stream =
882                    futures::stream::iter(chunks.into_iter().map(|chunk| Ok(chunk.into()))).boxed();
883                Ok(Response {
884                    llm_response: LlmResponse::Stream(stream),
885                    metadata: None,
886                })
887            }
888        };
889        (algo, serve)
890    }
891
892    #[tokio::test]
893    async fn run_returns_a_streamed_response_the_caller_aggregates() -> Result<()> {
894        // A streaming client -> its chunks flow through the promise and `Done`,
895        // and `run` returns the live stream untouched for the caller to fold.
896        let (orch, serve) = streaming_orch(vec![
897            LlmResponseChunk::MessageStart {
898                id: Some("m1".to_string()),
899                model: Some("stream/model".to_string()),
900            },
901            LlmResponseChunk::TextDelta {
902                index: 0,
903                text: "hel".to_string(),
904            },
905            LlmResponseChunk::TextDelta {
906                index: 0,
907                text: "lo".to_string(),
908            },
909            LlmResponseChunk::MessageStop {
910                reason: Some("stop".to_string()),
911            },
912        ]);
913        let (trace, response) = test_drive(orch, request(), serve).await?;
914        // The run handed back the live stream; the caller folds it to a buffered aggregate.
915        let agg = response
916            .llm_response
917            .into_agg()
918            .await
919            .map_err(|error| LibsyError::external("aggregating response stream", error))?;
920        assert_eq!(completion_text(&agg), "hello");
921        assert_eq!(agg.model.as_deref(), Some("stream/model"));
922        assert_eq!(trace.len(), 1);
923        Ok(())
924    }
925
926    #[tokio::test]
927    async fn aggregating_a_streamed_response_propagates_a_mid_stream_error() -> Result<()> {
928        // The run succeeds and returns the stream; the in-band `Error` chunk surfaces only
929        // when the caller aggregates it.
930        let (orch, serve) = streaming_orch(vec![
931            LlmResponseChunk::TextDelta {
932                index: 0,
933                text: "partial".to_string(),
934            },
935            LlmResponseChunk::StreamError {
936                message: "upstream exploded".to_string(),
937            },
938        ]);
939        let (_, response) = test_drive(orch, request(), serve).await?;
940        match response.llm_response.into_agg().await {
941            Ok(_) => panic!("expected a mid-stream error, got an aggregate"),
942            Err(err) => {
943                assert!(err.to_string().contains("upstream exploded"));
944                Ok(())
945            }
946        }
947    }
948
949    #[tokio::test]
950    async fn run_offloads_via_promise_then_finishes() -> Result<()> {
951        // Every call is offloaded via a promise the orchestrator surfaces as a
952        // `CallModel` step for us to fulfill.
953        let stream = orch(target_set(&["offload/model"])).run_stream(request());
954        tokio::pin!(stream);
955
956        let mut saw_call = false;
957        let mut final_completion = None;
958        while let Some(step) = stream.next().await {
959            match step? {
960                Step::CallModel(call) => {
961                    saw_call = true;
962                    // The decision rode along with the promise.
963                    assert_eq!(call.decision.selected_model_id(), "offload/model");
964                    // Fulfilling the promise is the "real" model call the caller makes.
965                    call.respond(Ok(Response {
966                        llm_response: LlmResponse::Agg(text_response(
967                            None,
968                            "fulfilled".to_string(),
969                        )),
970                        metadata: None,
971                    }))?;
972                }
973                Step::Decision(decision) => {
974                    assert_eq!(decision.selected_model_id(), "offload/model");
975                }
976                Step::Done(response) => {
977                    final_completion = Some(
978                        response
979                            .llm_response
980                            .as_agg()
981                            .map(completion_text)
982                            .unwrap_or_default(),
983                    );
984                }
985            }
986        }
987
988        assert!(saw_call, "expected a CallModel step before Done");
989        assert_eq!(
990            final_completion.ok_or_else(|| test_error("no Done step"))?,
991            "fulfilled"
992        );
993        Ok(())
994    }
995
996    #[tokio::test]
997    async fn a_driven_run_returns_the_trace_and_the_final_response() -> Result<()> {
998        let (trace, response) =
999            test_drive(orch(target_set(&["direct/model"])), request(), echo()).await?;
1000        // TestAlgo calls the first target; `echo` answers with its name.
1001        assert_eq!(
1002            response
1003                .llm_response
1004                .as_agg()
1005                .map(completion_text)
1006                .unwrap_or_default(),
1007            "direct/model"
1008        );
1009        assert_eq!(trace[0].selected_model_id(), "direct/model");
1010        Ok(())
1011    }
1012
1013    #[tokio::test(flavor = "multi_thread", worker_threads = 12)]
1014    async fn requests_are_processed_in_parallel() -> Result<()> {
1015        use std::time::Duration;
1016        use tokio::sync::Barrier;
1017
1018        const N: usize = 12;
1019
1020        // Serving blocks until all N concurrent calls have arrived. If requests were
1021        // serialized (one algorithm behind a `Mutex`), only one call could be in flight,
1022        // the barrier would never reach N, and the test would time out. It passes only
1023        // because the shared algorithm is driven concurrently across requests.
1024        let barrier = Arc::new(Barrier::new(N));
1025        // One shared algorithm driven by many concurrent requests.
1026        let algo = orch(target_set(&["m"]));
1027
1028        let mut handles = Vec::new();
1029        for _ in 0..N {
1030            let algo = algo.clone();
1031            let barrier = barrier.clone();
1032            let serve = move |decision: Decision, _request: Request| {
1033                let barrier = barrier.clone();
1034                async move {
1035                    barrier.wait().await;
1036                    Ok(reply(decision.selected_model_id()))
1037                }
1038            };
1039            handles.push(tokio::spawn(async move {
1040                test_drive(algo, request(), serve)
1041                    .await
1042                    .map(|(_, response)| {
1043                        response
1044                            .llm_response
1045                            .as_agg()
1046                            .map(completion_text)
1047                            .unwrap_or_default()
1048                    })
1049            }));
1050        }
1051
1052        for handle in handles {
1053            // The timeout turns a serialization deadlock into a failure, not a hang.
1054            let completion = tokio::time::timeout(Duration::from_secs(5), handle)
1055                .await
1056                .map_err(|error| LibsyError::external("waiting for test task", error))?
1057                .map_err(|source| LibsyError::external("joining a test task", source))??;
1058            assert_eq!(completion, "m");
1059        }
1060        Ok(())
1061    }
1062
1063    #[tokio::test]
1064    async fn offload_error_propagates_back_to_the_algorithm() -> Result<()> {
1065        // A client-less target offloads its call; we fulfill the promise with an
1066        // Err, which must flow back through `call_model_target` into the algorithm and
1067        // out as an error step — not a response.
1068        let stream = orch(target_set(&["offload/model"])).run_stream(request());
1069        tokio::pin!(stream);
1070
1071        let mut saw_error = false;
1072        while let Some(step) = stream.next().await {
1073            match step {
1074                Ok(Step::CallModel(call)) => {
1075                    call.respond(Err(test_error("upstream model call failed")))?;
1076                }
1077                Ok(Step::Decision(_)) => {}
1078                Ok(Step::Done(..)) => {
1079                    return Err(test_error(
1080                        "expected the offload error to propagate, got a response",
1081                    ));
1082                }
1083                Err(err) => {
1084                    // The algorithm's `call_model_target` saw the error via the promise.
1085                    assert!(err.to_string().contains("upstream model call failed"));
1086                    saw_error = true;
1087                }
1088            }
1089        }
1090
1091        assert!(saw_error, "expected an error step");
1092        Ok(())
1093    }
1094
1095    #[tokio::test]
1096    async fn dropping_the_stream_cancels_the_algorithm_task() -> Result<()> {
1097        use std::sync::atomic::{AtomicBool, Ordering};
1098        use std::time::Duration;
1099        use tokio::sync::mpsc;
1100
1101        // Sets a flag when dropped, so we can observe whether the algorithm task was
1102        // cancelled/dropped.
1103        struct DropGuard(Arc<AtomicBool>);
1104        impl Drop for DropGuard {
1105            fn drop(&mut self) {
1106                self.0.store(true, Ordering::SeqCst);
1107            }
1108        }
1109
1110        struct StuckAlgo {
1111            started: mpsc::UnboundedSender<()>,
1112            dropped: Arc<AtomicBool>,
1113        }
1114
1115        #[async_trait]
1116        impl Algorithm for StuckAlgo {
1117            fn name(&self) -> &str {
1118                "stuck"
1119            }
1120
1121            async fn route(
1122                self: Arc<Self>,
1123                _driver: Driver,
1124                _request: Request,
1125            ) -> Result<Response> {
1126                let _guard = DropGuard(self.dropped.clone());
1127                let _ = self.started.send(());
1128                // Await forever without ever touching the driver.
1129                std::future::pending::<()>().await;
1130                unreachable!()
1131            }
1132        }
1133
1134        let (started_tx, mut started_rx) = mpsc::unbounded_channel();
1135        let dropped = Arc::new(AtomicBool::new(false));
1136        let algo: Arc<dyn Algorithm> = Arc::new(StuckAlgo {
1137            started: started_tx,
1138            dropped: dropped.clone(),
1139        });
1140
1141        let stream = algo.run_stream(request());
1142        started_rx
1143            .recv()
1144            .await
1145            .ok_or_else(|| test_error("task never started"))?;
1146        drop(stream);
1147        tokio::time::sleep(Duration::from_millis(100)).await;
1148
1149        assert!(
1150            dropped.load(Ordering::SeqCst),
1151            "algorithm task was NOT cancelled after dropping the stream"
1152        );
1153        Ok(())
1154    }
1155
1156    #[tokio::test]
1157    async fn route_panic_surfaces_as_a_stream_error() -> Result<()> {
1158        // An algorithm whose task panics must surface an `Err` step carrying the panic
1159        // message, not abort the process from an unobserved detached task.
1160        struct Panicky;
1161
1162        #[async_trait]
1163        impl Algorithm for Panicky {
1164            fn name(&self) -> &str {
1165                "panicky"
1166            }
1167
1168            async fn route(
1169                self: Arc<Self>,
1170                _driver: Driver,
1171                _request: Request,
1172            ) -> Result<Response> {
1173                panic!("boom");
1174            }
1175        }
1176
1177        let algo: Arc<dyn Algorithm> = Arc::new(Panicky);
1178        let stream = algo.run_stream(request());
1179        tokio::pin!(stream);
1180
1181        let mut saw_error = false;
1182        while let Some(step) = stream.next().await {
1183            match step {
1184                Err(err) => {
1185                    // The panic message is preserved, not flattened into an opaque failure.
1186                    assert!(err.to_string().contains("algorithm task panicked: boom"));
1187                    saw_error = true;
1188                }
1189                Ok(_) => return Err(test_error("expected the panic to surface as an error step")),
1190            }
1191        }
1192
1193        assert!(saw_error, "expected an error step from the panicked task");
1194        Ok(())
1195    }
1196
1197    /// A panicking algorithm must publish its terminal step even when it left a `Driver`
1198    /// clone alive in another task. That clone holds the step channel open, so a run that
1199    /// merely unwound would never terminate and the consumer would wait forever.
1200    #[tokio::test]
1201    async fn a_panic_with_a_leaked_driver_clone_still_terminates_the_run() -> Result<()> {
1202        struct LeakyPanic;
1203
1204        #[async_trait]
1205        impl Algorithm for LeakyPanic {
1206            fn name(&self) -> &str {
1207                "leaky_panic"
1208            }
1209
1210            async fn route(self: Arc<Self>, driver: Driver, _request: Request) -> Result<Response> {
1211                tokio::spawn(async move {
1212                    // Outlives the panic below, keeping a sender clone alive.
1213                    let _keep_alive = driver;
1214                    std::future::pending::<()>().await;
1215                });
1216                tokio::task::yield_now().await;
1217                panic!("boom");
1218            }
1219        }
1220
1221        let algo: Arc<dyn Algorithm> = Arc::new(LeakyPanic);
1222        // The timeout turns the hang this guards against into a failure rather than a hang.
1223        let result = tokio::time::timeout(
1224            std::time::Duration::from_secs(1),
1225            test_drive(algo, request(), echo()),
1226        )
1227        .await
1228        .map_err(|error| LibsyError::external("waiting for the panicked run to end", error))?;
1229
1230        match result {
1231            Ok(_) => Err(test_error(
1232                "expected the panic to end the run with an error",
1233            )),
1234            Err(err) => {
1235                assert!(err.to_string().contains("algorithm task panicked: boom"));
1236                Ok(())
1237            }
1238        }
1239    }
1240
1241    #[tokio::test]
1242    async fn cancelling_run_cancels_the_algorithm_task() -> Result<()> {
1243        use std::sync::atomic::{AtomicBool, Ordering};
1244        use std::time::Duration;
1245        use tokio::sync::mpsc;
1246
1247        // Sets a flag when dropped, so we can observe whether the algorithm task was
1248        // cancelled once the `run` future driving it is dropped.
1249        struct DropGuard(Arc<AtomicBool>);
1250        impl Drop for DropGuard {
1251            fn drop(&mut self) {
1252                self.0.store(true, Ordering::SeqCst);
1253            }
1254        }
1255
1256        struct StuckAlgo {
1257            started: mpsc::UnboundedSender<()>,
1258            dropped: Arc<AtomicBool>,
1259        }
1260
1261        #[async_trait]
1262        impl Algorithm for StuckAlgo {
1263            fn name(&self) -> &str {
1264                "stuck"
1265            }
1266
1267            async fn route(
1268                self: Arc<Self>,
1269                _driver: Driver,
1270                _request: Request,
1271            ) -> Result<Response> {
1272                let _guard = DropGuard(self.dropped.clone());
1273                let _ = self.started.send(());
1274                // Hang forever without ever touching the driver, so only cancellation
1275                // (not a dropped step channel) can stop this task.
1276                std::future::pending::<()>().await;
1277                unreachable!()
1278            }
1279        }
1280
1281        let (started_tx, mut started_rx) = mpsc::unbounded_channel();
1282        let dropped = Arc::new(AtomicBool::new(false));
1283        let algo: Arc<dyn Algorithm> = Arc::new(StuckAlgo {
1284            started: started_tx,
1285            dropped: dropped.clone(),
1286        });
1287
1288        // Drive the run on its own task, wait until the algorithm task is up, then cancel
1289        // it — dropping its future (and the `run_stream` stream it holds).
1290        let run_task = tokio::spawn(async move { test_drive(algo, request(), echo()).await });
1291        started_rx
1292            .recv()
1293            .await
1294            .ok_or_else(|| test_error("task never started"))?;
1295        run_task.abort();
1296        tokio::time::sleep(Duration::from_millis(100)).await;
1297
1298        assert!(
1299            dropped.load(Ordering::SeqCst),
1300            "algorithm task was NOT cancelled after cancelling run"
1301        );
1302        Ok(())
1303    }
1304
1305    // --- first-wins hedging: `run` must not wait on losing speculative calls -------------
1306
1307    /// Offloads two targets concurrently and returns the first to resolve, dropping the
1308    /// loser's call (first-wins hedging).
1309    struct Hedge {
1310        winner: LlmTarget,
1311        loser: LlmTarget,
1312    }
1313
1314    #[async_trait]
1315    impl Algorithm for Hedge {
1316        fn name(&self) -> &str {
1317            "hedge"
1318        }
1319
1320        async fn route(self: Arc<Self>, driver: Driver, request: Request) -> Result<Response> {
1321            let dec_w = test_decision(self.winner.semantic_name.clone());
1322            let dec_l = test_decision(self.loser.semantic_name.clone());
1323            let win = driver.call_model(request.clone(), dec_w);
1324            let lose = driver.call_model(request, dec_l);
1325            // First to resolve wins; `select!` drops the losing future (and its promise).
1326            tokio::select! {
1327                res = win => res,
1328                res = lose => res,
1329            }
1330        }
1331    }
1332
1333    /// Builds a hedging algo and the `serve` that drives it: the winner is gated behind the
1334    /// loser starting (so the loser's serve is guaranteed in flight when the winner wins),
1335    /// and the loser finishes after `loser_delay` — or never, when `None`.
1336    fn hedge(loser_delay: Option<std::time::Duration>) -> (Arc<dyn Algorithm>, impl Serve) {
1337        let started = Arc::new(tokio::sync::Notify::new());
1338        let algo = Arc::new(Hedge {
1339            winner: LlmTarget {
1340                semantic_name: "winner".to_string(),
1341            },
1342            loser: LlmTarget {
1343                semantic_name: "loser".to_string(),
1344            },
1345        });
1346        let serve = move |decision: Decision, _request: Request| {
1347            let started = started.clone();
1348            async move {
1349                if decision.selected_model_id() == "loser" {
1350                    started.notify_one();
1351                    match loser_delay {
1352                        Some(delay) => tokio::time::sleep(delay).await,
1353                        None => std::future::pending::<()>().await,
1354                    }
1355                } else {
1356                    started.notified().await;
1357                }
1358                Ok(reply(decision.selected_model_id()))
1359            }
1360        };
1361        (algo, serve)
1362    }
1363
1364    #[tokio::test]
1365    async fn run_returns_the_winner_without_a_late_loser_overwriting_it() -> Result<()> {
1366        // The loser responds 50ms after the winner has already won. `run` must return the
1367        // winner, not the loser's `respond`-to-a-dropped-receiver error.
1368        let (algo, serve) = hedge(Some(std::time::Duration::from_millis(50)));
1369        let (_trace, response) = test_drive(algo, request(), serve).await?;
1370        assert_eq!(
1371            response
1372                .llm_response
1373                .as_agg()
1374                .map(completion_text)
1375                .unwrap_or_default(),
1376            "winner"
1377        );
1378        Ok(())
1379    }
1380
1381    #[tokio::test]
1382    async fn run_returns_the_winner_without_hanging_on_a_pending_loser() -> Result<()> {
1383        // The loser never resolves. `run` must return the winner promptly, not hang
1384        // waiting for the in-flight loser.
1385        let (algo, serve) = hedge(None);
1386        let run = test_drive(algo, request(), serve);
1387        let (_trace, response) = tokio::time::timeout(std::time::Duration::from_secs(1), run)
1388            .await
1389            .map_err(|error| LibsyError::external("waiting for pending loser", error))??;
1390        assert_eq!(
1391            response
1392                .llm_response
1393                .as_agg()
1394                .map(completion_text)
1395                .unwrap_or_default(),
1396            "winner"
1397        );
1398        Ok(())
1399    }
1400
1401    #[tokio::test]
1402    async fn run_surfaces_a_terminal_error_with_many_calls_in_flight() -> Result<()> {
1403        use std::sync::atomic::{AtomicUsize, Ordering};
1404
1405        // A large fan-out (10 matched the old, now-removed concurrency cap). The terminal
1406        // error must still reach the caller with all of these calls pending.
1407        const N: usize = 10;
1408
1409        // Fans out N calls, then errors as soon as all N are in flight — exercising a
1410        // terminal failure emitted while the offloaded calls are still pending.
1411        struct FanOutThenError {
1412            all_started: Arc<tokio::sync::Notify>,
1413            n: usize,
1414        }
1415
1416        #[async_trait]
1417        impl Algorithm for FanOutThenError {
1418            fn name(&self) -> &str {
1419                "fan_out_then_error"
1420            }
1421
1422            async fn route(self: Arc<Self>, driver: Driver, request: Request) -> Result<Response> {
1423                let offloads = futures::future::join_all((0..self.n).map(|i| {
1424                    let decision = test_decision(format!("m{i}"));
1425                    driver.call_model(request.clone(), decision)
1426                }));
1427                tokio::select! {
1428                    _ = offloads => Err(test_error("offloads unexpectedly completed")),
1429                    _ = self.all_started.notified() => {
1430                        Err(test_error("terminal error while calls pending"))
1431                    }
1432                }
1433            }
1434        }
1435
1436        let all_started = Arc::new(tokio::sync::Notify::new());
1437        let algo: Arc<dyn Algorithm> = Arc::new(FanOutThenError {
1438            all_started: all_started.clone(),
1439            n: N,
1440        });
1441
1442        // Serving enters each call; once all N are in flight it signals, then pends forever.
1443        let started = Arc::new(AtomicUsize::new(0));
1444        let serve = move |_decision: Decision, _request: Request| {
1445            let started = started.clone();
1446            let all_started = all_started.clone();
1447            async move {
1448                if started.fetch_add(1, Ordering::SeqCst) + 1 == N {
1449                    all_started.notify_one();
1450                }
1451                std::future::pending::<ServeResult>().await
1452            }
1453        };
1454
1455        // With the cap gone, the driver keeps polling the stream even with N calls in
1456        // flight, so the terminal error surfaces promptly instead of hanging.
1457        let run = test_drive(algo, request(), serve);
1458        let result = tokio::time::timeout(std::time::Duration::from_millis(500), run)
1459            .await
1460            .map_err(|error| {
1461                LibsyError::external("waiting for terminal error with full call cap", error)
1462            })?;
1463        match result {
1464            Ok(_) => Err(test_error("expected the terminal error, got a response")),
1465            Err(err) => {
1466                assert!(
1467                    err.to_string()
1468                        .contains("terminal error while calls pending")
1469                );
1470                Ok(())
1471            }
1472        }
1473    }
1474}