Fix automatic log level handling for sessionless connections (#917)

Author: cliffhall

src/server/index.test.ts

@@ -15,7 +15,8 @@ import {
   ListResourcesRequestSchema,
   ListToolsRequestSchema,
   SetLevelRequestSchema,
-  ErrorCode
+  ErrorCode,
+  LoggingMessageNotification
 } from "../types.js";
 import { Transport } from "../shared/transport.js";
 import { InMemoryTransport } from "../inMemory.js";
@@ -569,7 +570,7 @@ test("should allow elicitation reject and cancel without validation", async () =
     action: "decline",
   });
 
-  // Test cancel - should not validate  
+  // Test cancel - should not validate
   await expect(
     server.elicitInput({
       message: "Please provide your name",
@@ -861,3 +862,154 @@ test("should handle request timeout", async () => {
     code: ErrorCode.RequestTimeout,
   });
 });
+
+/*
+  Test automatic log level handling for transports with and without sessionId
+ */
+test("should respect log level for transport without sessionId", async () => {
+
+    const server = new Server(
+        {
+            name: "test server",
+            version: "1.0",
+        },
+        {
+            capabilities: {
+                prompts: {},
+                resources: {},
+                tools: {},
+                logging: {},
+            },
+            enforceStrictCapabilities: true,
+        },
+    );
+
+    const client = new Client(
+        {
+            name: "test client",
+            version: "1.0",
+        },
+    );
+
+    const [clientTransport, serverTransport] = InMemoryTransport.createLinkedPair();
+
+    await Promise.all([
+        client.connect(clientTransport),
+        server.connect(serverTransport),
+    ]);
+
+    expect(clientTransport.sessionId).toEqual(undefined);
+
+    // Client sets logging level to warning
+    await client.setLoggingLevel("warning");
+
+    // This one will make it through
+    const warningParams: LoggingMessageNotification["params"] = {
+        level: "warning",
+        logger: "test server",
+        data: "Warning message",
+    };
+
+    // This one will not
+    const debugParams: LoggingMessageNotification["params"] = {
+        level: "debug",
+        logger: "test server",
+        data: "Debug message",
+    };
+
+    // Test the one that makes it through
+    clientTransport.onmessage = jest.fn().mockImplementation((message) => {
+        expect(message).toEqual({
+            jsonrpc: "2.0",
+            method: "notifications/message",
+            params: warningParams
+        });
+    });
+
+    // This one will not make it through
+    await server.sendLoggingMessage(debugParams);
+    expect(clientTransport.onmessage).not.toHaveBeenCalled();
+
+    // This one will, triggering the above test in clientTransport.onmessage
+    await server.sendLoggingMessage(warningParams);
+    expect(clientTransport.onmessage).toHaveBeenCalled();
+
+});
+
+test("should respect log level for transport with sessionId", async () => {
+
+    const server = new Server(
+        {
+            name: "test server",
+            version: "1.0",
+        },
+        {
+            capabilities: {
+                prompts: {},
+                resources: {},
+                tools: {},
+                logging: {},
+            },
+            enforceStrictCapabilities: true,
+        },
+    );
+
+    const client = new Client(
+        {
+            name: "test client",
+            version: "1.0",
+        },
+    );
+
+    const [clientTransport, serverTransport] = InMemoryTransport.createLinkedPair();
+
+    // Add a session id to the transports
+    const SESSION_ID = "test-session-id";
+    clientTransport.sessionId = SESSION_ID;
+    serverTransport.sessionId = SESSION_ID;
+
+    expect(clientTransport.sessionId).toBeDefined();
+    expect(serverTransport.sessionId).toBeDefined();
+
+    await Promise.all([
+        client.connect(clientTransport),
+        server.connect(serverTransport),
+    ]);
+
+
+    // Client sets logging level to warning
+    await client.setLoggingLevel("warning");
+
+    // This one will make it through
+    const warningParams: LoggingMessageNotification["params"] = {
+        level: "warning",
+        logger: "test server",
+        data: "Warning message",
+    };
+
+    // This one will not
+    const debugParams: LoggingMessageNotification["params"] = {
+        level: "debug",
+        logger: "test server",
+        data: "Debug message",
+    };
+
+    // Test the one that makes it through
+    clientTransport.onmessage = jest.fn().mockImplementation((message) => {
+        expect(message).toEqual({
+            jsonrpc: "2.0",
+            method: "notifications/message",
+            params: warningParams
+        });
+    });
+
+    // This one will not make it through
+    await server.sendLoggingMessage(debugParams, SESSION_ID);
+    expect(clientTransport.onmessage).not.toHaveBeenCalled();
+
+    // This one will, triggering the above test in clientTransport.onmessage
+    await server.sendLoggingMessage(warningParams, SESSION_ID);
+    expect(clientTransport.onmessage).toHaveBeenCalled();
+
+});
+

src/server/index.ts

@@ -117,7 +117,7 @@ export class Server<
             const transportSessionId: string | undefined = extra.sessionId || extra.requestInfo?.headers['mcp-session-id'] as string || undefined;
             const { level } = request.params;
             const parseResult = LoggingLevelSchema.safeParse(level);
-            if (transportSessionId && parseResult.success) {
+            if (parseResult.success) {
                 this._loggingLevels.set(transportSessionId, parseResult.data);
             }
             return {};
@@ -126,15 +126,15 @@ export class Server<
   }
 
   // Map log levels by session id
-  private _loggingLevels = new Map<string, LoggingLevel>();
+  private _loggingLevels = new Map<string | undefined, LoggingLevel>();
 
   // Map LogLevelSchema to severity index
   private readonly LOG_LEVEL_SEVERITY = new Map(
     LoggingLevelSchema.options.map((level, index) => [level, index])
   );
 
   // Is a message with the given level ignored in the log level set for the given session id?
-  private isMessageIgnored = (level: LoggingLevel, sessionId: string): boolean => {
+  private isMessageIgnored = (level: LoggingLevel, sessionId?: string): boolean => {
     const currentLevel = this._loggingLevels.get(sessionId);
     return (currentLevel)
         ? this.LOG_LEVEL_SEVERITY.get(level)! < this.LOG_LEVEL_SEVERITY.get(currentLevel)!
@@ -398,7 +398,7 @@ export class Server<
      */
     async sendLoggingMessage(params: LoggingMessageNotification["params"], sessionId?: string) {
       if (this._capabilities.logging) {
-          if (!sessionId || !this.isMessageIgnored(params.level, sessionId)) {
+          if (!this.isMessageIgnored(params.level, sessionId)) {
               return this.notification({method: "notifications/message", params})
           }
       }