feat: add rest sse streaming adapter

2026-04-06 10:54:01 +03:00
parent 8204a59dac
commit bf56494336
10 changed files with 560 additions and 25 deletions
@@ -2,8 +2,8 @@ use std::collections::BTreeMap;

 use crank_adapter_graphql::{GraphqlAdapter, GraphqlRequest};
 use crank_adapter_grpc::{GrpcAdapter, GrpcRequest};
-use crank_adapter_rest::{RestAdapter, RestRequest};
-use crank_core::{ExecutionMode, Target};
+use crank_adapter_rest::{RestAdapter, RestRequest, RestWindowRequest};
+use crank_core::{ExecutionMode, Target, TransportBehavior};
 use serde_json::{Map, Value, json};

 use crate::{
@@ -60,7 +60,16 @@ impl RuntimeExecutor {
        }

        let prepared_request = self.prepare_request(operation, input)?;
-        let adapter_response = self.execute_adapter(operation, prepared_request).await?;
+        let adapter_response = if matches!(
+            streaming.transport_behavior,
+            TransportBehavior::ServerStream
+        ) && matches!(operation.target, Target::Rest(_))
+        {
+            self.execute_window_adapter(operation, prepared_request)
+                .await?
+        } else {
+            self.execute_adapter(operation, prepared_request).await?
+        };

        crate::aggregation::collect_window_result(&adapter_response.body, streaming)
    }
@@ -159,6 +168,48 @@ impl RuntimeExecutor {
            }
        }
    }
+
+    async fn execute_window_adapter(
+        &self,
+        operation: &RuntimeOperation,
+        prepared_request: PreparedRequest,
+    ) -> Result<AdapterResponse, RuntimeError> {
+        match &operation.target {
+            Target::Rest(target) => {
+                let Some(streaming) = operation.execution_config.streaming.as_ref() else {
+                    return Err(RuntimeError::MissingStreamingConfig {
+                        operation_id: operation.operation_id.as_str().to_owned(),
+                    });
+                };
+                let request = RestWindowRequest {
+                    request: RestRequest {
+                        path_params: prepared_request.path_params.clone(),
+                        query_params: prepared_request.query_params.clone(),
+                        headers: merge_headers(
+                            &target.static_headers,
+                            &operation.execution_config.headers,
+                            &prepared_request.headers,
+                        ),
+                        body: prepared_request.body.clone(),
+                        timeout_ms: streaming
+                            .upstream_timeout_ms
+                            .unwrap_or(operation.execution_config.timeout_ms),
+                    },
+                    window_duration_ms: streaming.window_duration_ms.unwrap_or_default(),
+                    max_items: streaming.max_items.map(|value| value.saturating_add(1)),
+                };
+                let response = self.rest_adapter.execute_window(target, &request).await?;
+
+                Ok(AdapterResponse {
+                    status_code: response.status_code,
+                    headers: response.headers,
+                    body: response.body,
+                    data: Value::Null,
+                })
+            }
+            _ => self.execute_adapter(operation, prepared_request).await,
+        }
+    }
 }

 impl PreparedRequest {
@@ -264,7 +315,11 @@ fn non_empty_payload(value: Option<Value>) -> Option<Value> {
 mod tests {
    use std::collections::BTreeMap;

-    use axum::{Json, Router, routing::post};
+    use axum::{
+        Json, Router,
+        response::sse::{Event, KeepAlive, Sse},
+        routing::post,
+    };
    use crank_adapter_grpc::test_support as grpc_test_support;
    use crank_core::{
        AggregationMode, DescriptorId, ExecutionConfig, ExecutionMode, GeneratedDraft,
@@ -274,6 +329,7 @@ mod tests {
    };
    use crank_mapping::{MappingRule, MappingSet};
    use crank_schema::{Schema, SchemaKind};
+    use futures_util::stream;
    use serde_json::{Value, json};
    use tokio::net::TcpListener;

@@ -378,7 +434,8 @@ mod tests {
    async fn executes_window_mode_with_raw_items() {
        let base_url = spawn_runtime_server().await;
        let executor = RuntimeExecutor::new();
-        let operation = test_window_operation(&base_url, AggregationMode::RawItems, None, None);
+        let operation =
+            test_window_snapshot_operation(&base_url, AggregationMode::RawItems, None, None);

        let result = executor
            .execute_window(&operation, &json!({}))
@@ -396,7 +453,7 @@ mod tests {
        let base_url = spawn_runtime_server().await;
        let executor = RuntimeExecutor::new();
        let operation =
-            test_window_operation(&base_url, AggregationMode::SummaryOnly, Some(10), None);
+            test_window_snapshot_operation(&base_url, AggregationMode::SummaryOnly, Some(10), None);

        let result = executor
            .execute_window(&operation, &json!({}))
@@ -411,7 +468,7 @@ mod tests {
    async fn executes_window_mode_with_truncation_and_redaction() {
        let base_url = spawn_runtime_server().await;
        let executor = RuntimeExecutor::new();
-        let operation = test_window_operation(
+        let operation = test_window_snapshot_operation(
            &base_url,
            AggregationMode::SummaryPlusSamples,
            Some(2),
@@ -434,7 +491,7 @@ mod tests {
    async fn propagates_timeout_in_window_mode() {
        let base_url = spawn_runtime_server().await;
        let executor = RuntimeExecutor::new();
-        let operation = test_slow_window_operation(&base_url);
+        let operation = test_slow_window_snapshot_operation(&base_url);

        let error = executor
            .execute_window(&operation, &json!({}))
@@ -444,11 +501,48 @@ mod tests {
        assert!(matches!(error, RuntimeError::RestAdapter(_)));
    }

+    #[tokio::test]
+    async fn executes_window_mode_with_rest_sse_stream() {
+        let base_url = spawn_runtime_server().await;
+        let executor = RuntimeExecutor::new();
+        let operation = test_window_sse_operation(&base_url, AggregationMode::RawItems, Some(2));
+
+        let result = executor
+            .execute_window(&operation, &json!({}))
+            .await
+            .unwrap();
+
+        assert_eq!(result.items.len(), 2);
+        assert_eq!(result.summary, Value::Null);
+        assert!(result.truncated);
+        assert!(result.has_more);
+        assert!(!result.window_complete);
+    }
+
+    #[tokio::test]
+    async fn completes_empty_window_for_idle_rest_sse_stream() {
+        let base_url = spawn_runtime_server().await;
+        let executor = RuntimeExecutor::new();
+        let operation = test_slow_window_sse_operation(&base_url);
+
+        let result = executor
+            .execute_window(&operation, &json!({}))
+            .await
+            .unwrap();
+
+        assert!(result.items.is_empty());
+        assert!(result.window_complete);
+        assert!(!result.truncated);
+        assert!(!result.has_more);
+    }
+
    async fn spawn_runtime_server() -> String {
        let app = Router::new()
            .route("/leads", post(create_lead))
            .route("/events", post(events_window))
-            .route("/slow-events", post(slow_events_window));
+            .route("/slow-events", post(slow_events_window))
+            .route("/events-sse", post(events_stream))
+            .route("/slow-events-sse", post(slow_events_stream));
        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
        let address = listener.local_addr().unwrap();

@@ -511,6 +605,29 @@ mod tests {
        events_window().await
    }

+    async fn events_stream()
+    -> Sse<impl futures_util::Stream<Item = Result<Event, std::convert::Infallible>>> {
+        let events = vec![
+            Ok(Event::default()
+                .data("{\"message\":\"disk pressure detected\",\"secret\":\"token-a\"}")),
+            Ok(Event::default()
+                .data("{\"message\":\"error budget exhausted\",\"secret\":\"token-b\"}")),
+            Ok(Event::default().data("{\"message\":\"node restarted\",\"secret\":\"token-c\"}")),
+        ];
+
+        Sse::new(stream::iter(events)).keep_alive(KeepAlive::default())
+    }
+
+    async fn slow_events_stream()
+    -> Sse<impl futures_util::Stream<Item = Result<Event, std::convert::Infallible>>> {
+        let delayed = stream::once(async {
+            tokio::time::sleep(std::time::Duration::from_millis(100)).await;
+            Ok(Event::default().data("{\"message\":\"late event\"}"))
+        });
+
+        Sse::new(delayed).keep_alive(KeepAlive::default())
+    }
+
    async fn graphql_handler(Json(payload): Json<Value>) -> Json<Value> {
        let email = payload
            .get("variables")
@@ -781,7 +898,7 @@ mod tests {
        })
    }

-    fn test_window_operation(
+    fn test_window_snapshot_operation(
        base_url: &str,
        aggregation_mode: AggregationMode,
        max_items: Option<u32>,
@@ -833,7 +950,7 @@ mod tests {
                protocol_options: None,
                streaming: Some(StreamingConfig {
                    mode: ExecutionMode::Window,
-                    transport_behavior: TransportBehavior::ServerStream,
+                    transport_behavior: TransportBehavior::RequestResponse,
                    window_duration_ms: Some(3_000),
                    poll_interval_ms: None,
                    upstream_timeout_ms: Some(1_000),
@@ -870,9 +987,9 @@ mod tests {
        })
    }

-    fn test_slow_window_operation(base_url: &str) -> RuntimeOperation {
+    fn test_slow_window_snapshot_operation(base_url: &str) -> RuntimeOperation {
        let mut operation =
-            test_window_operation(base_url, AggregationMode::RawItems, Some(10), None);
+            test_window_snapshot_operation(base_url, AggregationMode::RawItems, Some(10), None);

        if let Target::Rest(target) = &mut operation.target {
            target.path_template = "/slow-events".to_owned();
@@ -882,6 +999,47 @@ mod tests {
        operation
    }

+    fn test_window_sse_operation(
+        base_url: &str,
+        aggregation_mode: AggregationMode,
+        max_items: Option<u32>,
+    ) -> RuntimeOperation {
+        let mut operation =
+            test_window_snapshot_operation(base_url, aggregation_mode, max_items, None);
+
+        if let Target::Rest(target) = &mut operation.target {
+            target.path_template = "/events-sse".to_owned();
+        }
+
+        operation.execution_config.streaming = Some(StreamingConfig {
+            transport_behavior: TransportBehavior::ServerStream,
+            summary_path: None,
+            cursor_path: None,
+            done_path: Some("$.done".to_owned()),
+            ..operation.execution_config.streaming.clone().unwrap()
+        });
+
+        operation
+    }
+
+    fn test_slow_window_sse_operation(base_url: &str) -> RuntimeOperation {
+        let mut operation =
+            test_window_sse_operation(base_url, AggregationMode::RawItems, Some(10));
+
+        if let Target::Rest(target) = &mut operation.target {
+            target.path_template = "/slow-events-sse".to_owned();
+        }
+
+        operation
+            .execution_config
+            .streaming
+            .as_mut()
+            .expect("streaming config")
+            .window_duration_ms = Some(20);
+        operation.execution_config.timeout_ms = 1_000;
+        operation
+    }
+
    fn object_schema(field_name: &str, kind: SchemaKind) -> Schema {
        Schema {
            kind: SchemaKind::Object,