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:
@@ -229,14 +229,19 @@ pub(crate) async fn handle_gemini(
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
if !pending.is_empty() {
|
if !pending.is_empty() {
|
||||||
state
|
// Build a ToolRound from captured function calls + client results.
|
||||||
.mitm_store
|
// Accumulate with existing rounds for multi-round history rewriting.
|
||||||
.attach_tool_round_results(pending)
|
let last_calls = state.mitm_store.get_last_function_calls().await;
|
||||||
.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!(
|
info!(
|
||||||
count = results.len(),
|
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)"
|
||||||
);
|
);
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
@@ -269,14 +269,19 @@ pub(crate) async fn handle_responses(
|
|||||||
result: result_value,
|
result: result_value,
|
||||||
});
|
});
|
||||||
}
|
}
|
||||||
// Attach results to the latest open ToolRound (pushed by proxy.rs)
|
// Build a ToolRound from the MITM-captured function calls + client results.
|
||||||
state
|
// get_last_function_calls() has the calls from Google's previous response.
|
||||||
.mitm_store
|
// We take existing accumulated rounds and append this new round.
|
||||||
.attach_tool_round_results(pending)
|
let last_calls = state.mitm_store.get_last_function_calls().await;
|
||||||
.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!(
|
info!(
|
||||||
count = tool_results.len(),
|
count = tool_results.len(),
|
||||||
"Stored tool results for MITM injection (attached to tool round)"
|
"Stored tool results for MITM injection (built tool round)"
|
||||||
);
|
);
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
@@ -602,7 +602,7 @@ async fn handle_http_over_tls(
|
|||||||
let last_calls = store.get_last_function_calls().await;
|
let last_calls = store.get_last_function_calls().await;
|
||||||
let generation_params = store.get_generation_params().await;
|
let generation_params = store.get_generation_params().await;
|
||||||
let pending_image = store.take_pending_image().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()
|
let tool_ctx = if tools.is_some()
|
||||||
|| !pending_results.is_empty()
|
|| !pending_results.is_empty()
|
||||||
@@ -826,7 +826,6 @@ async fn handle_http_over_tls(
|
|||||||
.await;
|
.await;
|
||||||
}
|
}
|
||||||
store.set_last_function_calls(calls.clone()).await;
|
store.set_last_function_calls(calls.clone()).await;
|
||||||
store.push_tool_round_calls(calls.clone()).await;
|
|
||||||
info!(
|
info!(
|
||||||
"MITM: stored {} function call(s) from initial body",
|
"MITM: stored {} function call(s) from initial body",
|
||||||
calls.len()
|
calls.len()
|
||||||
@@ -903,7 +902,6 @@ async fn handle_http_over_tls(
|
|||||||
.await;
|
.await;
|
||||||
}
|
}
|
||||||
store.set_last_function_calls(calls.clone()).await;
|
store.set_last_function_calls(calls.clone()).await;
|
||||||
store.push_tool_round_calls(calls.clone()).await;
|
|
||||||
info!(
|
info!(
|
||||||
"MITM: stored {} function call(s) from body chunk",
|
"MITM: stored {} function call(s) from body chunk",
|
||||||
calls.len()
|
calls.len()
|
||||||
|
|||||||
@@ -484,6 +484,8 @@ impl MitmStore {
|
|||||||
pub async fn clear_tools(&self) {
|
pub async fn clear_tools(&self) {
|
||||||
*self.active_tools.write().await = None;
|
*self.active_tools.write().await = None;
|
||||||
*self.active_tool_config.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).
|
/// Set active tool config (Gemini toolConfig format).
|
||||||
@@ -537,6 +539,13 @@ impl MitmStore {
|
|||||||
std::mem::take(&mut *self.tool_rounds.write().await)
|
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).
|
/// Push a new tool round from Google's response (calls only, results empty).
|
||||||
/// Called by proxy.rs when the MITM intercepts functionCall parts.
|
/// Called by proxy.rs when the MITM intercepts functionCall parts.
|
||||||
pub async fn push_tool_round_calls(&self, calls: Vec<CapturedFunctionCall>) {
|
pub async fn push_tool_round_calls(&self, calls: Vec<CapturedFunctionCall>) {
|
||||||
|
|||||||
Reference in New Issue
Block a user