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