b3nw/zerogravity
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix: block ALL LS follow-up requests across connections

Browse Source 
Move the in-flight blocking check to the top of the LLM request flow,
BEFORE request modification. This catches follow-ups on ALL connections
(the LS opens multiple parallel TLS connections). Only the very first
modified request reaches Google — all others get fake STOP responses.

Previously, each new connection independently allowed one request
through before blocking, letting 4-5 requests leak per turn.
This commit is contained in:
Nikketryhard
2026-02-16 00:57:33 -06:00
parent a8f3c8915f
commit 3fdd0368a0
23 changed files with 992 additions and 568 deletions
Download Patch File Download Diff File Expand all files Collapse all files

163
src/mitm/modify.rs
View File

@@ -68,14 +68,14 @@ pub fn modify_request(body: &[u8], tool_ctx: Option<&ToolContext>) -> Option<Vec
"system instruction: keep <identity> only ({original_len} → {} chars, -{stripped})",
new_sys.len()
));
json["request"]["systemInstruction"]["parts"][0]["text"] =
Value::String(new_sys);
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(String::new());
changes.push(format!(
"system instruction: cleared ({original_len} chars)"
));
json["request"]["systemInstruction"]["parts"][0]["text"] = Value::String(String::new());
}
}
@@ -125,7 +125,11 @@ pub fn modify_request(body: &[u8], tool_ctx: Option<&ToolContext>) -> Option<Vec
let mut modified = text.clone();
// Strip conversation summaries block
if let Some(cleaned) = strip_between(&modified, "# Conversation History\n", "</conversation_summaries>") {
if let Some(cleaned) = strip_between(
&modified,
"# Conversation History\n",
"</conversation_summaries>",
) {
modified = cleaned;
}
@@ -147,7 +151,9 @@ pub fn modify_request(body: &[u8], tool_ctx: Option<&ToolContext>) -> Option<Vec
}
// Strip knowledge item blocks
if let Some(cleaned) = strip_between(&modified, "Here are the ", "</knowledge_item>") {
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;
@@ -202,7 +208,8 @@ pub fn modify_request(body: &[u8], tool_ctx: Option<&ToolContext>) -> Option<Vec
// Inject client-provided tools from ToolContext
if let Some(ref ctx) = tool_ctx {
if let Some(ref custom_tools) = ctx.tools {
let total_decls: usize = custom_tools.iter()
let total_decls: usize = custom_tools
.iter()
.filter_map(|t| t.get("functionDeclarations").and_then(|d| d.as_array()))
.map(|a| a.len())
.sum();
@@ -210,7 +217,10 @@ pub fn modify_request(body: &[u8], tool_ctx: Option<&ToolContext>) -> Option<Vec
tools.push(tool.clone());
}
has_custom_tools = true;
changes.push(format!("inject {} custom tool group(s)", custom_tools.len()));
changes.push(format!(
"inject {} custom tool group(s)",
custom_tools.len()
));
// Override LS's VALIDATED toolConfig → AUTO for custom tools.
// VALIDATED mode forces Google to validate function calls against a
@@ -218,16 +228,20 @@ pub fn modify_request(body: &[u8], tool_ctx: Option<&ToolContext>) -> Option<Vec
// that list, so they'd be rejected. AUTO lets the model freely choose
// between text and function calls.
if let Some(req) = json.get_mut("request").and_then(|v| v.as_object_mut()) {
let has_validated = req.get("toolConfig")
let has_validated = req
.get("toolConfig")
.and_then(|tc| tc.pointer("/functionCallingConfig/mode"))
.and_then(|m| m.as_str())
.map_or(false, |m| m == "VALIDATED");
if has_validated {
req.insert("toolConfig".to_string(), serde_json::json!({
"functionCallingConfig": {
"mode": "AUTO"
}
}));
req.insert(
"toolConfig".to_string(),
serde_json::json!({
"functionCallingConfig": {
"mode": "AUTO"
}
}),
);
changes.push("override toolConfig VALIDATED → AUTO".to_string());
}
}
@@ -243,7 +257,11 @@ pub fn modify_request(body: &[u8], tool_ctx: Option<&ToolContext>) -> Option<Vec
if STRIP_ALL_TOOLS && !has_custom_tools {
if let Some(req) = json.get_mut("request").and_then(|v| v.as_object_mut()) {
// Remove the empty tools array entirely
if req.get("tools").and_then(|v| v.as_array()).map_or(false, |a| a.is_empty()) {
if req
.get("tools")
.and_then(|v| v.as_array())
.map_or(false, |a| a.is_empty())
{
req.remove("tools");
changes.push("remove empty tools array".to_string());
}
@@ -266,7 +284,8 @@ pub fn modify_request(body: &[u8], tool_ctx: Option<&ToolContext>) -> Option<Vec
.as_ref()
.and_then(|ctx| ctx.tools.as_ref())
.map(|tools| {
tools.iter()
tools
.iter()
.filter_map(|t| t["functionDeclarations"].as_array())
.flatten()
.filter_map(|decl| decl["name"].as_str().map(|s| s.to_string()))
@@ -309,7 +328,9 @@ pub fn modify_request(body: &[u8], tool_ctx: Option<&ToolContext>) -> Option<Vec
.map_or(true, |parts| !parts.is_empty())
});
if stripped_fc > 0 {
changes.push(format!("strip {stripped_fc} functionCall/Response parts from history"));
changes.push(format!(
"strip {stripped_fc} functionCall/Response parts from history"
));
}
}
}
@@ -336,16 +357,22 @@ pub fn modify_request(body: &[u8], tool_ctx: Option<&ToolContext>) -> Option<Vec
for msg in contents.iter_mut() {
if msg["role"].as_str() == Some("model") {
if let Some(text) = msg["parts"][0]["text"].as_str() {
if text.contains("Tool call completed") || text.contains("Awaiting external tool result") {
if text.contains("Tool call completed")
|| text.contains("Awaiting external tool result")
{
// Replace with functionCall parts
let fc_parts: Vec<Value> = ctx.last_calls.iter().map(|fc| {
serde_json::json!({
"functionCall": {
"name": fc.name,
"args": fc.args,
}
let fc_parts: Vec<Value> = ctx
.last_calls
.iter()
.map(|fc| {
serde_json::json!({
"functionCall": {
"name": fc.name,
"args": fc.args,
}
})
})
}).collect();
.collect();
msg["parts"] = Value::Array(fc_parts);
changes.push("rewrite model turn with functionCall".to_string());
break;
@@ -355,29 +382,36 @@ pub fn modify_request(body: &[u8], tool_ctx: Option<&ToolContext>) -> Option<Vec
}
// Add functionResponse as a user turn before the last user message
let fn_response_parts: Vec<Value> = ctx.pending_results.iter().map(|r| {
serde_json::json!({
"functionResponse": {
"name": r.name,
"response": r.result,
}
let fn_response_parts: Vec<Value> = ctx
.pending_results
.iter()
.map(|r| {
serde_json::json!({
"functionResponse": {
"name": r.name,
"response": r.result,
}
})
})
}).collect();
.collect();
let fn_response_turn = serde_json::json!({
"role": "user",
"parts": fn_response_parts,
});
// Insert before the last user message
let last_user_idx = contents.iter().rposition(|msg| {
msg["role"].as_str() == Some("user")
});
let last_user_idx = contents
.iter()
.rposition(|msg| msg["role"].as_str() == Some("user"));
if let Some(idx) = last_user_idx {
contents.insert(idx, fn_response_turn);
} else {
contents.push(fn_response_turn);
}
changes.push(format!("inject {} functionResponse(s)", ctx.pending_results.len()));
changes.push(format!(
"inject {} functionResponse(s)",
ctx.pending_results.len()
));
}
}
}
@@ -420,8 +454,10 @@ pub fn modify_request(body: &[u8], tool_ctx: Option<&ToolContext>) -> Option<Vec
} else {
// Not wrapped in request — try top-level (public API format)
let gen_config = json.as_object_mut().and_then(|o| {
Some(o.entry("generationConfig")
.or_insert_with(|| serde_json::json!({})))
Some(
o.entry("generationConfig")
.or_insert_with(|| serde_json::json!({})),
)
});
if let Some(gc) = gen_config.and_then(|v| v.as_object_mut()) {
let thinking_config = gc
@@ -449,8 +485,10 @@ pub fn modify_request(body: &[u8], tool_ctx: Option<&ToolContext>) -> Option<Vec
if let Some(ref gp) = ctx.generation_params {
// Find or create generationConfig (same path as above)
let gc = if let Some(req) = json.get_mut("request").and_then(|v| v.as_object_mut()) {
Some(req.entry("generationConfig")
.or_insert_with(|| serde_json::json!({})))
Some(
req.entry("generationConfig")
.or_insert_with(|| serde_json::json!({})),
)
} else {
json.as_object_mut().map(|o| {
o.entry("generationConfig")
@@ -564,8 +602,6 @@ pub fn modify_request(body: &[u8], tool_ctx: Option<&ToolContext>) -> Option<Vec
changes.join(", ")
);
Some(modified_bytes)
}
@@ -832,8 +868,10 @@ mod tests {
let result: Value = serde_json::from_slice(&modified).unwrap();
// With no ToolContext, tools should be removed entirely
assert!(result["request"]["tools"].is_null() || result.pointer("/request/tools").is_none(),
"tools should be removed when no custom tools provided");
assert!(
result["request"]["tools"].is_null() || result.pointer("/request/tools").is_none(),
"tools should be removed when no custom tools provided"
);
}
#[test]
@@ -892,13 +930,23 @@ mod tests {
let contents = result["request"]["contents"].as_array().unwrap();
// 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");
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.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
@@ -921,8 +969,14 @@ mod tests {
#[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();
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");
}
}
@@ -977,7 +1031,9 @@ pub fn modify_response_chunk(chunk: &[u8]) -> Option<Vec<u8>> {
// Replace the JSON in the result string
result.replace_range(json_start..json_start + json_end, &new_json);
changed = true;
info!("MITM: rewrote functionCall in response → text placeholder for LS");
info!(
"MITM: rewrote functionCall in response → text placeholder for LS"
);
search_from = json_start + new_json.len();
continue;
}
@@ -1117,7 +1173,10 @@ fn rewrite_function_calls_in_response(json: &mut Value) -> bool {
}
// Try nested "response.candidates"
if let Some(candidates) = json.pointer_mut("/response/candidates").and_then(|v| v.as_array_mut()) {
if let Some(candidates) = json
.pointer_mut("/response/candidates")
.and_then(|v| v.as_array_mut())
{
changed |= rewrite_candidates(candidates);
}