machi_llm/
breaker_sampler.rs

1//! [`BreakerSampler`]: decorator that gates samples through a [`CircuitBreaker`].
2//!
3//! Maturity: **core**
4
5use std::sync::Arc;
6
7use async_trait::async_trait;
8use machi_types::{ErrorCode, MachiError};
9
10use crate::breaker::{Admission, BreakerOutcome, CircuitBreaker};
11use crate::sample::{SampleRequest, SampleResponse};
12use crate::sampler::LlmSampler;
13use crate::stream::SampleStream;
14
15/// Sampler wrapper that refuses traffic while the breaker is open.
16#[derive(Debug, Clone)]
17pub struct BreakerSampler<S> {
18    inner: Arc<S>,
19    breaker: Arc<CircuitBreaker>,
20    /// Stable key for multi-endpoint registries (metrics / logs).
21    endpoint: String,
22}
23
24impl<S> BreakerSampler<S> {
25    /// Wrap `inner` with a shared breaker.
26    #[must_use]
27    pub fn new(inner: Arc<S>, breaker: Arc<CircuitBreaker>, endpoint: impl Into<String>) -> Self {
28        Self {
29            inner,
30            breaker,
31            endpoint: endpoint.into(),
32        }
33    }
34
35    /// Endpoint label.
36    #[must_use]
37    pub fn endpoint(&self) -> &str {
38        &self.endpoint
39    }
40
41    /// Shared breaker.
42    #[must_use]
43    pub fn breaker(&self) -> &Arc<CircuitBreaker> {
44        &self.breaker
45    }
46}
47
48#[async_trait]
49impl<S: LlmSampler + 'static> LlmSampler for BreakerSampler<S> {
50    async fn sample(&self, request: SampleRequest) -> Result<SampleResponse, MachiError> {
51        self.admit()?;
52        match self.inner.sample(request).await {
53            Ok(r) => {
54                self.breaker.record(BreakerOutcome::Success);
55                Ok(r)
56            }
57            Err(e) => {
58                self.breaker.record(BreakerOutcome::Failure);
59                Err(e)
60            }
61        }
62    }
63
64    async fn sample_stream(&self, request: SampleRequest) -> Result<SampleStream, MachiError> {
65        self.admit()?;
66        match self.inner.sample_stream(request).await {
67            Ok(s) => {
68                self.breaker.record(BreakerOutcome::Success);
69                Ok(s)
70            }
71            Err(e) => {
72                self.breaker.record(BreakerOutcome::Failure);
73                Err(e)
74            }
75        }
76    }
77}
78
79impl<S> BreakerSampler<S> {
80    fn admit(&self) -> Result<(), MachiError> {
81        match self.breaker.check() {
82            Admission::Allow => Ok(()),
83            Admission::Reject { retry_after } => Err(MachiError::new(
84                ErrorCode::LlmProvider,
85                format!(
86                    "circuit breaker open for endpoint '{}'; retry after {}ms",
87                    self.endpoint,
88                    retry_after.as_millis()
89                ),
90            )
91            .with_retry(machi_types::RetryClass::Backoff)),
92        }
93    }
94}
machi_llm/breaker_sampler.rs

machi_llm/
breaker_sampler.rs
