diff --git a/TASKS.md b/TASKS.md index 0900790..ba233f1 100644 --- a/TASKS.md +++ b/TASKS.md @@ -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 diff --git a/apps/mcp-server/src/app.rs b/apps/mcp-server/src/app.rs index 2b9c2b2..367332e 100644 --- a/apps/mcp-server/src/app.rs +++ b/apps/mcp-server/src/app.rs @@ -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, path: &AgentRoutePath, headers: &HeaderMap, +) -> Result<(), RateLimitRejection> { + enforce_transport_rate_limit(state, path, headers) +} + +fn enforce_transport_rate_limit( + state: &Arc, + 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| { diff --git a/apps/mcp-server/src/main.rs b/apps/mcp-server/src/main.rs index b90a3b1..8c12a23 100644 --- a/apps/mcp-server/src/main.rs +++ b/apps/mcp-server/src/main.rs @@ -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;