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