Extract MCP rate limiting module

crates/crank-community-mcp/src/app.rs

@@ -8,7 +8,7 @@ use std::{
 use axum::{
     Json, Router,
     extract::{Path, State},
-    http::{HeaderMap, HeaderValue, StatusCode, header::RETRY_AFTER},
+    http::{HeaderMap, StatusCode},
     response::{IntoResponse, Response, sse::Event},
     routing::get,
 };
@@ -18,8 +18,8 @@ use crank_core::{
 };
 use crank_registry::{CreateInvocationLogRequest, PostgresRegistry, PublishedAgentTool};
 use crank_runtime::{
-    RateLimitRejection, RequestRateLimiter, ResolvedAuth, RuntimeError, RuntimeExecutionRequest,
+    RequestRateLimiter, ResolvedAuth, RuntimeError, RuntimeExecutionRequest, RuntimeExecutor,
-    RuntimeExecutor, RuntimeOperation, RuntimeRequestContext, SecretCrypto,
+    RuntimeOperation, RuntimeRequestContext, SecretCrypto,
 };
 use futures_util::stream;
 use serde::{Deserialize, Serialize};
@@ -29,8 +29,8 @@ use tracing::info;
 
 use crate::{
     access::{
-        bearer_token, credential_allows_security_level, hash_access_secret, require_machine_access,
+        credential_allows_security_level, require_machine_access, serialize_machine_access_mode,
-        serialize_machine_access_mode, serialize_security_level,
+        serialize_security_level,
     },
     auth::{SharedMachineCredentialVerifier, VerifiedMachineCredential},
     catalog::PublishedToolCatalog,
@@ -39,6 +39,10 @@ use crate::{
         jsonrpc_result, method_name, negotiated_protocol_version, params, request_id,
     },
     manifest::tool_definitions,
+    rate_limit::{
+        enforce_post_rate_limit, enforce_transport_rate_limit, rate_limited_jsonrpc_response,
+        rate_limited_status_response,
+    },
     session::{SessionState, SharedSessionStore},
     tool_error::{
         ToolErrorContract, generic_tool_error_contract, runtime_error_code,
@@ -59,7 +63,7 @@ pub(super) struct AppState {
     pub(super) registry: PostgresRegistry,
     catalog: PublishedToolCatalog,
     runtime: RuntimeExecutor,
-    api_rate_limiter: RequestRateLimiter,
+    pub(super) api_rate_limiter: RequestRateLimiter,
     secret_crypto: SecretCrypto,
     sessions: SharedSessionStore,
     pub(super) credential_verifier: SharedMachineCredentialVerifier,
@@ -810,80 +814,6 @@ async fn require_initialized_session(
     Ok(session)
 }
 
-async fn enforce_post_rate_limit(
-    state: &Arc<AppState>,
-    path: &AgentRoutePath,
-    headers: &HeaderMap,
-) -> Result<(), RateLimitRejection> {
-    enforce_transport_rate_limit(state, path, headers).await
-}
-
-async fn enforce_transport_rate_limit(
-    state: &Arc<AppState>,
-    path: &AgentRoutePath,
-    headers: &HeaderMap,
-) -> Result<(), RateLimitRejection> {
-    let key = rate_limit_key(path, headers);
-    state.api_rate_limiter.check(&key).await
-}
-
-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 rate_limited_status_response(rejection: RateLimitRejection) -> Response {
-    let mut response = StatusCode::TOO_MANY_REQUESTS.into_response();
-    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 internal_jsonrpc_error(message: &Value, error: impl std::fmt::Display) -> Response {
     transport_response(
         StatusCode::INTERNAL_SERVER_ERROR,

crates/crank-community-mcp/src/lib.rs

@@ -4,6 +4,7 @@ pub mod auth;
 pub mod catalog;
 pub mod jsonrpc;
 pub mod manifest;
+mod rate_limit;
 pub mod session;
 pub mod tool_error;
 mod transport;

crates/crank-community-mcp/src/rate_limit.rs

@@ -0,0 +1,90 @@
+use std::sync::Arc;
+
+use axum::{
+    http::{HeaderMap, HeaderValue, StatusCode, header::RETRY_AFTER},
+    response::{IntoResponse, Response},
+};
+use crank_runtime::RateLimitRejection;
+use serde_json::{Value, json};
+
+use crate::{
+    access::{bearer_token, hash_access_secret},
+    app::{AgentRoutePath, AppState},
+    jsonrpc::request_id,
+    transport::{ResponseMode, session_id_from_headers, transport_response},
+};
+
+pub(super) async fn enforce_post_rate_limit(
+    state: &Arc<AppState>,
+    path: &AgentRoutePath,
+    headers: &HeaderMap,
+) -> Result<(), RateLimitRejection> {
+    enforce_transport_rate_limit(state, path, headers).await
+}
+
+pub(super) async fn enforce_transport_rate_limit(
+    state: &Arc<AppState>,
+    path: &AgentRoutePath,
+    headers: &HeaderMap,
+) -> Result<(), RateLimitRejection> {
+    let key = rate_limit_key(path, headers);
+    state.api_rate_limiter.check(&key).await
+}
+
+pub(super) 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),
+    );
+    attach_retry_after_header(&mut response, rejection.retry_after_ms);
+    response
+}
+
+pub(super) fn rate_limited_status_response(rejection: RateLimitRejection) -> Response {
+    let mut response = StatusCode::TOO_MANY_REQUESTS.into_response();
+    attach_retry_after_header(&mut response, rejection.retry_after_ms);
+    response
+}
+
+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 attach_retry_after_header(response: &mut Response, retry_after_ms: u64) {
+    let retry_after_seconds = 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);
+    }
+}