Merge branch 'fawney19:main' into main

This commit is contained in:
Start
2026-07-09 15:11:44 +08:00
committed by GitHub
67 changed files with 3107 additions and 461 deletions
@@ -0,0 +1,169 @@
use std::collections::BTreeMap;
use std::sync::Arc;
use serde_json::Value;
use crate::ai_serving::transport::antigravity::{
build_antigravity_safe_v1internal_request, build_antigravity_static_identity_headers,
classify_local_antigravity_request_support, AntigravityEnvelopeRequestType,
AntigravityRequestAuth, AntigravityRequestAuthUnsupportedReason,
AntigravityRequestEnvelopeSupport, AntigravityRequestSideSupport,
AntigravityRequestSideUnsupportedReason,
};
use crate::ai_serving::transport::{
build_standard_provider_request_headers, GatewayProviderTransportSnapshot,
StandardProviderRequestHeaders, StandardProviderRequestHeadersInput,
};
use crate::AppState;
pub(crate) const ANTIGRAVITY_V1INTERNAL_ENVELOPE_NAME: &str = "antigravity:v1internal";
pub(crate) enum AntigravityV1InternalRequestError {
TransportUnsupported,
EnvelopeUnsupported,
UpstreamUrlUnavailable,
HeaderRulesApplyFailed,
}
pub(crate) struct AntigravityV1InternalRequestInput<'a> {
pub(crate) state: &'a AppState,
pub(crate) parts: &'a http::request::Parts,
pub(crate) transport: &'a Arc<GatewayProviderTransportSnapshot>,
pub(crate) trace_id: &'a str,
pub(crate) mapped_model: &'a str,
pub(crate) provider_api_format: &'a str,
pub(crate) auth_header: &'a str,
pub(crate) auth_value: &'a str,
pub(crate) request_headers: &'a http::HeaderMap,
pub(crate) original_request_body: &'a Value,
pub(crate) gemini_request_body: &'a Value,
pub(crate) upstream_is_stream: bool,
pub(crate) same_format: bool,
}
pub(crate) struct AntigravityV1InternalRequest {
pub(crate) transport: Arc<GatewayProviderTransportSnapshot>,
pub(crate) body: Value,
pub(crate) headers: StandardProviderRequestHeaders,
pub(crate) upstream_url: String,
}
pub(crate) async fn build_antigravity_v1internal_provider_request(
input: AntigravityV1InternalRequestInput<'_>,
) -> Result<AntigravityV1InternalRequest, AntigravityV1InternalRequestError> {
let payload = build_antigravity_v1internal_payload(
input.state,
input.transport,
input.trace_id,
input.mapped_model,
input.gemini_request_body,
)
.await?;
let upstream_url = crate::ai_serving::build_provider_transport_request_url_for_request_body(
&payload.transport,
input.provider_api_format,
Some(input.mapped_model),
input.upstream_is_stream,
input.parts.uri.query(),
None,
Some(&payload.body),
)
.ok_or(AntigravityV1InternalRequestError::UpstreamUrlUnavailable)?;
let extra_headers: BTreeMap<String, String> =
build_antigravity_static_identity_headers(&payload.auth);
let mut headers =
build_standard_provider_request_headers(StandardProviderRequestHeadersInput {
transport: &payload.transport,
provider_api_format: input.provider_api_format,
same_format: input.same_format,
headers: input.request_headers,
auth_header: input.auth_header,
auth_value: input.auth_value,
extra_headers: &extra_headers,
header_rules: payload.transport.endpoint.header_rules.as_ref(),
provider_request_body: &payload.body,
original_request_body: input.original_request_body,
upstream_is_stream: input.upstream_is_stream,
})
.ok_or(AntigravityV1InternalRequestError::HeaderRulesApplyFailed)?;
headers
.headers
.insert("accept".to_string(), "text/event-stream".to_string());
Ok(AntigravityV1InternalRequest {
transport: payload.transport,
body: payload.body,
headers,
upstream_url,
})
}
struct AntigravityV1InternalPayload {
transport: Arc<GatewayProviderTransportSnapshot>,
auth: AntigravityRequestAuth,
body: Value,
}
async fn build_antigravity_v1internal_payload(
state: &AppState,
transport: &Arc<GatewayProviderTransportSnapshot>,
trace_id: &str,
mapped_model: &str,
gemini_request_body: &Value,
) -> Result<AntigravityV1InternalPayload, AntigravityV1InternalRequestError> {
let mut resolved_transport = Arc::clone(transport);
let mut antigravity_support = classify_local_antigravity_request_support(
&resolved_transport,
gemini_request_body,
AntigravityEnvelopeRequestType::Agent,
);
if matches!(
antigravity_support,
AntigravityRequestSideSupport::Unsupported(
AntigravityRequestSideUnsupportedReason::UnsupportedAuth(
AntigravityRequestAuthUnsupportedReason::MissingProjectId
)
)
) {
if let Some(hydrated) = state
.hydrate_antigravity_project_metadata_for_transport(&resolved_transport)
.await
{
resolved_transport = Arc::new(hydrated);
antigravity_support = classify_local_antigravity_request_support(
&resolved_transport,
gemini_request_body,
AntigravityEnvelopeRequestType::Agent,
);
}
}
let auth = match antigravity_support {
AntigravityRequestSideSupport::Supported(spec) => spec.auth,
AntigravityRequestSideSupport::Unsupported(_) => {
return Err(AntigravityV1InternalRequestError::TransportUnsupported);
}
};
let body = match build_antigravity_safe_v1internal_request(
&auth,
trace_id,
mapped_model,
gemini_request_body,
AntigravityEnvelopeRequestType::Agent,
) {
AntigravityRequestEnvelopeSupport::Supported(envelope) => envelope,
AntigravityRequestEnvelopeSupport::Unsupported(_) => {
return Err(AntigravityV1InternalRequestError::EnvelopeUnsupported);
}
};
Ok(AntigravityV1InternalPayload {
transport: resolved_transport,
auth,
body,
})
}
@@ -1,6 +1,7 @@
use crate::ai_serving::{AiExecutionDecision, AiExecutionPlanPayload, GatewayControlDecision};
use crate::{AppState, GatewayError};
mod antigravity;
mod candidate_affinity_cache;
mod candidate_materialization;
mod candidate_metadata;
@@ -4,6 +4,10 @@ use std::sync::Arc;
use aether_contracts::ResolvedTransportProfile;
use serde_json::{json, Value};
use crate::ai_serving::planner::antigravity::{
build_antigravity_v1internal_provider_request, AntigravityV1InternalRequestError,
AntigravityV1InternalRequestInput, ANTIGRAVITY_V1INTERNAL_ENVELOPE_NAME,
};
use crate::ai_serving::planner::candidate_preparation::{
prepare_header_authenticated_candidate, prepare_header_authenticated_candidate_from_auth,
OauthPreparationContext,
@@ -27,6 +31,7 @@ use crate::ai_serving::planner::standard::{
build_local_openai_chat_upstream_url, request_body_build_failure_extra_data,
request_conversion_failure_extra_data,
};
use crate::ai_serving::transport::antigravity::is_antigravity_provider_transport;
use crate::ai_serving::transport::auth::resolve_local_openai_bearer_auth;
use crate::ai_serving::transport::kiro::{
build_kiro_provider_headers, build_kiro_provider_request_body,
@@ -183,6 +188,7 @@ pub(crate) async fn resolve_local_openai_chat_candidate_payload_parts(
.provider_type
.trim()
.eq_ignore_ascii_case("grok");
let is_antigravity = is_antigravity_provider_transport(transport);
let is_gemini_cli = is_gemini_cli_provider_transport(transport);
if is_grok && is_grok_text_provider_api_format(provider_api_format) {
@@ -549,7 +555,9 @@ pub(crate) async fn resolve_local_openai_chat_candidate_payload_parts(
transport,
conversion_kind,
) {
if !(is_gemini_cli && normalized_provider_api_format == "gemini:generate_content") {
if !((is_antigravity || is_gemini_cli)
&& normalized_provider_api_format == "gemini:generate_content")
{
mark_skipped_local_openai_chat_candidate(
state,
input,
@@ -757,6 +765,28 @@ pub(crate) async fn resolve_local_openai_chat_candidate_payload_parts(
)
.await);
}
if provider_api_format == "gemini:generate_content" && is_antigravity {
return Ok(build_antigravity_openai_chat_cross_format_payload_parts(
state,
parts,
trace_id,
body_json,
input,
eligible,
candidate_index,
candidate_id,
decision_kind,
transport,
provider_api_format.as_str(),
prepared_candidate.mapped_model,
prepared_candidate.auth_header,
prepared_candidate.auth_value,
provider_request_body,
upstream_is_stream,
redaction.redacted,
)
.await);
}
let Some(upstream_url) = build_cross_format_openai_chat_upstream_url(
parts,
@@ -857,6 +887,158 @@ pub(crate) async fn resolve_local_openai_chat_candidate_payload_parts(
}))
}
#[allow(clippy::too_many_arguments)]
async fn build_antigravity_openai_chat_cross_format_payload_parts(
state: &AppState,
parts: &http::request::Parts,
trace_id: &str,
original_body_json: &serde_json::Value,
input: &LocalOpenAiChatDecisionInput,
eligible: &EligibleLocalExecutionCandidate,
candidate_index: u32,
candidate_id: &str,
decision_kind: &str,
transport: &Arc<GatewayProviderTransportSnapshot>,
provider_api_format: &str,
mapped_model: String,
auth_header: String,
auth_value: String,
gemini_request_body: Value,
upstream_is_stream: bool,
request_redacted: bool,
) -> Option<LocalOpenAiChatCandidatePayloadParts> {
let candidate = &eligible.candidate;
let effective_headers = input.effective_headers(&parts.headers);
let resolved =
match build_antigravity_v1internal_provider_request(AntigravityV1InternalRequestInput {
state,
parts,
transport,
trace_id,
mapped_model: &mapped_model,
provider_api_format,
auth_header: &auth_header,
auth_value: &auth_value,
request_headers: effective_headers,
original_request_body: original_body_json,
gemini_request_body: &gemini_request_body,
upstream_is_stream,
same_format: false,
})
.await
{
Ok(resolved) => resolved,
Err(AntigravityV1InternalRequestError::TransportUnsupported) => {
mark_skipped_local_openai_chat_candidate(
state,
input,
trace_id,
candidate,
candidate_index,
candidate_id,
"transport_unsupported",
)
.await;
return None;
}
Err(AntigravityV1InternalRequestError::EnvelopeUnsupported) => {
mark_skipped_local_openai_chat_candidate_with_failure_diagnostic(
state,
input,
trace_id,
candidate,
candidate_index,
candidate_id,
"provider_request_body_build_failed",
CandidateFailureDiagnostic::envelope_build_failed(
"openai:chat",
provider_api_format,
"openai_chat_antigravity_envelope",
),
)
.await;
return None;
}
Err(AntigravityV1InternalRequestError::UpstreamUrlUnavailable) => {
mark_skipped_local_openai_chat_candidate_with_failure_diagnostic(
state,
input,
trace_id,
candidate,
candidate_index,
candidate_id,
"upstream_url_missing",
CandidateFailureDiagnostic::upstream_url_missing(
"openai:chat",
provider_api_format,
"openai_chat_antigravity_url",
),
)
.await;
return None;
}
Err(AntigravityV1InternalRequestError::HeaderRulesApplyFailed) => {
mark_skipped_local_openai_chat_candidate_with_failure_diagnostic(
state,
input,
trace_id,
candidate,
candidate_index,
candidate_id,
"transport_header_rules_apply_failed",
CandidateFailureDiagnostic::header_rules_apply_failed(
"openai:chat",
provider_api_format,
"openai_chat_antigravity_headers",
),
)
.await;
return None;
}
};
let mut provider_request_headers = resolved.headers.headers;
apply_codex_openai_responses_special_headers(
&mut provider_request_headers,
&resolved.body,
effective_headers,
resolved.transport.provider.provider_type.as_str(),
provider_api_format,
Some(trace_id),
resolved.transport.key.decrypted_auth_config.as_deref(),
);
request_identity_response_encoding_when_redacted(
&mut provider_request_headers,
request_redacted,
);
let resolved_report_kind = if decision_kind == OPENAI_CHAT_STREAM_PLAN_KIND {
"openai_chat_stream_success".to_string()
} else {
"openai_chat_sync_finalize".to_string()
};
let (execution_strategy, conversion_mode) =
ai_local_execution_contract_for_formats("openai:chat", provider_api_format);
Some(LocalOpenAiChatCandidatePayloadParts {
client_api_format: "openai:chat".to_string(),
auth_header: resolved.headers.auth_header,
auth_value: resolved.headers.auth_value,
mapped_model,
provider_api_format: provider_api_format.to_string(),
provider_request_body: resolved.body,
provider_request_headers,
upstream_url: resolved.upstream_url,
execution_strategy,
conversion_mode,
report_kind: resolved_report_kind,
envelope_name: Some(ANTIGRAVITY_V1INTERNAL_ENVELOPE_NAME),
transport: resolved.transport,
request_redacted,
transport_profile: None,
image_request_summary: None,
})
}
#[allow(clippy::too_many_arguments)]
async fn build_gemini_cli_openai_chat_cross_format_payload_parts(
state: &AppState,
@@ -1923,6 +2105,39 @@ mod tests {
}
}
fn sample_antigravity_transport() -> GatewayProviderTransportSnapshot {
let mut transport = sample_gemini_cli_transport();
transport.provider.name = "antigravity".to_string();
transport.provider.provider_type = "antigravity".to_string();
transport.endpoint.base_url = "https://antigravity.googleapis.com".to_string();
transport.endpoint.custom_path = None;
transport.key.auth_type = "oauth".to_string();
transport.key.decrypted_api_key = "__placeholder__".to_string();
transport.key.upstream_metadata = None;
transport.key.decrypted_auth_config = Some(
json!({
"provider_type": "antigravity",
"project_id": "test-antigravity-project",
"client_version": "1.2.3",
"session_id": "sess-antigravity-chat",
"access_token_import_temporary": true,
"headers": {
"Authorization": "Bearer imported-antigravity-token"
}
})
.to_string(),
);
transport
}
fn sample_antigravity_eligible() -> EligibleLocalExecutionCandidate {
let mut eligible = sample_gemini_cli_eligible();
eligible.candidate.provider_name = "antigravity".to_string();
eligible.candidate.provider_type = "antigravity".to_string();
eligible.transport = Arc::new(sample_antigravity_transport());
eligible
}
#[tokio::test]
async fn openai_chat_to_gemini_cli_wraps_cross_format_body_in_v1internal_envelope() {
let state = AppState::new().expect("state should build");
@@ -1989,6 +2204,96 @@ mod tests {
.is_some());
}
#[tokio::test]
async fn openai_chat_to_antigravity_wraps_cross_format_body_in_v1internal_envelope() {
let state = AppState::new().expect("state should build");
let request = http::Request::builder()
.method("POST")
.uri("/v1/chat/completions")
.header(http::header::CONTENT_TYPE, "application/json")
.body(())
.expect("request should build");
let (parts, _) = request.into_parts();
let body_json = json!({
"model": "gemini-3.5-flash",
"messages": [{"role": "user", "content": "hello"}],
"stream": true
});
let payload = resolve_local_openai_chat_candidate_payload_parts(
&state,
&parts,
"trace-openai-chat-antigravity",
&body_json,
&sample_input(),
None,
&sample_antigravity_eligible(),
0,
"candidate-0",
OPENAI_CHAT_STREAM_PLAN_KIND,
"openai_chat_stream_success",
true,
)
.await
.expect("candidate resolution should not fail")
.expect("antigravity candidate should build a payload");
assert_eq!(
payload.upstream_url,
"https://antigravity.googleapis.com/v1internal:streamGenerateContent?alt=sse"
);
assert_eq!(payload.envelope_name, Some("antigravity:v1internal"));
assert_eq!(
payload
.provider_request_headers
.get("authorization")
.map(String::as_str),
Some("Bearer imported-antigravity-token")
);
assert_eq!(
payload
.provider_request_headers
.get("x-client-name")
.map(String::as_str),
Some("antigravity")
);
assert_eq!(
payload
.provider_request_headers
.get("x-client-version")
.map(String::as_str),
Some("1.2.3")
);
assert_eq!(
payload
.provider_request_headers
.get("x-vscode-sessionid")
.map(String::as_str),
Some("sess-antigravity-chat")
);
assert_eq!(
payload.provider_request_body["project"],
"test-antigravity-project"
);
assert_eq!(
payload.provider_request_body["requestId"],
"trace-openai-chat-antigravity"
);
assert_eq!(payload.provider_request_body["model"], "gemini-2.5-pro");
assert_eq!(
payload.provider_request_body["userAgent"],
"antigravity/cli/1.0.16 (aidev_client; os_type=linux; arch=arm64; auth_method=consumer)"
);
assert_eq!(payload.provider_request_body["requestType"], "agent");
assert!(payload.provider_request_body.get("contents").is_none());
assert!(payload.provider_request_body["request"]
.get("contents")
.is_some());
assert!(payload.provider_request_body["request"]
.get("model")
.is_none());
}
#[test]
fn chatgpt_web_chat_image_bridge_body_uses_internal_web_shape() {
let body_json = json!({
@@ -31,11 +31,12 @@ pub(super) fn openai_chat_upstream_is_stream_for_candidate(
crate::ai_serving::transport::kiro::is_kiro_claude_messages_transport(
transport,
provider_api_format,
) || openai_chat_gemini_cli_client_stream_requires_upstream_streaming(
transport,
provider_api_format,
client_is_stream,
);
) || openai_chat_antigravity_requires_upstream_streaming(transport, provider_api_format)
|| openai_chat_gemini_cli_client_stream_requires_upstream_streaming(
transport,
provider_api_format,
client_is_stream,
);
resolve_upstream_is_stream_for_provider(
transport.endpoint.config.as_ref(),
transport.provider.provider_type.as_str(),
@@ -45,6 +46,15 @@ pub(super) fn openai_chat_upstream_is_stream_for_candidate(
)
}
fn openai_chat_antigravity_requires_upstream_streaming(
transport: &GatewayProviderTransportSnapshot,
provider_api_format: &str,
) -> bool {
crate::ai_serving::transport::antigravity::is_antigravity_provider_transport(transport)
&& crate::ai_serving::normalize_api_format_alias(provider_api_format)
== "gemini:generate_content"
}
fn openai_chat_gemini_cli_client_stream_requires_upstream_streaming(
transport: &GatewayProviderTransportSnapshot,
provider_api_format: &str,
@@ -201,4 +211,24 @@ mod tests {
false,
));
}
#[test]
fn openai_chat_policy_resolver_preserves_antigravity_streaming_envelope() {
let antigravity = sample_transport(
"antigravity",
"gemini:generate_content",
Some(json!({"upstream_stream_policy": "force_non_stream"})),
);
assert!(openai_chat_upstream_is_stream_for_candidate(
&antigravity,
"gemini:generate_content",
false,
));
assert!(openai_chat_upstream_is_stream_for_candidate(
&antigravity,
"gemini:generate_content",
true,
));
}
}
@@ -5,6 +5,10 @@ use aether_contracts::ResolvedTransportProfile;
use serde_json::{json, Value};
use tracing::debug;
use crate::ai_serving::planner::antigravity::{
build_antigravity_v1internal_provider_request, AntigravityV1InternalRequestError,
AntigravityV1InternalRequestInput, ANTIGRAVITY_V1INTERNAL_ENVELOPE_NAME,
};
use crate::ai_serving::planner::candidate_preparation::{
prepare_header_authenticated_candidate, prepare_header_authenticated_candidate_from_auth,
OauthPreparationContext,
@@ -29,13 +33,7 @@ use crate::ai_serving::planner::standard::{
build_local_openai_responses_upstream_url, request_body_build_failure_extra_data,
request_conversion_failure_extra_data,
};
use crate::ai_serving::transport::antigravity::{
build_antigravity_safe_v1internal_request, build_antigravity_static_identity_headers,
classify_local_antigravity_request_support, is_antigravity_provider_transport,
AntigravityEnvelopeRequestType, AntigravityRequestAuthUnsupportedReason,
AntigravityRequestEnvelopeSupport, AntigravityRequestSideSupport,
AntigravityRequestSideUnsupportedReason,
};
use crate::ai_serving::transport::antigravity::is_antigravity_provider_transport;
use crate::ai_serving::transport::auth::{
resolve_local_gemini_auth, resolve_local_openai_bearer_auth, resolve_local_standard_auth,
};
@@ -73,8 +71,6 @@ use super::support::{
};
use super::LocalOpenAiResponsesSpec;
const ANTIGRAVITY_ENVELOPE_NAME: &str = "antigravity:v1internal";
fn is_grok_text_provider_api_format(provider_api_format: &str) -> bool {
matches!(
crate::ai_serving::normalize_api_format_alias(provider_api_format).as_str(),
@@ -120,7 +116,7 @@ pub(crate) async fn resolve_local_openai_responses_candidate_payload_parts(
let provider_api_format = eligible.provider_api_format.as_str();
let normalized_provider_api_format =
crate::ai_serving::normalize_api_format_alias(provider_api_format);
let mut transport = Arc::clone(&eligible.transport);
let transport = Arc::clone(&eligible.transport);
let transport_profile = crate::ai_serving::transport::resolve_transport_profile(&transport);
let is_antigravity = is_antigravity_provider_transport(&transport);
let is_gemini_cli = is_gemini_cli_provider_transport(&transport);
@@ -415,82 +411,7 @@ pub(crate) async fn resolve_local_openai_responses_candidate_payload_parts(
provider_api_format,
Some(body_json),
);
let antigravity_auth = if is_antigravity {
let mut antigravity_support = classify_local_antigravity_request_support(
&transport,
&base_provider_request_body,
AntigravityEnvelopeRequestType::Agent,
);
if matches!(
antigravity_support,
AntigravityRequestSideSupport::Unsupported(
AntigravityRequestSideUnsupportedReason::UnsupportedAuth(
AntigravityRequestAuthUnsupportedReason::MissingProjectId
)
)
) {
if let Some(hydrated) = state
.hydrate_antigravity_project_metadata_for_transport(&transport)
.await
{
transport = Arc::new(hydrated);
antigravity_support = classify_local_antigravity_request_support(
&transport,
&base_provider_request_body,
AntigravityEnvelopeRequestType::Agent,
);
}
}
match antigravity_support {
AntigravityRequestSideSupport::Supported(spec) => Some(spec.auth),
AntigravityRequestSideSupport::Unsupported(_) => {
mark_skipped_local_openai_responses_candidate(
state,
input,
trace_id,
candidate,
candidate_index,
candidate_id,
"transport_unsupported",
)
.await;
return Ok(None);
}
}
} else {
None
};
let provider_request_body = if let Some(antigravity_auth) = antigravity_auth.as_ref() {
match build_antigravity_safe_v1internal_request(
antigravity_auth,
trace_id,
&mapped_model,
&base_provider_request_body,
AntigravityEnvelopeRequestType::Agent,
) {
AntigravityRequestEnvelopeSupport::Supported(envelope) => envelope,
AntigravityRequestEnvelopeSupport::Unsupported(_) => {
mark_skipped_local_openai_responses_candidate_with_failure_diagnostic(
state,
input,
trace_id,
candidate,
candidate_index,
candidate_id,
"provider_request_body_build_failed",
CandidateFailureDiagnostic::envelope_build_failed(
spec_metadata.api_format,
provider_api_format,
"openai_responses_antigravity_envelope",
),
)
.await;
return Ok(None);
}
}
} else {
base_provider_request_body
};
let provider_request_body = base_provider_request_body;
if let Some(kiro_auth) = kiro_auth.as_ref() {
return build_kiro_openai_responses_payload_parts(
@@ -562,6 +483,28 @@ pub(crate) async fn resolve_local_openai_responses_candidate_payload_parts(
)
.await);
}
if is_antigravity {
return Ok(build_antigravity_openai_responses_payload_parts(
state,
parts,
trace_id,
body_json,
input,
eligible,
candidate_index,
candidate_id,
spec_metadata.api_format,
&transport,
provider_api_format,
mapped_model,
auth_header,
auth_value,
provider_request_body,
upstream_is_stream,
redaction.redacted,
)
.await);
}
let Some(upstream_url) = (if is_grok && is_grok_text_provider_api_format(provider_api_format) {
Some(build_grok_upstream_url(&transport, GROK_CHAT_PATH))
@@ -598,10 +541,7 @@ pub(crate) async fn resolve_local_openai_responses_candidate_payload_parts(
.await;
return Ok(None);
};
let extra_headers = antigravity_auth
.as_ref()
.map(build_antigravity_static_identity_headers)
.unwrap_or_default();
let extra_headers = BTreeMap::new();
let resolved_headers = if is_grok && is_grok_text_provider_api_format(provider_api_format) {
let Some(headers) = build_grok_browser_headers(GrokHeaderInput {
transport: &transport,
@@ -728,13 +668,8 @@ pub(crate) async fn resolve_local_openai_responses_candidate_payload_parts(
upstream_url,
execution_strategy,
conversion_mode,
is_antigravity: is_antigravity
|| antigravity_auth.is_some() && ANTIGRAVITY_ENVELOPE_NAME == "antigravity:v1internal",
envelope_name: if is_antigravity || antigravity_auth.is_some() {
Some(ANTIGRAVITY_ENVELOPE_NAME)
} else {
None
},
is_antigravity: false,
envelope_name: None,
upstream_is_stream,
transport: Arc::clone(&transport),
transport_profile,
@@ -743,6 +678,160 @@ pub(crate) async fn resolve_local_openai_responses_candidate_payload_parts(
}))
}
#[allow(clippy::too_many_arguments)]
async fn build_antigravity_openai_responses_payload_parts(
state: &AppState,
parts: &http::request::Parts,
trace_id: &str,
original_body_json: &serde_json::Value,
input: &LocalOpenAiResponsesDecisionInput,
eligible: &EligibleLocalExecutionCandidate,
candidate_index: u32,
candidate_id: &str,
client_api_format: &str,
transport: &Arc<GatewayProviderTransportSnapshot>,
provider_api_format: &str,
mapped_model: String,
auth_header: String,
auth_value: String,
gemini_request_body: Value,
upstream_is_stream: bool,
request_redacted: bool,
) -> Option<LocalOpenAiResponsesCandidatePayloadParts> {
let candidate = &eligible.candidate;
let effective_headers = input.effective_headers(&parts.headers);
let resolved =
match build_antigravity_v1internal_provider_request(AntigravityV1InternalRequestInput {
state,
parts,
transport,
trace_id,
mapped_model: &mapped_model,
provider_api_format,
auth_header: &auth_header,
auth_value: &auth_value,
request_headers: effective_headers,
original_request_body: original_body_json,
gemini_request_body: &gemini_request_body,
upstream_is_stream,
same_format: api_format_alias_matches(provider_api_format, client_api_format),
})
.await
{
Ok(resolved) => resolved,
Err(AntigravityV1InternalRequestError::TransportUnsupported) => {
mark_skipped_local_openai_responses_candidate(
state,
input,
trace_id,
candidate,
candidate_index,
candidate_id,
"transport_unsupported",
)
.await;
return None;
}
Err(AntigravityV1InternalRequestError::EnvelopeUnsupported) => {
mark_skipped_local_openai_responses_candidate_with_failure_diagnostic(
state,
input,
trace_id,
candidate,
candidate_index,
candidate_id,
"provider_request_body_build_failed",
CandidateFailureDiagnostic::envelope_build_failed(
client_api_format,
provider_api_format,
"openai_responses_antigravity_envelope",
),
)
.await;
return None;
}
Err(AntigravityV1InternalRequestError::UpstreamUrlUnavailable) => {
mark_skipped_local_openai_responses_candidate_with_failure_diagnostic(
state,
input,
trace_id,
candidate,
candidate_index,
candidate_id,
"upstream_url_missing",
CandidateFailureDiagnostic::upstream_url_missing(
client_api_format,
provider_api_format,
"openai_responses_antigravity_url",
),
)
.await;
return None;
}
Err(AntigravityV1InternalRequestError::HeaderRulesApplyFailed) => {
mark_skipped_local_openai_responses_candidate_with_failure_diagnostic(
state,
input,
trace_id,
candidate,
candidate_index,
candidate_id,
"transport_header_rules_apply_failed",
CandidateFailureDiagnostic::header_rules_apply_failed(
client_api_format,
provider_api_format,
"openai_responses_antigravity_headers",
),
)
.await;
return None;
}
};
let mut provider_request_headers = resolved.headers.headers;
apply_codex_openai_responses_special_headers(
&mut provider_request_headers,
&resolved.body,
effective_headers,
resolved.transport.provider.provider_type.as_str(),
provider_api_format,
Some(trace_id),
resolved.transport.key.decrypted_auth_config.as_deref(),
);
apply_local_auth_config_header_overrides(
&mut provider_request_headers,
resolved.transport.key.decrypted_auth_config.as_deref(),
);
provider_request_headers.insert("accept".to_string(), "text/event-stream".to_string());
request_identity_response_encoding_when_redacted(
&mut provider_request_headers,
request_redacted,
);
let (execution_strategy, conversion_mode) =
ai_local_execution_contract_for_formats(client_api_format, provider_api_format);
Some(LocalOpenAiResponsesCandidatePayloadParts {
auth_header: resolved.headers.auth_header,
auth_value: resolved.headers.auth_value,
mapped_model,
provider_api_format: provider_api_format.to_string(),
provider_request_body: resolved.body,
provider_request_headers,
upstream_url: resolved.upstream_url,
execution_strategy,
conversion_mode,
is_antigravity: true,
envelope_name: Some(ANTIGRAVITY_V1INTERNAL_ENVELOPE_NAME),
upstream_is_stream,
transport: Arc::clone(&resolved.transport),
transport_profile: crate::ai_serving::transport::resolve_transport_profile(
&resolved.transport,
),
image_request_summary: None,
request_redacted,
})
}
#[allow(clippy::too_many_arguments)]
async fn build_gemini_cli_openai_responses_payload_parts(
state: &AppState,
+2
View File
@@ -1,6 +1,7 @@
pub(crate) fn normalized_signature(api_format: &str) -> Option<&'static str> {
match crate::ai_serving::normalize_api_format_alias(api_format).as_str() {
"gemini:generate_content" => Some("gemini:generate_content"),
"gemini:interactions" => Some("gemini:interactions"),
"gemini:embedding" => Some("gemini:embedding"),
"gemini:video" => Some("gemini:video"),
"gemini:files" => Some("gemini:files"),
@@ -11,6 +12,7 @@ pub(crate) fn normalized_signature(api_format: &str) -> Option<&'static str> {
pub(crate) fn local_path(api_format: &str) -> Option<&'static str> {
match crate::ai_serving::normalize_api_format_alias(api_format).as_str() {
"gemini" | "gemini:generate_content" => Some("/v1beta/models/{model}:{action}"),
"gemini:interactions" => Some("/v1/interactions"),
"gemini:embedding" => Some("/v1beta/models/{model}:{action}"),
"gemini:video" => Some("/v1beta/models/{model}:predictLongRunning"),
"gemini:files" => Some("/v1beta/files"),
@@ -17,13 +17,17 @@ const AI_POST_ROUTE_PATTERNS: &[&str] = &[
"/v1/responses/compact",
"/v1/images/generations",
"/v1/images/edits",
"/v1/interactions",
"/v1beta/interactions",
"/v1internal:loadCodeAssist",
"/v1internal:fetchAvailableModels",
"/v1internal:retrieveUserQuotaSummary",
"/v1internal:fetchUserInfo",
"/v1internal:fetchAdminControls",
"/v1internal:setUserSettings",
"/v1internal:listExperiments",
"/v1internal:recordCodeAssistMetrics",
"/v1internal:writeTrajectoryAcls",
"/v1internal:streamGenerateContent",
];
@@ -95,6 +99,12 @@ mod tests {
fn supports_data_api_endpoint_signatures_and_public_paths() {
for (api_format, family, kind, path) in [
("openai:embedding", "openai", "embedding", "/v1/embeddings"),
(
"gemini:interactions",
"gemini",
"interactions",
"/v1/interactions",
),
(
"gemini:embedding",
"gemini",
+2
View File
@@ -141,11 +141,13 @@ pub(crate) const RUST_FRONTDOOR_OWNED_ROUTE_PATTERNS: &[&str] = &[
"/v1beta/files/{path...}",
"/v1internal:loadCodeAssist",
"/v1internal:fetchAvailableModels",
"/v1internal:retrieveUserQuotaSummary",
"/v1internal:fetchUserInfo",
"/v1internal:fetchAdminControls",
"/v1internal:setUserSettings",
"/v1internal:listExperiments",
"/v1internal:recordCodeAssistMetrics",
"/v1internal:writeTrajectoryAcls",
"/v1internal:streamGenerateContent",
"/",
"/{*path}",
@@ -97,6 +97,17 @@ pub(super) fn classify_ai_public_route(
"openai:video",
true,
))
} else if method == http::Method::POST
&& matches!(normalized_path, "/v1/interactions" | "/v1beta/interactions")
{
Some(classified_with_request_auth_channel(
"ai_public",
"gemini",
"interactions",
"api_key",
"gemini:interactions",
true,
))
} else if method == http::Method::POST && is_gemini_models_route(normalized_path) {
if normalized_path.ends_with(":predictLongRunning") {
Some(classified(
@@ -182,11 +193,13 @@ fn classify_antigravity_v1internal_route(
let (route_kind, execution_runtime_candidate) = match action {
"loadCodeAssist" => ("load_code_assist", false),
"fetchAvailableModels" => ("fetch_available_models", false),
"retrieveUserQuotaSummary" => ("retrieve_user_quota_summary", 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),
"writeTrajectoryAcls" => ("write_trajectory_acls", false),
"streamGenerateContent" => ("stream_generate_content", true),
_ => return None,
};
@@ -274,6 +274,23 @@ fn classifies_gemini_batch_embed_contents_as_embedding_route() {
assert!(decision.is_execution_runtime_candidate());
}
#[test]
fn classifies_gemini_interactions_as_interactions_route() {
let headers = headers(&[("x-goog-api-key", "gemini-key")]);
let uri: Uri = "/v1/interactions".parse().expect("uri should parse");
let decision =
classify_control_route(&http::Method::POST, &uri, &headers).expect("route should classify");
assert_eq!(decision.route_family.as_deref(), Some("gemini"));
assert_eq!(decision.route_kind.as_deref(), Some("interactions"));
assert_eq!(decision.request_auth_channel.as_deref(), Some("api_key"));
assert_eq!(
decision.auth_endpoint_signature.as_deref(),
Some("gemini:interactions")
);
assert!(decision.is_execution_runtime_candidate());
}
#[test]
fn classifies_gemini_predict_long_running_as_video_route() {
let headers = headers(&[]);
@@ -302,6 +319,10 @@ fn classifies_antigravity_v1internal_control_plane_routes() {
for (path, route_kind) in [
("/v1internal:loadCodeAssist", "load_code_assist"),
("/v1internal:fetchAvailableModels", "fetch_available_models"),
(
"/v1internal:retrieveUserQuotaSummary",
"retrieve_user_quota_summary",
),
("/v1internal:fetchUserInfo", "fetch_user_info"),
("/v1internal:fetchAdminControls", "fetch_admin_controls"),
("/v1internal:setUserSettings", "set_user_settings"),
@@ -310,6 +331,7 @@ fn classifies_antigravity_v1internal_control_plane_routes() {
"/v1internal:recordCodeAssistMetrics",
"record_code_assist_metrics",
),
("/v1internal:writeTrajectoryAcls", "write_trajectory_acls"),
] {
let uri: Uri = path.parse().expect("uri should parse");
let decision = classify_control_route(&http::Method::POST, &uri, &headers)
@@ -179,6 +179,7 @@ pub(crate) fn resolve_core_sync_error_finalize_report_kind(
"openai_responses_compact_sync" => "openai_responses_compact_sync_finalize",
"claude_chat_sync" => "claude_chat_sync_finalize",
"gemini_chat_sync" => "gemini_chat_sync_finalize",
"gemini_interactions_sync" => "gemini_interactions_sync_finalize",
"claude_cli_sync" => "claude_cli_sync_finalize",
"gemini_cli_sync" => "gemini_cli_sync_finalize",
_ => return None,
@@ -324,6 +325,7 @@ pub(crate) fn resolve_core_stream_error_finalize_report_kind(
"openai_chat_stream" => "openai_chat_sync_finalize",
"claude_chat_stream" => "claude_chat_sync_finalize",
"gemini_chat_stream" => "gemini_chat_sync_finalize",
"gemini_interactions_stream" => "gemini_interactions_sync_finalize",
"openai_responses_stream" => "openai_responses_sync_finalize",
"openai_responses_compact_stream" => "openai_responses_compact_sync_finalize",
"claude_cli_stream" => "claude_cli_sync_finalize",
@@ -340,6 +342,7 @@ pub(crate) fn resolve_core_stream_direct_finalize_report_kind(plan_kind: &str) -
"openai_image_stream" => "openai_image_sync_finalize",
"claude_chat_stream" => "claude_chat_sync_finalize",
"gemini_chat_stream" => "gemini_chat_sync_finalize",
"gemini_interactions_stream" => "gemini_interactions_sync_finalize",
"openai_responses_stream" => "openai_responses_sync_finalize",
"openai_responses_compact_stream" => "openai_responses_compact_sync_finalize",
"claude_cli_stream" => "claude_cli_sync_finalize",
@@ -2637,7 +2637,20 @@ async fn execute_execution_runtime_sync_impl(
candidate_id,
)?;
if !usage_payload.report_kind.trim().is_empty() {
spawn_sync_report(state.clone(), usage_payload);
if status_code >= 400 {
let report_kind = usage_payload.report_kind.clone();
if let Err(err) = submit_sync_report(state, usage_payload).await {
warn!(
event_name = "local_sync_error_report_submit_failed",
log_type = "ops",
trace_id = %trace_id,
report_kind = %report_kind,
"gateway failed to submit local sync error report before returning response: {err:?}"
);
}
} else {
spawn_sync_report(state.clone(), usage_payload);
}
}
Ok(Some(response))
@@ -104,6 +104,10 @@ fn resolve_core_error_background_report_kind_maps_all_core_finalize_kinds() {
("openai_chat_sync_finalize", Some("openai_chat_sync_error")),
("claude_chat_sync_finalize", Some("claude_chat_sync_error")),
("gemini_chat_sync_finalize", Some("gemini_chat_sync_error")),
(
"gemini_interactions_sync_finalize",
Some("gemini_interactions_sync_error"),
),
(
"openai_responses_sync_finalize",
Some("openai_responses_sync_error"),
@@ -151,6 +155,10 @@ fn resolve_core_success_background_report_kind_maps_all_core_finalize_kinds() {
"gemini_chat_sync_finalize",
Some("gemini_chat_sync_success"),
),
(
"gemini_interactions_sync_finalize",
Some("gemini_interactions_sync_success"),
),
(
"openai_responses_sync_finalize",
Some("openai_responses_sync_success"),
@@ -36,8 +36,8 @@ use crate::provider_key_auth::{
use crate::provider_transport::antigravity::{
build_antigravity_safe_v1internal_request, build_antigravity_static_identity_headers,
classify_local_antigravity_request_support, AntigravityEnvelopeRequestType,
AntigravityRequestEnvelopeSupport, AntigravityRequestSideSupport,
AntigravityRequestSideUnsupportedReason,
AntigravityRequestAuthUnsupportedReason, AntigravityRequestEnvelopeSupport,
AntigravityRequestSideSupport, AntigravityRequestSideUnsupportedReason,
};
use crate::provider_transport::kiro::{
build_kiro_generate_assistant_response_url, build_kiro_provider_headers,
@@ -2061,7 +2061,7 @@ async fn provider_query_execute_openai_image_test_candidate(
trace_id: &str,
requested_model: &str,
) -> Result<ProviderQueryExecutionOutcome, GatewayError> {
let Some(transport) = state
let Some(mut transport) = state
.read_provider_transport_snapshot(&provider.id, &candidate.endpoint.id, &candidate.key.id)
.await?
else {
@@ -2429,7 +2429,7 @@ async fn provider_query_execute_antigravity_test_candidate(
trace_id: &str,
requested_model: &str,
) -> Result<ProviderQueryExecutionOutcome, GatewayError> {
let Some(transport) = state
let Some(mut transport) = state
.read_provider_transport_snapshot(&provider.id, &candidate.endpoint.id, &candidate.key.id)
.await?
else {
@@ -2463,11 +2463,34 @@ async fn provider_query_execute_antigravity_test_candidate(
));
};
let antigravity_spec = match classify_local_antigravity_request_support(
let mut antigravity_support = classify_local_antigravity_request_support(
&transport,
&base_provider_request_body,
AntigravityEnvelopeRequestType::EndpointTest,
);
if matches!(
antigravity_support,
AntigravityRequestSideSupport::Unsupported(
AntigravityRequestSideUnsupportedReason::UnsupportedAuth(
AntigravityRequestAuthUnsupportedReason::MissingProjectId
)
)
) {
if let Some(hydrated) = state
.app()
.hydrate_antigravity_project_metadata_for_transport(&transport)
.await
{
transport = hydrated;
antigravity_support = classify_local_antigravity_request_support(
&transport,
&base_provider_request_body,
AntigravityEnvelopeRequestType::EndpointTest,
);
}
}
let antigravity_spec = match antigravity_support {
AntigravityRequestSideSupport::Supported(spec) => spec,
AntigravityRequestSideSupport::Unsupported(reason) => {
let reason = provider_query_antigravity_unsupported_reason(reason);
@@ -2,7 +2,8 @@ use super::DEFAULT_PROVIDER_QUERY_TEST_MESSAGE;
use crate::handlers::admin::request::{AdminAppState, AdminGatewayProviderTransportSnapshot};
use crate::provider_transport::antigravity::{
classify_local_antigravity_request_support, AntigravityEnvelopeRequestType,
AntigravityRequestSideSupport, AntigravityRequestSideUnsupportedReason,
AntigravityRequestAuthUnsupportedReason, AntigravityRequestSideSupport,
AntigravityRequestSideUnsupportedReason,
};
use crate::provider_transport::kiro::supports_local_kiro_request_transport_with_network;
use serde_json::{json, Value};
@@ -82,14 +83,14 @@ pub(super) fn provider_query_standard_test_unsupported_reason(
api_format,
)
}
"gemini:generate_content" | "gemini:embedding"
"gemini:generate_content" | "gemini:embedding" | "gemini:interactions"
if crate::provider_transport::is_vertex_transport_context(transport) =>
{
aether_provider_transport::vertex::local_vertex_gemini_transport_unsupported_reason_with_network(
transport,
)
}
"gemini:generate_content" | "gemini:embedding" => {
"gemini:generate_content" | "gemini:embedding" | "gemini:interactions" => {
crate::provider_transport::policy::local_gemini_transport_unsupported_reason_with_network(
transport,
api_format,
@@ -130,9 +131,21 @@ pub(super) fn provider_query_antigravity_unsupported_reason(
AntigravityRequestSideUnsupportedReason::UnsupportedNetworkConfig => {
"transport_network_config_unsupported"
}
AntigravityRequestSideUnsupportedReason::UnsupportedAuth(_) => {
"transport_antigravity_auth_unsupported"
}
AntigravityRequestSideUnsupportedReason::UnsupportedAuth(
AntigravityRequestAuthUnsupportedReason::WrongProviderType,
) => "transport_provider_type_unsupported",
AntigravityRequestSideUnsupportedReason::UnsupportedAuth(
AntigravityRequestAuthUnsupportedReason::MissingAuthConfig,
) => "transport_antigravity_auth_config_missing",
AntigravityRequestSideUnsupportedReason::UnsupportedAuth(
AntigravityRequestAuthUnsupportedReason::InvalidAuthConfigJson,
) => "transport_antigravity_auth_config_invalid",
AntigravityRequestSideUnsupportedReason::UnsupportedAuth(
AntigravityRequestAuthUnsupportedReason::ComplexDynamicAuthConfig,
) => "transport_antigravity_auth_config_unsupported",
AntigravityRequestSideUnsupportedReason::UnsupportedAuth(
AntigravityRequestAuthUnsupportedReason::MissingProjectId,
) => "transport_antigravity_project_id_missing",
AntigravityRequestSideUnsupportedReason::UnsupportedEnvelope(_) => {
"transport_antigravity_envelope_unsupported"
}
@@ -250,6 +263,7 @@ pub(super) fn provider_query_test_adapter_for_provider_api_format(
| "openai:responses:compact"
| "claude:messages"
| "gemini:generate_content"
| "gemini:interactions"
| "openai:embedding"
| "gemini:embedding"
| "jina:embedding"
@@ -366,7 +380,7 @@ pub(super) fn provider_query_transport_supports_model_test_execution(
transport, api_format,
)
}
"gemini:generate_content" | "gemini:embedding" => {
"gemini:generate_content" | "gemini:embedding" | "gemini:interactions" => {
if crate::provider_transport::is_vertex_transport_context(transport) {
aether_provider_transport::vertex::supports_local_vertex_gemini_transport_with_network(transport)
} else {
@@ -120,6 +120,10 @@ fn provider_query_endpoint_route_payload(
("Gemini API", "gemini_native", "generateContent", "")
}
"gemini:generate_content" => ("Gemini native", "gemini_native", "generateContent", ""),
"gemini:interactions" if is_gemini_api => {
("Gemini API", "gemini_native", "interactions", "")
}
"gemini:interactions" => ("Gemini native", "gemini_native", "interactions", ""),
"openai:embedding" if is_vertex && is_openai_compat => (
"Vertex AI OpenAI-compatible",
"openai_compatible",
@@ -678,6 +678,10 @@ fn provider_query_test_adapter_routes_fixed_provider_endpoint_types() {
provider_query_test_adapter_for_provider_api_format("custom", "gemini:embedding"),
Some(ProviderQueryTestAdapter::Standard)
);
assert_eq!(
provider_query_test_adapter_for_provider_api_format("custom", "gemini:interactions"),
Some(ProviderQueryTestAdapter::Standard)
);
assert_eq!(
provider_query_test_adapter_for_provider_api_format(
"aliyun",
@@ -926,6 +926,9 @@ fn maybe_build_local_antigravity_v1internal_response(
"fetch_available_models" => {
Some(Json(build_antigravity_fetch_available_models_payload()).into_response())
}
"retrieve_user_quota_summary" => {
Some(Json(build_antigravity_retrieve_user_quota_summary_payload()).into_response())
}
"fetch_user_info" => {
Some(Json(build_antigravity_fetch_user_info_payload()).into_response())
}
@@ -938,6 +941,7 @@ fn maybe_build_local_antigravity_v1internal_response(
.into_response(),
),
"record_code_assist_metrics" => Some(Json(json!({})).into_response()),
"write_trajectory_acls" => Some(Json(json!({})).into_response()),
"set_user_settings" => Some(build_antigravity_set_user_settings_response(request_body)),
"stream_generate_content" => None,
_ => None,
@@ -1042,27 +1046,42 @@ fn build_antigravity_fetch_user_info_payload() -> Value {
})
}
fn build_antigravity_retrieve_user_quota_summary_payload() -> Value {
json!({
"description": "",
"groups": []
})
}
fn build_antigravity_default_user_settings_payload() -> Value {
json!({
"preferredModelId": "gemini-3.1-flash-lite"
"preferredModelId": "gemini-3.5-flash-low"
})
}
fn build_antigravity_fetch_available_models_payload() -> Value {
json!({
"models": {
"gemini-3.5-flash-low": antigravity_model_payload("gemini-3.5-flash-low", "Gemini 3.5 Flash Low"),
"gemini-3-flash-agent": antigravity_model_payload("gemini-3-flash-agent", "Gemini 3 Flash Agent"),
"gemini-pro-agent": antigravity_model_payload("gemini-pro-agent", "Gemini 3.1 Pro (High)"),
"gemini-3.1-pro-low": antigravity_model_payload("gemini-3.1-pro-low", "Gemini 3.1 Pro (Low)"),
"gemini-3-flash-agent": antigravity_model_payload("gemini-3-flash-agent", "Gemini 3.5 Flash (High)"),
"gemini-3.5-flash-low": antigravity_model_payload("gemini-3.5-flash-low", "Gemini 3.5 Flash (Medium)"),
"gemini-3.5-flash-extra-low": antigravity_model_payload("gemini-3.5-flash-extra-low", "Gemini 3.5 Flash (Low)"),
"claude-opus-4-6-thinking": antigravity_model_payload("claude-opus-4-6-thinking", "Claude Opus 4.6 (Thinking)"),
"claude-sonnet-4-6": antigravity_model_payload("claude-sonnet-4-6", "Claude Sonnet 4.6 (Thinking)"),
"gpt-oss-120b-medium": antigravity_model_payload("gpt-oss-120b-medium", "GPT-OSS 120B (Medium)"),
"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-flash": antigravity_model_payload("gemini-2.5-flash", "Gemini 3.1 Flash Lite"),
"gemini-2.5-flash-lite": antigravity_model_payload("gemini-2.5-flash-lite", "Gemini 3.1 Flash Lite"),
"gemini-2.5-flash-thinking": antigravity_model_payload("gemini-2.5-flash-thinking", "Gemini 3.1 Flash Lite"),
"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"),
"gemini-3.1-pro-high": antigravity_model_payload("gemini-3.1-pro-high", "Gemini 3.1 Pro (High)"),
"chat_20706": antigravity_model_payload("chat_20706", ""),
"chat_23310": antigravity_model_payload("chat_23310", ""),
"tab_flash_lite_preview": antigravity_model_payload("tab_flash_lite_preview", ""),
"tab_jump_flash_lite_preview": antigravity_model_payload("tab_jump_flash_lite_preview", ""),
"models/proactive-observer": antigravity_model_payload("models/proactive-observer", "Proactive Observer")
},
"agentModelSorts": [
@@ -1071,10 +1090,14 @@ fn build_antigravity_fetch_available_models_payload() -> Value {
"groups": [
{
"modelIds": [
"gemini-3.1-flash-lite",
"gemini-3.5-flash-low",
"gemini-3-flash-agent",
"gemini-3.5-flash-extra-low",
"gemini-3.1-pro-low",
"gemini-3.5-flash-low"
"gemini-pro-agent",
"claude-sonnet-4-6",
"claude-opus-4-6-thinking",
"gpt-oss-120b-medium"
]
}
]
@@ -1083,12 +1106,18 @@ fn build_antigravity_fetch_available_models_payload() -> Value {
"audioTranscriptionModelIds": ["models/proactive-observer"],
"commandModelIds": ["gemini-3-flash"],
"commitMessageModelIds": ["gemini-3.1-flash-lite"],
"defaultAgentModelId": "gemini-3.1-flash-lite",
"deprecatedModelIds": {},
"defaultAgentModelId": "gemini-3.5-flash-low",
"deprecatedModelIds": {
"gemini-3.1-pro-high": {
"newModelEnum": "MODEL_PLACEHOLDER_M16",
"newModelId": "gemini-pro-agent",
"oldModelEnum": "MODEL_PLACEHOLDER_M37"
}
},
"experimentIds": [],
"imageGenerationModelIds": ["gemini-3.1-flash-image"],
"mqueryModelIds": ["gemini-3.1-flash-lite"],
"tabModelIds": ["tab_flash_lite_preview", "tab_jump_flash_lite_preview"],
"tabModelIds": ["chat_20706", "chat_23310"],
"tieredModelIds": {
"flash": ["gemini-3-flash-agent"],
"flashLite": ["gemini-3.1-flash-lite"],
@@ -1100,6 +1129,11 @@ fn build_antigravity_fetch_available_models_payload() -> Value {
fn antigravity_model_payload(id: &str, display_name: &str) -> Value {
let model = match id {
"chat_20706" => "MODEL_CHAT_20706",
"chat_23310" => "MODEL_CHAT_23310",
"claude-opus-4-6-thinking" => "MODEL_PLACEHOLDER_M26",
"claude-sonnet-4-6" => "MODEL_PLACEHOLDER_M35",
"gpt-oss-120b-medium" => "MODEL_OPENAI_GPT_OSS_120B_MEDIUM",
"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",
@@ -1108,22 +1142,43 @@ fn antigravity_model_payload(id: &str, display_name: &str) -> Value {
"gemini-3-flash-agent" => "MODEL_PLACEHOLDER_M132",
"gemini-3.1-flash-image" => "MODEL_PLACEHOLDER_M21",
"gemini-3.1-flash-lite" => "MODEL_PLACEHOLDER_M50",
"gemini-pro-agent" => "MODEL_PLACEHOLDER_M16",
"gemini-3.1-pro-high" => "MODEL_PLACEHOLDER_M37",
"gemini-3.1-pro-low" => "MODEL_PLACEHOLDER_M36",
"gemini-3.5-flash-low" => "MODEL_PLACEHOLDER_M20",
"gemini-3.5-flash-extra-low" => "MODEL_PLACEHOLDER_M187",
"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",
let (api_provider, model_provider) = match id {
"claude-opus-4-6-thinking" | "claude-sonnet-4-6" => {
("API_PROVIDER_ANTHROPIC_VERTEX", "MODEL_PROVIDER_ANTHROPIC")
}
"gpt-oss-120b-medium" => ("API_PROVIDER_OPENAI_VERTEX", "MODEL_PROVIDER_OPENAI"),
"chat_20706" | "chat_23310" => ("API_PROVIDER_INTERNAL", "MODEL_PROVIDER_GOOGLE"),
_ => ("API_PROVIDER_GOOGLE_GEMINI", "MODEL_PROVIDER_GOOGLE"),
};
let mut payload = json!({
"apiProvider": api_provider,
"displayName": display_name,
"maxOutputTokens": 65536,
"maxTokens": 1048576,
"minThinkingBudget": 32,
"model": model,
"modelProvider": "MODEL_PROVIDER_GOOGLE",
"recommended": id == "gemini-3.1-flash-lite",
"modelProvider": model_provider,
"recommended": matches!(
id,
"gemini-3.5-flash-low"
| "gemini-3-flash-agent"
| "gemini-3.5-flash-extra-low"
| "gemini-3.1-pro-low"
| "gemini-pro-agent"
| "claude-sonnet-4-6"
| "claude-opus-4-6-thinking"
| "gpt-oss-120b-medium"
),
"supportedMimeTypes": {
"application/json": true,
"application/pdf": true,
@@ -1137,7 +1192,88 @@ fn antigravity_model_payload(id: &str, display_name: &str) -> Value {
"supportsVideo": true,
"thinkingBudget": 4000,
"tokenizerType": "LLAMA_WITH_SPECIAL"
})
});
if let Some(object) = payload.as_object_mut() {
match id {
"gemini-3-flash-agent" => {
object.insert("thinkingBudget".to_string(), json!(10000));
}
"gemini-pro-agent" => {
object.insert("maxOutputTokens".to_string(), json!(65535));
object.insert("thinkingBudget".to_string(), json!(10001));
}
"gemini-3.1-pro-high" => {
object.insert("maxOutputTokens".to_string(), json!(65535));
object.insert("minThinkingBudget".to_string(), json!(128));
object.insert("thinkingBudget".to_string(), json!(10001));
}
"gemini-3.5-flash-extra-low" => {
object.insert("thinkingBudget".to_string(), json!(1000));
object.insert("maxOutputTokens".to_string(), json!(65536));
}
"gemini-3.1-pro-low" => {
object.insert("maxOutputTokens".to_string(), json!(65535));
object.insert("minThinkingBudget".to_string(), json!(128));
object.insert("thinkingBudget".to_string(), json!(1001));
}
"claude-sonnet-4-6" | "claude-opus-4-6-thinking" => {
object.insert("maxTokens".to_string(), json!(250000));
object.insert("maxOutputTokens".to_string(), json!(64000));
object.insert("thinkingBudget".to_string(), json!(1024));
object.remove("minThinkingBudget");
object.remove("supportsVideo");
}
"gpt-oss-120b-medium" => {
object.insert("maxTokens".to_string(), json!(131072));
object.insert("maxOutputTokens".to_string(), json!(32768));
object.insert("thinkingBudget".to_string(), json!(8192));
object.remove("minThinkingBudget");
object.remove("supportsImages");
object.remove("supportsVideo");
}
"chat_20706" => {
object.insert("maxTokens".to_string(), json!(16384));
object.remove("displayName");
object.remove("maxOutputTokens");
object.remove("minThinkingBudget");
object.remove("recommended");
object.remove("supportsImages");
object.remove("supportsThinking");
object.remove("supportsVideo");
object.remove("thinkingBudget");
object.remove("tokenizerType");
object.remove("supportedMimeTypes");
}
"chat_23310" => {
object.insert("maxTokens".to_string(), json!(32768));
object.remove("displayName");
object.remove("maxOutputTokens");
object.remove("minThinkingBudget");
object.remove("recommended");
object.remove("supportsImages");
object.remove("supportsThinking");
object.remove("supportsVideo");
object.remove("thinkingBudget");
object.remove("tokenizerType");
object.remove("supportedMimeTypes");
}
"tab_flash_lite_preview" | "tab_jump_flash_lite_preview" => {
object.insert("maxTokens".to_string(), json!(16384));
object.insert("maxOutputTokens".to_string(), json!(4096));
object.remove("displayName");
object.remove("minThinkingBudget");
object.remove("recommended");
object.remove("supportsImages");
object.remove("supportsThinking");
object.remove("supportsVideo");
object.remove("thinkingBudget");
object.remove("tokenizerType");
object.remove("supportedMimeTypes");
}
_ => {}
}
}
payload
}
async fn maybe_build_local_gemini_video_operations_response(
@@ -102,7 +102,7 @@ pub(super) async fn maybe_build_local_test_connection_route_response(
let format_value = crate::ai_serving::normalize_api_format_alias(&format_value);
if !matches!(
format_value.as_str(),
"openai:chat" | "claude:messages" | "gemini:generate_content"
"openai:chat" | "claude:messages" | "gemini:generate_content" | "gemini:interactions"
) {
return None;
}
@@ -186,6 +186,10 @@ pub(super) async fn maybe_build_local_test_connection_route_response(
"parts": [{"text": "Health check"}],
}],
}),
"gemini:interactions" => json!({
"model": model,
"input": "Health check",
}),
_ => return None,
};
if !crate::provider_transport::apply_local_body_rules(
@@ -197,7 +201,7 @@ pub(super) async fn maybe_build_local_test_connection_route_response(
}
let oauth_auth = match format_value.as_str() {
"openai:chat" | "claude:messages" | "gemini:generate_content" => {
"openai:chat" | "claude:messages" | "gemini:generate_content" | "gemini:interactions" => {
match state.resolve_local_oauth_request_auth(&transport).await {
Ok(Some(crate::provider_transport::LocalResolvedOAuthRequestAuth::Header {
name,
@@ -218,7 +222,7 @@ pub(super) async fn maybe_build_local_test_connection_route_response(
crate::provider_transport::auth::resolve_local_standard_auth(&transport)
.or(oauth_auth.clone())
}
"gemini:generate_content" => {
"gemini:generate_content" | "gemini:interactions" => {
crate::provider_transport::auth::resolve_local_gemini_auth(&transport)
.or(oauth_auth.clone())
}
@@ -577,6 +577,44 @@ fn model_quota_window_snapshot(
Some(Value::Object(window))
}
fn canonical_antigravity_model_label(model_name: &str) -> Option<&'static str> {
match model_name.trim() {
"claude-opus-4-6-thinking" => Some("Claude Opus 4.6 (Thinking)"),
"claude-sonnet-4-6" | "claude-sonnet-4-6-thinking" => Some("Claude Sonnet 4.6 (Thinking)"),
"gemini-3-flash-agent" => Some("Gemini 3.5 Flash (High)"),
"gemini-3.5-flash-low" => Some("Gemini 3.5 Flash (Medium)"),
"gemini-3.5-flash-extra-low" => Some("Gemini 3.5 Flash (Low)"),
"gemini-3.1-pro-high" | "gemini-pro-agent" => Some("Gemini 3.1 Pro (High)"),
"gemini-3.1-pro-low" => Some("Gemini 3.1 Pro (Low)"),
"gemini-3.1-flash-image" => Some("Gemini 3.1 Flash Image"),
"gemini-3.1-flash-lite" => Some("Gemini 3.1 Flash Lite"),
"gemini-3-flash" => Some("Gemini 3 Flash"),
"gemini-2.5-pro" => Some("Gemini 2.5 Pro"),
"gemini-2.5-flash-thinking" | "gemini-2.5-flash" | "gemini-2.5-flash-lite" => {
Some("Gemini 3.1 Flash Lite")
}
"gpt-oss-120b-medium" => Some("GPT-OSS 120B (Medium)"),
"tab_flash_lite_preview" => Some("Tab Flash Lite Preview"),
"tab_jump_flash_lite_preview" => Some("Tab Jump Flash Lite Preview"),
"models/proactive-observer" => Some("Proactive Observer"),
_ => None,
}
}
fn antigravity_model_quota_window_snapshot(
model_name: &str,
item: &Map<String, Value>,
observed_at_unix_secs: Option<u64>,
) -> Option<Value> {
let mut window = model_quota_window_snapshot(model_name, item, observed_at_unix_secs)?;
if let Some(label) = canonical_antigravity_model_label(model_name) {
if let Some(window) = window.as_object_mut() {
window.insert("label".to_string(), json!(label));
}
}
Some(window)
}
fn provider_quota_metadata_string(
metadata: &Map<String, Value>,
fields: &[&str],
@@ -1466,7 +1504,7 @@ fn build_antigravity_quota_status_snapshot(
models
.iter()
.filter_map(|(model_name, item)| {
model_quota_window_snapshot(
antigravity_model_quota_window_snapshot(
model_name,
item.as_object()?,
observed_at_unix_secs,
@@ -3715,6 +3753,77 @@ mod tests {
);
}
#[test]
fn sync_provider_key_quota_status_snapshot_labels_antigravity_models_by_model_id() {
let upstream_metadata = json!({
"antigravity": {
"updated_at": 1_775_553_285u64,
"quota_by_model": {
"gemini-3.5-flash-extra-low": {
"display_name": "Gemini 3.5 Flash (Low)",
"remaining_fraction": 1.0
},
"gemini-3.5-flash-low": {
"display_name": "Gemini 3.5 Flash (Medium)",
"remaining_fraction": 0.75
},
"gemini-3-flash-agent": {
"display_name": "Gemini 3.5 Flash (High)",
"remaining_fraction": 0.5
},
"gemini-2.5-flash": {
"display_name": "Gemini 3.1 Flash Lite",
"remaining_fraction": 0.4
},
"claude-sonnet-4-6": {
"display_name": "Claude Sonnet 4.6 (Thinking)",
"remaining_fraction": 0.3
}
}
}
});
let payload = sync_provider_key_quota_status_snapshot(
None,
"antigravity",
Some(&upstream_metadata),
"refresh_api",
)
.expect("quota snapshot should sync");
let windows = payload["quota"]["windows"]
.as_array()
.expect("quota windows should exist");
let label_for_model = |model: &str| {
windows
.iter()
.filter_map(Value::as_object)
.find(|window| window.get("model") == Some(&json!(model)))
.and_then(|window| window.get("label"))
.cloned()
};
assert_eq!(
label_for_model("gemini-3.5-flash-extra-low"),
Some(json!("Gemini 3.5 Flash (Low)"))
);
assert_eq!(
label_for_model("gemini-3.5-flash-low"),
Some(json!("Gemini 3.5 Flash (Medium)"))
);
assert_eq!(
label_for_model("gemini-3-flash-agent"),
Some(json!("Gemini 3.5 Flash (High)"))
);
assert_eq!(
label_for_model("gemini-2.5-flash"),
Some(json!("Gemini 3.1 Flash Lite"))
);
assert_eq!(
label_for_model("claude-sonnet-4-6"),
Some(json!("Claude Sonnet 4.6 (Thinking)"))
);
}
#[test]
fn provider_key_status_snapshot_payload_backfills_account_block_from_oauth_invalid_reason() {
let mut key = sample_catalog_key();
@@ -1469,15 +1469,36 @@ async fn gateway_executes_openai_chat_cross_format_tool_use_upstream_stream_via_
}
#[test]
fn gateway_skips_openai_chat_antigravity_cross_format_sync_candidate_as_transport_unsupported() {
fn gateway_executes_openai_chat_antigravity_cross_format_sync_via_local_finalize_response() {
run_openai_chat_finalize_test(
"gateway_skips_openai_chat_antigravity_cross_format_sync_candidate_as_transport_unsupported",
gateway_skips_openai_chat_antigravity_cross_format_sync_candidate_as_transport_unsupported_impl,
"gateway_executes_openai_chat_antigravity_cross_format_sync_via_local_finalize_response",
gateway_executes_openai_chat_antigravity_cross_format_sync_via_local_finalize_response_impl,
);
}
async fn gateway_skips_openai_chat_antigravity_cross_format_sync_candidate_as_transport_unsupported_impl(
async fn gateway_executes_openai_chat_antigravity_cross_format_sync_via_local_finalize_response_impl(
) {
use base64::Engine as _;
#[derive(Debug, Clone)]
struct SeenExecutionRuntimeSyncRequest {
trace_id: String,
url: String,
accept: String,
authorization: String,
x_client_name: String,
x_client_version: String,
x_vscode_sessionid: String,
x_goog_api_client: String,
project: String,
request_id: String,
model: String,
user_agent: String,
request_type: String,
contents_len: usize,
request_has_model: bool,
}
fn hash_api_key(value: &str) -> String {
let mut hasher = Sha256::new();
hasher.update(value.as_bytes());
@@ -1582,11 +1603,11 @@ async fn gateway_skips_openai_chat_antigravity_cross_format_sync_candidate_as_tr
)
.expect("endpoint should build")
.with_transport_fields(
"https://generativelanguage.googleapis.com".to_string(),
"https://antigravity.googleapis.com".to_string(),
None,
None,
Some(2),
Some("/v1beta/models/claude-sonnet-4-5:generateContent".to_string()),
None,
None,
None,
None,
@@ -1595,6 +1616,11 @@ async fn gateway_skips_openai_chat_antigravity_cross_format_sync_candidate_as_tr
}
fn sample_provider_catalog_key() -> StoredProviderCatalogKey {
let encrypted_auth_config = encrypt_python_fernet_plaintext(
DEVELOPMENT_ENCRYPTION_KEY,
r#"{"provider_type":"antigravity","project_id":"project-antigravity-chat-local-1","client_version":"1.2.3","session_id":"sess-antigravity-chat-local-123","access_token_import_temporary":true,"headers":{"Authorization":"Bearer imported-antigravity-chat-token"}}"#,
)
.expect("auth config should encrypt");
StoredProviderCatalogKey::new(
"key-openai-chat-antigravity-finalize-local-1".to_string(),
"provider-openai-chat-antigravity-finalize-local-1".to_string(),
@@ -1606,12 +1632,9 @@ async fn gateway_skips_openai_chat_antigravity_cross_format_sync_candidate_as_tr
.expect("key should build")
.with_transport_fields(
Some(serde_json::json!(["gemini:generate_content"])),
encrypt_python_fernet_plaintext(
DEVELOPMENT_ENCRYPTION_KEY,
"sk-upstream-openai-chat-antigravity-finalize",
)
.expect("api key should encrypt"),
None,
encrypt_python_fernet_plaintext(DEVELOPMENT_ENCRYPTION_KEY, "__placeholder__")
.expect("api key should encrypt"),
Some(encrypted_auth_config),
None,
Some(serde_json::json!({"gemini:generate_content": 1})),
None,
@@ -1632,6 +1655,8 @@ async fn gateway_skips_openai_chat_antigravity_cross_format_sync_candidate_as_tr
let plan_hits_clone = Arc::clone(&plan_hits);
let public_hits = Arc::new(Mutex::new(0usize));
let public_hits_clone = Arc::clone(&public_hits);
let seen_execution_runtime = Arc::new(Mutex::new(None::<SeenExecutionRuntimeSyncRequest>));
let seen_execution_runtime_clone = Arc::clone(&seen_execution_runtime);
let request_candidate_repository = Arc::new(InMemoryRequestCandidateRepository::default());
let usage_repository = Arc::new(InMemoryUsageReadRepository::default());
@@ -1692,42 +1717,130 @@ async fn gateway_skips_openai_chat_antigravity_cross_format_sync_candidate_as_tr
let execution_runtime = Router::new().route(
"/v1/execute/sync",
any(|_request: Request| async move {
Json(json!({
"request_id": "trace-openai-chat-antigravity-sync-123",
"status_code": 200,
"headers": {
"content-type": "application/json"
},
"body": {
"json_body": {
"response": {
"responseId": "resp-antigravity-chat-sync-123",
"candidates": [{
"content": {
"parts": [
{"text": "Need a tool."},
{"functionCall": {"name": "get_weather", "args": {"city": "SF"}}}
],
"role": "model"
},
"finishReason": "STOP",
"index": 0
}],
"modelVersion": "claude-sonnet-4-5",
"usageMetadata": {
"promptTokenCount": 2,
"candidatesTokenCount": 3,
"totalTokenCount": 5
}
},
"responseId": "resp-antigravity-chat-sync-123"
any(move |request: Request| {
let seen_execution_runtime_inner = Arc::clone(&seen_execution_runtime_clone);
async move {
let (parts, body) = request.into_parts();
let raw_body = to_bytes(body, usize::MAX).await.expect("body should read");
let payload: serde_json::Value =
serde_json::from_slice(&raw_body).expect("execution runtime payload should parse");
*seen_execution_runtime_inner.lock().expect("mutex should lock") =
Some(SeenExecutionRuntimeSyncRequest {
trace_id: parts
.headers
.get(TRACE_ID_HEADER)
.and_then(|value| value.to_str().ok())
.unwrap_or_default()
.to_string(),
url: payload
.get("url")
.and_then(|value| value.as_str())
.unwrap_or_default()
.to_string(),
accept: payload
.get("headers")
.and_then(|value| value.get("accept"))
.and_then(|value| value.as_str())
.unwrap_or_default()
.to_string(),
authorization: payload
.get("headers")
.and_then(|value| value.get("authorization"))
.and_then(|value| value.as_str())
.unwrap_or_default()
.to_string(),
x_client_name: payload
.get("headers")
.and_then(|value| value.get("x-client-name"))
.and_then(|value| value.as_str())
.unwrap_or_default()
.to_string(),
x_client_version: payload
.get("headers")
.and_then(|value| value.get("x-client-version"))
.and_then(|value| value.as_str())
.unwrap_or_default()
.to_string(),
x_vscode_sessionid: payload
.get("headers")
.and_then(|value| value.get("x-vscode-sessionid"))
.and_then(|value| value.as_str())
.unwrap_or_default()
.to_string(),
x_goog_api_client: payload
.get("headers")
.and_then(|value| value.get("x-goog-api-client"))
.and_then(|value| value.as_str())
.unwrap_or_default()
.to_string(),
project: payload
.get("body")
.and_then(|value| value.get("json_body"))
.and_then(|value| value.get("project"))
.and_then(|value| value.as_str())
.unwrap_or_default()
.to_string(),
request_id: payload
.get("body")
.and_then(|value| value.get("json_body"))
.and_then(|value| value.get("requestId"))
.and_then(|value| value.as_str())
.unwrap_or_default()
.to_string(),
model: payload
.get("body")
.and_then(|value| value.get("json_body"))
.and_then(|value| value.get("model"))
.and_then(|value| value.as_str())
.unwrap_or_default()
.to_string(),
user_agent: payload
.get("body")
.and_then(|value| value.get("json_body"))
.and_then(|value| value.get("userAgent"))
.and_then(|value| value.as_str())
.unwrap_or_default()
.to_string(),
request_type: payload
.get("body")
.and_then(|value| value.get("json_body"))
.and_then(|value| value.get("requestType"))
.and_then(|value| value.as_str())
.unwrap_or_default()
.to_string(),
contents_len: payload
.get("body")
.and_then(|value| value.get("json_body"))
.and_then(|value| value.get("request"))
.and_then(|value| value.get("contents"))
.and_then(|value| value.as_array())
.map(Vec::len)
.unwrap_or_default(),
request_has_model: payload
.get("body")
.and_then(|value| value.get("json_body"))
.and_then(|value| value.get("request"))
.and_then(|value| value.get("model"))
.is_some(),
});
Json(json!({
"request_id": "trace-openai-chat-antigravity-sync-123",
"status_code": 200,
"headers": {
"content-type": "text/event-stream"
},
"body": {
"body_bytes_b64": base64::engine::general_purpose::STANDARD.encode(
concat!(
"data: {\"response\":{\"candidates\":[{\"content\":{\"parts\":[{\"text\":\"Hello Antigravity Chat\"}],\"role\":\"model\"},\"finishReason\":\"STOP\",\"index\":0}],\"modelVersion\":\"claude-sonnet-4-5\",\"usageMetadata\":{\"promptTokenCount\":2,\"candidatesTokenCount\":3,\"totalTokenCount\":5}},\"responseId\":\"resp-antigravity-chat-sync-123\"}\n\n"
)
)
},
"telemetry": {
"elapsed_ms": 29
}
},
"telemetry": {
"elapsed_ms": 29
}
}))
}))
}
}),
);
@@ -1779,31 +1892,84 @@ async fn gateway_skips_openai_chat_antigravity_cross_format_sync_candidate_as_tr
.await
.expect("request should succeed");
let elapsed = started_at.elapsed();
let response_status = response.status();
let execution_path = response
.headers()
.get(EXECUTION_PATH_HEADER)
.and_then(|value| value.to_str().ok())
.unwrap_or_default()
.to_string();
let miss_reason = response
.headers()
.get(LOCAL_EXECUTION_RUNTIME_MISS_REASON_HEADER)
.and_then(|value| value.to_str().ok())
.unwrap_or_default()
.to_string();
let response_body = response.text().await.expect("body should read");
assert_eq!(response_status, StatusCode::OK);
let response_json: serde_json::Value =
serde_json::from_str(&response_body).expect("body should parse");
assert_eq!(
response_status,
StatusCode::SERVICE_UNAVAILABLE,
"unexpected gateway response: path={execution_path} miss={miss_reason} body={response_body}"
response_json,
json!({
"id": "resp-local-stream",
"object": "chat.completion",
"model": "claude-sonnet-4-5",
"choices": [{
"index": 0,
"message": {
"role": "assistant",
"content": "Hello Antigravity Chat"
},
"finish_reason": "stop"
}],
"usage": {
"prompt_tokens": 2,
"completion_tokens": 3,
"total_tokens": 5
}
})
);
assert!(
elapsed < std::time::Duration::from_millis(3_500),
"response took unexpectedly long for local finalize path: elapsed={elapsed:?} finalize_hits={} report_hits={}",
*finalize_hits.lock().expect("mutex should lock"),
*report_hits.lock().expect("mutex should lock"),
);
let seen_execution_runtime_request = seen_execution_runtime
.lock()
.expect("mutex should lock")
.clone()
.expect("execution runtime request should be captured");
assert_eq!(
seen_execution_runtime_request.trace_id,
"trace-openai-chat-antigravity-sync-123"
);
assert_eq!(
execution_path, "local_execution_runtime_miss",
"unexpected execution path: miss={miss_reason} body={response_body}"
seen_execution_runtime_request.url,
"https://antigravity.googleapis.com/v1internal:streamGenerateContent?alt=sse"
);
assert_eq!(miss_reason, "all_candidates_skipped");
assert_eq!(seen_execution_runtime_request.accept, "text/event-stream");
assert_eq!(
seen_execution_runtime_request.authorization,
"Bearer imported-antigravity-chat-token"
);
assert_eq!(seen_execution_runtime_request.x_client_name, "antigravity");
assert_eq!(seen_execution_runtime_request.x_client_version, "1.2.3");
assert_eq!(
seen_execution_runtime_request.x_vscode_sessionid,
"sess-antigravity-chat-local-123"
);
assert_eq!(
seen_execution_runtime_request.x_goog_api_client,
"gl-node/18.18.2 fire/0.8.6 grpc/1.10.x"
);
assert_eq!(
seen_execution_runtime_request.project,
"project-antigravity-chat-local-1"
);
assert_eq!(
seen_execution_runtime_request.request_id,
"trace-openai-chat-antigravity-sync-123"
);
assert_eq!(seen_execution_runtime_request.model, "claude-sonnet-4-5");
assert_eq!(
seen_execution_runtime_request.user_agent,
aether_provider_transport::antigravity::ANTIGRAVITY_REQUEST_USER_AGENT
);
assert_eq!(seen_execution_runtime_request.request_type, "agent");
assert_eq!(seen_execution_runtime_request.contents_len, 1);
assert!(!seen_execution_runtime_request.request_has_model);
let mut stored_candidates = Vec::new();
for _ in 0..50 {
@@ -1819,11 +1985,8 @@ async fn gateway_skips_openai_chat_antigravity_cross_format_sync_candidate_as_tr
tokio::time::sleep(std::time::Duration::from_millis(10)).await;
}
assert_eq!(stored_candidates.len(), 1);
assert_eq!(stored_candidates[0].status, RequestCandidateStatus::Skipped);
assert_eq!(
stored_candidates[0].skip_reason.as_deref(),
Some("transport_provider_type_unsupported")
);
assert_eq!(stored_candidates[0].status, RequestCandidateStatus::Success);
assert_eq!(stored_candidates[0].skip_reason.as_deref(), None);
tokio::time::sleep(std::time::Duration::from_millis(100)).await;
assert_eq!(*report_hits.lock().expect("mutex should lock"), 0);
assert_eq!(*finalize_hits.lock().expect("mutex should lock"), 0);
@@ -1520,7 +1520,10 @@ async fn gateway_executes_openai_responses_antigravity_cross_format_upstream_str
seen_remote_execution_runtime_request.url,
"https://antigravity.googleapis.com/v1internal:streamGenerateContent?alt=sse"
);
assert_eq!(seen_remote_execution_runtime_request.accept, "*/*");
assert_eq!(
seen_remote_execution_runtime_request.accept,
"text/event-stream"
);
assert_eq!(
seen_remote_execution_runtime_request.authorization,
"Bearer refreshed-antigravity-cli-access-token"
@@ -1555,7 +1558,7 @@ async fn gateway_executes_openai_responses_antigravity_cross_format_upstream_str
);
assert_eq!(
seen_remote_execution_runtime_request.user_agent,
"antigravity"
aether_provider_transport::antigravity::ANTIGRAVITY_REQUEST_USER_AGENT
);
assert_eq!(seen_remote_execution_runtime_request.request_type, "agent");
assert_eq!(seen_remote_execution_runtime_request.contents_len, 1);
@@ -2143,7 +2143,7 @@ async fn gateway_executes_antigravity_gemini_cli_sync_upstream_stream_via_local_
);
assert_eq!(
seen_remote_execution_runtime_request.user_agent,
"antigravity"
aether_provider_transport::antigravity::ANTIGRAVITY_REQUEST_USER_AGENT
);
assert_eq!(seen_remote_execution_runtime_request.request_type, "agent");
assert_eq!(seen_remote_execution_runtime_request.contents_len, 0);
@@ -2031,7 +2031,10 @@ async fn gateway_executes_antigravity_gemini_cli_stream_via_local_decision_gate_
"trace-antigravity-cli-oauth-local-stream-123"
);
assert_eq!(seen_execution_runtime_request.model, "claude-sonnet-4-5");
assert_eq!(seen_execution_runtime_request.user_agent, "antigravity");
assert_eq!(
seen_execution_runtime_request.user_agent,
aether_provider_transport::antigravity::ANTIGRAVITY_REQUEST_USER_AGENT
);
assert_eq!(seen_execution_runtime_request.request_type, "agent");
assert_eq!(seen_execution_runtime_request.contents_len, 0);
assert!((seen_execution_runtime_request.exact_temperature - 0.2).abs() < f64::EPSILON);
@@ -2334,7 +2334,10 @@ async fn gateway_executes_antigravity_gemini_cli_sync_via_local_decision_gate_af
"trace-antigravity-cli-oauth-local-sync-123"
);
assert_eq!(seen_execution_runtime_request.model, "claude-sonnet-4-5");
assert_eq!(seen_execution_runtime_request.user_agent, "antigravity");
assert_eq!(
seen_execution_runtime_request.user_agent,
aether_provider_transport::antigravity::ANTIGRAVITY_REQUEST_USER_AGENT
);
assert_eq!(seen_execution_runtime_request.request_type, "agent");
assert_eq!(seen_execution_runtime_request.contents_len, 0);
assert!((seen_execution_runtime_request.exact_temperature - 0.2).abs() < f64::EPSILON);
@@ -23,6 +23,30 @@ use crate::constants::{
};
use crate::data::GatewayDataState;
const PROVIDER_QUOTA_TEST_STACK_BYTES: usize = 16 * 1024 * 1024;
fn run_provider_quota_test<F, Fut>(test_name: &'static str, make_future: F)
where
F: FnOnce() -> Fut + Send + 'static,
Fut: std::future::Future<Output = ()> + 'static,
{
let handle = std::thread::Builder::new()
.name(test_name.to_string())
.stack_size(PROVIDER_QUOTA_TEST_STACK_BYTES)
.spawn(move || {
let runtime = tokio::runtime::Builder::new_current_thread()
.enable_all()
.build()
.expect("test runtime should build");
runtime.block_on(make_future());
})
.expect("provider quota test thread should spawn");
if let Err(payload) = handle.join() {
std::panic::resume_unwind(payload);
}
}
#[tokio::test]
async fn gateway_refreshes_admin_provider_quota_locally_for_codex_with_trusted_admin_principal() {
#[derive(Debug, Clone)]
@@ -1748,8 +1772,15 @@ async fn gateway_reports_codex_quota_runtime_failures_locally_without_falling_ba
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_refreshes_admin_provider_quota_locally_for_antigravity_with_trusted_admin_principal(
#[test]
fn gateway_refreshes_admin_provider_quota_locally_for_antigravity_with_trusted_admin_principal() {
run_provider_quota_test(
"gateway_refreshes_admin_provider_quota_locally_for_antigravity_with_trusted_admin_principal",
gateway_refreshes_admin_provider_quota_locally_for_antigravity_with_trusted_admin_principal_inner,
);
}
async fn gateway_refreshes_admin_provider_quota_locally_for_antigravity_with_trusted_admin_principal_inner(
) {
#[derive(Debug, Clone)]
struct SeenExecutionRuntimeRequest {
@@ -4125,7 +4125,7 @@ async fn gateway_reports_transport_unsupported_reason_for_non_kiro_provider_impl
assert_eq!(
payload["error"],
json!(
"Rust local provider-query model test cannot execute endpoint format gemini:generate_content (transport_antigravity_auth_unsupported)"
"Rust local provider-query model test cannot execute endpoint format gemini:generate_content (transport_antigravity_auth_config_missing)"
)
);
@@ -4305,6 +4305,204 @@ async fn gateway_handles_antigravity_endpoint_test_model_locally_impl() {
execution_runtime_handle.abort();
}
#[test]
fn gateway_hydrates_antigravity_project_id_from_load_code_assist_for_test_model() {
run_provider_query_test(
"gateway_hydrates_antigravity_project_id_from_load_code_assist_for_test_model",
gateway_hydrates_antigravity_project_id_from_load_code_assist_for_test_model_impl,
);
}
async fn gateway_hydrates_antigravity_project_id_from_load_code_assist_for_test_model_impl() {
let seen_urls = Arc::new(Mutex::new(Vec::<String>::new()));
let seen_urls_clone = Arc::clone(&seen_urls);
let execution_runtime = Router::new().route(
"/v1/execute/sync",
any(move |Json(plan): Json<ExecutionPlan>| {
let seen_urls_inner = Arc::clone(&seen_urls_clone);
async move {
seen_urls_inner
.lock()
.expect("mutex should lock")
.push(plan.url.clone());
if plan.url == "https://daily-cloudcode-pa.googleapis.com/v1internal:loadCodeAssist"
{
assert_eq!(plan.model_name.as_deref(), Some("loadCodeAssist"));
assert_eq!(
plan.headers.get("authorization").map(String::as_str),
Some("Bearer cached-antigravity-token")
);
assert_eq!(
plan.body.json_body.as_ref().and_then(|body| body
.get("metadata")
.and_then(|metadata| metadata.get("pluginType"))),
Some(&json!("GEMINI"))
);
return Json(json!({
"request_id": plan.request_id,
"candidate_id": plan.candidate_id,
"status_code": 200,
"headers": {
"content-type": "application/json"
},
"body": {
"json_body": {
"cloudaicompanionProject": {
"id": "project-from-antigravity-load-code-assist"
},
"currentTier": {
"id": "free"
}
}
}
}));
}
assert_eq!(
plan.url,
"https://daily-cloudcode-pa.googleapis.com/v1internal:generateContent"
);
assert_eq!(plan.provider_id, "provider-antigravity");
assert_eq!(plan.endpoint_id, "endpoint-antigravity-gemini");
assert_eq!(plan.key_id, "key-antigravity-gemini");
assert_eq!(plan.provider_api_format, "gemini:generate_content");
assert!(!plan.stream);
assert_eq!(
plan.body.json_body.as_ref().unwrap()["project"],
json!("project-from-antigravity-load-code-assist")
);
assert_eq!(
plan.body.json_body.as_ref().unwrap()["requestType"],
json!("endpoint_test")
);
assert_eq!(
plan.body.json_body.as_ref().unwrap()["model"],
json!("gemini-3.1-flash-lite")
);
Json(json!({
"request_id": plan.request_id,
"candidate_id": plan.candidate_id,
"status_code": 200,
"headers": {
"content-type": "application/json"
},
"body": {
"json_body": {
"response": {
"candidates": [{
"content": {
"parts": [
{"text": "Hello from hydrated Antigravity"}
],
"role": "model"
},
"finishReason": "STOP",
"index": 0
}],
"modelVersion": "gemini-3.1-flash-lite",
"usageMetadata": {
"promptTokenCount": 2,
"candidatesTokenCount": 4,
"totalTokenCount": 6
}
},
"responseId": "resp-antigravity-hydrated-test-123"
}
},
"telemetry": {
"elapsed_ms": 23
}
}))
}
}),
);
let (execution_runtime_url, execution_runtime_handle) = start_server(execution_runtime).await;
let mut provider = sample_provider("provider-antigravity", "Antigravity", 10);
provider.provider_type = "antigravity".to_string();
let mut key = sample_key(
"key-antigravity-gemini",
"provider-antigravity",
"gemini:generate_content",
"cached-antigravity-token",
);
key.auth_type = "oauth".to_string();
key.encrypted_auth_config = Some(
aether_crypto::encrypt_python_fernet_plaintext(
DEVELOPMENT_ENCRYPTION_KEY,
r#"{"provider_type":"antigravity","refresh_token":"rt-antigravity-123"}"#,
)
.expect("auth config should encrypt"),
);
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![provider],
vec![sample_endpoint(
"endpoint-antigravity-gemini",
"provider-antigravity",
"gemini:generate_content",
"https://daily-cloudcode-pa.googleapis.com",
)],
vec![key],
));
let gateway = build_router_with_state(
build_state_with_execution_runtime_override(execution_runtime_url)
.with_data_state_for_tests(
GatewayDataState::with_provider_catalog_repository_for_tests(Arc::clone(
&provider_catalog_repository,
))
.with_encryption_key_for_tests(DEVELOPMENT_ENCRYPTION_KEY),
),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.post(format!("{gateway_url}/api/admin/provider-query/test-model"))
.header(GATEWAY_HEADER, "rust-phase3b")
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
.json(&json!({
"provider_id": "provider-antigravity",
"model": "gemini-3.1-flash-lite",
"api_format": "gemini:generate_content",
"message": "Say hello"
}))
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::OK);
let payload: serde_json::Value = response.json().await.expect("json body should parse");
assert_eq!(payload["success"], json!(true));
assert_eq!(
payload["data"]["response"]["choices"][0]["message"]["content"],
json!("Hello from hydrated Antigravity")
);
assert_eq!(
*seen_urls.lock().expect("mutex should lock"),
vec![
"https://daily-cloudcode-pa.googleapis.com/v1internal:loadCodeAssist".to_string(),
"https://daily-cloudcode-pa.googleapis.com/v1internal:generateContent".to_string(),
]
);
let reloaded = provider_catalog_repository
.list_keys_by_ids(&["key-antigravity-gemini".to_string()])
.await
.expect("key should reload");
assert_eq!(
reloaded[0]
.upstream_metadata
.as_ref()
.and_then(|metadata| metadata.get("antigravity"))
.and_then(|metadata| metadata.get("project_id")),
Some(&json!("project-from-antigravity-load-code-assist"))
);
gateway_handle.abort();
execution_runtime_handle.abort();
}
#[test]
fn gateway_prefers_supported_non_kiro_endpoint_when_api_format_is_omitted() {
run_provider_query_test(
@@ -1452,6 +1452,11 @@ async fn gateway_handles_admin_system_api_formats_locally_with_trusted_admin_pri
assert!(formats.iter().any(|item| item["value"] == "jina:embedding"));
assert!(formats.iter().any(|item| item["value"] == "jina:rerank"));
assert!(formats.iter().any(|item| item["value"] == "gemini:video"));
let gemini_interactions = formats
.iter()
.find(|item| item["value"] == "gemini:interactions")
.expect("gemini interactions format should exist");
assert_eq!(gemini_interactions["default_path"], "/v1/interactions");
let aliyun_embedding = formats
.iter()
.find(|item| item["value"] == "aliyun:multimodal_embedding")
+73 -6
View File
@@ -949,6 +949,10 @@ async fn gateway_handles_antigravity_v1internal_control_plane_without_proxying()
"/v1internal:fetchAvailableModels",
json!({"project": "aether-antigravity-local"}),
),
(
"/v1internal:retrieveUserQuotaSummary",
json!({"project": "aether-antigravity-local"}),
),
(
"/v1internal:fetchUserInfo",
json!({"project": "aether-antigravity-local"}),
@@ -966,6 +970,10 @@ async fn gateway_handles_antigravity_v1internal_control_plane_without_proxying()
"metrics": []
}),
),
(
"/v1internal:writeTrajectoryAcls",
json!({"trajectoryId": "trajectory-ant-123"}),
),
(
"/v1internal:setUserSettings",
json!({"userSettings": user_settings.clone()}),
@@ -1012,14 +1020,58 @@ async fn gateway_handles_antigravity_v1internal_control_plane_without_proxying()
);
}
"/v1internal:fetchAvailableModels" => {
assert_eq!(payload["defaultAgentModelId"], "gemini-3.1-flash-lite");
assert_eq!(payload["defaultAgentModelId"], "gemini-3.5-flash-low");
assert_eq!(
payload["tieredModelIds"]["flash"],
json!(["gemini-3-flash-agent"])
);
assert_eq!(
payload["tieredModelIds"]["pro"],
json!(["gemini-3.1-pro-low"])
);
assert_eq!(
payload["models"]["gemini-3-flash-agent"]["displayName"],
"Gemini 3.5 Flash (High)"
);
assert_eq!(
payload["models"]["gemini-3.5-flash-low"]["displayName"],
"Gemini 3.5 Flash Low"
"Gemini 3.5 Flash (Medium)"
);
assert_eq!(
payload["models"]["gemini-3.5-flash-extra-low"]["displayName"],
"Gemini 3.5 Flash (Low)"
);
assert_eq!(
payload["models"]["gemini-pro-agent"]["displayName"],
"Gemini 3.1 Pro (High)"
);
assert_eq!(
payload["models"]["claude-opus-4-6-thinking"]["displayName"],
"Claude Opus 4.6 (Thinking)"
);
assert_eq!(
payload["models"]["gpt-oss-120b-medium"]["displayName"],
"GPT-OSS 120B (Medium)"
);
assert_eq!(
payload["models"]["gemini-pro-agent"]["model"],
"MODEL_PLACEHOLDER_M16"
);
assert_eq!(
payload["models"]["gemini-3.1-pro-high"]["model"],
"MODEL_PLACEHOLDER_M37"
);
assert_eq!(
payload["models"]["gemini-3.5-flash-extra-low"]["model"],
"MODEL_PLACEHOLDER_M187"
);
assert_eq!(
payload["models"]["claude-sonnet-4-6"]["apiProvider"],
"API_PROVIDER_ANTHROPIC_VERTEX"
);
assert_eq!(
payload["models"]["gpt-oss-120b-medium"]["apiProvider"],
"API_PROVIDER_OPENAI_VERTEX"
);
assert_eq!(
payload["models"]["gemini-3.5-flash-low"]["apiProvider"],
@@ -1032,18 +1084,26 @@ async fn gateway_handles_antigravity_v1internal_control_plane_without_proxying()
assert_eq!(
payload["agentModelSorts"][0]["groups"][0]["modelIds"],
json!([
"gemini-3.1-flash-lite",
"gemini-3.5-flash-low",
"gemini-3-flash-agent",
"gemini-3.5-flash-extra-low",
"gemini-3.1-pro-low",
"gemini-3.5-flash-low"
"gemini-pro-agent",
"claude-sonnet-4-6",
"claude-opus-4-6-thinking",
"gpt-oss-120b-medium"
])
);
assert_eq!(payload["deprecatedModelIds"], json!({}));
assert_eq!(
payload["deprecatedModelIds"]["gemini-3.1-pro-high"]["newModelId"],
"gemini-pro-agent"
);
assert_eq!(payload["commandModelIds"], json!(["gemini-3-flash"]));
assert_eq!(
payload["imageGenerationModelIds"],
json!(["gemini-3.1-flash-image"])
);
assert_eq!(payload["tabModelIds"], json!(["chat_20706", "chat_23310"]));
assert_eq!(payload["mqueryModelIds"], json!(["gemini-3.1-flash-lite"]));
assert_eq!(
payload["webSearchModelIds"],
@@ -1058,9 +1118,13 @@ async fn gateway_handles_antigravity_v1internal_control_plane_without_proxying()
assert_eq!(payload["regionCode"], "US");
assert_eq!(
payload["userSettings"]["preferredModelId"],
"gemini-3.1-flash-lite"
"gemini-3.5-flash-low"
);
}
"/v1internal:retrieveUserQuotaSummary" => {
assert_eq!(payload["description"], "");
assert_eq!(payload["groups"], json!([]));
}
"/v1internal:fetchAdminControls" => {
assert_eq!(payload, json!({}));
}
@@ -1071,6 +1135,9 @@ async fn gateway_handles_antigravity_v1internal_control_plane_without_proxying()
"/v1internal:recordCodeAssistMetrics" => {
assert_eq!(payload, json!({}));
}
"/v1internal:writeTrajectoryAcls" => {
assert_eq!(payload, json!({}));
}
"/v1internal:setUserSettings" => {
assert_eq!(payload["userSettings"], user_settings);
}