mcp: throttle transport read requests
This commit is contained in:
@@ -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
|
||||
|
||||
|
||||
@@ -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| {
|
||||
|
||||
@@ -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(®istry, "sales-get-rate-limit", vec![]).await;
|
||||
let api_key = create_platform_api_key(
|
||||
®istry,
|
||||
"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;
|
||||
|
||||
Reference in New Issue
Block a user