From 4839cae3197430bc199260c7eaf5bb36406dd72b Mon Sep 17 00:00:00 2001 From: "a.tolmachev" Date: Sat, 2 May 2026 09:44:28 +0000 Subject: [PATCH] api: throttle streaming admin read polls --- apps/admin-api/src/app.rs | 131 ++++++++++++++++++- apps/admin-api/src/service.rs | 107 ++++++++++++++- apps/mcp-server/src/main.rs | 5 +- crates/crank-registry/src/postgres/mod.rs | 9 ++ crates/crank-registry/src/postgres/stream.rs | 37 ++++++ 5 files changed, 285 insertions(+), 4 deletions(-) diff --git a/apps/admin-api/src/app.rs b/apps/admin-api/src/app.rs index c668b63..2d9b159 100644 --- a/apps/admin-api/src/app.rs +++ b/apps/admin-api/src/app.rs @@ -231,10 +231,11 @@ mod tests { AsyncJobHandle, AsyncJobId, DescriptorId, ExecutionConfig, ExecutionMode, GraphqlOperationType, GraphqlTarget, GrpcTarget, HttpMethod, JobStatus, MembershipRole, OperationId, Protocol, RestTarget, SecretKind, SoapBindingStyle, SoapOperationMetadata, - SoapTarget, SoapVersion, Target, ToolDescription, TransportBehavior, WorkspaceId, + SoapTarget, SoapVersion, StreamSession, StreamStatus, Target, ToolDescription, + TransportBehavior, WorkspaceId, }; use crank_mapping::{MappingRule, MappingSet}; - use crank_registry::{CreateAsyncJobRequest, PostgresRegistry}; + use crank_registry::{CreateAsyncJobRequest, CreateStreamSessionRequest, PostgresRegistry}; use crank_runtime::SecretCrypto; use crank_schema::{Schema, SchemaKind}; use serde_json::{Value, json}; @@ -1810,6 +1811,132 @@ mod tests { ); } + #[tokio::test(flavor = "multi_thread")] + #[serial] + async fn rate_limits_rapid_stream_session_reads() { + let registry = test_registry().await; + let storage_root = test_storage_root("stream_session_rate_limit"); + let server_addr = grpc_test_support::spawn_unary_echo_server().await; + let base_url = spawn_admin_api(build_test_app(registry.clone(), storage_root)).await; + let client = authorized_client(&base_url).await; + + let created = client + .post(format!("{base_url}/operations")) + .json(&test_grpc_session_operation_payload( + &server_addr, + "echo_session_rate_limit", + )) + .send() + .await + .unwrap() + .json::() + .await + .unwrap(); + let operation_id = OperationId::new(created["operation_id"].as_str().unwrap().to_owned()); + let now = OffsetDateTime::now_utc(); + let session_id = crank_core::StreamSessionId::new("sess_rate_limited".to_owned()); + registry + .create_stream_session(CreateStreamSessionRequest { + session: &StreamSession { + id: session_id.clone(), + workspace_id: WorkspaceId::new(DEFAULT_WORKSPACE_ID), + agent_id: None, + operation_id, + protocol: Protocol::Grpc, + mode: ExecutionMode::Session, + status: StreamStatus::Running, + cursor: None, + state: json!({ "items": [] }), + expires_at: now + time::Duration::minutes(5), + last_poll_at: Some(now), + created_at: now, + closed_at: None, + }, + }) + .await + .unwrap(); + + let response = client + .get(format!("{base_url}/stream-sessions/{}", session_id.as_str())) + .send() + .await + .unwrap(); + let status = response.status(); + let body = response.json::().await.unwrap(); + + assert_eq!(status, reqwest::StatusCode::TOO_MANY_REQUESTS); + assert_eq!(body["error"]["code"], "rate_limited"); + assert_eq!( + body["error"]["message"], + "stream session poll rate limit exceeded" + ); + assert_eq!(body["error"]["context"]["session_id"], session_id.as_str()); + assert!(body["error"]["context"]["poll_after_ms"].as_u64().unwrap() > 0); + } + + #[tokio::test(flavor = "multi_thread")] + #[serial] + async fn rate_limits_rapid_async_job_result_polls() { + let registry = test_registry().await; + let storage_root = test_storage_root("async_job_rate_limit"); + let upstream_base_url = spawn_upstream_server().await; + let base_url = spawn_admin_api(build_test_app(registry.clone(), storage_root)).await; + let client = authorized_client(&base_url).await; + + let created = client + .post(format!("{base_url}/operations")) + .json(&test_rest_async_job_operation_payload( + &upstream_base_url, + "crm_async_job_rate_limit", + )) + .send() + .await + .unwrap() + .json::() + .await + .unwrap(); + let operation_id = OperationId::new(created["operation_id"].as_str().unwrap().to_owned()); + let now = OffsetDateTime::now_utc(); + let job_id = AsyncJobId::new("job_rate_limited".to_owned()); + registry + .create_async_job(CreateAsyncJobRequest { + job: &AsyncJobHandle { + id: job_id.clone(), + workspace_id: WorkspaceId::new(DEFAULT_WORKSPACE_ID), + agent_id: None, + operation_id, + status: JobStatus::Running, + progress: json!({ "pct": 25 }), + result: None, + error: None, + expires_at: Some(now + time::Duration::minutes(5)), + last_poll_at: Some(now), + created_at: now, + updated_at: now, + finished_at: None, + }, + }) + .await + .unwrap(); + + let response = client + .get(format!("{base_url}/async-jobs/{}/result", job_id.as_str())) + .send() + .await + .unwrap(); + let status = response.status(); + let body = response.json::().await.unwrap(); + + assert_eq!(status, reqwest::StatusCode::TOO_MANY_REQUESTS); + assert_eq!(body["error"]["code"], "rate_limited"); + assert_eq!( + body["error"]["message"], + "async job poll rate limit exceeded" + ); + assert_eq!(body["error"]["context"]["job_id"], job_id.as_str()); + assert!(body["error"]["context"]["poll_after_ms"].as_u64().unwrap() > 0); + } + #[tokio::test(flavor = "multi_thread")] #[serial] async fn returns_structured_context_for_missing_stream_session() { diff --git a/apps/admin-api/src/service.rs b/apps/admin-api/src/service.rs index 2ab1649..09df05f 100644 --- a/apps/admin-api/src/service.rs +++ b/apps/admin-api/src/service.rs @@ -1585,6 +1585,9 @@ impl AdminService { ) })?; ensure_stream_session_workspace(&session, workspace_id)?; + let session = self + .touch_stream_session_poll_if_ready(workspace_id, session) + .await?; Ok(stream_session_detail_view(session)) } @@ -1665,6 +1668,11 @@ impl AdminService { ) })?; ensure_async_job_workspace(&job, workspace_id)?; + let job = if matches!(job.status, JobStatus::Completed) { + job + } else { + self.touch_async_job_poll_if_ready(workspace_id, job).await? + }; Ok(async_job_detail_view(job)) } @@ -1710,6 +1718,7 @@ impl AdminService { ) })?; ensure_async_job_workspace(&job, workspace_id)?; + let job = self.touch_async_job_poll_if_ready(workspace_id, job).await?; match job.status { JobStatus::Completed => Ok(job.result.unwrap_or(Value::Null)), @@ -1745,6 +1754,102 @@ impl AdminService { } } + async fn touch_stream_session_poll_if_ready( + &self, + workspace_id: &WorkspaceId, + session: StreamSession, + ) -> Result { + let poll_interval_ms = self + .streaming_poll_interval_ms(workspace_id, &session.operation_id) + .await?; + let now = OffsetDateTime::now_utc(); + let remaining_delay_ms = session.remaining_poll_delay_ms(now, poll_interval_ms); + + if remaining_delay_ms > 0 { + return Err(ApiError::rate_limited_with_context( + "stream session poll rate limit exceeded", + json!({ + "session_id": session.id.as_str(), + "poll_after_ms": remaining_delay_ms, + }), + )); + } + + self.registry + .touch_stream_session_poll(&session.id, &now) + .await + .map_err(ApiError::from) + } + + async fn touch_async_job_poll_if_ready( + &self, + workspace_id: &WorkspaceId, + job: AsyncJobHandle, + ) -> Result { + let poll_interval_ms = self + .streaming_poll_interval_ms(workspace_id, &job.operation_id) + .await?; + let now = OffsetDateTime::now_utc(); + let remaining_delay_ms = job.remaining_poll_delay_ms(now, poll_interval_ms); + + if remaining_delay_ms > 0 { + return Err(ApiError::rate_limited_with_context( + "async job poll rate limit exceeded", + json!({ + "job_id": job.id.as_str(), + "poll_after_ms": remaining_delay_ms, + }), + )); + } + + self.registry + .touch_async_job_poll(&job.id, &now) + .await + .map_err(ApiError::from) + } + + async fn streaming_poll_interval_ms( + &self, + workspace_id: &WorkspaceId, + operation_id: &OperationId, + ) -> Result { + let summary = self + .registry + .get_operation_summary(workspace_id, operation_id) + .await? + .ok_or_else(|| { + ApiError::not_found_with_context( + format!("operation {} was not found", operation_id.as_str()), + json!({ "operation_id": operation_id.as_str() }), + ) + })?; + let version = self + .registry + .get_operation_version(workspace_id, operation_id, summary.current_draft_version) + .await? + .ok_or_else(|| { + ApiError::not_found_with_context( + format!( + "operation version {} for {} was not found", + summary.current_draft_version, + operation_id.as_str() + ), + json!({ + "operation_id": operation_id.as_str(), + "version": summary.current_draft_version, + }), + ) + })?; + + Ok(version + .snapshot + .execution_config + .streaming + .as_ref() + .and_then(|streaming| streaming.poll_interval_ms) + .unwrap_or(1_000)) + } + pub async fn list_operations( &self, workspace_id: &WorkspaceId, @@ -2354,7 +2459,7 @@ impl AdminService { batch_size, }), expires_at, - last_poll_at: Some(created_at), + last_poll_at: None, created_at, closed_at: None, }; diff --git a/apps/mcp-server/src/main.rs b/apps/mcp-server/src/main.rs index 8c12a23..641e924 100644 --- a/apps/mcp-server/src/main.rs +++ b/apps/mcp-server/src/main.rs @@ -989,7 +989,10 @@ mod tests { .await .unwrap(); - assert_eq!(second_response.status(), reqwest::StatusCode::TOO_MANY_REQUESTS); + assert_eq!( + second_response.status(), + reqwest::StatusCode::TOO_MANY_REQUESTS + ); assert!(second_response.headers().get(header::RETRY_AFTER).is_some()); } diff --git a/crates/crank-registry/src/postgres/mod.rs b/crates/crank-registry/src/postgres/mod.rs index f9151ea..2e64454 100644 --- a/crates/crank-registry/src/postgres/mod.rs +++ b/crates/crank-registry/src/postgres/mod.rs @@ -2027,6 +2027,15 @@ mod tests { .unwrap(); assert_eq!(closed.status, StreamStatus::Stopped); + let touched = registry + .touch_stream_session_poll(&session.id, ×tamp("2026-04-06T12:00:20Z")) + .await + .unwrap(); + assert_eq!( + touched.last_poll_at, + Some(timestamp("2026-04-06T12:00:20Z")) + ); + let updated = registry .update_stream_session_state(UpdateStreamSessionStateRequest { session_id: &session.id, diff --git a/crates/crank-registry/src/postgres/stream.rs b/crates/crank-registry/src/postgres/stream.rs index 559553b..74b204b 100644 --- a/crates/crank-registry/src/postgres/stream.rs +++ b/crates/crank-registry/src/postgres/stream.rs @@ -198,6 +198,43 @@ impl PostgresRegistry { } } + pub async fn touch_stream_session_poll( + &self, + id: &StreamSessionId, + now: &OffsetDateTime, + ) -> Result { + let row = sqlx::query( + "update stream_sessions + set last_poll_at = $2::timestamptz + where id = $1 + returning + id, + workspace_id, + agent_id, + operation_id, + protocol, + mode, + status, + cursor_json, + state_json, + expires_at, + last_poll_at, + created_at, + closed_at", + ) + .bind(id.as_str()) + .bind(now) + .fetch_optional(&self.pool) + .await?; + + match row.as_ref() { + Some(row) => map_stream_session(row), + None => Err(RegistryError::StreamSessionNotFound { + session_id: id.as_str().to_owned(), + }), + } + } + pub async fn list_stream_sessions( &self, filter: StreamSessionFilter<'_>,