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