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

Merge branch 'pr-562' into review/pr-562

Browse Source
This commit is contained in:
fawney19
2026-05-25 23:10:29 +08:00
parent 4fb8955bc2 b347f1816d
commit c03aebba3f
23 changed files with 1715 additions and 57 deletions
Download Patch File Download Diff File Expand all files Collapse all files
apps/aether-gateway/src/ai_serving/planner/passthrough/provider/family/payload.rs
+7 -1
View File
@@ -11,7 +11,8 @@ use crate::ai_serving::planner::materialization_policy::{
build_local_candidate_persistence_policy, LocalCandidatePersistencePolicyKind,
};
use crate::ai_serving::planner::report_context::{
build_local_execution_report_context, LocalExecutionReportContextParts,
build_local_execution_report_context, insert_native_client_envelope_name,
LocalExecutionReportContextParts,
};
use crate::ai_serving::planner::spec_metadata::local_same_format_provider_spec_metadata;
use crate::ai_serving::planner::CandidateFailureDiagnostic;
@@ -90,6 +91,11 @@ pub(crate) async fn maybe_build_local_same_format_provider_decision_payload_for_
"envelope_name".to_string(),
json!(super::super::ANTIGRAVITY_ENVELOPE_NAME),
);
insert_native_client_envelope_name(
&mut extra_fields,
super::super::ANTIGRAVITY_ENVELOPE_NAME,
parts.uri.path(),
);
}
let provider_api_format = resolved.provider_api_format.clone();
let effective_headers = input.effective_headers(&parts.headers);
apps/aether-gateway/src/ai_serving/planner/report_context.rs
+15
View File
@@ -198,6 +198,21 @@ pub(crate) fn insert_provider_stream_event_api_format(
insert_ai_provider_stream_event_api_format(extra_fields, provider_type);
}
pub(crate) fn insert_native_client_envelope_name(
extra_fields: &mut Map<String, Value>,
envelope_name: &str,
request_path: &str,
) {
if envelope_name.eq_ignore_ascii_case("antigravity:v1internal")
&& request_path == "/v1internal:streamGenerateContent"
{
extra_fields.insert(
"client_envelope_name".to_string(),
Value::String(envelope_name.to_string()),
);
}
}
fn merge_incoming_tls_fingerprint(extra_fields: &mut Map<String, Value>, incoming_tls: Value) {
let entry = extra_fields
.entry("tls_fingerprint".to_string())
apps/aether-gateway/src/ai_serving/planner/standard/family/payload.rs
+3 -1
View File
@@ -9,7 +9,8 @@ use crate::ai_serving::planner::materialization_policy::{
};
use crate::ai_serving::planner::passthrough::maybe_build_local_same_format_provider_decision_payload_for_candidate;
use crate::ai_serving::planner::report_context::{
build_local_execution_report_context, LocalExecutionReportContextParts,
build_local_execution_report_context, insert_native_client_envelope_name,
LocalExecutionReportContextParts,
};
use crate::ai_serving::planner::spec_metadata::local_standard_spec_metadata;
use crate::ai_serving::planner::CandidateFailureDiagnostic;
@@ -92,6 +93,7 @@ pub(super) async fn maybe_build_local_standard_decision_payload_for_candidate(
"envelope_name".to_string(),
serde_json::Value::String(envelope_name.to_string()),
);
insert_native_client_envelope_name(&mut extra_fields, envelope_name, parts.uri.path());
}
let (execution_strategy, conversion_mode) = ai_local_execution_contract_for_formats(
spec_metadata.api_format,
apps/aether-gateway/src/ai_serving/planner/standard/openai/chat/decision/payload.rs
+3 -2
View File
@@ -2,8 +2,8 @@ use crate::ai_serving::build_request_trace_proxy_value;
use crate::ai_serving::planner::common::OPENAI_CHAT_STREAM_PLAN_KIND;
use crate::ai_serving::planner::decision_input::apply_provider_request_routing_policy_to_decision;
use crate::ai_serving::planner::report_context::{
build_local_execution_report_context, insert_provider_stream_event_api_format,
LocalExecutionReportContextParts,
build_local_execution_report_context, insert_native_client_envelope_name,
insert_provider_stream_event_api_format, LocalExecutionReportContextParts,
};
use crate::ai_serving::planner::{
build_ai_execution_decision_response, AiExecutionDecisionResponseParts,
@@ -84,6 +84,7 @@ pub(crate) async fn maybe_build_local_openai_chat_decision_payload_for_candidate
"envelope_name".to_string(),
serde_json::Value::String(envelope_name.to_string()),
);
insert_native_client_envelope_name(&mut extra_fields, envelope_name, parts.uri.path());
}
insert_provider_stream_event_api_format(
&mut extra_fields,
apps/aether-gateway/src/ai_serving/planner/standard/openai/responses/decision/payload.rs
+3 -2
View File
@@ -4,8 +4,8 @@ use tracing::debug;
use crate::ai_serving::build_request_trace_proxy_value;
use crate::ai_serving::planner::decision_input::apply_provider_request_routing_policy_to_decision;
use crate::ai_serving::planner::report_context::{
build_local_execution_report_context, insert_provider_stream_event_api_format,
LocalExecutionReportContextParts,
build_local_execution_report_context, insert_native_client_envelope_name,
insert_provider_stream_event_api_format, LocalExecutionReportContextParts,
};
use crate::ai_serving::planner::spec_metadata::local_openai_responses_spec_metadata;
use crate::ai_serving::planner::{
@@ -80,6 +80,7 @@ pub(crate) async fn maybe_build_local_openai_responses_decision_payload_for_cand
}
if let Some(envelope_name) = resolved.envelope_name {
extra_fields.insert("envelope_name".to_string(), json!(envelope_name));
insert_native_client_envelope_name(&mut extra_fields, envelope_name, parts.uri.path());
}
if let Some(image_request_summary) = resolved.image_request_summary.as_ref() {
extra_fields.insert("image_request".to_string(), image_request_summary.clone());
apps/aether-gateway/src/api/ai/registry.rs
+8
View File
@@ -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] = &[
apps/aether-gateway/src/constants.rs
+9
View File
@@ -65,6 +65,7 @@ pub(crate) const TRUSTED_ADMIN_MANAGEMENT_TOKEN_ID_HEADER: &str =
"x-aether-admin-management-token-id";
pub(crate) const TRUSTED_RATE_LIMIT_PREFLIGHT_HEADER: &str = "x-aether-rate-limit-preflight";
pub(crate) const DEFAULT_USER_GROUP_CONFIG_KEY: &str = "default_user_group_id";
pub(crate) const ANTIGRAVITY_BEARER_BRIDGE_CONFIG_KEY: &str = "module.antigravity.bearer_bridge";
pub(crate) const BUILTIN_DEFAULT_USER_GROUP_ID: &str = "00000000-0000-0000-0000-000000000001";
pub(crate) const FRONTDOOR_REPLACEABLE_ROUTE_GROUPS: &[&str] = &["frontdoor_compat_router"];
@@ -136,6 +137,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}",
];
apps/aether-gateway/src/control/auth/credentials.rs
+40
View File
@@ -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();
apps/aether-gateway/src/control/auth/resolution.rs
+208 -7
View File
@@ -16,16 +16,48 @@ use crate::{AppState, GatewayError};
use super::super::GatewayControlDecision;
use super::credentials::{
build_auth_context_cache_key, current_unix_secs, extract_request_credentials,
extract_trusted_admin_headers,
extract_trusted_admin_headers, hash_api_key,
};
use super::gate::GatewayLocalAuthRejection;
use super::principal::derive_principal_candidate;
use super::types::{GatewayPrincipalCandidate, GatewayTrustedAuthHeaders};
use super::types::{
GatewayCredentialCarrier, GatewayPrincipalCandidate, GatewayTrustedAuthHeaders,
};
use crate::headers::header_value_str;
const AUTH_CONTEXT_CACHE_TTL: Duration = Duration::from_secs(60);
const AUTH_CONTEXT_CACHE_MAX_ENTRIES: usize = 256;
#[derive(Debug, Clone, Deserialize)]
struct AntigravityBearerBridgeConfig {
#[serde(default)]
enabled: bool,
#[serde(default)]
auth_user_id: String,
#[serde(default)]
auth_api_key_id: String,
#[serde(default)]
bearer_sha256_allowlist: Vec<String>,
#[serde(default)]
allow_unverified_google_bearer: bool,
}
impl AntigravityBearerBridgeConfig {
fn bearer_validation_mode(&self, raw_bearer: &str) -> Option<&'static str> {
if !self.bearer_sha256_allowlist.is_empty() {
let bearer_hash = hash_api_key(raw_bearer);
return self
.bearer_sha256_allowlist
.iter()
.any(|allowed| allowed.trim().eq_ignore_ascii_case(&bearer_hash))
.then_some("sha256_allowlist");
}
self.allow_unverified_google_bearer
.then_some("explicit_unverified")
}
}
#[derive(Debug, Clone, Deserialize, Serialize)]
pub(crate) struct GatewayControlAuthContext {
pub(crate) user_id: String,
@@ -607,14 +639,117 @@ pub(super) async fn resolve_data_backed_auth_context(
.await,
))
}
Some(
GatewayPrincipalCandidate::DeferredBearerToken { .. }
| GatewayPrincipalCandidate::DeferredCookieHeader { .. },
) => Ok(None),
Some(GatewayPrincipalCandidate::DeferredBearerToken { raw, carrier }) => {
if let Some(auth_context) = resolve_antigravity_bearer_bridge_auth_context(
state,
signature,
raw.as_str(),
carrier,
now_unix_secs,
)
.await?
{
return Ok(Some(auth_context));
}
Ok(None)
}
Some(GatewayPrincipalCandidate::DeferredCookieHeader { .. }) => Ok(None),
None => Ok(None),
}
}
async fn resolve_antigravity_bearer_bridge_auth_context(
state: &AppState,
auth_endpoint_signature: &str,
raw_bearer: &str,
carrier: GatewayCredentialCarrier,
now_unix_secs: u64,
) -> Result<Option<GatewayControlAuthContext>, GatewayError> {
if carrier != GatewayCredentialCarrier::AuthorizationBearer
|| !auth_endpoint_signature
.trim()
.eq_ignore_ascii_case("antigravity:v1internal")
{
return Ok(None);
}
let Some(config_value) = state
.read_system_config_json_value(crate::constants::ANTIGRAVITY_BEARER_BRIDGE_CONFIG_KEY)
.await?
else {
return Ok(None);
};
if config_value.is_null() {
return Ok(None);
}
let config: AntigravityBearerBridgeConfig =
serde_json::from_value(config_value).map_err(|err| {
GatewayError::Internal(format!(
"{} invalid: {err}",
crate::constants::ANTIGRAVITY_BEARER_BRIDGE_CONFIG_KEY
))
})?;
if !config.enabled {
return Ok(None);
}
let Some(validation_mode) = config.bearer_validation_mode(raw_bearer) else {
return Ok(None);
};
let user_id = config.auth_user_id.trim();
let api_key_id = config.auth_api_key_id.trim();
if user_id.is_empty() || api_key_id.is_empty() {
return Err(GatewayError::Internal(format!(
"{} requires auth_user_id and auth_api_key_id",
crate::constants::ANTIGRAVITY_BEARER_BRIDGE_CONFIG_KEY
)));
}
let snapshot = state
.data
.read_auth_api_key_snapshot(user_id, api_key_id, now_unix_secs)
.await
.map_err(|err| GatewayError::Internal(err.to_string()))?;
let Some(snapshot) = snapshot else {
return Ok(Some(GatewayControlAuthContext {
user_id: user_id.to_string(),
api_key_id: api_key_id.to_string(),
username: None,
api_key_name: None,
balance_remaining: None,
access_allowed: false,
user_rate_limit: None,
api_key_rate_limit: None,
api_key_is_standalone: false,
admin_bypass_limits: false,
local_rejection: Some(GatewayLocalAuthRejection::InvalidApiKey),
allowed_models: None,
ip_rules: None,
}));
};
let wallet_access = resolve_wallet_auth_gate(state, &snapshot).await?;
let auth_context = build_data_backed_auth_context(
state,
snapshot,
auth_endpoint_signature,
None,
None,
wallet_access,
)
.await;
info!(
event_name = "antigravity_bearer_bridge_auth_context_resolved",
log_type = "event",
validation_mode,
user_id = auth_context.user_id.as_str(),
api_key_id = auth_context.api_key_id.as_str(),
access_allowed = auth_context.access_allowed,
has_local_rejection = auth_context.local_rejection.is_some(),
"resolved Antigravity bearer bridge auth context"
);
Ok(Some(auth_context))
}
async fn resolve_trusted_auth_context(
state: &AppState,
auth_endpoint_signature: &str,
@@ -712,7 +847,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 +887,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 +1410,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";
apps/aether-gateway/src/control/route/ai.rs
+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,
))
}
apps/aether-gateway/src/control/tests/ai.rs
+83
View File
@@ -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());
}
apps/aether-gateway/src/execution_runtime/transport.rs
+4
View File
@@ -1889,6 +1889,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]
apps/aether-gateway/src/frontdoor_loop_guard.rs
+1
View File
@@ -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)
}
apps/aether-gateway/src/handlers/public/ai_public.rs
+247 -1
View File
@@ -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,238 @@ 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": [
antigravity_free_tier_payload(true),
antigravity_standard_tier_payload()
],
"cloudaicompanionProject": "aether-antigravity-local",
"currentTier": antigravity_free_tier_payload(false),
"gcpManaged": false,
"paidTier": antigravity_paid_tier_payload(),
"upgradeSubscriptionUri": "https://codeassist.google.com/upgrade"
})
}
fn antigravity_free_tier_payload(include_default_marker: bool) -> Value {
if include_default_marker {
json!({
"id": "free-tier",
"name": "Antigravity",
"description": "Gemini-powered code suggestions and chat in multiple IDEs",
"privacyNotice": {
"showNotice": false
},
"isDefault": true
})
} else {
json!({
"id": "free-tier",
"name": "Antigravity",
"description": "Gemini-powered code suggestions and chat in multiple IDEs",
"privacyNotice": {
"showNotice": false
},
"upgradeSubscriptionUri": "https://codeassist.google.com/upgrade",
"upgradeSubscriptionText": "Upgrade for higher Antigravity request limits",
"upgradeSubscriptionType": "GDP_HELIUM"
})
}
}
fn antigravity_standard_tier_payload() -> Value {
json!({
"id": "standard-tier",
"name": "Antigravity",
"description": "Unlimited coding assistant with the most powerful Gemini models",
"userDefinedCloudaicompanionProject": true,
"privacyNotice": {},
"usesGcpTos": true
})
}
fn antigravity_paid_tier_payload() -> Value {
json!({
"id": "g1-pro-tier",
"name": "Google AI Pro",
"description": "Google AI Pro",
"upgradeSubscriptionUri": "https://antigravity.google/g1-upgrade",
"upgradeSubscriptionText": "Upgrade for the highest Antigravity request limits"
})
}
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.1-flash-lite"
})
}
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-2.5-flash": antigravity_model_payload("gemini-2.5-flash", "Gemini 2.5 Flash"),
"gemini-2.5-flash-lite": antigravity_model_payload("gemini-2.5-flash-lite", "Gemini 2.5 Flash Lite"),
"gemini-2.5-flash-thinking": antigravity_model_payload("gemini-2.5-flash-thinking", "Gemini 2.5 Flash Thinking"),
"gemini-2.5-pro": antigravity_model_payload("gemini-2.5-pro", "Gemini 2.5 Pro"),
"gemini-3.1-flash-image": antigravity_model_payload("gemini-3.1-flash-image", "Gemini 3.1 Flash Image"),
"tab_flash_lite_preview": antigravity_model_payload("tab_flash_lite_preview", "Tab Flash Lite Preview"),
"tab_jump_flash_lite_preview": antigravity_model_payload("tab_jump_flash_lite_preview", "Tab Jump Flash Lite Preview"),
"models/proactive-observer": antigravity_model_payload("models/proactive-observer", "Proactive Observer")
},
"agentModelSorts": [
{
"displayName": "Recommended",
"groups": [
{
"modelIds": [
"gemini-3.1-flash-lite",
"gemini-3-flash-agent",
"gemini-3.1-pro-low",
"gemini-3.5-flash-low"
]
}
]
}
],
"audioTranscriptionModelIds": ["models/proactive-observer"],
"commandModelIds": ["gemini-3-flash"],
"commitMessageModelIds": ["gemini-3.1-flash-lite"],
"defaultAgentModelId": "gemini-3.1-flash-lite",
"deprecatedModelIds": {},
"experimentIds": [],
"imageGenerationModelIds": ["gemini-3.1-flash-image"],
"mqueryModelIds": ["gemini-3.1-flash-lite"],
"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.1-flash-lite"]
})
}
fn antigravity_model_payload(id: &str, display_name: &str) -> Value {
let model = match id {
"gemini-2.5-flash" => "MODEL_GOOGLE_GEMINI_2_5_FLASH",
"gemini-2.5-flash-lite" => "MODEL_GOOGLE_GEMINI_2_5_FLASH_LITE",
"gemini-2.5-flash-thinking" => "MODEL_GOOGLE_GEMINI_2_5_FLASH_THINKING",
"gemini-2.5-pro" => "MODEL_GOOGLE_GEMINI_2_5_PRO",
"gemini-3-flash" => "MODEL_PLACEHOLDER_M18",
"gemini-3-flash-agent" => "MODEL_PLACEHOLDER_M132",
"gemini-3.1-flash-image" => "MODEL_PLACEHOLDER_M21",
"gemini-3.1-flash-lite" => "MODEL_PLACEHOLDER_M50",
"gemini-3.1-pro-low" => "MODEL_PLACEHOLDER_M36",
"gemini-3.5-flash-low" => "MODEL_PLACEHOLDER_M20",
"models/proactive-observer" => "MODEL_PLACEHOLDER_M70",
"tab_flash_lite_preview" => "MODEL_PLACEHOLDER_M19",
"tab_jump_flash_lite_preview" => "MODEL_PLACEHOLDER_M28",
_ => "MODEL_PLACEHOLDER_M20",
};
json!({
"apiProvider": "API_PROVIDER_GOOGLE_GEMINI",
"displayName": display_name,
"maxOutputTokens": 65536,
"maxTokens": 1048576,
"minThinkingBudget": 32,
"model": model,
"modelProvider": "MODEL_PROVIDER_GOOGLE",
"recommended": id == "gemini-3.1-flash-lite",
"supportedMimeTypes": {
"application/json": true,
"application/pdf": true,
"image/jpeg": true,
"image/png": true,
"text/markdown": true,
"text/plain": true
},
"supportsImages": true,
"supportsThinking": true,
"supportsVideo": true,
"thinkingBudget": 4000,
"tokenizerType": "LLAMA_WITH_SPECIAL"
})
}
async fn maybe_build_local_gemini_video_operations_response(
state: &AppState,
request_context: &GatewayPublicRequestContext,
apps/aether-gateway/src/state/core.rs
+4 -2
View File
@@ -69,8 +69,10 @@ const SCHEDULER_AFFECTING_SYSTEM_CONFIG_KEYS: &[&str] = &[
"provider_priority_mode",
"scheduling_mode",
];
const AUTH_AFFECTING_SYSTEM_CONFIG_KEYS: &[&str] =
&[crate::constants::DEFAULT_USER_GROUP_CONFIG_KEY];
const AUTH_AFFECTING_SYSTEM_CONFIG_KEYS: &[&str] = &[
crate::constants::DEFAULT_USER_GROUP_CONFIG_KEY,
crate::constants::ANTIGRAVITY_BEARER_BRIDGE_CONFIG_KEY,
];
const FRONTDOOR_RPM_AFFECTING_SYSTEM_CONFIG_KEYS: &[&str] = &["rate_limit_per_minute"];
fn system_config_key_affects_scheduler(key: &str) -> bool {
apps/aether-gateway/src/tests/ai_execute/stream_provider_gemini/local_cli.rs
+407 -10
View File
@@ -1441,7 +1441,7 @@ async fn gateway_executes_antigravity_gemini_cli_stream_via_local_decision_gate_
false,
Some(serde_json::json!(["gemini", "antigravity"])),
Some(serde_json::json!(["gemini:generate_content"])),
Some(serde_json::json!(["gemini-cli"])),
Some(serde_json::json!(["gemini-cli", "gemini-3.1-flash-lite"])),
api_key_id.to_string(),
Some("default".to_string()),
true,
@@ -1452,12 +1452,23 @@ async fn gateway_executes_antigravity_gemini_cli_stream_via_local_decision_gate_
Some(4_102_444_800),
Some(serde_json::json!(["gemini", "antigravity"])),
Some(serde_json::json!(["gemini:generate_content"])),
Some(serde_json::json!(["gemini-cli"])),
Some(serde_json::json!(["gemini-cli", "gemini-3.1-flash-lite"])),
)
.expect("auth snapshot should build")
}
fn sample_candidate_row() -> StoredMinimalCandidateSelectionRow {
sample_candidate_row_for("gemini-cli", "1")
}
fn sample_native_antigravity_candidate_row() -> StoredMinimalCandidateSelectionRow {
sample_candidate_row_for("gemini-3.1-flash-lite", "native-1")
}
fn sample_candidate_row_for(
global_model_name: &str,
row_suffix: &str,
) -> StoredMinimalCandidateSelectionRow {
StoredMinimalCandidateSelectionRow {
provider_id: "provider-antigravity-cli-oauth-stream-local-1".to_string(),
provider_name: "antigravity".to_string(),
@@ -1478,9 +1489,11 @@ async fn gateway_executes_antigravity_gemini_cli_stream_via_local_decision_gate_
key_capabilities: None,
key_internal_priority: 5,
key_global_priority_by_format: Some(serde_json::json!({"gemini:generate_content": 1})),
model_id: "model-antigravity-cli-oauth-stream-local-1".to_string(),
global_model_id: "global-model-antigravity-cli-oauth-stream-local-1".to_string(),
global_model_name: "gemini-cli".to_string(),
model_id: format!("model-antigravity-cli-oauth-stream-local-{row_suffix}"),
global_model_id: format!(
"global-model-antigravity-cli-oauth-stream-local-{row_suffix}"
),
global_model_name: global_model_name.to_string(),
global_model_mappings: None,
global_model_supports_streaming: Some(true),
model_provider_model_name: "claude-sonnet-4-5".to_string(),
@@ -1800,6 +1813,7 @@ async fn gateway_executes_antigravity_gemini_cli_stream_via_local_decision_gate_
let candidate_selection_repository =
Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(vec![
sample_candidate_row(),
sample_native_antigravity_candidate_row(),
]));
let request_candidate_repository = Arc::new(InMemoryRequestCandidateRepository::default());
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
@@ -1818,17 +1832,26 @@ async fn gateway_executes_antigravity_gemini_cli_stream_via_local_decision_gate_
.with_token_url_for_tests("antigravity", format!("{refresh_url}/oauth/token")),
),
]);
let gateway_state = build_state_with_execution_runtime_override(execution_runtime_url.clone())
.with_data_state_for_tests(
let data_state =
crate::data::GatewayDataState::with_auth_candidate_selection_provider_catalog_and_request_candidate_repository_for_tests(
auth_repository,
candidate_selection_repository,
provider_catalog_repository,
Arc::clone(&request_candidate_repository),
DEVELOPMENT_ENCRYPTION_KEY,
),
)
.with_oauth_refresh_coordinator_for_tests(oauth_refresh);
)
.with_system_config_values_for_tests([(
crate::constants::ANTIGRAVITY_BEARER_BRIDGE_CONFIG_KEY.to_string(),
json!({
"enabled": true,
"auth_user_id": "user-antigravity-cli-oauth-stream-local-1",
"auth_api_key_id": "api-key-antigravity-cli-oauth-stream-local-1",
"allow_unverified_google_bearer": true
}),
)]);
let gateway_state = build_state_with_execution_runtime_override(execution_runtime_url.clone())
.with_data_state_for_tests(data_state)
.with_oauth_refresh_coordinator_for_tests(oauth_refresh);
let gateway = build_router_with_state(gateway_state);
let (gateway_url, gateway_handle) = start_server(gateway).await;
@@ -1950,6 +1973,380 @@ 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["responseId"],
"resp_antigravity_cli_local_stream_123"
);
assert_eq!(
inbound_response_json["response"]["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);
*seen_execution_runtime.lock().expect("mutex should lock") = None;
let bearer_only_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("user-agent", "antigravity/cli/1.0.2 linux/arm64")
.header(
TRACE_ID_HEADER,
"trace-antigravity-v1internal-bearer-only-stream-789",
)
.json(&json!({
"project": "client-side-project-should-not-leak",
"requestId": "client-v1internal-request-789",
"model": "gemini-cli",
"userAgent": "antigravity",
"requestType": "agent",
"request": {
"contents": [{
"role": "user",
"parts": [{"text": "bearer-only request"}]
}],
"generationConfig": {
"temperature": 0.5,
"thinkingConfig": {
"includeThoughts": true
}
},
"toolConfig": {
"functionCallingConfig": {
"mode": "NONE"
}
}
}
}))
.send()
.await
.expect("bearer-only antigravity request should succeed");
let bearer_only_status = bearer_only_response.status();
let bearer_only_miss_reason = bearer_only_response
.headers()
.get(crate::constants::LOCAL_EXECUTION_RUNTIME_MISS_REASON_HEADER)
.and_then(|value| value.to_str().ok())
.unwrap_or("-")
.to_string();
let bearer_only_response_body = bearer_only_response.text().await.expect("body should read");
assert_eq!(
bearer_only_status,
StatusCode::OK,
"unexpected bearer-only antigravity response body: {bearer_only_response_body}; miss_reason={bearer_only_miss_reason}"
);
let bearer_only_response_text = strip_sse_keepalive_comments(&bearer_only_response_body);
let bearer_only_payload = bearer_only_response_text
.trim()
.strip_prefix("data: ")
.expect("response should start with sse data prefix");
let bearer_only_response_json: serde_json::Value =
serde_json::from_str(bearer_only_payload).expect("stream payload should parse");
assert_eq!(
bearer_only_response_json["responseId"],
"resp_antigravity_cli_local_stream_123"
);
assert_eq!(
bearer_only_response_json["response"]["candidates"][0]["content"]["parts"][0]["text"],
"Hello Antigravity Stream"
);
let seen_bearer_only_execution_runtime_request = seen_execution_runtime
.lock()
.expect("mutex should lock")
.clone()
.expect("bearer-only inbound execution runtime stream should be captured");
assert_eq!(
seen_bearer_only_execution_runtime_request.trace_id,
"trace-antigravity-v1internal-bearer-only-stream-789"
);
assert_eq!(
seen_bearer_only_execution_runtime_request.url,
"https://antigravity.googleapis.com/v1internal:streamGenerateContent?alt=sse"
);
assert_eq!(
seen_bearer_only_execution_runtime_request.authorization,
"Bearer refreshed-antigravity-cli-stream-access-token"
);
assert_eq!(
seen_bearer_only_execution_runtime_request.request_id,
"client-v1internal-request-789"
);
assert_eq!(
seen_bearer_only_execution_runtime_request.request_type,
"agent"
);
assert_eq!(seen_bearer_only_execution_runtime_request.contents_len, 1);
assert!(
(seen_bearer_only_execution_runtime_request.exact_temperature - 0.5).abs() < f64::EPSILON
);
assert!(!seen_bearer_only_execution_runtime_request.request_has_model);
*seen_execution_runtime.lock().expect("mutex should lock") = None;
let native_model_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("user-agent", "antigravity/cli/1.0.2 linux/arm64")
.header(
TRACE_ID_HEADER,
"trace-antigravity-v1internal-native-model-stream-790",
)
.json(&json!({
"project": "client-side-project-should-not-leak",
"requestId": "client-v1internal-request-790",
"model": "gemini-3.1-flash-lite",
"userAgent": "antigravity",
"requestType": "agent",
"request": {
"contents": [{
"role": "user",
"parts": [{"text": "native antigravity model request"}]
}],
"generationConfig": {
"temperature": 0.6,
"thinkingConfig": {
"includeThoughts": true
}
},
"toolConfig": {
"functionCallingConfig": {
"mode": "NONE"
}
}
}
}))
.send()
.await
.expect("native-model antigravity request should succeed");
let native_model_status = native_model_response.status();
let native_model_miss_reason = native_model_response
.headers()
.get(crate::constants::LOCAL_EXECUTION_RUNTIME_MISS_REASON_HEADER)
.and_then(|value| value.to_str().ok())
.unwrap_or("-")
.to_string();
let native_model_response_body = native_model_response
.text()
.await
.expect("body should read");
assert_eq!(
native_model_status,
StatusCode::OK,
"unexpected native-model antigravity response body: {native_model_response_body}; miss_reason={native_model_miss_reason}"
);
let native_model_response_text = strip_sse_keepalive_comments(&native_model_response_body);
let native_model_payload = native_model_response_text
.trim()
.strip_prefix("data: ")
.expect("response should start with sse data prefix");
let native_model_response_json: serde_json::Value =
serde_json::from_str(native_model_payload).expect("stream payload should parse");
assert_eq!(
native_model_response_json["responseId"],
"resp_antigravity_cli_local_stream_123"
);
assert_eq!(
native_model_response_json["response"]["candidates"][0]["content"]["parts"][0]["text"],
"Hello Antigravity Stream"
);
let seen_native_model_execution_runtime_request = seen_execution_runtime
.lock()
.expect("mutex should lock")
.clone()
.expect("native-model inbound execution runtime stream should be captured");
assert_eq!(
seen_native_model_execution_runtime_request.trace_id,
"trace-antigravity-v1internal-native-model-stream-790"
);
assert_eq!(
seen_native_model_execution_runtime_request.url,
"https://antigravity.googleapis.com/v1internal:streamGenerateContent?alt=sse"
);
assert_eq!(
seen_native_model_execution_runtime_request.authorization,
"Bearer refreshed-antigravity-cli-stream-access-token"
);
assert_eq!(
seen_native_model_execution_runtime_request.model,
"claude-sonnet-4-5"
);
assert_eq!(
seen_native_model_execution_runtime_request.request_id,
"client-v1internal-request-790"
);
assert!(
(seen_native_model_execution_runtime_request.exact_temperature - 0.6).abs() < f64::EPSILON
);
assert!(!seen_native_model_execution_runtime_request.request_has_model);
if std::env::var("AETHER_REAL_AGY_CLI_SMOKE").ok().as_deref() == Some("1") {
*seen_execution_runtime.lock().expect("mutex should lock") = None;
let log_path = std::env::var("AETHER_REAL_AGY_CLI_LOG")
.unwrap_or_else(|_| "/tmp/aether-real-agy-cli-smoke.log".to_string());
let workdir = std::env::var("AETHER_REAL_AGY_CLI_WORKDIR")
.unwrap_or_else(|_| "/tmp/aether-real-agy-cli-work".to_string());
std::fs::create_dir_all(&workdir).expect("agy smoke workdir should create");
let gateway_url_for_agy = gateway_url.clone();
let log_path_for_agy = log_path.clone();
let workdir_for_agy = workdir.clone();
let output = tokio::task::spawn_blocking(move || {
std::process::Command::new("agy")
.arg("--log-file")
.arg(&log_path_for_agy)
.arg("-p")
.arg("Reply with AETHER_CLOSED_LOOP_OK only.")
.arg("--print-timeout")
.arg("45s")
.env("AGY_CLI_DISABLE_AUTO_UPDATE", "true")
.env("CLOUD_CODE_URL", &gateway_url_for_agy)
.current_dir(&workdir_for_agy)
.output()
})
.await
.expect("agy smoke blocking task should join")
.expect("agy smoke process should spawn");
let stdout = String::from_utf8_lossy(&output.stdout);
let stderr = String::from_utf8_lossy(&output.stderr);
let agy_log = std::fs::read_to_string(&log_path).unwrap_or_default();
let seen_agy_execution_runtime_snapshot = seen_execution_runtime
.lock()
.expect("mutex should lock")
.clone();
assert!(
output.status.success(),
"agy smoke failed: status={:?}\nseen_execution_runtime={seen_agy_execution_runtime_snapshot:?}\nstdout={stdout}\nstderr={stderr}\nlog={agy_log}",
output.status
);
assert!(
stdout.contains("Hello Antigravity Stream")
|| stdout.contains("AETHER_CLOSED_LOOP_OK"),
"agy smoke stdout did not contain the local runtime response: stdout={stdout}\nstderr={stderr}\nlog={agy_log}"
);
let seen_agy_execution_runtime_request = seen_agy_execution_runtime_snapshot
.expect("real agy smoke should reach execution runtime");
assert_eq!(
seen_agy_execution_runtime_request.url,
"https://antigravity.googleapis.com/v1internal:streamGenerateContent?alt=sse"
);
}
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"),
apps/aether-gateway/src/tests/frontdoor/ai.rs
+168
View File
@@ -700,6 +700,174 @@ 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.1-flash-lite",
"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"]["id"], "free-tier");
assert_eq!(payload["currentTier"]["name"], "Antigravity");
assert_eq!(payload["paidTier"]["id"], "g1-pro-tier");
assert_eq!(payload["gcpManaged"], false);
assert_eq!(payload["allowedTiers"][0]["id"], "free-tier");
assert_eq!(payload["allowedTiers"][0]["isDefault"], true);
assert_eq!(payload["allowedTiers"][1]["id"], "standard-tier");
assert_eq!(
payload["upgradeSubscriptionUri"],
"https://codeassist.google.com/upgrade"
);
}
"/v1internal:fetchAvailableModels" => {
assert_eq!(payload["defaultAgentModelId"], "gemini-3.1-flash-lite");
assert_eq!(
payload["tieredModelIds"]["flash"],
json!(["gemini-3-flash-agent"])
);
assert_eq!(
payload["models"]["gemini-3.5-flash-low"]["displayName"],
"Gemini 3.5 Flash Low"
);
assert_eq!(
payload["models"]["gemini-3.5-flash-low"]["apiProvider"],
"API_PROVIDER_GOOGLE_GEMINI"
);
assert_eq!(
payload["models"]["gemini-2.5-flash-lite"]["model"],
"MODEL_GOOGLE_GEMINI_2_5_FLASH_LITE"
);
assert_eq!(
payload["agentModelSorts"][0]["groups"][0]["modelIds"],
json!([
"gemini-3.1-flash-lite",
"gemini-3-flash-agent",
"gemini-3.1-pro-low",
"gemini-3.5-flash-low"
])
);
assert_eq!(payload["deprecatedModelIds"], json!({}));
assert_eq!(payload["commandModelIds"], json!(["gemini-3-flash"]));
assert_eq!(
payload["imageGenerationModelIds"],
json!(["gemini-3.1-flash-image"])
);
assert_eq!(payload["mqueryModelIds"], json!(["gemini-3.1-flash-lite"]));
assert_eq!(
payload["webSearchModelIds"],
json!(["gemini-3.1-flash-lite"])
);
assert_eq!(
payload["commitMessageModelIds"],
json!(["gemini-3.1-flash-lite"])
);
}
"/v1internal:fetchUserInfo" => {
assert_eq!(payload["regionCode"], "US");
assert_eq!(
payload["userSettings"]["preferredModelId"],
"gemini-3.1-flash-lite"
);
}
"/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));
apps/aether-gateway/src/tests/frontdoor/ops.rs
+9
View File
@@ -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"
crates/aether-ai-formats/src/formats/shared/routing.rs
+34
View File
@@ -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(
crates/aether-ai-formats/src/formats/shared/stream_rewrite.rs
+71
View File
@@ -47,6 +47,18 @@ pub fn resolve_finalize_stream_rewrite_mode(
.trim()
.to_ascii_lowercase();
if !needs_conversion
&& client_consumes_same_private_stream_envelope(
report_context,
envelope_name.as_str(),
provider_api_format.as_str(),
client_api_format.as_str(),
)
{
return model_directive_display_model_from_report_context(report_context)
.map(|_| FinalizeStreamRewriteMode::ModelDirectiveDisplay);
}
if needs_conversion
&& envelope_name.eq_ignore_ascii_case(KIRO_ENVELOPE_NAME)
&& provider_api_format == "claude:messages"
@@ -107,6 +119,26 @@ pub fn resolve_finalize_stream_rewrite_mode(
.then_some(FinalizeStreamRewriteMode::EnvelopeUnwrap)
}
fn client_consumes_same_private_stream_envelope(
report_context: &Value,
envelope_name: &str,
provider_api_format: &str,
client_api_format: &str,
) -> bool {
if envelope_name.is_empty()
|| provider_api_format != client_api_format
|| !provider_adaptation_should_unwrap_stream_envelope(envelope_name, provider_api_format)
{
return false;
}
report_context
.get("client_envelope_name")
.and_then(Value::as_str)
.is_some_and(|client_envelope_name| {
client_envelope_name.eq_ignore_ascii_case(envelope_name)
})
}
enum AiSurfaceStreamRewriteState {
EnvelopeUnwrap,
ModelDirectiveDisplay,
@@ -451,6 +483,45 @@ mod tests {
);
}
#[test]
fn resolves_no_rewriter_when_client_consumes_same_private_envelope() {
let report_context = json!({
"provider_api_format": "gemini:generate_content",
"client_api_format": "gemini:generate_content",
"envelope_name": "antigravity:v1internal",
"client_envelope_name": "antigravity:v1internal",
"needs_conversion": false,
});
assert_eq!(resolve_finalize_stream_rewrite_mode(&report_context), None);
assert!(maybe_build_ai_surface_stream_rewriter(Some(&report_context)).is_none());
}
#[test]
fn native_private_envelope_client_keeps_response_wrapper_for_model_display_rewrite() {
let report_context = json!({
"provider_api_format": "gemini:generate_content",
"client_api_format": "gemini:generate_content",
"envelope_name": "antigravity:v1internal",
"client_envelope_name": "antigravity:v1internal",
"model": "gemini-2.5-pro-high",
"mapped_model": "gemini-2.5-pro",
"needs_conversion": false,
});
let mut rewriter = maybe_build_ai_surface_stream_rewriter(Some(&report_context))
.expect("display-model rewriter should exist");
let output = rewriter
.push_chunk(
b"data: {\"response\":{\"modelVersion\":\"gemini-2.5-pro\",\"candidates\":[]},\"responseId\":\"resp_native_123\"}\n\n",
)
.expect("rewrite should succeed");
let output = String::from_utf8(output).expect("output should be utf8");
assert!(output.contains("\"response\":"));
assert!(output.contains("\"responseId\":\"resp_native_123\""));
assert!(output.contains("\"modelVersion\":\"gemini-2.5-pro-high\""));
assert!(!output.contains("_v1internal_response_id"));
}
#[test]
fn resolves_kiro_same_format_streams_to_kiro_mode() {
let report_context = json!({
crates/aether-ai-formats/src/provider_compat/private_envelope.rs
+69
View File
@@ -47,6 +47,13 @@ pub fn normalize_provider_private_report_context(report_context: Option<&Value>)
.get("provider_api_format")
.and_then(Value::as_str)
.unwrap_or_default();
if report_context_preserves_private_client_envelope(
report_context,
envelope_name,
provider_api_format,
) {
return Some(report_context.clone());
}
if provider_adaptation_descriptor_for_envelope(envelope_name, provider_api_format).is_none() {
return Some(report_context.clone());
}
@@ -64,6 +71,22 @@ pub fn normalize_provider_private_response_value(
{
return Some(data);
}
let envelope_name = report_context
.get("envelope_name")
.and_then(Value::as_str)
.unwrap_or_default();
let provider_api_format = report_context
.get("provider_api_format")
.and_then(Value::as_str)
.unwrap_or_default();
if report_context_preserves_private_client_envelope(
report_context,
envelope_name,
provider_api_format,
) {
return Some(data);
}
let mut unwrapped = match report_context.get("envelope_name").and_then(Value::as_str) {
Some(KIRO_ENVELOPE_NAME) => data,
Some(GEMINI_CLI_V1INTERNAL_ENVELOPE_NAME) => {
@@ -155,6 +178,13 @@ fn transform_provider_private_stream_line_with_event_state(
if !provider_adaptation_should_unwrap_stream_envelope(envelope_name, provider_api_format) {
return Ok(line);
}
if report_context_preserves_private_client_envelope(
report_context,
envelope_name,
provider_api_format,
) {
return Ok(line);
}
if envelope_name == WINDSURF_ENVELOPE_NAME && looks_like_windsurf_error(&body) {
return Ok(line);
}
@@ -305,6 +335,13 @@ pub fn maybe_build_provider_private_stream_normalizer<'a>(
.unwrap_or_default();
let descriptor =
provider_adaptation_descriptor_for_envelope(envelope_name, provider_api_format)?;
if report_context_preserves_private_client_envelope(
report_context,
envelope_name,
provider_api_format,
) {
return None;
}
let mode = if descriptor
.envelope_name
.eq_ignore_ascii_case(KIRO_ENVELOPE_NAME)
@@ -679,6 +716,24 @@ fn clear_private_envelope_context(report_context: &Value) -> Value {
normalized
}
fn report_context_preserves_private_client_envelope(
report_context: &Value,
envelope_name: &str,
provider_api_format: &str,
) -> bool {
if envelope_name.is_empty()
|| provider_adaptation_descriptor_for_envelope(envelope_name, provider_api_format).is_none()
{
return false;
}
report_context
.get("client_envelope_name")
.and_then(Value::as_str)
.is_some_and(|client_envelope_name| {
client_envelope_name.eq_ignore_ascii_case(envelope_name)
})
}
fn local_finalize_response_model(report_context: &Value) -> &str {
report_context
.get("mapped_model")
@@ -997,6 +1052,20 @@ mod tests {
assert!(output_text.contains("\"id\":\"call_get_weather_0\""));
}
#[test]
fn private_stream_normalizer_preserves_antigravity_native_client_envelope() {
let report_context = json!({
"has_envelope": true,
"provider_api_format": "gemini:generate_content",
"client_api_format": "gemini:generate_content",
"envelope_name": "antigravity:v1internal",
"client_envelope_name": "antigravity:v1internal",
"mapped_model": "claude-sonnet-4-5",
});
assert!(maybe_build_provider_private_stream_normalizer(Some(&report_context)).is_none());
}
#[test]
fn detects_sse_error_events_without_explicit_type_field() {
let body = br#"event: error
crates/aether-ai-serving/src/surface_spec.rs
+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"));
}
}
crates/aether-provider-transport/src/antigravity/request.rs
+275 -29
View File
@@ -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"
);
}
}