more tweaks

Author: nirinchev

src/common/logger.ts

@@ -64,7 +64,7 @@ export const LogId = {
     oidcFlow: mongoLogId(1_008_001),
 } as const;
 
-interface LogPayload {
+export interface LogPayload {
     id: MongoLogId;
     context: string;
     message: string;
@@ -152,6 +152,26 @@ export abstract class LoggerBase<T extends EventMap<T> = DefaultEventMap> extend
     public emergency(payload: LogPayload): void {
         this.log("emergency", payload);
     }
+
+    protected mapToMongoDBLogLevel(level: LogLevel): "info" | "warn" | "error" | "debug" | "fatal" {
+        switch (level) {
+            case "info":
+                return "info";
+            case "warning":
+                return "warn";
+            case "error":
+                return "error";
+            case "notice":
+            case "debug":
+                return "debug";
+            case "critical":
+            case "alert":
+            case "emergency":
+                return "fatal";
+            default:
+                return "info";
+        }
+    }
 }
 
 export class ConsoleLogger extends LoggerBase {
@@ -225,26 +245,6 @@ export class DiskLogger extends LoggerBase<{ initialized: [] }> {
 
         this.logWriter[mongoDBLevel]("MONGODB-MCP", id, context, message, payload.attributes);
     }
-
-    private mapToMongoDBLogLevel(level: LogLevel): "info" | "warn" | "error" | "debug" | "fatal" {
-        switch (level) {
-            case "info":
-                return "info";
-            case "warning":
-                return "warn";
-            case "error":
-                return "error";
-            case "notice":
-            case "debug":
-                return "debug";
-            case "critical":
-            case "alert":
-            case "emergency":
-                return "fatal";
-            default:
-                return "info";
-        }
-    }
 }
 
 export class McpLogger extends LoggerBase {
@@ -286,7 +286,11 @@ export class CompositeLogger extends LoggerBase {
     public log(level: LogLevel, payload: LogPayload): void {
         // Override the public method to avoid the base logger redacting the message payload
         for (const logger of this.loggers) {
-            logger.log(level, { ...payload, attributes: { ...this.attributes, ...payload.attributes } });
+            const attributes =
+                Object.keys(this.attributes).length > 0 || payload.attributes
+                    ? { ...this.attributes, ...payload.attributes }
+                    : undefined;
+            logger.log(level, { ...payload, attributes });
         }
     }
 

src/helpers/deviceId.ts

@@ -1,5 +1,6 @@
 import { getDeviceId } from "@mongodb-js/device-id";
-import nodeMachineId from "node-machine-id";
+// Important: don't import the default module as that doesn't work for cjs module resolution.
+import { machineId } from "node-machine-id";
 import type { LoggerBase } from "../common/logger.js";
 import { LogId } from "../common/logger.js";
 
@@ -17,7 +18,7 @@ export class DeviceId {
     private constructor(logger: LoggerBase, timeout: number = DEVICE_ID_TIMEOUT) {
         this.logger = logger;
         this.timeout = timeout;
-        this.getMachineId = (): Promise<string> => nodeMachineId.machineId(true);
+        this.getMachineId = (): Promise<string> => machineId(true);
         this.abortController = new AbortController();
 
         this.deviceIdPromise = DeviceId.UnknownDeviceId;

src/lib.ts

@@ -3,3 +3,5 @@ export { Telemetry } from "./telemetry/telemetry.js";
 export { Session, type SessionOptions } from "./common/session.js";
 export { type UserConfig, defaultUserConfig } from "./common/config.js";
 export { StreamableHttpRunner } from "./transports/streamableHttp.js";
+export { LoggerBase } from "./common/logger.js";
+export type { LogPayload, LoggerType, LogLevel } from "./common/logger.js";

src/transports/base.ts

@@ -16,9 +16,10 @@ export abstract class TransportRunnerBase {
 
     protected constructor(
         protected readonly userConfig: UserConfig,
-        private readonly driverOptions: DriverOptions
+        private readonly driverOptions: DriverOptions,
+        additionalLoggers: LoggerBase[]
     ) {
-        const loggers: LoggerBase[] = [];
+        const loggers: LoggerBase[] = [...additionalLoggers];
         if (this.userConfig.loggers.includes("stderr")) {
             loggers.push(new ConsoleLogger());
         }

src/transports/stdio.ts

@@ -1,3 +1,4 @@
+import type { LoggerBase } from "../common/logger.js";
 import { LogId } from "../common/logger.js";
 import type { Server } from "../server.js";
 import { TransportRunnerBase } from "./base.js";
@@ -54,8 +55,8 @@ export function createStdioTransport(): StdioServerTransport {
 export class StdioRunner extends TransportRunnerBase {
     private server: Server | undefined;
 
-    constructor(userConfig: UserConfig, driverOptions: DriverOptions) {
-        super(userConfig, driverOptions);
+    constructor(userConfig: UserConfig, driverOptions: DriverOptions, additionalLoggers: LoggerBase[] = []) {
+        super(userConfig, driverOptions, additionalLoggers);
     }
 
     async start(): Promise<void> {

src/transports/streamableHttp.ts

@@ -4,6 +4,7 @@ import { StreamableHTTPServerTransport } from "@modelcontextprotocol/sdk/server/
 import { isInitializeRequest } from "@modelcontextprotocol/sdk/types.js";
 import { TransportRunnerBase } from "./base.js";
 import type { DriverOptions, UserConfig } from "../common/config.js";
+import type { LoggerBase } from "../common/logger.js";
 import { LogId } from "../common/logger.js";
 import { randomUUID } from "crypto";
 import { SessionStore } from "../common/sessionStore.js";
@@ -30,8 +31,8 @@ export class StreamableHttpRunner extends TransportRunnerBase {
         throw new Error("Server is not started yet");
     }
 
-    constructor(userConfig: UserConfig, driverOptions: DriverOptions) {
-        super(userConfig, driverOptions);
+    constructor(userConfig: UserConfig, driverOptions: DriverOptions, additionalLoggers: LoggerBase[] = []) {
+        super(userConfig, driverOptions, additionalLoggers);
     }
 
     async start(): Promise<void> {

tests/integration/transports/streamableHttp.test.ts

@@ -1,8 +1,10 @@
 import { StreamableHttpRunner } from "../../../src/transports/streamableHttp.js";
 import { Client } from "@modelcontextprotocol/sdk/client/index.js";
 import { StreamableHTTPClientTransport } from "@modelcontextprotocol/sdk/client/streamableHttp.js";
-import { describe, expect, it, beforeAll, afterAll } from "vitest";
+import { describe, expect, it, beforeAll, afterAll, beforeEach } from "vitest";
 import { config, driverOptions } from "../../../src/common/config.js";
+import type { LoggerType, LogLevel, LogPayload } from "../../../src/common/logger.js";
+import { LoggerBase, LogId } from "../../../src/common/logger.js";
 
 describe("StreamableHttpRunner", () => {
     let runner: StreamableHttpRunner;
@@ -120,4 +122,35 @@ describe("StreamableHttpRunner", () => {
             }
         }
     });
+
+    describe("with custom logger", () => {
+        beforeEach(() => {
+            config.loggers = [];
+        });
+
+        class CustomLogger extends LoggerBase {
+            protected type?: LoggerType = "console";
+            public messages: { level: LogLevel; payload: LogPayload }[] = [];
+            protected logCore(level: LogLevel, payload: LogPayload): void {
+                this.messages.push({ level, payload });
+            }
+        }
+
+        it("can provide custom logger", async () => {
+            const logger = new CustomLogger();
+            const runner = new StreamableHttpRunner(config, driverOptions, [logger]);
+            await runner.start();
+
+            const messages = logger.messages;
+            expect(messages.length).toBeGreaterThan(0);
+
+            const serverStartedMessage = messages.filter(
+                (m) => m.payload.id === LogId.streamableHttpTransportStarted
+            )[0];
+            expect(serverStartedMessage).toBeDefined();
+            expect(serverStartedMessage?.payload.message).toContain("Server started on");
+            expect(serverStartedMessage?.payload.context).toBe("streamableHttpTransport");
+            expect(serverStartedMessage?.level).toBe("info");
+        });
+    });
 });