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//! Routing algorithms request external work through [`Driver`] and receive typed replies.
5//! Hosts execute that work, keeping routing policy independent of transport.
6
7use std::{
8    collections::HashMap, future::Future, panic::AssertUnwindSafe, pin::Pin, sync::Arc,
9    time::Instant,
10};
11
12use async_trait::async_trait;
13use futures::{FutureExt, Stream, StreamExt};
14use parking_lot::Mutex;
15use serde_json::Value;
16use tokio::sync::{mpsc, oneshot};
17use tokio_stream::wrappers::ReceiverStream;
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    Category, DecisionRequest, DecisionResponse, ModelId, Request, Response,
29};
30
31use crate::{DriverError, LibsyError, Result, State, observability};
32
33/// A boxed, `Send` stream of [`Step`]s — the output of
34/// [`Algorithm::run_stream`]. Boxed so the trait method that produces it keeps
35/// `Arc<dyn Algorithm>` object-safe.
36pub type StepStream = Pin<Box<dyn Stream<Item = Result<Step>> + Send>>;
37
38/// The models one algorithm run may use, grouped by [`Category`]. Within a
39/// category they are ordered best-first.
40///
41/// Delegated sub-agent work gets its own groups, reachable only through
42/// [`Driver::for_subagent`]. Keeping them separate is what stops a sub-agent's
43/// `capable` from resolving to the parent's, and stops the parent falling back
44/// onto a model only its sub-agents were given.
45///
46/// One run's driver clones all read the same value, so it is passed as
47/// `Arc<RuntimeModels>` rather than cloned per driver.
48#[derive(Clone, Debug, Default)]
49pub struct RuntimeModels {
50    /// When using subagents this is the parent agent category.
51    by_category: HashMap<Category, Vec<ModelId>>,
52    subagent: Option<HashMap<Category, Vec<ModelId>>>,
53}
54
55impl RuntimeModels {
56    /// The models available to the algorithm itself.
57    pub fn new(by_category: HashMap<Category, Vec<ModelId>>) -> Self {
58        Self {
59            by_category,
60            subagent: None,
61        }
62    }
63
64    /// Adds the groups used for delegated sub-agent work.
65    pub fn with_subagent(mut self, models: HashMap<Category, Vec<ModelId>>) -> Self {
66        self.subagent = Some(models);
67        self
68    }
69
70    /// The models in `category`, ordered best-first.
71    pub fn models_for(&self, category: &Category) -> &[ModelId] {
72        self.by_category.get(category).map_or(&[], Vec::as_slice)
73    }
74
75    /// The models delegated sub-agent work uses for `category`, ordered best-first.
76    pub fn subagent_models_for(&self, category: &Category) -> &[ModelId] {
77        self.subagent
78            .as_ref()
79            .and_then(|models| models.get(category))
80            .map_or(&[], Vec::as_slice)
81    }
82}
83
84impl From<HashMap<Category, Vec<ModelId>>> for RuntimeModels {
85    fn from(by_category: HashMap<Category, Vec<ModelId>>) -> Self {
86        Self::new(by_category)
87    }
88}
89
90/// Which of a [`RuntimeModels`]' groups a driver reads.
91#[derive(Clone, Copy)]
92enum Scope {
93    Parent,
94    Subagent,
95}
96
97/// An offloaded model call, surfaced inside [`Step::CallModel`].
98///
99/// The host reads the public fields and passes its model-call future to
100/// [`respond`](Self::respond), unblocking the algorithm's [`Driver::call_model`] on the
101/// other side. `switchyard-llm-client`'s `run` is the ready-made consumer that does this for
102/// you.
103///
104/// [`Driver::call_model`] stamps the first candidate model onto the request before publishing
105/// the call. A consumer that falls through to a later candidate must re-stamp it.
106pub struct CallModel {
107    /// The name of the algorithm that produced this call, so a host instrumenting the
108    /// calls it serves can attribute its own spans to the algorithm behind them.
109    pub algorithm: String,
110    /// The request to serve; its `model` is stamped with the first candidate.
111    pub request: Request,
112    /// Candidate models, tried in order until one answers. Never empty.
113    pub models: Vec<ModelId>,
114    /// Return client errors to the algorithm so it can apply its fallback policy.
115    pub recover_errors: bool,
116    /// How to send the response back to the algorithm. `None` once the call is recorded.
117    reply: Option<oneshot::Sender<Result<Response>>>,
118    started: Instant,
119    // Retain the originating span even if the waiting algorithm is cancelled first.
120    span: tracing::Span,
121}
122
123impl CallModel {
124    /// Run host work until it completes or the algorithm stops waiting for this call.
125    /// Dropping the waiting [`Driver::call_model`] future drops `work` and returns `Ok(())`.
126    /// Errors stop the host's run unless [`Self::recover_errors`] is enabled.
127    /// Use `std::future::ready(result)` for an already available response.
128    /// A delivered response stream owns its subsequent lifetime.
129    pub async fn respond(mut self, work: impl Future<Output = Result<Response>>) -> Result<()> {
130        let reply = self.reply.as_mut().ok_or(DriverError::ResponseDropped)?;
131        let result = tokio::select! {
132            biased;
133            _ = reply.closed() => return Ok(()),
134            result = work => result,
135        };
136        if self.reply.as_ref().is_some_and(oneshot::Sender::is_closed) {
137            return Ok(());
138        }
139        if !self.recover_errors
140            && let Err(error) = result
141        {
142            return self.fail(error);
143        }
144        self.record(result.is_ok());
145        if let Ok(response) = &result {
146            observability::record_llm_response(response, &self.span);
147        }
148        self.reply
149            .take()
150            .ok_or(DriverError::ResponseDropped)?
151            .send(result)
152            .ok();
153        Ok(())
154    }
155
156    /// Record a failed call and return its error to stop [`drive`].
157    /// Leaves the promise unfulfilled so the driver can cancel the algorithm.
158    pub fn fail(mut self, error: LibsyError) -> Result<()> {
159        self.reply = None;
160        self.record(false);
161        Err(error)
162    }
163
164    fn record(&self, is_ok: bool) {
165        observability::record_llm_call(
166            &self.algorithm,
167            self.models
168                .first()
169                .map(ModelId::as_str)
170                .unwrap_or("NoTargets"),
171            self.started.elapsed(),
172            is_ok,
173        );
174        self.span
175            .record("outcome", if is_ok { "ok" } else { "error" });
176    }
177}
178
179impl Drop for CallModel {
180    fn drop(&mut self) {
181        if self.reply.is_some() {
182            self.record(false);
183        }
184    }
185}
186
187/// A host-owned decision call. Dropping it without replying yields [`DriverError::ResponseDropped`].
188/// Completion or drop records one call and its duration, including time waiting for the host.
189pub struct CallDecision {
190    /// Algorithm name for attributing host telemetry.
191    pub algorithm: String,
192    /// Its model is set to [`Self::model`] before dispatch.
193    pub request: DecisionRequest,
194    /// Target ID for client lookup.
195    pub model: ModelId,
196    reply: Option<oneshot::Sender<Result<DecisionResponse>>>,
197    started: Instant,
198    // Retain the originating span even if the waiting algorithm is cancelled first.
199    span: tracing::Span,
200}
201
202impl CallDecision {
203    /// Run host work and return its response or error to the algorithm.
204    /// Dropping the waiting [`Driver::call_decision`] future drops `work` and returns `Ok(())`.
205    /// Use `std::future::ready(result)` for an already available response.
206    pub async fn respond(
207        mut self,
208        work: impl Future<Output = Result<DecisionResponse>>,
209    ) -> Result<()> {
210        let reply = self.reply.as_mut().ok_or(DriverError::ResponseDropped)?;
211        let result = tokio::select! {
212            biased;
213            _ = reply.closed() => return Ok(()),
214            result = work => result,
215        };
216        self.record(result.is_ok());
217        if let Ok(response) = &result {
218            observability::record_decision_response(response, &self.span);
219        }
220        self.reply
221            .take()
222            .ok_or(DriverError::ResponseDropped)?
223            .send(result)
224            .ok();
225        Ok(())
226    }
227
228    /// Returning this error from the host handler aborts [`drive`].
229    pub fn fail(mut self, error: LibsyError) -> Result<()> {
230        self.reply = None;
231        self.record(false);
232        Err(error)
233    }
234
235    fn record(&self, is_ok: bool) {
236        observability::record_decision_call(
237            &self.algorithm,
238            &self.model,
239            self.started.elapsed(),
240            is_ok,
241        );
242        self.span
243            .record("outcome", if is_ok { "ok" } else { "error" });
244    }
245}
246
247impl Drop for CallDecision {
248    fn drop(&mut self) {
249        if self.reply.is_some() {
250            self.record(false);
251        }
252    }
253}
254
255/// The terminal result of routing.
256pub struct RoutingOutcome {
257    /// Models selected by the algorithm, ordered best model first.
258    pub selected_model_ids: Vec<ModelId>,
259    /// The request after all routing-time rewrites, stamped with the selected model.
260    pub request: Request,
261    /// The request the fallback candidates receive instead of `request`, when routing
262    /// rewrote `request` for the selected model only.
263    /// Content-policy denials do not advance to these fallback candidates.
264    pub fallback_request: Option<Request>,
265    /// A response produced while routing, or `None` when the client must make the answer call.
266    pub response: Option<Response>,
267    /// Outcome identity and optional algorithm evidence.
268    ///
269    /// Constructors leave this empty; [`Algorithm::run_stream`] fills it before publishing a
270    /// successful outcome.
271    pub metadata: Option<crate::OutcomeMetadata>,
272    /// Routing state after the run, for hosts that store state outside the algorithm.
273    pub state: State,
274}
275
276impl RoutingOutcome {
277    /// The model the algorithm recommends, the best model for this request.
278    /// `LibsyError::NoTargets` if the algorithm selected no models, which should be impossible.
279    pub fn selected_model_id(&self) -> Result<&ModelId> {
280        self.selected_model_ids.first().ok_or(LibsyError::NoTargets)
281    }
282
283    /// The decision is that client should send this `request`. The `selected_model_id`
284    /// will be written into it by this function.
285    /// If that fails client should try the `fallback_models` in order.
286    pub fn route_to(
287        selected_model_id: ModelId,
288        fallback_models: Vec<ModelId>,
289        mut request: Request,
290    ) -> Self {
291        request.llm_request.model = Some(selected_model_id.to_string());
292        let mut selected_model_ids = Vec::with_capacity(1 + fallback_models.len());
293        selected_model_ids.push(selected_model_id);
294        selected_model_ids.extend(fallback_models);
295        Self {
296            selected_model_ids,
297            request,
298            fallback_request: None,
299            response: None,
300            metadata: None,
301            state: State::default(),
302        }
303    }
304
305    /// Algorithm generated the response as part of the routing decision. Here it is.
306    /// The `request` will have the `selected_model_id` written into it by this function.
307    pub fn answered(selected_model_id: ModelId, mut request: Request, response: Response) -> Self {
308        request.llm_request.model = Some(selected_model_id.to_string());
309        Self {
310            selected_model_ids: vec![selected_model_id],
311            request,
312            fallback_request: None,
313            response: Some(response),
314            metadata: None,
315            state: State::default(),
316        }
317    }
318}
319
320/// An algorithm's handle for requesting external work and awaiting the host's typed replies.
321#[derive(Clone)]
322pub struct Driver {
323    step_tx: mpsc::Sender<Result<Step>>,
324
325    /// The owning algorithm's telemetry label, stamped onto every call this driver publishes.
326    algorithm: String,
327
328    /// Run-scoped evidence shared by driver clones and attached only to a successful outcome.
329    evidence: Arc<Mutex<Option<Value>>>,
330
331    /// Caller-supplied routing state, read and updated when a request carries no session ID.
332    state: Arc<Mutex<State>>,
333
334    /// Every group this run may route over, shared by all driver clones.
335    models: Arc<RuntimeModels>,
336
337    /// Which of those groups this driver reads.
338    scope: Scope,
339}
340
341impl Driver {
342    /// Build an empty driver with its step channel ready. Created per call by
343    /// [`run_stream`](Algorithm::run_stream). Also returns the Step receiver.
344    pub(crate) fn new(
345        algorithm: &str,
346        models: Arc<RuntimeModels>,
347    ) -> (Self, mpsc::Receiver<Result<Step>>) {
348        // Capacity one keeps the algorithm paced by the stream consumer. It limits queued steps,
349        // not model calls already pulled from the stream, which can still run at the same time.
350        // A larger buffer would use more memory and let the algorithm run farther ahead with
351        // little benefit because reading a step is cheap compared with serving a model call.
352        let (step_tx, step_rx) = mpsc::channel(1);
353        (
354            Self {
355                step_tx,
356                algorithm: algorithm.to_string(),
357                evidence: Arc::new(Mutex::new(None)),
358                state: Arc::new(Mutex::new(State::default())),
359                models,
360                scope: Scope::Parent,
361            },
362            step_rx,
363        )
364    }
365
366    pub(crate) fn take_state(&self) -> State {
367        std::mem::take(&mut *self.state.lock())
368    }
369
370    pub(crate) fn set_state(&self, state: State) {
371        *self.state.lock() = state;
372    }
373
374    /// Replace the current run's evidence when a component makes the final decision.
375    pub(crate) fn set_evidence(&self, evidence: Value) {
376        *self.evidence.lock() = Some(evidence);
377    }
378
379    /// Supply fallback evidence without replacing a decision made earlier in the cascade.
380    pub(crate) fn set_evidence_if_empty(&self, evidence: Value) {
381        let mut current = self.evidence.lock();
382        if current.is_none() {
383            *current = Some(evidence);
384        }
385    }
386
387    /// Publish a model call and await the consumer's response.
388    ///
389    /// Errors if the stream is closed or the call failed.
390    /// The await is wrapped in a `libsy.llm_call` span measuring *fulfillment* as
391    /// the algorithm observes it (host queueing/serving included; a streamed
392    /// response resolves when its stream handle arrives). The host records call metrics,
393    /// outcome, and buffered token usage when it completes or drops the call, even if the
394    /// waiting algorithm has been cancelled. The provider call itself is the
395    /// host's, and is instrumented by whoever makes it.
396    pub async fn call_model(&self, request: Request, models: Vec<ModelId>) -> Result<Response> {
397        self.call_model_with_error_recovery(request, models, false)
398            .await
399    }
400
401    /// Allows a routing policy to handle a failed call when recovery is enabled.
402    #[tracing::instrument(
403        target = "libsy",
404        name = "libsy.llm_call",
405        skip_all,
406        fields(
407            algorithm = self.algorithm,
408            selected_model = %models.first().map(ModelId::as_str).unwrap_or("NoTargets"),
409            openinference.span.kind = "CHAIN",
410            outcome = tracing::field::Empty,
411            input_tokens = tracing::field::Empty,
412            output_tokens = tracing::field::Empty,
413            total_tokens = tracing::field::Empty,
414            reasoning_tokens = tracing::field::Empty,
415        )
416    )]
417    pub(crate) async fn call_model_with_error_recovery(
418        &self,
419        mut request: Request,
420        models: Vec<ModelId>,
421        recover_errors: bool,
422    ) -> Result<Response> {
423        let Some(selected_model_id) = models.first() else {
424            return Err(LibsyError::NoTargets);
425        };
426        request.llm_request.model = Some(selected_model_id.to_string());
427        let started = Instant::now();
428        let (reply, response) = oneshot::channel::<Result<Response>>();
429        let call = CallModel {
430            algorithm: self.algorithm.clone(),
431            request,
432            models,
433            recover_errors,
434            reply: Some(reply),
435            started,
436            span: tracing::Span::current(),
437        };
438        self.step_tx
439            .send(Ok(Step::CallModel(Box::new(call))))
440            .await
441            .map_err(|_| DriverError::StreamClosed)?;
442        response
443            .await
444            .map_err(|_| LibsyError::from(DriverError::ResponseDropped))?
445    }
446
447    /// Override the request's model with `model` and wait for the host to reply.
448    #[tracing::instrument(
449        target = "libsy",
450        name = "libsy.decision_call",
451        skip_all,
452        fields(
453            algorithm = self.algorithm,
454            selected_model = %model,
455            openinference.span.kind = "CHAIN",
456            outcome = tracing::field::Empty,
457            input_tokens = tracing::field::Empty,
458            output_tokens = tracing::field::Empty,
459            total_tokens = tracing::field::Empty,
460            reasoning_tokens = tracing::field::Empty,
461            gen_ai.response.id = tracing::field::Empty,
462            gen_ai.response.model = tracing::field::Empty,
463        ),
464    )]
465    pub async fn call_decision(
466        &self,
467        mut request: DecisionRequest,
468        model: ModelId,
469    ) -> Result<DecisionResponse> {
470        request.model = Some(model.clone());
471        let started = Instant::now();
472        let (reply, response) = oneshot::channel();
473        let call = CallDecision {
474            algorithm: self.algorithm.clone(),
475            request,
476            model,
477            reply: Some(reply),
478            started,
479            span: tracing::Span::current(),
480        };
481        self.step_tx
482            .send(Ok(Step::CallDecision(Box::new(call))))
483            .await
484            .map_err(|_| DriverError::StreamClosed)?;
485        response
486            .await
487            .map_err(|_| LibsyError::from(DriverError::ResponseDropped))?
488    }
489
490    /// The available models for this category, typically ordered best-first.
491    pub fn models_for(&self, category: &Category) -> &[ModelId] {
492        match self.scope {
493            Scope::Parent => self.models.models_for(category),
494            Scope::Subagent => self.models.subagent_models_for(category),
495        }
496    }
497
498    /// The first available model for `category`.
499    pub fn first_model_for(&self, category: &Category) -> Result<&ModelId> {
500        self.models_for(category)
501            .first()
502            .ok_or_else(|| LibsyError::AlgorithmError {
503                message: format!("no models available for category {}", category.as_str()),
504            })
505    }
506
507    /// A driver scoped to delegated sub-agent work: its categories are the
508    /// sub-agent's own, and the parent's are no longer reachable through it.
509    pub fn for_subagent(&self) -> Result<Self> {
510        if self.models.subagent.is_none() {
511            return Err(LibsyError::AlgorithmError {
512                message: "delegated work has no sub-agent models".to_string(),
513            });
514        }
515        Ok(Self {
516            scope: Scope::Subagent,
517            ..self.clone()
518        })
519    }
520
521    /// Emit the terminal step: [`Step::Done`] on `Ok`, or an `Err` stream
522    /// item on failure. Internal: called once by [`run_stream`](Algorithm::run_stream)
523    /// when the algorithm finishes.
524    pub(crate) async fn finish(&self, result: Result<RoutingOutcome>) -> Result<()> {
525        let result = result.map(|mut outcome| {
526            outcome.state = self.take_state();
527            let metadata = outcome.metadata.get_or_insert_with(|| {
528                crate::OutcomeMetadata::new(self.algorithm.clone(), self.evidence.lock().take())
529            });
530            observability::record_outcome(metadata, &outcome.selected_model_ids);
531            outcome
532        });
533        let selected_model = result
534            .as_ref()
535            .ok()
536            .and_then(|outcome| outcome.selected_model_id().ok().cloned());
537        let step = result.map(|outcome| Step::Done(Box::new(outcome)));
538        self.step_tx
539            .send(step)
540            .await
541            .map_err(|_| DriverError::StreamClosed)?;
542        if let Some(selected_model) = selected_model {
543            observability::record_decision(&self.algorithm, &selected_model);
544        }
545        Ok(())
546    }
547}
548
549/// Work requests and the terminal outcome of [`Algorithm::run_stream`].
550/// Each call carries its own reply channel, so the host can execute calls concurrently.
551pub enum Step {
552    /// The algorithm needs this model call performed. The host serves it and fulfills
553    /// it with [`CallModel::respond`]. Boxed: it is by far the largest variant.
554    CallModel(Box<CallModel>),
555    /// The host fulfills this decision call with [`CallDecision::respond`].
556    CallDecision(Box<CallDecision>),
557    /// The algorithm finished with its routing outcome — the last step of a run.
558    Done(Box<RoutingOutcome>),
559}
560
561/// Work dispatched through one host handler, with a typed reply channel per variant.
562/// [`drive`] handles terminal outcomes separately; the handler only executes calls.
563pub enum Call {
564    /// Fulfilled with an LLM [`Response`].
565    Model(Box<CallModel>),
566    /// Fulfilled with a [`DecisionResponse`].
567    Decision(Box<CallDecision>),
568}
569
570/// Drive [`Algorithm::run_stream`] to its routing outcome, serving calls concurrently.
571///
572/// The shared execution loop lets hosts supply clients or test doubles through `serve`
573/// without reimplementing stream handling and cancellation.
574///
575/// `serve` owns each [`Call`] and performs its I/O. Passing an error to `respond`
576/// lets the algorithm fall back; returning an error from `serve` aborts the run.
577///
578/// Dropping this future cancels the algorithm and pending handler futures.
579pub async fn drive<F, Fut>(
580    algorithm: Arc<dyn Algorithm>,
581    request: Request,
582    models: Arc<RuntimeModels>,
583    serve: F,
584) -> Result<RoutingOutcome>
585where
586    F: Fn(Call) -> Fut,
587    Fut: Future<Output = Result<()>>,
588{
589    drive_with_state(algorithm, request, models, State::default(), serve).await
590}
591
592/// [`drive`] starting from caller-supplied `state`.
593pub(crate) async fn drive_with_state<F, Fut>(
594    algorithm: Arc<dyn Algorithm>,
595    request: Request,
596    models: Arc<RuntimeModels>,
597    state: State,
598    serve: F,
599) -> Result<RoutingOutcome>
600where
601    F: Fn(Call) -> Fut,
602    Fut: Future<Output = Result<()>>,
603{
604    let stream = algorithm.run_stream(request, models, state);
605    tokio::pin!(stream);
606
607    let mut in_flight = futures::stream::FuturesUnordered::new();
608    let mut final_outcome: Option<RoutingOutcome> = None;
609
610    loop {
611        tokio::select! {
612            Some(result) = in_flight.next() => match result {
613                Ok(()) => {}, // Call completed successfully
614                Err(err) => return Err(err), // Call failed, propagate the error
615            },
616            step = stream.next() => {
617                match step {
618                    None => break, // stream has ended, no more steps
619                    Some(item) => match item? {
620                        Step::CallModel(call) => in_flight.push(serve(Call::Model(call))),
621                        Step::CallDecision(call) => {
622                            in_flight.push(serve(Call::Decision(call)));
623                        }
624                        Step::Done(outcome) => {
625                            final_outcome = Some(*outcome);
626                            break;
627                        }
628                    }
629                }
630            },
631        }
632    }
633    final_outcome.ok_or(LibsyError::MissingFinalResponse)
634}
635
636/// Recover the message from an algorithm's panic.
637fn panic_message(payload: &(dyn std::any::Any + Send)) -> String {
638    payload
639        .downcast_ref::<&'static str>()
640        .map(|message| (*message).to_string())
641        .or_else(|| payload.downcast_ref::<String>().cloned())
642        .unwrap_or_else(|| "unknown panic payload".to_string())
643}
644
645/// Abort guard
646struct AbortOnDrop(tokio::task::AbortHandle);
647
648impl Drop for AbortOnDrop {
649    fn drop(&mut self) {
650        self.0.abort();
651    }
652}
653
654/// Errors unless `targets` contains `name`.
655///
656/// Config target names must be resolved before an algorithm is built. This list contains
657/// model IDs, not target names.
658pub(crate) fn ensure_model_is_target(targets: &[ModelId], name: &ModelId) -> Result<()> {
659    targets
660        .iter()
661        .any(|target| target == name)
662        .then_some(())
663        .ok_or_else(|| LibsyError::TargetNotFound {
664            target: name.clone(),
665        })
666}
667
668/// Key for routing affinity: a root request by its session, a child request by its session
669/// and agent.
670#[derive(Clone, Hash, PartialEq, Eq)]
671pub(crate) enum RoutingIdentity {
672    /// Root request, keyed by session ID.
673    Session(String),
674    /// Child request, keyed by session and agent IDs.
675    Subagent { session: String, agent: String },
676}
677
678impl RoutingIdentity {
679    /// Builds a root or child identity from non-empty request metadata.
680    ///
681    /// A child request missing either ID returns `None`, so it keeps no routing history
682    /// rather than sharing the parent's.
683    pub(crate) fn from_request(request: &Request) -> Option<Self> {
684        let metadata = request.metadata.as_ref()?;
685        let session = metadata.session_id.as_deref().filter(|id| !id.is_empty())?;
686        if metadata.is_subagent {
687            let agent = metadata.agent_id.as_deref().filter(|id| !id.is_empty())?;
688            Some(Self::Subagent {
689                session: session.to_string(),
690                agent: agent.to_string(),
691            })
692        } else {
693            Some(Self::Session(session.to_string()))
694        }
695    }
696}
697
698/// Routing policy independent of how external calls are executed.
699/// Implement [`route`](Self::route), using [`Driver`] to request work from the host.
700/// Hosts can use [`drive`] for the shared execution loop or consume
701/// [`run_stream`](Self::run_stream) for custom execution.
702///
703/// Methods take `self: Arc<Self>`: one algorithm (`Arc<dyn Algorithm>`) is shared across
704/// requests and run concurrently, so it owns its thread-safety and any shared state.
705///
706/// # Concurrency
707///
708/// A host may run the same algorithm concurrently for many requests. Implementations
709/// must synchronize their own mutable shared state. Each call to [`run_stream`](Self::run_stream)
710/// creates an independent [`Driver`], so model-call promises and emitted [`Step`]s cannot
711/// cross between runs.
712///
713/// # Observability
714///
715/// [`run_stream`](Self::run_stream) creates a `libsy.run` span, and each offloaded model
716/// call creates a nested span for its call kind, with outcome and available token usage.
717/// Call metrics include host queueing and count unfulfilled drops as errors.
718/// Successful outcomes record their
719/// [`OutcomeMetadata::outcome_id`](crate::OutcomeMetadata::outcome_id) on `libsy.run`,
720/// alongside `selected_model_ids` (an ordered OpenTelemetry string array).
721/// `algorithm` and `switchyard.algorithm` retain the run's [`Algorithm::name`].
722/// Optional `evidence.source`, `evidence.verdict`, `evidence.trigger`, and
723/// `evidence.reason_code` are strings; `evidence.score`, `evidence.confidence`, and
724/// `evidence.threshold` are numbers. Unknown evidence fields are not exported.
725/// These fields are span attributes, never metric labels.
726///
727/// The run/call observability helpers retain `outcome` status and operational metrics,
728/// but omit error details and arbitrary request extra metadata. Algorithms and hosts
729/// may emit their own logs. Errors still reach the caller unchanged.
730/// The host controls the tracing subscriber and global OpenTelemetry
731/// meter provider; libsy installs no exporter and performs no telemetry network I/O.
732#[async_trait]
733pub trait Algorithm: Send + Sync + 'static {
734    /// Stable, low-cardinality name identifying this algorithm — the
735    /// `algorithm` attribute on every span, metric, and log line the crate
736    /// emits for its runs.
737    fn name(&self) -> &str;
738
739    /// Select a route, requesting external work through [`Driver`] as needed.
740    /// [`run_stream`](Self::run_stream) runs this method and exposes its work to the host.
741    async fn route(self: Arc<Self>, driver: Driver, request: Request) -> Result<RoutingOutcome>;
742
743    /// Expose work requests and the final outcome so the host can control execution.
744    ///
745    /// Each call waits for its host reply. A run ends with one terminal item — [`Step::Done`] on
746    /// success, an `Err` item on failure, including when the algorithm panics. Dropping
747    /// the stream aborts the spawned algorithm task.
748    ///
749    /// Every invocation owns a separate [`Driver`].
750    fn run_stream(
751        self: Arc<Self>,
752        request: Request,
753        models: Arc<RuntimeModels>,
754        state: State,
755    ) -> StepStream {
756        let (driver, step_rx) = Driver::new(self.name(), models);
757        driver.set_state(state);
758        let span = observability::run_span(self.name(), &request);
759        let handle = tokio::spawn(
760            async move {
761                let algorithm = self.name().to_string();
762                // Catch a panicking algorithm so the run still publishes a terminal step.
763                let route = AssertUnwindSafe(self.route(driver.clone(), request)).catch_unwind();
764                let result = observability::observe_run(&algorithm, async move {
765                    route.await.unwrap_or_else(|payload| {
766                        Err(LibsyError::AlgorithmError {
767                            message: format!(
768                                "algorithm task panicked: {}",
769                                panic_message(payload.as_ref())
770                            ),
771                        })
772                    })
773                })
774                .await;
775
776                let _ = driver.finish(result).await;
777            }
778            .instrument(span),
779        );
780        // Dropping the stream aborts the algorithm task when its consumer goes away.
781        let abort_guard = AbortOnDrop(handle.abort_handle());
782        Box::pin(ReceiverStream::new(step_rx).map(move |step| {
783            // link abort guard to stream
784            let _keep_alive = &abort_guard;
785            step
786        }))
787    }
788}
789
790#[cfg(test)]
791mod tests {
792    use std::collections::HashMap;
793
794    use super::*;
795    use crate::core::testing::{Serve, ServeResult, echo, reply, serve_decision, test_drive};
796    use futures::StreamExt;
797    use switchyard_protocol::{
798        LlmResponse, LlmResponseChunk, completion_text, text_request, text_response,
799    };
800
801    #[derive(Debug, thiserror::Error)]
802    #[error("{0}")]
803    struct TestError(&'static str);
804
805    fn test_error(message: &'static str) -> LibsyError {
806        LibsyError::external("test", TestError(message))
807    }
808
809    /// Trivial algo used only to exercise the orchestrator: calls the first target
810    /// and returns its response as the routing outcome.
811    struct TestAlgo {
812        target_set: Vec<ModelId>,
813    }
814
815    #[async_trait]
816    impl Algorithm for TestAlgo {
817        fn name(&self) -> &str {
818            "test"
819        }
820
821        async fn route(
822            self: Arc<Self>,
823            driver: Driver,
824            request: Request,
825        ) -> Result<RoutingOutcome> {
826            let target = self
827                .target_set
828                .first()
829                .ok_or(LibsyError::NoTargets)?
830                .clone();
831            let response = driver
832                .call_model(request.clone(), vec![target.clone()])
833                .await?;
834            driver.set_evidence(serde_json::json!({"source": "test"}));
835            driver.set_evidence_if_empty(serde_json::json!({"source": "ignored"}));
836            Ok(RoutingOutcome::answered(target, request, response))
837        }
838    }
839
840    /// Build a shared `TestAlgo` over the given target set.
841    fn orch(target_set: Vec<ModelId>) -> Arc<dyn Algorithm> {
842        Arc::new(TestAlgo { target_set })
843    }
844
845    fn request() -> Request {
846        Request {
847            llm_request: text_request(Some("auto".to_string()), "hi".to_string()),
848            raw_request: None,
849            metadata: None,
850        }
851    }
852
853    #[test]
854    fn routing_outcome_constructors_stamp_selection_and_preserve_payloads() {
855        let outcome = RoutingOutcome::route_to(
856            "selected".into(),
857            target_set(&["fallback-one", "fallback-two"]),
858            request(),
859        );
860
861        assert_eq!(
862            outcome.selected_model_ids,
863            target_set(&["selected", "fallback-one", "fallback-two"])
864        );
865        assert_eq!(outcome.request.model_id().as_deref(), Some("selected"));
866        assert!(outcome.response.is_none());
867        assert!(outcome.metadata.is_none());
868
869        let outcome = RoutingOutcome::route_to("only".into(), Vec::new(), request());
870        assert_eq!(outcome.selected_model_ids, target_set(&["only"]));
871
872        let outcome = RoutingOutcome::answered(
873            "answered".into(),
874            request(),
875            Response {
876                llm_response: LlmResponse::Agg(text_response(None, "existing")),
877                metadata: None,
878                upstream_headers: http::HeaderMap::new(),
879            },
880        );
881
882        assert_eq!(outcome.selected_model_ids, target_set(&["answered"]));
883        assert_eq!(outcome.request.model_id().as_deref(), Some("answered"));
884        assert_eq!(
885            outcome
886                .response
887                .as_ref()
888                .and_then(|response| response.llm_response.as_agg())
889                .map(completion_text),
890            Some("existing".to_string())
891        );
892    }
893
894    fn target_set(names: &[&str]) -> Vec<ModelId> {
895        names.iter().map(|name| ModelId::from(*name)).collect()
896    }
897
898    #[tokio::test]
899    async fn typed_driver_preserves_call_and_stream_boundaries() -> Result<()> {
900        tokio::time::timeout(std::time::Duration::from_secs(1), async {
901            // Distinct oneshots keep reverse-order replies paired with their producers, and a
902            // retained call remains pending until the host responds.
903            let (driver, mut step_rx) = Driver::new("test", Arc::new(RuntimeModels::default()));
904            let first_driver = driver.clone();
905            let mut first = tokio::spawn(async move {
906                first_driver
907                    .call_model(request(), vec![ModelId::from("first")])
908                    .await
909            });
910            let second = tokio::spawn(async move {
911                driver
912                    .call_model(request(), vec![ModelId::from("second")])
913                    .await
914            });
915
916            let mut calls = HashMap::new();
917            for _ in 0..2 {
918                let step = step_rx.recv().await.ok_or(DriverError::StreamClosed)??;
919                let Step::CallModel(call) = step else {
920                    return Err(test_error("expected a CallModel step"));
921                };
922                let selected_model = call
923                    .models
924                    .first()
925                    .ok_or_else(|| test_error("model call has no candidates"))?
926                    .to_string();
927                calls.insert(selected_model, call);
928            }
929            assert!(
930                tokio::time::timeout(std::time::Duration::from_millis(20), &mut first)
931                    .await
932                    .is_err(),
933                "call completed before the host responded"
934            );
935            calls
936                .remove("second")
937                .ok_or_else(|| test_error("missing second call"))?
938                .respond(std::future::ready(Ok(reply("second response")))).await?;
939            calls
940                .remove("first")
941                .ok_or_else(|| test_error("missing first call"))?
942                .respond(std::future::ready(Ok(reply("first response")))).await?;
943
944            let first_response = first
945                .await
946                .map_err(|source| LibsyError::external("joining a test task", source))??;
947            let second_response = second
948                .await
949                .map_err(|source| LibsyError::external("joining a test task", source))??;
950            assert_eq!(
951                first_response.llm_response.as_agg().map(completion_text),
952                Some("first response".to_string())
953            );
954            assert_eq!(
955                second_response.llm_response.as_agg().map(completion_text),
956                Some("second response".to_string())
957            );
958
959            // Dropping the host-facing promise closes only that call's reply channel.
960            let (driver, mut step_rx) = Driver::new("test", Arc::new(RuntimeModels::default()));
961            let producer = tokio::spawn(async move {
962                driver
963                    .call_model(request(), vec![ModelId::from("dropped")])
964                    .await
965            });
966            let step = step_rx.recv().await.ok_or(DriverError::StreamClosed)??;
967            let Step::CallModel(call) = step else {
968                return Err(test_error("expected a CallModel step"));
969            };
970            drop(call);
971            let result = producer
972                .await
973                .map_err(|source| LibsyError::external("joining a test task", source))?;
974            assert!(matches!(
975                result,
976                Err(LibsyError::Driver(DriverError::ResponseDropped))
977            ));
978
979            // A standalone driver reports the typed step receiver disappearing at its next call.
980            let (driver, step_rx) = Driver::new("test", Arc::new(RuntimeModels::default()));
981            drop(step_rx);
982            let result = driver
983                .call_model(request(), vec![ModelId::from("closed")])
984                .await;
985            assert!(matches!(
986                result,
987                Err(LibsyError::Driver(DriverError::StreamClosed))
988            ));
989
990            fn decision_request() -> DecisionRequest {
991                DecisionRequest {
992                    model: Some("overwritten".into()),
993                    context: serde_json::json!({"task": "choose a route"}),
994                    questions: Default::default(),
995                }
996            }
997
998            fn decision_response() -> DecisionResponse {
999                DecisionResponse {
1000                    id: Some("decision-1".to_string()),
1001                    model: Some("provider-model".into()),
1002                    answers: [(
1003                        "p_solve".to_string(),
1004                        switchyard_protocol::DecisionAnswer {
1005                            value: switchyard_protocol::DecisionValue::Boolean(
1006                                switchyard_protocol::BooleanEstimate::ProbabilityTrue(
1007                                    switchyard_protocol::Probability(0.8),
1008                                ),
1009                            ),
1010                            provider_confidence: None,
1011                        },
1012                    )]
1013                    .into(),
1014                    usage: Default::default(),
1015                }
1016            }
1017
1018            struct MixedCalls(&'static str);
1019
1020            #[async_trait]
1021            impl Algorithm for MixedCalls {
1022                fn name(&self) -> &str {
1023                    "mixed"
1024                }
1025
1026                async fn route(
1027                    self: Arc<Self>,
1028                    driver: Driver,
1029                    request: Request,
1030                ) -> Result<RoutingOutcome> {
1031                    let (llm, decision) = tokio::join!(
1032                        driver.call_model(request.clone(), vec!["llm".into()]),
1033                        driver.call_decision(decision_request(), "decision".into()),
1034                    );
1035                    assert_eq!(
1036                        llm?.llm_response.as_agg().map(completion_text),
1037                        Some("llm reply".into())
1038                    );
1039                    match self.0 {
1040                        "mock" => assert_eq!(decision?, DecisionResponse {
1041                            id: None,
1042                            model: Some("decision".into()),
1043                            answers: Default::default(),
1044                            usage: Default::default(),
1045                        }),
1046                        "reply" => assert_eq!(decision?, decision_response()),
1047                        "error" => assert!(matches!(
1048                            decision,
1049                            Err(LibsyError::AlgorithmError { message }) if message == "provider failed"
1050                        )),
1051                        "drop" => assert!(matches!(
1052                            decision,
1053                            Err(LibsyError::Driver(DriverError::ResponseDropped))
1054                        )),
1055                        "abort" => return std::future::pending().await,
1056                        _ => unreachable!(),
1057                    }
1058                    Ok(RoutingOutcome::route_to("answer".into(), vec![], request))
1059                }
1060            }
1061
1062            for mode in ["mock", "reply", "error", "drop", "abort"] {
1063                // Serial dispatch would deadlock here and fail the enclosing timeout.
1064                let barrier = Arc::new(tokio::sync::Barrier::new(2));
1065                let outcome = drive(
1066                    Arc::new(MixedCalls(mode)),
1067                    request(),
1068                    Arc::new(RuntimeModels::default()),
1069                    move |call| {
1070                        let barrier = barrier.clone();
1071                        async move {
1072                            barrier.wait().await;
1073                            let call = match call {
1074                                Call::Model(call) => return call.respond(std::future::ready(Ok(reply("llm reply")))).await,
1075                                Call::Decision(call) => *call,
1076                            };
1077                            assert_eq!(call.algorithm, "mixed");
1078                            assert_eq!(call.model, "decision");
1079                            assert_eq!(call.request.model, Some("decision".into()));
1080                            assert_eq!(call.request.context, decision_request().context);
1081                            match mode {
1082                                "mock" => serve_decision(call).await,
1083                                "reply" => call.respond(std::future::ready(Ok(decision_response()))).await,
1084                                "error" => call.respond(std::future::ready(Err(LibsyError::AlgorithmError {
1085                                    message: "provider failed".into(),
1086                                }))).await,
1087                                "drop" => {
1088                                    drop(call);
1089                                    Ok(())
1090                                }
1091                                "abort" => call.fail(test_error("host aborted")),
1092                                _ => unreachable!(),
1093                            }
1094                        }
1095                    },
1096                )
1097                .await;
1098                if mode == "abort" {
1099                    assert!(matches!(outcome, Err(LibsyError::External { .. })));
1100                } else {
1101                    assert_eq!(outcome?.selected_model_id()?, "answer");
1102                }
1103            }
1104
1105            let (driver, mut steps) = Driver::new("test", Arc::new(RuntimeModels::default()));
1106            let mut pending = Box::pin(driver.call_decision(decision_request(), "decision".into()));
1107            assert!(futures::poll!(&mut pending).is_pending());
1108            let Some(Ok(Step::CallDecision(call))) = steps.recv().await else {
1109                return Err(test_error("expected a decision call"));
1110            };
1111            drop(pending);
1112            call.respond(async { panic!("cancelled work must not be polled") }).await?;
1113            drop(steps);
1114            assert!(matches!(
1115                driver.call_decision(decision_request(), "decision".into()).await,
1116                Err(LibsyError::Driver(DriverError::StreamClosed))
1117            ));
1118            Ok(())
1119        })
1120        .await
1121        .map_err(|error| LibsyError::external("waiting for typed driver boundaries", error))?
1122    }
1123
1124    #[test]
1125    fn target_lookup_returns_the_missing_target() {
1126        let error = ensure_model_is_target(&target_set(&[]), &ModelId::from("missing")).err();
1127        assert!(matches!(
1128            error,
1129            Some(LibsyError::TargetNotFound { target }) if target == "missing"
1130        ));
1131    }
1132
1133    /// Build a single-target algo, plus a `serve` that answers it as a token stream
1134    /// replaying `chunks` in order (as `Ok` items).
1135    fn streaming_orch(chunks: Vec<LlmResponseChunk>) -> (Arc<dyn Algorithm>, impl Serve) {
1136        let algo = orch(target_set(&["stream/model"]));
1137        let serve = move |_target: ModelId, _request: Request| {
1138            let chunks = chunks.clone();
1139            async move {
1140                let stream =
1141                    futures::stream::iter(chunks.into_iter().map(|chunk| Ok(chunk.into()))).boxed();
1142                Ok(Response {
1143                    llm_response: LlmResponse::Stream(stream),
1144                    metadata: None,
1145                    upstream_headers: http::HeaderMap::new(),
1146                })
1147            }
1148        };
1149        (algo, serve)
1150    }
1151
1152    #[tokio::test]
1153    async fn run_returns_a_streamed_response_the_caller_aggregates() -> Result<()> {
1154        // A streaming client -> its chunks flow through the promise and `Done`,
1155        // and `run` returns the live stream untouched for the caller to fold.
1156        let (orch, serve) = streaming_orch(vec![
1157            LlmResponseChunk::MessageStart {
1158                id: Some("m1".to_string()),
1159                model: Some("stream/model".to_string()),
1160            },
1161            LlmResponseChunk::TextDelta {
1162                index: 0,
1163                text: "hel".to_string(),
1164            },
1165            LlmResponseChunk::TextDelta {
1166                index: 0,
1167                text: "lo".to_string(),
1168            },
1169            LlmResponseChunk::MessageStop {
1170                reason: Some("stop".to_string()),
1171            },
1172        ]);
1173        let (selected_model, response) = test_drive(orch, request(), serve).await?;
1174        // The run handed back the live stream; the caller folds it to a buffered aggregate.
1175        let agg = response
1176            .llm_response
1177            .into_agg()
1178            .await
1179            .map_err(|error| LibsyError::external("aggregating response stream", error))?;
1180        assert_eq!(completion_text(&agg), "hello");
1181        assert_eq!(agg.model.as_deref(), Some("stream/model"));
1182        assert_eq!(selected_model, "stream/model");
1183        Ok(())
1184    }
1185
1186    #[tokio::test]
1187    async fn aggregating_a_streamed_response_propagates_a_mid_stream_error() -> Result<()> {
1188        // The run succeeds and returns the stream; the in-band `Error` chunk surfaces only
1189        // when the caller aggregates it.
1190        let (orch, serve) = streaming_orch(vec![
1191            LlmResponseChunk::TextDelta {
1192                index: 0,
1193                text: "partial".to_string(),
1194            },
1195            LlmResponseChunk::StreamError {
1196                message: "upstream exploded".to_string(),
1197            },
1198        ]);
1199        let (_, response) = test_drive(orch, request(), serve).await?;
1200        match response.llm_response.into_agg().await {
1201            Ok(_) => panic!("expected a mid-stream error, got an aggregate"),
1202            Err(err) => {
1203                assert!(err.to_string().contains("upstream exploded"));
1204                Ok(())
1205            }
1206        }
1207    }
1208
1209    #[tokio::test]
1210    async fn run_offloads_via_promise_then_finishes() -> Result<()> {
1211        // Every call is offloaded via a promise the orchestrator surfaces as a
1212        // `CallModel` step for us to fulfill.
1213        let stream = orch(target_set(&["offload/model"])).run_stream(
1214            request(),
1215            Arc::new(RuntimeModels::default()),
1216            State::default(),
1217        );
1218        tokio::pin!(stream);
1219
1220        let mut saw_call = false;
1221        let mut final_completion = None;
1222        while let Some(step) = stream.next().await {
1223            match step? {
1224                Step::CallDecision(_) => return Err(test_error("unexpected decision call")),
1225                Step::CallModel(call) => {
1226                    saw_call = true;
1227                    assert_eq!(call.models, vec![ModelId::from("offload/model")]);
1228                    // Fulfilling the promise is the "real" model call the caller makes.
1229                    call.respond(std::future::ready(Ok(Response {
1230                        llm_response: LlmResponse::Agg(text_response(
1231                            None,
1232                            "fulfilled".to_string(),
1233                        )),
1234                        metadata: None,
1235                        upstream_headers: http::HeaderMap::new(),
1236                    })))
1237                    .await?;
1238                }
1239                Step::Done(outcome) => {
1240                    let metadata = outcome
1241                        .metadata
1242                        .as_ref()
1243                        .expect("run_stream should attach outcome metadata");
1244                    assert_eq!(metadata.algorithm, "test");
1245                    assert_eq!(
1246                        uuid::Uuid::parse_str(metadata.outcome_id())
1247                            .expect("outcome id should be a UUID")
1248                            .get_version_num(),
1249                        7
1250                    );
1251                    assert_eq!(
1252                        metadata.evidence,
1253                        Some(serde_json::json!({"source": "test"}))
1254                    );
1255                    let response = outcome
1256                        .response
1257                        .ok_or_else(|| test_error("expected an answered outcome"))?;
1258                    final_completion = Some(
1259                        response
1260                            .llm_response
1261                            .as_agg()
1262                            .map(completion_text)
1263                            .unwrap_or_default(),
1264                    );
1265                }
1266            }
1267        }
1268
1269        assert!(saw_call, "expected a CallModel step before Done");
1270        assert_eq!(
1271            final_completion.ok_or_else(|| test_error("no Done step"))?,
1272            "fulfilled"
1273        );
1274        Ok(())
1275    }
1276
1277    #[tokio::test(flavor = "multi_thread", worker_threads = 12)]
1278    async fn requests_are_processed_in_parallel() -> Result<()> {
1279        use std::time::Duration;
1280        use tokio::sync::Barrier;
1281
1282        const N: usize = 12;
1283
1284        // Serving blocks until all N concurrent calls have arrived. If requests were
1285        // serialized (one algorithm behind a `Mutex`), only one call could be in flight,
1286        // the barrier would never reach N, and the test would time out. It passes only
1287        // because the shared algorithm is driven concurrently across requests.
1288        let barrier = Arc::new(Barrier::new(N));
1289        // One shared algorithm driven by many concurrent requests.
1290        let algo = orch(target_set(&["m"]));
1291
1292        let mut handles = Vec::new();
1293        for _ in 0..N {
1294            let algo = algo.clone();
1295            let barrier = barrier.clone();
1296            let serve = move |target: ModelId, _request: Request| {
1297                let barrier = barrier.clone();
1298                async move {
1299                    barrier.wait().await;
1300                    Ok(reply(target))
1301                }
1302            };
1303            handles.push(tokio::spawn(async move {
1304                test_drive(algo, request(), serve)
1305                    .await
1306                    .map(|(_, response)| {
1307                        response
1308                            .llm_response
1309                            .as_agg()
1310                            .map(completion_text)
1311                            .unwrap_or_default()
1312                    })
1313            }));
1314        }
1315
1316        for handle in handles {
1317            // The timeout turns a serialization deadlock into a failure, not a hang.
1318            let completion = tokio::time::timeout(Duration::from_secs(5), handle)
1319                .await
1320                .map_err(|error| LibsyError::external("waiting for test task", error))?
1321                .map_err(|source| LibsyError::external("joining a test task", source))??;
1322            assert_eq!(completion, "m");
1323        }
1324        Ok(())
1325    }
1326
1327    #[tokio::test]
1328    async fn offload_error_propagates_back_to_the_algorithm() -> Result<()> {
1329        // A client-less target offloads its call; we fulfill the promise with an
1330        // Err, which must flow back through `call_model_target` into the algorithm and
1331        // out as an error step — not a response.
1332        let stream = orch(target_set(&["offload/model"])).run_stream(
1333            request(),
1334            Arc::new(RuntimeModels::default()),
1335            State::default(),
1336        );
1337        tokio::pin!(stream);
1338
1339        let mut saw_error = false;
1340        while let Some(step) = stream.next().await {
1341            match step {
1342                Ok(Step::CallDecision(_)) => return Err(test_error("unexpected decision call")),
1343                Ok(Step::CallModel(mut call)) => {
1344                    call.recover_errors = true;
1345                    call.respond(std::future::ready(Err(test_error(
1346                        "upstream model call failed",
1347                    ))))
1348                    .await?;
1349                }
1350                Ok(Step::Done(..)) => {
1351                    return Err(test_error(
1352                        "expected the offload error to propagate, got a response",
1353                    ));
1354                }
1355                Err(err) => {
1356                    // The algorithm's `call_model_target` saw the error via the promise.
1357                    assert!(err.to_string().contains("upstream model call failed"));
1358                    saw_error = true;
1359                }
1360            }
1361        }
1362
1363        assert!(saw_error, "expected an error step");
1364        Ok(())
1365    }
1366
1367    #[tokio::test]
1368    async fn dropping_the_stream_cancels_the_algorithm_task() -> Result<()> {
1369        use std::sync::atomic::{AtomicBool, Ordering};
1370        use std::time::Duration;
1371        use tokio::sync::mpsc;
1372
1373        // Sets a flag when dropped, so we can observe whether the algorithm task was
1374        // cancelled/dropped.
1375        struct DropGuard(Arc<AtomicBool>);
1376        impl Drop for DropGuard {
1377            fn drop(&mut self) {
1378                self.0.store(true, Ordering::SeqCst);
1379            }
1380        }
1381
1382        struct StuckAlgo {
1383            started: mpsc::UnboundedSender<()>,
1384            dropped: Arc<AtomicBool>,
1385        }
1386
1387        #[async_trait]
1388        impl Algorithm for StuckAlgo {
1389            fn name(&self) -> &str {
1390                "stuck"
1391            }
1392
1393            async fn route(
1394                self: Arc<Self>,
1395                _driver: Driver,
1396                _request: Request,
1397            ) -> Result<RoutingOutcome> {
1398                let _guard = DropGuard(self.dropped.clone());
1399                let _ = self.started.send(());
1400                // Await forever without ever touching the driver.
1401                std::future::pending::<()>().await;
1402                unreachable!()
1403            }
1404        }
1405
1406        let (started_tx, mut started_rx) = mpsc::unbounded_channel();
1407        let dropped = Arc::new(AtomicBool::new(false));
1408        let algo: Arc<dyn Algorithm> = Arc::new(StuckAlgo {
1409            started: started_tx,
1410            dropped: dropped.clone(),
1411        });
1412
1413        let stream = algo.run_stream(
1414            request(),
1415            Arc::new(RuntimeModels::default()),
1416            State::default(),
1417        );
1418        started_rx
1419            .recv()
1420            .await
1421            .ok_or_else(|| test_error("task never started"))?;
1422        drop(stream);
1423        tokio::time::sleep(Duration::from_millis(100)).await;
1424
1425        assert!(
1426            dropped.load(Ordering::SeqCst),
1427            "algorithm task was NOT cancelled after dropping the stream"
1428        );
1429        Ok(())
1430    }
1431
1432    #[tokio::test]
1433    async fn route_panic_surfaces_as_a_stream_error() -> Result<()> {
1434        // An algorithm whose task panics must surface an `Err` step carrying the panic
1435        // message, not abort the process from an unobserved detached task.
1436        struct Panicky;
1437
1438        #[async_trait]
1439        impl Algorithm for Panicky {
1440            fn name(&self) -> &str {
1441                "panicky"
1442            }
1443
1444            async fn route(
1445                self: Arc<Self>,
1446                _driver: Driver,
1447                _request: Request,
1448            ) -> Result<RoutingOutcome> {
1449                panic!("boom");
1450            }
1451        }
1452
1453        let algo: Arc<dyn Algorithm> = Arc::new(Panicky);
1454        let stream = algo.run_stream(
1455            request(),
1456            Arc::new(RuntimeModels::default()),
1457            State::default(),
1458        );
1459        tokio::pin!(stream);
1460
1461        let mut saw_error = false;
1462        while let Some(step) = stream.next().await {
1463            match step {
1464                Err(err) => {
1465                    // The panic message is preserved, not flattened into an opaque failure.
1466                    assert!(err.to_string().contains("algorithm task panicked: boom"));
1467                    saw_error = true;
1468                }
1469                Ok(_) => return Err(test_error("expected the panic to surface as an error step")),
1470            }
1471        }
1472
1473        assert!(saw_error, "expected an error step from the panicked task");
1474        Ok(())
1475    }
1476
1477    /// A panicking algorithm must publish its terminal step even when it left a `Driver`
1478    /// clone alive in another task. That clone holds the step channel open, so a run that
1479    /// merely unwound would never terminate and the consumer would wait forever.
1480    #[tokio::test]
1481    async fn a_panic_with_a_leaked_driver_clone_still_terminates_the_run() -> Result<()> {
1482        struct LeakyPanic;
1483
1484        #[async_trait]
1485        impl Algorithm for LeakyPanic {
1486            fn name(&self) -> &str {
1487                "leaky_panic"
1488            }
1489
1490            async fn route(
1491                self: Arc<Self>,
1492                driver: Driver,
1493                _request: Request,
1494            ) -> Result<RoutingOutcome> {
1495                tokio::spawn(async move {
1496                    // Outlives the panic below, keeping a sender clone alive.
1497                    let _keep_alive = driver;
1498                    std::future::pending::<()>().await;
1499                });
1500                tokio::task::yield_now().await;
1501                panic!("boom");
1502            }
1503        }
1504
1505        let algo: Arc<dyn Algorithm> = Arc::new(LeakyPanic);
1506        // The timeout turns the hang this guards against into a failure rather than a hang.
1507        let result = tokio::time::timeout(
1508            std::time::Duration::from_secs(1),
1509            test_drive(algo, request(), echo()),
1510        )
1511        .await
1512        .map_err(|error| LibsyError::external("waiting for the panicked run to end", error))?;
1513
1514        match result {
1515            Ok(_) => Err(test_error(
1516                "expected the panic to end the run with an error",
1517            )),
1518            Err(err) => {
1519                assert!(err.to_string().contains("algorithm task panicked: boom"));
1520                Ok(())
1521            }
1522        }
1523    }
1524
1525    #[tokio::test]
1526    async fn cancelling_run_cancels_the_algorithm_task() -> Result<()> {
1527        use std::sync::atomic::{AtomicBool, Ordering};
1528        use std::time::Duration;
1529        use tokio::sync::mpsc;
1530
1531        // Sets a flag when dropped, so we can observe whether the algorithm task was
1532        // cancelled once the `run` future driving it is dropped.
1533        struct DropGuard(Arc<AtomicBool>);
1534        impl Drop for DropGuard {
1535            fn drop(&mut self) {
1536                self.0.store(true, Ordering::SeqCst);
1537            }
1538        }
1539
1540        struct StuckAlgo {
1541            started: mpsc::UnboundedSender<()>,
1542            dropped: Arc<AtomicBool>,
1543        }
1544
1545        #[async_trait]
1546        impl Algorithm for StuckAlgo {
1547            fn name(&self) -> &str {
1548                "stuck"
1549            }
1550
1551            async fn route(
1552                self: Arc<Self>,
1553                _driver: Driver,
1554                _request: Request,
1555            ) -> Result<RoutingOutcome> {
1556                let _guard = DropGuard(self.dropped.clone());
1557                let _ = self.started.send(());
1558                // Hang forever without ever touching the driver, so only cancellation
1559                // (not a dropped step channel) can stop this task.
1560                std::future::pending::<()>().await;
1561                unreachable!()
1562            }
1563        }
1564
1565        let (started_tx, mut started_rx) = mpsc::unbounded_channel();
1566        let dropped = Arc::new(AtomicBool::new(false));
1567        let algo: Arc<dyn Algorithm> = Arc::new(StuckAlgo {
1568            started: started_tx,
1569            dropped: dropped.clone(),
1570        });
1571
1572        // Drive the run on its own task, wait until the algorithm task is up, then cancel
1573        // it — dropping its future (and the `run_stream` stream it holds).
1574        let run_task = tokio::spawn(async move { test_drive(algo, request(), echo()).await });
1575        started_rx
1576            .recv()
1577            .await
1578            .ok_or_else(|| test_error("task never started"))?;
1579        run_task.abort();
1580        tokio::time::sleep(Duration::from_millis(100)).await;
1581
1582        assert!(
1583            dropped.load(Ordering::SeqCst),
1584            "algorithm task was NOT cancelled after cancelling run"
1585        );
1586        Ok(())
1587    }
1588
1589    // --- first-wins hedging: `run` must not wait on losing speculative calls -------------
1590
1591    /// Offloads two targets concurrently and returns the first to resolve, dropping the
1592    /// loser's call (first-wins hedging).
1593    struct Hedge {
1594        winner: String,
1595        loser: String,
1596    }
1597
1598    #[async_trait]
1599    impl Algorithm for Hedge {
1600        fn name(&self) -> &str {
1601            "hedge"
1602        }
1603
1604        async fn route(
1605            self: Arc<Self>,
1606            driver: Driver,
1607            request: Request,
1608        ) -> Result<RoutingOutcome> {
1609            let outcome_request = request.clone();
1610            let win = driver.call_model(request.clone(), vec![self.winner.clone().into()]);
1611            let lose = driver.call_model(request, vec![self.loser.clone().into()]);
1612            // First to resolve wins; `select!` drops the losing future (and its promise).
1613            tokio::select! {
1614                res = win => Ok(RoutingOutcome::answered(
1615                    self.winner.clone().into(),
1616                    outcome_request,
1617                    res?,
1618                )),
1619                res = lose => Ok(RoutingOutcome::answered(
1620                    self.loser.clone().into(),
1621                    outcome_request,
1622                    res?,
1623                )),
1624            }
1625        }
1626    }
1627
1628    /// Builds a hedging algo and the `serve` that drives it: the winner is gated behind the
1629    /// loser starting (so the loser's serve is guaranteed in flight when the winner wins),
1630    /// and the loser finishes after `loser_delay` — or never, when `None`.
1631    fn hedge(loser_delay: Option<std::time::Duration>) -> (Arc<dyn Algorithm>, impl Serve) {
1632        let started = Arc::new(tokio::sync::Notify::new());
1633        let algo = Arc::new(Hedge {
1634            winner: "winner".to_string(),
1635            loser: "loser".to_string(),
1636        });
1637        let serve = move |target: ModelId, _request: Request| {
1638            let started = started.clone();
1639            async move {
1640                if target == "loser" {
1641                    started.notify_one();
1642                    match loser_delay {
1643                        Some(delay) => tokio::time::sleep(delay).await,
1644                        None => std::future::pending::<()>().await,
1645                    }
1646                } else {
1647                    started.notified().await;
1648                }
1649                Ok(reply(target))
1650            }
1651        };
1652        (algo, serve)
1653    }
1654
1655    #[tokio::test]
1656    async fn run_returns_the_winner_without_a_late_loser_overwriting_it() -> Result<()> {
1657        // The loser responds 50ms after the winner has already won. `run` must return the
1658        // winner, not the loser's `respond`-to-a-dropped-receiver error.
1659        let (algo, serve) = hedge(Some(std::time::Duration::from_millis(50)));
1660        let (_, response) = test_drive(algo, request(), serve).await?;
1661        assert_eq!(
1662            response
1663                .llm_response
1664                .as_agg()
1665                .map(completion_text)
1666                .unwrap_or_default(),
1667            "winner"
1668        );
1669        Ok(())
1670    }
1671
1672    #[tokio::test]
1673    async fn run_returns_the_winner_without_hanging_on_a_pending_loser() -> Result<()> {
1674        // The loser never resolves. `run` must return the winner promptly, not hang
1675        // waiting for the in-flight loser.
1676        let (algo, serve) = hedge(None);
1677        let run = test_drive(algo, request(), serve);
1678        let (_, response) = tokio::time::timeout(std::time::Duration::from_secs(1), run)
1679            .await
1680            .map_err(|error| LibsyError::external("waiting for pending loser", error))??;
1681        assert_eq!(
1682            response
1683                .llm_response
1684                .as_agg()
1685                .map(completion_text)
1686                .unwrap_or_default(),
1687            "winner"
1688        );
1689        Ok(())
1690    }
1691
1692    #[tokio::test]
1693    async fn run_surfaces_a_terminal_error_with_many_calls_in_flight() -> Result<()> {
1694        use std::sync::atomic::{AtomicUsize, Ordering};
1695
1696        // A large fan-out (10 matched the old, now-removed concurrency cap). The terminal
1697        // error must still reach the caller with all of these calls pending.
1698        const N: usize = 10;
1699
1700        // Fans out N calls, then errors as soon as all N are in flight — exercising a
1701        // terminal failure emitted while the offloaded calls are still pending.
1702        struct FanOutThenError {
1703            all_started: Arc<tokio::sync::Notify>,
1704            n: usize,
1705        }
1706
1707        #[async_trait]
1708        impl Algorithm for FanOutThenError {
1709            fn name(&self) -> &str {
1710                "fan_out_then_error"
1711            }
1712
1713            async fn route(
1714                self: Arc<Self>,
1715                driver: Driver,
1716                request: Request,
1717            ) -> Result<RoutingOutcome> {
1718                let offloads = futures::future::join_all(
1719                    (0..self.n)
1720                        .map(|i| driver.call_model(request.clone(), vec![format!("m{i}").into()])),
1721                );
1722                tokio::select! {
1723                    _ = offloads => Err(test_error("offloads unexpectedly completed")),
1724                    _ = self.all_started.notified() => {
1725                        Err(test_error("terminal error while calls pending"))
1726                    }
1727                }
1728            }
1729        }
1730
1731        let all_started = Arc::new(tokio::sync::Notify::new());
1732        let algo: Arc<dyn Algorithm> = Arc::new(FanOutThenError {
1733            all_started: all_started.clone(),
1734            n: N,
1735        });
1736
1737        // Serving enters each call; once all N are in flight it signals, then pends forever.
1738        let started = Arc::new(AtomicUsize::new(0));
1739        let serve = move |_target: ModelId, _request: Request| {
1740            let started = started.clone();
1741            let all_started = all_started.clone();
1742            async move {
1743                if started.fetch_add(1, Ordering::SeqCst) + 1 == N {
1744                    all_started.notify_one();
1745                }
1746                std::future::pending::<ServeResult>().await
1747            }
1748        };
1749
1750        // With the cap gone, the driver keeps polling the stream even with N calls in
1751        // flight, so the terminal error surfaces promptly instead of hanging.
1752        let run = test_drive(algo, request(), serve);
1753        let result = tokio::time::timeout(std::time::Duration::from_millis(500), run)
1754            .await
1755            .map_err(|error| {
1756                LibsyError::external("waiting for terminal error with full call cap", error)
1757            })?;
1758        match result {
1759            Ok(_) => Err(test_error("expected the terminal error, got a response")),
1760            Err(err) => {
1761                assert!(
1762                    err.to_string()
1763                        .contains("terminal error while calls pending")
1764                );
1765                Ok(())
1766            }
1767        }
1768    }
1769}