From a988af5534e36263bf2f03d4bea600c8633bd435 Mon Sep 17 00:00:00 2001 From: Himanshu Singh Date: Tue, 5 Aug 2025 13:18:06 +0200 Subject: [PATCH 01/50] feat: adds a session bound export manager --- package-lock.json | 5 +- package.json | 2 +- src/common/config.ts | 24 +- src/common/logger.ts | 4 + src/common/session.ts | 2 + src/common/sessionExportsManager.ts | 261 +++++++++++++++++ .../common/sessionExportsManager.test.ts | 273 ++++++++++++++++++ 7 files changed, 559 insertions(+), 12 deletions(-) create mode 100644 src/common/sessionExportsManager.ts create mode 100644 tests/integration/common/sessionExportsManager.test.ts diff --git a/package-lock.json b/package-lock.json index d6e7d7dd..d99d5a6f 100644 --- a/package-lock.json +++ b/package-lock.json @@ -27,6 +27,7 @@ "node-machine-id": "1.1.12", "oauth4webapi": "^3.6.0", "openapi-fetch": "^0.14.0", + "proper-lockfile": "^4.1.2", "yargs-parser": "^22.0.0", "zod": "^3.25.76" }, @@ -58,7 +59,6 @@ "openapi-types": "^12.1.3", "openapi-typescript": "^7.8.0", "prettier": "^3.6.2", - "proper-lockfile": "^4.1.2", "simple-git": "^3.28.0", "tsx": "^4.20.3", "typescript": "^5.8.3", @@ -8646,7 +8646,6 @@ "version": "4.2.11", "resolved": "https://registry.npmjs.org/graceful-fs/-/graceful-fs-4.2.11.tgz", "integrity": "sha512-RbJ5/jmFcNNCcDV5o9eTnBLJ/HszWV0P73bc+Ff4nS/rJj+YaS6IGyiOL0VoBYX+l1Wrl3k63h/KrH+nhJ0XvQ==", - "dev": true, "license": "ISC" }, "node_modules/graphemer": { @@ -11300,7 +11299,6 @@ "version": "4.1.2", "resolved": "https://registry.npmjs.org/proper-lockfile/-/proper-lockfile-4.1.2.tgz", "integrity": "sha512-TjNPblN4BwAWMXU8s9AEz4JmQxnD1NNL7bNOY/AKUzyamc379FWASUhc/K1pL2noVb+XmZKLL68cjzLsiOAMaA==", - "dev": true, "license": "MIT", "dependencies": { "graceful-fs": "^4.2.4", @@ -11733,7 +11731,6 @@ "version": "0.12.0", "resolved": "https://registry.npmjs.org/retry/-/retry-0.12.0.tgz", "integrity": "sha512-9LkiTwjUh6rT555DtE9rTX+BKByPfrMzEAtnlEtdEwr3Nkffwiihqe2bWADg+OQRjt9gl6ICdmB/ZFDCGAtSow==", - "dev": true, "license": "MIT", "engines": { "node": ">= 4" diff --git a/package.json b/package.json index 7bcd5a1e..2b1cf9f1 100644 --- a/package.json +++ b/package.json @@ -80,7 +80,6 @@ "openapi-types": "^12.1.3", "openapi-typescript": "^7.8.0", "prettier": "^3.6.2", - "proper-lockfile": "^4.1.2", "simple-git": "^3.28.0", "tsx": "^4.20.3", "typescript": "^5.8.3", @@ -108,6 +107,7 @@ "node-machine-id": "1.1.12", "oauth4webapi": "^3.6.0", "openapi-fetch": "^0.14.0", + "proper-lockfile": "^4.1.2", "yargs-parser": "^22.0.0", "zod": "^3.25.76" }, diff --git a/src/common/config.ts b/src/common/config.ts index cfcffb3d..0f99d855 100644 --- a/src/common/config.ts +++ b/src/common/config.ts @@ -19,6 +19,9 @@ export interface UserConfig { apiClientSecret?: string; telemetry: "enabled" | "disabled"; logPath: string; + exportPath: string; + exportTimeoutMs: number; + exportCleanupIntervalMs: number; connectionString?: string; connectOptions: ConnectOptions; disabledTools: Array; @@ -35,6 +38,9 @@ export interface UserConfig { const defaults: UserConfig = { apiBaseUrl: "https://cloud.mongodb.com/", logPath: getLogPath(), + exportPath: getExportPath(), + exportTimeoutMs: 300000, // 5 minutes + exportCleanupIntervalMs: 120000, // 2 minutes connectOptions: { readConcern: "local", readPreference: "secondaryPreferred", @@ -59,17 +65,21 @@ export const config = { ...getCliConfig(), }; -function getLogPath(): string { - const localDataPath = - process.platform === "win32" - ? path.join(process.env.LOCALAPPDATA || process.env.APPDATA || os.homedir(), "mongodb") - : path.join(os.homedir(), ".mongodb"); - - const logPath = path.join(localDataPath, "mongodb-mcp", ".app-logs"); +function getLocalDataPath(): string { + return process.platform === "win32" + ? path.join(process.env.LOCALAPPDATA || process.env.APPDATA || os.homedir(), "mongodb") + : path.join(os.homedir(), ".mongodb"); +} +function getLogPath(): string { + const logPath = path.join(getLocalDataPath(), "mongodb-mcp", ".app-logs"); return logPath; } +function getExportPath(): string { + return path.join(getLocalDataPath(), "mongodb-mcp", "exports"); +} + // Gets the config supplied by the user as environment variables. The variable names // are prefixed with `MDB_MCP_` and the keys match the UserConfig keys, but are converted // to SNAKE_UPPER_CASE. diff --git a/src/common/logger.ts b/src/common/logger.ts index 90bf97be..b33a08fc 100644 --- a/src/common/logger.ts +++ b/src/common/logger.ts @@ -49,6 +49,10 @@ export const LogId = { streamableHttpTransportSessionCloseNotificationFailure: mongoLogId(1_006_004), streamableHttpTransportRequestFailure: mongoLogId(1_006_005), streamableHttpTransportCloseFailure: mongoLogId(1_006_006), + + exportCleanupError: mongoLogId(1_007_001), + exportCreationError: mongoLogId(1_007_002), + exportReadError: mongoLogId(1_007_003), } as const; interface LogPayload { diff --git a/src/common/session.ts b/src/common/session.ts index 11d6e12b..a7358a9f 100644 --- a/src/common/session.ts +++ b/src/common/session.ts @@ -24,6 +24,8 @@ export type SessionEvents = { close: []; disconnect: []; "connection-error": [string]; + "export-expired": [string]; + "export-available": [string]; }; export class Session extends EventEmitter { diff --git a/src/common/sessionExportsManager.ts b/src/common/sessionExportsManager.ts new file mode 100644 index 00000000..754443ca --- /dev/null +++ b/src/common/sessionExportsManager.ts @@ -0,0 +1,261 @@ +import z from "zod"; +import path from "path"; +import fs from "fs/promises"; +import { createWriteStream } from "fs"; +import { lock } from "proper-lockfile"; +import { FindCursor } from "mongodb"; +import { EJSON, EJSONOptions } from "bson"; +import { Transform } from "stream"; +import { pipeline } from "stream/promises"; + +import { UserConfig } from "./config.js"; +import { Session } from "./session.js"; +import logger, { LogId } from "./logger.js"; + +export const jsonExportFormat = z.enum(["relaxed", "canonical"]); +export type JSONExportFormat = z.infer; + +export type Export = { + name: string; + uri: string; + createdAt: number; +}; + +export type SessionExportsManagerConfig = Pick< + UserConfig, + "exportPath" | "exportTimeoutMs" | "exportCleanupIntervalMs" +>; + +export class SessionExportsManager { + private mutableExports: Export[] = []; + private exportsCleanupInterval: NodeJS.Timeout; + private exportsCleanupInProgress: boolean = false; + + constructor( + private readonly session: Session, + private readonly config: SessionExportsManagerConfig + ) { + this.exportsCleanupInterval = setInterval( + () => void this.cleanupExpiredExports(), + this.config.exportCleanupIntervalMs + ); + } + + public close() { + clearInterval(this.exportsCleanupInterval); + } + + public exportNameToResourceURI(nameWithExtension: string): string { + if (!path.extname(nameWithExtension)) { + throw new Error("Provided export name has no extension"); + } + return `exported-data://${nameWithExtension}`; + } + + public exportsDirectoryPath(): string { + // If the session is not connected, we can't cannot work with exports + // for that session. + if (!this.session.sessionId) { + throw new Error("Cannot retrieve exports directory, no active session. Try to reconnect to the MCP server"); + } + + return path.join(this.config.exportPath, this.session.sessionId); + } + + public exportFilePath(exportsDirectoryPath: string, exportNameWithExtension: string): string { + if (!path.extname(exportNameWithExtension)) { + throw new Error("Provided export name has no extension"); + } + return path.join(exportsDirectoryPath, exportNameWithExtension); + } + + public listAvailableExports(): Export[] { + // Note that we don't account for ongoing cleanup or creation operation, + // by not acquiring a lock on read. That is because this we require this + // interface to be fast and just accurate enough for MCP completions + // API. + return this.mutableExports.filter(({ createdAt }) => { + return !this.isExportExpired(createdAt); + }); + } + + public async readExport(exportNameWithExtension: string): Promise { + try { + const exportsDirectoryPath = await this.ensureExportsDirectory(); + const exportFilePath = this.exportFilePath(exportsDirectoryPath, exportNameWithExtension); + if (await this.isExportFileExpired(exportFilePath)) { + throw new Error("Export has expired"); + } + + return await fs.readFile(exportFilePath, "utf8"); + } catch (error) { + logger.error( + LogId.exportReadError, + "Error when reading export", + error instanceof Error ? error.message : String(error) + ); + throw error; + } + } + + public async createJSONExport({ + input, + exportName, + jsonExportFormat, + }: { + input: FindCursor; + exportName: string; + jsonExportFormat: JSONExportFormat; + }): Promise { + try { + await this.withExportsLock(async (exportsDirectoryPath) => { + const exportNameWithExtension = this.withExtension(exportName, "json"); + const exportFilePath = path.join(exportsDirectoryPath, exportNameWithExtension); + const outputStream = createWriteStream(exportFilePath); + outputStream.write("["); + try { + const inputStream = input.stream(); + const ejsonOptions = this.getEJSONOptionsForFormat(jsonExportFormat); + await pipeline([inputStream, this.docToEJSONStream(ejsonOptions), outputStream]); + } finally { + outputStream.write("]\n"); + const resourceURI = this.exportNameToResourceURI(exportNameWithExtension); + this.mutableExports = [ + ...this.mutableExports, + { + createdAt: (await fs.stat(exportFilePath)).birthtimeMs, + name: exportNameWithExtension, + uri: resourceURI, + }, + ]; + this.session.emit("export-available", resourceURI); + void input.close(); + } + }); + } catch (error) { + logger.error( + LogId.exportCreationError, + "Error when generating JSON export", + error instanceof Error ? error.message : String(error) + ); + throw error; + } + } + + private getEJSONOptionsForFormat(format: JSONExportFormat): EJSONOptions | undefined { + if (format === "relaxed") { + return { + relaxed: true, + }; + } + return format === "canonical" + ? { + relaxed: false, + } + : undefined; + } + + private docToEJSONStream(ejsonOptions: EJSONOptions | undefined) { + let docsTransformed = 0; + return new Transform({ + objectMode: true, + transform: function (chunk: unknown, encoding, callback) { + ++docsTransformed; + try { + const doc: string = EJSON.stringify(chunk, undefined, 2, ejsonOptions); + const line = `${docsTransformed > 1 ? ",\n" : ""}${doc}`; + + callback(null, line); + } catch (err: unknown) { + callback(err as Error); + } + }, + final: function (callback) { + this.push("]"); + callback(null); + }, + }); + } + + private async cleanupExpiredExports(): Promise { + if (this.exportsCleanupInProgress) { + return; + } + + this.exportsCleanupInProgress = true; + try { + await this.withExportsLock(async (exportsDirectoryPath) => { + const exports = await this.listExportFiles(); + for (const exportName of exports) { + const exportPath = this.exportFilePath(exportsDirectoryPath, exportName); + if (await this.isExportFileExpired(exportPath)) { + await fs.unlink(exportPath); + this.mutableExports = this.mutableExports.filter(({ name }) => name !== exportName); + this.session.emit("export-expired", this.exportNameToResourceURI(exportName)); + } + } + }); + } catch (error) { + logger.error( + LogId.exportCleanupError, + "Error when cleaning up exports", + error instanceof Error ? error.message : String(error) + ); + } finally { + this.exportsCleanupInProgress = false; + } + } + + /** + * Small utility to centrally determine if an export is expired or not */ + private async isExportFileExpired(exportFilePath: string): Promise { + const stats = await fs.stat(exportFilePath); + return this.isExportExpired(stats.birthtimeMs); + } + + private isExportExpired(createdAt: number) { + return Date.now() - createdAt > this.config.exportTimeoutMs; + } + + /** + * Ensures the path ends with the provided extension */ + private withExtension(pathOrName: string, extension: string): string { + const extWithDot = extension.startsWith(".") ? extension : `.${extension}`; + if (path.extname(pathOrName) === extWithDot) { + return pathOrName; + } + return `${pathOrName}${extWithDot}`; + } + + /** + * Creates the session exports directory and returns the path */ + private async ensureExportsDirectory(): Promise { + const exportsDirectoryPath = this.exportsDirectoryPath(); + await fs.mkdir(exportsDirectoryPath, { recursive: true }); + return exportsDirectoryPath; + } + + /** + * Acquires a lock on the session exports directory. */ + private async withExportsLock(callback: (lockedPath: string) => Promise): Promise { + let releaseLock: (() => Promise) | undefined; + const exportsDirectoryPath = await this.ensureExportsDirectory(); + try { + releaseLock = await lock(exportsDirectoryPath, { retries: 10 }); + return await callback(exportsDirectoryPath); + } finally { + await releaseLock?.(); + } + } + + /** + * Lists exported files in the session export directory, while ignoring the + * hidden files and files without extensions. */ + private async listExportFiles(): Promise { + const exportsDirectory = await this.ensureExportsDirectory(); + const directoryContents = await fs.readdir(exportsDirectory, "utf8"); + return directoryContents.filter((maybeExportName) => { + return !maybeExportName.startsWith(".") && !!path.extname(maybeExportName); + }); + } +} diff --git a/tests/integration/common/sessionExportsManager.test.ts b/tests/integration/common/sessionExportsManager.test.ts new file mode 100644 index 00000000..7d6ed023 --- /dev/null +++ b/tests/integration/common/sessionExportsManager.test.ts @@ -0,0 +1,273 @@ +import path from "path"; +import fs from "fs/promises"; +import { beforeEach, describe, expect, it, vi } from "vitest"; +import { SessionExportsManager, SessionExportsManagerConfig } from "../../../src/common/sessionExportsManager.js"; +import { config } from "../../../src/common/config.js"; +import { Session } from "../../../src/common/session.js"; +import { ROOT_DIR } from "../../accuracy/sdk/constants.js"; +import { FindCursor, Long } from "mongodb"; +import { Readable } from "stream"; + +const dummySessionId = "1FOO"; +const dummyExportsPath = path.join(ROOT_DIR, "tests", "tmp", "exports"); +const dummySessionExportPath = path.join(dummyExportsPath, dummySessionId); +const exportsManagerConfig: SessionExportsManagerConfig = { + exportPath: dummyExportsPath, + exportTimeoutMs: config.exportTimeoutMs, + exportCleanupIntervalMs: config.exportCleanupIntervalMs, +} as const; +const dummyExportName = "foo.bar.json"; +const dummyExportPath = path.join(dummySessionExportPath, dummyExportName); + +async function createDummyExport() { + const content = "[]"; + await fs.mkdir(dummySessionExportPath, { recursive: true }); + await fs.writeFile(dummyExportPath, content); + return { + name: dummyExportName, + path: dummyExportPath, + content, + }; +} + +function createDummyFindCursor(dataArray: unknown[]): FindCursor { + let index = 0; + const readable = new Readable({ + objectMode: true, + read() { + if (index < dataArray.length) { + this.push(dataArray[index++]); + } else { + this.push(null); + } + }, + }); + + return { + stream() { + return readable; + }, + close() { + return Promise.resolve(readable.destroy()); + }, + } as unknown as FindCursor; +} + +async function fileExists(filePath: string) { + try { + await fs.access(filePath); + return true; + } catch { + return false; + } +} + +function timeout(ms: number) { + return new Promise((resolve) => setTimeout(resolve, ms)); +} + +describe("SessionExportsManager integration test", () => { + let session: Session; + let manager: SessionExportsManager; + + beforeEach(async () => { + manager?.close(); + await fs.rm(exportsManagerConfig.exportPath, { recursive: true, force: true }); + await fs.mkdir(exportsManagerConfig.exportPath, { recursive: true }); + session = new Session({ apiBaseUrl: "" }); + manager = new SessionExportsManager(session, exportsManagerConfig); + }); + + describe("#exportNameToResourceURI", function () { + it("should throw when export name has no extension", function () { + expect(() => manager.exportNameToResourceURI("name")).toThrow(); + }); + + it("should return a resource URI", function () { + expect(manager.exportNameToResourceURI("name.json")).toEqual("exported-data://name.json"); + }); + }); + + describe("#exportsDirectoryPath", function () { + it("should throw when session is not initialized", function () { + expect(() => manager.exportsDirectoryPath()).toThrow(); + }); + + it("should return a session path when session is initialized", function () { + session.sessionId = dummySessionId; + manager = new SessionExportsManager(session, exportsManagerConfig); + expect(manager.exportsDirectoryPath()).toEqual(path.join(exportsManagerConfig.exportPath, dummySessionId)); + }); + }); + + describe("#exportFilePath", function () { + it("should throw when export name has no extension", function () { + expect(() => manager.exportFilePath(dummySessionExportPath, "name")).toThrow(); + }); + + it("should return path to provided export file", function () { + expect(manager.exportFilePath(dummySessionExportPath, "mflix.movies.json")).toEqual( + path.join(dummySessionExportPath, "mflix.movies.json") + ); + }); + }); + + describe("#readExport", function () { + it("should throw when export name has no extension", async function () { + await expect(() => manager.readExport("name")).rejects.toThrow(); + }); + + it("should return the resource content", async function () { + const { name, content } = await createDummyExport(); + session.sessionId = dummySessionId; + manager = new SessionExportsManager(session, exportsManagerConfig); + expect(await manager.readExport(name)).toEqual(content); + }); + }); + + describe("#createJSONExport", function () { + let inputCursor: FindCursor; + beforeEach(() => { + void inputCursor?.close(); + inputCursor = createDummyFindCursor([ + { + name: "foo", + longNumber: Long.fromNumber(12), + }, + { + name: "bar", + longNumber: Long.fromNumber(123456), + }, + ]); + }); + + it.each([ + { cond: "when exportName does not contain extension", exportName: "foo.bar" }, + { cond: "when exportName contains extension", exportName: "foo.bar.json" }, + ])( + "$cond, should export relaxed json, update available exports and emit export-available event", + async function ({ exportName }) { + const emitSpy = vi.spyOn(session, "emit"); + session.sessionId = dummySessionId; + manager = new SessionExportsManager(session, exportsManagerConfig); + await manager.createJSONExport({ + input: inputCursor, + exportName, + jsonExportFormat: "relaxed", + }); + + // Updates available export + const availableExports = manager.listAvailableExports(); + expect(availableExports).toHaveLength(1); + expect(availableExports).toContainEqual( + expect.objectContaining({ + name: "foo.bar.json", + uri: "exported-data://foo.bar.json", + }) + ); + + // Emit event + expect(emitSpy).toHaveBeenCalledWith("export-available", "exported-data://foo.bar.json"); + + // Exports relaxed json + const jsonData = JSON.parse(await manager.readExport("foo.bar.json")) as unknown[]; + expect(jsonData).toContainEqual(expect.objectContaining({ name: "foo", longNumber: 12 })); + expect(jsonData).toContainEqual(expect.objectContaining({ name: "bar", longNumber: 123456 })); + } + ); + + it.each([ + { cond: "when exportName does not contain extension", exportName: "foo.bar" }, + { cond: "when exportName contains extension", exportName: "foo.bar.json" }, + ])( + "$cond, should export canonical json, update available exports and emit export-available event", + async function ({ exportName }) { + const emitSpy = vi.spyOn(session, "emit"); + session.sessionId = dummySessionId; + manager = new SessionExportsManager(session, exportsManagerConfig); + await manager.createJSONExport({ + input: inputCursor, + exportName, + jsonExportFormat: "canonical", + }); + + // Updates available export + const availableExports = manager.listAvailableExports(); + expect(availableExports).toHaveLength(1); + expect(availableExports).toContainEqual( + expect.objectContaining({ + name: "foo.bar.json", + uri: "exported-data://foo.bar.json", + }) + ); + + // Emit event + expect(emitSpy).toHaveBeenCalledWith("export-available", "exported-data://foo.bar.json"); + + // Exports relaxed json + const jsonData = JSON.parse(await manager.readExport("foo.bar.json")) as unknown[]; + expect(jsonData).toContainEqual( + expect.objectContaining({ name: "foo", longNumber: { $numberLong: "12" } }) + ); + expect(jsonData).toContainEqual( + expect.objectContaining({ name: "bar", longNumber: { $numberLong: "123456" } }) + ); + } + ); + }); + + describe("#cleanupExpiredExports", function () { + let input: FindCursor; + beforeEach(() => { + void input?.close(); + input = createDummyFindCursor([ + { + name: "foo", + longNumber: Long.fromNumber(12), + }, + { + name: "bar", + longNumber: Long.fromNumber(123456), + }, + ]); + }); + + it("should do nothing if session is not initialized", async function () { + const { path } = await createDummyExport(); + new SessionExportsManager(session, { + ...exportsManagerConfig, + exportTimeoutMs: 100, + exportCleanupIntervalMs: 50, + }); + + expect(await fileExists(path)).toEqual(true); + await timeout(200); + expect(await fileExists(path)).toEqual(true); + }); + + it("should cleanup expired exports if session is initialized", async function () { + session.sessionId = dummySessionId; + const manager = new SessionExportsManager(session, { + ...exportsManagerConfig, + exportTimeoutMs: 100, + exportCleanupIntervalMs: 50, + }); + await manager.createJSONExport({ + input, + exportName: dummyExportName, + jsonExportFormat: "relaxed", + }); + + expect(manager.listAvailableExports()).toContainEqual( + expect.objectContaining({ + name: "foo.bar.json", + uri: "exported-data://foo.bar.json", + }) + ); + expect(await fileExists(dummyExportPath)).toEqual(true); + await timeout(200); + expect(manager.listAvailableExports()).toEqual([]); + expect(await fileExists(dummyExportPath)).toEqual(false); + }); + }); +}); From dbf8ccf4a942bbf87c11a4ff784a9245f5988ddd Mon Sep 17 00:00:00 2001 From: Himanshu Singh Date: Tue, 5 Aug 2025 13:49:06 +0200 Subject: [PATCH 02/50] feat: adds an export tool also makes SessionExportManager available on MongoDBToolBase --- src/server.ts | 6 ++- src/tools/mongodb/mongodbTool.ts | 3 ++ src/tools/mongodb/read/export.ts | 85 ++++++++++++++++++++++++++++++++ src/transports/base.ts | 4 ++ 4 files changed, 97 insertions(+), 1 deletion(-) create mode 100644 src/tools/mongodb/read/export.ts diff --git a/src/server.ts b/src/server.ts index c23207a9..7f645882 100644 --- a/src/server.ts +++ b/src/server.ts @@ -13,9 +13,11 @@ import { type ServerCommand } from "./telemetry/types.js"; import { CallToolRequestSchema, CallToolResult } from "@modelcontextprotocol/sdk/types.js"; import assert from "assert"; import { ToolBase } from "./tools/tool.js"; +import { SessionExportsManager } from "./common/sessionExportsManager.js"; export interface ServerOptions { session: Session; + exportsManager: SessionExportsManager; userConfig: UserConfig; mcpServer: McpServer; telemetry: Telemetry; @@ -23,15 +25,17 @@ export interface ServerOptions { export class Server { public readonly session: Session; + public readonly exportsManager: SessionExportsManager; public readonly mcpServer: McpServer; private readonly telemetry: Telemetry; public readonly userConfig: UserConfig; public readonly tools: ToolBase[] = []; private readonly startTime: number; - constructor({ session, mcpServer, userConfig, telemetry }: ServerOptions) { + constructor({ session, exportsManager, mcpServer, userConfig, telemetry }: ServerOptions) { this.startTime = Date.now(); this.session = session; + this.exportsManager = exportsManager; this.telemetry = telemetry; this.mcpServer = mcpServer; this.userConfig = userConfig; diff --git a/src/tools/mongodb/mongodbTool.ts b/src/tools/mongodb/mongodbTool.ts index 708209f8..0d24aa72 100644 --- a/src/tools/mongodb/mongodbTool.ts +++ b/src/tools/mongodb/mongodbTool.ts @@ -5,6 +5,7 @@ import { CallToolResult } from "@modelcontextprotocol/sdk/types.js"; import { ErrorCodes, MongoDBError } from "../../common/errors.js"; import { LogId } from "../../common/logger.js"; import { Server } from "../../server.js"; +import { SessionExportsManager } from "../../common/sessionExportsManager.js"; export const DbOperationArgs = { database: z.string().describe("Database name"), @@ -13,6 +14,7 @@ export const DbOperationArgs = { export abstract class MongoDBToolBase extends ToolBase { private server?: Server; + public exportsManager?: SessionExportsManager; public category: ToolCategory = "mongodb"; protected async ensureConnected(): Promise { @@ -47,6 +49,7 @@ export abstract class MongoDBToolBase extends ToolBase { public register(server: Server): boolean { this.server = server; + this.exportsManager = server.exportsManager; return super.register(server); } diff --git a/src/tools/mongodb/read/export.ts b/src/tools/mongodb/read/export.ts new file mode 100644 index 00000000..ce721e10 --- /dev/null +++ b/src/tools/mongodb/read/export.ts @@ -0,0 +1,85 @@ +import { CallToolResult } from "@modelcontextprotocol/sdk/types.js"; +import { OperationType, ToolArgs } from "../../tool.js"; +import { DbOperationArgs, MongoDBToolBase } from "../mongodbTool.js"; +import { FindArgs } from "./find.js"; +import { jsonExportFormat } from "../../../common/sessionExportsManager.js"; +import z from "zod"; + +export class ExportTool extends MongoDBToolBase { + public name = "export"; + protected description = "Export a collection data or query results in the specified json format."; + protected argsShape = { + ...DbOperationArgs, + ...FindArgs, + limit: z.number().optional().describe("The maximum number of documents to return"), + jsonExportFormat: jsonExportFormat + .default("relaxed") + .describe( + [ + "The format to be used when exporting collection data as JSON with default being relaxed.", + "relaxed: A string format that emphasizes readability and interoperability at the expense of type preservation. That is, conversion from relaxed format to BSON can lose type information.", + "canonical: A string format that emphasizes type preservation at the expense of readability and interoperability. That is, conversion from canonical to BSON will generally preserve type information except in certain specific cases.", + ].join("\n") + ), + }; + public operationType: OperationType = "read"; + + protected async execute({ + database, + collection, + jsonExportFormat, + filter, + projection, + sort, + limit, + }: ToolArgs): Promise { + const provider = await this.ensureConnected(); + const findCursor = provider.find(database, collection, filter ?? {}, { + projection, + sort, + limit, + promoteValues: false, + bsonRegExp: true, + }); + const exportName = `${database}.${collection}.json`; + if (!this.exportsManager) { + throw new Error("Incorrect server configuration, export not possible!"); + } + + await this.exportsManager.createJSONExport({ input: findCursor, exportName, jsonExportFormat }); + const exportedResourceURI = this.exportsManager.exportNameToResourceURI(exportName); + const exportedResourcePath = this.exportsManager.exportFilePath( + this.exportsManager.exportsDirectoryPath(), + exportName + ); + const toolCallContent: CallToolResult["content"] = [ + // Not all the clients as of this commit understands how to + // parse a resource_link so we provide a text result for them to + // understand what to do with the result. + { + type: "text", + text: `Exported data for namespace ${database}.${collection} is available under resource URI - "${exportedResourceURI}".`, + }, + { + type: "resource_link", + name: exportName, + uri: exportedResourceURI, + description: "Resource URI for fetching exported data.", + mimeType: "application/json", + }, + ]; + + // This special case is to make it easier to work with exported data for + // stdio transport. + if (this.config.transport === "stdio") { + toolCallContent.push({ + type: "text", + text: `Optionally, the exported data can also be accessed under path - "${exportedResourcePath}"`, + }); + } + + return { + content: toolCallContent, + }; + } +} diff --git a/src/transports/base.ts b/src/transports/base.ts index 7052f1c4..62c79214 100644 --- a/src/transports/base.ts +++ b/src/transports/base.ts @@ -5,6 +5,7 @@ import { Session } from "../common/session.js"; import { Telemetry } from "../telemetry/telemetry.js"; import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js"; import { CompositeLogger, ConsoleLogger, DiskLogger, LoggerBase, McpLogger } from "../common/logger.js"; +import { SessionExportsManager } from "../common/sessionExportsManager.js"; export abstract class TransportRunnerBase { public logger: LoggerBase; @@ -46,11 +47,14 @@ export abstract class TransportRunnerBase { logger: new CompositeLogger(...loggers), }); + const exportsManager = new SessionExportsManager(session, userConfig); + const telemetry = Telemetry.create(session, userConfig); return new Server({ mcpServer, session, + exportsManager, telemetry, userConfig, }); From 3e8288eaf21c5a884085df59bbedf054477130de Mon Sep 17 00:00:00 2001 From: Himanshu Singh Date: Tue, 5 Aug 2025 15:39:15 +0200 Subject: [PATCH 03/50] chore: register export tool with server --- src/server.ts | 1 + src/tools/mongodb/tools.ts | 2 ++ 2 files changed, 3 insertions(+) diff --git a/src/server.ts b/src/server.ts index 7f645882..cf83607b 100644 --- a/src/server.ts +++ b/src/server.ts @@ -102,6 +102,7 @@ export class Server { } async close(): Promise { + this.exportsManager.close(); await this.telemetry.close(); await this.session.close(); await this.mcpServer.close(); diff --git a/src/tools/mongodb/tools.ts b/src/tools/mongodb/tools.ts index c74fdf29..00575ee0 100644 --- a/src/tools/mongodb/tools.ts +++ b/src/tools/mongodb/tools.ts @@ -18,6 +18,7 @@ import { DropCollectionTool } from "./delete/dropCollection.js"; import { ExplainTool } from "./metadata/explain.js"; import { CreateCollectionTool } from "./create/createCollection.js"; import { LogsTool } from "./metadata/logs.js"; +import { ExportTool } from "./read/export.js"; export const MongoDbTools = [ ConnectTool, @@ -40,4 +41,5 @@ export const MongoDbTools = [ ExplainTool, CreateCollectionTool, LogsTool, + ExportTool, ]; From b815a6f564989c1ce6c30d429cade82f943d0d8e Mon Sep 17 00:00:00 2001 From: Himanshu Singh Date: Tue, 5 Aug 2025 15:46:03 +0200 Subject: [PATCH 04/50] feat: adds exported-data resource --- src/resources/common/exported-data.ts | 114 ++++++++++++++++++++++++++ src/resources/resources.ts | 3 +- 2 files changed, 116 insertions(+), 1 deletion(-) create mode 100644 src/resources/common/exported-data.ts diff --git a/src/resources/common/exported-data.ts b/src/resources/common/exported-data.ts new file mode 100644 index 00000000..59b25e16 --- /dev/null +++ b/src/resources/common/exported-data.ts @@ -0,0 +1,114 @@ +import { + CompleteResourceTemplateCallback, + ListResourcesCallback, + ReadResourceTemplateCallback, + ResourceTemplate, +} from "@modelcontextprotocol/sdk/server/mcp.js"; +import { Server } from "../../server.js"; + +export class ExportedData { + private readonly name = "exported-data"; + private readonly description = "Data files exported in the current session."; + private readonly uri = "exported-data://{exportName}"; + + constructor(private readonly server: Server) { + this.server.session.on("export-available", (uri) => { + this.server.mcpServer.sendResourceListChanged(); + void this.server.mcpServer.server.sendResourceUpdated({ + uri, + }); + this.server.mcpServer.sendResourceListChanged(); + }); + this.server.session.on("export-expired", () => { + this.server.mcpServer.sendResourceListChanged(); + }); + } + + public register(): void { + this.server.mcpServer.registerResource( + this.name, + new ResourceTemplate(this.uri, { + /** + * A few clients have the capability of listing templated + * resources as well and this callback provides support for that + * */ + list: this.listResourcesCallback, + /** + * This is to provide auto completion when user starts typing in + * value for template variable, in our case, exportName */ + complete: { + exportName: this.autoCompleteExportName, + }, + }), + { description: this.description }, + this.readResourceCallback + ); + } + + private listResourcesCallback: ListResourcesCallback = () => { + const sessionId = this.server.session.sessionId; + if (!sessionId) { + // Note that we don't throw error here because this is a + // non-critical path and safe to return the most harmless value. + + // TODO: log warn here + return { resources: [] }; + } + + const sessionExports = this.server.exportsManager.listAvailableExports(); + return { + resources: sessionExports.map(({ name, uri }) => ({ + name: name, + description: this.exportNameToDescription(name), + uri: uri, + mimeType: "application/json", + })), + }; + }; + + private autoCompleteExportName: CompleteResourceTemplateCallback = (value) => { + const sessionId = this.server.session.sessionId; + if (!sessionId) { + // Note that we don't throw error here because this is a + // non-critical path and safe to return the most harmless value. + + // TODO: log warn here + return []; + } + + const sessionExports = this.server.exportsManager.listAvailableExports(); + return sessionExports.filter(({ name }) => name.startsWith(value)).map(({ name }) => name); + }; + + private readResourceCallback: ReadResourceTemplateCallback = async (uri, { exportName }) => { + const sessionId = this.server.session.sessionId; + if (!sessionId) { + throw new Error("Cannot retrieve exported data, session is not valid."); + } + + if (typeof exportName !== "string") { + throw new Error("Cannot retrieve exported data, exportName not provided."); + } + + return { + contents: [ + { + uri: this.server.exportsManager.exportNameToResourceURI(exportName), + text: await this.server.exportsManager.readExport(exportName), + mimeType: "application/json", + }, + ], + }; + }; + + private exportNameToDescription(exportName: string) { + const match = exportName.match(/^(.+)\.(\d+)\.json$/); + if (!match) return "Exported data for an unknown namespace."; + + const [, namespace, timestamp] = match; + if (!namespace || !timestamp) { + return "Exported data for an unknown namespace."; + } + return `Export from ${namespace} done on ${new Date(parseInt(timestamp)).toISOString()}`; + } +} diff --git a/src/resources/resources.ts b/src/resources/resources.ts index 40a17702..c27e4fc1 100644 --- a/src/resources/resources.ts +++ b/src/resources/resources.ts @@ -1,4 +1,5 @@ import { ConfigResource } from "./common/config.js"; import { DebugResource } from "./common/debug.js"; +import { ExportedData } from "./common/exported-data.js"; -export const Resources = [ConfigResource, DebugResource] as const; +export const Resources = [ConfigResource, DebugResource, ExportedData] as const; From 34c31fd4eb7c51324d84451001f1a46300b30c06 Mon Sep 17 00:00:00 2001 From: Himanshu Singh Date: Tue, 5 Aug 2025 17:58:55 +0200 Subject: [PATCH 05/50] chore: tests for export tool --- src/tools/mongodb/read/export.ts | 2 +- tests/integration/helpers.ts | 4 + .../tools/mongodb/read/export.test.ts | 302 ++++++++++++++++++ 3 files changed, 307 insertions(+), 1 deletion(-) create mode 100644 tests/integration/tools/mongodb/read/export.test.ts diff --git a/src/tools/mongodb/read/export.ts b/src/tools/mongodb/read/export.ts index ce721e10..b037d5e6 100644 --- a/src/tools/mongodb/read/export.ts +++ b/src/tools/mongodb/read/export.ts @@ -41,7 +41,7 @@ export class ExportTool extends MongoDBToolBase { promoteValues: false, bsonRegExp: true, }); - const exportName = `${database}.${collection}.json`; + const exportName = `${database}.${collection}.${Date.now()}.json`; if (!this.exportsManager) { throw new Error("Incorrect server configuration, export not possible!"); } diff --git a/tests/integration/helpers.ts b/tests/integration/helpers.ts index f5d8f05b..13cb90ea 100644 --- a/tests/integration/helpers.ts +++ b/tests/integration/helpers.ts @@ -10,6 +10,7 @@ import { config } from "../../src/common/config.js"; import { afterAll, afterEach, beforeAll, describe, expect, it, vi } from "vitest"; import { ConnectionManager } from "../../src/common/connectionManager.js"; import { CompositeLogger } from "../../src/common/logger.js"; +import { SessionExportsManager } from "../../src/common/sessionExportsManager.js"; interface ParameterInfo { name: string; @@ -75,8 +76,11 @@ export function setupIntegrationTest(getUserConfig: () => UserConfig): Integrati const telemetry = Telemetry.create(session, userConfig); + const exportsManager = new SessionExportsManager(session, userConfig); + mcpServer = new Server({ session, + exportsManager, userConfig, telemetry, mcpServer: new McpServer({ diff --git a/tests/integration/tools/mongodb/read/export.test.ts b/tests/integration/tools/mongodb/read/export.test.ts new file mode 100644 index 00000000..66ffc829 --- /dev/null +++ b/tests/integration/tools/mongodb/read/export.test.ts @@ -0,0 +1,302 @@ +import fs from "fs/promises"; +import { beforeEach, describe, expect, it } from "vitest"; +import { + databaseCollectionParameters, + validateThrowsForInvalidArguments, + validateToolMetadata, +} from "../../../helpers.js"; +import { describeWithMongoDB } from "../mongodbHelpers.js"; +import { CallToolResult } from "@modelcontextprotocol/sdk/types.js"; +import { Long } from "bson"; + +function contentWithTextResourceURI(content: CallToolResult["content"], namespace: string) { + return content.find((part) => { + return ( + part.type === "text" && + part.text.startsWith(`Exported data for namespace ${namespace} is available under resource URI -`) + ); + }); +} + +function contentWithResourceURILink(content: CallToolResult["content"], namespace: string) { + return content.find((part) => { + return part.type === "resource_link" && part.uri.startsWith(`exported-data://${namespace}`); + }); +} + +function contentWithExportPath(content: CallToolResult["content"]) { + return content.find((part) => { + return ( + part.type === "text" && + part.text.startsWith(`Optionally, the exported data can also be accessed under path -`) + ); + }); +} + +describeWithMongoDB("export tool", (integration) => { + validateToolMetadata( + integration, + "export", + "Export a collection data or query results in the specified json format.", + [ + ...databaseCollectionParameters, + + { + name: "filter", + description: "The query filter, matching the syntax of the query argument of db.collection.find()", + type: "object", + required: false, + }, + { + name: "jsonExportFormat", + description: [ + "The format to be used when exporting collection data as JSON with default being relaxed.", + "relaxed: A string format that emphasizes readability and interoperability at the expense of type preservation. That is, conversion from relaxed format to BSON can lose type information.", + "canonical: A string format that emphasizes type preservation at the expense of readability and interoperability. That is, conversion from canonical to BSON will generally preserve type information except in certain specific cases.", + ].join("\n"), + type: "string", + required: false, + }, + { + name: "limit", + description: "The maximum number of documents to return", + type: "number", + required: false, + }, + { + name: "projection", + description: "The projection, matching the syntax of the projection argument of db.collection.find()", + type: "object", + required: false, + }, + { + name: "sort", + description: + "A document, describing the sort order, matching the syntax of the sort argument of cursor.sort(). The keys of the object are the fields to sort on, while the values are the sort directions (1 for ascending, -1 for descending).", + type: "object", + required: false, + }, + ] + ); + + validateThrowsForInvalidArguments(integration, "export", [ + {}, + { database: 123, collection: "bar" }, + { database: "test", collection: [] }, + { database: "test", collection: "bar", filter: "{ $gt: { foo: 5 } }" }, + { database: "test", collection: "bar", projection: "name" }, + { database: "test", collection: "bar", limit: "10" }, + { database: "test", collection: "bar", sort: [], limit: 10 }, + ]); + + it("when provided with incorrect namespace, export should have empty data", async function () { + await integration.connectMcpClient(); + const response = await integration.mcpClient().callTool({ + name: "export", + arguments: { database: "non-existent", collection: "foos" }, + }); + + const content = response.content as CallToolResult["content"]; + const namespace = "non-existent.foos"; + expect(content).toHaveLength(3); + expect(contentWithTextResourceURI(content, namespace)).toBeDefined(); + expect(contentWithResourceURILink(content, namespace)).toBeDefined(); + + const localPathPart = contentWithExportPath(content); + expect(localPathPart).toBeDefined(); + + const [, localPath] = /"(.*)"/.exec(String(localPathPart?.text)) ?? []; + expect(localPath).toBeDefined(); + + expect(await fs.readFile(localPath as string, "utf8")).toEqual("[]"); + }); + + describe("with correct namespace", function () { + beforeEach(async () => { + const mongoClient = integration.mongoClient(); + await mongoClient + .db(integration.randomDbName()) + .collection("foo") + .insertMany([ + { name: "foo", longNumber: new Long(1234) }, + { name: "bar", bigInt: new Long(123412341234) }, + ]); + }); + + it("should export entire namespace when filter are empty", async function () { + await integration.connectMcpClient(); + const response = await integration.mcpClient().callTool({ + name: "export", + arguments: { database: integration.randomDbName(), collection: "foo" }, + }); + + const localPathPart = contentWithExportPath(response.content as CallToolResult["content"]); + expect(localPathPart).toBeDefined(); + const [, localPath] = /"(.*)"/.exec(String(localPathPart?.text)) ?? []; + expect(localPath).toBeDefined(); + + const exportedContent = JSON.parse(await fs.readFile(localPath as string, "utf8")) as Record< + string, + unknown + >[]; + expect(exportedContent).toHaveLength(2); + expect(exportedContent[0]?.name).toEqual("foo"); + expect(exportedContent[1]?.name).toEqual("bar"); + }); + + it("should export filter results namespace when there are filters", async function () { + await integration.connectMcpClient(); + const response = await integration.mcpClient().callTool({ + name: "export", + arguments: { database: integration.randomDbName(), collection: "foo", filter: { name: "foo" } }, + }); + + const localPathPart = contentWithExportPath(response.content as CallToolResult["content"]); + expect(localPathPart).toBeDefined(); + const [, localPath] = /"(.*)"/.exec(String(localPathPart?.text)) ?? []; + expect(localPath).toBeDefined(); + + const exportedContent = JSON.parse(await fs.readFile(localPath as string, "utf8")) as Record< + string, + unknown + >[]; + expect(exportedContent).toHaveLength(1); + expect(exportedContent[0]?.name).toEqual("foo"); + }); + + it("should export results limited to the provided limit", async function () { + await integration.connectMcpClient(); + const response = await integration.mcpClient().callTool({ + name: "export", + arguments: { database: integration.randomDbName(), collection: "foo", limit: 1 }, + }); + + const localPathPart = contentWithExportPath(response.content as CallToolResult["content"]); + expect(localPathPart).toBeDefined(); + const [, localPath] = /"(.*)"/.exec(String(localPathPart?.text)) ?? []; + expect(localPath).toBeDefined(); + + const exportedContent = JSON.parse(await fs.readFile(localPath as string, "utf8")) as Record< + string, + unknown + >[]; + expect(exportedContent).toHaveLength(1); + expect(exportedContent[0]?.name).toEqual("foo"); + }); + + it("should export results with sorted by the provided sort", async function () { + await integration.connectMcpClient(); + const response = await integration.mcpClient().callTool({ + name: "export", + arguments: { + database: integration.randomDbName(), + collection: "foo", + limit: 1, + sort: { longNumber: 1 }, + }, + }); + + const localPathPart = contentWithExportPath(response.content as CallToolResult["content"]); + expect(localPathPart).toBeDefined(); + const [, localPath] = /"(.*)"/.exec(String(localPathPart?.text)) ?? []; + expect(localPath).toBeDefined(); + + const exportedContent = JSON.parse(await fs.readFile(localPath as string, "utf8")) as Record< + string, + unknown + >[]; + expect(exportedContent).toHaveLength(1); + expect(exportedContent[0]?.name).toEqual("bar"); + }); + + it("should export results containing only projected fields", async function () { + await integration.connectMcpClient(); + const response = await integration.mcpClient().callTool({ + name: "export", + arguments: { + database: integration.randomDbName(), + collection: "foo", + limit: 1, + projection: { _id: 0, name: 1 }, + }, + }); + + const localPathPart = contentWithExportPath(response.content as CallToolResult["content"]); + expect(localPathPart).toBeDefined(); + const [, localPath] = /"(.*)"/.exec(String(localPathPart?.text)) ?? []; + expect(localPath).toBeDefined(); + + const exportedContent = JSON.parse(await fs.readFile(localPath as string, "utf8")) as Record< + string, + unknown + >[]; + expect(exportedContent).toEqual([ + { + name: "foo", + }, + ]); + }); + + it("should export relaxed json when provided jsonExportFormat is relaxed", async function () { + await integration.connectMcpClient(); + const response = await integration.mcpClient().callTool({ + name: "export", + arguments: { + database: integration.randomDbName(), + collection: "foo", + limit: 1, + projection: { _id: 0 }, + jsonExportFormat: "relaxed", + }, + }); + + const localPathPart = contentWithExportPath(response.content as CallToolResult["content"]); + expect(localPathPart).toBeDefined(); + const [, localPath] = /"(.*)"/.exec(String(localPathPart?.text)) ?? []; + expect(localPath).toBeDefined(); + + const exportedContent = JSON.parse(await fs.readFile(localPath as string, "utf8")) as Record< + string, + unknown + >[]; + expect(exportedContent).toEqual([ + { + name: "foo", + longNumber: 1234, + }, + ]); + }); + + it("should export canonical json when provided jsonExportFormat is canonical", async function () { + await integration.connectMcpClient(); + const response = await integration.mcpClient().callTool({ + name: "export", + arguments: { + database: integration.randomDbName(), + collection: "foo", + limit: 1, + projection: { _id: 0 }, + jsonExportFormat: "canonical", + }, + }); + + const localPathPart = contentWithExportPath(response.content as CallToolResult["content"]); + expect(localPathPart).toBeDefined(); + const [, localPath] = /"(.*)"/.exec(String(localPathPart?.text)) ?? []; + expect(localPath).toBeDefined(); + + const exportedContent = JSON.parse(await fs.readFile(localPath as string, "utf8")) as Record< + string, + unknown + >[]; + expect(exportedContent).toEqual([ + { + name: "foo", + longNumber: { + $numberLong: "1234", + }, + }, + ]); + }); + }); +}); From f3606c69afa097db9d78e64d5ecb11587fd9c6eb Mon Sep 17 00:00:00 2001 From: Himanshu Singh Date: Tue, 5 Aug 2025 18:58:43 +0200 Subject: [PATCH 06/50] chore: tests for exported-data resource --- src/common/logger.ts | 3 + src/common/sessionExportsManager.ts | 25 ++- src/resources/common/exported-data.ts | 114 ------------- src/resources/common/exportedData.ts | 155 ++++++++++++++++++ src/resources/resources.ts | 2 +- src/server.ts | 2 +- tests/integration/helpers.ts | 4 + .../resources/exportedData.test.ts | 102 ++++++++++++ .../common/sessionExportsManager.test.ts | 12 +- 9 files changed, 292 insertions(+), 127 deletions(-) delete mode 100644 src/resources/common/exported-data.ts create mode 100644 src/resources/common/exportedData.ts create mode 100644 tests/integration/resources/exportedData.test.ts rename tests/{integration => unit}/common/sessionExportsManager.test.ts (98%) diff --git a/src/common/logger.ts b/src/common/logger.ts index b33a08fc..9d1bd43d 100644 --- a/src/common/logger.ts +++ b/src/common/logger.ts @@ -53,6 +53,9 @@ export const LogId = { exportCleanupError: mongoLogId(1_007_001), exportCreationError: mongoLogId(1_007_002), exportReadError: mongoLogId(1_007_003), + exportCloseError: mongoLogId(1_007_004), + exportedDataListError: mongoLogId(1_007_005), + exportedDataAutoCompleteError: mongoLogId(1_007_006), } as const; interface LogPayload { diff --git a/src/common/sessionExportsManager.ts b/src/common/sessionExportsManager.ts index 754443ca..abdf50e0 100644 --- a/src/common/sessionExportsManager.ts +++ b/src/common/sessionExportsManager.ts @@ -41,8 +41,18 @@ export class SessionExportsManager { ); } - public close() { - clearInterval(this.exportsCleanupInterval); + public async close() { + try { + clearInterval(this.exportsCleanupInterval); + const exportsDirectory = this.exportsDirectoryPath(); + await fs.rm(exportsDirectory, { force: true, recursive: true }); + } catch (error) { + logger.error( + LogId.exportCloseError, + "Error while closing SessionExportManager", + error instanceof Error ? error.message : String(error) + ); + } } public exportNameToResourceURI(nameWithExtension: string): string { @@ -209,8 +219,15 @@ export class SessionExportsManager { /** * Small utility to centrally determine if an export is expired or not */ private async isExportFileExpired(exportFilePath: string): Promise { - const stats = await fs.stat(exportFilePath); - return this.isExportExpired(stats.birthtimeMs); + try { + const stats = await fs.stat(exportFilePath); + return this.isExportExpired(stats.birthtimeMs); + } catch (error) { + if ((error as NodeJS.ErrnoException).code === "ENOENT") { + throw new Error("Requested export does not exist!"); + } + throw error; + } } private isExportExpired(createdAt: number) { diff --git a/src/resources/common/exported-data.ts b/src/resources/common/exported-data.ts deleted file mode 100644 index 59b25e16..00000000 --- a/src/resources/common/exported-data.ts +++ /dev/null @@ -1,114 +0,0 @@ -import { - CompleteResourceTemplateCallback, - ListResourcesCallback, - ReadResourceTemplateCallback, - ResourceTemplate, -} from "@modelcontextprotocol/sdk/server/mcp.js"; -import { Server } from "../../server.js"; - -export class ExportedData { - private readonly name = "exported-data"; - private readonly description = "Data files exported in the current session."; - private readonly uri = "exported-data://{exportName}"; - - constructor(private readonly server: Server) { - this.server.session.on("export-available", (uri) => { - this.server.mcpServer.sendResourceListChanged(); - void this.server.mcpServer.server.sendResourceUpdated({ - uri, - }); - this.server.mcpServer.sendResourceListChanged(); - }); - this.server.session.on("export-expired", () => { - this.server.mcpServer.sendResourceListChanged(); - }); - } - - public register(): void { - this.server.mcpServer.registerResource( - this.name, - new ResourceTemplate(this.uri, { - /** - * A few clients have the capability of listing templated - * resources as well and this callback provides support for that - * */ - list: this.listResourcesCallback, - /** - * This is to provide auto completion when user starts typing in - * value for template variable, in our case, exportName */ - complete: { - exportName: this.autoCompleteExportName, - }, - }), - { description: this.description }, - this.readResourceCallback - ); - } - - private listResourcesCallback: ListResourcesCallback = () => { - const sessionId = this.server.session.sessionId; - if (!sessionId) { - // Note that we don't throw error here because this is a - // non-critical path and safe to return the most harmless value. - - // TODO: log warn here - return { resources: [] }; - } - - const sessionExports = this.server.exportsManager.listAvailableExports(); - return { - resources: sessionExports.map(({ name, uri }) => ({ - name: name, - description: this.exportNameToDescription(name), - uri: uri, - mimeType: "application/json", - })), - }; - }; - - private autoCompleteExportName: CompleteResourceTemplateCallback = (value) => { - const sessionId = this.server.session.sessionId; - if (!sessionId) { - // Note that we don't throw error here because this is a - // non-critical path and safe to return the most harmless value. - - // TODO: log warn here - return []; - } - - const sessionExports = this.server.exportsManager.listAvailableExports(); - return sessionExports.filter(({ name }) => name.startsWith(value)).map(({ name }) => name); - }; - - private readResourceCallback: ReadResourceTemplateCallback = async (uri, { exportName }) => { - const sessionId = this.server.session.sessionId; - if (!sessionId) { - throw new Error("Cannot retrieve exported data, session is not valid."); - } - - if (typeof exportName !== "string") { - throw new Error("Cannot retrieve exported data, exportName not provided."); - } - - return { - contents: [ - { - uri: this.server.exportsManager.exportNameToResourceURI(exportName), - text: await this.server.exportsManager.readExport(exportName), - mimeType: "application/json", - }, - ], - }; - }; - - private exportNameToDescription(exportName: string) { - const match = exportName.match(/^(.+)\.(\d+)\.json$/); - if (!match) return "Exported data for an unknown namespace."; - - const [, namespace, timestamp] = match; - if (!namespace || !timestamp) { - return "Exported data for an unknown namespace."; - } - return `Export from ${namespace} done on ${new Date(parseInt(timestamp)).toISOString()}`; - } -} diff --git a/src/resources/common/exportedData.ts b/src/resources/common/exportedData.ts new file mode 100644 index 00000000..aa1f995f --- /dev/null +++ b/src/resources/common/exportedData.ts @@ -0,0 +1,155 @@ +import { + CompleteResourceTemplateCallback, + ListResourcesCallback, + ReadResourceTemplateCallback, + ResourceTemplate, +} from "@modelcontextprotocol/sdk/server/mcp.js"; +import { Server } from "../../server.js"; +import logger, { LogId } from "../../common/logger.js"; + +export class ExportedData { + private readonly name = "exported-data"; + private readonly description = "Data files exported in the current session."; + private readonly uri = "exported-data://{exportName}"; + + constructor(private readonly server: Server) { + this.server.session.on("export-available", (uri) => { + this.server.mcpServer.sendResourceListChanged(); + void this.server.mcpServer.server.sendResourceUpdated({ + uri, + }); + this.server.mcpServer.sendResourceListChanged(); + }); + this.server.session.on("export-expired", () => { + this.server.mcpServer.sendResourceListChanged(); + }); + } + + public register(): void { + this.server.mcpServer.registerResource( + this.name, + new ResourceTemplate(this.uri, { + /** + * A few clients have the capability of listing templated + * resources as well and this callback provides support for that + * */ + list: this.listResourcesCallback, + /** + * This is to provide auto completion when user starts typing in + * value for template variable, in our case, exportName */ + complete: { + exportName: this.autoCompleteExportName, + }, + }), + { description: this.description }, + this.readResourceCallback + ); + } + + private listResourcesCallback: ListResourcesCallback = () => { + try { + const sessionId = this.server.session.sessionId; + if (!sessionId) { + // Note that we don't throw error here because this is a + // non-critical path and safe to return the most harmless value. + + // TODO: log warn here + return { resources: [] }; + } + + const sessionExports = this.server.exportsManager.listAvailableExports(); + return { + resources: sessionExports.map(({ name, uri }) => ({ + name: name, + description: this.exportNameToDescription(name), + uri: uri, + mimeType: "application/json", + })), + }; + } catch (error) { + logger.error( + LogId.exportedDataListError, + "Error when listing exported data resources", + error instanceof Error ? error.message : String(error) + ); + return { + resources: [], + }; + } + }; + + private autoCompleteExportName: CompleteResourceTemplateCallback = (value) => { + try { + const sessionId = this.server.session.sessionId; + if (!sessionId) { + // Note that we don't throw error here because this is a + // non-critical path and safe to return the most harmless value. + + // TODO: log warn here + return []; + } + + const sessionExports = this.server.exportsManager.listAvailableExports(); + return sessionExports.filter(({ name }) => name.startsWith(value)).map(({ name }) => name); + } catch (error) { + logger.error( + LogId.exportedDataAutoCompleteError, + "Error when autocompleting exported data", + error instanceof Error ? error.message : String(error) + ); + return []; + } + }; + + private readResourceCallback: ReadResourceTemplateCallback = async (uri, { exportName }) => { + try { + const sessionId = this.server.session.sessionId; + if (!sessionId) { + throw new Error("Cannot retrieve exported data, session is not valid."); + } + + if (typeof exportName !== "string") { + throw new Error("Cannot retrieve exported data, exportName not provided."); + } + + return { + contents: [ + { + uri: this.server.exportsManager.exportNameToResourceURI(exportName), + text: await this.server.exportsManager.readExport(exportName), + mimeType: "application/json", + }, + ], + }; + } catch (error) { + return { + contents: [ + { + uri: + typeof exportName === "string" + ? this.server.exportsManager.exportNameToResourceURI(exportName) + : this.uri, + text: `Error reading from ${this.uri}: ${error instanceof Error ? error.message : String(error)}`, + }, + ], + isError: true, + }; + } + }; + + private exportNameToDescription(exportName: string) { + const match = exportName.match(/^(.+)\.(\d+)\.json$/); + if (!match) return "Exported data for an unknown namespace."; + + const [, namespace, timestamp] = match; + if (!namespace) { + return "Exported data for an unknown namespace."; + } + + if (!timestamp) { + return `Export from ${namespace}.`; + } + + return `Export from ${namespace} done on ${new Date(parseInt(timestamp)).toLocaleString()}`; + } +} diff --git a/src/resources/resources.ts b/src/resources/resources.ts index c27e4fc1..24a129ab 100644 --- a/src/resources/resources.ts +++ b/src/resources/resources.ts @@ -1,5 +1,5 @@ import { ConfigResource } from "./common/config.js"; import { DebugResource } from "./common/debug.js"; -import { ExportedData } from "./common/exported-data.js"; +import { ExportedData } from "./common/exportedData.js"; export const Resources = [ConfigResource, DebugResource, ExportedData] as const; diff --git a/src/server.ts b/src/server.ts index cf83607b..5c6464e2 100644 --- a/src/server.ts +++ b/src/server.ts @@ -102,7 +102,7 @@ export class Server { } async close(): Promise { - this.exportsManager.close(); + await this.exportsManager.close(); await this.telemetry.close(); await this.session.close(); await this.mcpServer.close(); diff --git a/tests/integration/helpers.ts b/tests/integration/helpers.ts index 13cb90ea..b965b779 100644 --- a/tests/integration/helpers.ts +++ b/tests/integration/helpers.ts @@ -273,3 +273,7 @@ function validateToolAnnotations(tool: ToolInfo, name: string, description: stri expect(tool.annotations.destructiveHint).toBe(false); } } + +export function timeout(ms: number) { + return new Promise((resolve) => setTimeout(resolve, ms)); +} diff --git a/tests/integration/resources/exportedData.test.ts b/tests/integration/resources/exportedData.test.ts new file mode 100644 index 00000000..3f57473c --- /dev/null +++ b/tests/integration/resources/exportedData.test.ts @@ -0,0 +1,102 @@ +import { Long } from "bson"; +import { describe, expect, it, beforeEach } from "vitest"; +import { describeWithMongoDB } from "../tools/mongodb/mongodbHelpers.js"; +import { defaultTestConfig, timeout } from "../helpers.js"; +import { CallToolResult } from "@modelcontextprotocol/sdk/types.js"; + +describeWithMongoDB( + "exported-data resource", + (integration) => { + beforeEach(async () => { + const mongoClient = integration.mongoClient(); + await mongoClient + .db(integration.randomDbName()) + .collection("foo") + .insertMany([ + { name: "foo", longNumber: new Long(1234) }, + { name: "bar", bigInt: new Long(123412341234) }, + ]); + }); + + it("should be able to list resource template", async () => { + await integration.connectMcpClient(); + const response = await integration.mcpClient().listResourceTemplates(); + expect(response.resourceTemplates).toEqual([ + { + name: "exported-data", + uriTemplate: "exported-data://{exportName}", + description: "Data files exported in the current session.", + }, + ]); + }); + + describe("when requesting non-existent resource", () => { + it("should return an error", async () => { + await integration.connectMcpClient(); + const response = await integration.mcpClient().readResource({ + uri: "exported-data://foo.bar.json", + }); + expect(response.isError).toEqual(true); + expect(response.contents[0]?.uri).toEqual("exported-data://foo.bar.json"); + expect(response.contents[0]?.text).toEqual( + "Error reading from exported-data://{exportName}: Requested export does not exist!" + ); + }); + }); + + describe("when requesting an expired resource", () => { + it("should return an error", async () => { + await integration.connectMcpClient(); + const exportResponse = await integration.mcpClient().callTool({ + name: "export", + arguments: { database: integration.randomDbName(), collection: "foo" }, + }); + + const exportedResourceURI = (exportResponse as CallToolResult).content.find( + (part) => part.type === "resource_link" + )?.uri; + expect(exportedResourceURI).toBeDefined(); + + // wait for export expired + await timeout(200); + const response = await integration.mcpClient().readResource({ + uri: exportedResourceURI as string, + }); + expect(response.isError).toEqual(true); + expect(response.contents[0]?.uri).toEqual(exportedResourceURI); + expect(response.contents[0]?.text).toEqual( + "Error reading from exported-data://{exportName}: Export has expired" + ); + }); + }); + + describe("after requesting a fresh export", () => { + it("should be able to read the resource", async () => { + await integration.connectMcpClient(); + const exportResponse = await integration.mcpClient().callTool({ + name: "export", + arguments: { database: integration.randomDbName(), collection: "foo" }, + }); + + const exportedResourceURI = (exportResponse as CallToolResult).content.find( + (part) => part.type === "resource_link" + )?.uri; + expect(exportedResourceURI).toBeDefined(); + + const response = await integration.mcpClient().readResource({ + uri: exportedResourceURI as string, + }); + expect(response.isError).toBeFalsy(); + expect(response.contents[0]?.mimeType).toEqual("application/json"); + expect(response.contents[0]?.text).toContain("foo"); + }); + }); + }, + () => { + return { + ...defaultTestConfig, + exportTimeoutMs: 200, + exportCleanupIntervalMs: 100, + }; + } +); diff --git a/tests/integration/common/sessionExportsManager.test.ts b/tests/unit/common/sessionExportsManager.test.ts similarity index 98% rename from tests/integration/common/sessionExportsManager.test.ts rename to tests/unit/common/sessionExportsManager.test.ts index 7d6ed023..cbaeee0e 100644 --- a/tests/integration/common/sessionExportsManager.test.ts +++ b/tests/unit/common/sessionExportsManager.test.ts @@ -1,12 +1,14 @@ import path from "path"; import fs from "fs/promises"; +import { Readable } from "stream"; +import { FindCursor, Long } from "mongodb"; import { beforeEach, describe, expect, it, vi } from "vitest"; import { SessionExportsManager, SessionExportsManagerConfig } from "../../../src/common/sessionExportsManager.js"; + import { config } from "../../../src/common/config.js"; import { Session } from "../../../src/common/session.js"; import { ROOT_DIR } from "../../accuracy/sdk/constants.js"; -import { FindCursor, Long } from "mongodb"; -import { Readable } from "stream"; +import { timeout } from "../../integration/helpers.js"; const dummySessionId = "1FOO"; const dummyExportsPath = path.join(ROOT_DIR, "tests", "tmp", "exports"); @@ -62,16 +64,12 @@ async function fileExists(filePath: string) { } } -function timeout(ms: number) { - return new Promise((resolve) => setTimeout(resolve, ms)); -} - describe("SessionExportsManager integration test", () => { let session: Session; let manager: SessionExportsManager; beforeEach(async () => { - manager?.close(); + await manager?.close(); await fs.rm(exportsManagerConfig.exportPath, { recursive: true, force: true }); await fs.mkdir(exportsManagerConfig.exportPath, { recursive: true }); session = new Session({ apiBaseUrl: "" }); From 75ab89d14cb788ee6228b4fb1d22f5ad855a1f03 Mon Sep 17 00:00:00 2001 From: Himanshu Singh Date: Tue, 5 Aug 2025 19:20:43 +0200 Subject: [PATCH 07/50] chore: server should advertise listChanged capability --- src/server.ts | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/server.ts b/src/server.ts index 5c6464e2..aa1dd822 100644 --- a/src/server.ts +++ b/src/server.ts @@ -47,7 +47,7 @@ export class Server { this.registerResources(); await this.validateConfig(); - this.mcpServer.server.registerCapabilities({ logging: {}, resources: { subscribe: true, listChanged: true } }); + this.mcpServer.server.registerCapabilities({ logging: {}, resources: { listChanged: true } }); // TODO: Eventually we might want to make tools reactive too instead of relying on custom logic. this.registerTools(); From a38a45f55746c0d3524c1db43d9a7394cbdabed8 Mon Sep 17 00:00:00 2001 From: Himanshu Singh Date: Tue, 5 Aug 2025 19:20:59 +0200 Subject: [PATCH 08/50] chore: test for autocomplete API --- .../resources/exportedData.test.ts | 37 ++++++++++++++++--- 1 file changed, 31 insertions(+), 6 deletions(-) diff --git a/tests/integration/resources/exportedData.test.ts b/tests/integration/resources/exportedData.test.ts index 3f57473c..fca7592c 100644 --- a/tests/integration/resources/exportedData.test.ts +++ b/tests/integration/resources/exportedData.test.ts @@ -10,8 +10,8 @@ describeWithMongoDB( beforeEach(async () => { const mongoClient = integration.mongoClient(); await mongoClient - .db(integration.randomDbName()) - .collection("foo") + .db("db") + .collection("coll") .insertMany([ { name: "foo", longNumber: new Long(1234) }, { name: "bar", bigInt: new Long(123412341234) }, @@ -34,10 +34,10 @@ describeWithMongoDB( it("should return an error", async () => { await integration.connectMcpClient(); const response = await integration.mcpClient().readResource({ - uri: "exported-data://foo.bar.json", + uri: "exported-data://db.coll.json", }); expect(response.isError).toEqual(true); - expect(response.contents[0]?.uri).toEqual("exported-data://foo.bar.json"); + expect(response.contents[0]?.uri).toEqual("exported-data://db.coll.json"); expect(response.contents[0]?.text).toEqual( "Error reading from exported-data://{exportName}: Requested export does not exist!" ); @@ -49,7 +49,7 @@ describeWithMongoDB( await integration.connectMcpClient(); const exportResponse = await integration.mcpClient().callTool({ name: "export", - arguments: { database: integration.randomDbName(), collection: "foo" }, + arguments: { database: "db", collection: "coll" }, }); const exportedResourceURI = (exportResponse as CallToolResult).content.find( @@ -75,7 +75,7 @@ describeWithMongoDB( await integration.connectMcpClient(); const exportResponse = await integration.mcpClient().callTool({ name: "export", - arguments: { database: integration.randomDbName(), collection: "foo" }, + arguments: { database: "db", collection: "coll" }, }); const exportedResourceURI = (exportResponse as CallToolResult).content.find( @@ -90,6 +90,31 @@ describeWithMongoDB( expect(response.contents[0]?.mimeType).toEqual("application/json"); expect(response.contents[0]?.text).toContain("foo"); }); + + it("should be able to autocomplete the resource", async () => { + await integration.connectMcpClient(); + const exportResponse = await integration.mcpClient().callTool({ + name: "export", + arguments: { database: "big", collection: "coll" }, + }); + + const exportedResourceURI = (exportResponse as CallToolResult).content.find( + (part) => part.type === "resource_link" + )?.uri; + expect(exportedResourceURI).toBeDefined(); + + const completeResponse = await integration.mcpClient().complete({ + ref: { + type: "ref/resource", + uri: "exported-data://{exportName}", + }, + argument: { + name: "exportName", + value: "b", + }, + }); + expect(completeResponse.completion.total).toEqual(1); + }); }); }, () => { From 7a83e8f4982638ea217fd57a549d6f470e073625 Mon Sep 17 00:00:00 2001 From: Himanshu Singh Date: Tue, 5 Aug 2025 19:23:09 +0200 Subject: [PATCH 09/50] chore: fix broken test --- tests/integration/transports/stdio.test.ts | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/integration/transports/stdio.test.ts b/tests/integration/transports/stdio.test.ts index 6b08e4e6..aaa61d63 100644 --- a/tests/integration/transports/stdio.test.ts +++ b/tests/integration/transports/stdio.test.ts @@ -32,7 +32,7 @@ describeWithMongoDB("StdioRunner", (integration) => { const response = await client.listTools(); expect(response).toBeDefined(); expect(response.tools).toBeDefined(); - expect(response.tools).toHaveLength(20); + expect(response.tools).toHaveLength(21); const sortedTools = response.tools.sort((a, b) => a.name.localeCompare(b.name)); expect(sortedTools[0]?.name).toBe("aggregate"); From 2aacd8514d2210effc99a24638ea6ea598b0db14 Mon Sep 17 00:00:00 2001 From: Himanshu Singh Date: Wed, 6 Aug 2025 10:25:09 +0200 Subject: [PATCH 10/50] chore: handle pipeline failures explicitly --- src/common/logger.ts | 9 +- src/common/sessionExportsManager.ts | 46 ++-- .../unit/common/sessionExportsManager.test.ts | 204 +++++++++++++----- 3 files changed, 184 insertions(+), 75 deletions(-) diff --git a/src/common/logger.ts b/src/common/logger.ts index 9d1bd43d..98a7a829 100644 --- a/src/common/logger.ts +++ b/src/common/logger.ts @@ -52,10 +52,11 @@ export const LogId = { exportCleanupError: mongoLogId(1_007_001), exportCreationError: mongoLogId(1_007_002), - exportReadError: mongoLogId(1_007_003), - exportCloseError: mongoLogId(1_007_004), - exportedDataListError: mongoLogId(1_007_005), - exportedDataAutoCompleteError: mongoLogId(1_007_006), + exportCreationCleanupError: mongoLogId(1_007_003), + exportReadError: mongoLogId(1_007_004), + exportCloseError: mongoLogId(1_007_005), + exportedDataListError: mongoLogId(1_007_006), + exportedDataAutoCompleteError: mongoLogId(1_007_007), } as const; interface LogPayload { diff --git a/src/common/sessionExportsManager.ts b/src/common/sessionExportsManager.ts index abdf50e0..bc1bf53b 100644 --- a/src/common/sessionExportsManager.ts +++ b/src/common/sessionExportsManager.ts @@ -118,28 +118,44 @@ export class SessionExportsManager { jsonExportFormat: JSONExportFormat; }): Promise { try { + const exportNameWithExtension = this.withExtension(exportName, "json"); + const inputStream = input.stream(); + const ejsonDocStream = this.docToEJSONStream(this.getEJSONOptionsForFormat(jsonExportFormat)); await this.withExportsLock(async (exportsDirectoryPath) => { - const exportNameWithExtension = this.withExtension(exportName, "json"); const exportFilePath = path.join(exportsDirectoryPath, exportNameWithExtension); const outputStream = createWriteStream(exportFilePath); outputStream.write("["); + let pipeSuccessful = false; try { - const inputStream = input.stream(); - const ejsonOptions = this.getEJSONOptionsForFormat(jsonExportFormat); - await pipeline([inputStream, this.docToEJSONStream(ejsonOptions), outputStream]); + await pipeline([inputStream, ejsonDocStream, outputStream]); + pipeSuccessful = true; + } catch (pipelineError) { + // If the pipeline errors out then we might end up with + // partial and incorrect export so we remove it entirely. + await fs.unlink(exportFilePath).catch((error) => { + if ((error as NodeJS.ErrnoException).code !== "ENOENT") { + logger.error( + LogId.exportCreationCleanupError, + "Error when removing partial export", + error instanceof Error ? error.message : String(error) + ); + } + }); + throw pipelineError; } finally { - outputStream.write("]\n"); - const resourceURI = this.exportNameToResourceURI(exportNameWithExtension); - this.mutableExports = [ - ...this.mutableExports, - { - createdAt: (await fs.stat(exportFilePath)).birthtimeMs, - name: exportNameWithExtension, - uri: resourceURI, - }, - ]; - this.session.emit("export-available", resourceURI); void input.close(); + if (pipeSuccessful) { + const resourceURI = this.exportNameToResourceURI(exportNameWithExtension); + this.mutableExports = [ + ...this.mutableExports, + { + createdAt: (await fs.stat(exportFilePath)).birthtimeMs, + name: exportNameWithExtension, + uri: resourceURI, + }, + ]; + this.session.emit("export-available", resourceURI); + } } }); } catch (error) { diff --git a/tests/unit/common/sessionExportsManager.test.ts b/tests/unit/common/sessionExportsManager.test.ts index cbaeee0e..a869d166 100644 --- a/tests/unit/common/sessionExportsManager.test.ts +++ b/tests/unit/common/sessionExportsManager.test.ts @@ -1,6 +1,6 @@ import path from "path"; import fs from "fs/promises"; -import { Readable } from "stream"; +import { Readable, Transform } from "stream"; import { FindCursor, Long } from "mongodb"; import { beforeEach, describe, expect, it, vi } from "vitest"; import { SessionExportsManager, SessionExportsManagerConfig } from "../../../src/common/sessionExportsManager.js"; @@ -9,6 +9,7 @@ import { config } from "../../../src/common/config.js"; import { Session } from "../../../src/common/session.js"; import { ROOT_DIR } from "../../accuracy/sdk/constants.js"; import { timeout } from "../../integration/helpers.js"; +import { EJSON, EJSONOptions } from "bson"; const dummySessionId = "1FOO"; const dummyExportsPath = path.join(ROOT_DIR, "tests", "tmp", "exports"); @@ -18,16 +19,20 @@ const exportsManagerConfig: SessionExportsManagerConfig = { exportTimeoutMs: config.exportTimeoutMs, exportCleanupIntervalMs: config.exportCleanupIntervalMs, } as const; -const dummyExportName = "foo.bar.json"; -const dummyExportPath = path.join(dummySessionExportPath, dummyExportName); +function getDummyExportName(timestamp: number) { + return `foo.bar.${timestamp}.json`; +} +function getDummyExportPath(timestamp: number) { + return path.join(dummySessionExportPath, getDummyExportName(timestamp)); +} -async function createDummyExport() { +async function createDummyExport(timestamp: number) { const content = "[]"; await fs.mkdir(dummySessionExportPath, { recursive: true }); - await fs.writeFile(dummyExportPath, content); + await fs.writeFile(getDummyExportPath(timestamp), content); return { - name: dummyExportName, - path: dummyExportPath, + name: getDummyExportName(timestamp), + path: getDummyExportPath(timestamp), content, }; } @@ -76,54 +81,54 @@ describe("SessionExportsManager integration test", () => { manager = new SessionExportsManager(session, exportsManagerConfig); }); - describe("#exportNameToResourceURI", function () { - it("should throw when export name has no extension", function () { + describe("#exportNameToResourceURI", () => { + it("should throw when export name has no extension", () => { expect(() => manager.exportNameToResourceURI("name")).toThrow(); }); - it("should return a resource URI", function () { + it("should return a resource URI", () => { expect(manager.exportNameToResourceURI("name.json")).toEqual("exported-data://name.json"); }); }); - describe("#exportsDirectoryPath", function () { - it("should throw when session is not initialized", function () { + describe("#exportsDirectoryPath", () => { + it("should throw when session is not initialized", () => { expect(() => manager.exportsDirectoryPath()).toThrow(); }); - it("should return a session path when session is initialized", function () { + it("should return a session path when session is initialized", () => { session.sessionId = dummySessionId; manager = new SessionExportsManager(session, exportsManagerConfig); expect(manager.exportsDirectoryPath()).toEqual(path.join(exportsManagerConfig.exportPath, dummySessionId)); }); }); - describe("#exportFilePath", function () { - it("should throw when export name has no extension", function () { + describe("#exportFilePath", () => { + it("should throw when export name has no extension", () => { expect(() => manager.exportFilePath(dummySessionExportPath, "name")).toThrow(); }); - it("should return path to provided export file", function () { + it("should return path to provided export file", () => { expect(manager.exportFilePath(dummySessionExportPath, "mflix.movies.json")).toEqual( path.join(dummySessionExportPath, "mflix.movies.json") ); }); }); - describe("#readExport", function () { - it("should throw when export name has no extension", async function () { + describe("#readExport", () => { + it("should throw when export name has no extension", async () => { await expect(() => manager.readExport("name")).rejects.toThrow(); }); - it("should return the resource content", async function () { - const { name, content } = await createDummyExport(); + it("should return the resource content", async () => { + const { name, content } = await createDummyExport(Date.now()); session.sessionId = dummySessionId; manager = new SessionExportsManager(session, exportsManagerConfig); expect(await manager.readExport(name)).toEqual(content); }); }); - describe("#createJSONExport", function () { + describe("#createJSONExport", () => { let inputCursor: FindCursor; beforeEach(() => { void inputCursor?.close(); @@ -139,12 +144,47 @@ describe("SessionExportsManager integration test", () => { ]); }); - it.each([ - { cond: "when exportName does not contain extension", exportName: "foo.bar" }, - { cond: "when exportName contains extension", exportName: "foo.bar.json" }, - ])( - "$cond, should export relaxed json, update available exports and emit export-available event", - async function ({ exportName }) { + describe("when cursor is empty", () => { + it("should create an empty export", async () => { + inputCursor = createDummyFindCursor([]); + + const emitSpy = vi.spyOn(session, "emit"); + session.sessionId = dummySessionId; + manager = new SessionExportsManager(session, exportsManagerConfig); + const timestamp = Date.now(); + await manager.createJSONExport({ + input: inputCursor, + exportName: getDummyExportName(timestamp), + jsonExportFormat: "relaxed", + }); + + // Updates available export + const availableExports = manager.listAvailableExports(); + expect(availableExports).toHaveLength(1); + expect(availableExports).toContainEqual( + expect.objectContaining({ + name: getDummyExportName(timestamp), + uri: `exported-data://${getDummyExportName(timestamp)}`, + }) + ); + + // Emit event + expect(emitSpy).toHaveBeenCalledWith( + "export-available", + `exported-data://${getDummyExportName(timestamp)}` + ); + + // Exports relaxed json + const jsonData = JSON.parse(await manager.readExport(getDummyExportName(timestamp))) as unknown[]; + expect(jsonData).toEqual([]); + }); + }); + + describe.each([ + { cond: "when exportName does not contain extension", exportName: `foo.bar.${Date.now()}` }, + { cond: "when exportName contains extension", exportName: `foo.bar.${Date.now()}.json` }, + ])("$cond", ({ exportName }) => { + it("should export relaxed json, update available exports and emit export-available event", async () => { const emitSpy = vi.spyOn(session, "emit"); session.sessionId = dummySessionId; manager = new SessionExportsManager(session, exportsManagerConfig); @@ -154,32 +194,32 @@ describe("SessionExportsManager integration test", () => { jsonExportFormat: "relaxed", }); + const expectedExportName = exportName.endsWith(".json") ? exportName : `${exportName}.json`; // Updates available export const availableExports = manager.listAvailableExports(); expect(availableExports).toHaveLength(1); expect(availableExports).toContainEqual( expect.objectContaining({ - name: "foo.bar.json", - uri: "exported-data://foo.bar.json", + name: expectedExportName, + uri: `exported-data://${expectedExportName}`, }) ); // Emit event - expect(emitSpy).toHaveBeenCalledWith("export-available", "exported-data://foo.bar.json"); + expect(emitSpy).toHaveBeenCalledWith("export-available", `exported-data://${expectedExportName}`); // Exports relaxed json - const jsonData = JSON.parse(await manager.readExport("foo.bar.json")) as unknown[]; + const jsonData = JSON.parse(await manager.readExport(expectedExportName)) as unknown[]; expect(jsonData).toContainEqual(expect.objectContaining({ name: "foo", longNumber: 12 })); expect(jsonData).toContainEqual(expect.objectContaining({ name: "bar", longNumber: 123456 })); - } - ); - - it.each([ - { cond: "when exportName does not contain extension", exportName: "foo.bar" }, - { cond: "when exportName contains extension", exportName: "foo.bar.json" }, - ])( - "$cond, should export canonical json, update available exports and emit export-available event", - async function ({ exportName }) { + }); + }); + + describe.each([ + { cond: "when exportName does not contain extension", exportName: `foo.bar.${Date.now()}` }, + { cond: "when exportName contains extension", exportName: `foo.bar.${Date.now()}.json` }, + ])("$cond", ({ exportName }) => { + it("should export canonical json, update available exports and emit export-available event", async () => { const emitSpy = vi.spyOn(session, "emit"); session.sessionId = dummySessionId; manager = new SessionExportsManager(session, exportsManagerConfig); @@ -189,32 +229,81 @@ describe("SessionExportsManager integration test", () => { jsonExportFormat: "canonical", }); + const expectedExportName = exportName.endsWith(".json") ? exportName : `${exportName}.json`; // Updates available export const availableExports = manager.listAvailableExports(); expect(availableExports).toHaveLength(1); expect(availableExports).toContainEqual( expect.objectContaining({ - name: "foo.bar.json", - uri: "exported-data://foo.bar.json", + name: expectedExportName, + uri: `exported-data://${expectedExportName}`, }) ); // Emit event - expect(emitSpy).toHaveBeenCalledWith("export-available", "exported-data://foo.bar.json"); + expect(emitSpy).toHaveBeenCalledWith("export-available", `exported-data://${expectedExportName}`); // Exports relaxed json - const jsonData = JSON.parse(await manager.readExport("foo.bar.json")) as unknown[]; + const jsonData = JSON.parse(await manager.readExport(expectedExportName)) as unknown[]; expect(jsonData).toContainEqual( expect.objectContaining({ name: "foo", longNumber: { $numberLong: "12" } }) ); expect(jsonData).toContainEqual( expect.objectContaining({ name: "bar", longNumber: { $numberLong: "123456" } }) ); - } - ); + }); + }); + + describe("when transform stream throws an error", () => { + it("should remove the partial export and never make it available", async () => { + const emitSpy = vi.spyOn(session, "emit"); + session.sessionId = dummySessionId; + manager = new SessionExportsManager(session, exportsManagerConfig); + // eslint-disable-next-line @typescript-eslint/no-unsafe-member-access, @typescript-eslint/no-explicit-any + (manager as any).docToEJSONStream = function (ejsonOptions: EJSONOptions | undefined) { + let docsTransformed = 0; + return new Transform({ + objectMode: true, + transform: function (chunk: unknown, encoding, callback) { + ++docsTransformed; + try { + if (docsTransformed === 1) { + throw new Error("Could not transform the chunk!"); + } + const doc: string = EJSON.stringify(chunk, undefined, 2, ejsonOptions); + const line = `${docsTransformed > 1 ? ",\n" : ""}${doc}`; + + callback(null, line); + } catch (err: unknown) { + callback(err as Error); + } + }, + final: function (callback) { + this.push("]"); + callback(null); + }, + }); + }; + + const timestamp = Date.now(); + const exportName = getDummyExportName(timestamp); + const exportPath = getDummyExportPath(timestamp); + await expect(() => + manager.createJSONExport({ + input: inputCursor, + exportName, + jsonExportFormat: "relaxed", + }) + ).rejects.toThrow("Could not transform the chunk!"); + + expect(emitSpy).not.toHaveBeenCalled(); + expect(manager.listAvailableExports()).toEqual([]); + expect(await fileExists(exportPath)).toEqual(false); + }); + }); }); - describe("#cleanupExpiredExports", function () { + describe("#cleanupExpiredExports", () => { let input: FindCursor; beforeEach(() => { void input?.close(); @@ -230,8 +319,8 @@ describe("SessionExportsManager integration test", () => { ]); }); - it("should do nothing if session is not initialized", async function () { - const { path } = await createDummyExport(); + it("should do nothing if session is not initialized", async () => { + const { path } = await createDummyExport(Date.now()); new SessionExportsManager(session, { ...exportsManagerConfig, exportTimeoutMs: 100, @@ -243,8 +332,11 @@ describe("SessionExportsManager integration test", () => { expect(await fileExists(path)).toEqual(true); }); - it("should cleanup expired exports if session is initialized", async function () { + it("should cleanup expired exports if session is initialized", async () => { session.sessionId = dummySessionId; + const timestamp = Date.now(); + const exportName = getDummyExportName(timestamp); + const exportPath = getDummyExportPath(timestamp); const manager = new SessionExportsManager(session, { ...exportsManagerConfig, exportTimeoutMs: 100, @@ -252,20 +344,20 @@ describe("SessionExportsManager integration test", () => { }); await manager.createJSONExport({ input, - exportName: dummyExportName, + exportName, jsonExportFormat: "relaxed", }); expect(manager.listAvailableExports()).toContainEqual( expect.objectContaining({ - name: "foo.bar.json", - uri: "exported-data://foo.bar.json", + name: exportName, + uri: `exported-data://${exportName}`, }) ); - expect(await fileExists(dummyExportPath)).toEqual(true); - await timeout(200); + expect(await fileExists(exportPath)).toEqual(true); + await timeout(150); expect(manager.listAvailableExports()).toEqual([]); - expect(await fileExists(dummyExportPath)).toEqual(false); + expect(await fileExists(exportPath)).toEqual(false); }); }); }); From 90fbbd33247e78be95f81521a18e531a3d63a417 Mon Sep 17 00:00:00 2001 From: Himanshu Singh Date: Wed, 6 Aug 2025 10:36:31 +0200 Subject: [PATCH 11/50] chore: handle possible path traversals --- src/common/logger.ts | 1 + src/common/sessionExportsManager.ts | 42 ++++++++++++++++++---------- src/resources/common/exportedData.ts | 15 ++++++---- 3 files changed, 39 insertions(+), 19 deletions(-) diff --git a/src/common/logger.ts b/src/common/logger.ts index 98a7a829..488e454e 100644 --- a/src/common/logger.ts +++ b/src/common/logger.ts @@ -57,6 +57,7 @@ export const LogId = { exportCloseError: mongoLogId(1_007_005), exportedDataListError: mongoLogId(1_007_006), exportedDataAutoCompleteError: mongoLogId(1_007_007), + exportedDataSessionUninitialized: mongoLogId(1_007_008), } as const; interface LogPayload { diff --git a/src/common/sessionExportsManager.ts b/src/common/sessionExportsManager.ts index bc1bf53b..0d9c611e 100644 --- a/src/common/sessionExportsManager.ts +++ b/src/common/sessionExportsManager.ts @@ -26,8 +26,10 @@ export type SessionExportsManagerConfig = Pick< "exportPath" | "exportTimeoutMs" | "exportCleanupIntervalMs" >; +const MAX_LOCK_RETRIES = 10; + export class SessionExportsManager { - private mutableExports: Export[] = []; + private availableExports: Export[] = []; private exportsCleanupInterval: NodeJS.Timeout; private exportsCleanupInProgress: boolean = false; @@ -56,9 +58,7 @@ export class SessionExportsManager { } public exportNameToResourceURI(nameWithExtension: string): string { - if (!path.extname(nameWithExtension)) { - throw new Error("Provided export name has no extension"); - } + this.validateExportName(nameWithExtension); return `exported-data://${nameWithExtension}`; } @@ -73,9 +73,7 @@ export class SessionExportsManager { } public exportFilePath(exportsDirectoryPath: string, exportNameWithExtension: string): string { - if (!path.extname(exportNameWithExtension)) { - throw new Error("Provided export name has no extension"); - } + this.validateExportName(exportNameWithExtension); return path.join(exportsDirectoryPath, exportNameWithExtension); } @@ -84,13 +82,14 @@ export class SessionExportsManager { // by not acquiring a lock on read. That is because this we require this // interface to be fast and just accurate enough for MCP completions // API. - return this.mutableExports.filter(({ createdAt }) => { + return this.availableExports.filter(({ createdAt }) => { return !this.isExportExpired(createdAt); }); } public async readExport(exportNameWithExtension: string): Promise { try { + this.validateExportName(exportNameWithExtension); const exportsDirectoryPath = await this.ensureExportsDirectory(); const exportFilePath = this.exportFilePath(exportsDirectoryPath, exportNameWithExtension); if (await this.isExportFileExpired(exportFilePath)) { @@ -118,7 +117,9 @@ export class SessionExportsManager { jsonExportFormat: JSONExportFormat; }): Promise { try { - const exportNameWithExtension = this.withExtension(exportName, "json"); + const exportNameWithExtension = this.ensureExtension(exportName, "json"); + this.validateExportName(exportNameWithExtension); + const inputStream = input.stream(); const ejsonDocStream = this.docToEJSONStream(this.getEJSONOptionsForFormat(jsonExportFormat)); await this.withExportsLock(async (exportsDirectoryPath) => { @@ -146,8 +147,8 @@ export class SessionExportsManager { void input.close(); if (pipeSuccessful) { const resourceURI = this.exportNameToResourceURI(exportNameWithExtension); - this.mutableExports = [ - ...this.mutableExports, + this.availableExports = [ + ...this.availableExports, { createdAt: (await fs.stat(exportFilePath)).birthtimeMs, name: exportNameWithExtension, @@ -216,7 +217,7 @@ export class SessionExportsManager { const exportPath = this.exportFilePath(exportsDirectoryPath, exportName); if (await this.isExportFileExpired(exportPath)) { await fs.unlink(exportPath); - this.mutableExports = this.mutableExports.filter(({ name }) => name !== exportName); + this.availableExports = this.availableExports.filter(({ name }) => name !== exportName); this.session.emit("export-expired", this.exportNameToResourceURI(exportName)); } } @@ -232,6 +233,19 @@ export class SessionExportsManager { } } + /** + * Small utility to validate provided export name for path traversal or no + * extension */ + private validateExportName(nameWithExtension: string): void { + if (!path.extname(nameWithExtension)) { + throw new Error("Provided export name has no extension"); + } + + if (nameWithExtension.includes("..") || nameWithExtension.includes("/") || nameWithExtension.includes("\\")) { + throw new Error("Invalid export name: path traversal hinted"); + } + } + /** * Small utility to centrally determine if an export is expired or not */ private async isExportFileExpired(exportFilePath: string): Promise { @@ -252,7 +266,7 @@ export class SessionExportsManager { /** * Ensures the path ends with the provided extension */ - private withExtension(pathOrName: string, extension: string): string { + private ensureExtension(pathOrName: string, extension: string): string { const extWithDot = extension.startsWith(".") ? extension : `.${extension}`; if (path.extname(pathOrName) === extWithDot) { return pathOrName; @@ -274,7 +288,7 @@ export class SessionExportsManager { let releaseLock: (() => Promise) | undefined; const exportsDirectoryPath = await this.ensureExportsDirectory(); try { - releaseLock = await lock(exportsDirectoryPath, { retries: 10 }); + releaseLock = await lock(exportsDirectoryPath, { retries: MAX_LOCK_RETRIES }); return await callback(exportsDirectoryPath); } finally { await releaseLock?.(); diff --git a/src/resources/common/exportedData.ts b/src/resources/common/exportedData.ts index aa1f995f..34aa1ebf 100644 --- a/src/resources/common/exportedData.ts +++ b/src/resources/common/exportedData.ts @@ -18,7 +18,6 @@ export class ExportedData { void this.server.mcpServer.server.sendResourceUpdated({ uri, }); - this.server.mcpServer.sendResourceListChanged(); }); this.server.session.on("export-expired", () => { this.server.mcpServer.sendResourceListChanged(); @@ -52,8 +51,11 @@ export class ExportedData { if (!sessionId) { // Note that we don't throw error here because this is a // non-critical path and safe to return the most harmless value. - - // TODO: log warn here + logger.warning( + LogId.exportedDataSessionUninitialized, + "In ListResourcesCallback of exported-data resource", + "Session not initialized" + ); return { resources: [] }; } @@ -84,8 +86,11 @@ export class ExportedData { if (!sessionId) { // Note that we don't throw error here because this is a // non-critical path and safe to return the most harmless value. - - // TODO: log warn here + logger.warning( + LogId.exportedDataSessionUninitialized, + "In CompleteResourceTemplateCallback of exported-data resource", + "Session not initialized" + ); return []; } From 7ba831f5fe3991c2849bf9737ea38d455d18b391 Mon Sep 17 00:00:00 2001 From: Himanshu Singh Date: Wed, 6 Aug 2025 10:43:57 +0200 Subject: [PATCH 12/50] chore: add accuracy test for export tool --- tests/accuracy/export.test.ts | 69 +++++++++++++++++++++++++++++++++++ 1 file changed, 69 insertions(+) create mode 100644 tests/accuracy/export.test.ts diff --git a/tests/accuracy/export.test.ts b/tests/accuracy/export.test.ts new file mode 100644 index 00000000..9e1f0cff --- /dev/null +++ b/tests/accuracy/export.test.ts @@ -0,0 +1,69 @@ +import { describeAccuracyTests } from "./sdk/describeAccuracyTests.js"; +import { Matcher } from "./sdk/matcher.js"; + +describeAccuracyTests([ + { + prompt: "Export all the movies in 'mflix.movies' namespace.", + expectedToolCalls: [ + { + toolName: "export", + parameters: { + database: "mflix", + collection: "movies", + filter: Matcher.emptyObjectOrUndefined, + limit: Matcher.undefined, + }, + }, + ], + }, + { + prompt: "Export all the movies in 'mflix.movies' namespace with runtime less than 100.", + expectedToolCalls: [ + { + toolName: "export", + parameters: { + database: "mflix", + collection: "movies", + filter: { + runtime: { $lt: 100 }, + }, + }, + }, + ], + }, + { + prompt: "Export all the movie titles available in 'mflix.movies' namespace", + expectedToolCalls: [ + { + toolName: "export", + parameters: { + database: "mflix", + collection: "movies", + projection: { + title: 1, + _id: Matcher.anyOf( + Matcher.undefined, + Matcher.number((value) => value === 0) + ), + }, + filter: Matcher.emptyObjectOrUndefined, + }, + }, + ], + }, + { + prompt: "From the mflix.movies namespace, export the first 2 movies of Horror genre sorted ascending by their runtime", + expectedToolCalls: [ + { + toolName: "export", + parameters: { + database: "mflix", + collection: "movies", + filter: { genres: "Horror" }, + sort: { runtime: 1 }, + limit: 2, + }, + }, + ], + }, +]); From 6cc5f8fbce5ff3bab8717a36c59abd4ceff86a24 Mon Sep 17 00:00:00 2001 From: Himanshu Singh Date: Wed, 6 Aug 2025 10:50:16 +0200 Subject: [PATCH 13/50] chore: increase timeout to allow file to expire --- tests/unit/common/sessionExportsManager.test.ts | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/unit/common/sessionExportsManager.test.ts b/tests/unit/common/sessionExportsManager.test.ts index a869d166..e5d11852 100644 --- a/tests/unit/common/sessionExportsManager.test.ts +++ b/tests/unit/common/sessionExportsManager.test.ts @@ -355,7 +355,7 @@ describe("SessionExportsManager integration test", () => { }) ); expect(await fileExists(exportPath)).toEqual(true); - await timeout(150); + await timeout(200); expect(manager.listAvailableExports()).toEqual([]); expect(await fileExists(exportPath)).toEqual(false); }); From 0c82b2d8c706b2794337525b7f2f747b98b4998c Mon Sep 17 00:00:00 2001 From: Himanshu Singh Date: Wed, 6 Aug 2025 10:56:25 +0200 Subject: [PATCH 14/50] chore: increase the cleanupInterval to test expired resource --- tests/integration/resources/exportedData.test.ts | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/integration/resources/exportedData.test.ts b/tests/integration/resources/exportedData.test.ts index fca7592c..569a49f7 100644 --- a/tests/integration/resources/exportedData.test.ts +++ b/tests/integration/resources/exportedData.test.ts @@ -121,7 +121,7 @@ describeWithMongoDB( return { ...defaultTestConfig, exportTimeoutMs: 200, - exportCleanupIntervalMs: 100, + exportCleanupIntervalMs: 300, }; } ); From 2ac355b485f96f404b04fb596e98dd69321ca49b Mon Sep 17 00:00:00 2001 From: Himanshu Singh Date: Wed, 6 Aug 2025 13:23:29 +0200 Subject: [PATCH 15/50] chore: adds a header row for baseline info markdown --- scripts/accuracy/generateTestSummary.ts | 1 + 1 file changed, 1 insertion(+) diff --git a/scripts/accuracy/generateTestSummary.ts b/scripts/accuracy/generateTestSummary.ts index dc451caf..f328576a 100644 --- a/scripts/accuracy/generateTestSummary.ts +++ b/scripts/accuracy/generateTestSummary.ts @@ -230,6 +230,7 @@ function generateMarkdownBrief( markdownTexts.push( ...[ "## 📊 Baseline Comparison", + "| Metric | Value |", "|--------|-------|", `| **Baseline Commit** | \`${baselineInfo.commitSHA}\` |`, `| **Baseline Run ID** | \`${baselineInfo.accuracyRunId}\` |`, From 953751c567d86b83ca8510ac6f8fdbded0a61a59 Mon Sep 17 00:00:00 2001 From: Himanshu Singh Date: Wed, 6 Aug 2025 17:17:18 +0200 Subject: [PATCH 16/50] chore: sessionId is initialized on Session instantiation --- src/common/config.ts | 6 +- src/common/session.ts | 3 +- src/common/sessionExportsManager.ts | 10 +- src/server.ts | 2 - .../unit/common/sessionExportsManager.test.ts | 107 +++++++----------- 5 files changed, 47 insertions(+), 81 deletions(-) diff --git a/src/common/config.ts b/src/common/config.ts index 0f99d855..2367a3ad 100644 --- a/src/common/config.ts +++ b/src/common/config.ts @@ -19,7 +19,7 @@ export interface UserConfig { apiClientSecret?: string; telemetry: "enabled" | "disabled"; logPath: string; - exportPath: string; + exportsPath: string; exportTimeoutMs: number; exportCleanupIntervalMs: number; connectionString?: string; @@ -38,7 +38,7 @@ export interface UserConfig { const defaults: UserConfig = { apiBaseUrl: "https://cloud.mongodb.com/", logPath: getLogPath(), - exportPath: getExportPath(), + exportsPath: getExportsPath(), exportTimeoutMs: 300000, // 5 minutes exportCleanupIntervalMs: 120000, // 2 minutes connectOptions: { @@ -76,7 +76,7 @@ function getLogPath(): string { return logPath; } -function getExportPath(): string { +function getExportsPath(): string { return path.join(getLocalDataPath(), "mongodb-mcp", "exports"); } diff --git a/src/common/session.ts b/src/common/session.ts index a7358a9f..4d5f2820 100644 --- a/src/common/session.ts +++ b/src/common/session.ts @@ -10,6 +10,7 @@ import { } from "./connectionManager.js"; import { NodeDriverServiceProvider } from "@mongosh/service-provider-node-driver"; import { ErrorCodes, MongoDBError } from "./errors.js"; +import { ObjectId } from "bson"; export interface SessionOptions { apiBaseUrl: string; @@ -29,7 +30,7 @@ export type SessionEvents = { }; export class Session extends EventEmitter { - sessionId?: string; + readonly sessionId = new ObjectId().toString(); connectionManager: ConnectionManager; apiClient: ApiClient; agentRunner?: { diff --git a/src/common/sessionExportsManager.ts b/src/common/sessionExportsManager.ts index 0d9c611e..53872953 100644 --- a/src/common/sessionExportsManager.ts +++ b/src/common/sessionExportsManager.ts @@ -23,7 +23,7 @@ export type Export = { export type SessionExportsManagerConfig = Pick< UserConfig, - "exportPath" | "exportTimeoutMs" | "exportCleanupIntervalMs" + "exportsPath" | "exportTimeoutMs" | "exportCleanupIntervalMs" >; const MAX_LOCK_RETRIES = 10; @@ -63,13 +63,7 @@ export class SessionExportsManager { } public exportsDirectoryPath(): string { - // If the session is not connected, we can't cannot work with exports - // for that session. - if (!this.session.sessionId) { - throw new Error("Cannot retrieve exports directory, no active session. Try to reconnect to the MCP server"); - } - - return path.join(this.config.exportPath, this.session.sessionId); + return path.join(this.config.exportsPath, this.session.sessionId); } public exportFilePath(exportsDirectoryPath: string, exportNameWithExtension: string): string { diff --git a/src/server.ts b/src/server.ts index aa1dd822..6b990d71 100644 --- a/src/server.ts +++ b/src/server.ts @@ -5,7 +5,6 @@ import { AtlasTools } from "./tools/atlas/tools.js"; import { MongoDbTools } from "./tools/mongodb/tools.js"; import { Resources } from "./resources/resources.js"; import { LogId } from "./common/logger.js"; -import { ObjectId } from "mongodb"; import { Telemetry } from "./telemetry/telemetry.js"; import { UserConfig } from "./common/config.js"; import { type ServerEvent } from "./telemetry/types.js"; @@ -77,7 +76,6 @@ export class Server { this.mcpServer.server.oninitialized = (): void => { this.session.setAgentRunner(this.mcpServer.server.getClientVersion()); - this.session.sessionId = new ObjectId().toString(); this.session.logger.info({ id: LogId.serverInitialized, diff --git a/tests/unit/common/sessionExportsManager.test.ts b/tests/unit/common/sessionExportsManager.test.ts index e5d11852..12b82c19 100644 --- a/tests/unit/common/sessionExportsManager.test.ts +++ b/tests/unit/common/sessionExportsManager.test.ts @@ -11,28 +11,33 @@ import { ROOT_DIR } from "../../accuracy/sdk/constants.js"; import { timeout } from "../../integration/helpers.js"; import { EJSON, EJSONOptions } from "bson"; -const dummySessionId = "1FOO"; -const dummyExportsPath = path.join(ROOT_DIR, "tests", "tmp", "exports"); -const dummySessionExportPath = path.join(dummyExportsPath, dummySessionId); +const exportsPath = path.join(ROOT_DIR, "tests", "tmp", "exports"); const exportsManagerConfig: SessionExportsManagerConfig = { - exportPath: dummyExportsPath, + exportsPath, exportTimeoutMs: config.exportTimeoutMs, exportCleanupIntervalMs: config.exportCleanupIntervalMs, } as const; -function getDummyExportName(timestamp: number) { - return `foo.bar.${timestamp}.json`; -} -function getDummyExportPath(timestamp: number) { - return path.join(dummySessionExportPath, getDummyExportName(timestamp)); + +function getExportNameAndPath(sessionId: string, timestamp: number) { + const exportName = `foo.bar.${timestamp}.json`; + const sessionExportsPath = path.join(exportsPath, sessionId); + const exportPath = path.join(sessionExportsPath, exportName); + return { + sessionExportsPath, + exportName, + exportPath, + exportURI: `exported-data://${exportName}`, + }; } -async function createDummyExport(timestamp: number) { +async function createDummyExport(sessionId: string, timestamp: number) { const content = "[]"; - await fs.mkdir(dummySessionExportPath, { recursive: true }); - await fs.writeFile(getDummyExportPath(timestamp), content); + const { exportName, exportPath, sessionExportsPath } = getExportNameAndPath(sessionId, timestamp); + await fs.mkdir(sessionExportsPath, { recursive: true }); + await fs.writeFile(exportPath, content); return { - name: getDummyExportName(timestamp), - path: getDummyExportPath(timestamp), + exportName, + exportPath, content, }; } @@ -75,8 +80,8 @@ describe("SessionExportsManager integration test", () => { beforeEach(async () => { await manager?.close(); - await fs.rm(exportsManagerConfig.exportPath, { recursive: true, force: true }); - await fs.mkdir(exportsManagerConfig.exportPath, { recursive: true }); + await fs.rm(exportsManagerConfig.exportsPath, { recursive: true, force: true }); + await fs.mkdir(exportsManagerConfig.exportsPath, { recursive: true }); session = new Session({ apiBaseUrl: "" }); manager = new SessionExportsManager(session, exportsManagerConfig); }); @@ -92,25 +97,22 @@ describe("SessionExportsManager integration test", () => { }); describe("#exportsDirectoryPath", () => { - it("should throw when session is not initialized", () => { - expect(() => manager.exportsDirectoryPath()).toThrow(); - }); - it("should return a session path when session is initialized", () => { - session.sessionId = dummySessionId; manager = new SessionExportsManager(session, exportsManagerConfig); - expect(manager.exportsDirectoryPath()).toEqual(path.join(exportsManagerConfig.exportPath, dummySessionId)); + expect(manager.exportsDirectoryPath()).toEqual( + path.join(exportsManagerConfig.exportsPath, session.sessionId) + ); }); }); describe("#exportFilePath", () => { it("should throw when export name has no extension", () => { - expect(() => manager.exportFilePath(dummySessionExportPath, "name")).toThrow(); + expect(() => manager.exportFilePath("session-path", "name")).toThrow(); }); it("should return path to provided export file", () => { - expect(manager.exportFilePath(dummySessionExportPath, "mflix.movies.json")).toEqual( - path.join(dummySessionExportPath, "mflix.movies.json") + expect(manager.exportFilePath("session-path", "mflix.movies.json")).toEqual( + path.join("session-path", "mflix.movies.json") ); }); }); @@ -121,15 +123,16 @@ describe("SessionExportsManager integration test", () => { }); it("should return the resource content", async () => { - const { name, content } = await createDummyExport(Date.now()); - session.sessionId = dummySessionId; - manager = new SessionExportsManager(session, exportsManagerConfig); - expect(await manager.readExport(name)).toEqual(content); + const { exportName, content } = await createDummyExport(session.sessionId, Date.now()); + expect(await manager.readExport(exportName)).toEqual(content); }); }); describe("#createJSONExport", () => { let inputCursor: FindCursor; + let exportName: string; + let exportPath: string; + let exportURI: string; beforeEach(() => { void inputCursor?.close(); inputCursor = createDummyFindCursor([ @@ -142,6 +145,7 @@ describe("SessionExportsManager integration test", () => { longNumber: Long.fromNumber(123456), }, ]); + ({ exportName, exportPath, exportURI } = getExportNameAndPath(session.sessionId, Date.now())); }); describe("when cursor is empty", () => { @@ -149,12 +153,9 @@ describe("SessionExportsManager integration test", () => { inputCursor = createDummyFindCursor([]); const emitSpy = vi.spyOn(session, "emit"); - session.sessionId = dummySessionId; - manager = new SessionExportsManager(session, exportsManagerConfig); - const timestamp = Date.now(); await manager.createJSONExport({ input: inputCursor, - exportName: getDummyExportName(timestamp), + exportName, jsonExportFormat: "relaxed", }); @@ -163,19 +164,16 @@ describe("SessionExportsManager integration test", () => { expect(availableExports).toHaveLength(1); expect(availableExports).toContainEqual( expect.objectContaining({ - name: getDummyExportName(timestamp), - uri: `exported-data://${getDummyExportName(timestamp)}`, + name: exportName, + uri: exportURI, }) ); // Emit event - expect(emitSpy).toHaveBeenCalledWith( - "export-available", - `exported-data://${getDummyExportName(timestamp)}` - ); + expect(emitSpy).toHaveBeenCalledWith("export-available", exportURI); // Exports relaxed json - const jsonData = JSON.parse(await manager.readExport(getDummyExportName(timestamp))) as unknown[]; + const jsonData = JSON.parse(await manager.readExport(exportName)) as unknown[]; expect(jsonData).toEqual([]); }); }); @@ -186,8 +184,6 @@ describe("SessionExportsManager integration test", () => { ])("$cond", ({ exportName }) => { it("should export relaxed json, update available exports and emit export-available event", async () => { const emitSpy = vi.spyOn(session, "emit"); - session.sessionId = dummySessionId; - manager = new SessionExportsManager(session, exportsManagerConfig); await manager.createJSONExport({ input: inputCursor, exportName, @@ -221,8 +217,6 @@ describe("SessionExportsManager integration test", () => { ])("$cond", ({ exportName }) => { it("should export canonical json, update available exports and emit export-available event", async () => { const emitSpy = vi.spyOn(session, "emit"); - session.sessionId = dummySessionId; - manager = new SessionExportsManager(session, exportsManagerConfig); await manager.createJSONExport({ input: inputCursor, exportName, @@ -257,8 +251,6 @@ describe("SessionExportsManager integration test", () => { describe("when transform stream throws an error", () => { it("should remove the partial export and never make it available", async () => { const emitSpy = vi.spyOn(session, "emit"); - session.sessionId = dummySessionId; - manager = new SessionExportsManager(session, exportsManagerConfig); // eslint-disable-next-line @typescript-eslint/no-unsafe-member-access, @typescript-eslint/no-explicit-any (manager as any).docToEJSONStream = function (ejsonOptions: EJSONOptions | undefined) { let docsTransformed = 0; @@ -285,9 +277,6 @@ describe("SessionExportsManager integration test", () => { }); }; - const timestamp = Date.now(); - const exportName = getDummyExportName(timestamp); - const exportPath = getDummyExportPath(timestamp); await expect(() => manager.createJSONExport({ input: inputCursor, @@ -319,24 +308,8 @@ describe("SessionExportsManager integration test", () => { ]); }); - it("should do nothing if session is not initialized", async () => { - const { path } = await createDummyExport(Date.now()); - new SessionExportsManager(session, { - ...exportsManagerConfig, - exportTimeoutMs: 100, - exportCleanupIntervalMs: 50, - }); - - expect(await fileExists(path)).toEqual(true); - await timeout(200); - expect(await fileExists(path)).toEqual(true); - }); - it("should cleanup expired exports if session is initialized", async () => { - session.sessionId = dummySessionId; - const timestamp = Date.now(); - const exportName = getDummyExportName(timestamp); - const exportPath = getDummyExportPath(timestamp); + const { exportName, exportPath, exportURI } = getExportNameAndPath(session.sessionId, Date.now()); const manager = new SessionExportsManager(session, { ...exportsManagerConfig, exportTimeoutMs: 100, @@ -351,7 +324,7 @@ describe("SessionExportsManager integration test", () => { expect(manager.listAvailableExports()).toContainEqual( expect.objectContaining({ name: exportName, - uri: `exported-data://${exportName}`, + uri: exportURI, }) ); expect(await fileExists(exportPath)).toEqual(true); From ca4c4a5611e9f2cd8dd3ae0c5ca6bfa6ce2928e5 Mon Sep 17 00:00:00 2001 From: Himanshu Singh Date: Wed, 6 Aug 2025 19:44:17 +0200 Subject: [PATCH 17/50] chore: emit export related events on ExportsManager instance --- src/common/session.ts | 2 -- src/common/sessionExportsManager.ts | 13 ++++++++++--- src/resources/common/exportedData.ts | 4 ++-- tests/unit/common/sessionExportsManager.test.ts | 8 ++++---- 4 files changed, 16 insertions(+), 11 deletions(-) diff --git a/src/common/session.ts b/src/common/session.ts index 4d5f2820..e54bc3bd 100644 --- a/src/common/session.ts +++ b/src/common/session.ts @@ -25,8 +25,6 @@ export type SessionEvents = { close: []; disconnect: []; "connection-error": [string]; - "export-expired": [string]; - "export-available": [string]; }; export class Session extends EventEmitter { diff --git a/src/common/sessionExportsManager.ts b/src/common/sessionExportsManager.ts index 53872953..228fbd07 100644 --- a/src/common/sessionExportsManager.ts +++ b/src/common/sessionExportsManager.ts @@ -1,6 +1,7 @@ import z from "zod"; import path from "path"; import fs from "fs/promises"; +import EventEmitter from "events"; import { createWriteStream } from "fs"; import { lock } from "proper-lockfile"; import { FindCursor } from "mongodb"; @@ -28,7 +29,12 @@ export type SessionExportsManagerConfig = Pick< const MAX_LOCK_RETRIES = 10; -export class SessionExportsManager { +type SessionExportsManagerEvents = { + "export-expired": [string]; + "export-available": [string]; +}; + +export class SessionExportsManager extends EventEmitter { private availableExports: Export[] = []; private exportsCleanupInterval: NodeJS.Timeout; private exportsCleanupInProgress: boolean = false; @@ -37,6 +43,7 @@ export class SessionExportsManager { private readonly session: Session, private readonly config: SessionExportsManagerConfig ) { + super(); this.exportsCleanupInterval = setInterval( () => void this.cleanupExpiredExports(), this.config.exportCleanupIntervalMs @@ -149,7 +156,7 @@ export class SessionExportsManager { uri: resourceURI, }, ]; - this.session.emit("export-available", resourceURI); + this.emit("export-available", resourceURI); } } }); @@ -212,7 +219,7 @@ export class SessionExportsManager { if (await this.isExportFileExpired(exportPath)) { await fs.unlink(exportPath); this.availableExports = this.availableExports.filter(({ name }) => name !== exportName); - this.session.emit("export-expired", this.exportNameToResourceURI(exportName)); + this.emit("export-expired", this.exportNameToResourceURI(exportName)); } } }); diff --git a/src/resources/common/exportedData.ts b/src/resources/common/exportedData.ts index 34aa1ebf..71b70691 100644 --- a/src/resources/common/exportedData.ts +++ b/src/resources/common/exportedData.ts @@ -13,13 +13,13 @@ export class ExportedData { private readonly uri = "exported-data://{exportName}"; constructor(private readonly server: Server) { - this.server.session.on("export-available", (uri) => { + this.server.exportsManager.on("export-available", (uri) => { this.server.mcpServer.sendResourceListChanged(); void this.server.mcpServer.server.sendResourceUpdated({ uri, }); }); - this.server.session.on("export-expired", () => { + this.server.exportsManager.on("export-expired", () => { this.server.mcpServer.sendResourceListChanged(); }); } diff --git a/tests/unit/common/sessionExportsManager.test.ts b/tests/unit/common/sessionExportsManager.test.ts index 12b82c19..ff8980b5 100644 --- a/tests/unit/common/sessionExportsManager.test.ts +++ b/tests/unit/common/sessionExportsManager.test.ts @@ -152,7 +152,7 @@ describe("SessionExportsManager integration test", () => { it("should create an empty export", async () => { inputCursor = createDummyFindCursor([]); - const emitSpy = vi.spyOn(session, "emit"); + const emitSpy = vi.spyOn(manager, "emit"); await manager.createJSONExport({ input: inputCursor, exportName, @@ -183,7 +183,7 @@ describe("SessionExportsManager integration test", () => { { cond: "when exportName contains extension", exportName: `foo.bar.${Date.now()}.json` }, ])("$cond", ({ exportName }) => { it("should export relaxed json, update available exports and emit export-available event", async () => { - const emitSpy = vi.spyOn(session, "emit"); + const emitSpy = vi.spyOn(manager, "emit"); await manager.createJSONExport({ input: inputCursor, exportName, @@ -216,7 +216,7 @@ describe("SessionExportsManager integration test", () => { { cond: "when exportName contains extension", exportName: `foo.bar.${Date.now()}.json` }, ])("$cond", ({ exportName }) => { it("should export canonical json, update available exports and emit export-available event", async () => { - const emitSpy = vi.spyOn(session, "emit"); + const emitSpy = vi.spyOn(manager, "emit"); await manager.createJSONExport({ input: inputCursor, exportName, @@ -250,7 +250,7 @@ describe("SessionExportsManager integration test", () => { describe("when transform stream throws an error", () => { it("should remove the partial export and never make it available", async () => { - const emitSpy = vi.spyOn(session, "emit"); + const emitSpy = vi.spyOn(manager, "emit"); // eslint-disable-next-line @typescript-eslint/no-unsafe-member-access, @typescript-eslint/no-explicit-any (manager as any).docToEJSONStream = function (ejsonOptions: EJSONOptions | undefined) { let docsTransformed = 0; From 030d81a6d756de4d8bdd3d7a4408f4b1eff7220c Mon Sep 17 00:00:00 2001 From: Himanshu Singh Date: Wed, 6 Aug 2025 19:47:14 +0200 Subject: [PATCH 18/50] chore: SessionExportManager moved to Session instance --- src/common/logger.ts | 1 - src/common/session.ts | 4 ++ src/resources/common/exportedData.ts | 43 +++---------------- src/server.ts | 7 +-- src/tools/mongodb/mongodbTool.ts | 3 -- src/tools/mongodb/read/export.ts | 11 ++--- src/transports/base.ts | 4 -- tests/integration/helpers.ts | 4 -- .../resources/exportedData.test.ts | 3 +- 9 files changed, 18 insertions(+), 62 deletions(-) diff --git a/src/common/logger.ts b/src/common/logger.ts index 488e454e..98a7a829 100644 --- a/src/common/logger.ts +++ b/src/common/logger.ts @@ -57,7 +57,6 @@ export const LogId = { exportCloseError: mongoLogId(1_007_005), exportedDataListError: mongoLogId(1_007_006), exportedDataAutoCompleteError: mongoLogId(1_007_007), - exportedDataSessionUninitialized: mongoLogId(1_007_008), } as const; interface LogPayload { diff --git a/src/common/session.ts b/src/common/session.ts index e54bc3bd..47f41ea2 100644 --- a/src/common/session.ts +++ b/src/common/session.ts @@ -11,6 +11,8 @@ import { import { NodeDriverServiceProvider } from "@mongosh/service-provider-node-driver"; import { ErrorCodes, MongoDBError } from "./errors.js"; import { ObjectId } from "bson"; +import { SessionExportsManager } from "./sessionExportsManager.js"; +import { config } from "./config.js"; export interface SessionOptions { apiBaseUrl: string; @@ -29,6 +31,7 @@ export type SessionEvents = { export class Session extends EventEmitter { readonly sessionId = new ObjectId().toString(); + readonly exportsManager = new SessionExportsManager(this, config); connectionManager: ConnectionManager; apiClient: ApiClient; agentRunner?: { @@ -108,6 +111,7 @@ export class Session extends EventEmitter { async close(): Promise { await this.disconnect(); await this.apiClient.close(); + await this.exportsManager.close(); this.emit("close"); } diff --git a/src/resources/common/exportedData.ts b/src/resources/common/exportedData.ts index 71b70691..14598a30 100644 --- a/src/resources/common/exportedData.ts +++ b/src/resources/common/exportedData.ts @@ -13,13 +13,13 @@ export class ExportedData { private readonly uri = "exported-data://{exportName}"; constructor(private readonly server: Server) { - this.server.exportsManager.on("export-available", (uri) => { + this.server.session.exportsManager.on("export-available", (uri) => { this.server.mcpServer.sendResourceListChanged(); void this.server.mcpServer.server.sendResourceUpdated({ uri, }); }); - this.server.exportsManager.on("export-expired", () => { + this.server.session.exportsManager.on("export-expired", () => { this.server.mcpServer.sendResourceListChanged(); }); } @@ -47,19 +47,7 @@ export class ExportedData { private listResourcesCallback: ListResourcesCallback = () => { try { - const sessionId = this.server.session.sessionId; - if (!sessionId) { - // Note that we don't throw error here because this is a - // non-critical path and safe to return the most harmless value. - logger.warning( - LogId.exportedDataSessionUninitialized, - "In ListResourcesCallback of exported-data resource", - "Session not initialized" - ); - return { resources: [] }; - } - - const sessionExports = this.server.exportsManager.listAvailableExports(); + const sessionExports = this.server.session.exportsManager.listAvailableExports(); return { resources: sessionExports.map(({ name, uri }) => ({ name: name, @@ -82,19 +70,7 @@ export class ExportedData { private autoCompleteExportName: CompleteResourceTemplateCallback = (value) => { try { - const sessionId = this.server.session.sessionId; - if (!sessionId) { - // Note that we don't throw error here because this is a - // non-critical path and safe to return the most harmless value. - logger.warning( - LogId.exportedDataSessionUninitialized, - "In CompleteResourceTemplateCallback of exported-data resource", - "Session not initialized" - ); - return []; - } - - const sessionExports = this.server.exportsManager.listAvailableExports(); + const sessionExports = this.server.session.exportsManager.listAvailableExports(); return sessionExports.filter(({ name }) => name.startsWith(value)).map(({ name }) => name); } catch (error) { logger.error( @@ -108,11 +84,6 @@ export class ExportedData { private readResourceCallback: ReadResourceTemplateCallback = async (uri, { exportName }) => { try { - const sessionId = this.server.session.sessionId; - if (!sessionId) { - throw new Error("Cannot retrieve exported data, session is not valid."); - } - if (typeof exportName !== "string") { throw new Error("Cannot retrieve exported data, exportName not provided."); } @@ -120,8 +91,8 @@ export class ExportedData { return { contents: [ { - uri: this.server.exportsManager.exportNameToResourceURI(exportName), - text: await this.server.exportsManager.readExport(exportName), + uri: this.server.session.exportsManager.exportNameToResourceURI(exportName), + text: await this.server.session.exportsManager.readExport(exportName), mimeType: "application/json", }, ], @@ -132,7 +103,7 @@ export class ExportedData { { uri: typeof exportName === "string" - ? this.server.exportsManager.exportNameToResourceURI(exportName) + ? this.server.session.exportsManager.exportNameToResourceURI(exportName) : this.uri, text: `Error reading from ${this.uri}: ${error instanceof Error ? error.message : String(error)}`, }, diff --git a/src/server.ts b/src/server.ts index 6b990d71..de88d7be 100644 --- a/src/server.ts +++ b/src/server.ts @@ -12,11 +12,9 @@ import { type ServerCommand } from "./telemetry/types.js"; import { CallToolRequestSchema, CallToolResult } from "@modelcontextprotocol/sdk/types.js"; import assert from "assert"; import { ToolBase } from "./tools/tool.js"; -import { SessionExportsManager } from "./common/sessionExportsManager.js"; export interface ServerOptions { session: Session; - exportsManager: SessionExportsManager; userConfig: UserConfig; mcpServer: McpServer; telemetry: Telemetry; @@ -24,17 +22,15 @@ export interface ServerOptions { export class Server { public readonly session: Session; - public readonly exportsManager: SessionExportsManager; public readonly mcpServer: McpServer; private readonly telemetry: Telemetry; public readonly userConfig: UserConfig; public readonly tools: ToolBase[] = []; private readonly startTime: number; - constructor({ session, exportsManager, mcpServer, userConfig, telemetry }: ServerOptions) { + constructor({ session, mcpServer, userConfig, telemetry }: ServerOptions) { this.startTime = Date.now(); this.session = session; - this.exportsManager = exportsManager; this.telemetry = telemetry; this.mcpServer = mcpServer; this.userConfig = userConfig; @@ -100,7 +96,6 @@ export class Server { } async close(): Promise { - await this.exportsManager.close(); await this.telemetry.close(); await this.session.close(); await this.mcpServer.close(); diff --git a/src/tools/mongodb/mongodbTool.ts b/src/tools/mongodb/mongodbTool.ts index 0d24aa72..708209f8 100644 --- a/src/tools/mongodb/mongodbTool.ts +++ b/src/tools/mongodb/mongodbTool.ts @@ -5,7 +5,6 @@ import { CallToolResult } from "@modelcontextprotocol/sdk/types.js"; import { ErrorCodes, MongoDBError } from "../../common/errors.js"; import { LogId } from "../../common/logger.js"; import { Server } from "../../server.js"; -import { SessionExportsManager } from "../../common/sessionExportsManager.js"; export const DbOperationArgs = { database: z.string().describe("Database name"), @@ -14,7 +13,6 @@ export const DbOperationArgs = { export abstract class MongoDBToolBase extends ToolBase { private server?: Server; - public exportsManager?: SessionExportsManager; public category: ToolCategory = "mongodb"; protected async ensureConnected(): Promise { @@ -49,7 +47,6 @@ export abstract class MongoDBToolBase extends ToolBase { public register(server: Server): boolean { this.server = server; - this.exportsManager = server.exportsManager; return super.register(server); } diff --git a/src/tools/mongodb/read/export.ts b/src/tools/mongodb/read/export.ts index b037d5e6..5c4e1cd4 100644 --- a/src/tools/mongodb/read/export.ts +++ b/src/tools/mongodb/read/export.ts @@ -42,14 +42,11 @@ export class ExportTool extends MongoDBToolBase { bsonRegExp: true, }); const exportName = `${database}.${collection}.${Date.now()}.json`; - if (!this.exportsManager) { - throw new Error("Incorrect server configuration, export not possible!"); - } - await this.exportsManager.createJSONExport({ input: findCursor, exportName, jsonExportFormat }); - const exportedResourceURI = this.exportsManager.exportNameToResourceURI(exportName); - const exportedResourcePath = this.exportsManager.exportFilePath( - this.exportsManager.exportsDirectoryPath(), + await this.session.exportsManager.createJSONExport({ input: findCursor, exportName, jsonExportFormat }); + const exportedResourceURI = this.session.exportsManager.exportNameToResourceURI(exportName); + const exportedResourcePath = this.session.exportsManager.exportFilePath( + this.session.exportsManager.exportsDirectoryPath(), exportName ); const toolCallContent: CallToolResult["content"] = [ diff --git a/src/transports/base.ts b/src/transports/base.ts index 62c79214..7052f1c4 100644 --- a/src/transports/base.ts +++ b/src/transports/base.ts @@ -5,7 +5,6 @@ import { Session } from "../common/session.js"; import { Telemetry } from "../telemetry/telemetry.js"; import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js"; import { CompositeLogger, ConsoleLogger, DiskLogger, LoggerBase, McpLogger } from "../common/logger.js"; -import { SessionExportsManager } from "../common/sessionExportsManager.js"; export abstract class TransportRunnerBase { public logger: LoggerBase; @@ -47,14 +46,11 @@ export abstract class TransportRunnerBase { logger: new CompositeLogger(...loggers), }); - const exportsManager = new SessionExportsManager(session, userConfig); - const telemetry = Telemetry.create(session, userConfig); return new Server({ mcpServer, session, - exportsManager, telemetry, userConfig, }); diff --git a/tests/integration/helpers.ts b/tests/integration/helpers.ts index b965b779..a53e7456 100644 --- a/tests/integration/helpers.ts +++ b/tests/integration/helpers.ts @@ -10,7 +10,6 @@ import { config } from "../../src/common/config.js"; import { afterAll, afterEach, beforeAll, describe, expect, it, vi } from "vitest"; import { ConnectionManager } from "../../src/common/connectionManager.js"; import { CompositeLogger } from "../../src/common/logger.js"; -import { SessionExportsManager } from "../../src/common/sessionExportsManager.js"; interface ParameterInfo { name: string; @@ -76,11 +75,8 @@ export function setupIntegrationTest(getUserConfig: () => UserConfig): Integrati const telemetry = Telemetry.create(session, userConfig); - const exportsManager = new SessionExportsManager(session, userConfig); - mcpServer = new Server({ session, - exportsManager, userConfig, telemetry, mcpServer: new McpServer({ diff --git a/tests/integration/resources/exportedData.test.ts b/tests/integration/resources/exportedData.test.ts index 569a49f7..2fc2c5eb 100644 --- a/tests/integration/resources/exportedData.test.ts +++ b/tests/integration/resources/exportedData.test.ts @@ -44,7 +44,8 @@ describeWithMongoDB( }); }); - describe("when requesting an expired resource", () => { + // Skipping this currently as there seems to be a timing issue + describe.skip("when requesting an expired resource", () => { it("should return an error", async () => { await integration.connectMcpClient(); const exportResponse = await integration.mcpClient().callTool({ From c8dcf4552bad306c5038c08dbc314476551faaa8 Mon Sep 17 00:00:00 2001 From: Himanshu Singh Date: Thu, 7 Aug 2025 10:30:26 +0200 Subject: [PATCH 19/50] chore: handle exports management in memory --- src/common/session.ts | 15 +- src/common/sessionExportsManager.ts | 256 ++++++++---------- src/resources/common/exportedData.ts | 13 +- src/tools/mongodb/read/export.ts | 17 +- tests/integration/helpers.ts | 1 + .../resources/exportedData.test.ts | 11 +- .../unit/common/sessionExportsManager.test.ts | 112 ++++---- 7 files changed, 203 insertions(+), 222 deletions(-) diff --git a/src/common/session.ts b/src/common/session.ts index 47f41ea2..09f3d5ae 100644 --- a/src/common/session.ts +++ b/src/common/session.ts @@ -11,7 +11,7 @@ import { import { NodeDriverServiceProvider } from "@mongosh/service-provider-node-driver"; import { ErrorCodes, MongoDBError } from "./errors.js"; import { ObjectId } from "bson"; -import { SessionExportsManager } from "./sessionExportsManager.js"; +import { SessionExportsManager, SessionExportsManagerConfig } from "./sessionExportsManager.js"; import { config } from "./config.js"; export interface SessionOptions { @@ -19,6 +19,7 @@ export interface SessionOptions { apiClientId?: string; apiClientSecret?: string; connectionManager?: ConnectionManager; + exportsManagerConfig?: SessionExportsManagerConfig; logger: CompositeLogger; } @@ -31,7 +32,7 @@ export type SessionEvents = { export class Session extends EventEmitter { readonly sessionId = new ObjectId().toString(); - readonly exportsManager = new SessionExportsManager(this, config); + readonly exportsManager: SessionExportsManager; connectionManager: ConnectionManager; apiClient: ApiClient; agentRunner?: { @@ -41,7 +42,14 @@ export class Session extends EventEmitter { public logger: CompositeLogger; - constructor({ apiBaseUrl, apiClientId, apiClientSecret, connectionManager, logger }: SessionOptions) { + constructor({ + apiBaseUrl, + apiClientId, + apiClientSecret, + connectionManager, + logger, + exportsManagerConfig, + }: SessionOptions) { super(); this.logger = logger; @@ -55,6 +63,7 @@ export class Session extends EventEmitter { : undefined; this.apiClient = new ApiClient({ baseUrl: apiBaseUrl, credentials }, logger); + this.exportsManager = new SessionExportsManager(this.sessionId, exportsManagerConfig ?? config); this.connectionManager = connectionManager ?? new ConnectionManager(); this.connectionManager.on("connection-succeeded", () => this.emit("connect")); diff --git a/src/common/sessionExportsManager.ts b/src/common/sessionExportsManager.ts index 228fbd07..f317f955 100644 --- a/src/common/sessionExportsManager.ts +++ b/src/common/sessionExportsManager.ts @@ -3,14 +3,12 @@ import path from "path"; import fs from "fs/promises"; import EventEmitter from "events"; import { createWriteStream } from "fs"; -import { lock } from "proper-lockfile"; import { FindCursor } from "mongodb"; import { EJSON, EJSONOptions } from "bson"; import { Transform } from "stream"; import { pipeline } from "stream/promises"; import { UserConfig } from "./config.js"; -import { Session } from "./session.js"; import logger, { LogId } from "./logger.js"; export const jsonExportFormat = z.enum(["relaxed", "canonical"]); @@ -19,6 +17,7 @@ export type JSONExportFormat = z.infer; export type Export = { name: string; uri: string; + path: string; createdAt: number; }; @@ -27,23 +26,23 @@ export type SessionExportsManagerConfig = Pick< "exportsPath" | "exportTimeoutMs" | "exportCleanupIntervalMs" >; -const MAX_LOCK_RETRIES = 10; - type SessionExportsManagerEvents = { "export-expired": [string]; "export-available": [string]; }; export class SessionExportsManager extends EventEmitter { - private availableExports: Export[] = []; - private exportsCleanupInterval: NodeJS.Timeout; + private availableExports: Record = {}; private exportsCleanupInProgress: boolean = false; + private exportsCleanupInterval: NodeJS.Timeout; + private exportsDirectoryPath: string; constructor( - private readonly session: Session, + private readonly sessionId: string, private readonly config: SessionExportsManagerConfig ) { super(); + this.exportsDirectoryPath = path.join(this.config.exportsPath, sessionId); this.exportsCleanupInterval = setInterval( () => void this.cleanupExpiredExports(), this.config.exportCleanupIntervalMs @@ -53,8 +52,7 @@ export class SessionExportsManager extends EventEmitter { - return !this.isExportExpired(createdAt); - }); + return Object.values(this.availableExports).filter( + ({ createdAt }) => !isExportExpired(createdAt, this.config.exportTimeoutMs) + ); } - public async readExport(exportNameWithExtension: string): Promise { + public async readExport(exportName: string): Promise<{ + content: string; + exportURI: string; + }> { try { - this.validateExportName(exportNameWithExtension); - const exportsDirectoryPath = await this.ensureExportsDirectory(); - const exportFilePath = this.exportFilePath(exportsDirectoryPath, exportNameWithExtension); - if (await this.isExportFileExpired(exportFilePath)) { - throw new Error("Export has expired"); + const exportNameWithExtension = validateExportName(exportName); + const exportHandle = this.availableExports[exportNameWithExtension]; + if (!exportHandle) { + throw new Error("Requested export has either expired or does not exist!"); } - return await fs.readFile(exportFilePath, "utf8"); + const { path: exportPath, uri, createdAt } = exportHandle; + + if (isExportExpired(createdAt, this.config.exportTimeoutMs)) { + throw new Error("Requested export has expired!"); + } + + return { + exportURI: uri, + content: await fs.readFile(exportPath, "utf8"), + }; } catch (error) { logger.error( LogId.exportReadError, "Error when reading export", error instanceof Error ? error.message : String(error) ); + if ((error as NodeJS.ErrnoException).code === "ENOENT") { + throw new Error("Requested export does not exist!"); + } throw error; } } @@ -116,50 +110,53 @@ export class SessionExportsManager extends EventEmitter { + }): Promise<{ + exportURI: string; + exportPath: string; + }> { try { - const exportNameWithExtension = this.ensureExtension(exportName, "json"); - this.validateExportName(exportNameWithExtension); + const exportNameWithExtension = validateExportName(ensureExtension(exportName, "json")); + const exportURI = `exported-data://${encodeURIComponent(exportNameWithExtension)}`; + const exportFilePath = path.join(this.exportsDirectoryPath, exportNameWithExtension); + await fs.mkdir(this.exportsDirectoryPath, { recursive: true }); const inputStream = input.stream(); const ejsonDocStream = this.docToEJSONStream(this.getEJSONOptionsForFormat(jsonExportFormat)); - await this.withExportsLock(async (exportsDirectoryPath) => { - const exportFilePath = path.join(exportsDirectoryPath, exportNameWithExtension); - const outputStream = createWriteStream(exportFilePath); - outputStream.write("["); - let pipeSuccessful = false; - try { - await pipeline([inputStream, ejsonDocStream, outputStream]); - pipeSuccessful = true; - } catch (pipelineError) { - // If the pipeline errors out then we might end up with - // partial and incorrect export so we remove it entirely. - await fs.unlink(exportFilePath).catch((error) => { - if ((error as NodeJS.ErrnoException).code !== "ENOENT") { - logger.error( - LogId.exportCreationCleanupError, - "Error when removing partial export", - error instanceof Error ? error.message : String(error) - ); - } - }); - throw pipelineError; - } finally { - void input.close(); - if (pipeSuccessful) { - const resourceURI = this.exportNameToResourceURI(exportNameWithExtension); - this.availableExports = [ - ...this.availableExports, - { - createdAt: (await fs.stat(exportFilePath)).birthtimeMs, - name: exportNameWithExtension, - uri: resourceURI, - }, - ]; - this.emit("export-available", resourceURI); + const outputStream = createWriteStream(exportFilePath); + outputStream.write("["); + let pipeSuccessful = false; + try { + await pipeline([inputStream, ejsonDocStream, outputStream]); + pipeSuccessful = true; + return { + exportURI, + exportPath: exportFilePath, + }; + } catch (pipelineError) { + // If the pipeline errors out then we might end up with + // partial and incorrect export so we remove it entirely. + await fs.unlink(exportFilePath).catch((error) => { + if ((error as NodeJS.ErrnoException).code !== "ENOENT") { + logger.error( + LogId.exportCreationCleanupError, + "Error when removing partial export", + error instanceof Error ? error.message : String(error) + ); } + }); + throw pipelineError; + } finally { + void input.close(); + if (pipeSuccessful) { + this.availableExports[exportNameWithExtension] = { + name: exportNameWithExtension, + createdAt: Date.now(), + path: exportFilePath, + uri: exportURI, + }; + this.emit("export-available", exportURI); } - }); + } } catch (error) { logger.error( LogId.exportCreationError, @@ -211,18 +208,15 @@ export class SessionExportsManager extends EventEmitter { - const exports = await this.listExportFiles(); - for (const exportName of exports) { - const exportPath = this.exportFilePath(exportsDirectoryPath, exportName); - if (await this.isExportFileExpired(exportPath)) { - await fs.unlink(exportPath); - this.availableExports = this.availableExports.filter(({ name }) => name !== exportName); - this.emit("export-expired", this.exportNameToResourceURI(exportName)); - } + for (const { path: exportPath, createdAt, uri, name } of Object.values(exportsToBeConsidered)) { + if (isExportExpired(createdAt, this.config.exportTimeoutMs)) { + delete this.availableExports[name]; + await this.silentlyRemoveExport(exportPath); + this.emit("export-expired", uri); } - }); + } } catch (error) { logger.error( LogId.exportCleanupError, @@ -234,76 +228,50 @@ export class SessionExportsManager extends EventEmitter { + private async silentlyRemoveExport(exportPath: string) { try { - const stats = await fs.stat(exportFilePath); - return this.isExportExpired(stats.birthtimeMs); + await fs.unlink(exportPath); } catch (error) { - if ((error as NodeJS.ErrnoException).code === "ENOENT") { - throw new Error("Requested export does not exist!"); + // If the file does not exist or the containing directory itself + // does not exist then we can safely ignore that error anything else + // we need to flag. + if ((error as NodeJS.ErrnoException).code !== "ENOENT") { + logger.error( + LogId.exportCleanupError, + "Considerable error when removing export file", + error instanceof Error ? error.message : String(error) + ); } - throw error; } } +} - private isExportExpired(createdAt: number) { - return Date.now() - createdAt > this.config.exportTimeoutMs; +/** + * Ensures the path ends with the provided extension */ +export function ensureExtension(pathOrName: string, extension: string): string { + const extWithDot = extension.startsWith(".") ? extension : `.${extension}`; + if (pathOrName.endsWith(extWithDot)) { + return pathOrName; } + return `${pathOrName}${extWithDot}`; +} - /** - * Ensures the path ends with the provided extension */ - private ensureExtension(pathOrName: string, extension: string): string { - const extWithDot = extension.startsWith(".") ? extension : `.${extension}`; - if (path.extname(pathOrName) === extWithDot) { - return pathOrName; - } - return `${pathOrName}${extWithDot}`; +/** + * Small utility to decoding and validating provided export name for path + * traversal or no extension */ +export function validateExportName(nameWithExtension: string): string { + const decodedName = decodeURIComponent(nameWithExtension); + if (!path.extname(decodedName)) { + throw new Error("Provided export name has no extension"); } - /** - * Creates the session exports directory and returns the path */ - private async ensureExportsDirectory(): Promise { - const exportsDirectoryPath = this.exportsDirectoryPath(); - await fs.mkdir(exportsDirectoryPath, { recursive: true }); - return exportsDirectoryPath; + if (decodedName.includes("..") || decodedName.includes("/") || decodedName.includes("\\")) { + throw new Error("Invalid export name: path traversal hinted"); } - /** - * Acquires a lock on the session exports directory. */ - private async withExportsLock(callback: (lockedPath: string) => Promise): Promise { - let releaseLock: (() => Promise) | undefined; - const exportsDirectoryPath = await this.ensureExportsDirectory(); - try { - releaseLock = await lock(exportsDirectoryPath, { retries: MAX_LOCK_RETRIES }); - return await callback(exportsDirectoryPath); - } finally { - await releaseLock?.(); - } - } + return decodedName; +} - /** - * Lists exported files in the session export directory, while ignoring the - * hidden files and files without extensions. */ - private async listExportFiles(): Promise { - const exportsDirectory = await this.ensureExportsDirectory(); - const directoryContents = await fs.readdir(exportsDirectory, "utf8"); - return directoryContents.filter((maybeExportName) => { - return !maybeExportName.startsWith(".") && !!path.extname(maybeExportName); - }); - } +export function isExportExpired(createdAt: number, exportTimeoutMs: number) { + return Date.now() - createdAt > exportTimeoutMs; } diff --git a/src/resources/common/exportedData.ts b/src/resources/common/exportedData.ts index 14598a30..429b883c 100644 --- a/src/resources/common/exportedData.ts +++ b/src/resources/common/exportedData.ts @@ -82,17 +82,19 @@ export class ExportedData { } }; - private readResourceCallback: ReadResourceTemplateCallback = async (uri, { exportName }) => { + private readResourceCallback: ReadResourceTemplateCallback = async (url, { exportName }) => { try { if (typeof exportName !== "string") { throw new Error("Cannot retrieve exported data, exportName not provided."); } + const { content, exportURI } = await this.server.session.exportsManager.readExport(exportName); + return { contents: [ { - uri: this.server.session.exportsManager.exportNameToResourceURI(exportName), - text: await this.server.session.exportsManager.readExport(exportName), + uri: exportURI, + text: content, mimeType: "application/json", }, ], @@ -101,10 +103,7 @@ export class ExportedData { return { contents: [ { - uri: - typeof exportName === "string" - ? this.server.session.exportsManager.exportNameToResourceURI(exportName) - : this.uri, + uri: url.href, text: `Error reading from ${this.uri}: ${error instanceof Error ? error.message : String(error)}`, }, ], diff --git a/src/tools/mongodb/read/export.ts b/src/tools/mongodb/read/export.ts index 5c4e1cd4..0eca4efc 100644 --- a/src/tools/mongodb/read/export.ts +++ b/src/tools/mongodb/read/export.ts @@ -43,24 +43,23 @@ export class ExportTool extends MongoDBToolBase { }); const exportName = `${database}.${collection}.${Date.now()}.json`; - await this.session.exportsManager.createJSONExport({ input: findCursor, exportName, jsonExportFormat }); - const exportedResourceURI = this.session.exportsManager.exportNameToResourceURI(exportName); - const exportedResourcePath = this.session.exportsManager.exportFilePath( - this.session.exportsManager.exportsDirectoryPath(), - exportName - ); + const { exportURI, exportPath } = await this.session.exportsManager.createJSONExport({ + input: findCursor, + exportName, + jsonExportFormat, + }); const toolCallContent: CallToolResult["content"] = [ // Not all the clients as of this commit understands how to // parse a resource_link so we provide a text result for them to // understand what to do with the result. { type: "text", - text: `Exported data for namespace ${database}.${collection} is available under resource URI - "${exportedResourceURI}".`, + text: `Exported data for namespace ${database}.${collection} is available under resource URI - "${exportURI}".`, }, { type: "resource_link", name: exportName, - uri: exportedResourceURI, + uri: exportURI, description: "Resource URI for fetching exported data.", mimeType: "application/json", }, @@ -71,7 +70,7 @@ export class ExportTool extends MongoDBToolBase { if (this.config.transport === "stdio") { toolCallContent.push({ type: "text", - text: `Optionally, the exported data can also be accessed under path - "${exportedResourcePath}"`, + text: `Optionally, the exported data can also be accessed under path - "${exportPath}"`, }); } diff --git a/tests/integration/helpers.ts b/tests/integration/helpers.ts index a53e7456..1134921b 100644 --- a/tests/integration/helpers.ts +++ b/tests/integration/helpers.ts @@ -63,6 +63,7 @@ export function setupIntegrationTest(getUserConfig: () => UserConfig): Integrati apiClientSecret: userConfig.apiClientSecret, connectionManager, logger: new CompositeLogger(), + exportsManagerConfig: userConfig, }); // Mock hasValidAccessToken for tests diff --git a/tests/integration/resources/exportedData.test.ts b/tests/integration/resources/exportedData.test.ts index 2fc2c5eb..4aed5f2c 100644 --- a/tests/integration/resources/exportedData.test.ts +++ b/tests/integration/resources/exportedData.test.ts @@ -1,8 +1,8 @@ import { Long } from "bson"; import { describe, expect, it, beforeEach } from "vitest"; -import { describeWithMongoDB } from "../tools/mongodb/mongodbHelpers.js"; -import { defaultTestConfig, timeout } from "../helpers.js"; import { CallToolResult } from "@modelcontextprotocol/sdk/types.js"; +import { defaultTestConfig, timeout } from "../helpers.js"; +import { describeWithMongoDB } from "../tools/mongodb/mongodbHelpers.js"; describeWithMongoDB( "exported-data resource", @@ -39,13 +39,12 @@ describeWithMongoDB( expect(response.isError).toEqual(true); expect(response.contents[0]?.uri).toEqual("exported-data://db.coll.json"); expect(response.contents[0]?.text).toEqual( - "Error reading from exported-data://{exportName}: Requested export does not exist!" + "Error reading from exported-data://{exportName}: Requested export has either expired or does not exist!" ); }); }); - // Skipping this currently as there seems to be a timing issue - describe.skip("when requesting an expired resource", () => { + describe("when requesting an expired resource", () => { it("should return an error", async () => { await integration.connectMcpClient(); const exportResponse = await integration.mcpClient().callTool({ @@ -66,7 +65,7 @@ describeWithMongoDB( expect(response.isError).toEqual(true); expect(response.contents[0]?.uri).toEqual(exportedResourceURI); expect(response.contents[0]?.text).toEqual( - "Error reading from exported-data://{exportName}: Export has expired" + "Error reading from exported-data://{exportName}: Requested export has expired!" ); }); }); diff --git a/tests/unit/common/sessionExportsManager.test.ts b/tests/unit/common/sessionExportsManager.test.ts index ff8980b5..feb6f7ab 100644 --- a/tests/unit/common/sessionExportsManager.test.ts +++ b/tests/unit/common/sessionExportsManager.test.ts @@ -3,7 +3,13 @@ import fs from "fs/promises"; import { Readable, Transform } from "stream"; import { FindCursor, Long } from "mongodb"; import { beforeEach, describe, expect, it, vi } from "vitest"; -import { SessionExportsManager, SessionExportsManagerConfig } from "../../../src/common/sessionExportsManager.js"; +import { + ensureExtension, + isExportExpired, + SessionExportsManager, + SessionExportsManagerConfig, + validateExportName, +} from "../../../src/common/sessionExportsManager.js"; import { config } from "../../../src/common/config.js"; import { Session } from "../../../src/common/session.js"; @@ -30,18 +36,6 @@ function getExportNameAndPath(sessionId: string, timestamp: number) { }; } -async function createDummyExport(sessionId: string, timestamp: number) { - const content = "[]"; - const { exportName, exportPath, sessionExportsPath } = getExportNameAndPath(sessionId, timestamp); - await fs.mkdir(sessionExportsPath, { recursive: true }); - await fs.writeFile(exportPath, content); - return { - exportName, - exportPath, - content, - }; -} - function createDummyFindCursor(dataArray: unknown[]): FindCursor { let index = 0; const readable = new Readable({ @@ -74,7 +68,7 @@ async function fileExists(filePath: string) { } } -describe("SessionExportsManager integration test", () => { +describe("SessionExportsManager unit test", () => { let session: Session; let manager: SessionExportsManager; @@ -83,38 +77,7 @@ describe("SessionExportsManager integration test", () => { await fs.rm(exportsManagerConfig.exportsPath, { recursive: true, force: true }); await fs.mkdir(exportsManagerConfig.exportsPath, { recursive: true }); session = new Session({ apiBaseUrl: "" }); - manager = new SessionExportsManager(session, exportsManagerConfig); - }); - - describe("#exportNameToResourceURI", () => { - it("should throw when export name has no extension", () => { - expect(() => manager.exportNameToResourceURI("name")).toThrow(); - }); - - it("should return a resource URI", () => { - expect(manager.exportNameToResourceURI("name.json")).toEqual("exported-data://name.json"); - }); - }); - - describe("#exportsDirectoryPath", () => { - it("should return a session path when session is initialized", () => { - manager = new SessionExportsManager(session, exportsManagerConfig); - expect(manager.exportsDirectoryPath()).toEqual( - path.join(exportsManagerConfig.exportsPath, session.sessionId) - ); - }); - }); - - describe("#exportFilePath", () => { - it("should throw when export name has no extension", () => { - expect(() => manager.exportFilePath("session-path", "name")).toThrow(); - }); - - it("should return path to provided export file", () => { - expect(manager.exportFilePath("session-path", "mflix.movies.json")).toEqual( - path.join("session-path", "mflix.movies.json") - ); - }); + manager = session.exportsManager; }); describe("#readExport", () => { @@ -123,8 +86,17 @@ describe("SessionExportsManager integration test", () => { }); it("should return the resource content", async () => { - const { exportName, content } = await createDummyExport(session.sessionId, Date.now()); - expect(await manager.readExport(exportName)).toEqual(content); + const { exportName, exportURI } = getExportNameAndPath(session.sessionId, Date.now()); + const inputCursor = createDummyFindCursor([]); + await manager.createJSONExport({ + input: inputCursor, + exportName, + jsonExportFormat: "relaxed", + }); + expect(await manager.readExport(exportName)).toEqual({ + content: "[]", + exportURI, + }); }); }); @@ -173,7 +145,7 @@ describe("SessionExportsManager integration 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([]); }); }); @@ -205,7 +177,7 @@ describe("SessionExportsManager integration 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 })); }); @@ -238,7 +210,7 @@ describe("SessionExportsManager integration 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" } }) ); @@ -308,9 +280,9 @@ describe("SessionExportsManager integration test", () => { ]); }); - it("should cleanup expired exports if session is initialized", async () => { + it("should cleanup expired exports", async () => { const { exportName, exportPath, exportURI } = getExportNameAndPath(session.sessionId, Date.now()); - const manager = new SessionExportsManager(session, { + const manager = new SessionExportsManager(session.sessionId, { ...exportsManagerConfig, exportTimeoutMs: 100, exportCleanupIntervalMs: 50, @@ -334,3 +306,37 @@ describe("SessionExportsManager integration test", () => { }); }); }); + +describe("#ensureExtension", () => { + it("should append provided extension when not present", () => { + expect(ensureExtension("random", "json")).toEqual("random.json"); + expect(ensureExtension("random.1234", "json")).toEqual("random.1234.json"); + expect(ensureExtension("/random/random-file", "json")).toEqual("/random/random-file.json"); + }); + it("should not append provided when present", () => { + expect(ensureExtension("random.json", "json")).toEqual("random.json"); + expect(ensureExtension("random.1234.json", "json")).toEqual("random.1234.json"); + expect(ensureExtension("/random/random-file.json", "json")).toEqual("/random/random-file.json"); + }); +}); + +describe("#validateExportName", () => { + it("should return decoded name when name is valid", () => { + expect(validateExportName(encodeURIComponent("Test Name.json"))).toEqual("Test Name.json"); + }); + it("should throw when name is invalid", () => { + expect(() => validateExportName("NoExtension")).toThrow("Provided export name has no extension"); + expect(() => validateExportName("../something.json")).toThrow("Invalid export name: path traversal hinted"); + }); +}); + +describe("#isExportExpired", () => { + it("should return true if export is expired", () => { + const createdAt = Date.now() - 1000; + expect(isExportExpired(createdAt, 500)).toEqual(true); + }); + it("should return false if export is not expired", () => { + const createdAt = Date.now(); + expect(isExportExpired(createdAt, 500)).toEqual(false); + }); +}); From 6d2a3644b259a733866e56bc5201602be38ed46a Mon Sep 17 00:00:00 2001 From: Himanshu Singh Date: Thu, 7 Aug 2025 10:31:13 +0200 Subject: [PATCH 20/50] chore: move proper-lockfile back to dev --- package-lock.json | 5 ++++- package.json | 2 +- 2 files changed, 5 insertions(+), 2 deletions(-) diff --git a/package-lock.json b/package-lock.json index d99d5a6f..d6e7d7dd 100644 --- a/package-lock.json +++ b/package-lock.json @@ -27,7 +27,6 @@ "node-machine-id": "1.1.12", "oauth4webapi": "^3.6.0", "openapi-fetch": "^0.14.0", - "proper-lockfile": "^4.1.2", "yargs-parser": "^22.0.0", "zod": "^3.25.76" }, @@ -59,6 +58,7 @@ "openapi-types": "^12.1.3", "openapi-typescript": "^7.8.0", "prettier": "^3.6.2", + "proper-lockfile": "^4.1.2", "simple-git": "^3.28.0", "tsx": "^4.20.3", "typescript": "^5.8.3", @@ -8646,6 +8646,7 @@ "version": "4.2.11", "resolved": "https://registry.npmjs.org/graceful-fs/-/graceful-fs-4.2.11.tgz", "integrity": "sha512-RbJ5/jmFcNNCcDV5o9eTnBLJ/HszWV0P73bc+Ff4nS/rJj+YaS6IGyiOL0VoBYX+l1Wrl3k63h/KrH+nhJ0XvQ==", + "dev": true, "license": "ISC" }, "node_modules/graphemer": { @@ -11299,6 +11300,7 @@ "version": "4.1.2", "resolved": "https://registry.npmjs.org/proper-lockfile/-/proper-lockfile-4.1.2.tgz", "integrity": "sha512-TjNPblN4BwAWMXU8s9AEz4JmQxnD1NNL7bNOY/AKUzyamc379FWASUhc/K1pL2noVb+XmZKLL68cjzLsiOAMaA==", + "dev": true, "license": "MIT", "dependencies": { "graceful-fs": "^4.2.4", @@ -11731,6 +11733,7 @@ "version": "0.12.0", "resolved": "https://registry.npmjs.org/retry/-/retry-0.12.0.tgz", "integrity": "sha512-9LkiTwjUh6rT555DtE9rTX+BKByPfrMzEAtnlEtdEwr3Nkffwiihqe2bWADg+OQRjt9gl6ICdmB/ZFDCGAtSow==", + "dev": true, "license": "MIT", "engines": { "node": ">= 4" diff --git a/package.json b/package.json index 2b1cf9f1..7bcd5a1e 100644 --- a/package.json +++ b/package.json @@ -80,6 +80,7 @@ "openapi-types": "^12.1.3", "openapi-typescript": "^7.8.0", "prettier": "^3.6.2", + "proper-lockfile": "^4.1.2", "simple-git": "^3.28.0", "tsx": "^4.20.3", "typescript": "^5.8.3", @@ -107,7 +108,6 @@ "node-machine-id": "1.1.12", "oauth4webapi": "^3.6.0", "openapi-fetch": "^0.14.0", - "proper-lockfile": "^4.1.2", "yargs-parser": "^22.0.0", "zod": "^3.25.76" }, From 6d7efe6ea02589406228dc540dd03aa8b4f124ee Mon Sep 17 00:00:00 2001 From: Himanshu Singh Date: Thu, 7 Aug 2025 10:56:13 +0200 Subject: [PATCH 21/50] chore: adapt logger changes --- src/common/session.ts | 2 +- src/common/sessionExportsManager.ts | 67 ++++++++++--------- src/resources/common/exportedData.ts | 22 +++--- .../unit/common/sessionExportsManager.test.ts | 17 +++-- 4 files changed, 57 insertions(+), 51 deletions(-) diff --git a/src/common/session.ts b/src/common/session.ts index 09f3d5ae..8e6895df 100644 --- a/src/common/session.ts +++ b/src/common/session.ts @@ -63,7 +63,7 @@ export class Session extends EventEmitter { : undefined; this.apiClient = new ApiClient({ baseUrl: apiBaseUrl, credentials }, logger); - this.exportsManager = new SessionExportsManager(this.sessionId, exportsManagerConfig ?? config); + this.exportsManager = new SessionExportsManager(this.sessionId, exportsManagerConfig ?? config, logger); this.connectionManager = connectionManager ?? new ConnectionManager(); this.connectionManager.on("connection-succeeded", () => this.emit("connect")); diff --git a/src/common/sessionExportsManager.ts b/src/common/sessionExportsManager.ts index f317f955..0461f1ed 100644 --- a/src/common/sessionExportsManager.ts +++ b/src/common/sessionExportsManager.ts @@ -9,7 +9,7 @@ import { Transform } from "stream"; import { pipeline } from "stream/promises"; import { UserConfig } from "./config.js"; -import logger, { LogId } from "./logger.js"; +import { LoggerBase, LogId } from "./logger.js"; export const jsonExportFormat = z.enum(["relaxed", "canonical"]); export type JSONExportFormat = z.infer; @@ -38,8 +38,9 @@ export class SessionExportsManager extends EventEmitter { if ((error as NodeJS.ErrnoException).code !== "ENOENT") { - logger.error( - LogId.exportCreationCleanupError, - "Error when removing partial export", - error instanceof Error ? error.message : String(error) - ); + this.logger.error({ + id: LogId.exportCreationCleanupError, + context: "Error when removing partial export", + message: error instanceof Error ? error.message : String(error), + }); } }); throw pipelineError; @@ -158,11 +159,11 @@ export class SessionExportsManager extends EventEmitter name.startsWith(value)).map(({ name }) => name); } catch (error) { - logger.error( - LogId.exportedDataAutoCompleteError, - "Error when autocompleting exported data", - error instanceof Error ? error.message : String(error) - ); + this.server.session.logger.error({ + id: LogId.exportedDataAutoCompleteError, + context: "Error when autocompleting exported data", + message: error instanceof Error ? error.message : String(error), + }); return []; } }; diff --git a/tests/unit/common/sessionExportsManager.test.ts b/tests/unit/common/sessionExportsManager.test.ts index feb6f7ab..f5044e3c 100644 --- a/tests/unit/common/sessionExportsManager.test.ts +++ b/tests/unit/common/sessionExportsManager.test.ts @@ -16,6 +16,7 @@ import { Session } from "../../../src/common/session.js"; import { ROOT_DIR } from "../../accuracy/sdk/constants.js"; import { timeout } from "../../integration/helpers.js"; import { EJSON, EJSONOptions } from "bson"; +import { CompositeLogger } from "../../../src/common/logger.js"; const exportsPath = path.join(ROOT_DIR, "tests", "tmp", "exports"); const exportsManagerConfig: SessionExportsManagerConfig = { @@ -76,7 +77,7 @@ describe("SessionExportsManager unit test", () => { await manager?.close(); await fs.rm(exportsManagerConfig.exportsPath, { recursive: true, force: true }); await fs.mkdir(exportsManagerConfig.exportsPath, { recursive: true }); - session = new Session({ apiBaseUrl: "" }); + session = new Session({ apiBaseUrl: "", logger: new CompositeLogger() }); manager = session.exportsManager; }); @@ -282,11 +283,15 @@ describe("SessionExportsManager unit test", () => { it("should cleanup expired exports", async () => { const { exportName, exportPath, exportURI } = getExportNameAndPath(session.sessionId, Date.now()); - const manager = new SessionExportsManager(session.sessionId, { - ...exportsManagerConfig, - exportTimeoutMs: 100, - exportCleanupIntervalMs: 50, - }); + const manager = new SessionExportsManager( + session.sessionId, + { + ...exportsManagerConfig, + exportTimeoutMs: 100, + exportCleanupIntervalMs: 50, + }, + new CompositeLogger() + ); await manager.createJSONExport({ input, exportName, From d009b3cb8b4124bcdef3682f39330558ada68bc8 Mon Sep 17 00:00:00 2001 From: Himanshu Singh Date: Thu, 7 Aug 2025 11:00:45 +0200 Subject: [PATCH 22/50] chore: remove indents from exported json --- src/common/sessionExportsManager.ts | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/common/sessionExportsManager.ts b/src/common/sessionExportsManager.ts index 0461f1ed..c589a601 100644 --- a/src/common/sessionExportsManager.ts +++ b/src/common/sessionExportsManager.ts @@ -188,7 +188,7 @@ export class SessionExportsManager extends EventEmitter 1 ? ",\n" : ""}${doc}`; callback(null, line); From 4992fa606b9b41e0a4637f0d35d28f043a2a59b0 Mon Sep 17 00:00:00 2001 From: Himanshu Singh Date: Thu, 7 Aug 2025 11:58:28 +0200 Subject: [PATCH 23/50] chore: PR feedback - rename json to EJSON in tool meta - move availableExports as getter prop - make SessionExportsManager return Promise --- src/common/sessionExportsManager.ts | 26 +++++++++---------- src/resources/common/exportedData.ts | 8 +++--- src/tools/mongodb/read/export.ts | 4 +-- .../tools/mongodb/read/export.test.ts | 4 +-- .../unit/common/sessionExportsManager.test.ts | 12 ++++----- 5 files changed, 27 insertions(+), 27 deletions(-) diff --git a/src/common/sessionExportsManager.ts b/src/common/sessionExportsManager.ts index c589a601..53ca10b7 100644 --- a/src/common/sessionExportsManager.ts +++ b/src/common/sessionExportsManager.ts @@ -32,7 +32,7 @@ type SessionExportsManagerEvents = { }; export class SessionExportsManager extends EventEmitter { - private availableExports: Record = {}; + private sessionExports: Record = {}; private exportsCleanupInProgress: boolean = false; private exportsCleanupInterval: NodeJS.Timeout; private exportsDirectoryPath: string; @@ -50,7 +50,13 @@ export class SessionExportsManager extends EventEmitter !isExportExpired(createdAt, this.config.exportTimeoutMs) + ); + } + + public async close(): Promise { try { clearInterval(this.exportsCleanupInterval); await fs.rm(this.exportsDirectoryPath, { force: true, recursive: true }); @@ -63,19 +69,13 @@ export class SessionExportsManager extends EventEmitter !isExportExpired(createdAt, this.config.exportTimeoutMs) - ); - } - public async readExport(exportName: string): Promise<{ content: string; exportURI: string; }> { try { const exportNameWithExtension = validateExportName(exportName); - const exportHandle = this.availableExports[exportNameWithExtension]; + const exportHandle = this.sessionExports[exportNameWithExtension]; if (!exportHandle) { throw new Error("Requested export has either expired or does not exist!"); } @@ -149,7 +149,7 @@ export class SessionExportsManager extends EventEmitter { try { - const sessionExports = this.server.session.exportsManager.listAvailableExports(); return { - resources: sessionExports.map(({ name, uri }) => ({ + resources: this.server.session.exportsManager.availableExports.map(({ name, uri }) => ({ name: name, description: this.exportNameToDescription(name), uri: uri, @@ -70,8 +69,9 @@ export class ExportedData { private autoCompleteExportName: CompleteResourceTemplateCallback = (value) => { try { - const sessionExports = this.server.session.exportsManager.listAvailableExports(); - return sessionExports.filter(({ name }) => name.startsWith(value)).map(({ name }) => name); + return this.server.session.exportsManager.availableExports + .filter(({ name }) => name.startsWith(value)) + .map(({ name }) => name); } catch (error) { this.server.session.logger.error({ id: LogId.exportedDataAutoCompleteError, diff --git a/src/tools/mongodb/read/export.ts b/src/tools/mongodb/read/export.ts index 0eca4efc..d217486d 100644 --- a/src/tools/mongodb/read/export.ts +++ b/src/tools/mongodb/read/export.ts @@ -7,7 +7,7 @@ import z from "zod"; export class ExportTool extends MongoDBToolBase { public name = "export"; - protected description = "Export a collection data or query results in the specified json format."; + protected description = "Export a collection data or query results in the specified EJSON format."; protected argsShape = { ...DbOperationArgs, ...FindArgs, @@ -16,7 +16,7 @@ export class ExportTool extends MongoDBToolBase { .default("relaxed") .describe( [ - "The format to be used when exporting collection data as JSON with default being relaxed.", + "The format to be used when exporting collection data as EJSON with default being relaxed.", "relaxed: A string format that emphasizes readability and interoperability at the expense of type preservation. That is, conversion from relaxed format to BSON can lose type information.", "canonical: A string format that emphasizes type preservation at the expense of readability and interoperability. That is, conversion from canonical to BSON will generally preserve type information except in certain specific cases.", ].join("\n") diff --git a/tests/integration/tools/mongodb/read/export.test.ts b/tests/integration/tools/mongodb/read/export.test.ts index 66ffc829..4c175d9b 100644 --- a/tests/integration/tools/mongodb/read/export.test.ts +++ b/tests/integration/tools/mongodb/read/export.test.ts @@ -37,7 +37,7 @@ describeWithMongoDB("export tool", (integration) => { validateToolMetadata( integration, "export", - "Export a collection data or query results in the specified json format.", + "Export a collection data or query results in the specified EJSON format.", [ ...databaseCollectionParameters, @@ -50,7 +50,7 @@ describeWithMongoDB("export tool", (integration) => { { name: "jsonExportFormat", description: [ - "The format to be used when exporting collection data as JSON with default being relaxed.", + "The format to be used when exporting collection data as EJSON with default being relaxed.", "relaxed: A string format that emphasizes readability and interoperability at the expense of type preservation. That is, conversion from relaxed format to BSON can lose type information.", "canonical: A string format that emphasizes type preservation at the expense of readability and interoperability. That is, conversion from canonical to BSON will generally preserve type information except in certain specific cases.", ].join("\n"), diff --git a/tests/unit/common/sessionExportsManager.test.ts b/tests/unit/common/sessionExportsManager.test.ts index f5044e3c..20ff58b6 100644 --- a/tests/unit/common/sessionExportsManager.test.ts +++ b/tests/unit/common/sessionExportsManager.test.ts @@ -133,7 +133,7 @@ describe("SessionExportsManager unit test", () => { }); // Updates available export - const availableExports = manager.listAvailableExports(); + const availableExports = manager.availableExports; expect(availableExports).toHaveLength(1); expect(availableExports).toContainEqual( expect.objectContaining({ @@ -165,7 +165,7 @@ describe("SessionExportsManager unit test", () => { const expectedExportName = exportName.endsWith(".json") ? exportName : `${exportName}.json`; // Updates available export - const availableExports = manager.listAvailableExports(); + const availableExports = manager.availableExports; expect(availableExports).toHaveLength(1); expect(availableExports).toContainEqual( expect.objectContaining({ @@ -198,7 +198,7 @@ describe("SessionExportsManager unit test", () => { const expectedExportName = exportName.endsWith(".json") ? exportName : `${exportName}.json`; // Updates available export - const availableExports = manager.listAvailableExports(); + const availableExports = manager.availableExports; expect(availableExports).toHaveLength(1); expect(availableExports).toContainEqual( expect.objectContaining({ @@ -259,7 +259,7 @@ describe("SessionExportsManager unit test", () => { ).rejects.toThrow("Could not transform the chunk!"); expect(emitSpy).not.toHaveBeenCalled(); - expect(manager.listAvailableExports()).toEqual([]); + expect(manager.availableExports).toEqual([]); expect(await fileExists(exportPath)).toEqual(false); }); }); @@ -298,7 +298,7 @@ describe("SessionExportsManager unit test", () => { jsonExportFormat: "relaxed", }); - expect(manager.listAvailableExports()).toContainEqual( + expect(manager.availableExports).toContainEqual( expect.objectContaining({ name: exportName, uri: exportURI, @@ -306,7 +306,7 @@ describe("SessionExportsManager unit test", () => { ); expect(await fileExists(exportPath)).toEqual(true); await timeout(200); - expect(manager.listAvailableExports()).toEqual([]); + expect(manager.availableExports).toEqual([]); expect(await fileExists(exportPath)).toEqual(false); }); }); From 2762891451ef594448ef99372e5b7410009a2638 Mon Sep 17 00:00:00 2001 From: Himanshu Singh Date: Thu, 7 Aug 2025 13:51:30 +0200 Subject: [PATCH 24/50] chore: limit the exposed property of an export --- src/common/sessionExportsManager.ts | 83 ++++++++++++------- src/resources/common/exportedData.ts | 18 ++-- src/tools/mongodb/read/export.ts | 3 +- .../resources/exportedData.test.ts | 9 +- .../unit/common/sessionExportsManager.test.ts | 29 +++---- 5 files changed, 80 insertions(+), 62 deletions(-) diff --git a/src/common/sessionExportsManager.ts b/src/common/sessionExportsManager.ts index 53ca10b7..f438c387 100644 --- a/src/common/sessionExportsManager.ts +++ b/src/common/sessionExportsManager.ts @@ -14,13 +14,36 @@ import { LoggerBase, LogId } from "./logger.js"; export const jsonExportFormat = z.enum(["relaxed", "canonical"]); export type JSONExportFormat = z.infer; -export type Export = { - name: string; - uri: string; - path: string; - createdAt: number; +type StoredExport = { + exportName: string; + exportURI: string; + exportPath: string; + exportCreatedAt: number; }; +/** + * Ideally just exportName and exportURI should be made publicly available but + * we also make exportPath available because the export tool, also returns the + * exportPath in its response when the MCP server is running connected to stdio + * transport. The reasoning behind this is that a few clients, Cursor in + * particular, as of the date of this writing (7 August 2025) cannot refer to + * resource URIs which means they have no means to access the exported resource. + * As of this writing, majority of the usage of our MCP server is behind STDIO + * transport so we can assume that for most of the usages, if not all, the MCP + * server will be running on the same machine as of the MCP client and thus we + * can provide the local path to export so that these clients which do not still + * support parsing resource URIs, can still work with the exported data. We + * expect for clients to catch up and implement referencing resource URIs at + * which point it would be safe to remove the `exportPath` from the publicly + * exposed properties of an export. + * + * The editors that we would like to watch out for are Cursor and Windsurf as + * they don't yet support working with Resource URIs. + * + * Ref Cursor: https://forum.cursor.com/t/cursor-mcp-resource-feature-support/50987 + * JIRA: https://jira.mongodb.org/browse/MCP-104 */ +type AvailableExport = Pick; + export type SessionExportsManagerConfig = Pick< UserConfig, "exportsPath" | "exportTimeoutMs" | "exportCleanupIntervalMs" @@ -32,7 +55,7 @@ type SessionExportsManagerEvents = { }; export class SessionExportsManager extends EventEmitter { - private sessionExports: Record = {}; + private sessionExports: Record = {}; private exportsCleanupInProgress: boolean = false; private exportsCleanupInterval: NodeJS.Timeout; private exportsDirectoryPath: string; @@ -50,10 +73,14 @@ export class SessionExportsManager extends EventEmitter !isExportExpired(createdAt, this.config.exportTimeoutMs) - ); + public get availableExports(): AvailableExport[] { + return Object.values(this.sessionExports) + .filter(({ exportCreatedAt: createdAt }) => !isExportExpired(createdAt, this.config.exportTimeoutMs)) + .map(({ exportName, exportURI, exportPath }) => ({ + exportName, + exportURI, + exportPath, + })); } public async close(): Promise { @@ -69,10 +96,7 @@ export class SessionExportsManager extends EventEmitter { + public async readExport(exportName: string): Promise { try { const exportNameWithExtension = validateExportName(exportName); const exportHandle = this.sessionExports[exportNameWithExtension]; @@ -80,16 +104,13 @@ export class SessionExportsManager extends EventEmitter { + }): Promise { try { const exportNameWithExtension = validateExportName(ensureExtension(exportName, "json")); const exportURI = `exported-data://${encodeURIComponent(exportNameWithExtension)}`; @@ -130,6 +148,7 @@ export class SessionExportsManager extends EventEmitter { try { return { - resources: this.server.session.exportsManager.availableExports.map(({ name, uri }) => ({ - name: name, - description: this.exportNameToDescription(name), - uri: uri, + resources: this.server.session.exportsManager.availableExports.map(({ exportName, exportURI }) => ({ + name: exportName, + description: this.exportNameToDescription(exportName), + uri: exportURI, mimeType: "application/json", })), }; @@ -70,8 +70,8 @@ export class ExportedData { private autoCompleteExportName: CompleteResourceTemplateCallback = (value) => { try { return this.server.session.exportsManager.availableExports - .filter(({ name }) => name.startsWith(value)) - .map(({ name }) => name); + .filter(({ exportName }) => exportName.startsWith(value)) + .map(({ exportName }) => exportName); } catch (error) { this.server.session.logger.error({ id: LogId.exportedDataAutoCompleteError, @@ -88,12 +88,12 @@ export class ExportedData { throw new Error("Cannot retrieve exported data, exportName not provided."); } - const { content, exportURI } = await this.server.session.exportsManager.readExport(exportName); + const content = await this.server.session.exportsManager.readExport(exportName); return { contents: [ { - uri: exportURI, + uri: url.href, text: content, mimeType: "application/json", }, @@ -104,7 +104,7 @@ export class ExportedData { contents: [ { uri: url.href, - text: `Error reading from ${this.uri}: ${error instanceof Error ? error.message : String(error)}`, + text: `Error reading ${url.href}: ${error instanceof Error ? error.message : String(error)}`, }, ], isError: true, diff --git a/src/tools/mongodb/read/export.ts b/src/tools/mongodb/read/export.ts index d217486d..4e31d0a6 100644 --- a/src/tools/mongodb/read/export.ts +++ b/src/tools/mongodb/read/export.ts @@ -66,7 +66,8 @@ export class ExportTool extends MongoDBToolBase { ]; // This special case is to make it easier to work with exported data for - // stdio transport. + // clients that still cannot reference resources (Cursor). + // More information here: https://jira.mongodb.org/browse/MCP-104 if (this.config.transport === "stdio") { toolCallContent.push({ type: "text", diff --git a/tests/integration/resources/exportedData.test.ts b/tests/integration/resources/exportedData.test.ts index 4aed5f2c..e09ace85 100644 --- a/tests/integration/resources/exportedData.test.ts +++ b/tests/integration/resources/exportedData.test.ts @@ -32,14 +32,15 @@ describeWithMongoDB( describe("when requesting non-existent resource", () => { it("should return an error", async () => { + const exportURI = "exported-data://db.coll.json"; await integration.connectMcpClient(); const response = await integration.mcpClient().readResource({ - uri: "exported-data://db.coll.json", + uri: exportURI, }); expect(response.isError).toEqual(true); - expect(response.contents[0]?.uri).toEqual("exported-data://db.coll.json"); + expect(response.contents[0]?.uri).toEqual(exportURI); expect(response.contents[0]?.text).toEqual( - "Error reading from exported-data://{exportName}: Requested export has either expired or does not exist!" + `Error reading ${exportURI}: Requested export has either expired or does not exist!` ); }); }); @@ -65,7 +66,7 @@ describeWithMongoDB( expect(response.isError).toEqual(true); expect(response.contents[0]?.uri).toEqual(exportedResourceURI); expect(response.contents[0]?.text).toEqual( - "Error reading from exported-data://{exportName}: Requested export has expired!" + `Error reading ${exportedResourceURI}: Requested export has expired!` ); }); }); diff --git a/tests/unit/common/sessionExportsManager.test.ts b/tests/unit/common/sessionExportsManager.test.ts index 20ff58b6..69258da9 100644 --- a/tests/unit/common/sessionExportsManager.test.ts +++ b/tests/unit/common/sessionExportsManager.test.ts @@ -87,17 +87,14 @@ describe("SessionExportsManager unit test", () => { }); it("should return the resource content", async () => { - const { exportName, exportURI } = getExportNameAndPath(session.sessionId, Date.now()); + const { exportName } = getExportNameAndPath(session.sessionId, Date.now()); const inputCursor = createDummyFindCursor([]); await manager.createJSONExport({ input: inputCursor, exportName, jsonExportFormat: "relaxed", }); - expect(await manager.readExport(exportName)).toEqual({ - content: "[]", - exportURI, - }); + expect(await manager.readExport(exportName)).toEqual("[]"); }); }); @@ -137,8 +134,8 @@ describe("SessionExportsManager unit test", () => { expect(availableExports).toHaveLength(1); expect(availableExports).toContainEqual( expect.objectContaining({ - name: exportName, - uri: exportURI, + exportName, + exportURI, }) ); @@ -146,7 +143,7 @@ describe("SessionExportsManager unit test", () => { expect(emitSpy).toHaveBeenCalledWith("export-available", exportURI); // Exports relaxed json - const jsonData = JSON.parse((await manager.readExport(exportName)).content) as unknown[]; + const jsonData = JSON.parse(await manager.readExport(exportName)) as unknown[]; expect(jsonData).toEqual([]); }); }); @@ -169,8 +166,8 @@ describe("SessionExportsManager unit test", () => { expect(availableExports).toHaveLength(1); expect(availableExports).toContainEqual( expect.objectContaining({ - name: expectedExportName, - uri: `exported-data://${expectedExportName}`, + exportName: expectedExportName, + exportURI: `exported-data://${expectedExportName}`, }) ); @@ -178,7 +175,7 @@ describe("SessionExportsManager unit test", () => { expect(emitSpy).toHaveBeenCalledWith("export-available", `exported-data://${expectedExportName}`); // Exports relaxed json - const jsonData = JSON.parse((await manager.readExport(expectedExportName)).content) as unknown[]; + const jsonData = JSON.parse(await manager.readExport(expectedExportName)) as unknown[]; expect(jsonData).toContainEqual(expect.objectContaining({ name: "foo", longNumber: 12 })); expect(jsonData).toContainEqual(expect.objectContaining({ name: "bar", longNumber: 123456 })); }); @@ -202,8 +199,8 @@ describe("SessionExportsManager unit test", () => { expect(availableExports).toHaveLength(1); expect(availableExports).toContainEqual( expect.objectContaining({ - name: expectedExportName, - uri: `exported-data://${expectedExportName}`, + exportName: expectedExportName, + exportURI: `exported-data://${expectedExportName}`, }) ); @@ -211,7 +208,7 @@ describe("SessionExportsManager unit test", () => { expect(emitSpy).toHaveBeenCalledWith("export-available", `exported-data://${expectedExportName}`); // Exports relaxed json - const jsonData = JSON.parse((await manager.readExport(expectedExportName)).content) as unknown[]; + const jsonData = JSON.parse(await manager.readExport(expectedExportName)) as unknown[]; expect(jsonData).toContainEqual( expect.objectContaining({ name: "foo", longNumber: { $numberLong: "12" } }) ); @@ -300,8 +297,8 @@ describe("SessionExportsManager unit test", () => { expect(manager.availableExports).toContainEqual( expect.objectContaining({ - name: exportName, - uri: exportURI, + exportName, + exportURI, }) ); expect(await fileExists(exportPath)).toEqual(true); From 9f0d6414ccd5f223d07268e5bfd4950972807a25 Mon Sep 17 00:00:00 2001 From: Himanshu Singh Date: Thu, 7 Aug 2025 15:43:11 +0200 Subject: [PATCH 25/50] feat: makes export tool async --- src/common/sessionExportsManager.ts | 83 +++++++++--- src/tools/mongodb/read/export.ts | 8 +- .../resources/exportedData.test.ts | 10 +- .../tools/mongodb/read/export.test.ts | 26 +++- .../unit/common/sessionExportsManager.test.ts | 121 +++++++++++++++--- 5 files changed, 198 insertions(+), 50 deletions(-) diff --git a/src/common/sessionExportsManager.ts b/src/common/sessionExportsManager.ts index f438c387..96c3e3b1 100644 --- a/src/common/sessionExportsManager.ts +++ b/src/common/sessionExportsManager.ts @@ -14,12 +14,22 @@ import { LoggerBase, LogId } from "./logger.js"; export const jsonExportFormat = z.enum(["relaxed", "canonical"]); export type JSONExportFormat = z.infer; -type StoredExport = { +interface CommonExportData { exportName: string; exportURI: string; exportPath: string; +} + +interface ReadyExport extends CommonExportData { + exportStatus: "ready"; exportCreatedAt: number; -}; +} + +interface InProgressExport extends CommonExportData { + exportStatus: "in-progress"; +} + +type StoredExport = ReadyExport | InProgressExport; /** * Ideally just exportName and exportURI should be made publicly available but @@ -75,7 +85,12 @@ export class SessionExportsManager extends EventEmitter !isExportExpired(createdAt, this.config.exportTimeoutMs)) + .filter((sessionExport) => { + return ( + sessionExport.exportStatus === "ready" && + !isExportExpired(sessionExport.exportCreatedAt, this.config.exportTimeoutMs) + ); + }) .map(({ exportName, exportURI, exportPath }) => ({ exportName, exportURI, @@ -104,6 +119,10 @@ export class SessionExportsManager extends EventEmitter { + }): AvailableExport { try { const exportNameWithExtension = validateExportName(ensureExtension(exportName, "json")); const exportURI = `exported-data://${encodeURIComponent(exportNameWithExtension)}`; const exportFilePath = path.join(this.exportsDirectoryPath, exportNameWithExtension); + const inProgressExport: InProgressExport = (this.sessionExports[exportNameWithExtension] = { + exportName: exportNameWithExtension, + exportPath: exportFilePath, + exportURI: exportURI, + exportStatus: "in-progress", + }); + void this.startExport({ input, jsonExportFormat, inProgressExport }); + return inProgressExport; + } catch (error) { + this.logger.error({ + id: LogId.exportCreationError, + context: "Error when registering JSON export request", + message: error instanceof Error ? error.message : String(error), + }); + throw error; + } + } + + private async startExport({ + input, + jsonExportFormat, + inProgressExport, + }: { + input: FindCursor; + jsonExportFormat: JSONExportFormat; + inProgressExport: InProgressExport; + }): Promise { + try { await fs.mkdir(this.exportsDirectoryPath, { recursive: true }); const inputStream = input.stream(); const ejsonDocStream = this.docToEJSONStream(this.getEJSONOptionsForFormat(jsonExportFormat)); - const outputStream = createWriteStream(exportFilePath); + const outputStream = createWriteStream(inProgressExport.exportPath); outputStream.write("["); let pipeSuccessful = false; try { await pipeline([inputStream, ejsonDocStream, outputStream]); pipeSuccessful = true; - return { - exportName, - exportURI, - exportPath: exportFilePath, - }; } catch (pipelineError) { // If the pipeline errors out then we might end up with // partial and incorrect export so we remove it entirely. - await fs.unlink(exportFilePath).catch((error) => { + await fs.unlink(inProgressExport.exportPath).catch((error) => { if ((error as NodeJS.ErrnoException).code !== "ENOENT") { this.logger.error({ id: LogId.exportCreationCleanupError, @@ -164,17 +206,17 @@ export class SessionExportsManager extends EventEmitter sessionExport.exportStatus === "ready" + ); try { - for (const { exportPath, exportCreatedAt, exportURI, exportName } of Object.values(exportsForCleanup)) { + for (const { exportPath, exportCreatedAt, exportURI, exportName } of exportsForCleanup) { if (isExportExpired(exportCreatedAt, this.config.exportTimeoutMs)) { delete this.sessionExports[exportName]; await this.silentlyRemoveExport(exportPath); diff --git a/src/tools/mongodb/read/export.ts b/src/tools/mongodb/read/export.ts index 4e31d0a6..8e6be9dc 100644 --- a/src/tools/mongodb/read/export.ts +++ b/src/tools/mongodb/read/export.ts @@ -43,7 +43,7 @@ export class ExportTool extends MongoDBToolBase { }); const exportName = `${database}.${collection}.${Date.now()}.json`; - const { exportURI, exportPath } = await this.session.exportsManager.createJSONExport({ + const { exportURI, exportPath } = this.session.exportsManager.createJSONExport({ input: findCursor, exportName, jsonExportFormat, @@ -54,13 +54,13 @@ export class ExportTool extends MongoDBToolBase { // understand what to do with the result. { type: "text", - text: `Exported data for namespace ${database}.${collection} is available under resource URI - "${exportURI}".`, + text: `Data for namespace ${database}.${collection} is being exported and will be made available under resource URI - "${exportURI}".`, }, { type: "resource_link", name: exportName, uri: exportURI, - description: "Resource URI for fetching exported data.", + description: "Resource URI for fetching exported data once it is ready.", mimeType: "application/json", }, ]; @@ -71,7 +71,7 @@ export class ExportTool extends MongoDBToolBase { if (this.config.transport === "stdio") { toolCallContent.push({ type: "text", - text: `Optionally, the exported data can also be accessed under path - "${exportPath}"`, + text: `Optionally, when the export is finished, the exported data can also be accessed under path - "${exportPath}"`, }); } diff --git a/tests/integration/resources/exportedData.test.ts b/tests/integration/resources/exportedData.test.ts index e09ace85..0fe962cb 100644 --- a/tests/integration/resources/exportedData.test.ts +++ b/tests/integration/resources/exportedData.test.ts @@ -59,15 +59,13 @@ describeWithMongoDB( expect(exportedResourceURI).toBeDefined(); // wait for export expired - await timeout(200); + await timeout(250); const response = await integration.mcpClient().readResource({ uri: exportedResourceURI as string, }); expect(response.isError).toEqual(true); expect(response.contents[0]?.uri).toEqual(exportedResourceURI); - expect(response.contents[0]?.text).toEqual( - `Error reading ${exportedResourceURI}: Requested export has expired!` - ); + expect(response.contents[0]?.text).toMatch(`Error reading ${exportedResourceURI}:`); }); }); @@ -78,6 +76,8 @@ describeWithMongoDB( name: "export", arguments: { database: "db", collection: "coll" }, }); + // Small timeout to let export finish + await timeout(50); const exportedResourceURI = (exportResponse as CallToolResult).content.find( (part) => part.type === "resource_link" @@ -98,6 +98,8 @@ describeWithMongoDB( name: "export", arguments: { database: "big", collection: "coll" }, }); + // Small timeout to let export finish + await timeout(50); const exportedResourceURI = (exportResponse as CallToolResult).content.find( (part) => part.type === "resource_link" diff --git a/tests/integration/tools/mongodb/read/export.test.ts b/tests/integration/tools/mongodb/read/export.test.ts index 4c175d9b..c81438bf 100644 --- a/tests/integration/tools/mongodb/read/export.test.ts +++ b/tests/integration/tools/mongodb/read/export.test.ts @@ -2,6 +2,7 @@ import fs from "fs/promises"; import { beforeEach, describe, expect, it } from "vitest"; import { databaseCollectionParameters, + timeout, validateThrowsForInvalidArguments, validateToolMetadata, } from "../../../helpers.js"; @@ -11,10 +12,7 @@ import { Long } from "bson"; function contentWithTextResourceURI(content: CallToolResult["content"], namespace: string) { return content.find((part) => { - return ( - part.type === "text" && - part.text.startsWith(`Exported data for namespace ${namespace} is available under resource URI -`) - ); + return part.type === "text" && part.text.startsWith(`Data for namespace ${namespace}`); }); } @@ -28,7 +26,9 @@ function contentWithExportPath(content: CallToolResult["content"]) { return content.find((part) => { return ( part.type === "text" && - part.text.startsWith(`Optionally, the exported data can also be accessed under path -`) + part.text.startsWith( + `Optionally, when the export is finished, the exported data can also be accessed under path -` + ) ); }); } @@ -95,6 +95,8 @@ describeWithMongoDB("export tool", (integration) => { name: "export", arguments: { database: "non-existent", collection: "foos" }, }); + // Small timeout to let export finish + await timeout(10); const content = response.content as CallToolResult["content"]; const namespace = "non-existent.foos"; @@ -129,6 +131,8 @@ describeWithMongoDB("export tool", (integration) => { name: "export", arguments: { database: integration.randomDbName(), collection: "foo" }, }); + // Small timeout to let export finish + await timeout(10); const localPathPart = contentWithExportPath(response.content as CallToolResult["content"]); expect(localPathPart).toBeDefined(); @@ -150,6 +154,8 @@ describeWithMongoDB("export tool", (integration) => { name: "export", arguments: { database: integration.randomDbName(), collection: "foo", filter: { name: "foo" } }, }); + // Small timeout to let export finish + await timeout(10); const localPathPart = contentWithExportPath(response.content as CallToolResult["content"]); expect(localPathPart).toBeDefined(); @@ -170,6 +176,8 @@ describeWithMongoDB("export tool", (integration) => { name: "export", arguments: { database: integration.randomDbName(), collection: "foo", limit: 1 }, }); + // Small timeout to let export finish + await timeout(10); const localPathPart = contentWithExportPath(response.content as CallToolResult["content"]); expect(localPathPart).toBeDefined(); @@ -195,6 +203,8 @@ describeWithMongoDB("export tool", (integration) => { sort: { longNumber: 1 }, }, }); + // Small timeout to let export finish + await timeout(10); const localPathPart = contentWithExportPath(response.content as CallToolResult["content"]); expect(localPathPart).toBeDefined(); @@ -220,6 +230,8 @@ describeWithMongoDB("export tool", (integration) => { projection: { _id: 0, name: 1 }, }, }); + // Small timeout to let export finish + await timeout(10); const localPathPart = contentWithExportPath(response.content as CallToolResult["content"]); expect(localPathPart).toBeDefined(); @@ -249,6 +261,8 @@ describeWithMongoDB("export tool", (integration) => { jsonExportFormat: "relaxed", }, }); + // Small timeout to let export finish + await timeout(10); const localPathPart = contentWithExportPath(response.content as CallToolResult["content"]); expect(localPathPart).toBeDefined(); @@ -279,6 +293,8 @@ describeWithMongoDB("export tool", (integration) => { jsonExportFormat: "canonical", }, }); + // Small timeout to let export finish + await timeout(10); const localPathPart = contentWithExportPath(response.content as CallToolResult["content"]); expect(localPathPart).toBeDefined(); diff --git a/tests/unit/common/sessionExportsManager.test.ts b/tests/unit/common/sessionExportsManager.test.ts index 69258da9..1de514aa 100644 --- a/tests/unit/common/sessionExportsManager.test.ts +++ b/tests/unit/common/sessionExportsManager.test.ts @@ -37,14 +37,20 @@ function getExportNameAndPath(sessionId: string, timestamp: number) { }; } -function createDummyFindCursor(dataArray: unknown[]): FindCursor { +function createDummyFindCursor(dataArray: unknown[], chunkPushTimeoutMs?: number): FindCursor { let index = 0; const readable = new Readable({ objectMode: true, - read() { + async read() { if (index < dataArray.length) { + if (chunkPushTimeoutMs) { + await timeout(chunkPushTimeoutMs); + } this.push(dataArray[index++]); } else { + if (chunkPushTimeoutMs) { + await timeout(chunkPushTimeoutMs); + } this.push(null); } }, @@ -81,19 +87,62 @@ describe("SessionExportsManager unit test", () => { manager = session.exportsManager; }); + describe("#availableExport", () => { + it("should list only the exports that are in ready state", async () => { + // This export will finish in at-least 1 second + const { exportName: exportName1 } = getExportNameAndPath(session.sessionId, Date.now()); + manager.createJSONExport({ + input: createDummyFindCursor([{ name: "Test1" }], 1000), + exportName: exportName1, + jsonExportFormat: "relaxed", + }); + + // This export will finish way sooner than the first one + const { exportName: exportName2 } = getExportNameAndPath(session.sessionId, Date.now()); + manager.createJSONExport({ + input: createDummyFindCursor([{ name: "Test1" }]), + exportName: exportName2, + jsonExportFormat: "relaxed", + }); + + // Small timeout to let the second export finish + await timeout(10); + expect(manager.availableExports).toHaveLength(1); + expect(manager.availableExports[0]?.exportName).toEqual(exportName2); + }); + }); + describe("#readExport", () => { it("should throw when export name has no extension", async () => { await expect(() => manager.readExport("name")).rejects.toThrow(); }); - it("should return the resource content", async () => { + it("should throw if the resource is still being generated", async () => { + const { exportName } = getExportNameAndPath(session.sessionId, Date.now()); + const inputCursor = createDummyFindCursor([{ name: "Test1" }], 100); + manager.createJSONExport({ + input: inputCursor, + exportName, + jsonExportFormat: "relaxed", + }); + // Small timeout but it won't be sufficient and the export + // generation will still be in progress + await timeout(10); + await expect(() => manager.readExport(exportName)).rejects.toThrow( + "Requested export is still being generated!" + ); + }); + + it("should return the resource content if the resource is ready to be consumed", async () => { const { exportName } = getExportNameAndPath(session.sessionId, Date.now()); const inputCursor = createDummyFindCursor([]); - await manager.createJSONExport({ + manager.createJSONExport({ input: inputCursor, exportName, jsonExportFormat: "relaxed", }); + // Small timeout to account for async operation + await timeout(10); expect(await manager.readExport(exportName)).toEqual("[]"); }); }); @@ -123,11 +172,13 @@ describe("SessionExportsManager unit test", () => { inputCursor = createDummyFindCursor([]); const emitSpy = vi.spyOn(manager, "emit"); - await manager.createJSONExport({ + manager.createJSONExport({ input: inputCursor, exportName, jsonExportFormat: "relaxed", }); + // Small timeout to account for async operation + await timeout(10); // Updates available export const availableExports = manager.availableExports; @@ -154,11 +205,13 @@ describe("SessionExportsManager unit test", () => { ])("$cond", ({ exportName }) => { it("should export relaxed json, update available exports and emit export-available event", async () => { const emitSpy = vi.spyOn(manager, "emit"); - await manager.createJSONExport({ + manager.createJSONExport({ input: inputCursor, exportName, jsonExportFormat: "relaxed", }); + // Small timeout to account for async operation + await timeout(10); const expectedExportName = exportName.endsWith(".json") ? exportName : `${exportName}.json`; // Updates available export @@ -187,11 +240,13 @@ describe("SessionExportsManager unit test", () => { ])("$cond", ({ exportName }) => { it("should export canonical json, update available exports and emit export-available event", async () => { const emitSpy = vi.spyOn(manager, "emit"); - await manager.createJSONExport({ + manager.createJSONExport({ input: inputCursor, exportName, jsonExportFormat: "canonical", }); + // Small timeout to account for async operation + await timeout(50); const expectedExportName = exportName.endsWith(".json") ? exportName : `${exportName}.json`; // Updates available export @@ -218,7 +273,7 @@ describe("SessionExportsManager unit test", () => { }); }); - describe("when transform stream throws an error", () => { + describe("when there is an error in export generation", () => { it("should remove the partial export and never make it available", async () => { const emitSpy = vi.spyOn(manager, "emit"); // eslint-disable-next-line @typescript-eslint/no-unsafe-member-access, @typescript-eslint/no-explicit-any @@ -246,15 +301,18 @@ describe("SessionExportsManager unit test", () => { }, }); }; + manager.createJSONExport({ + input: inputCursor, + exportName, + jsonExportFormat: "relaxed", + }); + // Small timeout to account for async operation + await timeout(50); - await expect(() => - manager.createJSONExport({ - input: inputCursor, - exportName, - jsonExportFormat: "relaxed", - }) - ).rejects.toThrow("Could not transform the chunk!"); - + // Because the export was never populated in the available exports. + await expect(() => manager.readExport(exportName)).rejects.toThrow( + "Requested export has either expired or does not exist!" + ); expect(emitSpy).not.toHaveBeenCalled(); expect(manager.availableExports).toEqual([]); expect(await fileExists(exportPath)).toEqual(false); @@ -278,6 +336,32 @@ describe("SessionExportsManager unit test", () => { ]); }); + it("should not clean up in-progress exports", async () => { + const { exportName } = getExportNameAndPath(session.sessionId, Date.now()); + const manager = new SessionExportsManager( + session.sessionId, + { + ...exportsManagerConfig, + exportTimeoutMs: 100, + exportCleanupIntervalMs: 50, + }, + new CompositeLogger() + ); + manager.createJSONExport({ + input: createDummyFindCursor([{ name: "Test" }], 2000), + exportName, + jsonExportFormat: "relaxed", + }); + + // eslint-disable-next-line @typescript-eslint/no-explicit-any, @typescript-eslint/no-unsafe-member-access + expect((manager as any).sessionExports[exportName]?.exportStatus).toEqual("in-progress"); + + // After clean up interval the export should still be there + await timeout(200); + // eslint-disable-next-line @typescript-eslint/no-explicit-any, @typescript-eslint/no-unsafe-member-access + expect((manager as any).sessionExports[exportName]?.exportStatus).toEqual("in-progress"); + }); + it("should cleanup expired exports", async () => { const { exportName, exportPath, exportURI } = getExportNameAndPath(session.sessionId, Date.now()); const manager = new SessionExportsManager( @@ -289,11 +373,14 @@ describe("SessionExportsManager unit test", () => { }, new CompositeLogger() ); - await manager.createJSONExport({ + manager.createJSONExport({ input, exportName, jsonExportFormat: "relaxed", }); + // Small timeout to account for async operation and let export + // finish + await timeout(10); expect(manager.availableExports).toContainEqual( expect.objectContaining({ From bfa1deea9d16f96045da848a87ac831d863ecbbc Mon Sep 17 00:00:00 2001 From: Himanshu Singh Date: Thu, 7 Aug 2025 16:15:55 +0200 Subject: [PATCH 26/50] chore: more reliable SessionExportManager tests --- src/common/sessionExportsManager.ts | 2 +- .../unit/common/sessionExportsManager.test.ts | 102 ++++++++++-------- 2 files changed, 57 insertions(+), 47 deletions(-) diff --git a/src/common/sessionExportsManager.ts b/src/common/sessionExportsManager.ts index 96c3e3b1..d100a1f1 100644 --- a/src/common/sessionExportsManager.ts +++ b/src/common/sessionExportsManager.ts @@ -209,7 +209,6 @@ export class SessionExportsManager extends EventEmitter } { let index = 0; const readable = new Readable({ objectMode: true, @@ -56,14 +59,26 @@ function createDummyFindCursor(dataArray: unknown[], chunkPushTimeoutMs?: number }, }); + let notifyClose: () => Promise; + const cursorCloseNotification = new Promise((resolve) => { + notifyClose = async () => { + await timeout(10); + resolve(); + }; + }); + readable.once("close", () => void notifyClose?.()); + return { - stream() { - return readable; - }, - close() { - return Promise.resolve(readable.destroy()); - }, - } as unknown as FindCursor; + cursor: { + stream() { + return readable; + }, + close() { + return Promise.resolve(readable.destroy()); + }, + } as unknown as FindCursor, + cursorCloseNotification, + }; } async function fileExists(filePath: string) { @@ -92,21 +107,22 @@ describe("SessionExportsManager unit test", () => { // This export will finish in at-least 1 second const { exportName: exportName1 } = getExportNameAndPath(session.sessionId, Date.now()); manager.createJSONExport({ - input: createDummyFindCursor([{ name: "Test1" }], 1000), + input: createDummyFindCursor([{ name: "Test1" }], 1000).cursor, exportName: exportName1, jsonExportFormat: "relaxed", }); // This export will finish way sooner than the first one const { exportName: exportName2 } = getExportNameAndPath(session.sessionId, Date.now()); + const { cursor, cursorCloseNotification } = createDummyFindCursor([{ name: "Test1" }]); manager.createJSONExport({ - input: createDummyFindCursor([{ name: "Test1" }]), + input: cursor, exportName: exportName2, jsonExportFormat: "relaxed", }); // Small timeout to let the second export finish - await timeout(10); + await cursorCloseNotification; expect(manager.availableExports).toHaveLength(1); expect(manager.availableExports[0]?.exportName).toEqual(exportName2); }); @@ -119,15 +135,13 @@ describe("SessionExportsManager unit test", () => { it("should throw if the resource is still being generated", async () => { const { exportName } = getExportNameAndPath(session.sessionId, Date.now()); - const inputCursor = createDummyFindCursor([{ name: "Test1" }], 100); + const { cursor } = createDummyFindCursor([{ name: "Test1" }], 100); manager.createJSONExport({ - input: inputCursor, + input: cursor, exportName, jsonExportFormat: "relaxed", }); - // Small timeout but it won't be sufficient and the export - // generation will still be in progress - await timeout(10); + // note that we do not wait for cursor close await expect(() => manager.readExport(exportName)).rejects.toThrow( "Requested export is still being generated!" ); @@ -135,26 +149,26 @@ describe("SessionExportsManager unit test", () => { it("should return the resource content if the resource is ready to be consumed", async () => { const { exportName } = getExportNameAndPath(session.sessionId, Date.now()); - const inputCursor = createDummyFindCursor([]); + const { cursor, cursorCloseNotification } = createDummyFindCursor([]); manager.createJSONExport({ - input: inputCursor, + input: cursor, exportName, jsonExportFormat: "relaxed", }); - // Small timeout to account for async operation - await timeout(10); + await cursorCloseNotification; expect(await manager.readExport(exportName)).toEqual("[]"); }); }); describe("#createJSONExport", () => { - let inputCursor: FindCursor; + let cursor: FindCursor; + let cursorCloseNotification: Promise; let exportName: string; let exportPath: string; let exportURI: string; beforeEach(() => { - void inputCursor?.close(); - inputCursor = createDummyFindCursor([ + void cursor?.close(); + ({ cursor, cursorCloseNotification } = createDummyFindCursor([ { name: "foo", longNumber: Long.fromNumber(12), @@ -163,22 +177,21 @@ describe("SessionExportsManager unit test", () => { name: "bar", longNumber: Long.fromNumber(123456), }, - ]); + ])); ({ exportName, exportPath, exportURI } = getExportNameAndPath(session.sessionId, Date.now())); }); describe("when cursor is empty", () => { it("should create an empty export", async () => { - inputCursor = createDummyFindCursor([]); + const { cursor, cursorCloseNotification } = createDummyFindCursor([]); const emitSpy = vi.spyOn(manager, "emit"); manager.createJSONExport({ - input: inputCursor, + input: cursor, exportName, jsonExportFormat: "relaxed", }); - // Small timeout to account for async operation - await timeout(10); + await cursorCloseNotification; // Updates available export const availableExports = manager.availableExports; @@ -206,12 +219,11 @@ describe("SessionExportsManager unit test", () => { it("should export relaxed json, update available exports and emit export-available event", async () => { const emitSpy = vi.spyOn(manager, "emit"); manager.createJSONExport({ - input: inputCursor, + input: cursor, exportName, jsonExportFormat: "relaxed", }); - // Small timeout to account for async operation - await timeout(10); + await cursorCloseNotification; const expectedExportName = exportName.endsWith(".json") ? exportName : `${exportName}.json`; // Updates available export @@ -241,12 +253,11 @@ describe("SessionExportsManager unit test", () => { it("should export canonical json, update available exports and emit export-available event", async () => { const emitSpy = vi.spyOn(manager, "emit"); manager.createJSONExport({ - input: inputCursor, + input: cursor, exportName, jsonExportFormat: "canonical", }); - // Small timeout to account for async operation - await timeout(50); + await cursorCloseNotification; const expectedExportName = exportName.endsWith(".json") ? exportName : `${exportName}.json`; // Updates available export @@ -302,12 +313,11 @@ describe("SessionExportsManager unit test", () => { }); }; manager.createJSONExport({ - input: inputCursor, + input: cursor, exportName, jsonExportFormat: "relaxed", }); - // Small timeout to account for async operation - await timeout(50); + await cursorCloseNotification; // Because the export was never populated in the available exports. await expect(() => manager.readExport(exportName)).rejects.toThrow( @@ -321,10 +331,11 @@ describe("SessionExportsManager unit test", () => { }); describe("#cleanupExpiredExports", () => { - let input: FindCursor; + let cursor: FindCursor; + let cursorCloseNotification: Promise; beforeEach(() => { - void input?.close(); - input = createDummyFindCursor([ + void cursor?.close(); + ({ cursor, cursorCloseNotification } = createDummyFindCursor([ { name: "foo", longNumber: Long.fromNumber(12), @@ -333,7 +344,7 @@ describe("SessionExportsManager unit test", () => { name: "bar", longNumber: Long.fromNumber(123456), }, - ]); + ])); }); it("should not clean up in-progress exports", async () => { @@ -347,8 +358,9 @@ describe("SessionExportsManager unit test", () => { }, new CompositeLogger() ); + const { cursor } = createDummyFindCursor([{ name: "Test" }], 2000); manager.createJSONExport({ - input: createDummyFindCursor([{ name: "Test" }], 2000), + input: cursor, exportName, jsonExportFormat: "relaxed", }); @@ -374,13 +386,11 @@ describe("SessionExportsManager unit test", () => { new CompositeLogger() ); manager.createJSONExport({ - input, + input: cursor, exportName, jsonExportFormat: "relaxed", }); - // Small timeout to account for async operation and let export - // finish - await timeout(10); + await cursorCloseNotification; expect(manager.availableExports).toContainEqual( expect.objectContaining({ From af947d54dcefffc2230caeeee7cc09484bdb1859 Mon Sep 17 00:00:00 2001 From: Himanshu Singh Date: Thu, 7 Aug 2025 17:21:43 +0200 Subject: [PATCH 27/50] chore: more reliable export tool and exported-data resource tests --- tests/integration/helpers.ts | 15 +++- .../resources/exportedData.test.ts | 13 +-- .../tools/mongodb/read/export.test.ts | 83 +++++++++++-------- 3 files changed, 71 insertions(+), 40 deletions(-) diff --git a/tests/integration/helpers.ts b/tests/integration/helpers.ts index 1134921b..7910cdbf 100644 --- a/tests/integration/helpers.ts +++ b/tests/integration/helpers.ts @@ -2,7 +2,7 @@ import { Client } from "@modelcontextprotocol/sdk/client/index.js"; import { InMemoryTransport } from "./inMemoryTransport.js"; import { Server } from "../../src/server.js"; import { UserConfig } from "../../src/common/config.js"; -import { McpError } from "@modelcontextprotocol/sdk/types.js"; +import { McpError, ResourceUpdatedNotificationSchema } from "@modelcontextprotocol/sdk/types.js"; import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js"; import { Session } from "../../src/common/session.js"; import { Telemetry } from "../../src/telemetry/telemetry.js"; @@ -274,3 +274,16 @@ function validateToolAnnotations(tool: ToolInfo, name: string, description: stri export function timeout(ms: number) { return new Promise((resolve) => setTimeout(resolve, ms)); } + +/** + * Subscribes to the resources changed notification for the provided URI + */ +export function resourceChangedNotification(client: Client, uri: string): Promise { + return new Promise((resolve) => { + client.setNotificationHandler(ResourceUpdatedNotificationSchema, (notification) => { + if (notification.params.uri === uri) { + resolve(); + } + }); + }); +} diff --git a/tests/integration/resources/exportedData.test.ts b/tests/integration/resources/exportedData.test.ts index 0fe962cb..1c32ddbc 100644 --- a/tests/integration/resources/exportedData.test.ts +++ b/tests/integration/resources/exportedData.test.ts @@ -1,8 +1,9 @@ import { Long } from "bson"; import { describe, expect, it, beforeEach } from "vitest"; import { CallToolResult } from "@modelcontextprotocol/sdk/types.js"; -import { defaultTestConfig, timeout } from "../helpers.js"; +import { defaultTestConfig, resourceChangedNotification, timeout } from "../helpers.js"; import { describeWithMongoDB } from "../tools/mongodb/mongodbHelpers.js"; +import { contentWithResourceURILink } from "../tools/mongodb/read/export.test.js"; describeWithMongoDB( "exported-data resource", @@ -76,8 +77,9 @@ describeWithMongoDB( name: "export", arguments: { database: "db", collection: "coll" }, }); - // Small timeout to let export finish - await timeout(50); + const content = exportResponse.content as CallToolResult["content"]; + const exportURI = contentWithResourceURILink(content)?.uri as string; + await resourceChangedNotification(integration.mcpClient(), exportURI); const exportedResourceURI = (exportResponse as CallToolResult).content.find( (part) => part.type === "resource_link" @@ -98,8 +100,9 @@ describeWithMongoDB( name: "export", arguments: { database: "big", collection: "coll" }, }); - // Small timeout to let export finish - await timeout(50); + const content = exportResponse.content as CallToolResult["content"]; + const exportURI = contentWithResourceURILink(content)?.uri as string; + await resourceChangedNotification(integration.mcpClient(), exportURI); const exportedResourceURI = (exportResponse as CallToolResult).content.find( (part) => part.type === "resource_link" diff --git a/tests/integration/tools/mongodb/read/export.test.ts b/tests/integration/tools/mongodb/read/export.test.ts index c81438bf..f249f090 100644 --- a/tests/integration/tools/mongodb/read/export.test.ts +++ b/tests/integration/tools/mongodb/read/export.test.ts @@ -2,7 +2,7 @@ import fs from "fs/promises"; import { beforeEach, describe, expect, it } from "vitest"; import { databaseCollectionParameters, - timeout, + resourceChangedNotification, validateThrowsForInvalidArguments, validateToolMetadata, } from "../../../helpers.js"; @@ -10,19 +10,25 @@ import { describeWithMongoDB } from "../mongodbHelpers.js"; import { CallToolResult } from "@modelcontextprotocol/sdk/types.js"; import { Long } from "bson"; -function contentWithTextResourceURI(content: CallToolResult["content"], namespace: string) { +export function contentWithTextResourceURI( + content: CallToolResult["content"] +): CallToolResult["content"][number] | undefined { return content.find((part) => { - return part.type === "text" && part.text.startsWith(`Data for namespace ${namespace}`); + return part.type === "text" && part.text.startsWith(`Data for namespace`); }); } -function contentWithResourceURILink(content: CallToolResult["content"], namespace: string) { +export function contentWithResourceURILink( + content: CallToolResult["content"] +): CallToolResult["content"][number] | undefined { return content.find((part) => { - return part.type === "resource_link" && part.uri.startsWith(`exported-data://${namespace}`); + return part.type === "resource_link"; }); } -function contentWithExportPath(content: CallToolResult["content"]) { +export function contentWithExportPath( + content: CallToolResult["content"] +): CallToolResult["content"][number] | undefined { return content.find((part) => { return ( part.type === "text" && @@ -89,20 +95,22 @@ describeWithMongoDB("export tool", (integration) => { { database: "test", collection: "bar", sort: [], limit: 10 }, ]); - it("when provided with incorrect namespace, export should have empty data", async function () { + beforeEach(async () => { await integration.connectMcpClient(); + }); + + it("when provided with incorrect namespace, export should have empty data", async function () { const response = await integration.mcpClient().callTool({ name: "export", arguments: { database: "non-existent", collection: "foos" }, }); - // Small timeout to let export finish - await timeout(10); - const content = response.content as CallToolResult["content"]; - const namespace = "non-existent.foos"; + const exportURI = contentWithResourceURILink(content)?.uri as string; + await resourceChangedNotification(integration.mcpClient(), exportURI); + expect(content).toHaveLength(3); - expect(contentWithTextResourceURI(content, namespace)).toBeDefined(); - expect(contentWithResourceURILink(content, namespace)).toBeDefined(); + expect(contentWithTextResourceURI(content)).toBeDefined(); + expect(contentWithResourceURILink(content)).toBeDefined(); const localPathPart = contentWithExportPath(content); expect(localPathPart).toBeDefined(); @@ -131,10 +139,11 @@ describeWithMongoDB("export tool", (integration) => { name: "export", arguments: { database: integration.randomDbName(), collection: "foo" }, }); - // Small timeout to let export finish - await timeout(10); + const content = response.content as CallToolResult["content"]; + const exportURI = contentWithResourceURILink(content)?.uri as string; + await resourceChangedNotification(integration.mcpClient(), exportURI); - const localPathPart = contentWithExportPath(response.content as CallToolResult["content"]); + const localPathPart = contentWithExportPath(content); expect(localPathPart).toBeDefined(); const [, localPath] = /"(.*)"/.exec(String(localPathPart?.text)) ?? []; expect(localPath).toBeDefined(); @@ -154,10 +163,11 @@ describeWithMongoDB("export tool", (integration) => { name: "export", arguments: { database: integration.randomDbName(), collection: "foo", filter: { name: "foo" } }, }); - // Small timeout to let export finish - await timeout(10); + const content = response.content as CallToolResult["content"]; + const exportURI = contentWithResourceURILink(content)?.uri as string; + await resourceChangedNotification(integration.mcpClient(), exportURI); - const localPathPart = contentWithExportPath(response.content as CallToolResult["content"]); + const localPathPart = contentWithExportPath(content); expect(localPathPart).toBeDefined(); const [, localPath] = /"(.*)"/.exec(String(localPathPart?.text)) ?? []; expect(localPath).toBeDefined(); @@ -176,10 +186,11 @@ describeWithMongoDB("export tool", (integration) => { name: "export", arguments: { database: integration.randomDbName(), collection: "foo", limit: 1 }, }); - // Small timeout to let export finish - await timeout(10); + const content = response.content as CallToolResult["content"]; + const exportURI = contentWithResourceURILink(content)?.uri as string; + await resourceChangedNotification(integration.mcpClient(), exportURI); - const localPathPart = contentWithExportPath(response.content as CallToolResult["content"]); + const localPathPart = contentWithExportPath(content); expect(localPathPart).toBeDefined(); const [, localPath] = /"(.*)"/.exec(String(localPathPart?.text)) ?? []; expect(localPath).toBeDefined(); @@ -203,10 +214,11 @@ describeWithMongoDB("export tool", (integration) => { sort: { longNumber: 1 }, }, }); - // Small timeout to let export finish - await timeout(10); + const content = response.content as CallToolResult["content"]; + const exportURI = contentWithResourceURILink(content)?.uri as string; + await resourceChangedNotification(integration.mcpClient(), exportURI); - const localPathPart = contentWithExportPath(response.content as CallToolResult["content"]); + const localPathPart = contentWithExportPath(content); expect(localPathPart).toBeDefined(); const [, localPath] = /"(.*)"/.exec(String(localPathPart?.text)) ?? []; expect(localPath).toBeDefined(); @@ -230,10 +242,11 @@ describeWithMongoDB("export tool", (integration) => { projection: { _id: 0, name: 1 }, }, }); - // Small timeout to let export finish - await timeout(10); + const content = response.content as CallToolResult["content"]; + const exportURI = contentWithResourceURILink(content)?.uri as string; + await resourceChangedNotification(integration.mcpClient(), exportURI); - const localPathPart = contentWithExportPath(response.content as CallToolResult["content"]); + const localPathPart = contentWithExportPath(content); expect(localPathPart).toBeDefined(); const [, localPath] = /"(.*)"/.exec(String(localPathPart?.text)) ?? []; expect(localPath).toBeDefined(); @@ -261,10 +274,11 @@ describeWithMongoDB("export tool", (integration) => { jsonExportFormat: "relaxed", }, }); - // Small timeout to let export finish - await timeout(10); + const content = response.content as CallToolResult["content"]; + const exportURI = contentWithResourceURILink(content)?.uri as string; + await resourceChangedNotification(integration.mcpClient(), exportURI); - const localPathPart = contentWithExportPath(response.content as CallToolResult["content"]); + const localPathPart = contentWithExportPath(content); expect(localPathPart).toBeDefined(); const [, localPath] = /"(.*)"/.exec(String(localPathPart?.text)) ?? []; expect(localPath).toBeDefined(); @@ -293,10 +307,11 @@ describeWithMongoDB("export tool", (integration) => { jsonExportFormat: "canonical", }, }); - // Small timeout to let export finish - await timeout(10); + const content = response.content as CallToolResult["content"]; + const exportURI = contentWithResourceURILink(content)?.uri as string; + await resourceChangedNotification(integration.mcpClient(), exportURI); - const localPathPart = contentWithExportPath(response.content as CallToolResult["content"]); + const localPathPart = contentWithExportPath(content); expect(localPathPart).toBeDefined(); const [, localPath] = /"(.*)"/.exec(String(localPathPart?.text)) ?? []; expect(localPath).toBeDefined(); From 816fb0cf056ea2271543be3f7f7dbb1348f661f7 Mon Sep 17 00:00:00 2001 From: Himanshu Singh Date: Thu, 7 Aug 2025 19:21:59 +0200 Subject: [PATCH 28/50] chore: model Resources the same way as Tools --- src/common/sessionExportsManager.ts | 10 ++--- src/resources/common/config.ts | 18 ++++---- src/resources/common/debug.ts | 17 ++++---- src/resources/common/exportedData.ts | 29 +++++++------ src/resources/resource.ts | 42 ++++++++++++------- src/server.ts | 4 +- tests/integration/helpers.ts | 2 +- .../unit/common/sessionExportsManager.test.ts | 22 ++++++---- tests/unit/resources/common/debug.test.ts | 11 ++--- 9 files changed, 89 insertions(+), 66 deletions(-) diff --git a/src/common/sessionExportsManager.ts b/src/common/sessionExportsManager.ts index d100a1f1..f780cf44 100644 --- a/src/common/sessionExportsManager.ts +++ b/src/common/sessionExportsManager.ts @@ -241,11 +241,11 @@ export class SessionExportsManager extends EventEmitter { try { await fs.unlink(exportPath); } catch (error) { @@ -335,6 +335,6 @@ export function validateExportName(nameWithExtension: string): string { return decodedName; } -export function isExportExpired(createdAt: number, exportTimeoutMs: number) { +export function isExportExpired(createdAt: number, exportTimeoutMs: number): boolean { return Date.now() - createdAt > exportTimeoutMs; } diff --git a/src/resources/common/config.ts b/src/resources/common/config.ts index 2bd8a8aa..5a0570d4 100644 --- a/src/resources/common/config.ts +++ b/src/resources/common/config.ts @@ -1,13 +1,12 @@ import { ReactiveResource } from "../resource.js"; -import { config } from "../../common/config.js"; import type { UserConfig } from "../../common/config.js"; -import type { Server } from "../../server.js"; import type { Telemetry } from "../../telemetry/telemetry.js"; +import type { Session } from "../../lib.js"; export class ConfigResource extends ReactiveResource { - constructor(server: Server, telemetry: Telemetry) { - super( - { + constructor(session: Session, config: UserConfig, telemetry: Telemetry) { + super({ + resourceConfiguration: { name: "config", uri: "config://config", config: { @@ -15,13 +14,14 @@ export class ConfigResource extends ReactiveResource { "Server configuration, supplied by the user either as environment variables or as startup arguments", }, }, - { + options: { initial: { ...config }, events: [], }, - server, - telemetry - ); + session, + config, + telemetry, + }); } reduce(eventName: undefined, event: undefined): UserConfig { void eventName; diff --git a/src/resources/common/debug.ts b/src/resources/common/debug.ts index bf6cf5f5..40be3fc0 100644 --- a/src/resources/common/debug.ts +++ b/src/resources/common/debug.ts @@ -1,6 +1,6 @@ import { ReactiveResource } from "../resource.js"; -import type { Server } from "../../server.js"; import type { Telemetry } from "../../telemetry/telemetry.js"; +import { Session, UserConfig } from "../../lib.js"; type ConnectionStateDebuggingInformation = { readonly tag: "connected" | "connecting" | "disconnected" | "errored"; @@ -14,9 +14,9 @@ export class DebugResource extends ReactiveResource< ConnectionStateDebuggingInformation, readonly ["connect", "disconnect", "close", "connection-error"] > { - constructor(server: Server, telemetry: Telemetry) { - super( - { + constructor(session: Session, config: UserConfig, telemetry: Telemetry) { + super({ + resourceConfiguration: { name: "debug-mongodb", uri: "debug://mongodb", config: { @@ -24,13 +24,14 @@ export class DebugResource extends ReactiveResource< "Debugging information for MongoDB connectivity issues. Tracks the last connectivity error and attempt information.", }, }, - { + options: { initial: { tag: "disconnected" }, events: ["connect", "disconnect", "close", "connection-error"], }, - server, - telemetry - ); + session, + config, + telemetry, + }); } reduce( eventName: "connect" | "disconnect" | "close" | "connection-error", diff --git a/src/resources/common/exportedData.ts b/src/resources/common/exportedData.ts index daab4edb..b9744386 100644 --- a/src/resources/common/exportedData.ts +++ b/src/resources/common/exportedData.ts @@ -6,25 +6,28 @@ import { } from "@modelcontextprotocol/sdk/server/mcp.js"; import { Server } from "../../server.js"; import { LogId } from "../../common/logger.js"; +import { Session } from "../../common/session.js"; export class ExportedData { private readonly name = "exported-data"; private readonly description = "Data files exported in the current session."; private readonly uri = "exported-data://{exportName}"; + private server?: Server; - constructor(private readonly server: Server) { - this.server.session.exportsManager.on("export-available", (uri) => { - this.server.mcpServer.sendResourceListChanged(); - void this.server.mcpServer.server.sendResourceUpdated({ + constructor(private readonly session: Session) { + this.session.exportsManager.on("export-available", (uri) => { + this.server?.mcpServer.sendResourceListChanged(); + void this.server?.mcpServer.server.sendResourceUpdated({ uri, }); }); - this.server.session.exportsManager.on("export-expired", () => { - this.server.mcpServer.sendResourceListChanged(); + this.session.exportsManager.on("export-expired", () => { + this.server?.mcpServer.sendResourceListChanged(); }); } - public register(): void { + public register(server: Server): void { + this.server = server; this.server.mcpServer.registerResource( this.name, new ResourceTemplate(this.uri, { @@ -48,7 +51,7 @@ export class ExportedData { private listResourcesCallback: ListResourcesCallback = () => { try { return { - resources: this.server.session.exportsManager.availableExports.map(({ exportName, exportURI }) => ({ + resources: this.session.exportsManager.availableExports.map(({ exportName, exportURI }) => ({ name: exportName, description: this.exportNameToDescription(exportName), uri: exportURI, @@ -56,7 +59,7 @@ export class ExportedData { })), }; } catch (error) { - this.server.session.logger.error({ + this.session.logger.error({ id: LogId.exportedDataListError, context: "Error when listing exported data resources", message: error instanceof Error ? error.message : String(error), @@ -69,11 +72,11 @@ export class ExportedData { private autoCompleteExportName: CompleteResourceTemplateCallback = (value) => { try { - return this.server.session.exportsManager.availableExports + return this.session.exportsManager.availableExports .filter(({ exportName }) => exportName.startsWith(value)) .map(({ exportName }) => exportName); } catch (error) { - this.server.session.logger.error({ + this.session.logger.error({ id: LogId.exportedDataAutoCompleteError, context: "Error when autocompleting exported data", message: error instanceof Error ? error.message : String(error), @@ -88,7 +91,7 @@ export class ExportedData { throw new Error("Cannot retrieve exported data, exportName not provided."); } - const content = await this.server.session.exportsManager.readExport(exportName); + const content = await this.session.exportsManager.readExport(exportName); return { contents: [ @@ -112,7 +115,7 @@ export class ExportedData { } }; - private exportNameToDescription(exportName: string) { + private exportNameToDescription(exportName: string): string { const match = exportName.match(/^(.+)\.(\d+)\.json$/); if (!match) return "Exported data for an unknown namespace."; diff --git a/src/resources/resource.ts b/src/resources/resource.ts index 58ab13b8..c43c8380 100644 --- a/src/resources/resource.ts +++ b/src/resources/resource.ts @@ -20,28 +20,41 @@ export type ReactiveResourceOptions { - protected readonly session: Session; - protected readonly config: UserConfig; + protected server?: Server; + protected session: Session; + protected config: UserConfig; + protected telemetry: Telemetry; + protected current: Value; protected readonly name: string; protected readonly uri: string; protected readonly resourceConfig: ResourceMetadata; protected readonly events: RelevantEvents; - constructor( - resourceConfiguration: ResourceConfiguration, - options: ReactiveResourceOptions, - protected readonly server: Server, - protected readonly telemetry: Telemetry, - current?: Value - ) { + constructor({ + resourceConfiguration, + options, + session, + config, + telemetry, + current, + }: { + resourceConfiguration: ResourceConfiguration; + options: ReactiveResourceOptions; + session: Session; + config: UserConfig; + telemetry: Telemetry; + current?: Value; + }) { + this.session = session; + this.config = config; + this.telemetry = telemetry; + this.name = resourceConfiguration.name; this.uri = resourceConfiguration.uri; this.resourceConfig = resourceConfiguration.config; this.events = options.events; this.current = current ?? options.initial; - this.session = server.session; - this.config = server.userConfig; this.setupEventListeners(); } @@ -55,7 +68,8 @@ export abstract class ReactiveResource { try { - await this.server.mcpServer.server.sendResourceUpdated({ uri: this.uri }); - this.server.mcpServer.sendResourceListChanged(); + await this.server?.mcpServer.server.sendResourceUpdated({ uri: this.uri }); + this.server?.mcpServer.sendResourceListChanged(); } catch (error: unknown) { this.session.logger.warning({ id: LogId.resourceUpdateFailure, diff --git a/src/server.ts b/src/server.ts index de88d7be..e73d87c8 100644 --- a/src/server.ts +++ b/src/server.ts @@ -146,8 +146,8 @@ export class Server { private registerResources(): void { for (const resourceConstructor of Resources) { - const resource = new resourceConstructor(this, this.telemetry); - resource.register(); + const resource = new resourceConstructor(this.session, this.userConfig, this.telemetry); + resource.register(this); } } diff --git a/tests/integration/helpers.ts b/tests/integration/helpers.ts index 7910cdbf..31bb97a0 100644 --- a/tests/integration/helpers.ts +++ b/tests/integration/helpers.ts @@ -271,7 +271,7 @@ function validateToolAnnotations(tool: ToolInfo, name: string, description: stri } } -export function timeout(ms: number) { +export function timeout(ms: number): Promise { return new Promise((resolve) => setTimeout(resolve, ms)); } diff --git a/tests/unit/common/sessionExportsManager.test.ts b/tests/unit/common/sessionExportsManager.test.ts index aad3d9b0..7ac64114 100644 --- a/tests/unit/common/sessionExportsManager.test.ts +++ b/tests/unit/common/sessionExportsManager.test.ts @@ -25,7 +25,15 @@ const exportsManagerConfig: SessionExportsManagerConfig = { exportCleanupIntervalMs: config.exportCleanupIntervalMs, } as const; -function getExportNameAndPath(sessionId: string, timestamp: number) { +function getExportNameAndPath( + sessionId: string, + timestamp: number +): { + sessionExportsPath: string; + exportName: string; + exportPath: string; + exportURI: string; +} { const exportName = `foo.bar.${timestamp}.json`; const sessionExportsPath = path.join(exportsPath, sessionId); const exportPath = path.join(sessionExportsPath, exportName); @@ -44,7 +52,7 @@ function createDummyFindCursor( let index = 0; const readable = new Readable({ objectMode: true, - async read() { + async read(): Promise { if (index < dataArray.length) { if (chunkPushTimeoutMs) { await timeout(chunkPushTimeoutMs); @@ -61,7 +69,7 @@ function createDummyFindCursor( let notifyClose: () => Promise; const cursorCloseNotification = new Promise((resolve) => { - notifyClose = async () => { + notifyClose = async (): Promise => { await timeout(10); resolve(); }; @@ -81,7 +89,7 @@ function createDummyFindCursor( }; } -async function fileExists(filePath: string) { +async function fileExists(filePath: string): Promise { try { await fs.access(filePath); return true; @@ -288,11 +296,11 @@ describe("SessionExportsManager unit test", () => { it("should remove the partial export and never make it available", async () => { const emitSpy = vi.spyOn(manager, "emit"); // eslint-disable-next-line @typescript-eslint/no-unsafe-member-access, @typescript-eslint/no-explicit-any - (manager as any).docToEJSONStream = function (ejsonOptions: EJSONOptions | undefined) { + (manager as any).docToEJSONStream = function (ejsonOptions: EJSONOptions | undefined): Transform { let docsTransformed = 0; return new Transform({ objectMode: true, - transform: function (chunk: unknown, encoding, callback) { + transform: function (chunk: unknown, encoding, callback): void { ++docsTransformed; try { if (docsTransformed === 1) { @@ -306,7 +314,7 @@ describe("SessionExportsManager unit test", () => { callback(err as Error); } }, - final: function (callback) { + final: function (callback): void { this.push("]"); callback(null); }, diff --git a/tests/unit/resources/common/debug.test.ts b/tests/unit/resources/common/debug.test.ts index 4a2f704b..e47e7f67 100644 --- a/tests/unit/resources/common/debug.test.ts +++ b/tests/unit/resources/common/debug.test.ts @@ -1,21 +1,18 @@ import { beforeEach, describe, expect, it } from "vitest"; import { DebugResource } from "../../../../src/resources/common/debug.js"; import { Session } from "../../../../src/common/session.js"; -import { Server } from "../../../../src/server.js"; import { Telemetry } from "../../../../src/telemetry/telemetry.js"; import { config } from "../../../../src/common/config.js"; +import { CompositeLogger } from "../../../../src/common/logger.js"; describe("debug resource", () => { - // eslint-disable-next-line - const session = new Session({} as any); - // eslint-disable-next-line - const server = new Server({ session } as any); + const session = new Session({ apiBaseUrl: "", logger: new CompositeLogger() }); const telemetry = Telemetry.create(session, { ...config, telemetry: "disabled" }); - let debugResource: DebugResource = new DebugResource(server, telemetry); + let debugResource: DebugResource = new DebugResource(session, config, telemetry); beforeEach(() => { - debugResource = new DebugResource(server, telemetry); + debugResource = new DebugResource(session, config, telemetry); }); it("should be connected when a connected event happens", () => { From e0cc924a1badc95f56f667039e3e7ceb88423946 Mon Sep 17 00:00:00 2001 From: Himanshu Singh Date: Thu, 7 Aug 2025 20:29:44 +0200 Subject: [PATCH 29/50] chore: implements PR feedback 1. SessionExportManager directly injected in Session same as ConnectionManager 2. sessionId is not a property anymore on SessionExportManager 3. logs are enriched with export information 4. nested try/catch removed from SessionExportManager.startExport 5. Promise continuation removed by delegating to SessionExportManager.silentyRemoveExport 6. switch/case for getEJSONOptionsForFormat --- src/common/session.ts | 26 +++--- src/common/sessionExportsManager.ts | 92 +++++++++---------- src/transports/base.ts | 13 ++- tests/integration/helpers.ts | 10 +- tests/integration/telemetry.test.ts | 6 ++ tests/unit/common/session.test.ts | 8 +- .../unit/common/sessionExportsManager.test.ts | 13 ++- tests/unit/resources/common/debug.test.ts | 11 ++- 8 files changed, 112 insertions(+), 67 deletions(-) diff --git a/src/common/session.ts b/src/common/session.ts index 8e6895df..5cdb2690 100644 --- a/src/common/session.ts +++ b/src/common/session.ts @@ -10,17 +10,16 @@ import { } from "./connectionManager.js"; import { NodeDriverServiceProvider } from "@mongosh/service-provider-node-driver"; import { ErrorCodes, MongoDBError } from "./errors.js"; -import { ObjectId } from "bson"; -import { SessionExportsManager, SessionExportsManagerConfig } from "./sessionExportsManager.js"; -import { config } from "./config.js"; +import { SessionExportsManager } from "./sessionExportsManager.js"; export interface SessionOptions { apiBaseUrl: string; apiClientId?: string; apiClientSecret?: string; - connectionManager?: ConnectionManager; - exportsManagerConfig?: SessionExportsManagerConfig; logger: CompositeLogger; + sessionId: string; + exportsManager: SessionExportsManager; + connectionManager: ConnectionManager; } export type SessionEvents = { @@ -31,10 +30,10 @@ export type SessionEvents = { }; export class Session extends EventEmitter { - readonly sessionId = new ObjectId().toString(); + readonly sessionId: string; readonly exportsManager: SessionExportsManager; - connectionManager: ConnectionManager; - apiClient: ApiClient; + readonly connectionManager: ConnectionManager; + readonly apiClient: ApiClient; agentRunner?: { name: string; version: string; @@ -46,14 +45,14 @@ export class Session extends EventEmitter { apiBaseUrl, apiClientId, apiClientSecret, - connectionManager, logger, - exportsManagerConfig, + sessionId, + connectionManager, + exportsManager, }: SessionOptions) { super(); this.logger = logger; - const credentials: ApiClientCredentials | undefined = apiClientId && apiClientSecret ? { @@ -63,9 +62,10 @@ export class Session extends EventEmitter { : undefined; this.apiClient = new ApiClient({ baseUrl: apiBaseUrl, credentials }, logger); - this.exportsManager = new SessionExportsManager(this.sessionId, exportsManagerConfig ?? config, logger); - this.connectionManager = connectionManager ?? new ConnectionManager(); + this.sessionId = sessionId; + this.exportsManager = exportsManager; + this.connectionManager = connectionManager; this.connectionManager.on("connection-succeeded", () => this.emit("connect")); this.connectionManager.on("connection-timed-out", (error) => this.emit("connection-error", error.errorReason)); this.connectionManager.on("connection-closed", () => this.emit("disconnect")); diff --git a/src/common/sessionExportsManager.ts b/src/common/sessionExportsManager.ts index f780cf44..d4e9f27c 100644 --- a/src/common/sessionExportsManager.ts +++ b/src/common/sessionExportsManager.ts @@ -10,6 +10,7 @@ import { pipeline } from "stream/promises"; import { UserConfig } from "./config.js"; import { LoggerBase, LogId } from "./logger.js"; +import { MongoLogId } from "mongodb-log-writer"; export const jsonExportFormat = z.enum(["relaxed", "canonical"]); export type JSONExportFormat = z.infer; @@ -71,7 +72,7 @@ export class SessionExportsManager extends EventEmitter { + let pipeSuccessful = false; try { await fs.mkdir(this.exportsDirectoryPath, { recursive: true }); - const inputStream = input.stream(); - const ejsonDocStream = this.docToEJSONStream(this.getEJSONOptionsForFormat(jsonExportFormat)); const outputStream = createWriteStream(inProgressExport.exportPath); outputStream.write("["); - let pipeSuccessful = false; - try { - await pipeline([inputStream, ejsonDocStream, outputStream]); - pipeSuccessful = true; - } catch (pipelineError) { - // If the pipeline errors out then we might end up with - // partial and incorrect export so we remove it entirely. - await fs.unlink(inProgressExport.exportPath).catch((error) => { - if ((error as NodeJS.ErrnoException).code !== "ENOENT") { - this.logger.error({ - id: LogId.exportCreationCleanupError, - context: "Error when removing partial export", - message: error instanceof Error ? error.message : String(error), - }); - } - }); - delete this.sessionExports[inProgressExport.exportName]; - throw pipelineError; - } finally { - if (pipeSuccessful) { - this.sessionExports[inProgressExport.exportName] = { - ...inProgressExport, - exportCreatedAt: Date.now(), - exportStatus: "ready", - }; - this.emit("export-available", inProgressExport.exportURI); - } - void input.close(); - } + await pipeline([ + input.stream(), + this.docToEJSONStream(this.getEJSONOptionsForFormat(jsonExportFormat)), + outputStream, + ]); + pipeSuccessful = true; } catch (error) { this.logger.error({ id: LogId.exportCreationError, - context: "Error when generating JSON export", + context: `Error when generating JSON export for ${inProgressExport.exportName}`, message: error instanceof Error ? error.message : String(error), }); + + // If the pipeline errors out then we might end up with + // partial and incorrect export so we remove it entirely. + await this.silentlyRemoveExport( + inProgressExport.exportPath, + LogId.exportCreationCleanupError, + `Error when removing incomplete export ${inProgressExport.exportName}` + ); + delete this.sessionExports[inProgressExport.exportName]; + } finally { + if (pipeSuccessful) { + this.sessionExports[inProgressExport.exportName] = { + ...inProgressExport, + exportCreatedAt: Date.now(), + exportStatus: "ready", + }; + this.emit("export-available", inProgressExport.exportURI); + } + void input.close(); } } private getEJSONOptionsForFormat(format: JSONExportFormat): EJSONOptions | undefined { - if (format === "relaxed") { - return { - relaxed: true, - }; + switch (format) { + case "relaxed": + return { relaxed: true }; + case "canonical": + return { relaxed: false }; + default: + return undefined; } - return format === "canonical" - ? { - relaxed: false, - } - : undefined; } private docToEJSONStream(ejsonOptions: EJSONOptions | undefined): Transform { @@ -276,7 +270,11 @@ export class SessionExportsManager extends EventEmitter { + private async silentlyRemoveExport(exportPath: string, logId: MongoLogId, logContext: string): Promise { try { await fs.unlink(exportPath); } catch (error) { @@ -300,8 +298,8 @@ export class SessionExportsManager extends EventEmitter UserConfig): Integrati } ); + const logger = new CompositeLogger(); + const sessionId = new ObjectId().toString(); + const exportsManager = new SessionExportsManager(sessionId, userConfig, logger); const connectionManager = new ConnectionManager(); const session = new Session({ apiBaseUrl: userConfig.apiBaseUrl, apiClientId: userConfig.apiClientId, apiClientSecret: userConfig.apiClientSecret, + logger, + sessionId, + exportsManager, connectionManager, - logger: new CompositeLogger(), - exportsManagerConfig: userConfig, }); // Mock hasValidAccessToken for tests diff --git a/tests/integration/telemetry.test.ts b/tests/integration/telemetry.test.ts index 9be1b0ad..bbd537d2 100644 --- a/tests/integration/telemetry.test.ts +++ b/tests/integration/telemetry.test.ts @@ -5,6 +5,8 @@ import { config } from "../../src/common/config.js"; import nodeMachineId from "node-machine-id"; import { describe, expect, it } from "vitest"; import { CompositeLogger } from "../../src/common/logger.js"; +import { ConnectionManager } from "../../src/common/connectionManager.js"; +import { SessionExportsManager } from "../../src/common/sessionExportsManager.js"; describe("Telemetry", () => { it("should resolve the actual machine ID", async () => { @@ -12,10 +14,14 @@ describe("Telemetry", () => { const actualHashedId = createHmac("sha256", actualId.toUpperCase()).update("atlascli").digest("hex"); + const logger = new CompositeLogger(); const telemetry = Telemetry.create( new Session({ apiBaseUrl: "", logger: new CompositeLogger(), + sessionId: "1FOO", + exportsManager: new SessionExportsManager("1FOO", config, logger), + connectionManager: new ConnectionManager(), }), config ); diff --git a/tests/unit/common/session.test.ts b/tests/unit/common/session.test.ts index 592d60fe..dd34a010 100644 --- a/tests/unit/common/session.test.ts +++ b/tests/unit/common/session.test.ts @@ -3,6 +3,8 @@ import { NodeDriverServiceProvider } from "@mongosh/service-provider-node-driver import { Session } from "../../../src/common/session.js"; import { config } from "../../../src/common/config.js"; import { CompositeLogger } from "../../../src/common/logger.js"; +import { ConnectionManager } from "../../../src/common/connectionManager.js"; +import { SessionExportsManager } from "../../../src/common/sessionExportsManager.js"; vi.mock("@mongosh/service-provider-node-driver"); const MockNodeDriverServiceProvider = vi.mocked(NodeDriverServiceProvider); @@ -10,10 +12,14 @@ const MockNodeDriverServiceProvider = vi.mocked(NodeDriverServiceProvider); describe("Session", () => { let session: Session; beforeEach(() => { + const logger = new CompositeLogger(); session = new Session({ apiClientId: "test-client-id", apiBaseUrl: "https://api.test.com", - logger: new CompositeLogger(), + logger, + sessionId: "1FOO", + exportsManager: new SessionExportsManager("1FOO", config, logger), + connectionManager: new ConnectionManager(), }); MockNodeDriverServiceProvider.connect = vi.fn().mockResolvedValue({} as unknown as NodeDriverServiceProvider); diff --git a/tests/unit/common/sessionExportsManager.test.ts b/tests/unit/common/sessionExportsManager.test.ts index 7ac64114..613b22a4 100644 --- a/tests/unit/common/sessionExportsManager.test.ts +++ b/tests/unit/common/sessionExportsManager.test.ts @@ -15,8 +15,9 @@ import { config } from "../../../src/common/config.js"; import { Session } from "../../../src/common/session.js"; import { ROOT_DIR } from "../../accuracy/sdk/constants.js"; import { timeout } from "../../integration/helpers.js"; -import { EJSON, EJSONOptions } from "bson"; +import { EJSON, EJSONOptions, ObjectId } from "bson"; import { CompositeLogger } from "../../../src/common/logger.js"; +import { ConnectionManager } from "../../../src/common/connectionManager.js"; const exportsPath = path.join(ROOT_DIR, "tests", "tmp", "exports"); const exportsManagerConfig: SessionExportsManagerConfig = { @@ -106,7 +107,15 @@ describe("SessionExportsManager unit test", () => { await manager?.close(); await fs.rm(exportsManagerConfig.exportsPath, { recursive: true, force: true }); await fs.mkdir(exportsManagerConfig.exportsPath, { recursive: true }); - session = new Session({ apiBaseUrl: "", logger: new CompositeLogger() }); + const logger = new CompositeLogger(); + const sessionId = new ObjectId().toString(); + session = new Session({ + apiBaseUrl: "", + logger, + sessionId, + exportsManager: new SessionExportsManager(sessionId, config, logger), + connectionManager: new ConnectionManager(), + }); manager = session.exportsManager; }); diff --git a/tests/unit/resources/common/debug.test.ts b/tests/unit/resources/common/debug.test.ts index e47e7f67..0a3d8d32 100644 --- a/tests/unit/resources/common/debug.test.ts +++ b/tests/unit/resources/common/debug.test.ts @@ -4,9 +4,18 @@ import { Session } from "../../../../src/common/session.js"; import { Telemetry } from "../../../../src/telemetry/telemetry.js"; import { config } from "../../../../src/common/config.js"; import { CompositeLogger } from "../../../../src/common/logger.js"; +import { ConnectionManager } from "../../../../src/common/connectionManager.js"; +import { SessionExportsManager } from "../../../../src/common/sessionExportsManager.js"; describe("debug resource", () => { - const session = new Session({ apiBaseUrl: "", logger: new CompositeLogger() }); + const logger = new CompositeLogger(); + const session = new Session({ + apiBaseUrl: "", + logger, + sessionId: "1FOO", + exportsManager: new SessionExportsManager("1FOO", config, logger), + connectionManager: new ConnectionManager(), + }); const telemetry = Telemetry.create(session, { ...config, telemetry: "disabled" }); let debugResource: DebugResource = new DebugResource(session, config, telemetry); From 397e4600a0c10f840473681f7b3c66402b35cc4d Mon Sep 17 00:00:00 2001 From: Himanshu Singh Date: Thu, 7 Aug 2025 20:38:42 +0200 Subject: [PATCH 30/50] chore rename SessionExportsManager to ExportsManager --- ...ionExportsManager.ts => exportsManager.ts} | 37 +++++++++---------- src/common/session.ts | 6 +-- src/tools/mongodb/read/export.ts | 2 +- src/transports/base.ts | 4 +- tests/integration/helpers.ts | 4 +- tests/integration/telemetry.test.ts | 4 +- ...Manager.test.ts => exportsManager.test.ts} | 22 +++++------ tests/unit/common/session.test.ts | 4 +- tests/unit/resources/common/debug.test.ts | 4 +- 9 files changed, 42 insertions(+), 45 deletions(-) rename src/common/{sessionExportsManager.ts => exportsManager.ts} (90%) rename tests/unit/common/{sessionExportsManager.test.ts => exportsManager.test.ts} (96%) diff --git a/src/common/sessionExportsManager.ts b/src/common/exportsManager.ts similarity index 90% rename from src/common/sessionExportsManager.ts rename to src/common/exportsManager.ts index d4e9f27c..d5f4dd7a 100644 --- a/src/common/sessionExportsManager.ts +++ b/src/common/exportsManager.ts @@ -55,25 +55,22 @@ type StoredExport = ReadyExport | InProgressExport; * JIRA: https://jira.mongodb.org/browse/MCP-104 */ type AvailableExport = Pick; -export type SessionExportsManagerConfig = Pick< - UserConfig, - "exportsPath" | "exportTimeoutMs" | "exportCleanupIntervalMs" ->; +export type ExportsManagerConfig = Pick; -type SessionExportsManagerEvents = { +type ExportsManagerEvents = { "export-expired": [string]; "export-available": [string]; }; -export class SessionExportsManager extends EventEmitter { - private sessionExports: Record = {}; +export class ExportsManager extends EventEmitter { + private storedExports: Record = {}; private exportsCleanupInProgress: boolean = false; private exportsCleanupInterval: NodeJS.Timeout; private exportsDirectoryPath: string; constructor( sessionId: string, - private readonly config: SessionExportsManagerConfig, + private readonly config: ExportsManagerConfig, private readonly logger: LoggerBase ) { super(); @@ -85,11 +82,11 @@ export class SessionExportsManager extends EventEmitter { + return Object.values(this.storedExports) + .filter((storedExport) => { return ( - sessionExport.exportStatus === "ready" && - !isExportExpired(sessionExport.exportCreatedAt, this.config.exportTimeoutMs) + storedExport.exportStatus === "ready" && + !isExportExpired(storedExport.exportCreatedAt, this.config.exportTimeoutMs) ); }) .map(({ exportName, exportURI, exportPath }) => ({ @@ -106,7 +103,7 @@ export class SessionExportsManager extends EventEmitter { try { const exportNameWithExtension = validateExportName(exportName); - const exportHandle = this.sessionExports[exportNameWithExtension]; + const exportHandle = this.storedExports[exportNameWithExtension]; if (!exportHandle) { throw new Error("Requested export has either expired or does not exist!"); } @@ -157,7 +154,7 @@ export class SessionExportsManager extends EventEmitter sessionExport.exportStatus === "ready" + const exportsForCleanup = Object.values({ ...this.storedExports }).filter( + (storedExport): storedExport is ReadyExport => storedExport.exportStatus === "ready" ); try { for (const { exportPath, exportCreatedAt, exportURI, exportName } of exportsForCleanup) { if (isExportExpired(exportCreatedAt, this.config.exportTimeoutMs)) { - delete this.sessionExports[exportName]; + delete this.storedExports[exportName]; await this.silentlyRemoveExport( exportPath, LogId.exportCleanupError, diff --git a/src/common/session.ts b/src/common/session.ts index 5cdb2690..032375fd 100644 --- a/src/common/session.ts +++ b/src/common/session.ts @@ -10,7 +10,7 @@ import { } from "./connectionManager.js"; import { NodeDriverServiceProvider } from "@mongosh/service-provider-node-driver"; import { ErrorCodes, MongoDBError } from "./errors.js"; -import { SessionExportsManager } from "./sessionExportsManager.js"; +import { ExportsManager } from "./exportsManager.js"; export interface SessionOptions { apiBaseUrl: string; @@ -18,7 +18,7 @@ export interface SessionOptions { apiClientSecret?: string; logger: CompositeLogger; sessionId: string; - exportsManager: SessionExportsManager; + exportsManager: ExportsManager; connectionManager: ConnectionManager; } @@ -31,7 +31,7 @@ export type SessionEvents = { export class Session extends EventEmitter { readonly sessionId: string; - readonly exportsManager: SessionExportsManager; + readonly exportsManager: ExportsManager; readonly connectionManager: ConnectionManager; readonly apiClient: ApiClient; agentRunner?: { diff --git a/src/tools/mongodb/read/export.ts b/src/tools/mongodb/read/export.ts index 8e6be9dc..8a147345 100644 --- a/src/tools/mongodb/read/export.ts +++ b/src/tools/mongodb/read/export.ts @@ -2,7 +2,7 @@ import { CallToolResult } from "@modelcontextprotocol/sdk/types.js"; import { OperationType, ToolArgs } from "../../tool.js"; import { DbOperationArgs, MongoDBToolBase } from "../mongodbTool.js"; import { FindArgs } from "./find.js"; -import { jsonExportFormat } from "../../../common/sessionExportsManager.js"; +import { jsonExportFormat } from "../../../common/exportsManager.js"; import z from "zod"; export class ExportTool extends MongoDBToolBase { diff --git a/src/transports/base.ts b/src/transports/base.ts index d9e1f54b..9f2bca0d 100644 --- a/src/transports/base.ts +++ b/src/transports/base.ts @@ -6,7 +6,7 @@ import { Telemetry } from "../telemetry/telemetry.js"; import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js"; import { CompositeLogger, ConsoleLogger, DiskLogger, LoggerBase, McpLogger } from "../common/logger.js"; import { ObjectId } from "bson"; -import { SessionExportsManager } from "../common/sessionExportsManager.js"; +import { ExportsManager } from "../common/exportsManager.js"; import { ConnectionManager } from "../common/connectionManager.js"; export abstract class TransportRunnerBase { @@ -44,7 +44,7 @@ export abstract class TransportRunnerBase { const logger = new CompositeLogger(...loggers); const sessionId = new ObjectId().toString(); - const exportsManager = new SessionExportsManager(sessionId, userConfig, logger); + const exportsManager = new ExportsManager(sessionId, userConfig, logger); const connectionManager = new ConnectionManager(); const session = new Session({ diff --git a/tests/integration/helpers.ts b/tests/integration/helpers.ts index 542028fa..949c7974 100644 --- a/tests/integration/helpers.ts +++ b/tests/integration/helpers.ts @@ -11,7 +11,7 @@ import { config } from "../../src/common/config.js"; import { afterAll, afterEach, beforeAll, describe, expect, it, vi } from "vitest"; import { ConnectionManager } from "../../src/common/connectionManager.js"; import { CompositeLogger } from "../../src/common/logger.js"; -import { SessionExportsManager } from "../../src/common/sessionExportsManager.js"; +import { ExportsManager } from "../../src/common/exportsManager.js"; interface ParameterInfo { name: string; @@ -59,7 +59,7 @@ export function setupIntegrationTest(getUserConfig: () => UserConfig): Integrati const logger = new CompositeLogger(); const sessionId = new ObjectId().toString(); - const exportsManager = new SessionExportsManager(sessionId, userConfig, logger); + const exportsManager = new ExportsManager(sessionId, userConfig, logger); const connectionManager = new ConnectionManager(); const session = new Session({ diff --git a/tests/integration/telemetry.test.ts b/tests/integration/telemetry.test.ts index bbd537d2..c5547ac6 100644 --- a/tests/integration/telemetry.test.ts +++ b/tests/integration/telemetry.test.ts @@ -6,7 +6,7 @@ import nodeMachineId from "node-machine-id"; import { describe, expect, it } from "vitest"; import { CompositeLogger } from "../../src/common/logger.js"; import { ConnectionManager } from "../../src/common/connectionManager.js"; -import { SessionExportsManager } from "../../src/common/sessionExportsManager.js"; +import { ExportsManager } from "../../src/common/exportsManager.js"; describe("Telemetry", () => { it("should resolve the actual machine ID", async () => { @@ -20,7 +20,7 @@ describe("Telemetry", () => { apiBaseUrl: "", logger: new CompositeLogger(), sessionId: "1FOO", - exportsManager: new SessionExportsManager("1FOO", config, logger), + exportsManager: new ExportsManager("1FOO", config, logger), connectionManager: new ConnectionManager(), }), config diff --git a/tests/unit/common/sessionExportsManager.test.ts b/tests/unit/common/exportsManager.test.ts similarity index 96% rename from tests/unit/common/sessionExportsManager.test.ts rename to tests/unit/common/exportsManager.test.ts index 613b22a4..e03d2114 100644 --- a/tests/unit/common/sessionExportsManager.test.ts +++ b/tests/unit/common/exportsManager.test.ts @@ -6,10 +6,10 @@ import { beforeEach, describe, expect, it, vi } from "vitest"; import { ensureExtension, isExportExpired, - SessionExportsManager, - SessionExportsManagerConfig, + ExportsManager, + ExportsManagerConfig, validateExportName, -} from "../../../src/common/sessionExportsManager.js"; +} from "../../../src/common/exportsManager.js"; import { config } from "../../../src/common/config.js"; import { Session } from "../../../src/common/session.js"; @@ -20,7 +20,7 @@ import { CompositeLogger } from "../../../src/common/logger.js"; import { ConnectionManager } from "../../../src/common/connectionManager.js"; const exportsPath = path.join(ROOT_DIR, "tests", "tmp", "exports"); -const exportsManagerConfig: SessionExportsManagerConfig = { +const exportsManagerConfig: ExportsManagerConfig = { exportsPath, exportTimeoutMs: config.exportTimeoutMs, exportCleanupIntervalMs: config.exportCleanupIntervalMs, @@ -99,9 +99,9 @@ async function fileExists(filePath: string): Promise { } } -describe("SessionExportsManager unit test", () => { +describe("ExportsManager unit test", () => { let session: Session; - let manager: SessionExportsManager; + let manager: ExportsManager; beforeEach(async () => { await manager?.close(); @@ -113,7 +113,7 @@ describe("SessionExportsManager unit test", () => { apiBaseUrl: "", logger, sessionId, - exportsManager: new SessionExportsManager(sessionId, config, logger), + exportsManager: new ExportsManager(sessionId, config, logger), connectionManager: new ConnectionManager(), }); manager = session.exportsManager; @@ -366,7 +366,7 @@ describe("SessionExportsManager unit test", () => { it("should not clean up in-progress exports", async () => { const { exportName } = getExportNameAndPath(session.sessionId, Date.now()); - const manager = new SessionExportsManager( + const manager = new ExportsManager( session.sessionId, { ...exportsManagerConfig, @@ -383,17 +383,17 @@ describe("SessionExportsManager unit test", () => { }); // eslint-disable-next-line @typescript-eslint/no-explicit-any, @typescript-eslint/no-unsafe-member-access - expect((manager as any).sessionExports[exportName]?.exportStatus).toEqual("in-progress"); + expect((manager as any).storedExports[exportName]?.exportStatus).toEqual("in-progress"); // After clean up interval the export should still be there await timeout(200); // eslint-disable-next-line @typescript-eslint/no-explicit-any, @typescript-eslint/no-unsafe-member-access - expect((manager as any).sessionExports[exportName]?.exportStatus).toEqual("in-progress"); + expect((manager as any).storedExports[exportName]?.exportStatus).toEqual("in-progress"); }); it("should cleanup expired exports", async () => { const { exportName, exportPath, exportURI } = getExportNameAndPath(session.sessionId, Date.now()); - const manager = new SessionExportsManager( + const manager = new ExportsManager( session.sessionId, { ...exportsManagerConfig, diff --git a/tests/unit/common/session.test.ts b/tests/unit/common/session.test.ts index dd34a010..c860e97e 100644 --- a/tests/unit/common/session.test.ts +++ b/tests/unit/common/session.test.ts @@ -4,7 +4,7 @@ import { Session } from "../../../src/common/session.js"; import { config } from "../../../src/common/config.js"; import { CompositeLogger } from "../../../src/common/logger.js"; import { ConnectionManager } from "../../../src/common/connectionManager.js"; -import { SessionExportsManager } from "../../../src/common/sessionExportsManager.js"; +import { ExportsManager } from "../../../src/common/exportsManager.js"; vi.mock("@mongosh/service-provider-node-driver"); const MockNodeDriverServiceProvider = vi.mocked(NodeDriverServiceProvider); @@ -18,7 +18,7 @@ describe("Session", () => { apiBaseUrl: "https://api.test.com", logger, sessionId: "1FOO", - exportsManager: new SessionExportsManager("1FOO", config, logger), + exportsManager: new ExportsManager("1FOO", config, logger), connectionManager: new ConnectionManager(), }); diff --git a/tests/unit/resources/common/debug.test.ts b/tests/unit/resources/common/debug.test.ts index 0a3d8d32..997be48c 100644 --- a/tests/unit/resources/common/debug.test.ts +++ b/tests/unit/resources/common/debug.test.ts @@ -5,7 +5,7 @@ import { Telemetry } from "../../../../src/telemetry/telemetry.js"; import { config } from "../../../../src/common/config.js"; import { CompositeLogger } from "../../../../src/common/logger.js"; import { ConnectionManager } from "../../../../src/common/connectionManager.js"; -import { SessionExportsManager } from "../../../../src/common/sessionExportsManager.js"; +import { ExportsManager } from "../../../../src/common/exportsManager.js"; describe("debug resource", () => { const logger = new CompositeLogger(); @@ -13,7 +13,7 @@ describe("debug resource", () => { apiBaseUrl: "", logger, sessionId: "1FOO", - exportsManager: new SessionExportsManager("1FOO", config, logger), + exportsManager: new ExportsManager("1FOO", config, logger), connectionManager: new ConnectionManager(), }); const telemetry = Telemetry.create(session, { ...config, telemetry: "disabled" }); From 04cb6e71898c583c854a417dc8f3f17c28d8a35c Mon Sep 17 00:00:00 2001 From: Himanshu Singh Date: Thu, 7 Aug 2025 20:53:01 +0200 Subject: [PATCH 31/50] chore: remove validation from readExport but keep decodeURIComponent --- src/common/exportsManager.ts | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/common/exportsManager.ts b/src/common/exportsManager.ts index d5f4dd7a..b20ed3f1 100644 --- a/src/common/exportsManager.ts +++ b/src/common/exportsManager.ts @@ -111,8 +111,8 @@ export class ExportsManager extends EventEmitter { public async readExport(exportName: string): Promise { try { - const exportNameWithExtension = validateExportName(exportName); - const exportHandle = this.storedExports[exportNameWithExtension]; + exportName = decodeURIComponent(exportName); + const exportHandle = this.storedExports[exportName]; if (!exportHandle) { throw new Error("Requested export has either expired or does not exist!"); } From 1df2690f8568e22dca8e0316ee673c8c4e4386c1 Mon Sep 17 00:00:00 2001 From: Himanshu Singh Date: Thu, 7 Aug 2025 22:34:04 +0200 Subject: [PATCH 32/50] chore: publish resource updated notification only when there are subscriptions --- src/resources/common/exportedData.ts | 15 +++++------ src/resources/resource.ts | 6 ++--- src/server.ts | 40 ++++++++++++++++++++++++++-- tests/integration/helpers.ts | 1 + 4 files changed, 48 insertions(+), 14 deletions(-) diff --git a/src/resources/common/exportedData.ts b/src/resources/common/exportedData.ts index b9744386..3ec685ae 100644 --- a/src/resources/common/exportedData.ts +++ b/src/resources/common/exportedData.ts @@ -15,15 +15,12 @@ export class ExportedData { private server?: Server; constructor(private readonly session: Session) { - this.session.exportsManager.on("export-available", (uri) => { - this.server?.mcpServer.sendResourceListChanged(); - void this.server?.mcpServer.server.sendResourceUpdated({ - uri, - }); - }); - this.session.exportsManager.on("export-expired", () => { - this.server?.mcpServer.sendResourceListChanged(); - }); + const onExportChanged = (uri: string): void => { + this.server?.sendResourceListChanged(); + this.server?.sendResourceUpdated(uri); + }; + this.session.exportsManager.on("export-available", onExportChanged); + this.session.exportsManager.on("export-expired", onExportChanged); } public register(server: Server): void { diff --git a/src/resources/resource.ts b/src/resources/resource.ts index c43c8380..005e60c8 100644 --- a/src/resources/resource.ts +++ b/src/resources/resource.ts @@ -83,10 +83,10 @@ export abstract class ReactiveResource { + private triggerUpdate(): void { try { - await this.server?.mcpServer.server.sendResourceUpdated({ uri: this.uri }); - this.server?.mcpServer.sendResourceListChanged(); + this.server?.sendResourceListChanged(); + this.server?.sendResourceUpdated(this.uri); } catch (error: unknown) { this.session.logger.warning({ id: LogId.resourceUpdateFailure, diff --git a/src/server.ts b/src/server.ts index e73d87c8..cc8d30dd 100644 --- a/src/server.ts +++ b/src/server.ts @@ -9,7 +9,12 @@ import { Telemetry } from "./telemetry/telemetry.js"; import { UserConfig } from "./common/config.js"; import { type ServerEvent } from "./telemetry/types.js"; import { type ServerCommand } from "./telemetry/types.js"; -import { CallToolRequestSchema, CallToolResult } from "@modelcontextprotocol/sdk/types.js"; +import { + CallToolRequestSchema, + CallToolResult, + SubscribeRequestSchema, + UnsubscribeRequestSchema, +} from "@modelcontextprotocol/sdk/types.js"; import assert from "assert"; import { ToolBase } from "./tools/tool.js"; @@ -27,6 +32,7 @@ export class Server { public readonly userConfig: UserConfig; public readonly tools: ToolBase[] = []; private readonly startTime: number; + private readonly subscriptions = new Set(); constructor({ session, mcpServer, userConfig, telemetry }: ServerOptions) { this.startTime = Date.now(); @@ -42,7 +48,7 @@ export class Server { this.registerResources(); await this.validateConfig(); - this.mcpServer.server.registerCapabilities({ logging: {}, resources: { listChanged: true } }); + this.mcpServer.server.registerCapabilities({ logging: {}, resources: { listChanged: true, subscribe: true } }); // TODO: Eventually we might want to make tools reactive too instead of relying on custom logic. this.registerTools(); @@ -70,6 +76,26 @@ export class Server { return existingHandler(request, extra); }); + this.mcpServer.server.setRequestHandler(SubscribeRequestSchema, ({ params }) => { + this.subscriptions.add(params.uri); + this.session.logger.debug({ + id: LogId.serverInitialized, + context: "resources", + message: `Client subscribed to resource: ${params.uri}`, + }); + return {}; + }); + + this.mcpServer.server.setRequestHandler(UnsubscribeRequestSchema, ({ params }) => { + this.subscriptions.delete(params.uri); + this.session.logger.debug({ + id: LogId.serverInitialized, + context: "resources", + message: `Client unsubscribed from resource: ${params.uri}`, + }); + return {}; + }); + this.mcpServer.server.oninitialized = (): void => { this.session.setAgentRunner(this.mcpServer.server.getClientVersion()); @@ -101,6 +127,16 @@ export class Server { await this.mcpServer.close(); } + public sendResourceListChanged(): void { + this.mcpServer.sendResourceListChanged(); + } + + public sendResourceUpdated(uri: string): void { + if (this.subscriptions.has(uri)) { + void this.mcpServer.server.sendResourceUpdated({ uri }); + } + } + /** * Emits a server event * @param command - The server command (e.g., "start", "stop", "register", "deregister") diff --git a/tests/integration/helpers.ts b/tests/integration/helpers.ts index 949c7974..dea7f347 100644 --- a/tests/integration/helpers.ts +++ b/tests/integration/helpers.ts @@ -286,6 +286,7 @@ export function timeout(ms: number): Promise { */ export function resourceChangedNotification(client: Client, uri: string): Promise { return new Promise((resolve) => { + void client.subscribeResource({ uri }); client.setNotificationHandler(ResourceUpdatedNotificationSchema, (notification) => { if (notification.params.uri === uri) { resolve(); From d3d81d6e8f8fca100a16c610a15e97fcfdd21446 Mon Sep 17 00:00:00 2001 From: Himanshu Singh Date: Fri, 8 Aug 2025 12:10:52 +0200 Subject: [PATCH 33/50] chore: remove sideeffect from ctor --- src/common/exportsManager.ts | 31 +- src/transports/base.ts | 2 +- tests/integration/helpers.ts | 2 +- .../resources/exportedData.test.ts | 24 +- tests/integration/telemetry.test.ts | 2 +- .../tools/mongodb/read/export.test.ts | 535 +++++++++--------- tests/unit/common/exportsManager.test.ts | 17 +- tests/unit/common/session.test.ts | 2 +- tests/unit/resources/common/debug.test.ts | 2 +- 9 files changed, 329 insertions(+), 288 deletions(-) diff --git a/src/common/exportsManager.ts b/src/common/exportsManager.ts index b20ed3f1..84360b05 100644 --- a/src/common/exportsManager.ts +++ b/src/common/exportsManager.ts @@ -63,22 +63,17 @@ type ExportsManagerEvents = { }; export class ExportsManager extends EventEmitter { + private wasInitialized: boolean = false; private storedExports: Record = {}; private exportsCleanupInProgress: boolean = false; - private exportsCleanupInterval: NodeJS.Timeout; - private exportsDirectoryPath: string; + private exportsCleanupInterval?: NodeJS.Timeout; - constructor( - sessionId: string, + private constructor( + private readonly exportsDirectoryPath: string, private readonly config: ExportsManagerConfig, private readonly logger: LoggerBase ) { super(); - this.exportsDirectoryPath = path.join(this.config.exportsPath, sessionId); - this.exportsCleanupInterval = setInterval( - () => void this.cleanupExpiredExports(), - this.config.exportCleanupIntervalMs - ); } public get availableExports(): AvailableExport[] { @@ -96,6 +91,17 @@ export class ExportsManager extends EventEmitter { })); } + protected init(): void { + if (!this.wasInitialized) { + this.exportsCleanupInterval = setInterval( + () => void this.cleanupExpiredExports(), + this.config.exportCleanupIntervalMs + ); + + this.wasInitialized = true; + } + } + public async close(): Promise { try { clearInterval(this.exportsCleanupInterval); @@ -302,6 +308,13 @@ export class ExportsManager extends EventEmitter { } } } + + static init(sessionId: string, config: ExportsManagerConfig, logger: LoggerBase): ExportsManager { + const exportsDirectoryPath = path.join(config.exportsPath, sessionId); + const exportsManager = new ExportsManager(exportsDirectoryPath, config, logger); + exportsManager.init(); + return exportsManager; + } } /** diff --git a/src/transports/base.ts b/src/transports/base.ts index 9f2bca0d..ac95ae1b 100644 --- a/src/transports/base.ts +++ b/src/transports/base.ts @@ -44,7 +44,7 @@ export abstract class TransportRunnerBase { const logger = new CompositeLogger(...loggers); const sessionId = new ObjectId().toString(); - const exportsManager = new ExportsManager(sessionId, userConfig, logger); + const exportsManager = ExportsManager.init(sessionId, userConfig, logger); const connectionManager = new ConnectionManager(); const session = new Session({ diff --git a/tests/integration/helpers.ts b/tests/integration/helpers.ts index dea7f347..28193d10 100644 --- a/tests/integration/helpers.ts +++ b/tests/integration/helpers.ts @@ -59,7 +59,7 @@ export function setupIntegrationTest(getUserConfig: () => UserConfig): Integrati const logger = new CompositeLogger(); const sessionId = new ObjectId().toString(); - const exportsManager = new ExportsManager(sessionId, userConfig, logger); + const exportsManager = ExportsManager.init(sessionId, userConfig, logger); const connectionManager = new ConnectionManager(); const session = new Session({ diff --git a/tests/integration/resources/exportedData.test.ts b/tests/integration/resources/exportedData.test.ts index 1c32ddbc..b5e6039a 100644 --- a/tests/integration/resources/exportedData.test.ts +++ b/tests/integration/resources/exportedData.test.ts @@ -1,9 +1,19 @@ +import path from "path"; +import fs from "fs/promises"; import { Long } from "bson"; -import { describe, expect, it, beforeEach } from "vitest"; +import { describe, expect, it, beforeEach, afterAll } from "vitest"; import { CallToolResult } from "@modelcontextprotocol/sdk/types.js"; import { defaultTestConfig, resourceChangedNotification, timeout } from "../helpers.js"; import { describeWithMongoDB } from "../tools/mongodb/mongodbHelpers.js"; import { contentWithResourceURILink } from "../tools/mongodb/read/export.test.js"; +import { UserConfig } from "../../../src/lib.js"; + +const userConfig: UserConfig = { + ...defaultTestConfig, + exportsPath: path.join(path.dirname(defaultTestConfig.exportsPath), `exports-${Date.now()}`), + exportTimeoutMs: 200, + exportCleanupIntervalMs: 300, +}; describeWithMongoDB( "exported-data resource", @@ -19,6 +29,10 @@ describeWithMongoDB( ]); }); + afterAll(async () => { + await fs.rm(userConfig.exportsPath, { recursive: true, force: true }); + }); + it("should be able to list resource template", async () => { await integration.connectMcpClient(); const response = await integration.mcpClient().listResourceTemplates(); @@ -123,11 +137,5 @@ describeWithMongoDB( }); }); }, - () => { - return { - ...defaultTestConfig, - exportTimeoutMs: 200, - exportCleanupIntervalMs: 300, - }; - } + () => userConfig ); diff --git a/tests/integration/telemetry.test.ts b/tests/integration/telemetry.test.ts index c5547ac6..4a34d156 100644 --- a/tests/integration/telemetry.test.ts +++ b/tests/integration/telemetry.test.ts @@ -20,7 +20,7 @@ describe("Telemetry", () => { apiBaseUrl: "", logger: new CompositeLogger(), sessionId: "1FOO", - exportsManager: new ExportsManager("1FOO", config, logger), + exportsManager: ExportsManager.init("1FOO", config, logger), connectionManager: new ConnectionManager(), }), config diff --git a/tests/integration/tools/mongodb/read/export.test.ts b/tests/integration/tools/mongodb/read/export.test.ts index f249f090..0a30f45a 100644 --- a/tests/integration/tools/mongodb/read/export.test.ts +++ b/tests/integration/tools/mongodb/read/export.test.ts @@ -1,14 +1,22 @@ +import path from "path"; +import { Long } from "bson"; import fs from "fs/promises"; -import { beforeEach, describe, expect, it } from "vitest"; +import { afterAll, beforeEach, describe, expect, it } from "vitest"; +import { CallToolResult } from "@modelcontextprotocol/sdk/types.js"; import { databaseCollectionParameters, + defaultTestConfig, resourceChangedNotification, validateThrowsForInvalidArguments, validateToolMetadata, } from "../../../helpers.js"; import { describeWithMongoDB } from "../mongodbHelpers.js"; -import { CallToolResult } from "@modelcontextprotocol/sdk/types.js"; -import { Long } from "bson"; +import { UserConfig } from "../../../../../src/lib.js"; + +const userConfig: UserConfig = { + ...defaultTestConfig, + exportsPath: path.join(path.dirname(defaultTestConfig.exportsPath), `exports-${Date.now()}`), +}; export function contentWithTextResourceURI( content: CallToolResult["content"] @@ -39,295 +47,304 @@ export function contentWithExportPath( }); } -describeWithMongoDB("export tool", (integration) => { - validateToolMetadata( - integration, - "export", - "Export a collection data or query results in the specified EJSON format.", - [ - ...databaseCollectionParameters, - - { - name: "filter", - description: "The query filter, matching the syntax of the query argument of db.collection.find()", - type: "object", - required: false, - }, - { - name: "jsonExportFormat", - description: [ - "The format to be used when exporting collection data as EJSON with default being relaxed.", - "relaxed: A string format that emphasizes readability and interoperability at the expense of type preservation. That is, conversion from relaxed format to BSON can lose type information.", - "canonical: A string format that emphasizes type preservation at the expense of readability and interoperability. That is, conversion from canonical to BSON will generally preserve type information except in certain specific cases.", - ].join("\n"), - type: "string", - required: false, - }, - { - name: "limit", - description: "The maximum number of documents to return", - type: "number", - required: false, - }, - { - name: "projection", - description: "The projection, matching the syntax of the projection argument of db.collection.find()", - type: "object", - required: false, - }, - { - name: "sort", - description: - "A document, describing the sort order, matching the syntax of the sort argument of cursor.sort(). The keys of the object are the fields to sort on, while the values are the sort directions (1 for ascending, -1 for descending).", - type: "object", - required: false, - }, - ] - ); - - validateThrowsForInvalidArguments(integration, "export", [ - {}, - { database: 123, collection: "bar" }, - { database: "test", collection: [] }, - { database: "test", collection: "bar", filter: "{ $gt: { foo: 5 } }" }, - { database: "test", collection: "bar", projection: "name" }, - { database: "test", collection: "bar", limit: "10" }, - { database: "test", collection: "bar", sort: [], limit: 10 }, - ]); - - beforeEach(async () => { - await integration.connectMcpClient(); - }); - - it("when provided with incorrect namespace, export should have empty data", async function () { - const response = await integration.mcpClient().callTool({ - name: "export", - arguments: { database: "non-existent", collection: "foos" }, - }); - const content = response.content as CallToolResult["content"]; - const exportURI = contentWithResourceURILink(content)?.uri as string; - await resourceChangedNotification(integration.mcpClient(), exportURI); - - expect(content).toHaveLength(3); - expect(contentWithTextResourceURI(content)).toBeDefined(); - expect(contentWithResourceURILink(content)).toBeDefined(); - - const localPathPart = contentWithExportPath(content); - expect(localPathPart).toBeDefined(); +describeWithMongoDB( + "export tool", + (integration) => { + validateToolMetadata( + integration, + "export", + "Export a collection data or query results in the specified EJSON format.", + [ + ...databaseCollectionParameters, - const [, localPath] = /"(.*)"/.exec(String(localPathPart?.text)) ?? []; - expect(localPath).toBeDefined(); + { + name: "filter", + description: "The query filter, matching the syntax of the query argument of db.collection.find()", + type: "object", + required: false, + }, + { + name: "jsonExportFormat", + description: [ + "The format to be used when exporting collection data as EJSON with default being relaxed.", + "relaxed: A string format that emphasizes readability and interoperability at the expense of type preservation. That is, conversion from relaxed format to BSON can lose type information.", + "canonical: A string format that emphasizes type preservation at the expense of readability and interoperability. That is, conversion from canonical to BSON will generally preserve type information except in certain specific cases.", + ].join("\n"), + type: "string", + required: false, + }, + { + name: "limit", + description: "The maximum number of documents to return", + type: "number", + required: false, + }, + { + name: "projection", + description: + "The projection, matching the syntax of the projection argument of db.collection.find()", + type: "object", + required: false, + }, + { + name: "sort", + description: + "A document, describing the sort order, matching the syntax of the sort argument of cursor.sort(). The keys of the object are the fields to sort on, while the values are the sort directions (1 for ascending, -1 for descending).", + type: "object", + required: false, + }, + ] + ); - expect(await fs.readFile(localPath as string, "utf8")).toEqual("[]"); - }); + validateThrowsForInvalidArguments(integration, "export", [ + {}, + { database: 123, collection: "bar" }, + { database: "test", collection: [] }, + { database: "test", collection: "bar", filter: "{ $gt: { foo: 5 } }" }, + { database: "test", collection: "bar", projection: "name" }, + { database: "test", collection: "bar", limit: "10" }, + { database: "test", collection: "bar", sort: [], limit: 10 }, + ]); - describe("with correct namespace", function () { beforeEach(async () => { - const mongoClient = integration.mongoClient(); - await mongoClient - .db(integration.randomDbName()) - .collection("foo") - .insertMany([ - { name: "foo", longNumber: new Long(1234) }, - { name: "bar", bigInt: new Long(123412341234) }, - ]); - }); - - it("should export entire namespace when filter are empty", async function () { await integration.connectMcpClient(); - const response = await integration.mcpClient().callTool({ - name: "export", - arguments: { database: integration.randomDbName(), collection: "foo" }, - }); - const content = response.content as CallToolResult["content"]; - const exportURI = contentWithResourceURILink(content)?.uri as string; - await resourceChangedNotification(integration.mcpClient(), exportURI); - - const localPathPart = contentWithExportPath(content); - expect(localPathPart).toBeDefined(); - const [, localPath] = /"(.*)"/.exec(String(localPathPart?.text)) ?? []; - expect(localPath).toBeDefined(); - - const exportedContent = JSON.parse(await fs.readFile(localPath as string, "utf8")) as Record< - string, - unknown - >[]; - expect(exportedContent).toHaveLength(2); - expect(exportedContent[0]?.name).toEqual("foo"); - expect(exportedContent[1]?.name).toEqual("bar"); }); - it("should export filter results namespace when there are filters", async function () { - await integration.connectMcpClient(); - const response = await integration.mcpClient().callTool({ - name: "export", - arguments: { database: integration.randomDbName(), collection: "foo", filter: { name: "foo" } }, - }); - const content = response.content as CallToolResult["content"]; - const exportURI = contentWithResourceURILink(content)?.uri as string; - await resourceChangedNotification(integration.mcpClient(), exportURI); - - const localPathPart = contentWithExportPath(content); - expect(localPathPart).toBeDefined(); - const [, localPath] = /"(.*)"/.exec(String(localPathPart?.text)) ?? []; - expect(localPath).toBeDefined(); - - const exportedContent = JSON.parse(await fs.readFile(localPath as string, "utf8")) as Record< - string, - unknown - >[]; - expect(exportedContent).toHaveLength(1); - expect(exportedContent[0]?.name).toEqual("foo"); + afterAll(async () => { + await fs.rm(userConfig.exportsPath, { recursive: true, force: true }); }); - it("should export results limited to the provided limit", async function () { - await integration.connectMcpClient(); + it("when provided with incorrect namespace, export should have empty data", async function () { const response = await integration.mcpClient().callTool({ name: "export", - arguments: { database: integration.randomDbName(), collection: "foo", limit: 1 }, + arguments: { database: "non-existent", collection: "foos" }, }); const content = response.content as CallToolResult["content"]; const exportURI = contentWithResourceURILink(content)?.uri as string; await resourceChangedNotification(integration.mcpClient(), exportURI); - const localPathPart = contentWithExportPath(content); - expect(localPathPart).toBeDefined(); - const [, localPath] = /"(.*)"/.exec(String(localPathPart?.text)) ?? []; - expect(localPath).toBeDefined(); - - const exportedContent = JSON.parse(await fs.readFile(localPath as string, "utf8")) as Record< - string, - unknown - >[]; - expect(exportedContent).toHaveLength(1); - expect(exportedContent[0]?.name).toEqual("foo"); - }); - - it("should export results with sorted by the provided sort", async function () { - await integration.connectMcpClient(); - const response = await integration.mcpClient().callTool({ - name: "export", - arguments: { - database: integration.randomDbName(), - collection: "foo", - limit: 1, - sort: { longNumber: 1 }, - }, - }); - const content = response.content as CallToolResult["content"]; - const exportURI = contentWithResourceURILink(content)?.uri as string; - await resourceChangedNotification(integration.mcpClient(), exportURI); + expect(content).toHaveLength(3); + expect(contentWithTextResourceURI(content)).toBeDefined(); + expect(contentWithResourceURILink(content)).toBeDefined(); const localPathPart = contentWithExportPath(content); expect(localPathPart).toBeDefined(); + const [, localPath] = /"(.*)"/.exec(String(localPathPart?.text)) ?? []; expect(localPath).toBeDefined(); - const exportedContent = JSON.parse(await fs.readFile(localPath as string, "utf8")) as Record< - string, - unknown - >[]; - expect(exportedContent).toHaveLength(1); - expect(exportedContent[0]?.name).toEqual("bar"); + expect(await fs.readFile(localPath as string, "utf8")).toEqual("[]"); }); - it("should export results containing only projected fields", async function () { - await integration.connectMcpClient(); - const response = await integration.mcpClient().callTool({ - name: "export", - arguments: { - database: integration.randomDbName(), - collection: "foo", - limit: 1, - projection: { _id: 0, name: 1 }, - }, + describe("with correct namespace", function () { + beforeEach(async () => { + const mongoClient = integration.mongoClient(); + await mongoClient + .db(integration.randomDbName()) + .collection("foo") + .insertMany([ + { name: "foo", longNumber: new Long(1234) }, + { name: "bar", bigInt: new Long(123412341234) }, + ]); }); - const content = response.content as CallToolResult["content"]; - const exportURI = contentWithResourceURILink(content)?.uri as string; - await resourceChangedNotification(integration.mcpClient(), exportURI); - const localPathPart = contentWithExportPath(content); - expect(localPathPart).toBeDefined(); - const [, localPath] = /"(.*)"/.exec(String(localPathPart?.text)) ?? []; - expect(localPath).toBeDefined(); - - const exportedContent = JSON.parse(await fs.readFile(localPath as string, "utf8")) as Record< - string, - unknown - >[]; - expect(exportedContent).toEqual([ - { - name: "foo", - }, - ]); - }); + it("should export entire namespace when filter are empty", async function () { + await integration.connectMcpClient(); + const response = await integration.mcpClient().callTool({ + name: "export", + arguments: { database: integration.randomDbName(), collection: "foo" }, + }); + const content = response.content as CallToolResult["content"]; + const exportURI = contentWithResourceURILink(content)?.uri as string; + await resourceChangedNotification(integration.mcpClient(), exportURI); + + const localPathPart = contentWithExportPath(content); + expect(localPathPart).toBeDefined(); + const [, localPath] = /"(.*)"/.exec(String(localPathPart?.text)) ?? []; + expect(localPath).toBeDefined(); + + const exportedContent = JSON.parse(await fs.readFile(localPath as string, "utf8")) as Record< + string, + unknown + >[]; + expect(exportedContent).toHaveLength(2); + expect(exportedContent[0]?.name).toEqual("foo"); + expect(exportedContent[1]?.name).toEqual("bar"); + }); - it("should export relaxed json when provided jsonExportFormat is relaxed", async function () { - await integration.connectMcpClient(); - const response = await integration.mcpClient().callTool({ - name: "export", - arguments: { - database: integration.randomDbName(), - collection: "foo", - limit: 1, - projection: { _id: 0 }, - jsonExportFormat: "relaxed", - }, + it("should export filter results namespace when there are filters", async function () { + await integration.connectMcpClient(); + const response = await integration.mcpClient().callTool({ + name: "export", + arguments: { database: integration.randomDbName(), collection: "foo", filter: { name: "foo" } }, + }); + const content = response.content as CallToolResult["content"]; + const exportURI = contentWithResourceURILink(content)?.uri as string; + await resourceChangedNotification(integration.mcpClient(), exportURI); + + const localPathPart = contentWithExportPath(content); + expect(localPathPart).toBeDefined(); + const [, localPath] = /"(.*)"/.exec(String(localPathPart?.text)) ?? []; + expect(localPath).toBeDefined(); + + const exportedContent = JSON.parse(await fs.readFile(localPath as string, "utf8")) as Record< + string, + unknown + >[]; + expect(exportedContent).toHaveLength(1); + expect(exportedContent[0]?.name).toEqual("foo"); }); - const content = response.content as CallToolResult["content"]; - const exportURI = contentWithResourceURILink(content)?.uri as string; - await resourceChangedNotification(integration.mcpClient(), exportURI); - const localPathPart = contentWithExportPath(content); - expect(localPathPart).toBeDefined(); - const [, localPath] = /"(.*)"/.exec(String(localPathPart?.text)) ?? []; - expect(localPath).toBeDefined(); + it("should export results limited to the provided limit", async function () { + await integration.connectMcpClient(); + const response = await integration.mcpClient().callTool({ + name: "export", + arguments: { database: integration.randomDbName(), collection: "foo", limit: 1 }, + }); + const content = response.content as CallToolResult["content"]; + const exportURI = contentWithResourceURILink(content)?.uri as string; + await resourceChangedNotification(integration.mcpClient(), exportURI); + + const localPathPart = contentWithExportPath(content); + expect(localPathPart).toBeDefined(); + const [, localPath] = /"(.*)"/.exec(String(localPathPart?.text)) ?? []; + expect(localPath).toBeDefined(); + + const exportedContent = JSON.parse(await fs.readFile(localPath as string, "utf8")) as Record< + string, + unknown + >[]; + expect(exportedContent).toHaveLength(1); + expect(exportedContent[0]?.name).toEqual("foo"); + }); - const exportedContent = JSON.parse(await fs.readFile(localPath as string, "utf8")) as Record< - string, - unknown - >[]; - expect(exportedContent).toEqual([ - { - name: "foo", - longNumber: 1234, - }, - ]); - }); + it("should export results with sorted by the provided sort", async function () { + await integration.connectMcpClient(); + const response = await integration.mcpClient().callTool({ + name: "export", + arguments: { + database: integration.randomDbName(), + collection: "foo", + limit: 1, + sort: { longNumber: 1 }, + }, + }); + const content = response.content as CallToolResult["content"]; + const exportURI = contentWithResourceURILink(content)?.uri as string; + await resourceChangedNotification(integration.mcpClient(), exportURI); + + const localPathPart = contentWithExportPath(content); + expect(localPathPart).toBeDefined(); + const [, localPath] = /"(.*)"/.exec(String(localPathPart?.text)) ?? []; + expect(localPath).toBeDefined(); + + const exportedContent = JSON.parse(await fs.readFile(localPath as string, "utf8")) as Record< + string, + unknown + >[]; + expect(exportedContent).toHaveLength(1); + expect(exportedContent[0]?.name).toEqual("bar"); + }); - it("should export canonical json when provided jsonExportFormat is canonical", async function () { - await integration.connectMcpClient(); - const response = await integration.mcpClient().callTool({ - name: "export", - arguments: { - database: integration.randomDbName(), - collection: "foo", - limit: 1, - projection: { _id: 0 }, - jsonExportFormat: "canonical", - }, + it("should export results containing only projected fields", async function () { + await integration.connectMcpClient(); + const response = await integration.mcpClient().callTool({ + name: "export", + arguments: { + database: integration.randomDbName(), + collection: "foo", + limit: 1, + projection: { _id: 0, name: 1 }, + }, + }); + const content = response.content as CallToolResult["content"]; + const exportURI = contentWithResourceURILink(content)?.uri as string; + await resourceChangedNotification(integration.mcpClient(), exportURI); + + const localPathPart = contentWithExportPath(content); + expect(localPathPart).toBeDefined(); + const [, localPath] = /"(.*)"/.exec(String(localPathPart?.text)) ?? []; + expect(localPath).toBeDefined(); + + const exportedContent = JSON.parse(await fs.readFile(localPath as string, "utf8")) as Record< + string, + unknown + >[]; + expect(exportedContent).toEqual([ + { + name: "foo", + }, + ]); }); - const content = response.content as CallToolResult["content"]; - const exportURI = contentWithResourceURILink(content)?.uri as string; - await resourceChangedNotification(integration.mcpClient(), exportURI); - const localPathPart = contentWithExportPath(content); - expect(localPathPart).toBeDefined(); - const [, localPath] = /"(.*)"/.exec(String(localPathPart?.text)) ?? []; - expect(localPath).toBeDefined(); + it("should export relaxed json when provided jsonExportFormat is relaxed", async function () { + await integration.connectMcpClient(); + const response = await integration.mcpClient().callTool({ + name: "export", + arguments: { + database: integration.randomDbName(), + collection: "foo", + limit: 1, + projection: { _id: 0 }, + jsonExportFormat: "relaxed", + }, + }); + const content = response.content as CallToolResult["content"]; + const exportURI = contentWithResourceURILink(content)?.uri as string; + await resourceChangedNotification(integration.mcpClient(), exportURI); + + const localPathPart = contentWithExportPath(content); + expect(localPathPart).toBeDefined(); + const [, localPath] = /"(.*)"/.exec(String(localPathPart?.text)) ?? []; + expect(localPath).toBeDefined(); + + const exportedContent = JSON.parse(await fs.readFile(localPath as string, "utf8")) as Record< + string, + unknown + >[]; + expect(exportedContent).toEqual([ + { + name: "foo", + longNumber: 1234, + }, + ]); + }); - const exportedContent = JSON.parse(await fs.readFile(localPath as string, "utf8")) as Record< - string, - unknown - >[]; - expect(exportedContent).toEqual([ - { - name: "foo", - longNumber: { - $numberLong: "1234", + it("should export canonical json when provided jsonExportFormat is canonical", async function () { + await integration.connectMcpClient(); + const response = await integration.mcpClient().callTool({ + name: "export", + arguments: { + database: integration.randomDbName(), + collection: "foo", + limit: 1, + projection: { _id: 0 }, + jsonExportFormat: "canonical", }, - }, - ]); + }); + const content = response.content as CallToolResult["content"]; + const exportURI = contentWithResourceURILink(content)?.uri as string; + await resourceChangedNotification(integration.mcpClient(), exportURI); + + const localPathPart = contentWithExportPath(content); + expect(localPathPart).toBeDefined(); + const [, localPath] = /"(.*)"/.exec(String(localPathPart?.text)) ?? []; + expect(localPath).toBeDefined(); + + const exportedContent = JSON.parse(await fs.readFile(localPath as string, "utf8")) as Record< + string, + unknown + >[]; + expect(exportedContent).toEqual([ + { + name: "foo", + longNumber: { + $numberLong: "1234", + }, + }, + ]); + }); }); - }); -}); + }, + () => userConfig +); diff --git a/tests/unit/common/exportsManager.test.ts b/tests/unit/common/exportsManager.test.ts index e03d2114..63ae3e9e 100644 --- a/tests/unit/common/exportsManager.test.ts +++ b/tests/unit/common/exportsManager.test.ts @@ -2,7 +2,7 @@ import path from "path"; import fs from "fs/promises"; import { Readable, Transform } from "stream"; import { FindCursor, Long } from "mongodb"; -import { beforeEach, describe, expect, it, vi } from "vitest"; +import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; import { ensureExtension, isExportExpired, @@ -19,7 +19,7 @@ import { EJSON, EJSONOptions, ObjectId } from "bson"; import { CompositeLogger } from "../../../src/common/logger.js"; import { ConnectionManager } from "../../../src/common/connectionManager.js"; -const exportsPath = path.join(ROOT_DIR, "tests", "tmp", "exports"); +const exportsPath = path.join(ROOT_DIR, "tests", "tmp", `exports-${Date.now()}`); const exportsManagerConfig: ExportsManagerConfig = { exportsPath, exportTimeoutMs: config.exportTimeoutMs, @@ -104,8 +104,6 @@ describe("ExportsManager unit test", () => { let manager: ExportsManager; beforeEach(async () => { - await manager?.close(); - await fs.rm(exportsManagerConfig.exportsPath, { recursive: true, force: true }); await fs.mkdir(exportsManagerConfig.exportsPath, { recursive: true }); const logger = new CompositeLogger(); const sessionId = new ObjectId().toString(); @@ -113,12 +111,17 @@ describe("ExportsManager unit test", () => { apiBaseUrl: "", logger, sessionId, - exportsManager: new ExportsManager(sessionId, config, logger), + exportsManager: ExportsManager.init(sessionId, exportsManagerConfig, logger), connectionManager: new ConnectionManager(), }); manager = session.exportsManager; }); + afterEach(async () => { + await manager?.close(); + await fs.rm(exportsManagerConfig.exportsPath, { recursive: true, force: true }); + }); + describe("#availableExport", () => { it("should list only the exports that are in ready state", async () => { // This export will finish in at-least 1 second @@ -366,7 +369,7 @@ describe("ExportsManager unit test", () => { it("should not clean up in-progress exports", async () => { const { exportName } = getExportNameAndPath(session.sessionId, Date.now()); - const manager = new ExportsManager( + const manager = ExportsManager.init( session.sessionId, { ...exportsManagerConfig, @@ -393,7 +396,7 @@ describe("ExportsManager unit test", () => { it("should cleanup expired exports", async () => { const { exportName, exportPath, exportURI } = getExportNameAndPath(session.sessionId, Date.now()); - const manager = new ExportsManager( + const manager = ExportsManager.init( session.sessionId, { ...exportsManagerConfig, diff --git a/tests/unit/common/session.test.ts b/tests/unit/common/session.test.ts index c860e97e..e1cce9a0 100644 --- a/tests/unit/common/session.test.ts +++ b/tests/unit/common/session.test.ts @@ -18,7 +18,7 @@ describe("Session", () => { apiBaseUrl: "https://api.test.com", logger, sessionId: "1FOO", - exportsManager: new ExportsManager("1FOO", config, logger), + exportsManager: ExportsManager.init("1FOO", config, logger), connectionManager: new ConnectionManager(), }); diff --git a/tests/unit/resources/common/debug.test.ts b/tests/unit/resources/common/debug.test.ts index 997be48c..684adea4 100644 --- a/tests/unit/resources/common/debug.test.ts +++ b/tests/unit/resources/common/debug.test.ts @@ -13,7 +13,7 @@ describe("debug resource", () => { apiBaseUrl: "", logger, sessionId: "1FOO", - exportsManager: new ExportsManager("1FOO", config, logger), + exportsManager: ExportsManager.init("1FOO", config, logger), connectionManager: new ConnectionManager(), }); const telemetry = Telemetry.create(session, { ...config, telemetry: "disabled" }); From 2b2d615fb3274e61c24341c18490c7d06a571008 Mon Sep 17 00:00:00 2001 From: Himanshu Singh Date: Fri, 8 Aug 2025 15:11:11 +0200 Subject: [PATCH 34/50] chore: small improvements 1. outputStream.write moved to within the Transform.flush 2. won't send resource updated notification on export-expired event or it might trigger client to fetch expired exports. 3. added ObjectId to the file names to make them unique --- src/common/exportsManager.ts | 28 +++--- src/resources/common/exportedData.ts | 11 +-- src/tools/mongodb/read/export.ts | 9 +- tests/unit/common/exportsManager.test.ts | 108 ++++++++++++++++------- 4 files changed, 104 insertions(+), 52 deletions(-) diff --git a/src/common/exportsManager.ts b/src/common/exportsManager.ts index 84360b05..d625b48a 100644 --- a/src/common/exportsManager.ts +++ b/src/common/exportsManager.ts @@ -192,7 +192,6 @@ export class ExportsManager extends EventEmitter { try { await fs.mkdir(this.exportsDirectoryPath, { recursive: true }); const outputStream = createWriteStream(inProgressExport.exportPath); - outputStream.write("["); await pipeline([ input.stream(), this.docToEJSONStream(this.getEJSONOptionsForFormat(jsonExportFormat)), @@ -242,20 +241,27 @@ export class ExportsManager extends EventEmitter { let docsTransformed = 0; return new Transform({ objectMode: true, - transform: function (chunk: unknown, encoding, callback): void { - ++docsTransformed; + transform(chunk: unknown, encoding, callback): void { try { - const doc: string = EJSON.stringify(chunk, undefined, undefined, ejsonOptions); - const line = `${docsTransformed > 1 ? ",\n" : ""}${doc}`; - - callback(null, line); - } catch (err: unknown) { + 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); } }, - final: function (callback): void { - this.push("]"); - callback(null); + flush(callback): void { + if (docsTransformed === 0) { + this.push("[]"); + } else { + this.push("]"); + } + callback(); }, }); } diff --git a/src/resources/common/exportedData.ts b/src/resources/common/exportedData.ts index 3ec685ae..5f59311b 100644 --- a/src/resources/common/exportedData.ts +++ b/src/resources/common/exportedData.ts @@ -15,12 +15,13 @@ export class ExportedData { private server?: Server; constructor(private readonly session: Session) { - const onExportChanged = (uri: string): void => { + this.session.exportsManager.on("export-available", (uri: string): void => { this.server?.sendResourceListChanged(); this.server?.sendResourceUpdated(uri); - }; - this.session.exportsManager.on("export-available", onExportChanged); - this.session.exportsManager.on("export-expired", onExportChanged); + }); + this.session.exportsManager.on("export-expired", (): void => { + this.server?.sendResourceListChanged(); + }); } public register(server: Server): void { @@ -113,7 +114,7 @@ export class ExportedData { }; private exportNameToDescription(exportName: string): string { - const match = exportName.match(/^(.+)\.(\d+)\.json$/); + const match = exportName.match(/^(.+)\.(\d+)\.(.+)\.json$/); if (!match) return "Exported data for an unknown namespace."; const [, namespace, timestamp] = match; diff --git a/src/tools/mongodb/read/export.ts b/src/tools/mongodb/read/export.ts index 8a147345..b261e9e1 100644 --- a/src/tools/mongodb/read/export.ts +++ b/src/tools/mongodb/read/export.ts @@ -1,9 +1,10 @@ +import z from "zod"; +import { ObjectId } from "bson"; import { CallToolResult } from "@modelcontextprotocol/sdk/types.js"; import { OperationType, ToolArgs } from "../../tool.js"; import { DbOperationArgs, MongoDBToolBase } from "../mongodbTool.js"; import { FindArgs } from "./find.js"; import { jsonExportFormat } from "../../../common/exportsManager.js"; -import z from "zod"; export class ExportTool extends MongoDBToolBase { public name = "export"; @@ -41,7 +42,11 @@ export class ExportTool extends MongoDBToolBase { promoteValues: false, bsonRegExp: true, }); - const exportName = `${database}.${collection}.${Date.now()}.json`; + // The format is namespace.date.objectid.json + // - namespace to identify which namespace the export belongs to + // - date to identify when the export was generated + // - objectid for uniqueness of the names + const exportName = `${database}.${collection}.${Date.now()}.${new ObjectId().toString()}.json`; const { exportURI, exportPath } = this.session.exportsManager.createJSONExport({ input: findCursor, diff --git a/tests/unit/common/exportsManager.test.ts b/tests/unit/common/exportsManager.test.ts index 63ae3e9e..8466045a 100644 --- a/tests/unit/common/exportsManager.test.ts +++ b/tests/unit/common/exportsManager.test.ts @@ -28,14 +28,15 @@ const exportsManagerConfig: ExportsManagerConfig = { function getExportNameAndPath( sessionId: string, - timestamp: number + timestamp: number = Date.now(), + objectId: string = new ObjectId().toString() ): { sessionExportsPath: string; exportName: string; exportPath: string; exportURI: string; } { - const exportName = `foo.bar.${timestamp}.json`; + const exportName = `foo.bar.${timestamp}.${objectId}.json`; const sessionExportsPath = path.join(exportsPath, sessionId); const exportPath = path.join(sessionExportsPath, exportName); return { @@ -48,22 +49,21 @@ function getExportNameAndPath( function createDummyFindCursor( dataArray: unknown[], - chunkPushTimeoutMs?: number + beforeEachChunk?: (chunkIndex: number) => void | Promise ): { cursor: FindCursor; cursorCloseNotification: Promise } { let index = 0; const readable = new Readable({ objectMode: true, async read(): Promise { - if (index < dataArray.length) { - if (chunkPushTimeoutMs) { - await timeout(chunkPushTimeoutMs); + try { + await beforeEachChunk?.(index); + if (index < dataArray.length) { + this.push(dataArray[index++]); + } else { + this.push(null); } - this.push(dataArray[index++]); - } else { - if (chunkPushTimeoutMs) { - await timeout(chunkPushTimeoutMs); - } - this.push(null); + } catch (error) { + this.destroy(error as Error); } }, }); @@ -90,6 +90,13 @@ function createDummyFindCursor( }; } +function createDummyFindCursorWithDelay( + dataArray: unknown[], + delayMs: number +): { cursor: FindCursor; cursorCloseNotification: Promise } { + return createDummyFindCursor(dataArray, () => timeout(delayMs)); +} + async function fileExists(filePath: string): Promise { try { await fs.access(filePath); @@ -125,15 +132,15 @@ describe("ExportsManager unit test", () => { describe("#availableExport", () => { it("should list only the exports that are in ready state", async () => { // This export will finish in at-least 1 second - const { exportName: exportName1 } = getExportNameAndPath(session.sessionId, Date.now()); + const { exportName: exportName1 } = getExportNameAndPath(session.sessionId); manager.createJSONExport({ - input: createDummyFindCursor([{ name: "Test1" }], 1000).cursor, + input: createDummyFindCursorWithDelay([{ name: "Test1" }], 1000).cursor, exportName: exportName1, jsonExportFormat: "relaxed", }); // This export will finish way sooner than the first one - const { exportName: exportName2 } = getExportNameAndPath(session.sessionId, Date.now()); + const { exportName: exportName2 } = getExportNameAndPath(session.sessionId); const { cursor, cursorCloseNotification } = createDummyFindCursor([{ name: "Test1" }]); manager.createJSONExport({ input: cursor, @@ -154,8 +161,8 @@ describe("ExportsManager unit test", () => { }); it("should throw if the resource is still being generated", async () => { - const { exportName } = getExportNameAndPath(session.sessionId, Date.now()); - const { cursor } = createDummyFindCursor([{ name: "Test1" }], 100); + const { exportName } = getExportNameAndPath(session.sessionId); + const { cursor } = createDummyFindCursorWithDelay([{ name: "Test1" }], 100); manager.createJSONExport({ input: cursor, exportName, @@ -168,7 +175,7 @@ describe("ExportsManager unit test", () => { }); it("should return the resource content if the resource is ready to be consumed", async () => { - const { exportName } = getExportNameAndPath(session.sessionId, Date.now()); + const { exportName } = getExportNameAndPath(session.sessionId); const { cursor, cursorCloseNotification } = createDummyFindCursor([]); manager.createJSONExport({ input: cursor, @@ -198,7 +205,7 @@ describe("ExportsManager unit test", () => { longNumber: Long.fromNumber(123456), }, ])); - ({ exportName, exportPath, exportURI } = getExportNameAndPath(session.sessionId, Date.now())); + ({ exportName, exportPath, exportURI } = getExportNameAndPath(session.sessionId)); }); describe("when cursor is empty", () => { @@ -304,7 +311,7 @@ describe("ExportsManager unit test", () => { }); }); - describe("when there is an error in export generation", () => { + describe("when there is an error during stream transform", () => { it("should remove the partial export and never make it available", async () => { const emitSpy = vi.spyOn(manager, "emit"); // eslint-disable-next-line @typescript-eslint/no-unsafe-member-access, @typescript-eslint/no-explicit-any @@ -312,23 +319,29 @@ describe("ExportsManager unit test", () => { let docsTransformed = 0; return new Transform({ objectMode: true, - transform: function (chunk: unknown, encoding, callback): void { - ++docsTransformed; + transform(chunk: unknown, encoding, callback): void { try { - if (docsTransformed === 1) { + const doc = EJSON.stringify(chunk, undefined, undefined, ejsonOptions); + if (docsTransformed === 0) { + this.push("[" + doc); + } else if (docsTransformed === 1) { throw new Error("Could not transform the chunk!"); + } else { + this.push(",\n" + doc); } - const doc: string = EJSON.stringify(chunk, undefined, 2, ejsonOptions); - const line = `${docsTransformed > 1 ? ",\n" : ""}${doc}`; - - callback(null, line); - } catch (err: unknown) { + docsTransformed++; + callback(); + } catch (err) { callback(err as Error); } }, - final: function (callback): void { - this.push("]"); - callback(null); + flush(this: Transform, cb): void { + if (docsTransformed === 0) { + this.push("[]"); + } else { + this.push("]"); + } + cb(); }, }); }; @@ -348,6 +361,33 @@ describe("ExportsManager unit test", () => { expect(await fileExists(exportPath)).toEqual(false); }); }); + + describe("when there is an error on read stream", () => { + it("should remove the partial export and never make it available", async () => { + const emitSpy = vi.spyOn(manager, "emit"); + // A cursor that will make the read stream fail after the first chunk + const { cursor, cursorCloseNotification } = createDummyFindCursor([{ name: "Test1" }], (chunkIndex) => { + if (chunkIndex > 0) { + return Promise.reject(new Error("Connection timedout!")); + } + return Promise.resolve(); + }); + manager.createJSONExport({ + input: cursor, + exportName, + jsonExportFormat: "relaxed", + }); + await cursorCloseNotification; + + // Because the export was never populated in the available exports. + await expect(() => manager.readExport(exportName)).rejects.toThrow( + "Requested export has either expired or does not exist!" + ); + expect(emitSpy).not.toHaveBeenCalled(); + expect(manager.availableExports).toEqual([]); + expect(await fileExists(exportPath)).toEqual(false); + }); + }); }); describe("#cleanupExpiredExports", () => { @@ -368,7 +408,7 @@ describe("ExportsManager unit test", () => { }); it("should not clean up in-progress exports", async () => { - const { exportName } = getExportNameAndPath(session.sessionId, Date.now()); + const { exportName } = getExportNameAndPath(session.sessionId); const manager = ExportsManager.init( session.sessionId, { @@ -378,7 +418,7 @@ describe("ExportsManager unit test", () => { }, new CompositeLogger() ); - const { cursor } = createDummyFindCursor([{ name: "Test" }], 2000); + const { cursor } = createDummyFindCursorWithDelay([{ name: "Test" }], 2000); manager.createJSONExport({ input: cursor, exportName, @@ -395,7 +435,7 @@ describe("ExportsManager unit test", () => { }); it("should cleanup expired exports", async () => { - const { exportName, exportPath, exportURI } = getExportNameAndPath(session.sessionId, Date.now()); + const { exportName, exportPath, exportURI } = getExportNameAndPath(session.sessionId); const manager = ExportsManager.init( session.sessionId, { From 9998d8ef71482a8b6ff046fb0b51392b4463361f Mon Sep 17 00:00:00 2001 From: Himanshu Singh Date: Fri, 8 Aug 2025 15:14:04 +0200 Subject: [PATCH 35/50] chore: extra guard to avoid having same name export --- src/common/exportsManager.ts | 3 +++ 1 file changed, 3 insertions(+) diff --git a/src/common/exportsManager.ts b/src/common/exportsManager.ts index d625b48a..5366e5a5 100644 --- a/src/common/exportsManager.ts +++ b/src/common/exportsManager.ts @@ -158,6 +158,9 @@ export class ExportsManager extends EventEmitter { }): AvailableExport { try { const exportNameWithExtension = validateExportName(ensureExtension(exportName, "json")); + if (this.storedExports[exportNameWithExtension]) { + throw new Error("Export with same name is either already available or being generated."); + } const exportURI = `exported-data://${encodeURIComponent(exportNameWithExtension)}`; const exportFilePath = path.join(this.exportsDirectoryPath, exportNameWithExtension); const inProgressExport: InProgressExport = (this.storedExports[exportNameWithExtension] = { From 4fb1bbd8215b04437fdd5686a2c68c2505c92862 Mon Sep 17 00:00:00 2001 From: Himanshu Singh Date: Fri, 8 Aug 2025 15:53:22 +0200 Subject: [PATCH 36/50] chore: when closing ExportsManager should not accept any more requests --- src/common/exportsManager.ts | 17 ++++++++++++++++- tests/unit/common/exportsManager.test.ts | 22 ++++++++++++++++++++++ 2 files changed, 38 insertions(+), 1 deletion(-) diff --git a/src/common/exportsManager.ts b/src/common/exportsManager.ts index 5366e5a5..a0dec027 100644 --- a/src/common/exportsManager.ts +++ b/src/common/exportsManager.ts @@ -64,6 +64,7 @@ type ExportsManagerEvents = { export class ExportsManager extends EventEmitter { private wasInitialized: boolean = false; + private isShuttingDown: boolean = false; private storedExports: Record = {}; private exportsCleanupInProgress: boolean = false; private exportsCleanupInterval?: NodeJS.Timeout; @@ -77,6 +78,7 @@ export class ExportsManager extends EventEmitter { } public get availableExports(): AvailableExport[] { + this.assertIsNotShuttingDown(); return Object.values(this.storedExports) .filter((storedExport) => { return ( @@ -103,6 +105,11 @@ export class ExportsManager extends EventEmitter { } public async close(): Promise { + if (this.isShuttingDown) { + return; + } + + this.isShuttingDown = true; try { clearInterval(this.exportsCleanupInterval); await fs.rm(this.exportsDirectoryPath, { force: true, recursive: true }); @@ -117,6 +124,7 @@ export class ExportsManager extends EventEmitter { public async readExport(exportName: string): Promise { try { + this.assertIsNotShuttingDown(); exportName = decodeURIComponent(exportName); const exportHandle = this.storedExports[exportName]; if (!exportHandle) { @@ -157,6 +165,7 @@ export class ExportsManager extends EventEmitter { jsonExportFormat: JSONExportFormat; }): AvailableExport { try { + this.assertIsNotShuttingDown(); const exportNameWithExtension = validateExportName(ensureExtension(exportName, "json")); if (this.storedExports[exportNameWithExtension]) { throw new Error("Export with same name is either already available or being generated."); @@ -270,7 +279,7 @@ export class ExportsManager extends EventEmitter { } private async cleanupExpiredExports(): Promise { - if (this.exportsCleanupInProgress) { + if (this.exportsCleanupInProgress || this.isShuttingDown) { return; } @@ -318,6 +327,12 @@ export class ExportsManager extends EventEmitter { } } + private assertIsNotShuttingDown(): void { + if (this.isShuttingDown) { + throw new Error("ExportsManager is shutting down."); + } + } + static init(sessionId: string, config: ExportsManagerConfig, logger: LoggerBase): ExportsManager { const exportsDirectoryPath = path.join(config.exportsPath, sessionId); const exportsManager = new ExportsManager(exportsDirectoryPath, config, logger); diff --git a/tests/unit/common/exportsManager.test.ts b/tests/unit/common/exportsManager.test.ts index 8466045a..1bd0e493 100644 --- a/tests/unit/common/exportsManager.test.ts +++ b/tests/unit/common/exportsManager.test.ts @@ -130,6 +130,11 @@ describe("ExportsManager unit test", () => { }); describe("#availableExport", () => { + it("should throw if the manager is shutting down", () => { + void manager.close(); + expect(() => manager.availableExports).toThrow("ExportsManager is shutting down."); + }); + it("should list only the exports that are in ready state", async () => { // This export will finish in at-least 1 second const { exportName: exportName1 } = getExportNameAndPath(session.sessionId); @@ -156,6 +161,11 @@ describe("ExportsManager unit test", () => { }); describe("#readExport", () => { + it("should throw if the manager is shutting down", async () => { + void manager.close(); + await expect(() => manager.readExport("name")).rejects.toThrow("ExportsManager is shutting down."); + }); + it("should throw when export name has no extension", async () => { await expect(() => manager.readExport("name")).rejects.toThrow(); }); @@ -208,6 +218,18 @@ describe("ExportsManager unit test", () => { ({ exportName, exportPath, exportURI } = getExportNameAndPath(session.sessionId)); }); + it("should throw if the manager is shutting down", () => { + const { cursor } = createDummyFindCursor([]); + void manager.close(); + expect(() => + manager.createJSONExport({ + input: cursor, + exportName, + jsonExportFormat: "relaxed", + }) + ).toThrow(); + }); + describe("when cursor is empty", () => { it("should create an empty export", async () => { const { cursor, cursorCloseNotification } = createDummyFindCursor([]); From 50f9122cf3c94ff552d8d0827817b14f0e9714dc Mon Sep 17 00:00:00 2001 From: Himanshu Singh Date: Fri, 8 Aug 2025 16:41:54 +0200 Subject: [PATCH 37/50] chore: more tests for exportsManager --- tests/unit/common/exportsManager.test.ts | 16 ++++++++++++++++ 1 file changed, 16 insertions(+) diff --git a/tests/unit/common/exportsManager.test.ts b/tests/unit/common/exportsManager.test.ts index 1bd0e493..df1ffb36 100644 --- a/tests/unit/common/exportsManager.test.ts +++ b/tests/unit/common/exportsManager.test.ts @@ -230,6 +230,22 @@ describe("ExportsManager unit test", () => { ).toThrow(); }); + it("should throw if the same name export is requested more than once", () => { + const { cursor } = createDummyFindCursorWithDelay([{ name: 1 }, { name: 2 }], 100); + manager.createJSONExport({ + input: cursor, + exportName, + jsonExportFormat: "relaxed", + }); + expect(() => + manager.createJSONExport({ + input: cursor, + exportName, + jsonExportFormat: "relaxed", + }) + ).toThrow(); + }); + describe("when cursor is empty", () => { it("should create an empty export", async () => { const { cursor, cursorCloseNotification } = createDummyFindCursor([]); From c1d01177fa8a382226c4e66b476d0d4efce29a5e Mon Sep 17 00:00:00 2001 From: Himanshu Singh Date: Fri, 8 Aug 2025 17:12:06 +0200 Subject: [PATCH 38/50] chore: detect local run use-case better --- src/tools/mongodb/read/export.ts | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/src/tools/mongodb/read/export.ts b/src/tools/mongodb/read/export.ts index b261e9e1..d264f81b 100644 --- a/src/tools/mongodb/read/export.ts +++ b/src/tools/mongodb/read/export.ts @@ -73,7 +73,7 @@ export class ExportTool extends MongoDBToolBase { // This special case is to make it easier to work with exported data for // clients that still cannot reference resources (Cursor). // More information here: https://jira.mongodb.org/browse/MCP-104 - if (this.config.transport === "stdio") { + if (this.isServerRunningLocally()) { toolCallContent.push({ type: "text", text: `Optionally, when the export is finished, the exported data can also be accessed under path - "${exportPath}"`, @@ -84,4 +84,8 @@ export class ExportTool extends MongoDBToolBase { content: toolCallContent, }; } + + private isServerRunningLocally(): boolean { + return this.config.transport === "stdio" || ["127.0.0.1", "localhost"].includes(this.config.httpHost); + } } From 5bcac973b8e6b109ff3a60ed29bf61bd4f61cc14 Mon Sep 17 00:00:00 2001 From: Himanshu Singh Date: Fri, 8 Aug 2025 17:40:31 +0200 Subject: [PATCH 39/50] chore: ask for export description in the tool itself --- src/common/exportsManager.ts | 9 ++++- src/resources/common/exportedData.ts | 30 ++++----------- src/tools/mongodb/read/export.ts | 11 +++--- .../resources/exportedData.test.ts | 6 +-- .../tools/mongodb/read/export.test.ts | 37 ++++++++++++++++--- 5 files changed, 56 insertions(+), 37 deletions(-) diff --git a/src/common/exportsManager.ts b/src/common/exportsManager.ts index a0dec027..19a47d71 100644 --- a/src/common/exportsManager.ts +++ b/src/common/exportsManager.ts @@ -17,6 +17,7 @@ export type JSONExportFormat = z.infer; interface CommonExportData { exportName: string; + exportTitle: string; exportURI: string; exportPath: string; } @@ -53,7 +54,7 @@ type StoredExport = ReadyExport | InProgressExport; * * Ref Cursor: https://forum.cursor.com/t/cursor-mcp-resource-feature-support/50987 * JIRA: https://jira.mongodb.org/browse/MCP-104 */ -type AvailableExport = Pick; +type AvailableExport = Pick; export type ExportsManagerConfig = Pick; @@ -86,8 +87,9 @@ export class ExportsManager extends EventEmitter { !isExportExpired(storedExport.exportCreatedAt, this.config.exportTimeoutMs) ); }) - .map(({ exportName, exportURI, exportPath }) => ({ + .map(({ exportName, exportTitle, exportURI, exportPath }) => ({ exportName, + exportTitle, exportURI, exportPath, })); @@ -158,10 +160,12 @@ export class ExportsManager extends EventEmitter { public createJSONExport({ input, exportName, + exportTitle, jsonExportFormat, }: { input: FindCursor; exportName: string; + exportTitle: string; jsonExportFormat: JSONExportFormat; }): AvailableExport { try { @@ -174,6 +178,7 @@ export class ExportsManager extends EventEmitter { const exportFilePath = path.join(this.exportsDirectoryPath, exportNameWithExtension); const inProgressExport: InProgressExport = (this.storedExports[exportNameWithExtension] = { exportName: exportNameWithExtension, + exportTitle, exportPath: exportFilePath, exportURI: exportURI, exportStatus: "in-progress", diff --git a/src/resources/common/exportedData.ts b/src/resources/common/exportedData.ts index 5f59311b..67a3c014 100644 --- a/src/resources/common/exportedData.ts +++ b/src/resources/common/exportedData.ts @@ -49,12 +49,14 @@ export class ExportedData { private listResourcesCallback: ListResourcesCallback = () => { try { return { - resources: this.session.exportsManager.availableExports.map(({ exportName, exportURI }) => ({ - name: exportName, - description: this.exportNameToDescription(exportName), - uri: exportURI, - mimeType: "application/json", - })), + resources: this.session.exportsManager.availableExports.map( + ({ exportName, exportTitle, exportURI }) => ({ + name: exportName, + description: exportTitle, + uri: exportURI, + mimeType: "application/json", + }) + ), }; } catch (error) { this.session.logger.error({ @@ -112,20 +114,4 @@ export class ExportedData { }; } }; - - private exportNameToDescription(exportName: string): string { - const match = exportName.match(/^(.+)\.(\d+)\.(.+)\.json$/); - if (!match) return "Exported data for an unknown namespace."; - - const [, namespace, timestamp] = match; - if (!namespace) { - return "Exported data for an unknown namespace."; - } - - if (!timestamp) { - return `Export from ${namespace}.`; - } - - return `Export from ${namespace} done on ${new Date(parseInt(timestamp)).toLocaleString()}`; - } } diff --git a/src/tools/mongodb/read/export.ts b/src/tools/mongodb/read/export.ts index d264f81b..b61d77ea 100644 --- a/src/tools/mongodb/read/export.ts +++ b/src/tools/mongodb/read/export.ts @@ -10,6 +10,7 @@ export class ExportTool extends MongoDBToolBase { public name = "export"; protected description = "Export a collection data or query results in the specified EJSON format."; protected argsShape = { + exportTitle: z.string().describe("A short description to uniquely identify the export."), ...DbOperationArgs, ...FindArgs, limit: z.number().optional().describe("The maximum number of documents to return"), @@ -33,6 +34,7 @@ export class ExportTool extends MongoDBToolBase { projection, sort, limit, + exportTitle, }: ToolArgs): Promise { const provider = await this.ensureConnected(); const findCursor = provider.find(database, collection, filter ?? {}, { @@ -42,15 +44,14 @@ export class ExportTool extends MongoDBToolBase { promoteValues: false, bsonRegExp: true, }); - // The format is namespace.date.objectid.json - // - namespace to identify which namespace the export belongs to - // - date to identify when the export was generated - // - objectid for uniqueness of the names - const exportName = `${database}.${collection}.${Date.now()}.${new ObjectId().toString()}.json`; + const exportName = `${database}.${collection}.${new ObjectId().toString()}.json`; const { exportURI, exportPath } = this.session.exportsManager.createJSONExport({ input: findCursor, exportName, + exportTitle: + exportTitle || + `Export for namespace ${database}.${collection} requested on ${new Date().toLocaleString()}`, jsonExportFormat, }); const toolCallContent: CallToolResult["content"] = [ diff --git a/tests/integration/resources/exportedData.test.ts b/tests/integration/resources/exportedData.test.ts index b5e6039a..bc61b78e 100644 --- a/tests/integration/resources/exportedData.test.ts +++ b/tests/integration/resources/exportedData.test.ts @@ -65,7 +65,7 @@ describeWithMongoDB( await integration.connectMcpClient(); const exportResponse = await integration.mcpClient().callTool({ name: "export", - arguments: { database: "db", collection: "coll" }, + arguments: { database: "db", collection: "coll", exportTitle: "Export for db.coll" }, }); const exportedResourceURI = (exportResponse as CallToolResult).content.find( @@ -89,7 +89,7 @@ describeWithMongoDB( await integration.connectMcpClient(); const exportResponse = await integration.mcpClient().callTool({ name: "export", - arguments: { database: "db", collection: "coll" }, + arguments: { database: "db", collection: "coll", exportTitle: "Export for db.coll" }, }); const content = exportResponse.content as CallToolResult["content"]; const exportURI = contentWithResourceURILink(content)?.uri as string; @@ -112,7 +112,7 @@ describeWithMongoDB( await integration.connectMcpClient(); const exportResponse = await integration.mcpClient().callTool({ name: "export", - arguments: { database: "big", collection: "coll" }, + arguments: { database: "big", collection: "coll", exportTitle: "Export for big.coll" }, }); const content = exportResponse.content as CallToolResult["content"]; const exportURI = contentWithResourceURILink(content)?.uri as string; diff --git a/tests/integration/tools/mongodb/read/export.test.ts b/tests/integration/tools/mongodb/read/export.test.ts index 0a30f45a..343f3ef4 100644 --- a/tests/integration/tools/mongodb/read/export.test.ts +++ b/tests/integration/tools/mongodb/read/export.test.ts @@ -56,7 +56,12 @@ describeWithMongoDB( "Export a collection data or query results in the specified EJSON format.", [ ...databaseCollectionParameters, - + { + name: "exportTitle", + description: "A short description to uniquely identify the export.", + type: "string", + required: true, + }, { name: "filter", description: "The query filter, matching the syntax of the query argument of db.collection.find()", @@ -117,7 +122,11 @@ describeWithMongoDB( it("when provided with incorrect namespace, export should have empty data", async function () { const response = await integration.mcpClient().callTool({ name: "export", - arguments: { database: "non-existent", collection: "foos" }, + arguments: { + database: "non-existent", + collection: "foos", + exportTitle: "Export for non-existent.foos", + }, }); const content = response.content as CallToolResult["content"]; const exportURI = contentWithResourceURILink(content)?.uri as string; @@ -152,7 +161,11 @@ describeWithMongoDB( await integration.connectMcpClient(); const response = await integration.mcpClient().callTool({ name: "export", - arguments: { database: integration.randomDbName(), collection: "foo" }, + arguments: { + database: integration.randomDbName(), + collection: "foo", + exportTitle: `Export for ${integration.randomDbName()}.foo`, + }, }); const content = response.content as CallToolResult["content"]; const exportURI = contentWithResourceURILink(content)?.uri as string; @@ -176,7 +189,12 @@ describeWithMongoDB( await integration.connectMcpClient(); const response = await integration.mcpClient().callTool({ name: "export", - arguments: { database: integration.randomDbName(), collection: "foo", filter: { name: "foo" } }, + arguments: { + database: integration.randomDbName(), + collection: "foo", + filter: { name: "foo" }, + exportTitle: `Export for ${integration.randomDbName()}.foo`, + }, }); const content = response.content as CallToolResult["content"]; const exportURI = contentWithResourceURILink(content)?.uri as string; @@ -199,7 +217,12 @@ describeWithMongoDB( await integration.connectMcpClient(); const response = await integration.mcpClient().callTool({ name: "export", - arguments: { database: integration.randomDbName(), collection: "foo", limit: 1 }, + arguments: { + database: integration.randomDbName(), + collection: "foo", + limit: 1, + exportTitle: `Export for ${integration.randomDbName()}.foo`, + }, }); const content = response.content as CallToolResult["content"]; const exportURI = contentWithResourceURILink(content)?.uri as string; @@ -227,6 +250,7 @@ describeWithMongoDB( collection: "foo", limit: 1, sort: { longNumber: 1 }, + exportTitle: `Export for ${integration.randomDbName()}.foo`, }, }); const content = response.content as CallToolResult["content"]; @@ -255,6 +279,7 @@ describeWithMongoDB( collection: "foo", limit: 1, projection: { _id: 0, name: 1 }, + exportTitle: `Export for ${integration.randomDbName()}.foo`, }, }); const content = response.content as CallToolResult["content"]; @@ -287,6 +312,7 @@ describeWithMongoDB( limit: 1, projection: { _id: 0 }, jsonExportFormat: "relaxed", + exportTitle: `Export for ${integration.randomDbName()}.foo`, }, }); const content = response.content as CallToolResult["content"]; @@ -320,6 +346,7 @@ describeWithMongoDB( limit: 1, projection: { _id: 0 }, jsonExportFormat: "canonical", + exportTitle: `Export for ${integration.randomDbName()}.foo`, }, }); const content = response.content as CallToolResult["content"]; From 5d7657e5b0b5c96dad11db4b988600ed958edaff Mon Sep 17 00:00:00 2001 From: Himanshu Singh Date: Fri, 8 Aug 2025 17:49:51 +0200 Subject: [PATCH 40/50] chore: also apply exportTitle to tests --- tests/unit/common/exportsManager.test.ts | 14 ++++++++++++++ 1 file changed, 14 insertions(+) diff --git a/tests/unit/common/exportsManager.test.ts b/tests/unit/common/exportsManager.test.ts index df1ffb36..a02f800d 100644 --- a/tests/unit/common/exportsManager.test.ts +++ b/tests/unit/common/exportsManager.test.ts @@ -141,6 +141,7 @@ describe("ExportsManager unit test", () => { manager.createJSONExport({ input: createDummyFindCursorWithDelay([{ name: "Test1" }], 1000).cursor, exportName: exportName1, + exportTitle: "Some export", jsonExportFormat: "relaxed", }); @@ -150,6 +151,7 @@ describe("ExportsManager unit test", () => { manager.createJSONExport({ input: cursor, exportName: exportName2, + exportTitle: "Some export", jsonExportFormat: "relaxed", }); @@ -176,6 +178,7 @@ describe("ExportsManager unit test", () => { manager.createJSONExport({ input: cursor, exportName, + exportTitle: "Some export", jsonExportFormat: "relaxed", }); // note that we do not wait for cursor close @@ -190,6 +193,7 @@ describe("ExportsManager unit test", () => { manager.createJSONExport({ input: cursor, exportName, + exportTitle: "Some export", jsonExportFormat: "relaxed", }); await cursorCloseNotification; @@ -225,6 +229,7 @@ describe("ExportsManager unit test", () => { manager.createJSONExport({ input: cursor, exportName, + exportTitle: "Some export", jsonExportFormat: "relaxed", }) ).toThrow(); @@ -235,12 +240,14 @@ describe("ExportsManager unit test", () => { manager.createJSONExport({ input: cursor, exportName, + exportTitle: "Some export", jsonExportFormat: "relaxed", }); expect(() => manager.createJSONExport({ input: cursor, exportName, + exportTitle: "Some export", jsonExportFormat: "relaxed", }) ).toThrow(); @@ -254,6 +261,7 @@ describe("ExportsManager unit test", () => { manager.createJSONExport({ input: cursor, exportName, + exportTitle: "Some export", jsonExportFormat: "relaxed", }); await cursorCloseNotification; @@ -286,6 +294,7 @@ describe("ExportsManager unit test", () => { manager.createJSONExport({ input: cursor, exportName, + exportTitle: "Some export", jsonExportFormat: "relaxed", }); await cursorCloseNotification; @@ -320,6 +329,7 @@ describe("ExportsManager unit test", () => { manager.createJSONExport({ input: cursor, exportName, + exportTitle: "Some export", jsonExportFormat: "canonical", }); await cursorCloseNotification; @@ -386,6 +396,7 @@ describe("ExportsManager unit test", () => { manager.createJSONExport({ input: cursor, exportName, + exportTitle: "Some export", jsonExportFormat: "relaxed", }); await cursorCloseNotification; @@ -413,6 +424,7 @@ describe("ExportsManager unit test", () => { manager.createJSONExport({ input: cursor, exportName, + exportTitle: "Some export", jsonExportFormat: "relaxed", }); await cursorCloseNotification; @@ -460,6 +472,7 @@ describe("ExportsManager unit test", () => { manager.createJSONExport({ input: cursor, exportName, + exportTitle: "Some export", jsonExportFormat: "relaxed", }); @@ -486,6 +499,7 @@ describe("ExportsManager unit test", () => { manager.createJSONExport({ input: cursor, exportName, + exportTitle: "Some export", jsonExportFormat: "relaxed", }); await cursorCloseNotification; From 2136a88845c37e5c56a813d5776c194af28df9f3 Mon Sep 17 00:00:00 2001 From: Himanshu Singh Date: Fri, 8 Aug 2025 17:51:40 +0200 Subject: [PATCH 41/50] chore: move event subs to register in ExportedData resource --- src/resources/common/exportedData.ts | 17 ++++++++--------- 1 file changed, 8 insertions(+), 9 deletions(-) diff --git a/src/resources/common/exportedData.ts b/src/resources/common/exportedData.ts index 67a3c014..0c81f35e 100644 --- a/src/resources/common/exportedData.ts +++ b/src/resources/common/exportedData.ts @@ -14,15 +14,7 @@ export class ExportedData { private readonly uri = "exported-data://{exportName}"; private server?: Server; - constructor(private readonly session: Session) { - this.session.exportsManager.on("export-available", (uri: string): void => { - this.server?.sendResourceListChanged(); - this.server?.sendResourceUpdated(uri); - }); - this.session.exportsManager.on("export-expired", (): void => { - this.server?.sendResourceListChanged(); - }); - } + constructor(private readonly session: Session) {} public register(server: Server): void { this.server = server; @@ -44,6 +36,13 @@ export class ExportedData { { description: this.description }, this.readResourceCallback ); + this.session.exportsManager.on("export-available", (uri: string): void => { + server.sendResourceListChanged(); + server.sendResourceUpdated(uri); + }); + this.session.exportsManager.on("export-expired", (): void => { + server.sendResourceListChanged(); + }); } private listResourcesCallback: ListResourcesCallback = () => { From 701003f58a8fcdc7dd504010be0db5c7b78a7952 Mon Sep 17 00:00:00 2001 From: Himanshu Singh Date: Fri, 8 Aug 2025 18:09:46 +0200 Subject: [PATCH 42/50] chore: remove wasInitialized --- src/common/exportsManager.ts | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/src/common/exportsManager.ts b/src/common/exportsManager.ts index 19a47d71..19e825e4 100644 --- a/src/common/exportsManager.ts +++ b/src/common/exportsManager.ts @@ -64,7 +64,6 @@ type ExportsManagerEvents = { }; export class ExportsManager extends EventEmitter { - private wasInitialized: boolean = false; private isShuttingDown: boolean = false; private storedExports: Record = {}; private exportsCleanupInProgress: boolean = false; @@ -96,13 +95,11 @@ export class ExportsManager extends EventEmitter { } protected init(): void { - if (!this.wasInitialized) { + if (!this.exportsCleanupInterval) { this.exportsCleanupInterval = setInterval( () => void this.cleanupExpiredExports(), this.config.exportCleanupIntervalMs ); - - this.wasInitialized = true; } } From a3df32c79ec2e9e68858d0f1eb8cc403c5727cd6 Mon Sep 17 00:00:00 2001 From: Himanshu Singh Date: Fri, 8 Aug 2025 18:20:12 +0200 Subject: [PATCH 43/50] chore: disconnect sessionId and ExportsManager unique path id --- src/common/exportsManager.ts | 8 ++++++-- src/common/session.ts | 7 ++----- src/transports/base.ts | 5 +---- tests/integration/helpers.ts | 5 +---- tests/integration/telemetry.test.ts | 3 +-- tests/unit/common/exportsManager.test.ts | 12 +++++------- tests/unit/common/session.test.ts | 3 +-- tests/unit/resources/common/debug.test.ts | 3 +-- 8 files changed, 18 insertions(+), 28 deletions(-) diff --git a/src/common/exportsManager.ts b/src/common/exportsManager.ts index 19e825e4..2885fc30 100644 --- a/src/common/exportsManager.ts +++ b/src/common/exportsManager.ts @@ -4,7 +4,7 @@ import fs from "fs/promises"; import EventEmitter from "events"; import { createWriteStream } from "fs"; import { FindCursor } from "mongodb"; -import { EJSON, EJSONOptions } from "bson"; +import { EJSON, EJSONOptions, ObjectId } from "bson"; import { Transform } from "stream"; import { pipeline } from "stream/promises"; @@ -335,7 +335,11 @@ export class ExportsManager extends EventEmitter { } } - static init(sessionId: string, config: ExportsManagerConfig, logger: LoggerBase): ExportsManager { + static init( + config: ExportsManagerConfig, + logger: LoggerBase, + sessionId = new ObjectId().toString() + ): ExportsManager { const exportsDirectoryPath = path.join(config.exportsPath, sessionId); const exportsManager = new ExportsManager(exportsDirectoryPath, config, logger); exportsManager.init(); diff --git a/src/common/session.ts b/src/common/session.ts index 032375fd..444a747b 100644 --- a/src/common/session.ts +++ b/src/common/session.ts @@ -1,3 +1,4 @@ +import { ObjectId } from "bson"; import { ApiClient, ApiClientCredentials } from "./atlas/apiClient.js"; import { Implementation } from "@modelcontextprotocol/sdk/types.js"; import { CompositeLogger, LogId } from "./logger.js"; @@ -17,7 +18,6 @@ export interface SessionOptions { apiClientId?: string; apiClientSecret?: string; logger: CompositeLogger; - sessionId: string; exportsManager: ExportsManager; connectionManager: ConnectionManager; } @@ -30,7 +30,7 @@ export type SessionEvents = { }; export class Session extends EventEmitter { - readonly sessionId: string; + readonly sessionId: string = new ObjectId().toString(); readonly exportsManager: ExportsManager; readonly connectionManager: ConnectionManager; readonly apiClient: ApiClient; @@ -46,7 +46,6 @@ export class Session extends EventEmitter { apiClientId, apiClientSecret, logger, - sessionId, connectionManager, exportsManager, }: SessionOptions) { @@ -62,8 +61,6 @@ export class Session extends EventEmitter { : undefined; this.apiClient = new ApiClient({ baseUrl: apiBaseUrl, credentials }, logger); - - this.sessionId = sessionId; this.exportsManager = exportsManager; this.connectionManager = connectionManager; this.connectionManager.on("connection-succeeded", () => this.emit("connect")); diff --git a/src/transports/base.ts b/src/transports/base.ts index ac95ae1b..22a000cc 100644 --- a/src/transports/base.ts +++ b/src/transports/base.ts @@ -5,7 +5,6 @@ import { Session } from "../common/session.js"; import { Telemetry } from "../telemetry/telemetry.js"; import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js"; import { CompositeLogger, ConsoleLogger, DiskLogger, LoggerBase, McpLogger } from "../common/logger.js"; -import { ObjectId } from "bson"; import { ExportsManager } from "../common/exportsManager.js"; import { ConnectionManager } from "../common/connectionManager.js"; @@ -43,8 +42,7 @@ export abstract class TransportRunnerBase { } const logger = new CompositeLogger(...loggers); - const sessionId = new ObjectId().toString(); - const exportsManager = ExportsManager.init(sessionId, userConfig, logger); + const exportsManager = ExportsManager.init(userConfig, logger); const connectionManager = new ConnectionManager(); const session = new Session({ @@ -52,7 +50,6 @@ export abstract class TransportRunnerBase { apiClientId: userConfig.apiClientId, apiClientSecret: userConfig.apiClientSecret, logger, - sessionId, exportsManager, connectionManager, }); diff --git a/tests/integration/helpers.ts b/tests/integration/helpers.ts index 28193d10..738cbdfd 100644 --- a/tests/integration/helpers.ts +++ b/tests/integration/helpers.ts @@ -1,4 +1,3 @@ -import { ObjectId } from "bson"; import { Client } from "@modelcontextprotocol/sdk/client/index.js"; import { InMemoryTransport } from "./inMemoryTransport.js"; import { Server } from "../../src/server.js"; @@ -58,8 +57,7 @@ export function setupIntegrationTest(getUserConfig: () => UserConfig): Integrati ); const logger = new CompositeLogger(); - const sessionId = new ObjectId().toString(); - const exportsManager = ExportsManager.init(sessionId, userConfig, logger); + const exportsManager = ExportsManager.init(userConfig, logger); const connectionManager = new ConnectionManager(); const session = new Session({ @@ -67,7 +65,6 @@ export function setupIntegrationTest(getUserConfig: () => UserConfig): Integrati apiClientId: userConfig.apiClientId, apiClientSecret: userConfig.apiClientSecret, logger, - sessionId, exportsManager, connectionManager, }); diff --git a/tests/integration/telemetry.test.ts b/tests/integration/telemetry.test.ts index 4a34d156..95bc79c2 100644 --- a/tests/integration/telemetry.test.ts +++ b/tests/integration/telemetry.test.ts @@ -19,8 +19,7 @@ describe("Telemetry", () => { new Session({ apiBaseUrl: "", logger: new CompositeLogger(), - sessionId: "1FOO", - exportsManager: ExportsManager.init("1FOO", config, logger), + exportsManager: ExportsManager.init(config, logger), connectionManager: new ConnectionManager(), }), config diff --git a/tests/unit/common/exportsManager.test.ts b/tests/unit/common/exportsManager.test.ts index a02f800d..d00f2e40 100644 --- a/tests/unit/common/exportsManager.test.ts +++ b/tests/unit/common/exportsManager.test.ts @@ -113,12 +113,10 @@ describe("ExportsManager unit test", () => { beforeEach(async () => { await fs.mkdir(exportsManagerConfig.exportsPath, { recursive: true }); const logger = new CompositeLogger(); - const sessionId = new ObjectId().toString(); session = new Session({ apiBaseUrl: "", logger, - sessionId, - exportsManager: ExportsManager.init(sessionId, exportsManagerConfig, logger), + exportsManager: ExportsManager.init(exportsManagerConfig, logger), connectionManager: new ConnectionManager(), }); manager = session.exportsManager; @@ -460,13 +458,13 @@ describe("ExportsManager unit test", () => { it("should not clean up in-progress exports", async () => { const { exportName } = getExportNameAndPath(session.sessionId); const manager = ExportsManager.init( - session.sessionId, { ...exportsManagerConfig, exportTimeoutMs: 100, exportCleanupIntervalMs: 50, }, - new CompositeLogger() + new CompositeLogger(), + session.sessionId ); const { cursor } = createDummyFindCursorWithDelay([{ name: "Test" }], 2000); manager.createJSONExport({ @@ -488,13 +486,13 @@ describe("ExportsManager unit test", () => { it("should cleanup expired exports", async () => { const { exportName, exportPath, exportURI } = getExportNameAndPath(session.sessionId); const manager = ExportsManager.init( - session.sessionId, { ...exportsManagerConfig, exportTimeoutMs: 100, exportCleanupIntervalMs: 50, }, - new CompositeLogger() + new CompositeLogger(), + session.sessionId ); manager.createJSONExport({ input: cursor, diff --git a/tests/unit/common/session.test.ts b/tests/unit/common/session.test.ts index e1cce9a0..1d26d8d8 100644 --- a/tests/unit/common/session.test.ts +++ b/tests/unit/common/session.test.ts @@ -17,8 +17,7 @@ describe("Session", () => { apiClientId: "test-client-id", apiBaseUrl: "https://api.test.com", logger, - sessionId: "1FOO", - exportsManager: ExportsManager.init("1FOO", config, logger), + exportsManager: ExportsManager.init(config, logger), connectionManager: new ConnectionManager(), }); diff --git a/tests/unit/resources/common/debug.test.ts b/tests/unit/resources/common/debug.test.ts index 684adea4..8e798827 100644 --- a/tests/unit/resources/common/debug.test.ts +++ b/tests/unit/resources/common/debug.test.ts @@ -12,8 +12,7 @@ describe("debug resource", () => { const session = new Session({ apiBaseUrl: "", logger, - sessionId: "1FOO", - exportsManager: ExportsManager.init("1FOO", config, logger), + exportsManager: ExportsManager.init(config, logger), connectionManager: new ConnectionManager(), }); const telemetry = Telemetry.create(session, { ...config, telemetry: "disabled" }); From 85cb524f537a1c297e9b98991104514ad263322a Mon Sep 17 00:00:00 2001 From: Himanshu Singh Date: Fri, 8 Aug 2025 19:51:51 +0200 Subject: [PATCH 44/50] chore: track long running ops and abort them on close --- src/common/exportsManager.ts | 73 ++++++++++++++++++------ tests/unit/common/exportsManager.test.ts | 39 +++++++++++++ 2 files changed, 96 insertions(+), 16 deletions(-) diff --git a/src/common/exportsManager.ts b/src/common/exportsManager.ts index 2885fc30..c3f436f3 100644 --- a/src/common/exportsManager.ts +++ b/src/common/exportsManager.ts @@ -56,7 +56,11 @@ type StoredExport = ReadyExport | InProgressExport; * JIRA: https://jira.mongodb.org/browse/MCP-104 */ type AvailableExport = Pick; -export type ExportsManagerConfig = Pick; +export type ExportsManagerConfig = Pick & { + // The maximum number of milliseconds to wait for in-flight operations to + // settle before shutting down ExportsManager. + activeOpsDrainTimeoutMs?: number; +}; type ExportsManagerEvents = { "export-expired": [string]; @@ -64,10 +68,12 @@ type ExportsManagerEvents = { }; export class ExportsManager extends EventEmitter { - private isShuttingDown: boolean = false; private storedExports: Record = {}; private exportsCleanupInProgress: boolean = false; private exportsCleanupInterval?: NodeJS.Timeout; + private readonly shutdownController: AbortController = new AbortController(); + private readonly activeOperations: Set> = new Set(); + private readonly activeOpsDrainTimeoutMs: number; private constructor( private readonly exportsDirectoryPath: string, @@ -75,6 +81,7 @@ export class ExportsManager extends EventEmitter { private readonly logger: LoggerBase ) { super(); + this.activeOpsDrainTimeoutMs = this.config.activeOpsDrainTimeoutMs ?? 10_000; } public get availableExports(): AvailableExport[] { @@ -97,20 +104,19 @@ export class ExportsManager extends EventEmitter { protected init(): void { if (!this.exportsCleanupInterval) { this.exportsCleanupInterval = setInterval( - () => void this.cleanupExpiredExports(), + () => void this.trackOperation(this.cleanupExpiredExports()), this.config.exportCleanupIntervalMs ); } } - public async close(): Promise { - if (this.isShuttingDown) { + if (this.shutdownController.signal.aborted) { return; } - - this.isShuttingDown = true; try { clearInterval(this.exportsCleanupInterval); + this.shutdownController.abort(); + await this.waitForActiveOperationsToSettle(this.activeOpsDrainTimeoutMs); await fs.rm(this.exportsDirectoryPath, { force: true, recursive: true }); } catch (error) { this.logger.error({ @@ -140,7 +146,9 @@ export class ExportsManager extends EventEmitter { throw new Error("Requested export has expired!"); } - return await fs.readFile(exportPath, "utf8"); + return await this.trackOperation( + fs.readFile(exportPath, { encoding: "utf8", signal: this.shutdownController.signal }) + ); } catch (error) { this.logger.error({ id: LogId.exportReadError, @@ -181,7 +189,7 @@ export class ExportsManager extends EventEmitter { exportStatus: "in-progress", }); - void this.startExport({ input, jsonExportFormat, inProgressExport }); + void this.trackOperation(this.startExport({ input, jsonExportFormat, inProgressExport })); return inProgressExport; } catch (error) { this.logger.error({ @@ -206,11 +214,10 @@ export class ExportsManager extends EventEmitter { try { await fs.mkdir(this.exportsDirectoryPath, { recursive: true }); const outputStream = createWriteStream(inProgressExport.exportPath); - await pipeline([ - input.stream(), - this.docToEJSONStream(this.getEJSONOptionsForFormat(jsonExportFormat)), - outputStream, - ]); + await pipeline( + [input.stream(), this.docToEJSONStream(this.getEJSONOptionsForFormat(jsonExportFormat)), outputStream], + { signal: this.shutdownController.signal } + ); pipeSuccessful = true; } catch (error) { this.logger.error({ @@ -281,7 +288,7 @@ export class ExportsManager extends EventEmitter { } private async cleanupExpiredExports(): Promise { - if (this.exportsCleanupInProgress || this.isShuttingDown) { + if (this.exportsCleanupInProgress) { return; } @@ -291,6 +298,9 @@ export class ExportsManager extends EventEmitter { ); try { for (const { exportPath, exportCreatedAt, exportURI, exportName } of exportsForCleanup) { + if (this.shutdownController.signal.aborted) { + break; + } if (isExportExpired(exportCreatedAt, this.config.exportTimeoutMs)) { delete this.storedExports[exportName]; await this.silentlyRemoveExport( @@ -330,11 +340,42 @@ export class ExportsManager extends EventEmitter { } private assertIsNotShuttingDown(): void { - if (this.isShuttingDown) { + if (this.shutdownController.signal.aborted) { throw new Error("ExportsManager is shutting down."); } } + private async trackOperation(promise: Promise): Promise { + this.activeOperations.add(promise); + try { + return await promise; + } finally { + this.activeOperations.delete(promise); + } + } + + private async waitForActiveOperationsToSettle(timeoutMs: number): Promise { + const pendingPromises = Array.from(this.activeOperations); + if (pendingPromises.length === 0) { + return; + } + let timedOut = false; + const timeoutPromise = new Promise((resolve) => + setTimeout(() => { + timedOut = true; + resolve(); + }, timeoutMs) + ); + await Promise.race([Promise.allSettled(pendingPromises), timeoutPromise]); + if (timedOut && this.activeOperations.size > 0) { + this.logger.error({ + id: LogId.exportCloseError, + context: `Close timed out waiting for ${this.activeOperations.size} operation(s) to settle`, + message: "Proceeding to force cleanup after timeout", + }); + } + } + static init( config: ExportsManagerConfig, logger: LoggerBase, diff --git a/tests/unit/common/exportsManager.test.ts b/tests/unit/common/exportsManager.test.ts index d00f2e40..6c74c404 100644 --- a/tests/unit/common/exportsManager.test.ts +++ b/tests/unit/common/exportsManager.test.ts @@ -514,6 +514,45 @@ describe("ExportsManager unit test", () => { expect(await fileExists(exportPath)).toEqual(false); }); }); + + describe("#close", () => { + it("should abort ongoing export and remove partial file", async () => { + const { exportName, exportPath } = getExportNameAndPath(session.sessionId); + const { cursor } = createDummyFindCursorWithDelay([{ name: "Test" }], 2000); + manager.createJSONExport({ + input: cursor, + exportName, + exportTitle: "Some export", + jsonExportFormat: "relaxed", + }); + // Give the pipeline a brief moment to start and create the file + await timeout(50); + + await manager.close(); + + await expect(fileExists(exportPath)).resolves.toEqual(false); + }); + + it("should timeout shutdown wait when operations never settle", async () => { + await manager.close(); + const logger = new CompositeLogger(); + manager = ExportsManager.init({ ...exportsManagerConfig, activeOpsDrainTimeoutMs: 50 }, logger); + // eslint-disable-next-line @typescript-eslint/no-explicit-any, @typescript-eslint/no-unsafe-member-access, @typescript-eslint/no-unsafe-call + (manager as any).activeOperations.add( + new Promise(() => { + /* never resolves */ + }) + ); + + const start = Date.now(); + await manager.close(); + const elapsed = Date.now() - start; + + // Should not block indefinitely; should be close to the configured timeout but well under 1s + expect(elapsed).toBeGreaterThanOrEqual(45); + expect(elapsed).toBeLessThan(1000); + }); + }); }); describe("#ensureExtension", () => { From c0031ffaa8a9150ec8595563339916f961d8728b Mon Sep 17 00:00:00 2001 From: Himanshu Singh Date: Fri, 8 Aug 2025 23:55:07 +0200 Subject: [PATCH 45/50] chore: implement read write lock --- package-lock.json | 10 ++ package.json | 1 + src/common/exportsManager.ts | 211 ++++++++++++++--------- src/tools/mongodb/read/export.ts | 2 +- tests/unit/common/exportsManager.test.ts | 61 ++++--- 5 files changed, 179 insertions(+), 106 deletions(-) diff --git a/package-lock.json b/package-lock.json index d6e7d7dd..19772239 100644 --- a/package-lock.json +++ b/package-lock.json @@ -15,6 +15,7 @@ "@mongodb-js/devtools-proxy-support": "^0.5.1", "@mongosh/service-provider-node-driver": "^3.10.2", "@vitest/eslint-plugin": "^1.3.4", + "async-rwlock": "^1.1.1", "bson": "^6.10.4", "express": "^5.1.0", "lru-cache": "^11.1.0", @@ -6098,6 +6099,15 @@ "js-tokens": "^9.0.1" } }, + "node_modules/async-rwlock": { + "version": "1.1.1", + "resolved": "https://registry.npmjs.org/async-rwlock/-/async-rwlock-1.1.1.tgz", + "integrity": "sha512-K4ecpHLAc0Jul4dMb1KLpukblQmHxD5/HNgkTuO3sQ9oLLtVENCJVk7+b0wMj1K89cqnjGkTsAvOb83lCfoKwA==", + "license": "MIT", + "engines": { + "node": ">=8.0.0" + } + }, "node_modules/asynckit": { "version": "0.4.0", "resolved": "https://registry.npmjs.org/asynckit/-/asynckit-0.4.0.tgz", diff --git a/package.json b/package.json index 7bcd5a1e..df0b4f28 100644 --- a/package.json +++ b/package.json @@ -96,6 +96,7 @@ "@mongodb-js/devtools-proxy-support": "^0.5.1", "@mongosh/service-provider-node-driver": "^3.10.2", "@vitest/eslint-plugin": "^1.3.4", + "async-rwlock": "^1.1.1", "bson": "^6.10.4", "express": "^5.1.0", "lru-cache": "^11.1.0", diff --git a/src/common/exportsManager.ts b/src/common/exportsManager.ts index c3f436f3..6f632ced 100644 --- a/src/common/exportsManager.ts +++ b/src/common/exportsManager.ts @@ -7,10 +7,11 @@ import { FindCursor } from "mongodb"; import { EJSON, EJSONOptions, ObjectId } from "bson"; import { Transform } from "stream"; import { pipeline } from "stream/promises"; +import { MongoLogId } from "mongodb-log-writer"; +import { RWLock } from "async-rwlock"; import { UserConfig } from "./config.js"; import { LoggerBase, LogId } from "./logger.js"; -import { MongoLogId } from "mongodb-log-writer"; export const jsonExportFormat = z.enum(["relaxed", "canonical"]); export type JSONExportFormat = z.infer; @@ -60,6 +61,12 @@ export type ExportsManagerConfig = Pick { private readonly shutdownController: AbortController = new AbortController(); private readonly activeOperations: Set> = new Set(); private readonly activeOpsDrainTimeoutMs: number; + private readonly readTimeoutMs: number; + private readonly writeTimeoutMs: number; + private readonly exportLocks: Map = new Map(); private constructor( private readonly exportsDirectoryPath: string, @@ -82,6 +92,8 @@ export class ExportsManager extends EventEmitter { ) { super(); this.activeOpsDrainTimeoutMs = this.config.activeOpsDrainTimeoutMs ?? 10_000; + this.readTimeoutMs = this.config.readTimeout ?? 30_0000; // 30 seconds is the default timeout for an MCP request + this.writeTimeoutMs = this.config.writeTimeout ?? 120_000; // considering that writes can take time } public get availableExports(): AvailableExport[] { @@ -131,24 +143,24 @@ export class ExportsManager extends EventEmitter { try { this.assertIsNotShuttingDown(); exportName = decodeURIComponent(exportName); - const exportHandle = this.storedExports[exportName]; - if (!exportHandle) { - throw new Error("Requested export has either expired or does not exist!"); - } - - if (exportHandle.exportStatus === "in-progress") { - throw new Error("Requested export is still being generated!"); - } + return await this.withLock(exportName, "read", false, async (): Promise => { + const exportHandle = this.storedExports[exportName]; + if (!exportHandle) { + throw new Error("Requested export has either expired or does not exist!"); + } - const { exportPath, exportCreatedAt } = exportHandle; + // This won't happen anymore because of lock synchronization but + // keeping it here to make TS happy. + if (exportHandle.exportStatus === "in-progress") { + throw new Error("Requested export is still being generated!"); + } - if (isExportExpired(exportCreatedAt, this.config.exportTimeoutMs)) { - throw new Error("Requested export has expired!"); - } + const { exportPath } = exportHandle; - return await this.trackOperation( - fs.readFile(exportPath, { encoding: "utf8", signal: this.shutdownController.signal }) - ); + return await this.trackOperation( + fs.readFile(exportPath, { encoding: "utf8", signal: this.shutdownController.signal }) + ); + }); } catch (error) { this.logger.error({ id: LogId.exportReadError, @@ -162,7 +174,7 @@ export class ExportsManager extends EventEmitter { } } - public createJSONExport({ + public async createJSONExport({ input, exportName, exportTitle, @@ -172,25 +184,27 @@ export class ExportsManager extends EventEmitter { exportName: string; exportTitle: string; jsonExportFormat: JSONExportFormat; - }): AvailableExport { + }): Promise { try { this.assertIsNotShuttingDown(); const exportNameWithExtension = validateExportName(ensureExtension(exportName, "json")); - if (this.storedExports[exportNameWithExtension]) { - throw new Error("Export with same name is either already available or being generated."); - } - const exportURI = `exported-data://${encodeURIComponent(exportNameWithExtension)}`; - const exportFilePath = path.join(this.exportsDirectoryPath, exportNameWithExtension); - const inProgressExport: InProgressExport = (this.storedExports[exportNameWithExtension] = { - exportName: exportNameWithExtension, - exportTitle, - exportPath: exportFilePath, - exportURI: exportURI, - exportStatus: "in-progress", - }); + return await this.withLock(exportNameWithExtension, "write", false, (): AvailableExport => { + if (this.storedExports[exportNameWithExtension]) { + throw new Error("Export with same name is either already available or being generated."); + } + const exportURI = `exported-data://${encodeURIComponent(exportNameWithExtension)}`; + const exportFilePath = path.join(this.exportsDirectoryPath, exportNameWithExtension); + const inProgressExport: InProgressExport = (this.storedExports[exportNameWithExtension] = { + exportName: exportNameWithExtension, + exportTitle, + exportPath: exportFilePath, + exportURI: exportURI, + exportStatus: "in-progress", + }); - void this.trackOperation(this.startExport({ input, jsonExportFormat, inProgressExport })); - return inProgressExport; + void this.trackOperation(this.startExport({ input, jsonExportFormat, inProgressExport })); + return inProgressExport; + }); } catch (error) { this.logger.error({ id: LogId.exportCreationError, @@ -211,40 +225,46 @@ export class ExportsManager extends EventEmitter { inProgressExport: InProgressExport; }): Promise { let pipeSuccessful = false; - 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 } - ); - pipeSuccessful = true; - } catch (error) { - this.logger.error({ - id: LogId.exportCreationError, - context: `Error when generating JSON export for ${inProgressExport.exportName}`, - message: error instanceof Error ? error.message : String(error), - }); + await this.withLock(inProgressExport.exportName, "write", false, async (): Promise => { + 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 } + ); + pipeSuccessful = true; + } catch (error) { + this.logger.error({ + id: LogId.exportCreationError, + context: `Error when generating JSON export for ${inProgressExport.exportName}`, + message: error instanceof Error ? error.message : String(error), + }); - // If the pipeline errors out then we might end up with - // partial and incorrect export so we remove it entirely. - await this.silentlyRemoveExport( - inProgressExport.exportPath, - LogId.exportCreationCleanupError, - `Error when removing incomplete export ${inProgressExport.exportName}` - ); - delete this.storedExports[inProgressExport.exportName]; - } finally { - if (pipeSuccessful) { - this.storedExports[inProgressExport.exportName] = { - ...inProgressExport, - exportCreatedAt: Date.now(), - exportStatus: "ready", - }; - this.emit("export-available", inProgressExport.exportURI); + // If the pipeline errors out then we might end up with + // partial and incorrect export so we remove it entirely. + await this.silentlyRemoveExport( + inProgressExport.exportPath, + LogId.exportCreationCleanupError, + `Error when removing incomplete export ${inProgressExport.exportName}` + ); + delete this.storedExports[inProgressExport.exportName]; + } finally { + if (pipeSuccessful) { + this.storedExports[inProgressExport.exportName] = { + ...inProgressExport, + exportCreatedAt: Date.now(), + exportStatus: "ready", + }; + this.emit("export-available", inProgressExport.exportURI); + } + void input.close(); } - void input.close(); - } + }); } private getEJSONOptionsForFormat(format: JSONExportFormat): EJSONOptions | undefined { @@ -293,24 +313,26 @@ export class ExportsManager extends EventEmitter { } this.exportsCleanupInProgress = true; - const exportsForCleanup = Object.values({ ...this.storedExports }).filter( - (storedExport): storedExport is ReadyExport => storedExport.exportStatus === "ready" - ); try { - for (const { exportPath, exportCreatedAt, exportURI, exportName } of exportsForCleanup) { - if (this.shutdownController.signal.aborted) { - break; - } - if (isExportExpired(exportCreatedAt, this.config.exportTimeoutMs)) { - delete this.storedExports[exportName]; - await this.silentlyRemoveExport( - exportPath, - LogId.exportCleanupError, - `Considerable error when removing export ${exportName}` - ); - this.emit("export-expired", exportURI); - } - } + const exportsForCleanup = Object.values({ ...this.storedExports }).filter( + (storedExport): storedExport is ReadyExport => storedExport.exportStatus === "ready" + ); + + await Promise.allSettled( + exportsForCleanup.map(async ({ exportPath, exportCreatedAt, exportURI, exportName }) => { + if (isExportExpired(exportCreatedAt, this.config.exportTimeoutMs)) { + await this.withLock(exportName, "write", true, async (): Promise => { + delete this.storedExports[exportName]; + await this.silentlyRemoveExport( + exportPath, + LogId.exportCleanupError, + `Considerable error when removing export ${exportName}` + ); + this.emit("export-expired", exportURI); + }); + } + }) + ); } catch (error) { this.logger.error({ id: LogId.exportCleanupError, @@ -345,6 +367,33 @@ export class ExportsManager extends EventEmitter { } } + private async withLock( + exportName: string, + mode: "read" | "write", + finalize: boolean, + fn: () => T | Promise + ): Promise { + let lock = this.exportLocks.get(exportName); + if (!lock) { + lock = new RWLock(); + this.exportLocks.set(exportName, lock); + } + + try { + if (mode === "read") { + await lock.readLock(this.readTimeoutMs); + } else { + await lock.writeLock(this.writeTimeoutMs); + } + return await fn(); + } finally { + lock.unlock(); + if (finalize) { + this.exportLocks.delete(exportName); + } + } + } + private async trackOperation(promise: Promise): Promise { this.activeOperations.add(promise); try { diff --git a/src/tools/mongodb/read/export.ts b/src/tools/mongodb/read/export.ts index b61d77ea..9eaacba2 100644 --- a/src/tools/mongodb/read/export.ts +++ b/src/tools/mongodb/read/export.ts @@ -46,7 +46,7 @@ export class ExportTool extends MongoDBToolBase { }); const exportName = `${database}.${collection}.${new ObjectId().toString()}.json`; - const { exportURI, exportPath } = this.session.exportsManager.createJSONExport({ + const { exportURI, exportPath } = await this.session.exportsManager.createJSONExport({ input: findCursor, exportName, exportTitle: diff --git a/tests/unit/common/exportsManager.test.ts b/tests/unit/common/exportsManager.test.ts index 6c74c404..90c44e9e 100644 --- a/tests/unit/common/exportsManager.test.ts +++ b/tests/unit/common/exportsManager.test.ts @@ -136,7 +136,7 @@ describe("ExportsManager unit test", () => { it("should list only the exports that are in ready state", async () => { // This export will finish in at-least 1 second const { exportName: exportName1 } = getExportNameAndPath(session.sessionId); - manager.createJSONExport({ + await manager.createJSONExport({ input: createDummyFindCursorWithDelay([{ name: "Test1" }], 1000).cursor, exportName: exportName1, exportTitle: "Some export", @@ -146,7 +146,7 @@ describe("ExportsManager unit test", () => { // This export will finish way sooner than the first one const { exportName: exportName2 } = getExportNameAndPath(session.sessionId); const { cursor, cursorCloseNotification } = createDummyFindCursor([{ name: "Test1" }]); - manager.createJSONExport({ + await manager.createJSONExport({ input: cursor, exportName: exportName2, exportTitle: "Some export", @@ -170,25 +170,38 @@ describe("ExportsManager unit test", () => { await expect(() => manager.readExport("name")).rejects.toThrow(); }); - it("should throw if the resource is still being generated", async () => { + it("should wait if resource is still being generated", async () => { const { exportName } = getExportNameAndPath(session.sessionId); - const { cursor } = createDummyFindCursorWithDelay([{ name: "Test1" }], 100); - manager.createJSONExport({ + const { cursor } = createDummyFindCursorWithDelay([{ name: "Test1" }], 200); + await manager.createJSONExport({ input: cursor, exportName, exportTitle: "Some export", jsonExportFormat: "relaxed", }); // note that we do not wait for cursor close - await expect(() => manager.readExport(exportName)).rejects.toThrow( - "Requested export is still being generated!" - ); + expect(await manager.readExport(exportName)).toEqual(JSON.stringify([{ name: "Test1" }])); + }); + + it("should allow concurrent reads of the same resource", async () => { + const { exportName } = getExportNameAndPath(session.sessionId); + const { cursor } = createDummyFindCursorWithDelay([{ name: "Test1" }], 200); + await manager.createJSONExport({ + input: cursor, + exportName, + exportTitle: "Some export", + jsonExportFormat: "relaxed", + }); + // note that we do not wait for cursor close + expect( + await Promise.all([await manager.readExport(exportName), await manager.readExport(exportName)]) + ).toEqual([JSON.stringify([{ name: "Test1" }]), JSON.stringify([{ name: "Test1" }])]); }); it("should return the resource content if the resource is ready to be consumed", async () => { const { exportName } = getExportNameAndPath(session.sessionId); const { cursor, cursorCloseNotification } = createDummyFindCursor([]); - manager.createJSONExport({ + await manager.createJSONExport({ input: cursor, exportName, exportTitle: "Some export", @@ -220,35 +233,35 @@ describe("ExportsManager unit test", () => { ({ exportName, exportPath, exportURI } = getExportNameAndPath(session.sessionId)); }); - it("should throw if the manager is shutting down", () => { + it("should throw if the manager is shutting down", async () => { const { cursor } = createDummyFindCursor([]); void manager.close(); - expect(() => + await expect(() => manager.createJSONExport({ input: cursor, exportName, exportTitle: "Some export", jsonExportFormat: "relaxed", }) - ).toThrow(); + ).rejects.toThrow(); }); - it("should throw if the same name export is requested more than once", () => { + it("should throw if the same name export is requested more than once", async () => { const { cursor } = createDummyFindCursorWithDelay([{ name: 1 }, { name: 2 }], 100); - manager.createJSONExport({ + await manager.createJSONExport({ input: cursor, exportName, exportTitle: "Some export", jsonExportFormat: "relaxed", }); - expect(() => + await expect(() => manager.createJSONExport({ input: cursor, exportName, exportTitle: "Some export", jsonExportFormat: "relaxed", }) - ).toThrow(); + ).rejects.toThrow(); }); describe("when cursor is empty", () => { @@ -256,7 +269,7 @@ describe("ExportsManager unit test", () => { const { cursor, cursorCloseNotification } = createDummyFindCursor([]); const emitSpy = vi.spyOn(manager, "emit"); - manager.createJSONExport({ + await manager.createJSONExport({ input: cursor, exportName, exportTitle: "Some export", @@ -289,7 +302,7 @@ describe("ExportsManager unit test", () => { ])("$cond", ({ exportName }) => { it("should export relaxed json, update available exports and emit export-available event", async () => { const emitSpy = vi.spyOn(manager, "emit"); - manager.createJSONExport({ + await manager.createJSONExport({ input: cursor, exportName, exportTitle: "Some export", @@ -324,7 +337,7 @@ describe("ExportsManager unit test", () => { ])("$cond", ({ exportName }) => { it("should export canonical json, update available exports and emit export-available event", async () => { const emitSpy = vi.spyOn(manager, "emit"); - manager.createJSONExport({ + await manager.createJSONExport({ input: cursor, exportName, exportTitle: "Some export", @@ -391,7 +404,7 @@ describe("ExportsManager unit test", () => { }, }); }; - manager.createJSONExport({ + await manager.createJSONExport({ input: cursor, exportName, exportTitle: "Some export", @@ -419,7 +432,7 @@ describe("ExportsManager unit test", () => { } return Promise.resolve(); }); - manager.createJSONExport({ + await manager.createJSONExport({ input: cursor, exportName, exportTitle: "Some export", @@ -467,7 +480,7 @@ describe("ExportsManager unit test", () => { session.sessionId ); const { cursor } = createDummyFindCursorWithDelay([{ name: "Test" }], 2000); - manager.createJSONExport({ + await manager.createJSONExport({ input: cursor, exportName, exportTitle: "Some export", @@ -494,7 +507,7 @@ describe("ExportsManager unit test", () => { new CompositeLogger(), session.sessionId ); - manager.createJSONExport({ + await manager.createJSONExport({ input: cursor, exportName, exportTitle: "Some export", @@ -519,7 +532,7 @@ describe("ExportsManager unit test", () => { it("should abort ongoing export and remove partial file", async () => { const { exportName, exportPath } = getExportNameAndPath(session.sessionId); const { cursor } = createDummyFindCursorWithDelay([{ name: "Test" }], 2000); - manager.createJSONExport({ + await manager.createJSONExport({ input: cursor, exportName, exportTitle: "Some export", From 34aff9a17f256e87528fc0f9d9ac60fdec10d5e2 Mon Sep 17 00:00:00 2001 From: Himanshu Singh Date: Mon, 11 Aug 2025 11:43:13 +0200 Subject: [PATCH 46/50] chore: improve readability for locks and in-progress operations --- src/common/exportsManager.ts | 284 +++++++++++++---------- src/common/logger.ts | 1 + tests/unit/common/exportsManager.test.ts | 161 +++++++------ 3 files changed, 255 insertions(+), 191 deletions(-) diff --git a/src/common/exportsManager.ts b/src/common/exportsManager.ts index 6f632ced..8403f585 100644 --- a/src/common/exportsManager.ts +++ b/src/common/exportsManager.ts @@ -70,17 +70,18 @@ export type ExportsManagerConfig = Pick { private storedExports: Record = {}; private exportsCleanupInProgress: boolean = false; private exportsCleanupInterval?: NodeJS.Timeout; private readonly shutdownController: AbortController = new AbortController(); - private readonly activeOperations: Set> = new Set(); - private readonly activeOpsDrainTimeoutMs: number; private readonly readTimeoutMs: number; private readonly writeTimeoutMs: number; private readonly exportLocks: Map = new Map(); @@ -91,7 +92,6 @@ export class ExportsManager extends EventEmitter { private readonly logger: LoggerBase ) { super(); - this.activeOpsDrainTimeoutMs = this.config.activeOpsDrainTimeoutMs ?? 10_000; this.readTimeoutMs = this.config.readTimeout ?? 30_0000; // 30 seconds is the default timeout for an MCP request this.writeTimeoutMs = this.config.writeTimeout ?? 120_000; // considering that writes can take time } @@ -116,7 +116,7 @@ export class ExportsManager extends EventEmitter { protected init(): void { if (!this.exportsCleanupInterval) { this.exportsCleanupInterval = setInterval( - () => void this.trackOperation(this.cleanupExpiredExports()), + () => void this.cleanupExpiredExports(), this.config.exportCleanupIntervalMs ); } @@ -128,8 +128,8 @@ export class ExportsManager extends EventEmitter { try { clearInterval(this.exportsCleanupInterval); this.shutdownController.abort(); - await this.waitForActiveOperationsToSettle(this.activeOpsDrainTimeoutMs); await fs.rm(this.exportsDirectoryPath, { force: true, recursive: true }); + this.emit("closed"); } catch (error) { this.logger.error({ id: LogId.exportCloseError, @@ -143,33 +143,35 @@ export class ExportsManager extends EventEmitter { try { this.assertIsNotShuttingDown(); exportName = decodeURIComponent(exportName); - return await this.withLock(exportName, "read", false, async (): Promise => { - const exportHandle = this.storedExports[exportName]; - if (!exportHandle) { - throw new Error("Requested export has either expired or does not exist!"); - } + return await this.withLock( + { + exportName, + mode: "read", + callbackName: "readExport", + }, + async (): Promise => { + const exportHandle = this.storedExports[exportName]; + if (!exportHandle) { + throw new Error("Requested export has either expired or does not exist!"); + } - // This won't happen anymore because of lock synchronization but - // keeping it here to make TS happy. - if (exportHandle.exportStatus === "in-progress") { - throw new Error("Requested export is still being generated!"); - } + // This won't happen because of lock synchronization but + // keeping it here to make TS happy. + if (exportHandle.exportStatus === "in-progress") { + throw new Error("Requested export is still being generated!"); + } - const { exportPath } = exportHandle; + const { exportPath } = exportHandle; - return await this.trackOperation( - fs.readFile(exportPath, { encoding: "utf8", signal: this.shutdownController.signal }) - ); - }); + return fs.readFile(exportPath, { encoding: "utf8", signal: this.shutdownController.signal }); + } + ); } catch (error) { this.logger.error({ id: LogId.exportReadError, context: `Error when reading export - ${exportName}`, message: error instanceof Error ? error.message : String(error), }); - if ((error as NodeJS.ErrnoException).code === "ENOENT") { - throw new Error("Requested export does not exist!"); - } throw error; } } @@ -188,23 +190,32 @@ export class ExportsManager extends EventEmitter { try { this.assertIsNotShuttingDown(); const exportNameWithExtension = validateExportName(ensureExtension(exportName, "json")); - return await this.withLock(exportNameWithExtension, "write", false, (): AvailableExport => { - if (this.storedExports[exportNameWithExtension]) { - throw new Error("Export with same name is either already available or being generated."); - } - const exportURI = `exported-data://${encodeURIComponent(exportNameWithExtension)}`; - const exportFilePath = path.join(this.exportsDirectoryPath, exportNameWithExtension); - const inProgressExport: InProgressExport = (this.storedExports[exportNameWithExtension] = { + return await this.withLock( + { exportName: exportNameWithExtension, - exportTitle, - exportPath: exportFilePath, - exportURI: exportURI, - exportStatus: "in-progress", - }); - - void this.trackOperation(this.startExport({ input, jsonExportFormat, inProgressExport })); - return inProgressExport; - }); + mode: "write", + callbackName: "createJSONExport", + }, + (): Promise => { + if (this.storedExports[exportNameWithExtension]) { + return Promise.reject( + new Error("Export with same name is either already available or being generated.") + ); + } + const exportURI = `exported-data://${encodeURIComponent(exportNameWithExtension)}`; + const exportFilePath = path.join(this.exportsDirectoryPath, exportNameWithExtension); + const inProgressExport: InProgressExport = (this.storedExports[exportNameWithExtension] = { + exportName: exportNameWithExtension, + exportTitle, + exportPath: exportFilePath, + exportURI: exportURI, + exportStatus: "in-progress", + }); + + void this.startExport({ input, jsonExportFormat, inProgressExport }); + return Promise.resolve(inProgressExport); + } + ); } catch (error) { this.logger.error({ id: LogId.exportCreationError, @@ -224,47 +235,57 @@ export class ExportsManager extends EventEmitter { jsonExportFormat: JSONExportFormat; inProgressExport: InProgressExport; }): Promise { - let pipeSuccessful = false; - await this.withLock(inProgressExport.exportName, "write", false, async (): Promise => { - 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 } - ); - pipeSuccessful = true; - } catch (error) { - this.logger.error({ - id: LogId.exportCreationError, - context: `Error when generating JSON export for ${inProgressExport.exportName}`, - message: error instanceof Error ? error.message : String(error), - }); - - // If the pipeline errors out then we might end up with - // partial and incorrect export so we remove it entirely. - await this.silentlyRemoveExport( - inProgressExport.exportPath, - LogId.exportCreationCleanupError, - `Error when removing incomplete export ${inProgressExport.exportName}` - ); - delete this.storedExports[inProgressExport.exportName]; - } finally { - if (pipeSuccessful) { - this.storedExports[inProgressExport.exportName] = { - ...inProgressExport, - exportCreatedAt: Date.now(), - exportStatus: "ready", - }; - this.emit("export-available", inProgressExport.exportURI); + try { + await this.withLock( + { + exportName: inProgressExport.exportName, + mode: "write", + callbackName: "startExport", + }, + async (): Promise => { + let pipeSuccessful = false; + 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 } + ); + pipeSuccessful = true; + } catch (error) { + // If the pipeline errors out then we might end up with + // partial and incorrect export so we remove it entirely. + await this.silentlyRemoveExport( + inProgressExport.exportPath, + LogId.exportCreationCleanupError, + `Error when removing incomplete export ${inProgressExport.exportName}` + ); + delete this.storedExports[inProgressExport.exportName]; + throw error; + } finally { + if (pipeSuccessful) { + this.storedExports[inProgressExport.exportName] = { + ...inProgressExport, + exportCreatedAt: Date.now(), + exportStatus: "ready", + }; + this.emit("export-available", inProgressExport.exportURI); + } + void input.close(); + } } - void input.close(); - } - }); + ); + } catch (error) { + this.logger.error({ + id: LogId.exportCreationError, + context: `Error when generating JSON export for ${inProgressExport.exportName}`, + message: error instanceof Error ? error.message : String(error), + }); + } } private getEJSONOptionsForFormat(format: JSONExportFormat): EJSONOptions | undefined { @@ -321,15 +342,23 @@ export class ExportsManager extends EventEmitter { await Promise.allSettled( exportsForCleanup.map(async ({ exportPath, exportCreatedAt, exportURI, exportName }) => { if (isExportExpired(exportCreatedAt, this.config.exportTimeoutMs)) { - await this.withLock(exportName, "write", true, async (): Promise => { - delete this.storedExports[exportName]; - await this.silentlyRemoveExport( - exportPath, - LogId.exportCleanupError, - `Considerable error when removing export ${exportName}` - ); - this.emit("export-expired", exportURI); - }); + await this.withLock( + { + exportName, + mode: "write", + finalize: true, + callbackName: "cleanupExpiredExport", + }, + async (): Promise => { + delete this.storedExports[exportName]; + await this.silentlyRemoveExport( + exportPath, + LogId.exportCleanupError, + `Considerable error when removing export ${exportName}` + ); + this.emit("export-expired", exportURI); + } + ); } }) ); @@ -367,62 +396,67 @@ export class ExportsManager extends EventEmitter { } } - private async withLock( - exportName: string, - mode: "read" | "write", - finalize: boolean, - fn: () => T | Promise - ): Promise { + private async withLock>( + lockConfig: { + exportName: string; + mode: "read" | "write"; + finalize?: boolean; + callbackName?: string; + }, + callback: () => CallbackResult + ): Promise> { + const { exportName, mode, finalize = false, callbackName } = lockConfig; + const operationName = callbackName ? `${callbackName} - ${exportName}` : exportName; let lock = this.exportLocks.get(exportName); if (!lock) { lock = new RWLock(); this.exportLocks.set(exportName, lock); } - try { + let lockAcquired: boolean = false; + const acquireLock = async (): Promise => { if (mode === "read") { await lock.readLock(this.readTimeoutMs); } else { await lock.writeLock(this.writeTimeoutMs); } - return await fn(); + lockAcquired = true; + }; + + try { + await Promise.race([ + this.operationAbortedPromise(`Acquire ${mode} lock for ${operationName}`), + acquireLock(), + ]); + return await Promise.race([this.operationAbortedPromise(operationName), callback()]); } finally { - lock.unlock(); + if (lockAcquired) { + lock.unlock(); + } if (finalize) { this.exportLocks.delete(exportName); } } } - private async trackOperation(promise: Promise): Promise { - this.activeOperations.add(promise); - try { - return await promise; - } finally { - this.activeOperations.delete(promise); - } - } + private operationAbortedPromise(operationName?: string): Promise { + return new Promise((_, reject) => { + const rejectIfAborted = (): void => { + if (this.shutdownController.signal.aborted) { + // eslint-disable-next-line @typescript-eslint/no-unsafe-assignment + const abortReason = this.shutdownController.signal.reason; + const abortMessage = + typeof abortReason === "string" + ? abortReason + : `${operationName ?? "Operation"} aborted - ExportsManager shutting down!`; + reject(new OperationAbortedError(abortMessage)); + this.shutdownController.signal.removeEventListener("abort", rejectIfAborted); + } + }; - private async waitForActiveOperationsToSettle(timeoutMs: number): Promise { - const pendingPromises = Array.from(this.activeOperations); - if (pendingPromises.length === 0) { - return; - } - let timedOut = false; - const timeoutPromise = new Promise((resolve) => - setTimeout(() => { - timedOut = true; - resolve(); - }, timeoutMs) - ); - await Promise.race([Promise.allSettled(pendingPromises), timeoutPromise]); - if (timedOut && this.activeOperations.size > 0) { - this.logger.error({ - id: LogId.exportCloseError, - context: `Close timed out waiting for ${this.activeOperations.size} operation(s) to settle`, - message: "Proceeding to force cleanup after timeout", - }); - } + rejectIfAborted(); + this.shutdownController.signal.addEventListener("abort", rejectIfAborted); + }); } static init( diff --git a/src/common/logger.ts b/src/common/logger.ts index 98a7a829..0add105c 100644 --- a/src/common/logger.ts +++ b/src/common/logger.ts @@ -57,6 +57,7 @@ export const LogId = { exportCloseError: mongoLogId(1_007_005), exportedDataListError: mongoLogId(1_007_006), exportedDataAutoCompleteError: mongoLogId(1_007_007), + exportLockError: mongoLogId(1_007_008), } as const; interface LogPayload { diff --git a/tests/unit/common/exportsManager.test.ts b/tests/unit/common/exportsManager.test.ts index 90c44e9e..db096c72 100644 --- a/tests/unit/common/exportsManager.test.ts +++ b/tests/unit/common/exportsManager.test.ts @@ -12,13 +12,12 @@ import { } from "../../../src/common/exportsManager.js"; import { config } from "../../../src/common/config.js"; -import { Session } from "../../../src/common/session.js"; import { ROOT_DIR } from "../../accuracy/sdk/constants.js"; import { timeout } from "../../integration/helpers.js"; import { EJSON, EJSONOptions, ObjectId } from "bson"; import { CompositeLogger } from "../../../src/common/logger.js"; -import { ConnectionManager } from "../../../src/common/connectionManager.js"; +const logger = new CompositeLogger(); const exportsPath = path.join(ROOT_DIR, "tests", "tmp", `exports-${Date.now()}`); const exportsManagerConfig: ExportsManagerConfig = { exportsPath, @@ -26,24 +25,35 @@ const exportsManagerConfig: ExportsManagerConfig = { exportCleanupIntervalMs: config.exportCleanupIntervalMs, } as const; -function getExportNameAndPath( - sessionId: string, - timestamp: number = Date.now(), - objectId: string = new ObjectId().toString() -): { +function getExportNameAndPath({ + uniqueExportsId = new ObjectId().toString(), + uniqueFileId = new ObjectId().toString(), + database = "foo", + collection = "bar", +}: + | { + uniqueExportsId?: string; + uniqueFileId?: string; + database?: string; + collection?: string; + } + | undefined = {}): { sessionExportsPath: string; exportName: string; exportPath: string; exportURI: string; + uniqueExportsId: string; } { - const exportName = `foo.bar.${timestamp}.${objectId}.json`; - const sessionExportsPath = path.join(exportsPath, sessionId); + const exportName = `${database}.${collection}.${uniqueFileId}.json`; + // This is the exports directory for a session. + const sessionExportsPath = path.join(exportsPath, uniqueExportsId); const exportPath = path.join(sessionExportsPath, exportName); return { sessionExportsPath, exportName, exportPath, exportURI: `exported-data://${exportName}`, + uniqueExportsId, }; } @@ -106,24 +116,50 @@ async function fileExists(filePath: string): Promise { } } +function timeoutPromise(timeoutMS: number, context: string): Promise { + return new Promise((_, reject) => { + setTimeout(() => reject(new Error(`${context} - Timed out!`)), timeoutMS); + }); +} + +async function getExportAvailableNotifier( + expectedExportURI: string, + manager: ExportsManager, + timeoutMS = 10_000 +): Promise { + const exportAvailablePromise = new Promise((resolve) => { + manager.on("export-available", (exportURI) => { + if (expectedExportURI === exportURI) { + resolve(exportURI); + } + }); + }); + return await Promise.race([ + timeoutPromise(timeoutMS, `Waiting for export-available - ${expectedExportURI}`), + exportAvailablePromise, + ]); +} + describe("ExportsManager unit test", () => { - let session: Session; let manager: ExportsManager; + let managerClosedPromise: Promise; beforeEach(async () => { await fs.mkdir(exportsManagerConfig.exportsPath, { recursive: true }); - const logger = new CompositeLogger(); - session = new Session({ - apiBaseUrl: "", - logger, - exportsManager: ExportsManager.init(exportsManagerConfig, logger), - connectionManager: new ConnectionManager(), + manager = ExportsManager.init(exportsManagerConfig, logger); + + let notifyManagerClosed: () => void; + managerClosedPromise = new Promise((resolve): void => { + notifyManagerClosed = resolve; + }); + manager.once("closed", (): void => { + notifyManagerClosed(); }); - manager = session.exportsManager; }); afterEach(async () => { await manager?.close(); + await managerClosedPromise; await fs.rm(exportsManagerConfig.exportsPath, { recursive: true, force: true }); }); @@ -135,7 +171,7 @@ describe("ExportsManager unit test", () => { it("should list only the exports that are in ready state", async () => { // This export will finish in at-least 1 second - const { exportName: exportName1 } = getExportNameAndPath(session.sessionId); + const { exportName: exportName1, uniqueExportsId } = getExportNameAndPath(); await manager.createJSONExport({ input: createDummyFindCursorWithDelay([{ name: "Test1" }], 1000).cursor, exportName: exportName1, @@ -144,8 +180,9 @@ describe("ExportsManager unit test", () => { }); // This export will finish way sooner than the first one - const { exportName: exportName2 } = getExportNameAndPath(session.sessionId); - const { cursor, cursorCloseNotification } = createDummyFindCursor([{ name: "Test1" }]); + const { exportName: exportName2, exportURI } = getExportNameAndPath({ uniqueExportsId }); + const secondExportNotifier = getExportAvailableNotifier(exportURI, manager); + const { cursor } = createDummyFindCursor([{ name: "Test1" }]); await manager.createJSONExport({ input: cursor, exportName: exportName2, @@ -153,8 +190,7 @@ describe("ExportsManager unit test", () => { jsonExportFormat: "relaxed", }); - // Small timeout to let the second export finish - await cursorCloseNotification; + await secondExportNotifier; expect(manager.availableExports).toHaveLength(1); expect(manager.availableExports[0]?.exportName).toEqual(exportName2); }); @@ -166,50 +202,64 @@ describe("ExportsManager unit test", () => { await expect(() => manager.readExport("name")).rejects.toThrow("ExportsManager is shutting down."); }); - it("should throw when export name has no extension", async () => { - await expect(() => manager.readExport("name")).rejects.toThrow(); - }); - it("should wait if resource is still being generated", async () => { - const { exportName } = getExportNameAndPath(session.sessionId); + const { exportName } = getExportNameAndPath(); const { cursor } = createDummyFindCursorWithDelay([{ name: "Test1" }], 200); + // create only provides a readable handle but does not guarantee + // that resource is available for read await manager.createJSONExport({ input: cursor, exportName, exportTitle: "Some export", jsonExportFormat: "relaxed", }); - // note that we do not wait for cursor close + expect(await manager.readExport(exportName)).toEqual(JSON.stringify([{ name: "Test1" }])); }); - it("should allow concurrent reads of the same resource", async () => { - const { exportName } = getExportNameAndPath(session.sessionId); + it("should allow concurrent reads of the same in-progress resource", async () => { + const { exportName } = getExportNameAndPath(); const { cursor } = createDummyFindCursorWithDelay([{ name: "Test1" }], 200); + // create only provides a readable handle but does not guarantee + // that resource is available for read await manager.createJSONExport({ input: cursor, exportName, exportTitle: "Some export", jsonExportFormat: "relaxed", }); - // note that we do not wait for cursor close expect( await Promise.all([await manager.readExport(exportName), await manager.readExport(exportName)]) ).toEqual([JSON.stringify([{ name: "Test1" }]), JSON.stringify([{ name: "Test1" }])]); }); it("should return the resource content if the resource is ready to be consumed", async () => { - const { exportName } = getExportNameAndPath(session.sessionId); - const { cursor, cursorCloseNotification } = createDummyFindCursor([]); + const { exportName, exportURI } = getExportNameAndPath(); + const { cursor } = createDummyFindCursor([]); + const exportAvailableNotifier = getExportAvailableNotifier(exportURI, manager); await manager.createJSONExport({ input: cursor, exportName, exportTitle: "Some export", jsonExportFormat: "relaxed", }); - await cursorCloseNotification; + await exportAvailableNotifier; expect(await manager.readExport(exportName)).toEqual("[]"); }); + + it("should handle encoded name", async () => { + const { exportName, exportURI } = getExportNameAndPath({ database: "some database", collection: "coll" }); + const { cursor } = createDummyFindCursor([]); + const exportAvailableNotifier = getExportAvailableNotifier(encodeURI(exportURI), manager); + await manager.createJSONExport({ + input: cursor, + exportName: encodeURIComponent(exportName), + exportTitle: "Some export", + jsonExportFormat: "relaxed", + }); + await exportAvailableNotifier; + expect(await manager.readExport(encodeURIComponent(exportName))).toEqual("[]"); + }); }); describe("#createJSONExport", () => { @@ -230,7 +280,7 @@ describe("ExportsManager unit test", () => { longNumber: Long.fromNumber(123456), }, ])); - ({ exportName, exportPath, exportURI } = getExportNameAndPath(session.sessionId)); + ({ exportName, exportPath, exportURI } = getExportNameAndPath()); }); it("should throw if the manager is shutting down", async () => { @@ -247,21 +297,20 @@ describe("ExportsManager unit test", () => { }); it("should throw if the same name export is requested more than once", async () => { - const { cursor } = createDummyFindCursorWithDelay([{ name: 1 }, { name: 2 }], 100); await manager.createJSONExport({ - input: cursor, + input: createDummyFindCursor([{ name: 1 }, { name: 2 }]).cursor, exportName, - exportTitle: "Some export", + exportTitle: "Export title 1", jsonExportFormat: "relaxed", }); await expect(() => manager.createJSONExport({ - input: cursor, + input: createDummyFindCursor([{ name: 1 }, { name: 2 }]).cursor, exportName, - exportTitle: "Some export", + exportTitle: "Export title 2", jsonExportFormat: "relaxed", }) - ).rejects.toThrow(); + ).rejects.toThrow("Export with same name is either already available or being generated"); }); describe("when cursor is empty", () => { @@ -469,7 +518,7 @@ describe("ExportsManager unit test", () => { }); it("should not clean up in-progress exports", async () => { - const { exportName } = getExportNameAndPath(session.sessionId); + const { exportName, uniqueExportsId } = getExportNameAndPath(); const manager = ExportsManager.init( { ...exportsManagerConfig, @@ -477,7 +526,7 @@ describe("ExportsManager unit test", () => { exportCleanupIntervalMs: 50, }, new CompositeLogger(), - session.sessionId + uniqueExportsId ); const { cursor } = createDummyFindCursorWithDelay([{ name: "Test" }], 2000); await manager.createJSONExport({ @@ -497,7 +546,7 @@ describe("ExportsManager unit test", () => { }); it("should cleanup expired exports", async () => { - const { exportName, exportPath, exportURI } = getExportNameAndPath(session.sessionId); + const { exportName, exportPath, exportURI, uniqueExportsId } = getExportNameAndPath(); const manager = ExportsManager.init( { ...exportsManagerConfig, @@ -505,7 +554,7 @@ describe("ExportsManager unit test", () => { exportCleanupIntervalMs: 50, }, new CompositeLogger(), - session.sessionId + uniqueExportsId ); await manager.createJSONExport({ input: cursor, @@ -530,7 +579,7 @@ describe("ExportsManager unit test", () => { describe("#close", () => { it("should abort ongoing export and remove partial file", async () => { - const { exportName, exportPath } = getExportNameAndPath(session.sessionId); + const { exportName, exportPath } = getExportNameAndPath(); const { cursor } = createDummyFindCursorWithDelay([{ name: "Test" }], 2000); await manager.createJSONExport({ input: cursor, @@ -545,26 +594,6 @@ describe("ExportsManager unit test", () => { await expect(fileExists(exportPath)).resolves.toEqual(false); }); - - it("should timeout shutdown wait when operations never settle", async () => { - await manager.close(); - const logger = new CompositeLogger(); - manager = ExportsManager.init({ ...exportsManagerConfig, activeOpsDrainTimeoutMs: 50 }, logger); - // eslint-disable-next-line @typescript-eslint/no-explicit-any, @typescript-eslint/no-unsafe-member-access, @typescript-eslint/no-unsafe-call - (manager as any).activeOperations.add( - new Promise(() => { - /* never resolves */ - }) - ); - - const start = Date.now(); - await manager.close(); - const elapsed = Date.now() - start; - - // Should not block indefinitely; should be close to the configured timeout but well under 1s - expect(elapsed).toBeGreaterThanOrEqual(45); - expect(elapsed).toBeLessThan(1000); - }); }); }); From 602f1c199ad2382fa178113da5fa44b8d05ebf9c Mon Sep 17 00:00:00 2001 From: Himanshu Singh Date: Mon, 11 Aug 2025 11:51:40 +0200 Subject: [PATCH 47/50] chore: wait for expired --- tests/integration/resources/exportedData.test.ts | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/integration/resources/exportedData.test.ts b/tests/integration/resources/exportedData.test.ts index bc61b78e..67f663d9 100644 --- a/tests/integration/resources/exportedData.test.ts +++ b/tests/integration/resources/exportedData.test.ts @@ -74,7 +74,7 @@ describeWithMongoDB( expect(exportedResourceURI).toBeDefined(); // wait for export expired - await timeout(250); + await timeout(300); const response = await integration.mcpClient().readResource({ uri: exportedResourceURI as string, }); From 1039f4e61f709fc3b3955c8d770a92124e154820 Mon Sep 17 00:00:00 2001 From: Kevin Mas Ruiz Date: Mon, 11 Aug 2025 12:10:03 +0200 Subject: [PATCH 48/50] chore: remove usage of rwlocks as sessions are single user RWLocks are not necessary here because sessions are single user and we don't want the agent to wait until a resource is available, as it can take forever depending on the data set. --- package-lock.json | 10 - package.json | 1 - src/common/exportsManager.ts | 282 ++++++------------ .../resources/exportedData.test.ts | 2 +- tests/unit/common/exportsManager.test.ts | 29 +- 5 files changed, 96 insertions(+), 228 deletions(-) diff --git a/package-lock.json b/package-lock.json index 19772239..d6e7d7dd 100644 --- a/package-lock.json +++ b/package-lock.json @@ -15,7 +15,6 @@ "@mongodb-js/devtools-proxy-support": "^0.5.1", "@mongosh/service-provider-node-driver": "^3.10.2", "@vitest/eslint-plugin": "^1.3.4", - "async-rwlock": "^1.1.1", "bson": "^6.10.4", "express": "^5.1.0", "lru-cache": "^11.1.0", @@ -6099,15 +6098,6 @@ "js-tokens": "^9.0.1" } }, - "node_modules/async-rwlock": { - "version": "1.1.1", - "resolved": "https://registry.npmjs.org/async-rwlock/-/async-rwlock-1.1.1.tgz", - "integrity": "sha512-K4ecpHLAc0Jul4dMb1KLpukblQmHxD5/HNgkTuO3sQ9oLLtVENCJVk7+b0wMj1K89cqnjGkTsAvOb83lCfoKwA==", - "license": "MIT", - "engines": { - "node": ">=8.0.0" - } - }, "node_modules/asynckit": { "version": "0.4.0", "resolved": "https://registry.npmjs.org/asynckit/-/asynckit-0.4.0.tgz", diff --git a/package.json b/package.json index df0b4f28..7bcd5a1e 100644 --- a/package.json +++ b/package.json @@ -96,7 +96,6 @@ "@mongodb-js/devtools-proxy-support": "^0.5.1", "@mongosh/service-provider-node-driver": "^3.10.2", "@vitest/eslint-plugin": "^1.3.4", - "async-rwlock": "^1.1.1", "bson": "^6.10.4", "express": "^5.1.0", "lru-cache": "^11.1.0", diff --git a/src/common/exportsManager.ts b/src/common/exportsManager.ts index 8403f585..38420309 100644 --- a/src/common/exportsManager.ts +++ b/src/common/exportsManager.ts @@ -8,7 +8,6 @@ import { EJSON, EJSONOptions, ObjectId } from "bson"; import { Transform } from "stream"; import { pipeline } from "stream/promises"; import { MongoLogId } from "mongodb-log-writer"; -import { RWLock } from "async-rwlock"; import { UserConfig } from "./config.js"; import { LoggerBase, LogId } from "./logger.js"; @@ -57,17 +56,7 @@ type StoredExport = ReadyExport | InProgressExport; * JIRA: https://jira.mongodb.org/browse/MCP-104 */ type AvailableExport = Pick; -export type ExportsManagerConfig = Pick & { - // The maximum number of milliseconds to wait for in-flight operations to - // settle before shutting down ExportsManager. - activeOpsDrainTimeoutMs?: number; - - // The maximum number of milliseconds to wait before timing out queued reads - readTimeout?: number; - - // The maximum number of milliseconds to wait before timing out queued writes - writeTimeout?: number; -}; +export type ExportsManagerConfig = Pick; type ExportsManagerEvents = { closed: []; @@ -75,16 +64,11 @@ type ExportsManagerEvents = { "export-available": [string]; }; -class OperationAbortedError extends Error {} - export class ExportsManager extends EventEmitter { private storedExports: Record = {}; private exportsCleanupInProgress: boolean = false; private exportsCleanupInterval?: NodeJS.Timeout; private readonly shutdownController: AbortController = new AbortController(); - private readonly readTimeoutMs: number; - private readonly writeTimeoutMs: number; - private readonly exportLocks: Map = new Map(); private constructor( private readonly exportsDirectoryPath: string, @@ -92,8 +76,6 @@ export class ExportsManager extends EventEmitter { private readonly logger: LoggerBase ) { super(); - this.readTimeoutMs = this.config.readTimeout ?? 30_0000; // 30 seconds is the default timeout for an MCP request - this.writeTimeoutMs = this.config.writeTimeout ?? 120_000; // considering that writes can take time } public get availableExports(): AvailableExport[] { @@ -121,6 +103,7 @@ export class ExportsManager extends EventEmitter { ); } } + public async close(): Promise { if (this.shutdownController.signal.aborted) { return; @@ -143,29 +126,18 @@ export class ExportsManager extends EventEmitter { try { this.assertIsNotShuttingDown(); exportName = decodeURIComponent(exportName); - return await this.withLock( - { - exportName, - mode: "read", - callbackName: "readExport", - }, - async (): Promise => { - const exportHandle = this.storedExports[exportName]; - if (!exportHandle) { - throw new Error("Requested export has either expired or does not exist!"); - } + const exportHandle = this.storedExports[exportName]; + if (!exportHandle) { + throw new Error("Requested export has either expired or does not exist."); + } - // This won't happen because of lock synchronization but - // keeping it here to make TS happy. - if (exportHandle.exportStatus === "in-progress") { - throw new Error("Requested export is still being generated!"); - } + if (exportHandle.exportStatus === "in-progress") { + throw new Error("Requested export is still being generated. Try again later."); + } - const { exportPath } = exportHandle; + const { exportPath } = exportHandle; - return fs.readFile(exportPath, { encoding: "utf8", signal: this.shutdownController.signal }); - } - ); + return fs.readFile(exportPath, { encoding: "utf8", signal: this.shutdownController.signal }); } catch (error) { this.logger.error({ id: LogId.exportReadError, @@ -190,32 +162,23 @@ export class ExportsManager extends EventEmitter { try { this.assertIsNotShuttingDown(); const exportNameWithExtension = validateExportName(ensureExtension(exportName, "json")); - return await this.withLock( - { - exportName: exportNameWithExtension, - mode: "write", - callbackName: "createJSONExport", - }, - (): Promise => { - if (this.storedExports[exportNameWithExtension]) { - return Promise.reject( - new Error("Export with same name is either already available or being generated.") - ); - } - const exportURI = `exported-data://${encodeURIComponent(exportNameWithExtension)}`; - const exportFilePath = path.join(this.exportsDirectoryPath, exportNameWithExtension); - const inProgressExport: InProgressExport = (this.storedExports[exportNameWithExtension] = { - exportName: exportNameWithExtension, - exportTitle, - exportPath: exportFilePath, - exportURI: exportURI, - exportStatus: "in-progress", - }); - - void this.startExport({ input, jsonExportFormat, inProgressExport }); - return Promise.resolve(inProgressExport); - } - ); + if (this.storedExports[exportNameWithExtension]) { + return Promise.reject( + new Error("Export with same name is either already available or being generated.") + ); + } + const exportURI = `exported-data://${encodeURIComponent(exportNameWithExtension)}`; + const exportFilePath = path.join(this.exportsDirectoryPath, exportNameWithExtension); + const inProgressExport: InProgressExport = (this.storedExports[exportNameWithExtension] = { + exportName: exportNameWithExtension, + exportTitle, + exportPath: exportFilePath, + exportURI: exportURI, + exportStatus: "in-progress", + }); + + void this.startExport({ input, jsonExportFormat, inProgressExport }); + return Promise.resolve(inProgressExport); } catch (error) { this.logger.error({ id: LogId.exportCreationError, @@ -236,49 +199,41 @@ export class ExportsManager extends EventEmitter { inProgressExport: InProgressExport; }): Promise { try { - await this.withLock( - { - exportName: inProgressExport.exportName, - mode: "write", - callbackName: "startExport", - }, - async (): Promise => { - let pipeSuccessful = false; - 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 } - ); - pipeSuccessful = true; - } catch (error) { - // If the pipeline errors out then we might end up with - // partial and incorrect export so we remove it entirely. - await this.silentlyRemoveExport( - inProgressExport.exportPath, - LogId.exportCreationCleanupError, - `Error when removing incomplete export ${inProgressExport.exportName}` - ); - delete this.storedExports[inProgressExport.exportName]; - throw error; - } finally { - if (pipeSuccessful) { - this.storedExports[inProgressExport.exportName] = { - ...inProgressExport, - exportCreatedAt: Date.now(), - exportStatus: "ready", - }; - this.emit("export-available", inProgressExport.exportURI); - } - void input.close(); - } + let pipeSuccessful = false; + 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 } + ); + pipeSuccessful = true; + } catch (error) { + // If the pipeline errors out then we might end up with + // partial and incorrect export so we remove it entirely. + delete this.storedExports[inProgressExport.exportName]; + // do not block the user, just delete the file in the background + void this.silentlyRemoveExport( + inProgressExport.exportPath, + LogId.exportCreationCleanupError, + `Error when removing incomplete export ${inProgressExport.exportName}` + ); + throw error; + } finally { + if (pipeSuccessful) { + this.storedExports[inProgressExport.exportName] = { + ...inProgressExport, + exportCreatedAt: Date.now(), + exportStatus: "ready", + }; + this.emit("export-available", inProgressExport.exportURI); } - ); + void input.close(); + } } catch (error) { this.logger.error({ id: LogId.exportCreationError, @@ -335,33 +290,31 @@ export class ExportsManager extends EventEmitter { this.exportsCleanupInProgress = true; try { - const exportsForCleanup = Object.values({ ...this.storedExports }).filter( - (storedExport): storedExport is ReadyExport => storedExport.exportStatus === "ready" - ); + // first, unregister all exports that are expired, so they are not considered anymore for reading + const exportsForCleanup: ReadyExport[] = []; + for (const expiredExport of Object.values(this.storedExports)) { + if ( + expiredExport.exportStatus === "ready" && + isExportExpired(expiredExport.exportCreatedAt, this.config.exportTimeoutMs) + ) { + exportsForCleanup.push(expiredExport); + delete this.storedExports[expiredExport.exportName]; + } + } - await Promise.allSettled( - exportsForCleanup.map(async ({ exportPath, exportCreatedAt, exportURI, exportName }) => { - if (isExportExpired(exportCreatedAt, this.config.exportTimeoutMs)) { - await this.withLock( - { - exportName, - mode: "write", - finalize: true, - callbackName: "cleanupExpiredExport", - }, - async (): Promise => { - delete this.storedExports[exportName]; - await this.silentlyRemoveExport( - exportPath, - LogId.exportCleanupError, - `Considerable error when removing export ${exportName}` - ); - this.emit("export-expired", exportURI); - } - ); - } - }) - ); + // and then remove them (slow operation potentially) from disk. + const allDeletionPromises: Promise[] = []; + for (const { exportPath, exportName } of exportsForCleanup) { + allDeletionPromises.push( + this.silentlyRemoveExport( + exportPath, + LogId.exportCleanupError, + `Considerable error when removing export ${exportName}` + ) + ); + } + + await Promise.allSettled(allDeletionPromises); } catch (error) { this.logger.error({ id: LogId.exportCleanupError, @@ -396,69 +349,6 @@ export class ExportsManager extends EventEmitter { } } - private async withLock>( - lockConfig: { - exportName: string; - mode: "read" | "write"; - finalize?: boolean; - callbackName?: string; - }, - callback: () => CallbackResult - ): Promise> { - const { exportName, mode, finalize = false, callbackName } = lockConfig; - const operationName = callbackName ? `${callbackName} - ${exportName}` : exportName; - let lock = this.exportLocks.get(exportName); - if (!lock) { - lock = new RWLock(); - this.exportLocks.set(exportName, lock); - } - - let lockAcquired: boolean = false; - const acquireLock = async (): Promise => { - if (mode === "read") { - await lock.readLock(this.readTimeoutMs); - } else { - await lock.writeLock(this.writeTimeoutMs); - } - lockAcquired = true; - }; - - try { - await Promise.race([ - this.operationAbortedPromise(`Acquire ${mode} lock for ${operationName}`), - acquireLock(), - ]); - return await Promise.race([this.operationAbortedPromise(operationName), callback()]); - } finally { - if (lockAcquired) { - lock.unlock(); - } - if (finalize) { - this.exportLocks.delete(exportName); - } - } - } - - private operationAbortedPromise(operationName?: string): Promise { - return new Promise((_, reject) => { - const rejectIfAborted = (): void => { - if (this.shutdownController.signal.aborted) { - // eslint-disable-next-line @typescript-eslint/no-unsafe-assignment - const abortReason = this.shutdownController.signal.reason; - const abortMessage = - typeof abortReason === "string" - ? abortReason - : `${operationName ?? "Operation"} aborted - ExportsManager shutting down!`; - reject(new OperationAbortedError(abortMessage)); - this.shutdownController.signal.removeEventListener("abort", rejectIfAborted); - } - }; - - rejectIfAborted(); - this.shutdownController.signal.addEventListener("abort", rejectIfAborted); - }); - } - static init( config: ExportsManagerConfig, logger: LoggerBase, diff --git a/tests/integration/resources/exportedData.test.ts b/tests/integration/resources/exportedData.test.ts index 67f663d9..e8cb0d83 100644 --- a/tests/integration/resources/exportedData.test.ts +++ b/tests/integration/resources/exportedData.test.ts @@ -55,7 +55,7 @@ describeWithMongoDB( expect(response.isError).toEqual(true); expect(response.contents[0]?.uri).toEqual(exportURI); expect(response.contents[0]?.text).toEqual( - `Error reading ${exportURI}: Requested export has either expired or does not exist!` + `Error reading ${exportURI}: Requested export has either expired or does not exist.` ); }); }); diff --git a/tests/unit/common/exportsManager.test.ts b/tests/unit/common/exportsManager.test.ts index db096c72..d871933a 100644 --- a/tests/unit/common/exportsManager.test.ts +++ b/tests/unit/common/exportsManager.test.ts @@ -202,7 +202,7 @@ describe("ExportsManager unit test", () => { await expect(() => manager.readExport("name")).rejects.toThrow("ExportsManager is shutting down."); }); - it("should wait if resource is still being generated", async () => { + it("should notify the user if resource is still being generated", async () => { const { exportName } = getExportNameAndPath(); const { cursor } = createDummyFindCursorWithDelay([{ name: "Test1" }], 200); // create only provides a readable handle but does not guarantee @@ -214,23 +214,12 @@ describe("ExportsManager unit test", () => { jsonExportFormat: "relaxed", }); - expect(await manager.readExport(exportName)).toEqual(JSON.stringify([{ name: "Test1" }])); - }); - - it("should allow concurrent reads of the same in-progress resource", async () => { - const { exportName } = getExportNameAndPath(); - const { cursor } = createDummyFindCursorWithDelay([{ name: "Test1" }], 200); - // create only provides a readable handle but does not guarantee - // that resource is available for read - await manager.createJSONExport({ - input: cursor, - exportName, - exportTitle: "Some export", - jsonExportFormat: "relaxed", - }); - expect( - await Promise.all([await manager.readExport(exportName), await manager.readExport(exportName)]) - ).toEqual([JSON.stringify([{ name: "Test1" }]), JSON.stringify([{ name: "Test1" }])]); + try { + await manager.readExport(exportName); + throw new Error("Should have failed."); + } catch (err: unknown) { + expect(String(err)).toEqual("Error: Requested export is still being generated. Try again later."); + } }); it("should return the resource content if the resource is ready to be consumed", async () => { @@ -463,7 +452,7 @@ describe("ExportsManager unit test", () => { // Because the export was never populated in the available exports. await expect(() => manager.readExport(exportName)).rejects.toThrow( - "Requested export has either expired or does not exist!" + "Requested export has either expired or does not exist." ); expect(emitSpy).not.toHaveBeenCalled(); expect(manager.availableExports).toEqual([]); @@ -491,7 +480,7 @@ describe("ExportsManager unit test", () => { // Because the export was never populated in the available exports. await expect(() => manager.readExport(exportName)).rejects.toThrow( - "Requested export has either expired or does not exist!" + "Requested export has either expired or does not exist." ); expect(emitSpy).not.toHaveBeenCalled(); expect(manager.availableExports).toEqual([]); From be90c254aa40536737a233e2ffd7885433b6a70c Mon Sep 17 00:00:00 2001 From: Kevin Mas Ruiz Date: Mon, 11 Aug 2025 12:34:34 +0200 Subject: [PATCH 49/50] chore: add retries to a flaky test that depends on timeouts --- .../resources/exportedData.test.ts | 24 +++++++++++++------ 1 file changed, 17 insertions(+), 7 deletions(-) diff --git a/tests/integration/resources/exportedData.test.ts b/tests/integration/resources/exportedData.test.ts index e8cb0d83..d2ef6250 100644 --- a/tests/integration/resources/exportedData.test.ts +++ b/tests/integration/resources/exportedData.test.ts @@ -74,13 +74,23 @@ describeWithMongoDB( expect(exportedResourceURI).toBeDefined(); // wait for export expired - await timeout(300); - const response = await integration.mcpClient().readResource({ - uri: exportedResourceURI as string, - }); - expect(response.isError).toEqual(true); - expect(response.contents[0]?.uri).toEqual(exportedResourceURI); - expect(response.contents[0]?.text).toMatch(`Error reading ${exportedResourceURI}:`); + for (let tries = 0; tries < 10; tries++) { + await timeout(300); + const response = await integration.mcpClient().readResource({ + uri: exportedResourceURI as string, + }); + + // wait for an error from the MCP Server as it + // means the resource is not available anymore + if (response.isError === false) { + continue; + } + + expect(response.isError).toEqual(true); + expect(response.contents[0]?.uri).toEqual(exportedResourceURI); + expect(response.contents[0]?.text).toMatch(`Error reading ${exportedResourceURI}:`); + break; + } }); }); From 6c541d86f16189e74498424915acb5910c4e1503 Mon Sep 17 00:00:00 2001 From: Kevin Mas Ruiz Date: Mon, 11 Aug 2025 12:44:41 +0200 Subject: [PATCH 50/50] chore: js triple equals magic --- tests/integration/resources/exportedData.test.ts | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/integration/resources/exportedData.test.ts b/tests/integration/resources/exportedData.test.ts index d2ef6250..94710d87 100644 --- a/tests/integration/resources/exportedData.test.ts +++ b/tests/integration/resources/exportedData.test.ts @@ -82,7 +82,7 @@ describeWithMongoDB( // wait for an error from the MCP Server as it // means the resource is not available anymore - if (response.isError === false) { + if (response.isError !== true) { continue; }