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
+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| {