polyc_judgment/

fallback.rs

//! [`FallbackJudgment`]: compose a primary judgment backend with a fallback
//! one.

use async_trait::async_trait;

use crate::{JudgmentError, JudgmentProvider, JudgmentRequest, JudgmentResponse, JudgmentSource};

/// A judgment backend that tries a primary backend first and falls back to a
/// second backend when the primary fails in a way the fallback can fix.
///
/// Both backends answer through the shared [`JudgmentError`] taxonomy, so the
/// decision of whether to fall back is one match over `P`'s error, not a
/// per-backend special case.
pub struct FallbackJudgment<P, F> {
    primary: P,
    fallback: F,
}

impl<P, F> FallbackJudgment<P, F> {
    /// Composes `primary` with `fallback`.
    #[must_use]
    pub const fn new(primary: P, fallback: F) -> Self {
        Self { primary, fallback }
    }

    /// The primary backend.
    #[must_use]
    pub const fn primary(&self) -> &P {
        &self.primary
    }

    /// The fallback backend.
    #[must_use]
    pub const fn fallback(&self) -> &F {
        &self.fallback
    }
}

#[async_trait]
impl<P, F> JudgmentProvider for FallbackJudgment<P, F>
where
    P: JudgmentProvider<Error = JudgmentError>,
    F: JudgmentProvider<Error = JudgmentError>,
{
    type Error = JudgmentError;

    async fn judge(&self, request: JudgmentRequest) -> Result<JudgmentResponse, JudgmentError> {
        match self.primary.judge(request.clone()).await {
            // `source` is set here, not trusted from the primary backend: a
            // caller must be able to tell primary from fallback answers by
            // `FallbackJudgment`'s own routing, not by whether the backend
            // happened to fill the field correctly.
            Ok(response) => Ok(JudgmentResponse {
                source: JudgmentSource::Primary,
                ..response
            }),
            // A rate limit, an outage, a broken transport, a refused
            // credential, or an exhausted balance are all reasons the
            // primary backend specifically cannot answer right now — none
            // say the request itself is bad, so a different backend gets a
            // real chance. An expired or invalid gateway credential must not
            // silently take down the gate, so `Unauthorized` falls back here
            // too, unlike the two arms below; an out-of-credit gateway is
            // the same shape of failure from the gate's point of view, so
            // `Exhausted` falls back too.
            Err(
                JudgmentError::RateLimited { .. }
                | JudgmentError::Unavailable { .. }
                | JudgmentError::Transport(_)
                | JudgmentError::Unauthorized
                | JudgmentError::Exhausted,
            ) => self
                .fallback
                .judge(request)
                .await
                .map(|response| JudgmentResponse {
                    source: JudgmentSource::Fallback,
                    ..response
                }),
            // The request shape itself is rejected, or the primary's
            // response didn't decode — a fallback backend would be asked the
            // same rejected request or would see the same non-answer, so
            // falling back cannot fix either.
            Err(err @ (JudgmentError::Invalid(_) | JudgmentError::Malformed(_))) => Err(err),
        }
    }
}

#[cfg(test)]
mod tests {
    #![allow(clippy::pedantic, clippy::nursery, missing_docs)]

    use std::sync::atomic::{AtomicUsize, Ordering};

    use async_trait::async_trait;

    use super::super::{
        Answer, JudgmentError, JudgmentProvider, JudgmentRequest, JudgmentResponse, JudgmentSource,
        JudgmentUsage,
    };
    use super::FallbackJudgment;

    /// A scripted stub: answers with a fixed `Result`, counts its calls.
    struct Scripted {
        result: Result<f64, JudgmentError>,
        calls: AtomicUsize,
    }

    impl Scripted {
        fn ok(noul: f64) -> Self {
            Self {
                result: Ok(noul),
                calls: AtomicUsize::new(0),
            }
        }

        fn err(build: impl Fn() -> JudgmentError) -> Self {
            Self {
                result: Err(build()),
                calls: AtomicUsize::new(0),
            }
        }

        fn calls(&self) -> usize {
            self.calls.load(Ordering::SeqCst)
        }
    }

    #[async_trait]
    impl JudgmentProvider for Scripted {
        type Error = JudgmentError;

        async fn judge(&self, _: JudgmentRequest) -> Result<JudgmentResponse, JudgmentError> {
            self.calls.fetch_add(1, Ordering::SeqCst);
            match &self.result {
                Ok(noul) => {
                    let mut answers = std::collections::BTreeMap::new();
                    answers.insert("q".to_owned(), Answer::Noul { noul: *noul });
                    Ok(JudgmentResponse {
                        model: "scripted".to_owned(),
                        answers,
                        usage: JudgmentUsage::default(),
                        source: JudgmentSource::Primary,
                    })
                }
                Err(_) => Err(clone_err(&self.result)),
            }
        }
    }

    /// `JudgmentError` doesn't derive `Clone` (it wraps a boxed transport
    /// error) — rebuild an equivalent value from the scripted variant instead.
    fn clone_err(result: &Result<f64, JudgmentError>) -> JudgmentError {
        match result {
            Ok(_) => unreachable!("only called on the Err arm"),
            Err(JudgmentError::Invalid(s)) => JudgmentError::Invalid(s.clone()),
            Err(JudgmentError::Unauthorized) => JudgmentError::Unauthorized,
            Err(JudgmentError::Exhausted) => JudgmentError::Exhausted,
            Err(JudgmentError::RateLimited { retry_after }) => JudgmentError::RateLimited {
                retry_after: *retry_after,
            },
            Err(JudgmentError::Unavailable { status }) => {
                JudgmentError::Unavailable { status: *status }
            }
            Err(JudgmentError::Transport(_)) => {
                JudgmentError::Transport(Box::new(std::io::Error::other("scripted transport")))
            }
            Err(JudgmentError::Malformed(s)) => JudgmentError::Malformed(s.clone()),
        }
    }

    fn request() -> JudgmentRequest {
        JudgmentRequest {
            state: serde_json::json!({}),
            questions: std::collections::BTreeMap::new(),
        }
    }

    #[tokio::test]
    async fn primary_success_never_calls_the_fallback() {
        let primary = Scripted::ok(0.8);
        let fallback = Scripted::ok(0.1);
        let gate = FallbackJudgment::new(primary, fallback);
        let response = gate.judge(request()).await.expect("primary answered");
        assert_eq!(response.noul("q"), Some(0.8));
        assert_eq!(response.source, JudgmentSource::Primary);
        assert_eq!(gate.primary().calls(), 1);
        assert_eq!(gate.fallback().calls(), 0);
    }

    #[tokio::test]
    async fn primary_rate_limited_falls_back() {
        let primary = Scripted::err(|| JudgmentError::RateLimited { retry_after: None });
        let fallback = Scripted::ok(0.6);
        let gate = FallbackJudgment::new(primary, fallback);
        let response = gate.judge(request()).await.expect("fallback answered");
        assert_eq!(response.noul("q"), Some(0.6));
        assert_eq!(response.source, JudgmentSource::Fallback);
        assert_eq!(gate.primary().calls(), 1);
        assert_eq!(gate.fallback().calls(), 1);
    }

    #[tokio::test]
    async fn primary_unavailable_falls_back() {
        let primary = Scripted::err(|| JudgmentError::Unavailable { status: 503 });
        let fallback = Scripted::ok(0.6);
        let gate = FallbackJudgment::new(primary, fallback);
        let response = gate.judge(request()).await.expect("fallback answered");
        assert_eq!(response.source, JudgmentSource::Fallback);
        assert_eq!(gate.fallback().calls(), 1);
    }

    #[tokio::test]
    async fn primary_transport_failure_falls_back() {
        let primary = Scripted::err(|| {
            JudgmentError::Transport(Box::new(std::io::Error::other("connect refused")))
        });
        let fallback = Scripted::ok(0.6);
        let gate = FallbackJudgment::new(primary, fallback);
        let response = gate.judge(request()).await.expect("fallback answered");
        assert_eq!(response.source, JudgmentSource::Fallback);
        assert_eq!(gate.fallback().calls(), 1);
    }

    /// An expired or invalid gateway credential must not silently take down
    /// the gate, so `Unauthorized` also triggers fallback.
    #[tokio::test]
    async fn primary_unauthorized_falls_back() {
        let primary = Scripted::err(|| JudgmentError::Unauthorized);
        let fallback = Scripted::ok(0.6);
        let gate = FallbackJudgment::new(primary, fallback);
        let response = gate.judge(request()).await.expect("fallback answered");
        assert_eq!(response.source, JudgmentSource::Fallback);
        assert_eq!(gate.fallback().calls(), 1);
    }

    /// An exhausted gateway balance is an availability failure from the
    /// gate's point of view, so it falls back like an outage.
    #[tokio::test]
    async fn primary_exhausted_falls_back() {
        let primary = Scripted::err(|| JudgmentError::Exhausted);
        let fallback = Scripted::ok(0.6);
        let gate = FallbackJudgment::new(primary, fallback);
        let response = gate.judge(request()).await.expect("fallback answered");
        assert_eq!(response.source, JudgmentSource::Fallback);
        assert_eq!(gate.fallback().calls(), 1);
    }

    /// A bad-request shape is never fixed by a different backend.
    #[tokio::test]
    async fn primary_invalid_returns_as_is_without_a_fallback_call() {
        let primary = Scripted::err(|| JudgmentError::Invalid("bad state shape".to_owned()));
        let fallback = Scripted::ok(0.6);
        let gate = FallbackJudgment::new(primary, fallback);
        let err = gate
            .judge(request())
            .await
            .expect_err("no fallback for Invalid");
        assert!(matches!(err, JudgmentError::Invalid(_)), "{err:?}");
        assert_eq!(gate.fallback().calls(), 0);
    }

    #[tokio::test]
    async fn primary_malformed_returns_as_is_without_a_fallback_call() {
        let primary = Scripted::err(|| JudgmentError::Malformed("not json".to_owned()));
        let fallback = Scripted::ok(0.6);
        let gate = FallbackJudgment::new(primary, fallback);
        let err = gate
            .judge(request())
            .await
            .expect_err("no fallback for Malformed");
        assert!(matches!(err, JudgmentError::Malformed(_)), "{err:?}");
        assert_eq!(gate.fallback().calls(), 0);
    }
}