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    /// The first target's client that can serve `count_tokens` (an Anthropic
410    /// upstream), or `None` when no target has one. Used by an algorithm's
411    /// [`count_tokens_client`](crate::Algorithm::count_tokens_client).
412    pub fn count_tokens_client(&self) -> Option<Arc<dyn RoutedLlmClient>> {
413        self.targets.iter().find_map(|target| {
414            target
415                .llm_client
416                .as_ref()
417                .filter(|client| client.supports_count_tokens())
418                .cloned()
419        })
420    }
421}
422
423/// Bounds process-local overflow history. Dropping a live session's entry costs one
424/// rediscovered overflow, so the victim choice does not need to be exact.
425const MAX_EVICTION_SESSIONS: usize = 1_024;
426
427/// Per-session record of the targets that overflowed their context window.
428///
429/// A conversation only grows, so a target that could not fit one turn will not fit a
430/// later one; remembering it lets the next turn skip a call certain to fail. Requests
431/// without a session id are not tracked — there is nothing to remember them by.
432#[derive(Default)]
433pub(crate) struct SessionEvictions {
434    sessions: Mutex<HashMap<String, HashSet<String>>>,
435}
436
437impl SessionEvictions {
438    /// The targets `session` has already overflowed; empty for an untracked request.
439    fn evicted_in(&self, session: Option<&str>) -> Vec<String> {
440        let Some(session) = session else {
441            return Vec::new();
442        };
443        self.sessions
444            .lock()
445            .get(session)
446            .map(|targets| targets.iter().cloned().collect())
447            .unwrap_or_default()
448    }
449
450    /// Remembers that `target` overflowed in `session`, tracking at most
451    /// [`MAX_EVICTION_SESSIONS`] sessions.
452    fn record(&self, session: Option<&str>, target: &str) {
453        let Some(session) = session else { return };
454        let mut sessions = self.sessions.lock();
455        if sessions.len() >= MAX_EVICTION_SESSIONS
456            && !sessions.contains_key(session)
457            && let Some(oldest) = sessions.keys().next().cloned()
458        {
459            sessions.remove(&oldest);
460        }
461        sessions
462            .entry(session.to_string())
463            .or_default()
464            .insert(target.to_string());
465    }
466}
467
468/// How many of `targets` this request is still allowed to reach.
469fn eligible_targets(targets: &LlmTargetSet, ctx: &Context) -> usize {
470    targets
471        .targets()
472        .iter()
473        .filter(|t| !ctx.is_excluded(&t.semantic_name))
474        .count()
475}
476
477/// Bars the targets `session` has already overflowed from this request, so routing does
478/// not select one that is certain to fail again.
479pub(crate) fn exclude_evicted(
480    ctx: &mut Context,
481    targets: &LlmTargetSet,
482    evictions: &SessionEvictions,
483    session: Option<&str>,
484) {
485    for target in evictions.evicted_in(session) {
486        // Never seed the pool empty: a later turn may be small enough to serve, and the
487        // caller should get the upstream's answer rather than a routing error.
488        if eligible_targets(targets, ctx) <= 1 {
489            break;
490        }
491        ctx.exclude_target(target);
492    }
493}
494
495/// Calls `target`, falling back to the next eligible target in `targets` whenever one
496/// overflows its context window, until a call succeeds or every target has been tried.
497///
498/// Routing is deliberately not re-run: the fallback replaces the target in place, so the
499/// caller's request-side work and retained state still see exactly one turn.
500/// `fallback_decision` builds the [`Decision`] published for a `from -> to` hop, and each
501/// overflow is recorded against `session` so later turns skip that target outright.
502#[allow(clippy::too_many_arguments)]
503pub(crate) async fn call_llm_with_overflow_fallback(
504    mut ctx: Context,
505    driver: &Driver,
506    targets: &LlmTargetSet,
507    mut target: LlmTarget,
508    mut decision: Arc<dyn Decision>,
509    request: Request,
510    session: Option<&str>,
511    evictions: &SessionEvictions,
512    fallback_decision: impl Fn(&LlmTarget, &LlmTarget) -> Arc<dyn Decision>,
513) -> Result<Response> {
514    loop {
515        let result = driver
516            .call_llm_target(ctx.clone(), &target, request.clone(), decision.clone())
517            .await;
518        let Err(error) = result else { return result };
519        let LibsyError::ClientCall {
520            target: failed,
521            source: LlmClientError::ContextWindowExceeded { .. },
522        } = &error
523        else {
524            return Err(error);
525        };
526        // A target already excluded means the pool is spent; surface the client error
527        // so the caller still sees a context overflow rather than an internal failure.
528        if !ctx.exclude_target(failed) {
529            return Err(error);
530        }
531        evictions.record(session, failed);
532        let Ok(next) = targets.resolve_target(&target.semantic_name, &ctx) else {
533            return Err(error);
534        };
535        decision = fallback_decision(&target, &next);
536        target = next;
537        driver.info(ctx.clone(), decision.clone()).await?;
538    }
539}
540
541/// An optimization strategy. Implement [`create_run_task`](Self::create_run_task);
542/// callers drive it with the provided [`run`](Self::run) (serve calls, get the answer)
543/// or [`run_stream`](Self::run_stream) (drive the [`Step`] stream yourself).
544///
545/// Methods take `self: Arc<Self>`: one algorithm (`Arc<dyn Algorithm>`) is shared across
546/// requests and run concurrently, so it owns its thread-safety and any shared state.
547///
548/// # Concurrency
549///
550/// A host may run the same algorithm concurrently for many requests. Implementations
551/// must synchronize their own mutable shared state. Each call to [`run_stream`](Self::run_stream)
552/// creates an independent [`Driver`], so model-call promises and emitted [`Step`]s cannot
553/// cross between runs.
554///
555/// # Observability
556///
557/// [`run_stream`](Self::run_stream) creates a `libsy.run` span, and each offloaded model
558/// call creates a `libsy.llm_call` span. [`run`](Self::run) additionally wraps calls it
559/// serves through a target's default client in `libsy.client_call`. Decisions and failures
560/// are emitted through `tracing`; metrics use the global OpenTelemetry meter provider.
561#[async_trait]
562pub trait Algorithm: Send + Sync + 'static {
563    /// Stable, low-cardinality name identifying this algorithm — the
564    /// `algorithm` attribute on every span, metric, and log line the crate
565    /// emits for its runs.
566    fn name(&self) -> &str;
567
568    /// Run one request to completion: make model calls with [`Driver::call_llm_target`],
569    /// publish [`Decision`]s with [`Driver::info`], and return the final [`Response`].
570    /// The method an algorithm implements; [`run`](Self::run) / [`run_stream`](Self::run_stream)
571    /// drive it. `ctx` carries the request's cross-cutting values (today: the
572    /// algorithm's telemetry label in [`Context::values`]).
573    async fn create_run_task(
574        self: Arc<Self>,
575        ctx: Context,
576        driver: Driver,
577        request: Request,
578    ) -> Result<Response>;
579
580    /// Feed the algorithm agentic-stack events (tool results, budgets, etc.). The
581    /// reference algorithms ignore signals; a stateful algorithm updates its own
582    /// (interior-mutable) state. Takes `self: Arc<Self>` like the other run methods.
583    #[allow(unused_variables)]
584    async fn process_signals(self: Arc<Self>, signals: Signals) -> Result<()> {
585        Ok(())
586    }
587
588    /// The client [`count_tokens`](Self::count_tokens) forwards to: the first of
589    /// this algorithm's targets whose client can count tokens (an Anthropic
590    /// upstream). The default is `None` — an algorithm with no Anthropic target
591    /// does not support token counting.
592    ///
593    /// CAVEAT: this picks the *first* Anthropic target, not a routed one —
594    /// count_tokens is a direct passthrough, so it does not run the routing
595    /// cascade. For a route with several Anthropic tiers "first" is arbitrary;
596    /// choosing which tier count_tokens should reflect is deferred.
597    fn count_tokens_client(&self) -> Option<Arc<dyn RoutedLlmClient>> {
598        None
599    }
600
601    /// Count the tokens `request` would use — a **direct passthrough** to this
602    /// algorithm's Anthropic target (via
603    /// [`count_tokens_client`](Self::count_tokens_client)), **not** a routed
604    /// call. Token counting is a pre-flight estimate with no routing decision,
605    /// so it deliberately bypasses the classifier cascade (which runs only for
606    /// completions via [`run`](Self::run)). Returns the upstream's JSON
607    /// verbatim. Errors when the algorithm has no Anthropic target.
608    async fn count_tokens(&self, request: Request) -> Result<serde_json::Value> {
609        let client = self
610            .count_tokens_client()
611            .ok_or_else(|| LibsyError::AlgorithmError {
612                message: "no target supports count_tokens (needs an Anthropic upstream)"
613                    .to_string(),
614            })?;
615        client
616            .count_tokens(request)
617            .await
618            .map_err(|source| LibsyError::client_call("count_tokens", source))
619    }
620
621    /// Process a request to completion, returning a stream of [`Step`]s.
622    ///
623    /// The consumer must fulfill every [`Step::CallLlm`] before the algorithm can
624    /// continue. The bounded step channel applies backpressure when the consumer is
625    /// not polling. A successful run ends with [`Step::ReturnToAgent`]; a failure is
626    /// emitted as an `Err` item. Dropping the stream aborts the spawned algorithm task.
627    ///
628    /// Every invocation owns a separate [`Driver`]. `observer`, when present, receives
629    /// each completed model call and, after a successful routed run, its routing overhead.
630    fn run_stream(
631        self: Arc<Self>,
632        ctx: Context,
633        request: Request,
634        observer: Option<RunObserver>,
635    ) -> StepStream {
636        // Stamp the algorithm's telemetry label into the request context; the
637        // context rides on every driver call, so its telemetry is attributed.
638        let mut ctx = ctx;
639        ctx.values.insert(
640            observability::ALGORITHM_KEY.to_string(),
641            self.name().to_string(),
642        );
643        let driver = Driver::with_observer(observer);
644        let task_driver = driver.clone();
645        let task_ctx = ctx.clone();
646        let stream = task_driver.stream();
647        // One `libsy.run` span covers the whole algorithm task; the driver's
648        // `libsy.llm_call` spans and decision logs nest inside it via `tracing`'s
649        // contextual parenting.
650        let span = observability::run_span(self.name(), &request);
651        let observed_driver = task_driver.clone();
652        let handle = tokio::spawn(
653            async move {
654                observability::observe_run(
655                    task_ctx.clone(),
656                    observed_driver,
657                    self.create_run_task(task_ctx, task_driver, request),
658                )
659                .await
660            }
661            .instrument(span),
662        );
663        // Dropping the stream aborts the algorithm task when its consumer goes away.
664        let abort_guard = AbortOnDrop(handle.abort_handle());
665
666        let finish_driver = driver.clone();
667        let finish_ctx = ctx;
668        let tail: StepStream = Box::pin(
669            futures::stream::once(async move {
670                let result = match handle.await {
671                    Ok(response) => response,
672                    Err(source) => Err(LibsyError::AlgorithmTask { source }),
673                };
674                finish_driver.finish(finish_ctx, result).await
675            })
676            .filter_map(|finish_result| async move { finish_result.err().map(Err) }),
677        );
678
679        let stream: StepStream = Box::pin(stream);
680        Box::pin(futures::stream::select(stream, tail).map(move |step| {
681            // link abort guard to stream
682            let _keep_alive = &abort_guard;
683            step
684        }))
685    }
686
687    /// Process a request to completion, returning the final [`Response`] and the trace of
688    /// [`Decision`]s the algorithm made along the way.
689    ///
690    async fn run(
691        self: Arc<Self>,
692        ctx: Context,
693        request: Request,
694    ) -> Result<(Vec<Arc<dyn Decision>>, Response)> {
695        self.run_observed(ctx, request, None).await
696    }
697
698    /// Process a request to completion while reporting each model call to `observer`.
699    async fn run_observed(
700        self: Arc<Self>,
701        ctx: Context,
702        request: Request,
703        observer: Option<RunObserver>,
704    ) -> Result<(Vec<Arc<dyn Decision>>, Response)> {
705        // Serve one offloaded call with its target's default client. A failed *model*
706        // call is forwarded to the algorithm via `respond`; this errors only on an
707        // infrastructure failure (no default client, or the promise was dropped).
708        // `serve` makes the one API call libsy itself performs, so it gets its
709        // own `libsy.client_call` span.
710        #[tracing::instrument(
711            target = "libsy",
712            name = "libsy.client_call",
713            skip_all,
714            fields(
715                algorithm = observability::algorithm_label(&call.get_routed().ctx),
716                switchyard.algorithm = observability::algorithm_label(&call.get_routed().ctx),
717                switchyard.routing.tier = tracing::field::Empty,
718                selected_model = call.get_decision().selected_model(),
719                otel.kind = "client",
720                otel.name = %format_args!("chat {}", call.get_decision().selected_model()),
721                openinference.span.kind = "LLM",
722                gen_ai.operation.name = "chat",
723                gen_ai.request.model = call.get_decision().selected_model(),
724                gen_ai.request.stream = tracing::field::Empty,
725                gen_ai.request.temperature = tracing::field::Empty,
726                gen_ai.request.top_p = tracing::field::Empty,
727                gen_ai.request.top_k = tracing::field::Empty,
728                gen_ai.request.max_tokens = tracing::field::Empty,
729                gen_ai.request.reasoning.level = tracing::field::Empty,
730                gen_ai.output.type = tracing::field::Empty,
731                gen_ai.conversation.id = tracing::field::Empty,
732                server.address = tracing::field::Empty,
733                server.port = tracing::field::Empty,
734                gen_ai.response.id = tracing::field::Empty,
735                gen_ai.response.model = tracing::field::Empty,
736                gen_ai.usage.input_tokens = tracing::field::Empty,
737                gen_ai.usage.output_tokens = tracing::field::Empty,
738                gen_ai.usage.cache_read.input_tokens = tracing::field::Empty,
739                gen_ai.usage.cache_creation.input_tokens = tracing::field::Empty,
740                gen_ai.usage.reasoning.output_tokens = tracing::field::Empty,
741                outcome = tracing::field::Empty,
742                otel.status_code = tracing::field::Empty,
743                error.type = tracing::field::Empty,
744                error = tracing::field::Empty,
745            )
746        )]
747        async fn serve(call: CallLlmRequest) -> Result<()> {
748            let span = tracing::Span::current();
749            observability::record_gen_ai_request(&span, &call.get_routed().request.llm_request);
750            if let Some(tier) = call.get_decision().routing_tier() {
751                span.record("switchyard.routing.tier", tier);
752            }
753            if let Some(session_id) = call
754                .get_routed()
755                .request
756                .metadata
757                .as_ref()
758                .and_then(|metadata| metadata.session_id.as_deref())
759            {
760                span.record("gen_ai.conversation.id", session_id);
761            }
762            let routed = call.get_routed().clone();
763            let target = routed.decision.selected_model().to_string();
764            let client =
765                routed
766                    .default_client
767                    .clone()
768                    .ok_or_else(|| LibsyError::MissingClient {
769                        target: target.clone(),
770                    })?;
771            let result = client
772                .call(routed.ctx, routed.request, routed.decision)
773                .await
774                .map_err(|source| LibsyError::client_call(target, source));
775            let result = observability::observe_client_call(result);
776            call.respond(result)
777        }
778
779        let stream = self.run_stream(ctx, request, observer);
780        tokio::pin!(stream);
781
782        let mut trace: Vec<Arc<dyn Decision>> = Vec::new();
783        let mut in_flight = futures::stream::FuturesUnordered::new();
784        let mut final_response: Option<Response> = None;
785
786        loop {
787            tokio::select! {
788                Some(result) = in_flight.next() => match result {
789                    Ok(()) => {}, // CallLlm completed successfully
790                    Err(err) => return Err(err), // CallLlm failed, propagate the error
791                },
792                step = stream.next() => {
793                    match step {
794                        None => break, // stream has ended, no more steps
795                        Some(item) => match item? {
796                            Step::CallLlm(call) => in_flight.push(serve(*call)),
797                            Step::Decision(decision) => trace.push(decision),
798                            Step::ReturnToAgent(response) => {
799                                final_response = Some(*response);
800                                break;
801                            }
802                        }
803                    }
804                },
805            }
806        }
807        final_response
808            .map(|response| (trace, response))
809            .ok_or(LibsyError::MissingFinalResponse)
810    }
811}
812
813#[cfg(test)]
814mod tests {
815    use super::*;
816    use crate::text::{completion_text, text_request, text_response};
817    use futures::StreamExt;
818    use switchyard_protocol::{LlmResponse, LlmResponseChunk};
819
820    #[derive(Debug, thiserror::Error)]
821    #[error("{0}")]
822    struct TestError(&'static str);
823
824    fn test_error(message: &'static str) -> LibsyError {
825        LibsyError::external("test", TestError(message))
826    }
827
828    /// Mock client that echoes back the target name it was called with.
829    struct EchoClient;
830
831    #[async_trait]
832    impl RoutedLlmClient for EchoClient {
833        async fn call(
834            &self,
835            _ctx: Context,
836            _request: Request,
837            decision: Arc<dyn Decision>,
838        ) -> std::result::Result<Response, LlmClientError> {
839            // Echo back the model the algorithm routed to (the decision's selection).
840            Ok(Response {
841                llm_response: LlmResponse::Agg(text_response(
842                    None,
843                    decision.selected_model().to_string(),
844                )),
845                metadata: None,
846            })
847        }
848    }
849
850    /// Trivial decision + algo used only to exercise the orchestrator: calls the
851    /// first target and returns its response with a one-item trace.
852    struct TestDecision {
853        model: String,
854    }
855
856    impl Decision for TestDecision {
857        fn selected_model(&self) -> &str {
858            &self.model
859        }
860        fn reasoning(&self) -> Option<&str> {
861            None
862        }
863        fn as_any(&self) -> &dyn std::any::Any {
864            self
865        }
866    }
867
868    struct TestAlgo {
869        target_set: LlmTargetSet,
870    }
871
872    #[async_trait]
873    impl Algorithm for TestAlgo {
874        fn name(&self) -> &str {
875            "test"
876        }
877
878        async fn create_run_task(
879            self: Arc<Self>,
880            ctx: Context,
881            driver: Driver,
882            request: Request,
883        ) -> Result<Response> {
884            let target = self
885                .target_set
886                .targets()
887                .first()
888                .ok_or(LibsyError::NoTargets)?
889                .clone();
890            let decision: Arc<dyn Decision> = Arc::new(TestDecision {
891                model: target.semantic_name.clone(),
892            });
893            driver.info(ctx.clone(), decision.clone()).await?;
894            driver
895                .call_llm_target(ctx, &target, request, decision)
896                .await
897        }
898    }
899
900    /// Build a shared `TestAlgo` over the given target set.
901    fn orch(target_set: LlmTargetSet) -> Arc<dyn Algorithm> {
902        Arc::new(TestAlgo { target_set })
903    }
904
905    fn request() -> Request {
906        Request {
907            llm_request: text_request(Some("auto".to_string()), "hi".to_string()),
908            raw_request: None,
909            metadata: None,
910        }
911    }
912
913    /// `(name, has_client)` — `has_client: false` builds a target with no default client.
914    fn target_set(names: &[(&str, bool)]) -> LlmTargetSet {
915        let targets = names
916            .iter()
917            .map(|(name, has_client)| LlmTarget {
918                semantic_name: name.to_string(),
919                llm_client: has_client.then(|| Arc::new(EchoClient) as Arc<dyn RoutedLlmClient>),
920            })
921            .collect();
922        LlmTargetSet::new(targets)
923    }
924
925    #[tokio::test]
926    async fn observed_run_reports_one_successful_routed_call() -> Result<()> {
927        let observations = Arc::new(Mutex::new(Vec::new()));
928        let observed = observations.clone();
929        let observer: RunObserver = Arc::new(move |observation| observed.lock().push(observation));
930        let (_, response) = orch(target_set(&[("direct/model", true)]))
931            .run_observed(Context::default(), request(), Some(observer))
932            .await?;
933        assert_eq!(
934            response.llm_response.as_agg().map(completion_text),
935            Some("direct/model".to_string())
936        );
937        let observations = observations.lock();
938        assert_eq!(observations.len(), 2);
939        let RunObservation::LlmCall(observation) = &observations[0] else {
940            return Err(test_error("expected an LLM call observation"));
941        };
942        assert_eq!(observation.selected_model, "direct/model");
943        assert!(observation.is_routed);
944        assert!(observation.is_success);
945        assert!(observation.usage.is_some());
946        assert!(matches!(
947            observations[1],
948            RunObservation::RoutingOverhead(_)
949        ));
950        Ok(())
951    }
952
953    #[test]
954    fn target_lookup_returns_the_missing_target() {
955        let error = target_set(&[]).get_target("missing").err();
956        assert!(matches!(
957            error,
958            Some(LibsyError::TargetNotFound { target }) if target == "missing"
959        ));
960    }
961
962    /// Client that serves a call as a token stream — its `call` returns
963    /// [`LlmResponse::Stream`] replaying `chunks` in order (as `Ok` items).
964    struct StreamingClient {
965        chunks: Vec<LlmResponseChunk>,
966    }
967
968    #[async_trait]
969    impl RoutedLlmClient for StreamingClient {
970        async fn call(
971            &self,
972            _ctx: Context,
973            _request: Request,
974            _decision: Arc<dyn Decision>,
975        ) -> std::result::Result<Response, LlmClientError> {
976            let stream = futures::stream::iter(
977                self.chunks
978                    .clone()
979                    .into_iter()
980                    .map(|chunk| Ok(chunk.into())),
981            )
982            .boxed();
983            Ok(Response {
984                llm_response: LlmResponse::Stream(stream),
985                metadata: None,
986            })
987        }
988    }
989
990    /// Build a single-target algo whose one target streams `chunks`.
991    fn streaming_orch(chunks: Vec<LlmResponseChunk>) -> Arc<dyn Algorithm> {
992        let target = LlmTarget {
993            semantic_name: "stream/model".to_string(),
994            llm_client: Some(Arc::new(StreamingClient { chunks }) as Arc<dyn RoutedLlmClient>),
995        };
996        orch(LlmTargetSet::new(vec![target]))
997    }
998
999    #[tokio::test]
1000    async fn run_returns_a_streamed_response_the_caller_aggregates() -> Result<()> {
1001        // A streaming client -> its chunks flow through the promise and `ReturnToAgent`,
1002        // and `run` returns the live stream untouched for the caller to fold.
1003        let orch = streaming_orch(vec![
1004            LlmResponseChunk::MessageStart {
1005                id: Some("m1".to_string()),
1006                model: Some("stream/model".to_string()),
1007            },
1008            LlmResponseChunk::TextDelta {
1009                index: 0,
1010                text: "hel".to_string(),
1011            },
1012            LlmResponseChunk::TextDelta {
1013                index: 0,
1014                text: "lo".to_string(),
1015            },
1016            LlmResponseChunk::MessageStop {
1017                reason: Some("stop".to_string()),
1018            },
1019        ]);
1020        let (trace, response) = orch.run(Context::default(), request()).await?;
1021        // `run` handed back the live stream; the caller folds it to a buffered aggregate.
1022        let agg = response
1023            .llm_response
1024            .into_agg()
1025            .await
1026            .map_err(|error| LibsyError::external("aggregating response stream", error))?;
1027        assert_eq!(completion_text(&agg), "hello");
1028        assert_eq!(agg.model.as_deref(), Some("stream/model"));
1029        assert_eq!(trace.len(), 1);
1030        Ok(())
1031    }
1032
1033    #[tokio::test]
1034    async fn aggregating_a_streamed_response_propagates_a_mid_stream_error() -> Result<()> {
1035        // `run` succeeds and returns the stream; the in-band `Error` chunk surfaces only
1036        // when the caller aggregates it.
1037        let orch = streaming_orch(vec![
1038            LlmResponseChunk::TextDelta {
1039                index: 0,
1040                text: "partial".to_string(),
1041            },
1042            LlmResponseChunk::StreamError {
1043                message: "upstream exploded".to_string(),
1044            },
1045        ]);
1046        let (_, response) = orch.run(Context::default(), request()).await?;
1047        match response.llm_response.into_agg().await {
1048            Ok(_) => panic!("expected a mid-stream error, got an aggregate"),
1049            Err(err) => {
1050                assert!(err.to_string().contains("upstream exploded"));
1051                Ok(())
1052            }
1053        }
1054    }
1055
1056    #[tokio::test]
1057    async fn run_offloads_via_promise_then_returns_to_agent() -> Result<()> {
1058        // A client-less target -> its call is offloaded via a promise the
1059        // orchestrator surfaces as a `CallLlm` step for us to fulfill.
1060        let stream = orch(target_set(&[("offload/model", false)])).run_stream(
1061            Context::default(),
1062            request(),
1063            None,
1064        );
1065        tokio::pin!(stream);
1066
1067        let mut saw_call = false;
1068        let mut final_completion = None;
1069        while let Some(step) = stream.next().await {
1070            match step? {
1071                Step::CallLlm(call) => {
1072                    saw_call = true;
1073                    // The decision rode along with the promise.
1074                    assert_eq!(call.get_decision().selected_model(), "offload/model");
1075                    // Fulfilling the promise is the "real" model call the caller makes.
1076                    call.respond(Ok(Response {
1077                        llm_response: LlmResponse::Agg(text_response(
1078                            None,
1079                            "fulfilled".to_string(),
1080                        )),
1081                        metadata: None,
1082                    }))?;
1083                }
1084                Step::Decision(decision) => {
1085                    assert_eq!(decision.selected_model(), "offload/model");
1086                }
1087                Step::ReturnToAgent(response) => {
1088                    final_completion = Some(
1089                        response
1090                            .llm_response
1091                            .as_agg()
1092                            .map(completion_text)
1093                            .unwrap_or_default(),
1094                    );
1095                }
1096            }
1097        }
1098
1099        assert!(saw_call, "expected a CallLlm step before ReturnToAgent");
1100        assert_eq!(
1101            final_completion.ok_or_else(|| test_error("no ReturnToAgent step"))?,
1102            "fulfilled"
1103        );
1104        Ok(())
1105    }
1106
1107    #[tokio::test]
1108    async fn client_backed_target_offloads_with_a_default_client() -> Result<()> {
1109        // Every call now offloads to the stream; a client-backed target rides its
1110        // client along as `default_client` so the consumer can serve it by default.
1111        let stream = orch(target_set(&[("direct/model", true)])).run_stream(
1112            Context::default(),
1113            request(),
1114            None,
1115        );
1116        tokio::pin!(stream);
1117
1118        let mut final_completion = None;
1119        while let Some(step) = stream.next().await {
1120            match step? {
1121                Step::CallLlm(call) => {
1122                    let routed = call.get_routed().clone();
1123                    let client = routed
1124                        .default_client
1125                        .clone()
1126                        .ok_or_else(|| test_error("expected a default client"))?;
1127                    let target = routed.decision.selected_model().to_string();
1128                    let result = client
1129                        .call(routed.ctx, routed.request, routed.decision)
1130                        .await
1131                        .map_err(|error| LibsyError::client_call(target, error));
1132                    call.respond(result)?;
1133                }
1134                Step::Decision(_) => {}
1135                Step::ReturnToAgent(response) => {
1136                    final_completion = Some(
1137                        response
1138                            .llm_response
1139                            .as_agg()
1140                            .map(completion_text)
1141                            .unwrap_or_default(),
1142                    );
1143                }
1144            }
1145        }
1146
1147        // EchoClient echoes the model name back as the completion.
1148        assert_eq!(
1149            final_completion.ok_or_else(|| test_error("no ReturnToAgent"))?,
1150            "direct/model"
1151        );
1152        Ok(())
1153    }
1154
1155    #[tokio::test]
1156    async fn run_returns_the_response_when_all_targets_have_clients() -> Result<()> {
1157        // Every target has a client, so run serves every call via the
1158        // default client and returns the trace + final response.
1159        let (trace, response) = orch(target_set(&[("direct/model", true)]))
1160            .run(Context::default(), request())
1161            .await?;
1162        // TestAlgo calls the first target; EchoClient echoes its name.
1163        assert_eq!(
1164            response
1165                .llm_response
1166                .as_agg()
1167                .map(completion_text)
1168                .unwrap_or_default(),
1169            "direct/model"
1170        );
1171        assert_eq!(trace[0].selected_model(), "direct/model");
1172        Ok(())
1173    }
1174
1175    #[tokio::test]
1176    async fn run_errors_when_a_target_lacks_a_client() -> Result<()> {
1177        // A client-less target has no default client to serve its offloaded call, so
1178        // driving it to completion errors.
1179        let error = orch(target_set(&[("offload/model", false)]))
1180            .run(Context::default(), request())
1181            .await
1182            .err()
1183            .ok_or_else(|| test_error("expected a missing-client error"))?;
1184        assert!(matches!(
1185            error,
1186            LibsyError::MissingClient { target } if target == "offload/model"
1187        ));
1188        Ok(())
1189    }
1190
1191    #[tokio::test(flavor = "multi_thread", worker_threads = 12)]
1192    async fn requests_are_processed_in_parallel() -> Result<()> {
1193        use std::time::Duration;
1194        use tokio::sync::Barrier;
1195
1196        const N: usize = 12;
1197
1198        // A client that blocks until all N concurrent calls have arrived. If
1199        // requests were serialized (one algorithm behind a `Mutex`), only one
1200        // call could be in flight, the barrier would never reach N, and the test
1201        // would time out. It passes only because the shared algorithm is driven
1202        // concurrently across requests.
1203        struct BarrierClient {
1204            barrier: Arc<Barrier>,
1205        }
1206
1207        #[async_trait]
1208        impl RoutedLlmClient for BarrierClient {
1209            async fn call(
1210                &self,
1211                _ctx: Context,
1212                _request: Request,
1213                decision: Arc<dyn Decision>,
1214            ) -> std::result::Result<Response, LlmClientError> {
1215                self.barrier.wait().await;
1216                Ok(Response {
1217                    llm_response: LlmResponse::Agg(text_response(
1218                        None,
1219                        decision.selected_model().to_string(),
1220                    )),
1221                    metadata: None,
1222                })
1223            }
1224        }
1225
1226        let barrier = Arc::new(Barrier::new(N));
1227        let targets = LlmTargetSet::new(vec![LlmTarget {
1228            semantic_name: "m".to_string(),
1229            llm_client: Some(Arc::new(BarrierClient {
1230                barrier: barrier.clone(),
1231            })),
1232        }]);
1233        // One shared algorithm driven by many concurrent requests.
1234        let algo = orch(targets);
1235
1236        let mut handles = Vec::new();
1237        for _ in 0..N {
1238            let algo = algo.clone();
1239            handles.push(tokio::spawn(async move {
1240                algo.run(Context::default(), request())
1241                    .await
1242                    .map(|(_, response)| {
1243                        response
1244                            .llm_response
1245                            .as_agg()
1246                            .map(completion_text)
1247                            .unwrap_or_default()
1248                    })
1249            }));
1250        }
1251
1252        for handle in handles {
1253            // The timeout turns a serialization deadlock into a failure, not a hang.
1254            let completion = tokio::time::timeout(Duration::from_secs(5), handle)
1255                .await
1256                .map_err(|error| LibsyError::external("waiting for test task", error))?
1257                .map_err(|source| LibsyError::AlgorithmTask { source })??;
1258            assert_eq!(completion, "m");
1259        }
1260        Ok(())
1261    }
1262
1263    #[tokio::test]
1264    async fn offload_error_propagates_back_to_the_algorithm() -> Result<()> {
1265        // A client-less target offloads its call; we fulfill the promise with an
1266        // Err, which must flow back through `call_llm_target` into the algorithm and
1267        // out as an error step — not a response.
1268        let stream = orch(target_set(&[("offload/model", false)])).run_stream(
1269            Context::default(),
1270            request(),
1271            None,
1272        );
1273        tokio::pin!(stream);
1274
1275        let mut saw_error = false;
1276        while let Some(step) = stream.next().await {
1277            match step {
1278                Ok(Step::CallLlm(call)) => {
1279                    call.respond(Err(test_error("upstream model call failed")))?;
1280                }
1281                Ok(Step::Decision(_)) => {}
1282                Ok(Step::ReturnToAgent(..)) => {
1283                    return Err(test_error(
1284                        "expected the offload error to propagate, got a response",
1285                    ));
1286                }
1287                Err(err) => {
1288                    // The algorithm's `call_llm_target` saw the error via the promise.
1289                    assert!(err.to_string().contains("upstream model call failed"));
1290                    saw_error = true;
1291                }
1292            }
1293        }
1294
1295        assert!(saw_error, "expected an error step");
1296        Ok(())
1297    }
1298
1299    #[tokio::test]
1300    async fn dropping_the_stream_cancels_the_algorithm_task() -> Result<()> {
1301        use std::sync::atomic::{AtomicBool, Ordering};
1302        use std::time::Duration;
1303        use tokio::sync::mpsc;
1304
1305        // Sets a flag when dropped, so we can observe whether the algorithm task was
1306        // cancelled/dropped.
1307        struct DropGuard(Arc<AtomicBool>);
1308        impl Drop for DropGuard {
1309            fn drop(&mut self) {
1310                self.0.store(true, Ordering::SeqCst);
1311            }
1312        }
1313
1314        struct StuckAlgo {
1315            started: mpsc::UnboundedSender<()>,
1316            dropped: Arc<AtomicBool>,
1317        }
1318
1319        #[async_trait]
1320        impl Algorithm for StuckAlgo {
1321            fn name(&self) -> &str {
1322                "stuck"
1323            }
1324
1325            async fn create_run_task(
1326                self: Arc<Self>,
1327                _ctx: Context,
1328                _driver: Driver,
1329                _request: Request,
1330            ) -> Result<Response> {
1331                let _guard = DropGuard(self.dropped.clone());
1332                let _ = self.started.send(());
1333                // Await forever without ever touching the driver.
1334                std::future::pending::<()>().await;
1335                unreachable!()
1336            }
1337        }
1338
1339        let (started_tx, mut started_rx) = mpsc::unbounded_channel();
1340        let dropped = Arc::new(AtomicBool::new(false));
1341        let algo: Arc<dyn Algorithm> = Arc::new(StuckAlgo {
1342            started: started_tx,
1343            dropped: dropped.clone(),
1344        });
1345
1346        let stream = algo.run_stream(Context::default(), request(), None);
1347        started_rx
1348            .recv()
1349            .await
1350            .ok_or_else(|| test_error("task never started"))?;
1351        drop(stream);
1352        tokio::time::sleep(Duration::from_millis(100)).await;
1353
1354        assert!(
1355            dropped.load(Ordering::SeqCst),
1356            "algorithm task was NOT cancelled after dropping the stream"
1357        );
1358        Ok(())
1359    }
1360
1361    #[tokio::test]
1362    async fn create_run_task_panic_surfaces_as_a_stream_error() -> Result<()> {
1363        // An algorithm whose task panics must surface an `Err` step to the stream
1364        // consumer, not abort the process from an unobserved detached task.
1365        struct Panicky;
1366
1367        #[async_trait]
1368        impl Algorithm for Panicky {
1369            fn name(&self) -> &str {
1370                "panicky"
1371            }
1372
1373            async fn create_run_task(
1374                self: Arc<Self>,
1375                _ctx: Context,
1376                _driver: Driver,
1377                _request: Request,
1378            ) -> Result<Response> {
1379                panic!("boom");
1380            }
1381        }
1382
1383        let algo: Arc<dyn Algorithm> = Arc::new(Panicky);
1384        let stream = algo.run_stream(Context::default(), request(), None);
1385        tokio::pin!(stream);
1386
1387        let mut saw_error = false;
1388        while let Some(step) = stream.next().await {
1389            match step {
1390                Err(err) => {
1391                    assert!(matches!(err, LibsyError::AlgorithmTask { .. }));
1392                    saw_error = true;
1393                }
1394                Ok(_) => return Err(test_error("expected the panic to surface as an error step")),
1395            }
1396        }
1397
1398        assert!(saw_error, "expected an error step from the panicked task");
1399        Ok(())
1400    }
1401
1402    #[tokio::test]
1403    async fn run_returns_an_error_when_the_algorithm_task_panics() -> Result<()> {
1404        // The panic surfaces as an `Err` step inside `run_stream`; `run` propagates it
1405        // via `?`, so the caller gets an `Err` rather than a hang or a silent panic.
1406        struct Panicky;
1407
1408        #[async_trait]
1409        impl Algorithm for Panicky {
1410            fn name(&self) -> &str {
1411                "panicky"
1412            }
1413
1414            async fn create_run_task(
1415                self: Arc<Self>,
1416                _ctx: Context,
1417                _driver: Driver,
1418                _request: Request,
1419            ) -> Result<Response> {
1420                panic!("boom");
1421            }
1422        }
1423
1424        let algo: Arc<dyn Algorithm> = Arc::new(Panicky);
1425        match algo.run(Context::default(), request()).await {
1426            Ok(_) => Err(test_error(
1427                "expected run to surface the algorithm panic as an error",
1428            )),
1429            Err(err) => {
1430                assert!(matches!(err, LibsyError::AlgorithmTask { .. }));
1431                Ok(())
1432            }
1433        }
1434    }
1435
1436    #[tokio::test]
1437    async fn cancelling_run_cancels_the_algorithm_task() -> Result<()> {
1438        use std::sync::atomic::{AtomicBool, Ordering};
1439        use std::time::Duration;
1440        use tokio::sync::mpsc;
1441
1442        // Sets a flag when dropped, so we can observe whether the algorithm task was
1443        // cancelled once the `run` future driving it is dropped.
1444        struct DropGuard(Arc<AtomicBool>);
1445        impl Drop for DropGuard {
1446            fn drop(&mut self) {
1447                self.0.store(true, Ordering::SeqCst);
1448            }
1449        }
1450
1451        struct StuckAlgo {
1452            started: mpsc::UnboundedSender<()>,
1453            dropped: Arc<AtomicBool>,
1454        }
1455
1456        #[async_trait]
1457        impl Algorithm for StuckAlgo {
1458            fn name(&self) -> &str {
1459                "stuck"
1460            }
1461
1462            async fn create_run_task(
1463                self: Arc<Self>,
1464                _ctx: Context,
1465                _driver: Driver,
1466                _request: Request,
1467            ) -> Result<Response> {
1468                let _guard = DropGuard(self.dropped.clone());
1469                let _ = self.started.send(());
1470                // Hang forever without ever touching the driver, so only cancellation
1471                // (not a dropped step channel) can stop this task.
1472                std::future::pending::<()>().await;
1473                unreachable!()
1474            }
1475        }
1476
1477        let (started_tx, mut started_rx) = mpsc::unbounded_channel();
1478        let dropped = Arc::new(AtomicBool::new(false));
1479        let algo: Arc<dyn Algorithm> = Arc::new(StuckAlgo {
1480            started: started_tx,
1481            dropped: dropped.clone(),
1482        });
1483
1484        // Drive `run` on its own task, wait until the algorithm task is up, then cancel
1485        // `run` — dropping its future (and the `run_stream` stream it holds).
1486        let run_task = tokio::spawn(async move { algo.run(Context::default(), request()).await });
1487        started_rx
1488            .recv()
1489            .await
1490            .ok_or_else(|| test_error("task never started"))?;
1491        run_task.abort();
1492        tokio::time::sleep(Duration::from_millis(100)).await;
1493
1494        assert!(
1495            dropped.load(Ordering::SeqCst),
1496            "algorithm task was NOT cancelled after cancelling run"
1497        );
1498        Ok(())
1499    }
1500
1501    // --- first-wins hedging: `run` must not wait on losing speculative calls -------------
1502
1503    /// The loser: signals it has started serving (so the winner can then win with the
1504    /// loser's serve guaranteed in flight), then finishes late (`Some(delay)`) or never
1505    /// (`None`).
1506    struct LoserClient {
1507        started: Arc<tokio::sync::Notify>,
1508        delay: Option<std::time::Duration>,
1509    }
1510
1511    #[async_trait]
1512    impl RoutedLlmClient for LoserClient {
1513        async fn call(
1514            &self,
1515            _ctx: Context,
1516            _request: Request,
1517            decision: Arc<dyn Decision>,
1518        ) -> std::result::Result<Response, LlmClientError> {
1519            self.started.notify_one();
1520            match self.delay {
1521                Some(delay) => tokio::time::sleep(delay).await,
1522                None => std::future::pending::<()>().await,
1523            }
1524            Ok(Response {
1525                llm_response: LlmResponse::Agg(text_response(
1526                    None,
1527                    decision.selected_model().to_string(),
1528                )),
1529                metadata: None,
1530            })
1531        }
1532    }
1533
1534    /// The winner: waits until the loser's serve has started, then echoes immediately, so
1535    /// the loser's serve is guaranteed in flight when the winner wins.
1536    struct GatedEchoClient {
1537        gate: Arc<tokio::sync::Notify>,
1538    }
1539
1540    #[async_trait]
1541    impl RoutedLlmClient for GatedEchoClient {
1542        async fn call(
1543            &self,
1544            _ctx: Context,
1545            _request: Request,
1546            decision: Arc<dyn Decision>,
1547        ) -> std::result::Result<Response, LlmClientError> {
1548            self.gate.notified().await;
1549            Ok(Response {
1550                llm_response: LlmResponse::Agg(text_response(
1551                    None,
1552                    decision.selected_model().to_string(),
1553                )),
1554                metadata: None,
1555            })
1556        }
1557    }
1558
1559    /// Offloads two targets concurrently and returns the first to resolve, dropping the
1560    /// loser's call (first-wins hedging).
1561    struct Hedge {
1562        winner: LlmTarget,
1563        loser: LlmTarget,
1564    }
1565
1566    #[async_trait]
1567    impl Algorithm for Hedge {
1568        fn name(&self) -> &str {
1569            "hedge"
1570        }
1571
1572        async fn create_run_task(
1573            self: Arc<Self>,
1574            ctx: Context,
1575            driver: Driver,
1576            request: Request,
1577        ) -> Result<Response> {
1578            let dec_w: Arc<dyn Decision> = Arc::new(TestDecision {
1579                model: self.winner.semantic_name.clone(),
1580            });
1581            let dec_l: Arc<dyn Decision> = Arc::new(TestDecision {
1582                model: self.loser.semantic_name.clone(),
1583            });
1584            let win = driver.call_llm_target(ctx.clone(), &self.winner, request.clone(), dec_w);
1585            let lose = driver.call_llm_target(ctx, &self.loser, request, dec_l);
1586            // First to resolve wins; `select!` drops the losing future (and its promise).
1587            tokio::select! {
1588                res = win => res,
1589                res = lose => res,
1590            }
1591        }
1592    }
1593
1594    /// Builds a hedging algo whose winner is gated behind the loser starting, and whose
1595    /// loser finishes after `loser_delay` (or never, when `None`).
1596    fn hedge(loser_delay: Option<std::time::Duration>) -> Arc<dyn Algorithm> {
1597        let started = Arc::new(tokio::sync::Notify::new());
1598        let winner = LlmTarget {
1599            semantic_name: "winner".to_string(),
1600            llm_client: Some(Arc::new(GatedEchoClient {
1601                gate: started.clone(),
1602            })),
1603        };
1604        let loser = LlmTarget {
1605            semantic_name: "loser".to_string(),
1606            llm_client: Some(Arc::new(LoserClient {
1607                started,
1608                delay: loser_delay,
1609            })),
1610        };
1611        Arc::new(Hedge { winner, loser })
1612    }
1613
1614    #[tokio::test]
1615    async fn run_returns_the_winner_without_a_late_loser_overwriting_it() -> Result<()> {
1616        // The loser responds 50ms after the winner has already won. `run` must return the
1617        // winner, not the loser's `respond`-to-a-dropped-receiver error.
1618        let (_trace, response) = hedge(Some(std::time::Duration::from_millis(50)))
1619            .run(Context::default(), request())
1620            .await?;
1621        assert_eq!(
1622            response
1623                .llm_response
1624                .as_agg()
1625                .map(completion_text)
1626                .unwrap_or_default(),
1627            "winner"
1628        );
1629        Ok(())
1630    }
1631
1632    #[tokio::test]
1633    async fn run_returns_the_winner_without_hanging_on_a_pending_loser() -> Result<()> {
1634        // The loser never resolves. `run` must return the winner promptly, not hang
1635        // waiting for the in-flight loser.
1636        let run = hedge(None).run(Context::default(), request());
1637        let (_trace, response) = tokio::time::timeout(std::time::Duration::from_secs(1), run)
1638            .await
1639            .map_err(|error| LibsyError::external("waiting for pending loser", error))??;
1640        assert_eq!(
1641            response
1642                .llm_response
1643                .as_agg()
1644                .map(completion_text)
1645                .unwrap_or_default(),
1646            "winner"
1647        );
1648        Ok(())
1649    }
1650
1651    #[tokio::test]
1652    async fn run_surfaces_a_terminal_error_with_many_calls_in_flight() -> Result<()> {
1653        use std::sync::atomic::{AtomicUsize, Ordering};
1654
1655        // A large fan-out (10 matched the old, now-removed concurrency cap). The terminal
1656        // error must still reach the caller with all of these calls pending.
1657        const N: usize = 10;
1658
1659        // Enters each call; once all N are in flight, signals, then pends forever.
1660        struct EnterThenPend {
1661            started: Arc<AtomicUsize>,
1662            all_started: Arc<tokio::sync::Notify>,
1663            n: usize,
1664        }
1665
1666        #[async_trait]
1667        impl RoutedLlmClient for EnterThenPend {
1668            async fn call(
1669                &self,
1670                _ctx: Context,
1671                _request: Request,
1672                _decision: Arc<dyn Decision>,
1673            ) -> std::result::Result<Response, LlmClientError> {
1674                if self.started.fetch_add(1, Ordering::SeqCst) + 1 == self.n {
1675                    self.all_started.notify_one();
1676                }
1677                std::future::pending::<()>().await;
1678                unreachable!()
1679            }
1680        }
1681
1682        // Fans out N calls, then errors as soon as all N are in flight — exercising a
1683        // terminal failure emitted while the offloaded calls are still pending.
1684        struct FanOutThenError {
1685            target: LlmTarget,
1686            all_started: Arc<tokio::sync::Notify>,
1687            n: usize,
1688        }
1689
1690        #[async_trait]
1691        impl Algorithm for FanOutThenError {
1692            fn name(&self) -> &str {
1693                "fan_out_then_error"
1694            }
1695
1696            async fn create_run_task(
1697                self: Arc<Self>,
1698                ctx: Context,
1699                driver: Driver,
1700                request: Request,
1701            ) -> Result<Response> {
1702                let offloads = futures::future::join_all((0..self.n).map(|i| {
1703                    let decision: Arc<dyn Decision> = Arc::new(TestDecision {
1704                        model: format!("m{i}"),
1705                    });
1706                    driver.call_llm_target(ctx.clone(), &self.target, request.clone(), decision)
1707                }));
1708                tokio::select! {
1709                    _ = offloads => Err(test_error("offloads unexpectedly completed")),
1710                    _ = self.all_started.notified() => {
1711                        Err(test_error("terminal error while calls pending"))
1712                    }
1713                }
1714            }
1715        }
1716
1717        let all_started = Arc::new(tokio::sync::Notify::new());
1718        let target = LlmTarget {
1719            semantic_name: "pending".to_string(),
1720            llm_client: Some(Arc::new(EnterThenPend {
1721                started: Arc::new(AtomicUsize::new(0)),
1722                all_started: all_started.clone(),
1723                n: N,
1724            })),
1725        };
1726        let algo: Arc<dyn Algorithm> = Arc::new(FanOutThenError {
1727            target,
1728            all_started,
1729            n: N,
1730        });
1731
1732        // With the cap gone, `run` keeps polling the stream even with N calls in flight, so
1733        // the terminal error surfaces promptly instead of hanging.
1734        let run = algo.run(Context::default(), request());
1735        let result = tokio::time::timeout(std::time::Duration::from_millis(500), run)
1736            .await
1737            .map_err(|error| {
1738                LibsyError::external("waiting for terminal error with full call cap", error)
1739            })?;
1740        match result {
1741            Ok(_) => Err(test_error("expected the terminal error, got a response")),
1742            Err(err) => {
1743                assert!(
1744                    err.to_string()
1745                        .contains("terminal error while calls pending")
1746                );
1747                Ok(())
1748            }
1749        }
1750    }
1751}