From 66f28defe5b9c94a6102ee13b827cbda092d3b96 Mon Sep 17 00:00:00 2001 From: "a.tolmachev" Date: Fri, 1 May 2026 16:48:02 +0000 Subject: [PATCH] mcp: throttle async job poll requests --- ...c8a6d4068e33a5509f9e314dd0d97eac178b.json} | 14 ++- ...53d8a109b8aa766ab1c2b11f9b52e250adc8.json} | 14 ++- apps/admin-api/src/app.rs | 1 + apps/admin-api/src/service.rs | 1 + apps/mcp-server/src/app.rs | 79 +++++++++++- apps/mcp-server/src/main.rs | 115 +++++++++++++++++- crates/crank-core/src/stream_session.rs | 37 ++++++ crates/crank-registry/src/migrations.rs | 5 + crates/crank-registry/src/postgres/mod.rs | 11 ++ crates/crank-registry/src/postgres/stream.rs | 52 +++++++- 10 files changed, 312 insertions(+), 17 deletions(-) rename .sqlx/{query-30960f21f8c28a4190efacb6a68d5918da107b23a74b34d5dff0ee34e90a51b2.json => query-e86682005480df007b488b8543adc8a6d4068e33a5509f9e314dd0d97eac178b.json} (72%) rename .sqlx/{query-6ea069ced9ab34bbd92d2068e457d17c0f98483284e55548c6f1f24ac284ae11.json => query-f551a9cd253b7fd0dc9a387a141a53d8a109b8aa766ab1c2b11f9b52e250adc8.json} (78%) diff --git a/.sqlx/query-30960f21f8c28a4190efacb6a68d5918da107b23a74b34d5dff0ee34e90a51b2.json b/.sqlx/query-e86682005480df007b488b8543adc8a6d4068e33a5509f9e314dd0d97eac178b.json similarity index 72% rename from .sqlx/query-30960f21f8c28a4190efacb6a68d5918da107b23a74b34d5dff0ee34e90a51b2.json rename to .sqlx/query-e86682005480df007b488b8543adc8a6d4068e33a5509f9e314dd0d97eac178b.json index 4d4e703..19f81b5 100644 --- a/.sqlx/query-30960f21f8c28a4190efacb6a68d5918da107b23a74b34d5dff0ee34e90a51b2.json +++ b/.sqlx/query-e86682005480df007b488b8543adc8a6d4068e33a5509f9e314dd0d97eac178b.json @@ -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" } diff --git a/.sqlx/query-6ea069ced9ab34bbd92d2068e457d17c0f98483284e55548c6f1f24ac284ae11.json b/.sqlx/query-f551a9cd253b7fd0dc9a387a141a53d8a109b8aa766ab1c2b11f9b52e250adc8.json similarity index 78% rename from .sqlx/query-6ea069ced9ab34bbd92d2068e457d17c0f98483284e55548c6f1f24ac284ae11.json rename to .sqlx/query-f551a9cd253b7fd0dc9a387a141a53d8a109b8aa766ab1c2b11f9b52e250adc8.json index d5ce7bb..3ff8cb4 100644 --- a/.sqlx/query-6ea069ced9ab34bbd92d2068e457d17c0f98483284e55548c6f1f24ac284ae11.json +++ b/.sqlx/query-f551a9cd253b7fd0dc9a387a141a53d8a109b8aa766ab1c2b11f9b52e250adc8.json @@ -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" } diff --git a/apps/admin-api/src/app.rs b/apps/admin-api/src/app.rs index 443b945..fd82646 100644 --- a/apps/admin-api/src/app.rs +++ b/apps/admin-api/src/app.rs @@ -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, diff --git a/apps/admin-api/src/service.rs b/apps/admin-api/src/service.rs index 2534436..2ab1649 100644 --- a/apps/admin-api/src/service.rs +++ b/apps/admin-api/src/service.rs @@ -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, diff --git a/apps/mcp-server/src/app.rs b/apps/mcp-server/src/app.rs index 6dbecdc..09a25d5 100644 --- a/apps/mcp-server/src/app.rs +++ b/apps/mcp-server/src/app.rs @@ -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, diff --git a/apps/mcp-server/src/main.rs b/apps/mcp-server/src/main.rs index 698557f..f0e7851 100644 --- a/apps/mcp-server/src/main.rs +++ b/apps/mcp-server/src/main.rs @@ -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(®istry, &operation, "sales-async-rate").await; + let api_key = create_platform_api_key( + ®istry, + "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") diff --git a/crates/crank-core/src/stream_session.rs b/crates/crank-core/src/stream_session.rs index e849e94..96c1209 100644 --- a/crates/crank-core/src/stream_session.rs +++ b/crates/crank-core/src/stream_session.rs @@ -114,6 +114,9 @@ pub struct AsyncJobHandle { #[serde(skip_serializing_if = "Option::is_none")] #[serde(with = "time::serde::rfc3339::option")] pub expires_at: Option, + #[serde(skip_serializing_if = "Option::is_none")] + #[serde(with = "time::serde::rfc3339::option")] + pub last_poll_at: Option, #[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 + ); + } } diff --git a/crates/crank-registry/src/migrations.rs b/crates/crank-registry/src/migrations.rs index a590b93..bb83717 100644 --- a/crates/crank-registry/src/migrations.rs +++ b/crates/crank-registry/src/migrations.rs @@ -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)", diff --git a/crates/crank-registry/src/postgres/mod.rs b/crates/crank-registry/src/postgres/mod.rs index f1f3c83..f9151ea 100644 --- a/crates/crank-registry/src/postgres/mod.rs +++ b/crates/crank-registry/src/postgres/mod.rs @@ -556,6 +556,7 @@ fn map_async_job(row: &PgRow) -> Result { .try_get::>, _>("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, ×tamp("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, diff --git a/crates/crank-registry/src/postgres/stream.rs b/crates/crank-registry/src/postgres/stream.rs index a036af4..559553b 100644 --- a/crates/crank-registry/src/postgres/stream.rs +++ b/crates/crank-registry/src/postgres/stream.rs @@ -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 { + 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,