mirror of
https://github.com/fawney19/Aether.git
synced 2026-10-05 00:47:48 +08:00
Add Antigravity v1internal gateway adapter
This commit is contained in:
MMEXA
@@ -5,6 +5,7 @@ use super::auth::{AntigravityRequestAuth, ANTIGRAVITY_REQUEST_USER_AGENT};
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub enum AntigravityEnvelopeRequestType {
|
||||
Agent,
|
||||
Checkpoint,
|
||||
EndpointTest,
|
||||
}
|
||||
|
||||
@@ -12,6 +13,7 @@ impl AntigravityEnvelopeRequestType {
|
||||
fn as_str(self) -> &'static str {
|
||||
match self {
|
||||
Self::Agent => "agent",
|
||||
Self::Checkpoint => "checkpoint",
|
||||
Self::EndpointTest => "endpoint_test",
|
||||
}
|
||||
}
|
||||
@@ -29,7 +31,6 @@ pub enum AntigravityRequestEnvelopeUnsupportedReason {
|
||||
MissingContents,
|
||||
MissingRequestId,
|
||||
MissingModel,
|
||||
ComplexEnvelopeTransform,
|
||||
}
|
||||
|
||||
pub fn classify_antigravity_safe_request_body(
|
||||
@@ -38,12 +39,9 @@ pub fn classify_antigravity_safe_request_body(
|
||||
let Value::Object(map) = request_body else {
|
||||
return Err(AntigravityRequestEnvelopeUnsupportedReason::NonObjectBody);
|
||||
};
|
||||
if !map.contains_key("contents") {
|
||||
if !map.contains_key("contents") && existing_v1internal_request_object(map).is_none() {
|
||||
return Err(AntigravityRequestEnvelopeUnsupportedReason::MissingContents);
|
||||
}
|
||||
if contains_blocked_request_features(request_body) {
|
||||
return Err(AntigravityRequestEnvelopeUnsupportedReason::ComplexEnvelopeTransform);
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
@@ -75,6 +73,27 @@ pub fn build_antigravity_safe_v1internal_request(
|
||||
);
|
||||
};
|
||||
|
||||
if let Some(existing_request) = existing_v1internal_request_object(source) {
|
||||
let mut inner_request: Map<String, Value> = existing_request.clone();
|
||||
inner_request.remove("model");
|
||||
inner_request.remove("safetySettings");
|
||||
inner_request.remove("safety_settings");
|
||||
let request_id = non_empty_string_field(source, "requestId").unwrap_or(request_id);
|
||||
let user_agent =
|
||||
non_empty_string_field(source, "userAgent").unwrap_or(ANTIGRAVITY_REQUEST_USER_AGENT);
|
||||
let request_type =
|
||||
existing_v1internal_request_type(source).unwrap_or_else(|| request_type.as_str());
|
||||
|
||||
return AntigravityRequestEnvelopeSupport::Supported(serde_json::json!({
|
||||
"project": auth.project_id,
|
||||
"requestId": request_id,
|
||||
"request": Value::Object(inner_request),
|
||||
"model": model,
|
||||
"userAgent": user_agent,
|
||||
"requestType": request_type,
|
||||
}));
|
||||
}
|
||||
|
||||
let mut inner_request: Map<String, Value> = source.clone();
|
||||
inner_request.remove("model");
|
||||
inner_request.remove("safetySettings");
|
||||
@@ -90,31 +109,258 @@ pub fn build_antigravity_safe_v1internal_request(
|
||||
}))
|
||||
}
|
||||
|
||||
fn contains_blocked_request_features(value: &Value) -> bool {
|
||||
match value {
|
||||
Value::Object(map) => map.iter().any(|(key, inner)| {
|
||||
is_blocked_request_key(key.as_str()) || contains_blocked_request_features(inner)
|
||||
}),
|
||||
Value::Array(items) => items.iter().any(contains_blocked_request_features),
|
||||
_ => false,
|
||||
fn existing_v1internal_request_object(source: &Map<String, Value>) -> Option<&Map<String, Value>> {
|
||||
source
|
||||
.get("request")
|
||||
.and_then(Value::as_object)
|
||||
.filter(|request| request.contains_key("contents"))
|
||||
}
|
||||
|
||||
fn non_empty_string_field<'a>(source: &'a Map<String, Value>, key: &str) -> Option<&'a str> {
|
||||
source
|
||||
.get(key)
|
||||
.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
}
|
||||
|
||||
fn existing_v1internal_request_type(source: &Map<String, Value>) -> Option<&str> {
|
||||
match non_empty_string_field(source, "requestType")? {
|
||||
"agent" => Some("agent"),
|
||||
"checkpoint" => Some("checkpoint"),
|
||||
"endpoint_test" => Some("endpoint_test"),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
fn is_blocked_request_key(key: &str) -> bool {
|
||||
matches!(
|
||||
key.trim(),
|
||||
"systemInstruction"
|
||||
| "system_instruction"
|
||||
| "tools"
|
||||
| "toolConfig"
|
||||
| "tool_config"
|
||||
| "thinkingConfig"
|
||||
| "thinking_config"
|
||||
| "imageConfig"
|
||||
| "image_config"
|
||||
| "functionCall"
|
||||
| "function_call"
|
||||
| "functionResponse"
|
||||
| "function_response"
|
||||
)
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use serde_json::json;
|
||||
|
||||
use super::{
|
||||
build_antigravity_safe_v1internal_request, classify_antigravity_safe_request_body,
|
||||
AntigravityEnvelopeRequestType, AntigravityRequestAuth, AntigravityRequestEnvelopeSupport,
|
||||
};
|
||||
|
||||
fn sample_auth() -> AntigravityRequestAuth {
|
||||
AntigravityRequestAuth {
|
||||
project_id: "project-ant-123".to_string(),
|
||||
client_version: None,
|
||||
session_id: None,
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn real_agent_request_preserves_antigravity_agent_fields() {
|
||||
let request_body = json!({
|
||||
"model": "client-side-model-should-not-be-nested",
|
||||
"contents": [
|
||||
{
|
||||
"role": "user",
|
||||
"parts": [
|
||||
{ "text": "Reply with OK only." }
|
||||
]
|
||||
}
|
||||
],
|
||||
"systemInstruction": {
|
||||
"role": "user",
|
||||
"parts": [
|
||||
{ "text": "Antigravity agent system prompt" }
|
||||
]
|
||||
},
|
||||
"generationConfig": {
|
||||
"maxOutputTokens": 8192,
|
||||
"thinkingConfig": {
|
||||
"includeThoughts": true,
|
||||
"thinkingBudget": 4000
|
||||
}
|
||||
},
|
||||
"toolConfig": {
|
||||
"functionCallingConfig": {
|
||||
"mode": "VALIDATED"
|
||||
}
|
||||
},
|
||||
"tools": [
|
||||
{
|
||||
"functionDeclarations": [
|
||||
{
|
||||
"name": "run_command",
|
||||
"description": "Run a command",
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"cmd": { "type": "string" }
|
||||
},
|
||||
"required": ["cmd"]
|
||||
}
|
||||
}
|
||||
]
|
||||
}
|
||||
],
|
||||
"labels": {
|
||||
"trajectory_id": "trajectory-123",
|
||||
"used_claude": "false"
|
||||
},
|
||||
"sessionId": "session-ant-123",
|
||||
"safetySettings": [
|
||||
{ "category": "HARM_CATEGORY_UNSPECIFIED" }
|
||||
]
|
||||
});
|
||||
|
||||
assert_eq!(
|
||||
classify_antigravity_safe_request_body(&request_body),
|
||||
Ok(())
|
||||
);
|
||||
|
||||
let envelope = match build_antigravity_safe_v1internal_request(
|
||||
&sample_auth(),
|
||||
"request-ant-agent-123",
|
||||
"gemini-3.5-flash-low",
|
||||
&request_body,
|
||||
AntigravityEnvelopeRequestType::Agent,
|
||||
) {
|
||||
AntigravityRequestEnvelopeSupport::Supported(envelope) => envelope,
|
||||
AntigravityRequestEnvelopeSupport::Unsupported(reason) => {
|
||||
panic!("real agent envelope should be supported: {reason:?}")
|
||||
}
|
||||
};
|
||||
|
||||
assert_eq!(envelope["project"], "project-ant-123");
|
||||
assert_eq!(envelope["requestId"], "request-ant-agent-123");
|
||||
assert_eq!(envelope["model"], "gemini-3.5-flash-low");
|
||||
assert_eq!(envelope["userAgent"], "antigravity");
|
||||
assert_eq!(envelope["requestType"], "agent");
|
||||
assert!(envelope["request"].get("model").is_none());
|
||||
assert!(envelope["request"].get("safetySettings").is_none());
|
||||
assert_eq!(
|
||||
envelope["request"]["systemInstruction"]["parts"][0]["text"],
|
||||
"Antigravity agent system prompt"
|
||||
);
|
||||
assert_eq!(
|
||||
envelope["request"]["generationConfig"]["thinkingConfig"]["thinkingBudget"],
|
||||
4000
|
||||
);
|
||||
assert_eq!(
|
||||
envelope["request"]["toolConfig"]["functionCallingConfig"]["mode"],
|
||||
"VALIDATED"
|
||||
);
|
||||
assert_eq!(
|
||||
envelope["request"]["tools"][0]["functionDeclarations"][0]["name"],
|
||||
"run_command"
|
||||
);
|
||||
assert_eq!(
|
||||
envelope["request"]["labels"]["trajectory_id"],
|
||||
"trajectory-123"
|
||||
);
|
||||
assert_eq!(envelope["request"]["sessionId"], "session-ant-123");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn checkpoint_request_type_builds_checkpoint_envelope() {
|
||||
let request_body = json!({
|
||||
"contents": [
|
||||
{
|
||||
"role": "user",
|
||||
"parts": [
|
||||
{ "text": "checkpoint context" }
|
||||
]
|
||||
}
|
||||
],
|
||||
"generationConfig": {
|
||||
"maxOutputTokens": 8192,
|
||||
"thinkingConfig": {
|
||||
"includeThoughts": true,
|
||||
"thinkingBudget": 4000
|
||||
}
|
||||
},
|
||||
"toolConfig": {
|
||||
"functionCallingConfig": {
|
||||
"mode": "NONE"
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
let envelope = match build_antigravity_safe_v1internal_request(
|
||||
&sample_auth(),
|
||||
"request-ant-checkpoint-123",
|
||||
"gemini-3.5-flash-low",
|
||||
&request_body,
|
||||
AntigravityEnvelopeRequestType::Checkpoint,
|
||||
) {
|
||||
AntigravityRequestEnvelopeSupport::Supported(envelope) => envelope,
|
||||
AntigravityRequestEnvelopeSupport::Unsupported(reason) => {
|
||||
panic!("checkpoint envelope should be supported: {reason:?}")
|
||||
}
|
||||
};
|
||||
|
||||
assert_eq!(envelope["requestType"], "checkpoint");
|
||||
assert_eq!(
|
||||
envelope["request"]["toolConfig"]["functionCallingConfig"]["mode"],
|
||||
"NONE"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn existing_v1internal_envelope_is_not_double_wrapped() {
|
||||
let request_body = json!({
|
||||
"project": "client-side-project",
|
||||
"requestId": "client-request-id-123",
|
||||
"model": "gemini-3.5-flash-low",
|
||||
"userAgent": "antigravity",
|
||||
"requestType": "checkpoint",
|
||||
"request": {
|
||||
"contents": [
|
||||
{
|
||||
"role": "user",
|
||||
"parts": [
|
||||
{ "text": "checkpoint context" }
|
||||
]
|
||||
}
|
||||
],
|
||||
"generationConfig": {
|
||||
"thinkingConfig": {
|
||||
"includeThoughts": true
|
||||
}
|
||||
},
|
||||
"toolConfig": {
|
||||
"functionCallingConfig": {
|
||||
"mode": "NONE"
|
||||
}
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
assert_eq!(
|
||||
classify_antigravity_safe_request_body(&request_body),
|
||||
Ok(())
|
||||
);
|
||||
|
||||
let envelope = match build_antigravity_safe_v1internal_request(
|
||||
&sample_auth(),
|
||||
"trace-request-id-should-not-overwrite-client-id",
|
||||
"mapped-antigravity-model",
|
||||
&request_body,
|
||||
AntigravityEnvelopeRequestType::Agent,
|
||||
) {
|
||||
AntigravityRequestEnvelopeSupport::Supported(envelope) => envelope,
|
||||
AntigravityRequestEnvelopeSupport::Unsupported(reason) => {
|
||||
panic!("existing v1internal envelope should be supported: {reason:?}")
|
||||
}
|
||||
};
|
||||
|
||||
assert_eq!(envelope["project"], "project-ant-123");
|
||||
assert_eq!(envelope["requestId"], "client-request-id-123");
|
||||
assert_eq!(envelope["model"], "mapped-antigravity-model");
|
||||
assert_eq!(envelope["userAgent"], "antigravity");
|
||||
assert_eq!(envelope["requestType"], "checkpoint");
|
||||
assert!(envelope["request"].get("request").is_none());
|
||||
assert_eq!(
|
||||
envelope["request"]["contents"][0]["parts"][0]["text"],
|
||||
"checkpoint context"
|
||||
);
|
||||
assert_eq!(
|
||||
envelope["request"]["toolConfig"]["functionCallingConfig"]["mode"],
|
||||
"NONE"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user