mirror of
https://github.com/fawney19/Aether.git
synced 2026-10-11 11:49:50 +08:00
refactor(workspace): enforce layered crate boundaries
This commit is contained in:
elky
@@ -0,0 +1,406 @@
|
||||
use std::collections::BTreeMap;
|
||||
|
||||
use serde_json::Value;
|
||||
|
||||
use super::super::snapshot::GatewayProviderTransportSnapshot;
|
||||
|
||||
pub const ANTIGRAVITY_PROVIDER_TYPE: &str = "antigravity";
|
||||
pub const ANTIGRAVITY_REQUEST_USER_AGENT: &str =
|
||||
"antigravity/cli/1.0.16 (aidev_client; os_type=linux; arch=arm64; auth_method=consumer)";
|
||||
const ANTIGRAVITY_CLIENT_NAME: &str = "antigravity";
|
||||
const ANTIGRAVITY_GOOG_API_CLIENT: &str = "gl-node/18.18.2 fire/0.8.6 grpc/1.10.x";
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct AntigravityRequestAuth {
|
||||
pub project_id: String,
|
||||
pub client_version: Option<String>,
|
||||
pub session_id: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub enum AntigravityRequestAuthSupport {
|
||||
Supported(AntigravityRequestAuth),
|
||||
Unsupported(AntigravityRequestAuthUnsupportedReason),
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub enum AntigravityRequestAuthUnsupportedReason {
|
||||
WrongProviderType,
|
||||
MissingAuthConfig,
|
||||
InvalidAuthConfigJson,
|
||||
ComplexDynamicAuthConfig,
|
||||
MissingProjectId,
|
||||
}
|
||||
|
||||
pub fn resolve_local_antigravity_request_auth(
|
||||
transport: &GatewayProviderTransportSnapshot,
|
||||
) -> AntigravityRequestAuthSupport {
|
||||
if !transport
|
||||
.provider
|
||||
.provider_type
|
||||
.trim()
|
||||
.eq_ignore_ascii_case(ANTIGRAVITY_PROVIDER_TYPE)
|
||||
{
|
||||
return AntigravityRequestAuthSupport::Unsupported(
|
||||
AntigravityRequestAuthUnsupportedReason::WrongProviderType,
|
||||
);
|
||||
}
|
||||
|
||||
let Some(raw_auth_config) = transport
|
||||
.key
|
||||
.decrypted_auth_config
|
||||
.as_deref()
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
else {
|
||||
return AntigravityRequestAuthSupport::Unsupported(
|
||||
AntigravityRequestAuthUnsupportedReason::MissingAuthConfig,
|
||||
);
|
||||
};
|
||||
|
||||
let Ok(auth_config) = serde_json::from_str::<Value>(raw_auth_config) else {
|
||||
return AntigravityRequestAuthSupport::Unsupported(
|
||||
AntigravityRequestAuthUnsupportedReason::InvalidAuthConfigJson,
|
||||
);
|
||||
};
|
||||
|
||||
if contains_blocked_auth_fields(&auth_config) {
|
||||
return AntigravityRequestAuthSupport::Unsupported(
|
||||
AntigravityRequestAuthUnsupportedReason::ComplexDynamicAuthConfig,
|
||||
);
|
||||
}
|
||||
|
||||
let upstream_metadata = transport.key.upstream_metadata.as_ref();
|
||||
let Some(project_id) = find_antigravity_string(
|
||||
upstream_metadata,
|
||||
&auth_config,
|
||||
ANTIGRAVITY_PROJECT_ID_PATHS,
|
||||
) else {
|
||||
return AntigravityRequestAuthSupport::Unsupported(
|
||||
AntigravityRequestAuthUnsupportedReason::MissingProjectId,
|
||||
);
|
||||
};
|
||||
|
||||
let client_version = find_antigravity_string(
|
||||
upstream_metadata,
|
||||
&auth_config,
|
||||
ANTIGRAVITY_CLIENT_VERSION_PATHS,
|
||||
);
|
||||
let session_id = find_antigravity_string(
|
||||
upstream_metadata,
|
||||
&auth_config,
|
||||
ANTIGRAVITY_SESSION_ID_PATHS,
|
||||
);
|
||||
|
||||
AntigravityRequestAuthSupport::Supported(AntigravityRequestAuth {
|
||||
project_id,
|
||||
client_version,
|
||||
session_id,
|
||||
})
|
||||
}
|
||||
|
||||
pub fn build_antigravity_static_identity_headers(
|
||||
auth: &AntigravityRequestAuth,
|
||||
) -> BTreeMap<String, String> {
|
||||
build_antigravity_static_client_headers(
|
||||
auth.client_version.as_deref(),
|
||||
auth.session_id.as_deref(),
|
||||
)
|
||||
}
|
||||
|
||||
pub fn build_antigravity_static_client_headers(
|
||||
client_version: Option<&str>,
|
||||
session_id: Option<&str>,
|
||||
) -> BTreeMap<String, String> {
|
||||
let mut headers = BTreeMap::from([
|
||||
(
|
||||
String::from("x-client-name"),
|
||||
String::from(ANTIGRAVITY_CLIENT_NAME),
|
||||
),
|
||||
(
|
||||
String::from("x-goog-api-client"),
|
||||
String::from(ANTIGRAVITY_GOOG_API_CLIENT),
|
||||
),
|
||||
(
|
||||
String::from("user-agent"),
|
||||
String::from(ANTIGRAVITY_REQUEST_USER_AGENT),
|
||||
),
|
||||
]);
|
||||
|
||||
if let Some(client_version) = client_version
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
{
|
||||
headers.insert(String::from("x-client-version"), client_version.to_string());
|
||||
}
|
||||
if let Some(session_id) = session_id.map(str::trim).filter(|value| !value.is_empty()) {
|
||||
headers.insert(String::from("x-vscode-sessionid"), session_id.to_string());
|
||||
}
|
||||
|
||||
headers
|
||||
}
|
||||
|
||||
const ANTIGRAVITY_PROJECT_ID_PATHS: &[&[&str]] = &[
|
||||
&["project_id"],
|
||||
&["projectId"],
|
||||
&["project", "id"],
|
||||
&["project", "project_id"],
|
||||
&["project", "projectId"],
|
||||
&["cloudaicompanionProject"],
|
||||
&["cloudaicompanionProject", "id"],
|
||||
&["cloudAiCompanionProject"],
|
||||
&["cloudAiCompanionProject", "id"],
|
||||
&["antigravity", "project_id"],
|
||||
&["antigravity", "projectId"],
|
||||
&["antigravity", "project", "id"],
|
||||
&["antigravity", "cloudaicompanionProject"],
|
||||
&["antigravity", "cloudaicompanionProject", "id"],
|
||||
&["antigravity", "cloudAiCompanionProject"],
|
||||
&["antigravity", "cloudAiCompanionProject", "id"],
|
||||
&["metadata", "project_id"],
|
||||
&["metadata", "projectId"],
|
||||
&["metadata", "cloudaicompanionProject"],
|
||||
&["metadata", "cloudaicompanionProject", "id"],
|
||||
&["metadata", "cloudAiCompanionProject"],
|
||||
&["metadata", "cloudAiCompanionProject", "id"],
|
||||
];
|
||||
|
||||
const ANTIGRAVITY_CLIENT_VERSION_PATHS: &[&[&str]] = &[
|
||||
&["client_version"],
|
||||
&["clientVersion"],
|
||||
&["antigravity", "client_version"],
|
||||
&["antigravity", "clientVersion"],
|
||||
&["metadata", "client_version"],
|
||||
&["metadata", "clientVersion"],
|
||||
];
|
||||
|
||||
const ANTIGRAVITY_SESSION_ID_PATHS: &[&[&str]] = &[
|
||||
&["session_id"],
|
||||
&["sessionId"],
|
||||
&["antigravity", "session_id"],
|
||||
&["antigravity", "sessionId"],
|
||||
&["metadata", "session_id"],
|
||||
&["metadata", "sessionId"],
|
||||
];
|
||||
|
||||
fn find_antigravity_string(
|
||||
upstream_metadata: Option<&Value>,
|
||||
auth_config: &Value,
|
||||
paths: &[&[&str]],
|
||||
) -> Option<String> {
|
||||
upstream_metadata
|
||||
.and_then(|metadata| find_string_by_paths(metadata, paths))
|
||||
.or_else(|| find_string_by_paths(auth_config, paths))
|
||||
}
|
||||
|
||||
fn find_string_by_paths(value: &Value, paths: &[&[&str]]) -> Option<String> {
|
||||
for path in paths {
|
||||
let mut current = value;
|
||||
let mut matched = true;
|
||||
for segment in *path {
|
||||
let Some(next) = current.get(*segment) else {
|
||||
matched = false;
|
||||
break;
|
||||
};
|
||||
current = next;
|
||||
}
|
||||
if !matched {
|
||||
continue;
|
||||
}
|
||||
if let Some(string) = current
|
||||
.as_str()
|
||||
.map(str::trim)
|
||||
.filter(|item| !item.is_empty())
|
||||
{
|
||||
return Some(string.to_string());
|
||||
}
|
||||
if let Some(string) = current
|
||||
.as_object()
|
||||
.and_then(|object| {
|
||||
object
|
||||
.get("id")
|
||||
.or_else(|| object.get("project_id"))
|
||||
.or_else(|| object.get("projectId"))
|
||||
})
|
||||
.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
.filter(|item| !item.is_empty())
|
||||
{
|
||||
return Some(string.to_string());
|
||||
}
|
||||
}
|
||||
|
||||
None
|
||||
}
|
||||
|
||||
fn contains_blocked_auth_fields(value: &Value) -> bool {
|
||||
match value {
|
||||
Value::Object(map) => map.iter().any(|(key, inner)| {
|
||||
is_blocked_auth_key(key.as_str()) || contains_blocked_auth_fields(inner)
|
||||
}),
|
||||
Value::Array(items) => items.iter().any(contains_blocked_auth_fields),
|
||||
_ => false,
|
||||
}
|
||||
}
|
||||
|
||||
fn is_blocked_auth_key(key: &str) -> bool {
|
||||
matches!(
|
||||
key.trim().to_ascii_lowercase().as_str(),
|
||||
"private_key"
|
||||
| "privateKey"
|
||||
| "private_key_id"
|
||||
| "privateKeyId"
|
||||
| "service_account"
|
||||
| "serviceAccount"
|
||||
| "service_account_json"
|
||||
| "serviceAccountJson"
|
||||
| "service_account_key"
|
||||
| "serviceAccountKey"
|
||||
| "credential_source"
|
||||
| "credentialSource"
|
||||
| "token_url"
|
||||
| "tokenUrl"
|
||||
| "auth_uri"
|
||||
| "authUri"
|
||||
| "subject"
|
||||
| "audience"
|
||||
)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use serde_json::json;
|
||||
|
||||
use super::{
|
||||
build_antigravity_static_client_headers, resolve_local_antigravity_request_auth,
|
||||
AntigravityRequestAuth, AntigravityRequestAuthSupport, ANTIGRAVITY_REQUEST_USER_AGENT,
|
||||
};
|
||||
use crate::snapshot::{
|
||||
GatewayProviderTransportEndpoint, GatewayProviderTransportKey,
|
||||
GatewayProviderTransportProvider, GatewayProviderTransportSnapshot,
|
||||
};
|
||||
|
||||
fn sample_transport(auth_config: &str) -> GatewayProviderTransportSnapshot {
|
||||
GatewayProviderTransportSnapshot {
|
||||
provider: GatewayProviderTransportProvider {
|
||||
id: "provider-1".to_string(),
|
||||
name: "Antigravity".to_string(),
|
||||
provider_type: "antigravity".to_string(),
|
||||
website: None,
|
||||
is_active: true,
|
||||
keep_priority_on_conversion: false,
|
||||
enable_format_conversion: true,
|
||||
concurrent_limit: None,
|
||||
max_retries: None,
|
||||
proxy: None,
|
||||
request_timeout_secs: None,
|
||||
stream_first_byte_timeout_secs: None,
|
||||
config: None,
|
||||
},
|
||||
endpoint: GatewayProviderTransportEndpoint {
|
||||
id: "endpoint-1".to_string(),
|
||||
provider_id: "provider-1".to_string(),
|
||||
api_format: "gemini:generate_content".to_string(),
|
||||
api_family: Some("gemini".to_string()),
|
||||
endpoint_kind: Some("generate_content".to_string()),
|
||||
is_active: true,
|
||||
base_url: "https://daily-cloudcode-pa.googleapis.com".to_string(),
|
||||
header_rules: None,
|
||||
body_rules: None,
|
||||
max_retries: None,
|
||||
custom_path: None,
|
||||
config: None,
|
||||
format_acceptance_config: None,
|
||||
proxy: None,
|
||||
},
|
||||
key: GatewayProviderTransportKey {
|
||||
id: "key-1".to_string(),
|
||||
provider_id: "provider-1".to_string(),
|
||||
name: "key".to_string(),
|
||||
auth_type: "oauth".to_string(),
|
||||
is_active: true,
|
||||
api_formats: Some(vec!["gemini:generate_content".to_string()]),
|
||||
auth_type_by_format: None,
|
||||
allow_auth_channel_mismatch_formats: None,
|
||||
allowed_models: None,
|
||||
capabilities: None,
|
||||
rate_multipliers: None,
|
||||
global_priority_by_format: None,
|
||||
expires_at_unix_secs: None,
|
||||
proxy: None,
|
||||
fingerprint: None,
|
||||
upstream_metadata: None,
|
||||
decrypted_api_key: "__placeholder__".to_string(),
|
||||
decrypted_auth_config: Some(auth_config.to_string()),
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn resolves_cloudaicompanion_project_object_from_auth_config() {
|
||||
let transport = sample_transport(
|
||||
r#"{
|
||||
"provider_type":"antigravity",
|
||||
"refresh_token":"rt",
|
||||
"cloudaicompanionProject":{"id":"project-from-auth-config"}
|
||||
}"#,
|
||||
);
|
||||
|
||||
assert_eq!(
|
||||
resolve_local_antigravity_request_auth(&transport),
|
||||
AntigravityRequestAuthSupport::Supported(AntigravityRequestAuth {
|
||||
project_id: "project-from-auth-config".to_string(),
|
||||
client_version: None,
|
||||
session_id: None,
|
||||
})
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn resolves_identity_from_antigravity_upstream_metadata() {
|
||||
let mut transport = sample_transport(
|
||||
r#"{
|
||||
"provider_type":"antigravity",
|
||||
"refresh_token":"rt"
|
||||
}"#,
|
||||
);
|
||||
transport.key.upstream_metadata = Some(json!({
|
||||
"antigravity": {
|
||||
"project_id": "project-from-metadata",
|
||||
"client_version": "1.99.0",
|
||||
"session_id": "session-from-metadata"
|
||||
}
|
||||
}));
|
||||
|
||||
assert_eq!(
|
||||
resolve_local_antigravity_request_auth(&transport),
|
||||
AntigravityRequestAuthSupport::Supported(AntigravityRequestAuth {
|
||||
project_id: "project-from-metadata".to_string(),
|
||||
client_version: Some("1.99.0".to_string()),
|
||||
session_id: Some("session-from-metadata".to_string()),
|
||||
})
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn static_client_headers_use_native_antigravity_cli_user_agent() {
|
||||
let headers = build_antigravity_static_client_headers(Some("1.0.16"), Some("session-abc"));
|
||||
|
||||
assert_eq!(
|
||||
headers.get("user-agent").map(String::as_str),
|
||||
Some(ANTIGRAVITY_REQUEST_USER_AGENT)
|
||||
);
|
||||
assert_eq!(
|
||||
headers.get("x-client-name").map(String::as_str),
|
||||
Some("antigravity")
|
||||
);
|
||||
assert_eq!(
|
||||
headers.get("x-client-version").map(String::as_str),
|
||||
Some("1.0.16")
|
||||
);
|
||||
assert_eq!(
|
||||
headers.get("x-vscode-sessionid").map(String::as_str),
|
||||
Some("session-abc")
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,25 @@
|
||||
mod auth;
|
||||
mod policy;
|
||||
mod request;
|
||||
mod url;
|
||||
|
||||
pub use auth::{
|
||||
build_antigravity_static_client_headers, build_antigravity_static_identity_headers,
|
||||
resolve_local_antigravity_request_auth, AntigravityRequestAuth, AntigravityRequestAuthSupport,
|
||||
AntigravityRequestAuthUnsupportedReason, ANTIGRAVITY_PROVIDER_TYPE,
|
||||
ANTIGRAVITY_REQUEST_USER_AGENT,
|
||||
};
|
||||
pub use policy::{
|
||||
classify_local_antigravity_request_support, is_antigravity_provider_transport,
|
||||
AntigravityRequestSideSpec, AntigravityRequestSideSupport,
|
||||
AntigravityRequestSideUnsupportedReason,
|
||||
};
|
||||
pub use request::{
|
||||
build_antigravity_safe_v1internal_request, classify_antigravity_safe_request_body,
|
||||
AntigravityEnvelopeRequestType, AntigravityRequestEnvelopeSupport,
|
||||
AntigravityRequestEnvelopeUnsupportedReason,
|
||||
};
|
||||
pub use url::{
|
||||
build_antigravity_v1internal_url, AntigravityRequestUrlAction,
|
||||
ANTIGRAVITY_V1INTERNAL_PATH_TEMPLATE,
|
||||
};
|
||||
@@ -0,0 +1,116 @@
|
||||
use serde_json::Value;
|
||||
|
||||
use super::super::snapshot::GatewayProviderTransportSnapshot;
|
||||
use super::auth::{
|
||||
resolve_local_antigravity_request_auth, AntigravityRequestAuth, AntigravityRequestAuthSupport,
|
||||
AntigravityRequestAuthUnsupportedReason, ANTIGRAVITY_PROVIDER_TYPE,
|
||||
};
|
||||
use super::request::{
|
||||
classify_antigravity_safe_request_body, AntigravityEnvelopeRequestType,
|
||||
AntigravityRequestEnvelopeUnsupportedReason,
|
||||
};
|
||||
use crate::rules::{body_rules_have_enabled_rules, header_rules_have_enabled_rules};
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct AntigravityRequestSideSpec {
|
||||
pub auth: AntigravityRequestAuth,
|
||||
pub request_type: AntigravityEnvelopeRequestType,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub enum AntigravityRequestSideSupport {
|
||||
Supported(AntigravityRequestSideSpec),
|
||||
Unsupported(AntigravityRequestSideUnsupportedReason),
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub enum AntigravityRequestSideUnsupportedReason {
|
||||
InactiveTransport,
|
||||
WrongProviderType,
|
||||
UnsupportedApiFormat,
|
||||
UnsupportedCustomPath,
|
||||
UnsupportedHeaderRules,
|
||||
UnsupportedBodyRules,
|
||||
UnsupportedNetworkConfig,
|
||||
UnsupportedAuth(AntigravityRequestAuthUnsupportedReason),
|
||||
UnsupportedEnvelope(AntigravityRequestEnvelopeUnsupportedReason),
|
||||
}
|
||||
|
||||
pub fn is_antigravity_provider_transport(transport: &GatewayProviderTransportSnapshot) -> bool {
|
||||
transport
|
||||
.provider
|
||||
.provider_type
|
||||
.trim()
|
||||
.eq_ignore_ascii_case(ANTIGRAVITY_PROVIDER_TYPE)
|
||||
}
|
||||
|
||||
pub fn classify_local_antigravity_request_support(
|
||||
transport: &GatewayProviderTransportSnapshot,
|
||||
request_body: &Value,
|
||||
request_type: AntigravityEnvelopeRequestType,
|
||||
) -> AntigravityRequestSideSupport {
|
||||
if !transport.provider.is_active || !transport.endpoint.is_active || !transport.key.is_active {
|
||||
return AntigravityRequestSideSupport::Unsupported(
|
||||
AntigravityRequestSideUnsupportedReason::InactiveTransport,
|
||||
);
|
||||
}
|
||||
if !is_antigravity_provider_transport(transport) {
|
||||
return AntigravityRequestSideSupport::Unsupported(
|
||||
AntigravityRequestSideUnsupportedReason::WrongProviderType,
|
||||
);
|
||||
}
|
||||
|
||||
let endpoint_format =
|
||||
aether_ai_formats::normalize_api_format_alias(&transport.endpoint.api_format);
|
||||
if endpoint_format != "gemini:generate_content" {
|
||||
return AntigravityRequestSideSupport::Unsupported(
|
||||
AntigravityRequestSideUnsupportedReason::UnsupportedApiFormat,
|
||||
);
|
||||
}
|
||||
if transport
|
||||
.endpoint
|
||||
.custom_path
|
||||
.as_deref()
|
||||
.is_some_and(|value| !value.trim().is_empty())
|
||||
{
|
||||
return AntigravityRequestSideSupport::Unsupported(
|
||||
AntigravityRequestSideUnsupportedReason::UnsupportedCustomPath,
|
||||
);
|
||||
}
|
||||
if header_rules_have_enabled_rules(transport.endpoint.header_rules.as_ref()) {
|
||||
return AntigravityRequestSideSupport::Unsupported(
|
||||
AntigravityRequestSideUnsupportedReason::UnsupportedHeaderRules,
|
||||
);
|
||||
}
|
||||
if body_rules_have_enabled_rules(transport.endpoint.body_rules.as_ref()) {
|
||||
return AntigravityRequestSideSupport::Unsupported(
|
||||
AntigravityRequestSideUnsupportedReason::UnsupportedBodyRules,
|
||||
);
|
||||
}
|
||||
if transport.provider.proxy.is_some()
|
||||
|| transport.endpoint.proxy.is_some()
|
||||
|| transport.key.proxy.is_some()
|
||||
|| transport.key.fingerprint.is_some()
|
||||
{
|
||||
return AntigravityRequestSideSupport::Unsupported(
|
||||
AntigravityRequestSideUnsupportedReason::UnsupportedNetworkConfig,
|
||||
);
|
||||
}
|
||||
|
||||
let auth = match resolve_local_antigravity_request_auth(transport) {
|
||||
AntigravityRequestAuthSupport::Supported(auth) => auth,
|
||||
AntigravityRequestAuthSupport::Unsupported(reason) => {
|
||||
return AntigravityRequestSideSupport::Unsupported(
|
||||
AntigravityRequestSideUnsupportedReason::UnsupportedAuth(reason),
|
||||
);
|
||||
}
|
||||
};
|
||||
|
||||
if let Err(reason) = classify_antigravity_safe_request_body(request_body) {
|
||||
return AntigravityRequestSideSupport::Unsupported(
|
||||
AntigravityRequestSideUnsupportedReason::UnsupportedEnvelope(reason),
|
||||
);
|
||||
}
|
||||
|
||||
AntigravityRequestSideSupport::Supported(AntigravityRequestSideSpec { auth, request_type })
|
||||
}
|
||||
@@ -0,0 +1,367 @@
|
||||
use serde_json::{Map, Value};
|
||||
|
||||
use super::auth::{AntigravityRequestAuth, ANTIGRAVITY_REQUEST_USER_AGENT};
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub enum AntigravityEnvelopeRequestType {
|
||||
Agent,
|
||||
Checkpoint,
|
||||
EndpointTest,
|
||||
}
|
||||
|
||||
impl AntigravityEnvelopeRequestType {
|
||||
fn as_str(self) -> &'static str {
|
||||
match self {
|
||||
Self::Agent => "agent",
|
||||
Self::Checkpoint => "checkpoint",
|
||||
Self::EndpointTest => "endpoint_test",
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq)]
|
||||
pub enum AntigravityRequestEnvelopeSupport {
|
||||
Supported(Value),
|
||||
Unsupported(AntigravityRequestEnvelopeUnsupportedReason),
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub enum AntigravityRequestEnvelopeUnsupportedReason {
|
||||
NonObjectBody,
|
||||
MissingContents,
|
||||
MissingRequestId,
|
||||
MissingModel,
|
||||
}
|
||||
|
||||
pub fn classify_antigravity_safe_request_body(
|
||||
request_body: &Value,
|
||||
) -> Result<(), AntigravityRequestEnvelopeUnsupportedReason> {
|
||||
let Value::Object(map) = request_body else {
|
||||
return Err(AntigravityRequestEnvelopeUnsupportedReason::NonObjectBody);
|
||||
};
|
||||
if !map.contains_key("contents") && existing_v1internal_request_object(map).is_none() {
|
||||
return Err(AntigravityRequestEnvelopeUnsupportedReason::MissingContents);
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub fn build_antigravity_safe_v1internal_request(
|
||||
auth: &AntigravityRequestAuth,
|
||||
request_id: &str,
|
||||
model: &str,
|
||||
request_body: &Value,
|
||||
request_type: AntigravityEnvelopeRequestType,
|
||||
) -> AntigravityRequestEnvelopeSupport {
|
||||
if request_id.trim().is_empty() {
|
||||
return AntigravityRequestEnvelopeSupport::Unsupported(
|
||||
AntigravityRequestEnvelopeUnsupportedReason::MissingRequestId,
|
||||
);
|
||||
}
|
||||
if model.trim().is_empty() {
|
||||
return AntigravityRequestEnvelopeSupport::Unsupported(
|
||||
AntigravityRequestEnvelopeUnsupportedReason::MissingModel,
|
||||
);
|
||||
}
|
||||
if let Err(reason) = classify_antigravity_safe_request_body(request_body) {
|
||||
return AntigravityRequestEnvelopeSupport::Unsupported(reason);
|
||||
}
|
||||
|
||||
let Value::Object(source) = request_body else {
|
||||
return AntigravityRequestEnvelopeSupport::Unsupported(
|
||||
AntigravityRequestEnvelopeUnsupportedReason::NonObjectBody,
|
||||
);
|
||||
};
|
||||
|
||||
if let Some(existing_request) = existing_v1internal_request_object(source) {
|
||||
let mut inner_request: Map<String, Value> = existing_request.clone();
|
||||
inner_request.remove("model");
|
||||
inner_request.remove("safetySettings");
|
||||
inner_request.remove("safety_settings");
|
||||
let request_id = non_empty_string_field(source, "requestId").unwrap_or(request_id);
|
||||
let user_agent =
|
||||
non_empty_string_field(source, "userAgent").unwrap_or(ANTIGRAVITY_REQUEST_USER_AGENT);
|
||||
let request_type =
|
||||
existing_v1internal_request_type(source).unwrap_or_else(|| request_type.as_str());
|
||||
|
||||
return AntigravityRequestEnvelopeSupport::Supported(serde_json::json!({
|
||||
"project": auth.project_id,
|
||||
"requestId": request_id,
|
||||
"request": Value::Object(inner_request),
|
||||
"model": model,
|
||||
"userAgent": user_agent,
|
||||
"requestType": request_type,
|
||||
}));
|
||||
}
|
||||
|
||||
let mut inner_request: Map<String, Value> = source.clone();
|
||||
inner_request.remove("model");
|
||||
inner_request.remove("safetySettings");
|
||||
inner_request.remove("safety_settings");
|
||||
|
||||
AntigravityRequestEnvelopeSupport::Supported(serde_json::json!({
|
||||
"project": auth.project_id,
|
||||
"requestId": request_id,
|
||||
"request": Value::Object(inner_request),
|
||||
"model": model,
|
||||
"userAgent": ANTIGRAVITY_REQUEST_USER_AGENT,
|
||||
"requestType": request_type.as_str(),
|
||||
}))
|
||||
}
|
||||
|
||||
fn existing_v1internal_request_object(source: &Map<String, Value>) -> Option<&Map<String, Value>> {
|
||||
source
|
||||
.get("request")
|
||||
.and_then(Value::as_object)
|
||||
.filter(|request| request.contains_key("contents"))
|
||||
}
|
||||
|
||||
fn non_empty_string_field<'a>(source: &'a Map<String, Value>, key: &str) -> Option<&'a str> {
|
||||
source
|
||||
.get(key)
|
||||
.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
}
|
||||
|
||||
fn existing_v1internal_request_type(source: &Map<String, Value>) -> Option<&str> {
|
||||
match non_empty_string_field(source, "requestType")? {
|
||||
"agent" => Some("agent"),
|
||||
"checkpoint" => Some("checkpoint"),
|
||||
"endpoint_test" => Some("endpoint_test"),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use serde_json::json;
|
||||
|
||||
use super::{
|
||||
build_antigravity_safe_v1internal_request, classify_antigravity_safe_request_body,
|
||||
AntigravityEnvelopeRequestType, AntigravityRequestAuth, AntigravityRequestEnvelopeSupport,
|
||||
};
|
||||
use crate::antigravity::ANTIGRAVITY_REQUEST_USER_AGENT;
|
||||
|
||||
fn sample_auth() -> AntigravityRequestAuth {
|
||||
AntigravityRequestAuth {
|
||||
project_id: "project-ant-123".to_string(),
|
||||
client_version: None,
|
||||
session_id: None,
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn real_agent_request_preserves_antigravity_agent_fields() {
|
||||
let request_body = json!({
|
||||
"model": "client-side-model-should-not-be-nested",
|
||||
"contents": [
|
||||
{
|
||||
"role": "user",
|
||||
"parts": [
|
||||
{ "text": "Reply with OK only." }
|
||||
]
|
||||
}
|
||||
],
|
||||
"systemInstruction": {
|
||||
"role": "user",
|
||||
"parts": [
|
||||
{ "text": "Antigravity agent system prompt" }
|
||||
]
|
||||
},
|
||||
"generationConfig": {
|
||||
"maxOutputTokens": 8192,
|
||||
"thinkingConfig": {
|
||||
"includeThoughts": true,
|
||||
"thinkingBudget": 4000
|
||||
}
|
||||
},
|
||||
"toolConfig": {
|
||||
"functionCallingConfig": {
|
||||
"mode": "VALIDATED"
|
||||
}
|
||||
},
|
||||
"tools": [
|
||||
{
|
||||
"functionDeclarations": [
|
||||
{
|
||||
"name": "run_command",
|
||||
"description": "Run a command",
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"cmd": { "type": "string" }
|
||||
},
|
||||
"required": ["cmd"]
|
||||
}
|
||||
}
|
||||
]
|
||||
}
|
||||
],
|
||||
"labels": {
|
||||
"trajectory_id": "trajectory-123",
|
||||
"used_claude": "false"
|
||||
},
|
||||
"sessionId": "session-ant-123",
|
||||
"safetySettings": [
|
||||
{ "category": "HARM_CATEGORY_UNSPECIFIED" }
|
||||
]
|
||||
});
|
||||
|
||||
assert_eq!(
|
||||
classify_antigravity_safe_request_body(&request_body),
|
||||
Ok(())
|
||||
);
|
||||
|
||||
let envelope = match build_antigravity_safe_v1internal_request(
|
||||
&sample_auth(),
|
||||
"request-ant-agent-123",
|
||||
"gemini-3.5-flash-low",
|
||||
&request_body,
|
||||
AntigravityEnvelopeRequestType::Agent,
|
||||
) {
|
||||
AntigravityRequestEnvelopeSupport::Supported(envelope) => envelope,
|
||||
AntigravityRequestEnvelopeSupport::Unsupported(reason) => {
|
||||
panic!("real agent envelope should be supported: {reason:?}")
|
||||
}
|
||||
};
|
||||
|
||||
assert_eq!(envelope["project"], "project-ant-123");
|
||||
assert_eq!(envelope["requestId"], "request-ant-agent-123");
|
||||
assert_eq!(envelope["model"], "gemini-3.5-flash-low");
|
||||
assert_eq!(envelope["userAgent"], ANTIGRAVITY_REQUEST_USER_AGENT);
|
||||
assert_eq!(envelope["requestType"], "agent");
|
||||
assert!(envelope["request"].get("model").is_none());
|
||||
assert!(envelope["request"].get("safetySettings").is_none());
|
||||
assert_eq!(
|
||||
envelope["request"]["systemInstruction"]["parts"][0]["text"],
|
||||
"Antigravity agent system prompt"
|
||||
);
|
||||
assert_eq!(
|
||||
envelope["request"]["generationConfig"]["thinkingConfig"]["thinkingBudget"],
|
||||
4000
|
||||
);
|
||||
assert_eq!(
|
||||
envelope["request"]["toolConfig"]["functionCallingConfig"]["mode"],
|
||||
"VALIDATED"
|
||||
);
|
||||
assert_eq!(
|
||||
envelope["request"]["tools"][0]["functionDeclarations"][0]["name"],
|
||||
"run_command"
|
||||
);
|
||||
assert_eq!(
|
||||
envelope["request"]["labels"]["trajectory_id"],
|
||||
"trajectory-123"
|
||||
);
|
||||
assert_eq!(envelope["request"]["sessionId"], "session-ant-123");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn checkpoint_request_type_builds_checkpoint_envelope() {
|
||||
let request_body = json!({
|
||||
"contents": [
|
||||
{
|
||||
"role": "user",
|
||||
"parts": [
|
||||
{ "text": "checkpoint context" }
|
||||
]
|
||||
}
|
||||
],
|
||||
"generationConfig": {
|
||||
"maxOutputTokens": 8192,
|
||||
"thinkingConfig": {
|
||||
"includeThoughts": true,
|
||||
"thinkingBudget": 4000
|
||||
}
|
||||
},
|
||||
"toolConfig": {
|
||||
"functionCallingConfig": {
|
||||
"mode": "NONE"
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
let envelope = match build_antigravity_safe_v1internal_request(
|
||||
&sample_auth(),
|
||||
"request-ant-checkpoint-123",
|
||||
"gemini-3.5-flash-low",
|
||||
&request_body,
|
||||
AntigravityEnvelopeRequestType::Checkpoint,
|
||||
) {
|
||||
AntigravityRequestEnvelopeSupport::Supported(envelope) => envelope,
|
||||
AntigravityRequestEnvelopeSupport::Unsupported(reason) => {
|
||||
panic!("checkpoint envelope should be supported: {reason:?}")
|
||||
}
|
||||
};
|
||||
|
||||
assert_eq!(envelope["requestType"], "checkpoint");
|
||||
assert_eq!(
|
||||
envelope["request"]["toolConfig"]["functionCallingConfig"]["mode"],
|
||||
"NONE"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn existing_v1internal_envelope_is_not_double_wrapped() {
|
||||
let request_body = json!({
|
||||
"project": "client-side-project",
|
||||
"requestId": "client-request-id-123",
|
||||
"model": "gemini-3.5-flash-low",
|
||||
"userAgent": "antigravity",
|
||||
"requestType": "checkpoint",
|
||||
"request": {
|
||||
"contents": [
|
||||
{
|
||||
"role": "user",
|
||||
"parts": [
|
||||
{ "text": "checkpoint context" }
|
||||
]
|
||||
}
|
||||
],
|
||||
"generationConfig": {
|
||||
"thinkingConfig": {
|
||||
"includeThoughts": true
|
||||
}
|
||||
},
|
||||
"toolConfig": {
|
||||
"functionCallingConfig": {
|
||||
"mode": "NONE"
|
||||
}
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
assert_eq!(
|
||||
classify_antigravity_safe_request_body(&request_body),
|
||||
Ok(())
|
||||
);
|
||||
|
||||
let envelope = match build_antigravity_safe_v1internal_request(
|
||||
&sample_auth(),
|
||||
"trace-request-id-should-not-overwrite-client-id",
|
||||
"mapped-antigravity-model",
|
||||
&request_body,
|
||||
AntigravityEnvelopeRequestType::Agent,
|
||||
) {
|
||||
AntigravityRequestEnvelopeSupport::Supported(envelope) => envelope,
|
||||
AntigravityRequestEnvelopeSupport::Unsupported(reason) => {
|
||||
panic!("existing v1internal envelope should be supported: {reason:?}")
|
||||
}
|
||||
};
|
||||
|
||||
assert_eq!(envelope["project"], "project-ant-123");
|
||||
assert_eq!(envelope["requestId"], "client-request-id-123");
|
||||
assert_eq!(envelope["model"], "mapped-antigravity-model");
|
||||
assert_eq!(envelope["userAgent"], "antigravity");
|
||||
assert_eq!(envelope["requestType"], "checkpoint");
|
||||
assert!(envelope["request"].get("request").is_none());
|
||||
assert_eq!(
|
||||
envelope["request"]["contents"][0]["parts"][0]["text"],
|
||||
"checkpoint context"
|
||||
);
|
||||
assert_eq!(
|
||||
envelope["request"]["toolConfig"]["functionCallingConfig"]["mode"],
|
||||
"NONE"
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,73 @@
|
||||
use std::collections::BTreeMap;
|
||||
|
||||
use url::form_urlencoded;
|
||||
|
||||
pub const ANTIGRAVITY_V1INTERNAL_PATH_TEMPLATE: &str = "/v1internal:{action}";
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub enum AntigravityRequestUrlAction {
|
||||
GenerateContent,
|
||||
StreamGenerateContent,
|
||||
}
|
||||
|
||||
impl AntigravityRequestUrlAction {
|
||||
fn as_str(self) -> &'static str {
|
||||
match self {
|
||||
Self::GenerateContent => "generateContent",
|
||||
Self::StreamGenerateContent => "streamGenerateContent",
|
||||
}
|
||||
}
|
||||
|
||||
fn is_stream(self) -> bool {
|
||||
matches!(self, Self::StreamGenerateContent)
|
||||
}
|
||||
}
|
||||
|
||||
pub fn build_antigravity_v1internal_url(
|
||||
base_url: &str,
|
||||
action: AntigravityRequestUrlAction,
|
||||
query: Option<&BTreeMap<String, String>>,
|
||||
) -> Option<String> {
|
||||
let trimmed_base = base_url.trim();
|
||||
if trimmed_base.is_empty() {
|
||||
return None;
|
||||
}
|
||||
|
||||
let path = ANTIGRAVITY_V1INTERNAL_PATH_TEMPLATE.replace("{action}", action.as_str());
|
||||
let mut url = format!("{}{}", trimmed_base.trim_end_matches('/'), path);
|
||||
|
||||
let mut params = BTreeMap::new();
|
||||
if let Some(query) = query {
|
||||
for (key, value) in query {
|
||||
let key = key.trim();
|
||||
let value = value.trim();
|
||||
if key.is_empty()
|
||||
|| value.is_empty()
|
||||
|| key.eq_ignore_ascii_case("beta")
|
||||
|| key.eq_ignore_ascii_case("key")
|
||||
{
|
||||
continue;
|
||||
}
|
||||
params.insert(key.to_string(), value.to_string());
|
||||
}
|
||||
}
|
||||
if action.is_stream() {
|
||||
params
|
||||
.entry(String::from("alt"))
|
||||
.or_insert_with(|| String::from("sse"));
|
||||
}
|
||||
|
||||
if !params.is_empty() {
|
||||
let mut serializer = form_urlencoded::Serializer::new(String::new());
|
||||
for (key, value) in params {
|
||||
serializer.append_pair(key.as_str(), value.as_str());
|
||||
}
|
||||
let query_string = serializer.finish();
|
||||
if !query_string.is_empty() {
|
||||
url.push('?');
|
||||
url.push_str(&query_string);
|
||||
}
|
||||
}
|
||||
|
||||
Some(url)
|
||||
}
|
||||
Reference in New Issue
Block a user