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}