Skip to content

fix: use ejson parsing for stdio messages #218

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 3 commits into from
May 8, 2025
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
47 changes: 47 additions & 0 deletions src/helpers/EJsonTransport.ts
Original file line number Diff line number Diff line change
@@ -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();
Comment on lines +42 to +44
Copy link
Preview

Copilot AI May 7, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[nitpick] Accessing a private property via server["_readBuffer"] may lead to maintainability issues if the underlying implementation changes. Consider exposing a dedicated method or refactoring the transport to safely update the read buffer.

Suggested change
export function createEJsonTransport(): StdioServerTransport {
const server = new StdioServerTransport();
server["_readBuffer"] = new EJsonReadBuffer();
class EJsonStdioServerTransport extends StdioServerTransport {
setReadBuffer(buffer: EJsonReadBuffer): void {
(this as any)._readBuffer = buffer; // Accessing private property safely within the subclass
}
}
export function createEJsonTransport(): StdioServerTransport {
const server = new EJsonStdioServerTransport();
server.setReadBuffer(new EJsonReadBuffer());

Copilot uses AI. Check for mistakes.


return server;
}
4 changes: 2 additions & 2 deletions src/index.ts
Original file line number Diff line number Diff line change
@@ -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({
Expand All @@ -29,7 +29,7 @@ try {
userConfig: config,
});

const transport = new StdioServerTransport();
const transport = createEJsonTransport();

await server.connect(transport);
} catch (error: unknown) {
Expand Down
3 changes: 2 additions & 1 deletion tests/integration/helpers.ts
Original file line number Diff line number Diff line change
Expand Up @@ -227,6 +227,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();
}
28 changes: 28 additions & 0 deletions tests/integration/tools/mongodb/read/find.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@ import {
validateToolMetadata,
validateThrowsForInvalidArguments,
getResponseElements,
expectDefined,
} from "../../../helpers.js";
import { describeWithMongoDB, validateAutoConnectBehavior } from "../mongodbHelpers.js";

Expand Down Expand Up @@ -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", () => {
Expand Down
71 changes: 71 additions & 0 deletions tests/unit/EJsonTransport.test.ts
Original file line number Diff line number Diff line change
@@ -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);
});
});
});
Loading