From 0a30c95ee1693d90575ef7af3d4653f6b21601b1 Mon Sep 17 00:00:00 2001 From: Nitanshu Lokhande <56120084+nlok5923@users.noreply.github.com> Date: Sat, 21 Feb 2026 16:09:53 +0530 Subject: [PATCH] Feat: add rate limiting for built-in tools (closes #171) (#276) * feat: add rate limiting for built-in tools (closes #171) Extend the Tool trait with an optional rate_limit_config() method and wire a shared sliding-window RateLimiter into the tool execution path in worker.rs so that per-tool per-user limits are enforced at runtime. - Add ToolRateLimitConfig struct (requests_per_minute / requests_per_hour) and rate_limit_config() default method to the Tool trait - Extract shared RateLimiter from tools/wasm/ into tools/rate_limiter.rs; WASM rate_limiter.rs now re-exports from the shared module - Add RateLimited error variant to crate::error::ToolError - Register RateLimiter on ToolRegistry and check limits in execute_tool_inner - Apply conservative configs to high-impact tools: ShellTool 30 rpm / 300 rph HttpTool 30 rpm / 500 rph WriteFileTool 20 rpm / 200 rph ApplyPatchTool 20 rpm / 200 rph MemoryWriteTool 20 rpm / 200 rph CreateJobTool 5 rpm / 30 rph Co-Authored-By: Claude Sonnet 4.6 * refactor: address Gemini review comments on rate limiter - worker.rs: collapse nested if-let into a single `if let ... && let ...` (clippy::collapsible_if) - rate_limiter.rs: extract check_internal(record: bool) helper to DRY up check_and_record / check (were identical except for the increment step) - rate_limiter.rs: replace magic numbers 60 / 3600 with MINUTE_SECS / HOUR_SECS constants Co-Authored-By: Claude Sonnet 4.6 --------- Co-authored-by: Claude Sonnet 4.6 Co-authored-by: firat.sertgoz --- src/agent/worker.rs | 18 +- src/error.rs | 6 + src/tools/builtin/file.rs | 8 + src/tools/builtin/http.rs | 4 + src/tools/builtin/job.rs | 4 + src/tools/builtin/memory.rs | 4 + src/tools/builtin/shell.rs | 4 + src/tools/mod.rs | 4 +- src/tools/rate_limiter.rs | 397 ++++++++++++++++++++++++++++++ src/tools/registry.rs | 9 + src/tools/tool.rs | 49 ++++ src/tools/wasm/capabilities.rs | 40 +--- src/tools/wasm/rate_limiter.rs | 424 +-------------------------------- 13 files changed, 514 insertions(+), 457 deletions(-) create mode 100644 src/tools/rate_limiter.rs diff --git a/src/agent/worker.rs b/src/agent/worker.rs index a8a26d87..e953f705 100644 --- a/src/agent/worker.rs +++ b/src/agent/worker.rs @@ -18,6 +18,7 @@ use crate::llm::{ }; use crate::safety::SafetyLayer; use crate::tools::ToolRegistry; +use crate::tools::rate_limiter::RateLimitResult; /// Shared dependencies for worker execution. /// @@ -439,9 +440,24 @@ Report when the job is complete or if you encounter issues you cannot resolve."# .into()); } - // Fetch job context early so we have the real user_id for hooks + // Fetch job context early so we have the real user_id for hooks and rate limiting let job_ctx = deps.context_manager.get_context(job_id).await?; + // Check per-tool rate limit before running hooks or executing (cheaper check first) + if let Some(config) = tool.rate_limit_config() + && let RateLimitResult::Limited { retry_after, .. } = deps + .tools + .rate_limiter() + .check_and_record(&job_ctx.user_id, tool_name, &config) + .await + { + return Err(crate::error::ToolError::RateLimited { + name: tool_name.to_string(), + retry_after: Some(retry_after), + } + .into()); + } + // Run BeforeToolCall hook let params = { use crate::hooks::{HookError, HookEvent, HookOutcome}; diff --git a/src/error.rs b/src/error.rs index e46b7fe7..8ff80704 100644 --- a/src/error.rs +++ b/src/error.rs @@ -202,6 +202,12 @@ pub enum ToolError { #[error("Tool {name} requires authentication")] AuthRequired { name: String }, + #[error("Tool {name} is rate limited, retry after {retry_after:?}")] + RateLimited { + name: String, + retry_after: Option, + }, + #[error("Tool builder failed: {0}")] BuilderFailed(String), } diff --git a/src/tools/builtin/file.rs b/src/tools/builtin/file.rs index aafeb409..a7ff799d 100644 --- a/src/tools/builtin/file.rs +++ b/src/tools/builtin/file.rs @@ -385,6 +385,10 @@ impl Tool for WriteFileTool { fn domain(&self) -> ToolDomain { ToolDomain::Container } + + fn rate_limit_config(&self) -> Option { + Some(crate::tools::tool::ToolRateLimitConfig::new(20, 200)) + } } /// List directory contents tool. @@ -710,6 +714,10 @@ impl Tool for ApplyPatchTool { fn domain(&self) -> ToolDomain { ToolDomain::Container } + + fn rate_limit_config(&self) -> Option { + Some(crate::tools::tool::ToolRateLimitConfig::new(20, 200)) + } } #[cfg(test)] diff --git a/src/tools/builtin/http.rs b/src/tools/builtin/http.rs index 9f591ea4..cfe3fd54 100644 --- a/src/tools/builtin/http.rs +++ b/src/tools/builtin/http.rs @@ -457,6 +457,10 @@ impl Tool for HttpTool { // Default: outbound HTTP still needs approval unless auto-approved ApprovalRequirement::UnlessAutoApproved } + + fn rate_limit_config(&self) -> Option { + Some(crate::tools::tool::ToolRateLimitConfig::new(30, 500)) + } } #[cfg(test)] diff --git a/src/tools/builtin/job.rs b/src/tools/builtin/job.rs index 45e45a94..6d1befda 100644 --- a/src/tools/builtin/job.rs +++ b/src/tools/builtin/job.rs @@ -715,6 +715,10 @@ impl Tool for CreateJobTool { } } + fn rate_limit_config(&self) -> Option { + Some(crate::tools::tool::ToolRateLimitConfig::new(5, 30)) + } + async fn execute( &self, params: serde_json::Value, diff --git a/src/tools/builtin/memory.rs b/src/tools/builtin/memory.rs index cdcf504d..ea48da70 100644 --- a/src/tools/builtin/memory.rs +++ b/src/tools/builtin/memory.rs @@ -280,6 +280,10 @@ impl Tool for MemoryWriteTool { fn requires_sanitization(&self) -> bool { false // Internal tool } + + fn rate_limit_config(&self) -> Option { + Some(crate::tools::tool::ToolRateLimitConfig::new(20, 200)) + } } /// Tool for reading workspace files. diff --git a/src/tools/builtin/shell.rs b/src/tools/builtin/shell.rs index fc1ccb70..61620fce 100644 --- a/src/tools/builtin/shell.rs +++ b/src/tools/builtin/shell.rs @@ -725,6 +725,10 @@ impl Tool for ShellTool { fn domain(&self) -> ToolDomain { ToolDomain::Container } + + fn rate_limit_config(&self) -> Option { + Some(crate::tools::tool::ToolRateLimitConfig::new(30, 300)) + } } /// Truncate output to fit within limits (UTF-8 safe). diff --git a/src/tools/mod.rs b/src/tools/mod.rs index 60614b2d..ee2ad6a3 100644 --- a/src/tools/mod.rs +++ b/src/tools/mod.rs @@ -10,6 +10,7 @@ pub mod builder; pub mod builtin; pub mod mcp; +pub mod rate_limiter; pub mod wasm; mod registry; @@ -20,5 +21,6 @@ pub use builder::{ LlmSoftwareBuilder, SoftwareBuilder, SoftwareType, Template, TemplateEngine, TemplateType, TestCase, TestHarness, TestResult, TestSuite, ValidationError, ValidationResult, WasmValidator, }; +pub use rate_limiter::RateLimiter; pub use registry::ToolRegistry; -pub use tool::{ApprovalRequirement, Tool, ToolDomain, ToolError, ToolOutput}; +pub use tool::{ApprovalRequirement, Tool, ToolDomain, ToolError, ToolOutput, ToolRateLimitConfig}; diff --git a/src/tools/rate_limiter.rs b/src/tools/rate_limiter.rs new file mode 100644 index 00000000..0d94037a --- /dev/null +++ b/src/tools/rate_limiter.rs @@ -0,0 +1,397 @@ +//! Shared rate limiter for built-in and WASM tool invocations. +//! +//! Provides per-tool, per-user rate limiting using a sliding window counter. +//! Built-in tools (shell, http, file write, etc.) are throttled here before +//! `tool.execute()` is called in the agent loop. WASM tools re-export these +//! types for HTTP-level rate limiting inside host functions. +//! +//! # Rate Limit Algorithm +//! +//! Uses a simplified sliding window counter: +//! - Track request counts for current minute and hour windows +//! - Reset counters when window expires +//! - Increment counter and check against limits +//! +//! # Persistence +//! +//! Rate limit state is in-memory only. Limits reset on process restart. +//! This is acceptable for v1; future versions may persist to the database. + +use std::collections::HashMap; +use std::time::{Duration, Instant}; + +use tokio::sync::RwLock; + +use crate::tools::tool::ToolRateLimitConfig; + +const MINUTE_SECS: u64 = 60; +const HOUR_SECS: u64 = 3600; + +/// Result of a rate limit check. +#[derive(Debug, Clone)] +pub enum RateLimitResult { + /// Request is allowed. + Allowed { + /// Remaining requests in the current minute. + remaining_minute: u32, + /// Remaining requests in the current hour. + remaining_hour: u32, + }, + /// Request is rate limited. + Limited { + /// When the rate limit will reset. + retry_after: Duration, + /// Which limit was exceeded. + limit_type: LimitType, + }, +} + +impl RateLimitResult { + pub fn is_allowed(&self) -> bool { + matches!(self, RateLimitResult::Allowed { .. }) + } +} + +/// Which rate limit was exceeded. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum LimitType { + PerMinute, + PerHour, +} + +impl std::fmt::Display for LimitType { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + LimitType::PerMinute => write!(f, "per-minute"), + LimitType::PerHour => write!(f, "per-hour"), + } + } +} + +/// State for a single rate limit window. +#[derive(Debug, Clone)] +struct WindowState { + window_start: Instant, + count: u32, +} + +impl WindowState { + fn new() -> Self { + Self { + window_start: Instant::now(), + count: 0, + } + } + + /// Check if the window has expired and reset if needed. + fn maybe_reset(&mut self, window_duration: Duration) { + if self.window_start.elapsed() >= window_duration { + self.window_start = Instant::now(); + self.count = 0; + } + } + + /// Time until window resets. + fn time_until_reset(&self, window_duration: Duration) -> Duration { + let elapsed = self.window_start.elapsed(); + if elapsed >= window_duration { + Duration::ZERO + } else { + window_duration - elapsed + } + } +} + +/// Rate limit state for a single (user, tool) pair. +#[derive(Debug)] +struct ToolRateLimitState { + minute_window: WindowState, + hour_window: WindowState, +} + +impl ToolRateLimitState { + fn new() -> Self { + Self { + minute_window: WindowState::new(), + hour_window: WindowState::new(), + } + } +} + +/// In-memory rate limiter for tool invocations. +/// +/// Keyed by `(user_id, tool_name)` so each user has independent limits. +/// Shared via `Arc` — a single instance lives in `ToolRegistry` and is +/// checked before every built-in tool execution. +pub struct RateLimiter { + state: RwLock>, +} + +impl RateLimiter { + /// Create a new rate limiter. + pub fn new() -> Self { + Self { + state: RwLock::new(HashMap::new()), + } + } + + /// Shared logic: reset windows, check limits, and optionally record the request. + async fn check_internal( + &self, + user_id: &str, + tool_name: &str, + config: &ToolRateLimitConfig, + record: bool, + ) -> RateLimitResult { + let key = (user_id.to_string(), tool_name.to_string()); + + let mut state = self.state.write().await; + let tool_state = state.entry(key).or_insert_with(ToolRateLimitState::new); + + // Reset windows if expired. + tool_state + .minute_window + .maybe_reset(Duration::from_secs(MINUTE_SECS)); + tool_state + .hour_window + .maybe_reset(Duration::from_secs(HOUR_SECS)); + + // Check minute limit. + if tool_state.minute_window.count >= config.requests_per_minute { + return RateLimitResult::Limited { + retry_after: tool_state + .minute_window + .time_until_reset(Duration::from_secs(MINUTE_SECS)), + limit_type: LimitType::PerMinute, + }; + } + + // Check hour limit. + if tool_state.hour_window.count >= config.requests_per_hour { + return RateLimitResult::Limited { + retry_after: tool_state + .hour_window + .time_until_reset(Duration::from_secs(HOUR_SECS)), + limit_type: LimitType::PerHour, + }; + } + + if record { + tool_state.minute_window.count += 1; + tool_state.hour_window.count += 1; + } + + RateLimitResult::Allowed { + remaining_minute: config.requests_per_minute - tool_state.minute_window.count, + remaining_hour: config.requests_per_hour - tool_state.hour_window.count, + } + } + + /// Check if a request is allowed and record it if so. + pub async fn check_and_record( + &self, + user_id: &str, + tool_name: &str, + config: &ToolRateLimitConfig, + ) -> RateLimitResult { + self.check_internal(user_id, tool_name, config, true).await + } + + /// Check without recording (for preview/estimation). + pub async fn check( + &self, + user_id: &str, + tool_name: &str, + config: &ToolRateLimitConfig, + ) -> RateLimitResult { + self.check_internal(user_id, tool_name, config, false).await + } + + /// Get current usage for a (user, tool) pair. + pub async fn get_usage(&self, user_id: &str, tool_name: &str) -> Option<(u32, u32)> { + let key = (user_id.to_string(), tool_name.to_string()); + let state = self.state.read().await; + state + .get(&key) + .map(|s| (s.minute_window.count, s.hour_window.count)) + } + + /// Clear rate limit state for a specific (user, tool) pair. + pub async fn clear(&self, user_id: &str, tool_name: &str) { + let key = (user_id.to_string(), tool_name.to_string()); + self.state.write().await.remove(&key); + } + + /// Clear all rate limit state. + pub async fn clear_all(&self) { + self.state.write().await.clear(); + } +} + +impl Default for RateLimiter { + fn default() -> Self { + Self::new() + } +} + +/// Error when rate limited. +#[derive(Debug, Clone, thiserror::Error)] +#[error("Rate limited ({limit_type}), retry after {retry_after:?}")] +pub struct RateLimitError { + pub retry_after: Duration, + pub limit_type: LimitType, +} + +impl From for Result<(), RateLimitError> { + fn from(result: RateLimitResult) -> Self { + match result { + RateLimitResult::Allowed { .. } => Ok(()), + RateLimitResult::Limited { + retry_after, + limit_type, + } => Err(RateLimitError { + retry_after, + limit_type, + }), + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::tools::tool::ToolRateLimitConfig; + + #[tokio::test] + async fn test_allowed_within_limits() { + let limiter = RateLimiter::new(); + let config = ToolRateLimitConfig::new(10, 100); + + let result = limiter.check_and_record("user1", "shell", &config).await; + + match result { + RateLimitResult::Allowed { + remaining_minute, + remaining_hour, + } => { + assert_eq!(remaining_minute, 9); + assert_eq!(remaining_hour, 99); + } + _ => panic!("Expected allowed"), + } + } + + #[tokio::test] + async fn test_minute_limit_exceeded() { + let limiter = RateLimiter::new(); + let config = ToolRateLimitConfig::new(2, 100); + + // Use up the minute limit + limiter.check_and_record("user1", "shell", &config).await; + limiter.check_and_record("user1", "shell", &config).await; + + // Third request should be limited + let result = limiter.check_and_record("user1", "shell", &config).await; + + match result { + RateLimitResult::Limited { + limit_type, + retry_after, + } => { + assert_eq!(limit_type, LimitType::PerMinute); + assert!(retry_after.as_secs() <= 60); + } + _ => panic!("Expected limited"), + } + } + + #[tokio::test] + async fn test_hour_limit_exceeded() { + let limiter = RateLimiter::new(); + let config = ToolRateLimitConfig::new(100, 2); + + // Use up the hour limit + limiter.check_and_record("user1", "shell", &config).await; + limiter.check_and_record("user1", "shell", &config).await; + + // Third request should be limited + let result = limiter.check_and_record("user1", "shell", &config).await; + + match result { + RateLimitResult::Limited { limit_type, .. } => { + assert_eq!(limit_type, LimitType::PerHour); + } + _ => panic!("Expected limited"), + } + } + + #[tokio::test] + async fn test_user_isolation() { + let limiter = RateLimiter::new(); + let config = ToolRateLimitConfig::new(1, 10); + + // User1 uses their limit + limiter.check_and_record("user1", "shell", &config).await; + let result1 = limiter.check_and_record("user1", "shell", &config).await; + + // User2 should still have their limit + let result2 = limiter.check_and_record("user2", "shell", &config).await; + + assert!(!result1.is_allowed()); + assert!(result2.is_allowed()); + } + + #[tokio::test] + async fn test_tool_isolation() { + let limiter = RateLimiter::new(); + let config = ToolRateLimitConfig::new(1, 10); + + // shell uses its limit + limiter.check_and_record("user1", "shell", &config).await; + let result1 = limiter.check_and_record("user1", "shell", &config).await; + + // http should still have its limit + let result2 = limiter.check_and_record("user1", "http", &config).await; + + assert!(!result1.is_allowed()); + assert!(result2.is_allowed()); + } + + #[tokio::test] + async fn test_get_usage() { + let limiter = RateLimiter::new(); + let config = ToolRateLimitConfig::new(30, 300); + + limiter.check_and_record("user1", "shell", &config).await; + limiter.check_and_record("user1", "shell", &config).await; + limiter.check_and_record("user1", "shell", &config).await; + + let usage = limiter.get_usage("user1", "shell").await; + assert_eq!(usage, Some((3, 3))); + } + + #[tokio::test] + async fn test_clear() { + let limiter = RateLimiter::new(); + let config = ToolRateLimitConfig::new(1, 10); + + limiter.check_and_record("user1", "shell", &config).await; + let result1 = limiter.check_and_record("user1", "shell", &config).await; + assert!(!result1.is_allowed()); + + limiter.clear("user1", "shell").await; + + let result2 = limiter.check_and_record("user1", "shell", &config).await; + assert!(result2.is_allowed()); + } + + #[tokio::test] + async fn test_read_only_tools_have_no_config() { + // Read-only tools return None from rate_limit_config() — + // verified in the individual tool tests, but assert the config + // type we'd use for write tools has sensible defaults here. + let write_config = ToolRateLimitConfig::new(20, 200); + assert_eq!(write_config.requests_per_minute, 20); + assert_eq!(write_config.requests_per_hour, 200); + } +} diff --git a/src/tools/registry.rs b/src/tools/registry.rs index 36d31eaa..f17a73a0 100644 --- a/src/tools/registry.rs +++ b/src/tools/registry.rs @@ -22,6 +22,7 @@ use crate::tools::builtin::{ SkillListTool, SkillRemoveTool, SkillSearchTool, TimeTool, ToolActivateTool, ToolAuthTool, ToolInstallTool, ToolListTool, ToolRemoveTool, ToolSearchTool, WriteFileTool, }; +use crate::tools::rate_limiter::RateLimiter; use crate::tools::tool::{Tool, ToolDomain}; use crate::tools::wasm::{ Capabilities, OAuthRefreshConfig, ResourceLimits, SharedCredentialRegistry, WasmError, @@ -77,6 +78,8 @@ pub struct ToolRegistry { credential_registry: Option>, /// Secrets store for credential injection (shared with HTTP tool). secrets_store: Option>, + /// Shared rate limiter for built-in tool invocations. + rate_limiter: RateLimiter, } impl ToolRegistry { @@ -87,6 +90,7 @@ impl ToolRegistry { builtin_names: RwLock::new(std::collections::HashSet::new()), credential_registry: None, secrets_store: None, + rate_limiter: RateLimiter::new(), } } @@ -106,6 +110,11 @@ impl ToolRegistry { self.credential_registry.as_ref() } + /// Get the shared rate limiter for checking built-in tool limits. + pub fn rate_limiter(&self) -> &RateLimiter { + &self.rate_limiter + } + /// Register a tool. Rejects dynamic tools that try to shadow a built-in name. pub async fn register(&self, tool: Arc) { let name = tool.name().to_string(); diff --git a/src/tools/tool.rs b/src/tools/tool.rs index 6cc0d4d2..c4e4a6b9 100644 --- a/src/tools/tool.rs +++ b/src/tools/tool.rs @@ -28,6 +28,39 @@ impl ApprovalRequirement { } } +/// Per-tool rate limit configuration for built-in tool invocations. +/// +/// Controls how many times a tool can be invoked per user, per time window. +/// Read-only tools (echo, time, json, file_read, etc.) should NOT be rate limited. +/// Write/external tools (shell, http, file_write, memory_write, create_job) should be. +#[derive(Debug, Clone)] +pub struct ToolRateLimitConfig { + /// Maximum invocations per minute. + pub requests_per_minute: u32, + /// Maximum invocations per hour. + pub requests_per_hour: u32, +} + +impl ToolRateLimitConfig { + /// Create a config with explicit limits. + pub fn new(requests_per_minute: u32, requests_per_hour: u32) -> Self { + Self { + requests_per_minute, + requests_per_hour, + } + } +} + +impl Default for ToolRateLimitConfig { + /// Default: 60 requests/minute, 1000 requests/hour (generous for WASM HTTP). + fn default() -> Self { + Self { + requests_per_minute: 60, + requests_per_hour: 1000, + } + } +} + /// Where a tool should execute: orchestrator process or inside a container. /// /// Orchestrator tools run in the main agent process (memory access, job mgmt, etc). @@ -206,6 +239,22 @@ pub trait Tool: Send + Sync { ToolDomain::Orchestrator } + /// Per-invocation rate limit for this tool. + /// + /// Return `Some(config)` to throttle how often this tool can be called per user. + /// Read-only tools (echo, time, json, file_read, memory_search, etc.) should + /// return `None`. Write/external tools (shell, http, file_write, memory_write, + /// create_job) should return sensible limits to prevent runaway agents. + /// + /// Rate limits are per-user, per-tool, and in-memory (reset on restart). + /// This is orthogonal to `requires_approval()` — a tool can be both + /// approval-gated and rate limited. Rate limit is checked first (cheaper). + /// + /// Default: `None` (no rate limiting). + fn rate_limit_config(&self) -> Option { + None + } + /// Get the tool schema for LLM function calling. fn schema(&self) -> ToolSchema { ToolSchema { diff --git a/src/tools/wasm/capabilities.rs b/src/tools/wasm/capabilities.rs index 3b43d15c..088d7e18 100644 --- a/src/tools/wasm/capabilities.rs +++ b/src/tools/wasm/capabilities.rs @@ -302,41 +302,11 @@ impl SecretsCapability { } } -/// Rate limiting configuration. -#[derive(Debug, Clone)] -pub struct RateLimitConfig { - /// Maximum requests per minute. - pub requests_per_minute: u32, - /// Maximum requests per hour. - pub requests_per_hour: u32, -} - -impl Default for RateLimitConfig { - fn default() -> Self { - Self { - requests_per_minute: 60, - requests_per_hour: 1000, - } - } -} - -impl RateLimitConfig { - /// Create a restrictive rate limit. - pub fn restrictive() -> Self { - Self { - requests_per_minute: 10, - requests_per_hour: 100, - } - } - - /// Create a permissive rate limit. - pub fn permissive() -> Self { - Self { - requests_per_minute: 120, - requests_per_hour: 5000, - } - } -} +/// Rate limiting configuration for WASM tool HTTP calls. +/// +/// Type alias for `ToolRateLimitConfig` from the shared rate limiter module. +/// WASM capabilities use it to configure per-tool HTTP request limits. +pub use crate::tools::tool::ToolRateLimitConfig as RateLimitConfig; #[cfg(test)] mod tests { diff --git a/src/tools/wasm/rate_limiter.rs b/src/tools/wasm/rate_limiter.rs index c79056c7..687c69d4 100644 --- a/src/tools/wasm/rate_limiter.rs +++ b/src/tools/wasm/rate_limiter.rs @@ -1,422 +1,6 @@ -//! Rate limiting for WASM tool operations. +//! WASM-tool rate limiting — re-exports the shared rate limiter. //! -//! Provides per-tool rate limiting for HTTP requests and tool invocations. -//! Uses a sliding window algorithm for smooth rate enforcement. -//! -//! # Rate Limit Algorithm -//! -//! Uses a simplified sliding window counter: -//! - Track request counts for current minute and hour windows -//! - Reset counters when window expires -//! - Increment counter and check against limits -//! -//! # Persistence -//! -//! Rate limit state can be persisted to PostgreSQL for cross-process -//! rate limiting (useful for distributed deployments). +//! The implementation lives in `crate::tools::rate_limiter`. WASM host +//! functions import the types from here so existing call-sites don't change. -use std::collections::HashMap; -use std::time::{Duration, Instant}; - -use tokio::sync::RwLock; - -use crate::tools::wasm::capabilities::RateLimitConfig; - -/// Result of a rate limit check. -#[derive(Debug, Clone)] -pub enum RateLimitResult { - /// Request is allowed. - Allowed { - /// Remaining requests in the current minute. - remaining_minute: u32, - /// Remaining requests in the current hour. - remaining_hour: u32, - }, - /// Request is rate limited. - Limited { - /// When the rate limit will reset. - retry_after: Duration, - /// Which limit was exceeded. - limit_type: LimitType, - }, -} - -impl RateLimitResult { - pub fn is_allowed(&self) -> bool { - matches!(self, RateLimitResult::Allowed { .. }) - } -} - -/// Which rate limit was exceeded. -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub enum LimitType { - PerMinute, - PerHour, -} - -impl std::fmt::Display for LimitType { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - match self { - LimitType::PerMinute => write!(f, "per-minute"), - LimitType::PerHour => write!(f, "per-hour"), - } - } -} - -/// State for a single rate limit window. -#[derive(Debug, Clone)] -struct WindowState { - window_start: Instant, - count: u32, -} - -impl WindowState { - fn new() -> Self { - Self { - window_start: Instant::now(), - count: 0, - } - } - - /// Check if the window has expired and reset if needed. - fn maybe_reset(&mut self, window_duration: Duration) { - if self.window_start.elapsed() >= window_duration { - self.window_start = Instant::now(); - self.count = 0; - } - } - - /// Time until window resets. - fn time_until_reset(&self, window_duration: Duration) -> Duration { - let elapsed = self.window_start.elapsed(); - if elapsed >= window_duration { - Duration::ZERO - } else { - window_duration - elapsed - } - } -} - -/// Rate limit state for a single tool. -#[derive(Debug)] -struct ToolRateLimitState { - minute_window: WindowState, - hour_window: WindowState, -} - -impl ToolRateLimitState { - fn new() -> Self { - Self { - minute_window: WindowState::new(), - hour_window: WindowState::new(), - } - } -} - -/// In-memory rate limiter for WASM tools. -pub struct RateLimiter { - /// State per (user_id, tool_name). - state: RwLock>, -} - -impl RateLimiter { - /// Create a new rate limiter. - pub fn new() -> Self { - Self { - state: RwLock::new(HashMap::new()), - } - } - - /// Check if a request is allowed and record it if so. - pub async fn check_and_record( - &self, - user_id: &str, - tool_name: &str, - config: &RateLimitConfig, - ) -> RateLimitResult { - let key = (user_id.to_string(), tool_name.to_string()); - - let mut state = self.state.write().await; - let tool_state = state.entry(key).or_insert_with(ToolRateLimitState::new); - - // Reset windows if expired - tool_state - .minute_window - .maybe_reset(Duration::from_secs(60)); - tool_state - .hour_window - .maybe_reset(Duration::from_secs(3600)); - - // Check minute limit - if tool_state.minute_window.count >= config.requests_per_minute { - return RateLimitResult::Limited { - retry_after: tool_state - .minute_window - .time_until_reset(Duration::from_secs(60)), - limit_type: LimitType::PerMinute, - }; - } - - // Check hour limit - if tool_state.hour_window.count >= config.requests_per_hour { - return RateLimitResult::Limited { - retry_after: tool_state - .hour_window - .time_until_reset(Duration::from_secs(3600)), - limit_type: LimitType::PerHour, - }; - } - - // Record the request - tool_state.minute_window.count += 1; - tool_state.hour_window.count += 1; - - RateLimitResult::Allowed { - remaining_minute: config.requests_per_minute - tool_state.minute_window.count, - remaining_hour: config.requests_per_hour - tool_state.hour_window.count, - } - } - - /// Check without recording (for preview/estimation). - pub async fn check( - &self, - user_id: &str, - tool_name: &str, - config: &RateLimitConfig, - ) -> RateLimitResult { - let key = (user_id.to_string(), tool_name.to_string()); - - let mut state = self.state.write().await; - let tool_state = state.entry(key).or_insert_with(ToolRateLimitState::new); - - // Reset windows if expired - tool_state - .minute_window - .maybe_reset(Duration::from_secs(60)); - tool_state - .hour_window - .maybe_reset(Duration::from_secs(3600)); - - // Check minute limit - if tool_state.minute_window.count >= config.requests_per_minute { - return RateLimitResult::Limited { - retry_after: tool_state - .minute_window - .time_until_reset(Duration::from_secs(60)), - limit_type: LimitType::PerMinute, - }; - } - - // Check hour limit - if tool_state.hour_window.count >= config.requests_per_hour { - return RateLimitResult::Limited { - retry_after: tool_state - .hour_window - .time_until_reset(Duration::from_secs(3600)), - limit_type: LimitType::PerHour, - }; - } - - RateLimitResult::Allowed { - remaining_minute: config.requests_per_minute - tool_state.minute_window.count, - remaining_hour: config.requests_per_hour - tool_state.hour_window.count, - } - } - - /// Get current usage for a tool. - pub async fn get_usage(&self, user_id: &str, tool_name: &str) -> Option<(u32, u32)> { - let key = (user_id.to_string(), tool_name.to_string()); - let state = self.state.read().await; - - state - .get(&key) - .map(|s| (s.minute_window.count, s.hour_window.count)) - } - - /// Clear rate limit state for a tool (for testing or manual reset). - pub async fn clear(&self, user_id: &str, tool_name: &str) { - let key = (user_id.to_string(), tool_name.to_string()); - self.state.write().await.remove(&key); - } - - /// Clear all rate limit state. - pub async fn clear_all(&self) { - self.state.write().await.clear(); - } -} - -impl Default for RateLimiter { - fn default() -> Self { - Self::new() - } -} - -/// Error when rate limited. -#[derive(Debug, Clone, thiserror::Error)] -#[error("Rate limited ({limit_type}), retry after {retry_after:?}")] -pub struct RateLimitError { - pub retry_after: Duration, - pub limit_type: LimitType, -} - -impl From for Result<(), RateLimitError> { - fn from(result: RateLimitResult) -> Self { - match result { - RateLimitResult::Allowed { .. } => Ok(()), - RateLimitResult::Limited { - retry_after, - limit_type, - } => Err(RateLimitError { - retry_after, - limit_type, - }), - } - } -} - -#[cfg(test)] -mod tests { - use crate::tools::wasm::capabilities::RateLimitConfig; - use crate::tools::wasm::rate_limiter::{LimitType, RateLimitResult, RateLimiter}; - - #[tokio::test] - async fn test_allowed_within_limits() { - let limiter = RateLimiter::new(); - let config = RateLimitConfig { - requests_per_minute: 10, - requests_per_hour: 100, - }; - - let result = limiter.check_and_record("user1", "tool1", &config).await; - - match result { - RateLimitResult::Allowed { - remaining_minute, - remaining_hour, - } => { - assert_eq!(remaining_minute, 9); - assert_eq!(remaining_hour, 99); - } - _ => panic!("Expected allowed"), - } - } - - #[tokio::test] - async fn test_minute_limit_exceeded() { - let limiter = RateLimiter::new(); - let config = RateLimitConfig { - requests_per_minute: 2, - requests_per_hour: 100, - }; - - // Use up the minute limit - limiter.check_and_record("user1", "tool1", &config).await; - limiter.check_and_record("user1", "tool1", &config).await; - - // Third request should be limited - let result = limiter.check_and_record("user1", "tool1", &config).await; - - match result { - RateLimitResult::Limited { - limit_type, - retry_after, - } => { - assert_eq!(limit_type, LimitType::PerMinute); - assert!(retry_after.as_secs() <= 60); - } - _ => panic!("Expected limited"), - } - } - - #[tokio::test] - async fn test_hour_limit_exceeded() { - let limiter = RateLimiter::new(); - let config = RateLimitConfig { - requests_per_minute: 100, // High minute limit - requests_per_hour: 2, // Low hour limit - }; - - // Use up the hour limit - limiter.check_and_record("user1", "tool1", &config).await; - limiter.check_and_record("user1", "tool1", &config).await; - - // Third request should be limited - let result = limiter.check_and_record("user1", "tool1", &config).await; - - match result { - RateLimitResult::Limited { limit_type, .. } => { - assert_eq!(limit_type, LimitType::PerHour); - } - _ => panic!("Expected limited"), - } - } - - #[tokio::test] - async fn test_user_isolation() { - let limiter = RateLimiter::new(); - let config = RateLimitConfig { - requests_per_minute: 1, - requests_per_hour: 10, - }; - - // User1 uses their limit - limiter.check_and_record("user1", "tool1", &config).await; - let result1 = limiter.check_and_record("user1", "tool1", &config).await; - - // User2 should still have their limit - let result2 = limiter.check_and_record("user2", "tool1", &config).await; - - assert!(!result1.is_allowed()); - assert!(result2.is_allowed()); - } - - #[tokio::test] - async fn test_tool_isolation() { - let limiter = RateLimiter::new(); - let config = RateLimitConfig { - requests_per_minute: 1, - requests_per_hour: 10, - }; - - // Tool1 uses its limit - limiter.check_and_record("user1", "tool1", &config).await; - let result1 = limiter.check_and_record("user1", "tool1", &config).await; - - // Tool2 should still have its limit - let result2 = limiter.check_and_record("user1", "tool2", &config).await; - - assert!(!result1.is_allowed()); - assert!(result2.is_allowed()); - } - - #[tokio::test] - async fn test_get_usage() { - let limiter = RateLimiter::new(); - let config = RateLimitConfig::default(); - - limiter.check_and_record("user1", "tool1", &config).await; - limiter.check_and_record("user1", "tool1", &config).await; - limiter.check_and_record("user1", "tool1", &config).await; - - let usage = limiter.get_usage("user1", "tool1").await; - assert_eq!(usage, Some((3, 3))); - } - - #[tokio::test] - async fn test_clear() { - let limiter = RateLimiter::new(); - let config = RateLimitConfig { - requests_per_minute: 1, - requests_per_hour: 10, - }; - - limiter.check_and_record("user1", "tool1", &config).await; - let result1 = limiter.check_and_record("user1", "tool1", &config).await; - assert!(!result1.is_allowed()); - - limiter.clear("user1", "tool1").await; - - let result2 = limiter.check_and_record("user1", "tool1", &config).await; - assert!(result2.is_allowed()); - } -} +pub use crate::tools::rate_limiter::{LimitType, RateLimitError, RateLimitResult, RateLimiter};