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