Skip to main content

switchyard_libsy/algorithms/
llm_class.rs

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