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/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,