From f76d12ad843b8103e2da0231d935ee188a6bedc3 Mon Sep 17 00:00:00 2001 From: Dale Seo <5466341+DaleSeo@users.noreply.github.com> Date: Wed, 17 Jun 2026 16:42:26 -0400 Subject: [PATCH] feat: add SEP-414 trace context meta accessors --- crates/rmcp/Cargo.toml | 5 + crates/rmcp/src/model/meta.rs | 142 ++++++++++++++++++++++++ crates/rmcp/tests/test_trace_context.rs | 85 ++++++++++++++ typos.toml | 4 + 4 files changed, 236 insertions(+) create mode 100644 crates/rmcp/tests/test_trace_context.rs create mode 100644 typos.toml diff --git a/crates/rmcp/Cargo.toml b/crates/rmcp/Cargo.toml index 8a5bd63a3..8425f69f4 100644 --- a/crates/rmcp/Cargo.toml +++ b/crates/rmcp/Cargo.toml @@ -288,6 +288,11 @@ name = "test_custom_request" required-features = ["server", "client"] path = "tests/test_custom_request.rs" +[[test]] +name = "test_trace_context" +required-features = ["server", "client"] +path = "tests/test_trace_context.rs" + [[test]] name = "test_prompt_macros" required-features = ["server", "client"] diff --git a/crates/rmcp/src/model/meta.rs b/crates/rmcp/src/model/meta.rs index 186db6a24..5a22f9c8e 100644 --- a/crates/rmcp/src/model/meta.rs +++ b/crates/rmcp/src/model/meta.rs @@ -46,6 +46,34 @@ pub trait RequestParamsMeta { } } } + /// Get the W3C `traceparent` value from meta, if present (SEP-414) + fn traceparent(&self) -> Option<&str> { + self.meta().and_then(|m| m.get_traceparent()) + } + /// Set the W3C `traceparent` value in meta (SEP-414) + fn set_traceparent(&mut self, value: &str) { + self.meta_or_default().set_traceparent(value); + } + /// Get the W3C `tracestate` value from meta, if present (SEP-414) + fn tracestate(&self) -> Option<&str> { + self.meta().and_then(|m| m.get_tracestate()) + } + /// Set the W3C `tracestate` value in meta (SEP-414) + fn set_tracestate(&mut self, value: &str) { + self.meta_or_default().set_tracestate(value); + } + /// Get the W3C `baggage` value from meta, if present (SEP-414) + fn baggage(&self) -> Option<&str> { + self.meta().and_then(|m| m.get_baggage()) + } + /// Set the W3C `baggage` value in meta (SEP-414) + fn set_baggage(&mut self, value: &str) { + self.meta_or_default().set_baggage(value); + } + /// Get a mutable reference to meta, inserting an empty one if absent. + fn meta_or_default(&mut self) -> &mut Meta { + self.meta_mut().get_or_insert_with(Meta::new) + } } /// Trait for task-augmented request params that contain both `_meta` and `task` fields. @@ -198,6 +226,12 @@ variant_extension! { #[expect(clippy::exhaustive_structs, reason = "intentionally exhaustive")] pub struct Meta(pub JsonObject); const PROGRESS_TOKEN_FIELD: &str = "progressToken"; +/// Reserved `_meta` key for the W3C Trace Context `traceparent` value (SEP-414). +pub const TRACEPARENT_FIELD: &str = "traceparent"; +/// Reserved `_meta` key for the W3C Trace Context `tracestate` value (SEP-414). +pub const TRACESTATE_FIELD: &str = "tracestate"; +/// Reserved `_meta` key for the W3C Baggage value (SEP-414). +pub const BAGGAGE_FIELD: &str = "baggage"; impl Meta { pub fn new() -> Self { Self(JsonObject::new()) @@ -247,6 +281,58 @@ impl Meta { }; } + /// Read a string-valued `_meta` field, or `None` if absent or not a string. + fn get_str(&self, field: &str) -> Option<&str> { + self.0.get(field).and_then(Value::as_str) + } + + /// Write a string-valued `_meta` field. + fn set_str(&mut self, field: &str, value: impl Into) { + self.0 + .insert(field.to_string(), Value::String(value.into())); + } + + /// Get the W3C `traceparent` value (SEP-414), if present. + pub fn get_traceparent(&self) -> Option<&str> { + self.get_str(TRACEPARENT_FIELD) + } + + /// Set the W3C `traceparent` value (SEP-414). + /// + /// ``` + /// use rmcp::model::Meta; + /// + /// let mut meta = Meta::new(); + /// meta.set_traceparent("00-0af7651916cd43dd8448eb211c80319c-00f067aa0ba902b7-01"); + /// assert_eq!( + /// meta.get_traceparent(), + /// Some("00-0af7651916cd43dd8448eb211c80319c-00f067aa0ba902b7-01"), + /// ); + /// ``` + pub fn set_traceparent(&mut self, value: impl Into) { + self.set_str(TRACEPARENT_FIELD, value); + } + + /// Get the W3C `tracestate` value (SEP-414), if present. + pub fn get_tracestate(&self) -> Option<&str> { + self.get_str(TRACESTATE_FIELD) + } + + /// Set the W3C `tracestate` value (SEP-414). + pub fn set_tracestate(&mut self, value: impl Into) { + self.set_str(TRACESTATE_FIELD, value); + } + + /// Get the W3C `baggage` value (SEP-414), if present. + pub fn get_baggage(&self) -> Option<&str> { + self.get_str(BAGGAGE_FIELD) + } + + /// Set the W3C `baggage` value (SEP-414). + pub fn set_baggage(&mut self, value: impl Into) { + self.set_str(BAGGAGE_FIELD, value); + } + pub fn extend(&mut self, other: Meta) { for (k, v) in other.0.into_iter() { self.0.insert(k, v); @@ -288,3 +374,59 @@ where } } } + +#[cfg(test)] +mod tests { + use super::*; + + #[derive(Default)] + struct Params { + meta: Option, + } + + impl RequestParamsMeta for Params { + fn meta(&self) -> Option<&Meta> { + self.meta.as_ref() + } + fn meta_mut(&mut self) -> &mut Option { + &mut self.meta + } + } + + const TRACEPARENT: &str = "00-0af7651916cd43dd8448eb211c80319c-00f067aa0ba902b7-01"; + + #[test] + fn trace_context_round_trip() { + let mut meta = Meta::new(); + meta.set_traceparent(TRACEPARENT); + meta.set_tracestate("vendor1=value1,vendor2=value2"); + meta.set_baggage("userId=alice,region=us-east-1"); + assert_eq!(meta.get_traceparent(), Some(TRACEPARENT)); + assert_eq!(meta.get_tracestate(), Some("vendor1=value1,vendor2=value2")); + assert_eq!(meta.get_baggage(), Some("userId=alice,region=us-east-1")); + } + + #[test] + fn absent_field_is_none() { + let meta = Meta::new(); + assert_eq!(meta.get_traceparent(), None); + assert_eq!(meta.get_tracestate(), None); + assert_eq!(meta.get_baggage(), None); + } + + #[test] + fn non_string_value_is_none() { + let mut meta = Meta::new(); + meta.0 + .insert(TRACEPARENT_FIELD.to_string(), Value::from(42)); + assert_eq!(meta.get_traceparent(), None); + } + + #[test] + fn trait_setter_inserts_meta_when_absent() { + let mut params = Params::default(); + assert_eq!(params.traceparent(), None); + params.set_traceparent(TRACEPARENT); + assert_eq!(params.traceparent(), Some(TRACEPARENT)); + } +} diff --git a/crates/rmcp/tests/test_trace_context.rs b/crates/rmcp/tests/test_trace_context.rs new file mode 100644 index 000000000..50214b715 --- /dev/null +++ b/crates/rmcp/tests/test_trace_context.rs @@ -0,0 +1,85 @@ +#![cfg(not(feature = "local"))] +//! SEP-414: the reserved trace-context `_meta` keys survive a client→server round trip unchanged. +use std::sync::Arc; + +use rmcp::{ + RoleServer, ServerHandler, ServiceExt, + model::{ClientRequest, CustomRequest, CustomResult, Meta}, + service::{PeerRequestOptions, RequestContext}, +}; +use serde_json::json; +use tokio::sync::{Mutex, Notify}; + +const TRACEPARENT: &str = "00-0af7651916cd43dd8448eb211c80319c-00f067aa0ba902b7-01"; +const TRACESTATE: &str = "vendor1=value1,vendor2=value2"; +const BAGGAGE: &str = "userId=alice,region=us-east-1"; + +/// Records the `_meta` it receives on the incoming request so the test can assert passthrough. +struct TraceCapturingServer { + receive_signal: Arc, + seen: Arc>>, +} + +impl ServerHandler for TraceCapturingServer { + async fn on_custom_request( + &self, + _request: CustomRequest, + context: RequestContext, + ) -> Result { + *self.seen.lock().await = Some(context.meta); + self.receive_signal.notify_one(); + Ok(CustomResult::new(json!({ "status": "ok" }))) + } +} + +#[tokio::test] +async fn trace_context_meta_survives_round_trip() -> anyhow::Result<()> { + let (server_transport, client_transport) = tokio::io::duplex(4096); + let receive_signal = Arc::new(Notify::new()); + let seen = Arc::new(Mutex::new(None)); + + { + let receive_signal = receive_signal.clone(); + let seen = seen.clone(); + tokio::spawn(async move { + let server = TraceCapturingServer { + receive_signal, + seen, + } + .serve(server_transport) + .await?; + server.waiting().await?; + anyhow::Ok(()) + }); + } + + let client = ().serve(client_transport).await?; + + // Client attaches trace context to the outgoing request's `_meta`. + let mut meta = Meta::new(); + meta.set_traceparent(TRACEPARENT); + meta.set_tracestate(TRACESTATE); + meta.set_baggage(BAGGAGE); + + let mut options = PeerRequestOptions::no_options(); + options.meta = Some(meta); + client + .send_cancellable_request( + ClientRequest::CustomRequest(CustomRequest::new("requests/trace-test", None)), + options, + ) + .await? + .await_response() + .await?; + + tokio::time::timeout(std::time::Duration::from_secs(5), receive_signal.notified()).await?; + + // Server saw the reserved keys unchanged (alongside the injected progressToken). + let seen = seen.lock().await.take().expect("server observed meta"); + assert_eq!(seen.get_traceparent(), Some(TRACEPARENT)); + assert_eq!(seen.get_tracestate(), Some(TRACESTATE)); + assert_eq!(seen.get_baggage(), Some(BAGGAGE)); + + client.cancel().await?; + Ok(()) +} diff --git a/typos.toml b/typos.toml new file mode 100644 index 000000000..38fcb1fd0 --- /dev/null +++ b/typos.toml @@ -0,0 +1,4 @@ +# W3C `traceparent` example values (SEP-414) embed hex spans like `0ba9` that the +# spell checker misreads as typos; ignore the canonical traceparent format. +[default] +extend-ignore-re = ["00-[0-9a-f]{32}-[0-9a-f]{16}-[0-9a-f]{2}"]