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(provider, &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 base_url = crate::llm_config::resolve_base_url(&def);
103    let client = match options.client {
104        Some(client) => client,
105        None => match crate::egress::install_ssrf_guard_with_private_host_allowlist(
106            reqwest::Client::builder()
107                .timeout(Duration::from_secs(DEFAULT_HEALTHCHECK_TIMEOUT_SECS))
108                .redirect(crate::egress::redirect_policy(
109                    "llm_healthcheck_redirect",
110                    5,
111                )),
112            &crate::egress::configured_provider_private_allow_host(&base_url)
113                .into_iter()
114                .collect::<Vec<_>>(),
115        )
116        .build()
117        {
118            Ok(client) => client,
119            Err(error) => {
120                let mut metadata = base_metadata("client_build_failed");
121                metadata.insert("provider".to_string(), json!(provider));
122                return ProviderHealthcheckResult::new(
123                    provider,
124                    false,
125                    format!("{provider} healthcheck failed: {error}"),
126                    metadata,
127                );
128            }
129        },
130    };
131
132    let mut request = client.request(method.clone(), &url);
133    if let Some(api_key) = auth.api_key.as_deref() {
134        request = apply_auth_headers(request, api_key, Some(&def));
135    }
136    for (name, value) in &def.extra_headers {
137        request = request.header(name, value);
138    }
139    if let Some(body) = &healthcheck.body {
140        request = request
141            .header(reqwest::header::CONTENT_TYPE, "application/json")
142            .body(body.clone());
143    }
144
145    match request.send().await {
146        Ok(response) => {
147            let status = response.status();
148            let status_code = status.as_u16();
149            let valid = status.is_success();
150            let body_text = response.text().await.unwrap_or_default();
151            let mut metadata = base_metadata(if valid { "ok" } else { "http_status" });
152            metadata.insert("provider".to_string(), json!(provider));
153            metadata.insert("status".to_string(), json!(status_code));
154            metadata.insert("url".to_string(), json!(url));
155            metadata.insert("method".to_string(), json!(method.as_str()));
156            if !valid && !body_text.is_empty() {
157                metadata.insert("body".to_string(), json!(body_snippet(&body_text)));
158            }
159
160            let message = if valid {
161                format!("{provider} is reachable (HTTP {status_code})")
162            } else {
163                let suffix = body_snippet(&body_text);
164                if suffix.is_empty() {
165                    format!("{provider} returned HTTP {status_code}")
166                } else {
167                    format!("{provider} returned HTTP {status_code}: {suffix}")
168                }
169            };
170
171            ProviderHealthcheckResult::new(provider, valid, message, metadata)
172        }
173        Err(error) => {
174            let mut metadata = base_metadata("request_failed");
175            metadata.insert("provider".to_string(), json!(provider));
176            metadata.insert(
177                "url".to_string(),
178                json!(crate::egress::redact_diagnostic_text(&url)),
179            );
180            metadata.insert("method".to_string(), json!(method.as_str()));
181            let error = crate::egress::redact_reqwest_error(&error);
182            ProviderHealthcheckResult::new(
183                provider,
184                false,
185                format!("{provider} healthcheck failed: {error}"),
186                metadata,
187            )
188        }
189    }
190}
191
192pub fn build_healthcheck_url(def: &ProviderDef, healthcheck: &HealthcheckDef) -> String {
193    if let Some(url) = &healthcheck.url {
194        return url.clone();
195    }
196
197    let base = llm_config::resolve_base_url(def);
198    let path = healthcheck.path.as_deref().unwrap_or("");
199    join_base_and_path(&base, path)
200}
201
202pub(crate) fn join_base_and_path(base: &str, path: &str) -> String {
203    if path.starts_with('/') {
204        format!("{}{}", base.trim_end_matches('/'), path)
205    } else if path.is_empty() {
206        base.to_string()
207    } else {
208        format!("{}/{}", base.trim_end_matches('/'), path)
209    }
210}
211
212#[derive(Debug, Clone)]
213struct ResolvedHealthcheckAuth {
214    requires_auth: bool,
215    api_key: Option<String>,
216    candidates: Vec<String>,
217}
218
219fn resolve_healthcheck_auth(
220    provider: &str,
221    def: &ProviderDef,
222    api_key_override: Option<String>,
223) -> ResolvedHealthcheckAuth {
224    let candidates = auth_env_candidates(&def.auth_env);
225    if let Some(api_key) = api_key_override.and_then(non_empty) {
226        return ResolvedHealthcheckAuth {
227            requires_auth: true,
228            api_key: Some(api_key),
229            candidates,
230        };
231    }
232
233    let auth = super::provider_auth::resolve_provider_auth(provider);
234    let requires_auth = matches!(
235        auth.status.credential_status,
236        super::ProviderCredentialStatus::Ok | super::ProviderCredentialStatus::Missing
237    );
238    ResolvedHealthcheckAuth {
239        requires_auth,
240        api_key: auth.into_api_key(),
241        candidates,
242    }
243}
244
245fn auth_env_candidates(auth_env: &AuthEnv) -> Vec<String> {
246    match auth_env {
247        AuthEnv::None => Vec::new(),
248        AuthEnv::Single(env) => vec![env.clone()],
249        AuthEnv::Multiple(envs) => envs.clone(),
250    }
251}
252
253fn non_empty(value: String) -> Option<String> {
254    let trimmed = value.trim();
255    if trimmed.is_empty() {
256        None
257    } else {
258        Some(trimmed.to_string())
259    }
260}
261
262fn base_metadata(reason: &str) -> BTreeMap<String, JsonValue> {
263    BTreeMap::from([("reason".to_string(), json!(reason))])
264}
265
266fn body_snippet(body: &str) -> String {
267    let mut snippet = String::new();
268    for ch in body.chars().take(BODY_SNIPPET_LIMIT) {
269        snippet.push(ch);
270    }
271    snippet
272}
273
274#[cfg(test)]
275mod tests {
276    use crate::http::framing::{http_content_length_from_header_lines, TEST_HTTP_MAX_BODY_BYTES};
277    use std::io::{Read, Write};
278    use std::net::TcpListener;
279    use std::sync::{Arc, Mutex};
280
281    use super::*;
282
283    fn provider_with_healthcheck(base_url: String, healthcheck: HealthcheckDef) -> ProviderDef {
284        ProviderDef {
285            base_url,
286            auth_style: "bearer".to_string(),
287            auth_env: AuthEnv::Single("HARN_TEST_PROVIDER_KEY".to_string()),
288            extra_headers: BTreeMap::from([("x-extra".to_string(), "extra-value".to_string())]),
289            chat_endpoint: "/chat/completions".to_string(),
290            healthcheck: Some(healthcheck),
291            ..Default::default()
292        }
293    }
294
295    fn install_provider(name: &str, provider: ProviderDef) {
296        let mut config = llm_config::ProvidersConfig::default();
297        config.providers.insert(name.to_string(), provider);
298        llm_config::set_user_overrides(Some(config));
299    }
300
301    struct ScopedEnv {
302        name: &'static str,
303        previous: Option<String>,
304    }
305
306    impl ScopedEnv {
307        fn set(name: &'static str, value: &str) -> Self {
308            let previous = std::env::var(name).ok();
309            unsafe { std::env::set_var(name, value) };
310            Self { name, previous }
311        }
312    }
313
314    impl Drop for ScopedEnv {
315        fn drop(&mut self) {
316            unsafe {
317                match &self.previous {
318                    Some(value) => std::env::set_var(self.name, value),
319                    None => std::env::remove_var(self.name),
320                }
321            }
322        }
323    }
324
325    fn spawn_healthcheck_stub(
326        status: u16,
327        body: &'static str,
328        captured: Arc<Mutex<Option<String>>>,
329    ) -> (String, std::thread::JoinHandle<()>) {
330        let listener = TcpListener::bind("127.0.0.1:0").expect("bind healthcheck stub");
331        let addr = listener.local_addr().expect("stub addr");
332
333        // Block on `accept()` directly. The earlier nonblocking poll +
334        // 30s wall-clock deadline introduced two failure modes under
335        // CI fan-out: (1) a 10ms tick stretched under load, delaying
336        // accept past the client's connect timeout; (2) the deadline
337        // wall-clock could elapse before the polling thread woke,
338        // panicking even though the client was already connected.
339        // Blocking accept is deterministic and cannot starve. The
340        // test always sends a request, so the accept always returns;
341        // if a test panics before sending, the leaked thread is
342        // captured by nextest's `leak-timeout = "fail"` rather than
343        // hidden behind a synthetic panic.
344        let handle = std::thread::spawn(move || {
345            let (mut stream, _) = listener
346                .accept()
347                .unwrap_or_else(|e| panic!("healthcheck stub accept failed: {e}"));
348            stream
349                .set_read_timeout(Some(std::time::Duration::from_secs(30)))
350                .ok();
351            stream
352                .set_write_timeout(Some(std::time::Duration::from_secs(30)))
353                .ok();
354
355            let mut bytes = Vec::new();
356            let mut buf = [0u8; 4096];
357            loop {
358                let n = stream.read(&mut buf).expect("read request");
359                if n == 0 {
360                    break;
361                }
362                bytes.extend_from_slice(&buf[..n]);
363                let request = String::from_utf8_lossy(&bytes);
364                if request_complete(&request) {
365                    break;
366                }
367            }
368            *captured.lock().expect("capture request") =
369                Some(String::from_utf8_lossy(&bytes).to_string());
370
371            let response = format!(
372                "HTTP/1.1 {status} OK\r\ncontent-type: application/json\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{body}",
373                body.len()
374            );
375            stream
376                .write_all(response.as_bytes())
377                .expect("write response");
378        });
379
380        (format!("http://{addr}"), handle)
381    }
382
383    fn request_complete(request: &str) -> bool {
384        let Some((headers, body)) = request.split_once("\r\n\r\n") else {
385            return false;
386        };
387        let content_length = match http_content_length_from_header_lines(
388            headers.lines(),
389            TEST_HTTP_MAX_BODY_BYTES,
390        ) {
391            Ok(content_length) => content_length,
392            Err(_) => return true,
393        };
394        body.len() >= content_length
395    }
396
397    #[test]
398    fn request_complete_stops_on_oversized_content_length() {
399        let request = format!(
400            "POST /probe HTTP/1.1\r\nContent-Length: {}\r\n\r\n",
401            TEST_HTTP_MAX_BODY_BYTES + 1
402        );
403        assert!(request_complete(&request));
404    }
405
406    #[tokio::test(flavor = "current_thread")]
407    #[allow(clippy::await_holding_lock)]
408    async fn sends_configured_probe_request_with_candidate_key() {
409        let _guard = crate::llm::env_guard();
410        let captured = Arc::new(Mutex::new(None));
411        let (base_url, server) = spawn_healthcheck_stub(200, r#"{"ok":true}"#, captured.clone());
412        install_provider(
413            "acme",
414            provider_with_healthcheck(
415                base_url.clone(),
416                HealthcheckDef {
417                    method: "POST".to_string(),
418                    path: Some("probe".to_string()),
419                    url: None,
420                    body: Some(r#"{"ping":true}"#.to_string()),
421                },
422            ),
423        );
424
425        let result = run_provider_healthcheck_with_options(
426            "acme",
427            ProviderHealthcheckOptions {
428                api_key: Some("candidate-key".to_string()),
429                client: None,
430            },
431        )
432        .await;
433        server.join().expect("stub server");
434        llm_config::clear_user_overrides();
435
436        assert!(result.valid);
437        assert_eq!(result.provider, "acme");
438        assert_eq!(result.metadata["status"], json!(200));
439        assert_eq!(result.metadata["method"], json!("POST"));
440        assert_eq!(result.metadata["url"], json!(format!("{base_url}/probe")));
441
442        let request = captured
443            .lock()
444            .expect("captured request")
445            .clone()
446            .expect("request");
447        assert!(request.starts_with("POST /probe HTTP/1.1\r\n"));
448        assert!(request.contains("authorization: Bearer candidate-key\r\n"));
449        assert!(request.contains("x-extra: extra-value\r\n"));
450        assert!(request.ends_with(r#"{"ping":true}"#));
451    }
452
453    #[tokio::test(flavor = "current_thread")]
454    #[allow(clippy::await_holding_lock)]
455    async fn reports_missing_credentials_without_network() {
456        let _guard = crate::llm::env_guard();
457        unsafe {
458            std::env::remove_var("HARN_TEST_PROVIDER_KEY");
459        }
460        install_provider(
461            "acme-missing-key",
462            provider_with_healthcheck(
463                "http://127.0.0.1:9".to_string(),
464                HealthcheckDef {
465                    method: "GET".to_string(),
466                    path: Some("/models".to_string()),
467                    url: None,
468                    body: None,
469                },
470            ),
471        );
472
473        let result = run_provider_healthcheck("acme-missing-key").await;
474        llm_config::clear_user_overrides();
475
476        assert!(!result.valid);
477        assert_eq!(result.metadata["reason"], json!("missing_credentials"));
478        assert_eq!(
479            result.metadata["auth_env"],
480            json!(["HARN_TEST_PROVIDER_KEY"])
481        );
482        assert!(result.message.contains("Missing credentials"));
483    }
484
485    #[test]
486    fn healthcheck_auth_uses_dispatch_secret_and_platform_resolution() {
487        let _guard = crate::llm::env_guard();
488        let _providers = ScopedEnv::set("HARN_SECRET_PROVIDERS", "env");
489        let _reference = ScopedEnv::set(
490            "HARN_TEST_PROVIDER_KEY",
491            "harn-secret://provider/healthcheck-key",
492        );
493        let _secret = ScopedEnv::set("HARN_SECRET_PROVIDER_HEALTHCHECK_KEY", "resolved-key");
494
495        let definition = provider_with_healthcheck(
496            "https://example.invalid".to_string(),
497            HealthcheckDef {
498                method: "GET".to_string(),
499                path: Some("/health".to_string()),
500                url: None,
501                body: None,
502            },
503        );
504        install_provider("acme-secret", definition.clone());
505        let secret_auth = resolve_healthcheck_auth("acme-secret", &definition, None);
506        assert!(secret_auth.requires_auth);
507        assert_eq!(secret_auth.api_key.as_deref(), Some("resolved-key"));
508
509        let managed = ProviderDef {
510            credential_resolution: "platform_managed".to_string(),
511            ..definition
512        };
513        install_provider("acme-managed", managed.clone());
514        let managed_auth = resolve_healthcheck_auth("acme-managed", &managed, None);
515        assert!(!managed_auth.requires_auth);
516        assert!(managed_auth.api_key.is_none());
517
518        llm_config::clear_user_overrides();
519    }
520
521    #[tokio::test(flavor = "current_thread")]
522    #[allow(clippy::await_holding_lock)]
523    async fn returns_stable_failure_shape_for_http_errors() {
524        let _guard = crate::llm::env_guard();
525        let captured = Arc::new(Mutex::new(None));
526        let (base_url, server) = spawn_healthcheck_stub(401, r#"{"error":"bad key"}"#, captured);
527        install_provider(
528            "acme-auth",
529            provider_with_healthcheck(
530                base_url,
531                HealthcheckDef {
532                    method: "GET".to_string(),
533                    path: Some("/models".to_string()),
534                    url: None,
535                    body: None,
536                },
537            ),
538        );
539
540        let result = run_provider_healthcheck_with_options(
541            "acme-auth",
542            ProviderHealthcheckOptions {
543                api_key: Some("bad-key".to_string()),
544                client: None,
545            },
546        )
547        .await;
548        server.join().expect("stub server");
549        llm_config::clear_user_overrides();
550
551        assert!(!result.valid);
552        assert_eq!(result.provider, "acme-auth");
553        assert_eq!(result.metadata["reason"], json!("http_status"));
554        assert_eq!(result.metadata["status"], json!(401));
555        assert_eq!(result.metadata["body"], json!(r#"{"error":"bad key"}"#));
556    }
557
558    #[test]
559    fn default_external_provider_catalog_has_healthchecks() {
560        for provider in [
561            "openrouter",
562            "anthropic",
563            "openai",
564            "huggingface",
565            "together",
566        ] {
567            let config = llm_config::provider_config(provider)
568                .unwrap_or_else(|| panic!("missing provider {provider}"));
569            let healthcheck = config
570                .healthcheck
571                .as_ref()
572                .unwrap_or_else(|| panic!("missing healthcheck for {provider}"));
573            assert!(!healthcheck.method.is_empty());
574            assert!(healthcheck.path.is_some() || healthcheck.url.is_some());
575        }
576    }
577}