bsodfather/crank
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

mcp: add ingress request throttling

Browse Source
This commit is contained in:
a.tolmachev
2026-05-01 16:57:54 +00:00
parent 66f28defe5
commit 411d662676
6 changed files with 372 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files
apps/mcp-server/src/app.rs
+70 -3
View File
@@ -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| {
apps/mcp-server/src/main.rs
+142 -3
View File
@@ -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;