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

将 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

View File

@@ -16,12 +16,12 @@ use aether_data_contracts::repository::provider_catalog::{
StoredProviderCatalogEndpoint, StoredProviderCatalogKey, StoredProviderCatalogProvider,
};
use aether_model_fetch::{
aggregate_models_for_cache, build_models_fetch_execution_plan,
endpoint_supports_rust_models_fetch, extract_error_message, parse_models_response,
aggregate_models_for_cache, fetch_models_from_transports, json_string_list,
preset_models_for_provider, selected_models_fetch_endpoints,
};
use axum::{body::Body, http::Response, response::IntoResponse, Json};
use serde_json::{json, Value};
use std::collections::{BTreeMap, BTreeSet};
use std::collections::BTreeSet;
pub(crate) const ADMIN_PROVIDER_QUERY_LOCAL_TEST_MODEL_MESSAGE: &str =
"Rust local provider-query model test is not configured";
@@ -31,22 +31,15 @@ const ADMIN_PROVIDER_QUERY_NO_ACTIVE_ENDPOINT_DETAIL: &str =
"No active endpoints found for this provider";
const ADMIN_PROVIDER_QUERY_NO_MODELS_FROM_ENDPOINT_DETAIL: &str =
"No models returned from any endpoint";
const PROVIDER_QUERY_FETCH_FORMAT_PRIORITY: &[&[&str]] = &[
&[
"openai:chat",
"openai:responses",
"openai:cli",
"openai:compact",
],
&["claude:chat", "claude:cli"],
&["gemini:chat", "gemini:cli"],
];
const ADMIN_PROVIDER_QUERY_NO_MODELS_FROM_KEY_DETAIL: &str = "No models returned from any key";
const ANTIGRAVITY_PROVIDER_CACHE_KEY_PREFIX: &str = "upstream_models_provider:";
#[derive(Debug)]
struct ProviderQueryKeyFetchResult {
models: Vec<Value>,
error: Option<String>,
from_cache: bool,
has_success: bool,
}
fn provider_query_provider_payload(provider: &StoredProviderCatalogProvider) -> Value {
@@ -66,64 +59,6 @@ fn provider_query_key_display_name(key: &StoredProviderCatalogKey) -> String {
}
}
fn provider_query_normalize_api_format(value: &str) -> String {
value.trim().to_ascii_lowercase()
}
fn provider_query_selected_fetch_endpoints(
endpoints: &[StoredProviderCatalogEndpoint],
key: &StoredProviderCatalogKey,
) -> Vec<StoredProviderCatalogEndpoint> {
let allowed_api_formats = key
.api_formats
.as_ref()
.and_then(Value::as_array)
.map(|items| {
items
.iter()
.filter_map(Value::as_str)
.map(provider_query_normalize_api_format)
.filter(|value| !value.is_empty())
.collect::<BTreeSet<_>>()
})
.filter(|items| !items.is_empty());
let mut by_format = BTreeMap::<String, StoredProviderCatalogEndpoint>::new();
for endpoint in endpoints.iter().filter(|endpoint| endpoint.is_active) {
let api_format = provider_query_normalize_api_format(&endpoint.api_format);
if api_format.is_empty() || !endpoint_supports_rust_models_fetch(&api_format) {
continue;
}
if allowed_api_formats
.as_ref()
.is_some_and(|formats| !formats.contains(&api_format))
{
continue;
}
by_format.insert(api_format, endpoint.clone());
}
// 与 Python 版本保持一致:同族优先使用 chat 端点,其次才回退到其他抓取格式。
let covered_formats = PROVIDER_QUERY_FETCH_FORMAT_PRIORITY
.iter()
.flat_map(|items| items.iter().copied())
.collect::<BTreeSet<_>>();
let mut selected = PROVIDER_QUERY_FETCH_FORMAT_PRIORITY
.iter()
.filter_map(|candidates| {
candidates
.iter()
.find_map(|api_format| by_format.remove(*api_format))
})
.collect::<Vec<_>>();
selected.extend(
by_format
.into_iter()
.filter(|(api_format, _)| !covered_formats.contains(api_format.as_str()))
.map(|(_, endpoint)| endpoint),
);
selected
}
async fn provider_query_read_cached_models(
state: &AdminAppState<'_>,
provider_id: &str,
@@ -147,38 +82,97 @@ async fn provider_query_read_cached_models(
Some(aggregate_models_for_cache(&parsed))
}
async fn provider_query_fetch_models_from_transport(
async fn provider_query_read_provider_cached_models(
state: &AdminAppState<'_>,
transport: &crate::provider_transport::GatewayProviderTransportSnapshot,
) -> Result<Vec<Value>, String> {
let plan = build_models_fetch_execution_plan(state.app(), transport).await?;
let result = execution_runtime::execute_execution_runtime_sync_plan(state.app(), None, &plan)
provider_id: &str,
) -> Option<Vec<Value>> {
let runner = state.app().redis_kv_runner()?;
let cache_key = runner.keyspace().key(&format!(
"{ANTIGRAVITY_PROVIDER_CACHE_KEY_PREFIX}{provider_id}"
));
let mut connection = runner
.client()
.get_multiplexed_async_connection()
.await
.map_err(|err| format!("{err:?}"))?;
.ok()?;
let raw = redis::cmd("GET")
.arg(&cache_key)
.query_async::<Option<String>>(&mut connection)
.await
.ok()??;
let parsed = serde_json::from_str::<Vec<Value>>(&raw).ok()?;
Some(aggregate_models_for_cache(&parsed))
}
if result.status_code != 200 {
let message = result
.body
.as_ref()
.and_then(|body| body.json_body.as_ref())
.and_then(extract_error_message)
.or_else(|| {
result.error.as_ref().and_then(|error| {
let message = error.message.trim();
(!message.is_empty()).then_some(message.to_string())
async fn provider_query_write_provider_cached_models(
state: &AdminAppState<'_>,
provider_id: &str,
models: &[Value],
) {
let Some(runner) = state.app().redis_kv_runner() else {
return;
};
let Ok(serialized) = serde_json::to_string(&aggregate_models_for_cache(models)) else {
return;
};
let cache_key = format!("{ANTIGRAVITY_PROVIDER_CACHE_KEY_PREFIX}{provider_id}");
let _ = runner
.setex(
&cache_key,
&serialized,
Some(aether_model_fetch::model_fetch_interval_minutes().saturating_mul(60)),
)
.await;
}
fn provider_query_antigravity_tier_weight(raw_auth_config: Option<&str>) -> i32 {
raw_auth_config
.and_then(|value| serde_json::from_str::<Value>(value).ok())
.and_then(|value| value.get("tier").cloned())
.and_then(|value| value.as_str().map(ToOwned::to_owned))
.map(|tier| match tier.trim().to_ascii_lowercase().as_str() {
"ultra" => 3,
"pro" => 2,
"free" => 1,
_ => 0,
})
.unwrap_or(0)
}
async fn provider_query_sort_antigravity_keys(
state: &AdminAppState<'_>,
provider: &StoredProviderCatalogProvider,
endpoints: &[StoredProviderCatalogEndpoint],
keys: Vec<StoredProviderCatalogKey>,
) -> Result<Vec<StoredProviderCatalogKey>, GatewayError> {
let mut ranked = Vec::new();
for key in keys {
let availability = if key.oauth_invalid_at_unix_secs.is_some() {
0
} else {
1
};
let tier_weight = if let Some(endpoint) = selected_models_fetch_endpoints(endpoints, &key)
.into_iter()
.next()
{
state
.app()
.read_provider_transport_snapshot(&provider.id, &endpoint.id, &key.id)
.await?
.map(|transport| {
provider_query_antigravity_tier_weight(
transport.key.decrypted_auth_config.as_deref(),
)
})
})
.unwrap_or_else(|| format!("upstream returned status {}", result.status_code));
return Err(message);
.unwrap_or(0)
} else {
0
};
ranked.push(((availability, tier_weight), key));
}
let body_json = result
.body
.as_ref()
.and_then(|body| body.json_body.as_ref())
.ok_or_else(|| "models fetch response body is missing JSON payload".to_string())?;
let parsed = parse_models_response(&transport.endpoint.api_format, body_json)?;
Ok(parsed.cached_models)
ranked.sort_by(|left, right| right.0.cmp(&left.0));
Ok(ranked.into_iter().map(|(_, key)| key).collect())
}
async fn provider_query_fetch_models_for_key(
@@ -196,20 +190,30 @@ async fn provider_query_fetch_models_for_key(
models: cached_models,
error: None,
from_cache: true,
has_success: true,
});
}
}
let selected_endpoints = provider_query_selected_fetch_endpoints(endpoints, key);
let selected_endpoints = selected_models_fetch_endpoints(endpoints, key);
if selected_endpoints.is_empty() {
if let Some(models) = preset_models_for_provider(&provider.provider_type) {
return Ok(ProviderQueryKeyFetchResult {
models: aggregate_models_for_cache(&models),
error: None,
from_cache: false,
has_success: true,
});
}
return Ok(ProviderQueryKeyFetchResult {
models: Vec::new(),
error: Some(ADMIN_PROVIDER_QUERY_NO_ACTIVE_ENDPOINT_DETAIL.to_string()),
from_cache: false,
has_success: false,
});
}
let mut all_models = Vec::new();
let mut transports = Vec::new();
let mut all_errors = Vec::new();
for endpoint in selected_endpoints {
let Some(transport) = state
@@ -223,14 +227,34 @@ async fn provider_query_fetch_models_for_key(
));
continue;
};
match provider_query_fetch_models_from_transport(state, &transport).await {
Ok(models) => all_models.extend(models),
Err(err) => all_errors.push(err),
}
transports.push(transport);
}
let unique_models = aggregate_models_for_cache(&all_models);
if !unique_models.is_empty() {
if transports.is_empty() {
return Ok(ProviderQueryKeyFetchResult {
models: Vec::new(),
error: Some(all_errors.join("; ")),
from_cache: false,
has_success: false,
});
}
let outcome = match fetch_models_from_transports(state.app(), &transports).await {
Ok(outcome) => outcome,
Err(err) => {
all_errors.push(err);
return Ok(ProviderQueryKeyFetchResult {
models: Vec::new(),
error: Some(all_errors.join("; ")),
from_cache: false,
has_success: false,
});
}
};
all_errors.extend(outcome.errors);
let unique_models = aggregate_models_for_cache(&outcome.cached_models);
if outcome.has_success && !unique_models.is_empty() {
<AppState as ModelFetchRuntimeState>::write_upstream_models_cache(
state.app(),
&provider.id,
@@ -253,6 +277,7 @@ async fn provider_query_fetch_models_for_key(
models: unique_models,
error,
from_cache: false,
has_success: outcome.has_success,
})
}
@@ -317,7 +342,10 @@ pub(crate) async fn build_admin_provider_query_models_response(
.into_response());
}
let active_keys = keys.iter().filter(|key| key.is_active).collect::<Vec<_>>();
let active_keys = keys
.into_iter()
.filter(|key| key.is_active)
.collect::<Vec<_>>();
if active_keys.is_empty() {
return Ok(build_admin_provider_query_bad_request_response(
ADMIN_PROVIDER_QUERY_NO_ACTIVE_API_KEY_DETAIL,
@@ -325,11 +353,45 @@ pub(crate) async fn build_admin_provider_query_models_response(
}
let active_key_count = active_keys.len();
if provider
.provider_type
.trim()
.eq_ignore_ascii_case("antigravity")
&& !force_refresh
{
if let Some(models) = provider_query_read_provider_cached_models(state, &provider.id).await
{
return Ok(Json(json!({
"success": !models.is_empty(),
"data": {
"models": models,
"error": serde_json::Value::Null,
"from_cache": true,
"keys_total": active_key_count,
"keys_cached": active_key_count,
"keys_fetched": 0,
},
"provider": provider_query_provider_payload(&provider),
}))
.into_response());
}
}
let ordered_keys = if provider
.provider_type
.trim()
.eq_ignore_ascii_case("antigravity")
{
provider_query_sort_antigravity_keys(state, &provider, &endpoints, active_keys).await?
} else {
active_keys
};
let mut all_models = Vec::new();
let mut all_errors = Vec::new();
let mut cache_hit_count = 0usize;
let mut fetch_count = 0usize;
for key in active_keys {
for key in &ordered_keys {
let result =
provider_query_fetch_models_for_key(state, &provider, &endpoints, key, force_refresh)
.await?;
@@ -346,9 +408,25 @@ pub(crate) async fn build_admin_provider_query_models_response(
} else {
fetch_count += 1;
}
if provider
.provider_type
.trim()
.eq_ignore_ascii_case("antigravity")
&& result.has_success
{
break;
}
}
let models = aggregate_models_for_cache(&all_models);
if provider
.provider_type
.trim()
.eq_ignore_ascii_case("antigravity")
&& !models.is_empty()
{
provider_query_write_provider_cached_models(state, &provider.id, &models).await;
}
let success = !models.is_empty();
let mut error = if all_errors.is_empty() {
None
@@ -356,7 +434,7 @@ pub(crate) async fn build_admin_provider_query_models_response(
Some(all_errors.join("; "))
};
if !success && error.is_none() {
error = Some("No models returned from any key".to_string());
error = Some(ADMIN_PROVIDER_QUERY_NO_MODELS_FROM_KEY_DETAIL.to_string());
}
Ok(Json(json!({