mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-11 05:30:19 +08:00
Add Antigravity v1internal gateway adapter
This commit is contained in:
@@ -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] = &[
|
||||
|
||||
@@ -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";
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user