1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
use std::task::{Context, Poll};
use tower::Layer;
use tower::Service;
use super::types::{LlmRequest, LlmResponse};
use crate::client::BoxFuture;
use crate::error::{LiterLlmError, Result};
/// Tower [`Layer`] that routes to a fallback service when the primary service
/// returns an error.
///
/// Only transient errors trigger the fallback — specifically:
/// [`LiterLlmError::RateLimited`], [`LiterLlmError::ServiceUnavailable`],
/// [`LiterLlmError::Timeout`], and [`LiterLlmError::ServerError`].
/// Authentication or bad-request errors are propagated directly without
/// consulting the fallback because retrying on a different service would
/// produce the same result.
#[cfg_attr(alef, alef(skip))]
pub struct FallbackLayer<F> {
fallback: F,
}
#[cfg_attr(alef, alef(skip))]
impl<F> FallbackLayer<F> {
/// Create a new fallback layer with the given fallback service.
#[must_use]
pub fn new(fallback: F) -> Self {
Self { fallback }
}
}
impl<S, F> Layer<S> for FallbackLayer<F>
where
F: Clone,
{
type Service = FallbackService<S, F>;
fn layer(&self, primary: S) -> Self::Service {
FallbackService {
primary,
// Clone the fallback so the produced service owns it independently.
fallback: self.fallback.clone(),
}
}
}
/// Tower service produced by [`FallbackLayer`].
#[cfg_attr(alef, alef(skip))]
pub struct FallbackService<S, F> {
primary: S,
fallback: F,
}
impl<S, F> Clone for FallbackService<S, F>
where
S: Clone,
F: Clone,
{
fn clone(&self) -> Self {
Self {
primary: self.primary.clone(),
fallback: self.fallback.clone(),
}
}
}
impl<S, F> Service<LlmRequest> for FallbackService<S, F>
where
S: Service<LlmRequest, Response = LlmResponse, Error = LiterLlmError> + Send + 'static,
S::Future: Send + 'static,
F: Service<LlmRequest, Response = LlmResponse, Error = LiterLlmError> + Clone + Send + 'static,
F::Future: Send + 'static,
{
type Response = LlmResponse;
type Error = LiterLlmError;
type Future = BoxFuture<'static, Result<LlmResponse>>;
fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<()>> {
// Tower contract: poll_ready should prepare exactly one service for a
// subsequent call. Ideally we would only poll the primary here and
// poll the fallback lazily in `call`. However, because `call` takes
// `&mut self` and must return a `'static` future (no reference to
// `self`), we cannot hold a mutable borrow across the await point.
// For our concrete use case (DefaultClient is always ready), polling
// both here is not harmful — neither service blocks and both remain
// ready until the next call. Callers that compose non-trivially-ready
// services should use a dedicated load-balancing layer instead.
match self.primary.poll_ready(cx) {
Poll::Pending => return Poll::Pending,
Poll::Ready(Err(e)) => return Poll::Ready(Err(e)),
Poll::Ready(Ok(())) => {}
}
self.fallback.poll_ready(cx)
}
fn call(&mut self, req: LlmRequest) -> Self::Future {
// Clone the request so it can be replayed on the fallback if needed.
let fallback_req = req.clone();
let primary_fut = self.primary.call(req);
// `poll_ready` readied `self.fallback` for exactly one call.
// We move the readied service into the async block (so the future is
// 'static) and replace it with a fresh clone for the *next* call cycle.
// Tower's contract guarantees at most one `call` per `poll_ready`, so
// the fresh clone is not used until `poll_ready` runs again.
let fresh = self.fallback.clone();
let mut readied_fallback = std::mem::replace(&mut self.fallback, fresh);
Box::pin(async move {
match primary_fut.await {
Ok(resp) => Ok(resp),
Err(e) if e.is_transient() => {
tracing::warn!(
error = %e,
"primary service failed with transient error; trying fallback"
);
readied_fallback.call(fallback_req).await
}
Err(e) => Err(e),
}
})
}
}