aimux_core/
router.rs

1//! `RouterModel` — a composite model that routes each call to one child and
2//! optionally falls back (RFC-0021).
3//!
4//! `RouterModel` implements [`LanguageModel`], so `generate_text` /
5//! `stream_text` and every binding use it unchanged. Routing decisions live in
6//! the pluggable [`Router`] trait (pure decision: prompt + child list → index);
7//! execution + fallback live here. Built-in strategies: [`RuleRouter`],
8//! [`WeightedRouter`].
9
10use async_trait::async_trait;
11
12use crate::composite::ChildModel;
13use crate::error::AiMuxError;
14use crate::language_model::LanguageModel;
15use crate::language_model_message::LanguageModelPrompt;
16use crate::options::CallOptions;
17use crate::result::{GenerateResult, StreamResult};
18use crate::retry;
19
20// ─────────────────────────────────────────────────────────────────────────────
21// Router trait
22// ─────────────────────────────────────────────────────────────────────────────
23
24/// Routing strategy: a pure decision over the prompt + child list. It does NOT
25/// execute the call — [`RouterModel`] owns execution and fallback.
26///
27/// Built-in implementations: [`RuleRouter`], [`WeightedRouter`]. Users can
28/// implement this to inject learned classifiers (e.g. RouteLLM via `ort`) —
29/// learned routing is intentionally out of core (see RFC-0021 §6.1).
30pub trait Router: Send + Sync {
31    /// Choose a child-model index. `Err` means "no child can serve this prompt".
32    ///
33    /// # Errors
34    ///
35    /// `Err` means no child model can serve this prompt (all children filtered
36    /// out).
37    fn route(
38        &self,
39        prompt: &LanguageModelPrompt,
40        models: &[ChildModel],
41    ) -> Result<usize, AiMuxError>;
42}
43
44/// Static-priority router: always pick child 0 (the primary); fallback walks the
45/// rest in array order. Equivalent to "primary + backups".
46pub struct RuleRouter;
47
48impl Router for RuleRouter {
49    fn route(
50        &self,
51        _prompt: &LanguageModelPrompt,
52        models: &[ChildModel],
53    ) -> Result<usize, AiMuxError> {
54        if models.is_empty() {
55            return Err(AiMuxError::Other("router: no models configured".into()));
56        }
57        Ok(0)
58    }
59}
60
61/// Weighted router: pick the child with the highest weight. On ties the
62/// **earliest** index wins (so all-equal weights behave like `RuleRouter` —
63/// always child 0). Missing trailing weights default to `0.0`. NaN at index > 0
64/// loses to any finite weight (NaN at index 0 wins only because nothing
65/// compares greater than NaN — avoid NaN weights). To route by cost
66/// (lowest-cost first), pass reciprocals or negative weights.
67pub struct WeightedRouter {
68    weights: Vec<f64>,
69}
70
71impl WeightedRouter {
72    /// Weights are positional — `weights[i]` applies to `models[i]`. Missing
73    /// trailing weights default to `0.0`.
74    #[must_use]
75    pub fn new(weights: Vec<f64>) -> Self {
76        Self { weights }
77    }
78}
79
80impl Router for WeightedRouter {
81    fn route(
82        &self,
83        _prompt: &LanguageModelPrompt,
84        models: &[ChildModel],
85    ) -> Result<usize, AiMuxError> {
86        if models.is_empty() {
87            return Err(AiMuxError::Other("router: no models configured".into()));
88        }
89        // Pick the child with the highest weight. On ties, prefer the earliest
90        // index (matches `RuleRouter`'s "child 0 first" expectation when all
91        // weights are equal). NaN weights lose to any finite weight.
92        let mut best_idx = 0;
93        let mut best_weight = *self.weights.first().unwrap_or(&0.0);
94        for i in 1..models.len() {
95            let w = *self.weights.get(i).unwrap_or(&0.0);
96            if w > best_weight {
97                best_idx = i;
98                best_weight = w;
99            }
100        }
101        Ok(best_idx)
102    }
103}
104
105// ─────────────────────────────────────────────────────────────────────────────
106// RouterModel
107// ─────────────────────────────────────────────────────────────────────────────
108
109/// When a routed call fails, should we try the next child?
110#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
111pub enum FallbackPolicy {
112    /// On error, try the remaining children in array order (default).
113    #[default]
114    OnError,
115    /// The chosen child's failure is final.
116    None,
117}
118
119#[derive(Debug, Clone)]
120pub struct RouterConfig {
121    /// `provider()` value (default `"router"`).
122    pub provider_name: String,
123    /// `model_id()` value (default `"router"`).
124    pub model_id: String,
125}
126
127impl Default for RouterConfig {
128    fn default() -> Self {
129        Self {
130            provider_name: "router".into(),
131            model_id: "router".into(),
132        }
133    }
134}
135
136/// A composite model that routes each call to one child and (optionally) falls
137/// back across the rest on error.
138///
139/// Streaming retries setup on the routed child, but does not fall back or retry
140/// after setup: once stream parts are visible to the user, either would risk
141/// duplicating tokens. See RFC-0021 §3.3.
142pub struct RouterModel {
143    models: Vec<ChildModel>,
144    router: Box<dyn Router>,
145    fallback: FallbackPolicy,
146    config: RouterConfig,
147}
148
149impl RouterModel {
150    #[must_use]
151    pub fn new(
152        models: Vec<ChildModel>,
153        router: Box<dyn Router>,
154        fallback: FallbackPolicy,
155        config: RouterConfig,
156    ) -> Self {
157        Self {
158            models,
159            router,
160            fallback,
161            config,
162        }
163    }
164
165    /// Try every child except `exclude` in array order; return the first `Ok`,
166    /// else the last error. `primary_err` seeds the error returned when there
167    /// are no fallback candidates (e.g. a single-child router) so the real
168    /// failure is never lost to a generic "all models failed".
169    async fn fallback_generate(
170        &self,
171        exclude: usize,
172        options: &CallOptions,
173        primary_err: AiMuxError,
174    ) -> Result<GenerateResult, AiMuxError> {
175        let mut last_err = Some(primary_err);
176        for (i, m) in self.models.iter().enumerate() {
177            if i == exclude {
178                continue;
179            }
180            let retries = retry::prepare_retries(
181                options.max_retries,
182                m.retry_config(),
183                options.abort_signal.clone(),
184            );
185            let child_options = options.for_step(self.step_label(i));
186            match retries
187                .retry(|| {
188                    if let Some(context) = &child_options.recording_context {
189                        let _ = context.start_attempt();
190                    }
191                    m.do_generate(&child_options)
192                })
193                .await
194            {
195                Ok(r) => return Ok(r),
196                Err(e) => last_err = Some(e),
197            }
198        }
199        Err(last_err.expect("seeded primary_err makes last_err always Some"))
200    }
201
202    fn step_label(&self, idx: usize) -> String {
203        let m = &self.models[idx];
204        format!("router[{idx}]:{}/{}", m.provider(), m.model_id())
205    }
206
207    /// Validate a `Router`-returned index before indexing `self.models`. A
208    /// buggy/hostile `Router` (user-implementable trait) must not panic the
209    /// process — surface it as `InvalidArgument` instead.
210    fn check_index(&self, idx: usize, from: &str) -> Result<usize, AiMuxError> {
211        if idx < self.models.len() {
212            Ok(idx)
213        } else {
214            Err(AiMuxError::InvalidArgument(format!(
215                "{from}: router returned out-of-bounds index {idx} (models: {})",
216                self.models.len()
217            )))
218        }
219    }
220}
221
222#[async_trait]
223impl LanguageModel for RouterModel {
224    fn provider(&self) -> &str {
225        &self.config.provider_name
226    }
227
228    fn model_id(&self) -> &str {
229        &self.config.model_id
230    }
231
232    fn retry_config(&self) -> retry::RetryConfig {
233        // Retrying the composite would rerun routing and previously attempted
234        // children; each child is retried at its own execution boundary.
235        retry::RetryConfig {
236            max_retries: 0,
237            ..retry::RetryConfig::default()
238        }
239    }
240
241    async fn do_generate(&self, options: &CallOptions) -> Result<GenerateResult, AiMuxError> {
242        let raw = self.router.route(&options.prompt, &self.models)?;
243        let idx = self.check_index(raw, "router")?;
244        let model = &self.models[idx];
245        let retries = retry::prepare_retries(
246            options.max_retries,
247            model.retry_config(),
248            options.abort_signal.clone(),
249        );
250        let child_options = options.for_step(self.step_label(idx));
251        match retries
252            .retry(|| {
253                if let Some(context) = &child_options.recording_context {
254                    let _ = context.start_attempt();
255                }
256                model.do_generate(&child_options)
257            })
258            .await
259        {
260            Ok(r) => Ok(r),
261            Err(e) => {
262                if self.fallback == FallbackPolicy::OnError {
263                    self.fallback_generate(idx, options, e).await
264                } else {
265                    Err(e)
266                }
267            }
268        }
269    }
270
271    async fn do_stream(&self, options: &CallOptions) -> Result<StreamResult, AiMuxError> {
272        let raw = self.router.route(&options.prompt, &self.models)?;
273        let idx = self.check_index(raw, "router")?;
274        let model = &self.models[idx];
275        let retries = retry::prepare_retries(
276            options.max_retries,
277            model.retry_config(),
278            options.abort_signal.clone(),
279        );
280        let child_options = options.for_step(self.step_label(idx));
281        // Only setup can be retried. The returned stream is passed through, so
282        // failures after setup never invoke another child or duplicate output.
283        retries
284            .retry(|| {
285                if let Some(context) = &child_options.recording_context {
286                    let _ = context.start_attempt();
287                }
288                model.do_stream(&child_options)
289            })
290            .await
291    }
292}
293
294#[cfg(test)]
295mod tests {
296    use super::*;
297    use crate::recording::{
298        HttpExchange, OutcomeRecord, ProviderRecord, Recorder, RecordingContext, ResponseRecord,
299    };
300    use crate::result::GenerateResult;
301    use crate::stream_part::StreamPart;
302    use crate::types::{FinishReason, FinishReasonUnified, Usage};
303    use async_trait::async_trait;
304    use futures::stream;
305    use std::sync::Arc;
306    use std::sync::atomic::{AtomicUsize, Ordering};
307
308    /// A mock child that either succeeds with `text` or always fails.
309    struct MockChild {
310        name: &'static str,
311        text: String,
312        fail: bool,
313        retry_failures: usize,
314        calls: Arc<AtomicUsize>,
315    }
316
317    #[async_trait]
318    impl LanguageModel for MockChild {
319        fn provider(&self) -> &str {
320            "mock"
321        }
322        fn model_id(&self) -> &str {
323            self.name
324        }
325        async fn do_generate(&self, _options: &CallOptions) -> Result<GenerateResult, AiMuxError> {
326            let attempt = self.calls.fetch_add(1, Ordering::SeqCst) + 1;
327            if self.fail {
328                return Err(AiMuxError::Other(format!("{} always fails", self.name)));
329            }
330            if attempt <= self.retry_failures {
331                return Err(retryable_error(self.name));
332            }
333            Ok(GenerateResult {
334                content: vec![crate::result::GenerateContent::Text {
335                    text: self.text.clone(),
336                    provider_metadata: None,
337                }],
338                finish_reason: FinishReason {
339                    unified: FinishReasonUnified::Stop,
340                    raw: None,
341                },
342                usage: Usage::default(),
343                warnings: vec![],
344                provider_metadata: None,
345                response: Default::default(),
346                request_body: None,
347                response_headers: None,
348            })
349        }
350        async fn do_stream(&self, _options: &CallOptions) -> Result<StreamResult, AiMuxError> {
351            let attempt = self.calls.fetch_add(1, Ordering::SeqCst) + 1;
352            if self.fail {
353                return Err(AiMuxError::Other(format!("{} always fails", self.name)));
354            }
355            if attempt <= self.retry_failures {
356                return Err(retryable_error(self.name));
357            }
358            let parts: Vec<Result<StreamPart, AiMuxError>> = vec![
359                Ok(StreamPart::StreamStart { warnings: vec![] }),
360                Ok(StreamPart::TextDelta {
361                    id: "t1".into(),
362                    delta: self.text.clone(),
363                    provider_metadata: None,
364                }),
365                Ok(StreamPart::Finish {
366                    finish_reason: FinishReason {
367                        unified: FinishReasonUnified::Stop,
368                        raw: None,
369                    },
370                    usage: Usage::default(),
371                    provider_metadata: None,
372                }),
373            ];
374            Ok(StreamResult {
375                stream: Box::pin(stream::iter(parts)),
376                request_body: None,
377                response_headers: None,
378            })
379        }
380    }
381
382    fn retryable_error(name: &str) -> AiMuxError {
383        AiMuxError::ApiCall(Box::new(crate::ApiCallError {
384            status_code: Some(503),
385            response_headers: Some(std::collections::HashMap::from([(
386                "retry-after-ms".into(),
387                "0".into(),
388            )])),
389            is_retryable: true,
390            ..crate::ApiCallError::new(
391                format!("{name} retryable failure"),
392                "https://example.test",
393                serde_json::json!({}),
394            )
395        }))
396    }
397
398    fn retry_child(
399        name: &'static str,
400        text: &str,
401        retry_failures: usize,
402    ) -> (ChildModel, Arc<AtomicUsize>) {
403        let calls = Arc::new(AtomicUsize::new(0));
404        (
405            Arc::new(MockChild {
406                name,
407                text: text.into(),
408                fail: false,
409                retry_failures,
410                calls: calls.clone(),
411            }),
412            calls,
413        )
414    }
415
416    struct CountingRuleRouter(Arc<AtomicUsize>);
417
418    impl Router for CountingRuleRouter {
419        fn route(
420            &self,
421            prompt: &LanguageModelPrompt,
422            models: &[ChildModel],
423        ) -> Result<usize, AiMuxError> {
424            self.0.fetch_add(1, Ordering::SeqCst);
425            RuleRouter.route(prompt, models)
426        }
427    }
428
429    fn child(name: &'static str, text: &str, fail: bool) -> ChildModel {
430        Arc::new(MockChild {
431            name,
432            text: text.into(),
433            fail,
434            retry_failures: 0,
435            calls: Arc::new(AtomicUsize::new(0)),
436        })
437    }
438
439    fn opts() -> CallOptions {
440        CallOptions {
441            prompt: vec![],
442            ..Default::default()
443        }
444    }
445
446    #[tokio::test]
447    async fn routes_to_first_with_rule_router() {
448        let router = RouterModel::new(
449            vec![child("a", "from-a", false), child("b", "from-b", false)],
450            Box::new(RuleRouter),
451            FallbackPolicy::OnError,
452            RouterConfig::default(),
453        );
454        let r = router.do_generate(&opts()).await.unwrap();
455        let text = crate::composite::extract_text(&r.content);
456        assert_eq!(text, "from-a");
457    }
458
459    #[tokio::test]
460    async fn weighted_router_picks_max() {
461        let router = RouterModel::new(
462            vec![child("a", "from-a", false), child("b", "from-b", false)],
463            Box::new(WeightedRouter::new(vec![1.0, 5.0])),
464            FallbackPolicy::None,
465            RouterConfig::default(),
466        );
467        let r = router.do_generate(&opts()).await.unwrap();
468        assert_eq!(crate::composite::extract_text(&r.content), "from-b");
469    }
470
471    #[tokio::test]
472    async fn fallback_on_error_walks_remaining() {
473        // Primary fails, second fails, third succeeds.
474        let router = RouterModel::new(
475            vec![
476                child("a", "from-a", true),
477                child("b", "from-b", true),
478                child("c", "from-c", false),
479            ],
480            Box::new(RuleRouter),
481            FallbackPolicy::OnError,
482            RouterConfig::default(),
483        );
484        let r = router.do_generate(&opts()).await.unwrap();
485        assert_eq!(crate::composite::extract_text(&r.content), "from-c");
486    }
487
488    #[tokio::test]
489    async fn fallback_none_propagates_error() {
490        let router = RouterModel::new(
491            vec![child("a", "from-a", true), child("b", "from-b", false)],
492            Box::new(RuleRouter),
493            FallbackPolicy::None,
494            RouterConfig::default(),
495        );
496        let err = router.do_generate(&opts()).await.unwrap_err();
497        assert!(err.to_string().contains("always fails"));
498    }
499
500    #[tokio::test]
501    async fn all_fail_returns_last_error() {
502        let router = RouterModel::new(
503            vec![child("a", "from-a", true), child("b", "from-b", true)],
504            Box::new(RuleRouter),
505            FallbackPolicy::OnError,
506            RouterConfig::default(),
507        );
508        let err = router.do_generate(&opts()).await.unwrap_err();
509        assert!(err.to_string().contains("always fails"));
510    }
511
512    #[tokio::test]
513    async fn stream_routes_without_fallback() {
514        // Primary is selected by RuleRouter; even though it would fail, stream
515        // does not fall back. Use a healthy primary to confirm the happy path.
516        let router = RouterModel::new(
517            vec![child("a", "from-a", false), child("b", "from-b", false)],
518            Box::new(RuleRouter),
519            FallbackPolicy::OnError,
520            RouterConfig::default(),
521        );
522        let result = router.do_stream(&opts()).await.unwrap();
523        use futures::StreamExt;
524        let mut s = result.stream;
525        let mut text = String::new();
526        while let Some(part) = s.next().await {
527            if let Ok(StreamPart::TextDelta { delta, .. }) = part {
528                text.push_str(&delta);
529            }
530        }
531        assert_eq!(text, "from-a");
532    }
533
534    #[tokio::test]
535    async fn single_child_failure_preserves_real_error() {
536        // B1 regression guard: a single failing child must surface its real
537        // error, not a generic "all models failed".
538        let router = RouterModel::new(
539            vec![child("solo", "solo-out", true)],
540            Box::new(RuleRouter),
541            FallbackPolicy::OnError,
542            RouterConfig::default(),
543        );
544        let err = router.do_generate(&opts()).await.unwrap_err();
545        assert!(
546            err.to_string().contains("solo always fails"),
547            "expected the real child error, got: {err}"
548        );
549    }
550
551    #[tokio::test]
552    async fn single_child_success_works() {
553        let router = RouterModel::new(
554            vec![child("solo", "solo-out", false)],
555            Box::new(RuleRouter),
556            FallbackPolicy::OnError,
557            RouterConfig::default(),
558        );
559        let r = router.do_generate(&opts()).await.unwrap();
560        assert_eq!(crate::composite::extract_text(&r.content), "solo-out");
561    }
562
563    #[tokio::test]
564    async fn out_of_bounds_router_index_is_invalid_argument() {
565        // A buggy/hostile Router returning an OOB index must not panic; it
566        // surfaces as InvalidArgument.
567        struct OobRouter;
568        impl Router for OobRouter {
569            fn route(
570                &self,
571                _prompt: &LanguageModelPrompt,
572                _models: &[ChildModel],
573            ) -> Result<usize, AiMuxError> {
574                Ok(99)
575            }
576        }
577        let router = RouterModel::new(
578            vec![child("a", "from-a", false)],
579            Box::new(OobRouter),
580            FallbackPolicy::OnError,
581            RouterConfig::default(),
582        );
583        let err = router.do_generate(&opts()).await.unwrap_err();
584        match err {
585            AiMuxError::InvalidArgument(msg) => {
586                assert!(msg.contains("out-of-bounds"), "got: {msg}");
587            }
588            other => panic!("expected InvalidArgument, got {other:?}"),
589        }
590    }
591
592    #[tokio::test]
593    async fn weighted_router_ties_pick_earliest() {
594        // S2 regression guard: equal weights resolve to the lowest index.
595        let router = RouterModel::new(
596            vec![child("a", "from-a", false), child("b", "from-b", false)],
597            Box::new(WeightedRouter::new(vec![1.0, 1.0])),
598            FallbackPolicy::None,
599            RouterConfig::default(),
600        );
601        let r = router.do_generate(&opts()).await.unwrap();
602        assert_eq!(crate::composite::extract_text(&r.content), "from-a");
603    }
604
605    #[tokio::test]
606    async fn empty_models_is_error() {
607        let router = RouterModel::new(
608            vec![],
609            Box::new(RuleRouter),
610            FallbackPolicy::OnError,
611            RouterConfig::default(),
612        );
613        let err = router.do_generate(&opts()).await.unwrap_err();
614        assert!(err.to_string().contains("no models"));
615    }
616
617    #[tokio::test]
618    async fn stream_does_not_fallback_when_routed_child_fails() {
619        // G3: the key streaming invariant — even with FallbackPolicy::OnError,
620        // a failing primary is NOT retried via the backup during do_stream.
621        // (StreamStart is already emitted by the time a retry would matter.)
622        let router = RouterModel::new(
623            vec![
624                child("primary", "primary-out", true),
625                child("backup", "backup-out", false),
626            ],
627            Box::new(RuleRouter),
628            FallbackPolicy::OnError, // would fall back in do_generate, must NOT here
629            RouterConfig::default(),
630        );
631        let err = router.do_stream(&opts()).await.unwrap_err();
632        assert!(
633            err.to_string().contains("primary always fails"),
634            "do_stream must surface the primary error, not fall back; got: {err}"
635        );
636    }
637
638    #[tokio::test]
639    async fn usable_as_dyn_language_model_via_generate_text() {
640        // G7: the RFCs' headline claim — composite models drop into the
641        // generate_text free function unchanged via &dyn LanguageModel.
642        let router = RouterModel::new(
643            vec![child("a", "works through trait object", false)],
644            Box::new(RuleRouter),
645            FallbackPolicy::OnError,
646            RouterConfig::default(),
647        );
648        let r = crate::generate::generate_text(
649            &router,
650            "hello",
651            crate::generate::GenerateTextOptions::default(),
652        )
653        .await
654        .unwrap();
655        assert_eq!(r.text, "works through trait object");
656    }
657
658    #[tokio::test]
659    async fn retries_each_child_before_falling_back_without_rerouting() {
660        let (primary, primary_calls) = retry_child("primary", "primary", usize::MAX);
661        let (backup, backup_calls) = retry_child("backup", "backup", 1);
662        let route_calls = Arc::new(AtomicUsize::new(0));
663        let router = RouterModel::new(
664            vec![primary, backup],
665            Box::new(CountingRuleRouter(route_calls.clone())),
666            FallbackPolicy::OnError,
667            RouterConfig::default(),
668        );
669        assert_eq!(router.retry_config().max_retries, 0);
670
671        let result = crate::generate::generate_text(
672            &router,
673            "hello",
674            crate::generate::GenerateTextOptions {
675                max_retries: Some(2),
676                ..Default::default()
677            },
678        )
679        .await
680        .unwrap();
681
682        assert_eq!(result.text, "backup");
683        assert_eq!(route_calls.load(Ordering::SeqCst), 1);
684        assert_eq!(primary_calls.load(Ordering::SeqCst), 3);
685        assert_eq!(backup_calls.load(Ordering::SeqCst), 2);
686    }
687
688    /// Children must receive a step-labeled child recording context so their
689    /// exchanges and retry attempts are recorded as nested steps.
690    #[tokio::test]
691    async fn children_receive_step_labeled_recording_contexts() {
692        struct StepProbe {
693            seen: std::sync::Mutex<Vec<Option<String>>>,
694        }
695
696        #[async_trait]
697        impl LanguageModel for StepProbe {
698            fn provider(&self) -> &str {
699                "mock"
700            }
701            fn model_id(&self) -> &str {
702                "probe"
703            }
704            async fn do_generate(
705                &self,
706                options: &CallOptions,
707            ) -> Result<GenerateResult, AiMuxError> {
708                self.seen.lock().unwrap().push(
709                    options
710                        .recording_context
711                        .as_ref()
712                        .and_then(RecordingContext::step),
713                );
714                Err(AiMuxError::Other("probe fails".into()))
715            }
716            async fn do_stream(&self, _: &CallOptions) -> Result<StreamResult, AiMuxError> {
717                unreachable!()
718            }
719        }
720
721        struct NoopRecorder;
722        impl Recorder for NoopRecorder {
723            fn record_input(&self, _: &str, _: &CallOptions, _: &str, _: &str) {}
724            fn record_provider(&self, _: &str, _: &ProviderRecord) {}
725            fn record_exchange(&self, _: &str, _: &HttpExchange) {}
726            fn record_exchange_update(
727                &self,
728                _: &str,
729                _: u32,
730                _: u32,
731                _: &ResponseRecord,
732                _: Option<String>,
733            ) {
734            }
735            fn record_outcome(&self, _: &str, _: &OutcomeRecord) {}
736            fn flush(&self) {}
737        }
738
739        let probe = Arc::new(StepProbe {
740            seen: std::sync::Mutex::new(Vec::new()),
741        });
742        let router = RouterModel::new(
743            vec![probe.clone()],
744            Box::new(CountingRuleRouter(Arc::new(AtomicUsize::new(0)))),
745            FallbackPolicy::OnError,
746            RouterConfig::default(),
747        );
748
749        let options = CallOptions {
750            recording_context: Some(RecordingContext::new("call-1", Arc::new(NoopRecorder))),
751            ..CallOptions::default()
752        };
753        let _ = router.do_generate(&options).await;
754
755        let seen = probe.seen.lock().unwrap();
756        assert!(!seen.is_empty());
757        assert_eq!(seen[0].as_deref(), Some("router[0]:mock/probe"));
758    }
759
760    /// Even when every child fails and the caller asked for per-call retries,
761    /// the composite operation must run exactly once: the exhausted child's
762    /// `Retry` error passes through the outer retry loop instead of replaying
763    /// routing and fallback.
764    #[tokio::test]
765    async fn all_children_failing_does_not_replay_the_composite() {
766        let (primary, primary_calls) = retry_child("primary", "primary", usize::MAX);
767        let (backup, backup_calls) = retry_child("backup", "backup", usize::MAX);
768        let route_calls = Arc::new(AtomicUsize::new(0));
769        let router = RouterModel::new(
770            vec![primary, backup],
771            Box::new(CountingRuleRouter(route_calls.clone())),
772            FallbackPolicy::OnError,
773            RouterConfig::default(),
774        );
775
776        let error = crate::generate::generate_text(
777            &router,
778            "hello",
779            crate::generate::GenerateTextOptions {
780                max_retries: Some(2),
781                ..Default::default()
782            },
783        )
784        .await
785        .unwrap_err();
786
787        assert!(matches!(error, crate::AiMuxError::Retry(_)), "{error:?}");
788        assert_eq!(route_calls.load(Ordering::SeqCst), 1);
789        assert_eq!(primary_calls.load(Ordering::SeqCst), 3);
790        assert_eq!(backup_calls.load(Ordering::SeqCst), 3);
791    }
792
793    #[tokio::test]
794    async fn retries_stream_setup_on_selected_child_without_fallback_or_rerouting() {
795        let (primary, primary_stream_calls) = retry_child("primary", "primary", 1);
796        let (backup, backup_stream_calls) = retry_child("backup", "backup", 0);
797        let route_calls = Arc::new(AtomicUsize::new(0));
798        let router = RouterModel::new(
799            vec![primary, backup],
800            Box::new(CountingRuleRouter(route_calls.clone())),
801            FallbackPolicy::OnError,
802            RouterConfig::default(),
803        );
804
805        let _result = crate::generate::stream_text(
806            &router,
807            "hello",
808            crate::generate::GenerateTextOptions {
809                max_retries: Some(2),
810                ..Default::default()
811            },
812        )
813        .await
814        .unwrap();
815
816        assert_eq!(route_calls.load(Ordering::SeqCst), 1);
817        assert_eq!(primary_stream_calls.load(Ordering::SeqCst), 2);
818        assert_eq!(backup_stream_calls.load(Ordering::SeqCst), 0);
819    }
820}
aimux_core/router.rs

aimux_core/
router.rs
