mcp: add ingress request throttling

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<String>,
     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<AppState>,
+    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| {