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