mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-02 09:20:22 +08:00
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
@@ -2,21 +2,18 @@ use std::collections::HashMap;
|
||||
use std::time::Duration;
|
||||
use std::time::{SystemTime, UNIX_EPOCH};
|
||||
|
||||
use aether_contracts::ExecutionResult;
|
||||
use aether_data_contracts::repository::provider_catalog::{
|
||||
StoredProviderCatalogEndpoint, StoredProviderCatalogKey, StoredProviderCatalogProvider,
|
||||
};
|
||||
use aether_model_fetch::{
|
||||
apply_model_filters, build_models_fetch_execution_plan, extract_error_message,
|
||||
json_string_list, model_fetch_interval_minutes, model_fetch_startup_delay_seconds,
|
||||
model_fetch_startup_enabled, parse_models_response, select_models_fetch_endpoint,
|
||||
apply_model_filters, fetch_models_from_transports, json_string_list, merge_upstream_metadata,
|
||||
model_fetch_interval_minutes, model_fetch_startup_delay_seconds, model_fetch_startup_enabled,
|
||||
preset_models_for_provider, selected_models_fetch_endpoints,
|
||||
sync_provider_model_whitelist_associations, ModelFetchAssociationStore, ModelFetchRunSummary,
|
||||
ModelsFetchSuccess,
|
||||
};
|
||||
use serde_json::{json, Value};
|
||||
use tracing::{debug, info, warn};
|
||||
|
||||
use crate::provider_transport::GatewayProviderTransportSnapshot;
|
||||
use crate::{AppState, GatewayError};
|
||||
|
||||
pub(crate) mod state;
|
||||
@@ -26,8 +23,8 @@ use self::state::ModelFetchRuntimeState;
|
||||
#[derive(Debug, Clone)]
|
||||
struct SelectedFetchTarget {
|
||||
provider: StoredProviderCatalogProvider,
|
||||
endpoint: StoredProviderCatalogEndpoint,
|
||||
key: StoredProviderCatalogKey,
|
||||
endpoints: Vec<StoredProviderCatalogEndpoint>,
|
||||
}
|
||||
|
||||
pub(crate) fn spawn_model_fetch_worker(state: AppState) -> Option<tokio::task::JoinHandle<()>> {
|
||||
@@ -131,38 +128,12 @@ where
|
||||
if !key.is_active || !key.auto_fetch_models {
|
||||
continue;
|
||||
}
|
||||
if let Some(endpoint) = select_models_fetch_endpoint(&endpoints, &key) {
|
||||
targets.push(SelectedFetchTarget {
|
||||
provider: provider.clone(),
|
||||
endpoint,
|
||||
key,
|
||||
});
|
||||
} else {
|
||||
targets.push(SelectedFetchTarget {
|
||||
provider: provider.clone(),
|
||||
endpoint: StoredProviderCatalogEndpoint::new(
|
||||
"__unsupported__".to_string(),
|
||||
provider.id.clone(),
|
||||
"__unsupported__".to_string(),
|
||||
None,
|
||||
None,
|
||||
false,
|
||||
)
|
||||
.expect("unsupported sentinel endpoint should build")
|
||||
.with_transport_fields(
|
||||
"https://unsupported.invalid".to_string(),
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
)
|
||||
.expect("unsupported sentinel endpoint transport should build"),
|
||||
key,
|
||||
});
|
||||
}
|
||||
let selected_endpoints = selected_models_fetch_endpoints(&endpoints, &key);
|
||||
targets.push(SelectedFetchTarget {
|
||||
provider: provider.clone(),
|
||||
key,
|
||||
endpoints: selected_endpoints,
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
@@ -215,7 +186,34 @@ async fn fetch_and_persist_key_models(
|
||||
target: &SelectedFetchTarget,
|
||||
) -> Result<KeyFetchDisposition, GatewayError> {
|
||||
let now_unix_secs = now_unix_secs();
|
||||
if target.endpoint.api_format == "__unsupported__" {
|
||||
if target.endpoints.is_empty() {
|
||||
if let Some(models) = preset_models_for_provider(&target.provider.provider_type) {
|
||||
let fetched_model_ids = models
|
||||
.iter()
|
||||
.filter_map(|model| model.get("id"))
|
||||
.filter_map(Value::as_str)
|
||||
.map(ToOwned::to_owned)
|
||||
.collect::<Vec<_>>();
|
||||
let filtered_models = apply_model_filters(
|
||||
&fetched_model_ids,
|
||||
json_string_list(target.key.locked_models.as_ref()),
|
||||
json_string_list(target.key.model_include_patterns.as_ref()),
|
||||
json_string_list(target.key.model_exclude_patterns.as_ref()),
|
||||
);
|
||||
persist_key_fetch_success(state, &target.key, now_unix_secs, &filtered_models, None)
|
||||
.await?;
|
||||
state
|
||||
.write_upstream_models_cache(&target.provider.id, &target.key.id, &models)
|
||||
.await;
|
||||
sync_provider_model_whitelist_associations(
|
||||
state,
|
||||
&target.provider.id,
|
||||
&filtered_models,
|
||||
)
|
||||
.await
|
||||
.map_err(GatewayError::Internal)?;
|
||||
return Ok(KeyFetchDisposition::Succeeded);
|
||||
}
|
||||
persist_key_fetch_failure(
|
||||
state,
|
||||
&target.key,
|
||||
@@ -226,10 +224,25 @@ async fn fetch_and_persist_key_models(
|
||||
return Ok(KeyFetchDisposition::Skipped);
|
||||
}
|
||||
|
||||
let Some(transport) = state
|
||||
.read_provider_transport_snapshot(&target.provider.id, &target.endpoint.id, &target.key.id)
|
||||
.await?
|
||||
else {
|
||||
let mut transports = Vec::new();
|
||||
for endpoint in &target.endpoints {
|
||||
match state
|
||||
.read_provider_transport_snapshot(&target.provider.id, &endpoint.id, &target.key.id)
|
||||
.await?
|
||||
{
|
||||
Some(transport) => transports.push(transport),
|
||||
None => {
|
||||
warn!(
|
||||
provider_id = %target.provider.id,
|
||||
endpoint_id = %endpoint.id,
|
||||
key_id = %target.key.id,
|
||||
"gateway model fetch transport snapshot unavailable"
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if transports.is_empty() {
|
||||
persist_key_fetch_failure(
|
||||
state,
|
||||
&target.key,
|
||||
@@ -238,9 +251,9 @@ async fn fetch_and_persist_key_models(
|
||||
)
|
||||
.await?;
|
||||
return Ok(KeyFetchDisposition::Skipped);
|
||||
};
|
||||
}
|
||||
|
||||
let result = match execute_models_fetch_request(state, &transport).await {
|
||||
let result = match fetch_models_from_transports(state, &transports).await {
|
||||
Ok(result) => result,
|
||||
Err(err) => {
|
||||
persist_key_fetch_failure(state, &target.key, now_unix_secs, err.clone()).await?;
|
||||
@@ -254,6 +267,22 @@ async fn fetch_and_persist_key_models(
|
||||
}
|
||||
};
|
||||
|
||||
if !result.has_success {
|
||||
let error = if result.errors.is_empty() {
|
||||
"Upstream models fetch failed".to_string()
|
||||
} else {
|
||||
result.errors.join("; ")
|
||||
};
|
||||
persist_key_fetch_failure(state, &target.key, now_unix_secs, error.clone()).await?;
|
||||
warn!(
|
||||
provider_id = %target.provider.id,
|
||||
key_id = %target.key.id,
|
||||
message = %error,
|
||||
"gateway model fetch failed"
|
||||
);
|
||||
return Ok(KeyFetchDisposition::Failed);
|
||||
}
|
||||
|
||||
let filtered_models = apply_model_filters(
|
||||
&result.fetched_model_ids,
|
||||
json_string_list(target.key.locked_models.as_ref()),
|
||||
@@ -261,7 +290,14 @@ async fn fetch_and_persist_key_models(
|
||||
json_string_list(target.key.model_exclude_patterns.as_ref()),
|
||||
);
|
||||
|
||||
persist_key_fetch_success(state, &target.key, now_unix_secs, &filtered_models).await?;
|
||||
persist_key_fetch_success(
|
||||
state,
|
||||
&target.key,
|
||||
now_unix_secs,
|
||||
&filtered_models,
|
||||
result.upstream_metadata.as_ref(),
|
||||
)
|
||||
.await?;
|
||||
state
|
||||
.write_upstream_models_cache(&target.provider.id, &target.key.id, &result.cached_models)
|
||||
.await;
|
||||
@@ -271,41 +307,6 @@ async fn fetch_and_persist_key_models(
|
||||
Ok(KeyFetchDisposition::Succeeded)
|
||||
}
|
||||
|
||||
async fn execute_models_fetch_request(
|
||||
state: &(impl ModelFetchRuntimeState + ?Sized),
|
||||
transport: &GatewayProviderTransportSnapshot,
|
||||
) -> Result<ModelsFetchSuccess, String> {
|
||||
let plan = build_models_fetch_execution_plan(state, transport).await?;
|
||||
|
||||
let result = state
|
||||
.execute_execution_runtime_sync_plan(&plan)
|
||||
.await
|
||||
.map_err(|err| format!("{err:?}"))?;
|
||||
|
||||
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())
|
||||
})
|
||||
})
|
||||
.unwrap_or_else(|| format!("upstream returned status {}", result.status_code));
|
||||
return Err(message);
|
||||
}
|
||||
|
||||
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())?;
|
||||
parse_models_response(&transport.endpoint.api_format, body_json)
|
||||
}
|
||||
|
||||
async fn persist_key_fetch_failure(
|
||||
state: &(impl ModelFetchRuntimeState + ?Sized),
|
||||
key: &StoredProviderCatalogKey,
|
||||
@@ -325,6 +326,7 @@ async fn persist_key_fetch_success(
|
||||
key: &StoredProviderCatalogKey,
|
||||
now_unix_secs: u64,
|
||||
allowed_models: &[String],
|
||||
upstream_metadata: Option<&Value>,
|
||||
) -> Result<(), GatewayError> {
|
||||
let mut updated = key.clone();
|
||||
updated.allowed_models = if allowed_models.is_empty() {
|
||||
@@ -332,6 +334,12 @@ async fn persist_key_fetch_success(
|
||||
} else {
|
||||
Some(json!(allowed_models))
|
||||
};
|
||||
if let Some(upstream_metadata) = upstream_metadata {
|
||||
updated.upstream_metadata = Some(merge_upstream_metadata(
|
||||
updated.upstream_metadata.as_ref(),
|
||||
upstream_metadata,
|
||||
));
|
||||
}
|
||||
updated.last_models_fetch_at_unix_secs = Some(now_unix_secs);
|
||||
updated.last_models_fetch_error = None;
|
||||
updated.updated_at_unix_secs = Some(now_unix_secs);
|
||||
@@ -345,3 +353,533 @@ fn now_unix_secs() -> u64 {
|
||||
.unwrap_or_default()
|
||||
.as_secs()
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::{perform_model_fetch_once_with_state, state::ModelFetchRuntimeState};
|
||||
use aether_contracts::{ExecutionPlan, ExecutionResult, ProxySnapshot};
|
||||
use aether_data_contracts::repository::global_models::{
|
||||
AdminGlobalModelListQuery, AdminProviderModelListQuery, StoredAdminGlobalModelPage,
|
||||
StoredAdminProviderModel, UpsertAdminProviderModelRecord,
|
||||
};
|
||||
use aether_data_contracts::repository::provider_catalog::{
|
||||
StoredProviderCatalogEndpoint, StoredProviderCatalogKey, StoredProviderCatalogProvider,
|
||||
};
|
||||
use aether_model_fetch::{ModelFetchAssociationStore, ModelFetchTransportRuntime};
|
||||
use async_trait::async_trait;
|
||||
use serde_json::{json, Value};
|
||||
use std::collections::{HashMap, VecDeque};
|
||||
use std::sync::{Arc, Mutex};
|
||||
|
||||
use crate::provider_transport::LocalResolvedOAuthRequestAuth;
|
||||
use crate::GatewayError;
|
||||
use aether_provider_transport::snapshot::{
|
||||
GatewayProviderTransportEndpoint, GatewayProviderTransportKey,
|
||||
GatewayProviderTransportProvider, GatewayProviderTransportSnapshot,
|
||||
};
|
||||
|
||||
#[derive(Clone, Default)]
|
||||
struct TestState {
|
||||
providers: Arc<Vec<StoredProviderCatalogProvider>>,
|
||||
endpoints: Arc<Vec<StoredProviderCatalogEndpoint>>,
|
||||
keys: Arc<Mutex<Vec<StoredProviderCatalogKey>>>,
|
||||
transports: Arc<HashMap<(String, String, String), GatewayProviderTransportSnapshot>>,
|
||||
execution_results: Arc<Mutex<VecDeque<ExecutionResult>>>,
|
||||
cached_models: Arc<Mutex<HashMap<(String, String), Vec<Value>>>>,
|
||||
}
|
||||
|
||||
impl TestState {
|
||||
fn new(
|
||||
providers: Vec<StoredProviderCatalogProvider>,
|
||||
endpoints: Vec<StoredProviderCatalogEndpoint>,
|
||||
keys: Vec<StoredProviderCatalogKey>,
|
||||
transports: HashMap<(String, String, String), GatewayProviderTransportSnapshot>,
|
||||
execution_results: Vec<ExecutionResult>,
|
||||
) -> Self {
|
||||
Self {
|
||||
providers: Arc::new(providers),
|
||||
endpoints: Arc::new(endpoints),
|
||||
keys: Arc::new(Mutex::new(keys)),
|
||||
transports: Arc::new(transports),
|
||||
execution_results: Arc::new(Mutex::new(VecDeque::from(execution_results))),
|
||||
cached_models: Arc::new(Mutex::new(HashMap::new())),
|
||||
}
|
||||
}
|
||||
|
||||
fn key(&self, key_id: &str) -> StoredProviderCatalogKey {
|
||||
self.keys
|
||||
.lock()
|
||||
.expect("keys mutex")
|
||||
.iter()
|
||||
.find(|key| key.id == key_id)
|
||||
.cloned()
|
||||
.expect("key should exist")
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl ModelFetchTransportRuntime for TestState {
|
||||
async fn resolve_local_oauth_request_auth(
|
||||
&self,
|
||||
transport: &GatewayProviderTransportSnapshot,
|
||||
) -> Result<Option<LocalResolvedOAuthRequestAuth>, String> {
|
||||
if transport.key.auth_type.trim().eq_ignore_ascii_case("oauth") {
|
||||
return Ok(Some(LocalResolvedOAuthRequestAuth::Header {
|
||||
name: "authorization".to_string(),
|
||||
value: "Bearer oauth-token".to_string(),
|
||||
}));
|
||||
}
|
||||
Ok(None)
|
||||
}
|
||||
|
||||
async fn resolve_model_fetch_proxy(
|
||||
&self,
|
||||
_transport: &GatewayProviderTransportSnapshot,
|
||||
) -> Option<ProxySnapshot> {
|
||||
None
|
||||
}
|
||||
|
||||
async fn execute_model_fetch_execution_plan(
|
||||
&self,
|
||||
_plan: &ExecutionPlan,
|
||||
) -> Result<ExecutionResult, String> {
|
||||
self.execution_results
|
||||
.lock()
|
||||
.expect("execution result mutex")
|
||||
.pop_front()
|
||||
.ok_or_else(|| "missing execution result".to_string())
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl ModelFetchAssociationStore for TestState {
|
||||
type Error = String;
|
||||
|
||||
fn has_global_model_reader(&self) -> bool {
|
||||
false
|
||||
}
|
||||
|
||||
fn has_global_model_writer(&self) -> bool {
|
||||
false
|
||||
}
|
||||
|
||||
fn model_fetch_internal_error(&self, message: String) -> Self::Error {
|
||||
message
|
||||
}
|
||||
|
||||
async fn list_admin_provider_models(
|
||||
&self,
|
||||
_query: &AdminProviderModelListQuery,
|
||||
) -> Result<Vec<StoredAdminProviderModel>, Self::Error> {
|
||||
Ok(Vec::new())
|
||||
}
|
||||
|
||||
async fn list_admin_global_models(
|
||||
&self,
|
||||
_query: &AdminGlobalModelListQuery,
|
||||
) -> Result<StoredAdminGlobalModelPage, Self::Error> {
|
||||
Ok(StoredAdminGlobalModelPage {
|
||||
items: Vec::new(),
|
||||
total: 0,
|
||||
})
|
||||
}
|
||||
|
||||
async fn create_admin_provider_model(
|
||||
&self,
|
||||
_record: &UpsertAdminProviderModelRecord,
|
||||
) -> Result<Option<StoredAdminProviderModel>, Self::Error> {
|
||||
Ok(None)
|
||||
}
|
||||
|
||||
async fn list_provider_catalog_keys_by_provider_ids(
|
||||
&self,
|
||||
provider_ids: &[String],
|
||||
) -> Result<Vec<StoredProviderCatalogKey>, Self::Error> {
|
||||
Ok(self
|
||||
.keys
|
||||
.lock()
|
||||
.expect("keys mutex")
|
||||
.iter()
|
||||
.filter(|key| {
|
||||
provider_ids
|
||||
.iter()
|
||||
.any(|provider_id| provider_id == &key.provider_id)
|
||||
})
|
||||
.cloned()
|
||||
.collect())
|
||||
}
|
||||
|
||||
async fn delete_admin_provider_model(
|
||||
&self,
|
||||
_provider_id: &str,
|
||||
_model_id: &str,
|
||||
) -> Result<bool, Self::Error> {
|
||||
Ok(false)
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl ModelFetchRuntimeState for TestState {
|
||||
fn has_provider_catalog_data_reader(&self) -> bool {
|
||||
true
|
||||
}
|
||||
|
||||
fn has_provider_catalog_data_writer(&self) -> bool {
|
||||
true
|
||||
}
|
||||
|
||||
async fn list_provider_catalog_providers(
|
||||
&self,
|
||||
_active_only: bool,
|
||||
) -> Result<Vec<StoredProviderCatalogProvider>, GatewayError> {
|
||||
Ok(self.providers.as_ref().clone())
|
||||
}
|
||||
|
||||
async fn list_provider_catalog_endpoints_by_provider_ids(
|
||||
&self,
|
||||
provider_ids: &[String],
|
||||
) -> Result<Vec<StoredProviderCatalogEndpoint>, GatewayError> {
|
||||
Ok(self
|
||||
.endpoints
|
||||
.iter()
|
||||
.filter(|endpoint| {
|
||||
provider_ids
|
||||
.iter()
|
||||
.any(|provider_id| provider_id == &endpoint.provider_id)
|
||||
})
|
||||
.cloned()
|
||||
.collect())
|
||||
}
|
||||
|
||||
async fn read_provider_transport_snapshot(
|
||||
&self,
|
||||
provider_id: &str,
|
||||
endpoint_id: &str,
|
||||
key_id: &str,
|
||||
) -> Result<Option<GatewayProviderTransportSnapshot>, GatewayError> {
|
||||
Ok(self
|
||||
.transports
|
||||
.get(&(
|
||||
provider_id.to_string(),
|
||||
endpoint_id.to_string(),
|
||||
key_id.to_string(),
|
||||
))
|
||||
.cloned())
|
||||
}
|
||||
|
||||
async fn execute_execution_runtime_sync_plan(
|
||||
&self,
|
||||
_plan: &ExecutionPlan,
|
||||
) -> Result<ExecutionResult, GatewayError> {
|
||||
Err(GatewayError::Internal(
|
||||
"execute_execution_runtime_sync_plan should not be called".to_string(),
|
||||
))
|
||||
}
|
||||
|
||||
async fn update_provider_catalog_key(
|
||||
&self,
|
||||
key: &StoredProviderCatalogKey,
|
||||
) -> Result<(), GatewayError> {
|
||||
let mut keys = self.keys.lock().expect("keys mutex");
|
||||
let Some(slot) = keys.iter_mut().find(|item| item.id == key.id) else {
|
||||
return Err(GatewayError::Internal("key not found".to_string()));
|
||||
};
|
||||
*slot = key.clone();
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn write_upstream_models_cache(
|
||||
&self,
|
||||
provider_id: &str,
|
||||
key_id: &str,
|
||||
cached_models: &[Value],
|
||||
) {
|
||||
self.cached_models.lock().expect("cache mutex").insert(
|
||||
(provider_id.to_string(), key_id.to_string()),
|
||||
cached_models.to_vec(),
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
fn sample_provider(provider_id: &str, provider_type: &str) -> StoredProviderCatalogProvider {
|
||||
StoredProviderCatalogProvider::new(
|
||||
provider_id.to_string(),
|
||||
provider_id.to_string(),
|
||||
None,
|
||||
provider_type.to_string(),
|
||||
)
|
||||
.expect("provider should build")
|
||||
.with_transport_fields(true, false, false, None, None, None, None, None, None)
|
||||
}
|
||||
|
||||
fn sample_endpoint(
|
||||
endpoint_id: &str,
|
||||
provider_id: &str,
|
||||
api_format: &str,
|
||||
) -> StoredProviderCatalogEndpoint {
|
||||
StoredProviderCatalogEndpoint::new(
|
||||
endpoint_id.to_string(),
|
||||
provider_id.to_string(),
|
||||
api_format.to_string(),
|
||||
None,
|
||||
None,
|
||||
true,
|
||||
)
|
||||
.expect("endpoint should build")
|
||||
.with_transport_fields(
|
||||
"https://cloudcode-pa.googleapis.com".to_string(),
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
)
|
||||
.expect("endpoint transport should build")
|
||||
}
|
||||
|
||||
fn sample_key(
|
||||
key_id: &str,
|
||||
provider_id: &str,
|
||||
auth_type: &str,
|
||||
api_formats: &[&str],
|
||||
) -> StoredProviderCatalogKey {
|
||||
let mut key = StoredProviderCatalogKey::new(
|
||||
key_id.to_string(),
|
||||
provider_id.to_string(),
|
||||
"primary".to_string(),
|
||||
auth_type.to_string(),
|
||||
None,
|
||||
true,
|
||||
)
|
||||
.expect("key should build")
|
||||
.with_transport_fields(
|
||||
Some(json!(api_formats)),
|
||||
"encrypted".to_string(),
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
)
|
||||
.expect("key transport should build");
|
||||
key.auto_fetch_models = true;
|
||||
key
|
||||
}
|
||||
|
||||
fn sample_transport(
|
||||
provider_type: &str,
|
||||
provider_id: &str,
|
||||
endpoint_id: &str,
|
||||
key_id: &str,
|
||||
api_format: &str,
|
||||
auth_type: &str,
|
||||
decrypted_auth_config: Option<&str>,
|
||||
) -> GatewayProviderTransportSnapshot {
|
||||
GatewayProviderTransportSnapshot {
|
||||
provider: GatewayProviderTransportProvider {
|
||||
id: provider_id.to_string(),
|
||||
name: provider_id.to_string(),
|
||||
provider_type: provider_type.to_string(),
|
||||
website: None,
|
||||
is_active: true,
|
||||
keep_priority_on_conversion: false,
|
||||
enable_format_conversion: false,
|
||||
concurrent_limit: None,
|
||||
max_retries: None,
|
||||
proxy: None,
|
||||
request_timeout_secs: None,
|
||||
stream_first_byte_timeout_secs: None,
|
||||
config: None,
|
||||
},
|
||||
endpoint: GatewayProviderTransportEndpoint {
|
||||
id: endpoint_id.to_string(),
|
||||
provider_id: provider_id.to_string(),
|
||||
api_format: api_format.to_string(),
|
||||
api_family: None,
|
||||
endpoint_kind: None,
|
||||
is_active: true,
|
||||
base_url: "https://cloudcode-pa.googleapis.com".to_string(),
|
||||
header_rules: None,
|
||||
body_rules: None,
|
||||
max_retries: None,
|
||||
custom_path: None,
|
||||
config: None,
|
||||
format_acceptance_config: None,
|
||||
proxy: None,
|
||||
},
|
||||
key: GatewayProviderTransportKey {
|
||||
id: key_id.to_string(),
|
||||
provider_id: provider_id.to_string(),
|
||||
name: "primary".to_string(),
|
||||
auth_type: auth_type.to_string(),
|
||||
is_active: true,
|
||||
api_formats: Some(vec![api_format.to_string()]),
|
||||
allowed_models: None,
|
||||
capabilities: None,
|
||||
rate_multipliers: None,
|
||||
global_priority_by_format: None,
|
||||
expires_at_unix_secs: None,
|
||||
proxy: None,
|
||||
fingerprint: None,
|
||||
decrypted_api_key: "secret".to_string(),
|
||||
decrypted_auth_config: decrypted_auth_config.map(ToOwned::to_owned),
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
fn execution_result(body: Value) -> ExecutionResult {
|
||||
ExecutionResult {
|
||||
request_id: "req-1".to_string(),
|
||||
candidate_id: None,
|
||||
status_code: 200,
|
||||
headers: Default::default(),
|
||||
body: Some(aether_contracts::ResponseBody {
|
||||
json_body: Some(body),
|
||||
body_bytes_b64: None,
|
||||
}),
|
||||
telemetry: None,
|
||||
error: None,
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn model_fetch_uses_preset_models_without_endpoint() {
|
||||
let provider = sample_provider("provider-codex", "codex");
|
||||
let key = sample_key("key-codex", "provider-codex", "api_key", &["openai:cli"]);
|
||||
let state = TestState::new(vec![provider], vec![], vec![key], HashMap::new(), vec![]);
|
||||
|
||||
let summary = perform_model_fetch_once_with_state(&state)
|
||||
.await
|
||||
.expect("fetch should succeed");
|
||||
|
||||
assert_eq!(summary.attempted, 1);
|
||||
assert_eq!(summary.succeeded, 1);
|
||||
let updated = state.key("key-codex");
|
||||
let allowed_models = updated
|
||||
.allowed_models
|
||||
.and_then(|value| value.as_array().cloned())
|
||||
.expect("allowed_models should be set");
|
||||
assert!(allowed_models.iter().any(|model| model == "gpt-5.4"));
|
||||
assert!(state
|
||||
.cached_models
|
||||
.lock()
|
||||
.expect("cache mutex")
|
||||
.contains_key(&("provider-codex".to_string(), "key-codex".to_string())));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn model_fetch_merges_antigravity_metadata_and_preserves_reset_time() {
|
||||
let provider = sample_provider("provider-antigravity", "antigravity");
|
||||
let endpoint = sample_endpoint(
|
||||
"endpoint-antigravity",
|
||||
"provider-antigravity",
|
||||
"gemini:chat",
|
||||
);
|
||||
let mut key = sample_key(
|
||||
"key-antigravity",
|
||||
"provider-antigravity",
|
||||
"oauth",
|
||||
&["gemini:chat"],
|
||||
);
|
||||
key.upstream_metadata = Some(json!({
|
||||
"antigravity": {
|
||||
"quota_by_model": {
|
||||
"gemini-2.5-pro": {
|
||||
"reset_time": "2026-04-12T00:00:00Z"
|
||||
}
|
||||
}
|
||||
}
|
||||
}));
|
||||
let transport = sample_transport(
|
||||
"antigravity",
|
||||
"provider-antigravity",
|
||||
"endpoint-antigravity",
|
||||
"key-antigravity",
|
||||
"gemini:chat",
|
||||
"oauth",
|
||||
Some(r#"{"project_id":"project-1","client_version":"1.2.3","session_id":"sess-1"}"#),
|
||||
);
|
||||
let state = TestState::new(
|
||||
vec![provider],
|
||||
vec![endpoint],
|
||||
vec![key],
|
||||
HashMap::from([(
|
||||
(
|
||||
"provider-antigravity".to_string(),
|
||||
"endpoint-antigravity".to_string(),
|
||||
"key-antigravity".to_string(),
|
||||
),
|
||||
transport,
|
||||
)]),
|
||||
vec![execution_result(json!({
|
||||
"models": {
|
||||
"gemini-2.5-pro": {
|
||||
"displayName": "Gemini 2.5 Pro",
|
||||
"quotaInfo": {
|
||||
"remainingFraction": 0.25
|
||||
}
|
||||
}
|
||||
}
|
||||
}))],
|
||||
);
|
||||
|
||||
let summary = perform_model_fetch_once_with_state(&state)
|
||||
.await
|
||||
.expect("fetch should succeed");
|
||||
|
||||
assert_eq!(summary.succeeded, 1);
|
||||
let updated = state.key("key-antigravity");
|
||||
assert_eq!(updated.allowed_models, Some(json!(["gemini-2.5-pro"])));
|
||||
assert_eq!(
|
||||
updated
|
||||
.upstream_metadata
|
||||
.as_ref()
|
||||
.and_then(|value| value.get("antigravity"))
|
||||
.and_then(|value| value.get("quota_by_model"))
|
||||
.and_then(|value| value.get("gemini-2.5-pro"))
|
||||
.and_then(|value| value.get("reset_time")),
|
||||
Some(&json!("2026-04-12T00:00:00Z"))
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn model_fetch_failure_keeps_existing_allowed_models() {
|
||||
let provider = sample_provider("provider-openai", "openai");
|
||||
let endpoint = sample_endpoint(
|
||||
"endpoint-openai-responses",
|
||||
"provider-openai",
|
||||
"openai:responses",
|
||||
);
|
||||
let mut key = sample_key(
|
||||
"key-openai-responses",
|
||||
"provider-openai",
|
||||
"api_key",
|
||||
&["openai:responses"],
|
||||
);
|
||||
key.allowed_models = Some(json!(["gpt-old"]));
|
||||
let state = TestState::new(
|
||||
vec![provider],
|
||||
vec![endpoint],
|
||||
vec![key],
|
||||
HashMap::new(),
|
||||
vec![],
|
||||
);
|
||||
|
||||
let summary = perform_model_fetch_once_with_state(&state)
|
||||
.await
|
||||
.expect("fetch should finish");
|
||||
|
||||
assert_eq!(summary.succeeded, 0);
|
||||
assert_eq!(summary.skipped, 1);
|
||||
let updated = state.key("key-openai-responses");
|
||||
assert_eq!(updated.allowed_models, Some(json!(["gpt-old"])));
|
||||
assert_eq!(
|
||||
updated.last_models_fetch_error.as_deref(),
|
||||
Some("No supported endpoint for Rust models fetch")
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user