fix: block ALL LS follow-up requests, deduplicate function calls
- Add request_in_flight flag to MitmStore, set immediately when first LLM request is forwarded with custom tools active - Block ALL subsequent LS requests (agentic loop + internal flash-lite) with fake SSE responses instead of waiting for response_complete - Fix function call deduplication: drain() accumulator after storing to prevent 3x duplicate tool calls across SSE chunks - Clear all stale state (response, thinking, function calls, errors) at the start of each streaming request - Handle response_complete with no content (thoughtSignature-only) gracefully with timeout instead of infinite hang
This commit is contained in:
@@ -130,6 +130,13 @@ pub struct MitmStore {
|
||||
/// Simple flag: set when a functionCall is captured, cleared when consumed.
|
||||
/// Used to block follow-up requests regardless of cascade identification.
|
||||
has_active_function_call: Arc<AtomicBool>,
|
||||
/// Persistent flag: set when a function call is captured, cleared ONLY when
|
||||
/// a tool result is submitted. Prevents the LS from making follow-up API
|
||||
/// calls during the entire tool execution cycle.
|
||||
awaiting_tool_result: Arc<AtomicBool>,
|
||||
/// Set when the MITM forwards the first LLM request with custom tools.
|
||||
/// Blocks ALL subsequent LS requests until the API handler clears it.
|
||||
request_in_flight: Arc<AtomicBool>,
|
||||
|
||||
// ── Tool call support ────────────────────────────────────────────────
|
||||
/// Active tool definitions (Gemini format) for MITM injection.
|
||||
@@ -205,6 +212,8 @@ impl MitmStore {
|
||||
stats: Arc::new(RwLock::new(MitmStats::default())),
|
||||
pending_function_calls: Arc::new(RwLock::new(HashMap::new())),
|
||||
has_active_function_call: Arc::new(AtomicBool::new(false)),
|
||||
awaiting_tool_result: Arc::new(AtomicBool::new(false)),
|
||||
request_in_flight: Arc::new(AtomicBool::new(false)),
|
||||
active_tools: Arc::new(RwLock::new(None)),
|
||||
active_tool_config: Arc::new(RwLock::new(None)),
|
||||
pending_tool_results: Arc::new(RwLock::new(Vec::new())),
|
||||
@@ -343,6 +352,7 @@ impl MitmStore {
|
||||
let mut pending = self.pending_function_calls.write().await;
|
||||
pending.entry(key).or_default().push(fc);
|
||||
self.has_active_function_call.store(true, Ordering::SeqCst);
|
||||
self.awaiting_tool_result.store(true, Ordering::SeqCst);
|
||||
}
|
||||
|
||||
/// Check if there's an active (unclaimed) function call.
|
||||
@@ -355,6 +365,18 @@ impl MitmStore {
|
||||
self.has_active_function_call.store(false, Ordering::SeqCst);
|
||||
}
|
||||
|
||||
/// Check if we're awaiting a tool result (blocks LS follow-up requests).
|
||||
/// This persists across function call consumption — only cleared when
|
||||
/// actual tool results are submitted.
|
||||
pub fn is_awaiting_tool_result(&self) -> bool {
|
||||
self.awaiting_tool_result.load(Ordering::SeqCst)
|
||||
}
|
||||
|
||||
/// Clear the awaiting-tool-result flag (called when tool results arrive).
|
||||
pub fn clear_awaiting_tool_result(&self) {
|
||||
self.awaiting_tool_result.store(false, Ordering::SeqCst);
|
||||
}
|
||||
|
||||
|
||||
/// Take any pending function calls (ignoring cascade ID).
|
||||
pub async fn take_any_function_calls(&self) -> Option<Vec<CapturedFunctionCall>> {
|
||||
@@ -467,10 +489,21 @@ impl MitmStore {
|
||||
/// Async version of clear_response.
|
||||
pub async fn clear_response_async(&self) {
|
||||
self.response_complete.store(false, Ordering::SeqCst);
|
||||
self.request_in_flight.store(false, Ordering::SeqCst);
|
||||
*self.captured_response_text.write().await = None;
|
||||
*self.captured_thinking_text.write().await = None;
|
||||
}
|
||||
|
||||
/// Mark the request as in-flight (first LLM request forwarded).
|
||||
pub fn mark_request_in_flight(&self) {
|
||||
self.request_in_flight.store(true, Ordering::SeqCst);
|
||||
}
|
||||
|
||||
/// Check if a request is currently in-flight.
|
||||
pub fn is_request_in_flight(&self) -> bool {
|
||||
self.request_in_flight.load(Ordering::SeqCst)
|
||||
}
|
||||
|
||||
// ── Thinking text capture ────────────────────────────────────────────
|
||||
|
||||
/// Set (replace) the captured thinking text.
|
||||
|
||||
Reference in New Issue
Block a user