From 4cd5d5b13279333a3426c83f20f862e796f46929 Mon Sep 17 00:00:00 2001 From: "a.tolmachev" Date: Fri, 1 May 2026 16:36:50 +0000 Subject: [PATCH] mcp: throttle rapid stream session polls --- apps/mcp-server/src/app.rs | 20 ++++- apps/mcp-server/src/main.rs | 106 ++++++++++++++++++++++++ crates/crank-core/src/stream_session.rs | 57 +++++++++++++ 3 files changed, 182 insertions(+), 1 deletion(-) diff --git a/apps/mcp-server/src/app.rs b/apps/mcp-server/src/app.rs index 8380a15..6dbecdc 100644 --- a/apps/mcp-server/src/app.rs +++ b/apps/mcp-server/src/app.rs @@ -811,7 +811,7 @@ async fn handle_session_start_call( batch_size, }), expires_at, - last_poll_at: Some(now), + last_poll_at: None, created_at: now, closed_at: None, }; @@ -996,6 +996,24 @@ async fn handle_session_poll_call( ); } + let poll_after_ms = streaming.poll_interval_ms.unwrap_or(1_000); + let remaining_delay_ms = loaded.remaining_poll_delay_ms(now, poll_after_ms); + if remaining_delay_ms > 0 { + return tool_error_response( + message, + response_mode, + &session.protocol_version, + "stream_session_poll_rate_limited", + format!( + "stream session {} must wait before the next poll", + control.session_id + ), + Some(json!({ + "poll_after_ms": remaining_delay_ms + })), + ); + } + let mut state_payload: StoredSessionState = match serde_json::from_value(loaded.state.clone()) { Ok(value) => value, Err(error) => { diff --git a/apps/mcp-server/src/main.rs b/apps/mcp-server/src/main.rs index 42ddb6a..698557f 100644 --- a/apps/mcp-server/src/main.rs +++ b/apps/mcp-server/src/main.rs @@ -1455,6 +1455,112 @@ mod tests { ); } + #[tokio::test] + async fn rejects_rapid_repeat_session_poll() { + let registry = test_registry().await; + let server_addr = grpc_test_support::spawn_unary_echo_server().await; + let operation = test_grpc_session_operation(&server_addr, "echo_stream_session_rate"); + + 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(®istry, &operation, "sales-session-rate").await; + let api_key = create_platform_api_key( + ®istry, + "mcp-session-rate", + &[PlatformApiKeyScope::Read, PlatformApiKeyScope::Write], + ) + .await; + + let base_url = spawn_mcp_server(build_test_app( + registry, + Duration::from_millis(0), + Some("https://crank.example.com".to_owned()), + )) + .await; + let client = reqwest::Client::new(); + let mcp_url = agent_mcp_url(&base_url, "sales-session-rate"); + let initialized_session = initialize_session(&client, &mcp_url, &api_key).await; + + let start_response = post_jsonrpc( + &client, + &mcp_url, + &api_key, + Some(&initialized_session), + json!({ + "jsonrpc": "2.0", + "id": 1, + "method": "tools/call", + "params": { + "name": "echo_stream_session_rate_start", + "arguments": { "message": "hello" } + } + }), + ) + .await; + let session_id = start_response["result"]["structuredContent"]["session_id"] + .as_str() + .unwrap() + .to_owned(); + + let first_poll = post_jsonrpc( + &client, + &mcp_url, + &api_key, + Some(&initialized_session), + json!({ + "jsonrpc": "2.0", + "id": 2, + "method": "tools/call", + "params": { + "name": "echo_stream_session_rate_poll", + "arguments": { "session_id": session_id } + } + }), + ) + .await; + assert_eq!(first_poll["result"]["isError"], false); + + let second_poll = post_jsonrpc( + &client, + &mcp_url, + &api_key, + Some(&initialized_session), + json!({ + "jsonrpc": "2.0", + "id": 3, + "method": "tools/call", + "params": { + "name": "echo_stream_session_rate_poll", + "arguments": { "session_id": session_id } + } + }), + ) + .await; + + assert_eq!( + second_poll["result"]["structuredContent"]["error"]["code"], + json!("stream_session_poll_rate_limited") + ); + let poll_after_ms = + second_poll["result"]["structuredContent"]["error"]["context"]["poll_after_ms"] + .as_u64() + .unwrap(); + assert!((1..=250).contains(&poll_after_ms)); + } + #[tokio::test] async fn rejects_cross_agent_async_job_access() { let registry = test_registry().await; diff --git a/crates/crank-core/src/stream_session.rs b/crates/crank-core/src/stream_session.rs index 70c90d6..e849e94 100644 --- a/crates/crank-core/src/stream_session.rs +++ b/crates/crank-core/src/stream_session.rs @@ -58,6 +58,15 @@ impl StreamSession { !self.status.is_terminal() && !self.is_expired(now) } + pub fn remaining_poll_delay_ms(&self, now: OffsetDateTime, poll_interval_ms: u64) -> u64 { + let Some(last_poll_at) = self.last_poll_at else { + return 0; + }; + + let elapsed_ms = (now - last_poll_at).whole_milliseconds().max(0) as u64; + poll_interval_ms.saturating_sub(elapsed_ms) + } + pub fn mark_polled(&mut self, now: OffsetDateTime) { self.last_poll_at = Some(now); } @@ -183,6 +192,54 @@ mod tests { assert!(!session.can_poll(timestamp("2026-04-06T12:03:00Z"))); } + #[test] + fn first_poll_has_no_required_delay() { + let session = StreamSession { + id: StreamSessionId::new("stream_01"), + workspace_id: WorkspaceId::new("ws_01"), + agent_id: None, + operation_id: OperationId::new("op_01"), + protocol: Protocol::Rest, + mode: ExecutionMode::Session, + status: StreamStatus::Running, + cursor: None, + state: json!({"cursor":"abc"}), + expires_at: timestamp("2026-04-06T12:05:00Z"), + last_poll_at: None, + created_at: timestamp("2026-04-06T12:00:00Z"), + closed_at: None, + }; + + assert_eq!( + session.remaining_poll_delay_ms(timestamp("2026-04-06T12:00:00.100Z"), 250), + 0 + ); + } + + #[test] + fn rapid_repeat_poll_reports_remaining_delay() { + let session = StreamSession { + id: StreamSessionId::new("stream_01"), + workspace_id: WorkspaceId::new("ws_01"), + agent_id: None, + operation_id: OperationId::new("op_01"), + protocol: Protocol::Rest, + mode: ExecutionMode::Session, + status: StreamStatus::Running, + cursor: None, + state: json!({"cursor":"abc"}), + expires_at: timestamp("2026-04-06T12:05:00Z"), + last_poll_at: Some(timestamp("2026-04-06T12:01:00Z")), + created_at: timestamp("2026-04-06T12:00:00Z"), + closed_at: None, + }; + + assert_eq!( + session.remaining_poll_delay_ms(timestamp("2026-04-06T12:01:00.100Z"), 250), + 150 + ); + } + #[test] fn async_job_tracks_finish_and_failure() { let mut job = AsyncJobHandle {