i/Aether
1
0
Fork 0
mirror of https://github.com/fawney19/Aether.git synced 2026-09-02 01:10:23 +08:00
Code Issues Packages Projects Releases Wiki Activity

feat(model-fetch): 对齐 Rust 上游模型抓取行为到 Python 语义

Browse Source 
将 Rust 版上游可用模型抓取逻辑收敛到 Python 版行为，统一后台自动抓模
与管理员 provider-query 的模型发现路径，消除标准 /models、固定模型目录、
Antigravity、Vertex AI 等 provider 在两端实现上的分叉。

核心变更：
- 在 aether-model-fetch 中引入统一抓模策略层
- 覆盖标准 /models、Vertex API Key、Vertex Service Account、
  Antigravity fetchAvailableModels、固定模型目录五类抓模路径
- 将 provider-query 与后台自动抓模都切换到共享抓模入口，避免重复拼接
  URL、headers 和 provider 特判逻辑

标准模型抓取对齐：
- 按 Python 语义调整抓模优先级：
  openai:chat > openai:cli > openai:compact
  claude:chat > claude:cli
  gemini:chat > gemini:cli
- 从抓模候选中移除 openai:responses
- 为 openai:cli/openai:compact、claude:cli、gemini:* 补齐 Python 同款
  User-Agent / 浏览器指纹请求头
- Claude 抓模保留 after_id 分页语义
- Gemini 抓模统一为 v1beta/models?key=... 语义

provider-query 对齐：
- 返回结果改为按 model id 聚合，并合并/排序 api_formats
- 最终模型列表按 model id 排序，行为与 Python 保持一致
- 固定目录 provider（codex/kiro/claude_code/gemini_cli）不再依赖活跃
  endpoint，即使无 endpoint 也能返回预设模型目录
- Antigravity 多 key 查询改为按账户可用性 + tier 排序，首个成功结果即
  停止，并接入 provider 级缓存
- 仅配置 openai:responses 的 provider 不再被视为抓模成功路径

自动抓模对齐：
- 自动抓模成功时写入 allowed_models、upstream_models cache，并同步
  upstream_metadata
- upstream_metadata 合并逻辑对齐 Python，对 quota_by_model 做模型级合并，
  并保留已有 reset_time
- 自动抓模失败时不覆盖已有 allowed_models
- 固定目录 provider 在无 endpoint 场景下也可成功更新 allowed_models

Antigravity 对齐：
- 使用 POST /v1internal:fetchAvailableModels 抓取可用模型
- 按 Python 规则处理 URL fallback 和 429/404/408/5xx fallback 状态
- 强制要求 auth_config.project_id
- 过滤 Python 黑名单模型
- 解析并持久化 upstream_metadata.antigravity.quota_by_model

Vertex AI 对齐：
- API Key 模式仅抓取 publishers/google/models
- Service Account 模式新增 JWT token exchange，并按 Python region 顺序
  抓取 google + anthropic publishers
- 模型 owned_by / display_name / api_format 推断与 Python 对齐
- 软 404 处理行为与 Python 收敛

Gemini CLI / 固定目录对齐：
- Gemini CLI 改为返回 Python 预设模型目录
- 在可用时通过 loadCodeAssist 补充 plan_type/project_id 元数据
- Codex/Kiro/Claude Code 改为共享固定模型目录实现

测试：
- 扩展 aether-model-fetch 单元测试，覆盖格式优先级、openai:responses 排除、
  请求头、Claude 分页、Gemini query auth、固定目录与 metadata 合并
- 调整 provider-query 控制面测试到 Python 语义
- 新增自动抓模运行时测试，覆盖固定目录成功、Antigravity metadata 合并、
  失败保留旧 allowed_models

验证：
- cargo nextest run -p aether-model-fetch --lib
- cargo nextest run -p aether-gateway control::admin::provider_query model_fetch::runtime::tests
This commit is contained in:
AAEE86
2026-04-12 10:14:29 +08:00
parent def8135118
commit ab82841426
10 changed files with 2815 additions and 338 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
crates/aether-model-fetch/Cargo.toml
View File

@@ -11,8 +11,11 @@ aether-data-contracts.workspace = true
aether-provider-transport.workspace = true
aether-scheduler-core.workspace = true
async-trait.workspace = true
base64.workspace = true
regex.workspace = true
rsa = "0.9.10"
serde_json.workspace = true
sha2 = { workspace = true, features = ["oid"] }
uuid.workspace = true
[dev-dependencies]

12
crates/aether-model-fetch/src/lib.rs
View File

@@ -1,6 +1,7 @@
mod association_sync;
mod config;
mod logic;
mod strategy;
mod transport;
pub use association_sync::{
@@ -12,6 +13,13 @@ pub use config::{
pub use logic::{
aggregate_models_for_cache, apply_model_filters, build_models_fetch_url,
endpoint_supports_rust_models_fetch, extract_error_message, json_string_list,
parse_models_response, select_models_fetch_endpoint, ModelFetchRunSummary, ModelsFetchSuccess,
merge_upstream_metadata, parse_models_response, parse_models_response_page,
preset_models_for_provider, provider_type_uses_preset_models, select_models_fetch_endpoint,
selected_models_fetch_endpoints, ModelFetchRunSummary, ModelsFetchPage, ModelsFetchSuccess,
};
pub use strategy::{fetch_models_from_transports, ModelsFetchOutcome};
pub use transport::{
build_antigravity_fetch_available_models_plan, build_gemini_cli_load_code_assist_plan,
build_models_fetch_execution_plan, build_standard_models_fetch_execution_plan,
build_vertex_models_fetch_execution_plan, ModelFetchTransportRuntime,
};
pub use transport::{build_models_fetch_execution_plan, ModelFetchTransportRuntime};

329
crates/aether-model-fetch/src/logic.rs
View File

@@ -3,9 +3,14 @@ use std::collections::{BTreeMap, BTreeSet};
use aether_data_contracts::repository::provider_catalog::{
StoredProviderCatalogEndpoint, StoredProviderCatalogKey,
};
use aether_provider_transport::provider_types::provider_type_supports_model_fetch;
use regex::Regex;
use serde_json::Value;
use serde_json::{json, Value};
const MODEL_FETCH_FORMAT_PRIORITY: &[&[&str]] = &[
&["openai:chat", "openai:cli", "openai:compact"],
&["claude:chat", "claude:cli"],
&["gemini:chat", "gemini:cli"],
];
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct ModelFetchRunSummary {
@@ -21,6 +26,14 @@ pub struct ModelsFetchSuccess {
pub cached_models: Vec<Value>,
}
#[derive(Debug, Clone, PartialEq)]
pub struct ModelsFetchPage {
pub fetched_model_ids: Vec<String>,
pub cached_models: Vec<Value>,
pub has_more: bool,
pub next_after_id: Option<String>,
}
pub fn extract_error_message(value: &Value) -> Option<String> {
value
.get("error")
@@ -41,21 +54,20 @@ 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)> {
let api_format = normalize_api_format(endpoint_api_format);
if !provider_type_supports_model_fetch(provider_type) {
if !endpoint_supports_rust_models_fetch(&api_format) {
return None;
}
let url = 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)
} else {
return None;
None
}?;
Some((url, api_format))
}
@@ -64,13 +76,38 @@ pub fn parse_models_response(
endpoint_api_format: &str,
body: &Value,
) -> Result<ModelsFetchSuccess, String> {
let parsed = parse_models_response_page(endpoint_api_format, body)?;
Ok(ModelsFetchSuccess {
fetched_model_ids: parsed.fetched_model_ids,
cached_models: parsed.cached_models,
})
}
pub fn parse_models_response_page(
endpoint_api_format: &str,
body: &Value,
) -> Result<ModelsFetchPage, String> {
let api_format = normalize_api_format(endpoint_api_format);
let mut cached_models = Vec::new();
let mut fetched_model_ids = Vec::new();
let mut seen = BTreeSet::new();
let mut has_more = false;
let mut next_after_id = None;
if api_format.starts_with("openai:") || api_format.starts_with("claude:") {
let items = if let Some(items) = body.get("data").and_then(Value::as_array) {
has_more = body
.get("has_more")
.and_then(Value::as_bool)
.unwrap_or(false);
if api_format.starts_with("claude:") && has_more {
next_after_id = body
.get("last_id")
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToOwned::to_owned);
}
items
} else if let Some(items) = body.as_array() {
items
@@ -117,29 +154,54 @@ pub fn parse_models_response(
return Err("models response parser does not support this provider format".to_string());
}
Ok(ModelsFetchSuccess {
Ok(ModelsFetchPage {
fetched_model_ids,
cached_models,
has_more,
next_after_id,
})
}
pub fn selected_models_fetch_endpoints(
endpoints: &[StoredProviderCatalogEndpoint],
key: &StoredProviderCatalogKey,
) -> Vec<StoredProviderCatalogEndpoint> {
let key_formats = json_string_list(key.api_formats.as_ref())
.into_iter()
.map(|value| normalize_api_format(&value))
.collect::<BTreeSet<_>>();
let mut by_format = BTreeMap::<String, StoredProviderCatalogEndpoint>::new();
for endpoint in endpoints.iter().filter(|endpoint| endpoint.is_active) {
let api_format = normalize_api_format(&endpoint.api_format);
if api_format.is_empty() || !endpoint_supports_rust_models_fetch(&api_format) {
continue;
}
if !key_formats.is_empty() && !key_formats.contains(&api_format) {
continue;
}
by_format
.entry(api_format)
.or_insert_with(|| endpoint.clone());
}
MODEL_FETCH_FORMAT_PRIORITY
.iter()
.filter_map(|candidates| {
candidates
.iter()
.find_map(|api_format| by_format.remove(*api_format))
})
.collect()
}
pub fn select_models_fetch_endpoint(
endpoints: &[StoredProviderCatalogEndpoint],
key: &StoredProviderCatalogKey,
) -> Option<StoredProviderCatalogEndpoint> {
let key_formats = json_string_list(key.api_formats.as_ref())
selected_models_fetch_endpoints(endpoints, key)
.into_iter()
.map(|value| normalize_api_format(&value))
.collect::<BTreeSet<_>>();
endpoints
.iter()
.filter(|endpoint| endpoint.is_active)
.find(|endpoint| {
let api_format = normalize_api_format(&endpoint.api_format);
(key_formats.is_empty() || key_formats.contains(&api_format))
&& endpoint_supports_rust_models_fetch(&endpoint.api_format)
})
.cloned()
.next()
}
pub fn endpoint_supports_rust_models_fetch(api_format: &str) -> bool {
@@ -148,7 +210,6 @@ pub fn endpoint_supports_rust_models_fetch(api_format: &str) -> bool {
api_format.as_str(),
"openai:chat"
| "openai:cli"
| "openai:responses"
| "openai:compact"
| "claude:chat"
| "claude:cli"
@@ -157,6 +218,100 @@ 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"
)
}
#[rustfmt::skip]
pub fn preset_models_for_provider(provider_type: &str) -> Option<Vec<Value>> {
let models = match provider_type.trim().to_ascii_lowercase().as_str() {
"gemini_cli" => vec![
preset_model("gemini-2.5-pro", "google", "Gemini 2.5 Pro", "gemini:cli"),
preset_model("gemini-2.5-flash", "google", "Gemini 2.5 Flash", "gemini:cli"),
preset_model("gemini-3-pro-preview", "google", "Gemini 3 Pro Preview", "gemini:cli"),
preset_model("gemini-3-flash-preview", "google", "Gemini 3 Flash Preview", "gemini:cli"),
preset_model("gemini-3.1-pro-preview", "google", "Gemini 3.1 Pro Preview", "gemini:cli"),
],
"kiro" => vec![
preset_model("claude-sonnet-4.5", "anthropic", "Claude Sonnet 4.5", "claude:cli"),
preset_model("claude-sonnet-4.6", "anthropic", "Claude Sonnet 4.6", "claude:cli"),
preset_model("claude-opus-4.5", "anthropic", "Claude Opus 4.5", "claude:cli"),
preset_model("claude-opus-4.6", "anthropic", "Claude Opus 4.6", "claude:cli"),
preset_model("claude-haiku-4.5", "anthropic", "Claude Haiku 4.5", "claude:cli"),
],
"claude_code" => vec![
preset_model("claude-opus-4-5-20251101", "anthropic", "Claude Opus 4.5", "claude:cli"),
preset_model("claude-opus-4-6", "anthropic", "Claude Opus 4.6", "claude:cli"),
preset_model("claude-sonnet-4-6", "anthropic", "Claude Sonnet 4.6", "claude:cli"),
preset_model("claude-sonnet-4-5-20250929", "anthropic", "Claude Sonnet 4.5", "claude:cli"),
preset_model("claude-haiku-4-5-20251001", "anthropic", "Claude Haiku 4.5", "claude:cli"),
],
"codex" => vec![
preset_model("gpt-5", "openai", "GPT-5", "openai:cli"),
preset_model("gpt-5-codex", "openai", "GPT-5 Codex", "openai:cli"),
preset_model("gpt-5-codex-mini", "openai", "GPT-5 Codex Mini", "openai:cli"),
preset_model("gpt-5.1", "openai", "GPT-5.1", "openai:cli"),
preset_model("gpt-5.1-codex", "openai", "GPT-5.1 Codex", "openai:cli"),
preset_model("gpt-5.1-codex-mini", "openai", "GPT-5.1 Codex Mini", "openai:cli"),
preset_model("gpt-5.1-codex-max", "openai", "GPT-5.1 Codex Max", "openai:cli"),
preset_model("gpt-5.2", "openai", "GPT-5.2", "openai:cli"),
preset_model("gpt-5.2-codex", "openai", "GPT-5.2 Codex", "openai:cli"),
preset_model("gpt-5.3-codex", "openai", "GPT-5.3 Codex", "openai:cli"),
preset_model("gpt-5.4", "openai", "GPT-5.4", "openai:cli"),
],
_ => return None,
};
Some(models)
}
pub fn merge_upstream_metadata(current: Option<&Value>, incoming: &Value) -> Value {
let mut merged = current
.and_then(Value::as_object)
.cloned()
.unwrap_or_default();
let Some(incoming_object) = incoming.as_object() else {
return Value::Object(merged);
};
for (namespace, value) in incoming_object {
let mut next_value = value.clone();
if let (Some(next_namespace), Some(old_namespace)) = (
next_value.as_object_mut(),
merged.get(namespace).and_then(Value::as_object),
) {
if let (Some(new_quota), Some(old_quota)) = (
next_namespace
.get_mut("quota_by_model")
.and_then(Value::as_object_mut),
old_namespace
.get("quota_by_model")
.and_then(Value::as_object),
) {
for (model_id, new_info) in new_quota.iter_mut() {
let Some(new_info_object) = new_info.as_object_mut() else {
continue;
};
let Some(old_info_object) = old_quota.get(model_id).and_then(Value::as_object)
else {
continue;
};
if !new_info_object.contains_key("reset_time") {
if let Some(reset_time) = old_info_object.get("reset_time") {
new_info_object.insert("reset_time".to_string(), reset_time.clone());
}
}
}
}
}
merged.insert(namespace.clone(), next_value);
}
Value::Object(merged)
}
pub fn apply_model_filters(
fetched_model_ids: &[String],
locked_models: Vec<String>,
@@ -211,7 +366,6 @@ pub fn json_string_list(value: Option<&Value>) -> Vec<String> {
pub fn aggregate_models_for_cache(models: &[Value]) -> Vec<Value> {
let mut aggregated = BTreeMap::<String, serde_json::Map<String, Value>>::new();
let mut order = Vec::<String>::new();
for model in models {
let Some(object) = model.as_object() else {
@@ -227,7 +381,6 @@ pub fn aggregate_models_for_cache(models: &[Value]) -> Vec<Value> {
};
let entry = aggregated.entry(model_id.to_string()).or_insert_with(|| {
order.push(model_id.to_string());
let mut cloned = object.clone();
cloned.remove("api_format");
cloned
@@ -286,11 +439,7 @@ pub fn aggregate_models_for_cache(models: &[Value]) -> Vec<Value> {
}
}
order
.into_iter()
.filter_map(|model_id| aggregated.remove(&model_id))
.map(Value::Object)
.collect()
aggregated.into_values().map(Value::Object).collect()
}
fn build_v1_models_url(base_url: &str) -> Option<String> {
@@ -347,10 +496,37 @@ fn normalize_cached_model(item: &Value, model_id: &str, api_format: &str) -> Val
"api_formats".to_string(),
Value::Array(vec![Value::String(api_format.to_string())]),
);
if api_format.starts_with("gemini:") {
object
.entry("owned_by".to_string())
.or_insert_with(|| Value::String("google".to_string()));
if !object.contains_key("display_name") {
let display_name = item
.get("displayName")
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.unwrap_or(model_id);
object.insert(
"display_name".to_string(),
Value::String(display_name.to_string()),
);
}
}
object.remove("api_format");
Value::Object(object)
}
fn preset_model(model_id: &str, owned_by: &str, display_name: &str, api_format: &str) -> Value {
json!({
"id": model_id,
"object": "model",
"owned_by": owned_by,
"display_name": display_name,
"api_formats": [api_format],
})
}
fn wildcard_matches(pattern: &str, model_id: &str) -> bool {
let mut regex = String::from("^");
for ch in pattern.chars() {
@@ -379,7 +555,8 @@ mod tests {
use super::{
aggregate_models_for_cache, apply_model_filters, build_gemini_models_url,
build_models_fetch_url, parse_models_response, select_models_fetch_endpoint,
build_models_fetch_url, merge_upstream_metadata, parse_models_response,
parse_models_response_page, preset_models_for_provider, selected_models_fetch_endpoints,
};
fn sample_endpoint(
@@ -410,7 +587,11 @@ mod tests {
.expect("endpoint transport should build")
}
fn sample_key(provider_id: &str, key_id: &str) -> StoredProviderCatalogKey {
fn sample_key(
provider_id: &str,
key_id: &str,
api_formats: &[&str],
) -> StoredProviderCatalogKey {
StoredProviderCatalogKey::new(
key_id.to_string(),
provider_id.to_string(),
@@ -421,7 +602,7 @@ mod tests {
)
.expect("key should build")
.with_transport_fields(
Some(json!(["openai:chat"])),
Some(json!(api_formats)),
"encrypted".to_string(),
None,
None,
@@ -453,12 +634,15 @@ mod tests {
}
#[test]
fn aggregate_models_for_cache_merges_api_formats_by_model_id() {
fn aggregate_models_for_cache_merges_api_formats_and_sorts_by_model_id() {
let aggregated = aggregate_models_for_cache(&[
json!({"id":"gpt-5","api_formats":["openai:chat"]}),
json!({"id":"gpt-5","api_formats":["openai:cli"]}),
json!({"id":"zeta","api_formats":["openai:chat"]}),
json!({"id":"alpha","api_formats":["openai:cli"]}),
json!({"id":"alpha","api_formats":["openai:chat"]}),
]);
assert_eq!(aggregated.len(), 1);
assert_eq!(aggregated.len(), 2);
assert_eq!(aggregated[0]["id"], "alpha");
assert_eq!(aggregated[1]["id"], "zeta");
assert_eq!(
aggregated[0]["api_formats"],
json!(["openai:chat", "openai:cli"])
@@ -488,9 +672,9 @@ mod tests {
}
#[test]
fn build_models_fetch_url_rejects_provider_types_without_fetch_support() {
fn build_models_fetch_url_excludes_openai_responses() {
assert_eq!(
build_models_fetch_url("vertex_ai", "gemini:chat", "https://example.com"),
build_models_fetch_url("openai", "openai:responses", "https://example.com"),
None
);
}
@@ -510,9 +694,30 @@ mod tests {
}
#[test]
fn select_models_fetch_endpoint_respects_key_api_formats() {
let key = sample_key("provider-1", "key-1");
fn parse_models_response_page_reads_claude_pagination_state() {
let parsed = parse_models_response_page(
"claude:chat",
&json!({
"data": [{"id": "claude-sonnet-4"}],
"has_more": true,
"last_id": "cursor-2"
}),
)
.expect("response should parse");
assert!(parsed.has_more);
assert_eq!(parsed.next_after_id.as_deref(), Some("cursor-2"));
}
#[test]
fn selected_models_fetch_endpoints_prefers_chat_and_excludes_responses() {
let key = sample_key("provider-1", "key-1", &["openai:chat", "openai:responses"]);
let endpoints = vec![
sample_endpoint(
"provider-1",
"endpoint-responses",
"openai:responses",
"https://example.com",
),
sample_endpoint(
"provider-1",
"endpoint-cli",
@@ -526,8 +731,50 @@ mod tests {
"https://example.com",
),
];
let selected =
select_models_fetch_endpoint(&endpoints, &key).expect("endpoint should be selected");
assert_eq!(selected.id, "endpoint-chat");
let selected = selected_models_fetch_endpoints(&endpoints, &key);
assert_eq!(selected.len(), 1);
assert_eq!(selected[0].id, "endpoint-chat");
}
#[test]
fn merge_upstream_metadata_keeps_existing_reset_time_for_returned_models() {
let merged = merge_upstream_metadata(
Some(&json!({
"antigravity": {
"quota_by_model": {
"gemini-2.5-pro": {
"remaining_fraction": 0.3,
"reset_time": "2026-04-12T00:00:00Z"
},
"stale-model": {
"remaining_fraction": 0.1,
"reset_time": "old"
}
}
}
})),
&json!({
"antigravity": {
"quota_by_model": {
"gemini-2.5-pro": {
"remaining_fraction": 0.6
}
}
}
}),
);
assert_eq!(
merged["antigravity"]["quota_by_model"]["gemini-2.5-pro"]["reset_time"],
"2026-04-12T00:00:00Z"
);
assert!(merged["antigravity"]["quota_by_model"]
.get("stale-model")
.is_none());
}
#[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"));
}
}

1063
crates/aether-model-fetch/src/strategy.rs Normal file
View File

File diff suppressed because it is too large Load Diff

643
crates/aether-model-fetch/src/transport.rs
View File

@@ -1,20 +1,47 @@
use std::collections::BTreeMap;
use aether_contracts::{ExecutionPlan, ProxySnapshot, RequestBody};
use aether_provider_transport::auth::{
resolve_local_gemini_auth, resolve_local_openai_chat_auth, resolve_local_standard_auth,
use aether_contracts::{ExecutionPlan, ExecutionResult, ProxySnapshot, RequestBody};
use aether_provider_transport::antigravity::{
build_antigravity_static_identity_headers, resolve_local_antigravity_request_auth,
AntigravityRequestAuthSupport, ANTIGRAVITY_REQUEST_USER_AGENT,
};
use aether_provider_transport::auth::{
ensure_upstream_auth_header, resolve_local_gemini_auth, resolve_local_openai_chat_auth,
resolve_local_standard_auth,
};
use aether_provider_transport::url::build_passthrough_path_url;
use aether_provider_transport::vertex::resolve_local_vertex_api_key_query_auth;
use aether_provider_transport::{
apply_local_header_rules, ensure_upstream_auth_header, resolve_transport_execution_timeouts,
resolve_transport_tls_profile, GatewayProviderTransportSnapshot, LocalResolvedOAuthRequestAuth,
apply_local_header_rules, resolve_transport_execution_timeouts, resolve_transport_tls_profile,
GatewayProviderTransportSnapshot, LocalResolvedOAuthRequestAuth,
};
use async_trait::async_trait;
use serde_json::json;
use crate::build_models_fetch_url;
const OPENAI_CLI_USER_AGENT: &str = "openai-codex/1.0";
const CLAUDE_CLI_USER_AGENT: &str = "claude-code/1.0.1";
const GEMINI_CLI_USER_AGENT: &str = "GeminiCLI/0.1.5 (Windows; AMD64)";
const CLAUDE_VERSION_HEADER: &str = "2023-06-01";
const ANTIGRAVITY_FETCH_PROVIDER_API_FORMAT: &str = "antigravity:fetch_available_models";
const GEMINI_CLI_LOAD_CODE_ASSIST_PROVIDER_API_FORMAT: &str = "gemini_cli:load_code_assist";
const BROWSER_FINGERPRINT_HEADERS: &[(&str, &str)] = &[
(
"user-agent",
"Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/140.0.7339.249 Electron/38.7.0 Safari/537.36",
),
("accept", "application/json"),
("accept-encoding", "gzip, deflate, br"),
("accept-language", "zh-CN"),
("sec-ch-ua", "\"Not=A?Brand\";v=\"24\", \"Chromium\";v=\"140\""),
("sec-ch-ua-mobile", "?0"),
("sec-ch-ua-platform", "\"macOS\""),
("sec-fetch-site", "cross-site"),
("sec-fetch-mode", "cors"),
("sec-fetch-dest", "empty"),
];
#[async_trait]
pub trait ModelFetchTransportRuntime: Send + Sync {
async fn resolve_local_oauth_request_auth(
@@ -26,113 +53,415 @@ pub trait ModelFetchTransportRuntime: Send + Sync {
&self,
transport: &GatewayProviderTransportSnapshot,
) -> Option<ProxySnapshot>;
async fn execute_model_fetch_execution_plan(
&self,
plan: &ExecutionPlan,
) -> Result<ExecutionResult, String>;
}
pub async fn build_models_fetch_execution_plan(
runtime: &(impl ModelFetchTransportRuntime + ?Sized),
transport: &GatewayProviderTransportSnapshot,
) -> Result<ExecutionPlan, String> {
let (upstream_url, provider_api_format) = build_models_fetch_url(
&transport.provider.provider_type,
&transport.endpoint.api_format,
&transport.endpoint.base_url,
)
.ok_or_else(|| "Rust models fetch does not support this provider format yet".to_string())?;
let (auth_header_name, auth_header_value) = resolve_models_fetch_auth(runtime, transport)
.await?
.ok_or_else(|| {
"Rust models fetch auth resolution is not supported for this key".to_string()
})?;
build_standard_models_fetch_execution_plan(runtime, transport, None).await
}
let mut headers = BTreeMap::from([(auth_header_name.clone(), auth_header_value.clone())]);
if !apply_local_header_rules(
&mut headers,
transport.endpoint.header_rules.as_ref(),
&[auth_header_name.as_str()],
&json!({}),
None,
) {
return Err("Endpoint header_rules application failed".to_string());
pub async fn build_standard_models_fetch_execution_plan(
runtime: &(impl ModelFetchTransportRuntime + ?Sized),
transport: &GatewayProviderTransportSnapshot,
after_id: Option<&str>,
) -> 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);
let mut protected_headers = Vec::<String>::new();
if api_format.starts_with("openai:") || api_format.starts_with("claude:") {
let (auth_header_name, auth_header_value) =
resolve_standard_header_auth(runtime, transport)
.await?
.ok_or_else(|| {
"Rust models fetch auth resolution is not supported for this key".to_string()
})?;
protected_headers.push(auth_header_name.clone());
headers.insert(auth_header_name.clone(), auth_header_value.clone());
headers = apply_fetch_header_rules(transport, headers, &protected_headers)?;
ensure_upstream_auth_header(&mut headers, &auth_header_name, &auth_header_value);
} else {
headers = apply_fetch_header_rules(transport, headers, &protected_headers)?;
}
ensure_upstream_auth_header(&mut headers, &auth_header_name, &auth_header_value);
let upstream_url = build_standard_models_fetch_url(transport, after_id)?;
build_execution_plan(
runtime,
transport,
"GET",
upstream_url,
headers,
None,
RequestBody {
json_body: None,
body_bytes_b64: None,
body_ref: None,
},
provider_api_format.clone(),
provider_api_format,
Some("models".to_string()),
)
.await
}
pub async fn build_antigravity_fetch_available_models_plan(
runtime: &(impl ModelFetchTransportRuntime + ?Sized),
transport: &GatewayProviderTransportSnapshot,
base_url: &str,
project_id: &str,
) -> Result<ExecutionPlan, String> {
let authorization = resolve_oauth_header_auth(runtime, transport)
.await?
.ok_or_else(|| "Antigravity fetch requires OAuth authorization header".to_string())?;
let identity_auth = match resolve_local_antigravity_request_auth(transport) {
AntigravityRequestAuthSupport::Supported(auth) => auth,
AntigravityRequestAuthSupport::Unsupported(reason) => {
return Err(format!(
"Antigravity fetch auth resolution is not supported: {reason:?}"
))
}
};
let mut headers = build_antigravity_static_identity_headers(&identity_auth);
headers.insert(authorization.0.clone(), authorization.1.clone());
headers.insert("content-type".to_string(), "application/json".to_string());
headers.insert("accept".to_string(), "application/json".to_string());
headers
.entry("user-agent".to_string())
.or_insert_with(|| ANTIGRAVITY_REQUEST_USER_AGENT.to_string());
let protected_headers = vec![authorization.0];
headers = apply_fetch_header_rules(transport, headers, &protected_headers)?;
let url = format!(
"{}{}",
base_url.trim_end_matches('/'),
"/v1internal:fetchAvailableModels"
);
build_execution_plan(
runtime,
transport,
"POST",
url,
headers,
Some("application/json".to_string()),
RequestBody::from_json(json!({ "project": project_id })),
"gemini:chat".to_string(),
ANTIGRAVITY_FETCH_PROVIDER_API_FORMAT.to_string(),
Some("fetchAvailableModels".to_string()),
)
.await
}
pub async fn build_gemini_cli_load_code_assist_plan(
runtime: &(impl ModelFetchTransportRuntime + ?Sized),
transport: &GatewayProviderTransportSnapshot,
) -> Result<ExecutionPlan, String> {
let authorization = resolve_bearer_or_oauth_header_auth(runtime, transport)
.await?
.ok_or_else(|| "GeminiCLI loadCodeAssist requires bearer or OAuth auth".to_string())?;
let mut headers = BTreeMap::from([
("user-agent".to_string(), GEMINI_CLI_USER_AGENT.to_string()),
("accept-encoding".to_string(), "identity".to_string()),
("content-type".to_string(), "application/json".to_string()),
]);
headers.insert(authorization.0.clone(), authorization.1.clone());
let protected_headers = vec![authorization.0];
headers = apply_fetch_header_rules(transport, headers, &protected_headers)?;
build_execution_plan(
runtime,
transport,
"POST",
"https://cloudcode-pa.googleapis.com/v1internal:loadCodeAssist".to_string(),
headers,
Some("application/json".to_string()),
RequestBody::from_json(json!({
"metadata": {
"ideType": "ANTIGRAVITY",
"platform": "PLATFORM_UNSPECIFIED",
"pluginType": "GEMINI",
}
})),
"gemini:cli".to_string(),
GEMINI_CLI_LOAD_CODE_ASSIST_PROVIDER_API_FORMAT.to_string(),
Some("loadCodeAssist".to_string()),
)
.await
}
pub async fn build_vertex_models_fetch_execution_plan(
runtime: &(impl ModelFetchTransportRuntime + ?Sized),
transport: &GatewayProviderTransportSnapshot,
url: &str,
api_format: &str,
auth_header: Option<(String, String)>,
) -> Result<ExecutionPlan, String> {
let mut headers = standard_models_fetch_headers(api_format);
let mut protected_headers = Vec::<String>::new();
if let Some((name, value)) = auth_header {
protected_headers.push(name.clone());
headers.insert(name.clone(), value.clone());
headers = apply_fetch_header_rules(transport, headers, &protected_headers)?;
ensure_upstream_auth_header(&mut headers, &name, &value);
} else {
headers = apply_fetch_header_rules(transport, headers, &protected_headers)?;
}
build_execution_plan(
runtime,
transport,
"GET",
url.trim().to_string(),
headers,
None,
RequestBody {
json_body: None,
body_bytes_b64: None,
body_ref: None,
},
api_format.to_string(),
api_format.to_string(),
Some("models".to_string()),
)
.await
}
async fn build_execution_plan(
runtime: &(impl ModelFetchTransportRuntime + ?Sized),
transport: &GatewayProviderTransportSnapshot,
method: &str,
url: String,
headers: BTreeMap<String, String>,
content_type: Option<String>,
body: RequestBody,
client_api_format: String,
provider_api_format: String,
model_name: Option<String>,
) -> Result<ExecutionPlan, String> {
Ok(ExecutionPlan {
request_id: format!("req-model-fetch-{}", transport.key.id),
request_id: format!(
"req-model-fetch-{}-{}",
transport.key.id,
provider_api_format.replace(':', "-")
),
candidate_id: None,
provider_name: Some(transport.provider.name.clone()),
provider_id: transport.provider.id.clone(),
endpoint_id: transport.endpoint.id.clone(),
key_id: transport.key.id.clone(),
method: "GET".to_string(),
url: upstream_url,
method: method.to_string(),
url,
headers,
content_type: None,
content_type,
content_encoding: None,
body: RequestBody {
json_body: None,
body_bytes_b64: None,
body_ref: None,
},
body,
stream: false,
client_api_format: provider_api_format.clone(),
client_api_format,
provider_api_format,
model_name: None,
model_name,
proxy: runtime.resolve_model_fetch_proxy(transport).await,
tls_profile: resolve_transport_tls_profile(transport),
timeouts: resolve_transport_execution_timeouts(transport),
})
}
async fn resolve_models_fetch_auth(
async fn resolve_standard_header_auth(
runtime: &(impl ModelFetchTransportRuntime + ?Sized),
transport: &GatewayProviderTransportSnapshot,
) -> Result<Option<(String, String)>, String> {
if transport.key.auth_type.trim().eq_ignore_ascii_case("oauth")
|| transport.key.auth_type.trim().eq_ignore_ascii_case("kiro")
{
return match runtime.resolve_local_oauth_request_auth(transport).await {
Ok(Some(LocalResolvedOAuthRequestAuth::Header { name, value })) => {
Ok(Some((name, value)))
}
Ok(Some(LocalResolvedOAuthRequestAuth::Kiro(_))) => Ok(None),
Ok(None) => Ok(None),
Err(err) => Err(err),
};
return resolve_oauth_header_auth(runtime, transport).await;
}
if let Some(auth) = resolve_local_openai_chat_auth(transport) {
return Ok(Some(auth));
let api_format = transport.endpoint.api_format.trim().to_ascii_lowercase();
if api_format.starts_with("openai:") {
return Ok(resolve_local_openai_chat_auth(transport));
}
if let Some(auth) = resolve_local_standard_auth(transport) {
return Ok(Some(auth));
}
if let Some(auth) = resolve_local_gemini_auth(transport) {
return Ok(Some(auth));
}
if let Some(query_auth) = resolve_local_vertex_api_key_query_auth(transport) {
let url = build_passthrough_path_url(
&transport.endpoint.base_url,
"/v1/publishers/google/models",
Some(&format!("{}={}", query_auth.name, query_auth.value)),
&[],
);
if url.is_some() {
return Ok(None);
}
if api_format.starts_with("claude:") {
return Ok(resolve_local_standard_auth(transport));
}
Ok(None)
}
async fn resolve_oauth_header_auth(
runtime: &(impl ModelFetchTransportRuntime + ?Sized),
transport: &GatewayProviderTransportSnapshot,
) -> Result<Option<(String, String)>, String> {
match runtime.resolve_local_oauth_request_auth(transport).await {
Ok(Some(LocalResolvedOAuthRequestAuth::Header { name, value })) => Ok(Some((name, value))),
Ok(Some(LocalResolvedOAuthRequestAuth::Kiro(_))) => Ok(None),
Ok(None) => Ok(None),
Err(err) => Err(err),
}
}
async fn resolve_bearer_or_oauth_header_auth(
runtime: &(impl ModelFetchTransportRuntime + ?Sized),
transport: &GatewayProviderTransportSnapshot,
) -> Result<Option<(String, String)>, String> {
if let Some(auth) = resolve_oauth_header_auth(runtime, transport).await? {
return Ok(Some(auth));
}
if let Some((name, value)) = resolve_local_openai_chat_auth(transport) {
return Ok(Some((name, value)));
}
if transport
.key
.auth_type
.trim()
.eq_ignore_ascii_case("bearer")
{
let secret = transport.key.decrypted_api_key.trim();
if !secret.is_empty() {
return Ok(Some((
"authorization".to_string(),
format!("Bearer {secret}"),
)));
}
}
Ok(None)
}
fn apply_fetch_header_rules(
transport: &GatewayProviderTransportSnapshot,
mut headers: BTreeMap<String, String>,
protected_headers: &[String],
) -> Result<BTreeMap<String, String>, String> {
let protected = protected_headers
.iter()
.map(String::as_str)
.collect::<Vec<_>>();
if !apply_local_header_rules(
&mut headers,
transport.endpoint.header_rules.as_ref(),
&protected,
&json!({}),
None,
) {
return Err("Endpoint header_rules application failed".to_string());
}
Ok(headers)
}
fn standard_models_fetch_headers(api_format: &str) -> BTreeMap<String, String> {
let api_format = api_format.trim().to_ascii_lowercase();
match api_format.as_str() {
"openai:cli" | "openai:compact" => {
BTreeMap::from([("user-agent".to_string(), OPENAI_CLI_USER_AGENT.to_string())])
}
"claude:chat" => BTreeMap::from([(
"anthropic-version".to_string(),
CLAUDE_VERSION_HEADER.to_string(),
)]),
"claude:cli" => BTreeMap::from([
("user-agent".to_string(), CLAUDE_CLI_USER_AGENT.to_string()),
(
"anthropic-version".to_string(),
CLAUDE_VERSION_HEADER.to_string(),
),
]),
"gemini:chat" => BROWSER_FINGERPRINT_HEADERS
.iter()
.map(|(key, value)| (key.to_string(), value.to_string()))
.collect(),
"gemini:cli" => {
let mut headers = BROWSER_FINGERPRINT_HEADERS
.iter()
.map(|(key, value)| (key.to_string(), value.to_string()))
.collect::<BTreeMap<_, _>>();
headers.insert("user-agent".to_string(), GEMINI_CLI_USER_AGENT.to_string());
headers
}
_ => BTreeMap::new(),
}
}
fn build_standard_models_fetch_url(
transport: &GatewayProviderTransportSnapshot,
after_id: Option<&str>,
) -> Result<String, String> {
let api_format = transport.endpoint.api_format.trim().to_ascii_lowercase();
if api_format.starts_with("gemini:") {
let secret = resolve_local_vertex_api_key_query_auth(transport)
.map(|auth| auth.value)
.or_else(|| {
resolve_local_gemini_auth(transport).and_then(|(name, value)| {
name.eq_ignore_ascii_case("x-goog-api-key").then_some(value)
})
})
.or_else(|| {
let secret = transport.key.decrypted_api_key.trim();
(!secret.is_empty()).then_some(secret.to_string())
})
.ok_or_else(|| "Gemini models fetch requires an API key".to_string())?;
let (url, _) = build_models_fetch_url(
&transport.provider.provider_type,
&transport.endpoint.api_format,
&transport.endpoint.base_url,
)
.ok_or_else(|| "Rust models fetch does not support this provider format yet".to_string())?;
return Ok(append_query_param(url, "key", &secret));
}
let (mut url, _) = build_models_fetch_url(
&transport.provider.provider_type,
&transport.endpoint.api_format,
&transport.endpoint.base_url,
)
.ok_or_else(|| "Rust models fetch does not support this provider format yet".to_string())?;
if api_format.starts_with("claude:") {
url = append_query_param(url, "limit", "100");
if let Some(after_id) = after_id.map(str::trim).filter(|value| !value.is_empty()) {
url = append_query_param(url, "after_id", after_id);
}
}
Ok(url)
}
fn append_query_param(mut url: String, key: &str, value: &str) -> String {
if key.trim().is_empty() || value.trim().is_empty() {
return url;
}
let separator = if url.contains('?') { '&' } else { '?' };
url.push(separator);
url.push_str(key.trim());
url.push('=');
url.push_str(value.trim());
url
}
#[cfg(test)]
mod tests {
use aether_contracts::ProxySnapshot;
use aether_contracts::{ExecutionPlan, ExecutionResult, ProxySnapshot};
use aether_provider_transport::snapshot::{
GatewayProviderTransportEndpoint, GatewayProviderTransportKey,
GatewayProviderTransportProvider, GatewayProviderTransportSnapshot,
};
use async_trait::async_trait;
use serde_json::json;
use super::{build_models_fetch_execution_plan, ModelFetchTransportRuntime};
use super::{
build_antigravity_fetch_available_models_plan, build_gemini_cli_load_code_assist_plan,
build_models_fetch_execution_plan, build_standard_models_fetch_execution_plan,
build_vertex_models_fetch_execution_plan, ModelFetchTransportRuntime,
};
struct TestRuntime {
oauth_auth: Option<aether_provider_transport::LocalResolvedOAuthRequestAuth>,
@@ -155,14 +484,25 @@ mod tests {
) -> Option<ProxySnapshot> {
self.proxy.clone()
}
async fn execute_model_fetch_execution_plan(
&self,
_plan: &ExecutionPlan,
) -> Result<ExecutionResult, String> {
unreachable!("tests only validate plan construction")
}
}
fn sample_transport(api_format: &str, auth_type: &str) -> GatewayProviderTransportSnapshot {
fn sample_transport(
provider_type: &str,
api_format: &str,
auth_type: &str,
) -> GatewayProviderTransportSnapshot {
GatewayProviderTransportSnapshot {
provider: GatewayProviderTransportProvider {
id: "provider-1".to_string(),
name: "Provider One".to_string(),
provider_type: "openai".to_string(),
provider_type: provider_type.to_string(),
website: None,
is_active: true,
keep_priority_on_conversion: false,
@@ -196,7 +536,7 @@ mod tests {
name: "key".to_string(),
auth_type: auth_type.to_string(),
is_active: true,
api_formats: None,
api_formats: Some(vec![api_format.to_string()]),
allowed_models: None,
capabilities: None,
rate_multipliers: None,
@@ -205,35 +545,86 @@ mod tests {
proxy: None,
fingerprint: None,
decrypted_api_key: "secret".to_string(),
decrypted_auth_config: None,
decrypted_auth_config: Some(
r#"{"project_id":"project-1","client_version":"1.2.3","session_id":"sess-1"}"#
.to_string(),
),
},
}
}
#[tokio::test]
async fn builds_openai_models_fetch_plan_from_transport_snapshot() {
async fn builds_openai_cli_models_fetch_plan_with_cli_user_agent() {
let runtime = TestRuntime {
oauth_auth: None,
proxy: None,
};
let plan = build_models_fetch_execution_plan(
&runtime,
&sample_transport("openai:chat", "api_key"),
)
.await
.expect("plan");
let mut transport = sample_transport("openai", "openai:cli", "api_key");
transport.key.decrypted_auth_config = None;
let plan = build_models_fetch_execution_plan(&runtime, &transport)
.await
.expect("plan");
assert_eq!(plan.method, "GET");
assert_eq!(plan.url, "https://example.com/v1/models");
assert_eq!(
plan.headers.get("user-agent").map(String::as_str),
Some("openai-codex/1.0")
);
assert_eq!(
plan.headers.get("authorization").map(String::as_str),
Some("Bearer secret")
);
assert_eq!(plan.provider_api_format, "openai:chat");
}
#[tokio::test]
async fn builds_oauth_models_fetch_plan_from_runtime_auth() {
async fn builds_claude_models_fetch_plan_with_pagination() {
let runtime = TestRuntime {
oauth_auth: None,
proxy: None,
};
let mut transport = sample_transport("custom", "claude:chat", "api_key");
transport.key.decrypted_auth_config = None;
let plan =
build_standard_models_fetch_execution_plan(&runtime, &transport, Some("cursor-1"))
.await
.expect("plan");
assert_eq!(
plan.url,
"https://example.com/v1/models?limit=100&after_id=cursor-1"
);
assert_eq!(
plan.headers.get("anthropic-version").map(String::as_str),
Some("2023-06-01")
);
assert_eq!(
plan.headers.get("x-api-key").map(String::as_str),
Some("secret")
);
}
#[tokio::test]
async fn builds_gemini_models_fetch_plan_with_browser_headers_and_query_auth() {
let runtime = TestRuntime {
oauth_auth: None,
proxy: None,
};
let mut transport = sample_transport("custom", "gemini:chat", "api_key");
transport.key.decrypted_auth_config = None;
let plan = build_models_fetch_execution_plan(&runtime, &transport)
.await
.expect("plan");
assert_eq!(plan.url, "https://example.com/v1beta/models?key=secret");
assert!(plan.headers.contains_key("sec-ch-ua"));
assert_eq!(
plan.headers.get("accept-language").map(String::as_str),
Some("zh-CN")
);
}
#[tokio::test]
async fn builds_antigravity_fetch_available_models_plan() {
let runtime = TestRuntime {
oauth_auth: Some(
aether_provider_transport::LocalResolvedOAuthRequestAuth::Header {
@@ -241,27 +632,85 @@ mod tests {
value: "Bearer oauth-token".to_string(),
},
),
proxy: Some(ProxySnapshot {
enabled: Some(true),
mode: Some("fixed".to_string()),
node_id: None,
label: None,
url: Some("http://proxy.internal".to_string()),
extra: None,
}),
proxy: None,
};
let plan =
build_models_fetch_execution_plan(&runtime, &sample_transport("openai:chat", "oauth"))
.await
.expect("plan");
let transport = sample_transport("antigravity", "gemini:chat", "oauth");
let plan = build_antigravity_fetch_available_models_plan(
&runtime,
&transport,
"https://daily-cloudcode-pa.sandbox.googleapis.com",
"project-1",
)
.await
.expect("plan");
assert_eq!(plan.method, "POST");
assert_eq!(
plan.headers.get("authorization").map(String::as_str),
Some("Bearer oauth-token")
plan.url,
"https://daily-cloudcode-pa.sandbox.googleapis.com/v1internal:fetchAvailableModels"
);
assert_eq!(
plan.proxy.as_ref().and_then(|proxy| proxy.url.as_deref()),
Some("http://proxy.internal")
plan.provider_api_format,
"antigravity:fetch_available_models"
);
assert_eq!(
plan.body
.json_body
.as_ref()
.and_then(|value| value.get("project")),
Some(&json!("project-1"))
);
}
#[tokio::test]
async fn builds_gemini_cli_load_code_assist_plan() {
let runtime = TestRuntime {
oauth_auth: Some(
aether_provider_transport::LocalResolvedOAuthRequestAuth::Header {
name: "authorization".to_string(),
value: "Bearer oauth-token".to_string(),
},
),
proxy: None,
};
let transport = sample_transport("gemini_cli", "gemini:cli", "oauth");
let plan = build_gemini_cli_load_code_assist_plan(&runtime, &transport)
.await
.expect("plan");
assert_eq!(plan.method, "POST");
assert_eq!(
plan.url,
"https://cloudcode-pa.googleapis.com/v1internal:loadCodeAssist"
);
assert_eq!(
plan.headers.get("user-agent").map(String::as_str),
Some("GeminiCLI/0.1.5 (Windows; AMD64)")
);
}
#[tokio::test]
async fn builds_vertex_models_fetch_plan_with_auth_override() {
let runtime = TestRuntime {
oauth_auth: None,
proxy: None,
};
let mut transport = sample_transport("vertex_ai", "claude:chat", "api_key");
transport.key.decrypted_auth_config = None;
let plan = build_vertex_models_fetch_execution_plan(
&runtime,
&transport,
"https://aiplatform.googleapis.com/v1/publishers/google/models?key=secret",
"gemini:chat",
None,
)
.await
.expect("plan");
assert_eq!(
plan.url,
"https://aiplatform.googleapis.com/v1/publishers/google/models?key=secret"
);
assert!(plan.headers.contains_key("sec-ch-ua"));
}
}