mcp: throttle transport read requests

This commit is contained in:
a.tolmachev
2026-05-02 08:33:07 +00:00
parent fdad20a8d1
commit c597557a5e
3 changed files with 77 additions and 7 deletions
+6 -7
View File
@@ -2,20 +2,19 @@
## Current
### `feat/distributed-mcp-session-store`
### `feat/runtime-rate-limiting-and-backpressure`
Status: in_progress
DoD:
- transport sessions are stored behind a shared store abstraction
- mcp-server no longer depends directly on process-local session storage implementation
- in-memory session store remains available for tests and local fallback
- production mcp-server uses a Postgres-backed transport session store
- expired transport sessions are evicted on read and no longer leak indefinitely
- runtime concurrency limits reject overload deterministically
- mcp-server transport requests are throttled on ingress across `POST`, `GET`, and `DELETE`
- streaming poll pressure is bounded for stream sessions and async jobs
- admin-api ingress requests are throttled with stable `429` behavior
## Next
- `feat/runtime-rate-limiting-and-backpressure`
- `feat/distributed-mcp-session-store`
## Backlog
+25
View File
@@ -179,6 +179,10 @@ async fn mcp_get(
return status.into_response();
}
if let Err(rejection) = enforce_transport_rate_limit(&state, &path, &headers) {
return rate_limited_status_response(rejection);
}
let session_id = match session_id_from_headers(&headers) {
Ok(Some(session_id)) => session_id,
Ok(None) => return StatusCode::BAD_REQUEST.into_response(),
@@ -223,6 +227,10 @@ async fn mcp_delete(
return status.into_response();
}
if let Err(rejection) = enforce_transport_rate_limit(&state, &path, &headers) {
return rate_limited_status_response(rejection);
}
match session_id_from_headers(&headers) {
Ok(Some(session_id)) => match state.sessions.get(&session_id).await {
Ok(Some(session))
@@ -1832,6 +1840,14 @@ fn enforce_post_rate_limit(
state: &Arc<AppState>,
path: &AgentRoutePath,
headers: &HeaderMap,
) -> Result<(), RateLimitRejection> {
enforce_transport_rate_limit(state, path, headers)
}
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)
@@ -1885,6 +1901,15 @@ fn rate_limited_jsonrpc_response(
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 allows_scope(scopes: &[PlatformApiKeyScope], required_scope: PlatformApiKeyScope) -> bool {
match required_scope {
PlatformApiKeyScope::Read => scopes.iter().any(|scope| {
+46
View File
@@ -947,6 +947,52 @@ mod tests {
assert_eq!(response.status(), reqwest::StatusCode::NOT_FOUND);
}
#[tokio::test]
async fn rejects_rapid_transport_get_requests_with_429() {
let registry = test_registry().await;
publish_agent_with_bindings(&registry, "sales-get-rate-limit", vec![]).await;
let api_key = create_platform_api_key(
&registry,
"mcp-get-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, 2).unwrap(),
))
.await;
let client = reqwest::Client::new();
let mcp_url = agent_mcp_url(&base_url, "sales-get-rate-limit");
let initialized_session = initialize_session(&client, &mcp_url, &api_key).await;
let first_response = client
.get(&mcp_url)
.header(header::ACCEPT, "text/event-stream")
.header(header::AUTHORIZATION, format!("Bearer {api_key}"))
.header("MCP-Session-Id", &initialized_session)
.header("MCP-Protocol-Version", "2025-11-25")
.send()
.await
.unwrap();
assert_eq!(first_response.status(), reqwest::StatusCode::OK);
let second_response = client
.get(&mcp_url)
.header(header::ACCEPT, "text/event-stream")
.header(header::AUTHORIZATION, format!("Bearer {api_key}"))
.header("MCP-Session-Id", &initialized_session)
.header("MCP-Protocol-Version", "2025-11-25")
.send()
.await
.unwrap();
assert_eq!(second_response.status(), reqwest::StatusCode::TOO_MANY_REQUESTS);
assert!(second_response.headers().get(header::RETRY_AFTER).is_some());
}
#[tokio::test]
async fn get_requires_session_header() {
let registry = test_registry().await;