From 65989430d3873da4823dfa99a8ff9db073df2d37 Mon Sep 17 00:00:00 2001 From: centdix <40307056+centdix@users.noreply.github.com> Date: Wed, 7 Jan 2026 17:02:51 +0100 Subject: [PATCH] upgrade rmcp crate (#7518) --- backend/Cargo.lock | 54 ++++++++++++++++--- backend/windmill-api/Cargo.toml | 2 +- backend/windmill-api/src/lib.rs | 19 +++---- backend/windmill-api/src/mcp/mod.rs | 2 +- backend/windmill-api/src/mcp/server.rs | 41 ++++---------- .../src/mcp/tools/endpoint_tools.rs | 1 + backend/windmill-common/Cargo.toml | 2 +- backend/windmill-common/src/mcp_client.rs | 2 +- backend/windmill-worker/Cargo.toml | 1 - backend/windmill-worker/src/ai/utils.rs | 4 +- 10 files changed, 73 insertions(+), 55 deletions(-) diff --git a/backend/Cargo.lock b/backend/Cargo.lock index fdc816ccfa..048b9d5e27 100644 --- a/backend/Cargo.lock +++ b/backend/Cargo.lock @@ -2656,6 +2656,16 @@ dependencies = [ "darling_macro 0.21.3", ] +[[package]] +name = "darling" +version = "0.23.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "25ae13da2f202d56bd7f91c25fba009e7717a1e4a1cc98a76d844b65ae912e9d" +dependencies = [ + "darling_core 0.23.0", + "darling_macro 0.23.0", +] + [[package]] name = "darling_core" version = "0.13.4" @@ -2712,6 +2722,19 @@ dependencies = [ "syn 2.0.111", ] +[[package]] +name = "darling_core" +version = "0.23.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9865a50f7c335f53564bb694ef660825eb8610e0a53d3e11bf1b0d3df31e03b0" +dependencies = [ + "ident_case", + "proc-macro2", + "quote", + "strsim 0.11.1", + "syn 2.0.111", +] + [[package]] name = "darling_macro" version = "0.13.4" @@ -2756,6 +2779,17 @@ dependencies = [ "syn 2.0.111", ] +[[package]] +name = "darling_macro" +version = "0.23.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ac3984ec7bd6cfa798e62b4a642426a5be0e68f9401cfc2a01e3fa9ea2fcdb8d" +dependencies = [ + "darling_core 0.23.0", + "quote", + "syn 2.0.111", +] + [[package]] name = "dashmap" version = "5.5.3" @@ -9504,6 +9538,12 @@ version = "1.0.15" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "57c0d7b74b563b49d38dae00a0c37d4d6de9b432382b2892f0574ddcae73fd0a" +[[package]] +name = "pastey" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b867cad97c0791bbd3aaa6472142568c6c9e8f71937e98379f584cfb0cf35bec" + [[package]] name = "path-clean" version = "0.1.0" @@ -10902,10 +10942,11 @@ dependencies = [ [[package]] name = "rmcp" -version = "0.8.5" +version = "0.12.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e5947688160b56fb6c827e3c20a72c90392a1d7e9dec74749197aa1780ac42ca" +checksum = "528d42f8176e6e5e71ea69182b17d1d0a19a6b3b894b564678b74cd7cab13cfa" dependencies = [ + "async-trait", "base64 0.22.1", "bytes", "chrono", @@ -10913,7 +10954,7 @@ dependencies = [ "http 1.4.0", "http-body 1.0.1", "http-body-util", - "paste", + "pastey", "pin-project-lite", "rand 0.9.0", "reqwest 0.12.24", @@ -10933,11 +10974,11 @@ dependencies = [ [[package]] name = "rmcp-macros" -version = "0.8.5" +version = "0.12.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "01263441d3f8635c628e33856c468b96ebbce1af2d3699ea712ca71432d4ee7a" +checksum = "e3f81daaa494eb8e985c9462f7d6ce1ab05e5299f48aafd76cdd3d8b060e6f59" dependencies = [ - "darling 0.21.3", + "darling 0.23.0", "proc-macro2", "quote", "serde_json", @@ -15867,7 +15908,6 @@ dependencies = [ "regex", "reqwest 0.12.24", "reqwest-middleware", - "rmcp", "rust_decimal", "serde", "serde_json", diff --git a/backend/windmill-api/Cargo.toml b/backend/windmill-api/Cargo.toml index 78ff288a1c..3e352d8b66 100644 --- a/backend/windmill-api/Cargo.toml +++ b/backend/windmill-api/Cargo.toml @@ -40,7 +40,7 @@ mcp = ["dep:rmcp"] python = [] [dependencies] -rmcp = { version = "0.8.1", features=["transport-streamable-http-server", "transport-streamable-http-server-session", "transport-worker"], optional = true } +rmcp = { version = "0.12.0", features=["transport-streamable-http-server", "transport-streamable-http-server-session", "transport-worker"], optional = true } windmill-queue.workspace = true windmill-common = { workspace = true, default-features = false } windmill-audit.workspace = true diff --git a/backend/windmill-api/src/lib.rs b/backend/windmill-api/src/lib.rs index 1e3b46f223..eeee4af1cc 100644 --- a/backend/windmill-api/src/lib.rs +++ b/backend/windmill-api/src/lib.rs @@ -19,10 +19,8 @@ use crate::oauth2_oss::SlackVerifier; use crate::smtp_server_oss::SmtpServer; #[cfg(feature = "mcp")] -use crate::mcp::{extract_and_store_workspace_id, setup_mcp_server, shutdown_mcp_server}; +use crate::mcp::{extract_and_store_workspace_id, setup_mcp_server}; use crate::triggers::start_all_listeners; -#[cfg(feature = "mcp")] -use rmcp::transport::streamable_http_server::session::local::LocalSessionManager; use tower_http::catch_panic::CatchPanicLayer; use crate::tracing_init::MyOnFailure; @@ -400,14 +398,17 @@ pub async fn run_server( // Setup MCP server #[allow(unused_variables)] - let (mcp_router, mcp_session_manager) = { + let (mcp_router, mcp_cancellation_token) = { #[cfg(feature = "mcp")] if server_mode || mcp_mode { - let (mcp_router, mcp_session_manager) = setup_mcp_server().await?; + let (mcp_router, mcp_cancellation_token) = setup_mcp_server().await?; let mcp_middleware = axum::middleware::from_fn(extract_and_store_workspace_id); - (mcp_router.layer(mcp_middleware), Some(mcp_session_manager)) + ( + mcp_router.layer(mcp_middleware), + Some(mcp_cancellation_token), + ) } else { - (Router::new(), Option::>::None) + (Router::new(), None) } #[cfg(not(feature = "mcp"))] @@ -757,8 +758,8 @@ pub async fn run_server( tracing::info!("Graceful shutdown of server"); #[cfg(feature = "mcp")] - if let Some(mcp_session_manager) = mcp_session_manager { - shutdown_mcp_server(mcp_session_manager).await; + if let Some(mcp_cancellation_token) = mcp_cancellation_token { + mcp_cancellation_token.cancel(); tracing::info!("MCP server shutdown"); } }); diff --git a/backend/windmill-api/src/mcp/mod.rs b/backend/windmill-api/src/mcp/mod.rs index e7223aac8d..e1cd855547 100644 --- a/backend/windmill-api/src/mcp/mod.rs +++ b/backend/windmill-api/src/mcp/mod.rs @@ -8,4 +8,4 @@ pub mod tools; pub mod utils; // Re-export main components -pub use server::{extract_and_store_workspace_id, setup_mcp_server, shutdown_mcp_server, list_tools_service}; \ No newline at end of file +pub use server::{extract_and_store_workspace_id, list_tools_service, setup_mcp_server}; diff --git a/backend/windmill-api/src/mcp/server.rs b/backend/windmill-api/src/mcp/server.rs index 0e1a585383..4270729f6e 100644 --- a/backend/windmill-api/src/mcp/server.rs +++ b/backend/windmill-api/src/mcp/server.rs @@ -18,6 +18,7 @@ use rmcp::{ }; use serde_json::Value; use tokio::try_join; +use tokio_util::sync::CancellationToken; use windmill_common::db::UserDB; use windmill_common::worker::to_raw_value; use windmill_common::{utils::StripPath, DB}; @@ -47,7 +48,7 @@ use axum::{ extract::Path, http::Request, middleware::Next, response::Response, routing::get, Json, Router, }; use rmcp::transport::streamable_http_server::{ - session::local::LocalSessionManager, SessionManager, StreamableHttpService, + session::local::LocalSessionManager, StreamableHttpService, }; use windmill_common::error::JsonResult; @@ -126,6 +127,7 @@ impl Runner { idempotent_hint: Some(false), // Are not guaranteed to be idempotent open_world_hint: Some(true), // Can interact with external services }), + meta: None, }) } } @@ -471,7 +473,7 @@ impl ServerHandler for Runner { ); } - Ok(ListToolsResult { tools, next_cursor: None }) + Ok(ListToolsResult { tools, next_cursor: None, meta: None }) } fn get_info(&self) -> ServerInfo { @@ -498,7 +500,7 @@ impl ServerHandler for Runner { _request: Option, _context: RequestContext, ) -> Result { - Ok(ListResourcesResult { resources: vec![], next_cursor: None }) + Ok(ListResourcesResult { resources: vec![], next_cursor: None, meta: None }) } async fn list_prompts( @@ -530,11 +532,13 @@ pub async fn extract_and_store_workspace_id( } /// Setup the MCP server with HTTP transport -pub async fn setup_mcp_server() -> anyhow::Result<(Router, Arc)> { +pub async fn setup_mcp_server() -> anyhow::Result<(Router, CancellationToken)> { + let cancellation_token = CancellationToken::new(); let session_manager = Arc::new(LocalSessionManager::default()); let service_config = StreamableHttpServerConfig { sse_keep_alive: Some(Duration::from_secs(15)), stateful_mode: false, + cancellation_token: cancellation_token.clone(), }; let service = StreamableHttpService::new( || Ok(Runner::new()), @@ -543,34 +547,7 @@ pub async fn setup_mcp_server() -> anyhow::Result<(Router, Arc) { - let session_ids_to_close = { - let sessions_map = session_manager.sessions.read().await; - sessions_map.keys().cloned().collect::>() - }; - - if !session_ids_to_close.is_empty() { - tracing::info!( - "Closing {} active MCP session(s)...", - session_ids_to_close.len() - ); - let close_futures = session_ids_to_close - .iter() - .map(|session_id| { - let manager_clone = session_manager.clone(); - async move { - if let Err(_) = manager_clone.close_session(session_id).await { - tracing::warn!("Error closing MCP session"); - } - } - }) - .collect::>(); - futures::future::join_all(close_futures).await; - } + Ok((router, cancellation_token)) } /// HTTP handler to list MCP tools as JSON diff --git a/backend/windmill-api/src/mcp/tools/endpoint_tools.rs b/backend/windmill-api/src/mcp/tools/endpoint_tools.rs index 882d9704f0..0ed1f493e4 100644 --- a/backend/windmill-api/src/mcp/tools/endpoint_tools.rs +++ b/backend/windmill-api/src/mcp/tools/endpoint_tools.rs @@ -61,6 +61,7 @@ pub fn endpoint_tool_to_mcp_tool(tool: &EndpointTool) -> Tool { output_schema: None, icons: None, annotations: Some(annotations), + meta: None, } } diff --git a/backend/windmill-common/Cargo.toml b/backend/windmill-common/Cargo.toml index 1e8230e2e3..e9a642e59b 100644 --- a/backend/windmill-common/Cargo.toml +++ b/backend/windmill-common/Cargo.toml @@ -108,7 +108,7 @@ tempfile.workspace = true systemstat.workspace = true size.workspace = true globset.workspace = true -rmcp = { version = "0.8.1", features = ["client", "transport-streamable-http-client", "transport-streamable-http-client-reqwest"] } +rmcp = { version = "0.12.0", features = ["client", "transport-streamable-http-client", "transport-streamable-http-client-reqwest"] } opentelemetry-semantic-conventions = { workspace = true, optional = true } opentelemetry-otlp = { workspace = true, optional = true } diff --git a/backend/windmill-common/src/mcp_client.rs b/backend/windmill-common/src/mcp_client.rs index e12aebe5bd..ee57622902 100644 --- a/backend/windmill-common/src/mcp_client.rs +++ b/backend/windmill-common/src/mcp_client.rs @@ -5,7 +5,7 @@ use reqwest::header::{HeaderMap, HeaderName, HeaderValue}; use serde_json::{json, Value}; use std::str::FromStr; -use rmcp::model::Tool as McpTool; +pub use rmcp::model::Tool as McpTool; use rmcp::{ model::{ CallToolRequestParam, ClientCapabilities, ClientInfo, Implementation, diff --git a/backend/windmill-worker/Cargo.toml b/backend/windmill-worker/Cargo.toml index 19ebb88db5..137abe3d62 100644 --- a/backend/windmill-worker/Cargo.toml +++ b/backend/windmill-worker/Cargo.toml @@ -57,7 +57,6 @@ windmill-parser-sql.workspace = true windmill-parser-graphql.workspace = true windmill-parser-php = { workspace = true, optional = true } windmill-git-sync.workspace = true -rmcp = { version = "0.8.1", features = ["client", "transport-streamable-http-client", "transport-streamable-http-client-reqwest"] } aws-sdk-bedrockruntime.workspace = true aws-config.workspace = true aws-credential-types.workspace = true diff --git a/backend/windmill-worker/src/ai/utils.rs b/backend/windmill-worker/src/ai/utils.rs index c1dae5a73c..17b6f9fce0 100644 --- a/backend/windmill-worker/src/ai/utils.rs +++ b/backend/windmill-worker/src/ai/utils.rs @@ -20,7 +20,7 @@ use windmill_common::{ }; use windmill_common::{ flows::FlowModuleValue, - mcp_client::{McpClient, McpResource, McpToolSource}, + mcp_client::{McpClient, McpResource, McpTool, McpToolSource}, }; use windmill_queue::{flow_status::get_step_of_flow_status, MiniPulledJob}; @@ -352,7 +352,7 @@ pub async fn cleanup_mcp_clients(mcp_clients: HashMap>) { /// Convert raw MCP tools to Windmill Tool format with source tracking fn convert_mcp_tools_to_windmill_tools( - mcp_tools: &[rmcp::model::Tool], + mcp_tools: &[McpTool], resource_name: &str, resource_path: &str, ) -> Result, Error> {