use std::fmt;
use serde::{Deserialize, Serialize};
use crate::limits::{
MAX_SPAN_ID_LEN, MAX_TRACE_ID_LEN, WireValidate, WireValidationError, check_opt_string_len,
check_string_len,
};
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[cfg_attr(feature = "schema", derive(schemars::JsonSchema))]
pub struct TraceContext {
pub trace_id: String,
#[serde(skip_serializing_if = "Option::is_none", default)]
pub span_id: Option<String>,
}
impl TraceContext {
pub fn generate() -> Self {
Self {
trace_id: uuid::Uuid::new_v4().simple().to_string(),
span_id: None,
}
}
}
impl Default for TraceContext {
fn default() -> Self {
Self::generate()
}
}
impl fmt::Display for TraceContext {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
match &self.span_id {
Some(span) => write!(f, "{}:{}", self.trace_id, span),
None => f.write_str(&self.trace_id),
}
}
}
impl WireValidate for TraceContext {
fn wire_validate(&self) -> Result<(), WireValidationError> {
check_string_len(&self.trace_id, MAX_TRACE_ID_LEN, "trace_context.trace_id")?;
check_opt_string_len(&self.span_id, MAX_SPAN_ID_LEN, "trace_context.span_id")?;
Ok(())
}
}
pub fn current_trace_context() -> TraceContext {
TraceContext::generate()
}
#[cfg(test)]
mod tests {
#![expect(
clippy::assertions_on_result_states,
reason = "test assertions — is_ok/is_err provides readable failure messages"
)]
use super::*;
#[test]
fn generate_produces_valid_trace_id() {
let ctx = TraceContext::generate();
assert_eq!(ctx.trace_id.len(), 32, "trace_id must be 32 hex chars");
assert!(
ctx.trace_id.chars().all(|c| c.is_ascii_hexdigit()),
"trace_id must contain only hex characters"
);
assert!(ctx.span_id.is_none(), "generated context has no span_id");
}
#[test]
fn generate_produces_unique_ids() {
let a = TraceContext::generate();
let b = TraceContext::generate();
assert_ne!(a.trace_id, b.trace_id, "two generated IDs must differ");
}
#[test]
fn display_without_span_id() {
let ctx = TraceContext {
trace_id: "0123456789abcdef0123456789abcdef".to_string(),
span_id: None,
};
assert_eq!(ctx.to_string(), "0123456789abcdef0123456789abcdef");
}
#[test]
fn display_with_span_id() {
let ctx = TraceContext {
trace_id: "0123456789abcdef0123456789abcdef".to_string(),
span_id: Some("fedcba9876543210".to_string()),
};
assert_eq!(
ctx.to_string(),
"0123456789abcdef0123456789abcdef:fedcba9876543210"
);
}
#[test]
fn serde_roundtrip_with_span_id() {
let ctx = TraceContext {
trace_id: "0123456789abcdef0123456789abcdef".to_string(),
span_id: Some("fedcba9876543210".to_string()),
};
let json = serde_json::to_string(&ctx).unwrap();
assert!(json.contains("span_id"));
let deserialized: TraceContext = serde_json::from_str(&json).unwrap();
assert_eq!(deserialized, ctx);
}
#[test]
fn serde_roundtrip_without_span_id() {
let ctx = TraceContext {
trace_id: "0123456789abcdef0123456789abcdef".to_string(),
span_id: None,
};
let json = serde_json::to_string(&ctx).unwrap();
assert!(!json.contains("span_id"), "None span_id must be omitted");
let deserialized: TraceContext = serde_json::from_str(&json).unwrap();
assert_eq!(deserialized, ctx);
}
#[test]
fn deserialize_missing_span_id() {
let json = r#"{"trace_id":"0123456789abcdef0123456789abcdef"}"#;
let ctx: TraceContext = serde_json::from_str(json).unwrap();
assert_eq!(ctx.trace_id, "0123456789abcdef0123456789abcdef");
assert!(ctx.span_id.is_none());
}
#[test]
fn wire_validate_valid() {
let ctx = TraceContext {
trace_id: "0123456789abcdef0123456789abcdef".to_string(),
span_id: Some("fedcba9876543210".to_string()),
};
assert!(ctx.wire_validate().is_ok());
}
#[test]
fn wire_validate_trace_id_too_long() {
let ctx = TraceContext {
trace_id: "a".repeat(33),
span_id: None,
};
let err = ctx.wire_validate().unwrap_err();
assert_eq!(err.field, "trace_context.trace_id");
}
#[test]
fn wire_validate_span_id_too_long() {
let ctx = TraceContext {
trace_id: "0123456789abcdef0123456789abcdef".to_string(),
span_id: Some("a".repeat(17)),
};
let err = ctx.wire_validate().unwrap_err();
assert_eq!(err.field, "trace_context.span_id");
}
#[test]
fn current_trace_context_generates_valid() {
let ctx = current_trace_context();
assert!(ctx.wire_validate().is_ok());
}
}