switchyard_libsy/core/

algorithm.rs

// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
// SPDX-License-Identifier: Apache-2.0

//! Routing algorithms request external work through [`Driver`] and receive typed replies.
//! Hosts execute that work, keeping routing policy independent of transport.

use std::{
    collections::HashMap, future::Future, panic::AssertUnwindSafe, pin::Pin, sync::Arc,
    time::Instant,
};

use async_trait::async_trait;
use futures::{FutureExt, Stream, StreamExt};
use parking_lot::Mutex;
use serde_json::Value;
use tokio::sync::{mpsc, oneshot};
use tokio_stream::wrappers::ReceiverStream;
use tracing::Instrument;

/// The request/response protocol types come from [`switchyard_protocol`].
/// [`switchyard_protocol::LlmRequest`] is the normalized request;
/// [`switchyard_protocol::AggLlmResponse`] is the buffered response;
/// [`switchyard_protocol::LlmResponseChunk`] is normalized streaming content;
/// [`switchyard_protocol::LlmResponseStreamEvent`] is its host/algorithm envelope; and
/// [`switchyard_protocol::LlmResponse`] carries either a live
/// [`switchyard_protocol::LlmResponseStream`] or the terminal aggregate.
use switchyard_protocol::{
    AggLlmResponse, Category, ContentBlock, DecisionRequest, DecisionResponse, ModelId, Request,
    Response, StopReason,
};

use crate::{DriverError, LibsyError, Result, State, observability};

/// A boxed, `Send` stream of [`Step`]s — the output of
/// [`Algorithm::run_stream`]. Boxed so the trait method that produces it keeps
/// `Arc<dyn Algorithm>` object-safe.
pub type StepStream = Pin<Box<dyn Stream<Item = Result<Step>> + Send>>;

/// The models one algorithm run may use, grouped by [`Category`]. Within a
/// category they are ordered best-first.
///
/// Delegated sub-agent work gets its own groups, reachable only through
/// [`Driver::for_subagent`]. Keeping them separate is what stops a sub-agent's
/// `capable` from resolving to the parent's, and stops the parent falling back
/// onto a model only its sub-agents were given.
///
/// One run's driver clones all read the same value, so it is passed as
/// `Arc<RuntimeModels>` rather than cloned per driver.
#[derive(Clone, Debug, Default)]
pub struct RuntimeModels {
    /// When using subagents this is the parent agent category.
    by_category: HashMap<Category, Vec<ModelId>>,
    subagent: Option<HashMap<Category, Vec<ModelId>>>,
    has_target_restriction: bool,
    provider_models: HashMap<ModelId, ModelId>,
}

impl RuntimeModels {
    /// The models available to the algorithm itself.
    pub fn new(by_category: HashMap<Category, Vec<ModelId>>) -> Self {
        Self {
            by_category,
            subagent: None,
            has_target_restriction: false,
            provider_models: HashMap::new(),
        }
    }

    /// Supplies provider identities for telemetry while routing keeps using target keys.
    pub fn with_provider_models(mut self, models: HashMap<ModelId, ModelId>) -> Self {
        self.provider_models = models;
        self
    }

    /// Returns the provider model for a routing key, or the key for direct library callers.
    pub fn provider_model<'a>(&'a self, target: &'a ModelId) -> &'a ModelId {
        self.provider_models.get(target).unwrap_or(target)
    }

    /// Adds the groups used for delegated sub-agent work.
    pub fn with_subagent(mut self, models: HashMap<Category, Vec<ModelId>>) -> Self {
        self.subagent = Some(models);
        self
    }

    /// Rejects routing calls and completion candidates outside the configured model groups.
    /// Enforced by the driver and [`drive`]; hosts bind IDs to permitted deployments.
    pub fn with_target_restriction(mut self) -> Self {
        self.has_target_restriction = true;
        self
    }

    fn check_targets(&self, targets: &[ModelId]) -> Result<()> {
        if !self.has_target_restriction {
            return Ok(());
        }
        // Include delegated groups: the parent driver publishes subagent outcomes too.
        for target in targets {
            if !self
                .by_category
                .values()
                .chain(self.subagent.iter().flat_map(HashMap::values))
                .any(|models| models.contains(target))
            {
                return Err(LibsyError::TargetNotFound {
                    target: target.clone(),
                });
            }
        }
        Ok(())
    }

    fn check_request_targets(&self, request: &Request, targets: &[ModelId]) -> Result<()> {
        self.check_targets(targets)?;
        if self.has_target_restriction {
            let selected = targets.first().ok_or(LibsyError::NoTargets)?;
            if request.llm_request.model.as_deref() != Some(selected.as_str()) {
                return Err(LibsyError::AlgorithmError {
                    message: "request model does not match selected target".into(),
                });
            }
        }
        Ok(())
    }

    /// The models in `category`, ordered best-first.
    pub fn models_for(&self, category: &Category) -> &[ModelId] {
        self.by_category.get(category).map_or(&[], Vec::as_slice)
    }

    /// The models delegated sub-agent work uses for `category`, ordered best-first.
    pub fn subagent_models_for(&self, category: &Category) -> &[ModelId] {
        self.subagent
            .as_ref()
            .and_then(|models| models.get(category))
            .map_or(&[], Vec::as_slice)
    }
}

impl From<HashMap<Category, Vec<ModelId>>> for RuntimeModels {
    fn from(by_category: HashMap<Category, Vec<ModelId>>) -> Self {
        Self::new(by_category)
    }
}

/// Which of a [`RuntimeModels`]' groups a driver reads.
#[derive(Clone, Copy)]
enum Scope {
    Parent,
    Subagent,
}

/// An offloaded model call, surfaced inside [`Step::CallModel`].
///
/// The host reads the public fields and passes its model-call future to
/// [`respond`](Self::respond), unblocking the algorithm's [`Driver::call_model`] on the
/// other side. `switchyard-llm-client`'s `run` is the ready-made consumer that does this for
/// you.
///
/// [`Driver::call_model`] stamps the first candidate model onto the request before publishing
/// the call. A consumer that falls through to a later candidate must re-stamp it.
pub struct CallModel {
    /// The name of the algorithm that produced this call, so a host instrumenting the
    /// calls it serves can attribute its own spans to the algorithm behind them.
    pub algorithm: String,
    /// The request to serve; its `model` is stamped with the first candidate.
    pub request: Request,
    /// Candidate models, tried in order until one answers. Never empty.
    pub models: Vec<ModelId>,
    /// Return client errors to the algorithm so it can apply its fallback policy.
    pub recover_errors: bool,
    telemetry_model: ModelId,
    /// How to send the response back to the algorithm. `None` once the call is recorded.
    reply: Option<oneshot::Sender<Result<Response>>>,
    started: Instant,
    // Retain the originating span even if the waiting algorithm is cancelled first.
    span: tracing::Span,
}

impl CallModel {
    /// Run host work until it completes or the algorithm stops waiting for this call.
    /// Dropping the waiting [`Driver::call_model`] future drops `work` and returns `Ok(())`.
    /// Errors stop the host's run unless [`Self::recover_errors`] is enabled.
    /// Use `std::future::ready(result)` for an already available response.
    /// A delivered response stream owns its subsequent lifetime.
    pub async fn respond(mut self, work: impl Future<Output = Result<Response>>) -> Result<()> {
        let reply = self.reply.as_mut().ok_or(DriverError::ResponseDropped)?;
        let result = tokio::select! {
            biased;
            _ = reply.closed() => return Ok(()),
            result = work => result,
        };
        if self.reply.as_ref().is_some_and(oneshot::Sender::is_closed) {
            return Ok(());
        }
        if !self.recover_errors
            && let Err(error) = result
        {
            return self.fail(error);
        }
        self.record(result.is_ok());
        if let Ok(response) = &result {
            observability::record_llm_response(response, &self.span);
        }
        self.reply
            .take()
            .ok_or(DriverError::ResponseDropped)?
            .send(result)
            .ok();
        Ok(())
    }

    /// Record a failed call and return its error to stop [`drive`].
    /// Leaves the promise unfulfilled so the driver can cancel the algorithm.
    pub fn fail(mut self, error: LibsyError) -> Result<()> {
        self.reply = None;
        self.record(false);
        Err(error)
    }

    fn record(&self, is_ok: bool) {
        observability::record_llm_call(
            &self.algorithm,
            self.telemetry_model.as_str(),
            self.started.elapsed(),
            is_ok,
        );
        self.span
            .record("outcome", if is_ok { "ok" } else { "error" });
    }
}

impl Drop for CallModel {
    fn drop(&mut self) {
        if self.reply.is_some() {
            self.record(false);
        }
    }
}

/// A host-owned decision call. Dropping it without replying yields [`DriverError::ResponseDropped`].
/// Completion or drop records one call and its duration, including time waiting for the host.
pub struct CallDecision {
    /// Algorithm name for attributing host telemetry.
    pub algorithm: String,
    /// Its model is set to [`Self::model`] before dispatch.
    pub request: DecisionRequest,
    /// Target ID for client lookup.
    pub model: ModelId,
    reply: Option<oneshot::Sender<Result<DecisionResponse>>>,
    started: Instant,
    // Retain the originating span even if the waiting algorithm is cancelled first.
    span: tracing::Span,
}

impl CallDecision {
    /// Run host work and return its response or error to the algorithm.
    /// Dropping the waiting [`Driver::call_decision`] future drops `work` and returns `Ok(())`.
    /// Use `std::future::ready(result)` for an already available response.
    pub async fn respond(
        mut self,
        work: impl Future<Output = Result<DecisionResponse>>,
    ) -> Result<()> {
        let reply = self.reply.as_mut().ok_or(DriverError::ResponseDropped)?;
        let result = tokio::select! {
            biased;
            _ = reply.closed() => return Ok(()),
            result = work => result,
        };
        self.record(result.is_ok());
        if let Ok(response) = &result {
            observability::record_decision_response(response, &self.span);
        }
        self.reply
            .take()
            .ok_or(DriverError::ResponseDropped)?
            .send(result)
            .ok();
        Ok(())
    }

    /// Returning this error from the host handler aborts [`drive`].
    pub fn fail(mut self, error: LibsyError) -> Result<()> {
        self.reply = None;
        self.record(false);
        Err(error)
    }

    fn record(&self, is_ok: bool) {
        observability::record_decision_call(
            &self.algorithm,
            &self.model,
            self.started.elapsed(),
            is_ok,
        );
        self.span
            .record("outcome", if is_ok { "ok" } else { "error" });
    }
}

impl Drop for CallDecision {
    fn drop(&mut self) {
        if self.reply.is_some() {
            self.record(false);
        }
    }
}

/// The terminal result of routing.
pub struct RoutingOutcome {
    /// Models selected by the algorithm, ordered best model first.
    pub selected_model_ids: Vec<ModelId>,
    /// The request after all routing-time rewrites, stamped with the selected model.
    pub request: Request,
    /// The request the fallback candidates receive instead of `request`, when routing
    /// rewrote `request` for the selected model only.
    /// Content-policy denials do not advance to these fallback candidates.
    pub fallback_request: Option<Request>,
    /// Candidate-specific requests, including model-specific instructions.
    pub candidate_requests: HashMap<ModelId, Request>,
    /// A response produced while routing, or `None` when the client must make the answer call.
    pub response: Option<Response>,
    /// Outcome identity and optional algorithm evidence.
    ///
    /// Constructors leave this empty; [`Algorithm::run_stream`] fills it before publishing a
    /// successful outcome.
    pub metadata: Option<crate::OutcomeMetadata>,
    /// Routing state after the run, for hosts that store state outside the algorithm.
    pub state: State,
}

impl RoutingOutcome {
    /// The model the algorithm recommends, the best model for this request.
    /// `LibsyError::NoTargets` if the algorithm selected no models, which should be impossible.
    pub fn selected_model_id(&self) -> Result<&ModelId> {
        self.selected_model_ids.first().ok_or(LibsyError::NoTargets)
    }

    /// The decision is that client should send this `request`. The `selected_model_id`
    /// will be written into it by this function.
    /// If that fails client should try the `fallback_models` in order.
    pub fn route_to(
        selected_model_id: ModelId,
        fallback_models: Vec<ModelId>,
        mut request: Request,
    ) -> Self {
        request.llm_request.model = Some(selected_model_id.to_string());
        let mut selected_model_ids = Vec::with_capacity(1 + fallback_models.len());
        selected_model_ids.push(selected_model_id);
        selected_model_ids.extend(fallback_models);
        Self {
            selected_model_ids,
            request,
            fallback_request: None,
            candidate_requests: HashMap::new(),
            response: None,
            metadata: None,
            state: State::default(),
        }
    }

    /// Algorithm generated the response as part of the routing decision. Here it is.
    /// The `request` will have the `selected_model_id` written into it by this function.
    pub fn answered(selected_model_id: ModelId, mut request: Request, response: Response) -> Self {
        request.llm_request.model = Some(selected_model_id.to_string());
        Self {
            selected_model_ids: vec![selected_model_id],
            request,
            fallback_request: None,
            candidate_requests: HashMap::new(),
            response: Some(response),
            metadata: None,
            state: State::default(),
        }
    }
}

/// What the host saw of the answer it served for a routed request.
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub struct AnswerSummary {
    /// The answer ended its turn with visible text and requested no tool call.
    pub final_reply: bool,
}

impl AnswerSummary {
    /// Summarizes a completed answer.
    pub fn of(response: &AggLlmResponse) -> Self {
        let final_reply = response.outputs.iter().all(|output| {
            matches!(output.stop_reason, None | Some(StopReason::EndTurn))
                && !output
                    .content
                    .iter()
                    .any(|block| matches!(block, ContentBlock::ToolCall(_)))
        }) && response.outputs.iter().any(|output| {
            output.role == switchyard_protocol::Role::Assistant
                && output.content.iter().any(
                    |block| matches!(block, ContentBlock::Text { text } if !text.trim().is_empty()),
                )
        });
        Self { final_reply }
    }
}

#[derive(Default)]
struct RunEvidence {
    outcome: Option<Value>,
    judgments: Vec<crate::JudgeDecision>,
}

/// An algorithm's handle for requesting external work and awaiting the host's typed replies.
#[derive(Clone)]
pub struct Driver {
    step_tx: mpsc::Sender<Result<Step>>,

    /// The owning algorithm's telemetry label, stamped onto every call this driver publishes.
    algorithm: String,

    /// Run-scoped evidence shared by driver clones and attached only to a successful outcome.
    evidence: Arc<Mutex<RunEvidence>>,

    /// Caller-supplied routing state, read and updated when a request carries no session ID.
    state: Arc<Mutex<State>>,

    /// Every group this run may route over, shared by all driver clones.
    models: Arc<RuntimeModels>,

    /// Which of those groups this driver reads.
    scope: Scope,
    judge: Option<ModelId>,
}

impl Driver {
    /// Build an empty driver with its step channel ready. Created per call by
    /// [`run_stream`](Algorithm::run_stream). Also returns the Step receiver.
    pub(crate) fn new(
        algorithm: &str,
        models: Arc<RuntimeModels>,
    ) -> (Self, mpsc::Receiver<Result<Step>>) {
        // Capacity one keeps the algorithm paced by the stream consumer. It limits queued steps,
        // not model calls already pulled from the stream, which can still run at the same time.
        // A larger buffer would use more memory and let the algorithm run farther ahead with
        // little benefit because reading a step is cheap compared with serving a model call.
        let (step_tx, step_rx) = mpsc::channel(1);
        (
            Self {
                step_tx,
                algorithm: algorithm.to_string(),
                evidence: Arc::new(Mutex::new(RunEvidence::default())),
                state: Arc::new(Mutex::new(State::default())),
                models,
                scope: Scope::Parent,
                judge: None,
            },
            step_rx,
        )
    }

    pub(crate) fn take_state(&self) -> State {
        std::mem::take(&mut *self.state.lock())
    }

    pub(crate) fn set_state(&self, state: State) {
        *self.state.lock() = state;
    }

    /// Replace the current run's evidence when a component makes the final decision.
    pub(crate) fn set_evidence(&self, evidence: Value) {
        self.evidence.lock().outcome = Some(evidence);
    }

    /// Supply fallback evidence without replacing a decision made earlier in the cascade.
    pub(crate) fn set_evidence_if_empty(&self, evidence: Value) {
        let mut current = self.evidence.lock();
        if current.outcome.is_none() {
            current.outcome = Some(evidence);
        }
    }

    pub(crate) fn record_judgment(&self, selected: Option<crate::Score>, evidence: Option<Value>) {
        let (target, category) = selected
            .map(|score| (Some(score.target), score.category))
            .unwrap_or_default();
        self.evidence.lock().judgments.push(crate::JudgeDecision {
            target,
            category: category.map(|category| category.as_str().to_string()),
            evidence,
        });
    }

    /// Publish a model call and await the consumer's response.
    ///
    /// Errors if the stream is closed or the call failed.
    /// The await is wrapped in a `libsy.llm_call` span measuring *fulfillment* as
    /// the algorithm observes it (host queueing/serving included; a streamed
    /// response resolves when its stream handle arrives). The host records call metrics,
    /// outcome, and buffered token usage when it completes or drops the call, even if the
    /// waiting algorithm has been cancelled. The provider call itself is the
    /// host's, and is instrumented by whoever makes it.
    pub async fn call_model(&self, request: Request, models: Vec<ModelId>) -> Result<Response> {
        self.call_model_with_error_recovery(request, models, false)
            .await
    }

    /// Allows a routing policy to handle a failed call when recovery is enabled.
    #[tracing::instrument(
        target = "libsy",
        name = "libsy.llm_call",
        skip_all,
        fields(
            algorithm = self.algorithm,
            selected_model = %models.first().map(|model| self.provider_model(model).as_str()).unwrap_or("NoTargets"),
            openinference.span.kind = "CHAIN",
            outcome = tracing::field::Empty,
            input_tokens = tracing::field::Empty,
            output_tokens = tracing::field::Empty,
            total_tokens = tracing::field::Empty,
            reasoning_tokens = tracing::field::Empty,
        )
    )]
    pub(crate) async fn call_model_with_error_recovery(
        &self,
        mut request: Request,
        models: Vec<ModelId>,
        recover_errors: bool,
    ) -> Result<Response> {
        let Some(selected_model_id) = models.first() else {
            return Err(LibsyError::NoTargets);
        };
        self.models.check_targets(&models)?;
        request.llm_request.model = Some(selected_model_id.to_string());
        let started = Instant::now();
        let (reply, response) = oneshot::channel::<Result<Response>>();
        let call = CallModel {
            algorithm: self.algorithm.clone(),
            request,
            telemetry_model: self.provider_model(selected_model_id).clone(),
            models,
            recover_errors,
            reply: Some(reply),
            started,
            span: tracing::Span::current(),
        };
        self.step_tx
            .send(Ok(Step::CallModel(Box::new(call))))
            .await
            .map_err(|_| DriverError::StreamClosed)?;
        response
            .await
            .map_err(|_| LibsyError::from(DriverError::ResponseDropped))?
    }

    /// Override the request's model with `model` and wait for the host to reply.
    #[tracing::instrument(
        target = "libsy",
        name = "libsy.decision_call",
        skip_all,
        fields(
            algorithm = self.algorithm,
            selected_model = %model,
            openinference.span.kind = "CHAIN",
            outcome = tracing::field::Empty,
            input_tokens = tracing::field::Empty,
            output_tokens = tracing::field::Empty,
            total_tokens = tracing::field::Empty,
            reasoning_tokens = tracing::field::Empty,
            gen_ai.response.id = tracing::field::Empty,
            gen_ai.response.model = tracing::field::Empty,
        ),
    )]
    pub async fn call_decision(
        &self,
        mut request: DecisionRequest,
        model: ModelId,
    ) -> Result<DecisionResponse> {
        self.models.check_targets(std::slice::from_ref(&model))?;
        request.model = Some(model.clone());
        let started = Instant::now();
        let (reply, response) = oneshot::channel();
        let call = CallDecision {
            algorithm: self.algorithm.clone(),
            request,
            model,
            reply: Some(reply),
            started,
            span: tracing::Span::current(),
        };
        self.step_tx
            .send(Ok(Step::CallDecision(Box::new(call))))
            .await
            .map_err(|_| DriverError::StreamClosed)?;
        response
            .await
            .map_err(|_| LibsyError::from(DriverError::ResponseDropped))?
    }

    /// Provider identity for telemetry; algorithms still route on the original target key.
    pub fn provider_model<'a>(&'a self, target: &'a ModelId) -> &'a ModelId {
        self.models.provider_model(target)
    }

    /// The available models for this category, typically ordered best-first.
    pub fn models_for(&self, category: &Category) -> &[ModelId] {
        if *category == Category::Judge
            && let Some(judge) = &self.judge
        {
            return std::slice::from_ref(judge);
        }
        match self.scope {
            Scope::Parent => self.models.models_for(category),
            Scope::Subagent => self.models.subagent_models_for(category),
        }
    }

    pub(crate) fn for_judge(&self, judge: ModelId) -> Self {
        Self {
            judge: Some(judge),
            ..self.clone()
        }
    }

    /// The first available model for `category`.
    pub fn first_model_for(&self, category: &Category) -> Result<&ModelId> {
        self.models_for(category)
            .first()
            .ok_or_else(|| LibsyError::AlgorithmError {
                message: format!("no models available for category {}", category.as_str()),
            })
    }

    /// A driver scoped to delegated sub-agent work: its categories are the
    /// sub-agent's own, and the parent's are no longer reachable through it.
    pub fn for_subagent(&self) -> Result<Self> {
        if self.models.subagent.is_none() {
            return Err(LibsyError::AlgorithmError {
                message: "delegated work has no sub-agent models".to_string(),
            });
        }
        Ok(Self {
            scope: Scope::Subagent,
            ..self.clone()
        })
    }

    /// Emit the terminal step: [`Step::Done`] on `Ok`, or an `Err` stream
    /// item on failure. Internal: called once by [`run_stream`](Algorithm::run_stream)
    /// when the algorithm finishes.
    pub(crate) async fn finish(&self, result: Result<RoutingOutcome>) -> Result<()> {
        let result = result.map(|mut outcome| {
            outcome.state = self.take_state();
            let evidence = std::mem::take(&mut *self.evidence.lock());
            let metadata = outcome.metadata.get_or_insert_with(|| {
                crate::OutcomeMetadata::new(self.algorithm.clone(), evidence.outcome)
            });
            metadata.judgments.extend(evidence.judgments);
            let provider_models = outcome
                .selected_model_ids
                .iter()
                .map(|target| self.provider_model(target).clone())
                .collect::<Vec<_>>();
            observability::record_outcome(metadata, &provider_models);
            outcome
        });
        let selected_model = result
            .as_ref()
            .ok()
            .and_then(|outcome| outcome.selected_model_id().ok().cloned());
        let step = result.map(|outcome| Step::Done(Box::new(outcome)));
        self.step_tx
            .send(step)
            .await
            .map_err(|_| DriverError::StreamClosed)?;
        if let Some(selected_model) = selected_model {
            observability::record_decision(&self.algorithm, self.provider_model(&selected_model));
        }
        Ok(())
    }
}

/// Work requests and the terminal outcome of [`Algorithm::run_stream`].
/// Each call carries its own reply channel, so the host can execute calls concurrently.
pub enum Step {
    /// The algorithm needs this model call performed. The host serves it and fulfills
    /// it with [`CallModel::respond`]. Boxed: it is by far the largest variant.
    CallModel(Box<CallModel>),
    /// The host fulfills this decision call with [`CallDecision::respond`].
    CallDecision(Box<CallDecision>),
    /// The algorithm finished with its routing outcome — the last step of a run.
    Done(Box<RoutingOutcome>),
}

/// Work dispatched through one host handler, with a typed reply channel per variant.
/// [`drive`] handles terminal outcomes separately; the handler only executes calls.
pub enum Call {
    /// Fulfilled with an LLM [`Response`].
    Model(Box<CallModel>),
    /// Fulfilled with a [`DecisionResponse`].
    Decision(Box<CallDecision>),
}

/// Drive [`Algorithm::run_stream`] to its routing outcome, serving calls concurrently.
///
/// The shared execution loop lets hosts supply clients or test doubles through `serve`
/// without reimplementing stream handling and cancellation.
///
/// `serve` owns each [`Call`] and performs its I/O. Passing an error to `respond`
/// lets the algorithm fall back; returning an error from `serve` aborts the run.
///
/// Dropping this future cancels the algorithm and pending handler futures.
pub async fn drive<F, Fut>(
    algorithm: Arc<dyn Algorithm>,
    request: Request,
    models: Arc<RuntimeModels>,
    serve: F,
) -> Result<RoutingOutcome>
where
    F: Fn(Call) -> Fut,
    Fut: Future<Output = Result<()>>,
{
    drive_with_state(algorithm, request, models, State::default(), serve).await
}

/// [`drive`] starting from caller-supplied `state`.
pub(crate) async fn drive_with_state<F, Fut>(
    algorithm: Arc<dyn Algorithm>,
    request: Request,
    models: Arc<RuntimeModels>,
    state: State,
    serve: F,
) -> Result<RoutingOutcome>
where
    F: Fn(Call) -> Fut,
    Fut: Future<Output = Result<()>>,
{
    let stream = algorithm.run_stream(request, Arc::clone(&models), state);
    tokio::pin!(stream);

    let mut in_flight = futures::stream::FuturesUnordered::new();
    let mut final_outcome: Option<RoutingOutcome> = None;

    loop {
        tokio::select! {
            Some(result) = in_flight.next() => match result {
                Ok(()) => {}, // Call completed successfully
                Err(err) => return Err(err), // Call failed, propagate the error
            },
            step = stream.next() => {
                match step {
                    None => break, // stream has ended, no more steps
                    Some(item) => match item? {
                        Step::CallModel(call) => {
                            models.check_request_targets(&call.request, &call.models)?;
                            in_flight.push(serve(Call::Model(call)));
                        }
                        Step::CallDecision(call) => {
                            models.check_targets(std::slice::from_ref(&call.model))?;
                            if models.has_target_restriction
                                && call.request.model.as_ref() != Some(&call.model)
                            {
                                return Err(LibsyError::AlgorithmError {
                                    message: "decision payload model does not match selected target".into(),
                                });
                            }
                            in_flight.push(serve(Call::Decision(call)));
                        }
                        Step::Done(outcome) => {
                            if outcome.response.is_none() {
                                models.check_request_targets(&outcome.request, &outcome.selected_model_ids)?;
                            }
                            final_outcome = Some(*outcome);
                            break;
                        }
                    }
                }
            },
        }
    }
    final_outcome.ok_or(LibsyError::MissingFinalResponse)
}

/// Recover the message from an algorithm's panic.
fn panic_message(payload: &(dyn std::any::Any + Send)) -> String {
    payload
        .downcast_ref::<&'static str>()
        .map(|message| (*message).to_string())
        .or_else(|| payload.downcast_ref::<String>().cloned())
        .unwrap_or_else(|| "unknown panic payload".to_string())
}

/// Abort guard
struct AbortOnDrop(tokio::task::AbortHandle);

impl Drop for AbortOnDrop {
    fn drop(&mut self) {
        self.0.abort();
    }
}

/// Errors unless `targets` contains `name`.
///
/// Config target names must be resolved before an algorithm is built. This list contains
/// model IDs, not target names.
pub(crate) fn ensure_model_is_target(targets: &[ModelId], name: &ModelId) -> Result<()> {
    targets
        .iter()
        .any(|target| target == name)
        .then_some(())
        .ok_or_else(|| LibsyError::TargetNotFound {
            target: name.clone(),
        })
}

/// Key for routing affinity: a root request by its session, a child request by its session
/// and agent.
#[derive(Clone, Hash, PartialEq, Eq)]
pub(crate) enum RoutingIdentity {
    /// Root request, keyed by session ID.
    Session(String),
    /// Child request, keyed by session and agent IDs.
    Subagent { session: String, agent: String },
}

impl RoutingIdentity {
    /// Builds a root or child identity from non-empty request metadata.
    ///
    /// A child request missing either ID returns `None`, so it keeps no routing history
    /// rather than sharing the parent's.
    pub(crate) fn from_request(request: &Request) -> Option<Self> {
        let metadata = request.metadata.as_ref()?;
        let session = metadata.session_id.as_deref().filter(|id| !id.is_empty())?;
        if metadata.is_subagent {
            let agent = metadata.agent_id.as_deref().filter(|id| !id.is_empty())?;
            Some(Self::Subagent {
                session: session.to_string(),
                agent: agent.to_string(),
            })
        } else {
            Some(Self::Session(session.to_string()))
        }
    }
}

/// Routing policy independent of how external calls are executed.
/// Implement [`route`](Self::route), using [`Driver`] to request work from the host.
/// Hosts can use [`drive`] for the shared execution loop or consume
/// [`run_stream`](Self::run_stream) for custom execution.
///
/// Methods take `self: Arc<Self>`: one algorithm (`Arc<dyn Algorithm>`) is shared across
/// requests and run concurrently, so it owns its thread-safety and any shared state.
///
/// # Concurrency
///
/// A host may run the same algorithm concurrently for many requests. Implementations
/// must synchronize their own mutable shared state. Each call to [`run_stream`](Self::run_stream)
/// creates an independent [`Driver`], so model-call promises and emitted [`Step`]s cannot
/// cross between runs.
///
/// # Observability
///
/// [`run_stream`](Self::run_stream) creates a `libsy.run` span, and each offloaded model
/// call creates a nested span for its call kind, with outcome and available token usage.
/// Call metrics include host queueing and count unfulfilled drops as errors.
/// Successful outcomes record their
/// [`OutcomeMetadata::outcome_id`](crate::OutcomeMetadata::outcome_id) on `libsy.run`,
/// alongside `selected_model_ids` (an ordered OpenTelemetry string array).
/// `algorithm` and `switchyard.algorithm` retain the run's [`Algorithm::name`].
/// Optional `evidence.source`, `evidence.verdict`, `evidence.trigger`, and
/// `evidence.reason_code` are strings; `evidence.score`, `evidence.confidence`,
/// `evidence.threshold`, `evidence.capable_score`, and `evidence.capable_threshold`
/// are numbers. Unknown evidence fields are not exported.
/// These fields are span attributes, never metric labels.
///
/// The run/call observability helpers retain `outcome` status and operational metrics,
/// but omit error details and arbitrary request extra metadata. Algorithms and hosts
/// may emit their own logs. Errors still reach the caller unchanged.
/// The host controls the tracing subscriber and global OpenTelemetry
/// meter provider; libsy installs no exporter and performs no telemetry network I/O.
#[async_trait]
pub trait Algorithm: Send + Sync + 'static {
    /// Stable, low-cardinality name identifying this algorithm — the
    /// `algorithm` attribute on every span, metric, and log line the crate
    /// emits for its runs.
    fn name(&self) -> &str;

    /// Select a route, requesting external work through [`Driver`] as needed.
    /// [`run_stream`](Self::run_stream) runs this method and exposes its work to the host.
    async fn route(self: Arc<Self>, driver: Driver, request: Request) -> Result<RoutingOutcome>;

    /// Whether the host should summarize served answers for [`Self::observe_answer`].
    fn tracks_answers(&self) -> bool {
        false
    }

    /// Records the answer the host served for a request in `session_id`.
    ///
    /// The host reports a streamed answer before forwarding its final chunk, so the
    /// session's next request is routed with this answer already recorded.
    async fn observe_answer(&self, _session_id: &str, _answer: AnswerSummary) {}

    /// Expose work requests and the final outcome so the host can control execution.
    ///
    /// Each call waits for its host reply. A run ends with one terminal item — [`Step::Done`] on
    /// success, an `Err` item on failure, including when the algorithm panics. Dropping
    /// the stream aborts the spawned algorithm task.
    ///
    /// Every invocation owns a separate [`Driver`].
    fn run_stream(
        self: Arc<Self>,
        request: Request,
        models: Arc<RuntimeModels>,
        state: State,
    ) -> StepStream {
        let (driver, step_rx) = Driver::new(self.name(), models);
        driver.set_state(state);
        let span = observability::run_span(self.name(), &request);
        let handle = tokio::spawn(
            async move {
                let algorithm = self.name().to_string();
                // Catch a panicking algorithm so the run still publishes a terminal step.
                let route = AssertUnwindSafe(self.route(driver.clone(), request)).catch_unwind();
                let result = observability::observe_run(&algorithm, async {
                    let outcome = route.await.unwrap_or_else(|payload| {
                        Err(LibsyError::AlgorithmError {
                            message: format!(
                                "algorithm task panicked: {}",
                                panic_message(payload.as_ref())
                            ),
                        })
                    })?;
                    if outcome.response.is_none() {
                        driver
                            .models
                            .check_request_targets(&outcome.request, &outcome.selected_model_ids)?;
                    }
                    Ok(outcome)
                })
                .await;

                let _ = driver.finish(result).await;
            }
            .instrument(span),
        );
        // Dropping the stream aborts the algorithm task when its consumer goes away.
        let abort_guard = AbortOnDrop(handle.abort_handle());
        Box::pin(ReceiverStream::new(step_rx).map(move |step| {
            // link abort guard to stream
            let _keep_alive = &abort_guard;
            step
        }))
    }
}

#[cfg(test)]
mod tests {
    use std::collections::HashMap;

    use super::*;
    use crate::core::testing::{Serve, ServeResult, echo, reply, serve_decision, test_drive};
    use futures::StreamExt;
    use switchyard_protocol::{
        LlmResponse, LlmResponseChunk, completion_text, text_request, text_response,
    };

    #[test]
    fn a_final_answer_has_text_and_no_unfinished_or_tool_output() {
        let mut answer = text_response(None, "done");
        assert!(AnswerSummary::of(&answer).final_reply);
        answer.outputs[0].stop_reason = Some(StopReason::MaxTokens);
        assert!(!AnswerSummary::of(&answer).final_reply);
        answer.outputs[0].stop_reason = Some(StopReason::EndTurn);
        let mut tool = answer.outputs[0].clone();
        tool.content = vec![ContentBlock::ToolCall(switchyard_protocol::ToolCall {
            id: "call".into(),
            name: "bash".into(),
            arguments: serde_json::json!({}),
        })];
        answer.outputs.push(tool);
        assert!(!AnswerSummary::of(&answer).final_reply);
        answer.outputs.pop();
        answer.outputs[0].content = vec![ContentBlock::Text { text: "  ".into() }];
        assert!(!AnswerSummary::of(&answer).final_reply);
        assert!(!AnswerSummary::of(&AggLlmResponse::default()).final_reply);
    }

    #[derive(Debug, thiserror::Error)]
    #[error("{0}")]
    struct TestError(&'static str);

    fn test_error(message: &'static str) -> LibsyError {
        LibsyError::external("test", TestError(message))
    }

    /// Trivial algo used only to exercise the orchestrator: calls the first target
    /// and returns its response as the routing outcome.
    struct TestAlgo {
        target_set: Vec<ModelId>,
    }

    #[async_trait]
    impl Algorithm for TestAlgo {
        fn name(&self) -> &str {
            "test"
        }

        async fn route(
            self: Arc<Self>,
            driver: Driver,
            request: Request,
        ) -> Result<RoutingOutcome> {
            let target = self
                .target_set
                .first()
                .ok_or(LibsyError::NoTargets)?
                .clone();
            let response = driver
                .call_model(request.clone(), vec![target.clone()])
                .await?;
            driver.record_judgment(
                Some(crate::Score {
                    target: target.clone(),
                    category: Some(Category::Capable),
                    confidence: 1.0,
                }),
                Some(serde_json::json!({"source": "judge"})),
            );
            driver.set_evidence(serde_json::json!({"source": "test"}));
            driver.set_evidence_if_empty(serde_json::json!({"source": "ignored"}));
            Ok(RoutingOutcome::answered(target, request, response))
        }
    }

    /// Build a shared `TestAlgo` over the given target set.
    fn orch(target_set: Vec<ModelId>) -> Arc<dyn Algorithm> {
        Arc::new(TestAlgo { target_set })
    }

    fn request() -> Request {
        Request {
            llm_request: text_request(Some("auto".to_string()), "hi".to_string()),
            raw_request: None,
            metadata: None,
        }
    }

    #[test]
    fn routing_outcome_constructors_stamp_selection_and_preserve_payloads() {
        let outcome = RoutingOutcome::route_to(
            "selected".into(),
            target_set(&["fallback-one", "fallback-two"]),
            request(),
        );

        assert_eq!(
            outcome.selected_model_ids,
            target_set(&["selected", "fallback-one", "fallback-two"])
        );
        assert_eq!(outcome.request.model_id().as_deref(), Some("selected"));
        assert!(outcome.response.is_none());
        assert!(outcome.metadata.is_none());

        let outcome = RoutingOutcome::route_to("only".into(), Vec::new(), request());
        assert_eq!(outcome.selected_model_ids, target_set(&["only"]));

        let outcome = RoutingOutcome::answered(
            "answered".into(),
            request(),
            Response {
                llm_response: LlmResponse::Agg(text_response(None, "existing")),
                metadata: None,
                upstream_headers: http::HeaderMap::new(),
            },
        );

        assert_eq!(outcome.selected_model_ids, target_set(&["answered"]));
        assert_eq!(outcome.request.model_id().as_deref(), Some("answered"));
        assert_eq!(
            outcome
                .response
                .as_ref()
                .and_then(|response| response.llm_response.as_agg())
                .map(completion_text),
            Some("existing".to_string())
        );
    }

    fn target_set(names: &[&str]) -> Vec<ModelId> {
        names.iter().map(|name| ModelId::from(*name)).collect()
    }

    enum TargetProbe {
        Model(&'static [&'static str]),
        Decision(&'static str),
        Completion(&'static [&'static str]),
        MismatchedCompletion,
    }

    #[async_trait]
    impl Algorithm for TargetProbe {
        fn name(&self) -> &str {
            "target_probe"
        }

        async fn route(
            self: Arc<Self>,
            driver: Driver,
            request: Request,
        ) -> Result<RoutingOutcome> {
            match &*self {
                Self::Model(models) => {
                    let response = driver
                        .call_model(request.clone(), target_set(models))
                        .await?;
                    Ok(RoutingOutcome::answered(
                        models[0].into(),
                        request,
                        response,
                    ))
                }
                Self::Decision(model) => {
                    driver
                        .call_decision(
                            DecisionRequest {
                                model: None,
                                input: Value::Null,
                                questions: Default::default(),
                            },
                            (*model).into(),
                        )
                        .await?;
                    Ok(RoutingOutcome::route_to("allowed".into(), vec![], request))
                }
                Self::Completion(models) => Ok(RoutingOutcome::route_to(
                    models[0].into(),
                    target_set(&models[1..]),
                    request,
                )),
                Self::MismatchedCompletion => {
                    let mut outcome = RoutingOutcome::route_to("allowed".into(), vec![], request);
                    outcome.request.llm_request.model = Some("outside".into());
                    Ok(outcome)
                }
            }
        }
    }

    struct RewrittenSteps(Arc<TargetProbe>, fn(&mut Step));

    #[async_trait]
    impl Algorithm for RewrittenSteps {
        fn name(&self) -> &str {
            "rewritten_steps"
        }

        async fn route(
            self: Arc<Self>,
            driver: Driver,
            request: Request,
        ) -> Result<RoutingOutcome> {
            Arc::clone(&self.0).route(driver, request).await
        }

        fn run_stream(
            self: Arc<Self>,
            request: Request,
            models: Arc<RuntimeModels>,
            state: State,
        ) -> StepStream {
            Box::pin(
                Arc::clone(&self.0)
                    .run_stream(request, models, state)
                    .map(move |step| {
                        step.map(|mut step| {
                            (self.1)(&mut step);
                            step
                        })
                    }),
            )
        }
    }

    fn restricted_models() -> RuntimeModels {
        RuntimeModels::new(HashMap::from([
            (Category::Any, target_set(&["allowed", "fallback"])),
            (Category::Judge, target_set(&["judge"])),
        ]))
        .with_target_restriction()
        .with_subagent(HashMap::from([(Category::Any, target_set(&["subagent"]))]))
    }

    #[tokio::test]
    async fn target_restriction_checks_default_streams() {
        for probe in [
            TargetProbe::Model(&["outside"]),
            TargetProbe::Model(&["allowed", "outside"]),
            TargetProbe::Decision("outside"),
            TargetProbe::Completion(&["outside"]),
            TargetProbe::Completion(&["allowed", "outside"]),
        ] {
            let mut stream = Arc::new(probe).run_stream(
                request(),
                Arc::new(restricted_models()),
                State::default(),
            );
            assert!(matches!(
                stream.next().await,
                Some(Err(LibsyError::TargetNotFound { target })) if target == "outside"
            ));
            assert!(stream.next().await.is_none());
        }
        let mut stream = Arc::new(TargetProbe::MismatchedCompletion).run_stream(
            request(),
            Arc::new(restricted_models()),
            State::default(),
        );
        assert!(matches!(
            stream.next().await,
            Some(Err(LibsyError::AlgorithmError { .. }))
        ));

        let mut stream = Arc::new(TargetProbe::Model(&["outside"])).run_stream(
            request(),
            Arc::new(RuntimeModels::default()),
            State::default(),
        );
        assert!(matches!(stream.next().await, Some(Ok(Step::CallModel(_)))));
    }

    #[tokio::test]
    async fn target_restriction_rechecks_rewritten_steps() {
        for probe in [
            TargetProbe::Model(&["allowed"]),
            TargetProbe::Decision("judge"),
            TargetProbe::Completion(&["allowed"]),
        ] {
            let probe = Arc::new(probe);
            let run = |rewrite| {
                drive(
                    Arc::new(RewrittenSteps(Arc::clone(&probe), rewrite)),
                    request(),
                    Arc::new(restricted_models()),
                    |_| async { panic!("invalid work must not reach the host") },
                )
            };
            let result = run(|step| match step {
                Step::CallModel(call) => call.models.push("outside".into()),
                Step::CallDecision(call) => call.model = "outside".into(),
                Step::Done(outcome) => outcome.selected_model_ids.push("outside".into()),
            })
            .await;
            assert!(
                matches!(result, Err(LibsyError::TargetNotFound { target }) if target == "outside")
            );

            let result = run(|step| match step {
                Step::CallModel(call) => call.request.llm_request.model = Some("outside".into()),
                Step::CallDecision(call) => call.request.model = Some("allowed".into()),
                Step::Done(outcome) => outcome.request.llm_request.model = Some("outside".into()),
            })
            .await;
            assert!(matches!(result, Err(LibsyError::AlgorithmError { .. })));
        }
    }

    #[tokio::test]
    async fn target_restriction_preserves_allowed_work() -> Result<()> {
        for (probe, expected) in [
            (
                TargetProbe::Model(&["allowed", "fallback"]),
                &["allowed"][..],
            ),
            (TargetProbe::Model(&["judge"]), &["judge"][..]),
            (TargetProbe::Decision("judge"), &["allowed"][..]),
            (
                TargetProbe::Completion(&["allowed", "fallback"]),
                &["allowed", "fallback"][..],
            ),
            (TargetProbe::Completion(&["subagent"]), &["subagent"][..]),
        ] {
            let outcome = drive(
                Arc::new(probe),
                request(),
                Arc::new(restricted_models()),
                |call| async move {
                    match call {
                        Call::Model(call) => {
                            call.respond(std::future::ready(Ok(reply("ok")))).await
                        }
                        Call::Decision(call) => serve_decision(*call).await,
                    }
                },
            )
            .await?;
            assert_eq!(outcome.selected_model_ids, target_set(expected));
        }
        let outcome = drive(
            Arc::new(crate::Noop {}),
            request(),
            Arc::new(RuntimeModels::default().with_target_restriction()),
            |_| async { panic!("Noop must not request model calls") },
        )
        .await?;
        assert!(outcome.response.is_some());
        Ok(())
    }

    #[tokio::test]
    async fn typed_driver_preserves_call_and_stream_boundaries() -> Result<()> {
        tokio::time::timeout(std::time::Duration::from_secs(1), async {
            // Distinct oneshots keep reverse-order replies paired with their producers, and a
            // retained call remains pending until the host responds.
            let (driver, mut step_rx) = Driver::new("test", Arc::new(RuntimeModels::default()));
            let first_driver = driver.clone();
            let mut first = tokio::spawn(async move {
                first_driver
                    .call_model(request(), vec![ModelId::from("first")])
                    .await
            });
            let second = tokio::spawn(async move {
                driver
                    .call_model(request(), vec![ModelId::from("second")])
                    .await
            });

            let mut calls = HashMap::new();
            for _ in 0..2 {
                let step = step_rx.recv().await.ok_or(DriverError::StreamClosed)??;
                let Step::CallModel(call) = step else {
                    return Err(test_error("expected a CallModel step"));
                };
                let selected_model = call
                    .models
                    .first()
                    .ok_or_else(|| test_error("model call has no candidates"))?
                    .to_string();
                calls.insert(selected_model, call);
            }
            assert!(
                tokio::time::timeout(std::time::Duration::from_millis(20), &mut first)
                    .await
                    .is_err(),
                "call completed before the host responded"
            );
            calls
                .remove("second")
                .ok_or_else(|| test_error("missing second call"))?
                .respond(std::future::ready(Ok(reply("second response")))).await?;
            calls
                .remove("first")
                .ok_or_else(|| test_error("missing first call"))?
                .respond(std::future::ready(Ok(reply("first response")))).await?;

            let first_response = first
                .await
                .map_err(|source| LibsyError::external("joining a test task", source))??;
            let second_response = second
                .await
                .map_err(|source| LibsyError::external("joining a test task", source))??;
            assert_eq!(
                first_response.llm_response.as_agg().map(completion_text),
                Some("first response".to_string())
            );
            assert_eq!(
                second_response.llm_response.as_agg().map(completion_text),
                Some("second response".to_string())
            );

            // Dropping the host-facing promise closes only that call's reply channel.
            let (driver, mut step_rx) = Driver::new("test", Arc::new(RuntimeModels::default()));
            let producer = tokio::spawn(async move {
                driver
                    .call_model(request(), vec![ModelId::from("dropped")])
                    .await
            });
            let step = step_rx.recv().await.ok_or(DriverError::StreamClosed)??;
            let Step::CallModel(call) = step else {
                return Err(test_error("expected a CallModel step"));
            };
            drop(call);
            let result = producer
                .await
                .map_err(|source| LibsyError::external("joining a test task", source))?;
            assert!(matches!(
                result,
                Err(LibsyError::Driver(DriverError::ResponseDropped))
            ));

            // A standalone driver reports the typed step receiver disappearing at its next call.
            let (driver, step_rx) = Driver::new("test", Arc::new(RuntimeModels::default()));
            drop(step_rx);
            let result = driver
                .call_model(request(), vec![ModelId::from("closed")])
                .await;
            assert!(matches!(
                result,
                Err(LibsyError::Driver(DriverError::StreamClosed))
            ));

            fn decision_request() -> DecisionRequest {
                DecisionRequest {
                    model: Some("overwritten".into()),
                    input: serde_json::json!({"task": "choose a route"}),
                    questions: Default::default(),
                }
            }

            fn decision_response() -> DecisionResponse {
                DecisionResponse {
                    id: Some("decision-1".to_string()),
                    model: Some("provider-model".into()),
                    answers: [(
                        "p_solve".to_string(),
                        switchyard_protocol::DecisionAnswer {
                            value: switchyard_protocol::DecisionValue::Boolean(
                                switchyard_protocol::BooleanEstimate::ProbabilityTrue(
                                    switchyard_protocol::Probability(0.8),
                                ),
                            ),
                            provider_confidence: None,
                        },
                    )]
                    .into(),
                    usage: Default::default(),
                }
            }

            struct MixedCalls(&'static str);

            #[async_trait]
            impl Algorithm for MixedCalls {
                fn name(&self) -> &str {
                    "mixed"
                }

                async fn route(
                    self: Arc<Self>,
                    driver: Driver,
                    request: Request,
                ) -> Result<RoutingOutcome> {
                    let (llm, decision) = tokio::join!(
                        driver.call_model(request.clone(), vec!["llm".into()]),
                        driver.call_decision(decision_request(), "decision".into()),
                    );
                    assert_eq!(
                        llm?.llm_response.as_agg().map(completion_text),
                        Some("llm reply".into())
                    );
                    match self.0 {
                        "mock" => assert_eq!(decision?, DecisionResponse {
                            id: None,
                            model: Some("decision".into()),
                            answers: Default::default(),
                            usage: Default::default(),
                        }),
                        "reply" => assert_eq!(decision?, decision_response()),
                        "error" => assert!(matches!(
                            decision,
                            Err(LibsyError::AlgorithmError { message }) if message == "provider failed"
                        )),
                        "drop" => assert!(matches!(
                            decision,
                            Err(LibsyError::Driver(DriverError::ResponseDropped))
                        )),
                        "abort" => return std::future::pending().await,
                        _ => unreachable!(),
                    }
                    Ok(RoutingOutcome::route_to("answer".into(), vec![], request))
                }
            }

            for mode in ["mock", "reply", "error", "drop", "abort"] {
                // Serial dispatch would deadlock here and fail the enclosing timeout.
                let barrier = Arc::new(tokio::sync::Barrier::new(2));
                let outcome = drive(
                    Arc::new(MixedCalls(mode)),
                    request(),
                    Arc::new(RuntimeModels::default()),
                    move |call| {
                        let barrier = barrier.clone();
                        async move {
                            barrier.wait().await;
                            let call = match call {
                                Call::Model(call) => return call.respond(std::future::ready(Ok(reply("llm reply")))).await,
                                Call::Decision(call) => *call,
                            };
                            assert_eq!(call.algorithm, "mixed");
                            assert_eq!(call.model, "decision");
                            assert_eq!(call.request.model, Some("decision".into()));
                            assert_eq!(call.request.input, decision_request().input);
                            match mode {
                                "mock" => serve_decision(call).await,
                                "reply" => call.respond(std::future::ready(Ok(decision_response()))).await,
                                "error" => call.respond(std::future::ready(Err(LibsyError::AlgorithmError {
                                    message: "provider failed".into(),
                                }))).await,
                                "drop" => {
                                    drop(call);
                                    Ok(())
                                }
                                "abort" => call.fail(test_error("host aborted")),
                                _ => unreachable!(),
                            }
                        }
                    },
                )
                .await;
                if mode == "abort" {
                    assert!(matches!(outcome, Err(LibsyError::External { .. })));
                } else {
                    assert_eq!(outcome?.selected_model_id()?, "answer");
                }
            }

            let (driver, mut steps) = Driver::new("test", Arc::new(RuntimeModels::default()));
            let mut pending = Box::pin(driver.call_decision(decision_request(), "decision".into()));
            assert!(futures::poll!(&mut pending).is_pending());
            let Some(Ok(Step::CallDecision(call))) = steps.recv().await else {
                return Err(test_error("expected a decision call"));
            };
            drop(pending);
            call.respond(async { panic!("cancelled work must not be polled") }).await?;
            drop(steps);
            assert!(matches!(
                driver.call_decision(decision_request(), "decision".into()).await,
                Err(LibsyError::Driver(DriverError::StreamClosed))
            ));
            Ok(())
        })
        .await
        .map_err(|error| LibsyError::external("waiting for typed driver boundaries", error))?
    }

    #[test]
    fn target_lookup_returns_the_missing_target() {
        let error = ensure_model_is_target(&target_set(&[]), &ModelId::from("missing")).err();
        assert!(matches!(
            error,
            Some(LibsyError::TargetNotFound { target }) if target == "missing"
        ));
    }

    /// Build a single-target algo, plus a `serve` that answers it as a token stream
    /// replaying `chunks` in order (as `Ok` items).
    fn streaming_orch(chunks: Vec<LlmResponseChunk>) -> (Arc<dyn Algorithm>, impl Serve) {
        let algo = orch(target_set(&["stream/model"]));
        let serve = move |_target: ModelId, _request: Request| {
            let chunks = chunks.clone();
            async move {
                let stream =
                    futures::stream::iter(chunks.into_iter().map(|chunk| Ok(chunk.into()))).boxed();
                Ok(Response {
                    llm_response: LlmResponse::Stream(stream),
                    metadata: None,
                    upstream_headers: http::HeaderMap::new(),
                })
            }
        };
        (algo, serve)
    }

    #[tokio::test]
    async fn run_returns_a_streamed_response_the_caller_aggregates() -> Result<()> {
        // A streaming client -> its chunks flow through the promise and `Done`,
        // and `run` returns the live stream untouched for the caller to fold.
        let (orch, serve) = streaming_orch(vec![
            LlmResponseChunk::MessageStart {
                id: Some("m1".to_string()),
                model: Some("stream/model".to_string()),
            },
            LlmResponseChunk::TextDelta {
                index: 0,
                text: "hel".to_string(),
            },
            LlmResponseChunk::TextDelta {
                index: 0,
                text: "lo".to_string(),
            },
            LlmResponseChunk::MessageStop {
                reason: Some("stop".to_string()),
            },
        ]);
        let (selected_model, response) = test_drive(orch, request(), serve).await?;
        // The run handed back the live stream; the caller folds it to a buffered aggregate.
        let agg = response
            .llm_response
            .into_agg()
            .await
            .map_err(|error| LibsyError::external("aggregating response stream", error))?;
        assert_eq!(completion_text(&agg), "hello");
        assert_eq!(agg.model.as_deref(), Some("stream/model"));
        assert_eq!(selected_model, "stream/model");
        Ok(())
    }

    #[tokio::test]
    async fn aggregating_a_streamed_response_propagates_a_mid_stream_error() -> Result<()> {
        // The run succeeds and returns the stream; the in-band `Error` chunk surfaces only
        // when the caller aggregates it.
        let (orch, serve) = streaming_orch(vec![
            LlmResponseChunk::TextDelta {
                index: 0,
                text: "partial".to_string(),
            },
            LlmResponseChunk::StreamError {
                message: "upstream exploded".to_string(),
            },
        ]);
        let (_, response) = test_drive(orch, request(), serve).await?;
        match response.llm_response.into_agg().await {
            Ok(_) => panic!("expected a mid-stream error, got an aggregate"),
            Err(err) => {
                assert!(err.to_string().contains("upstream exploded"));
                Ok(())
            }
        }
    }

    #[tokio::test]
    async fn run_offloads_via_promise_then_finishes() -> Result<()> {
        // Every call is offloaded via a promise the orchestrator surfaces as a
        // `CallModel` step for us to fulfill.
        let stream = orch(target_set(&["offload/model"])).run_stream(
            request(),
            Arc::new(RuntimeModels::default()),
            State::default(),
        );
        tokio::pin!(stream);

        let mut saw_call = false;
        let mut final_completion = None;
        while let Some(step) = stream.next().await {
            match step? {
                Step::CallDecision(_) => return Err(test_error("unexpected decision call")),
                Step::CallModel(call) => {
                    saw_call = true;
                    assert_eq!(call.models, vec![ModelId::from("offload/model")]);
                    // Fulfilling the promise is the "real" model call the caller makes.
                    call.respond(std::future::ready(Ok(Response {
                        llm_response: LlmResponse::Agg(text_response(
                            None,
                            "fulfilled".to_string(),
                        )),
                        metadata: None,
                        upstream_headers: http::HeaderMap::new(),
                    })))
                    .await?;
                }
                Step::Done(outcome) => {
                    let metadata = outcome
                        .metadata
                        .as_ref()
                        .expect("run_stream should attach outcome metadata");
                    assert_eq!(metadata.algorithm, "test");
                    assert_eq!(
                        uuid::Uuid::parse_str(metadata.outcome_id())
                            .expect("outcome id should be a UUID")
                            .get_version_num(),
                        7
                    );
                    assert_eq!(
                        metadata.evidence,
                        Some(serde_json::json!({"source": "test"}))
                    );
                    assert_eq!(
                        metadata.judgments[0].evidence,
                        Some(serde_json::json!({"source": "judge"}))
                    );
                    let response = outcome
                        .response
                        .ok_or_else(|| test_error("expected an answered outcome"))?;
                    final_completion = Some(
                        response
                            .llm_response
                            .as_agg()
                            .map(completion_text)
                            .unwrap_or_default(),
                    );
                }
            }
        }

        assert!(saw_call, "expected a CallModel step before Done");
        assert_eq!(
            final_completion.ok_or_else(|| test_error("no Done step"))?,
            "fulfilled"
        );
        Ok(())
    }

    #[tokio::test(flavor = "multi_thread", worker_threads = 12)]
    async fn requests_are_processed_in_parallel() -> Result<()> {
        use std::time::Duration;
        use tokio::sync::Barrier;

        const N: usize = 12;

        // Serving blocks until all N concurrent calls have arrived. If requests were
        // serialized (one algorithm behind a `Mutex`), only one call could be in flight,
        // the barrier would never reach N, and the test would time out. It passes only
        // because the shared algorithm is driven concurrently across requests.
        let barrier = Arc::new(Barrier::new(N));
        // One shared algorithm driven by many concurrent requests.
        let algo = orch(target_set(&["m"]));

        let mut handles = Vec::new();
        for _ in 0..N {
            let algo = algo.clone();
            let barrier = barrier.clone();
            let serve = move |target: ModelId, _request: Request| {
                let barrier = barrier.clone();
                async move {
                    barrier.wait().await;
                    Ok(reply(target))
                }
            };
            handles.push(tokio::spawn(async move {
                test_drive(algo, request(), serve)
                    .await
                    .map(|(_, response)| {
                        response
                            .llm_response
                            .as_agg()
                            .map(completion_text)
                            .unwrap_or_default()
                    })
            }));
        }

        for handle in handles {
            // The timeout turns a serialization deadlock into a failure, not a hang.
            let completion = tokio::time::timeout(Duration::from_secs(5), handle)
                .await
                .map_err(|error| LibsyError::external("waiting for test task", error))?
                .map_err(|source| LibsyError::external("joining a test task", source))??;
            assert_eq!(completion, "m");
        }
        Ok(())
    }

    #[tokio::test]
    async fn offload_error_propagates_back_to_the_algorithm() -> Result<()> {
        // A client-less target offloads its call; we fulfill the promise with an
        // Err, which must flow back through `call_model_target` into the algorithm and
        // out as an error step — not a response.
        let stream = orch(target_set(&["offload/model"])).run_stream(
            request(),
            Arc::new(RuntimeModels::default()),
            State::default(),
        );
        tokio::pin!(stream);

        let mut saw_error = false;
        while let Some(step) = stream.next().await {
            match step {
                Ok(Step::CallDecision(_)) => return Err(test_error("unexpected decision call")),
                Ok(Step::CallModel(mut call)) => {
                    call.recover_errors = true;
                    call.respond(std::future::ready(Err(test_error(
                        "upstream model call failed",
                    ))))
                    .await?;
                }
                Ok(Step::Done(..)) => {
                    return Err(test_error(
                        "expected the offload error to propagate, got a response",
                    ));
                }
                Err(err) => {
                    // The algorithm's `call_model_target` saw the error via the promise.
                    assert!(err.to_string().contains("upstream model call failed"));
                    saw_error = true;
                }
            }
        }

        assert!(saw_error, "expected an error step");
        Ok(())
    }

    #[tokio::test]
    async fn dropping_the_stream_cancels_the_algorithm_task() -> Result<()> {
        use std::sync::atomic::{AtomicBool, Ordering};
        use std::time::Duration;
        use tokio::sync::mpsc;

        // Sets a flag when dropped, so we can observe whether the algorithm task was
        // cancelled/dropped.
        struct DropGuard(Arc<AtomicBool>);
        impl Drop for DropGuard {
            fn drop(&mut self) {
                self.0.store(true, Ordering::SeqCst);
            }
        }

        struct StuckAlgo {
            started: mpsc::UnboundedSender<()>,
            dropped: Arc<AtomicBool>,
        }

        #[async_trait]
        impl Algorithm for StuckAlgo {
            fn name(&self) -> &str {
                "stuck"
            }

            async fn route(
                self: Arc<Self>,
                _driver: Driver,
                _request: Request,
            ) -> Result<RoutingOutcome> {
                let _guard = DropGuard(self.dropped.clone());
                let _ = self.started.send(());
                // Await forever without ever touching the driver.
                std::future::pending::<()>().await;
                unreachable!()
            }
        }

        let (started_tx, mut started_rx) = mpsc::unbounded_channel();
        let dropped = Arc::new(AtomicBool::new(false));
        let algo: Arc<dyn Algorithm> = Arc::new(StuckAlgo {
            started: started_tx,
            dropped: dropped.clone(),
        });

        let stream = algo.run_stream(
            request(),
            Arc::new(RuntimeModels::default()),
            State::default(),
        );
        started_rx
            .recv()
            .await
            .ok_or_else(|| test_error("task never started"))?;
        drop(stream);
        tokio::time::sleep(Duration::from_millis(100)).await;

        assert!(
            dropped.load(Ordering::SeqCst),
            "algorithm task was NOT cancelled after dropping the stream"
        );
        Ok(())
    }

    #[tokio::test]
    async fn route_panic_surfaces_as_a_stream_error() -> Result<()> {
        // An algorithm whose task panics must surface an `Err` step carrying the panic
        // message, not abort the process from an unobserved detached task.
        struct Panicky;

        #[async_trait]
        impl Algorithm for Panicky {
            fn name(&self) -> &str {
                "panicky"
            }

            async fn route(
                self: Arc<Self>,
                _driver: Driver,
                _request: Request,
            ) -> Result<RoutingOutcome> {
                panic!("boom");
            }
        }

        let algo: Arc<dyn Algorithm> = Arc::new(Panicky);
        let stream = algo.run_stream(
            request(),
            Arc::new(RuntimeModels::default()),
            State::default(),
        );
        tokio::pin!(stream);

        let mut saw_error = false;
        while let Some(step) = stream.next().await {
            match step {
                Err(err) => {
                    // The panic message is preserved, not flattened into an opaque failure.
                    assert!(err.to_string().contains("algorithm task panicked: boom"));
                    saw_error = true;
                }
                Ok(_) => return Err(test_error("expected the panic to surface as an error step")),
            }
        }

        assert!(saw_error, "expected an error step from the panicked task");
        Ok(())
    }

    /// A panicking algorithm must publish its terminal step even when it left a `Driver`
    /// clone alive in another task. That clone holds the step channel open, so a run that
    /// merely unwound would never terminate and the consumer would wait forever.
    #[tokio::test]
    async fn a_panic_with_a_leaked_driver_clone_still_terminates_the_run() -> Result<()> {
        struct LeakyPanic;

        #[async_trait]
        impl Algorithm for LeakyPanic {
            fn name(&self) -> &str {
                "leaky_panic"
            }

            async fn route(
                self: Arc<Self>,
                driver: Driver,
                _request: Request,
            ) -> Result<RoutingOutcome> {
                tokio::spawn(async move {
                    // Outlives the panic below, keeping a sender clone alive.
                    let _keep_alive = driver;
                    std::future::pending::<()>().await;
                });
                tokio::task::yield_now().await;
                panic!("boom");
            }
        }

        let algo: Arc<dyn Algorithm> = Arc::new(LeakyPanic);
        // The timeout turns the hang this guards against into a failure rather than a hang.
        let result = tokio::time::timeout(
            std::time::Duration::from_secs(1),
            test_drive(algo, request(), echo()),
        )
        .await
        .map_err(|error| LibsyError::external("waiting for the panicked run to end", error))?;

        match result {
            Ok(_) => Err(test_error(
                "expected the panic to end the run with an error",
            )),
            Err(err) => {
                assert!(err.to_string().contains("algorithm task panicked: boom"));
                Ok(())
            }
        }
    }

    #[tokio::test]
    async fn cancelling_run_cancels_the_algorithm_task() -> Result<()> {
        use std::sync::atomic::{AtomicBool, Ordering};
        use std::time::Duration;
        use tokio::sync::mpsc;

        // Sets a flag when dropped, so we can observe whether the algorithm task was
        // cancelled once the `run` future driving it is dropped.
        struct DropGuard(Arc<AtomicBool>);
        impl Drop for DropGuard {
            fn drop(&mut self) {
                self.0.store(true, Ordering::SeqCst);
            }
        }

        struct StuckAlgo {
            started: mpsc::UnboundedSender<()>,
            dropped: Arc<AtomicBool>,
        }

        #[async_trait]
        impl Algorithm for StuckAlgo {
            fn name(&self) -> &str {
                "stuck"
            }

            async fn route(
                self: Arc<Self>,
                _driver: Driver,
                _request: Request,
            ) -> Result<RoutingOutcome> {
                let _guard = DropGuard(self.dropped.clone());
                let _ = self.started.send(());
                // Hang forever without ever touching the driver, so only cancellation
                // (not a dropped step channel) can stop this task.
                std::future::pending::<()>().await;
                unreachable!()
            }
        }

        let (started_tx, mut started_rx) = mpsc::unbounded_channel();
        let dropped = Arc::new(AtomicBool::new(false));
        let algo: Arc<dyn Algorithm> = Arc::new(StuckAlgo {
            started: started_tx,
            dropped: dropped.clone(),
        });

        // Drive the run on its own task, wait until the algorithm task is up, then cancel
        // it — dropping its future (and the `run_stream` stream it holds).
        let run_task = tokio::spawn(async move { test_drive(algo, request(), echo()).await });
        started_rx
            .recv()
            .await
            .ok_or_else(|| test_error("task never started"))?;
        run_task.abort();
        tokio::time::sleep(Duration::from_millis(100)).await;

        assert!(
            dropped.load(Ordering::SeqCst),
            "algorithm task was NOT cancelled after cancelling run"
        );
        Ok(())
    }

    // --- first-wins hedging: `run` must not wait on losing speculative calls -------------

    /// Offloads two targets concurrently and returns the first to resolve, dropping the
    /// loser's call (first-wins hedging).
    struct Hedge {
        winner: String,
        loser: String,
    }

    #[async_trait]
    impl Algorithm for Hedge {
        fn name(&self) -> &str {
            "hedge"
        }

        async fn route(
            self: Arc<Self>,
            driver: Driver,
            request: Request,
        ) -> Result<RoutingOutcome> {
            let outcome_request = request.clone();
            let win = driver.call_model(request.clone(), vec![self.winner.clone().into()]);
            let lose = driver.call_model(request, vec![self.loser.clone().into()]);
            // First to resolve wins; `select!` drops the losing future (and its promise).
            tokio::select! {
                res = win => Ok(RoutingOutcome::answered(
                    self.winner.clone().into(),
                    outcome_request,
                    res?,
                )),
                res = lose => Ok(RoutingOutcome::answered(
                    self.loser.clone().into(),
                    outcome_request,
                    res?,
                )),
            }
        }
    }

    /// Builds a hedging algo and the `serve` that drives it: the winner is gated behind the
    /// loser starting (so the loser's serve is guaranteed in flight when the winner wins),
    /// and the loser finishes after `loser_delay` — or never, when `None`.
    fn hedge(loser_delay: Option<std::time::Duration>) -> (Arc<dyn Algorithm>, impl Serve) {
        let started = Arc::new(tokio::sync::Notify::new());
        let algo = Arc::new(Hedge {
            winner: "winner".to_string(),
            loser: "loser".to_string(),
        });
        let serve = move |target: ModelId, _request: Request| {
            let started = started.clone();
            async move {
                if target == "loser" {
                    started.notify_one();
                    match loser_delay {
                        Some(delay) => tokio::time::sleep(delay).await,
                        None => std::future::pending::<()>().await,
                    }
                } else {
                    started.notified().await;
                }
                Ok(reply(target))
            }
        };
        (algo, serve)
    }

    #[tokio::test]
    async fn run_returns_the_winner_without_a_late_loser_overwriting_it() -> Result<()> {
        // The loser responds 50ms after the winner has already won. `run` must return the
        // winner, not the loser's `respond`-to-a-dropped-receiver error.
        let (algo, serve) = hedge(Some(std::time::Duration::from_millis(50)));
        let (_, response) = test_drive(algo, request(), serve).await?;
        assert_eq!(
            response
                .llm_response
                .as_agg()
                .map(completion_text)
                .unwrap_or_default(),
            "winner"
        );
        Ok(())
    }

    #[tokio::test]
    async fn run_returns_the_winner_without_hanging_on_a_pending_loser() -> Result<()> {
        // The loser never resolves. `run` must return the winner promptly, not hang
        // waiting for the in-flight loser.
        let (algo, serve) = hedge(None);
        let run = test_drive(algo, request(), serve);
        let (_, response) = tokio::time::timeout(std::time::Duration::from_secs(1), run)
            .await
            .map_err(|error| LibsyError::external("waiting for pending loser", error))??;
        assert_eq!(
            response
                .llm_response
                .as_agg()
                .map(completion_text)
                .unwrap_or_default(),
            "winner"
        );
        Ok(())
    }

    #[tokio::test]
    async fn run_surfaces_a_terminal_error_with_many_calls_in_flight() -> Result<()> {
        use std::sync::atomic::{AtomicUsize, Ordering};

        // A large fan-out (10 matched the old, now-removed concurrency cap). The terminal
        // error must still reach the caller with all of these calls pending.
        const N: usize = 10;

        // Fans out N calls, then errors as soon as all N are in flight — exercising a
        // terminal failure emitted while the offloaded calls are still pending.
        struct FanOutThenError {
            all_started: Arc<tokio::sync::Notify>,
            n: usize,
        }

        #[async_trait]
        impl Algorithm for FanOutThenError {
            fn name(&self) -> &str {
                "fan_out_then_error"
            }

            async fn route(
                self: Arc<Self>,
                driver: Driver,
                request: Request,
            ) -> Result<RoutingOutcome> {
                let offloads = futures::future::join_all(
                    (0..self.n)
                        .map(|i| driver.call_model(request.clone(), vec![format!("m{i}").into()])),
                );
                tokio::select! {
                    _ = offloads => Err(test_error("offloads unexpectedly completed")),
                    _ = self.all_started.notified() => {
                        Err(test_error("terminal error while calls pending"))
                    }
                }
            }
        }

        let all_started = Arc::new(tokio::sync::Notify::new());
        let algo: Arc<dyn Algorithm> = Arc::new(FanOutThenError {
            all_started: all_started.clone(),
            n: N,
        });

        // Serving enters each call; once all N are in flight it signals, then pends forever.
        let started = Arc::new(AtomicUsize::new(0));
        let serve = move |_target: ModelId, _request: Request| {
            let started = started.clone();
            let all_started = all_started.clone();
            async move {
                if started.fetch_add(1, Ordering::SeqCst) + 1 == N {
                    all_started.notify_one();
                }
                std::future::pending::<ServeResult>().await
            }
        };

        // With the cap gone, the driver keeps polling the stream even with N calls in
        // flight, so the terminal error surfaces promptly instead of hanging.
        let run = test_drive(algo, request(), serve);
        let result = tokio::time::timeout(std::time::Duration::from_millis(500), run)
            .await
            .map_err(|error| {
                LibsyError::external("waiting for terminal error with full call cap", error)
            })?;
        match result {
            Ok(_) => Err(test_error("expected the terminal error, got a response")),
            Err(err) => {
                assert!(
                    err.to_string()
                        .contains("terminal error while calls pending")
                );
                Ok(())
            }
        }
    }
}