mirror of
https://github.com/fawney19/Aether.git
synced 2026-10-04 08:27:46 +08:00
Merge commit 'refs/pr/530'
This commit is contained in:
@@ -31,6 +31,7 @@ pub(super) struct AdminProviderOAuthBatchImportEntry {
|
||||
pub user_id: Option<String>,
|
||||
pub email: Option<String>,
|
||||
pub account_name: Option<String>,
|
||||
pub project_id: Option<String>,
|
||||
pub sso_rw_token: Option<String>,
|
||||
pub cf_cookies: Option<String>,
|
||||
pub cf_clearance: Option<String>,
|
||||
@@ -76,6 +77,20 @@ fn coerce_admin_provider_oauth_import_str(value: Option<&serde_json::Value>) ->
|
||||
.map(ToOwned::to_owned)
|
||||
}
|
||||
|
||||
fn coerce_admin_provider_oauth_import_project_id(
|
||||
value: Option<&serde_json::Value>,
|
||||
) -> Option<String> {
|
||||
match value {
|
||||
Some(serde_json::Value::Object(object)) => coerce_admin_provider_oauth_import_str(
|
||||
object
|
||||
.get("id")
|
||||
.or_else(|| object.get("project_id"))
|
||||
.or_else(|| object.get("projectId")),
|
||||
),
|
||||
other => coerce_admin_provider_oauth_import_str(other),
|
||||
}
|
||||
}
|
||||
|
||||
fn json_import_expiry_value(value: Option<&serde_json::Value>) -> Option<u64> {
|
||||
let value = value?;
|
||||
json_u64_value(Some(value)).or_else(|| {
|
||||
@@ -175,6 +190,7 @@ fn extract_admin_provider_oauth_batch_import_entry(
|
||||
user_id: grok_cookie_value(raw_token, "x-userid"),
|
||||
email: None,
|
||||
account_name: None,
|
||||
project_id: None,
|
||||
sso_rw_token: grok_cookie_value(raw_token, "sso-rw"),
|
||||
cf_cookies: grok_cookie_profile(raw_token),
|
||||
cf_clearance: grok_cookie_value(raw_token, "cf_clearance"),
|
||||
@@ -308,6 +324,13 @@ fn extract_admin_provider_oauth_batch_import_entry(
|
||||
.get("account_name")
|
||||
.or_else(|| object.get("accountName")),
|
||||
);
|
||||
let project_id = coerce_admin_provider_oauth_import_project_id(
|
||||
object
|
||||
.get("project_id")
|
||||
.or_else(|| object.get("projectId"))
|
||||
.or_else(|| object.get("cloudaicompanionProject"))
|
||||
.or_else(|| object.get("cloudAiCompanionProject")),
|
||||
);
|
||||
let sso_rw_token = coerce_admin_provider_oauth_import_str(
|
||||
object
|
||||
.get("sso_rw_token")
|
||||
@@ -355,6 +378,7 @@ fn extract_admin_provider_oauth_batch_import_entry(
|
||||
user_id,
|
||||
email,
|
||||
account_name,
|
||||
project_id,
|
||||
sso_rw_token,
|
||||
cf_cookies,
|
||||
cf_clearance,
|
||||
@@ -445,6 +469,7 @@ fn parse_error_entry(error: String) -> AdminProviderOAuthBatchImportEntry {
|
||||
user_id: None,
|
||||
email: None,
|
||||
account_name: None,
|
||||
project_id: None,
|
||||
sso_rw_token: None,
|
||||
cf_cookies: None,
|
||||
cf_clearance: None,
|
||||
@@ -464,6 +489,19 @@ pub(super) fn apply_admin_provider_oauth_batch_import_hints(
|
||||
auth_config: &mut serde_json::Map<String, serde_json::Value>,
|
||||
) {
|
||||
let provider_type = provider_type.trim().to_ascii_lowercase();
|
||||
if provider_type == "gemini_cli" {
|
||||
if let Some(project_id) = entry.project_id.as_ref() {
|
||||
auth_config
|
||||
.entry("project_id".to_string())
|
||||
.or_insert_with(|| json!(project_id));
|
||||
}
|
||||
if let Some(plan_type) = entry.plan_type.as_ref() {
|
||||
auth_config
|
||||
.entry("plan_type".to_string())
|
||||
.or_insert_with(|| json!(plan_type));
|
||||
}
|
||||
return;
|
||||
}
|
||||
if !matches!(provider_type.as_str(), "codex" | "chatgpt_web" | "grok") {
|
||||
return;
|
||||
}
|
||||
@@ -614,7 +652,10 @@ pub(super) fn build_admin_provider_oauth_batch_task_state(
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::parse_admin_provider_oauth_batch_import_entries;
|
||||
use super::{
|
||||
apply_admin_provider_oauth_batch_import_hints,
|
||||
parse_admin_provider_oauth_batch_import_entries,
|
||||
};
|
||||
use base64::{engine::general_purpose::URL_SAFE_NO_PAD, Engine as _};
|
||||
use serde_json::json;
|
||||
|
||||
@@ -750,6 +791,38 @@ mod tests {
|
||||
assert_eq!(entries[0].pool_tier.as_deref(), Some("heavy"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parses_gemini_cli_project_id_hint() {
|
||||
let entries = parse_admin_provider_oauth_batch_import_entries(
|
||||
"gemini_cli",
|
||||
r#"[{"refresh_token":"rt-1","projectId":"project-gemini-cli-1","planType":"free"}]"#,
|
||||
);
|
||||
|
||||
assert_eq!(entries.len(), 1);
|
||||
assert_eq!(entries[0].refresh_token.as_deref(), Some("rt-1"));
|
||||
assert_eq!(
|
||||
entries[0].project_id.as_deref(),
|
||||
Some("project-gemini-cli-1")
|
||||
);
|
||||
assert_eq!(entries[0].plan_type.as_deref(), Some("free"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn applies_gemini_cli_project_id_hint_to_auth_config() {
|
||||
let entries = parse_admin_provider_oauth_batch_import_entries(
|
||||
"gemini_cli",
|
||||
r#"{"refreshToken":"rt-1","cloudaicompanionProject":{"id":"project-gemini-cli-2"}}"#,
|
||||
);
|
||||
let mut auth_config = serde_json::Map::new();
|
||||
|
||||
apply_admin_provider_oauth_batch_import_hints("gemini_cli", &entries[0], &mut auth_config);
|
||||
|
||||
assert_eq!(
|
||||
auth_config.get("project_id"),
|
||||
Some(&json!("project-gemini-cli-2"))
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parses_windsurf_json_credentials_for_native_import() {
|
||||
let entries = parse_admin_provider_oauth_batch_import_entries(
|
||||
|
||||
@@ -4,6 +4,7 @@ use std::pin::Pin;
|
||||
use super::antigravity::refresh_antigravity_provider_quota_locally;
|
||||
use super::chatgpt_web::refresh_chatgpt_web_provider_quota_locally;
|
||||
use super::codex::refresh_codex_provider_quota_locally;
|
||||
use super::gemini_cli::refresh_gemini_cli_provider_quota_locally;
|
||||
use super::grok::refresh_grok_provider_quota_locally;
|
||||
use super::kiro::refresh_kiro_provider_quota_locally;
|
||||
use super::windsurf::refresh_windsurf_provider_quota_locally;
|
||||
@@ -35,6 +36,10 @@ const PROVIDER_QUOTA_REFRESH_HANDLERS: &[(&str, ProviderQuotaRefreshHandler)] =
|
||||
refresh_chatgpt_web_provider_quota_locally_boxed,
|
||||
),
|
||||
("codex", refresh_codex_provider_quota_locally_boxed),
|
||||
(
|
||||
"gemini_cli",
|
||||
refresh_gemini_cli_provider_quota_locally_boxed,
|
||||
),
|
||||
("grok", refresh_grok_provider_quota_locally_boxed),
|
||||
("kiro", refresh_kiro_provider_quota_locally_boxed),
|
||||
("windsurf", refresh_windsurf_provider_quota_locally_boxed),
|
||||
@@ -106,6 +111,22 @@ fn refresh_codex_provider_quota_locally_boxed<'a>(
|
||||
))
|
||||
}
|
||||
|
||||
fn refresh_gemini_cli_provider_quota_locally_boxed<'a>(
|
||||
state: &'a AdminAppState<'a>,
|
||||
provider: &'a StoredProviderCatalogProvider,
|
||||
endpoint: &'a StoredProviderCatalogEndpoint,
|
||||
keys: Vec<StoredProviderCatalogKey>,
|
||||
proxy_override: Option<ProxySnapshot>,
|
||||
) -> ProviderQuotaRefreshFuture<'a> {
|
||||
Box::pin(refresh_gemini_cli_provider_quota_locally(
|
||||
state,
|
||||
provider,
|
||||
endpoint,
|
||||
keys,
|
||||
proxy_override,
|
||||
))
|
||||
}
|
||||
|
||||
fn refresh_kiro_provider_quota_locally_boxed<'a>(
|
||||
state: &'a AdminAppState<'a>,
|
||||
provider: &'a StoredProviderCatalogProvider,
|
||||
|
||||
@@ -0,0 +1,271 @@
|
||||
use super::shared::{
|
||||
build_provider_quota_execution_plan, build_quota_snapshot_payload,
|
||||
default_provider_quota_execution_timeouts, execute_provider_quota_plan,
|
||||
extract_execution_error_message, oauth_refresh_auto_removed_result,
|
||||
persist_provider_quota_refresh_state, quota_key_auto_removed,
|
||||
quota_refresh_success_invalid_state, ProviderQuotaExecutionOutcome,
|
||||
};
|
||||
use crate::handlers::admin::request::{AdminAppState, AdminGatewayProviderTransportSnapshot};
|
||||
use crate::GatewayError;
|
||||
use aether_admin::provider::quota::parse_gemini_cli_retrieve_user_quota_response;
|
||||
use aether_contracts::ProxySnapshot;
|
||||
use aether_data_contracts::repository::provider_catalog::{
|
||||
StoredProviderCatalogEndpoint, StoredProviderCatalogKey, StoredProviderCatalogProvider,
|
||||
};
|
||||
use aether_provider_pool::build_gemini_cli_pool_quota_request;
|
||||
use serde_json::json;
|
||||
use std::time::{SystemTime, UNIX_EPOCH};
|
||||
|
||||
async fn execute_gemini_cli_quota_plan(
|
||||
state: &AdminAppState<'_>,
|
||||
transport: &AdminGatewayProviderTransportSnapshot,
|
||||
authorization: (String, String),
|
||||
project_id: &str,
|
||||
proxy_override: Option<&ProxySnapshot>,
|
||||
) -> Result<ProviderQuotaExecutionOutcome, GatewayError> {
|
||||
let proxy = match proxy_override {
|
||||
Some(proxy) => Some(proxy.clone()),
|
||||
None => {
|
||||
state
|
||||
.resolve_transport_proxy_snapshot_with_tunnel_affinity(transport)
|
||||
.await
|
||||
}
|
||||
};
|
||||
let timeouts = state
|
||||
.resolve_transport_execution_timeouts(transport)
|
||||
.or(Some(default_provider_quota_execution_timeouts(
|
||||
proxy.as_ref(),
|
||||
)));
|
||||
let spec = build_gemini_cli_pool_quota_request(
|
||||
&transport.key.id,
|
||||
&transport.endpoint.base_url,
|
||||
authorization,
|
||||
project_id,
|
||||
);
|
||||
let plan = build_provider_quota_execution_plan(
|
||||
transport,
|
||||
spec,
|
||||
proxy,
|
||||
state.resolve_transport_profile(transport),
|
||||
timeouts,
|
||||
);
|
||||
|
||||
execute_provider_quota_plan(state, transport, plan, "gemini_cli").await
|
||||
}
|
||||
|
||||
pub(crate) async fn refresh_gemini_cli_provider_quota_locally(
|
||||
state: &AdminAppState<'_>,
|
||||
provider: &StoredProviderCatalogProvider,
|
||||
endpoint: &StoredProviderCatalogEndpoint,
|
||||
keys: Vec<StoredProviderCatalogKey>,
|
||||
proxy_override: Option<ProxySnapshot>,
|
||||
) -> Result<Option<serde_json::Value>, GatewayError> {
|
||||
let mut results = Vec::new();
|
||||
let mut success_count = 0usize;
|
||||
let mut failed_count = 0usize;
|
||||
let mut auto_removed_count = 0usize;
|
||||
|
||||
for key in keys {
|
||||
let mut transport = match state
|
||||
.read_provider_transport_snapshot(&provider.id, &endpoint.id, &key.id)
|
||||
.await?
|
||||
{
|
||||
Some(transport) => transport,
|
||||
None => {
|
||||
failed_count += 1;
|
||||
results.push(json!({
|
||||
"key_id": key.id,
|
||||
"key_name": key.name,
|
||||
"status": "error",
|
||||
"message": "Provider transport snapshot unavailable",
|
||||
}));
|
||||
continue;
|
||||
}
|
||||
};
|
||||
|
||||
let authorization = match state.resolve_local_oauth_header_auth(&transport).await? {
|
||||
Some(auth) => auth,
|
||||
_ => {
|
||||
if quota_key_auto_removed(state, &key.id).await? {
|
||||
auto_removed_count += 1;
|
||||
results.push(oauth_refresh_auto_removed_result(&key));
|
||||
continue;
|
||||
}
|
||||
failed_count += 1;
|
||||
results.push(json!({
|
||||
"key_id": key.id,
|
||||
"key_name": key.name,
|
||||
"status": "error",
|
||||
"message": "缺少 OAuth 认证信息,请先授权/刷新 Token",
|
||||
}));
|
||||
continue;
|
||||
}
|
||||
};
|
||||
|
||||
let project_id = match crate::provider_transport::resolve_gemini_cli_project_id(&transport)
|
||||
{
|
||||
Some(project_id) => Some(project_id),
|
||||
None => state
|
||||
.app()
|
||||
.hydrate_gemini_cli_project_metadata_for_transport(&transport)
|
||||
.await
|
||||
.and_then(|hydrated| {
|
||||
let project_id =
|
||||
crate::provider_transport::resolve_gemini_cli_project_id(&hydrated);
|
||||
transport = hydrated;
|
||||
project_id
|
||||
}),
|
||||
};
|
||||
let Some(project_id) = project_id else {
|
||||
failed_count += 1;
|
||||
results.push(json!({
|
||||
"key_id": key.id,
|
||||
"key_name": key.name,
|
||||
"status": "error",
|
||||
"message": "缺少 Gemini CLI project_id,loadCodeAssist 未返回可用项目信息",
|
||||
}));
|
||||
continue;
|
||||
};
|
||||
|
||||
let result = match execute_gemini_cli_quota_plan(
|
||||
state,
|
||||
&transport,
|
||||
authorization,
|
||||
&project_id,
|
||||
proxy_override.as_ref(),
|
||||
)
|
||||
.await?
|
||||
{
|
||||
ProviderQuotaExecutionOutcome::Response(result) => result,
|
||||
ProviderQuotaExecutionOutcome::Failure(detail) => {
|
||||
failed_count += 1;
|
||||
results.push(json!({
|
||||
"key_id": key.id,
|
||||
"key_name": key.name,
|
||||
"status": "error",
|
||||
"message": format!("retrieveUserQuota 请求执行失败: {detail}"),
|
||||
"status_code": 502,
|
||||
}));
|
||||
continue;
|
||||
}
|
||||
};
|
||||
|
||||
let now_unix_secs = SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.ok()
|
||||
.map(|duration| duration.as_secs())
|
||||
.unwrap_or(0);
|
||||
let mut metadata_update = None::<serde_json::Value>;
|
||||
let (mut oauth_invalid_at_unix_secs, mut oauth_invalid_reason) =
|
||||
quota_refresh_success_invalid_state(&key);
|
||||
let mut status = "error".to_string();
|
||||
let mut message = None::<String>;
|
||||
|
||||
if result.status_code == 200 {
|
||||
if let Some(body_json) = result
|
||||
.body
|
||||
.as_ref()
|
||||
.and_then(|body| body.json_body.as_ref())
|
||||
{
|
||||
metadata_update =
|
||||
parse_gemini_cli_retrieve_user_quota_response(body_json, now_unix_secs)
|
||||
.map(|metadata| json!({ "gemini_cli": metadata }));
|
||||
if metadata_update.is_some() {
|
||||
status = "success".to_string();
|
||||
} else {
|
||||
status = "no_metadata".to_string();
|
||||
message = Some("响应中未包含配额 buckets".to_string());
|
||||
}
|
||||
} else {
|
||||
status = "no_metadata".to_string();
|
||||
message = Some("响应中未包含配额信息".to_string());
|
||||
}
|
||||
} else {
|
||||
let err_msg = extract_execution_error_message(&result);
|
||||
message = Some(match err_msg.as_deref() {
|
||||
Some(detail) if !detail.is_empty() => {
|
||||
format!(
|
||||
"retrieveUserQuota 返回状态码 {}: {}",
|
||||
result.status_code, detail
|
||||
)
|
||||
}
|
||||
_ => format!("retrieveUserQuota 返回状态码 {}", result.status_code),
|
||||
});
|
||||
if result.status_code == 403 {
|
||||
let reason = err_msg
|
||||
.clone()
|
||||
.filter(|value| !value.trim().is_empty())
|
||||
.unwrap_or_else(|| "账户访问被禁止".to_string());
|
||||
oauth_invalid_at_unix_secs = Some(now_unix_secs);
|
||||
oauth_invalid_reason = Some(format!("账户访问被禁止: {reason}"));
|
||||
metadata_update = Some(json!({
|
||||
"gemini_cli": {
|
||||
"is_forbidden": true,
|
||||
"forbidden_reason": reason,
|
||||
"forbidden_at": now_unix_secs,
|
||||
"updated_at": now_unix_secs,
|
||||
}
|
||||
}));
|
||||
status = "forbidden".to_string();
|
||||
}
|
||||
}
|
||||
|
||||
if !persist_provider_quota_refresh_state(
|
||||
state,
|
||||
&key.id,
|
||||
metadata_update.as_ref(),
|
||||
oauth_invalid_at_unix_secs,
|
||||
oauth_invalid_reason,
|
||||
None,
|
||||
)
|
||||
.await?
|
||||
{
|
||||
failed_count += 1;
|
||||
results.push(json!({
|
||||
"key_id": key.id,
|
||||
"key_name": key.name,
|
||||
"status": "error",
|
||||
"message": "Key 状态写入失败",
|
||||
}));
|
||||
continue;
|
||||
}
|
||||
|
||||
if status == "success" {
|
||||
success_count += 1;
|
||||
} else {
|
||||
failed_count += 1;
|
||||
}
|
||||
|
||||
let mut payload = serde_json::Map::new();
|
||||
payload.insert("key_id".to_string(), json!(key.id));
|
||||
payload.insert("key_name".to_string(), json!(key.name));
|
||||
payload.insert("status".to_string(), json!(status));
|
||||
if let Some(message) = message {
|
||||
payload.insert("message".to_string(), json!(message));
|
||||
}
|
||||
if let Some(metadata) = metadata_update
|
||||
.as_ref()
|
||||
.and_then(|value| value.get("gemini_cli"))
|
||||
.cloned()
|
||||
{
|
||||
payload.insert("metadata".to_string(), metadata);
|
||||
}
|
||||
if let Some(quota_snapshot) = build_quota_snapshot_payload(
|
||||
"gemini_cli",
|
||||
key.status_snapshot.as_ref(),
|
||||
metadata_update.as_ref(),
|
||||
) {
|
||||
payload.insert("quota_snapshot".to_string(), quota_snapshot);
|
||||
}
|
||||
results.push(serde_json::Value::Object(payload));
|
||||
}
|
||||
|
||||
Ok(Some(json!({
|
||||
"success": success_count,
|
||||
"failed": failed_count,
|
||||
"total": results.len(),
|
||||
"results": results,
|
||||
"message": format!("已处理 {} 个 Key", results.len()),
|
||||
"auto_removed": auto_removed_count,
|
||||
})))
|
||||
}
|
||||
@@ -2,6 +2,7 @@ pub(crate) mod antigravity;
|
||||
pub(crate) mod chatgpt_web;
|
||||
pub(crate) mod codex;
|
||||
pub(crate) mod dispatch;
|
||||
pub(crate) mod gemini_cli;
|
||||
pub(crate) mod grok;
|
||||
pub(crate) mod kiro;
|
||||
pub(crate) mod shared;
|
||||
|
||||
@@ -782,6 +782,18 @@ fn admin_pool_build_grok_account_quota_from_snapshot(
|
||||
fn admin_pool_build_gemini_cli_account_quota_from_snapshot(
|
||||
quota_snapshot: &serde_json::Map<String, serde_json::Value>,
|
||||
) -> Option<String> {
|
||||
if let Some(credits) = quota_snapshot
|
||||
.get("credits")
|
||||
.and_then(serde_json::Value::as_object)
|
||||
{
|
||||
if let Some(remaining) = admin_pool_json_to_f64(credits.get("remaining")) {
|
||||
return Some(format!(
|
||||
"AI Credits 剩余 {}",
|
||||
admin_pool_format_quota_value(remaining)
|
||||
));
|
||||
}
|
||||
}
|
||||
|
||||
let now = chrono::Utc::now().timestamp();
|
||||
let mut active = admin_pool_quota_windows(quota_snapshot)
|
||||
.into_iter()
|
||||
|
||||
@@ -15,7 +15,7 @@ use crate::ai_serving::{
|
||||
ANTIGRAVITY_V1INTERNAL_ENVELOPE_NAME, GEMINI_CHAT_SYNC_FINALIZE_REPORT_KIND,
|
||||
OPENAI_IMAGE_SYNC_FINALIZE_REPORT_KIND,
|
||||
};
|
||||
use crate::clock::current_unix_ms;
|
||||
use crate::clock::{current_unix_ms, current_unix_secs};
|
||||
use crate::execution_runtime;
|
||||
use crate::handlers::admin::provider::shared::model_test_capabilities::{
|
||||
admin_provider_model_supports_image_generation, admin_provider_model_test_capabilities_payload,
|
||||
@@ -64,7 +64,7 @@ use aether_data_contracts::repository::provider_catalog::{
|
||||
};
|
||||
use aether_model_fetch::{
|
||||
aggregate_models_for_cache, fetch_models_from_transports, json_string_list,
|
||||
preset_models_for_provider, selected_models_fetch_endpoints,
|
||||
merge_upstream_metadata, preset_models_for_provider, selected_models_fetch_endpoints,
|
||||
};
|
||||
use axum::{
|
||||
body::{to_bytes, Body},
|
||||
@@ -516,6 +516,18 @@ async fn provider_query_fetch_models_for_key(
|
||||
)
|
||||
.await;
|
||||
}
|
||||
if let Some(upstream_metadata) = outcome.upstream_metadata.as_ref() {
|
||||
let merged_metadata =
|
||||
merge_upstream_metadata(key.upstream_metadata.as_ref(), upstream_metadata);
|
||||
state
|
||||
.app()
|
||||
.update_provider_catalog_key_upstream_metadata(
|
||||
&key.id,
|
||||
Some(&merged_metadata),
|
||||
Some(current_unix_secs()),
|
||||
)
|
||||
.await?;
|
||||
}
|
||||
|
||||
if unique_models.is_empty() && !all_errors.is_empty() {
|
||||
if let Some(fallback) = provider_query_codex_preset_fallback(provider) {
|
||||
|
||||
@@ -1576,12 +1576,18 @@ fn provider_query_aggregate_standard_stream_sync_response(
|
||||
fn provider_query_standard_execution_response_body(
|
||||
provider_api_format: &str,
|
||||
result: &aether_contracts::ExecutionResult,
|
||||
report_context: Option<&Value>,
|
||||
) -> Option<Value> {
|
||||
let body = provider_query_execution_json_body(result).or_else(|| {
|
||||
provider_query_decode_execution_body(result).and_then(|body| {
|
||||
provider_query_aggregate_standard_stream_sync_response(provider_api_format, &body)
|
||||
})
|
||||
})?;
|
||||
let body = report_context
|
||||
.and_then(|context| {
|
||||
crate::ai_serving::api::normalize_provider_private_response_value(body.clone(), context)
|
||||
})
|
||||
.unwrap_or(body);
|
||||
if result.status_code < 400
|
||||
&& provider_query_normalize_api_format_alias(provider_api_format)
|
||||
== "gemini:generate_content"
|
||||
@@ -2750,7 +2756,7 @@ async fn provider_query_execute_standard_test_candidate(
|
||||
route_path: &str,
|
||||
trace_id: &str,
|
||||
) -> Result<ProviderQueryExecutionOutcome, GatewayError> {
|
||||
let Some(transport) = state
|
||||
let Some(mut transport) = state
|
||||
.read_provider_transport_snapshot(&provider.id, &candidate.endpoint.id, &candidate.key.id)
|
||||
.await?
|
||||
else {
|
||||
@@ -2963,6 +2969,56 @@ async fn provider_query_execute_standard_test_candidate(
|
||||
upstream_is_stream,
|
||||
require_body_stream_field,
|
||||
);
|
||||
if crate::provider_transport::is_gemini_cli_provider_transport(&transport)
|
||||
&& normalized_provider_api_format == "gemini:generate_content"
|
||||
{
|
||||
let project_id = match crate::provider_transport::resolve_gemini_cli_project_id(&transport)
|
||||
{
|
||||
Some(project_id) => Some(project_id),
|
||||
None => state
|
||||
.app()
|
||||
.hydrate_gemini_cli_project_metadata_for_transport(&transport)
|
||||
.await
|
||||
.and_then(|hydrated| {
|
||||
let project_id =
|
||||
crate::provider_transport::resolve_gemini_cli_project_id(&hydrated);
|
||||
transport = hydrated;
|
||||
project_id
|
||||
}),
|
||||
};
|
||||
let Some(project_id) = project_id else {
|
||||
return Ok(provider_query_skipped_execution_outcome(
|
||||
provider_request_body,
|
||||
"Gemini CLI project_id is unavailable for v1internal request",
|
||||
));
|
||||
};
|
||||
provider_request_body = match crate::provider_transport::build_gemini_cli_v1internal_request(
|
||||
project_id.as_str(),
|
||||
trace_id,
|
||||
request_model,
|
||||
&provider_request_body,
|
||||
) {
|
||||
crate::provider_transport::GeminiCliRequestEnvelopeSupport::Supported(envelope) => {
|
||||
envelope
|
||||
}
|
||||
crate::provider_transport::GeminiCliRequestEnvelopeSupport::Unsupported(_) => {
|
||||
return Ok(provider_query_skipped_execution_outcome(
|
||||
provider_request_body,
|
||||
"Gemini CLI v1internal envelope could not be built",
|
||||
));
|
||||
}
|
||||
};
|
||||
}
|
||||
let private_report_context =
|
||||
(crate::provider_transport::is_gemini_cli_provider_transport(&transport)
|
||||
&& normalized_provider_api_format == "gemini:generate_content")
|
||||
.then(|| {
|
||||
json!({
|
||||
"has_envelope": true,
|
||||
"envelope_name": crate::provider_transport::GEMINI_CLI_V1INTERNAL_ENVELOPE_NAME,
|
||||
"provider_api_format": provider_api_format,
|
||||
})
|
||||
});
|
||||
|
||||
let uses_vertex_query_auth =
|
||||
crate::provider_transport::uses_vertex_api_key_query_auth(&transport, provider_api_format);
|
||||
@@ -3083,6 +3139,13 @@ async fn provider_query_execute_standard_test_candidate(
|
||||
request_headers
|
||||
.entry("content-type".to_string())
|
||||
.or_insert_with(|| "application/json".to_string());
|
||||
if crate::provider_transport::is_gemini_cli_provider_transport(&transport)
|
||||
&& normalized_provider_api_format == "gemini:generate_content"
|
||||
{
|
||||
request_headers
|
||||
.entry("user-agent".to_string())
|
||||
.or_insert_with(|| crate::provider_transport::GEMINI_CLI_USER_AGENT.to_string());
|
||||
}
|
||||
let protected_headers = if uses_vertex_query_auth {
|
||||
vec!["content-type"]
|
||||
} else {
|
||||
@@ -3160,7 +3223,11 @@ async fn provider_query_execute_standard_test_candidate(
|
||||
.execute_execution_runtime_sync_plan(Some(trace_id), &plan)
|
||||
.await?;
|
||||
let response_body = if result.status_code < 400 {
|
||||
provider_query_standard_execution_response_body(provider_api_format, &result)
|
||||
provider_query_standard_execution_response_body(
|
||||
provider_api_format,
|
||||
&result,
|
||||
private_report_context.as_ref(),
|
||||
)
|
||||
} else {
|
||||
result.body.as_ref().and_then(|body| body.json_body.clone())
|
||||
};
|
||||
|
||||
@@ -51,6 +51,7 @@ fn sample_openai_image_transport(provider_type: &str) -> AdminGatewayProviderTra
|
||||
expires_at_unix_secs: None,
|
||||
proxy: None,
|
||||
fingerprint: None,
|
||||
upstream_metadata: None,
|
||||
decrypted_api_key: String::new(),
|
||||
decrypted_auth_config: Some(
|
||||
json!({
|
||||
@@ -185,7 +186,7 @@ fn provider_query_execution_json_body_decodes_stream_encoded_json_response() {
|
||||
Some(body.clone())
|
||||
);
|
||||
assert_eq!(
|
||||
provider_query_standard_execution_response_body("openai:image", &result),
|
||||
provider_query_standard_execution_response_body("openai:image", &result, None),
|
||||
Some(body)
|
||||
);
|
||||
}
|
||||
@@ -390,7 +391,7 @@ fn provider_query_standard_test_aggregates_responses_stream_body() {
|
||||
error: None,
|
||||
};
|
||||
|
||||
let body = provider_query_standard_execution_response_body("openai:responses", &result)
|
||||
let body = provider_query_standard_execution_response_body("openai:responses", &result, None)
|
||||
.expect("stream body should aggregate");
|
||||
|
||||
assert_eq!(body["model"], json!("gpt-5.4-mini"));
|
||||
@@ -422,7 +423,7 @@ fn provider_query_standard_test_aggregates_responses_image_generation_call() {
|
||||
error: None,
|
||||
};
|
||||
|
||||
let body = provider_query_standard_execution_response_body("openai:responses", &result)
|
||||
let body = provider_query_standard_execution_response_body("openai:responses", &result, None)
|
||||
.expect("responses image stream body should aggregate");
|
||||
|
||||
assert_eq!(body["output"][0]["type"], json!("image_generation_call"));
|
||||
@@ -584,10 +585,12 @@ fn provider_query_standard_test_rejects_gemini_success_without_visible_output()
|
||||
error: None,
|
||||
};
|
||||
|
||||
assert!(
|
||||
provider_query_standard_execution_response_body("gemini:generate_content", &result)
|
||||
.is_none()
|
||||
);
|
||||
assert!(provider_query_standard_execution_response_body(
|
||||
"gemini:generate_content",
|
||||
&result,
|
||||
None
|
||||
)
|
||||
.is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
|
||||
@@ -295,6 +295,25 @@ impl<'a> AdminAppState<'a> {
|
||||
use crate::handlers::public::{admin_requested_force_stream, normalize_admin_base_url};
|
||||
use aether_admin::provider::endpoints as admin_provider_endpoints_pure;
|
||||
let (fields, payload) = patch.into_parts();
|
||||
let provider_type = provider.provider_type.trim().to_ascii_lowercase();
|
||||
|
||||
if provider_type == "gemini_cli"
|
||||
&& [
|
||||
"base_url",
|
||||
"custom_path",
|
||||
"header_rules",
|
||||
"body_rules",
|
||||
"max_retries",
|
||||
"is_active",
|
||||
"config",
|
||||
"proxy",
|
||||
"format_acceptance_config",
|
||||
]
|
||||
.iter()
|
||||
.any(|field| fields.contains(field))
|
||||
{
|
||||
return Err("Gemini CLI Endpoint 由系统固定管理,不允许修改".to_string());
|
||||
}
|
||||
|
||||
if self.provider_type_is_fixed(&provider.provider_type)
|
||||
&& (fields.contains("base_url") || fields.contains("custom_path"))
|
||||
@@ -326,7 +345,6 @@ impl<'a> AdminAppState<'a> {
|
||||
&update_fields,
|
||||
)?;
|
||||
|
||||
let provider_type = provider.provider_type.trim().to_ascii_lowercase();
|
||||
if provider_type == "codex"
|
||||
&& crate::ai_serving::is_openai_responses_format(&existing_endpoint.api_format)
|
||||
{
|
||||
|
||||
@@ -580,6 +580,52 @@ fn provider_quota_metadata_string(
|
||||
})
|
||||
}
|
||||
|
||||
fn provider_quota_metadata_value_by_path<'a>(
|
||||
metadata: &'a Map<String, Value>,
|
||||
path: &[&str],
|
||||
) -> Option<&'a Value> {
|
||||
let (first, rest) = path.split_first()?;
|
||||
let mut current = metadata.get(*first)?;
|
||||
for segment in rest {
|
||||
current = current.as_object()?.get(*segment)?;
|
||||
}
|
||||
Some(current)
|
||||
}
|
||||
|
||||
fn provider_quota_metadata_number_by_paths(
|
||||
metadata: &Map<String, Value>,
|
||||
paths: &[&[&str]],
|
||||
) -> Option<f64> {
|
||||
paths.iter().find_map(|path| {
|
||||
provider_quota_metadata_value_by_path(metadata, path)
|
||||
.and_then(admin_provider_quota_pure::coerce_json_f64)
|
||||
.filter(|value| value.is_finite())
|
||||
})
|
||||
}
|
||||
|
||||
fn provider_quota_metadata_bool_by_paths(
|
||||
metadata: &Map<String, Value>,
|
||||
paths: &[&[&str]],
|
||||
) -> Option<bool> {
|
||||
paths.iter().find_map(|path| {
|
||||
provider_quota_metadata_value_by_path(metadata, path)
|
||||
.and_then(admin_provider_quota_pure::coerce_json_bool)
|
||||
})
|
||||
}
|
||||
|
||||
fn provider_quota_metadata_string_by_paths(
|
||||
metadata: &Map<String, Value>,
|
||||
paths: &[&[&str]],
|
||||
) -> Option<String> {
|
||||
paths.iter().find_map(|path| {
|
||||
provider_quota_metadata_value_by_path(metadata, path)
|
||||
.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.map(ToOwned::to_owned)
|
||||
})
|
||||
}
|
||||
|
||||
fn quota_windows_usage_ratio(windows: &[Value]) -> Option<f64> {
|
||||
windows
|
||||
.iter()
|
||||
@@ -1531,12 +1577,144 @@ fn build_grok_quota_status_snapshot(
|
||||
}))
|
||||
}
|
||||
|
||||
fn gemini_cli_plan_type(metadata: &Map<String, Value>) -> Option<String> {
|
||||
provider_quota_metadata_string(metadata, &["plan_type", "tier", "plan"]).or_else(|| {
|
||||
provider_quota_metadata_string_by_paths(
|
||||
metadata,
|
||||
&[
|
||||
&["paidTier", "id"],
|
||||
&["paidTier", "tierType"],
|
||||
&["paidTier", "name"],
|
||||
&["currentTier", "id"],
|
||||
&["currentTier", "tierType"],
|
||||
&["currentTier", "name"],
|
||||
],
|
||||
)
|
||||
})
|
||||
}
|
||||
|
||||
fn gemini_cli_credits_status_snapshot(metadata: &Map<String, Value>) -> Option<Value> {
|
||||
let remaining = provider_quota_metadata_number_by_paths(
|
||||
metadata,
|
||||
&[
|
||||
&["credits", "remaining"],
|
||||
&["credits", "remainingCredits"],
|
||||
&["credits", "available"],
|
||||
&["credits", "availableCredits"],
|
||||
&["credits", "balance"],
|
||||
&["remainingCredits"],
|
||||
&["availableCredits"],
|
||||
&["paidTier", "remainingCredits"],
|
||||
&["paidTier", "availableCredits"],
|
||||
&["currentTier", "remainingCredits"],
|
||||
&["currentTier", "availableCredits"],
|
||||
],
|
||||
);
|
||||
let balance = provider_quota_metadata_number_by_paths(
|
||||
metadata,
|
||||
&[
|
||||
&["credits", "balance"],
|
||||
&["credits", "remaining"],
|
||||
&["credits", "available"],
|
||||
&["paidTier", "availableCredits"],
|
||||
&["currentTier", "availableCredits"],
|
||||
&["availableCredits"],
|
||||
&["remainingCredits"],
|
||||
],
|
||||
)
|
||||
.or(remaining);
|
||||
let consumed = provider_quota_metadata_number_by_paths(
|
||||
metadata,
|
||||
&[
|
||||
&["credits", "consumed"],
|
||||
&["credits", "consumedCredits"],
|
||||
&["consumedCredits"],
|
||||
&["paidTier", "consumedCredits"],
|
||||
&["currentTier", "consumedCredits"],
|
||||
],
|
||||
);
|
||||
let total = provider_quota_metadata_number_by_paths(
|
||||
metadata,
|
||||
&[
|
||||
&["credits", "total"],
|
||||
&["credits", "totalCredits"],
|
||||
&["totalCredits"],
|
||||
&["paidTier", "totalCredits"],
|
||||
&["currentTier", "totalCredits"],
|
||||
],
|
||||
);
|
||||
let unlimited = provider_quota_metadata_bool_by_paths(
|
||||
metadata,
|
||||
&[
|
||||
&["credits", "unlimited"],
|
||||
&["credits_unlimited"],
|
||||
&["paidTier", "unlimited"],
|
||||
&["currentTier", "unlimited"],
|
||||
],
|
||||
);
|
||||
let explicit_has_credits = provider_quota_metadata_bool_by_paths(
|
||||
metadata,
|
||||
&[
|
||||
&["credits", "has_credits"],
|
||||
&["has_credits"],
|
||||
&["paidTier", "hasCredits"],
|
||||
&["currentTier", "hasCredits"],
|
||||
],
|
||||
);
|
||||
let trace_id = provider_quota_metadata_string_by_paths(
|
||||
metadata,
|
||||
&[
|
||||
&["credits", "trace_id"],
|
||||
&["credits", "traceId"],
|
||||
&["trace_id"],
|
||||
&["traceId"],
|
||||
],
|
||||
);
|
||||
let updated_at = provider_quota_metadata_value_by_path(metadata, &["credits", "updated_at"])
|
||||
.or_else(|| provider_quota_metadata_value_by_path(metadata, &["credits", "updatedAt"]))
|
||||
.or_else(|| metadata.get("updated_at"))
|
||||
.and_then(|value| provider_quota_timestamp_unix_secs(Some(value)));
|
||||
|
||||
if remaining.is_none()
|
||||
&& balance.is_none()
|
||||
&& consumed.is_none()
|
||||
&& total.is_none()
|
||||
&& unlimited.is_none()
|
||||
&& explicit_has_credits.is_none()
|
||||
&& trace_id.is_none()
|
||||
{
|
||||
return None;
|
||||
}
|
||||
|
||||
let has_credits = explicit_has_credits
|
||||
.or_else(|| unlimited.filter(|value| *value))
|
||||
.or_else(|| remaining.or(balance).map(|value| value > 0.0));
|
||||
let mut credits = Map::new();
|
||||
credits.insert("has_credits".to_string(), json!(has_credits));
|
||||
credits.insert("balance".to_string(), json!(balance));
|
||||
credits.insert("remaining".to_string(), json!(remaining.or(balance)));
|
||||
credits.insert("consumed".to_string(), json!(consumed));
|
||||
credits.insert("total".to_string(), json!(total));
|
||||
credits.insert("unlimited".to_string(), json!(unlimited));
|
||||
credits.insert("trace_id".to_string(), json!(trace_id));
|
||||
credits.insert("updated_at".to_string(), json!(updated_at));
|
||||
Some(Value::Object(credits))
|
||||
}
|
||||
|
||||
fn build_gemini_cli_quota_status_snapshot(
|
||||
upstream_metadata: Option<&Value>,
|
||||
source: &str,
|
||||
) -> Option<Value> {
|
||||
let metadata = provider_quota_metadata_bucket(upstream_metadata, "gemini_cli")?;
|
||||
let observed_at_unix_secs = provider_quota_timestamp_unix_secs(metadata.get("updated_at"));
|
||||
let credits = gemini_cli_credits_status_snapshot(metadata);
|
||||
let plan_type = gemini_cli_plan_type(metadata);
|
||||
let observed_at_unix_secs = provider_quota_timestamp_unix_secs(metadata.get("updated_at"))
|
||||
.or_else(|| {
|
||||
credits
|
||||
.as_ref()
|
||||
.and_then(|value| value.get("updated_at"))
|
||||
.and_then(|value| provider_quota_timestamp_unix_secs(Some(value)))
|
||||
});
|
||||
let windows = provider_quota_model_bucket(metadata)
|
||||
.map(|models| {
|
||||
models
|
||||
@@ -1552,7 +1730,11 @@ fn build_gemini_cli_quota_status_snapshot(
|
||||
})
|
||||
.unwrap_or_default();
|
||||
|
||||
if windows.is_empty() && observed_at_unix_secs.is_none() {
|
||||
if windows.is_empty()
|
||||
&& observed_at_unix_secs.is_none()
|
||||
&& credits.is_none()
|
||||
&& plan_type.is_none()
|
||||
{
|
||||
return None;
|
||||
}
|
||||
|
||||
@@ -1616,7 +1798,8 @@ fn build_gemini_cli_quota_status_snapshot(
|
||||
"updated_at": observed_at_unix_secs,
|
||||
"reset_at": reset_at,
|
||||
"reset_seconds": reset_seconds,
|
||||
"plan_type": serde_json::Value::Null,
|
||||
"plan_type": plan_type,
|
||||
"credits": credits,
|
||||
"windows": windows,
|
||||
}))
|
||||
}
|
||||
@@ -2683,6 +2866,55 @@ mod tests {
|
||||
assert_eq!(auto.get("used_value"), Some(&json!(90.0)));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn provider_key_status_snapshot_payload_backfills_gemini_cli_account_credits() {
|
||||
let mut key = sample_catalog_key();
|
||||
key.upstream_metadata = Some(json!({
|
||||
"gemini_cli": {
|
||||
"updated_at": 1_778_067_246u64,
|
||||
"plan_type": "g1-pro-tier",
|
||||
"paidTier": {
|
||||
"availableCredits": 123.5,
|
||||
"consumedCredits": 7.0,
|
||||
"totalCredits": 200.0
|
||||
},
|
||||
"quota_by_model": {
|
||||
"gemini-2.5-pro": {
|
||||
"display_name": "Gemini 2.5 Pro",
|
||||
"remaining_fraction": 0.75,
|
||||
"is_exhausted": false
|
||||
}
|
||||
}
|
||||
}
|
||||
}));
|
||||
|
||||
let payload = provider_key_status_snapshot_payload(&key, "gemini_cli");
|
||||
let quota = payload
|
||||
.get("quota")
|
||||
.and_then(Value::as_object)
|
||||
.expect("quota snapshot should be object");
|
||||
let credits = quota
|
||||
.get("credits")
|
||||
.and_then(Value::as_object)
|
||||
.expect("credits snapshot should exist");
|
||||
let windows = quota
|
||||
.get("windows")
|
||||
.and_then(Value::as_array)
|
||||
.expect("Gemini CLI model windows should exist");
|
||||
|
||||
assert_eq!(quota.get("provider_type"), Some(&json!("gemini_cli")));
|
||||
assert_eq!(quota.get("code"), Some(&json!("ok")));
|
||||
assert_eq!(quota.get("plan_type"), Some(&json!("g1-pro-tier")));
|
||||
assert_eq!(quota.get("updated_at"), Some(&json!(1_778_067_246u64)));
|
||||
assert_eq!(credits.get("remaining"), Some(&json!(123.5)));
|
||||
assert_eq!(credits.get("balance"), Some(&json!(123.5)));
|
||||
assert_eq!(credits.get("consumed"), Some(&json!(7.0)));
|
||||
assert_eq!(credits.get("total"), Some(&json!(200.0)));
|
||||
assert_eq!(credits.get("has_credits"), Some(&json!(true)));
|
||||
assert_eq!(windows.len(), 1);
|
||||
assert_eq!(windows[0].get("remaining_ratio"), Some(&json!(0.75)));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn provider_key_status_snapshot_payload_backfills_windsurf_daily_and_weekly_quota() {
|
||||
let mut key = sample_catalog_key();
|
||||
|
||||
Reference in New Issue
Block a user