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