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
.sqlx/query-30960f21f8c28a4190efacb6a68d5918da107b23a74b34d5dff0ee34e90a51b2.json → .sqlx/query-e86682005480df007b488b8543adc8a6d4068e33a5509f9e314dd0d97eac178b.json
Generated
+10 -4
View File
@@ -1,6 +1,6 @@
{
"db_name": "PostgreSQL",
"query": "select\n id,\n workspace_id,\n agent_id,\n operation_id,\n status,\n progress_json,\n result_json,\n error_json,\n expires_at as \"expires_at: OffsetDateTime\",\n created_at as \"created_at!: OffsetDateTime\",\n updated_at as \"updated_at!: OffsetDateTime\",\n finished_at as \"finished_at: OffsetDateTime\"\n from async_jobs\n where workspace_id = $1\n and ($2::text is null or agent_id = $2)\n and ($3::text is null or operation_id = $3)\n and ($4::text is null or status = $4)\n order by updated_at desc\n limit $5",
"query": "select\n id,\n workspace_id,\n agent_id,\n operation_id,\n status,\n progress_json,\n result_json,\n error_json,\n expires_at as \"expires_at: OffsetDateTime\",\n last_poll_at as \"last_poll_at: OffsetDateTime\",\n created_at as \"created_at!: OffsetDateTime\",\n updated_at as \"updated_at!: OffsetDateTime\",\n finished_at as \"finished_at: OffsetDateTime\"\n from async_jobs\n where workspace_id = $1\n and ($2::text is null or agent_id = $2)\n and ($3::text is null or operation_id = $3)\n and ($4::text is null or status = $4)\n order by updated_at desc\n limit $5",
"describe": {
"columns": [
{
@@ -50,16 +50,21 @@
},
{
"ordinal": 9,
"name": "created_at!: OffsetDateTime",
"name": "last_poll_at: OffsetDateTime",
"type_info": "Timestamptz"
},
{
"ordinal": 10,
"name": "updated_at!: OffsetDateTime",
"name": "created_at!: OffsetDateTime",
"type_info": "Timestamptz"
},
{
"ordinal": 11,
"name": "updated_at!: OffsetDateTime",
"type_info": "Timestamptz"
},
{
"ordinal": 12,
"name": "finished_at: OffsetDateTime",
"type_info": "Timestamptz"
}
@@ -83,10 +88,11 @@
true,
true,
true,
true,
false,
false,
true
]
},
"hash": "30960f21f8c28a4190efacb6a68d5918da107b23a74b34d5dff0ee34e90a51b2"
"hash": "e86682005480df007b488b8543adc8a6d4068e33a5509f9e314dd0d97eac178b"
}
.sqlx/query-6ea069ced9ab34bbd92d2068e457d17c0f98483284e55548c6f1f24ac284ae11.json → .sqlx/query-f551a9cd253b7fd0dc9a387a141a53d8a109b8aa766ab1c2b11f9b52e250adc8.json
Generated
+10 -4
View File
@@ -1,6 +1,6 @@
{
"db_name": "PostgreSQL",
"query": "select\n id,\n workspace_id,\n agent_id,\n operation_id,\n status,\n progress_json,\n result_json,\n error_json,\n expires_at as \"expires_at: OffsetDateTime\",\n created_at as \"created_at!: OffsetDateTime\",\n updated_at as \"updated_at!: OffsetDateTime\",\n finished_at as \"finished_at: OffsetDateTime\"\n from async_jobs\n where id = $1",
"query": "select\n id,\n workspace_id,\n agent_id,\n operation_id,\n status,\n progress_json,\n result_json,\n error_json,\n expires_at as \"expires_at: OffsetDateTime\",\n last_poll_at as \"last_poll_at: OffsetDateTime\",\n created_at as \"created_at!: OffsetDateTime\",\n updated_at as \"updated_at!: OffsetDateTime\",\n finished_at as \"finished_at: OffsetDateTime\"\n from async_jobs\n where id = $1",
"describe": {
"columns": [
{
@@ -50,16 +50,21 @@
},
{
"ordinal": 9,
"name": "created_at!: OffsetDateTime",
"name": "last_poll_at: OffsetDateTime",
"type_info": "Timestamptz"
},
{
"ordinal": 10,
"name": "updated_at!: OffsetDateTime",
"name": "created_at!: OffsetDateTime",
"type_info": "Timestamptz"
},
{
"ordinal": 11,
"name": "updated_at!: OffsetDateTime",
"type_info": "Timestamptz"
},
{
"ordinal": 12,
"name": "finished_at: OffsetDateTime",
"type_info": "Timestamptz"
}
@@ -79,10 +84,11 @@
true,
true,
true,
true,
false,
false,
true
]
},
"hash": "6ea069ced9ab34bbd92d2068e457d17c0f98483284e55548c6f1f24ac284ae11"
"hash": "f551a9cd253b7fd0dc9a387a141a53d8a109b8aa766ab1c2b11f9b52e250adc8"
}
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")
crates/crank-core/src/stream_session.rs
+37
View File
@@ -114,6 +114,9 @@ pub struct AsyncJobHandle {
#[serde(skip_serializing_if = "Option::is_none")]
#[serde(with = "time::serde::rfc3339::option")]
pub expires_at: Option<OffsetDateTime>,
#[serde(skip_serializing_if = "Option::is_none")]
#[serde(with = "time::serde::rfc3339::option")]
pub last_poll_at: Option<OffsetDateTime>,
#[serde(with = "time::serde::rfc3339")]
pub created_at: OffsetDateTime,
#[serde(with = "time::serde::rfc3339")]
@@ -132,6 +135,15 @@ impl AsyncJobHandle {
matches!(self.status, JobStatus::Created | JobStatus::Running)
}
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_finished(&mut self, now: OffsetDateTime, result: Value) {
self.status = JobStatus::Completed;
self.result = Some(result);
@@ -252,6 +264,7 @@ mod tests {
result: None,
error: None,
expires_at: Some(timestamp("2026-04-06T12:05:00Z")),
last_poll_at: None,
created_at: timestamp("2026-04-06T12:00:00Z"),
updated_at: timestamp("2026-04-06T12:00:00Z"),
finished_at: None,
@@ -278,4 +291,28 @@ mod tests {
Some(timestamp("2026-04-06T12:02:00Z"))
);
}
#[test]
fn async_job_reports_remaining_poll_delay() {
let job = AsyncJobHandle {
id: AsyncJobId::new("job_01"),
workspace_id: WorkspaceId::new("ws_01"),
agent_id: None,
operation_id: OperationId::new("op_01"),
status: JobStatus::Running,
progress: json!({"percent": 60}),
result: None,
error: None,
expires_at: Some(timestamp("2026-04-06T12:05:00Z")),
last_poll_at: Some(timestamp("2026-04-06T12:01:00Z")),
created_at: timestamp("2026-04-06T12:00:00Z"),
updated_at: timestamp("2026-04-06T12:00:00Z"),
finished_at: None,
};
assert_eq!(
job.remaining_poll_delay_ms(timestamp("2026-04-06T12:01:00.100Z"), 250),
150
);
}
}
crates/crank-registry/src/migrations.rs
+5
View File
@@ -558,6 +558,7 @@ pub async fn apply_postgres(pool: &PgPool) -> Result<(), sqlx::Error> {
result_json jsonb null,
error_json jsonb null,
expires_at timestamptz null,
last_poll_at timestamptz null,
created_at timestamptz not null,
updated_at timestamptz not null,
finished_at timestamptz null
@@ -566,6 +567,10 @@ pub async fn apply_postgres(pool: &PgPool) -> Result<(), sqlx::Error> {
.execute(pool)
.await?;
query("alter table async_jobs add column if not exists last_poll_at timestamptz null")
.execute(pool)
.await?;
query(
"create index if not exists async_jobs_workspace_status_idx
on async_jobs(workspace_id, status, updated_at desc)",
crates/crank-registry/src/postgres/mod.rs
+11
View File
@@ -556,6 +556,7 @@ fn map_async_job(row: &PgRow) -> Result<AsyncJobHandle, RegistryError> {
.try_get::<Option<Json<Value>>, _>("error_json")?
.map(|value| value.0),
expires_at: row.try_get("expires_at")?,
last_poll_at: row.try_get("last_poll_at")?,
created_at: row.try_get("created_at")?,
updated_at: row.try_get("updated_at")?,
finished_at: row.try_get("finished_at")?,
@@ -2171,6 +2172,15 @@ mod tests {
let loaded = registry.get_async_job(&job.id).await.unwrap().unwrap();
assert_eq!(loaded.status, JobStatus::Completed);
let touched = registry
.touch_async_job_poll(&job.id, &timestamp("2026-04-06T12:02:00Z"))
.await
.unwrap();
assert_eq!(
touched.last_poll_at,
Some(timestamp("2026-04-06T12:02:00Z"))
);
let page = registry
.list_async_jobs(AsyncJobFilter {
workspace_id: &test_workspace_id(),
@@ -2511,6 +2521,7 @@ mod tests {
result: None,
error: None,
expires_at: Some(timestamp("2026-04-06T12:05:00Z")),
last_poll_at: None,
created_at: timestamp("2026-04-06T12:00:00Z"),
updated_at: timestamp("2026-04-06T12:00:00Z"),
finished_at: None,
crates/crank-registry/src/postgres/stream.rs
+48 -4
View File
@@ -302,12 +302,13 @@ impl PostgresRegistry {
result_json,
error_json,
expires_at,
last_poll_at,
created_at,
updated_at,
finished_at
) values (
$1, $2, $3, $4, $5, $6::jsonb, $7::jsonb, $8::jsonb, $9::timestamptz,
$10::timestamptz, $11::timestamptz, $12::timestamptz
$10::timestamptz, $11::timestamptz, $12::timestamptz, $13::timestamptz
)",
)
.bind(request.job.id.as_str())
@@ -319,6 +320,7 @@ impl PostgresRegistry {
.bind(request.job.result.clone().map(Json))
.bind(request.job.error.clone().map(Json))
.bind(request.job.expires_at)
.bind(request.job.last_poll_at)
.bind(request.job.created_at)
.bind(request.job.updated_at)
.bind(request.job.finished_at)
@@ -343,6 +345,7 @@ impl PostgresRegistry {
result_json,
error_json,
expires_at as \"expires_at: OffsetDateTime\",
last_poll_at as \"last_poll_at: OffsetDateTime\",
created_at as \"created_at!: OffsetDateTime\",
updated_at as \"updated_at!: OffsetDateTime\",
finished_at as \"finished_at: OffsetDateTime\"
@@ -364,6 +367,7 @@ impl PostgresRegistry {
result: row.result_json,
error: row.error_json,
expires_at: row.expires_at,
last_poll_at: row.last_poll_at,
created_at: row.created_at,
updated_at: row.updated_at,
finished_at: row.finished_at,
@@ -383,9 +387,9 @@ impl PostgresRegistry {
set status = $3,
progress_json = $4::jsonb,
result_json = $5::jsonb,
error_json = $6::jsonb,
expires_at = coalesce($7::timestamptz, expires_at),
updated_at = $8::timestamptz,
error_json = $6::jsonb,
expires_at = coalesce($7::timestamptz, expires_at),
updated_at = $8::timestamptz,
finished_at = case
when $9::timestamptz is null then finished_at
else $9::timestamptz
@@ -402,6 +406,7 @@ impl PostgresRegistry {
result_json,
error_json,
expires_at,
last_poll_at,
created_at,
updated_at,
finished_at",
@@ -436,6 +441,43 @@ impl PostgresRegistry {
}
}
pub async fn touch_async_job_poll(
&self,
id: &AsyncJobId,
now: &OffsetDateTime,
) -> Result<AsyncJobHandle, RegistryError> {
let row = sqlx::query(
"update async_jobs
set last_poll_at = $2::timestamptz
where id = $1
returning
id,
workspace_id,
agent_id,
operation_id,
status,
progress_json,
result_json,
error_json,
expires_at,
last_poll_at,
created_at,
updated_at,
finished_at",
)
.bind(id.as_str())
.bind(now)
.fetch_optional(&self.pool)
.await?;
match row.as_ref() {
Some(row) => map_async_job(row),
None => Err(RegistryError::AsyncJobNotFound {
job_id: id.as_str().to_owned(),
}),
}
}
pub async fn cancel_async_job(
&self,
id: &AsyncJobId,
@@ -492,6 +534,7 @@ impl PostgresRegistry {
result_json,
error_json,
expires_at as \"expires_at: OffsetDateTime\",
last_poll_at as \"last_poll_at: OffsetDateTime\",
created_at as \"created_at!: OffsetDateTime\",
updated_at as \"updated_at!: OffsetDateTime\",
finished_at as \"finished_at: OffsetDateTime\"
@@ -524,6 +567,7 @@ impl PostgresRegistry {
result: row.result_json,
error: row.error_json,
expires_at: row.expires_at,
last_poll_at: row.last_poll_at,
created_at: row.created_at,
updated_at: row.updated_at,
finished_at: row.finished_at,