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;
@@ -19,7 +19,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,
@@ -35,39 +37,53 @@ mod route_tests;
 use crate::error::to_error_response;
 use crate::platform::{build_runtime_services, open_kv_store, UnavailableKvStore};
 
-#[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:?}");
 
     // Build the auction orchestrator once at startup
     let orchestrator = match build_orchestrator(&settings) {
-        Ok(orchestrator) => orchestrator,
+        Ok(o) => o,
         Err(e) => {
             log::error!("Failed to build auction orchestrator: {:?}", e);
-            return Ok(to_error_response(&e));
+            to_error_response(&e).send_to_client();
+            return;
         }
     };
 
     let integration_registry = match IntegrationRegistry::new(&settings) {
         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;
         }
     };
 
@@ -78,13 +94,17 @@ fn main(req: Request) -> Result<Response, Error> {
         as std::sync::Arc<dyn trusted_server_core::platform::PlatformKvStore>;
     let runtime_services = build_runtime_services(&req, kv_store);
 
-    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,
         &runtime_services,
         req,
-    ))
+    )) {
+        response.send_to_client();
+    }
 }
 
 async fn route_request(
@@ -93,7 +113,7 @@ async fn route_request(
     integration_registry: &IntegrationRegistry,
     runtime_services: &RuntimeServices,
     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).
@@ -115,14 +135,14 @@ async fn route_request(
     match enforce_basic_auth(settings, &req) {
         Ok(Some(mut response)) => {
             finalize_response(settings, geo_info.as_ref(), &mut response);
-            return Ok(response);
+            return Some(response);
         }
         Ok(None) => {}
         Err(e) => {
             log::error!("Failed to evaluate basic auth: {:?}", e);
             let mut response = to_error_response(&e);
             finalize_response(settings, geo_info.as_ref(), &mut response);
-            return Ok(response);
+            return Some(response);
         }
     }
 
@@ -193,7 +213,32 @@ async fn route_request(
                         &publisher_services,
                         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 committed. 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)
@@ -210,7 +255,7 @@ async fn route_request(
 
     finalize_response(settings, geo_info.as_ref(), &mut response);
 
-    Ok(response)
+    Some(response)
 }
 
 fn runtime_services_for_consent_route(

crates/trusted-server-core/src/integrations/registry.rs

@@ -730,6 +730,15 @@ impl IntegrationRegistry {
         self.inner.script_rewriters.clone()
     }
 
+    /// Check whether any HTML post-processors are registered.
+    ///
+    /// Cheaper than [`html_post_processors()`](Self::html_post_processors) when
+    /// only the presence check is needed — avoids cloning `Vec<Arc<…>>`.
+    #[must_use]
+    pub fn has_html_post_processors(&self) -> bool {
+        !self.inner.html_post_processors.is_empty()
+    }
+
     /// Expose registered HTML post-processors.
     #[must_use]
     pub fn html_post_processors(&self) -> Vec<Arc<dyn IntegrationHtmlPostProcessor>> {