fix: forge dummy STOP response to LS on functionCall capture
When the MITM detects a functionCall in Google's response AND custom tools are active, send a forged clean text response to the LS instead of the real one. This prevents the LS from seeing function calls for tools it doesn't manage, eliminating the retry loop entirely. The real function call data is captured in MitmStore and returned to the client (OpenCode) through the completions handler. Also removes the complex chunked-encoding response rewriting approach in favor of this simpler forge-and-break strategy.
This commit is contained in:
Nikketryhard
@@ -694,3 +694,196 @@ mod tests {
|
||||
assert_eq!(result, "keep this and this");
|
||||
}
|
||||
}
|
||||
|
||||
// ─── Response modification ──────────────────────────────────────────────────
|
||||
|
||||
/// Rewrite an SSE response chunk to replace `functionCall` parts with text,
|
||||
/// so the LS doesn't see tool calls for tools it doesn't manage.
|
||||
///
|
||||
/// The MITM intercept layer has already captured the function call data
|
||||
/// (via `parse_streaming_chunk`) before this function runs, so we're not
|
||||
/// losing any information — just hiding it from the LS.
|
||||
///
|
||||
/// Handles HTTP chunked transfer encoding framing (size\r\n...data...\r\n).
|
||||
///
|
||||
/// Returns `Some(modified_bytes)` if the chunk was rewritten, `None` if no
|
||||
/// change was needed.
|
||||
pub fn modify_response_chunk(chunk: &[u8]) -> Option<Vec<u8>> {
|
||||
let text = std::str::from_utf8(chunk).ok()?;
|
||||
|
||||
// Quick check — no point parsing if no functionCall present
|
||||
if !text.contains("functionCall") {
|
||||
return None;
|
||||
}
|
||||
|
||||
// Strategy: find each `data: {json}` SSE event in the raw text (which may
|
||||
// be wrapped in chunked encoding). Parse the JSON, rewrite functionCall
|
||||
// parts, and rebuild the chunked frame with updated sizes.
|
||||
|
||||
// First, dechunk: extract SSE data lines from chunked encoding
|
||||
// Chunked format: <hex-size>\r\n<data>\r\n
|
||||
// We'll work on the whole text, finding "data: " prefixed JSON objects
|
||||
let mut result = text.to_string();
|
||||
let mut changed = false;
|
||||
|
||||
// Find all `data: {...}` patterns (SSE events with JSON)
|
||||
// Use a simple approach: find "data: {" and match to the end of JSON
|
||||
let mut search_from = 0;
|
||||
while let Some(data_pos) = result[search_from..].find("data: {") {
|
||||
let abs_pos = search_from + data_pos;
|
||||
let json_start = abs_pos + 6; // skip "data: "
|
||||
|
||||
// Find the end of this JSON object by finding the matching closing brace
|
||||
if let Some(json_end) = find_json_end(&result[json_start..]) {
|
||||
let json_str = &result[json_start..json_start + json_end];
|
||||
|
||||
if json_str.contains("functionCall") {
|
||||
if let Ok(mut json) = serde_json::from_str::<Value>(json_str) {
|
||||
if rewrite_function_calls_in_response(&mut json) {
|
||||
if let Ok(new_json) = serde_json::to_string(&json) {
|
||||
// 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");
|
||||
search_from = json_start + new_json.len();
|
||||
continue;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
search_from = json_start + json_end;
|
||||
} else {
|
||||
search_from = json_start;
|
||||
}
|
||||
}
|
||||
|
||||
if !changed {
|
||||
return None;
|
||||
}
|
||||
|
||||
// Rechunk: if the original was chunked, we need to recalculate chunk sizes
|
||||
// The format is: <hex-size>\r\n<payload>\r\n
|
||||
// We'll rebuild the chunked encoding from scratch
|
||||
if text.contains("\r\n") && text.chars().next().map_or(false, |c| c.is_ascii_hexdigit()) {
|
||||
// This looks like chunked encoding — rebuild it
|
||||
// Extract the payload (everything between first \r\n and last \r\n)
|
||||
let rechunked = rechunk_response(&result);
|
||||
Some(rechunked.into_bytes())
|
||||
} else {
|
||||
Some(result.into_bytes())
|
||||
}
|
||||
}
|
||||
|
||||
/// Find the end of a JSON object starting at the given string.
|
||||
/// Returns the index past the closing brace.
|
||||
fn find_json_end(s: &str) -> Option<usize> {
|
||||
let mut depth = 0i32;
|
||||
let mut in_string = false;
|
||||
let mut escape = false;
|
||||
|
||||
for (i, c) in s.char_indices() {
|
||||
if escape {
|
||||
escape = false;
|
||||
continue;
|
||||
}
|
||||
if c == '\\' && in_string {
|
||||
escape = true;
|
||||
continue;
|
||||
}
|
||||
if c == '"' {
|
||||
in_string = !in_string;
|
||||
continue;
|
||||
}
|
||||
if in_string {
|
||||
continue;
|
||||
}
|
||||
if c == '{' {
|
||||
depth += 1;
|
||||
} else if c == '}' {
|
||||
depth -= 1;
|
||||
if depth == 0 {
|
||||
return Some(i + 1);
|
||||
}
|
||||
}
|
||||
}
|
||||
None
|
||||
}
|
||||
|
||||
/// Rebuild chunked encoding from a modified response body.
|
||||
/// Takes the full text (which contains old chunk sizes) and rebuilds
|
||||
/// with correct sizes.
|
||||
fn rechunk_response(text: &str) -> String {
|
||||
// Extract the actual SSE data lines (skip chunk size lines)
|
||||
let mut payload = String::new();
|
||||
for line in text.split('\n') {
|
||||
let trimmed = line.trim_end_matches('\r');
|
||||
// Skip lines that are purely hex chunk sizes
|
||||
if trimmed.is_empty() {
|
||||
continue;
|
||||
}
|
||||
if trimmed.chars().all(|c| c.is_ascii_hexdigit()) && !trimmed.is_empty() {
|
||||
continue;
|
||||
}
|
||||
// Skip "0" (chunked terminator)
|
||||
if trimmed == "0" {
|
||||
continue;
|
||||
}
|
||||
payload.push_str(line);
|
||||
if !line.ends_with('\n') {
|
||||
payload.push('\n');
|
||||
}
|
||||
}
|
||||
|
||||
// Wrap in a single chunk
|
||||
let payload_bytes = payload.as_bytes();
|
||||
format!("{:x}\r\n{}\r\n", payload_bytes.len(), payload)
|
||||
}
|
||||
|
||||
/// Rewrite a parsed SSE JSON object: replace `functionCall` parts with
|
||||
/// text placeholder and change `finishReason` from `MALFORMED_FUNCTION_CALL`
|
||||
/// or any non-STOP reason to `STOP`.
|
||||
///
|
||||
/// Handles both Gemini public API format (`{"candidates":[...]}`) and
|
||||
/// internal LS format (`{"response":{"candidates":[...]}}`).
|
||||
fn rewrite_function_calls_in_response(json: &mut Value) -> bool {
|
||||
let mut changed = false;
|
||||
|
||||
// Helper to rewrite candidates array in-place
|
||||
fn rewrite_candidates(candidates: &mut Vec<Value>) -> bool {
|
||||
let mut changed = false;
|
||||
for candidate in candidates.iter_mut() {
|
||||
if let Some(parts) = candidate
|
||||
.pointer_mut("/content/parts")
|
||||
.and_then(|v| v.as_array_mut())
|
||||
{
|
||||
for part in parts.iter_mut() {
|
||||
if part.get("functionCall").is_some() {
|
||||
*part = serde_json::json!({
|
||||
"text": "Tool call completed. Awaiting external tool result."
|
||||
});
|
||||
changed = true;
|
||||
}
|
||||
}
|
||||
}
|
||||
if let Some(reason) = candidate.get("finishReason").and_then(|v| v.as_str()) {
|
||||
if reason != "STOP" {
|
||||
candidate["finishReason"] = Value::String("STOP".to_string());
|
||||
changed = true;
|
||||
}
|
||||
}
|
||||
}
|
||||
changed
|
||||
}
|
||||
|
||||
// Try direct "candidates" first
|
||||
if let Some(candidates) = json.get_mut("candidates").and_then(|v| v.as_array_mut()) {
|
||||
changed |= rewrite_candidates(candidates);
|
||||
}
|
||||
|
||||
// Try nested "response.candidates"
|
||||
if let Some(candidates) = json.pointer_mut("/response/candidates").and_then(|v| v.as_array_mut()) {
|
||||
changed |= rewrite_candidates(candidates);
|
||||
}
|
||||
|
||||
changed
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user