From 861c79a008c7ef530029ae3d6aa5d286035f86d3 Mon Sep 17 00:00:00 2001 From: muxammadreza <137672463+muxammadreza@users.noreply.github.com> Date: Wed, 12 Aug 2026 18:56:20 -0700 Subject: [PATCH 1/6] fix: preserve open capability and subscription fields --- crates/rmcp/src/handler/server.rs | 10 ++-- crates/rmcp/src/model.rs | 21 +++++++- crates/rmcp/src/model/capabilities.rs | 52 +++++++++++++++++++ crates/rmcp/tests/test_subscriptions.rs | 49 +++++++++++++++++ crates/rmcp/tests/test_subscriptions_model.rs | 18 +++++++ 5 files changed, 144 insertions(+), 6 deletions(-) diff --git a/crates/rmcp/src/handler/server.rs b/crates/rmcp/src/handler/server.rs index 6dc7883ed..d6d7aa2fa 100644 --- a/crates/rmcp/src/handler/server.rs +++ b/crates/rmcp/src/handler/server.rs @@ -155,9 +155,8 @@ impl Service for H { ); }; let server_info = self.get_info(); - let advertised = requested.supported_by(&server_info.capabilities); let handler_accepted = requested.intersection(&candidate); - let accepted = handler_accepted.intersection(&advertised); + let accepted = handler_accepted.supported_by(&server_info.capabilities); if accepted != handler_accepted { tracing::debug!( requested_resource_count = requested @@ -405,9 +404,10 @@ macro_rules! server_handler_methods { /// Return the subset of a requested notification filter this server accepts. /// /// Returning `None` leaves `subscriptions/listen` unimplemented. The SDK - /// intersects the returned filter with both `requested` and the notification - /// capabilities advertised by [`Self::get_info`] before acknowledging it. - /// Categories that were not requested or advertised are always removed. + /// intersects the returned filter with `requested`, then filters the core + /// notification categories against capabilities advertised by [`Self::get_info`]. + /// Extension-owned fields returned by this handler are preserved when the + /// client requested the same field because the SDK does not own their schema. fn accepted_subscription_filter( &self, requested: &SubscriptionFilter, diff --git a/crates/rmcp/src/model.rs b/crates/rmcp/src/model.rs index 6a1409870..34a117e8f 100644 --- a/crates/rmcp/src/model.rs +++ b/crates/rmcp/src/model.rs @@ -1927,6 +1927,13 @@ pub struct SubscriptionFilter { #[serde(default, skip_serializing_if = "Option::is_none")] #[cfg_attr(feature = "schemars", schemars(with = "Vec"))] pub resource_subscriptions: Option>, + /// Extension-owned subscription filter entries not modeled by the core SDK. + /// + /// The server handler remains responsible for validating and accepting + /// these values. Core capability filtering preserves handler-accepted + /// extension fields instead of interpreting their schemas. + #[serde(flatten)] + pub additional_fields: JsonObject, } impl SubscriptionFilter { @@ -1955,6 +1962,12 @@ impl SubscriptionFilter { }) }) .filter(|uris: &Vec| !uris.is_empty()); + let additional_fields = other + .additional_fields + .iter() + .filter(|(key, _)| self.additional_fields.contains_key(*key)) + .map(|(key, value)| (key.clone(), value.clone())) + .collect(); Self { tools_list_changed: (self.tools_list_changed == Some(true) && other.tools_list_changed == Some(true)) @@ -1966,6 +1979,7 @@ impl SubscriptionFilter { && other.resources_list_changed == Some(true)) .then_some(true), resource_subscriptions, + additional_fields, } } @@ -1986,7 +2000,11 @@ impl SubscriptionFilter { .is_some_and(|requested| requested.contains(uri)) }) }); - booleans_are_subset && resources_are_subset + let additional_fields_are_subset = self + .additional_fields + .keys() + .all(|key| other.additional_fields.contains_key(key)); + booleans_are_subset && resources_are_subset && additional_fields_are_subset } /// Return the requested notification types advertised by server capabilities. @@ -2016,6 +2034,7 @@ impl SubscriptionFilter { .is_some_and(|resources| resources.subscribe == Some(true)) .then(|| self.resource_subscriptions.clone()) .flatten(), + additional_fields: self.additional_fields.clone(), } } } diff --git a/crates/rmcp/src/model/capabilities.rs b/crates/rmcp/src/model/capabilities.rs index f014f569f..dc003548d 100644 --- a/crates/rmcp/src/model/capabilities.rs +++ b/crates/rmcp/src/model/capabilities.rs @@ -192,6 +192,13 @@ pub struct ClientCapabilities { /// Capability to handle elicitation requests from servers for interactive user input #[serde(skip_serializing_if = "Option::is_none")] pub elicitation: Option, + /// Additional capability entries not yet modeled by this SDK. + /// + /// MCP capability objects are open sets. Retaining these values allows + /// intermediaries to preserve additive or extension-defined capabilities + /// without requiring a new SDK release for every field. + #[serde(flatten)] + pub additional_capabilities: JsonObject, } impl ClientCapabilities { @@ -240,6 +247,13 @@ pub struct ServerCapabilities { pub resources: Option, #[serde(skip_serializing_if = "Option::is_none")] pub tools: Option, + /// Additional capability entries not yet modeled by this SDK. + /// + /// MCP capability objects are open sets. Retaining these values allows + /// intermediaries to preserve additive or extension-defined capabilities + /// without requiring a new SDK release for every field. + #[serde(flatten)] + pub additional_capabilities: JsonObject, } impl ServerCapabilities { @@ -277,6 +291,7 @@ macro_rules! builder { pub fn build(self) -> $Target { $Target { $( $f: self.$f, )* + ..Default::default() } } } @@ -677,4 +692,41 @@ mod test { serde_json::json!({}) ); } + + #[test] + fn client_capabilities_preserve_unknown_top_level_fields() { + let input = serde_json::json!({ + "sampling": {}, + "com.example/futureCoreCapability": { + "mode": "future", + "nested": [1, {"enabled": true}] + } + }); + + let capabilities: ClientCapabilities = serde_json::from_value(input.clone()).unwrap(); + let output = serde_json::to_value(capabilities).unwrap(); + + assert_eq!( + output["com.example/futureCoreCapability"], + input["com.example/futureCoreCapability"] + ); + } + + #[test] + fn server_capabilities_preserve_unknown_top_level_fields() { + let input = serde_json::json!({ + "tools": {"listChanged": true}, + "com.example/futureCoreCapability": { + "arbitrary": ["data", 42] + } + }); + + let capabilities: ServerCapabilities = serde_json::from_value(input.clone()).unwrap(); + let output = serde_json::to_value(capabilities).unwrap(); + + assert_eq!( + output["com.example/futureCoreCapability"], + input["com.example/futureCoreCapability"] + ); + } } diff --git a/crates/rmcp/tests/test_subscriptions.rs b/crates/rmcp/tests/test_subscriptions.rs index 8ceef93cd..e7f6eb8a0 100644 --- a/crates/rmcp/tests/test_subscriptions.rs +++ b/crates/rmcp/tests/test_subscriptions.rs @@ -117,6 +117,33 @@ impl ServerHandler for ToolsAndPromptsServer { } } +struct ExtensionFilterServer; + +impl ServerHandler for ExtensionFilterServer { + fn accepted_subscription_filter( + &self, + _requested: &SubscriptionFilter, + ) -> Option { + Some( + serde_json::from_value(serde_json::json!({ + "taskIds": ["task-a"], + "com.example/filter": {"channels": ["alpha"]} + })) + .expect("extension filter candidate"), + ) + } + + async fn listen(&self, context: SubscriptionContext) -> Result<(), rmcp::ErrorData> { + let accepted = serde_json::to_value(context.accepted()).expect("serialize accepted filter"); + assert_eq!(accepted["taskIds"], serde_json::json!(["task-a"])); + assert_eq!( + accepted["com.example/filter"], + serde_json::json!({"channels": ["alpha"]}) + ); + Ok(()) + } +} + struct ResourceSubscriptionServer; impl ServerHandler for ResourceSubscriptionServer { @@ -400,6 +427,28 @@ async fn modern_client( .map_err(Into::into) } +#[tokio::test] +async fn extension_subscription_acknowledges_handler_accepted_subset() -> anyhow::Result<()> { + let client = modern_client(ExtensionFilterServer).await?; + let requested: SubscriptionFilter = serde_json::from_value(serde_json::json!({ + "taskIds": ["task-a", "task-b"], + "com.example/filter": {"channels": ["alpha", "beta"]} + }))?; + let mut subscription = client.listen(requested).await?; + + assert_eq!( + serde_json::to_value(subscription.acknowledged())?, + serde_json::json!({ + "taskIds": ["task-a"], + "com.example/filter": {"channels": ["alpha"]} + }) + ); + assert!(subscription.next().await?.is_none()); + + client.cancel().await?; + Ok(()) +} + #[tokio::test] async fn listen_exposes_acknowledged_filter_and_graceful_result() -> anyhow::Result<()> { let client = modern_client(ToolsOnlyServer).await?; diff --git a/crates/rmcp/tests/test_subscriptions_model.rs b/crates/rmcp/tests/test_subscriptions_model.rs index 8d79fb119..70091ab78 100644 --- a/crates/rmcp/tests/test_subscriptions_model.rs +++ b/crates/rmcp/tests/test_subscriptions_model.rs @@ -38,6 +38,24 @@ fn subscription_filter_subset_is_order_independent_and_ignores_false_flags() { assert!(accepted.is_subset_of(&requested)); } +#[test] +fn subscription_filter_preserves_extension_owned_fields_through_serde() { + let input = json!({ + "taskIds": ["task-a", "task-b"], + "com.example/filter": { + "channels": ["alpha", "beta"], + "nested": {"enabled": true} + } + }); + + let filter: SubscriptionFilter = + serde_json::from_value(input.clone()).expect("deserialize extension filter"); + let output = serde_json::to_value(filter).expect("serialize extension filter"); + + assert_eq!(output["taskIds"], input["taskIds"]); + assert_eq!(output["com.example/filter"], input["com.example/filter"]); +} + #[test] fn subscription_filter_omits_empty_resource_intersection() { let requested = SubscriptionFilter::builder() From d304c4138da6375646fe2493517d5d04be22f322 Mon Sep 17 00:00:00 2001 From: muxammadreza <137672463+muxammadreza@users.noreply.github.com> Date: Wed, 12 Aug 2026 20:05:45 -0700 Subject: [PATCH 2/6] fix: scope discovery cache by client contract --- crates/rmcp/Cargo.toml | 1 + crates/rmcp/src/service/client.rs | 64 ++++++++++++++++++++++++++----- 2 files changed, 56 insertions(+), 9 deletions(-) diff --git a/crates/rmcp/Cargo.toml b/crates/rmcp/Cargo.toml index 9dd27a251..51bc9cf7c 100644 --- a/crates/rmcp/Cargo.toml +++ b/crates/rmcp/Cargo.toml @@ -48,6 +48,7 @@ rustdoc-args = ["--cfg", "docsrs"] async-trait = { version = "0.1.89", optional = true } serde = { version = "1.0", features = ["derive", "rc"] } serde_json = "1.0" +serde_json_canonicalizer = "0.3.2" thiserror = "2" tokio = { version = "1", features = ["sync", "macros", "rt", "time"] } futures = "0.3" diff --git a/crates/rmcp/src/service/client.rs b/crates/rmcp/src/service/client.rs index 520410fb1..053ab6a05 100644 --- a/crates/rmcp/src/service/client.rs +++ b/crates/rmcp/src/service/client.rs @@ -1033,13 +1033,25 @@ const RESOURCE_LIST_CACHE_PREFIX: &str = "resources/list:"; const RESOURCE_TEMPLATE_LIST_CACHE_PREFIX: &str = "resources/templates/list:"; const RESOURCE_READ_CACHE_PREFIX: &str = "resources/read:"; -// Cache keys are built only from the request method plus the parameters that -// affect the result (SEP-2549). Request `_meta` (progress tokens, trace -// context, etc.) does not affect the result, so it is deliberately excluded to -// avoid fragmenting the cache across otherwise-identical requests. -fn discover_cache_key() -> String { - // `server/discover` carries no result-affecting parameters. - DISCOVER_CACHE_PREFIX.to_string() +// Cache keys are built only from the request method plus the parameters or +// reserved request metadata that affect the result (SEP-2549 / SEP-2575). +// Trace context and client identity are deliberately excluded; discovery may +// vary by protocol revision and the complete open client capability document. +fn discover_cache_key(meta: &RequestMetaObject) -> String { + let mut result_affecting_meta = serde_json::Map::new(); + for key in [ + "io.modelcontextprotocol/protocolVersion", + "io.modelcontextprotocol/clientCapabilities", + ] { + if let Some(value) = meta.get(key) { + result_affecting_meta.insert(key.to_string(), value.clone()); + } + } + let canonical = + serde_json_canonicalizer::to_vec(&serde_json::Value::Object(result_affecting_meta)) + .expect("request metadata is already valid JSON"); + let canonical = String::from_utf8(canonical).expect("canonical JSON is UTF-8"); + format!("{DISCOVER_CACHE_PREFIX}{canonical}") } fn list_response_cache_key(prefix: &str, params: &Option) -> String { @@ -1267,7 +1279,7 @@ impl Peer { /// The high-level client currently exposes this peer only after initialization; /// pre-initialization probing is planned as follow-up work. pub async fn discover(&self, meta: RequestMetaObject) -> Result { - let cache_key = discover_cache_key(); + let cache_key = discover_cache_key(&meta); if let Some(ServerResult::DiscoverResult(result)) = self.cached_response(&cache_key).await { return Ok(result); } @@ -2387,11 +2399,45 @@ mod tests { assert_eq!(peer.list_tools(params).await.unwrap(), expected); } + #[test] + fn discover_cache_key_tracks_result_affecting_request_metadata_canonically() { + let mut first = RequestMetaObject::new(); + first.set_protocol_version(ProtocolVersion::V_2026_07_28); + first.set_client_capabilities( + serde_json::from_value(serde_json::json!({ + "sampling": {}, + "com.example/future": {"b": 2, "a": 1} + })) + .expect("open client capability fixture"), + ); + let mut reordered = RequestMetaObject::new(); + reordered.set_protocol_version(ProtocolVersion::V_2026_07_28); + reordered.set_client_capabilities( + serde_json::from_value(serde_json::json!({ + "com.example/future": {"a": 1, "b": 2}, + "sampling": {} + })) + .expect("reordered open client capability fixture"), + ); + let mut changed = RequestMetaObject::new(); + changed.set_protocol_version(ProtocolVersion::V_2026_07_28); + changed.set_client_capabilities( + serde_json::from_value(serde_json::json!({ + "sampling": {}, + "com.example/future": {"a": 1, "b": 3} + })) + .expect("changed open client capability fixture"), + ); + + assert_eq!(discover_cache_key(&first), discover_cache_key(&reordered)); + assert_ne!(discover_cache_key(&first), discover_cache_key(&changed)); + } + #[tokio::test] async fn discover_serves_a_fresh_cached_response_without_transport_io() { let peer = disconnected_peer(); let meta = RequestMetaObject::default(); - let key = discover_cache_key(); + let key = discover_cache_key(&meta); let expected = DiscoverResult::new(vec![ProtocolVersion::default()], Default::default()) .with_server_info(crate::model::Implementation::from_build_env()) .with_ttl_ms(5_000) From 5b59948c4e7d8508a3a9f3a0b1082c043cd195f6 Mon Sep 17 00:00:00 2001 From: muxammadreza <137672463+muxammadreza@users.noreply.github.com> Date: Thu, 13 Aug 2026 09:46:58 -0700 Subject: [PATCH 3/6] feat: expose contextual subscription acceptance --- crates/rmcp/src/handler/server.rs | 28 ++++++++++++++- crates/rmcp/tests/test_subscriptions.rs | 47 +++++++++++++++++++++++-- 2 files changed, 72 insertions(+), 3 deletions(-) diff --git a/crates/rmcp/src/handler/server.rs b/crates/rmcp/src/handler/server.rs index d6d7aa2fa..18f7841b2 100644 --- a/crates/rmcp/src/handler/server.rs +++ b/crates/rmcp/src/handler/server.rs @@ -149,7 +149,10 @@ impl Service for H { Err(McpError::method_not_found::()) } else { let requested = request.params.notifications; - let Some(candidate) = self.accepted_subscription_filter(&requested) else { + let Some(candidate) = self + .accept_subscription_filter(requested.clone(), context.clone()) + .await? + else { return Err( McpError::method_not_found::(), ); @@ -414,6 +417,21 @@ macro_rules! server_handler_methods { ) -> Option { None } + /// Asynchronously accept a requested notification filter with full request context. + /// + /// Gateways and other policy-aware servers can override this hook when subscription + /// authorization depends on request-scoped client capabilities or other contextual + /// metadata. The default preserves source compatibility by delegating to + /// [`Self::accepted_subscription_filter`]. Returning an error rejects the request + /// before the acknowledgement notification is emitted. + fn accept_subscription_filter( + &self, + requested: SubscriptionFilter, + _context: RequestContext, + ) -> impl Future, McpError>> + MaybeSendFuture + '_ { + let accepted = self.accepted_subscription_filter(&requested); + std::future::ready(Ok(accepted)) + } /// Run one established subscription until it is cancelled or closed gracefully. /// /// The SDK sends the acknowledgment before invoking this method. Returning @@ -693,6 +711,14 @@ macro_rules! impl_server_handler_for_wrapper { (**self).accepted_subscription_filter(requested) } + fn accept_subscription_filter( + &self, + requested: SubscriptionFilter, + context: RequestContext, + ) -> impl Future, McpError>> + MaybeSendFuture + '_ { + (**self).accept_subscription_filter(requested, context) + } + fn listen( &self, context: SubscriptionContext, diff --git a/crates/rmcp/tests/test_subscriptions.rs b/crates/rmcp/tests/test_subscriptions.rs index e7f6eb8a0..ff534918f 100644 --- a/crates/rmcp/tests/test_subscriptions.rs +++ b/crates/rmcp/tests/test_subscriptions.rs @@ -24,8 +24,8 @@ use rmcp::{ SubscriptionsListenResult, }, service::{ - NotificationContext, RequestContext, RoleClient, RoleServer, SubscriptionContext, - SubscriptionEnd, SubscriptionSendError, SubscriptionSink, + NotificationContext, RequestContext, RoleClient, RoleServer, ServiceError, + SubscriptionContext, SubscriptionEnd, SubscriptionSendError, SubscriptionSink, }, }; use tokio::sync::{Mutex, Notify}; @@ -144,6 +144,27 @@ impl ServerHandler for ExtensionFilterServer { } } +struct ContextRejectingFilterServer; + +impl ServerHandler for ContextRejectingFilterServer { + async fn accept_subscription_filter( + &self, + _requested: SubscriptionFilter, + context: RequestContext, + ) -> Result, rmcp::ErrorData> { + assert_eq!( + context.protocol_version(), + Some(ProtocolVersion::V_2026_07_28), + "subscription acceptance must receive the request-scoped protocol context" + ); + Err(rmcp::ErrorData::new( + rmcp::model::ErrorCode(-32003), + "Missing required client capability", + Some(serde_json::json!({"requiredCapabilities": {"extensions": {}}})), + )) + } +} + struct ResourceSubscriptionServer; impl ServerHandler for ResourceSubscriptionServer { @@ -449,6 +470,28 @@ async fn extension_subscription_acknowledges_handler_accepted_subset() -> anyhow Ok(()) } +#[tokio::test] +async fn subscription_acceptance_can_reject_with_request_context_before_acknowledgment() +-> anyhow::Result<()> { + let client = modern_client(ContextRejectingFilterServer).await?; + let mut requested = SubscriptionFilter::new(); + requested + .additional_fields + .insert("taskIds".to_string(), serde_json::json!(["task-a"])); + + let error = client + .listen(requested) + .await + .expect_err("context-aware subscription acceptance should reject before acknowledgment"); + let ServiceError::McpError(error) = error else { + anyhow::bail!("expected MCP error, got {error:?}"); + }; + assert_eq!(error.code, rmcp::model::ErrorCode(-32003)); + + client.cancel().await?; + Ok(()) +} + #[tokio::test] async fn listen_exposes_acknowledged_filter_and_graceful_result() -> anyhow::Result<()> { let client = modern_client(ToolsOnlyServer).await?; From 9078c499b0102def9115b1acaab0b864024960e7 Mon Sep 17 00:00:00 2001 From: muxammadreza <137672463+muxammadreza@users.noreply.github.com> Date: Wed, 12 Aug 2026 21:21:36 -0700 Subject: [PATCH 4/6] fix: preserve explicit discovery client context --- crates/rmcp/src/service/client.rs | 63 ++++++++++++++++++++++++++++--- 1 file changed, 58 insertions(+), 5 deletions(-) diff --git a/crates/rmcp/src/service/client.rs b/crates/rmcp/src/service/client.rs index 053ab6a05..db271e576 100644 --- a/crates/rmcp/src/service/client.rs +++ b/crates/rmcp/src/service/client.rs @@ -1284,11 +1284,17 @@ impl Peer { return Ok(result); } let generation = self.capture_response_cache_generation().await; - let mut request = DiscoverRequest::new(DiscoverRequestParams {}); - request.extensions.insert(meta); - let result = self - .send_request(ClientRequest::DiscoverRequest(request)) - .await; + let request = DiscoverRequest::new(DiscoverRequestParams {}); + let result = match self + .send_request_with_option( + ClientRequest::DiscoverRequest(request), + PeerRequestOptions::no_options().with_meta(meta), + ) + .await + { + Ok(handle) => handle.await_response().await, + Err(error) => Err(error), + }; let result = match result { Ok(result) => result, Err(error) => { @@ -2399,6 +2405,53 @@ mod tests { assert_eq!(peer.list_tools(params).await.unwrap(), expected); } + #[tokio::test] + async fn discover_explicit_client_context_overrides_channel_startup_metadata() { + let (peer, mut outbound) = + Peer::::new(Arc::new(AtomicU32RequestIdProvider::default()), None); + let startup_capabilities = serde_json::from_value(serde_json::json!({ + "extensions": {"com.example/startup": {}} + })) + .expect("startup capabilities"); + peer.set_client_request_metadata(ClientRequestMetadata { + protocol_version: ProtocolVersion::V_2026_07_28, + client_info: crate::model::Implementation::from_build_env(), + client_capabilities: startup_capabilities, + }); + let contextual_capabilities = serde_json::from_value(serde_json::json!({ + "extensions": {"com.example/contextual": {"enabled": true}} + })) + .expect("contextual capabilities"); + let meta = RequestMetaObject::with_client_context( + ProtocolVersion::V_2026_07_28, + crate::model::Implementation::from_build_env(), + contextual_capabilities, + ); + + let discover = tokio::spawn({ + let peer = peer.clone(); + async move { peer.discover(meta).await } + }); + let PeerSinkMessage::Request { request, .. } = + outbound.recv().await.expect("discover request") + else { + panic!("expected discover request"); + }; + let ClientRequest::DiscoverRequest(request) = request else { + panic!("expected server/discover request"); + }; + let capabilities = request + .extensions + .get::() + .expect("request metadata") + .client_capabilities() + .expect("request-scoped client capabilities"); + let raw = serde_json::to_value(capabilities).expect("serialize client capabilities"); + assert!(raw["extensions"].get("com.example/startup").is_none()); + assert_eq!(raw["extensions"]["com.example/contextual"]["enabled"], true); + discover.abort(); + } + #[test] fn discover_cache_key_tracks_result_affecting_request_metadata_canonically() { let mut first = RequestMetaObject::new(); From c60f2481138e0e2dd2f845db437f08fa81d78c46 Mon Sep 17 00:00:00 2001 From: muxammadreza <137672463+muxammadreza@users.noreply.github.com> Date: Thu, 13 Aug 2026 10:32:40 -0700 Subject: [PATCH 5/6] fix: filter task status subscriptions by task id --- crates/rmcp/src/service/client.rs | 13 +++- crates/rmcp/src/service/server.rs | 26 +++++--- crates/rmcp/tests/test_subscriptions.rs | 83 +++++++++++++++++++++++-- 3 files changed, 108 insertions(+), 14 deletions(-) diff --git a/crates/rmcp/src/service/client.rs b/crates/rmcp/src/service/client.rs index db271e576..43f1eb416 100644 --- a/crates/rmcp/src/service/client.rs +++ b/crates/rmcp/src/service/client.rs @@ -557,11 +557,22 @@ impl Subscription { .resource_subscriptions .as_ref() .is_some_and(|uris| uris.contains(&update.params.uri)), + ServerNotification::TaskStatusNotification(notification) => { + let task_id = notification.params.task.task.task_id.as_str(); + self.acknowledged + .additional_fields + .get("taskIds") + .and_then(serde_json::Value::as_array) + .is_some_and(|task_ids| { + task_ids + .iter() + .any(|accepted| accepted.as_str() == Some(task_id)) + }) + } ServerNotification::SubscriptionsAcknowledgedNotification(_) | ServerNotification::CancelledNotification(_) | ServerNotification::ProgressNotification(_) | ServerNotification::LoggingMessageNotification(_) - | ServerNotification::TaskStatusNotification(_) | ServerNotification::CustomNotification(_) => false, } } diff --git a/crates/rmcp/src/service/server.rs b/crates/rmcp/src/service/server.rs index 8efbd34be..1f7793bda 100644 --- a/crates/rmcp/src/service/server.rs +++ b/crates/rmcp/src/service/server.rs @@ -235,15 +235,23 @@ impl SubscriptionSink { "notifications/message", )); } - // SEP-2663 task status notifications are not yet routable through - // `subscriptions/listen`: `SubscriptionFilter` has no `taskIds` - // field yet (the upstream conformance check for this flow is also - // still skipped, pending the subscriptions/listen rewrite). - // Clients currently observe task state by polling `tasks/get`. - ServerNotification::TaskStatusNotification(_) => { - return Err(SubscriptionSendError::UnsupportedNotification( - "notifications/tasks", - )); + ServerNotification::TaskStatusNotification(notification) => { + let task_id = notification.params.task.task.task_id.as_str(); + let accepted = self + .accepted + .additional_fields + .get("taskIds") + .and_then(serde_json::Value::as_array) + .is_some_and(|task_ids| { + task_ids + .iter() + .any(|accepted| accepted.as_str() == Some(task_id)) + }); + if !accepted { + return Err(SubscriptionSendError::NotificationNotAccepted( + "notifications/tasks", + )); + } } ServerNotification::CustomNotification(_) => { return Err(SubscriptionSendError::UnsupportedNotification( diff --git a/crates/rmcp/tests/test_subscriptions.rs b/crates/rmcp/tests/test_subscriptions.rs index ff534918f..b789f4f04 100644 --- a/crates/rmcp/tests/test_subscriptions.rs +++ b/crates/rmcp/tests/test_subscriptions.rs @@ -17,11 +17,12 @@ use std::{ use rmcp::{ ClientHandler, ClientServiceExt, ServerHandler, ServiceExt, model::{ - ClientNotification, ClientRequest, DiscoverResult, GetMeta, Implementation, - NotificationMetaObject, PromptListChangedNotification, ProtocolVersion, ServerCapabilities, - ServerInfo, ServerNotification, ServerResult, SubscriptionFilter, + ClientNotification, ClientRequest, DetailedTask, DiscoverResult, GetMeta, Implementation, + JsonObject, NotificationMetaObject, PromptListChangedNotification, ProtocolVersion, + ServerCapabilities, ServerInfo, ServerNotification, ServerResult, SubscriptionFilter, SubscriptionsAcknowledgedNotification, SubscriptionsAcknowledgedNotificationParams, - SubscriptionsListenResult, + SubscriptionsListenResult, TASKS_EXTENSION_ID, Task, TaskPayload, TaskStatus, + TaskStatusNotification, TaskStatusNotificationParams, }, service::{ NotificationContext, RequestContext, RoleClient, RoleServer, ServiceError, @@ -165,6 +166,58 @@ impl ServerHandler for ContextRejectingFilterServer { } } +struct TaskStatusFilterServer; + +impl TaskStatusFilterServer { + fn notification(task_id: &str) -> ServerNotification { + ServerNotification::TaskStatusNotification(TaskStatusNotification::new( + TaskStatusNotificationParams::new(DetailedTask::new( + Task::new( + task_id, + TaskStatus::Working, + "2026-08-13T16:00:00Z", + "2026-08-13T16:00:00Z", + ), + TaskPayload::Working, + )), + )) + } +} + +impl ServerHandler for TaskStatusFilterServer { + fn get_info(&self) -> ServerInfo { + let mut capabilities = ServerCapabilities::default(); + capabilities.extensions = Some( + [(TASKS_EXTENSION_ID.to_string(), JsonObject::new())] + .into_iter() + .collect(), + ); + ServerInfo::new(capabilities) + } + + fn accepted_subscription_filter( + &self, + requested: &SubscriptionFilter, + ) -> Option { + Some(requested.clone()) + } + + async fn listen(&self, context: SubscriptionContext) -> Result<(), rmcp::ErrorData> { + context + .sink() + .send(Self::notification("task-a")) + .await + .expect("accepted task status notification"); + assert!(matches!( + context.sink().send(Self::notification("task-b")).await, + Err(SubscriptionSendError::NotificationNotAccepted( + "notifications/tasks" + )) + )); + Ok(()) + } +} + struct ResourceSubscriptionServer; impl ServerHandler for ResourceSubscriptionServer { @@ -492,6 +545,28 @@ async fn subscription_acceptance_can_reject_with_request_context_before_acknowle Ok(()) } +#[tokio::test] +async fn task_status_notifications_are_enforced_by_accepted_task_ids() -> anyhow::Result<()> { + let client = modern_client(TaskStatusFilterServer).await?; + let mut requested = SubscriptionFilter::new(); + requested + .additional_fields + .insert("taskIds".to_string(), serde_json::json!(["task-a"])); + let mut subscription = client.listen(requested).await?; + + let notification = tokio::time::timeout(Duration::from_secs(5), subscription.next()) + .await?? + .expect("accepted task status notification"); + let ServerNotification::TaskStatusNotification(notification) = notification else { + anyhow::bail!("expected task status notification"); + }; + assert_eq!(notification.params.task.task.task_id, "task-a"); + assert!(subscription.next().await?.is_none()); + + client.cancel().await?; + Ok(()) +} + #[tokio::test] async fn listen_exposes_acknowledged_filter_and_graceful_result() -> anyhow::Result<()> { let client = modern_client(ToolsOnlyServer).await?; From 3e42416b0927676e908691efa9a24dd0db887668 Mon Sep 17 00:00:00 2001 From: muxammadreza <137672463+muxammadreza@users.noreply.github.com> Date: Thu, 13 Aug 2026 13:15:26 -0700 Subject: [PATCH 6/6] feat: support extension notifications on subscriptions --- crates/rmcp/src/service/client.rs | 42 ++++++++++++++-- crates/rmcp/src/service/server.rs | 50 ++++++++++++++----- crates/rmcp/tests/test_subscriptions.rs | 64 ++++++++++++++++++++++--- 3 files changed, 132 insertions(+), 24 deletions(-) diff --git a/crates/rmcp/src/service/client.rs b/crates/rmcp/src/service/client.rs index 43f1eb416..ba38d67e6 100644 --- a/crates/rmcp/src/service/client.rs +++ b/crates/rmcp/src/service/client.rs @@ -2,7 +2,7 @@ #![expect(deprecated)] pub(super) mod cache; -use std::{borrow::Cow, num::NonZeroUsize, sync::Arc, time::Duration}; +use std::{borrow::Cow, collections::HashSet, num::NonZeroUsize, sync::Arc, time::Duration}; use cache::CacheGeneration; pub use cache::{ClientCacheConfig, MAX_CLIENT_CACHE_TTL}; @@ -376,6 +376,7 @@ pub enum SubscriptionEnd { pub struct Subscription { id: RequestId, acknowledged: SubscriptionFilter, + custom_methods: HashSet, notifications: tokio::sync::mpsc::Receiver, request: Option>, end: Option, @@ -572,8 +573,10 @@ impl Subscription { ServerNotification::SubscriptionsAcknowledgedNotification(_) | ServerNotification::CancelledNotification(_) | ServerNotification::ProgressNotification(_) - | ServerNotification::LoggingMessageNotification(_) - | ServerNotification::CustomNotification(_) => false, + | ServerNotification::LoggingMessageNotification(_) => false, + ServerNotification::CustomNotification(notification) => { + self.custom_methods.contains(notification.method.as_str()) + } } } @@ -1201,6 +1204,29 @@ impl Peer { self.listen_with_channel_capacity_inner( notifications, DEFAULT_SUBSCRIPTION_CHANNEL_CAPACITY, + HashSet::new(), + ) + .await + } + + /// Open a subscription that explicitly permits the listed custom notification methods. + /// + /// RMCP cannot infer extension method ownership from open-world subscription filter fields. + /// Callers must therefore opt in to each extension notification method they have negotiated + /// and validated. Core notification filtering remains unchanged. + pub async fn listen_with_custom_methods( + &self, + notifications: SubscriptionFilter, + custom_methods: I, + ) -> Result + where + I: IntoIterator, + S: Into, + { + self.listen_with_channel_capacity_inner( + notifications, + DEFAULT_SUBSCRIPTION_CHANNEL_CAPACITY, + custom_methods.into_iter().map(Into::into).collect(), ) .await } @@ -1219,14 +1245,19 @@ impl Peer { notifications: SubscriptionFilter, channel_capacity: NonZeroUsize, ) -> Result { - self.listen_with_channel_capacity_inner(notifications, channel_capacity.get()) - .await + self.listen_with_channel_capacity_inner( + notifications, + channel_capacity.get(), + HashSet::new(), + ) + .await } async fn listen_with_channel_capacity_inner( &self, notifications: SubscriptionFilter, channel_capacity: usize, + custom_methods: HashSet, ) -> Result { let request = ClientRequest::SubscriptionsListenRequest(SubscriptionsListenRequest::new( SubscriptionsListenRequestParams::new(notifications.clone()), @@ -1265,6 +1296,7 @@ impl Peer { Ok(Subscription { id, acknowledged: accepted, + custom_methods, notifications: subscription_notifications, request: Some(handle), end: None, diff --git a/crates/rmcp/src/service/server.rs b/crates/rmcp/src/service/server.rs index 1f7793bda..c0b3e5a3f 100644 --- a/crates/rmcp/src/service/server.rs +++ b/crates/rmcp/src/service/server.rs @@ -15,13 +15,14 @@ use crate::{ model::{ CancelledNotification, CancelledNotificationParam, ClientInfo, ClientJsonRpcMessage, ClientNotification, ClientRequest, ClientResult, CreateMessageRequest, - CreateMessageRequestParams, CreateMessageResult, EmptyResult, ErrorData, ListRootsRequest, - ListRootsResult, LoggingMessageNotification, LoggingMessageNotificationParam, - ProgressNotification, ProgressNotificationParam, PromptListChangedNotification, - ProtocolVersion, ResourceListChangedNotification, ResourceUpdatedNotification, - ResourceUpdatedNotificationParam, ServerInfo, ServerNotification, ServerRequest, - ServerResult, SubscriptionFilter, SubscriptionsAcknowledgedNotification, - SubscriptionsAcknowledgedNotificationParams, ToolListChangedNotification, + CreateMessageRequestParams, CreateMessageResult, CustomNotification, EmptyResult, + ErrorData, ListRootsRequest, ListRootsResult, LoggingMessageNotification, + LoggingMessageNotificationParam, ProgressNotification, ProgressNotificationParam, + PromptListChangedNotification, ProtocolVersion, ResourceListChangedNotification, + ResourceUpdatedNotification, ResourceUpdatedNotificationParam, ServerInfo, + ServerNotification, ServerRequest, ServerResult, SubscriptionFilter, + SubscriptionsAcknowledgedNotification, SubscriptionsAcknowledgedNotificationParams, + ToolListChangedNotification, }, transport::DynamicTransportError, }; @@ -144,6 +145,17 @@ pub struct SubscriptionSink { } impl SubscriptionSink { + async fn send_scoped( + &self, + mut notification: ServerNotification, + ) -> Result<(), SubscriptionSendError> { + notification + .get_meta_mut() + .set_subscription_id(self.id.clone()); + self.peer.send_notification(notification).await?; + Ok(()) + } + fn new( peer: Peer, id: RequestId, @@ -176,7 +188,7 @@ impl SubscriptionSink { /// ends, a filter error for disallowed notifications, or a transport error. pub async fn send( &self, - mut notification: ServerNotification, + notification: ServerNotification, ) -> Result<(), SubscriptionSendError> { if self.active.is_cancelled() { return Err(SubscriptionSendError::SubscriptionClosed); @@ -260,11 +272,23 @@ impl SubscriptionSink { } } - notification - .get_meta_mut() - .set_subscription_id(self.id.clone()); - self.peer.send_notification(notification).await?; - Ok(()) + self.send_scoped(notification).await + } + + /// Send a custom extension notification on this subscription stream. + /// + /// RMCP cannot infer open-world extension method ownership from arbitrary + /// subscription filter fields. The caller is responsible for verifying that + /// the method is owned by an extension accepted for this subscription. + pub async fn send_custom_notification( + &self, + notification: CustomNotification, + ) -> Result<(), SubscriptionSendError> { + if self.active.is_cancelled() { + return Err(SubscriptionSendError::SubscriptionClosed); + } + self.send_scoped(ServerNotification::CustomNotification(notification)) + .await } /// Send `notifications/tools/list_changed`. diff --git a/crates/rmcp/tests/test_subscriptions.rs b/crates/rmcp/tests/test_subscriptions.rs index b789f4f04..16fdd8258 100644 --- a/crates/rmcp/tests/test_subscriptions.rs +++ b/crates/rmcp/tests/test_subscriptions.rs @@ -17,12 +17,12 @@ use std::{ use rmcp::{ ClientHandler, ClientServiceExt, ServerHandler, ServiceExt, model::{ - ClientNotification, ClientRequest, DetailedTask, DiscoverResult, GetMeta, Implementation, - JsonObject, NotificationMetaObject, PromptListChangedNotification, ProtocolVersion, - ServerCapabilities, ServerInfo, ServerNotification, ServerResult, SubscriptionFilter, - SubscriptionsAcknowledgedNotification, SubscriptionsAcknowledgedNotificationParams, - SubscriptionsListenResult, TASKS_EXTENSION_ID, Task, TaskPayload, TaskStatus, - TaskStatusNotification, TaskStatusNotificationParams, + ClientNotification, ClientRequest, CustomNotification, DetailedTask, DiscoverResult, + GetMeta, Implementation, JsonObject, NotificationMetaObject, PromptListChangedNotification, + ProtocolVersion, ServerCapabilities, ServerInfo, ServerNotification, ServerResult, + SubscriptionFilter, SubscriptionsAcknowledgedNotification, + SubscriptionsAcknowledgedNotificationParams, SubscriptionsListenResult, TASKS_EXTENSION_ID, + Task, TaskPayload, TaskStatus, TaskStatusNotification, TaskStatusNotificationParams, }, service::{ NotificationContext, RequestContext, RoleClient, RoleServer, ServiceError, @@ -33,6 +33,30 @@ use tokio::sync::{Mutex, Notify}; struct ToolsOnlyServer; +struct ExtensionNotificationServer; + +impl ServerHandler for ExtensionNotificationServer { + fn accepted_subscription_filter( + &self, + requested: &SubscriptionFilter, + ) -> Option { + Some(requested.clone()) + } + + async fn listen(&self, context: SubscriptionContext) -> Result<(), rmcp::ErrorData> { + context + .sink() + .send_custom_notification(CustomNotification::new( + "com.example/changed", + Some(serde_json::json!({ "revision": 1 })), + )) + .await + .expect("send explicitly authorized extension notification"); + context.cancelled().await; + Ok(()) + } +} + #[derive(Clone)] struct CountingClient { tool_changes: Arc, @@ -523,6 +547,34 @@ async fn extension_subscription_acknowledges_handler_accepted_subset() -> anyhow Ok(()) } +#[tokio::test] +async fn custom_subscription_notifications_require_an_explicit_method_allowlist() +-> anyhow::Result<()> { + let client = modern_client(ExtensionNotificationServer).await?; + let mut requested = SubscriptionFilter::new(); + requested.additional_fields.insert( + "com.example/channels".to_string(), + serde_json::json!(["alpha"]), + ); + + let mut subscription = client + .listen_with_custom_methods(requested, ["com.example/changed"]) + .await?; + let Some(ServerNotification::CustomNotification(notification)) = subscription.next().await? + else { + panic!("expected explicitly allowed custom subscription notification"); + }; + assert_eq!(notification.method, "com.example/changed"); + assert_eq!( + notification.params, + Some(serde_json::json!({ "revision": 1 })) + ); + + subscription.cancel().await?; + client.cancel().await?; + Ok(()) +} + #[tokio::test] async fn subscription_acceptance_can_reject_with_request_context_before_acknowledgment() -> anyhow::Result<()> {