feat: extensions to context (#102)

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

View file

@ -200,6 +200,15 @@ pub struct Request<M = String, P = JsonObject> {
pub extensions: Extensions, 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)] #[derive(Debug, Clone)]
pub struct RequestOptionalParam<M = String, P = JsonObject> { pub struct RequestOptionalParam<M = String, P = JsonObject> {
pub method: M, pub method: M,
@ -220,6 +229,14 @@ pub struct RequestNoParam<M = String> {
pub extensions: Extensions, 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)] #[derive(Debug, Clone)]
pub struct Notification<M = String, P = JsonObject> { pub struct Notification<M = String, P = JsonObject> {
pub method: M, pub method: M,

View file

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

View file

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

View file

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