diff --git a/src/common/exportsManager.ts b/src/common/exportsManager.ts index ad90f8cc2..f8ce94501 100644 --- a/src/common/exportsManager.ts +++ b/src/common/exportsManager.ts @@ -27,6 +27,7 @@ interface CommonExportData { interface ReadyExport extends CommonExportData { exportStatus: "ready"; exportCreatedAt: number; + docsTransformed: number; } interface InProgressExport extends CommonExportData { @@ -124,7 +125,7 @@ export class ExportsManager extends EventEmitter { } } - public async readExport(exportName: string): Promise { + public async readExport(exportName: string): Promise<{ content: string; docsTransformed: number }> { try { this.assertIsNotShuttingDown(); exportName = decodeAndNormalize(exportName); @@ -137,9 +138,12 @@ export class ExportsManager extends EventEmitter { throw new Error("Requested export is still being generated. Try again later."); } - const { exportPath } = exportHandle; + const { exportPath, docsTransformed } = exportHandle; - return fs.readFile(exportPath, { encoding: "utf8", signal: this.shutdownController.signal }); + return { + content: await fs.readFile(exportPath, { encoding: "utf8", signal: this.shutdownController.signal }), + docsTransformed, + }; } catch (error) { this.logger.error({ id: LogId.exportReadError, @@ -202,17 +206,15 @@ export class ExportsManager extends EventEmitter { }): Promise { try { let pipeSuccessful = false; + let docsTransformed = 0; try { await fs.mkdir(this.exportsDirectoryPath, { recursive: true }); const outputStream = createWriteStream(inProgressExport.exportPath); - await pipeline( - [ - input.stream(), - this.docToEJSONStream(this.getEJSONOptionsForFormat(jsonExportFormat)), - outputStream, - ], - { signal: this.shutdownController.signal } - ); + const ejsonTransform = this.docToEJSONStream(this.getEJSONOptionsForFormat(jsonExportFormat)); + await pipeline([input.stream(), ejsonTransform, outputStream], { + signal: this.shutdownController.signal, + }); + docsTransformed = ejsonTransform.docsTransformed; pipeSuccessful = true; } catch (error) { // If the pipeline errors out then we might end up with @@ -231,6 +233,7 @@ export class ExportsManager extends EventEmitter { ...inProgressExport, exportCreatedAt: Date.now(), exportStatus: "ready", + docsTransformed, }; this.emit("export-available", inProgressExport.exportURI); } @@ -256,33 +259,39 @@ export class ExportsManager extends EventEmitter { } } - private docToEJSONStream(ejsonOptions: EJSONOptions | undefined): Transform { + private docToEJSONStream(ejsonOptions: EJSONOptions | undefined): Transform & { docsTransformed: number } { let docsTransformed = 0; - return new Transform({ - objectMode: true, - transform(chunk: unknown, encoding, callback): void { - try { - const doc = EJSON.stringify(chunk, undefined, undefined, ejsonOptions); + const result = Object.assign( + new Transform({ + objectMode: true, + transform(chunk: unknown, encoding, callback): void { + try { + const doc = EJSON.stringify(chunk, undefined, undefined, ejsonOptions); + if (docsTransformed === 0) { + this.push("[" + doc); + } else { + this.push(",\n" + doc); + } + docsTransformed++; + callback(); + } catch (err) { + callback(err as Error); + } + }, + flush(callback): void { if (docsTransformed === 0) { - this.push("[" + doc); + this.push("[]"); } else { - this.push(",\n" + doc); + this.push("]"); } - docsTransformed++; + result.docsTransformed = docsTransformed; callback(); - } catch (err) { - callback(err as Error); - } - }, - flush(callback): void { - if (docsTransformed === 0) { - this.push("[]"); - } else { - this.push("]"); - } - callback(); - }, - }); + }, + }), + { docsTransformed } + ); + + return result; } private async cleanupExpiredExports(): Promise { diff --git a/src/resources/common/exportedData.ts b/src/resources/common/exportedData.ts index 2ae4ba80e..7fed7dbab 100644 --- a/src/resources/common/exportedData.ts +++ b/src/resources/common/exportedData.ts @@ -7,6 +7,7 @@ import { ResourceTemplate } from "@modelcontextprotocol/sdk/server/mcp.js"; import type { Server } from "../../server.js"; import { LogId } from "../../common/logger.js"; import type { Session } from "../../common/session.js"; +import { formatUntrustedData } from "../../tools/tool.js"; export class ExportedData { private readonly name = "exported-data"; @@ -95,13 +96,17 @@ export class ExportedData { throw new Error("Cannot retrieve exported data, exportName not provided."); } - const content = await this.session.exportsManager.readExport(exportName); + const { content, docsTransformed } = await this.session.exportsManager.readExport(exportName); + + const text = formatUntrustedData(`The exported data contains ${docsTransformed} documents.`, content) + .map((t) => t.text) + .join("\n"); return { contents: [ { uri: url.href, - text: content, + text, mimeType: "application/json", }, ], diff --git a/tests/integration/resources/exportedData.test.ts b/tests/integration/resources/exportedData.test.ts index 394bed254..6e361bf03 100644 --- a/tests/integration/resources/exportedData.test.ts +++ b/tests/integration/resources/exportedData.test.ts @@ -1,9 +1,9 @@ import path from "path"; import fs from "fs/promises"; -import { Long } from "bson"; +import { EJSON, Long, ObjectId } from "bson"; import { describe, expect, it, beforeEach, afterAll } from "vitest"; import type { CallToolResult } from "@modelcontextprotocol/sdk/types.js"; -import { defaultTestConfig, resourceChangedNotification, timeout } from "../helpers.js"; +import { defaultTestConfig, getDataFromUntrustedContent, resourceChangedNotification, timeout } from "../helpers.js"; import { describeWithMongoDB } from "../tools/mongodb/mongodbHelpers.js"; import { contentWithResourceURILink } from "../tools/mongodb/read/export.test.js"; import type { UserConfig } from "../../../src/lib.js"; @@ -18,15 +18,17 @@ const userConfig: UserConfig = { describeWithMongoDB( "exported-data resource", (integration) => { + let docs: { _id: ObjectId; name: string; longNumber?: Long; bigInt?: Long }[]; + let collection: string; + beforeEach(async () => { const mongoClient = integration.mongoClient(); - await mongoClient - .db("db") - .collection("coll") - .insertMany([ - { name: "foo", longNumber: new Long(1234) }, - { name: "bar", bigInt: new Long(123412341234) }, - ]); + collection = new ObjectId().toString(); + docs = [ + { name: "foo", longNumber: new Long(1234), _id: new ObjectId() }, + { name: "bar", bigInt: new Long(123412341234), _id: new ObjectId() }, + ]; + await mongoClient.db("db").collection(collection).insertMany(docs); }); afterAll(async () => { @@ -67,7 +69,7 @@ describeWithMongoDB( name: "export", arguments: { database: "db", - collection: "coll", + collection, exportTitle: "Export for db.coll", exportTarget: [{ name: "find", arguments: {} }], }, @@ -106,7 +108,7 @@ describeWithMongoDB( name: "export", arguments: { database: "db", - collection: "coll", + collection, exportTitle: "Export for db.coll", exportTarget: [{ name: "find", arguments: {} }], }, @@ -125,7 +127,16 @@ describeWithMongoDB( }); expect(response.isError).toBeFalsy(); expect(response.contents[0]?.mimeType).toEqual("application/json"); - expect(response.contents[0]?.text).toContain("foo"); + + expect(response.contents[0]?.text).toContain(`The exported data contains ${docs.length} documents.`); + expect(response.contents[0]?.text).toContain(" { @@ -134,7 +145,7 @@ describeWithMongoDB( name: "export", arguments: { database: "big", - collection: "coll", + collection, exportTitle: "Export for big.coll", exportTarget: [{ name: "find", arguments: {} }], }, diff --git a/tests/unit/common/exportsManager.test.ts b/tests/unit/common/exportsManager.test.ts index 81759e0ad..bfc1eba18 100644 --- a/tests/unit/common/exportsManager.test.ts +++ b/tests/unit/common/exportsManager.test.ts @@ -235,7 +235,9 @@ describe("ExportsManager unit test", () => { jsonExportFormat: "relaxed", }); await exportAvailableNotifier; - expect(await manager.readExport(exportName)).toEqual("[]"); + const { content, docsTransformed } = await manager.readExport(exportName); + expect(content).toEqual("[]"); + expect(docsTransformed).toEqual(0); }); it("should handle encoded name", async () => { @@ -249,7 +251,9 @@ describe("ExportsManager unit test", () => { jsonExportFormat: "relaxed", }); await exportAvailableNotifier; - expect(await manager.readExport(encodeURIComponent(exportName))).toEqual("[]"); + const { content, docsTransformed } = await manager.readExport(encodeURIComponent(exportName)); + expect(content).toEqual("[]"); + expect(docsTransformed).toEqual(0); }); }); @@ -332,7 +336,7 @@ describe("ExportsManager unit test", () => { expect(emitSpy).toHaveBeenCalledWith("export-available", exportURI); // Exports relaxed json - const jsonData = JSON.parse(await manager.readExport(exportName)) as unknown[]; + const jsonData = JSON.parse((await manager.readExport(exportName)).content) as unknown[]; expect(jsonData).toEqual([]); }); }); @@ -366,7 +370,7 @@ describe("ExportsManager unit test", () => { expect(emitSpy).toHaveBeenCalledWith("export-available", `exported-data://${expectedExportName}`); // Exports relaxed json - const jsonData = JSON.parse(await manager.readExport(expectedExportName)) as unknown[]; + const jsonData = JSON.parse((await manager.readExport(expectedExportName)).content) as unknown[]; expect(jsonData).toContainEqual(expect.objectContaining({ name: "foo", longNumber: 12 })); expect(jsonData).toContainEqual(expect.objectContaining({ name: "bar", longNumber: 123456 })); }); @@ -401,7 +405,7 @@ describe("ExportsManager unit test", () => { expect(emitSpy).toHaveBeenCalledWith("export-available", `exported-data://${expectedExportName}`); // Exports relaxed json - const jsonData = JSON.parse(await manager.readExport(expectedExportName)) as unknown[]; + const jsonData = JSON.parse((await manager.readExport(expectedExportName)).content) as unknown[]; expect(jsonData).toContainEqual( expect.objectContaining({ name: "foo", longNumber: { $numberLong: "12" } }) );