Skip to main content

zentinel_proxy/proxy/
http_trait.rs

1//! ProxyHttp trait implementation for ZentinelProxy.
2//!
3//! This module contains the Pingora ProxyHttp trait implementation which defines
4//! the core request/response lifecycle handling.
5
6use async_trait::async_trait;
7use bytes::Bytes;
8use pingora::http::ResponseHeader;
9use pingora::prelude::*;
10use pingora::protocols::Digest;
11use pingora::proxy::{ProxyHttp, Session};
12use pingora::upstreams::peer::Peer;
13use pingora_cache::{
14    CacheKey, CacheMeta, ForcedFreshness, HitHandler, NoCacheReason, RespCacheable,
15};
16use pingora_timeout::sleep;
17use std::os::unix::io::RawFd;
18use std::time::Duration;
19use tracing::{debug, error, info, trace, warn};
20
21use crate::cache::{get_cache_eviction, get_cache_lock, get_cache_storage};
22use crate::disk_cache::DiskHitHandler;
23use crate::hybrid_cache::HybridHitHandler;
24use crate::inference::{
25    extract_inference_content, is_sse_response, PromptInjectionResult, StreamingTokenCounter,
26};
27use crate::logging::{AccessLogEntry, AuditEventType, AuditLogEntry};
28use crate::rate_limit::HeaderAccessor;
29use crate::routing::RequestInfo;
30
31use super::context::{FallbackReason, RequestContext};
32use super::fallback::FallbackEvaluator;
33use super::fallback_metrics::get_fallback_metrics;
34use super::model_routing;
35use super::model_routing_metrics::get_model_routing_metrics;
36use super::ZentinelProxy;
37
38/// Helper type for rate limiting when we don't need header access
39struct NoHeaderAccessor;
40impl HeaderAccessor for NoHeaderAccessor {
41    fn get_header(&self, _name: &str) -> Option<String> {
42        None
43    }
44}
45
46#[async_trait]
47impl ProxyHttp for ZentinelProxy {
48    type CTX = RequestContext;
49
50    fn new_ctx(&self) -> Self::CTX {
51        RequestContext::new()
52    }
53
54    fn fail_to_connect(
55        &self,
56        _session: &mut Session,
57        peer: &HttpPeer,
58        ctx: &mut Self::CTX,
59        e: Box<Error>,
60    ) -> Box<Error> {
61        error!(
62            correlation_id = %ctx.trace_id,
63            route_id = ctx.route_id.as_deref().unwrap_or("unknown"),
64            upstream = ctx.upstream.as_deref().unwrap_or("unknown"),
65            peer_address = %peer.address(),
66            error = %e,
67            "Failed to connect to upstream peer"
68        );
69        self.log_manager.log_request_error(
70            "error",
71            "Failed to connect to upstream peer",
72            &ctx.trace_id,
73            ctx.route_id.as_deref(),
74            ctx.upstream.as_deref(),
75            Some(format!("peer={} error={}", peer.address(), e)),
76        );
77        // Custom error pages are handled in response_filter
78        e
79    }
80
81    /// Early request filter - runs before upstream selection
82    /// Used to handle builtin routes that don't need an upstream connection
83    async fn early_request_filter(
84        &self,
85        session: &mut Session,
86        ctx: &mut Self::CTX,
87    ) -> Result<(), Box<Error>> {
88        // Extract request info for routing
89        let req_header = session.req_header();
90        let method = req_header.method.as_str();
91        let path = req_header.uri.path();
92        let host = crate::http_helpers::extract_request_host(req_header);
93
94        // Handle ACME HTTP-01 challenges before any other processing
95        if let Some(ref challenge_manager) = self.acme_challenges {
96            if let Some(token) = crate::acme::ChallengeManager::extract_token(path) {
97                if let Some(key_authorization) = challenge_manager.get_response(token) {
98                    debug!(
99                        token = %token,
100                        "Serving ACME HTTP-01 challenge response"
101                    );
102
103                    // Build response
104                    let mut resp = ResponseHeader::build(200, None)?;
105                    resp.insert_header("Content-Type", "text/plain")?;
106                    resp.insert_header("Content-Length", key_authorization.len().to_string())?;
107
108                    // Send response
109                    session.write_response_header(Box::new(resp), false).await?;
110                    session
111                        .write_response_body(Some(Bytes::from(key_authorization)), true)
112                        .await?;
113
114                    // Return error to signal request is complete
115                    return Err(Error::explain(
116                        ErrorType::InternalError,
117                        "ACME challenge served",
118                    ));
119                } else {
120                    // Token not found - could be a stale request or attack
121                    warn!(
122                        token = %token,
123                        "ACME challenge token not found"
124                    );
125                }
126            }
127        }
128
129        ctx.method = method.to_string();
130        ctx.path = path.to_string();
131        ctx.host = Some(host.to_string());
132
133        // Match route to determine service type
134        let route_match = {
135            let route_matcher = self.route_matcher.read();
136            let mut request_info = RequestInfo::new(method, path, host);
137
138            // Include headers for header-based route matching (Gateway API)
139            if route_matcher.needs_headers() {
140                request_info = request_info
141                    .with_headers(RequestInfo::build_headers(req_header.headers.iter()));
142            }
143
144            match route_matcher.match_request(&request_info) {
145                Some(m) => m,
146                None => return Ok(()), // No matching route, let upstream_peer handle it
147            }
148        };
149
150        ctx.trace_id = self.get_trace_id(session);
151        ctx.route_id = Some(route_match.route_id.to_string());
152        ctx.route_config = Some(route_match.config.clone());
153
154        // Parse incoming W3C trace context if present
155        if let Some(traceparent) = req_header.headers.get(crate::otel::TRACEPARENT_HEADER) {
156            if let Ok(s) = traceparent.to_str() {
157                ctx.trace_context = crate::otel::TraceContext::parse_traceparent(s);
158            }
159        }
160
161        // Start OpenTelemetry request span if tracing is enabled
162        if let Some(tracer) = crate::otel::get_tracer() {
163            ctx.otel_span = Some(tracer.start_span(method, path, ctx.trace_context.as_ref()));
164        }
165
166        // Check if this is a builtin handler route
167        if route_match.config.service_type == zentinel_config::ServiceType::Builtin {
168            trace!(
169                correlation_id = %ctx.trace_id,
170                route_id = %route_match.route_id,
171                builtin_handler = ?route_match.config.builtin_handler,
172                "Handling builtin route in early_request_filter"
173            );
174
175            // Handle the builtin route directly
176            let handled = self
177                .handle_builtin_route(session, ctx, &route_match)
178                .await?;
179
180            if handled {
181                // Return error to signal that request is complete (Pingora will not continue)
182                return Err(Error::explain(
183                    ErrorType::InternalError,
184                    "Builtin handler complete",
185                ));
186            }
187        }
188
189        Ok(())
190    }
191
192    async fn upstream_peer(
193        &self,
194        session: &mut Session,
195        ctx: &mut Self::CTX,
196    ) -> Result<Box<HttpPeer>, Box<Error>> {
197        // Track active request
198        self.reload_coordinator.inc_requests();
199
200        // Cache global config once per request (avoids repeated Arc clones)
201        if ctx.config.is_none() {
202            ctx.config = Some(self.config_manager.current());
203        }
204
205        // Cache client address for logging if not already set
206        if ctx.client_ip.is_empty() {
207            ctx.client_ip = session
208                .client_addr()
209                .map(|a| a.to_string())
210                .unwrap_or_else(|| "unknown".to_string());
211        }
212
213        let req_header = session.req_header();
214
215        // Cache request info for access logging if not already set
216        if ctx.method.is_empty() {
217            ctx.method = req_header.method.to_string();
218            ctx.path = req_header.uri.path().to_string();
219            ctx.query = req_header.uri.query().map(|q| q.to_string());
220            ctx.host = Some(crate::http_helpers::extract_request_host(req_header).to_string());
221        }
222        ctx.user_agent = req_header
223            .headers
224            .get("user-agent")
225            .and_then(|v| v.to_str().ok())
226            .map(|s| s.to_string());
227        ctx.referer = req_header
228            .headers
229            .get("referer")
230            .and_then(|v| v.to_str().ok())
231            .map(|s| s.to_string());
232
233        trace!(
234            correlation_id = %ctx.trace_id,
235            client_ip = %ctx.client_ip,
236            "Request received, initializing context"
237        );
238
239        // Use cached route info if already set by early_request_filter
240        let route_match = if let Some(ref route_config) = ctx.route_config {
241            let route_id = ctx.route_id.as_deref().unwrap_or("");
242            crate::routing::RouteMatch {
243                route_id: zentinel_common::RouteId::new(route_id),
244                config: route_config.clone(),
245            }
246        } else {
247            // Match route using sync RwLock (scoped to ensure lock is released before async ops)
248            let (match_result, route_duration) = {
249                let route_matcher = self.route_matcher.read();
250                let host = ctx.host.as_deref().unwrap_or("");
251
252                // Build request info (zero-copy for common case)
253                let mut request_info = RequestInfo::new(&ctx.method, &ctx.path, host);
254
255                // Only build headers HashMap if any route needs header matching
256                if route_matcher.needs_headers() {
257                    request_info = request_info
258                        .with_headers(RequestInfo::build_headers(req_header.headers.iter()));
259                }
260
261                // Only parse query params if any route needs query param matching
262                if route_matcher.needs_query_params() {
263                    request_info =
264                        request_info.with_query_params(RequestInfo::parse_query_params(&ctx.path));
265                }
266
267                trace!(
268                    correlation_id = %ctx.trace_id,
269                    method = %request_info.method,
270                    path = %request_info.path,
271                    host = %request_info.host,
272                    "Built request info for route matching"
273                );
274
275                let route_start = std::time::Instant::now();
276                let route_match = route_matcher.match_request(&request_info).ok_or_else(|| {
277                    warn!(
278                        correlation_id = %ctx.trace_id,
279                        method = %request_info.method,
280                        path = %request_info.path,
281                        host = %request_info.host,
282                        "No matching route found for request"
283                    );
284                    self.log_manager.log_request_error(
285                        "warn",
286                        "No matching route found for request",
287                        &ctx.trace_id,
288                        None,
289                        None,
290                        Some(format!(
291                            "method={} path={} host={}",
292                            request_info.method, request_info.path, request_info.host
293                        )),
294                    );
295                    Error::explain(ErrorType::HTTPStatus(404), "No matching route found")
296                })?;
297                let route_duration = route_start.elapsed();
298                // Lock is dropped here when block ends
299                (route_match, route_duration)
300            };
301
302            ctx.route_id = Some(match_result.route_id.to_string());
303            ctx.route_config = Some(match_result.config.clone());
304
305            // Set trace_id if not already set by early_request_filter
306            if ctx.trace_id.is_empty() {
307                ctx.trace_id = self.get_trace_id(session);
308
309                // Parse incoming W3C trace context if present
310                if let Some(traceparent) = req_header.headers.get(crate::otel::TRACEPARENT_HEADER) {
311                    if let Ok(s) = traceparent.to_str() {
312                        ctx.trace_context = crate::otel::TraceContext::parse_traceparent(s);
313                    }
314                }
315
316                // Start OpenTelemetry request span if tracing is enabled
317                if let Some(tracer) = crate::otel::get_tracer() {
318                    ctx.otel_span =
319                        Some(tracer.start_span(&ctx.method, &ctx.path, ctx.trace_context.as_ref()));
320                }
321            }
322
323            trace!(
324                correlation_id = %ctx.trace_id,
325                route_id = %match_result.route_id,
326                route_duration_us = route_duration.as_micros(),
327                service_type = ?match_result.config.service_type,
328                "Route matched"
329            );
330            match_result
331        };
332
333        // Check if this is a builtin handler route (no upstream needed)
334        if route_match.config.service_type == zentinel_config::ServiceType::Builtin {
335            trace!(
336                correlation_id = %ctx.trace_id,
337                route_id = %route_match.route_id,
338                builtin_handler = ?route_match.config.builtin_handler,
339                "Route type is builtin, skipping upstream"
340            );
341            // Mark as builtin route for later processing in request_filter
342            ctx.upstream = Some(format!("_builtin_{}", route_match.route_id));
343            // Return error to skip upstream connection for builtin routes
344            return Err(Error::explain(
345                ErrorType::InternalError,
346                "Builtin handler handled in request_filter",
347            ));
348        }
349
350        // Check if this is a static file route
351        if route_match.config.service_type == zentinel_config::ServiceType::Static {
352            trace!(
353                correlation_id = %ctx.trace_id,
354                route_id = %route_match.route_id,
355                "Route type is static, checking for static server"
356            );
357            // Static routes don't need an upstream
358            if self
359                .static_servers
360                .get(route_match.route_id.as_str())
361                .await
362                .is_some()
363            {
364                // Mark this as a static route for later processing
365                ctx.upstream = Some(format!("_static_{}", route_match.route_id));
366                info!(
367                    correlation_id = %ctx.trace_id,
368                    route_id = %route_match.route_id,
369                    path = %ctx.path,
370                    "Serving static file"
371                );
372                // Return error to avoid upstream connection for static routes
373                return Err(Error::explain(
374                    ErrorType::InternalError,
375                    "Static file serving handled in request_filter",
376                ));
377            }
378        }
379
380        // === Model-based routing (for inference routes) ===
381        // Check if model routing is configured and select upstream based on model
382        let mut model_routing_applied = false;
383        if let Some(ref inference) = route_match.config.inference {
384            if let Some(ref model_routing) = inference.model_routing {
385                // Try to extract model from headers (fast path - no body parsing needed)
386                let model = model_routing::extract_model_from_headers(&req_header.headers);
387
388                if let Some(ref model_name) = model {
389                    // Find upstream for this model
390                    if let Some(routing_result) =
391                        model_routing::find_upstream_for_model(model_routing, model_name)
392                    {
393                        debug!(
394                            correlation_id = %ctx.trace_id,
395                            route_id = %route_match.route_id,
396                            model = %model_name,
397                            upstream = %routing_result.upstream,
398                            is_default = routing_result.is_default,
399                            provider_override = ?routing_result.provider,
400                            "Model-based routing selected upstream"
401                        );
402
403                        ctx.record_model_routing(
404                            &routing_result.upstream,
405                            Some(model_name.clone()),
406                            routing_result.provider,
407                        );
408                        model_routing_applied = true;
409
410                        // Record metrics
411                        if let Some(metrics) = get_model_routing_metrics() {
412                            metrics.record_model_routed(
413                                route_match.route_id.as_str(),
414                                model_name,
415                                &routing_result.upstream,
416                            );
417                            if routing_result.is_default {
418                                metrics.record_default_upstream(route_match.route_id.as_str());
419                            }
420                            if let Some(provider) = routing_result.provider {
421                                metrics.record_provider_override(
422                                    route_match.route_id.as_str(),
423                                    &routing_result.upstream,
424                                    provider.as_str(),
425                                );
426                            }
427                        }
428                    }
429                } else if let Some(ref default_upstream) = model_routing.default_upstream {
430                    // No model in headers, use default upstream
431                    debug!(
432                        correlation_id = %ctx.trace_id,
433                        route_id = %route_match.route_id,
434                        upstream = %default_upstream,
435                        "Model-based routing using default upstream (no model header)"
436                    );
437                    ctx.record_model_routing(default_upstream, None, None);
438                    model_routing_applied = true;
439
440                    // Record metrics for no model header case
441                    if let Some(metrics) = get_model_routing_metrics() {
442                        metrics.record_no_model_header(route_match.route_id.as_str());
443                    }
444                }
445            }
446        }
447
448        // Regular route with upstream (if model routing didn't set it)
449        if !model_routing_applied {
450            if let Some(ref upstream) = route_match.config.upstream {
451                ctx.upstream = Some(upstream.clone());
452                trace!(
453                    correlation_id = %ctx.trace_id,
454                    route_id = %route_match.route_id,
455                    upstream = %upstream,
456                    "Upstream configured for route"
457                );
458            } else {
459                // Route matched but has no valid upstream (e.g. invalid backend
460                // kind, denied cross-namespace ref). Return HTTP 500 per Gateway
461                // API spec — the route exists but cannot be fulfilled.
462                warn!(
463                    correlation_id = %ctx.trace_id,
464                    route_id = %route_match.route_id,
465                    "Route has no upstream configured, returning 500"
466                );
467                crate::http_helpers::write_error(
468                    session,
469                    500,
470                    "Internal Server Error",
471                    "text/plain",
472                )
473                .await?;
474                return Err(Error::explain(
475                    ErrorType::HTTPStatus(500),
476                    "Route has no valid upstream",
477                ));
478            }
479        }
480
481        // === Fallback routing evaluation (pre-request) ===
482        // Check if fallback should be triggered due to health or budget conditions
483        if let Some(ref fallback_config) = route_match.config.fallback {
484            let upstream_name = ctx.upstream.as_ref().unwrap();
485
486            // Check if primary upstream is healthy
487            let is_healthy = if let Some(pool) = self.upstream_pools.get(upstream_name).await {
488                pool.has_healthy_targets().await
489            } else {
490                false // Pool not found, treat as unhealthy
491            };
492
493            // Check if budget is exhausted (for inference routes)
494            let is_budget_exhausted = ctx.inference_budget_exhausted;
495
496            // Get model name for model mapping (inference routes)
497            let current_model = ctx.inference_model.as_deref();
498
499            // Create fallback evaluator
500            let evaluator = FallbackEvaluator::new(
501                fallback_config,
502                ctx.tried_upstreams(),
503                ctx.fallback_attempt,
504            );
505
506            // Evaluate pre-request fallback conditions
507            if let Some(decision) = evaluator.should_fallback_before_request(
508                upstream_name,
509                is_healthy,
510                is_budget_exhausted,
511                current_model,
512            ) {
513                info!(
514                    correlation_id = %ctx.trace_id,
515                    route_id = %route_match.route_id,
516                    from_upstream = %upstream_name,
517                    to_upstream = %decision.next_upstream,
518                    reason = %decision.reason,
519                    fallback_attempt = ctx.fallback_attempt + 1,
520                    "Triggering fallback routing"
521                );
522
523                // Record fallback metrics
524                if let Some(metrics) = get_fallback_metrics() {
525                    metrics.record_fallback_attempt(
526                        route_match.route_id.as_str(),
527                        upstream_name,
528                        &decision.next_upstream,
529                        &decision.reason,
530                    );
531                }
532
533                // Record fallback in context
534                ctx.record_fallback(decision.reason, &decision.next_upstream);
535
536                // Apply model mapping if present
537                if let Some((original, mapped)) = decision.model_mapping {
538                    // Record model mapping metrics
539                    if let Some(metrics) = get_fallback_metrics() {
540                        metrics.record_model_mapping(
541                            route_match.route_id.as_str(),
542                            &original,
543                            &mapped,
544                        );
545                    }
546
547                    ctx.record_model_mapping(original, mapped);
548                    trace!(
549                        correlation_id = %ctx.trace_id,
550                        original_model = ?ctx.model_mapping_applied().map(|m| &m.0),
551                        mapped_model = ?ctx.model_mapping_applied().map(|m| &m.1),
552                        "Applied model mapping for fallback"
553                    );
554                }
555            }
556        }
557
558        debug!(
559            correlation_id = %ctx.trace_id,
560            route_id = %route_match.route_id,
561            upstream = ?ctx.upstream,
562            method = %req_header.method,
563            path = %req_header.uri.path(),
564            host = ctx.host.as_deref().unwrap_or("-"),
565            client_ip = %ctx.client_ip,
566            "Processing request"
567        );
568
569        // Get upstream pool (skip for static routes)
570        if ctx
571            .upstream
572            .as_ref()
573            .is_some_and(|u| u.starts_with("_static_"))
574        {
575            // Static routes are handled in request_filter, should not reach here
576            return Err(Error::explain(
577                ErrorType::InternalError,
578                "Static route should be handled in request_filter",
579            ));
580        }
581
582        let upstream_name = ctx
583            .upstream
584            .as_ref()
585            .ok_or_else(|| Error::explain(ErrorType::InternalError, "No upstream configured"))?;
586
587        trace!(
588            correlation_id = %ctx.trace_id,
589            upstream = %upstream_name,
590            "Looking up upstream pool"
591        );
592
593        let pool = self
594            .upstream_pools
595            .get(upstream_name)
596            .await
597            .ok_or_else(|| {
598                error!(
599                    correlation_id = %ctx.trace_id,
600                    upstream = %upstream_name,
601                    "Upstream pool not found"
602                );
603                self.log_manager.log_request_error(
604                    "error",
605                    "Upstream pool not found",
606                    &ctx.trace_id,
607                    ctx.route_id.as_deref(),
608                    Some(upstream_name),
609                    None,
610                );
611                Error::explain(
612                    ErrorType::InternalError,
613                    format!("Upstream pool '{}' not found", upstream_name),
614                )
615            })?;
616
617        // Select peer from pool with retries
618        let max_retries = route_match
619            .config
620            .retry_policy
621            .as_ref()
622            .map(|r| r.max_attempts)
623            .unwrap_or(1);
624
625        trace!(
626            correlation_id = %ctx.trace_id,
627            upstream = %upstream_name,
628            max_retries = max_retries,
629            "Starting upstream peer selection"
630        );
631
632        let mut last_error = None;
633        let selection_start = std::time::Instant::now();
634
635        for attempt in 1..=max_retries {
636            ctx.upstream_attempts = attempt;
637
638            trace!(
639                correlation_id = %ctx.trace_id,
640                upstream = %upstream_name,
641                attempt = attempt,
642                max_retries = max_retries,
643                "Attempting to select upstream peer"
644            );
645
646            match pool.select_peer_with_metadata(None).await {
647                Ok((mut peer, metadata)) => {
648                    let selection_duration = selection_start.elapsed();
649                    // Track active request for drain lifecycle
650                    pool.increment_active();
651                    // Store selected peer address for feedback reporting in logging()
652                    let peer_addr = peer.address().to_string();
653                    ctx.selected_upstream_address = Some(peer_addr.clone());
654
655                    // Copy sticky session metadata to context for response_filter
656                    if metadata.contains_key("sticky_session_new") {
657                        ctx.sticky_session_new_assignment = true;
658                        ctx.sticky_session_set_cookie =
659                            metadata.get("sticky_set_cookie_header").cloned();
660                        ctx.sticky_target_index = metadata
661                            .get("sticky_target_index")
662                            .and_then(|s| s.parse().ok());
663
664                        trace!(
665                            correlation_id = %ctx.trace_id,
666                            sticky_target_index = ?ctx.sticky_target_index,
667                            "New sticky session assignment, will set cookie"
668                        );
669                    }
670
671                    debug!(
672                        correlation_id = %ctx.trace_id,
673                        upstream = %upstream_name,
674                        peer_address = %peer_addr,
675                        attempt = attempt,
676                        selection_duration_us = selection_duration.as_micros(),
677                        sticky_session_hit = metadata.contains_key("sticky_session_hit"),
678                        sticky_session_new = ctx.sticky_session_new_assignment,
679                        "Selected upstream peer"
680                    );
681                    // Apply per-route policy timeout (lowest priority)
682                    if let Some(ref rc) = ctx.route_config {
683                        if let Some(timeout_secs) = rc.policies.timeout_secs {
684                            peer.options.read_timeout = Some(Duration::from_secs(timeout_secs));
685                        }
686                    }
687
688                    // Apply filter timeout overrides (higher priority, overwrites policy)
689                    if let Some(connect_secs) = ctx.filter_connect_timeout_secs {
690                        peer.options.connection_timeout = Some(Duration::from_secs(connect_secs));
691                    }
692                    if let Some(upstream_secs) = ctx.filter_upstream_timeout_secs {
693                        peer.options.read_timeout = Some(Duration::from_secs(upstream_secs));
694                    }
695
696                    return Ok(Box::new(peer));
697                }
698                Err(e) => {
699                    warn!(
700                        correlation_id = %ctx.trace_id,
701                        upstream = %upstream_name,
702                        attempt = attempt,
703                        max_retries = max_retries,
704                        error = %e,
705                        "Failed to select upstream peer"
706                    );
707                    last_error = Some(e);
708
709                    if attempt < max_retries {
710                        // Exponential backoff (using pingora-timeout for efficiency)
711                        let backoff = Duration::from_millis(100 * 2_u64.pow(attempt - 1));
712                        trace!(
713                            correlation_id = %ctx.trace_id,
714                            backoff_ms = backoff.as_millis(),
715                            "Backing off before retry"
716                        );
717                        sleep(backoff).await;
718                    }
719                }
720            }
721        }
722
723        let selection_duration = selection_start.elapsed();
724        error!(
725            correlation_id = %ctx.trace_id,
726            upstream = %upstream_name,
727            attempts = max_retries,
728            selection_duration_ms = selection_duration.as_millis(),
729            last_error = ?last_error,
730            "All upstream selection attempts failed"
731        );
732        self.log_manager.log_request_error(
733            "error",
734            "All upstream selection attempts failed",
735            &ctx.trace_id,
736            ctx.route_id.as_deref(),
737            Some(upstream_name),
738            Some(format!("attempts={} error={:?}", max_retries, last_error)),
739        );
740
741        // Record exhausted metric if fallback was used but all upstreams failed
742        if ctx.used_fallback() {
743            if let Some(metrics) = get_fallback_metrics() {
744                metrics.record_fallback_exhausted(ctx.route_id.as_deref().unwrap_or("unknown"));
745            }
746        }
747
748        Err(Error::explain(
749            ErrorType::InternalError,
750            format!("All upstream attempts failed: {:?}", last_error),
751        ))
752    }
753
754    async fn request_filter(
755        &self,
756        session: &mut Session,
757        ctx: &mut Self::CTX,
758    ) -> Result<bool, Box<Error>> {
759        trace!(
760            correlation_id = %ctx.trace_id,
761            route_id = ctx.route_id.as_deref().unwrap_or("unknown"),
762            "Starting request filter phase"
763        );
764
765        // Apply per-listener timeouts from config
766        if let Some(server_addr) = session.downstream_session.server_addr() {
767            let server_addr_str = server_addr.to_string();
768            let config = ctx
769                .config
770                .get_or_insert_with(|| self.config_manager.current());
771            for listener in &config.listeners {
772                if listener.address == server_addr_str {
773                    // Apply downstream read timeout
774                    session.downstream_session.set_read_timeout(Some(
775                        std::time::Duration::from_secs(listener.request_timeout_secs),
776                    ));
777                    // Store keepalive for response phase
778                    ctx.listener_keepalive_timeout_secs = Some(listener.keepalive_timeout_secs);
779                    break;
780                }
781            }
782        }
783
784        // Check rate limiting early (before other processing)
785        // Fast path: skip if no rate limiting is configured for this route
786        if let Some(route_id) = ctx.route_id.as_deref() {
787            if self.rate_limit_manager.has_route_limiter(route_id) {
788                let rate_result = self.rate_limit_manager.check(
789                    route_id,
790                    &ctx.client_ip,
791                    &ctx.path,
792                    Option::<&NoHeaderAccessor>::None,
793                );
794
795                // Store rate limit info for response headers (even if allowed)
796                if rate_result.limit > 0 {
797                    ctx.rate_limit_info = Some(super::context::RateLimitHeaderInfo {
798                        limit: rate_result.limit,
799                        remaining: rate_result.remaining,
800                        reset_at: rate_result.reset_at,
801                    });
802                }
803
804                if !rate_result.allowed {
805                    use zentinel_config::RateLimitAction;
806
807                    match rate_result.action {
808                        RateLimitAction::Reject => {
809                            warn!(
810                                correlation_id = %ctx.trace_id,
811                                route_id = route_id,
812                                client_ip = %ctx.client_ip,
813                                limiter = %rate_result.limiter,
814                                limit = rate_result.limit,
815                                remaining = rate_result.remaining,
816                                "Request rate limited"
817                            );
818                            self.metrics.record_blocked_request("rate_limited");
819
820                            // Audit log the rate limit
821                            let audit_entry = AuditLogEntry::rate_limited(
822                                &ctx.trace_id,
823                                &ctx.method,
824                                &ctx.path,
825                                &ctx.client_ip,
826                                &rate_result.limiter,
827                            )
828                            .with_route_id(route_id)
829                            .with_status_code(rate_result.status_code);
830                            self.log_manager.log_audit(&audit_entry);
831
832                            // Send rate limit response with headers
833                            let body = rate_result
834                                .message
835                                .unwrap_or_else(|| "Rate limit exceeded".to_string());
836
837                            // Build response with rate limit headers
838                            let retry_after = rate_result.reset_at.saturating_sub(
839                                std::time::SystemTime::now()
840                                    .duration_since(std::time::UNIX_EPOCH)
841                                    .unwrap_or_default()
842                                    .as_secs(),
843                            );
844                            crate::http_helpers::write_rate_limit_error(
845                                session,
846                                rate_result.status_code,
847                                &body,
848                                rate_result.limit,
849                                rate_result.remaining,
850                                rate_result.reset_at,
851                                retry_after,
852                            )
853                            .await?;
854                            return Ok(true); // Request complete, don't continue
855                        }
856                        RateLimitAction::LogOnly => {
857                            debug!(
858                                correlation_id = %ctx.trace_id,
859                                route_id = route_id,
860                                "Rate limit exceeded (log only mode)"
861                            );
862                            // Continue processing
863                        }
864                        RateLimitAction::Delay => {
865                            // Apply delay if suggested by rate limiter
866                            if let Some(delay_ms) = rate_result.suggested_delay_ms {
867                                // Cap delay at the configured maximum
868                                let actual_delay = delay_ms.min(rate_result.max_delay_ms);
869
870                                if actual_delay > 0 {
871                                    debug!(
872                                        correlation_id = %ctx.trace_id,
873                                        route_id = route_id,
874                                        suggested_delay_ms = delay_ms,
875                                        max_delay_ms = rate_result.max_delay_ms,
876                                        actual_delay_ms = actual_delay,
877                                        "Applying rate limit delay"
878                                    );
879
880                                    tokio::time::sleep(std::time::Duration::from_millis(
881                                        actual_delay,
882                                    ))
883                                    .await;
884                                }
885                            }
886                            // Continue processing after delay
887                        }
888                    }
889                }
890            }
891        }
892
893        // Inference rate limiting (token-based, for LLM/AI routes)
894        // This runs after regular rate limiting and checks service type
895        if let Some(route_id) = ctx.route_id.as_deref() {
896            if let Some(ref route_config) = ctx.route_config {
897                if route_config.service_type == zentinel_config::ServiceType::Inference
898                    && self.inference_rate_limit_manager.has_route(route_id)
899                {
900                    // For inference rate limiting, we need access to the request body
901                    // to estimate tokens. We'll use buffered body if available.
902                    let headers = &session.req_header().headers;
903
904                    // Try to get buffered body, or use empty (will estimate from headers only)
905                    let body = ctx.body_buffer.as_slice();
906
907                    // Use client IP as the rate limit key (could be enhanced to use API key header)
908                    let rate_limit_key = &ctx.client_ip;
909
910                    if let Some(check_result) = self.inference_rate_limit_manager.check(
911                        route_id,
912                        rate_limit_key,
913                        headers,
914                        body,
915                    ) {
916                        // Store inference rate limiting context for recording actual tokens later
917                        ctx.inference_rate_limit_enabled = true;
918                        ctx.inference_estimated_tokens = check_result.estimated_tokens;
919                        ctx.inference_rate_limit_key = Some(rate_limit_key.to_string());
920                        ctx.inference_model = check_result.model.clone();
921
922                        if !check_result.is_allowed() {
923                            let retry_after_ms = check_result.retry_after_ms();
924                            let retry_after_secs = retry_after_ms.div_ceil(1000);
925
926                            warn!(
927                                correlation_id = %ctx.trace_id,
928                                route_id = route_id,
929                                client_ip = %ctx.client_ip,
930                                estimated_tokens = check_result.estimated_tokens,
931                                model = ?check_result.model,
932                                retry_after_ms = retry_after_ms,
933                                "Inference rate limit exceeded (tokens)"
934                            );
935                            self.metrics
936                                .record_blocked_request("inference_rate_limited");
937
938                            // Audit log the token rate limit
939                            let audit_entry = AuditLogEntry::new(
940                                &ctx.trace_id,
941                                AuditEventType::RateLimitExceeded,
942                                &ctx.method,
943                                &ctx.path,
944                                &ctx.client_ip,
945                            )
946                            .with_route_id(route_id)
947                            .with_status_code(429)
948                            .with_reason(format!(
949                                "Token rate limit exceeded: estimated {} tokens, model={:?}",
950                                check_result.estimated_tokens, check_result.model
951                            ));
952                            self.log_manager.log_audit(&audit_entry);
953
954                            // Send 429 response with appropriate headers
955                            let body = "Token rate limit exceeded";
956                            let reset_at = std::time::SystemTime::now()
957                                .duration_since(std::time::UNIX_EPOCH)
958                                .unwrap_or_default()
959                                .as_secs()
960                                + retry_after_secs;
961
962                            // Use simplified error write for inference rate limit
963                            crate::http_helpers::write_rate_limit_error(
964                                session,
965                                429,
966                                body,
967                                0, // No request limit
968                                0, // No remaining
969                                reset_at,
970                                retry_after_secs,
971                            )
972                            .await?;
973                            return Ok(true); // Request complete, don't continue
974                        }
975
976                        trace!(
977                            correlation_id = %ctx.trace_id,
978                            route_id = route_id,
979                            estimated_tokens = check_result.estimated_tokens,
980                            model = ?check_result.model,
981                            "Inference rate limit check passed"
982                        );
983
984                        // Check budget tracking (cumulative per-period limits)
985                        if self.inference_rate_limit_manager.has_budget(route_id) {
986                            ctx.inference_budget_enabled = true;
987
988                            if let Some(budget_result) =
989                                self.inference_rate_limit_manager.check_budget(
990                                    route_id,
991                                    rate_limit_key,
992                                    check_result.estimated_tokens,
993                                )
994                            {
995                                if !budget_result.is_allowed() {
996                                    let retry_after_secs = budget_result.retry_after_secs();
997
998                                    warn!(
999                                        correlation_id = %ctx.trace_id,
1000                                        route_id = route_id,
1001                                        client_ip = %ctx.client_ip,
1002                                        estimated_tokens = check_result.estimated_tokens,
1003                                        retry_after_secs = retry_after_secs,
1004                                        "Token budget exhausted"
1005                                    );
1006
1007                                    ctx.inference_budget_exhausted = true;
1008                                    self.metrics.record_blocked_request("budget_exhausted");
1009
1010                                    // Audit log the budget exhaustion
1011                                    let audit_entry = AuditLogEntry::new(
1012                                        &ctx.trace_id,
1013                                        AuditEventType::RateLimitExceeded,
1014                                        &ctx.method,
1015                                        &ctx.path,
1016                                        &ctx.client_ip,
1017                                    )
1018                                    .with_route_id(route_id)
1019                                    .with_status_code(429)
1020                                    .with_reason("Token budget exhausted".to_string());
1021                                    self.log_manager.log_audit(&audit_entry);
1022
1023                                    // Send 429 response with budget headers
1024                                    let body = "Token budget exhausted";
1025                                    let reset_at = std::time::SystemTime::now()
1026                                        .duration_since(std::time::UNIX_EPOCH)
1027                                        .unwrap_or_default()
1028                                        .as_secs()
1029                                        + retry_after_secs;
1030
1031                                    crate::http_helpers::write_rate_limit_error(
1032                                        session,
1033                                        429,
1034                                        body,
1035                                        0,
1036                                        0,
1037                                        reset_at,
1038                                        retry_after_secs,
1039                                    )
1040                                    .await?;
1041                                    return Ok(true);
1042                                }
1043
1044                                // Capture budget status for response headers
1045                                let remaining = match &budget_result {
1046                                    zentinel_common::budget::BudgetCheckResult::Allowed {
1047                                        remaining,
1048                                    } => *remaining as i64,
1049                                    zentinel_common::budget::BudgetCheckResult::Soft {
1050                                        remaining,
1051                                        ..
1052                                    } => *remaining,
1053                                    _ => 0,
1054                                };
1055                                ctx.inference_budget_remaining = Some(remaining);
1056
1057                                // Get period reset time from budget status
1058                                if let Some(status) = self
1059                                    .inference_rate_limit_manager
1060                                    .budget_status(route_id, rate_limit_key)
1061                                {
1062                                    ctx.inference_budget_period_reset = Some(status.period_end);
1063                                }
1064
1065                                trace!(
1066                                    correlation_id = %ctx.trace_id,
1067                                    route_id = route_id,
1068                                    budget_remaining = remaining,
1069                                    "Token budget check passed"
1070                                );
1071                            }
1072                        }
1073
1074                        // Check if cost attribution is enabled
1075                        if self
1076                            .inference_rate_limit_manager
1077                            .has_cost_attribution(route_id)
1078                        {
1079                            ctx.inference_cost_enabled = true;
1080                        }
1081                    }
1082                }
1083            }
1084        }
1085
1086        // Prompt injection guardrail (for inference routes)
1087        if let Some(ref route_config) = ctx.route_config {
1088            if let Some(ref inference) = route_config.inference {
1089                if let Some(ref guardrails) = inference.guardrails {
1090                    if let Some(ref pi_config) = guardrails.prompt_injection {
1091                        if pi_config.enabled && !ctx.body_buffer.is_empty() {
1092                            ctx.guardrails_enabled = true;
1093
1094                            // Extract content from request body
1095                            if let Some(content) = extract_inference_content(&ctx.body_buffer) {
1096                                let result = self
1097                                    .guardrail_processor
1098                                    .check_prompt_injection(
1099                                        pi_config,
1100                                        &content,
1101                                        ctx.inference_model.as_deref(),
1102                                        ctx.route_id.as_deref(),
1103                                        &ctx.trace_id,
1104                                    )
1105                                    .await;
1106
1107                                match result {
1108                                    PromptInjectionResult::Blocked {
1109                                        status,
1110                                        message,
1111                                        detections,
1112                                    } => {
1113                                        warn!(
1114                                            correlation_id = %ctx.trace_id,
1115                                            route_id = ctx.route_id.as_deref().unwrap_or("unknown"),
1116                                            detection_count = detections.len(),
1117                                            "Prompt injection detected, blocking request"
1118                                        );
1119
1120                                        self.metrics.record_blocked_request("prompt_injection");
1121
1122                                        // Store detection categories for logging
1123                                        ctx.guardrail_detection_categories =
1124                                            detections.iter().map(|d| d.category.clone()).collect();
1125
1126                                        // Audit log the block
1127                                        let audit_entry = AuditLogEntry::new(
1128                                            &ctx.trace_id,
1129                                            AuditEventType::Blocked,
1130                                            &ctx.method,
1131                                            &ctx.path,
1132                                            &ctx.client_ip,
1133                                        )
1134                                        .with_route_id(ctx.route_id.as_deref().unwrap_or("unknown"))
1135                                        .with_status_code(status)
1136                                        .with_reason("Prompt injection detected".to_string());
1137                                        self.log_manager.log_audit(&audit_entry);
1138
1139                                        // Send error response
1140                                        crate::http_helpers::write_json_error(
1141                                            session,
1142                                            status,
1143                                            "prompt_injection_blocked",
1144                                            Some(&message),
1145                                        )
1146                                        .await?;
1147                                        return Ok(true);
1148                                    }
1149                                    PromptInjectionResult::Detected { detections } => {
1150                                        // Log but allow
1151                                        warn!(
1152                                            correlation_id = %ctx.trace_id,
1153                                            route_id = ctx.route_id.as_deref().unwrap_or("unknown"),
1154                                            detection_count = detections.len(),
1155                                            "Prompt injection detected (logged only)"
1156                                        );
1157                                        ctx.guardrail_detection_categories =
1158                                            detections.iter().map(|d| d.category.clone()).collect();
1159                                    }
1160                                    PromptInjectionResult::Warning { detections } => {
1161                                        // Set flag for response header
1162                                        ctx.guardrail_warning = true;
1163                                        ctx.guardrail_detection_categories =
1164                                            detections.iter().map(|d| d.category.clone()).collect();
1165                                        debug!(
1166                                            correlation_id = %ctx.trace_id,
1167                                            "Prompt injection warning set"
1168                                        );
1169                                    }
1170                                    PromptInjectionResult::Clean => {
1171                                        trace!(
1172                                            correlation_id = %ctx.trace_id,
1173                                            "No prompt injection detected"
1174                                        );
1175                                    }
1176                                    PromptInjectionResult::Error { message } => {
1177                                        // Already logged in processor, just trace here
1178                                        trace!(
1179                                            correlation_id = %ctx.trace_id,
1180                                            error = %message,
1181                                            "Prompt injection check error (failure mode applied)"
1182                                        );
1183                                    }
1184                                }
1185                            }
1186                        }
1187                    }
1188                }
1189            }
1190        }
1191
1192        // Geo filtering
1193        if let Some(route_id) = ctx.route_id.as_deref() {
1194            if let Some(ref route_config) = ctx.route_config {
1195                for filter_id in &route_config.filters {
1196                    if let Some(result) = self.geo_filter_manager.check(filter_id, &ctx.client_ip) {
1197                        // Store country code for response header
1198                        ctx.geo_country_code = result.country_code.clone();
1199                        ctx.geo_lookup_performed = true;
1200
1201                        if !result.allowed {
1202                            warn!(
1203                                correlation_id = %ctx.trace_id,
1204                                route_id = route_id,
1205                                client_ip = %ctx.client_ip,
1206                                country = ?result.country_code,
1207                                filter_id = %filter_id,
1208                                "Request blocked by geo filter"
1209                            );
1210                            self.metrics.record_blocked_request("geo_blocked");
1211
1212                            // Audit log the geo block
1213                            let audit_entry = AuditLogEntry::new(
1214                                &ctx.trace_id,
1215                                AuditEventType::Blocked,
1216                                &ctx.method,
1217                                &ctx.path,
1218                                &ctx.client_ip,
1219                            )
1220                            .with_route_id(route_id)
1221                            .with_status_code(result.status_code)
1222                            .with_reason(format!(
1223                                "Geo blocked: country={}, filter={}",
1224                                result.country_code.as_deref().unwrap_or("unknown"),
1225                                filter_id
1226                            ));
1227                            self.log_manager.log_audit(&audit_entry);
1228
1229                            // Send geo block response
1230                            let body = result
1231                                .block_message
1232                                .unwrap_or_else(|| "Access denied".to_string());
1233
1234                            crate::http_helpers::write_error(
1235                                session,
1236                                result.status_code,
1237                                &body,
1238                                "text/plain",
1239                            )
1240                            .await?;
1241                            return Ok(true); // Request complete, don't continue
1242                        }
1243
1244                        // Only check first geo filter that matches
1245                        break;
1246                    }
1247                }
1248            }
1249        }
1250
1251        // Route-level filters (CORS preflight, Timeout, Log)
1252        // Clone the Arc to avoid borrow conflict between &Config and &mut ctx
1253        let config_for_filters = std::sync::Arc::clone(
1254            ctx.config
1255                .get_or_insert_with(|| self.config_manager.current()),
1256        );
1257        if super::filters::apply_request_filters(session, ctx, &config_for_filters).await? {
1258            return Ok(true); // Filter handled request (e.g. CORS preflight)
1259        }
1260
1261        // Check for WebSocket upgrade requests
1262        let is_websocket_upgrade = session
1263            .req_header()
1264            .headers
1265            .get(http::header::UPGRADE)
1266            .map(|v| v.as_bytes().eq_ignore_ascii_case(b"websocket"))
1267            .unwrap_or(false);
1268
1269        if is_websocket_upgrade {
1270            ctx.is_websocket_upgrade = true;
1271
1272            // Check if route allows WebSocket upgrades
1273            if let Some(ref route_config) = ctx.route_config {
1274                if !route_config.websocket {
1275                    warn!(
1276                        correlation_id = %ctx.trace_id,
1277                        route_id = ctx.route_id.as_deref().unwrap_or("unknown"),
1278                        client_ip = %ctx.client_ip,
1279                        "WebSocket upgrade rejected: not enabled for route"
1280                    );
1281
1282                    self.metrics.record_blocked_request("websocket_not_enabled");
1283
1284                    // Audit log the rejection
1285                    let audit_entry = AuditLogEntry::new(
1286                        &ctx.trace_id,
1287                        AuditEventType::Blocked,
1288                        &ctx.method,
1289                        &ctx.path,
1290                        &ctx.client_ip,
1291                    )
1292                    .with_route_id(ctx.route_id.as_deref().unwrap_or("unknown"))
1293                    .with_action("websocket_rejected")
1294                    .with_reason("WebSocket not enabled for route");
1295                    self.log_manager.log_audit(&audit_entry);
1296
1297                    // Send 403 Forbidden response
1298                    crate::http_helpers::write_error(
1299                        session,
1300                        403,
1301                        "WebSocket not enabled for this route",
1302                        "text/plain",
1303                    )
1304                    .await?;
1305                    return Ok(true); // Request complete, don't continue
1306                }
1307
1308                debug!(
1309                    correlation_id = %ctx.trace_id,
1310                    route_id = ctx.route_id.as_deref().unwrap_or("unknown"),
1311                    "WebSocket upgrade request allowed"
1312                );
1313
1314                // Check for WebSocket frame inspection
1315                if route_config.websocket_inspection {
1316                    // Check for compression negotiation - skip inspection if permessage-deflate
1317                    let has_compression = session
1318                        .req_header()
1319                        .headers
1320                        .get("Sec-WebSocket-Extensions")
1321                        .and_then(|v| v.to_str().ok())
1322                        .map(|s| s.contains("permessage-deflate"))
1323                        .unwrap_or(false);
1324
1325                    if has_compression {
1326                        debug!(
1327                            correlation_id = %ctx.trace_id,
1328                            route_id = ctx.route_id.as_deref().unwrap_or("unknown"),
1329                            "WebSocket inspection skipped: permessage-deflate negotiated"
1330                        );
1331                        ctx.websocket_skip_inspection = true;
1332                    } else {
1333                        ctx.websocket_inspection_enabled = true;
1334
1335                        // Get agents that handle WebSocketFrame events
1336                        ctx.websocket_inspection_agents = self.agent_manager.get_agents_for_event(
1337                            zentinel_agent_protocol::EventType::WebSocketFrame,
1338                        );
1339
1340                        debug!(
1341                            correlation_id = %ctx.trace_id,
1342                            route_id = ctx.route_id.as_deref().unwrap_or("unknown"),
1343                            agent_count = ctx.websocket_inspection_agents.len(),
1344                            "WebSocket frame inspection enabled"
1345                        );
1346                    }
1347                }
1348            }
1349        }
1350
1351        // Use cached route config from upstream_peer (avoids duplicate route matching)
1352        // Handle static file and builtin routes
1353        if let Some(route_config) = ctx.route_config.clone() {
1354            if route_config.service_type == zentinel_config::ServiceType::Static {
1355                trace!(
1356                    correlation_id = %ctx.trace_id,
1357                    route_id = ctx.route_id.as_deref().unwrap_or("unknown"),
1358                    "Handling static file route"
1359                );
1360                // Create a minimal RouteMatch for the handler
1361                let route_match = crate::routing::RouteMatch {
1362                    route_id: zentinel_common::RouteId::new(ctx.route_id.as_deref().unwrap_or("")),
1363                    config: route_config.clone(),
1364                };
1365                return self.handle_static_route(session, ctx, &route_match).await;
1366            } else if route_config.service_type == zentinel_config::ServiceType::Builtin {
1367                trace!(
1368                    correlation_id = %ctx.trace_id,
1369                    route_id = ctx.route_id.as_deref().unwrap_or("unknown"),
1370                    builtin_handler = ?route_config.builtin_handler,
1371                    "Handling builtin route"
1372                );
1373                // Create a minimal RouteMatch for the handler
1374                let route_match = crate::routing::RouteMatch {
1375                    route_id: zentinel_common::RouteId::new(ctx.route_id.as_deref().unwrap_or("")),
1376                    config: route_config.clone(),
1377                };
1378                return self.handle_builtin_route(session, ctx, &route_match).await;
1379            }
1380        }
1381
1382        // API validation for API routes
1383        if let Some(route_id) = ctx.route_id.clone() {
1384            if let Some(validator) = self.validators.get(&route_id).await {
1385                trace!(
1386                    correlation_id = %ctx.trace_id,
1387                    route_id = %route_id,
1388                    "Running API schema validation"
1389                );
1390                if let Some(result) = self
1391                    .validate_api_request(session, ctx, &route_id, &validator)
1392                    .await?
1393                {
1394                    debug!(
1395                        correlation_id = %ctx.trace_id,
1396                        route_id = %route_id,
1397                        validation_passed = result,
1398                        "API validation complete"
1399                    );
1400                    return Ok(result);
1401                }
1402            }
1403        }
1404
1405        // Get client address before mutable borrow
1406        let client_addr = session
1407            .client_addr()
1408            .map(|a| format!("{}", a))
1409            .unwrap_or_else(|| "unknown".to_string());
1410        let client_port = session.client_addr().map(|_| 0).unwrap_or(0);
1411
1412        let req_header = session.req_header_mut();
1413
1414        // Add correlation ID header
1415        req_header
1416            .insert_header("X-Correlation-Id", &ctx.trace_id)
1417            .ok();
1418        req_header.insert_header("X-Forwarded-By", "Zentinel").ok();
1419
1420        // Use cached config (set in upstream_peer, or fetch now if needed)
1421        let config = ctx
1422            .config
1423            .get_or_insert_with(|| self.config_manager.current());
1424
1425        // Enforce header limits (fast path: skip if limits are very high)
1426        const HEADER_LIMIT_THRESHOLD: usize = 1024 * 1024; // 1MB = effectively unlimited
1427
1428        // Header count check - O(1)
1429        let header_count = req_header.headers.len();
1430        if config.limits.max_header_count < HEADER_LIMIT_THRESHOLD
1431            && header_count > config.limits.max_header_count
1432        {
1433            warn!(
1434                correlation_id = %ctx.trace_id,
1435                header_count = header_count,
1436                limit = config.limits.max_header_count,
1437                "Request blocked: exceeds header count limit"
1438            );
1439
1440            self.metrics.record_blocked_request("header_count_exceeded");
1441            return Err(Error::explain(ErrorType::InternalError, "Too many headers"));
1442        }
1443
1444        // Header size check - O(n), skip if limit is very high
1445        if config.limits.max_header_size_bytes < HEADER_LIMIT_THRESHOLD {
1446            let total_header_size: usize = req_header
1447                .headers
1448                .iter()
1449                .map(|(k, v)| k.as_str().len() + v.len())
1450                .sum();
1451
1452            if total_header_size > config.limits.max_header_size_bytes {
1453                warn!(
1454                    correlation_id = %ctx.trace_id,
1455                    header_size = total_header_size,
1456                    limit = config.limits.max_header_size_bytes,
1457                    "Request blocked: exceeds header size limit"
1458                );
1459
1460                self.metrics.record_blocked_request("header_size_exceeded");
1461                return Err(Error::explain(
1462                    ErrorType::InternalError,
1463                    "Headers too large",
1464                ));
1465            }
1466        }
1467
1468        // Process through external agents
1469        trace!(
1470            correlation_id = %ctx.trace_id,
1471            "Processing request through agents"
1472        );
1473        if let Err(e) = self
1474            .process_agents(session, ctx, &client_addr, client_port)
1475            .await
1476        {
1477            // Check if this is an HTTPStatus error (e.g., agent block or fail-closed)
1478            // In that case, we need to send a proper HTTP response instead of just closing the connection
1479            if let ErrorType::HTTPStatus(status) = e.etype() {
1480                // Extract the message from the error (the context part after "HTTPStatus context:")
1481                let error_msg = e.to_string();
1482                let body = error_msg
1483                    .split("context:")
1484                    .nth(1)
1485                    .map(|s| s.trim())
1486                    .unwrap_or("Request blocked");
1487                debug!(
1488                    correlation_id = %ctx.trace_id,
1489                    status = status,
1490                    body = %body,
1491                    "Sending HTTP error response for agent block"
1492                );
1493                crate::http_helpers::write_error(session, *status, body, "text/plain").await?;
1494                return Ok(true); // Request complete, don't continue to upstream
1495            }
1496            // For other errors, propagate them
1497            return Err(e);
1498        }
1499
1500        trace!(
1501            correlation_id = %ctx.trace_id,
1502            "Request filter phase complete, forwarding to upstream"
1503        );
1504
1505        Ok(false) // Continue processing
1506    }
1507
1508    /// Process incoming request body chunks.
1509    /// Used for body size enforcement and WAF/agent inspection.
1510    ///
1511    /// Supports two modes:
1512    /// - **Buffer mode** (default): Buffer chunks until end of stream or limit, then send to agents
1513    /// - **Stream mode**: Send each chunk immediately to agents as it arrives
1514    async fn request_body_filter(
1515        &self,
1516        _session: &mut Session,
1517        body: &mut Option<Bytes>,
1518        end_of_stream: bool,
1519        ctx: &mut Self::CTX,
1520    ) -> Result<(), Box<Error>> {
1521        use zentinel_config::BodyStreamingMode;
1522
1523        // Handle WebSocket frame inspection (client -> server)
1524        if ctx.is_websocket_upgrade {
1525            if let Some(ref handler) = ctx.websocket_handler {
1526                let result = handler.process_client_data(body.take()).await;
1527                match result {
1528                    crate::websocket::ProcessResult::Forward(data) => {
1529                        *body = data;
1530                    }
1531                    crate::websocket::ProcessResult::Close(reason) => {
1532                        warn!(
1533                            correlation_id = %ctx.trace_id,
1534                            code = reason.code,
1535                            reason = %reason.reason,
1536                            "WebSocket connection closed by agent (client->server)"
1537                        );
1538                        // Return an error to close the connection
1539                        return Err(Error::explain(
1540                            ErrorType::InternalError,
1541                            format!("WebSocket closed: {} {}", reason.code, reason.reason),
1542                        ));
1543                    }
1544                }
1545            }
1546            // Skip normal body processing for WebSocket
1547            return Ok(());
1548        }
1549
1550        // Track request body size
1551        let chunk_len = body.as_ref().map(|b| b.len()).unwrap_or(0);
1552        if chunk_len > 0 {
1553            ctx.request_body_bytes += chunk_len as u64;
1554
1555            trace!(
1556                correlation_id = %ctx.trace_id,
1557                chunk_size = chunk_len,
1558                total_body_bytes = ctx.request_body_bytes,
1559                end_of_stream = end_of_stream,
1560                streaming_mode = ?ctx.request_body_streaming_mode,
1561                "Processing request body chunk"
1562            );
1563
1564            // Check body size limit (use cached config)
1565            let config = ctx
1566                .config
1567                .get_or_insert_with(|| self.config_manager.current());
1568            if ctx.request_body_bytes > config.limits.max_body_size_bytes as u64 {
1569                warn!(
1570                    correlation_id = %ctx.trace_id,
1571                    body_bytes = ctx.request_body_bytes,
1572                    limit = config.limits.max_body_size_bytes,
1573                    "Request body size limit exceeded"
1574                );
1575                self.metrics.record_blocked_request("body_size_exceeded");
1576                return Err(Error::explain(
1577                    ErrorType::InternalError,
1578                    "Request body too large",
1579                ));
1580            }
1581        }
1582
1583        // Body inspection for agents (WAF, etc.)
1584        if ctx.body_inspection_enabled && !ctx.body_inspection_agents.is_empty() {
1585            let config = ctx
1586                .config
1587                .get_or_insert_with(|| self.config_manager.current());
1588            let max_inspection_bytes = config
1589                .waf
1590                .as_ref()
1591                .map(|w| w.body_inspection.max_inspection_bytes as u64)
1592                .unwrap_or(1024 * 1024);
1593
1594            match ctx.request_body_streaming_mode {
1595                BodyStreamingMode::Stream => {
1596                    // Stream mode: send each chunk immediately
1597                    if body.is_some() {
1598                        self.process_body_chunk_streaming(body, end_of_stream, ctx)
1599                            .await?;
1600                    } else if end_of_stream && ctx.agent_needs_more {
1601                        // Send final empty chunk to signal end
1602                        self.process_body_chunk_streaming(body, end_of_stream, ctx)
1603                            .await?;
1604                    }
1605                }
1606                BodyStreamingMode::Hybrid { buffer_threshold } => {
1607                    // Hybrid mode: buffer up to threshold, then stream
1608                    if ctx.body_bytes_inspected < buffer_threshold as u64 {
1609                        // Still in buffering phase
1610                        if let Some(ref chunk) = body {
1611                            let bytes_to_buffer = std::cmp::min(
1612                                chunk.len(),
1613                                (buffer_threshold as u64 - ctx.body_bytes_inspected) as usize,
1614                            );
1615                            ctx.body_buffer.extend_from_slice(&chunk[..bytes_to_buffer]);
1616                            ctx.body_bytes_inspected += bytes_to_buffer as u64;
1617
1618                            // If we've reached threshold or end of stream, switch to streaming
1619                            if ctx.body_bytes_inspected >= buffer_threshold as u64 || end_of_stream
1620                            {
1621                                // Send buffered content first
1622                                self.send_buffered_body_to_agents(
1623                                    end_of_stream && chunk.len() == bytes_to_buffer,
1624                                    ctx,
1625                                )
1626                                .await?;
1627                                ctx.body_buffer.clear();
1628
1629                                // If there's remaining data in this chunk, stream it
1630                                if bytes_to_buffer < chunk.len() {
1631                                    let remaining = chunk.slice(bytes_to_buffer..);
1632                                    let mut remaining_body = Some(remaining);
1633                                    self.process_body_chunk_streaming(
1634                                        &mut remaining_body,
1635                                        end_of_stream,
1636                                        ctx,
1637                                    )
1638                                    .await?;
1639                                }
1640                            }
1641                        }
1642                    } else {
1643                        // Past threshold, stream directly
1644                        self.process_body_chunk_streaming(body, end_of_stream, ctx)
1645                            .await?;
1646                    }
1647                }
1648                BodyStreamingMode::Buffer => {
1649                    // Buffer mode: collect chunks until ready to send
1650                    if let Some(ref chunk) = body {
1651                        if ctx.body_bytes_inspected < max_inspection_bytes {
1652                            let bytes_to_inspect = std::cmp::min(
1653                                chunk.len() as u64,
1654                                max_inspection_bytes - ctx.body_bytes_inspected,
1655                            ) as usize;
1656
1657                            ctx.body_buffer
1658                                .extend_from_slice(&chunk[..bytes_to_inspect]);
1659                            ctx.body_bytes_inspected += bytes_to_inspect as u64;
1660
1661                            trace!(
1662                                correlation_id = %ctx.trace_id,
1663                                bytes_inspected = ctx.body_bytes_inspected,
1664                                max_inspection_bytes = max_inspection_bytes,
1665                                buffer_size = ctx.body_buffer.len(),
1666                                "Buffering body for agent inspection"
1667                            );
1668                        }
1669                    }
1670
1671                    // Send when complete or limit reached
1672                    let should_send =
1673                        end_of_stream || ctx.body_bytes_inspected >= max_inspection_bytes;
1674                    if should_send && !ctx.body_buffer.is_empty() {
1675                        self.send_buffered_body_to_agents(end_of_stream, ctx)
1676                            .await?;
1677                        ctx.body_buffer.clear();
1678                    }
1679                }
1680            }
1681        }
1682
1683        if end_of_stream {
1684            trace!(
1685                correlation_id = %ctx.trace_id,
1686                total_body_bytes = ctx.request_body_bytes,
1687                bytes_inspected = ctx.body_bytes_inspected,
1688                "Request body complete"
1689            );
1690        }
1691
1692        Ok(())
1693    }
1694
1695    async fn response_filter(
1696        &self,
1697        session: &mut Session,
1698        upstream_response: &mut ResponseHeader,
1699        ctx: &mut Self::CTX,
1700    ) -> Result<(), Box<Error>> {
1701        let status = upstream_response.status.as_u16();
1702        let duration = ctx.elapsed();
1703
1704        trace!(
1705            correlation_id = %ctx.trace_id,
1706            status = status,
1707            "Starting response filter phase"
1708        );
1709
1710        // Handle WebSocket 101 Switching Protocols
1711        if status == 101 && ctx.is_websocket_upgrade {
1712            if ctx.websocket_inspection_enabled && !ctx.websocket_skip_inspection {
1713                // Create WebSocket inspector and handler with metrics
1714                let inspector = crate::websocket::WebSocketInspector::with_metrics(
1715                    self.agent_manager.clone(),
1716                    ctx.route_id
1717                        .clone()
1718                        .unwrap_or_else(|| "unknown".to_string()),
1719                    ctx.trace_id.clone(),
1720                    ctx.client_ip.clone(),
1721                    100, // 100ms timeout per frame inspection
1722                    Some(self.metrics.clone()),
1723                );
1724
1725                let handler = crate::websocket::WebSocketHandler::new(
1726                    std::sync::Arc::new(inspector),
1727                    1024 * 1024, // 1MB max frame size
1728                );
1729
1730                ctx.websocket_handler = Some(std::sync::Arc::new(handler));
1731
1732                info!(
1733                    correlation_id = %ctx.trace_id,
1734                    route_id = ctx.route_id.as_deref().unwrap_or("unknown"),
1735                    agent_count = ctx.websocket_inspection_agents.len(),
1736                    "WebSocket upgrade successful, frame inspection enabled"
1737                );
1738            } else if ctx.websocket_skip_inspection {
1739                debug!(
1740                    correlation_id = %ctx.trace_id,
1741                    route_id = ctx.route_id.as_deref().unwrap_or("unknown"),
1742                    "WebSocket upgrade successful, inspection skipped (compression negotiated)"
1743                );
1744            } else {
1745                debug!(
1746                    correlation_id = %ctx.trace_id,
1747                    route_id = ctx.route_id.as_deref().unwrap_or("unknown"),
1748                    "WebSocket upgrade successful"
1749                );
1750            }
1751        }
1752
1753        // Add correlation ID to response
1754        upstream_response.insert_header("X-Correlation-Id", &ctx.trace_id)?;
1755
1756        // Add rate limit headers if rate limiting was applied
1757        if let Some(ref rate_info) = ctx.rate_limit_info {
1758            upstream_response.insert_header("X-RateLimit-Limit", rate_info.limit.to_string())?;
1759            upstream_response
1760                .insert_header("X-RateLimit-Remaining", rate_info.remaining.to_string())?;
1761            upstream_response.insert_header("X-RateLimit-Reset", rate_info.reset_at.to_string())?;
1762        }
1763
1764        // Add token budget headers if budget tracking was enabled
1765        if ctx.inference_budget_enabled {
1766            if let Some(remaining) = ctx.inference_budget_remaining {
1767                upstream_response.insert_header("X-Budget-Remaining", remaining.to_string())?;
1768            }
1769            if let Some(period_reset) = ctx.inference_budget_period_reset {
1770                // Format as ISO 8601 timestamp
1771                let reset_datetime = chrono::DateTime::from_timestamp(period_reset as i64, 0)
1772                    .map(|dt| dt.format("%Y-%m-%dT%H:%M:%SZ").to_string())
1773                    .unwrap_or_else(|| period_reset.to_string());
1774                upstream_response.insert_header("X-Budget-Period-Reset", reset_datetime)?;
1775            }
1776        }
1777
1778        // Add GeoIP country header if geo lookup was performed
1779        if let Some(ref country_code) = ctx.geo_country_code {
1780            upstream_response.insert_header("X-GeoIP-Country", country_code)?;
1781        }
1782
1783        // Apply route-specific response header modifications (policies)
1784        if let Some(ref route_config) = ctx.route_config {
1785            let mods = &route_config.policies.response_headers;
1786            // Rename runs before set/add/remove
1787            for (old_name, new_name) in &mods.rename {
1788                if let Some(value) = upstream_response
1789                    .headers
1790                    .get(old_name)
1791                    .and_then(|v| v.to_str().ok())
1792                {
1793                    let owned = value.to_string();
1794                    upstream_response
1795                        .insert_header(new_name.clone(), &owned)
1796                        .ok();
1797                    upstream_response.remove_header(old_name);
1798                }
1799            }
1800            for (name, value) in &mods.set {
1801                upstream_response
1802                    .insert_header(name.clone(), value.as_str())
1803                    .ok();
1804            }
1805            for (name, value) in &mods.add {
1806                upstream_response
1807                    .append_header(name.clone(), value.as_str())
1808                    .ok();
1809            }
1810            for name in &mods.remove {
1811                upstream_response.remove_header(name);
1812            }
1813        }
1814
1815        // Inject Cache-Status header (RFC 9211) if enabled
1816        if let Some(ref cache_status) = ctx.cache_status {
1817            let status_header_enabled = ctx
1818                .config
1819                .as_ref()
1820                .and_then(|c| c.cache.as_ref())
1821                .map(|c| c.status_header)
1822                .unwrap_or(false);
1823
1824            if status_header_enabled {
1825                let cache_name = ctx
1826                    .config
1827                    .as_ref()
1828                    .and_then(|c| c.cache.as_ref())
1829                    .map(|c| c.status_header_name.as_str())
1830                    .unwrap_or("zentinel");
1831
1832                let value = match cache_status {
1833                    super::context::CacheStatus::HitMemory => {
1834                        format!("{cache_name}; hit; detail=memory")
1835                    }
1836                    super::context::CacheStatus::HitDisk => {
1837                        format!("{cache_name}; hit; detail=disk")
1838                    }
1839                    super::context::CacheStatus::Hit => format!("{cache_name}; hit"),
1840                    super::context::CacheStatus::HitStale => format!("{cache_name}; fwd=stale"),
1841                    super::context::CacheStatus::Miss => format!("{cache_name}; fwd=miss"),
1842                    super::context::CacheStatus::Bypass(reason) => match *reason {
1843                        "method" => format!("{cache_name}; fwd=bypass; detail=method"),
1844                        "disabled" => format!("{cache_name}; fwd=bypass; detail=disabled"),
1845                        "no-route" => format!("{cache_name}; fwd=bypass; detail=no-route"),
1846                        _ => format!("{cache_name}; fwd=bypass"),
1847                    },
1848                };
1849                upstream_response.insert_header("Cache-Status", &value).ok();
1850            }
1851        }
1852
1853        // Apply response-phase route filters (Headers, CORS, Compress, Log)
1854        if let Some(config) = ctx.config.as_ref().map(std::sync::Arc::clone) {
1855            super::filters::apply_response_filters(upstream_response, ctx, &config);
1856        }
1857
1858        // Enable Pingora response compression if Compress filter marked it eligible
1859        if ctx.compress_enabled {
1860            session.upstream_compression.adjust_level(6);
1861        }
1862
1863        // Apply per-listener keepalive timeout
1864        if let Some(keepalive_secs) = ctx.listener_keepalive_timeout_secs {
1865            session
1866                .downstream_session
1867                .set_keepalive(Some(keepalive_secs));
1868        }
1869
1870        // Add sticky session cookie if a new assignment was made
1871        if ctx.sticky_session_new_assignment {
1872            if let Some(ref set_cookie_header) = ctx.sticky_session_set_cookie {
1873                upstream_response.insert_header("Set-Cookie", set_cookie_header)?;
1874                trace!(
1875                    correlation_id = %ctx.trace_id,
1876                    sticky_target_index = ?ctx.sticky_target_index,
1877                    "Added sticky session Set-Cookie header"
1878                );
1879            }
1880        }
1881
1882        // Add guardrail warning header if prompt injection was detected (warn mode)
1883        if ctx.guardrail_warning {
1884            upstream_response.insert_header("X-Guardrail-Warning", "prompt-injection-detected")?;
1885        }
1886
1887        // Add fallback routing headers if fallback was used
1888        if ctx.used_fallback() {
1889            upstream_response.insert_header("X-Fallback-Used", "true")?;
1890
1891            if let Some(ref upstream) = ctx.upstream {
1892                upstream_response.insert_header("X-Fallback-Upstream", upstream)?;
1893            }
1894
1895            if let Some(ref reason) = ctx.fallback_reason {
1896                upstream_response.insert_header("X-Fallback-Reason", reason.to_string())?;
1897            }
1898
1899            if let Some(ref original) = ctx.original_upstream {
1900                upstream_response.insert_header("X-Original-Upstream", original)?;
1901            }
1902
1903            if let Some(ref mapping) = ctx.model_mapping_applied {
1904                upstream_response
1905                    .insert_header("X-Model-Mapping", format!("{} -> {}", mapping.0, mapping.1))?;
1906            }
1907
1908            trace!(
1909                correlation_id = %ctx.trace_id,
1910                fallback_attempt = ctx.fallback_attempt,
1911                fallback_upstream = ctx.upstream.as_deref().unwrap_or("unknown"),
1912                original_upstream = ctx.original_upstream.as_deref().unwrap_or("unknown"),
1913                "Added fallback response headers"
1914            );
1915
1916            // Record fallback success metrics for successful responses (2xx/3xx)
1917            if status < 400 {
1918                if let Some(metrics) = get_fallback_metrics() {
1919                    metrics.record_fallback_success(
1920                        ctx.route_id.as_deref().unwrap_or("unknown"),
1921                        ctx.upstream.as_deref().unwrap_or("unknown"),
1922                    );
1923                }
1924            }
1925        }
1926
1927        // Initialize streaming token counter for SSE responses on inference routes
1928        if ctx.inference_rate_limit_enabled {
1929            // Check if this is an SSE response
1930            let content_type = upstream_response
1931                .headers
1932                .get("content-type")
1933                .and_then(|ct| ct.to_str().ok());
1934
1935            if is_sse_response(content_type) {
1936                // Get provider from route config
1937                let provider = ctx
1938                    .route_config
1939                    .as_ref()
1940                    .and_then(|r| r.inference.as_ref())
1941                    .map(|i| i.provider)
1942                    .unwrap_or_default();
1943
1944                ctx.inference_streaming_response = true;
1945                ctx.inference_streaming_counter = Some(StreamingTokenCounter::new(
1946                    provider,
1947                    ctx.inference_model.clone(),
1948                ));
1949
1950                trace!(
1951                    correlation_id = %ctx.trace_id,
1952                    content_type = ?content_type,
1953                    model = ?ctx.inference_model,
1954                    "Initialized streaming token counter for SSE response"
1955                );
1956            }
1957        }
1958
1959        // Process response headers through agents (for agents that subscribe to ResponseHeaders events)
1960        if !ctx.route_agent_ids.is_empty() {
1961            let agent_ids = ctx.route_agent_ids.clone();
1962            let mut resp_headers_map: std::collections::HashMap<String, Vec<String>> =
1963                std::collections::HashMap::with_capacity(upstream_response.headers.len());
1964            for (name, value) in upstream_response.headers.iter() {
1965                resp_headers_map
1966                    .entry(name.as_str().to_string())
1967                    .or_default()
1968                    .push(value.to_str().unwrap_or("").to_string());
1969            }
1970
1971            let agent_ctx = crate::agents::AgentCallContext {
1972                correlation_id: zentinel_common::CorrelationId::from_string(&ctx.trace_id),
1973                metadata: zentinel_agent_protocol::RequestMetadata {
1974                    correlation_id: ctx.trace_id.clone(),
1975                    request_id: uuid::Uuid::new_v4().to_string(),
1976                    client_ip: ctx.client_ip.clone(),
1977                    client_port: 0,
1978                    server_name: ctx.host.clone(),
1979                    protocol: "HTTP/1.1".to_string(),
1980                    tls_version: None,
1981                    tls_cipher: None,
1982                    route_id: ctx.route_id.clone(),
1983                    upstream_id: ctx.upstream.clone(),
1984                    timestamp: chrono::Utc::now().to_rfc3339(),
1985                    traceparent: ctx.traceparent(),
1986                },
1987                route_id: ctx.route_id.clone(),
1988                upstream_id: ctx.upstream.clone(),
1989                request_body: None,
1990                response_body: None,
1991            };
1992
1993            match self
1994                .agent_manager
1995                .process_response_headers(&agent_ctx, status, &resp_headers_map, &agent_ids)
1996                .await
1997            {
1998                Ok(decision) => {
1999                    // Apply response header modifications from agent
2000                    for op in &decision.response_headers {
2001                        match op {
2002                            zentinel_agent_protocol::HeaderOp::Set { name, value } => {
2003                                upstream_response
2004                                    .insert_header(name.clone(), value.as_str())
2005                                    .ok();
2006                            }
2007                            zentinel_agent_protocol::HeaderOp::Add { name, value } => {
2008                                upstream_response
2009                                    .append_header(name.clone(), value.as_str())
2010                                    .ok();
2011                            }
2012                            zentinel_agent_protocol::HeaderOp::Remove { name } => {
2013                                upstream_response.remove_header(name);
2014                            }
2015                        }
2016                    }
2017
2018                    // Check if any agent subscribes to response body events
2019                    let has_body_agents = self
2020                        .agent_manager
2021                        .any_agent_handles_event(
2022                            &agent_ids,
2023                            zentinel_agent_protocol::EventType::ResponseBodyChunk,
2024                        )
2025                        .await;
2026                    if has_body_agents {
2027                        ctx.response_agent_processing_enabled = true;
2028                        // Since agent may replace the body, Content-Length is invalid.
2029                        // Use Connection: close to signal end-of-body to the client.
2030                        upstream_response.insert_header("Connection", "close").ok();
2031                        session.downstream_session.set_keepalive(None);
2032                        debug!(
2033                            correlation_id = %ctx.trace_id,
2034                            "Enabling response body agent processing (agent subscribes to ResponseBody)"
2035                        );
2036                    }
2037
2038                    debug!(
2039                        correlation_id = %ctx.trace_id,
2040                        response_headers_modified = !decision.response_headers.is_empty(),
2041                        needs_body = ctx.response_agent_processing_enabled,
2042                        "Response headers processed through agents"
2043                    );
2044                }
2045                Err(e) => {
2046                    warn!(
2047                        correlation_id = %ctx.trace_id,
2048                        error = %e,
2049                        "Agent response header processing failed, continuing without agent"
2050                    );
2051                }
2052            }
2053        }
2054
2055        // Generate custom error pages for error responses
2056        if status >= 400 {
2057            trace!(
2058                correlation_id = %ctx.trace_id,
2059                status = status,
2060                "Handling error response"
2061            );
2062            self.handle_error_response(upstream_response, ctx).await?;
2063        }
2064
2065        // Record metrics
2066        self.metrics.record_request(
2067            ctx.route_id.as_deref().unwrap_or("unknown"),
2068            &ctx.method,
2069            status,
2070            duration,
2071        );
2072
2073        // Record OpenTelemetry span status
2074        if let Some(ref mut span) = ctx.otel_span {
2075            span.set_status(status);
2076            if let Some(ref upstream) = ctx.upstream {
2077                span.set_upstream(upstream, "");
2078            }
2079            if status >= 500 {
2080                span.record_error(&format!("HTTP {}", status));
2081            }
2082        }
2083
2084        // Record passive health check
2085        if let Some(ref upstream) = ctx.upstream {
2086            let success = status < 500;
2087
2088            trace!(
2089                correlation_id = %ctx.trace_id,
2090                upstream = %upstream,
2091                success = success,
2092                status = status,
2093                "Recording passive health check result"
2094            );
2095
2096            let error_msg = if !success {
2097                Some(format!("HTTP {}", status))
2098            } else {
2099                None
2100            };
2101            self.passive_health
2102                .record_outcome(upstream, success, error_msg.as_deref())
2103                .await;
2104
2105            // Report to upstream pool
2106            if let Some(pool) = self.upstream_pools.get(upstream).await {
2107                pool.report_result(upstream, success).await;
2108            }
2109
2110            if !success {
2111                warn!(
2112                    correlation_id = %ctx.trace_id,
2113                    upstream = %upstream,
2114                    status = status,
2115                    "Upstream returned error status"
2116                );
2117            }
2118        }
2119
2120        // Final request completion log
2121        if status >= 500 {
2122            error!(
2123                correlation_id = %ctx.trace_id,
2124                route_id = ctx.route_id.as_deref().unwrap_or("unknown"),
2125                upstream = ctx.upstream.as_deref().unwrap_or("none"),
2126                method = %ctx.method,
2127                path = %ctx.path,
2128                status = status,
2129                duration_ms = duration.as_millis(),
2130                attempts = ctx.upstream_attempts,
2131                "Request completed with server error"
2132            );
2133            self.log_manager.log_request_error(
2134                "error",
2135                "Request completed with server error",
2136                &ctx.trace_id,
2137                ctx.route_id.as_deref(),
2138                ctx.upstream.as_deref(),
2139                Some(format!(
2140                    "status={} method={} path={} duration_ms={}",
2141                    status,
2142                    ctx.method,
2143                    ctx.path,
2144                    duration.as_millis()
2145                )),
2146            );
2147        } else if status >= 400 {
2148            warn!(
2149                correlation_id = %ctx.trace_id,
2150                route_id = ctx.route_id.as_deref().unwrap_or("unknown"),
2151                upstream = ctx.upstream.as_deref().unwrap_or("none"),
2152                method = %ctx.method,
2153                path = %ctx.path,
2154                status = status,
2155                duration_ms = duration.as_millis(),
2156                "Request completed with client error"
2157            );
2158            self.log_manager.log_request_error(
2159                "warn",
2160                "Request completed with client error",
2161                &ctx.trace_id,
2162                ctx.route_id.as_deref(),
2163                ctx.upstream.as_deref(),
2164                Some(format!(
2165                    "status={} method={} path={}",
2166                    status, ctx.method, ctx.path
2167                )),
2168            );
2169        } else {
2170            debug!(
2171                correlation_id = %ctx.trace_id,
2172                route_id = ctx.route_id.as_deref().unwrap_or("unknown"),
2173                upstream = ctx.upstream.as_deref().unwrap_or("none"),
2174                method = %ctx.method,
2175                path = %ctx.path,
2176                status = status,
2177                duration_ms = duration.as_millis(),
2178                attempts = ctx.upstream_attempts,
2179                "Request completed"
2180            );
2181        }
2182
2183        Ok(())
2184    }
2185
2186    /// Modify the request before sending to upstream.
2187    /// Used for header modifications, adding authentication, etc.
2188    async fn upstream_request_filter(
2189        &self,
2190        _session: &mut Session,
2191        upstream_request: &mut pingora::http::RequestHeader,
2192        ctx: &mut Self::CTX,
2193    ) -> Result<()>
2194    where
2195        Self::CTX: Send + Sync,
2196    {
2197        trace!(
2198            correlation_id = %ctx.trace_id,
2199            route_id = ctx.route_id.as_deref().unwrap_or("unknown"),
2200            "Applying upstream request modifications"
2201        );
2202
2203        // Add trace ID header for upstream correlation
2204        upstream_request
2205            .insert_header("X-Trace-Id", &ctx.trace_id)
2206            .ok();
2207
2208        // Add W3C traceparent header for distributed tracing
2209        if let Some(ref span) = ctx.otel_span {
2210            let sampled = ctx
2211                .trace_context
2212                .as_ref()
2213                .map(|c| c.sampled)
2214                .unwrap_or(true);
2215            let traceparent =
2216                crate::otel::create_traceparent(&span.trace_id, &span.span_id, sampled);
2217            upstream_request
2218                .insert_header(crate::otel::TRACEPARENT_HEADER, &traceparent)
2219                .ok();
2220        }
2221
2222        // Add request metadata headers
2223        upstream_request
2224            .insert_header("X-Forwarded-By", "Zentinel")
2225            .ok();
2226
2227        // Apply route-specific request header modifications
2228        // Note: Pingora's IntoCaseHeaderName requires owned String for header names,
2229        // so we clone names but pass values by reference to avoid cloning both.
2230        if let Some(ref route_config) = ctx.route_config {
2231            let mods = &route_config.policies.request_headers;
2232
2233            // Rename runs before set/add/remove
2234            for (old_name, new_name) in &mods.rename {
2235                if let Some(value) = upstream_request
2236                    .headers
2237                    .get(old_name)
2238                    .and_then(|v| v.to_str().ok())
2239                {
2240                    let owned = value.to_string();
2241                    upstream_request
2242                        .insert_header(new_name.clone(), &owned)
2243                        .ok();
2244                    upstream_request.remove_header(old_name);
2245                }
2246            }
2247
2248            // Set headers (overwrite existing)
2249            for (name, value) in &mods.set {
2250                upstream_request
2251                    .insert_header(name.clone(), value.as_str())
2252                    .ok();
2253            }
2254
2255            // Add headers (append)
2256            for (name, value) in &mods.add {
2257                upstream_request
2258                    .append_header(name.clone(), value.as_str())
2259                    .ok();
2260            }
2261
2262            // Remove headers
2263            for name in &mods.remove {
2264                upstream_request.remove_header(name);
2265            }
2266
2267            trace!(
2268                correlation_id = %ctx.trace_id,
2269                "Applied request header modifications"
2270            );
2271        }
2272
2273        // Apply request-phase Headers filters
2274        if let Some(ref config) = ctx.config {
2275            super::filters::apply_request_headers_filters(upstream_request, ctx, config);
2276        }
2277
2278        // Remove sensitive headers that shouldn't go to upstream
2279        upstream_request.remove_header("X-Internal-Token");
2280        upstream_request.remove_header("Authorization-Internal");
2281
2282        // === Traffic Mirroring / Shadowing ===
2283        // Check if this route has shadow configuration
2284        if let Some(ref route_config) = ctx.route_config {
2285            if let Some(ref shadow_config) = route_config.shadow {
2286                // Get snapshot of upstream pools for shadow manager
2287                let pools_snapshot = self.upstream_pools.snapshot().await;
2288                let upstream_pools = std::sync::Arc::new(pools_snapshot);
2289
2290                // Get route ID for metrics labeling
2291                let route_id = ctx
2292                    .route_id
2293                    .clone()
2294                    .unwrap_or_else(|| "unknown".to_string());
2295
2296                // Create shadow manager
2297                let shadow_manager = crate::shadow::ShadowManager::new(
2298                    upstream_pools,
2299                    shadow_config.clone(),
2300                    Some(std::sync::Arc::clone(&self.metrics)),
2301                    route_id,
2302                );
2303
2304                // Check if we should shadow this request (sampling + header check)
2305                if shadow_manager.should_shadow(upstream_request) {
2306                    trace!(
2307                        correlation_id = %ctx.trace_id,
2308                        shadow_upstream = %shadow_config.upstream,
2309                        percentage = shadow_config.percentage,
2310                        "Shadowing request"
2311                    );
2312
2313                    // Clone headers for shadow request
2314                    let shadow_headers = upstream_request.clone();
2315
2316                    // Create request context for shadow (simplified from proxy context)
2317                    let shadow_ctx = crate::upstream::RequestContext {
2318                        client_ip: ctx.client_ip.parse().ok(),
2319                        headers: std::collections::HashMap::new(), // Empty for now
2320                        path: ctx.path.clone(),
2321                        method: ctx.method.clone(),
2322                    };
2323
2324                    // Determine if we should buffer the body
2325                    let buffer_body = shadow_config.buffer_body
2326                        && crate::shadow::should_buffer_method(&ctx.method);
2327
2328                    if buffer_body {
2329                        // Body buffering requested - defer shadow request until body is available
2330                        // Store shadow info in context; will be fired in logging phase
2331                        // or when request body filter completes
2332                        trace!(
2333                            correlation_id = %ctx.trace_id,
2334                            "Deferring shadow request until body is buffered"
2335                        );
2336                        ctx.shadow_pending = Some(crate::proxy::context::ShadowPendingRequest {
2337                            headers: shadow_headers,
2338                            manager: std::sync::Arc::new(shadow_manager),
2339                            request_ctx: shadow_ctx,
2340                            include_body: true,
2341                        });
2342                        // Enable body inspection to capture the body for shadow
2343                        // (only if not already enabled for other reasons)
2344                        if !ctx.body_inspection_enabled {
2345                            ctx.body_inspection_enabled = true;
2346                            // Set a reasonable buffer limit from shadow config
2347                            // (body_buffer will accumulate chunks)
2348                        }
2349                    } else {
2350                        // No body buffering needed - fire shadow request immediately
2351                        shadow_manager.shadow_request(shadow_headers, None, shadow_ctx);
2352                        ctx.shadow_sent = true;
2353                    }
2354                }
2355            }
2356        }
2357
2358        Ok(())
2359    }
2360
2361    /// Process response body chunks from upstream.
2362    /// Used for response size tracking and WAF inspection.
2363    ///
2364    /// Note: Response body inspection is currently buffered only (streaming mode not supported
2365    /// for responses due to Pingora's synchronous filter design).
2366    fn response_body_filter(
2367        &self,
2368        _session: &mut Session,
2369        body: &mut Option<Bytes>,
2370        end_of_stream: bool,
2371        ctx: &mut Self::CTX,
2372    ) -> Result<Option<Duration>, Box<Error>> {
2373        // Handle WebSocket frame inspection (server -> client)
2374        // Note: This filter is synchronous, so we use block_in_place for async agent calls
2375        if ctx.is_websocket_upgrade {
2376            if let Some(ref handler) = ctx.websocket_handler {
2377                let handler = handler.clone();
2378                let data = body.take();
2379
2380                // Use block_in_place to run async handler from sync context
2381                // This is safe because Pingora uses a multi-threaded tokio runtime
2382                let result = tokio::task::block_in_place(|| {
2383                    tokio::runtime::Handle::current()
2384                        .block_on(async { handler.process_server_data(data).await })
2385                });
2386
2387                match result {
2388                    crate::websocket::ProcessResult::Forward(data) => {
2389                        *body = data;
2390                    }
2391                    crate::websocket::ProcessResult::Close(reason) => {
2392                        warn!(
2393                            correlation_id = %ctx.trace_id,
2394                            code = reason.code,
2395                            reason = %reason.reason,
2396                            "WebSocket connection closed by agent (server->client)"
2397                        );
2398                        // For sync filter, we can't return an error that closes the connection
2399                        // Instead, inject a close frame
2400                        let close_frame =
2401                            crate::websocket::WebSocketFrame::close(reason.code, &reason.reason);
2402                        let codec = crate::websocket::WebSocketCodec::new(1024 * 1024);
2403                        if let Ok(encoded) = codec.encode_frame(&close_frame, false) {
2404                            *body = Some(Bytes::from(encoded));
2405                        }
2406                    }
2407                }
2408            }
2409            // Skip normal body processing for WebSocket
2410            return Ok(None);
2411        }
2412
2413        // Process response body through agents (for agents that subscribe to ResponseBody events)
2414        if ctx.response_agent_processing_enabled && !ctx.route_agent_ids.is_empty() {
2415            if let Some(ref chunk) = body {
2416                ctx.response_agent_body_buffer.extend_from_slice(chunk);
2417            }
2418
2419            if end_of_stream {
2420                let agent_ids = ctx.route_agent_ids.clone();
2421                let buffer = std::mem::take(&mut ctx.response_agent_body_buffer);
2422                let chunk_index = 0u32;
2423                let total_size = Some(buffer.len());
2424                let trace_id = ctx.trace_id.clone();
2425                let client_ip = ctx.client_ip.clone();
2426                let host = ctx.host.clone();
2427                let route_id = ctx.route_id.clone();
2428                let upstream_id = ctx.upstream.clone();
2429                let traceparent = ctx.traceparent();
2430                let agent_mgr = self.agent_manager.clone();
2431
2432                // Use block_in_place to run async agent call from sync context
2433                // This is safe because Pingora uses a multi-threaded tokio runtime
2434                let result = tokio::task::block_in_place(|| {
2435                    tokio::runtime::Handle::current().block_on(async {
2436                        let agent_ctx = crate::agents::AgentCallContext {
2437                            correlation_id: zentinel_common::CorrelationId::from_string(&trace_id),
2438                            metadata: zentinel_agent_protocol::RequestMetadata {
2439                                correlation_id: trace_id.clone(),
2440                                request_id: uuid::Uuid::new_v4().to_string(),
2441                                client_ip,
2442                                client_port: 0,
2443                                server_name: host,
2444                                protocol: "HTTP/1.1".to_string(),
2445                                tls_version: None,
2446                                tls_cipher: None,
2447                                route_id: route_id.clone(),
2448                                upstream_id: upstream_id.clone(),
2449                                timestamp: chrono::Utc::now().to_rfc3339(),
2450                                traceparent,
2451                            },
2452                            route_id,
2453                            upstream_id,
2454                            request_body: None,
2455                            response_body: None,
2456                        };
2457
2458                        agent_mgr
2459                            .process_response_body_streaming(
2460                                &agent_ctx,
2461                                &buffer,
2462                                true, // is_last
2463                                chunk_index,
2464                                buffer.len(),
2465                                total_size,
2466                                &agent_ids,
2467                            )
2468                            .await
2469                    })
2470                });
2471
2472                match result {
2473                    Ok(decision) => {
2474                        // Apply response body mutation if present
2475                        if let Some(mutation) = decision.response_body_mutation {
2476                            if let Some(ref data) = mutation.data {
2477                                if !data.is_empty() {
2478                                    // Decode base64 body from agent
2479                                    if let Ok(decoded) = base64::Engine::decode(
2480                                        &base64::engine::general_purpose::STANDARD,
2481                                        data,
2482                                    ) {
2483                                        debug!(
2484                                            correlation_id = %ctx.trace_id,
2485                                            original_size = buffer.len(),
2486                                            new_size = decoded.len(),
2487                                            "Agent replaced response body"
2488                                        );
2489                                        *body = Some(Bytes::from(decoded));
2490                                        ctx.response_agent_body_complete = true;
2491                                    } else {
2492                                        warn!(
2493                                            correlation_id = %ctx.trace_id,
2494                                            "Failed to decode agent response body mutation (invalid base64)"
2495                                        );
2496                                    }
2497                                }
2498                                // Empty data means drop the chunk — leave body as-is
2499                            }
2500                            // None data means pass through unchanged
2501                        }
2502
2503                        // Apply any additional response header modifications
2504                        // Note: Cannot modify response headers here (sync context, headers already sent)
2505                        // Header mods should be done in response_filter via ResponseHeaders event
2506                    }
2507                    Err(e) => {
2508                        warn!(
2509                            correlation_id = %ctx.trace_id,
2510                            error = %e,
2511                            "Agent response body processing failed, passing through original"
2512                        );
2513                    }
2514                }
2515            } else if !end_of_stream {
2516                // Buffer chunks — suppress output until we have the full body
2517                *body = None;
2518                return Ok(None);
2519            }
2520        }
2521
2522        // Track response body size
2523        if let Some(ref chunk) = body {
2524            ctx.response_bytes += chunk.len() as u64;
2525
2526            trace!(
2527                correlation_id = %ctx.trace_id,
2528                chunk_size = chunk.len(),
2529                total_response_bytes = ctx.response_bytes,
2530                end_of_stream = end_of_stream,
2531                "Processing response body chunk"
2532            );
2533
2534            // Process SSE chunks for streaming token counting
2535            if let Some(ref mut counter) = ctx.inference_streaming_counter {
2536                let result = counter.process_chunk(chunk);
2537
2538                if result.content.is_some() || result.is_done {
2539                    trace!(
2540                        correlation_id = %ctx.trace_id,
2541                        has_content = result.content.is_some(),
2542                        is_done = result.is_done,
2543                        chunks_processed = counter.chunks_processed(),
2544                        accumulated_content_len = counter.content().len(),
2545                        "Processed SSE chunk for token counting"
2546                    );
2547                }
2548            }
2549
2550            // Response body inspection (buffered mode only)
2551            // Note: Streaming mode for response bodies is not currently supported
2552            // due to Pingora's synchronous response_body_filter design
2553            if ctx.response_body_inspection_enabled
2554                && !ctx.response_body_inspection_agents.is_empty()
2555            {
2556                let config = ctx
2557                    .config
2558                    .get_or_insert_with(|| self.config_manager.current());
2559                let max_inspection_bytes = config
2560                    .waf
2561                    .as_ref()
2562                    .map(|w| w.body_inspection.max_inspection_bytes as u64)
2563                    .unwrap_or(1024 * 1024);
2564
2565                if ctx.response_body_bytes_inspected < max_inspection_bytes {
2566                    let bytes_to_inspect = std::cmp::min(
2567                        chunk.len() as u64,
2568                        max_inspection_bytes - ctx.response_body_bytes_inspected,
2569                    ) as usize;
2570
2571                    // Buffer for later processing (during logging phase)
2572                    // Response body inspection happens asynchronously and results
2573                    // are logged rather than blocking the response
2574                    ctx.response_body_bytes_inspected += bytes_to_inspect as u64;
2575                    ctx.response_body_chunk_index += 1;
2576
2577                    trace!(
2578                        correlation_id = %ctx.trace_id,
2579                        bytes_inspected = ctx.response_body_bytes_inspected,
2580                        max_inspection_bytes = max_inspection_bytes,
2581                        chunk_index = ctx.response_body_chunk_index,
2582                        "Tracking response body for inspection"
2583                    );
2584                }
2585            }
2586        }
2587
2588        if end_of_stream {
2589            trace!(
2590                correlation_id = %ctx.trace_id,
2591                total_response_bytes = ctx.response_bytes,
2592                response_bytes_inspected = ctx.response_body_bytes_inspected,
2593                "Response body complete"
2594            );
2595        }
2596
2597        // Return None to indicate no delay needed
2598        Ok(None)
2599    }
2600
2601    /// Called when a connection to upstream is established or reused.
2602    /// Logs connection reuse statistics for observability.
2603    async fn connected_to_upstream(
2604        &self,
2605        _session: &mut Session,
2606        reused: bool,
2607        peer: &HttpPeer,
2608        #[cfg(unix)] _fd: RawFd,
2609        #[cfg(windows)] _sock: std::os::windows::io::RawSocket,
2610        digest: Option<&Digest>,
2611        ctx: &mut Self::CTX,
2612    ) -> Result<(), Box<Error>> {
2613        // Track connection reuse for metrics
2614        ctx.connection_reused = reused;
2615
2616        // Log connection establishment/reuse
2617        if reused {
2618            trace!(
2619                correlation_id = %ctx.trace_id,
2620                upstream = ctx.upstream.as_deref().unwrap_or("unknown"),
2621                peer_address = %peer.address(),
2622                "Reusing existing upstream connection"
2623            );
2624        } else {
2625            debug!(
2626                correlation_id = %ctx.trace_id,
2627                upstream = ctx.upstream.as_deref().unwrap_or("unknown"),
2628                peer_address = %peer.address(),
2629                ssl = digest.as_ref().map(|d| d.ssl_digest.is_some()).unwrap_or(false),
2630                "Established new upstream connection"
2631            );
2632        }
2633
2634        Ok(())
2635    }
2636
2637    // =========================================================================
2638    // HTTP Caching - Pingora Cache Integration
2639    // =========================================================================
2640
2641    /// Decide if the request should use caching.
2642    ///
2643    /// This method is called early in the request lifecycle to determine if
2644    /// the response should be served from cache or if the response should
2645    /// be cached.
2646    fn request_cache_filter(&self, session: &mut Session, ctx: &mut Self::CTX) -> Result<()> {
2647        // Check if route has caching enabled
2648        let route_id = match ctx.route_id.as_deref() {
2649            Some(id) => id,
2650            None => {
2651                trace!(
2652                    correlation_id = %ctx.trace_id,
2653                    "Cache filter: no route ID, skipping cache"
2654                );
2655                return Ok(());
2656            }
2657        };
2658
2659        // Check if caching is enabled for this route
2660        if !self.cache_manager.is_enabled(route_id) {
2661            ctx.cache_status = Some(super::context::CacheStatus::Bypass("disabled"));
2662            trace!(
2663                correlation_id = %ctx.trace_id,
2664                route_id = %route_id,
2665                "Cache disabled for route"
2666            );
2667            return Ok(());
2668        }
2669
2670        // Check if method is cacheable (typically GET/HEAD)
2671        if !self
2672            .cache_manager
2673            .is_method_cacheable(route_id, &ctx.method)
2674        {
2675            ctx.cache_status = Some(super::context::CacheStatus::Bypass("method"));
2676            trace!(
2677                correlation_id = %ctx.trace_id,
2678                route_id = %route_id,
2679                method = %ctx.method,
2680                "Method not cacheable"
2681            );
2682            return Ok(());
2683        }
2684
2685        // Check if path is excluded from caching (by extension or pattern)
2686        if !self.cache_manager.is_path_cacheable(route_id, &ctx.path) {
2687            ctx.cache_status = Some(super::context::CacheStatus::Bypass("excluded"));
2688            trace!(
2689                correlation_id = %ctx.trace_id,
2690                route_id = %route_id,
2691                path = %ctx.path,
2692                "Path excluded from caching"
2693            );
2694            return Ok(());
2695        }
2696
2697        // Enable caching for this request using Pingora's cache infrastructure
2698        debug!(
2699            correlation_id = %ctx.trace_id,
2700            route_id = %route_id,
2701            method = %ctx.method,
2702            path = %ctx.path,
2703            "Enabling HTTP caching for request"
2704        );
2705
2706        // Get static references to cache infrastructure
2707        let storage = get_cache_storage();
2708        let eviction = get_cache_eviction();
2709        let cache_lock = get_cache_lock();
2710
2711        // Enable the cache with storage, eviction, and lock
2712        session.cache.enable(
2713            storage,
2714            Some(eviction),
2715            None, // predictor - optional
2716            Some(cache_lock),
2717            None, // option overrides
2718        );
2719
2720        // Mark request as cache-eligible in context
2721        ctx.cache_eligible = true;
2722
2723        trace!(
2724            correlation_id = %ctx.trace_id,
2725            route_id = %route_id,
2726            cache_enabled = session.cache.enabled(),
2727            "Cache enabled for request"
2728        );
2729
2730        Ok(())
2731    }
2732
2733    /// Generate the cache key for this request.
2734    ///
2735    /// The cache key uniquely identifies the cached response. It typically
2736    /// includes the method, host, path, and potentially query parameters.
2737    fn cache_key_callback(&self, session: &Session, ctx: &mut Self::CTX) -> Result<CacheKey> {
2738        let req_header = session.req_header();
2739        let method = req_header.method.as_str();
2740        let path = req_header.uri.path();
2741        let host = ctx.host.as_deref().unwrap_or("unknown");
2742        let query = req_header.uri.query();
2743
2744        // Generate cache key using our cache manager
2745        let key_string = crate::cache::CacheManager::generate_cache_key(method, host, path, query);
2746
2747        trace!(
2748            correlation_id = %ctx.trace_id,
2749            cache_key = %key_string,
2750            "Generated cache key"
2751        );
2752
2753        // Generate cache key from request URI (namespace empty, user_tag empty)
2754        // CacheKey::default() was removed in Pingora 0.8.0
2755        Ok(CacheKey::new("", format!("{}", req_header.uri), ""))
2756    }
2757
2758    /// Called when a cache miss occurs.
2759    ///
2760    /// This is called when the cache lookup found no matching entry.
2761    /// We can use this to log and track cache misses.
2762    fn cache_miss(&self, session: &mut Session, ctx: &mut Self::CTX) {
2763        // Let Pingora handle the cache miss
2764        session.cache.cache_miss();
2765
2766        ctx.cache_status = Some(super::context::CacheStatus::Miss);
2767
2768        // Track statistics
2769        if let Some(route_id) = ctx.route_id.as_deref() {
2770            self.cache_manager.stats().record_miss();
2771
2772            trace!(
2773                correlation_id = %ctx.trace_id,
2774                route_id = %route_id,
2775                path = %ctx.path,
2776                "Cache miss"
2777            );
2778        }
2779    }
2780
2781    /// Called after a successful cache lookup.
2782    ///
2783    /// This filter allows inspecting the cached response before serving it.
2784    /// Returns `None` to serve the cached response, or a `ForcedFreshness`
2785    /// to override the freshness decision.
2786    async fn cache_hit_filter(
2787        &self,
2788        session: &mut Session,
2789        meta: &CacheMeta,
2790        hit_handler: &mut HitHandler,
2791        is_fresh: bool,
2792        ctx: &mut Self::CTX,
2793    ) -> Result<Option<ForcedFreshness>>
2794    where
2795        Self::CTX: Send + Sync,
2796    {
2797        // Check if this cache entry should be invalidated due to a purge request
2798        let req_header = session.req_header();
2799        let method = req_header.method.as_str();
2800        let path = req_header.uri.path();
2801        let host = req_header.uri.host().unwrap_or("localhost");
2802        let query = req_header.uri.query();
2803
2804        // Generate the cache key for this request
2805        let cache_key = crate::cache::CacheManager::generate_cache_key(method, host, path, query);
2806
2807        // Check if this key should be invalidated
2808        if self.cache_manager.should_invalidate(&cache_key) {
2809            info!(
2810                correlation_id = %ctx.trace_id,
2811                route_id = ctx.route_id.as_deref().unwrap_or("unknown"),
2812                cache_key = %cache_key,
2813                "Cache entry invalidated by purge request"
2814            );
2815            // Force expiration so the entry is refetched from upstream
2816            return Ok(Some(ForcedFreshness::ForceExpired));
2817        }
2818
2819        // Track cache hit statistics
2820        if is_fresh {
2821            // Detect which tier served the hit via downcast
2822            let is_disk_hit = hit_handler
2823                .as_any()
2824                .downcast_ref::<HybridHitHandler>()
2825                .is_some()
2826                || hit_handler
2827                    .as_any()
2828                    .downcast_ref::<DiskHitHandler>()
2829                    .is_some();
2830
2831            let stats = self.cache_manager.stats();
2832            if is_disk_hit {
2833                ctx.cache_status = Some(super::context::CacheStatus::HitDisk);
2834                stats.record_disk_hit();
2835            } else {
2836                ctx.cache_status = Some(super::context::CacheStatus::HitMemory);
2837                stats.record_memory_hit();
2838            }
2839            stats.record_hit();
2840
2841            debug!(
2842                correlation_id = %ctx.trace_id,
2843                route_id = ctx.route_id.as_deref().unwrap_or("unknown"),
2844                is_fresh = is_fresh,
2845                tier = if is_disk_hit { "disk" } else { "memory" },
2846                "Cache hit (fresh)"
2847            );
2848        } else {
2849            ctx.cache_status = Some(super::context::CacheStatus::HitStale);
2850
2851            trace!(
2852                correlation_id = %ctx.trace_id,
2853                route_id = ctx.route_id.as_deref().unwrap_or("unknown"),
2854                is_fresh = is_fresh,
2855                "Cache hit (stale)"
2856            );
2857        }
2858
2859        // Serve the cached response without invalidation
2860        Ok(None)
2861    }
2862
2863    /// Decide if the response should be cached.
2864    ///
2865    /// Called after receiving the response from upstream to determine
2866    /// if it should be stored in the cache.
2867    fn response_cache_filter(
2868        &self,
2869        _session: &Session,
2870        resp: &ResponseHeader,
2871        ctx: &mut Self::CTX,
2872    ) -> Result<RespCacheable> {
2873        let route_id = match ctx.route_id.as_deref() {
2874            Some(id) => id,
2875            None => {
2876                return Ok(RespCacheable::Uncacheable(NoCacheReason::Custom(
2877                    "no_route",
2878                )));
2879            }
2880        };
2881
2882        // Check if caching is enabled for this route
2883        if !self.cache_manager.is_enabled(route_id) {
2884            return Ok(RespCacheable::Uncacheable(NoCacheReason::Custom(
2885                "disabled",
2886            )));
2887        }
2888
2889        let status = resp.status.as_u16();
2890
2891        // Check if status code is cacheable
2892        if !self.cache_manager.is_status_cacheable(route_id, status) {
2893            trace!(
2894                correlation_id = %ctx.trace_id,
2895                route_id = %route_id,
2896                status = status,
2897                "Status code not cacheable"
2898            );
2899            return Ok(RespCacheable::Uncacheable(NoCacheReason::OriginNotCache));
2900        }
2901
2902        // Check Cache-Control header for no-store, no-cache, private
2903        if let Some(cache_control) = resp.headers.get("cache-control") {
2904            if let Ok(cc_str) = cache_control.to_str() {
2905                if crate::cache::CacheManager::is_no_cache(cc_str) {
2906                    trace!(
2907                        correlation_id = %ctx.trace_id,
2908                        route_id = %route_id,
2909                        cache_control = %cc_str,
2910                        "Response has no-cache directive"
2911                    );
2912                    return Ok(RespCacheable::Uncacheable(NoCacheReason::OriginNotCache));
2913                }
2914            }
2915        }
2916
2917        // Calculate TTL from Cache-Control or use default
2918        let cache_control = resp
2919            .headers
2920            .get("cache-control")
2921            .and_then(|v| v.to_str().ok());
2922        let ttl = self.cache_manager.calculate_ttl(route_id, cache_control);
2923
2924        if ttl.is_zero() {
2925            trace!(
2926                correlation_id = %ctx.trace_id,
2927                route_id = %route_id,
2928                "TTL is zero, not caching"
2929            );
2930            return Ok(RespCacheable::Uncacheable(NoCacheReason::OriginNotCache));
2931        }
2932
2933        // Get route cache config for stale settings
2934        let config = self
2935            .cache_manager
2936            .get_route_config(route_id)
2937            .unwrap_or_default();
2938
2939        // Create timestamps for cache metadata
2940        let now = std::time::SystemTime::now();
2941        let fresh_until = now + ttl;
2942
2943        // Clone the response header for storage
2944        let header = resp.clone();
2945
2946        // Create CacheMeta with proper timestamps and TTLs
2947        let cache_meta = CacheMeta::new(
2948            fresh_until,
2949            now,
2950            config.stale_while_revalidate_secs as u32,
2951            config.stale_if_error_secs as u32,
2952            header,
2953        );
2954
2955        // Track the cache store
2956        self.cache_manager.stats().record_store();
2957
2958        debug!(
2959            correlation_id = %ctx.trace_id,
2960            route_id = %route_id,
2961            status = status,
2962            ttl_secs = ttl.as_secs(),
2963            stale_while_revalidate_secs = config.stale_while_revalidate_secs,
2964            stale_if_error_secs = config.stale_if_error_secs,
2965            "Caching response"
2966        );
2967
2968        Ok(RespCacheable::Cacheable(cache_meta))
2969    }
2970
2971    /// Decide whether to serve stale content on error or during revalidation.
2972    ///
2973    /// This implements stale-while-revalidate and stale-if-error semantics.
2974    fn should_serve_stale(
2975        &self,
2976        _session: &mut Session,
2977        ctx: &mut Self::CTX,
2978        error: Option<&Error>,
2979    ) -> bool {
2980        let route_id = match ctx.route_id.as_deref() {
2981            Some(id) => id,
2982            None => return false,
2983        };
2984
2985        // Get route cache config for stale settings
2986        let config = match self.cache_manager.get_route_config(route_id) {
2987            Some(c) => c,
2988            None => return false,
2989        };
2990
2991        // If there's an upstream error, use stale-if-error
2992        if let Some(e) = error {
2993            // Only serve stale for upstream errors
2994            if e.esource() == &pingora::ErrorSource::Upstream {
2995                debug!(
2996                    correlation_id = %ctx.trace_id,
2997                    route_id = %route_id,
2998                    error = %e,
2999                    stale_if_error_secs = config.stale_if_error_secs,
3000                    "Considering stale-if-error"
3001                );
3002                return config.stale_if_error_secs > 0;
3003            }
3004        }
3005
3006        // During stale-while-revalidate (error is None)
3007        if error.is_none() && config.stale_while_revalidate_secs > 0 {
3008            trace!(
3009                correlation_id = %ctx.trace_id,
3010                route_id = %route_id,
3011                stale_while_revalidate_secs = config.stale_while_revalidate_secs,
3012                "Allowing stale-while-revalidate"
3013            );
3014            return true;
3015        }
3016
3017        false
3018    }
3019
3020    /// Handle Range header for byte-range requests (streaming support).
3021    ///
3022    /// This method is called when a Range header is present in the request.
3023    /// It allows proper handling of:
3024    /// - Video streaming (HTML5 video seeking)
3025    /// - Large file downloads with resume support
3026    /// - Partial content delivery
3027    ///
3028    /// Uses Pingora's built-in range handling with route-specific logging.
3029    fn range_header_filter(
3030        &self,
3031        session: &mut Session,
3032        response: &mut ResponseHeader,
3033        ctx: &mut Self::CTX,
3034    ) -> pingora_proxy::RangeType
3035    where
3036        Self::CTX: Send + Sync,
3037    {
3038        // Check if route supports range requests
3039        let supports_range = ctx.route_config.as_ref().is_none_or(|config| {
3040            // Static file routes and media routes should support range requests
3041            matches!(
3042                config.service_type,
3043                zentinel_config::ServiceType::Static | zentinel_config::ServiceType::Web
3044            )
3045        });
3046
3047        if !supports_range {
3048            trace!(
3049                correlation_id = %ctx.trace_id,
3050                route_id = ctx.route_id.as_deref().unwrap_or("unknown"),
3051                "Range request not supported for this route type"
3052            );
3053            return pingora_proxy::RangeType::None;
3054        }
3055
3056        // Use Pingora's built-in range header parsing and handling
3057        let range_type = pingora_proxy::range_header_filter(session.req_header(), response, None);
3058
3059        match &range_type {
3060            pingora_proxy::RangeType::None => {
3061                trace!(
3062                    correlation_id = %ctx.trace_id,
3063                    route_id = ctx.route_id.as_deref().unwrap_or("unknown"),
3064                    "No range request or not applicable"
3065                );
3066            }
3067            pingora_proxy::RangeType::Single(range) => {
3068                trace!(
3069                    correlation_id = %ctx.trace_id,
3070                    route_id = ctx.route_id.as_deref().unwrap_or("unknown"),
3071                    range_start = range.start,
3072                    range_end = range.end,
3073                    "Processing single-range request"
3074                );
3075            }
3076            pingora_proxy::RangeType::Multi(multi) => {
3077                trace!(
3078                    correlation_id = %ctx.trace_id,
3079                    route_id = ctx.route_id.as_deref().unwrap_or("unknown"),
3080                    range_count = multi.ranges.len(),
3081                    "Processing multi-range request"
3082                );
3083            }
3084            pingora_proxy::RangeType::Invalid => {
3085                debug!(
3086                    correlation_id = %ctx.trace_id,
3087                    route_id = ctx.route_id.as_deref().unwrap_or("unknown"),
3088                    "Invalid range header"
3089                );
3090            }
3091        }
3092
3093        range_type
3094    }
3095
3096    /// Handle fatal proxy errors by generating custom error pages.
3097    /// Called when the proxy itself fails to process the request.
3098    async fn fail_to_proxy(
3099        &self,
3100        session: &mut Session,
3101        e: &Error,
3102        ctx: &mut Self::CTX,
3103    ) -> pingora_proxy::FailToProxy
3104    where
3105        Self::CTX: Send + Sync,
3106    {
3107        let error_code = match e.etype() {
3108            // Connection errors
3109            ErrorType::ConnectRefused => 503,
3110            ErrorType::ConnectTimedout => 504,
3111            ErrorType::ConnectNoRoute => 502,
3112
3113            // Timeout errors
3114            ErrorType::ReadTimedout => 504,
3115            ErrorType::WriteTimedout => 504,
3116
3117            // TLS errors
3118            ErrorType::TLSHandshakeFailure => 502,
3119            ErrorType::InvalidCert => 502,
3120
3121            // Protocol errors
3122            ErrorType::InvalidHTTPHeader => 400,
3123            ErrorType::H2Error => 502,
3124
3125            // Resource errors
3126            ErrorType::ConnectProxyFailure => 502,
3127            ErrorType::ConnectionClosed => 502,
3128
3129            // Explicit HTTP status (e.g., from agent fail-closed blocking)
3130            ErrorType::HTTPStatus(status) => *status,
3131
3132            // Internal errors indicate server-side configuration issues (e.g.,
3133            // invalid/missing backends). Gateway API spec requires 500 for these.
3134            ErrorType::InternalError => 500,
3135
3136            // Default to 502 for unknown errors
3137            _ => 502,
3138        };
3139
3140        error!(
3141            correlation_id = %ctx.trace_id,
3142            route_id = ctx.route_id.as_deref().unwrap_or("unknown"),
3143            upstream = ctx.upstream.as_deref().unwrap_or("unknown"),
3144            error_type = ?e.etype(),
3145            error = %e,
3146            error_code = error_code,
3147            "Proxy error occurred"
3148        );
3149
3150        // Record the error in metrics
3151        self.metrics
3152            .record_blocked_request(&format!("proxy_error_{}", error_code));
3153
3154        // Write error response to ensure client receives a proper HTTP response
3155        // This is necessary because some errors occur before the upstream connection
3156        // is established, and Pingora may not send a response automatically
3157        let error_message = match error_code {
3158            400 => "Bad Request",
3159            502 => "Bad Gateway",
3160            503 => "Service Unavailable",
3161            504 => "Gateway Timeout",
3162            _ => "Internal Server Error",
3163        };
3164
3165        // Build a minimal error response body
3166        let body = format!(
3167            r#"{{"error":"{} {}","trace_id":"{}"}}"#,
3168            error_code, error_message, ctx.trace_id
3169        );
3170
3171        // Write the response header
3172        let mut header = pingora::http::ResponseHeader::build(error_code, None).unwrap();
3173        header
3174            .insert_header("Content-Type", "application/json")
3175            .ok();
3176        header
3177            .insert_header("Content-Length", body.len().to_string())
3178            .ok();
3179        header
3180            .insert_header("X-Correlation-Id", ctx.trace_id.as_str())
3181            .ok();
3182        header.insert_header("Connection", "close").ok();
3183
3184        // Write headers and body
3185        if let Err(write_err) = session.write_response_header(Box::new(header), false).await {
3186            warn!(
3187                correlation_id = %ctx.trace_id,
3188                error = %write_err,
3189                "Failed to write error response header"
3190            );
3191        } else {
3192            // Write the body
3193            if let Err(write_err) = session
3194                .write_response_body(Some(bytes::Bytes::from(body)), true)
3195                .await
3196            {
3197                warn!(
3198                    correlation_id = %ctx.trace_id,
3199                    error = %write_err,
3200                    "Failed to write error response body"
3201                );
3202            }
3203        }
3204
3205        // Return the error response info
3206        // can_reuse_downstream: false since we already wrote and closed the response
3207        pingora_proxy::FailToProxy {
3208            error_code,
3209            can_reuse_downstream: false,
3210        }
3211    }
3212
3213    /// Handle errors that occur during proxying after upstream connection is established.
3214    ///
3215    /// This method enables retry logic and circuit breaker integration.
3216    /// It's called when an error occurs during the request/response exchange
3217    /// with the upstream server.
3218    fn error_while_proxy(
3219        &self,
3220        peer: &HttpPeer,
3221        session: &mut Session,
3222        e: Box<Error>,
3223        ctx: &mut Self::CTX,
3224        client_reused: bool,
3225    ) -> Box<Error> {
3226        let error_type = e.etype().clone();
3227        let upstream_id = ctx.upstream.as_deref().unwrap_or("unknown");
3228
3229        // Classify error for retry decisions
3230        let is_retryable = matches!(
3231            error_type,
3232            ErrorType::ConnectTimedout
3233                | ErrorType::ReadTimedout
3234                | ErrorType::WriteTimedout
3235                | ErrorType::ConnectionClosed
3236                | ErrorType::ConnectRefused
3237        );
3238
3239        // Log the error with context
3240        warn!(
3241            correlation_id = %ctx.trace_id,
3242            route_id = ctx.route_id.as_deref().unwrap_or("unknown"),
3243            upstream = %upstream_id,
3244            peer_address = %peer.address(),
3245            error_type = ?error_type,
3246            error = %e,
3247            client_reused = client_reused,
3248            is_retryable = is_retryable,
3249            "Error during proxy operation"
3250        );
3251
3252        // Record failure with circuit breaker via upstream pool
3253        // This is done asynchronously since we can't await in a sync fn
3254        let peer_address = peer.address().to_string();
3255        let upstream_pools = self.upstream_pools.clone();
3256        let upstream_id_owned = upstream_id.to_string();
3257        tokio::spawn(async move {
3258            if let Some(pool) = upstream_pools.get(&upstream_id_owned).await {
3259                pool.report_result(&peer_address, false).await;
3260            }
3261        });
3262
3263        // Metrics tracking
3264        self.metrics
3265            .record_blocked_request(&format!("proxy_error_{:?}", error_type));
3266
3267        // Create enhanced error with retry information
3268        let mut enhanced_error = e.more_context(format!(
3269            "Upstream: {}, Peer: {}, Attempts: {}",
3270            upstream_id,
3271            peer.address(),
3272            ctx.upstream_attempts
3273        ));
3274
3275        // Determine if retry should be attempted:
3276        // - Only retry if error is retryable type
3277        // - Only retry reused connections if buffer isn't truncated
3278        // - Track retry metrics
3279        if is_retryable {
3280            let can_retry = if client_reused {
3281                // For reused connections, check if retry buffer is intact
3282                !session.as_ref().retry_buffer_truncated()
3283            } else {
3284                // Fresh connections can always retry
3285                true
3286            };
3287
3288            enhanced_error.retry.decide_reuse(can_retry);
3289
3290            if can_retry {
3291                debug!(
3292                    correlation_id = %ctx.trace_id,
3293                    upstream = %upstream_id,
3294                    error_type = ?error_type,
3295                    "Error is retryable, will attempt retry"
3296                );
3297            }
3298        } else {
3299            // Non-retryable error - don't retry
3300            enhanced_error.retry.decide_reuse(false);
3301        }
3302
3303        enhanced_error
3304    }
3305
3306    async fn logging(&self, session: &mut Session, _error: Option<&Error>, ctx: &mut Self::CTX) {
3307        // Decrement active requests
3308        self.reload_coordinator.dec_requests();
3309
3310        // === Fire pending shadow request (if body buffering was enabled) ===
3311        if !ctx.shadow_sent {
3312            if let Some(shadow_pending) = ctx.shadow_pending.take() {
3313                let body = if shadow_pending.include_body && !ctx.body_buffer.is_empty() {
3314                    // Clone the buffered body for the shadow request
3315                    Some(ctx.body_buffer.clone())
3316                } else {
3317                    None
3318                };
3319
3320                trace!(
3321                    correlation_id = %ctx.trace_id,
3322                    body_size = body.as_ref().map(|b| b.len()).unwrap_or(0),
3323                    "Firing deferred shadow request with buffered body"
3324                );
3325
3326                shadow_pending.manager.shadow_request(
3327                    shadow_pending.headers,
3328                    body,
3329                    shadow_pending.request_ctx,
3330                );
3331                ctx.shadow_sent = true;
3332            }
3333        }
3334
3335        let duration = ctx.elapsed();
3336
3337        // Get response status
3338        let status = session
3339            .response_written()
3340            .map(|r| r.status.as_u16())
3341            .unwrap_or(0);
3342
3343        // Report result to load balancer for adaptive LB feedback
3344        // This enables latency-aware weight adjustment
3345        if let (Some(ref peer_addr), Some(ref upstream_id)) =
3346            (&ctx.selected_upstream_address, &ctx.upstream)
3347        {
3348            // Success = status code < 500 (client errors are not upstream failures)
3349            let success = status > 0 && status < 500;
3350
3351            if let Some(pool) = self.upstream_pools.get(upstream_id).await {
3352                pool.report_result_with_latency(peer_addr, success, Some(duration))
3353                    .await;
3354                pool.decrement_active();
3355                trace!(
3356                    correlation_id = %ctx.trace_id,
3357                    upstream = %upstream_id,
3358                    peer_address = %peer_addr,
3359                    success = success,
3360                    duration_ms = duration.as_millis(),
3361                    status = status,
3362                    "Reported result to adaptive load balancer"
3363                );
3364            }
3365
3366            // Track warmth for inference routes (cold model detection)
3367            if ctx.inference_rate_limit_enabled && success {
3368                let cold_detected = self.warmth_tracker.record_request(peer_addr, duration);
3369                if cold_detected {
3370                    debug!(
3371                        correlation_id = %ctx.trace_id,
3372                        upstream = %upstream_id,
3373                        peer_address = %peer_addr,
3374                        duration_ms = duration.as_millis(),
3375                        "Cold model detected on inference upstream"
3376                    );
3377                }
3378            }
3379        }
3380
3381        // Record actual token usage for inference rate limiting
3382        // This adjusts the token bucket based on actual vs estimated tokens
3383        if ctx.inference_rate_limit_enabled {
3384            if let (Some(route_id), Some(ref rate_limit_key)) =
3385                (ctx.route_id.as_deref(), &ctx.inference_rate_limit_key)
3386            {
3387                // Try to extract actual tokens from response headers
3388                let response_headers = session
3389                    .response_written()
3390                    .map(|r| &r.headers)
3391                    .cloned()
3392                    .unwrap_or_default();
3393
3394                // For streaming responses, finalize the streaming token counter
3395                let streaming_result = if ctx.inference_streaming_response {
3396                    ctx.inference_streaming_counter
3397                        .as_ref()
3398                        .map(|counter| counter.finalize())
3399                } else {
3400                    None
3401                };
3402
3403                // Log streaming token count info
3404                if let Some(ref result) = streaming_result {
3405                    debug!(
3406                        correlation_id = %ctx.trace_id,
3407                        output_tokens = result.output_tokens,
3408                        input_tokens = ?result.input_tokens,
3409                        source = ?result.source,
3410                        content_length = result.content_length,
3411                        "Finalized streaming token count"
3412                    );
3413                }
3414
3415                // PII detection guardrail (for streaming inference responses)
3416                if ctx.inference_streaming_response {
3417                    if let Some(ref route_config) = ctx.route_config {
3418                        if let Some(ref inference) = route_config.inference {
3419                            if let Some(ref guardrails) = inference.guardrails {
3420                                if let Some(ref pii_config) = guardrails.pii_detection {
3421                                    if pii_config.enabled {
3422                                        // Get accumulated content from streaming counter
3423                                        if let Some(ref counter) = ctx.inference_streaming_counter {
3424                                            let response_content = counter.content();
3425                                            if !response_content.is_empty() {
3426                                                let pii_result = self
3427                                                    .guardrail_processor
3428                                                    .check_pii(
3429                                                        pii_config,
3430                                                        response_content,
3431                                                        ctx.route_id.as_deref(),
3432                                                        &ctx.trace_id,
3433                                                    )
3434                                                    .await;
3435
3436                                                match pii_result {
3437                                                    crate::inference::PiiCheckResult::Detected {
3438                                                        detections,
3439                                                        redacted_content: _,
3440                                                    } => {
3441                                                        warn!(
3442                                                            correlation_id = %ctx.trace_id,
3443                                                            route_id = ctx.route_id.as_deref().unwrap_or("unknown"),
3444                                                            detection_count = detections.len(),
3445                                                            "PII detected in inference response"
3446                                                        );
3447
3448                                                        // Store detection categories for logging
3449                                                        ctx.pii_detection_categories = detections
3450                                                            .iter()
3451                                                            .map(|d| d.category.clone())
3452                                                            .collect();
3453
3454                                                        // Record metrics for each category
3455                                                        for detection in &detections {
3456                                                            self.metrics.record_pii_detected(
3457                                                                ctx.route_id.as_deref().unwrap_or("unknown"),
3458                                                                &detection.category,
3459                                                            );
3460                                                        }
3461                                                    }
3462                                                    crate::inference::PiiCheckResult::Clean => {
3463                                                        trace!(
3464                                                            correlation_id = %ctx.trace_id,
3465                                                            "No PII detected in response"
3466                                                        );
3467                                                    }
3468                                                    crate::inference::PiiCheckResult::Error { message } => {
3469                                                        debug!(
3470                                                            correlation_id = %ctx.trace_id,
3471                                                            error = %message,
3472                                                            "PII detection check failed"
3473                                                        );
3474                                                    }
3475                                                }
3476                                            }
3477                                        }
3478                                    }
3479                                }
3480                            }
3481                        }
3482                    }
3483                }
3484
3485                // Response body would require buffering, which is expensive
3486                // For non-streaming, most LLM APIs provide token counts in headers
3487                // For streaming, we use the accumulated SSE content
3488                let empty_body: &[u8] = &[];
3489
3490                if let Some(actual_estimate) = self.inference_rate_limit_manager.record_actual(
3491                    route_id,
3492                    rate_limit_key,
3493                    &response_headers,
3494                    empty_body,
3495                    ctx.inference_estimated_tokens,
3496                ) {
3497                    // Use streaming result if available and header extraction failed
3498                    let (actual_tokens, source_info) = if let Some(ref streaming) = streaming_result
3499                    {
3500                        // Prefer API-provided counts from streaming, otherwise use tiktoken count
3501                        if let Some(total_tokens) = streaming.total_tokens {
3502                            (total_tokens, "streaming_api")
3503                        } else if actual_estimate.source == crate::inference::TokenSource::Estimated
3504                        {
3505                            // Header extraction failed, use streaming tiktoken count
3506                            // Estimate total by adding input estimate + output from streaming
3507                            let total = ctx.inference_input_tokens + streaming.output_tokens;
3508                            (total, "streaming_tiktoken")
3509                        } else {
3510                            (actual_estimate.tokens, "headers")
3511                        }
3512                    } else {
3513                        (actual_estimate.tokens, "headers")
3514                    };
3515
3516                    ctx.inference_actual_tokens = Some(actual_tokens);
3517
3518                    debug!(
3519                        correlation_id = %ctx.trace_id,
3520                        route_id = route_id,
3521                        estimated_tokens = ctx.inference_estimated_tokens,
3522                        actual_tokens = actual_tokens,
3523                        source = source_info,
3524                        streaming_response = ctx.inference_streaming_response,
3525                        model = ?ctx.inference_model,
3526                        "Recorded actual inference tokens"
3527                    );
3528
3529                    // Record budget usage with actual tokens (if budget tracking enabled)
3530                    if ctx.inference_budget_enabled {
3531                        let alerts = self.inference_rate_limit_manager.record_budget(
3532                            route_id,
3533                            rate_limit_key,
3534                            actual_tokens,
3535                        );
3536
3537                        // Log any budget alerts that fired
3538                        for alert in alerts.iter() {
3539                            warn!(
3540                                correlation_id = %ctx.trace_id,
3541                                route_id = route_id,
3542                                tenant = %alert.tenant,
3543                                threshold_pct = alert.threshold * 100.0,
3544                                tokens_used = alert.tokens_used,
3545                                tokens_limit = alert.tokens_limit,
3546                                "Token budget alert threshold crossed"
3547                            );
3548                        }
3549
3550                        // Update context with remaining budget
3551                        if let Some(status) = self
3552                            .inference_rate_limit_manager
3553                            .budget_status(route_id, rate_limit_key)
3554                        {
3555                            ctx.inference_budget_remaining = Some(status.tokens_remaining as i64);
3556                        }
3557                    }
3558
3559                    // Calculate cost if cost attribution is enabled
3560                    if ctx.inference_cost_enabled {
3561                        if let Some(model) = ctx.inference_model.as_deref() {
3562                            // Use streaming result for more accurate input/output split if available
3563                            let (input_tokens, output_tokens) = if let Some(ref streaming) =
3564                                streaming_result
3565                            {
3566                                // Streaming gives us accurate output tokens
3567                                let input =
3568                                    streaming.input_tokens.unwrap_or(ctx.inference_input_tokens);
3569                                let output = streaming.output_tokens;
3570                                (input, output)
3571                            } else {
3572                                // Fallback: estimate output from total - input
3573                                let input = ctx.inference_input_tokens;
3574                                let output = actual_tokens.saturating_sub(input);
3575                                (input, output)
3576                            };
3577                            ctx.inference_output_tokens = output_tokens;
3578
3579                            if let Some(cost_result) = self
3580                                .inference_rate_limit_manager
3581                                .calculate_cost(route_id, model, input_tokens, output_tokens)
3582                            {
3583                                ctx.inference_request_cost = Some(cost_result.total_cost);
3584
3585                                trace!(
3586                                    correlation_id = %ctx.trace_id,
3587                                    route_id = route_id,
3588                                    model = model,
3589                                    input_tokens = input_tokens,
3590                                    output_tokens = output_tokens,
3591                                    total_cost = cost_result.total_cost,
3592                                    currency = %cost_result.currency,
3593                                    "Calculated inference request cost"
3594                                );
3595                            }
3596                        }
3597                    }
3598                }
3599            }
3600        }
3601
3602        // Write to access log file if configured (check sampling before allocating entry)
3603        if self.log_manager.should_log_access(status) {
3604            let access_entry = AccessLogEntry {
3605                timestamp: chrono::Utc::now().to_rfc3339(),
3606                trace_id: ctx.trace_id.clone(),
3607                method: ctx.method.clone(),
3608                path: ctx.path.clone(),
3609                query: ctx.query.clone(),
3610                protocol: "HTTP/1.1".to_string(),
3611                status,
3612                body_bytes: ctx.response_bytes,
3613                duration_ms: duration.as_millis() as u64,
3614                client_ip: ctx.client_ip.clone(),
3615                user_agent: ctx.user_agent.clone(),
3616                referer: ctx.referer.clone(),
3617                host: ctx.host.clone(),
3618                route_id: ctx.route_id.clone(),
3619                upstream: ctx.upstream.clone(),
3620                upstream_attempts: ctx.upstream_attempts,
3621                instance_id: self.app_state.instance_id.clone(),
3622                namespace: ctx.namespace.clone(),
3623                service: ctx.service.clone(),
3624                // New fields
3625                body_bytes_sent: ctx.response_bytes,
3626                upstream_addr: ctx.selected_upstream_address.clone(),
3627                connection_reused: ctx.connection_reused,
3628                rate_limit_hit: status == 429,
3629                geo_country: ctx.geo_country_code.clone(),
3630            };
3631            self.log_manager.log_access(&access_entry);
3632        }
3633
3634        // Log to tracing at debug level (avoid allocations if debug disabled)
3635        if tracing::enabled!(tracing::Level::DEBUG) {
3636            // Pingora 0.8.0: upstream_write_pending_time for upload diagnostics
3637            let write_pending_ms = session.upstream_write_pending_time().as_millis() as u64;
3638            debug!(
3639                trace_id = %ctx.trace_id,
3640                method = %ctx.method,
3641                path = %ctx.path,
3642                route_id = ?ctx.route_id,
3643                upstream = ?ctx.upstream,
3644                status = status,
3645                duration_ms = duration.as_millis() as u64,
3646                upstream_write_pending_ms = write_pending_ms,
3647                upstream_attempts = ctx.upstream_attempts,
3648                error = ?_error.map(|e| e.to_string()),
3649                "Request completed"
3650            );
3651        }
3652
3653        // Log WebSocket upgrades at info level
3654        if ctx.is_websocket_upgrade && status == 101 {
3655            info!(
3656                trace_id = %ctx.trace_id,
3657                route_id = ?ctx.route_id,
3658                upstream = ?ctx.upstream,
3659                client_ip = %ctx.client_ip,
3660                "WebSocket connection established"
3661            );
3662        }
3663
3664        // End OpenTelemetry span
3665        if let Some(span) = ctx.otel_span.take() {
3666            span.end();
3667        }
3668    }
3669}
3670
3671// =============================================================================
3672// Helper methods for body streaming (not part of ProxyHttp trait)
3673// =============================================================================
3674
3675impl ZentinelProxy {
3676    /// Process a single body chunk in streaming mode.
3677    async fn process_body_chunk_streaming(
3678        &self,
3679        body: &mut Option<Bytes>,
3680        end_of_stream: bool,
3681        ctx: &mut RequestContext,
3682    ) -> Result<(), Box<Error>> {
3683        // Clone the chunk data to avoid borrowing issues when mutating body later
3684        let chunk_data: Vec<u8> = body.as_ref().map(|b| b.to_vec()).unwrap_or_default();
3685        let chunk_index = ctx.request_body_chunk_index;
3686        ctx.request_body_chunk_index += 1;
3687        ctx.body_bytes_inspected += chunk_data.len() as u64;
3688
3689        debug!(
3690            correlation_id = %ctx.trace_id,
3691            chunk_index = chunk_index,
3692            chunk_size = chunk_data.len(),
3693            end_of_stream = end_of_stream,
3694            "Streaming body chunk to agents"
3695        );
3696
3697        // Create agent call context
3698        let agent_ctx = crate::agents::AgentCallContext {
3699            correlation_id: zentinel_common::CorrelationId::from_string(&ctx.trace_id),
3700            metadata: zentinel_agent_protocol::RequestMetadata {
3701                correlation_id: ctx.trace_id.clone(),
3702                request_id: ctx.trace_id.clone(),
3703                client_ip: ctx.client_ip.clone(),
3704                client_port: 0,
3705                server_name: ctx.host.clone(),
3706                protocol: "HTTP/1.1".to_string(),
3707                tls_version: None,
3708                tls_cipher: None,
3709                route_id: ctx.route_id.clone(),
3710                upstream_id: ctx.upstream.clone(),
3711                timestamp: chrono::Utc::now().to_rfc3339(),
3712                traceparent: ctx.traceparent(),
3713            },
3714            route_id: ctx.route_id.clone(),
3715            upstream_id: ctx.upstream.clone(),
3716            request_body: None, // Not used in streaming mode
3717            response_body: None,
3718        };
3719
3720        let agent_ids = ctx.body_inspection_agents.clone();
3721        let total_size = None; // Unknown in streaming mode
3722
3723        match self
3724            .agent_manager
3725            .process_request_body_streaming(
3726                &agent_ctx,
3727                &chunk_data,
3728                end_of_stream,
3729                chunk_index,
3730                ctx.body_bytes_inspected as usize,
3731                total_size,
3732                &agent_ids,
3733            )
3734            .await
3735        {
3736            Ok(decision) => {
3737                // Track if agent needs more data
3738                ctx.agent_needs_more = decision.needs_more;
3739
3740                // Apply body mutation if present
3741                if let Some(ref mutation) = decision.request_body_mutation {
3742                    if !mutation.is_pass_through() {
3743                        if mutation.is_drop() {
3744                            // Drop the chunk
3745                            *body = None;
3746                            trace!(
3747                                correlation_id = %ctx.trace_id,
3748                                chunk_index = chunk_index,
3749                                "Agent dropped body chunk"
3750                            );
3751                        } else if let Some(ref new_data) = mutation.data {
3752                            // Replace chunk with mutated content
3753                            *body = Some(Bytes::from(new_data.clone()));
3754                            trace!(
3755                                correlation_id = %ctx.trace_id,
3756                                chunk_index = chunk_index,
3757                                original_size = chunk_data.len(),
3758                                new_size = new_data.len(),
3759                                "Agent mutated body chunk"
3760                            );
3761                        }
3762                    }
3763                }
3764
3765                // Check decision (only final if needs_more is false)
3766                if !decision.needs_more && !decision.is_allow() {
3767                    warn!(
3768                        correlation_id = %ctx.trace_id,
3769                        action = ?decision.action,
3770                        "Agent blocked request body"
3771                    );
3772                    self.metrics.record_blocked_request("agent_body_inspection");
3773
3774                    let (status, message) = match &decision.action {
3775                        crate::agents::AgentAction::Block { status, body, .. } => (
3776                            *status,
3777                            body.clone().unwrap_or_else(|| "Blocked".to_string()),
3778                        ),
3779                        _ => (403, "Forbidden".to_string()),
3780                    };
3781
3782                    return Err(Error::explain(ErrorType::HTTPStatus(status), message));
3783                }
3784
3785                trace!(
3786                    correlation_id = %ctx.trace_id,
3787                    needs_more = decision.needs_more,
3788                    "Agent processed body chunk"
3789                );
3790            }
3791            Err(e) => {
3792                let fail_closed = ctx
3793                    .route_config
3794                    .as_ref()
3795                    .map(|r| r.policies.failure_mode == zentinel_config::FailureMode::Closed)
3796                    .unwrap_or(false);
3797
3798                if fail_closed {
3799                    error!(
3800                        correlation_id = %ctx.trace_id,
3801                        error = %e,
3802                        "Agent streaming body inspection failed, blocking (fail-closed)"
3803                    );
3804                    self.log_manager.log_request_error(
3805                        "error",
3806                        "Agent streaming body inspection failed, blocking (fail-closed)",
3807                        &ctx.trace_id,
3808                        ctx.route_id.as_deref(),
3809                        ctx.upstream.as_deref(),
3810                        Some(format!("error={}", e)),
3811                    );
3812                    return Err(Error::explain(
3813                        ErrorType::HTTPStatus(503),
3814                        "Service unavailable",
3815                    ));
3816                } else {
3817                    warn!(
3818                        correlation_id = %ctx.trace_id,
3819                        error = %e,
3820                        "Agent streaming body inspection failed, allowing (fail-open)"
3821                    );
3822                    self.log_manager.log_request_error(
3823                        "warn",
3824                        "Agent streaming body inspection failed, allowing (fail-open)",
3825                        &ctx.trace_id,
3826                        ctx.route_id.as_deref(),
3827                        ctx.upstream.as_deref(),
3828                        Some(format!("error={}", e)),
3829                    );
3830                }
3831            }
3832        }
3833
3834        Ok(())
3835    }
3836
3837    /// Send buffered body to agents (buffer mode).
3838    async fn send_buffered_body_to_agents(
3839        &self,
3840        end_of_stream: bool,
3841        ctx: &mut RequestContext,
3842    ) -> Result<(), Box<Error>> {
3843        debug!(
3844            correlation_id = %ctx.trace_id,
3845            buffer_size = ctx.body_buffer.len(),
3846            end_of_stream = end_of_stream,
3847            agent_count = ctx.body_inspection_agents.len(),
3848            decompression_enabled = ctx.decompression_enabled,
3849            "Sending buffered body to agents for inspection"
3850        );
3851
3852        // Decompress body if enabled and we have a supported encoding
3853        let body_for_inspection = if ctx.decompression_enabled {
3854            if let Some(ref encoding) = ctx.body_content_encoding {
3855                let config = crate::decompression::DecompressionConfig {
3856                    max_ratio: ctx.max_decompression_ratio,
3857                    max_output_bytes: ctx.max_decompression_bytes,
3858                };
3859
3860                match crate::decompression::decompress_body(&ctx.body_buffer, encoding, &config) {
3861                    Ok(result) => {
3862                        ctx.body_was_decompressed = true;
3863                        self.metrics
3864                            .record_decompression_success(encoding, result.ratio);
3865                        debug!(
3866                            correlation_id = %ctx.trace_id,
3867                            encoding = %encoding,
3868                            compressed_size = result.compressed_size,
3869                            decompressed_size = result.decompressed_size,
3870                            ratio = result.ratio,
3871                            "Body decompressed for agent inspection"
3872                        );
3873                        result.data
3874                    }
3875                    Err(e) => {
3876                        // Record failure metric
3877                        let failure_reason = match &e {
3878                            crate::decompression::DecompressionError::RatioExceeded { .. } => {
3879                                "ratio_exceeded"
3880                            }
3881                            crate::decompression::DecompressionError::SizeExceeded { .. } => {
3882                                "size_exceeded"
3883                            }
3884                            crate::decompression::DecompressionError::InvalidData { .. } => {
3885                                "invalid_data"
3886                            }
3887                            crate::decompression::DecompressionError::UnsupportedEncoding {
3888                                ..
3889                            } => "unsupported",
3890                            crate::decompression::DecompressionError::IoError(_) => "io_error",
3891                        };
3892                        self.metrics
3893                            .record_decompression_failure(encoding, failure_reason);
3894
3895                        // Decompression failed - decide based on failure mode
3896                        let fail_closed = ctx
3897                            .route_config
3898                            .as_ref()
3899                            .map(|r| {
3900                                r.policies.failure_mode == zentinel_config::FailureMode::Closed
3901                            })
3902                            .unwrap_or(false);
3903
3904                        if fail_closed {
3905                            error!(
3906                                correlation_id = %ctx.trace_id,
3907                                error = %e,
3908                                encoding = %encoding,
3909                                "Decompression failed, blocking (fail-closed)"
3910                            );
3911                            return Err(Error::explain(
3912                                ErrorType::HTTPStatus(400),
3913                                "Invalid compressed body",
3914                            ));
3915                        } else {
3916                            warn!(
3917                                correlation_id = %ctx.trace_id,
3918                                error = %e,
3919                                encoding = %encoding,
3920                                "Decompression failed, sending compressed body (fail-open)"
3921                            );
3922                            ctx.body_buffer.clone()
3923                        }
3924                    }
3925                }
3926            } else {
3927                ctx.body_buffer.clone()
3928            }
3929        } else {
3930            ctx.body_buffer.clone()
3931        };
3932
3933        let agent_ctx = crate::agents::AgentCallContext {
3934            correlation_id: zentinel_common::CorrelationId::from_string(&ctx.trace_id),
3935            metadata: zentinel_agent_protocol::RequestMetadata {
3936                correlation_id: ctx.trace_id.clone(),
3937                request_id: ctx.trace_id.clone(),
3938                client_ip: ctx.client_ip.clone(),
3939                client_port: 0,
3940                server_name: ctx.host.clone(),
3941                protocol: "HTTP/1.1".to_string(),
3942                tls_version: None,
3943                tls_cipher: None,
3944                route_id: ctx.route_id.clone(),
3945                upstream_id: ctx.upstream.clone(),
3946                timestamp: chrono::Utc::now().to_rfc3339(),
3947                traceparent: ctx.traceparent(),
3948            },
3949            route_id: ctx.route_id.clone(),
3950            upstream_id: ctx.upstream.clone(),
3951            request_body: Some(body_for_inspection.clone()),
3952            response_body: None,
3953        };
3954
3955        let agent_ids = ctx.body_inspection_agents.clone();
3956        match self
3957            .agent_manager
3958            .process_request_body(&agent_ctx, &body_for_inspection, end_of_stream, &agent_ids)
3959            .await
3960        {
3961            Ok(decision) => {
3962                if !decision.is_allow() {
3963                    warn!(
3964                        correlation_id = %ctx.trace_id,
3965                        action = ?decision.action,
3966                        "Agent blocked request body"
3967                    );
3968                    self.metrics.record_blocked_request("agent_body_inspection");
3969
3970                    let (status, message) = match &decision.action {
3971                        crate::agents::AgentAction::Block { status, body, .. } => (
3972                            *status,
3973                            body.clone().unwrap_or_else(|| "Blocked".to_string()),
3974                        ),
3975                        _ => (403, "Forbidden".to_string()),
3976                    };
3977
3978                    return Err(Error::explain(ErrorType::HTTPStatus(status), message));
3979                }
3980
3981                trace!(
3982                    correlation_id = %ctx.trace_id,
3983                    "Agent allowed request body"
3984                );
3985            }
3986            Err(e) => {
3987                let fail_closed = ctx
3988                    .route_config
3989                    .as_ref()
3990                    .map(|r| r.policies.failure_mode == zentinel_config::FailureMode::Closed)
3991                    .unwrap_or(false);
3992
3993                if fail_closed {
3994                    error!(
3995                        correlation_id = %ctx.trace_id,
3996                        error = %e,
3997                        "Agent body inspection failed, blocking (fail-closed)"
3998                    );
3999                    self.log_manager.log_request_error(
4000                        "error",
4001                        "Agent body inspection failed, blocking (fail-closed)",
4002                        &ctx.trace_id,
4003                        ctx.route_id.as_deref(),
4004                        ctx.upstream.as_deref(),
4005                        Some(format!("error={}", e)),
4006                    );
4007                    return Err(Error::explain(
4008                        ErrorType::HTTPStatus(503),
4009                        "Service unavailable",
4010                    ));
4011                } else {
4012                    warn!(
4013                        correlation_id = %ctx.trace_id,
4014                        error = %e,
4015                        "Agent body inspection failed, allowing (fail-open)"
4016                    );
4017                    self.log_manager.log_request_error(
4018                        "warn",
4019                        "Agent body inspection failed, allowing (fail-open)",
4020                        &ctx.trace_id,
4021                        ctx.route_id.as_deref(),
4022                        ctx.upstream.as_deref(),
4023                        Some(format!("error={}", e)),
4024                    );
4025                }
4026            }
4027        }
4028
4029        Ok(())
4030    }
4031}