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

Author: himanshusinghs

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<SessionEvents> {
-    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<SessionEvents> {
         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<SessionEvents> {
                 : 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"));

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<typeof jsonExportFormat>;
@@ -71,7 +72,7 @@ export class SessionExportsManager extends EventEmitter<SessionExportsManagerEve
     private exportsDirectoryPath: string;
 
     constructor(
-        readonly sessionId: string,
+        sessionId: string,
         private readonly config: SessionExportsManagerConfig,
         private readonly logger: LoggerBase
     ) {
@@ -133,7 +134,7 @@ export class SessionExportsManager extends EventEmitter<SessionExportsManagerEve
         } catch (error) {
             this.logger.error({
                 id: LogId.exportReadError,
-                context: "Error when reading export",
+                context: `Error when reading export - ${exportName}`,
                 message: error instanceof Error ? error.message : String(error),
             });
             if ((error as NodeJS.ErrnoException).code === "ENOENT") {
@@ -184,61 +185,54 @@ export class SessionExportsManager extends EventEmitter<SessionExportsManagerEve
         jsonExportFormat: JSONExportFormat;
         inProgressExport: InProgressExport;
     }): Promise<void> {
+        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<SessionExportsManagerEve
             for (const { exportPath, exportCreatedAt, exportURI, exportName } of exportsForCleanup) {
                 if (isExportExpired(exportCreatedAt, this.config.exportTimeoutMs)) {
                     delete this.sessionExports[exportName];
-                    await this.silentlyRemoveExport(exportPath);
+                    await this.silentlyRemoveExport(
+                        exportPath,
+                        LogId.exportCleanupError,
+                        `Considerable error when removing export ${exportName}`
+                    );
                     this.emit("export-expired", exportURI);
                 }
             }
@@ -291,7 +289,7 @@ export class SessionExportsManager extends EventEmitter<SessionExportsManagerEve
         }
     }
 
-    private async silentlyRemoveExport(exportPath: string): Promise<void> {
+    private async silentlyRemoveExport(exportPath: string, logId: MongoLogId, logContext: string): Promise<void> {
         try {
             await fs.unlink(exportPath);
         } catch (error) {
@@ -300,8 +298,8 @@ export class SessionExportsManager extends EventEmitter<SessionExportsManagerEve
             // we need to flag.
             if ((error as NodeJS.ErrnoException).code !== "ENOENT") {
                 this.logger.error({
-                    id: LogId.exportCleanupError,
-                    context: "Considerable error when removing export file",
+                    id: logId,
+                    context: logContext,
                     message: error instanceof Error ? error.message : String(error),
                 });
             }

src/transports/base.ts

@@ -5,6 +5,9 @@ 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 { SessionExportsManager } from "../common/sessionExportsManager.js";
+import { ConnectionManager } from "../common/connectionManager.js";
 
 export abstract class TransportRunnerBase {
     public logger: LoggerBase;
@@ -39,11 +42,19 @@ export abstract class TransportRunnerBase {
             loggers.push(new McpLogger(mcpServer));
         }
 
+        const logger = new CompositeLogger(...loggers);
+        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: new CompositeLogger(...loggers),
+            logger,
+            sessionId,
+            exportsManager,
+            connectionManager,
         });
 
         const telemetry = Telemetry.create(session, userConfig);

tests/integration/helpers.ts

@@ -1,3 +1,4 @@
+import { ObjectId } from "bson";
 import { Client } from "@modelcontextprotocol/sdk/client/index.js";
 import { InMemoryTransport } from "./inMemoryTransport.js";
 import { Server } from "../../src/server.js";
@@ -10,6 +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";
 
 interface ParameterInfo {
     name: string;
@@ -55,15 +57,19 @@ export function setupIntegrationTest(getUserConfig: () => 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

tests/integration/telemetry.test.ts

@@ -5,17 +5,23 @@ 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 () => {
         const actualId: string = await nodeMachineId.machineId(true);
 
         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
         );

tests/unit/common/session.test.ts

@@ -3,17 +3,23 @@ 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);
 
 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);

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;
     });
 

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);