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};
29use crate::core::classifier::{Classification, Classifier};
30use crate::core::state::State;
31use crate::{LibsyError, Result};
32use switchyard_protocol::{ModelId, 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: &ModelId) -> 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: ModelId,
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(capable: ModelId, efficient: ModelId, config: StageRouterConfig) -> Result<Self> {
130        Ok(Self {
131            route: build_route(capable, efficient, config)?,
132        })
133    }
134}
135
136#[async_trait]
137impl Algorithm for StageRouter {
138    fn name(&self) -> &str {
139        STAGE_ROUTER
140    }
141
142    async fn route(self: Arc<Self>, driver: Driver, request: Request) -> Result<Response> {
143        self.route.execute(driver, request).await
144    }
145}
146
147/// Wires the cascade the wrapper drives.
148fn build_route(
149    capable: ModelId,
150    efficient: ModelId,
151    config: StageRouterConfig,
152) -> Result<FallThrough<State>> {
153    if !(0.0..=1.0).contains(&config.confidence_threshold) {
154        return Err(LibsyError::AlgorithmError {
155            message: format!(
156                "confidence_threshold must be between 0 and 1, got {}",
157                config.confidence_threshold
158            ),
159        });
160    }
161    // The tiers are a fixed pair; their targets are whatever the deployment calls
162    // them, and the classifier scores onto those names.
163    let targets = StageTargets::new(capable.clone(), efficient.clone());
164    // The picker's mode fixes the fallback tier up front, so the terminal
165    // classifier is a constant rather than a per-turn lookup.
166    let fall_open = targets.name(config.mode.default_tier()).to_string();
167
168    let mut classifier = StageClassifier::new(targets, config.mode, config.confidence_threshold);
169    if let Some(notes) = config.handoff_notes {
170        classifier = classifier.with_handoff_notes(notes);
171    }
172    let signals = ToolSignalProcessor {
173        recent_window: config.recent_window.unwrap_or(DEFAULT_RECENT_WINDOW),
174    };
175    let target_set = vec![capable.clone(), efficient.clone()];
176    let mut router = FallThrough::<State>::new_with_state(target_set)
177        .with_name(STAGE_ROUTER)
178        .with_processor(Arc::new(signals))
179        .with_classifier(Arc::new(classifier));
180    if let Some(fallback) = config.llm_fallback {
181        // The capability judge takes its tiers in the same order the capability
182        // route passes them: efficient first, capable second.
183        router = router.with_classifier(Arc::new(SourceStamp {
184            inner: Arc::new(LlmTaskClassifier::new(LlmClassifierConfig::Capability {
185                judge_target: fallback.judge_target,
186                efficient_target: efficient,
187                capable_target: capable,
188                config: fallback.config,
189            })?),
190            source: DecisionSource::LlmClassifier,
191        }));
192    }
193    // Nothing behind this, so the turn lands on the picker's default tier —
194    // including when the judge could not tell.
195    router = router.with_classifier(Arc::new(SourceStamp {
196        inner: Arc::new(DefaultTarget::new(fall_open)),
197        source: DecisionSource::FallOpen,
198    }));
199    // Runs on the post-decision hook, so it applies to the target the cascade
200    // settled on, whichever classifier picked it. With no prompts configured it
201    // is a no-op, so there is nothing to branch on.
202    router = router.with_processor(Arc::new(SystemPromptProcessor::new(config.tier_prompts)));
203    Ok(router)
204}
205
206#[cfg(test)]
207mod tests {
208    use std::sync::Arc;
209
210    use async_trait::async_trait;
211    use parking_lot::Mutex;
212    use serde_json::json;
213    use switchyard_protocol::{
214        ContentBlock, LlmRequest, Message, Role, ToolCall, ToolResult, WireFormat,
215    };
216
217    use super::*;
218    use crate::algorithms::util::stage::DECISION_SOURCE_KEY;
219    use crate::core::classifier::Score;
220    use crate::core::state::StateValue;
221    use crate::core::testing::{Serve, reply, test_drive};
222    use switchyard_protocol::{Decision, Metadata, Response};
223
224    /// A classifier that always picks `target`, standing in for a cascade member.
225    struct Fixed(&'static str);
226
227    #[async_trait]
228    impl Classifier<State> for Fixed {
229        async fn score(
230            &self,
231            _state: &mut State,
232            _request: &mut Request,
233            _driver: Option<&Driver>,
234        ) -> Result<(Classification, Option<Response>)> {
235            Ok((
236                Classification::Scores(vec![Score {
237                    target: ModelId::from(self.0),
238                    confidence: 1.0,
239                }]),
240                None,
241            ))
242        }
243    }
244
245    /// A classifier that never decides.
246    struct Abstains;
247
248    #[async_trait]
249    impl Classifier<State> for Abstains {
250        async fn score(
251            &self,
252            _state: &mut State,
253            _request: &mut Request,
254            _driver: Option<&Driver>,
255        ) -> Result<(Classification, Option<Response>)> {
256            Ok((Classification::Ambiguous(vec![]), None))
257        }
258    }
259
260    async fn stamped(inner: Arc<dyn Classifier<State>>) -> Result<Option<String>> {
261        let stamp = SourceStamp {
262            inner,
263            source: DecisionSource::LlmClassifier,
264        };
265        let mut state = State::default();
266        stamp
267            .score(&mut state, &mut Request::default(), None)
268            .await?;
269        Ok(match state.extra.get(DECISION_SOURCE_KEY) {
270            Some(StateValue::String(source)) => Some(source.clone()),
271            _ => None,
272        })
273    }
274
275    #[tokio::test]
276    async fn a_deciding_classifier_is_credited_with_the_turn() -> Result<()> {
277        assert_eq!(
278            stamped(Arc::new(Fixed("strong"))).await?.as_deref(),
279            Some("llm-classifier")
280        );
281        Ok(())
282    }
283
284    #[tokio::test]
285    async fn an_abstaining_classifier_claims_nothing() -> Result<()> {
286        // It passed the turn on, so the next classifier is the one that decided.
287        assert_eq!(stamped(Arc::new(Abstains)).await?, None);
288        Ok(())
289    }
290
291    fn config() -> StageRouterConfig {
292        StageRouterConfig::new(PickerMode::EfficientFirst, 0.5)
293    }
294
295    #[test]
296    fn rejects_an_out_of_range_confidence_threshold() {
297        let mut config = config();
298        config.confidence_threshold = 1.5;
299        assert!(matches!(
300            StageRouter::new(ModelId::from("strong"), ModelId::from("weak"), config),
301            Err(LibsyError::AlgorithmError { .. })
302        ));
303    }
304
305    #[test]
306    fn rejects_an_out_of_range_judge_threshold() {
307        let mut config = config();
308        config.llm_fallback = Some(LlmFallback {
309            judge_target: ModelId::from("judge"),
310            config: TaskClassifierConfig {
311                base_threshold: -0.1,
312                ..Default::default()
313            },
314        });
315        assert!(matches!(
316            StageRouter::new(ModelId::from("strong"), ModelId::from("weak"), config),
317            Err(LibsyError::AlgorithmError { .. })
318        ));
319    }
320
321    #[test]
322    fn builds_over_both_tiers() -> Result<()> {
323        let router = StageRouter::new(ModelId::from("strong"), ModelId::from("weak"), config())?;
324        assert_eq!(router.name(), STAGE_ROUTER);
325        Ok(())
326    }
327
328    // ── routing integration tests ────────────────────────────────────────────
329
330    const ESCALATION: &str = "the previous model was stalling; pick up the diagnosis";
331    const JUDGE: &str = "judge";
332
333    #[derive(Clone, Debug)]
334    struct Call {
335        target: String,
336        messages: Vec<String>,
337        is_answer_call: bool,
338    }
339
340    /// Records what each target receives.
341    #[derive(Default)]
342    struct Recorder {
343        calls: Mutex<Vec<Call>>,
344        judge_p_solve: Mutex<f64>,
345    }
346
347    impl Recorder {
348        fn routed(&self) -> Vec<Call> {
349            self.calls
350                .lock()
351                .iter()
352                .filter(|call| call.is_answer_call)
353                .cloned()
354                .collect()
355        }
356
357        /// Serves every call, recording it. The judge target gets a structured verdict
358        /// back so the fallback classifier has an answer without a real model.
359        fn serve(self: &Arc<Self>) -> impl Serve {
360            let recorder = Arc::clone(self);
361            move |decision: Decision, request: Request| {
362                let recorder = Arc::clone(&recorder);
363                async move {
364                    let target = decision.selected_model_id().to_string();
365                    recorder.calls.lock().push(Call {
366                        target: target.clone(),
367                        messages: request
368                            .llm_request
369                            .messages
370                            .iter()
371                            .filter_map(|message| message.text_content("|"))
372                            .collect(),
373                        is_answer_call: decision.is_answer_call(),
374                    });
375                    let completion = if target == JUDGE {
376                        let p_solve = *recorder.judge_p_solve.lock();
377                        format!(
378                            r#"{{"crux":"bounded task","primary_rule":"SUP-1","capability_boundary":"supported","p_solve":{p_solve}}}"#
379                        )
380                    } else {
381                        target
382                    };
383                    Ok(reply(completion))
384                }
385            }
386        }
387    }
388
389    fn recording_router(config: StageRouterConfig) -> Result<Arc<StageRouter>> {
390        Ok(Arc::new(StageRouter::new(
391            ModelId::from("strong"),
392            ModelId::from("weak"),
393            config,
394        )?))
395    }
396
397    fn config_with_notes() -> StageRouterConfig {
398        let mut c = config();
399        c.handoff_notes = Some(HandoffNoteConfig::new(ESCALATION, None, true));
400        c
401    }
402
403    fn config_with_judge(recorder: &Arc<Recorder>, p_solve: f64) -> StageRouterConfig {
404        *recorder.judge_p_solve.lock() = p_solve;
405        let mut c = config();
406        c.llm_fallback = Some(LlmFallback {
407            judge_target: ModelId::from(JUDGE),
408            config: TaskClassifierConfig {
409                base_threshold: 0.5,
410                recent_turn_window: Some(3),
411                ..Default::default()
412            },
413        });
414        c
415    }
416
417    fn turn_request(failed: bool) -> Request {
418        let content = if failed {
419            "fatal runtime error: out of memory"
420        } else {
421            "ok"
422        };
423        Request {
424            llm_request: LlmRequest {
425                model: Some("auto".to_string()),
426                messages: vec![
427                    Message::text(Role::User, "fix the build"),
428                    Message {
429                        role: Role::Assistant,
430                        content: vec![ContentBlock::ToolCall(ToolCall {
431                            id: "call_1".to_string(),
432                            name: "Bash".to_string(),
433                            arguments: json!({"command": "cargo test"}),
434                        })],
435                    },
436                    Message {
437                        role: Role::Tool,
438                        content: vec![ContentBlock::ToolResult(ToolResult {
439                            tool_call_id: "call_1".to_string(),
440                            content: vec![ContentBlock::Text {
441                                text: content.to_string(),
442                            }],
443                            is_error: Some(failed),
444                        })],
445                    },
446                ],
447                ..LlmRequest::default()
448            },
449            raw_request: Some(json!({
450                "model": "auto",
451                "messages": [
452                    {"role": "user", "content": "fix the build"},
453                    {"role": "assistant", "tool_calls": [{"id": "call_1", "type": "function",
454                        "function": {"name": "Bash", "arguments": "{\"command\": \"cargo test\"}"}}]},
455                    {"role": "tool", "tool_call_id": "call_1", "content": content},
456                ],
457            })),
458            metadata: Some(Metadata {
459                wire_format: Some(WireFormat::OpenAiChat),
460                session_id: Some("session-1".to_string()),
461                ..Default::default()
462            }),
463        }
464    }
465
466    #[tokio::test]
467    async fn a_signal_driven_escalation_hands_the_note_to_the_model() -> Result<()> {
468        let recorder = Arc::new(Recorder::default());
469        let router = recording_router(config_with_notes())?;
470
471        test_drive(router.clone(), turn_request(false), recorder.serve()).await?;
472        test_drive(router.clone(), turn_request(true), recorder.serve()).await?;
473
474        let calls = recorder.routed();
475        assert_eq!(calls[0].target, "weak");
476        assert_eq!(calls[1].target, "strong");
477        assert!(
478            !calls[0].messages.iter().any(|t| t.contains(ESCALATION)),
479            "steady-state turn should carry no note: {:?}",
480            calls[0].messages
481        );
482        assert!(
483            calls[1]
484                .messages
485                .last()
486                .is_some_and(|t| t.ends_with(ESCALATION)),
487            "escalating turn should carry the note last: {:?}",
488            calls[1].messages
489        );
490        Ok(())
491    }
492
493    #[tokio::test]
494    async fn the_judge_decides_a_turn_the_signals_leave_undecided() -> Result<()> {
495        let recorder = Arc::new(Recorder::default());
496        let router = recording_router(config_with_judge(&recorder, 0.1))?;
497
498        let (trace, _) = test_drive(router.clone(), turn_request(false), recorder.serve()).await?;
499
500        let calls = recorder.calls.lock();
501        assert!(
502            calls
503                .iter()
504                .any(|call| call.target == JUDGE && !call.is_answer_call),
505            "the judge should be recorded as a routing side call"
506        );
507        assert!(
508            calls
509                .iter()
510                .any(|call| call.target == "strong" && call.is_answer_call),
511            "the selected target should be recorded as an answer call"
512        );
513        drop(calls);
514        assert!(
515            trace
516                .last()
517                .and_then(|decision| decision.reasoning())
518                .is_some_and(|reasoning| reasoning.contains("routing tier: strong"))
519        );
520        Ok(())
521    }
522
523    #[tokio::test]
524    async fn a_decisive_signal_never_reaches_the_judge() -> Result<()> {
525        let recorder = Arc::new(Recorder::default());
526        let router = recording_router(config_with_judge(&recorder, 0.9))?;
527
528        test_drive(router.clone(), turn_request(true), recorder.serve()).await?;
529
530        assert!(
531            !recorder.calls.lock().iter().any(|c| c.target == JUDGE),
532            "a resolved turn should not pay for a judge call"
533        );
534        assert_eq!(recorder.routed()[0].target, "strong");
535        Ok(())
536    }
537
538    #[tokio::test]
539    async fn the_judges_verdict_is_not_pinned_to_the_session() -> Result<()> {
540        let recorder = Arc::new(Recorder::default());
541        let router = recording_router(config_with_judge(&recorder, 0.1))?;
542
543        test_drive(router.clone(), turn_request(false), recorder.serve()).await?;
544        *recorder.judge_p_solve.lock() = 0.9;
545        test_drive(router.clone(), turn_request(false), recorder.serve()).await?;
546
547        let routed = recorder.routed();
548        assert_eq!(routed[0].target, "strong");
549        assert_eq!(routed[1].target, "weak");
550        assert_eq!(
551            recorder
552                .calls
553                .lock()
554                .iter()
555                .filter(|c| c.target == JUDGE)
556                .count(),
557            2,
558            "each undecided turn is its own question"
559        );
560        Ok(())
561    }
562
563    #[tokio::test]
564    async fn a_judge_that_cannot_tell_lands_on_the_picker_default() -> Result<()> {
565        let recorder = Arc::new(Recorder::default());
566        let router = recording_router(config_with_judge(&recorder, 42.0))?;
567
568        test_drive(router.clone(), turn_request(false), recorder.serve()).await?;
569
570        assert_eq!(recorder.routed()[0].target, "weak");
571        Ok(())
572    }
573
574    #[tokio::test]
575    async fn the_judge_reads_the_window_it_was_configured_with() -> Result<()> {
576        let recorder = Arc::new(Recorder::default());
577        let router = recording_router(config_with_judge(&recorder, 0.9))?;
578
579        test_drive(router.clone(), turn_request(false), recorder.serve()).await?;
580
581        let judged = recorder
582            .calls
583            .lock()
584            .iter()
585            .find(|c| c.target == JUDGE)
586            .map(|c| c.messages.join("|"));
587        let Some(judged) = judged else {
588            panic!("the judge was never called");
589        };
590        assert!(
591            judged.contains("fix the build"),
592            "the judge should see the opening task: {judged}"
593        );
594        Ok(())
595    }
596}
switchyard_libsy/algorithms/stage.rs

switchyard_libsy/algorithms/
stage.rs
