Extract MCP rate limiting module

2026-06-21 09:48:30 +00:00
parent 5a52a7a565
commit 35841359b3
3 changed files with 101 additions and 80 deletions
@@ -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);
+    }
+}