Skip to main content

harn_vm/llm/
healthcheck.rs

1use std::collections::BTreeMap;
2use std::time::Duration;
3
4use reqwest::Method;
5use serde::{Deserialize, Serialize};
6use serde_json::{json, Value as JsonValue};
7
8use crate::llm_config::{self, AuthEnv, HealthcheckDef, ProviderDef};
9
10use super::api::apply_auth_headers;
11
12const DEFAULT_HEALTHCHECK_TIMEOUT_SECS: u64 = 5;
13const BODY_SNIPPET_LIMIT: usize = 1000;
14
15#[derive(Debug, Clone, Default)]
16pub struct ProviderHealthcheckOptions {
17    /// Candidate API key to validate. When unset, Harn resolves credentials
18    /// from the provider's configured environment variables.
19    pub api_key: Option<String>,
20    /// Optional client override for hosts that need custom transport policy.
21    pub client: Option<reqwest::Client>,
22}
23
24#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
25pub struct ProviderHealthcheckResult {
26    pub provider: String,
27    pub valid: bool,
28    pub message: String,
29    pub metadata: BTreeMap<String, JsonValue>,
30}
31
32impl ProviderHealthcheckResult {
33    fn new(
34        provider: impl Into<String>,
35        valid: bool,
36        message: impl Into<String>,
37        metadata: BTreeMap<String, JsonValue>,
38    ) -> Self {
39        Self {
40            provider: provider.into(),
41            valid,
42            message: message.into(),
43            metadata,
44        }
45    }
46}
47
48pub async fn run_provider_healthcheck(provider: &str) -> ProviderHealthcheckResult {
49    run_provider_healthcheck_with_options(provider, ProviderHealthcheckOptions::default()).await
50}
51
52pub async fn run_provider_healthcheck_with_options(
53    provider: &str,
54    options: ProviderHealthcheckOptions,
55) -> ProviderHealthcheckResult {
56    let provider = if provider.trim().is_empty() {
57        "anthropic"
58    } else {
59        provider.trim()
60    };
61
62    let Some(def) = llm_config::provider_config(provider) else {
63        let mut metadata = base_metadata("unknown_provider");
64        metadata.insert("provider".to_string(), json!(provider));
65        return ProviderHealthcheckResult::new(
66            provider,
67            false,
68            format!("Unknown provider: {provider}"),
69            metadata,
70        );
71    };
72
73    let Some(healthcheck) = def.healthcheck.as_ref() else {
74        let mut metadata = base_metadata("no_healthcheck");
75        metadata.insert("provider".to_string(), json!(provider));
76        return ProviderHealthcheckResult::new(
77            provider,
78            false,
79            format!("No healthcheck configured for {provider}"),
80            metadata,
81        );
82    };
83
84    let auth = resolve_healthcheck_auth(&def, options.api_key);
85    if auth.requires_auth && auth.api_key.is_none() {
86        let mut metadata = base_metadata("missing_credentials");
87        metadata.insert("provider".to_string(), json!(provider));
88        metadata.insert("auth_env".to_string(), json!(auth.candidates));
89        return ProviderHealthcheckResult::new(
90            provider,
91            false,
92            format!(
93                "Missing credentials for {provider}: set {} or pass an api_key",
94                auth.candidates.join(", ")
95            ),
96            metadata,
97        );
98    }
99
100    let url = build_healthcheck_url(&def, healthcheck);
101    let method = Method::from_bytes(healthcheck.method.as_bytes()).unwrap_or(Method::GET);
102    let client = match options.client {
103        Some(client) => client,
104        None => match reqwest::Client::builder()
105            .timeout(Duration::from_secs(DEFAULT_HEALTHCHECK_TIMEOUT_SECS))
106            .redirect(crate::egress::redirect_policy(
107                "llm_healthcheck_redirect",
108                5,
109            ))
110            .build()
111        {
112            Ok(client) => client,
113            Err(error) => {
114                let mut metadata = base_metadata("client_build_failed");
115                metadata.insert("provider".to_string(), json!(provider));
116                return ProviderHealthcheckResult::new(
117                    provider,
118                    false,
119                    format!("{provider} healthcheck failed: {error}"),
120                    metadata,
121                );
122            }
123        },
124    };
125
126    let mut request = client.request(method.clone(), &url);
127    if let Some(api_key) = auth.api_key.as_deref() {
128        request = apply_auth_headers(request, api_key, Some(&def));
129    }
130    for (name, value) in &def.extra_headers {
131        request = request.header(name, value);
132    }
133    if let Some(body) = &healthcheck.body {
134        request = request
135            .header(reqwest::header::CONTENT_TYPE, "application/json")
136            .body(body.clone());
137    }
138
139    match request.send().await {
140        Ok(response) => {
141            let status = response.status();
142            let status_code = status.as_u16();
143            let valid = status.is_success();
144            let body_text = response.text().await.unwrap_or_default();
145            let mut metadata = base_metadata(if valid { "ok" } else { "http_status" });
146            metadata.insert("provider".to_string(), json!(provider));
147            metadata.insert("status".to_string(), json!(status_code));
148            metadata.insert("url".to_string(), json!(url));
149            metadata.insert("method".to_string(), json!(method.as_str()));
150            if !valid && !body_text.is_empty() {
151                metadata.insert("body".to_string(), json!(body_snippet(&body_text)));
152            }
153
154            let message = if valid {
155                format!("{provider} is reachable (HTTP {status_code})")
156            } else {
157                let suffix = body_snippet(&body_text);
158                if suffix.is_empty() {
159                    format!("{provider} returned HTTP {status_code}")
160                } else {
161                    format!("{provider} returned HTTP {status_code}: {suffix}")
162                }
163            };
164
165            ProviderHealthcheckResult::new(provider, valid, message, metadata)
166        }
167        Err(error) => {
168            let mut metadata = base_metadata("request_failed");
169            metadata.insert("provider".to_string(), json!(provider));
170            metadata.insert(
171                "url".to_string(),
172                json!(crate::egress::redact_diagnostic_text(&url)),
173            );
174            metadata.insert("method".to_string(), json!(method.as_str()));
175            let error = crate::egress::redact_reqwest_error(&error);
176            ProviderHealthcheckResult::new(
177                provider,
178                false,
179                format!("{provider} healthcheck failed: {error}"),
180                metadata,
181            )
182        }
183    }
184}
185
186pub fn build_healthcheck_url(def: &ProviderDef, healthcheck: &HealthcheckDef) -> String {
187    if let Some(url) = &healthcheck.url {
188        return url.clone();
189    }
190
191    let base = llm_config::resolve_base_url(def);
192    let path = healthcheck.path.as_deref().unwrap_or("");
193    if path.starts_with('/') {
194        format!("{}{}", base.trim_end_matches('/'), path)
195    } else if path.is_empty() {
196        base
197    } else {
198        format!("{}/{}", base.trim_end_matches('/'), path)
199    }
200}
201
202#[derive(Debug, Clone)]
203struct ResolvedHealthcheckAuth {
204    requires_auth: bool,
205    api_key: Option<String>,
206    candidates: Vec<String>,
207}
208
209fn resolve_healthcheck_auth(
210    def: &ProviderDef,
211    api_key_override: Option<String>,
212) -> ResolvedHealthcheckAuth {
213    let candidates = auth_env_candidates(&def.auth_env);
214    if def.auth_style == "none" || matches!(def.auth_env, AuthEnv::None) {
215        let api_key = api_key_override.and_then(non_empty);
216        return ResolvedHealthcheckAuth {
217            requires_auth: api_key.is_some(),
218            api_key,
219            candidates,
220        };
221    }
222
223    let api_key = api_key_override
224        .and_then(non_empty)
225        .or_else(|| resolve_api_key_from_env(&def.auth_env));
226    ResolvedHealthcheckAuth {
227        requires_auth: true,
228        api_key,
229        candidates,
230    }
231}
232
233fn auth_env_candidates(auth_env: &AuthEnv) -> Vec<String> {
234    match auth_env {
235        AuthEnv::None => Vec::new(),
236        AuthEnv::Single(env) => vec![env.clone()],
237        AuthEnv::Multiple(envs) => envs.clone(),
238    }
239}
240
241fn resolve_api_key_from_env(auth_env: &AuthEnv) -> Option<String> {
242    match auth_env {
243        AuthEnv::None => None,
244        AuthEnv::Single(env) => std::env::var(env).ok().and_then(non_empty),
245        AuthEnv::Multiple(envs) => envs
246            .iter()
247            .find_map(|env| std::env::var(env).ok().and_then(non_empty)),
248    }
249}
250
251fn non_empty(value: String) -> Option<String> {
252    let trimmed = value.trim();
253    if trimmed.is_empty() {
254        None
255    } else {
256        Some(trimmed.to_string())
257    }
258}
259
260fn base_metadata(reason: &str) -> BTreeMap<String, JsonValue> {
261    BTreeMap::from([("reason".to_string(), json!(reason))])
262}
263
264fn body_snippet(body: &str) -> String {
265    let mut snippet = String::new();
266    for ch in body.chars().take(BODY_SNIPPET_LIMIT) {
267        snippet.push(ch);
268    }
269    snippet
270}
271
272#[cfg(test)]
273mod tests {
274    use crate::http::framing::{http_content_length_from_header_lines, TEST_HTTP_MAX_BODY_BYTES};
275    use std::io::{Read, Write};
276    use std::net::TcpListener;
277    use std::sync::{Arc, Mutex};
278
279    use super::*;
280
281    fn provider_with_healthcheck(base_url: String, healthcheck: HealthcheckDef) -> ProviderDef {
282        ProviderDef {
283            base_url,
284            auth_style: "bearer".to_string(),
285            auth_env: AuthEnv::Single("HARN_TEST_PROVIDER_KEY".to_string()),
286            extra_headers: BTreeMap::from([("x-extra".to_string(), "extra-value".to_string())]),
287            chat_endpoint: "/chat/completions".to_string(),
288            healthcheck: Some(healthcheck),
289            ..Default::default()
290        }
291    }
292
293    fn install_provider(name: &str, provider: ProviderDef) {
294        let mut config = llm_config::ProvidersConfig::default();
295        config.providers.insert(name.to_string(), provider);
296        llm_config::set_user_overrides(Some(config));
297    }
298
299    fn spawn_healthcheck_stub(
300        status: u16,
301        body: &'static str,
302        captured: Arc<Mutex<Option<String>>>,
303    ) -> (String, std::thread::JoinHandle<()>) {
304        let listener = TcpListener::bind("127.0.0.1:0").expect("bind healthcheck stub");
305        let addr = listener.local_addr().expect("stub addr");
306
307        // Block on `accept()` directly. The earlier nonblocking poll +
308        // 30s wall-clock deadline introduced two failure modes under
309        // CI fan-out: (1) a 10ms tick stretched under load, delaying
310        // accept past the client's connect timeout; (2) the deadline
311        // wall-clock could elapse before the polling thread woke,
312        // panicking even though the client was already connected.
313        // Blocking accept is deterministic and cannot starve. The
314        // test always sends a request, so the accept always returns;
315        // if a test panics before sending, the leaked thread is
316        // captured by nextest's `leak-timeout = "fail"` rather than
317        // hidden behind a synthetic panic.
318        let handle = std::thread::spawn(move || {
319            let (mut stream, _) = listener
320                .accept()
321                .unwrap_or_else(|e| panic!("healthcheck stub accept failed: {e}"));
322            stream
323                .set_read_timeout(Some(std::time::Duration::from_secs(30)))
324                .ok();
325            stream
326                .set_write_timeout(Some(std::time::Duration::from_secs(30)))
327                .ok();
328
329            let mut bytes = Vec::new();
330            let mut buf = [0u8; 4096];
331            loop {
332                let n = stream.read(&mut buf).expect("read request");
333                if n == 0 {
334                    break;
335                }
336                bytes.extend_from_slice(&buf[..n]);
337                let request = String::from_utf8_lossy(&bytes);
338                if request_complete(&request) {
339                    break;
340                }
341            }
342            *captured.lock().expect("capture request") =
343                Some(String::from_utf8_lossy(&bytes).to_string());
344
345            let response = format!(
346                "HTTP/1.1 {status} OK\r\ncontent-type: application/json\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{body}",
347                body.len()
348            );
349            stream
350                .write_all(response.as_bytes())
351                .expect("write response");
352        });
353
354        (format!("http://{addr}"), handle)
355    }
356
357    fn request_complete(request: &str) -> bool {
358        let Some((headers, body)) = request.split_once("\r\n\r\n") else {
359            return false;
360        };
361        let content_length = match http_content_length_from_header_lines(
362            headers.lines(),
363            TEST_HTTP_MAX_BODY_BYTES,
364        ) {
365            Ok(content_length) => content_length,
366            Err(_) => return true,
367        };
368        body.len() >= content_length
369    }
370
371    #[test]
372    fn request_complete_stops_on_oversized_content_length() {
373        let request = format!(
374            "POST /probe HTTP/1.1\r\nContent-Length: {}\r\n\r\n",
375            TEST_HTTP_MAX_BODY_BYTES + 1
376        );
377        assert!(request_complete(&request));
378    }
379
380    #[tokio::test(flavor = "current_thread")]
381    #[allow(clippy::await_holding_lock)]
382    async fn sends_configured_probe_request_with_candidate_key() {
383        let _guard = crate::llm::env_guard();
384        let captured = Arc::new(Mutex::new(None));
385        let (base_url, server) = spawn_healthcheck_stub(200, r#"{"ok":true}"#, captured.clone());
386        install_provider(
387            "acme",
388            provider_with_healthcheck(
389                base_url.clone(),
390                HealthcheckDef {
391                    method: "POST".to_string(),
392                    path: Some("probe".to_string()),
393                    url: None,
394                    body: Some(r#"{"ping":true}"#.to_string()),
395                },
396            ),
397        );
398
399        let result = run_provider_healthcheck_with_options(
400            "acme",
401            ProviderHealthcheckOptions {
402                api_key: Some("candidate-key".to_string()),
403                client: None,
404            },
405        )
406        .await;
407        server.join().expect("stub server");
408        llm_config::clear_user_overrides();
409
410        assert!(result.valid);
411        assert_eq!(result.provider, "acme");
412        assert_eq!(result.metadata["status"], json!(200));
413        assert_eq!(result.metadata["method"], json!("POST"));
414        assert_eq!(result.metadata["url"], json!(format!("{base_url}/probe")));
415
416        let request = captured
417            .lock()
418            .expect("captured request")
419            .clone()
420            .expect("request");
421        assert!(request.starts_with("POST /probe HTTP/1.1\r\n"));
422        assert!(request.contains("authorization: Bearer candidate-key\r\n"));
423        assert!(request.contains("x-extra: extra-value\r\n"));
424        assert!(request.ends_with(r#"{"ping":true}"#));
425    }
426
427    #[tokio::test(flavor = "current_thread")]
428    #[allow(clippy::await_holding_lock)]
429    async fn reports_missing_credentials_without_network() {
430        let _guard = crate::llm::env_guard();
431        unsafe {
432            std::env::remove_var("HARN_TEST_PROVIDER_KEY");
433        }
434        install_provider(
435            "acme-missing-key",
436            provider_with_healthcheck(
437                "http://127.0.0.1:9".to_string(),
438                HealthcheckDef {
439                    method: "GET".to_string(),
440                    path: Some("/models".to_string()),
441                    url: None,
442                    body: None,
443                },
444            ),
445        );
446
447        let result = run_provider_healthcheck("acme-missing-key").await;
448        llm_config::clear_user_overrides();
449
450        assert!(!result.valid);
451        assert_eq!(result.metadata["reason"], json!("missing_credentials"));
452        assert_eq!(
453            result.metadata["auth_env"],
454            json!(["HARN_TEST_PROVIDER_KEY"])
455        );
456        assert!(result.message.contains("Missing credentials"));
457    }
458
459    #[tokio::test(flavor = "current_thread")]
460    #[allow(clippy::await_holding_lock)]
461    async fn returns_stable_failure_shape_for_http_errors() {
462        let _guard = crate::llm::env_guard();
463        let captured = Arc::new(Mutex::new(None));
464        let (base_url, server) = spawn_healthcheck_stub(401, r#"{"error":"bad key"}"#, captured);
465        install_provider(
466            "acme-auth",
467            provider_with_healthcheck(
468                base_url,
469                HealthcheckDef {
470                    method: "GET".to_string(),
471                    path: Some("/models".to_string()),
472                    url: None,
473                    body: None,
474                },
475            ),
476        );
477
478        let result = run_provider_healthcheck_with_options(
479            "acme-auth",
480            ProviderHealthcheckOptions {
481                api_key: Some("bad-key".to_string()),
482                client: None,
483            },
484        )
485        .await;
486        server.join().expect("stub server");
487        llm_config::clear_user_overrides();
488
489        assert!(!result.valid);
490        assert_eq!(result.provider, "acme-auth");
491        assert_eq!(result.metadata["reason"], json!("http_status"));
492        assert_eq!(result.metadata["status"], json!(401));
493        assert_eq!(result.metadata["body"], json!(r#"{"error":"bad key"}"#));
494    }
495
496    #[test]
497    fn default_external_provider_catalog_has_healthchecks() {
498        for provider in [
499            "openrouter",
500            "anthropic",
501            "openai",
502            "huggingface",
503            "together",
504        ] {
505            let config = llm_config::provider_config(provider)
506                .unwrap_or_else(|| panic!("missing provider {provider}"));
507            let healthcheck = config
508                .healthcheck
509                .as_ref()
510                .unwrap_or_else(|| panic!("missing healthcheck for {provider}"));
511            assert!(!healthcheck.method.is_empty());
512            assert!(healthcheck.path.is_some() || healthcheck.url.is_some());
513        }
514    }
515}