mcp: throttle async job poll requests

2026-05-01 16:48:02 +00:00
parent 4cd5d5b132
commit 66f28defe5
10 changed files with 312 additions and 17 deletions
@@ -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")