diff --git a/.github/workflows/conformance.yml b/.github/workflows/conformance.yml index 297fad8f3..817c7d4ed 100644 --- a/.github/workflows/conformance.yml +++ b/.github/workflows/conformance.yml @@ -226,6 +226,14 @@ jobs: --scenario sep-2322-client-request-state \ -o conformance-client-results/mrtr + - name: Run draft SEP-2575 client scenario + run: | + npx -y "@modelcontextprotocol/conformance@${DRAFT_CONFORMANCE_VERSION}" client \ + --command "$(pwd)/target/debug/conformance-client" \ + --scenario request-metadata \ + --spec-version draft \ + -o conformance-client-results/sep-2575 + - name: Upload results if: always() uses: actions/upload-artifact@v7 diff --git a/README.md b/README.md index 68f3343c9..0f8367377 100644 --- a/README.md +++ b/README.md @@ -74,6 +74,44 @@ async fn main() -> Result<(), Box> { ``` +### Client lifecycle modes + +`serve()` uses the legacy MCP lifecycle: the client sends `initialize`, receives +the negotiated server information, and then sends `notifications/initialized`. +Use [`ClientServiceExt::serve_with_lifecycle`](crates/rmcp/src/service/client.rs) to +select another lifecycle explicitly: + +```rust, ignore +use rmcp::{ClientInfo, ClientLifecycleMode, ClientServiceExt, ProtocolVersion}; + +// Start directly with server/discover and include client metadata on every request. +let client = ClientInfo::default() + .serve_with_lifecycle( + transport, + ClientLifecycleMode::Discover { + preferred_versions: vec![ProtocolVersion::V_2026_07_28], + }, + ) + .await?; + +// Or probe the discover lifecycle and fall back when a legacy server reports +// that server/discover is not implemented. +let client = ClientInfo::default() + .serve_with_lifecycle( + transport, + ClientLifecycleMode::Auto { + preferred_versions: vec![ProtocolVersion::V_2026_07_28], + legacy_version: Some(ProtocolVersion::V_2025_11_25), + }, + ) + .await?; +``` + +`ClientLifecycleMode::Initialize` is equivalent to the existing `serve()` behavior. +Discover startup does not send `notifications/initialized`; discovery completes +startup, and each subsequent request carries its protocol version, client +information, and capabilities in `_meta`. + ### Build a Server
@@ -812,10 +850,11 @@ impl ServerHandler for MyServer { ### Initialized notification -Clients send `initialized` after the handshake completes: +Legacy clients send `initialized` after the `initialize` handshake completes. +Clients using `ClientLifecycleMode::Discover` do not send this notification: ```rust -// Sent automatically by rmcp during the serve() handshake. +// Sent automatically by rmcp during the legacy serve() handshake. // Servers handle it via: impl ServerHandler for MyServer { async fn on_initialized( diff --git a/conformance/src/bin/client.rs b/conformance/src/bin/client.rs index cd4671f1f..63a116ff4 100644 --- a/conformance/src/bin/client.rs +++ b/conformance/src/bin/client.rs @@ -1,7 +1,7 @@ use rmcp::{ - ClientHandler, ErrorData, RoleClient, ServiceExt, + ClientHandler, ClientLifecycleMode, ClientServiceExt, ErrorData, RoleClient, ServiceExt, model::*, - service::{RequestContext, serve_directly}, + service::RequestContext, transport::{ AuthClient, AuthorizationManager, StreamableHttpClientTransport, auth::{AuthorizationCallback, OAuthState}, @@ -846,33 +846,6 @@ async fn run_elicitation_defaults_client(server_url: &str) -> anyhow::Result<()> Ok(()) } -/// A minimal stateless client transport: every outgoing message is one HTTP -/// POST and the JSON response body (if any) is queued for `receive()`. -/// -/// The SEP-2322 client scenario's mock server speaks the stateless lifecycle -/// (no `initialize` handshake, plain JSON responses), which the session-based -/// `StreamableHttpClientTransport` cannot do. The transport is harness -/// plumbing; the behavior under test — the SDK's MRTR retry driver — runs -/// unchanged on top of it. -struct StatelessHttpTransport { - http: reqwest::Client, - uri: std::sync::Arc, - tx: tokio::sync::mpsc::Sender, - rx: tokio::sync::mpsc::Receiver, -} - -impl StatelessHttpTransport { - fn new(uri: &str) -> Self { - let (tx, rx) = tokio::sync::mpsc::channel(16); - Self { - http: reqwest::Client::new(), - uri: uri.into(), - tx, - rx, - } - } -} - fn conformance_protocol_version() -> ProtocolVersion { std::env::var("MCP_CONFORMANCE_PROTOCOL_VERSION") .ok() @@ -880,62 +853,22 @@ fn conformance_protocol_version() -> ProtocolVersion { .unwrap_or(ProtocolVersion::V_2026_07_28) } -impl rmcp::transport::Transport for StatelessHttpTransport { - type Error = std::io::Error; - - fn send( - &mut self, - item: rmcp::model::ClientJsonRpcMessage, - ) -> impl std::future::Future> + Send + 'static { - let http = self.http.clone(); - let uri = self.uri.clone(); - let tx = self.tx.clone(); - async move { - let response = http - .post(uri.as_ref()) - .header( - "MCP-Protocol-Version", - conformance_protocol_version().as_str(), - ) - .json(&item) - .send() - .await - .map_err(std::io::Error::other)?; - match response.json::().await { - Ok(message) => { - let _ = tx.send(message).await; - } - Err(_) => { - // No JSON-RPC body (e.g. 202/204 for notifications). - } - } - Ok(()) +/// Runs draft stateless scenarios through the public discover lifecycle and +/// Streamable HTTP transport. +async fn run_discover_client(server_url: &str) -> anyhow::Result<()> { + let mut preferred_versions = vec![conformance_protocol_version()]; + for version in ProtocolVersion::KNOWN_VERSIONS.iter().rev() { + if !preferred_versions.contains(version) { + preferred_versions.push(version.clone()); } } - - async fn receive(&mut self) -> Option { - self.rx.recv().await - } - - async fn close(&mut self) -> Result<(), Self::Error> { - Ok(()) - } -} - -/// Runs a client using the draft stateless lifecycle. -/// -/// Stateless servers do not implement the `initialize` handshake, so this -/// uses `serve_directly`. The protocol version comes from -/// `MCP_CONFORMANCE_PROTOCOL_VERSION` (defaulting to `2026-07-28`) and is -/// used for both peer configuration and outgoing HTTP request headers. -/// -/// Lists available tools and calls each one, allowing the SDK's high-level -/// tool-call handling to process any request retries. -async fn run_stateless_client(server_url: &str) -> anyhow::Result<()> { - let transport = StatelessHttpTransport::new(server_url); - let peer_info = InitializeResult::new(ServerCapabilities::builder().enable_tools().build()) - .with_protocol_version(conformance_protocol_version()); - let client = serve_directly(FullClientHandler, transport, Some(peer_info)); + let transport = StreamableHttpClientTransport::from_uri(server_url); + let client = FullClientHandler + .serve_with_lifecycle( + transport, + ClientLifecycleMode::Discover { preferred_versions }, + ) + .await?; let tools = client.list_tools(Default::default()).await?; tracing::debug!("Listed {} tools", tools.tools.len()); @@ -989,14 +922,14 @@ async fn main() -> anyhow::Result<()> { match scenario.as_str() { // Non-auth scenarios "initialize" => run_basic_client(&server_url).await?, - "json-schema-ref-no-deref" => run_stateless_client(&server_url).await?, + "json-schema-ref-no-deref" => run_discover_client(&server_url).await?, "tools_call" => run_tools_call_client(&server_url, &ctx).await?, "elicitation-sep1034-client-defaults" => { run_elicitation_defaults_client(&server_url).await? } "sse-retry" => run_sse_retry_client(&server_url).await?, "request-metadata" | "sep-2322-client-request-state" => { - run_stateless_client(&server_url).await? + run_discover_client(&server_url).await? } "http-standard-headers" | "http-custom-headers" | "http-invalid-tool-headers" => { run_tools_call_client(&server_url, &ctx).await? diff --git a/crates/rmcp/Cargo.toml b/crates/rmcp/Cargo.toml index a963c156d..6004f9323 100644 --- a/crates/rmcp/Cargo.toml +++ b/crates/rmcp/Cargo.toml @@ -302,6 +302,26 @@ name = "test_protocol_version_negotiation" required-features = ["server", "client"] path = "tests/test_protocol_version_negotiation.rs" +[[test]] +name = "test_client_lifecycle_modes" +required-features = ["client", "server"] +path = "tests/test_client_lifecycle_modes.rs" + +[[test]] +name = "test_stateless_server_requests" +required-features = ["client", "server"] +path = "tests/test_stateless_server_requests.rs" + +[[test]] +name = "test_discover_http_client_startup" +required-features = [ + "client", + "reqwest", + "transport-streamable-http-client-reqwest", + "transport-streamable-http-server", +] +path = "tests/test_discover_http_client_startup.rs" + [[test]] name = "test_streamable_http_standard_headers" required-features = ["server", "client", "transport-streamable-http-server", "reqwest"] diff --git a/crates/rmcp/src/handler/server.rs b/crates/rmcp/src/handler/server.rs index c4db90f37..689bcfc78 100644 --- a/crates/rmcp/src/handler/server.rs +++ b/crates/rmcp/src/handler/server.rs @@ -41,22 +41,30 @@ impl Service for H { )); } } - if matches!(&request, ClientRequest::DiscoverRequest(_)) { - if requested_version.is_none() { + // Self-contained metadata is required only when the request itself uses + // the inline lifecycle: a discover opener, a session that started without + // `initialize`, or a request that declares 2026-07-28+ in its own _meta. + // Sessions that negotiated via `initialize` (or `serve_directly`) keep the + // session model and may omit per-request metadata. + let requires_request_metadata = uses_inline_negotiation + && (matches!(&request, ClientRequest::DiscoverRequest(_)) + || context.peer.request_metadata_required() + || requested_version.as_ref().is_some_and(|version| { + version.as_str() >= ProtocolVersion::V_2026_07_28.as_str() + })); + if requires_request_metadata { + // Inline lifecycle requests are defined by the 2026-07-28 protocol. + // Validate that lifecycle contract even when a request selects an + // older application protocol version. + let missing = context + .meta + .missing_required_keys(&ProtocolVersion::V_2026_07_28); + if !missing.is_empty() { return Err(McpError::invalid_params( - "server/discover requires protocolVersion in request _meta", - None, - )); - } - if context.meta.client_info().is_none() { - return Err(McpError::invalid_params( - "server/discover requires clientInfo in request _meta", - None, - )); - } - if context.meta.client_capabilities().is_none() { - return Err(McpError::invalid_params( - "server/discover requires clientCapabilities in request _meta", + format!( + "request _meta is missing or has malformed required fields: {}", + missing.join(", ") + ), None, )); } diff --git a/crates/rmcp/src/lib.rs b/crates/rmcp/src/lib.rs index 78ba2e27f..ca195a6e8 100644 --- a/crates/rmcp/src/lib.rs +++ b/crates/rmcp/src/lib.rs @@ -16,10 +16,13 @@ pub use handler::client::ClientHandler; pub use handler::server::ServerHandler; #[cfg(feature = "server")] pub use handler::server::wrapper::Json; +#[cfg(feature = "client")] +pub use service::{ + ClientLifecycleMode, ClientServiceExt, RoleClient, select_protocol_version, serve_client, + serve_client_with_lifecycle, +}; #[cfg(any(feature = "client", feature = "server"))] pub use service::{Peer, Service, ServiceError, ServiceExt}; -#[cfg(feature = "client")] -pub use service::{RoleClient, select_protocol_version, serve_client}; #[cfg(feature = "server")] pub use service::{RoleServer, serve_server}; diff --git a/crates/rmcp/src/model/meta.rs b/crates/rmcp/src/model/meta.rs index 9b1e1b0f6..c51fc6d7c 100644 --- a/crates/rmcp/src/model/meta.rs +++ b/crates/rmcp/src/model/meta.rs @@ -439,6 +439,19 @@ impl RequestMetaObject { meta } + /// Create request metadata with the client context SEP-2575 requires on every request. + pub fn with_client_context( + protocol_version: ProtocolVersion, + client_info: Implementation, + client_capabilities: ClientCapabilities, + ) -> Self { + let mut meta = Self::new(); + meta.set_protocol_version(protocol_version); + meta.set_client_info(client_info); + meta.set_client_capabilities(client_capabilities); + meta + } + pub(crate) fn static_empty() -> &'static Self { static EMPTY: std::sync::OnceLock = std::sync::OnceLock::new(); EMPTY.get_or_init(Default::default) diff --git a/crates/rmcp/src/service.rs b/crates/rmcp/src/service.rs index 6a7aa53e7..69c3f2105 100644 --- a/crates/rmcp/src/service.rs +++ b/crates/rmcp/src/service.rs @@ -1,3 +1,5 @@ +use std::sync::OnceLock; + use futures::FutureExt; #[cfg(not(feature = "local"))] use futures::future::BoxFuture; @@ -51,9 +53,10 @@ use crate::model::ServerNotification; use crate::{ error::ErrorData as McpError, model::{ - CancelledNotification, CancelledNotificationParam, Extensions, GetExtensions, GetMeta, - JsonRpcError, JsonRpcMessage, JsonRpcNotification, JsonRpcRequest, JsonRpcResponse, - NotificationMetaObject, NumberOrString, ProgressToken, RequestId, RequestMetaObject, + CancelledNotification, CancelledNotificationParam, ClientCapabilities, Extensions, + GetExtensions, GetMeta, Implementation, JsonRpcError, JsonRpcMessage, JsonRpcNotification, + JsonRpcRequest, JsonRpcResponse, NotificationMetaObject, NumberOrString, ProgressToken, + ProtocolVersion, RequestId, RequestMetaObject, }, transport::{DynamicTransportError, IntoTransport, Transport}, }; @@ -515,6 +518,13 @@ pub(crate) enum PeerSinkMessage { }, } +#[derive(Debug, Clone)] +pub(crate) struct ClientRequestMetadata { + pub protocol_version: ProtocolVersion, + pub client_info: Implementation, + pub client_capabilities: ClientCapabilities, +} + /// An interface to fetch the remote client or server /// /// For general purpose, call [`Peer::send_request`] or [`Peer::send_notification`] to send message to remote peer. @@ -527,6 +537,8 @@ pub struct Peer { progress_token_provider: Arc, progress_timeout_watchers: ProgressTimeoutWatchers, info: Arc>>>, + client_request_metadata: Arc>, + request_metadata_required: Arc, } impl std::fmt::Debug for Peer { @@ -563,6 +575,14 @@ impl PeerRequestOptions { } } + /// Adds request metadata while preserving any other configured options. + /// + /// Explicit values take precedence over discover-lifecycle metadata defaults. + pub fn with_meta(mut self, meta: RequestMetaObject) -> Self { + self.meta = Some(meta); + self + } + pub fn reset_timeout_on_progress(mut self) -> Self { self.reset_timeout_on_progress = true; self @@ -588,6 +608,8 @@ impl Peer { progress_token_provider: Arc::new(AtomicU32ProgressTokenProvider::default()), progress_timeout_watchers: Default::default(), info: Arc::new(std::sync::RwLock::new(peer_info.map(Arc::new))), + client_request_metadata: Default::default(), + request_metadata_required: Default::default(), }, rx, ) @@ -625,6 +647,12 @@ impl Peer { ) -> Result, ServiceError> { let id = self.request_id_provider.next_request_id(); let progress_token = self.progress_token_provider.next_progress_token(); + if let Some(metadata) = self.client_request_metadata.get() { + let meta = request.get_meta_mut(); + meta.set_protocol_version(metadata.protocol_version.clone()); + meta.set_client_info(metadata.client_info.clone()); + meta.set_client_capabilities(metadata.client_capabilities.clone()); + } if let Some(meta) = options.meta.clone() { request.get_meta_mut().extend(meta); } @@ -703,6 +731,21 @@ impl Peer { *self.info.write().expect("peer info lock poisoned") = Some(Arc::new(info)); } + pub(crate) fn set_client_request_metadata(&self, metadata: ClientRequestMetadata) { + let result = self.client_request_metadata.set(metadata); + debug_assert!(result.is_ok(), "client request metadata set more than once"); + } + + pub(crate) fn require_request_metadata(&self) { + self.request_metadata_required + .store(true, std::sync::atomic::Ordering::Release); + } + + pub(crate) fn request_metadata_required(&self) -> bool { + self.request_metadata_required + .load(std::sync::atomic::Ordering::Acquire) + } + pub fn is_transport_closed(&self) -> bool { self.tx.is_closed() } @@ -883,11 +926,35 @@ impl RequestContext { #[cfg(feature = "server")] impl RequestContext { - /// The protocol version the client negotiated, or `None` before peer info is recorded. + /// The current request's protocol version, falling back to legacy handshake state. pub fn protocol_version(&self) -> Option { - self.peer - .peer_info() - .map(|info| info.protocol_version.clone()) + self.meta.protocol_version().or_else(|| { + self.peer + .peer_info() + .map(|info| info.protocol_version.clone()) + }) + } + + /// The current request's client implementation, falling back only for legacy sessions. + pub fn client_info(&self) -> Option { + if self.peer.request_metadata_required() { + self.meta.client_info() + } else { + self.meta + .client_info() + .or_else(|| self.peer.peer_info().map(|info| info.client_info.clone())) + } + } + + /// The current request's client capabilities, falling back only for legacy sessions. + pub fn client_capabilities(&self) -> Option { + if self.peer.request_metadata_required() { + self.meta.client_capabilities() + } else { + self.meta + .client_capabilities() + .or_else(|| self.peer.peer_info().map(|info| info.capabilities.clone())) + } } } diff --git a/crates/rmcp/src/service/client.rs b/crates/rmcp/src/service/client.rs index d922f59f9..9ba3024d0 100644 --- a/crates/rmcp/src/service/client.rs +++ b/crates/rmcp/src/service/client.rs @@ -54,6 +54,17 @@ pub enum ClientInitializeError { #[error("JSON-RPC error: {0}")] JsonRpcError(ErrorData), + #[error( + "no compatible protocol version (client: {client_supported:?}, server: {server_supported:?})" + )] + NoCompatibleProtocolVersion { + client_supported: Vec, + server_supported: Vec, + }, + + #[error("discover startup requires at least one preferred protocol version")] + NoPreferredProtocolVersion, + #[error("Cancelled")] Cancelled, } @@ -177,6 +188,43 @@ impl ServiceRole for RoleClient { pub type ServerSink = Peer; +/// Selects how a client establishes its MCP lifecycle. +/// +/// Existing [`ServiceExt::serve`] behavior remains legacy initialization. +#[derive(Debug, Clone, PartialEq, Eq)] +#[non_exhaustive] +pub enum ClientLifecycleMode { + /// Use the legacy `initialize` / `notifications/initialized` handshake. + Initialize, + /// Use `server/discover` and send self-contained per-request metadata. + Discover { + preferred_versions: Vec, + }, + /// Probe with `server/discover`, falling back only when the peer proves it is legacy. + Auto { + preferred_versions: Vec, + legacy_version: Option, + }, +} + +/// Client-specific lifecycle entry points. +pub trait ClientServiceExt: Service + Sized { + fn serve_with_lifecycle( + self, + transport: T, + lifecycle: ClientLifecycleMode, + ) -> impl Future, ClientInitializeError>> + + MaybeSendFuture + where + T: IntoTransport, + E: std::error::Error + Send + Sync + 'static, + { + serve_client_with_lifecycle(self, transport, lifecycle) + } +} + +impl> ClientServiceExt for S {} + impl> ServiceExt for S { fn serve_with_ct( self, @@ -202,7 +250,13 @@ where T: IntoTransport, E: std::error::Error + Send + Sync + 'static, { - serve_client_with_ct(service, transport, Default::default()).await + serve_client_with_lifecycle_and_ct( + service, + transport, + ClientLifecycleMode::Initialize, + Default::default(), + ) + .await } pub async fn serve_client_with_ct( @@ -210,13 +264,41 @@ pub async fn serve_client_with_ct( transport: T, ct: CancellationToken, ) -> Result, ClientInitializeError> +where + S: Service, + T: IntoTransport, + E: std::error::Error + Send + Sync + 'static, +{ + serve_client_with_lifecycle_and_ct(service, transport, ClientLifecycleMode::Initialize, ct) + .await +} + +pub async fn serve_client_with_lifecycle( + service: S, + transport: T, + lifecycle: ClientLifecycleMode, +) -> Result, ClientInitializeError> +where + S: Service, + T: IntoTransport, + E: std::error::Error + Send + Sync + 'static, +{ + serve_client_with_lifecycle_and_ct(service, transport, lifecycle, Default::default()).await +} + +pub async fn serve_client_with_lifecycle_and_ct( + service: S, + transport: T, + lifecycle: ClientLifecycleMode, + ct: CancellationToken, +) -> Result, ClientInitializeError> where S: Service, T: IntoTransport, E: std::error::Error + Send + Sync + 'static, { tokio::select! { - result = serve_client_with_ct_inner(service, transport.into_transport(), ct.clone()) => { result } + result = serve_client_with_ct_inner(service, transport.into_transport(), lifecycle, ct.clone()) => { result } _ = ct.cancelled() => { Err(ClientInitializeError::Cancelled) } @@ -226,6 +308,7 @@ where async fn serve_client_with_ct_inner( service: S, transport: T, + lifecycle: ClientLifecycleMode, ct: CancellationToken, ) -> Result, ClientInitializeError> where @@ -234,12 +317,71 @@ where { let mut transport = transport.into_transport(); let id_provider = >::default(); + let (peer, peer_rx) = Peer::new(id_provider.clone(), None); + let client_info = service.get_info(); - // service + match lifecycle { + ClientLifecycleMode::Initialize => { + legacy_startup(&service, &mut transport, &id_provider, &peer, client_info).await?; + } + ClientLifecycleMode::Discover { preferred_versions } => { + discover_startup( + &service, + &mut transport, + &id_provider, + &peer, + &client_info, + preferred_versions, + ) + .await?; + } + ClientLifecycleMode::Auto { + preferred_versions, + legacy_version, + } => { + let discover_result = discover_startup( + &service, + &mut transport, + &id_provider, + &peer, + &client_info, + preferred_versions, + ) + .await; + match discover_result { + Ok(()) => {} + Err(ClientInitializeError::JsonRpcError(error)) + if error.code == crate::model::ErrorCode::METHOD_NOT_FOUND => + { + let mut legacy_info = client_info; + if let Some(version) = legacy_version { + legacy_info.protocol_version = version; + } + legacy_startup(&service, &mut transport, &id_provider, &peer, legacy_info) + .await?; + } + Err(error) => return Err(error), + } + } + } + Ok(serve_inner(service, transport, peer, peer_rx, ct)) +} + +async fn legacy_startup( + service: &S, + transport: &mut T, + id_provider: &Arc, + peer: &Peer, + client_info: ClientInfo, +) -> Result<(), ClientInitializeError> +where + S: Service, + T: Transport + 'static, +{ let id = id_provider.next_request_id(); let init_request = InitializeRequest { method: Default::default(), - params: service.get_info(), + params: client_info, extensions: Default::default(), }; transport @@ -253,15 +395,8 @@ where context: "send initialize request".into(), })?; - let (peer, peer_rx) = Peer::new(id_provider, None); - - let (response, response_id) = expect_response( - &mut transport, - "initialize response", - &service, - peer.clone(), - ) - .await?; + let (response, response_id) = + expect_response(transport, "initialize response", service, peer.clone()).await?; if id != response_id { return Err(ClientInitializeError::ConflictInitResponseId( @@ -285,7 +420,115 @@ where transport.send(notification).await.map_err(|error| { ClientInitializeError::transport::(error, "send initialized notification") })?; - Ok(serve_inner(service, transport, peer, peer_rx, ct)) + Ok(()) +} + +async fn discover_startup( + service: &S, + transport: &mut T, + id_provider: &Arc, + peer: &Peer, + client_info: &ClientInfo, + preferred_versions: Vec, +) -> Result<(), ClientInitializeError> +where + S: Service, + T: Transport + 'static, +{ + if preferred_versions.is_empty() { + return Err(ClientInitializeError::NoPreferredProtocolVersion); + } + + let mut attempted = Vec::new(); + let mut candidate = preferred_versions[0].clone(); + loop { + attempted.push(candidate.clone()); + + let meta = RequestMetaObject::with_client_context( + candidate.clone(), + client_info.client_info.clone(), + client_info.capabilities.clone(), + ); + let mut discover = DiscoverRequest::new(DiscoverRequestParams {}); + discover.extensions.insert(meta); + let id = id_provider.next_request_id(); + transport + .send(ClientJsonRpcMessage::request( + ClientRequest::DiscoverRequest(discover), + id.clone(), + )) + .await + .map_err(|error| { + ClientInitializeError::transport::(error, "send discover request") + })?; + + match expect_response(transport, "discover response", service, peer.clone()).await { + Ok((ServerResult::DiscoverResult(result), response_id)) => { + if response_id != id { + return Err(ClientInitializeError::ConflictInitResponseId( + id, + response_id, + )); + } + let Some(selected) = + select_protocol_version(&preferred_versions, &result.supported_versions) + else { + return Err(ClientInitializeError::NoCompatibleProtocolVersion { + client_supported: preferred_versions, + server_supported: result.supported_versions, + }); + }; + peer.set_peer_info(ServerInfo { + protocol_version: selected.clone(), + capabilities: result.capabilities, + server_info: result.server_info, + instructions: result.instructions, + meta: result.meta, + }); + peer.set_client_request_metadata(ClientRequestMetadata { + protocol_version: selected, + client_info: client_info.client_info.clone(), + client_capabilities: client_info.capabilities.clone(), + }); + return Ok(()); + } + Ok((response, _)) => { + return Err(ClientInitializeError::ExpectedInitResult(Some(response))); + } + Err(ClientInitializeError::JsonRpcError(error)) + if error.code == crate::model::ErrorCode::UNSUPPORTED_PROTOCOL_VERSION => + { + let supported = error + .data + .as_ref() + .and_then(|data| data.get("supported")) + .cloned() + .and_then(|value| serde_json::from_value::>(value).ok()) + .unwrap_or_default(); + let may_retry_current = attempted + .iter() + .filter(|version| *version == &candidate) + .count() + == 1; + let next = preferred_versions + .iter() + .find(|version| { + supported.contains(version) + && (!attempted.contains(version) + || (may_retry_current && *version == &candidate)) + }) + .cloned(); + let Some(next) = next else { + return Err(ClientInitializeError::NoCompatibleProtocolVersion { + client_supported: preferred_versions, + server_supported: supported, + }); + }; + candidate = next; + } + Err(error) => return Err(error), + } + } } macro_rules! method { diff --git a/crates/rmcp/src/service/server.rs b/crates/rmcp/src/service/server.rs index 8c7a87dda..e24ae085f 100644 --- a/crates/rmcp/src/service/server.rs +++ b/crates/rmcp/src/service/server.rs @@ -217,12 +217,41 @@ where } }; - let ClientRequest::InitializeRequest(peer_info) = &request else { - return Err(ServerInitializeError::ExpectedInitializeRequest(Some( - ClientJsonRpcMessage::request(request, id), - ))); + let initialize_request = match request { + ClientRequest::InitializeRequest(request) => request, + mut request => { + if !request + .get_meta() + .missing_required_keys(&ProtocolVersion::V_2026_07_28) + .is_empty() + { + return Err(ServerInitializeError::ExpectedInitializeRequest(Some( + ClientJsonRpcMessage::request(request, id), + ))); + } + let (peer, peer_rx) = Peer::new(id_provider, None); + peer.require_request_metadata(); + let context = RequestContext { + ct: ct.child_token(), + id: id.clone(), + meta: std::mem::take(request.get_meta_mut()), + extensions: std::mem::take(request.extensions_mut()), + peer: peer.clone(), + }; + let response = match service.handle_request(request, context).await { + Ok(result) => ServerJsonRpcMessage::response(result, id), + Err(error) => ServerJsonRpcMessage::error(error, Some(id)), + }; + transport.send(response).await.map_err(|error| { + ServerInitializeError::transport::(error, "sending negotiated request response") + })?; + return Ok(serve_inner(service, transport, peer, peer_rx, ct)); + } }; - let (peer, peer_rx) = Peer::new(id_provider, Some(peer_info.params.clone())); + let requested_protocol_version = initialize_request.params.protocol_version.clone(); + let mut negotiated_peer_info = initialize_request.params.clone(); + let (peer, peer_rx) = Peer::new(id_provider, Some(negotiated_peer_info.clone())); + let request = ClientRequest::InitializeRequest(initialize_request); let context = RequestContext { ct: ct.child_token(), id: id.clone(), @@ -231,7 +260,7 @@ where peer: peer.clone(), }; // Send initialize response - let init_response = service.handle_request(request.clone(), context).await; + let init_response = service.handle_request(request, context).await; let mut init_response = match init_response { Ok(ServerResult::InitializeResult(init_response)) => init_response, Ok(result) => { @@ -247,13 +276,10 @@ where return Err(ServerInitializeError::InitializeFailed(e)); } }; - init_response.protocol_version = negotiate_protocol_version( - &peer_info.params.protocol_version, - init_response.protocol_version, - ); + init_response.protocol_version = + negotiate_protocol_version(&requested_protocol_version, init_response.protocol_version); // Update peer_info so context.protocol_version() reflects the negotiated // version in all subsequent request handlers. - let mut negotiated_peer_info = peer_info.params.clone(); negotiated_peer_info.protocol_version = init_response.protocol_version.clone(); peer.set_peer_info(negotiated_peer_info); transport diff --git a/crates/rmcp/src/transport/streamable_http_client.rs b/crates/rmcp/src/transport/streamable_http_client.rs index 8fe6d7621..8ba3e423c 100644 --- a/crates/rmcp/src/transport/streamable_http_client.rs +++ b/crates/rmcp/src/transport/streamable_http_client.rs @@ -19,8 +19,9 @@ use super::common::client_side_sse::{ use crate::{ RoleClient, model::{ - ClientJsonRpcMessage, ClientNotification, ErrorData, InitializedNotification, JsonObject, - ProtocolVersion, RequestId, ServerJsonRpcMessage, ServerResult, + ClientJsonRpcMessage, ClientNotification, ClientRequest, ErrorData, GetMeta, + InitializedNotification, JsonObject, ProtocolVersion, RequestId, ServerJsonRpcMessage, + ServerResult, }, transport::{ common::{client_side_sse::SseAutoReconnectStream, mcp_headers}, @@ -29,6 +30,7 @@ use crate::{ }; type BoxedSseStream = BoxStream<'static, Result>; +const SESSION_CLEANUP_TIMEOUT: Duration = Duration::from_secs(5); fn build_request_headers( base: &HashMap, @@ -58,6 +60,27 @@ fn build_request_headers( headers } +fn request_version_headers( + base: &HashMap, + message: &ClientJsonRpcMessage, + fallback: &ProtocolVersion, + tool_cache: &HashMap>, +) -> (ProtocolVersion, HashMap) { + let version = match message { + ClientJsonRpcMessage::Request(request) => request + .request + .get_meta() + .protocol_version() + .unwrap_or_else(|| fallback.clone()), + _ => fallback.clone(), + }; + let mut headers = build_request_headers(base, message, tool_cache, &version); + if let Ok(value) = HeaderValue::from_str(version.as_str()) { + headers.insert(HeaderName::from_static("mcp-protocol-version"), value); + } + (version, headers) +} + fn cache_tools_from_response( cache: &mut HashMap>, message: &ServerJsonRpcMessage, @@ -184,7 +207,10 @@ pub enum StreamableHttpProtocolError { MissingSessionIdInResponse, } -#[allow(clippy::large_enum_variant)] +#[expect( + clippy::large_enum_variant, + reason = "boxing the streaming response would add an allocation to the common response path" +)] #[non_exhaustive] pub enum StreamableHttpPostResponse { Accepted, @@ -651,6 +677,67 @@ impl StreamableHttpClientWorker { Ok(()) } + fn spawn_common_stream( + streams: &mut tokio::task::JoinSet>>, + client: C, + session_id: Arc, + config: &StreamableHttpClientTransportConfig, + protocol_headers: HashMap, + sse_worker_tx: tokio::sync::mpsc::Sender, + transport_task_ct: CancellationToken, + ) { + let uri = config.uri.clone(); + let auth_header = config.auth_header.clone(); + let retry_config = config.retry_config.clone(); + let reconnect_uri = config.uri.clone(); + let reconnect_auth_header = config.auth_header.clone(); + let max_sse_event_size = config.max_sse_event_size; + + streams.spawn(async move { + match client + .get_stream_with_max_sse_event_size( + uri, + session_id.clone(), + None, + auth_header, + protocol_headers.clone(), + max_sse_event_size, + ) + .await + { + Ok(stream) => { + let sse_stream = SseAutoReconnectStream::new( + stream, + StreamableHttpClientReconnect { + client, + session_id, + uri: reconnect_uri, + auth_header: reconnect_auth_header, + custom_headers: protocol_headers, + max_sse_event_size, + }, + retry_config, + ); + Self::execute_sse_stream( + sse_stream, + sse_worker_tx, + false, + transport_task_ct.child_token(), + ) + .await + } + Err(StreamableHttpError::ServerDoesNotSupportSse) => { + tracing::debug!("server doesn't support sse, skip common stream"); + Ok(()) + } + Err(error) => { + tracing::error!("fail to get common stream: {error}"); + Err(error) + } + } + }); + } + /// Performs a transparent re-initialization handshake after a session-expired 404. /// /// Takes an owned clone of the client (avoiding `&self` across `.await` so the @@ -667,8 +754,14 @@ impl StreamableHttpClientWorker { auth_header: Option, custom_headers: HashMap, max_sse_event_size: usize, - ) -> Result<(Option>, HashMap), StreamableHttpError> - { + ) -> Result< + ( + Option>, + ProtocolVersion, + HashMap, + ), + StreamableHttpError, + > { let (init_msg, new_session_id_str) = client .post_message_with_max_sse_event_size( uri.clone(), @@ -712,7 +805,7 @@ impl StreamableHttpClientWorker { .await? .expect_accepted_or_json::()?; - Ok((new_session_id, new_protocol_headers)) + Ok((new_session_id, negotiated_version, new_protocol_headers)) } } @@ -743,17 +836,33 @@ impl Worker for StreamableHttpClientWorker { let _drop_guard = transport_task_ct.clone().drop_guard(); let WorkerSendRequest { responder, - message: initialize_request, + message: startup_request, } = context.recv_from_handler().await?; - let saved_init_request = initialize_request.clone(); + let is_legacy_startup = matches!( + &startup_request, + ClientJsonRpcMessage::Request(request) + if matches!(&request.request, ClientRequest::InitializeRequest(_)) + ); + let mut saved_init_request = is_legacy_startup.then(|| startup_request.clone()); + let empty_tool_cache = HashMap::new(); + let (bootstrap_version, bootstrap_headers) = if is_legacy_startup { + (ProtocolVersion::default(), config.custom_headers.clone()) + } else { + request_version_headers( + &config.custom_headers, + &startup_request, + &ProtocolVersion::default(), + &empty_tool_cache, + ) + }; let (message, session_id) = match self .client .post_message_with_max_sse_event_size( config.uri.clone(), - initialize_request, + startup_request, None, config.auth_header.clone(), - config.custom_headers.clone(), + bootstrap_headers.clone(), config.max_sse_event_size, ) .await @@ -784,8 +893,11 @@ impl Worker for StreamableHttpClientWorker { } None }; - let (negotiated_version, mut protocol_headers) = - negotiate_version_headers(&message, config.custom_headers.clone()); + let (mut negotiated_version, mut protocol_headers) = if is_legacy_startup { + negotiate_version_headers(&message, config.custom_headers.clone()) + } else { + (bootstrap_version, bootstrap_headers) + }; // SEP-2243: tool input schemas (name -> schema) cached from tools/list responses, // used to promote annotated tools/call arguments to Mcp-Param-* headers. let mut tool_header_cache: HashMap> = HashMap::new(); @@ -800,33 +912,37 @@ impl Worker for StreamableHttpClientWorker { }); context.send_to_handler(message).await?; - let initialized_notification = context.recv_from_handler().await?; - // expect a initialized response - let initialized_headers = build_request_headers( - &protocol_headers, - &initialized_notification.message, - &tool_header_cache, - &negotiated_version, - ); - self.client - .post_message_with_max_sse_event_size( - config.uri.clone(), - initialized_notification.message, - session_id.clone(), - config.auth_header.clone(), - initialized_headers, - config.max_sse_event_size, - ) - .await - .map_err(WorkerQuitReason::fatal_context( - "send initialized notification", - ))? - .expect_accepted_or_json::() - .map_err(WorkerQuitReason::fatal_context( - "process initialized notification response", - ))?; - let _ = initialized_notification.responder.send(Ok(())); - #[allow(clippy::large_enum_variant)] + if is_legacy_startup { + let initialized_notification = context.recv_from_handler().await?; + let initialized_headers = build_request_headers( + &protocol_headers, + &initialized_notification.message, + &tool_header_cache, + &negotiated_version, + ); + self.client + .post_message_with_max_sse_event_size( + config.uri.clone(), + initialized_notification.message, + session_id.clone(), + config.auth_header.clone(), + initialized_headers, + config.max_sse_event_size, + ) + .await + .map_err(WorkerQuitReason::fatal_context( + "send initialized notification", + ))? + .expect_accepted_or_json::() + .map_err(WorkerQuitReason::fatal_context( + "process initialized notification response", + ))?; + let _ = initialized_notification.responder.send(Ok(())); + } + #[expect( + clippy::large_enum_variant, + reason = "the event is short-lived and boxing would add allocation in the event loop" + )] enum Event { ClientMessage(WorkerSendRequest), ServerMessage(ServerJsonRpcMessage), @@ -834,63 +950,17 @@ impl Worker for StreamableHttpClientWorker { } let mut streams = tokio::task::JoinSet::new(); let mut pending_stream_response_ids = HashSet::new(); + let mut awaiting_fallback_initialized = false; if let Some(session_id) = &session_id { - let client = self.client.clone(); - let uri = config.uri.clone(); - let session_id = session_id.clone(); - let auth_header = config.auth_header.clone(); - let retry_config = self.config.retry_config.clone(); - let sse_worker_tx = sse_worker_tx.clone(); - let transport_task_ct = transport_task_ct.clone(); - let config_uri = config.uri.clone(); - let config_auth_header = config.auth_header.clone(); - let spawn_headers = protocol_headers.clone(); - let max_sse_event_size = config.max_sse_event_size; - - streams.spawn(async move { - match client - .get_stream_with_max_sse_event_size( - uri.clone(), - session_id.clone(), - None, - auth_header.clone(), - spawn_headers.clone(), - max_sse_event_size, - ) - .await - { - Ok(stream) => { - let sse_stream = SseAutoReconnectStream::new( - stream, - StreamableHttpClientReconnect { - client: client.clone(), - session_id: session_id.clone(), - uri: config_uri, - auth_header: config_auth_header, - custom_headers: spawn_headers, - max_sse_event_size, - }, - retry_config, - ); - Self::execute_sse_stream( - sse_stream, - sse_worker_tx, - false, - transport_task_ct.child_token(), - ) - .await - } - Err(StreamableHttpError::ServerDoesNotSupportSse) => { - tracing::debug!("server doesn't support sse, skip common stream"); - Ok(()) - } - Err(e) => { - // fail to get common stream - tracing::error!("fail to get common stream: {e}"); - Err(e) - } - } - }); + Self::spawn_common_stream( + &mut streams, + self.client.clone(), + session_id.clone(), + &config, + protocol_headers.clone(), + sse_worker_tx.clone(), + transport_task_ct.clone(), + ); } // Main event loop - capture exit reason so we can do cleanup before returning let loop_result: Result<(), WorkerQuitReason> = 'main_loop: loop { @@ -926,16 +996,112 @@ impl Worker for StreamableHttpClientWorker { match event { Event::ClientMessage(send_request) => { let WorkerSendRequest { message, responder } = send_request; + let is_fallback_initialize = saved_init_request.is_none() + && matches!( + &message, + ClientJsonRpcMessage::Request(request) + if matches!( + &request.request, + ClientRequest::InitializeRequest(_) + ) + ); + if is_fallback_initialize { + saved_init_request = Some(message.clone()); + // Servers do not assign sessions to `server/discover`, so a + // fallback initialize starts from a clean slate: no session + // ID, no cleanup state, and no streams to tear down. + debug_assert!( + session_id.is_none() + && session_cleanup_info.is_none() + && streams.is_empty(), + "discover bootstrap must not create session state" + ); + + let response = self + .client + .post_message_with_max_sse_event_size( + config.uri.clone(), + message, + None, + config.auth_header.clone(), + config.custom_headers.clone(), + config.max_sse_event_size, + ) + .await; + let response = match response { + Ok(response) => { + let _ = responder.send(Ok(())); + response + } + Err(error) => { + let _ = responder.send(Err(error)); + continue; + } + }; + let (initialize_response, new_session_id) = response + .expect_initialized::() + .await + .map_err(WorkerQuitReason::fatal_context( + "process fallback initialize response", + ))?; + session_id = new_session_id.map(Arc::from); + if session_id.is_none() && !config.allow_stateless { + return Err(WorkerQuitReason::fatal( + StreamableHttpError::::MissingSessionIdInResponse, + "process fallback initialize response", + )); + } + (negotiated_version, protocol_headers) = negotiate_version_headers( + &initialize_response, + config.custom_headers.clone(), + ); + session_cleanup_info = + session_id.as_ref().map(|session_id| SessionCleanupInfo { + client: self.client.clone(), + uri: config.uri.clone(), + session_id: session_id.clone(), + auth_header: config.auth_header.clone(), + protocol_headers: protocol_headers.clone(), + }); + context.send_to_handler(initialize_response).await?; + awaiting_fallback_initialized = true; + continue; + } + let request_id = Self::client_request_id(&message); + let inline_version = match &message { + ClientJsonRpcMessage::Request(request) => { + request.request.get_meta().protocol_version() + } + _ => None, + }; + let is_initialized_notification = matches!( + &message, + ClientJsonRpcMessage::Notification(notification) + if matches!( + ¬ification.notification, + ClientNotification::InitializedNotification(_) + ) + ); // Pass a clone to the first attempt so `message` is retained for a // potential re-init retry. `post_message` takes ownership and the // trait cannot be changed, so the clone is unavoidable. - let request_headers = build_request_headers( + let (request_version, request_headers) = request_version_headers( &protocol_headers, &message, - &tool_header_cache, &negotiated_version, + &tool_header_cache, ); + if inline_version.is_some() { + negotiated_version = request_version.clone(); + if let Ok(value) = HeaderValue::from_str(request_version.as_str()) { + protocol_headers + .insert(HeaderName::from_static("mcp-protocol-version"), value); + } + if let Some(cleanup) = &mut session_cleanup_info { + cleanup.protocol_headers = protocol_headers.clone(); + } + } let response = self .client .post_message_with_max_sse_event_size( @@ -949,9 +1115,10 @@ impl Worker for StreamableHttpClientWorker { .await; let send_result = match response { Err(StreamableHttpError::SessionExpired) => { - if !config.reinit_on_expired_session { - Err(StreamableHttpError::SessionExpired) - } else { + if let Some(saved_init_request) = saved_init_request + .as_ref() + .filter(|_| config.reinit_on_expired_session) + { // The server discarded the session (HTTP 404). Perform a // fresh handshake once and replay the original message. tracing::info!( @@ -967,7 +1134,11 @@ impl Worker for StreamableHttpClientWorker { ) .await { - Ok((new_session_id, new_protocol_headers)) => { + Ok(( + new_session_id, + new_negotiated_version, + new_protocol_headers, + )) => { // Old streams hold the stale session ID. Stop them first // so no late stale-session messages can arrive after the // pending requests below are completed. @@ -990,6 +1161,7 @@ impl Worker for StreamableHttpClientWorker { .await?; session_id = new_session_id; + negotiated_version = new_negotiated_version; protocol_headers = new_protocol_headers; session_cleanup_info = session_id.as_ref().map(|sid| SessionCleanupInfo { @@ -1001,71 +1173,22 @@ impl Worker for StreamableHttpClientWorker { }); if let Some(new_sid) = &session_id { - let client = self.client.clone(); - let uri = config.uri.clone(); - let new_sid = new_sid.clone(); - let auth_header = config.auth_header.clone(); - let retry_config = self.config.retry_config.clone(); - let sse_tx = sse_worker_tx.clone(); - let task_ct = transport_task_ct.clone(); - let config_uri = config.uri.clone(); - let config_auth = config.auth_header.clone(); - let spawn_headers = protocol_headers.clone(); - let max_sse_event_size = config.max_sse_event_size; - streams.spawn(async move { - match client - .get_stream_with_max_sse_event_size( - uri, - new_sid.clone(), - None, - auth_header.clone(), - spawn_headers.clone(), - max_sse_event_size, - ) - .await - { - Ok(stream) => { - let sse_stream = SseAutoReconnectStream::new( - stream, - StreamableHttpClientReconnect { - client: client.clone(), - session_id: new_sid, - uri: config_uri, - auth_header: config_auth, - custom_headers: spawn_headers, - max_sse_event_size, - }, - retry_config, - ); - Self::execute_sse_stream( - sse_stream, - sse_tx, - false, - task_ct.child_token(), - ) - .await - } - Err(StreamableHttpError::ServerDoesNotSupportSse) => { - tracing::debug!( - "server doesn't support sse after re-init" - ); - Ok(()) - } - Err(e) => { - tracing::error!( - "fail to get common stream after re-init: {e}" - ); - Err(e) - } - } - }); + Self::spawn_common_stream( + &mut streams, + self.client.clone(), + new_sid.clone(), + &config, + protocol_headers.clone(), + sse_worker_tx.clone(), + transport_task_ct.clone(), + ); } - let retry_headers = build_request_headers( + let (_, retry_headers) = request_version_headers( &protocol_headers, &message, - &tool_header_cache, &negotiated_version, + &tool_header_cache, ); let retry_response = self .client @@ -1126,7 +1249,9 @@ impl Worker for StreamableHttpClientWorker { } Err(reinit_err) => Err(reinit_err), } - } // else enable_reinit_on_expired_session + } else { + Err(StreamableHttpError::SessionExpired) + } } Err(e) => Err(e), Ok(StreamableHttpPostResponse::Accepted) => { @@ -1167,6 +1292,23 @@ impl Worker for StreamableHttpClientWorker { Ok(()) } }; + if send_result.is_ok() + && awaiting_fallback_initialized + && is_initialized_notification + { + if let Some(session_id) = &session_id { + Self::spawn_common_stream( + &mut streams, + self.client.clone(), + session_id.clone(), + &config, + protocol_headers.clone(), + sse_worker_tx.clone(), + transport_task_ct.clone(), + ); + } + awaiting_fallback_initialized = false; + } let _ = responder.send(send_result); } Event::ServerMessage(json_rpc_message) => { @@ -1194,7 +1336,6 @@ impl Worker for StreamableHttpClientWorker { // Cleanup session before returning (ensures close() waits for session deletion) // Use a timeout to prevent indefinite hangs if the server is unresponsive if let Some(cleanup) = session_cleanup_info { - const SESSION_CLEANUP_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(5); let cleanup_session_id = cleanup.session_id.clone(); match tokio::time::timeout( SESSION_CLEANUP_TIMEOUT, diff --git a/crates/rmcp/tests/test_client_lifecycle_modes.rs b/crates/rmcp/tests/test_client_lifecycle_modes.rs new file mode 100644 index 000000000..66638c1f2 --- /dev/null +++ b/crates/rmcp/tests/test_client_lifecycle_modes.rs @@ -0,0 +1,312 @@ +#![cfg(all(feature = "client", feature = "server", not(feature = "local")))] + +use rmcp::{ + ClientHandler, ClientLifecycleMode, ClientServiceExt, ServerHandler, ServiceExt, + model::{ + ClientJsonRpcMessage, ClientRequest, DiscoverResult, ErrorCode, ErrorData, GetMeta, + Implementation, InitializeResult, ProtocolVersion, ServerCapabilities, + ServerJsonRpcMessage, ServerResult, + }, + service::PeerRequestOptions, + transport::{IntoTransport, Transport}, +}; + +#[derive(Clone, Default)] +struct DiscoverClient; + +impl ClientHandler for DiscoverClient {} + +#[derive(Clone, Default)] +struct StatelessServer; + +impl ServerHandler for StatelessServer {} + +#[tokio::test] +async fn high_level_server_accepts_discover_startup_without_initialize() { + let (server_transport, client_transport) = tokio::io::duplex(4096); + let server_task = tokio::spawn(async move { + StatelessServer + .serve(server_transport) + .await + .expect("server should accept discover") + }); + + let client = DiscoverClient + .serve_with_lifecycle( + client_transport, + ClientLifecycleMode::Discover { + preferred_versions: vec![ProtocolVersion::V_2026_07_28], + }, + ) + .await + .expect("discover client should start"); + client.list_tools(None).await.expect("list tools"); + client.cancel().await.expect("cancel client"); + let server = server_task.await.expect("server task"); + server.cancel().await.expect("cancel server"); +} + +#[tokio::test] +async fn discover_startup_omits_initialize() { + let (server_transport, client_transport) = tokio::io::duplex(4096); + let mut server = IntoTransport::::into_transport(server_transport); + let server_task = tokio::spawn(async move { + let ClientJsonRpcMessage::Request(request) = + server.receive().await.expect("expected discover request") + else { + panic!("expected request"); + }; + assert!(matches!(request.request, ClientRequest::DiscoverRequest(_))); + let meta = request.request.get_meta(); + assert_eq!(meta.protocol_version(), Some(ProtocolVersion::V_2026_07_28)); + assert!(meta.client_info().is_some()); + assert!(meta.client_capabilities().is_some()); + + server + .send(ServerJsonRpcMessage::response( + ServerResult::DiscoverResult(DiscoverResult::new( + vec![ProtocolVersion::V_2026_07_28], + ServerCapabilities::default(), + Implementation::new("discover-server", "1.0.0"), + )), + request.id, + )) + .await + .expect("send discover response"); + + let ClientJsonRpcMessage::Request(request) = + server.receive().await.expect("expected normal request") + else { + panic!("expected request"); + }; + assert!(!matches!( + request.request, + ClientRequest::InitializeRequest(_) + )); + let meta = request.request.get_meta(); + assert_eq!(meta.protocol_version(), Some(ProtocolVersion::V_2025_11_25)); + assert!(meta.client_info().is_some()); + assert!(meta.client_capabilities().is_some()); + assert_eq!( + meta.get("example.test/extension"), + Some(&serde_json::json!(7)) + ); + server + .send(ServerJsonRpcMessage::response( + ServerResult::ListToolsResult(Default::default()), + request.id, + )) + .await + .expect("send tools response"); + }); + + let client = DiscoverClient + .serve_with_lifecycle( + client_transport, + ClientLifecycleMode::Discover { + preferred_versions: vec![ProtocolVersion::V_2026_07_28], + }, + ) + .await + .expect("discover client should start"); + let mut caller_meta = rmcp::model::RequestMetaObject::new(); + caller_meta.insert("example.test/extension".into(), serde_json::json!(7)); + caller_meta.set_protocol_version(ProtocolVersion::V_2025_11_25); + client + .send_request_with_option( + ClientRequest::ListToolsRequest(rmcp::model::ListToolsRequest { + method: Default::default(), + params: None, + extensions: Default::default(), + }), + PeerRequestOptions::default().with_meta(caller_meta), + ) + .await + .expect("send list tools") + .await_response() + .await + .expect("list tools response"); + client.cancel().await.expect("cancel client"); + server_task.await.expect("server task"); +} + +#[tokio::test] +async fn auto_startup_falls_back_after_discover_method_not_found() { + let (server_transport, client_transport) = tokio::io::duplex(4096); + let mut server = IntoTransport::::into_transport(server_transport); + let server_task = tokio::spawn(async move { + let ClientJsonRpcMessage::Request(discover) = + server.receive().await.expect("expected discover request") + else { + panic!("expected request"); + }; + assert!(matches!( + discover.request, + ClientRequest::DiscoverRequest(_) + )); + server + .send(ServerJsonRpcMessage::error( + ErrorData::new(ErrorCode::METHOD_NOT_FOUND, "Method not found", None), + Some(discover.id), + )) + .await + .expect("send method-not-found"); + + let ClientJsonRpcMessage::Request(initialize) = + server.receive().await.expect("expected initialize request") + else { + panic!("expected request"); + }; + assert!(matches!( + initialize.request, + ClientRequest::InitializeRequest(_) + )); + server + .send(ServerJsonRpcMessage::response( + ServerResult::InitializeResult( + InitializeResult::new(ServerCapabilities::default()), + ), + initialize.id, + )) + .await + .expect("send initialize response"); + assert!(matches!( + server.receive().await, + Some(ClientJsonRpcMessage::Notification(_)) + )); + }); + + let client = DiscoverClient + .serve_with_lifecycle( + client_transport, + ClientLifecycleMode::Auto { + preferred_versions: vec![ProtocolVersion::V_2026_07_28], + legacy_version: Some(ProtocolVersion::V_2025_11_25), + }, + ) + .await + .expect("auto client should fall back"); + client.cancel().await.expect("cancel client"); + server_task.await.expect("server task"); +} + +#[tokio::test] +async fn discover_startup_retries_a_mutually_supported_version() { + let unsupported: ProtocolVersion = + serde_json::from_value(serde_json::json!("2099-01-01")).unwrap(); + let (server_transport, client_transport) = tokio::io::duplex(4096); + let mut server = IntoTransport::::into_transport(server_transport); + let server_task = tokio::spawn(async move { + let ClientJsonRpcMessage::Request(first) = + server.receive().await.expect("expected first discover") + else { + panic!("expected request"); + }; + assert_eq!( + first.request.get_meta().protocol_version(), + Some(unsupported.clone()) + ); + server + .send(ServerJsonRpcMessage::error( + ErrorData::unsupported_protocol_version( + unsupported, + &[ProtocolVersion::V_2026_07_28], + ), + Some(first.id), + )) + .await + .expect("send unsupported error"); + + let ClientJsonRpcMessage::Request(second) = + server.receive().await.expect("expected retry discover") + else { + panic!("expected request"); + }; + assert_eq!( + second.request.get_meta().protocol_version(), + Some(ProtocolVersion::V_2026_07_28) + ); + server + .send(ServerJsonRpcMessage::response( + ServerResult::DiscoverResult(DiscoverResult::new( + vec![ProtocolVersion::V_2026_07_28], + ServerCapabilities::default(), + Implementation::new("discover-server", "1.0.0"), + )), + second.id, + )) + .await + .expect("send discover response"); + }); + + let client = DiscoverClient + .serve_with_lifecycle( + client_transport, + ClientLifecycleMode::Discover { + preferred_versions: vec![ + serde_json::from_value(serde_json::json!("2099-01-01")).unwrap(), + ProtocolVersion::V_2026_07_28, + ], + }, + ) + .await + .expect("discover client should retry"); + client.cancel().await.expect("cancel client"); + server_task.await.expect("server task"); +} + +#[tokio::test] +async fn discover_startup_retries_current_version_once_when_server_reports_it_supported() { + let (server_transport, client_transport) = tokio::io::duplex(4096); + let mut server = IntoTransport::::into_transport(server_transport); + let server_task = tokio::spawn(async move { + let ClientJsonRpcMessage::Request(first) = + server.receive().await.expect("expected first discover") + else { + panic!("expected request"); + }; + server + .send(ServerJsonRpcMessage::error( + ErrorData::unsupported_protocol_version( + ProtocolVersion::V_2026_07_28, + &[ProtocolVersion::V_2026_07_28], + ), + Some(first.id), + )) + .await + .expect("send unsupported error"); + + let ClientJsonRpcMessage::Request(second) = + server.receive().await.expect("expected retry discover") + else { + panic!("expected request"); + }; + assert_eq!( + second.request.get_meta().protocol_version(), + Some(ProtocolVersion::V_2026_07_28) + ); + server + .send(ServerJsonRpcMessage::response( + ServerResult::DiscoverResult(DiscoverResult::new( + vec![ProtocolVersion::V_2026_07_28], + ServerCapabilities::default(), + Implementation::new("discover-server", "1.0.0"), + )), + second.id, + )) + .await + .expect("send discover response"); + }); + + let client = DiscoverClient + .serve_with_lifecycle( + client_transport, + ClientLifecycleMode::Discover { + preferred_versions: vec![ProtocolVersion::V_2026_07_28], + }, + ) + .await + .expect("discover client should retry once"); + client.cancel().await.expect("cancel client"); + server_task.await.expect("server task"); +} diff --git a/crates/rmcp/tests/test_discover_http_client_startup.rs b/crates/rmcp/tests/test_discover_http_client_startup.rs new file mode 100644 index 000000000..6ce1dbe40 --- /dev/null +++ b/crates/rmcp/tests/test_discover_http_client_startup.rs @@ -0,0 +1,137 @@ +#![cfg(all( + not(feature = "local"), + feature = "client", + feature = "reqwest", + feature = "transport-streamable-http-server" +))] + +use std::borrow::Cow; + +use rmcp::{ + ClientLifecycleMode, ClientServiceExt, ServerHandler, + model::{ClientInfo, DiscoverResult, ErrorCode, ErrorData, ProtocolVersion}, + service::{MaybeSendFuture, RequestContext, RoleServer}, + transport::{ + StreamableHttpClientTransport, + streamable_http_client::StreamableHttpClientTransportConfig, + streamable_http_server::{ + StreamableHttpServerConfig, StreamableHttpService, session::local::LocalSessionManager, + }, + }, +}; +use tokio_util::sync::CancellationToken; + +#[derive(Clone, Default)] +struct DiscoverHttpServer; + +impl ServerHandler for DiscoverHttpServer { + fn supported_protocol_versions(&self) -> Cow<'static, [ProtocolVersion]> { + Cow::Borrowed(&[ProtocolVersion::V_2026_07_28]) + } +} + +#[derive(Clone, Default)] +struct LegacyHttpServer; + +impl ServerHandler for LegacyHttpServer { + fn discover( + &self, + _context: RequestContext, + ) -> impl Future> + MaybeSendFuture + '_ { + std::future::ready(Err(ErrorData::new( + ErrorCode::METHOD_NOT_FOUND, + "Method not found", + None, + ))) + } +} + +#[tokio::test] +async fn discover_http_client_bootstraps_headers_without_initialize() { + let ct = CancellationToken::new(); + let service: StreamableHttpService = + StreamableHttpService::new( + || Ok(DiscoverHttpServer), + Default::default(), + StreamableHttpServerConfig::default() + .with_stateful_mode(false) + .with_json_response(true) + .with_cancellation_token(ct.child_token()), + ); + let router = axum::Router::new().nest_service("/mcp", service); + let listener = tokio::net::TcpListener::bind("127.0.0.1:0") + .await + .expect("listener should bind"); + let address = listener.local_addr().expect("listener address"); + let server = tokio::spawn({ + let ct = ct.clone(); + async move { + let _ = axum::serve(listener, router) + .with_graceful_shutdown(async move { ct.cancelled_owned().await }) + .await; + } + }); + + let transport = StreamableHttpClientTransport::from_config( + StreamableHttpClientTransportConfig::with_uri(format!("http://{address}/mcp")), + ); + let client = ClientInfo::default() + .serve_with_lifecycle( + transport, + ClientLifecycleMode::Discover { + preferred_versions: vec![ProtocolVersion::V_2026_07_28], + }, + ) + .await + .expect("discover HTTP client should start"); + client.list_tools(None).await.expect("list tools"); + client.cancel().await.expect("cancel client"); + + ct.cancel(); + server.await.expect("server task"); +} + +#[tokio::test] +async fn auto_http_client_falls_back_to_stateful_legacy_startup() { + let ct = CancellationToken::new(); + let service: StreamableHttpService = + StreamableHttpService::new( + || Ok(LegacyHttpServer), + Default::default(), + StreamableHttpServerConfig::default() + .with_json_response(true) + .with_cancellation_token(ct.child_token()), + ); + let router = axum::Router::new().nest_service("/mcp", service); + let listener = tokio::net::TcpListener::bind("127.0.0.1:0") + .await + .expect("listener should bind"); + let address = listener.local_addr().expect("listener address"); + let server = tokio::spawn({ + let ct = ct.clone(); + async move { + let _ = axum::serve(listener, router) + .with_graceful_shutdown(async move { ct.cancelled_owned().await }) + .await; + } + }); + + let transport = StreamableHttpClientTransport::from_config( + StreamableHttpClientTransportConfig::with_uri(format!("http://{address}/mcp")), + ); + let client = ClientInfo::default() + .serve_with_lifecycle( + transport, + ClientLifecycleMode::Auto { + preferred_versions: vec![ProtocolVersion::V_2026_07_28], + legacy_version: Some(ProtocolVersion::V_2025_11_25), + }, + ) + .await + .expect("auto HTTP client should fall back"); + client.list_tools(None).await.expect("list tools"); + client.cancel().await.expect("cancel client"); + + ct.cancel(); + server.await.expect("server task"); +} diff --git a/crates/rmcp/tests/test_meta_helpers.rs b/crates/rmcp/tests/test_meta_helpers.rs index a45c6a4e2..5f9128a42 100644 --- a/crates/rmcp/tests/test_meta_helpers.rs +++ b/crates/rmcp/tests/test_meta_helpers.rs @@ -30,6 +30,29 @@ fn meta_setters_store_sep_2575_values() { assert_eq!(meta.get(META_KEY_LOG_LEVEL), Some(&json!("warning"))); } +#[test] +fn with_client_context_sets_all_required_sep_2575_fields() { + let meta = RequestMetaObject::with_client_context( + ProtocolVersion::V_2026_07_28, + Implementation::new("test-client", "1.0.0"), + ClientCapabilities::default(), + ); + + assert!( + meta.missing_required_keys(&ProtocolVersion::V_2026_07_28) + .is_empty() + ); + assert_eq!(meta.protocol_version(), Some(ProtocolVersion::V_2026_07_28)); + assert_eq!( + meta.client_info(), + Some(Implementation::new("test-client", "1.0.0")) + ); + assert_eq!( + meta.client_capabilities(), + Some(ClientCapabilities::default()) + ); +} + #[test] fn meta_accessors_decode_wire_values() { let meta: RequestMetaObject = serde_json::from_value(json!({ diff --git a/crates/rmcp/tests/test_stateless_server_requests.rs b/crates/rmcp/tests/test_stateless_server_requests.rs new file mode 100644 index 000000000..4404e08cf --- /dev/null +++ b/crates/rmcp/tests/test_stateless_server_requests.rs @@ -0,0 +1,196 @@ +#![cfg(all(feature = "server", not(feature = "local")))] + +use std::sync::{Arc, Mutex}; + +use rmcp::{ + ServerHandler, ServiceExt, + model::{ + ClientCapabilities, ClientJsonRpcMessage, ClientRequest, DiscoverRequest, + DiscoverRequestParams, ErrorCode, ErrorData, Implementation, ListToolsRequest, + ListToolsResult, PaginatedRequestParams, ProtocolVersion, RequestId, RequestMetaObject, + ServerJsonRpcMessage, + }, + service::{MaybeSendFuture, RequestContext, RoleServer, ServerInitializeError}, + transport::{IntoTransport, Transport}, +}; + +#[derive(Clone, Default)] +struct StatelessServer; + +impl ServerHandler for StatelessServer {} + +fn complete_meta() -> RequestMetaObject { + complete_meta_for("stateless-client") +} + +fn complete_meta_for(client_name: &str) -> RequestMetaObject { + let mut meta = RequestMetaObject::new(); + meta.set_protocol_version(ProtocolVersion::V_2026_07_28); + meta.set_client_info(Implementation::new(client_name, "1.0.0")); + meta.set_client_capabilities(ClientCapabilities::default()); + meta +} + +fn list_tools_request(meta: RequestMetaObject) -> ClientJsonRpcMessage { + let mut request = ListToolsRequest { + method: Default::default(), + params: None, + extensions: Default::default(), + }; + request.extensions.insert(meta); + ClientJsonRpcMessage::request( + ClientRequest::ListToolsRequest(request), + RequestId::Number(1), + ) +} + +#[tokio::test] +async fn stateless_server_rejects_missing_metadata_on_every_request() { + let (server_transport, client_transport) = tokio::io::duplex(4096); + let server_task = tokio::spawn(async move { + StatelessServer + .serve(server_transport) + .await + .expect("server should start") + }); + let mut client = IntoTransport::::into_transport(client_transport); + + let mut discover = DiscoverRequest::new(DiscoverRequestParams {}); + discover.extensions.insert(complete_meta()); + client + .send(ClientJsonRpcMessage::request( + ClientRequest::DiscoverRequest(discover), + RequestId::Number(1), + )) + .await + .expect("send discover"); + assert!(matches!( + client.receive().await, + Some(ServerJsonRpcMessage::Response(_)) + )); + + client + .send(ClientJsonRpcMessage::request( + ClientRequest::ListToolsRequest(ListToolsRequest { + method: Default::default(), + params: None, + extensions: Default::default(), + }), + RequestId::Number(2), + )) + .await + .expect("send list tools"); + let Some(ServerJsonRpcMessage::Error(error)) = client.receive().await else { + panic!("expected invalid params"); + }; + assert_eq!(error.error.code, ErrorCode::INVALID_PARAMS); + + server_task + .await + .expect("server task") + .cancel() + .await + .expect("cancel server"); +} + +#[derive(Clone)] +struct ContextServer { + seen_clients: Arc>>, +} + +impl ServerHandler for ContextServer { + fn list_tools( + &self, + _request: Option, + context: RequestContext, + ) -> impl Future> + MaybeSendFuture + '_ { + let seen_clients = self.seen_clients.clone(); + async move { + seen_clients + .lock() + .expect("seen clients lock") + .push(context.client_info().expect("current client info").name); + Ok(ListToolsResult::default()) + } + } +} + +#[tokio::test] +async fn stateless_server_uses_each_requests_client_context() { + let seen_clients = Arc::new(Mutex::new(Vec::new())); + let handler = ContextServer { + seen_clients: seen_clients.clone(), + }; + let (server_transport, client_transport) = tokio::io::duplex(4096); + let server_task = tokio::spawn(async move { + handler + .serve(server_transport) + .await + .expect("server should start") + }); + let mut client = IntoTransport::::into_transport(client_transport); + + client + .send(list_tools_request(complete_meta_for("first-client"))) + .await + .expect("send first request"); + assert!(matches!( + client.receive().await, + Some(ServerJsonRpcMessage::Response(_)) + )); + + let mut second = list_tools_request(complete_meta_for("second-client")); + if let ClientJsonRpcMessage::Request(request) = &mut second { + request.id = RequestId::Number(2); + } + client.send(second).await.expect("send second request"); + assert!(matches!( + client.receive().await, + Some(ServerJsonRpcMessage::Response(_)) + )); + + assert_eq!( + *seen_clients.lock().expect("seen clients lock"), + ["first-client", "second-client"] + ); + server_task + .await + .expect("server task") + .cancel() + .await + .expect("cancel server"); +} + +#[tokio::test] +async fn stateless_server_rejects_malformed_metadata_opener() { + let (server_transport, client_transport) = tokio::io::duplex(4096); + let server_task = tokio::spawn(async move { StatelessServer.serve(server_transport).await }); + let mut client = IntoTransport::::into_transport(client_transport); + + let mut request = ListToolsRequest { + method: Default::default(), + params: None, + extensions: Default::default(), + }; + let malformed: RequestMetaObject = serde_json::from_value(serde_json::json!({ + "io.modelcontextprotocol/protocolVersion": "2026-07-28", + "io.modelcontextprotocol/clientInfo": "wrong", + "io.modelcontextprotocol/clientCapabilities": null + })) + .unwrap(); + request.extensions.insert(malformed); + client + .send(ClientJsonRpcMessage::request( + ClientRequest::ListToolsRequest(request), + RequestId::Number(1), + )) + .await + .expect("send list tools"); + let Err(error) = server_task.await.expect("server task") else { + panic!("malformed opener should not start a session"); + }; + assert!(matches!( + error, + ServerInitializeError::ExpectedInitializeRequest(Some(_)) + )); +} diff --git a/crates/rmcp/tests/test_streamable_http_stale_session.rs b/crates/rmcp/tests/test_streamable_http_stale_session.rs index 9aa0309ad..137460f7b 100644 --- a/crates/rmcp/tests/test_streamable_http_stale_session.rs +++ b/crates/rmcp/tests/test_streamable_http_stale_session.rs @@ -16,7 +16,7 @@ use rmcp::{ ServiceError, ServiceExt, model::{ CallToolRequestParams, ClientInfo, ClientJsonRpcMessage, ClientRequest, ErrorCode, - ErrorData, InitializeResult, PingRequest, RequestId, ServerCapabilities, + ErrorData, InitializeResult, PingRequest, ProtocolVersion, RequestId, ServerCapabilities, ServerJsonRpcMessage, ServerResult, }, transport::{ @@ -90,7 +90,7 @@ impl StreamableHttpClient for ReinitDropsAcceptedResponseClient { message: ClientJsonRpcMessage, _session_id: Option>, _auth_header: Option, - _custom_headers: HashMap, + custom_headers: HashMap, ) -> Result> { let mut state = self.state.lock().await; match state @@ -100,15 +100,23 @@ impl StreamableHttpClient for ReinitDropsAcceptedResponseClient { { MockPost::Initialize => { state.session_counter += 1; + let protocol_version = if state.session_counter == 1 { + ProtocolVersion::V_2025_11_25 + } else { + ProtocolVersion::V_2025_06_18 + }; let id = match message { ClientJsonRpcMessage::Request(request) => request.id, other => panic!("expected initialize request, got {other:?}"), }; Ok(StreamableHttpPostResponse::Json( ServerJsonRpcMessage::response( - ServerResult::InitializeResult(InitializeResult::new( - ServerCapabilities::builder().enable_tools().build(), - )), + ServerResult::InitializeResult( + InitializeResult::new( + ServerCapabilities::builder().enable_tools().build(), + ) + .with_protocol_version(protocol_version), + ), id, ), Some(format!("session-{}", state.session_counter)), @@ -123,6 +131,12 @@ impl StreamableHttpClient for ReinitDropsAcceptedResponseClient { } MockPost::Accepted => { if state.posts.is_empty() { + assert_eq!( + custom_headers + .get(&HeaderName::from_static("mcp-protocol-version")) + .and_then(|value| value.to_str().ok()), + Some(ProtocolVersion::V_2025_06_18.as_str()) + ); self.final_retry_accepted.add_permits(1); } else { self.initial_request_accepted.add_permits(1); diff --git a/docs/readme/README.zh-cn.md b/docs/readme/README.zh-cn.md index 84393849a..55518f095 100644 --- a/docs/readme/README.zh-cn.md +++ b/docs/readme/README.zh-cn.md @@ -74,6 +74,41 @@ async fn main() -> Result<(), Box> { ```
+### 客户端生命周期模式 + +`serve()` 使用传统 MCP 生命周期：客户端发送 `initialize`，接收协商后的服务端信息， +然后发送 `notifications/initialized`。如需显式选择其他生命周期，请使用 +[`ClientServiceExt::serve_with_lifecycle`](../../crates/rmcp/src/service/client.rs)： + +```rust, ignore +use rmcp::{ClientInfo, ClientLifecycleMode, ClientServiceExt, ProtocolVersion}; + +// 直接通过 server/discover 启动，并在每个请求中携带客户端元数据。 +let client = ClientInfo::default() + .serve_with_lifecycle( + transport, + ClientLifecycleMode::Discover { + preferred_versions: vec![ProtocolVersion::V_2026_07_28], + }, + ) + .await?; + +// 或先尝试发现生命周期；当传统服务端报告未实现 server/discover 时回退。 +let client = ClientInfo::default() + .serve_with_lifecycle( + transport, + ClientLifecycleMode::Auto { + preferred_versions: vec![ProtocolVersion::V_2026_07_28], + legacy_version: Some(ProtocolVersion::V_2025_11_25), + }, + ) + .await?; +``` + +`ClientLifecycleMode::Initialize` 等同于现有的 `serve()` 行为。发现启动不会发送 +`notifications/initialized`；发现过程即完成启动，后续每个请求都会在 `_meta` +中携带协议版本、客户端信息和客户端能力。 + ### 构建服务端
@@ -804,10 +839,11 @@ impl ServerHandler for MyServer { ### 初始化通知 -客户端在握手完成后发送 `initialized` 通知： +传统客户端在 `initialize` 握手完成后发送 `initialized` 通知。 +使用 `ClientLifecycleMode::Discover` 的客户端不会发送此通知： ```rust -// 在 serve() 握手过程中由 rmcp 自动发送。 +// 在传统 serve() 握手过程中由 rmcp 自动发送。 // 服务端通过以下方式处理： impl ServerHandler for MyServer { async fn on_initialized(