@posthog/mcp 0.0.9 → 0.1.16

package/src/extensions/instrumentation.ts

@@ -0,0 +1,415 @@
+import {
+  InitializeRequestSchema,
+  ListToolsRequestSchema,
+  type ListToolsResult,
+} from '@modelcontextprotocol/sdk/types.js'
+import type { CompatibleRequestHandlerExtra, MCPAnalyticsData, MCPRequestLike, MCPServerLike, McpEvent } from '../types'
+import { addContextParameterToTools, getContextDescription, isContextEnabled } from './context-parameters'
+import {
+  addConversationIdToTools,
+  type ConversationIdResolution,
+  canInjectConversationIdPromptBack,
+  cloneRequestWithoutConversationId,
+  injectConversationIdPromptBack,
+  resolveConversationId,
+} from './conversation-id'
+import { captureEvent } from './capture'
+import { MCPAnalyticsEventType } from './event-types'
+import { captureException } from './exceptions'
+import { resolveToolCallIntent, setEventIntent, setExplicitContextIntent } from './intent'
+import { getServerTrackingData, handleIdentify } from './internal'
+import { log } from './logger'
+import { buildCapturedMcpParameters } from './mcp-payloads'
+import { getServerSessionId } from './session'
+import { GET_MORE_TOOLS_NAME, getReportMissingToolDescriptor } from './tools'
+import { applyResolvedMetadata, isToolResultError } from './tracing-helpers'
+
+/**
+ * Single instrumentation core shared by the low-level (`Server`) and high-level
+ * (`McpServer`) wrappers. The two entrypoints differ only in how they reach the
+ * underlying tool — they both funnel the tool-call lifecycle through
+ * {@link captureToolCall}, so error handling, conversation-id minting, session
+ * attribution, and event capture are defined exactly once.
+ */
+
+type MCPRequestHandler = (request: MCPRequestLike, extra?: CompatibleRequestHandlerExtra) => Promise<unknown>
+
+/**
+ * Runs the underlying tool. Receives the request with the SDK-injected
+ * `conversation_id` stripped; an adapter is free to ignore it (the high-level
+ * path strips arguments inside the wrapped callback instead).
+ */
+type ToolExecutor = (downstreamRequest: MCPRequestLike) => Promise<unknown>
+
+interface TraceToolCallParams {
+  server: MCPServerLike
+  data: MCPAnalyticsData
+  request: MCPRequestLike
+  extra?: CompatibleRequestHandlerExtra
+  execute: ToolExecutor
+  /**
+   * Event type to capture. Defaults to a tool call; the `get_more_tools` virtual
+   * tool passes `mcpMissingCapability` so it records a capability gap rather than
+   * a tool invocation.
+   */
+  eventType?: MCPAnalyticsEventType
+  /**
+   * When set, used verbatim as the captured intent (source `context_parameter`)
+   * instead of running `resolveToolCallIntent`. Used by the `get_more_tools`
+   * virtual tool, which carries its intent in the `context` argument.
+   */
+  explicitContextIntent?: string
+  /**
+   * Optional accessor for an error the executor captured out-of-band. The
+   * high-level SDK turns thrown tool errors into `isError: true` results before
+   * they reach us, so the wrapped callback stashes the original error and we
+   * read it here to capture the real stack rather than the result envelope.
+   */
+  takeCapturedError?: () => unknown
+}
+
+/**
+ * The shared tool-call lifecycle: resolve conversation id, build + enrich the
+ * analytics event, run the tool, then capture success/failure.
+ *
+ * Analytics is isolated from the tool path on both sides — a failure while
+ * preparing or publishing the event can never change what the tool returns or
+ * throws, and the tool's own errors are always re-thrown to the caller.
+ */
+export async function captureToolCall(params: TraceToolCallParams): Promise<unknown> {
+  const { server, data, request, extra, execute, eventType, explicitContextIntent, takeCapturedError } = params
+
+  const conversation = resolveConversationId(
+    data.options.enableConversationId ?? false,
+    request.params?.arguments,
+    request.params?.name
+  )
+  const downstreamRequest = conversation.conversationId ? cloneRequestWithoutConversationId(request) : request
+
+  // 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(
+    server,
+    data,
+    request,
+    downstreamRequest,
+    extra,
+    startTime,
+    conversation,
+    eventType ?? MCPAnalyticsEventType.mcpToolsCall
+  )
+  if (event && explicitContextIntent) {
+    setExplicitContextIntent(event, explicitContextIntent)
+  }
+
+  let result: unknown
+  try {
+    result = await execute(downstreamRequest)
+  } catch (error) {
+    publishFailedToolEvent(server, event, error, startTime, conversation)
+    throw error
+  }
+
+  const finalResult = applyConversationPromptBack(event, result, conversation)
+  publishSuccessfulToolEvent(server, event, finalResult, startTime, takeCapturedError)
+  return finalResult
+}
+
+async function prepareToolCallEvent(
+  server: MCPServerLike,
+  data: MCPAnalyticsData,
+  request: MCPRequestLike,
+  downstreamRequest: MCPRequestLike,
+  extra: CompatibleRequestHandlerExtra | undefined,
+  startTime: Date,
+  conversation: ConversationIdResolution,
+  eventType: MCPAnalyticsEventType
+): Promise<McpEvent | null> {
+  try {
+    const sessionId = getServerSessionId(server, extra)
+    await handleIdentify(server, data, sessionId, request, extra)
+
+    const toolName = request.params?.name
+    const event: McpEvent = {
+      sessionId,
+      conversationId: conversation.conversationId,
+      resourceName: toolName || 'Unknown Tool Name',
+      parameters: buildCapturedMcpParameters(downstreamRequest),
+      eventType,
+      timestamp: startTime,
+      toolDescription: toolName ? data.toolDescriptions.get(toolName) : undefined,
+    }
+
+    await applyResolvedMetadata(event, data, request, extra)
+    setEventIntent(event, await resolveToolCallIntent(data, request, extra))
+    return event
+  } catch (error) {
+    log(
+      `Warning: PostHog MCP analytics instrumentation failed for tool ${request.params?.name}, the tool will still run - ${error}`
+    )
+    return null
+  }
+}
+
+/**
+ * When we minted a conversation id, append the prompt-back so the agent echoes
+ * it on subsequent calls. If the result can't carry it, clear the id off the
+ * event so analytics doesn't show an orphan the agent never received.
+ */
+function applyConversationPromptBack(
+  event: McpEvent | null,
+  result: unknown,
+  conversation: ConversationIdResolution
+): unknown {
+  if (!conversation.minted) {
+    return result
+  }
+  if (canInjectConversationIdPromptBack(result)) {
+    return injectConversationIdPromptBack(result, conversation.conversationId)
+  }
+  if (event) {
+    event.conversationId = undefined
+  }
+  return result
+}
+
+function publishSuccessfulToolEvent(
+  server: MCPServerLike,
+  event: McpEvent | null,
+  result: unknown,
+  startTime: Date,
+  takeCapturedError?: () => unknown
+): void {
+  if (!event) {
+    return
+  }
+  try {
+    if (isToolResultError(result)) {
+      event.isError = true
+      const capturedError = takeCapturedError?.()
+      event.error = captureException(capturedError ?? result)
+    } else {
+      event.isError = false
+    }
+    event.response = result
+    event.duration = Date.now() - startTime.getTime()
+    captureEvent(server, event)
+  } catch (error) {
+    log(`Warning: PostHog MCP analytics failed to publish tool event - ${error}`)
+  }
+}
+
+function publishFailedToolEvent(
+  server: MCPServerLike,
+  event: McpEvent | null,
+  error: unknown,
+  startTime: Date,
+  conversation: ConversationIdResolution
+): void {
+  if (!event) {
+    return
+  }
+  try {
+    if (conversation.minted) {
+      event.conversationId = undefined
+    }
+    event.isError = true
+    event.error = captureException(error)
+    event.duration = Date.now() - startTime.getTime()
+    captureEvent(server, event)
+  } catch (publishError) {
+    log(`Warning: PostHog MCP analytics failed to publish failed tool event - ${publishError}`)
+  }
+}
+
+// --- tools/list -----------------------------------------------------------
+
+type MCPServerWithCapabilities = MCPServerLike & {
+  _capabilities?: {
+    tools?: unknown
+  }
+}
+
+const listToolsTracingSetup = new WeakMap<MCPServerLike, boolean>()
+
+/**
+ * Wraps the server's `tools/list` handler so each listing is captured and the
+ * SDK-managed tools (context parameter, conversation id, `get_more_tools`) are
+ * injected. Idempotent per server. Works for both low-level and high-level
+ * servers — the high-level wrapper passes its underlying `server`.
+ */
+export function instrumentToolsListHandler(server: MCPServerLike): void {
+  if (!(server as MCPServerWithCapabilities)._capabilities?.tools) {
+    return
+  }
+  if (listToolsTracingSetup.get(server)) {
+    return
+  }
+
+  const originalListToolsHandler = server._requestHandlers.get('tools/list')
+  if (!originalListToolsHandler) {
+    return
+  }
+
+  try {
+    server.setRequestHandler(
+      ListToolsRequestSchema,
+      async (request, extra) => await handleListToolsRequest(server, originalListToolsHandler, request, extra)
+    )
+    listToolsTracingSetup.set(server, true)
+  } catch (error) {
+    log(`Warning: Failed to override list tools handler - ${error}`)
+  }
+}
+
+async function handleListToolsRequest(
+  server: MCPServerLike,
+  originalListToolsHandler: MCPRequestHandler,
+  request: MCPRequestLike,
+  extra?: CompatibleRequestHandlerExtra
+): Promise<{ tools: ListToolsResult['tools'] }> {
+  const data = getServerTrackingData(server)
+  const startTime = new Date()
+  const event: McpEvent = {
+    sessionId: getServerSessionId(server, extra),
+    parameters: buildCapturedMcpParameters(request),
+    eventType: MCPAnalyticsEventType.mcpToolsList,
+    timestamp: startTime,
+  }
+
+  if (data) {
+    await applyResolvedMetadata(event, data, request, extra)
+  }
+
+  const tools = await getTracedToolsList(server, originalListToolsHandler, request, extra, event)
+
+  if (!data) {
+    log(
+      'Warning: PostHog MCP analytics is unable to find server tracking data. Please ensure you have called instrument(server, options) before using tool calls.'
+    )
+    return { tools }
+  }
+
+  if (tools.length === 0) {
+    log(
+      'Warning: No tools found in the original list. This is likely due to the tools not being registered before PostHog MCP analytics.instrument().'
+    )
+    event.error = captureException('No tools were sent to MCP client.')
+    event.isError = true
+    event.duration = Date.now() - startTime.getTime()
+    captureEvent(server, event)
+    return { tools }
+  }
+
+  event.response = { tools }
+  event.listedToolNames = collectListedToolNames(tools)
+  event.isError = false
+  event.duration = Date.now() - startTime.getTime()
+  captureEvent(server, event)
+  return { tools }
+}
+
+function collectListedToolNames(tools: ListToolsResult['tools'] | undefined): string[] | undefined {
+  if (!tools || tools.length === 0) {
+    return
+  }
+  const names = tools.map((tool) => tool?.name).filter((name): name is string => typeof name === 'string')
+  return names.length > 0 ? names : undefined
+}
+
+async function getTracedToolsList(
+  server: MCPServerLike,
+  originalListToolsHandler: MCPRequestHandler,
+  request: MCPRequestLike,
+  extra: CompatibleRequestHandlerExtra | undefined,
+  event: McpEvent
+): Promise<ListToolsResult['tools']> {
+  try {
+    const data = getServerTrackingData(server)
+    const originalResponse = (await originalListToolsHandler(request, extra)) as ListToolsResult
+    let tools = originalResponse.tools || []
+
+    if (data && isContextEnabled(data.options.context)) {
+      tools = addContextParameterToTools(tools, getContextDescription(data.options.context))
+    }
+
+    if (data?.options.enableConversationId) {
+      tools = addConversationIdToTools(tools)
+    }
+
+    if (data?.options.reportMissing) {
+      const alreadyPresent = tools.some((tool) => tool?.name === GET_MORE_TOOLS_NAME)
+      if (!alreadyPresent) {
+        tools.push(getReportMissingToolDescriptor())
+      }
+    }
+
+    if (data) {
+      cacheToolDescriptions(data.toolDescriptions, tools)
+    }
+
+    return tools
+  } catch (error) {
+    log(
+      `Warning: Original list tools handler failed, this suggests an error PostHog MCP analytics did not cause - ${error}`
+    )
+    event.error = captureException(error)
+    event.isError = true
+    event.duration = event.timestamp ? Date.now() - event.timestamp.getTime() : 0
+    captureEvent(server, event)
+    throw error
+  }
+}
+
+export function cacheToolDescriptions(cache: Map<string, string>, tools: ListToolsResult['tools'] | undefined): void {
+  if (!tools) {
+    return
+  }
+  for (const tool of tools) {
+    if (tool?.name && typeof tool.description === 'string') {
+      cache.set(tool.name, tool.description)
+    }
+  }
+}
+
+// --- initialize -----------------------------------------------------------
+
+/**
+ * Wraps the server's `initialize` handler so the connection handshake is
+ * captured (and identity resolved) before the original handler runs.
+ */
+export function instrumentInitializeHandler(server: MCPServerLike): void {
+  const originalInitializeHandler = server._requestHandlers.get('initialize')
+  if (!originalInitializeHandler) {
+    return
+  }
+
+  server.setRequestHandler(InitializeRequestSchema, async (request, extra) => {
+    const data = getServerTrackingData(server)
+    if (!data) {
+      log(
+        '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 originalInitializeHandler(request, extra)
+    }
+
+    const sessionId = getServerSessionId(server, extra)
+    await handleIdentify(server, data, sessionId, request, extra)
+
+    const event: McpEvent = {
+      sessionId,
+      resourceName: request.params?.name || 'Unknown Tool Name',
+      eventType: MCPAnalyticsEventType.mcpInitialize,
+      parameters: buildCapturedMcpParameters(request),
+      timestamp: new Date(),
+    }
+
+    await applyResolvedMetadata(event, data, request, extra)
+
+    const result = await originalInitializeHandler(request, extra)
+    event.response = result
+    captureEvent(server, event)
+    return result
+  })
+}

package/src/extensions/intent.ts

@@ -0,0 +1,78 @@
+import type {
+  CompatibleRequestHandlerExtra,
+  MCPAnalyticsData,
+  MCPAnalyticsIntentSource,
+  MCPRequestLike,
+  McpEvent,
+} from '../types'
+import { isContextEnabled } from './context-parameters'
+import { log } from './logger'
+
+interface ResolvedIntent {
+  intent: string
+  source: MCPAnalyticsIntentSource
+}
+
+function getContextArgument(request: MCPRequestLike): string | undefined {
+  const context = request.params?.arguments?.context
+  return typeof context === 'string' && context.trim() ? context : undefined
+}
+
+function normalizeIntent(intent: string | null | undefined): string | null {
+  if (typeof intent !== 'string') {
+    return null
+  }
+
+  const trimmed = intent.trim()
+  return trimmed ? trimmed : null
+}
+
+async function runIntentFallback(
+  data: MCPAnalyticsData,
+  request: MCPRequestLike,
+  extra?: CompatibleRequestHandlerExtra
+): Promise<ResolvedIntent | null> {
+  if (!data.options.intentFallback) {
+    return null
+  }
+
+  try {
+    const intent = normalizeIntent(await data.options.intentFallback(request, extra))
+    return intent ? { intent, source: 'inferred' } : null
+  } catch (error) {
+    log(`intentFallback callback error: ${error}`)
+    return null
+  }
+}
+
+export async function resolveToolCallIntent(
+  data: MCPAnalyticsData,
+  request: MCPRequestLike,
+  extra?: CompatibleRequestHandlerExtra
+): Promise<ResolvedIntent | null> {
+  const contextArgument = getContextArgument(request)
+  if (isContextEnabled(data.options.context) && request.params?.name !== 'get_more_tools' && contextArgument) {
+    return { intent: contextArgument, source: 'context_parameter' }
+  }
+
+  return await runIntentFallback(data, request, extra)
+}
+
+export function setEventIntent(event: McpEvent, resolvedIntent: ResolvedIntent | null): void {
+  if (!resolvedIntent) {
+    return
+  }
+
+  event.userIntent = resolvedIntent.intent
+  event.userIntentSource = resolvedIntent.source
+}
+
+export function setExplicitContextIntent(event: McpEvent, context: string): void {
+  const intent = normalizeIntent(context)
+  if (!intent) {
+    return
+  }
+
+  event.userIntent = intent
+  event.userIntentSource = 'context_parameter'
+}

package/src/extensions/internal.ts

@@ -0,0 +1,193 @@
+import type {
+  CompatibleRequestHandlerExtra,
+  MCPAnalyticsData,
+  MCPRequestLike,
+  MCPServerLike,
+  McpEvent,
+  UserIdentity,
+} from '../types'
+import { MCPAnalyticsEventType } from './event-types'
+import { log } from './logger'
+import { captureEvent } from './capture'
+
+/**
+ * Bounded LRU cache for session identities, capped at `maxSize` entries so a
+ * long-lived server can't accumulate identities for unboundedly many sessions.
+ * One instance lives on each server's tracking data — it is NOT shared across
+ * server instances, so identities never bleed between servers.
+ */
+export class IdentityCache {
+  private readonly _cache = new Map<string, UserIdentity>()
+  private readonly _maxSize: number
+
+  constructor(maxSize = 1000) {
+    this._maxSize = maxSize
+  }
+
+  get(sessionId: string): UserIdentity | undefined {
+    const identity = this._cache.get(sessionId)
+    if (identity === undefined) {
+      return
+    }
+    // Touch: re-insert so it counts as most-recently-used.
+    this._cache.delete(sessionId)
+    this._cache.set(sessionId, identity)
+    return identity
+  }
+
+  set(sessionId: string, identity: UserIdentity): void {
+    this._cache.delete(sessionId)
+
+    if (this._cache.size >= this._maxSize) {
+      const oldestKey = this._cache.keys().next().value
+      if (oldestKey !== undefined) {
+        this._cache.delete(oldestKey)
+      }
+    }
+
+    this._cache.set(sessionId, identity)
+  }
+
+  has(sessionId: string): boolean {
+    return this._cache.has(sessionId)
+  }
+
+  size(): number {
+    return this._cache.size
+  }
+}
+
+const _serverTracking = new WeakMap<MCPServerLike, MCPAnalyticsData>()
+
+export function getServerTrackingData(server: MCPServerLike): MCPAnalyticsData | undefined {
+  return _serverTracking.get(server)
+}
+
+export function setServerTrackingData(server: MCPServerLike, data: MCPAnalyticsData): void {
+  _serverTracking.set(server, data)
+}
+
+export function areIdentitiesEqual(a: UserIdentity, b: UserIdentity): boolean {
+  if (a.distinctId !== b.distinctId) {
+    return false
+  }
+  if (JSON.stringify(a.groups ?? {}) !== JSON.stringify(b.groups ?? {})) {
+    return false
+  }
+
+  const aProps = a.properties || {}
+  const bProps = b.properties || {}
+
+  const aKeys = Object.keys(aProps)
+  const bKeys = Object.keys(bProps)
+
+  if (aKeys.length !== bKeys.length) {
+    return false
+  }
+
+  for (const key of aKeys) {
+    if (!(key in bProps)) {
+      return false
+    }
+    if (JSON.stringify(aProps[key]) !== JSON.stringify(bProps[key])) {
+      return false
+    }
+  }
+
+  return true
+}
+
+export function mergeIdentities(previous: UserIdentity | undefined, next: UserIdentity): UserIdentity {
+  if (!previous) {
+    return next
+  }
+
+  return {
+    distinctId: next.distinctId,
+    properties: {
+      ...(previous.properties || {}),
+      ...(next.properties || {}),
+    },
+    groups: next.groups ?? previous.groups,
+  }
+}
+
+/**
+ * Resolves the optional `identify` callback, dedupes against the server's identity
+ * cache, and publishes an `$identify` event only when the identity has materially changed.
+ */
+export async function handleIdentify(
+  server: MCPServerLike,
+  data: MCPAnalyticsData,
+  sessionId: string,
+  request: MCPRequestLike,
+  extra?: CompatibleRequestHandlerExtra
+): Promise<void> {
+  if (!data.options.identify) {
+    return
+  }
+
+  const identifyEvent: McpEvent = {
+    sessionId,
+    resourceName: getRequestResourceName(request),
+    eventType: MCPAnalyticsEventType.identify,
+    parameters: { request, extra },
+    timestamp: new Date(),
+  }
+
+  try {
+    const identityResult =
+      typeof data.options.identify === 'function' ? await data.options.identify(request, extra) : data.options.identify
+
+    if (identityResult) {
+      const previousIdentity = data.identifiedSessions.get(sessionId)
+      const mergedIdentity = mergeIdentities(previousIdentity, identityResult)
+      const hasChanged = !(previousIdentity && areIdentitiesEqual(previousIdentity, mergedIdentity))
+
+      data.identifiedSessions.set(sessionId, mergedIdentity)
+
+      if (hasChanged) {
+        log(`Identified session ${sessionId} with identity: ${JSON.stringify(mergedIdentity)}`)
+        captureEvent(server, identifyEvent)
+      }
+    } else {
+      log(`Warning: Supplied identify function returned null for session ${sessionId}`)
+    }
+  } catch (error) {
+    log(`Error: User supplied identify function threw an error while identifying session ${sessionId} - ${error}`)
+  }
+}
+
+/**
+ * Resolves the `eventProperties` callback. Returns null when no callback is
+ * configured, or when it returns nullish or throws — a throw never propagates
+ * into the tool path.
+ */
+export async function resolveEventProperties(
+  data: MCPAnalyticsData,
+  request: MCPRequestLike,
+  extra?: CompatibleRequestHandlerExtra
+): Promise<Record<string, unknown> | null> {
+  if (!data.options.eventProperties) {
+    return null
+  }
+  try {
+    return (await data.options.eventProperties(request, extra)) ?? null
+  } catch (e) {
+    log(`eventProperties callback error: ${e}`)
+    return null
+  }
+}
+
+function getRequestResourceName(request: unknown): string {
+  if (!request || typeof request !== 'object' || !('params' in request)) {
+    return 'Unknown'
+  }
+
+  const params = request.params
+  if (!params || typeof params !== 'object' || !('name' in params)) {
+    return 'Unknown'
+  }
+
+  return typeof params.name === 'string' ? params.name : 'Unknown'
+}

package/src/extensions/logger.ts

@@ -0,0 +1,27 @@
+/**
+ * STDIO-safe logger.
+ *
+ * MCP servers running over the STDIO transport use stdout/stderr to exchange
+ * protocol messages, so we cannot use the default `console.log`. We accept a
+ * `logger` option on the public API; when omitted, log calls are silently
+ * dropped. Errors that affect tracking should still be observable in apps that
+ * want them, so the consumer can plug in any function (e.g. `console.error`
+ * for non-STDIO transports, a file logger, etc.).
+ */
+export type LoggerFn = (message: string) => void
+
+let activeLogger: LoggerFn | undefined
+
+export function setLogger(logger: LoggerFn | undefined): void {
+  activeLogger = logger
+}
+
+export function log(message: string): void {
+  if (activeLogger) {
+    try {
+      activeLogger(message)
+    } catch {
+      // never let logging blow up the tracking pipeline
+    }
+  }
+}