Improve Codex model fetching

This commit is contained in:
fawney19
2026-05-04 01:16:56 +08:00
parent fb642f7d62
commit 1d62722d47
5 changed files with 363 additions and 48 deletions

View File

@@ -58,7 +58,7 @@ pub fn extract_error_message(value: &Value) -> Option<String> {
}
pub fn build_models_fetch_url(
_provider_type: &str,
provider_type: &str,
endpoint_api_format: &str,
base_url: &str,
) -> Option<(String, String)> {
@@ -66,7 +66,10 @@ pub fn build_models_fetch_url(
if !endpoint_supports_rust_models_fetch(&api_format) {
return None;
}
let url = if api_format.starts_with("openai:") || api_format.starts_with("claude:") {
let provider_type = provider_type.trim().to_ascii_lowercase();
let url = if provider_type == "codex" && api_format.starts_with("openai:") {
build_codex_models_url(base_url)
} else if api_format.starts_with("openai:") || api_format.starts_with("claude:") {
build_v1_models_url(base_url)
} else if api_format.starts_with("gemini:") {
build_gemini_models_url(base_url)
@@ -115,23 +118,20 @@ pub fn parse_models_response_page(
items
} else if let Some(items) = body.as_array() {
items
} else if let Some(items) = body.get("models").and_then(Value::as_array) {
items
} else {
return Err("models response is missing data array".to_string());
};
for item in items {
let Some(model_id) = item
.get("id")
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
else {
let Some(model_id) = model_id_from_openai_like_item(item) else {
continue;
};
if !seen.insert(model_id.to_string()) {
if !seen.insert(model_id.clone()) {
continue;
}
fetched_model_ids.push(model_id.to_string());
cached_models.push(normalize_cached_model(item, model_id, &api_format));
fetched_model_ids.push(model_id.clone());
cached_models.push(normalize_cached_model(item, &model_id, &api_format));
}
} else if api_format.starts_with("gemini:") {
let items = body
@@ -229,7 +229,7 @@ pub fn endpoint_supports_rust_models_fetch(api_format: &str) -> bool {
pub fn provider_type_uses_preset_models(provider_type: &str) -> bool {
matches!(
provider_type.trim().to_ascii_lowercase().as_str(),
"codex" | "kiro" | "claude_code" | "gemini_cli"
"kiro" | "claude_code" | "gemini_cli"
)
}
@@ -258,34 +258,11 @@ pub fn preset_models_for_provider(provider_type: &str) -> Option<Vec<Value>> {
preset_model("claude-haiku-4-5-20251001", "anthropic", "Claude Haiku 4.5", "claude:messages"),
],
"codex" => vec![
preset_model("gpt-5", "openai", "GPT-5", "openai:responses"),
preset_model("gpt-image-1", "openai", "GPT Image 1", "openai:image"),
preset_model("gpt-image-1.5", "openai", "GPT Image 1.5", "openai:image"),
preset_model("gpt-image-1-mini", "openai", "GPT Image 1 Mini", "openai:image"),
preset_model("gpt-image-2", "openai", "GPT Image 2", "openai:image"),
preset_model("chatgpt-image-latest", "openai", "ChatGPT Image Latest", "openai:image"),
preset_model("dall-e-2", "openai", "DALL-E 2", "openai:image"),
preset_model("dall-e-3", "openai", "DALL-E 3", "openai:image"),
preset_model("gpt-5-codex", "openai", "GPT-5 Codex", "openai:responses"),
preset_model("gpt-5-codex-mini", "openai", "GPT-5 Codex Mini", "openai:responses"),
preset_model("gpt-5.1", "openai", "GPT-5.1", "openai:responses"),
preset_model("gpt-5.1-codex", "openai", "GPT-5.1 Codex", "openai:responses"),
preset_model(
"gpt-5.1-codex-mini",
"openai",
"GPT-5.1 Codex Mini",
"openai:responses",
),
preset_model(
"gpt-5.1-codex-max",
"openai",
"GPT-5.1 Codex Max",
"openai:responses",
),
preset_model("gpt-5.2", "openai", "GPT-5.2", "openai:responses"),
preset_model("gpt-5.2-codex", "openai", "GPT-5.2 Codex", "openai:responses"),
preset_model("gpt-5.3-codex", "openai", "GPT-5.3 Codex", "openai:responses"),
preset_model("gpt-5.5", "openai", "GPT-5.5", "openai:responses"),
preset_model("gpt-5.4", "openai", "GPT-5.4", "openai:responses"),
preset_model("gpt-5.4-mini", "openai", "GPT-5.4 Mini", "openai:responses"),
preset_model("gpt-5.3-codex", "openai", "GPT-5.3 Codex", "openai:responses"),
preset_model("gpt-5.3-codex-spark", "openai", "GPT-5.3 Codex Spark", "openai:responses"),
],
_ => return None,
};
@@ -485,6 +462,37 @@ fn build_v1_models_url(base_url: &str) -> Option<String> {
Some(url)
}
fn build_codex_models_url(base_url: &str) -> Option<String> {
let (trimmed_base_url, query) = split_url_query(base_url);
let trimmed_base_url = trimmed_base_url.trim_end_matches('/');
if trimmed_base_url.is_empty() {
return None;
}
let mut url = if trimmed_base_url.ends_with("/models") {
trimmed_base_url.to_string()
} else {
format!("{trimmed_base_url}/models")
};
let mut has_client_version = false;
if let Some(query) = query.filter(|value| !value.trim().is_empty()) {
has_client_version = query.split('&').any(|part| {
part.split_once('=')
.map(|(key, _)| key)
.unwrap_or(part)
.trim()
.eq_ignore_ascii_case("client_version")
});
url.push('?');
url.push_str(query);
}
if !has_client_version {
let separator = if url.contains('?') { '&' } else { '?' };
url.push(separator);
url.push_str("client_version=0.128.0-alpha.1");
}
Some(url)
}
fn build_gemini_models_url(base_url: &str) -> Option<String> {
let (trimmed_base_url, base_query) = split_url_query(base_url);
let trimmed_base_url = trimmed_base_url.trim_end_matches('/');
@@ -506,6 +514,24 @@ fn build_gemini_models_url(base_url: &str) -> Option<String> {
Some(url)
}
fn model_id_from_openai_like_item(item: &Value) -> Option<String> {
if let Some(value) = item
.as_str()
.map(str::trim)
.filter(|value| !value.is_empty())
{
return Some(value.trim_start_matches("models/").to_string());
}
["id", "model", "slug", "name"].iter().find_map(|field| {
item.get(*field)
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.map(|value| value.trim_start_matches("models/").to_string())
})
}
fn split_url_query(base_url: &str) -> (&str, Option<&str>) {
let trimmed = base_url.trim();
trimmed
@@ -707,6 +733,22 @@ mod tests {
);
}
#[test]
fn build_models_fetch_url_uses_codex_backend_models_endpoint() {
assert_eq!(
build_models_fetch_url(
"codex",
"openai:responses",
"https://chatgpt.com/backend-api/codex"
),
Some((
"https://chatgpt.com/backend-api/codex/models?client_version=0.128.0-alpha.1"
.to_string(),
"openai:responses".to_string()
))
);
}
#[test]
fn parse_models_response_normalizes_openai_payload() {
let parsed = parse_models_response(
@@ -721,6 +763,23 @@ mod tests {
);
}
#[test]
fn parse_models_response_accepts_codex_models_array_payload() {
let parsed = parse_models_response(
"openai:responses",
&json!({"models": [{"id": "gpt-5-codex"}, {"slug": "gpt-5.4"}]}),
)
.expect("response should parse");
assert_eq!(
parsed.fetched_model_ids,
vec!["gpt-5-codex".to_string(), "gpt-5.4".to_string()]
);
assert_eq!(
parsed.cached_models[0]["api_formats"],
json!(["openai:responses"])
);
}
#[test]
fn parse_models_response_page_reads_claude_pagination_state() {
let parsed = parse_models_response_page(
@@ -822,6 +881,19 @@ mod tests {
#[test]
fn preset_models_cover_codex_catalog() {
let models = preset_models_for_provider("codex").expect("preset models should exist");
assert!(models.iter().any(|model| model["id"] == "gpt-5.4"));
let model_ids = models
.iter()
.map(|model| model["id"].as_str().expect("model id"))
.collect::<Vec<_>>();
assert_eq!(
model_ids,
vec![
"gpt-5.5",
"gpt-5.4",
"gpt-5.4-mini",
"gpt-5.3-codex",
"gpt-5.3-codex-spark",
]
);
}
}

View File

@@ -55,6 +55,9 @@ pub async fn fetch_models_from_transports(
.trim()
.to_ascii_lowercase();
if let Some(models) = preset_models_for_provider(&provider_type) {
if provider_type == "codex" {
return fetch_standard_models(runtime, transports).await;
}
if provider_type == "gemini_cli" {
return fetch_gemini_cli_models(runtime, first_transport, models).await;
}
@@ -1079,13 +1082,14 @@ mod tests {
GatewayProviderTransportProvider, GatewayProviderTransportSnapshot,
};
use async_trait::async_trait;
use serde_json::json;
use serde_json::{json, Value};
use crate::fetch_models_from_transports;
use crate::transport::ModelFetchTransportRuntime;
struct TestRuntime {
executed_urls: Arc<Mutex<Vec<String>>>,
response_body: Value,
}
#[async_trait]
@@ -1119,11 +1123,7 @@ mod tests {
status_code: 200,
headers: BTreeMap::new(),
body: Some(ResponseBody {
json_body: Some(json!({
"models": [{
"name": "publishers/google/models/gemini-3.1-pro-preview"
}]
})),
json_body: Some(self.response_body.clone()),
body_bytes_b64: None,
}),
telemetry: None,
@@ -1188,12 +1188,31 @@ mod tests {
}
}
fn sample_codex_transport() -> GatewayProviderTransportSnapshot {
let mut transport = sample_custom_aiplatform_transport();
transport.provider.provider_type = "codex".to_string();
transport.provider.name = "Codex".to_string();
transport.endpoint.api_format = "openai:responses".to_string();
transport.endpoint.api_family = Some("openai".to_string());
transport.endpoint.endpoint_kind = Some("responses".to_string());
transport.endpoint.base_url = "https://chatgpt.com/backend-api/codex".to_string();
transport.endpoint.custom_path = Some("/responses".to_string());
transport.key.api_formats = Some(vec!["openai:responses".to_string()]);
transport.key.decrypted_api_key = "access-token".to_string();
transport
}
#[tokio::test]
async fn custom_aiplatform_transport_uses_vertex_models_fetch_path_and_normalizes_chat_format()
{
let executed_urls = Arc::new(Mutex::new(Vec::new()));
let runtime = TestRuntime {
executed_urls: Arc::clone(&executed_urls),
response_body: json!({
"models": [{
"name": "publishers/google/models/gemini-3.1-pro-preview"
}]
}),
};
let outcome =
fetch_models_from_transports(&runtime, &[sample_custom_aiplatform_transport()])
@@ -1212,4 +1231,28 @@ mod tests {
Some("gemini:generate_content")
);
}
#[tokio::test]
async fn codex_transport_fetches_upstream_models_instead_of_preset_catalog() {
let executed_urls = Arc::new(Mutex::new(Vec::new()));
let runtime = TestRuntime {
executed_urls: Arc::clone(&executed_urls),
response_body: json!({
"models": [{
"id": "gpt-5.4-upstream"
}]
}),
};
let outcome = fetch_models_from_transports(&runtime, &[sample_codex_transport()])
.await
.expect("models fetch should succeed");
let urls = executed_urls.lock().expect("executed_urls lock");
assert_eq!(
urls.as_slice(),
&["https://chatgpt.com/backend-api/codex/models?client_version=0.128.0-alpha.1"]
);
assert_eq!(outcome.fetched_model_ids, vec!["gpt-5.4-upstream"]);
assert_eq!(outcome.cached_models.len(), 1);
}
}

View File

@@ -15,7 +15,7 @@ use aether_provider_transport::{
GatewayProviderTransportSnapshot, LocalResolvedOAuthRequestAuth,
};
use async_trait::async_trait;
use serde_json::json;
use serde_json::{json, Value};
use crate::build_models_fetch_url;
@@ -85,7 +85,13 @@ pub async fn build_standard_models_fetch_execution_plan(
) -> Result<ExecutionPlan, String> {
let api_format = transport.endpoint.api_format.trim().to_ascii_lowercase();
let provider_api_format = api_format.clone();
let mut headers = standard_models_fetch_headers(&api_format, &transport.provider.provider_type);
let provider_type = transport.provider.provider_type.trim().to_ascii_lowercase();
let is_codex_openai_models_fetch =
provider_type == "codex" && api_format.starts_with("openai:");
let mut headers = standard_models_fetch_headers(&api_format, &provider_type);
if is_codex_openai_models_fetch {
headers.insert("accept".to_string(), "application/json".to_string());
}
let mut protected_headers = Vec::<String>::new();
if api_format.starts_with("openai:") || api_format.starts_with("claude:") {
@@ -101,6 +107,16 @@ pub async fn build_standard_models_fetch_execution_plan(
&auth_header_name,
&auth_header_value,
);
if is_codex_openai_models_fetch {
if let Some(account_id) = extract_codex_account_id(transport) {
insert_non_empty_auth_header(
&mut headers,
&mut protected_headers,
"chatgpt-account-id",
&account_id,
);
}
}
headers = apply_fetch_header_rules(transport, headers, &protected_headers)?;
ensure_upstream_auth_header(&mut headers, &auth_header_name, &auth_header_value);
} else {
@@ -481,6 +497,23 @@ fn append_query_param(mut url: String, key: &str, value: &str) -> String {
url
}
fn extract_codex_account_id(transport: &GatewayProviderTransportSnapshot) -> Option<String> {
let raw = transport.key.decrypted_auth_config.as_deref()?.trim();
if raw.is_empty() {
return None;
}
serde_json::from_str::<Value>(raw).ok().and_then(|value| {
value
.get("account_id")
.or_else(|| value.get("chatgpt_account_id"))
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToOwned::to_owned)
})
}
fn insert_non_empty_auth_header(
headers: &mut BTreeMap<String, String>,
protected_headers: &mut Vec<String>,
@@ -648,6 +681,43 @@ mod tests {
);
}
#[tokio::test]
async fn builds_codex_models_fetch_plan_with_account_header() {
let runtime = TestRuntime {
oauth_auth: Some(
aether_provider_transport::LocalResolvedOAuthRequestAuth::Header {
name: "authorization".to_string(),
value: "Bearer access-token".to_string(),
},
),
proxy: None,
};
let mut transport = sample_transport("codex", "openai:responses", "oauth");
transport.endpoint.base_url = "https://chatgpt.com/backend-api/codex".to_string();
transport.key.decrypted_auth_config = Some(r#"{"account_id":"account-1"}"#.to_string());
let plan = build_models_fetch_execution_plan(&runtime, &transport)
.await
.expect("plan");
assert_eq!(
plan.url,
"https://chatgpt.com/backend-api/codex/models?client_version=0.128.0-alpha.1"
);
assert_eq!(
plan.headers.get("authorization").map(String::as_str),
Some("Bearer access-token")
);
assert_eq!(
plan.headers.get("chatgpt-account-id").map(String::as_str),
Some("account-1")
);
assert_eq!(
plan.headers.get("accept").map(String::as_str),
Some("application/json")
);
}
#[tokio::test]
async fn builds_claude_models_fetch_plan_with_pagination() {
let runtime = TestRuntime {