Skip to main content

switchyard_libsy/algorithms/
composite.rs

1// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2// SPDX-License-Identifier: Apache-2.0
3
4//! Routing that stacks a judge over a stage router.
5//!
6//! The judge runs as a [`Processor`]: it sets configuration the stage router reads,
7//! and picks no target itself.
8
9use std::collections::HashMap;
10use std::sync::Arc;
11
12use async_trait::async_trait;
13use parking_lot::Mutex;
14
15use super::fall_through::FallThrough;
16use super::llm_class::{
17    CustomClassifierConfig, LlmClassifierConfig, LlmTaskClassifier, TaskClassifierConfig,
18};
19use super::stage::{StageRouterConfig, build_stage_route};
20use super::util::affinity::{ClassifyTrigger, evict_if_full, has_new_user_turn, retention_key};
21use super::util::stage::{PickerMode, Tier, fall_open_tier, set_fall_open};
22use crate::core::algorithm::{Algorithm, Driver, RoutingIdentity};
23use crate::core::classifier::{Classification, Classifier};
24use crate::core::processor::{Event, Processor};
25use crate::core::state::State;
26use crate::{LibsyError, Result};
27use switchyard_protocol::{Category, Request, Response};
28
29const COMPOSITE: &str = "composite";
30
31/// Sets the stage router's fall-open tier from a judge verdict.
32///
33/// Retains the tier per routing identity, so it survives requests that carry no
34/// session ID when `message_hash_fallback` is on. The retained tier is replayed
35/// into state on every request so the cascade below reads it.
36struct TierSetter {
37    judge: Arc<dyn Classifier>,
38    trigger: ClassifyTrigger,
39    message_hash_fallback: bool,
40    default_tier: Tier,
41    tiers: Mutex<HashMap<RoutingIdentity, Tier>>,
42}
43
44impl TierSetter {
45    fn tier_for(category: Option<Category>) -> Option<Tier> {
46        match category {
47            Some(Category::Capable) => Some(Tier::Capable),
48            Some(Category::Efficient) => Some(Tier::Efficient),
49            _ => None,
50        }
51    }
52
53    /// Two requests for one identity can both pass this and both judge, since a
54    /// judge call sits between here and [`retain`](Self::retain). The later wins.
55    fn is_due(&self, identity: Option<&RoutingIdentity>, state: &State, request: &Request) -> bool {
56        match self.trigger {
57            ClassifyTrigger::UserTurn => has_new_user_turn(&request.llm_request.messages),
58            // Unkeyed requests with no carried tier cannot be told apart, so every one
59            // is a new session.
60            ClassifyTrigger::NewSession => {
61                identity.is_none_or(|identity| !self.tiers.lock().contains_key(identity))
62                    && fall_open_tier(state).is_none()
63            }
64            // Rejected by the constructor, and only in the enum for the standalone route.
65            ClassifyTrigger::EveryRequest => true,
66        }
67    }
68
69    fn retain(&self, identity: RoutingIdentity, tier: Tier) {
70        let mut tiers = self.tiers.lock();
71        // A user turn re-decides, so it overwrites. A session keeps its first verdict,
72        // matching how affinity retains an assignment.
73        let writable = self.trigger == ClassifyTrigger::UserTurn || !tiers.contains_key(&identity);
74        if writable {
75            evict_if_full(&mut tiers);
76            tiers.insert(identity, tier);
77        }
78    }
79}
80
81#[async_trait]
82impl Processor for TierSetter {
83    async fn process(&self, state: &mut State, event: Event<'_>) -> Result<()> {
84        let Event::Request { request, driver } = event else {
85            return Ok(());
86        };
87        let identity = retention_key(request, self.message_hash_fallback);
88        if self.is_due(identity.as_ref(), state, request) {
89            let (classification, _) = self.judge.score(state, request, driver).await?;
90            if let Some(winner) = classification.argmax(false)?
91                && let Some(tier) = Self::tier_for(winner.category)
92            {
93                set_fall_open(state, tier);
94                if let Some(identity) = identity {
95                    self.retain(identity, tier);
96                }
97                return Ok(());
98            }
99            if self.trigger == ClassifyTrigger::NewSession
100                && let Some(identity) = identity.as_ref()
101            {
102                // A failed opening judgment is still the session's one judgment.
103                self.retain(identity.clone(), self.default_tier);
104            }
105        }
106        // Either not this request's turn or the judge had no verdict, so the last
107        // tier stands rather than dropping back to the picker default.
108        if let Some(tier) = identity
109            .and_then(|identity| self.tiers.lock().get(&identity).copied())
110            .or_else(|| fall_open_tier(state))
111        {
112            set_fall_open(state, tier);
113            driver.set_evidence_if_empty(serde_json::json!({"source": "retained"}));
114        }
115        Ok(())
116    }
117}
118
119/// Classifier settings for the composite router's retained tier.
120pub enum CompositeJudgeConfig {
121    /// Selects a tier from a capability verdict.
122    Capability(TaskClassifierConfig),
123    /// Selects a tier using a custom schema and policy, retaining the fallback on failure.
124    Custom(CustomClassifierConfig),
125}
126
127/// A judge stacked over a stage router.
128pub struct CompositeRouterConfig {
129    /// Judge settings, including how often `classify_trigger` runs it.
130    pub judge: CompositeJudgeConfig,
131    /// Serves the turns, with the tier the judge picked as its fall-open default.
132    pub stage: StageRouterConfig,
133}
134
135/// Runs a stage router with a tier the judge picks.
136pub struct CompositeRouter {
137    route: FallThrough,
138}
139
140impl CompositeRouter {
141    /// Stacks the judge over a stage router across the same tier pair.
142    ///
143    /// Errors on a configuration either algorithm rejects and on `every_request`.
144    ///
145    /// A stage router carrying its own judge is allowed, but that judge sits ahead
146    /// of the fall-open tier and so answers most of the turns this one set a tier for.
147    pub fn new(config: CompositeRouterConfig) -> Result<Self> {
148        let setter = Self::tier_setter(config.judge, config.stage.mode.default_tier())?;
149        let route = build_stage_route(config.stage, setter.trigger == ClassifyTrigger::UserTurn)?
150            .with_name(COMPOSITE)
151            .with_processor(Arc::new(setter));
152        Ok(Self { route })
153    }
154
155    /// Builds two independent stage profiles selected by the judge.
156    pub fn new_with_profiles(
157        judge: CompositeJudgeConfig,
158        mut profile_0: StageRouterConfig,
159        mut profile_1: StageRouterConfig,
160    ) -> Result<Self> {
161        for profile in [&mut profile_0, &mut profile_1] {
162            if profile.mode == PickerMode::CapableFirst {
163                profile.efficient_hold_turns.get_or_insert(-1);
164            }
165        }
166        let setter = Self::tier_setter(judge, Tier::Efficient)?;
167        let reset_on_user_turn = setter.trigger == ClassifyTrigger::UserTurn;
168        let profiles = StageProfiles([
169            build_stage_route(profile_0, reset_on_user_turn)?,
170            build_stage_route(profile_1, reset_on_user_turn)?,
171        ]);
172        let route = FallThrough::new_with_state()
173            .with_name(COMPOSITE)
174            .with_processor(Arc::new(setter))
175            .with_classifier(Arc::new(profiles));
176        Ok(Self { route })
177    }
178
179    fn tier_setter(judge: CompositeJudgeConfig, default_tier: Tier) -> Result<TierSetter> {
180        let (trigger, message_hash_fallback) = match &judge {
181            CompositeJudgeConfig::Capability(judge) => {
182                (judge.classify_trigger, judge.message_hash_fallback)
183            }
184            CompositeJudgeConfig::Custom(config) => {
185                (config.classify_trigger, config.message_hash_fallback)
186            }
187        };
188        if trigger == ClassifyTrigger::EveryRequest {
189            return Err(LibsyError::AlgorithmError {
190                message: "composite: classify_trigger must be user_turn or new_session".to_string(),
191            });
192        }
193        let judge: Arc<dyn Classifier> = match judge {
194            CompositeJudgeConfig::Capability(judge) => {
195                // Composite owns the trigger and retained verdict.
196                Arc::new(LlmTaskClassifier::new(LlmClassifierConfig::Capability {
197                    config: TaskClassifierConfig {
198                        classify_trigger: ClassifyTrigger::EveryRequest,
199                        message_hash_fallback: false,
200                        ..judge
201                    },
202                })?)
203            }
204            CompositeJudgeConfig::Custom(config) => {
205                LlmTaskClassifier::custom_classifier(&config, true)?
206            }
207        };
208        Ok(TierSetter {
209            judge,
210            trigger,
211            message_hash_fallback,
212            default_tier,
213            tiers: Mutex::new(HashMap::new()),
214        })
215    }
216}
217
218struct StageProfiles([FallThrough; 2]);
219
220#[async_trait]
221impl Classifier for StageProfiles {
222    async fn score(
223        &self,
224        state: &mut State,
225        request: &mut Request,
226        driver: &Driver,
227    ) -> Result<(Classification, Option<Response>)> {
228        let (profile, prefix) = match fall_open_tier(state) {
229            Some(Tier::Capable) => (&self.0[1], "composite.profile_1."),
230            _ => (&self.0[0], "composite.profile_0."),
231        };
232        // Each profile retains its own holds, including in caller-carried state.
233        let mut profile_state = State {
234            extra: state
235                .extra
236                .iter()
237                .filter_map(|(key, value)| {
238                    key.strip_prefix(prefix)
239                        .map(|key| (key.to_string(), value.clone()))
240                })
241                .collect(),
242            ..State::default()
243        };
244        let result = profile.route(&mut profile_state, driver, request).await;
245        state.extra.retain(|key, _| !key.starts_with(prefix));
246        state.extra.extend(
247            profile_state
248                .extra
249                .into_iter()
250                .map(|(key, value)| (format!("{prefix}{key}"), value)),
251        );
252        state.tool_signals = profile_state.tool_signals;
253        state.target_prompts = profile_state.target_prompts;
254        let (score, response) = result?;
255        Ok((Classification::Scores(vec![score]), response))
256    }
257}
258
259#[async_trait]
260impl Algorithm for CompositeRouter {
261    fn name(&self) -> &str {
262        COMPOSITE
263    }
264
265    async fn route(
266        self: Arc<Self>,
267        driver: Driver,
268        request: Request,
269    ) -> Result<crate::RoutingOutcome> {
270        self.route.execute(driver, request).await
271    }
272}
273
274#[cfg(test)]
275mod tests {
276    use crate::{CapabilityJudgeConfig, LlmCapabilityConfig};
277    use std::collections::HashMap;
278    use std::sync::Arc;
279
280    use switchyard_protocol::{Category, LlmClientError, Message, ModelId, Role};
281
282    use super::*;
283    use crate::algorithms::util::tier_fixtures::{JUDGE, Recorder, turn_request};
284    use crate::core::testing::test_drive_with_models;
285
286    fn runtime_models() -> HashMap<Category, Vec<ModelId>> {
287        [
288            (Category::Judge, vec![ModelId::from(JUDGE)]),
289            (Category::Efficient, vec![ModelId::from("weak")]),
290            (Category::Capable, vec![ModelId::from("strong")]),
291            (
292                Category::Any,
293                vec![ModelId::from("strong"), ModelId::from("weak")],
294            ),
295        ]
296        .into()
297    }
298
299    fn user_turn_request() -> Request {
300        let mut request = turn_request(false);
301        request
302            .llm_request
303            .messages
304            .push(Message::text(Role::User, "now rewrite the parser"));
305        request
306    }
307
308    /// The same request shape with no session ID, so only the hash can key it.
309    fn unkeyed(mut request: Request) -> Request {
310        if let Some(metadata) = request.metadata.as_mut() {
311            metadata.session_id = None;
312        }
313        request
314    }
315
316    fn hash_keyed_router() -> Result<Arc<CompositeRouter>> {
317        Ok(Arc::new(CompositeRouter::new(CompositeRouterConfig {
318            judge: CompositeJudgeConfig::Capability(TaskClassifierConfig {
319                judge: CapabilityJudgeConfig::Llm(LlmCapabilityConfig {
320                    base_threshold: 0.5,
321                    ..LlmCapabilityConfig::default()
322                }),
323                classify_trigger: ClassifyTrigger::UserTurn,
324                message_hash_fallback: true,
325                ..Default::default()
326            }),
327            stage: StageRouterConfig::new(PickerMode::EfficientFirst, 0.5),
328        })?))
329    }
330
331    fn router() -> Result<Arc<CompositeRouter>> {
332        router_with_trigger(ClassifyTrigger::UserTurn)
333    }
334
335    #[tokio::test]
336    async fn capable_first_defaults_to_pinning_efficient_after_a_switch() -> Result<()> {
337        for capable_first in [false, true] {
338            for turns in [None, Some(0), Some(1), Some(-1)] {
339                let mut stage = StageRouterConfig::new(PickerMode::EfficientFirst, 0.5);
340                stage.efficient_hold_turns = turns;
341                let router = router_with_stage(ClassifyTrigger::UserTurn, stage)?;
342                let recorder = Arc::new(Recorder::default());
343                *recorder.judge_p_solve.lock() = if capable_first { 0.0 } else { 1.0 };
344                let mut initial = turn_request(false);
345                initial.llm_request.messages.truncate(1);
346                let mut write = turn_request(false);
347                if let switchyard_protocol::ContentBlock::ToolCall(call) =
348                    &mut write.llm_request.messages[1].content[0]
349                {
350                    call.name = "Write".into();
351                }
352                for request in [
353                    initial,
354                    turn_request(true),
355                    write.clone(),
356                    write.clone(),
357                    write,
358                    turn_request(true),
359                    turn_request(true),
360                    turn_request(true),
361                ] {
362                    test_drive_with_models(
363                        router.clone(),
364                        request,
365                        runtime_models(),
366                        recorder.serve(),
367                    )
368                    .await?;
369                }
370                let actual: Vec<_> = recorder
371                    .routed()
372                    .into_iter()
373                    .map(|call| call.target)
374                    .collect();
375                let initial = if capable_first { "strong" } else { "weak" };
376                let mut expected = vec![initial, "strong", "strong", "strong", "weak"];
377                let count = turns.unwrap_or(if capable_first { -1 } else { 0 });
378                expected.extend((0..3).map(|i| {
379                    if count == -1 || i < count {
380                        "weak"
381                    } else {
382                        "strong"
383                    }
384                }));
385                assert_eq!(
386                    actual, expected,
387                    "capable_first={capable_first}, turns={turns:?}"
388                );
389                assert_eq!(recorder.judge_calls(), 1);
390            }
391        }
392        Ok(())
393    }
394
395    fn router_with_trigger(classify_trigger: ClassifyTrigger) -> Result<Arc<CompositeRouter>> {
396        router_with_stage(
397            classify_trigger,
398            StageRouterConfig::new(PickerMode::EfficientFirst, 0.5),
399        )
400    }
401
402    fn router_with_stage(
403        classify_trigger: ClassifyTrigger,
404        stage: StageRouterConfig,
405    ) -> Result<Arc<CompositeRouter>> {
406        Ok(Arc::new(CompositeRouter::new(CompositeRouterConfig {
407            judge: CompositeJudgeConfig::Capability(TaskClassifierConfig {
408                judge: CapabilityJudgeConfig::Llm(LlmCapabilityConfig {
409                    base_threshold: 0.5,
410                    ..LlmCapabilityConfig::default()
411                }),
412                classify_trigger,
413                ..Default::default()
414            }),
415            stage,
416        })?))
417    }
418
419    #[tokio::test]
420    async fn profiles_reset_holds_and_keep_prompts_when_the_judge_switches() -> Result<()> {
421        for (has_session_id, picker, first_target, first_prompt) in [
422            (false, PickerMode::EfficientFirst, "weak", "EF efficient"),
423            (true, PickerMode::EfficientFirst, "weak", "EF efficient"),
424            (true, PickerMode::CapableFirst, "strong", "EF capable"),
425        ] {
426            let mut efficient = StageRouterConfig::new(picker, 0.5);
427            efficient.capable_hold_turns = -1;
428            efficient.capable_system_prompt = Some("EF capable".into());
429            efficient.efficient_system_prompt = Some("EF efficient".into());
430            let mut capable = StageRouterConfig::new(PickerMode::CapableFirst, 0.0);
431            capable.capable_hold_turns = 0;
432            capable.capable_system_prompt = Some("CF capable".into());
433            capable.efficient_system_prompt = Some("CF efficient".into());
434            for stage in [&mut efficient, &mut capable] {
435                stage.llm_fallback = Some(crate::algorithms::stage::LlmFallback {
436                    target: None,
437                    config: Default::default(),
438                });
439            }
440            let router = Arc::new(CompositeRouter::new_with_profiles(
441                CompositeJudgeConfig::Capability(TaskClassifierConfig {
442                    judge: CapabilityJudgeConfig::Llm(LlmCapabilityConfig {
443                        base_threshold: 0.5,
444                        ..Default::default()
445                    }),
446                    classify_trigger: ClassifyTrigger::UserTurn,
447                    ..Default::default()
448                }),
449                efficient,
450                capable,
451            )?);
452            let recorder = Arc::new(Recorder::default());
453            let mut write = turn_request(false);
454            if let switchyard_protocol::ContentBlock::ToolCall(call) =
455                &mut write.llm_request.messages[1].content[0]
456            {
457                call.name = "Write".into();
458            }
459            let mut state = State::default();
460            for (p_solve, request, target, prompt) in [
461                (1.0, user_turn_request(), first_target, first_prompt),
462                (1.0, turn_request(true), "strong", "EF capable"),
463                (0.0, user_turn_request(), "strong", "CF capable"),
464                (0.0, write, "weak", "CF efficient"),
465                (1.0, user_turn_request(), first_target, first_prompt),
466                (0.0, user_turn_request(), "strong", "CF capable"),
467            ] {
468                *recorder.judge_p_solve.lock() = p_solve;
469                let request = if has_session_id {
470                    request
471                } else {
472                    unkeyed(request)
473                };
474                let outcome = crate::core::testing::test_drive_with_state(
475                    router.clone(),
476                    request,
477                    runtime_models(),
478                    state,
479                    recorder.serve(),
480                )
481                .await?;
482                assert_eq!(outcome.selected_model_id()?.as_str(), target);
483                assert_eq!(outcome.request.llm_request.instructions.len(), 1);
484                assert_eq!(
485                    outcome.request.llm_request.instructions[0].content,
486                    vec![switchyard_protocol::ContentBlock::Text {
487                        text: prompt.into()
488                    }]
489                );
490                state = outcome.state;
491            }
492            assert_eq!(recorder.judge_calls(), 4);
493        }
494        Ok(())
495    }
496
497    #[test]
498    fn rejects_every_request_as_a_trigger() {
499        for judge in [
500            CompositeJudgeConfig::Capability(TaskClassifierConfig::default()),
501            CompositeJudgeConfig::Custom(CustomClassifierConfig::new(
502                "Choose a tier",
503                serde_json::json!({"type": "object"}),
504                crate::CustomClassifierPolicy::target_selector("/decision/target"),
505            )),
506        ] {
507            let error = CompositeRouter::new(CompositeRouterConfig {
508                judge,
509                stage: StageRouterConfig::new(PickerMode::EfficientFirst, 0.5),
510            })
511            .err()
512            .expect("every_request must be rejected for both classifier modes");
513            assert!(matches!(error, LibsyError::AlgorithmError { message }
514                if message == "composite: classify_trigger must be user_turn or new_session"));
515        }
516    }
517
518    fn custom_judge(trigger: ClassifyTrigger) -> CompositeJudgeConfig {
519        let mut config = CustomClassifierConfig::new(
520            "Select a tier",
521            serde_json::json!({"type": "object"}),
522            crate::CustomClassifierPolicy::target_selector("/decision/target"),
523        );
524        config.classify_trigger = trigger;
525        CompositeJudgeConfig::Custom(config)
526    }
527
528    #[tokio::test]
529    async fn a_failed_session_judgment_retains_the_configured_stage_default() -> Result<()> {
530        for custom in [false, true] {
531            for (mode, expected) in [
532                (PickerMode::EfficientFirst, "weak"),
533                (PickerMode::CapableFirst, "strong"),
534            ] {
535                let judge = if custom {
536                    custom_judge(ClassifyTrigger::NewSession)
537                } else {
538                    CompositeJudgeConfig::Capability(TaskClassifierConfig {
539                        classify_trigger: ClassifyTrigger::NewSession,
540                        ..Default::default()
541                    })
542                };
543                let router = Arc::new(CompositeRouter::new(CompositeRouterConfig {
544                    judge,
545                    stage: StageRouterConfig::new(mode, 0.5),
546                })?);
547                let calls = Arc::new(Mutex::new(0));
548                let serve = {
549                    let calls = calls.clone();
550                    move |target: ModelId, _request: Request| {
551                        let completion = if target.as_str() == JUDGE {
552                            *calls.lock() += 1;
553                            "not JSON".to_string()
554                        } else {
555                            target.to_string()
556                        };
557                        async move { Ok(crate::core::testing::reply(completion)) }
558                    }
559                };
560                for request in [
561                    user_turn_request(),
562                    turn_request(false),
563                    user_turn_request(),
564                ] {
565                    let (selected, _) = test_drive_with_models(
566                        router.clone(),
567                        request,
568                        runtime_models(),
569                        serve.clone(),
570                    )
571                    .await?;
572                    assert_eq!(selected.as_str(), expected);
573                }
574                assert_eq!(
575                    *calls.lock(),
576                    1,
577                    "new_session must retain a failed judgment"
578                );
579            }
580        }
581        Ok(())
582    }
583
584    #[tokio::test]
585    async fn a_custom_judge_sets_the_tier_and_retains_it_on_failure() -> Result<()> {
586        let router = Arc::new(CompositeRouter::new(CompositeRouterConfig {
587            judge: custom_judge(ClassifyTrigger::UserTurn),
588            stage: StageRouterConfig::new(PickerMode::EfficientFirst, 0.5),
589        })?);
590        let replies = Arc::new(Mutex::new(
591            [
592                Some(r#"{"decision":{"target":"capable"}}"#),
593                None,
594                Some("not JSON"),
595                Some(r#"{"decision":{"target":"efficient"}}"#),
596            ]
597            .into_iter(),
598        ));
599        let serve = {
600            let replies = replies.clone();
601            move |target: ModelId, _request: Request| {
602                let completion = if target.as_str() == JUDGE {
603                    replies
604                        .lock()
605                        .next()
606                        .expect("unexpected judge call")
607                        .map(str::to_string)
608                } else {
609                    Some(target.to_string())
610                };
611                async move {
612                    completion
613                        .map(crate::core::testing::reply)
614                        .ok_or_else(|| LlmClientError::General("judge unavailable".into()))
615                }
616            }
617        };
618        for (request, expected) in [
619            (user_turn_request(), "strong"),
620            (turn_request(false), "strong"),
621            (user_turn_request(), "strong"),
622            (user_turn_request(), "strong"),
623            (user_turn_request(), "weak"),
624        ] {
625            let (selected, _) =
626                test_drive_with_models(router.clone(), request, runtime_models(), serve.clone())
627                    .await?;
628            assert_eq!(selected.as_str(), expected);
629        }
630        assert!(replies.lock().next().is_none());
631        Ok(())
632    }
633
634    /// A route may point its judge at the same target it serves capable turns on.
635    /// The tier then cannot be recovered by looking the served model up in the
636    /// runtime groups — one id, two categories — so it comes off the verdict itself.
637    #[tokio::test]
638    async fn a_judge_sharing_the_capable_model_still_latches_the_tier() -> Result<()> {
639        let models: HashMap<Category, Vec<ModelId>> = [
640            (Category::Judge, vec![ModelId::from("strong")]),
641            (Category::Capable, vec![ModelId::from("strong")]),
642            (Category::Efficient, vec![ModelId::from("weak")]),
643            (
644                Category::Any,
645                vec![ModelId::from("strong"), ModelId::from("weak")],
646            ),
647        ]
648        .into();
649        // The judge runs first as a request-side processor, so the opening call is its own.
650        let calls = Arc::new(Mutex::new(0u32));
651        let serve = {
652            let calls = Arc::clone(&calls);
653            move |target: ModelId, _request: Request| {
654                let calls = Arc::clone(&calls);
655                async move {
656                    let mut calls = calls.lock();
657                    *calls += 1;
658                    let completion = if *calls == 1 {
659                        // p_solve below the threshold: the judge does not trust the
660                        // efficient tier, so the verdict is capable.
661                        r#"{"crux":"bounded task","primary_rule":"SUP-1","capability_boundary":"supported","p_solve":0.1}"#.to_string()
662                    } else {
663                        target.to_string()
664                    };
665                    Ok(crate::core::testing::reply(completion))
666                }
667            }
668        };
669        let router = router()?;
670
671        test_drive_with_models(
672            router.clone(),
673            user_turn_request(),
674            models.clone(),
675            serve.clone(),
676        )
677        .await?;
678        // A tool step is not a user turn, so nothing re-judges and the latched tier decides.
679        let (selected, _) =
680            test_drive_with_models(router, turn_request(false), models, serve).await?;
681
682        assert_eq!(selected, ModelId::from("strong"));
683        Ok(())
684    }
685
686    #[tokio::test]
687    async fn a_session_without_an_id_keys_on_the_message_hash() -> Result<()> {
688        let recorder = Arc::new(Recorder::default());
689        *recorder.judge_p_solve.lock() = 0.1;
690        let router = hash_keyed_router()?;
691
692        test_drive_with_models(
693            router.clone(),
694            unkeyed(user_turn_request()),
695            runtime_models(),
696            recorder.serve(),
697        )
698        .await?;
699        test_drive_with_models(
700            router.clone(),
701            unkeyed(turn_request(false)),
702            runtime_models(),
703            recorder.serve(),
704        )
705        .await?;
706
707        assert_eq!(
708            recorder.judge_calls(),
709            1,
710            "a tool step is not a user turn, session id or not"
711        );
712        assert_eq!(
713            recorder.routed()[1].target,
714            "strong",
715            "and the tier survives the tool step"
716        );
717        Ok(())
718    }
719
720    #[tokio::test]
721    async fn the_judge_sets_the_tier_once_a_turn_and_the_signals_run_within_it() -> Result<()> {
722        let recorder = Arc::new(Recorder::default());
723        *recorder.judge_p_solve.lock() = 0.1;
724        let router = router()?;
725
726        test_drive_with_models(
727            router.clone(),
728            user_turn_request(),
729            runtime_models(),
730            recorder.serve(),
731        )
732        .await?;
733        test_drive_with_models(
734            router.clone(),
735            turn_request(false),
736            runtime_models(),
737            recorder.serve(),
738        )
739        .await?;
740
741        let routed = recorder.routed();
742        assert_eq!(
743            routed[0].target, "strong",
744            "a quiet turn falls open to the verdict"
745        );
746        assert_eq!(
747            routed[1].target, "strong",
748            "which holds across the tool steps after it"
749        );
750        assert_eq!(
751            recorder.judge_calls(),
752            1,
753            "a tool step is not a new user turn"
754        );
755        Ok(())
756    }
757
758    #[tokio::test]
759    async fn caller_state_carries_the_judged_tier_without_a_session_id() -> Result<()> {
760        for trigger in [ClassifyTrigger::UserTurn, ClassifyTrigger::NewSession] {
761            let recorder = Arc::new(Recorder::default());
762            *recorder.judge_p_solve.lock() = 0.1;
763            let router = router_with_trigger(trigger)?;
764            let run = |request, state| {
765                crate::core::testing::test_drive_with_state(
766                    router.clone(),
767                    request,
768                    runtime_models(),
769                    state,
770                    recorder.serve(),
771                )
772            };
773
774            let first = run(unkeyed(user_turn_request()), State::default()).await?;
775            assert_eq!(first.selected_model_id()?, &ModelId::from("strong"));
776
777            let continued = run(unkeyed(turn_request(false)), first.state).await?;
778            assert_eq!(continued.selected_model_id()?, &ModelId::from("strong"));
779            assert_eq!(
780                continued.metadata.and_then(|m| m.evidence),
781                Some(serde_json::json!({"source": "retained"}))
782            );
783            assert_eq!(recorder.judge_calls(), 1, "{trigger:?}");
784        }
785        Ok(())
786    }
787}