llm_kernel/llm/
router.rs

1//! Cost-aware routing and fallback chain for [`LLMClient`].
2//!
3//! [`RouterClient`] holds an ordered set of [`Backend`]s and routes each
4//! request according to a [`RoutingStrategy`]. On a transient error it falls
5//! through to the next backend; a permanent client error short-circuits. This
6//! is the orchestration layer for cross-provider / cross-model resilience.
7//!
8//! # Separation of concerns
9//!
10//! The `llm` module composes resilience from focused decorators rather than
11//! folding every behavior into one client:
12//!
13//! | Concern | Decorator |
14//! |---------|-----------|
15//! | rate-limit (429) / 5xx backoff | [`RetryClient`](crate::llm::retry::RetryClient) |
16//! | routing / cross-provider fallback | [`RouterClient`] |
17//! | request/response observation | [`MiddlewareClient`](crate::llm::middleware::MiddlewareClient) |
18//! | response caching | [`CacheClient`](crate::llm::cache::CacheClient) |
19//!
20//! Because every decorator implements [`LLMClient`], they compose freely. A
21//! resilient, observed, cached multi-backend stack reads inside-out:
22//!
23//! ```ignore
24//! use llm_kernel::llm::{
25//!     CacheClient, LLMClient, MiddlewareClient, NoopMiddleware, OpenAIClient,
26//!     RetryClient, RouterClient, RoutingStrategy, Backend, RetryConfig,
27//! };
28//!
29//! let cheap = Backend::new(OpenAIClient::from_key("gpt-4o-mini", "sk-...")?, Some((0.15, 0.60)));
30//! let powerful = Backend::new(OpenAIClient::from_key("gpt-4o", "sk-...")?, Some((2.50, 10.00)));
31//!
32//! // Route across backends; retry each transiently; observe; cache.
33//! let stack = CacheClient::new(
34//!     MiddlewareClient::new(
35//!         RouterClient::new(
36//!             vec![cheap, powerful]
37//!                 .into_iter()
38//!                 .map(|b| Backend { client: RetryClient::new(b.client, RetryConfig::default()), ..b })
39//!                 .collect(),
40//!             RoutingStrategy::LowestCost,
41//!         )?,
42//!         NoopMiddleware,
43//!     ),
44//!     store,
45//! );
46//! ```
47//!
48//! Streaming cannot fall through once a stream is established. So
49//! [`RouterClient::stream_complete`] tries each backend in order and returns
50//! the first stream it can establish; an error raised *before* the stream
51//! starts (connection, 403) falls through to the next backend just like
52//! [`LLMClient::complete`]. Once a stream is returned, errors raised *during*
53//! streaming are not retried — wrap individual backends in
54//! [`RetryClient`](crate::llm::retry::RetryClient) for transient resilience.
55
56use std::cmp::Ordering;
57
58use async_trait::async_trait;
59
60use crate::error::{KernelError, Result};
61use crate::llm::client::LLMClient;
62use crate::llm::types::{LLMRequest, LLMResponse, LLMStream};
63
64/// A single backend in a [`RouterClient`] chain.
65///
66/// `cost_per_1m` is the `(input, output)` price per 1,000,000 tokens in USD,
67/// used by [`RoutingStrategy::LowestCost`] to order backends. `priority`
68/// (ascending) is used by [`RoutingStrategy::Fallback`]. Ties in either keep
69/// insertion order (stable sort).
70pub struct Backend<C: LLMClient> {
71    /// The underlying client.
72    pub client: C,
73    /// `(input_per_1m, output_per_1m)` in USD. `None` = unknown → sorts last.
74    pub cost_per_1m: Option<(f64, f64)>,
75    /// Lower = tried first. Defaults to `0`.
76    pub priority: usize,
77    label: Option<String>,
78}
79
80impl<C: LLMClient> Backend<C> {
81    /// Create a backend with no cost metadata and default priority.
82    pub fn new(client: C, cost_per_1m: Option<(f64, f64)>) -> Self {
83        Self {
84            client,
85            cost_per_1m,
86            priority: 0,
87            label: None,
88        }
89    }
90
91    /// Set the fallback priority (lower = tried first).
92    pub fn with_priority(mut self, priority: usize) -> Self {
93        self.priority = priority;
94        self
95    }
96
97    /// Override the label used for this backend (otherwise the underlying
98    /// client's [`LLMClient::model_name`] is used).
99    pub fn with_label(mut self, label: impl Into<String>) -> Self {
100        self.label = Some(label.into());
101        self
102    }
103
104    fn name(&self) -> &str {
105        self.label
106            .as_deref()
107            .unwrap_or_else(|| self.client.model_name())
108    }
109
110    /// Average of input+output per-1m cost. `None` and non-finite (NaN) values
111    /// normalize to `f64::MAX` so they sort last under [`RoutingStrategy::LowestCost`].
112    fn cost_rank(&self) -> f64 {
113        let raw = self
114            .cost_per_1m
115            .map(|(input, output)| (input + output) / 2.0)
116            .unwrap_or(f64::MAX);
117        if raw.is_nan() { f64::MAX } else { raw }
118    }
119}
120
121/// How a [`RouterClient`] orders and falls through its backends.
122#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
123pub enum RoutingStrategy {
124    /// Try backends in ascending `priority` order; on error fall through.
125    /// Ties keep insertion order.
126    #[default]
127    Fallback,
128    /// Try backends in ascending cost order; on error fall through to the next
129    /// cheapest. Backends with unknown (or non-finite) cost sort last.
130    LowestCost,
131}
132
133/// Compute the try-order of `backends` under `strategy`. Stable: ties keep
134/// insertion order.
135fn compute_order<C: LLMClient>(backends: &[Backend<C>], strategy: RoutingStrategy) -> Vec<usize> {
136    let mut idx: Vec<usize> = (0..backends.len()).collect();
137    match strategy {
138        RoutingStrategy::Fallback => idx.sort_by_key(|&i| backends[i].priority),
139        RoutingStrategy::LowestCost => idx.sort_by(|&a, &b| {
140            backends[a]
141                .cost_rank()
142                .partial_cmp(&backends[b].cost_rank())
143                .unwrap_or(Ordering::Equal)
144        }),
145    }
146    idx
147}
148
149/// Whether `err` is a permanent client-side failure that no backend would
150/// recover from. The router short-circuits on these instead of falling through.
151///
152/// `408` (request timeout) and `429` (rate-limit) are excluded: they are
153/// transient and another provider may accept the request immediately.
154/// Non-HTTP errors (network, auth-token fetch, etc.) are treated as transient.
155fn is_permanent_client_error(err: &KernelError) -> bool {
156    matches!(
157        err,
158        KernelError::Http { status, .. }
159            if (400..500).contains(status) && *status != 408 && *status != 429
160    )
161}
162
163/// A [`LLMClient`] that routes requests across multiple backends with fallback.
164///
165/// Construct with [`RouterClient::new`]. Each `complete` tries the backends in
166/// the order dictated by the [`RoutingStrategy`]; the first success is returned
167/// and any remaining backends are skipped. If every backend errors, the last
168/// error is returned.
169///
170/// Fall-through is **error-class aware**: transient errors (5xx, rate-limit
171/// `429`, timeout `408`, network) move on to the next backend, but permanent
172/// client errors (4xx except `408`/`429` — a misconfigured request, invalid
173/// model, auth failure) short-circuit immediately, since every backend would
174/// reject the same request. Wrap each backend in a
175/// [`RetryClient`](crate::llm::retry::RetryClient) to absorb transient errors
176/// (rate-limit, 5xx) before the router sees them.
177pub struct RouterClient<C: LLMClient> {
178    backends: Vec<Backend<C>>,
179    strategy: RoutingStrategy,
180    /// Precomputed try-order; immutable after [`RouterClient::new`].
181    order: Vec<usize>,
182}
183
184impl<C: LLMClient> RouterClient<C> {
185    /// Create a router over `backends`.
186    ///
187    /// Returns [`KernelError::Config`] if `backends` is empty.
188    pub fn new(backends: Vec<Backend<C>>, strategy: RoutingStrategy) -> Result<Self> {
189        if backends.is_empty() {
190            return Err(KernelError::Config(
191                "RouterClient requires at least one backend".into(),
192            ));
193        }
194        let order = compute_order(&backends, strategy);
195        Ok(Self {
196            backends,
197            strategy,
198            order,
199        })
200    }
201
202    /// The backends in insertion order.
203    pub fn backends(&self) -> &[Backend<C>] {
204        &self.backends
205    }
206
207    /// The active routing strategy.
208    pub fn strategy(&self) -> RoutingStrategy {
209        self.strategy
210    }
211}
212
213#[async_trait]
214impl<C: LLMClient> LLMClient for RouterClient<C> {
215    async fn complete(&self, request: LLMRequest) -> Result<LLMResponse> {
216        let mut last_err: Option<KernelError> = None;
217        for &i in &self.order {
218            match self.backends[i].client.complete(request.clone()).await {
219                Ok(response) => return Ok(response),
220                Err(err) => {
221                    if is_permanent_client_error(&err) {
222                        return Err(err);
223                    }
224                    last_err = Some(err);
225                }
226            }
227        }
228        Err(last_err
229            .unwrap_or_else(|| KernelError::Config("RouterClient: no backends tried".into())))
230    }
231
232    fn model_name(&self) -> &str {
233        // The preferred backend is the first in routing order.
234        let first = self.order.first().copied().unwrap_or(0);
235        self.backends[first].name()
236    }
237
238    async fn stream_complete(&self, request: LLMRequest) -> Result<LLMStream> {
239        // A stream cannot fall through once established, but an error raised
240        // *before* the stream starts (connection, 403) falls through to the
241        // next backend just like `complete`. Permanent 4xx short-circuits.
242        let mut last_err: Option<KernelError> = None;
243        for &i in &self.order {
244            match self.backends[i]
245                .client
246                .stream_complete(request.clone())
247                .await
248            {
249                Ok(stream) => return Ok(stream),
250                Err(err) => {
251                    if is_permanent_client_error(&err) {
252                        return Err(err);
253                    }
254                    last_err = Some(err);
255                }
256            }
257        }
258        Err(last_err
259            .unwrap_or_else(|| KernelError::Config("RouterClient: no stream established".into())))
260    }
261}
262
263#[cfg(test)]
264mod tests {
265    use super::*;
266    use std::sync::{
267        Arc, Mutex,
268        atomic::{AtomicU32, Ordering},
269    };
270
271    use crate::llm::types::StreamEvent;
272    use tokio::sync::mpsc;
273    use tokio_stream::wrappers::ReceiverStream;
274
275    /// A mock client that returns a fixed sequence of `complete` results
276    /// (drained one per call) and records `stream_complete` invocations.
277    struct MockClient {
278        name: &'static str,
279        responses: Mutex<Vec<Result<LLMResponse>>>,
280        stream_calls: Arc<AtomicU32>,
281        /// If set, `stream_complete` returns an HTTP error with this status
282        /// instead of a stream — used to exercise pre-stream fallback.
283        /// (`KernelError` is not `Clone`, so this is a status to reconstruct.)
284        stream_fail_status: Option<u16>,
285    }
286
287    impl MockClient {
288        fn new(name: &'static str, responses: Vec<Result<LLMResponse>>) -> Self {
289            Self {
290                name,
291                responses: Mutex::new(responses),
292                stream_calls: Arc::new(AtomicU32::new(0)),
293                stream_fail_status: None,
294            }
295        }
296
297        fn stream_call_count(&self) -> Arc<AtomicU32> {
298            Arc::clone(&self.stream_calls)
299        }
300
301        /// Make `stream_complete` fail with the given HTTP status every call.
302        fn with_failing_stream(mut self, status: u16) -> Self {
303            self.stream_fail_status = Some(status);
304            self
305        }
306    }
307
308    #[async_trait]
309    impl LLMClient for MockClient {
310        async fn complete(&self, _request: LLMRequest) -> Result<LLMResponse> {
311            let mut responses = self.responses.lock().unwrap();
312            if responses.is_empty() {
313                panic!("MockClient({}): no more responses queued", self.name);
314            }
315            responses.remove(0)
316        }
317
318        fn model_name(&self) -> &str {
319            self.name
320        }
321
322        async fn stream_complete(&self, _request: LLMRequest) -> Result<LLMStream> {
323            self.stream_calls.fetch_add(1, Ordering::Relaxed);
324            if let Some(status) = self.stream_fail_status {
325                return Err(KernelError::Http {
326                    status,
327                    message: "stream unavailable".into(),
328                });
329            }
330            // A single Done event is enough to exercise the delegation path.
331            let (tx, rx) = mpsc::channel::<Result<StreamEvent>>(16);
332            let _ = tx.send(Ok(StreamEvent::Done)).await;
333            Ok(Box::pin(ReceiverStream::new(rx)))
334        }
335    }
336
337    fn ok_from(name: &str) -> Result<LLMResponse> {
338        Ok(LLMResponse {
339            content: format!("from-{name}"),
340            model: name.into(),
341            ..Default::default()
342        })
343    }
344
345    fn server_error() -> Result<LLMResponse> {
346        Err(KernelError::Http {
347            status: 500,
348            message: "boom".into(),
349        })
350    }
351
352    fn client_error(status: u16) -> Result<LLMResponse> {
353        Err(KernelError::Http {
354            status,
355            message: "bad request".into(),
356        })
357    }
358
359    // --- complete() behavior -------------------------------------------------
360
361    #[tokio::test]
362    async fn fallback_falls_through_to_secondary_on_error() {
363        let primary = Backend::new(MockClient::new("primary", vec![server_error()]), None);
364        let secondary = Backend::new(
365            MockClient::new("secondary", vec![ok_from("secondary")]),
366            None,
367        );
368        let router = RouterClient::new(vec![primary, secondary], RoutingStrategy::Fallback)
369            .expect("non-empty");
370
371        let resp = router
372            .complete(LLMRequest::builder().build())
373            .await
374            .expect("secondary succeeds");
375        assert_eq!(resp.content, "from-secondary");
376    }
377
378    #[tokio::test]
379    async fn permanent_4xx_short_circuits_no_fallback() {
380        // A 400 from the primary is permanent — the secondary (no responses
381        // queued) must never be contacted.
382        let primary = Backend::new(MockClient::new("primary", vec![client_error(400)]), None);
383        let secondary = Backend::new(MockClient::new("secondary", vec![]), None);
384        let router = RouterClient::new(vec![primary, secondary], RoutingStrategy::Fallback)
385            .expect("non-empty");
386
387        let err = router
388            .complete(LLMRequest::builder().build())
389            .await
390            .expect_err("400 short-circuits");
391        assert!(matches!(err, KernelError::Http { status: 400, .. }));
392    }
393
394    #[tokio::test]
395    async fn transient_429_falls_through() {
396        // 429 (rate-limit) is transient — fall through to the secondary.
397        let primary = Backend::new(MockClient::new("primary", vec![client_error(429)]), None);
398        let secondary = Backend::new(
399            MockClient::new("secondary", vec![ok_from("secondary")]),
400            None,
401        );
402        let router = RouterClient::new(vec![primary, secondary], RoutingStrategy::Fallback)
403            .expect("non-empty");
404
405        let resp = router
406            .complete(LLMRequest::builder().build())
407            .await
408            .expect("429 falls through");
409        assert_eq!(resp.content, "from-secondary");
410    }
411
412    #[tokio::test]
413    async fn fallback_respects_priority_order() {
414        // Insert "secondary" first, but give it a higher priority value so the
415        // router must try the (later-inserted) primary first.
416        let secondary = Backend::new(MockClient::new("secondary", vec![server_error()]), None)
417            .with_priority(10);
418        let primary = Backend::new(MockClient::new("primary", vec![ok_from("primary")]), None)
419            .with_priority(1);
420
421        let router = RouterClient::new(vec![secondary, primary], RoutingStrategy::Fallback)
422            .expect("non-empty");
423        let resp = router
424            .complete(LLMRequest::builder().build())
425            .await
426            .expect("primary succeeds");
427        assert_eq!(resp.content, "from-primary");
428    }
429
430    #[tokio::test]
431    async fn lowest_cost_orders_by_cost() {
432        // Insert expensive first; the router must still try cheap first.
433        let expensive = Backend::new(
434            MockClient::new("expensive", vec![server_error(), ok_from("expensive")]),
435            Some((10.0, 30.0)),
436        );
437        let cheap = Backend::new(
438            MockClient::new("cheap", vec![ok_from("cheap")]),
439            Some((0.15, 0.60)),
440        );
441        let router = RouterClient::new(vec![expensive, cheap], RoutingStrategy::LowestCost)
442            .expect("non-empty");
443
444        let resp = router
445            .complete(LLMRequest::builder().build())
446            .await
447            .expect("cheap succeeds first");
448        assert_eq!(resp.content, "from-cheap");
449    }
450
451    #[tokio::test]
452    async fn lowest_cost_falls_through_to_next_cheapest() {
453        // Cheapest errors; the next-cheapest must be tried and succeed.
454        let cheap = Backend::new(
455            MockClient::new("cheap", vec![server_error()]),
456            Some((0.15, 0.60)),
457        );
458        let mid = Backend::new(
459            MockClient::new("mid", vec![ok_from("mid")]),
460            Some((1.0, 2.0)),
461        );
462        let router =
463            RouterClient::new(vec![mid, cheap], RoutingStrategy::LowestCost).expect("non-empty");
464
465        let resp = router
466            .complete(LLMRequest::builder().build())
467            .await
468            .expect("mid succeeds after cheap fails");
469        assert_eq!(resp.content, "from-mid");
470    }
471
472    #[tokio::test]
473    async fn unknown_cost_sorts_last() {
474        // The `unknown` backend has no responses queued: if it is tried first,
475        // MockClient panics. The finite-cost `priced` backend must precede it
476        // in LowestCost order, succeed, and short-circuit before `unknown`.
477        let unknown = Backend::new(MockClient::new("unknown", vec![]), None);
478        let priced = Backend::new(
479            MockClient::new("priced", vec![ok_from("priced")]),
480            Some((1.0, 2.0)),
481        );
482        let router = RouterClient::new(vec![unknown, priced], RoutingStrategy::LowestCost)
483            .expect("non-empty");
484        let resp = router
485            .complete(LLMRequest::builder().build())
486            .await
487            .expect("priced tried first");
488        assert_eq!(resp.content, "from-priced");
489    }
490
491    #[tokio::test]
492    async fn nan_cost_is_treated_as_unknown() {
493        // A NaN cost must not corrupt the sort; it must sort last like `None`.
494        let nan_cost = Backend::new(MockClient::new("nan", vec![]), Some((f64::NAN, 1.0)));
495        let priced = Backend::new(
496            MockClient::new("priced", vec![ok_from("priced")]),
497            Some((1.0, 2.0)),
498        );
499        let router = RouterClient::new(vec![nan_cost, priced], RoutingStrategy::LowestCost)
500            .expect("non-empty");
501        let resp = router
502            .complete(LLMRequest::builder().build())
503            .await
504            .expect("priced precedes NaN-cost backend");
505        assert_eq!(resp.content, "from-priced");
506    }
507
508    #[tokio::test]
509    async fn all_backends_fail_returns_last_error() {
510        let primary = Backend::new(MockClient::new("primary", vec![server_error()]), None);
511        let secondary = Backend::new(
512            MockClient::new("secondary", vec![Err(KernelError::LlmApi("down".into()))]),
513            None,
514        );
515        let router = RouterClient::new(vec![primary, secondary], RoutingStrategy::Fallback)
516            .expect("non-empty");
517
518        let err = router
519            .complete(LLMRequest::builder().build())
520            .await
521            .expect_err("both fail");
522        // Last-tried backend was "secondary" → LlmApi, not the primary's Http 500.
523        assert!(matches!(err, KernelError::LlmApi(_)));
524    }
525
526    #[tokio::test]
527    async fn first_success_short_circuits_remaining_backends() {
528        // If the primary succeeds, the secondary must never be called — it has
529        // no responses queued, so a stray call would panic.
530        let primary = Backend::new(MockClient::new("primary", vec![ok_from("primary")]), None);
531        let secondary = Backend::new(MockClient::new("secondary", vec![]), None);
532        let router = RouterClient::new(vec![primary, secondary], RoutingStrategy::Fallback)
533            .expect("non-empty");
534
535        let resp = router
536            .complete(LLMRequest::builder().build())
537            .await
538            .expect("primary succeeds");
539        assert_eq!(resp.content, "from-primary");
540    }
541
542    // --- ordering invariants --------------------------------------------------
543
544    #[tokio::test]
545    async fn empty_backends_is_rejected() {
546        let result = RouterClient::<MockClient>::new(vec![], RoutingStrategy::Fallback);
547        assert!(result.is_err());
548    }
549
550    #[tokio::test]
551    async fn ties_keep_insertion_order() {
552        // Equal priority: insertion order must win. primary (inserted first)
553        // succeeds; secondary (no responses) must not be tried.
554        let primary = Backend::new(MockClient::new("primary", vec![ok_from("primary")]), None)
555            .with_priority(5);
556        let secondary = Backend::new(MockClient::new("secondary", vec![]), None).with_priority(5);
557        let router = RouterClient::new(vec![primary, secondary], RoutingStrategy::Fallback)
558            .expect("non-empty");
559        let resp = router
560            .complete(LLMRequest::builder().build())
561            .await
562            .expect("insertion-order primary first");
563        assert_eq!(resp.content, "from-primary");
564    }
565
566    #[tokio::test]
567    async fn equal_cost_keeps_insertion_order() {
568        let first = Backend::new(
569            MockClient::new("first", vec![ok_from("first")]),
570            Some((1.0, 1.0)),
571        );
572        let second = Backend::new(MockClient::new("second", vec![]), Some((1.0, 1.0)));
573        let router =
574            RouterClient::new(vec![first, second], RoutingStrategy::LowestCost).expect("non-empty");
575        let resp = router
576            .complete(LLMRequest::builder().build())
577            .await
578            .expect("insertion-order first");
579        assert_eq!(resp.content, "from-first");
580    }
581
582    #[test]
583    fn compute_order_lowest_cost_ranks_unknown_and_nan_last() {
584        // indices: 0=nan, 1=unknown(None), 2=mid, 3=cheap.
585        let backends = [
586            Backend::new(MockClient::new("nan", vec![]), Some((f64::NAN, 1.0))),
587            Backend::new(MockClient::new("unknown", vec![]), None),
588            Backend::new(MockClient::new("mid", vec![]), Some((1.0, 2.0))),
589            Backend::new(MockClient::new("cheap", vec![]), Some((0.1, 0.2))),
590        ];
591        let order = compute_order(&backends, RoutingStrategy::LowestCost);
592        // cheap(0.15) < mid(1.5) < nan(MAX) ≤ unknown(MAX); the two MAX entries
593        // tie and stable sort keeps insertion order, so nan(idx 0) precedes
594        // unknown(idx 1).
595        assert_eq!(order, vec![3, 2, 0, 1]);
596    }
597
598    #[test]
599    fn compute_order_fallback_uses_priority_with_ties() {
600        // indices: 0=p5, 1=p1, 2=p5 → [1, 0, 2] (p1 first; p5 tie keeps order).
601        let backends = [
602            Backend::new(MockClient::new("a", vec![]), None).with_priority(5),
603            Backend::new(MockClient::new("b", vec![]), None).with_priority(1),
604            Backend::new(MockClient::new("c", vec![]), None).with_priority(5),
605        ];
606        let order = compute_order(&backends, RoutingStrategy::Fallback);
607        assert_eq!(order, vec![1, 0, 2]);
608    }
609
610    // --- model_name / labels --------------------------------------------------
611
612    #[tokio::test]
613    async fn model_name_reflects_preferred_backend() {
614        let primary = Backend::new(MockClient::new("primary", vec![ok_from("primary")]), None)
615            .with_priority(1);
616        let secondary = Backend::new(
617            MockClient::new("secondary", vec![ok_from("secondary")]),
618            None,
619        )
620        .with_priority(5);
621        let router = RouterClient::new(vec![secondary, primary], RoutingStrategy::Fallback)
622            .expect("non-empty");
623        assert_eq!(router.model_name(), "primary");
624    }
625
626    #[tokio::test]
627    async fn with_label_overrides_model_name() {
628        let backend = Backend::new(MockClient::new("gpt-4o", vec![ok_from("gpt-4o")]), None)
629            .with_label("fallback-alias");
630        let router =
631            RouterClient::new(vec![backend], RoutingStrategy::Fallback).expect("non-empty");
632        assert_eq!(router.model_name(), "fallback-alias");
633    }
634
635    // --- stream delegation ----------------------------------------------------
636
637    #[tokio::test]
638    async fn stream_returns_preferred_and_skips_rest() {
639        // The preferred backend establishes a stream, so the secondary must
640        // never be contacted.
641        let primary_mock = MockClient::new("primary", vec![]);
642        let primary_calls = primary_mock.stream_call_count();
643        let secondary_mock = MockClient::new("secondary", vec![]);
644        let secondary_calls = secondary_mock.stream_call_count();
645
646        let primary = Backend::new(primary_mock, None).with_priority(1);
647        let secondary = Backend::new(secondary_mock, None).with_priority(5);
648        let router = RouterClient::new(vec![secondary, primary], RoutingStrategy::Fallback)
649            .expect("non-empty");
650
651        let stream = router
652            .stream_complete(LLMRequest::builder().build())
653            .await
654            .expect("preferred backend streams");
655        drop(stream); // drop the stream; we only assert delegation.
656
657        assert_eq!(
658            primary_calls.load(Ordering::Relaxed),
659            1,
660            "preferred streamed once"
661        );
662        assert_eq!(
663            secondary_calls.load(Ordering::Relaxed),
664            0,
665            "secondary not contacted"
666        );
667    }
668
669    #[tokio::test]
670    async fn stream_falls_through_on_pre_stream_error() {
671        // Primary fails to establish a stream (5xx before first byte); the
672        // secondary must be tried and succeed.
673        let primary_mock = MockClient::new("primary", vec![]).with_failing_stream(503);
674        let primary_calls = primary_mock.stream_call_count();
675        let secondary_mock = MockClient::new("secondary", vec![]);
676        let secondary_calls = secondary_mock.stream_call_count();
677
678        let primary = Backend::new(primary_mock, None);
679        let secondary = Backend::new(secondary_mock, None);
680        let router = RouterClient::new(vec![primary, secondary], RoutingStrategy::Fallback)
681            .expect("non-empty");
682
683        let stream = router
684            .stream_complete(LLMRequest::builder().build())
685            .await
686            .expect("secondary establishes stream");
687        drop(stream);
688
689        assert_eq!(primary_calls.load(Ordering::Relaxed), 1);
690        assert_eq!(secondary_calls.load(Ordering::Relaxed), 1);
691    }
692
693    #[tokio::test]
694    async fn stream_permanent_4xx_short_circuits() {
695        // A 403 from the primary is permanent — the secondary must not be
696        // contacted on the streaming path either.
697        let primary_mock = MockClient::new("primary", vec![]).with_failing_stream(403);
698        let primary_calls = primary_mock.stream_call_count();
699        let secondary_mock = MockClient::new("secondary", vec![]);
700        let secondary_calls = secondary_mock.stream_call_count();
701
702        let primary = Backend::new(primary_mock, None);
703        let secondary = Backend::new(secondary_mock, None);
704        let router = RouterClient::new(vec![primary, secondary], RoutingStrategy::Fallback)
705            .expect("non-empty");
706
707        let err = router
708            .stream_complete(LLMRequest::builder().build())
709            .await
710            .err()
711            .expect("403 short-circuits");
712        assert!(matches!(err, KernelError::Http { status: 403, .. }));
713        assert_eq!(primary_calls.load(Ordering::Relaxed), 1);
714        assert_eq!(secondary_calls.load(Ordering::Relaxed), 0);
715    }
716}
llm_kernel/llm/router.rs

llm_kernel/llm/
router.rs
