api: throttle streaming admin read polls

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::<Value>()
+            .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::<Value>().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::<Value>()
+            .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::<Value>().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() {