Skip to main content

agentic_core/
readiness.rs

1use std::time::Duration;
2
3use tracing::info;
4
5use crate::config::Config;
6use crate::error::Error;
7
8fn checked_duration_seconds(name: &str, value: f64) -> Result<Duration, Error> {
9    if !value.is_finite() || value <= 0.0 {
10        return Err(Error::Config(format!(
11            "{name} must be a finite number > 0 (got {value})"
12        )));
13    }
14    Duration::try_from_secs_f64(value)
15        .map_err(|_| Error::Config(format!("{name} must be representable as a Duration (got {value})")))
16}
17
18fn timeout_error(url: &str, timeout_s: f64) -> Error {
19    Error::LlmTimeout {
20        url: url.to_owned(),
21        timeout_s,
22    }
23}
24
25/// Poll LLM `/health` until it responds 200 or the timeout is reached.
26///
27/// # Errors
28///
29/// Returns an error if the LLM does not become ready within the configured timeout.
30pub async fn wait_llm_ready(config: &Config) -> Result<(), Error> {
31    let base = config.llm_api_base.trim_end_matches('/');
32    let url = format!("{base}/health");
33
34    let mut headers = reqwest::header::HeaderMap::new();
35    if let Some(key) = config.openai_api_key.as_deref() {
36        let trimmed = key.trim();
37        if !trimmed.is_empty() {
38            headers.insert(
39                reqwest::header::AUTHORIZATION,
40                reqwest::header::HeaderValue::from_str(&format!("Bearer {trimmed}"))?,
41            );
42        }
43    }
44
45    let client = reqwest::Client::builder()
46        .timeout(Duration::from_secs(2))
47        .default_headers(headers)
48        .build()
49        .map_err(Error::HttpClient)?;
50
51    let timeout = checked_duration_seconds("llm_ready_timeout_s", config.llm_ready_timeout_s)?;
52    let interval = checked_duration_seconds("llm_ready_interval_s", config.llm_ready_interval_s)?;
53    let start = tokio::time::Instant::now();
54    let mut last_notice = Duration::ZERO;
55
56    loop {
57        let remaining = timeout
58            .checked_sub(start.elapsed())
59            .ok_or_else(|| timeout_error(&url, config.llm_ready_timeout_s))?;
60        if remaining.is_zero() {
61            return Err(timeout_error(&url, config.llm_ready_timeout_s));
62        }
63
64        match tokio::time::timeout(remaining, client.get(&url).send()).await {
65            Ok(Ok(resp)) if resp.status().as_u16() == 200 => return Ok(()),
66            _ => {}
67        }
68
69        let elapsed = start.elapsed();
70        if elapsed.saturating_sub(last_notice) >= interval {
71            last_notice = elapsed;
72            info!("waiting for LLM ({}s elapsed): {url}", elapsed.as_secs());
73        }
74
75        let remaining = timeout
76            .checked_sub(start.elapsed())
77            .ok_or_else(|| timeout_error(&url, config.llm_ready_timeout_s))?;
78        if remaining.is_zero() {
79            return Err(timeout_error(&url, config.llm_ready_timeout_s));
80        }
81
82        tokio::time::sleep(interval.min(remaining)).await;
83    }
84}
85
86#[cfg(test)]
87mod tests {
88    use std::time::Duration;
89
90    use super::{checked_duration_seconds, timeout_error};
91
92    #[test]
93    fn checked_duration_rejects_non_positive() {
94        assert!(checked_duration_seconds("v", 0.0).is_err());
95        assert!(checked_duration_seconds("v", -1.0).is_err());
96    }
97
98    #[test]
99    fn checked_duration_rejects_nan() {
100        assert!(checked_duration_seconds("v", f64::NAN).is_err());
101    }
102
103    #[test]
104    fn checked_duration_rejects_infinite() {
105        assert!(checked_duration_seconds("v", f64::INFINITY).is_err());
106    }
107
108    #[test]
109    fn checked_duration_rejects_too_large_finite() {
110        assert!(checked_duration_seconds("v", 1e50).is_err());
111    }
112
113    #[test]
114    fn checked_duration_accepts_positive_finite() {
115        let duration = checked_duration_seconds("v", 0.25).unwrap();
116        assert_eq!(duration.as_millis(), 250);
117    }
118
119    #[test]
120    fn timeout_error_preserves_inputs() {
121        let err = timeout_error("http://127.0.0.1:8000/health", 0.5);
122        match err {
123            crate::error::Error::LlmTimeout { url, timeout_s } => {
124                assert_eq!(url, "http://127.0.0.1:8000/health");
125                assert!((timeout_s - 0.5).abs() < f64::EPSILON);
126            }
127            other => panic!("expected timeout error, got {other:?}"),
128        }
129    }
130
131    #[test]
132    fn interval_sleep_is_capped_by_remaining_timeout() {
133        let interval = Duration::from_secs(2);
134        let remaining = Duration::from_millis(100);
135        assert_eq!(interval.min(remaining), Duration::from_millis(100));
136    }
137}