Skip to content

Commit 2816bba

Browse files
committed
fix: Use JSON for stdio, similiar to the HTTP Transport and use EJSON.deserialize when necessary
We can't leverage using EJSON as the underlying transport because it diverges from the HTTP Transport, and because not all EJSON is equivalent to it's underlying JSON. For example, take this sample into consideration: ```js $> JSON.stringify({ long: 23098479807234 }) <= {"long":23098479807234} $> EJSON.stringify({ long: 23098479807234 }) // relaxed = true <= {"long":23098479807234} $> EJSON.stringify({ long: 23098479807234 }, { relaxed: false }) <= "long":{"$numberLong":"23098479807234"}} ``` This behaviour forbids us on using `{relaxed: false}` always, as it breaks the shape of the underlying JSON. Moreover, if we use dates, it gets worse, as we lose type info in the protocol when serialising/deserializing. The approach here is to use a custom `zod` type that leverages on EJSON.deserialize, a function that reads a plain EJSON object and transforms it to the underlying types (for example, $date to js.Date). This is important because the Node.js driver does not work with EJSON, it works with actual JS objects (like Date), so by using this approach documents and queries using non-primitive types like BinData or Date work as expected.
1 parent dd36b1a commit 2816bba

File tree

13 files changed

+63
-146
lines changed

13 files changed

+63
-146
lines changed

src/tools/args.ts

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
11
import { z, type ZodString } from "zod";
2+
import { EJSON } from "bson";
23

34
const NO_UNICODE_REGEX = /^[\x20-\x7E]*$/;
45
export const NO_UNICODE_ERROR = "String cannot contain special characters or Unicode symbols";
@@ -68,3 +69,13 @@ export const AtlasArgs = {
6869
password: (): z.ZodString =>
6970
z.string().min(1, "Password is required").max(100, "Password must be 100 characters or less"),
7071
};
72+
73+
function toEJSON<T extends object | undefined>(value: T): T {
74+
if (!value) {
75+
return value;
76+
}
77+
78+
return EJSON.deserialize(value, { relaxed: false }) as T;
79+
}
80+
81+
export const zEJSON = () => z.object({}).passthrough().transform(toEJSON);

src/tools/mongodb/create/insertMany.ts

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2,14 +2,15 @@ import { z } from "zod";
22
import type { CallToolResult } from "@modelcontextprotocol/sdk/types.js";
33
import { DbOperationArgs, MongoDBToolBase } from "../mongodbTool.js";
44
import type { ToolArgs, OperationType } from "../../tool.js";
5+
import { zEJSON } from "../../args.js";
56

67
export class InsertManyTool extends MongoDBToolBase {
78
public name = "insert-many";
89
protected description = "Insert an array of documents into a MongoDB collection";
910
protected argsShape = {
1011
...DbOperationArgs,
1112
documents: z
12-
.array(z.object({}).passthrough().describe("An individual MongoDB document"))
13+
.array(zEJSON().describe("An individual MongoDB document"))
1314
.describe(
1415
"The array of documents to insert, matching the syntax of the document argument of db.collection.insertMany()"
1516
),

src/tools/mongodb/delete/deleteMany.ts

Lines changed: 2 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,18 +1,16 @@
1-
import { z } from "zod";
21
import type { CallToolResult } from "@modelcontextprotocol/sdk/types.js";
32
import { DbOperationArgs, MongoDBToolBase } from "../mongodbTool.js";
43
import type { ToolArgs, OperationType } from "../../tool.js";
54
import { checkIndexUsage } from "../../../helpers/indexCheck.js";
65
import { EJSON } from "bson";
6+
import { zEJSON } from "../../args.js";
77

88
export class DeleteManyTool extends MongoDBToolBase {
99
public name = "delete-many";
1010
protected description = "Removes all documents that match the filter from a MongoDB collection";
1111
protected argsShape = {
1212
...DbOperationArgs,
13-
filter: z
14-
.object({})
15-
.passthrough()
13+
filter: zEJSON()
1614
.optional()
1715
.describe(
1816
"The query filter, specifying the deletion criteria. Matches the syntax of the filter argument of db.collection.deleteMany()"

src/tools/mongodb/read/aggregate.ts

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -6,9 +6,10 @@ import { formatUntrustedData } from "../../tool.js";
66
import { checkIndexUsage } from "../../../helpers/indexCheck.js";
77
import { EJSON } from "bson";
88
import { ErrorCodes, MongoDBError } from "../../../common/errors.js";
9+
import { zEJSON } from "../../args.js";
910

1011
export const AggregateArgs = {
11-
pipeline: z.array(z.object({}).passthrough()).describe("An array of aggregation stages to execute"),
12+
pipeline: z.array(zEJSON()).describe("An array of aggregation stages to execute"),
1213
};
1314

1415
export class AggregateTool extends MongoDBToolBase {

src/tools/mongodb/read/count.ts

Lines changed: 2 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,13 +1,11 @@
11
import type { CallToolResult } from "@modelcontextprotocol/sdk/types.js";
22
import { DbOperationArgs, MongoDBToolBase } from "../mongodbTool.js";
33
import type { ToolArgs, OperationType } from "../../tool.js";
4-
import { z } from "zod";
54
import { checkIndexUsage } from "../../../helpers/indexCheck.js";
5+
import { zEJSON } from "../../args.js";
66

77
export const CountArgs = {
8-
query: z
9-
.object({})
10-
.passthrough()
8+
query: zEJSON()
119
.optional()
1210
.describe(
1311
"A filter/query parameter. Allows users to filter the documents to count. Matches the syntax of the filter argument of db.collection.count()."

src/tools/mongodb/read/find.ts

Lines changed: 2 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -6,11 +6,10 @@ import { formatUntrustedData } from "../../tool.js";
66
import type { SortDirection } from "mongodb";
77
import { checkIndexUsage } from "../../../helpers/indexCheck.js";
88
import { EJSON } from "bson";
9+
import { zEJSON } from "../../args.js";
910

1011
export const FindArgs = {
11-
filter: z
12-
.object({})
13-
.passthrough()
12+
filter: zEJSON()
1413
.optional()
1514
.describe("The query filter, matching the syntax of the query argument of db.collection.find()"),
1615
projection: z

src/tools/mongodb/update/updateMany.ts

Lines changed: 5 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -3,23 +3,21 @@ import type { CallToolResult } from "@modelcontextprotocol/sdk/types.js";
33
import { DbOperationArgs, MongoDBToolBase } from "../mongodbTool.js";
44
import type { ToolArgs, OperationType } from "../../tool.js";
55
import { checkIndexUsage } from "../../../helpers/indexCheck.js";
6+
import { zEJSON } from "../../args.js";
67

78
export class UpdateManyTool extends MongoDBToolBase {
89
public name = "update-many";
910
protected description = "Updates all documents that match the specified filter for a collection";
1011
protected argsShape = {
1112
...DbOperationArgs,
12-
filter: z
13-
.object({})
14-
.passthrough()
13+
filter: zEJSON()
1514
.optional()
1615
.describe(
1716
"The selection criteria for the update, matching the syntax of the filter argument of db.collection.updateOne()"
1817
),
19-
update: z
20-
.object({})
21-
.passthrough()
22-
.describe("An update document describing the modifications to apply using update operator expressions"),
18+
update: zEJSON().describe(
19+
"An update document describing the modifications to apply using update operator expressions"
20+
),
2321
upsert: z
2422
.boolean()
2523
.optional()

src/transports/stdio.ts

Lines changed: 1 addition & 49 deletions
Original file line numberDiff line numberDiff line change
@@ -1,55 +1,8 @@
1-
import { EJSON } from "bson";
2-
import type { JSONRPCMessage } from "@modelcontextprotocol/sdk/types.js";
3-
import { JSONRPCMessageSchema } from "@modelcontextprotocol/sdk/types.js";
41
import { StdioServerTransport } from "@modelcontextprotocol/sdk/server/stdio.js";
52
import { LogId } from "../common/logger.js";
63
import type { Server } from "../server.js";
74
import { TransportRunnerBase, type TransportRunnerConfig } from "./base.js";
85

9-
// This is almost a copy of ReadBuffer from @modelcontextprotocol/sdk
10-
// but it uses EJSON.parse instead of JSON.parse to handle BSON types
11-
export class EJsonReadBuffer {
12-
private _buffer?: Buffer;
13-
14-
append(chunk: Buffer): void {
15-
this._buffer = this._buffer ? Buffer.concat([this._buffer, chunk]) : chunk;
16-
}
17-
18-
readMessage(): JSONRPCMessage | null {
19-
if (!this._buffer) {
20-
return null;
21-
}
22-
23-
const index = this._buffer.indexOf("\n");
24-
if (index === -1) {
25-
return null;
26-
}
27-
28-
const line = this._buffer.toString("utf8", 0, index).replace(/\r$/, "");
29-
this._buffer = this._buffer.subarray(index + 1);
30-
31-
// This is using EJSON.parse instead of JSON.parse to handle BSON types
32-
return JSONRPCMessageSchema.parse(EJSON.parse(line));
33-
}
34-
35-
clear(): void {
36-
this._buffer = undefined;
37-
}
38-
}
39-
40-
// This is a hacky workaround for https://github.com/mongodb-js/mongodb-mcp-server/issues/211
41-
// The underlying issue is that StdioServerTransport uses JSON.parse to deserialize
42-
// messages, but that doesn't handle bson types, such as ObjectId when serialized as EJSON.
43-
//
44-
// This function creates a StdioServerTransport and replaces the internal readBuffer with EJsonReadBuffer
45-
// that uses EJson.parse instead.
46-
export function createStdioTransport(): StdioServerTransport {
47-
const server = new StdioServerTransport();
48-
server["_readBuffer"] = new EJsonReadBuffer();
49-
50-
return server;
51-
}
52-
536
export class StdioRunner extends TransportRunnerBase {
547
private server: Server | undefined;
558

@@ -60,8 +13,7 @@ export class StdioRunner extends TransportRunnerBase {
6013
async start(): Promise<void> {
6114
try {
6215
this.server = await this.setupServer();
63-
64-
const transport = createStdioTransport();
16+
const transport = new StdioServerTransport();
6517

6618
await this.server.connect(transport);
6719
} catch (error: unknown) {

tests/integration/indexCheck.test.ts

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -80,7 +80,7 @@ describe("IndexCheck integration tests", () => {
8080
arguments: {
8181
database: integration.randomDbName(),
8282
collection: "find-test-collection",
83-
filter: { _id: docs[0]?._id }, // Uses _id index (IDHACK)
83+
filter: { _id: { $oid: docs[0]?._id } }, // Uses _id index (IDHACK)
8484
},
8585
});
8686

tests/integration/tools/mongodb/create/insertMany.test.ts

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -76,7 +76,7 @@ describeWithMongoDB("insertMany tool", (integration) => {
7676
arguments: {
7777
database: integration.randomDbName(),
7878
collection: "coll1",
79-
documents: [{ prop1: "value1", _id: insertedIds[0] }],
79+
documents: [{ prop1: "value1", _id: { $oid: insertedIds[0] } }],
8080
},
8181
});
8282

0 commit comments

Comments
 (0)