mcp: throttle rapid stream session polls
This commit is contained in:
@@ -811,7 +811,7 @@ async fn handle_session_start_call(
|
|||||||
batch_size,
|
batch_size,
|
||||||
}),
|
}),
|
||||||
expires_at,
|
expires_at,
|
||||||
last_poll_at: Some(now),
|
last_poll_at: None,
|
||||||
created_at: now,
|
created_at: now,
|
||||||
closed_at: None,
|
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()) {
|
let mut state_payload: StoredSessionState = match serde_json::from_value(loaded.state.clone()) {
|
||||||
Ok(value) => value,
|
Ok(value) => value,
|
||||||
Err(error) => {
|
Err(error) => {
|
||||||
|
|||||||
@@ -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]
|
#[tokio::test]
|
||||||
async fn rejects_cross_agent_async_job_access() {
|
async fn rejects_cross_agent_async_job_access() {
|
||||||
let registry = test_registry().await;
|
let registry = test_registry().await;
|
||||||
|
|||||||
@@ -58,6 +58,15 @@ impl StreamSession {
|
|||||||
!self.status.is_terminal() && !self.is_expired(now)
|
!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) {
|
pub fn mark_polled(&mut self, now: OffsetDateTime) {
|
||||||
self.last_poll_at = Some(now);
|
self.last_poll_at = Some(now);
|
||||||
}
|
}
|
||||||
@@ -183,6 +192,54 @@ mod tests {
|
|||||||
assert!(!session.can_poll(timestamp("2026-04-06T12:03:00Z")));
|
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]
|
#[test]
|
||||||
fn async_job_tracks_finish_and_failure() {
|
fn async_job_tracks_finish_and_failure() {
|
||||||
let mut job = AsyncJobHandle {
|
let mut job = AsyncJobHandle {
|
||||||
|
|||||||
Reference in New Issue
Block a user