i/Aether
1
0
Fork 0
mirror of https://github.com/fawney19/Aether.git synced 2026-10-04 08:27:46 +08:00
Code Issues Packages Projects Releases Wiki Activity

fix(providers): refresh Antigravity grouped quotas

Browse Source
This commit is contained in:
ZheFox
2026-09-03 10:46:44 +08:00
parent cc6f5e89b6
commit 45c840b8d3
20 changed files with 976 additions and 103 deletions
Download Patch File Download Diff File Expand all files Collapse all files
apps/aether-gateway/src/tests/control/admin/endpoints/quota.rs
+133 -31
View File
@@ -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();