feat: makes export tool async

himanshusinghs · himanshusinghs · commit c8709d217698 · 2025-08-07T15:43:11.000+02:00
diff --git 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<typeof jsonExportFormat>;
 
-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<SessionExportsManagerEve
 
     public get availableExports(): AvailableExport[] {
         return Object.values(this.sessionExports)
-            .filter(({ exportCreatedAt: createdAt }) => !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<SessionExportsManagerEve
                 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!");
+            }
+
             const { exportPath, exportCreatedAt } = exportHandle;
 
             if (isExportExpired(exportCreatedAt, this.config.exportTimeoutMs)) {
@@ -124,38 +143,61 @@ export class SessionExportsManager extends EventEmitter<SessionExportsManagerEve
         }
     }
 
-    public async createJSONExport({
+    public createJSONExport({
         input,
         exportName,
         jsonExportFormat,
     }: {
         input: FindCursor;
         exportName: string;
         jsonExportFormat: JSONExportFormat;
-    }): Promise<AvailableExport> {
+    }): 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<void> {
+        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<SessionExportsManagerEve
                         });
                     }
                 });
+                delete this.sessionExports[inProgressExport.exportName];
                 throw pipelineError;
             } finally {
                 void input.close();
                 if (pipeSuccessful) {
-                    this.sessionExports[exportNameWithExtension] = {
-                        exportName: exportNameWithExtension,
+                    this.sessionExports[inProgressExport.exportName] = {
+                        ...inProgressExport,
                         exportCreatedAt: Date.now(),
-                        exportPath: exportFilePath,
-                        exportURI: exportURI,
+                        exportStatus: "ready",
                     };
-                    this.emit("export-available", exportURI);
+                    this.emit("export-available", inProgressExport.exportURI);
                 }
             }
         } catch (error) {
@@ -183,7 +225,6 @@ export class SessionExportsManager extends EventEmitter<SessionExportsManagerEve
                 context: "Error when generating JSON export",
                 message: error instanceof Error ? error.message : String(error),
             });
-            throw error;
         }
     }
 
@@ -228,9 +269,11 @@ export class SessionExportsManager extends EventEmitter<SessionExportsManagerEve
         }
 
         this.exportsCleanupInProgress = true;
-        const exportsForCleanup = { ...this.sessionExports };
+        const exportsForCleanup = Object.values({ ...this.sessionExports }).filter(
+            (sessionExport): sessionExport is ReadyExport => 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
@@ -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
@@ -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
@@ -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
