use super::models::V1InternalRequest;
use serde_json::{json, Value};
pub fn wrap_request(
body: &Value,
project_id: &str,
mapped_model: &str,
session_id: Option<&str>,
) -> Value {
let original_model = body
.get("model")
.and_then(|v| v.as_str())
.unwrap_or(mapped_model);
let final_model_name = if !mapped_model.is_empty() {
mapped_model
} else {
original_model
};
let mut inner_request = body.clone();
crate::proxy::mappers::common_utils::deep_clean_undefined(&mut inner_request);
let is_target_claude = crate::proxy::common::model_mapping::is_claude_model(final_model_name);
if let Some(contents) = inner_request
.get_mut("contents")
.and_then(|c| c.as_array_mut())
{
for content in contents {
let mut name_counters: std::collections::HashMap<String, usize> =
std::collections::HashMap::new();
if let Some(parts) = content.get_mut("parts").and_then(|p| p.as_array_mut()) {
for part in parts {
if let Some(obj) = part.as_object_mut() {
if let Some(fc) = obj.get_mut("functionCall") {
if fc.get("id").is_none() && is_target_claude {
let name =
fc.get("name").and_then(|n| n.as_str()).unwrap_or("unknown");
let count = name_counters.entry(name.to_string()).or_insert(0);
let call_id = format!("call_{}_{}", name, count);
*count += 1;
fc.as_object_mut()
.unwrap()
.insert("id".to_string(), json!(call_id));
tracing::debug!("[Gemini-Wrap] Request stage: Injected missing call_id '{}' for Claude model", call_id);
}
}
if let Some(fr) = obj.get_mut("functionResponse") {
if fr.get("id").is_none() && is_target_claude {
let name =
fr.get("name").and_then(|n| n.as_str()).unwrap_or("unknown");
let count = name_counters.entry(name.to_string()).or_insert(0);
let call_id = format!("call_{}_{}", name, count);
*count += 1;
fr.as_object_mut()
.unwrap()
.insert("id".to_string(), json!(call_id));
tracing::debug!("[Gemini-Wrap] Request stage: Injected synced response_id '{}' for Claude model", call_id);
}
}
if obj.contains_key("functionCall") && obj.get("thoughtSignature").is_none()
{
if let Some(s_id) = session_id {
if let Some(sig) = crate::proxy::SignatureCache::global()
.get_session_signature(s_id)
{
obj.insert("thoughtSignature".to_string(), json!(sig));
tracing::debug!("[Gemini-Wrap] Injected signature (len: {}) for session: {}", sig.len(), s_id);
}
}
}
}
}
}
}
}
let lower_model = final_model_name.to_lowercase();
if lower_model.contains("flash")
|| lower_model.contains("pro")
|| lower_model.contains("thinking")
{
let gen_config = inner_request
.as_object_mut()
.unwrap()
.entry("generationConfig")
.or_insert(json!({}))
.as_object_mut()
.unwrap();
if gen_config.get("thinkingConfig").is_none() {
let should_inject =
crate::proxy::common::model_mapping::model_supports_thinking(&lower_model);
if should_inject {
tracing::debug!(
"[Gemini-Wrap] Auto-injecting default thinkingConfig for {}",
final_model_name
);
gen_config.insert(
"thinkingConfig".to_string(),
json!({
"includeThoughts": true,
"thinkingBudget": 24576
}),
);
}
}
if let Some(thinking_config) = gen_config.get_mut("thinkingConfig") {
if let Some(budget_val) = thinking_config.get("thinkingBudget") {
if let Some(budget) = budget_val.as_u64() {
let tb_config = crate::proxy::config::get_thinking_budget_config();
let final_budget = match tb_config.mode {
crate::proxy::config::ThinkingBudgetMode::Passthrough => {
tracing::debug!(
"[Gemini-Wrap] Passthrough mode: keeping budget {} for model {}",
budget,
final_model_name
);
budget
}
crate::proxy::config::ThinkingBudgetMode::Custom => {
let val = tb_config.custom_value as u64;
if val > 24576 {
tracing::warn!(
"[Gemini-Wrap] Custom mode: capping thinking_budget from {} to 24576 for model {}",
val, final_model_name
);
24576
} else {
if val != budget {
tracing::debug!(
"[Gemini-Wrap] Custom mode: overriding {} with {} for model {}",
budget, val, final_model_name
);
}
val
}
}
crate::proxy::config::ThinkingBudgetMode::Auto => {
if budget > 24576 {
tracing::info!(
"[Gemini-Wrap] Auto mode: capping thinking_budget from {} to 24576 for model {}",
budget, final_model_name
);
24576
} else {
budget
}
}
};
if final_budget != budget {
thinking_config["thinkingBudget"] = json!(final_budget);
}
}
}
}
}
let tools_val: Option<Vec<Value>> = inner_request
.get("tools")
.and_then(|t| t.as_array())
.cloned();
let size = body.get("size").and_then(|v| v.as_str());
let quality = body.get("quality").and_then(|v| v.as_str());
let config = crate::proxy::mappers::common_utils::resolve_request_config(
original_model,
final_model_name,
&tools_val,
size,
quality,
Some(body),
);
if let Some(tools) = inner_request.get_mut("tools") {
if let Some(tools_arr) = tools.as_array_mut() {
for tool in tools_arr {
if let Some(decls) = tool.get_mut("functionDeclarations") {
if let Some(decls_arr) = decls.as_array_mut() {
decls_arr.retain(|decl| {
if let Some(name) = decl.get("name").and_then(|v| v.as_str()) {
if name == "web_search" || name == "google_search" {
return false;
}
}
true
});
for decl in decls_arr {
if let Some(decl_obj) = decl.as_object_mut() {
if let Some(params_json_schema) =
decl_obj.remove("parametersJsonSchema")
{
let mut params = params_json_schema;
crate::proxy::common::json_schema::clean_json_schema(
&mut params,
);
decl_obj.insert("parameters".to_string(), params);
} else if let Some(params) = decl_obj.get_mut("parameters") {
crate::proxy::common::json_schema::clean_json_schema(params);
}
}
}
}
}
}
}
}
tracing::debug!(
" Gemini Wrap: original='{}', mapped='{}', final='{}', type='{}'",
original_model,
final_model_name,
config.final_model,
config.request_type
);
if config.inject_google_search {
crate::proxy::mappers::common_utils::inject_google_search_tool(&mut inner_request);
}
if let Some(image_config) = config.image_config {
if let Some(obj) = inner_request.as_object_mut() {
obj.remove("tools");
obj.remove("systemInstruction");
if let Some(contents) = obj.get_mut("contents").and_then(|c| c.as_array_mut()) {
for content in contents {
if let Some(c_obj) = content.as_object_mut() {
if !c_obj.contains_key("role") {
c_obj.insert("role".to_string(), json!("user"));
}
}
}
}
let gen_config = obj.entry("generationConfig").or_insert_with(|| json!({}));
if let Some(gen_obj) = gen_config.as_object_mut() {
gen_obj.remove("responseMimeType");
gen_obj.remove("responseModalities");
gen_obj.insert("imageConfig".to_string(), image_config);
}
}
} else {
let antigravity_identity = "You are Antigravity, a powerful agentic AI coding assistant designed by the Google Deepmind team working on Advanced Agentic Coding.\n\
You are pair programming with a USER to solve their coding task. The task may require creating a new codebase, modifying or debugging an existing codebase, or simply answering a question.\n\
**Absolute paths only**\n\
**Proactiveness**";
if let Some(system_instruction) = inner_request.get_mut("systemInstruction") {
if let Some(obj) = system_instruction.as_object_mut() {
if !obj.contains_key("role") {
obj.insert("role".to_string(), json!("user"));
}
}
if let Some(parts) = system_instruction.get_mut("parts") {
if let Some(parts_array) = parts.as_array_mut() {
let has_antigravity = parts_array
.first()
.and_then(|p| p.get("text"))
.and_then(|t| t.as_str())
.map(|s| s.contains("You are Antigravity"))
.unwrap_or(false);
if !has_antigravity {
parts_array.insert(0, json!({"text": antigravity_identity}));
}
}
}
} else {
inner_request["systemInstruction"] = json!({
"role": "user",
"parts": [{"text": antigravity_identity}]
});
}
}
let final_request = V1InternalRequest {
project: project_id.to_string(),
request_id: format!("agent-{}", uuid::Uuid::new_v4()),
request: inner_request,
model: config.final_model,
user_agent: "antigravity".to_string(),
request_type: config.request_type,
};
serde_json::to_value(final_request).unwrap_or_else(|_| json!({}))
}
#[cfg(test)]
mod test_fixes {
use super::*;
use crate::proxy::common::model_mapping::MODEL_GEMINI_PRO_ALIAS;
use serde_json::json;
#[test]
fn test_wrap_request_with_signature() {
let session_id = "test-session-sig";
let signature = "test-signature-must-be-longer-than-fifty-characters-to-be-cached-by-signature-cache-12345";
crate::proxy::SignatureCache::global().cache_session_signature(
session_id,
signature.to_string(),
1,
);
let body = json!({
"model": MODEL_GEMINI_PRO_ALIAS,
"contents": [{
"role": "user",
"parts": [{
"functionCall": {
"name": "get_weather",
"args": {"location": "London"}
}
}]
}]
});
let result = wrap_request(&body, "proj", MODEL_GEMINI_PRO_ALIAS, Some(session_id));
let injected_sig = result["request"]["contents"][0]["parts"][0]["thoughtSignature"]
.as_str()
.unwrap();
assert_eq!(injected_sig, signature);
}
}
pub fn unwrap_response(response: &Value) -> Value {
response.get("response").unwrap_or(response).clone()
}
pub fn inject_ids_to_response(response: &mut Value, model_name: &str) {
if !crate::proxy::common::model_mapping::is_claude_model(model_name) {
return;
}
if let Some(candidates) = response
.get_mut("candidates")
.and_then(|c| c.as_array_mut())
{
for candidate in candidates {
if let Some(parts) = candidate
.get_mut("content")
.and_then(|c| c.get_mut("parts"))
.and_then(|p| p.as_array_mut())
{
let mut name_counters: std::collections::HashMap<String, usize> =
std::collections::HashMap::new();
for part in parts {
if let Some(fc) = part.get_mut("functionCall").and_then(|f| f.as_object_mut()) {
if fc.get("id").is_none() {
let name = fc.get("name").and_then(|n| n.as_str()).unwrap_or("unknown");
let count = name_counters.entry(name.to_string()).or_insert(0);
let call_id = format!("call_{}_{}", name, count);
*count += 1;
fc.insert("id".to_string(), json!(call_id));
tracing::debug!("[Gemini-Wrap] Response stage: Injected synthetic call_id '{}' for client", call_id);
}
}
}
}
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::proxy::common::model_mapping::{
MODEL_GEMINI_3_FLASH, MODEL_GEMINI_3_FLASH_THINKING, MODEL_GEMINI_3_PRO,
MODEL_GEMINI_3_PRO_IMAGE, MODEL_GEMINI_3_PRO_PREVIEW, MODEL_GEMINI_PRO_ALIAS,
};
use serde_json::json;
#[test]
fn test_wrap_request() {
let body = json!({
"model": MODEL_GEMINI_3_FLASH,
"contents": [{"role": "user", "parts": [{"text": "Hi"}]}]
});
let result = wrap_request(&body, "test-project", MODEL_GEMINI_3_FLASH, None);
assert_eq!(result["project"], "test-project");
assert_eq!(result["model"], MODEL_GEMINI_3_FLASH);
assert!(result["requestId"].as_str().unwrap().starts_with("agent-"));
}
#[test]
fn test_unwrap_response() {
let wrapped = json!({
"response": {
"candidates": [{"content": {"parts": [{"text": "Hello"}]}}]
}
});
let result = unwrap_response(&wrapped);
assert!(result.get("candidates").is_some());
assert!(result.get("response").is_none());
}
#[test]
fn test_antigravity_identity_injection_with_role() {
let body = json!({
"model": MODEL_GEMINI_PRO_ALIAS,
"messages": []
});
let result = wrap_request(&body, "test-proj", MODEL_GEMINI_PRO_ALIAS, None);
let _sys = result
.get("request")
.unwrap()
.get("systemInstruction")
.unwrap();
}
#[test]
fn test_gemini_flash_thinking_budget_capping() {
let _budget_guard = crate::proxy::config::lock_thinking_budget_for_test();
let body = json!({
"model": MODEL_GEMINI_3_FLASH_THINKING,
"generationConfig": {
"thinkingConfig": {
"includeThoughts": true,
"thinkingBudget": 32000
}
}
});
let result = wrap_request(&body, "test-proj", MODEL_GEMINI_3_FLASH_THINKING, None);
let req = result.get("request").unwrap();
let gen_config = req.get("generationConfig").unwrap();
let budget = gen_config["thinkingConfig"]["thinkingBudget"]
.as_u64()
.unwrap();
assert_eq!(budget, 24576);
let body_pro = json!({
"model": MODEL_GEMINI_3_PRO,
"generationConfig": {
"thinkingConfig": {
"includeThoughts": true,
"thinkingBudget": 32000
}
}
});
let result_pro = wrap_request(&body_pro, "test-proj", MODEL_GEMINI_3_PRO, None);
let budget_pro = result_pro["request"]["generationConfig"]["thinkingConfig"]
["thinkingBudget"]
.as_u64()
.unwrap();
assert_eq!(budget_pro, 24576);
}
#[test]
fn test_user_instruction_preservation() {
let body = json!({
"model": MODEL_GEMINI_PRO_ALIAS,
"systemInstruction": {
"role": "user",
"parts": [{"text": "User custom prompt"}]
}
});
let result = wrap_request(&body, "test-proj", MODEL_GEMINI_PRO_ALIAS, None);
let sys = result
.get("request")
.unwrap()
.get("systemInstruction")
.unwrap();
let parts = sys.get("parts").unwrap().as_array().unwrap();
assert_eq!(parts.len(), 2);
assert!(parts[0]
.get("text")
.unwrap()
.as_str()
.unwrap()
.contains("You are Antigravity"));
assert_eq!(
parts[1].get("text").unwrap().as_str().unwrap(),
"User custom prompt"
);
}
#[test]
fn test_duplicate_prevention() {
let body = json!({
"model": MODEL_GEMINI_PRO_ALIAS,
"systemInstruction": {
"parts": [{"text": "You are Antigravity..."}]
}
});
let result = wrap_request(&body, "test-proj", MODEL_GEMINI_PRO_ALIAS, None);
let sys = result
.get("request")
.unwrap()
.get("systemInstruction")
.unwrap();
let parts = sys.get("parts").unwrap().as_array().unwrap();
assert_eq!(parts.len(), 1);
}
#[test]
fn test_image_generation_with_reference_images() {
let mut parts = Vec::new();
parts.push(json!({"text": "Generate a variation"}));
for _ in 0..14 {
parts.push(json!({
"inlineData": {
"mimeType": "image/jpeg",
"data": "base64data..."
}
}));
}
let body = json!({
"model": MODEL_GEMINI_3_PRO_IMAGE,
"contents": [{"parts": parts}]
});
let result = wrap_request(&body, "test-proj", MODEL_GEMINI_3_PRO_IMAGE, None);
let request = result.get("request").unwrap();
let contents = request.get("contents").unwrap().as_array().unwrap();
let result_parts = contents[0].get("parts").unwrap().as_array().unwrap();
assert_eq!(result_parts.len(), 15);
}
#[test]
fn test_gemini_pro_thinking_budget_processing() {
let _budget_guard = crate::proxy::config::lock_thinking_budget_for_test();
use crate::proxy::config::{
update_thinking_budget_config, ThinkingBudgetConfig, ThinkingBudgetMode,
};
update_thinking_budget_config(ThinkingBudgetConfig {
mode: ThinkingBudgetMode::Custom,
custom_value: 1024,
});
let body = json!({
"model": MODEL_GEMINI_3_PRO_PREVIEW,
"generationConfig": {
"thinkingConfig": {
"includeThoughts": true,
"thinkingBudget": 32000
}
}
});
let result = wrap_request(&body, "test-proj", MODEL_GEMINI_3_PRO_PREVIEW, None);
let req = result.get("request").unwrap();
let gen_config = req.get("generationConfig").unwrap();
let budget = gen_config["thinkingConfig"]["thinkingBudget"]
.as_u64()
.unwrap();
assert_eq!(
budget, 1024,
"Budget should be overridden to 1024 by custom config, proving logic execution"
);
}
#[test]
fn test_gemini_pro_auto_inject_thinking() {
let _budget_guard = crate::proxy::config::lock_thinking_budget_for_test();
crate::proxy::config::update_thinking_budget_config(
crate::proxy::config::ThinkingBudgetConfig {
mode: crate::proxy::config::ThinkingBudgetMode::Auto,
custom_value: 24576,
},
);
let body = json!({
"model": MODEL_GEMINI_3_PRO_PREVIEW,
"generationConfig": {}
});
let result = wrap_request(&body, "test-proj", MODEL_GEMINI_3_PRO_PREVIEW, None);
let req = result.get("request").unwrap();
let gen_config = req.get("generationConfig").unwrap();
assert!(
gen_config.get("thinkingConfig").is_some(),
"Should auto-inject thinkingConfig for gemini-3-pro"
);
let budget = gen_config["thinkingConfig"]["thinkingBudget"]
.as_u64()
.unwrap();
assert_eq!(budget, 24576);
}
#[test]
fn test_openai_image_params_support() {
let body_1 = json!({
"model": MODEL_GEMINI_3_PRO_IMAGE,
"size": "1920x1080",
"quality": "hd",
"prompt": "Test"
});
let result_1 = wrap_request(&body_1, "test-proj", MODEL_GEMINI_3_PRO_IMAGE, None);
let req_1 = result_1.get("request").unwrap();
let gen_config_1 = req_1.get("generationConfig").unwrap();
let image_config_1 = gen_config_1.get("imageConfig").unwrap();
assert_eq!(image_config_1["aspectRatio"], "16:9");
assert_eq!(image_config_1["imageSize"], "4K");
let body_2 = json!({
"model": MODEL_GEMINI_3_PRO_IMAGE,
"size": "1:1",
"quality": "standard",
"prompt": "Test"
});
let result_2 = wrap_request(&body_2, "test-proj", MODEL_GEMINI_3_PRO_IMAGE, None);
let req_2 = result_2.get("request").unwrap();
let image_config_2 = req_2["generationConfig"]["imageConfig"]
.as_object()
.unwrap();
assert_eq!(image_config_2["aspectRatio"], "1:1");
assert_eq!(image_config_2["imageSize"], "1K");
}
}