@rivetkit/engine 1.0.0 → 2.2.1-pr.4600.b74ff3b

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

@@ -1,358 +1,66 @@
 use anyhow::{Context, Result, bail, ensure};
 use bytes::Bytes;
 use futures_util::{SinkExt, StreamExt};
-use http_body_util::{BodyExt, Full};
-use hyper::{Request, Response, StatusCode, body::Incoming as BodyIncoming, header::HeaderName};
+use http_body_util::{BodyExt, Full, Limited};
+use hyper::{
+	Request, Response, StatusCode,
+	body::Incoming as BodyIncoming,
+	header::{HeaderName, HeaderValue},
+};
 use hyper_tungstenite;
 use hyper_util::{client::legacy::Client, rt::TokioExecutor};
 use moka::future::Cache;
-use rand;
-use rivet_api_builder::{ErrorResponse, RawErrorResponse};
+use rand::seq::SliceRandom;
 use rivet_api_builder::{RequestIds, X_RIVET_RAY_ID};
-use rivet_error::{INTERNAL_ERROR, RivetError};
-use rivet_metrics::KeyValue;
 use rivet_util::Id;
-use serde_json;
+use tracing_opentelemetry::OpenTelemetrySpanExt;
 
 use rivet_runner_protocol as protocol;
 use std::{
-	borrow::Cow,
-	collections::{HashMap as StdHashMap, HashSet},
-	net::SocketAddr,
+	net::{IpAddr, SocketAddr},
 	sync::Arc,
 	time::{Duration, Instant},
 };
 use tokio::sync::Mutex;
 use tokio::time::timeout;
-use tokio_tungstenite::tungstenite::{
-	client::IntoClientRequest,
-	protocol::{CloseFrame, frame::coding::CloseCode},
-};
+use tokio_tungstenite::tungstenite::client::IntoClientRequest;
 use tracing::Instrument;
 use url::Url;
 
+use crate::RouteTarget;
+use crate::request_context::RequestContext;
+use crate::response_body::ResponseBody;
+use crate::route::{
+	CacheKeyFn, DEFAULT_ROUTE_TIMEOUT, ResolveRouteOutput, RouteCache, RoutingFn, RoutingOutput,
+};
+use crate::utils::{InFlightCounter, RateLimiter};
 use crate::{
-	WebSocketHandle,
-	custom_serve::{CustomServeTrait, HibernationResult},
-	errors, metrics,
-	request_context::RequestContext,
-	task_group::TaskGroup,
+	WebSocketHandle, custom_serve::HibernationResult, errors, metrics, task_group::TaskGroup, utils,
 };
 
-const X_RIVET_TARGET: HeaderName = HeaderName::from_static("x-rivet-target");
-const X_RIVET_ACTOR: HeaderName = HeaderName::from_static("x-rivet-actor");
-const X_RIVET_TOKEN: HeaderName = HeaderName::from_static("x-rivet-token");
 pub const X_FORWARDED_FOR: HeaderName = HeaderName::from_static("x-forwarded-for");
 pub const X_RIVET_ERROR: HeaderName = HeaderName::from_static("x-rivet-error");
 
-const ROUTE_CACHE_TTL: Duration = Duration::from_secs(60 * 10); // 10 minutes
 const PROXY_STATE_CACHE_TTL: Duration = Duration::from_secs(60 * 60); // 1 hour
 const WEBSOCKET_CLOSE_LINGER: Duration = Duration::from_millis(100); // Keep TCP connection open briefly after WebSocket close
 
-/// Response body type that can handle both streaming and buffered responses
-#[derive(Debug)]
-pub enum ResponseBody {
-	/// Buffered response body
-	Full(Full<Bytes>),
-	/// Streaming response body
-	Incoming(BodyIncoming),
-}
-
-impl http_body::Body for ResponseBody {
-	type Data = Bytes;
-	type Error = Box<dyn std::error::Error + Send + Sync>;
-
-	fn poll_frame(
-		self: std::pin::Pin<&mut Self>,
-		cx: &mut std::task::Context<'_>,
-	) -> std::task::Poll<Option<Result<http_body::Frame<Self::Data>, Self::Error>>> {
-		match self.get_mut() {
-			ResponseBody::Full(body) => {
-				let pin = std::pin::Pin::new(body);
-				match pin.poll_frame(cx) {
-					std::task::Poll::Ready(Some(Ok(frame))) => {
-						std::task::Poll::Ready(Some(Ok(frame)))
-					}
-					std::task::Poll::Ready(Some(Err(e))) => {
-						std::task::Poll::Ready(Some(Err(Box::new(e))))
-					}
-					std::task::Poll::Ready(None) => std::task::Poll::Ready(None),
-					std::task::Poll::Pending => std::task::Poll::Pending,
-				}
-			}
-			ResponseBody::Incoming(body) => {
-				let pin = std::pin::Pin::new(body);
-				match pin.poll_frame(cx) {
-					std::task::Poll::Ready(Some(Ok(frame))) => {
-						std::task::Poll::Ready(Some(Ok(frame)))
-					}
-					std::task::Poll::Ready(Some(Err(e))) => {
-						std::task::Poll::Ready(Some(Err(Box::new(e))))
-					}
-					std::task::Poll::Ready(None) => std::task::Poll::Ready(None),
-					std::task::Poll::Pending => std::task::Poll::Pending,
-				}
-			}
-		}
-	}
-
-	fn is_end_stream(&self) -> bool {
-		match self {
-			ResponseBody::Full(body) => body.is_end_stream(),
-			ResponseBody::Incoming(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(),
-		}
-	}
-}
-
-// Routing types
-#[derive(Clone, Debug)]
-pub struct RouteTarget {
-	pub actor_id: Option<Id>,
-	pub host: String,
-	pub port: u16,
-	pub path: String,
-}
-
-#[derive(Clone, Debug)]
-pub struct RoutingTimeout {
-	pub routing_timeout: u64, // in seconds
-}
-
-#[derive(Clone, Debug)]
-pub struct RouteConfig {
-	pub targets: Vec<RouteTarget>,
-	pub timeout: RoutingTimeout,
-}
-
-#[derive(Clone)]
-pub enum RoutingOutput {
-	/// Return the data to route to.
-	Route(RouteConfig),
-	/// Return a custom response.
-	Response(StructuredResponse),
-	/// Return a custom serve handler.
-	CustomServe(Arc<dyn CustomServeTrait>),
-}
-
-#[derive(Clone, Debug)]
-pub struct StructuredResponse {
-	pub status: StatusCode,
-	pub message: Cow<'static, str>,
-	pub docs: Option<Cow<'static, str>>,
-}
-
-impl StructuredResponse {
-	pub fn build_response(&self) -> Result<Response<ResponseBody>> {
-		let mut body = StdHashMap::new();
-		body.insert("message", self.message.clone().into_owned());
-
-		if let Some(docs) = &self.docs {
-			body.insert("docs", docs.clone().into_owned());
-		}
-
-		let body_json = serde_json::to_string(&body)?;
-		let bytes = Bytes::from(body_json);
-
-		let response = Response::builder()
-			.status(self.status)
-			.header(hyper::header::CONTENT_TYPE, "application/json")
-			.body(ResponseBody::Full(Full::new(bytes)))?;
-
-		Ok(response)
-	}
-}
-
-#[derive(Clone)]
-enum ResolveRouteOutput {
-	Target(RouteTarget),
-	Response(StructuredResponse),
-	CustomServe(Arc<dyn CustomServeTrait>),
-}
-
-/// Enum defining the type of port the request came in on
-#[derive(Clone, Debug, PartialEq)]
-pub enum PortType {
-	Http,
-	Https,
-}
-
-pub type RoutingFn = Arc<
-	dyn for<'a> Fn(
-			&'a str,
-			&'a str,
-			PortType,
-			&'a hyper::HeaderMap,
-		) -> futures::future::BoxFuture<'a, Result<RoutingOutput>>
-		+ Send
-		+ Sync,
->;
-
-pub type CacheKeyFn = Arc<
-	dyn for<'a> Fn(
-			&'a str,
-			&'a str,
-			&'a hyper::Method,
-			PortType,
-			&'a hyper::HeaderMap,
-		) -> Result<u64>
-		+ Send
-		+ Sync,
->;
-
-#[derive(Clone, Debug)]
-pub struct MiddlewareConfig {
-	pub rate_limit: RateLimitConfig,
-	pub max_in_flight: MaxInFlightConfig,
-	pub retry: RetryConfig,
-	pub timeout: TimeoutConfig,
-}
-
-#[derive(Clone, Debug)]
-pub struct RateLimitConfig {
-	pub requests: u64,
-	pub period: u64, // in seconds
-}
-
-#[derive(Clone, Debug)]
-pub struct MaxInFlightConfig {
-	pub amount: usize,
-}
-
-#[derive(Clone, Debug)]
-pub struct RetryConfig {
-	pub max_attempts: u32,
-	pub initial_interval: u64, // in milliseconds
-}
-
-#[derive(Clone, Debug)]
-pub struct TimeoutConfig {
-	pub request_timeout: u64, // in seconds
-}
-
-#[derive(Clone, Debug)]
-pub enum MiddlewareResponse {
-	Ok(MiddlewareConfig),
-	NotFound,
-}
-
-pub type MiddlewareFn = Arc<
-	dyn for<'a> Fn(
-			&'a Id,
-			&'a hyper::HeaderMap,
-		) -> futures::future::BoxFuture<'a, Result<MiddlewareResponse>>
-		+ Send
-		+ Sync,
->;
-
-// Cache for routing results
-struct RouteCache {
-	cache: Cache<u64, RouteConfig>,
-}
-
-impl RouteCache {
-	fn new() -> Self {
-		Self {
-			cache: Cache::builder()
-				.max_capacity(10_000)
-				.time_to_live(ROUTE_CACHE_TTL)
-				.build(),
-		}
-	}
-
-	#[tracing::instrument(skip_all)]
-	async fn get(&self, key: &u64) -> Option<RouteConfig> {
-		self.cache.get(key).await
-	}
-
-	#[tracing::instrument(skip_all)]
-	async fn insert(&self, key: u64, result: RouteConfig) {
-		self.cache.insert(key, result).await;
-
-		metrics::ROUTE_CACHE_COUNT.record(self.cache.entry_count(), &[]);
-	}
-}
-
-// Rate limiter
-struct RateLimiter {
-	requests_remaining: u64,
-	reset_time: Instant,
-	requests_limit: u64,
-	period: Duration,
-}
-
-impl RateLimiter {
-	fn new(requests: u64, period_seconds: u64) -> Self {
-		Self {
-			requests_remaining: requests,
-			reset_time: Instant::now() + Duration::from_secs(period_seconds),
-			requests_limit: requests,
-			period: Duration::from_secs(period_seconds),
-		}
-	}
-
-	fn try_acquire(&mut self) -> bool {
-		let now = Instant::now();
-
-		// Check if we need to reset the counter
-		if now >= self.reset_time {
-			self.requests_remaining = self.requests_limit;
-			self.reset_time = now + self.period;
-		}
-
-		// Try to consume a request
-		if self.requests_remaining > 0 {
-			self.requests_remaining -= 1;
-			true
-		} else {
-			false
-		}
-	}
-}
-
-// In-flight requests counter
-struct InFlightCounter {
-	count: usize,
-	max: usize,
-}
-
-impl InFlightCounter {
-	fn new(max: usize) -> Self {
-		Self { count: 0, max }
-	}
-
-	fn try_acquire(&mut self) -> bool {
-		if self.count < self.max {
-			self.count += 1;
-			true
-		} else {
-			false
-		}
-	}
-
-	fn release(&mut self) {
-		self.count = self.count.saturating_sub(1);
-	}
-}
-
 // State shared across all request handlers
 pub struct ProxyState {
 	config: rivet_config::Config,
 	routing_fn: RoutingFn,
 	cache_key_fn: CacheKeyFn,
-	middleware_fn: MiddlewareFn,
+	// NOTE: Using the hyper legacy client is the only option currently.
+	// This is what reqwest uses under the hood. Eventually we'll migrate to h3 once it's ready.
+	client: Client<
+		hyper_tls::HttpsConnector<hyper_util::client::legacy::connect::HttpConnector>,
+		Full<Bytes>,
+	>,
 	route_cache: RouteCache,
-	rate_limiters: Cache<(Id, std::net::IpAddr), Arc<Mutex<RateLimiter>>>,
-	in_flight_counters: Cache<(Id, std::net::IpAddr), Arc<Mutex<InFlightCounter>>>,
-	in_flight_requests: Arc<Mutex<HashSet<protocol::RequestId>>>,
-	port_type: PortType,
-	clickhouse_inserter: Option<clickhouse_inserter::ClickHouseInserterHandle>,
+	// We use moka::Cache instead of scc::HashMap because it automatically handles TTL and capacity
+	rate_limiters: Cache<std::net::IpAddr, Arc<Mutex<RateLimiter>>>,
+	in_flight_counters: Cache<std::net::IpAddr, Arc<Mutex<InFlightCounter>>>,
+	in_flight_requests: Cache<protocol::RequestId, ()>,
+
 	tasks: Arc<TaskGroup>,
 }
 
@@ -361,15 +69,17 @@ impl ProxyState {
 		config: rivet_config::Config,
 		routing_fn: RoutingFn,
 		cache_key_fn: CacheKeyFn,
-		middleware_fn: MiddlewareFn,
-		port_type: PortType,
-		clickhouse_inserter: Option<clickhouse_inserter::ClickHouseInserterHandle>,
 	) -> Self {
+		let https_connector = hyper_tls::HttpsConnector::new();
+		let client = Client::builder(TokioExecutor::new())
+			.pool_idle_timeout(Duration::from_secs(30))
+			.build(https_connector);
+
 		Self {
 			config,
 			routing_fn,
 			cache_key_fn,
-			middleware_fn,
+			client,
 			route_cache: RouteCache::new(),
 			rate_limiters: Cache::builder()
 				.max_capacity(10_000)
@@ -379,119 +89,98 @@ impl ProxyState {
 				.max_capacity(10_000)
 				.time_to_live(PROXY_STATE_CACHE_TTL)
 				.build(),
-			in_flight_requests: Arc::new(Mutex::new(HashSet::new())),
-			port_type,
-			clickhouse_inserter,
+			in_flight_requests: Cache::builder().max_capacity(10_000_000).build(),
 			tasks: TaskGroup::new(),
 		}
 	}
 
-	#[tracing::instrument(skip(self))]
+	#[tracing::instrument(skip_all)]
 	async fn resolve_route(
 		&self,
-		hostname: &str,
-		path: &str,
-		method: &hyper::Method,
-		port_type: PortType,
-		headers: &hyper::HeaderMap,
+		req_ctx: &mut RequestContext,
 		ignore_cache: bool,
 	) -> Result<ResolveRouteOutput> {
-		// Extract just the hostname, stripping the port if present
-		let hostname_only = hostname.split(':').next().unwrap_or(hostname);
-
 		tracing::debug!(
-			hostname = %hostname_only,
-			path = %path,
-			method = %method,
-			port_type = ?port_type,
+			hostname = %req_ctx.hostname,
+			path = %req_ctx.path,
+			method = %req_ctx.method,
 			"Resolving route for request"
 		);
 
-		let cache_key =
-			(self.cache_key_fn)(hostname_only, &path, method, port_type.clone(), headers)?;
+		let cache_key = (self.cache_key_fn)(req_ctx)?;
 
 		// Check cache first
-		if !ignore_cache {
-			if let Some(result) = self.route_cache.get(&cache_key).await {
-				// Choose a random target from the cached targets
-				if let Some(target) = choose_random_target(&result.targets) {
-					return Ok(ResolveRouteOutput::Target(target.clone()));
-				}
-			}
-		}
+		let cache_res = if !ignore_cache {
+			self.route_cache.get(&cache_key).await
+		} else {
+			None
+		};
 
-		// Not in cache, call routing function with a default timeout
-		// Default 15 seconds, routing functions should have their own internal timeouts that are shorter
-		let default_timeout = Duration::from_secs(15);
+		let res = if let Some(res) = cache_res {
+			res
+		} else {
+			// Not in cache, call routing function with a default timeout
+			// Default 15 seconds, routing functions should have their own internal timeouts that are shorter
+			tracing::debug!(
+				hostname = %req_ctx.hostname,
+				path = %req_ctx.path,
+				cache_hit = false,
+				timeout_seconds = DEFAULT_ROUTE_TIMEOUT.as_secs(),
+				"Cache miss, calling routing function"
+			);
 
-		tracing::debug!(
-			hostname = %hostname_only,
-			path = %path,
-			cache_hit = false,
-			timeout_seconds = default_timeout.as_secs(),
-			"Cache miss, calling routing function"
-		);
+			let routing_res = timeout(DEFAULT_ROUTE_TIMEOUT, (self.routing_fn)(req_ctx))
+				.await
+				.map_err(|_| {
+					errors::RequestTimeout {
+						timeout_seconds: DEFAULT_ROUTE_TIMEOUT.as_secs(),
+					}
+					.build()
+				})??;
 
-		let res = timeout(
-			default_timeout,
-			(self.routing_fn)(hostname_only, path, port_type, headers),
-		)
-		.await
-		.map_err(|_| {
-			errors::RequestTimeout {
-				timeout_seconds: default_timeout.as_secs(),
-			}
-			.build()
-		})?;
+			// TODO: Disable route caching for now, determine edge cases with gateway
+			// // Cache the result
+			// self.route_cache
+			// 	.insert(cache_key, routing_res.clone())
+			// 	.await;
+			// tracing::debug!("Added route to cache");
 
-		match res? {
+			routing_res
+		};
+
+		match res {
 			RoutingOutput::Route(result) => {
 				tracing::debug!(
-					hostname = %hostname_only,
-					path = %path,
+					hostname = %req_ctx.hostname,
+					path = %req_ctx.path,
 					targets_count = result.targets.len(),
-					timeout_secs = result.timeout.routing_timeout,
 					"Received routing result"
 				);
 
-				// Cache the result
-				self.route_cache.insert(cache_key, result.clone()).await;
-				tracing::debug!("Added route to cache");
-
 				// Choose a random target
 				if let Some(target) = choose_random_target(&result.targets) {
 					tracing::debug!(
-						hostname = %hostname_only,
-						path = %path,
+						hostname = %req_ctx.hostname,
+						path = %req_ctx.path,
 						target_host = %target.host,
 						target_port = target.port,
 						target_path = %target.path,
-						actor_id = ?target.actor_id,
 						"Selected target for request"
 					);
 					Ok(ResolveRouteOutput::Target(target.clone()))
 				} else {
 					tracing::warn!(
-						hostname = %hostname_only,
-						path = %path,
+						hostname = %req_ctx.hostname,
+						path = %req_ctx.path,
 						"No route targets available from result"
 					);
 					Err(errors::NoRouteTargets.build())
 				}
 			}
-			RoutingOutput::Response(response) => {
-				tracing::debug!(
-					hostname = %hostname_only,
-					path = %path,
-					status = ?response.status,
-					"Routing returned custom response"
-				);
-				Ok(ResolveRouteOutput::Response(response))
-			}
 			RoutingOutput::CustomServe(handler) => {
 				tracing::debug!(
-					hostname = %hostname_only,
-					path = %path,
+					hostname = %req_ctx.hostname,
+					path = %req_ctx.path,
 					"Routing returned custom serve handler"
 				);
 				Ok(ResolveRouteOutput::CustomServe(handler))
@@ -499,183 +188,109 @@ impl ProxyState {
 		}
 	}
 
+	/// Returns true if the rate limit was hit.
 	#[tracing::instrument(skip_all)]
-	async fn get_middleware_config(
-		&self,
-		actor_id: &Id,
-		headers: &hyper::HeaderMap,
-	) -> Result<MiddlewareConfig> {
-		// Call the middleware function with a timeout
-		let default_timeout = Duration::from_secs(5); // Default 5 seconds
-
-		let middleware_result =
-			timeout(default_timeout, (self.middleware_fn)(actor_id, headers)).await;
-
-		match middleware_result {
-			Ok(result) => match result? {
-				MiddlewareResponse::Ok(config) => Ok(config),
-				MiddlewareResponse::NotFound => {
-					// Default values if middleware not found for this actor
-					Ok(MiddlewareConfig {
-						rate_limit: RateLimitConfig {
-							requests: 100, // 100 requests
-							period: 60,    // per 60 seconds
-						},
-						max_in_flight: MaxInFlightConfig {
-							amount: 20, // 20 concurrent requests
-						},
-						retry: RetryConfig {
-							max_attempts: 3,       // 3 retry attempts
-							initial_interval: 100, // 100ms initial interval
-						},
-						timeout: TimeoutConfig {
-							request_timeout: 30, // 30 seconds for requests
-						},
-					})
-				}
-			},
-			Err(_) => {
-				// Default values if middleware times out
-				Ok(MiddlewareConfig {
-					rate_limit: RateLimitConfig {
-						requests: 100, // 100 requests
-						period: 60,    // per 60 seconds
-					},
-					max_in_flight: MaxInFlightConfig {
-						amount: 20, // 20 concurrent requests
-					},
-					retry: RetryConfig {
-						max_attempts: 3,       // 3 retry attempts
-						initial_interval: 100, // 100ms initial interval
-					},
-					timeout: TimeoutConfig {
-						request_timeout: 30, // 30 seconds for requests
-					},
-				})
-			}
-		}
-	}
-
-	#[tracing::instrument(skip_all)]
-	async fn check_rate_limit(
-		&self,
-		ip_addr: std::net::IpAddr,
-		actor_id: &Option<Id>,
-		headers: &hyper::HeaderMap,
-	) -> Result<bool> {
-		let Some(actor_id) = *actor_id else {
-			// No rate limiting when actor_id is None
-			return Ok(true);
-		};
-
-		// Get actor-specific middleware config
-		let middleware_config = self.get_middleware_config(&actor_id, headers).await?;
-
-		let cache_key = (actor_id, ip_addr);
-
+	async fn check_rate_limit(&self, req_ctx: &RequestContext) -> Result<bool> {
 		// Get existing limiter or create a new one
-		let limiter_arc = if let Some(existing_limiter) = self.rate_limiters.get(&cache_key).await {
-			existing_limiter
-		} else {
-			let new_limiter = Arc::new(Mutex::new(RateLimiter::new(
-				middleware_config.rate_limit.requests,
-				middleware_config.rate_limit.period,
-			)));
-			self.rate_limiters
-				.insert(cache_key, new_limiter.clone())
-				.await;
-			metrics::RATE_LIMITER_COUNT.record(self.rate_limiters.entry_count(), &[]);
-			new_limiter
-		};
+		let limiter_arc =
+			if let Some(existing_limiter) = self.rate_limiters.get(&req_ctx.client_ip).await {
+				existing_limiter
+			} else {
+				let new_limiter = Arc::new(Mutex::new(RateLimiter::new(
+					req_ctx.rate_limit.requests,
+					req_ctx.rate_limit.period,
+				)));
+				self.rate_limiters
+					.insert(req_ctx.client_ip, new_limiter.clone())
+					.await;
+				metrics::RATE_LIMITER_COUNT.set(self.rate_limiters.entry_count() as i64);
+				new_limiter
+			};
 
 		// Try to acquire from the limiter
-		let result = {
+		let acquired = {
 			let mut limiter = limiter_arc.lock().await;
 			limiter.try_acquire()
 		};
 
-		Ok(result)
+		Ok(!acquired)
 	}
 
+	/// Returns true if the counter could not be acquired.
 	#[tracing::instrument(skip_all)]
-	async fn acquire_in_flight(
-		&self,
-		ip_addr: std::net::IpAddr,
-		actor_id: &Option<Id>,
-		headers: &hyper::HeaderMap,
-	) -> Result<Option<protocol::RequestId>> {
-		// Check in-flight limit if actor_id is present
-		if let Some(actor_id) = *actor_id {
-			// Get actor-specific middleware config
-			let middleware_config = self.get_middleware_config(&actor_id, headers).await?;
-
-			let cache_key = (actor_id, ip_addr);
-
-			// Get existing counter or create a new one
-			let counter_arc = if let Some(existing_counter) =
-				self.in_flight_counters.get(&cache_key).await
-			{
+	async fn acquire_in_flight(&self, req_ctx: &mut RequestContext) -> Result<bool> {
+		let cache_key = req_ctx.client_ip;
+
+		// Get existing counter or create a new one
+		let counter_arc =
+			if let Some(existing_counter) = self.in_flight_counters.get(&cache_key).await {
 				existing_counter
 			} else {
 				let new_counter = Arc::new(Mutex::new(InFlightCounter::new(
-					middleware_config.max_in_flight.amount,
+					req_ctx.max_in_flight.amount,
 				)));
 				self.in_flight_counters
 					.insert(cache_key, new_counter.clone())
 					.await;
-				metrics::IN_FLIGHT_COUNTER_COUNT.record(self.in_flight_counters.entry_count(), &[]);
+				metrics::IN_FLIGHT_COUNTER_COUNT.set(self.in_flight_counters.entry_count() as i64);
 				new_counter
 			};
 
-			// Try to acquire from the counter
-			let acquired = {
-				let mut counter = counter_arc.lock().await;
-				counter.try_acquire()
-			};
+		// Try to acquire from the counter
+		let acquired = {
+			let mut counter = counter_arc.lock().await;
+			counter.try_acquire()
+		};
 
-			if !acquired {
-				return Ok(None); // Rate limited
-			}
+		if !acquired {
+			return Ok(true); // Rate limited
 		}
 
 		// Generate unique request ID
-		let request_id = Some(self.generate_unique_request_id().await?);
-		Ok(request_id)
+		req_ctx.in_flight_request_id = Some(self.generate_unique_in_flight_request_id().await?);
+
+		Ok(false)
 	}
 
 	#[tracing::instrument(skip_all)]
 	async fn release_in_flight(
 		&self,
-		ip_addr: std::net::IpAddr,
-		actor_id: &Option<Id>,
-		request_id: protocol::RequestId,
+		client_ip: IpAddr,
+		in_flight_request_id: Option<protocol::RequestId>,
 	) {
-		// Release in-flight counter if actor_id is present
-		if let Some(actor_id) = *actor_id {
-			let cache_key = (actor_id, ip_addr);
-			if let Some(counter_arc) = self.in_flight_counters.get(&cache_key).await {
-				let mut counter = counter_arc.lock().await;
-				counter.release();
-			}
+		if let Some(counter_arc) = self.in_flight_counters.get(&client_ip).await {
+			let mut counter = counter_arc.lock().await;
+			counter.release();
 		}
 
-		// Release request ID
-		let mut requests = self.in_flight_requests.lock().await;
-		requests.remove(&request_id);
+		if let Some(in_flight_request_id) = in_flight_request_id {
+			// Release request ID
+			self.in_flight_requests
+				.invalidate(&in_flight_request_id)
+				.await;
+			metrics::IN_FLIGHT_REQUEST_COUNT.set(self.in_flight_requests.entry_count() as i64);
+		}
 	}
 
 	/// Generate a unique request ID that is not currently in flight
-	async fn generate_unique_request_id(&self) -> Result<protocol::RequestId> {
+	async fn generate_unique_in_flight_request_id(&self) -> Result<protocol::RequestId> {
 		const MAX_TRIES: u32 = 100;
-		let mut requests = self.in_flight_requests.lock().await;
 
 		for attempt in 0..MAX_TRIES {
 			let request_id = protocol::util::generate_request_id();
+			let mut inserted = false;
 
 			// Check if this ID is already in use
-			if !requests.contains(&request_id) {
-				// Insert the ID and return it
-				requests.insert(request_id);
+			self.in_flight_requests
+				.entry(request_id)
+				.or_insert_with(async {
+					inserted = true;
+				})
+				.await;
+
+			if inserted {
+				metrics::IN_FLIGHT_REQUEST_COUNT.set(self.in_flight_requests.entry_count() as i64);
+
 				return Ok(request_id);
 			}
 
@@ -697,91 +312,66 @@ impl ProxyState {
 
 // Helper function to choose a random target from a list of targets
 fn choose_random_target(targets: &[RouteTarget]) -> Option<&RouteTarget> {
-	if targets.is_empty() {
-		return None;
-	}
-
-	// Use a simple random index selection
-	let random_index = rand::random::<usize>() % targets.len();
-	targets.get(random_index)
+	targets.choose(&mut rand::thread_rng())
 }
 
 // Proxy service
 pub struct ProxyService {
 	state: Arc<ProxyState>,
 	remote_addr: SocketAddr,
-	// Note: Using the hyper legacy client is the only option currently.
-	// This is what reqwest uses under the hood. Eventually we'll migrate to h3 once it's ready.
-	client: Client<hyper_util::client::legacy::connect::HttpConnector, Full<Bytes>>,
+	connection_start: Instant,
 }
 
 impl ProxyService {
 	pub fn new(state: Arc<ProxyState>, remote_addr: SocketAddr) -> Self {
-		// Create a client with the hyper-util legacy client
-		let client = Client::builder(TokioExecutor::new())
-			.pool_idle_timeout(Duration::from_secs(30))
-			.build_http();
-
 		Self {
 			state,
 			remote_addr,
-			client,
+			connection_start: Instant::now(),
 		}
 	}
 
-	// Calculate backoff duration for a given retry attempt
-	pub fn calculate_backoff(attempt: u32, initial_interval: u64) -> Duration {
-		Duration::from_millis(initial_interval * 2u64.pow(attempt - 1))
-	}
+	/// Process an individual request.
+	#[tracing::instrument(name = "guard_request", skip_all, fields(ray_id, req_id, uri=%req.uri()))]
+	pub async fn process(&self, mut req: Request<BodyIncoming>) -> Result<Response<ResponseBody>> {
+		let start_time = Instant::now();
 
-	#[tracing::instrument(skip_all)]
-	async fn handle_request(
-		&self,
-		req: Request<BodyIncoming>,
-		start_time: Instant,
-		request_context: &mut RequestContext,
-	) -> Result<Response<ResponseBody>> {
+		let request_ids = RequestIds::new(self.state.config.dc_label());
+		req.extensions_mut().insert(request_ids);
+
+		let current_span = tracing::Span::current();
+
+		current_span.record("req_id", request_ids.req_id.to_string());
+		current_span.record("ray_id", request_ids.ray_id.to_string());
+
+		// Extract request information for logging and analytics before consuming the request
+		let incoming_ray_id = req
+			.headers()
+			.get(X_RIVET_RAY_ID)
+			.and_then(|h| h.to_str().ok())
+			.and_then(|id| Id::parse(id).ok());
 		let host = req
 			.headers()
 			.get(hyper::header::HOST)
 			.and_then(|h| h.to_str().ok())
-			.unwrap_or("unknown");
-
+			.unwrap_or("unknown")
+			.to_string();
+		let uri_string = req.uri().to_string();
 		let path = req
 			.uri()
 			.path_and_query()
 			.map(|x| x.to_string())
 			.unwrap_or_else(|| req.uri().path().to_string());
+		let method = req.method().clone();
 
-		// Set request body size in analytics (will be updated with actual size later)
-
-		let target_res = self
-			.state
-			.resolve_route(
-				host,
-				&path,
-				req.method(),
-				self.state.port_type.clone(),
-				req.headers(),
-				false,
-			)
-			.await;
-
-		let duration_secs = start_time.elapsed().as_secs_f64();
-		metrics::RESOLVE_ROUTE_DURATION.record(duration_secs, &[]);
-
-		// Resolve target
-		let target = target_res?;
-		if let ResolveRouteOutput::Response(response) = &target {
-			// Return the custom response
-			return response.build_response();
-		}
+		current_span.set_attribute("http.request.method", method.to_string());
+		current_span.set_attribute("http.path", uri_string.clone());
 
-		let actor_id = if let ResolveRouteOutput::Target(target) = &target {
-			target.actor_id
-		} else {
-			None
-		};
+		let user_agent = req
+			.headers()
+			.get(hyper::header::USER_AGENT)
+			.and_then(|h| h.to_str().ok())
+			.map(|s| s.to_string());
 
 		// Extract IP address from X-Forwarded-For header or fall back to remote_addr
 		let client_ip = req
@@ -795,74 +385,327 @@ impl ProxyService {
 			.and_then(|ip_str| ip_str.parse::<std::net::IpAddr>().ok())
 			.unwrap_or_else(|| self.remote_addr.ip());
 
-		// Apply rate limiting
-		if !self
-			.state
-			.check_rate_limit(client_ip, &actor_id, req.headers())
-			.await?
-		{
-			return Err(errors::RateLimit {
-				actor_id,
-				method: req.method().to_string(),
-				path: path.clone(),
-				ip: client_ip.to_string(),
-			}
-			.build());
-		}
+		let is_websocket = hyper_tungstenite::is_upgrade_request(&req);
+		let mut req_ctx = RequestContext::new(
+			self.remote_addr,
+			request_ids.ray_id,
+			request_ids.req_id,
+			host,
+			path,
+			req.method().clone(),
+			req.headers().clone(),
+			is_websocket,
+			client_ip,
+			start_time,
+		);
 
-		// Acquire in-flight limit and generate request ID
-		let request_id = match self
-			.state
-			.acquire_in_flight(client_ip, &actor_id, req.headers())
-			.await?
-		{
-			Some(id) => id,
-			None => {
-				return Err(errors::RateLimit {
-					actor_id,
-					method: req.method().to_string(),
-					path: path.clone(),
-					ip: client_ip.to_string(),
-				}
-				.build());
-			}
-		};
+		// TLS information would be set here if available (for HTTPS connections)
+		// This requires TLS connection introspection and is marked for future enhancement
 
-		// Increment metrics
-		metrics::PROXY_REQUEST_PENDING.add(1, &[]);
-		metrics::PROXY_REQUEST_TOTAL.add(1, &[]);
+		// Debug log request information with structured fields (Apache-like access log)
+		tracing::debug!(
+			?incoming_ray_id,
+			ray_id=?req_ctx.ray_id,
+			req_id=?req_ctx.req_id,
+			method=%req_ctx.method,
+			path=%req_ctx.path,
+			host=%req_ctx.host,
+			remote_addr=%req_ctx.remote_addr,
+			uri=%uri_string,
+			user_agent=?user_agent,
+			"Request received"
+		);
 
-		// Update request context with target info
-		if let Some(actor_id) = actor_id {
-			request_context.service_actor_id = Some(actor_id);
+		// Used for ws error proxying later
+		let mut mock_req_builder = Request::builder()
+			.method(req.method().clone())
+			.uri(req.uri().clone())
+			.version(req.version().clone());
+		if let Some(headers) = mock_req_builder.headers_mut() {
+			*headers = req.headers().clone();
 		}
+		if let Some(extensions) = mock_req_builder.extensions_mut() {
+			*extensions = req.extensions().clone();
+		}
+		let mock_req = mock_req_builder.body(())?;
 
-		let res = if hyper_tungstenite::is_upgrade_request(&req) {
-			self.handle_websocket_upgrade(req, target, request_context, client_ip, actor_id)
-				.await
-		} else {
-			self.handle_http_request(req, target, request_context, client_ip, actor_id)
-				.await
-		};
-
-		let status = match &res {
-			Ok(resp) => resp.status().as_u16().to_string(),
-			Err(_) => "error".to_string(),
-		};
-
-		// Record metrics
-		let duration_secs = start_time.elapsed().as_secs_f64();
-		metrics::PROXY_REQUEST_DURATION
-			.record(duration_secs, &[KeyValue::new("status", status.clone())]);
+		// Process the request
+		let mut res = match self.handle_request(req, &mut req_ctx).await {
+			Ok(res) => res,
+			Err(err) => {
+				// Log the error
+				tracing::error!(?err, "Request failed");
 
-		metrics::PROXY_REQUEST_PENDING.add(-1, &[]);
+				metrics::PROXY_REQUEST_ERROR_TOTAL
+					.with_label_values(&[&err.to_string()])
+					.inc();
 
-		// Release in-flight counter and request ID when done
-		let state_clone = self.state.clone();
-		tokio::spawn(
+				// If we receive an error during a websocket request, we attempt to open the websocket anyway
+				// so we can send the error via websocket instead of http. Most websocket clients don't handle
+				// HTTP errors in a meaningful way resulting in unhelpful errors for the user
+				if is_websocket {
+					tracing::debug!("Upgrading client connection to WebSocket for error proxy");
+					match hyper_tungstenite::upgrade(mock_req, None) {
+						Ok((client_response, client_ws)) => {
+							tracing::debug!("Client WebSocket upgrade for error proxy successful");
+
+							self.state.tasks.spawn(
+								async move {
+									let ws_handle = match WebSocketHandle::new(client_ws).await {
+										Ok(ws_handle) => ws_handle,
+										Err(err) => {
+											tracing::debug!(
+												?err,
+												"failed initiating websocket handle for error proxy"
+											);
+											return;
+										}
+									};
+									let frame = utils::err_to_close_frame(err, request_ids.ray_id);
+
+									// Manual conversion to handle different tungstenite versions
+									let code_num: u16 = frame.code.into();
+									let reason = frame.reason.clone();
+
+									if let Err(err) = ws_handle
+										.send(tokio_tungstenite::tungstenite::Message::Close(Some(
+											tokio_tungstenite::tungstenite::protocol::CloseFrame {
+												code: code_num.into(),
+												reason,
+											},
+										)))
+										.await
+									{
+										tracing::debug!(
+											?err,
+											"failed sending websocket error proxy"
+										);
+									}
+
+									// Flush to ensure close frame is sent
+									if let Err(err) = ws_handle.flush().await {
+										tracing::debug!(
+											?err,
+											"failed flushing websocket in error proxy"
+										);
+									}
+
+									// Keep TCP connection open briefly to allow client to process close
+									tokio::time::sleep(WEBSOCKET_CLOSE_LINGER).await;
+								}
+								.instrument(tracing::info_span!("ws_error_proxy_task")),
+							);
+
+							// Return the response that will upgrade the client connection
+							// For proper WebSocket handshaking, we need to preserve the original response
+							// structure but convert it to our expected return type without modifying its content
+							tracing::debug!(
+								"Returning WebSocket upgrade response for error proxy to client"
+							);
+							// Extract the parts from the response but preserve all headers and status
+							let (mut parts, _) = client_response.into_parts();
+
+							// Add Sec-WebSocket-Protocol header to the response
+							// Many WebSocket clients (e.g. node-ws & Cloudflare) require a protocol in the response
+							parts.headers.insert(
+								"sec-websocket-protocol",
+								hyper::header::HeaderValue::from_static("rivet"),
+							);
+
+							// Create a new response with an empty body - WebSocket upgrades don't need a body
+							Response::from_parts(
+								parts,
+								ResponseBody::Full(Full::<Bytes>::new(Bytes::new())),
+							)
+						}
+						Err(err) => {
+							tracing::error!(
+								?err,
+								"Failed to upgrade client WebSocket for error proxy"
+							);
+
+							utils::err_into_response(
+								errors::ConnectionError {
+									error_message: format!(
+										"Failed to upgrade client WebSocket for error proxy: {}",
+										err
+									),
+									remote_addr: req_ctx.remote_addr.to_string(),
+								}
+								.build(),
+							)?
+						}
+					}
+				} else {
+					utils::err_into_response(err)?
+				}
+			}
+		};
+
+		if is_websocket && res.status() != StatusCode::SWITCHING_PROTOCOLS {
+			tracing::debug!("returned non-101 response to websocket");
+		}
+
+		// Add ray_id to response headers
+		if let Ok(ray_id_value) = request_ids.ray_id.to_string().parse() {
+			if let Some(existing_ray_id_value) = res
+				.headers()
+				.get(X_RIVET_RAY_ID)
+				.and_then(|h| h.to_str().ok())
+			{
+				if ray_id_value != existing_ray_id_value {
+					tracing::warn!(
+						expected_ray_id=%request_ids.ray_id,
+						received_ray_id=%existing_ray_id_value,
+						"downstream service set ray id header to a different value",
+					);
+				}
+			}
+
+			res.headers_mut().insert(X_RIVET_RAY_ID, ray_id_value);
+		}
+
+		// Add cors headers to response
+		if let Some(cors) = &req_ctx.cors {
+			let headers = res.headers_mut();
+
+			headers.insert(
+				"access-control-allow-origin",
+				HeaderValue::from_str(&cors.allow_origin)?,
+			);
+			headers.insert(
+				"access-control-allow-credentials",
+				HeaderValue::from_static(if cors.allow_credentials {
+					"true"
+				} else {
+					"false"
+				}),
+			);
+			headers.insert(
+				"access-control-expose-headers",
+				HeaderValue::from_str(&cors.expose_headers)?,
+			);
+
+			if let Some(allow_methods) = &cors.allow_methods {
+				headers.insert(
+					"access-control-allow-methods",
+					HeaderValue::from_str(allow_methods)?,
+				);
+			}
+
+			if let Some(allow_headers) = &cors.allow_headers {
+				headers.insert(
+					"access-control-allow-headers",
+					HeaderValue::from_str(allow_headers)?,
+				);
+			}
+
+			if let Some(max_age) = &cors.max_age {
+				headers.insert(
+					"access-control-max-age",
+					HeaderValue::from_str(&max_age.to_string())?,
+				);
+			}
+
+			// Add Vary header to prevent cache poisoning when echoing origin
+			if cors.allow_origin != "*" {
+				headers.insert("vary", HeaderValue::from_static("Origin"));
+			}
+		}
+
+		// Set span status code
+		let status = res.status().as_u16();
+		current_span.set_attribute("http.response.status_code", status as i64);
+
+		let content_length = res
+			.headers()
+			.get(hyper::header::CONTENT_LENGTH)
+			.and_then(|h| h.to_str().ok())
+			.and_then(|s| s.parse::<usize>().ok())
+			.unwrap_or(0);
+
+		// Log information about the completed request
+		tracing::debug!(
+			?incoming_ray_id,
+			ray_id=?req_ctx.ray_id,
+			req_id=?req_ctx.req_id,
+			method = %req_ctx.method,
+			path = %req_ctx.path,
+			host = %req_ctx.host,
+			remote_addr = %req_ctx.remote_addr,
+			status = %status,
+			content_length = %content_length,
+			"Request completed"
+		);
+
+		Ok(res)
+	}
+
+	#[tracing::instrument(skip_all)]
+	async fn handle_request(
+		&self,
+		req: Request<BodyIncoming>,
+		req_ctx: &mut RequestContext,
+	) -> Result<Response<ResponseBody>> {
+		// Resolve target
+		let target_res = self.state.resolve_route(req_ctx, false).await;
+
+		let duration_secs = req_ctx.start_time.elapsed().as_secs_f64();
+		metrics::RESOLVE_ROUTE_DURATION.observe(duration_secs);
+
+		let target = target_res?;
+
+		// Apply rate limiting
+		if self.state.check_rate_limit(req_ctx).await? {
+			return Err(errors::RateLimit {
+				method: req_ctx.method.to_string(),
+				path: req_ctx.path.clone(),
+				ip: req_ctx.client_ip.to_string(),
+			}
+			.build());
+		}
+
+		// Acquire in-flight limit and generate protocol request ID
+		if self.state.acquire_in_flight(req_ctx).await? {
+			return Err(errors::RateLimit {
+				method: req_ctx.method.to_string(),
+				path: req_ctx.path.clone(),
+				ip: req_ctx.client_ip.to_string(),
+			}
+			.build());
+		}
+
+		// Increment metrics
+		metrics::PROXY_REQUEST_PENDING.inc();
+		metrics::PROXY_REQUEST_TOTAL.inc();
+
+		let res = if hyper_tungstenite::is_upgrade_request(&req) {
+			self.handle_websocket_upgrade(req, req_ctx, target).await
+		} else {
+			self.handle_http_request(req, req_ctx, target).await
+		};
+
+		let status = match &res {
+			Ok(resp) => resp.status().as_u16().to_string(),
+			Err(_) => "error".to_string(),
+		};
+
+		// Record metrics
+		let duration_secs = req_ctx.start_time.elapsed().as_secs_f64();
+		metrics::PROXY_REQUEST_DURATION
+			.with_label_values(&[status])
+			.observe(duration_secs);
+
+		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, &actor_id, request_id)
+					.release_in_flight(client_ip, in_flight_request_id)
 					.await;
 			}
 			.instrument(tracing::info_span!("release_in_flight_task")),
@@ -875,95 +718,40 @@ impl ProxyService {
 	async fn handle_http_request(
 		&self,
 		req: Request<BodyIncoming>,
+		req_ctx: &mut RequestContext,
 		resolved_route: ResolveRouteOutput,
-		request_context: &mut RequestContext,
-		client_ip: std::net::IpAddr,
-		actor_id: Option<Id>,
 	) -> Result<Response<ResponseBody>> {
-		// Get middleware config for this actor if it exists
-		let middleware_config = if let ResolveRouteOutput::Target(target) = &resolved_route
-			&& let Some(actor_id) = &target.actor_id
-		{
-			self.state
-				.get_middleware_config(actor_id, req.headers())
-				.await?
-		} else {
-			// Default middleware config for targets without actor_id
-			MiddlewareConfig {
-				rate_limit: RateLimitConfig {
-					requests: 100, // 100 requests
-					period: 60,    // per 60 seconds
-				},
-				max_in_flight: MaxInFlightConfig {
-					amount: 20, // 20 concurrent requests
-				},
-				retry: RetryConfig {
-					max_attempts: 3,       // 3 retry attempts
-					initial_interval: 100, // 100ms initial interval
-				},
-				timeout: TimeoutConfig {
-					request_timeout: 30, // 30 seconds for requests
-				},
-			}
-		};
-
-		let host = req
-			.headers()
-			.get(hyper::header::HOST)
-			.and_then(|h| h.to_str().ok())
-			.unwrap_or("unknown")
-			.to_string();
-
-		let path = req
-			.uri()
-			.path_and_query()
-			.map(|x| x.to_string())
-			.unwrap_or_else(|| req.uri().path().to_string());
-
-		// Set up retry with backoff from middleware config
-		let max_attempts = middleware_config.retry.max_attempts;
-		let initial_interval = middleware_config.retry.initial_interval;
-		let timeout_duration = Duration::from_secs(middleware_config.timeout.request_timeout);
+		// Set up retry with backoff
+		let timeout_duration = Duration::from_secs(req_ctx.timeout.request_timeout);
 
 		match resolved_route {
 			ResolveRouteOutput::Target(mut target) => {
-				// Set service IP from target
-				if let Ok(target_ip) =
-					format!("{}:{}", target.host, target.port).parse::<std::net::SocketAddr>()
-				{
-					request_context.service_ip = Some(target_ip.ip());
-				}
-
 				// Read the request body before proceeding with retries
 				let (req_parts, body) = req.into_parts();
-				let req_body = match http_body_util::BodyExt::collect(body).await {
-					Ok(collected) => collected.to_bytes(),
-					Err(err) => {
-						tracing::debug!(?err, "Failed to read request body");
-						Bytes::new()
-					}
-				};
-
-				// Set actual request body size in analytics
-				request_context.client_request_body_bytes = Some(req_body.len() as u64);
+				let req_body =
+					Limited::new(body, self.state.config.guard().http_max_request_body_size())
+						.collect()
+						.await
+						.map_err(|err| errors::InvalidRequestBody(err.to_string()).build())?
+						.to_bytes();
 
 				// Use a value-returning loop to handle both errors and successful responses
 				let mut attempts = 0;
-				while attempts < max_attempts {
+				while attempts < req_ctx.retry.max_attempts {
 					attempts += 1;
 
 					// Use the common function to build request parts
-					let builder = self
-						.proxied_request_builder(&req_parts, &target)
+					let builder = utils::proxied_request_builder(&req_parts, req_ctx, &target)
 						.map_err(|err| errors::HttpRequestBuildFailed(err.to_string()).build())?;
 
 					// Create the final request with body
 					let proxied_req = builder
-						.body(Full::<Bytes>::new(req_body.clone()))
+						// NOTE: the `Bytes` type is cheaply cloneable, this is not resource intensive
+						.body(Full::new(req_body.clone()))
 						.map_err(|err| errors::RequestBuildError(err.to_string()).build())?;
 
 					// Send the request with timeout
-					let res = timeout(timeout_duration, self.client.request(proxied_req))
+					let res = timeout(timeout_duration, self.state.client.request(proxied_req))
 						.await
 						.map_err(|_| {
 							errors::RequestTimeout {
@@ -975,29 +763,23 @@ impl ProxyService {
 					match res {
 						Ok(resp) => {
 							// Check if this is a retryable response
-							if should_retry_request_inner(resp.status(), resp.headers()) {
+							if utils::should_retry_request_inner(resp.status(), resp.headers()) {
 								// Request connect error, might retry
 								tracing::debug!(
 									"Request attempt {attempts} failed (service unavailable)"
 								);
 
 								// Use backoff and continue
-								let backoff = Self::calculate_backoff(attempts, initial_interval);
+								let backoff = utils::calculate_backoff(
+									attempts,
+									req_ctx.retry.initial_interval,
+								);
 								tokio::time::sleep(backoff).await;
 
 								// Resolve target again, this time ignoring cache. This makes sure
 								// we always re-fetch the route on error
-								let ResolveRouteOutput::Target(new_target) = self
-									.state
-									.resolve_route(
-										&host,
-										&path,
-										&req_parts.method,
-										self.state.port_type.clone(),
-										&req_parts.headers,
-										true,
-									)
-									.await?
+								let ResolveRouteOutput::Target(new_target) =
+									self.state.resolve_route(req_ctx, true).await?
 								else {
 									bail!("resolved route does not match Target");
 								};
@@ -1019,28 +801,27 @@ impl ProxyService {
 								// For streaming responses, pass through the body without buffering
 								tracing::debug!("Detected streaming response, preserving stream");
 
-								// We can't easily calculate response size for streaming, so set it to None
-								request_context.guard_response_body_bytes = None;
-
 								let streaming_body = ResponseBody::Incoming(body);
 								return Ok(Response::from_parts(parts, streaming_body));
 							} else {
 								// For non-streaming responses, buffer as before
-								let body_bytes = match BodyExt::collect(body).await {
-									Ok(collected) => collected.to_bytes(),
-									Err(_) => Bytes::new(),
-								};
-
-								// Set actual response body size in analytics
-								request_context.guard_response_body_bytes =
-									Some(body_bytes.len() as u64);
+								let body_bytes = Limited::new(
+									body,
+									self.state.config.guard().http_max_request_body_size(),
+								)
+								.collect()
+								.await
+								.map_err(|err| {
+									errors::InvalidResponseBody(err.to_string()).build()
+								})?
+								.to_bytes();
 
 								let full_body = ResponseBody::Full(Full::new(body_bytes));
 								return Ok(Response::from_parts(parts, full_body));
 							}
 						}
 						Err(err) => {
-							if !err.is_connect() || attempts >= max_attempts {
+							if !err.is_connect() || attempts >= req_ctx.retry.max_attempts {
 								tracing::error!(
 									?err,
 									?target,
@@ -1057,22 +838,16 @@ impl ProxyService {
 								tracing::debug!(?err, "Request attempt {attempts} failed");
 
 								// Use backoff and continue
-								let backoff = Self::calculate_backoff(attempts, initial_interval);
+								let backoff = utils::calculate_backoff(
+									attempts,
+									req_ctx.retry.initial_interval,
+								);
 								tokio::time::sleep(backoff).await;
 
 								// Resolve target again, this time ignoring cache. This makes sure
 								// we always re-fetch the route on error
-								let ResolveRouteOutput::Target(new_target) = self
-									.state
-									.resolve_route(
-										&host,
-										&path,
-										&req_parts.method,
-										self.state.port_type.clone(),
-										&req_parts.headers,
-										true,
-									)
-									.await?
+								let ResolveRouteOutput::Target(new_target) =
+									self.state.resolve_route(req_ctx, true).await?
 								else {
 									bail!("resolved route does not match Target");
 								};
@@ -1086,77 +861,40 @@ impl ProxyService {
 
 				// If we get here, all attempts failed
 				return Err(errors::RetryAttemptsExceeded {
-					attempts: max_attempts,
+					attempts: req_ctx.retry.max_attempts,
 				}
 				.build());
 			}
-			ResolveRouteOutput::Response(_) => {
-				unreachable!()
-			}
 			ResolveRouteOutput::CustomServe(mut handler) => {
-				let req_headers = req.headers().clone();
-				let req_method = req.method().clone();
-
-				// Acquire in-flight limit and generate request ID
-				let request_id = match self
-					.state
-					.acquire_in_flight(client_ip, &actor_id, &req_headers)
-					.await?
-				{
-					Some(id) => id,
-					None => {
-						return Err(errors::RateLimit {
-							actor_id,
-							method: req_method.to_string(),
-							path: path.clone(),
-							ip: client_ip.to_string(),
-						}
-						.build());
-					}
-				};
-
 				// Collect request body
 				let (req_parts, body) = req.into_parts();
-				let collected_body = match http_body_util::BodyExt::collect(body).await {
-					Ok(collected) => collected.to_bytes(),
-					Err(err) => {
-						tracing::debug!(?err, "Failed to read request body");
-						Bytes::new()
-					}
-				};
-				let req_collected = hyper::Request::from_parts(
-					req_parts,
-					Full::<Bytes>::new(collected_body.clone()),
-				);
+				let req_body =
+					Limited::new(body, self.state.config.guard().http_max_request_body_size())
+						.collect()
+						.await
+						.map_err(|err| errors::InvalidRequestBody(err.to_string()).build())?
+						.to_bytes();
+				let req_collected =
+					hyper::Request::from_parts(req_parts, Full::<Bytes>::new(req_body));
 
 				// Attempt request
 				let mut attempts = 0;
-				while attempts < max_attempts {
+				while attempts < req_ctx.retry.max_attempts {
 					attempts += 1;
 
-					let res = handler
-						.handle_request(req_collected.clone(), request_context, request_id)
-						.await;
-					if should_retry_request(&res) {
+					let res = handler.handle_request(req_collected.clone(), req_ctx).await;
+					if utils::should_retry_request(&res) {
 						// Request connect error, might retry
 						tracing::debug!("Request attempt {attempts} failed (service unavailable)");
 
 						// Use backoff and continue
-						let backoff = Self::calculate_backoff(attempts, initial_interval);
+						let backoff =
+							utils::calculate_backoff(attempts, req_ctx.retry.initial_interval);
 						tokio::time::sleep(backoff).await;
 
 						// Refresh route (ignore cache) so subsequent requests can hit new target
-						let ResolveRouteOutput::CustomServe(new_handler) = self
-							.state
-							.resolve_route(
-								&host,
-								&path,
-								req_collected.method(),
-								self.state.port_type.clone(),
-								&req_headers,
-								true,
-							)
-							.await?
+						let ResolveRouteOutput::CustomServe(new_handler) =
+							self.state.resolve_route(req_ctx, true).await?
 						else {
 							bail!("resolved route does not match CustomServe");
 						};
@@ -1167,7 +905,7 @@ impl ProxyService {
 
 					// Release in-flight counter and request ID before returning
 					self.state
-						.release_in_flight(client_ip, &actor_id, request_id)
+						.release_in_flight(req_ctx.client_ip, req_ctx.in_flight_request_id)
 						.await;
 					return res;
 				}
@@ -1175,130 +913,33 @@ impl ProxyService {
 				// If we get here, all attempts failed
 				// Release in-flight counter and request ID before returning error
 				self.state
-					.release_in_flight(client_ip, &actor_id, request_id)
+					.release_in_flight(req_ctx.client_ip, req_ctx.in_flight_request_id)
 					.await;
 				return Err(errors::RetryAttemptsExceeded {
-					attempts: max_attempts,
+					attempts: req_ctx.retry.max_attempts,
 				}
 				.build());
 			}
 		}
 	}
 
-	/// Modifies the incoming request before it is proxied.
-	fn proxied_request_builder(
-		&self,
-		req_parts: &hyper::http::request::Parts,
-		target: &RouteTarget,
-	) -> Result<hyper::http::request::Builder> {
-		let scheme = if target.port == 443 { "https" } else { "http" };
-
-		// Bracket raw IPv6 hosts
-		let host = if target.host.contains(':') && !target.host.starts_with('[') {
-			format!("[{}]", target.host)
-		} else {
-			target.host.clone()
-		};
-
-		// Ensure path starts with a leading slash
-		let path = if target.path.starts_with('/') {
-			target.path.clone()
-		} else {
-			format!("/{}", target.path)
-		};
-
-		let url = Url::parse(&format!("{scheme}://{host}:{}{}", target.port, path))
-			.context("invalid scheme/host/port when building URL")?;
-
-		// Build the proxied request
-		let mut builder = hyper::Request::builder()
-			.method(req_parts.method.clone())
-			.uri(url.to_string());
-
-		// Modify proxy headers
-		let headers = builder
-			.headers_mut()
-			.expect("request builder unexpectedly in error state");
-
-		headers.remove(X_RIVET_TARGET);
-		headers.remove(X_RIVET_ACTOR);
-		headers.remove(X_RIVET_TOKEN);
-
-		add_proxy_headers_with_addr(headers, &req_parts.headers, self.remote_addr)?;
-
-		Ok(builder)
-	}
-
 	#[tracing::instrument(skip_all)]
 	async fn handle_websocket_upgrade(
 		&self,
 		req: Request<BodyIncoming>,
+		req_ctx: &mut RequestContext,
 		target: ResolveRouteOutput,
-		request_context: &mut RequestContext,
-		client_ip: std::net::IpAddr,
-		actor_id: Option<Id>,
 	) -> Result<Response<ResponseBody>> {
-		// Parsed for retries later
-		let req_host = req
-			.headers()
-			.get(hyper::header::HOST)
-			.and_then(|h| h.to_str().ok())
-			.unwrap_or("unknown")
-			.to_string();
-		let req_path = req
-			.uri()
-			.path_and_query()
-			.map(|x| x.to_string())
-			.unwrap_or_else(|| req.uri().path().to_string());
-
-		// Capture headers and method before request is consumed
-		let req_headers = req.headers().clone();
-		let req_method = req.method().clone();
-		let ray_id = req.extensions().get::<RequestIds>().map(|x| x.ray_id);
-
-		// Get middleware config for this actor if it exists
-		let middleware_config = match &actor_id {
-			Some(actor_id) => {
-				self.state
-					.get_middleware_config(actor_id, &req_headers)
-					.await?
-			}
-			None => {
-				// Default middleware config for targets without actor_id
-				tracing::debug!("Using default middleware config (no actor_id)");
-				MiddlewareConfig {
-					rate_limit: RateLimitConfig {
-						requests: 100, // 100 requests
-						period: 60,    // per 60 seconds
-					},
-					max_in_flight: MaxInFlightConfig {
-						amount: 20, // 20 concurrent requests
-					},
-					retry: RetryConfig {
-						max_attempts: 3,       // 3 retry attempts
-						initial_interval: 100, // 100ms initial interval
-					},
-					timeout: TimeoutConfig {
-						request_timeout: 30, // 30 seconds for requests
-					},
-				}
-			}
-		};
-
-		// Set up retry with backoff from middleware config
-		let max_attempts = middleware_config.retry.max_attempts;
-		let initial_interval = middleware_config.retry.initial_interval;
-
 		// Log the headers for debugging
 		tracing::debug!("WebSocket upgrade request headers:");
-		for (name, value) in req.headers() {
+		for (name, value) in &req_ctx.headers {
 			if let Ok(val) = value.to_str() {
 				tracing::debug!("  {}: {}", name, val);
 			}
 		}
 
 		// Handle WebSocket upgrade properly with hyper_tungstenite
-		tracing::debug!(%req_path, "Upgrading client connection to WebSocket");
+		tracing::debug!(path=%req_ctx.path, "Upgrading client connection to WebSocket");
 		let (client_response, client_ws) = match hyper_tungstenite::upgrade(req, None) {
 			Ok(x) => {
 				tracing::debug!("Client WebSocket upgrade successful");
@@ -1308,7 +949,7 @@ impl ProxyService {
 				tracing::error!(?err, "Failed to upgrade client WebSocket");
 				return Err(errors::ConnectionError {
 					error_message: format!("Failed to upgrade client WebSocket: {}", err),
-					remote_addr: self.remote_addr.to_string(),
+					remote_addr: req_ctx.remote_addr.to_string(),
 				}
 				.build());
 			}
@@ -1327,14 +968,17 @@ impl ProxyService {
 
 		// Clone needed values for the spawned task
 		let state = self.state.clone();
-		let remote_addr = self.remote_addr;
 
 		// Spawn a new task to handle the WebSocket bidirectional communication
 		match target {
 			ResolveRouteOutput::Target(mut target) => {
 				tracing::debug!("Spawning task to handle WebSocket communication");
+				let mut req_ctx = req_ctx.clone();
+
 				self.state.tasks.spawn(
 					async move {
+						let req_ctx = &mut req_ctx;
+
 						// Set up a timeout for the entire operation
 						let timeout_duration = Duration::from_secs(30); // 30 seconds timeout
 						tracing::debug!(
@@ -1369,7 +1013,7 @@ impl ProxyService {
 
 						// Now attempt to connect to the upstream server
 						tracing::debug!("Attempting connect to upstream WebSocket");
-						while attempts < max_attempts {
+						while attempts < req_ctx.retry.max_attempts {
 							attempts += 1;
 
 							// Build the WebSocket URL using the url crate to properly handle IPv6 addresses
@@ -1414,7 +1058,7 @@ impl ProxyService {
 							tracing::debug!(
 								"WebSocket request attempt {}/{} to {}",
 								attempts,
-								max_attempts,
+								req_ctx.retry.max_attempts,
 								target_url
 							);
 
@@ -1428,10 +1072,9 @@ impl ProxyService {
 							};
 
 							// Add proxy headers to the websocket request
-							if let Err(err) = add_proxy_headers_with_addr(
+							if let Err(err) = utils::add_proxy_headers_with_addr(
 								ws_request.headers_mut(),
-								&req_headers,
-								remote_addr,
+								req_ctx,
 							) {
 								tracing::error!(
 									?err,
@@ -1485,21 +1128,23 @@ impl ProxyService {
 							}
 
 							// Check if we've reached max attempts
-							if attempts >= max_attempts {
+							if attempts >= req_ctx.retry.max_attempts {
 								tracing::debug!(
 									"All {} WebSocket connection attempts failed",
-									max_attempts
+									req_ctx.retry.max_attempts
 								);
 
 								// Send a close message to the client since we can't connect to upstream
-								tracing::debug!(
-									"Sending close message to client due to upstream connection failure"
+								let err = errors::RetryAttemptsExceeded { attempts }.build();
+								tracing::warn!(
+									?err,
+									"sending close message to client due to upstream connection failure"
 								);
 								let (mut client_sink, _) = client_ws.split();
 								match client_sink
-									.send(to_hyper_close(Some(err_to_close_frame(
-										errors::RetryAttemptsExceeded { attempts }.build(),
-										ray_id,
+									.send(utils::to_hyper_close(Some(utils::err_to_close_frame(
+										err,
+										req_ctx.ray_id,
 									))))
 									.await
 								{
@@ -1532,7 +1177,8 @@ impl ProxyService {
 							}
 
 							// Use backoff for the next attempt
-							let backoff = Self::calculate_backoff(attempts, initial_interval);
+							let backoff =
+								utils::calculate_backoff(attempts, req_ctx.retry.initial_interval);
 							tracing::debug!(
 								"Waiting for {:?} before next connection attempt",
 								backoff
@@ -1542,41 +1188,20 @@ impl ProxyService {
 
 							// Resolve target again, this time ignoring cache. This makes sure
 							// we always re-fetch the route on error
-							let new_target = state
-								.resolve_route(
-									&req_host,
-									&req_path,
-									&req_method,
-									state.port_type.clone(),
-									&req_headers,
-									true,
-								)
-								.await;
+							let new_target = state.resolve_route(req_ctx, true).await;
 
 							match new_target {
 								Ok(ResolveRouteOutput::Target(new_target)) => {
 									target = new_target;
 								}
-								Ok(ResolveRouteOutput::Response(response)) => {
-									tracing::debug!(
-										status=?response.status,
-										message=?response.message,
-										docs=?response.docs,
-										"got response instead of websocket target",
-									);
-
-									// Close the WebSocket connection with the response message
-									let _ = client_ws
-										.close(Some(str_to_close_frame(response.message.as_ref())))
-										.await;
-									return;
-								}
 								Ok(ResolveRouteOutput::CustomServe(_)) => {
+									let err = errors::WebSocketTargetChanged.build();
+									tracing::warn!(
+										?err,
+										"websocket target changed to custom serve"
+									);
 									let _ = client_ws
-										.close(Some(err_to_close_frame(
-											errors::WebSocketTargetChanged.build(),
-											ray_id,
-										)))
+										.close(Some(utils::err_to_close_frame(err, req_ctx.ray_id)))
 										.await;
 									return;
 								}
@@ -1658,7 +1283,7 @@ impl ProxyService {
 														// Signal shutdown to other direction
 														let _ = shutdown_tx.send(true);
 
-														to_hyper_close(frame)
+														utils::to_hyper_close(frame)
 													},
 													hyper_tungstenite::tungstenite::Message::Frame(_) => {
 														// Skip frames - they're an implementation detail
@@ -1809,7 +1434,7 @@ impl ProxyService {
 														// Signal shutdown to other direction
 														let _ = shutdown_tx.send(true);
 
-														to_hyper_close(frame)
+														utils::to_hyper_close(frame)
 													},
 													tokio_tungstenite::tungstenite::Message::Frame(_) => {
 														// Skip frames - they're an implementation detail
@@ -1879,7 +1504,7 @@ impl ProxyService {
 
 							// Try to send a close frame - ignore errors as the connection might already be closed
 							tracing::trace!("Attempting to send close message to client");
-							match sink.send(to_hyper_close(None)).await {
+							match sink.send(utils::to_hyper_close(None)).await {
 								Ok(_) => {
 									tracing::trace!("Close message sent to client successfully")
 								}
@@ -1908,34 +1533,14 @@ impl ProxyService {
 					.instrument(tracing::info_span!("handle_ws_task_target")),
 				);
 			}
-			ResolveRouteOutput::Response(_) => unreachable!(),
 			ResolveRouteOutput::CustomServe(mut handler) => {
-				tracing::debug!(%req_path, "Spawning task to handle WebSocket communication");
-				let mut request_context = request_context.clone();
-				let req_headers = req_headers.clone();
+				tracing::debug!(path=%req_ctx.path, "Spawning task to handle WebSocket communication");
 				let state = self.state.clone();
-				let req_path = req_path.clone();
-				let req_host = req_host.clone();
-				let req_method = req_method.clone();
+				let mut req_ctx = req_ctx.clone();
 
 				self.state.tasks.spawn(
 					async move {
-						let request_id = match state
-							.acquire_in_flight(client_ip, &actor_id, &req_headers)
-							.await?
-						{
-							Some(id) => id,
-							None => {
-								return Err(errors::RateLimit {
-									actor_id,
-									method: req_method.to_string(),
-									path: req_path.clone(),
-									ip: client_ip.to_string(),
-								}
-								.build()
-								.into());
-							}
-						};
+						let req_ctx = &mut req_ctx;
 						let mut ws_hibernation_close = false;
 						let mut after_hibernation = false;
 						let mut attempts = 0u32;
@@ -1946,14 +1551,7 @@ impl ProxyService {
 
 						loop {
 							match handler
-								.handle_websocket(
-									ws_handle.clone(),
-									&req_headers,
-									&req_path,
-									&mut request_context,
-									request_id,
-									after_hibernation,
-								)
+								.handle_websocket(req_ctx, ws_handle.clone(), after_hibernation)
 								.await
 							{
 								Ok(close_frame) => {
@@ -1962,7 +1560,7 @@ impl ProxyService {
 									// Send graceful close. This may fail if client already sent
 									// close frame, which is normal.
 									tracing::debug!(?close_frame, "sending close frame to client");
-									match ws_handle.send(to_hyper_close(close_frame)).await {
+									match ws_handle.send(utils::to_hyper_close(close_frame)).await {
 										Ok(_) => {
 											tracing::debug!("close frame sent successfully");
 										}
@@ -1997,7 +1595,7 @@ impl ProxyService {
 									tracing::debug!(?err, "websocket handler error");
 
 									// Denotes that the connection did not fail, but the downstream has closed
-									let ws_hibernate = is_ws_hibernate(&err);
+									let ws_hibernate = utils::is_ws_hibernate(&err);
 
 									if ws_hibernate {
 										attempts = 0;
@@ -2021,8 +1619,8 @@ impl ProxyService {
 										//   (starting with the message that caused the hibernation to end)
 										let res = handler
 											.handle_websocket_hibernation(
+												req_ctx,
 												ws_handle.clone(),
-												request_id,
 											)
 											.await?;
 
@@ -2036,20 +1634,21 @@ impl ProxyService {
 
 											ws_hibernation_close = true;
 										}
-									} else if attempts > max_attempts
-										|| !is_retryable_ws_error(&err)
+									} else if attempts > req_ctx.retry.max_attempts
+										|| !utils::is_retryable_ws_error(&err)
 									{
 										tracing::debug!(
+											?err,
 											?attempts,
-											?max_attempts,
-											"WebSocket failed"
+											max_attempts=?req_ctx.retry.max_attempts,
+											"websocket failed"
 										);
 
 										// Close WebSocket with error
 										ws_handle
-											.send(to_hyper_close(Some(err_to_close_frame(
-												err, ray_id,
-											))))
+											.send(utils::to_hyper_close(Some(
+												utils::err_to_close_frame(err, req_ctx.ray_id),
+											)))
 											.await?;
 
 										// Flush to ensure close frame is sent
@@ -2060,9 +1659,9 @@ impl ProxyService {
 
 										break;
 									} else {
-										let backoff = ProxyService::calculate_backoff(
+										let backoff = utils::calculate_backoff(
 											attempts,
-											initial_interval,
+											req_ctx.retry.initial_interval,
 										);
 
 										tracing::debug!(
@@ -2075,40 +1674,21 @@ impl ProxyService {
 									}
 
 									// Retry route resolution
-									match state
-										.resolve_route(
-											&req_host,
-											&req_path,
-											&req_method,
-											state.port_type.clone(),
-											&req_headers,
-											true,
-										)
-										.await
-									{
+									match state.resolve_route(req_ctx, true).await {
 										Ok(ResolveRouteOutput::CustomServe(new_handler)) => {
 											handler = new_handler;
 											continue;
 										}
-										Ok(ResolveRouteOutput::Response(response)) => {
-											ws_handle
-												.send(to_hyper_close(Some(str_to_close_frame(
-													response.message.as_ref(),
-												))))
-												.await?;
-
-											// Flush to ensure close frame is sent
-											ws_handle.flush().await?;
-
-											// Keep TCP connection open briefly to allow client to process close
-											tokio::time::sleep(WEBSOCKET_CLOSE_LINGER).await;
-										}
 										Ok(ResolveRouteOutput::Target(_)) => {
+											let err = errors::WebSocketTargetChanged.build();
+											tracing::warn!(
+												?err,
+												"websocket target changed to target"
+											);
 											ws_handle
-												.send(to_hyper_close(Some(err_to_close_frame(
-													errors::WebSocketTargetChanged.build(),
-													ray_id,
-												))))
+												.send(utils::to_hyper_close(Some(
+													utils::err_to_close_frame(err, req_ctx.ray_id),
+												)))
 												.await?;
 
 											// Flush to ensure close frame is sent
@@ -2120,10 +1700,14 @@ impl ProxyService {
 											break;
 										}
 										Err(err) => {
+											tracing::warn!(
+												?err,
+												"closing websocket due to route resolution error"
+											);
 											ws_handle
-												.send(to_hyper_close(Some(err_to_close_frame(
-													err, ray_id,
-												))))
+												.send(utils::to_hyper_close(Some(
+													utils::err_to_close_frame(err, req_ctx.ray_id),
+												)))
 												.await?;
 
 											// Flush to ensure close frame is sent
@@ -2141,7 +1725,7 @@ impl ProxyService {
 
 						// Release in-flight counter and request ID when task completes
 						state
-							.release_in_flight(client_ip, &actor_id, request_id)
+							.release_in_flight(req_ctx.client_ip, req_ctx.in_flight_request_id)
 							.await;
 
 						Ok(())
@@ -2173,318 +1757,12 @@ impl ProxyService {
 	}
 }
 
-impl ProxyService {
-	// Process an individual request
-	#[tracing::instrument(name = "guard_request", skip_all, fields(ray_id, req_id))]
-	pub async fn process(&self, mut req: Request<BodyIncoming>) -> Result<Response<ResponseBody>> {
-		let start_time = Instant::now();
-
-		let request_ids = RequestIds::new(self.state.config.dc_label());
-		req.extensions_mut().insert(request_ids);
-
-		tracing::Span::current()
-			.record("req_id", request_ids.req_id.to_string())
-			.record("ray_id", request_ids.ray_id.to_string());
-
-		// Create request context for analytics tracking
-		let mut request_context =
-			RequestContext::new(self.state.clickhouse_inserter.clone(), request_ids);
-
-		// Extract request information for logging and analytics before consuming the request
-		let incoming_ray_id = req
-			.headers()
-			.get(X_RIVET_RAY_ID)
-			.and_then(|h| h.to_str().ok())
-			.and_then(|id| Id::parse(id).ok());
-		let host = req
-			.headers()
-			.get(hyper::header::HOST)
-			.and_then(|h| h.to_str().ok())
-			.unwrap_or("unknown")
-			.to_string();
-		let uri_string = req.uri().to_string();
-		let path = req
-			.uri()
-			.path_and_query()
-			.map(|x| x.to_string())
-			.unwrap_or_else(|| req.uri().path().to_string());
-		let method = req.method().clone();
-
-		let user_agent = req
-			.headers()
-			.get(hyper::header::USER_AGENT)
-			.and_then(|h| h.to_str().ok())
-			.map(|s| s.to_string());
-
-		// Populate request context with available data
-		request_context.client_ip = Some(self.remote_addr.ip());
-		request_context.client_request_host = Some(host.clone());
-		request_context.client_request_method = Some(method.to_string());
-		request_context.client_request_path = Some(req.uri().path().to_string());
-		request_context.client_request_protocol = Some(format!("{:?}", req.version()));
-		request_context.client_request_scheme =
-			Some(req.uri().scheme_str().unwrap_or("http").to_string());
-		request_context.client_request_uri = Some(path.clone());
-		request_context.client_src_port = Some(self.remote_addr.port());
-
-		if let Some(referer) = req
-			.headers()
-			.get(hyper::header::REFERER)
-			.and_then(|h| h.to_str().ok())
-		{
-			request_context.client_request_referer = Some(referer.to_string());
-		}
-
-		if let Some(ua) = &user_agent {
-			request_context.client_request_user_agent = Some(ua.clone());
-		}
-
-		if let Some(requested_with) = req
-			.headers()
-			.get("x-requested-with")
-			.and_then(|h| h.to_str().ok())
-		{
-			request_context.client_x_requested_with = Some(requested_with.to_string());
-		}
-
-		// TLS information would be set here if available (for HTTPS connections)
-		// This requires TLS connection introspection and is marked for future enhancement
-
-		// Debug log request information with structured fields (Apache-like access log)
-		tracing::debug!(
-			?incoming_ray_id,
-			ray_id=?request_ids.ray_id,
-			req_id=?request_ids.req_id,
-			method = %method,
-			path = %path,
-			host = %host,
-			remote_addr = %self.remote_addr,
-			uri = %uri_string,
-			protocol = ?self.state.port_type,
-			user_agent = ?user_agent,
-			"Request received"
-		);
-
-		let is_websocket = hyper_tungstenite::is_upgrade_request(&req);
-
-		// Used for ws error proxying later
-		let mut mock_req_builder = Request::builder()
-			.method(req.method().clone())
-			.uri(req.uri().clone())
-			.version(req.version().clone());
-		if let Some(headers) = mock_req_builder.headers_mut() {
-			*headers = req.headers().clone();
-		}
-		if let Some(extensions) = mock_req_builder.extensions_mut() {
-			*extensions = req.extensions().clone();
-		}
-		let mock_req = mock_req_builder.body(())?;
-
-		// Process the request
-		let mut res = match self
-			.handle_request(req, start_time, &mut request_context)
-			.await
-		{
-			Ok(res) => res,
-			Err(err) => {
-				// Log the error
-				tracing::error!(?err, "Request failed");
-
-				metrics::PROXY_REQUEST_ERROR
-					.add(1, &[KeyValue::new("error_type", err.to_string())]);
-
-				// If we receive an error during a websocket request, we attempt to open the websocket anyway
-				// so we can send the error via websocket instead of http. Most websocket clients don't handle
-				// HTTP errors in a meaningful way resulting in unhelpful errors for the user
-				if is_websocket {
-					tracing::debug!("Upgrading client connection to WebSocket for error proxy");
-					match hyper_tungstenite::upgrade(mock_req, None) {
-						Ok((client_response, client_ws)) => {
-							tracing::debug!("Client WebSocket upgrade for error proxy successful");
-
-							self.state.tasks.spawn(
-								async move {
-									let ws_handle = match WebSocketHandle::new(client_ws).await {
-										Ok(ws_handle) => ws_handle,
-										Err(err) => {
-											tracing::debug!(
-												?err,
-												"failed initiating websocket handle for error proxy"
-											);
-											return;
-										}
-									};
-									let frame = err_to_close_frame(err, Some(request_ids.ray_id));
-
-									// Manual conversion to handle different tungstenite versions
-									let code_num: u16 = frame.code.into();
-									let reason = frame.reason.clone();
-
-									if let Err(err) = ws_handle
-										.send(tokio_tungstenite::tungstenite::Message::Close(Some(
-											tokio_tungstenite::tungstenite::protocol::CloseFrame {
-												code: code_num.into(),
-												reason,
-											},
-										)))
-										.await
-									{
-										tracing::debug!(
-											?err,
-											"failed sending websocket error proxy"
-										);
-									}
-								}
-								.instrument(
-									tracing::info_span!("ws_error_proxy_task", ?request_ids.ray_id),
-								),
-							);
-
-							// Return the response that will upgrade the client connection
-							// For proper WebSocket handshaking, we need to preserve the original response
-							// structure but convert it to our expected return type without modifying its content
-							tracing::debug!(
-								"Returning WebSocket upgrade response for error proxy to client"
-							);
-							// Extract the parts from the response but preserve all headers and status
-							let (mut parts, _) = client_response.into_parts();
-
-							// Add Sec-WebSocket-Protocol header to the response
-							// Many WebSocket clients (e.g. node-ws & Cloudflare) require a protocol in the response
-							parts.headers.insert(
-								"sec-websocket-protocol",
-								hyper::header::HeaderValue::from_static("rivet"),
-							);
-
-							// Create a new response with an empty body - WebSocket upgrades don't need a body
-							Response::from_parts(
-								parts,
-								ResponseBody::Full(Full::<Bytes>::new(Bytes::new())),
-							)
-						}
-						Err(err) => {
-							tracing::error!(
-								?err,
-								"Failed to upgrade client WebSocket for error proxy"
-							);
-
-							err_into_response(
-								errors::ConnectionError {
-									error_message: format!(
-										"Failed to upgrade client WebSocket for error proxy: {}",
-										err
-									),
-									remote_addr: self.remote_addr.to_string(),
-								}
-								.build(),
-							)?
-						}
-					}
-				} else {
-					err_into_response(err)?
-				}
-			}
-		};
-
-		if is_websocket && res.status() != StatusCode::SWITCHING_PROTOCOLS {
-			tracing::debug!("returned non-101 response to websocket");
-		}
-
-		// Add ray_id to response headers
-		if let Ok(ray_id_value) = request_ids.ray_id.to_string().parse() {
-			if let Some(existing_ray_id_value) = res
-				.headers()
-				.get(X_RIVET_RAY_ID)
-				.and_then(|h| h.to_str().ok())
-			{
-				if ray_id_value != existing_ray_id_value {
-					tracing::warn!(
-						expected_ray_id=%request_ids.ray_id,
-						received_ray_id=%existing_ray_id_value,
-						"downstream service set ray id header to a different value",
-					);
-				}
-			}
-
-			res.headers_mut().insert(X_RIVET_RAY_ID, ray_id_value);
-		}
-
-		let status = res.status().as_u16();
-
-		// Update request context with response details
-		request_context.guard_response_status = Some(status);
-		request_context.service_response_status = Some(status);
-
-		if let Some(content_type) = res
-			.headers()
-			.get(hyper::header::CONTENT_TYPE)
-			.and_then(|h| h.to_str().ok())
-		{
-			request_context.guard_response_content_type = Some(content_type.to_string());
-		}
-
-		if let Some(expires) = res
-			.headers()
-			.get(hyper::header::EXPIRES)
-			.and_then(|h| h.to_str().ok())
-		{
-			request_context.service_response_http_expires = Some(expires.to_string());
-		}
-
-		if let Some(last_modified) = res
-			.headers()
-			.get(hyper::header::LAST_MODIFIED)
-			.and_then(|h| h.to_str().ok())
-		{
-			request_context.service_response_http_last_modified = Some(last_modified.to_string());
-		}
-
-		// Set timing information
-		request_context.service_response_duration_ms =
-			Some(start_time.elapsed().as_millis() as u32);
-
-		// Insert analytics event asynchronously
-		let mut context_clone = request_context.clone();
-		tokio::spawn(
-			async move {
-				if let Err(error) = context_clone.insert_event().await {
-					tracing::warn!(?error, "failed to insert guard analytics event");
-				}
-			}
-			.instrument(tracing::info_span!("insert_event_task")),
-		);
-
-		let content_length = res
-			.headers()
-			.get(hyper::header::CONTENT_LENGTH)
-			.and_then(|h| h.to_str().ok())
-			.and_then(|s| s.parse::<usize>().ok())
-			.unwrap_or(0);
-
-		// Log information about the completed request
-		tracing::debug!(
-			?incoming_ray_id,
-			ray_id=?request_ids.ray_id,
-			req_id=?request_ids.req_id,
-			method = %method,
-			path = %path,
-			host = %host,
-			remote_addr = %self.remote_addr,
-			status = %status,
-			content_length = %content_length,
-			"Request completed"
-		);
-
-		Ok(res)
-	}
-}
-
 impl Clone for ProxyService {
 	fn clone(&self) -> Self {
 		Self {
 			state: self.state.clone(),
 			remote_addr: self.remote_addr,
-			client: self.client.clone(),
+			connection_start: self.connection_start,
 		}
 	}
 }
@@ -2499,18 +1777,8 @@ impl ProxyServiceFactory {
 		config: rivet_config::Config,
 		routing_fn: RoutingFn,
 		cache_key_fn: CacheKeyFn,
-		middleware_fn: MiddlewareFn,
-		port_type: PortType,
-		clickhouse_inserter: Option<clickhouse_inserter::ClickHouseInserterHandle>,
 	) -> Self {
-		let state = Arc::new(ProxyState::new(
-			config,
-			routing_fn,
-			cache_key_fn,
-			middleware_fn,
-			port_type,
-			clickhouse_inserter,
-		));
+		let state = Arc::new(ProxyState::new(config, routing_fn, cache_key_fn));
 		Self { state }
 	}
 
@@ -2522,178 +1790,8 @@ impl ProxyServiceFactory {
 	pub async fn wait_idle(&self) {
 		self.state.tasks.wait_idle().await
 	}
-}
-
-fn add_proxy_headers_with_addr(
-	headers: &mut hyper::HeaderMap,
-	original_headers: &hyper::HeaderMap,
-	remote_addr: SocketAddr,
-) -> Result<()> {
-	// Copy headers except Host
-	for (key, value) in original_headers.iter() {
-		if key != hyper::header::HOST {
-			headers.insert(key.clone(), value.clone());
-		}
-	}
-
-	// Add X-Forwarded-For header
-	if let Some(existing) = original_headers.get(X_FORWARDED_FOR) {
-		if let Ok(forwarded) = existing.to_str() {
-			if !forwarded.contains(&remote_addr.ip().to_string()) {
-				headers.insert(
-					X_FORWARDED_FOR,
-					hyper::header::HeaderValue::from_str(&format!(
-						"{}, {}",
-						forwarded,
-						remote_addr.ip()
-					))?,
-				);
-			}
-		}
-	} else {
-		headers.insert(
-			X_FORWARDED_FOR,
-			hyper::header::HeaderValue::from_str(&remote_addr.ip().to_string())?,
-		);
-	}
-
-	Ok(())
-}
-
-fn err_into_response(err: anyhow::Error) -> Result<Response<ResponseBody>> {
-	let (status, error_response) =
-		if let Some(rivet_err) = err.chain().find_map(|x| x.downcast_ref::<RivetError>()) {
-			let status = match (rivet_err.group(), rivet_err.code()) {
-				("api", "not_found") => StatusCode::NOT_FOUND,
-				("api", "unauthorized") => StatusCode::UNAUTHORIZED,
-				("api", "forbidden") => StatusCode::FORBIDDEN,
-				("guard", "rate_limit") => StatusCode::TOO_MANY_REQUESTS,
-				("guard", "upstream_error") => StatusCode::BAD_GATEWAY,
-				("guard", "routing_error") => StatusCode::BAD_GATEWAY,
-				("guard", "request_timeout") => StatusCode::GATEWAY_TIMEOUT,
-				("guard", "retry_attempts_exceeded") => StatusCode::BAD_GATEWAY,
-				("guard", "actor_not_found") => StatusCode::NOT_FOUND,
-				("guard", "actor_destroyed") => StatusCode::NOT_FOUND,
-				("guard", "service_unavailable") => StatusCode::SERVICE_UNAVAILABLE,
-				("guard", "actor_ready_timeout") => StatusCode::SERVICE_UNAVAILABLE,
-				("guard", "no_route") => StatusCode::NOT_FOUND,
-				_ => StatusCode::BAD_REQUEST,
-			};
-
-			(status, ErrorResponse::from(rivet_err))
-		} else if let Some(raw_err) = err
-			.chain()
-			.find_map(|x| x.downcast_ref::<RawErrorResponse>())
-		{
-			(raw_err.0, raw_err.1.clone())
-		} else {
-			(
-				StatusCode::INTERNAL_SERVER_ERROR,
-				ErrorResponse::from(&RivetError {
-					schema: &rivet_error::INTERNAL_ERROR,
-					meta: None,
-					message: None,
-				}),
-			)
-		};
-
-	let body_json = serde_json::to_vec(&error_response)?;
-	let bytes = Bytes::from(body_json);
-
-	Response::builder()
-		.status(status)
-		.header(hyper::header::CONTENT_TYPE, "application/json")
-		.body(ResponseBody::Full(Full::new(bytes)))
-		.map_err(Into::into)
-}
-
-fn should_retry_request(res: &Result<Response<ResponseBody>>) -> bool {
-	match res {
-		Ok(resp) => should_retry_request_inner(resp.status(), resp.headers()),
-		Err(err) => {
-			if let Some(rivet_err) = err.chain().find_map(|x| x.downcast_ref::<RivetError>()) {
-				rivet_err.group() == "guard" && rivet_err.code() == "service_unavailable"
-			} else {
-				false
-			}
-		}
-	}
-}
 
-// Determine if a response should trigger a retry: 503 + x-rivet-error
-fn should_retry_request_inner(status: StatusCode, headers: &hyper::HeaderMap) -> bool {
-	status == StatusCode::SERVICE_UNAVAILABLE && headers.contains_key(X_RIVET_ERROR)
-}
-
-// Determine if a websocket error is retryable (e.g., transient UPS/tunnel issues)
-fn is_retryable_ws_error(err: &anyhow::Error) -> bool {
-	if let Some(rivet_err) = err.chain().find_map(|x| x.downcast_ref::<RivetError>()) {
-		rivet_err.group() == "guard" && rivet_err.code() == "websocket_service_unavailable"
-	} else {
-		false
-	}
-}
-
-pub fn is_ws_hibernate(err: &anyhow::Error) -> bool {
-	if let Some(rivet_err) = err.chain().find_map(|x| x.downcast_ref::<RivetError>()) {
-		rivet_err.group() == "guard" && rivet_err.code() == "websocket_service_hibernate"
-	} else {
-		false
-	}
-}
-
-fn str_to_close_frame(err: &str) -> CloseFrame {
-	// NOTE: reason cannot be more than 123 bytes as per the WS protocol spec
-	let reason = rivet_util::safe_slice(err, 0, 123).into();
-
-	CloseFrame {
-		code: CloseCode::Error,
-		reason,
-	}
-}
-
-fn err_to_close_frame(err: anyhow::Error, ray_id: Option<Id>) -> CloseFrame {
-	let rivet_err = err
-		.chain()
-		.find_map(|x| x.downcast_ref::<RivetError>())
-		.cloned()
-		.unwrap_or_else(|| RivetError::from(&INTERNAL_ERROR));
-
-	let code = match (rivet_err.group(), rivet_err.code()) {
-		("ws", "connection_closed") | ("ws", "eviction") => CloseCode::Normal,
-		_ => CloseCode::Error,
-	};
-
-	let reason = if let Some(ray_id) = ray_id {
-		format!("{}.{}#{}", rivet_err.group(), rivet_err.code(), ray_id)
-	} else {
-		format!("{}.{}", rivet_err.group(), rivet_err.code())
-	};
-
-	// NOTE: reason cannot be more than 123 bytes as per the WS protocol
-	let reason = rivet_util::safe_slice(&reason, 0, 123).into();
-
-	CloseFrame { code, reason }
-}
-
-fn to_hyper_close(frame: Option<CloseFrame>) -> hyper_tungstenite::tungstenite::Message {
-	if let Some(frame) = frame {
-		// Manual conversion to handle different tungstenite versions
-		let code_num: u16 = frame.code.into();
-		let reason = frame.reason.clone();
-
-		tokio_tungstenite::tungstenite::Message::Close(Some(
-			tokio_tungstenite::tungstenite::protocol::CloseFrame {
-				code: code_num.into(),
-				reason,
-			},
-		))
-	} else {
-		tokio_tungstenite::tungstenite::Message::Close(Some(
-			tokio_tungstenite::tungstenite::protocol::CloseFrame {
-				code: CloseCode::Normal,
-				reason: "ws.closed".into(),
-			},
-		))
+	pub fn remaining_tasks(&self) -> usize {
+		self.state.tasks.remaining_tasks()
 	}
 }