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

Fix native Antigravity stream envelope handling

Browse Source
This commit is contained in:
MMEXA
2026-05-25 09:39:55 +08:00
parent 28c3a5dbe4
commit b347f1816d
13 changed files with 770 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/constants.rs
+1
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"];
apps/aether-gateway/src/control/auth/resolution.rs
+141 -6
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,
apps/aether-gateway/src/handlers/public/ai_public.rs
+118 -20
View File
@@ -900,7 +900,7 @@ fn maybe_build_local_antigravity_v1internal_response(
"list_experiments" => Some(
Json(json!({
"experimentIds": [],
"flags": {}
"flags": []
}))
.into_response(),
),
@@ -943,12 +943,62 @@ fn build_antigravity_set_user_settings_response(request_body: Option<&Bytes>) ->
fn build_antigravity_load_code_assist_payload() -> Value {
json!({
"allowedTiers": ["free"],
"allowedTiers": [
antigravity_free_tier_payload(true),
antigravity_standard_tier_payload()
],
"cloudaicompanionProject": "aether-antigravity-local",
"currentTier": "free",
"currentTier": antigravity_free_tier_payload(false),
"gcpManaged": false,
"paidTier": false,
"upgradeSubscriptionUri": ""
"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"
})
}
@@ -961,7 +1011,7 @@ fn build_antigravity_fetch_user_info_payload() -> Value {
fn build_antigravity_default_user_settings_payload() -> Value {
json!({
"preferredModelId": "gemini-3.5-flash-low"
"preferredModelId": "gemini-3.1-flash-lite"
})
}
@@ -973,39 +1023,87 @@ fn build_antigravity_fetch_available_models_payload() -> Value {
"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": [
"gemini-3.5-flash-low",
"gemini-3-flash-agent",
"gemini-3.1-pro-low",
"gemini-3.1-flash-lite"
{
"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-flash"],
"defaultAgentModelId": "gemini-3.5-flash-low",
"deprecatedModelIds": [],
"commitMessageModelIds": ["gemini-3.1-flash-lite"],
"defaultAgentModelId": "gemini-3.1-flash-lite",
"deprecatedModelIds": {},
"experimentIds": [],
"imageGenerationModelIds": ["gemini-3.1-flash-image"],
"mqueryModelIds": ["gemini-3-flash"],
"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"
"flash": ["gemini-3-flash-agent"],
"flashLite": ["gemini-3.1-flash-lite"],
"pro": ["gemini-3.1-pro-low"]
},
"webSearchModelIds": ["gemini-3-flash"]
"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!({
"id": id,
"displayName": display_name
"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"
})
}
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
+288 -12
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;
@@ -2012,11 +2035,11 @@ async fn gateway_executes_antigravity_gemini_cli_stream_via_local_decision_gate_
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"],
inbound_response_json["responseId"],
"resp_antigravity_cli_local_stream_123"
);
assert_eq!(
inbound_response_json["candidates"][0]["content"]["parts"][0]["text"],
inbound_response_json["response"]["candidates"][0]["content"]["parts"][0]["text"],
"Hello Antigravity Stream"
);
@@ -2061,6 +2084,259 @@ async fn gateway_executes_antigravity_gemini_cli_stream_via_local_decision_gate_
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
apps/aether-gateway/src/tests/frontdoor/ai.rs
+47 -11
View File
@@ -721,7 +721,7 @@ async fn gateway_handles_antigravity_v1internal_control_plane_without_proxying()
let client = reqwest::Client::new();
let user_settings = json!({
"preferredModelId": "gemini-3.5-flash-low",
"preferredModelId": "gemini-3.1-flash-lite",
"theme": "dark"
});
let requests = vec![
@@ -783,30 +783,66 @@ async fn gateway_handles_antigravity_v1internal_control_plane_without_proxying()
payload["cloudaicompanionProject"],
"aether-antigravity-local"
);
assert_eq!(payload["currentTier"], "free");
assert_eq!(payload["paidTier"], false);
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"], json!(["free"]));
assert_eq!(payload["upgradeSubscriptionUri"], "");
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.5-flash-low");
assert_eq!(payload["tieredModelIds"]["flash"], "gemini-3-flash-agent");
assert_eq!(payload["defaultAgentModelId"], "gemini-3.1-flash-lite");
assert_eq!(
payload["models"]["gemini-3.5-flash-low"]["id"],
"gemini-3.5-flash-low"
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.5-flash-low"
"gemini-3.1-flash-lite"
);
}
"/v1internal:fetchAdminControls" => {
@@ -814,7 +850,7 @@ async fn gateway_handles_antigravity_v1internal_control_plane_without_proxying()
}
"/v1internal:listExperiments" => {
assert_eq!(payload["experimentIds"], json!([]));
assert_eq!(payload["flags"], json!({}));
assert_eq!(payload["flags"], json!([]));
}
"/v1internal:recordCodeAssistMetrics" => {
assert_eq!(payload, json!({}));
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