ironflow_core/
trace_context.rs

1//! W3C Trace Context propagation for workflow distributed tracing.
2//!
3//! Provides [`WorkflowTraceContext`] to generate and parse W3C `traceparent`
4//! headers, enabling correlation between Ironflow workflow spans and
5//! downstream service spans (LLM providers, MCP servers, etc.).
6//!
7//! The `traceparent` header follows the
8//! [W3C Trace Context](https://www.w3.org/TR/trace-context/) format:
9//!
10//! ```text
11//! 00-{trace_id}-{span_id}-{trace_flags}
12//!  |     |          |          |
13//!  |     |          |          +-- 2 hex (01 = sampled)
14//!  |     |          +-- 16 hex (8 bytes)
15//!  |     +-- 32 hex (16 bytes)
16//!  +-- version (always "00")
17//! ```
18//!
19//! # Examples
20//!
21//! ```
22//! use ironflow_core::trace_context::WorkflowTraceContext;
23//!
24//! // Create a root context and emit the traceparent header.
25//! let root = WorkflowTraceContext::new_root();
26//! let header = root.to_traceparent();
27//! assert!(header.starts_with("00-"));
28//!
29//! // Derive a child span (preserves trace_id, new span_id).
30//! let child = root.child();
31//! assert_eq!(child.trace_id(), root.trace_id());
32//! assert_ne!(child.span_id(), root.span_id());
33//!
34//! // Parse an incoming traceparent header.
35//! let parsed = WorkflowTraceContext::from_traceparent(&header).unwrap();
36//! assert_eq!(parsed.trace_id(), root.trace_id());
37//! ```
38
39use std::fmt;
40use std::process;
41use std::sync::atomic::{AtomicU64, Ordering};
42use std::time::{SystemTime, UNIX_EPOCH};
43
44use serde::{Deserialize, Serialize};
45use sha2::{Digest, Sha256};
46use thiserror::Error;
47
48static SPAN_COUNTER: AtomicU64 = AtomicU64::new(0);
49
50/// Errors that can occur when parsing a `traceparent` header.
51///
52/// # Examples
53///
54/// ```
55/// use ironflow_core::trace_context::{WorkflowTraceContext, TraceContextError};
56///
57/// let err = WorkflowTraceContext::from_traceparent("bad").unwrap_err();
58/// assert!(matches!(err, TraceContextError::InvalidFormat { .. }));
59/// ```
60#[derive(Debug, Error, PartialEq, Eq)]
61pub enum TraceContextError {
62    /// The header does not have the expected 4-field format.
63    #[error("invalid traceparent format: expected 4 dash-separated fields, got {field_count}")]
64    InvalidFormat {
65        /// Number of fields found.
66        field_count: usize,
67    },
68
69    /// The trace-id field is not valid lowercase hex of the expected length.
70    #[error("invalid trace-id: expected 32 lowercase hex chars, got \"{value}\"")]
71    InvalidTraceId {
72        /// The raw value found.
73        value: String,
74    },
75
76    /// The span-id (parent-id) field is not valid lowercase hex of the expected length.
77    #[error("invalid span-id: expected 16 lowercase hex chars, got \"{value}\"")]
78    InvalidSpanId {
79        /// The raw value found.
80        value: String,
81    },
82
83    /// The trace-id is all zeros, which is invalid per the W3C spec.
84    #[error("trace-id must not be all zeros")]
85    ZeroTraceId,
86
87    /// The span-id is all zeros, which is invalid per the W3C spec.
88    #[error("span-id must not be all zeros")]
89    ZeroSpanId,
90}
91
92/// W3C Trace Context for distributed tracing across workflow steps.
93///
94/// Each `WorkflowTraceContext` carries a `trace_id` (32 lowercase hex chars)
95/// and a `span_id` (16 lowercase hex chars). Use [`new_root`](Self::new_root)
96/// to start a new trace, [`child`](Self::child) to create a child span, and
97/// [`to_traceparent`](Self::to_traceparent) to emit the W3C header.
98///
99/// # Examples
100///
101/// ```
102/// use ironflow_core::trace_context::WorkflowTraceContext;
103///
104/// let ctx = WorkflowTraceContext::new_root();
105/// assert_eq!(ctx.trace_id().len(), 32);
106/// assert_eq!(ctx.span_id().len(), 16);
107/// assert_eq!(ctx.to_traceparent().len(), 55); // 2 + 1 + 32 + 1 + 16 + 1 + 2
108/// ```
109#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
110pub struct WorkflowTraceContext {
111    trace_id: String,
112    span_id: String,
113}
114
115impl WorkflowTraceContext {
116    /// Create a new root trace context with a random trace-id and span-id.
117    ///
118    /// The trace-id is derived from the current timestamp, process ID, and an
119    /// atomic counter to ensure uniqueness without requiring a CSPRNG.
120    ///
121    /// # Examples
122    ///
123    /// ```
124    /// use ironflow_core::trace_context::WorkflowTraceContext;
125    ///
126    /// let ctx = WorkflowTraceContext::new_root();
127    /// assert_eq!(ctx.trace_id().len(), 32);
128    /// assert_eq!(ctx.span_id().len(), 16);
129    /// ```
130    pub fn new_root() -> Self {
131        let trace_id = generate_trace_id();
132        let span_id = generate_span_id();
133        Self { trace_id, span_id }
134    }
135
136    /// Create a trace context derived from a workflow run ID.
137    ///
138    /// The `run_id` is hashed (SHA-256) to produce a deterministic 32-hex
139    /// trace-id. This allows correlating all spans of a given workflow run
140    /// under a single trace. A fresh span-id is generated for this context.
141    ///
142    /// # Examples
143    ///
144    /// ```
145    /// use ironflow_core::trace_context::WorkflowTraceContext;
146    ///
147    /// let ctx = WorkflowTraceContext::from_workflow_run_id("run-abc-123");
148    /// assert_eq!(ctx.trace_id().len(), 32);
149    ///
150    /// // Same run_id always produces the same trace_id.
151    /// let ctx2 = WorkflowTraceContext::from_workflow_run_id("run-abc-123");
152    /// assert_eq!(ctx.trace_id(), ctx2.trace_id());
153    /// ```
154    ///
155    /// # Panics
156    ///
157    /// Panics if `run_id` is empty.
158    pub fn from_workflow_run_id(run_id: &str) -> Self {
159        assert!(!run_id.is_empty(), "run_id must not be empty");
160        let mut hasher = Sha256::new();
161        hasher.update(run_id.as_bytes());
162        let hash = hasher.finalize();
163        let trace_id = hex_encode(&hash[..16]);
164        let span_id = generate_span_id();
165        Self { trace_id, span_id }
166    }
167
168    /// Create a child context that shares this trace-id but has a new span-id.
169    ///
170    /// Use this when a workflow step fans out to sub-steps: each child gets
171    /// its own span-id while remaining part of the same trace.
172    ///
173    /// # Examples
174    ///
175    /// ```
176    /// use ironflow_core::trace_context::WorkflowTraceContext;
177    ///
178    /// let parent = WorkflowTraceContext::new_root();
179    /// let child = parent.child();
180    /// assert_eq!(child.trace_id(), parent.trace_id());
181    /// assert_ne!(child.span_id(), parent.span_id());
182    /// ```
183    pub fn child(&self) -> Self {
184        Self {
185            trace_id: self.trace_id.clone(),
186            span_id: generate_span_id(),
187        }
188    }
189
190    /// Format as a W3C `traceparent` header value.
191    ///
192    /// The output follows the format `00-{trace_id}-{span_id}-01`, where
193    /// `01` indicates the trace is sampled.
194    ///
195    /// # Examples
196    ///
197    /// ```
198    /// use ironflow_core::trace_context::WorkflowTraceContext;
199    ///
200    /// let ctx = WorkflowTraceContext::from_traceparent(
201    ///     "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01"
202    /// ).unwrap();
203    /// assert_eq!(
204    ///     ctx.to_traceparent(),
205    ///     "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01"
206    /// );
207    /// ```
208    pub fn to_traceparent(&self) -> String {
209        format!("00-{}-{}-01", self.trace_id, self.span_id)
210    }
211
212    /// Parse a W3C `traceparent` header into a `WorkflowTraceContext`.
213    ///
214    /// Accepts any version field (not just `"00"`), but always emits
215    /// version `"00"` when calling [`to_traceparent`](Self::to_traceparent).
216    /// Extra fields beyond the 4th are silently ignored.
217    ///
218    /// # Errors
219    ///
220    /// Returns [`TraceContextError`] if the header is malformed:
221    /// - Fewer than 4 dash-separated fields
222    /// - trace-id is not exactly 32 lowercase hex characters
223    /// - span-id is not exactly 16 lowercase hex characters
224    /// - trace-id or span-id is all zeros
225    ///
226    /// # Examples
227    ///
228    /// ```
229    /// use ironflow_core::trace_context::WorkflowTraceContext;
230    ///
231    /// let ctx = WorkflowTraceContext::from_traceparent(
232    ///     "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01"
233    /// ).unwrap();
234    /// assert_eq!(ctx.trace_id(), "4bf92f3577b34da6a3ce929d0e0e4736");
235    /// assert_eq!(ctx.span_id(), "00f067aa0ba902b7");
236    /// ```
237    pub fn from_traceparent(header: &str) -> Result<Self, TraceContextError> {
238        let parts: Vec<&str> = header.split('-').collect();
239        if parts.len() < 4 {
240            return Err(TraceContextError::InvalidFormat {
241                field_count: parts.len(),
242            });
243        }
244
245        let trace_id = parts[1];
246        let span_id = parts[2];
247
248        if trace_id.len() != 32 || !trace_id.chars().all(|c| c.is_ascii_hexdigit()) {
249            return Err(TraceContextError::InvalidTraceId {
250                value: trace_id.to_string(),
251            });
252        }
253
254        if trace_id.chars().all(|c| c == '0') {
255            return Err(TraceContextError::ZeroTraceId);
256        }
257
258        if span_id.len() != 16 || !span_id.chars().all(|c| c.is_ascii_hexdigit()) {
259            return Err(TraceContextError::InvalidSpanId {
260                value: span_id.to_string(),
261            });
262        }
263
264        if span_id.chars().all(|c| c == '0') {
265            return Err(TraceContextError::ZeroSpanId);
266        }
267
268        Ok(Self {
269            trace_id: trace_id.to_lowercase(),
270            span_id: span_id.to_lowercase(),
271        })
272    }
273
274    /// Return the trace-id (32 lowercase hex characters).
275    pub fn trace_id(&self) -> &str {
276        &self.trace_id
277    }
278
279    /// Return the span-id (16 lowercase hex characters).
280    pub fn span_id(&self) -> &str {
281        &self.span_id
282    }
283}
284
285impl fmt::Display for WorkflowTraceContext {
286    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
287        write!(f, "{}", self.to_traceparent())
288    }
289}
290
291fn generate_trace_id() -> String {
292    let nanos = SystemTime::now()
293        .duration_since(UNIX_EPOCH)
294        .unwrap_or_default()
295        .as_nanos();
296    let counter = SPAN_COUNTER.fetch_add(1, Ordering::Relaxed);
297    let pid = process::id();
298
299    let mut hasher = Sha256::new();
300    hasher.update(nanos.to_le_bytes());
301    hasher.update(counter.to_le_bytes());
302    hasher.update(pid.to_le_bytes());
303    hasher.update(b"trace");
304    let hash = hasher.finalize();
305    hex_encode(&hash[..16])
306}
307
308fn generate_span_id() -> String {
309    let nanos = SystemTime::now()
310        .duration_since(UNIX_EPOCH)
311        .unwrap_or_default()
312        .as_nanos();
313    let counter = SPAN_COUNTER.fetch_add(1, Ordering::Relaxed);
314    let pid = process::id();
315
316    let mut hasher = Sha256::new();
317    hasher.update(nanos.to_le_bytes());
318    hasher.update(counter.to_le_bytes());
319    hasher.update(pid.to_le_bytes());
320    hasher.update(b"span");
321    let hash = hasher.finalize();
322    hex_encode(&hash[..8])
323}
324
325fn hex_encode(bytes: &[u8]) -> String {
326    bytes.iter().map(|b| format!("{b:02x}")).collect()
327}
328
329#[cfg(test)]
330mod tests {
331    use super::*;
332
333    #[test]
334    fn new_root_produces_valid_ids() {
335        let ctx = WorkflowTraceContext::new_root();
336        assert_eq!(ctx.trace_id().len(), 32, "trace_id must be 32 hex chars");
337        assert_eq!(ctx.span_id().len(), 16, "span_id must be 16 hex chars");
338        assert!(
339            ctx.trace_id().chars().all(|c| c.is_ascii_hexdigit()),
340            "trace_id must be valid hex"
341        );
342        assert!(
343            ctx.span_id().chars().all(|c| c.is_ascii_hexdigit()),
344            "span_id must be valid hex"
345        );
346    }
347
348    #[test]
349    fn new_root_produces_unique_contexts() {
350        let a = WorkflowTraceContext::new_root();
351        let b = WorkflowTraceContext::new_root();
352        assert_ne!(a.trace_id(), b.trace_id());
353    }
354
355    #[test]
356    fn from_workflow_run_id_deterministic() {
357        let a = WorkflowTraceContext::from_workflow_run_id("run-abc-123");
358        let b = WorkflowTraceContext::from_workflow_run_id("run-abc-123");
359        assert_eq!(
360            a.trace_id(),
361            b.trace_id(),
362            "same run_id must produce same trace_id"
363        );
364        assert_eq!(a.trace_id().len(), 32);
365        assert!(a.trace_id().chars().all(|c| c.is_ascii_hexdigit()));
366    }
367
368    #[test]
369    fn from_workflow_run_id_different_inputs() {
370        let a = WorkflowTraceContext::from_workflow_run_id("run-1");
371        let b = WorkflowTraceContext::from_workflow_run_id("run-2");
372        assert_ne!(a.trace_id(), b.trace_id());
373    }
374
375    #[test]
376    #[should_panic(expected = "run_id must not be empty")]
377    fn from_workflow_run_id_empty_panics() {
378        WorkflowTraceContext::from_workflow_run_id("");
379    }
380
381    #[test]
382    fn child_preserves_trace_id() {
383        let parent = WorkflowTraceContext::new_root();
384        let child = parent.child();
385        assert_eq!(child.trace_id(), parent.trace_id());
386        assert_ne!(child.span_id(), parent.span_id());
387        assert_eq!(child.span_id().len(), 16);
388    }
389
390    #[test]
391    fn child_children_are_unique() {
392        let parent = WorkflowTraceContext::new_root();
393        let c1 = parent.child();
394        let c2 = parent.child();
395        assert_ne!(c1.span_id(), c2.span_id());
396        assert_eq!(c1.trace_id(), c2.trace_id());
397    }
398
399    #[test]
400    fn to_traceparent_format() {
401        let ctx = WorkflowTraceContext::new_root();
402        let header = ctx.to_traceparent();
403
404        assert!(header.starts_with("00-"), "must start with version 00");
405        assert!(header.ends_with("-01"), "must end with trace-flags 01");
406        assert_eq!(header.len(), 55, "00-<32>-<16>-01 = 55 chars");
407
408        let parts: Vec<&str> = header.split('-').collect();
409        assert_eq!(parts.len(), 4);
410        assert_eq!(parts[0], "00");
411        assert_eq!(parts[1], ctx.trace_id());
412        assert_eq!(parts[2], ctx.span_id());
413        assert_eq!(parts[3], "01");
414    }
415
416    #[test]
417    fn display_matches_to_traceparent() {
418        let ctx = WorkflowTraceContext::new_root();
419        assert_eq!(format!("{ctx}"), ctx.to_traceparent());
420    }
421
422    #[test]
423    fn from_traceparent_valid() {
424        let ctx = WorkflowTraceContext::from_traceparent(
425            "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01",
426        )
427        .unwrap();
428        assert_eq!(ctx.trace_id(), "4bf92f3577b34da6a3ce929d0e0e4736");
429        assert_eq!(ctx.span_id(), "00f067aa0ba902b7");
430    }
431
432    #[test]
433    fn from_traceparent_accepts_other_versions() {
434        let ctx = WorkflowTraceContext::from_traceparent(
435            "ff-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01",
436        )
437        .unwrap();
438        assert_eq!(ctx.trace_id(), "4bf92f3577b34da6a3ce929d0e0e4736");
439    }
440
441    #[test]
442    fn from_traceparent_ignores_extra_fields() {
443        let ctx = WorkflowTraceContext::from_traceparent(
444            "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01-extra-stuff",
445        )
446        .unwrap();
447        assert_eq!(ctx.trace_id(), "4bf92f3577b34da6a3ce929d0e0e4736");
448        assert_eq!(ctx.span_id(), "00f067aa0ba902b7");
449    }
450
451    #[test]
452    fn from_traceparent_uppercase_hex_normalized() {
453        let ctx = WorkflowTraceContext::from_traceparent(
454            "00-4BF92F3577B34DA6A3CE929D0E0E4736-00F067AA0BA902B7-01",
455        )
456        .unwrap();
457        assert_eq!(ctx.trace_id(), "4bf92f3577b34da6a3ce929d0e0e4736");
458        assert_eq!(ctx.span_id(), "00f067aa0ba902b7");
459    }
460
461    #[test]
462    fn from_traceparent_invalid_too_few_fields() {
463        let err = WorkflowTraceContext::from_traceparent("00-abc").unwrap_err();
464        assert!(matches!(
465            err,
466            TraceContextError::InvalidFormat { field_count: 2 }
467        ));
468    }
469
470    #[test]
471    fn from_traceparent_invalid_trace_id_length() {
472        let err = WorkflowTraceContext::from_traceparent("00-abc-00f067aa0ba902b7-01").unwrap_err();
473        assert!(matches!(err, TraceContextError::InvalidTraceId { .. }));
474    }
475
476    #[test]
477    fn from_traceparent_invalid_trace_id_hex() {
478        let err = WorkflowTraceContext::from_traceparent(
479            "00-zzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzz-00f067aa0ba902b7-01",
480        )
481        .unwrap_err();
482        assert!(matches!(err, TraceContextError::InvalidTraceId { .. }));
483    }
484
485    #[test]
486    fn from_traceparent_invalid_span_id_length() {
487        let err =
488            WorkflowTraceContext::from_traceparent("00-4bf92f3577b34da6a3ce929d0e0e4736-abc-01")
489                .unwrap_err();
490        assert!(matches!(err, TraceContextError::InvalidSpanId { .. }));
491    }
492
493    #[test]
494    fn from_traceparent_zero_trace_id() {
495        let err = WorkflowTraceContext::from_traceparent(
496            "00-00000000000000000000000000000000-00f067aa0ba902b7-01",
497        )
498        .unwrap_err();
499        assert!(matches!(err, TraceContextError::ZeroTraceId));
500    }
501
502    #[test]
503    fn from_traceparent_zero_span_id() {
504        let err = WorkflowTraceContext::from_traceparent(
505            "00-4bf92f3577b34da6a3ce929d0e0e4736-0000000000000000-01",
506        )
507        .unwrap_err();
508        assert!(matches!(err, TraceContextError::ZeroSpanId));
509    }
510
511    #[test]
512    fn from_traceparent_empty_string() {
513        let err = WorkflowTraceContext::from_traceparent("").unwrap_err();
514        assert!(matches!(err, TraceContextError::InvalidFormat { .. }));
515    }
516
517    #[test]
518    fn roundtrip_to_from_traceparent() {
519        let original = WorkflowTraceContext::new_root();
520        let header = original.to_traceparent();
521        let parsed = WorkflowTraceContext::from_traceparent(&header).unwrap();
522        assert_eq!(parsed.trace_id(), original.trace_id());
523        assert_eq!(parsed.span_id(), original.span_id());
524        assert_eq!(parsed, original);
525    }
526
527    #[test]
528    fn roundtrip_from_workflow_run_id() {
529        let ctx = WorkflowTraceContext::from_workflow_run_id("my-run-42");
530        let header = ctx.to_traceparent();
531        let parsed = WorkflowTraceContext::from_traceparent(&header).unwrap();
532        assert_eq!(parsed.trace_id(), ctx.trace_id());
533        assert_eq!(parsed.span_id(), ctx.span_id());
534    }
535
536    #[test]
537    fn serde_roundtrip() {
538        let ctx = WorkflowTraceContext::new_root();
539        let json = serde_json::to_string(&ctx).unwrap();
540        let back: WorkflowTraceContext = serde_json::from_str(&json).unwrap();
541        assert_eq!(back, ctx);
542    }
543
544    #[test]
545    fn serde_contains_expected_fields() {
546        let ctx = WorkflowTraceContext::from_traceparent(
547            "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01",
548        )
549        .unwrap();
550        let json: serde_json::Value = serde_json::to_value(&ctx).unwrap();
551        assert_eq!(json["trace_id"], "4bf92f3577b34da6a3ce929d0e0e4736");
552        assert_eq!(json["span_id"], "00f067aa0ba902b7");
553    }
554}
ironflow_core/trace_context.rs

ironflow_core/
trace_context.rs
