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