bsodfather/crank
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

mcp: throttle async job poll requests

Browse Source
This commit is contained in:
a.tolmachev
2026-05-01 16:48:02 +00:00
parent 4cd5d5b132
commit 66f28defe5
10 changed files with 312 additions and 17 deletions
Download Patch File Download Diff File Expand all files Collapse all files
apps/admin-api/src/app.rs
+1
View File
@@ -1776,6 +1776,7 @@ mod tests {
result: None,
error: None,
expires_at: Some(now + time::Duration::minutes(5)),
last_poll_at: None,
created_at: now,
updated_at: now,
finished_at: None,
apps/admin-api/src/service.rs
+1
View File
@@ -2410,6 +2410,7 @@ impl AdminService {
result: None,
error: None,
expires_at: Some(expires_at),
last_poll_at: None,
created_at,
updated_at: created_at,
finished_at: None,
apps/mcp-server/src/app.rs
+75 -4
View File
@@ -1205,6 +1205,7 @@ async fn handle_async_job_start_call(
now,
streaming.max_session_lifetime_ms.unwrap_or(300_000),
)),
last_poll_at: None,
created_at: now,
updated_at: now,
finished_at: None,
@@ -1325,7 +1326,7 @@ async fn handle_async_job_status_call(
session: &SessionState,
message: &Value,
response_mode: ResponseMode,
_tool: PublishedAgentTool,
tool: PublishedAgentTool,
arguments: Value,
) -> Response {
let control: AsyncJobControlArgs = match serde_json::from_value(arguments.clone()) {
@@ -1360,7 +1361,7 @@ async fn handle_async_job_status_call(
Err(error) => return internal_jsonrpc_error(message, error),
};
if !async_job_belongs_to_tool(&job, &_tool) {
if !async_job_belongs_to_tool(&job, &tool) {
return tool_error_response(
message,
response_mode,
@@ -1371,6 +1372,41 @@ async fn handle_async_job_status_call(
);
}
let Some(streaming) = tool.operation.execution_config.streaming.as_ref() else {
return tool_error_response(
message,
response_mode,
&session.protocol_version,
"streaming_config_error",
"streaming config is required for async job tools".to_owned(),
None,
);
};
let now = OffsetDateTime::now_utc();
let poll_after_ms = streaming.poll_interval_ms.unwrap_or(1_000);
let remaining_delay_ms = job.remaining_poll_delay_ms(now, poll_after_ms);
if remaining_delay_ms > 0 {
return tool_error_response(
message,
response_mode,
&session.protocol_version,
"async_job_poll_rate_limited",
format!(
"async job {} must wait before the next poll",
control.job_id
),
Some(json!({
"poll_after_ms": remaining_delay_ms,
})),
);
}
let job = match state.registry.touch_async_job_poll(&job.id, &now).await {
Ok(job) => job,
Err(error) => return internal_jsonrpc_error(message, error),
};
success_tool_response(
message,
response_mode,
@@ -1390,7 +1426,7 @@ async fn handle_async_job_result_call(
session: &SessionState,
message: &Value,
response_mode: ResponseMode,
_tool: PublishedAgentTool,
tool: PublishedAgentTool,
arguments: Value,
) -> Response {
let control: AsyncJobControlArgs = match serde_json::from_value(arguments.clone()) {
@@ -1425,7 +1461,7 @@ async fn handle_async_job_result_call(
Err(error) => return internal_jsonrpc_error(message, error),
};
if !async_job_belongs_to_tool(&job, &_tool) {
if !async_job_belongs_to_tool(&job, &tool) {
return tool_error_response(
message,
response_mode,
@@ -1436,6 +1472,41 @@ async fn handle_async_job_result_call(
);
}
let Some(streaming) = tool.operation.execution_config.streaming.as_ref() else {
return tool_error_response(
message,
response_mode,
&session.protocol_version,
"streaming_config_error",
"streaming config is required for async job tools".to_owned(),
None,
);
};
let now = OffsetDateTime::now_utc();
let poll_after_ms = streaming.poll_interval_ms.unwrap_or(1_000);
let remaining_delay_ms = job.remaining_poll_delay_ms(now, poll_after_ms);
if remaining_delay_ms > 0 {
return tool_error_response(
message,
response_mode,
&session.protocol_version,
"async_job_poll_rate_limited",
format!(
"async job {} must wait before the next poll",
control.job_id
),
Some(json!({
"poll_after_ms": remaining_delay_ms,
})),
);
}
let job = match state.registry.touch_async_job_poll(&job.id, &now).await {
Ok(job) => job,
Err(error) => return internal_jsonrpc_error(message, error),
};
match job.status {
JobStatus::Completed => success_tool_response(
message,
apps/mcp-server/src/main.rs
+114 -1
View File
@@ -1561,6 +1561,117 @@ mod tests {
assert!((1..=250).contains(&poll_after_ms));
}
#[tokio::test]
async fn rejects_rapid_repeat_async_job_status_poll() {
let registry = test_registry().await;
let upstream_base_url = spawn_upstream_server().await;
let operation = test_rest_async_job_operation(&upstream_base_url, "crm_async_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(&registry, &operation, "sales-async-rate").await;
let api_key = create_platform_api_key(
&registry,
"mcp-async-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-async-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": "crm_async_rate_start",
"arguments": {
"email": "user@example.com"
}
}
}),
)
.await;
let job_id = start_response["result"]["structuredContent"]["job_id"]
.as_str()
.unwrap()
.to_owned();
let first_status = post_jsonrpc(
&client,
&mcp_url,
&api_key,
Some(&initialized_session),
json!({
"jsonrpc": "2.0",
"id": 2,
"method": "tools/call",
"params": {
"name": "crm_async_rate_status",
"arguments": {
"job_id": job_id
}
}
}),
)
.await;
assert_eq!(first_status["result"]["isError"], false);
let second_status = post_jsonrpc(
&client,
&mcp_url,
&api_key,
Some(&initialized_session),
json!({
"jsonrpc": "2.0",
"id": 3,
"method": "tools/call",
"params": {
"name": "crm_async_rate_status",
"arguments": {
"job_id": job_id
}
}
}),
)
.await;
assert_eq!(
second_status["result"]["structuredContent"]["error"]["code"],
json!("async_job_poll_rate_limited")
);
let poll_after_ms =
second_status["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;
@@ -1770,9 +1881,11 @@ mod tests {
break;
}
sleep(Duration::from_millis(25)).await;
sleep(Duration::from_millis(250)).await;
}
sleep(Duration::from_millis(250)).await;
assert_eq!(
status_response["result"]["structuredContent"]["status"],
json!("completed")