Skip to main content

switchyard_libsy/algorithms/
llm_class.rs

1// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2// SPDX-License-Identifier: Apache-2.0
3
4//! Judge-backed capability, escalation, and custom-policy routing.
5
6use std::collections::HashSet;
7use std::sync::Arc;
8
9use async_trait::async_trait;
10use serde::{Deserialize, Deserializer};
11use serde_json::Value;
12use switchyard_protocol::{Category, ContentBlock, Message, Role};
13
14mod decision;
15pub use decision::{DecisionChoice, DecisionJudgeConfig, DirectClassifierConfig};
16use decision::{DecisionClassifier, DirectClassifier};
17
18use super::escalation;
19use super::fall_through::FallThrough;
20use super::util::DEFAULT_JUDGE_MAX_OUTPUT_TOKENS;
21use super::util::affinity::{AffinityRouter, ClassifyTrigger, is_user_turn};
22use super::util::classifier_contract::{
23    ClassifierContract, ClassifierContractConfig, ClassifierResponseFormat,
24};
25use super::util::escalation::EscalationJudgeConfig;
26use super::util::llm_judge::{
27    ClassifierInput, JsonSchemaDecoder, JudgeClassifier, JudgePolicy, JudgeRuntimeConfig,
28    SerdeDecoder, StructuredJudge,
29};
30use super::util::target_selector::TargetSelectorPolicy;
31use crate::core::algorithm::{Algorithm, Driver};
32use crate::core::classifier::{Classification, Classifier, Score};
33use crate::core::state::State;
34use crate::{LibsyError, Result};
35use switchyard_protocol::{Request, Response};
36
37const PROMPT_TEMPLATE: &str = include_str!("../prompts/capability-classifier/prompt.md");
38const SCHEMA_TEMPLATE: &str = include_str!("../prompts/capability-classifier/schema.json");
39/// Telemetry label for this algorithm's spans, metrics, and logs.
40const ALGORITHM_NAME: &str = "llm_task_classifier";
41
42/// Restates the task after the conversation the judge is asked to route.
43///
44/// The rubric leads the request, which works while the payload is a single short
45/// message. Once `recent_turn_window` includes real conversation, those instructions
46/// sit far from the generation point and the judge sometimes *answers* the
47/// conversation instead of classifying it — the reply then fails to parse, the
48/// verdict is unavailable, and the turn silently falls back. Repeating the task at
49/// the end is what reliably prevents that.
50const TRAILING_ROUTING_INSTRUCTION: &str =
51    "Route the conversation above. Output ONLY the routing JSON object, nothing else.";
52
53#[derive(Deserialize)]
54#[serde(deny_unknown_fields)]
55struct TaskClassifierVerdict {
56    crux: String,
57    primary_rule: String,
58    capability_boundary: String,
59    p_solve: f64,
60}
61
62impl TaskClassifierVerdict {
63    /// Rejects malformed or internally inconsistent verdicts before policy evaluation.
64    fn is_valid(&self) -> bool {
65        (0.0..=1.0).contains(&self.p_solve)
66            && !self.crux.trim().is_empty()
67            && matches!(
68                (
69                    self.primary_rule.as_str(),
70                    self.capability_boundary.as_str()
71                ),
72                ("SUP-1" | "SUP-2" | "SUP-3" | "SUP-4" | "SUP-5", "supported")
73                    | ("UNC-1" | "UNC-2", "uncertain")
74                    | ("LIM-1" | "LIM-2", "unsupported")
75                    | ("none", "unmatched")
76            )
77    }
78
79    /// Returns the number of threshold steps assigned to this capability boundary.
80    ///
81    /// `unmatched_steps` comes from the route configuration: a request that matches no
82    /// capability rule has no evidence behind its solve probability, so operators can
83    /// require more confidence for it than for an uncertain verdict.
84    fn boundary_steps(&self, unmatched_steps: u8) -> Option<u8> {
85        match self.capability_boundary.as_str() {
86            "supported" => Some(0),
87            "uncertain" => Some(1),
88            "unmatched" => Some(unmatched_steps),
89            "unsupported" => Some(2),
90            _ => None,
91        }
92    }
93}
94
95/// Keeps the opening task and the last `recent_turn_window` turns after it. A
96/// window of `0` keeps the task alone.
97///
98/// Inline system and developer messages are kept, in their original order,
99/// ahead of the task.
100///
101/// Selects by reference and clones only what survives — a coding-agent
102/// conversation carries every tool result, so cloning it whole to keep a window
103/// would copy the transcript on each judged turn.
104fn trim_messages(messages: &[Message], recent_turn_window: usize) -> Vec<Message> {
105    let is_instruction = |message: &Message| matches!(message.role, Role::System | Role::Developer);
106    let mut kept: Vec<&Message> = messages.iter().filter(|m| is_instruction(m)).collect();
107    let Some(task) = messages.iter().position(is_user_turn) else {
108        return kept.into_iter().cloned().collect();
109    };
110    kept.push(&messages[task]);
111
112    let tail: Vec<&Message> = messages[task + 1..]
113        .iter()
114        .filter(|m| !is_instruction(m))
115        .collect();
116    kept.extend(&tail[window_start(&tail, recent_turn_window)..]);
117    kept.into_iter().cloned().collect()
118}
119
120/// The first index of the trailing window.
121///
122/// Counting messages alone can start the window between an assistant tool call and the
123/// result answering it, leaving the judge a result whose call id was never introduced. The
124/// start therefore moves back to the nearest one that keeps every tool pair whole.
125///
126/// One newest-to-oldest pass carries the ids still waiting for a call. Direction is what
127/// makes it correct: ids repeat across a conversation, and in this order a call is only
128/// ever seen after the results it could answer, so a later call — already passed — clears
129/// nothing. A result whose call sits before the opening task, which trimming never reaches,
130/// keeps the set non-empty to the end and falls back to the counted start, so an unpairable
131/// result costs one pass and cannot widen the window to the whole conversation.
132fn window_start(tail: &[&Message], recent_turn_window: usize) -> usize {
133    let counted = tail.len().saturating_sub(recent_turn_window);
134    // An empty window holds no result to pair, and the loop below never visits its start.
135    if counted == tail.len() {
136        return counted;
137    }
138    let mut unpaired: HashSet<&str> = HashSet::new();
139    for (start, message) in tail.iter().enumerate().rev() {
140        // Blocks reverse too, so a call answers a result only when it precedes it inside
141        // one message as well as across messages.
142        for block in message.content.iter().rev() {
143            match block {
144                ContentBlock::ToolResult(result) => {
145                    unpaired.insert(result.tool_call_id.as_str());
146                }
147                ContentBlock::ToolCall(call) => {
148                    unpaired.remove(call.id.as_str());
149                }
150                _ => {}
151            }
152        }
153        if start <= counted && unpaired.is_empty() {
154            return start;
155        }
156    }
157    counted
158}
159
160/// Keeps the opening task and the latest user follow-up when they differ.
161fn task_messages(messages: &[Message]) -> Vec<Message> {
162    // Decoders also use the user role for tool results. Select ordinary user content
163    // first, so a tool result cannot replace the opening task or latest follow-up.
164    let is_task_content = |block: &ContentBlock| {
165        !matches!(
166            block,
167            ContentBlock::ToolCall(_)
168                | ContentBlock::ToolResult(_)
169                | ContentBlock::Reasoning { .. }
170        )
171    };
172    let mut user_messages = messages
173        .iter()
174        .filter(|message| is_user_turn(message) && message.content.iter().any(is_task_content));
175    let Some(opening_task) = user_messages.next() else {
176        return Vec::new();
177    };
178    [Some(opening_task), user_messages.next_back()]
179        .into_iter()
180        .flatten()
181        .map(|message| {
182            Message::new(
183                Role::User,
184                message
185                    .content
186                    .iter()
187                    .filter(|block| is_task_content(block))
188                    .cloned()
189                    .collect(),
190            )
191        })
192        .collect()
193}
194
195/// Selects the task messages shown to capability and custom-schema classifiers.
196struct TaskInput {
197    recent_turn_window: Option<usize>,
198}
199
200impl TaskInput {
201    /// Selects conversation content without adding judge instructions.
202    fn messages(&self, request: &Request) -> Vec<Message> {
203        // The default preserves the whole-task anchor and latest user update. A
204        // configured window widens that to the surrounding conversation.
205        let mut messages = match self.recent_turn_window {
206            Some(window) => trim_messages(&request.llm_request.messages, window),
207            None => task_messages(&request.llm_request.messages),
208        };
209        // Reasoning is provider-private and not required to classify the task. Some
210        // upstreams also reject an unsigned reasoning item replayed without the
211        // opaque state it was issued with, so it cannot travel through a windowed
212        // classifier request as ordinary history.
213        for message in &mut messages {
214            message
215                .content
216                .retain(|block| !matches!(block, ContentBlock::Reasoning { .. }));
217        }
218        messages.retain(|message| !message.content.is_empty());
219        messages
220    }
221}
222
223impl ClassifierInput for TaskInput {
224    fn build_messages(&self, _state: &State, request: &Request) -> Vec<Message> {
225        let mut messages = self.messages(request);
226        // Only the windowed path carries assistant turns and tool traffic for the judge
227        // to be distracted by. The default path is user task messages only — the anchor
228        // and the latest follow-up — so there is nothing there to outrank.
229        if self.recent_turn_window.is_some() {
230            messages.push(Message::text(
231                Role::User,
232                TRAILING_ROUTING_INSTRUCTION.to_string(),
233            ));
234        }
235        messages
236    }
237}
238
239struct TaskClassifierPolicy {
240    base_threshold: f64,
241    threshold_step: f64,
242    unmatched_steps: u8,
243}
244
245impl TaskClassifierPolicy {
246    fn new(config: &LlmCapabilityConfig) -> Self {
247        Self {
248            base_threshold: config.base_threshold,
249            threshold_step: config.threshold_step,
250            unmatched_steps: config.unmatched_steps,
251        }
252    }
253
254    /// Returns the required solve probability for one validated verdict.
255    fn threshold(&self, verdict: &TaskClassifierVerdict) -> Option<f64> {
256        let steps = verdict.boundary_steps(self.unmatched_steps)?;
257        Some(self.base_threshold + f64::from(steps) * self.threshold_step)
258    }
259}
260
261impl JudgePolicy for TaskClassifierPolicy {
262    type Verdict = TaskClassifierVerdict;
263
264    fn to_classification(
265        &self,
266        verdict: Option<&Self::Verdict>,
267        driver: &Driver,
268    ) -> Result<Classification> {
269        // Judge output is untrusted. An absent, invalid, or inconsistent verdict is
270        // ambiguous so the surrounding router applies its configured fallback.
271        let Some(verdict) = verdict.filter(|verdict| verdict.is_valid()) else {
272            return Ok(Classification::Ambiguous(vec![]));
273        };
274        // A usable verdict below the capability threshold is still a decision: the judge
275        // does not trust the efficient tier with this task.
276        let Some(threshold) = self.threshold(verdict) else {
277            return Ok(Classification::Ambiguous(vec![]));
278        };
279        let category = if verdict.p_solve >= threshold
280            || (threshold - verdict.p_solve).abs() <= f64::EPSILON
281        {
282            Category::Efficient
283        } else {
284            Category::Capable
285        };
286        // The chosen category may have no models configured. That is ambiguous, not an
287        // error, so the surrounding router applies its configured fallback.
288        let Some(target) = driver.models_for(&category).first().cloned() else {
289            return Ok(Classification::Ambiguous(vec![]));
290        };
291        Ok(Classification::Scores(vec![Score {
292            target,
293            confidence: 1.0,
294            category: Some(category),
295        }]))
296    }
297}
298
299/// Maps valid verdicts to scores, invalid verdicts to a reason, and leaves absent verdicts alone.
300fn capability_evidence(
301    policy: &TaskClassifierPolicy,
302    verdict: Option<&TaskClassifierVerdict>,
303) -> Option<Value> {
304    let verdict = verdict?;
305    let Some(threshold) = verdict
306        .is_valid()
307        .then(|| policy.threshold(verdict))
308        .flatten()
309    else {
310        return Some(serde_json::json!({
311            "source": "fail_open",
312            "reason_code": "invalid_verdict",
313        }));
314    };
315    Some(serde_json::json!({
316        "source": "llm-classifier",
317        "score": verdict.p_solve,
318        "threshold": threshold,
319    }))
320}
321
322/// Settings shared by capability judges and their surrounding route.
323#[derive(Clone, Debug)]
324pub struct TaskClassifierConfig {
325    /// The judge's prediction contract and routing policy.
326    pub judge: CapabilityJudgeConfig,
327    /// Routes to the capable tier on judge client failures and deadlines. Defaults to true.
328    pub fail_open: bool,
329    /// How often the classifier re-decides this session's target.
330    pub classify_trigger: ClassifyTrigger,
331    /// Uses the first user message as the SessionKey for sticky routing when session metadata is unavailable.
332    pub message_hash_fallback: bool,
333    /// `None` judges the opening task and latest user follow-up. `Some(n)` also
334    /// includes the surrounding conversation using the last `n` turns.
335    pub recent_turn_window: Option<usize>,
336}
337
338/// Selects how a capability judge obtains and interprets its prediction.
339#[derive(Clone, Debug)]
340pub enum CapabilityJudgeConfig {
341    /// A structured LLM verdict with a solve probability and capability boundary.
342    Llm(LlmCapabilityConfig),
343    /// A Choice score for the capable candidate succeeding where the efficient one fails.
344    Decision(DecisionJudgeConfig),
345}
346
347impl Default for CapabilityJudgeConfig {
348    fn default() -> Self {
349        Self::Llm(LlmCapabilityConfig::default())
350    }
351}
352
353/// Prompt, output, and solve-probability settings for an LLM capability judge.
354#[derive(Clone, Debug)]
355pub struct LlmCapabilityConfig {
356    /// Lowest solve probability that routes a supported task to the efficient target.
357    pub base_threshold: f64,
358    /// Amount added per capability-boundary step: zero for supported, one for
359    /// uncertain, `unmatched_steps` for unmatched, and two for unsupported.
360    pub threshold_step: f64,
361    /// Threshold steps applied to an unmatched verdict (no capability rule applies).
362    ///
363    /// Defaults to 1, the same as an uncertain verdict. 2 treats it like unsupported.
364    pub unmatched_steps: u8,
365    /// Prompt and verdict contract settings for the classifier judge.
366    pub contract: ClassifierContractConfig,
367    /// Maximum completion tokens available to the classifier verdict.
368    pub max_output_tokens: u64,
369    /// Whole-consultation bound on the judge call, in milliseconds. Covers the model
370    /// call and the response drain; on expiry the judge is treated as unavailable and
371    /// follows the route's `fail_open` setting. `None` is unbounded.
372    pub judge_deadline_ms: Option<u64>,
373}
374
375impl Default for LlmCapabilityConfig {
376    fn default() -> Self {
377        Self {
378            base_threshold: 0.0,
379            threshold_step: 0.0,
380            unmatched_steps: DEFAULT_UNMATCHED_STEPS,
381            contract: ClassifierContractConfig::default(),
382            max_output_tokens: DEFAULT_JUDGE_MAX_OUTPUT_TOKENS,
383            judge_deadline_ms: None,
384        }
385    }
386}
387
388/// Flat serialized shape that maps prompt settings into the runtime contract.
389#[derive(Deserialize)]
390#[serde(deny_unknown_fields)]
391struct TaskClassifierConfigWire {
392    base_threshold: f64,
393    #[serde(default = "default_fail_open")]
394    fail_open: bool,
395    #[serde(default)]
396    threshold_step: f64,
397    #[serde(default = "default_unmatched_steps")]
398    unmatched_steps: u8,
399    #[serde(default)]
400    classify_trigger: ClassifyTrigger,
401    #[serde(default)]
402    message_hash_fallback: bool,
403    #[serde(default)]
404    recent_turn_window: Option<usize>,
405    #[serde(default)]
406    prompt: Option<String>,
407    #[serde(default)]
408    response_format_type: ClassifierResponseFormat,
409    #[serde(default = "default_judge_max_output_tokens")]
410    max_output_tokens: u64,
411    #[serde(default)]
412    judge_deadline_ms: Option<u64>,
413}
414
415impl<'de> Deserialize<'de> for TaskClassifierConfig {
416    fn deserialize<D>(deserializer: D) -> std::result::Result<Self, D::Error>
417    where
418        D: Deserializer<'de>,
419    {
420        let wire = TaskClassifierConfigWire::deserialize(deserializer)?;
421        let mut contract = ClassifierContractConfig::default();
422        if let Some(prompt) = wire.prompt {
423            contract = contract.with_prompt(prompt);
424        }
425        contract = contract.with_response_format_type(wire.response_format_type);
426        Ok(Self {
427            judge: CapabilityJudgeConfig::Llm(LlmCapabilityConfig {
428                base_threshold: wire.base_threshold,
429                threshold_step: wire.threshold_step,
430                unmatched_steps: wire.unmatched_steps,
431                contract,
432                max_output_tokens: wire.max_output_tokens,
433                judge_deadline_ms: wire.judge_deadline_ms,
434            }),
435            fail_open: wire.fail_open,
436            classify_trigger: wire.classify_trigger,
437            message_hash_fallback: wire.message_hash_fallback,
438            recent_turn_window: wire.recent_turn_window,
439        })
440    }
441}
442
443const fn default_fail_open() -> bool {
444    true
445}
446
447/// Unmatched verdicts share the uncertain step unless the route says otherwise.
448pub const DEFAULT_UNMATCHED_STEPS: u8 = 1;
449
450const fn default_unmatched_steps() -> u8 {
451    DEFAULT_UNMATCHED_STEPS
452}
453
454const fn default_judge_max_output_tokens() -> u64 {
455    DEFAULT_JUDGE_MAX_OUTPUT_TOKENS
456}
457
458impl Default for TaskClassifierConfig {
459    fn default() -> Self {
460        Self {
461            judge: CapabilityJudgeConfig::default(),
462            fail_open: default_fail_open(),
463            classify_trigger: ClassifyTrigger::default(),
464            message_hash_fallback: false,
465            recent_turn_window: None,
466        }
467    }
468}
469
470impl LlmCapabilityConfig {
471    fn validate(&self) -> Result<()> {
472        if !(0.0..=1.0).contains(&self.base_threshold) {
473            return Err(LibsyError::AlgorithmError {
474                message: format!(
475                    "base_threshold must be between 0 and 1, got {}",
476                    self.base_threshold
477                ),
478            });
479        }
480        if !self.threshold_step.is_finite() || self.threshold_step < 0.0 {
481            return Err(LibsyError::AlgorithmError {
482                message: format!(
483                    "threshold_step must be finite and greater than or equal to 0, got {}",
484                    self.threshold_step
485                ),
486            });
487        }
488        if self.unmatched_steps > 2 {
489            return Err(LibsyError::AlgorithmError {
490                message: format!(
491                    "unmatched_steps must be 0, 1, or 2, got {}",
492                    self.unmatched_steps
493                ),
494            });
495        }
496        let unsupported_threshold = self.base_threshold + 2.0 * self.threshold_step;
497        if unsupported_threshold > 1.0 && unsupported_threshold - 1.0 > f64::EPSILON {
498            return Err(LibsyError::AlgorithmError {
499                message: format!(
500                    "base_threshold + 2 * threshold_step must be at most 1, got {unsupported_threshold}"
501                ),
502            });
503        }
504        if self.max_output_tokens == 0 {
505            return Err(LibsyError::AlgorithmError {
506                message: "max_output_tokens must be at least 1".to_string(),
507            });
508        }
509        Ok(())
510    }
511}
512
513impl TaskClassifierConfig {
514    fn validate(&self) -> Result<()> {
515        match &self.judge {
516            CapabilityJudgeConfig::Llm(config) => config.validate()?,
517            CapabilityJudgeConfig::Decision(config) => config.validate()?,
518        }
519        // 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.
520        if self.message_hash_fallback && self.classify_trigger == ClassifyTrigger::EveryRequest {
521            return Err(LibsyError::AlgorithmError {
522                message:
523                    "message_hash_fallback requires classify_trigger = new_session or user_turn"
524                        .to_string(),
525            });
526        }
527        Ok(())
528    }
529}
530
531/// Policy that maps a custom classifier verdict to a routing target.
532#[derive(Clone, Debug)]
533pub enum CustomClassifierPolicy {
534    /// Resolves a JSON Pointer and treats its string value as a model category.
535    TargetSelector {
536        /// JSON Pointer evaluated against each schema-validated verdict.
537        selector: String,
538    },
539}
540
541impl CustomClassifierPolicy {
542    /// Creates a policy that selects a model category through a JSON Pointer.
543    pub fn target_selector(selector: impl Into<String>) -> Self {
544        Self::TargetSelector {
545            selector: selector.into(),
546        }
547    }
548}
549
550/// Settings for a classifier whose JSON Schema and target-selection policy are user supplied.
551#[derive(Clone, Debug)]
552pub struct CustomClassifierConfig {
553    /// System prompt sent to the classifier judge.
554    pub prompt: String,
555    /// Inner JSON Schema placed inside the provider's structured-output wrapper.
556    pub response_schema: Value,
557    /// Deterministic policy applied after the verdict passes schema validation.
558    pub policy: CustomClassifierPolicy,
559    /// How often the classifier re-decides this session's target.
560    pub classify_trigger: ClassifyTrigger,
561    /// Uses the first user message when session metadata is unavailable.
562    pub message_hash_fallback: bool,
563    /// Trailing conversation turns shown to the classifier judge.
564    pub recent_turn_window: Option<usize>,
565    /// Maximum completion tokens available to the classifier verdict.
566    pub max_output_tokens: u64,
567    /// Whole-consultation bound on the judge call, in milliseconds. `None` is unbounded.
568    pub judge_deadline_ms: Option<u64>,
569}
570
571impl CustomClassifierConfig {
572    /// Creates a custom-schema classifier contract with conservative runtime defaults.
573    pub fn new(
574        prompt: impl Into<String>,
575        response_schema: Value,
576        policy: CustomClassifierPolicy,
577    ) -> Self {
578        Self {
579            prompt: prompt.into(),
580            response_schema,
581            policy,
582            classify_trigger: ClassifyTrigger::default(),
583            message_hash_fallback: false,
584            recent_turn_window: None,
585            max_output_tokens: DEFAULT_JUDGE_MAX_OUTPUT_TOKENS,
586            judge_deadline_ms: None,
587        }
588    }
589
590    fn validate(&self) -> Result<()> {
591        if self.max_output_tokens == 0 {
592            return Err(LibsyError::AlgorithmError {
593                message: "max_output_tokens must be at least 1".to_string(),
594            });
595        }
596        // 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.
597        if self.message_hash_fallback && self.classify_trigger == ClassifyTrigger::EveryRequest {
598            return Err(LibsyError::AlgorithmError {
599                message:
600                    "message_hash_fallback requires classify_trigger = new_session or user_turn"
601                        .to_string(),
602            });
603        }
604        Ok(())
605    }
606}
607
608enum CustomPolicyRuntime {
609    TargetSelector(TargetSelectorPolicy),
610}
611
612impl JudgePolicy for CustomPolicyRuntime {
613    type Verdict = Value;
614
615    fn to_classification(
616        &self,
617        verdict: Option<&Self::Verdict>,
618        driver: &Driver,
619    ) -> Result<Classification> {
620        match self {
621            Self::TargetSelector(policy) => policy.to_classification(verdict, driver),
622        }
623    }
624}
625
626/// Builds the affinity router a trigger calls for, if any.
627fn affinity_router(
628    trigger: ClassifyTrigger,
629    message_hash_fallback: bool,
630) -> Option<Arc<AffinityRouter>> {
631    let router = match trigger {
632        ClassifyTrigger::EveryRequest => return None,
633        ClassifyTrigger::NewSession => AffinityRouter::new(),
634        ClassifyTrigger::UserTurn => AffinityRouter::new().with_release_on_user_turn(),
635    };
636    let router = if message_hash_fallback {
637        router.with_message_hash_fallback()
638    } else {
639        router
640    };
641    Some(Arc::new(router))
642}
643
644/// Routes requests through a capability, escalation, custom, or direct classifier mode.
645pub struct LlmTaskClassifier {
646    route: FallThrough,
647    /// Classifier used when this router is embedded in another cascade.
648    inner: Arc<dyn Classifier>,
649}
650
651struct ClassifierRouteConfig {
652    default_target: Category,
653    classify_trigger: ClassifyTrigger,
654    message_hash_fallback: bool,
655}
656
657/// Terminal classifier for a cascade whose classifiers may all abstain.
658/// Closes a cascade with the first runtime model in `category`.
659pub struct DefaultCategoryClassifier(pub Category);
660
661#[async_trait]
662impl Classifier for DefaultCategoryClassifier {
663    async fn score(
664        &self,
665        _state: &mut State,
666        _request: &mut Request,
667        driver: &Driver,
668    ) -> Result<(Classification, Option<Response>)> {
669        let target = driver.first_model_for(&self.0)?;
670        driver.set_evidence_if_empty(serde_json::json!({"source": "fall_open"}));
671        Ok((
672            Classification::Scores(vec![Score {
673                target: target.clone(),
674                confidence: 0.0,
675                category: Some(self.0.clone()),
676            }]),
677            None,
678        ))
679    }
680}
681
682/// Complete construction settings for one LLM classifier mode.
683#[derive(Clone)]
684#[non_exhaustive]
685pub enum LlmClassifierConfig {
686    /// Routes between efficient and capable targets from a task-level verdict.
687    Capability {
688        /// Capability classifier settings.
689        config: TaskClassifierConfig,
690    },
691    /// Judges efficient responses and escalates after a confirmed streak.
692    Escalation {
693        /// Prompt and verdict contract settings for the escalation judge.
694        contract: ClassifierContractConfig,
695        /// Escalation policy settings.
696        config: EscalationJudgeConfig,
697        /// Maximum completion tokens available to the escalation verdict.
698        max_output_tokens: u64,
699        /// Whole-consultation bound on the escalation judge call, in milliseconds.
700        judge_deadline_ms: Option<u64>,
701    },
702    /// Routes among model categories using a user-supplied schema and policy.
703    Custom {
704        /// Category selected when the judge does not produce a usable verdict.
705        default_target: Category,
706        /// Custom classifier settings.
707        config: CustomClassifierConfig,
708    },
709    /// Routes to the model category a decision service selects.
710    Direct {
711        /// Category used when the decision call fails or returns an invalid choice.
712        default_target: Category,
713        /// Direct selection settings.
714        config: DirectClassifierConfig,
715    },
716}
717
718impl LlmTaskClassifier {
719    /// Builds the classifier mode described by `config`.
720    ///
721    /// # Errors
722    ///
723    /// Returns an error when the selected mode's targets, contract, policy, or runtime
724    /// settings are invalid.
725    pub fn new(config: LlmClassifierConfig) -> Result<Self> {
726        match config {
727            LlmClassifierConfig::Capability { config } => Self::build_capability(config),
728            LlmClassifierConfig::Escalation {
729                contract,
730                config,
731                max_output_tokens,
732                judge_deadline_ms,
733            } => Self::build_escalation(contract, config, max_output_tokens, judge_deadline_ms),
734            LlmClassifierConfig::Custom {
735                default_target,
736                config,
737            } => Self::build_custom(default_target, config),
738            LlmClassifierConfig::Direct {
739                default_target,
740                config,
741            } => Self::build_direct(default_target, config),
742        }
743    }
744
745    fn build_capability(config: TaskClassifierConfig) -> Result<Self> {
746        config.validate()?;
747        let classify_trigger = config.classify_trigger;
748        let message_hash_fallback = config.message_hash_fallback;
749        let input = TaskInput {
750            recent_turn_window: config.recent_turn_window,
751        };
752        let classifier: Arc<dyn Classifier> = match config.judge {
753            CapabilityJudgeConfig::Llm(judge) => Arc::new(
754                JudgeClassifier::new(
755                    StructuredJudge::new(
756                        input,
757                        Self::load_capability_contract(&judge.contract)?,
758                        SerdeDecoder::new(),
759                        JudgeRuntimeConfig::new(judge.max_output_tokens)?
760                            .with_deadline_ms(judge.judge_deadline_ms)?,
761                    ),
762                    TaskClassifierPolicy::new(&judge),
763                )
764                .with_error_recovery(config.fail_open)
765                .with_evidence(capability_evidence),
766            ),
767            CapabilityJudgeConfig::Decision(judge) => {
768                Arc::new(DecisionClassifier::new(judge, input, config.fail_open)?)
769            }
770        };
771        Self::from_classifier(
772            classifier,
773            ClassifierRouteConfig {
774                default_target: Category::Capable,
775                classify_trigger,
776                message_hash_fallback,
777            },
778        )
779    }
780
781    fn build_custom(default_target: Category, config: CustomClassifierConfig) -> Result<Self> {
782        Self::from_classifier(
783            Self::custom_classifier(&config, false)?,
784            ClassifierRouteConfig {
785                default_target,
786                classify_trigger: config.classify_trigger,
787                message_hash_fallback: config.message_hash_fallback,
788            },
789        )
790    }
791
792    fn build_direct(default_target: Category, config: DirectClassifierConfig) -> Result<Self> {
793        if !config.options.contains(&default_target) {
794            return Err(LibsyError::AlgorithmError {
795                message: "direct default_target must be one of the options".into(),
796            });
797        }
798        let route = ClassifierRouteConfig {
799            default_target,
800            classify_trigger: config.classify_trigger,
801            message_hash_fallback: config.message_hash_fallback,
802        };
803        Self::from_classifier(Arc::new(DirectClassifier::new(config)?), route)
804    }
805
806    /// The judge face alone, for composite's own trigger and retained tier.
807    pub(crate) fn custom_classifier(
808        config: &CustomClassifierConfig,
809        fail_open: bool,
810    ) -> Result<Arc<dyn Classifier>> {
811        config.validate()?;
812        let contract =
813            ClassifierContract::from_inner_schema(&config.prompt, config.response_schema.clone())?;
814        let policy = match &config.policy {
815            CustomClassifierPolicy::TargetSelector { selector } => {
816                CustomPolicyRuntime::TargetSelector(TargetSelectorPolicy::new(selector)?)
817            }
818        };
819        let classifier = JudgeClassifier::new(
820            StructuredJudge::new(
821                TaskInput {
822                    recent_turn_window: config.recent_turn_window,
823                },
824                contract,
825                JsonSchemaDecoder::new(),
826                JudgeRuntimeConfig::new(config.max_output_tokens)?
827                    .with_deadline_ms(config.judge_deadline_ms)?,
828            ),
829            policy,
830        )
831        .with_error_recovery(fail_open);
832        let classifier = if fail_open {
833            classifier.with_evidence(|_, verdict| {
834                verdict.map(|_| serde_json::json!({"source": "llm-classifier"}))
835            })
836        } else {
837            classifier
838        };
839        Ok(Arc::new(classifier))
840    }
841
842    fn build_escalation(
843        contract_config: ClassifierContractConfig,
844        config: EscalationJudgeConfig,
845        max_output_tokens: u64,
846        judge_deadline_ms: Option<u64>,
847    ) -> Result<Self> {
848        let inner = escalation::build_classifier(
849            contract_config,
850            config,
851            max_output_tokens,
852            judge_deadline_ms,
853        )?;
854        Ok(Self {
855            route: FallThrough::new_with_state()
856                .with_name(ALGORITHM_NAME)
857                .with_classifier(Arc::clone(&inner)),
858            inner,
859        })
860    }
861
862    /// Loads the packaged capability-classifier contract.
863    fn load_capability_contract(config: &ClassifierContractConfig) -> Result<ClassifierContract> {
864        ClassifierContract::from_config(config, PROMPT_TEMPLATE, SCHEMA_TEMPLATE)
865    }
866
867    /// Keeps affinity and fallback ordering identical across judge-backed modes.
868    fn from_classifier(inner: Arc<dyn Classifier>, config: ClassifierRouteConfig) -> Result<Self> {
869        // 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.
870        if config.message_hash_fallback && config.classify_trigger == ClassifyTrigger::EveryRequest
871        {
872            return Err(LibsyError::AlgorithmError {
873                message:
874                    "message_hash_fallback requires classify_trigger = new_session or user_turn"
875                        .to_string(),
876            });
877        }
878        // Affinity comes first so a retained assignment short-circuits the judge call.
879        let mut route = FallThrough::new_with_state().with_name(ALGORITHM_NAME);
880        if let Some(affinity) =
881            affinity_router(config.classify_trigger, config.message_hash_fallback).as_ref()
882        {
883            // Both roles must share one `Arc` so the classifier reads what the processor wrote.
884            route = route
885                .with_processor(affinity.clone())
886                .with_classifier(affinity.clone());
887        }
888        let fallback = DefaultCategoryClassifier(config.default_target);
889        Ok(Self {
890            route: route
891                .with_classifier(inner.clone())
892                .with_classifier(Arc::new(fallback)),
893            inner,
894        })
895    }
896}
897
898#[async_trait]
899impl Classifier for LlmTaskClassifier {
900    async fn score(
901        &self,
902        state: &mut State,
903        request: &mut Request,
904        driver: &Driver,
905    ) -> Result<(Classification, Option<Response>)> {
906        self.inner.score(state, request, driver).await
907    }
908}
909
910#[async_trait]
911impl Algorithm for LlmTaskClassifier {
912    fn name(&self) -> &str {
913        "llm_task_classifier"
914    }
915
916    async fn route(
917        self: Arc<Self>,
918        driver: Driver,
919        request: Request,
920    ) -> Result<crate::RoutingOutcome> {
921        self.route.execute(driver, request).await
922    }
923}
924
925#[cfg(test)]
926mod tests {
927    use std::collections::HashMap;
928    use std::sync::Arc;
929
930    use parking_lot::Mutex;
931    use serde_json::Value;
932
933    use super::*;
934    use switchyard_protocol::{
935        ContentBlock, InstructionBlock, LlmClientError, LlmRequest, Metadata, ModelId, ToolCall,
936        ToolResult, completion_text, text_request, text_response,
937    };
938
939    use crate::algorithms::util::llm_judge::Judge;
940    use crate::core::testing::{Serve, test_drive_with_models};
941    use switchyard_protocol::{LlmResponse, Response};
942
943    const TEST_THRESHOLD: f64 = 0.5;
944
945    type CapabilityJudge = StructuredJudge<TaskInput, SerdeDecoder<TaskClassifierVerdict>>;
946
947    fn llm_config(base_threshold: f64) -> LlmCapabilityConfig {
948        LlmCapabilityConfig {
949            base_threshold,
950            ..LlmCapabilityConfig::default()
951        }
952    }
953
954    fn test_config(base_threshold: f64) -> TaskClassifierConfig {
955        TaskClassifierConfig {
956            judge: CapabilityJudgeConfig::Llm(llm_config(base_threshold)),
957            ..TaskClassifierConfig::default()
958        }
959    }
960
961    fn policy() -> TaskClassifierPolicy {
962        TaskClassifierPolicy::new(&llm_config(TEST_THRESHOLD))
963    }
964
965    fn runtime_models() -> HashMap<Category, Vec<ModelId>> {
966        [
967            (Category::Judge, vec![ModelId::from("judge")]),
968            (Category::Efficient, vec![ModelId::from("efficient")]),
969            (Category::Capable, vec![ModelId::from("capable")]),
970            (
971                Category::Any,
972                vec![ModelId::from("efficient"), ModelId::from("capable")],
973            ),
974        ]
975        .into()
976    }
977
978    fn policy_driver() -> Driver {
979        Driver::new("test", Arc::new(runtime_models().into())).0
980    }
981
982    fn verdict(
983        p_solve: f64,
984        capability_boundary: &str,
985        primary_rule: &str,
986    ) -> TaskClassifierVerdict {
987        TaskClassifierVerdict {
988            crux: "test crux".to_string(),
989            primary_rule: primary_rule.to_string(),
990            capability_boundary: capability_boundary.to_string(),
991            p_solve,
992        }
993    }
994
995    fn selected(
996        policy: &TaskClassifierPolicy,
997        verdict: Option<&TaskClassifierVerdict>,
998    ) -> Result<ModelId> {
999        policy
1000            .to_classification(verdict, &policy_driver())?
1001            .argmax(false)?
1002            .map(|score| score.target)
1003            .ok_or_else(|| LibsyError::AlgorithmError {
1004                message: "policy abstained".to_string(),
1005            })
1006    }
1007
1008    /// Records what each target received; answers the judge with a supported verdict and
1009    /// every other target with a plain completion.
1010    #[derive(Default)]
1011    struct Recorder {
1012        calls: Mutex<Vec<String>>,
1013        call_roles: Mutex<Vec<(String, bool)>>,
1014        judge_max_output_tokens: Mutex<Vec<Option<u64>>>,
1015        judge_system_prompts: Mutex<Vec<String>>,
1016    }
1017
1018    impl Recorder {
1019        fn calls(&self) -> Vec<String> {
1020            self.calls.lock().clone()
1021        }
1022
1023        fn call_roles(&self) -> Vec<(String, bool)> {
1024            self.call_roles.lock().clone()
1025        }
1026
1027        fn judge_max_output_tokens(&self) -> Vec<Option<u64>> {
1028            self.judge_max_output_tokens.lock().clone()
1029        }
1030
1031        fn judge_system_prompts(&self) -> Vec<String> {
1032            self.judge_system_prompts.lock().clone()
1033        }
1034
1035        fn serve(self: &Arc<Self>) -> impl Serve {
1036            let recorder = Arc::clone(self);
1037            move |model: ModelId, request: Request| {
1038                let recorder = Arc::clone(&recorder);
1039                async move {
1040                    let model = model.to_string();
1041                    recorder.calls.lock().push(model.clone());
1042                    recorder
1043                        .call_roles
1044                        .lock()
1045                        .push((model.clone(), model != "judge"));
1046                    let completion = if model == "judge" {
1047                        recorder
1048                            .judge_max_output_tokens
1049                            .lock()
1050                            .push(request.llm_request.output.max_output_tokens);
1051                        recorder.judge_system_prompts.lock().extend(
1052                            request
1053                                .llm_request
1054                                .instructions
1055                                .first()
1056                                .and_then(|instruction| {
1057                                    instruction.content.iter().find_map(|b| {
1058                                        if let ContentBlock::Text { text } = b {
1059                                            Some(text.clone())
1060                                        } else {
1061                                            None
1062                                        }
1063                                    })
1064                                }),
1065                        );
1066                        r#"{"crux":"bounded task","primary_rule":"SUP-1","capability_boundary":"supported","p_solve":0.9}"#.to_string()
1067                    } else {
1068                        format!("answer from {model}")
1069                    };
1070                    Ok(Response {
1071                        llm_response: LlmResponse::Agg(text_response(None, completion)),
1072                        metadata: request.metadata,
1073                        upstream_headers: http::HeaderMap::new(),
1074                    })
1075                }
1076            }
1077        }
1078    }
1079
1080    /// The judge times out; every other target answers normally.
1081    fn unreachable_judge() -> impl Serve {
1082        |model: ModelId, request: Request| async move {
1083            let model = model.to_string();
1084            if model == "judge" {
1085                return Err(LlmClientError::Timeout {
1086                    source: Box::new(std::io::Error::other("judge unreachable")),
1087                });
1088            }
1089            Ok(Response {
1090                llm_response: LlmResponse::Agg(text_response(None, format!("answer from {model}"))),
1091                metadata: request.metadata,
1092                upstream_headers: http::HeaderMap::new(),
1093            })
1094        }
1095    }
1096
1097    /// The judge stalls before answering; every other target answers normally.
1098    fn slow_judge(delay: std::time::Duration) -> impl Serve {
1099        move |model: ModelId, request: Request| {
1100            let model = model.to_string();
1101            let slow = model == "judge";
1102            async move {
1103                if slow {
1104                    tokio::time::sleep(delay).await;
1105                }
1106                Ok(Response {
1107                    llm_response: LlmResponse::Agg(text_response(
1108                        None,
1109                        format!("answer from {model}"),
1110                    )),
1111                    metadata: request.metadata,
1112                    upstream_headers: http::HeaderMap::new(),
1113                })
1114            }
1115        }
1116    }
1117
1118    /// The judge accepts the call but its stream never carries a chunk.
1119    fn stalled_judge_stream() -> impl Serve {
1120        use futures::StreamExt;
1121        |model: ModelId, request: Request| async move {
1122            let model = model.to_string();
1123            let llm_response = if model == "judge" {
1124                LlmResponse::Stream(futures::stream::pending().boxed())
1125            } else {
1126                LlmResponse::Agg(text_response(None, format!("answer from {model}")))
1127            };
1128            Ok(Response {
1129                llm_response,
1130                metadata: request.metadata,
1131                upstream_headers: http::HeaderMap::new(),
1132            })
1133        }
1134    }
1135
1136    fn deadline_router(fail_open: bool) -> Result<Arc<LlmTaskClassifier>> {
1137        Ok(Arc::new(LlmTaskClassifier::new(
1138            LlmClassifierConfig::Capability {
1139                config: TaskClassifierConfig {
1140                    judge: CapabilityJudgeConfig::Llm(LlmCapabilityConfig {
1141                        judge_deadline_ms: Some(10),
1142                        ..llm_config(TEST_THRESHOLD)
1143                    }),
1144                    fail_open,
1145                    ..TaskClassifierConfig::default()
1146                },
1147            },
1148        )?))
1149    }
1150
1151    fn router() -> Result<Arc<LlmTaskClassifier>> {
1152        Ok(Arc::new(LlmTaskClassifier::new(
1153            LlmClassifierConfig::Capability {
1154                config: test_config(TEST_THRESHOLD),
1155            },
1156        )?))
1157    }
1158
1159    fn classify_request() -> Request {
1160        Request {
1161            llm_request: text_request(Some("auto".to_string()), "classify this task"),
1162            raw_request: None,
1163            metadata: None,
1164        }
1165    }
1166
1167    fn classify_session_request() -> Request {
1168        Request {
1169            metadata: Some(Metadata {
1170                session_id: Some("session-1".to_string()),
1171                ..Metadata::default()
1172            }),
1173            ..classify_request()
1174        }
1175    }
1176
1177    fn classify_follow_up_request() -> Request {
1178        let mut request = classify_request();
1179        request
1180            .llm_request
1181            .messages
1182            .push(Message::text(Role::Assistant, "I will add the test."));
1183        request.llm_request.messages.push(Message::text(
1184            Role::User,
1185            "Now run the test suite and report the result.",
1186        ));
1187        request
1188    }
1189
1190    #[tokio::test]
1191    async fn an_unreachable_judge_routes_capable_instead_of_failing_the_request() -> Result<()> {
1192        let router = router()?;
1193
1194        let (selected_model, response) = test_drive_with_models(
1195            router,
1196            classify_request(),
1197            runtime_models(),
1198            unreachable_judge(),
1199        )
1200        .await?;
1201
1202        assert_eq!(selected_model, "capable");
1203        assert_eq!(
1204            response.llm_response.as_agg().map(completion_text),
1205            Some("answer from capable".to_string())
1206        );
1207        Ok(())
1208    }
1209
1210    #[tokio::test]
1211    async fn a_judge_past_its_deadline_routes_capable() -> Result<()> {
1212        let (selected_model, response) = test_drive_with_models(
1213            deadline_router(true)?,
1214            classify_request(),
1215            runtime_models(),
1216            slow_judge(std::time::Duration::from_millis(500)),
1217        )
1218        .await?;
1219
1220        assert_eq!(selected_model, "capable");
1221        assert_eq!(
1222            response.llm_response.as_agg().map(completion_text),
1223            Some("answer from capable".to_string())
1224        );
1225        Ok(())
1226    }
1227
1228    #[tokio::test]
1229    async fn a_stalled_judge_stream_is_cut_by_the_deadline() -> Result<()> {
1230        // The judge returned headers promptly, so only a bound on the whole
1231        // consultation — the drain included — can end this turn.
1232        let (selected_model, _) = test_drive_with_models(
1233            deadline_router(true)?,
1234            classify_request(),
1235            runtime_models(),
1236            stalled_judge_stream(),
1237        )
1238        .await?;
1239
1240        assert_eq!(selected_model, "capable");
1241        Ok(())
1242    }
1243
1244    #[tokio::test]
1245    async fn a_judge_past_its_deadline_stops_when_not_failing_open() -> Result<()> {
1246        let outcome = test_drive_with_models(
1247            deadline_router(false)?,
1248            classify_request(),
1249            runtime_models(),
1250            slow_judge(std::time::Duration::from_millis(500)),
1251        )
1252        .await;
1253        let error = outcome
1254            .err()
1255            .expect("a deadline expiry with fail_open = false must stop the request");
1256
1257        assert!(
1258            matches!(
1259                error,
1260                LibsyError::ClientCall {
1261                    source: LlmClientError::Timeout { .. },
1262                    ..
1263                }
1264            ),
1265            "expected a client timeout, got {error:?}"
1266        );
1267        Ok(())
1268    }
1269
1270    #[test]
1271    fn a_zero_judge_deadline_is_rejected_in_every_mode() {
1272        let capability = LlmClassifierConfig::Capability {
1273            config: TaskClassifierConfig {
1274                judge: CapabilityJudgeConfig::Llm(LlmCapabilityConfig {
1275                    judge_deadline_ms: Some(0),
1276                    ..llm_config(TEST_THRESHOLD)
1277                }),
1278                ..TaskClassifierConfig::default()
1279            },
1280        };
1281        assert!(
1282            LlmTaskClassifier::new(capability).is_err(),
1283            "capability mode must reject judge_deadline_ms = 0"
1284        );
1285
1286        let mut custom = CustomClassifierConfig::new(
1287            "Route by topic.",
1288            serde_json::json!({"type": "object"}),
1289            CustomClassifierPolicy::target_selector("/target"),
1290        );
1291        custom.judge_deadline_ms = Some(0);
1292        assert!(
1293            LlmTaskClassifier::new(LlmClassifierConfig::Custom {
1294                default_target: Category::Capable,
1295                config: custom,
1296            })
1297            .is_err(),
1298            "custom mode must reject judge_deadline_ms = 0"
1299        );
1300
1301        let escalation = LlmClassifierConfig::Escalation {
1302            contract: ClassifierContractConfig::default(),
1303            config: EscalationJudgeConfig::default(),
1304            max_output_tokens: DEFAULT_JUDGE_MAX_OUTPUT_TOKENS,
1305            judge_deadline_ms: Some(0),
1306        };
1307        assert!(
1308            LlmTaskClassifier::new(escalation).is_err(),
1309            "escalation mode must reject judge_deadline_ms = 0"
1310        );
1311    }
1312
1313    #[test]
1314    fn classifier_config_parses_a_judge_deadline() {
1315        let config: TaskClassifierConfig = serde_json::from_value(serde_json::json!({
1316            "base_threshold": 0.5,
1317            "judge_deadline_ms": 250,
1318        }))
1319        .expect("a configured judge deadline should parse");
1320        let CapabilityJudgeConfig::Llm(judge) = config.judge else {
1321            panic!("expected the LLM judge variant");
1322        };
1323        assert_eq!(judge.judge_deadline_ms, Some(250));
1324
1325        let config: TaskClassifierConfig = serde_json::from_value(serde_json::json!({
1326            "base_threshold": 0.5,
1327        }))
1328        .expect("an unset judge deadline should parse");
1329        let CapabilityJudgeConfig::Llm(judge) = config.judge else {
1330            panic!("expected the LLM judge variant");
1331        };
1332        assert_eq!(judge.judge_deadline_ms, None);
1333    }
1334
1335    #[tokio::test]
1336    async fn classifier_judges_each_request_without_affinity() -> Result<()> {
1337        let recorder = Arc::new(Recorder::default());
1338        let router = router()?;
1339        let request = classify_request();
1340        let models = runtime_models();
1341
1342        test_drive_with_models(
1343            router.clone(),
1344            request.clone(),
1345            models.clone(),
1346            recorder.serve(),
1347        )
1348        .await?;
1349        test_drive_with_models(router, request, models, recorder.serve()).await?;
1350
1351        assert_eq!(
1352            recorder.calls(),
1353            vec!["judge", "efficient", "judge", "efficient"]
1354        );
1355        assert_eq!(
1356            recorder.call_roles(),
1357            vec![
1358                ("judge".to_string(), false),
1359                ("efficient".to_string(), true),
1360                ("judge".to_string(), false),
1361                ("efficient".to_string(), true),
1362            ]
1363        );
1364        Ok(())
1365    }
1366
1367    #[tokio::test]
1368    async fn classifier_config_sets_the_judge_completion_cap() -> Result<()> {
1369        let recorder = Arc::new(Recorder::default());
1370        let router = Arc::new(LlmTaskClassifier::new(LlmClassifierConfig::Capability {
1371            config: TaskClassifierConfig {
1372                judge: CapabilityJudgeConfig::Llm(LlmCapabilityConfig {
1373                    max_output_tokens: 512,
1374                    ..llm_config(TEST_THRESHOLD)
1375                }),
1376                ..test_config(TEST_THRESHOLD)
1377            },
1378        })?);
1379
1380        test_drive_with_models(
1381            router,
1382            classify_request(),
1383            runtime_models(),
1384            recorder.serve(),
1385        )
1386        .await?;
1387
1388        assert_eq!(recorder.judge_max_output_tokens(), vec![Some(512)]);
1389        Ok(())
1390    }
1391
1392    #[tokio::test]
1393    async fn classifier_config_overrides_the_packaged_prompt() -> Result<()> {
1394        let recorder = Arc::new(Recorder::default());
1395        let router = Arc::new(LlmTaskClassifier::new(LlmClassifierConfig::Capability {
1396            config: TaskClassifierConfig {
1397                judge: CapabilityJudgeConfig::Llm(LlmCapabilityConfig {
1398                    contract: ClassifierContractConfig::default()
1399                        .with_prompt("Custom capability rubric."),
1400                    ..llm_config(TEST_THRESHOLD)
1401                }),
1402                ..test_config(TEST_THRESHOLD)
1403            },
1404        })?);
1405
1406        test_drive_with_models(
1407            router,
1408            classify_request(),
1409            runtime_models(),
1410            recorder.serve(),
1411        )
1412        .await?;
1413
1414        let prompts = recorder.judge_system_prompts();
1415        assert_eq!(prompts.len(), 1);
1416        assert_eq!(prompts[0], "Custom capability rubric.");
1417        Ok(())
1418    }
1419
1420    #[tokio::test]
1421    async fn classifier_config_enables_new_session_trigger() -> Result<()> {
1422        let recorder = Arc::new(Recorder::default());
1423        let router = Arc::new(LlmTaskClassifier::new(LlmClassifierConfig::Capability {
1424            config: TaskClassifierConfig {
1425                classify_trigger: ClassifyTrigger::NewSession,
1426                ..test_config(TEST_THRESHOLD)
1427            },
1428        })?);
1429
1430        let request = classify_session_request();
1431        let models = runtime_models();
1432        test_drive_with_models(
1433            router.clone(),
1434            request.clone(),
1435            models.clone(),
1436            recorder.serve(),
1437        )
1438        .await?;
1439        test_drive_with_models(router, request, models, recorder.serve()).await?;
1440
1441        assert_eq!(recorder.calls(), vec!["judge", "efficient", "efficient"]);
1442        Ok(())
1443    }
1444
1445    #[tokio::test]
1446    async fn classifier_config_reuses_message_hash_affinity_for_a_follow_up() -> Result<()> {
1447        let recorder = Arc::new(Recorder::default());
1448        let router = Arc::new(LlmTaskClassifier::new(LlmClassifierConfig::Capability {
1449            config: TaskClassifierConfig {
1450                classify_trigger: ClassifyTrigger::NewSession,
1451                message_hash_fallback: true,
1452                recent_turn_window: None,
1453                ..test_config(TEST_THRESHOLD)
1454            },
1455        })?);
1456
1457        let models = runtime_models();
1458        test_drive_with_models(
1459            router.clone(),
1460            classify_request(),
1461            models.clone(),
1462            recorder.serve(),
1463        )
1464        .await?;
1465        test_drive_with_models(
1466            router,
1467            classify_follow_up_request(),
1468            models,
1469            recorder.serve(),
1470        )
1471        .await?;
1472
1473        assert_eq!(recorder.calls(), vec!["judge", "efficient", "efficient"]);
1474        Ok(())
1475    }
1476
1477    #[tokio::test]
1478    async fn one_classifier_uses_each_requests_runtime_models() -> Result<()> {
1479        let router = Arc::new(LlmTaskClassifier::new(LlmClassifierConfig::Capability {
1480            config: TaskClassifierConfig {
1481                classify_trigger: ClassifyTrigger::NewSession,
1482                ..test_config(TEST_THRESHOLD)
1483            },
1484        })?);
1485        let calls = Arc::new(Mutex::new(Vec::new()));
1486        let serve = |calls: Arc<Mutex<Vec<String>>>| {
1487            move |model: ModelId, _request: Request| {
1488                let calls = Arc::clone(&calls);
1489                async move {
1490                    calls.lock().push(model.to_string());
1491                    let text = if model.as_str().starts_with("judge-") {
1492                        r#"{"crux":"bounded task","primary_rule":"SUP-1","capability_boundary":"supported","p_solve":0.9}"#.to_string()
1493                    } else {
1494                        model.to_string()
1495                    };
1496                    Ok(Response {
1497                        llm_response: LlmResponse::Agg(text_response(None, text)),
1498                        metadata: None,
1499                        upstream_headers: Default::default(),
1500                    })
1501                }
1502            }
1503        };
1504        let models = |suffix: &str| -> HashMap<Category, Vec<ModelId>> {
1505            [
1506                (
1507                    Category::Judge,
1508                    vec![ModelId::from(format!("judge-{suffix}"))],
1509                ),
1510                (
1511                    Category::Efficient,
1512                    vec![ModelId::from(format!("efficient-{suffix}"))],
1513                ),
1514                (
1515                    Category::Capable,
1516                    vec![ModelId::from(format!("capable-{suffix}"))],
1517                ),
1518                (
1519                    Category::Any,
1520                    vec![
1521                        ModelId::from(format!("efficient-{suffix}")),
1522                        ModelId::from(format!("capable-{suffix}")),
1523                    ],
1524                ),
1525            ]
1526            .into()
1527        };
1528        let request = classify_session_request();
1529
1530        let (first, _) = test_drive_with_models(
1531            router.clone(),
1532            request.clone(),
1533            models("a"),
1534            serve(Arc::clone(&calls)),
1535        )
1536        .await?;
1537        let (second, _) =
1538            test_drive_with_models(router, request, models("b"), serve(Arc::clone(&calls))).await?;
1539
1540        assert_eq!(first, "efficient-a");
1541        assert_eq!(second, "efficient-b");
1542        assert_eq!(
1543            &*calls.lock(),
1544            &["judge-a", "efficient-a", "judge-b", "efficient-b"]
1545        );
1546        Ok(())
1547    }
1548
1549    #[tokio::test]
1550    async fn decision_judge_routes_on_capable_probability() -> Result<()> {
1551        use crate::{Call, RuntimeModels, drive};
1552        use serde_json::json;
1553        use std::collections::BTreeMap;
1554        use std::sync::atomic::{AtomicUsize, Ordering};
1555        use switchyard_protocol::{
1556            BooleanEstimate, DecisionAnswer, DecisionResponse, DecisionValue, Probability,
1557            ProviderConfidence,
1558        };
1559
1560        let settings = DecisionJudgeConfig {
1561            cutoff: Some(0.4),
1562            ..DecisionJudgeConfig::default()
1563        };
1564        let mut request = classify_session_request();
1565        request.llm_request.messages.push(Message::new(
1566            Role::Assistant,
1567            vec![
1568                ContentBlock::Reasoning {
1569                    text: "private reasoning".into(),
1570                    signature: None,
1571                    details: vec![],
1572                },
1573                ContentBlock::Text {
1574                    text: "visible progress".into(),
1575                },
1576            ],
1577        ));
1578        request
1579            .llm_request
1580            .messages
1581            .push(Message::text(Role::User, "latest follow-up"));
1582        let original_messages = request.llm_request.messages.clone();
1583        let expected_task = json!([
1584            original_messages[0].clone(),
1585            Message::text(Role::Assistant, "visible progress"),
1586            Message::text(Role::User, "latest follow-up"),
1587        ]);
1588
1589        for (name, score, fail_open, expected) in [
1590            ("above", 0.45, true, Some("capable")),
1591            ("equal", 0.4, true, Some("efficient")),
1592            ("below", 0.39, false, Some("efficient")),
1593            ("missing answer", 0.0, true, Some("capable")),
1594            ("no distribution", 0.0, true, Some("capable")),
1595            ("wrong type", 0.0, true, Some("capable")),
1596            ("out of range", 1.1, true, Some("capable")),
1597            ("nonfinite", f64::NAN, true, Some("capable")),
1598            ("missing probability", 0.0, true, Some("capable")),
1599            ("missing answer", 0.0, false, None),
1600            ("no distribution", 0.0, false, None),
1601            ("wrong type", 0.0, false, None),
1602            ("out of range", 1.1, false, None),
1603            ("negative", -0.1, false, None),
1604            ("nonfinite", f64::NAN, false, None),
1605            ("infinite", f64::INFINITY, false, None),
1606            ("missing probability", 0.0, false, None),
1607            ("provider error", 0.0, true, Some("capable")),
1608            ("provider error", 0.0, false, None),
1609            ("dropped reply", 0.0, true, Some("capable")),
1610            ("custom choices", 0.45, false, Some("capable")),
1611        ] {
1612            let mut config = settings.clone();
1613            if name == "below" {
1614                config.instructions = Some("Choose capable when the task needs it.".into());
1615            }
1616            if name == "custom choices" {
1617                config.choices[0].id = "complex".into();
1618                config.choices[1].id = "simple".into();
1619                config.apply_cutoff_to = "complex".into();
1620            }
1621            let override_instructions = config.instructions.clone();
1622            let choice_ids = [config.choices[0].id.clone(), config.choices[1].id.clone()];
1623            let router: Arc<dyn Algorithm> =
1624                Arc::new(LlmTaskClassifier::new(LlmClassifierConfig::Capability {
1625                    config: TaskClassifierConfig {
1626                        judge: CapabilityJudgeConfig::Decision(config),
1627                        fail_open,
1628                        recent_turn_window: None,
1629                        classify_trigger: ClassifyTrigger::NewSession,
1630                        ..TaskClassifierConfig::default()
1631                    },
1632                })?);
1633            let calls = AtomicUsize::new(0);
1634            let serve = |call| {
1635                let calls = &calls;
1636                let expected_task = &expected_task;
1637                let override_instructions = &override_instructions;
1638                let [capable_id, efficient_id] = choice_ids.clone();
1639                async move {
1640                    let Call::Decision(call) = call else {
1641                        panic!("judge must use the decision step")
1642                    };
1643                    calls.fetch_add(1, Ordering::SeqCst);
1644                    assert_eq!(call.model, "judge");
1645                    assert_eq!(call.request.model, Some("judge".into()));
1646                    assert_eq!(
1647                        serde_json::to_value(&call.request.input).unwrap(),
1648                        *expected_task
1649                    );
1650                    assert_eq!(call.request.questions.len(), 1);
1651                    let question = &call.request.questions["route"];
1652                    let switchyard_protocol::DecisionKind::Choice { options } = &question.kind
1653                    else {
1654                        panic!("expected a choice")
1655                    };
1656                    assert_eq!(
1657                        options
1658                            .iter()
1659                            .map(|option| option.id.as_str())
1660                            .collect::<Vec<_>>(),
1661                        [capable_id.as_str(), efficient_id.as_str()]
1662                    );
1663                    if let Some(instructions) = override_instructions {
1664                        assert_eq!(&question.instructions, instructions);
1665                    } else {
1666                        assert!(
1667                            question
1668                                .instructions
1669                                .as_str()
1670                                .is_some_and(|text| !text.is_empty())
1671                        );
1672                    }
1673                    if name == "provider error" {
1674                        return call
1675                            .respond(std::future::ready(Err(LibsyError::client_call(
1676                                "judge",
1677                                LlmClientError::General("private provider body".into()),
1678                            ))))
1679                            .await;
1680                    }
1681                    if name == "dropped reply" {
1682                        drop(call);
1683                        return Ok(());
1684                    }
1685                    let value = if name == "wrong type" {
1686                        DecisionValue::Boolean(BooleanEstimate::Value(true))
1687                    } else {
1688                        DecisionValue::Choice {
1689                            selected: efficient_id.clone(),
1690                            probabilities: (name != "no distribution").then(|| {
1691                                BTreeMap::from([
1692                                    (
1693                                        if name == "missing probability" {
1694                                            "other".into()
1695                                        } else {
1696                                            capable_id.clone()
1697                                        },
1698                                        Probability(score),
1699                                    ),
1700                                    (efficient_id.clone(), Probability(1.0 - score)),
1701                                ])
1702                            }),
1703                        }
1704                    };
1705                    let answers = if name == "missing answer" {
1706                        BTreeMap::new()
1707                    } else {
1708                        BTreeMap::from([(
1709                            "route".into(),
1710                            DecisionAnswer {
1711                                value,
1712                                provider_confidence: Some(ProviderConfidence(0.99)),
1713                            },
1714                        )])
1715                    };
1716                    call.respond(std::future::ready(Ok(DecisionResponse {
1717                        id: None,
1718                        model: Some("provider-judge".into()),
1719                        answers,
1720                        usage: Default::default(),
1721                    })))
1722                    .await
1723                }
1724            };
1725            let models = Arc::new(RuntimeModels::new(runtime_models()));
1726            let result = drive(router.clone(), request.clone(), models.clone(), &serve).await;
1727            assert_eq!(calls.load(Ordering::SeqCst), 1, "{name}");
1728            let Some(expected) = expected else {
1729                if name == "provider error" {
1730                    assert!(
1731                        matches!(result, Err(LibsyError::ClientCall { .. })),
1732                        "{name}"
1733                    );
1734                } else {
1735                    assert!(
1736                        matches!(result, Err(LibsyError::InvalidDecision { target, .. }) if target == "judge"),
1737                        "expected invalid decision error: {name}"
1738                    );
1739                }
1740                continue;
1741            };
1742            let outcome = result?;
1743            assert_eq!(outcome.selected_model_id()?, expected, "{name}");
1744            assert!(outcome.response.is_none());
1745            assert_eq!(outcome.request.llm_request.messages, original_messages);
1746            let metadata = outcome.metadata.expect("routing metadata");
1747            assert_eq!(metadata.judgments.len(), 1);
1748            let evidence = metadata.evidence.expect("routing evidence");
1749            if matches!(name, "above" | "equal" | "below" | "custom choices") {
1750                assert_eq!(metadata.judgments[0].target.as_deref(), Some(expected));
1751                assert_eq!(evidence["source"], "decision_classifier");
1752                assert_eq!(evidence["verdict"], expected);
1753                assert_eq!(evidence["capable_threshold"], 0.4);
1754                assert_eq!(evidence["capable_score"], score);
1755            } else {
1756                let reason = match name {
1757                    "provider error" => "client_error",
1758                    "dropped reply" => "call_error",
1759                    _ => "invalid_verdict",
1760                };
1761                assert_eq!(
1762                    evidence,
1763                    json!({"source": "fail_open", "reason_code": reason})
1764                );
1765            }
1766            let retained = drive(router, request.clone(), models, &serve).await?;
1767            assert_eq!(retained.selected_model_id()?, expected);
1768            assert!(
1769                retained
1770                    .metadata
1771                    .expect("retained metadata")
1772                    .judgments
1773                    .is_empty()
1774            );
1775            assert_eq!(
1776                calls.load(Ordering::SeqCst),
1777                1,
1778                "affinity should skip the judge: {name}"
1779            );
1780        }
1781        for cutoff in [-0.1, 1.1, f64::NAN] {
1782            let mut config = settings.clone();
1783            config.cutoff = Some(cutoff);
1784            assert!(
1785                LlmTaskClassifier::new(LlmClassifierConfig::Capability {
1786                    config: TaskClassifierConfig {
1787                        judge: CapabilityJudgeConfig::Decision(config),
1788                        ..TaskClassifierConfig::default()
1789                    },
1790                })
1791                .is_err()
1792            );
1793        }
1794        Ok(())
1795    }
1796
1797    #[tokio::test]
1798    async fn direct_judge_falls_back_to_the_default_target() -> Result<()> {
1799        use crate::{Call, RuntimeModels, drive};
1800        use std::collections::BTreeMap;
1801        use switchyard_protocol::{DecisionAnswer, DecisionResponse, DecisionValue};
1802
1803        for (reply, fail_open, expected) in [
1804            (Some("unknown"), true, Some("efficient")),
1805            (Some("unknown"), false, None),
1806            (None, true, Some("efficient")),
1807            (None, false, None),
1808        ] {
1809            let router: Arc<dyn Algorithm> =
1810                Arc::new(LlmTaskClassifier::new(LlmClassifierConfig::Direct {
1811                    default_target: Category::Efficient,
1812                    config: DirectClassifierConfig {
1813                        options: vec![Category::Efficient, Category::Capable],
1814                        fail_open,
1815                        classify_trigger: ClassifyTrigger::EveryRequest,
1816                        message_hash_fallback: false,
1817                        recent_turn_window: None,
1818                    },
1819                })?);
1820            let serve = |call| async move {
1821                let Call::Decision(call) = call else {
1822                    panic!("judge must use the decision step")
1823                };
1824                let response = match reply {
1825                    Some(selected) => Ok(DecisionResponse {
1826                        id: None,
1827                        model: None,
1828                        answers: BTreeMap::from([(
1829                            "route".into(),
1830                            DecisionAnswer {
1831                                value: DecisionValue::Choice {
1832                                    selected: selected.into(),
1833                                    probabilities: None,
1834                                },
1835                                provider_confidence: None,
1836                            },
1837                        )]),
1838                        usage: Default::default(),
1839                    }),
1840                    None => Err(LibsyError::client_call(
1841                        "judge",
1842                        LlmClientError::General("unavailable".into()),
1843                    )),
1844                };
1845                call.respond(std::future::ready(response)).await
1846            };
1847            let models = Arc::new(RuntimeModels::new(runtime_models()));
1848            let result = drive(router, classify_request(), models, &serve).await;
1849            match expected {
1850                Some(expected) => assert_eq!(result?.selected_model_id()?, expected),
1851                None => assert!(result.is_err()),
1852            }
1853        }
1854        Ok(())
1855    }
1856
1857    #[test]
1858    fn the_threshold_boundary_is_inclusive() -> Result<()> {
1859        let policy = policy();
1860        let at_threshold = verdict(0.5, "supported", "SUP-1");
1861        let below_threshold = verdict(0.49, "supported", "SUP-1");
1862        assert_eq!(selected(&policy, Some(&at_threshold))?, "efficient");
1863        assert_eq!(selected(&policy, Some(&below_threshold))?, "capable");
1864        Ok(())
1865    }
1866
1867    #[test]
1868    fn the_threshold_moves_the_routing_boundary() -> Result<()> {
1869        let borderline = verdict(0.5, "supported", "SUP-1");
1870        let strict = TaskClassifierPolicy::new(&llm_config(0.9));
1871        let lenient = TaskClassifierPolicy::new(&llm_config(0.1));
1872        assert_eq!(selected(&strict, Some(&borderline))?, "capable");
1873        assert_eq!(selected(&lenient, Some(&borderline))?, "efficient");
1874        Ok(())
1875    }
1876
1877    #[test]
1878    fn classifier_config_rejects_unknown_fields() {
1879        let error = serde_json::from_value::<TaskClassifierConfig>(serde_json::json!({
1880            "base_threshold": 0.5,
1881            "classifier_magic": true,
1882        }))
1883        .expect_err("unknown classifier fields must be rejected");
1884
1885        assert!(
1886            error
1887                .to_string()
1888                .contains("unknown field `classifier_magic`"),
1889            "{error}"
1890        );
1891    }
1892
1893    #[test]
1894    fn invalid_classifier_config_is_rejected() -> Result<()> {
1895        for bad in [1.5, -0.1, f64::NAN, f64::INFINITY] {
1896            assert!(
1897                LlmTaskClassifier::new(LlmClassifierConfig::Capability {
1898                    config: test_config(bad),
1899                })
1900                .is_err(),
1901                "base threshold {bad} should be rejected"
1902            );
1903        }
1904        for config in [
1905            TaskClassifierConfig {
1906                judge: CapabilityJudgeConfig::Llm(LlmCapabilityConfig {
1907                    base_threshold: 0.5,
1908                    threshold_step: -0.1,
1909                    ..LlmCapabilityConfig::default()
1910                }),
1911                ..TaskClassifierConfig::default()
1912            },
1913            TaskClassifierConfig {
1914                judge: CapabilityJudgeConfig::Llm(LlmCapabilityConfig {
1915                    base_threshold: 0.8,
1916                    threshold_step: 0.11,
1917                    ..LlmCapabilityConfig::default()
1918                }),
1919                ..TaskClassifierConfig::default()
1920            },
1921            TaskClassifierConfig {
1922                judge: CapabilityJudgeConfig::Llm(LlmCapabilityConfig {
1923                    base_threshold: 0.5,
1924                    ..LlmCapabilityConfig::default()
1925                }),
1926                message_hash_fallback: true,
1927                ..TaskClassifierConfig::default()
1928            },
1929            TaskClassifierConfig {
1930                judge: CapabilityJudgeConfig::Llm(LlmCapabilityConfig {
1931                    base_threshold: 0.5,
1932                    max_output_tokens: 0,
1933                    ..LlmCapabilityConfig::default()
1934                }),
1935                ..TaskClassifierConfig::default()
1936            },
1937        ] {
1938            assert!(LlmTaskClassifier::new(LlmClassifierConfig::Capability { config }).is_err());
1939        }
1940        for base_threshold in [0.0, 1.0] {
1941            LlmTaskClassifier::new(LlmClassifierConfig::Capability {
1942                config: test_config(base_threshold),
1943            })?;
1944        }
1945        Ok(())
1946    }
1947
1948    #[test]
1949    fn message_hash_fallback_accepts_retaining_triggers() -> Result<()> {
1950        // Only `every_request` is rejected: it retains no target, so a fallback
1951        // identity has nothing to key on. Both retaining triggers may key the
1952        // retained target on a message hash when the caller sends no session id.
1953        for trigger in [ClassifyTrigger::NewSession, ClassifyTrigger::UserTurn] {
1954            let config = TaskClassifierConfig {
1955                judge: CapabilityJudgeConfig::Llm(LlmCapabilityConfig {
1956                    base_threshold: 0.5,
1957                    ..LlmCapabilityConfig::default()
1958                }),
1959                classify_trigger: trigger,
1960                message_hash_fallback: true,
1961                ..TaskClassifierConfig::default()
1962            };
1963            LlmTaskClassifier::new(LlmClassifierConfig::Capability { config }).map_err(
1964                |error| LibsyError::AlgorithmError {
1965                    message: format!("{trigger:?} with message_hash_fallback rejected: {error}"),
1966                },
1967            )?;
1968        }
1969        let every_request = TaskClassifierConfig {
1970            judge: CapabilityJudgeConfig::Llm(LlmCapabilityConfig {
1971                base_threshold: 0.5,
1972                ..LlmCapabilityConfig::default()
1973            }),
1974            classify_trigger: ClassifyTrigger::EveryRequest,
1975            message_hash_fallback: true,
1976            ..TaskClassifierConfig::default()
1977        };
1978        assert!(
1979            LlmTaskClassifier::new(LlmClassifierConfig::Capability {
1980                config: every_request
1981            })
1982            .is_err(),
1983            "every_request with message_hash_fallback should stay rejected"
1984        );
1985        Ok(())
1986    }
1987
1988    #[test]
1989    fn an_unusable_verdict_is_ambiguous() -> Result<()> {
1990        let policy = policy();
1991        let inconsistent_rule = TaskClassifierVerdict {
1992            capability_boundary: "uncertain".to_string(),
1993            ..verdict(1.0, "supported", "SUP-1")
1994        };
1995        let empty_crux = TaskClassifierVerdict {
1996            crux: "  ".to_string(),
1997            ..verdict(1.0, "supported", "SUP-1")
1998        };
1999        let unusable = [
2000            Some(verdict(1.1, "supported", "SUP-1")),
2001            Some(inconsistent_rule),
2002            Some(empty_crux),
2003            None,
2004        ];
2005        for verdict in unusable {
2006            let classification = policy.to_classification(verdict.as_ref(), &policy_driver())?;
2007            assert!(matches!(classification, Classification::Ambiguous(_)));
2008            assert!(classification.argmax(false)?.is_none());
2009            assert!(classification.argmax(true)?.is_none());
2010        }
2011        Ok(())
2012    }
2013
2014    #[test]
2015    fn capability_boundaries_apply_monotonic_threshold_steps() -> Result<()> {
2016        let policy = TaskClassifierPolicy::new(&LlmCapabilityConfig {
2017            threshold_step: 0.1,
2018            ..llm_config(0.4)
2019        });
2020
2021        assert_eq!(
2022            selected(&policy, Some(&verdict(0.4, "supported", "SUP-2")))?,
2023            "efficient"
2024        );
2025        assert_eq!(
2026            selected(&policy, Some(&verdict(0.49, "uncertain", "UNC-1")))?,
2027            "capable"
2028        );
2029        assert_eq!(
2030            selected(&policy, Some(&verdict(0.5, "uncertain", "UNC-1")))?,
2031            "efficient"
2032        );
2033        assert_eq!(
2034            selected(&policy, Some(&verdict(0.5, "unmatched", "none")))?,
2035            "efficient"
2036        );
2037        assert_eq!(
2038            selected(&policy, Some(&verdict(0.59, "unsupported", "LIM-1")))?,
2039            "capable"
2040        );
2041        assert_eq!(
2042            selected(&policy, Some(&verdict(0.6, "unsupported", "LIM-1")))?,
2043            "efficient"
2044        );
2045        Ok(())
2046    }
2047
2048    #[test]
2049    fn unmatched_steps_raise_only_the_unmatched_threshold() -> Result<()> {
2050        let policy = TaskClassifierPolicy::new(&LlmCapabilityConfig {
2051            threshold_step: 0.1,
2052            unmatched_steps: 2,
2053            ..llm_config(0.75)
2054        });
2055
2056        // 0.85 clears one step but not two: unmatched now routes to the capable target.
2057        assert_eq!(
2058            selected(&policy, Some(&verdict(0.85, "unmatched", "none")))?,
2059            "capable"
2060        );
2061        assert_eq!(
2062            selected(&policy, Some(&verdict(0.95, "unmatched", "none")))?,
2063            "efficient"
2064        );
2065        // Other boundaries keep their steps.
2066        assert_eq!(
2067            selected(&policy, Some(&verdict(0.85, "uncertain", "UNC-1")))?,
2068            "efficient"
2069        );
2070        assert_eq!(
2071            selected(&policy, Some(&verdict(0.75, "supported", "SUP-2")))?,
2072            "efficient"
2073        );
2074        Ok(())
2075    }
2076
2077    #[test]
2078    fn unmatched_steps_default_to_one_and_reject_values_above_two() {
2079        let parsed: TaskClassifierConfig =
2080            serde_json::from_value(serde_json::json!({ "base_threshold": 0.5 }))
2081                .expect("config without unmatched_steps parses");
2082        let CapabilityJudgeConfig::Llm(judge) = &parsed.judge else {
2083            panic!("a flat capability config builds an LLM judge");
2084        };
2085        assert_eq!(judge.unmatched_steps, 1);
2086
2087        let invalid = TaskClassifierConfig {
2088            judge: CapabilityJudgeConfig::Llm(LlmCapabilityConfig {
2089                unmatched_steps: 3,
2090                ..llm_config(0.5)
2091            }),
2092            ..TaskClassifierConfig::default()
2093        };
2094        assert!(invalid.validate().is_err());
2095    }
2096
2097    /// The text of each message a judge with `recent_turn_window` would be sent.
2098    /// The no-window case is covered by `capability_judge_builds_a_structured_request`.
2099    fn capability_judge(recent_turn_window: Option<usize>) -> Result<CapabilityJudge> {
2100        Ok(StructuredJudge::new(
2101            TaskInput { recent_turn_window },
2102            LlmTaskClassifier::load_capability_contract(&ClassifierContractConfig::default())?,
2103            SerdeDecoder::new(),
2104            JudgeRuntimeConfig::new(DEFAULT_JUDGE_MAX_OUTPUT_TOKENS)?,
2105        ))
2106    }
2107
2108    fn judged_contents(recent_turn_window: usize) -> Result<Vec<String>> {
2109        let judge = capability_judge(Some(recent_turn_window))?;
2110        let request = Request {
2111            llm_request: LlmRequest {
2112                messages: vec![
2113                    Message::text(Role::System, "client instructions"),
2114                    Message::text(
2115                        Role::User,
2116                        "<environment_context>setup</environment_context>",
2117                    ),
2118                    Message::text(Role::User, "initial task"),
2119                    Message::text(Role::Assistant, "old response"),
2120                    Message::text(Role::User, "old follow-up"),
2121                    Message::text(Role::Assistant, "recent 1"),
2122                    Message::text(Role::User, "recent 2"),
2123                ],
2124                ..LlmRequest::default()
2125            },
2126            raw_request: None,
2127            metadata: None,
2128        };
2129        Ok(judge
2130            .build_request(&State::default(), &request)
2131            .llm_request
2132            .messages
2133            .iter()
2134            .filter_map(|message| message.text_content("\n"))
2135            .collect())
2136    }
2137
2138    #[test]
2139    fn a_window_widens_the_judge_to_the_surrounding_conversation() -> Result<()> {
2140        // Client instructions and the opening task, plus the last two turns.
2141        let contents = judged_contents(2)?;
2142        assert!(contents.contains(&"client instructions".to_string()));
2143        assert!(contents.contains(&"initial task".to_string()));
2144        assert!(contents.contains(&"recent 1".to_string()));
2145        assert!(contents.contains(&"recent 2".to_string()));
2146        assert!(!contents.contains(&"old response".to_string()));
2147        Ok(())
2148    }
2149
2150    #[test]
2151    fn a_zero_window_keeps_only_the_instructions_and_the_task() -> Result<()> {
2152        let contents = judged_contents(0)?;
2153        assert!(contents.contains(&"client instructions".to_string()));
2154        assert!(contents.contains(&"initial task".to_string()));
2155        assert!(!contents.contains(&"recent 2".to_string()));
2156        Ok(())
2157    }
2158
2159    fn tool_call(id: &str) -> Message {
2160        Message::new(
2161            Role::Assistant,
2162            vec![ContentBlock::ToolCall(ToolCall {
2163                id: id.to_string(),
2164                name: "search".to_string(),
2165                arguments: Value::Null,
2166            })],
2167        )
2168    }
2169
2170    fn tool_result(id: &str) -> Message {
2171        Message::new(
2172            Role::Tool,
2173            vec![ContentBlock::ToolResult(ToolResult {
2174                tool_call_id: id.to_string(),
2175                content: vec![ContentBlock::Text {
2176                    text: "tool output".to_string(),
2177                }],
2178                is_error: None,
2179            })],
2180        )
2181    }
2182
2183    #[test]
2184    fn default_task_input_keeps_user_content_around_tool_results() {
2185        let context = Message::text(
2186            Role::User,
2187            "<environment_context><current_date>2026-10-31</current_date></environment_context>",
2188        );
2189        let mut result = tool_result("call-1");
2190        result.role = Role::User;
2191        let mut mixed = result.clone();
2192        mixed.content.push(ContentBlock::Text {
2193            text: "latest follow-up".to_string(),
2194        });
2195        let input = TaskInput {
2196            recent_turn_window: None,
2197        };
2198        let mut request = Request {
2199            llm_request: LlmRequest {
2200                messages: vec![
2201                    context.clone(),
2202                    result.clone(),
2203                    Message::text(Role::User, "initial task"),
2204                    tool_call("call-1"),
2205                    mixed,
2206                    result.clone(),
2207                    context,
2208                ],
2209                ..LlmRequest::default()
2210            },
2211            ..Request::default()
2212        };
2213        assert_eq!(
2214            input.build_messages(&State::default(), &request),
2215            vec![
2216                Message::text(Role::User, "initial task"),
2217                Message::text(Role::User, "latest follow-up"),
2218            ]
2219        );
2220        request.llm_request.messages = vec![result];
2221        assert!(input.build_messages(&State::default(), &request).is_empty());
2222    }
2223
2224    /// A count-based window can begin on a tool result, which leaves the call that
2225    /// introduced its id outside the window and the classifier history invalid.
2226    #[test]
2227    fn trimming_keeps_the_call_that_introduced_a_kept_tool_result() {
2228        let messages = vec![
2229            Message::text(Role::System, "client instructions"),
2230            Message::text(Role::User, "initial task"),
2231            Message::text(Role::Assistant, "old response"),
2232            tool_call("call-1"),
2233            tool_result("call-1"),
2234            Message::text(Role::Assistant, "recent 1"),
2235            Message::text(Role::User, "recent 2"),
2236            Message::text(Role::Assistant, "recent 3"),
2237            Message::text(Role::User, "recent 4"),
2238        ];
2239
2240        // The five-message tail begins exactly on the tool result.
2241        let kept = trim_messages(&messages, 5);
2242
2243        assert_eq!(
2244            kept,
2245            vec![
2246                Message::text(Role::System, "client instructions"),
2247                Message::text(Role::User, "initial task"),
2248                tool_call("call-1"),
2249                tool_result("call-1"),
2250                Message::text(Role::Assistant, "recent 1"),
2251                Message::text(Role::User, "recent 2"),
2252                Message::text(Role::Assistant, "recent 3"),
2253                Message::text(Role::User, "recent 4"),
2254            ]
2255        );
2256    }
2257
2258    /// Ids repeat across a conversation, so a later call must not stand in for the one that
2259    /// answers an earlier result.
2260    #[test]
2261    fn trimming_pairs_a_repeated_id_with_the_call_that_precedes_it() {
2262        let messages = vec![
2263            Message::text(Role::System, "client instructions"),
2264            Message::text(Role::User, "initial task"),
2265            tool_call("x"),
2266            tool_result("x"),
2267            Message::text(Role::Assistant, "later"),
2268            tool_call("x"),
2269            tool_result("x"),
2270        ];
2271
2272        // The four-message tail begins on the first result, whose own call sits one earlier.
2273        let kept = trim_messages(&messages, 4);
2274
2275        assert_eq!(
2276            kept,
2277            vec![
2278                Message::text(Role::System, "client instructions"),
2279                Message::text(Role::User, "initial task"),
2280                tool_call("x"),
2281                tool_result("x"),
2282                Message::text(Role::Assistant, "later"),
2283                tool_call("x"),
2284                tool_result("x"),
2285            ]
2286        );
2287    }
2288
2289    /// A result whose call precedes the opening task can never be paired, because trimming
2290    /// never reaches behind the task. The window must not widen hunting for it.
2291    #[test]
2292    fn trimming_keeps_the_counted_window_when_a_result_cannot_be_paired() {
2293        let messages = vec![
2294            Message::text(Role::System, "client instructions"),
2295            tool_call("orphan"),
2296            Message::text(Role::User, "initial task"),
2297            Message::text(Role::Assistant, "old response"),
2298            tool_result("orphan"),
2299            Message::text(Role::Assistant, "recent 1"),
2300            Message::text(Role::User, "recent 2"),
2301        ];
2302
2303        let kept = trim_messages(&messages, 3);
2304
2305        assert_eq!(
2306            kept,
2307            vec![
2308                Message::text(Role::System, "client instructions"),
2309                Message::text(Role::User, "initial task"),
2310                tool_result("orphan"),
2311                Message::text(Role::Assistant, "recent 1"),
2312                Message::text(Role::User, "recent 2"),
2313            ]
2314        );
2315    }
2316
2317    #[test]
2318    fn a_window_restates_the_routing_instruction_last() -> Result<()> {
2319        let contents = judged_contents(2)?;
2320
2321        // Last, not merely present: the instruction only works in final position,
2322        // after the conversation content it is meant to outrank.
2323        assert_eq!(
2324            contents.last().map(String::as_str),
2325            Some(TRAILING_ROUTING_INSTRUCTION)
2326        );
2327        Ok(())
2328    }
2329
2330    /// Reasoning is provider-private and some upstreams reject it replayed without the
2331    /// opaque state it was issued with, so it must not reach a windowed classifier
2332    /// request. Visible assistant text and complete tool pairs still do, and a turn
2333    /// whose only content was reasoning is dropped rather than left behind empty.
2334    #[test]
2335    fn a_window_drops_reasoning_but_keeps_visible_text_and_tool_pairs() {
2336        let messages = vec![
2337            Message::text(Role::System, "client instructions"),
2338            Message::text(Role::User, "initial task"),
2339            Message::new(
2340                Role::Assistant,
2341                vec![
2342                    ContentBlock::Reasoning {
2343                        text: "private chain of thought".to_string(),
2344                        signature: None,
2345                        details: Vec::new(),
2346                    },
2347                    ContentBlock::Text {
2348                        text: "visible answer".to_string(),
2349                    },
2350                ],
2351            ),
2352            tool_call("call-1"),
2353            tool_result("call-1"),
2354            Message::new(
2355                Role::Assistant,
2356                vec![ContentBlock::Reasoning {
2357                    text: "reasoning-only turn".to_string(),
2358                    signature: None,
2359                    details: Vec::new(),
2360                }],
2361            ),
2362            Message::text(Role::User, "follow-up"),
2363        ];
2364        let request = Request {
2365            llm_request: LlmRequest {
2366                messages,
2367                ..LlmRequest::default()
2368            },
2369            raw_request: None,
2370            metadata: None,
2371        };
2372
2373        let built = TaskInput {
2374            recent_turn_window: Some(10),
2375        }
2376        .build_messages(&State::default(), &request);
2377
2378        assert!(
2379            !built
2380                .iter()
2381                .flat_map(|message| &message.content)
2382                .any(|block| matches!(block, ContentBlock::Reasoning { .. })),
2383            "{built:?}"
2384        );
2385        assert!(
2386            built
2387                .iter()
2388                .any(|message| message.text_content("\n").as_deref() == Some("visible answer"))
2389        );
2390        assert!(built.contains(&tool_call("call-1")));
2391        assert!(built.contains(&tool_result("call-1")));
2392        // Six of the seven fixtures survive — the reasoning-only turn is gone entirely
2393        // rather than left behind empty — plus the trailing routing instruction.
2394        assert_eq!(built.len(), 7);
2395        assert!(built.iter().all(|message| !message.content.is_empty()));
2396    }
2397
2398    #[test]
2399    fn the_default_path_is_left_unchanged() -> Result<()> {
2400        // No window means no conversation to be distracted by, so the default
2401        // request shape stays exactly as it was.
2402        let judge = capability_judge(None)?;
2403        let request = Request {
2404            llm_request: LlmRequest {
2405                messages: vec![Message::text(Role::User, "the task")],
2406                ..LlmRequest::default()
2407            },
2408            raw_request: None,
2409            metadata: None,
2410        };
2411
2412        let built = judge.build_request(&State::default(), &request);
2413
2414        // The task message alone; the rubric reaches the judge as an instruction block
2415        // rather than a message, so nothing was dropped by it not being counted here.
2416        assert_eq!(built.llm_request.messages.len(), 1);
2417        assert!(!built.llm_request.instructions.is_empty());
2418        assert!(
2419            !built
2420                .llm_request
2421                .messages
2422                .iter()
2423                .filter_map(|message| message.text_content("\n"))
2424                .any(|text| text.contains(TRAILING_ROUTING_INSTRUCTION))
2425        );
2426        Ok(())
2427    }
2428
2429    #[test]
2430    fn capability_judge_builds_a_structured_request() -> Result<()> {
2431        let judge = capability_judge(None)?;
2432        let request = Request {
2433            llm_request: LlmRequest {
2434                model: Some("inbound".to_string()),
2435                messages: vec![
2436                    Message::text(Role::System, "client instructions"),
2437                    Message::text(Role::Developer, "client developer instructions"),
2438                    Message::text(Role::User, "initial task"),
2439                    Message::text(Role::Assistant, "old response"),
2440                    Message::text(Role::User, "old follow-up"),
2441                    Message::text(Role::Assistant, "recent 1"),
2442                    Message::text(Role::User, "recent 2"),
2443                    Message::text(Role::Assistant, "recent 3"),
2444                    Message::text(Role::User, "recent 4"),
2445                    Message::text(Role::Assistant, "recent 5"),
2446                ],
2447                ..LlmRequest::default()
2448            },
2449            raw_request: None,
2450            metadata: None,
2451        };
2452        let judge_request = judge.build_request(&State::default(), &request);
2453
2454        assert_eq!(judge_request.llm_request.model, request.llm_request.model);
2455        assert_eq!(judge_request.llm_request.instructions.len(), 1);
2456        assert_eq!(judge_request.llm_request.instructions[0].role, Role::System);
2457        assert_eq!(
2458            judge_request.llm_request.instructions[0].content,
2459            InstructionBlock {
2460                role: Role::System,
2461                content: Message::text(Role::System, judge.contract().system_prompt()).content,
2462            }
2463            .content,
2464        );
2465        assert_eq!(judge_request.llm_request.messages.len(), 2);
2466        let contents = judge_request
2467            .llm_request
2468            .messages
2469            .iter()
2470            .filter_map(|message| message.text_content("\n"))
2471            .collect::<Vec<_>>();
2472        assert!(contents.contains(&"recent 4".to_string()));
2473        assert!(contents.contains(&"initial task".to_string()));
2474        assert!(!contents.contains(&"recent 5".to_string()));
2475        assert!(!contents.contains(&"client instructions".to_string()));
2476        assert_eq!(
2477            judge_request.llm_request.output.response_format,
2478            Some(judge.contract().response_format().clone())
2479        );
2480        assert_eq!(
2481            judge_request.llm_request.output.max_output_tokens,
2482            Some(DEFAULT_JUDGE_MAX_OUTPUT_TOKENS)
2483        );
2484        Ok(())
2485    }
2486
2487    fn sample_value(spec: &Value) -> Value {
2488        if let Some(first) = spec
2489            .get("enum")
2490            .and_then(Value::as_array)
2491            .and_then(|values| values.first())
2492        {
2493            return first.clone();
2494        }
2495        match spec.get("type").and_then(Value::as_str) {
2496            Some("number") => serde_json::json!(0.5),
2497            Some("boolean") => serde_json::json!(false),
2498            _ => serde_json::json!("sample"),
2499        }
2500    }
2501
2502    fn schema_shaped_verdict(schema: &Value) -> Result<String> {
2503        let properties = schema
2504            .pointer("/json_schema/schema/properties")
2505            .and_then(Value::as_object)
2506            .ok_or_else(|| LibsyError::AlgorithmError {
2507                message: "packaged schema declares no properties".to_string(),
2508            })?;
2509        Ok(Value::Object(
2510            properties
2511                .iter()
2512                .map(|(name, spec)| (name.clone(), sample_value(spec)))
2513                .collect(),
2514        )
2515        .to_string())
2516    }
2517
2518    /// Built from the schema so a property added there fails here rather than silently
2519    /// rejecting every production verdict.
2520    #[test]
2521    fn every_schema_property_round_trips_through_the_judge_parser() -> Result<()> {
2522        let contract =
2523            LlmTaskClassifier::load_capability_contract(&ClassifierContractConfig::default())?;
2524        let schema = contract.response_format();
2525        let reply = schema_shaped_verdict(schema)?;
2526        let judge: CapabilityJudge = StructuredJudge::new(
2527            TaskInput {
2528                recent_turn_window: None,
2529            },
2530            contract,
2531            SerdeDecoder::new(),
2532            JudgeRuntimeConfig::new(DEFAULT_JUDGE_MAX_OUTPUT_TOKENS)?,
2533        );
2534
2535        let verdict = judge.parse(&text_response(None, reply))?;
2536
2537        assert!(verdict.is_valid());
2538        assert!((0.0..=1.0).contains(&verdict.p_solve));
2539        Ok(())
2540    }
2541
2542    #[test]
2543    fn packaged_prompt_keeps_the_schema_in_the_structured_request() -> Result<()> {
2544        let contract =
2545            LlmTaskClassifier::load_capability_contract(&ClassifierContractConfig::default())?;
2546        let prompt = contract.system_prompt();
2547        let schema_name = contract
2548            .response_format()
2549            .pointer("/json_schema/name")
2550            .and_then(Value::as_str)
2551            .ok_or_else(|| LibsyError::AlgorithmError {
2552                message: "packaged response schema has no name".to_string(),
2553            })?;
2554        assert_eq!(schema_name, "CapabilityClassifierDecision");
2555        assert!(prompt.contains("SUP-1 [supported]"));
2556        assert!(prompt.contains("SUP-5 [supported]"));
2557        assert!(!prompt.contains("{{RESPONSE_SCHEMA}}"));
2558        assert!(!prompt.contains("\"type\": \"object\""));
2559        assert!(!prompt.contains("\"json_schema\""));
2560        assert!(!prompt.contains(schema_name));
2561        let rule_values = contract
2562            .response_format()
2563            .pointer("/json_schema/schema/properties/primary_rule/enum")
2564            .and_then(Value::as_array)
2565            .ok_or_else(|| LibsyError::AlgorithmError {
2566                message: "rendered response schema has no primary rule enum".to_string(),
2567            })?;
2568        assert!(
2569            rule_values
2570                .iter()
2571                .any(|value| value.as_str() == Some("SUP-1"))
2572        );
2573        assert!(
2574            rule_values
2575                .iter()
2576                .any(|value| value.as_str() == Some("none"))
2577        );
2578        Ok(())
2579    }
2580}