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/// Connections are never pooled: the probe runs every few seconds, which is close
44/// to the idle keep-alive timeout of common upstream servers (uvicorn closes idle
45/// connections after 5 s). Reusing a pooled connection the upstream has already
46/// closed fails with `hyper::Error(IncompleteMessage)` and flaps `/ready` even
47/// though the upstream is healthy.
48///
49/// # Errors
50///
51/// Returns an error when the HTTP client cannot be constructed.
52pub fn llm_readiness_client() -> Result<reqwest::Client, Error> {
53    reqwest::Client::builder()
54        .redirect(reqwest::redirect::Policy::none())
55        .pool_max_idle_per_host(0)
56        .build()
57        .map_err(Error::HttpClient)
58}
59
60/// Probe the inference service's `/health` endpoint once.
61///
62/// # Errors
63///
64/// Returns an error when the configured bearer credential cannot be represented
65/// as an HTTP header.
66pub async fn probe_llm_readiness(
67    client: &reqwest::Client,
68    llm_api_base: &str,
69    openai_api_key: Option<&str>,
70    timeout: Duration,
71) -> Result<LlmReadiness, Error> {
72    let base = llm_api_base.trim_end_matches('/');
73    let url = format!("{base}/health");
74    let mut request = client.get(url);
75    if let Some(key) = openai_api_key.map(str::trim).filter(|key| !key.is_empty()) {
76        let value = reqwest::header::HeaderValue::from_str(&format!("Bearer {key}"))?;
77        request = request.header(reqwest::header::AUTHORIZATION, value);
78    }
79
80    Ok(match tokio::time::timeout(timeout, request.send()).await {
81        Ok(Ok(response)) if response.status().is_success() => LlmReadiness::Ready,
82        Ok(Ok(response)) => LlmReadiness::Rejected(response.status()),
83        Ok(Err(error)) => LlmReadiness::Unreachable(error),
84        Err(_) => LlmReadiness::TimedOut,
85    })
86}
87
88/// Poll LLM `/health` until it responds successfully or the timeout is reached.
89///
90/// # Errors
91///
92/// Returns an error if the LLM does not become ready within the configured timeout.
93pub async fn wait_llm_ready(config: &Config) -> Result<(), Error> {
94    let base = config.llm_api_base.trim_end_matches('/');
95    let url = format!("{base}/health");
96
97    let client = llm_readiness_client()?;
98
99    let timeout = checked_duration_seconds("llm_ready_timeout_s", config.llm_ready_timeout_s)?;
100    let interval = checked_duration_seconds("llm_ready_interval_s", config.llm_ready_interval_s)?;
101    let start = tokio::time::Instant::now();
102    let mut last_notice = Duration::ZERO;
103
104    loop {
105        let remaining = timeout
106            .checked_sub(start.elapsed())
107            .ok_or_else(|| timeout_error(&url, config.llm_ready_timeout_s))?;
108        if remaining.is_zero() {
109            return Err(timeout_error(&url, config.llm_ready_timeout_s));
110        }
111
112        if matches!(
113            probe_llm_readiness(
114                &client,
115                &config.llm_api_base,
116                config.openai_api_key.as_deref(),
117                LLM_READINESS_PROBE_TIMEOUT.min(remaining),
118            )
119            .await?,
120            LlmReadiness::Ready
121        ) {
122            return Ok(());
123        }
124
125        let elapsed = start.elapsed();
126        if elapsed.saturating_sub(last_notice) >= interval {
127            last_notice = elapsed;
128            info!("waiting for LLM ({}s elapsed): {url}", elapsed.as_secs());
129        }
130
131        let remaining = timeout
132            .checked_sub(start.elapsed())
133            .ok_or_else(|| timeout_error(&url, config.llm_ready_timeout_s))?;
134        if remaining.is_zero() {
135            return Err(timeout_error(&url, config.llm_ready_timeout_s));
136        }
137
138        tokio::time::sleep(interval.min(remaining)).await;
139    }
140}
141
142#[cfg(test)]
143mod tests {
144    use std::sync::Arc;
145    use std::sync::atomic::{AtomicUsize, Ordering};
146    use std::time::Duration;
147
148    use axum::Router;
149    use axum::http::HeaderMap;
150    use axum::response::{IntoResponse, Redirect};
151    use axum::routing::get;
152    use http::StatusCode;
153    use tokio::net::TcpListener;
154
155    use super::{checked_duration_seconds, probe_llm_readiness, timeout_error, wait_llm_ready};
156
157    fn test_config(llm_api_base: String) -> crate::config::Config {
158        crate::config::Config {
159            llm_api_base,
160            openai_api_key: Some("test-key".to_owned()),
161            llm_ready_timeout_s: 0.5,
162            llm_ready_interval_s: 0.01,
163            skip_llm_ready_check: false,
164            db_url: None,
165            postgres: crate::config::PostgresConfig::default(),
166            sqlite: crate::config::SqliteConfig::default(),
167            tools: crate::config::ToolRuntimeConfig::default(),
168        }
169    }
170
171    async fn spawn_upstream(app: Router) -> (String, tokio::task::JoinHandle<()>) {
172        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
173        let addr = listener.local_addr().unwrap();
174        let handle = tokio::spawn(async move { axum::serve(listener, app).await.unwrap() });
175        (format!("http://{addr}"), handle)
176    }
177
178    #[test]
179    fn checked_duration_rejects_non_positive() {
180        assert!(checked_duration_seconds("v", 0.0).is_err());
181        assert!(checked_duration_seconds("v", -1.0).is_err());
182    }
183
184    #[test]
185    fn checked_duration_rejects_nan() {
186        assert!(checked_duration_seconds("v", f64::NAN).is_err());
187    }
188
189    #[test]
190    fn checked_duration_rejects_infinite() {
191        assert!(checked_duration_seconds("v", f64::INFINITY).is_err());
192    }
193
194    #[test]
195    fn checked_duration_rejects_too_large_finite() {
196        assert!(checked_duration_seconds("v", 1e50).is_err());
197    }
198
199    #[test]
200    fn checked_duration_accepts_positive_finite() {
201        let duration = checked_duration_seconds("v", 0.25).unwrap();
202        assert_eq!(duration.as_millis(), 250);
203    }
204
205    #[test]
206    fn timeout_error_preserves_inputs() {
207        let err = timeout_error("http://127.0.0.1:8000/health", 0.5);
208        match err {
209            crate::error::Error::LlmTimeout { url, timeout_s } => {
210                assert_eq!(url, "http://127.0.0.1:8000/health");
211                assert!((timeout_s - 0.5).abs() < f64::EPSILON);
212            }
213            other => panic!("expected timeout error, got {other:?}"),
214        }
215    }
216
217    #[test]
218    fn interval_sleep_is_capped_by_remaining_timeout() {
219        let interval = Duration::from_secs(2);
220        let remaining = Duration::from_millis(100);
221        assert_eq!(interval.min(remaining), Duration::from_millis(100));
222    }
223
224    /// Regression test for readiness flapping: the probe must not reuse a pooled
225    /// keep-alive connection between runs, because upstreams such as uvicorn close
226    /// idle connections on roughly the same cadence as the probe and a reused dead
227    /// socket fails with `IncompleteMessage` while the upstream is healthy.
228    #[tokio::test]
229    async fn readiness_client_opens_a_new_connection_per_probe() {
230        use tokio::io::{AsyncReadExt, AsyncWriteExt};
231
232        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
233        let url = format!("http://{}", listener.local_addr().unwrap());
234        let connections = Arc::new(AtomicUsize::new(0));
235        let upstream = tokio::spawn({
236            let connections = Arc::clone(&connections);
237            async move {
238                loop {
239                    let (mut socket, _) = listener.accept().await.unwrap();
240                    connections.fetch_add(1, Ordering::SeqCst);
241                    tokio::spawn(async move {
242                        // Serve keep-alive responses for as many requests as arrive on this socket.
243                        let mut buffer = [0_u8; 2048];
244                        while let Ok(read) = socket.read(&mut buffer).await {
245                            if read == 0 {
246                                break;
247                            }
248                            let response = "HTTP/1.1 204 No Content\r\nconnection: keep-alive\r\n\r\n";
249                            if socket.write_all(response.as_bytes()).await.is_err() {
250                                break;
251                            }
252                        }
253                    });
254                }
255            }
256        });
257
258        let client = super::llm_readiness_client().unwrap();
259        for _ in 0..3 {
260            let readiness = probe_llm_readiness(&client, &url, None, Duration::from_secs(1))
261                .await
262                .unwrap();
263            assert!(matches!(readiness, super::LlmReadiness::Ready), "{readiness:?}");
264        }
265        upstream.abort();
266
267        assert_eq!(
268            connections.load(Ordering::SeqCst),
269            3,
270            "each probe must open its own TCP connection instead of reusing a pooled one"
271        );
272    }
273
274    #[tokio::test]
275    async fn probe_rejects_invalid_bearer_header_before_network_io() {
276        let error = probe_llm_readiness(
277            &reqwest::Client::new(),
278            "http://127.0.0.1:1",
279            Some("invalid\nkey"),
280            Duration::from_secs(1),
281        )
282        .await
283        .unwrap_err();
284
285        assert!(matches!(error, crate::error::Error::InvalidHeader(_)));
286    }
287
288    #[tokio::test]
289    async fn wait_llm_ready_retries_with_authentication_until_success() {
290        let requests = Arc::new(AtomicUsize::new(0));
291        let app = Router::new().route(
292            "/health",
293            get({
294                let requests = Arc::clone(&requests);
295                move |headers: HeaderMap| {
296                    let requests = Arc::clone(&requests);
297                    async move {
298                        if headers.get("authorization").and_then(|value| value.to_str().ok()) != Some("Bearer test-key")
299                        {
300                            return StatusCode::UNAUTHORIZED;
301                        }
302                        if requests.fetch_add(1, Ordering::SeqCst) == 0 {
303                            StatusCode::SERVICE_UNAVAILABLE
304                        } else {
305                            StatusCode::NO_CONTENT
306                        }
307                    }
308                }
309            }),
310        );
311        let (url, upstream) = spawn_upstream(app).await;
312
313        wait_llm_ready(&test_config(url)).await.unwrap();
314
315        assert!(requests.load(Ordering::SeqCst) >= 2);
316        upstream.abort();
317    }
318
319    #[tokio::test]
320    async fn wait_llm_ready_rejects_redirects_until_final_timeout() {
321        let app = Router::new()
322            .route("/health", get(|| async { Redirect::temporary("/login") }))
323            .route("/login", get(|| async { StatusCode::OK.into_response() }));
324        let (url, upstream) = spawn_upstream(app).await;
325        let mut config = test_config(url.clone());
326        config.llm_ready_timeout_s = 0.05;
327
328        let error = wait_llm_ready(&config).await.unwrap_err();
329
330        match error {
331            crate::error::Error::LlmTimeout {
332                url: timed_out_url,
333                timeout_s,
334            } => {
335                assert_eq!(timed_out_url, format!("{url}/health"));
336                assert!((timeout_s - 0.05).abs() < f64::EPSILON);
337            }
338            other => panic!("expected timeout error, got {other:?}"),
339        }
340        upstream.abort();
341    }
342}
agentic_core/readiness.rs

agentic_core/
readiness.rs
