refactor: wrap setRequestHandler for high-level servers

Author: kashishhora

src/modules/tracingV2.ts

@@ -233,10 +233,9 @@ function addMCPcatToolsToServer(server: HighLevelMCPServerLike): void {
 function addTracingToToolCallback(
   tool: RegisteredTool,
   toolName: string,
-  server: HighLevelMCPServerLike,
+  _server: HighLevelMCPServerLike,
 ): RegisteredTool {
   const originalCallback = tool.callback;
-  const lowLevelServer = server.server as MCPServerLike;
 
   // Check if this callback has already been wrapped
   if (wrappedCallbacks.has(originalCallback)) {
@@ -277,156 +276,168 @@ function addTracingToToolCallback(
       return args;
     };
 
-    try {
-      const data = getServerTrackingData(lowLevelServer);
-      if (!data) {
-        writeToLog(
-          "Warning: MCPCat is unable to find server tracking data. Please ensure you have called track(server, options) before using tool calls.",
-        );
+    // Remove context from args before calling original callback
+    // BUT keep it for get_more_tools since it's a required parameter
+    const cleanedArgs =
+      toolName === "get_more_tools" ? args : removeContextFromArgs(args);
+
+    // Call original callback with cleaned args
+    if (cleanedArgs === undefined) {
+      return await (
+        originalCallback as (
+          extra: CompatibleRequestHandlerExtra,
+        ) => Promise<CallToolResult>
+      )(extra);
+    } else {
+      return await (
+        originalCallback as (
+          args: any,
+          extra: CompatibleRequestHandlerExtra,
+        ) => Promise<CallToolResult>
+      )(cleanedArgs, extra);
+    }
+  };
 
-        // Remove context from args before calling original callback
-        // BUT keep it for get_more_tools since it's a required parameter
-        const cleanedArgs =
-          toolName === "get_more_tools" ? args : removeContextFromArgs(args);
-
-        // Call with original params
-        return await (cleanedArgs === undefined
-          ? (
-              originalCallback as (
-                extra: CompatibleRequestHandlerExtra,
-              ) => Promise<CallToolResult>
-            )(extra)
-          : (
-              originalCallback as (
-                args: any,
-                extra: CompatibleRequestHandlerExtra,
-              ) => Promise<CallToolResult>
-            )(cleanedArgs, extra));
-      }
+  // Mark the original callback as wrapped
+  wrappedCallbacks.set(originalCallback, true);
 
-      const sessionId = getServerSessionId(lowLevelServer, extra);
+  // Mark the wrapped callback as well (in case it gets re-wrapped)
+  wrappedCallbacks.set(wrappedCallback, true);
 
-      // Create a request-like object for compatibility with existing code
-      const request = {
-        params: {
-          name: toolName,
-          arguments: args,
-        },
-      };
+  // Create a new tool object with the wrapped callback
+  const wrappedTool = {
+    ...tool,
+    callback: wrappedCallback as RegisteredTool["callback"],
+  };
 
-      let event: UnredactedEvent = {
-        sessionId: sessionId,
-        resourceName: toolName,
-        parameters: {
-          request: request,
-          extra: extra,
-        },
-        eventType: PublishEventRequestEventTypeEnum.mcpToolsCall,
-        timestamp: new Date(),
-        redactionFn: data.options.redactSensitiveInformation,
-      };
+  // Mark the tool as processed
+  (wrappedTool as any)[MCPCAT_PROCESSED] = true;
 
-      try {
-        // Try to identify the session if identify function is provided
-        await handleIdentify(lowLevelServer, data, request, extra);
+  return wrappedTool;
+}
+
+function setupToolsCallHandlerWrapping(server: HighLevelMCPServerLike): void {
+  const lowLevelServer = server.server as MCPServerLike;
+
+  // Check if tools/call handler already exists
+  const existingHandler = lowLevelServer._requestHandlers.get("tools/call");
+  if (existingHandler) {
+    const wrappedHandler = createToolsCallWrapper(
+      existingHandler,
+      lowLevelServer,
+    );
+    lowLevelServer._requestHandlers.set("tools/call", wrappedHandler);
+  }
+
+  // Intercept future calls to setRequestHandler for tools registered after track()
+  const originalSetRequestHandler =
+    lowLevelServer.setRequestHandler.bind(lowLevelServer);
+
+  lowLevelServer.setRequestHandler = function (
+    requestSchema: any,
+    handler: any,
+  ) {
+    const method = requestSchema?.shape?.method?.value;
+
+    // Only wrap tools/call handler
+    if (method === "tools/call") {
+      const wrappedHandler = createToolsCallWrapper(handler, lowLevelServer);
+      return originalSetRequestHandler(requestSchema, wrappedHandler);
+    }
+
+    // Pass through all other handlers unchanged
+    return originalSetRequestHandler(requestSchema, handler);
+  } as any;
+}
 
-        // Update event sessionId in case handleIdentify reconnected to a different session
+function createToolsCallWrapper(
+  originalHandler: any,
+  server: MCPServerLike,
+): any {
+  return async (request: any, extra: any) => {
+    const startTime = new Date();
+    let shouldPublishEvent = false;
+    let event: UnredactedEvent | null = null;
+
+    try {
+      const data = getServerTrackingData(server);
+
+      if (!data) {
+        writeToLog(
+          "Warning: MCPCat is unable to find server tracking data. Please ensure you have called track(server, options) before using tool calls.",
+        );
+      } else {
+        shouldPublishEvent = true;
+
+        const sessionId = getServerSessionId(server, extra);
+
+        event = {
+          sessionId,
+          resourceName: request.params?.name || "Unknown Tool",
+          parameters: { request, extra },
+          eventType: PublishEventRequestEventTypeEnum.mcpToolsCall,
+          timestamp: startTime,
+          redactionFn: data.options.redactSensitiveInformation,
+        };
+
+        // Identify user session
+        await handleIdentify(server, data, request, extra);
         event.sessionId = data.sessionId;
 
-        // Extract context for userIntent if present
-        if (args && typeof args === "object" && "context" in args) {
-          event.userIntent = args.context;
+        // Extract context for userIntent
+        if (
+          data.options.enableToolCallContext &&
+          request.params?.arguments?.context
+        ) {
+          event.userIntent = request.params.arguments.context;
         }
+      }
+    } catch (error) {
+      // If tracing setup fails, log it but continue with tool execution
+      writeToLog(
+        `Warning: MCPCat tracing failed for tool ${request.params?.name}, falling back to original handler - ${error}`,
+      );
+    }
 
-        // Remove context from args before calling original callback
-        // BUT keep it for get_more_tools since it's a required parameter
-        const cleanedArgs =
-          toolName === "get_more_tools" ? args : removeContextFromArgs(args);
-
-        let result = await (cleanedArgs === undefined
-          ? (
-              originalCallback as (
-                extra: CompatibleRequestHandlerExtra,
-              ) => Promise<CallToolResult>
-            )(extra)
-          : (
-              originalCallback as (
-                args: any,
-                extra: CompatibleRequestHandlerExtra,
-              ) => Promise<CallToolResult>
-            )(cleanedArgs, extra));
-
-        // Check if the result indicates an error
+    // Execute the tool (this should always happen, even if tracing setup failed)
+    try {
+      const result = await originalHandler(request, extra);
+
+      if (event && shouldPublishEvent) {
+        // Check for execution errors (SDK converts them to CallToolResult)
         if (isToolResultError(result)) {
           event.isError = true;
           event.error = captureException(result);
         }
 
         event.response = result;
-        event.duration =
-          (event.timestamp &&
-            new Date().getTime() - event.timestamp.getTime()) ||
-          undefined;
-        publishEvent(lowLevelServer, event);
-        return result;
-      } catch (error) {
+        event.duration = new Date().getTime() - startTime.getTime();
+        publishEvent(server, event);
+      }
+
+      return result;
+    } catch (error) {
+      // Validation errors, unknown tool, disabled tool
+      if (event && shouldPublishEvent) {
         event.isError = true;
         event.error = captureException(error);
-        event.duration =
-          (event.timestamp &&
-            new Date().getTime() - event.timestamp.getTime()) ||
-          undefined;
-        publishEvent(lowLevelServer, event);
-        throw error;
+        event.duration = new Date().getTime() - startTime.getTime();
+        publishEvent(server, event);
       }
-    } catch (error) {
-      // If any error occurs in our tracing code, log it and call the original callback
-      writeToLog(
-        `Warning: MCPCat tracing failed for tool ${toolName}, falling back to original callback - ${error}`,
-      );
 
-      // Remove context from args before calling original callback
-      // BUT keep it for get_more_tools since it's a required parameter
-      const cleanedArgs =
-        toolName === "get_more_tools" ? args : removeContextFromArgs(args);
-
-      return await (cleanedArgs === undefined
-        ? (
-            originalCallback as (
-              extra: CompatibleRequestHandlerExtra,
-            ) => Promise<CallToolResult>
-          )(extra)
-        : (
-            originalCallback as (
-              args: any,
-              extra: CompatibleRequestHandlerExtra,
-            ) => Promise<CallToolResult>
-          )(cleanedArgs, extra));
+      // Re-throw so Protocol converts to JSONRPC error response
+      throw error;
     }
   };
-
-  // Mark the original callback as wrapped
-  wrappedCallbacks.set(originalCallback, true);
-
-  // Mark the wrapped callback as well (in case it gets re-wrapped)
-  wrappedCallbacks.set(wrappedCallback, true);
-
-  // Create a new tool object with the wrapped callback
-  const wrappedTool = {
-    ...tool,
-    callback: wrappedCallback as RegisteredTool["callback"],
-  };
-
-  // Mark the tool as processed
-  (wrappedTool as any)[MCPCAT_PROCESSED] = true;
-
-  return wrappedTool;
 }
 
 export function setupTracking(server: HighLevelMCPServerLike): void {
   try {
     const mcpcatData = getServerTrackingData(server.server);
 
+    // Setup handler wrapping before any tools are registered
+    setupToolsCallHandlerWrapping(server);
+
     setupInitializeTracing(server);
     // Modify existing tools to include context parameters in their inputSchemas
     if (mcpcatData?.options.enableToolCallContext) {

src/tests/context-parameters.test.ts

@@ -315,7 +315,8 @@ describe("Context Parameters", () => {
       const listEvent = events.find(
         (e) =>
           e.eventType === PublishEventRequestEventTypeEnum.mcpToolsCall &&
-          e.resourceName === "list_todos",
+          e.resourceName === "list_todos" &&
+          !e.isError, // Find the successful event, not the validation errors
       );
 
       expect(listEvent).toBeDefined();
@@ -396,5 +397,69 @@ describe("Context Parameters", () => {
         );
       });
     });
+
+    it("should remove context parameter before calling tool callback", async () => {
+      // Variable to capture what arguments the tool callback actually receives
+      let capturedCallbackArguments: any = null;
+
+      // Register a test tool that captures its arguments
+      const { z } = await import("zod");
+      server.tool(
+        "test_context_removal",
+        "Test tool that captures callback arguments",
+        {
+          testParam: z.string().describe("A test parameter"),
+        },
+        async (args: any) => {
+          // Capture exactly what arguments this callback receives
+          capturedCallbackArguments = { ...args };
+          return {
+            content: [
+              {
+                type: "text",
+                text: "Arguments captured",
+              },
+            ],
+          };
+        },
+      );
+
+      // Enable tracking with context parameters
+      await track(server, {
+        projectId: "test-project",
+        enableTracing: true,
+      });
+
+      // Call the test tool WITH context parameter
+      const result = await client.request(
+        {
+          method: "tools/call",
+          params: {
+            name: "test_context_removal",
+            arguments: {
+              testParam: "test-value",
+              context: "This context should be removed before callback",
+            },
+          },
+        },
+        CallToolResultSchema,
+      );
+
+      // Wait for processing
+      await new Promise((resolve) => setTimeout(resolve, 100));
+
+      // The tool call should succeed (successful calls have undefined isError)
+      expect(result).toBeDefined();
+      expect(result.isError).not.toBe(true);
+
+      // Verify that the callback received the testParam
+      expect(capturedCallbackArguments).not.toBeNull();
+      expect(capturedCallbackArguments).toHaveProperty("testParam");
+      expect(capturedCallbackArguments.testParam).toBe("test-value");
+
+      // This is the key assertion: context should NOT be in the arguments
+      // that the tool callback received (it should have been removed by the wrapper)
+      expect(capturedCallbackArguments).not.toHaveProperty("context");
+    });
   });
 });