mirror of
https://github.com/fawney19/Aether.git
synced 2026-10-07 18:07:47 +08:00
fix(providers): refresh Antigravity grouped quotas
This commit is contained in:
@@ -2741,7 +2741,7 @@ mod tests {
|
||||
.provider_request_headers
|
||||
.get("x-client-version")
|
||||
.map(String::as_str),
|
||||
Some("1.2.3")
|
||||
Some("4.3.0")
|
||||
);
|
||||
assert_eq!(
|
||||
payload
|
||||
@@ -2761,7 +2761,7 @@ mod tests {
|
||||
assert_eq!(payload.provider_request_body["model"], "gemini-2.5-pro");
|
||||
assert_eq!(
|
||||
payload.provider_request_body["userAgent"],
|
||||
"antigravity/cli/1.0.16 (aidev_client; os_type=linux; arch=arm64; auth_method=consumer)"
|
||||
"vscode/1.X.X (Antigravity/4.3.0)"
|
||||
);
|
||||
assert_eq!(payload.provider_request_body["requestType"], "agent");
|
||||
assert!(payload.provider_request_body.get("contents").is_none());
|
||||
|
||||
@@ -7,15 +7,20 @@ use super::shared::{
|
||||
};
|
||||
use crate::handlers::admin::request::{AdminAppState, AdminGatewayProviderTransportSnapshot};
|
||||
use crate::GatewayError;
|
||||
use aether_admin::provider::quota::parse_antigravity_usage_response;
|
||||
use aether_admin::provider::quota::{
|
||||
parse_antigravity_quota_summary_response, parse_antigravity_usage_response,
|
||||
};
|
||||
use aether_contracts::ProxySnapshot;
|
||||
use aether_data_contracts::repository::provider_catalog::{
|
||||
StoredProviderCatalogEndpoint, StoredProviderCatalogKey, StoredProviderCatalogProvider,
|
||||
};
|
||||
use aether_provider_pool::build_antigravity_pool_quota_request;
|
||||
use aether_provider_pool::{
|
||||
build_antigravity_pool_quota_request, build_antigravity_pool_quota_summary_request,
|
||||
};
|
||||
use serde_json::json;
|
||||
use std::collections::BTreeMap;
|
||||
use std::time::{SystemTime, UNIX_EPOCH};
|
||||
use tracing::warn;
|
||||
|
||||
async fn execute_antigravity_quota_plan(
|
||||
state: &AdminAppState<'_>,
|
||||
@@ -55,6 +60,79 @@ async fn execute_antigravity_quota_plan(
|
||||
execute_provider_quota_plan(state, transport, plan, "antigravity").await
|
||||
}
|
||||
|
||||
async fn fetch_antigravity_quota_summary_best_effort(
|
||||
state: &AdminAppState<'_>,
|
||||
transport: &AdminGatewayProviderTransportSnapshot,
|
||||
authorization: (String, String),
|
||||
project_id: &str,
|
||||
identity_headers: BTreeMap<String, String>,
|
||||
proxy_override: Option<&ProxySnapshot>,
|
||||
) -> Option<serde_json::Value> {
|
||||
let mut request_project_id = Some(project_id);
|
||||
|
||||
loop {
|
||||
let proxy = match proxy_override {
|
||||
Some(proxy) => Some(proxy.clone()),
|
||||
None => {
|
||||
state
|
||||
.resolve_transport_proxy_snapshot_with_tunnel_affinity(transport)
|
||||
.await
|
||||
}
|
||||
};
|
||||
let timeouts = Some(resolve_provider_quota_execution_timeouts(
|
||||
state.resolve_transport_execution_timeouts(transport),
|
||||
proxy.as_ref(),
|
||||
));
|
||||
let spec = build_antigravity_pool_quota_summary_request(
|
||||
&transport.key.id,
|
||||
&transport.endpoint.base_url,
|
||||
authorization.clone(),
|
||||
request_project_id,
|
||||
identity_headers.clone(),
|
||||
);
|
||||
let plan = build_provider_quota_execution_plan(
|
||||
transport,
|
||||
spec,
|
||||
proxy,
|
||||
state.resolve_transport_profile(transport),
|
||||
timeouts,
|
||||
);
|
||||
let outcome = match execute_provider_quota_plan(state, transport, plan, "antigravity").await
|
||||
{
|
||||
Ok(outcome) => outcome,
|
||||
Err(error) => {
|
||||
warn!(error = ?error, "Antigravity grouped quota request failed");
|
||||
return None;
|
||||
}
|
||||
};
|
||||
let result = match outcome {
|
||||
ProviderQuotaExecutionOutcome::Response(result) => result,
|
||||
ProviderQuotaExecutionOutcome::Failure(detail) => {
|
||||
warn!(detail = %detail, "Antigravity grouped quota execution failed");
|
||||
return None;
|
||||
}
|
||||
};
|
||||
|
||||
if result.status_code == 200 {
|
||||
return result
|
||||
.body
|
||||
.as_ref()
|
||||
.and_then(|body| body.json_body.as_ref())
|
||||
.and_then(parse_antigravity_quota_summary_response);
|
||||
}
|
||||
if result.status_code == 403 && request_project_id.is_some() {
|
||||
request_project_id = None;
|
||||
continue;
|
||||
}
|
||||
|
||||
warn!(
|
||||
status_code = result.status_code,
|
||||
"Antigravity grouped quota request returned a non-success status"
|
||||
);
|
||||
return None;
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn refresh_antigravity_provider_quota_locally(
|
||||
state: &AdminAppState<'_>,
|
||||
provider: &StoredProviderCatalogProvider,
|
||||
@@ -130,9 +208,9 @@ pub(crate) async fn refresh_antigravity_provider_quota_locally(
|
||||
let result = match execute_antigravity_quota_plan(
|
||||
state,
|
||||
&transport,
|
||||
authorization,
|
||||
authorization.clone(),
|
||||
&project_id,
|
||||
identity_headers,
|
||||
identity_headers.clone(),
|
||||
proxy_override.as_ref(),
|
||||
)
|
||||
.await?
|
||||
@@ -168,9 +246,31 @@ pub(crate) async fn refresh_antigravity_provider_quota_locally(
|
||||
.as_ref()
|
||||
.and_then(|body| body.json_body.as_ref())
|
||||
{
|
||||
metadata_update = parse_antigravity_usage_response(body_json, now_unix_secs)
|
||||
.map(|metadata| json!({ "antigravity": metadata }));
|
||||
if metadata_update.is_some() {
|
||||
if let Some(mut metadata) =
|
||||
parse_antigravity_usage_response(body_json, now_unix_secs)
|
||||
{
|
||||
if let Some(metadata) = metadata.as_object_mut() {
|
||||
metadata.insert("project_id".to_string(), json!(project_id));
|
||||
}
|
||||
if let Some(quota_groups) = fetch_antigravity_quota_summary_best_effort(
|
||||
state,
|
||||
&transport,
|
||||
authorization,
|
||||
&project_id,
|
||||
identity_headers,
|
||||
proxy_override.as_ref(),
|
||||
)
|
||||
.await
|
||||
{
|
||||
if let Some(metadata) = metadata.as_object_mut() {
|
||||
metadata.insert("quota_groups".to_string(), quota_groups);
|
||||
metadata.insert(
|
||||
"quota_groups_updated_at".to_string(),
|
||||
json!(now_unix_secs),
|
||||
);
|
||||
}
|
||||
}
|
||||
metadata_update = Some(json!({ "antigravity": metadata }));
|
||||
status = "success".to_string();
|
||||
} else {
|
||||
status = "no_metadata".to_string();
|
||||
|
||||
@@ -654,6 +654,112 @@ fn antigravity_model_quota_window_snapshot(
|
||||
Some(window)
|
||||
}
|
||||
|
||||
fn antigravity_grouped_quota_window_snapshots(
|
||||
metadata: &Map<String, Value>,
|
||||
observed_at_unix_secs: Option<u64>,
|
||||
) -> Vec<Value> {
|
||||
let mut windows = Vec::new();
|
||||
let Some(groups) = metadata.get("quota_groups").and_then(Value::as_array) else {
|
||||
return windows;
|
||||
};
|
||||
|
||||
for (group_index, group) in groups.iter().filter_map(Value::as_object).enumerate() {
|
||||
let group_code = group
|
||||
.get("group_id")
|
||||
.or_else(|| group.get("groupId"))
|
||||
.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.map(ToOwned::to_owned)
|
||||
.unwrap_or_else(|| format!("group:{group_index}"));
|
||||
let group_label = group
|
||||
.get("display_name")
|
||||
.or_else(|| group.get("displayName"))
|
||||
.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.map(ToOwned::to_owned)
|
||||
.unwrap_or_else(|| format!("Quota group {}", group_index + 1));
|
||||
let group_description = group
|
||||
.get("description")
|
||||
.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.map(ToOwned::to_owned);
|
||||
|
||||
for (bucket_index, bucket) in group
|
||||
.get("buckets")
|
||||
.and_then(Value::as_array)
|
||||
.into_iter()
|
||||
.flatten()
|
||||
.filter_map(Value::as_object)
|
||||
.enumerate()
|
||||
{
|
||||
let bucket_id = bucket
|
||||
.get("bucket_id")
|
||||
.or_else(|| bucket.get("bucketId"))
|
||||
.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.map(ToOwned::to_owned)
|
||||
.unwrap_or_else(|| format!("bucket-{}", bucket_index + 1));
|
||||
let Some(mut window) =
|
||||
model_quota_window_snapshot(&bucket_id, bucket, observed_at_unix_secs)
|
||||
else {
|
||||
continue;
|
||||
};
|
||||
let Some(window) = window.as_object_mut() else {
|
||||
continue;
|
||||
};
|
||||
let period = bucket
|
||||
.get("window")
|
||||
.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.map(ToOwned::to_owned);
|
||||
let bucket_detail = bucket
|
||||
.get("display_name")
|
||||
.or_else(|| bucket.get("displayName"))
|
||||
.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.map(ToOwned::to_owned)
|
||||
.or_else(|| period.clone());
|
||||
let bucket_label = bucket_detail
|
||||
.filter(|detail| !detail.eq_ignore_ascii_case(&group_label))
|
||||
.map(|detail| format!("{group_label} · {detail}"))
|
||||
.unwrap_or_else(|| group_label.clone());
|
||||
let description = bucket
|
||||
.get("description")
|
||||
.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.map(ToOwned::to_owned)
|
||||
.or_else(|| group_description.clone());
|
||||
|
||||
window.insert(
|
||||
"code".to_string(),
|
||||
json!(format!("group:{group_index}:{bucket_id}")),
|
||||
);
|
||||
window.insert("label".to_string(), json!(bucket_label));
|
||||
window.insert("scope".to_string(), json!("quota_group"));
|
||||
window.remove("model");
|
||||
window.insert("quota_group".to_string(), json!(group_code));
|
||||
window.insert("quota_group_label".to_string(), json!(group_label));
|
||||
window.insert("bucket_id".to_string(), json!(bucket_id));
|
||||
if let Some(period) = period {
|
||||
window.insert("window".to_string(), json!(period));
|
||||
}
|
||||
if let Some(description) = description {
|
||||
window.insert("description".to_string(), json!(description));
|
||||
}
|
||||
windows.push(Value::Object(window.clone()));
|
||||
}
|
||||
}
|
||||
|
||||
windows
|
||||
}
|
||||
|
||||
fn provider_quota_metadata_string(
|
||||
metadata: &Map<String, Value>,
|
||||
fields: &[&str],
|
||||
@@ -1622,7 +1728,7 @@ fn build_antigravity_quota_status_snapshot(
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.map(ToOwned::to_owned);
|
||||
let windows = provider_quota_model_bucket(metadata)
|
||||
let mut windows = provider_quota_model_bucket(metadata)
|
||||
.map(|models| {
|
||||
models
|
||||
.iter()
|
||||
@@ -1636,6 +1742,13 @@ fn build_antigravity_quota_status_snapshot(
|
||||
.collect::<Vec<_>>()
|
||||
})
|
||||
.unwrap_or_default();
|
||||
let grouped_observed_at_unix_secs =
|
||||
provider_quota_timestamp_unix_secs(metadata.get("quota_groups_updated_at"))
|
||||
.or(observed_at_unix_secs);
|
||||
windows.extend(antigravity_grouped_quota_window_snapshots(
|
||||
metadata,
|
||||
grouped_observed_at_unix_secs,
|
||||
));
|
||||
|
||||
if windows.is_empty() && observed_at_unix_secs.is_none() && !is_forbidden {
|
||||
return None;
|
||||
@@ -3132,6 +3245,52 @@ mod tests {
|
||||
assert_eq!(spark_weekly.get("remaining_ratio"), Some(&json!(0.95)));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn codex_wham_snapshot_does_not_duplicate_normalized_spark_windows() {
|
||||
let codex = admin_provider_quota_pure::parse_codex_wham_usage_response(
|
||||
&json!({
|
||||
"plan_type": "plus",
|
||||
"rate_limit": {
|
||||
"primary_window": {"used_percent": 25.0},
|
||||
"secondary_window": {"used_percent": 10.0}
|
||||
},
|
||||
"additional_rate_limits": [{
|
||||
"limit_name": "GPT-5.3-Codex-Spark",
|
||||
"rate_limit": {
|
||||
"primary_window": {
|
||||
"used_percent": 40.0,
|
||||
"limit_window_seconds": 18_000
|
||||
},
|
||||
"secondary_window": {
|
||||
"used_percent": 5.0,
|
||||
"limit_window_seconds": 604_800
|
||||
}
|
||||
}
|
||||
}]
|
||||
}),
|
||||
1_777_000_000,
|
||||
)
|
||||
.expect("Codex WHAM quota should parse");
|
||||
let upstream_metadata = json!({"codex": codex});
|
||||
let payload = sync_provider_key_quota_status_snapshot(
|
||||
None,
|
||||
"codex",
|
||||
Some(&upstream_metadata),
|
||||
"refresh_api",
|
||||
)
|
||||
.expect("Codex quota snapshot should sync");
|
||||
let windows = payload["quota"]["windows"]
|
||||
.as_array()
|
||||
.expect("quota windows should exist");
|
||||
let spark_codes = windows
|
||||
.iter()
|
||||
.filter_map(|window| window.get("code").and_then(Value::as_str))
|
||||
.filter(|code| code.starts_with("spark_"))
|
||||
.collect::<Vec<_>>();
|
||||
|
||||
assert_eq!(spark_codes, vec!["spark_5h", "spark_weekly"]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn provider_key_status_snapshot_payload_keeps_codex_free_window_quota_available() {
|
||||
let mut key = sample_catalog_key();
|
||||
@@ -4166,6 +4325,68 @@ mod tests {
|
||||
assert_eq!(claude_window.get("reset_seconds"), Some(&json!(12_011u64)));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn sync_provider_key_quota_status_snapshot_materializes_antigravity_group_windows() {
|
||||
let upstream_metadata = json!({
|
||||
"antigravity": {
|
||||
"updated_at": 1_777_000_000u64,
|
||||
"models": {
|
||||
"gemini-3.7-flash-tiered": {
|
||||
"remaining_fraction": 0.9
|
||||
}
|
||||
},
|
||||
"quota_groups": [{
|
||||
"display_name": "Claude and GPT models",
|
||||
"description": "Shared quota",
|
||||
"buckets": [{
|
||||
"bucket_id": "3p-5h",
|
||||
"window": "5h",
|
||||
"remaining_fraction": 0.25,
|
||||
"reset_time": "2026-05-05T05:00:00Z",
|
||||
"display_name": "5 hour"
|
||||
}, {
|
||||
"bucket_id": "3p-weekly",
|
||||
"window": "weekly",
|
||||
"remaining_fraction": 0.8,
|
||||
"reset_time": "2026-05-11T00:00:00Z"
|
||||
}]
|
||||
}]
|
||||
}
|
||||
});
|
||||
|
||||
let payload = sync_provider_key_quota_status_snapshot(
|
||||
None,
|
||||
"antigravity",
|
||||
Some(&upstream_metadata),
|
||||
"refresh_api",
|
||||
)
|
||||
.expect("Antigravity quota snapshot should sync");
|
||||
let windows = payload["quota"]["windows"]
|
||||
.as_array()
|
||||
.expect("quota windows should exist");
|
||||
let five_hour = windows
|
||||
.iter()
|
||||
.find(|window| window["code"] == "group:0:3p-5h")
|
||||
.expect("5h grouped quota window should exist");
|
||||
let weekly = windows
|
||||
.iter()
|
||||
.find(|window| window["code"] == "group:0:3p-weekly")
|
||||
.expect("weekly grouped quota window should exist");
|
||||
|
||||
assert_eq!(windows.len(), 3);
|
||||
assert_eq!(five_hour["scope"], json!("quota_group"));
|
||||
assert_eq!(five_hour["bucket_id"], json!("3p-5h"));
|
||||
assert_eq!(five_hour["label"], json!("Claude and GPT models · 5 hour"));
|
||||
assert_eq!(
|
||||
five_hour["quota_group_label"],
|
||||
json!("Claude and GPT models")
|
||||
);
|
||||
assert_eq!(five_hour["remaining_ratio"], json!(0.25));
|
||||
assert_eq!(five_hour["used_ratio"], json!(0.75));
|
||||
assert!(five_hour.get("model").is_none());
|
||||
assert_eq!(weekly["window"], json!("weekly"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn provider_key_status_snapshot_payload_backfills_account_block_from_oauth_invalid_reason() {
|
||||
let mut key = sample_catalog_key();
|
||||
|
||||
@@ -1949,7 +1949,7 @@ async fn gateway_executes_openai_chat_antigravity_cross_format_sync_via_local_fi
|
||||
"Bearer imported-antigravity-chat-token"
|
||||
);
|
||||
assert_eq!(seen_execution_runtime_request.x_client_name, "antigravity");
|
||||
assert_eq!(seen_execution_runtime_request.x_client_version, "1.2.3");
|
||||
assert_eq!(seen_execution_runtime_request.x_client_version, "4.3.0");
|
||||
assert_eq!(
|
||||
seen_execution_runtime_request.x_vscode_sessionid,
|
||||
"sess-antigravity-chat-local-123"
|
||||
|
||||
@@ -1537,7 +1537,7 @@ async fn gateway_executes_openai_responses_antigravity_cross_format_upstream_str
|
||||
);
|
||||
assert_eq!(
|
||||
seen_remote_execution_runtime_request.x_client_version,
|
||||
"1.2.3"
|
||||
"4.3.0"
|
||||
);
|
||||
assert_eq!(
|
||||
seen_remote_execution_runtime_request.x_vscode_sessionid,
|
||||
|
||||
@@ -2123,7 +2123,7 @@ async fn gateway_executes_antigravity_gemini_cli_sync_upstream_stream_via_local_
|
||||
);
|
||||
assert_eq!(
|
||||
seen_remote_execution_runtime_request.x_client_version,
|
||||
"1.2.3"
|
||||
"4.3.0"
|
||||
);
|
||||
assert_eq!(
|
||||
seen_remote_execution_runtime_request.x_vscode_sessionid,
|
||||
|
||||
@@ -2017,7 +2017,7 @@ async fn gateway_executes_antigravity_gemini_cli_stream_via_local_decision_gate_
|
||||
"Bearer refreshed-antigravity-cli-stream-access-token"
|
||||
);
|
||||
assert_eq!(seen_execution_runtime_request.x_client_name, "antigravity");
|
||||
assert_eq!(seen_execution_runtime_request.x_client_version, "1.2.3");
|
||||
assert_eq!(seen_execution_runtime_request.x_client_version, "4.3.0");
|
||||
assert_eq!(
|
||||
seen_execution_runtime_request.x_vscode_sessionid,
|
||||
"sess-antigravity-stream-local-123"
|
||||
|
||||
@@ -2321,7 +2321,7 @@ async fn gateway_executes_antigravity_gemini_cli_sync_via_local_decision_gate_af
|
||||
"Bearer refreshed-antigravity-cli-access-token"
|
||||
);
|
||||
assert_eq!(seen_execution_runtime_request.x_client_name, "antigravity");
|
||||
assert_eq!(seen_execution_runtime_request.x_client_version, "1.2.3");
|
||||
assert_eq!(seen_execution_runtime_request.x_client_version, "4.3.0");
|
||||
assert_eq!(
|
||||
seen_execution_runtime_request.x_vscode_sessionid,
|
||||
"sess-antigravity-local-123"
|
||||
|
||||
@@ -2264,6 +2264,8 @@ async fn gateway_refreshes_admin_provider_quota_locally_for_antigravity_with_tru
|
||||
struct SeenExecutionRuntimeRequest {
|
||||
url: String,
|
||||
authorization: String,
|
||||
user_agent: String,
|
||||
x_client_version: String,
|
||||
provider_api_format: String,
|
||||
request_body: Option<serde_json::Value>,
|
||||
}
|
||||
@@ -2281,7 +2283,7 @@ async fn gateway_refreshes_admin_provider_quota_locally_for_antigravity_with_tru
|
||||
}),
|
||||
);
|
||||
|
||||
let seen_execution_runtime = Arc::new(Mutex::new(None::<SeenExecutionRuntimeRequest>));
|
||||
let seen_execution_runtime = Arc::new(Mutex::new(Vec::<SeenExecutionRuntimeRequest>::new()));
|
||||
let seen_execution_runtime_clone = Arc::clone(&seen_execution_runtime);
|
||||
let execution_runtime = Router::new().route(
|
||||
"/v1/execute/sync",
|
||||
@@ -2294,26 +2296,30 @@ async fn gateway_refreshes_admin_provider_quota_locally_for_antigravity_with_tru
|
||||
.expect("body should read"),
|
||||
)
|
||||
.expect("plan should parse");
|
||||
*seen_execution_runtime_inner
|
||||
let request_body = plan.body.json_body.clone();
|
||||
seen_execution_runtime_inner
|
||||
.lock()
|
||||
.expect("mutex should lock") = Some(SeenExecutionRuntimeRequest {
|
||||
url: plan.url.clone(),
|
||||
authorization: plan
|
||||
.headers
|
||||
.get("authorization")
|
||||
.cloned()
|
||||
.unwrap_or_default(),
|
||||
provider_api_format: plan.provider_api_format.clone(),
|
||||
request_body: plan.body.json_body.clone(),
|
||||
});
|
||||
let result = aether_contracts::ExecutionResult {
|
||||
request_id: plan.request_id,
|
||||
candidate_id: None,
|
||||
status_code: 200,
|
||||
headers: BTreeMap::new(),
|
||||
response_observation: None,
|
||||
body: Some(aether_contracts::ResponseBody {
|
||||
json_body: Some(json!({
|
||||
.expect("mutex should lock")
|
||||
.push(SeenExecutionRuntimeRequest {
|
||||
url: plan.url.clone(),
|
||||
authorization: plan
|
||||
.headers
|
||||
.get("authorization")
|
||||
.cloned()
|
||||
.unwrap_or_default(),
|
||||
user_agent: plan.headers.get("user-agent").cloned().unwrap_or_default(),
|
||||
x_client_version: plan
|
||||
.headers
|
||||
.get("x-client-version")
|
||||
.cloned()
|
||||
.unwrap_or_default(),
|
||||
provider_api_format: plan.provider_api_format.clone(),
|
||||
request_body: request_body.clone(),
|
||||
});
|
||||
let (status_code, json_body) = match plan.provider_api_format.as_str() {
|
||||
"antigravity:fetch_available_models" => (
|
||||
200,
|
||||
json!({
|
||||
"models": {
|
||||
"claude-sonnet-4": {
|
||||
"displayName": "Claude Sonnet 4",
|
||||
@@ -2326,7 +2332,47 @@ async fn gateway_refreshes_admin_provider_quota_locally_for_antigravity_with_tru
|
||||
"displayName": "Gemini 2.5 Pro"
|
||||
}
|
||||
}
|
||||
})),
|
||||
}),
|
||||
),
|
||||
"antigravity:retrieve_user_quota_summary"
|
||||
if request_body
|
||||
.as_ref()
|
||||
.and_then(|body| body.get("project"))
|
||||
.is_some() =>
|
||||
{
|
||||
(403, json!({"error": {"message": "project not accepted"}}))
|
||||
}
|
||||
"antigravity:retrieve_user_quota_summary" => (
|
||||
200,
|
||||
json!({
|
||||
"groups": [{
|
||||
"displayName": "Claude and GPT models",
|
||||
"description": "Shared quota",
|
||||
"buckets": [{
|
||||
"bucketId": "3p-5h",
|
||||
"window": "5h",
|
||||
"remainingFraction": 0.25,
|
||||
"resetTime": "2026-05-05T05:00:00Z",
|
||||
"displayName": "5 hour"
|
||||
}, {
|
||||
"bucketId": "3p-weekly",
|
||||
"window": "weekly",
|
||||
"remainingFraction": 0.8,
|
||||
"resetTime": "2026-05-11T00:00:00Z"
|
||||
}]
|
||||
}]
|
||||
}),
|
||||
),
|
||||
unexpected => panic!("unexpected quota request format: {unexpected}"),
|
||||
};
|
||||
let result = aether_contracts::ExecutionResult {
|
||||
request_id: plan.request_id,
|
||||
candidate_id: None,
|
||||
status_code,
|
||||
headers: BTreeMap::new(),
|
||||
response_observation: None,
|
||||
body: Some(aether_contracts::ResponseBody {
|
||||
json_body: Some(json_body),
|
||||
body_bytes_b64: None,
|
||||
}),
|
||||
telemetry: None,
|
||||
@@ -2429,31 +2475,55 @@ async fn gateway_refreshes_admin_provider_quota_locally_for_antigravity_with_tru
|
||||
payload["results"][0]["quota_snapshot"]["windows"]
|
||||
.as_array()
|
||||
.map(Vec::len),
|
||||
Some(1usize)
|
||||
Some(3usize)
|
||||
);
|
||||
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
|
||||
|
||||
let seen_execution_runtime_request = seen_execution_runtime
|
||||
let seen_execution_runtime_requests = seen_execution_runtime
|
||||
.lock()
|
||||
.expect("mutex should lock")
|
||||
.clone()
|
||||
.expect("execution runtime request should be captured");
|
||||
.clone();
|
||||
assert_eq!(seen_execution_runtime_requests.len(), 3);
|
||||
let fetch_models_request = &seen_execution_runtime_requests[0];
|
||||
assert_eq!(
|
||||
seen_execution_runtime_request.url,
|
||||
fetch_models_request.url,
|
||||
"https://daily-cloudcode-pa.googleapis.com/v1internal:fetchAvailableModels"
|
||||
);
|
||||
assert_eq!(fetch_models_request.authorization, "Bearer ya29.ant-token");
|
||||
assert_eq!(
|
||||
seen_execution_runtime_request.authorization,
|
||||
"Bearer ya29.ant-token"
|
||||
fetch_models_request.user_agent,
|
||||
"vscode/1.X.X (Antigravity/4.3.0)"
|
||||
);
|
||||
assert_eq!(fetch_models_request.x_client_version, "4.3.0");
|
||||
assert_eq!(
|
||||
seen_execution_runtime_request.provider_api_format,
|
||||
fetch_models_request.provider_api_format,
|
||||
"antigravity:fetch_available_models"
|
||||
);
|
||||
assert_eq!(
|
||||
seen_execution_runtime_request.request_body,
|
||||
fetch_models_request.request_body,
|
||||
Some(json!({ "project": "project-ant-123" }))
|
||||
);
|
||||
let grouped_with_project = &seen_execution_runtime_requests[1];
|
||||
let grouped_without_project = &seen_execution_runtime_requests[2];
|
||||
assert_eq!(
|
||||
grouped_with_project.url,
|
||||
"https://daily-cloudcode-pa.googleapis.com/v1internal:retrieveUserQuotaSummary"
|
||||
);
|
||||
assert_eq!(
|
||||
grouped_with_project.provider_api_format,
|
||||
"antigravity:retrieve_user_quota_summary"
|
||||
);
|
||||
assert_eq!(
|
||||
grouped_with_project.request_body,
|
||||
Some(json!({"project": "project-ant-123"}))
|
||||
);
|
||||
assert_eq!(grouped_without_project.url, grouped_with_project.url);
|
||||
assert_eq!(grouped_without_project.request_body, Some(json!({})));
|
||||
assert!(seen_execution_runtime_requests.iter().all(|request| {
|
||||
request.authorization == "Bearer ya29.ant-token"
|
||||
&& request.user_agent == "vscode/1.X.X (Antigravity/4.3.0)"
|
||||
&& request.x_client_version == "4.3.0"
|
||||
}));
|
||||
|
||||
let reloaded = provider_catalog_repository
|
||||
.list_keys_by_ids(&["key-antigravity-a".to_string()])
|
||||
@@ -2481,6 +2551,26 @@ async fn gateway_refreshes_admin_provider_quota_locally_for_antigravity_with_tru
|
||||
.and_then(|value| value.get("used_percent")),
|
||||
Some(&json!(75.0))
|
||||
);
|
||||
assert_eq!(
|
||||
reloaded[0]
|
||||
.upstream_metadata
|
||||
.as_ref()
|
||||
.and_then(|value| value.pointer("/antigravity/quota_groups/0/buckets/0/bucket_id")),
|
||||
Some(&json!("3p-5h"))
|
||||
);
|
||||
assert_eq!(
|
||||
reloaded[0]
|
||||
.upstream_metadata
|
||||
.as_ref()
|
||||
.and_then(|value| value.pointer("/antigravity/project_id")),
|
||||
Some(&json!("project-ant-123"))
|
||||
);
|
||||
assert!(reloaded[0]
|
||||
.upstream_metadata
|
||||
.as_ref()
|
||||
.and_then(|value| value.pointer("/antigravity/quota_groups_updated_at"))
|
||||
.and_then(serde_json::Value::as_u64)
|
||||
.is_some());
|
||||
assert_eq!(
|
||||
reloaded[0]
|
||||
.status_snapshot
|
||||
@@ -2505,7 +2595,19 @@ async fn gateway_refreshes_admin_provider_quota_locally_for_antigravity_with_tru
|
||||
.and_then(|value| value.get("windows"))
|
||||
.and_then(|value| value.as_array())
|
||||
.map(Vec::len),
|
||||
Some(1usize)
|
||||
Some(3usize)
|
||||
);
|
||||
assert_eq!(
|
||||
reloaded[0]
|
||||
.status_snapshot
|
||||
.as_ref()
|
||||
.and_then(|value| value.pointer("/quota/windows"))
|
||||
.and_then(serde_json::Value::as_array)
|
||||
.and_then(|windows| windows
|
||||
.iter()
|
||||
.find(|window| window["code"] == "group:0:3p-5h"))
|
||||
.and_then(|window| window.get("remaining_ratio")),
|
||||
Some(&json!(0.25))
|
||||
);
|
||||
|
||||
gateway_handle.abort();
|
||||
|
||||
Reference in New Issue
Block a user