Skip to main content

agentic_server/handler/http/
models.rs

1use std::future::Future;
2use std::sync::OnceLock;
3
4use axum::extract::{Query, State};
5use axum::http::HeaderMap;
6use axum::response::{IntoResponse, Response};
7use http::StatusCode;
8use serde_json::{Value, json};
9use tracing::{debug, info, warn};
10
11use agentic_core::proxy::{ProxyBody, ProxyResponse, error_response, proxy_get};
12use agentic_core::readiness::{LLM_READINESS_PROBE_TIMEOUT, LlmReadiness, probe_llm_readiness};
13
14use super::super::common::convert_response;
15use crate::app::AppState;
16
17/// Static fields shared by every Codex `ModelInfo` entry.
18///
19/// Built once on first use; cloned per model and patched with the per-model
20/// values (`slug`, `display_name`, `auto_review_model_override`,
21/// `supports_reasoning_summaries`, `input_modalities`, and optionally
22/// `context_window` / `max_context_window`).
23fn codex_model_template() -> &'static Value {
24    static TEMPLATE: OnceLock<Value> = OnceLock::new();
25    TEMPLATE.get_or_init(|| {
26        json!({
27            "supported_in_api": true,
28            "priority": 1,
29            "shell_type": "shell_command",
30            "visibility": "list",
31            "base_instructions": "",
32            "supported_reasoning_levels": [
33                {"effort": "low",    "description": "Fast responses with lighter reasoning"},
34                {"effort": "medium", "description": "Balances speed and reasoning depth"},
35                {"effort": "high",   "description": "Greater reasoning depth for complex problems"}
36            ],
37            "default_reasoning_summary": "auto",
38            "support_verbosity": false,
39            "default_verbosity": null,
40            "apply_patch_tool_type": "freeform",
41            "web_search_tool_type": "text",
42            "truncation_policy": {"mode": "bytes", "limit": 100_000},
43            "supports_parallel_tool_calls": true,
44            "supports_image_detail_original": false,
45            "effective_context_window_percent": 95,
46            "experimental_supported_tools": [],
47            "supports_search_tool": false,
48            "use_responses_lite": false,
49            "tool_mode": null,
50            "multi_agent_version": null,
51        })
52    })
53}
54
55/// Transform a single upstream model entry into a Codex `ModelInfo` object.
56///
57/// Returns `None` when the entry has no `id` field (malformed upstream data).
58fn upstream_model_to_codex(m: &Value) -> Option<Value> {
59    let id = m["id"].as_str()?.to_owned();
60    let display_name = m.get("name").and_then(Value::as_str).unwrap_or(&id).to_owned();
61    // vLLM uses max_model_len; other providers may use context_length
62    let context_length = m["max_model_len"].as_i64().or_else(|| m["context_length"].as_i64());
63    // Single pass over capabilities for both flags
64    let (supports_reasoning, supports_image) = m["capabilities"].as_array().map_or((false, false), |c| {
65        c.iter().fold((false, false), |(r, i), v| {
66            let s = v.as_str();
67            (r || s == Some("reasoning"), i || s == Some("image"))
68        })
69    });
70    let input_modalities = if supports_image {
71        json!(["text", "image"])
72    } else {
73        json!(["text"])
74    };
75
76    let mut model = codex_model_template().clone();
77    let obj = model.as_object_mut().expect("template is object");
78    obj.insert("slug".into(), json!(id));
79    obj.insert("display_name".into(), json!(display_name));
80    obj.insert("auto_review_model_override".into(), json!(id));
81    obj.insert("supports_reasoning_summaries".into(), json!(supports_reasoning));
82    obj.insert("input_modalities".into(), input_modalities);
83    if let Some(ctx) = context_length {
84        obj.insert("context_window".into(), json!(ctx));
85        obj.insert("max_context_window".into(), json!(ctx));
86    }
87
88    Some(model)
89}
90
91/// Build the Codex `ModelsResponse` from a raw upstream vLLM models payload.
92fn build_codex_models_response(upstream_bytes: &[u8]) -> Value {
93    let models: Vec<Value> = serde_json::from_slice::<Value>(upstream_bytes)
94        .ok()
95        .and_then(|mut v| match v["data"].take() {
96            Value::Array(arr) => Some(arr),
97            _ => None,
98        })
99        .into_iter()
100        .flatten()
101        .filter_map(|m| upstream_model_to_codex(&m))
102        .collect();
103    json!({ "models": models })
104}
105
106pub async fn health() -> impl IntoResponse {
107    StatusCode::OK
108}
109
110async fn upstream_is_ready(state: &AppState) -> bool {
111    match probe_llm_readiness(
112        &state.llm_readiness_client,
113        &state.llm_api_base,
114        state.openai_api_key.as_deref(),
115        LLM_READINESS_PROBE_TIMEOUT,
116    )
117    .await
118    {
119        Ok(LlmReadiness::Ready) => true,
120        Ok(LlmReadiness::Rejected(status)) => {
121            debug!(http.status = %status, "LLM backend not ready");
122            false
123        }
124        Ok(LlmReadiness::Unreachable(error)) => {
125            debug!(error = ?error, "LLM backend unreachable");
126            false
127        }
128        Ok(LlmReadiness::TimedOut) => {
129            debug!("LLM backend readiness check timed out");
130            false
131        }
132        Ok(_) => {
133            debug!("LLM backend returned an unsupported readiness state");
134            false
135        }
136        Err(error) => {
137            debug!(error = ?error, "LLM backend readiness configuration invalid");
138            false
139        }
140    }
141}
142
143async fn configured_upstream_is_ready(state: &AppState) -> bool {
144    state.skip_llm_ready_check || upstream_is_ready(state).await
145}
146
147async fn dependencies_are_ready(
148    storage_ready: impl Future<Output = bool>,
149    upstream_ready: impl Future<Output = bool>,
150) -> bool {
151    tokio::pin!(storage_ready, upstream_ready);
152
153    tokio::select! {
154        storage_ready = &mut storage_ready => {
155            if storage_ready {
156                upstream_ready.await
157            } else {
158                debug!("database persistence not ready");
159                false
160            }
161        }
162        upstream_ready = &mut upstream_ready => {
163            if upstream_ready {
164                let storage_ready = storage_ready.await;
165                if !storage_ready {
166                    debug!("database persistence not ready");
167                }
168                storage_ready
169            } else {
170                false
171            }
172        }
173    }
174}
175
176pub async fn ready(State(state): State<AppState>) -> impl IntoResponse {
177    let Some(probe) = state.readiness_tracker.try_start_probe() else {
178        let cached_ready = state.readiness_tracker.last_result().unwrap_or(false);
179        debug!(
180            readiness.ready = cached_ready,
181            "returning cached readiness while dependency probe is in progress"
182        );
183        return if cached_ready {
184            StatusCode::OK
185        } else {
186            StatusCode::SERVICE_UNAVAILABLE
187        };
188    };
189    let dependencies_ready = dependencies_are_ready(
190        state.exec_ctx.storage_ready(std::time::Duration::from_secs(1)),
191        configured_upstream_is_ready(&state),
192    )
193    .await;
194
195    if probe.finish(dependencies_ready) {
196        if dependencies_ready {
197            info!(readiness.ready = true, "gateway dependencies ready");
198        } else {
199            warn!(readiness.ready = false, "gateway dependencies not ready");
200        }
201    }
202
203    if dependencies_ready {
204        StatusCode::OK
205    } else {
206        StatusCode::SERVICE_UNAVAILABLE
207    }
208}
209
210/// Query parameters for GET /v1/models.
211///
212/// Codex CLI appends `?client_version=<ver>` to identify itself; its presence
213/// triggers transformation to the Codex `ModelsResponse` shape.
214#[derive(serde::Deserialize)]
215pub struct ModelsParams {
216    client_version: Option<String>,
217}
218
219/// GET /v1/models — Codex-aware model list.
220///
221/// When `?client_version` is present (Codex CLI), fetches vLLM's model list via
222/// [`proxy_get`] and transforms it into the Codex `ModelsResponse` shape
223/// (`{ "models": [...] }` with rich metadata). Without `client_version`, the
224/// upstream response is returned unchanged via [`proxy_get`].
225pub async fn models(State(state): State<AppState>, headers: HeaderMap, Query(params): Query<ModelsParams>) -> Response {
226    let upstream = proxy_get("/v1/models", &headers, &state.proxy_state).await;
227
228    if params.client_version.is_none() {
229        return convert_response(upstream);
230    }
231
232    let ProxyBody::Full(upstream_bytes) = upstream.body else {
233        return convert_response(error_response(
234            http::StatusCode::BAD_GATEWAY,
235            "upstream_unavailable",
236            "unexpected streaming response from /v1/models",
237        ));
238    };
239
240    if !upstream.status.is_success() {
241        return convert_response(ProxyResponse {
242            body: ProxyBody::Full(upstream_bytes),
243            ..upstream
244        });
245    }
246
247    axum::Json(build_codex_models_response(&upstream_bytes)).into_response()
248}
249
250#[cfg(test)]
251mod tests {
252    use std::future;
253    use std::time::Duration;
254
255    use super::dependencies_are_ready;
256    use crate::app::ReadinessTracker;
257
258    #[test]
259    fn readiness_tracker_reports_only_state_transitions() {
260        let tracker = ReadinessTracker::default();
261
262        assert_eq!(tracker.last_result(), None);
263        assert!(tracker.try_start_probe().unwrap().finish(false));
264        assert_eq!(tracker.last_result(), Some(false));
265        assert!(!tracker.try_start_probe().unwrap().finish(false));
266        assert!(tracker.try_start_probe().unwrap().finish(true));
267        assert_eq!(tracker.last_result(), Some(true));
268        assert!(!tracker.try_start_probe().unwrap().finish(true));
269        assert!(tracker.try_start_probe().unwrap().finish(false));
270        assert_eq!(tracker.last_result(), Some(false));
271    }
272
273    #[test]
274    fn unfinished_readiness_probe_releases_permit_without_changing_result() {
275        let tracker = ReadinessTracker::default();
276        assert!(tracker.try_start_probe().unwrap().finish(true));
277
278        let unfinished = tracker.try_start_probe().unwrap();
279        drop(unfinished);
280
281        assert_eq!(tracker.last_result(), Some(true));
282        assert!(tracker.try_start_probe().is_some());
283    }
284
285    #[tokio::test(start_paused = true)]
286    async fn dependency_check_fails_fast_when_upstream_is_unready() {
287        let result = tokio::time::timeout(
288            Duration::from_secs(1),
289            dependencies_are_ready(future::pending(), future::ready(false)),
290        )
291        .await
292        .expect("upstream failure must win before the pending storage check");
293
294        assert!(!result);
295    }
296}