Merge branch 'fawney19:main' into main

This commit is contained in:
Start
2026-07-09 15:11:44 +08:00
committed by GitHub
67 changed files with 3107 additions and 461 deletions
@@ -36,8 +36,8 @@ use crate::provider_key_auth::{
use crate::provider_transport::antigravity::{
build_antigravity_safe_v1internal_request, build_antigravity_static_identity_headers,
classify_local_antigravity_request_support, AntigravityEnvelopeRequestType,
AntigravityRequestEnvelopeSupport, AntigravityRequestSideSupport,
AntigravityRequestSideUnsupportedReason,
AntigravityRequestAuthUnsupportedReason, AntigravityRequestEnvelopeSupport,
AntigravityRequestSideSupport, AntigravityRequestSideUnsupportedReason,
};
use crate::provider_transport::kiro::{
build_kiro_generate_assistant_response_url, build_kiro_provider_headers,
@@ -2061,7 +2061,7 @@ async fn provider_query_execute_openai_image_test_candidate(
trace_id: &str,
requested_model: &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 {
@@ -2429,7 +2429,7 @@ async fn provider_query_execute_antigravity_test_candidate(
trace_id: &str,
requested_model: &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 {
@@ -2463,11 +2463,34 @@ async fn provider_query_execute_antigravity_test_candidate(
));
};
let antigravity_spec = match classify_local_antigravity_request_support(
let mut antigravity_support = classify_local_antigravity_request_support(
&transport,
&base_provider_request_body,
AntigravityEnvelopeRequestType::EndpointTest,
);
if matches!(
antigravity_support,
AntigravityRequestSideSupport::Unsupported(
AntigravityRequestSideUnsupportedReason::UnsupportedAuth(
AntigravityRequestAuthUnsupportedReason::MissingProjectId
)
)
) {
if let Some(hydrated) = state
.app()
.hydrate_antigravity_project_metadata_for_transport(&transport)
.await
{
transport = hydrated;
antigravity_support = classify_local_antigravity_request_support(
&transport,
&base_provider_request_body,
AntigravityEnvelopeRequestType::EndpointTest,
);
}
}
let antigravity_spec = match antigravity_support {
AntigravityRequestSideSupport::Supported(spec) => spec,
AntigravityRequestSideSupport::Unsupported(reason) => {
let reason = provider_query_antigravity_unsupported_reason(reason);
@@ -2,7 +2,8 @@ use super::DEFAULT_PROVIDER_QUERY_TEST_MESSAGE;
use crate::handlers::admin::request::{AdminAppState, AdminGatewayProviderTransportSnapshot};
use crate::provider_transport::antigravity::{
classify_local_antigravity_request_support, AntigravityEnvelopeRequestType,
AntigravityRequestSideSupport, AntigravityRequestSideUnsupportedReason,
AntigravityRequestAuthUnsupportedReason, AntigravityRequestSideSupport,
AntigravityRequestSideUnsupportedReason,
};
use crate::provider_transport::kiro::supports_local_kiro_request_transport_with_network;
use serde_json::{json, Value};
@@ -82,14 +83,14 @@ pub(super) fn provider_query_standard_test_unsupported_reason(
api_format,
)
}
"gemini:generate_content" | "gemini:embedding"
"gemini:generate_content" | "gemini:embedding" | "gemini:interactions"
if crate::provider_transport::is_vertex_transport_context(transport) =>
{
aether_provider_transport::vertex::local_vertex_gemini_transport_unsupported_reason_with_network(
transport,
)
}
"gemini:generate_content" | "gemini:embedding" => {
"gemini:generate_content" | "gemini:embedding" | "gemini:interactions" => {
crate::provider_transport::policy::local_gemini_transport_unsupported_reason_with_network(
transport,
api_format,
@@ -130,9 +131,21 @@ pub(super) fn provider_query_antigravity_unsupported_reason(
AntigravityRequestSideUnsupportedReason::UnsupportedNetworkConfig => {
"transport_network_config_unsupported"
}
AntigravityRequestSideUnsupportedReason::UnsupportedAuth(_) => {
"transport_antigravity_auth_unsupported"
}
AntigravityRequestSideUnsupportedReason::UnsupportedAuth(
AntigravityRequestAuthUnsupportedReason::WrongProviderType,
) => "transport_provider_type_unsupported",
AntigravityRequestSideUnsupportedReason::UnsupportedAuth(
AntigravityRequestAuthUnsupportedReason::MissingAuthConfig,
) => "transport_antigravity_auth_config_missing",
AntigravityRequestSideUnsupportedReason::UnsupportedAuth(
AntigravityRequestAuthUnsupportedReason::InvalidAuthConfigJson,
) => "transport_antigravity_auth_config_invalid",
AntigravityRequestSideUnsupportedReason::UnsupportedAuth(
AntigravityRequestAuthUnsupportedReason::ComplexDynamicAuthConfig,
) => "transport_antigravity_auth_config_unsupported",
AntigravityRequestSideUnsupportedReason::UnsupportedAuth(
AntigravityRequestAuthUnsupportedReason::MissingProjectId,
) => "transport_antigravity_project_id_missing",
AntigravityRequestSideUnsupportedReason::UnsupportedEnvelope(_) => {
"transport_antigravity_envelope_unsupported"
}
@@ -250,6 +263,7 @@ pub(super) fn provider_query_test_adapter_for_provider_api_format(
| "openai:responses:compact"
| "claude:messages"
| "gemini:generate_content"
| "gemini:interactions"
| "openai:embedding"
| "gemini:embedding"
| "jina:embedding"
@@ -366,7 +380,7 @@ pub(super) fn provider_query_transport_supports_model_test_execution(
transport, api_format,
)
}
"gemini:generate_content" | "gemini:embedding" => {
"gemini:generate_content" | "gemini:embedding" | "gemini:interactions" => {
if crate::provider_transport::is_vertex_transport_context(transport) {
aether_provider_transport::vertex::supports_local_vertex_gemini_transport_with_network(transport)
} else {
@@ -120,6 +120,10 @@ fn provider_query_endpoint_route_payload(
("Gemini API", "gemini_native", "generateContent", "")
}
"gemini:generate_content" => ("Gemini native", "gemini_native", "generateContent", ""),
"gemini:interactions" if is_gemini_api => {
("Gemini API", "gemini_native", "interactions", "")
}
"gemini:interactions" => ("Gemini native", "gemini_native", "interactions", ""),
"openai:embedding" if is_vertex && is_openai_compat => (
"Vertex AI OpenAI-compatible",
"openai_compatible",
@@ -678,6 +678,10 @@ fn provider_query_test_adapter_routes_fixed_provider_endpoint_types() {
provider_query_test_adapter_for_provider_api_format("custom", "gemini:embedding"),
Some(ProviderQueryTestAdapter::Standard)
);
assert_eq!(
provider_query_test_adapter_for_provider_api_format("custom", "gemini:interactions"),
Some(ProviderQueryTestAdapter::Standard)
);
assert_eq!(
provider_query_test_adapter_for_provider_api_format(
"aliyun",
@@ -926,6 +926,9 @@ fn maybe_build_local_antigravity_v1internal_response(
"fetch_available_models" => {
Some(Json(build_antigravity_fetch_available_models_payload()).into_response())
}
"retrieve_user_quota_summary" => {
Some(Json(build_antigravity_retrieve_user_quota_summary_payload()).into_response())
}
"fetch_user_info" => {
Some(Json(build_antigravity_fetch_user_info_payload()).into_response())
}
@@ -938,6 +941,7 @@ fn maybe_build_local_antigravity_v1internal_response(
.into_response(),
),
"record_code_assist_metrics" => Some(Json(json!({})).into_response()),
"write_trajectory_acls" => Some(Json(json!({})).into_response()),
"set_user_settings" => Some(build_antigravity_set_user_settings_response(request_body)),
"stream_generate_content" => None,
_ => None,
@@ -1042,27 +1046,42 @@ fn build_antigravity_fetch_user_info_payload() -> Value {
})
}
fn build_antigravity_retrieve_user_quota_summary_payload() -> Value {
json!({
"description": "",
"groups": []
})
}
fn build_antigravity_default_user_settings_payload() -> Value {
json!({
"preferredModelId": "gemini-3.1-flash-lite"
"preferredModelId": "gemini-3.5-flash-low"
})
}
fn build_antigravity_fetch_available_models_payload() -> Value {
json!({
"models": {
"gemini-3.5-flash-low": antigravity_model_payload("gemini-3.5-flash-low", "Gemini 3.5 Flash Low"),
"gemini-3-flash-agent": antigravity_model_payload("gemini-3-flash-agent", "Gemini 3 Flash Agent"),
"gemini-pro-agent": antigravity_model_payload("gemini-pro-agent", "Gemini 3.1 Pro (High)"),
"gemini-3.1-pro-low": antigravity_model_payload("gemini-3.1-pro-low", "Gemini 3.1 Pro (Low)"),
"gemini-3-flash-agent": antigravity_model_payload("gemini-3-flash-agent", "Gemini 3.5 Flash (High)"),
"gemini-3.5-flash-low": antigravity_model_payload("gemini-3.5-flash-low", "Gemini 3.5 Flash (Medium)"),
"gemini-3.5-flash-extra-low": antigravity_model_payload("gemini-3.5-flash-extra-low", "Gemini 3.5 Flash (Low)"),
"claude-opus-4-6-thinking": antigravity_model_payload("claude-opus-4-6-thinking", "Claude Opus 4.6 (Thinking)"),
"claude-sonnet-4-6": antigravity_model_payload("claude-sonnet-4-6", "Claude Sonnet 4.6 (Thinking)"),
"gpt-oss-120b-medium": antigravity_model_payload("gpt-oss-120b-medium", "GPT-OSS 120B (Medium)"),
"gemini-3.1-flash-lite": antigravity_model_payload("gemini-3.1-flash-lite", "Gemini 3.1 Flash Lite"),
"gemini-3.1-pro-low": antigravity_model_payload("gemini-3.1-pro-low", "Gemini 3.1 Pro Low"),
"gemini-3-flash": antigravity_model_payload("gemini-3-flash", "Gemini 3 Flash"),
"gemini-2.5-flash": antigravity_model_payload("gemini-2.5-flash", "Gemini 2.5 Flash"),
"gemini-2.5-flash-lite": antigravity_model_payload("gemini-2.5-flash-lite", "Gemini 2.5 Flash Lite"),
"gemini-2.5-flash-thinking": antigravity_model_payload("gemini-2.5-flash-thinking", "Gemini 2.5 Flash Thinking"),
"gemini-2.5-flash": antigravity_model_payload("gemini-2.5-flash", "Gemini 3.1 Flash Lite"),
"gemini-2.5-flash-lite": antigravity_model_payload("gemini-2.5-flash-lite", "Gemini 3.1 Flash Lite"),
"gemini-2.5-flash-thinking": antigravity_model_payload("gemini-2.5-flash-thinking", "Gemini 3.1 Flash Lite"),
"gemini-2.5-pro": antigravity_model_payload("gemini-2.5-pro", "Gemini 2.5 Pro"),
"gemini-3.1-flash-image": antigravity_model_payload("gemini-3.1-flash-image", "Gemini 3.1 Flash Image"),
"tab_flash_lite_preview": antigravity_model_payload("tab_flash_lite_preview", "Tab Flash Lite Preview"),
"tab_jump_flash_lite_preview": antigravity_model_payload("tab_jump_flash_lite_preview", "Tab Jump Flash Lite Preview"),
"gemini-3.1-pro-high": antigravity_model_payload("gemini-3.1-pro-high", "Gemini 3.1 Pro (High)"),
"chat_20706": antigravity_model_payload("chat_20706", ""),
"chat_23310": antigravity_model_payload("chat_23310", ""),
"tab_flash_lite_preview": antigravity_model_payload("tab_flash_lite_preview", ""),
"tab_jump_flash_lite_preview": antigravity_model_payload("tab_jump_flash_lite_preview", ""),
"models/proactive-observer": antigravity_model_payload("models/proactive-observer", "Proactive Observer")
},
"agentModelSorts": [
@@ -1071,10 +1090,14 @@ fn build_antigravity_fetch_available_models_payload() -> Value {
"groups": [
{
"modelIds": [
"gemini-3.1-flash-lite",
"gemini-3.5-flash-low",
"gemini-3-flash-agent",
"gemini-3.5-flash-extra-low",
"gemini-3.1-pro-low",
"gemini-3.5-flash-low"
"gemini-pro-agent",
"claude-sonnet-4-6",
"claude-opus-4-6-thinking",
"gpt-oss-120b-medium"
]
}
]
@@ -1083,12 +1106,18 @@ fn build_antigravity_fetch_available_models_payload() -> Value {
"audioTranscriptionModelIds": ["models/proactive-observer"],
"commandModelIds": ["gemini-3-flash"],
"commitMessageModelIds": ["gemini-3.1-flash-lite"],
"defaultAgentModelId": "gemini-3.1-flash-lite",
"deprecatedModelIds": {},
"defaultAgentModelId": "gemini-3.5-flash-low",
"deprecatedModelIds": {
"gemini-3.1-pro-high": {
"newModelEnum": "MODEL_PLACEHOLDER_M16",
"newModelId": "gemini-pro-agent",
"oldModelEnum": "MODEL_PLACEHOLDER_M37"
}
},
"experimentIds": [],
"imageGenerationModelIds": ["gemini-3.1-flash-image"],
"mqueryModelIds": ["gemini-3.1-flash-lite"],
"tabModelIds": ["tab_flash_lite_preview", "tab_jump_flash_lite_preview"],
"tabModelIds": ["chat_20706", "chat_23310"],
"tieredModelIds": {
"flash": ["gemini-3-flash-agent"],
"flashLite": ["gemini-3.1-flash-lite"],
@@ -1100,6 +1129,11 @@ fn build_antigravity_fetch_available_models_payload() -> Value {
fn antigravity_model_payload(id: &str, display_name: &str) -> Value {
let model = match id {
"chat_20706" => "MODEL_CHAT_20706",
"chat_23310" => "MODEL_CHAT_23310",
"claude-opus-4-6-thinking" => "MODEL_PLACEHOLDER_M26",
"claude-sonnet-4-6" => "MODEL_PLACEHOLDER_M35",
"gpt-oss-120b-medium" => "MODEL_OPENAI_GPT_OSS_120B_MEDIUM",
"gemini-2.5-flash" => "MODEL_GOOGLE_GEMINI_2_5_FLASH",
"gemini-2.5-flash-lite" => "MODEL_GOOGLE_GEMINI_2_5_FLASH_LITE",
"gemini-2.5-flash-thinking" => "MODEL_GOOGLE_GEMINI_2_5_FLASH_THINKING",
@@ -1108,22 +1142,43 @@ fn antigravity_model_payload(id: &str, display_name: &str) -> Value {
"gemini-3-flash-agent" => "MODEL_PLACEHOLDER_M132",
"gemini-3.1-flash-image" => "MODEL_PLACEHOLDER_M21",
"gemini-3.1-flash-lite" => "MODEL_PLACEHOLDER_M50",
"gemini-pro-agent" => "MODEL_PLACEHOLDER_M16",
"gemini-3.1-pro-high" => "MODEL_PLACEHOLDER_M37",
"gemini-3.1-pro-low" => "MODEL_PLACEHOLDER_M36",
"gemini-3.5-flash-low" => "MODEL_PLACEHOLDER_M20",
"gemini-3.5-flash-extra-low" => "MODEL_PLACEHOLDER_M187",
"models/proactive-observer" => "MODEL_PLACEHOLDER_M70",
"tab_flash_lite_preview" => "MODEL_PLACEHOLDER_M19",
"tab_jump_flash_lite_preview" => "MODEL_PLACEHOLDER_M28",
_ => "MODEL_PLACEHOLDER_M20",
};
json!({
"apiProvider": "API_PROVIDER_GOOGLE_GEMINI",
let (api_provider, model_provider) = match id {
"claude-opus-4-6-thinking" | "claude-sonnet-4-6" => {
("API_PROVIDER_ANTHROPIC_VERTEX", "MODEL_PROVIDER_ANTHROPIC")
}
"gpt-oss-120b-medium" => ("API_PROVIDER_OPENAI_VERTEX", "MODEL_PROVIDER_OPENAI"),
"chat_20706" | "chat_23310" => ("API_PROVIDER_INTERNAL", "MODEL_PROVIDER_GOOGLE"),
_ => ("API_PROVIDER_GOOGLE_GEMINI", "MODEL_PROVIDER_GOOGLE"),
};
let mut payload = json!({
"apiProvider": api_provider,
"displayName": display_name,
"maxOutputTokens": 65536,
"maxTokens": 1048576,
"minThinkingBudget": 32,
"model": model,
"modelProvider": "MODEL_PROVIDER_GOOGLE",
"recommended": id == "gemini-3.1-flash-lite",
"modelProvider": model_provider,
"recommended": matches!(
id,
"gemini-3.5-flash-low"
| "gemini-3-flash-agent"
| "gemini-3.5-flash-extra-low"
| "gemini-3.1-pro-low"
| "gemini-pro-agent"
| "claude-sonnet-4-6"
| "claude-opus-4-6-thinking"
| "gpt-oss-120b-medium"
),
"supportedMimeTypes": {
"application/json": true,
"application/pdf": true,
@@ -1137,7 +1192,88 @@ fn antigravity_model_payload(id: &str, display_name: &str) -> Value {
"supportsVideo": true,
"thinkingBudget": 4000,
"tokenizerType": "LLAMA_WITH_SPECIAL"
})
});
if let Some(object) = payload.as_object_mut() {
match id {
"gemini-3-flash-agent" => {
object.insert("thinkingBudget".to_string(), json!(10000));
}
"gemini-pro-agent" => {
object.insert("maxOutputTokens".to_string(), json!(65535));
object.insert("thinkingBudget".to_string(), json!(10001));
}
"gemini-3.1-pro-high" => {
object.insert("maxOutputTokens".to_string(), json!(65535));
object.insert("minThinkingBudget".to_string(), json!(128));
object.insert("thinkingBudget".to_string(), json!(10001));
}
"gemini-3.5-flash-extra-low" => {
object.insert("thinkingBudget".to_string(), json!(1000));
object.insert("maxOutputTokens".to_string(), json!(65536));
}
"gemini-3.1-pro-low" => {
object.insert("maxOutputTokens".to_string(), json!(65535));
object.insert("minThinkingBudget".to_string(), json!(128));
object.insert("thinkingBudget".to_string(), json!(1001));
}
"claude-sonnet-4-6" | "claude-opus-4-6-thinking" => {
object.insert("maxTokens".to_string(), json!(250000));
object.insert("maxOutputTokens".to_string(), json!(64000));
object.insert("thinkingBudget".to_string(), json!(1024));
object.remove("minThinkingBudget");
object.remove("supportsVideo");
}
"gpt-oss-120b-medium" => {
object.insert("maxTokens".to_string(), json!(131072));
object.insert("maxOutputTokens".to_string(), json!(32768));
object.insert("thinkingBudget".to_string(), json!(8192));
object.remove("minThinkingBudget");
object.remove("supportsImages");
object.remove("supportsVideo");
}
"chat_20706" => {
object.insert("maxTokens".to_string(), json!(16384));
object.remove("displayName");
object.remove("maxOutputTokens");
object.remove("minThinkingBudget");
object.remove("recommended");
object.remove("supportsImages");
object.remove("supportsThinking");
object.remove("supportsVideo");
object.remove("thinkingBudget");
object.remove("tokenizerType");
object.remove("supportedMimeTypes");
}
"chat_23310" => {
object.insert("maxTokens".to_string(), json!(32768));
object.remove("displayName");
object.remove("maxOutputTokens");
object.remove("minThinkingBudget");
object.remove("recommended");
object.remove("supportsImages");
object.remove("supportsThinking");
object.remove("supportsVideo");
object.remove("thinkingBudget");
object.remove("tokenizerType");
object.remove("supportedMimeTypes");
}
"tab_flash_lite_preview" | "tab_jump_flash_lite_preview" => {
object.insert("maxTokens".to_string(), json!(16384));
object.insert("maxOutputTokens".to_string(), json!(4096));
object.remove("displayName");
object.remove("minThinkingBudget");
object.remove("recommended");
object.remove("supportsImages");
object.remove("supportsThinking");
object.remove("supportsVideo");
object.remove("thinkingBudget");
object.remove("tokenizerType");
object.remove("supportedMimeTypes");
}
_ => {}
}
}
payload
}
async fn maybe_build_local_gemini_video_operations_response(
@@ -102,7 +102,7 @@ pub(super) async fn maybe_build_local_test_connection_route_response(
let format_value = crate::ai_serving::normalize_api_format_alias(&format_value);
if !matches!(
format_value.as_str(),
"openai:chat" | "claude:messages" | "gemini:generate_content"
"openai:chat" | "claude:messages" | "gemini:generate_content" | "gemini:interactions"
) {
return None;
}
@@ -186,6 +186,10 @@ pub(super) async fn maybe_build_local_test_connection_route_response(
"parts": [{"text": "Health check"}],
}],
}),
"gemini:interactions" => json!({
"model": model,
"input": "Health check",
}),
_ => return None,
};
if !crate::provider_transport::apply_local_body_rules(
@@ -197,7 +201,7 @@ pub(super) async fn maybe_build_local_test_connection_route_response(
}
let oauth_auth = match format_value.as_str() {
"openai:chat" | "claude:messages" | "gemini:generate_content" => {
"openai:chat" | "claude:messages" | "gemini:generate_content" | "gemini:interactions" => {
match state.resolve_local_oauth_request_auth(&transport).await {
Ok(Some(crate::provider_transport::LocalResolvedOAuthRequestAuth::Header {
name,
@@ -218,7 +222,7 @@ pub(super) async fn maybe_build_local_test_connection_route_response(
crate::provider_transport::auth::resolve_local_standard_auth(&transport)
.or(oauth_auth.clone())
}
"gemini:generate_content" => {
"gemini:generate_content" | "gemini:interactions" => {
crate::provider_transport::auth::resolve_local_gemini_auth(&transport)
.or(oauth_auth.clone())
}
@@ -577,6 +577,44 @@ fn model_quota_window_snapshot(
Some(Value::Object(window))
}
fn canonical_antigravity_model_label(model_name: &str) -> Option<&'static str> {
match model_name.trim() {
"claude-opus-4-6-thinking" => Some("Claude Opus 4.6 (Thinking)"),
"claude-sonnet-4-6" | "claude-sonnet-4-6-thinking" => Some("Claude Sonnet 4.6 (Thinking)"),
"gemini-3-flash-agent" => Some("Gemini 3.5 Flash (High)"),
"gemini-3.5-flash-low" => Some("Gemini 3.5 Flash (Medium)"),
"gemini-3.5-flash-extra-low" => Some("Gemini 3.5 Flash (Low)"),
"gemini-3.1-pro-high" | "gemini-pro-agent" => Some("Gemini 3.1 Pro (High)"),
"gemini-3.1-pro-low" => Some("Gemini 3.1 Pro (Low)"),
"gemini-3.1-flash-image" => Some("Gemini 3.1 Flash Image"),
"gemini-3.1-flash-lite" => Some("Gemini 3.1 Flash Lite"),
"gemini-3-flash" => Some("Gemini 3 Flash"),
"gemini-2.5-pro" => Some("Gemini 2.5 Pro"),
"gemini-2.5-flash-thinking" | "gemini-2.5-flash" | "gemini-2.5-flash-lite" => {
Some("Gemini 3.1 Flash Lite")
}
"gpt-oss-120b-medium" => Some("GPT-OSS 120B (Medium)"),
"tab_flash_lite_preview" => Some("Tab Flash Lite Preview"),
"tab_jump_flash_lite_preview" => Some("Tab Jump Flash Lite Preview"),
"models/proactive-observer" => Some("Proactive Observer"),
_ => None,
}
}
fn antigravity_model_quota_window_snapshot(
model_name: &str,
item: &Map<String, Value>,
observed_at_unix_secs: Option<u64>,
) -> Option<Value> {
let mut window = model_quota_window_snapshot(model_name, item, observed_at_unix_secs)?;
if let Some(label) = canonical_antigravity_model_label(model_name) {
if let Some(window) = window.as_object_mut() {
window.insert("label".to_string(), json!(label));
}
}
Some(window)
}
fn provider_quota_metadata_string(
metadata: &Map<String, Value>,
fields: &[&str],
@@ -1466,7 +1504,7 @@ fn build_antigravity_quota_status_snapshot(
models
.iter()
.filter_map(|(model_name, item)| {
model_quota_window_snapshot(
antigravity_model_quota_window_snapshot(
model_name,
item.as_object()?,
observed_at_unix_secs,
@@ -3715,6 +3753,77 @@ mod tests {
);
}
#[test]
fn sync_provider_key_quota_status_snapshot_labels_antigravity_models_by_model_id() {
let upstream_metadata = json!({
"antigravity": {
"updated_at": 1_775_553_285u64,
"quota_by_model": {
"gemini-3.5-flash-extra-low": {
"display_name": "Gemini 3.5 Flash (Low)",
"remaining_fraction": 1.0
},
"gemini-3.5-flash-low": {
"display_name": "Gemini 3.5 Flash (Medium)",
"remaining_fraction": 0.75
},
"gemini-3-flash-agent": {
"display_name": "Gemini 3.5 Flash (High)",
"remaining_fraction": 0.5
},
"gemini-2.5-flash": {
"display_name": "Gemini 3.1 Flash Lite",
"remaining_fraction": 0.4
},
"claude-sonnet-4-6": {
"display_name": "Claude Sonnet 4.6 (Thinking)",
"remaining_fraction": 0.3
}
}
}
});
let payload = sync_provider_key_quota_status_snapshot(
None,
"antigravity",
Some(&upstream_metadata),
"refresh_api",
)
.expect("quota snapshot should sync");
let windows = payload["quota"]["windows"]
.as_array()
.expect("quota windows should exist");
let label_for_model = |model: &str| {
windows
.iter()
.filter_map(Value::as_object)
.find(|window| window.get("model") == Some(&json!(model)))
.and_then(|window| window.get("label"))
.cloned()
};
assert_eq!(
label_for_model("gemini-3.5-flash-extra-low"),
Some(json!("Gemini 3.5 Flash (Low)"))
);
assert_eq!(
label_for_model("gemini-3.5-flash-low"),
Some(json!("Gemini 3.5 Flash (Medium)"))
);
assert_eq!(
label_for_model("gemini-3-flash-agent"),
Some(json!("Gemini 3.5 Flash (High)"))
);
assert_eq!(
label_for_model("gemini-2.5-flash"),
Some(json!("Gemini 3.1 Flash Lite"))
);
assert_eq!(
label_for_model("claude-sonnet-4-6"),
Some(json!("Claude Sonnet 4.6 (Thinking)"))
);
}
#[test]
fn provider_key_status_snapshot_payload_backfills_account_block_from_oauth_invalid_reason() {
let mut key = sample_catalog_key();