diff --git a/Cargo.lock b/Cargo.lock index a95c9a85a5..f3f595f431 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2038,8 +2038,18 @@ version = "0.23.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "25ae13da2f202d56bd7f91c25fba009e7717a1e4a1cc98a76d844b65ae912e9d" dependencies = [ - "darling_core", - "darling_macro", + "darling_core 0.23.0", + "darling_macro 0.23.0", +] + +[[package]] +name = "darling" +version = "0.24.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ed17f5901b6630b993ca003def43f2f8ef4014fc13b047b57aad617ff32bc2ec" +dependencies = [ + "darling_core 0.24.1", + "darling_macro 0.24.1", ] [[package]] @@ -2055,17 +2065,41 @@ dependencies = [ "syn 2.0.119", ] +[[package]] +name = "darling_core" +version = "0.24.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6837e2cf7485aaae18f86181d2f0e9a7ed297a025e220aeabf63fdebd3a2ddff" +dependencies = [ + "ident_case", + "proc-macro2", + "quote", + "strsim", + "syn 3.0.3", +] + [[package]] name = "darling_macro" version = "0.23.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ac3984ec7bd6cfa798e62b4a642426a5be0e68f9401cfc2a01e3fa9ea2fcdb8d" dependencies = [ - "darling_core", + "darling_core 0.23.0", "quote", "syn 2.0.119", ] +[[package]] +name = "darling_macro" +version = "0.24.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2ac7135c3ef02b2f7833bbeb1be5ba7f966dcde8a87c6b87f65a778d71a02785" +dependencies = [ + "darling_core 0.24.1", + "quote", + "syn 3.0.3", +] + [[package]] name = "data-encoding" version = "2.11.0" @@ -5854,7 +5888,7 @@ name = "mysql-common-derive" version = "0.32.2" source = "git+https://github.com/t8y2/rust_mysql_common.git?rev=47740d85a8277167a5717afac785d37b46a36296#47740d85a8277167a5717afac785d37b46a36296" dependencies = [ - "darling", + "darling 0.23.0", "heck 0.5.0", "manyhow", "num-bigint", @@ -7877,18 +7911,19 @@ dependencies = [ [[package]] name = "rmcp" -version = "2.2.0" +version = "3.5.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "14db48ee17a9ba61810ab1a9c1beb7d06d8136ae39ac25a1137f10d357af01af" +checksum = "ab5023804cbe0304cbe06197ac16e645a0a0d0ffa2953bf735c15c2fa0213350" dependencies = [ "async-trait", - "base64 0.22.1", + "base64 0.23.0", "bytes", "chrono", "futures", "http 1.5.0", "http-body 1.1.0", "http-body-util", + "indexmap 2.14.0", "pastey", "pin-project-lite", "rand 0.10.2", @@ -7909,15 +7944,15 @@ dependencies = [ [[package]] name = "rmcp-macros" -version = "2.2.0" +version = "3.5.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "783d787bf21813b285f13019adc49e11af501c658890c1e519f31f937c68b7e3" +checksum = "a1da498700030856ba63114bceefa30ba21f5e278cb975e5b9de3529fe842b40" dependencies = [ - "darling", + "darling 0.24.1", "proc-macro2", "quote", "serde_json", - "syn 2.0.119", + "syn 3.0.3", ] [[package]] @@ -8776,7 +8811,7 @@ version = "3.21.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "84d57bc0c8b9a17920c178daa6bb924850d54a9c97ab45194bb8c17ad66bb660" dependencies = [ - "darling", + "darling 0.23.0", "proc-macro2", "quote", "syn 2.0.119", diff --git a/crates/dbx-mcp/Cargo.toml b/crates/dbx-mcp/Cargo.toml index 03e9f5cfee..c9ee4575c8 100644 --- a/crates/dbx-mcp/Cargo.toml +++ b/crates/dbx-mcp/Cargo.toml @@ -32,7 +32,7 @@ log = "0.4" mysql_async = { version = "0.37", default-features = false, features = ["default-rustls", "client_ed25519", "chrono", "rust_decimal"] } axum = "0.8" chrono = "0.4" -rmcp = { version = "2.2.0", features = ["client", "transport-io", "transport-streamable-http-server"] } +rmcp = { version = "3.5.1", features = ["client", "transport-io", "transport-streamable-http-server"] } reqwest = { version = "0.12", default-features = false, features = ["json", "rustls-tls", "socks"] } rustls = { version = "0.23", features = ["aws-lc-rs"] } schemars = "1" @@ -49,7 +49,7 @@ url = "2" security-framework = { version = "3.7", optional = true } [dev-dependencies] -rmcp = { version = "2.2.0", features = ["transport-streamable-http-client-reqwest"] } +rmcp = { version = "3.5.1", features = ["transport-streamable-http-client-reqwest"] } rustls-pemfile = "2.2" tempfile = "3" tokio = { version = "1", features = ["net", "io-util"] } diff --git a/crates/dbx-mcp/src/http.rs b/crates/dbx-mcp/src/http.rs index 8b71f5c0f2..d75884bb11 100644 --- a/crates/dbx-mcp/src/http.rs +++ b/crates/dbx-mcp/src/http.rs @@ -15,7 +15,8 @@ use crate::{ diagnostics::health, http_auth::{authorize_request, HttpAuth}, runtime::HttpRuntimeConfig, - DbxBackend, DbxMcpServer, McpScope, + server::{PendingSalesforceWrites, PluginToolsMode, SALESFORCE_WRITE_CONFIRM_TTL}, + DbxBackend, DbxMcpServer, McpScope, McpSessionStore, }; /// Builds a protected Streamable HTTP MCP router for embedding in an existing @@ -28,7 +29,17 @@ pub fn streamable_http_router( allowed_hosts: Vec, web_mode: bool, ) -> Result { - build_streamable_http_router(backend, path, auth, allowed_hosts, web_mode, None, Default::default()) + build_streamable_http_router( + backend, + path, + auth, + allowed_hosts, + web_mode, + None, + Default::default(), + McpSessionStore::new(), + PendingSalesforceWrites::new(SALESFORCE_WRITE_CONFIRM_TTL), + ) } fn build_streamable_http_router( @@ -39,6 +50,8 @@ fn build_streamable_http_router( web_mode: bool, cancellation: Option, session_manager: Arc, + sessions: Arc, + pending_salesforce_writes: Arc, ) -> Result { auth.set_allowed_hosts(allowed_hosts.clone())?; // Web settings update the shared policy without rebuilding the router. @@ -53,8 +66,23 @@ fn build_streamable_http_router( } let server_backend = backend.clone(); let scope = McpScope::from_env(); + let plugin_tools_mode = PluginToolsMode::from_env(); + // Legacy (`< 2026-07-28`) conversations are driven through one long-lived + // service, but stateless `2026-07-28` requests, tool-schema discovery, and + // session restoration each call the factory again. Session and + // pending-write state belongs to the endpoint, not to one instance, so a + // stateless request can still find the session an earlier request opened. let service: StreamableHttpService = StreamableHttpService::new( - move || Ok(DbxMcpServer::with_runtime_options(server_backend.clone(), scope.clone(), web_mode)), + move || { + Ok(DbxMcpServer::with_shared_state( + server_backend.clone(), + scope.clone(), + web_mode, + plugin_tools_mode, + sessions.clone(), + pending_salesforce_writes.clone(), + )) + }, session_manager, rmcp_config, ); @@ -110,14 +138,18 @@ pub async fn serve_streamable_http_on_listener( listener: tokio::net::TcpListener, ) -> io::Result<()> { let session_manager = Arc::new(LocalSessionManager::default()); + let sessions = McpSessionStore::new(); + let pending_salesforce_writes = PendingSalesforceWrites::new(SALESFORCE_WRITE_CONFIRM_TTL); let mcp_router = build_streamable_http_router( - backend, + backend.clone(), &config.path, config.auth, config.allowed_hosts, false, Some(cancellation.child_token()), session_manager.clone(), + sessions.clone(), + pending_salesforce_writes, ) .map_err(io::Error::other)?; let router = Router::new().route("/healthz", get(health)).route("/readyz", get(health)).merge(mcp_router); @@ -129,16 +161,29 @@ pub async fn serve_streamable_http_on_listener( cancellation.cancelled().await; }) .await; - close_local_sessions_bounded(&session_manager).await; + close_http_sessions_bounded(&backend, &session_manager, &sessions).await; result } -async fn close_local_sessions_bounded(session_manager: &Arc) { +/// Drain protocol sessions, then the stateful sessions they opened. +/// +/// Sessions live on the endpoint-wide store, not on the dropped per-session +/// service, so they have to be rolled back explicitly here. +async fn close_http_sessions_bounded( + backend: &Arc, + session_manager: &Arc, + sessions: &Arc, +) { let session_ids = session_manager.sessions.read().await.keys().cloned().collect::>(); let cleanup = async { for session_id in session_ids { let _ = session_manager.close_session(&session_id).await; } + let leftover = sessions.take_all_active().await; + if !leftover.is_empty() { + let server = DbxMcpServer::with_runtime_options(backend.clone(), McpScope::from_env(), false); + server.close_backend_sessions_best_effort(leftover).await; + } }; if tokio::time::timeout(std::time::Duration::from_secs(10), cleanup).await.is_err() { log::warn!("Timed out draining MCP HTTP protocol sessions during shutdown"); @@ -155,12 +200,13 @@ mod tests { storage::McpGlobalPolicy, }; use rmcp::{ - model::CallToolRequestParams, + model::{CallToolRequestParams, ProtocolVersion}, service::ServiceExt, transport::{ streamable_http_client::StreamableHttpClientTransportConfig, streamable_http_server::session::local::SessionConfig, StreamableHttpClientTransport, }, + ClientLifecycleMode, ClientServiceExt, }; use serde_json::{json, Value}; use tokio::io::{AsyncReadExt, AsyncWriteExt}; @@ -168,6 +214,7 @@ mod tests { use super::*; use crate::{ backend::DbxBackend, + server::{PendingSalesforceWrites, SALESFORCE_WRITE_CONFIRM_TTL}, transaction::{ TransactionIo, TransactionIoError, TransactionIoSuccess, TransactionOwner, TransactionOwnerConfig, }, @@ -292,6 +339,19 @@ mod tests { async fn add_connection_for_mcp(&self, config: ConnectionConfig) -> Result { Ok(config) } + + /// Model the pool release a real backend performs: the rollback itself + /// comes from the transaction owner the session carries, and the + /// disconnect counter records that the pinned pool was disposed. + async fn close_client_session( + &self, + _connection_id: &str, + _database: &str, + _client_session_id: &str, + ) -> Result { + self.disconnects.fetch_add(1, Ordering::SeqCst); + Ok(true) + } async fn duplicate_connection_for_mcp( &self, _source_id: &str, @@ -308,7 +368,7 @@ mod tests { async fn start_http_test_server( backend: Arc, keep_alive: std::time::Duration, - ) -> (String, Arc, CancellationToken, tokio::task::JoinHandle<()>) { + ) -> (String, Arc, Arc, CancellationToken, tokio::task::JoinHandle<()>) { // A single default provider avoids the "No rustls crypto provider is // configured" panic when tests build reqwest clients in workspace // builds where multiple rustls crypto features are present; the @@ -322,26 +382,31 @@ mod tests { local_manager.session_config = session_config; let manager = Arc::new(local_manager); let cancellation = CancellationToken::new(); + let sessions = McpSessionStore::new(); let router = build_streamable_http_router( - backend, + backend.clone(), "/mcp", HttpAuth::new("http-test-token".to_string(), Vec::::new(), true).unwrap(), vec![address.to_string()], false, Some(cancellation.child_token()), manager.clone(), + sessions.clone(), + PendingSalesforceWrites::new(SALESFORCE_WRITE_CONFIRM_TTL), ) .unwrap(); let shutdown = cancellation.clone(); let shutdown_manager = manager.clone(); + let shutdown_sessions = sessions.clone(); + let shutdown_backend: Arc = backend; let task = tokio::spawn(async move { axum::serve(listener, router) .with_graceful_shutdown(async move { shutdown.cancelled().await }) .await .unwrap(); - close_local_sessions_bounded(&shutdown_manager).await; + close_http_sessions_bounded(&shutdown_backend, &shutdown_manager, &shutdown_sessions).await; }); - (format!("http://{address}/mcp"), manager, cancellation, task) + (format!("http://{address}/mcp"), manager, sessions, cancellation, task) } async fn open_active_transaction(url: &str) -> (rmcp::service::RunningService, String) { @@ -374,10 +439,14 @@ mod tests { (client, session_id) } + /// Wait until the backend connection the session pinned was rolled back and + /// disposed. Sessions are reclaimed by an explicit `dbx_close_session`, the + /// idle TTL, or server shutdown, so every caller triggers one of those + /// first; ending the HTTP transport session deliberately does not. async fn wait_for_disposal(backend: &HttpTestBackend) { tokio::time::timeout(std::time::Duration::from_secs(1), async { loop { - if backend.disconnects.load(Ordering::SeqCst) == 1 + if backend.disconnects.load(Ordering::SeqCst) >= 1 && backend.sql.lock().unwrap().iter().any(|sql| sql == "ROLLBACK") { break; @@ -386,7 +455,19 @@ mod tests { } }) .await - .expect("HTTP session cleanup must roll back and disconnect the owner"); + .expect("session cleanup must roll back and disconnect the owner"); + } + + /// Close the stateful session the client opened by its own id. + async fn close_inner_session(client: &rmcp::service::RunningService, session_id: &str) { + let closed = client + .call_tool( + CallToolRequestParams::new("dbx_close_session") + .with_arguments(serde_json::from_value(json!({ "session_id": session_id })).unwrap()), + ) + .await + .unwrap(); + assert_ne!(closed.is_error, Some(true), "close session failed: {closed:?}"); } async fn open_and_drop_raw_sse(url: &str, outer_session_id: &str) { @@ -425,12 +506,16 @@ mod tests { drop(stream); } + /// Ending the HTTP transport session must still be possible without + /// touching the stateful session the agent opened on it: HTTP DELETE is a + /// transport-level operation, and a stateless `2026-07-28` agent has no + /// transport session at all. #[tokio::test] - async fn authenticated_http_delete_rolls_back_and_disconnects_inner_owner() { + async fn authenticated_http_delete_keeps_the_stateful_session_until_it_is_closed() { let backend = Arc::new(HttpTestBackend::new()); - let (url, manager, cancellation, server_task) = + let (url, manager, _sessions, cancellation, server_task) = start_http_test_server(backend.clone(), std::time::Duration::from_secs(30)).await; - let (client, _) = open_active_transaction(&url).await; + let (client, inner_session_id) = open_active_transaction(&url).await; let outer_session_id = manager.sessions.read().await.keys().next().unwrap().to_string(); let response = reqwest::Client::new() @@ -441,6 +526,9 @@ mod tests { .await .unwrap(); assert!(response.status().is_success(), "DELETE returned {}", response.status()); + assert_eq!(backend.disconnects.load(Ordering::SeqCst), 0); + + close_inner_session(&client, &inner_session_id).await; wait_for_disposal(&backend).await; drop(client); @@ -448,21 +536,26 @@ mod tests { server_task.await.unwrap(); } + /// The transport idle timeout reclaims the protocol session, not the + /// stateful session an agent opened through it. #[tokio::test] - async fn http_inactivity_expiry_rolls_back_and_disconnects_inner_owner() { + async fn http_inactivity_expiry_reclaims_the_transport_session() { let backend = Arc::new(HttpTestBackend::new()); - let (url, manager, cancellation, server_task) = + let (url, manager, _sessions, cancellation, server_task) = start_http_test_server(backend.clone(), std::time::Duration::from_millis(500)).await; - let (client, _) = open_active_transaction(&url).await; + let (client, inner_session_id) = open_active_transaction(&url).await; - wait_for_disposal(&backend).await; - tokio::time::timeout(std::time::Duration::from_secs(1), async { + tokio::time::timeout(std::time::Duration::from_secs(2), async { while !manager.sessions.read().await.is_empty() { tokio::task::yield_now().await; } }) .await .expect("expired outer HTTP session must be removed"); + assert_eq!(backend.disconnects.load(Ordering::SeqCst), 0); + + close_inner_session(&client, &inner_session_id).await; + wait_for_disposal(&backend).await; drop(client); cancellation.cancel(); @@ -472,7 +565,7 @@ mod tests { #[tokio::test] async fn transient_http_connections_preserve_one_outer_session() { let backend = Arc::new(HttpTestBackend::new()); - let (url, manager, cancellation, server_task) = + let (url, manager, _sessions, cancellation, server_task) = start_http_test_server(backend.clone(), std::time::Duration::from_secs(30)).await; let (client, inner_session_id) = open_active_transaction(&url).await; let outer_session_id = manager.sessions.read().await.keys().next().unwrap().to_string(); @@ -507,16 +600,133 @@ mod tests { .await .unwrap(); assert!(response.status().is_success()); + assert_eq!(backend.disconnects.load(Ordering::SeqCst), 0); + + close_inner_session(&client, &inner_session_id).await; wait_for_disposal(&backend).await; drop(client); cancellation.cancel(); server_task.await.unwrap(); } + /// The same endpoint must serve both protocol generations: a legacy agent + /// gets the stateful `initialize` handshake and a session id, while a modern + /// agent gets stateless discovery on `2026-07-28` with no session at all. + #[tokio::test] + async fn http_serves_stateful_legacy_and_stateless_2026_agents() { + let backend = Arc::new(HttpTestBackend::new()); + let (url, manager, _sessions, cancellation, server_task) = + start_http_test_server(backend.clone(), std::time::Duration::from_secs(30)).await; + + let legacy = () + .serve(StreamableHttpClientTransport::from_config( + StreamableHttpClientTransportConfig::with_uri(url.clone()).auth_header("http-test-token"), + )) + .await + .expect("legacy initialize client"); + let legacy_info = legacy.peer_info().expect("legacy initialize info"); + assert_eq!(legacy_info.protocol_version, ProtocolVersion::V_2025_11_25); + let listed = legacy.list_all_tools().await.expect("legacy tools/list"); + assert!(!listed.is_empty()); + assert_eq!(manager.sessions.read().await.len(), 1, "legacy agents keep a stateful session"); + + // `Discover` mode never sends `initialize`, so this only succeeds if the + // server answers `server/discover` for the modern lifecycle. + let modern = ClientServiceExt::serve_with_lifecycle( + (), + StreamableHttpClientTransport::from_config( + StreamableHttpClientTransportConfig::with_uri(url.clone()).auth_header("http-test-token"), + ), + ClientLifecycleMode::Discover { preferred_versions: vec![ProtocolVersion::V_2026_07_28] }, + ) + .await + .expect("modern discover client"); + let modern_info = modern.peer_info().expect("discover info"); + assert_eq!(modern_info.protocol_version, ProtocolVersion::V_2026_07_28); + let modern_tools = modern.list_all_tools().await.expect("modern tools/list"); + assert!(!modern_tools.is_empty()); + assert_eq!(manager.sessions.read().await.len(), 1, "2026-07-28 discovery stays stateless"); + + let _ = legacy.cancel().await; + let _ = modern.cancel().await; + cancellation.cancel(); + server_task.await.unwrap(); + assert_eq!(manager.sessions.read().await.len(), 0, "both agents released their transport state"); + } + + /// `2026-07-28` has no protocol-level session, so every request arrives at a + /// freshly built `DbxMcpServer`. A session handle minted by one request must + /// therefore still resolve in the next one, otherwise transactions silently + /// break for modern agents while still working for legacy ones. + #[tokio::test] + async fn stateless_2026_requests_share_db_session_state() { + let backend = Arc::new(HttpTestBackend::new()); + let (url, _manager, _sessions, cancellation, server_task) = + start_http_test_server(backend.clone(), std::time::Duration::from_secs(30)).await; + + let modern = ClientServiceExt::serve_with_lifecycle( + (), + StreamableHttpClientTransport::from_config( + StreamableHttpClientTransportConfig::with_uri(url).auth_header("http-test-token"), + ), + ClientLifecycleMode::Discover { preferred_versions: vec![ProtocolVersion::V_2026_07_28] }, + ) + .await + .expect("modern discover client"); + + let opened = modern + .call_tool( + CallToolRequestParams::new("dbx_open_session").with_arguments( + serde_json::from_value(json!({ + "connection_id": "mysql", + "database": "app", + "enable_transactions": true + })) + .unwrap(), + ), + ) + .await + .unwrap(); + assert_ne!(opened.is_error, Some(true), "open session failed: {opened:?}"); + let session_id = opened.structured_content.as_ref().unwrap()["session_id"].as_str().unwrap().to_string(); + + let begun = modern + .call_tool( + CallToolRequestParams::new("dbx_begin_transaction") + .with_arguments(serde_json::from_value(json!({ "session_id": session_id.clone() })).unwrap()), + ) + .await + .unwrap(); + assert_ne!(begun.is_error, Some(true), "begin transaction failed: {begun:?}"); + + // A separate HTTP request (and therefore a separate server instance) + // must find the session opened two requests ago. + let query = modern + .call_tool( + CallToolRequestParams::new("dbx_execute_query").with_arguments( + serde_json::from_value(json!({ + "connection_id": "mysql", + "database": "app", + "session_id": session_id, + "sql": "SELECT 1" + })) + .unwrap(), + ), + ) + .await + .unwrap(); + assert_ne!(query.is_error, Some(true), "query on shared session failed: {query:?}"); + assert_eq!(query.structured_content.as_ref().unwrap()["transaction_state"], "active"); + + let _ = modern.cancel().await; + cancellation.cancel(); + server_task.await.unwrap(); + } + #[tokio::test] async fn http_service_shutdown_rolls_back_and_disconnects_inner_owner() { let backend = Arc::new(HttpTestBackend::new()); - let (url, _manager, cancellation, server_task) = + let (url, _manager, _sessions, cancellation, server_task) = start_http_test_server(backend.clone(), std::time::Duration::from_secs(30)).await; let (client, _) = open_active_transaction(&url).await; diff --git a/crates/dbx-mcp/src/server.rs b/crates/dbx-mcp/src/server.rs index 786d70dc30..e789e4e16b 100644 --- a/crates/dbx-mcp/src/server.rs +++ b/crates/dbx-mcp/src/server.rs @@ -3,9 +3,9 @@ use std::{collections::HashMap, sync::Arc, time::Duration, time::Instant}; use rmcp::{ handler::server::{router::tool::ToolRouter, wrapper::Parameters}, model::{ - CallToolResult, ContentBlock, ErrorData, Implementation, ListResourceTemplatesResult, ListResourcesResult, - ReadResourceRequestParams, ReadResourceResult, Resource, ResourceContents, ResourceTemplate, - ServerCapabilities, ServerInfo, + CallToolResponse, CallToolResult, ContentBlock, ErrorData, Implementation, ListResourceTemplatesResult, + ListResourcesResult, ReadResourceRequestParams, ReadResourceResponse, ReadResourceResult, Resource, + ResourceContents, ResourceTemplate, ServerCapabilities, ServerConfig, }, schemars, tool, tool_handler, tool_router, ServerHandler, }; @@ -463,7 +463,7 @@ pub struct SalesforceApplyWriteRequest { /// agent to show the summary and for a human to answer, short enough that a /// token left in a transcript cannot be replayed much later against a changed /// org. -const SALESFORCE_WRITE_CONFIRM_TTL: Duration = Duration::from_secs(300); +pub(crate) const SALESFORCE_WRITE_CONFIRM_TTL: Duration = Duration::from_secs(300); /// Ceiling on simultaneously pending confirmations. Prepared writes live in /// process memory, so an agent that prepares in a loop must not be able to grow /// the map without bound. @@ -493,13 +493,13 @@ struct PendingSalesforceWrite { } #[derive(Debug)] -struct PendingSalesforceWrites { +pub(crate) struct PendingSalesforceWrites { ttl: Duration, entries: tokio::sync::Mutex>, } impl PendingSalesforceWrites { - fn new(ttl: Duration) -> Arc { + pub(crate) fn new(ttl: Duration) -> Arc { Arc::new(Self { ttl, entries: tokio::sync::Mutex::new(HashMap::new()) }) } @@ -802,6 +802,46 @@ impl DbxMcpServer { scope: McpScope, web_mode: bool, plugin_tools_mode: PluginToolsMode, + ) -> Self { + Self::build( + backend, + scope, + web_mode, + plugin_tools_mode, + McpSessionStore::new(), + PendingSalesforceWrites::new(SALESFORCE_WRITE_CONFIRM_TTL), + ) + } + + /// Same as [`with_plugin_tools_mode`], but reuses session and pending-write + /// state that outlives this instance. + /// + /// The Streamable HTTP transport drives a `Stateful` (protocol `< 2026-07-28`) + /// conversation through one long-lived service, but it calls the service + /// factory **per request** for every other mode: stateless `2026-07-28` + /// requests, tool-schema discovery, and session restoration. Tools that hand + /// state back to the model — `dbx_open_session` and the transaction tools, + /// `dbx_salesforce_prepare_write` and its confirmation — must therefore keep + /// that state in `Arc`s owned by the transport, not in the per-instance + /// defaults, or the next request would not find its own session. + pub(crate) fn with_shared_state( + backend: Arc, + scope: McpScope, + web_mode: bool, + plugin_tools_mode: PluginToolsMode, + sessions: Arc, + pending_salesforce_writes: Arc, + ) -> Self { + Self::build(backend, scope, web_mode, plugin_tools_mode, sessions, pending_salesforce_writes) + } + + fn build( + backend: Arc, + scope: McpScope, + web_mode: bool, + plugin_tools_mode: PluginToolsMode, + sessions: Arc, + pending_salesforce_writes: Arc, ) -> Self { // The workspace enables more than one rustls crypto feature through // transitive dependencies. Native MCP runs outside the desktop/web @@ -831,14 +871,7 @@ impl DbxMcpServer { tool_router.disable_route("dbx_plugin_tools"); tool_router.disable_route("dbx_plugin_call"); } - Self { - backend, - scope, - plugin_tools_mode, - sessions: McpSessionStore::new(), - pending_salesforce_writes: PendingSalesforceWrites::new(SALESFORCE_WRITE_CONFIRM_TTL), - tool_router, - } + Self { backend, scope, plugin_tools_mode, sessions, pending_salesforce_writes, tool_router } } fn spawn_session_cleanup(&self, session: McpSession) -> tokio::sync::oneshot::Receiver { @@ -879,7 +912,7 @@ impl DbxMcpServer { result_rx } - async fn close_backend_sessions_best_effort(&self, sessions: Vec) { + pub(crate) async fn close_backend_sessions_best_effort(&self, sessions: Vec) { let cleanups = sessions.into_iter().map(|session| self.spawn_session_cleanup(session)).collect::>(); for cleanup in cleanups { let _ = cleanup.await; @@ -3646,7 +3679,7 @@ impl ServerHandler for DbxMcpServer { &self, request: rmcp::model::CallToolRequestParams, context: rmcp::service::RequestContext, - ) -> Result { + ) -> Result { let started = Instant::now(); let args = serde_json::Value::Object(request.arguments.clone().unwrap_or_default()); let arg = |name: &str| args.get(name).and_then(|value| value.as_str()).unwrap_or_default().to_string(); @@ -3677,26 +3710,33 @@ impl ServerHandler for DbxMcpServer { let is_plugin_tool = crate::plugin_tools::is_plugin_tool_name(&tool_name); CALL_HISTORY.scope(std::sync::Mutex::new(entry), async { let cancellation = context.ct.clone(); - let result = if is_plugin_tool { + let result: Result = if is_plugin_tool { tokio::select! { - result = self.call_plugin_tool_dispatch(&tool_name, args) => result, - _ = cancellation.cancelled() => Ok(tool_error("REQUEST_CANCELLED", "The MCP request was cancelled.")), + result = self.call_plugin_tool_dispatch(&tool_name, args) => result.map(Into::into), + _ = cancellation.cancelled() => Ok(tool_error("REQUEST_CANCELLED", "The MCP request was cancelled.").into()), } } else { let tcc = rmcp::handler::server::tool::ToolCallContext::new(self, request, context); tokio::select! { result = self.tool_router.call(tcc) => result, - _ = cancellation.cancelled() => Ok(tool_error("REQUEST_CANCELLED", "The MCP request was cancelled.")), + _ = cancellation.cancelled() => Ok(tool_error("REQUEST_CANCELLED", "The MCP request was cancelled.").into()), } }; let mut entry = CALL_HISTORY.with(|entry| entry.lock().unwrap().clone()); entry.execution_time_ms = started.elapsed().as_millis(); - entry.success = result.as_ref().is_ok_and(|result| result.is_error != Some(true)); + entry.success = result.as_ref().is_ok_and(|response| match response { + CallToolResponse::Complete(result) => result.is_error != Some(true), + // MRTR and task results are never terminal tool failures. + _ => true, + }); // Keep the complete response within a bounded, redacted payload so // the history detail view can be used for troubleshooting without // allowing a large result to grow the local database indefinitely. + // Only a completed call has a concrete payload to archive; interim + // MRTR or task results carry a handle the client resolves later. entry.mcp_response_json = Some(match &result { - Ok(result) => bounded_history_response(result), + Ok(CallToolResponse::Complete(result)) => bounded_history_response(result), + Ok(response) => format!("{response:?}"), Err(error) => bounded_history_response(error), }); if !entry.success { @@ -3716,8 +3756,8 @@ impl ServerHandler for DbxMcpServer { }).await } - fn get_info(&self) -> ServerInfo { - ServerInfo::new(ServerCapabilities::builder().enable_tools().enable_resources().build()) + fn get_info(&self) -> ServerConfig { + ServerConfig::new(ServerCapabilities::builder().enable_tools().enable_resources().build()) .with_server_info(Implementation::new("dbx", env!("CARGO_PKG_VERSION"))) .with_instructions("Use DBX connections to inspect schemas and query databases safely.") } @@ -3739,7 +3779,14 @@ impl ServerHandler for DbxMcpServer { }) .into_iter() .collect(); - Ok(ListResourcesResult { resources, meta: None, next_cursor: None }) + Ok(ListResourcesResult { + result_type: None, + meta: None, + next_cursor: None, + ttl_ms: None, + cache_scope: None, + resources, + }) } async fn list_resource_templates( @@ -3774,14 +3821,21 @@ impl ServerHandler for DbxMcpServer { .with_mime_type("text/markdown"), ); } - Ok(ListResourceTemplatesResult { resource_templates, meta: None, next_cursor: None }) + Ok(ListResourceTemplatesResult { + result_type: None, + meta: None, + next_cursor: None, + ttl_ms: None, + cache_scope: None, + resource_templates, + }) } async fn read_resource( &self, request: ReadResourceRequestParams, _context: rmcp::service::RequestContext, - ) -> Result { + ) -> Result { let uri = request.uri; let result = match parse_dbx_resource_uri(&uri)? { DbxResourceRequest::Connections => self.list_connections(Parameters(ListConnectionsRequest {})).await, @@ -3809,7 +3863,7 @@ impl ServerHandler for DbxMcpServer { .await } }; - resource_result_from_tool(uri, result) + resource_result_from_tool(uri, result).map(Into::into) } /// Hide tools the global policy disallows from the advertised list, the @@ -3824,7 +3878,14 @@ impl ServerHandler for DbxMcpServer { _request: Option, _context: rmcp::service::RequestContext, ) -> Result { - Ok(rmcp::model::ListToolsResult { tools: self.policy_filtered_tools().await, meta: None, next_cursor: None }) + Ok(rmcp::model::ListToolsResult { + result_type: None, + meta: None, + next_cursor: None, + ttl_ms: None, + cache_scope: None, + tools: self.policy_filtered_tools().await, + }) } } diff --git a/crates/dbx-mcp/src/session.rs b/crates/dbx-mcp/src/session.rs index 0f90806f3e..06706d7f29 100644 --- a/crates/dbx-mcp/src/session.rs +++ b/crates/dbx-mcp/src/session.rs @@ -80,6 +80,21 @@ impl McpSessionStore { Arc::new(Self { idle_ttl: session_idle_ttl_from_env(), ..Self::default() }) } + /// Remove every active session so the caller can close its backend pool. + /// + /// One [`McpSessionStore`] is shared by every request of a Streamable HTTP + /// endpoint, so the sessions an agent opened outlive the `DbxMcpServer` + /// instance that served the request which created them. Its HTTP transport + /// session ending therefore cannot release them, and a stateless + /// `2026-07-28` agent never has one to end. Sessions are reclaimed by the + /// idle TTL or by an explicit `dbx_close_session`; server shutdown uses this + /// method to roll back whatever is still open. + pub async fn take_all_active(&self) -> Vec { + let mut state = self.state.lock().await; + let closing = state.closing.values().cloned().collect::>(); + closing.into_iter().chain(state.active.drain().map(|(_, session)| session)).collect() + } + /// Open a new session bound to `connection_id` + `database`. /// /// Returns `Err` with a human-readable message when the session cap is diff --git a/crates/dbx-mcp/src/transport.rs b/crates/dbx-mcp/src/transport.rs index 669bc032d8..370fdb5b3b 100644 --- a/crates/dbx-mcp/src/transport.rs +++ b/crates/dbx-mcp/src/transport.rs @@ -1,12 +1,33 @@ +//! Compatibility shim for the `server/discover` probe used by agents that +//! predate the final `2026-07-28` lifecycle. +//! +//! rmcp implements `server/discover` natively, but it rejects a probe that omits +//! the `_meta` fields the spec made mandatory (`protocolVersion` and +//! `clientCapabilities`) with `-32602 Invalid Params`. Its client-side +//! classification treats that code as a *modern* server rejection, so a client +//! that probes discovery before falling back to `initialize` does not fall back +//! and the connection fails. +//! +//! Before the SDK upgrade DBX answered every `server/discover` request with +//! `-32601 Method not found`, which those clients do treat as "this server is +//! legacy, use `initialize`". This shim keeps that negotiated behaviour for a +//! bare probe while leaving a well-formed probe to the SDK, so modern agents +//! still discover `2026-07-28` and older ones keep their established fallback. + use std::future::Future; use rmcp::{ - model::{ClientJsonRpcMessage, ClientRequest, ErrorCode, ErrorData, ServerJsonRpcMessage}, + model::{ + ClientJsonRpcMessage, ClientRequest, ErrorCode, ErrorData, GetExtensions, RequestMetaObject, + ServerJsonRpcMessage, + }, service::{RxJsonRpcMessage, TxJsonRpcMessage}, transport::{IntoTransport, Transport}, RoleServer, }; +/// Wrap a server transport so a bare `server/discover` probe is answered with +/// `METHOD_NOT_FOUND` instead of the SDK's `INVALID_PARAMS`. pub fn with_legacy_discovery_fallback(transport: T) -> impl Transport + 'static where T: IntoTransport, @@ -35,23 +56,14 @@ where async fn receive(&mut self) -> Option> { loop { let message = self.inner.receive().await?; - let discovery_request_id = match &message { - ClientJsonRpcMessage::Request(request) - if matches!( - &request.request, - ClientRequest::CustomRequest(custom) if custom.method == "server/discover" - ) => - { - Some(request.id.clone()) - } - _ => None, - }; - let Some(request_id) = discovery_request_id else { + let Some(request_id) = bare_discovery_probe_id(&message) else { return Some(message); }; - // Reject discovery without closing so new clients can fall back to the legacy initialize flow. + // Reject the probe without closing the transport so the client can + // fall back to the legacy `initialize` handshake, exactly as it did + // before the SDK learned about `server/discover`. let error = ErrorData::new(ErrorCode::METHOD_NOT_FOUND, "Method not found", None); if self.inner.send(ServerJsonRpcMessage::error(error, Some(request_id))).await.is_err() { return None; @@ -63,3 +75,59 @@ where self.inner.close() } } + +/// The request id of a `server/discover` probe that carries no request +/// metadata at all. +/// +/// Returns `None` for every other message, including a probe that carries +/// `_meta`, so the SDK answers those. Deciding on the presence of `_meta` rather +/// than on its completeness keeps the SDK in charge of validating a +/// well-intentioned modern request. +fn bare_discovery_probe_id(message: &ClientJsonRpcMessage) -> Option { + let ClientJsonRpcMessage::Request(request) = message else { + return None; + }; + let ClientRequest::DiscoverRequest(discover) = &request.request else { + return None; + }; + // The SDK keeps a request's `_meta` in the request extensions after + // deserializing it, so an absent or empty entry means the client probed + // discovery without any per-request metadata. + let has_meta = discover.extensions().get::().is_some_and(|meta| !meta.0.is_empty()); + (!has_meta).then(|| request.id.clone()) +} + +#[cfg(test)] +mod tests { + use rmcp::model::{ + ClientJsonRpcMessage, ClientRequest, DiscoverRequest, DiscoverRequestParams, GetExtensions, JsonRpcRequest, + NumberOrString, RequestMetaObject, + }; + + use super::bare_discovery_probe_id; + + fn discover_probe(id: i64, meta: Option) -> ClientJsonRpcMessage { + let mut discover = DiscoverRequest::new(DiscoverRequestParams {}); + if let Some(meta) = meta { + discover.extensions_mut().insert(meta); + } + ClientJsonRpcMessage::Request(JsonRpcRequest::new( + NumberOrString::Number(id), + ClientRequest::DiscoverRequest(discover), + )) + } + + #[test] + fn a_bare_discovery_probe_is_reported_so_it_can_be_rejected_as_legacy() { + let message = discover_probe(1, None); + assert_eq!(bare_discovery_probe_id(&message), Some(NumberOrString::Number(1))); + } + + #[test] + fn a_probe_with_request_metadata_is_left_to_the_sdk() { + let mut meta = RequestMetaObject::default(); + meta.insert("io.modelcontextprotocol/protocolVersion".to_string(), serde_json::json!("2026-07-28")); + let message = discover_probe(7, Some(meta)); + assert_eq!(bare_discovery_probe_id(&message), None); + } +} diff --git a/docs/content/docs/mcp.cn.mdx b/docs/content/docs/mcp.cn.mdx index c8d31b6614..3085b7a2d2 100644 --- a/docs/content/docs/mcp.cn.mdx +++ b/docs/content/docs/mcp.cn.mdx @@ -233,7 +233,7 @@ Resource 的发现和读取会复用对应 Tool 的白名单，以及连接、 该 opt-in 模式仅支持本地原生 MySQL 连接。DBX Web/`DBX_WEB_URL`、外部驱动 profile、兼容引擎及其他数据库类型会返回 `TRANSACTION_UNSUPPORTED`，不会静默退化为自动提交。事务会话每次操作只接受一条经过解析的 `SELECT`（含锁定读）、`INSERT`、`UPDATE`、`DELETE` 或 `REPLACE`；原始事务控制、DDL/临时 DDL、`CALL`、`LOAD`、表锁、`SET`、`USE`、XA、文件/输出子句、可执行注释、优化器 hint 及多语句均会被拒绝。完整批次持有同一个操作槽位，并为每条语句返回执行后的事务状态。 -响应包含 `transaction_state`（`idle`、`active` 或 `unknown`），并在适用时包含 `transaction_outcome`（`committed`、`rolled_back` 或 `unknown`）。普通 MySQL 语句错误可能仍让事务保持 active；DBX 会在同一连接上重新探测服务器状态后再返回。超时、取消、传输故障或丢失 `COMMIT` 确认会把结果永久标为 `unknown`，禁止该会话继续执行 SQL，并在不重放语句或 `COMMIT` 的前提下销毁连接。显式正值的连接查询超时同样用于事务操作；当连接配置为“无限制”（`0`）时，事务 owner 仍使用默认 300 秒操作上限，确保 busy 操作保持有界。关闭、过期、协议会话删除和服务退出会对状态已知的 active 事务执行有界回滚，然后销毁连接。 +响应包含 `transaction_state`（`idle`、`active` 或 `unknown`），并在适用时包含 `transaction_outcome`（`committed`、`rolled_back` 或 `unknown`）。普通 MySQL 语句错误可能仍让事务保持 active；DBX 会在同一连接上重新探测服务器状态后再返回。超时、取消、传输故障或丢失 `COMMIT` 确认会把结果永久标为 `unknown`，禁止该会话继续执行 SQL，并在不重放语句或 `COMMIT` 的前提下销毁连接。显式正值的连接查询超时同样用于事务操作；当连接配置为“无限制”（`0`）时，事务 owner 仍使用默认 300 秒操作上限，确保 busy 操作保持有界。关闭、过期和服务退出会对状态已知的 active 事务执行有界回滚，然后销毁连接。会话的回收途径是 `dbx_close_session`、空闲超时或 MCP 服务退出；用 `DELETE` 结束 HTTP 传输会话不会回收 Agent 通过它打开的有状态会话，而无状态的 `2026-07-28` Agent 本身就没有传输会话可结束。 回滚保证依赖 InnoDB 等事务型存储。触发器、存储函数、非事务表及外部副作用可能产生 MySQL 无法回滚的效果。DBX 会在排队完成后重新检查当前连接/数据库范围、工具 allowlist、只读与生产保护以及已确认 SQL。即使写权限被撤销，仍允许原会话执行回滚和关闭以完成清理。 @@ -288,6 +288,15 @@ Desktop 可以选择直接托管 Streamable HTTP MCP 服务。该服务**默认 本机客户端建议保留默认回环监听。绑定到局域网地址时，必须开启 **允许远程访问**，并填写精确的允许 `Host`；浏览器客户端还必须填写精确的 `Origin`。客户端通过非默认端口访问时，允许的 Host 应写成 `host:port`。轮换 Token 会立即使旧凭证失效。 +同一个端点同时响应两代 MCP 协议，任何版本的 Agent 都无需特殊配置即可连接： + +| Agent | 握手方式 | 传输状态 | +| --- | --- | --- | +| 旧版客户端 | `initialize`（协议 `2025-11-25` 及更早） | 有状态：响应返回 `Mcp-Session-Id`，GET/DELETE 管理流 | +| 新版客户端 | `server/discover`（协议 `2026-07-28`） | 无状态：每个请求在 `_meta` 中自带协议版本与能力，不使用会话头 | + +协商到 `2026-07-28` 的新版 Agent 以无状态方式服务，同一监听端口上的旧版 Agent 仍保有会话。DBX 的有状态会话（`dbx_open_session`）在两种情况下都可用，因为 `session_id` 是由 Agent 在后续调用中显式回传的句柄。 + ### 本地连接 这是默认模式。MCP 使用 DBX 中保存的连接配置。原生可执行的连接不要求 DBX 桌面端保持运行；需要 Desktop bridge 或已安装 Agent/驱动的连接仍依赖对应运行时。 diff --git a/docs/content/docs/mcp.mdx b/docs/content/docs/mcp.mdx index f8821f8aaa..da4e9cd0e7 100644 --- a/docs/content/docs/mcp.mdx +++ b/docs/content/docs/mcp.mdx @@ -233,7 +233,7 @@ Set `enable_transactions: true` when calling `dbx_open_session` to reserve one p This opt-in mode is available only for local native MySQL connections. DBX Web/`DBX_WEB_URL`, external driver profiles, compatibility engines, and other database types return `TRANSACTION_UNSUPPORTED` instead of falling back to auto-commit. Transaction-session SQL accepts one parsed `SELECT` (including locking reads), `INSERT`, `UPDATE`, `DELETE`, or `REPLACE` statement per operation. It rejects raw transaction control, DDL and temporary DDL, `CALL`, `LOAD`, table locks, `SET`, `USE`, XA, file/output clauses, executable comments, optimizer hints, and multi-statements. Complete batches hold one operation slot and report the resulting transaction state for every statement. -Responses include `transaction_state` (`idle`, `active`, or `unknown`) and, when known, `transaction_outcome` (`committed`, `rolled_back`, or `unknown`). A normal MySQL statement error can leave the transaction active; DBX probes fresh state on the same connection before reporting it. A timeout, cancellation, transport failure, or lost `COMMIT` acknowledgement makes the outcome `unknown`, permanently blocks more SQL on that session, and disposes the connection without replaying the statement or `COMMIT`. A positive connection query timeout applies to transaction operations; when the connection is configured as unlimited (`0`), transaction owners still apply their default 300-second operation cap so busy work remains bounded. Closing, expiry, protocol-session deletion, and server shutdown attempt bounded rollback for a known active transaction and then dispose the connection. +Responses include `transaction_state` (`idle`, `active`, or `unknown`) and, when known, `transaction_outcome` (`committed`, `rolled_back`, or `unknown`). A normal MySQL statement error can leave the transaction active; DBX probes fresh state on the same connection before reporting it. A timeout, cancellation, transport failure, or lost `COMMIT` acknowledgement makes the outcome `unknown`, permanently blocks more SQL on that session, and disposes the connection without replaying the statement or `COMMIT`. A positive connection query timeout applies to transaction operations; when the connection is configured as unlimited (`0`), transaction owners still apply their default 300-second operation cap so busy work remains bounded. Closing, expiry, and server shutdown attempt bounded rollback for a known active transaction and then dispose the connection. Sessions are reclaimed by `dbx_close_session`, by the idle timeout, or by MCP server shutdown; ending an HTTP transport session with `DELETE` does not reclaim the stateful session an agent opened through it, and a stateless `2026-07-28` agent has no transport session to end. Rollback guarantees require transactional storage such as InnoDB. Triggers, stored functions, nontransactional tables, and external side effects can have effects that MySQL cannot roll back. DBX rechecks the current connection/database scope, tool allowlist, read-only and production protections, and confirmed SQL after queue waiting. Rollback and close remain available for cleanup after write permission is revoked. @@ -288,6 +288,15 @@ Use the generated values in an HTTP-capable MCP client. The transport-specific f Keep the default loopback binding for local clients. Binding to a LAN address requires turning on **Allow remote access** and setting the exact allowed `Host` authorities; browser clients also need their exact `Origin` values. Use `host:port` in an allowed Host when the client sends a non-default port. Rotating the token immediately invalidates the previous credential. +DBX answers both MCP protocol generations on the same endpoint, so any agent can connect without a version-specific setup: + +| Agent | Handshake | Transport state | +| --- | --- | --- | +| Older clients | `initialize` (protocol `2025-11-25` and earlier) | Stateful: the response carries `Mcp-Session-Id`, and GET/DELETE manage the stream | +| Modern clients | `server/discover` (protocol `2026-07-28`) | Stateless: every request carries its own protocol version and capabilities in `_meta`, with no session header | + +A modern agent negotiating `2026-07-28` is served statelessly while an older agent keeps its session on the same listener. Stateful DBX sessions (`dbx_open_session`) work in both cases because the session id is an explicit handle that the agent passes back on later calls. + ### Local connections MCP uses the connections saved in DBX. Connections that execute natively do not require the Desktop app to stay open; bridge-backed or installed Agent/driver paths still depend on their runtime. diff --git a/packages/mcp-server/tests/bin.test.ts b/packages/mcp-server/tests/bin.test.ts index 19d3b9304b..a9d83bd844 100644 --- a/packages/mcp-server/tests/bin.test.ts +++ b/packages/mcp-server/tests/bin.test.ts @@ -19,11 +19,19 @@ type InitializeResponse = { }; }; +type DiscoverResponse = { + id: number; + result: { + resultType?: string; + supportedVersions?: string[]; + capabilities?: Record; + _meta?: Record; + }; +}; + type ErrorResponse = { id: number; - error: { - code: number; - }; + error: { code: number }; }; test("hides the native server console window on Windows", async () => { @@ -32,7 +40,7 @@ test("hides the native server console window on Windows", async () => { assert.match(launcher, /windowsHide:\s*true/); }); -test("falls back from server discovery when invoked through an npm-style symlink", async () => { +test("answers server discovery and still serves the legacy initialize flow through an npm-style symlink", async () => { const bin = await symlinkedMcpServer(); let child: ChildProcessWithoutNullStreams | undefined; try { @@ -45,7 +53,10 @@ test("falls back from server discovery when invoked through an npm-style symlink }, }); - const discoveryPromise = readJsonRpcResponse(child, 5000); + // A client that probes discovery without the `_meta` the final spec made + // mandatory must still see the pre-upgrade `-32601`, which is the signal it + // treats as "legacy server, use initialize" and falls back on. + const bareProbePromise = readJsonRpcResponse(child, 5000); child.stdin.write( encodeMessage({ jsonrpc: "2.0", @@ -54,16 +65,45 @@ test("falls back from server discovery when invoked through an npm-style symlink params: {}, }), ); + const bareProbe = await bareProbePromise; + assert.equal(bareProbe.id, 1); + assert.equal(bareProbe.error.code, -32601); - const discovery = await discoveryPromise; - assert.equal(discovery.id, 1); - assert.equal(discovery.error.code, -32601); - - const responsePromise = readJsonRpcResponse(child, 5000); + // Modern agents start with `server/discover` (MCP 2026-07-28) and must get a + // real discovery result rather than the legacy `-32601` rejection. Per the + // 2026-07-28 spec every request carries its protocol version and client + // capabilities in `_meta`. + const discoveryPromise = readJsonRpcResponse(child, 5000); child.stdin.write( encodeMessage({ jsonrpc: "2.0", id: 2, + method: "server/discover", + params: { + _meta: { + "io.modelcontextprotocol/protocolVersion": "2026-07-28", + "io.modelcontextprotocol/clientInfo": { name: "dbx-test", version: "0.0.0" }, + "io.modelcontextprotocol/clientCapabilities": {}, + }, + }, + }), + ); + + const discovery = await discoveryPromise; + assert.equal(discovery.id, 2); + assert.ok(discovery.result, `discovery failed: ${JSON.stringify(discovery)}`); + assert.equal(discovery.result._meta?.["io.modelcontextprotocol/serverInfo"]?.name, "dbx"); + assert.ok( + discovery.result.supportedVersions?.includes("2026-07-28"), + `discovery must advertise 2026-07-28, got ${JSON.stringify(discovery.result.supportedVersions)}`, + ); + + // Legacy agents keep using the initialize handshake, unchanged. + const responsePromise = readJsonRpcResponse(child, 5000); + child.stdin.write( + encodeMessage({ + jsonrpc: "2.0", + id: 3, method: "initialize", params: { protocolVersion: "2024-11-05", @@ -75,7 +115,7 @@ test("falls back from server discovery when invoked through an npm-style symlink const response = await responsePromise; - assert.equal(response.id, 2); + assert.equal(response.id, 3); assert.equal(response.result.serverInfo.name, "dbx"); } finally { child?.kill(); diff --git a/scripts/dbx-mcp-check.ps1 b/scripts/dbx-mcp-check.ps1 new file mode 100644 index 0000000000..d46143369a --- /dev/null +++ b/scripts/dbx-mcp-check.ps1 @@ -0,0 +1,152 @@ +<# +.SYNOPSIS + Manually verify the dbx-mcp stdio handshake paths on Windows. +.DESCRIPTION + Runs the four paths that matter after the MCP protocol work: + 1. bare `server/discover` (no _meta) -> -32601, the legacy fallback signal + 2. `initialize` at 2025-11-25 -> legacy stateful session + 3. `server/discover` with _meta -> advertises 2026-07-28 + 4. `tools/call` with no initialize -> served statelessly +.EXAMPLE + .\dbx-mcp-check.ps1 -Binary .\target\debug\dbx-mcp.exe +.EXAMPLE + # against your real DBX connections + .\dbx-mcp-check.ps1 -Binary .\target\debug\dbx-mcp.exe -DataDir "$env:APPDATA\com.dbx.app" +#> +param( + [string]$Binary = ".\target\debug\dbx-mcp.exe", + [string]$DataDir = "", + [int]$TimeoutMs = 15000 +) + +$ErrorActionPreference = "Stop" + +if (-not (Test-Path $Binary)) { + Write-Error "binary not found: $Binary`nBuild it first:`n cargo build -p dbx-mcp --no-default-features --features sqlite-bundled" +} +$Binary = (Resolve-Path $Binary).Path + +$explicitDataDir = [bool]$DataDir +if (-not $DataDir) { + $DataDir = Join-Path $env:TEMP ("dbx-mcp-check-" + [guid]::NewGuid().ToString("N").Substring(0, 8)) + New-Item -ItemType Directory -Force -Path $DataDir | Out-Null +} +elseif (-not (Test-Path -LiteralPath $DataDir)) { + Write-Error "data directory does not exist: $DataDir`nPass -DataDir only with a directory that contains dbx.db, or omit it to test with an empty store." +} +elseif (-not (Test-Path -LiteralPath (Join-Path $DataDir "dbx.db")) -and -not (Test-Path -LiteralPath (Join-Path $DataDir "storage.db"))) { + Write-Warning "no dbx.db found in $DataDir; dbx_list_connections may report no connections" +} +$DataDir = (Resolve-Path -LiteralPath $DataDir).Path + +$env:DBX_DATA_DIR = $DataDir +$env:RUST_LOG = "warn" + +$meta = @{ + "io.modelcontextprotocol/protocolVersion" = "2026-07-28" + "io.modelcontextprotocol/clientInfo" = @{ name = "manual-check"; version = "0" } + "io.modelcontextprotocol/clientCapabilities" = @{} +} + +$failures = New-Object System.Collections.Generic.List[string] +$proc = $null + +function Invoke-Mcp { + param([hashtable]$Message, [switch]$NoReply) + $line = ($Message | ConvertTo-Json -Depth 12 -Compress) + $proc.StandardInput.WriteLine($line) + $proc.StandardInput.Flush() + if ($NoReply) { return $null } + $read = $proc.StandardOutput.ReadLineAsync() + if (-not $read.Wait($TimeoutMs)) { + $proc.Kill() + throw "timed out waiting for a reply to '$($Message.method)'; stderr: $($proc.StandardError.ReadToEnd())" + } + return ($read.Result | ConvertFrom-Json) +} + +function Invoke-Request { + param([int]$Id, [string]$Method, [hashtable]$Params) + $paramsJson = ($Params | ConvertTo-Json -Depth 12 -Compress) + if (-not $paramsJson -or $paramsJson -eq "{}") { $paramsJson = "{}" } + $raw = "{`"jsonrpc`":`"2.0`",`"id`":$Id,`"method`":`"$Method`",`"params`":$paramsJson}" + $proc.StandardInput.WriteLine($raw) + $proc.StandardInput.Flush() + $read = $proc.StandardOutput.ReadLineAsync() + if (-not $read.Wait($TimeoutMs)) { + $proc.Kill() + throw "timed out waiting for '$Method'; stderr: $($proc.StandardError.ReadToEnd())" + } + return ($read.Result | ConvertFrom-Json) +} + +function Invoke-Notification { + param([string]$Method) + $proc.StandardInput.WriteLine("{`"jsonrpc`":`"2.0`",`"method`":`"$Method`"}") + $proc.StandardInput.Flush() +} + +try { + $proc = New-Object System.Diagnostics.Process + $proc.StartInfo.FileName = $Binary + $proc.StartInfo.UseShellExecute = $false + $proc.StartInfo.RedirectStandardInput = $true + $proc.StartInfo.RedirectStandardOutput = $true + $proc.StartInfo.RedirectStandardError = $true + $proc.StartInfo.CreateNoWindow = $true + [void]$proc.Start() + + # 1. Bare discovery probe: agents older than the final 2026 lifecycle read + # -32601 as "this server is legacy, use initialize". + $bare = Invoke-Request -Id 1 -Method "server/discover" -Params @{} + $code = $bare.error.code + Write-Host "[1] bare server/discover -> error $code (expected -32601)" + if ($code -ne -32601) { $failures.Add("bare probe returned $code, expected -32601") } + + # 2. Legacy initialize still negotiates a pre-2026 version. + $init = Invoke-Request -Id 2 -Method "initialize" -Params @{ + protocolVersion = "2025-11-25" + capabilities = @{} + clientInfo = @{ name = "manual-check"; version = "0" } + } + $version = $init.result.protocolVersion + $server = $init.result.serverInfo.name + Write-Host "[2] initialize 2025-11-25 -> protocolVersion=$version serverInfo=$server" + if ($version -ne "2025-11-25" -or $server -ne "dbx") { $failures.Add("legacy initialize returned $($init | ConvertTo-Json -Compress)") } + + Invoke-Notification -Method "notifications/initialized" + $tools = (Invoke-Request -Id 3 -Method "tools/list" -Params @{}).result.tools + Write-Host " legacy tools/list -> $($tools.Count) tools" + if (-not $tools -or $tools.Count -eq 0) { $failures.Add("legacy tools/list returned no tools") } + + # 3. A well-formed probe must reach the SDK and advertise 2026-07-28. + $discover = Invoke-Request -Id 4 -Method "server/discover" -Params @{ _meta = $meta } + $versions = $discover.result.supportedVersions + Write-Host "[3] server/discover + _meta -> versions=$($versions -join ', ')" + $infoName = $discover.result._meta."io.modelcontextprotocol/serverInfo".name + Write-Host " serverInfo=$infoName ttlMs=$($discover.result.ttlMs) cacheScope=$($discover.result.cacheScope)" + if ($versions -notcontains "2026-07-28") { $failures.Add("discovery did not advertise 2026-07-28") } + + # 4. Stateless tools/call: no initialize at all, metadata per request. + $call = Invoke-Request -Id 5 -Method "tools/call" -Params @{ + name = "dbx_list_connections" + arguments = @{} + _meta = $meta + } + $text = $call.result.content[0].text + Write-Host "[4] stateless tools/call -> isError=$($call.result.isError) text='$($text -replace "`r?`n", ' ')'" + if ($call.result.isError) { $failures.Add("stateless tools/call failed: $($call | ConvertTo-Json -Compress)") } +} +finally { + if ($proc) { + try { if (-not $proc.HasExited) { $proc.Kill() } } catch { } + } +} + +if ($failures.Count -gt 0) { + Write-Host "`nFAILED:" -ForegroundColor Red + $failures | ForEach-Object { Write-Host " - $_" -ForegroundColor Red } + exit 1 +} +Write-Host "`nAll checks passed." -ForegroundColor Green +exit 0 diff --git a/scripts/manual-mcp-check.py b/scripts/manual-mcp-check.py new file mode 100644 index 0000000000..3f1a578abf --- /dev/null +++ b/scripts/manual-mcp-check.py @@ -0,0 +1,136 @@ +#!/usr/bin/env python3 +"""Manually verify the dbx-mcp stdio handshake paths without extra tooling. + +Usage: + python3 scripts/manual-mcp-check.py [path/to/dbx-mcp] + +Runs the four paths that matter after the MCP protocol work: + + 1. bare `server/discover` (no `_meta`) -> -32601, the legacy fallback signal + 2. `initialize` at 2025-11-25 -> legacy stateful session + 3. `server/discover` with `_meta` -> advertises 2026-07-28 + 4. `tools/call` with no `initialize` -> served statelessly +""" + +import json +import os +import subprocess +import sys +import tempfile + +DEFAULT_BINARY = "target/debug/dbx-mcp" + +META_2026 = { + "io.modelcontextprotocol/protocolVersion": "2026-07-28", + "io.modelcontextprotocol/clientInfo": {"name": "manual-check", "version": "0"}, + "io.modelcontextprotocol/clientCapabilities": {}, +} + + +class Probe: + def __init__(self, binary): + self.data_dir = tempfile.mkdtemp(prefix="dbx-mcp-manual-") + env = {**os.environ, "DBX_DATA_DIR": self.data_dir, "RUST_LOG": "warn"} + self.process = subprocess.Popen( + [binary], + stdin=subprocess.PIPE, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + env=env, + text=True, + bufsize=1, + ) + + def send(self, **message): + """Send one request and return the matching response.""" + self.process.stdin.write(json.dumps({"jsonrpc": "2.0", **message}) + "\n") + self.process.stdin.flush() + line = self.process.stdout.readline() + if not line.strip(): + raise RuntimeError(f"no response for {message.get('method')}; stderr: {self.stderr()}") + return json.loads(line) + + def notify(self, method): + self.process.stdin.write(json.dumps({"jsonrpc": "2.0", "method": method}) + "\n") + self.process.stdin.flush() + + def stderr(self): + if self.process.poll() is None: + return "" + return self.process.stderr.read() + + def close(self): + self.process.kill() + + +def main(): + binary = sys.argv[1] if len(sys.argv) > 1 else DEFAULT_BINARY + if not os.path.exists(binary): + if os.path.exists(os.path.join(os.path.dirname(__file__), "..", binary)): + binary = os.path.join(os.path.dirname(__file__), "..", binary) + else: + sys.exit(f"binary not found: {binary}\nBuild it first:\n cargo build -p dbx-mcp --no-default-features --features sqlite-bundled") + + failures = [] + probe = Probe(binary) + try: + # 1. Bare discovery probe: agents older than the final 2026 lifecycle use + # -32601 as the "this server is legacy, use initialize" signal. + bare = probe.send(id=1, method="server/discover", params={}) + code = bare.get("error", {}).get("code") + print(f"[1] bare server/discover -> error {code} (expected -32601)") + if code != -32601: + failures.append(f"bare probe returned {code}, expected -32601") + + # 2. Legacy initialize still negotiates a pre-2026 version. + init = probe.send( + id=2, + method="initialize", + params={"protocolVersion": "2025-11-25", "capabilities": {}, "clientInfo": {"name": "manual-check", "version": "0"}}, + ) + version = init.get("result", {}).get("protocolVersion") + server = init.get("result", {}).get("serverInfo", {}).get("name") + print(f"[2] initialize 2025-11-25 -> protocolVersion={version} serverInfo={server}") + if version != "2025-11-25" or server != "dbx": + failures.append(f"legacy initialize returned {init}") + + probe.notify("notifications/initialized") + legacy_tools = probe.send(id=3, method="tools/list", params={}).get("result", {}).get("tools", []) + print(f" legacy tools/list -> {len(legacy_tools)} tools") + if not legacy_tools: + failures.append("legacy tools/list returned no tools") + + # 3. A well-formed probe must reach the SDK and advertise 2026-07-28. + discover = probe.send(id=4, method="server/discover", params={"_meta": META_2026}) + result = discover.get("result", {}) + versions = result.get("supportedVersions", []) + print(f"[3] server/discover + _meta -> versions={versions}") + print(f" serverInfo={result.get('_meta', {}).get('io.modelcontextprotocol/serverInfo', {}).get('name')}" + f" ttlMs={result.get('ttlMs')} cacheScope={result.get('cacheScope')}") + if "2026-07-28" not in versions: + failures.append(f"discovery did not advertise 2026-07-28: {discover}") + + # 4. Stateless tools/call: no initialize at all, metadata per request. + call = probe.send( + id=5, + method="tools/call", + params={"name": "dbx_list_connections", "arguments": {}, "_meta": META_2026}, + ) + call_result = call.get("result", {}) + text = (call_result.get("content") or [{}])[0].get("text", "") + print(f"[4] stateless tools/call -> isError={call_result.get('isError')} text={text[:60]!r}") + if call_result.get("isError"): + failures.append(f"stateless tools/call failed: {call}") + finally: + probe.close() + + if failures: + print("\nFAILED:") + for failure in failures: + print(f" - {failure}") + sys.exit(1) + print("\nAll checks passed.") + + +if __name__ == "__main__": + main()