diff --git a/crates/aisix-core/src/config.rs b/crates/aisix-core/src/config.rs index 458ead8ac..e385dde98 100644 --- a/crates/aisix-core/src/config.rs +++ b/crates/aisix-core/src/config.rs @@ -764,6 +764,15 @@ impl RequestIdConfig { if RESERVED_UPSTREAM_HEADERS.contains(&name.as_str()) { return Err(s.clone()); } + // W3C trace-context headers can't be request-id sources + // either (AISIX-Cloud#1279): the resolved id is echoed to + // the caller, logged, and sent upstream — adopting the + // inbound `traceparent`/`tracestate` value would disclose + // the caller's trace context through all three, walking + // around the never-forward guard. + if matches!(name.as_str(), "traceparent" | "tracestate") { + return Err(s.clone()); + } Ok(name) }) .collect() diff --git a/crates/aisix-gateway/src/upstream_headers.rs b/crates/aisix-gateway/src/upstream_headers.rs index 09772ca3b..4be155f7a 100644 --- a/crates/aisix-gateway/src/upstream_headers.rs +++ b/crates/aisix-gateway/src/upstream_headers.rs @@ -63,6 +63,18 @@ const NEVER_FORWARD_HEADERS: &[&str] = &[ "te", "trailer", "transfer-encoding", + // W3C trace context (AISIX-Cloud#1279): the caller's `traceparent` + // names a span in the CALLER's trace — relayed to a provider it links + // the provider's internal telemetry into the caller's tracing backend + // and leaks the caller's trace ids to a third party. The gateway is a + // fresh hop; provider-side propagation, if ever wanted, is a separate + // opt-in that injects the gateway's own context, never the inbound + // value. Listed here rather than in RESERVED_UPSTREAM_HEADERS so an + // operator's deliberate `default_headers` entry for a trusted + // first-party upstream still works — only the client's copy is + // blocked. + "traceparent", + "tracestate", "upgrade", ]; @@ -370,6 +382,32 @@ mod tests { assert_eq!(headers["x-keep"], "yes"); } + // AISIX-Cloud#1279: the caller's W3C trace context must never reach a + // provider — not under a `*` glob, and not even when an allowlist names + // the headers outright. The gateway is a fresh tracing hop; a future + // provider-side propagation opt-in injects the gateway's OWN context. + #[test] + fn trace_context_headers_are_never_forwarded() { + for allowlist in [&["*"][..], &["traceparent", "tracestate"][..]] { + let patterns: Vec<&str> = allowlist.to_vec(); + let r = overrides(&[], &patterns); + let mut headers = HeaderMap::new(); + let inbound = client(&[ + ( + "traceparent", + "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01", + ), + ("tracestate", "vendor=x"), + ]); + let ctx = UpstreamHeaderContext::from_overrides(Some(&r)).with_client_headers(&inbound); + apply_request_headers(&mut headers, &ctx); + assert!( + headers.is_empty(), + "trace context leaked under {allowlist:?}: {headers:?}" + ); + } + } + #[test] fn allowlist_matches_exact_names_and_a_single_glob() { let r = overrides(&[], &["Anthropic-Beta", "x-trace-*"]); diff --git a/crates/aisix-obs/src/lib.rs b/crates/aisix-obs/src/lib.rs index 9e78c1ee4..7f3f9c4ac 100644 --- a/crates/aisix-obs/src/lib.rs +++ b/crates/aisix-obs/src/lib.rs @@ -20,6 +20,7 @@ pub mod metrics; pub mod otlp; pub mod otlp_http_sink; pub mod sink; +pub mod trace; pub mod usage; use std::io::IsTerminal as _; @@ -41,6 +42,10 @@ pub use sink::{ PipelineConfig, SinkAck, SinkCapabilities, SinkContent, SinkError, SinkHandle, SinkHealth, SinkPipeline, SinkRecord, SinkResult, SinkStatsSnapshot, SCHEMA_VERSION, }; +pub use trace::{ + parse_traceparent, screen_tracestate, RemoteTraceContext, RequestTraceBundle, SpanEmit, + SpanRole, TraceEmission, TRACEPARENT_HEADER, TRACESTATE_HEADER, +}; pub use usage::{UsageEvent, UsageSink}; #[derive(Debug, thiserror::Error)] diff --git a/crates/aisix-obs/src/otlp_http_sink.rs b/crates/aisix-obs/src/otlp_http_sink.rs index d0059f3b1..ea7f87b8e 100644 --- a/crates/aisix-obs/src/otlp_http_sink.rs +++ b/crates/aisix-obs/src/otlp_http_sink.rs @@ -132,10 +132,19 @@ impl OtlpHttpFanOut { /// `content_mode = full`; every other exporter — and the CP telemetry /// path, which is not in this loop — receives the shared metadata-only /// record, so prompt/response can never leak there. + /// + /// `trace` is the event's span snapshot (AISIX-Cloud#1279) — ids and + /// boundaries frozen at emission time by the request's trace bundle. + /// It rides every record unchanged (`SinkRecord::trace` is + /// serde-skipped, so only the OTLP encoder reads it), which is what + /// makes span ids byte-identical across delivery retries and across + /// exporters. `None` (a caller predating the bundle, or a synthetic + /// event) falls back to the legacy single-span encoding. pub fn fan_out<'a, I>( &self, event: &UsageEvent, content: Option<&CapturedContent>, + trace: Option<&crate::trace::TraceEmission>, exporters: I, ) where I: IntoIterator, @@ -246,11 +255,10 @@ impl OtlpHttpFanOut { // A content-bearing record for an SLS exporter that opted into // full capture (and only when the handler captured content); the // shared metadata-only record for everyone else. - let record = content_record(&exp.kind, event, content).unwrap_or_else(|| { - Arc::clone( - metadata_record - .get_or_insert_with(|| Arc::new(SinkRecord::metadata_only(event.clone()))), - ) + let record = content_record(&exp.kind, event, content, trace).unwrap_or_else(|| { + Arc::clone(metadata_record.get_or_insert_with(|| { + Arc::new(SinkRecord::metadata_only(event.clone()).with_trace(trace.cloned())) + })) }); handle.try_enqueue(record); } @@ -393,6 +401,7 @@ fn content_record( kind: &ExporterKind, event: &UsageEvent, content: Option<&CapturedContent>, + trace: Option<&crate::trace::TraceEmission>, ) -> Option> { // The exporters that opt into full content capture share the // `SlsContentMode` model; pull each one's (mode, cap). Any other kind @@ -410,7 +419,9 @@ fn content_record( let mut sc = SinkContent::capture(&captured.prompt, &captured.response, max_bytes as usize); sc.truncated = sc.truncated || captured.truncated; Some(Arc::new( - SinkRecord::metadata_only(event.clone()).with_content(sc), + SinkRecord::metadata_only(event.clone()) + .with_content(sc) + .with_trace(trace.cloned()), )) } @@ -532,12 +543,12 @@ impl ObservabilitySink for OtlpSink { if batch.is_empty() { return Ok(SinkAck::default()); } - // One export request carrying every record's span — one POST, one + // One export request carrying every record's spans — one POST, one // atomic retry unit (vs. the per-event fan-out's N spawns). let spans: Vec = batch .records .iter() - .map(|record| build_otlp_span(record, &self.name)) + .flat_map(|record| build_otlp_spans(record, &self.name)) .collect(); let body = otlp_export_request(spans); let bytes = serde_json::to_vec(&body) @@ -631,13 +642,141 @@ fn build_otlp_span(record: &SinkRecord, exporter_name: &str) -> Value { let latency_nanos = (event.upstream_latency_ms as u128).saturating_mul(1_000_000); let start_unix_nano = end_unix_nano.saturating_sub(latency_nanos); - // Status: OK (1) for 2xx, ERROR (2) otherwise. - let status_code = if (200..300).contains(&event.status_code) { + json!({ + "traceId": trace_id, + "spanId": span_id, + "name": span_name(event), + // SPAN_KIND_CLIENT for all three: the gateway is the client of a + // remote LLM, agent or MCP server. The conventions name INTERNAL for + // a tool executed in-process, which is not what this span describes. + "kind": 3, + "startTimeUnixNano": start_unix_nano.to_string(), + "endTimeUnixNano": end_unix_nano.to_string(), + "attributes": event_attributes(record, exporter_name), + "status": { "code": span_status(event.status_code) }, + }) +} + +/// Encode one sink record into its OTLP span set. +/// +/// A record carrying a [`crate::trace::TraceEmission`] (AISIX-Cloud#1279) +/// yields the request's span hierarchy — SERVER → logical GenAI CLIENT → +/// attempt CLIENT — with ids and nanosecond boundaries frozen at emission +/// time, so a delivery retry re-encodes byte-identical spans and every +/// exporter agrees on the trace. A record without one (a caller predating +/// the bundle, or a synthetic event) falls back to the legacy single-span +/// shape of [`build_otlp_span`]. +fn build_otlp_spans(record: &SinkRecord, exporter_name: &str) -> Vec { + use crate::trace::SpanRole; + + let Some(trace) = &record.trace else { + return vec![build_otlp_span(record, exporter_name)]; + }; + let event = &record.usage; + let trace_id = trace.trace_id.to_hex(); + let status = span_status(event.status_code); + let carrier = trace.carrier_role(); + trace + .spans + .iter() + .map(|span| { + // The carrier span — the attempt when present, else the + // logical CLIENT, else SERVER — gets the event's full + // attribute set (and captured content); the structural spans + // above it carry only what a consumer needs to join them. + let attributes = if Some(span.role) == carrier { + event_attributes(record, exporter_name) + } else { + structural_attributes(event, exporter_name) + }; + // OTLP SpanKind: SERVER (2) for the inbound request; CLIENT + // (3) for the logical GenAI operation and each upstream + // attempt. + let kind = match span.role { + SpanRole::Server => 2, + SpanRole::Logical | SpanRole::Attempt => 3, + }; + let mut json = json!({ + "traceId": trace_id, + "spanId": span.span_id.to_hex(), + "name": span_name(event), + "kind": kind, + "startTimeUnixNano": span.start_unix_nano.to_string(), + "endTimeUnixNano": span.end_unix_nano.to_string(), + "attributes": attributes, + "status": { "code": status }, + "flags": span_flags(span, trace), + }); + if let Some(parent) = span.parent_span_id { + json["parentSpanId"] = Value::String(parent.to_hex()); + } + // The caller's tracestate belongs to the propagated context; + // it rides the SERVER span, the one that continues it. + if span.role == SpanRole::Server { + if let Some(state) = &trace.tracestate { + json["traceState"] = Value::String(state.clone()); + } + } + json + }) + .collect() +} + +/// OTLP span `flags` (uint32): the low byte is the span's W3C trace-flags +/// — always `sampled`, since encoding IS the export decision, plus the +/// inbound octet's `random` bit (0x02) on a SERVER span continuing a +/// remote trace — and `HAS_IS_REMOTE` (bit 8) on every span, with +/// `IS_REMOTE` (bit 9) on a SERVER span whose parent arrived over the +/// wire in `traceparent`. +fn span_flags(span: &crate::trace::SpanEmit, trace: &crate::trace::TraceEmission) -> u32 { + const SAMPLED: u32 = 0x01; + const W3C_RANDOM: u32 = 0x02; + const HAS_IS_REMOTE: u32 = 0x100; + const IS_REMOTE: u32 = 0x200; + let remote = span.role == crate::trace::SpanRole::Server + && span.parent_span_id.is_some() + && trace.remote_flags.is_some(); + let mut flags = SAMPLED | HAS_IS_REMOTE; + if remote { + flags |= IS_REMOTE; + // The random bit describes the trace id, which the remote parent + // minted — preserve its claim; never invent it for local roots. + flags |= u32::from(trace.remote_flags.unwrap_or(0)) & W3C_RANDOM; + } + flags +} + +/// The joining attributes a structural (non-carrier) span carries: enough +/// to group a trace's spans by request and operation without duplicating +/// the event's full attribute set onto every level of the hierarchy. +fn structural_attributes(event: &UsageEvent, exporter_name: &str) -> Vec { + let mut attributes = vec![ + attr_string("gen_ai.system", "aisix"), + attr_string("gen_ai.operation.name", operation_name(event)), + attr_string("aisix.exporter_name", exporter_name), + attr_string("aisix.request_id", &event.request_id), + attr_int("http.response.status_code", event.status_code as i64), + ]; + if !event.requested_model.is_empty() { + attributes.push(attr_string("gen_ai.request.model", &event.requested_model)); + } + attributes +} + +/// Status: OK (1) for 2xx, ERROR (2) otherwise. +fn span_status(status: u16) -> i64 { + if (200..300).contains(&status) { 1 } else { 2 - }; + } +} +/// The event's full attribute set — every wire field the OTLP allowlist +/// exports, plus captured content when the record carries it. Attached to +/// the legacy single span and to the hierarchy's carrier span. +fn event_attributes(record: &SinkRecord, exporter_name: &str) -> Vec { + let event = &record.usage; let mut attributes = vec![ attr_string("gen_ai.system", "aisix"), attr_string("gen_ai.operation.name", operation_name(event)), @@ -834,19 +973,7 @@ fn build_otlp_span(record: &SinkRecord, exporter_name: &str) -> Value { } } - json!({ - "traceId": trace_id, - "spanId": span_id, - "name": span_name(event), - // SPAN_KIND_CLIENT for all three: the gateway is the client of a - // remote LLM, agent or MCP server. The conventions name INTERNAL for - // a tool executed in-process, which is not what this span describes. - "kind": 3, - "startTimeUnixNano": start_unix_nano.to_string(), - "endTimeUnixNano": end_unix_nano.to_string(), - "attributes": attributes, - "status": { "code": status_code }, - }) + attributes } /// The OpenTelemetry GenAI operation this event describes. @@ -1140,6 +1267,105 @@ mod tests { }) } + /// A three-attempt failover bundle with a remote parent, for the + /// hierarchy-encoding assertions (AISIX-Cloud#1279). + fn hierarchy_record() -> SinkRecord { + let (trace_id, parent_span_id, flags) = crate::trace::parse_traceparent( + "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01", + ) + .unwrap(); + let bundle = + crate::trace::RequestTraceBundle::new(Some(crate::trace::RemoteTraceContext { + trace_id, + parent_span_id, + flags, + tracestate: Some("vendor=x".into()), + })); + bundle.start_attempt(0); + bundle.end_attempt(0); + bundle.start_attempt(1); + bundle.end_attempt(1); + let mut event = sample_event(); + event.attempt_index = 1; + SinkRecord::metadata_only(event).with_trace(Some(bundle.emission(true, 1, 250, true))) + } + + /// The hierarchy contract (AISIX-Cloud#1279): a terminal record with a + /// remote parent encodes SERVER (kind 2, parented to the inbound + /// traceparent, carrying traceState) → logical CLIENT → attempt + /// CLIENT, all under the caller's trace id, with the event's full + /// attribute set on the attempt span only. + #[test] + fn hierarchy_record_encodes_server_logical_attempt_spans() { + let record = hierarchy_record(); + let spans = build_otlp_spans(&record, "exp-1"); + assert_eq!(spans.len(), 3); + + for span in &spans { + assert_eq!(span["traceId"], "4bf92f3577b34da6a3ce929d0e0e4736"); + } + let server = &spans[0]; + let logical = &spans[1]; + let attempt = &spans[2]; + assert_eq!(server["kind"], 2); + assert_eq!(server["parentSpanId"], "00f067aa0ba902b7"); + assert_eq!(server["traceState"], "vendor=x"); + // SAMPLED | HAS_IS_REMOTE | IS_REMOTE. + assert_eq!(server["flags"], 0x301); + + assert_eq!(logical["kind"], 3); + assert_eq!(logical["parentSpanId"], server["spanId"]); + assert_eq!(logical["flags"], 0x101); + assert!(logical.get("traceState").is_none()); + + assert_eq!(attempt["kind"], 3); + assert_eq!(attempt["parentSpanId"], logical["spanId"]); + + // The event's full attribute set rides the attempt span; the + // structural spans carry only the joining keys. + let keys = |span: &Value| -> Vec { + span["attributes"] + .as_array() + .unwrap() + .iter() + .map(|a| a["key"].as_str().unwrap().to_string()) + .collect() + }; + assert!(keys(attempt).contains(&"gen_ai.response.id".to_string())); + assert!(!keys(server).contains(&"gen_ai.response.id".to_string())); + assert!(keys(server).contains(&"aisix.request_id".to_string())); + + // Timestamps bracket: server ⊇ logical ⊇ attempt. + let nanos = + |span: &Value, key: &str| -> u128 { span[key].as_str().unwrap().parse().unwrap() }; + assert!(nanos(server, "startTimeUnixNano") <= nanos(logical, "startTimeUnixNano")); + assert!(nanos(logical, "endTimeUnixNano") <= nanos(server, "endTimeUnixNano")); + assert!(nanos(logical, "startTimeUnixNano") <= nanos(attempt, "startTimeUnixNano")); + assert!(nanos(attempt, "endTimeUnixNano") <= nanos(logical, "endTimeUnixNano")); + } + + /// The regenerate-per-retry defect this PR removes: encoding the same + /// record twice — what the delivery retry loop does — must produce + /// byte-identical ids and timestamps. + #[test] + fn re_encoding_a_trace_record_is_byte_identical() { + let record = hierarchy_record(); + let a = build_otlp_spans(&record, "exp-1"); + let b = build_otlp_spans(&record, "exp-1"); + assert_eq!(a, b); + } + + /// A record without a trace snapshot (a synthetic event, or a caller + /// predating the bundle) keeps the legacy single-span shape. + #[test] + fn traceless_record_falls_back_to_the_legacy_single_span() { + let record = SinkRecord::metadata_only(sample_event()); + let spans = build_otlp_spans(&record, "exp-1"); + assert_eq!(spans.len(), 1); + assert_eq!(spans[0]["kind"], 3); + assert!(spans[0].get("parentSpanId").is_none()); + } + fn otlp_kind(content_mode: SlsContentMode, max_bytes: u32) -> ExporterKind { ExporterKind::OtlpHttp(OtlpHttpConfig { endpoint: "https://x/v1/traces".into(), @@ -1174,7 +1400,7 @@ mod tests { // otlp on its default (metadata_only) never carries content, even // when content was captured. let otlp = otlp_kind(SlsContentMode::MetadataOnly, 1024); - assert!(content_record(&otlp, &event, Some(&captured)).is_none()); + assert!(content_record(&otlp, &event, Some(&captured), None).is_none()); // object_store never carries content either — content_record gates on // the AliyunSls variant, so any other kind gets metadata-only even when @@ -1190,19 +1416,19 @@ mod tests { })) .unwrap() .kind; - assert!(content_record(&objstore, &event, Some(&captured)).is_none()); + assert!(content_record(&objstore, &event, Some(&captured), None).is_none()); // sls metadata_only → no content. let meta = sls_kind(SlsContentMode::MetadataOnly, 1024); - assert!(content_record(&meta, &event, Some(&captured)).is_none()); + assert!(content_record(&meta, &event, Some(&captured), None).is_none()); // sls full but nothing captured → falls back to metadata. let full = sls_kind(SlsContentMode::Full, 1024); - assert!(content_record(&full, &event, None).is_none()); + assert!(content_record(&full, &event, None, None).is_none()); // sls full + captured content → a content-bearing record that still // carries the metadata. - let rec = content_record(&full, &event, Some(&captured)) + let rec = content_record(&full, &event, Some(&captured), None) .expect("full-capture sls with content yields a content record"); let c = rec.content.as_ref().expect("content attached"); assert_eq!(c.prompt, "the prompt"); @@ -1216,7 +1442,13 @@ mod tests { response: "ok".into(), truncated: false, }; - let rec = content_record(&sls_kind(SlsContentMode::Full, 16), &event, Some(&big)).unwrap(); + let rec = content_record( + &sls_kind(SlsContentMode::Full, 16), + &event, + Some(&big), + None, + ) + .unwrap(); let c = rec.content.as_ref().unwrap(); assert_eq!(c.prompt.len(), 16); assert!(c.truncated, "oversize content must flag truncated"); @@ -1228,7 +1460,7 @@ mod tests { response: "short".into(), truncated: true, }; - let rec = content_record(&full, &event, Some(&pre)).unwrap(); + let rec = content_record(&full, &event, Some(&pre), None).unwrap(); assert!( rec.content.as_ref().unwrap().truncated, "source truncation must propagate" @@ -1238,31 +1470,42 @@ mod tests { // full + captured content → a content-bearing record (same shared // `SlsContentMode` plumbing). let dd_meta = datadog_kind(SlsContentMode::MetadataOnly, 1024); - assert!(content_record(&dd_meta, &event, Some(&captured)).is_none()); + assert!(content_record(&dd_meta, &event, Some(&captured), None).is_none()); let dd_full = datadog_kind(SlsContentMode::Full, 1024); - assert!(content_record(&dd_full, &event, None).is_none()); - let rec = content_record(&dd_full, &event, Some(&captured)) + assert!(content_record(&dd_full, &event, None, None).is_none()); + let rec = content_record(&dd_full, &event, Some(&captured), None) .expect("full-capture datadog with content yields a content record"); let c = rec.content.as_ref().expect("content attached"); assert_eq!(c.prompt, "the prompt"); assert_eq!(c.response, "the response"); // Per-exporter cap truncates a datadog record too. - let rec = - content_record(&datadog_kind(SlsContentMode::Full, 16), &event, Some(&big)).unwrap(); + let rec = content_record( + &datadog_kind(SlsContentMode::Full, 16), + &event, + Some(&big), + None, + ) + .unwrap(); assert_eq!(rec.content.as_ref().unwrap().prompt.len(), 16); assert!(rec.content.as_ref().unwrap().truncated); // otlp_http behaves identically (#519 B.2): full + captured content → // a content-bearing record; full with nothing captured → metadata. let otlp_full = otlp_kind(SlsContentMode::Full, 1024); - assert!(content_record(&otlp_full, &event, None).is_none()); - let rec = content_record(&otlp_full, &event, Some(&captured)) + assert!(content_record(&otlp_full, &event, None, None).is_none()); + let rec = content_record(&otlp_full, &event, Some(&captured), None) .expect("full-capture otlp with content yields a content record"); let c = rec.content.as_ref().expect("content attached"); assert_eq!(c.prompt, "the prompt"); assert_eq!(c.response, "the response"); // Per-exporter cap truncates an otlp record too. - let rec = content_record(&otlp_kind(SlsContentMode::Full, 16), &event, Some(&big)).unwrap(); + let rec = content_record( + &otlp_kind(SlsContentMode::Full, 16), + &event, + Some(&big), + None, + ) + .unwrap(); assert_eq!(rec.content.as_ref().unwrap().prompt.len(), 16); assert!(rec.content.as_ref().unwrap().truncated); } @@ -1295,6 +1538,7 @@ mod tests { &otlp_kind(SlsContentMode::Full, 1024), &event, Some(&captured), + None, ) .unwrap(); let span = build_otlp_span(&rec, "x"); @@ -1317,7 +1561,13 @@ mod tests { response: "ok".into(), truncated: false, }; - let rec = content_record(&otlp_kind(SlsContentMode::Full, 16), &event, Some(&big)).unwrap(); + let rec = content_record( + &otlp_kind(SlsContentMode::Full, 16), + &event, + Some(&big), + None, + ) + .unwrap(); let span = build_otlp_span(&rec, "x"); let attrs = span["attributes"].as_array().unwrap(); let flag = attrs @@ -1959,7 +2209,7 @@ mod tests { // can't easily count spawned tasks, but if the call returned // and the test process didn't hang, we're good. let f = OtlpHttpFanOut::new(); - f.fan_out(&sample_event(), None, std::iter::empty()); + f.fan_out(&sample_event(), None, None, std::iter::empty()); } #[test] @@ -1973,7 +2223,7 @@ mod tests { let mut exp = sample_exporter(); exp.enabled = false; let f = OtlpHttpFanOut::new(); - f.fan_out(&sample_event(), None, std::iter::once(&exp)); + f.fan_out(&sample_event(), None, None, std::iter::once(&exp)); } #[tokio::test] @@ -2006,7 +2256,7 @@ mod tests { let control = mk("control", &ctrl_srv.uri(), None); let f = OtlpHttpFanOut::new(); - f.fan_out(&sample_event(), None, [&zero, &control]); + f.fan_out(&sample_event(), None, None, [&zero, &control]); // The control's span lands after the 1s flush… let deadline = std::time::Instant::now() + std::time::Duration::from_secs(5); @@ -2051,7 +2301,7 @@ mod tests { .unwrap(); let f = OtlpHttpFanOut::new(); - f.fan_out(&sample_event(), None, std::iter::once(&exp)); + f.fan_out(&sample_event(), None, None, std::iter::once(&exp)); // Poll for the batched POST (flush_interval is 1s). let deadline = std::time::Instant::now() + std::time::Duration::from_secs(5); diff --git a/crates/aisix-obs/src/sink/record.rs b/crates/aisix-obs/src/sink/record.rs index ed37da8ae..ada453f8f 100644 --- a/crates/aisix-obs/src/sink/record.rs +++ b/crates/aisix-obs/src/sink/record.rs @@ -97,6 +97,15 @@ pub struct SinkRecord { /// Opt-in captured content; omitted entirely under `metadata_only`. #[serde(skip_serializing_if = "Option::is_none")] pub content: Option, + /// The event's span snapshot (AISIX-Cloud#1279): ids + boundaries + /// frozen at emission time, so the encoder — which runs inside the + /// delivery retry loop — reproduces byte-identical ids on every retry. + /// `serde(skip)`: this is OTLP structure, not record data; the + /// Datadog/SLS/object-store wire shapes are unchanged. (The public + /// `UsageEvent::trace_id` correlation key rides the flattened usage + /// fields instead.) + #[serde(skip)] + pub trace: Option, } impl SinkRecord { @@ -106,6 +115,7 @@ impl SinkRecord { schema_version: SCHEMA_VERSION, usage, content: None, + trace: None, } } @@ -114,6 +124,12 @@ impl SinkRecord { self.content = Some(content); self } + + /// Attach the request's span snapshot (AISIX-Cloud#1279). + pub fn with_trace(mut self, trace: Option) -> Self { + self.trace = trace; + self + } } /// A batch of records the pipeline hands to a sink in one delivery. @@ -156,6 +172,23 @@ mod tests { assert_eq!(json["request_id"], "req-1"); } + /// The public correlation key (AISIX-Cloud#1279): `trace_id` rides the + /// flattened usage fields on every sink wire when set — the exact key + /// the CP column and `{trace_id}` link template consume — and is + /// absent (→ SQL NULL) when the emitting path predates the bundle. + #[test] + fn trace_id_rides_the_flattened_wire_when_set_and_is_absent_when_empty() { + let rec = SinkRecord::metadata_only(UsageEvent { + trace_id: "4bf92f3577b34da6a3ce929d0e0e4736".into(), + ..UsageEvent::default() + }); + let json = serde_json::to_value(&rec).unwrap(); + assert_eq!(json["trace_id"], "4bf92f3577b34da6a3ce929d0e0e4736"); + + let bare = serde_json::to_value(SinkRecord::metadata_only(UsageEvent::default())).unwrap(); + assert!(bare.get("trace_id").is_none()); + } + #[test] fn full_content_record_carries_prompt_and_response() { let rec = SinkRecord::metadata_only(UsageEvent::default()).with_content(SinkContent { diff --git a/crates/aisix-obs/src/trace.rs b/crates/aisix-obs/src/trace.rs new file mode 100644 index 000000000..6a46177d1 --- /dev/null +++ b/crates/aisix-obs/src/trace.rs @@ -0,0 +1,887 @@ +//! Request-owned trace identity for the OTLP export path (AISIX-Cloud#1279). +//! +//! Before this module every exported span minted fresh random trace/span +//! ids inside the encoder — which runs inside the sink pipeline's retry +//! loop, so a transient delivery failure re-encoded the batch and a +//! receiver that dedups on span id saw N distinct spans for one attempt. +//! And with every span its own root, the per-attempt events of one +//! request (#655) joined only on the `aisix.request_id` attribute, which +//! no trace waterfall understands. +//! +//! The fix is ownership: a [`RequestTraceBundle`] is minted ONCE per +//! request (by the proxy's request-id middleware, the same single mint +//! point the request id itself uses), carries the trace id, the SERVER +//! span id, the logical GenAI span id and one span id per upstream +//! attempt, and rides into the sink layer as an immutable +//! [`TraceEmission`] snapshot on each record — so a delivery retry +//! re-encodes byte-identical ids, and every exporter sees the same trace. +//! +//! Wire vocabulary follows W3C Trace Context Level 1 +//! (): `traceparent` version 00, +//! 16-byte trace id / 8-byte parent id, lowercase hex. Ingestion is +//! strict — anything malformed degrades to a locally-rooted trace and the +//! caller's `tracestate` is discarded with it, so an upstream proxy that +//! mangles the header can never smuggle a zero id or mixed-case hex into +//! the export. + +use std::sync::atomic::{AtomicBool, Ordering}; +use std::sync::Mutex; +use std::time::{Instant, SystemTime, UNIX_EPOCH}; + +/// The W3C `traceparent` request header (Trace Context Level 1 §3.2). +pub const TRACEPARENT_HEADER: &str = "traceparent"; +/// The W3C `tracestate` request header (Trace Context Level 1 §3.3). +pub const TRACESTATE_HEADER: &str = "tracestate"; + +/// Longest accepted `tracestate` value. The spec permits up to 32 list +/// members; 512 bytes covers every real-world vendor chain while keeping +/// a hostile caller from riding kilobytes into every exported span. +const MAX_TRACESTATE_BYTES: usize = 512; + +/// A 16-byte W3C trace id. All-zero is invalid per spec and unrepresentable +/// through the public constructors. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct TraceId(pub [u8; 16]); + +/// An 8-byte W3C span id. All-zero is invalid per spec and unrepresentable +/// through the public constructors. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct SpanId(pub [u8; 8]); + +impl TraceId { + /// 16 random bytes. UUIDv4 gives 122 random bits, which is the same + /// source the request-id mint uses; the 6 version/variant bits are an + /// acceptable loss for an id whose only job is uniqueness. + pub fn random() -> Self { + Self(*uuid::Uuid::new_v4().as_bytes()) + } + + /// 32 lowercase-hex chars, the OTLP/JSON and W3C wire form. + pub fn to_hex(self) -> String { + hex_lower(&self.0) + } +} + +impl SpanId { + pub fn random() -> Self { + let b = uuid::Uuid::new_v4(); + let b = b.as_bytes(); + Self([b[0], b[1], b[2], b[3], b[4], b[5], b[6], b[7]]) + } + + /// 16 lowercase-hex chars, the OTLP/JSON and W3C wire form. + pub fn to_hex(self) -> String { + hex_lower(&self.0) + } +} + +fn hex_lower(bytes: &[u8]) -> String { + use std::fmt::Write as _; + let mut s = String::with_capacity(bytes.len() * 2); + for b in bytes { + // Infallible for String. + let _ = write!(s, "{b:02x}"); + } + s +} + +/// A validated inbound W3C trace context — the caller's own tracing, +/// which the request's SERVER span parents under so the caller's backend +/// shows the gateway inside the caller's trace. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct RemoteTraceContext { + pub trace_id: TraceId, + pub parent_span_id: SpanId, + /// The verbatim `trace-flags` octet. The encoder ORs its W3C-defined + /// non-sampled bits (today only `random`, 0x02) into the SERVER span's + /// exported flags; the `sampled` bit is deliberately NOT consulted by + /// the exporter's own sampling — the operator's `sample_rate` stays + /// authoritative, so an inbound `sampled=1` cannot force export past + /// `sample_rate=0`. + pub flags: u8, + /// The caller's `tracestate`, kept only when the `traceparent` it + /// travelled with was valid (spec: state is meaningless without its + /// parent) and it passed the charset/length screen. + pub tracestate: Option, +} + +/// Strictly parse one `traceparent` header value (W3C Trace Context +/// Level 1 §3.2). Accepted shape is exactly the version-00 form: +/// +/// ```text +/// 00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01 +/// ``` +/// +/// Anything else — wrong length, uppercase hex, an unknown or `ff` +/// version, an all-zero trace or parent id — yields `None` and the +/// request gets a locally-rooted trace. The spec permits leniently +/// parsing future versions; a trust boundary does not: a value this +/// gateway will attach its own telemetry under has to be one it fully +/// understands. +pub fn parse_traceparent(value: &str) -> Option<(TraceId, SpanId, u8)> { + let b = value.as_bytes(); + if b.len() != 55 || b[2] != b'-' || b[35] != b'-' || b[52] != b'-' { + return None; + } + // Version: exactly "00". This also rejects "ff" (forbidden by spec) + // and any future version we have not reviewed. + if &value[0..2] != "00" { + return None; + } + let trace_id = hex_bytes::<16>(&value[3..35])?; + let parent_id = hex_bytes::<8>(&value[36..52])?; + let flags = hex_bytes::<1>(&value[53..55])?[0]; + if trace_id == [0u8; 16] || parent_id == [0u8; 8] { + return None; + } + Some((TraceId(trace_id), SpanId(parent_id), flags)) +} + +/// Decode exactly `2 * N` lowercase-hex chars. Uppercase is rejected — +/// the W3C wire form is lowercase-only and a mixed-case value signals a +/// producer we should not trust to have composed the rest correctly. +fn hex_bytes(s: &str) -> Option<[u8; N]> { + let b = s.as_bytes(); + if b.len() != N * 2 { + return None; + } + let mut out = [0u8; N]; + for (i, chunk) in b.chunks_exact(2).enumerate() { + let hi = hex_nibble(chunk[0])?; + let lo = hex_nibble(chunk[1])?; + out[i] = (hi << 4) | lo; + } + Some(out) +} + +fn hex_nibble(c: u8) -> Option { + match c { + b'0'..=b'9' => Some(c - b'0'), + b'a'..=b'f' => Some(c - b'a' + 10), + _ => None, + } +} + +/// Screen a combined `tracestate` value: printable ASCII (0x20..=0x7E), +/// non-empty, within [`MAX_TRACESTATE_BYTES`]. The list-member grammar is +/// deliberately not enforced beyond that — the value is forwarded opaque +/// into the exported span's `traceState`, never interpreted. +pub fn screen_tracestate(value: &str) -> Option { + let trimmed = value.trim(); + if trimmed.is_empty() + || trimmed.len() > MAX_TRACESTATE_BYTES + || !trimmed.bytes().all(|b| (0x20..=0x7E).contains(&b)) + { + return None; + } + Some(trimmed.to_owned()) +} + +/// Wall-clock nanosecond timestamps measured on a monotonic base. +/// +/// The wall clock is read ONCE, at request start; every later timestamp is +/// that base plus a monotonic `Instant` offset. This is what makes child +/// spans provably bracketed by their parents — a wall-clock step (NTP) +/// mid-request cannot reorder them — and what replaces the old +/// `occurred_at - latency` reconstruction, whose whole-second stamp put +/// ±1s of error on every span's absolute placement. +#[derive(Debug, Clone, Copy)] +pub struct TraceClock { + base_unix_nanos: u64, + base: Instant, +} + +impl TraceClock { + pub fn start() -> Self { + let base_unix_nanos = SystemTime::now() + .duration_since(UNIX_EPOCH) + .map(|d| d.as_nanos().min(u64::MAX as u128) as u64) + .unwrap_or(0); + Self { + base_unix_nanos, + base: Instant::now(), + } + } + + /// Nanoseconds since the Unix epoch, now. + pub fn now_unix_nanos(&self) -> u64 { + self.base_unix_nanos + .saturating_add(self.base.elapsed().as_nanos().min(u64::MAX as u128) as u64) + } +} + +/// One upstream attempt's span identity and real dispatch boundaries. +#[derive(Debug, Clone, Copy)] +struct AttemptSpan { + span_id: SpanId, + start_unix_nano: u64, + /// Stamped when the attempt settles ([`RequestTraceBundle::end_attempt`]); + /// `None` for an attempt still in flight at emission time. + end_unix_nano: Option, +} + +/// The request-owned trace identity: minted once at request start, +/// mutated only through the attempt chokepoints, snapshotted into an +/// immutable [`TraceEmission`] per exported usage event. +#[derive(Debug)] +pub struct RequestTraceBundle { + trace_id: TraceId, + server_span_id: SpanId, + /// The logical GenAI operation span — one per request, covering every + /// retry/failover attempt per the GenAI semantic conventions. + logical_span_id: SpanId, + remote: Option, + clock: TraceClock, + server_start_unix_nano: u64, + attempts: Mutex>, + /// Set by the first terminal emission. A second terminal call (a + /// double-emit bug upstream of here) degrades to an attempt-only + /// emission instead of shipping a duplicate SERVER span. + terminal_emitted: AtomicBool, + /// Set when a non-terminal emission synthesized a sub-call span (the + /// ensemble trunk bypass) — those spans are parented under the logical + /// span id, so the terminal emission MUST export that parent even when + /// the terminal event itself never dispatched (an ensemble whose judge + /// failed or whose output was blocked): otherwise every already-shipped + /// sub-call span dangles. + synthesized_children: AtomicBool, +} + +impl RequestTraceBundle { + /// Mint the request's trace identity. `remote` is the caller's + /// validated `traceparent`/`tracestate`, if any: it decides the trace + /// id (the caller's trace continues through the gateway) and the + /// SERVER span's parent. Without it the request roots a fresh trace. + pub fn new(remote: Option) -> Self { + let clock = TraceClock::start(); + Self { + trace_id: remote + .as_ref() + .map(|r| r.trace_id) + .unwrap_or_else(TraceId::random), + server_span_id: SpanId::random(), + logical_span_id: SpanId::random(), + remote, + server_start_unix_nano: clock.now_unix_nanos(), + clock, + attempts: Mutex::new(Vec::new()), + terminal_emitted: AtomicBool::new(false), + synthesized_children: AtomicBool::new(false), + } + } + + /// The trace id every span of this request shares, as 32 lowercase-hex + /// chars — the public correlation key (`UsageEvent::trace_id`). + pub fn trace_id_hex(&self) -> String { + self.trace_id.to_hex() + } + + /// Whether this request continues a caller-supplied trace (a valid + /// inbound `traceparent`). + pub fn has_remote_parent(&self) -> bool { + self.remote.is_some() + } + + /// Mark attempt `index` as dispatched now: mint its span id and stamp + /// its start. Called from `RoutingTelemetry::begin_attempt`, the + /// single place an attempt begins. Indexes arrive sequentially; a gap + /// (a dispatch path that skipped `begin_attempt`) is backfilled at the + /// same instant so the vec stays index-addressable. + pub fn start_attempt(&self, index: u32) { + let now = self.clock.now_unix_nanos(); + // Poisoning is survivable here: the guarded vec cannot be left + // logically inconsistent (worst case one backfilled entry), and + // this path runs from Drop guards where a panic could abort. + let mut attempts = self + .attempts + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + while attempts.len() <= index as usize { + attempts.push(AttemptSpan { + span_id: SpanId::random(), + start_unix_nano: now, + end_unix_nano: None, + }); + } + } + + /// Stamp attempt `index`'s end, once. Called from + /// `RoutingTelemetry::record`, the single place an attempt settles. + /// For the winning streaming attempt this fires at commit time; the + /// terminal emission later extends the end to the real stream end via + /// the event's own measured duration (see [`Self::emission`]). + pub fn end_attempt(&self, index: u32) { + let now = self.clock.now_unix_nanos(); + let mut attempts = self + .attempts + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + if let Some(a) = attempts.get_mut(index as usize) { + a.end_unix_nano.get_or_insert(now); + } + } + + /// Snapshot the spans one usage event carries. + /// + /// Non-terminal (`terminal = false`): the event describes one failed / + /// superseded attempt — emit that attempt's CLIENT span alone; the + /// SERVER and logical spans ride the terminal event. + /// + /// Terminal (`terminal = true`, first such call): the request is done — + /// stamp the SERVER span's end at NOW (the response body's real + /// EOF/drop, since every terminal emit runs from the family's Drop + /// guard or tail), emit SERVER + logical + this event's attempt span. + /// The event's attempt end is extended to `start + upstream_latency` + /// when the handler measured a longer duration than the commit-time + /// stamp — the streaming case, where `record()` fires at commit but + /// the upstream keeps streaming. + /// + /// A request that never dispatched an attempt (`attempts` empty): + /// - `dispatched_upstream` — a family that dispatches without + /// per-attempt tracking (MCP / A2A / jobs / realtime, the ensemble + /// trunk bypass): one CLIENT span under SERVER, placed by the + /// handler's own measured duration ending now. + /// - otherwise — a cache hit, a guardrail block, a pre-dispatch error: + /// the SERVER span alone. No fictitious upstream CLIENT span. + /// + /// `dispatched_upstream` is the CALLER's statement of fact and the + /// ONLY signal consulted — never inferred from the event's latency: a + /// cache hit and a pre-dispatch error legitimately carry the handler's + /// elapsed time in `upstream_latency_ms` (which an inference would + /// misread as an upstream call), while a real sub-millisecond upstream + /// call legitimately measures 0ms (which an inference would drop, + /// orphaning any spans parented under the never-exported logical + /// span). A 0ms dispatched call yields a truthful zero-length span. + pub fn emission( + &self, + terminal: bool, + attempt_index: u32, + upstream_latency_ms: u32, + dispatched_upstream: bool, + ) -> TraceEmission { + let now = self.clock.now_unix_nanos(); + let latency_nanos = u64::from(upstream_latency_ms).saturating_mul(1_000_000); + let attempts = self + .attempts + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + let mut spans = Vec::new(); + + // First terminal call wins; a duplicate degrades to attempt-only. + // (`&&` short-circuits, so a non-terminal call never touches the + // flag.) The original request is kept separately: a degraded + // duplicate must not fall into the synthesize-a-sub-call branch + // below and invent a second upstream span. + let terminal_requested = terminal; + let terminal = terminal && !self.terminal_emitted.swap(true, Ordering::AcqRel); + + let attempt_span = attempts.get(attempt_index as usize).map(|a| { + let mut end = a.end_unix_nano.unwrap_or(now); + if terminal { + // The winning streaming attempt: `record()` stamped the + // commit, the handler measured the full stream. Both are + // real monotonic measurements from the same dispatch + // boundary; take the longer, clamped to now. + end = end + .max(a.start_unix_nano.saturating_add(latency_nanos)) + .min(now); + } + SpanEmit { + role: SpanRole::Attempt, + span_id: a.span_id, + parent_span_id: Some(self.logical_span_id), + start_unix_nano: a.start_unix_nano, + end_unix_nano: end, + } + }); + // A non-terminal event from a dispatch path outside the attempt + // chokepoints — an ensemble panel member or judge sub-call, the + // known trunk bypass — still describes real upstream work: + // synthesize its CLIENT span (placed by the handler's own measured + // duration, ending now) so the sub-call is not lost from the + // trace. The id is minted once per emission snapshot, so delivery + // retries and every exporter still agree on it. + let attempt_span = attempt_span.or_else(|| { + (!terminal_requested && dispatched_upstream).then(|| { + // Remember that a span parented under the logical id has + // shipped — the terminal emission must export that parent + // even if the terminal event itself never dispatched. + self.synthesized_children.store(true, Ordering::Release); + SpanEmit { + role: SpanRole::Attempt, + span_id: SpanId::random(), + parent_span_id: Some(self.logical_span_id), + start_unix_nano: now + .saturating_sub(latency_nanos) + .max(self.server_start_unix_nano), + end_unix_nano: now, + } + }) + }); + + if terminal { + spans.push(SpanEmit { + role: SpanRole::Server, + span_id: self.server_span_id, + parent_span_id: self.remote.as_ref().map(|r| r.parent_span_id), + start_unix_nano: self.server_start_unix_nano, + end_unix_nano: now, + }); + if let Some(first) = attempts.first() { + // Logical span: first dispatch → last settle (extended to + // this terminal event's own measured end). + let last_end = attempts + .iter() + .filter_map(|a| a.end_unix_nano) + .max() + .unwrap_or(now) + .max(attempt_span.as_ref().map(|s| s.end_unix_nano).unwrap_or(0)) + .min(now); + spans.push(SpanEmit { + role: SpanRole::Logical, + span_id: self.logical_span_id, + parent_span_id: Some(self.server_span_id), + start_unix_nano: first.start_unix_nano, + end_unix_nano: last_end, + }); + } else if dispatched_upstream { + // Dispatched upstream without per-attempt tracking: one + // CLIENT span placed by the handler's measured duration, + // clamped inside the SERVER span. + spans.push(SpanEmit { + role: SpanRole::Logical, + span_id: self.logical_span_id, + parent_span_id: Some(self.server_span_id), + start_unix_nano: now + .saturating_sub(latency_nanos) + .max(self.server_start_unix_nano), + end_unix_nano: now, + }); + } else if self.synthesized_children.load(Ordering::Acquire) { + // The terminal event itself never dispatched — an ensemble + // whose judge failed or whose synthesized answer was + // blocked — but sub-call spans parented under the logical + // id already shipped: export their parent (spanning the + // request, which brackets every child) so they don't + // dangle. + spans.push(SpanEmit { + role: SpanRole::Logical, + span_id: self.logical_span_id, + parent_span_id: Some(self.server_span_id), + start_unix_nano: self.server_start_unix_nano, + end_unix_nano: now, + }); + } + // else: cache hit / block / pre-dispatch error — SERVER only. + } + + spans.extend(attempt_span); + + TraceEmission { + trace_id: self.trace_id, + remote_flags: self.remote.as_ref().map(|r| r.flags), + tracestate: self.remote.as_ref().and_then(|r| r.tracestate.clone()), + spans, + } + } +} + +/// Which structural role a span plays in the exported hierarchy. Decides +/// OTLP `kind` and which attribute set the encoder attaches; deliberately +/// NOT exported as a span attribute — consumers distinguish spans by +/// `kind` and parent linkage, the vocabulary trace backends already have. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum SpanRole { + /// The inbound HTTP request: one per request, OTLP `SPAN_KIND_SERVER`. + Server, + /// The logical GenAI operation covering every attempt, OTLP + /// `SPAN_KIND_CLIENT`. Doubles as the single upstream-call span for + /// families without per-attempt tracking. + Logical, + /// One upstream dispatch attempt, OTLP `SPAN_KIND_CLIENT`. + Attempt, +} + +/// One span's identity and boundaries, ready to encode. +#[derive(Debug, Clone, Copy)] +pub struct SpanEmit { + pub role: SpanRole, + pub span_id: SpanId, + /// `None` only for a SERVER span with no valid inbound `traceparent` + /// — the trace's local root. + pub parent_span_id: Option, + pub start_unix_nano: u64, + pub end_unix_nano: u64, +} + +/// The immutable per-event snapshot the sink layer carries +/// (`SinkRecord::trace`). Built once at emission time, so the encoder — +/// which runs inside the delivery retry loop — reproduces byte-identical +/// ids and timestamps on every retry and across every exporter. +#[derive(Debug, Clone)] +pub struct TraceEmission { + pub trace_id: TraceId, + /// The verbatim inbound `trace-flags` octet when the request continued + /// a remote trace; `None` for a locally-rooted one. + pub remote_flags: Option, + /// The caller's screened `tracestate`, exported on the SERVER span. + pub tracestate: Option, + pub spans: Vec, +} + +impl TraceEmission { + /// The event's carrier span — the one the encoder attaches the usage + /// event's full attribute set to: the attempt span when present, else + /// the logical span, else the SERVER span. + pub fn carrier_role(&self) -> Option { + [SpanRole::Attempt, SpanRole::Logical, SpanRole::Server] + .into_iter() + .find(|role| self.spans.iter().any(|s| s.role == *role)) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + const VALID: &str = "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01"; + + #[test] + fn parses_a_valid_v00_traceparent() { + let (trace, parent, flags) = parse_traceparent(VALID).expect("valid header"); + assert_eq!(trace.to_hex(), "4bf92f3577b34da6a3ce929d0e0e4736"); + assert_eq!(parent.to_hex(), "00f067aa0ba902b7"); + assert_eq!(flags, 0x01); + } + + #[test] + fn rejects_malformed_traceparents() { + for (case, why) in [ + ("", "empty"), + ( + "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7", + "missing flags", + ), + ( + "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01-extra", + "trailing segment", + ), + ( + "ff-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01", + "forbidden version ff", + ), + ( + "01-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01", + "unreviewed future version", + ), + ( + "00-00000000000000000000000000000000-00f067aa0ba902b7-01", + "all-zero trace id", + ), + ( + "00-4bf92f3577b34da6a3ce929d0e0e4736-0000000000000000-01", + "all-zero parent id", + ), + ( + "00-4BF92F3577B34DA6A3CE929D0E0E4736-00f067aa0ba902b7-01", + "uppercase hex", + ), + ( + "00-4bf92f3577b34da6a3ce929d0e0e473g-00f067aa0ba902b7-01", + "non-hex char", + ), + ( + "00_4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01", + "wrong delimiter", + ), + ( + "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-1", + "one-char flags", + ), + ] { + assert!(parse_traceparent(case).is_none(), "must reject: {why}"); + } + } + + #[test] + fn tracestate_screen_keeps_sane_values_and_drops_hostile_ones() { + assert_eq!( + screen_tracestate("congo=t61rcWkgMzE,rojo=00f067aa0ba902b7").as_deref(), + Some("congo=t61rcWkgMzE,rojo=00f067aa0ba902b7") + ); + assert_eq!( + screen_tracestate(" vendor=x ").as_deref(), + Some("vendor=x") + ); + assert!(screen_tracestate("").is_none()); + assert!(screen_tracestate(" ").is_none()); + assert!(screen_tracestate("bad\nvalue").is_none(), "control chars"); + assert!(screen_tracestate("non-ascii-日本").is_none()); + assert!(screen_tracestate(&"x".repeat(MAX_TRACESTATE_BYTES + 1)).is_none()); + } + + #[test] + fn clock_is_monotonic_and_epoch_anchored() { + let clock = TraceClock::start(); + let a = clock.now_unix_nanos(); + let b = clock.now_unix_nanos(); + assert!(b >= a); + // Sanity: after 2020-01-01 in nanos. + assert!(a > 1_577_836_800_000_000_000); + } + + fn bundle_with_remote() -> RequestTraceBundle { + let (trace_id, parent_span_id, flags) = parse_traceparent(VALID).unwrap(); + RequestTraceBundle::new(Some(RemoteTraceContext { + trace_id, + parent_span_id, + flags, + tracestate: Some("vendor=x".into()), + })) + } + + #[test] + fn remote_parent_decides_trace_id_and_server_parent() { + let bundle = bundle_with_remote(); + assert_eq!(bundle.trace_id_hex(), "4bf92f3577b34da6a3ce929d0e0e4736"); + let em = bundle.emission(true, 0, 0, false); + let server = em + .spans + .iter() + .find(|s| s.role == SpanRole::Server) + .expect("terminal emission carries the SERVER span"); + assert_eq!( + server.parent_span_id.map(SpanId::to_hex).as_deref(), + Some("00f067aa0ba902b7") + ); + assert_eq!(em.remote_flags, Some(0x01)); + assert_eq!(em.tracestate.as_deref(), Some("vendor=x")); + } + + #[test] + fn local_root_has_no_parent_and_a_random_trace() { + let bundle = RequestTraceBundle::new(None); + let em = bundle.emission(true, 0, 0, false); + let server = em + .spans + .iter() + .find(|s| s.role == SpanRole::Server) + .unwrap(); + assert!(server.parent_span_id.is_none()); + assert!(em.remote_flags.is_none()); + assert_ne!(bundle.trace_id_hex(), "00000000000000000000000000000000"); + } + + /// The failover shape the whole design exists for: SERVER → logical → + /// N attempt children, ids stable across snapshots, children bracketed + /// by parents. + #[test] + fn failover_emissions_share_one_hierarchy() { + let bundle = RequestTraceBundle::new(None); + bundle.start_attempt(0); + bundle.end_attempt(0); + bundle.start_attempt(1); + bundle.end_attempt(1); + + // The failed attempt's event (non-terminal). + let failed = bundle.emission(false, 0, 40, true); + assert_eq!(failed.spans.len(), 1, "attempt span only"); + let failed_span = failed.spans[0]; + assert_eq!(failed_span.role, SpanRole::Attempt); + + // The winner's terminal event. + let terminal = bundle.emission(true, 1, 60, true); + let server = terminal + .spans + .iter() + .find(|s| s.role == SpanRole::Server) + .unwrap(); + let logical = terminal + .spans + .iter() + .find(|s| s.role == SpanRole::Logical) + .unwrap(); + let attempt = terminal + .spans + .iter() + .find(|s| s.role == SpanRole::Attempt) + .unwrap(); + + assert_eq!(failed.trace_id, terminal.trace_id); + assert_eq!(logical.parent_span_id, Some(server.span_id)); + assert_eq!(attempt.parent_span_id, Some(logical.span_id)); + assert_eq!(failed_span.parent_span_id, Some(logical.span_id)); + assert_ne!(failed_span.span_id, attempt.span_id); + + // Bracketing: logical inside server, attempts inside logical. + assert!(server.start_unix_nano <= logical.start_unix_nano); + assert!(logical.end_unix_nano <= server.end_unix_nano); + assert!(logical.start_unix_nano <= failed_span.start_unix_nano); + assert!(attempt.end_unix_nano <= logical.end_unix_nano); + } + + /// Ids must be identical across repeated snapshots — the delivery + /// retry re-encodes from the same emission, but a second EXPORTER also + /// snapshots independently, and both must agree. + #[test] + fn snapshots_are_id_stable() { + let bundle = RequestTraceBundle::new(None); + bundle.start_attempt(0); + bundle.end_attempt(0); + let a = bundle.emission(false, 0, 10, true); + let b = bundle.emission(false, 0, 10, true); + assert_eq!(a.trace_id, b.trace_id); + assert_eq!(a.spans[0].span_id.to_hex(), b.spans[0].span_id.to_hex()); + assert_eq!(a.spans[0].start_unix_nano, b.spans[0].start_unix_nano); + } + + /// The streaming winner: `record()` stamped commit-time, the handler + /// measured the full stream — the longer real measurement wins, but + /// never past the emission instant (a claimed duration cannot push a + /// span end into the future). + #[test] + fn terminal_emission_extends_the_winning_attempt_to_the_measured_duration() { + let bundle = RequestTraceBundle::new(None); + bundle.start_attempt(0); + bundle.end_attempt(0); // commit-time stamp, ~0ms after start + std::thread::sleep(std::time::Duration::from_millis(15)); + // Terminal emission happens at real stream end; the handler + // measured 10ms of streaming after the commit stamp. + let em = bundle.emission(true, 0, 10, true); + let attempt = em + .spans + .iter() + .find(|s| s.role == SpanRole::Attempt) + .unwrap(); + let dur = attempt.end_unix_nano - attempt.start_unix_nano; + assert!( + dur >= 10_000_000, + "attempt span must cover the measured 10ms stream, got {dur}ns" + ); + let server = em + .spans + .iter() + .find(|s| s.role == SpanRole::Server) + .unwrap(); + assert!( + attempt.end_unix_nano <= server.end_unix_nano, + "a measured duration must never push a span past the emission instant" + ); + } + + /// Cache hits and pre-dispatch blocks: SERVER span alone — no + /// fictitious upstream CLIENT span. + #[test] + fn no_attempts_and_no_latency_emits_server_only() { + let bundle = RequestTraceBundle::new(None); + let em = bundle.emission(true, 0, 0, false); + assert_eq!(em.spans.len(), 1); + assert_eq!(em.spans[0].role, SpanRole::Server); + assert_eq!(em.carrier_role(), Some(SpanRole::Server)); + } + + /// The finding the `dispatched_upstream` parameter exists for: a cache + /// hit and a pre-dispatch block carry the HANDLER's elapsed time in + /// `upstream_latency_ms`, so a latency-based inference would fabricate + /// an upstream CLIENT span for a request that contacted no provider. + #[test] + fn undispatched_terminal_with_nonzero_latency_still_emits_server_only() { + let bundle = RequestTraceBundle::new(None); + let em = bundle.emission(true, 0, 40, false); + assert_eq!(em.spans.len(), 1); + assert_eq!(em.spans[0].role, SpanRole::Server); + } + + /// A duplicate terminal from a family without per-attempt tracking + /// degrades to NOTHING — it must not fall into the synthesize branch + /// and invent a second upstream span for a request that dispatched + /// once. + #[test] + fn duplicate_untracked_terminal_does_not_emit_a_synthetic_attempt() { + let bundle = RequestTraceBundle::new(None); + let first = bundle.emission(true, 0, 25, true); + assert_eq!(first.spans.len(), 2); + let duplicate = bundle.emission(true, 0, 25, true); + assert!(duplicate.spans.is_empty()); + } + + /// A real sub-millisecond upstream call must still get its span — the + /// caller's `dispatched_upstream` is the only signal, never latency. + #[test] + fn dispatched_zero_latency_still_emits_the_client_child() { + let bundle = RequestTraceBundle::new(None); + let em = bundle.emission(true, 0, 0, true); + assert_eq!(em.spans.len(), 2); + assert!(em.spans.iter().any(|s| s.role == SpanRole::Logical)); + } + + /// A failed ensemble: panel sub-call spans shipped (non-terminal, + /// parented under the logical id), then the judge failed and the outer + /// error arm emits the terminal WITHOUT having dispatched. The + /// already-shipped children's parent must still be exported. + #[test] + fn undispatched_terminal_still_exports_the_parent_of_shipped_subcall_spans() { + let bundle = RequestTraceBundle::new(None); + // Two panel members shipped their synthesized spans. + let member = bundle.emission(false, 0, 12, true); + assert_eq!(member.spans.len(), 1); + let member_parent = member.spans[0].parent_span_id.expect("parented"); + let _ = bundle.emission(false, 1, 15, true); + + // The pre-dispatch-shaped terminal (judge failed → outer error arm). + let terminal = bundle.emission(true, 0, 40, false); + let server = terminal + .spans + .iter() + .find(|s| s.role == SpanRole::Server) + .expect("server span"); + let logical = terminal + .spans + .iter() + .find(|s| s.role == SpanRole::Logical) + .expect("the shipped children's parent must be exported"); + assert_eq!(logical.span_id, member_parent); + assert_eq!(logical.parent_span_id, Some(server.span_id)); + assert!(logical.start_unix_nano <= member.spans[0].start_unix_nano); + assert!(member.spans[0].end_unix_nano <= logical.end_unix_nano); + } + + /// Families that dispatch without per-attempt tracking (MCP / A2A / + /// jobs): one CLIENT span under SERVER, placed by measured duration. + #[test] + fn no_attempts_with_latency_emits_one_client_child() { + let bundle = RequestTraceBundle::new(None); + let em = bundle.emission(true, 0, 25, true); + let server = em + .spans + .iter() + .find(|s| s.role == SpanRole::Server) + .unwrap(); + let logical = em + .spans + .iter() + .find(|s| s.role == SpanRole::Logical) + .unwrap(); + assert_eq!(em.spans.len(), 2); + assert_eq!(logical.parent_span_id, Some(server.span_id)); + assert!(logical.start_unix_nano >= server.start_unix_nano); + assert!(logical.end_unix_nano <= server.end_unix_nano); + assert_eq!(em.carrier_role(), Some(SpanRole::Logical)); + } + + /// A double terminal emit (an upstream bug) must not ship two SERVER + /// spans. + #[test] + fn second_terminal_emission_degrades_to_attempt_only() { + let bundle = RequestTraceBundle::new(None); + bundle.start_attempt(0); + bundle.end_attempt(0); + let first = bundle.emission(true, 0, 10, true); + assert!(first.spans.iter().any(|s| s.role == SpanRole::Server)); + let second = bundle.emission(true, 0, 10, true); + assert!( + !second.spans.iter().any(|s| s.role == SpanRole::Server), + "duplicate SERVER span" + ); + } +} diff --git a/crates/aisix-obs/src/usage.rs b/crates/aisix-obs/src/usage.rs index 003d14a44..155116c7c 100644 --- a/crates/aisix-obs/src/usage.rs +++ b/crates/aisix-obs/src/usage.rs @@ -52,6 +52,17 @@ pub struct UsageEvent { /// `x-aisix-call-id` response header carries so logs join. pub request_id: String, + /// The W3C trace id every OTLP span of this request shares, as 32 + /// lowercase-hex chars (AISIX-Cloud#1279). The public correlation key + /// between a usage row and a trace backend: the dashboard's + /// `trace_ui_url_template` can substitute it as `{trace_id}` to link a + /// log row straight to the trace. Deliberately the ONLY trace field on + /// this wire contract — span ids and boundaries stay on the exporter + /// path (`SinkRecord::trace`). Empty when the emitting path predates + /// the trace bundle. + #[serde(default, skip_serializing_if = "String::is_empty")] + pub trace_id: String, + /// Wall-clock time the upstream call completed, RFC 3339 (UTC). pub occurred_at: String, diff --git a/crates/aisix-proxy/src/a2a.rs b/crates/aisix-proxy/src/a2a.rs index c7b11e460..952aa3aad 100644 --- a/crates/aisix-proxy/src/a2a.rs +++ b/crates/aisix-proxy/src/a2a.rs @@ -136,13 +136,18 @@ pub async fn a2a_endpoint( .get::() .map(|r| r.0.clone()) .unwrap_or_else(new_request_id); + // The request's trace bundle (AISIX-Cloud#1279), minted beside the id. + let trace = request + .extensions() + .get::>() + .cloned(); let api_key_id = auth.entry.id.clone(); let http_method = request.method().clone(); // `dispatch` takes the key by value; the terminal emit below still needs // the caller's team / user labels (the handle is an `Arc` clone). let caller_auth = auth.clone(); - let response = dispatch(auth, &agent, &state, request, &request_id).await; + let response = dispatch(auth, &agent, &state, request, &request_id, trace).await; let elapsed = started.elapsed(); let status = response.status().as_u16(); @@ -192,6 +197,7 @@ async fn dispatch( state: &ProxyState, request: Request, request_id: &str, + trace: Option>, ) -> Response { // Resolve the agent from the live snapshot. A disabled agent is treated as // absent — not served, same as a missing one. @@ -287,6 +293,8 @@ async fn dispatch( &call, response.status().as_u16(), Duration::ZERO, + trace.as_ref(), + /* dispatched */ false, ); return response; } @@ -299,6 +307,7 @@ async fn dispatch( state, &snapshot, request_id, + trace.clone(), upstream, value, call, @@ -331,6 +340,8 @@ async fn dispatch( &call, StatusCode::OK.as_u16(), latency, + trace.as_ref(), + /* dispatched */ true, ); axum::Json(response_value).into_response() } @@ -346,6 +357,8 @@ async fn dispatch( &call, status.as_u16(), latency, + trace.as_ref(), + /* dispatched */ true, ); a2a_error_response(rpc_id, status, &err.to_string()) } @@ -380,6 +393,9 @@ struct StreamUsageOnDrop { /// caller and is downgraded if the stream later faults — once the headers /// are out, the usage event is the only place that can say so. status: u16, + /// The request's trace bundle (AISIX-Cloud#1279) — the Drop emit is + /// the request's terminal emission, so it carries the terminal spans. + trace: Option>, _concurrency: aisix_ratelimit::StreamConcurrencyGuard, } @@ -411,6 +427,8 @@ impl Drop for StreamUsageOnDrop { &self.call, status, self.started.elapsed(), + self.trace.as_ref(), + /* dispatched */ true, ); } } @@ -427,6 +445,7 @@ async fn dispatch_stream( state: &ProxyState, snapshot: &aisix_core::AisixSnapshot, request_id: &str, + trace: Option>, upstream: aisix_a2a::A2aUpstream, request: serde_json::Value, call: A2aCall, @@ -452,6 +471,8 @@ async fn dispatch_stream( &call, status.as_u16(), started.elapsed(), + trace.as_ref(), + /* dispatched */ true, ); return a2a_error_response(rpc_id, status, &err.to_string()); } @@ -465,6 +486,7 @@ async fn dispatch_stream( call, started, status: StatusCode::OK.as_u16(), + trace, _concurrency: reservation.into_stream_hold(), }; let agent_label = agent.to_string(); @@ -701,6 +723,10 @@ fn emit_a2a_usage( call: &A2aCall, status_code: u16, latency: Duration, + trace: Option<&std::sync::Arc>, + // Whether the call reached the upstream agent — false for a quota + // rejection, which refuses before any upstream contact. + dispatched: bool, ) { // No model resolves on this endpoint, so the estimator falls back to its // default encoding — the same thing it does for any non-OpenAI model. @@ -777,22 +803,24 @@ fn emit_a2a_usage( ); // As with MCP: an agent call has no model and no ProviderKey, so the // attribution labels are the placeholder (AISIX-Cloud#1317). - state - .usage_sink - .try_emit("a2a", event.clone(), aisix_obs::UsageEventLabels::default()); - let exporters = crate::usage_attr::live_exporters(state, snap); // Opt-in content capture, on the same terms as every other endpoint: only // an exporter configured for full content sees the words, and they never - // travel to the control plane — the usage event above carries counts - // only. The captured text is the message parts, not the JSON-RPC - // envelopes, so it reads as a prompt and a completion rather than as - // protocol scaffolding. + // travel to the control plane — the usage event carries counts only. The + // captured text is the message parts, not the JSON-RPC envelopes, so it + // reads as a prompt and a completion rather than as protocol scaffolding. + let exporters = crate::usage_attr::live_exporters(state, snap); let captured = content_capture_cap(exporters.iter().map(|e| &e.value)) .map(|cap| CapturedContent::new(&call.text.request, &response_text, cap as usize)); - state.otlp_fan_out.fan_out( - &event, + crate::usage_attr::emit_usage( + state, + snap, + "a2a", + event, + aisix_obs::UsageEventLabels::default(), captured.as_ref(), - exporters.iter().map(|e| &e.value), + trace, + /* terminal */ true, + dispatched, ); } diff --git a/crates/aisix-proxy/src/attempt.rs b/crates/aisix-proxy/src/attempt.rs index fe9fd1dc6..5a650e6bc 100644 --- a/crates/aisix-proxy/src/attempt.rs +++ b/crates/aisix-proxy/src/attempt.rs @@ -90,6 +90,11 @@ pub(crate) struct RoutingTelemetry { /// THIS group have to fall back", so the group is the useful key and /// the target it moved to is the second label. requested_model: String, + /// The request's trace bundle (AISIX-Cloud#1279). [`Self::begin_attempt`] + /// and [`Self::record`] stamp each attempt's span start/end into it — + /// the same chokepoints the per-attempt metrics use, so a dispatch + /// path cannot record an attempt without its span boundaries. + trace: Option>, } impl RoutingTelemetry { @@ -103,6 +108,15 @@ impl RoutingTelemetry { } } + /// Attach the request's trace bundle so attempt boundaries land on it. + pub fn with_trace( + mut self, + trace: Option>, + ) -> Self { + self.trace = trace; + self + } + /// Record one resolved attempt: emit its per-attempt metrics, then /// keep the record for the per-attempt usage events. /// @@ -127,6 +141,12 @@ impl RoutingTelemetry { &rec.target_model, ); } + // The attempt settled — stamp its span end (AISIX-Cloud#1279). For + // a streaming winner this is commit time; the terminal emission + // extends it to the measured stream end. + if let Some(trace) = &self.trace { + trace.end_attempt(rec.index); + } self.attempts.push(rec); } @@ -144,6 +164,11 @@ impl RoutingTelemetry { "retry" }; self.last_target = Some(display_name.to_string()); + // Mint the attempt's span id and stamp its start at the real + // dispatch boundary (AISIX-Cloud#1279). + if let Some(trace) = &self.trace { + trace.start_attempt(index); + } (index, kind) } diff --git a/crates/aisix-proxy/src/audio.rs b/crates/aisix-proxy/src/audio.rs index 07cdaf3dc..dc8143deb 100644 --- a/crates/aisix-proxy/src/audio.rs +++ b/crates/aisix-proxy/src/audio.rs @@ -2023,16 +2023,19 @@ fn emit_usage_event( crate::usage_attr::apply_pk_telemetry(&mut event, pk); // Handler label "audio" — bucketed prometheus counter (#408). crate::usage_attr::apply_jwt_identity(&mut event, client.jwt.as_ref()); - let usage_model = crate::usage_attr::usage_event_model_label(snap, &event.requested_model); - state.usage_sink.try_emit( + let usage_model = + crate::usage_attr::usage_event_model_label(snap, &event.requested_model).into_owned(); + crate::usage_attr::emit_usage( + state, + snap, "audio", - event.clone(), + event, crate::usage_attr::usage_event_labels(&usage_model, pk), + content, + client.trace.as_ref(), + /* terminal */ true, + /* dispatched */ true, ); - let exporters = crate::usage_attr::live_exporters(state, snap); - state - .otlp_fan_out - .fan_out(&event, content, exporters.iter().map(|e| &e.value)); // Speech (TTS) reports no tokens at all, so this is a no-op there; the // transcription routes report them when the model supplies a usage block. let owned_caller = crate::request_metrics::Caller::from_api_key_id(snap, api_key_id); diff --git a/crates/aisix-proxy/src/chat.rs b/crates/aisix-proxy/src/chat.rs index 2881fd326..61c0c58f7 100644 --- a/crates/aisix-proxy/src/chat.rs +++ b/crates/aisix-proxy/src/chat.rs @@ -212,6 +212,9 @@ pub async fn chat_completions( // The winner's success event carries the content. /* content_for_last */ None, + // ...and the terminal trace spans. + /* terminal_last */ + false, ); // The streaming path wires the WINNER's telemetry into the SSE // stream's on_complete callback (it has to wait for the @@ -283,6 +286,11 @@ pub async fn chat_completions( /* guardrail_blocked */ false, &client, success.captured_content.take(), + /* terminal */ true, + // A cache hit records no attempt AND contacted no + // upstream; any other no-attempt success did. + /* dispatched */ + winner.is_some() || success.cache_hit_layer.is_none(), ); } // Inject x-ratelimit-* headers so OpenAI SDK clients see the @@ -494,6 +502,11 @@ pub async fn chat_completions( &applied_guardrails, &routing, content_for_last, + // All-failed (no charge): the last failed attempt is the + // request's terminal emission. With a charge, the + // billed-then-blocked event below is terminal instead. + /* terminal_last */ + charge.is_none(), ); // Terminal event: // - output-blocked (charge present): the WINNING attempt was @@ -568,6 +581,10 @@ pub async fn chat_completions( guardrail_blocked, &client, failure_content.take(), + /* terminal */ true, + // Billed-then-blocked: the upstream answered. + /* dispatched */ + true, ); } None if routing.attempts.is_empty() => { @@ -601,11 +618,17 @@ pub async fn chat_completions( guardrail_blocked, &client, failure_content.take(), + /* terminal */ true, + // Pre-dispatch failure: no upstream was contacted; + // `elapsed` is handler time, not an upstream call. + /* dispatched */ + false, ); } None => { // All attempts failed; each was emitted above (the last - // one carrying the captured request body). + // one carrying the captured request body — and, per + // AISIX-Cloud#1279, the trace's terminal spans). } } err.into_response() @@ -1478,7 +1501,8 @@ async fn dispatch( .unwrap_or(&[]); let is_routing_request = virtual_entry.value.routing.is_some() || virtual_entry.value.is_semantic(); - let mut stream_routing = RoutingTelemetry::for_request(&virtual_entry.value.display_name); + let mut stream_routing = RoutingTelemetry::for_request(&virtual_entry.value.display_name) + .with_trace(client.trace.clone()); let mut last_err: Option = None; // The un-flattened per-target quota rejection, kept only while it // is the loop's LATEST failure: on exhaustion it is surfaced @@ -2057,6 +2081,11 @@ async fn dispatch( )), _ => None, }, + // The stream's Drop-guard emit IS the request's end — + // body EOF or client drop. + /* terminal */ + true, + /* dispatched */ true, ); // #1002: comp.total_tokens is the cache-inclusive total (an // Anthropic upstream bridged to an OpenAI-shape client folds @@ -2605,7 +2634,8 @@ async fn dispatch( .unwrap_or(&[]); let is_routing_request = virtual_entry.value.routing.is_some() || virtual_entry.value.is_semantic(); - let mut routing = RoutingTelemetry::for_request(&virtual_entry.value.display_name); + let mut routing = RoutingTelemetry::for_request(&virtual_entry.value.display_name) + .with_trace(client.trace.clone()); // The winning target's own model-layer reservation (routing dispatch // only) — folded into `reservation` at the commit point below so the // member's TPM/TPD bills with the request-level layers @@ -3320,6 +3350,8 @@ async fn dispatch_ensemble( blocked, client, /* content */ None, + /* terminal */ false, + /* dispatched */ true, ); }; @@ -3691,6 +3723,8 @@ async fn dispatch_ensemble( comp.guardrail_blocked, &client_for_telem, /* content */ None, + /* terminal */ false, + /* dispatched */ true, ); } let judge_pk = @@ -3747,6 +3781,12 @@ async fn dispatch_ensemble( comp.guardrail_blocked, &client_for_telem, /* content */ None, + // The streaming ensemble emits nothing after the judge + // — this IS the request's terminal emission, same as + // the single-model stream guard (AISIX-Cloud#1279). + /* terminal */ + true, + /* dispatched */ true, ); // SLO histograms (AISIX-Cloud#1011): the handler's // record_success is stream-gated, so the ensemble stream @@ -3898,7 +3938,8 @@ async fn dispatch_ensemble( blocked: bool, bypass: &str, redactions: &crate::redact::RedactionCounts, - hits: &[aisix_core::GuardrailMonitorHit]| { + hits: &[aisix_core::GuardrailMonitorHit], + terminal_judge: bool| { for (index, member) in outcome.panel.iter().enumerate() { emit_panel_member(member, index, blocked, bypass); } @@ -3953,6 +3994,15 @@ async fn dispatch_ensemble( blocked, client, /* content */ None, + // The judge event is the ensemble's final emission. On the + // success path it is the request's terminal event (nothing + // else emits — `telemetry_handled_by_stream` suppresses the + // handler's own emit); on the blocked path the outer error + // arm emits the real terminal carrying the caller's status, + // so the judge stays non-terminal there (AISIX-Cloud#1279). + /* terminal */ + terminal_judge, + /* dispatched */ true, ); }; @@ -3997,6 +4047,7 @@ async fn dispatch_ensemble( &bypass_reason.clone().unwrap_or_default(), &input_redactions, &ensemble_monitor_hits, + /* terminal_judge */ false, ); return Err(DispatchFailure::new( Some(model_id.to_string()), @@ -4028,6 +4079,7 @@ async fn dispatch_ensemble( &bypass_reason.clone().unwrap_or_default(), &ensemble_redactions, &ensemble_monitor_hits, + /* terminal_judge */ true, ); // The synthesized answer is the client-facing response, rendered with @@ -4228,6 +4280,16 @@ fn emit_usage_event( guardrail_blocked: bool, client: &ClientContext, content: Option, + // Whether this event ends the request (AISIX-Cloud#1279): the terminal + // event carries the trace's SERVER + logical spans; a non-terminal one + // (a failed attempt, an ensemble panel member / judge sub-call) carries + // its own attempt span alone. + terminal: bool, + // Whether the work this event describes actually reached an upstream. + // A cache hit and a pre-dispatch error pass false: their events carry + // the HANDLER's elapsed time, which must not fabricate an upstream + // CLIENT span. + dispatched: bool, ) { // Per-PK telemetry attribution tags. An unresolved key (the // pre-dispatch error paths) yields default (all empty / false) tags → @@ -4298,29 +4360,30 @@ fn emit_usage_event( ..Default::default() }; crate::usage_attr::apply_jwt_identity(&mut event, client.jwt.as_ref()); - // Handler label "chat" matches the documented enumeration for - // `aisix_usage_events_emitted_total` (#408). Keep `&'static str` - // so prometheus cardinality stays bounded. - let usage_model = crate::usage_attr::usage_event_model_label(snap, &event.requested_model); - state.usage_sink.try_emit( - "chat", - event.clone(), - crate::usage_attr::usage_event_labels(&usage_model, pk), - ); // Guardrail outcome counters (#379). Recorded here — the one place every // chat path (success / error / streaming / cache-hit) funnels through — // from the same guardrail fields the UsageEvent carries. state .metrics .record_guardrail_outcome(guardrail_blocked, &event.guardrail_bypassed_reason); - // Per-env OTLP/HTTP fan-out. The snapshot's exporter table is - // empty for envs that haven't configured any, so this is a cheap - // no-op on the common path. Spawned tasks own the POST work and - // never block the request return. - let exporters = crate::usage_attr::live_exporters(state, snap); - state - .otlp_fan_out - .fan_out(&event, content.as_ref(), exporters.iter().map(|e| &e.value)); + // Handler label "chat" matches the documented enumeration for + // `aisix_usage_events_emitted_total` (#408). Keep `&'static str` + // so prometheus cardinality stays bounded. Both emit legs (CP sink + + // per-env exporter fan-out) go through the shared chokepoint, which + // also snapshots the trace spans (AISIX-Cloud#1279). + let usage_model = + crate::usage_attr::usage_event_model_label(snap, &event.requested_model).into_owned(); + crate::usage_attr::emit_usage( + state, + snap, + "chat", + event, + crate::usage_attr::usage_event_labels(&usage_model, pk), + content.as_ref(), + client.trace.as_ref(), + terminal, + dispatched, + ); } /// Defence-in-depth sanitiser for operator-defined `ProviderKey @@ -4456,6 +4519,11 @@ fn emit_failed_attempts( // success-path caller) stay content-less to avoid duplicating a large // prompt across N events. mut content_for_last: Option, + // AISIX-Cloud#1279: on the all-failed path the LAST failed attempt's + // event is the request's terminal emission (there is no other), so it + // carries the trace's SERVER + logical spans. False on the success + // path, where the winner's event is the terminal one. + terminal_last: bool, ) { let last_failed = routing.attempts.iter().rposition(|a| !a.success); for (i, rec) in routing @@ -4498,6 +4566,11 @@ fn emit_failed_attempts( /* guardrail_blocked */ false, client, content, + /* terminal */ terminal_last && Some(i) == last_failed, + // The record's own network-boundary fact: a rate-limit-refused + // attempt never reached an upstream. + /* dispatched */ + rec.dispatched, ); } } diff --git a/crates/aisix-proxy/src/client_ip.rs b/crates/aisix-proxy/src/client_ip.rs index c81c0b7d0..a79bed597 100644 --- a/crates/aisix-proxy/src/client_ip.rs +++ b/crates/aisix-proxy/src/client_ip.rs @@ -154,6 +154,13 @@ pub struct ClientContext { /// extractor; `None` when the request authenticated with the key's /// plaintext. pub jwt: Option>, + /// The request's trace identity (AISIX-Cloud#1279), minted by + /// `ensure_request_id` beside the [`RequestId`]. Handlers thread it + /// into their attempt telemetry and usage-event emission so every + /// exported span of the request shares one trace. `None` only when + /// the middleware isn't in the chain (a handler unit test with a + /// bare router) — emission then falls back to the legacy flat span. + pub trace: Option>, } /// Resolve the caller's address from the peer plus the trusted-proxy @@ -224,6 +231,10 @@ where .extensions .get::>() .cloned(), + trace: parts + .extensions + .get::>() + .cloned(), }) } } diff --git a/crates/aisix-proxy/src/completions.rs b/crates/aisix-proxy/src/completions.rs index c56659af1..97bfdfa8b 100644 --- a/crates/aisix-proxy/src/completions.rs +++ b/crates/aisix-proxy/src/completions.rs @@ -729,16 +729,19 @@ fn emit_usage_event( }; crate::usage_attr::apply_pk_telemetry(&mut event, pk); crate::usage_attr::apply_jwt_identity(&mut event, client.jwt.as_ref()); - let usage_model = crate::usage_attr::usage_event_model_label(snap, &event.requested_model); - state.usage_sink.try_emit( + let usage_model = + crate::usage_attr::usage_event_model_label(snap, &event.requested_model).into_owned(); + crate::usage_attr::emit_usage( + state, + snap, "completions", - event.clone(), + event, crate::usage_attr::usage_event_labels(&usage_model, pk), + content, + client.trace.as_ref(), + /* terminal */ true, + /* dispatched */ true, ); - let exporters = crate::usage_attr::live_exporters(state, snap); - state - .otlp_fan_out - .fan_out(&event, content, exporters.iter().map(|e| &e.value)); let owned_caller = crate::request_metrics::Caller::from_api_key_id(snap, api_key_id); crate::request_metrics::record_usage( state, diff --git a/crates/aisix-proxy/src/embeddings.rs b/crates/aisix-proxy/src/embeddings.rs index c0d6ed8d5..b1962598c 100644 --- a/crates/aisix-proxy/src/embeddings.rs +++ b/crates/aisix-proxy/src/embeddings.rs @@ -688,19 +688,19 @@ fn emit_usage_event( crate::usage_attr::apply_pk_telemetry(&mut event, pk); // Handler label "embeddings" — bucketed prometheus counter (#408). crate::usage_attr::apply_jwt_identity(&mut event, client.jwt.as_ref()); - let usage_model = crate::usage_attr::usage_event_model_label(snap, &event.requested_model); - state.usage_sink.try_emit( + let usage_model = + crate::usage_attr::usage_event_model_label(snap, &event.requested_model).into_owned(); + crate::usage_attr::emit_usage( + state, + snap, "embeddings", - event.clone(), + event, crate::usage_attr::usage_event_labels(&usage_model, pk), + content, + client.trace.as_ref(), + /* terminal */ true, + /* dispatched */ true, ); - // Per-env OTLP/HTTP fan-out — same shape as chat.rs:1334. The - // snapshot's exporter table is empty for envs that haven't - // configured any, so this is a cheap no-op on the common path. - let exporters = crate::usage_attr::live_exporters(state, snap); - state - .otlp_fan_out - .fan_out(&event, content, exporters.iter().map(|e| &e.value)); let owned_caller = crate::request_metrics::Caller::from_api_key_id(snap, api_key_id); crate::request_metrics::record_usage( state, diff --git a/crates/aisix-proxy/src/images.rs b/crates/aisix-proxy/src/images.rs index ad902855e..95f9c5f85 100644 --- a/crates/aisix-proxy/src/images.rs +++ b/crates/aisix-proxy/src/images.rs @@ -520,16 +520,19 @@ fn emit_usage_event( crate::usage_attr::apply_pk_telemetry(&mut event, pk); // Handler label "images" — bucketed prometheus counter (#408). crate::usage_attr::apply_jwt_identity(&mut event, client.jwt.as_ref()); - let usage_model = crate::usage_attr::usage_event_model_label(snap, &event.requested_model); - state.usage_sink.try_emit( + let usage_model = + crate::usage_attr::usage_event_model_label(snap, &event.requested_model).into_owned(); + crate::usage_attr::emit_usage( + state, + snap, "images", - event.clone(), + event, crate::usage_attr::usage_event_labels(&usage_model, pk), + content, + client.trace.as_ref(), + /* terminal */ true, + /* dispatched */ true, ); - let exporters = crate::usage_attr::live_exporters(state, snap); - state - .otlp_fan_out - .fan_out(&event, content, exporters.iter().map(|e| &e.value)); let owned_caller = crate::request_metrics::Caller::from_api_key_id(snap, api_key_id); crate::request_metrics::record_usage( state, diff --git a/crates/aisix-proxy/src/jobs.rs b/crates/aisix-proxy/src/jobs.rs index 28863fe23..b3e80ad1d 100644 --- a/crates/aisix-proxy/src/jobs.rs +++ b/crates/aisix-proxy/src/jobs.rs @@ -601,16 +601,19 @@ fn emit_job_usage_event( let pk = crate::usage_attr::ResolvedPk::resolve(snap, &target.pk_entry.id); crate::usage_attr::apply_pk_telemetry(&mut event, &pk); crate::usage_attr::apply_jwt_identity(&mut event, auth.jwt.as_ref()); - let usage_model = crate::usage_attr::usage_event_model_label(snap, &event.requested_model); - state.usage_sink.try_emit( + let usage_model = + crate::usage_attr::usage_event_model_label(snap, &event.requested_model).into_owned(); + crate::usage_attr::emit_usage( + state, + snap, label, - event.clone(), + event, crate::usage_attr::usage_event_labels(&usage_model, &pk), + None, + client.trace.as_ref(), + /* terminal */ true, + /* dispatched */ true, ); - let exporters = crate::usage_attr::live_exporters(state, snap); - state - .otlp_fan_out - .fan_out(&event, None, exporters.iter().map(|e| &e.value)); } #[allow(clippy::too_many_arguments)] @@ -1851,15 +1854,19 @@ async fn attribute_batch_usage( // Attribution names the identity that observed completion — the // same caller the event's api_key_id already reflects. crate::usage_attr::apply_jwt_identity(&mut event, jwt); - let usage_model = crate::usage_attr::usage_event_model_label(&snap, &event.requested_model); + let usage_model = + crate::usage_attr::usage_event_model_label(&snap, &event.requested_model).into_owned(); state.usage_sink.try_emit( "batch", event.clone(), crate::usage_attr::usage_event_labels(&usage_model, &pk), ); + // A background poll attributes usage after the fact — there is no + // live request and therefore no trace bundle; the exporter falls + // back to the legacy flat span (AISIX-Cloud#1279). state .otlp_fan_out - .fan_out(&event, None, exporters.iter().map(|e| &e.value)); + .fan_out(&event, None, None, exporters.iter().map(|e| &e.value)); tracing::info!( batch_id = %raw_batch_id, provider_model = %provider_model, diff --git a/crates/aisix-proxy/src/mcp.rs b/crates/aisix-proxy/src/mcp.rs index 99bf0485a..9130e4bfe 100644 --- a/crates/aisix-proxy/src/mcp.rs +++ b/crates/aisix-proxy/src/mcp.rs @@ -165,6 +165,11 @@ async fn serve(state: ProxyState, request: Request, scope: Option) -> Re .get::() .map(|r| r.0.clone()) .unwrap_or_else(new_request_id); + // The request's trace bundle (AISIX-Cloud#1279), minted beside the id. + let trace = request + .extensions() + .get::>() + .cloned(); let api_key_id = auth.entry.id.clone(); let method = request.method().clone(); // `dispatch` takes the key by value; the terminal emit below still needs @@ -178,6 +183,7 @@ async fn serve(state: ProxyState, request: Request, scope: Option) -> Re &state, request, &request_id, + trace.as_ref(), ) .await; @@ -235,6 +241,7 @@ async fn dispatch( state: &ProxyState, request: Request, request_id: &str, + trace: Option<&std::sync::Arc>, ) -> Response { // One snapshot for the whole request: the scoped-server resolution below // and the gateway construction further down must see the same resource @@ -347,6 +354,8 @@ async fn dispatch( Duration::ZERO, false, Vec::new(), + trace, + /* dispatched */ false, ); return response; } @@ -418,6 +427,8 @@ async fn dispatch( Duration::ZERO, true, monitor_hits, + trace, + /* dispatched */ false, ); return jsonrpc_guardrail_block(rpc_id, "tool call", guardrail_name.as_deref()); } @@ -497,6 +508,8 @@ async fn dispatch( latency, true, monitor_hits, + trace, + /* dispatched */ true, ); return jsonrpc_guardrail_block(rpc_id, "tool result", guardrail_name.as_deref()); } @@ -517,6 +530,8 @@ async fn dispatch( latency, false, monitor_hits, + trace, + /* dispatched */ true, ); } response @@ -638,6 +653,11 @@ fn emit_tool_call_usage( latency: Duration, guardrail_blocked: bool, guardrail_monitor_hits: Vec, + trace: Option<&std::sync::Arc>, + // Whether the tool call reached its upstream MCP server — false for a + // quota rejection or an input-guardrail block, which refuse before any + // upstream contact and must not export an upstream CLIENT span. + dispatched: bool, ) { let mut event = UsageEvent { request_id: request_id.to_string(), @@ -660,17 +680,21 @@ fn emit_tool_call_usage( // A tool call resolves neither a model nor a ProviderKey, so the // attribution labels are the placeholder — present so this family has // ONE label set across every handler (AISIX-Cloud#1317). - state - .usage_sink - .try_emit("mcp", event.clone(), aisix_obs::UsageEventLabels::default()); - // #698: fan the event out to the per-env OTLP/SLS/Datadog exporters like - // every other emitter — pre-fix MCP usage reached only the CP sink, so + // #698: both emit legs (CP sink + per-env exporter fan-out) go through + // the shared chokepoint — pre-fix MCP usage reached only the CP sink, so // exporters never saw /mcp traffic. No content capture (tool args/results // are a separate surface from prompt/response). - let exporters = crate::usage_attr::live_exporters(state, snap); - state - .otlp_fan_out - .fan_out(&event, None, exporters.iter().map(|e| &e.value)); + crate::usage_attr::emit_usage( + state, + snap, + "mcp", + event, + aisix_obs::UsageEventLabels::default(), + None, + trace, + /* terminal */ true, + dispatched, + ); } /// Reject a request whose `MCP-Protocol-Version` header names a version diff --git a/crates/aisix-proxy/src/messages.rs b/crates/aisix-proxy/src/messages.rs index 21cf88782..ebbbcd64c 100644 --- a/crates/aisix-proxy/src/messages.rs +++ b/crates/aisix-proxy/src/messages.rs @@ -241,6 +241,9 @@ pub async fn messages( // The winner's success event carries the content. /* content_for_last */ None, + // ...and the terminal trace spans. + /* terminal_last */ + false, ); if !usage_handled_by_stream { // Winning-attempt classification (#655). Direct models have @@ -288,6 +291,8 @@ pub async fn messages( redaction_counts.clone(), monitor_hits.clone(), captured_content, + /* terminal */ true, + /* dispatched */ true, ); } response @@ -385,6 +390,10 @@ pub async fn messages( &applied_guardrails, &routing, content_for_last, + // All-failed: the last failed attempt is the terminal + // emission; the pre-dispatch branch below covers empty. + /* terminal_last */ + !routing.attempts.is_empty(), ); // Pre-dispatch failure (model-not-found, auth, budget, guardrail // block before any upstream attempt) records no attempts — emit a @@ -418,6 +427,10 @@ pub async fn messages( redaction_counts.clone(), monitor_hits.clone(), failure_content.take(), + /* terminal */ true, + // Pre-dispatch failure: no upstream was contacted. + /* dispatched */ + false, ); } // /v1/messages must return Anthropic-shape error envelope @@ -453,6 +466,10 @@ fn emit_failed_attempts_anthropic( // the one whose status the caller saw. Other attempts (and the // success-path caller) stay content-less. mut content_for_last: Option, + // AISIX-Cloud#1279: on the all-failed path the LAST failed attempt's + // event is the request's terminal emission, so it carries the trace's + // SERVER + logical spans. False on the success path. + terminal_last: bool, ) { let last_failed = routing.attempts.iter().rposition(|a| !a.success); for (i, rec) in routing @@ -493,6 +510,11 @@ fn emit_failed_attempts_anthropic( crate::redact::RedactionCounts::new(), Vec::new(), content, + /* terminal */ terminal_last && Some(i) == last_failed, + // The record's own network-boundary fact: a rate-limit-refused + // attempt never reached an upstream. + /* dispatched */ + rec.dispatched, ); } } @@ -679,7 +701,8 @@ async fn dispatch( // pre-dispatch 4xx); if this endpoint ever grows semantic support, // widen this flag or the deferred policies are silently skipped. let is_routing_request = model_entry.value.routing.is_some(); - let mut routing = RoutingTelemetry::for_request(&model_entry.value.display_name); + let mut routing = RoutingTelemetry::for_request(&model_entry.value.display_name) + .with_trace(client.trace.clone()); // Walk targets, failing over to the next only on a retryable upstream // failure. A 4xx / config error is returned as-is — retrying other @@ -1445,6 +1468,10 @@ async fn anthropic_passthrough_dispatch( )), _ => None, }, + // Drop-guard emit at stream end = the request's end. + /* terminal */ + true, + /* dispatched */ true, ); }, ); @@ -2089,6 +2116,10 @@ async fn cross_provider_dispatch( )), _ => None, }, + // Drop-guard emit at stream end = the request's end. + /* terminal */ + true, + /* dispatched */ true, ); }, ); @@ -2781,6 +2812,13 @@ fn emit_anthropic_usage_event( // output merged. guardrail_monitor_hits: Vec, content: Option, + // Whether this event ends the request (AISIX-Cloud#1279): the terminal + // event carries the trace's SERVER + logical spans; a failed attempt's + // event carries its own attempt span alone. + terminal: bool, + // Whether the work this event describes actually reached an upstream — + // false for a pre-dispatch failure, whose `elapsed` is handler time. + dispatched: bool, ) { // Per-PK telemetry attribution (#302 M17 / AISIX-Cloud#436). // Same shape as chat.rs's emit_usage_event — look up the @@ -2830,15 +2868,19 @@ fn emit_anthropic_usage_event( // path. Bucketed prometheus counter (#408). crate::usage_attr::apply_jwt_identity(&mut event, client.jwt.as_ref()); let usage_model = crate::usage_attr::usage_event_model_label(snap, &event.requested_model); - state.usage_sink.try_emit( + // The metric code below still reads `event`, so the chokepoint gets + // its own copy (it stamps `trace_id` on the emitted one). + crate::usage_attr::emit_usage( + state, + snap, "messages", event.clone(), crate::usage_attr::usage_event_labels(&usage_model, pk), + content.as_ref(), + client.trace.as_ref(), + terminal, + dispatched, ); - let exporters = crate::usage_attr::live_exporters(state, snap); - state - .otlp_fan_out - .fan_out(&event, content.as_ref(), exporters.iter().map(|e| &e.value)); // Cache-inclusive canonical total: Anthropic reports cache tokens as // counters separate from prompt_tokens, so prompt+completion undercounts // cached traffic (#995/#906). Shared by the LLM-usage total metric and the diff --git a/crates/aisix-proxy/src/passthrough_route.rs b/crates/aisix-proxy/src/passthrough_route.rs index 193e7b684..a80d8f8ef 100644 --- a/crates/aisix-proxy/src/passthrough_route.rs +++ b/crates/aisix-proxy/src/passthrough_route.rs @@ -110,6 +110,13 @@ const ALWAYS_STRIP: &[&str] = &[ "proxy-authorization", "te", "trailer", + "traceparent", + // W3C trace context (AISIX-Cloud#1279): stripped unconditionally, same + // as the standard pipeline's never-forward set — the caller's trace + // ids are not the provider's to see, and passthrough forwarded them + // verbatim before this entry existed. A future provider-side + // propagation opt-in would inject the gateway's own context instead. + "tracestate", "transfer-encoding", "upgrade", // Gateway-owned correlation id: the dispatch sets its own value, and @@ -402,6 +409,7 @@ pub async fn entry( &usage_model, &crate::usage_attr::ResolvedPk::unresolved(), ), + client.trace.as_ref(), ); error.into_response() } @@ -806,6 +814,7 @@ async fn dispatch( let mut telemetry = RouteTelemetry { state: state.clone(), route_name: route.name.clone(), + trace: client.trace.clone(), provider_label: pk_entry .as_ref() .map(|pk| pk.value.provider.to_ascii_lowercase()) @@ -1972,6 +1981,9 @@ struct RouteTelemetry { state: ProxyState, route_name: String, provider_label: String, + /// The request's trace bundle (AISIX-Cloud#1279) — the Drop emit is + /// the request's terminal emission, so it carries the terminal spans. + trace: Option>, pk_id: String, method: Method, path: String, @@ -2142,17 +2154,15 @@ impl RouteTelemetry { event.auth_type = "anonymous".to_string(); } let usage_model = crate::usage_attr::usage_event_model_label( - &self.state.snapshot.load(), + // The snapshot loaded above: a config swap between two loads + // would make this label disagree with the emit's attribution. + &snapshot, &event.requested_model, - ); - self.state.usage_sink.try_emit( - "passthrough_route", - event.clone(), - crate::usage_attr::usage_event_labels(&usage_model, &pk), - ); + ) + .into_owned(); // Captured content rides ONLY on the exporter fan-out, per the - // content_mode invariant (never the CP telemetry path above). + // content_mode invariant (never the CP telemetry path). let content = match (&self.captured_prompt, self.content_cap) { (Some(prompt), Some(cap)) => Some(aisix_obs::CapturedContent::new( prompt, @@ -2161,11 +2171,18 @@ impl RouteTelemetry { )), _ => None, }; - let exporters = crate::usage_attr::live_exporters(&self.state, &snapshot); - self.state.otlp_fan_out.fan_out( - &event, + crate::usage_attr::emit_usage( + &self.state, + &snapshot, + "passthrough_route", + event, + crate::usage_attr::usage_event_labels(&usage_model, &pk), content.as_ref(), - exporters.iter().map(|e| &e.value), + self.trace.as_ref(), + // The Drop emit is the request's end — body EOF or client drop. + /* terminal */ + true, + /* dispatched */ true, ); } } diff --git a/crates/aisix-proxy/src/realtime.rs b/crates/aisix-proxy/src/realtime.rs index ced6bc4c3..616d6bf15 100644 --- a/crates/aisix-proxy/src/realtime.rs +++ b/crates/aisix-proxy/src/realtime.rs @@ -776,16 +776,19 @@ async fn run_session( }; crate::usage_attr::apply_pk_telemetry(&mut event, &pk); crate::usage_attr::apply_jwt_identity(&mut event, auth.jwt.as_ref()); - let usage_model = crate::usage_attr::usage_event_model_label(&snap, &event.requested_model); - state.usage_sink.try_emit( + let usage_model = + crate::usage_attr::usage_event_model_label(&snap, &event.requested_model).into_owned(); + crate::usage_attr::emit_usage( + &state, + &snap, "realtime", event.clone(), crate::usage_attr::usage_event_labels(&usage_model, &pk), + None, + client.trace.as_ref(), + /* terminal */ true, + /* dispatched */ true, ); - let exporters = crate::usage_attr::live_exporters(&state, &snap); - state - .otlp_fan_out - .fan_out(&event, None, exporters.iter().map(|e| &e.value)); // A realtime session bills real tokens against a real model, and this is // the only place that knows the session's totals. Its cost is resolved // here too, unlike the other endpoints, so it is the one non-chat surface diff --git a/crates/aisix-proxy/src/request_id.rs b/crates/aisix-proxy/src/request_id.rs index 5af0368ef..a8e5353d6 100644 --- a/crates/aisix-proxy/src/request_id.rs +++ b/crates/aisix-proxy/src/request_id.rs @@ -129,6 +129,19 @@ pub(crate) async fn ensure_request_id( .unwrap_or_else(new_request_id); request.extensions_mut().insert(RequestId(id.clone())); + // The request's trace identity (AISIX-Cloud#1279), minted at the same + // single point as the request id so every usage event and exported + // span of this request shares one trace. A valid inbound W3C + // `traceparent` continues the caller's trace; anything malformed — + // or a duplicated header — degrades to a locally-rooted one, and the + // caller's `tracestate` is only kept alongside a valid parent. + let remote = remote_trace_context(request.headers()); + request + .extensions_mut() + .insert(std::sync::Arc::new(aisix_obs::RequestTraceBundle::new( + remote, + ))); + let span = tracing::info_span!("request", request_id = %id); let mut response = next.run(request).instrument(span).await; @@ -142,6 +155,51 @@ pub(crate) async fn ensure_request_id( response } +/// The caller's W3C trace context, when it arrived intact. +/// +/// Strict on purpose — this value decides what the gateway's own +/// telemetry is filed under. Exactly one `traceparent` header (the spec +/// invalidates a duplicated one), parsed by the strict v00 grammar in +/// `aisix_obs::parse_traceparent`; `tracestate` (spec: multiple headers +/// combine as a comma-joined list) is screened and kept only when the +/// parent it travelled with was valid. Invalid → `None` → a +/// locally-rooted trace, never a rejected request. +fn remote_trace_context(headers: &axum::http::HeaderMap) -> Option { + let mut parents = headers.get_all(aisix_obs::TRACEPARENT_HEADER).iter(); + let parent = parents.next()?; + if parents.next().is_some() { + tracing::debug!("multiple traceparent headers; starting a local trace"); + return None; + } + let Some((trace_id, parent_span_id, flags)) = + parent.to_str().ok().and_then(aisix_obs::parse_traceparent) + else { + tracing::debug!("unusable traceparent header; starting a local trace"); + return None; + }; + let tracestate = { + // All-or-nothing: joining only the readable values would forward a + // list the caller never sent. One unreadable value discards the + // whole tracestate (the traceparent stays valid — the spec treats + // state as optional refinement). + let values: Result, _> = headers + .get_all(aisix_obs::TRACESTATE_HEADER) + .iter() + .map(|v| v.to_str()) + .collect(); + match values { + Ok(values) if !values.is_empty() => aisix_obs::screen_tracestate(&values.join(",")), + _ => None, + } + }; + Some(aisix_obs::RemoteTraceContext { + trace_id, + parent_span_id, + flags, + tracestate, + }) +} + /// Stream adapter that enters `span` for the duration of every /// `poll_next`. struct InSpan { diff --git a/crates/aisix-proxy/src/rerank.rs b/crates/aisix-proxy/src/rerank.rs index cf4b128b3..504f4d66d 100644 --- a/crates/aisix-proxy/src/rerank.rs +++ b/crates/aisix-proxy/src/rerank.rs @@ -701,16 +701,19 @@ fn emit_usage_event( // chat / messages / responses / embeddings (AISIX-Cloud#867 parity). crate::usage_attr::apply_pk_telemetry(&mut event, pk); crate::usage_attr::apply_jwt_identity(&mut event, client.jwt.as_ref()); - let usage_model = crate::usage_attr::usage_event_model_label(snap, &event.requested_model); - state.usage_sink.try_emit( + let usage_model = + crate::usage_attr::usage_event_model_label(snap, &event.requested_model).into_owned(); + crate::usage_attr::emit_usage( + state, + snap, "rerank", - event.clone(), + event, crate::usage_attr::usage_event_labels(&usage_model, pk), + content, + client.trace.as_ref(), + /* terminal */ true, + /* dispatched */ true, ); - let exporters = crate::usage_attr::live_exporters(state, snap); - state - .otlp_fan_out - .fan_out(&event, content, exporters.iter().map(|e| &e.value)); let owned_caller = crate::request_metrics::Caller::from_api_key_id(snap, api_key_id); crate::request_metrics::record_usage( state, diff --git a/crates/aisix-proxy/src/responses.rs b/crates/aisix-proxy/src/responses.rs index 59645e0f0..6119914d4 100644 --- a/crates/aisix-proxy/src/responses.rs +++ b/crates/aisix-proxy/src/responses.rs @@ -284,6 +284,9 @@ pub async fn responses( // The winner's success event carries the content. /* content_for_last */ None, + // The winner's event carries the terminal spans. + /* terminal_last */ + false, ); // Issue #404: emit UsageEvent so cp-api's budget ledger // and customer-facing /logs analytics see /v1/responses @@ -350,6 +353,8 @@ pub async fn responses( redaction_counts.clone(), monitor_hits.clone(), success.captured_content.as_ref(), + /* terminal */ true, + /* dispatched */ true, ); } } @@ -443,6 +448,10 @@ pub async fn responses( &client, &routing, content_for_last, + // All-failed: the last failed attempt is the terminal + // emission; the pre-dispatch branch below covers empty. + /* terminal_last */ + !routing.attempts.is_empty(), ); // Pre-dispatch failure (model-not-found, auth, budget) records no // attempts — emit a single terminal event carrying the failure @@ -470,6 +479,10 @@ pub async fn responses( redaction_counts.clone(), monitor_hits.clone(), failure_content.take(), + /* terminal */ true, + // Pre-dispatch failure: no upstream contacted. + /* dispatched */ + false, ); } err.into_response() @@ -645,7 +658,8 @@ async fn dispatch( // pre-dispatch 4xx); if this endpoint ever grows semantic support, // widen this flag or the deferred policies are silently skipped. let is_routing_request = model_entry.value.routing.is_some(); - let mut routing = RoutingTelemetry::for_request(&model_entry.value.display_name); + let mut routing = RoutingTelemetry::for_request(&model_entry.value.display_name) + .with_trace(client.trace.clone()); // Walk the targets, failing over on a retryable failure. Streaming and // non-streaming share this loop: the per-target dispatch branches // internally and, for streaming, only returns Ok once the first chunk @@ -1605,6 +1619,8 @@ async fn responses_to_target( input_redactions.clone(), monitor_hits, captured_content.as_ref(), + /* terminal */ true, + /* dispatched */ true, ); }, ); @@ -2100,6 +2116,8 @@ async fn responses_cross_provider_to_target( merged }, captured_content.as_ref(), + /* terminal */ true, + /* dispatched */ true, ); }, ); @@ -2920,6 +2938,11 @@ fn emit_usage_event( // Captured request/response content for content-capturing exporters // (AISIX-Cloud#947). Forwarded only to `fan_out`, never to the CP sink. content: Option<&CapturedContent>, + // Whether this event ends the request (AISIX-Cloud#1279). + terminal: bool, + // Whether the work this event describes actually reached an upstream — + // false for a pre-dispatch failure, whose `elapsed` is handler time. + dispatched: bool, ) { let tags = pk.telemetry_tags(); let mut event = UsageEvent { @@ -2962,15 +2985,17 @@ fn emit_usage_event( }; crate::usage_attr::apply_jwt_identity(&mut event, client.jwt.as_ref()); let usage_model = crate::usage_attr::usage_event_model_label(snap, &event.requested_model); - state.usage_sink.try_emit( + crate::usage_attr::emit_usage( + state, + snap, "responses", event.clone(), crate::usage_attr::usage_event_labels(&usage_model, pk), + content, + client.trace.as_ref(), + terminal, + dispatched, ); - let exporters = crate::usage_attr::live_exporters(state, snap); - state - .otlp_fan_out - .fan_out(&event, content, exporters.iter().map(|e| &e.value)); // AISIX-Cloud#1044: token volume by inbound client type × model. Codex // traffic arrives on /v1/responses, so leaving this endpoint out of the // by-client series made an allowlisted client invisible in it. All three @@ -3033,6 +3058,12 @@ fn emit_zero_token_event( // AISIX-Cloud#1013: captured (post-mask) request body for failed // requests. Forwarded only to `fan_out`, never to the CP sink. content: Option, + // Whether this event ends the request (AISIX-Cloud#1279): true for a + // pre-dispatch / all-failed terminal, false for a superseded attempt. + terminal: bool, + // Whether this attempt actually reached an upstream — false on the + // pre-dispatch path, true for a recorded failed attempt. + dispatched: bool, ) { let pk = ResolvedPk::resolve(snap, provider_key_id); let tags = pk.telemetry_tags(); @@ -3063,15 +3094,17 @@ fn emit_zero_token_event( }; crate::usage_attr::apply_jwt_identity(&mut event, client.jwt.as_ref()); let usage_model = crate::usage_attr::usage_event_model_label(snap, &event.requested_model); - state.usage_sink.try_emit( + crate::usage_attr::emit_usage( + state, + snap, "responses", event.clone(), crate::usage_attr::usage_event_labels(&usage_model, &pk), + content.as_ref(), + client.trace.as_ref(), + terminal, + dispatched, ); - let exporters = crate::usage_attr::live_exporters(state, snap); - state - .otlp_fan_out - .fan_out(&event, content.as_ref(), exporters.iter().map(|e| &e.value)); } /// Emit one zero-token `UsageEvent` per FAILED attempt of a `/v1/responses` @@ -3090,6 +3123,10 @@ fn emit_failed_attempts( // the one whose status the caller saw. Other attempts (and the // success-path caller) stay content-less. mut content_for_last: Option, + // AISIX-Cloud#1279: on the all-failed path the LAST failed attempt's + // event is the request's terminal emission, so it carries the trace's + // SERVER + logical spans. False on the success path. + terminal_last: bool, ) { let last_failed = routing.attempts.iter().rposition(|a| !a.success); for (i, rec) in routing @@ -3122,6 +3159,11 @@ fn emit_failed_attempts( crate::redact::RedactionCounts::new(), Vec::new(), content, + /* terminal */ terminal_last && Some(i) == last_failed, + // The record's own network-boundary fact: a rate-limit-refused + // attempt never reached an upstream. + /* dispatched */ + rec.dispatched, ); } } diff --git a/crates/aisix-proxy/src/usage_attr.rs b/crates/aisix-proxy/src/usage_attr.rs index c386edf4b..c796b6d49 100644 --- a/crates/aisix-proxy/src/usage_attr.rs +++ b/crates/aisix-proxy/src/usage_attr.rs @@ -384,7 +384,14 @@ pub(crate) fn emit_error_usage_event( let attributed = crate::attribution::current().unwrap_or_default(); let pk = ResolvedPk::resolve(snap, &attributed.provider_key_id); let model = usage_event_model_label(snap, requested_model); - emit_prepared_usage_event(state, snap, label, event, usage_event_labels(&model, &pk)); + emit_prepared_usage_event( + state, + snap, + label, + event, + usage_event_labels(&model, &pk), + client.trace.as_ref(), + ); } /// The [`emit_error_usage_event`] event without the emission, for a caller @@ -464,12 +471,58 @@ pub(crate) fn emit_prepared_usage_event( label: &'static str, event: UsageEvent, labels: UsageEventLabels<'_>, + trace: Option<&std::sync::Arc>, ) { + // An error event is the request's terminal (and only) emission. Its + // builder leaves `upstream_latency_ms` at 0, so no upstream span is + // derivable either way — `dispatched: false` states the common case + // (most error events never reached a provider). + emit_usage(state, snap, label, event, labels, None, trace, true, false); +} + +/// THE emission chokepoint (AISIX-Cloud#1279): every usage event leaves +/// through here, so the CP telemetry leg and the exporter fan-out cannot +/// drift, and the trace snapshot is taken in exactly one place. +/// +/// `trace` is the request's bundle (`ClientContext::trace`); `terminal` +/// says whether this event ends the request — the terminal event carries +/// the SERVER + logical spans (ending the SERVER span at NOW, the real +/// body EOF/drop), a non-terminal one carries its own attempt span alone. +/// `dispatched` is the caller's statement that this event describes work +/// that actually reached an upstream — a cache hit and a pre-dispatch +/// error pass `false` so no fictitious upstream CLIENT span is fabricated +/// from the handler-elapsed time their events carry. The event's public +/// `trace_id` is stamped here so the CP row and the exported spans always +/// agree. +#[allow(clippy::too_many_arguments)] +pub(crate) fn emit_usage( + state: &ProxyState, + snap: &AisixSnapshot, + label: &'static str, + mut event: UsageEvent, + labels: UsageEventLabels<'_>, + content: Option<&aisix_obs::CapturedContent>, + trace: Option<&std::sync::Arc>, + terminal: bool, + dispatched: bool, +) { + let emission = trace.map(|bundle| { + event.trace_id = bundle.trace_id_hex(); + bundle.emission( + terminal, + event.attempt_index, + event.upstream_latency_ms, + dispatched, + ) + }); state.usage_sink.try_emit(label, event.clone(), labels); let exporters = live_exporters(state, snap); - state - .otlp_fan_out - .fan_out(&event, None, exporters.iter().map(|e| &e.value)); + state.otlp_fan_out.fan_out( + &event, + content, + emission.as_ref(), + exporters.iter().map(|e| &e.value), + ); } #[cfg(test)] diff --git a/crates/aisix-proxy/src/videos.rs b/crates/aisix-proxy/src/videos.rs index 827cb409b..3bc2644e1 100644 --- a/crates/aisix-proxy/src/videos.rs +++ b/crates/aisix-proxy/src/videos.rs @@ -2011,16 +2011,19 @@ fn emit_submit_usage_event( let pk = crate::usage_attr::ResolvedPk::resolve(snap, provider_key_id); crate::usage_attr::apply_pk_telemetry(&mut event, &pk); crate::usage_attr::apply_jwt_identity(&mut event, client.jwt.as_ref()); - let usage_model = crate::usage_attr::usage_event_model_label(snap, &event.requested_model); - state.usage_sink.try_emit( + let usage_model = + crate::usage_attr::usage_event_model_label(snap, &event.requested_model).into_owned(); + crate::usage_attr::emit_usage( + state, + snap, "videos", - event.clone(), + event, crate::usage_attr::usage_event_labels(&usage_model, &pk), + None, + client.trace.as_ref(), + /* terminal */ true, + /* dispatched */ true, ); - let exporters = crate::usage_attr::live_exporters(state, snap); - state - .otlp_fan_out - .fan_out(&event, None, exporters.iter().map(|e| &e.value)); } #[cfg(test)] diff --git a/tests/e2e/src/cases/audio-duration-cost-basis-e2e.test.ts b/tests/e2e/src/cases/audio-duration-cost-basis-e2e.test.ts index 0a5aadf99..cd33201e6 100644 --- a/tests/e2e/src/cases/audio-duration-cost-basis-e2e.test.ts +++ b/tests/e2e/src/cases/audio-duration-cost-basis-e2e.test.ts @@ -109,7 +109,14 @@ async function waitSpan( ): Promise> { const deadline = Date.now() + timeoutMs; while (Date.now() < deadline) { - const hit = recv.spanAttrs.find((a) => a["aisix.request_id"] === requestId); + // Attempt carriers only: since AISIX-Cloud#1279 the export also ships + // the trace's structural SERVER / logical spans, which share + // `aisix.request_id` but carry no per-attempt usage attributes. + const hit = recv.spanAttrs.find( + (a) => + a["aisix.request_id"] === requestId && + a["aisix.attempt_index"] !== undefined, + ); if (hit) return hit; await new Promise((r) => setTimeout(r, 50)); } diff --git a/tests/e2e/src/cases/claim-mapping-e2e.test.ts b/tests/e2e/src/cases/claim-mapping-e2e.test.ts index 3cfad7ea7..acb021207 100644 --- a/tests/e2e/src/cases/claim-mapping-e2e.test.ts +++ b/tests/e2e/src/cases/claim-mapping-e2e.test.ts @@ -104,7 +104,14 @@ async function waitForSpan( ): Promise> { const deadline = Date.now() + timeoutMs; while (Date.now() < deadline) { - const hit = recv.spans.find((a) => a["aisix.request_id"] === requestId); + // Attempt carriers only: since AISIX-Cloud#1279 the export also ships + // the trace's structural SERVER / logical spans, which share + // `aisix.request_id` but carry no per-attempt usage attributes. + const hit = recv.spans.find( + (a) => + a["aisix.request_id"] === requestId && + a["aisix.attempt_index"] !== undefined, + ); if (hit) return hit; await new Promise((r) => setTimeout(r, 50)); } diff --git a/tests/e2e/src/cases/client-request-id-e2e.test.ts b/tests/e2e/src/cases/client-request-id-e2e.test.ts index 8a528c8f9..af6610ae1 100644 --- a/tests/e2e/src/cases/client-request-id-e2e.test.ts +++ b/tests/e2e/src/cases/client-request-id-e2e.test.ts @@ -99,8 +99,13 @@ async function waitForSpans( ): Promise>> { const deadline = Date.now() + timeoutMs; for (;;) { + // Attempt carriers only: since AISIX-Cloud#1279 the export also ships + // the trace's structural SERVER / logical spans, which share + // `aisix.request_id` but carry no per-attempt usage attributes. const hits = recv.spans.filter( - (a) => a["aisix.request_id"] === requestId, + (a) => + a["aisix.request_id"] === requestId && + a["aisix.attempt_index"] !== undefined, ); if (hits.length >= count) return hits; if (Date.now() >= deadline) { diff --git a/tests/e2e/src/cases/datadog-exporter-e2e.test.ts b/tests/e2e/src/cases/datadog-exporter-e2e.test.ts index bacac95cd..1ad5e9f65 100644 --- a/tests/e2e/src/cases/datadog-exporter-e2e.test.ts +++ b/tests/e2e/src/cases/datadog-exporter-e2e.test.ts @@ -1,4 +1,5 @@ import { createHash } from "node:crypto"; +import { startMockOtlp, type MockOtlp } from "../harness/otlp-mock.js"; import { createServer, type IncomingMessage, type Server } from "node:http"; import { gunzipSync } from "node:zlib"; import { afterAll, beforeAll, describe, expect, test } from "vitest"; @@ -247,6 +248,15 @@ describe("datadog exporter e2e (#688): DP delivers a gzip JSON intake to Datadog // Default privacy posture: operational metadata only, never content. content_mode: "metadata_only", }); + // A second exporter on the OTLP wire, to pin that the public + // trace_id both sinks carry names the SAME trace (AISIX-Cloud#1279). + const otlp: MockOtlp = await startMockOtlp(); + await seed.createObservabilityExporter({ + name: "mock-datadog-otlp-side", + enabled: true, + kind: "otlp_http", + endpoint: otlp.url, + }); await seedRouting(seed, upstream, "datadog-exporter-model"); await waitConfigPropagation(async () => { @@ -298,6 +308,29 @@ describe("datadog exporter e2e (#688): DP delivers a gzip JSON intake to Datadog // The API key must appear ONLY in the header — never anywhere in the body. expect(intake.bodyText.includes(DD_API_KEY)).toBe(false); + + // The public trace_id (AISIX-Cloud#1279): present on the log wire as + // 32 lowercase hex, and EQUAL to the trace the OTLP exporter shipped + // for the same request — the correlation contract the control + // plane's {trace_id} link template depends on. + const logTraceId = log["aisix.trace_id"]; + expect(logTraceId).toMatch(/^[0-9a-f]{32}$/); + const requestId = log["aisix.request_id"]; + const deadline = Date.now() + 10_000; + while ( + Date.now() < deadline && + !otlp.spans.some((s) => s.attributes["aisix.request_id"] === requestId) + ) { + await new Promise((r) => setTimeout(r, 50)); + } + const spans = otlp.spans.filter( + (s) => s.attributes["aisix.request_id"] === requestId, + ); + expect(spans.length).toBeGreaterThan(0); + for (const span of spans) { + expect(span.traceId).toBe(logTraceId); + } + await otlp.close(); // Metadata-only posture: no captured prompt / response fields at all. expect(log["gen_ai.prompt"]).toBeUndefined(); expect(log["gen_ai.completion"]).toBeUndefined(); diff --git a/tests/e2e/src/cases/latency-guardrail-holdback-e2e.test.ts b/tests/e2e/src/cases/latency-guardrail-holdback-e2e.test.ts index 2bf55dee6..d794478a0 100644 --- a/tests/e2e/src/cases/latency-guardrail-holdback-e2e.test.ts +++ b/tests/e2e/src/cases/latency-guardrail-holdback-e2e.test.ts @@ -116,7 +116,14 @@ async function waitForSpan( ): Promise> { const deadline = Date.now() + timeoutMs; while (Date.now() < deadline) { - const hit = recv.spans.find((a) => a["aisix.request_id"] === requestId); + // Attempt carriers only: since AISIX-Cloud#1279 the export also ships + // the trace's structural SERVER / logical spans, which share + // `aisix.request_id` but carry no per-attempt usage attributes. + const hit = recv.spans.find( + (a) => + a["aisix.request_id"] === requestId && + a["aisix.attempt_index"] !== undefined, + ); if (hit) return hit; await new Promise((r) => setTimeout(r, 50)); } diff --git a/tests/e2e/src/cases/latency-upstream-downstream-e2e.test.ts b/tests/e2e/src/cases/latency-upstream-downstream-e2e.test.ts index eb2b35a9f..85f6b710f 100644 --- a/tests/e2e/src/cases/latency-upstream-downstream-e2e.test.ts +++ b/tests/e2e/src/cases/latency-upstream-downstream-e2e.test.ts @@ -98,7 +98,14 @@ async function waitForSpan( ): Promise> { const deadline = Date.now() + timeoutMs; while (Date.now() < deadline) { - const hit = recv.spans.find((a) => a["aisix.request_id"] === requestId); + // Attempt carriers only: since AISIX-Cloud#1279 the export also ships + // the trace's structural SERVER / logical spans, which share + // `aisix.request_id` but carry no per-attempt usage attributes. + const hit = recv.spans.find( + (a) => + a["aisix.request_id"] === requestId && + a["aisix.attempt_index"] !== undefined, + ); if (hit) return hit; await new Promise((r) => setTimeout(r, 50)); } diff --git a/tests/e2e/src/cases/per-attempt-telemetry-e2e.test.ts b/tests/e2e/src/cases/per-attempt-telemetry-e2e.test.ts index e626f5e04..c2e03f511 100644 --- a/tests/e2e/src/cases/per-attempt-telemetry-e2e.test.ts +++ b/tests/e2e/src/cases/per-attempt-telemetry-e2e.test.ts @@ -103,16 +103,23 @@ async function waitForAttempts( timeoutMs = 10_000, ): Promise>> { const deadline = Date.now() + timeoutMs; - while (Date.now() < deadline) { - const hits = recv.spanAttrs.filter( - (a) => a["aisix.request_id"] === requestId, + // Attempt carriers only: since AISIX-Cloud#1279 the export also ships + // the structural SERVER / logical spans of the trace hierarchy, which + // share `aisix.request_id` but carry no per-attempt attributes. + const matching = () => + recv.spanAttrs.filter( + (a) => + a["aisix.request_id"] === requestId && + a["aisix.attempt_index"] !== undefined, ); + while (Date.now() < deadline) { + const hits = matching(); if (hits.length >= count) return hits; await new Promise((r) => setTimeout(r, 50)); } throw new Error( `expected ${count} attempt spans for request_id=${requestId}, ` + - `saw ${recv.spanAttrs.filter((a) => a["aisix.request_id"] === requestId).length}`, + `saw ${matching().length}`, ); } @@ -124,8 +131,13 @@ async function waitForAttemptSpans( timeoutMs = 10_000, ): Promise; latencyMs: number }>> { const deadline = Date.now() + timeoutMs; + // Attempt carriers only — see `waitForAttempts`. const matching = () => - recv.spans.filter((s) => s.attrs["aisix.request_id"] === requestId); + recv.spans.filter( + (s) => + s.attrs["aisix.request_id"] === requestId && + s.attrs["aisix.attempt_index"] !== undefined, + ); while (Date.now() < deadline) { const hits = matching(); if (hits.length >= count) { diff --git a/tests/e2e/src/cases/realtime-ws-e2e.test.ts b/tests/e2e/src/cases/realtime-ws-e2e.test.ts index 9e49aa6ca..866c8693b 100644 --- a/tests/e2e/src/cases/realtime-ws-e2e.test.ts +++ b/tests/e2e/src/cases/realtime-ws-e2e.test.ts @@ -295,8 +295,12 @@ describe("realtime e2e: /v1/realtime WebSocket relay (#721)", () => { let span; const deadline = Date.now() + 10_000; while (Date.now() < deadline) { + // Attempt carrier only — the AISIX-Cloud#1279 hierarchy also ships + // structural spans sharing the request id without usage attributes. span = otlp.spans.find( - (s) => s.attributes["aisix.request_id"] === refusal.requestId, + (s) => + s.attributes["aisix.request_id"] === refusal.requestId && + s.attributes["aisix.attempt_index"] !== undefined, ); if (span) break; await new Promise((r) => setTimeout(r, 100)); diff --git a/tests/e2e/src/cases/trace-hierarchy-e2e.test.ts b/tests/e2e/src/cases/trace-hierarchy-e2e.test.ts new file mode 100644 index 000000000..d82431d73 --- /dev/null +++ b/tests/e2e/src/cases/trace-hierarchy-e2e.test.ts @@ -0,0 +1,877 @@ +import { createHash } from "node:crypto"; +import { afterAll, beforeAll, describe, expect, test } from "vitest"; +import { + EtcdClient, + SeedClient, + spawnApp, + startOpenAiUpstream, + waitConfigPropagation, + type OpenAiUpstream, + type SpawnedApp, +} from "../harness/index.js"; +import { + startMockOtlp, + type CapturedSpan, + type MockOtlp, +} from "../harness/otlp-mock.js"; + +// E2E (AISIX-Cloud#1279): the OTLP export is a real trace, not a bag of +// disconnected spans. One request — including a failover with a failed +// attempt — exports ONE trace: an HTTP SERVER span for the inbound +// request, a logical GenAI CLIENT span covering every upstream attempt, +// and one CLIENT child per attempt, all sharing a trace id, with +// nanosecond timestamps bracketed parent-over-child. +// +// The trust rules ride along: a valid inbound W3C `traceparent` continues +// the caller's trace (the SERVER span parents under it and carries the +// caller's `tracestate`); anything malformed starts a local trace; an +// inbound sampled=1 flag cannot force export past an exporter's +// `sample_rate: 0`; and the caller's trace context never reaches the +// upstream provider — on the standard pipeline even under a +// `forward_client_headers: ["*"]` glob, and on `/passthrough/*` where +// every unlisted header is otherwise forwarded verbatim. +// +// Delivery integrity: a transiently failing receiver sees the SAME span +// ids again on the retry — ids are frozen at emission time, not minted +// per encode. + +const CALLER_PLAINTEXT = "sk-trace-hierarchy-caller"; +const CALLER_KEY_HASH = createHash("sha256") + .update(CALLER_PLAINTEXT) + .digest("hex"); + +/** A well-formed W3C traceparent the tests reuse. */ +const REMOTE_TRACE_ID = "4bf92f3577b34da6a3ce929d0e0e4736"; +const REMOTE_PARENT_ID = "00f067aa0ba902b7"; +const VALID_TRACEPARENT = `00-${REMOTE_TRACE_ID}-${REMOTE_PARENT_ID}-01`; + +/** OTLP SpanKind values (opentelemetry-proto trace.proto). */ +const KIND_SERVER = 2; +const KIND_CLIENT = 3; + +async function waitForSpans( + recv: MockOtlp, + requestId: string, + count: number, + timeoutMs = 10_000, +): Promise { + const deadline = Date.now() + timeoutMs; + const matching = () => + recv.spans.filter((s) => s.attributes["aisix.request_id"] === requestId); + while (Date.now() < deadline) { + const hits = matching(); + if (hits.length >= count) return hits; + await new Promise((r) => setTimeout(r, 50)); + } + throw new Error( + `expected ${count} spans for request_id=${requestId}, saw ${matching().length}`, + ); +} + +const nanos = (value: string): bigint => BigInt(value); + +/** + * Cover the exporter pipelines' independent ~1s flush phases before + * trusting a negative or exact-count assertion: a late batch (or a + * duplicate one, if emission regressed) can land up to a flush interval + * plus a retry after the first spans arrive. + */ +const settle = () => new Promise((r) => setTimeout(r, 2_500)); + +describe("trace hierarchy e2e (AISIX-Cloud#1279)", () => { + let etcdReachable = false; + let app: SpawnedApp | undefined; + let seed: SeedClient | undefined; + const upstreams: OpenAiUpstream[] = []; + const receivers: MockOtlp[] = []; + + beforeAll(async () => { + const etcd = new EtcdClient(); + etcdReachable = await etcd.ping(); + if (!etcdReachable) return; + app = await spawnApp(); + seed = new SeedClient(etcd, app.etcdPrefix); + await seed.createApiKey({ + key_hash: CALLER_KEY_HASH, + allowed_models: ["*"], + // The passthrough leg addresses a route, not a model. + allowed_routes: ["*"], + }); + }); + + afterAll(async () => { + await app?.exit(); + await Promise.all(upstreams.map((u) => u.close())); + await Promise.all(receivers.map((r) => r.close())); + }); + + async function propagate(): Promise { + const canary = `sk-canary-trace-${Date.now()}-${Math.random()}`; + await seed!.createApiKey({ + key_hash: createHash("sha256").update(canary).digest("hex"), + allowed_models: ["*"], + }); + await waitConfigPropagation(async () => { + const res = await fetch(`${app!.proxyUrl}/v1/models`, { + headers: { authorization: `Bearer ${canary}` }, + }); + await res.text(); + return res.status === 200; + }); + } + + async function createOpenAiModel( + displayName: string, + upstream: OpenAiUpstream, + extraPk: Record = {}, + ): Promise { + const providerKey = await seed!.createProviderKey({ + display_name: `${displayName}-pk`, + secret: "sk-mock", + api_base: `${upstream.baseUrl}/v1`, + ...extraPk, + }); + await seed!.createModel({ + display_name: displayName, + provider: "openai", + model_name: "gpt-4o-mini", + provider_key_id: providerKey.id, + }); + } + + function okUpstreamBody(id: string) { + return { + id, + object: "chat.completion", + created: Math.floor(Date.now() / 1000), + model: "gpt-4o-mini", + choices: [ + { + index: 0, + message: { role: "assistant", content: "ok" }, + finish_reason: "stop", + }, + ], + usage: { prompt_tokens: 3, completion_tokens: 4, total_tokens: 7 }, + }; + } + + async function driveChat( + model: string, + headers: Record = {}, + ): Promise { + const res = await fetch(`${app!.proxyUrl}/v1/chat/completions`, { + method: "POST", + headers: { + authorization: `Bearer ${CALLER_PLAINTEXT}`, + "content-type": "application/json", + ...headers, + }, + body: JSON.stringify({ + model, + messages: [{ role: "user", content: "trace me" }], + }), + }); + expect(res.status).toBe(200); + const requestId = res.headers.get("x-aisix-request-id"); + expect(requestId).toBeTruthy(); + await res.text(); + return requestId!; + } + + test("a failover request exports one trace: SERVER → logical CLIENT → attempt children", async (ctx) => { + if (!etcdReachable || !app || !seed) { + ctx.skip(); + return; + } + const otlp = await startMockOtlp(); + receivers.push(otlp); + await seed.createObservabilityExporter({ + name: "trace-hier-otlp", + enabled: true, + kind: "otlp_http", + endpoint: otlp.url, + }); + + const primary = await startOpenAiUpstream({ + status: 502, + errorBody: { error: { message: "primary down", type: "server_error" } }, + }); + const secondary = await startOpenAiUpstream({ + nonStreamBody: okUpstreamBody("cmpl-trace-hier"), + }); + upstreams.push(primary, secondary); + await createOpenAiModel("trace-hier-primary", primary); + await createOpenAiModel("trace-hier-secondary", secondary); + await seed.createModel({ + display_name: "trace-hier-virtual", + routing: { + strategy: "failover", + targets: [ + { model: "trace-hier-primary" }, + { model: "trace-hier-secondary" }, + ], + retries: 0, + max_fallbacks: 1, + }, + }); + await propagate(); + + const requestId = await driveChat("trace-hier-virtual"); + + // 2 attempts + logical + server. + await waitForSpans(otlp, requestId, 4); + await settle(); + const spans = otlp.spans.filter( + (s) => s.attributes["aisix.request_id"] === requestId, + ); + expect(spans).toHaveLength(4); + + // One trace: every span shares one non-zero trace id, all ids distinct. + const traceIds = new Set(spans.map((s) => s.traceId)); + expect(traceIds.size).toBe(1); + expect([...traceIds][0]).toMatch(/^[0-9a-f]{32}$/); + expect([...traceIds][0]).not.toBe("0".repeat(32)); + expect(new Set(spans.map((s) => s.spanId)).size).toBe(4); + + const server = spans.find((s) => s.kind === KIND_SERVER); + expect(server, "exactly one SERVER span").toBeTruthy(); + // No inbound traceparent → the SERVER span roots the trace. + expect(server!.parentSpanId).toBe(""); + + const clients = spans.filter((s) => s.kind === KIND_CLIENT); + expect(clients).toHaveLength(3); + const logical = clients.find((s) => s.parentSpanId === server!.spanId); + expect(logical, "one CLIENT span parented to SERVER").toBeTruthy(); + + const attempts = clients + .filter((s) => s.parentSpanId === logical!.spanId) + .sort( + (a, b) => + Number(a.attributes["aisix.attempt_index"]) - + Number(b.attributes["aisix.attempt_index"]), + ); + expect(attempts).toHaveLength(2); + expect(attempts[0].attributes["aisix.attempt_kind"]).toBe("initial"); + expect(attempts[0].attributes["aisix.error_class"]).toBe("upstream_status"); + expect(attempts[1].attributes["aisix.attempt_kind"]).toBe("fallback"); + expect(attempts[1].attributes["gen_ai.usage.input_tokens"]).toBe(3); + + // Nanosecond bracketing: server ⊇ logical ⊇ each attempt, and the + // failed attempt starts no later than the fallback. + expect(nanos(server!.startTimeUnixNano)).toBeLessThanOrEqual( + nanos(logical!.startTimeUnixNano), + ); + expect(nanos(logical!.endTimeUnixNano)).toBeLessThanOrEqual( + nanos(server!.endTimeUnixNano), + ); + for (const attempt of attempts) { + expect(nanos(logical!.startTimeUnixNano)).toBeLessThanOrEqual( + nanos(attempt.startTimeUnixNano), + ); + expect(nanos(attempt.endTimeUnixNano)).toBeLessThanOrEqual( + nanos(logical!.endTimeUnixNano), + ); + } + expect(nanos(attempts[0].startTimeUnixNano)).toBeLessThanOrEqual( + nanos(attempts[1].startTimeUnixNano), + ); + }); + + test("a valid inbound traceparent parents the SERVER span and continues the caller's trace", async (ctx) => { + if (!etcdReachable || !app || !seed) { + ctx.skip(); + return; + } + const otlp = await startMockOtlp(); + receivers.push(otlp); + await seed.createObservabilityExporter({ + name: "trace-remote-otlp", + enabled: true, + kind: "otlp_http", + endpoint: otlp.url, + }); + const upstream = await startOpenAiUpstream({ + nonStreamBody: okUpstreamBody("cmpl-trace-remote"), + }); + upstreams.push(upstream); + await createOpenAiModel("trace-remote-direct", upstream); + await propagate(); + + const requestId = await driveChat("trace-remote-direct", { + traceparent: VALID_TRACEPARENT, + tracestate: "vendor=x", + }); + + // Direct model, one attempt: server + logical + attempt. + const spans = await waitForSpans(otlp, requestId, 3); + for (const span of spans) { + expect(span.traceId).toBe(REMOTE_TRACE_ID); + } + const server = spans.find((s) => s.kind === KIND_SERVER)!; + expect(server.parentSpanId).toBe(REMOTE_PARENT_ID); + expect(server.traceState).toBe("vendor=x"); + // Span flags: sampled (0x01) | HAS_IS_REMOTE (0x100) | IS_REMOTE (0x200). + expect(server.flags & 0x200).toBe(0x200); + }); + + test("a malformed traceparent starts a local trace and discards tracestate", async (ctx) => { + if (!etcdReachable || !app || !seed) { + ctx.skip(); + return; + } + const otlp = await startMockOtlp(); + receivers.push(otlp); + await seed.createObservabilityExporter({ + name: "trace-badparent-otlp", + enabled: true, + kind: "otlp_http", + endpoint: otlp.url, + }); + const upstream = await startOpenAiUpstream({ + nonStreamBody: okUpstreamBody("cmpl-trace-bad"), + }); + upstreams.push(upstream); + await createOpenAiModel("trace-bad-direct", upstream); + await propagate(); + + for (const bad of [ + // Forbidden version. + `ff-${REMOTE_TRACE_ID}-${REMOTE_PARENT_ID}-01`, + // All-zero trace id. + `00-${"0".repeat(32)}-${REMOTE_PARENT_ID}-01`, + // Uppercase hex. + `00-${REMOTE_TRACE_ID.toUpperCase()}-${REMOTE_PARENT_ID}-01`, + ]) { + const requestId = await driveChat("trace-bad-direct", { + traceparent: bad, + tracestate: "vendor=x", + }); + const spans = await waitForSpans(otlp, requestId, 3); + const server = spans.find((s) => s.kind === KIND_SERVER)!; + expect(server.parentSpanId, `local root for ${bad}`).toBe(""); + expect(spans[0].traceId).not.toBe(REMOTE_TRACE_ID); + expect(server.traceState).toBe(""); + } + }); + + test("an inbound sampled=1 flag cannot override an exporter's sample_rate: 0", async (ctx) => { + if (!etcdReachable || !app || !seed) { + ctx.skip(); + return; + } + const zeroRecv = await startMockOtlp(); + const controlRecv = await startMockOtlp(); + receivers.push(zeroRecv, controlRecv); + await seed.createObservabilityExporter({ + name: "trace-sample-zero", + enabled: true, + kind: "otlp_http", + endpoint: zeroRecv.url, + sample_rate: 0.0, + }); + await seed.createObservabilityExporter({ + name: "trace-sample-control", + enabled: true, + kind: "otlp_http", + endpoint: controlRecv.url, + sample_rate: 1.0, + }); + const upstream = await startOpenAiUpstream({ + nonStreamBody: okUpstreamBody("cmpl-trace-sample"), + }); + upstreams.push(upstream); + await createOpenAiModel("trace-sample-direct", upstream); + await propagate(); + + const requestId = await driveChat("trace-sample-direct", { + traceparent: VALID_TRACEPARENT, // flags 01 = sampled + }); + + // The control exporter proves the export pipeline ran end-to-end AND + // that the inbound traceparent was accepted (otherwise this test says + // nothing about the sampled flag) — the SERVER span continues it. + const control = await waitForSpans(controlRecv, requestId, 3); + const controlServer = control.find((s) => s.kind === KIND_SERVER)!; + expect(controlServer.parentSpanId).toBe(REMOTE_PARENT_ID); + // ...and only after covering the zero exporter's own flush phase is + // its silence a sampling decision rather than slowness: the caller's + // sampled=1 did not force its way in. + await settle(); + expect( + zeroRecv.spans.filter( + (s) => s.attributes["aisix.request_id"] === requestId, + ), + ).toHaveLength(0); + }); + + test("the caller's trace context never reaches the provider — even under a forward_client_headers glob", async (ctx) => { + if (!etcdReachable || !app || !seed) { + ctx.skip(); + return; + } + const upstream = await startOpenAiUpstream({ + nonStreamBody: okUpstreamBody("cmpl-trace-leak"), + }); + upstreams.push(upstream); + await createOpenAiModel("trace-leak-direct", upstream, { + request: { forward_client_headers: ["*"] }, + }); + await propagate(); + + await driveChat("trace-leak-direct", { + traceparent: VALID_TRACEPARENT, + tracestate: "vendor=x", + "x-keep": "forwarded", + }); + + expect(upstream.receivedRequests).toHaveLength(1); + const seen = upstream.receivedRequests[0].headers; + // The glob genuinely forwarded the benign header... + expect(seen["x-keep"]).toBe("forwarded"); + // ...and the trace context stayed behind. + expect(seen["traceparent"]).toBeUndefined(); + expect(seen["tracestate"]).toBeUndefined(); + }); + + test("passthrough strips the caller's trace context unconditionally", async (ctx) => { + if (!etcdReachable || !app || !seed) { + ctx.skip(); + return; + } + const upstream = await startOpenAiUpstream({ + nonStreamBody: okUpstreamBody("cmpl-trace-pt"), + }); + upstreams.push(upstream); + const pk = await seed.createProviderKey({ + display_name: "trace-pt-pk", + secret: "sk-mock", + api_base: `${upstream.baseUrl}/v1`, + }); + await seed.createPassthroughRoute({ + name: "trace-pt-route", + path_prefix: "/passthrough/trace-pt", + target_url: `${upstream.baseUrl}/v1`, + provider_key_id: pk.id, + }); + await propagate(); + + const res = await fetch( + `${app.proxyUrl}/passthrough/trace-pt/v1/chat/completions`, + { + method: "POST", + headers: { + authorization: `Bearer ${CALLER_PLAINTEXT}`, + "content-type": "application/json", + traceparent: VALID_TRACEPARENT, + tracestate: "vendor=x", + "x-custom-app": "kept", + }, + body: JSON.stringify({ + model: "gpt-4o-mini", + messages: [{ role: "user", content: "tunnel me" }], + }), + }, + ); + expect(res.status).toBe(200); + await res.text(); + + expect(upstream.receivedRequests).toHaveLength(1); + const seen = upstream.receivedRequests[0].headers; + // Passthrough forwards unlisted headers verbatim... + expect(seen["x-custom-app"]).toBe("kept"); + // ...which is exactly why the trace context needs the explicit strip. + expect(seen["traceparent"]).toBeUndefined(); + expect(seen["tracestate"]).toBeUndefined(); + }); + + test("a streamed request exports the same three-span hierarchy from the drop-guard emit", async (ctx) => { + if (!etcdReachable || !app || !seed) { + ctx.skip(); + return; + } + const otlp = await startMockOtlp(); + receivers.push(otlp); + await seed.createObservabilityExporter({ + name: "trace-stream-otlp", + enabled: true, + kind: "otlp_http", + endpoint: otlp.url, + }); + const upstream = await startOpenAiUpstream({ + streamEvents: [ + JSON.stringify({ + id: "chatcmpl-trace-stream", + object: "chat.completion.chunk", + created: Math.floor(Date.now() / 1000), + model: "gpt-4o-mini", + choices: [{ index: 0, delta: { content: "hi" }, finish_reason: null }], + }), + JSON.stringify({ + id: "chatcmpl-trace-stream", + object: "chat.completion.chunk", + created: Math.floor(Date.now() / 1000), + model: "gpt-4o-mini", + choices: [{ index: 0, delta: {}, finish_reason: "stop" }], + usage: { prompt_tokens: 3, completion_tokens: 4, total_tokens: 7 }, + }), + "[DONE]", + ], + }); + upstreams.push(upstream); + await createOpenAiModel("trace-stream-direct", upstream); + await propagate(); + + const res = await fetch(`${app.proxyUrl}/v1/chat/completions`, { + method: "POST", + headers: { + authorization: `Bearer ${CALLER_PLAINTEXT}`, + "content-type": "application/json", + }, + body: JSON.stringify({ + model: "trace-stream-direct", + messages: [{ role: "user", content: "stream me" }], + stream: true, + }), + }); + expect(res.status).toBe(200); + const requestId = res.headers.get("x-aisix-request-id"); + expect(requestId).toBeTruthy(); + // Drain — the terminal emit fires from the stream's drop guard. + await res.text(); + + const spans = await waitForSpans(otlp, requestId!, 3); + expect(spans).toHaveLength(3); + const server = spans.find((s) => s.kind === KIND_SERVER)!; + const logical = spans.find( + (s) => s.kind === KIND_CLIENT && s.parentSpanId === server.spanId, + )!; + const attempt = spans.find((s) => s.parentSpanId === logical.spanId)!; + expect(attempt.attributes["gen_ai.usage.output_tokens"]).toBe(4); + expect(nanos(attempt.endTimeUnixNano)).toBeLessThanOrEqual( + nanos(server.endTimeUnixNano), + ); + }); + + test("an all-failed failover exports one SERVER span with the failed attempts under it", async (ctx) => { + if (!etcdReachable || !app || !seed) { + ctx.skip(); + return; + } + const otlp = await startMockOtlp(); + receivers.push(otlp); + await seed.createObservabilityExporter({ + name: "trace-allfail-otlp", + enabled: true, + kind: "otlp_http", + endpoint: otlp.url, + }); + const first = await startOpenAiUpstream({ + status: 502, + errorBody: { error: { message: "down 1", type: "server_error" } }, + }); + const second = await startOpenAiUpstream({ + status: 502, + errorBody: { error: { message: "down 2", type: "server_error" } }, + }); + upstreams.push(first, second); + await createOpenAiModel("trace-allfail-a", first); + await createOpenAiModel("trace-allfail-b", second); + await seed.createModel({ + display_name: "trace-allfail-virtual", + routing: { + strategy: "failover", + targets: [{ model: "trace-allfail-a" }, { model: "trace-allfail-b" }], + retries: 0, + max_fallbacks: 1, + }, + }); + await propagate(); + + const res = await fetch(`${app.proxyUrl}/v1/chat/completions`, { + method: "POST", + headers: { + authorization: `Bearer ${CALLER_PLAINTEXT}`, + "content-type": "application/json", + }, + body: JSON.stringify({ + model: "trace-allfail-virtual", + messages: [{ role: "user", content: "everything is down" }], + }), + }); + expect(res.status).toBeGreaterThanOrEqual(500); + const requestId = res.headers.get("x-aisix-request-id"); + expect(requestId).toBeTruthy(); + await res.text(); + + // Exactly SERVER + logical + two failed attempts: no terminal event + // exists on this shape, so the LAST failed attempt's emission must + // carry the structural spans — and only once. + await waitForSpans(otlp, requestId!, 4); + await settle(); + const spans = otlp.spans.filter( + (s) => s.attributes["aisix.request_id"] === requestId, + ); + expect(spans).toHaveLength(4); + expect(spans.filter((s) => s.kind === KIND_SERVER)).toHaveLength(1); + const server = spans.find((s) => s.kind === KIND_SERVER)!; + const logical = spans.find( + (s) => s.kind === KIND_CLIENT && s.parentSpanId === server.spanId, + )!; + const attempts = spans.filter((s) => s.parentSpanId === logical.spanId); + expect(attempts).toHaveLength(2); + for (const attempt of attempts) { + expect(attempt.attributes["aisix.error_class"]).toBe("upstream_status"); + } + }); + + test("/v1/messages and /v1/responses export the same hierarchy shape", async (ctx) => { + if (!etcdReachable || !app || !seed) { + ctx.skip(); + return; + } + const otlp = await startMockOtlp(); + receivers.push(otlp); + await seed.createObservabilityExporter({ + name: "trace-family-otlp", + enabled: true, + kind: "otlp_http", + endpoint: otlp.url, + }); + const upstream = await startOpenAiUpstream({ + nonStreamBody: okUpstreamBody("cmpl-trace-family"), + }); + upstreams.push(upstream); + await createOpenAiModel("trace-family-direct", upstream); + await propagate(); + + const drive = async (path: string, body: unknown) => { + const res = await fetch(`${app!.proxyUrl}${path}`, { + method: "POST", + headers: { + authorization: `Bearer ${CALLER_PLAINTEXT}`, + "content-type": "application/json", + }, + body: JSON.stringify(body), + }); + expect(res.status, path).toBe(200); + const requestId = res.headers.get("x-aisix-request-id"); + expect(requestId, path).toBeTruthy(); + await res.text(); + return requestId!; + }; + + for (const [path, body] of [ + [ + "/v1/messages", + { + model: "trace-family-direct", + max_tokens: 32, + messages: [{ role: "user", content: "trace me" }], + }, + ], + ["/v1/responses", { model: "trace-family-direct", input: "trace me" }], + ] as const) { + const requestId = await drive(path, body); + const spans = await waitForSpans(otlp, requestId, 3); + expect(spans, path).toHaveLength(3); + const server = spans.find((s) => s.kind === KIND_SERVER)!; + expect(server, path).toBeTruthy(); + const logical = spans.find( + (s) => s.kind === KIND_CLIENT && s.parentSpanId === server.spanId, + )!; + expect(logical, path).toBeTruthy(); + const attempt = spans.find((s) => s.parentSpanId === logical.spanId)!; + expect(attempt, path).toBeTruthy(); + expect(new Set(spans.map((s) => s.traceId)).size, path).toBe(1); + } + }); + + test("an ensemble request exports one SERVER span and no dangling parents", async (ctx) => { + if (!etcdReachable || !app || !seed) { + ctx.skip(); + return; + } + const otlp = await startMockOtlp(); + receivers.push(otlp); + await seed.createObservabilityExporter({ + name: "trace-ens-otlp", + enabled: true, + kind: "otlp_http", + endpoint: otlp.url, + }); + const member = await startOpenAiUpstream({ + nonStreamBody: okUpstreamBody("cmpl-trace-ens-member"), + }); + const judge = await startOpenAiUpstream({ + nonStreamBody: okUpstreamBody("cmpl-trace-ens-judge"), + }); + upstreams.push(member, judge); + await createOpenAiModel("trace-ens-member-a", member); + await createOpenAiModel("trace-ens-member-b", member); + await createOpenAiModel("trace-ens-judge", judge); + await seed.createModel({ + display_name: "trace-ens-virtual", + ensemble: { + panel: [{ model: "trace-ens-member-a" }, { model: "trace-ens-member-b" }], + judge: { model: "trace-ens-judge" }, + min_responses: 2, + }, + }); + await propagate(); + + const requestId = await driveChat("trace-ens-virtual"); + + // Two panel sub-calls + the judge's terminal emission (SERVER + + // logical carrier) = four spans. + await waitForSpans(otlp, requestId, 4); + await settle(); + const spans = otlp.spans.filter( + (s) => s.attributes["aisix.request_id"] === requestId, + ); + expect(spans).toHaveLength(4); + expect(spans.filter((s) => s.kind === KIND_SERVER)).toHaveLength(1); + // One trace: parent resolution below is meaningful only within it. + expect(new Set(spans.map((s) => s.traceId)).size).toBe(1); + + // Every parent resolves within the exported trace — the orphan-span + // regression this test exists to prevent. + const ids = new Set(spans.map((s) => s.spanId)); + for (const span of spans) { + if (span.parentSpanId !== "") { + expect( + ids.has(span.parentSpanId), + `dangling parent ${span.parentSpanId} on ${span.name}`, + ).toBe(true); + } + } + const panel = spans.filter( + (s) => s.attributes["aisix.attempt_kind"] === "panel", + ); + expect(panel).toHaveLength(2); + }); + + test("a cache hit exports the SERVER span alone — no fictitious upstream span", async (ctx) => { + if (!etcdReachable || !app || !seed) { + ctx.skip(); + return; + } + const otlp = await startMockOtlp(); + receivers.push(otlp); + await seed.createObservabilityExporter({ + name: "trace-cache-otlp", + enabled: true, + kind: "otlp_http", + endpoint: otlp.url, + }); + const upstream = await startOpenAiUpstream({ + nonStreamBody: okUpstreamBody("cmpl-trace-cache"), + }); + upstreams.push(upstream); + await createOpenAiModel("trace-cache-direct", upstream); + await seed.createCachePolicy({ + name: "trace-cache-policy", + enabled: true, + applies_to: "all", + }); + await propagate(); + + const driveOnce = async () => { + const res = await fetch(`${app!.proxyUrl}/v1/chat/completions`, { + method: "POST", + headers: { + authorization: `Bearer ${CALLER_PLAINTEXT}`, + "content-type": "application/json", + }, + body: JSON.stringify({ + model: "trace-cache-direct", + messages: [{ role: "user", content: "cache me exactly" }], + }), + }); + expect(res.status).toBe(200); + const outcome = res.headers.get("x-aisix-cache"); + const requestId = res.headers.get("x-aisix-request-id"); + await res.text(); + return { outcome, requestId: requestId! }; + }; + + // First call misses and populates; poll until a call reports a hit + // (the write is asynchronous). + await driveOnce(); + let hit: { outcome: string | null; requestId: string } | undefined; + const deadline = Date.now() + 10_000; + while (Date.now() < deadline) { + const attempt = await driveOnce(); + if (attempt.outcome === "hit") { + hit = attempt; + break; + } + await new Promise((r) => setTimeout(r, 100)); + } + expect(hit, "no cache hit within 10s").toBeTruthy(); + + await waitForSpans(otlp, hit!.requestId, 1); + await settle(); + const spans = otlp.spans.filter( + (s) => s.attributes["aisix.request_id"] === hit!.requestId, + ); + expect(spans).toHaveLength(1); + expect(spans[0].kind).toBe(KIND_SERVER); + expect(spans[0].parentSpanId).toBe(""); + }); + + test("a transient receiver failure re-delivers byte-identical span ids", async (ctx) => { + if (!etcdReachable || !app || !seed) { + ctx.skip(); + return; + } + // First POST → 503 (spans still recorded); retry → 200. + const otlp = await startMockOtlp({ failFirst: 1 }); + receivers.push(otlp); + await seed.createObservabilityExporter({ + name: "trace-retry-otlp", + enabled: true, + kind: "otlp_http", + endpoint: otlp.url, + }); + const upstream = await startOpenAiUpstream({ + nonStreamBody: okUpstreamBody("cmpl-trace-retry"), + }); + upstreams.push(upstream); + await createOpenAiModel("trace-retry-direct", upstream); + await propagate(); + + const requestId = await driveChat("trace-retry-direct"); + + // Wait for the retried delivery: the same spans must arrive twice. + const deadline = Date.now() + 15_000; + const byPost = () => { + const mine = otlp.spans.filter( + (s) => s.attributes["aisix.request_id"] === requestId, + ); + const posts = new Map(); + for (const s of mine) { + posts.set(s.postIndex, [...(posts.get(s.postIndex) ?? []), s]); + } + return posts; + }; + while (Date.now() < deadline && byPost().size < 2) { + await new Promise((r) => setTimeout(r, 100)); + } + const posts = [...byPost().values()]; + expect(posts.length).toBeGreaterThanOrEqual(2); + + const idSet = (spans: CapturedSpan[]) => + spans + .map( + (s) => + `${s.traceId}/${s.spanId}/${s.parentSpanId}/${s.kind}/` + + `${s.startTimeUnixNano}/${s.endTimeUnixNano}`, + ) + .sort() + .join("|"); + // The 503'd delivery and its retry carry byte-identical ids and + // timestamps — the regenerate-per-encode defect is gone. + expect(idSet(posts[1])).toBe(idSet(posts[0])); + }); +}); diff --git a/tests/e2e/src/cases/ttft-first-frame-e2e.test.ts b/tests/e2e/src/cases/ttft-first-frame-e2e.test.ts index 395e7a792..3b05d8dc9 100644 --- a/tests/e2e/src/cases/ttft-first-frame-e2e.test.ts +++ b/tests/e2e/src/cases/ttft-first-frame-e2e.test.ts @@ -104,7 +104,14 @@ async function waitForSpan( ): Promise> { const deadline = Date.now() + timeoutMs; while (Date.now() < deadline) { - const hit = recv.spans.find((a) => a["aisix.request_id"] === requestId); + // Attempt carriers only: since AISIX-Cloud#1279 the export also ships + // the trace's structural SERVER / logical spans, which share + // `aisix.request_id` but carry no per-attempt usage attributes. + const hit = recv.spans.find( + (a) => + a["aisix.request_id"] === requestId && + a["aisix.attempt_index"] !== undefined, + ); if (hit) return hit; await new Promise((r) => setTimeout(r, 50)); } diff --git a/tests/e2e/src/harness/otlp-mock.ts b/tests/e2e/src/harness/otlp-mock.ts index e12a913df..0fd4231fb 100644 --- a/tests/e2e/src/harness/otlp-mock.ts +++ b/tests/e2e/src/harness/otlp-mock.ts @@ -11,11 +11,39 @@ import { createServer, type Server } from "node:http"; export interface CapturedSpan { name: string; attributes: Record; + /** 32-hex W3C trace id; empty when the body omitted it. */ + traceId: string; + /** 16-hex span id; empty when the body omitted it. */ + spanId: string; + /** 16-hex parent span id; empty for a root span. */ + parentSpanId: string; + /** OTLP SpanKind (2 = SERVER, 3 = CLIENT); 0 when omitted. */ + kind: number; + /** OTLP span flags (uint32); 0 when omitted. */ + flags: number; + /** W3C tracestate carried on the span; empty when absent. */ + traceState: string; + /** Nanosecond boundaries, as strings (OTLP/JSON int64 encoding). */ + startTimeUnixNano: string; + endTimeUnixNano: string; + /** 0-based index of the POST this span arrived in (delivery-retry tests). */ + postIndex: number; +} + +export interface MockOtlpOptions { + /** + * Respond to the first N POSTs with 503 (their spans are still recorded), + * then accept. Models a transiently failing receiver so a test can assert + * the sink's delivery retry re-sends byte-identical spans. + */ + failFirst?: number; } export interface MockOtlp { url: string; spans: CapturedSpan[]; + /** Number of POSTs received so far (including ones answered 503). */ + posts: number; /** * Bodies this receiver could not parse, empty on a healthy run. Without it a * malformed export and a never-sent one look identical: both end as "no @@ -40,19 +68,32 @@ function anyValue(value: Record): string | number | boolean { return ""; } -export async function startMockOtlp(): Promise { +export async function startMockOtlp( + options: MockOtlpOptions = {}, +): Promise { const spans: CapturedSpan[] = []; const parseFailures: string[] = []; + const state = { posts: 0, failuresLeft: options.failFirst ?? 0 }; const server: Server = createServer((req, res) => { const chunks: Buffer[] = []; req.on("data", (chunk: Buffer) => chunks.push(chunk)); req.on("end", () => { + const postIndex = state.posts; + state.posts += 1; try { const body = JSON.parse(Buffer.concat(chunks).toString("utf8")) as { resourceSpans?: { scopeSpans?: { spans?: { name: string; + traceId?: string; + spanId?: string; + parentSpanId?: string; + kind?: number; + flags?: number; + traceState?: string; + startTimeUnixNano?: string; + endTimeUnixNano?: string; attributes?: { key: string; value: Record }[]; }[]; }[]; @@ -65,7 +106,19 @@ export async function startMockOtlp(): Promise { for (const attr of span.attributes ?? []) { attributes[attr.key] = anyValue(attr.value); } - spans.push({ name: span.name, attributes }); + spans.push({ + name: span.name, + attributes, + traceId: span.traceId ?? "", + spanId: span.spanId ?? "", + parentSpanId: span.parentSpanId ?? "", + kind: span.kind ?? 0, + flags: span.flags ?? 0, + traceState: span.traceState ?? "", + startTimeUnixNano: span.startTimeUnixNano ?? "0", + endTimeUnixNano: span.endTimeUnixNano ?? "0", + postIndex, + }); } } } @@ -75,6 +128,12 @@ export async function startMockOtlp(): Promise { // test as an indistinguishable "no span arrived". parseFailures.push(String(err)); } + if (state.failuresLeft > 0) { + state.failuresLeft -= 1; + res.writeHead(503, { "content-type": "application/json" }); + res.end("{}"); + return; + } res.writeHead(200, { "content-type": "application/json" }); res.end("{}"); }); @@ -87,6 +146,9 @@ export async function startMockOtlp(): Promise { return { url: `http://127.0.0.1:${address.port}/v1/traces`, spans, + get posts() { + return state.posts; + }, parseFailures, close: () => new Promise((resolve) => server.close(() => resolve())), };