switchyard_libsy/algorithms/
stage.rs

1// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2// SPDX-License-Identifier: Apache-2.0
3
4//! Signal-driven stage routing for coding agents.
5//!
6//! [`StageRouter`] is the assembled algorithm: a [`FallThrough`] pre-wired with
7//! the tool-signal processor that reads each turn's tool results and the
8//! [`StageClassifier`] that scores them onto the capable/efficient tiers. The cascade
9//! is an internal detail — callers drive the algorithm, not its parts.
10//!
11//! Signals do not decide every turn. An under-threshold turn abstains and falls
12//! through to the optional [`LlmTaskClassifier`] — the capability route's judge,
13//! joined in unchanged — and then to the picker's default tier, or to an
14//! override a decider ahead of stage set in its place.
15//!
16use std::sync::Arc;
17
18use async_trait::async_trait;
19
20use super::fall_through::FallThrough;
21use super::llm_class::{LlmClassifierConfig, LlmTaskClassifier, TaskClassifierConfig};
22use super::util::prompts::prepend_system_prompt;
23use super::util::stage::{
24    DecisionSource, HandoffNoteConfig, PickerMode, StageClassifier, Tier, fall_open_tier,
25    record_decision_source, record_routing_decision,
26};
27use super::util::tool_signals::{DEFAULT_RECENT_WINDOW, ToolSemantics, ToolSignalProcessor};
28use crate::core::algorithm::{Algorithm, Driver};
29use crate::core::classifier::{Classification, Classifier, Score};
30use crate::core::processor::{Event, Processor};
31use crate::core::state::State;
32use crate::{LibsyError, Result};
33use switchyard_protocol::{Category, Request, Response};
34
35/// Telemetry name for a router this module assembles.
36const STAGE_ROUTER: &str = "stage_router";
37
38struct TierPromptProcessor {
39    capable: Option<String>,
40    efficient: Option<String>,
41}
42
43#[async_trait]
44impl Processor<State> for TierPromptProcessor {
45    async fn process(&self, _state: &mut State, event: Event<'_>) -> Result<()> {
46        if let Event::Decision {
47            request, category, ..
48        } = event
49        {
50            let prompt = match category {
51                Some(Category::Capable) => self.capable.as_deref(),
52                Some(Category::Efficient) => self.efficient.as_deref(),
53                _ => None,
54            };
55            if let Some(prompt) = prompt {
56                prepend_system_prompt(request, prompt);
57            }
58        }
59        Ok(())
60    }
61}
62
63/// Attributes a turn to the classifier it wraps, when that classifier decides it.
64///
65/// The classifiers themselves are composition-agnostic and write no state; only
66/// this router knows where each sits in its cascade.
67struct SourceStamp {
68    inner: Arc<dyn Classifier<State>>,
69    source: DecisionSource,
70}
71
72#[async_trait]
73impl Classifier<State> for SourceStamp {
74    async fn score(
75        &self,
76        state: &mut State,
77        request: &mut Request,
78        driver: &Driver,
79    ) -> Result<(Classification, Option<Response>)> {
80        let (classification, served) = self.inner.score(state, request, driver).await?;
81        // An abstaining classifier passes the turn on, so it is not its to claim.
82        if let Some(winner) = classification.argmax(false)? {
83            record_decision_source(state, self.source);
84            record_routing_decision(self.source, &winner.target);
85            driver.set_evidence_if_empty(serde_json::json!({
86                "source": self.source.as_str(),
87            }));
88        }
89        Ok((classification, served))
90    }
91}
92
93/// Closes the cascade at zero confidence: a fallback, not a judgement.
94struct FallOpen {
95    default_tier: Tier,
96}
97
98#[async_trait]
99impl Classifier<State> for FallOpen {
100    async fn score(
101        &self,
102        state: &mut State,
103        _request: &mut Request,
104        driver: &Driver,
105    ) -> Result<(Classification, Option<Response>)> {
106        let tier = fall_open_tier(state).unwrap_or(self.default_tier);
107        let category = match tier {
108            Tier::Capable => Category::Capable,
109            Tier::Efficient => Category::Efficient,
110        };
111        let target = driver.first_model_for(&category)?.clone();
112        Ok((
113            Classification::Scores(vec![Score {
114                target,
115                confidence: 0.0,
116                category: Some(category),
117            }]),
118            None,
119        ))
120    }
121}
122
123/// The capability judge a stage router falls through to.
124pub struct LlmFallback {
125    /// Judge configuration. `recent_turn_window` is worth setting to this router's
126    /// `recent_window` so the judge reads the same span the signal scorer scored.
127    /// Note: `classify_trigger = new_session` and `message_hash_fallback` have no effect here —
128    /// the judge runs as a cascade classifier, not a standalone algorithm.
129    pub config: TaskClassifierConfig,
130}
131
132/// How a stage router scores turns, and what it hands the model it picks.
133pub struct StageRouterConfig {
134    /// Tier a turn falls open to when the scorer is not confident.
135    pub mode: PickerMode,
136    /// How much corroboration a decisive pick needs, in `[0.0, 1.0]`.
137    pub confidence_threshold: f64,
138    /// Trailing tool results the signals are computed over. `None` uses
139    /// [`DEFAULT_RECENT_WINDOW`].
140    pub recent_window: Option<usize>,
141    /// Exact tool-name semantics added to the built-in coding vocabulary.
142    pub tool_semantics: ToolSemantics,
143    /// Requests to keep on the capable tier after an escalation. A clean test
144    /// pass clears the hold early.
145    pub capable_hold_turns: u32,
146    /// Note handed to the model on a signal-driven escalation, and on a
147    /// hand-back to the efficient tier when a de-escalation note is configured.
148    pub handoff_notes: Option<HandoffNoteConfig>,
149    /// System prompt handed to the runtime capable model.
150    pub capable_system_prompt: Option<String>,
151    /// System prompt handed to the runtime efficient model.
152    pub efficient_system_prompt: Option<String>,
153    /// Capability judge consulted on turns the signals leave undecided. It uses
154    /// the runtime judge model and the standalone capability route's settings.
155    pub llm_fallback: Option<LlmFallback>,
156}
157
158impl StageRouterConfig {
159    /// The signal-only configuration: no notes, no per-tier prompts, no judge.
160    /// Set the optional fields to add them.
161    pub fn new(mode: PickerMode, confidence_threshold: f64) -> Self {
162        Self {
163            mode,
164            confidence_threshold,
165            recent_window: None,
166            tool_semantics: ToolSemantics::default(),
167            capable_hold_turns: 2,
168            handoff_notes: None,
169            capable_system_prompt: None,
170            efficient_system_prompt: None,
171            llm_fallback: None,
172        }
173    }
174}
175
176/// Routes coding-agent turns between a capable and an efficient tier: tool signals
177/// decide first, an optional capability judge takes the turns they cannot, and
178/// the picker's default tier closes the cascade so a turn is never left unrouted.
179pub struct StageRouter {
180    route: FallThrough<State>,
181}
182
183impl StageRouter {
184    /// Routes between the runtime `capable` and `efficient` models. The judge,
185    /// when configured, is called through the runtime `judge` model.
186    ///
187    /// Errors if either threshold in `config` is outside `[0.0, 1.0]`.
188    pub fn new(config: StageRouterConfig) -> Result<Self> {
189        Ok(Self {
190            route: build_stage_route(config)?,
191        })
192    }
193}
194
195#[async_trait]
196impl Algorithm for StageRouter {
197    fn name(&self) -> &str {
198        STAGE_ROUTER
199    }
200
201    async fn route(
202        self: Arc<Self>,
203        driver: Driver,
204        request: Request,
205    ) -> Result<crate::RoutingOutcome> {
206        self.route.execute(driver, request).await
207    }
208}
209
210/// Wires the cascade the wrapper drives. Exposed so a composition above can
211/// stack a prelude onto it.
212pub(crate) fn build_stage_route(config: StageRouterConfig) -> Result<FallThrough<State>> {
213    if !(0.0..=1.0).contains(&config.confidence_threshold) {
214        return Err(LibsyError::AlgorithmError {
215            message: format!(
216                "confidence_threshold must be between 0 and 1, got {}",
217                config.confidence_threshold
218            ),
219        });
220    }
221    config.tool_semantics.validate()?;
222    let default_tier = config.mode.default_tier();
223    let fall_open = FallOpen { default_tier };
224
225    let mut classifier = StageClassifier::new(config.mode, config.confidence_threshold)
226        .with_capable_hold_turns(config.capable_hold_turns);
227    if let Some(notes) = config.handoff_notes {
228        classifier = classifier.with_handoff_notes(notes);
229    }
230    let signals = ToolSignalProcessor {
231        recent_window: config.recent_window.unwrap_or(DEFAULT_RECENT_WINDOW),
232        tool_semantics: config.tool_semantics,
233    };
234    let mut router = FallThrough::<State>::new_with_state()
235        .with_name(STAGE_ROUTER)
236        .with_processor(Arc::new(signals))
237        .with_classifier(Arc::new(classifier));
238    if let Some(fallback) = config.llm_fallback {
239        router = router.with_classifier(Arc::new(SourceStamp {
240            inner: Arc::new(LlmTaskClassifier::new(LlmClassifierConfig::Capability {
241                config: fallback.config,
242            })?),
243            source: DecisionSource::LlmClassifier,
244        }));
245    }
246    // Nothing behind this, so no turn is left unrouted.
247    router = router.with_classifier(Arc::new(SourceStamp {
248        inner: Arc::new(fall_open),
249        source: DecisionSource::FallOpen,
250    }));
251    if config.capable_system_prompt.is_some() || config.efficient_system_prompt.is_some() {
252        router = router.with_processor(Arc::new(TierPromptProcessor {
253            capable: config.capable_system_prompt,
254            efficient: config.efficient_system_prompt,
255        }));
256    }
257    Ok(router)
258}
259
260#[cfg(test)]
261mod tests {
262    use std::collections::HashMap;
263    use std::sync::Arc;
264
265    use async_trait::async_trait;
266    use parking_lot::Mutex;
267
268    use super::*;
269    use crate::algorithms::util::stage::{DECISION_SOURCE_KEY, clear_fall_open, set_fall_open};
270    use crate::algorithms::util::tier_fixtures::{JUDGE, Recorder, turn_request};
271    use crate::core::state::StateValue;
272    use crate::core::testing::{empty_driver, test_drive_with_models};
273    use switchyard_protocol::{Category, ModelId};
274
275    /// A classifier that always picks `target`, standing in for a cascade member.
276    struct Fixed(&'static str);
277
278    #[async_trait]
279    impl Classifier<State> for Fixed {
280        async fn score(
281            &self,
282            _state: &mut State,
283            _request: &mut Request,
284            _driver: &Driver,
285        ) -> Result<(Classification, Option<Response>)> {
286            Ok((
287                Classification::Scores(vec![Score {
288                    target: ModelId::from(self.0),
289                    confidence: 1.0,
290                    category: None,
291                }]),
292                None,
293            ))
294        }
295    }
296
297    /// A classifier that never decides.
298    struct Abstains;
299
300    #[async_trait]
301    impl Classifier<State> for Abstains {
302        async fn score(
303            &self,
304            _state: &mut State,
305            _request: &mut Request,
306            _driver: &Driver,
307        ) -> Result<(Classification, Option<Response>)> {
308            Ok((Classification::Ambiguous(vec![]), None))
309        }
310    }
311
312    async fn stamped(inner: Arc<dyn Classifier<State>>) -> Result<Option<String>> {
313        let stamp = SourceStamp {
314            inner,
315            source: DecisionSource::LlmClassifier,
316        };
317        let mut state = State::default();
318        stamp
319            .score(&mut state, &mut Request::default(), &empty_driver())
320            .await?;
321        Ok(match state.extra.get(DECISION_SOURCE_KEY) {
322            Some(StateValue::String(source)) => Some(source.clone()),
323            _ => None,
324        })
325    }
326
327    #[tokio::test]
328    async fn a_deciding_classifier_is_credited_with_the_turn() -> Result<()> {
329        assert_eq!(
330            stamped(Arc::new(Fixed("strong"))).await?.as_deref(),
331            Some("llm-classifier")
332        );
333        Ok(())
334    }
335
336    #[tokio::test]
337    async fn an_abstaining_classifier_claims_nothing() -> Result<()> {
338        // It passed the turn on, so the next classifier is the one that decided.
339        assert_eq!(stamped(Arc::new(Abstains)).await?, None);
340        Ok(())
341    }
342
343    fn config() -> StageRouterConfig {
344        StageRouterConfig::new(PickerMode::EfficientFirst, 0.5)
345    }
346
347    fn runtime_models() -> HashMap<Category, Vec<ModelId>> {
348        runtime_models_for("strong", "weak")
349    }
350
351    fn runtime_models_for(capable: &str, efficient: &str) -> HashMap<Category, Vec<ModelId>> {
352        [
353            (Category::Judge, vec![ModelId::from(JUDGE)]),
354            (Category::Efficient, vec![ModelId::from(efficient)]),
355            (Category::Capable, vec![ModelId::from(capable)]),
356            (
357                Category::Any,
358                vec![ModelId::from(capable), ModelId::from(efficient)],
359            ),
360        ]
361        .into()
362    }
363
364    #[test]
365    fn rejects_an_out_of_range_confidence_threshold() {
366        let mut config = config();
367        config.confidence_threshold = 1.5;
368        assert!(matches!(
369            StageRouter::new(config),
370            Err(LibsyError::AlgorithmError { .. })
371        ));
372    }
373
374    #[test]
375    fn rejects_an_out_of_range_judge_threshold() {
376        let mut config = config();
377        config.llm_fallback = Some(LlmFallback {
378            config: TaskClassifierConfig {
379                base_threshold: -0.1,
380                ..Default::default()
381            },
382        });
383        assert!(matches!(
384            StageRouter::new(config),
385            Err(LibsyError::AlgorithmError { .. })
386        ));
387    }
388
389    #[test]
390    fn builds() -> Result<()> {
391        let router = StageRouter::new(config())?;
392        assert_eq!(router.name(), STAGE_ROUTER);
393        Ok(())
394    }
395
396    // ── routing integration tests ────────────────────────────────────────────
397
398    const ESCALATION: &str = "the previous model was stalling; pick up the diagnosis";
399    fn recording_router(config: StageRouterConfig) -> Result<Arc<StageRouter>> {
400        Ok(Arc::new(StageRouter::new(config)?))
401    }
402
403    fn config_with_notes() -> StageRouterConfig {
404        let mut c = config();
405        c.handoff_notes = Some(HandoffNoteConfig::new(ESCALATION, None, true));
406        c
407    }
408
409    fn config_with_judge(recorder: &Arc<Recorder>, p_solve: f64) -> StageRouterConfig {
410        *recorder.judge_p_solve.lock() = p_solve;
411        let mut c = config();
412        c.llm_fallback = Some(LlmFallback {
413            config: TaskClassifierConfig {
414                base_threshold: 0.5,
415                recent_turn_window: Some(3),
416                ..Default::default()
417            },
418        });
419        c
420    }
421
422    /// Stands in for a decider ahead of stage: sets the override once, clears it
423    /// once, and leaves the turns between alone.
424    #[derive(Default)]
425    struct TierDecider {
426        requests: Mutex<u32>,
427    }
428
429    #[async_trait]
430    impl Processor<State> for TierDecider {
431        async fn process(&self, state: &mut State, event: Event<'_>) -> Result<()> {
432            if matches!(event, Event::Request { .. }) {
433                let mut requests = self.requests.lock();
434                *requests += 1;
435                match *requests {
436                    1 => set_fall_open(state, Tier::Efficient),
437                    4 => clear_fall_open(state),
438                    _ => {}
439                }
440            }
441            Ok(())
442        }
443    }
444
445    #[tokio::test]
446    async fn an_override_replaces_the_picker_default_and_leaves_the_signals_alone() -> Result<()> {
447        let recorder = Arc::new(Recorder::default());
448        // The picker would fall open to "strong"; the override says "weak".
449        let config = StageRouterConfig::new(PickerMode::CapableFirst, 0.5);
450        let route: Arc<dyn Algorithm> =
451            Arc::new(build_stage_route(config)?.with_processor(Arc::new(TierDecider::default())));
452
453        let models = runtime_models();
454        for is_turn_failed in [false, false, true, false] {
455            test_drive_with_models(
456                route.clone(),
457                turn_request(is_turn_failed),
458                models.clone(),
459                recorder.serve(),
460            )
461            .await?;
462        }
463
464        let routed = recorder.routed();
465        assert_eq!(
466            routed[0].target, "weak",
467            "an undecided turn takes the override"
468        );
469        assert_eq!(
470            routed[1].target, "weak",
471            "which outlives the turn that set it"
472        );
473        assert_eq!(
474            routed[2].target, "strong",
475            "a critical failure still reaches the signals"
476        );
477        assert_eq!(
478            routed[3].target, "strong",
479            "clearing restores the picker default"
480        );
481        Ok(())
482    }
483
484    #[tokio::test]
485    async fn a_signal_driven_escalation_hands_the_note_to_the_model() -> Result<()> {
486        let recorder = Arc::new(Recorder::default());
487        let router = recording_router(config_with_notes())?;
488
489        test_drive_with_models(
490            router.clone(),
491            turn_request(false),
492            runtime_models(),
493            recorder.serve(),
494        )
495        .await?;
496        test_drive_with_models(
497            router.clone(),
498            turn_request(true),
499            runtime_models_for("runtime-strong", "runtime-weak"),
500            recorder.serve(),
501        )
502        .await?;
503
504        let calls = recorder.routed();
505        assert_eq!(calls[0].target, "weak");
506        assert_eq!(calls[1].target, "runtime-strong");
507        assert!(
508            !calls[0].messages.iter().any(|t| t.contains(ESCALATION)),
509            "steady-state turn should carry no note: {:?}",
510            calls[0].messages
511        );
512        assert!(
513            calls[1]
514                .messages
515                .last()
516                .is_some_and(|t| t.ends_with(ESCALATION)),
517            "escalating turn should carry the note last: {:?}",
518            calls[1].messages
519        );
520        Ok(())
521    }
522
523    #[tokio::test]
524    async fn the_judge_decides_a_turn_the_signals_leave_undecided() -> Result<()> {
525        let recorder = Arc::new(Recorder::default());
526        let router = recording_router(config_with_judge(&recorder, 0.1))?;
527
528        let (selected_model, _) = test_drive_with_models(
529            router.clone(),
530            turn_request(false),
531            runtime_models(),
532            recorder.serve(),
533        )
534        .await?;
535
536        let calls = recorder.calls.lock();
537        assert!(
538            calls.iter().any(|call| call.target == JUDGE),
539            "the judge should be recorded as a routing side call"
540        );
541        assert!(
542            calls.iter().any(|call| call.target == "strong"),
543            "the selected target should be recorded as an answer call"
544        );
545        drop(calls);
546        assert_eq!(selected_model, "strong");
547        Ok(())
548    }
549
550    #[tokio::test]
551    async fn a_decisive_signal_never_reaches_the_judge() -> Result<()> {
552        let recorder = Arc::new(Recorder::default());
553        let router = recording_router(config_with_judge(&recorder, 0.9))?;
554
555        test_drive_with_models(
556            router.clone(),
557            turn_request(true),
558            runtime_models(),
559            recorder.serve(),
560        )
561        .await?;
562
563        assert!(
564            !recorder.calls.lock().iter().any(|c| c.target == JUDGE),
565            "a resolved turn should not pay for a judge call"
566        );
567        assert_eq!(recorder.routed()[0].target, "strong");
568        Ok(())
569    }
570
571    #[tokio::test]
572    async fn the_judges_verdict_is_not_pinned_to_the_session() -> Result<()> {
573        let recorder = Arc::new(Recorder::default());
574        let router = recording_router(config_with_judge(&recorder, 0.1))?;
575
576        test_drive_with_models(
577            router.clone(),
578            turn_request(false),
579            runtime_models(),
580            recorder.serve(),
581        )
582        .await?;
583        *recorder.judge_p_solve.lock() = 0.9;
584        test_drive_with_models(
585            router.clone(),
586            turn_request(false),
587            runtime_models(),
588            recorder.serve(),
589        )
590        .await?;
591
592        let routed = recorder.routed();
593        assert_eq!(routed[0].target, "strong");
594        assert_eq!(routed[1].target, "weak");
595        assert_eq!(
596            recorder
597                .calls
598                .lock()
599                .iter()
600                .filter(|c| c.target == JUDGE)
601                .count(),
602            2,
603            "each undecided turn is its own question"
604        );
605        Ok(())
606    }
607
608    #[tokio::test]
609    async fn a_judge_that_cannot_tell_lands_on_the_picker_default() -> Result<()> {
610        let recorder = Arc::new(Recorder::default());
611        let router = recording_router(config_with_judge(&recorder, 42.0))?;
612
613        test_drive_with_models(
614            router.clone(),
615            turn_request(false),
616            runtime_models(),
617            recorder.serve(),
618        )
619        .await?;
620
621        assert_eq!(recorder.routed()[0].target, "weak");
622        Ok(())
623    }
624
625    #[tokio::test]
626    async fn the_judge_reads_the_window_it_was_configured_with() -> Result<()> {
627        let recorder = Arc::new(Recorder::default());
628        let router = recording_router(config_with_judge(&recorder, 0.9))?;
629
630        test_drive_with_models(
631            router.clone(),
632            turn_request(false),
633            runtime_models(),
634            recorder.serve(),
635        )
636        .await?;
637
638        let judged = recorder
639            .calls
640            .lock()
641            .iter()
642            .find(|c| c.target == JUDGE)
643            .map(|c| c.messages.join("|"));
644        let Some(judged) = judged else {
645            panic!("the judge was never called");
646        };
647        assert!(
648            judged.contains("fix the build"),
649            "the judge should see the opening task: {judged}"
650        );
651        Ok(())
652    }
653}
switchyard_libsy/algorithms/stage.rs

switchyard_libsy/algorithms/
stage.rs
