Skip to main content

tower_mcp/transport/
subscriptions.rs

1//! Shared transport helpers for final-protocol subscriptions.
2
3use std::time::Duration;
4
5use crate::context::ServerNotification;
6use crate::protocol::{
7    Implementation, JsonRpcNotification, JsonRpcResponse, NotificationMeta, RequestId, ResultType,
8    SubscriptionFilter, SubscriptionsAcknowledgedParams, SubscriptionsListenResult,
9    SubscriptionsListenResultMeta, notifications,
10};
11
12/// Narrow a requested filter to what this server will actually honor.
13///
14/// The acknowledgement reports this filter back to the client, so anything
15/// dropped here is a promise the server declines to make. `tasks_enabled`
16/// reflects whether the server opted into the Tasks extension: without it
17/// there is nothing to notify about, so the task IDs are dropped rather than
18/// acknowledged and then silently ignored.
19pub(crate) fn accepted_subscription_filter(
20    requested: SubscriptionFilter,
21    tasks_enabled: bool,
22) -> SubscriptionFilter {
23    SubscriptionFilter {
24        tools_list_changed: requested.tools_list_changed.filter(|enabled| *enabled),
25        prompts_list_changed: requested.prompts_list_changed.filter(|enabled| *enabled),
26        resources_list_changed: requested.resources_list_changed.filter(|enabled| *enabled),
27        resource_subscriptions: requested.resource_subscriptions,
28        task_ids: requested.task_ids.filter(|_| tasks_enabled),
29    }
30}
31
32pub(crate) fn subscription_matches(
33    notification: &ServerNotification,
34    filter: &SubscriptionFilter,
35) -> bool {
36    match notification {
37        ServerNotification::ToolsListChanged => filter.tools_list_changed == Some(true),
38        ServerNotification::PromptsListChanged => filter.prompts_list_changed == Some(true),
39        ServerNotification::ResourcesListChanged => filter.resources_list_changed == Some(true),
40        ServerNotification::ResourceUpdated { uri } => filter
41            .resource_subscriptions
42            .as_ref()
43            .is_some_and(|subscriptions| subscriptions.iter().any(|item| item == uri)),
44        // A task is named individually rather than opted into as a class, so
45        // an unlisted task ID never matches even a subscriber that asked for
46        // every other notification type.
47        ServerNotification::FinalTaskStatusChanged(params) => filter
48            .task_ids
49            .as_ref()
50            .is_some_and(|ids| ids.iter().any(|id| id == params.task.task_id())),
51        _ => false,
52    }
53}
54
55pub(crate) fn tagged_subscription_notification(
56    notification: &ServerNotification,
57    subscription_id: &RequestId,
58) -> Option<String> {
59    let json = crate::transport::stdio::serialize_notification(notification)?;
60    let mut value: serde_json::Value = serde_json::from_str(&json).ok()?;
61    let object = value.as_object_mut()?;
62    let params = object
63        .entry("params")
64        .or_insert_with(|| serde_json::json!({}))
65        .as_object_mut()?;
66    let meta = params
67        .entry("_meta")
68        .or_insert_with(|| serde_json::json!({}))
69        .as_object_mut()?;
70    meta.insert(
71        "io.modelcontextprotocol/subscriptionId".to_string(),
72        serde_json::to_value(subscription_id).ok()?,
73    );
74    serde_json::to_string(&value).ok()
75}
76
77pub(crate) fn subscription_acknowledgment(
78    subscription_id: RequestId,
79    notifications: SubscriptionFilter,
80) -> JsonRpcNotification {
81    JsonRpcNotification::new(notifications::SUBSCRIPTIONS_ACKNOWLEDGED).with_params(
82        serde_json::to_value(SubscriptionsAcknowledgedParams {
83            meta: Some(NotificationMeta {
84                subscription_id: Some(subscription_id),
85            }),
86            notifications,
87        })
88        .expect("subscription acknowledgment is serializable"),
89    )
90}
91
92/// Why a `subscriptions/listen` stream ended.
93#[derive(Debug, Clone, Copy, PartialEq, Eq)]
94#[non_exhaustive]
95pub enum SubscriptionCloseReason {
96    /// The client cancelled the subscription with `notifications/cancelled`.
97    Cancelled,
98    /// The client connection or response stream dropped.
99    Disconnected,
100    /// The server drained the stream gracefully (shutdown or an explicit
101    /// close), sending the terminal `SubscriptionsListenResult` first where
102    /// the transport still could.
103    Drained,
104}
105
106/// Terminal record of one `subscriptions/listen` stream.
107#[derive(Debug, Clone)]
108#[non_exhaustive]
109pub struct SubscriptionClose {
110    /// JSON-RPC id of the `subscriptions/listen` request that opened the
111    /// stream, the same id middleware observed at acceptance.
112    pub subscription_id: RequestId,
113    /// Why the stream ended.
114    pub reason: SubscriptionCloseReason,
115    /// Time from acceptance to close.
116    pub duration: Duration,
117}
118
119/// Observes the terminal half of `subscriptions/listen` streams.
120///
121/// The request half of the observation boundary is ordinary
122/// `Service<RouterRequest>` middleware: transports dispatch the listen
123/// request through the service before upgrading, so a layer sees acceptance
124/// and rejection like any other method. What that boundary cannot express is
125/// the stream's end, which happens long after the service call returns. This
126/// hook carries exactly that remainder and nothing else: implement it for
127/// ledger or audit records that need terminal reason and duration, and pair
128/// it with a layer for the request half.
129///
130/// Attach with [`McpRouter::with_subscription_observer`]; every transport
131/// that owns listen streams (stdio, generic stdio, bidirectional stdio,
132/// channel, HTTP) reports through it. Calls are made from transport
133/// internals, so implementations must be fast and non-blocking.
134///
135/// [`McpRouter::with_subscription_observer`]: crate::McpRouter::with_subscription_observer
136pub trait SubscriptionObserver: Send + Sync {
137    /// A `subscriptions/listen` stream reached its end.
138    fn on_close(&self, close: SubscriptionClose);
139}
140
141pub(crate) fn subscription_complete_response(
142    subscription_id: RequestId,
143    server_info: Option<Implementation>,
144) -> JsonRpcResponse {
145    let result = SubscriptionsListenResult {
146        result_type: ResultType::Complete,
147        meta: SubscriptionsListenResultMeta {
148            subscription_id: subscription_id.clone(),
149            server_info,
150        },
151    };
152    JsonRpcResponse::result(
153        subscription_id,
154        serde_json::to_value(result).expect("subscription result is serializable"),
155    )
156}
157
158#[cfg(test)]
159mod tests {
160    use super::*;
161    use crate::tasks::{DetailedTask, TaskMetadata, TaskStatusNotificationParams};
162
163    fn task_notification(task_id: &str) -> ServerNotification {
164        ServerNotification::FinalTaskStatusChanged(TaskStatusNotificationParams {
165            task: DetailedTask::working(TaskMetadata::new(
166                task_id.to_string(),
167                "2026-07-28T00:00:00Z".to_string(),
168                "2026-07-28T00:00:01Z".to_string(),
169                Some(60_000),
170            )),
171            meta: None,
172        })
173    }
174
175    fn subscribed_to(task_ids: &[&str]) -> SubscriptionFilter {
176        SubscriptionFilter {
177            task_ids: Some(task_ids.iter().map(|id| id.to_string()).collect()),
178            ..SubscriptionFilter::default()
179        }
180    }
181
182    #[test]
183    fn task_notifications_match_only_the_named_task_ids() {
184        let filter = subscribed_to(&["task-a", "task-b"]);
185        assert!(subscription_matches(&task_notification("task-a"), &filter));
186        assert!(subscription_matches(&task_notification("task-b"), &filter));
187        assert!(!subscription_matches(&task_notification("task-c"), &filter));
188    }
189
190    #[test]
191    fn a_broad_subscription_still_excludes_unnamed_tasks() {
192        // Tasks are named individually rather than opted into as a class, so
193        // asking for every other notification type grants nothing here.
194        let filter = SubscriptionFilter {
195            tools_list_changed: Some(true),
196            prompts_list_changed: Some(true),
197            resources_list_changed: Some(true),
198            resource_subscriptions: Some(vec!["file:///everything".to_string()]),
199            task_ids: None,
200        };
201        assert!(!subscription_matches(&task_notification("task-a"), &filter));
202    }
203
204    #[test]
205    fn accepted_filter_declines_task_ids_when_the_server_has_no_tasks() {
206        let requested = subscribed_to(&["task-a"]);
207
208        let accepted = accepted_subscription_filter(requested.clone(), true);
209        assert_eq!(
210            accepted.task_ids.as_deref(),
211            Some(&["task-a".to_string()][..])
212        );
213
214        // The acknowledgement reports what the server agreed to honor, so a
215        // server without the extension must not echo the IDs back.
216        let declined = accepted_subscription_filter(requested, false);
217        assert!(declined.task_ids.is_none());
218    }
219
220    #[test]
221    fn task_notifications_serialize_as_notifications_tasks() {
222        let json = crate::transport::stdio::serialize_notification(&task_notification("task-a"))
223            .expect("task notification is serializable");
224        let value: serde_json::Value = serde_json::from_str(&json).unwrap();
225        assert_eq!(value["method"], "notifications/tasks");
226        // The DetailedTask is flattened into params rather than nested.
227        assert_eq!(value["params"]["taskId"], "task-a");
228        assert_eq!(value["params"]["status"], "working");
229        assert_eq!(value["params"]["ttlMs"], 60_000);
230    }
231}