vtcode_core/cli/
rate_limiter.rs

1//! Rate limiter for API requests and tool calls to prevent abuse and rate limiting
2
3use anyhow::Result;
4use std::sync::atomic::{AtomicUsize, Ordering};
5use std::sync::{Arc, Mutex};
6use std::time::{Duration, Instant};
7
8/// Rate limiter to prevent API abuse and rate limiting
9#[derive(Debug)]
10pub struct RateLimiter {
11    /// Maximum requests per minute
12    requests_per_minute: usize,
13    /// Timestamp of recent requests (for sliding window)
14    request_times: Arc<Mutex<Vec<Instant>>>,
15    /// Current tool call count
16    tool_call_count: Arc<AtomicUsize>,
17    /// Maximum tool calls allowed
18    max_tool_calls: usize,
19}
20
21impl RateLimiter {
22    /// Create a new rate limiter
23    pub fn new(requests_per_minute: usize, max_tool_calls: usize) -> Self {
24        Self {
25            requests_per_minute,
26            request_times: Arc::new(Mutex::new(Vec::new())),
27            tool_call_count: Arc::new(AtomicUsize::new(0)),
28            max_tool_calls,
29        }
30    }
31
32    /// Check if we can make an API request, blocking if necessary
33    pub async fn wait_for_api_request(&self) -> Result<()> {
34        loop {
35            let wait_time = {
36                let mut request_times = self.request_times.lock().unwrap();
37
38                let now = Instant::now();
39                let one_minute_ago = now - Duration::from_secs(60);
40                request_times.retain(|&time| time > one_minute_ago);
41
42                if request_times.len() < self.requests_per_minute {
43                    request_times.push(now);
44                    return Ok(());
45                }
46
47                Duration::from_secs(60).saturating_sub(request_times[0].elapsed())
48            };
49
50            if !wait_time.is_zero() {
51                tokio::time::sleep(wait_time).await;
52            }
53        }
54    }
55
56    /// Check if we can make a tool call
57    pub fn can_make_tool_call(&self) -> bool {
58        self.tool_call_count.load(Ordering::Relaxed) < self.max_tool_calls
59    }
60
61    /// Increment the tool call count
62    pub fn increment_tool_call(&self) {
63        self.tool_call_count.fetch_add(1, Ordering::Relaxed);
64    }
65
66    /// Get the current tool call count
67    pub fn get_tool_call_count(&self) -> usize {
68        self.tool_call_count.load(Ordering::Relaxed)
69    }
70
71    /// Reset tool call count for new session
72    pub fn reset_tool_calls(&self) {
73        self.tool_call_count.store(0, Ordering::Relaxed);
74    }
75
76    /// Get the current request count in the sliding window
77    pub fn get_current_request_count(&self) -> usize {
78        let request_times = self.request_times.lock().unwrap();
79        let one_minute_ago = Instant::now() - Duration::from_secs(60);
80        request_times
81            .iter()
82            .filter(|&&time| time > one_minute_ago)
83            .count()
84    }
85}
vtcode_core/cli/rate_limiter.rs

vtcode_core/cli/
rate_limiter.rs
