fix: bypass LS entirely when custom tools are active

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

View File

@@ -88,6 +88,13 @@ pub struct MitmStore {
call_id_to_name: Arc<RwLock<HashMap<String, String>>>,
/// Last captured function calls (for conversation history rewriting).
last_function_calls: Arc<RwLock<Vec<CapturedFunctionCall>>>,
// ── Direct response capture (bypasses LS) ────────────────────────────
/// Captured response text from MITM when custom tools are active.
/// The completions handler reads this instead of polling LS steps.
captured_response_text: Arc<RwLock<Option<String>>>,
/// Whether the captured response is complete (finishReason received).
response_complete: Arc<AtomicBool>,
}
/// Aggregate statistics across all intercepted traffic.
@@ -126,6 +133,8 @@ impl MitmStore {
pending_tool_results: Arc::new(RwLock::new(Vec::new())),
call_id_to_name: Arc::new(RwLock::new(HashMap::new())),
last_function_calls: Arc::new(RwLock::new(Vec::new())),
captured_response_text: Arc::new(RwLock::new(None)),
response_complete: Arc::new(AtomicBool::new(false)),
}
}
@@ -354,4 +363,56 @@ impl MitmStore {
pub async fn get_last_function_calls(&self) -> Vec<CapturedFunctionCall> {
self.last_function_calls.read().await.clone()
}
// ── Direct response capture (bypass LS) ──────────────────────────────
/// Append text to the captured response.
pub async fn append_response_text(&self, text: &str) {
let mut resp = self.captured_response_text.write().await;
if let Some(ref mut existing) = *resp {
existing.push_str(text);
} else {
*resp = Some(text.to_string());
}
}
/// Set (replace) the captured response text.
pub async fn set_response_text(&self, text: &str) {
*self.captured_response_text.write().await = Some(text.to_string());
}
/// Take the captured response text (consumes it).
pub async fn take_response_text(&self) -> Option<String> {
self.captured_response_text.write().await.take()
}
/// Peek at the captured response text without consuming it.
pub async fn peek_response_text(&self) -> Option<String> {
self.captured_response_text.read().await.clone()
}
/// Mark the response as complete.
pub fn mark_response_complete(&self) {
self.response_complete.store(true, Ordering::SeqCst);
}
/// Check if the response is complete.
pub fn is_response_complete(&self) -> bool {
self.response_complete.load(Ordering::SeqCst)
}
/// Clear captured response state (call at start of new request).
pub fn clear_response(&self) {
self.response_complete.store(false, Ordering::SeqCst);
// Can't use async in sync fn, so we spawn a task... or just use try_write
if let Ok(mut resp) = self.captured_response_text.try_write() {
*resp = None;
}
}
/// Async version of clear_response.
pub async fn clear_response_async(&self) {
self.response_complete.store(false, Ordering::SeqCst);
*self.captured_response_text.write().await = None;
}
}