diff --git a/codex-rs/codex-mcp/src/connection_manager/resources.rs b/codex-rs/codex-mcp/src/connection_manager/resources.rs index 6762375147..c84cc6b43f 100644 --- a/codex-rs/codex-mcp/src/connection_manager/resources.rs +++ b/codex-rs/codex-mcp/src/connection_manager/resources.rs @@ -156,7 +156,10 @@ impl McpConnectionSet { .with_context(|| format!("resources/read failed for `{server}` ({uri})")) } - async fn client_by_name(&self, name: &str) -> Result<(ManagedClient, Option)> { + pub(crate) async fn client_by_name( + &self, + name: &str, + ) -> Result<(ManagedClient, Option)> { let view = self .servers .get(name) diff --git a/codex-rs/codex-mcp/src/lib.rs b/codex-rs/codex-mcp/src/lib.rs index add33ffede..32a3b65811 100644 --- a/codex-rs/codex-mcp/src/lib.rs +++ b/codex-rs/codex-mcp/src/lib.rs @@ -7,6 +7,10 @@ pub use elicitation::ElicitationLifecycle; pub use elicitation::ElicitationReviewRequest; pub use elicitation::ElicitationReviewer; pub use elicitation::ElicitationReviewerHandle; +pub use resource_client::McpEventCatalogSnapshot; +pub use resource_client::McpEventDefinition; +pub use resource_client::McpEventNotification; +pub use resource_client::McpEventStream; pub use resource_client::McpResourceClient; pub use resource_client::McpResourceClientCacheKey; pub use resource_client::McpResourcePage; diff --git a/codex-rs/codex-mcp/src/resource_client.rs b/codex-rs/codex-mcp/src/resource_client.rs index eb4fa7bc34..89ec66eb91 100644 --- a/codex-rs/codex-mcp/src/resource_client.rs +++ b/codex-rs/codex-mcp/src/resource_client.rs @@ -1,15 +1,26 @@ use std::sync::Arc; use std::sync::Weak; +use std::time::Duration; use anyhow::Context; use anyhow::Result; +use anyhow::anyhow; use codex_protocol::mcp::Resource; use codex_protocol::mcp::ResourceContent; +use codex_rmcp_client::CancellableEventStreamRequest; use rmcp::model::PaginatedRequestParams; use rmcp::model::ReadResourceRequestParams; +use rmcp::model::ServerResult; +use rmcp::service::ServiceError; +use serde::Deserialize; +use serde_json::Map; +use serde_json::Value; +use serde_json::json; +use tokio::runtime::Handle; use crate::McpRuntime; use crate::connection_manager::McpConnectionSet; +use crate::mcp::CODEX_APPS_MCP_SERVER_NAME; /// One page of resources returned by an MCP server. #[derive(Clone, Debug, PartialEq)] @@ -27,7 +38,105 @@ pub struct McpResourceReadResult { pub contents: Vec, } -/// Access to MCP resources through the latest runtime. +/// An event advertised by an MCP server. +#[derive(Clone, Debug, Deserialize, PartialEq)] +#[serde(rename_all = "camelCase")] +pub struct McpEventDefinition { + pub name: String, + pub description: String, + pub delivery: Vec, + pub input_schema: Value, + pub payload_schema: Value, +} + +/// Events returned from one stable MCP connection generation. +pub struct McpEventCatalogSnapshot { + pub cache_key: McpResourceClientCacheKey, + pub events: Vec, +} + +/// One unmodified lifecycle notification from an MCP event subscription. +#[derive(Clone, Debug, PartialEq)] +pub struct McpEventNotification { + pub method: String, + pub params: Option, +} + +/// Owns an MCP event subscription and cancels its request when dropped. +pub struct McpEventStream { + request: Option, + runtime_handle: Handle, + _connections: Arc, +} + +impl McpEventStream { + /// Receives the next raw lifecycle notification for this subscription. + pub async fn recv(&mut self) -> Result> { + let Some(request) = self.request.as_mut() else { + return Ok(None); + }; + + tokio::select! { + biased; + + notification = request.notifications.recv() => { + match notification { + Some(notification) => Ok(Some(McpEventNotification { + method: notification.method, + params: notification.params, + })), + None => { + let response = (&mut request.handle.rx).await; + self.request = None; + match response { + Ok(Ok(_)) | Ok(Err(ServiceError::Cancelled { .. })) => Ok(None), + Ok(Err(error)) => Err(error.into()), + Err(error) => Err(error.into()), + } + } + } + } + response = &mut request.handle.rx => { + self.request = None; + + match response { + Ok(Ok(_)) | Ok(Err(ServiceError::Cancelled { .. })) => Ok(None), + Ok(Err(error)) => Err(error.into()), + Err(error) => Err(error.into()), + } + } + } + } +} + +impl Drop for McpEventStream { + fn drop(&mut self) { + if let Some(CancellableEventStreamRequest { + handle, + notifications, + }) = self.request.take() + { + drop(notifications); + let connections = Arc::clone(&self._connections); + self.runtime_handle.spawn(async move { + let _ = tokio::time::timeout( + Duration::from_secs(30), + handle.cancel(Some("event subscription closed".to_string())), + ) + .await; + drop(connections); + }); + } + } +} + +#[derive(Deserialize)] +#[serde(rename_all = "camelCase")] +struct McpEventListResult { + events: Vec, +} + +/// Access to MCP resources and event subscriptions through the latest runtime. #[derive(Clone)] pub struct McpResourceClient { runtime: Arc, @@ -110,6 +219,63 @@ impl McpResourceClient { .collect::>>()?; Ok(McpResourceReadResult { contents }) } + + /// Lists the events advertised by the hosted Plugin Runtime. + pub async fn list_events(&self) -> Result { + let connections = self.runtime.latest_connections(); + let cache_key = McpResourceClientCacheKey(Arc::downgrade(&connections)); + let (managed, request_timeout) = connections + .client_by_name(CODEX_APPS_MCP_SERVER_NAME) + .await?; + let result = managed + .client + .send_custom_request_with_timeout("events/list", /*params*/ None, request_timeout) + .await + .context("events/list failed for hosted Plugin Runtime")?; + let ServerResult::CustomResult(result) = result else { + return Err(anyhow!("events/list returned an unexpected MCP result")); + }; + let result = result + .result_as::() + .context("events/list returned invalid event definitions")?; + + Ok(McpEventCatalogSnapshot { + cache_key, + events: result.events, + }) + } + + /// Opens an MCP event subscription with the supplied event arguments. + pub async fn open_event_stream( + &self, + event_name: &str, + arguments: &Value, + request_meta: Option<&Map>, + ) -> Result { + let mut params = json!({ + "name": event_name, + "arguments": arguments, + }); + if let Some(request_meta) = request_meta { + params["_meta"] = Value::Object(request_meta.clone()); + } + + let connections = self.runtime.latest_connections(); + let (managed, _) = connections + .client_by_name(CODEX_APPS_MCP_SERVER_NAME) + .await?; + let request = managed + .client + .send_event_stream_request(Some(params)) + .await + .context("events/stream failed for hosted Plugin Runtime")?; + + Ok(McpEventStream { + request: Some(request), + runtime_handle: Handle::current(), + _connections: connections, + }) + } } fn resource_from_rmcp(resource: rmcp::model::Resource) -> Result { diff --git a/codex-rs/rmcp-client/src/event_notification_transport.rs b/codex-rs/rmcp-client/src/event_notification_transport.rs new file mode 100644 index 0000000000..a0ff2cd5bb --- /dev/null +++ b/codex-rs/rmcp-client/src/event_notification_transport.rs @@ -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, + available_bytes: Arc, +} + +pub struct EventNotificationReceiver { + notifications: mpsc::UnboundedReceiver, + available_bytes: Arc, +} + +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 { + 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( + transport: T, +) -> impl Transport + 'static +where + T: IntoTransport, + E: std::error::Error + Send + Sync + 'static, +{ + EventNotificationTransport { + inner: transport.into_transport(), + routes: Arc::default(), + } +} + +struct EventNotificationTransport { + inner: T, + routes: Arc>>, +} + +impl Transport for EventNotificationTransport +where + T: Transport + 'static, +{ + type Error = T::Error; + + fn send( + &mut self, + mut message: TxJsonRpcMessage, + ) -> impl Future> + 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::() + .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> { + 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) -> Option<&RequestId> { + match message { + JsonRpcMessage::Response(response) => Some(&response.id), + JsonRpcMessage::Error(error) => error.id.as_ref(), + JsonRpcMessage::Request(_) | JsonRpcMessage::Notification(_) => None, + } +} diff --git a/codex-rs/rmcp-client/src/http_client_adapter.rs b/codex-rs/rmcp-client/src/http_client_adapter.rs index e1473001d9..2564aef3a1 100644 --- a/codex-rs/rmcp-client/src/http_client_adapter.rs +++ b/codex-rs/rmcp-client/src/http_client_adapter.rs @@ -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, default_headers: HeaderMap, auth_provider: Option, + event_stream_cancellations: Arc>>>, has_configured_headers: bool, redirect_mode: StreamableHttpRedirectMode, initialize_deadline: Arc>>, } +struct EventStreamCancellation { + request_id: RequestId, + cancellations: Arc>>>, +} + +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) = ¬ification.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, + modern_session: bool, ) -> BoxStream<'static, std::result::Result> { - 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 { diff --git a/codex-rs/rmcp-client/src/lib.rs b/codex-rs/rmcp-client/src/lib.rs index e476e4b029..1aef70fbd7 100644 --- a/codex-rs/rmcp-client/src/lib.rs +++ b/codex-rs/rmcp-client/src/lib.rs @@ -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; diff --git a/codex-rs/rmcp-client/src/rmcp_client.rs b/codex-rs/rmcp-client/src/rmcp_client.rs index 7a6a7d5d98..ddf2a73f51 100644 --- a/codex-rs/rmcp-client/src/rmcp_client.rs +++ b/codex-rs/rmcp-client/src/rmcp_client.rs @@ -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, } +/// An active Plugin Runtime event request and its request-scoped notifications. +pub struct CancellableEventStreamRequest { + pub handle: RequestHandle, + 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, + ) -> Result { + self.send_custom_request_with_timeout(method, params, /*timeout*/ None) + .await + } + + pub async fn send_custom_request_with_timeout( + &self, + method: &str, + params: Option, + timeout: Option, ) -> Result { 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, + ) -> Result { + 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>> { 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, ), diff --git a/codex-rs/rmcp-client/tests/mcp_events.rs b/codex-rs/rmcp-client/tests/mcp_events.rs new file mode 100644 index 0000000000..3b015a01d2 --- /dev/null +++ b/codex-rs/rmcp-client/tests/mcp_events.rs @@ -0,0 +1,242 @@ +mod streamable_http_test_support; + +use std::convert::Infallible; +use std::time::Duration; + +use anyhow::Context as _; +use axum::Json; +use axum::Router; +use axum::http::StatusCode; +use axum::response::IntoResponse; +use axum::response::sse::Event; +use axum::response::sse::Sse; +use axum::routing::post; +use futures::StreamExt as _; +use futures::stream; +use pretty_assertions::assert_eq; +use serde_json::Value; +use serde_json::json; +use tokio::net::TcpListener; +use tokio::sync::Notify; +use tokio::sync::mpsc; +use tokio::time::timeout; + +use streamable_http_test_support::create_client; + +struct StreamClosed { + event_name: String, + closed: mpsc::UnboundedSender, +} + +impl Drop for StreamClosed { + fn drop(&mut self) { + let _ = self.closed.send(self.event_name.clone()); + } +} + +fn initialize_response(message: &Value) -> axum::response::Response { + Json(json!({ + "jsonrpc": "2.0", + "id": message["id"], + "result": { + "protocolVersion": "2025-06-18", + "capabilities": {}, + "serverInfo": {"name": "plugin-runtime", "version": "1.0.0"}, + }, + })) + .into_response() +} + +fn plugin_runtime_notifications(event_name: &str, request_id: Value) -> [Value; 2] { + let metadata = json!({"io.modelcontextprotocol/subscriptionId": request_id}); + [ + json!({ + "jsonrpc": "2.0", + "method": "notifications/events/active", + "params": {"cursor": null, "truncated": false, "_meta": metadata}, + }), + json!({ + "jsonrpc": "2.0", + "method": "notifications/events/event", + "params": { + "eventId": format!("event-{event_name}"), + "name": event_name, + "timestamp": "2026-08-07T12:00:00Z", + "data": {"issue": 42}, + "cursor": null, + "_meta": metadata, + }, + }), + ] +} + +#[tokio::test] +async fn plugin_runtime_event_streams_are_isolated_and_cancel_locally() -> anyhow::Result<()> { + let (stream_closed_tx, mut stream_closed_rx) = mpsc::unbounded_channel::(); + let router = Router::new().route( + "/mcp", + post(move |Json(message): Json| { + let stream_closed_tx = stream_closed_tx.clone(); + async move { + match message["method"].as_str() { + Some("initialize") => initialize_response(&message), + Some("notifications/initialized") => StatusCode::ACCEPTED.into_response(), + Some("events/stream") => { + let event_name = message["params"]["name"] + .as_str() + .expect("event stream request must contain a name") + .to_string(); + assert_eq!( + message.pointer("/params/arguments"), + Some(&json!({"project": "codex"})) + ); + assert!(message["params"].get("cursor").is_none()); + + let closed = StreamClosed { + event_name: event_name.clone(), + closed: stream_closed_tx, + }; + let events = stream::iter( + plugin_runtime_notifications(&event_name, message["id"].clone()) + .into_iter() + .map(|notification| { + Ok::<_, Infallible>( + Event::default() + .event("message") + .data(notification.to_string()), + ) + }), + ) + .chain(stream::pending()) + .map(move |event| { + let _ = &closed; + event + }); + + Sse::new(events).into_response() + } + Some("notifications/cancelled") => { + panic!("event cancellation should close the local stream without a POST") + } + method => panic!("unexpected Plugin Runtime request: {method:?}"), + } + } + }), + ); + + let listener = TcpListener::bind("127.0.0.1:0").await?; + let base_url = format!("http://{}", listener.local_addr()?); + let server = tokio::spawn(async move { axum::serve(listener, router).await }); + let client = create_client(&base_url).await?; + + let mut first = client + .send_event_stream_request(Some(json!({ + "name": "github.pull_request.opened", + "arguments": {"project": "codex"}, + }))) + .await?; + let mut second = client + .send_event_stream_request(Some(json!({ + "name": "gmail.message.received", + "arguments": {"project": "codex"}, + }))) + .await?; + + for (event_name, request) in [ + ("github.pull_request.opened", &mut first), + ("gmail.message.received", &mut second), + ] { + let active = timeout(Duration::from_secs(5), request.notifications.recv()) + .await? + .context("event stream closed before activation")?; + assert_eq!(active.method, "notifications/events/active"); + assert_eq!( + active.params, + Some(json!({"cursor": null, "truncated": false})) + ); + + let event = timeout(Duration::from_secs(5), request.notifications.recv()) + .await? + .context("event stream closed before delivery")?; + assert_eq!(event.method, "notifications/events/event"); + assert_eq!( + event.params.as_ref().and_then(|params| params.get("name")), + Some(&json!(event_name)) + ); + } + + first + .handle + .cancel(Some("event subscription closed".to_string())) + .await?; + let closed = timeout(Duration::from_secs(5), stream_closed_rx.recv()) + .await? + .context("cancelled stream did not close")?; + assert_eq!(closed, "github.pull_request.opened"); + + second + .handle + .cancel(Some("event subscription closed".to_string())) + .await?; + let closed = timeout(Duration::from_secs(5), stream_closed_rx.recv()) + .await? + .context("cancelled stream did not close")?; + assert_eq!(closed, "gmail.message.received"); + + client.shutdown().await; + server.abort(); + let _ = server.await; + Ok(()) +} + +#[tokio::test] +async fn plugin_runtime_event_stream_times_out_before_stalled_headers() -> anyhow::Result<()> { + let request_started = std::sync::Arc::new(Notify::new()); + let server_request_started = std::sync::Arc::clone(&request_started); + let router = Router::new().route( + "/mcp", + post(move |Json(message): Json| { + let request_started = std::sync::Arc::clone(&server_request_started); + async move { + match message["method"].as_str() { + Some("initialize") => initialize_response(&message), + Some("notifications/initialized") => StatusCode::ACCEPTED.into_response(), + Some("events/stream") => { + request_started.notify_one(); + std::future::pending().await + } + method => panic!("unexpected Plugin Runtime request: {method:?}"), + } + } + }), + ); + + let listener = TcpListener::bind("127.0.0.1:0").await?; + let base_url = format!("http://{}", listener.local_addr()?); + let server = tokio::spawn(async move { axum::serve(listener, router).await }); + let client = create_client(&base_url).await?; + let request = client + .send_event_stream_request(Some(json!({ + "name": "github.pull_request.opened", + "arguments": {}, + }))) + .await?; + + timeout(Duration::from_secs(5), request_started.notified()) + .await + .context("event stream request did not reach the server")?; + tokio::time::pause(); + tokio::time::advance(Duration::from_secs(31)).await; + let error = request + .handle + .rx + .await? + .expect_err("stalled response headers must time out"); + tokio::time::resume(); + assert!(error.to_string().contains("timed out")); + + client.shutdown().await; + server.abort(); + let _ = server.await; + Ok(()) +}