fix(mcp): handle discovery probes before initialization

crates/dbx-mcp/src/lib.rs

@@ -2,8 +2,10 @@ pub mod backend;
 pub mod paths;
 pub mod server;
 pub mod session;
+pub mod transport;
 
 pub use backend::{ConnectionSummary, DbxBackend, LocalBackend, WebBackend};
 pub use dbx_core::mongo_shell as mongo;
 pub use server::{DbxMcpServer, McpScope};
 pub use session::McpSessionStore;
+pub use transport::with_legacy_discovery_fallback;

crates/dbx-mcp/src/main.rs

@@ -1,6 +1,6 @@
 use std::sync::Arc;
 
-use dbx_mcp::{DbxBackend, DbxMcpServer, LocalBackend, WebBackend};
+use dbx_mcp::{with_legacy_discovery_fallback, DbxBackend, DbxMcpServer, LocalBackend, WebBackend};
 use rmcp::ServiceExt;
 
 #[tokio::main]
@@ -14,7 +14,8 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
         let db_path = dbx_mcp::paths::storage_db_path().map_err(std::io::Error::other)?;
         Arc::new(LocalBackend::open(&db_path).await.map_err(std::io::Error::other)?)
     };
-    let service = DbxMcpServer::new(backend).serve(rmcp::transport::stdio()).await?;
+    let transport = with_legacy_discovery_fallback(rmcp::transport::stdio());
+    let service = DbxMcpServer::new(backend).serve(transport).await?;
     service.waiting().await?;
     Ok(())
 }

crates/dbx-mcp/src/transport.rs

@@ -0,0 +1,67 @@
+use std::future::Future;
+
+use rmcp::{
+    model::{ClientJsonRpcMessage, ClientRequest, ErrorCode, ErrorData, ServerJsonRpcMessage},
+    service::{RxJsonRpcMessage, TxJsonRpcMessage},
+    transport::{IntoTransport, Transport},
+    RoleServer,
+};
+
+pub fn with_legacy_discovery_fallback<T, E, A>(transport: T) -> impl Transport<RoleServer, Error = E> + 'static
+where
+    T: IntoTransport<RoleServer, E, A>,
+    E: std::error::Error + Send + Sync + 'static,
+{
+    LegacyDiscoveryFallback { inner: transport.into_transport() }
+}
+
+struct LegacyDiscoveryFallback<T> {
+    inner: T,
+}
+
+impl<T> Transport<RoleServer> for LegacyDiscoveryFallback<T>
+where
+    T: Transport<RoleServer>,
+{
+    type Error = T::Error;
+
+    fn send(
+        &mut self,
+        item: TxJsonRpcMessage<RoleServer>,
+    ) -> impl Future<Output = Result<(), Self::Error>> + Send + 'static {
+        self.inner.send(item)
+    }
+
+    fn receive(&mut self) -> impl Future<Output = Option<RxJsonRpcMessage<RoleServer>>> + Send {
+        async move {
+            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 {
+                    return Some(message);
+                };
+
+                // Reject discovery without closing so new clients can fall back to the legacy initialize flow.
+                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;
+                }
+            }
+        }
+    }
+
+    fn close(&mut self) -> impl Future<Output = Result<(), Self::Error>> + Send {
+        self.inner.close()
+    }
+}

packages/mcp-server/tests/bin.test.ts

@@ -19,7 +19,14 @@ type InitializeResponse = {
   };
 };
 
-test("responds to initialize when invoked through an npm-style symlink", async () => {
+type ErrorResponse = {
+  id: number;
+  error: {
+    code: number;
+  };
+};
+
+test("falls back from server discovery when invoked through an npm-style symlink", async () => {
   const bin = await symlinkedMcpServer();
   let child: ChildProcessWithoutNullStreams | undefined;
   try {
@@ -32,11 +39,25 @@ test("responds to initialize when invoked through an npm-style symlink", async (
       },
     });
 
-    const responsePromise = readJsonRpcResponse(child, 5000);
+    const discoveryPromise = readJsonRpcResponse<ErrorResponse>(child, 5000);
     child.stdin.write(
       encodeMessage({
         jsonrpc: "2.0",
         id: 1,
+        method: "server/discover",
+        params: {},
+      }),
+    );
+
+    const discovery = await discoveryPromise;
+    assert.equal(discovery.id, 1);
+    assert.equal(discovery.error.code, -32601);
+
+    const responsePromise = readJsonRpcResponse<InitializeResponse>(child, 5000);
+    child.stdin.write(
+      encodeMessage({
+        jsonrpc: "2.0",
+        id: 2,
         method: "initialize",
         params: {
           protocolVersion: "2024-11-05",
@@ -48,7 +69,7 @@ test("responds to initialize when invoked through an npm-style symlink", async (
 
     const response = await responsePromise;
 
-    assert.equal(response.id, 1);
+    assert.equal(response.id, 2);
     assert.equal(response.result.serverInfo.name, "dbx");
   } finally {
     child?.kill();
@@ -67,7 +88,7 @@ function encodeMessage(payload: unknown): string {
   return `${JSON.stringify(payload)}\n`;
 }
 
-function readJsonRpcResponse(child: ChildProcessWithoutNullStreams, timeoutMs: number): Promise<InitializeResponse> {
+function readJsonRpcResponse<T>(child: ChildProcessWithoutNullStreams, timeoutMs: number): Promise<T> {
   return new Promise((resolve, reject) => {
     let stdout = "";
     let stderr = "";