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. The judge is
14//! asked per turn and its verdict is never pinned to the session.
15//!
16use std::sync::Arc;
17
18use async_trait::async_trait;
19
20use super::fall_through::{DefaultTarget, FallThrough};
21use super::llm_class::{LlmClassifierConfig, LlmTaskClassifier, TaskClassifierConfig};
22use super::util::prompts::{SystemPromptProcessor, TargetPrompts};
23use super::util::stage::{
24    DecisionSource, HandoffNoteConfig, PickerMode, StageClassifier, StageTargets,
25    record_decision_source, record_routing_decision,
26};
27use super::util::tool_signals::{DEFAULT_RECENT_WINDOW, ToolSignalProcessor};
28use crate::core::algorithm::{Algorithm, Driver, LlmTarget, LlmTargetSet};
29use crate::core::classifier::{Classification, Classifier};
30use crate::core::state::State;
31use crate::{LibsyError, Result};
32use switchyard_protocol::{Request, Response};
33
34/// Telemetry name for a router this module assembles.
35const STAGE_ROUTER: &str = "stage_router";
36
37/// Attributes a turn to the classifier it wraps, when that classifier decides it.
38///
39/// The classifiers themselves are composition-agnostic and write no state; only
40/// this router knows where each sits in its cascade.
41struct SourceStamp {
42    inner: Arc<dyn Classifier<State>>,
43    source: DecisionSource,
44}
45
46#[async_trait]
47impl Classifier<State> for SourceStamp {
48    fn routing_tier(&self, selected_model_id: &str) -> Option<&'static str> {
49        self.inner.routing_tier(selected_model_id)
50    }
51
52    async fn score(
53        &self,
54        state: &mut State,
55        request: &mut Request,
56        driver: Option<&Driver>,
57    ) -> Result<(Classification, Option<Response>)> {
58        let (classification, served) = self.inner.score(state, request, driver).await?;
59        // An abstaining classifier passes the turn on, so it is not its to claim.
60        if let Some(winner) = classification.argmax(false)? {
61            record_decision_source(state, self.source);
62            record_routing_decision(self.source, &winner.target);
63        }
64        Ok((classification, served))
65    }
66}
67
68/// The capability judge a stage router falls through to.
69pub struct LlmFallback {
70    /// Target the judge model is called through. It is not a routing
71    /// destination, so it does not belong in the router's target set.
72    pub judge_target: LlmTarget,
73    /// Judge configuration. `recent_turn_window` is worth setting to this router's
74    /// `recent_window` so the judge reads the same span the signal scorer scored.
75    /// Note: `session_affinity` and `message_hash_fallback` have no effect here —
76    /// the judge runs as a cascade classifier, not a standalone algorithm.
77    pub config: TaskClassifierConfig,
78}
79
80/// How a stage router scores turns, and what it hands the model it picks.
81pub struct StageRouterConfig {
82    /// Tier a turn falls open to when the scorer is not confident.
83    pub mode: PickerMode,
84    /// How much corroboration a decisive pick needs, in `[0.0, 1.0]`.
85    pub confidence_threshold: f64,
86    /// Trailing tool results the signals are computed over. `None` uses
87    /// [`DEFAULT_RECENT_WINDOW`].
88    pub recent_window: Option<usize>,
89    /// Note handed to the model on a signal-driven escalation, and on a
90    /// hand-back to the efficient tier when a de-escalation note is configured.
91    pub handoff_notes: Option<HandoffNoteConfig>,
92    /// System prompts keyed by target, handed over on every turn that target
93    /// serves. Empty by default.
94    pub tier_prompts: TargetPrompts,
95    /// Capability judge consulted on turns the signals leave undecided — the
96    /// judge's own target, plus the same configuration the standalone capability
97    /// route takes.
98    pub llm_fallback: Option<LlmFallback>,
99}
100
101impl StageRouterConfig {
102    /// The signal-only configuration: no notes, no per-tier prompts, no judge.
103    /// Set the optional fields to add them.
104    pub fn new(mode: PickerMode, confidence_threshold: f64) -> Self {
105        Self {
106            mode,
107            confidence_threshold,
108            recent_window: None,
109            handoff_notes: None,
110            tier_prompts: TargetPrompts::default(),
111            llm_fallback: None,
112        }
113    }
114}
115
116/// Routes coding-agent turns between a capable and an efficient tier: tool signals
117/// decide first, an optional capability judge takes the turns they cannot, and
118/// the picker's default tier closes the cascade so a turn is never left unrouted.
119pub struct StageRouter {
120    route: FallThrough<State>,
121}
122
123impl StageRouter {
124    /// Routes between the `capable` and `efficient` targets. The
125    /// judge, when configured, is called through its own target and is not a
126    /// routing destination.
127    ///
128    /// Errors if either threshold in `config` is outside `[0.0, 1.0]`.
129    pub fn new(
130        capable: LlmTarget,
131        efficient: LlmTarget,
132        config: StageRouterConfig,
133    ) -> Result<Self> {
134        Ok(Self {
135            route: build_route(capable, efficient, config)?,
136        })
137    }
138}
139
140#[async_trait]
141impl Algorithm for StageRouter {
142    fn name(&self) -> &str {
143        STAGE_ROUTER
144    }
145
146    async fn route(self: Arc<Self>, driver: Driver, request: Request) -> Result<Response> {
147        self.route.execute(driver, request).await
148    }
149}
150
151/// Wires the cascade the wrapper drives.
152fn build_route(
153    capable: LlmTarget,
154    efficient: LlmTarget,
155    config: StageRouterConfig,
156) -> Result<FallThrough<State>> {
157    if !(0.0..=1.0).contains(&config.confidence_threshold) {
158        return Err(LibsyError::AlgorithmError {
159            message: format!(
160                "confidence_threshold must be between 0 and 1, got {}",
161                config.confidence_threshold
162            ),
163        });
164    }
165    // The tiers are a fixed pair; their targets are whatever the deployment calls
166    // them, and the classifier scores onto those names.
167    let targets = StageTargets::new(
168        capable.semantic_name.clone(),
169        efficient.semantic_name.clone(),
170    );
171    // The picker's mode fixes the fallback tier up front, so the terminal
172    // classifier is a constant rather than a per-turn lookup.
173    let fall_open = targets.name(config.mode.default_tier()).to_string();
174
175    let mut classifier = StageClassifier::new(targets, config.mode, config.confidence_threshold);
176    if let Some(notes) = config.handoff_notes {
177        classifier = classifier.with_handoff_notes(notes);
178    }
179    let signals = ToolSignalProcessor {
180        recent_window: config.recent_window.unwrap_or(DEFAULT_RECENT_WINDOW),
181    };
182    let target_set = LlmTargetSet::new(vec![capable.clone(), efficient.clone()]);
183    let mut router = FallThrough::<State>::new_with_state(target_set)
184        .with_name(STAGE_ROUTER)
185        .with_processor(Arc::new(signals))
186        .with_classifier(Arc::new(classifier));
187    if let Some(fallback) = config.llm_fallback {
188        // The capability judge takes its tiers in the same order the capability
189        // route passes them: efficient first, capable second.
190        router = router.with_classifier(Arc::new(SourceStamp {
191            inner: Arc::new(LlmTaskClassifier::new(LlmClassifierConfig::Capability {
192                judge_target: fallback.judge_target,
193                efficient_target: efficient,
194                capable_target: capable,
195                config: fallback.config,
196            })?),
197            source: DecisionSource::LlmClassifier,
198        }));
199    }
200    // Nothing behind this, so the turn lands on the picker's default tier —
201    // including when the judge could not tell.
202    router = router.with_classifier(Arc::new(SourceStamp {
203        inner: Arc::new(DefaultTarget::new(fall_open)),
204        source: DecisionSource::FallOpen,
205    }));
206    // Runs on the post-decision hook, so it applies to the target the cascade
207    // settled on, whichever classifier picked it. With no prompts configured it
208    // is a no-op, so there is nothing to branch on.
209    router = router.with_processor(Arc::new(SystemPromptProcessor::new(config.tier_prompts)));
210    Ok(router)
211}
212
213#[cfg(test)]
214mod tests {
215    use std::sync::Arc;
216
217    use async_trait::async_trait;
218    use parking_lot::Mutex;
219    use serde_json::json;
220    use switchyard_protocol::{
221        ContentBlock, LlmRequest, Message, Role, ToolCall, ToolResult, WireFormat,
222    };
223
224    use super::*;
225    use crate::algorithms::util::stage::DECISION_SOURCE_KEY;
226    use crate::core::algorithm::LlmTarget;
227    use crate::core::classifier::Score;
228    use crate::core::state::StateValue;
229    use crate::core::testing::{Serve, reply, test_drive};
230    use switchyard_protocol::{Decision, Metadata, Response};
231
232    fn tier_target(name: &str) -> LlmTarget {
233        LlmTarget {
234            semantic_name: name.to_string(),
235        }
236    }
237
238    /// A classifier that always picks `target`, standing in for a cascade member.
239    struct Fixed(&'static str);
240
241    #[async_trait]
242    impl Classifier<State> for Fixed {
243        async fn score(
244            &self,
245            _state: &mut State,
246            _request: &mut Request,
247            _driver: Option<&Driver>,
248        ) -> Result<(Classification, Option<Response>)> {
249            Ok((
250                Classification::Scores(vec![Score {
251                    target: self.0.to_string(),
252                    confidence: 1.0,
253                }]),
254                None,
255            ))
256        }
257    }
258
259    /// A classifier that never decides.
260    struct Abstains;
261
262    #[async_trait]
263    impl Classifier<State> for Abstains {
264        async fn score(
265            &self,
266            _state: &mut State,
267            _request: &mut Request,
268            _driver: Option<&Driver>,
269        ) -> Result<(Classification, Option<Response>)> {
270            Ok((Classification::Ambiguous(vec![]), None))
271        }
272    }
273
274    async fn stamped(inner: Arc<dyn Classifier<State>>) -> Result<Option<String>> {
275        let stamp = SourceStamp {
276            inner,
277            source: DecisionSource::LlmClassifier,
278        };
279        let mut state = State::default();
280        stamp
281            .score(&mut state, &mut Request::default(), None)
282            .await?;
283        Ok(match state.extra.get(DECISION_SOURCE_KEY) {
284            Some(StateValue::String(source)) => Some(source.clone()),
285            _ => None,
286        })
287    }
288
289    #[tokio::test]
290    async fn a_deciding_classifier_is_credited_with_the_turn() -> Result<()> {
291        assert_eq!(
292            stamped(Arc::new(Fixed("strong"))).await?.as_deref(),
293            Some("llm-classifier")
294        );
295        Ok(())
296    }
297
298    #[tokio::test]
299    async fn an_abstaining_classifier_claims_nothing() -> Result<()> {
300        // It passed the turn on, so the next classifier is the one that decided.
301        assert_eq!(stamped(Arc::new(Abstains)).await?, None);
302        Ok(())
303    }
304
305    fn config() -> StageRouterConfig {
306        StageRouterConfig::new(PickerMode::EfficientFirst, 0.5)
307    }
308
309    #[test]
310    fn rejects_an_out_of_range_confidence_threshold() {
311        let mut config = config();
312        config.confidence_threshold = 1.5;
313        assert!(matches!(
314            StageRouter::new(tier_target("strong"), tier_target("weak"), config),
315            Err(LibsyError::AlgorithmError { .. })
316        ));
317    }
318
319    #[test]
320    fn rejects_an_out_of_range_judge_threshold() {
321        let mut config = config();
322        config.llm_fallback = Some(LlmFallback {
323            judge_target: LlmTarget {
324                semantic_name: "judge".to_string(),
325            },
326            config: TaskClassifierConfig {
327                base_threshold: -0.1,
328                ..Default::default()
329            },
330        });
331        assert!(matches!(
332            StageRouter::new(tier_target("strong"), tier_target("weak"), config),
333            Err(LibsyError::AlgorithmError { .. })
334        ));
335    }
336
337    #[test]
338    fn builds_over_both_tiers() -> Result<()> {
339        let router = StageRouter::new(tier_target("strong"), tier_target("weak"), config())?;
340        assert_eq!(router.name(), STAGE_ROUTER);
341        Ok(())
342    }
343
344    // ── routing integration tests ────────────────────────────────────────────
345
346    const ESCALATION: &str = "the previous model was stalling; pick up the diagnosis";
347    const JUDGE: &str = "judge";
348
349    #[derive(Clone, Debug)]
350    struct Call {
351        target: String,
352        messages: Vec<String>,
353        is_answer_call: bool,
354    }
355
356    /// Records what each target receives.
357    #[derive(Default)]
358    struct Recorder {
359        calls: Mutex<Vec<Call>>,
360        judge_p_solve: Mutex<f64>,
361    }
362
363    impl Recorder {
364        fn routed(&self) -> Vec<Call> {
365            self.calls
366                .lock()
367                .iter()
368                .filter(|call| call.is_answer_call)
369                .cloned()
370                .collect()
371        }
372
373        /// Serves every call, recording it. The judge target gets a structured verdict
374        /// back so the fallback classifier has an answer without a real model.
375        fn serve(self: &Arc<Self>) -> impl Serve {
376            let recorder = Arc::clone(self);
377            move |decision: Decision, request: Request| {
378                let recorder = Arc::clone(&recorder);
379                async move {
380                    let target = decision.selected_model_id().to_string();
381                    recorder.calls.lock().push(Call {
382                        target: target.clone(),
383                        messages: request
384                            .llm_request
385                            .messages
386                            .iter()
387                            .filter_map(|message| message.text_content("|"))
388                            .collect(),
389                        is_answer_call: decision.is_answer_call(),
390                    });
391                    let completion = if target == JUDGE {
392                        let p_solve = *recorder.judge_p_solve.lock();
393                        format!(
394                            r#"{{"crux":"bounded task","primary_rule":"SUP-1","capability_boundary":"supported","p_solve":{p_solve}}}"#
395                        )
396                    } else {
397                        target
398                    };
399                    Ok(reply(completion))
400                }
401            }
402        }
403    }
404
405    fn recording_target(name: &str) -> LlmTarget {
406        LlmTarget {
407            semantic_name: name.to_string(),
408        }
409    }
410
411    fn recording_router(config: StageRouterConfig) -> Result<Arc<StageRouter>> {
412        Ok(Arc::new(StageRouter::new(
413            recording_target("strong"),
414            recording_target("weak"),
415            config,
416        )?))
417    }
418
419    fn config_with_notes() -> StageRouterConfig {
420        let mut c = config();
421        c.handoff_notes = Some(HandoffNoteConfig::new(ESCALATION, None, true));
422        c
423    }
424
425    fn config_with_judge(recorder: &Arc<Recorder>, p_solve: f64) -> StageRouterConfig {
426        *recorder.judge_p_solve.lock() = p_solve;
427        let mut c = config();
428        c.llm_fallback = Some(LlmFallback {
429            judge_target: recording_target(JUDGE),
430            config: TaskClassifierConfig {
431                base_threshold: 0.5,
432                recent_turn_window: Some(3),
433                ..Default::default()
434            },
435        });
436        c
437    }
438
439    fn turn_request(failed: bool) -> Request {
440        let content = if failed {
441            "fatal runtime error: out of memory"
442        } else {
443            "ok"
444        };
445        Request {
446            llm_request: LlmRequest {
447                model: Some("auto".to_string()),
448                messages: vec![
449                    Message::text(Role::User, "fix the build"),
450                    Message {
451                        role: Role::Assistant,
452                        content: vec![ContentBlock::ToolCall(ToolCall {
453                            id: "call_1".to_string(),
454                            name: "Bash".to_string(),
455                            arguments: json!({"command": "cargo test"}),
456                        })],
457                    },
458                    Message {
459                        role: Role::Tool,
460                        content: vec![ContentBlock::ToolResult(ToolResult {
461                            tool_call_id: "call_1".to_string(),
462                            content: vec![ContentBlock::Text {
463                                text: content.to_string(),
464                            }],
465                            is_error: Some(failed),
466                        })],
467                    },
468                ],
469                ..LlmRequest::default()
470            },
471            raw_request: Some(json!({
472                "model": "auto",
473                "messages": [
474                    {"role": "user", "content": "fix the build"},
475                    {"role": "assistant", "tool_calls": [{"id": "call_1", "type": "function",
476                        "function": {"name": "Bash", "arguments": "{\"command\": \"cargo test\"}"}}]},
477                    {"role": "tool", "tool_call_id": "call_1", "content": content},
478                ],
479            })),
480            metadata: Some(Metadata {
481                wire_format: Some(WireFormat::OpenAiChat),
482                session_id: Some("session-1".to_string()),
483                ..Default::default()
484            }),
485        }
486    }
487
488    #[tokio::test]
489    async fn a_signal_driven_escalation_hands_the_note_to_the_model() -> Result<()> {
490        let recorder = Arc::new(Recorder::default());
491        let router = recording_router(config_with_notes())?;
492
493        test_drive(router.clone(), turn_request(false), recorder.serve()).await?;
494        test_drive(router.clone(), turn_request(true), recorder.serve()).await?;
495
496        let calls = recorder.routed();
497        assert_eq!(calls[0].target, "weak");
498        assert_eq!(calls[1].target, "strong");
499        assert!(
500            !calls[0].messages.iter().any(|t| t.contains(ESCALATION)),
501            "steady-state turn should carry no note: {:?}",
502            calls[0].messages
503        );
504        assert!(
505            calls[1]
506                .messages
507                .last()
508                .is_some_and(|t| t.ends_with(ESCALATION)),
509            "escalating turn should carry the note last: {:?}",
510            calls[1].messages
511        );
512        Ok(())
513    }
514
515    #[tokio::test]
516    async fn the_judge_decides_a_turn_the_signals_leave_undecided() -> Result<()> {
517        let recorder = Arc::new(Recorder::default());
518        let router = recording_router(config_with_judge(&recorder, 0.1))?;
519
520        let (trace, _) = test_drive(router.clone(), turn_request(false), recorder.serve()).await?;
521
522        let calls = recorder.calls.lock();
523        assert!(
524            calls
525                .iter()
526                .any(|call| call.target == JUDGE && !call.is_answer_call),
527            "the judge should be recorded as a routing side call"
528        );
529        assert!(
530            calls
531                .iter()
532                .any(|call| call.target == "strong" && call.is_answer_call),
533            "the selected target should be recorded as an answer call"
534        );
535        drop(calls);
536        assert!(
537            trace
538                .last()
539                .and_then(|decision| decision.reasoning())
540                .is_some_and(|reasoning| reasoning.contains("routing tier: strong"))
541        );
542        Ok(())
543    }
544
545    #[tokio::test]
546    async fn a_decisive_signal_never_reaches_the_judge() -> Result<()> {
547        let recorder = Arc::new(Recorder::default());
548        let router = recording_router(config_with_judge(&recorder, 0.9))?;
549
550        test_drive(router.clone(), turn_request(true), recorder.serve()).await?;
551
552        assert!(
553            !recorder.calls.lock().iter().any(|c| c.target == JUDGE),
554            "a resolved turn should not pay for a judge call"
555        );
556        assert_eq!(recorder.routed()[0].target, "strong");
557        Ok(())
558    }
559
560    #[tokio::test]
561    async fn the_judges_verdict_is_not_pinned_to_the_session() -> Result<()> {
562        let recorder = Arc::new(Recorder::default());
563        let router = recording_router(config_with_judge(&recorder, 0.1))?;
564
565        test_drive(router.clone(), turn_request(false), recorder.serve()).await?;
566        *recorder.judge_p_solve.lock() = 0.9;
567        test_drive(router.clone(), turn_request(false), recorder.serve()).await?;
568
569        let routed = recorder.routed();
570        assert_eq!(routed[0].target, "strong");
571        assert_eq!(routed[1].target, "weak");
572        assert_eq!(
573            recorder
574                .calls
575                .lock()
576                .iter()
577                .filter(|c| c.target == JUDGE)
578                .count(),
579            2,
580            "each undecided turn is its own question"
581        );
582        Ok(())
583    }
584
585    #[tokio::test]
586    async fn a_judge_that_cannot_tell_lands_on_the_picker_default() -> Result<()> {
587        let recorder = Arc::new(Recorder::default());
588        let router = recording_router(config_with_judge(&recorder, 42.0))?;
589
590        test_drive(router.clone(), turn_request(false), recorder.serve()).await?;
591
592        assert_eq!(recorder.routed()[0].target, "weak");
593        Ok(())
594    }
595
596    #[tokio::test]
597    async fn the_judge_reads_the_window_it_was_configured_with() -> Result<()> {
598        let recorder = Arc::new(Recorder::default());
599        let router = recording_router(config_with_judge(&recorder, 0.9))?;
600
601        test_drive(router.clone(), turn_request(false), recorder.serve()).await?;
602
603        let judged = recorder
604            .calls
605            .lock()
606            .iter()
607            .find(|c| c.target == JUDGE)
608            .map(|c| c.messages.join("|"));
609        let Some(judged) = judged else {
610            panic!("the judge was never called");
611        };
612        assert!(
613            judged.contains("fix the build"),
614            "the judge should see the opening task: {judged}"
615        );
616        Ok(())
617    }
618}
switchyard_libsy/algorithms/stage.rs

switchyard_libsy/algorithms/
stage.rs
