Add Antigravity v1internal gateway adapter

This commit is contained in:
MMEXA
2026-05-25 06:58:15 +08:00
parent 505d9fd8bc
commit 28c3a5dbe4
15 changed files with 978 additions and 33 deletions
@@ -17,6 +17,14 @@ const AI_POST_ROUTE_PATTERNS: &[&str] = &[
"/v1/responses/compact",
"/v1/images/generations",
"/v1/images/edits",
"/v1internal:loadCodeAssist",
"/v1internal:fetchAvailableModels",
"/v1internal:fetchUserInfo",
"/v1internal:fetchAdminControls",
"/v1internal:setUserSettings",
"/v1internal:listExperiments",
"/v1internal:recordCodeAssistMetrics",
"/v1internal:streamGenerateContent",
];
const AI_ANY_ROUTE_PATTERNS: &[&str] = &[
+8
View File
@@ -136,6 +136,14 @@ pub(crate) const RUST_FRONTDOOR_OWNED_ROUTE_PATTERNS: &[&str] = &[
"/upload/v1beta/files",
"/v1beta/files",
"/v1beta/files/{path...}",
"/v1internal:loadCodeAssist",
"/v1internal:fetchAvailableModels",
"/v1internal:fetchUserInfo",
"/v1internal:fetchAdminControls",
"/v1internal:setUserSettings",
"/v1internal:listExperiments",
"/v1internal:recordCodeAssistMetrics",
"/v1internal:streamGenerateContent",
"/",
"/{*path}",
];
@@ -186,6 +186,9 @@ fn select_primary_credential(
if signature.starts_with("gemini:") {
return select_gemini_credential(bundle);
}
if signature.starts_with("antigravity:") {
return select_antigravity_credential(bundle);
}
if signature.starts_with("claude:") {
return select_claude_messages_credential(bundle);
}
@@ -196,6 +199,20 @@ fn select_primary_credential(
select_generic_credential(bundle)
}
fn select_antigravity_credential(
bundle: &GatewayCredentialBundle,
) -> Option<GatewayPrimaryCredential> {
first_provider_api_key(
bundle,
&[
GatewayCredentialCarrier::XApiKey,
GatewayCredentialCarrier::ApiKey,
],
)
.or_else(|| first_bearer_token(bundle))
.or_else(|| select_cookie_credential(bundle))
}
fn select_openai_credential(bundle: &GatewayCredentialBundle) -> Option<GatewayPrimaryCredential> {
first_provider_api_key(
bundle,
@@ -459,6 +476,29 @@ mod tests {
);
}
#[test]
fn prefers_antigravity_aether_api_key_over_google_bearer() {
let mut headers = http::HeaderMap::new();
headers.insert(
http::header::AUTHORIZATION,
"Bearer google-oauth-access-token".parse().unwrap(),
);
headers.insert("x-api-key", "sk-aether-antigravity".parse().unwrap());
let extracted = extract_request_credentials(
&headers,
&uri("/v1internal:streamGenerateContent?alt=sse"),
"antigravity:v1internal",
);
assert_eq!(
extracted.primary,
Some(GatewayPrimaryCredential::ProviderApiKey {
raw: "sk-aether-antigravity".to_string(),
carrier: GatewayCredentialCarrier::XApiKey,
})
);
}
#[test]
fn prefers_gemini_query_key_over_header_key() {
let mut headers = http::HeaderMap::new();
@@ -712,7 +712,12 @@ async fn build_data_backed_auth_context(
})
} else if snapshot
.effective_allowed_api_formats()
.is_some_and(|allowed| !contains_api_format_or_alias(allowed, auth_endpoint_signature))
.is_some_and(|allowed| {
!contains_api_format_or_alias(
allowed,
auth_gate_api_format(auth_endpoint_signature).as_str(),
)
})
{
Some(GatewayLocalAuthRejection::ApiFormatNotAllowed {
api_format: auth_endpoint_signature.to_string(),
@@ -747,6 +752,15 @@ fn normalize_api_format_alias(value: &str) -> String {
crate::ai_serving::normalize_api_format_alias(value)
}
fn auth_gate_api_format(auth_endpoint_signature: &str) -> String {
let normalized = normalize_api_format_alias(auth_endpoint_signature);
if normalized == "antigravity:v1internal" {
"gemini:generate_content".to_string()
} else {
normalized
}
}
fn api_format_matches(left: &str, right: &str) -> bool {
aether_scheduler_core::api_format_matches_allowed_value(left, right)
}
@@ -1261,6 +1275,58 @@ mod tests {
assert_eq!(auth_context.local_rejection, None);
}
#[tokio::test]
async fn data_backed_auth_context_allows_antigravity_v1internal_for_gemini_generate_content_keys(
) {
let api_key = "sk-test-antigravity-v1internal";
let mut snapshot = sample_snapshot("key-ant-v1internal", "user-ant-v1internal");
snapshot.user_allowed_providers = Some(vec!["antigravity".to_string()]);
snapshot.api_key_allowed_providers = Some(vec!["antigravity".to_string()]);
snapshot.user_allowed_api_formats = Some(vec!["gemini:generate_content".to_string()]);
snapshot.api_key_allowed_api_formats = Some(vec!["gemini:generate_content".to_string()]);
let repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![(
Some(hash_api_key(api_key)),
snapshot,
)]));
let provider_catalog = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![sample_provider(
"provider-antigravity-1",
"Antigravity",
"antigravity",
)],
vec![sample_endpoint(
"endpoint-antigravity-1",
"provider-antigravity-1",
"gemini:generate_content",
)],
Vec::new(),
));
let data = GatewayDataState::with_auth_api_key_reader_for_tests(repository)
.with_provider_catalog_reader(provider_catalog);
let state = AppState::new()
.expect("state should build")
.with_data_state_for_tests(data);
let mut headers = HeaderMap::new();
headers.insert("x-api-key", api_key.parse().unwrap());
headers.insert(
http::header::AUTHORIZATION,
"Bearer google-oauth-access-token".parse().unwrap(),
);
let auth_context = resolve_data_backed_auth_context(
&state,
&headers,
&uri("/v1internal:streamGenerateContent?alt=sse"),
Some("antigravity:v1internal"),
)
.await
.expect("resolution should succeed")
.expect("auth context should exist");
assert_eq!(auth_context.local_rejection, None);
}
#[tokio::test]
async fn data_backed_auth_context_allows_provider_id_for_convertible_endpoint_format() {
let api_key = "sk-test-provider-convertible-endpoint";
+34 -1
View File
@@ -8,7 +8,9 @@ pub(super) fn classify_ai_public_route(
normalized_path: &str,
headers: &http::HeaderMap,
) -> Option<ClassifiedRoute> {
if method == http::Method::POST && normalized_path == "/v1/chat/completions" {
if let Some(route) = classify_antigravity_v1internal_route(method, normalized_path) {
Some(route)
} else if method == http::Method::POST && normalized_path == "/v1/chat/completions" {
Some(classified(
"ai_public",
"openai",
@@ -156,3 +158,34 @@ pub(super) fn classify_ai_public_route(
None
}
}
fn classify_antigravity_v1internal_route(
method: &http::Method,
normalized_path: &str,
) -> Option<ClassifiedRoute> {
if method != http::Method::POST {
return None;
}
let action = normalized_path.strip_prefix("/v1internal:")?;
let (route_kind, execution_runtime_candidate) = match action {
"loadCodeAssist" => ("load_code_assist", false),
"fetchAvailableModels" => ("fetch_available_models", false),
"fetchUserInfo" => ("fetch_user_info", false),
"fetchAdminControls" => ("fetch_admin_controls", false),
"setUserSettings" => ("set_user_settings", false),
"listExperiments" => ("list_experiments", false),
"recordCodeAssistMetrics" => ("record_code_assist_metrics", false),
"streamGenerateContent" => ("stream_generate_content", true),
_ => return None,
};
Some(classified_with_request_auth_channel(
"ai_public",
"antigravity",
route_kind,
"bearer_like",
"antigravity:v1internal",
execution_runtime_candidate,
))
}
@@ -274,3 +274,86 @@ fn classifies_gemini_predict_long_running_as_video_route() {
);
assert!(decision.is_execution_runtime_candidate());
}
#[test]
fn classifies_antigravity_v1internal_control_plane_routes() {
let headers = headers(&[
("authorization", "Bearer ant-access-token"),
("user-agent", "antigravity/cli/1.0.2 linux/arm64"),
]);
for (path, route_kind) in [
("/v1internal:loadCodeAssist", "load_code_assist"),
("/v1internal:fetchAvailableModels", "fetch_available_models"),
("/v1internal:fetchUserInfo", "fetch_user_info"),
("/v1internal:fetchAdminControls", "fetch_admin_controls"),
("/v1internal:setUserSettings", "set_user_settings"),
("/v1internal:listExperiments", "list_experiments"),
(
"/v1internal:recordCodeAssistMetrics",
"record_code_assist_metrics",
),
] {
let uri: Uri = path.parse().expect("uri should parse");
let decision = classify_control_route(&http::Method::POST, &uri, &headers)
.expect("route should classify");
assert_eq!(decision.route_class.as_deref(), Some("ai_public"));
assert_eq!(decision.route_family.as_deref(), Some("antigravity"));
assert_eq!(decision.route_kind.as_deref(), Some(route_kind));
assert_eq!(
decision.request_auth_channel.as_deref(),
Some("bearer_like")
);
assert_eq!(
decision.auth_endpoint_signature.as_deref(),
Some("antigravity:v1internal")
);
assert!(
!decision.is_execution_runtime_candidate(),
"control-plane route {path} must be handled by local facade before execution runtime"
);
}
}
#[test]
fn classifies_antigravity_stream_generate_content_as_execution_route() {
let headers = headers(&[
("authorization", "Bearer ant-access-token"),
("user-agent", "antigravity/cli/1.0.2 linux/arm64"),
]);
let uri: Uri = "/v1internal:streamGenerateContent?alt=sse"
.parse()
.expect("uri should parse");
let decision =
classify_control_route(&http::Method::POST, &uri, &headers).expect("route should classify");
assert_eq!(decision.route_class.as_deref(), Some("ai_public"));
assert_eq!(decision.route_family.as_deref(), Some("antigravity"));
assert_eq!(
decision.route_kind.as_deref(),
Some("stream_generate_content")
);
assert_eq!(
decision.request_auth_channel.as_deref(),
Some("bearer_like")
);
assert_eq!(
decision.auth_endpoint_signature.as_deref(),
Some("antigravity:v1internal")
);
assert!(decision.is_execution_runtime_candidate());
}
#[test]
fn rejects_unknown_antigravity_v1internal_route() {
let headers = headers(&[
("authorization", "Bearer ant-access-token"),
("user-agent", "antigravity/cli/1.0.2 linux/arm64"),
]);
let uri: Uri = "/v1internal:deleteEverything"
.parse()
.expect("uri should parse");
assert!(classify_control_route(&http::Method::POST, &uri, &headers).is_none());
}
@@ -1626,6 +1626,10 @@ mod tests {
8084,
"http://localhost:8084/v1/responses"
));
assert!(gateway_frontdoor_self_loop_guard_matches_with_port(
8084,
"http://localhost:8084/v1internal:streamGenerateContent?alt=sse"
));
}
#[test]
@@ -48,6 +48,7 @@ pub(crate) fn frontdoor_self_loop_public_ai_path(path: &str) -> bool {
) || path.starts_with("/v1/videos/")
|| path.starts_with("/v1beta/files/")
|| path.starts_with("/v1beta/operations/")
|| path.starts_with("/v1internal:")
|| is_gemini_generation_path(path)
}
@@ -52,6 +52,12 @@ const OPENAI_RERANK_TOP_N_DETAIL: &str = "Rerank request top_n must be a positiv
const OPENAI_RERANK_CHAT_PAYLOAD_DETAIL: &str =
"Rerank request must use query/documents, not chat messages";
const OPENAI_RERANK_STREAM_UNSUPPORTED_DETAIL: &str = "Rerank requests do not support streaming";
const ANTIGRAVITY_USER_SETTINGS_MISSING_BODY_DETAIL: &str =
"Antigravity setUserSettings request body is required";
const ANTIGRAVITY_USER_SETTINGS_INVALID_JSON_DETAIL: &str =
"Antigravity setUserSettings request JSON body is invalid";
const ANTIGRAVITY_USER_SETTINGS_INVALID_DETAIL: &str =
"Antigravity setUserSettings request must include object userSettings";
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
enum OpenAiImageOperation {
@@ -103,7 +109,9 @@ pub(crate) fn ai_public_local_requires_buffered_body(
&& request_context.request_path == "/v1/embeddings")
|| (decision.route_family.as_deref() == Some("openai")
&& decision.route_kind.as_deref() == Some("rerank")
&& request_context.request_path == "/v1/rerank"))
&& request_context.request_path == "/v1/rerank")
|| (decision.route_family.as_deref() == Some("antigravity")
&& decision.route_kind.as_deref() != Some("stream_generate_content")))
})
}
@@ -133,6 +141,12 @@ pub(crate) async fn maybe_build_local_ai_public_response(
return Some(response);
}
if let Some(response) =
maybe_build_local_antigravity_v1internal_response(request_context, request_body)
{
return Some(response);
}
maybe_build_local_gemini_video_operations_response(state, request_context, decision).await
}
@@ -861,6 +875,140 @@ fn maybe_build_local_claude_count_tokens_response(
Some(Json(json!({ "input_tokens": input_tokens })).into_response())
}
fn maybe_build_local_antigravity_v1internal_response(
request_context: &GatewayPublicRequestContext,
request_body: Option<&Bytes>,
) -> Option<Response<Body>> {
let decision = request_context.control_decision.as_ref()?;
if decision.route_family.as_deref() != Some("antigravity")
|| request_context.request_method != http::Method::POST
{
return None;
}
match decision.route_kind.as_deref()? {
"load_code_assist" => {
Some(Json(build_antigravity_load_code_assist_payload()).into_response())
}
"fetch_available_models" => {
Some(Json(build_antigravity_fetch_available_models_payload()).into_response())
}
"fetch_user_info" => {
Some(Json(build_antigravity_fetch_user_info_payload()).into_response())
}
"fetch_admin_controls" => Some(Json(json!({})).into_response()),
"list_experiments" => Some(
Json(json!({
"experimentIds": [],
"flags": {}
}))
.into_response(),
),
"record_code_assist_metrics" => Some(Json(json!({})).into_response()),
"set_user_settings" => Some(build_antigravity_set_user_settings_response(request_body)),
"stream_generate_content" => None,
_ => None,
}
}
fn build_antigravity_set_user_settings_response(request_body: Option<&Bytes>) -> Response<Body> {
let Some(request_body) = request_body else {
return build_ai_public_error_response(
http::StatusCode::BAD_REQUEST,
ANTIGRAVITY_USER_SETTINGS_MISSING_BODY_DETAIL,
);
};
let payload = match serde_json::from_slice::<Value>(request_body) {
Ok(payload) => payload,
Err(_) => {
return build_ai_public_error_response(
http::StatusCode::BAD_REQUEST,
ANTIGRAVITY_USER_SETTINGS_INVALID_JSON_DETAIL,
);
}
};
let Some(user_settings) = payload
.get("userSettings")
.filter(|value| value.is_object())
.cloned()
else {
return build_ai_public_error_response(
http::StatusCode::BAD_REQUEST,
ANTIGRAVITY_USER_SETTINGS_INVALID_DETAIL,
);
};
Json(json!({ "userSettings": user_settings })).into_response()
}
fn build_antigravity_load_code_assist_payload() -> Value {
json!({
"allowedTiers": ["free"],
"cloudaicompanionProject": "aether-antigravity-local",
"currentTier": "free",
"gcpManaged": false,
"paidTier": false,
"upgradeSubscriptionUri": ""
})
}
fn build_antigravity_fetch_user_info_payload() -> Value {
json!({
"regionCode": "US",
"userSettings": build_antigravity_default_user_settings_payload()
})
}
fn build_antigravity_default_user_settings_payload() -> Value {
json!({
"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-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-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"),
"models/proactive-observer": antigravity_model_payload("models/proactive-observer", "Proactive Observer")
},
"agentModelSorts": [
"gemini-3.5-flash-low",
"gemini-3-flash-agent",
"gemini-3.1-pro-low",
"gemini-3.1-flash-lite"
],
"audioTranscriptionModelIds": ["models/proactive-observer"],
"commandModelIds": ["gemini-3-flash"],
"commitMessageModelIds": ["gemini-3-flash"],
"defaultAgentModelId": "gemini-3.5-flash-low",
"deprecatedModelIds": [],
"experimentIds": [],
"imageGenerationModelIds": ["gemini-3.1-flash-image"],
"mqueryModelIds": ["gemini-3-flash"],
"tabModelIds": ["tab_flash_lite_preview", "tab_jump_flash_lite_preview"],
"tieredModelIds": {
"flash": "gemini-3-flash-agent",
"flashLite": "gemini-3.1-flash-lite",
"pro": "gemini-3.1-pro-low"
},
"webSearchModelIds": ["gemini-3-flash"]
})
}
fn antigravity_model_payload(id: &str, display_name: &str) -> Value {
json!({
"id": id,
"displayName": display_name
})
}
async fn maybe_build_local_gemini_video_operations_response(
state: &AppState,
request_context: &GatewayPublicRequestContext,
@@ -1950,6 +1950,127 @@ async fn gateway_executes_antigravity_gemini_cli_stream_via_local_decision_gate_
assert_eq!(stored_candidates.len(), 1);
assert_eq!(stored_candidates[0].status, RequestCandidateStatus::Success);
*seen_execution_runtime.lock().expect("mutex should lock") = None;
let inbound_response = reqwest::Client::new()
.post(format!(
"{gateway_url}/v1internal:streamGenerateContent?alt=sse"
))
.header(http::header::CONTENT_TYPE, "application/json")
.header("authorization", "Bearer google-antigravity-access-token")
.header("x-api-key", client_api_key)
.header("user-agent", "antigravity/cli/1.0.2 linux/arm64")
.header(
TRACE_ID_HEADER,
"trace-antigravity-v1internal-inbound-stream-456",
)
.json(&json!({
"project": "client-side-project-should-not-leak",
"requestId": "client-v1internal-request-456",
"model": "gemini-cli",
"userAgent": "antigravity",
"requestType": "checkpoint",
"request": {
"contents": [{
"role": "user",
"parts": [{"text": "checkpoint context"}]
}],
"generationConfig": {
"temperature": 0.4,
"thinkingConfig": {
"includeThoughts": true
}
},
"toolConfig": {
"functionCallingConfig": {
"mode": "NONE"
}
}
}
}))
.send()
.await
.expect("inbound antigravity request should succeed");
let inbound_status = inbound_response.status();
let inbound_miss_reason = inbound_response
.headers()
.get(crate::constants::LOCAL_EXECUTION_RUNTIME_MISS_REASON_HEADER)
.and_then(|value| value.to_str().ok())
.unwrap_or("-")
.to_string();
let inbound_response_body = inbound_response.text().await.expect("body should read");
assert_eq!(
inbound_status,
StatusCode::OK,
"unexpected inbound antigravity response body: {inbound_response_body}; miss_reason={inbound_miss_reason}"
);
let inbound_response_text = strip_sse_keepalive_comments(&inbound_response_body);
let inbound_payload = inbound_response_text
.trim()
.strip_prefix("data: ")
.expect("response should start with sse data prefix");
let inbound_response_json: serde_json::Value =
serde_json::from_str(inbound_payload).expect("stream payload should parse");
assert_eq!(
inbound_response_json["_v1internal_response_id"],
"resp_antigravity_cli_local_stream_123"
);
assert_eq!(
inbound_response_json["candidates"][0]["content"]["parts"][0]["text"],
"Hello Antigravity Stream"
);
let seen_inbound_execution_runtime_request = seen_execution_runtime
.lock()
.expect("mutex should lock")
.clone()
.expect("inbound execution runtime stream should be captured");
assert_eq!(
seen_inbound_execution_runtime_request.trace_id,
"trace-antigravity-v1internal-inbound-stream-456"
);
assert_eq!(
seen_inbound_execution_runtime_request.url,
"https://antigravity.googleapis.com/v1internal:streamGenerateContent?alt=sse"
);
assert_eq!(
seen_inbound_execution_runtime_request.authorization,
"Bearer refreshed-antigravity-cli-stream-access-token"
);
assert_eq!(
seen_inbound_execution_runtime_request.project,
"project-antigravity-stream-local-1"
);
assert_eq!(
seen_inbound_execution_runtime_request.request_id,
"client-v1internal-request-456"
);
assert_eq!(
seen_inbound_execution_runtime_request.model,
"claude-sonnet-4-5"
);
assert_eq!(
seen_inbound_execution_runtime_request.user_agent,
"antigravity"
);
assert_eq!(
seen_inbound_execution_runtime_request.request_type,
"checkpoint"
);
assert_eq!(seen_inbound_execution_runtime_request.contents_len, 1);
assert!((seen_inbound_execution_runtime_request.exact_temperature - 0.4).abs() < f64::EPSILON);
assert!(!seen_inbound_execution_runtime_request.request_has_model);
let inbound_stored_candidates = request_candidate_repository
.list_by_request_id("trace-antigravity-v1internal-inbound-stream-456")
.await
.expect("inbound request candidate trace should read");
assert_eq!(inbound_stored_candidates.len(), 1);
assert_eq!(
inbound_stored_candidates[0].status,
RequestCandidateStatus::Success
);
tokio::time::sleep(std::time::Duration::from_millis(100)).await;
assert!(
!*seen_report.lock().expect("mutex should lock"),
@@ -700,6 +700,138 @@ async fn gateway_rejects_invalid_claude_count_tokens_payload_without_hitting_fal
fallback_probe_handle.abort();
}
#[tokio::test]
async fn gateway_handles_antigravity_v1internal_control_plane_without_proxying() {
let fallback_probe_hits = Arc::new(Mutex::new(0usize));
let fallback_probe_hits_clone = Arc::clone(&fallback_probe_hits);
let fallback_probe = Router::new().route(
"/{*path}",
any(move |_request: Request| {
let fallback_probe_hits_inner = Arc::clone(&fallback_probe_hits_clone);
async move {
*fallback_probe_hits_inner.lock().expect("mutex should lock") += 1;
(StatusCode::OK, Json(json!({"proxied": true}))).into_response()
}
}),
);
let (_unused_fallback_probe_url, fallback_probe_handle) = start_server(fallback_probe).await;
let gateway = build_router_with_state(AppState::new().expect("gateway should build"));
let (gateway_url, gateway_handle) = start_server(gateway).await;
let client = reqwest::Client::new();
let user_settings = json!({
"preferredModelId": "gemini-3.5-flash-low",
"theme": "dark"
});
let requests = vec![
(
"/v1internal:loadCodeAssist",
json!({"metadata": {"ideType": "ANTIGRAVITY_CLI"}}),
),
(
"/v1internal:fetchAvailableModels",
json!({"project": "aether-antigravity-local"}),
),
(
"/v1internal:fetchUserInfo",
json!({"project": "aether-antigravity-local"}),
),
(
"/v1internal:fetchAdminControls",
json!({"project": "aether-antigravity-local"}),
),
("/v1internal:listExperiments", json!({})),
(
"/v1internal:recordCodeAssistMetrics",
json!({
"project": "aether-antigravity-local",
"requestId": "opaque-request-id",
"metrics": []
}),
),
(
"/v1internal:setUserSettings",
json!({"userSettings": user_settings.clone()}),
),
];
for (path, request_body) in requests {
let response = client
.post(format!("{gateway_url}{path}"))
.header("authorization", "Bearer ant-access-token")
.header("user-agent", "antigravity/cli/1.0.2 linux/arm64")
.json(&request_body)
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::OK, "path {path}");
assert_eq!(
response
.headers()
.get(EXECUTION_PATH_HEADER)
.and_then(|value| value.to_str().ok()),
Some(EXECUTION_PATH_LOCAL_AI_PUBLIC),
"path {path}"
);
let payload: serde_json::Value = response.json().await.expect("json body should parse");
match path {
"/v1internal:loadCodeAssist" => {
assert_eq!(
payload["cloudaicompanionProject"],
"aether-antigravity-local"
);
assert_eq!(payload["currentTier"], "free");
assert_eq!(payload["paidTier"], false);
assert_eq!(payload["gcpManaged"], false);
assert_eq!(payload["allowedTiers"], json!(["free"]));
assert_eq!(payload["upgradeSubscriptionUri"], "");
}
"/v1internal:fetchAvailableModels" => {
assert_eq!(payload["defaultAgentModelId"], "gemini-3.5-flash-low");
assert_eq!(payload["tieredModelIds"]["flash"], "gemini-3-flash-agent");
assert_eq!(
payload["models"]["gemini-3.5-flash-low"]["id"],
"gemini-3.5-flash-low"
);
assert_eq!(payload["commandModelIds"], json!(["gemini-3-flash"]));
assert_eq!(
payload["imageGenerationModelIds"],
json!(["gemini-3.1-flash-image"])
);
}
"/v1internal:fetchUserInfo" => {
assert_eq!(payload["regionCode"], "US");
assert_eq!(
payload["userSettings"]["preferredModelId"],
"gemini-3.5-flash-low"
);
}
"/v1internal:fetchAdminControls" => {
assert_eq!(payload, json!({}));
}
"/v1internal:listExperiments" => {
assert_eq!(payload["experimentIds"], json!([]));
assert_eq!(payload["flags"], json!({}));
}
"/v1internal:recordCodeAssistMetrics" => {
assert_eq!(payload, json!({}));
}
"/v1internal:setUserSettings" => {
assert_eq!(payload["userSettings"], user_settings);
}
other => panic!("unexpected path {other}"),
}
}
assert_eq!(*fallback_probe_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();
fallback_probe_handle.abort();
}
#[tokio::test]
async fn gateway_does_not_locally_reject_image_model_name_on_chat_completions() {
let fallback_probe_hits = Arc::new(Mutex::new(0usize));
@@ -168,6 +168,15 @@ async fn gateway_exposes_frontdoor_manifest_without_proxying_upstream() {
assert!(owned_routes
.iter()
.any(|value| value == "/v1beta/files/{path...}"));
assert!(owned_routes
.iter()
.any(|value| value == "/v1internal:loadCodeAssist"));
assert!(owned_routes
.iter()
.any(|value| value == "/v1internal:fetchAvailableModels"));
assert!(owned_routes
.iter()
.any(|value| value == "/v1internal:streamGenerateContent"));
assert_eq!(
payload["rust_frontdoor"]["internal_gateway"]["status"],
"rust_native_control_plane"
@@ -69,6 +69,14 @@ pub fn resolve_execution_runtime_stream_plan_kind(
));
}
if route_family == Some("antigravity")
&& route_kind == Some("stream_generate_content")
&& *method == Method::POST
&& path == "/v1internal:streamGenerateContent"
{
return Some(GEMINI_CLI_STREAM_PLAN_KIND);
}
if route_family == Some("openai")
&& is_openai_responses_route_kind(route_kind)
&& *method == Method::POST
@@ -679,6 +687,32 @@ mod tests {
);
}
#[test]
fn resolves_antigravity_v1internal_stream_plan_kind_as_gemini_cli_stream() {
assert_eq!(
resolve_execution_runtime_stream_plan_kind(
Some("ai_public"),
Some("antigravity"),
Some("stream_generate_content"),
Some("bearer_like"),
&Method::POST,
"/v1internal:streamGenerateContent",
),
Some(GEMINI_CLI_STREAM_PLAN_KIND)
);
assert_eq!(
resolve_execution_runtime_sync_plan_kind(
Some("ai_public"),
Some("antigravity"),
Some("stream_generate_content"),
Some("bearer_like"),
&Method::POST,
"/v1internal:streamGenerateContent",
),
None
);
}
#[test]
fn stream_path_detection_handles_gemini_method_paths_with_query() {
assert!(request_path_implies_stream_request(
+13 -1
View File
@@ -149,7 +149,8 @@ pub fn extract_ai_requested_model_from_request_path(
) -> Option<String> {
match family {
AiRequestedModelFamily::Standard => extract_ai_standard_requested_model(body_json),
AiRequestedModelFamily::Gemini => extract_ai_gemini_model_from_path(request_path),
AiRequestedModelFamily::Gemini => extract_ai_gemini_model_from_path(request_path)
.or_else(|| extract_ai_standard_requested_model(body_json)),
}
}
@@ -233,4 +234,15 @@ mod tests {
Some("gemini-2.5-pro")
);
}
#[test]
fn gemini_requested_model_parser_uses_body_model_when_path_has_no_model() {
let requested_model = extract_ai_requested_model_from_request_path(
"/v1internal:streamGenerateContent",
&serde_json::json!({ "model": " gemini-cli " }),
AiRequestedModelFamily::Gemini,
);
assert_eq!(requested_model.as_deref(), Some("gemini-cli"));
}
}
@@ -5,6 +5,7 @@ use super::auth::{AntigravityRequestAuth, ANTIGRAVITY_REQUEST_USER_AGENT};
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum AntigravityEnvelopeRequestType {
Agent,
Checkpoint,
EndpointTest,
}
@@ -12,6 +13,7 @@ impl AntigravityEnvelopeRequestType {
fn as_str(self) -> &'static str {
match self {
Self::Agent => "agent",
Self::Checkpoint => "checkpoint",
Self::EndpointTest => "endpoint_test",
}
}
@@ -29,7 +31,6 @@ pub enum AntigravityRequestEnvelopeUnsupportedReason {
MissingContents,
MissingRequestId,
MissingModel,
ComplexEnvelopeTransform,
}
pub fn classify_antigravity_safe_request_body(
@@ -38,12 +39,9 @@ pub fn classify_antigravity_safe_request_body(
let Value::Object(map) = request_body else {
return Err(AntigravityRequestEnvelopeUnsupportedReason::NonObjectBody);
};
if !map.contains_key("contents") {
if !map.contains_key("contents") && existing_v1internal_request_object(map).is_none() {
return Err(AntigravityRequestEnvelopeUnsupportedReason::MissingContents);
}
if contains_blocked_request_features(request_body) {
return Err(AntigravityRequestEnvelopeUnsupportedReason::ComplexEnvelopeTransform);
}
Ok(())
}
@@ -75,6 +73,27 @@ pub fn build_antigravity_safe_v1internal_request(
);
};
if let Some(existing_request) = existing_v1internal_request_object(source) {
let mut inner_request: Map<String, Value> = existing_request.clone();
inner_request.remove("model");
inner_request.remove("safetySettings");
inner_request.remove("safety_settings");
let request_id = non_empty_string_field(source, "requestId").unwrap_or(request_id);
let user_agent =
non_empty_string_field(source, "userAgent").unwrap_or(ANTIGRAVITY_REQUEST_USER_AGENT);
let request_type =
existing_v1internal_request_type(source).unwrap_or_else(|| request_type.as_str());
return AntigravityRequestEnvelopeSupport::Supported(serde_json::json!({
"project": auth.project_id,
"requestId": request_id,
"request": Value::Object(inner_request),
"model": model,
"userAgent": user_agent,
"requestType": request_type,
}));
}
let mut inner_request: Map<String, Value> = source.clone();
inner_request.remove("model");
inner_request.remove("safetySettings");
@@ -90,31 +109,258 @@ pub fn build_antigravity_safe_v1internal_request(
}))
}
fn contains_blocked_request_features(value: &Value) -> bool {
match value {
Value::Object(map) => map.iter().any(|(key, inner)| {
is_blocked_request_key(key.as_str()) || contains_blocked_request_features(inner)
}),
Value::Array(items) => items.iter().any(contains_blocked_request_features),
_ => false,
fn existing_v1internal_request_object(source: &Map<String, Value>) -> Option<&Map<String, Value>> {
source
.get("request")
.and_then(Value::as_object)
.filter(|request| request.contains_key("contents"))
}
fn non_empty_string_field<'a>(source: &'a Map<String, Value>, key: &str) -> Option<&'a str> {
source
.get(key)
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
}
fn existing_v1internal_request_type(source: &Map<String, Value>) -> Option<&str> {
match non_empty_string_field(source, "requestType")? {
"agent" => Some("agent"),
"checkpoint" => Some("checkpoint"),
"endpoint_test" => Some("endpoint_test"),
_ => None,
}
}
fn is_blocked_request_key(key: &str) -> bool {
matches!(
key.trim(),
"systemInstruction"
| "system_instruction"
| "tools"
| "toolConfig"
| "tool_config"
| "thinkingConfig"
| "thinking_config"
| "imageConfig"
| "image_config"
| "functionCall"
| "function_call"
| "functionResponse"
| "function_response"
)
#[cfg(test)]
mod tests {
use serde_json::json;
use super::{
build_antigravity_safe_v1internal_request, classify_antigravity_safe_request_body,
AntigravityEnvelopeRequestType, AntigravityRequestAuth, AntigravityRequestEnvelopeSupport,
};
fn sample_auth() -> AntigravityRequestAuth {
AntigravityRequestAuth {
project_id: "project-ant-123".to_string(),
client_version: None,
session_id: None,
}
}
#[test]
fn real_agent_request_preserves_antigravity_agent_fields() {
let request_body = json!({
"model": "client-side-model-should-not-be-nested",
"contents": [
{
"role": "user",
"parts": [
{ "text": "Reply with OK only." }
]
}
],
"systemInstruction": {
"role": "user",
"parts": [
{ "text": "Antigravity agent system prompt" }
]
},
"generationConfig": {
"maxOutputTokens": 8192,
"thinkingConfig": {
"includeThoughts": true,
"thinkingBudget": 4000
}
},
"toolConfig": {
"functionCallingConfig": {
"mode": "VALIDATED"
}
},
"tools": [
{
"functionDeclarations": [
{
"name": "run_command",
"description": "Run a command",
"parameters": {
"type": "object",
"properties": {
"cmd": { "type": "string" }
},
"required": ["cmd"]
}
}
]
}
],
"labels": {
"trajectory_id": "trajectory-123",
"used_claude": "false"
},
"sessionId": "session-ant-123",
"safetySettings": [
{ "category": "HARM_CATEGORY_UNSPECIFIED" }
]
});
assert_eq!(
classify_antigravity_safe_request_body(&request_body),
Ok(())
);
let envelope = match build_antigravity_safe_v1internal_request(
&sample_auth(),
"request-ant-agent-123",
"gemini-3.5-flash-low",
&request_body,
AntigravityEnvelopeRequestType::Agent,
) {
AntigravityRequestEnvelopeSupport::Supported(envelope) => envelope,
AntigravityRequestEnvelopeSupport::Unsupported(reason) => {
panic!("real agent envelope should be supported: {reason:?}")
}
};
assert_eq!(envelope["project"], "project-ant-123");
assert_eq!(envelope["requestId"], "request-ant-agent-123");
assert_eq!(envelope["model"], "gemini-3.5-flash-low");
assert_eq!(envelope["userAgent"], "antigravity");
assert_eq!(envelope["requestType"], "agent");
assert!(envelope["request"].get("model").is_none());
assert!(envelope["request"].get("safetySettings").is_none());
assert_eq!(
envelope["request"]["systemInstruction"]["parts"][0]["text"],
"Antigravity agent system prompt"
);
assert_eq!(
envelope["request"]["generationConfig"]["thinkingConfig"]["thinkingBudget"],
4000
);
assert_eq!(
envelope["request"]["toolConfig"]["functionCallingConfig"]["mode"],
"VALIDATED"
);
assert_eq!(
envelope["request"]["tools"][0]["functionDeclarations"][0]["name"],
"run_command"
);
assert_eq!(
envelope["request"]["labels"]["trajectory_id"],
"trajectory-123"
);
assert_eq!(envelope["request"]["sessionId"], "session-ant-123");
}
#[test]
fn checkpoint_request_type_builds_checkpoint_envelope() {
let request_body = json!({
"contents": [
{
"role": "user",
"parts": [
{ "text": "checkpoint context" }
]
}
],
"generationConfig": {
"maxOutputTokens": 8192,
"thinkingConfig": {
"includeThoughts": true,
"thinkingBudget": 4000
}
},
"toolConfig": {
"functionCallingConfig": {
"mode": "NONE"
}
}
});
let envelope = match build_antigravity_safe_v1internal_request(
&sample_auth(),
"request-ant-checkpoint-123",
"gemini-3.5-flash-low",
&request_body,
AntigravityEnvelopeRequestType::Checkpoint,
) {
AntigravityRequestEnvelopeSupport::Supported(envelope) => envelope,
AntigravityRequestEnvelopeSupport::Unsupported(reason) => {
panic!("checkpoint envelope should be supported: {reason:?}")
}
};
assert_eq!(envelope["requestType"], "checkpoint");
assert_eq!(
envelope["request"]["toolConfig"]["functionCallingConfig"]["mode"],
"NONE"
);
}
#[test]
fn existing_v1internal_envelope_is_not_double_wrapped() {
let request_body = json!({
"project": "client-side-project",
"requestId": "client-request-id-123",
"model": "gemini-3.5-flash-low",
"userAgent": "antigravity",
"requestType": "checkpoint",
"request": {
"contents": [
{
"role": "user",
"parts": [
{ "text": "checkpoint context" }
]
}
],
"generationConfig": {
"thinkingConfig": {
"includeThoughts": true
}
},
"toolConfig": {
"functionCallingConfig": {
"mode": "NONE"
}
}
}
});
assert_eq!(
classify_antigravity_safe_request_body(&request_body),
Ok(())
);
let envelope = match build_antigravity_safe_v1internal_request(
&sample_auth(),
"trace-request-id-should-not-overwrite-client-id",
"mapped-antigravity-model",
&request_body,
AntigravityEnvelopeRequestType::Agent,
) {
AntigravityRequestEnvelopeSupport::Supported(envelope) => envelope,
AntigravityRequestEnvelopeSupport::Unsupported(reason) => {
panic!("existing v1internal envelope should be supported: {reason:?}")
}
};
assert_eq!(envelope["project"], "project-ant-123");
assert_eq!(envelope["requestId"], "client-request-id-123");
assert_eq!(envelope["model"], "mapped-antigravity-model");
assert_eq!(envelope["userAgent"], "antigravity");
assert_eq!(envelope["requestType"], "checkpoint");
assert!(envelope["request"].get("request").is_none());
assert_eq!(
envelope["request"]["contents"][0]["parts"][0]["text"],
"checkpoint context"
);
assert_eq!(
envelope["request"]["toolConfig"]["functionCallingConfig"]["mode"],
"NONE"
);
}
}