@posthog/mcp 0.10.3 → 0.10.5

package/src/extensions/analytics-parameters.ts

@@ -1,4 +1,21 @@
+import type { AnalyticsParameterOwnership } from '../types'
+import { getObjectShape, isZodRawShapeCompat } from './mcp-sdk-compat'
+
+const JSON_SCHEMA_KEYWORDS = [
+  '$defs',
+  '$ref',
+  'additionalProperties',
+  'allOf',
+  'anyOf',
+  'definitions',
+  'oneOf',
+  'properties',
+  'required',
+  'type',
+]
+
 export interface AnalyticsInjectableJsonSchema {
+  $ref?: unknown
   additionalProperties?: boolean
   allOf?: unknown
   anyOf?: unknown
@@ -19,5 +36,54 @@ export function canInjectAnalyticsParameter(
   schema: AnalyticsInjectableJsonSchema | undefined,
   parameterName: string
 ): boolean {
-  return !hasAnalyticsParameter(schema, parameterName) && !schema?.oneOf && !schema?.allOf && !schema?.anyOf
+  return (
+    !hasAnalyticsParameter(schema, parameterName) && !schema?.$ref && !schema?.oneOf && !schema?.allOf && !schema?.anyOf
+  )
+}
+
+export function analyticsOwnsParameter(inputSchema: unknown, parameterName: string): boolean {
+  if (!inputSchema || typeof inputSchema !== 'object') {
+    return canInjectAnalyticsParameter(undefined, parameterName)
+  }
+
+  if (isZodRawShapeCompat(inputSchema)) {
+    return canInjectAnalyticsParameter({ properties: inputSchema }, parameterName)
+  }
+
+  const shape = getObjectShape(inputSchema)
+  if (shape) {
+    return canInjectAnalyticsParameter({ properties: shape }, parameterName)
+  }
+
+  const schema = inputSchema as Record<string, unknown>
+  if (JSON_SCHEMA_KEYWORDS.some((keyword) => Object.prototype.hasOwnProperty.call(schema, keyword))) {
+    return canInjectAnalyticsParameter(schema, parameterName)
+  }
+
+  if ('_def' in schema || '_zod' in schema || '~standard' in schema) {
+    // The MCP SDK advertises non-object schemas as an empty object schema.
+    return canInjectAnalyticsParameter(undefined, parameterName)
+  }
+
+  return canInjectAnalyticsParameter(schema, parameterName)
+}
+
+export function getAnalyticsParameterOwnership(inputSchema: unknown): AnalyticsParameterOwnership {
+  return {
+    context: analyticsOwnsParameter(inputSchema, 'context'),
+    conversationId: analyticsOwnsParameter(inputSchema, 'conversation_id'),
+  }
+}
+
+export function stripOwnedAnalyticsArguments(args: unknown, ownership: AnalyticsParameterOwnership): unknown {
+  let cleanedArgs = args
+  if (ownership.context && cleanedArgs && typeof cleanedArgs === 'object' && 'context' in cleanedArgs) {
+    const { context: _context, ...rest } = cleanedArgs
+    cleanedArgs = rest
+  }
+  if (ownership.conversationId && cleanedArgs && typeof cleanedArgs === 'object' && 'conversation_id' in cleanedArgs) {
+    const { conversation_id: _conversationId, ...rest } = cleanedArgs
+    cleanedArgs = rest
+  }
+  return cleanedArgs
 }

package/src/extensions/capture.ts

@@ -6,7 +6,7 @@
 import type { MCPServerLike, McpEvent } from '../types'
 import { MCPAnalyticsEventType } from './event-types'
 import { getServerTrackingData } from './internal'
-import { log } from './logger'
+import type { LoggerFn } from './logger'
 import { getSessionInfo } from './session'
 
 /**
@@ -19,10 +19,10 @@ import { getSessionInfo } from './session'
  * Auto-capture callers (tool calls, listings, identify) intentionally ignore the
  * return value, keeping the tool path isolated from analytics latency/errors.
  */
-export function captureEvent(server: MCPServerLike, eventInput: McpEvent): Promise<void> | undefined {
+export function captureEvent(server: MCPServerLike, eventInput: McpEvent, logger: LoggerFn): Promise<void> | undefined {
   const data = getServerTrackingData(server)
   if (!data) {
-    log('Warning: Server tracking data not found. Event will not be published.')
+    logger('Warning: Server tracking data not found. Event will not be published.')
     return
   }
 

package/src/extensions/compatibility.ts

@@ -4,7 +4,7 @@
 // Licensed under the MIT License: https://github.com/agentcathq/agentcat-typescript-sdk/blob/main/LICENSE
 
 import type { HighLevelMCPServerLike, MCPServerLike } from '../types'
-import { log } from './logger'
+import type { LoggerFn } from './logger'
 
 type ServerRecord = Record<string, unknown>
 
@@ -22,8 +22,8 @@ type ServerRecord = Record<string, unknown>
  */
 
 // Function to log compatibility information
-export function logCompatibilityWarning(): void {
-  log(
+export function logCompatibilityWarning(logger: LoggerFn): void {
+  logger(
     'PostHog MCP analytics SDK Compatibility: This version only supports Model Context Protocol TypeScript SDK v1.11 and above. Please upgrade if using an older version.'
   )
 }
@@ -41,9 +41,9 @@ export function isLowLevelServer(server: unknown): server is ServerRecord {
 }
 
 // Type guard function that validates server compatibility and returns typed server
-export function isCompatibleServerType(server: unknown): MCPServerLike | HighLevelMCPServerLike {
+export function isCompatibleServerType(server: unknown, logger: LoggerFn): MCPServerLike | HighLevelMCPServerLike {
   if (!server || typeof server !== 'object') {
-    logCompatibilityWarning()
+    logCompatibilityWarning(logger)
     throw new Error(
       "PostHog MCP analytics SDK compatibility error: Server must be an object. Ensure you're using MCP SDK v1.11 or higher."
     )
@@ -52,13 +52,13 @@ export function isCompatibleServerType(server: unknown): MCPServerLike | HighLev
   if (isHighLevelServer(server)) {
     // Validate high-level server requirements
     if (!server._registeredTools || typeof server._registeredTools !== 'object') {
-      logCompatibilityWarning()
+      logCompatibilityWarning(logger)
       throw new Error(
         'PostHog MCP analytics SDK compatibility error: High-level server must have _registeredTools object. This requires MCP SDK v1.11 or higher.'
       )
     }
     if (typeof server.tool !== 'function') {
-      logCompatibilityWarning()
+      logCompatibilityWarning(logger)
       throw new Error(
         'PostHog MCP analytics SDK compatibility error: High-level server must have tool() method. This requires MCP SDK v1.11 or higher.'
       )
@@ -66,19 +66,19 @@ export function isCompatibleServerType(server: unknown): MCPServerLike | HighLev
 
     // Validate the underlying low-level server
     const targetServer = server.server
-    validateLowLevelServer(targetServer)
+    validateLowLevelServer(targetServer, logger)
 
     return server as unknown as HighLevelMCPServerLike
   }
   // Direct low-level server validation
-  validateLowLevelServer(server)
+  validateLowLevelServer(server, logger)
   return server as MCPServerLike
 }
 
 // Helper function to validate low-level server requirements
-function validateLowLevelServer(server: unknown): void {
+function validateLowLevelServer(server: unknown, logger: LoggerFn): void {
   if (!server || typeof server !== 'object') {
-    logCompatibilityWarning()
+    logCompatibilityWarning(logger)
     throw new Error(
       "PostHog MCP analytics SDK compatibility error: Server must be an object. Ensure you're using MCP SDK v1.11 or higher."
     )
@@ -87,14 +87,14 @@ function validateLowLevelServer(server: unknown): void {
   const serverRecord = server as ServerRecord
 
   if (typeof serverRecord.setRequestHandler !== 'function') {
-    logCompatibilityWarning()
+    logCompatibilityWarning(logger)
     throw new Error(
       'PostHog MCP analytics SDK compatibility error: Server must have a setRequestHandler method. This requires MCP SDK v1.11 or higher.'
     )
   }
 
   if (!(serverRecord._requestHandlers && serverRecord._requestHandlers instanceof Map)) {
-    logCompatibilityWarning()
+    logCompatibilityWarning(logger)
     throw new Error(
       'PostHog MCP analytics SDK compatibility error: Server._requestHandlers is not accessible. This requires MCP SDK v1.11 or higher.'
     )
@@ -102,14 +102,14 @@ function validateLowLevelServer(server: unknown): void {
 
   // Validate that _requestHandlers contains functions with compatible signatures
   if (typeof serverRecord._requestHandlers.get !== 'function') {
-    logCompatibilityWarning()
+    logCompatibilityWarning(logger)
     throw new Error(
       'PostHog MCP analytics SDK compatibility error: Server._requestHandlers must be a Map with a get method. This requires MCP SDK v1.11 or higher.'
     )
   }
 
   if (typeof serverRecord.getClientVersion !== 'function') {
-    logCompatibilityWarning()
+    logCompatibilityWarning(logger)
     throw new Error(
       'PostHog MCP analytics SDK compatibility error: Server.getClientVersion must be a function. This requires MCP SDK v1.11 or higher.'
     )
@@ -120,7 +120,7 @@ function validateLowLevelServer(server: unknown): void {
     typeof serverRecord._serverInfo !== 'object' ||
     !('name' in serverRecord._serverInfo)
   ) {
-    logCompatibilityWarning()
+    logCompatibilityWarning(logger)
     throw new Error(
       'PostHog MCP analytics SDK compatibility error: Server._serverInfo is not accessible or missing name. This requires MCP SDK v1.11 or higher.'
     )

package/src/extensions/context-parameters.ts

@@ -10,7 +10,7 @@ import {
   type AnalyticsInjectableJsonSchema,
 } from './analytics-parameters'
 import { DEFAULT_CONTEXT_PARAMETER_DESCRIPTION } from './constants'
-import { log } from './logger'
+import { log, type LoggerFn } from './logger'
 
 export interface ContextInjectableTool {
   inputSchema?: AnalyticsInjectableJsonSchema
@@ -33,12 +33,13 @@ export function getContextDescription(context: MCPAnalyticsOptions['context']):
  *
  * Skips injection (with warning) for:
  * - Tools that already have a 'context' parameter
- * - Complex schemas (oneOf/allOf/anyOf) that can't safely have properties added
+ * - Complex schemas (oneOf/allOf/anyOf/$ref) that can't safely have properties added
  * - Schemas with additionalProperties: false
  */
 export function addContextParameterToTool<TTool extends ContextInjectableTool>(
   tool: TTool,
-  contextDescriptionOverride?: string
+  contextDescriptionOverride?: string,
+  logger: LoggerFn = log
 ): TTool {
   // Create a shallow copy of the tool to avoid modifying the original
   const modifiedTool = { ...tool }
@@ -47,9 +48,9 @@ export function addContextParameterToTool<TTool extends ContextInjectableTool>(
 
   if (!canInjectAnalyticsParameter(schema, 'context')) {
     if (hasAnalyticsParameter(schema, 'context')) {
-      log(`WARN: Tool "${toolName}" already has 'context' parameter. Skipping context injection.`)
+      logger(`WARN: Tool "${toolName}" already has 'context' parameter. Skipping context injection.`)
     } else {
-      log(`WARN: Tool "${toolName}" has complex schema (oneOf/allOf/anyOf). Skipping context injection.`)
+      logger(`WARN: Tool "${toolName}" has complex schema (oneOf/allOf/anyOf/$ref). Skipping context injection.`)
     }
     return modifiedTool
   }
@@ -104,13 +105,14 @@ export function addContextParameterToTool<TTool extends ContextInjectableTool>(
 
 export function addContextParameterToTools<TTool extends ContextInjectableTool>(
   tools: TTool[],
-  contextDescriptionOverride?: string
+  contextDescriptionOverride?: string,
+  logger: LoggerFn = log
 ): TTool[] {
   return tools.map((tool) => {
     // Skip get_more_tools - it has its own special context parameter
     if (tool.name === 'get_more_tools') {
       return tool
     }
-    return addContextParameterToTool(tool, contextDescriptionOverride)
+    return addContextParameterToTool(tool, contextDescriptionOverride, logger)
   })
 }

package/src/extensions/conversation-id.ts

@@ -10,7 +10,7 @@ import {
   type AnalyticsInjectableJsonSchema,
 } from './analytics-parameters'
 import { DEFAULT_CONVERSATION_ID_DESCRIPTION } from './constants'
-import { log } from './logger'
+import { log, type LoggerFn } from './logger'
 import { GET_MORE_TOOLS_NAME } from './tools'
 
 export const CONVERSATION_ID_PARAM_NAME = 'conversation_id'
@@ -21,18 +21,23 @@ export interface ConversationIdInjectableTool {
   [key: string]: unknown
 }
 
-export function addConversationIdToTool<TTool extends ConversationIdInjectableTool>(tool: TTool): TTool {
+export function addConversationIdToTool<TTool extends ConversationIdInjectableTool>(
+  tool: TTool,
+  logger: LoggerFn = log
+): TTool {
   const modifiedTool = { ...tool }
   const toolName = tool.name || 'unknown'
   const schema = modifiedTool.inputSchema as AnalyticsInjectableJsonSchema | undefined
 
   if (!canInjectAnalyticsParameter(schema, CONVERSATION_ID_PARAM_NAME)) {
     if (hasAnalyticsParameter(schema, CONVERSATION_ID_PARAM_NAME)) {
-      log(
+      logger(
         `WARN: Tool "${toolName}" already has '${CONVERSATION_ID_PARAM_NAME}' parameter. Skipping conversation_id injection.`
       )
     } else {
-      log(`WARN: Tool "${toolName}" has complex schema (oneOf/allOf/anyOf). Skipping conversation_id injection.`)
+      logger(
+        `WARN: Tool "${toolName}" has complex schema (oneOf/allOf/anyOf/$ref). Skipping conversation_id injection.`
+      )
     }
     return modifiedTool
   }
@@ -67,13 +72,14 @@ export function addConversationIdToTool<TTool extends ConversationIdInjectableTo
 
 export function addConversationIdToTools<TTool extends ConversationIdInjectableTool>(
   tools: TTool[],
-  missingCapabilityToolName: string = GET_MORE_TOOLS_NAME
+  missingCapabilityToolName: string = GET_MORE_TOOLS_NAME,
+  logger: LoggerFn = log
 ): TTool[] {
   return tools.map((tool) => {
     if (tool.name === missingCapabilityToolName) {
       return tool
     }
-    return addConversationIdToTool(tool)
+    return addConversationIdToTool(tool, logger)
   })
 }
 
@@ -131,25 +137,6 @@ export function extractConversationId(args: unknown): string | undefined {
   return trimmed.length > 0 ? trimmed : undefined
 }
 
-export function cloneRequestWithoutConversationId<
-  TRequest extends { params?: { arguments?: unknown; [k: string]: unknown } },
->(request: TRequest): TRequest {
-  if (!request.params || typeof request.params !== 'object') {
-    return request
-  }
-  const args = request.params.arguments
-  if (!(args && typeof args === 'object')) {
-    return request
-  }
-  return {
-    ...request,
-    params: {
-      ...request.params,
-      arguments: stripConversationId(args) as typeof request.params.arguments,
-    },
-  }
-}
-
 export function stripConversationId(args: unknown): unknown {
   if (!args || typeof args !== 'object' || !(CONVERSATION_ID_PARAM_NAME in (args as Record<string, unknown>))) {
     return args

package/src/extensions/instrument-highlevel.ts

@@ -11,19 +11,16 @@ import type {
   RegisteredTool,
   ToolCallback,
 } from '../types'
-import { canInjectAnalyticsParameter } from './analytics-parameters'
-import { stripConversationId } from './conversation-id'
+import {
+  analyticsOwnsParameter,
+  getAnalyticsParameterOwnership,
+  stripOwnedAnalyticsArguments,
+} from './analytics-parameters'
 import { isContextEnabled } from './context-parameters'
 import { MCPAnalyticsEventType } from './event-types'
 import { getServerTrackingData } from './internal'
-import { log } from './logger'
-import {
-  createWrappedTool,
-  getObjectShape,
-  getToolFunction,
-  hasToolFunction,
-  isZodRawShapeCompat,
-} from './mcp-sdk-compat'
+import type { LoggerFn } from './logger'
+import { createWrappedTool, getToolFunction, hasToolFunction } from './mcp-sdk-compat'
 import { handleReportMissing, resolveMissingCapabilityToolName } from './tools'
 import {
   handleInitializeRequest,
@@ -31,6 +28,7 @@ import {
   patchRequestHandlers,
   captureToolCall,
   readToolMetaCategory,
+  type HandlerPatch,
 } from './instrumentation'
 import { getContextArgument } from './tracing-helpers'
 
@@ -52,55 +50,24 @@ function isCallbackUpdate(value: unknown): value is { callback: ToolCallback } {
   return !!value && typeof value === 'object' && 'callback' in value && typeof value.callback === 'function'
 }
 
-function analyticsOwnsParameter(inputSchema: unknown, parameterName: string): boolean {
-  if (!inputSchema || typeof inputSchema !== 'object') {
-    return canInjectAnalyticsParameter(undefined, parameterName)
-  }
-
-  if (isZodRawShapeCompat(inputSchema)) {
-    return canInjectAnalyticsParameter({ properties: inputSchema }, parameterName)
-  }
-
-  const shape = getObjectShape(inputSchema)
-  if (shape) {
-    return canInjectAnalyticsParameter({ properties: shape }, parameterName)
-  }
-
-  const schema = inputSchema as Record<string, unknown>
-  if ('_def' in schema || '_zod' in schema || '~standard' in schema) {
-    // The MCP SDK advertises non-object schemas as an empty object schema.
-    return canInjectAnalyticsParameter(undefined, parameterName)
-  }
-
-  return canInjectAnalyticsParameter(schema, parameterName)
-}
-
-function stripOwnedAnalyticsArguments(
-  args: unknown,
-  ownership: { context: boolean; conversationId: boolean }
-): unknown {
-  let cleanedArgs = args
-  if (ownership.context && cleanedArgs && typeof cleanedArgs === 'object' && 'context' in cleanedArgs) {
-    const { context: _context, ...rest } = cleanedArgs
-    cleanedArgs = rest
-  }
-  return ownership.conversationId ? stripConversationId(cleanedArgs) : cleanedArgs
-}
-
 function addTracingToToolRegistry(
   tools: Record<string, RegisteredTool>,
-  server: HighLevelMCPServerLike
+  server: HighLevelMCPServerLike,
+  logger: LoggerFn
 ): Record<string, RegisteredTool> {
   return Object.fromEntries(
-    Object.entries(tools).map(([name, tool]) => [name, addTracingToToolCallbackInternal(tool, name, server)])
+    Object.entries(tools).map(([name, tool]) => [
+      name,
+      addTracingToToolCallbackInternal(tool, name, server, () => tool.inputSchema, logger),
+    ])
   )
 }
 
-function setupListenerToRegisteredTools(server: HighLevelMCPServerLike): void {
+function setupListenerToRegisteredTools(server: HighLevelMCPServerLike, logger: LoggerFn): void {
   try {
     const data = getServerTrackingData(server.server as MCPServerLike)
     if (!data) {
-      log('Warning: Cannot setup listener - no tracking data found')
+      logger('Warning: Cannot setup listener - no tracking data found')
       return
     }
 
@@ -112,19 +79,19 @@ function setupListenerToRegisteredTools(server: HighLevelMCPServerLike): void {
               data.toolDescriptions.set(property, value.description)
             }
             if ((value as ProcessedRegisteredTool)[MCP_ANALYTICS_PROCESSED]) {
-              log(`Tool ${String(property)} already processed, skipping proxy wrapping`)
+              logger(`Tool ${String(property)} already processed, skipping proxy wrapping`)
               return Reflect.set(target, property, value)
             }
 
             if (wrappedCallbacks.has(getToolFunction(value))) {
-              log(`Tool ${String(property)} callback already wrapped, skipping proxy wrapping`)
+              logger(`Tool ${String(property)} callback already wrapped, skipping proxy wrapping`)
               return Reflect.set(target, property, value)
             }
 
             // The MCP SDK's registry copy is currently stale after update(), but keep this lookup lazy so
             // ownership follows the current schema once the registry reflects live tool state.
             const getCurrentInputSchema = () => value.inputSchema
-            const nextValue = addTracingToToolCallbackInternal(value, property, server, getCurrentInputSchema)
+            const nextValue = addTracingToToolCallbackInternal(value, property, server, getCurrentInputSchema, logger)
 
             if (typeof nextValue.update === 'function') {
               const originalUpdate = nextValue.update
@@ -136,7 +103,8 @@ function setupListenerToRegisteredTools(server: HighLevelMCPServerLike): void {
                       { callback: updateObj.callback },
                       property,
                       server,
-                      getCurrentInputSchema
+                      getCurrentInputSchema,
+                      logger
                     )
                     updateObj.callback = getToolFunction(wrappedTool)
                   }
@@ -149,7 +117,7 @@ function setupListenerToRegisteredTools(server: HighLevelMCPServerLike): void {
 
           return Reflect.set(target, property, value)
         } catch (error) {
-          log(`Warning: Error in proxy set handler for tool ${String(property)} - ${error}`)
+          logger(`Warning: Error in proxy set handler for tool ${String(property)} - ${error}`)
           return Reflect.set(target, property, value)
         }
       },
@@ -170,9 +138,9 @@ function setupListenerToRegisteredTools(server: HighLevelMCPServerLike): void {
     const originalTools = server._registeredTools || {}
     server._registeredTools = new Proxy(originalTools, handler)
 
-    log('Successfully set up listener for new tool registrations')
+    logger('Successfully set up listener for new tool registrations')
   } catch (error) {
-    log(`Warning: Failed to setup listener for registered tools - ${error}`)
+    logger(`Warning: Failed to setup listener for registered tools - ${error}`)
   }
 }
 
@@ -189,19 +157,20 @@ function addTracingToToolCallbackInternal(
   tool: RegisteredTool,
   toolName: string,
   server: HighLevelMCPServerLike,
-  getCurrentInputSchema: () => unknown = () => tool.inputSchema
+  getCurrentInputSchema: () => unknown,
+  logger: LoggerFn
 ): RegisteredTool {
   const originalCallback = getToolFunction(tool)
   const options = getServerTrackingData(server.server as MCPServerLike)?.options
   const missingToolName = resolveMissingCapabilityToolName(options)
 
   if (wrappedCallbacks.has(originalCallback)) {
-    log(`Tool ${toolName} callback already wrapped, skipping re-wrap`)
+    logger(`Tool ${toolName} callback already wrapped, skipping re-wrap`)
     return tool
   }
 
   if ((tool as ProcessedRegisteredTool)[MCP_ANALYTICS_PROCESSED]) {
-    log(`Tool ${toolName} already processed, skipping re-wrap`)
+    logger(`Tool ${toolName} already processed, skipping re-wrap`)
     return tool
   }
 
@@ -258,14 +227,16 @@ function addTracingToToolCallbackInternal(
 }
 
 async function handleToolCallRequest(
+  highLevelServer: HighLevelMCPServerLike,
   server: MCPServerLike,
   originalCallToolHandler: MCPRequestHandler,
   request: MCPRequest,
-  extra: MCPRequestExtra
+  extra: MCPRequestExtra,
+  logger: LoggerFn
 ): Promise<unknown> {
   const data = getServerTrackingData(server)
   if (!data) {
-    log(
+    logger(
       'Warning: PostHog MCP analytics is unable to find server tracking data. Please ensure you have called instrument(server, options) before using tool calls.'
     )
     return await originalCallToolHandler(request, extra)
@@ -280,20 +251,23 @@ async function handleToolCallRequest(
       extra,
       eventType: MCPAnalyticsEventType.mcpMissingCapability,
       explicitContextIntent: context,
-      execute: async () => handleReportMissing({ context }),
+      execute: async () => handleReportMissing({ context }, data.logger),
     })
   }
 
-  // The high-level handler re-derives arguments from the original request and
-  // strips the injected params inside the wrapped callback, so we hand it the
-  // original request rather than the conversation-stripped one. Errors thrown
-  // by the tool are stashed on `extra` by the callback wrapper; surface them.
+  const toolName = request.params?.name
+  const registeredTool = toolName ? highLevelServer._registeredTools[toolName] : undefined
+
+  // Strip SDK-owned arguments before the MCP SDK validates the registered Zod
+  // schema. The callback wrapper repeats the ownership-aware cleanup as a
+  // defensive fallback for direct callback invocation.
   return await captureToolCall({
     server,
     data,
     request,
     extra,
-    execute: () => originalCallToolHandler(request, extra),
+    execute: (downstreamRequest) => originalCallToolHandler(downstreamRequest as MCPRequest, extra),
+    parameterOwnership: registeredTool ? getAnalyticsParameterOwnership(registeredTool.inputSchema) : undefined,
     takeCapturedError: () => {
       const captured = extra?.__mcp_analytics_error
       if (extra) {
@@ -304,29 +278,32 @@ async function handleToolCallRequest(
   })
 }
 
-export function instrumentHighLevelServer(server: HighLevelMCPServerLike): void {
+export function instrumentHighLevelServer(server: HighLevelMCPServerLike, logger: LoggerFn): void {
   try {
     const lowLevelServer = server.server
     const mcpAnalyticsData = getServerTrackingData(lowLevelServer)
 
     // Patch already existing handlers, and patch setRequestHandler to capture dynamically created handlers.
-    const handlers = {
-      initialize: handleInitializeRequest,
-      'tools/list': handleListToolsRequest,
-      'tools/call': handleToolCallRequest,
+    const handlers: Record<string, HandlerPatch> = {
+      initialize: (trackedServer, originalHandler, request, extra) =>
+        handleInitializeRequest(trackedServer, originalHandler, request, extra, logger),
+      'tools/list': (trackedServer, originalHandler, request, extra) =>
+        handleListToolsRequest(trackedServer, originalHandler, request, extra, logger),
+      'tools/call': (trackedServer, originalHandler, request, extra) =>
+        handleToolCallRequest(server, trackedServer, originalHandler, request, extra, logger),
     }
     patchRequestHandlers(lowLevelServer, handlers)
 
-    server._registeredTools = addTracingToToolRegistry(server._registeredTools, server)
+    server._registeredTools = addTracingToToolRegistry(server._registeredTools, server, logger)
 
     if (mcpAnalyticsData) {
       seedToolDescriptionsFromRegistry(mcpAnalyticsData.toolDescriptions, server._registeredTools)
       seedToolCategoriesFromRegistry(mcpAnalyticsData.toolCategories, server._registeredTools)
     }
 
-    setupListenerToRegisteredTools(server)
+    setupListenerToRegisteredTools(server, logger)
   } catch (error) {
-    log(`Warning: Failed to setup tool call instrumentation - ${error}`)
+    logger(`Warning: Failed to setup tool call instrumentation - ${error}`)
   }
 }