switchyard_libsy/algorithms/

llm_class.rs

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

//! Judge-backed capability, escalation, and custom-policy routing.

use std::collections::HashSet;
use std::sync::Arc;

use async_trait::async_trait;
use serde::{Deserialize, Deserializer};
use serde_json::Value;
use switchyard_protocol::{Category, ContentBlock, Message, Role};

mod decision;
use decision::DecisionClassifier;
pub use decision::{DecisionJudgeConfig, DecisionPolicy, DecisionRule};

use super::escalation;
use super::fall_through::FallThrough;
use super::util::DEFAULT_JUDGE_MAX_OUTPUT_TOKENS;
use super::util::affinity::{AffinityRouter, ClassifyTrigger};
use super::util::classifier_contract::{
    ClassifierContract, ClassifierContractConfig, ClassifierResponseFormat,
};
use super::util::escalation::EscalationJudgeConfig;
use super::util::llm_judge::{
    ClassifierInput, JsonSchemaDecoder, JudgeClassifier, JudgePolicy, JudgeRuntimeConfig,
    SerdeDecoder, StructuredJudge,
};
use super::util::target_selector::TargetSelectorPolicy;
use crate::core::algorithm::{Algorithm, Driver};
use crate::core::classifier::{Classification, Classifier, Score};
use crate::core::state::State;
use crate::{LibsyError, Result};
use switchyard_protocol::{Request, Response};

const PROMPT_TEMPLATE: &str = include_str!("../prompts/capability-classifier/prompt.md");
const SCHEMA_TEMPLATE: &str = include_str!("../prompts/capability-classifier/schema.json");
/// Telemetry label for this algorithm's spans, metrics, and logs.
const ALGORITHM_NAME: &str = "llm_task_classifier";

/// Restates the task after the conversation the judge is asked to route.
///
/// The rubric leads the request, which works while the payload is a single short
/// message. Once `recent_turn_window` includes real conversation, those instructions
/// sit far from the generation point and the judge sometimes *answers* the
/// conversation instead of classifying it — the reply then fails to parse, the
/// verdict is unavailable, and the turn silently falls back. Repeating the task at
/// the end is what reliably prevents that.
const TRAILING_ROUTING_INSTRUCTION: &str =
    "Route the conversation above. Output ONLY the routing JSON object, nothing else.";

#[derive(Deserialize)]
#[serde(deny_unknown_fields)]
struct TaskClassifierVerdict {
    crux: String,
    primary_rule: String,
    capability_boundary: String,
    p_solve: f64,
}

impl TaskClassifierVerdict {
    /// Rejects malformed or internally inconsistent verdicts before policy evaluation.
    fn is_valid(&self) -> bool {
        (0.0..=1.0).contains(&self.p_solve)
            && !self.crux.trim().is_empty()
            && matches!(
                (
                    self.primary_rule.as_str(),
                    self.capability_boundary.as_str()
                ),
                ("SUP-1" | "SUP-2" | "SUP-3" | "SUP-4" | "SUP-5", "supported")
                    | ("UNC-1" | "UNC-2", "uncertain")
                    | ("LIM-1" | "LIM-2", "unsupported")
                    | ("none", "unmatched")
            )
    }

    /// Returns the number of threshold steps assigned to this capability boundary.
    fn boundary_steps(&self) -> Option<u8> {
        match self.capability_boundary.as_str() {
            "supported" => Some(0),
            "uncertain" | "unmatched" => Some(1),
            "unsupported" => Some(2),
            _ => None,
        }
    }
}

/// Keeps the opening task and the last `recent_turn_window` turns after it. A
/// window of `0` keeps the task alone.
///
/// Inbound decoders normalize client system and developer content into
/// `LlmRequest::instructions`, so it never reaches this list.
///
/// Selects by reference and clones only what survives — a coding-agent
/// conversation carries every tool result, so cloning it whole to keep a window
/// would copy the transcript on each judged turn.
fn trim_messages(messages: &[Message], recent_turn_window: usize) -> Vec<Message> {
    let is_instruction = |message: &Message| matches!(message.role, Role::System | Role::Developer);
    let mut kept: Vec<&Message> = messages.iter().filter(|m| is_instruction(m)).collect();
    let Some(task) = messages.iter().position(|m| m.role == Role::User) else {
        return kept.into_iter().cloned().collect();
    };
    kept.push(&messages[task]);

    let tail: Vec<&Message> = messages[task + 1..]
        .iter()
        .filter(|m| !is_instruction(m))
        .collect();
    kept.extend(&tail[window_start(&tail, recent_turn_window)..]);
    kept.into_iter().cloned().collect()
}

/// The first index of the trailing window.
///
/// Counting messages alone can start the window between an assistant tool call and the
/// result answering it, leaving the judge a result whose call id was never introduced. The
/// start therefore moves back to the nearest one that keeps every tool pair whole.
///
/// One newest-to-oldest pass carries the ids still waiting for a call. Direction is what
/// makes it correct: ids repeat across a conversation, and in this order a call is only
/// ever seen after the results it could answer, so a later call — already passed — clears
/// nothing. A result whose call sits before the opening task, which trimming never reaches,
/// keeps the set non-empty to the end and falls back to the counted start, so an unpairable
/// result costs one pass and cannot widen the window to the whole conversation.
fn window_start(tail: &[&Message], recent_turn_window: usize) -> usize {
    let counted = tail.len().saturating_sub(recent_turn_window);
    // An empty window holds no result to pair, and the loop below never visits its start.
    if counted == tail.len() {
        return counted;
    }
    let mut unpaired: HashSet<&str> = HashSet::new();
    for (start, message) in tail.iter().enumerate().rev() {
        // Blocks reverse too, so a call answers a result only when it precedes it inside
        // one message as well as across messages.
        for block in message.content.iter().rev() {
            match block {
                ContentBlock::ToolResult(result) => {
                    unpaired.insert(result.tool_call_id.as_str());
                }
                ContentBlock::ToolCall(call) => {
                    unpaired.remove(call.id.as_str());
                }
                _ => {}
            }
        }
        if start <= counted && unpaired.is_empty() {
            return start;
        }
    }
    counted
}

/// Keeps the opening task and the latest user follow-up when they differ.
fn task_messages(messages: &[Message]) -> Vec<Message> {
    // Decoders also use the user role for tool results. Select ordinary user content
    // first, so a tool result cannot replace the opening task or latest follow-up.
    let is_task_content = |block: &ContentBlock| {
        !matches!(
            block,
            ContentBlock::ToolCall(_)
                | ContentBlock::ToolResult(_)
                | ContentBlock::Reasoning { .. }
        )
    };
    let mut user_messages = messages.iter().filter(|message| {
        message.role == Role::User && message.content.iter().any(is_task_content)
    });
    let Some(opening_task) = user_messages.next() else {
        return Vec::new();
    };
    [Some(opening_task), user_messages.next_back()]
        .into_iter()
        .flatten()
        .map(|message| Message {
            role: Role::User,
            content: message
                .content
                .iter()
                .filter(|block| is_task_content(block))
                .cloned()
                .collect(),
        })
        .collect()
}

/// Selects the task messages shown to capability and custom-schema classifiers.
struct TaskInput {
    recent_turn_window: Option<usize>,
}

impl TaskInput {
    /// Selects conversation content without adding judge instructions.
    fn messages(&self, request: &Request) -> Vec<Message> {
        // The default preserves the whole-task anchor and latest user update. A
        // configured window widens that to the surrounding conversation.
        let mut messages = match self.recent_turn_window {
            Some(window) => trim_messages(&request.llm_request.messages, window),
            None => task_messages(&request.llm_request.messages),
        };
        // Reasoning is provider-private and not required to classify the task. Some
        // upstreams also reject an unsigned reasoning item replayed without the
        // opaque state it was issued with, so it cannot travel through a windowed
        // classifier request as ordinary history.
        for message in &mut messages {
            message
                .content
                .retain(|block| !matches!(block, ContentBlock::Reasoning { .. }));
        }
        messages.retain(|message| !message.content.is_empty());
        messages
    }
}

impl ClassifierInput for TaskInput {
    fn build_messages(&self, _state: &State, request: &Request) -> Vec<Message> {
        let mut messages = self.messages(request);
        // Only the windowed path carries assistant turns and tool traffic for the judge
        // to be distracted by. The default path is user task messages only — the anchor
        // and the latest follow-up — so there is nothing there to outrank.
        if self.recent_turn_window.is_some() {
            messages.push(Message::text(
                Role::User,
                TRAILING_ROUTING_INSTRUCTION.to_string(),
            ));
        }
        messages
    }
}

struct TaskClassifierPolicy {
    base_threshold: f64,
    threshold_step: f64,
}

impl TaskClassifierPolicy {
    fn new(config: &LlmCapabilityConfig) -> Self {
        Self {
            base_threshold: config.base_threshold,
            threshold_step: config.threshold_step,
        }
    }

    /// Returns the required solve probability for one validated verdict.
    fn threshold(&self, verdict: &TaskClassifierVerdict) -> Option<f64> {
        Some(self.base_threshold + f64::from(verdict.boundary_steps()?) * self.threshold_step)
    }
}

impl JudgePolicy for TaskClassifierPolicy {
    type Verdict = TaskClassifierVerdict;

    fn to_classification(
        &self,
        verdict: Option<&Self::Verdict>,
        driver: &Driver,
    ) -> Result<Classification> {
        // Judge output is untrusted. An absent, invalid, or inconsistent verdict is
        // ambiguous so the surrounding router applies its configured fallback.
        let Some(verdict) = verdict.filter(|verdict| verdict.is_valid()) else {
            return Ok(Classification::Ambiguous(vec![]));
        };
        // A usable verdict below the capability threshold is still a decision: the judge
        // does not trust the efficient tier with this task.
        let Some(threshold) = self.threshold(verdict) else {
            return Ok(Classification::Ambiguous(vec![]));
        };
        let category = if verdict.p_solve >= threshold
            || (threshold - verdict.p_solve).abs() <= f64::EPSILON
        {
            Category::Efficient
        } else {
            Category::Capable
        };
        // The chosen category may have no models configured. That is ambiguous, not an
        // error, so the surrounding router applies its configured fallback.
        let Some(target) = driver.models_for(&category).first().cloned() else {
            return Ok(Classification::Ambiguous(vec![]));
        };
        Ok(Classification::Scores(vec![Score {
            target,
            confidence: 1.0,
            category: Some(category),
        }]))
    }
}

/// Maps valid verdicts to scores, invalid verdicts to a reason, and leaves absent verdicts alone.
fn capability_evidence(
    policy: &TaskClassifierPolicy,
    verdict: Option<&TaskClassifierVerdict>,
) -> Option<Value> {
    let verdict = verdict?;
    let Some(threshold) = verdict
        .is_valid()
        .then(|| policy.threshold(verdict))
        .flatten()
    else {
        return Some(serde_json::json!({
            "source": "fail_open",
            "reason_code": "invalid_verdict",
        }));
    };
    Some(serde_json::json!({
        "source": "llm-classifier",
        "score": verdict.p_solve,
        "threshold": threshold,
    }))
}

/// Settings shared by capability judges and their surrounding route.
#[derive(Clone, Debug)]
pub struct TaskClassifierConfig {
    /// The judge's prediction contract and routing policy.
    pub judge: CapabilityJudgeConfig,
    /// Uses fallback routing on judge client failures and deadlines. Defaults to true.
    pub fail_open: bool,
    /// How often the classifier re-decides this session's target.
    pub classify_trigger: ClassifyTrigger,
    /// Uses the first user message as the SessionKey for sticky routing when session metadata is unavailable.
    pub message_hash_fallback: bool,
    /// `None` judges the opening task and latest user follow-up. `Some(n)` also
    /// includes the surrounding conversation using the last `n` turns.
    pub recent_turn_window: Option<usize>,
}

/// Selects how a capability judge obtains and interprets its prediction.
#[derive(Clone, Debug)]
pub enum CapabilityJudgeConfig {
    /// A structured LLM verdict with a solve probability and capability boundary.
    Llm(LlmCapabilityConfig),
    /// Candidate probabilities evaluated by an ordered routing policy.
    Decision(DecisionJudgeConfig),
}

impl Default for CapabilityJudgeConfig {
    fn default() -> Self {
        Self::Llm(LlmCapabilityConfig::default())
    }
}

/// Prompt, output, and solve-probability settings for an LLM capability judge.
#[derive(Clone, Debug)]
pub struct LlmCapabilityConfig {
    /// Lowest solve probability that routes a supported task to the efficient target.
    pub base_threshold: f64,
    /// Amount added per capability-boundary step: zero for supported, one for
    /// uncertain or unmatched, and two for unsupported.
    pub threshold_step: f64,
    /// Prompt and verdict contract settings for the classifier judge.
    pub contract: ClassifierContractConfig,
    /// Maximum completion tokens available to the classifier verdict.
    pub max_output_tokens: u64,
}

impl Default for LlmCapabilityConfig {
    fn default() -> Self {
        Self {
            base_threshold: 0.0,
            threshold_step: 0.0,
            contract: ClassifierContractConfig::default(),
            max_output_tokens: DEFAULT_JUDGE_MAX_OUTPUT_TOKENS,
        }
    }
}

/// Flat serialized shape that maps prompt settings into the runtime contract.
#[derive(Deserialize)]
#[serde(deny_unknown_fields)]
struct TaskClassifierConfigWire {
    base_threshold: f64,
    #[serde(default = "default_fail_open")]
    fail_open: bool,
    #[serde(default)]
    threshold_step: f64,
    #[serde(default)]
    classify_trigger: ClassifyTrigger,
    #[serde(default)]
    message_hash_fallback: bool,
    #[serde(default)]
    recent_turn_window: Option<usize>,
    #[serde(default)]
    prompt: Option<String>,
    #[serde(default)]
    response_format_type: ClassifierResponseFormat,
    #[serde(default = "default_judge_max_output_tokens")]
    max_output_tokens: u64,
}

impl<'de> Deserialize<'de> for TaskClassifierConfig {
    fn deserialize<D>(deserializer: D) -> std::result::Result<Self, D::Error>
    where
        D: Deserializer<'de>,
    {
        let wire = TaskClassifierConfigWire::deserialize(deserializer)?;
        let mut contract = ClassifierContractConfig::default();
        if let Some(prompt) = wire.prompt {
            contract = contract.with_prompt(prompt);
        }
        contract = contract.with_response_format_type(wire.response_format_type);
        Ok(Self {
            judge: CapabilityJudgeConfig::Llm(LlmCapabilityConfig {
                base_threshold: wire.base_threshold,
                threshold_step: wire.threshold_step,
                contract,
                max_output_tokens: wire.max_output_tokens,
            }),
            fail_open: wire.fail_open,
            classify_trigger: wire.classify_trigger,
            message_hash_fallback: wire.message_hash_fallback,
            recent_turn_window: wire.recent_turn_window,
        })
    }
}

const fn default_fail_open() -> bool {
    true
}

const fn default_judge_max_output_tokens() -> u64 {
    DEFAULT_JUDGE_MAX_OUTPUT_TOKENS
}

impl Default for TaskClassifierConfig {
    fn default() -> Self {
        Self {
            judge: CapabilityJudgeConfig::default(),
            fail_open: default_fail_open(),
            classify_trigger: ClassifyTrigger::default(),
            message_hash_fallback: false,
            recent_turn_window: None,
        }
    }
}

impl LlmCapabilityConfig {
    fn validate(&self) -> Result<()> {
        if !(0.0..=1.0).contains(&self.base_threshold) {
            return Err(LibsyError::AlgorithmError {
                message: format!(
                    "base_threshold must be between 0 and 1, got {}",
                    self.base_threshold
                ),
            });
        }
        if !self.threshold_step.is_finite() || self.threshold_step < 0.0 {
            return Err(LibsyError::AlgorithmError {
                message: format!(
                    "threshold_step must be finite and greater than or equal to 0, got {}",
                    self.threshold_step
                ),
            });
        }
        let unsupported_threshold = self.base_threshold + 2.0 * self.threshold_step;
        if unsupported_threshold > 1.0 && unsupported_threshold - 1.0 > f64::EPSILON {
            return Err(LibsyError::AlgorithmError {
                message: format!(
                    "base_threshold + 2 * threshold_step must be at most 1, got {unsupported_threshold}"
                ),
            });
        }
        if self.max_output_tokens == 0 {
            return Err(LibsyError::AlgorithmError {
                message: "max_output_tokens must be at least 1".to_string(),
            });
        }
        Ok(())
    }
}

impl TaskClassifierConfig {
    fn validate(&self) -> Result<()> {
        match &self.judge {
            CapabilityJudgeConfig::Llm(config) => config.validate()?,
            CapabilityJudgeConfig::Decision(config) => config.validate()?,
        }
        // Only `every_request` is rejected: it retains no target, so a fallback identity has nothing to key on. Both retaining triggers can key the retained target on a message hash when the caller sends no session id.
        if self.message_hash_fallback && self.classify_trigger == ClassifyTrigger::EveryRequest {
            return Err(LibsyError::AlgorithmError {
                message:
                    "message_hash_fallback requires classify_trigger = new_session or user_turn"
                        .to_string(),
            });
        }
        Ok(())
    }
}

/// Policy that maps a custom classifier verdict to a routing target.
#[derive(Clone, Debug)]
pub enum CustomClassifierPolicy {
    /// Resolves a JSON Pointer and treats its string value as a model category.
    TargetSelector {
        /// JSON Pointer evaluated against each schema-validated verdict.
        selector: String,
    },
}

impl CustomClassifierPolicy {
    /// Creates a policy that selects a model category through a JSON Pointer.
    pub fn target_selector(selector: impl Into<String>) -> Self {
        Self::TargetSelector {
            selector: selector.into(),
        }
    }
}

/// Settings for a classifier whose JSON Schema and target-selection policy are user supplied.
#[derive(Clone, Debug)]
pub struct CustomClassifierConfig {
    /// System prompt sent to the classifier judge.
    pub prompt: String,
    /// Inner JSON Schema placed inside the provider's structured-output wrapper.
    pub response_schema: Value,
    /// Deterministic policy applied after the verdict passes schema validation.
    pub policy: CustomClassifierPolicy,
    /// How often the classifier re-decides this session's target.
    pub classify_trigger: ClassifyTrigger,
    /// Uses the first user message when session metadata is unavailable.
    pub message_hash_fallback: bool,
    /// Trailing conversation turns shown to the classifier judge.
    pub recent_turn_window: Option<usize>,
    /// Maximum completion tokens available to the classifier verdict.
    pub max_output_tokens: u64,
}

impl CustomClassifierConfig {
    /// Creates a custom-schema classifier contract with conservative runtime defaults.
    pub fn new(
        prompt: impl Into<String>,
        response_schema: Value,
        policy: CustomClassifierPolicy,
    ) -> Self {
        Self {
            prompt: prompt.into(),
            response_schema,
            policy,
            classify_trigger: ClassifyTrigger::default(),
            message_hash_fallback: false,
            recent_turn_window: None,
            max_output_tokens: DEFAULT_JUDGE_MAX_OUTPUT_TOKENS,
        }
    }

    fn validate(&self) -> Result<()> {
        if self.max_output_tokens == 0 {
            return Err(LibsyError::AlgorithmError {
                message: "max_output_tokens must be at least 1".to_string(),
            });
        }
        // Only `every_request` is rejected: it retains no target, so a fallback identity has nothing to key on. Both retaining triggers can key the retained target on a message hash when the caller sends no session id.
        if self.message_hash_fallback && self.classify_trigger == ClassifyTrigger::EveryRequest {
            return Err(LibsyError::AlgorithmError {
                message:
                    "message_hash_fallback requires classify_trigger = new_session or user_turn"
                        .to_string(),
            });
        }
        Ok(())
    }
}

enum CustomPolicyRuntime {
    TargetSelector(TargetSelectorPolicy),
}

impl JudgePolicy for CustomPolicyRuntime {
    type Verdict = Value;

    fn to_classification(
        &self,
        verdict: Option<&Self::Verdict>,
        driver: &Driver,
    ) -> Result<Classification> {
        match self {
            Self::TargetSelector(policy) => policy.to_classification(verdict, driver),
        }
    }
}

/// Builds the affinity router a trigger calls for, if any.
fn affinity_router(
    trigger: ClassifyTrigger,
    message_hash_fallback: bool,
) -> Option<Arc<AffinityRouter>> {
    let router = match trigger {
        ClassifyTrigger::EveryRequest => return None,
        ClassifyTrigger::NewSession => AffinityRouter::new(),
        ClassifyTrigger::UserTurn => AffinityRouter::new().with_release_on_user_turn(),
    };
    let router = if message_hash_fallback {
        router.with_message_hash_fallback()
    } else {
        router
    };
    Some(Arc::new(router))
}

/// Routes requests through a capability, escalation, or custom classifier mode.
pub struct LlmTaskClassifier {
    route: FallThrough,
    /// Classifier used when this router is embedded in another cascade.
    inner: Arc<dyn Classifier>,
}

struct ClassifierRouteConfig {
    default_target: Category,
    classify_trigger: ClassifyTrigger,
    message_hash_fallback: bool,
}

/// Terminal classifier for a cascade whose classifiers may all abstain.
/// Closes a cascade with the first runtime model in `category`.
pub struct DefaultCategoryClassifier(pub Category);

#[async_trait]
impl Classifier for DefaultCategoryClassifier {
    async fn score(
        &self,
        _state: &mut State,
        _request: &mut Request,
        driver: &Driver,
    ) -> Result<(Classification, Option<Response>)> {
        let target = driver.first_model_for(&self.0)?;
        driver.set_evidence_if_empty(serde_json::json!({"source": "fall_open"}));
        Ok((
            Classification::Scores(vec![Score {
                target: target.clone(),
                confidence: 0.0,
                category: Some(self.0.clone()),
            }]),
            None,
        ))
    }
}

/// Complete construction settings for one LLM classifier mode.
#[derive(Clone)]
#[non_exhaustive]
pub enum LlmClassifierConfig {
    /// Routes between efficient and capable targets from a task-level verdict.
    Capability {
        /// Capability classifier settings.
        config: TaskClassifierConfig,
    },
    /// Judges efficient responses and escalates after a confirmed streak.
    Escalation {
        /// Prompt and verdict contract settings for the escalation judge.
        contract: ClassifierContractConfig,
        /// Escalation policy settings.
        config: EscalationJudgeConfig,
        /// Maximum completion tokens available to the escalation verdict.
        max_output_tokens: u64,
    },
    /// Routes among model categories using a user-supplied schema and policy.
    Custom {
        /// Category selected when the judge does not produce a usable verdict.
        default_target: Category,
        /// Custom classifier settings.
        config: CustomClassifierConfig,
    },
}

impl LlmTaskClassifier {
    /// Builds the classifier mode described by `config`.
    ///
    /// # Errors
    ///
    /// Returns an error when the selected mode's targets, contract, policy, or runtime
    /// settings are invalid.
    pub fn new(config: LlmClassifierConfig) -> Result<Self> {
        match config {
            LlmClassifierConfig::Capability { config } => Self::build_capability(config),
            LlmClassifierConfig::Escalation {
                contract,
                config,
                max_output_tokens,
            } => Self::build_escalation(contract, config, max_output_tokens),
            LlmClassifierConfig::Custom {
                default_target,
                config,
            } => Self::build_custom(default_target, config),
        }
    }

    fn build_capability(config: TaskClassifierConfig) -> Result<Self> {
        config.validate()?;
        let classify_trigger = config.classify_trigger;
        let message_hash_fallback = config.message_hash_fallback;
        let input = TaskInput {
            recent_turn_window: config.recent_turn_window,
        };
        let classifier: Arc<dyn Classifier> = match config.judge {
            CapabilityJudgeConfig::Llm(judge) => Arc::new(
                JudgeClassifier::new(
                    StructuredJudge::new(
                        input,
                        Self::load_capability_contract(&judge.contract)?,
                        SerdeDecoder::new(),
                        JudgeRuntimeConfig::new(judge.max_output_tokens)?,
                    ),
                    TaskClassifierPolicy::new(&judge),
                )
                .with_error_recovery(config.fail_open)
                .with_evidence(capability_evidence),
            ),
            CapabilityJudgeConfig::Decision(judge) => {
                Arc::new(DecisionClassifier::new(judge, input, config.fail_open))
            }
        };
        Self::from_classifier(
            classifier,
            ClassifierRouteConfig {
                default_target: Category::Capable,
                classify_trigger,
                message_hash_fallback,
            },
        )
    }

    fn build_custom(default_target: Category, config: CustomClassifierConfig) -> Result<Self> {
        Self::from_classifier(
            Self::custom_classifier(&config, false)?,
            ClassifierRouteConfig {
                default_target,
                classify_trigger: config.classify_trigger,
                message_hash_fallback: config.message_hash_fallback,
            },
        )
    }

    /// The judge face alone, for composite's own trigger and retained tier.
    pub(crate) fn custom_classifier(
        config: &CustomClassifierConfig,
        fail_open: bool,
    ) -> Result<Arc<dyn Classifier>> {
        config.validate()?;
        let contract =
            ClassifierContract::from_inner_schema(&config.prompt, config.response_schema.clone())?;
        let policy = match &config.policy {
            CustomClassifierPolicy::TargetSelector { selector } => {
                CustomPolicyRuntime::TargetSelector(TargetSelectorPolicy::new(selector)?)
            }
        };
        let classifier = JudgeClassifier::new(
            StructuredJudge::new(
                TaskInput {
                    recent_turn_window: config.recent_turn_window,
                },
                contract,
                JsonSchemaDecoder::new(),
                JudgeRuntimeConfig::new(config.max_output_tokens)?,
            ),
            policy,
        )
        .with_error_recovery(fail_open);
        let classifier = if fail_open {
            classifier.with_evidence(|_, verdict| {
                verdict.map(|_| serde_json::json!({"source": "llm-classifier"}))
            })
        } else {
            classifier
        };
        Ok(Arc::new(classifier))
    }

    fn build_escalation(
        contract_config: ClassifierContractConfig,
        config: EscalationJudgeConfig,
        max_output_tokens: u64,
    ) -> Result<Self> {
        let inner = escalation::build_classifier(contract_config, config, max_output_tokens)?;
        Ok(Self {
            route: FallThrough::new_with_state()
                .with_name(ALGORITHM_NAME)
                .with_classifier(Arc::clone(&inner)),
            inner,
        })
    }

    /// Loads the packaged capability-classifier contract.
    fn load_capability_contract(config: &ClassifierContractConfig) -> Result<ClassifierContract> {
        ClassifierContract::from_config(config, PROMPT_TEMPLATE, SCHEMA_TEMPLATE)
    }

    /// Keeps affinity and fallback ordering identical across judge-backed modes.
    fn from_classifier(inner: Arc<dyn Classifier>, config: ClassifierRouteConfig) -> Result<Self> {
        // Only `every_request` is rejected: it retains no target, so a fallback identity has nothing to key on. Both retaining triggers can key the retained target on a message hash when the caller sends no session id.
        if config.message_hash_fallback && config.classify_trigger == ClassifyTrigger::EveryRequest
        {
            return Err(LibsyError::AlgorithmError {
                message:
                    "message_hash_fallback requires classify_trigger = new_session or user_turn"
                        .to_string(),
            });
        }
        // Affinity comes first so a retained assignment short-circuits the judge call.
        let mut route = FallThrough::new_with_state().with_name(ALGORITHM_NAME);
        if let Some(affinity) =
            affinity_router(config.classify_trigger, config.message_hash_fallback).as_ref()
        {
            // Both roles must share one `Arc` so the classifier reads what the processor wrote.
            route = route
                .with_processor(affinity.clone())
                .with_classifier(affinity.clone());
        }
        let fallback = DefaultCategoryClassifier(config.default_target);
        Ok(Self {
            route: route
                .with_classifier(inner.clone())
                .with_classifier(Arc::new(fallback)),
            inner,
        })
    }
}

#[async_trait]
impl Classifier for LlmTaskClassifier {
    async fn score(
        &self,
        state: &mut State,
        request: &mut Request,
        driver: &Driver,
    ) -> Result<(Classification, Option<Response>)> {
        self.inner.score(state, request, driver).await
    }
}

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

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

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

    use parking_lot::Mutex;
    use serde_json::Value;

    use super::*;
    use switchyard_protocol::{
        ContentBlock, InstructionBlock, LlmClientError, LlmRequest, Metadata, ModelId, ToolCall,
        ToolResult, completion_text, text_request, text_response,
    };

    use crate::algorithms::util::llm_judge::Judge;
    use crate::core::testing::{Serve, test_drive_with_models};
    use switchyard_protocol::{LlmResponse, Response};

    const TEST_THRESHOLD: f64 = 0.5;

    type CapabilityJudge = StructuredJudge<TaskInput, SerdeDecoder<TaskClassifierVerdict>>;

    fn llm_config(base_threshold: f64) -> LlmCapabilityConfig {
        LlmCapabilityConfig {
            base_threshold,
            ..LlmCapabilityConfig::default()
        }
    }

    fn test_config(base_threshold: f64) -> TaskClassifierConfig {
        TaskClassifierConfig {
            judge: CapabilityJudgeConfig::Llm(llm_config(base_threshold)),
            ..TaskClassifierConfig::default()
        }
    }

    fn policy() -> TaskClassifierPolicy {
        TaskClassifierPolicy::new(&llm_config(TEST_THRESHOLD))
    }

    fn runtime_models() -> HashMap<Category, Vec<ModelId>> {
        [
            (Category::Judge, vec![ModelId::from("judge")]),
            (Category::Efficient, vec![ModelId::from("efficient")]),
            (Category::Capable, vec![ModelId::from("capable")]),
            (
                Category::Any,
                vec![ModelId::from("efficient"), ModelId::from("capable")],
            ),
        ]
        .into()
    }

    fn policy_driver() -> Driver {
        Driver::new("test", Arc::new(runtime_models().into())).0
    }

    fn verdict(
        p_solve: f64,
        capability_boundary: &str,
        primary_rule: &str,
    ) -> TaskClassifierVerdict {
        TaskClassifierVerdict {
            crux: "test crux".to_string(),
            primary_rule: primary_rule.to_string(),
            capability_boundary: capability_boundary.to_string(),
            p_solve,
        }
    }

    fn selected(
        policy: &TaskClassifierPolicy,
        verdict: Option<&TaskClassifierVerdict>,
    ) -> Result<ModelId> {
        policy
            .to_classification(verdict, &policy_driver())?
            .argmax(false)?
            .map(|score| score.target)
            .ok_or_else(|| LibsyError::AlgorithmError {
                message: "policy abstained".to_string(),
            })
    }

    /// Records what each target received; answers the judge with a supported verdict and
    /// every other target with a plain completion.
    #[derive(Default)]
    struct Recorder {
        calls: Mutex<Vec<String>>,
        call_roles: Mutex<Vec<(String, bool)>>,
        judge_max_output_tokens: Mutex<Vec<Option<u64>>>,
        judge_system_prompts: Mutex<Vec<String>>,
    }

    impl Recorder {
        fn calls(&self) -> Vec<String> {
            self.calls.lock().clone()
        }

        fn call_roles(&self) -> Vec<(String, bool)> {
            self.call_roles.lock().clone()
        }

        fn judge_max_output_tokens(&self) -> Vec<Option<u64>> {
            self.judge_max_output_tokens.lock().clone()
        }

        fn judge_system_prompts(&self) -> Vec<String> {
            self.judge_system_prompts.lock().clone()
        }

        fn serve(self: &Arc<Self>) -> impl Serve {
            let recorder = Arc::clone(self);
            move |model: ModelId, request: Request| {
                let recorder = Arc::clone(&recorder);
                async move {
                    let model = model.to_string();
                    recorder.calls.lock().push(model.clone());
                    recorder
                        .call_roles
                        .lock()
                        .push((model.clone(), model != "judge"));
                    let completion = if model == "judge" {
                        recorder
                            .judge_max_output_tokens
                            .lock()
                            .push(request.llm_request.output.max_output_tokens);
                        recorder.judge_system_prompts.lock().extend(
                            request
                                .llm_request
                                .instructions
                                .first()
                                .and_then(|instruction| {
                                    instruction.content.iter().find_map(|b| {
                                        if let ContentBlock::Text { text } = b {
                                            Some(text.clone())
                                        } else {
                                            None
                                        }
                                    })
                                }),
                        );
                        r#"{"crux":"bounded task","primary_rule":"SUP-1","capability_boundary":"supported","p_solve":0.9}"#.to_string()
                    } else {
                        format!("answer from {model}")
                    };
                    Ok(Response {
                        llm_response: LlmResponse::Agg(text_response(None, completion)),
                        metadata: request.metadata,
                        upstream_headers: http::HeaderMap::new(),
                    })
                }
            }
        }
    }

    /// The judge times out; every other target answers normally.
    fn unreachable_judge() -> impl Serve {
        |model: ModelId, request: Request| async move {
            let model = model.to_string();
            if model == "judge" {
                return Err(LlmClientError::Timeout {
                    source: Box::new(std::io::Error::other("judge unreachable")),
                });
            }
            Ok(Response {
                llm_response: LlmResponse::Agg(text_response(None, format!("answer from {model}"))),
                metadata: request.metadata,
                upstream_headers: http::HeaderMap::new(),
            })
        }
    }

    fn router() -> Result<Arc<LlmTaskClassifier>> {
        Ok(Arc::new(LlmTaskClassifier::new(
            LlmClassifierConfig::Capability {
                config: test_config(TEST_THRESHOLD),
            },
        )?))
    }

    fn classify_request() -> Request {
        Request {
            llm_request: text_request(Some("auto".to_string()), "classify this task"),
            raw_request: None,
            metadata: None,
        }
    }

    fn classify_session_request() -> Request {
        Request {
            metadata: Some(Metadata {
                session_id: Some("session-1".to_string()),
                ..Metadata::default()
            }),
            ..classify_request()
        }
    }

    fn classify_follow_up_request() -> Request {
        let mut request = classify_request();
        request
            .llm_request
            .messages
            .push(Message::text(Role::Assistant, "I will add the test."));
        request.llm_request.messages.push(Message::text(
            Role::User,
            "Now run the test suite and report the result.",
        ));
        request
    }

    #[tokio::test]
    async fn an_unreachable_judge_routes_capable_instead_of_failing_the_request() -> Result<()> {
        let router = router()?;

        let (selected_model, response) = test_drive_with_models(
            router,
            classify_request(),
            runtime_models(),
            unreachable_judge(),
        )
        .await?;

        assert_eq!(selected_model, "capable");
        assert_eq!(
            response.llm_response.as_agg().map(completion_text),
            Some("answer from capable".to_string())
        );
        Ok(())
    }

    #[tokio::test]
    async fn classifier_judges_each_request_without_affinity() -> Result<()> {
        let recorder = Arc::new(Recorder::default());
        let router = router()?;
        let request = classify_request();
        let models = runtime_models();

        test_drive_with_models(
            router.clone(),
            request.clone(),
            models.clone(),
            recorder.serve(),
        )
        .await?;
        test_drive_with_models(router, request, models, recorder.serve()).await?;

        assert_eq!(
            recorder.calls(),
            vec!["judge", "efficient", "judge", "efficient"]
        );
        assert_eq!(
            recorder.call_roles(),
            vec![
                ("judge".to_string(), false),
                ("efficient".to_string(), true),
                ("judge".to_string(), false),
                ("efficient".to_string(), true),
            ]
        );
        Ok(())
    }

    #[tokio::test]
    async fn classifier_config_sets_the_judge_completion_cap() -> Result<()> {
        let recorder = Arc::new(Recorder::default());
        let router = Arc::new(LlmTaskClassifier::new(LlmClassifierConfig::Capability {
            config: TaskClassifierConfig {
                judge: CapabilityJudgeConfig::Llm(LlmCapabilityConfig {
                    max_output_tokens: 512,
                    ..llm_config(TEST_THRESHOLD)
                }),
                ..test_config(TEST_THRESHOLD)
            },
        })?);

        test_drive_with_models(
            router,
            classify_request(),
            runtime_models(),
            recorder.serve(),
        )
        .await?;

        assert_eq!(recorder.judge_max_output_tokens(), vec![Some(512)]);
        Ok(())
    }

    #[tokio::test]
    async fn classifier_config_overrides_the_packaged_prompt() -> Result<()> {
        let recorder = Arc::new(Recorder::default());
        let router = Arc::new(LlmTaskClassifier::new(LlmClassifierConfig::Capability {
            config: TaskClassifierConfig {
                judge: CapabilityJudgeConfig::Llm(LlmCapabilityConfig {
                    contract: ClassifierContractConfig::default()
                        .with_prompt("Custom capability rubric."),
                    ..llm_config(TEST_THRESHOLD)
                }),
                ..test_config(TEST_THRESHOLD)
            },
        })?);

        test_drive_with_models(
            router,
            classify_request(),
            runtime_models(),
            recorder.serve(),
        )
        .await?;

        let prompts = recorder.judge_system_prompts();
        assert_eq!(prompts.len(), 1);
        assert_eq!(prompts[0], "Custom capability rubric.");
        Ok(())
    }

    #[tokio::test]
    async fn classifier_config_enables_new_session_trigger() -> Result<()> {
        let recorder = Arc::new(Recorder::default());
        let router = Arc::new(LlmTaskClassifier::new(LlmClassifierConfig::Capability {
            config: TaskClassifierConfig {
                classify_trigger: ClassifyTrigger::NewSession,
                ..test_config(TEST_THRESHOLD)
            },
        })?);

        let request = classify_session_request();
        let models = runtime_models();
        test_drive_with_models(
            router.clone(),
            request.clone(),
            models.clone(),
            recorder.serve(),
        )
        .await?;
        test_drive_with_models(router, request, models, recorder.serve()).await?;

        assert_eq!(recorder.calls(), vec!["judge", "efficient", "efficient"]);
        Ok(())
    }

    #[tokio::test]
    async fn classifier_config_reuses_message_hash_affinity_for_a_follow_up() -> Result<()> {
        let recorder = Arc::new(Recorder::default());
        let router = Arc::new(LlmTaskClassifier::new(LlmClassifierConfig::Capability {
            config: TaskClassifierConfig {
                classify_trigger: ClassifyTrigger::NewSession,
                message_hash_fallback: true,
                recent_turn_window: None,
                ..test_config(TEST_THRESHOLD)
            },
        })?);

        let models = runtime_models();
        test_drive_with_models(
            router.clone(),
            classify_request(),
            models.clone(),
            recorder.serve(),
        )
        .await?;
        test_drive_with_models(
            router,
            classify_follow_up_request(),
            models,
            recorder.serve(),
        )
        .await?;

        assert_eq!(recorder.calls(), vec!["judge", "efficient", "efficient"]);
        Ok(())
    }

    #[tokio::test]
    async fn one_classifier_uses_each_requests_runtime_models() -> Result<()> {
        let router = Arc::new(LlmTaskClassifier::new(LlmClassifierConfig::Capability {
            config: TaskClassifierConfig {
                classify_trigger: ClassifyTrigger::NewSession,
                ..test_config(TEST_THRESHOLD)
            },
        })?);
        let calls = Arc::new(Mutex::new(Vec::new()));
        let serve = |calls: Arc<Mutex<Vec<String>>>| {
            move |model: ModelId, _request: Request| {
                let calls = Arc::clone(&calls);
                async move {
                    calls.lock().push(model.to_string());
                    let text = if model.as_str().starts_with("judge-") {
                        r#"{"crux":"bounded task","primary_rule":"SUP-1","capability_boundary":"supported","p_solve":0.9}"#.to_string()
                    } else {
                        model.to_string()
                    };
                    Ok(Response {
                        llm_response: LlmResponse::Agg(text_response(None, text)),
                        metadata: None,
                        upstream_headers: Default::default(),
                    })
                }
            }
        };
        let models = |suffix: &str| -> HashMap<Category, Vec<ModelId>> {
            [
                (
                    Category::Judge,
                    vec![ModelId::from(format!("judge-{suffix}"))],
                ),
                (
                    Category::Efficient,
                    vec![ModelId::from(format!("efficient-{suffix}"))],
                ),
                (
                    Category::Capable,
                    vec![ModelId::from(format!("capable-{suffix}"))],
                ),
                (
                    Category::Any,
                    vec![
                        ModelId::from(format!("efficient-{suffix}")),
                        ModelId::from(format!("capable-{suffix}")),
                    ],
                ),
            ]
            .into()
        };
        let request = classify_session_request();

        let (first, _) = test_drive_with_models(
            router.clone(),
            request.clone(),
            models("a"),
            serve(Arc::clone(&calls)),
        )
        .await?;
        let (second, _) =
            test_drive_with_models(router, request, models("b"), serve(Arc::clone(&calls))).await?;

        assert_eq!(first, "efficient-a");
        assert_eq!(second, "efficient-b");
        assert_eq!(
            &*calls.lock(),
            &["judge-a", "efficient-a", "judge-b", "efficient-b"]
        );
        Ok(())
    }

    #[tokio::test]
    async fn decision_judge_routes_on_ordered_policy() -> Result<()> {
        use crate::{Call, RuntimeModels, drive};
        use serde_json::json;
        use std::collections::BTreeMap;
        use std::sync::atomic::{AtomicUsize, Ordering};
        use switchyard_protocol::{
            BooleanEstimate, DecisionAnswer, DecisionResponse, DecisionValue, Probability,
            ProviderConfidence,
        };

        let settings = DecisionJudgeConfig {
            policy: DecisionPolicy {
                rules: vec![DecisionRule {
                    candidate: "a".into(),
                    cutoff: 0.4,
                }],
                fallback: "b".into(),
            },
            instructions: None,
            candidates: BTreeMap::from([
                ("a".into(), "capable".into()),
                ("b".into(), "efficient".into()),
            ]),
            evidence: json!({
                "candidate_descriptions": {"a": "Profile A", "b": "Profile B", "c": "Profile C"},
                "reference_cases": [{"task": "Independent example", "observed_solved": {"a": true, "b": null, "c": false}}],
                "summaries": [],
            }),
        };
        let ordered_policy = DecisionPolicy {
            rules: vec![
                DecisionRule {
                    candidate: "a".into(),
                    cutoff: 0.6,
                },
                DecisionRule {
                    candidate: "c".into(),
                    cutoff: 0.2,
                },
            ],
            fallback: "b".into(),
        };
        let mut request = classify_session_request();
        request.llm_request.messages.push(Message {
            role: Role::Assistant,
            content: vec![
                ContentBlock::Reasoning {
                    text: "private reasoning".into(),
                    signature: None,
                    details: vec![],
                },
                ContentBlock::Text {
                    text: "visible progress".into(),
                },
            ],
        });
        request
            .llm_request
            .messages
            .push(Message::text(Role::User, "latest follow-up"));
        let original_messages = request.llm_request.messages.clone();
        let expected_task = json!([
            original_messages[0].clone(),
            Message::text(Role::Assistant, "visible progress"),
            Message::text(Role::User, "latest follow-up"),
        ]);

        for (name, score, fail_open, expected) in [
            ("above", 0.45, true, Some("capable")),
            ("equal", 0.4, true, Some("efficient")),
            ("below", 0.39, false, Some("efficient")),
            ("missing answer", 0.0, true, Some("efficient")),
            ("no distribution", 0.0, true, Some("efficient")),
            ("wrong type", 0.0, true, Some("efficient")),
            ("out of range", 1.1, true, Some("efficient")),
            ("nonfinite", f64::NAN, false, Some("efficient")),
            ("provider error", 0.0, true, Some("efficient")),
            ("provider error", 0.0, false, None),
            ("dropped reply", 0.0, true, Some("efficient")),
            ("ordered first", 0.65, true, Some("capable")),
            ("ordered second", 0.6, false, Some("third")),
            ("ordered fallback", 0.6, false, Some("efficient")),
            ("ordered no distribution", 0.0, true, Some("efficient")),
            ("ordered provider error", 0.0, true, Some("efficient")),
            ("ordered provider error", 0.0, false, None),
            ("ordered invalid", f64::NAN, true, Some("efficient")),
        ] {
            let mut config = settings.clone();
            let ordered = name.starts_with("ordered");
            if ordered {
                config.policy = ordered_policy.clone();
                config.candidates.insert("c".into(), "third".into());
            }
            if name == "below" {
                config.instructions =
                    Some(json!({"policy": "Choose a candidate using the supplied evidence."}));
            }
            let override_instructions = config.instructions.clone();
            let router: Arc<dyn Algorithm> =
                Arc::new(LlmTaskClassifier::new(LlmClassifierConfig::Capability {
                    config: TaskClassifierConfig {
                        judge: CapabilityJudgeConfig::Decision(config),
                        fail_open,
                        recent_turn_window: Some(3),
                        classify_trigger: ClassifyTrigger::NewSession,
                        ..TaskClassifierConfig::default()
                    },
                })?);
            let calls = AtomicUsize::new(0);
            let serve = |call| {
                let calls = &calls;
                let expected_task = &expected_task;
                let settings = &settings;
                let override_instructions = &override_instructions;
                async move {
                    let Call::Decision(call) = call else {
                        panic!("judge must use the decision step")
                    };
                    calls.fetch_add(1, Ordering::SeqCst);
                    assert_eq!(call.model, "judge");
                    assert_eq!(call.request.model, Some("judge".into()));
                    let labels = if ordered {
                        vec!["a", "b", "c"]
                    } else {
                        vec!["a", "b"]
                    };
                    let context = json!({
                        "task": expected_task,
                        "candidates": labels,
                        "evidence": settings.evidence,
                    });
                    assert_eq!(call.request.context, context);
                    assert_eq!(call.request.questions.len(), 1);
                    let question = &call.request.questions["route"];
                    let switchyard_protocol::DecisionKind::Choice { options } = &question.kind
                    else {
                        panic!("expected a choice")
                    };
                    assert_eq!(
                        options
                            .iter()
                            .map(|option| option.id.as_str())
                            .collect::<Vec<_>>(),
                        labels
                    );
                    if let Some(instructions) = override_instructions {
                        assert_eq!(&question.instructions, instructions);
                    } else {
                        assert!(question.instructions.as_str().is_some());
                    }
                    if name.ends_with("provider error") {
                        return call
                            .respond(std::future::ready(Err(LibsyError::client_call(
                                "judge",
                                LlmClientError::General("private provider body".into()),
                            ))))
                            .await;
                    }
                    if name == "dropped reply" {
                        drop(call);
                        return Ok(());
                    }
                    let value = if name == "wrong type" {
                        DecisionValue::Boolean(BooleanEstimate::Value(true))
                    } else {
                        DecisionValue::Choice {
                            selected: "b".into(),
                            probabilities: (!name.ends_with("no distribution")).then(|| {
                                if ordered {
                                    let third = if name == "ordered fallback" { 0.2 } else { 0.3 };
                                    return BTreeMap::from([
                                        ("a".into(), Probability(score)),
                                        ("b".into(), Probability(1.0 - score - third)),
                                        ("c".into(), Probability(third)),
                                    ]);
                                }
                                BTreeMap::from([
                                    ("a".into(), Probability(score)),
                                    ("b".into(), Probability(1.0 - score)),
                                ])
                            }),
                        }
                    };
                    let answers = if name == "missing answer" {
                        BTreeMap::new()
                    } else {
                        BTreeMap::from([(
                            "route".into(),
                            DecisionAnswer {
                                value,
                                provider_confidence: Some(ProviderConfidence(0.99)),
                            },
                        )])
                    };
                    call.respond(std::future::ready(Ok(DecisionResponse {
                        id: None,
                        model: Some("provider-judge".into()),
                        answers,
                        usage: Default::default(),
                    })))
                    .await
                }
            };
            let mut runtime = runtime_models();
            if ordered {
                runtime.remove(&Category::Capable);
                runtime.remove(&Category::Efficient);
                runtime
                    .get_mut(&Category::Any)
                    .unwrap()
                    .push("third".into());
            }
            let models = Arc::new(RuntimeModels::new(runtime));
            let result = drive(router.clone(), request.clone(), models.clone(), &serve).await;
            assert_eq!(calls.load(Ordering::SeqCst), 1, "{name}");
            let Some(expected) = expected else {
                assert!(
                    matches!(result, Err(LibsyError::ClientCall { .. })),
                    "{name}"
                );
                continue;
            };
            let outcome = result?;
            assert_eq!(outcome.selected_model_id()?, expected, "{name}");
            assert!(outcome.response.is_none());
            assert_eq!(outcome.request.llm_request.messages, original_messages);
            let evidence = outcome
                .metadata
                .and_then(|metadata| metadata.evidence)
                .expect("routing evidence");
            if matches!(
                name,
                "above"
                    | "equal"
                    | "below"
                    | "ordered first"
                    | "ordered second"
                    | "ordered fallback"
            ) {
                assert_eq!(evidence["verdict"], "ordered_cutoffs");
                assert_eq!(evidence["source"], "decision_classifier");
                if name == "above" {
                    assert_eq!(evidence["threshold"], settings.policy.rules[0].cutoff);
                    assert_eq!(evidence["target"], "a");
                }
            } else {
                let reason = match name {
                    name if name.ends_with("provider error") => "client_error",
                    "dropped reply" => "call_error",
                    _ => "invalid_verdict",
                };
                assert_eq!(
                    evidence,
                    json!({"source": "fail_open", "reason_code": reason})
                );
            }
            let retained = drive(router, request.clone(), models, &serve).await?;
            assert_eq!(retained.selected_model_id()?, expected);
            assert_eq!(
                calls.load(Ordering::SeqCst),
                1,
                "affinity should skip the judge: {name}"
            );
        }
        for cutoff in [-0.1, 1.1, f64::NAN] {
            let mut config = settings.clone();
            config.policy.rules[0].cutoff = cutoff;
            assert!(
                LlmTaskClassifier::new(LlmClassifierConfig::Capability {
                    config: TaskClassifierConfig {
                        judge: CapabilityJudgeConfig::Decision(config),
                        ..TaskClassifierConfig::default()
                    },
                })
                .is_err()
            );
        }
        for field in ["candidate", "fallback", "rules"] {
            let mut config = settings.clone();
            let policy = &mut config.policy;
            match field {
                "candidate" => policy.rules[0].candidate = "missing".into(),
                "fallback" => policy.fallback = "missing".into(),
                _ => policy.rules.clear(),
            }
            assert!(config.validate().is_err());
        }
        let mut duplicate = settings.clone();
        duplicate
            .candidates
            .insert("duplicate-a".into(), "capable".into());
        assert!(duplicate.validate().is_err());
        Ok(())
    }

    #[test]
    fn the_threshold_boundary_is_inclusive() -> Result<()> {
        let policy = policy();
        let at_threshold = verdict(0.5, "supported", "SUP-1");
        let below_threshold = verdict(0.49, "supported", "SUP-1");
        assert_eq!(selected(&policy, Some(&at_threshold))?, "efficient");
        assert_eq!(selected(&policy, Some(&below_threshold))?, "capable");
        Ok(())
    }

    #[test]
    fn the_threshold_moves_the_routing_boundary() -> Result<()> {
        let borderline = verdict(0.5, "supported", "SUP-1");
        let strict = TaskClassifierPolicy::new(&llm_config(0.9));
        let lenient = TaskClassifierPolicy::new(&llm_config(0.1));
        assert_eq!(selected(&strict, Some(&borderline))?, "capable");
        assert_eq!(selected(&lenient, Some(&borderline))?, "efficient");
        Ok(())
    }

    #[test]
    fn classifier_config_rejects_unknown_fields() {
        let error = serde_json::from_value::<TaskClassifierConfig>(serde_json::json!({
            "base_threshold": 0.5,
            "classifier_magic": true,
        }))
        .expect_err("unknown classifier fields must be rejected");

        assert!(
            error
                .to_string()
                .contains("unknown field `classifier_magic`"),
            "{error}"
        );
    }

    #[test]
    fn invalid_classifier_config_is_rejected() -> Result<()> {
        for bad in [1.5, -0.1, f64::NAN, f64::INFINITY] {
            assert!(
                LlmTaskClassifier::new(LlmClassifierConfig::Capability {
                    config: test_config(bad),
                })
                .is_err(),
                "base threshold {bad} should be rejected"
            );
        }
        for config in [
            TaskClassifierConfig {
                judge: CapabilityJudgeConfig::Llm(LlmCapabilityConfig {
                    base_threshold: 0.5,
                    threshold_step: -0.1,
                    ..LlmCapabilityConfig::default()
                }),
                ..TaskClassifierConfig::default()
            },
            TaskClassifierConfig {
                judge: CapabilityJudgeConfig::Llm(LlmCapabilityConfig {
                    base_threshold: 0.8,
                    threshold_step: 0.11,
                    ..LlmCapabilityConfig::default()
                }),
                ..TaskClassifierConfig::default()
            },
            TaskClassifierConfig {
                judge: CapabilityJudgeConfig::Llm(LlmCapabilityConfig {
                    base_threshold: 0.5,
                    ..LlmCapabilityConfig::default()
                }),
                message_hash_fallback: true,
                ..TaskClassifierConfig::default()
            },
            TaskClassifierConfig {
                judge: CapabilityJudgeConfig::Llm(LlmCapabilityConfig {
                    base_threshold: 0.5,
                    max_output_tokens: 0,
                    ..LlmCapabilityConfig::default()
                }),
                ..TaskClassifierConfig::default()
            },
        ] {
            assert!(LlmTaskClassifier::new(LlmClassifierConfig::Capability { config }).is_err());
        }
        for base_threshold in [0.0, 1.0] {
            LlmTaskClassifier::new(LlmClassifierConfig::Capability {
                config: test_config(base_threshold),
            })?;
        }
        Ok(())
    }

    #[test]
    fn message_hash_fallback_accepts_retaining_triggers() -> Result<()> {
        // Only `every_request` is rejected: it retains no target, so a fallback
        // identity has nothing to key on. Both retaining triggers may key the
        // retained target on a message hash when the caller sends no session id.
        for trigger in [ClassifyTrigger::NewSession, ClassifyTrigger::UserTurn] {
            let config = TaskClassifierConfig {
                judge: CapabilityJudgeConfig::Llm(LlmCapabilityConfig {
                    base_threshold: 0.5,
                    ..LlmCapabilityConfig::default()
                }),
                classify_trigger: trigger,
                message_hash_fallback: true,
                ..TaskClassifierConfig::default()
            };
            LlmTaskClassifier::new(LlmClassifierConfig::Capability { config }).map_err(
                |error| LibsyError::AlgorithmError {
                    message: format!("{trigger:?} with message_hash_fallback rejected: {error}"),
                },
            )?;
        }
        let every_request = TaskClassifierConfig {
            judge: CapabilityJudgeConfig::Llm(LlmCapabilityConfig {
                base_threshold: 0.5,
                ..LlmCapabilityConfig::default()
            }),
            classify_trigger: ClassifyTrigger::EveryRequest,
            message_hash_fallback: true,
            ..TaskClassifierConfig::default()
        };
        assert!(
            LlmTaskClassifier::new(LlmClassifierConfig::Capability {
                config: every_request
            })
            .is_err(),
            "every_request with message_hash_fallback should stay rejected"
        );
        Ok(())
    }

    #[test]
    fn an_unusable_verdict_is_ambiguous() -> Result<()> {
        let policy = policy();
        let inconsistent_rule = TaskClassifierVerdict {
            capability_boundary: "uncertain".to_string(),
            ..verdict(1.0, "supported", "SUP-1")
        };
        let empty_crux = TaskClassifierVerdict {
            crux: "  ".to_string(),
            ..verdict(1.0, "supported", "SUP-1")
        };
        let unusable = [
            Some(verdict(1.1, "supported", "SUP-1")),
            Some(inconsistent_rule),
            Some(empty_crux),
            None,
        ];
        for verdict in unusable {
            let classification = policy.to_classification(verdict.as_ref(), &policy_driver())?;
            assert!(matches!(classification, Classification::Ambiguous(_)));
            assert!(classification.argmax(false)?.is_none());
            assert!(classification.argmax(true)?.is_none());
        }
        Ok(())
    }

    #[test]
    fn capability_boundaries_apply_monotonic_threshold_steps() -> Result<()> {
        let policy = TaskClassifierPolicy::new(&LlmCapabilityConfig {
            threshold_step: 0.1,
            ..llm_config(0.4)
        });

        assert_eq!(
            selected(&policy, Some(&verdict(0.4, "supported", "SUP-2")))?,
            "efficient"
        );
        assert_eq!(
            selected(&policy, Some(&verdict(0.49, "uncertain", "UNC-1")))?,
            "capable"
        );
        assert_eq!(
            selected(&policy, Some(&verdict(0.5, "uncertain", "UNC-1")))?,
            "efficient"
        );
        assert_eq!(
            selected(&policy, Some(&verdict(0.5, "unmatched", "none")))?,
            "efficient"
        );
        assert_eq!(
            selected(&policy, Some(&verdict(0.59, "unsupported", "LIM-1")))?,
            "capable"
        );
        assert_eq!(
            selected(&policy, Some(&verdict(0.6, "unsupported", "LIM-1")))?,
            "efficient"
        );
        Ok(())
    }

    /// The text of each message a judge with `recent_turn_window` would be sent.
    /// The no-window case is covered by `capability_judge_builds_a_structured_request`.
    fn capability_judge(recent_turn_window: Option<usize>) -> Result<CapabilityJudge> {
        Ok(StructuredJudge::new(
            TaskInput { recent_turn_window },
            LlmTaskClassifier::load_capability_contract(&ClassifierContractConfig::default())?,
            SerdeDecoder::new(),
            JudgeRuntimeConfig::new(DEFAULT_JUDGE_MAX_OUTPUT_TOKENS)?,
        ))
    }

    fn judged_contents(recent_turn_window: usize) -> Result<Vec<String>> {
        let judge = capability_judge(Some(recent_turn_window))?;
        let request = Request {
            llm_request: LlmRequest {
                messages: vec![
                    Message::text(Role::System, "client instructions"),
                    Message::text(Role::User, "initial task"),
                    Message::text(Role::Assistant, "old response"),
                    Message::text(Role::User, "old follow-up"),
                    Message::text(Role::Assistant, "recent 1"),
                    Message::text(Role::User, "recent 2"),
                ],
                ..LlmRequest::default()
            },
            raw_request: None,
            metadata: None,
        };
        Ok(judge
            .build_request(&State::default(), &request)
            .llm_request
            .messages
            .iter()
            .filter_map(|message| message.text_content("\n"))
            .collect())
    }

    #[test]
    fn a_window_widens_the_judge_to_the_surrounding_conversation() -> Result<()> {
        // Client instructions and the opening task, plus the last two turns.
        let contents = judged_contents(2)?;
        assert!(contents.contains(&"client instructions".to_string()));
        assert!(contents.contains(&"initial task".to_string()));
        assert!(contents.contains(&"recent 1".to_string()));
        assert!(contents.contains(&"recent 2".to_string()));
        assert!(!contents.contains(&"old response".to_string()));
        Ok(())
    }

    #[test]
    fn a_zero_window_keeps_only_the_instructions_and_the_task() -> Result<()> {
        let contents = judged_contents(0)?;
        assert!(contents.contains(&"client instructions".to_string()));
        assert!(contents.contains(&"initial task".to_string()));
        assert!(!contents.contains(&"recent 2".to_string()));
        Ok(())
    }

    fn tool_call(id: &str) -> Message {
        Message {
            role: Role::Assistant,
            content: vec![ContentBlock::ToolCall(ToolCall {
                id: id.to_string(),
                name: "search".to_string(),
                arguments: Value::Null,
            })],
        }
    }

    fn tool_result(id: &str) -> Message {
        Message {
            role: Role::Tool,
            content: vec![ContentBlock::ToolResult(ToolResult {
                tool_call_id: id.to_string(),
                content: vec![ContentBlock::Text {
                    text: "tool output".to_string(),
                }],
                is_error: None,
            })],
        }
    }

    #[test]
    fn default_task_input_keeps_user_content_around_tool_results() {
        let mut result = tool_result("call-1");
        result.role = Role::User;
        let mut mixed = result.clone();
        mixed.content.push(ContentBlock::Text {
            text: "latest follow-up".to_string(),
        });
        let input = TaskInput {
            recent_turn_window: None,
        };
        let mut request = Request {
            llm_request: LlmRequest {
                messages: vec![
                    result.clone(),
                    Message::text(Role::User, "initial task"),
                    tool_call("call-1"),
                    mixed,
                    result.clone(),
                ],
                ..LlmRequest::default()
            },
            ..Request::default()
        };
        assert_eq!(
            input.build_messages(&State::default(), &request),
            vec![
                Message::text(Role::User, "initial task"),
                Message::text(Role::User, "latest follow-up"),
            ]
        );
        request.llm_request.messages = vec![result];
        assert!(input.build_messages(&State::default(), &request).is_empty());
    }

    /// A count-based window can begin on a tool result, which leaves the call that
    /// introduced its id outside the window and the classifier history invalid.
    #[test]
    fn trimming_keeps_the_call_that_introduced_a_kept_tool_result() {
        let messages = vec![
            Message::text(Role::System, "client instructions"),
            Message::text(Role::User, "initial task"),
            Message::text(Role::Assistant, "old response"),
            tool_call("call-1"),
            tool_result("call-1"),
            Message::text(Role::Assistant, "recent 1"),
            Message::text(Role::User, "recent 2"),
            Message::text(Role::Assistant, "recent 3"),
            Message::text(Role::User, "recent 4"),
        ];

        // The five-message tail begins exactly on the tool result.
        let kept = trim_messages(&messages, 5);

        assert_eq!(
            kept,
            vec![
                Message::text(Role::System, "client instructions"),
                Message::text(Role::User, "initial task"),
                tool_call("call-1"),
                tool_result("call-1"),
                Message::text(Role::Assistant, "recent 1"),
                Message::text(Role::User, "recent 2"),
                Message::text(Role::Assistant, "recent 3"),
                Message::text(Role::User, "recent 4"),
            ]
        );
    }

    /// Ids repeat across a conversation, so a later call must not stand in for the one that
    /// answers an earlier result.
    #[test]
    fn trimming_pairs_a_repeated_id_with_the_call_that_precedes_it() {
        let messages = vec![
            Message::text(Role::System, "client instructions"),
            Message::text(Role::User, "initial task"),
            tool_call("x"),
            tool_result("x"),
            Message::text(Role::Assistant, "later"),
            tool_call("x"),
            tool_result("x"),
        ];

        // The four-message tail begins on the first result, whose own call sits one earlier.
        let kept = trim_messages(&messages, 4);

        assert_eq!(
            kept,
            vec![
                Message::text(Role::System, "client instructions"),
                Message::text(Role::User, "initial task"),
                tool_call("x"),
                tool_result("x"),
                Message::text(Role::Assistant, "later"),
                tool_call("x"),
                tool_result("x"),
            ]
        );
    }

    /// A result whose call precedes the opening task can never be paired, because trimming
    /// never reaches behind the task. The window must not widen hunting for it.
    #[test]
    fn trimming_keeps_the_counted_window_when_a_result_cannot_be_paired() {
        let messages = vec![
            Message::text(Role::System, "client instructions"),
            tool_call("orphan"),
            Message::text(Role::User, "initial task"),
            Message::text(Role::Assistant, "old response"),
            tool_result("orphan"),
            Message::text(Role::Assistant, "recent 1"),
            Message::text(Role::User, "recent 2"),
        ];

        let kept = trim_messages(&messages, 3);

        assert_eq!(
            kept,
            vec![
                Message::text(Role::System, "client instructions"),
                Message::text(Role::User, "initial task"),
                tool_result("orphan"),
                Message::text(Role::Assistant, "recent 1"),
                Message::text(Role::User, "recent 2"),
            ]
        );
    }

    #[test]
    fn a_window_restates_the_routing_instruction_last() -> Result<()> {
        let contents = judged_contents(2)?;

        // Last, not merely present: the instruction only works in final position,
        // after the conversation content it is meant to outrank.
        assert_eq!(
            contents.last().map(String::as_str),
            Some(TRAILING_ROUTING_INSTRUCTION)
        );
        Ok(())
    }

    /// Reasoning is provider-private and some upstreams reject it replayed without the
    /// opaque state it was issued with, so it must not reach a windowed classifier
    /// request. Visible assistant text and complete tool pairs still do, and a turn
    /// whose only content was reasoning is dropped rather than left behind empty.
    #[test]
    fn a_window_drops_reasoning_but_keeps_visible_text_and_tool_pairs() {
        let messages = vec![
            Message::text(Role::System, "client instructions"),
            Message::text(Role::User, "initial task"),
            Message {
                role: Role::Assistant,
                content: vec![
                    ContentBlock::Reasoning {
                        text: "private chain of thought".to_string(),
                        signature: None,
                        details: Vec::new(),
                    },
                    ContentBlock::Text {
                        text: "visible answer".to_string(),
                    },
                ],
            },
            tool_call("call-1"),
            tool_result("call-1"),
            Message {
                role: Role::Assistant,
                content: vec![ContentBlock::Reasoning {
                    text: "reasoning-only turn".to_string(),
                    signature: None,
                    details: Vec::new(),
                }],
            },
            Message::text(Role::User, "follow-up"),
        ];
        let request = Request {
            llm_request: LlmRequest {
                messages,
                ..LlmRequest::default()
            },
            raw_request: None,
            metadata: None,
        };

        let built = TaskInput {
            recent_turn_window: Some(10),
        }
        .build_messages(&State::default(), &request);

        assert!(
            !built
                .iter()
                .flat_map(|message| &message.content)
                .any(|block| matches!(block, ContentBlock::Reasoning { .. })),
            "{built:?}"
        );
        assert!(
            built
                .iter()
                .any(|message| message.text_content("\n").as_deref() == Some("visible answer"))
        );
        assert!(built.contains(&tool_call("call-1")));
        assert!(built.contains(&tool_result("call-1")));
        // Six of the seven fixtures survive — the reasoning-only turn is gone entirely
        // rather than left behind empty — plus the trailing routing instruction.
        assert_eq!(built.len(), 7);
        assert!(built.iter().all(|message| !message.content.is_empty()));
    }

    #[test]
    fn the_default_path_is_left_unchanged() -> Result<()> {
        // No window means no conversation to be distracted by, so the default
        // request shape stays exactly as it was.
        let judge = capability_judge(None)?;
        let request = Request {
            llm_request: LlmRequest {
                messages: vec![Message::text(Role::User, "the task")],
                ..LlmRequest::default()
            },
            raw_request: None,
            metadata: None,
        };

        let built = judge.build_request(&State::default(), &request);

        // The task message alone; the rubric reaches the judge as an instruction block
        // rather than a message, so nothing was dropped by it not being counted here.
        assert_eq!(built.llm_request.messages.len(), 1);
        assert!(!built.llm_request.instructions.is_empty());
        assert!(
            !built
                .llm_request
                .messages
                .iter()
                .filter_map(|message| message.text_content("\n"))
                .any(|text| text.contains(TRAILING_ROUTING_INSTRUCTION))
        );
        Ok(())
    }

    #[test]
    fn capability_judge_builds_a_structured_request() -> Result<()> {
        let judge = capability_judge(None)?;
        let request = Request {
            llm_request: LlmRequest {
                model: Some("inbound".to_string()),
                messages: vec![
                    Message::text(Role::System, "client instructions"),
                    Message::text(Role::Developer, "client developer instructions"),
                    Message::text(Role::User, "initial task"),
                    Message::text(Role::Assistant, "old response"),
                    Message::text(Role::User, "old follow-up"),
                    Message::text(Role::Assistant, "recent 1"),
                    Message::text(Role::User, "recent 2"),
                    Message::text(Role::Assistant, "recent 3"),
                    Message::text(Role::User, "recent 4"),
                    Message::text(Role::Assistant, "recent 5"),
                ],
                ..LlmRequest::default()
            },
            raw_request: None,
            metadata: None,
        };
        let judge_request = judge.build_request(&State::default(), &request);

        assert_eq!(judge_request.llm_request.model, request.llm_request.model);
        assert_eq!(judge_request.llm_request.instructions.len(), 1);
        assert_eq!(judge_request.llm_request.instructions[0].role, Role::System);
        assert_eq!(
            judge_request.llm_request.instructions[0].content,
            InstructionBlock {
                role: Role::System,
                content: Message::text(Role::System, judge.contract().system_prompt()).content,
            }
            .content,
        );
        assert_eq!(judge_request.llm_request.messages.len(), 2);
        let contents = judge_request
            .llm_request
            .messages
            .iter()
            .filter_map(|message| message.text_content("\n"))
            .collect::<Vec<_>>();
        assert!(contents.contains(&"recent 4".to_string()));
        assert!(contents.contains(&"initial task".to_string()));
        assert!(!contents.contains(&"recent 5".to_string()));
        assert!(!contents.contains(&"client instructions".to_string()));
        assert_eq!(
            judge_request.llm_request.output.response_format,
            Some(judge.contract().response_format().clone())
        );
        assert_eq!(
            judge_request.llm_request.output.max_output_tokens,
            Some(DEFAULT_JUDGE_MAX_OUTPUT_TOKENS)
        );
        Ok(())
    }

    fn sample_value(spec: &Value) -> Value {
        if let Some(first) = spec
            .get("enum")
            .and_then(Value::as_array)
            .and_then(|values| values.first())
        {
            return first.clone();
        }
        match spec.get("type").and_then(Value::as_str) {
            Some("number") => serde_json::json!(0.5),
            Some("boolean") => serde_json::json!(false),
            _ => serde_json::json!("sample"),
        }
    }

    fn schema_shaped_verdict(schema: &Value) -> Result<String> {
        let properties = schema
            .pointer("/json_schema/schema/properties")
            .and_then(Value::as_object)
            .ok_or_else(|| LibsyError::AlgorithmError {
                message: "packaged schema declares no properties".to_string(),
            })?;
        Ok(Value::Object(
            properties
                .iter()
                .map(|(name, spec)| (name.clone(), sample_value(spec)))
                .collect(),
        )
        .to_string())
    }

    /// Built from the schema so a property added there fails here rather than silently
    /// rejecting every production verdict.
    #[test]
    fn every_schema_property_round_trips_through_the_judge_parser() -> Result<()> {
        let contract =
            LlmTaskClassifier::load_capability_contract(&ClassifierContractConfig::default())?;
        let schema = contract.response_format();
        let reply = schema_shaped_verdict(schema)?;
        let judge: CapabilityJudge = StructuredJudge::new(
            TaskInput {
                recent_turn_window: None,
            },
            contract,
            SerdeDecoder::new(),
            JudgeRuntimeConfig::new(DEFAULT_JUDGE_MAX_OUTPUT_TOKENS)?,
        );

        let verdict = judge.parse(&text_response(None, reply))?;

        assert!(verdict.is_valid());
        assert!((0.0..=1.0).contains(&verdict.p_solve));
        Ok(())
    }

    #[test]
    fn packaged_prompt_keeps_the_schema_in_the_structured_request() -> Result<()> {
        let contract =
            LlmTaskClassifier::load_capability_contract(&ClassifierContractConfig::default())?;
        let prompt = contract.system_prompt();
        let schema_name = contract
            .response_format()
            .pointer("/json_schema/name")
            .and_then(Value::as_str)
            .ok_or_else(|| LibsyError::AlgorithmError {
                message: "packaged response schema has no name".to_string(),
            })?;
        assert_eq!(schema_name, "CapabilityClassifierDecision");
        assert!(prompt.contains("SUP-1 [supported]"));
        assert!(prompt.contains("SUP-5 [supported]"));
        assert!(!prompt.contains("{{RESPONSE_SCHEMA}}"));
        assert!(!prompt.contains("\"type\": \"object\""));
        assert!(!prompt.contains("\"json_schema\""));
        assert!(!prompt.contains(schema_name));
        let rule_values = contract
            .response_format()
            .pointer("/json_schema/schema/properties/primary_rule/enum")
            .and_then(Value::as_array)
            .ok_or_else(|| LibsyError::AlgorithmError {
                message: "rendered response schema has no primary rule enum".to_string(),
            })?;
        assert!(
            rule_values
                .iter()
                .any(|value| value.as_str() == Some("SUP-1"))
        );
        assert!(
            rule_values
                .iter()
                .any(|value| value.as_str() == Some("none"))
        );
        Ok(())
    }
}