diff --git a/.env.example b/.env.example index ded7fa7..b2156c1 100644 --- a/.env.example +++ b/.env.example @@ -16,6 +16,8 @@ CRANK_PUBLISH_BIND=127.0.0.1 CRANK_ADMIN_BIND=0.0.0.0:3001 CRANK_MCP_BIND=0.0.0.0:3002 CRANK_MCP_REFRESH_MS=5000 +CRANK_MCP_RATE_LIMIT_RPS=60 +CRANK_MCP_RATE_LIMIT_BURST=120 CRANK_RUNTIME_MAX_CONCURRENT_UNARY=64 CRANK_RUNTIME_MAX_CONCURRENT_WINDOW=16 CRANK_RUNTIME_MAX_CONCURRENT_SESSIONS=16 diff --git a/apps/mcp-server/src/app.rs b/apps/mcp-server/src/app.rs index 09a25d5..5ea21f0 100644 --- a/apps/mcp-server/src/app.rs +++ b/apps/mcp-server/src/app.rs @@ -10,7 +10,7 @@ use axum::{ extract::{Path, State}, http::{ HeaderMap, HeaderValue, StatusCode, - header::{self, ACCEPT, AUTHORIZATION, HeaderName}, + header::{self, ACCEPT, AUTHORIZATION, HeaderName, RETRY_AFTER}, }, response::{ IntoResponse, Response, @@ -30,8 +30,8 @@ use crank_registry::{ UpdateStreamSessionStateRequest, }; use crank_runtime::{ - ResolvedAuth, RuntimeError, RuntimeExecutor, RuntimeOperation, RuntimeRequestContext, - SecretCrypto, + RateLimitRejection, RequestRateLimiter, ResolvedAuth, RuntimeError, RuntimeExecutor, + RuntimeOperation, RuntimeRequestContext, SecretCrypto, }; use futures_util::stream; use serde::{Deserialize, Serialize}; @@ -66,6 +66,7 @@ pub struct AppState { registry: PostgresRegistry, catalog: PublishedToolCatalog, runtime: RuntimeExecutor, + api_rate_limiter: RequestRateLimiter, secret_crypto: SecretCrypto, sessions: SessionStore, allowed_origins: AllowedOrigins, @@ -129,11 +130,13 @@ pub fn build_app( public_base_url: Option, secret_crypto: SecretCrypto, runtime: RuntimeExecutor, + api_rate_limiter: RequestRateLimiter, ) -> Router { let state = Arc::new(AppState { registry: registry.clone(), catalog: PublishedToolCatalog::new(registry, refresh_interval), runtime, + api_rate_limiter, secret_crypto, sessions: SessionStore::new(), allowed_origins: AllowedOrigins::new(public_base_url), @@ -272,6 +275,13 @@ async fn mcp_post( } }; + if let Err(rejection) = enforce_post_rate_limit(&state, &path, &headers) { + return with_request_id_header( + rate_limited_jsonrpc_response(&message, response_mode, &protocol_version, rejection), + &transport_request_id, + ); + } + if let Some(session_id) = headers.get(HEADER_MCP_SESSION_ID) { if let Ok(session_id) = session_id.to_str() { if let Some(session) = state.sessions.get(session_id).await { @@ -1783,6 +1793,63 @@ fn bearer_token(headers: &HeaderMap) -> Option<&str> { Some(token) } +fn enforce_post_rate_limit( + state: &Arc, + path: &AgentRoutePath, + headers: &HeaderMap, +) -> Result<(), RateLimitRejection> { + let key = rate_limit_key(path, headers); + state.api_rate_limiter.check(&key) +} + +fn rate_limit_key(path: &AgentRoutePath, headers: &HeaderMap) -> String { + if let Ok(Some(session_id)) = session_id_from_headers(headers) { + return format!( + "session:{}:{}:{}", + path.workspace_slug, path.agent_slug, session_id + ); + } + + if let Some(secret) = bearer_token(headers) { + return format!("api_key:{}", hash_access_secret(secret)); + } + + format!("workspace:{}:anonymous", path.workspace_slug) +} + +fn rate_limited_jsonrpc_response( + message: &Value, + response_mode: ResponseMode, + protocol_version: &str, + rejection: RateLimitRejection, +) -> Response { + let payload = json!({ + "jsonrpc": "2.0", + "id": request_id(message), + "error": { + "code": -32029, + "message": "request rate limit exceeded", + "data": { + "code": "request_rate_limited", + "retry_after_ms": rejection.retry_after_ms, + } + } + }); + + let mut response = transport_response( + StatusCode::TOO_MANY_REQUESTS, + payload, + response_mode, + None, + Some(protocol_version), + ); + let retry_after_seconds = rejection.retry_after_ms.div_ceil(1000); + if let Ok(value) = HeaderValue::from_str(&retry_after_seconds.to_string()) { + response.headers_mut().insert(RETRY_AFTER, value); + } + response +} + fn allows_scope(scopes: &[PlatformApiKeyScope], required_scope: PlatformApiKeyScope) -> bool { match required_scope { PlatformApiKeyScope::Read => scopes.iter().any(|scope| { diff --git a/apps/mcp-server/src/main.rs b/apps/mcp-server/src/main.rs index f0e7851..90f3af0 100644 --- a/apps/mcp-server/src/main.rs +++ b/apps/mcp-server/src/main.rs @@ -6,7 +6,9 @@ mod session; use std::{env, net::SocketAddr, time::Duration}; use crank_registry::{PostgresPoolConfig, PostgresRegistry}; -use crank_runtime::{RuntimeExecutor, RuntimeLimits, SecretCrypto}; +use crank_runtime::{ + RequestRateLimitConfig, RequestRateLimiter, RuntimeExecutor, RuntimeLimits, SecretCrypto, +}; use sqlx::postgres::PgConnectOptions; use tokio::net::TcpListener; use tracing::info; @@ -32,6 +34,7 @@ async fn main() -> Result<(), Box> { let socket_addr: SocketAddr = bind_addr.parse()?; let pool_config = PostgresPoolConfig::from_env()?; let runtime_limits = RuntimeLimits::from_env()?; + let api_rate_limit = mcp_api_rate_limit_config_from_env()?; let registry = PostgresRegistry::connect_with_options_and_pool_config( database_options_from_env()?, pool_config, @@ -39,7 +42,14 @@ async fn main() -> Result<(), Box> { .await?; let secret_crypto = SecretCrypto::new(&env::var("CRANK_MASTER_KEY")?)?; let runtime = RuntimeExecutor::with_limits(runtime_limits); - let app = build_app(registry, refresh_interval, base_url, secret_crypto, runtime); + let app = build_app( + registry, + refresh_interval, + base_url, + secret_crypto, + runtime, + RequestRateLimiter::new(api_rate_limit), + ); let listener = TcpListener::bind(socket_addr).await?; info!( @@ -47,6 +57,8 @@ async fn main() -> Result<(), Box> { runtime_max_concurrent_window = runtime_limits.max_concurrent_window, runtime_max_concurrent_sessions = runtime_limits.max_concurrent_sessions, runtime_max_concurrent_jobs = runtime_limits.max_concurrent_jobs, + mcp_rate_limit_rps = api_rate_limit.requests_per_second, + mcp_rate_limit_burst = api_rate_limit.burst, max_connections = pool_config.max_connections, min_connections = pool_config.min_connections, acquire_timeout_ms = pool_config.acquire_timeout_ms, @@ -83,6 +95,20 @@ fn database_options_from_env() -> Result Result> +{ + let requests_per_second = env::var("CRANK_MCP_RATE_LIMIT_RPS") + .ok() + .and_then(|value| value.parse::().ok()) + .unwrap_or(60); + let burst = env::var("CRANK_MCP_RATE_LIMIT_BURST") + .ok() + .and_then(|value| value.parse::().ok()) + .unwrap_or(120); + + Ok(RequestRateLimitConfig::new(requests_per_second, burst)?) +} + #[cfg(test)] mod tests { use std::{ @@ -112,7 +138,9 @@ mod tests { CreateAgentRequest, CreatePlatformApiKeyRequest, ListInvocationLogsQuery, PostgresRegistry, PublishAgentRequest, PublishRequest, }; - use crank_runtime::{RuntimeExecutor, SecretCrypto}; + use crank_runtime::{ + RequestRateLimitConfig, RequestRateLimiter, RuntimeExecutor, SecretCrypto, + }; use crank_schema::{Schema, SchemaKind}; use futures_util::stream; use serde_json::{Value, json}; @@ -173,6 +201,20 @@ mod tests { registry: PostgresRegistry, refresh_interval: Duration, public_base_url: Option, + ) -> axum::Router { + build_test_app_with_rate_limit( + registry, + refresh_interval, + public_base_url, + RequestRateLimitConfig::new(10_000, 10_000).unwrap(), + ) + } + + fn build_test_app_with_rate_limit( + registry: PostgresRegistry, + refresh_interval: Duration, + public_base_url: Option, + rate_limit_config: RequestRateLimitConfig, ) -> axum::Router { build_app( registry, @@ -180,6 +222,7 @@ mod tests { public_base_url, SecretCrypto::new("test-master-key").unwrap(), RuntimeExecutor::new(), + RequestRateLimiter::new(rate_limit_config), ) } @@ -1672,6 +1715,102 @@ mod tests { assert!((1..=250).contains(&poll_after_ms)); } + #[tokio::test] + async fn rejects_rapid_initialize_requests_with_429() { + let registry = test_registry().await; + let upstream_base_url = spawn_upstream_server().await; + let operation = test_operation(&upstream_base_url, "crm_rate_limited"); + + registry + .create_operation(&test_workspace_id(), &operation, Some("alice")) + .await + .unwrap(); + registry + .publish_operation(PublishRequest { + workspace_id: &test_workspace_id(), + operation_id: &operation.id, + version: 1, + published_at: &OffsetDateTime::parse("2026-03-26T10:00:00Z", &Rfc3339).unwrap(), + published_by: Some("alice"), + }) + .await + .unwrap(); + publish_agent_for_operation(®istry, &operation, "sales-rate-limited").await; + let api_key = create_platform_api_key( + ®istry, + "mcp-rate-limit", + &[PlatformApiKeyScope::Read, PlatformApiKeyScope::Write], + ) + .await; + + let base_url = spawn_mcp_server(build_test_app_with_rate_limit( + registry, + Duration::from_millis(0), + Some("https://crank.example.com".to_owned()), + RequestRateLimitConfig::new(1, 1).unwrap(), + )) + .await; + let client = reqwest::Client::new(); + let mcp_url = agent_mcp_url(&base_url, "sales-rate-limited"); + + let first_response = client + .post(&mcp_url) + .header(header::ACCEPT, "application/json, text/event-stream") + .header(header::AUTHORIZATION, format!("Bearer {api_key}")) + .header("x-request-id", "req_rate_limit_01") + .json(&json!({ + "jsonrpc": "2.0", + "id": 1, + "method": "initialize", + "params": { + "protocolVersion": "2025-11-25" + } + })) + .send() + .await + .unwrap(); + assert_eq!(first_response.status(), reqwest::StatusCode::OK); + + let second_response = client + .post(&mcp_url) + .header(header::ACCEPT, "application/json, text/event-stream") + .header(header::AUTHORIZATION, format!("Bearer {api_key}")) + .header("x-request-id", "req_rate_limit_02") + .json(&json!({ + "jsonrpc": "2.0", + "id": 2, + "method": "initialize", + "params": { + "protocolVersion": "2025-11-25" + } + })) + .send() + .await + .unwrap(); + + assert_eq!( + second_response.status(), + reqwest::StatusCode::TOO_MANY_REQUESTS + ); + assert_eq!( + second_response.headers()["x-request-id"].to_str().unwrap(), + "req_rate_limit_02" + ); + assert_eq!( + second_response.headers()["retry-after"].to_str().unwrap(), + "1" + ); + + let payload = second_response.json::().await.unwrap(); + assert_eq!(payload["error"]["code"], json!(-32029)); + assert_eq!( + payload["error"]["data"]["code"], + json!("request_rate_limited") + ); + let retry_after_ms = payload["error"]["data"]["retry_after_ms"].as_u64().unwrap(); + assert!((1..=1000).contains(&retry_after_ms)); + } + #[tokio::test] async fn rejects_cross_agent_async_job_access() { let registry = test_registry().await; diff --git a/crates/crank-runtime/src/lib.rs b/crates/crank-runtime/src/lib.rs index 383752c..463ea99 100644 --- a/crates/crank-runtime/src/lib.rs +++ b/crates/crank-runtime/src/lib.rs @@ -4,6 +4,7 @@ mod error; mod executor; mod limits; mod model; +mod rate_limit; mod redaction; mod request_context; mod secret_crypto; @@ -14,6 +15,9 @@ pub use error::RuntimeError; pub use executor::RuntimeExecutor; pub use limits::{RuntimeLimits, RuntimeLimitsConfigError}; pub use model::{AdapterResponse, PreparedRequest, RuntimeOperation}; +pub use rate_limit::{ + RateLimitRejection, RequestRateLimitConfig, RequestRateLimitConfigError, RequestRateLimiter, +}; pub use request_context::RuntimeRequestContext; pub use secret_crypto::SecretCrypto; pub use streaming::WindowExecutionResult; diff --git a/crates/crank-runtime/src/rate_limit.rs b/crates/crank-runtime/src/rate_limit.rs new file mode 100644 index 0000000..a411942 --- /dev/null +++ b/crates/crank-runtime/src/rate_limit.rs @@ -0,0 +1,145 @@ +use std::{ + collections::HashMap, + sync::{Arc, Mutex}, + time::{Duration, Instant}, +}; + +use thiserror::Error; + +const STALE_KEY_TTL: Duration = Duration::from_secs(300); + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct RequestRateLimitConfig { + pub requests_per_second: u32, + pub burst: u32, +} + +impl RequestRateLimitConfig { + pub fn new(requests_per_second: u32, burst: u32) -> Result { + if requests_per_second == 0 { + return Err(RequestRateLimitConfigError::InvalidRequestsPerSecond( + requests_per_second, + )); + } + if burst == 0 { + return Err(RequestRateLimitConfigError::InvalidBurst(burst)); + } + + Ok(Self { + requests_per_second, + burst, + }) + } +} + +#[derive(Debug, Error, PartialEq, Eq)] +pub enum RequestRateLimitConfigError { + #[error("requests_per_second must be greater than zero, got {0}")] + InvalidRequestsPerSecond(u32), + #[error("burst must be greater than zero, got {0}")] + InvalidBurst(u32), +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct RateLimitRejection { + pub retry_after_ms: u64, +} + +#[derive(Clone, Debug)] +pub struct RequestRateLimiter { + config: RequestRateLimitConfig, + states: Arc>>, +} + +#[derive(Clone, Copy, Debug)] +struct BucketState { + tokens: f64, + last_refill: Instant, + last_seen: Instant, +} + +impl RequestRateLimiter { + pub fn new(config: RequestRateLimitConfig) -> Self { + Self { + config, + states: Arc::new(Mutex::new(HashMap::new())), + } + } + + pub fn check(&self, key: &str) -> Result<(), RateLimitRejection> { + self.check_at(key, Instant::now()) + } + + pub fn check_at(&self, key: &str, now: Instant) -> Result<(), RateLimitRejection> { + let mut states = self + .states + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + states.retain(|_, state| now.duration_since(state.last_seen) <= STALE_KEY_TTL); + + let burst = self.config.burst as f64; + let requests_per_second = self.config.requests_per_second as f64; + let state = states.entry(key.to_owned()).or_insert(BucketState { + tokens: burst, + last_refill: now, + last_seen: now, + }); + + let elapsed = now.duration_since(state.last_refill).as_secs_f64(); + state.tokens = (state.tokens + elapsed * requests_per_second).min(burst); + state.last_refill = now; + state.last_seen = now; + + if state.tokens >= 1.0 { + state.tokens -= 1.0; + return Ok(()); + } + + let missing_tokens = (1.0 - state.tokens).max(0.0); + let retry_after_ms = ((missing_tokens / requests_per_second) * 1000.0) + .ceil() + .max(1.0) as u64; + Err(RateLimitRejection { retry_after_ms }) + } +} + +#[cfg(test)] +mod tests { + use std::time::{Duration, Instant}; + + use super::{RequestRateLimitConfig, RequestRateLimitConfigError, RequestRateLimiter}; + + #[test] + fn rejects_zero_requests_per_second() { + assert_eq!( + RequestRateLimitConfig::new(0, 1), + Err(RequestRateLimitConfigError::InvalidRequestsPerSecond(0)) + ); + } + + #[test] + fn rejects_zero_burst() { + assert_eq!( + RequestRateLimitConfig::new(1, 0), + Err(RequestRateLimitConfigError::InvalidBurst(0)) + ); + } + + #[test] + fn allows_burst_then_rejects_until_refilled() { + let limiter = RequestRateLimiter::new(RequestRateLimitConfig::new(2, 2).unwrap()); + let start = Instant::now(); + + assert!(limiter.check_at("key", start).is_ok()); + assert!(limiter.check_at("key", start).is_ok()); + + let rejection = limiter.check_at("key", start).unwrap_err(); + assert_eq!(rejection.retry_after_ms, 500); + + assert!( + limiter + .check_at("key", start + Duration::from_millis(500)) + .is_ok() + ); + } +} diff --git a/docs/runtime-config.md b/docs/runtime-config.md index 111aa45..866a22b 100644 --- a/docs/runtime-config.md +++ b/docs/runtime-config.md @@ -73,6 +73,8 @@ var/crank/ - `CRANK_ADMIN_BIND` - `CRANK_MCP_BIND` - `CRANK_MCP_REFRESH_MS` +- `CRANK_MCP_RATE_LIMIT_RPS` +- `CRANK_MCP_RATE_LIMIT_BURST` - `CRANK_RUNTIME_MAX_CONCURRENT_UNARY` - `CRANK_RUNTIME_MAX_CONCURRENT_WINDOW` - `CRANK_RUNTIME_MAX_CONCURRENT_SESSIONS` @@ -89,6 +91,8 @@ var/crank/ Стартовое значение для refresh published tools: - `CRANK_MCP_REFRESH_MS=5000` +- `CRANK_MCP_RATE_LIMIT_RPS=60` +- `CRANK_MCP_RATE_LIMIT_BURST=120` Стартовые значения для runtime concurrency limits: @@ -198,6 +202,11 @@ runtime env-переменной. Это значит, что: - `CRANK_RUNTIME_MAX_CONCURRENT_SESSIONS=16` - `CRANK_RUNTIME_MAX_CONCURRENT_JOBS=16` +Для MCP transport ingress throttling используются явные defaults: + +- `CRANK_MCP_RATE_LIMIT_RPS=60` +- `CRANK_MCP_RATE_LIMIT_BURST=120` + `CRANK_DATABASE_URL` допускается только как backward-compatible fallback для локальных тестов и переходного периода, но не как основная deployment-модель.