diff --git a/apis/src/correlation/mod.rs b/apis/src/correlation/mod.rs new file mode 100644 index 0000000000..7b28cb9313 --- /dev/null +++ b/apis/src/correlation/mod.rs @@ -0,0 +1,433 @@ +// SPDX-License-Identifier: MIT +// Copyright (c) 2026 Praxis Contributors + +//! Request correlation shared by forwarded and delegated calls. +//! +//! A single AI request reaches the backend over two different +//! paths. The **forwarded** request is the client's own request, +//! proxied upstream. A **delegated** call is one the proxy +//! originates itself while handling that request — a Files API +//! metadata fetch, a vector-store search. Without shared +//! correlation the two appear at the backend as unrelated traffic, +//! so a slow or failing request cannot be attributed to a layer. +//! +//! [`TraceContext`] resolves the identifiers once per downstream +//! request. Every hop then draws its own span from that context, +//! so all legs share a trace-id while remaining individually +//! measurable. +//! +//! # Resolution +//! +//! The request ID follows the same precedence the `request_id` +//! core builtin uses when echoing an ID onto the response: +//! +//! 1. A client-supplied header on the downstream request. +//! 2. An ID injected into `extra_request_headers` by an earlier filter. +//! 3. A freshly generated ID. +//! +//! Step 2 matters: filters that inject headers write to +//! `extra_request_headers` (a pending mutation applied to the +//! upstream request) rather than back into `ctx.request.headers`, +//! so reading the downstream headers alone finds an injected value +//! only when the client happened to supply one. It is also why the +//! `request_id` builtin must be registered *before* any filter that +//! resolves correlation: it reads only the downstream headers, so a +//! later `request_id` generates a second, different ID that wins on +//! the forwarded request while the delegated calls keep the first. +//! +//! # Agreement across hops +//! +//! Header inspection alone is not enough to make every hop of one +//! request agree. Delegated callouts can run in the `StreamBuffer` +//! pre-read phase, ahead of header-phase filters, and see only the +//! original downstream headers — never a value a filter injected. +//! +//! The resolved [`TraceContext`] is therefore shared through request +//! extensions: whichever hop asks first initializes it, and every +//! later hop reuses it regardless of filter order. + +use http::{HeaderMap, HeaderName, HeaderValue}; +use praxis_filter::HttpFilterContext; +use tracing::{debug, trace}; + +// ----------------------------------------------------------------------------- +// Constants +// ----------------------------------------------------------------------------- + +/// Header carrying the request correlation ID. +const REQUEST_ID: HeaderName = HeaderName::from_static("x-request-id"); + +/// Header carrying W3C trace context. +const TRACEPARENT: HeaderName = HeaderName::from_static("traceparent"); + +/// Trace-context version this proxy emits. +const VERSION: &str = "00"; + +/// Trace-flags emitted when starting a new trace: sampled. +const SAMPLED: &str = "01"; + +/// The only trace-flags bit defined by the current specification. +const SAMPLED_BIT: u8 = 0x01; + +/// Hex length of a W3C trace-id. +const TRACE_ID_LEN: usize = 32; + +/// Hex length of a W3C span-id. +const SPAN_ID_LEN: usize = 16; + +/// Number of fixed fields every `traceparent` version begins with. +const BASE_FIELDS: usize = 4; + +/// Last-resort trace-id when sanitization would yield all zeros, +/// which W3C Trace Context section 2.2.2 forbids. +const FALLBACK_TRACE_ID: &str = "00000000000000000000000000000001"; + +/// Last-resort span-id, for the same reason as [`FALLBACK_TRACE_ID`]. +const FALLBACK_SPAN_ID: &str = "0000000000000001"; + +// ----------------------------------------------------------------------------- +// TraceContext +// ----------------------------------------------------------------------------- + +/// Correlation identifiers for one downstream request. +/// +/// Resolved once, then used to stamp every hop the request makes. +/// Each hop gets its own span-id under the shared trace-id. +/// +/// Shared through [`RequestExtensions`] rather than through header +/// mutations, because filter ordering does not guarantee who resolves +/// first. Delegated callouts made from the `StreamBuffer` pre-read +/// phase run before header-phase filters and see only the original +/// downstream headers, so a value injected by a filter would not +/// reach them. Whichever hop resolves first initializes the shared +/// context; every later hop reuses it. +/// +/// [`RequestExtensions`]: praxis_filter::RequestExtensions +#[derive(Clone)] +pub struct TraceContext { + /// 2 lowercase hex characters of W3C trace-flags. + flags: String, + + /// Value for the `x-request-id` header. + request_id: String, + + /// 32 lowercase hex characters shared by every hop. + trace_id: String, +} + +impl TraceContext { + /// Read the request's shared trace context. + /// + /// Falls back to resolving a fresh context when no hop has + /// initialized one yet, so a callout still carries correlation + /// headers in a chain that never called [`Self::get_or_init`]. + /// + /// The fallback resolves *without storing*, so repeated calls on a + /// request with no client-supplied `traceparent` produce different + /// trace-ids. Callers that make more than one hop must therefore + /// call [`Self::get_or_init`] first — every callsite does today, + /// from the filter's `on_request_body` entry point. + pub(crate) fn from_filter_context(ctx: &HttpFilterContext<'_>) -> Self { + if let Some(shared) = ctx.extensions.get::() { + return shared.clone(); + } + + debug!( + "no shared trace context in extensions; resolving independently \ + — call `get_or_init` first for stable correlation" + ); + Self::resolve(ctx) + } + + /// Return the request's shared trace context, resolving and + /// storing it if this is the first hop to ask. + /// + /// Prefer this wherever the context is available mutably: it is + /// what makes every hop of one request agree on a trace-id. + pub fn get_or_init(ctx: &mut HttpFilterContext<'_>) -> Self { + if let Some(shared) = ctx.extensions.get::() { + return shared.clone(); + } + + let resolved = Self::resolve(ctx); + ctx.extensions.insert(resolved.clone()); + resolved + } + + /// Build the correlation headers for one hop. + #[must_use] + pub fn headers_for_hop(&self, ctx: &HttpFilterContext<'_>) -> [(HeaderName, String); 2] { + [ + (REQUEST_ID, self.request_id.clone()), + (TRACEPARENT, self.traceparent_for_hop(ctx)), + ] + } + + /// The resolved request correlation ID. + #[must_use] + pub fn request_id(&self) -> &str { + &self.request_id + } + + /// Resolve identifiers from the request, ignoring any shared + /// context. + fn resolve(ctx: &HttpFilterContext<'_>) -> Self { + let request_id = resolve_request_id(ctx); + let (trace_id, flags) = resolve_trace(ctx); + + trace!(%request_id, %trace_id, "resolved correlation for request"); + + Self { + flags, + request_id, + trace_id, + } + } + + /// Build a `traceparent` for one hop, with a fresh span-id. + /// + /// Called once per outbound leg. The forwarded request and each + /// delegated call are separate spans of the same trace, which is + /// what makes their latencies separable. + fn traceparent_for_hop(&self, ctx: &HttpFilterContext<'_>) -> String { + let span_id = generate_span_id(ctx); + let Self { flags, trace_id, .. } = self; + format!("{VERSION}-{trace_id}-{span_id}-{flags}") + } +} + +// ----------------------------------------------------------------------------- +// Correlation +// ----------------------------------------------------------------------------- + +/// Correlation headers for the delegated calls of one request. +/// +/// Unlike an operator-configured header allowlist, these are +/// injected unconditionally. Correlation that depends on per-filter +/// configuration is correlation that silently goes missing. +pub(crate) struct Correlation { + /// Value for the `x-request-id` header. + request_id: HeaderValue, + + /// Value for the `traceparent` header. + traceparent: HeaderValue, +} + +impl Correlation { + /// Insert correlation headers into an outbound callout map. + /// + /// Inserts rather than appends, so correlation overwrites any + /// same-named header copied from the downstream request. + pub(crate) fn apply(&self, map: &mut HeaderMap) { + map.insert(REQUEST_ID, self.request_id.clone()); + map.insert(TRACEPARENT, self.traceparent.clone()); + } + + /// Resolve correlation headers for this request's delegated calls. + pub(crate) fn from_filter_context(ctx: &HttpFilterContext<'_>) -> Self { + let context = TraceContext::from_filter_context(ctx); + + // A generated ID is hex; a client-supplied one already passed + // header parsing to reach us. Fall back rather than fail the + // callout on an unrepresentable value. + let to_value = |raw: &str| HeaderValue::from_str(raw).unwrap_or_else(|_| HeaderValue::from_static("")); + + Self { + request_id: to_value(context.request_id()), + traceparent: to_value(&context.traceparent_for_hop(ctx)), + } + } +} + +// ----------------------------------------------------------------------------- +// InboundTrace +// ----------------------------------------------------------------------------- + +/// Trace-id and flags carried forward from a valid `traceparent`. +struct InboundTrace { + /// 2 lowercase hex characters. + flags: String, + + /// 32 lowercase hex characters. + trace_id: String, +} + +// ----------------------------------------------------------------------------- +// Utility Functions +// ----------------------------------------------------------------------------- + +/// Generate a 16-hex-character span-id for one hop. +/// +/// The generator's trailing sequence counter advances on every +/// call, so each hop's span-id differs even when drawn within the +/// same microsecond. +fn generate_span_id(ctx: &HttpFilterContext<'_>) -> String { + span_id_from(&generate_trace_id(ctx)) +} + +/// Generate a 32-hex-character trace-id. +/// +/// The core ID generator emits exactly 32 hex characters +/// (`{micros:012x}{seed:08x}{seq:012x}`), which is the W3C +/// trace-id width. +fn generate_trace_id(ctx: &HttpFilterContext<'_>) -> String { + sanitize_trace_id(&ctx.id_generator.generate(ctx.time_source)) +} + +/// Look up a header injected by an earlier filter into +/// `extra_request_headers`. +fn injected_header(ctx: &HttpFilterContext<'_>, name: &str) -> Option { + ctx.extra_request_headers + .iter() + .find(|(header, _)| header.eq_ignore_ascii_case(name)) + .map(|(_, value)| value.clone()) +} + +/// Check that a hex string is entirely zeros. +fn is_all_zero(value: &str) -> bool { + value.bytes().all(|b| b == b'0') +} + +/// Check that every character is a lowercase hex digit. +fn is_lower_hex(value: &str) -> bool { + value.bytes().all(|b| b.is_ascii_digit() || (b'a'..=b'f').contains(&b)) +} + +/// Clear every trace-flags bit this version does not define. +/// +/// Only the sampled bit is specified today. Re-emitting reserved bits +/// under version `00` would give them a meaning the specification has +/// not assigned, so they are dropped rather than carried downstream. +fn mask_flags(flags: &str) -> String { + let bits = u8::from_str_radix(flags, 16).unwrap_or(0); + format!("{:02x}", bits & SAMPLED_BIT) +} + +/// Parse and validate a W3C `traceparent`. +/// +/// Returns `None` for anything not matching +/// `---` with lowercase hex +/// fields, an all-zero trace-id or span-id, or the forbidden +/// version `ff`. +/// +/// A higher version may append fields after the flags. Those are +/// ignored and the base fields are still continued, as the +/// [versioning rules] require — a trace must survive a future-version +/// rollout upstream. Version `00` accepts no trailing fields. +/// +/// [versioning rules]: https://www.w3.org/TR/trace-context/#versioning-of-traceparent +fn parse_traceparent(value: &str) -> Option { + let fields: Vec<&str> = value.split('-').collect(); + let [version, trace_id, span_id, flags] = fields.get(..BASE_FIELDS)? else { + return None; + }; + + if version.len() != 2 || !is_lower_hex(version) || *version == "ff" { + return None; + } + // Only a future version may carry extension fields; version 00 is + // exactly four. + if fields.len() > BASE_FIELDS && *version == VERSION { + return None; + } + if trace_id.len() != TRACE_ID_LEN || !is_lower_hex(trace_id) || is_all_zero(trace_id) { + return None; + } + if span_id.len() != SPAN_ID_LEN || !is_lower_hex(span_id) || is_all_zero(span_id) { + return None; + } + if flags.len() != 2 || !is_lower_hex(flags) { + return None; + } + + Some(InboundTrace { + flags: mask_flags(flags), + trace_id: (*trace_id).to_owned(), + }) +} + +/// Resolve the request ID, generating one if nothing upstream +/// supplied or injected it. +fn resolve_request_id(ctx: &HttpFilterContext<'_>) -> String { + if let Some(client_id) = ctx.request.headers.get(&REQUEST_ID).and_then(|v| v.to_str().ok()) { + return client_id.to_owned(); + } + + if let Some(injected) = injected_header(ctx, REQUEST_ID.as_str()) { + return injected; + } + + ctx.id_generator.generate(ctx.time_source) +} + +/// Resolve the trace-id and flags shared by every hop. +/// +/// Continues a valid inbound or injected trace; otherwise starts a +/// new sampled one. A value that fails validation is discarded +/// rather than carried forward — it is client-controlled input that +/// would otherwise reach the telemetry backend unchecked. +fn resolve_trace(ctx: &HttpFilterContext<'_>) -> (String, String) { + let inbound = ctx + .request + .headers + .get(&TRACEPARENT) + .and_then(|v| v.to_str().ok()) + .and_then(parse_traceparent) + .or_else(|| { + injected_header(ctx, TRACEPARENT.as_str()) + .as_deref() + .and_then(parse_traceparent) + }); + + if let Some(InboundTrace { flags, trace_id }) = inbound { + return (trace_id, flags); + } + + (generate_trace_id(ctx), SAMPLED.to_owned()) +} + +/// Coerce a generated ID into a W3C-valid trace-id. +/// +/// Defensive: keeps emitting a well-formed trace-id if the core +/// generator's format ever changes. An all-zero value is invalid per +/// W3C Trace Context section 2.2.2, so sanitization that would erase +/// every digit falls back to a fixed non-zero ID rather than emitting +/// one a collector must reject. +fn sanitize_trace_id(id: &str) -> String { + if id.len() == TRACE_ID_LEN && is_lower_hex(id) && !is_all_zero(id) { + return id.to_owned(); + } + + let sanitized = format!("{id:0>TRACE_ID_LEN$.TRACE_ID_LEN$}") + .to_ascii_lowercase() + .replace(|c: char| !c.is_ascii_hexdigit(), "0"); + if sanitized.len() != TRACE_ID_LEN || is_all_zero(&sanitized) { + return FALLBACK_TRACE_ID.to_owned(); + } + sanitized +} + +/// Take the span-id from the tail of a trace-id, guarding the same +/// all-zero case as [`sanitize_trace_id`]. +fn span_id_from(trace_id: &str) -> String { + let span_id = trace_id.get(TRACE_ID_LEN - SPAN_ID_LEN..).unwrap_or(""); + if span_id.len() != SPAN_ID_LEN || is_all_zero(span_id) { + return FALLBACK_SPAN_ID.to_owned(); + } + span_id.to_owned() +} + +// ----------------------------------------------------------------------------- +// Tests +// ----------------------------------------------------------------------------- + +#[cfg(test)] +#[expect(clippy::allow_attributes, reason = "blanket test suppressions")] +#[allow( + clippy::unwrap_used, + clippy::expect_used, + clippy::indexing_slicing, + clippy::panic, + reason = "tests" +)] +mod tests; diff --git a/apis/src/correlation/tests.rs b/apis/src/correlation/tests.rs new file mode 100644 index 0000000000..130a424e1c --- /dev/null +++ b/apis/src/correlation/tests.rs @@ -0,0 +1,383 @@ +// SPDX-License-Identifier: MIT +// Copyright (c) 2026 Praxis Contributors + +use std::borrow::Cow; + +use http::{HeaderMap, HeaderValue, Method}; + +use super::*; +use crate::test_utils::{make_filter_context, make_request}; + +// ----------------------------------------------------------------------------- +// Request ID resolution +// ----------------------------------------------------------------------------- + +#[test] +fn prefers_client_supplied_request_id() { + let req = request_with(&[("x-request-id", "client-abc")]); + let ctx = make_filter_context(&req); + + let map = applied(&Correlation::from_filter_context(&ctx)); + + assert_eq!( + header(&map, "x-request-id"), + "client-abc", + "client-supplied request ID should be forwarded unchanged" + ); +} + +#[test] +fn uses_id_injected_by_request_id_filter() { + let req = request_with(&[]); + let mut ctx = make_filter_context(&req); + // The request_id core builtin injects here, not into + // ctx.request.headers — the case a downstream-headers-only + // lookup would miss. + ctx.extra_request_headers + .push((Cow::Borrowed("X-Request-ID"), "generated-by-filter".to_owned())); + + let map = applied(&Correlation::from_filter_context(&ctx)); + + assert_eq!( + header(&map, "x-request-id"), + "generated-by-filter", + "injected request ID should be picked up from extra_request_headers" + ); +} + +#[test] +fn generates_request_id_when_no_source_available() { + let req = request_with(&[]); + let ctx = make_filter_context(&req); + + let map = applied(&Correlation::from_filter_context(&ctx)); + let id = header(&map, "x-request-id"); + + assert_eq!(id.len(), 32, "generated request ID should be 32 hex chars: {id}"); + assert!(is_lower_hex(&id), "generated request ID should be lowercase hex: {id}"); +} + +// ----------------------------------------------------------------------------- +// Traceparent +// ----------------------------------------------------------------------------- + +#[test] +fn continues_valid_inbound_trace_under_new_span() { + let inbound = "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01"; + let req = request_with(&[("traceparent", inbound)]); + let ctx = make_filter_context(&req); + + let map = applied(&Correlation::from_filter_context(&ctx)); + let outbound = header(&map, "traceparent"); + let parts: Vec<&str> = outbound.split('-').collect(); + + assert_eq!(parts.len(), 4, "traceparent should have four fields: {outbound}"); + assert_eq!( + parts[1], "4bf92f3577b34da6a3ce929d0e0e4736", + "trace-id should be carried forward" + ); + assert_eq!(parts[3], "01", "inbound trace flags should be preserved"); + assert_ne!( + parts[2], "00f067aa0ba902b7", + "delegation hop should emit its own span-id" + ); + assert_eq!(parts[2].len(), 16, "span-id should be 16 hex chars: {outbound}"); +} + +#[test] +fn starts_new_trace_when_absent() { + let req = request_with(&[]); + let ctx = make_filter_context(&req); + + let map = applied(&Correlation::from_filter_context(&ctx)); + let outbound = header(&map, "traceparent"); + let parts: Vec<&str> = outbound.split('-').collect(); + + assert_eq!(parts.len(), 4, "traceparent should have four fields: {outbound}"); + assert_eq!(parts[0], "00", "version should be 00"); + assert_eq!(parts[1].len(), 32, "trace-id should be 32 hex chars: {outbound}"); + assert_eq!(parts[2].len(), 16, "span-id should be 16 hex chars: {outbound}"); + assert_eq!(parts[3], "01", "new traces should be marked sampled"); + assert!(!is_all_zero(parts[1]), "trace-id must not be all zeros"); + assert!(!is_all_zero(parts[2]), "span-id must not be all zeros"); +} + +#[test] +fn discards_malformed_inbound_traceparent() { + // Client-controlled input must not reach the telemetry backend + // unchecked. + let malformed = [ + "not-a-traceparent", + "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7", + "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01-extra", + "00-00000000000000000000000000000000-00f067aa0ba902b7-01", + "00-4bf92f3577b34da6a3ce929d0e0e4736-0000000000000000-01", + "ff-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01", + "00-4BF92F3577B34DA6A3CE929D0E0E4736-00f067aa0ba902b7-01", + "00-4bf92f3577b34da6a3ce929d0e0e473-00f067aa0ba902b7-01", + "00-zzf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01", + ]; + + for value in malformed { + let req = request_with(&[("traceparent", value)]); + let ctx = make_filter_context(&req); + + let map = applied(&Correlation::from_filter_context(&ctx)); + let outbound = header(&map, "traceparent"); + + assert_ne!( + outbound, value, + "malformed traceparent should not be forwarded: {value}" + ); + let parts: Vec<&str> = outbound.split('-').collect(); + assert_eq!(parts.len(), 4, "replacement should be well-formed: {outbound}"); + assert_eq!(parts[1].len(), 32, "replacement trace-id should be 32 hex: {outbound}"); + } +} + +#[test] +fn parse_traceparent_accepts_unsampled_flags() { + let parsed = parse_traceparent("00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-00") + .expect("unsampled but well-formed traceparent should parse"); + + assert_eq!(parsed.flags, "00", "unsampled flags should be preserved"); +} + +#[test] +fn parse_traceparent_accepts_future_versions() { + let parsed = parse_traceparent("01-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01") + .expect("future version should still yield a usable trace-id"); + + assert_eq!( + parsed.trace_id, "4bf92f3577b34da6a3ce929d0e0e4736", + "trace-id should be extracted from a future version" + ); +} + +#[test] +fn parse_traceparent_ignores_future_version_extension_fields() { + // A higher version may append fields after the flags. Dropping the + // trace because of them would restart traces during an upstream + // version rollout. + let parsed = parse_traceparent("01-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01-extension") + .expect("extension fields of a future version should be ignored"); + + assert_eq!( + parsed.trace_id, "4bf92f3577b34da6a3ce929d0e0e4736", + "base fields of a future version should still be continued" + ); + assert_eq!(parsed.flags, "01", "base flags should survive the extension field"); +} + +#[test] +fn masks_trace_flags_this_version_does_not_define() { + // Only the sampled bit is specified; the rest are reserved and + // must not be re-emitted with an invented meaning. + for (inbound, expected) in [("ff", "01"), ("fe", "00"), ("03", "01"), ("02", "00")] { + let value = format!("00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-{inbound}"); + let parsed = parse_traceparent(&value).expect("well-formed flags should parse"); + + assert_eq!( + parsed.flags, expected, + "flags {inbound} should be masked to the sampled bit" + ); + } +} + +#[test] +fn masked_flags_reach_the_outbound_header() { + let req = request_with(&[("traceparent", "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-ff")]); + let ctx = make_filter_context(&req); + + let map = applied(&Correlation::from_filter_context(&ctx)); + let outbound = header(&map, "traceparent"); + + assert!( + outbound.ends_with("-01"), + "reserved flag bits must not be forwarded: {outbound}" + ); +} + +// ----------------------------------------------------------------------------- +// Defensive ID sanitization +// ----------------------------------------------------------------------------- + +#[test] +fn sanitized_ids_are_never_all_zero() { + // Inputs a future core generator could plausibly produce, none of + // which may yield the all-zero ID W3C section 2.2.2 forbids. + let degenerate = ["", "0", "00000000000000000000000000000000", "----", "zzzz"]; + + for id in degenerate { + let trace_id = sanitize_trace_id(id); + assert_eq!(trace_id.len(), TRACE_ID_LEN, "trace-id should be 32 chars: {id:?}"); + assert!(is_lower_hex(&trace_id), "trace-id should be lowercase hex: {trace_id}"); + assert!(!is_all_zero(&trace_id), "trace-id must not be all zeros: {id:?}"); + + let span_id = span_id_from(&trace_id); + assert_eq!(span_id.len(), SPAN_ID_LEN, "span-id should be 16 chars: {id:?}"); + assert!(!is_all_zero(&span_id), "span-id must not be all zeros: {id:?}"); + } +} + +#[test] +fn sanitize_preserves_a_well_formed_generated_id() { + let id = "4bf92f3577b34da6a3ce929d0e0e4736"; + + assert_eq!(sanitize_trace_id(id), id, "a valid generated ID should pass through"); +} + +#[test] +fn span_id_falls_back_when_the_trace_id_tail_is_zero() { + // A trace-id that is itself valid can still end in 16 zeros. + let trace_id = "4bf92f3577b34da60000000000000000"; + + let span_id = span_id_from(trace_id); + + assert!(!is_all_zero(&span_id), "span-id must not be all zeros: {span_id}"); + assert_eq!(span_id.len(), SPAN_ID_LEN, "span-id should be 16 chars: {span_id}"); +} + +// ----------------------------------------------------------------------------- +// Forwarded and delegated legs share one trace +// ----------------------------------------------------------------------------- + +#[test] +fn delegated_call_joins_trace_injected_for_the_forwarded_request() { + // What the trace_context filter injects for the upstream hop. + let req = request_with(&[]); + let mut ctx = make_filter_context(&req); + let forwarded = TraceContext::from_filter_context(&ctx); + let forwarded_traceparent = forwarded.traceparent_for_hop(&ctx); + ctx.extra_request_headers + .push((Cow::Borrowed("traceparent"), forwarded_traceparent.clone())); + ctx.extra_request_headers + .push((Cow::Borrowed("X-Request-ID"), forwarded.request_id().to_owned())); + + // What a delegated call resolves later in the same request. + let map = applied(&Correlation::from_filter_context(&ctx)); + let delegated_traceparent = header(&map, "traceparent"); + + let forwarded_parts: Vec<&str> = forwarded_traceparent.split('-').collect(); + let delegated_parts: Vec<&str> = delegated_traceparent.split('-').collect(); + + assert_eq!( + forwarded_parts[1], delegated_parts[1], + "forwarded and delegated legs must share a trace-id" + ); + assert_ne!( + forwarded_parts[2], delegated_parts[2], + "each leg must be its own span so their latencies stay separable" + ); + assert_eq!( + header(&map, "x-request-id"), + forwarded.request_id(), + "both legs must carry the same request ID" + ); +} + +#[test] +fn init_makes_repeated_resolutions_share_one_trace() { + let req = request_with(&[]); + let mut ctx = make_filter_context(&req); + drop(TraceContext::get_or_init(&mut ctx)); + + // Two resolution phases of one filter, e.g. current input and + // rehydrated history. + let first = applied(&Correlation::from_filter_context(&ctx)); + let second = applied(&Correlation::from_filter_context(&ctx)); + + let first_parts: Vec = header(&first, "traceparent").split('-').map(str::to_owned).collect(); + let second_parts: Vec = header(&second, "traceparent").split('-').map(str::to_owned).collect(); + + assert_eq!( + first_parts[1], second_parts[1], + "callouts of one request must share a trace-id once initialized" + ); + assert_eq!( + header(&first, "x-request-id"), + header(&second, "x-request-id"), + "callouts of one request must share a request ID" + ); +} + +#[test] +fn client_trace_reaches_both_legs_unchanged() { + let inbound = "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01"; + let req = request_with(&[("traceparent", inbound)]); + let ctx = make_filter_context(&req); + + let forwarded = TraceContext::from_filter_context(&ctx); + let forwarded_traceparent = forwarded.traceparent_for_hop(&ctx); + let map = applied(&Correlation::from_filter_context(&ctx)); + + for value in [&forwarded_traceparent, &header(&map, "traceparent")] { + let parts: Vec<&str> = value.split('-').collect(); + assert_eq!( + parts[1], "4bf92f3577b34da6a3ce929d0e0e4736", + "client trace-id should reach both legs: {value}" + ); + } +} + +// ----------------------------------------------------------------------------- +// Application onto the callout map +// ----------------------------------------------------------------------------- + +#[test] +fn correlation_overwrites_forwarded_value_of_same_name() { + let req = request_with(&[("x-request-id", "client-abc")]); + let ctx = make_filter_context(&req); + let correlation = Correlation::from_filter_context(&ctx); + + // Simulate an operator listing x-request-id in forward_headers + // and a stale value already present in the callout map. + let mut map = HeaderMap::new(); + map.insert( + HeaderName::from_static("x-request-id"), + HeaderValue::from_static("stale"), + ); + correlation.apply(&mut map); + + assert_eq!( + header(&map, "x-request-id"), + "client-abc", + "correlation should overwrite, not append" + ); + assert_eq!( + map.get_all("x-request-id").iter().count(), + 1, + "correlation should not leave duplicate header values" + ); +} + +// ----------------------------------------------------------------------------- +// Test Utilities +// ----------------------------------------------------------------------------- + +/// Apply correlation to an empty map and return it. +fn applied(correlation: &Correlation) -> HeaderMap { + let mut map = HeaderMap::new(); + correlation.apply(&mut map); + map +} + +/// Read a header as a string. +fn header(map: &HeaderMap, name: &str) -> String { + map.get(name) + .and_then(|v| v.to_str().ok()) + .unwrap_or_default() + .to_owned() +} + +/// Build a request carrying the given headers. +fn request_with(headers: &[(&'static str, &str)]) -> praxis_filter::Request { + let mut req = make_request(Method::POST, "/v1/responses"); + for (name, value) in headers { + req.headers.insert( + HeaderName::from_static(name), + HeaderValue::from_str(value).expect("valid test header value"), + ); + } + req +} diff --git a/apis/src/lib.rs b/apis/src/lib.rs index 06f1679385..428a3db42a 100644 --- a/apis/src/lib.rs +++ b/apis/src/lib.rs @@ -11,6 +11,7 @@ pub mod anthropic; pub mod classifier; +pub mod correlation; pub mod json_body; pub(crate) mod mcp_client; pub mod openai; diff --git a/apis/src/openai/api_client/mod.rs b/apis/src/openai/api_client/mod.rs index d37cba0ec4..c71e05836b 100644 --- a/apis/src/openai/api_client/mod.rs +++ b/apis/src/openai/api_client/mod.rs @@ -29,6 +29,7 @@ pub(crate) use self::{ error::ApiClientError, url::{resource_url, validate_base_url, validate_forward_headers}, }; +pub(crate) use crate::correlation::Correlation; use crate::subrequest::{self, SubRequest, SubRequestClient, SubRequestError, SubResponse}; /// Configuration for constructing an [`ApiClient`]. @@ -126,12 +127,14 @@ impl ApiClient { } /// Send a GET request and parse the response body as JSON. + /// + /// `callout_headers` comes from [`Self::callout_headers`]. pub(crate) async fn get_json( &self, url: String, - request_headers: &HeaderMap, + callout_headers: &HeaderMap, ) -> Result { - let headers = self.build_header_map(request_headers); + let headers = callout_headers.clone(); let response = self.execute_url(&url, http::Method::GET, headers, Bytes::new()).await?; serde_json::from_slice(&response.body).map_err(|e| ApiClientError::DecodeFailed { detail: format!("JSON decode failed: {e}"), @@ -144,13 +147,13 @@ impl ApiClient { &self, url: String, body: &serde_json::Value, - request_headers: &HeaderMap, + callout_headers: &HeaderMap, ) -> Result { let serialized = serde_json::to_vec(body).map_err(|e| ApiClientError::DecodeFailed { detail: format!("request body serialization failed: {e}"), })?; - let response = self.post_json_bytes(url, serialized, request_headers).await?; + let response = self.post_json_bytes(url, serialized, callout_headers).await?; serde_json::from_slice(&response).map_err(|e| ApiClientError::DecodeFailed { detail: format!("JSON decode failed: {e}"), @@ -163,9 +166,9 @@ impl ApiClient { &self, url: String, body: Vec, - request_headers: &HeaderMap, + callout_headers: &HeaderMap, ) -> Result { - let mut headers = self.build_header_map(request_headers); + let mut headers = callout_headers.clone(); headers.remove(http::header::CONTENT_TYPE); headers.insert( http::header::CONTENT_TYPE, @@ -187,10 +190,10 @@ impl ApiClient { pub(crate) async fn get_bytes( &self, url: &str, - request_headers: &HeaderMap, + callout_headers: &HeaderMap, max_bytes: usize, ) -> Result { - let headers = self.build_header_map(request_headers); + let headers = callout_headers.clone(); let request = SubRequest { method: http::Method::GET, uri: http::Uri::default(), @@ -223,14 +226,27 @@ impl ApiClient { headers } - /// Build a [`HeaderMap`] from forwarded headers. - fn build_header_map(&self, request_headers: &HeaderMap) -> HeaderMap { - let mut map = HeaderMap::new(); + /// Build the header set sent on every callout for one + /// downstream request. + /// + /// Combines the operator-configured `forward_headers` allowlist + /// with correlation headers, which are injected unconditionally + /// so a delegated call is always attributable to the downstream + /// request that caused it. + /// + /// Resolved once per downstream request at the filter boundary + /// and threaded into each callout, so the allowlist copy does + /// not repeat per file reference. + pub(crate) fn callout_headers(&self, request_headers: &HeaderMap, correlation: &Correlation) -> HeaderMap { + let mut map = HeaderMap::with_capacity(self.forward_header_names.len().saturating_add(2)); for name in &self.forward_header_names { if let Some(value) = request_headers.get(name) { map.insert(name.clone(), value.clone()); } } + // After the allowlist, so correlation wins over a stale + // client-supplied value of the same name. + correlation.apply(&mut map); map } diff --git a/apis/src/openai/responses/file_resolve/config.rs b/apis/src/openai/responses/file_resolve/config.rs index 751f64e441..2e19f2db73 100644 --- a/apis/src/openai/responses/file_resolve/config.rs +++ b/apis/src/openai/responses/file_resolve/config.rs @@ -82,7 +82,9 @@ pub(crate) struct FileResolveConfig { /// Headers to forward from the original request to the /// Files API for authentication and tenant isolation. No - /// downstream headers are forwarded by default. + /// downstream headers are forwarded by default. The + /// correlation headers `x-request-id` and `traceparent` are + /// always sent and do not need to be listed here. #[serde(default)] pub forward_headers: Vec, diff --git a/apis/src/openai/responses/file_resolve/mod.rs b/apis/src/openai/responses/file_resolve/mod.rs index 3dd1117515..1c7db05f80 100644 --- a/apis/src/openai/responses/file_resolve/mod.rs +++ b/apis/src/openai/responses/file_resolve/mod.rs @@ -311,14 +311,25 @@ async fn resolve_and_rewrite( body: &mut Option, parsed: &mut serde_json::Value, ) -> Result { + // Establish the request's shared trace context before the first + // callout. This filter's callouts can run in the pre-read phase, + // ahead of header-phase filters, so this is often the hop that + // initializes it rather than the one that inherits it. + drop(crate::correlation::TraceContext::get_or_init(ctx)); + + // One header set for every callout this filter makes: current + // input and rehydrated history resolve under a single delegation + // span, not one span per resolution phase. + let callout_headers = filter.client.callout_headers(ctx); + let mut budget = filter.client.resolution_budget(); - let count = match resolve_current_input(filter, ctx, parsed, &mut budget).await { + let count = match resolve_current_input(filter, parsed, &callout_headers, &mut budget).await { Ok(count) => count, Err(e) => return Ok(reject_resolve_error(&e)), }; if count == 0 { trace!("no file_id references found"); - if let Err(e) = update_state(filter, ctx, None, &mut budget).await { + if let Err(e) = update_state(filter, ctx, None, &callout_headers, &mut budget).await { return Ok(reject_resolve_error(&e)); } if let Some(rejection) = reject_oversized_state_body(ctx, filter.config.max_body_bytes)? { @@ -331,7 +342,7 @@ async fn resolve_and_rewrite( if let Some(rejection) = rewrite_body(body, parsed, filter.config.max_body_bytes, filter.name())? { return Ok(rejection); } - if let Err(e) = update_state(filter, ctx, Some(parsed), &mut budget).await { + if let Err(e) = update_state(filter, ctx, Some(parsed), &callout_headers, &mut budget).await { return Ok(reject_resolve_error(&e)); } if let Some(rejection) = reject_oversized_state_body(ctx, filter.config.max_body_bytes)? { @@ -359,15 +370,15 @@ fn reject_oversized_state_body( /// Resolve references in the request body's current input. async fn resolve_current_input( filter: &FileResolveFilter, - ctx: &HttpFilterContext<'_>, parsed: &mut serde_json::Value, + callout_headers: &http::HeaderMap, budget: &mut ResolutionBudget, ) -> Result { Box::pin(resolve_input_with_budget( parsed, &filter.client, filter.config.on_missing, - &ctx.request.headers, + callout_headers, filter.url_resolver.as_ref(), budget, )) @@ -379,6 +390,7 @@ async fn update_state( filter: &FileResolveFilter, ctx: &mut HttpFilterContext<'_>, resolved_body: Option<&serde_json::Value>, + callout_headers: &http::HeaderMap, budget: &mut ResolutionBudget, ) -> Result<(), ResolveError> { match resolved_body { @@ -389,6 +401,7 @@ async fn update_state( &filter.client, filter.config.on_missing, filter.url_resolver.as_ref(), + callout_headers, budget, )) .await @@ -399,6 +412,7 @@ async fn update_state( &filter.client, filter.config.on_missing, filter.url_resolver.as_ref(), + callout_headers, budget, )) .await @@ -437,6 +451,7 @@ async fn sync_state_with_budget( client: &FilesApiClient, on_missing: OnMissing, url_resolver: Option<&FileUrlResolver>, + callout_headers: &http::HeaderMap, budget: &mut ResolutionBudget, ) -> Result<(), ResolveError> { let Some(state) = ctx.extensions.get_mut::() else { @@ -453,7 +468,7 @@ async fn sync_state_with_budget( let resolver = HistoryResolver { client, on_missing, - request_headers: &ctx.request.headers, + request_headers: callout_headers, url_resolver, }; @@ -477,16 +492,28 @@ async fn sync_state( on_missing: OnMissing, ) -> Result<(), ResolveError> { let mut budget = client.resolution_budget(); - sync_state_with_budget(ctx, resolved_body, client, on_missing, None, &mut budget).await + let callout_headers = client.callout_headers(ctx); + sync_state_with_budget( + ctx, + resolved_body, + client, + on_missing, + None, + &callout_headers, + &mut budget, + ) + .await } /// Resolve file references in rehydrated history when the /// current input did not require a body rewrite. +#[expect(clippy::too_many_arguments, reason = "threading resolver through state sync")] async fn resolve_state_history( ctx: &mut HttpFilterContext<'_>, client: &FilesApiClient, on_missing: OnMissing, url_resolver: Option<&FileUrlResolver>, + callout_headers: &http::HeaderMap, budget: &mut ResolutionBudget, ) -> Result<(), ResolveError> { let Some(state) = ctx.extensions.get_mut::() else { @@ -497,7 +524,7 @@ async fn resolve_state_history( let resolver = HistoryResolver { client, on_missing, - request_headers: &ctx.request.headers, + request_headers: callout_headers, url_resolver, }; diff --git a/apis/src/openai/responses/file_resolve/resolve.rs b/apis/src/openai/responses/file_resolve/resolve.rs index 95f144d50c..5ef6338e92 100644 --- a/apis/src/openai/responses/file_resolve/resolve.rs +++ b/apis/src/openai/responses/file_resolve/resolve.rs @@ -26,7 +26,7 @@ use super::{ config::OnMissing, resolve_url::{FileUrlResolver, redact_url}, }; -use crate::openai::api_client::{ApiClient, ApiClientError}; +use crate::openai::api_client::{ApiClient, ApiClientError, Correlation}; /// Files API path prefix used in resource URL construction. const FILES_PATH_PREFIX: &str = "v1/files"; @@ -348,6 +348,18 @@ impl FilesApiClient { } } + /// Build the header set sent on every Files API callout for one + /// downstream request. + /// + /// Resolved once per request at the filter boundary — the only + /// place with access to the filter context — then threaded into + /// each callout, so correlation is established once rather than + /// per file reference. + pub(crate) fn callout_headers(&self, ctx: &praxis_filter::HttpFilterContext<'_>) -> http::HeaderMap { + self.client + .callout_headers(&ctx.request.headers, &Correlation::from_filter_context(ctx)) + } + /// Create request-scoped resolution limits and cache state. pub(crate) fn resolution_budget(&self) -> ResolutionBudget { ResolutionBudget { diff --git a/apis/src/openai/responses/file_resolve/tests.rs b/apis/src/openai/responses/file_resolve/tests.rs index e2d0980e8f..16c8235d29 100644 --- a/apis/src/openai/responses/file_resolve/tests.rs +++ b/apis/src/openai/responses/file_resolve/tests.rs @@ -639,6 +639,60 @@ async fn resolves_history_when_current_input_has_no_file_id() { } } +#[tokio::test] +async fn current_input_and_history_callouts_share_one_span() { + // Distinct file ids, so the request-wide cache cannot hide a + // second correlation resolution behind a cache hit: current input + // and rehydrated history each make real callouts. + let (files_api_url, seen) = start_recording_files_api_stub(); + let filter = make_filter_for_url(&files_api_url); + let req = Box::leak(Box::new(crate::test_utils::make_request( + http::Method::POST, + "/v1/responses", + ))); + let mut ctx = crate::test_utils::make_filter_context(req); + ctx.set_metadata("openai_responses_format.format", "openai_responses"); + + let request_body = json!({ + "input": [{ + "type": "message", + "role": "user", + "content": [{"type": "input_file", "file_id": "file-current"}] + }] + }); + let history = json!({ + "type": "message", + "role": "user", + "content": [{"type": "input_file", "file_id": "file-history"}] + }); + let mut state = ResponsesState::from_request_body(request_body.clone()); + state.messages.insert(0, history.clone()); + state.persisted_messages.insert(0, history); + ctx.extensions.insert(state); + let mut body = Some(Bytes::from(serde_json::to_vec(&request_body).unwrap())); + + let action = filter.on_request_body(&mut ctx, &mut body, true).await.unwrap(); + assert!( + matches!(action, FilterAction::Continue), + "resolution across input and history should continue the request" + ); + + let traceparents = seen.lock().unwrap().clone(); + assert!( + traceparents.len() > 1, + "expected callouts for both current input and history, saw {traceparents:?}" + ); + let spans: std::collections::HashSet<&str> = traceparents + .iter() + .map(|value| value.split('-').nth(2).unwrap_or_default()) + .collect(); + assert_eq!( + spans.len(), + 1, + "file resolution is one delegation hop, so every phase's callouts must share a span-id, saw {traceparents:?}" + ); +} + #[tokio::test] async fn mirrored_history_has_independent_inline_budget() { let files_api_url = start_files_api_stub(); @@ -659,10 +713,18 @@ async fn mirrored_history_has_independent_inline_budget() { state.persisted_messages.push(history); ctx.extensions.insert(state); let mut budget = client.resolution_budget(); + let callout_headers = client.callout_headers(&ctx); - resolve_state_history(&mut ctx, &client, OnMissing::Reject, None, &mut budget) - .await - .unwrap(); + resolve_state_history( + &mut ctx, + &client, + OnMissing::Reject, + None, + &callout_headers, + &mut budget, + ) + .await + .unwrap(); let state = ctx.extensions.get::().unwrap(); assert_eq!(state.messages[0]["content"][0]["file_data"], "aGlzdG9yeQ=="); @@ -813,6 +875,58 @@ fn start_files_api_stub() -> String { format!("http://{address}") } +/// Files API stub that also records the `traceparent` of every +/// callout it serves. +fn start_recording_files_api_stub() -> (String, std::sync::Arc>>) { + let listener = TcpListener::bind("127.0.0.1:0").unwrap(); + let address = listener.local_addr().unwrap(); + let seen = std::sync::Arc::new(std::sync::Mutex::new(Vec::new())); + let recorder = std::sync::Arc::clone(&seen); + + std::thread::spawn(move || { + for stream in listener.incoming().flatten() { + let recorder = std::sync::Arc::clone(&recorder); + std::thread::spawn(move || serve_and_record(stream, &recorder)); + } + }); + + (format!("http://{address}"), seen) +} + +fn serve_and_record(mut stream: std::net::TcpStream, seen: &std::sync::Mutex>) { + let mut request = [0_u8; 4096]; + let read = stream.read(&mut request).unwrap(); + let raw = String::from_utf8_lossy(&request[..read]).into_owned(); + + let traceparent = raw + .lines() + .filter_map(|line| line.split_once(':')) + .find(|(name, _)| name.trim().eq_ignore_ascii_case("traceparent")) + .map(|(_, value)| value.trim().to_owned()) + .unwrap_or_default(); + seen.lock().unwrap().push(traceparent); + + let path = raw + .lines() + .next() + .and_then(|line| line.split_whitespace().nth(1)) + .unwrap(); + let (content_type, body): (&str, &[u8]) = if path.ends_with("/content") { + ("text/plain", b"history") + } else { + ( + "application/json", + br#"{"id":"file-history","filename":"history.txt","content_type":"text/plain","bytes":7}"#, + ) + }; + let headers = format!( + "HTTP/1.1 200 OK\r\nContent-Type: {content_type}\r\nContent-Length: {}\r\nConnection: close\r\n\r\n", + body.len() + ); + stream.write_all(headers.as_bytes()).unwrap(); + stream.write_all(body).unwrap(); +} + fn serve_file_request(mut stream: std::net::TcpStream) { let mut request = [0_u8; 4096]; let read = stream.read(&mut request).unwrap(); diff --git a/apis/src/openai/responses/file_search_callout/client.rs b/apis/src/openai/responses/file_search_callout/client.rs index f2ca9fa691..85e3c230c4 100644 --- a/apis/src/openai/responses/file_search_callout/client.rs +++ b/apis/src/openai/responses/file_search_callout/client.rs @@ -14,7 +14,10 @@ use http::HeaderMap; use serde::{Deserialize, Serialize, de::Visitor}; use serde_json::Value; -use crate::openai::{api_client::ApiClient, responses::config_validation::FailureMode}; +use crate::openai::{ + api_client::{ApiClient, Correlation}, + responses::config_validation::FailureMode, +}; // ----------------------------------------------------------------------------- // Constants @@ -379,6 +382,22 @@ impl FileSearchClient { } } + /// Build the header set sent on every vector-store callout for + /// one downstream request. + /// + /// Resolved once per request at the filter boundary — the only + /// place with access to the filter context — then threaded into + /// each callout of the fan-out, so correlation is established + /// once rather than per vector store. + pub(crate) fn callout_headers(&self, ctx: &praxis_filter::HttpFilterContext<'_>) -> HeaderMap { + // On IRR continuation iterations the synthetic request lacks client + // credentials, so the allowlist is applied to the original request's + // headers rather than the synthetic ones. + let request_headers = super::callout_request_headers(ctx); + self.api_client + .callout_headers(&request_headers, &Correlation::from_filter_context(ctx)) + } + /// Search multiple vector stores with bounded concurrency and aggregation. #[expect( clippy::too_many_lines, @@ -388,7 +407,7 @@ impl FileSearchClient { &self, specs: &[SearchSpec<'_>], call_count: usize, - request_headers: &HeaderMap, + callout_headers: &HeaderMap, ) -> SearchBatch { let mut batch = SearchBatch::new(call_count); let mut consumed_response_bytes = 0_usize; @@ -424,7 +443,7 @@ impl FileSearchClient { }; let futures = chunk .iter() - .map(|spec| self.search_one(spec, execution_started, Arc::clone(&admission), request_headers)); + .map(|spec| self.search_one(spec, execution_started, Arc::clone(&admission), callout_headers)); let chunk_results = futures::future::join_all(futures).await; let chunk_failed = merge_chunk_results( &mut batch, @@ -461,13 +480,13 @@ impl FileSearchClient { spec: &SearchSpec<'_>, execution_started: Instant, response_admission: Arc, - request_headers: &HeaderMap, + callout_headers: &HeaderMap, ) -> Result { deadline_remaining(self.timeout, execution_started, spec.store_id)?; let request = self.build_request(spec, execution_started)?; deadline_remaining(self.timeout, execution_started, spec.store_id)?; let body = self - .execute_request(request, spec.store_id, execution_started, request_headers) + .execute_request(request, spec.store_id, execution_started, callout_headers) .await?; parse_response_body_with_deadline( body, @@ -533,13 +552,13 @@ impl FileSearchClient { request: PreparedSearchRequest, store_id: &str, execution_started: Instant, - request_headers: &HeaderMap, + callout_headers: &HeaderMap, ) -> Result { let remaining = deadline_remaining(self.timeout, execution_started, store_id)?; tokio::time::timeout( remaining, self.api_client - .post_json_bytes(request.url, request.body, request_headers), + .post_json_bytes(request.url, request.body, callout_headers), ) .await .map_err(|_elapsed| execution_deadline_error(store_id)) diff --git a/apis/src/openai/responses/file_search_callout/mod.rs b/apis/src/openai/responses/file_search_callout/mod.rs index 88165b80e9..2275959602 100644 --- a/apis/src/openai/responses/file_search_callout/mod.rs +++ b/apis/src/openai/responses/file_search_callout/mod.rs @@ -228,7 +228,7 @@ impl FileSearchCalloutFilter { clippy::too_many_lines, reason = "separates global, per-call, and transport planning failures" )] - async fn execute_plan(&self, plan: &SearchPlan, request_headers: &HeaderMap) -> SearchBatch { + async fn execute_plan(&self, plan: &SearchPlan, callout_headers: &HeaderMap) -> SearchBatch { if let Some(message) = plan.planning_error { return SearchBatch::with_failures( plan.calls.len(), @@ -257,7 +257,7 @@ impl FileSearchCalloutFilter { let mut batch = if specs.is_empty() { SearchBatch::new(plan.calls.len()) } else { - self.client.search(&specs, plan.calls.len(), request_headers).await + self.client.search(&specs, plan.calls.len(), callout_headers).await }; batch.failures.extend(planning_failures); batch @@ -283,33 +283,61 @@ impl FileSearchCalloutFilter { ))) } + /// Build the header set for this request's callouts, establishing + /// the shared trace context first. + /// + /// Delegated callouts can run in the `StreamBuffer` pre-read + /// phase, ahead of header-phase filters, so this filter is often + /// the hop that initializes the request's trace context rather + /// than the one that inherits it. + /// + /// One header set covers the whole fan-out: the callouts of one + /// request share a delegation span, and the operator's + /// `forward_headers` allowlist is applied once here rather than + /// the raw downstream headers reaching the vector store. + fn prepare_callout_headers(&self, ctx: &mut HttpFilterContext<'_>) -> HeaderMap { + drop(crate::correlation::TraceContext::get_or_init(ctx)); + self.client.callout_headers(ctx) + } + /// Execute pending calls before the next inference body is serialized. - #[expect( - clippy::too_many_lines, - reason = "linear sequence: plan → callout → apply → size check" - )] async fn execute_pending(&self, ctx: &mut HttpFilterContext<'_>) -> Result { if let Some(rejection) = unsupported_streaming_rejection(ctx) { return Ok(rejection); } + // Checked before planning, so trace context is established + // only for requests that actually call out, and before the + // plan borrows `ctx` immutably. + if !ctx.extensions.get::().is_some_and(has_pending_calls) { + return Ok(FilterAction::Continue); + } + let callout_headers = self.prepare_callout_headers(ctx); + let Some(state) = ctx.extensions.get::() else { return Ok(FilterAction::Continue); }; let plan = build_search_plan(state); - if !plan.has_pending_calls { - return Ok(FilterAction::Continue); - } - let hdrs = callout_request_headers(ctx); - let batch = self.execute_plan(&plan, &hdrs).await; + let batch = self.execute_plan(&plan, &callout_headers).await; if let Some(rejection) = self.failure_rejection(&batch) { return Ok(rejection); } + self.commit_batch(ctx, &plan, &batch) + } + + /// Commit a completed batch into request state and advance the + /// iteration counter. + fn commit_batch( + &self, + ctx: &mut HttpFilterContext<'_>, + plan: &SearchPlan, + batch: &SearchBatch, + ) -> Result { let framework_bytes = retained_iteration_bytes(ctx); let state = ctx .extensions .get_mut::() .ok_or_else(|| -> FilterError { "openai_file_search_callout: ResponsesState disappeared".into() })?; - if let Err(rejection) = Self::apply_batch(state, &plan, &batch) { + if let Err(rejection) = Self::apply_batch(state, plan, batch) { return Ok(rejection); } if !continuation_state_fits(framework_bytes, state, self.max_state_bytes, 0) { @@ -873,9 +901,6 @@ struct SearchPlan { /// Metadata filters shared by every search spec. filters: Option, - /// Whether the response contained any pending call before local caps. - has_pending_calls: bool, - /// Maximum number of aggregate results per call. max_num_results: Option, @@ -1049,10 +1074,18 @@ struct FileSearchToolDef { vector_store_count: usize, } +/// Return whether the state holds any file-search call awaiting a +/// callout. +/// +/// Shared with the filter boundary, which needs the answer before +/// planning borrows the state, so the two cannot drift apart. +fn has_pending_calls(state: &ResponsesState) -> bool { + state.output_items().iter().any(is_pending_file_search_call) +} + /// Build an owned plan for every pending call before applying the fan-out cap. fn build_search_plan(state: &ResponsesState) -> SearchPlan { let tool = extract_file_search_tool_def(&state.tools); - let has_pending_calls = state.output_items().iter().any(is_pending_file_search_call); let call_budget = remaining_file_search_call_budget(state); let mut calls = pending_calls(state, tool.vector_store_count, call_budget); let spec_coordinates = schedule_searches(&mut calls, tool.vector_store_ids.len()); @@ -1060,7 +1093,6 @@ fn build_search_plan(state: &ResponsesState) -> SearchPlan { SearchPlan { calls, filters: tool.filters, - has_pending_calls, max_num_results: tool.max_num_results, planning_error: tool.planning_error, ranking_options: tool.ranking_options, diff --git a/apis/src/openai/responses/file_search_callout/tests.rs b/apis/src/openai/responses/file_search_callout/tests.rs index ad998bb801..4d86e6c833 100644 --- a/apis/src/openai/responses/file_search_callout/tests.rs +++ b/apis/src/openai/responses/file_search_callout/tests.rs @@ -1338,6 +1338,52 @@ async fn searches_multiple_stores_concurrently() { ); } +/// Callouts must be built from the configured allowlist plus +/// correlation, never from the raw downstream header map. +/// +/// The vector-store search is a delegated call across a security +/// boundary: passing the downstream headers through would leak every +/// client header the operator did not opt into — credentials +/// included — and would carry no correlation. +#[tokio::test] +async fn search_callouts_apply_the_forward_allowlist_and_carry_correlation() { + let server = MockServer::json(200, &json!({"data": []})); + let filter = make_filter(server.port, "forward_headers:\n - authorization\n"); + let mut ctx = make_context_with_headers( + Some(one_pending_state(&["vs-a"])), + &[ + ("authorization", "Bearer allowlisted"), + ("cookie", "session=not-allowlisted"), + ], + ); + + assert!(matches!( + filter.on_request(&mut ctx).await.unwrap(), + FilterAction::Continue + )); + + let requests = server.requests(); + assert_eq!(requests.len(), 1); + let sent = requests[0].to_ascii_lowercase(); + + assert!( + sent.contains("authorization: bearer allowlisted"), + "allowlisted header should reach the vector store: {sent}" + ); + assert!( + !sent.contains("cookie"), + "headers outside the allowlist must not reach the vector store: {sent}" + ); + assert!( + sent.contains("x-request-id:"), + "delegated call must carry correlation: {sent}" + ); + assert!( + sent.contains("traceparent:"), + "delegated call must carry trace context: {sent}" + ); +} + #[tokio::test] async fn aggregate_results_are_score_sorted_and_limited_to_top_k() { let server = MockServer::routes([ @@ -1834,10 +1880,21 @@ fn make_concrete_filter(port: u16, extra: &str) -> FileSearchCalloutFilter { } fn make_context(state: Option) -> HttpFilterContext<'static> { - let request = Box::leak(Box::new(crate::test_utils::make_request( - http::Method::POST, - "/v1/responses", - ))); + make_context_with_headers(state, &[]) +} + +fn make_context_with_headers( + state: Option, + downstream_headers: &[(&str, &str)], +) -> HttpFilterContext<'static> { + let mut built = crate::test_utils::make_request(http::Method::POST, "/v1/responses"); + for (name, value) in downstream_headers { + built.headers.insert( + http::HeaderName::from_bytes(name.as_bytes()).unwrap(), + http::HeaderValue::from_str(value).unwrap(), + ); + } + let request = Box::leak(Box::new(built)); let mut ctx = crate::test_utils::make_filter_context(request); ctx.set_metadata("openai_responses_format.stream", "false"); if let Some(state) = state { diff --git a/docs/filters/openai_file_resolve.md b/docs/filters/openai_file_resolve.md index 02cfabc4ef..4cce1662d1 100644 --- a/docs/filters/openai_file_resolve.md +++ b/docs/filters/openai_file_resolve.md @@ -18,7 +18,7 @@ This filter resolves references inside Responses requests; it does not proxy cli | `allow_private_files_api_url` | bool | no | Allow `files_api_url` to target private, loopback, link-local, or DNS-name hosts. Default `false` rejects SSRF-sensitive targets; set to `true` in development or when the Files API is an internal service on a private network. | | `allow_pre_security_callout` | bool | no | Allow Files API callouts from the `StreamBuffer` pre-read phase, before header-phase security filters execute. This must be explicitly enabled only when an outer trust boundary authenticates and authorizes requests before they reach this listener. Forwarded headers are the original downstream values, not mutations from request filters. | | `files_api_url` | string | yes | Base URL of the Files API endpoint. Example: `http://files-api:8321` | -| `forward_headers` | string[] | no | Headers to forward from the original request to the Files API for authentication and tenant isolation. No downstream headers are forwarded by default. | +| `forward_headers` | string[] | no | Headers to forward from the original request to the Files API for authentication and tenant isolation. No downstream headers are forwarded by default. The correlation headers `x-request-id` and `traceparent` are always sent and do not need to be listed here. | | `max_body_bytes` | integer | no | Maximum body size in bytes for `StreamBuffer` mode. | | `max_file_references` | integer | no | Maximum number of distinct content-part / `file_id` pairs to resolve in one request, including rehydrated history. | | `on_missing` | `continue` \| `reject` | no | Behavior when a `file_id` reference cannot be fetched. Does not apply to `file_url`: a failed `file_url` fetch is always rejected, regardless of this setting. | diff --git a/docs/filters/reference.md b/docs/filters/reference.md index 32452bc2f3..9a7875e488 100644 --- a/docs/filters/reference.md +++ b/docs/filters/reference.md @@ -92,3 +92,9 @@ see the [Praxis core filter reference][core-ref]. |--------|-------------| | [`token_count`](token_count.md) | Extracts token usage from AI inference responses and writes unified counts to [`filter_metadata`]. | | [`token_usage_headers`](token_usage_headers.md) | Injects `Praxis-Token-Input`, `Praxis-Token-Output`, and `Praxis-Token-Total` headers into downstream responses when token usage data is present in [`filter_metadata`]. | + +### Trace Context + +| Filter | Description | +|--------|-------------| +| [`trace_context`](trace_context.md) | Propagates correlation and W3C trace context to the upstream request. | diff --git a/docs/filters/trace_context.md b/docs/filters/trace_context.md new file mode 100644 index 0000000000..3c74016ef2 --- /dev/null +++ b/docs/filters/trace_context.md @@ -0,0 +1,16 @@ + + + +# `trace_context` + +Propagates correlation and W3C trace context to the upstream request. + +## Configuration Notes + +Register early: filters that make delegated calls read what this injects. The one filter that must run even earlier is the `request_id` core builtin, if configured — it reads only the client's headers, so running it after this filter mints a second ID that reaches the backend on the forwarded request while the delegated calls keep the first. + +## Example + +```yaml +filter: trace_context +``` diff --git a/examples/README.md b/examples/README.md index 94a41b2ddc..16c05c1258 100644 --- a/examples/README.md +++ b/examples/README.md @@ -38,6 +38,7 @@ before sending requests. | [time-to-first-token.yaml](configs/time-to-first-token.yaml) | Measures the elapsed time from request receipt to the first non-empty SSE body chunk and records a praxis_ai_ttft_seconds Prometheus histogram labeled by model | | [token-counting.yaml](configs/token-counting.yaml) | Extracts token usage from AI inference responses (streaming and non-streaming) and makes counts available to downstream filters via filter metadata as token.input, token.output, and token.total | | [token-usage-headers.yaml](configs/token-usage-headers.yaml) | Inject Praxis-Token-Input, Praxis-Token-Output, and Praxis-Token-Total headers into downstream responses when token counts are available in filter metadata | +| [trace-context.yaml](configs/trace-context.yaml) | Propagates x-request-id and W3C traceparent across every leg of an AI request, so a single request is traceable end to end | ### Anthropic diff --git a/examples/configs/trace-context.yaml b/examples/configs/trace-context.yaml new file mode 100644 index 0000000000..2e7517de70 --- /dev/null +++ b/examples/configs/trace-context.yaml @@ -0,0 +1,90 @@ +# Trace Context Propagation +# +# Propagates x-request-id and W3C traceparent across every leg of an +# AI request, so a single request is traceable end to end. +# +# An AI request reaches the backend over two paths: +# +# forwarded client -> proxy -> backend (the request itself) +# delegated proxy -> backend (callouts the proxy makes +# while handling it, such as +# Files API lookups) +# +# The trace_context filter stamps the forwarded request. Filters that +# make delegated calls inject the same identifiers automatically, +# without configuration, and draw their own span-id from the shared +# trace-id. Both legs land in one trace while remaining individually +# measurable. +# +# A valid inbound traceparent is continued. An absent or malformed one +# starts a new sampled trace: client-supplied values are validated +# rather than forwarded. +# +# Placement: register trace_context early. Filters that make delegated +# calls read the injected values, so anything running before it +# resolves trace context independently and lands in a different trace. +# +# The one exception is the request_id core builtin: when it is used, it +# must come before trace_context. It reads only the client's headers, +# so running it later mints a second ID that wins on the forwarded +# request while the delegated calls keep the first. Ordered as below, +# trace_context adopts the ID request_id injected and every leg agrees. +# +# Usage: +# cargo run -p praxis-ai-proxy -- -c examples/configs/trace-context.yaml +# +# curl -X POST http://localhost:8080/v1/responses \ +# -H "Content-Type: application/json" \ +# -H "traceparent: 00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01" \ +# -d '{"model":"gpt-4o","input":[{"role":"user","content":[{"type":"input_file","file_id":"file-abc"}]}]}' +# +# The Files API callout and the forwarded /v1/responses request both +# arrive carrying trace-id 4bf92f3577b34da6a3ce929d0e0e4736 under +# different span-ids. + +listeners: + - name: ai-gateway + address: "127.0.0.1:8080" + filter_chains: [traced] + +filter_chains: + - name: traced + filters: + # Before trace_context: it adopts the ID injected here, so the + # forwarded request, the delegated calls, and the echoed + # response header all carry one x-request-id. + - filter: request_id + + # Early: everything downstream reads what this injects. + - filter: trace_context + + - filter: openai_responses_format + on_invalid: continue + headers: + format: x-praxis-ai-format + model: x-praxis-ai-model + + # Delegated calls. Correlation headers are sent unconditionally, + # so forward_headers only needs the application headers. + - filter: openai_file_resolve + files_api_url: "http://127.0.0.1:8321" + allow_private_files_api_url: true + allow_pre_security_callout: true + forward_headers: + - authorization + on_missing: reject + timeout_ms: 10000 + + - filter: router + routes: + - path_prefix: "/v1/files" + cluster: files-api + - path_prefix: "/" + cluster: inference-backend + + - filter: load_balancer + clusters: + - name: files-api + endpoints: ["127.0.0.1:8321"] + - name: inference-backend + endpoints: ["127.0.0.1:8321"] diff --git a/filters/src/lib.rs b/filters/src/lib.rs index 9508f2a42a..7bc801e522 100644 --- a/filters/src/lib.rs +++ b/filters/src/lib.rs @@ -18,6 +18,7 @@ mod register; pub mod routing; mod time_to_first_token; mod token_usage; +mod trace_context; pub use agentic::{a2a::A2aFilter, mcp::McpFilter}; pub use guardrails::AiGuardrailsFilter; @@ -27,6 +28,7 @@ pub use register::{build_ai_registry, register_ai_filters}; pub use routing::{CredentialInjectFilter, IntelligentRouteFilter, ProviderRouteFilter}; pub use time_to_first_token::TimeToFirstTokenFilter; pub use token_usage::{TokenCountFilter, TokenUsageHeadersFilter}; +pub use trace_context::TraceContextFilter; // ----------------------------------------------------------------------------- // Test Utilities diff --git a/filters/src/register.rs b/filters/src/register.rs index 836b4326ac..20726d16dc 100644 --- a/filters/src/register.rs +++ b/filters/src/register.rs @@ -9,6 +9,7 @@ use praxis_filter::FilterRegistry; use crate::{ A2aFilter, AiGuardrailsFilter, CredentialInjectFilter, IntelligentRouteFilter, McpFilter, ModelToHeaderFilter, PromptEnrichFilter, ProviderRouteFilter, TimeToFirstTokenFilter, TokenCountFilter, TokenUsageHeadersFilter, + TraceContextFilter, }; /// Register all in-tree AI HTTP filters into `registry`. @@ -93,6 +94,10 @@ fn register_general_ai_filters(registry: &mut FilterRegistry) { @register registry, http "time_to_first_token" => TimeToFirstTokenFilter::from_config ); + praxis_filter::register_filters!( + @register registry, + http "trace_context" => TraceContextFilter::from_config + ); } /// Register intelligent routing filters. @@ -360,6 +365,7 @@ mod tests { "openai_responses_validate", "responses_to_chat_completions", "a2a", + "trace_context", "intelligent_route", "provider_route", "credential_inject", diff --git a/filters/src/trace_context/mod.rs b/filters/src/trace_context/mod.rs new file mode 100644 index 0000000000..2dc7afb3bb --- /dev/null +++ b/filters/src/trace_context/mod.rs @@ -0,0 +1,182 @@ +// SPDX-License-Identifier: MIT +// Copyright (c) 2026 Praxis Contributors + +//! W3C trace-context propagation for forwarded requests. +//! +//! Stamps `x-request-id` and `traceparent` onto the request sent +//! upstream, so the forwarded request joins the same trace as the +//! delegated calls made while handling it. +//! +//! Without this filter only delegated calls carry trace context, +//! which traces the file lookups but not the inference request they +//! belong to. +//! +//! # Placement +//! +//! Register early in the chain. Filters that make delegated calls +//! read the injected values, so anything running before this filter +//! resolves trace context independently and lands in a different +//! trace. +//! +//! One filter must come *earlier* still: the `request_id` core +//! builtin, when it is configured at all. That builtin reads only the +//! client's headers, so it generates a second, unrelated ID when it +//! runs after this filter. Pending header mutations are applied in +//! chain order with last-write-wins, so the forwarded request would +//! carry the builtin's ID while the delegated calls and the echoed +//! response header keep this filter's — exactly the split correlation +//! the filter exists to prevent. With `request_id` first, this filter +//! adopts the ID it injected and every leg agrees. +//! +//! The mismatch is detected at response time and logged, since a +//! filter cannot see what the chain places after it. +//! +//! # Behavior +//! +//! A valid inbound `traceparent` is continued: its trace-id and +//! flags carry forward under a fresh span-id for the upstream hop. +//! An absent or malformed one starts a new sampled trace — client- +//! supplied values are validated rather than forwarded, since they +//! would otherwise reach the telemetry backend unchecked. +//! +//! `x-request-id` follows the same precedence as the `request_id` +//! core builtin: client-supplied, then injected, then generated. +//! Running both filters is safe when `request_id` runs first; this +//! one then reuses whatever it injected. +//! +//! # YAML +//! +//! ```yaml +//! filter: trace_context +//! ``` + +use std::borrow::Cow; + +use async_trait::async_trait; +use praxis_ai_apis::correlation::TraceContext; +use praxis_filter::{EmptyFilterConfig, FilterAction, FilterError, HttpFilter, HttpFilterContext, parse_filter_config}; +use tracing::{debug, warn}; + +// ----------------------------------------------------------------------------- +// Constants +// ----------------------------------------------------------------------------- + +/// Header carrying the request correlation ID. +const REQUEST_ID: &str = "x-request-id"; + +// ----------------------------------------------------------------------------- +// TraceContextFilter +// ----------------------------------------------------------------------------- + +/// Propagates correlation and W3C trace context to the upstream +/// request. +/// +/// Register early: filters that make delegated calls read what this +/// injects. The one filter that must run even earlier is the +/// `request_id` core builtin, if configured — it reads only the +/// client's headers, so running it after this filter mints a second +/// ID that reaches the backend on the forwarded request while the +/// delegated calls keep the first. +/// +/// # YAML +/// +/// ```yaml +/// filter: trace_context +/// ``` +pub struct TraceContextFilter; + +impl TraceContextFilter { + /// Create a filter from parsed YAML config. + /// + /// # Errors + /// + /// Returns [`FilterError`] if the YAML config is invalid. + pub fn from_config(config: &serde_yaml::Value) -> Result, FilterError> { + let _: EmptyFilterConfig = parse_filter_config("trace_context", config)?; + Ok(Box::new(Self)) + } +} + +#[async_trait] +impl HttpFilter for TraceContextFilter { + fn name(&self) -> &'static str { + "trace_context" + } + + async fn on_request(&self, ctx: &mut HttpFilterContext<'_>) -> Result { + let context = TraceContext::get_or_init(ctx); + let headers = context.headers_for_hop(ctx); + + for (name, value) in headers { + // Replace any value this filter already injected, so a + // re-entered chain does not accumulate duplicates. A + // client-supplied header lives on ctx.request.headers + // and is untouched here; the pipeline applies these + // pending mutations over it. + ctx.extra_request_headers + .retain(|(existing, _)| !existing.eq_ignore_ascii_case(name.as_str())); + ctx.extra_request_headers + .push((Cow::Owned(name.as_str().to_owned()), value)); + } + + debug!(request_id = %context.request_id(), "propagating trace context upstream"); + + Ok(FilterAction::Continue) + } + + async fn on_response(&self, ctx: &mut HttpFilterContext<'_>) -> Result { + // Every request filter has run by now, so this is the first + // point at which a later filter's competing request ID is + // visible. + let Some(shared) = ctx.extensions.get::() else { + return Ok(FilterAction::Continue); + }; + + if let Some(forwarded) = competing_request_id(ctx, shared.request_id()) { + warn!( + correlated = %shared.request_id(), + forwarded = %forwarded, + "another filter injected a different x-request-id; the forwarded request and the \ + delegated calls are in different correlation IDs. Register `request_id` before \ + `trace_context`." + ); + } + + Ok(FilterAction::Continue) + } +} + +// ----------------------------------------------------------------------------- +// Utility Functions +// ----------------------------------------------------------------------------- + +/// Find a pending `x-request-id` that disagrees with the shared +/// context. +/// +/// Pending mutations are applied last-write-wins, so the value that +/// actually reaches the backend is the final one. Returns `None` when +/// every pending value agrees with the correlated ID. +fn competing_request_id(ctx: &HttpFilterContext<'_>, correlated: &str) -> Option { + let forwarded = ctx + .extra_request_headers + .iter() + .rfind(|(name, _)| name.eq_ignore_ascii_case(REQUEST_ID)) + .map(|(_, value)| value.clone())?; + + (forwarded != correlated).then_some(forwarded) +} + +// ----------------------------------------------------------------------------- +// Tests +// ----------------------------------------------------------------------------- + +#[cfg(test)] +#[expect(clippy::allow_attributes, reason = "blanket test suppressions")] +#[allow( + clippy::unwrap_used, + clippy::expect_used, + clippy::indexing_slicing, + clippy::panic, + reason = "tests" +)] +mod tests; diff --git a/filters/src/trace_context/tests.rs b/filters/src/trace_context/tests.rs new file mode 100644 index 0000000000..7d168e88ba --- /dev/null +++ b/filters/src/trace_context/tests.rs @@ -0,0 +1,226 @@ +// SPDX-License-Identifier: MIT +// Copyright (c) 2026 Praxis Contributors + +use http::{HeaderName, HeaderValue, Method}; +use praxis_filter::Request; + +use super::*; +use crate::test_utils::{make_filter_context, make_request}; + +#[tokio::test] +async fn injects_both_correlation_headers_upstream() { + let filter = make_filter(); + let req = request_with(&[]); + let mut ctx = make_filter_context(&req); + + let action = filter.on_request(&mut ctx).await.unwrap(); + + assert!(matches!(action, FilterAction::Continue), "filter should continue"); + assert_eq!( + injected(&ctx, "x-request-id").len(), + 32, + "should inject a generated request ID" + ); + let traceparent = injected(&ctx, "traceparent"); + let parts: Vec<&str> = traceparent.split('-').collect(); + assert_eq!(parts.len(), 4, "traceparent should be well-formed: {traceparent}"); + assert_eq!(parts[1].len(), 32, "trace-id should be 32 hex: {traceparent}"); + assert_eq!(parts[2].len(), 16, "span-id should be 16 hex: {traceparent}"); +} + +#[tokio::test] +async fn continues_client_supplied_trace() { + let filter = make_filter(); + let req = request_with(&[("traceparent", "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01")]); + let mut ctx = make_filter_context(&req); + + drop(filter.on_request(&mut ctx).await.unwrap()); + + let traceparent = injected(&ctx, "traceparent"); + let parts: Vec<&str> = traceparent.split('-').collect(); + assert_eq!( + parts[1], "4bf92f3577b34da6a3ce929d0e0e4736", + "client trace-id should be continued" + ); + assert_ne!(parts[2], "00f067aa0ba902b7", "upstream hop should get its own span-id"); +} + +#[tokio::test] +async fn reuses_request_id_injected_by_request_id_filter() { + let filter = make_filter(); + let req = request_with(&[]); + let mut ctx = make_filter_context(&req); + ctx.extra_request_headers + .push((Cow::Borrowed("X-Request-ID"), "from-request-id-filter".to_owned())); + + drop(filter.on_request(&mut ctx).await.unwrap()); + + assert_eq!( + injected(&ctx, "x-request-id"), + "from-request-id-filter", + "should reuse the ID the request_id builtin injected" + ); + assert_eq!( + ctx.extra_request_headers + .iter() + .filter(|(name, _)| name.eq_ignore_ascii_case("x-request-id")) + .count(), + 1, + "should not duplicate the request ID header" + ); +} + +#[tokio::test] +async fn request_id_registered_first_keeps_every_leg_on_one_id() { + // The supported order: request_id injects, trace_context adopts. + // What the backend sees on the forwarded request is then the same + // ID the delegated calls resolve from the shared context. + let filter = make_filter(); + let req = request_with(&[]); + let mut ctx = make_filter_context(&req); + ctx.extra_request_headers + .push((Cow::Borrowed("X-Request-ID"), "from-request-id-filter".to_owned())); + + drop(filter.on_request(&mut ctx).await.unwrap()); + let shared = ctx.extensions.get::().unwrap().request_id().to_owned(); + + assert_eq!( + injected(&ctx, "x-request-id"), + shared, + "forwarded request and shared context must carry one ID" + ); + assert!( + competing_request_id(&ctx, &shared).is_none(), + "the supported order must not report a conflict" + ); +} + +#[tokio::test] +async fn detects_a_request_id_injected_after_this_filter() { + // The unsupported order: request_id runs later and generates its + // own ID, which wins on the forwarded request under last-write-wins + // while the delegated calls keep the correlated one. + let filter = make_filter(); + let req = request_with(&[]); + let mut ctx = make_filter_context(&req); + + drop(filter.on_request(&mut ctx).await.unwrap()); + let shared = ctx.extensions.get::().unwrap().request_id().to_owned(); + ctx.extra_request_headers + .push((Cow::Borrowed("X-Request-ID"), "generated-later".to_owned())); + + assert_eq!( + competing_request_id(&ctx, &shared), + Some("generated-later".to_owned()), + "a later filter's competing request ID should be reported" + ); + + let action = filter.on_response(&mut ctx).await.unwrap(); + assert!( + matches!(action, FilterAction::Continue), + "reporting the conflict must not fail the response" + ); +} + +#[tokio::test] +async fn does_not_accumulate_duplicates_when_run_twice() { + let filter = make_filter(); + let req = request_with(&[]); + let mut ctx = make_filter_context(&req); + + drop(filter.on_request(&mut ctx).await.unwrap()); + let first = injected(&ctx, "traceparent"); + drop(filter.on_request(&mut ctx).await.unwrap()); + + assert_eq!( + ctx.extra_request_headers.len(), + 2, + "re-running should replace, not append: {:?}", + ctx.extra_request_headers + ); + let second = injected(&ctx, "traceparent"); + let first_trace: Vec<&str> = first.split('-').collect(); + let second_trace: Vec<&str> = second.split('-').collect(); + assert_eq!( + first_trace[1], second_trace[1], + "trace-id should be stable across re-entry" + ); +} + +#[tokio::test] +async fn discards_malformed_client_traceparent() { + let filter = make_filter(); + let req = request_with(&[("traceparent", "00-not-a-valid-trace-id-01")]); + let mut ctx = make_filter_context(&req); + + drop(filter.on_request(&mut ctx).await.unwrap()); + + let traceparent = injected(&ctx, "traceparent"); + assert_ne!( + traceparent, "00-not-a-valid-trace-id-01", + "malformed value must not be forwarded" + ); + let parts: Vec<&str> = traceparent.split('-').collect(); + assert_eq!(parts.len(), 4, "replacement should be well-formed: {traceparent}"); + assert_eq!(parts[1].len(), 32, "replacement trace-id should be 32 hex"); +} + +#[tokio::test] +async fn response_phase_is_quiet_without_a_shared_context() { + let filter = make_filter(); + let req = request_with(&[]); + let mut ctx = make_filter_context(&req); + + let action = filter.on_response(&mut ctx).await.unwrap(); + + assert!( + matches!(action, FilterAction::Continue), + "a response with no request-phase context should continue" + ); +} + +#[test] +fn from_config_rejects_unknown_fields() { + let yaml: serde_yaml::Value = serde_yaml::from_str("unexpected: true").expect("valid yaml"); + assert!( + TraceContextFilter::from_config(&yaml).is_err(), + "unknown config fields should be rejected" + ); +} + +#[test] +fn filter_name_is_stable() { + let filter = make_filter(); + assert_eq!(filter.name(), "trace_context", "filter name should be trace_context"); +} + +// ----------------------------------------------------------------------------- +// Test Utilities +// ----------------------------------------------------------------------------- + +/// Read an injected header from the pending upstream mutations. +fn injected(ctx: &HttpFilterContext<'_>, name: &str) -> String { + ctx.extra_request_headers + .iter() + .find(|(header, _)| header.eq_ignore_ascii_case(name)) + .map(|(_, value)| value.clone()) + .unwrap_or_default() +} + +/// Build a filter from empty YAML config. +fn make_filter() -> Box { + let yaml: serde_yaml::Value = serde_yaml::from_str("{}").expect("valid empty config"); + TraceContextFilter::from_config(&yaml).expect("filter should build") +} + +/// Build a request carrying the given headers. +fn request_with(headers: &[(&'static str, &str)]) -> Request { + let mut req = make_request(Method::POST, "/v1/responses"); + for (name, value) in headers { + req.headers.insert( + HeaderName::from_static(name), + HeaderValue::from_str(value).expect("valid test header value"), + ); + } + req +} diff --git a/tests/integration/tests/suite/examples/mod.rs b/tests/integration/tests/suite/examples/mod.rs index fe69c83ecf..bca29fba32 100644 --- a/tests/integration/tests/suite/examples/mod.rs +++ b/tests/integration/tests/suite/examples/mod.rs @@ -45,6 +45,7 @@ mod time_to_first_token; mod token_count; mod token_counting; mod token_usage_headers; +mod trace_context; mod vector_stores_routing; mod vllm_agentic_api; mod web_search; diff --git a/tests/integration/tests/suite/examples/trace_context.rs b/tests/integration/tests/suite/examples/trace_context.rs new file mode 100644 index 0000000000..acf3ac6499 --- /dev/null +++ b/tests/integration/tests/suite/examples/trace_context.rs @@ -0,0 +1,269 @@ +// SPDX-License-Identifier: MIT +// Copyright (c) 2026 Praxis Contributors + +//! Integration tests for the `trace_context` example config. +//! +//! Verifies the acceptance criterion of the correlation work: one +//! client request is traceable across both legs it produces — the +//! delegated Files API callout the proxy originates itself, and the +//! forwarded inference request — under a single trace-id with +//! distinct span-ids per leg. + +use std::{ + collections::{HashMap, HashSet}, + io::{Read as _, Write as _}, + net::{TcpListener, TcpStream}, + sync::{Arc, Mutex}, + time::Duration, +}; + +use praxis_test_utils::{free_port, http_send, json_post, load_example_config, parse_status, start_proxy}; + +/// Client trace the proxy is expected to continue. +const CLIENT_TRACEPARENT: &str = "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01"; + +/// Trace-id embedded in [`CLIENT_TRACEPARENT`]. +const CLIENT_TRACE_ID: &str = "4bf92f3577b34da6a3ce929d0e0e4736"; + +/// Span-id embedded in [`CLIENT_TRACEPARENT`]. +const CLIENT_SPAN_ID: &str = "00f067aa0ba902b7"; + +/// Responses request referencing a file, which triggers a delegated +/// Files API callout before the request is forwarded upstream. +const REQUEST_BODY: &str = r#"{ + "model": "gpt-4.1", + "input": [ + { + "type": "message", + "role": "user", + "content": [{"type": "input_file", "file_id": "file-abc"}] + } + ] +}"#; + +/// File metadata returned by the stub. +const FILE_METADATA: &str = r#"{"id":"file-abc","object":"file","bytes":13,"created_at":1750000000,"filename":"test.txt","purpose":"user_data"}"#; + +/// File content returned by the stub. +const FILE_CONTENT: &str = "Hello, world!"; + +/// One request seen by the stub backend. +#[derive(Clone)] +struct SeenRequest { + /// Request line target, e.g. `/v1/files/file-abc`. + path: String, + /// Lowercased header names mapped to their values. + headers: HashMap, +} + +impl SeenRequest { + /// Read a header value by case-insensitive name. + fn header(&self, name: &str) -> Option<&str> { + self.headers.get(&name.to_ascii_lowercase()).map(String::as_str) + } + + /// The `traceparent` this leg carried, or a panic naming the leg. + fn traceparent(&self) -> &str { + self.header("traceparent") + .unwrap_or_else(|| panic!("leg {} must carry traceparent", self.path)) + } +} + +/// Extract field `index` of a `traceparent` value. +fn field(traceparent: &str, index: usize) -> &str { + traceparent.split('-').nth(index).unwrap_or_default() +} + +/// A backend that answers Files API and inference paths while +/// recording every request it sees. +/// +/// Path-aware rather than sequential, because the proxy readiness +/// probe issues its own request before the traced one. +fn start_recording_backend() -> (u16, Arc>>) { + let listener = TcpListener::bind("127.0.0.1:0").expect("stub should bind"); + let port = listener.local_addr().expect("stub should have an address").port(); + let seen: Arc>> = Arc::new(Mutex::new(Vec::new())); + let recorder = Arc::clone(&seen); + + std::thread::spawn(move || { + for stream in listener.incoming().flatten() { + let recorder = Arc::clone(&recorder); + std::thread::spawn(move || handle_request(stream, &recorder)); + } + }); + + (port, seen) +} + +/// Serve one request and record it. +fn handle_request(mut stream: TcpStream, seen: &Arc>>) { + stream + .set_read_timeout(Some(Duration::from_secs(5))) + .expect("read timeout should apply"); + + let mut data = Vec::new(); + let mut buf = [0_u8; 4096]; + loop { + match stream.read(&mut buf) { + Ok(0) | Err(_) => break, + Ok(n) => data.extend_from_slice(&buf[..n]), + } + if data.windows(4).any(|w| w == b"\r\n\r\n") { + break; + } + } + + let raw = String::from_utf8_lossy(&data); + let path = raw + .lines() + .next() + .and_then(|line| line.split_whitespace().nth(1)) + .unwrap_or("/") + .to_owned(); + + let headers = raw + .lines() + .skip(1) + .filter_map(|line| line.split_once(':')) + .map(|(name, value)| (name.trim().to_ascii_lowercase(), value.trim().to_owned())) + .collect(); + + // The readiness probe is not part of the traced request. + if path != "/" { + seen.lock() + .expect("recorder mutex should not be poisoned") + .push(SeenRequest { + path: path.clone(), + headers, + }); + } + + let (content_type, body) = if path.ends_with("/content") { + ("text/plain", FILE_CONTENT) + } else if path.starts_with("/v1/files") { + ("application/json", FILE_METADATA) + } else { + ("application/json", r#"{"id":"resp_1","object":"response","output":[]}"#) + }; + + let header = format!( + "HTTP/1.1 200 OK\r\nContent-Type: {content_type}\r\nContent-Length: {}\r\nConnection: close\r\n\r\n", + body.len() + ); + let _sent = stream.write_all(header.as_bytes()); + let _sent = stream.write_all(body.as_bytes()); +} + +/// Send one file-referencing request through the example config and +/// return the legs the backend saw. +fn capture_legs(client_traceparent: Option<&str>) -> Vec { + let (backend_port, seen) = start_recording_backend(); + let proxy_port = free_port(); + + let config = load_example_config( + "trace-context.yaml", + proxy_port, + HashMap::from([("127.0.0.1:8321", backend_port)]), + ); + let proxy = start_proxy(&config); + + let mut request = json_post("/v1/responses", REQUEST_BODY); + if let Some(traceparent) = client_traceparent { + request = request.replace("\r\n\r\n", &format!("\r\ntraceparent: {traceparent}\r\n\r\n")); + } + + let raw = http_send(proxy.addr(), &request); + assert_eq!(parse_status(&raw), 200, "traced request should succeed"); + + let legs = seen.lock().expect("recorder mutex should not be poisoned").clone(); + assert!( + legs.iter().any(|leg| leg.path.starts_with("/v1/files")), + "expected a delegated Files API callout, saw {:?}", + legs.iter().map(|leg| &leg.path).collect::>() + ); + assert!( + legs.iter().any(|leg| leg.path == "/v1/responses"), + "expected the forwarded inference request, saw {:?}", + legs.iter().map(|leg| &leg.path).collect::>() + ); + legs +} + +#[test] +fn example_config_correlates_delegated_and_forwarded_legs() { + let legs = capture_legs(None); + + let paths: Vec<&str> = legs.iter().map(|leg| leg.path.as_str()).collect(); + + let trace_ids: HashSet<&str> = legs.iter().map(|leg| field(leg.traceparent(), 1)).collect(); + assert_eq!( + trace_ids.len(), + 1, + "all legs must share one trace-id, saw {trace_ids:?} across {paths:?}" + ); + assert!( + !trace_ids.iter().next().is_some_and(|id| id.is_empty()), + "trace-id should be populated" + ); + + // The delegated callouts of one request share a span: correlation + // is resolved once at the filter boundary, so file resolution is + // one delegation hop regardless of how many files it fetches. The + // forwarded request must be a distinct span, which is what keeps + // delegation latency separable from inference latency. + let delegated_spans: HashSet<&str> = legs + .iter() + .filter(|leg| leg.path.starts_with("/v1/files")) + .map(|leg| field(leg.traceparent(), 2)) + .collect(); + let forwarded_spans: HashSet<&str> = legs + .iter() + .filter(|leg| leg.path == "/v1/responses") + .map(|leg| field(leg.traceparent(), 2)) + .collect(); + + assert_eq!( + delegated_spans.len(), + 1, + "delegated callouts of one request should share a span, saw {delegated_spans:?} across {paths:?}" + ); + assert!( + delegated_spans.is_disjoint(&forwarded_spans), + "forwarded request must be its own span so delegation latency stays separable: \ + delegated={delegated_spans:?} forwarded={forwarded_spans:?}" + ); + + let request_ids: HashSet<&str> = legs + .iter() + .map(|leg| { + leg.header("x-request-id") + .unwrap_or_else(|| panic!("leg {} must carry x-request-id", leg.path)) + }) + .collect(); + assert_eq!( + request_ids.len(), + 1, + "all legs must share one request ID, saw {request_ids:?} across {paths:?}" + ); +} + +#[test] +fn example_config_continues_client_supplied_trace() { + let legs = capture_legs(Some(CLIENT_TRACEPARENT)); + + for leg in &legs { + let traceparent = leg.traceparent(); + assert_eq!( + field(traceparent, 1), + CLIENT_TRACE_ID, + "client trace-id should be continued on leg {}", + leg.path + ); + assert_ne!( + field(traceparent, 2), + CLIENT_SPAN_ID, + "leg {} must mint its own span-id, not reuse the client's", + leg.path + ); + } +}