feat: Implement request generation counter and state management to prevent stale data and unblock Language Server for follow-up requests.

src/mitm/store.rs

@@ -4,7 +4,7 @@
 
 use serde::{Deserialize, Serialize};
 use std::collections::HashMap;
-use std::sync::atomic::{AtomicBool, Ordering};
+use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
 use std::sync::Arc;
 use tokio::sync::RwLock;
 use tracing::{debug, info};
@@ -137,6 +137,9 @@ pub struct MitmStore {
     /// 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>,
+    /// Generation counter — incremented each time a new completions turn starts.
+    /// Used to discard stale data from leaked LS connections.
+    request_generation: Arc<AtomicU64>,
 
     // ── Tool call support ────────────────────────────────────────────────
     /// Active tool definitions (Gemini format) for MITM injection.
@@ -214,6 +217,7 @@ impl MitmStore {
             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)),
+            request_generation: Arc::new(AtomicU64::new(0)),
             active_tools: Arc::new(RwLock::new(None)),
             active_tool_config: Arc::new(RwLock::new(None)),
             pending_tool_results: Arc::new(RwLock::new(Vec::new())),
@@ -483,17 +487,22 @@ impl MitmStore {
         self.response_complete.load(Ordering::SeqCst)
     }
 
-    /// Async version of clear_response.
+    /// Async version of clear_response. Bumps generation counter.
     pub async fn clear_response_async(&self) {
         self.response_complete.store(false, Ordering::SeqCst);
         self.request_in_flight.store(false, Ordering::SeqCst);
+        self.request_generation.fetch_add(1, 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);
+    /// Atomically try to mark request as in-flight.
+    /// Returns true if this caller won the race (was first to set it).
+    /// Returns false if already in-flight (someone else set it first).
+    pub fn try_mark_request_in_flight(&self) -> bool {
+        self.request_in_flight
+            .compare_exchange(false, true, Ordering::SeqCst, Ordering::SeqCst)
+            .is_ok()
     }
 
     /// Check if a request is currently in-flight.
@@ -501,6 +510,26 @@ impl MitmStore {
         self.request_in_flight.load(Ordering::SeqCst)
     }
 
+    /// Clear the in-flight flag so the LS can make follow-up requests.
+    pub fn clear_request_in_flight(&self) {
+        self.request_in_flight.store(false, Ordering::SeqCst);
+    }
+
+    /// Reset response_complete so we can wait for the next response.
+    pub fn clear_response_complete(&self) {
+        self.response_complete.store(false, Ordering::SeqCst);
+    }
+
+    /// Get current generation number.
+    pub fn current_generation(&self) -> u64 {
+        self.request_generation.load(Ordering::SeqCst)
+    }
+
+    /// Bump generation counter (invalidates all pending data from old generation).
+    pub fn bump_generation(&self) -> u64 {
+        self.request_generation.fetch_add(1, Ordering::SeqCst) + 1
+    }
+
     // ── Thinking text capture ────────────────────────────────────────────
 
     /// Set (replace) the captured thinking text.