feat: add SEP-414 trace context meta accessors (#910)
This commit is contained in:
parent
496902b9cf
commit
ee2d81f4d0
4 changed files with 236 additions and 0 deletions
|
|
@ -304,6 +304,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"]
|
||||
|
|
|
|||
|
|
@ -47,6 +47,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.
|
||||
|
|
@ -207,6 +235,12 @@ impl Meta {
|
|||
const META_KEY_CLIENT_INFO: &str = "io.modelcontextprotocol/clientInfo";
|
||||
const META_KEY_CLIENT_CAPABILITIES: &str = "io.modelcontextprotocol/clientCapabilities";
|
||||
const META_KEY_LOG_LEVEL: &str = "io.modelcontextprotocol/logLevel";
|
||||
/// Reserved `_meta` key for the W3C Trace Context `traceparent` value (SEP-414).
|
||||
const TRACEPARENT_FIELD: &str = "traceparent";
|
||||
/// Reserved `_meta` key for the W3C Trace Context `tracestate` value (SEP-414).
|
||||
const TRACESTATE_FIELD: &str = "tracestate";
|
||||
/// Reserved `_meta` key for the W3C Baggage value (SEP-414).
|
||||
const BAGGAGE_FIELD: &str = "baggage";
|
||||
|
||||
pub fn new() -> Self {
|
||||
Self(JsonObject::new())
|
||||
|
|
@ -304,6 +338,58 @@ impl Meta {
|
|||
self.insert_serialized(Self::META_KEY_LOG_LEVEL, log_level);
|
||||
}
|
||||
|
||||
/// 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<String>) {
|
||||
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(Self::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<String>) {
|
||||
self.set_str(Self::TRACEPARENT_FIELD, value);
|
||||
}
|
||||
|
||||
/// Get the W3C `tracestate` value (SEP-414), if present.
|
||||
pub fn get_tracestate(&self) -> Option<&str> {
|
||||
self.get_str(Self::TRACESTATE_FIELD)
|
||||
}
|
||||
|
||||
/// Set the W3C `tracestate` value (SEP-414).
|
||||
pub fn set_tracestate(&mut self, value: impl Into<String>) {
|
||||
self.set_str(Self::TRACESTATE_FIELD, value);
|
||||
}
|
||||
|
||||
/// Get the W3C `baggage` value (SEP-414), if present.
|
||||
pub fn get_baggage(&self) -> Option<&str> {
|
||||
self.get_str(Self::BAGGAGE_FIELD)
|
||||
}
|
||||
|
||||
/// Set the W3C `baggage` value (SEP-414).
|
||||
pub fn set_baggage(&mut self, value: impl Into<String>) {
|
||||
self.set_str(Self::BAGGAGE_FIELD, value);
|
||||
}
|
||||
|
||||
pub fn extend(&mut self, other: Meta) {
|
||||
for (k, v) in other.0.into_iter() {
|
||||
self.0.insert(k, v);
|
||||
|
|
@ -361,3 +447,59 @@ where
|
|||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[derive(Default)]
|
||||
struct Params {
|
||||
meta: Option<Meta>,
|
||||
}
|
||||
|
||||
impl RequestParamsMeta for Params {
|
||||
fn meta(&self) -> Option<&Meta> {
|
||||
self.meta.as_ref()
|
||||
}
|
||||
fn meta_mut(&mut self) -> &mut Option<Meta> {
|
||||
&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(Meta::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));
|
||||
}
|
||||
}
|
||||
|
|
|
|||
85
crates/rmcp/tests/test_trace_context.rs
Normal file
85
crates/rmcp/tests/test_trace_context.rs
Normal file
|
|
@ -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<Notify>,
|
||||
seen: Arc<Mutex<Option<Meta>>>,
|
||||
}
|
||||
|
||||
impl ServerHandler for TraceCapturingServer {
|
||||
async fn on_custom_request(
|
||||
&self,
|
||||
_request: CustomRequest,
|
||||
context: RequestContext<RoleServer>,
|
||||
) -> Result<CustomResult, rmcp::ErrorData> {
|
||||
*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(())
|
||||
}
|
||||
4
typos.toml
Normal file
4
typos.toml
Normal file
|
|
@ -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}"]
|
||||
Loading…
Reference in a new issue