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