mcp: add ingress request throttling

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<dyn std::error::Error>> {
     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<dyn std::error::Error>> {
     .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<dyn std::error::Error>> {
         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<PgConnectOptions, Box<dyn std::error::E
         .password(&password))
 }
 
+fn mcp_api_rate_limit_config_from_env() -> Result<RequestRateLimitConfig, Box<dyn std::error::Error>>
+{
+    let requests_per_second = env::var("CRANK_MCP_RATE_LIMIT_RPS")
+        .ok()
+        .and_then(|value| value.parse::<u32>().ok())
+        .unwrap_or(60);
+    let burst = env::var("CRANK_MCP_RATE_LIMIT_BURST")
+        .ok()
+        .and_then(|value| value.parse::<u32>().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<String>,
+    ) -> 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<String>,
+        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(&registry, &operation, "sales-rate-limited").await;
+        let api_key = create_platform_api_key(
+            &registry,
+            "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::<Value>().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;