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