mirror of
https://github.com/t8y2/dbx.git
synced 2026-10-19 16:16:45 +00:00
feat(mcp): serve stateful and stateless agents on both protocol generations
This commit is contained in:
KeaneQy
Generated
+47
-12
@@ -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",
|
||||
|
||||
@@ -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"] }
|
||||
|
||||
+233
-23
@@ -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<String>,
|
||||
web_mode: bool,
|
||||
) -> Result<Router, String> {
|
||||
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<CancellationToken>,
|
||||
session_manager: Arc<LocalSessionManager>,
|
||||
sessions: Arc<McpSessionStore>,
|
||||
pending_salesforce_writes: Arc<PendingSalesforceWrites>,
|
||||
) -> Result<Router, String> {
|
||||
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<DbxMcpServer, LocalSessionManager> = 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<LocalSessionManager>) {
|
||||
/// 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<dyn DbxBackend>,
|
||||
session_manager: &Arc<LocalSessionManager>,
|
||||
sessions: &Arc<McpSessionStore>,
|
||||
) {
|
||||
let session_ids = session_manager.sessions.read().await.keys().cloned().collect::<Vec<_>>();
|
||||
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<ConnectionConfig, String> {
|
||||
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<bool, String> {
|
||||
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<HttpTestBackend>,
|
||||
keep_alive: std::time::Duration,
|
||||
) -> (String, Arc<LocalSessionManager>, CancellationToken, tokio::task::JoinHandle<()>) {
|
||||
) -> (String, Arc<LocalSessionManager>, Arc<McpSessionStore>, 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::<String>::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<dyn DbxBackend> = 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<rmcp::RoleClient, ()>, 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<rmcp::RoleClient, ()>, 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;
|
||||
|
||||
|
||||
@@ -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<HashMap<String, PendingSalesforceWrite>>,
|
||||
}
|
||||
|
||||
impl PendingSalesforceWrites {
|
||||
fn new(ttl: Duration) -> Arc<Self> {
|
||||
pub(crate) fn new(ttl: Duration) -> Arc<Self> {
|
||||
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<dyn DbxBackend>,
|
||||
scope: McpScope,
|
||||
web_mode: bool,
|
||||
plugin_tools_mode: PluginToolsMode,
|
||||
sessions: Arc<McpSessionStore>,
|
||||
pending_salesforce_writes: Arc<PendingSalesforceWrites>,
|
||||
) -> Self {
|
||||
Self::build(backend, scope, web_mode, plugin_tools_mode, sessions, pending_salesforce_writes)
|
||||
}
|
||||
|
||||
fn build(
|
||||
backend: Arc<dyn DbxBackend>,
|
||||
scope: McpScope,
|
||||
web_mode: bool,
|
||||
plugin_tools_mode: PluginToolsMode,
|
||||
sessions: Arc<McpSessionStore>,
|
||||
pending_salesforce_writes: Arc<PendingSalesforceWrites>,
|
||||
) -> 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<SessionCleanupResult> {
|
||||
@@ -879,7 +912,7 @@ impl DbxMcpServer {
|
||||
result_rx
|
||||
}
|
||||
|
||||
async fn close_backend_sessions_best_effort(&self, sessions: Vec<McpSession>) {
|
||||
pub(crate) async fn close_backend_sessions_best_effort(&self, sessions: Vec<McpSession>) {
|
||||
let cleanups = sessions.into_iter().map(|session| self.spawn_session_cleanup(session)).collect::<Vec<_>>();
|
||||
for cleanup in cleanups {
|
||||
let _ = cleanup.await;
|
||||
@@ -3646,7 +3679,7 @@ impl ServerHandler for DbxMcpServer {
|
||||
&self,
|
||||
request: rmcp::model::CallToolRequestParams,
|
||||
context: rmcp::service::RequestContext<rmcp::RoleServer>,
|
||||
) -> Result<CallToolResult, rmcp::ErrorData> {
|
||||
) -> Result<CallToolResponse, rmcp::ErrorData> {
|
||||
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<CallToolResponse, rmcp::ErrorData> = 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<rmcp::service::RoleServer>,
|
||||
) -> Result<ReadResourceResult, ErrorData> {
|
||||
) -> Result<ReadResourceResponse, ErrorData> {
|
||||
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<rmcp::model::PaginatedRequestParams>,
|
||||
_context: rmcp::service::RequestContext<rmcp::service::RoleServer>,
|
||||
) -> Result<rmcp::model::ListToolsResult, rmcp::ErrorData> {
|
||||
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,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -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<McpSession> {
|
||||
let mut state = self.state.lock().await;
|
||||
let closing = state.closing.values().cloned().collect::<Vec<_>>();
|
||||
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
|
||||
|
||||
@@ -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<T, E, A>(transport: T) -> impl Transport<RoleServer, Error = E> + 'static
|
||||
where
|
||||
T: IntoTransport<RoleServer, E, A>,
|
||||
@@ -35,23 +56,14 @@ where
|
||||
async fn receive(&mut self) -> Option<RxJsonRpcMessage<RoleServer>> {
|
||||
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<rmcp::model::RequestId> {
|
||||
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::<RequestMetaObject>().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<RequestMetaObject>) -> 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);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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 本身就没有传输会话可结束。
|
||||
|
||||
<Callout type="warn">回滚保证依赖 InnoDB 等事务型存储。触发器、存储函数、非事务表及外部副作用可能产生 MySQL 无法回滚的效果。DBX 会在排队完成后重新检查当前连接/数据库范围、工具 allowlist、只读与生产保护以及已确认 SQL。即使写权限被撤销，仍允许原会话执行回滚和关闭以完成清理。</Callout>
|
||||
|
||||
@@ -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/驱动的连接仍依赖对应运行时。
|
||||
|
||||
@@ -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.
|
||||
|
||||
<Callout type="warn">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.</Callout>
|
||||
|
||||
@@ -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.
|
||||
|
||||
+51
-11
@@ -19,11 +19,19 @@ type InitializeResponse = {
|
||||
};
|
||||
};
|
||||
|
||||
type DiscoverResponse = {
|
||||
id: number;
|
||||
result: {
|
||||
resultType?: string;
|
||||
supportedVersions?: string[];
|
||||
capabilities?: Record<string, unknown>;
|
||||
_meta?: Record<string, { name?: string }>;
|
||||
};
|
||||
};
|
||||
|
||||
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<ErrorResponse>(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<ErrorResponse>(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<InitializeResponse>(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<DiscoverResponse>(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<InitializeResponse>(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();
|
||||
|
||||
@@ -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
|
||||
@@ -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()
|
||||
Reference in New Issue
Block a user