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

fix: bypass LS entirely when custom tools are active

Browse Source 
When custom tools are set, don't forward ANY response from Google
to the LS. Instead, capture text and function calls directly into
MitmStore. The completions handler reads from MitmStore.

This eliminates the LS multi-turn loop (5 requests, 30+ seconds)
that occurred because the LS kept processing responses internally.
Tool calls now return in ~1.3s instead of timing out.
This commit is contained in:
Nikketryhard
2026-02-15 00:54:40 -06:00
parent ec1c0c700d
commit 50b53097bc
3 changed files with 229 additions and 36 deletions
Download Patch File Download Diff File Expand all files Collapse all files

94
src/mitm/proxy.rs
View File

@@ -737,6 +737,8 @@ async fn handle_http_over_tls(
// Parse ORIGINAL initial body for MITM interception
let mut has_function_call = false;
let bypass_ls = modify_requests && store.get_tools().await.is_some();
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);
@@ -750,41 +752,38 @@ async fn handle_http_over_tls(
store.set_last_function_calls(streaming_acc.function_calls.clone()).await;
info!("MITM: stored {} function call(s) from initial body", streaming_acc.function_calls.len());
}
// Capture response text directly into MitmStore
if bypass_ls && !streaming_acc.response_text.is_empty() {
store.set_response_text(&streaming_acc.response_text).await;
}
if bypass_ls && streaming_acc.is_complete {
store.mark_response_complete();
}
}
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 bypass_ls {
if has_function_call {
info!("MITM: functionCall captured → NOT forwarding to LS (bypass mode)");
store.mark_response_complete();
break;
}
if let Err(e) = client.write_all(dummy_chunk.as_bytes()).await {
warn!(error = %e, "MITM: write dummy body failed");
// Don't forward to LS — just continue reading chunks
// Send headers only so upstream doesn't close
if let Some(cl) = response_content_length {
if response_body_buf.len() >= cl {
store.mark_response_complete();
break;
}
}
// Done — don't forward the real response
break;
if is_chunked && has_chunked_terminator(&response_body_buf) {
store.mark_response_complete();
break;
}
continue;
}
// Normal path: forward headers+body as-is
// Normal path (no custom tools): forward headers+body as-is
if let Err(e) = client.write_all(&header_buf).await {
warn!(error = %e, "MITM: write to client failed");
break;
@@ -804,14 +803,15 @@ async fn handle_http_over_tls(
}
// ── Response body interception ────────────────────────────────
// Parse ORIGINAL chunk for MITM interception (captures functionCalls)
let mut chunk_has_fc = false;
let bypass_ls = modify_requests && store.get_tools().await.is_some();
if is_streaming_response {
let s = String::from_utf8_lossy(chunk);
parse_streaming_chunk(&s, &mut streaming_acc);
chunk_has_fc = !streaming_acc.function_calls.is_empty();
// Immediately store captured function calls — don't wait for loop end
// Immediately store captured function calls
if chunk_has_fc {
for fc in &streaming_acc.function_calls {
store.record_function_call(cascade_hint.as_deref(), fc.clone()).await;
@@ -819,13 +819,35 @@ async fn handle_http_over_tls(
store.set_last_function_calls(streaming_acc.function_calls.clone()).await;
info!("MITM: stored {} function call(s) from body chunk", streaming_acc.function_calls.len());
}
// Capture response text directly into MitmStore
if bypass_ls && !streaming_acc.response_text.is_empty() {
store.set_response_text(&streaming_acc.response_text).await;
}
if bypass_ls && streaming_acc.is_complete {
store.mark_response_complete();
}
}
// If functionCall detected + 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");
let _ = client.write_all(b"0\r\n\r\n").await;
break;
if bypass_ls {
if chunk_has_fc || streaming_acc.is_complete {
info!("MITM: response captured → NOT forwarding to LS (bypass mode)");
store.mark_response_complete();
break;
}
// Keep reading chunks without forwarding to LS
response_body_buf.extend_from_slice(chunk);
if let Some(cl) = response_content_length {
if response_body_buf.len() >= cl {
store.mark_response_complete();
break;
}
}
if is_chunked && has_chunked_terminator(&response_body_buf) {
store.mark_response_complete();
break;
}
continue;
}
// Normal path: forward chunk to client (LS)