Skip to main content

switchyard_libsy/algorithms/
llm_class.rs

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