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        )
423    )]
424    pub(crate) async fn call_model_with_error_recovery(
425        &self,
426        mut request: Request,
427        models: Vec<ModelId>,
428        recover_errors: bool,
429    ) -> Result<Response> {
430        let Some(selected_model_id) = models.first() else {
431            return Err(LibsyError::NoTargets);
432        };
433        self.models.check_targets(&models)?;
434        request.llm_request.model = Some(selected_model_id.to_string());
435        let started = Instant::now();
436        let (reply, response) = oneshot::channel::<Result<Response>>();
437        let call = CallModel {
438            algorithm: self.algorithm.clone(),
439            request,
440            models,
441            recover_errors,
442            reply: Some(reply),
443            started,
444            span: tracing::Span::current(),
445        };
446        self.step_tx
447            .send(Ok(Step::CallModel(Box::new(call))))
448            .await
449            .map_err(|_| DriverError::StreamClosed)?;
450        response
451            .await
452            .map_err(|_| LibsyError::from(DriverError::ResponseDropped))?
453    }
454
455    /// Override the request's model with `model` and wait for the host to reply.
456    #[tracing::instrument(
457        target = "libsy",
458        name = "libsy.decision_call",
459        skip_all,
460        fields(
461            algorithm = self.algorithm,
462            selected_model = %model,
463            openinference.span.kind = "CHAIN",
464            outcome = tracing::field::Empty,
465            input_tokens = tracing::field::Empty,
466            output_tokens = tracing::field::Empty,
467            total_tokens = tracing::field::Empty,
468            reasoning_tokens = tracing::field::Empty,
469            gen_ai.response.id = tracing::field::Empty,
470            gen_ai.response.model = tracing::field::Empty,
471        ),
472    )]
473    pub async fn call_decision(
474        &self,
475        mut request: DecisionRequest,
476        model: ModelId,
477    ) -> Result<DecisionResponse> {
478        self.models.check_targets(std::slice::from_ref(&model))?;
479        request.model = Some(model.clone());
480        let started = Instant::now();
481        let (reply, response) = oneshot::channel();
482        let call = CallDecision {
483            algorithm: self.algorithm.clone(),
484            request,
485            model,
486            reply: Some(reply),
487            started,
488            span: tracing::Span::current(),
489        };
490        self.step_tx
491            .send(Ok(Step::CallDecision(Box::new(call))))
492            .await
493            .map_err(|_| DriverError::StreamClosed)?;
494        response
495            .await
496            .map_err(|_| LibsyError::from(DriverError::ResponseDropped))?
497    }
498
499    /// The available models for this category, typically ordered best-first.
500    pub fn models_for(&self, category: &Category) -> &[ModelId] {
501        match self.scope {
502            Scope::Parent => self.models.models_for(category),
503            Scope::Subagent => self.models.subagent_models_for(category),
504        }
505    }
506
507    /// The first available model for `category`.
508    pub fn first_model_for(&self, category: &Category) -> Result<&ModelId> {
509        self.models_for(category)
510            .first()
511            .ok_or_else(|| LibsyError::AlgorithmError {
512                message: format!("no models available for category {}", category.as_str()),
513            })
514    }
515
516    /// A driver scoped to delegated sub-agent work: its categories are the
517    /// sub-agent's own, and the parent's are no longer reachable through it.
518    pub fn for_subagent(&self) -> Result<Self> {
519        if self.models.subagent.is_none() {
520            return Err(LibsyError::AlgorithmError {
521                message: "delegated work has no sub-agent models".to_string(),
522            });
523        }
524        Ok(Self {
525            scope: Scope::Subagent,
526            ..self.clone()
527        })
528    }
529
530    /// Emit the terminal step: [`Step::Done`] on `Ok`, or an `Err` stream
531    /// item on failure. Internal: called once by [`run_stream`](Algorithm::run_stream)
532    /// when the algorithm finishes.
533    pub(crate) async fn finish(&self, result: Result<RoutingOutcome>) -> Result<()> {
534        let result = result.map(|mut outcome| {
535            let metadata = outcome.metadata.get_or_insert_with(|| {
536                crate::OutcomeMetadata::new(self.algorithm.clone(), self.evidence.lock().take())
537            });
538            observability::record_outcome(metadata, &outcome.selected_model_ids);
539            outcome
540        });
541        let selected_model = result
542            .as_ref()
543            .ok()
544            .and_then(|outcome| outcome.selected_model_id().ok().cloned());
545        let step = result.map(|outcome| Step::Done(Box::new(outcome)));
546        self.step_tx
547            .send(step)
548            .await
549            .map_err(|_| DriverError::StreamClosed)?;
550        if let Some(selected_model) = selected_model {
551            observability::record_decision(&self.algorithm, &selected_model);
552        }
553        Ok(())
554    }
555}
556
557/// Work requests and the terminal outcome of [`Algorithm::run_stream`].
558/// Each call carries its own reply channel, so the host can execute calls concurrently.
559pub enum Step {
560    /// The algorithm needs this model call performed. The host serves it and fulfills
561    /// it with [`CallModel::respond`]. Boxed: it is by far the largest variant.
562    CallModel(Box<CallModel>),
563    /// The host fulfills this decision call with [`CallDecision::respond`].
564    CallDecision(Box<CallDecision>),
565    /// The algorithm finished with its routing outcome — the last step of a run.
566    Done(Box<RoutingOutcome>),
567}
568
569/// Work dispatched through one host handler, with a typed reply channel per variant.
570/// [`drive`] handles terminal outcomes separately; the handler only executes calls.
571pub enum Call {
572    /// Fulfilled with an LLM [`Response`].
573    Model(Box<CallModel>),
574    /// Fulfilled with a [`DecisionResponse`].
575    Decision(Box<CallDecision>),
576}
577
578/// Drive [`Algorithm::run_stream`] to its routing outcome, serving calls concurrently.
579///
580/// The shared execution loop lets hosts supply clients or test doubles through `serve`
581/// without reimplementing stream handling and cancellation.
582///
583/// `serve` owns each [`Call`] and performs its I/O. Passing an error to `respond`
584/// lets the algorithm fall back; returning an error from `serve` aborts the run.
585///
586/// Dropping this future cancels the algorithm and pending handler futures.
587pub async fn drive<F, Fut>(
588    algorithm: Arc<dyn Algorithm>,
589    request: Request,
590    models: Arc<RuntimeModels>,
591    serve: F,
592) -> Result<RoutingOutcome>
593where
594    F: Fn(Call) -> Fut,
595    Fut: Future<Output = Result<()>>,
596{
597    let stream = algorithm.run_stream(request, Arc::clone(&models));
598    tokio::pin!(stream);
599
600    let mut in_flight = futures::stream::FuturesUnordered::new();
601    let mut final_outcome: Option<RoutingOutcome> = None;
602
603    loop {
604        tokio::select! {
605            Some(result) = in_flight.next() => match result {
606                Ok(()) => {}, // Call completed successfully
607                Err(err) => return Err(err), // Call failed, propagate the error
608            },
609            step = stream.next() => {
610                match step {
611                    None => break, // stream has ended, no more steps
612                    Some(item) => match item? {
613                        Step::CallModel(call) => {
614                            models.check_targets(&call.models)?;
615                            in_flight.push(serve(Call::Model(call)));
616                        }
617                        Step::CallDecision(call) => {
618                            models.check_targets(std::slice::from_ref(&call.model))?;
619                            in_flight.push(serve(Call::Decision(call)));
620                        }
621                        Step::Done(outcome) => {
622                            if outcome.response.is_none() {
623                                models.check_targets(&outcome.selected_model_ids)?;
624                            }
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(self: Arc<Self>, request: Request, models: Arc<RuntimeModels>) -> StepStream {
751        let (driver, step_rx) = Driver::new(self.name(), models);
752        let span = observability::run_span(self.name(), &request);
753        let handle = tokio::spawn(
754            async move {
755                let algorithm = self.name().to_string();
756                // Catch a panicking algorithm so the run still publishes a terminal step.
757                let route = AssertUnwindSafe(self.route(driver.clone(), request)).catch_unwind();
758                let result = observability::observe_run(&algorithm, async {
759                    let outcome = route.await.unwrap_or_else(|payload| {
760                        Err(LibsyError::AlgorithmError {
761                            message: format!(
762                                "algorithm task panicked: {}",
763                                panic_message(payload.as_ref())
764                            ),
765                        })
766                    })?;
767                    if outcome.response.is_none() {
768                        driver.models.check_targets(&outcome.selected_model_ids)?;
769                    }
770                    Ok(outcome)
771                })
772                .await;
773
774                let _ = driver.finish(result).await;
775            }
776            .instrument(span),
777        );
778        // Dropping the stream aborts the algorithm task when its consumer goes away.
779        let abort_guard = AbortOnDrop(handle.abort_handle());
780        Box::pin(ReceiverStream::new(step_rx).map(move |step| {
781            // link abort guard to stream
782            let _keep_alive = &abort_guard;
783            step
784        }))
785    }
786}
787
788#[cfg(test)]
789mod tests {
790    use std::collections::HashMap;
791
792    use super::*;
793    use crate::core::testing::{Serve, ServeResult, echo, reply, serve_decision, test_drive};
794    use futures::StreamExt;
795    use switchyard_protocol::{
796        LlmResponse, LlmResponseChunk, completion_text, text_request, text_response,
797    };
798
799    #[derive(Debug, thiserror::Error)]
800    #[error("{0}")]
801    struct TestError(&'static str);
802
803    fn test_error(message: &'static str) -> LibsyError {
804        LibsyError::external("test", TestError(message))
805    }
806
807    /// Trivial algo used only to exercise the orchestrator: calls the first target
808    /// and returns its response as the routing outcome.
809    struct TestAlgo {
810        target_set: Vec<ModelId>,
811    }
812
813    #[async_trait]
814    impl Algorithm for TestAlgo {
815        fn name(&self) -> &str {
816            "test"
817        }
818
819        async fn route(
820            self: Arc<Self>,
821            driver: Driver,
822            request: Request,
823        ) -> Result<RoutingOutcome> {
824            let target = self
825                .target_set
826                .first()
827                .ok_or(LibsyError::NoTargets)?
828                .clone();
829            let response = driver
830                .call_model(request.clone(), vec![target.clone()])
831                .await?;
832            driver.set_evidence(serde_json::json!({"source": "test"}));
833            driver.set_evidence_if_empty(serde_json::json!({"source": "ignored"}));
834            Ok(RoutingOutcome::answered(target, request, response))
835        }
836    }
837
838    /// Build a shared `TestAlgo` over the given target set.
839    fn orch(target_set: Vec<ModelId>) -> Arc<dyn Algorithm> {
840        Arc::new(TestAlgo { target_set })
841    }
842
843    fn request() -> Request {
844        Request {
845            llm_request: text_request(Some("auto".to_string()), "hi".to_string()),
846            raw_request: None,
847            metadata: None,
848        }
849    }
850
851    #[test]
852    fn routing_outcome_constructors_stamp_selection_and_preserve_payloads() {
853        let outcome = RoutingOutcome::route_to(
854            "selected".into(),
855            target_set(&["fallback-one", "fallback-two"]),
856            request(),
857        );
858
859        assert_eq!(
860            outcome.selected_model_ids,
861            target_set(&["selected", "fallback-one", "fallback-two"])
862        );
863        assert_eq!(outcome.request.model_id().as_deref(), Some("selected"));
864        assert!(outcome.response.is_none());
865        assert!(outcome.metadata.is_none());
866
867        let outcome = RoutingOutcome::route_to("only".into(), Vec::new(), request());
868        assert_eq!(outcome.selected_model_ids, target_set(&["only"]));
869
870        let outcome = RoutingOutcome::answered(
871            "answered".into(),
872            request(),
873            Response {
874                llm_response: LlmResponse::Agg(text_response(None, "existing")),
875                metadata: None,
876                upstream_headers: http::HeaderMap::new(),
877            },
878        );
879
880        assert_eq!(outcome.selected_model_ids, target_set(&["answered"]));
881        assert_eq!(outcome.request.model_id().as_deref(), Some("answered"));
882        assert_eq!(
883            outcome
884                .response
885                .as_ref()
886                .and_then(|response| response.llm_response.as_agg())
887                .map(completion_text),
888            Some("existing".to_string())
889        );
890    }
891
892    fn target_set(names: &[&str]) -> Vec<ModelId> {
893        names.iter().map(|name| ModelId::from(*name)).collect()
894    }
895
896    enum TargetProbe {
897        Model(Vec<ModelId>),
898        Decision(ModelId),
899        Completion(Vec<ModelId>),
900        Subagent,
901    }
902
903    #[async_trait]
904    impl Algorithm for TargetProbe {
905        fn name(&self) -> &str {
906            "target_probe"
907        }
908
909        async fn route(
910            self: Arc<Self>,
911            driver: Driver,
912            request: Request,
913        ) -> Result<RoutingOutcome> {
914            match &*self {
915                Self::Model(models) => {
916                    let response = driver.call_model(request.clone(), models.clone()).await?;
917                    Ok(RoutingOutcome::answered(
918                        models[0].clone(),
919                        request,
920                        response,
921                    ))
922                }
923                Self::Decision(model) => {
924                    driver
925                        .call_decision(
926                            DecisionRequest {
927                                model: None,
928                                context: Value::Null,
929                                questions: Default::default(),
930                            },
931                            model.clone(),
932                        )
933                        .await?;
934                    Ok(RoutingOutcome::route_to("allowed".into(), vec![], request))
935                }
936                Self::Completion(models) => Ok(RoutingOutcome::route_to(
937                    models[0].clone(),
938                    models[1..].to_vec(),
939                    request,
940                )),
941                Self::Subagent => Ok(RoutingOutcome::route_to(
942                    driver
943                        .for_subagent()?
944                        .first_model_for(&Category::Any)?
945                        .clone(),
946                    vec![],
947                    request,
948                )),
949            }
950        }
951    }
952
953    struct RewrittenSteps(Arc<TargetProbe>);
954
955    #[async_trait]
956    impl Algorithm for RewrittenSteps {
957        fn name(&self) -> &str {
958            "rewritten_steps"
959        }
960
961        async fn route(
962            self: Arc<Self>,
963            driver: Driver,
964            request: Request,
965        ) -> Result<RoutingOutcome> {
966            Arc::clone(&self.0).route(driver, request).await
967        }
968
969        fn run_stream(self: Arc<Self>, request: Request, models: Arc<RuntimeModels>) -> StepStream {
970            Box::pin(Arc::clone(&self.0).run_stream(request, models).map(|step| {
971                step.map(|mut step| {
972                    match &mut step {
973                        Step::CallModel(call) => call.models.push("outside".into()),
974                        Step::CallDecision(call) => call.model = "outside".into(),
975                        Step::Done(outcome) => outcome.selected_model_ids.push("outside".into()),
976                    }
977                    step
978                })
979            }))
980        }
981    }
982
983    fn restricted_models() -> RuntimeModels {
984        RuntimeModels::new(HashMap::from([
985            (Category::Any, target_set(&["allowed", "fallback"])),
986            (Category::Judge, target_set(&["judge"])),
987        ]))
988        .with_target_restriction()
989        .with_subagent(HashMap::from([(Category::Any, target_set(&["subagent"]))]))
990    }
991
992    // No forbidden call or completion candidate reaches the host.
993    #[tokio::test]
994    async fn target_restriction_blocks_undeclared_work() {
995        for probe in [
996            TargetProbe::Model(target_set(&["outside"])),
997            TargetProbe::Model(target_set(&["allowed", "outside"])),
998            TargetProbe::Decision("outside".into()),
999            TargetProbe::Completion(target_set(&["outside"])),
1000            TargetProbe::Completion(target_set(&["allowed", "outside"])),
1001        ] {
1002            let mut stream = Arc::new(probe).run_stream(request(), Arc::new(restricted_models()));
1003            assert!(matches!(
1004                stream.next().await,
1005                Some(Err(LibsyError::TargetNotFound { target })) if target == "outside"
1006            ));
1007            assert!(stream.next().await.is_none());
1008        }
1009        // A stream override can rewrite valid steps after the default driver checked them.
1010        for probe in [
1011            TargetProbe::Model(target_set(&["allowed"])),
1012            TargetProbe::Decision("judge".into()),
1013            TargetProbe::Completion(target_set(&["allowed"])),
1014        ] {
1015            let result = drive(
1016                Arc::new(RewrittenSteps(Arc::new(probe))),
1017                request(),
1018                Arc::new(restricted_models()),
1019                |_| async { panic!("undeclared work must not reach the host") },
1020            )
1021            .await;
1022            assert!(
1023                matches!(result, Err(LibsyError::TargetNotFound { target }) if target == "outside")
1024            );
1025        }
1026    }
1027
1028    // Judges, fallbacks, delegated targets and local answers remain valid.
1029    #[tokio::test]
1030    async fn target_restriction_preserves_allowed_work() -> Result<()> {
1031        for (probe, expected) in [
1032            (
1033                TargetProbe::Model(target_set(&["allowed", "fallback"])),
1034                &["allowed"][..],
1035            ),
1036            (TargetProbe::Model(target_set(&["judge"])), &["judge"][..]),
1037            (TargetProbe::Decision("judge".into()), &["allowed"][..]),
1038            (
1039                TargetProbe::Completion(target_set(&["allowed", "fallback"])),
1040                &["allowed", "fallback"][..],
1041            ),
1042            (TargetProbe::Subagent, &["subagent"][..]),
1043        ] {
1044            let outcome = drive(
1045                Arc::new(probe),
1046                request(),
1047                Arc::new(restricted_models()),
1048                |call| async move {
1049                    match call {
1050                        Call::Model(call) => {
1051                            call.respond(std::future::ready(Ok(reply("ok")))).await
1052                        }
1053                        Call::Decision(call) => serve_decision(*call).await,
1054                    }
1055                },
1056            )
1057            .await?;
1058            assert_eq!(outcome.selected_model_ids, target_set(expected));
1059        }
1060        let outcome = drive(
1061            Arc::new(crate::Noop {}),
1062            request(),
1063            Arc::new(RuntimeModels::default().with_target_restriction()),
1064            |_| async { panic!("Noop must not request model calls") },
1065        )
1066        .await?;
1067        assert!(outcome.response.is_some());
1068        Ok(())
1069    }
1070
1071    #[tokio::test]
1072    async fn target_restriction_is_opt_in() {
1073        let mut stream = orch(target_set(&["outside"]))
1074            .run_stream(request(), Arc::new(RuntimeModels::default()));
1075        assert!(matches!(stream.next().await, Some(Ok(Step::CallModel(_)))));
1076    }
1077
1078    #[tokio::test]
1079    async fn typed_driver_preserves_call_and_stream_boundaries() -> Result<()> {
1080        tokio::time::timeout(std::time::Duration::from_secs(1), async {
1081            // Distinct oneshots keep reverse-order replies paired with their producers, and a
1082            // retained call remains pending until the host responds.
1083            let (driver, mut step_rx) = Driver::new("test", Arc::new(RuntimeModels::default()));
1084            let first_driver = driver.clone();
1085            let mut first = tokio::spawn(async move {
1086                first_driver
1087                    .call_model(request(), vec![ModelId::from("first")])
1088                    .await
1089            });
1090            let second = tokio::spawn(async move {
1091                driver
1092                    .call_model(request(), vec![ModelId::from("second")])
1093                    .await
1094            });
1095
1096            let mut calls = HashMap::new();
1097            for _ in 0..2 {
1098                let step = step_rx.recv().await.ok_or(DriverError::StreamClosed)??;
1099                let Step::CallModel(call) = step else {
1100                    return Err(test_error("expected a CallModel step"));
1101                };
1102                let selected_model = call
1103                    .models
1104                    .first()
1105                    .ok_or_else(|| test_error("model call has no candidates"))?
1106                    .to_string();
1107                calls.insert(selected_model, call);
1108            }
1109            assert!(
1110                tokio::time::timeout(std::time::Duration::from_millis(20), &mut first)
1111                    .await
1112                    .is_err(),
1113                "call completed before the host responded"
1114            );
1115            calls
1116                .remove("second")
1117                .ok_or_else(|| test_error("missing second call"))?
1118                .respond(std::future::ready(Ok(reply("second response")))).await?;
1119            calls
1120                .remove("first")
1121                .ok_or_else(|| test_error("missing first call"))?
1122                .respond(std::future::ready(Ok(reply("first response")))).await?;
1123
1124            let first_response = first
1125                .await
1126                .map_err(|source| LibsyError::external("joining a test task", source))??;
1127            let second_response = second
1128                .await
1129                .map_err(|source| LibsyError::external("joining a test task", source))??;
1130            assert_eq!(
1131                first_response.llm_response.as_agg().map(completion_text),
1132                Some("first response".to_string())
1133            );
1134            assert_eq!(
1135                second_response.llm_response.as_agg().map(completion_text),
1136                Some("second response".to_string())
1137            );
1138
1139            // Dropping the host-facing promise closes only that call's reply channel.
1140            let (driver, mut step_rx) = Driver::new("test", Arc::new(RuntimeModels::default()));
1141            let producer = tokio::spawn(async move {
1142                driver
1143                    .call_model(request(), vec![ModelId::from("dropped")])
1144                    .await
1145            });
1146            let step = step_rx.recv().await.ok_or(DriverError::StreamClosed)??;
1147            let Step::CallModel(call) = step else {
1148                return Err(test_error("expected a CallModel step"));
1149            };
1150            drop(call);
1151            let result = producer
1152                .await
1153                .map_err(|source| LibsyError::external("joining a test task", source))?;
1154            assert!(matches!(
1155                result,
1156                Err(LibsyError::Driver(DriverError::ResponseDropped))
1157            ));
1158
1159            // A standalone driver reports the typed step receiver disappearing at its next call.
1160            let (driver, step_rx) = Driver::new("test", Arc::new(RuntimeModels::default()));
1161            drop(step_rx);
1162            let result = driver
1163                .call_model(request(), vec![ModelId::from("closed")])
1164                .await;
1165            assert!(matches!(
1166                result,
1167                Err(LibsyError::Driver(DriverError::StreamClosed))
1168            ));
1169
1170            fn decision_request() -> DecisionRequest {
1171                DecisionRequest {
1172                    model: Some("overwritten".into()),
1173                    context: serde_json::json!({"task": "choose a route"}),
1174                    questions: Default::default(),
1175                }
1176            }
1177
1178            fn decision_response() -> DecisionResponse {
1179                DecisionResponse {
1180                    id: Some("decision-1".to_string()),
1181                    model: Some("provider-model".into()),
1182                    answers: [(
1183                        "p_solve".to_string(),
1184                        switchyard_protocol::DecisionAnswer {
1185                            value: switchyard_protocol::DecisionValue::Boolean(
1186                                switchyard_protocol::BooleanEstimate::ProbabilityTrue(
1187                                    switchyard_protocol::Probability(0.8),
1188                                ),
1189                            ),
1190                            provider_confidence: None,
1191                        },
1192                    )]
1193                    .into(),
1194                    usage: Default::default(),
1195                }
1196            }
1197
1198            struct MixedCalls(&'static str);
1199
1200            #[async_trait]
1201            impl Algorithm for MixedCalls {
1202                fn name(&self) -> &str {
1203                    "mixed"
1204                }
1205
1206                async fn route(
1207                    self: Arc<Self>,
1208                    driver: Driver,
1209                    request: Request,
1210                ) -> Result<RoutingOutcome> {
1211                    let (llm, decision) = tokio::join!(
1212                        driver.call_model(request.clone(), vec!["llm".into()]),
1213                        driver.call_decision(decision_request(), "decision".into()),
1214                    );
1215                    assert_eq!(
1216                        llm?.llm_response.as_agg().map(completion_text),
1217                        Some("llm reply".into())
1218                    );
1219                    match self.0 {
1220                        "mock" => assert_eq!(decision?, DecisionResponse {
1221                            id: None,
1222                            model: Some("decision".into()),
1223                            answers: Default::default(),
1224                            usage: Default::default(),
1225                        }),
1226                        "reply" => assert_eq!(decision?, decision_response()),
1227                        "error" => assert!(matches!(
1228                            decision,
1229                            Err(LibsyError::AlgorithmError { message }) if message == "provider failed"
1230                        )),
1231                        "drop" => assert!(matches!(
1232                            decision,
1233                            Err(LibsyError::Driver(DriverError::ResponseDropped))
1234                        )),
1235                        "abort" => return std::future::pending().await,
1236                        _ => unreachable!(),
1237                    }
1238                    Ok(RoutingOutcome::route_to("answer".into(), vec![], request))
1239                }
1240            }
1241
1242            for mode in ["mock", "reply", "error", "drop", "abort"] {
1243                // Serial dispatch would deadlock here and fail the enclosing timeout.
1244                let barrier = Arc::new(tokio::sync::Barrier::new(2));
1245                let outcome = drive(
1246                    Arc::new(MixedCalls(mode)),
1247                    request(),
1248                    Arc::new(RuntimeModels::default()),
1249                    move |call| {
1250                        let barrier = barrier.clone();
1251                        async move {
1252                            barrier.wait().await;
1253                            let call = match call {
1254                                Call::Model(call) => return call.respond(std::future::ready(Ok(reply("llm reply")))).await,
1255                                Call::Decision(call) => *call,
1256                            };
1257                            assert_eq!(call.algorithm, "mixed");
1258                            assert_eq!(call.model, "decision");
1259                            assert_eq!(call.request.model, Some("decision".into()));
1260                            assert_eq!(call.request.context, decision_request().context);
1261                            match mode {
1262                                "mock" => serve_decision(call).await,
1263                                "reply" => call.respond(std::future::ready(Ok(decision_response()))).await,
1264                                "error" => call.respond(std::future::ready(Err(LibsyError::AlgorithmError {
1265                                    message: "provider failed".into(),
1266                                }))).await,
1267                                "drop" => {
1268                                    drop(call);
1269                                    Ok(())
1270                                }
1271                                "abort" => call.fail(test_error("host aborted")),
1272                                _ => unreachable!(),
1273                            }
1274                        }
1275                    },
1276                )
1277                .await;
1278                if mode == "abort" {
1279                    assert!(matches!(outcome, Err(LibsyError::External { .. })));
1280                } else {
1281                    assert_eq!(outcome?.selected_model_id()?, "answer");
1282                }
1283            }
1284
1285            let (driver, mut steps) = Driver::new("test", Arc::new(RuntimeModels::default()));
1286            let mut pending = Box::pin(driver.call_decision(decision_request(), "decision".into()));
1287            assert!(futures::poll!(&mut pending).is_pending());
1288            let Some(Ok(Step::CallDecision(call))) = steps.recv().await else {
1289                return Err(test_error("expected a decision call"));
1290            };
1291            drop(pending);
1292            call.respond(async { panic!("cancelled work must not be polled") }).await?;
1293            drop(steps);
1294            assert!(matches!(
1295                driver.call_decision(decision_request(), "decision".into()).await,
1296                Err(LibsyError::Driver(DriverError::StreamClosed))
1297            ));
1298            Ok(())
1299        })
1300        .await
1301        .map_err(|error| LibsyError::external("waiting for typed driver boundaries", error))?
1302    }
1303
1304    #[test]
1305    fn target_lookup_returns_the_missing_target() {
1306        let error = ensure_model_is_target(&target_set(&[]), &ModelId::from("missing")).err();
1307        assert!(matches!(
1308            error,
1309            Some(LibsyError::TargetNotFound { target }) if target == "missing"
1310        ));
1311    }
1312
1313    /// Build a single-target algo, plus a `serve` that answers it as a token stream
1314    /// replaying `chunks` in order (as `Ok` items).
1315    fn streaming_orch(chunks: Vec<LlmResponseChunk>) -> (Arc<dyn Algorithm>, impl Serve) {
1316        let algo = orch(target_set(&["stream/model"]));
1317        let serve = move |_target: ModelId, _request: Request| {
1318            let chunks = chunks.clone();
1319            async move {
1320                let stream =
1321                    futures::stream::iter(chunks.into_iter().map(|chunk| Ok(chunk.into()))).boxed();
1322                Ok(Response {
1323                    llm_response: LlmResponse::Stream(stream),
1324                    metadata: None,
1325                    upstream_headers: http::HeaderMap::new(),
1326                })
1327            }
1328        };
1329        (algo, serve)
1330    }
1331
1332    #[tokio::test]
1333    async fn run_returns_a_streamed_response_the_caller_aggregates() -> Result<()> {
1334        // A streaming client -> its chunks flow through the promise and `Done`,
1335        // and `run` returns the live stream untouched for the caller to fold.
1336        let (orch, serve) = streaming_orch(vec![
1337            LlmResponseChunk::MessageStart {
1338                id: Some("m1".to_string()),
1339                model: Some("stream/model".to_string()),
1340            },
1341            LlmResponseChunk::TextDelta {
1342                index: 0,
1343                text: "hel".to_string(),
1344            },
1345            LlmResponseChunk::TextDelta {
1346                index: 0,
1347                text: "lo".to_string(),
1348            },
1349            LlmResponseChunk::MessageStop {
1350                reason: Some("stop".to_string()),
1351            },
1352        ]);
1353        let (selected_model, response) = test_drive(orch, request(), serve).await?;
1354        // The run handed back the live stream; the caller folds it to a buffered aggregate.
1355        let agg = response
1356            .llm_response
1357            .into_agg()
1358            .await
1359            .map_err(|error| LibsyError::external("aggregating response stream", error))?;
1360        assert_eq!(completion_text(&agg), "hello");
1361        assert_eq!(agg.model.as_deref(), Some("stream/model"));
1362        assert_eq!(selected_model, "stream/model");
1363        Ok(())
1364    }
1365
1366    #[tokio::test]
1367    async fn aggregating_a_streamed_response_propagates_a_mid_stream_error() -> Result<()> {
1368        // The run succeeds and returns the stream; the in-band `Error` chunk surfaces only
1369        // when the caller aggregates it.
1370        let (orch, serve) = streaming_orch(vec![
1371            LlmResponseChunk::TextDelta {
1372                index: 0,
1373                text: "partial".to_string(),
1374            },
1375            LlmResponseChunk::StreamError {
1376                message: "upstream exploded".to_string(),
1377            },
1378        ]);
1379        let (_, response) = test_drive(orch, request(), serve).await?;
1380        match response.llm_response.into_agg().await {
1381            Ok(_) => panic!("expected a mid-stream error, got an aggregate"),
1382            Err(err) => {
1383                assert!(err.to_string().contains("upstream exploded"));
1384                Ok(())
1385            }
1386        }
1387    }
1388
1389    #[tokio::test]
1390    async fn run_offloads_via_promise_then_finishes() -> Result<()> {
1391        // Every call is offloaded via a promise the orchestrator surfaces as a
1392        // `CallModel` step for us to fulfill.
1393        let stream = orch(target_set(&["offload/model"]))
1394            .run_stream(request(), Arc::new(RuntimeModels::default()));
1395        tokio::pin!(stream);
1396
1397        let mut saw_call = false;
1398        let mut final_completion = None;
1399        while let Some(step) = stream.next().await {
1400            match step? {
1401                Step::CallDecision(_) => return Err(test_error("unexpected decision call")),
1402                Step::CallModel(call) => {
1403                    saw_call = true;
1404                    assert_eq!(call.models, vec![ModelId::from("offload/model")]);
1405                    // Fulfilling the promise is the "real" model call the caller makes.
1406                    call.respond(std::future::ready(Ok(Response {
1407                        llm_response: LlmResponse::Agg(text_response(
1408                            None,
1409                            "fulfilled".to_string(),
1410                        )),
1411                        metadata: None,
1412                        upstream_headers: http::HeaderMap::new(),
1413                    })))
1414                    .await?;
1415                }
1416                Step::Done(outcome) => {
1417                    let metadata = outcome
1418                        .metadata
1419                        .as_ref()
1420                        .expect("run_stream should attach outcome metadata");
1421                    assert_eq!(metadata.algorithm, "test");
1422                    assert_eq!(
1423                        uuid::Uuid::parse_str(metadata.outcome_id())
1424                            .expect("outcome id should be a UUID")
1425                            .get_version_num(),
1426                        7
1427                    );
1428                    assert_eq!(
1429                        metadata.evidence,
1430                        Some(serde_json::json!({"source": "test"}))
1431                    );
1432                    let response = outcome
1433                        .response
1434                        .ok_or_else(|| test_error("expected an answered outcome"))?;
1435                    final_completion = Some(
1436                        response
1437                            .llm_response
1438                            .as_agg()
1439                            .map(completion_text)
1440                            .unwrap_or_default(),
1441                    );
1442                }
1443            }
1444        }
1445
1446        assert!(saw_call, "expected a CallModel step before Done");
1447        assert_eq!(
1448            final_completion.ok_or_else(|| test_error("no Done step"))?,
1449            "fulfilled"
1450        );
1451        Ok(())
1452    }
1453
1454    #[tokio::test(flavor = "multi_thread", worker_threads = 12)]
1455    async fn requests_are_processed_in_parallel() -> Result<()> {
1456        use std::time::Duration;
1457        use tokio::sync::Barrier;
1458
1459        const N: usize = 12;
1460
1461        // Serving blocks until all N concurrent calls have arrived. If requests were
1462        // serialized (one algorithm behind a `Mutex`), only one call could be in flight,
1463        // the barrier would never reach N, and the test would time out. It passes only
1464        // because the shared algorithm is driven concurrently across requests.
1465        let barrier = Arc::new(Barrier::new(N));
1466        // One shared algorithm driven by many concurrent requests.
1467        let algo = orch(target_set(&["m"]));
1468
1469        let mut handles = Vec::new();
1470        for _ in 0..N {
1471            let algo = algo.clone();
1472            let barrier = barrier.clone();
1473            let serve = move |target: ModelId, _request: Request| {
1474                let barrier = barrier.clone();
1475                async move {
1476                    barrier.wait().await;
1477                    Ok(reply(target))
1478                }
1479            };
1480            handles.push(tokio::spawn(async move {
1481                test_drive(algo, request(), serve)
1482                    .await
1483                    .map(|(_, response)| {
1484                        response
1485                            .llm_response
1486                            .as_agg()
1487                            .map(completion_text)
1488                            .unwrap_or_default()
1489                    })
1490            }));
1491        }
1492
1493        for handle in handles {
1494            // The timeout turns a serialization deadlock into a failure, not a hang.
1495            let completion = tokio::time::timeout(Duration::from_secs(5), handle)
1496                .await
1497                .map_err(|error| LibsyError::external("waiting for test task", error))?
1498                .map_err(|source| LibsyError::external("joining a test task", source))??;
1499            assert_eq!(completion, "m");
1500        }
1501        Ok(())
1502    }
1503
1504    #[tokio::test]
1505    async fn offload_error_propagates_back_to_the_algorithm() -> Result<()> {
1506        // A client-less target offloads its call; we fulfill the promise with an
1507        // Err, which must flow back through `call_model_target` into the algorithm and
1508        // out as an error step — not a response.
1509        let stream = orch(target_set(&["offload/model"]))
1510            .run_stream(request(), Arc::new(RuntimeModels::default()));
1511        tokio::pin!(stream);
1512
1513        let mut saw_error = false;
1514        while let Some(step) = stream.next().await {
1515            match step {
1516                Ok(Step::CallDecision(_)) => return Err(test_error("unexpected decision call")),
1517                Ok(Step::CallModel(mut call)) => {
1518                    call.recover_errors = true;
1519                    call.respond(std::future::ready(Err(test_error(
1520                        "upstream model call failed",
1521                    ))))
1522                    .await?;
1523                }
1524                Ok(Step::Done(..)) => {
1525                    return Err(test_error(
1526                        "expected the offload error to propagate, got a response",
1527                    ));
1528                }
1529                Err(err) => {
1530                    // The algorithm's `call_model_target` saw the error via the promise.
1531                    assert!(err.to_string().contains("upstream model call failed"));
1532                    saw_error = true;
1533                }
1534            }
1535        }
1536
1537        assert!(saw_error, "expected an error step");
1538        Ok(())
1539    }
1540
1541    #[tokio::test]
1542    async fn dropping_the_stream_cancels_the_algorithm_task() -> Result<()> {
1543        use std::sync::atomic::{AtomicBool, Ordering};
1544        use std::time::Duration;
1545        use tokio::sync::mpsc;
1546
1547        // Sets a flag when dropped, so we can observe whether the algorithm task was
1548        // cancelled/dropped.
1549        struct DropGuard(Arc<AtomicBool>);
1550        impl Drop for DropGuard {
1551            fn drop(&mut self) {
1552                self.0.store(true, Ordering::SeqCst);
1553            }
1554        }
1555
1556        struct StuckAlgo {
1557            started: mpsc::UnboundedSender<()>,
1558            dropped: Arc<AtomicBool>,
1559        }
1560
1561        #[async_trait]
1562        impl Algorithm for StuckAlgo {
1563            fn name(&self) -> &str {
1564                "stuck"
1565            }
1566
1567            async fn route(
1568                self: Arc<Self>,
1569                _driver: Driver,
1570                _request: Request,
1571            ) -> Result<RoutingOutcome> {
1572                let _guard = DropGuard(self.dropped.clone());
1573                let _ = self.started.send(());
1574                // Await forever without ever touching the driver.
1575                std::future::pending::<()>().await;
1576                unreachable!()
1577            }
1578        }
1579
1580        let (started_tx, mut started_rx) = mpsc::unbounded_channel();
1581        let dropped = Arc::new(AtomicBool::new(false));
1582        let algo: Arc<dyn Algorithm> = Arc::new(StuckAlgo {
1583            started: started_tx,
1584            dropped: dropped.clone(),
1585        });
1586
1587        let stream = algo.run_stream(request(), Arc::new(RuntimeModels::default()));
1588        started_rx
1589            .recv()
1590            .await
1591            .ok_or_else(|| test_error("task never started"))?;
1592        drop(stream);
1593        tokio::time::sleep(Duration::from_millis(100)).await;
1594
1595        assert!(
1596            dropped.load(Ordering::SeqCst),
1597            "algorithm task was NOT cancelled after dropping the stream"
1598        );
1599        Ok(())
1600    }
1601
1602    #[tokio::test]
1603    async fn route_panic_surfaces_as_a_stream_error() -> Result<()> {
1604        // An algorithm whose task panics must surface an `Err` step carrying the panic
1605        // message, not abort the process from an unobserved detached task.
1606        struct Panicky;
1607
1608        #[async_trait]
1609        impl Algorithm for Panicky {
1610            fn name(&self) -> &str {
1611                "panicky"
1612            }
1613
1614            async fn route(
1615                self: Arc<Self>,
1616                _driver: Driver,
1617                _request: Request,
1618            ) -> Result<RoutingOutcome> {
1619                panic!("boom");
1620            }
1621        }
1622
1623        let algo: Arc<dyn Algorithm> = Arc::new(Panicky);
1624        let stream = algo.run_stream(request(), Arc::new(RuntimeModels::default()));
1625        tokio::pin!(stream);
1626
1627        let mut saw_error = false;
1628        while let Some(step) = stream.next().await {
1629            match step {
1630                Err(err) => {
1631                    // The panic message is preserved, not flattened into an opaque failure.
1632                    assert!(err.to_string().contains("algorithm task panicked: boom"));
1633                    saw_error = true;
1634                }
1635                Ok(_) => return Err(test_error("expected the panic to surface as an error step")),
1636            }
1637        }
1638
1639        assert!(saw_error, "expected an error step from the panicked task");
1640        Ok(())
1641    }
1642
1643    /// A panicking algorithm must publish its terminal step even when it left a `Driver`
1644    /// clone alive in another task. That clone holds the step channel open, so a run that
1645    /// merely unwound would never terminate and the consumer would wait forever.
1646    #[tokio::test]
1647    async fn a_panic_with_a_leaked_driver_clone_still_terminates_the_run() -> Result<()> {
1648        struct LeakyPanic;
1649
1650        #[async_trait]
1651        impl Algorithm for LeakyPanic {
1652            fn name(&self) -> &str {
1653                "leaky_panic"
1654            }
1655
1656            async fn route(
1657                self: Arc<Self>,
1658                driver: Driver,
1659                _request: Request,
1660            ) -> Result<RoutingOutcome> {
1661                tokio::spawn(async move {
1662                    // Outlives the panic below, keeping a sender clone alive.
1663                    let _keep_alive = driver;
1664                    std::future::pending::<()>().await;
1665                });
1666                tokio::task::yield_now().await;
1667                panic!("boom");
1668            }
1669        }
1670
1671        let algo: Arc<dyn Algorithm> = Arc::new(LeakyPanic);
1672        // The timeout turns the hang this guards against into a failure rather than a hang.
1673        let result = tokio::time::timeout(
1674            std::time::Duration::from_secs(1),
1675            test_drive(algo, request(), echo()),
1676        )
1677        .await
1678        .map_err(|error| LibsyError::external("waiting for the panicked run to end", error))?;
1679
1680        match result {
1681            Ok(_) => Err(test_error(
1682                "expected the panic to end the run with an error",
1683            )),
1684            Err(err) => {
1685                assert!(err.to_string().contains("algorithm task panicked: boom"));
1686                Ok(())
1687            }
1688        }
1689    }
1690
1691    #[tokio::test]
1692    async fn cancelling_run_cancels_the_algorithm_task() -> Result<()> {
1693        use std::sync::atomic::{AtomicBool, Ordering};
1694        use std::time::Duration;
1695        use tokio::sync::mpsc;
1696
1697        // Sets a flag when dropped, so we can observe whether the algorithm task was
1698        // cancelled once the `run` future driving it is dropped.
1699        struct DropGuard(Arc<AtomicBool>);
1700        impl Drop for DropGuard {
1701            fn drop(&mut self) {
1702                self.0.store(true, Ordering::SeqCst);
1703            }
1704        }
1705
1706        struct StuckAlgo {
1707            started: mpsc::UnboundedSender<()>,
1708            dropped: Arc<AtomicBool>,
1709        }
1710
1711        #[async_trait]
1712        impl Algorithm for StuckAlgo {
1713            fn name(&self) -> &str {
1714                "stuck"
1715            }
1716
1717            async fn route(
1718                self: Arc<Self>,
1719                _driver: Driver,
1720                _request: Request,
1721            ) -> Result<RoutingOutcome> {
1722                let _guard = DropGuard(self.dropped.clone());
1723                let _ = self.started.send(());
1724                // Hang forever without ever touching the driver, so only cancellation
1725                // (not a dropped step channel) can stop this task.
1726                std::future::pending::<()>().await;
1727                unreachable!()
1728            }
1729        }
1730
1731        let (started_tx, mut started_rx) = mpsc::unbounded_channel();
1732        let dropped = Arc::new(AtomicBool::new(false));
1733        let algo: Arc<dyn Algorithm> = Arc::new(StuckAlgo {
1734            started: started_tx,
1735            dropped: dropped.clone(),
1736        });
1737
1738        // Drive the run on its own task, wait until the algorithm task is up, then cancel
1739        // it — dropping its future (and the `run_stream` stream it holds).
1740        let run_task = tokio::spawn(async move { test_drive(algo, request(), echo()).await });
1741        started_rx
1742            .recv()
1743            .await
1744            .ok_or_else(|| test_error("task never started"))?;
1745        run_task.abort();
1746        tokio::time::sleep(Duration::from_millis(100)).await;
1747
1748        assert!(
1749            dropped.load(Ordering::SeqCst),
1750            "algorithm task was NOT cancelled after cancelling run"
1751        );
1752        Ok(())
1753    }
1754
1755    // --- first-wins hedging: `run` must not wait on losing speculative calls -------------
1756
1757    /// Offloads two targets concurrently and returns the first to resolve, dropping the
1758    /// loser's call (first-wins hedging).
1759    struct Hedge {
1760        winner: String,
1761        loser: String,
1762    }
1763
1764    #[async_trait]
1765    impl Algorithm for Hedge {
1766        fn name(&self) -> &str {
1767            "hedge"
1768        }
1769
1770        async fn route(
1771            self: Arc<Self>,
1772            driver: Driver,
1773            request: Request,
1774        ) -> Result<RoutingOutcome> {
1775            let outcome_request = request.clone();
1776            let win = driver.call_model(request.clone(), vec![self.winner.clone().into()]);
1777            let lose = driver.call_model(request, vec![self.loser.clone().into()]);
1778            // First to resolve wins; `select!` drops the losing future (and its promise).
1779            tokio::select! {
1780                res = win => Ok(RoutingOutcome::answered(
1781                    self.winner.clone().into(),
1782                    outcome_request,
1783                    res?,
1784                )),
1785                res = lose => Ok(RoutingOutcome::answered(
1786                    self.loser.clone().into(),
1787                    outcome_request,
1788                    res?,
1789                )),
1790            }
1791        }
1792    }
1793
1794    /// Builds a hedging algo and the `serve` that drives it: the winner is gated behind the
1795    /// loser starting (so the loser's serve is guaranteed in flight when the winner wins),
1796    /// and the loser finishes after `loser_delay` — or never, when `None`.
1797    fn hedge(loser_delay: Option<std::time::Duration>) -> (Arc<dyn Algorithm>, impl Serve) {
1798        let started = Arc::new(tokio::sync::Notify::new());
1799        let algo = Arc::new(Hedge {
1800            winner: "winner".to_string(),
1801            loser: "loser".to_string(),
1802        });
1803        let serve = move |target: ModelId, _request: Request| {
1804            let started = started.clone();
1805            async move {
1806                if target == "loser" {
1807                    started.notify_one();
1808                    match loser_delay {
1809                        Some(delay) => tokio::time::sleep(delay).await,
1810                        None => std::future::pending::<()>().await,
1811                    }
1812                } else {
1813                    started.notified().await;
1814                }
1815                Ok(reply(target))
1816            }
1817        };
1818        (algo, serve)
1819    }
1820
1821    #[tokio::test]
1822    async fn run_returns_the_winner_without_a_late_loser_overwriting_it() -> Result<()> {
1823        // The loser responds 50ms after the winner has already won. `run` must return the
1824        // winner, not the loser's `respond`-to-a-dropped-receiver error.
1825        let (algo, serve) = hedge(Some(std::time::Duration::from_millis(50)));
1826        let (_, response) = test_drive(algo, request(), serve).await?;
1827        assert_eq!(
1828            response
1829                .llm_response
1830                .as_agg()
1831                .map(completion_text)
1832                .unwrap_or_default(),
1833            "winner"
1834        );
1835        Ok(())
1836    }
1837
1838    #[tokio::test]
1839    async fn run_returns_the_winner_without_hanging_on_a_pending_loser() -> Result<()> {
1840        // The loser never resolves. `run` must return the winner promptly, not hang
1841        // waiting for the in-flight loser.
1842        let (algo, serve) = hedge(None);
1843        let run = test_drive(algo, request(), serve);
1844        let (_, response) = tokio::time::timeout(std::time::Duration::from_secs(1), run)
1845            .await
1846            .map_err(|error| LibsyError::external("waiting for pending loser", error))??;
1847        assert_eq!(
1848            response
1849                .llm_response
1850                .as_agg()
1851                .map(completion_text)
1852                .unwrap_or_default(),
1853            "winner"
1854        );
1855        Ok(())
1856    }
1857
1858    #[tokio::test]
1859    async fn run_surfaces_a_terminal_error_with_many_calls_in_flight() -> Result<()> {
1860        use std::sync::atomic::{AtomicUsize, Ordering};
1861
1862        // A large fan-out (10 matched the old, now-removed concurrency cap). The terminal
1863        // error must still reach the caller with all of these calls pending.
1864        const N: usize = 10;
1865
1866        // Fans out N calls, then errors as soon as all N are in flight — exercising a
1867        // terminal failure emitted while the offloaded calls are still pending.
1868        struct FanOutThenError {
1869            all_started: Arc<tokio::sync::Notify>,
1870            n: usize,
1871        }
1872
1873        #[async_trait]
1874        impl Algorithm for FanOutThenError {
1875            fn name(&self) -> &str {
1876                "fan_out_then_error"
1877            }
1878
1879            async fn route(
1880                self: Arc<Self>,
1881                driver: Driver,
1882                request: Request,
1883            ) -> Result<RoutingOutcome> {
1884                let offloads = futures::future::join_all(
1885                    (0..self.n)
1886                        .map(|i| driver.call_model(request.clone(), vec![format!("m{i}").into()])),
1887                );
1888                tokio::select! {
1889                    _ = offloads => Err(test_error("offloads unexpectedly completed")),
1890                    _ = self.all_started.notified() => {
1891                        Err(test_error("terminal error while calls pending"))
1892                    }
1893                }
1894            }
1895        }
1896
1897        let all_started = Arc::new(tokio::sync::Notify::new());
1898        let algo: Arc<dyn Algorithm> = Arc::new(FanOutThenError {
1899            all_started: all_started.clone(),
1900            n: N,
1901        });
1902
1903        // Serving enters each call; once all N are in flight it signals, then pends forever.
1904        let started = Arc::new(AtomicUsize::new(0));
1905        let serve = move |_target: ModelId, _request: Request| {
1906            let started = started.clone();
1907            let all_started = all_started.clone();
1908            async move {
1909                if started.fetch_add(1, Ordering::SeqCst) + 1 == N {
1910                    all_started.notify_one();
1911                }
1912                std::future::pending::<ServeResult>().await
1913            }
1914        };
1915
1916        // With the cap gone, the driver keeps polling the stream even with N calls in
1917        // flight, so the terminal error surfaces promptly instead of hanging.
1918        let run = test_drive(algo, request(), serve);
1919        let result = tokio::time::timeout(std::time::Duration::from_millis(500), run)
1920            .await
1921            .map_err(|error| {
1922                LibsyError::external("waiting for terminal error with full call cap", error)
1923            })?;
1924        match result {
1925            Ok(_) => Err(test_error("expected the terminal error, got a response")),
1926            Err(err) => {
1927                assert!(
1928                    err.to_string()
1929                        .contains("terminal error while calls pending")
1930                );
1931                Ok(())
1932            }
1933        }
1934    }
1935}