From f1f87d751a3beb20ecd93c2742b8a61af35d38a4 Mon Sep 17 00:00:00 2001 From: daniel Date: Thu, 15 Oct 2026 11:41:27 +0100 Subject: [PATCH 1/3] feat(mcp)!: support acknowledged notification subscriptions Add dual-era connection lifecycle, owned subscription filters and resource restoration, and fenced catalog refresh APIs. Upgrade rmcp to 3.5.1 and bundle MCP and umbrella crate 0.11.0 versions. BREAKING CHANGE: McpConnection::capabilities now returns an owned snapshot so callers observe current capabilities after authentication reconnect. --- Cargo.lock | 17 +- book/src/ch17-mcp.md | 40 +- crates/agentkit-integration-tests/Cargo.toml | 4 +- .../src/http_mcp_server.rs | 20 +- crates/agentkit-mcp/Cargo.toml | 6 +- crates/agentkit-mcp/README.md | 45 ++ crates/agentkit-mcp/src/lib.rs | 750 ++++++++++++++---- crates/agentkit-mcp/src/subscriptions.rs | 493 ++++++++++++ crates/agentkit-mcp/tests/catalog_refresh.rs | 303 +++++++ crates/agentkit-mcp/tests/modern_http.rs | 433 ++++++++++ crates/agentkit-mcp/tests/subscriptions.rs | 236 ++++++ crates/agentkit/Cargo.toml | 4 +- docs/mcp.md | 47 +- 13 files changed, 2218 insertions(+), 180 deletions(-) create mode 100644 crates/agentkit-mcp/src/subscriptions.rs create mode 100644 crates/agentkit-mcp/tests/catalog_refresh.rs create mode 100644 crates/agentkit-mcp/tests/modern_http.rs create mode 100644 crates/agentkit-mcp/tests/subscriptions.rs diff --git a/Cargo.lock b/Cargo.lock index 3b6516d..3f07f15 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -20,7 +20,7 @@ dependencies = [ [[package]] name = "agentkit" -version = "0.10.12" +version = "0.11.0" dependencies = [ "agentkit-acp", "agentkit-adapter-completions", @@ -233,7 +233,7 @@ dependencies = [ [[package]] name = "agentkit-mcp" -version = "0.10.7" +version = "0.11.0" dependencies = [ "agentkit-capabilities", "agentkit-core", @@ -2703,9 +2703,9 @@ dependencies = [ [[package]] name = "process-wrap" -version = "9.1.0" +version = "10.0.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2e842efad9119158434d193c6682e2ebee4b44d6ad801d7b349623b3f57cdf55" +checksum = "1f21b97672d2dc848e7b25701ab4535618b92f4861c13cc3f7f7bed52ad3c8da" dependencies = [ "futures", "indexmap 2.14.0", @@ -3006,9 +3006,9 @@ dependencies = [ [[package]] name = "rmcp" -version = "3.1.2" +version = "3.5.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c8dddc5b1924b9a59fba420166160ca2c4663a4e01803e52eda33070f56d63c8" +checksum = "ab5023804cbe0304cbe06197ac16e645a0a0d0ffa2953bf735c15c2fa0213350" dependencies = [ "async-trait", "base64 0.23.1", @@ -3018,6 +3018,7 @@ dependencies = [ "http", "http-body", "http-body-util", + "indexmap 2.14.0", "oauth2", "pastey", "pin-project-lite", @@ -3041,9 +3042,9 @@ dependencies = [ [[package]] name = "rmcp-macros" -version = "3.1.2" +version = "3.5.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6898e24cd16342b59bfa8a53c2c04b9cf62fc8a2cfea57b9c038b09984bfc521" +checksum = "a1da498700030856ba63114bceefa30ba21f5e278cb975e5b9de3529fe842b40" dependencies = [ "darling 0.24.0", "proc-macro2", diff --git a/book/src/ch17-mcp.md b/book/src/ch17-mcp.md index 58200d8..b377280 100644 --- a/book/src/ch17-mcp.md +++ b/book/src/ch17-mcp.md @@ -32,6 +32,40 @@ The MCP specification moves quickly. Because agentkit-mcp re-exports [`rmcp::mod The same logic applies to transports: stdio and Streamable HTTP are rmcp implementations driven declaratively here; future rmcp transports are reachable through `McpConnection::from_running_service_with_events` without touching `McpTransportBinding`. +## Modern subscriptions and catalog refresh + +Built-in transports negotiate with RMCP's `Auto` lifecycle (`server/discover` +with legacy initialization fallback). Check `protocol_version()` and +`capabilities()` on the connection instead of inferring support from HTTP or +stdio. Adopting an existing service with +`McpConnection::from_running_service_with_events` does not perform asynchronous +setup: call `discover().await` to establish its modern notification stream. + +Modern discovery waits for a `subscriptions/listen` acknowledgment. +`subscription_status().await` reports the requested and accepted filters, and +`subscription_changes().await` watches their lifecycle. Partial catalog +acknowledgments are visible rather than treated as full acceptance. Resource +subscribe/unsubscribe operations replace the complete filter; a rejected URI +set leaves the previous active stream in place. Legacy peers keep the separate +`resources/subscribe` and `resources/unsubscribe` methods. + +Stream notifications use the existing `subscribe_events()` broadcast. For catalog +reconciliation, use `catalog_changes()` and `catalog_revision()`; watch updates +coalesce and survive slow broadcast consumers. The host must schedule +`refresh_changed_catalogs()`, or prepare/fetch/publish refresh work outside its +manager lock. Invalidated, failed, or canceled work must not publish a stale +snapshot. + +A transport failure triggers bounded subscription recovery and dirties the +catalog. Graceful completion yields `Closed { graceful: true, .. }` without +retry. Authentication reconnect rebuilds subscriptions on the new peer and +retains committed resource URIs; `connection_epoch()` fences old work. Await +`close()` to shut down streams before transport cleanup. + +**Migration (agentkit-mcp 0.11.0):** `capabilities()` returns an owned snapshot +instead of a borrowed reference. Read it again after authentication reconnect +for the current peer capabilities; do not retain a lock guard across awaits. + ## MCP in the capability model MCP servers expose three capability types, which map directly to agentkit's capability layer: @@ -132,7 +166,7 @@ while let Ok(event) = events.recv().await { `set_logging_level(LoggingLevel::Info)` negotiates the minimum severity the server emits. `subscribe_resource` / `unsubscribe_resource` toggle per-URI watch on resources that support it. -List-changed events are _also_ delivered through the legacy `McpServerNotification` mpsc receiver that `McpServerManager::refresh_changed_catalogs` drains to re-run discovery — the two channels coexist: events for live UI/observability, mpsc for catalog re-sync. +List-changed events also advance the coalescing catalog revision watch used by `McpServerManager::refresh_changed_catalogs`. The compatibility `McpServerNotification` receiver remains available; catalog reconciliation does not require draining the event broadcast. ## Auth handled outside the loop @@ -163,11 +197,11 @@ The manager owns server lifecycle: | `connect_server(id)` / `connect_all()` | Open a connection and run discovery | | `connect_all_settled()` | Attempt all registered servers concurrently, install successes, and return per-server failures | | `refresh_server(id)` | Re-run discovery and emit per-tool/resource/prompt diff events | -| `refresh_changed_catalogs()` | Drain pending `*list_changed` notifications and refresh affected catalogs | +| `refresh_changed_catalogs()` | Reconcile catalogs marked dirty by coalesced invalidations | | `disconnect_server(id)` | Close the connection, drop tools from the federated catalog, emit `ServerDisconnected` | | `subscribe_catalog_events()` | Broadcast receiver for `McpCatalogEvent` (server connect / disconnect / tool added / removed / changed / refresh failed / auth changed) | -Register a server with `with_server_options(config, McpServerOptions::new().with_timeout(duration))` to bound connecting to that server. The timeout covers connection establishment — transport setup and the MCP initialize handshake — plus initial discovery, and applies to refresh discovery on its own; failures surface as `McpError::Timeout`. +Register a server with `with_server_options(config, McpServerOptions::new().with_timeout(duration))` to bound connecting to that server. The timeout covers connection establishment — transport setup and MCP lifecycle negotiation — plus initial discovery, and applies to refresh discovery on its own; failures surface as `McpError::Timeout`. ## Federating MCP into the agent diff --git a/crates/agentkit-integration-tests/Cargo.toml b/crates/agentkit-integration-tests/Cargo.toml index 3ac01a4..e24777c 100644 --- a/crates/agentkit-integration-tests/Cargo.toml +++ b/crates/agentkit-integration-tests/Cargo.toml @@ -14,7 +14,7 @@ doctest = false [dependencies] agentkit-core = { version = "0.10.5", path = "../agentkit-core" } agentkit-loop = { version = "0.10.5", path = "../agentkit-loop" } -agentkit-mcp = { version = "0.10.5", path = "../agentkit-mcp" } +agentkit-mcp = { version = "0.11.0", path = "../agentkit-mcp" } agentkit-task-manager = { version = "0.10.5", path = "../agentkit-task-manager" } agentkit-tool-compose = { version = "0.10.5", path = "../agentkit-tool-compose" } agentkit-tools-core = { version = "0.10.5", path = "../agentkit-tools-core", features = ["schemars"] } @@ -23,7 +23,7 @@ async-trait.workspace = true schemars = { workspace = true, features = ["derive"] } axum = { version = "0.8", default-features = false, features = ["http1", "tokio"] } pretty_assertions = "1.4" -rmcp = { version = "3.1.2", default-features = false, features = ["client", "server", "transport-async-rw", "transport-streamable-http-server", "macros"] } +rmcp = { version = "=3.5.1", default-features = false, features = ["client", "server", "transport-async-rw", "transport-streamable-http-server", "macros"] } ron = "0.8" serde.workspace = true serde_json.workspace = true diff --git a/crates/agentkit-integration-tests/src/http_mcp_server.rs b/crates/agentkit-integration-tests/src/http_mcp_server.rs index bbb4985..37c7db7 100644 --- a/crates/agentkit-integration-tests/src/http_mcp_server.rs +++ b/crates/agentkit-integration-tests/src/http_mcp_server.rs @@ -14,8 +14,9 @@ use rmcp::{ RoleServer, ServerHandler, model::{ CallToolRequestParams, CallToolResponse, CallToolResult, ContentBlock as Content, - ErrorData as RmcpError, InitializeRequestParams, InitializeResult, ListToolsResult, - PaginatedRequestParams, ServerCapabilities, ServerInfo, Tool, + DiscoverResult, ErrorData as RmcpError, InitializeRequestParams, InitializeResult, + ListToolsResult, PaginatedRequestParams, ProtocolVersion, ServerCapabilities, ServerConfig, + Tool, }, service::{NotificationContext, Peer, RequestContext}, transport::streamable_http_server::{ @@ -49,13 +50,24 @@ impl MutableMcpServer { } impl ServerHandler for MutableMcpServer { - fn get_info(&self) -> ServerInfo { - ServerInfo::new( + fn get_info(&self) -> ServerConfig { + ServerConfig::new( ServerCapabilities::builder() .enable_tools() .enable_tool_list_changed() .build(), ) + .with_protocol_version(ProtocolVersion::V_2025_11_25) + } + + async fn discover( + &self, + _context: RequestContext, + ) -> Result { + // This fixture deliberately exercises the initialize-era protocol. + Err(RmcpError::method_not_found::< + rmcp::model::DiscoverRequestMethod, + >()) } async fn initialize( diff --git a/crates/agentkit-mcp/Cargo.toml b/crates/agentkit-mcp/Cargo.toml index d754d96..9e7aa13 100644 --- a/crates/agentkit-mcp/Cargo.toml +++ b/crates/agentkit-mcp/Cargo.toml @@ -4,7 +4,7 @@ homepage.workspace = true name = "agentkit-mcp" readme = "README.md" repository.workspace = true -version = "0.10.7" +version = "0.11.0" edition.workspace = true license.workspace = true rust-version.workspace = true @@ -20,7 +20,7 @@ async-trait.workspace = true futures-util.workspace = true http = "1" reqwest.workspace = true -rmcp = { version = "3.1.2", default-features = false, features = ["client", "transport-child-process", "transport-streamable-http-client-reqwest", "reqwest", "auth"] } +rmcp = { version = "=3.5.1", default-features = false, features = ["client", "transport-child-process", "transport-streamable-http-client-reqwest", "reqwest", "auth"] } serde.workspace = true serde_json.workspace = true sse-stream = "0.2" @@ -30,5 +30,5 @@ tokio = { workspace = true, features = ["sync", "time"] } [dev-dependencies] axum = "0.8" -rmcp = { version = "3.1.2", default-features = false, features = ["client", "server", "transport-async-rw", "macros", "elicitation"] } +rmcp = { version = "=3.5.1", default-features = false, features = ["client", "server", "transport-async-rw", "macros", "elicitation"] } tokio = { workspace = true, features = ["macros", "rt-multi-thread", "io-util", "time", "net"] } diff --git a/crates/agentkit-mcp/README.md b/crates/agentkit-mcp/README.md index dc5fe66..1b17562 100644 --- a/crates/agentkit-mcp/README.md +++ b/crates/agentkit-mcp/README.md @@ -69,6 +69,51 @@ Use `connect_all_settled().await` when startup should be best effort: it attempt For slow or unresponsive servers, register with `with_server_options(config, McpServerOptions::new().with_timeout(duration))`. The timeout bounds connection establishment — transport setup and the MCP initialize handshake — together with initial discovery, and bounds refresh discovery on its own; exceeding it returns `McpError::Timeout`. +## Negotiation and notification subscriptions + +Built-in stdio and Streamable HTTP connections use RMCP's `Auto` lifecycle: probe +`server/discover`, then fall back to legacy initialization when needed. Inspect +`connection.protocol_version()` and the cloned `connection.capabilities()` for +the negotiated peer, rather than assuming a protocol from the transport. + +For modern peers, discovery establishes a connection-owned `subscriptions/listen` +stream and waits for its acknowledgment. The requested filter covers advertised +tool, resource, and prompt catalogs. `subscription_status().await` exposes both +the requested and acknowledged filters: a server may accept only part of the +catalog filter. `subscription_changes().await` returns a watch receiver once a +modern stream exists; legacy peers return `None`. Adopting an existing RMCP +service with `from_running_service_with_events` is synchronous: call +`discover().await` to establish its modern subscription lazily. + +`subscribe_resource(uri)` and `unsubscribe_resource(uri)` replace the modern +stream's complete filter. A resource update is committed only after the server +acknowledges the requested URI set; rejection preserves the previous active +filter and stream. Legacy peers retain `resources/subscribe` and +`resources/unsubscribe`. Streamed notifications feed the same +`subscribe_events()` broadcast used by legacy notifications. + +Catalog invalidations also advance `catalog_revision()` and wake +`catalog_changes()`, a coalescing watch channel that does not require draining +broadcast events. Watch wakeups are invalidations, not newly published catalogs. +Hosts should call `McpServerManager::refresh_changed_catalogs()` to reconcile +snapshots, or use `prepare_changed_catalogs()`, fetch work outside their manager +lock, and `publish_catalog_refresh(...)` to publish it. Revision and connection +epoch checks fence obsolete results; canceled or failed fetches remain dirty. +A host must schedule this refresh work itself; event delivery does not silently +refresh the host's tool registry. + +On transport failure, the modern subscription worker retries with bounded +backoff and marks catalogs dirty for reconciliation. Graceful server completion +is observable as `Closed { graceful: true, .. }` and is not retried. Authentication +reconnect replaces the subscription on the new peer, retaining committed resource +URIs. `connection_epoch()` changes when the transport is replaced or closed. +Call `close().await` to stop owned streams before transport cleanup; dropping the +connection also cancels its workers. + +**Migration (agentkit-mcp 0.11.0):** `capabilities()` now returns an owned snapshot, +not a borrowed reference. Call it again after authentication reconnect to obtain +the current peer capabilities; no capability lock guard needs to span an await. + ## Tool response timeouts `McpConnection::call_tool_with_timeout(name, arguments, Duration)` is an opt-in, diff --git a/crates/agentkit-mcp/src/lib.rs b/crates/agentkit-mcp/src/lib.rs index 2952f41..ff71949 100644 --- a/crates/agentkit-mcp/src/lib.rs +++ b/crates/agentkit-mcp/src/lib.rs @@ -21,6 +21,7 @@ use std::collections::{BTreeMap, BTreeSet, HashMap}; use std::fmt; +use std::sync::atomic::{AtomicBool, AtomicU64, Ordering}; use std::sync::{Arc, RwLock}; use std::time::Duration; @@ -42,12 +43,12 @@ use async_trait::async_trait; use futures_util::future::{join_all, try_join_all}; use futures_util::stream::BoxStream; use http::{HeaderName, HeaderValue}; -use rmcp::ServiceExt; use rmcp::handler::client::ClientHandler; use rmcp::model as rmcp_model; use rmcp::service::{ ClientInitializeError, Peer, PeerRequestOptions, RoleClient, RunningService, ServiceError, }; +use rmcp::service::{ClientLifecycleMode, ClientServiceExt}; use rmcp::transport::streamable_http_client::{ AuthRequiredError, InsufficientScopeError, StreamableHttpClient as RmcpStreamableHttpClient, StreamableHttpClientTransportConfig as RmcpStreamableHttpClientTransportConfig, @@ -60,7 +61,10 @@ use serde::{Deserialize, Serialize}; use serde_json::{Value, json}; use sse_stream::{Error as SseError, Sse}; use thiserror::Error; -use tokio::sync::{Mutex, broadcast, mpsc}; +use tokio::sync::{Mutex, broadcast, mpsc, watch}; + +mod subscriptions; +pub use subscriptions::McpSubscriptionStatus; /// Re-exports of the rmcp wire-protocol types this crate now surfaces directly /// instead of wrapping. Pull these in to pattern-match on tool annotations, @@ -78,13 +82,13 @@ pub use rmcp::model::{ FormElicitationCapability as McpFormElicitationCapability, GetPromptResult, ImageContent, ImageContent as RawImageContent, Implementation as McpImplementation, ProgressNotificationParam as McpProgressNotificationParam, Prompt as McpPrompt, PromptArgument, - PromptMessage, ReadResourceResult, Resource as McpRawResource, Resource as McpResource, - ResourceContents as McpResourceContents, + PromptMessage, ProtocolVersion as McpProtocolVersion, ReadResourceResult, + Resource as McpRawResource, Resource as McpResource, ResourceContents as McpResourceContents, ResourceUpdatedNotificationParam as McpResourceUpdatedNotificationParam, Role, Role as PromptMessageRole, RootsCapabilities as McpRootsCapabilities, - SamplingCapability as McpSamplingCapability, TextContent as RawTextContent, TextContent, - Tool as McpTool, ToolAnnotations as McpToolAnnotations, - UrlElicitationCapability as McpUrlElicitationCapability, + SamplingCapability as McpSamplingCapability, SubscriptionFilter as McpSubscriptionFilter, + TextContent as RawTextContent, TextContent, Tool as McpTool, + ToolAnnotations as McpToolAnnotations, UrlElicitationCapability as McpUrlElicitationCapability, }; // RMCP 3 still supports these legacy MCP surfaces. Keep exposing them for API @@ -946,8 +950,8 @@ pub struct LoggingCapability {} /// Server-originated catalog notifications observed by [`McpClientHandler`]. /// -/// Drained by [`McpConnection`] inside -/// [`McpServerManager::refresh_changed_catalogs`] to trigger re-discovery of +/// Consumed by [`McpConnection`] to advance its coalesced catalog revision. +/// Hosts use [`McpServerManager::refresh_changed_catalogs`] to reconcile /// the affected capability lists. For richer push-style consumption (progress, /// logging, resource updates, cancellation), subscribe via /// [`McpConnection::subscribe_events`] and pattern-match on @@ -1038,9 +1042,9 @@ const DEFAULT_EVENTS_CAPACITY: usize = 128; /// Channels paired with an [`McpClientHandler`] returned by /// [`McpHandlerConfig::build`]. /// -/// `notifications` is the legacy mpsc receiver consumed by the catalog refresh -/// path inside [`McpServerManager::refresh_changed_catalogs`]. `events` is the -/// broadcast sender that surfaces every [`McpServerEvent`] — clone it once and +/// `notifications` preserves the compatibility mpsc stream for standalone +/// handler consumers. Connections track coalesced revisions directly and release +/// this receiver on adoption. `events` is the broadcast sender that surfaces every [`McpServerEvent`] — clone it once and /// pass it into [`McpConnection::from_running_service_with_events`] when /// adopting an externally constructed [`rmcp::service::RunningService`]. If the /// adopted connection also needs adapter-time hooks from the same @@ -1064,7 +1068,9 @@ pub struct McpClientChannels { /// the connection must preserve adapter-time hooks from the config. #[derive(Clone)] pub struct McpClientHandler { - info: rmcp_model::ClientInfo, + info: rmcp_model::ClientConfig, + revision: watch::Sender, + active: Arc, notifications: mpsc::UnboundedSender, events: broadcast::Sender, sampling: Option>, @@ -1138,6 +1144,9 @@ impl ClientHandler for McpClientHandler { params: rmcp_model::ProgressNotificationParam, _context: rmcp::service::NotificationContext, ) -> impl Future + rmcp::service::MaybeSendFuture + '_ { + if !self.active.load(Ordering::Acquire) { + return std::future::ready(()); + } let _ = self.events.send(McpServerEvent::Progress(params)); std::future::ready(()) } @@ -1148,6 +1157,9 @@ impl ClientHandler for McpClientHandler { params: rmcp_model::LoggingMessageNotificationParam, _context: rmcp::service::NotificationContext, ) -> impl Future + rmcp::service::MaybeSendFuture + '_ { + if !self.active.load(Ordering::Acquire) { + return std::future::ready(()); + } let _ = self.events.send(McpServerEvent::Logging(params)); std::future::ready(()) } @@ -1157,6 +1169,9 @@ impl ClientHandler for McpClientHandler { params: rmcp_model::ResourceUpdatedNotificationParam, _context: rmcp::service::NotificationContext, ) -> impl Future + rmcp::service::MaybeSendFuture + '_ { + if !self.active.load(Ordering::Acquire) { + return std::future::ready(()); + } let _ = self.events.send(McpServerEvent::ResourceUpdated(params)); std::future::ready(()) } @@ -1166,6 +1181,9 @@ impl ClientHandler for McpClientHandler { params: rmcp_model::CancelledNotificationParam, _context: rmcp::service::NotificationContext, ) -> impl Future + rmcp::service::MaybeSendFuture + '_ { + if !self.active.load(Ordering::Acquire) { + return std::future::ready(()); + } let _ = self.events.send(McpServerEvent::Cancelled(params)); std::future::ready(()) } @@ -1174,6 +1192,11 @@ impl ClientHandler for McpClientHandler { &self, _context: rmcp::service::NotificationContext, ) -> impl Future + rmcp::service::MaybeSendFuture + '_ { + if !self.active.load(Ordering::Acquire) { + return std::future::ready(()); + } + self.revision + .send_modify(|revision| *revision = revision.saturating_add(1)); let _ = self.notifications.send(McpServerNotification::ToolsChanged); let _ = self.events.send(McpServerEvent::ToolListChanged); std::future::ready(()) @@ -1183,6 +1206,11 @@ impl ClientHandler for McpClientHandler { &self, _context: rmcp::service::NotificationContext, ) -> impl Future + rmcp::service::MaybeSendFuture + '_ { + if !self.active.load(Ordering::Acquire) { + return std::future::ready(()); + } + self.revision + .send_modify(|revision| *revision = revision.saturating_add(1)); let _ = self .notifications .send(McpServerNotification::ResourcesChanged); @@ -1194,6 +1222,11 @@ impl ClientHandler for McpClientHandler { &self, _context: rmcp::service::NotificationContext, ) -> impl Future + rmcp::service::MaybeSendFuture + '_ { + if !self.active.load(Ordering::Acquire) { + return std::future::ready(()); + } + self.revision + .send_modify(|revision| *revision = revision.saturating_add(1)); let _ = self .notifications .send(McpServerNotification::PromptsChanged); @@ -1201,7 +1234,7 @@ impl ClientHandler for McpClientHandler { std::future::ready(()) } - fn get_info(&self) -> rmcp_model::ClientInfo { + fn get_info(&self) -> rmcp_model::ClientConfig { self.info.clone() } } @@ -1335,7 +1368,9 @@ impl McpHandlerConfig { } let handler = McpClientHandler { - info: rmcp_model::ClientInfo::new( + revision: watch::channel(1).0, + active: Arc::new(AtomicBool::new(true)), + info: rmcp_model::ClientConfig::new( capabilities, rmcp_model::Implementation::new("agentkit-mcp", env!("CARGO_PKG_VERSION")) .with_title("agentkit MCP client"), @@ -1358,6 +1393,43 @@ impl McpHandlerConfig { } } +/// Notification state is scoped to a connection. Each invalidation advances a +/// watch revision; catalog publishers acknowledge only the revision they fetched. +#[derive(Clone)] +struct McpNotificationSink { + revision: watch::Sender, + active: Arc, + events: broadcast::Sender, +} + +impl McpNotificationSink { + fn mark_catalog_dirty(&self) { + self.revision + .send_modify(|revision| *revision = revision.saturating_add(1)); + } + + fn dispatch(&self, notification: rmcp_model::ServerNotification) { + if !self.active.load(Ordering::Acquire) { + return; + } + use rmcp_model::ServerNotification as N; + let event = match notification { + N::ToolListChangedNotification(_) => McpServerEvent::ToolListChanged, + N::ResourceListChangedNotification(_) => McpServerEvent::ResourceListChanged, + N::PromptListChangedNotification(_) => McpServerEvent::PromptListChanged, + N::ResourceUpdatedNotification(update) => { + let _ = self + .events + .send(McpServerEvent::ResourceUpdated(update.params)); + return; + } + _ => return, + }; + self.mark_catalog_dirty(); + let _ = self.events.send(event); + } +} + /// A live connection to a single MCP server, wrapping an /// [`rmcp::service::RunningService`]. pub struct McpConnection { @@ -1366,10 +1438,15 @@ pub struct McpConnection { inner: Mutex, peer: RwLock>, auth: Mutex>, - notifications: Mutex>, + sink: McpNotificationSink, + subscriptions: Mutex>, + lifecycle: Mutex<()>, + epoch: AtomicU64, + closed: AtomicBool, + desired_resources: Mutex>, events: broadcast::Sender, handler_config: McpHandlerConfig, - capabilities: McpServerCapabilities, + capabilities: RwLock, } /// The result of replaying an MCP operation after auth resolution. @@ -1386,7 +1463,7 @@ pub enum McpOperationResult { } impl McpConnection { - /// Connects to an MCP server, performs the rmcp `initialize` handshake, + /// Connects to an MCP server with modern discovery or legacy initialization, /// and returns a ready-to-use connection. No sampling / elicitation / /// roots responders are wired; use [`Self::connect_with_handler`] when /// the server may issue those requests. @@ -1422,17 +1499,31 @@ impl McpConnection { }; let peer = service.peer().clone(); - Ok(Self { + let revision = service.service().revision.clone(); + let sink = McpNotificationSink { + revision, + active: service.service().active.clone(), + events: events_tx.clone(), + }; + drop(notification_rx); + let connection = Self { server_id: config.id.clone(), config: Some(config.clone()), inner: Mutex::new(service), peer: RwLock::new(peer), auth: Mutex::new(auth.cloned()), - notifications: Mutex::new(notification_rx), + sink, + subscriptions: Mutex::new(None), + lifecycle: Mutex::new(()), + epoch: AtomicU64::new(1), + closed: AtomicBool::new(false), + desired_resources: Mutex::new(BTreeSet::new()), events: events_tx, handler_config, - capabilities, - }) + capabilities: RwLock::new(capabilities), + }; + connection.ensure_subscriptions().await?; + Ok(connection) } /// Adopts an externally constructed [`rmcp::service::RunningService`] as @@ -1440,9 +1531,9 @@ impl McpConnection { /// /// Use this when you need a transport rmcp supports but /// [`McpTransportBinding`] does not (in-memory pipes for tests, websockets, - /// custom IO). Pair the service with the notification receiver returned by - /// [`McpHandlerConfig::build`] so list-change notifications stay - /// observable. + /// custom IO). Pass the notification receiver returned by + /// [`McpHandlerConfig::build`]; the handler tracks catalog revisions directly + /// and the compatibility receiver is released on adoption. /// /// The connection has no [`McpServerConfig`] attached, so reconnect-on-auth /// is unavailable; [`resolve_auth`](Self::resolve_auth) only updates stored @@ -1498,29 +1589,45 @@ impl McpConnection { .map(|info| rmcp_server_capabilities_to_agentkit(&info.capabilities)) .unwrap_or_default(); let peer = service.peer().clone(); + let revision = service.service().revision.clone(); + let sink = McpNotificationSink { + revision, + active: service.service().active.clone(), + events: events.clone(), + }; + drop(notifications); Self { server_id: server_id.into(), config: None, inner: Mutex::new(service), peer: RwLock::new(peer), auth: Mutex::new(None), - notifications: Mutex::new(notifications), + sink, + subscriptions: Mutex::new(None), + lifecycle: Mutex::new(()), + epoch: AtomicU64::new(1), + closed: AtomicBool::new(false), + desired_resources: Mutex::new(BTreeSet::new()), events, handler_config, - capabilities, + capabilities: RwLock::new(capabilities), } } async fn reconnect_inner(&self, auth: Option<&MetadataMap>) -> Result<(), McpError> { + let _lifecycle = self.lifecycle.lock().await; + if self.closed.load(Ordering::Acquire) { + return Err(McpError::Transport("MCP connection is closed".into())); + } let Some(config) = self.config.clone() else { return Ok(()); }; - let (handler, channels) = self.handler_config.build_with(self.events.clone()); - let McpClientChannels { - notifications: notification_rx, - .. - } = channels; - let (service, _capabilities) = match &config.transport { + let (mut handler, channels) = self.handler_config.build_with(self.events.clone()); + handler.revision = self.sink.revision.clone(); + handler.active = Arc::new(AtomicBool::new(false)); + let mut candidate_sink = self.sink.clone(); + candidate_sink.active = handler.active.clone(); + let (service, capabilities) = match &config.transport { McpTransportBinding::Stdio(binding) => { connect_rmcp_stdio(&config, binding, handler).await? } @@ -1529,9 +1636,73 @@ impl McpConnection { } }; let new_peer = service.peer().clone(); - *self.notifications.lock().await = notification_rx; - *self.inner.lock().await = service; - *self.peer.write().expect("MCP peer lock poisoned") = new_peer; + let modern = new_peer + .peer_info() + .is_some_and(|info| !info.protocol_version.has_initialize()); + let resources = match self.subscriptions.lock().await.as_ref() { + Some(controller) => controller.settled_resources().await, + None => self.desired_resources.lock().await.clone(), + }; + let controller = if modern { + let mut filter = filter_for_capabilities(&capabilities); + if !resources.is_empty() { + filter.resource_subscriptions = Some(resources.iter().cloned().collect()); + } + Some( + subscriptions::SubscriptionController::start( + new_peer.clone(), + filter, + candidate_sink.clone(), + ) + .await?, + ) + } else { + // Legacy subscriptions are session-scoped and must be replayed on + // the candidate before retiring the authoritative old transport. + for uri in &resources { + #[allow(deprecated)] + new_peer + .subscribe(rmcp_model::SubscribeRequestParams::new(uri.clone())) + .await + .map_err(rmcp_service_error)?; + } + None + }; + drop(channels.notifications); + // Acquire every async publication guard before changing state. Cancellation + // during preparation leaves the old connection and controller untouched. + let mut inner = self.inner.lock().await; + let mut active = self.subscriptions.lock().await; + let mut desired_resources = self.desired_resources.lock().await; + let next_epoch = self + .epoch + .load(Ordering::Acquire) + .checked_add(1) + .filter(|epoch| *epoch < u64::MAX) + .ok_or_else(|| McpError::Transport("MCP connection generation exhausted".into()))?; + let (mut previous, old_controller) = { + let mut peer = self.peer.write().expect("MCP peer lock poisoned"); + let mut current_capabilities = self + .capabilities + .write() + .expect("MCP capabilities lock poisoned"); + self.epoch.store(next_epoch, Ordering::Release); + inner.service().active.store(false, Ordering::Release); + let previous = std::mem::replace(&mut *inner, service); + candidate_sink.active.store(true, Ordering::Release); + *peer = new_peer; + *current_capabilities = capabilities; + *desired_resources = resources; + (previous, std::mem::replace(&mut *active, controller)) + }; + drop(desired_resources); + drop(active); + drop(inner); + self.mark_catalog_dirty(); + if let Some(mut controller) = old_controller { + controller.shutdown().await; + } + let _ = previous.close().await; Ok(()) } @@ -1544,9 +1715,71 @@ impl McpConnection { &self.server_id } - /// Returns the capabilities advertised by the server during `initialize`. - pub fn capabilities(&self) -> &McpServerCapabilities { - &self.capabilities + /// Returns the capabilities from the current connection, including auth reconnects. + pub fn capabilities(&self) -> McpServerCapabilities { + self.capabilities + .read() + .expect("MCP capabilities lock poisoned") + .clone() + } + + /// Protocol selected by discovery or the legacy initialization handshake. + pub fn protocol_version(&self) -> Option { + self.peer() + .peer_info() + .map(|info| info.protocol_version.clone()) + } + + /// Receives coalesced catalog invalidations; revision one denotes initial discovery. + pub fn catalog_changes(&self) -> watch::Receiver { + self.sink.revision.subscribe() + } + + /// Changes whenever authentication replaces the transport or the connection closes. + pub fn connection_epoch(&self) -> u64 { + self.epoch.load(Ordering::Acquire) + } + + /// Current catalog invalidation revision. A successful fetch acknowledges only its captured revision. + pub fn catalog_revision(&self) -> u64 { + *self.sink.revision.borrow() + } + + fn mark_catalog_dirty(&self) { + self.sink.mark_catalog_dirty(); + } + + fn modern(&self) -> bool { + self.protocol_version() + .is_some_and(|version| !version.has_initialize()) + } + + fn subscription_filter(&self) -> McpSubscriptionFilter { + let capabilities = self.capabilities(); + filter_for_capabilities(&capabilities) + } + + async fn ensure_subscriptions(&self) -> Result<(), McpError> { + let _lifecycle = self.lifecycle.lock().await; + if self.closed.load(Ordering::Acquire) { + return Err(McpError::Transport("MCP connection is closed".into())); + } + let mut subscriptions = self.subscriptions.lock().await; + if self.modern() && subscriptions.is_none() { + let mut filter = self.subscription_filter(); + let resources = self.desired_resources.lock().await; + if !resources.is_empty() { + filter.resource_subscriptions = Some(resources.iter().cloned().collect()); + } + let controller = subscriptions::SubscriptionController::start( + self.peer(), + filter, + self.sink.clone(), + ) + .await?; + *subscriptions = Some(controller); + } + Ok(()) } /// Returns the [`McpHandlerConfig`] this connection was built with. @@ -1562,8 +1795,8 @@ impl McpConnection { /// [`McpHandlerConfig::with_events_capacity`], defaults to /// `DEFAULT_EVENTS_CAPACITY`) before slow consumers are signalled with /// [`broadcast::error::RecvError::Lagged`]. Catalog `*ListChanged` events - /// are also delivered through the legacy [`McpServerNotification`] - /// receiver consumed by [`McpServerManager::refresh_changed_catalogs`]. + /// also advance [`Self::catalog_revision`] independently of broadcast lag. + /// Hosts reconcile them with [`McpServerManager::refresh_changed_catalogs`]. pub fn subscribe_events(&self) -> broadcast::Receiver { self.events.subscribe() } @@ -1575,32 +1808,46 @@ impl McpConnection { #[allow(deprecated)] pub async fn subscribe_resource(&self, uri: impl Into) -> Result<(), McpError> { let uri = uri.into(); + self.ensure_subscriptions().await?; + let _lifecycle = self.lifecycle.lock().await; + if let Some(controller) = self.subscriptions.lock().await.as_ref() { + return controller.subscribe_resource(uri).await; + } self.peer() .subscribe(rmcp_model::SubscribeRequestParams::new(uri.clone())) .await .map_err(|error| { rmcp_operation_error( &self.server_id, - McpMethod::ResourcesSubscribe { uri }, + McpMethod::ResourcesSubscribe { uri: uri.clone() }, error, ) - }) + })?; + self.desired_resources.lock().await.insert(uri); + Ok(()) } /// Cancels a previous [`Self::subscribe_resource`] subscription. #[allow(deprecated)] pub async fn unsubscribe_resource(&self, uri: impl Into) -> Result<(), McpError> { let uri = uri.into(); + self.ensure_subscriptions().await?; + let _lifecycle = self.lifecycle.lock().await; + if let Some(controller) = self.subscriptions.lock().await.as_ref() { + return controller.unsubscribe_resource(uri).await; + } self.peer() .unsubscribe(rmcp_model::UnsubscribeRequestParams::new(uri.clone())) .await .map_err(|error| { rmcp_operation_error( &self.server_id, - McpMethod::ResourcesUnsubscribe { uri }, + McpMethod::ResourcesUnsubscribe { uri: uri.clone() }, error, ) - }) + })?; + self.desired_resources.lock().await.remove(&uri); + Ok(()) } /// Negotiates the minimum severity the server should emit through @@ -1642,12 +1889,43 @@ impl McpConnection { .map_err(rmcp_service_error) } - /// Gracefully closes the underlying rmcp service. - /// - /// For Streamable HTTP this drives the rmcp transport to issue a `DELETE` - /// against the negotiated session, releasing server-side state. + /// Current acknowledged modern subscription, or `None` for legacy peers. + pub async fn subscription_status(&self) -> Option { + self.subscriptions + .lock() + .await + .as_ref() + .map(|controller| controller.status()) + } + + /// Watches establishment, accepted filters, disconnects, and recovery. + pub async fn subscription_changes(&self) -> Option> { + self.subscriptions + .lock() + .await + .as_ref() + .map(|controller| controller.subscribe_status()) + } + + /// Closes owned subscriptions before shutting down the underlying transport. + /// Legacy HTTP sessions retain their protocol-specific DELETE cleanup. pub async fn close(&self) -> Result<(), McpError> { + let _lifecycle = self.lifecycle.lock().await; + // Acquire every guard before initiating shutdown. Once closed becomes + // visible, transport cancellation and worker stop need no more awaits. let mut inner = self.inner.lock().await; + let mut subscriptions = self.subscriptions.lock().await; + self.closed.store(true, Ordering::Release); + inner.service().active.store(false, Ordering::Release); + self.epoch + .store(self.connection_epoch().saturating_add(1), Ordering::Release); + let controller = subscriptions.take(); + inner.cancellation_token().cancel(); + drop(subscriptions); + self.mark_catalog_dirty(); + if let Some(mut controller) = controller { + controller.shutdown().await; + } inner .close() .await @@ -1668,32 +1946,36 @@ impl McpConnection { } } let snapshot = auth_slot.clone(); - drop(auth_slot); + // Keep auth ordering through reconnect: later resolutions cannot publish + // an older transport after a newer credential selection. // Only reconnect if we have a config to reconnect with. Without one // (e.g. constructed via [`from_running_service`]) the auth is stored // but not pushed to the live transport. if self.config.is_some() { self.reconnect_inner(snapshot.as_ref()).await?; } + drop(auth_slot); Ok(()) } /// Discovers tools, resources, and prompts that the server advertised. pub async fn discover(&self) -> Result { + self.ensure_subscriptions().await?; + let capabilities = self.capabilities(); let tools = async { - match self.capabilities.tools { + match capabilities.tools { Some(_) => self.list_tools().await, None => Ok(Vec::new()), } }; let resources = async { - match self.capabilities.resources { + match capabilities.resources { Some(_) => self.list_resources().await, None => Ok(Vec::new()), } }; let prompts = async { - match self.capabilities.prompts { + match capabilities.prompts { Some(_) => self.list_prompts().await, None => Ok(Vec::new()), } @@ -1708,15 +1990,6 @@ impl McpConnection { }) } - async fn drain_notifications(&self) -> Vec { - let mut notifications = self.notifications.lock().await; - let mut drained = Vec::new(); - while let Ok(notification) = notifications.try_recv() { - drained.push(notification); - } - drained - } - /// Lists all tools advertised by the connected MCP server. pub async fn list_tools(&self) -> Result, McpError> { self.peer() @@ -1910,6 +2183,21 @@ impl McpConnection { } } +fn filter_for_capabilities(capabilities: &McpServerCapabilities) -> McpSubscriptionFilter { + let mut filter = McpSubscriptionFilter::default(); + filter.tools_list_changed = capabilities.tools.as_ref().map(|_| true); + filter.resources_list_changed = capabilities.resources.as_ref().map(|_| true); + filter.prompts_list_changed = capabilities.prompts.as_ref().map(|_| true); + filter +} + +fn client_lifecycle() -> ClientLifecycleMode { + ClientLifecycleMode::Auto { + preferred_versions: vec![rmcp_model::ProtocolVersion::V_2026_07_28], + legacy_version: Some(rmcp_model::ProtocolVersion::V_2025_11_25), + } +} + async fn connect_rmcp_stdio( config: &McpServerConfig, binding: &StdioTransportConfig, @@ -1929,7 +2217,7 @@ async fn connect_rmcp_stdio( .map_err(McpError::Io)?; let service = handler - .serve(transport) + .serve_with_lifecycle(transport, client_lifecycle()) .await .map_err(|error| rmcp_initialize_error(config, error))?; let capabilities = service @@ -1961,11 +2249,15 @@ async fn connect_rmcp_streamable_http( DynHttpClient(client.clone()), rmcp_config, ); - handler.serve(transport).await + handler + .serve_with_lifecycle(transport, client_lifecycle()) + .await } None => { let transport = StreamableHttpClientTransport::from_config(rmcp_config); - handler.serve(transport).await + handler + .serve_with_lifecycle(transport, client_lifecycle()) + .await } }; let service = result.map_err(|error| rmcp_initialize_error(config, error))?; @@ -2325,6 +2617,64 @@ impl McpServerOptions { } } +/// Prepared catalog discovery. Fetch outside the manager lock, then publish +/// under the lock. Dropping this value never acknowledges pending changes. +pub struct McpCatalogRefresh { + server_id: McpServerId, + connection: Arc, + revision: u64, + epoch: u64, + publication: u64, + options: McpServerOptions, +} + +impl McpCatalogRefresh { + /// Returns the target server. + pub fn server_id(&self) -> &McpServerId { + &self.server_id + } + + /// Performs discovery without borrowing the manager. + pub async fn fetch(self) -> McpCatalogRefreshResult { + let snapshot = + McpServerManager::discover_with_options(&self.connection, &self.options).await; + McpCatalogRefreshResult { + request: self, + snapshot, + } + } +} + +/// Completed discovery with its identity and generation publication fences. +pub struct McpCatalogRefreshResult { + request: McpCatalogRefresh, + snapshot: Result, +} + +/// Wake sources captured under the manager lock. Rebuild after each wake to +/// include newly connected servers. Waiting does not borrow the manager. +pub struct McpCatalogRefreshWake { + receivers: Vec>, +} + +impl McpCatalogRefreshWake { + /// Waits for a pending change or a connection lifecycle change. + pub async fn changed(&mut self) { + let waits = self.receivers.iter_mut().map(|receiver| { + Box::pin(async move { + let _ = receiver.changed().await; + }) + }); + let _ = futures_util::future::select_all(waits).await; + } +} + +#[derive(Default)] +struct McpCatalogPublication { + revision: u64, + publication: u64, +} + /// Manages the lifecycle of one or more MCP servers. pub struct McpServerManager { configs: BTreeMap, @@ -2340,6 +2690,8 @@ pub struct McpServerManager { /// [`CatalogWriter`] on connect/disconnect/refresh without rebuilding /// the whole catalog. server_tools: BTreeMap>, + catalog_publications: BTreeMap, + catalog_wake_tx: tokio::sync::watch::Sender, } impl Default for McpServerManager { @@ -2356,6 +2708,8 @@ impl Default for McpServerManager { handler_config: McpHandlerConfig::default(), catalog_writer, server_tools: BTreeMap::new(), + catalog_publications: BTreeMap::new(), + catalog_wake_tx: tokio::sync::watch::channel(0).0, } } } @@ -2522,23 +2876,28 @@ impl McpServerManager { snapshot, namespace: self.namespace.clone(), }; - self.install_connection(server_id.clone(), handle.clone()) - .await; - self.register_server_tools(server_id, &handle.snapshot); - self.emit_catalog_event(McpCatalogEvent::ServerConnected { - server_id: server_id.clone(), - }); + let previous = self.install_connection(server_id.clone(), handle.clone()); + if let Some(previous) = previous { + let _ = previous.connection.close().await; + } Ok(handle) } - /// Installs a freshly connected handle, gracefully closing any previous - /// connection it replaces so transports (and, for Streamable HTTP, - /// server-side sessions) are not leaked. Close failures on the replaced - /// connection are ignored — it is being discarded either way. - async fn install_connection(&mut self, server_id: McpServerId, handle: McpServerHandle) { - if let Some(previous) = self.connections.insert(server_id, handle) { - let _ = previous.connection.close().await; - } + /// Commits a connection and its complete catalog without suspension. The + /// caller closes the returned old handle only after this state is consistent. + fn install_connection( + &mut self, + server_id: McpServerId, + handle: McpServerHandle, + ) -> Option { + self.catalog_publications.remove(&server_id); + let snapshot = handle.snapshot.clone(); + let previous = self.connections.insert(server_id.clone(), handle); + self.register_server_tools(&server_id, &snapshot); + self.catalog_wake_tx + .send_modify(|revision| *revision = revision.saturating_add(1)); + self.emit_catalog_event(McpCatalogEvent::ServerConnected { server_id }); + previous } /// Connects all registered servers concurrently. @@ -2584,18 +2943,16 @@ impl McpServerManager { let results = try_join_all(futures).await?; let mut handles = Vec::with_capacity(results.len()); - let mut connected: Vec<(McpServerId, McpDiscoverySnapshot)> = - Vec::with_capacity(results.len()); + let mut replaced = Vec::new(); for (server_id, handle) in results { - connected.push((server_id.clone(), handle.snapshot.clone())); - self.install_connection(server_id, handle.clone()).await; + if let Some(previous) = self.install_connection(server_id, handle.clone()) { + replaced.push(previous); + } handles.push(handle); } - for (server_id, snapshot) in &connected { - self.register_server_tools(server_id, snapshot); - } - for (server_id, _) in connected { - self.emit_catalog_event(McpCatalogEvent::ServerConnected { server_id }); + // Every successful installation is consistent before any cleanup await. + for previous in replaced { + let _ = previous.connection.close().await; } Ok(handles) } @@ -2697,13 +3054,14 @@ impl McpServerManager { }); let results = join_all(futures).await; - let mut connected_snapshots = Vec::new(); + let mut replaced = Vec::new(); for (server_id, result) in results { match result { Ok(handle) => { - connected_snapshots.push((server_id.clone(), handle.snapshot.clone())); - self.install_connection(server_id, handle.clone()).await; + if let Some(previous) = self.install_connection(server_id, handle.clone()) { + replaced.push(previous); + } connected.push(handle); } Err(error) => { @@ -2712,11 +3070,8 @@ impl McpServerManager { } } - for (server_id, snapshot) in &connected_snapshots { - self.register_server_tools(server_id, snapshot); - } - for (server_id, _) in connected_snapshots { - self.emit_catalog_event(McpCatalogEvent::ServerConnected { server_id }); + for previous in replaced { + let _ = previous.connection.close().await; } McpConnectAllSettled { @@ -2725,86 +3080,165 @@ impl McpServerManager { } } - /// Re-discovers capabilities for a connected server. - pub async fn refresh_server( - &mut self, + /// Prepares explicit refresh without network I/O. Marks the server dirty + /// so failed or cancelled work can be retried by the changed-catalog path. + pub fn prepare_catalog_refresh( + &self, server_id: &McpServerId, - ) -> Result { + ) -> Result { let handle = self .connections - .get_mut(server_id) + .get(server_id) .ok_or_else(|| McpError::UnknownServer(server_id.to_string()))?; - let options = self.options.get(server_id).cloned().unwrap_or_default(); - let previous = handle.snapshot.clone(); - let snapshot = match Self::discover_with_options(&handle.connection, &options).await { + handle.connection.mark_catalog_dirty(); + Ok(self.catalog_refresh_request(server_id, handle)) + } + + fn catalog_refresh_request( + &self, + server_id: &McpServerId, + handle: &McpServerHandle, + ) -> McpCatalogRefresh { + McpCatalogRefresh { + server_id: server_id.clone(), + connection: handle.connection.clone(), + revision: handle.connection.catalog_revision(), + epoch: handle.connection.connection_epoch(), + publication: self + .catalog_publications + .get(server_id) + .map_or(0, |state| state.publication), + options: self.options.get(server_id).cloned().unwrap_or_default(), + } + } + + /// Coalesces notifications into one request per dirty server. Preparation + /// never acknowledges changes, so dropping a request leaves it retryable. + pub fn prepare_changed_catalogs(&self) -> Vec { + self.connections + .iter() + .filter_map(|(server_id, handle)| { + let published = self + .catalog_publications + .get(server_id) + .map_or(0, |state| state.revision); + let revision = handle.connection.catalog_revision(); + (revision == u64::MAX || revision != published) + .then(|| self.catalog_refresh_request(server_id, handle)) + }) + .collect() + } + + /// Captures wake sources while locked. Await outside the lock and rebuild + /// after each wake. Already-dirty servers wake immediately. + pub fn catalog_refresh_wake(&self) -> McpCatalogRefreshWake { + let mut receivers = vec![self.catalog_wake_tx.subscribe()]; + for (server_id, handle) in &self.connections { + let mut receiver = handle.connection.catalog_changes(); + let published = self + .catalog_publications + .get(server_id) + .map_or(0, |state| state.revision); + let revision = *receiver.borrow(); + if revision == u64::MAX || revision != published { + receiver.mark_changed(); + } + receivers.push(receiver); + } + McpCatalogRefreshWake { receivers } + } + + /// Publishes only while connection identity, connection epoch, notification + /// revision and prior publication still match. Superseded results are ignored. + /// Failure retains the previous catalog and dirty revision for retry. + pub fn publish_catalog_refresh( + &mut self, + result: McpCatalogRefreshResult, + ) -> Result, McpError> { + let McpCatalogRefreshResult { request, snapshot } = result; + let Some(handle) = self.connections.get(&request.server_id) else { + return Ok(Vec::new()); + }; + let publication = self + .catalog_publications + .get(&request.server_id) + .map_or(0, |state| state.publication); + if !Arc::ptr_eq(&handle.connection, &request.connection) + || handle.connection.connection_epoch() != request.epoch + || handle.connection.catalog_revision() != request.revision + || publication != request.publication + { + return Ok(Vec::new()); + } + let snapshot = match snapshot { Ok(snapshot) => snapshot, Err(error) => { self.emit_catalog_event(McpCatalogEvent::RefreshFailed { - server_id: server_id.clone(), + server_id: request.server_id, message: error.to_string(), }); return Err(error); } }; - handle.snapshot = snapshot.clone(); - let events = diff_discovery_snapshots(server_id, &previous, &snapshot); - if !events.is_empty() { - self.apply_catalog_events(server_id, &snapshot, &events); - for event in events { - self.emit_catalog_event(event); - } + let events = diff_discovery_snapshots(&request.server_id, &handle.snapshot, &snapshot); + self.apply_catalog_events(&request.server_id, &snapshot, &events); + self.connections + .get_mut(&request.server_id) + .expect("identity checked above") + .snapshot = snapshot; + self.catalog_publications.insert( + request.server_id, + McpCatalogPublication { + revision: request.revision, + publication: publication.saturating_add(1), + }, + ); + for event in &events { + self.emit_catalog_event(event.clone()); } - Ok(snapshot) + Ok(events) } - /// Processes pending server list-change notifications. + /// Re-discovers capabilities. Prefer prepare/fetch/publish for shared managers. + pub async fn refresh_server( + &mut self, + server_id: &McpServerId, + ) -> Result { + let request = self.prepare_catalog_refresh(server_id)?; + self.publish_catalog_refresh(request.fetch().await)?; + Ok(self + .connections + .get(server_id) + .expect("server remains connected") + .snapshot + .clone()) + } + + /// Refreshes dirty servers concurrently, isolating failures. All completed + /// results are processed before the first error is returned. pub async fn refresh_changed_catalogs(&mut self) -> Result, McpError> { - let server_ids = self.connections.keys().cloned().collect::>(); + let results = join_all( + self.prepare_changed_catalogs() + .into_iter() + .map(McpCatalogRefresh::fetch), + ) + .await; let mut emitted = Vec::new(); - - for server_id in server_ids { - let Some(connection) = self - .connections - .get(&server_id) - .map(McpServerHandle::connection) - else { - continue; - }; - let notifications = connection.drain_notifications().await; - if notifications.is_empty() { - continue; - } - - let handle = self - .connections - .get_mut(&server_id) - .ok_or_else(|| McpError::UnknownServer(server_id.to_string()))?; - let options = self.options.get(&server_id).cloned().unwrap_or_default(); - let previous = handle.snapshot.clone(); - let snapshot = match Self::discover_with_options(&handle.connection, &options).await { - Ok(snapshot) => snapshot, + let mut first_error = None; + for result in results { + match self.publish_catalog_refresh(result) { + Ok(events) => emitted.extend(events), Err(error) => { - let event = McpCatalogEvent::RefreshFailed { - server_id: server_id.clone(), - message: error.to_string(), - }; - self.emit_catalog_event(event.clone()); - emitted.push(event); - return Err(error); - } - }; - handle.snapshot = snapshot.clone(); - let events = diff_discovery_snapshots(&server_id, &previous, &snapshot); - if !events.is_empty() { - self.apply_catalog_events(&server_id, &snapshot, &events); - for event in events { - self.emit_catalog_event(event.clone()); - emitted.push(event); + if first_error.is_none() { + first_error = Some(error); + } } } } - - Ok(emitted) + match first_error { + Some(error) => Err(error), + None => Ok(emitted), + } } /// Disconnects a server and removes it from active connections. @@ -2819,12 +3253,14 @@ impl McpServerManager { let Some(handle) = self.connections.remove(server_id) else { return Err(McpError::UnknownServer(server_id.to_string())); }; - let close_result = handle.connection.close().await; + self.catalog_publications.remove(server_id); + self.catalog_wake_tx + .send_modify(|revision| *revision = revision.saturating_add(1)); self.unregister_server_tools(server_id); self.emit_catalog_event(McpCatalogEvent::ServerDisconnected { server_id: server_id.clone(), }); - close_result + handle.connection.close().await } /// Removes a server's registration entirely: its config, options, and diff --git a/crates/agentkit-mcp/src/subscriptions.rs b/crates/agentkit-mcp/src/subscriptions.rs new file mode 100644 index 0000000..e150aac --- /dev/null +++ b/crates/agentkit-mcp/src/subscriptions.rs @@ -0,0 +1,493 @@ +//! Connection-owned subscriptions. Only the worker commits desired state. +use std::{collections::BTreeSet, time::Duration}; + +use rmcp::{ + model::SubscriptionFilter, + service::{Peer, RoleClient, Subscription, SubscriptionEnd}, +}; +use tokio::{ + sync::{mpsc, oneshot, watch}, + task::JoinHandle, +}; + +use crate::{McpError, McpNotificationSink}; + +const ESTABLISHMENT_TIMEOUT: Duration = Duration::from_secs(10); +const MAX_RETRIES: u32 = 5; + +/// Current state of the connection's modern MCP notification subscription. +#[derive(Clone, Debug, PartialEq)] +pub enum McpSubscriptionStatus { + /// Establishing a fresh stream; catalog contents may need reconciliation. + Pending, + /// The server acknowledged this stream. Unaccepted catalog flags remain visible. + Active { + /// Complete filter sent to the server. + requested: SubscriptionFilter, + /// Filter accepted by the server. + acknowledged: SubscriptionFilter, + }, + /// The worker has stopped. A graceful server completion is not retried. + Closed { + /// Whether the server returned the final listen result. + graceful: bool, + /// Failure diagnostic, if this was not an intentional shutdown. + error: Option, + }, +} + +enum ResourceChange { + Barrier, + Add(String), + Remove(String), +} + +struct Update { + change: ResourceChange, + reply: oneshot::Sender>, +} + +/// Owns a stream and the only task allowed to replace its committed filter. +pub(super) struct SubscriptionController { + commands: mpsc::Sender, + status: watch::Receiver, + requested: watch::Receiver>, + stop: watch::Sender, + worker: Option>, +} + +impl SubscriptionController { + /// Return only after the initial stream has been acknowledged. + pub(super) async fn start( + peer: Peer, + base_filter: SubscriptionFilter, + sink: McpNotificationSink, + ) -> Result { + let subscription = establish(&peer, base_filter.clone()).await?; + let initial = active_status(&base_filter, &subscription); + let (status_tx, status) = watch::channel(initial); + let (requested_tx, requested) = watch::channel(resource_set(&base_filter)); + let (commands, updates) = mpsc::channel(32); + let (stop, mut stopped) = watch::channel(false); + let worker = tokio::spawn(async move { + // Dropping run drops every in-flight RMCP request/stream, which + // unregisters it and sends best-effort cancellation in RMCP. + tokio::select! { + result = run(peer, base_filter, subscription, sink, updates, &status_tx, &requested_tx) => { + status_tx.send_replace(result); + } + _ = stopped.changed() => { + status_tx.send_replace(McpSubscriptionStatus::Closed { + graceful: false, error: None, + }); + } + } + }); + Ok(Self { + commands, + status, + requested, + stop, + worker: Some(worker), + }) + } + + pub(super) fn status(&self) -> McpSubscriptionStatus { + self.status.borrow().clone() + } + + pub(super) fn subscribe_status(&self) -> watch::Receiver { + self.status.clone() + } + + // Lifecycle replacement waits behind updates already owned by this worker. + // A canceled API caller may have released its connection guard while its + // command still awaits acknowledgment. + pub(super) async fn settled_resources(&self) -> BTreeSet { + let _ = self.change_resources(ResourceChange::Barrier).await; + self.requested_resources() + } + + pub(super) fn requested_resources(&self) -> BTreeSet { + self.requested.borrow().clone() + } + + pub(super) async fn subscribe_resource(&self, uri: String) -> Result<(), McpError> { + self.change_resources(ResourceChange::Add(uri)).await + } + + pub(super) async fn unsubscribe_resource(&self, uri: String) -> Result<(), McpError> { + self.change_resources(ResourceChange::Remove(uri)).await + } + + /// Once enqueued, the worker owns the change despite caller cancellation. + async fn change_resources(&self, change: ResourceChange) -> Result<(), McpError> { + let (reply, response) = oneshot::channel(); + self.commands + .send(Update { change, reply }) + .await + .map_err(|_| closed_error())?; + response.await.map_err(|_| closed_error())? + } + + /// Cancel all outstanding establishment/receive work and join the task. + pub(super) async fn shutdown(&mut self) { + self.stop.send_replace(true); + // Keep ownership until await completes: cancellation of shutdown still + // leaves Drop able to abort the worker. + if let Some(worker) = self.worker.as_mut() { + let _ = worker.await; + } + self.worker = None; + } +} + +impl Drop for SubscriptionController { + fn drop(&mut self) { + self.stop.send_replace(true); + if let Some(worker) = self.worker.take() { + worker.abort(); + } + } +} + +// Non-security jitter avoids synchronized reconnect bursts without adding an RNG +// dependency. The bounded retry count keeps the shift and total delay finite. +fn retry_delay(attempt: u32) -> Duration { + let jitter = std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap_or_default() + .subsec_nanos() as u64 + % 100; + Duration::from_millis(100 * (1 << attempt) + jitter) +} + +fn resource_set(filter: &SubscriptionFilter) -> BTreeSet { + filter + .resource_subscriptions + .iter() + .flatten() + .cloned() + .collect() +} + +fn closed_error() -> McpError { + McpError::Transport("subscription controller closed".into()) +} + +async fn establish( + peer: &Peer, + filter: SubscriptionFilter, +) -> Result { + tokio::time::timeout(ESTABLISHMENT_TIMEOUT, peer.listen(filter)) + .await + .map_err(|_| McpError::Timeout { + operation: "subscriptions/listen acknowledgment", + duration: ESTABLISHMENT_TIMEOUT, + })? + .map_err(|error| McpError::Transport(error.to_string())) +} + +fn verify_resources( + requested: &SubscriptionFilter, + accepted: &SubscriptionFilter, +) -> Result<(), McpError> { + let accepted = accepted + .resource_subscriptions + .as_deref() + .unwrap_or_default(); + if requested + .resource_subscriptions + .as_deref() + .unwrap_or_default() + .iter() + .all(|uri| accepted.contains(uri)) + { + Ok(()) + } else { + Err(McpError::Protocol( + "server did not acknowledge every requested resource URI".into(), + )) + } +} + +fn active_status( + filter: &SubscriptionFilter, + subscription: &Subscription, +) -> McpSubscriptionStatus { + McpSubscriptionStatus::Active { + requested: filter.clone(), + acknowledged: subscription.acknowledged().clone(), + } +} + +fn terminal(subscription: &Subscription) -> Option { + match subscription.end() { + Some(SubscriptionEnd::Graceful(_)) => Some(McpSubscriptionStatus::Closed { + graceful: true, + error: None, + }), + Some(SubscriptionEnd::Cancelled) => Some(McpSubscriptionStatus::Closed { + graceful: false, + error: None, + }), + _ => None, + } +} + +async fn run( + peer: Peer, + mut filter: SubscriptionFilter, + mut active: Subscription, + sink: McpNotificationSink, + mut commands: mpsc::Receiver, + status: &watch::Sender, + requested: &watch::Sender>, +) -> McpSubscriptionStatus { + loop { + tokio::select! { + command = commands.recv() => { + let Some(command) = command else { + return McpSubscriptionStatus::Closed { graceful: false, error: None }; + }; + let mut resources = resource_set(&filter); + match command.change { + ResourceChange::Barrier => { + let _ = command.reply.send(Ok(())); + continue; + } + ResourceChange::Add(uri) => { resources.insert(uri); } + ResourceChange::Remove(uri) => { resources.remove(&uri); } + } + let mut proposed = filter.clone(); + proposed.resource_subscriptions = if resources.is_empty() { + None + } else { + Some(resources.iter().cloned().collect()) + }; + // Drain the original stream during establishment: it remains + // authoritative until the replacement ACK has been validated. + let replacement = establish(&peer, proposed.clone()); + tokio::pin!(replacement); + let mut original_ended = false; + let result = loop { + tokio::select! { + result = &mut replacement => break result, + notification = active.next(), if !original_ended => { + match notification { + Ok(Some(notification)) => sink.dispatch(notification), + _ => { + original_ended = true; + sink.mark_catalog_dirty(); + } + } + } + } + }; + let result = result.and_then(|replacement| { + verify_resources(&proposed, replacement.acknowledged())?; + // No awaits or callbacks in this commit. Dropping the old + // RMCP stream cancels it without touching the new stream. + active = replacement; + filter = proposed; + requested.send_replace(resources); + status.send_replace(active_status(&filter, &active)); + Ok(()) + }); + let _ = command.reply.send(result); + // If the original ended during a rejected update, process its + // terminal state/recovery on the next next() call. + } + notification = active.next() => { + match notification { + Ok(Some(notification)) => sink.dispatch(notification), + outcome => { + if let Some(terminal) = terminal(&active) { + return terminal; + } + sink.mark_catalog_dirty(); + status.send_replace(McpSubscriptionStatus::Pending); + let mut last_error = match outcome { + Err(error) => error.to_string(), + _ => "subscription ended abruptly".into(), + }; + let mut recovered = None; + for attempt in 0..MAX_RETRIES { + tokio::time::sleep(retry_delay(attempt)).await; + match establish(&peer, filter.clone()).await.and_then(|stream| { + verify_resources(&filter, stream.acknowledged())?; + Ok(stream) + }) { + Ok(stream) => { recovered = Some(stream); break; } + Err(error) => last_error = error.to_string(), + } + } + let Some(stream) = recovered else { + return McpSubscriptionStatus::Closed { graceful: false, error: Some(last_error) }; + }; + active = stream; + status.send_replace(active_status(&filter, &active)); + } + } + } + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn rejected_ack_does_not_mutate_requested_filter() { + let mut requested = SubscriptionFilter::default(); + requested.resource_subscriptions = Some(vec!["resource:a".into(), "resource:b".into()]); + let mut accepted = SubscriptionFilter::default(); + accepted.resource_subscriptions = Some(vec!["resource:a".into()]); + assert!(verify_resources(&requested, &accepted).is_err()); + assert_eq!(requested.resource_subscriptions.as_ref().unwrap().len(), 2); + assert!(verify_resources(&requested, &requested).is_ok()); + } + + #[test] + fn empty_resource_set_accepts_absent_resource_filter() { + assert!( + verify_resources( + &SubscriptionFilter::default(), + &SubscriptionFilter::default() + ) + .is_ok() + ); + } + + struct ResourceServer; + + impl rmcp::ServerHandler for ResourceServer { + fn get_info(&self) -> rmcp::model::ServerConfig { + rmcp::model::ServerConfig::new( + rmcp::model::ServerCapabilities::builder() + .enable_resources() + .enable_resources_subscribe() + .build(), + ) + } + + fn accepted_subscription_filter( + &self, + requested: &SubscriptionFilter, + ) -> Option { + let mut accepted = requested.clone(); + if let Some(uris) = accepted.resource_subscriptions.as_mut() { + uris.retain(|uri| uri != "rejected"); + } + Some(accepted) + } + + async fn listen( + &self, + _context: rmcp::service::SubscriptionContext, + ) -> Result<(), rmcp::ErrorData> { + std::future::pending().await + } + } + + async fn controller() -> ( + SubscriptionController, + rmcp::service::RunningService, + ) { + use rmcp::{ClientServiceExt, ServiceExt}; + let (server_transport, client_transport) = tokio::io::duplex(16 * 1024); + tokio::spawn(async move { + let server = ResourceServer.serve(server_transport).await.unwrap(); + let _ = server.waiting().await; + }); + let client = () + .serve_with_lifecycle( + client_transport, + rmcp::ClientLifecycleMode::Discover { + preferred_versions: vec![rmcp::model::ProtocolVersion::V_2026_07_28], + }, + ) + .await + .unwrap(); + let (revision, _) = watch::channel(0); + let (events, _) = tokio::sync::broadcast::channel(16); + let sink = McpNotificationSink { + active: std::sync::Arc::new(std::sync::atomic::AtomicBool::new(true)), + revision, + events, + }; + let controller = SubscriptionController::start( + client.peer().clone(), + SubscriptionFilter::default(), + sink, + ) + .await + .unwrap(); + (controller, client) + } + + #[tokio::test] + async fn rejected_replacement_preserves_committed_stream_and_resources() { + let (mut controller, client) = controller().await; + controller + .subscribe_resource("accepted".into()) + .await + .unwrap(); + let before = controller.status(); + assert!( + controller + .subscribe_resource("rejected".into()) + .await + .is_err() + ); + assert_eq!(controller.status(), before); + assert_eq!( + controller.requested_resources(), + BTreeSet::from(["accepted".into()]) + ); + controller + .subscribe_resource("another".into()) + .await + .unwrap(); + assert_eq!(controller.requested_resources().len(), 2); + controller.shutdown().await; + assert!(matches!( + controller.status(), + McpSubscriptionStatus::Closed { error: None, .. } + )); + let _ = client.cancel().await; + } + + #[tokio::test] + async fn cancelled_caller_does_not_lose_enqueued_incremental_change() { + let (mut controller, client) = controller().await; + { + let update = controller.subscribe_resource("first".into()); + tokio::pin!(update); + assert!(futures_util::poll!(&mut update).is_pending()); + } + let settled = tokio::time::timeout(Duration::from_secs(2), controller.settled_resources()) + .await + .unwrap(); + assert_eq!(settled, BTreeSet::from(["first".into()])); + controller + .subscribe_resource("second".into()) + .await + .unwrap(); + assert_eq!( + controller.requested_resources(), + BTreeSet::from(["first".into(), "second".into()]) + ); + controller + .unsubscribe_resource("first".into()) + .await + .unwrap(); + assert_eq!( + controller.requested_resources(), + BTreeSet::from(["second".into()]) + ); + controller.shutdown().await; + let _ = client.cancel().await; + } +} diff --git a/crates/agentkit-mcp/tests/catalog_refresh.rs b/crates/agentkit-mcp/tests/catalog_refresh.rs new file mode 100644 index 0000000..69d1297 --- /dev/null +++ b/crates/agentkit-mcp/tests/catalog_refresh.rs @@ -0,0 +1,303 @@ +//! Catalog publication through the public manager API and a JSON-RPC HTTP boundary. +use agentkit_mcp::{ + AuthOperation, AuthRequest, AuthResolution, McpCatalogEvent, McpServerConfig, McpServerId, + McpServerManager, McpTransportBinding, StreamableHttpTransportConfig, +}; +use agentkit_tools_core::ToolSource; +use axum::{ + Json, Router, + extract::State, + http::{HeaderMap, StatusCode}, + response::{IntoResponse, Response}, + routing::post, +}; +use serde_json::{Value, json}; +use std::sync::{ + Arc, + atomic::{AtomicBool, AtomicU64, Ordering}, +}; +use std::time::Duration; + +#[derive(Clone, Default)] +struct Endpoint { + fail: Arc, + version: Arc, + lists: Arc, + authenticated: Arc, + resource_entered: Arc, + resource_release: Arc, +} + +async fn rpc( + State(state): State, + headers: HeaderMap, + Json(message): Json, +) -> Response { + if headers + .get("authorization") + .is_some_and(|value| value == "Bearer refreshed-token") + { + state.authenticated.store(true, Ordering::SeqCst); + } + let Some(id) = message.get("id") else { + return StatusCode::ACCEPTED.into_response(); + }; + let result = match message["method"].as_str().unwrap_or_default() { + "server/discover" => return (StatusCode::BAD_REQUEST, Json(json!({"jsonrpc":"2.0","id":id,"error":{"code":-32601,"message":"legacy endpoint"}}))).into_response(), + "initialize" => { + json!({"protocolVersion":"2024-11-05", "capabilities":{"tools":{}}, "serverInfo":{"name":"catalog-test","version":"1"}}) + } + "tools/list" => { + state.lists.fetch_add(1, Ordering::SeqCst); + if state.fail.load(Ordering::SeqCst) { + return Json(json!({"jsonrpc":"2.0", "id":id, "error":{"code":-32603,"message":"discovery failed"}})).into_response(); + } + json!({"tools":[{"name":format!("tool{}",state.version.load(Ordering::SeqCst)),"inputSchema":{"type":"object"}}]}) + } + "resources/subscribe" => { + state.resource_entered.notify_one(); + state.resource_release.notified().await; + json!({}) + } + "resources/list" => json!({"resources":[]}), + "prompts/list" => json!({"prompts":[]}), + _ => json!({}), + }; + Json(json!({"jsonrpc":"2.0","id":id,"result":result})).into_response() +} + +async fn endpoint(name: &str) -> (Endpoint, McpServerConfig, tokio::task::JoinHandle<()>) { + let state = Endpoint::default(); + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let url = format!("http://{}/mcp", listener.local_addr().unwrap()); + let app = Router::new() + .route("/mcp", post(rpc)) + .with_state(state.clone()); + let task = tokio::spawn(async move { + axum::serve(listener, app).await.unwrap(); + }); + ( + state, + McpServerConfig::new( + name, + McpTransportBinding::StreamableHttp(StreamableHttpTransportConfig::new(url)), + ), + task, + ) +} + +#[tokio::test] +async fn cancelled_and_superseded_work_stays_dirty_and_coalesces() { + let (state, config, task) = endpoint("catalog").await; + let id = McpServerId::new("catalog"); + let mut manager = McpServerManager::new().with_server(config); + manager.connect_server(&id).await.unwrap(); + manager.refresh_changed_catalogs().await.unwrap(); + assert!(manager.prepare_changed_catalogs().is_empty()); + + let mut wake = manager.catalog_refresh_wake(); + drop(manager.prepare_catalog_refresh(&id).unwrap()); + tokio::time::timeout(Duration::from_secs(2), wake.changed()) + .await + .unwrap(); + assert_eq!(manager.prepare_changed_catalogs().len(), 1); + let stale = manager + .prepare_changed_catalogs() + .pop() + .unwrap() + .fetch() + .await; + state.version.store(1, Ordering::SeqCst); + // A new invalidation after fetch fences the completed old discovery. + drop(manager.prepare_catalog_refresh(&id).unwrap()); + assert!(manager.publish_catalog_refresh(stale).unwrap().is_empty()); + assert_eq!( + manager.connected_server(&id).unwrap().snapshot().tools[0].name, + "tool0" + ); + let before = state.lists.load(Ordering::SeqCst); + drop(manager.prepare_catalog_refresh(&id).unwrap()); + drop(manager.prepare_catalog_refresh(&id).unwrap()); + manager.refresh_changed_catalogs().await.unwrap(); + assert_eq!(state.lists.load(Ordering::SeqCst), before + 1); + assert_eq!( + manager.connected_server(&id).unwrap().snapshot().tools[0].name, + "tool1" + ); + assert!(manager.prepare_changed_catalogs().is_empty()); + manager.disconnect_server(&id).await.unwrap(); + task.abort(); +} + +#[tokio::test] +async fn failed_server_does_not_block_other_publications() { + let (bad, bad_config, bad_task) = endpoint("a-bad").await; + let (good, good_config, good_task) = endpoint("z-good").await; + let mut manager = McpServerManager::new() + .with_server(bad_config) + .with_server(good_config); + manager.connect_all().await.unwrap(); + bad.fail.store(true, Ordering::SeqCst); + good.version.store(2, Ordering::SeqCst); + assert!(manager.refresh_changed_catalogs().await.is_err()); + let good_id = McpServerId::new("z-good"); + assert_eq!( + manager.connected_server(&good_id).unwrap().snapshot().tools[0].name, + "tool2" + ); + let pending = manager.prepare_changed_catalogs(); + assert_eq!(pending.len(), 1); + assert_eq!(pending[0].server_id(), &McpServerId::new("a-bad")); + bad.fail.store(false, Ordering::SeqCst); + manager.refresh_changed_catalogs().await.unwrap(); + assert!(manager.prepare_changed_catalogs().is_empty()); + bad_task.abort(); + good_task.abort(); +} + +#[tokio::test] +async fn replacement_rejects_old_connection_results() { + let (state, config, task) = endpoint("replace").await; + let id = McpServerId::new("replace"); + let mut manager = McpServerManager::new().with_server(config); + manager.connect_server(&id).await.unwrap(); + let old = manager.prepare_catalog_refresh(&id).unwrap().fetch().await; + state.version.store(3, Ordering::SeqCst); + manager.connect_server(&id).await.unwrap(); + assert!(manager.publish_catalog_refresh(old).unwrap().is_empty()); + assert_eq!( + manager.connected_server(&id).unwrap().snapshot().tools[0].name, + "tool3" + ); + manager.disconnect_server(&id).await.unwrap(); + task.abort(); +} + +#[tokio::test] +async fn auth_reconnect_fences_discovery_on_same_connection_arc() { + let (state, config, task) = endpoint("auth-reconnect").await; + let id = McpServerId::new("auth-reconnect"); + let mut manager = McpServerManager::new().with_server(config); + let handle = manager.connect_server(&id).await.unwrap(); + let connection = handle.connection(); + let old_epoch = connection.connection_epoch(); + let old = manager.prepare_catalog_refresh(&id).unwrap().fetch().await; + state.version.store(4, Ordering::SeqCst); + let request = AuthRequest { + id: "refresh-auth".into(), + provider: "test".into(), + operation: AuthOperation::McpConnect { + server_id: id.to_string(), + metadata: Default::default(), + }, + challenge: Default::default(), + }; + let mut credentials = agentkit_core::MetadataMap::new(); + credentials.insert("bearer_token".into(), json!("refreshed-token")); + tokio::time::timeout( + Duration::from_secs(5), + connection.resolve_auth(AuthResolution::provided(request, credentials)), + ) + .await + .unwrap() + .unwrap(); + assert!(state.authenticated.load(Ordering::SeqCst)); + assert!(Arc::ptr_eq( + &connection, + &manager.connected_server(&id).unwrap().connection() + )); + assert!(connection.connection_epoch() > old_epoch); + assert!(manager.publish_catalog_refresh(old).unwrap().is_empty()); + assert_eq!( + manager.connected_server(&id).unwrap().snapshot().tools[0].name, + "tool0" + ); + assert_eq!(manager.prepare_changed_catalogs().len(), 1); + manager.refresh_changed_catalogs().await.unwrap(); + assert_eq!( + manager.connected_server(&id).unwrap().snapshot().tools[0].name, + "tool4" + ); + assert!(manager.prepare_changed_catalogs().is_empty()); + manager.disconnect_server(&id).await.unwrap(); + task.abort(); +} + +// The legacy resources/subscribe request holds the connection lifecycle guard +// until this HTTP boundary replies. It establishes an actual cleanup await, +// rather than relying on a timer to cancel before or after the manager commit. +async fn hold_connection_lifecycle( + state: &Endpoint, + manager: &McpServerManager, + id: &McpServerId, +) -> tokio::task::JoinHandle<()> { + let connection = manager.connected_server(id).unwrap().connection(); + let pending = tokio::spawn(async move { + connection.subscribe_resource("memo:held").await.unwrap(); + }); + tokio::time::timeout(Duration::from_secs(5), state.resource_entered.notified()) + .await + .unwrap(); + pending +} + +#[tokio::test] +async fn cancelled_disconnect_has_already_removed_catalog_and_handle() { + let (state, config, task) = endpoint("disconnect-cancel").await; + let id = McpServerId::new("disconnect-cancel"); + let mut manager = McpServerManager::new().with_server(config); + manager.connect_server(&id).await.unwrap(); + let source = manager.source(); + assert_eq!(source.specs().len(), 1); + let pending = hold_connection_lifecycle(&state, &manager, &id).await; + let mut events = manager.subscribe_catalog_events(); + let mut disconnect = Box::pin(manager.disconnect_server(&id)); + tokio::time::timeout(Duration::from_secs(5), async { + tokio::select! { + result = &mut disconnect => panic!("cleanup unexpectedly completed: {result:?}"), + event = events.recv() => assert!(matches!(event.unwrap(), McpCatalogEvent::ServerDisconnected { .. })), + } + }).await.unwrap(); + drop(disconnect); + assert!(manager.connected_server(&id).is_none()); + assert!(source.specs().is_empty()); + state.resource_release.notify_one(); + pending.await.unwrap(); + task.abort(); +} + +#[tokio::test] +async fn cancelled_replacement_has_already_installed_new_catalog_and_handle() { + let (state, config, task) = endpoint("replace-cancel").await; + let id = McpServerId::new("replace-cancel"); + let mut manager = McpServerManager::new().with_server(config); + manager.connect_server(&id).await.unwrap(); + let source = manager.source(); + let old_connection = manager.connected_server(&id).unwrap().connection(); + let pending = hold_connection_lifecycle(&state, &manager, &id).await; + state.version.store(9, Ordering::SeqCst); + let mut events = manager.subscribe_catalog_events(); + let mut replace = Box::pin(manager.connect_server(&id)); + tokio::time::timeout(Duration::from_secs(5), async { + tokio::select! { + _ = &mut replace => panic!("cleanup unexpectedly completed"), + event = events.recv() => assert!(matches!(event.unwrap(), McpCatalogEvent::ServerConnected { .. })), + } + }).await.unwrap(); + drop(replace); + assert!(!Arc::ptr_eq( + &old_connection, + &manager.connected_server(&id).unwrap().connection() + )); + assert_eq!( + manager.connected_server(&id).unwrap().snapshot().tools[0].name, + "tool9" + ); + let names: Vec<_> = source.specs().into_iter().map(|spec| spec.name.0).collect(); + assert_eq!(names, ["mcp_replace-cancel_tool9"]); + state.resource_release.notify_one(); + pending.await.unwrap(); + old_connection.close().await.unwrap(); + manager.disconnect_server(&id).await.unwrap(); + task.abort(); +} diff --git a/crates/agentkit-mcp/tests/modern_http.rs b/crates/agentkit-mcp/tests/modern_http.rs new file mode 100644 index 0000000..0e99eb0 --- /dev/null +++ b/crates/agentkit-mcp/tests/modern_http.rs @@ -0,0 +1,433 @@ +//! Built-in Auto lifecycle through an actual HTTP JSON-RPC and POST SSE boundary. +use std::{convert::Infallible, future::Future, sync::Arc, time::Duration}; + +use agentkit_mcp::{ + AuthOperation, AuthRequest, AuthResolution, McpConnection, McpServerConfig, McpServerEvent, + McpSubscriptionStatus, McpTransportBinding, StreamableHttpTransportConfig, +}; +use axum::{ + Json, Router, + body::{Body, Bytes}, + extract::State, + http::{HeaderMap, StatusCode}, + response::{IntoResponse, Response}, + routing::post, +}; +use rmcp::model::{DiscoverResult, ProtocolVersion, ServerCapabilities}; +use serde_json::{Value, json}; +use tokio::sync::{Mutex, mpsc}; + +async fn bounded(future: impl Future) -> T { + tokio::time::timeout(Duration::from_secs(5), future) + .await + .expect("HTTP MCP operation timed out") +} + +struct Listen { + id: Value, + filter: Value, + headers: HeaderMap, + frames: mpsc::UnboundedSender, +} +impl Listen { + fn notify(&self, method: &str, params: Value) { + let mut params = params; + params["_meta"] = json!({"io.modelcontextprotocol/subscriptionId":self.id}); + self.frames + .send(frame( + json!({"jsonrpc":"2.0", "method":method, "params":params}), + )) + .unwrap(); + } +} +fn frame(message: Value) -> Bytes { + Bytes::from(format!("event: message\ndata: {message}\n\n")) +} + +#[derive(Clone)] +struct Endpoint { + listens: mpsc::UnboundedSender, + requests: Arc>>, +} + +async fn rpc( + State(state): State, + headers: HeaderMap, + Json(message): Json, +) -> Response { + state + .requests + .lock() + .await + .push((headers.clone(), message.clone())); + let Some(id) = message.get("id") else { + return StatusCode::ACCEPTED.into_response(); + }; + let token = headers + .get("authorization") + .and_then(|v| v.to_str().ok()) + .unwrap_or(""); + let result = match message["method"].as_str().unwrap_or_default() { + "server/discover" => { + if token.starts_with("Bearer legacy") { + return Json(json!({"jsonrpc":"2.0", "id":id, "error":{ + "code":-32601, "message":"legacy server has no discovery" + }})) + .into_response(); + } + if token == "Bearer fail" { + return Json(json!({"jsonrpc":"2.0","id":id,"error":{"code":-32603,"message":"test discovery failure"}})).into_response(); + } + let capabilities = ServerCapabilities::builder() + .enable_tools() + .enable_tool_list_changed() + .enable_resources() + .enable_resources_list_changed() + .enable_resources_subscribe(); + let capabilities = if token == "Bearer replacement" { + capabilities + .enable_prompts() + .enable_prompts_list_changed() + .build() + } else { + capabilities.build() + }; + serde_json::to_value(DiscoverResult::new( + vec![ProtocolVersion::V_2026_07_28], + capabilities, + )) + .unwrap() + } + "subscriptions/listen" => { + let (frames, receiver) = mpsc::unbounded_channel(); + let filter = message["params"]["notifications"].clone(); + if token != "Bearer hold" { + frames.send(frame(json!({"jsonrpc":"2.0", "method":"notifications/subscriptions/acknowledged", + "params":{"notifications":filter,"_meta":{"io.modelcontextprotocol/subscriptionId":id}}}))).unwrap(); + } + state + .listens + .send(Listen { + id: id.clone(), + filter, + headers, + frames, + }) + .unwrap_or_else(|_| panic!("fixture receiver alive")); + let stream = futures_util::stream::unfold(receiver, |mut receiver| async move { + receiver + .recv() + .await + .map(|bytes| (Ok::<_, Infallible>(bytes), receiver)) + }); + return ( + [("content-type", "text/event-stream")], + Body::from_stream(stream), + ) + .into_response(); + } + "initialize" => { + if token == "Bearer fail" { + return Json(json!({"jsonrpc":"2.0", "id":id, "error":{ + "code":-32603, "message":"test initialization failure" + }})) + .into_response(); + } + json!({ + "protocolVersion":"2024-11-05", + "capabilities":{"tools":{}, "resources":{"subscribe":true}}, + "serverInfo":{"name":"legacy-http-test", "version":"1"} + }) + } + "resources/subscribe" => { + if token == "Bearer legacy-fail" { + return Json(json!({"jsonrpc":"2.0", "id":id, "error":{ + "code":-32603, "message":"test subscription replay failure" + }})) + .into_response(); + } + json!({}) + } + "resources/unsubscribe" => json!({}), + "tools/list" => json!({"tools":[]}), + "resources/list" => json!({"resources":[]}), + "prompts/list" => json!({"prompts":[]}), + _ => return Json( + json!({"jsonrpc":"2.0","id":id,"error":{"code":-32601,"message":"unexpected method"}}), + ) + .into_response(), + }; + Json(json!({"jsonrpc":"2.0", "id":id, "result":result})).into_response() +} + +struct Fixture { + config: McpServerConfig, + listens: mpsc::UnboundedReceiver, + requests: Arc>>, + task: tokio::task::JoinHandle<()>, +} +impl Drop for Fixture { + fn drop(&mut self) { + self.task.abort(); + } +} +impl Fixture { + async fn start() -> Self { + let (tx, listens) = mpsc::unbounded_channel(); + let requests = Arc::new(Mutex::new(Vec::new())); + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let url = format!("http://{}/mcp", listener.local_addr().unwrap()); + let app = Router::new().route("/mcp", post(rpc)).with_state(Endpoint { + listens: tx, + requests: requests.clone(), + }); + let task = tokio::spawn(async move { + axum::serve(listener, app).await.unwrap(); + }); + let config = McpServerConfig::new( + "http-modern", + McpTransportBinding::StreamableHttp( + StreamableHttpTransportConfig::new(url).with_bearer_token("initial"), + ), + ); + Self { + config, + listens, + requests, + task, + } + } + async fn listen(&mut self) -> Listen { + bounded(self.listens.recv()).await.unwrap() + } +} + +fn auth(token: &str) -> AuthResolution { + AuthResolution::Provided { + request: AuthRequest { + id: "test".into(), + provider: "test".into(), + operation: AuthOperation::McpConnect { + server_id: "http-modern".into(), + metadata: Default::default(), + }, + challenge: Default::default(), + }, + credentials: [("bearer_token".into(), json!(token))] + .into_iter() + .collect(), + } +} + +#[tokio::test] +async fn auto_http_discovery_stream_metadata_and_auth_replacement_restore_resources() { + bounded(async { + let mut fixture = Fixture::start().await; + let connection = McpConnection::connect(&fixture.config).await.unwrap(); + let initial = fixture.listen().await; + assert_eq!(connection.protocol_version(), Some(ProtocolVersion::V_2026_07_28)); + assert!(matches!(connection.subscription_status().await, Some(McpSubscriptionStatus::Active { .. }))); + assert!(connection.capabilities().prompts.is_none()); + connection.discover().await.unwrap(); + let requests = fixture.requests.lock().await; + assert_eq!(requests[0].1["method"], "server/discover"); + assert!(requests.iter().all(|(_, message)| message["method"] != "initialize")); + let (headers, request) = requests.iter().find(|(_, message)| message["method"] == "subscriptions/listen").unwrap(); + assert_eq!(headers.get("mcp-protocol-version").unwrap(), "2026-07-28"); + assert_eq!(headers.get("mcp-method").unwrap(), "subscriptions/listen"); + let meta = &request["params"]["_meta"]; + assert_eq!(meta["io.modelcontextprotocol/protocolVersion"], "2026-07-28"); + assert!(meta["io.modelcontextprotocol/clientInfo"].is_object()); + assert!(meta["io.modelcontextprotocol/clientCapabilities"].is_object()); + drop(requests); + connection.subscribe_resource("memo:keep").await.unwrap(); + let subscribed = fixture.listen().await; + bounded(initial.frames.closed()).await; + let epoch = connection.connection_epoch(); + let revision = connection.catalog_revision(); + connection.resolve_auth(auth("replacement")).await.unwrap(); + let replacement = fixture.listen().await; + assert_eq!(replacement.headers.get("authorization").unwrap(), "Bearer replacement"); + assert_eq!(replacement.filter["resourceSubscriptions"], json!(["memo:keep"])); + assert_eq!(replacement.filter["promptsListChanged"], true); + assert!(connection.capabilities().prompts.is_some()); + assert!(connection.connection_epoch() > epoch); + assert!(connection.catalog_revision() > revision); + bounded(subscribed.frames.closed()).await; + let mut events = connection.subscribe_events(); + replacement.notify("notifications/resources/updated", json!({"uri":"memo:keep"})); + assert!(matches!(events.recv().await.unwrap(), McpServerEvent::ResourceUpdated(params) if params.uri == "memo:keep")); + connection.close().await.unwrap(); + bounded(replacement.frames.closed()).await; + }).await; +} + +#[tokio::test] +async fn failed_and_cancelled_auth_reconnect_leave_old_http_stream_usable() { + bounded(async { + let mut fixture = Fixture::start().await; + let connection = Arc::new(McpConnection::connect(&fixture.config).await.unwrap()); + let active = fixture.listen().await; + let epoch = connection.connection_epoch(); + let capabilities = connection.capabilities(); + let status = connection.subscription_status().await; + assert!(connection.resolve_auth(auth("fail")).await.is_err()); + assert_eq!(connection.connection_epoch(), epoch); + assert_eq!(connection.capabilities(), capabilities); + assert_eq!(connection.subscription_status().await, status); + let reconnect = { + let connection = connection.clone(); + tokio::spawn(async move { connection.resolve_auth(auth("hold")).await }) + }; + let pending = fixture.listen().await; + reconnect.abort(); + assert!(reconnect.await.unwrap_err().is_cancelled()); + bounded(pending.frames.closed()).await; + assert_eq!(connection.connection_epoch(), epoch); + assert_eq!(connection.capabilities(), capabilities); + assert_eq!(connection.subscription_status().await, status); + let mut events = connection.subscribe_events(); + active.notify("notifications/tools/list_changed", json!({})); + assert!(matches!( + events.recv().await.unwrap(), + McpServerEvent::ToolListChanged + )); + connection.discover().await.unwrap(); + connection.close().await.unwrap(); + bounded(active.frames.closed()).await; + }) + .await; +} + +#[tokio::test] +async fn modern_to_legacy_and_legacy_to_legacy_replay_committed_resource_uris() { + bounded(async { + let mut fixture = Fixture::start().await; + let connection = McpConnection::connect(&fixture.config).await.unwrap(); + let initial = fixture.listen().await; + connection.subscribe_resource("memo:keep").await.unwrap(); + let subscribed = fixture.listen().await; + bounded(initial.frames.closed()).await; + let mut epoch = connection.connection_epoch(); + for token in ["legacy-first", "legacy-second"] { + connection.resolve_auth(auth(token)).await.unwrap(); + assert_eq!( + connection.protocol_version(), + Some(ProtocolVersion::V_2024_11_05) + ); + assert!(connection.subscription_status().await.is_none()); + assert!(connection.connection_epoch() > epoch); + epoch = connection.connection_epoch(); + connection.discover().await.unwrap(); + let requests = fixture.requests.lock().await; + let authorization = format!("Bearer {token}"); + let candidate: Vec<_> = requests + .iter() + .filter(|(headers, _)| { + headers + .get("authorization") + .and_then(|value| value.to_str().ok()) + == Some(authorization.as_str()) + }) + .map(|(_, request)| request) + .collect(); + let initialized = candidate + .iter() + .position(|request| request["method"] == "initialize") + .unwrap(); + let replayed = candidate + .iter() + .position(|request| request["method"] == "resources/subscribe") + .unwrap(); + assert!(initialized < replayed); + assert_eq!(candidate[replayed]["params"]["uri"], "memo:keep"); + assert!( + candidate + .iter() + .all(|request| request["method"] != "subscriptions/listen") + ); + } + bounded(subscribed.frames.closed()).await; + // A replay rejection during legacy-to-legacy reconnect must not retire + // the second legacy transport or lose its committed resource set. + assert!(connection.resolve_auth(auth("legacy-fail")).await.is_err()); + assert_eq!(connection.connection_epoch(), epoch); + assert_eq!( + connection.protocol_version(), + Some(ProtocolVersion::V_2024_11_05) + ); + connection.list_resources().await.unwrap(); + { + let requests = fixture.requests.lock().await; + let (headers, request) = requests.last().unwrap(); + assert_eq!(request["method"], "resources/list"); + assert_eq!( + headers.get("authorization").unwrap(), + "Bearer legacy-second" + ); + } + // Crossing back also proves the committed URI set survived both legacy commits. + connection.resolve_auth(auth("replacement")).await.unwrap(); + let modern = fixture.listen().await; + assert_eq!(modern.filter["resourceSubscriptions"], json!(["memo:keep"])); + connection.close().await.unwrap(); + bounded(modern.frames.closed()).await; + }) + .await; +} + +#[tokio::test] +async fn rejected_legacy_resource_replay_keeps_authoritative_modern_connection() { + bounded(async { + let mut fixture = Fixture::start().await; + let connection = McpConnection::connect(&fixture.config).await.unwrap(); + let initial = fixture.listen().await; + connection.subscribe_resource("memo:keep").await.unwrap(); + let active = fixture.listen().await; + bounded(initial.frames.closed()).await; + let epoch = connection.connection_epoch(); + let capabilities = connection.capabilities(); + let status = connection.subscription_status().await; + assert!(connection.resolve_auth(auth("legacy-fail")).await.is_err()); + assert_eq!(connection.connection_epoch(), epoch); + assert_eq!(connection.protocol_version(), Some(ProtocolVersion::V_2026_07_28)); + assert_eq!(connection.capabilities(), capabilities); + assert_eq!(connection.subscription_status().await, status); + let requests = fixture.requests.lock().await; + assert!(requests.iter().any(|(headers, request)| { + headers.get("authorization").is_some_and(|value| value == "Bearer legacy-fail") + && request["method"] == "resources/subscribe" + && request["params"]["uri"] == "memo:keep" + })); + drop(requests); + let mut events = connection.subscribe_events(); + active.notify("notifications/resources/updated", json!({"uri":"memo:keep"})); + assert!(matches!(events.recv().await.unwrap(), McpServerEvent::ResourceUpdated(params) if params.uri == "memo:keep")); + connection.discover().await.unwrap(); + connection.close().await.unwrap(); + bounded(active.frames.closed()).await; + }).await; +} + +#[tokio::test] +async fn cancelling_initiated_close_terminates_http_stream_with_connection_retained() { + bounded(async { + let mut fixture = Fixture::start().await; + let connection = McpConnection::connect(&fixture.config).await.unwrap(); + let active = fixture.listen().await; + let epoch = connection.connection_epoch(); + let mut closing = Box::pin(connection.close()); + // One poll initiates shutdown with uncontended guards, then yields on worker cleanup. + // Drop precisely there rather than racing a spawned task against an arbitrary sleep. + assert!(futures_util::poll!(closing.as_mut()).is_pending()); + assert!( + connection.connection_epoch() > epoch, + "shutdown initiated before cancellation" + ); + drop(closing); + bounded(active.frames.closed()).await; + // Keep the connection alive throughout: stream disposal cannot be attributed to Drop. + assert!(connection.discover().await.is_err()); + }) + .await; +} diff --git a/crates/agentkit-mcp/tests/subscriptions.rs b/crates/agentkit-mcp/tests/subscriptions.rs new file mode 100644 index 0000000..4e7c531 --- /dev/null +++ b/crates/agentkit-mcp/tests/subscriptions.rs @@ -0,0 +1,236 @@ +//! Real modern RMCP peers: discovery, acknowledged streams, and public invalidations. +use std::{future::Future, time::Duration}; + +use agentkit_mcp::{ + McpConnection, McpHandlerConfig, McpServerEvent, McpServerId, McpSubscriptionStatus, +}; +use rmcp::model::{ProtocolVersion, ServerCapabilities, ServerConfig, SubscriptionFilter}; +use rmcp::service::{SubscriptionContext, SubscriptionSink}; +use rmcp::{ClientLifecycleMode, ClientServiceExt, ServerHandler, ServiceExt}; +use tokio::sync::{mpsc, oneshot}; + +async fn bounded(future: impl Future) -> T { + tokio::time::timeout(Duration::from_secs(15), future) + .await + .expect("MCP operation timed out") +} + +struct Stream { + sink: SubscriptionSink, + finish: oneshot::Sender<()>, + ended: oneshot::Receiver<()>, +} + +#[derive(Clone)] +struct Server { + streams: mpsc::UnboundedSender, + partial: bool, +} + +impl ServerHandler for Server { + fn get_info(&self) -> ServerConfig { + ServerConfig::new( + ServerCapabilities::builder() + .enable_tools() + .enable_tool_list_changed() + .enable_resources() + .enable_resources_list_changed() + .enable_resources_subscribe() + .enable_prompts() + .enable_prompts_list_changed() + .build(), + ) + } + + fn accepted_subscription_filter( + &self, + requested: &SubscriptionFilter, + ) -> Option { + let mut accepted = requested.clone(); + if self.partial { + accepted.prompts_list_changed = None; + } + if let Some(resources) = &mut accepted.resource_subscriptions { + resources.retain(|uri| uri != "memo:rejected"); + } + Some(accepted) + } + + async fn listen(&self, context: SubscriptionContext) -> Result<(), rmcp::ErrorData> { + let (finish, finished) = oneshot::channel(); + let (ended_tx, ended) = oneshot::channel(); + let _ = self.streams.send(Stream { + sink: context.sink().clone(), + finish, + ended, + }); + tokio::select! { + _ = context.cancelled() => {}, + _ = finished => {}, + } + let _ = ended_tx.send(()); + Ok(()) + } +} + +async fn connect(partial: bool, auto: bool) -> (McpConnection, mpsc::UnboundedReceiver) { + let (streams, incoming) = mpsc::unbounded_channel(); + let (server_io, client_io) = tokio::io::duplex(16 * 1024); + tokio::spawn(async move { + let running = Server { streams, partial }.serve(server_io).await.unwrap(); + let _ = running.waiting().await; + }); + let (handler, channels) = McpHandlerConfig::new().build(); + let versions = vec![ProtocolVersion::V_2026_07_28]; + let lifecycle = if auto { + ClientLifecycleMode::Auto { + preferred_versions: versions, + legacy_version: None, + } + } else { + ClientLifecycleMode::Discover { + preferred_versions: versions, + } + }; + let service = bounded(handler.serve_with_lifecycle(client_io, lifecycle)) + .await + .unwrap(); + let connection = McpConnection::from_running_service_with_events( + McpServerId::new("modern"), + service, + channels.notifications, + channels.events, + ); + (connection, incoming) +} + +#[tokio::test] +async fn auto_and_discover_establish_acknowledged_streams_and_bridge_catalog_events() { + bounded(async { + for auto in [false, true] { + let (connection, mut streams) = connect(false, auto).await; + assert_eq!( + connection.protocol_version(), + Some(ProtocolVersion::V_2026_07_28) + ); + assert!(connection.capabilities().tools.is_some()); + assert!( + connection.subscription_status().await.is_none(), + "adoption is synchronous" + ); + connection.discover().await.unwrap(); + let stream = streams.recv().await.unwrap(); + let Some(McpSubscriptionStatus::Active { + requested, + acknowledged, + }) = connection.subscription_status().await + else { + panic!("discovery must await acknowledgment"); + }; + assert_eq!(requested, acknowledged); + assert_eq!(acknowledged.tools_list_changed, Some(true)); + assert_eq!(acknowledged.resources_list_changed, Some(true)); + assert_eq!(acknowledged.prompts_list_changed, Some(true)); + let mut events = connection.subscribe_events(); + let mut dirty = connection.catalog_changes(); + let revision = connection.catalog_revision(); + stream.sink.notify_tool_list_changed().await.unwrap(); + assert!(matches!( + events.recv().await.unwrap(), + McpServerEvent::ToolListChanged + )); + stream.sink.notify_resource_list_changed().await.unwrap(); + assert!(matches!( + events.recv().await.unwrap(), + McpServerEvent::ResourceListChanged + )); + stream.sink.notify_prompt_list_changed().await.unwrap(); + assert!(matches!( + events.recv().await.unwrap(), + McpServerEvent::PromptListChanged + )); + dirty.changed().await.unwrap(); + assert!(connection.catalog_revision() >= revision + 3); + let epoch = connection.connection_epoch(); + connection.close().await.unwrap(); + // RMCP drains server handler responses for five seconds after transport + // EOF before dropping RunningService and canceling its request contexts. + // Client transport cancellation need not deliver a wire-level request + // cancellation first, so keep a separate bound above that SDK drain. + tokio::time::timeout(Duration::from_secs(7), stream.ended) + .await + .expect("server subscription handler remained alive after transport EOF") + .expect("server subscription handler must observe cancellation"); + assert!(connection.connection_epoch() > epoch); + } + }) + .await; +} + +#[tokio::test] +async fn resource_replacement_rejection_preserves_active_filter_and_stream() { + bounded(async { + let (connection, mut streams) = connect(false, false).await; + connection.discover().await.unwrap(); + let _initial = streams.recv().await.unwrap(); + connection.subscribe_resource("memo:accepted").await.unwrap(); + let accepted = streams.recv().await.unwrap(); + let before = connection.subscription_status().await.unwrap(); + assert!(connection.subscribe_resource("memo:rejected").await.is_err()); + let _rejected = streams.recv().await.unwrap(); + assert_eq!(connection.subscription_status().await.unwrap(), before); + let mut events = connection.subscribe_events(); + accepted.sink.notify_resource_updated("memo:accepted").await.unwrap(); + assert!(matches!(events.recv().await.unwrap(), McpServerEvent::ResourceUpdated(params) if params.uri == "memo:accepted")); + connection.unsubscribe_resource("memo:accepted").await.unwrap(); + let _unsubscribed = streams.recv().await.unwrap(); + let Some(McpSubscriptionStatus::Active { requested, acknowledged }) = connection.subscription_status().await else { panic!("active replacement"); }; + assert_eq!(requested, acknowledged); + assert!(requested.resource_subscriptions.as_ref().is_none_or(Vec::is_empty)); + connection.close().await.unwrap(); + }).await; +} + +#[tokio::test] +async fn partial_catalog_acknowledgment_and_graceful_completion_are_observable() { + bounded(async { + let (connection, mut streams) = connect(true, false).await; + connection.discover().await.unwrap(); + let stream = streams.recv().await.unwrap(); + let Some(McpSubscriptionStatus::Active { + requested, + acknowledged, + }) = connection.subscription_status().await + else { + panic!("active stream"); + }; + assert_eq!(requested.prompts_list_changed, Some(true)); + assert_eq!(acknowledged.prompts_list_changed, None); + let mut status = connection.subscription_changes().await.unwrap(); + stream.finish.send(()).unwrap(); + loop { + status.changed().await.unwrap(); + if matches!( + *status.borrow_and_update(), + McpSubscriptionStatus::Closed { graceful: true, .. } + ) { + break; + } + } + assert!( + streams.try_recv().is_err(), + "graceful completion must not restart" + ); + connection.close().await.unwrap(); + // A fresh connection starts clean instead of inheriting the closed stream. + let (next, mut next_streams) = connect(false, true).await; + next.discover().await.unwrap(); + let _stream = next_streams.recv().await.unwrap(); + assert!(matches!( + next.subscription_status().await, + Some(McpSubscriptionStatus::Active { .. }) + )); + next.close().await.unwrap(); + }) + .await; +} diff --git a/crates/agentkit/Cargo.toml b/crates/agentkit/Cargo.toml index bf96fda..38d5234 100644 --- a/crates/agentkit/Cargo.toml +++ b/crates/agentkit/Cargo.toml @@ -4,7 +4,7 @@ homepage.workspace = true name = "agentkit" readme = "README.md" repository.workspace = true -version = "0.10.12" +version = "0.11.0" edition.workspace = true license.workspace = true rust-version.workspace = true @@ -16,7 +16,7 @@ agentkit-compaction = { version = "0.10.5", path = "../agentkit-compaction", opt agentkit-context = { version = "0.10.5", path = "../agentkit-context", optional = true } agentkit-core = { version = "0.10.5", path = "../agentkit-core", optional = true } agentkit-loop = { version = "0.10.11", path = "../agentkit-loop", optional = true } -agentkit-mcp = { version = "0.10.5", path = "../agentkit-mcp", optional = true } +agentkit-mcp = { version = "0.11.0", path = "../agentkit-mcp", optional = true } agentkit-plugins = { version = "0.10.6", path = "../agentkit-plugins", optional = true } agentkit-adapter-completions = { version = "0.10.8", path = "../agentkit-adapter-completions", optional = true } agentkit-provider-anthropic = { version = "0.10.9", path = "../agentkit-provider-anthropic", optional = true } diff --git a/docs/mcp.md b/docs/mcp.md index fda90bb..182a75e 100644 --- a/docs/mcp.md +++ b/docs/mcp.md @@ -91,7 +91,7 @@ MCP server capabilities may change over time, but hosts should not be forced to MCP is bidirectional — servers can issue `sampling/createMessage`, `elicitation/create`, and `roots/list` requests back into the client. The host installs `McpSamplingResponder` / `McpElicitationResponder` / `McpRootsProvider` implementations and the client only advertises the matching `ClientCapabilities` entry when a responder is wired in. -Server-pushed notifications (progress, logging, resource updates, list-changed, cancellation) are delivered as `McpServerEvent` over a `tokio::sync::broadcast` channel obtained via `McpConnection::subscribe_events`. List-changed events are _also_ delivered through the legacy `McpServerNotification` mpsc receiver consumed by `refresh_changed_catalogs` — the two channels coexist: events for live UI/observability, mpsc for catalog re-sync. +Server-pushed notifications (progress, logging, resource updates, list-changed, cancellation) are delivered as `McpServerEvent` over a `tokio::sync::broadcast` channel obtained via `McpConnection::subscribe_events`. List-changed events also advance a coalescing catalog revision watch used for catalog reconciliation. The compatibility `McpServerNotification` receiver remains available; hosts do not need to drain broadcast events to keep catalog invalidations. ## Main boundary @@ -202,6 +202,51 @@ pub enum McpServerEvent { Subscribers fan out from `McpConnection::subscribe_events()`. `McpConnection::set_logging_level(...)` negotiates the minimum severity the server emits; `subscribe_resource(uri)` / `unsubscribe_resource(uri)` toggle per-URI watch on resources that support it. +## Negotiation and notification subscriptions + +Built-in stdio and Streamable HTTP connections use RMCP's `Auto` lifecycle: probe +`server/discover`, then fall back to legacy initialization when needed. Inspect +`connection.protocol_version()` and the cloned `connection.capabilities()` for +the negotiated peer, rather than assuming a protocol from the transport. + +For modern peers, discovery establishes a connection-owned `subscriptions/listen` +stream and waits for its acknowledgment. The requested filter covers advertised +tool, resource, and prompt catalogs. `subscription_status().await` exposes both +the requested and acknowledged filters: a server may accept only part of the +catalog filter. `subscription_changes().await` returns a watch receiver once a +modern stream exists; legacy peers return `None`. Adopting an existing RMCP +service with `from_running_service_with_events` is synchronous: call +`discover().await` to establish its modern subscription lazily. + +`subscribe_resource(uri)` and `unsubscribe_resource(uri)` replace the modern +stream's complete filter. A resource update is committed only after the server +acknowledges the requested URI set; rejection preserves the previous active +filter and stream. Legacy peers retain `resources/subscribe` and +`resources/unsubscribe`. Streamed notifications feed the same +`subscribe_events()` broadcast used by legacy notifications. + +Catalog invalidations also advance `catalog_revision()` and wake +`catalog_changes()`, a coalescing watch channel that does not require draining +broadcast events. Watch wakeups are invalidations, not newly published catalogs. +Hosts should call `McpServerManager::refresh_changed_catalogs()` to reconcile +snapshots, or use `prepare_changed_catalogs()`, fetch work outside their manager +lock, and `publish_catalog_refresh(...)` to publish it. Revision and connection +epoch checks fence obsolete results; canceled or failed fetches remain dirty. +A host must schedule this refresh work itself; event delivery does not silently +refresh the host's tool registry. + +On transport failure, the modern subscription worker retries with bounded +backoff and marks catalogs dirty for reconciliation. Graceful server completion +is observable as `Closed { graceful: true, .. }` and is not retried. Authentication +reconnect replaces the subscription on the new peer, retaining committed resource +URIs. `connection_epoch()` changes when the transport is replaced or closed. +Call `close().await` to stop owned streams before transport cleanup; dropping the +connection also cancels its workers. + +**Migration (agentkit-mcp 0.11.0):** `capabilities()` now returns an owned snapshot, +not a borrowed reference. Call it again after authentication reconnect to obtain +the current peer capabilities; no capability lock guard needs to span an await. + ### 8. Errors `McpError` covers transport failure, protocol errors, auth challenges (`AuthRequired(Box)`), invocation errors, and unknown-server lookup misses. It implements `From for rmcp::model::ErrorData` so responders can return agentkit errors and have them surface to the server as JSON-RPC errors. From b6af4a0902cb5d5a9365b05e04dd64e1fbbc2fb2 Mon Sep 17 00:00:00 2001 From: daniel Date: Thu, 15 Oct 2026 11:53:21 +0100 Subject: [PATCH 2/3] fix(mcp): preserve subscriptions across replay and recovery Reject partial resource replay acknowledgments before replacing the authoritative connection, and invalidate catalogs after stream recovery. Add HTTP-boundary regressions and explicit legacy/modern reference interop modes without weakening legacy session and resumption checks. --- crates/agentkit-mcp/src/subscriptions.rs | 4 + crates/agentkit-mcp/tests/modern_http.rs | 154 +++++++++++++++++- examples/mcp-reference-interop/Cargo.toml | 2 +- examples/mcp-reference-interop/README.md | 20 ++- .../src/bin/mcp-reference-server.rs | 76 +++++++-- examples/mcp-reference-interop/src/lib.rs | 4 + examples/mcp-reference-interop/tests/live.rs | 85 +++++++++- 7 files changed, 318 insertions(+), 27 deletions(-) diff --git a/crates/agentkit-mcp/src/subscriptions.rs b/crates/agentkit-mcp/src/subscriptions.rs index e150aac..4314bc3 100644 --- a/crates/agentkit-mcp/src/subscriptions.rs +++ b/crates/agentkit-mcp/src/subscriptions.rs @@ -64,6 +64,7 @@ impl SubscriptionController { sink: McpNotificationSink, ) -> Result { let subscription = establish(&peer, base_filter.clone()).await?; + verify_resources(&base_filter, subscription.acknowledged())?; let initial = active_status(&base_filter, &subscription); let (status_tx, status) = watch::channel(initial); let (requested_tx, requested) = watch::channel(resource_set(&base_filter)); @@ -326,6 +327,9 @@ async fn run( return McpSubscriptionStatus::Closed { graceful: false, error: Some(last_error) }; }; active = stream; + // A host may have refreshed during backoff. Reconcile the + // entire notification gap after this stream is active. + sink.mark_catalog_dirty(); status.send_replace(active_status(&filter, &active)); } } diff --git a/crates/agentkit-mcp/tests/modern_http.rs b/crates/agentkit-mcp/tests/modern_http.rs index 0e99eb0..e43e476 100644 --- a/crates/agentkit-mcp/tests/modern_http.rs +++ b/crates/agentkit-mcp/tests/modern_http.rs @@ -1,9 +1,18 @@ //! Built-in Auto lifecycle through an actual HTTP JSON-RPC and POST SSE boundary. -use std::{convert::Infallible, future::Future, sync::Arc, time::Duration}; +use std::{ + convert::Infallible, + future::Future, + sync::{ + Arc, + atomic::{AtomicBool, AtomicU64, Ordering}, + }, + time::Duration, +}; use agentkit_mcp::{ AuthOperation, AuthRequest, AuthResolution, McpConnection, McpServerConfig, McpServerEvent, - McpSubscriptionStatus, McpTransportBinding, StreamableHttpTransportConfig, + McpServerId, McpServerManager, McpSubscriptionStatus, McpTransportBinding, + StreamableHttpTransportConfig, }; use axum::{ Json, Router, @@ -30,6 +39,20 @@ struct Listen { frames: mpsc::UnboundedSender, } impl Listen { + fn acknowledge(&self, filter: Value) { + self.notify( + "notifications/subscriptions/acknowledged", + json!({"notifications":filter}), + ); + } + + fn fail(&self) { + self.frames + .send(frame(json!({"jsonrpc":"2.0", "id":self.id, + "error":{"code":-32603,"message":"listen failed; ordinary RPC remains healthy"}}))) + .unwrap(); + } + fn notify(&self, method: &str, params: Value) { let mut params = params; params["_meta"] = json!({"io.modelcontextprotocol/subscriptionId":self.id}); @@ -46,6 +69,8 @@ fn frame(message: Value) -> Bytes { #[derive(Clone)] struct Endpoint { + hold_ack: Arc, + tool_version: Arc, listens: mpsc::UnboundedSender, requests: Arc>>, } @@ -101,7 +126,7 @@ async fn rpc( "subscriptions/listen" => { let (frames, receiver) = mpsc::unbounded_channel(); let filter = message["params"]["notifications"].clone(); - if token != "Bearer hold" { + if token != "Bearer hold" && !state.hold_ack.load(Ordering::SeqCst) { frames.send(frame(json!({"jsonrpc":"2.0", "method":"notifications/subscriptions/acknowledged", "params":{"notifications":filter,"_meta":{"io.modelcontextprotocol/subscriptionId":id}}}))).unwrap(); } @@ -149,7 +174,9 @@ async fn rpc( json!({}) } "resources/unsubscribe" => json!({}), - "tools/list" => json!({"tools":[]}), + "tools/list" => { + json!({"tools":[{"name":format!("tool{}", state.tool_version.load(Ordering::SeqCst)),"inputSchema":{"type":"object"}}]}) + } "resources/list" => json!({"resources":[]}), "prompts/list" => json!({"prompts":[]}), _ => return Json( @@ -161,6 +188,8 @@ async fn rpc( } struct Fixture { + hold_ack: Arc, + tool_version: Arc, config: McpServerConfig, listens: mpsc::UnboundedReceiver, requests: Arc>>, @@ -174,10 +203,14 @@ impl Drop for Fixture { impl Fixture { async fn start() -> Self { let (tx, listens) = mpsc::unbounded_channel(); + let hold_ack = Arc::new(AtomicBool::new(false)); + let tool_version = Arc::new(AtomicU64::new(0)); let requests = Arc::new(Mutex::new(Vec::new())); let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); let url = format!("http://{}/mcp", listener.local_addr().unwrap()); let app = Router::new().route("/mcp", post(rpc)).with_state(Endpoint { + hold_ack: hold_ack.clone(), + tool_version: tool_version.clone(), listens: tx, requests: requests.clone(), }); @@ -191,6 +224,8 @@ impl Fixture { ), ); Self { + hold_ack, + tool_version, config, listens, requests, @@ -431,3 +466,114 @@ async fn cancelling_initiated_close_terminates_http_stream_with_connection_retai }) .await; } + +#[tokio::test] +async fn auth_reconnect_partial_ack_preserves_committed_http_subscription() { + bounded(async { + let mut fixture = Fixture::start().await; + let connection = Arc::new(McpConnection::connect(&fixture.config).await.unwrap()); + let initial = fixture.listen().await; + connection.subscribe_resource("memo:first").await.unwrap(); + let first = fixture.listen().await; + bounded(initial.frames.closed()).await; + connection.subscribe_resource("memo:second").await.unwrap(); + let active = fixture.listen().await; + bounded(first.frames.closed()).await; + let epoch = connection.connection_epoch(); + let capabilities = connection.capabilities(); + let status = connection.subscription_status().await; + let mut events = connection.subscribe_events(); + + // Gate only the candidate's actual POST SSE acknowledgement. Discovery + // advertises different capabilities, so neither discovery nor a partial + // acknowledgement may commit any part of this authentication change. + fixture.hold_ack.store(true, Ordering::SeqCst); + let reconnect = { + let connection = connection.clone(); + tokio::spawn(async move { connection.resolve_auth(auth("replacement")).await }) + }; + let candidate = fixture.listen().await; + assert_eq!(candidate.headers.get("authorization").unwrap(), "Bearer replacement"); + assert_eq!(candidate.filter["resourceSubscriptions"], json!(["memo:first", "memo:second"])); + let mut partial = candidate.filter.clone(); + partial["resourceSubscriptions"] = json!(["memo:first"]); + candidate.acknowledge(partial); + assert!(bounded(reconnect).await.unwrap().is_err()); + bounded(candidate.frames.closed()).await; + assert_eq!(connection.connection_epoch(), epoch); + assert_eq!(connection.capabilities(), capabilities); + assert_eq!(connection.subscription_status().await, status); + assert!(!active.frames.is_closed()); + for uri in ["memo:first", "memo:second"] { + active.notify("notifications/resources/updated", json!({"uri":uri})); + assert!(matches!(events.recv().await.unwrap(), McpServerEvent::ResourceUpdated(params) if params.uri == uri)); + } + connection.list_resources().await.unwrap(); + let requests = fixture.requests.lock().await; + let (headers, request) = requests.last().unwrap(); + assert_eq!(request["method"], "resources/list"); + assert_eq!(headers.get("authorization").unwrap(), "Bearer initial"); + drop(requests); + connection.close().await.unwrap(); + bounded(active.frames.closed()).await; + }).await; +} + +#[tokio::test] +async fn recovered_http_listen_redirties_catalog_after_refresh_during_gap() { + bounded(async { + let mut fixture = Fixture::start().await; + let id = McpServerId::new("http-modern"); + let mut manager = McpServerManager::new().with_server(fixture.config.clone()); + manager.connect_server(&id).await.unwrap(); + let active = fixture.listen().await; + let connection = manager.connected_server(&id).unwrap().connection(); + manager.refresh_changed_catalogs().await.unwrap(); + assert!(manager.prepare_changed_catalogs().is_empty()); + assert_eq!( + manager.connected_server(&id).unwrap().snapshot().tools[0].name, + "tool0" + ); + + // A JSON-RPC error terminates the listen, not the HTTP endpoint. Hold + // recovery at the next POST SSE ACK while ordinary catalog RPCs work. + fixture.hold_ack.store(true, Ordering::SeqCst); + active.fail(); + let recovery = fixture.listen().await; + assert_eq!( + connection.subscription_status().await, + Some(McpSubscriptionStatus::Pending) + ); + manager.refresh_changed_catalogs().await.unwrap(); + assert!(manager.prepare_changed_catalogs().is_empty()); + assert_eq!( + manager.connected_server(&id).unwrap().snapshot().tools[0].name, + "tool0" + ); + assert_eq!( + connection.subscription_status().await, + Some(McpSubscriptionStatus::Pending) + ); + + // Change the remote catalog after the host has consumed the first dirty + // mark. Send no list-changed notification, only the recovery ACK. + fixture.tool_version.store(1, Ordering::SeqCst); + recovery.acknowledge(recovery.filter.clone()); + while !matches!( + connection.subscription_status().await, + Some(McpSubscriptionStatus::Active { .. }) + ) { + tokio::task::yield_now().await; + } + assert_eq!(manager.prepare_changed_catalogs().len(), 1); + manager.refresh_changed_catalogs().await.unwrap(); + assert_eq!( + manager.connected_server(&id).unwrap().snapshot().tools[0].name, + "tool1" + ); + assert!(manager.prepare_changed_catalogs().is_empty()); + manager.disconnect_server(&id).await.unwrap(); + bounded(recovery.frames.closed()).await; + }) + .await; +} diff --git a/examples/mcp-reference-interop/Cargo.toml b/examples/mcp-reference-interop/Cargo.toml index 5adf601..f175b9e 100644 --- a/examples/mcp-reference-interop/Cargo.toml +++ b/examples/mcp-reference-interop/Cargo.toml @@ -9,7 +9,7 @@ version.workspace = true [dependencies] agentkit-mcp = { path = "../../crates/agentkit-mcp" } axum = "0.8.4" -rmcp = { version = "3.1.2", features = ["server", "transport-streamable-http-server"] } +rmcp = { version = "=3.5.1", features = ["server", "transport-streamable-http-server"] } schemars = "1.0.4" reqwest.workspace = true serde = { workspace = true, features = ["derive"] } diff --git a/examples/mcp-reference-interop/README.md b/examples/mcp-reference-interop/README.md index b12625e..d1445f5 100644 --- a/examples/mcp-reference-interop/README.md +++ b/examples/mcp-reference-interop/README.md @@ -3,21 +3,23 @@ Live Streamable HTTP interop checks for `agentkit-mcp` against the official Rust MCP SDK, `rmcp`. -This example does not mock the transport layer. It launches two `rmcp`-backed +This example does not mock the transport layer. It launches three `rmcp`-backed reference servers and probes them through the real `agentkit-mcp` client: -- stateful Streamable HTTP with SSE responses -- stateless Streamable HTTP with JSON responses +- legacy (2025-11-25) stateful Streamable HTTP with SSE responses +- legacy (2025-11-25) stateless Streamable HTTP with JSON responses +- modern (2026-07-28) sessionless Streamable HTTP with acknowledged subscriptions Each probe verifies: -- `initialize` and capability discovery +- modern `server/discover` or legacy `initialize` and capability discovery - tool listing and `tools/call` - resource listing and `resources/read` - prompt listing and `prompts/get` - negotiated protocol header propagation - stateful session header reuse and `DELETE` on close -- interrupted SSE response resume via `Last-Event-ID` +- legacy interrupted SSE response resume via `Last-Event-ID` +- modern `subscriptions/listen` without initialization, GET streams, or session cleanup ## Run @@ -33,6 +35,12 @@ Probe the stateless JSON-response server: cargo run -p mcp-reference-interop -- stateless ``` +Probe the modern subscription server: + +```bash +cargo run -p mcp-reference-interop -- modern +``` + ## Live Tests The integration tests are self-contained and run with Cargo alone. @@ -57,5 +65,5 @@ cargo test -p mcp-reference-interop -- --nocapture ## Upstream Targets -- Rust SDK crate: `rmcp = 3.1.2` +- Rust SDK crate: `rmcp = 3.5.1` (resolved by the workspace lockfile) - Upstream repository: diff --git a/examples/mcp-reference-interop/src/bin/mcp-reference-server.rs b/examples/mcp-reference-interop/src/bin/mcp-reference-server.rs index 095b610..a18c702 100644 --- a/examples/mcp-reference-interop/src/bin/mcp-reference-server.rs +++ b/examples/mcp-reference-interop/src/bin/mcp-reference-server.rs @@ -19,9 +19,10 @@ use rmcp::{ wrapper::Parameters, }, model::{ - GetPromptResult, Implementation, ListResourcesResult, PaginatedRequestParams, - PromptMessage, ReadResourceRequestParams, ReadResourceResponse, ReadResourceResult, - Resource, ResourceContents, Role, ServerCapabilities, ServerInfo, + DiscoverResult, GetPromptResult, Implementation, ListResourcesResult, + PaginatedRequestParams, PromptMessage, ProtocolVersion, ReadResourceRequestParams, + ReadResourceResponse, ReadResourceResult, Resource, ResourceContents, Role, + ServerCapabilities, ServerConfig, }, prompt, prompt_handler, prompt_router, tool, tool_handler, tool_router, transport::streamable_http_server::{ @@ -41,11 +42,13 @@ struct AppState { enum ServerMode { Stateful, StatelessJson, + Modern, } impl ServerMode { fn parse(value: &str) -> Option { match value { + "modern" => Some(Self::Modern), "stateful" => Some(Self::Stateful), "stateless-json" => Some(Self::StatelessJson), _ => None, @@ -55,6 +58,7 @@ impl ServerMode { fn tool_prefix(&self) -> &'static str { match self { Self::Stateful => "rmcp-stateful", + Self::Modern => "rmcp-modern", Self::StatelessJson => "rmcp-stateless", } } @@ -62,6 +66,7 @@ impl ServerMode { fn resource_text(&self) -> &'static str { match self { Self::Stateful => "rust sdk stateful fixture resource", + Self::Modern => "rust sdk modern fixture resource", Self::StatelessJson => "rust sdk stateless fixture resource", } } @@ -69,6 +74,7 @@ impl ServerMode { fn prompt_text(&self, name: &str) -> String { match self { Self::Stateful => format!("rust sdk stateful prompt for {name}"), + Self::Modern => format!("rust sdk modern prompt for {name}"), Self::StatelessJson => format!("rust sdk stateless prompt for {name}"), } } @@ -76,6 +82,7 @@ impl ServerMode { fn server_name(&self) -> &'static str { match self { Self::Stateful => "agentkit-rmcp-stateful", + Self::Modern => "agentkit-rmcp-modern", Self::StatelessJson => "agentkit-rmcp-stateless-json", } } @@ -84,6 +91,7 @@ impl ServerMode { let config = StreamableHttpServerConfig::default(); match self { Self::Stateful => config, + Self::Modern => config.with_legacy_session_mode(false), Self::StatelessJson => config .with_legacy_session_mode(false) .with_json_response(true) @@ -144,15 +152,59 @@ impl ReferenceServer { #[tool_handler(router = self.tool_router)] #[prompt_handler(router = self.prompt_router)] impl ServerHandler for ReferenceServer { - fn get_info(&self) -> ServerInfo { - ServerInfo::new( - ServerCapabilities::builder() - .enable_tools() - .enable_prompts() - .enable_resources() - .build(), - ) - .with_server_info(Implementation::new(self.mode.server_name(), "1.0.0")) + fn get_info(&self) -> ServerConfig { + let mut capabilities = ServerCapabilities::builder() + .enable_tools() + .enable_prompts() + .enable_resources(); + if self.mode == ServerMode::Modern { + capabilities = capabilities + .enable_tool_list_changed() + .enable_prompts_list_changed() + .enable_resources_list_changed(); + } + ServerConfig::new(capabilities.build()) + .with_server_info(Implementation::new(self.mode.server_name(), "1.0.0")) + .with_protocol_version(if self.mode == ServerMode::Modern { + ProtocolVersion::V_2026_07_28 + } else { + ProtocolVersion::V_2025_11_25 + }) + } + + async fn discover( + &self, + _context: rmcp::service::RequestContext, + ) -> Result { + if self.mode != ServerMode::Modern { + // Session and resumption modes deliberately implement the legacy era. + return Err(ErrorData::method_not_found::< + rmcp::model::DiscoverRequestMethod, + >()); + } + Ok(DiscoverResult::from_server_info( + vec![ProtocolVersion::V_2026_07_28], + self.get_info(), + )) + } + + fn accepted_subscription_filter( + &self, + requested: &rmcp::model::SubscriptionFilter, + ) -> Option { + (self.mode == ServerMode::Modern).then(|| requested.clone()) + } + + async fn listen(&self, context: rmcp::service::SubscriptionContext) -> Result<(), ErrorData> { + if self.mode != ServerMode::Modern { + return Err(ErrorData::method_not_found::< + rmcp::model::SubscriptionsListenRequestMethod, + >()); + } + // RMCP acknowledges the accepted subset before invoking this handler. + // Keep the POST stream alive until this client's subscription ends. + context.cancelled().await; + Ok(()) } async fn list_resources( diff --git a/examples/mcp-reference-interop/src/lib.rs b/examples/mcp-reference-interop/src/lib.rs index 79dd179..1381015 100644 --- a/examples/mcp-reference-interop/src/lib.rs +++ b/examples/mcp-reference-interop/src/lib.rs @@ -32,11 +32,13 @@ pub struct RecordedRequest { pub enum ReferenceImplementation { RustSdkStatefulSse, RustSdkStatelessJson, + RustSdkModern, } impl ReferenceImplementation { pub fn parse(value: &str) -> Option { match value { + "rust-modern" | "modern" => Some(Self::RustSdkModern), "rust-stateful" | "stateful" => Some(Self::RustSdkStatefulSse), "rust-stateless" | "stateless" | "json" => Some(Self::RustSdkStatelessJson), _ => None, @@ -45,6 +47,7 @@ impl ReferenceImplementation { pub fn as_str(&self) -> &'static str { match self { + Self::RustSdkModern => "rust-modern", Self::RustSdkStatefulSse => "rust-stateful", Self::RustSdkStatelessJson => "rust-stateless", } @@ -52,6 +55,7 @@ impl ReferenceImplementation { fn server_mode_arg(&self) -> &'static str { match self { + Self::RustSdkModern => "modern", Self::RustSdkStatefulSse => "stateful", Self::RustSdkStatelessJson => "stateless-json", } diff --git a/examples/mcp-reference-interop/tests/live.rs b/examples/mcp-reference-interop/tests/live.rs index 4a4d862..c75879f 100644 --- a/examples/mcp-reference-interop/tests/live.rs +++ b/examples/mcp-reference-interop/tests/live.rs @@ -18,6 +18,57 @@ use serde_json::{Value, json}; use tokio::sync::Mutex; use tokio_util::sync::CancellationToken; +#[tokio::test] +async fn rust_sdk_modern_subscription_streamable_http_interop() { + let result = probe_reference_implementation_with_binary( + ReferenceImplementation::RustSdkModern, + env!("CARGO_BIN_EXE_mcp-reference-server"), + ) + .await + .unwrap(); + assert_eq!(result.tool_output, "rmcp-modern:streamable-http-ok"); + assert_eq!(result.resource_text, "rust sdk modern fixture resource"); + assert_eq!(result.prompt_text, "rust sdk modern prompt for Ada"); + let requests = capture_transport_exchange_with_binary( + ReferenceImplementation::RustSdkModern, + env!("CARGO_BIN_EXE_mcp-reference-server"), + ) + .await + .unwrap(); + find_jsonrpc_request(&requests, "server/discover"); + let listen = find_jsonrpc_request(&requests, "subscriptions/listen"); + assert_eq!( + listen + .headers + .get("mcp-protocol-version") + .map(String::as_str), + Some("2026-07-28") + ); + assert_eq!( + listen.headers.get("mcp-method").map(String::as_str), + Some("subscriptions/listen") + ); + let body: Value = serde_json::from_str(&listen.body).unwrap(); + assert_eq!(body["params"]["notifications"]["toolsListChanged"], true); + assert_eq!( + body["params"]["notifications"]["resourcesListChanged"], + true + ); + assert_eq!(body["params"]["notifications"]["promptsListChanged"], true); + assert_eq!( + body["params"]["_meta"]["io.modelcontextprotocol/protocolVersion"], + "2026-07-28" + ); + assert!( + requests.iter().all(|request| { + !request.headers.contains_key("mcp-session-id") + && !matches!(request.method.as_str(), "GET" | "DELETE") + && request.jsonrpc_method.as_deref() != Some("initialize") + }), + "modern exchange must not establish or resume a legacy session: {requests:#?}" + ); +} + #[tokio::test] async fn rust_sdk_stateful_streamable_http_interop() { let result = probe_reference_implementation_with_binary( @@ -83,6 +134,8 @@ async fn rust_sdk_stateful_transport_propagates_session_protocol_and_delete() { .await .unwrap(); + let discover = find_jsonrpc_request(&requests, "server/discover"); + assert!(!discover.headers.contains_key("mcp-session-id")); let initialize = find_jsonrpc_request(&requests, "initialize"); assert!(!initialize.headers.contains_key("mcp-session-id")); @@ -97,13 +150,21 @@ async fn rust_sdk_stateful_transport_propagates_session_protocol_and_delete() { let negotiated_protocol = post_requests .iter() - .find(|request| request.jsonrpc_method.as_deref() != Some("initialize")) + .find(|request| { + !matches!( + request.jsonrpc_method.as_deref(), + Some("initialize" | "server/discover") + ) + }) .and_then(|request| request.headers.get("mcp-protocol-version")) .cloned() .expect("stateful exchange should advertise a negotiated protocol version"); for request in post_requests { - if request.jsonrpc_method.as_deref() == Some("initialize") { + if matches!( + request.jsonrpc_method.as_deref(), + Some("initialize" | "server/discover") + ) { continue; } @@ -155,7 +216,12 @@ async fn rust_sdk_stateless_transport_skips_session_and_delete() { let negotiated_protocol = requests .iter() .filter(|request| request.method == "POST") - .filter(|request| request.jsonrpc_method.as_deref() != Some("initialize")) + .filter(|request| { + !matches!( + request.jsonrpc_method.as_deref(), + Some("initialize" | "server/discover") + ) + }) .find_map(|request| request.headers.get("mcp-protocol-version")) .cloned() .expect("stateless exchange should advertise a negotiated protocol version"); @@ -163,7 +229,12 @@ async fn rust_sdk_stateless_transport_skips_session_and_delete() { for request in requests .iter() .filter(|request| request.method == "POST") - .filter(|request| request.jsonrpc_method.as_deref() != Some("initialize")) + .filter(|request| { + !matches!( + request.jsonrpc_method.as_deref(), + Some("initialize" | "server/discover") + ) + }) { assert_eq!( request @@ -293,6 +364,12 @@ async fn resume_post( let request = serde_json::from_str::(&body).unwrap(); match request.get("method").and_then(Value::as_str) { + Some("server/discover") => response_with_headers( + StatusCode::BAD_REQUEST, + &[("content-type", "application/json")], + json!({"jsonrpc":"2.0", "id":request["id"], "error":{"code":-32601,"message":"legacy fixture"}}) + .to_string(), + ), Some("initialize") => response_with_headers( StatusCode::OK, &[ From b3e46cf1587f6d8cb1c2f67d3e4a773dab46ffce Mon Sep 17 00:00:00 2001 From: daniel Date: Thu, 15 Oct 2026 13:02:07 +0100 Subject: [PATCH 3/3] docs(mcp): explain subscription ownership in harness design --- book/src/ch17-mcp.md | 50 ++++++++++++++++---------------------------- docs/mcp.md | 4 ---- 2 files changed, 18 insertions(+), 36 deletions(-) diff --git a/book/src/ch17-mcp.md b/book/src/ch17-mcp.md index b377280..6a3cf55 100644 --- a/book/src/ch17-mcp.md +++ b/book/src/ch17-mcp.md @@ -34,37 +34,23 @@ The same logic applies to transports: stdio and Streamable HTTP are rmcp impleme ## Modern subscriptions and catalog refresh -Built-in transports negotiate with RMCP's `Auto` lifecycle (`server/discover` -with legacy initialization fallback). Check `protocol_version()` and -`capabilities()` on the connection instead of inferring support from HTTP or -stdio. Adopting an existing service with -`McpConnection::from_running_service_with_events` does not perform asynchronous -setup: call `discover().await` to establish its modern notification stream. - -Modern discovery waits for a `subscriptions/listen` acknowledgment. -`subscription_status().await` reports the requested and accepted filters, and -`subscription_changes().await` watches their lifecycle. Partial catalog -acknowledgments are visible rather than treated as full acceptance. Resource -subscribe/unsubscribe operations replace the complete filter; a rejected URI -set leaves the previous active stream in place. Legacy peers keep the separate -`resources/subscribe` and `resources/unsubscribe` methods. - -Stream notifications use the existing `subscribe_events()` broadcast. For catalog -reconciliation, use `catalog_changes()` and `catalog_revision()`; watch updates -coalesce and survive slow broadcast consumers. The host must schedule -`refresh_changed_catalogs()`, or prepare/fetch/publish refresh work outside its -manager lock. Invalidated, failed, or canceled work must not publish a stale -snapshot. - -A transport failure triggers bounded subscription recovery and dirties the -catalog. Graceful completion yields `Closed { graceful: true, .. }` without -retry. Authentication reconnect rebuilds subscriptions on the new peer and -retains committed resource URIs; `connection_epoch()` fences old work. Await -`close()` to shut down streams before transport cleanup. - -**Migration (agentkit-mcp 0.11.0):** `capabilities()` returns an owned snapshot -instead of a borrowed reference. Read it again after authentication reconnect -for the current peer capabilities; do not retain a lock guard across awaits. +Discovery gives a harness a snapshot of a server's tools, resources, and prompts. +A subscription does not replace that snapshot: its notifications are +invalidations, telling the harness that its view may be stale. + +The host owns reconciliation. It decides when to fetch and publish a fresh +snapshot, coalescing repeated invalidations instead of refreshing once per event. +After a notification stream is restored, reconcile again to cover changes missed +during the gap. If a refresh fails, keep the last usable snapshot and leave the +refresh pending. + +Bind refresh results to the connection that produced them. Authentication changes +or server replacement can make in-flight work obsolete; a result from an old +connection must never overwrite the catalog for its replacement. + +Subscription lifecycle belongs to the connection, not to an individual agent +turn. Close subscriptions with their connection and rebuild them when it is +replaced, so background streams cannot outlive the server they describe. ## MCP in the capability model @@ -197,7 +183,7 @@ The manager owns server lifecycle: | `connect_server(id)` / `connect_all()` | Open a connection and run discovery | | `connect_all_settled()` | Attempt all registered servers concurrently, install successes, and return per-server failures | | `refresh_server(id)` | Re-run discovery and emit per-tool/resource/prompt diff events | -| `refresh_changed_catalogs()` | Reconcile catalogs marked dirty by coalesced invalidations | +| `refresh_changed_catalogs()` | Reconcile catalogs marked dirty by coalesced invalidations | | `disconnect_server(id)` | Close the connection, drop tools from the federated catalog, emit `ServerDisconnected` | | `subscribe_catalog_events()` | Broadcast receiver for `McpCatalogEvent` (server connect / disconnect / tool added / removed / changed / refresh failed / auth changed) | diff --git a/docs/mcp.md b/docs/mcp.md index 182a75e..c36e0f2 100644 --- a/docs/mcp.md +++ b/docs/mcp.md @@ -243,10 +243,6 @@ URIs. `connection_epoch()` changes when the transport is replaced or closed. Call `close().await` to stop owned streams before transport cleanup; dropping the connection also cancels its workers. -**Migration (agentkit-mcp 0.11.0):** `capabilities()` now returns an owned snapshot, -not a borrowed reference. Call it again after authentication reconnect to obtain -the current peer capabilities; no capability lock guard needs to span an await. - ### 8. Errors `McpError` covers transport failure, protocol errors, auth challenges (`AuthRequired(Box)`), invocation errors, and unknown-server lookup misses. It implements `From for rmcp::model::ErrorData` so responders can return agentkit errors and have them surface to the server as JSON-RPC errors.