Merge branch 'main' into telemetry-update

Author: blva

src/helpers/EJsonTransport.ts

@@ -0,0 +1,47 @@
+import { JSONRPCMessage, JSONRPCMessageSchema } from "@modelcontextprotocol/sdk/types.js";
+import { EJSON } from "bson";
+import { StdioServerTransport } from "@modelcontextprotocol/sdk/server/stdio.js";
+
+// This is almost a copy of ReadBuffer from @modelcontextprotocol/sdk
+// but it uses EJSON.parse instead of JSON.parse to handle BSON types
+export class EJsonReadBuffer {
+    private _buffer?: Buffer;
+
+    append(chunk: Buffer): void {
+        this._buffer = this._buffer ? Buffer.concat([this._buffer, chunk]) : chunk;
+    }
+
+    readMessage(): JSONRPCMessage | null {
+        if (!this._buffer) {
+            return null;
+        }
+
+        const index = this._buffer.indexOf("\n");
+        if (index === -1) {
+            return null;
+        }
+
+        const line = this._buffer.toString("utf8", 0, index).replace(/\r$/, "");
+        this._buffer = this._buffer.subarray(index + 1);
+
+        // This is using EJSON.parse instead of JSON.parse to handle BSON types
+        return JSONRPCMessageSchema.parse(EJSON.parse(line));
+    }
+
+    clear(): void {
+        this._buffer = undefined;
+    }
+}
+
+// This is a hacky workaround for https://github.com/mongodb-js/mongodb-mcp-server/issues/211
+// The underlying issue is that StdioServerTransport uses JSON.parse to deserialize
+// messages, but that doesn't handle bson types, such as ObjectId when serialized as EJSON.
+//
+// This function creates a StdioServerTransport and replaces the internal readBuffer with EJsonReadBuffer
+// that uses EJson.parse instead.
+export function createEJsonTransport(): StdioServerTransport {
+    const server = new StdioServerTransport();
+    server["_readBuffer"] = new EJsonReadBuffer();
+
+    return server;
+}

src/index.ts

@@ -1,13 +1,13 @@
 #!/usr/bin/env node
 
-import { StdioServerTransport } from "@modelcontextprotocol/sdk/server/stdio.js";
 import logger, { LogId } from "./logger.js";
 import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js";
 import { config } from "./config.js";
 import { Session } from "./session.js";
 import { Server } from "./server.js";
 import { packageInfo } from "./helpers/packageInfo.js";
 import { Telemetry } from "./telemetry/telemetry.js";
+import { createEJsonTransport } from "./helpers/EJsonTransport.js";
 
 try {
     const session = new Session({
@@ -29,7 +29,7 @@ try {
         userConfig: config,
     });
 
-    const transport = new StdioServerTransport();
+    const transport = createEJsonTransport();
 
     await server.connect(transport);
 } catch (error: unknown) {

src/tools/atlas/atlasTool.ts

@@ -1,7 +1,9 @@
-import { ToolBase, ToolCategory, TelemetryToolMetadata } from "../tool.js";
+import { ToolBase, ToolCategory, TelemetryToolMetadata, ToolArgs } from "../tool.js";
 import { ToolCallback } from "@modelcontextprotocol/sdk/server/mcp.js";
+import { CallToolResult } from "@modelcontextprotocol/sdk/types.js";
 import logger, { LogId } from "../../logger.js";
 import { z } from "zod";
+import { ApiClientError } from "../../common/atlas/apiClientError.js";
 
 export abstract class AtlasToolBase extends ToolBase {
     protected category: ToolCategory = "atlas";
@@ -13,6 +15,50 @@ export abstract class AtlasToolBase extends ToolBase {
         return super.verifyAllowed();
     }
 
+    protected handleError(
+        error: unknown,
+        args: ToolArgs<typeof this.argsShape>
+    ): Promise<CallToolResult> | CallToolResult {
+        if (error instanceof ApiClientError) {
+            const statusCode = error.response.status;
+
+            if (statusCode === 401) {
+                return {
+                    content: [
+                        {
+                            type: "text",
+                            text: `Unable to authenticate with MongoDB Atlas, API error: ${error.message}
+
+Hint: Your API credentials may be invalid, expired or lack permissions. 
+Please check your Atlas API credentials and ensure they have the appropriate permissions.
+For more information on setting up API keys, visit: https://www.mongodb.com/docs/atlas/configure-api-access/`,
+                        },
+                    ],
+                    isError: true,
+                };
+            }
+
+            if (statusCode === 403) {
+                return {
+                    content: [
+                        {
+                            type: "text",
+                            text: `Received a Forbidden API Error: ${error.message}
+                            
+You don't have sufficient permissions to perform this action in MongoDB Atlas
+Please ensure your API key has the necessary roles assigned.
+For more information on Atlas API access roles, visit: https://www.mongodb.com/docs/atlas/api/service-accounts-overview/`,
+                        },
+                    ],
+                    isError: true,
+                };
+            }
+        }
+
+        // For other types of errors, use the default error handling from the base class
+        return super.handleError(error, args);
+    }
+
     /**
      *
      * Resolves the tool metadata from the arguments passed to the tool

tests/integration/helpers.ts

@@ -235,6 +235,7 @@ export function validateThrowsForInvalidArguments(
 }
 
 /** Expects the argument being defined and asserts it */
-export function expectDefined<T>(arg: T): asserts arg is Exclude<T, undefined> {
+export function expectDefined<T>(arg: T): asserts arg is Exclude<T, undefined | null> {
     expect(arg).toBeDefined();
+    expect(arg).not.toBeNull();
 }

tests/integration/tools/mongodb/read/find.test.ts

@@ -4,6 +4,7 @@ import {
     validateToolMetadata,
     validateThrowsForInvalidArguments,
     getResponseElements,
+    expectDefined,
 } from "../../../helpers.js";
 import { describeWithMongoDB, validateAutoConnectBehavior } from "../mongodbHelpers.js";
 
@@ -171,6 +172,33 @@ describeWithMongoDB("find tool", (integration) => {
                 expect(JSON.parse(elements[i + 1].text).value).toEqual(i);
             }
         });
+
+        it("can find objects by $oid", async () => {
+            await integration.connectMcpClient();
+
+            const fooObject = await integration
+                .mongoClient()
+                .db(integration.randomDbName())
+                .collection("foo")
+                .findOne();
+            expectDefined(fooObject);
+
+            const response = await integration.mcpClient().callTool({
+                name: "find",
+                arguments: {
+                    database: integration.randomDbName(),
+                    collection: "foo",
+                    filter: { _id: fooObject._id },
+                },
+            });
+
+            const elements = getResponseElements(response.content);
+            expect(elements).toHaveLength(2);
+            expect(elements[0].text).toEqual('Found 1 documents in the collection "foo":');
+
+            // eslint-disable-next-line @typescript-eslint/no-unsafe-member-access
+            expect(JSON.parse(elements[1].text).value).toEqual(fooObject.value);
+        });
     });
 
     validateAutoConnectBehavior(integration, "find", () => {

tests/unit/EJsonTransport.test.ts

@@ -0,0 +1,71 @@
+import { Decimal128, MaxKey, MinKey, ObjectId, Timestamp, UUID } from "bson";
+import { createEJsonTransport, EJsonReadBuffer } from "../../src/helpers/EJsonTransport.js";
+import { JSONRPCMessage } from "@modelcontextprotocol/sdk/types.js";
+import { AuthInfo } from "@modelcontextprotocol/sdk/server/auth/types.js";
+import { StdioServerTransport } from "@modelcontextprotocol/sdk/server/stdio.js";
+import { Readable } from "stream";
+import { ReadBuffer } from "@modelcontextprotocol/sdk/shared/stdio.js";
+
+describe("EJsonTransport", () => {
+    let transport: StdioServerTransport;
+    beforeEach(async () => {
+        transport = createEJsonTransport();
+        await transport.start();
+    });
+
+    afterEach(async () => {
+        await transport.close();
+    });
+
+    it("ejson deserializes messages", () => {
+        const messages: { message: JSONRPCMessage; extra?: { authInfo?: AuthInfo } }[] = [];
+        transport.onmessage = (
+            message,
+            extra?: {
+                authInfo?: AuthInfo;
+            }
+        ) => {
+            messages.push({ message, extra });
+        };
+
+        (transport["_stdin"] as Readable).emit(
+            "data",
+            Buffer.from(
+                '{"jsonrpc":"2.0","id":1,"method":"testMethod","params":{"oid":{"$oid":"681b741f13aa74a0687b5110"},"uuid":{"$uuid":"f81d4fae-7dec-11d0-a765-00a0c91e6bf6"},"date":{"$date":"2025-05-07T14:54:23.973Z"},"decimal":{"$numberDecimal":"1234567890987654321"},"int32":123,"maxKey":{"$maxKey":1},"minKey":{"$minKey":1},"timestamp":{"$timestamp":{"t":123,"i":456}}}}\n',
+                "utf-8"
+            )
+        );
+
+        expect(messages.length).toBe(1);
+        const message = messages[0].message;
+
+        expect(message).toEqual({
+            jsonrpc: "2.0",
+            id: 1,
+            method: "testMethod",
+            params: {
+                oid: new ObjectId("681b741f13aa74a0687b5110"),
+                uuid: new UUID("f81d4fae-7dec-11d0-a765-00a0c91e6bf6"),
+                date: new Date(Date.parse("2025-05-07T14:54:23.973Z")),
+                decimal: new Decimal128("1234567890987654321"),
+                int32: 123,
+                maxKey: new MaxKey(),
+                minKey: new MinKey(),
+                timestamp: new Timestamp({ t: 123, i: 456 }),
+            },
+        });
+    });
+
+    it("has _readBuffer field of type EJsonReadBuffer", () => {
+        expect(transport["_readBuffer"]).toBeDefined();
+        expect(transport["_readBuffer"]).toBeInstanceOf(EJsonReadBuffer);
+    });
+
+    describe("standard StdioServerTransport", () => {
+        it("has a _readBuffer field", () => {
+            const standardTransport = new StdioServerTransport();
+            expect(standardTransport["_readBuffer"]).toBeDefined();
+            expect(standardTransport["_readBuffer"]).toBeInstanceOf(ReadBuffer);
+        });
+    });
+});