modelcontextprotocol/rust-sdk
Watch
1
0
Fork
You've already forked rust-sdk
 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions 1

feat: extensions to context (#102)

Browse source
This commit is contained in:
Eitan Yarmush 2025-04-07 22:01:17 -04:00 • committed by GitHub
commit 6a423048fa
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 43 additions and 9 deletions
Download patch file Download diff file Expand all files Collapse all files

17
crates/rmcp/src/model.rs
View file

@ -200,6 +200,15 @@ pub struct Request<M = String, P = JsonObject> {
pub extensions: Extensions,
}
impl<M, P> GetExtensions for Request<M, P> {
fn extensions(&self) -> &Extensions {
&self.extensions
}
fn extensions_mut(&mut self) -> &mut Extensions {
&mut self.extensions
}
}
#[derive(Debug, Clone)]
pub struct RequestOptionalParam<M = String, P = JsonObject> {
pub method: M,
@ -220,6 +229,14 @@ pub struct RequestNoParam<M = String> {
pub extensions: Extensions,
}
impl<M> GetExtensions for RequestNoParam<M> {
fn extensions(&self) -> &Extensions {
&self.extensions
}
fn extensions_mut(&mut self) -> &mut Extensions {
&mut self.extensions
}
}
#[derive(Debug, Clone)]
pub struct Notification<M = String, P = JsonObject> {
pub method: M,

11
crates/rmcp/src/model/meta.rs
View file

@ -13,21 +13,26 @@ pub trait GetMeta {
fn get_meta(&self) -> &Meta;
}
pub trait GetExtensions {
fn extensions(&self) -> &Extensions;
fn extensions_mut(&mut self) -> &mut Extensions;
}
macro_rules! variant_extension {
(
$Enum: ident {
$($variant: ident)*
}
) => {
impl $Enum {
pub fn extensions(&self) -> &Extensions {
impl GetExtensions for $Enum {
fn extensions(&self) -> &Extensions {
match self {
$(
$Enum::$variant(v) => &v.extensions,
)*
}
}
pub fn extensions_mut(&mut self) -> &mut Extensions {
fn extensions_mut(&mut self) -> &mut Extensions {
match self {
$(
$Enum::$variant(v) => &mut v.extensions,

14
crates/rmcp/src/service.rs
View file

@ -4,10 +4,10 @@ use thiserror::Error;
use crate::{
error::Error as McpError,
model::{
CancelledNotification, CancelledNotificationParam, GetMeta, JsonRpcBatchRequestItem,
JsonRpcBatchResponseItem, JsonRpcError, JsonRpcMessage, JsonRpcNotification,
JsonRpcRequest, JsonRpcResponse, Meta, NumberOrString, ProgressToken, RequestId,
ServerJsonRpcMessage,
CancelledNotification, CancelledNotificationParam, Extensions, GetExtensions, GetMeta,
JsonRpcBatchRequestItem, JsonRpcBatchResponseItem, JsonRpcError, JsonRpcMessage,
JsonRpcNotification, JsonRpcRequest, JsonRpcResponse, Meta, NumberOrString, ProgressToken,
RequestId, ServerJsonRpcMessage,
},
transport::IntoTransport,
};
@ -59,12 +59,12 @@ impl<T> TransferObject for T where
#[allow(private_bounds, reason = "there's no the third implementation")]
pub trait ServiceRole: std::fmt::Debug + Send + Sync + 'static + Copy + Clone {
type Req: TransferObject + GetMeta;
type Req: TransferObject + GetMeta + GetExtensions;
type Resp: TransferObject;
type Not: TryInto<CancelledNotification, Error = Self::Not>
+ From<CancelledNotification>
+ TransferObject;
type PeerReq: TransferObject + GetMeta;
type PeerReq: TransferObject + GetMeta + GetExtensions;
type PeerResp: TransferObject;
type PeerNot: TryInto<CancelledNotification, Error = Self::PeerNot>
+ From<CancelledNotification>
@ -471,6 +471,7 @@ pub struct RequestContext<R: ServiceRole> {
pub ct: CancellationToken,
pub id: RequestId,
pub meta: Meta,
pub extensions: Extensions,
/// An interface to fetch the remote client or server
pub peer: Peer<R>,
}
@ -667,6 +668,7 @@ where
id: id.clone(),
peer: peer.clone(),
meta: request.get_meta().clone(),
extensions: request.extensions().clone(),
};
tokio::spawn(async move {
let result = service.handle_request(request, context).await;

10
crates/rmcp/tests/test_message_protocol.rs
View file

@ -71,6 +71,7 @@ async fn test_context_inclusion_integration() -> anyhow::Result<()> {
ct: CancellationToken::new(),
id: NumberOrString::Number(1),
meta: Default::default(),
extensions: Default::default(),
},
)
.await?;
@ -112,6 +113,7 @@ async fn test_context_inclusion_integration() -> anyhow::Result<()> {
ct: CancellationToken::new(),
id: NumberOrString::Number(2),
meta: Default::default(),
extensions: Default::default(),
},
)
.await?;
@ -153,6 +155,7 @@ async fn test_context_inclusion_integration() -> anyhow::Result<()> {
ct: CancellationToken::new(),
id: NumberOrString::Number(3),
meta: Default::default(),
extensions: Default::default(),
},
)
.await?;
@ -214,6 +217,7 @@ async fn test_context_inclusion_ignored_integration() -> anyhow::Result<()> {
ct: CancellationToken::new(),
id: NumberOrString::Number(1),
meta: Meta::default(),
extensions: Default::default(),
},
)
.await?;
@ -280,6 +284,7 @@ async fn test_message_sequence_integration() -> anyhow::Result<()> {
ct: CancellationToken::new(),
id: NumberOrString::Number(1),
meta: Meta::default(),
extensions: Default::default(),
},
)
.await?;
@ -354,6 +359,7 @@ async fn test_message_sequence_validation_integration() -> anyhow::Result<()> {
ct: CancellationToken::new(),
id: NumberOrString::Number(1),
meta: Meta::default(),
extensions: Default::default(),
},
)
.await?;
@ -387,6 +393,7 @@ async fn test_message_sequence_validation_integration() -> anyhow::Result<()> {
ct: CancellationToken::new(),
id: NumberOrString::Number(2),
meta: Meta::default(),
extensions: Default::default(),
},
)
.await;
@ -439,6 +446,7 @@ async fn test_selective_context_handling_integration() -> anyhow::Result<()> {
ct: CancellationToken::new(),
id: NumberOrString::Number(1),
meta: Meta::default(),
extensions: Default::default(),
},
)
.await?;
@ -478,6 +486,7 @@ async fn test_selective_context_handling_integration() -> anyhow::Result<()> {
ct: CancellationToken::new(),
id: NumberOrString::Number(2),
meta: Meta::default(),
extensions: Default::default(),
},
)
.await?;
@ -534,6 +543,7 @@ async fn test_context_inclusion() -> anyhow::Result<()> {
ct: CancellationToken::new(),
id: NumberOrString::Number(1),
meta: Meta::default(),
extensions: Default::default(),
},
)
.await?;