Skip to main content

github_copilot_sdk/
trace_context.rs

1//! W3C Trace Context propagation for distributed tracing.
2//!
3//! The GitHub Copilot CLI propagates [W3C Trace Context] headers (`traceparent`
4//! and `tracestate`) so SDK consumers can correlate spans created by the
5//! CLI with their own observability pipelines.
6//!
7//! Two injection paths are supported:
8//!
9//! - **Per-turn override** via [`MessageOptions::traceparent`] /
10//!   [`MessageOptions::tracestate`](crate::types::MessageOptions::tracestate),
11//!   which take precedence when set.
12//! - **Ambient callback** via
13//!   [`ClientOptions::on_get_trace_context`](crate::ClientOptions::on_get_trace_context),
14//!   which the SDK invokes before `session.create`, `session.resume`, and
15//!   `session.send` whenever the per-turn override is absent.
16//!
17//! [W3C Trace Context]: https://www.w3.org/TR/trace-context/
18//! [`MessageOptions::traceparent`]: crate::types::MessageOptions::traceparent
19
20use async_trait::async_trait;
21
22/// W3C Trace Context headers propagated to and from the GitHub Copilot CLI.
23///
24/// `traceparent` carries the trace and parent-span identifiers; `tracestate`
25/// carries vendor-specific extensions. Either field may be `None` when the
26/// caller has nothing to propagate; in that case the corresponding wire
27/// field is omitted.
28#[derive(Debug, Clone, Default, PartialEq, Eq)]
29#[non_exhaustive]
30pub struct TraceContext {
31    /// `traceparent` HTTP header value.
32    pub traceparent: Option<String>,
33    /// `tracestate` HTTP header value.
34    pub tracestate: Option<String>,
35}
36
37impl TraceContext {
38    /// Construct an empty [`TraceContext`]; both fields default to unset
39    /// (the SDK skips trace-context injection on the wire).
40    pub fn new() -> Self {
41        Self::default()
42    }
43
44    /// Construct a [`TraceContext`] from a `traceparent` header value, with
45    /// no `tracestate`.
46    ///
47    /// Equivalent to `TraceContext::new().with_traceparent(value)`; kept
48    /// for ergonomics in the common single-header case.
49    pub fn from_traceparent(traceparent: impl Into<String>) -> Self {
50        Self::new().with_traceparent(traceparent)
51    }
52
53    /// Set or replace the `traceparent` header value, returning `self` for
54    /// chaining.
55    pub fn with_traceparent(mut self, traceparent: impl Into<String>) -> Self {
56        self.traceparent = Some(traceparent.into());
57        self
58    }
59
60    /// Set or replace the `tracestate` header value, returning `self` for
61    /// chaining.
62    pub fn with_tracestate(mut self, tracestate: impl Into<String>) -> Self {
63        self.tracestate = Some(tracestate.into());
64        self
65    }
66
67    /// Returns `true` when neither `traceparent` nor `tracestate` is set.
68    pub fn is_empty(&self) -> bool {
69        self.traceparent.is_none() && self.tracestate.is_none()
70    }
71}
72
73/// Async provider that returns the current [`TraceContext`] for outbound
74/// session RPCs.
75///
76/// Set via
77/// [`ClientOptions::on_get_trace_context`](crate::ClientOptions::on_get_trace_context).
78/// The SDK invokes [`get_trace_context`](Self::get_trace_context) before
79/// each `session.create`, `session.resume`, and `session.send` whenever
80/// the call site does not carry a per-turn override.
81///
82/// Implementations should handle errors internally and return
83/// [`TraceContext::default()`] to skip injection — no `Result` return type
84/// is exposed because trace propagation is a best-effort observability
85/// feature, not a correctness-critical RPC parameter.
86#[async_trait]
87pub trait TraceContextProvider: Send + Sync + 'static {
88    /// Return the current trace context, or [`TraceContext::default()`] to
89    /// skip injection.
90    async fn get_trace_context(&self) -> TraceContext;
91}
92
93/// Inject `traceparent` / `tracestate` from `ctx` into the JSON `params`
94/// object if either field is set. No-op when both are `None`.
95pub(crate) fn inject_trace_context(params: &mut serde_json::Value, ctx: &TraceContext) {
96    if let Some(tp) = &ctx.traceparent {
97        params["traceparent"] = serde_json::Value::String(tp.clone());
98    }
99    if let Some(ts) = &ctx.tracestate {
100        params["tracestate"] = serde_json::Value::String(ts.clone());
101    }
102}
103
104#[cfg(test)]
105mod tests {
106    use super::TraceContext;
107
108    #[test]
109    fn new_yields_empty_context() {
110        let ctx = TraceContext::new();
111        assert!(ctx.is_empty());
112        assert!(ctx.traceparent.is_none());
113        assert!(ctx.tracestate.is_none());
114    }
115
116    #[test]
117    fn builder_composes_traceparent_and_tracestate() {
118        let ctx = TraceContext::new()
119            .with_traceparent("00-trace-span-01")
120            .with_tracestate("vendor=key");
121        assert_eq!(ctx.traceparent.as_deref(), Some("00-trace-span-01"));
122        assert_eq!(ctx.tracestate.as_deref(), Some("vendor=key"));
123        assert!(!ctx.is_empty());
124    }
125
126    #[test]
127    fn from_traceparent_matches_builder() {
128        let direct = TraceContext::from_traceparent("00-trace-span-01");
129        let chained = TraceContext::new().with_traceparent("00-trace-span-01");
130        assert_eq!(direct, chained);
131    }
132}