Add MCP event discovery and subscriptions (#37494)

## What changed - Expose hosted Plugin Runtime event definitions through `McpResourceClient::list_events`. - Add cancellable `events/stream` subscriptions that route lifecycle notifications to the matching request and cancel when the stream is dropped. - Bound event notification and queue sizes, time out stalled response headers, and close streamable HTTP responses locally on cancellation. ## Testing - Add streamable HTTP integration coverage for isolated concurrent subscriptions, local cancellation, and stalled response-header timeouts. GitOrigin-RevId: 98c6fffa8f004eb2638c1b0705f659222cbbdfdb
2026-11-13 14:46:14 +00:00 · 2026-08-07 19:36:52 +00:00
1 parent 2b1811e562
commit 41014b11bd
@@ -0,0 +1,263 @@
+use std::collections::HashMap;
+use std::future::Future;
+use std::sync::Arc;
+use std::sync::Mutex;
+use std::sync::PoisonError;
+
+use rmcp::model::ClientRequest;
+use rmcp::model::CustomNotification;
+use rmcp::model::JsonRpcMessage;
+use rmcp::model::RequestId;
+use rmcp::model::ServerNotification;
+use rmcp::service::RoleClient;
+use rmcp::service::RxJsonRpcMessage;
+use rmcp::service::TxJsonRpcMessage;
+use rmcp::transport::IntoTransport;
+use rmcp::transport::Transport;
+use tokio::sync::OwnedSemaphorePermit;
+use tokio::sync::Semaphore;
+use tokio::sync::mpsc;
+use tracing::warn;
+
+pub(crate) const MAX_EVENT_NOTIFICATION_BYTES: usize = 1024 * 1024;
+const MAX_QUEUED_EVENT_NOTIFICATION_BYTES: usize = 2 * 1024 * 1024;
+
+#[derive(Clone)]
+pub(crate) struct EventNotificationSender {
+    notifications: mpsc::UnboundedSender<QueuedEventNotification>,
+    available_bytes: Arc<Semaphore>,
+}
+
+pub struct EventNotificationReceiver {
+    notifications: mpsc::UnboundedReceiver<QueuedEventNotification>,
+    available_bytes: Arc<Semaphore>,
+}
+
+struct QueuedEventNotification {
+    notification: CustomNotification,
+    _bytes: OwnedSemaphorePermit,
+}
+
+pub(crate) fn event_notification_channel() -> (EventNotificationSender, EventNotificationReceiver) {
+    let (notifications_tx, notifications_rx) = mpsc::unbounded_channel();
+    let available_bytes = Arc::new(Semaphore::new(MAX_QUEUED_EVENT_NOTIFICATION_BYTES));
+
+    (
+        EventNotificationSender {
+            notifications: notifications_tx,
+            available_bytes: Arc::clone(&available_bytes),
+        },
+        EventNotificationReceiver {
+            notifications: notifications_rx,
+            available_bytes,
+        },
+    )
+}
+
+impl EventNotificationSender {
+    async fn send(
+        &self,
+        notification: CustomNotification,
+        notification_bytes: usize,
+    ) -> Result<(), ()> {
+        let bytes = u32::try_from(notification_bytes).map_err(|_| ())?;
+        let permit = Arc::clone(&self.available_bytes)
+            .acquire_many_owned(bytes)
+            .await
+            .map_err(|_| ())?;
+        self.notifications
+            .send(QueuedEventNotification {
+                notification,
+                _bytes: permit,
+            })
+            .map_err(|_| ())
+    }
+
+    fn close(&self) {
+        self.available_bytes.close();
+    }
+}
+
+impl EventNotificationReceiver {
+    pub async fn recv(&mut self) -> Option<CustomNotification> {
+        self.notifications
+            .recv()
+            .await
+            .map(|queued| queued.notification)
+    }
+}
+
+impl Drop for EventNotificationReceiver {
+    fn drop(&mut self) {
+        self.available_bytes.close();
+    }
+}
+
+/// Consumes Plugin Runtime event notifications before rmcp can log their payloads.
+pub(crate) fn capture_event_notifications<T, E, A>(
+    transport: T,
+) -> impl Transport<RoleClient, Error = E> + 'static
+where
+    T: IntoTransport<RoleClient, E, A>,
+    E: std::error::Error + Send + Sync + 'static,
+{
+    EventNotificationTransport {
+        inner: transport.into_transport(),
+        routes: Arc::default(),
+    }
+}
+
+struct EventNotificationTransport<T> {
+    inner: T,
+    routes: Arc<Mutex<HashMap<RequestId, EventNotificationSender>>>,
+}
+
+impl<T> Transport<RoleClient> for EventNotificationTransport<T>
+where
+    T: Transport<RoleClient> + 'static,
+{
+    type Error = T::Error;
+
+    fn send(
+        &mut self,
+        mut message: TxJsonRpcMessage<RoleClient>,
+    ) -> impl Future<Output = Result<(), Self::Error>> + Send + 'static {
+        let request_route = match &mut message {
+            JsonRpcMessage::Request(envelope) => match &mut envelope.request {
+                ClientRequest::CustomRequest(request) if request.method == "events/stream" => {
+                    request
+                        .extensions
+                        .remove::<EventNotificationSender>()
+                        .map(|sender| (envelope.id.clone(), sender))
+                }
+                _ => None,
+            },
+            JsonRpcMessage::Notification(envelope) => {
+                if let rmcp::model::ClientNotification::CancelledNotification(cancelled) =
+                    &envelope.notification
+                    && let Some(request_id) = cancelled.params.request_id.as_ref()
+                    && let Some(route) = self
+                        .routes
+                        .lock()
+                        .unwrap_or_else(PoisonError::into_inner)
+                        .remove(request_id)
+                {
+                    route.close();
+                }
+                None
+            }
+            JsonRpcMessage::Response(_) | JsonRpcMessage::Error(_) => None,
+        };
+
+        let route_id = request_route
+            .as_ref()
+            .map(|(request_id, _)| request_id.clone());
+        if let Some((request_id, sender)) = request_route
+            && let Some(previous) = self
+                .routes
+                .lock()
+                .unwrap_or_else(PoisonError::into_inner)
+                .insert(request_id, sender)
+        {
+            previous.close();
+        }
+
+        let routes = Arc::clone(&self.routes);
+        let send = self.inner.send(message);
+        async move {
+            let result = send.await;
+            if result.is_err()
+                && let Some(route_id) = route_id
+                && let Some(route) = routes
+                    .lock()
+                    .unwrap_or_else(PoisonError::into_inner)
+                    .remove(&route_id)
+            {
+                route.close();
+            }
+            result
+        }
+    }
+
+    async fn receive(&mut self) -> Option<RxJsonRpcMessage<RoleClient>> {
+        loop {
+            let message = self.inner.receive().await?;
+            if let Some(request_id) = response_id(&message) {
+                if let Some(route) = self
+                    .routes
+                    .lock()
+                    .unwrap_or_else(PoisonError::into_inner)
+                    .remove(request_id)
+                {
+                    route.close();
+                }
+                return Some(message);
+            }
+
+            let JsonRpcMessage::Notification(envelope) = &message else {
+                return Some(message);
+            };
+            let ServerNotification::CustomNotification(notification) = &envelope.notification
+            else {
+                return Some(message);
+            };
+            if !notification.method.starts_with("notifications/events/") {
+                return Some(message);
+            }
+            let Some(subscription_id) =
+                rmcp::model::GetMeta::get_meta(notification).subscription_id()
+            else {
+                continue;
+            };
+
+            let route = self
+                .routes
+                .lock()
+                .unwrap_or_else(PoisonError::into_inner)
+                .get(&subscription_id)
+                .cloned();
+            let Some(route) = route else {
+                continue;
+            };
+            let notification_bytes = serde_json::to_vec(&message)
+                .map(|message| message.len())
+                .unwrap_or(usize::MAX);
+            if notification_bytes > MAX_EVENT_NOTIFICATION_BYTES {
+                warn!(
+                    notification_bytes,
+                    "discarding oversized MCP event notification"
+                );
+                continue;
+            }
+            if route
+                .send(notification.clone(), notification_bytes)
+                .await
+                .is_err()
+                && let Some(route) = self
+                    .routes
+                    .lock()
+                    .unwrap_or_else(PoisonError::into_inner)
+                    .remove(&subscription_id)
+            {
+                route.close();
+            }
+        }
+    }
+
+    async fn close(&mut self) -> Result<(), Self::Error> {
+        let routes =
+            std::mem::take(&mut *self.routes.lock().unwrap_or_else(PoisonError::into_inner));
+        for route in routes.into_values() {
+            route.close();
+        }
+        self.inner.close().await
+    }
+}
+
+fn response_id(message: &RxJsonRpcMessage<RoleClient>) -> Option<&RequestId> {
+    match message {
+        JsonRpcMessage::Response(response) => Some(&response.id),
+        JsonRpcMessage::Error(error) => error.id.as_ref(),
+        JsonRpcMessage::Request(_) | JsonRpcMessage::Notification(_) => None,
+    }
+}
@@ -12,6 +12,7 @@ use std::io;
 use std::sync::Arc;
 use std::sync::Mutex;
 use std::sync::PoisonError;
+use std::time::Duration;
 use std::time::Instant;

 use bytes::Bytes;
@@ -41,6 +42,7 @@ use rmcp::model::ErrorCode;
 use rmcp::model::ErrorData;
 use rmcp::model::JsonRpcMessage;
 use rmcp::model::ProtocolVersion;
+use rmcp::model::RequestId;
 use rmcp::model::ServerJsonRpcMessage;
 use rmcp::model::ServerResult;
 use rmcp::transport::common::http_header::HEADER_MCP_PROTOCOL_VERSION;
@@ -51,7 +53,9 @@ use rmcp::transport::streamable_http_client::StreamableHttpError;
 use rmcp::transport::streamable_http_client::StreamableHttpPostResponse;
 use sse_stream::Sse;
 use sse_stream::SseStream;
+use tokio::sync::oneshot;

+use crate::event_notification_transport::MAX_EVENT_NOTIFICATION_BYTES;
 use crate::incoming_jsonrpc::deserialize_incoming_jsonrpc_message;
 use crate::incoming_jsonrpc::normalize_sse_jsonrpc_message;
 use crate::local_stdio_transport::MAX_MCP_STDIO_LINE_BYTES;
@@ -65,6 +69,7 @@ const JSON_MIME_TYPE: &str = "application/json";
 const HEADER_SESSION_ID: &str = "Mcp-Session-Id";
 const NON_JSON_RESPONSE_BODY_PREVIEW_BYTES: usize = 8_192;
 const LEGACY_HTTP_PREVALIDATION_ERROR_CODE: ErrorCode = ErrorCode(-32000);
+const EVENT_STREAM_RESPONSE_TIMEOUT: Duration = Duration::from_secs(30);

 #[derive(Clone, Copy, Debug, PartialEq, Eq)]
 pub enum StreamableHttpRedirectMode {
@@ -77,11 +82,26 @@ pub(crate) struct StreamableHttpClientAdapter {
    http_client: Arc<dyn HttpClient>,
    default_headers: HeaderMap,
    auth_provider: Option<SharedAuthProvider>,
+    event_stream_cancellations: Arc<Mutex<HashMap<RequestId, oneshot::Sender<()>>>>,
    has_configured_headers: bool,
    redirect_mode: StreamableHttpRedirectMode,
    initialize_deadline: Arc<Mutex<Option<Instant>>>,
 }

+struct EventStreamCancellation {
+    request_id: RequestId,
+    cancellations: Arc<Mutex<HashMap<RequestId, oneshot::Sender<()>>>>,
+}
+
+impl Drop for EventStreamCancellation {
+    fn drop(&mut self) {
+        self.cancellations
+            .lock()
+            .unwrap_or_else(PoisonError::into_inner)
+            .remove(&self.request_id);
+    }
+}
+
 #[derive(Debug, thiserror::Error)]
 pub(crate) enum StreamableHttpClientAdapterError {
    #[error("streamable HTTP session expired with 404 Not Found")]
@@ -107,6 +127,7 @@ impl StreamableHttpClientAdapter {
            http_client,
            default_headers,
            auth_provider,
+            event_stream_cancellations: Arc::default(),
            has_configured_headers,
            redirect_mode,
            initialize_deadline,
@@ -163,13 +184,18 @@ impl StreamableHttpClient for StreamableHttpClientAdapter {
            )?;
        }

-        let maximum_response_bytes = (mcp_method.as_deref() == Some(DiscoverRequestMethod::VALUE)
-            || headers
-                .get(HEADER_MCP_PROTOCOL_VERSION)
-                .and_then(|value| value.to_str().ok())
-                == Some(ProtocolVersion::V_2026_07_28.as_str()))
-        .then_some(MAX_MCP_STDIO_LINE_BYTES);
-        let redirect_policy = if mcp_method.as_deref() == Some(DiscoverRequestMethod::VALUE) {
+        let is_discovery_request = mcp_method.as_deref() == Some(DiscoverRequestMethod::VALUE);
+        let is_event_stream_request = mcp_method.as_deref() == Some("events/stream");
+        let uses_modern_protocol = headers
+            .get(HEADER_MCP_PROTOCOL_VERSION)
+            .and_then(|value| value.to_str().ok())
+            == Some(ProtocolVersion::V_2026_07_28.as_str());
+        let maximum_response_bytes = if is_event_stream_request {
+            Some(MAX_EVENT_NOTIFICATION_BYTES)
+        } else {
+            (is_discovery_request || uses_modern_protocol).then_some(MAX_MCP_STDIO_LINE_BYTES)
+        };
+        let redirect_policy = if is_discovery_request {
            HttpRedirectPolicy::Stop
        } else {
            self.redirect_policy(&headers)
@@ -197,19 +223,40 @@ impl StreamableHttpClient for StreamableHttpClientAdapter {

        let body = serde_json::to_vec(&message).map_err(StreamableHttpError::Deserialize)?;
        let has_authorization_header = headers.contains_key(AUTHORIZATION);
-        let response = self
-            .http_client
-            .http_request_stream(HttpRequestParams {
-                method: "POST".to_string(),
-                url: uri.to_string(),
-                headers: protocol_headers(&headers),
-                body: Some(body.into()),
-                timeout_ms,
-                redirect_policy,
-                request_id: "buffered-request".to_string(),
-                stream_response: true,
-            })
-            .await;
+        if let JsonRpcMessage::Notification(notification) = &message
+            && let ClientNotification::CancelledNotification(cancelled) = &notification.notification
+            && let Some(request_id) = cancelled.params.request_id.as_ref()
+            && let Some(cancellation) = self
+                .event_stream_cancellations
+                .lock()
+                .unwrap_or_else(PoisonError::into_inner)
+                .remove(request_id)
+        {
+            let _ = cancellation.send(());
+            return Ok(StreamableHttpPostResponse::Accepted);
+        }
+
+        let request = self.http_client.http_request_stream(HttpRequestParams {
+            method: "POST".to_string(),
+            url: uri.to_string(),
+            headers: protocol_headers(&headers),
+            body: Some(body.into()),
+            timeout_ms,
+            redirect_policy,
+            request_id: "buffered-request".to_string(),
+            stream_response: true,
+        });
+        let response = if is_event_stream_request {
+            tokio::time::timeout(EVENT_STREAM_RESPONSE_TIMEOUT, request)
+                .await
+                .map_err(|_| {
+                    StreamableHttpError::UnexpectedServerResponse(
+                        "timed out waiting for MCP event stream response headers".into(),
+                    )
+                })?
+        } else {
+            request.await
+        };
        let (response, mut body_stream) = match response {
            Ok(response) => response,
            Err(error) => {
@@ -309,7 +356,11 @@ impl StreamableHttpClient for StreamableHttpClientAdapter {
        }
        match content_type.as_deref() {
            Some(content_type) if content_type.starts_with(EVENT_STREAM_MIME_TYPE) => {
-                let mut event_stream = sse_stream_from_body(body_stream, maximum_response_bytes);
+                let mut event_stream = sse_stream_from_body(
+                    body_stream,
+                    maximum_response_bytes,
+                    is_discovery_request || uses_modern_protocol,
+                );
                if mcp_method.as_deref() == Some(DiscoverRequestMethod::VALUE) {
                    while let Some(event) = event_stream.next().await {
                        let event = event.map_err(StreamableHttpError::Sse)?;
@@ -340,6 +391,35 @@ impl StreamableHttpClient for StreamableHttpClientAdapter {
                        "empty sse stream".into(),
                    ));
                }
+                if is_event_stream_request && let JsonRpcMessage::Request(request) = &message {
+                    let (cancel, cancelled) = oneshot::channel();
+                    let cancellation = EventStreamCancellation {
+                        request_id: request.id.clone(),
+                        cancellations: Arc::clone(&self.event_stream_cancellations),
+                    };
+                    cancellation
+                        .cancellations
+                        .lock()
+                        .unwrap_or_else(PoisonError::into_inner)
+                        .insert(cancellation.request_id.clone(), cancel);
+
+                    event_stream = stream::unfold(
+                        Some((event_stream, cancelled, cancellation)),
+                        |state| async move {
+                            let (mut event_stream, mut cancelled, cancellation) = state?;
+
+                            tokio::select! {
+                                biased;
+
+                                _ = &mut cancelled => None,
+                                event = event_stream.next() => event.map(|event| {
+                                    (event, Some((event_stream, cancelled, cancellation)))
+                                }),
+                            }
+                        },
+                    )
+                    .boxed();
+                }
                Ok(StreamableHttpPostResponse::Sse(event_stream, session_id))
            }
            Some(content_type) if content_type.starts_with(JSON_MIME_TYPE) => {
@@ -507,12 +587,16 @@ impl StreamableHttpClient for StreamableHttpClientAdapter {
            }
        }

-        let maximum_response_bytes = headers
+        let uses_modern_protocol = headers
            .get(HEADER_MCP_PROTOCOL_VERSION)
            .and_then(|value| value.to_str().ok())
-            .is_some_and(|version| version == ProtocolVersion::V_2026_07_28.as_str())
-            .then_some(MAX_MCP_STDIO_LINE_BYTES);
-        Ok(sse_stream_from_body(body_stream, maximum_response_bytes))
+            .is_some_and(|version| version == ProtocolVersion::V_2026_07_28.as_str());
+        let maximum_response_bytes = uses_modern_protocol.then_some(MAX_MCP_STDIO_LINE_BYTES);
+        Ok(sse_stream_from_body(
+            body_stream,
+            maximum_response_bytes,
+            uses_modern_protocol,
+        ))
    }
 }

@@ -863,8 +947,8 @@ async fn collect_body(
 fn sse_stream_from_body(
    body_stream: HttpResponseBodyStream,
    maximum_event_bytes: Option<usize>,
+    modern_session: bool,
 ) -> BoxStream<'static, std::result::Result<Sse, sse_stream::Error>> {
-    let modern_session = maximum_event_bytes.is_some();
    SseStream::from_bytes_stream(stream::unfold(
        (body_stream, SseEventSizeLimit::new(maximum_event_bytes)),
        |(mut body_stream, mut size_limit)| async move {
@@ -1,5 +1,6 @@
 mod auth_status;
 mod elicitation_client_service;
+mod event_notification_transport;
 mod executor_process_transport;
 mod http_client_adapter;
 mod in_process_transport;
@@ -24,6 +25,7 @@ pub use auth_status::determine_streamable_http_auth_status;
 pub use auth_status::determine_streamable_http_auth_status_from_credentials;
 pub use auth_status::discover_streamable_http_oauth;
 pub use codex_protocol::protocol::McpAuthStatus;
+pub use event_notification_transport::EventNotificationReceiver;
 pub use http_client_adapter::StreamableHttpRedirectMode;
 pub use in_process_transport::InProcessTransportFactory;
 pub use oauth::StoredOAuthCredentialSnapshot;
@@ -40,6 +42,7 @@ pub use perform_oauth_login::perform_oauth_login_return_url;
 pub use perform_oauth_login::perform_oauth_login_silent;
 pub use protocol_mode::McpProtocolMode;
 pub use rmcp::model::ElicitationAction;
+pub use rmcp_client::CancellableEventStreamRequest;
 pub use rmcp_client::Elicitation;
 pub use rmcp_client::ElicitationResponse;
 pub use rmcp_client::ListToolsWithConnectorIdResult;
@@ -50,6 +50,7 @@ use rmcp::model::ServerResult;
 use rmcp::model::Tool;
 use rmcp::service::ClientCacheConfig;
 use rmcp::service::ClientServiceExt;
+use rmcp::service::RequestHandle;
 use rmcp::service::RoleClient;
 use rmcp::service::RunningService;
 use rmcp::transport::AuthorizationManager;
@@ -70,6 +71,8 @@ use tracing::instrument;
 use tracing::warn;

 use crate::elicitation_client_service::ElicitationClientService;
+use crate::event_notification_transport::capture_event_notifications;
+use crate::event_notification_transport::event_notification_channel;
 use crate::http_client_adapter::StreamableHttpClientAdapter;
 use crate::http_client_adapter::StreamableHttpClientAdapterError;
 use crate::http_client_adapter::StreamableHttpRedirectMode;
@@ -344,6 +347,12 @@ pub struct ListToolsWithConnectorIdResult {
    pub tools: Vec<ToolWithConnectorId>,
 }

+/// An active Plugin Runtime event request and its request-scoped notifications.
+pub struct CancellableEventStreamRequest {
+    pub handle: RequestHandle<RoleClient>,
+    pub notifications: crate::EventNotificationReceiver,
+}
+
 /// MCP client implemented on top of the official `rmcp` SDK.
 /// https://github.com/modelcontextprotocol/rust-sdk
 pub struct RmcpClient {
@@ -845,10 +854,20 @@ impl RmcpClient {
        &self,
        method: &str,
        params: Option<serde_json::Value>,
+    ) -> Result<ServerResult> {
+        self.send_custom_request_with_timeout(method, params, /*timeout*/ None)
+            .await
+    }
+
+    pub async fn send_custom_request_with_timeout(
+        &self,
+        method: &str,
+        params: Option<serde_json::Value>,
+        timeout: Option<Duration>,
    ) -> Result<ServerResult> {
        self.refresh_oauth_if_needed().await?;
        let response = self
-            .run_service_operation("requests/custom", /*timeout*/ None, move |service| {
+            .run_service_operation("requests/custom", timeout, move |service| {
                let params = params.clone();
                async move {
                    service
@@ -864,6 +883,29 @@ impl RmcpClient {
        Ok(response)
    }

+    /// Starts a Plugin Runtime event stream without waiting for its final response.
+    pub async fn send_event_stream_request(
+        &self,
+        params: Option<serde_json::Value>,
+    ) -> Result<CancellableEventStreamRequest> {
+        let service = self.service().await?;
+        let (sender, notifications) = event_notification_channel();
+        let mut request = CustomRequest::new("events/stream", params);
+        request.extensions.insert(sender);
+        let handle = service
+            .peer()
+            .send_cancellable_request(
+                ClientRequest::CustomRequest(request),
+                rmcp::service::PeerRequestOptions::no_options(),
+            )
+            .await?;
+
+        Ok(CancellableEventStreamRequest {
+            handle,
+            notifications,
+        })
+    }
+
    async fn service(&self) -> Result<Arc<RunningService<RoleClient, ElicitationClientService>>> {
        let guard = self.state.lock().await;
        match &*guard {
@@ -1139,7 +1181,7 @@ impl RmcpClient {
            ),
            PendingTransport::StreamableHttp { transport } => (
                client_service
-                    .serve_with_lifecycle(transport, lifecycle)
+                    .serve_with_lifecycle(capture_event_notifications(transport), lifecycle)
                    .boxed(),
                None,
            ),