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:
Nikketryhard
@@ -1,34 +1,16 @@
|
||||
//! Request body modification for intercepted LLM API calls.
|
||||
//!
|
||||
//! Strips redundant/verbose sections from the Google Gemini API request
|
||||
//! to reduce token usage while keeping the request looking legitimate.
|
||||
//! Nothing structural changes — just trimming fat.
|
||||
//! Aggressively strips everything except identity and actual conversation
|
||||
//! from the Gemini API request. No integrity checks exist on the request
|
||||
//! body — Google validates OAuth, project, model, and JSON structure only.
|
||||
|
||||
use regex::Regex;
|
||||
use serde_json::Value;
|
||||
use tracing::info;
|
||||
|
||||
/// Whether to strip ALL tool definitions (default: true).
|
||||
/// The model generates responses fine without them — tools are only
|
||||
/// needed by the Antigravity webview, not by our proxy.
|
||||
/// Strip ALL tool definitions.
|
||||
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.
|
||||
/// Returns the modified JSON bytes, or None if modification wasn't possible.
|
||||
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 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
|
||||
.pointer_mut("/request/systemInstruction/parts/0/text")
|
||||
.and_then(|v| v.as_str())
|
||||
.map(|s| s.to_string())
|
||||
{
|
||||
let mut modified = sys.clone();
|
||||
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)"));
|
||||
}
|
||||
}
|
||||
let original_len = sys.len();
|
||||
|
||||
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"] =
|
||||
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
|
||||
.pointer_mut("/request/contents")
|
||||
.and_then(|v| v.as_array_mut())
|
||||
{
|
||||
let before = contents.len();
|
||||
|
||||
// Remove messages matching strip prefixes
|
||||
// Remove messages that are pure Antigravity context injection
|
||||
contents.retain(|msg| {
|
||||
if let Some(text) = msg["parts"][0]["text"].as_str() {
|
||||
for prefix in STRIP_CONTENT_PREFIXES {
|
||||
if text.starts_with(prefix) {
|
||||
return false;
|
||||
}
|
||||
// Strip user_information (OS, workspace paths)
|
||||
if text.starts_with("<user_information>") {
|
||||
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
|
||||
});
|
||||
|
||||
// Strip conversation summaries from remaining messages
|
||||
// These appear as "# Conversation History\nHere are the conversation IDs..."
|
||||
// For remaining messages, strip embedded metadata
|
||||
for msg in contents.iter_mut() {
|
||||
if let Some(text) = msg["parts"][0]["text"].as_str().map(|s| s.to_string()) {
|
||||
if let Some(start) = text.find("# Conversation History\n") {
|
||||
// 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()
|
||||
};
|
||||
let mut modified = text.clone();
|
||||
|
||||
if trimmed.len() < text.len() {
|
||||
let saved = text.len() - trimmed.len();
|
||||
changes.push(format!("strip conversation summaries ({saved} chars)"));
|
||||
msg["parts"][0]["text"] = Value::String(trimmed);
|
||||
// Strip conversation summaries block
|
||||
if let Some(cleaned) = strip_between(&modified, "# Conversation History\n", "</conversation_summaries>") {
|
||||
modified = cleaned;
|
||||
}
|
||||
|
||||
// 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();
|
||||
if removed_msgs > 0 {
|
||||
changes.push(format!("remove {removed_msgs} content messages"));
|
||||
// Remove now-empty messages
|
||||
contents.retain(|msg| {
|
||||
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 let Some(tools) = json
|
||||
.pointer_mut("/request/tools")
|
||||
@@ -151,21 +176,56 @@ pub fn modify_request(body: &[u8]) -> Option<Vec<u8>> {
|
||||
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.
|
||||
/// 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> {
|
||||
let mut result = Vec::with_capacity(data.len());
|
||||
let mut pos = 0;
|
||||
|
||||
while pos < data.len() {
|
||||
// Find end of chunk size line
|
||||
let line_end = match data[pos..].windows(2).position(|w| w == b"\r\n") {
|
||||
Some(p) => pos + p,
|
||||
None => break,
|
||||
};
|
||||
|
||||
// Parse hex chunk size (ignore chunk extensions after ';')
|
||||
let size_str = std::str::from_utf8(&data[pos..line_end])
|
||||
.unwrap_or("")
|
||||
.split(';')
|
||||
@@ -174,16 +234,15 @@ pub fn dechunk(data: &[u8]) -> Vec<u8> {
|
||||
.trim();
|
||||
|
||||
let chunk_size = match usize::from_str_radix(size_str, 16) {
|
||||
Ok(0) => break, // Terminal chunk
|
||||
Ok(0) => break,
|
||||
Ok(n) => n,
|
||||
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());
|
||||
result.extend_from_slice(&data[data_start..data_end]);
|
||||
|
||||
// Skip past data + trailing \r\n
|
||||
pos = data_end + 2;
|
||||
}
|
||||
|
||||
@@ -246,7 +305,6 @@ mod tests {
|
||||
"tools": [
|
||||
{"functionDeclarations": [{"name": "view_file", "description": "view", "parameters": {}}]},
|
||||
{"functionDeclarations": [{"name": "browser_subagent", "description": "browse", "parameters": {}}]},
|
||||
{"functionDeclarations": [{"name": "grep_search", "description": "grep", "parameters": {}}]},
|
||||
],
|
||||
"generationConfig": {}
|
||||
},
|
||||
@@ -262,8 +320,8 @@ mod tests {
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_modify_strips_system_sections() {
|
||||
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>";
|
||||
fn test_modify_keeps_only_identity() {
|
||||
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!({
|
||||
"project": "test",
|
||||
"requestId": "test/1",
|
||||
@@ -285,20 +343,24 @@ mod tests {
|
||||
.unwrap();
|
||||
|
||||
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("lots of web dev stuff"));
|
||||
assert!(!new_sys.contains("communication_style"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_modify_strips_empty_user_rules() {
|
||||
fn test_modify_strips_context_messages() {
|
||||
let body = serde_json::json!({
|
||||
"project": "test",
|
||||
"requestId": "test/1",
|
||||
"request": {
|
||||
"contents": [
|
||||
{"role": "user", "parts": [{"text": "<user_rules>\nThe user has not defined any custom rules.\n</user_rules>"}]},
|
||||
{"role": "user", "parts": [{"text": "hello world"}]},
|
||||
{"role": "user", "parts": [{"text": "<user_information>\nLinux\n</user_information>"}]},
|
||||
{"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": [],
|
||||
"generationConfig": {}
|
||||
@@ -311,7 +373,39 @@ mod tests {
|
||||
let result: Value = serde_json::from_slice(&modified).unwrap();
|
||||
|
||||
let contents = result["request"]["contents"].as_array().unwrap();
|
||||
assert_eq!(contents.len(), 1);
|
||||
assert_eq!(contents[0]["parts"][0]["text"].as_str().unwrap(), "hello world");
|
||||
// Should have removed user_information, user_rules, workflows (3 messages)
|
||||
// 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