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::{BTreeMap, BTreeSet, HashSet};
7use std::sync::Arc;
8
9use async_trait::async_trait;
10use serde::{Deserialize, Deserializer};
11use serde_json::Value;
12use switchyard_protocol::{ContentBlock, Message, ModelId, Role};
13
14use super::escalation;
15use super::fall_through::{DefaultTarget, FallThrough};
16use super::util::DEFAULT_JUDGE_MAX_OUTPUT_TOKENS;
17use super::util::affinity::{AffinityRouter, ClassifyTrigger};
18use super::util::classifier_contract::{
19    ClassifierContract, ClassifierContractConfig, ClassifierResponseFormat,
20};
21use super::util::escalation::EscalationJudgeConfig;
22use super::util::llm_judge::{
23    ClassifierInput, JsonSchemaDecoder, JudgeClassifier, JudgePolicy, JudgeRuntimeConfig,
24    SerdeDecoder, StructuredJudge,
25};
26use super::util::target_selector::TargetSelectorPolicy;
27use crate::core::algorithm::{self, Algorithm, Driver};
28use crate::core::classifier::{Classification, Classifier, Score};
29use crate::core::state::State;
30use crate::{LibsyError, Result};
31use switchyard_protocol::{Request, Response};
32
33const PROMPT_TEMPLATE: &str = include_str!("../prompts/capability-classifier/prompt.md");
34const SCHEMA_TEMPLATE: &str = include_str!("../prompts/capability-classifier/schema.json");
35/// Telemetry label for this algorithm's spans, metrics, and logs.
36const ALGORITHM_NAME: &str = "llm_task_classifier";
37
38/// Restates the task after the conversation the judge is asked to route.
39///
40/// The rubric leads the request, which works while the payload is a single short
41/// message. Once `recent_turn_window` includes real conversation, those instructions
42/// sit far from the generation point and the judge sometimes *answers* the
43/// conversation instead of classifying it — the reply then fails to parse, the
44/// verdict is unavailable, and the turn silently falls back. Repeating the task at
45/// the end is what reliably prevents that.
46const TRAILING_ROUTING_INSTRUCTION: &str =
47    "Route the conversation above. Output ONLY the routing JSON object, nothing else.";
48
49#[derive(Deserialize)]
50#[serde(deny_unknown_fields)]
51struct TaskClassifierVerdict {
52    crux: String,
53    primary_rule: String,
54    capability_boundary: String,
55    p_solve: f64,
56}
57
58impl TaskClassifierVerdict {
59    /// Rejects malformed or internally inconsistent verdicts before policy evaluation.
60    fn is_valid(&self) -> bool {
61        (0.0..=1.0).contains(&self.p_solve)
62            && !self.crux.trim().is_empty()
63            && matches!(
64                (
65                    self.primary_rule.as_str(),
66                    self.capability_boundary.as_str()
67                ),
68                ("SUP-1" | "SUP-2" | "SUP-3" | "SUP-4" | "SUP-5", "supported")
69                    | ("UNC-1" | "UNC-2", "uncertain")
70                    | ("LIM-1" | "LIM-2", "unsupported")
71                    | ("none", "unmatched")
72            )
73    }
74
75    /// Returns the number of threshold steps assigned to this capability boundary.
76    fn boundary_steps(&self) -> Option<u8> {
77        match self.capability_boundary.as_str() {
78            "supported" => Some(0),
79            "uncertain" | "unmatched" => Some(1),
80            "unsupported" => Some(2),
81            _ => None,
82        }
83    }
84}
85
86/// Keeps the opening task and the last `recent_turn_window` turns after it. A
87/// window of `0` keeps the task alone.
88///
89/// Inbound decoders normalize client system and developer content into
90/// `LlmRequest::instructions`, so it never reaches this list.
91///
92/// Selects by reference and clones only what survives — a coding-agent
93/// conversation carries every tool result, so cloning it whole to keep a window
94/// would copy the transcript on each judged turn.
95fn trim_messages(messages: &[Message], recent_turn_window: usize) -> Vec<Message> {
96    let is_instruction = |message: &Message| matches!(message.role, Role::System | Role::Developer);
97    let mut kept: Vec<&Message> = messages.iter().filter(|m| is_instruction(m)).collect();
98    let Some(task) = messages.iter().position(|m| m.role == Role::User) else {
99        return kept.into_iter().cloned().collect();
100    };
101    kept.push(&messages[task]);
102
103    let tail: Vec<&Message> = messages[task + 1..]
104        .iter()
105        .filter(|m| !is_instruction(m))
106        .collect();
107    kept.extend(&tail[window_start(&tail, recent_turn_window)..]);
108    kept.into_iter().cloned().collect()
109}
110
111/// The first index of the trailing window.
112///
113/// Counting messages alone can start the window between an assistant tool call and the
114/// result answering it, leaving the judge a result whose call id was never introduced. The
115/// start therefore moves back to the nearest one that keeps every tool pair whole.
116///
117/// One newest-to-oldest pass carries the ids still waiting for a call. Direction is what
118/// makes it correct: ids repeat across a conversation, and in this order a call is only
119/// ever seen after the results it could answer, so a later call — already passed — clears
120/// nothing. A result whose call sits before the opening task, which trimming never reaches,
121/// keeps the set non-empty to the end and falls back to the counted start, so an unpairable
122/// result costs one pass and cannot widen the window to the whole conversation.
123fn window_start(tail: &[&Message], recent_turn_window: usize) -> usize {
124    let counted = tail.len().saturating_sub(recent_turn_window);
125    // An empty window holds no result to pair, and the loop below never visits its start.
126    if counted == tail.len() {
127        return counted;
128    }
129    let mut unpaired: HashSet<&str> = HashSet::new();
130    for (start, message) in tail.iter().enumerate().rev() {
131        // Blocks reverse too, so a call answers a result only when it precedes it inside
132        // one message as well as across messages.
133        for block in message.content.iter().rev() {
134            match block {
135                ContentBlock::ToolResult(result) => {
136                    unpaired.insert(result.tool_call_id.as_str());
137                }
138                ContentBlock::ToolCall(call) => {
139                    unpaired.remove(call.id.as_str());
140                }
141                _ => {}
142            }
143        }
144        if start <= counted && unpaired.is_empty() {
145            return start;
146        }
147    }
148    counted
149}
150
151/// Keeps the opening task and the latest user follow-up when they differ.
152fn task_messages(messages: &[Message]) -> Vec<Message> {
153    let mut user_messages = messages.iter().filter(|message| message.role == Role::User);
154    let Some(opening_task) = user_messages.next() else {
155        return Vec::new();
156    };
157    match user_messages.next_back() {
158        Some(latest_follow_up) => vec![opening_task.clone(), latest_follow_up.clone()],
159        None => vec![opening_task.clone()],
160    }
161}
162
163/// Selects the task messages shown to capability and custom-schema classifiers.
164struct TaskInput {
165    recent_turn_window: Option<usize>,
166}
167
168impl ClassifierInput for TaskInput {
169    fn build_messages(&self, _state: &State, request: &Request) -> Vec<Message> {
170        // The default preserves the whole-task anchor and latest user update. A
171        // configured window widens that to the surrounding conversation.
172        let mut messages = match self.recent_turn_window {
173            Some(window) => trim_messages(&request.llm_request.messages, window),
174            None => task_messages(&request.llm_request.messages),
175        };
176        // Only the windowed path carries assistant turns and tool traffic for the judge
177        // to be distracted by. The default path is user task messages only — the anchor
178        // and the latest follow-up — so there is nothing there to outrank.
179        if self.recent_turn_window.is_some() {
180            messages.push(Message::text(
181                Role::User,
182                TRAILING_ROUTING_INSTRUCTION.to_string(),
183            ));
184        }
185        messages
186    }
187}
188
189type CapabilityJudge = StructuredJudge<TaskInput, SerdeDecoder<TaskClassifierVerdict>>;
190
191struct TaskClassifierPolicy {
192    efficient_target: ModelId,
193    capable_target: ModelId,
194    base_threshold: f64,
195    threshold_step: f64,
196}
197
198impl TaskClassifierPolicy {
199    fn new(
200        efficient_target: impl Into<ModelId>,
201        capable_target: impl Into<ModelId>,
202        config: &TaskClassifierConfig,
203    ) -> Self {
204        Self {
205            efficient_target: efficient_target.into(),
206            capable_target: capable_target.into(),
207            base_threshold: config.base_threshold,
208            threshold_step: config.threshold_step,
209        }
210    }
211
212    /// Returns the required solve probability for one validated verdict.
213    fn threshold(&self, verdict: &TaskClassifierVerdict) -> Option<f64> {
214        Some(self.base_threshold + f64::from(verdict.boundary_steps()?) * self.threshold_step)
215    }
216}
217
218impl JudgePolicy for TaskClassifierPolicy {
219    type Verdict = TaskClassifierVerdict;
220
221    fn to_classification(&self, verdict: Option<&Self::Verdict>) -> Classification {
222        // Judge output is untrusted. An absent, invalid, or inconsistent verdict is
223        // ambiguous so the surrounding router applies its configured fallback.
224        let Some(verdict) = verdict.filter(|verdict| verdict.is_valid()) else {
225            return Classification::Ambiguous(vec![]);
226        };
227        // A usable verdict below the capability threshold is still a decision: the judge
228        // does not trust the efficient tier with this task.
229        let Some(threshold) = self.threshold(verdict) else {
230            return Classification::Ambiguous(vec![]);
231        };
232        let target = if verdict.p_solve >= threshold
233            || (threshold - verdict.p_solve).abs() <= f64::EPSILON
234        {
235            &self.efficient_target
236        } else {
237            &self.capable_target
238        };
239        Classification::Scores(vec![Score {
240            target: target.clone(),
241            confidence: 1.0,
242        }])
243    }
244}
245
246#[derive(Clone, Debug)]
247/// Settings that control capability classifier prompting and routing.
248pub struct TaskClassifierConfig {
249    /// Lowest solve probability that routes a supported task to the efficient target.
250    pub base_threshold: f64,
251    /// Amount added per capability-boundary step.
252    ///
253    /// Supported verdicts use `base_threshold`, uncertain and unmatched verdicts use one
254    /// step, and unsupported verdicts use two steps.
255    pub threshold_step: f64,
256    /// How often the classifier re-decides this session's target.
257    pub classify_trigger: ClassifyTrigger,
258    /// Uses the first user message as the SessionKey for sticky routing when session metadata is unavailable.
259    pub message_hash_fallback: bool,
260    /// Trailing conversation turns the judge sees on top of the client
261    /// instructions and the opening task.
262    ///
263    /// `None` (the default) judges the opening task and latest user follow-up.
264    /// `Some(n)` widens that to the client instructions, the opening task, and
265    /// the last `n` turns after it.
266    pub recent_turn_window: Option<usize>,
267    /// Prompt and verdict contract settings for the classifier judge.
268    pub contract: ClassifierContractConfig,
269    /// Maximum completion tokens available to the classifier verdict.
270    pub max_output_tokens: u64,
271}
272
273/// Flat serialized shape that maps prompt settings into the runtime contract.
274#[derive(Deserialize)]
275#[serde(deny_unknown_fields)]
276struct TaskClassifierConfigWire {
277    base_threshold: f64,
278    #[serde(default)]
279    threshold_step: f64,
280    #[serde(default)]
281    classify_trigger: ClassifyTrigger,
282    #[serde(default)]
283    message_hash_fallback: bool,
284    #[serde(default)]
285    recent_turn_window: Option<usize>,
286    #[serde(default)]
287    prompt: Option<String>,
288    #[serde(default)]
289    response_format_type: ClassifierResponseFormat,
290    #[serde(default = "default_judge_max_output_tokens")]
291    max_output_tokens: u64,
292}
293
294impl<'de> Deserialize<'de> for TaskClassifierConfig {
295    fn deserialize<D>(deserializer: D) -> std::result::Result<Self, D::Error>
296    where
297        D: Deserializer<'de>,
298    {
299        let wire = TaskClassifierConfigWire::deserialize(deserializer)?;
300        let mut contract = ClassifierContractConfig::default();
301        if let Some(prompt) = wire.prompt {
302            contract = contract.with_prompt(prompt);
303        }
304        contract = contract.with_response_format_type(wire.response_format_type);
305        Ok(Self {
306            base_threshold: wire.base_threshold,
307            threshold_step: wire.threshold_step,
308            classify_trigger: wire.classify_trigger,
309            message_hash_fallback: wire.message_hash_fallback,
310            recent_turn_window: wire.recent_turn_window,
311            contract,
312            max_output_tokens: wire.max_output_tokens,
313        })
314    }
315}
316
317const fn default_judge_max_output_tokens() -> u64 {
318    DEFAULT_JUDGE_MAX_OUTPUT_TOKENS
319}
320
321impl Default for TaskClassifierConfig {
322    fn default() -> Self {
323        Self {
324            base_threshold: 0.0,
325            threshold_step: 0.0,
326            classify_trigger: ClassifyTrigger::default(),
327            message_hash_fallback: false,
328            recent_turn_window: None,
329            contract: ClassifierContractConfig::default(),
330            max_output_tokens: DEFAULT_JUDGE_MAX_OUTPUT_TOKENS,
331        }
332    }
333}
334
335impl TaskClassifierConfig {
336    /// Validates routing thresholds before the classifier is constructed.
337    fn validate(&self) -> Result<()> {
338        if !(0.0..=1.0).contains(&self.base_threshold) {
339            return Err(LibsyError::AlgorithmError {
340                message: format!(
341                    "base_threshold must be between 0 and 1, got {}",
342                    self.base_threshold
343                ),
344            });
345        }
346        if !self.threshold_step.is_finite() || self.threshold_step < 0.0 {
347            return Err(LibsyError::AlgorithmError {
348                message: format!(
349                    "threshold_step must be finite and greater than or equal to 0, got {}",
350                    self.threshold_step
351                ),
352            });
353        }
354        let unsupported_threshold = self.base_threshold + 2.0 * self.threshold_step;
355        if unsupported_threshold > 1.0 && unsupported_threshold - 1.0 > f64::EPSILON {
356            return Err(LibsyError::AlgorithmError {
357                message: format!(
358                    "base_threshold + 2 * threshold_step must be at most 1, got {unsupported_threshold}"
359                ),
360            });
361        }
362        if self.max_output_tokens == 0 {
363            return Err(LibsyError::AlgorithmError {
364                message: "max_output_tokens must be at least 1".to_string(),
365            });
366        }
367        if self.message_hash_fallback && self.classify_trigger != ClassifyTrigger::NewSession {
368            return Err(LibsyError::AlgorithmError {
369                message: "message_hash_fallback requires classify_trigger = new_session"
370                    .to_string(),
371            });
372        }
373        Ok(())
374    }
375}
376
377/// Policy that maps a custom classifier verdict to a routing target.
378#[derive(Clone, Debug)]
379pub enum CustomClassifierPolicy {
380    /// Resolves a JSON Pointer and treats its string value as a configured target label.
381    TargetSelector {
382        /// JSON Pointer evaluated against each schema-validated verdict.
383        selector: String,
384    },
385}
386
387impl CustomClassifierPolicy {
388    /// Creates a policy that selects a target label through a JSON Pointer.
389    pub fn target_selector(selector: impl Into<String>) -> Self {
390        Self::TargetSelector {
391            selector: selector.into(),
392        }
393    }
394}
395
396/// Settings for a classifier whose JSON Schema and target-selection policy are user supplied.
397#[derive(Clone, Debug)]
398pub struct CustomClassifierConfig {
399    /// System prompt sent to the classifier judge.
400    pub prompt: String,
401    /// Inner JSON Schema placed inside the provider's structured-output wrapper.
402    pub response_schema: Value,
403    /// Deterministic policy applied after the verdict passes schema validation.
404    pub policy: CustomClassifierPolicy,
405    /// How often the classifier re-decides this session's target.
406    pub classify_trigger: ClassifyTrigger,
407    /// Uses the first user message when session metadata is unavailable.
408    pub message_hash_fallback: bool,
409    /// Trailing conversation turns shown to the classifier judge.
410    pub recent_turn_window: Option<usize>,
411    /// Maximum completion tokens available to the classifier verdict.
412    pub max_output_tokens: u64,
413}
414
415impl CustomClassifierConfig {
416    /// Creates a custom-schema classifier contract with conservative runtime defaults.
417    pub fn new(
418        prompt: impl Into<String>,
419        response_schema: Value,
420        policy: CustomClassifierPolicy,
421    ) -> Self {
422        Self {
423            prompt: prompt.into(),
424            response_schema,
425            policy,
426            classify_trigger: ClassifyTrigger::default(),
427            message_hash_fallback: false,
428            recent_turn_window: None,
429            max_output_tokens: DEFAULT_JUDGE_MAX_OUTPUT_TOKENS,
430        }
431    }
432
433    fn validate(&self) -> Result<()> {
434        if self.max_output_tokens == 0 {
435            return Err(LibsyError::AlgorithmError {
436                message: "max_output_tokens must be at least 1".to_string(),
437            });
438        }
439        if self.message_hash_fallback && self.classify_trigger != ClassifyTrigger::NewSession {
440            return Err(LibsyError::AlgorithmError {
441                message: "message_hash_fallback requires classify_trigger = new_session"
442                    .to_string(),
443            });
444        }
445        Ok(())
446    }
447}
448
449enum CustomPolicyRuntime {
450    TargetSelector(TargetSelectorPolicy),
451}
452
453impl JudgePolicy for CustomPolicyRuntime {
454    type Verdict = Value;
455
456    fn to_classification(&self, verdict: Option<&Self::Verdict>) -> Classification {
457        match self {
458            Self::TargetSelector(policy) => policy.to_classification(verdict),
459        }
460    }
461}
462
463struct TaskClassifier {
464    classifier: JudgeClassifier<CapabilityJudge, TaskClassifierPolicy>,
465    efficient_target: ModelId,
466    capable_target: ModelId,
467}
468
469/// Builds the affinity router a trigger calls for, if any.
470fn affinity_router(
471    trigger: ClassifyTrigger,
472    message_hash_fallback: bool,
473) -> Option<Arc<AffinityRouter>> {
474    let router = match trigger {
475        ClassifyTrigger::EveryRequest => return None,
476        ClassifyTrigger::NewSession => AffinityRouter::new(),
477        ClassifyTrigger::UserTurn => AffinityRouter::new().with_release_on_user_turn(),
478    };
479    let router = if message_hash_fallback {
480        router.with_message_hash_fallback()
481    } else {
482        router
483    };
484    Some(Arc::new(router))
485}
486
487/// Routes requests through a capability, escalation, or custom classifier mode.
488pub struct LlmTaskClassifier {
489    route: FallThrough<State>,
490    /// Classifier used when this router is embedded in another cascade.
491    inner: Arc<dyn Classifier<State>>,
492}
493
494struct ClassifierRouteConfig {
495    default_target: ModelId,
496    classify_trigger: ClassifyTrigger,
497    message_hash_fallback: bool,
498}
499
500/// Complete construction settings for one LLM classifier mode.
501#[derive(Clone)]
502#[non_exhaustive]
503pub enum LlmClassifierConfig {
504    /// Routes between efficient and capable targets from a task-level verdict.
505    Capability {
506        /// Target that produces classifier verdicts.
507        judge_target: ModelId,
508        /// Target used when the efficient tier can handle the task.
509        efficient_target: ModelId,
510        /// Target used when the task needs the capable tier.
511        capable_target: ModelId,
512        /// Capability classifier settings.
513        config: TaskClassifierConfig,
514    },
515    /// Judges efficient responses and escalates after a confirmed streak.
516    Escalation {
517        /// Target that produces escalation verdicts.
518        judge_target: ModelId,
519        /// Target called before each escalation decision.
520        efficient_target: ModelId,
521        /// Target used after escalation is confirmed.
522        capable_target: ModelId,
523        /// Prompt and verdict contract settings for the escalation judge.
524        contract: ClassifierContractConfig,
525        /// Escalation policy settings.
526        config: EscalationJudgeConfig,
527        /// Maximum completion tokens available to the escalation verdict.
528        max_output_tokens: u64,
529    },
530    /// Routes among named targets using a user-supplied schema and policy.
531    Custom {
532        /// Target that produces classifier verdicts.
533        judge_target: ModelId,
534        /// User-facing labels paired with their resolved routing targets.
535        targets: Vec<(String, ModelId)>,
536        /// Label selected when the judge does not produce a usable verdict.
537        default_target: String,
538        /// Custom classifier settings.
539        config: CustomClassifierConfig,
540    },
541}
542
543impl LlmTaskClassifier {
544    /// Builds the classifier mode described by `config`.
545    ///
546    /// # Errors
547    ///
548    /// Returns an error when the selected mode's targets, contract, policy, or runtime
549    /// settings are invalid.
550    pub fn new(config: LlmClassifierConfig) -> Result<Self> {
551        match config {
552            LlmClassifierConfig::Capability {
553                judge_target,
554                efficient_target,
555                capable_target,
556                config,
557            } => Self::build_capability(judge_target, efficient_target, capable_target, config),
558            LlmClassifierConfig::Escalation {
559                judge_target,
560                efficient_target,
561                capable_target,
562                contract,
563                config,
564                max_output_tokens,
565            } => Self::build_escalation(
566                judge_target,
567                efficient_target,
568                capable_target,
569                contract,
570                config,
571                max_output_tokens,
572            ),
573            LlmClassifierConfig::Custom {
574                judge_target,
575                targets,
576                default_target,
577                config,
578            } => Self::build_custom(judge_target, targets, default_target, config),
579        }
580    }
581
582    fn build_capability(
583        judge_target: ModelId,
584        efficient_target: ModelId,
585        capable_target: ModelId,
586        config: TaskClassifierConfig,
587    ) -> Result<Self> {
588        config.validate()?;
589        let contract = Self::load_capability_contract(&config.contract)?;
590        let targets = vec![efficient_target.clone(), capable_target.clone()];
591        let classify_trigger = config.classify_trigger;
592        let message_hash_fallback = config.message_hash_fallback;
593        let classifier = Arc::new(TaskClassifier {
594            classifier: JudgeClassifier::new(
595                StructuredJudge::new(
596                    TaskInput {
597                        recent_turn_window: config.recent_turn_window,
598                    },
599                    contract,
600                    SerdeDecoder::new(),
601                    JudgeRuntimeConfig::new(config.max_output_tokens)?,
602                ),
603                judge_target.clone(),
604                TaskClassifierPolicy::new(
605                    efficient_target.clone(),
606                    capable_target.clone(),
607                    &config,
608                ),
609            ),
610            efficient_target: efficient_target.clone(),
611            capable_target: capable_target.clone(),
612        });
613        let inner: Arc<dyn Classifier<State>> = classifier.clone();
614        Self::from_classifier(
615            targets,
616            inner,
617            ClassifierRouteConfig {
618                default_target: classifier.capable_target.clone(),
619                classify_trigger,
620                message_hash_fallback,
621            },
622        )
623    }
624
625    fn build_custom(
626        judge_target: ModelId,
627        targets: Vec<(String, ModelId)>,
628        default_target: String,
629        config: CustomClassifierConfig,
630    ) -> Result<Self> {
631        config.validate()?;
632        if targets.len() < 2 {
633            return Err(LibsyError::AlgorithmError {
634                message: "custom classifier requires at least two targets".to_string(),
635            });
636        }
637
638        let mut labels = BTreeSet::new();
639        let mut resolved_names = BTreeSet::new();
640        let mut target_map = BTreeMap::new();
641        let mut resolved_targets = Vec::with_capacity(targets.len());
642        for (label, target) in targets {
643            if label.trim().is_empty() || label.trim() != label {
644                return Err(LibsyError::AlgorithmError {
645                    message: "custom classifier target labels must be non-empty and have no surrounding whitespace"
646                        .to_string(),
647                });
648            }
649            if !labels.insert(label.clone()) {
650                return Err(LibsyError::AlgorithmError {
651                    message: format!("custom classifier target label {label:?} is duplicated"),
652                });
653            }
654            if !resolved_names.insert(target.clone()) {
655                return Err(LibsyError::AlgorithmError {
656                    message: format!("custom classifier resolved target {target:?} is duplicated"),
657                });
658            }
659            target_map.insert(label, target.clone());
660            resolved_targets.push(target);
661        }
662        let default_name =
663            target_map
664                .get(&default_target)
665                .cloned()
666                .ok_or_else(|| LibsyError::AlgorithmError {
667                    message: format!(
668                        "default_target {default_target:?} must be one of the configured targets"
669                    ),
670                })?;
671
672        let CustomClassifierConfig {
673            prompt,
674            response_schema,
675            policy,
676            classify_trigger,
677            message_hash_fallback,
678            recent_turn_window,
679            max_output_tokens,
680        } = config;
681        let contract = ClassifierContract::from_inner_schema(&prompt, response_schema)?;
682        let policy = match policy {
683            CustomClassifierPolicy::TargetSelector { selector } => {
684                CustomPolicyRuntime::TargetSelector(TargetSelectorPolicy::new(
685                    selector, target_map,
686                )?)
687            }
688        };
689        let classifier: Arc<dyn Classifier<State>> = Arc::new(JudgeClassifier::new(
690            StructuredJudge::new(
691                TaskInput { recent_turn_window },
692                contract,
693                JsonSchemaDecoder::new(),
694                JudgeRuntimeConfig::new(max_output_tokens)?,
695            ),
696            judge_target,
697            policy,
698        ));
699
700        Self::from_classifier(
701            resolved_targets,
702            classifier,
703            ClassifierRouteConfig {
704                default_target: default_name,
705                classify_trigger,
706                message_hash_fallback,
707            },
708        )
709    }
710
711    fn build_escalation(
712        judge_target: ModelId,
713        efficient_target: ModelId,
714        capable_target: ModelId,
715        contract_config: ClassifierContractConfig,
716        config: EscalationJudgeConfig,
717        max_output_tokens: u64,
718    ) -> Result<Self> {
719        let inner = escalation::build_classifier(
720            judge_target,
721            &efficient_target,
722            &capable_target,
723            contract_config,
724            config,
725            max_output_tokens,
726        )?;
727        let targets = vec![capable_target, efficient_target];
728        Ok(Self {
729            route: FallThrough::<State>::new_with_state(targets)
730                .with_name(ALGORITHM_NAME)
731                .with_classifier(Arc::clone(&inner)),
732            inner,
733        })
734    }
735
736    /// Loads the packaged capability-classifier contract.
737    fn load_capability_contract(config: &ClassifierContractConfig) -> Result<ClassifierContract> {
738        ClassifierContract::from_config(config, PROMPT_TEMPLATE, SCHEMA_TEMPLATE)
739    }
740
741    /// Keeps affinity and fallback ordering identical across judge-backed modes.
742    fn from_classifier(
743        targets: Vec<ModelId>,
744        inner: Arc<dyn Classifier<State>>,
745        config: ClassifierRouteConfig,
746    ) -> Result<Self> {
747        algorithm::ensure_model_is_target(&targets, &config.default_target)?;
748        if config.message_hash_fallback && config.classify_trigger != ClassifyTrigger::NewSession {
749            return Err(LibsyError::AlgorithmError {
750                message: "message_hash_fallback requires classify_trigger = new_session"
751                    .to_string(),
752            });
753        }
754        // Affinity comes first so a retained assignment short-circuits the judge call.
755        let mut route = FallThrough::<State>::new_with_state(targets).with_name(ALGORITHM_NAME);
756        if let Some(affinity) =
757            affinity_router(config.classify_trigger, config.message_hash_fallback).as_ref()
758        {
759            // Both roles must share one `Arc` so the classifier reads what the processor wrote.
760            route = route
761                .with_processor(affinity.clone())
762                .with_classifier(affinity.clone());
763        }
764        let fallback = DefaultTarget::new(config.default_target);
765        Ok(Self {
766            route: route
767                .with_classifier(inner.clone())
768                .with_classifier(Arc::new(fallback)),
769            inner,
770        })
771    }
772}
773
774#[async_trait]
775impl Classifier<State> for TaskClassifier {
776    fn routing_tier(&self, selected_model_id: &ModelId) -> Option<&'static str> {
777        if self.efficient_target == self.capable_target {
778            None
779        } else if *selected_model_id == self.efficient_target {
780            Some("weak")
781        } else if *selected_model_id == self.capable_target {
782            Some("strong")
783        } else {
784            None
785        }
786    }
787
788    async fn score(
789        &self,
790        state: &mut State,
791        request: &mut Request,
792        driver: Option<&Driver>,
793    ) -> Result<(Classification, Option<Response>)> {
794        self.classifier.score(state, request, driver).await
795    }
796}
797
798#[async_trait]
799impl Classifier<State> for LlmTaskClassifier {
800    fn routing_tier(&self, selected_model_id: &ModelId) -> Option<&'static str> {
801        self.inner.routing_tier(selected_model_id)
802    }
803
804    async fn score(
805        &self,
806        state: &mut State,
807        request: &mut Request,
808        driver: Option<&Driver>,
809    ) -> Result<(Classification, Option<Response>)> {
810        self.inner.score(state, request, driver).await
811    }
812}
813
814#[async_trait]
815impl Algorithm for LlmTaskClassifier {
816    fn name(&self) -> &str {
817        "llm_task_classifier"
818    }
819
820    async fn route(
821        self: Arc<Self>,
822        driver: Driver,
823        request: Request,
824    ) -> Result<crate::RoutingOutcome> {
825        self.route.execute(driver, request).await
826    }
827}
828
829#[cfg(test)]
830mod tests {
831    use std::sync::Arc;
832
833    use parking_lot::Mutex;
834    use serde_json::Value;
835
836    use super::*;
837    use switchyard_protocol::{
838        ContentBlock, InstructionBlock, LlmClientError, LlmRequest, Metadata, ToolCall, ToolResult,
839        completion_text, text_request, text_response,
840    };
841
842    use crate::algorithms::util::llm_judge::Judge;
843    use crate::core::testing::{Serve, test_drive};
844    use switchyard_protocol::{LlmResponse, Response};
845
846    const TEST_THRESHOLD: f64 = 0.5;
847
848    fn test_config(base_threshold: f64) -> TaskClassifierConfig {
849        TaskClassifierConfig {
850            base_threshold,
851            ..TaskClassifierConfig::default()
852        }
853    }
854
855    fn policy() -> TaskClassifierPolicy {
856        TaskClassifierPolicy::new("efficient", "capable", &test_config(TEST_THRESHOLD))
857    }
858
859    fn verdict(
860        p_solve: f64,
861        capability_boundary: &str,
862        primary_rule: &str,
863    ) -> TaskClassifierVerdict {
864        TaskClassifierVerdict {
865            crux: "test crux".to_string(),
866            primary_rule: primary_rule.to_string(),
867            capability_boundary: capability_boundary.to_string(),
868            p_solve,
869        }
870    }
871
872    fn selected(
873        policy: &TaskClassifierPolicy,
874        verdict: Option<&TaskClassifierVerdict>,
875    ) -> Result<ModelId> {
876        policy
877            .to_classification(verdict)
878            .argmax(false)?
879            .map(|score| score.target)
880            .ok_or_else(|| LibsyError::AlgorithmError {
881                message: "policy abstained".to_string(),
882            })
883    }
884
885    /// Records what each target received; answers the judge with a supported verdict and
886    /// every other target with a plain completion.
887    #[derive(Default)]
888    struct Recorder {
889        calls: Mutex<Vec<String>>,
890        call_roles: Mutex<Vec<(String, bool)>>,
891        judge_max_output_tokens: Mutex<Vec<Option<u64>>>,
892        judge_system_prompts: Mutex<Vec<String>>,
893    }
894
895    impl Recorder {
896        fn calls(&self) -> Vec<String> {
897            self.calls.lock().clone()
898        }
899
900        fn call_roles(&self) -> Vec<(String, bool)> {
901            self.call_roles.lock().clone()
902        }
903
904        fn judge_max_output_tokens(&self) -> Vec<Option<u64>> {
905            self.judge_max_output_tokens.lock().clone()
906        }
907
908        fn judge_system_prompts(&self) -> Vec<String> {
909            self.judge_system_prompts.lock().clone()
910        }
911
912        fn serve(self: &Arc<Self>) -> impl Serve {
913            let recorder = Arc::clone(self);
914            move |model: ModelId, request: Request| {
915                let recorder = Arc::clone(&recorder);
916                async move {
917                    let model = model.to_string();
918                    recorder.calls.lock().push(model.clone());
919                    recorder
920                        .call_roles
921                        .lock()
922                        .push((model.clone(), model != "judge"));
923                    let completion = if model == "judge" {
924                        recorder
925                            .judge_max_output_tokens
926                            .lock()
927                            .push(request.llm_request.output.max_output_tokens);
928                        recorder.judge_system_prompts.lock().extend(
929                            request
930                                .llm_request
931                                .instructions
932                                .first()
933                                .and_then(|instruction| {
934                                    instruction.content.iter().find_map(|b| {
935                                        if let ContentBlock::Text { text } = b {
936                                            Some(text.clone())
937                                        } else {
938                                            None
939                                        }
940                                    })
941                                }),
942                        );
943                        r#"{"crux":"bounded task","primary_rule":"SUP-1","capability_boundary":"supported","p_solve":0.9}"#.to_string()
944                    } else {
945                        format!("answer from {model}")
946                    };
947                    Ok(Response {
948                        llm_response: LlmResponse::Agg(text_response(None, completion)),
949                        metadata: request.metadata,
950                    })
951                }
952            }
953        }
954    }
955
956    /// The judge times out; every other target answers normally.
957    fn unreachable_judge() -> impl Serve {
958        |model: ModelId, request: Request| async move {
959            let model = model.to_string();
960            if model == "judge" {
961                return Err(LlmClientError::Timeout {
962                    source: Box::new(std::io::Error::other("judge unreachable")),
963                });
964            }
965            Ok(Response {
966                llm_response: LlmResponse::Agg(text_response(None, format!("answer from {model}"))),
967                metadata: request.metadata,
968            })
969        }
970    }
971
972    fn router() -> Result<Arc<LlmTaskClassifier>> {
973        Ok(Arc::new(LlmTaskClassifier::new(
974            LlmClassifierConfig::Capability {
975                judge_target: ModelId::from("judge"),
976                efficient_target: ModelId::from("efficient"),
977                capable_target: ModelId::from("capable"),
978                config: test_config(TEST_THRESHOLD),
979            },
980        )?))
981    }
982
983    fn classify_request() -> Request {
984        Request {
985            llm_request: text_request(Some("auto".to_string()), "classify this task"),
986            raw_request: None,
987            metadata: None,
988        }
989    }
990
991    fn classify_session_request() -> Request {
992        Request {
993            metadata: Some(Metadata {
994                session_id: Some("session-1".to_string()),
995                ..Metadata::default()
996            }),
997            ..classify_request()
998        }
999    }
1000
1001    fn classify_follow_up_request() -> Request {
1002        let mut request = classify_request();
1003        request
1004            .llm_request
1005            .messages
1006            .push(Message::text(Role::Assistant, "I will add the test."));
1007        request.llm_request.messages.push(Message::text(
1008            Role::User,
1009            "Now run the test suite and report the result.",
1010        ));
1011        request
1012    }
1013
1014    #[tokio::test]
1015    async fn an_unreachable_judge_routes_capable_instead_of_failing_the_request() -> Result<()> {
1016        let router = router()?;
1017
1018        let (selected_model, response) =
1019            test_drive(router, classify_request(), unreachable_judge()).await?;
1020
1021        assert_eq!(selected_model, "capable");
1022        assert_eq!(
1023            response.llm_response.as_agg().map(completion_text),
1024            Some("answer from capable".to_string())
1025        );
1026        Ok(())
1027    }
1028
1029    #[tokio::test]
1030    async fn classifier_judges_each_request_without_affinity() -> Result<()> {
1031        let recorder = Arc::new(Recorder::default());
1032        let router = router()?;
1033        let request = classify_request;
1034
1035        test_drive(router.clone(), request(), recorder.serve()).await?;
1036        test_drive(router.clone(), request(), recorder.serve()).await?;
1037
1038        assert_eq!(
1039            recorder.calls(),
1040            vec!["judge", "efficient", "judge", "efficient"]
1041        );
1042        assert_eq!(
1043            recorder.call_roles(),
1044            vec![
1045                ("judge".to_string(), false),
1046                ("efficient".to_string(), true),
1047                ("judge".to_string(), false),
1048                ("efficient".to_string(), true),
1049            ]
1050        );
1051        Ok(())
1052    }
1053
1054    #[tokio::test]
1055    async fn classifier_config_sets_the_judge_completion_cap() -> Result<()> {
1056        let recorder = Arc::new(Recorder::default());
1057        let router = Arc::new(LlmTaskClassifier::new(LlmClassifierConfig::Capability {
1058            judge_target: ModelId::from("judge"),
1059            efficient_target: ModelId::from("efficient"),
1060            capable_target: ModelId::from("capable"),
1061            config: TaskClassifierConfig {
1062                max_output_tokens: 512,
1063                ..test_config(TEST_THRESHOLD)
1064            },
1065        })?);
1066
1067        test_drive(router, classify_request(), recorder.serve()).await?;
1068
1069        assert_eq!(recorder.judge_max_output_tokens(), vec![Some(512)]);
1070        Ok(())
1071    }
1072
1073    #[tokio::test]
1074    async fn classifier_config_overrides_the_packaged_prompt() -> Result<()> {
1075        let recorder = Arc::new(Recorder::default());
1076        let router = Arc::new(LlmTaskClassifier::new(LlmClassifierConfig::Capability {
1077            judge_target: ModelId::from("judge"),
1078            efficient_target: ModelId::from("efficient"),
1079            capable_target: ModelId::from("capable"),
1080            config: TaskClassifierConfig {
1081                contract: ClassifierContractConfig::default()
1082                    .with_prompt("Custom capability rubric."),
1083                ..test_config(TEST_THRESHOLD)
1084            },
1085        })?);
1086
1087        test_drive(router, classify_request(), recorder.serve()).await?;
1088
1089        let prompts = recorder.judge_system_prompts();
1090        assert_eq!(prompts.len(), 1);
1091        assert_eq!(prompts[0], "Custom capability rubric.");
1092        Ok(())
1093    }
1094
1095    #[tokio::test]
1096    async fn classifier_config_enables_new_session_trigger() -> Result<()> {
1097        let recorder = Arc::new(Recorder::default());
1098        let router = Arc::new(LlmTaskClassifier::new(LlmClassifierConfig::Capability {
1099            judge_target: ModelId::from("judge"),
1100            efficient_target: ModelId::from("efficient"),
1101            capable_target: ModelId::from("capable"),
1102            config: TaskClassifierConfig {
1103                classify_trigger: ClassifyTrigger::NewSession,
1104                ..test_config(TEST_THRESHOLD)
1105            },
1106        })?);
1107
1108        let session_request = classify_session_request;
1109        test_drive(router.clone(), session_request(), recorder.serve()).await?;
1110        test_drive(router.clone(), session_request(), recorder.serve()).await?;
1111
1112        assert_eq!(recorder.calls(), vec!["judge", "efficient", "efficient"]);
1113        Ok(())
1114    }
1115
1116    #[tokio::test]
1117    async fn classifier_config_reuses_message_hash_affinity_for_a_follow_up() -> Result<()> {
1118        let recorder = Arc::new(Recorder::default());
1119        let router = Arc::new(LlmTaskClassifier::new(LlmClassifierConfig::Capability {
1120            judge_target: ModelId::from("judge"),
1121            efficient_target: ModelId::from("efficient"),
1122            capable_target: ModelId::from("capable"),
1123            config: TaskClassifierConfig {
1124                classify_trigger: ClassifyTrigger::NewSession,
1125                message_hash_fallback: true,
1126                recent_turn_window: None,
1127                ..test_config(TEST_THRESHOLD)
1128            },
1129        })?);
1130
1131        test_drive(router.clone(), classify_request(), recorder.serve()).await?;
1132        test_drive(
1133            router.clone(),
1134            classify_follow_up_request(),
1135            recorder.serve(),
1136        )
1137        .await?;
1138
1139        assert_eq!(recorder.calls(), vec!["judge", "efficient", "efficient"]);
1140        Ok(())
1141    }
1142
1143    #[test]
1144    fn the_threshold_boundary_is_inclusive() -> Result<()> {
1145        let policy = policy();
1146        let at_threshold = verdict(0.5, "supported", "SUP-1");
1147        let below_threshold = verdict(0.49, "supported", "SUP-1");
1148        assert_eq!(selected(&policy, Some(&at_threshold))?, "efficient");
1149        assert_eq!(selected(&policy, Some(&below_threshold))?, "capable");
1150        Ok(())
1151    }
1152
1153    #[test]
1154    fn the_threshold_moves_the_routing_boundary() -> Result<()> {
1155        let borderline = verdict(0.5, "supported", "SUP-1");
1156        let strict = TaskClassifierPolicy::new("efficient", "capable", &test_config(0.9));
1157        let lenient = TaskClassifierPolicy::new("efficient", "capable", &test_config(0.1));
1158        assert_eq!(selected(&strict, Some(&borderline))?, "capable");
1159        assert_eq!(selected(&lenient, Some(&borderline))?, "efficient");
1160        Ok(())
1161    }
1162
1163    #[test]
1164    fn classifier_config_rejects_unknown_fields() {
1165        let error = serde_json::from_value::<TaskClassifierConfig>(serde_json::json!({
1166            "base_threshold": 0.5,
1167            "classifier_magic": true,
1168        }))
1169        .expect_err("unknown classifier fields must be rejected");
1170
1171        assert!(
1172            error
1173                .to_string()
1174                .contains("unknown field `classifier_magic`"),
1175            "{error}"
1176        );
1177    }
1178
1179    #[test]
1180    fn invalid_classifier_config_is_rejected() -> Result<()> {
1181        for bad in [1.5, -0.1, f64::NAN, f64::INFINITY] {
1182            assert!(
1183                LlmTaskClassifier::new(LlmClassifierConfig::Capability {
1184                    judge_target: ModelId::from("judge"),
1185                    efficient_target: ModelId::from("e"),
1186                    capable_target: ModelId::from("c"),
1187                    config: test_config(bad),
1188                })
1189                .is_err(),
1190                "base threshold {bad} should be rejected"
1191            );
1192        }
1193        for config in [
1194            TaskClassifierConfig {
1195                base_threshold: 0.5,
1196                threshold_step: -0.1,
1197                ..TaskClassifierConfig::default()
1198            },
1199            TaskClassifierConfig {
1200                base_threshold: 0.8,
1201                threshold_step: 0.11,
1202                ..TaskClassifierConfig::default()
1203            },
1204            TaskClassifierConfig {
1205                base_threshold: 0.5,
1206                message_hash_fallback: true,
1207                ..TaskClassifierConfig::default()
1208            },
1209            TaskClassifierConfig {
1210                base_threshold: 0.5,
1211                max_output_tokens: 0,
1212                ..TaskClassifierConfig::default()
1213            },
1214        ] {
1215            assert!(
1216                LlmTaskClassifier::new(LlmClassifierConfig::Capability {
1217                    judge_target: ModelId::from("judge"),
1218                    efficient_target: ModelId::from("e"),
1219                    capable_target: ModelId::from("c"),
1220                    config,
1221                })
1222                .is_err()
1223            );
1224        }
1225        for base_threshold in [0.0, 1.0] {
1226            LlmTaskClassifier::new(LlmClassifierConfig::Capability {
1227                judge_target: ModelId::from("judge"),
1228                efficient_target: ModelId::from("e"),
1229                capable_target: ModelId::from("c"),
1230                config: test_config(base_threshold),
1231            })?;
1232        }
1233        Ok(())
1234    }
1235
1236    #[test]
1237    fn an_unusable_verdict_is_ambiguous() -> Result<()> {
1238        let policy = policy();
1239        let inconsistent_rule = TaskClassifierVerdict {
1240            capability_boundary: "uncertain".to_string(),
1241            ..verdict(1.0, "supported", "SUP-1")
1242        };
1243        let empty_crux = TaskClassifierVerdict {
1244            crux: "  ".to_string(),
1245            ..verdict(1.0, "supported", "SUP-1")
1246        };
1247        let unusable = [
1248            Some(verdict(1.1, "supported", "SUP-1")),
1249            Some(inconsistent_rule),
1250            Some(empty_crux),
1251            None,
1252        ];
1253        for verdict in unusable {
1254            let classification = policy.to_classification(verdict.as_ref());
1255            assert!(matches!(classification, Classification::Ambiguous(_)));
1256            assert!(classification.argmax(false)?.is_none());
1257            assert!(classification.argmax(true)?.is_none());
1258        }
1259        Ok(())
1260    }
1261
1262    #[test]
1263    fn capability_boundaries_apply_monotonic_threshold_steps() -> Result<()> {
1264        let policy = TaskClassifierPolicy::new(
1265            "efficient",
1266            "capable",
1267            &TaskClassifierConfig {
1268                threshold_step: 0.1,
1269                ..test_config(0.4)
1270            },
1271        );
1272
1273        assert_eq!(
1274            selected(&policy, Some(&verdict(0.4, "supported", "SUP-2")))?,
1275            "efficient"
1276        );
1277        assert_eq!(
1278            selected(&policy, Some(&verdict(0.49, "uncertain", "UNC-1")))?,
1279            "capable"
1280        );
1281        assert_eq!(
1282            selected(&policy, Some(&verdict(0.5, "uncertain", "UNC-1")))?,
1283            "efficient"
1284        );
1285        assert_eq!(
1286            selected(&policy, Some(&verdict(0.5, "unmatched", "none")))?,
1287            "efficient"
1288        );
1289        assert_eq!(
1290            selected(&policy, Some(&verdict(0.59, "unsupported", "LIM-1")))?,
1291            "capable"
1292        );
1293        assert_eq!(
1294            selected(&policy, Some(&verdict(0.6, "unsupported", "LIM-1")))?,
1295            "efficient"
1296        );
1297        Ok(())
1298    }
1299
1300    /// The text of each message a judge with `recent_turn_window` would be sent.
1301    /// The no-window case is covered by `capability_judge_builds_a_structured_request`.
1302    fn capability_judge(recent_turn_window: Option<usize>) -> Result<CapabilityJudge> {
1303        Ok(StructuredJudge::new(
1304            TaskInput { recent_turn_window },
1305            LlmTaskClassifier::load_capability_contract(&ClassifierContractConfig::default())?,
1306            SerdeDecoder::new(),
1307            JudgeRuntimeConfig::new(DEFAULT_JUDGE_MAX_OUTPUT_TOKENS)?,
1308        ))
1309    }
1310
1311    fn judged_contents(recent_turn_window: usize) -> Result<Vec<String>> {
1312        let judge = capability_judge(Some(recent_turn_window))?;
1313        let request = Request {
1314            llm_request: LlmRequest {
1315                messages: vec![
1316                    Message::text(Role::System, "client instructions"),
1317                    Message::text(Role::User, "initial task"),
1318                    Message::text(Role::Assistant, "old response"),
1319                    Message::text(Role::User, "old follow-up"),
1320                    Message::text(Role::Assistant, "recent 1"),
1321                    Message::text(Role::User, "recent 2"),
1322                ],
1323                ..LlmRequest::default()
1324            },
1325            raw_request: None,
1326            metadata: None,
1327        };
1328        Ok(judge
1329            .build_request(&State::default(), &request)
1330            .llm_request
1331            .messages
1332            .iter()
1333            .filter_map(|message| message.text_content("\n"))
1334            .collect())
1335    }
1336
1337    #[test]
1338    fn a_window_widens_the_judge_to_the_surrounding_conversation() -> Result<()> {
1339        // Client instructions and the opening task, plus the last two turns.
1340        let contents = judged_contents(2)?;
1341        assert!(contents.contains(&"client instructions".to_string()));
1342        assert!(contents.contains(&"initial task".to_string()));
1343        assert!(contents.contains(&"recent 1".to_string()));
1344        assert!(contents.contains(&"recent 2".to_string()));
1345        assert!(!contents.contains(&"old response".to_string()));
1346        Ok(())
1347    }
1348
1349    #[test]
1350    fn a_zero_window_keeps_only_the_instructions_and_the_task() -> Result<()> {
1351        let contents = judged_contents(0)?;
1352        assert!(contents.contains(&"client instructions".to_string()));
1353        assert!(contents.contains(&"initial task".to_string()));
1354        assert!(!contents.contains(&"recent 2".to_string()));
1355        Ok(())
1356    }
1357
1358    fn tool_call(id: &str) -> Message {
1359        Message {
1360            role: Role::Assistant,
1361            content: vec![ContentBlock::ToolCall(ToolCall {
1362                id: id.to_string(),
1363                name: "search".to_string(),
1364                arguments: Value::Null,
1365            })],
1366        }
1367    }
1368
1369    fn tool_result(id: &str) -> Message {
1370        Message {
1371            role: Role::Tool,
1372            content: vec![ContentBlock::ToolResult(ToolResult {
1373                tool_call_id: id.to_string(),
1374                content: vec![ContentBlock::Text {
1375                    text: "tool output".to_string(),
1376                }],
1377                is_error: None,
1378            })],
1379        }
1380    }
1381
1382    /// A count-based window can begin on a tool result, which leaves the call that
1383    /// introduced its id outside the window and the classifier history invalid.
1384    #[test]
1385    fn trimming_keeps_the_call_that_introduced_a_kept_tool_result() {
1386        let messages = vec![
1387            Message::text(Role::System, "client instructions"),
1388            Message::text(Role::User, "initial task"),
1389            Message::text(Role::Assistant, "old response"),
1390            tool_call("call-1"),
1391            tool_result("call-1"),
1392            Message::text(Role::Assistant, "recent 1"),
1393            Message::text(Role::User, "recent 2"),
1394            Message::text(Role::Assistant, "recent 3"),
1395            Message::text(Role::User, "recent 4"),
1396        ];
1397
1398        // The five-message tail begins exactly on the tool result.
1399        let kept = trim_messages(&messages, 5);
1400
1401        assert_eq!(
1402            kept,
1403            vec![
1404                Message::text(Role::System, "client instructions"),
1405                Message::text(Role::User, "initial task"),
1406                tool_call("call-1"),
1407                tool_result("call-1"),
1408                Message::text(Role::Assistant, "recent 1"),
1409                Message::text(Role::User, "recent 2"),
1410                Message::text(Role::Assistant, "recent 3"),
1411                Message::text(Role::User, "recent 4"),
1412            ]
1413        );
1414    }
1415
1416    /// Ids repeat across a conversation, so a later call must not stand in for the one that
1417    /// answers an earlier result.
1418    #[test]
1419    fn trimming_pairs_a_repeated_id_with_the_call_that_precedes_it() {
1420        let messages = vec![
1421            Message::text(Role::System, "client instructions"),
1422            Message::text(Role::User, "initial task"),
1423            tool_call("x"),
1424            tool_result("x"),
1425            Message::text(Role::Assistant, "later"),
1426            tool_call("x"),
1427            tool_result("x"),
1428        ];
1429
1430        // The four-message tail begins on the first result, whose own call sits one earlier.
1431        let kept = trim_messages(&messages, 4);
1432
1433        assert_eq!(
1434            kept,
1435            vec![
1436                Message::text(Role::System, "client instructions"),
1437                Message::text(Role::User, "initial task"),
1438                tool_call("x"),
1439                tool_result("x"),
1440                Message::text(Role::Assistant, "later"),
1441                tool_call("x"),
1442                tool_result("x"),
1443            ]
1444        );
1445    }
1446
1447    /// A result whose call precedes the opening task can never be paired, because trimming
1448    /// never reaches behind the task. The window must not widen hunting for it.
1449    #[test]
1450    fn trimming_keeps_the_counted_window_when_a_result_cannot_be_paired() {
1451        let messages = vec![
1452            Message::text(Role::System, "client instructions"),
1453            tool_call("orphan"),
1454            Message::text(Role::User, "initial task"),
1455            Message::text(Role::Assistant, "old response"),
1456            tool_result("orphan"),
1457            Message::text(Role::Assistant, "recent 1"),
1458            Message::text(Role::User, "recent 2"),
1459        ];
1460
1461        let kept = trim_messages(&messages, 3);
1462
1463        assert_eq!(
1464            kept,
1465            vec![
1466                Message::text(Role::System, "client instructions"),
1467                Message::text(Role::User, "initial task"),
1468                tool_result("orphan"),
1469                Message::text(Role::Assistant, "recent 1"),
1470                Message::text(Role::User, "recent 2"),
1471            ]
1472        );
1473    }
1474
1475    #[test]
1476    fn a_window_restates_the_routing_instruction_last() -> Result<()> {
1477        let contents = judged_contents(2)?;
1478
1479        // Last, not merely present: the instruction only works in final position,
1480        // after the conversation content it is meant to outrank.
1481        assert_eq!(
1482            contents.last().map(String::as_str),
1483            Some(TRAILING_ROUTING_INSTRUCTION)
1484        );
1485        Ok(())
1486    }
1487
1488    #[test]
1489    fn the_default_path_is_left_unchanged() -> Result<()> {
1490        // No window means no conversation to be distracted by, so the default
1491        // request shape stays exactly as it was.
1492        let judge = capability_judge(None)?;
1493        let request = Request {
1494            llm_request: LlmRequest {
1495                messages: vec![Message::text(Role::User, "the task")],
1496                ..LlmRequest::default()
1497            },
1498            raw_request: None,
1499            metadata: None,
1500        };
1501
1502        let built = judge.build_request(&State::default(), &request);
1503
1504        // The task message alone; the rubric reaches the judge as an instruction block
1505        // rather than a message, so nothing was dropped by it not being counted here.
1506        assert_eq!(built.llm_request.messages.len(), 1);
1507        assert!(!built.llm_request.instructions.is_empty());
1508        assert!(
1509            !built
1510                .llm_request
1511                .messages
1512                .iter()
1513                .filter_map(|message| message.text_content("\n"))
1514                .any(|text| text.contains(TRAILING_ROUTING_INSTRUCTION))
1515        );
1516        Ok(())
1517    }
1518
1519    #[test]
1520    fn capability_judge_builds_a_structured_request() -> Result<()> {
1521        let judge = capability_judge(None)?;
1522        let request = Request {
1523            llm_request: LlmRequest {
1524                model: Some("inbound".to_string()),
1525                messages: vec![
1526                    Message::text(Role::System, "client instructions"),
1527                    Message::text(Role::Developer, "client developer instructions"),
1528                    Message::text(Role::User, "initial task"),
1529                    Message::text(Role::Assistant, "old response"),
1530                    Message::text(Role::User, "old follow-up"),
1531                    Message::text(Role::Assistant, "recent 1"),
1532                    Message::text(Role::User, "recent 2"),
1533                    Message::text(Role::Assistant, "recent 3"),
1534                    Message::text(Role::User, "recent 4"),
1535                    Message::text(Role::Assistant, "recent 5"),
1536                ],
1537                ..LlmRequest::default()
1538            },
1539            raw_request: None,
1540            metadata: None,
1541        };
1542        let judge_request = judge.build_request(&State::default(), &request);
1543
1544        assert_eq!(judge_request.llm_request.model, request.llm_request.model);
1545        assert_eq!(judge_request.llm_request.instructions.len(), 1);
1546        assert_eq!(judge_request.llm_request.instructions[0].role, Role::System);
1547        assert_eq!(
1548            judge_request.llm_request.instructions[0].content,
1549            InstructionBlock {
1550                role: Role::System,
1551                content: Message::text(Role::System, judge.contract().system_prompt()).content,
1552            }
1553            .content,
1554        );
1555        assert_eq!(judge_request.llm_request.messages.len(), 2);
1556        let contents = judge_request
1557            .llm_request
1558            .messages
1559            .iter()
1560            .filter_map(|message| message.text_content("\n"))
1561            .collect::<Vec<_>>();
1562        assert!(contents.contains(&"recent 4".to_string()));
1563        assert!(contents.contains(&"initial task".to_string()));
1564        assert!(!contents.contains(&"recent 5".to_string()));
1565        assert!(!contents.contains(&"client instructions".to_string()));
1566        assert_eq!(
1567            judge_request.llm_request.output.response_format,
1568            Some(judge.contract().response_format().clone())
1569        );
1570        assert_eq!(
1571            judge_request.llm_request.output.max_output_tokens,
1572            Some(DEFAULT_JUDGE_MAX_OUTPUT_TOKENS)
1573        );
1574        Ok(())
1575    }
1576
1577    fn sample_value(spec: &Value) -> Value {
1578        if let Some(first) = spec
1579            .get("enum")
1580            .and_then(Value::as_array)
1581            .and_then(|values| values.first())
1582        {
1583            return first.clone();
1584        }
1585        match spec.get("type").and_then(Value::as_str) {
1586            Some("number") => serde_json::json!(0.5),
1587            Some("boolean") => serde_json::json!(false),
1588            _ => serde_json::json!("sample"),
1589        }
1590    }
1591
1592    fn schema_shaped_verdict(schema: &Value) -> Result<String> {
1593        let properties = schema
1594            .pointer("/json_schema/schema/properties")
1595            .and_then(Value::as_object)
1596            .ok_or_else(|| LibsyError::AlgorithmError {
1597                message: "packaged schema declares no properties".to_string(),
1598            })?;
1599        Ok(Value::Object(
1600            properties
1601                .iter()
1602                .map(|(name, spec)| (name.clone(), sample_value(spec)))
1603                .collect(),
1604        )
1605        .to_string())
1606    }
1607
1608    /// Built from the schema so a property added there fails here rather than silently
1609    /// rejecting every production verdict.
1610    #[test]
1611    fn every_schema_property_round_trips_through_the_judge_parser() -> Result<()> {
1612        let contract =
1613            LlmTaskClassifier::load_capability_contract(&ClassifierContractConfig::default())?;
1614        let schema = contract.response_format();
1615        let reply = schema_shaped_verdict(schema)?;
1616        let judge: CapabilityJudge = StructuredJudge::new(
1617            TaskInput {
1618                recent_turn_window: None,
1619            },
1620            contract,
1621            SerdeDecoder::new(),
1622            JudgeRuntimeConfig::new(DEFAULT_JUDGE_MAX_OUTPUT_TOKENS)?,
1623        );
1624
1625        let verdict = judge.parse(&text_response(None, reply))?;
1626
1627        assert!(verdict.is_valid());
1628        assert!((0.0..=1.0).contains(&verdict.p_solve));
1629        Ok(())
1630    }
1631
1632    #[test]
1633    fn packaged_prompt_keeps_the_schema_in_the_structured_request() -> Result<()> {
1634        let contract =
1635            LlmTaskClassifier::load_capability_contract(&ClassifierContractConfig::default())?;
1636        let prompt = contract.system_prompt();
1637        let schema_name = contract
1638            .response_format()
1639            .pointer("/json_schema/name")
1640            .and_then(Value::as_str)
1641            .ok_or_else(|| LibsyError::AlgorithmError {
1642                message: "packaged response schema has no name".to_string(),
1643            })?;
1644        assert_eq!(schema_name, "CapabilityClassifierDecision");
1645        assert!(prompt.contains("SUP-1 [supported]"));
1646        assert!(prompt.contains("SUP-5 [supported]"));
1647        assert!(!prompt.contains("{{RESPONSE_SCHEMA}}"));
1648        assert!(!prompt.contains("\"type\": \"object\""));
1649        assert!(!prompt.contains("\"json_schema\""));
1650        assert!(!prompt.contains(schema_name));
1651        let rule_values = contract
1652            .response_format()
1653            .pointer("/json_schema/schema/properties/primary_rule/enum")
1654            .and_then(Value::as_array)
1655            .ok_or_else(|| LibsyError::AlgorithmError {
1656                message: "rendered response schema has no primary rule enum".to_string(),
1657            })?;
1658        assert!(
1659            rule_values
1660                .iter()
1661                .any(|value| value.as_str() == Some("SUP-1"))
1662        );
1663        assert!(
1664            rule_values
1665                .iter()
1666                .any(|value| value.as_str() == Some("none"))
1667        );
1668        Ok(())
1669    }
1670}