Phase 2: Stream responses to client via StreamingBody

crates/trusted-server-adapter-fastly/src/main.rs

@@ -1,6 +1,6 @@
 use error_stack::Report;
 use fastly::http::Method;
-use fastly::{Error, Request, Response};
+use fastly::{Request, Response};
 use log_fastly::Logger;
 
 use trusted_server_core::auction::endpoints::handle_auction;
@@ -18,7 +18,9 @@ use trusted_server_core::proxy::{
     handle_first_party_click, handle_first_party_proxy, handle_first_party_proxy_rebuild,
     handle_first_party_proxy_sign,
 };
-use trusted_server_core::publisher::{handle_publisher_request, handle_tsjs_dynamic};
+use trusted_server_core::publisher::{
+    handle_publisher_request, handle_tsjs_dynamic, stream_publisher_body, PublisherResponse,
+};
 use trusted_server_core::request_signing::{
     handle_deactivate_key, handle_rotate_key, handle_trusted_server_discovery,
     handle_verify_signature,
@@ -29,21 +31,33 @@ use trusted_server_core::settings_data::get_settings;
 mod error;
 use crate::error::to_error_response;
 
-#[fastly::main]
-fn main(req: Request) -> Result<Response, Error> {
+/// Entry point for the Fastly Compute program.
+///
+/// Uses an undecorated `main()` with `Request::from_client()` instead of
+/// `#[fastly::main]` so we can call `stream_to_client()` or `send_to_client()`
+/// explicitly. `#[fastly::main]` is syntactic sugar that auto-calls
+/// `send_to_client()` on the returned `Response`, which is incompatible with
+/// streaming.
+fn main() {
     init_logger();
 
+    let req = Request::from_client();
+
     // Keep the health probe independent from settings loading and routing so
     // readiness checks still get a cheap liveness response during startup.
     if req.get_method() == Method::GET && req.get_path() == "/health" {
-        return Ok(Response::from_status(200).with_body_text_plain("ok"));
+        Response::from_status(200)
+            .with_body_text_plain("ok")
+            .send_to_client();
+        return;
     }
 
     let settings = match get_settings() {
         Ok(s) => s,
         Err(e) => {
             log::error!("Failed to load settings: {:?}", e);
-            return Ok(to_error_response(&e));
+            to_error_response(&e).send_to_client();
+            return;
         }
     };
     log::debug!("Settings {settings:?}");
@@ -55,24 +69,29 @@ fn main(req: Request) -> Result<Response, Error> {
         Ok(r) => r,
         Err(e) => {
             log::error!("Failed to create integration registry: {:?}", e);
-            return Ok(to_error_response(&e));
+            to_error_response(&e).send_to_client();
+            return;
         }
     };
 
-    futures::executor::block_on(route_request(
+    // route_request may send the response directly (streaming path) or
+    // return it for us to send (buffered path).
+    if let Some(response) = futures::executor::block_on(route_request(
         &settings,
         &orchestrator,
         &integration_registry,
         req,
-    ))
+    )) {
+        response.send_to_client();
+    }
 }
 
 async fn route_request(
     settings: &Settings,
     orchestrator: &AuctionOrchestrator,
     integration_registry: &IntegrationRegistry,
     mut req: Request,
-) -> Result<Response, Error> {
+) -> Option<Response> {
     // Strip client-spoofable forwarded headers at the edge.
     // On Fastly this service IS the first proxy — these headers from
     // clients are untrusted and can hijack URL rewriting (see #409).
@@ -83,7 +102,7 @@ async fn route_request(
 
     if let Some(mut response) = enforce_basic_auth(settings, &req) {
         finalize_response(settings, geo_info.as_ref(), &mut response);
-        return Ok(response);
+        return Some(response);
     }
 
     // Get path and method for routing
@@ -139,7 +158,32 @@ async fn route_request(
             );
 
             match handle_publisher_request(settings, integration_registry, req) {
-                Ok(response) => Ok(response),
+                Ok(PublisherResponse::Stream {
+                    mut response,
+                    body,
+                    params,
+                }) => {
+                    // Streaming path: finalize headers, then stream body to client.
+                    finalize_response(settings, geo_info.as_ref(), &mut response);
+                    let mut streaming_body = response.stream_to_client();
+                    if let Err(e) = stream_publisher_body(
+                        body,
+                        &mut streaming_body,
+                        &params,
+                        settings,
+                        integration_registry,
+                    ) {
+                        // Headers already sent (200 OK). Log and abort — client
+                        // sees a truncated response. Standard proxy behavior.
+                        log::error!("Streaming processing failed: {e:?}");
+                        drop(streaming_body);
+                    } else if let Err(e) = streaming_body.finish() {
+                        log::error!("Failed to finish streaming body: {e}");
+                    }
+                    // Response already sent via stream_to_client()
+                    return None;
+                }
+                Ok(PublisherResponse::Buffered(response)) => Ok(response),
                 Err(e) => {
                     log::error!("Failed to proxy to publisher origin: {:?}", e);
                     Err(e)
@@ -153,7 +197,7 @@ async fn route_request(
 
     finalize_response(settings, geo_info.as_ref(), &mut response);
 
-    Ok(response)
+    Some(response)
 }
 
 /// Applies all standard response headers: geo, version, staging, and configured headers.