diff --git a/Cargo.toml b/Cargo.toml index 772e82e..e6494f0 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -41,6 +41,10 @@ which = "7" schemars = { version = "0.8", optional = true, features = ["derive"] } +# OpenTelemetry (optional — enabled via `opentelemetry` feature) +opentelemetry = { version = "0.31", optional = true } +opentelemetry_sdk = { version = "0.31", optional = true } + [dev-dependencies] tokio-test = "0.4" pretty_assertions = "1" @@ -54,6 +58,7 @@ default = [] e2e = [] # Enable E2E tests that require real Copilot CLI snapshots = [] # Enable snapshot-based conformance tests schemars = ["dep:schemars"] # Derive tool JSON schemas from Rust types +opentelemetry = ["dep:opentelemetry", "dep:opentelemetry_sdk"] # Auto-propagate W3C Trace Context via OpenTelemetry [package.metadata.docs.rs] rustdoc-args = ["--cfg", "docsrs"] diff --git a/README.md b/README.md index 01dd4b6..19609bf 100644 --- a/README.md +++ b/README.md @@ -201,6 +201,16 @@ let client = Client::builder() .build()?; ``` +#### Trace Context Propagation + +Enable the `opentelemetry` feature to automatically propagate W3C Trace Context (`traceparent`/`tracestate`) from your app into every JSON-RPC call to the CLI. This links your application spans with the CLI's internal spans in the same distributed trace. + +```toml +copilot-sdk = { version = "0.1", features = ["opentelemetry"] } +``` + +No wiring is required. Before each `session.create`, `session.resume`, and `session.send` call, the SDK reads the current `opentelemetry::Context` using a local W3C `TraceContextPropagator` and injects the headers into the request. If no span is active, nothing is injected. + ### BYOK (Bring Your Own Key) Use your own API keys with compatible providers, with custom model listing: diff --git a/src/client.rs b/src/client.rs index 13c5de3..894ee71 100644 --- a/src/client.rs +++ b/src/client.rs @@ -705,7 +705,11 @@ impl Client { } // Build the request - let params = serde_json::to_value(&config)?; + #[allow(unused_mut)] + let mut params = serde_json::to_value(&config)?; + + #[cfg(feature = "opentelemetry")] + crate::otel::inject_trace_context(&mut params); // Send the request let result = self.invoke("session.create", Some(params)).await?; @@ -761,6 +765,9 @@ impl Client { let mut params = serde_json::to_value(&config)?; params["sessionId"] = json!(session_id); + #[cfg(feature = "opentelemetry")] + crate::otel::inject_trace_context(&mut params); + // Send the request let result = self.invoke("session.resume", Some(params)).await?; diff --git a/src/events.rs b/src/events.rs index 6748266..1221f93 100644 --- a/src/events.rs +++ b/src/events.rs @@ -572,6 +572,10 @@ pub struct ExternalToolRequestedData { pub tool_call_id: Option, /// Arguments to pass to the tool handler. pub arguments: Option, + /// W3C traceparent header propagated from the CLI (if tracing is active). + pub traceparent: Option, + /// W3C tracestate header propagated from the CLI (if tracing is active). + pub tracestate: Option, } /// Data for `permission.requested` event (protocol v3 broadcast model). @@ -979,6 +983,32 @@ mod tests { } } + #[test] + fn test_parse_external_tool_requested_with_trace_context() { + let json = json!({ + "id": "evt_trace", + "timestamp": "2024-01-15T10:30:02Z", + "type": "external_tool.requested", + "data": { + "requestId": "req_t1", + "toolName": "echo", + "toolCallId": "call_t1", + "arguments": { "text": "hi" }, + "traceparent": "00-abc123-def456-01", + "tracestate": "vendor=value" + } + }); + + let event = SessionEvent::from_json(&json).unwrap(); + match &event.data { + SessionEventData::ExternalToolRequested(data) => { + assert_eq!(data.traceparent.as_deref(), Some("00-abc123-def456-01")); + assert_eq!(data.tracestate.as_deref(), Some("vendor=value")); + } + other => panic!("Expected ExternalToolRequested, got {other:?}"), + } + } + #[test] fn test_parse_permission_requested() { let json = json!({ diff --git a/src/lib.rs b/src/lib.rs index 29499e7..fec45d0 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -45,6 +45,9 @@ pub mod tools; pub mod transport; pub mod types; +#[cfg(feature = "opentelemetry")] +pub(crate) mod otel; + // Re-export tool utilities pub use tools::define_tool; diff --git a/src/otel.rs b/src/otel.rs new file mode 100644 index 0000000..ecbf7dc --- /dev/null +++ b/src/otel.rs @@ -0,0 +1,93 @@ +// Copyright (c) 2026 Elias Bachaalany +// SPDX-License-Identifier: MIT + +//! Built-in W3C Trace Context propagation for OpenTelemetry. +//! +//! When the `opentelemetry` feature is enabled, [`inject_trace_context`] is +//! called before each outgoing `session.create`, `session.resume`, and +//! `session.send` JSON-RPC request. It reads `traceparent` / `tracestate` +//! from the current [`opentelemetry::Context`] using a local +//! [`TraceContextPropagator`] and merges them into the params object so the +//! CLI's spans join the same distributed trace. +//! +//! A local propagator is used on purpose: the application does **not** need +//! to call `opentelemetry::global::set_text_map_propagator(...)` — the +//! default global propagator is a no-op and would silently drop the context. +//! +//! If no OpenTelemetry context is active (no `TracerProvider` set up, or no +//! span on the stack), this is a silent no-op — no errors, no panics. + +use opentelemetry::propagation::{Injector, TextMapPropagator}; +use opentelemetry::Context; +use opentelemetry_sdk::propagation::TraceContextPropagator; + +/// Inject W3C `traceparent` / `tracestate` from the current OpenTelemetry +/// context into a JSON-RPC params object. +/// +/// Does nothing if the context carries no span or if `params` is not a JSON +/// object. +pub(crate) fn inject_trace_context(params: &mut serde_json::Value) { + let serde_json::Value::Object(map) = params else { + return; + }; + TraceContextPropagator::new().inject_context(&Context::current(), &mut JsonMapInjector(map)); +} + +struct JsonMapInjector<'a>(&'a mut serde_json::Map); + +impl Injector for JsonMapInjector<'_> { + fn set(&mut self, key: &str, value: String) { + if key == "traceparent" || key == "tracestate" { + self.0 + .insert(key.to_string(), serde_json::Value::String(value)); + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use opentelemetry::trace::{ + SpanContext, SpanId, TraceContextExt, TraceFlags, TraceId, TraceState, + }; + + #[test] + fn inject_without_active_span_is_noop() { + let mut params = serde_json::json!({ "sessionId": "s1" }); + inject_trace_context(&mut params); + assert_eq!(params, serde_json::json!({ "sessionId": "s1" })); + } + + #[test] + fn inject_into_non_object_is_noop() { + let mut params = serde_json::json!("not an object"); + inject_trace_context(&mut params); + assert_eq!(params, serde_json::json!("not an object")); + } + + #[test] + fn inject_with_active_span_writes_traceparent() { + // Synthetic remote span context attached to the current OTel context. + // The local TraceContextPropagator should produce a well-formed W3C + // traceparent without any global propagator being registered. + let trace_id = TraceId::from_hex("4bf92f3577b34da6a3ce929d0e0e4736").unwrap(); + let span_id = SpanId::from_hex("00f067aa0ba902b7").unwrap(); + let span_ctx = SpanContext::new( + trace_id, + span_id, + TraceFlags::SAMPLED, + true, + TraceState::default(), + ); + let cx = Context::current().with_remote_span_context(span_ctx); + let _guard = cx.attach(); + + let mut params = serde_json::json!({ "sessionId": "s1" }); + inject_trace_context(&mut params); + assert_eq!(params["sessionId"], "s1"); + assert_eq!( + params["traceparent"], + "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01" + ); + } +} diff --git a/src/session.rs b/src/session.rs index e3d30aa..418fb60 100644 --- a/src/session.rs +++ b/src/session.rs @@ -380,13 +380,17 @@ impl Session { /// Returns the message ID. pub async fn send(&self, options: impl Into) -> Result { let options = options.into(); - let params = serde_json::json!({ + #[allow(unused_mut)] + let mut params = serde_json::json!({ "sessionId": self.session_id, "prompt": options.prompt, "attachments": options.attachments, "mode": options.mode, }); + #[cfg(feature = "opentelemetry")] + crate::otel::inject_trace_context(&mut params); + let result = (self.invoke_fn)("session.send", Some(params)).await?; result diff --git a/src/types.rs b/src/types.rs index e1cfb4e..f87efb0 100644 --- a/src/types.rs +++ b/src/types.rs @@ -157,6 +157,12 @@ pub struct ToolInvocation { pub tool_name: String, #[serde(default)] pub arguments: Option, + /// W3C Trace Context `traceparent` from the CLI's execute_tool span. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub traceparent: Option, + /// W3C Trace Context `tracestate` from the CLI's execute_tool span. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub tracestate: Option, } impl ToolInvocation { @@ -2110,4 +2116,34 @@ mod tests { assert_eq!(value["clientName"], "my-cli"); assert_eq!(value["agent"], "helper"); } + + #[test] + fn test_tool_invocation_deserialize_with_trace_fields() { + let json = serde_json::json!({ + "sessionId": "s1", + "toolName": "my_tool", + "toolCallId": "call-1", + "arguments": {"key": "value"}, + "traceparent": "00-abc-def-01", + "tracestate": "vendor=val" + }); + let invocation: ToolInvocation = serde_json::from_value(json).unwrap(); + assert_eq!(invocation.tool_name, "my_tool"); + assert_eq!(invocation.traceparent.as_deref(), Some("00-abc-def-01")); + assert_eq!(invocation.tracestate.as_deref(), Some("vendor=val")); + } + + #[test] + fn test_tool_invocation_deserialize_without_trace_fields() { + let json = serde_json::json!({ + "sessionId": "s1", + "toolName": "my_tool", + "toolCallId": "call-1", + "arguments": {"key": "value"} + }); + let invocation: ToolInvocation = serde_json::from_value(json).unwrap(); + assert_eq!(invocation.tool_name, "my_tool"); + assert!(invocation.traceparent.is_none()); + assert!(invocation.tracestate.is_none()); + } }