fix(codex): serve versioned dynamic model catalogs

This commit is contained in:
zhefox
2026-08-14 18:41:44 +08:00
parent a1d64e5239
commit 810c3dfe2b
39 changed files with 6147 additions and 362 deletions
@@ -1241,7 +1241,7 @@ async fn refresh_codex_provider_quota_locally_with_reset_fence(
continue;
}
let persisted_key = state
.read_provider_catalog_keys_by_ids(&[key.id.clone()])
.read_provider_catalog_keys_by_ids(std::slice::from_ref(&key.id))
.await?
.into_iter()
.next();
@@ -22,6 +22,8 @@ pub(crate) use self::system_modules_helpers::{
serialize_public_capability, supported_capability_names, PUBLIC_CAPABILITY_DEFINITIONS,
};
#[cfg(test)]
pub(crate) use self::support::filter_eligible_model_rows;
pub(crate) use self::support::{
build_api_key_install_session_response, build_proxy_node_install_session_response,
build_unhandled_public_support_response, matches_model_mapping_for_models,
@@ -52,6 +52,8 @@ mod support_user_me;
mod support_wallet;
pub(crate) use self::support_announcements::maybe_build_local_admin_announcements_response;
#[cfg(test)]
pub(crate) use self::support_models::filter_eligible_model_rows;
pub(crate) use self::support_models::matches_model_mapping_for_models;
use self::support_announcements::{
@@ -10,6 +10,8 @@ mod models_route;
mod models_shared;
pub(crate) use self::models_responses::build_models_auth_error_response;
#[cfg(test)]
pub(crate) use self::models_shared::filter_eligible_model_rows;
pub(crate) use self::models_shared::{matches_model_mapping_for_models, models_api_format};
pub(super) async fn maybe_build_local_models_response(
@@ -102,8 +102,15 @@ pub(super) fn build_empty_models_list_response(api_format: &str) -> Response<Bod
}
}
pub(super) fn build_codex_models_list_response(models: Vec<serde_json::Value>) -> Response<Body> {
Json(json!({ "models": models })).into_response()
pub(super) fn build_codex_models_list_response(
models: Vec<serde_json::Value>,
etag: Option<&str>,
) -> Response<Body> {
let mut response = Json(json!({ "models": models })).into_response();
if let Some(etag) = etag.and_then(|value| http::HeaderValue::from_str(value).ok()) {
response.headers_mut().insert(http::header::ETAG, etag);
}
response
}
pub(super) fn build_openai_models_list_response(
@@ -1,4 +1,4 @@
use std::collections::{BTreeMap, BTreeSet};
use std::collections::BTreeSet;
use std::fmt::Debug;
use std::future::Future;
use std::time::{Duration, SystemTime, UNIX_EPOCH};
@@ -25,8 +25,18 @@ use super::{query_param_value, AppState, GatewayPublicRequestContext};
#[cfg(not(test))]
const MODELS_ROUTE_READ_TIMEOUT: Duration = Duration::from_secs(5);
#[cfg(test)]
const MODELS_ROUTE_READ_TIMEOUT: Duration = Duration::from_millis(50);
const MODELS_ROUTE_READ_TIMEOUT: Duration = Duration::from_secs(1);
const CODEX_MODELS_QUERY_API_FORMATS: &[&str] = &["openai:responses"];
const CODEX_MODELS_MAX_RESPONSE_MODELS: usize = 512;
const CODEX_MODELS_MAX_RESPONSE_JSON_BYTES: usize = 8 * 1024 * 1024;
fn codex_projected_catalog_fits_response_limits(cards: &[Value]) -> bool {
if cards.len() > CODEX_MODELS_MAX_RESPONSE_MODELS {
return false;
}
serde_json::to_vec(&serde_json::json!({ "models": cards }))
.is_ok_and(|body| body.len() <= CODEX_MODELS_MAX_RESPONSE_JSON_BYTES)
}
async fn await_models_route_read<T, E, Fut>(operation: &'static str, future: Fut) -> Option<T>
where
@@ -115,124 +125,130 @@ fn is_codex_provider_row(row: &StoredMinimalCandidateSelectionRow) -> bool {
row.provider_type.trim().eq_ignore_ascii_case("codex")
}
fn codex_model_card_is_complete(card: &serde_json::Map<String, Value>) -> bool {
card.get("slug").and_then(Value::as_str).is_some()
&& card.get("display_name").and_then(Value::as_str).is_some()
&& card
.get("supported_reasoning_levels")
.and_then(Value::as_array)
.is_some()
&& card.get("shell_type").and_then(Value::as_str).is_some()
&& card.get("visibility").and_then(Value::as_str).is_some()
&& card
.get("supported_in_api")
.and_then(Value::as_bool)
.is_some()
&& card.get("priority").and_then(Value::as_i64).is_some()
&& card
.get("base_instructions")
.and_then(Value::as_str)
.is_some()
&& card
.get("supports_reasoning_summary_parameter")
.is_none_or(Value::is_boolean)
&& card
.get("support_verbosity")
.and_then(Value::as_bool)
.is_some()
&& card
.get("truncation_policy")
.and_then(Value::as_object)
.is_some()
&& card
.get("supports_parallel_tool_calls")
.and_then(Value::as_bool)
.is_some()
&& card
.get("experimental_supported_tools")
.and_then(Value::as_array)
.is_some()
}
fn project_codex_model_card(
cached_models: &[Value],
source_model: &str,
global_model: &str,
) -> Option<Value> {
let mut card = cached_models
.iter()
.find(|model| {
model.get("id").and_then(Value::as_str) == Some(source_model)
|| model.get("slug").and_then(Value::as_str) == Some(source_model)
})?
.as_object()?
.clone();
if !codex_model_card_is_complete(&card) {
return None;
}
card.remove("id");
card.remove("api_formats");
card.insert("slug".to_string(), Value::String(global_model.to_string()));
Some(Value::Object(card))
}
async fn load_codex_model_cards(
state: &AppState,
rows: &[StoredMinimalCandidateSelectionRow],
) -> Vec<Value> {
let cache_keys = rows
targets: &[crate::model_fetch::CodexCatalogTarget],
client_version: crate::model_fetch::NormalizedCodexClientVersion,
) -> (Vec<Value>, Option<String>) {
let catalogs = crate::model_fetch::load_codex_catalogs(state, targets, &client_version).await;
for target in catalogs.stale_targets() {
let state = state.clone();
let target = target.clone();
let client_version = client_version.clone();
tokio::spawn(async move {
crate::model_fetch::refresh_codex_catalog_target(&state, &target, &client_version)
.await;
});
}
if !catalogs.is_complete() {
warn!(
event_name = "codex_catalog_aggregate_incomplete",
client_version = %client_version.as_str(),
target_count = targets.len(),
"Codex catalog aggregation was incomplete; returning an empty remote catalog so the client can use its bundled fallback"
);
return (Vec::new(), None);
}
let mut seen_global_models = BTreeSet::new();
let possible_inference_catalogs = rows
.iter()
.filter(|row| is_codex_provider_row(row))
.map(|row| format!("upstream_models:{}:{}", row.provider_id, row.key_id))
.collect::<BTreeSet<_>>()
.into_iter()
.collect::<Vec<_>>();
let cached_values = await_models_route_read(
"codex_models_cache",
state.runtime_state.kv_get_many(&cache_keys),
)
.await
.unwrap_or_default();
let cached_models_by_key = cache_keys
.into_iter()
.zip(cached_values)
.filter_map(|(key, raw)| {
let models = serde_json::from_str::<Vec<Value>>(raw.as_deref()?).ok()?;
Some((key, models))
})
.collect::<BTreeMap<_, _>>();
let mut seen_global_models = BTreeSet::new();
.map(|row| (row.provider_id.clone(), row.key_id.clone()))
.collect::<BTreeSet<_>>();
let expected_global_models = rows
.iter()
.filter(|row| is_codex_provider_row(row))
.map(|row| row.global_model_name.clone())
.collect::<BTreeSet<_>>();
let mut cards = Vec::new();
for row in rows.iter().filter(|row| is_codex_provider_row(row)) {
if seen_global_models.contains(&row.global_model_name) {
continue;
}
let cache_key = format!("upstream_models:{}:{}", row.provider_id, row.key_id);
let Some(cached_models) = cached_models_by_key.get(&cache_key) else {
let Some(snapshot) = catalogs.snapshot(&row.provider_id, &row.key_id) else {
continue;
};
let source_model =
aether_scheduler_core::select_provider_model_name(row, "openai:responses");
let Some(card) = project_codex_model_card(
cached_models,
let Some(card) = crate::ai_serving::project_codex_catalog_model_card(
&snapshot.models,
source_model.as_str(),
row.global_model_name.as_str(),
) else {
warn!(
event_name = "codex_catalog_authorized_model_missing",
provider_id = %row.provider_id,
key_id = %row.key_id,
client_version = %client_version.as_str(),
source_model = %source_model,
global_model = %row.global_model_name,
"authorized Codex model was not present in this upstream catalog mapping"
);
continue;
};
seen_global_models.insert(row.global_model_name.clone());
cards.push(card);
if cards.len() > CODEX_MODELS_MAX_RESPONSE_MODELS {
warn!(
event_name = "codex_catalog_aggregate_model_limit",
client_version = %client_version.as_str(),
model_count = cards.len(),
limit = CODEX_MODELS_MAX_RESPONSE_MODELS,
"Codex projected catalog exceeded the aggregate model limit; returning an empty remote catalog"
);
return (Vec::new(), None);
}
}
cards
let missing_model_count = expected_global_models
.difference(&seen_global_models)
.count();
if missing_model_count > 0 {
warn!(
event_name = "codex_catalog_authorized_models_incomplete",
client_version = %client_version.as_str(),
expected_model_count = expected_global_models.len(),
projected_model_count = cards.len(),
missing_model_count,
"Codex upstream catalogs omitted authorized mappings; returning an empty remote catalog so the client can use its bundled fallback"
);
return (Vec::new(), None);
}
if !codex_projected_catalog_fits_response_limits(&cards) {
warn!(
event_name = "codex_catalog_aggregate_body_limit",
client_version = %client_version.as_str(),
model_count = cards.len(),
limit_bytes = CODEX_MODELS_MAX_RESPONSE_JSON_BYTES,
"Codex projected catalog exceeded the aggregate response body limit; returning an empty remote catalog"
);
return (Vec::new(), None);
}
if cards.is_empty() {
return (cards, None);
}
let etag = if possible_inference_catalogs.len() == 1 {
possible_inference_catalogs
.iter()
.next()
.and_then(|(provider_id, key_id)| catalogs.snapshot(provider_id, key_id))
.and_then(|snapshot| snapshot.etag.clone())
} else {
None
};
(cards, etag)
}
struct ModelRowsForClientFormat {
rows: Vec<StoredMinimalCandidateSelectionRow>,
codex_catalog_targets: Vec<crate::model_fetch::CodexCatalogTarget>,
}
async fn list_model_rows_for_client_format(
state: &AppState,
api_format: &str,
auth_snapshot: Option<&crate::data::auth::GatewayAuthApiKeySnapshot>,
) -> Option<Vec<StoredMinimalCandidateSelectionRow>> {
) -> Option<ModelRowsForClientFormat> {
let mut collected = Vec::new();
let query_api_formats = if is_codex_models_api_format(api_format) {
CODEX_MODELS_QUERY_API_FORMATS
@@ -254,9 +270,16 @@ async fn list_model_rows_for_client_format(
}
if is_codex_models_api_format(api_format) {
collected.retain(is_codex_provider_row);
Some(sort_model_rows(collected))
let codex_catalog_targets = crate::model_fetch::codex_catalog_targets(&collected);
Some(ModelRowsForClientFormat {
rows: sort_model_rows(collected),
codex_catalog_targets,
})
} else {
Some(sort_and_dedup_model_rows(collected))
Some(ModelRowsForClientFormat {
rows: sort_and_dedup_model_rows(collected),
codex_catalog_targets: Vec::new(),
})
}
}
@@ -296,6 +319,9 @@ pub(super) async fn maybe_build_local_models_route_response(
}
let auth_context = decision.auth_context.as_ref()?;
if !auth_context.access_allowed || auth_context.local_rejection.is_some() {
return Some(build_models_auth_error_response(api_format));
}
let now_unix_secs = SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap_or_default()
@@ -318,11 +344,23 @@ pub(super) async fn maybe_build_local_models_route_response(
))
}
};
let auth_snapshot = auth_snapshot.as_ref();
let Some(auth_snapshot) = auth_snapshot.as_ref() else {
warn!(
event_name = "models_route_auth_snapshot_missing",
user_id = %auth_context.user_id,
api_key_id = %auth_context.api_key_id,
"gateway models route rejected a request whose authenticated API key snapshot disappeared"
);
return Some(build_models_auth_error_response(api_format));
};
if !auth_snapshot.currently_usable {
return Some(build_models_auth_error_response(api_format));
}
let auth_snapshot = Some(auth_snapshot);
match decision.route_kind.as_deref() {
Some("list") => {
let rows =
let listed =
match list_model_rows_for_client_format(state, api_format, auth_snapshot).await {
Some(rows) => rows,
None => {
@@ -332,12 +370,34 @@ pub(super) async fn maybe_build_local_models_route_response(
))
}
};
let rows = listed.rows;
if rows.is_empty() {
return Some(build_empty_models_list_response(api_format));
}
if is_codex_models_api_format(api_format) {
let models = load_codex_model_cards(state, &rows).await;
return Some(build_codex_models_list_response(models));
let raw_client_version = query_param_value(
request_context.request_query_string.as_deref(),
"client_version",
);
let client_version = crate::model_fetch::normalize_codex_client_version(
raw_client_version.as_deref(),
);
if client_version.used_fallback() {
warn!(
event_name = "codex_catalog_invalid_client_version",
raw_length = raw_client_version.as_ref().map_or(0, String::len),
fallback_version = %client_version.as_str(),
"invalid Codex client_version used the bounded fallback version"
);
}
let (models, etag) = load_codex_model_cards(
state,
&rows,
&listed.codex_catalog_targets,
client_version,
)
.await;
return Some(build_codex_models_list_response(models, etag.as_deref()));
}
let response = match api_format {
"claude:messages" => {
@@ -410,3 +470,32 @@ pub(super) async fn maybe_build_local_models_route_response(
_ => Some(build_models_auth_error_response(api_format)),
}
}
#[cfg(test)]
mod tests {
use serde_json::json;
use super::{
codex_projected_catalog_fits_response_limits, CODEX_MODELS_MAX_RESPONSE_JSON_BYTES,
CODEX_MODELS_MAX_RESPONSE_MODELS,
};
#[test]
fn projected_codex_catalog_enforces_aggregate_count_and_body_limits() {
assert!(codex_projected_catalog_fits_response_limits(&[json!({
"slug": "gpt-future-dynamic",
"model_messages": {"instructions_template": "opaque"}
})]));
let too_many = (0..=CODEX_MODELS_MAX_RESPONSE_MODELS)
.map(|index| json!({"slug": format!("gpt-future-{index}")}))
.collect::<Vec<_>>();
assert!(!codex_projected_catalog_fits_response_limits(&too_many));
let oversized = vec![json!({
"slug": "gpt-future-oversized",
"future_capability": "x".repeat(CODEX_MODELS_MAX_RESPONSE_JSON_BYTES)
})];
assert!(!codex_projected_catalog_fits_response_limits(&oversized));
}
}
@@ -194,7 +194,7 @@ fn row_exposes_global_model_for_models(
false
}
pub(super) fn filter_eligible_model_rows(
pub(crate) fn filter_eligible_model_rows(
rows: Vec<StoredMinimalCandidateSelectionRow>,
auth_snapshot: Option<&crate::data::auth::GatewayAuthApiKeySnapshot>,
api_format: &str,