Skip to main content

agent_engine/runtime/openai/
ping.rs

1//! Model ping / health check.
2//!
3//! Sends a minimal chat completion (`max_tokens: 1`, message `"hi"`) to each
4//! configured model in parallel and classifies the response.
5
6use std::time::{Duration, Instant};
7
8use serde_json::json;
9
10use super::registry;
11use super::types::ProviderConfig;
12use agent_core::auth::broker;
13
14#[derive(Debug, Clone, Copy, PartialEq, Eq)]
15pub enum PingStatus {
16    Online,
17    RateLimited,
18    Unauthorized,
19    NotFound,
20    Error,
21    Timeout,
22}
23
24impl PingStatus {
25    pub fn icon(&self) -> &'static str {
26        match self {
27            PingStatus::Online => "✅",
28            PingStatus::RateLimited => "⏳",
29            PingStatus::Unauthorized => "🔒",
30            PingStatus::NotFound => "❌",
31            PingStatus::Error => "⚠️",
32            PingStatus::Timeout => "⌛",
33        }
34    }
35
36    pub fn label(&self) -> &'static str {
37        match self {
38            PingStatus::Online => "online",
39            PingStatus::RateLimited => "429 rate limited",
40            PingStatus::Unauthorized => "401 unauthorized",
41            PingStatus::NotFound => "404 not found",
42            PingStatus::Error => "error",
43            PingStatus::Timeout => "timeout",
44        }
45    }
46}
47
48#[derive(Debug, Clone)]
49pub struct PingResult {
50    pub provider_key: String,
51    pub model_id: String,
52    pub status: PingStatus,
53    pub latency_ms: u64,
54}
55
56const TIMEOUT: Duration = Duration::from_secs(10);
57
58pub async fn ping_model(cfg: &ProviderConfig, provider_key: &str) -> PingResult {
59    let body = json!({
60        "model": cfg.model,
61        "messages": [{"role": "user", "content": "hi"}],
62        "max_tokens": 1,
63        "stream": false,
64    });
65
66    let start = Instant::now();
67    // Credential applied broker-side; the ping path never sees a key.
68    let broker_handle = broker::global_broker();
69    let fut = broker_handle.proxy(broker::ProxyRequest {
70        provider: cfg.provider.clone(),
71        method: broker::ProxyMethod::Post,
72        path: "/chat/completions".to_string(),
73        body: Some(body),
74        stream: false,
75        body_bytes: None,
76    });
77
78    let status = match tokio::time::timeout(TIMEOUT, fut).await {
79        Err(_) => PingStatus::Timeout,
80        Ok(Err(_)) => PingStatus::Error,
81        Ok(Ok(resp)) => match resp.status {
82            200..=299 => PingStatus::Online,
83            401 | 403 => PingStatus::Unauthorized,
84            404 => PingStatus::NotFound,
85            429 => PingStatus::RateLimited,
86            _ => PingStatus::Error,
87        },
88    };
89
90    PingResult {
91        provider_key: provider_key.to_string(),
92        model_id: cfg.model.clone(),
93        status,
94        latency_ms: start.elapsed().as_millis() as u64,
95    }
96}
97
98/// Ping every model of every configured provider in parallel.
99/// Results are sent through `tx` as they arrive (not batched).
100pub async fn ping_all_configured(
101    tx: tokio::sync::mpsc::UnboundedSender<(String, PingStatus, u64)>,
102) {
103    let specs = registry::providers();
104    let mut handles = Vec::new();
105
106    for spec in specs {
107        let Some(base_cfg) = registry::resolve_provider_model(spec.key, spec.default_model) else {
108            continue;
109        };
110        for (model_id, _label, _tier) in spec.models {
111            let cfg = ProviderConfig {
112                base_url: base_cfg.base_url.clone(),
113                model: (*model_id).to_string(),
114                provider: base_cfg.provider.clone(),
115            };
116            let key = spec.key.to_string();
117            let tx = tx.clone();
118            handles.push(tokio::spawn(async move {
119                let result = ping_model(&cfg, &key).await;
120                let full_key = format!("{}/{}", result.provider_key, result.model_id);
121                let _ = tx.send((full_key, result.status, result.latency_ms));
122            }));
123        }
124    }
125
126    for h in handles {
127        let _ = h.await;
128    }
129    // tx drops here — receiver sees None and knows all pings are done
130}