feat: aggressive request stripping — keep only identity + conversation
Strip everything from intercepted LLM requests except: - <identity> section in system instruction - Actual conversation turns (user messages + model responses) Removed: tool_calling, web_app_dev, knowledge_discovery, persistent_context, skills, ephemeral_message, communication_style, user_information, user_rules, MEMORY, workflows, mcp_servers, conversation_summaries, ADDITIONAL_METADATA, Step Id prefixes. Expected reduction: ~92% (63KB → ~5KB for simple requests).
This commit is contained in:
@@ -1,34 +1,16 @@
|
|||||||
//! Request body modification for intercepted LLM API calls.
|
//! Request body modification for intercepted LLM API calls.
|
||||||
//!
|
//!
|
||||||
//! Strips redundant/verbose sections from the Google Gemini API request
|
//! Aggressively strips everything except identity and actual conversation
|
||||||
//! to reduce token usage while keeping the request looking legitimate.
|
//! from the Gemini API request. No integrity checks exist on the request
|
||||||
//! Nothing structural changes — just trimming fat.
|
//! body — Google validates OAuth, project, model, and JSON structure only.
|
||||||
|
|
||||||
|
use regex::Regex;
|
||||||
use serde_json::Value;
|
use serde_json::Value;
|
||||||
use tracing::info;
|
use tracing::info;
|
||||||
|
|
||||||
/// Whether to strip ALL tool definitions (default: true).
|
/// Strip ALL tool definitions.
|
||||||
/// The model generates responses fine without them — tools are only
|
|
||||||
/// needed by the Antigravity webview, not by our proxy.
|
|
||||||
const STRIP_ALL_TOOLS: bool = true;
|
const STRIP_ALL_TOOLS: bool = true;
|
||||||
|
|
||||||
/// System instruction sections to STRIP (matched by XML tag name).
|
|
||||||
/// These are verbose instructional manuals that add tokens but don't
|
|
||||||
/// meaningfully affect output quality for coding tasks.
|
|
||||||
const STRIP_SYSTEM_SECTIONS: &[&str] = &[
|
|
||||||
"web_application_development",
|
|
||||||
"knowledge_discovery",
|
|
||||||
"persistent_context",
|
|
||||||
"skills",
|
|
||||||
];
|
|
||||||
|
|
||||||
/// Content message patterns to strip entirely.
|
|
||||||
/// These appear as separate `contents[]` entries with recognizable prefixes.
|
|
||||||
const STRIP_CONTENT_PREFIXES: &[&str] = &[
|
|
||||||
"<user_rules>\nThe user has not defined any custom rules.",
|
|
||||||
"<workflows>\n",
|
|
||||||
];
|
|
||||||
|
|
||||||
/// Modify a streamGenerateContent request body in-place.
|
/// Modify a streamGenerateContent request body in-place.
|
||||||
/// Returns the modified JSON bytes, or None if modification wasn't possible.
|
/// Returns the modified JSON bytes, or None if modification wasn't possible.
|
||||||
pub fn modify_request(body: &[u8]) -> Option<Vec<u8>> {
|
pub fn modify_request(body: &[u8]) -> Option<Vec<u8>> {
|
||||||
@@ -37,83 +19,126 @@ pub fn modify_request(body: &[u8]) -> Option<Vec<u8>> {
|
|||||||
let original_size = body.len();
|
let original_size = body.len();
|
||||||
let mut changes: Vec<String> = Vec::new();
|
let mut changes: Vec<String> = Vec::new();
|
||||||
|
|
||||||
// ── 1. Strip verbose system instruction sections ──────────────────────
|
// ── 1. System instruction: keep ONLY <identity>, nuke everything else ──
|
||||||
if let Some(sys) = json
|
if let Some(sys) = json
|
||||||
.pointer_mut("/request/systemInstruction/parts/0/text")
|
.pointer_mut("/request/systemInstruction/parts/0/text")
|
||||||
.and_then(|v| v.as_str())
|
.and_then(|v| v.as_str())
|
||||||
.map(|s| s.to_string())
|
.map(|s| s.to_string())
|
||||||
{
|
{
|
||||||
let mut modified = sys.clone();
|
let original_len = sys.len();
|
||||||
for section in STRIP_SYSTEM_SECTIONS {
|
|
||||||
let pattern = format!("<{section}>");
|
|
||||||
let end_pattern = format!("</{section}>");
|
|
||||||
if let (Some(start), Some(end)) = (modified.find(&pattern), modified.find(&end_pattern))
|
|
||||||
{
|
|
||||||
let end_pos = end + end_pattern.len();
|
|
||||||
let removed = end_pos - start;
|
|
||||||
modified = format!("{}{}", &modified[..start], &modified[end_pos..]);
|
|
||||||
changes.push(format!("strip <{section}> ({removed} chars)"));
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
if modified.len() != sys.len() {
|
// Extract <identity>...</identity> block
|
||||||
|
let identity = extract_xml_section(&sys, "identity");
|
||||||
|
|
||||||
|
if let Some(identity_text) = identity {
|
||||||
|
let new_sys = format!("<identity>\n{}\n</identity>", identity_text.trim());
|
||||||
|
let stripped = original_len - new_sys.len();
|
||||||
|
if stripped > 0 {
|
||||||
|
changes.push(format!(
|
||||||
|
"system instruction: keep <identity> only ({original_len} → {} chars, -{stripped})",
|
||||||
|
new_sys.len()
|
||||||
|
));
|
||||||
|
json["request"]["systemInstruction"]["parts"][0]["text"] =
|
||||||
|
Value::String(new_sys);
|
||||||
|
}
|
||||||
|
} else {
|
||||||
|
// No identity tag found — clear the whole thing
|
||||||
|
changes.push(format!("system instruction: cleared ({original_len} chars)"));
|
||||||
json["request"]["systemInstruction"]["parts"][0]["text"] =
|
json["request"]["systemInstruction"]["parts"][0]["text"] =
|
||||||
Value::String(modified);
|
Value::String(String::new());
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
// ── 2. Strip bloated content messages ─────────────────────────────────
|
// ── 2. Content messages: keep only actual conversation turns ───────────
|
||||||
if let Some(contents) = json
|
if let Some(contents) = json
|
||||||
.pointer_mut("/request/contents")
|
.pointer_mut("/request/contents")
|
||||||
.and_then(|v| v.as_array_mut())
|
.and_then(|v| v.as_array_mut())
|
||||||
{
|
{
|
||||||
let before = contents.len();
|
let before = contents.len();
|
||||||
|
|
||||||
// Remove messages matching strip prefixes
|
// Remove messages that are pure Antigravity context injection
|
||||||
contents.retain(|msg| {
|
contents.retain(|msg| {
|
||||||
if let Some(text) = msg["parts"][0]["text"].as_str() {
|
if let Some(text) = msg["parts"][0]["text"].as_str() {
|
||||||
for prefix in STRIP_CONTENT_PREFIXES {
|
// Strip user_information (OS, workspace paths)
|
||||||
if text.starts_with(prefix) {
|
if text.starts_with("<user_information>") {
|
||||||
return false;
|
return false;
|
||||||
}
|
}
|
||||||
|
// Strip user_rules / MEMORY blocks
|
||||||
|
if text.starts_with("<user_rules>") {
|
||||||
|
return false;
|
||||||
|
}
|
||||||
|
// Strip workflows
|
||||||
|
if text.starts_with("<workflows>") {
|
||||||
|
return false;
|
||||||
|
}
|
||||||
|
// Strip MCP servers block
|
||||||
|
if text.starts_with("<mcp_servers>") {
|
||||||
|
return false;
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
true
|
true
|
||||||
});
|
});
|
||||||
|
|
||||||
// Strip conversation summaries from remaining messages
|
// For remaining messages, strip embedded metadata
|
||||||
// These appear as "# Conversation History\nHere are the conversation IDs..."
|
|
||||||
for msg in contents.iter_mut() {
|
for msg in contents.iter_mut() {
|
||||||
if let Some(text) = msg["parts"][0]["text"].as_str().map(|s| s.to_string()) {
|
if let Some(text) = msg["parts"][0]["text"].as_str().map(|s| s.to_string()) {
|
||||||
if let Some(start) = text.find("# Conversation History\n") {
|
let mut modified = text.clone();
|
||||||
// Find the end of the conversation summaries block
|
|
||||||
let end_marker = "</conversation_summaries>";
|
|
||||||
let trimmed = if let Some(end) = text.find(end_marker) {
|
|
||||||
let end_pos = end + end_marker.len();
|
|
||||||
// Find next non-whitespace after end marker
|
|
||||||
let rest = text[end_pos..].trim_start();
|
|
||||||
format!("{}{}", &text[..start], rest)
|
|
||||||
} else {
|
|
||||||
// No end marker — just cut from "# Conversation History" onward
|
|
||||||
text[..start].trim_end().to_string()
|
|
||||||
};
|
|
||||||
|
|
||||||
if trimmed.len() < text.len() {
|
// Strip conversation summaries block
|
||||||
let saved = text.len() - trimmed.len();
|
if let Some(cleaned) = strip_between(&modified, "# Conversation History\n", "</conversation_summaries>") {
|
||||||
changes.push(format!("strip conversation summaries ({saved} chars)"));
|
modified = cleaned;
|
||||||
msg["parts"][0]["text"] = Value::String(trimmed);
|
}
|
||||||
|
|
||||||
|
// Strip <ADDITIONAL_METADATA> blocks (cursor pos, open files, etc.)
|
||||||
|
if let Some(cleaned) = strip_xml_section(&modified, "ADDITIONAL_METADATA") {
|
||||||
|
modified = cleaned;
|
||||||
|
}
|
||||||
|
|
||||||
|
// Strip <EPHEMERAL_MESSAGE> blocks
|
||||||
|
if let Some(cleaned) = strip_xml_section(&modified, "EPHEMERAL_MESSAGE") {
|
||||||
|
modified = cleaned;
|
||||||
|
}
|
||||||
|
|
||||||
|
// Strip "Step Id: N\n" prefixes
|
||||||
|
if modified.starts_with("Step Id:") {
|
||||||
|
if let Some(newline_pos) = modified.find('\n') {
|
||||||
|
modified = modified[newline_pos + 1..].to_string();
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// Strip knowledge item blocks
|
||||||
|
if let Some(cleaned) = strip_between(&modified, "Here are the ", "</knowledge_item>") {
|
||||||
|
// Only strip if it's about knowledge items
|
||||||
|
if cleaned.len() < modified.len() && modified.contains("knowledge item") {
|
||||||
|
modified = cleaned;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// Clean up excessive whitespace from stripping
|
||||||
|
let modified = collapse_newlines(&modified);
|
||||||
|
|
||||||
|
if modified.len() < text.len() {
|
||||||
|
msg["parts"][0]["text"] = Value::String(modified);
|
||||||
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
let removed_msgs = before - contents.len();
|
// Remove now-empty messages
|
||||||
if removed_msgs > 0 {
|
contents.retain(|msg| {
|
||||||
changes.push(format!("remove {removed_msgs} content messages"));
|
if let Some(text) = msg["parts"][0]["text"].as_str() {
|
||||||
|
!text.trim().is_empty()
|
||||||
|
} else {
|
||||||
|
true
|
||||||
|
}
|
||||||
|
});
|
||||||
|
|
||||||
|
let removed = before - contents.len();
|
||||||
|
if removed > 0 {
|
||||||
|
changes.push(format!("remove {removed}/{before} content messages"));
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
// ── 3. Strip tool definitions ────────────────────────────────────────
|
// ── 3. Strip all tool definitions ────────────────────────────────────
|
||||||
if STRIP_ALL_TOOLS {
|
if STRIP_ALL_TOOLS {
|
||||||
if let Some(tools) = json
|
if let Some(tools) = json
|
||||||
.pointer_mut("/request/tools")
|
.pointer_mut("/request/tools")
|
||||||
@@ -151,21 +176,56 @@ pub fn modify_request(body: &[u8]) -> Option<Vec<u8>> {
|
|||||||
Some(modified_bytes)
|
Some(modified_bytes)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// Extract the inner text of an XML-style section.
|
||||||
|
fn extract_xml_section(text: &str, tag: &str) -> Option<String> {
|
||||||
|
let open = format!("<{tag}>");
|
||||||
|
let close = format!("</{tag}>");
|
||||||
|
let start = text.find(&open)?;
|
||||||
|
let end = text.find(&close)?;
|
||||||
|
let inner_start = start + open.len();
|
||||||
|
if inner_start >= end {
|
||||||
|
return None;
|
||||||
|
}
|
||||||
|
Some(text[inner_start..end].to_string())
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Strip an XML-style section and return the modified text.
|
||||||
|
fn strip_xml_section(text: &str, tag: &str) -> Option<String> {
|
||||||
|
let open = format!("<{tag}>");
|
||||||
|
let close = format!("</{tag}>");
|
||||||
|
let start = text.find(&open)?;
|
||||||
|
let end = text.find(&close)?;
|
||||||
|
let end_pos = end + close.len();
|
||||||
|
Some(format!("{}{}", &text[..start], &text[end_pos..]))
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Strip everything between two markers (inclusive of markers).
|
||||||
|
fn strip_between(text: &str, start_marker: &str, end_marker: &str) -> Option<String> {
|
||||||
|
let start = text.find(start_marker)?;
|
||||||
|
let end = text.find(end_marker)?;
|
||||||
|
let end_pos = end + end_marker.len();
|
||||||
|
// Skip any trailing whitespace after end marker
|
||||||
|
let rest = text[end_pos..].trim_start();
|
||||||
|
Some(format!("{}{}", &text[..start], rest))
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Collapse 3+ consecutive newlines into 2.
|
||||||
|
fn collapse_newlines(text: &str) -> String {
|
||||||
|
let re = Regex::new(r"\n{3,}").unwrap();
|
||||||
|
re.replace_all(text, "\n\n").to_string()
|
||||||
|
}
|
||||||
|
|
||||||
/// Dechunk an HTTP chunked-encoded body into raw bytes.
|
/// Dechunk an HTTP chunked-encoded body into raw bytes.
|
||||||
/// Input: "hex_size\r\n data\r\n hex_size\r\n data\r\n 0\r\n\r\n"
|
|
||||||
/// Output: concatenated data segments.
|
|
||||||
pub fn dechunk(data: &[u8]) -> Vec<u8> {
|
pub fn dechunk(data: &[u8]) -> Vec<u8> {
|
||||||
let mut result = Vec::with_capacity(data.len());
|
let mut result = Vec::with_capacity(data.len());
|
||||||
let mut pos = 0;
|
let mut pos = 0;
|
||||||
|
|
||||||
while pos < data.len() {
|
while pos < data.len() {
|
||||||
// Find end of chunk size line
|
|
||||||
let line_end = match data[pos..].windows(2).position(|w| w == b"\r\n") {
|
let line_end = match data[pos..].windows(2).position(|w| w == b"\r\n") {
|
||||||
Some(p) => pos + p,
|
Some(p) => pos + p,
|
||||||
None => break,
|
None => break,
|
||||||
};
|
};
|
||||||
|
|
||||||
// Parse hex chunk size (ignore chunk extensions after ';')
|
|
||||||
let size_str = std::str::from_utf8(&data[pos..line_end])
|
let size_str = std::str::from_utf8(&data[pos..line_end])
|
||||||
.unwrap_or("")
|
.unwrap_or("")
|
||||||
.split(';')
|
.split(';')
|
||||||
@@ -174,16 +234,15 @@ pub fn dechunk(data: &[u8]) -> Vec<u8> {
|
|||||||
.trim();
|
.trim();
|
||||||
|
|
||||||
let chunk_size = match usize::from_str_radix(size_str, 16) {
|
let chunk_size = match usize::from_str_radix(size_str, 16) {
|
||||||
Ok(0) => break, // Terminal chunk
|
Ok(0) => break,
|
||||||
Ok(n) => n,
|
Ok(n) => n,
|
||||||
Err(_) => break,
|
Err(_) => break,
|
||||||
};
|
};
|
||||||
|
|
||||||
let data_start = line_end + 2; // skip \r\n
|
let data_start = line_end + 2;
|
||||||
let data_end = (data_start + chunk_size).min(data.len());
|
let data_end = (data_start + chunk_size).min(data.len());
|
||||||
result.extend_from_slice(&data[data_start..data_end]);
|
result.extend_from_slice(&data[data_start..data_end]);
|
||||||
|
|
||||||
// Skip past data + trailing \r\n
|
|
||||||
pos = data_end + 2;
|
pos = data_end + 2;
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -246,7 +305,6 @@ mod tests {
|
|||||||
"tools": [
|
"tools": [
|
||||||
{"functionDeclarations": [{"name": "view_file", "description": "view", "parameters": {}}]},
|
{"functionDeclarations": [{"name": "view_file", "description": "view", "parameters": {}}]},
|
||||||
{"functionDeclarations": [{"name": "browser_subagent", "description": "browse", "parameters": {}}]},
|
{"functionDeclarations": [{"name": "browser_subagent", "description": "browse", "parameters": {}}]},
|
||||||
{"functionDeclarations": [{"name": "grep_search", "description": "grep", "parameters": {}}]},
|
|
||||||
],
|
],
|
||||||
"generationConfig": {}
|
"generationConfig": {}
|
||||||
},
|
},
|
||||||
@@ -262,8 +320,8 @@ mod tests {
|
|||||||
}
|
}
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
fn test_modify_strips_system_sections() {
|
fn test_modify_keeps_only_identity() {
|
||||||
let sys_text = "<identity>I am an AI</identity>\n<web_application_development>lots of web dev stuff here</web_application_development>\n<communication_style>be helpful</communication_style>";
|
let sys_text = "<identity>\nYou are a helpful AI.\n</identity>\n\n<tool_calling>\nUse absolute paths.\n</tool_calling>\n<web_application_development>\nlots of web dev stuff\n</web_application_development>\n<communication_style>\nbe helpful\n</communication_style>";
|
||||||
let body = serde_json::json!({
|
let body = serde_json::json!({
|
||||||
"project": "test",
|
"project": "test",
|
||||||
"requestId": "test/1",
|
"requestId": "test/1",
|
||||||
@@ -285,20 +343,24 @@ mod tests {
|
|||||||
.unwrap();
|
.unwrap();
|
||||||
|
|
||||||
assert!(new_sys.contains("<identity>"));
|
assert!(new_sys.contains("<identity>"));
|
||||||
assert!(new_sys.contains("<communication_style>"));
|
assert!(new_sys.contains("You are a helpful AI."));
|
||||||
|
assert!(!new_sys.contains("tool_calling"));
|
||||||
assert!(!new_sys.contains("web_application_development"));
|
assert!(!new_sys.contains("web_application_development"));
|
||||||
assert!(!new_sys.contains("lots of web dev stuff"));
|
assert!(!new_sys.contains("communication_style"));
|
||||||
}
|
}
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
fn test_modify_strips_empty_user_rules() {
|
fn test_modify_strips_context_messages() {
|
||||||
let body = serde_json::json!({
|
let body = serde_json::json!({
|
||||||
"project": "test",
|
"project": "test",
|
||||||
"requestId": "test/1",
|
"requestId": "test/1",
|
||||||
"request": {
|
"request": {
|
||||||
"contents": [
|
"contents": [
|
||||||
{"role": "user", "parts": [{"text": "<user_rules>\nThe user has not defined any custom rules.\n</user_rules>"}]},
|
{"role": "user", "parts": [{"text": "<user_information>\nLinux\n</user_information>"}]},
|
||||||
{"role": "user", "parts": [{"text": "hello world"}]},
|
{"role": "user", "parts": [{"text": "<user_rules>\nno rules\n</user_rules>"}]},
|
||||||
|
{"role": "user", "parts": [{"text": "<workflows>\nsome workflows\n</workflows>"}]},
|
||||||
|
{"role": "user", "parts": [{"text": "Step Id: 0\n\n<USER_REQUEST>\nSay hello\n</USER_REQUEST>\n<ADDITIONAL_METADATA>\ncursor stuff\n</ADDITIONAL_METADATA>"}]},
|
||||||
|
{"role": "model", "parts": [{"text": "Hello!"}]},
|
||||||
],
|
],
|
||||||
"tools": [],
|
"tools": [],
|
||||||
"generationConfig": {}
|
"generationConfig": {}
|
||||||
@@ -311,7 +373,39 @@ mod tests {
|
|||||||
let result: Value = serde_json::from_slice(&modified).unwrap();
|
let result: Value = serde_json::from_slice(&modified).unwrap();
|
||||||
|
|
||||||
let contents = result["request"]["contents"].as_array().unwrap();
|
let contents = result["request"]["contents"].as_array().unwrap();
|
||||||
assert_eq!(contents.len(), 1);
|
// Should have removed user_information, user_rules, workflows (3 messages)
|
||||||
assert_eq!(contents[0]["parts"][0]["text"].as_str().unwrap(), "hello world");
|
// Kept: USER_REQUEST message (with ADDITIONAL_METADATA stripped) + model response
|
||||||
|
assert_eq!(contents.len(), 2, "should keep only user request + model response");
|
||||||
|
|
||||||
|
// Check USER_REQUEST message had metadata stripped
|
||||||
|
let user_msg = contents[0]["parts"][0]["text"].as_str().unwrap();
|
||||||
|
assert!(user_msg.contains("Say hello"), "should keep user request");
|
||||||
|
assert!(!user_msg.contains("ADDITIONAL_METADATA"), "should strip metadata");
|
||||||
|
assert!(!user_msg.contains("cursor stuff"), "should strip cursor info");
|
||||||
|
assert!(!user_msg.starts_with("Step Id:"), "should strip step id");
|
||||||
|
|
||||||
|
// Model response kept intact
|
||||||
|
assert_eq!(contents[1]["parts"][0]["text"].as_str().unwrap(), "Hello!");
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_extract_xml_section() {
|
||||||
|
let text = "before <identity>\nI am AI\n</identity> after";
|
||||||
|
let result = extract_xml_section(text, "identity").unwrap();
|
||||||
|
assert_eq!(result, "\nI am AI\n");
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_strip_xml_section() {
|
||||||
|
let text = "before <META>\nstuff\n</META> after";
|
||||||
|
let result = strip_xml_section(text, "META").unwrap();
|
||||||
|
assert_eq!(result, "before after");
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_strip_between() {
|
||||||
|
let text = "keep this # Conversation History\nlots of stuff\n</conversation_summaries>\nand this";
|
||||||
|
let result = strip_between(text, "# Conversation History\n", "</conversation_summaries>").unwrap();
|
||||||
|
assert_eq!(result, "keep this and this");
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
Reference in New Issue
Block a user