@posthog/mcp 0.10.5 → 0.10.7

package/src/extensions/instrumentation.ts

@@ -12,6 +12,7 @@ import type {
   MCPServerLike,
   McpEvent,
   ServerClientInfoLike,
+  SessionInfo,
 } from '../types'
 import { getAnalyticsParameterOwnership, stripOwnedAnalyticsArguments } from './analytics-parameters'
 import { addContextParameterToTools, getContextDescription, isContextEnabled } from './context-parameters'
@@ -27,11 +28,11 @@ import { captureEvent } from './capture'
 import { MCPAnalyticsEventType } from './event-types'
 import { captureException } from './exceptions'
 import { resolveToolCallIntent, setEventIntent, setExplicitContextIntent } from './intent'
-import { getServerTrackingData, handleIdentify, setServerTrackingData } from './internal'
+import { getServerTrackingData, handleIdentify, setServerTrackingData, withIdentity } from './internal'
 import type { LoggerFn } from './logger'
 import { buildCapturedMcpParameters } from './mcp-payloads'
 import { getLiteralValue, getObjectShape } from './mcp-sdk-compat'
-import { getSessionId, newSessionId } from './session'
+import { getSessionId, getSessionInfo, newSessionId } from './session'
 import { encodeSessionId, readMcpSessionHeader, writeSessionIdToTransport } from './session-token'
 import { getReportMissingToolDescriptor, resolveMissingCapabilityToolName } from './tools'
 import { applyResolvedMetadata, isToolResultError } from './tracing-helpers'
@@ -98,20 +99,27 @@ export async function captureToolCall(params: TraceToolCallParams): Promise<unkn
     explicitContextIntent,
     takeCapturedError,
   } = params
-
-  const ownership = getActiveAnalyticsParameterOwnership(data, request.params?.name, parameterOwnership)
+  const resolvedEventType = eventType ?? MCPAnalyticsEventType.mcpToolsCall
+  const ownership = getActiveAnalyticsParameterOwnership(
+    data,
+    request.params?.name,
+    parameterOwnership,
+    resolvedEventType === MCPAnalyticsEventType.mcpMissingCapability
+  )
   const conversation = resolveConversationId(
     ownership.conversationId,
     request.params?.arguments,
     request.params?.name,
-    resolveMissingCapabilityToolName(data.options)
+    resolvedEventType === MCPAnalyticsEventType.mcpMissingCapability
+      ? resolveMissingCapabilityToolName(data.options)
+      : ''
   )
   const downstreamRequest = cloneRequestWithoutOwnedAnalyticsArguments(request, ownership)
 
   // Prepare the event in isolation: if identity/metadata/intent resolution
   // throws, we drop instrumentation for this call but still run the tool.
   const startTime = new Date()
-  const event = await prepareToolCallEvent(
+  const preparedEvent = await prepareToolCallEvent(
     server,
     data,
     request,
@@ -120,32 +128,37 @@ export async function captureToolCall(params: TraceToolCallParams): Promise<unkn
     startTime,
     conversation,
     ownership,
-    eventType ?? MCPAnalyticsEventType.mcpToolsCall
+    resolvedEventType
   )
-  if (event && explicitContextIntent) {
-    setExplicitContextIntent(event, explicitContextIntent)
+  if (preparedEvent && explicitContextIntent) {
+    setExplicitContextIntent(preparedEvent.event, explicitContextIntent)
   }
 
   let result: unknown
   try {
     result = await execute(downstreamRequest)
   } catch (error) {
-    publishFailedToolEvent(server, event, error, startTime, conversation, data.logger)
+    publishFailedToolEvent(server, preparedEvent, error, startTime, conversation, data.logger)
     throw error
   }
 
-  const finalResult = applyConversationPromptBack(event, result, conversation)
-  publishSuccessfulToolEvent(server, event, finalResult, startTime, data.logger, takeCapturedError)
+  const finalResult = applyConversationPromptBack(preparedEvent?.event ?? null, result, conversation)
+  publishSuccessfulToolEvent(server, preparedEvent, finalResult, startTime, data.logger, takeCapturedError)
   return finalResult
 }
 
+interface PreparedToolEvent {
+  event: McpEvent
+  requestAttribution: SessionInfo
+}
+
 function getActiveAnalyticsParameterOwnership(
   data: MCPAnalyticsData,
   toolName: string | undefined,
-  override?: AnalyticsParameterOwnership
+  override: AnalyticsParameterOwnership | undefined,
+  isMissingCapabilityTool: boolean
 ): AnalyticsParameterOwnership {
   const ownership = override ?? (toolName ? data.toolAnalyticsParameterOwnership.get(toolName) : undefined)
-  const isMissingCapabilityTool = toolName === resolveMissingCapabilityToolName(data.options)
   return {
     context: !isMissingCapabilityTool && isContextEnabled(data.options.context) && ownership?.context === true,
     conversationId:
@@ -181,10 +194,14 @@ async function prepareToolCallEvent(
   conversation: ConversationIdResolution,
   ownership: AnalyticsParameterOwnership,
   eventType: MCPAnalyticsEventType
-): Promise<McpEvent | null> {
+): Promise<PreparedToolEvent | null> {
   try {
     const sessionId = getSessionId(server, extra)
-    await handleIdentify(server, data, sessionId, request, extra)
+    // Snapshot token/client/protocol metadata synchronously, before identify or
+    // metadata callbacks can yield and let another request replace shared state.
+    const sessionInfo = getSessionInfo(server, data, sessionId)
+    const identity = await handleIdentify(server, data, sessionId, request, sessionInfo, extra)
+    const requestAttribution = withIdentity(sessionInfo, identity)
 
     const toolName = request.params?.name
     const event: McpEvent = {
@@ -204,7 +221,7 @@ async function prepareToolCallEvent(
 
     await applyResolvedMetadata(event, data, request, extra)
     setEventIntent(event, await resolveToolCallIntent(data, request, ownership.context, extra))
-    return event
+    return { event, requestAttribution }
   } catch (error) {
     data.logger(
       `Warning: PostHog MCP analytics instrumentation failed for tool ${request.params?.name}, the tool will still run - ${error}`
@@ -237,15 +254,16 @@ function applyConversationPromptBack(
 
 function publishSuccessfulToolEvent(
   server: MCPServerLike,
-  event: McpEvent | null,
+  preparedEvent: PreparedToolEvent | null,
   result: unknown,
   startTime: Date,
   logger: LoggerFn,
   takeCapturedError?: () => unknown
 ): void {
-  if (!event) {
+  if (!preparedEvent) {
     return
   }
+  const { event, requestAttribution } = preparedEvent
   try {
     if (isToolResultError(result)) {
       event.isError = true
@@ -256,7 +274,7 @@ function publishSuccessfulToolEvent(
     }
     event.response = result
     event.duration = Date.now() - startTime.getTime()
-    captureEvent(server, event, logger)
+    captureEvent(server, event, logger, requestAttribution)
   } catch (error) {
     logger(`Warning: PostHog MCP analytics failed to publish tool event - ${error}`)
   }
@@ -264,15 +282,16 @@ function publishSuccessfulToolEvent(
 
 function publishFailedToolEvent(
   server: MCPServerLike,
-  event: McpEvent | null,
+  preparedEvent: PreparedToolEvent | null,
   error: unknown,
   startTime: Date,
   conversation: ConversationIdResolution,
   logger: LoggerFn
 ): void {
-  if (!event) {
+  if (!preparedEvent) {
     return
   }
+  const { event, requestAttribution } = preparedEvent
   try {
     if (conversation.minted) {
       event.conversationId = undefined
@@ -280,7 +299,7 @@ function publishFailedToolEvent(
     event.isError = true
     event.error = captureException(error)
     event.duration = Date.now() - startTime.getTime()
-    captureEvent(server, event, logger)
+    captureEvent(server, event, logger, requestAttribution)
   } catch (publishError) {
     logger(`Warning: PostHog MCP analytics failed to publish failed tool event - ${publishError}`)
   }
@@ -307,11 +326,27 @@ export type HandlerPatch = (
  * that register handlers post-construction work — e.g. `@rekog/mcp-nest` hands a
  * bare server to instrument() and only then registers its handlers.
  */
+const originalRequestHandlers = new WeakMap<MCPServerLike, Map<string, MCPRequestHandler>>()
+
+function rememberOriginalRequestHandler(
+  server: MCPServerLike,
+  handlerName: string,
+  originalHandler: MCPRequestHandler
+): void {
+  let handlers = originalRequestHandlers.get(server)
+  if (!handlers) {
+    handlers = new Map()
+    originalRequestHandlers.set(server, handlers)
+  }
+  handlers.set(handlerName, originalHandler)
+}
+
 export function patchRequestHandlers(server: MCPServerLike, patches: Record<string, HandlerPatch>): void {
   // Monkey patch existing handlers.
   for (const [handlerName, patch] of Object.entries(patches)) {
     const originalHandler = server._requestHandlers.get(handlerName)
     if (originalHandler) {
+      rememberOriginalRequestHandler(server, handlerName, originalHandler)
       server._requestHandlers.set(handlerName, (request, extra) => patch(server, originalHandler, request, extra))
     }
   }
@@ -322,14 +357,55 @@ export function patchRequestHandlers(server: MCPServerLike, patches: Record<stri
     const shape = getObjectShape(requestSchema)
     const handlerName = shape?.method ? getLiteralValue(shape.method) : undefined
     const patch = typeof handlerName === 'string' ? patches[handlerName] : undefined
-    if (!patch) {
+    if (!patch || typeof handlerName !== 'string') {
       return originalSetRequestHandler(requestSchema, originalHandler)
     }
 
-    return originalSetRequestHandler(requestSchema, (request, extra) => patch(server, originalHandler, request, extra))
+    // Register first so the MCP SDK's request/result validation stays inside
+    // our analytics wrapper, matching handlers that existed before instrument().
+    const result = originalSetRequestHandler(requestSchema, originalHandler)
+    const registeredHandler = server._requestHandlers.get(handlerName)
+    if (registeredHandler) {
+      rememberOriginalRequestHandler(server, handlerName, registeredHandler)
+      server._requestHandlers.set(handlerName, (request, extra) => patch(server, registeredHandler, request, extra))
+    }
+    return result
   }) as MCPServerLike['setRequestHandler']
 }
 
+/**
+ * Checks the server's raw listing for a real owner of a candidate virtual tool.
+ * This does not depend on a previous client request and does not call the
+ * instrumented list wrapper, so it neither injects PostHog tools nor captures a
+ * synthetic tools/list event. `undefined` fails open to the real dispatcher.
+ */
+export async function isToolAdvertised(
+  server: MCPServerLike,
+  toolName: string,
+  extra: CompatibleRequestHandlerExtra | undefined,
+  logger: LoggerFn
+): Promise<boolean | undefined> {
+  const listHandler = originalRequestHandlers.get(server)?.get('tools/list')
+  if (!listHandler || !server._requestHandlers.has('tools/list')) {
+    return undefined
+  }
+
+  try {
+    const response = (await listHandler({ method: 'tools/list', params: {} }, extra)) as ListToolsResult
+    if (!response || !Array.isArray(response.tools)) {
+      return undefined
+    }
+    // Match the page the current list instrumentation can expose. Pagination
+    // passthrough is handled separately from missing-capability ownership.
+    return response.tools.some((tool) => tool?.name === toolName)
+  } catch (error) {
+    logger(
+      `Warning: PostHog MCP analytics could not determine whether "${toolName}" is advertised; delegating to the server - ${error}`
+    )
+    return undefined
+  }
+}
+
 /**
  * Captures each `tools/list` and injects the SDK-managed tools (context
  * parameter, conversation id, `get_more_tools`) into the returned list.
@@ -343,8 +419,12 @@ export async function handleListToolsRequest(
 ): Promise<{ tools: ListToolsResult['tools'] }> {
   const data = getServerTrackingData(server)
   const startTime = new Date()
+  const sessionId = getSessionId(server, extra)
+  // Snapshot before metadata resolution or the list handler can yield to a
+  // concurrent request using the same instrumented server.
+  const requestAttribution = getSessionInfo(server, data, sessionId)
   const event: McpEvent = {
-    sessionId: getSessionId(server, extra),
+    sessionId,
     parameters: buildCapturedMcpParameters(request),
     eventType: MCPAnalyticsEventType.mcpToolsList,
     timestamp: startTime,
@@ -355,7 +435,15 @@ export async function handleListToolsRequest(
     await applyResolvedMetadata(event, data, request, extra)
   }
 
-  const tools = await getTracedToolsList(server, originalListToolsHandler, request, extra, event, logger)
+  const tools = await getTracedToolsList(
+    server,
+    originalListToolsHandler,
+    request,
+    extra,
+    event,
+    logger,
+    requestAttribution
+  )
 
   if (!data) {
     logger(
@@ -371,7 +459,7 @@ export async function handleListToolsRequest(
     event.error = captureException('No tools were sent to MCP client.')
     event.isError = true
     event.duration = Date.now() - startTime.getTime()
-    captureEvent(server, event, data.logger)
+    captureEvent(server, event, data.logger, requestAttribution)
     return { tools }
   }
 
@@ -379,7 +467,7 @@ export async function handleListToolsRequest(
   event.listedToolNames = collectListedToolNames(tools)
   event.isError = false
   event.duration = Date.now() - startTime.getTime()
-  captureEvent(server, event, data.logger)
+  captureEvent(server, event, data.logger, requestAttribution)
   return { tools }
 }
 
@@ -409,30 +497,39 @@ async function getTracedToolsList(
   request: MCPRequestLike,
   extra: CompatibleRequestHandlerExtra | undefined,
   event: McpEvent,
-  logger: LoggerFn
+  logger: LoggerFn,
+  requestAttribution: SessionInfo
 ): Promise<ListToolsResult['tools']> {
   try {
     const data = getServerTrackingData(server)
     const originalResponse = (await originalListToolsHandler(request, extra)) as ListToolsResult
-    let tools = originalResponse.tools || []
+    // Injection must not mutate arrays reused or frozen by the server.
+    let tools = [...(originalResponse.tools || [])]
 
     if (data) {
       cacheToolAnalyticsParameterOwnership(data.toolAnalyticsParameterOwnership, tools)
     }
-
     if (data && isContextEnabled(data.options.context)) {
       tools = addContextParameterToTools(tools, getContextDescription(data.options.context), data.logger)
     }
 
-    if (data?.options.enableConversationId) {
-      tools = addConversationIdToTools(tools, resolveMissingCapabilityToolName(data.options), data.logger)
-    }
-
-    if (data?.options.reportMissing) {
+    if (data) {
       const missingToolName = resolveMissingCapabilityToolName(data.options)
-      const alreadyPresent = tools.some((tool) => tool?.name === missingToolName)
-      if (!alreadyPresent) {
-        tools.push(getReportMissingToolDescriptor(missingToolName))
+      let injectedMissingCapabilityTool = false
+      if (data.options.reportMissing) {
+        const alreadyPresent = tools.some((tool) => tool?.name === missingToolName)
+        if (alreadyPresent) {
+          data.logger(
+            `Warning: Cannot inject missing-capability tool "${missingToolName}" because a real tool already uses that name. The real tool will not be intercepted.`
+          )
+        } else {
+          tools.push(getReportMissingToolDescriptor(missingToolName))
+          injectedMissingCapabilityTool = true
+        }
+      }
+
+      if (data.options.enableConversationId) {
+        tools = addConversationIdToTools(tools, missingToolName, injectedMissingCapabilityTool, data.logger)
       }
     }
 
@@ -449,7 +546,7 @@ async function getTracedToolsList(
     event.error = captureException(error)
     event.isError = true
     event.duration = event.timestamp ? Date.now() - event.timestamp.getTime() : 0
-    captureEvent(server, event, logger)
+    captureEvent(server, event, logger, requestAttribution)
     throw error
   }
 }
@@ -558,9 +655,10 @@ function mintStatelessSessionOnInitialize(
  */
 function upgradeMintedTokenToNegotiated(
   server: MCPServerLike,
-  data: MCPAnalyticsData,
   mintedSessionId: string,
-  negotiatedProtocolVersion: string | undefined
+  sessionInfo: SessionInfo,
+  negotiatedProtocolVersion: string | undefined,
+  logger: LoggerFn
 ): void {
   try {
     const transport = server.transport
@@ -569,13 +667,13 @@ function upgradeMintedTokenToNegotiated(
     }
     const token = encodeSessionId({
       sessionId: mintedSessionId,
-      clientName: data.sessionInfo.clientName,
-      clientVersion: data.sessionInfo.clientVersion,
+      clientName: sessionInfo.clientName,
+      clientVersion: sessionInfo.clientVersion,
       protocolVersion: negotiatedProtocolVersion,
     })
     writeSessionIdToTransport(transport, token)
   } catch (error) {
-    data.logger(`Warning: PostHog MCP analytics failed to upgrade the stateless session token - ${error}`)
+    logger(`Warning: PostHog MCP analytics failed to upgrade the stateless session token - ${error}`)
   }
 }
 
@@ -617,7 +715,17 @@ export async function handleInitializeRequest(
   // Mint first so the `$mcp_initialize` event below already carries the minted id.
   const mintedSessionId = mintStatelessSessionOnInitialize(server, data, request, extra)
   const sessionId = getSessionId(server, extra)
-  await handleIdentify(server, data, sessionId, request, extra)
+  // Snapshot before identify, metadata, or the initialize handler can yield to
+  // another request using the same instrumented server.
+  const sessionInfo = getSessionInfo(server, data, sessionId)
+  const initializeClientInfo = readInitializeClientInfo(request)
+  const requestSessionInfo: SessionInfo = {
+    ...sessionInfo,
+    clientName: initializeClientInfo?.name ?? sessionInfo.clientName,
+    clientVersion: initializeClientInfo?.version ?? sessionInfo.clientVersion,
+  }
+  const identity = await handleIdentify(server, data, sessionId, request, requestSessionInfo, extra)
+  const requestAttribution = withIdentity(requestSessionInfo, identity)
 
   const event: McpEvent = {
     sessionId,
@@ -640,12 +748,19 @@ export async function handleInitializeRequest(
   // re-mint the token so pods replaying it report the negotiated version too.
   const negotiatedProtocolVersion = readProtocolVersion(result, request)
   event.protocolVersion = negotiatedProtocolVersion
-  data.sessionInfo.protocolVersion = negotiatedProtocolVersion
-  setServerTrackingData(server, data)
+  // Do not let a delayed initialize overwrite whichever session became current
+  // while its callbacks or the original handler were awaiting.
+  if (data.sessionId === sessionId) {
+    data.sessionInfo = { ...data.sessionInfo, protocolVersion: negotiatedProtocolVersion }
+    setServerTrackingData(server, data)
+  }
   if (mintedSessionId) {
-    upgradeMintedTokenToNegotiated(server, data, mintedSessionId, negotiatedProtocolVersion)
+    upgradeMintedTokenToNegotiated(server, mintedSessionId, requestSessionInfo, negotiatedProtocolVersion, data.logger)
   }
-  captureEvent(server, event, data.logger)
+  captureEvent(server, event, data.logger, {
+    ...requestAttribution,
+    protocolVersion: negotiatedProtocolVersion,
+  })
   return result
 }
 

package/src/extensions/intent.ts

@@ -56,12 +56,7 @@ export async function resolveToolCallIntent(
   extra?: CompatibleRequestHandlerExtra
 ): Promise<ResolvedIntent | null> {
   const contextArgument = getContextArgument(request)
-  if (
-    analyticsOwnsContext &&
-    isContextEnabled(data.options.context) &&
-    request.params?.name !== 'get_more_tools' &&
-    contextArgument
-  ) {
+  if (analyticsOwnsContext && isContextEnabled(data.options.context) && contextArgument) {
     return { intent: contextArgument, source: 'context_parameter' }
   }
 

package/src/extensions/internal.ts

@@ -9,6 +9,7 @@ import type {
   MCPRequestLike,
   MCPServerLike,
   McpEvent,
+  SessionInfo,
   UserIdentity,
 } from '../types'
 import { MCPAnalyticsEventType } from './event-types'
@@ -128,10 +129,13 @@ export async function handleIdentify(
   data: MCPAnalyticsData,
   sessionId: string,
   request: MCPRequestLike,
+  requestAttribution: SessionInfo,
   extra?: CompatibleRequestHandlerExtra
-): Promise<void> {
+): Promise<UserIdentity | undefined> {
+  const identityBeforeRequest = data.identifiedSessions.get(sessionId)
+  const sessionSourceBeforeIdentify = data.sessionSource
   if (!data.options.identify) {
-    return
+    return identityBeforeRequest
   }
 
   const identifyEvent: McpEvent = {
@@ -162,23 +166,34 @@ export async function handleIdentify(
       // announced); revisit with the stateless-by-default rework.
       const changed = previousIdentity !== undefined && !areIdentitiesEqual(previousIdentity, mergedIdentity)
       const firstSeen = previousIdentity === undefined
-      const announcedAtInitialize = data.sessionSource === 'token' && request.method !== 'initialize'
+      const announcedAtInitialize = sessionSourceBeforeIdentify === 'token' && request.method !== 'initialize'
       const shouldPublish = changed || (firstSeen && !announcedAtInitialize)
 
       data.identifiedSessions.set(sessionId, mergedIdentity)
 
       if (shouldPublish) {
         data.logger(`Identified session ${sessionId}`)
-        captureEvent(server, identifyEvent, data.logger)
+        captureEvent(server, identifyEvent, data.logger, withIdentity(requestAttribution, mergedIdentity))
       }
-    } else {
-      data.logger(`Warning: Supplied identify function returned null for session ${sessionId}`)
+      return mergedIdentity
     }
+
+    data.logger(`Warning: Supplied identify function returned null for session ${sessionId}`)
   } catch (error) {
     data.logger(
       `Error: User supplied identify function threw an error while identifying session ${sessionId} - ${error}`
     )
   }
+  return identityBeforeRequest
+}
+
+export function withIdentity(sessionInfo: SessionInfo, identity: UserIdentity | undefined): SessionInfo {
+  return {
+    ...sessionInfo,
+    identifyActorGivenId: identity?.distinctId,
+    identifyActorData: identity?.properties || {},
+    identifyActorGroups: identity?.groups,
+  }
 }
 
 /**

package/src/extensions/session.ts

@@ -125,5 +125,5 @@ export function getSessionInfo(
 
   data.sessionInfo = sessionInfo
   setServerTrackingData(server, data)
-  return data.sessionInfo
+  return { ...sessionInfo }
 }

package/src/version.ts

@@ -1 +1 @@
-export const version = '0.10.5'
+export const version = '0.10.7'