feat(http): stream HTTP bodies end-to-end

.github/workflows/publish.yaml

@@ -561,6 +561,7 @@ jobs:
           pnpm build -F rivetkit
           pnpm build -F rivetkit -F '@rivetkit/*' -F '!@rivetkit/shared-data' -F '!@rivetkit/engine-frontend' -F '!@rivetkit/mcp-hub' -F '!@rivetkit/rivetkit-napi' -F '!@rivetkit/rivetkit-wasm'
           pnpm build -F '@rivet-dev/vercel-world'
+          pnpm build -F '@rivet-dev/flue'
 
       # ---- shared publish (runs for all triggers) ----
       - name: Finalize package versions for publish

engine/packages/config/src/config/pegboard.rs

@@ -105,6 +105,10 @@ pub struct Pegboard {
 	pub gateway_websocket_open_timeout_ms: Option<u64>,
 	/// Timeout for response to start in milliseconds.
 	pub gateway_response_start_timeout_ms: Option<u64>,
+	/// Timeout between streaming HTTP response chunks in milliseconds.
+	///
+	/// Disabled when unset so long-lived streams such as SSE may remain idle.
+	pub gateway_response_chunk_idle_timeout_ms: Option<u64>,
 	/// Ping interval for gateway updates in milliseconds.
 	pub gateway_update_ping_interval_ms: Option<u64>,
 	/// GC interval for in-flight requests in milliseconds.
@@ -280,6 +284,10 @@ impl Pegboard {
 			.unwrap_or(5 * 60 * 1000)
 	}
 
+	pub fn gateway_response_chunk_idle_timeout_ms(&self) -> Option<u64> {
+		self.gateway_response_chunk_idle_timeout_ms
+	}
+
 	pub fn gateway_update_ping_interval_ms(&self) -> u64 {
 		self.gateway_update_ping_interval_ms.unwrap_or(3_000)
 	}

engine/packages/guard-core/src/custom_serve.rs

@@ -2,7 +2,7 @@ use anyhow::{Result, bail};
 use async_trait::async_trait;
 use bytes::Bytes;
 use http_body_util::Full;
-use hyper::{Request, Response};
+use hyper::{Request, Response, body::Incoming as BodyIncoming};
 use tokio_tungstenite::tungstenite::protocol::frame::CloseFrame;
 
 use crate::WebSocketHandle;
@@ -17,13 +17,28 @@ pub enum HibernationResult {
 /// Trait for custom request serving logic that can handle both HTTP and WebSocket requests
 #[async_trait]
 pub trait CustomServeTrait: Send + Sync {
+	/// Returns true when this service wants the original request body stream.
+	/// The default buffered path keeps retry semantics for existing custom routes.
+	fn streams_request_body(&self) -> bool {
+		false
+	}
+
 	/// Handle a regular HTTP request
 	async fn handle_request(
 		&self,
 		req: Request<Full<Bytes>>,
 		req_ctx: &mut RequestContext,
 	) -> Result<Response<ResponseBody>>;
 
+	/// Handle a regular HTTP request with the original inbound body stream.
+	async fn handle_streaming_request(
+		&self,
+		_req: Request<BodyIncoming>,
+		_req_ctx: &mut RequestContext,
+	) -> Result<Response<ResponseBody>> {
+		bail!("service does not support streaming request bodies");
+	}
+
 	/// Handle a WebSocket connection after upgrade. Supports connection retries.
 	async fn handle_websocket(
 		&self,

engine/packages/guard-core/src/proxy_service.rs

@@ -727,7 +727,8 @@ impl ProxyService {
 		metrics::PROXY_REQUEST_PENDING.inc();
 		metrics::PROXY_REQUEST_TOTAL.inc();
 
-		let res = if hyper_tungstenite::is_upgrade_request(&req) {
+		let is_websocket = hyper_tungstenite::is_upgrade_request(&req);
+		let res = if is_websocket {
 			self.handle_websocket_upgrade(req, req_ctx, target).await
 		} else {
 			self.handle_http_request(req, req_ctx, target).await
@@ -746,20 +747,42 @@ impl ProxyService {
 
 		metrics::PROXY_REQUEST_PENDING.dec();
 
-		// Release in-flight counter and request ID when done
-		let state_clone = self.state.clone();
 		let client_ip = req_ctx.client_ip;
 		let in_flight_request_id = req_ctx.in_flight_request_id;
-		tokio::spawn(
-			async move {
-				state_clone
-					.release_in_flight(client_ip, in_flight_request_id)
-					.await;
-			}
-			.instrument(tracing::info_span!("release_in_flight_task")),
-		);
+		let state = self.state.clone();
+		let runtime = tokio::runtime::Handle::current();
+		// HTTP capacity remains held until the response reaches EOF, errors, or is dropped.
+		// WebSocket upgrades retain the existing immediate release behavior. Capturing the
+		// runtime handle lets the response body's synchronous Drop path schedule async cleanup.
+		let release_in_flight = move || {
+			runtime.spawn(
+				async move {
+					state
+						.release_in_flight(client_ip, in_flight_request_id)
+						.await;
+				}
+				.instrument(tracing::info_span!("release_in_flight_task")),
+			);
+		};
 
-		res
+		if is_websocket {
+			release_in_flight();
+			res
+		} else {
+			match res {
+				Ok(response) => {
+					let (parts, body) = response.into_parts();
+					Ok(Response::from_parts(
+						parts,
+						body.with_completion(release_in_flight),
+					))
+				}
+				Err(err) => {
+					release_in_flight();
+					Err(err)
+				}
+			}
+		}
 	}
 
 	#[tracing::instrument(skip_all)]
@@ -936,6 +959,10 @@ impl ProxyService {
 				.build());
 			}
 			ResolveRouteOutput::CustomServe(mut handler) => {
+				if handler.streams_request_body() {
+					return handler.handle_streaming_request(req, req_ctx).await;
+				}
+
 				// Collect request body
 				let (req_parts, body) = req.into_parts();
 				let req_body =
@@ -1002,18 +1029,10 @@ impl ProxyService {
 						continue;
 					}
 
-					// Release in-flight counter and request ID before returning
-					self.state
-						.release_in_flight(req_ctx.client_ip, req_ctx.in_flight_request_id)
-						.await;
 					return res;
 				}
 
 				// If we get here, all attempts failed
-				// Release in-flight counter and request ID before returning error
-				self.state
-					.release_in_flight(req_ctx.client_ip, req_ctx.in_flight_request_id)
-					.await;
 				return Err(errors::RetryAttemptsExceeded {
 					attempts: req_ctx.retry.max_attempts,
 					last_error_code,

engine/packages/guard-core/src/response_body.rs

@@ -1,6 +1,32 @@
 use bytes::Bytes;
 use http_body_util::Full;
 use hyper::body::Incoming as BodyIncoming;
+use tokio::sync::mpsc;
+
+pub type ResponseBodyError = Box<dyn std::error::Error + Send + Sync>;
+
+#[doc(hidden)]
+pub struct CompletionGuard(Option<Box<dyn FnOnce() + Send + 'static>>);
+
+impl CompletionGuard {
+	fn complete(&mut self) {
+		if let Some(callback) = self.0.take() {
+			callback();
+		}
+	}
+}
+
+impl std::fmt::Debug for CompletionGuard {
+	fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+		f.debug_struct("CompletionGuard").finish_non_exhaustive()
+	}
+}
+
+impl Drop for CompletionGuard {
+	fn drop(&mut self) {
+		self.complete();
+	}
+}
 
 /// Response body type that can handle both streaming and buffered responses
 #[derive(Debug)]
@@ -9,11 +35,27 @@ pub enum ResponseBody {
 	Full(Full<Bytes>),
 	/// Streaming response body
 	Incoming(BodyIncoming),
+	/// Channel-backed streaming response body
+	Channel(mpsc::Receiver<Result<Bytes, ResponseBodyError>>),
+	/// Body carrying a callback that runs at EOF, error, or drop.
+	WithCompletion {
+		body: Box<ResponseBody>,
+		completion: CompletionGuard,
+	},
+}
+
+impl ResponseBody {
+	pub(crate) fn with_completion(self, callback: impl FnOnce() + Send + 'static) -> Self {
+		Self::WithCompletion {
+			body: Box::new(self),
+			completion: CompletionGuard(Some(Box::new(callback))),
+		}
+	}
 }
 
 impl http_body::Body for ResponseBody {
 	type Data = Bytes;
-	type Error = Box<dyn std::error::Error + Send + Sync>;
+	type Error = ResponseBodyError;
 
 	fn poll_frame(
 		self: std::pin::Pin<&mut Self>,
@@ -46,20 +88,46 @@ impl http_body::Body for ResponseBody {
 					std::task::Poll::Pending => std::task::Poll::Pending,
 				}
 			}
+			ResponseBody::Channel(rx) => match rx.poll_recv(cx) {
+				std::task::Poll::Ready(Some(Ok(bytes))) => {
+					std::task::Poll::Ready(Some(Ok(http_body::Frame::data(bytes))))
+				}
+				std::task::Poll::Ready(Some(Err(err))) => std::task::Poll::Ready(Some(Err(err))),
+				std::task::Poll::Ready(None) => std::task::Poll::Ready(None),
+				std::task::Poll::Pending => std::task::Poll::Pending,
+			},
+			ResponseBody::WithCompletion { body, completion } => {
+				let result = std::pin::Pin::new(body.as_mut()).poll_frame(cx);
+				if matches!(
+					&result,
+					std::task::Poll::Ready(None) | std::task::Poll::Ready(Some(Err(_)))
+				) {
+					completion.complete();
+				}
+				result
+			}
 		}
 	}
 
 	fn is_end_stream(&self) -> bool {
 		match self {
 			ResponseBody::Full(body) => body.is_end_stream(),
 			ResponseBody::Incoming(body) => body.is_end_stream(),
+			ResponseBody::Channel(rx) => rx.is_closed() && rx.is_empty(),
+			ResponseBody::WithCompletion { body, .. } => body.is_end_stream(),
 		}
 	}
 
 	fn size_hint(&self) -> http_body::SizeHint {
 		match self {
 			ResponseBody::Full(body) => body.size_hint(),
 			ResponseBody::Incoming(body) => body.size_hint(),
+			ResponseBody::Channel(_) => http_body::SizeHint::default(),
+			ResponseBody::WithCompletion { body, .. } => body.size_hint(),
 		}
 	}
 }
+
+#[cfg(test)]
+#[path = "../tests/support/response_body.rs"]
+mod tests;

engine/packages/guard-core/tests/response_body.rs

@@ -0,0 +1,30 @@
+use bytes::Bytes;
+use http_body_util::BodyExt;
+use rivet_guard_core::ResponseBody;
+use tokio::sync::mpsc;
+
+#[tokio::test]
+async fn channel_body_yields_sent_chunks() {
+	let (tx, rx) = mpsc::channel(2);
+	tx.send(Ok(Bytes::from_static(b"hello "))).await.unwrap();
+	tx.send(Ok(Bytes::from_static(b"world"))).await.unwrap();
+	drop(tx);
+
+	let collected = ResponseBody::Channel(rx).collect().await.unwrap();
+
+	assert_eq!(collected.to_bytes(), Bytes::from_static(b"hello world"));
+}
+
+#[tokio::test]
+async fn channel_body_surfaces_errors() {
+	let (tx, rx) = mpsc::channel(1);
+	tx.send(Err(std::io::Error::other("stream failed").into()))
+		.await
+		.unwrap();
+	drop(tx);
+
+	let mut body = ResponseBody::Channel(rx);
+	let frame = body.frame().await.expect("expected frame");
+
+	assert!(frame.is_err());
+}

engine/packages/guard-core/tests/support/response_body.rs

@@ -0,0 +1,57 @@
+use std::sync::{
+	Arc,
+	atomic::{AtomicUsize, Ordering},
+};
+
+use bytes::Bytes;
+use http_body_util::{BodyExt, Full};
+use tokio::sync::mpsc;
+
+use super::{ResponseBody, ResponseBodyError};
+
+fn completion_counter() -> (Arc<AtomicUsize>, impl FnOnce() + Send + 'static) {
+	let count = Arc::new(AtomicUsize::new(0));
+	let callback_count = count.clone();
+	(count, move || {
+		callback_count.fetch_add(1, Ordering::SeqCst);
+	})
+}
+
+#[tokio::test]
+async fn completion_runs_once_at_eof() {
+	let (count, callback) = completion_counter();
+	let body = ResponseBody::Full(Full::new(Bytes::from_static(b"done"))).with_completion(callback);
+
+	assert_eq!(
+		body.collect().await.unwrap().to_bytes(),
+		Bytes::from_static(b"done")
+	);
+	assert_eq!(count.load(Ordering::SeqCst), 1);
+}
+
+#[tokio::test]
+async fn completion_runs_once_on_error() {
+	let (tx, rx) = mpsc::channel(1);
+	tx.send(Err::<Bytes, ResponseBodyError>("stream failed".into()))
+		.await
+		.unwrap();
+	drop(tx);
+
+	let (count, callback) = completion_counter();
+	let mut body = ResponseBody::Channel(rx).with_completion(callback);
+
+	assert!(body.frame().await.unwrap().is_err());
+	assert_eq!(count.load(Ordering::SeqCst), 1);
+	drop(body);
+	assert_eq!(count.load(Ordering::SeqCst), 1);
+}
+
+#[test]
+fn completion_runs_once_on_drop() {
+	let (_tx, rx) = mpsc::channel(1);
+	let (count, callback) = completion_counter();
+	let body = ResponseBody::Channel(rx).with_completion(callback);
+
+	drop(body);
+	assert_eq!(count.load(Ordering::SeqCst), 1);
+}