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

switchyard_libsy/algorithms/
stage.rs
