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

fix: forge dummy STOP response to LS on functionCall capture

Browse Source 
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
2026-02-15 00:15:00 -06:00
parent 19ff784cae
commit 7c44729ace
2 changed files with 256 additions and 11 deletions
Download Patch File Download Diff File Expand all files Collapse all files

74
src/mitm/proxy.rs
View File

@@ -735,15 +735,54 @@ async fn handle_http_over_tls(
// Save body for usage parsing
response_body_buf.extend_from_slice(&header_buf[hdr_end..]);
// Forward to client immediately
if let Err(e) = client.write_all(&header_buf).await {
warn!(error = %e, "MITM: write to client failed");
break;
}
// Parse ORIGINAL initial body for MITM interception
let mut has_function_call = false;
if is_streaming_response && hdr_end < header_buf.len() {
let body = String::from_utf8_lossy(&header_buf[hdr_end..]);
parse_streaming_chunk(&body, &mut streaming_acc);
has_function_call = body.contains("functionCall");
}
// If we detected a functionCall AND custom tools are active,
// forge a dummy "STOP" response for the LS so it doesn't
// freak out and retry. The real function call data is already
// captured in MitmStore.
if has_function_call && modify_requests && store.get_tools().await.is_some() {
info!("MITM: functionCall detected → sending dummy STOP response to LS");
// Build a clean SSE response the LS will accept
let dummy_json = serde_json::json!({
"response": {
"candidates": [{
"content": {
"role": "model",
"parts": [{"text": "Tool call completed. Awaiting external tool result."}]
},
"finishReason": "STOP"
}],
"modelVersion": "gemini-3-flash"
},
"metadata": {}
});
let dummy_data = format!("data: {}\r\n\r\n", serde_json::to_string(&dummy_json).unwrap());
let dummy_chunk = format!("{:x}\r\n{}\r\n0\r\n\r\n", dummy_data.len(), dummy_data);
// Send headers (from original response) + dummy body
let headers_only = &header_buf[..hdr_end];
if let Err(e) = client.write_all(headers_only).await {
warn!(error = %e, "MITM: write headers failed");
}
if let Err(e) = client.write_all(dummy_chunk.as_bytes()).await {
warn!(error = %e, "MITM: write dummy body failed");
}
// Done — don't forward the real response
break;
}
// Normal path: forward headers+body as-is
if let Err(e) = client.write_all(&header_buf).await {
warn!(error = %e, "MITM: write to client failed");
break;
}
if let Some(cl) = response_content_length {
@@ -759,17 +798,30 @@ async fn handle_http_over_tls(
continue;
}
// Forward to client immediately
// ── Response body interception ────────────────────────────────
// Parse ORIGINAL chunk for MITM interception (captures functionCalls)
let mut chunk_has_fc = false;
if is_streaming_response {
let s = String::from_utf8_lossy(chunk);
parse_streaming_chunk(&s, &mut streaming_acc);
chunk_has_fc = s.contains("functionCall");
}
// If functionCall in body chunk + custom tools → send dummy + stop
if chunk_has_fc && modify_requests && store.get_tools().await.is_some() {
info!("MITM: functionCall in body chunk → sending chunked terminator to LS");
// Send the chunked terminator to end the stream
let _ = client.write_all(b"0\r\n\r\n").await;
break;
}
// Normal path: forward chunk to client (LS)
if let Err(e) = client.write_all(chunk).await {
warn!(error = %e, "MITM: write to client failed");
break;
}
response_body_buf.extend_from_slice(chunk);
if is_streaming_response {
let s = String::from_utf8_lossy(chunk);
parse_streaming_chunk(&s, &mut streaming_acc);
}
if let Some(cl) = response_content_length {
if response_body_buf.len() >= cl { break; }
}