agentic_core/
readiness.rs

1use std::time::Duration;
2
3use tracing::info;
4
5use crate::config::Config;
6use crate::error::Error;
7
8/// Maximum duration of one inference-service readiness probe.
9pub const LLM_READINESS_PROBE_TIMEOUT: Duration = Duration::from_secs(2);
10
11fn checked_duration_seconds(name: &str, value: f64) -> Result<Duration, Error> {
12    if !value.is_finite() || value <= 0.0 {
13        return Err(Error::Config(format!(
14            "{name} must be a finite number > 0 (got {value})"
15        )));
16    }
17    Duration::try_from_secs_f64(value)
18        .map_err(|_| Error::Config(format!("{name} must be representable as a Duration (got {value})")))
19}
20
21fn timeout_error(url: &str, timeout_s: f64) -> Error {
22    Error::LlmTimeout {
23        url: url.to_owned(),
24        timeout_s,
25    }
26}
27
28/// Result of a single bounded inference-service health probe.
29#[derive(Debug)]
30#[non_exhaustive]
31pub enum LlmReadiness {
32    Ready,
33    Rejected(reqwest::StatusCode),
34    Unreachable(reqwest::Error),
35    TimedOut,
36}
37
38/// Build the dedicated HTTP client used for inference-service health probes.
39///
40/// The client rejects redirects so an authentication page or generic UI cannot
41/// turn an unsuccessful `/health` response into a false-positive readiness result.
42///
43/// # Errors
44///
45/// Returns an error when the HTTP client cannot be constructed.
46pub fn llm_readiness_client() -> Result<reqwest::Client, Error> {
47    reqwest::Client::builder()
48        .redirect(reqwest::redirect::Policy::none())
49        .build()
50        .map_err(Error::HttpClient)
51}
52
53/// Probe the inference service's `/health` endpoint once.
54///
55/// # Errors
56///
57/// Returns an error when the configured bearer credential cannot be represented
58/// as an HTTP header.
59pub async fn probe_llm_readiness(
60    client: &reqwest::Client,
61    llm_api_base: &str,
62    openai_api_key: Option<&str>,
63    timeout: Duration,
64) -> Result<LlmReadiness, Error> {
65    let base = llm_api_base.trim_end_matches('/');
66    let url = format!("{base}/health");
67    let mut request = client.get(url);
68    if let Some(key) = openai_api_key.map(str::trim).filter(|key| !key.is_empty()) {
69        let value = reqwest::header::HeaderValue::from_str(&format!("Bearer {key}"))?;
70        request = request.header(reqwest::header::AUTHORIZATION, value);
71    }
72
73    Ok(match tokio::time::timeout(timeout, request.send()).await {
74        Ok(Ok(response)) if response.status().is_success() => LlmReadiness::Ready,
75        Ok(Ok(response)) => LlmReadiness::Rejected(response.status()),
76        Ok(Err(error)) => LlmReadiness::Unreachable(error),
77        Err(_) => LlmReadiness::TimedOut,
78    })
79}
80
81/// Poll LLM `/health` until it responds successfully or the timeout is reached.
82///
83/// # Errors
84///
85/// Returns an error if the LLM does not become ready within the configured timeout.
86pub async fn wait_llm_ready(config: &Config) -> Result<(), Error> {
87    let base = config.llm_api_base.trim_end_matches('/');
88    let url = format!("{base}/health");
89
90    let client = llm_readiness_client()?;
91
92    let timeout = checked_duration_seconds("llm_ready_timeout_s", config.llm_ready_timeout_s)?;
93    let interval = checked_duration_seconds("llm_ready_interval_s", config.llm_ready_interval_s)?;
94    let start = tokio::time::Instant::now();
95    let mut last_notice = Duration::ZERO;
96
97    loop {
98        let remaining = timeout
99            .checked_sub(start.elapsed())
100            .ok_or_else(|| timeout_error(&url, config.llm_ready_timeout_s))?;
101        if remaining.is_zero() {
102            return Err(timeout_error(&url, config.llm_ready_timeout_s));
103        }
104
105        if matches!(
106            probe_llm_readiness(
107                &client,
108                &config.llm_api_base,
109                config.openai_api_key.as_deref(),
110                LLM_READINESS_PROBE_TIMEOUT.min(remaining),
111            )
112            .await?,
113            LlmReadiness::Ready
114        ) {
115            return Ok(());
116        }
117
118        let elapsed = start.elapsed();
119        if elapsed.saturating_sub(last_notice) >= interval {
120            last_notice = elapsed;
121            info!("waiting for LLM ({}s elapsed): {url}", elapsed.as_secs());
122        }
123
124        let remaining = timeout
125            .checked_sub(start.elapsed())
126            .ok_or_else(|| timeout_error(&url, config.llm_ready_timeout_s))?;
127        if remaining.is_zero() {
128            return Err(timeout_error(&url, config.llm_ready_timeout_s));
129        }
130
131        tokio::time::sleep(interval.min(remaining)).await;
132    }
133}
134
135#[cfg(test)]
136mod tests {
137    use std::sync::Arc;
138    use std::sync::atomic::{AtomicUsize, Ordering};
139    use std::time::Duration;
140
141    use axum::Router;
142    use axum::http::HeaderMap;
143    use axum::response::{IntoResponse, Redirect};
144    use axum::routing::get;
145    use http::StatusCode;
146    use tokio::net::TcpListener;
147
148    use super::{checked_duration_seconds, probe_llm_readiness, timeout_error, wait_llm_ready};
149
150    fn test_config(llm_api_base: String) -> crate::config::Config {
151        crate::config::Config {
152            llm_api_base,
153            openai_api_key: Some("test-key".to_owned()),
154            llm_ready_timeout_s: 0.5,
155            llm_ready_interval_s: 0.01,
156            skip_llm_ready_check: false,
157            db_url: None,
158            postgres: crate::config::PostgresConfig::default(),
159            sqlite: crate::config::SqliteConfig::default(),
160        }
161    }
162
163    async fn spawn_upstream(app: Router) -> (String, tokio::task::JoinHandle<()>) {
164        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
165        let addr = listener.local_addr().unwrap();
166        let handle = tokio::spawn(async move { axum::serve(listener, app).await.unwrap() });
167        (format!("http://{addr}"), handle)
168    }
169
170    #[test]
171    fn checked_duration_rejects_non_positive() {
172        assert!(checked_duration_seconds("v", 0.0).is_err());
173        assert!(checked_duration_seconds("v", -1.0).is_err());
174    }
175
176    #[test]
177    fn checked_duration_rejects_nan() {
178        assert!(checked_duration_seconds("v", f64::NAN).is_err());
179    }
180
181    #[test]
182    fn checked_duration_rejects_infinite() {
183        assert!(checked_duration_seconds("v", f64::INFINITY).is_err());
184    }
185
186    #[test]
187    fn checked_duration_rejects_too_large_finite() {
188        assert!(checked_duration_seconds("v", 1e50).is_err());
189    }
190
191    #[test]
192    fn checked_duration_accepts_positive_finite() {
193        let duration = checked_duration_seconds("v", 0.25).unwrap();
194        assert_eq!(duration.as_millis(), 250);
195    }
196
197    #[test]
198    fn timeout_error_preserves_inputs() {
199        let err = timeout_error("http://127.0.0.1:8000/health", 0.5);
200        match err {
201            crate::error::Error::LlmTimeout { url, timeout_s } => {
202                assert_eq!(url, "http://127.0.0.1:8000/health");
203                assert!((timeout_s - 0.5).abs() < f64::EPSILON);
204            }
205            other => panic!("expected timeout error, got {other:?}"),
206        }
207    }
208
209    #[test]
210    fn interval_sleep_is_capped_by_remaining_timeout() {
211        let interval = Duration::from_secs(2);
212        let remaining = Duration::from_millis(100);
213        assert_eq!(interval.min(remaining), Duration::from_millis(100));
214    }
215
216    #[tokio::test]
217    async fn probe_rejects_invalid_bearer_header_before_network_io() {
218        let error = probe_llm_readiness(
219            &reqwest::Client::new(),
220            "http://127.0.0.1:1",
221            Some("invalid\nkey"),
222            Duration::from_secs(1),
223        )
224        .await
225        .unwrap_err();
226
227        assert!(matches!(error, crate::error::Error::InvalidHeader(_)));
228    }
229
230    #[tokio::test]
231    async fn wait_llm_ready_retries_with_authentication_until_success() {
232        let requests = Arc::new(AtomicUsize::new(0));
233        let app = Router::new().route(
234            "/health",
235            get({
236                let requests = Arc::clone(&requests);
237                move |headers: HeaderMap| {
238                    let requests = Arc::clone(&requests);
239                    async move {
240                        if headers.get("authorization").and_then(|value| value.to_str().ok()) != Some("Bearer test-key")
241                        {
242                            return StatusCode::UNAUTHORIZED;
243                        }
244                        if requests.fetch_add(1, Ordering::SeqCst) == 0 {
245                            StatusCode::SERVICE_UNAVAILABLE
246                        } else {
247                            StatusCode::NO_CONTENT
248                        }
249                    }
250                }
251            }),
252        );
253        let (url, upstream) = spawn_upstream(app).await;
254
255        wait_llm_ready(&test_config(url)).await.unwrap();
256
257        assert!(requests.load(Ordering::SeqCst) >= 2);
258        upstream.abort();
259    }
260
261    #[tokio::test]
262    async fn wait_llm_ready_rejects_redirects_until_final_timeout() {
263        let app = Router::new()
264            .route("/health", get(|| async { Redirect::temporary("/login") }))
265            .route("/login", get(|| async { StatusCode::OK.into_response() }));
266        let (url, upstream) = spawn_upstream(app).await;
267        let mut config = test_config(url.clone());
268        config.llm_ready_timeout_s = 0.05;
269
270        let error = wait_llm_ready(&config).await.unwrap_err();
271
272        match error {
273            crate::error::Error::LlmTimeout {
274                url: timed_out_url,
275                timeout_s,
276            } => {
277                assert_eq!(timed_out_url, format!("{url}/health"));
278                assert!((timeout_s - 0.05).abs() < f64::EPSILON);
279            }
280            other => panic!("expected timeout error, got {other:?}"),
281        }
282        upstream.abort();
283    }
284}
agentic_core/readiness.rs

agentic_core/
readiness.rs
