From ba96534ead23c190522ff5f6392a8862b1b97e70 Mon Sep 17 00:00:00 2001 From: Nikketryhard Date: Mon, 16 Feb 2026 19:21:03 -0600 Subject: [PATCH] fix: prevent tool_rounds cross-cascade contamination causing hangs MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Root cause: proxy.rs eagerly pushed tool rounds via push_tool_round_calls when intercepting Google's functionCall response. These stale rounds leaked into LS follow-up requests, producing malformed history that Google timed out on (60s 'no upstream response'). Changes: - Remove push_tool_round_calls from proxy.rs response interception - proxy.rs: use get_tool_rounds (non-destructive) instead of take_tool_rounds so accumulated rounds persist across multiple LS requests per cascade - responses.rs/gemini.rs: build rounds via take+push+set pattern — each handler accumulates its own rounds from get_last_function_calls + results - completions.rs: unchanged (set_tool_rounds replaces from messages) - clear_tools: also clears tool_rounds to prevent stale data between sessions - store.rs: add get_tool_rounds (non-destructive clone) method --- src/api/gemini.rs | 15 ++++++++++----- src/api/responses.rs | 17 +++++++++++------ src/mitm/proxy.rs | 4 +--- src/mitm/store.rs | 9 +++++++++ 4 files changed, 31 insertions(+), 14 deletions(-) diff --git a/src/api/gemini.rs b/src/api/gemini.rs index d6a80d6..119c5f0 100644 --- a/src/api/gemini.rs +++ b/src/api/gemini.rs @@ -229,14 +229,19 @@ pub(crate) async fn handle_gemini( } } if !pending.is_empty() { - state - .mitm_store - .attach_tool_round_results(pending) - .await; + // Build a ToolRound from captured function calls + client results. + // Accumulate with existing rounds for multi-round history rewriting. + let last_calls = state.mitm_store.get_last_function_calls().await; + let mut rounds = state.mitm_store.take_tool_rounds().await; + rounds.push(crate::mitm::store::ToolRound { + calls: last_calls, + results: pending, + }); + state.mitm_store.set_tool_rounds(rounds).await; } info!( count = results.len(), - "Stored Gemini-native tool results for MITM injection (attached to tool round)" + "Stored Gemini-native tool results for MITM injection (built tool round)" ); } diff --git a/src/api/responses.rs b/src/api/responses.rs index 99d7edd..8fb2b7e 100644 --- a/src/api/responses.rs +++ b/src/api/responses.rs @@ -269,14 +269,19 @@ pub(crate) async fn handle_responses( result: result_value, }); } - // Attach results to the latest open ToolRound (pushed by proxy.rs) - state - .mitm_store - .attach_tool_round_results(pending) - .await; + // Build a ToolRound from the MITM-captured function calls + client results. + // get_last_function_calls() has the calls from Google's previous response. + // We take existing accumulated rounds and append this new round. + let last_calls = state.mitm_store.get_last_function_calls().await; + let mut rounds = state.mitm_store.take_tool_rounds().await; + rounds.push(crate::mitm::store::ToolRound { + calls: last_calls, + results: pending, + }); + state.mitm_store.set_tool_rounds(rounds).await; info!( count = tool_results.len(), - "Stored tool results for MITM injection (attached to tool round)" + "Stored tool results for MITM injection (built tool round)" ); } diff --git a/src/mitm/proxy.rs b/src/mitm/proxy.rs index f8fb6ed..632cc55 100644 --- a/src/mitm/proxy.rs +++ b/src/mitm/proxy.rs @@ -602,7 +602,7 @@ async fn handle_http_over_tls( let last_calls = store.get_last_function_calls().await; let generation_params = store.get_generation_params().await; let pending_image = store.take_pending_image().await; - let tool_rounds = store.take_tool_rounds().await; + let tool_rounds = store.get_tool_rounds().await; let tool_ctx = if tools.is_some() || !pending_results.is_empty() @@ -826,7 +826,6 @@ async fn handle_http_over_tls( .await; } store.set_last_function_calls(calls.clone()).await; - store.push_tool_round_calls(calls.clone()).await; info!( "MITM: stored {} function call(s) from initial body", calls.len() @@ -903,7 +902,6 @@ async fn handle_http_over_tls( .await; } store.set_last_function_calls(calls.clone()).await; - store.push_tool_round_calls(calls.clone()).await; info!( "MITM: stored {} function call(s) from body chunk", calls.len() diff --git a/src/mitm/store.rs b/src/mitm/store.rs index bcc8a55..4342275 100644 --- a/src/mitm/store.rs +++ b/src/mitm/store.rs @@ -484,6 +484,8 @@ impl MitmStore { pub async fn clear_tools(&self) { *self.active_tools.write().await = None; *self.active_tool_config.write().await = None; + // Also clear accumulated tool rounds to prevent stale data + self.tool_rounds.write().await.clear(); } /// Set active tool config (Gemini toolConfig format). @@ -537,6 +539,13 @@ impl MitmStore { std::mem::take(&mut *self.tool_rounds.write().await) } + /// Get (non-destructive clone) multi-round tool call history. + /// Used by proxy.rs to read rounds without consuming them, so they + /// persist across multiple LS requests in the same cascade. + pub async fn get_tool_rounds(&self) -> Vec { + self.tool_rounds.read().await.clone() + } + /// Push a new tool round from Google's response (calls only, results empty). /// Called by proxy.rs when the MITM intercepts functionCall parts. pub async fn push_tool_round_calls(&self, calls: Vec) {