From 32f02d6456e804b02cbed5e7b382a731ce0e89e3 Mon Sep 17 00:00:00 2001 From: Nikketryhard Date: Mon, 16 Feb 2026 19:11:38 -0600 Subject: [PATCH] fix: extend multi-round tool history to responses and gemini endpoints MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - proxy.rs: push_tool_round_calls alongside set_last_function_calls when Google responds with functionCall — accumulates rounds - responses.rs: attach_tool_round_results to pair tool results with the correct round instead of flat add_tool_result - gemini.rs: same attach_tool_round_results integration - store.rs: add push_tool_round_calls and attach_tool_round_results methods for cross-request round accumulation - Legacy add_tool_result kept for backward compat alongside new path --- src/api/gemini.rs | 18 +++++++++++++++--- src/api/responses.rs | 18 +++++++++++++++--- src/mitm/proxy.rs | 2 ++ src/mitm/store.rs | 28 ++++++++++++++++++++++++++++ 4 files changed, 60 insertions(+), 6 deletions(-) diff --git a/src/api/gemini.rs b/src/api/gemini.rs index 5eea228..d6a80d6 100644 --- a/src/api/gemini.rs +++ b/src/api/gemini.rs @@ -209,22 +209,34 @@ pub(crate) async fn handle_gemini( // Handle tool results (Gemini format: functionResponse) if let Some(ref results) = body.tool_results { + let mut pending: Vec = Vec::new(); for r in results { if let Some(fr) = r.get("functionResponse") { let name = fr["name"].as_str().unwrap_or("unknown").to_string(); let response = fr.get("response").cloned().unwrap_or(serde_json::json!({})); + // Legacy compat state .mitm_store .add_tool_result(PendingToolResult { - name, - result: response, + name: name.clone(), + result: response.clone(), }) .await; + pending.push(PendingToolResult { + name, + result: response, + }); } } + if !pending.is_empty() { + state + .mitm_store + .attach_tool_round_results(pending) + .await; + } info!( count = results.len(), - "Stored Gemini-native tool results for MITM injection" + "Stored Gemini-native tool results for MITM injection (attached to tool round)" ); } diff --git a/src/api/responses.rs b/src/api/responses.rs index ac62828..99d7edd 100644 --- a/src/api/responses.rs +++ b/src/api/responses.rs @@ -242,6 +242,7 @@ pub(crate) async fn handle_responses( // Handle tool result submission (function_call_output in input) let is_tool_result_turn = !tool_results.is_empty(); if is_tool_result_turn { + let mut pending: Vec = Vec::new(); for tr in &tool_results { // Look up function name from call_id let name = state @@ -254,17 +255,28 @@ pub(crate) async fn handle_responses( let result_value = serde_json::from_str::(&tr.output) .unwrap_or_else(|_| serde_json::json!({"result": tr.output})); + // Also store as pending (legacy compat) state .mitm_store .add_tool_result(PendingToolResult { - name, - result: result_value, + name: name.clone(), + result: result_value.clone(), }) .await; + + pending.push(PendingToolResult { + name, + result: result_value, + }); } + // Attach results to the latest open ToolRound (pushed by proxy.rs) + state + .mitm_store + .attach_tool_round_results(pending) + .await; info!( count = tool_results.len(), - "Stored tool results for MITM injection" + "Stored tool results for MITM injection (attached to tool round)" ); } diff --git a/src/mitm/proxy.rs b/src/mitm/proxy.rs index 93bcedf..f8fb6ed 100644 --- a/src/mitm/proxy.rs +++ b/src/mitm/proxy.rs @@ -826,6 +826,7 @@ 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() @@ -902,6 +903,7 @@ 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 467fd15..bcc8a55 100644 --- a/src/mitm/store.rs +++ b/src/mitm/store.rs @@ -537,6 +537,34 @@ impl MitmStore { std::mem::take(&mut *self.tool_rounds.write().await) } + /// 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) { + if !calls.is_empty() { + self.tool_rounds.write().await.push(ToolRound { + calls, + results: Vec::new(), + }); + } + } + + /// Attach tool results to the latest incomplete tool round (one with empty results). + /// Called by responses.rs/gemini.rs when the client sends tool results. + /// If there's no open round, creates a legacy round with no calls. + pub async fn attach_tool_round_results(&self, results: Vec) { + let mut rounds = self.tool_rounds.write().await; + // Find the last round that has no results yet + if let Some(round) = rounds.iter_mut().rev().find(|r| r.results.is_empty()) { + round.results = results; + } else { + // No open round — probably a race or legacy path, create standalone + rounds.push(ToolRound { + calls: Vec::new(), + results, + }); + } + } + // ── Direct response capture (bypass LS) ────────────────────────────── /// Set (replace) the captured response text.