fix: prevent tool_rounds cross-cascade contamination causing hangs

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
This commit is contained in:
Nikketryhard
2026-02-16 19:21:03 -06:00
parent 32f02d6456
commit ba96534ead
4 changed files with 31 additions and 14 deletions

View File

@@ -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)"
);
}

View File

@@ -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)"
);
}

View File

@@ -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()

View File

@@ -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<ToolRound> {
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<CapturedFunctionCall>) {