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

harn_vm/llm/
healthcheck.rs
