i/Aether
1
0
Fork 0
mirror of https://github.com/fawney19/Aether.git synced 2026-09-03 01:40:21 +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

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"));
}
}