diff --git a/.changeset/swift-hoops-yell.md b/.changeset/swift-hoops-yell.md new file mode 100644 index 0000000000..34b3276f1b --- /dev/null +++ b/.changeset/swift-hoops-yell.md @@ -0,0 +1,5 @@ +--- +'@posthog/mcp': minor +--- + +Record safe tool input field names for automatic and custom MCP servers. Unknown names are `[redacted]` by default; `shouldRecordInputKey` replaces that rule. diff --git a/packages/mcp/docs/ARCHITECTURE.md b/packages/mcp/docs/ARCHITECTURE.md index 575dcdf47b..e0040984a2 100644 --- a/packages/mcp/docs/ARCHITECTURE.md +++ b/packages/mcp/docs/ARCHITECTURE.md @@ -74,6 +74,60 @@ The pipeline lives in an exported `processMcpEvent()` function in `src/extension 4. **`beforeSend`** — each fully-built PostHog payload (`{ event, distinct_id, properties }`) is passed through `options.beforeSend(event)` (sync or async) right before dispatch — so it runs **once per emitted event**, including the `$exception` sibling. Returning the (possibly mutated) payload sends it; returning a nullish value drops it; a throw drops that event (and is logged). This is the seam for customer redaction or property tweaks. 5. **Dispatch** — each surviving event is handed to the user's `posthog-node` client via `posthog.capture()`. Batching, retries, and flushing are owned by that client. The host calls `posthog.shutdown()` to drain — the SDK installs no process-signal handlers and owns no client lifecycle. +### Tool input field names + +Automatic tool-call events include `$mcp_input_keys` on success and failure. +The SDK reads the original arguments before validation can remove unknown fields. +It records up to 20 top-level field names, sorted, without their values. +By default, only names declared by the server's input schema remain visible. +Unknown names and names longer than 64 characters are replaced by one `[redacted]` entry, the same marker the SDK uses for other hidden data. +Declared names come first, so `[redacted]` appears only when the 20-name limit leaves space. + +The `shouldRecordInputKey(key, { declared })` option replaces the default rule, for automatic capture and as the helper's third argument. +Return `true` to record a name; any other result, or a throw, records `[redacted]`. +The 64-character limit, the 20-name limit, and declared-names-first ordering still apply. +Use it when your server can accept that a caller-chosen name reaches analytics, for example to see misspelled parameter names: + +```ts +instrument(server, posthog, { + shouldRecordInputKey: (key, { declared }) => declared || /^[A-Za-z0-9_.-]+$/.test(key), +}) +``` +SDK argument names (`context`, `llm_model`, and `conversation_id`) are omitted unless the application schema declares them. +Non-object arguments do not produce this property. + +High-level servers use the registered tool's schema. +Low-level servers use schemas from prior `tools/list` responses on the same server instance. +Before a listing, or when a schema cannot be inspected, every name is hidden behind `[redacted]`. +The helper supports top-level JSON Schema properties, Zod raw shapes, and Zod object schemas, including objects wrapped by refinements, transforms, preprocessors, pipes, and optional, nullable, default, catch, or readonly wrappers. +A pipe reports the names of its input schema. +It does not resolve JSON Schema references or inspect fields inside unions. + +Custom dispatchers use the same helper through the existing `properties` argument: + +```ts +import { getToolInputProperties, PostHogMCP } from '@posthog/mcp' + +const posthog = new PostHogMCP(process.env.POSTHOG_PROJECT_TOKEN) +await posthog.register({ $mcp_server_build: 'example-build' }) + +const properties = getToolInputProperties(rawArguments, originalTool.inputSchema) +posthog.captureToolCall({ toolName, isError: false, properties }) +``` + +Compute these properties before argument normalization, and include them in both success and error events. +Pass a schema owned by the server, never one supplied by the caller. +Custom command formats must extract the actual tool arguments and schema before calling the helper. +Alternative field names must appear in the supplied schema, or pass `shouldRecordInputKey`, to remain visible. +Do not report which alternative names a call used through server-specific `$mcp_*` properties. +Alias telemetry is planned SDK follow-up work: the server will pass its own alias map to the helper, and the helper will add `$mcp_input_aliases_used` (for example `["experimentId:id"]`) without exposing unknown names. +The SDK does not normalize arguments or infer which alternative a server accepted. + +The helper adds no request values to the event. +Existing parameter and response capture remains unchanged. +Use `beforeSend` to remove `$mcp_input_keys` when needed (`before_send` on the underlying PostHog client). +No session store or additional network request is required. + ## 4. Session & identity ### Shared event properties diff --git a/packages/mcp/src/__tests__/beforeSend.test.ts b/packages/mcp/src/__tests__/beforeSend.test.ts index a9d853d7fd..5c7dc9d4d9 100644 --- a/packages/mcp/src/__tests__/beforeSend.test.ts +++ b/packages/mcp/src/__tests__/beforeSend.test.ts @@ -60,6 +60,7 @@ describe('beforeSend option', () => { beforeSend: (event) => { if (event.properties.$mcp_parameters) { event.properties.$mcp_parameters = '[redacted]' + delete event.properties.$mcp_input_keys } return event }, @@ -70,6 +71,25 @@ describe('beforeSend option', () => { const toolCall = capture.findCapturesByEvent('$mcp_tool_call')[0] expect(toolCall.properties.$mcp_parameters).toBe('[redacted]') + expect(toolCall.properties).not.toHaveProperty('$mcp_input_keys') + }) + + it('lets shouldRecordInputKey record undeclared argument names', async () => { + instrument(server, fakePostHog(), { + shouldRecordInputKey: (key, { declared }) => declared || /^[A-Za-z_]+$/.test(key), + }) + + await client.request( + { + method: 'tools/call', + params: { name: 'add_todo', arguments: { text: 'secret-value', todoText: 'x', 'person@example.com': 1 } }, + }, + CallToolResultSchema + ) + await new Promise((r) => setTimeout(r, 50)) + + const toolCall = capture.findCapturesByEvent('$mcp_tool_call')[0] + expect(toolCall.properties.$mcp_input_keys).toEqual(['text', 'todoText', '[redacted]']) }) it('drops an event when beforeSend returns null', async () => { diff --git a/packages/mcp/src/__tests__/concurrent-attribution.test.ts b/packages/mcp/src/__tests__/concurrent-attribution.test.ts index dcd8c79942..18eeac4fc7 100644 --- a/packages/mcp/src/__tests__/concurrent-attribution.test.ts +++ b/packages/mcp/src/__tests__/concurrent-attribution.test.ts @@ -655,4 +655,47 @@ describe('concurrent request attribution', () => { $mcp_is_error: true, }) }) + + it('keeps listed input schemas scoped to the request session', async () => { + const listAStarted = deferred() + const releaseListA = deferred() + const server = createServer({}) + let listIndex = 0 + server.setRequestHandler(ListToolsRequestSchema, async () => { + listIndex += 1 + if (listIndex === 1) { + listAStarted.resolve() + await releaseListA.promise + return { + tools: [{ name: 'echo', inputSchema: { type: 'object', properties: { labelA: { type: 'string' } } } }], + } + } + return { + tools: [{ name: 'echo', inputSchema: { type: 'object', properties: { requestLabel: { type: 'string' } } } }], + } + }) + const tokenA = encodeSessionId({ sessionId: 'ses_a' }) + const tokenB = encodeSessionId({ sessionId: 'ses_b' }) + + const listA = invokeListTools(server, { + requestInfo: { headers: { 'mcp-session-id': tokenA } }, + }) + await listAStarted.promise + await invokeListTools(server, { + requestInfo: { headers: { 'mcp-session-id': tokenB } }, + }) + releaseListA.resolve() + await listA + + await invokeTool(server, 'B', { + requestInfo: { headers: { 'mcp-session-id': tokenB } }, + }) + await flushCaptures() + + const toolCall = capture.findCapturesByEvent('$mcp_tool_call')[0] + expect(toolCall.properties).toMatchObject({ + $session_id: 'ses_b', + $mcp_input_keys: ['requestLabel'], + }) + }) }) diff --git a/packages/mcp/src/__tests__/error-capture.test.ts b/packages/mcp/src/__tests__/error-capture.test.ts index 0f278d2791..8b74ca1cd4 100644 --- a/packages/mcp/src/__tests__/error-capture.test.ts +++ b/packages/mcp/src/__tests__/error-capture.test.ts @@ -132,7 +132,7 @@ describe('error capture on the tool-call path', () => { ) }, 'calculate', - { op: 'modulo', a: 10, b: 3, context: 'test' }, + { op: 'modulo', a: 10, b: 3, context: 'test', private_identifier: true }, ], [ 'missing required parameter', @@ -154,6 +154,7 @@ describe('error capture on the tool-call path', () => { await new Promise((r) => setTimeout(r, 50)) const event = capture.findEventsByResourceName(toolName).find((e) => e.isError) const exception = event?.error?.$exception_list?.[0] + expect(event?.properties?.$mcp_input_keys).toEqual(toolName === 'calculate' ? ['a', 'b', 'op', '[redacted]'] : []) expect(exception?.value).toMatch(/Invalid|required/i) expect(['McpError', 'Error', undefined]).toContain(exception?.type) }) diff --git a/packages/mcp/src/__tests__/identity-cache.test.ts b/packages/mcp/src/__tests__/identity-cache.test.ts index 9b6d98cccb..47dd10aef8 100644 --- a/packages/mcp/src/__tests__/identity-cache.test.ts +++ b/packages/mcp/src/__tests__/identity-cache.test.ts @@ -1,4 +1,4 @@ -import { IdentityCache } from '../extensions/internal' +import { BoundedCache, IdentityCache } from '../extensions/internal' describe('IdentityCache', () => { it('stores and retrieves identities by session id', () => { @@ -31,3 +31,17 @@ describe('IdentityCache', () => { expect(serverB.get('ses_shared')).toBeUndefined() }) }) + +describe('BoundedCache', () => { + it('evicts the least-recently-used entry', () => { + const cache = new BoundedCache(2) + cache.set('one', 1) + cache.set('two', 2) + expect(cache.get('one')).toBe(1) + cache.set('three', 3) + + expect(cache.get('two')).toBeUndefined() + expect(cache.get('one')).toBe(1) + expect(cache.get('three')).toBe(3) + }) +}) diff --git a/packages/mcp/src/__tests__/instrument-lowlevel.test.ts b/packages/mcp/src/__tests__/instrument-lowlevel.test.ts index 6b3fc99735..74147bb94a 100644 --- a/packages/mcp/src/__tests__/instrument-lowlevel.test.ts +++ b/packages/mcp/src/__tests__/instrument-lowlevel.test.ts @@ -455,17 +455,21 @@ describe('Low-level Server tracing (e2e)', () => { await eventCapture.stop() }) - it('captures a single $mcp_tool_call for a successful call', async () => { + it.each([false, true])('captures safe input names with a prior listing: %s', async (listed) => { const { server, client, connect, cleanup } = await setupLowLevelServer() try { instrument(server, fakePostHog()) await connect() + if (listed) await client.request({ method: 'tools/list', params: {} }, ListToolsResultSchema) const result = await client.request( - { method: 'tools/call', params: { name: 'echo', arguments: { text: 'hi' } } }, + { + method: 'tools/call', + params: { name: 'echo', arguments: { text: 'hi', private_identifier: true, context: 'example' } }, + }, CallToolResultSchema ) - await new Promise((r) => setTimeout(r, 50)) + await vi.waitFor(() => expect(eventCapture.findCapturesByEvent('$mcp_tool_call')).toHaveLength(1)) expect((result.content as { text: string }[])[0].text).toBe('echo: hi') @@ -473,6 +477,7 @@ describe('Low-level Server tracing (e2e)', () => { expect(toolCalls).toHaveLength(1) const props = toolCalls[0].properties expect(props.$mcp_tool_name).toBe('echo') + expect(props.$mcp_input_keys).toEqual(listed ? ['text', '[redacted]'] : ['[redacted]']) expect(props.$mcp_resource_name).toBe('echo') expect(props.$mcp_is_error).toBe(false) expect(props.$mcp_duration_ms).toEqual(expect.any(Number)) diff --git a/packages/mcp/src/__tests__/posthog-mcp.test.ts b/packages/mcp/src/__tests__/posthog-mcp.test.ts index 29a49fdbe8..9b76d95147 100644 --- a/packages/mcp/src/__tests__/posthog-mcp.test.ts +++ b/packages/mcp/src/__tests__/posthog-mcp.test.ts @@ -1,4 +1,4 @@ -import { getMoreToolsResult, PostHogMCP } from '../index' +import { getMoreToolsResult, getToolInputProperties, PostHogMCP } from '../index' import { PostHogMCPAnalyticsEvent, PostHogMCPAnalyticsProperty } from '../extensions/constants' import { GET_MORE_TOOLS_NAME } from '../extensions/tools' import type { PostHogCaptureEvent } from '../extensions/posthog-events' @@ -53,7 +53,14 @@ describe('PostHogMCP', () => { distinctId: 'user-123', sessionId: 'session-abc', groups: { organization: 'org-1', project: 'proj-1' }, - properties: { $mcp_client_name: 'claude-code', custom_flag: true }, + properties: { + $mcp_client_name: 'claude-code', + custom_flag: true, + ...getToolInputProperties( + { query: 'example-value', private_identifier: true }, + { properties: { query: {} } } + ), + }, }) await tick() @@ -70,6 +77,9 @@ describe('PostHogMCP', () => { expect(p.$groups).toEqual({ organization: 'org-1', project: 'proj-1' }) expect(p.$mcp_client_name).toBe('claude-code') expect(p.custom_flag).toBe(true) + expect(p.$mcp_input_keys).toEqual(['query', '[redacted]']) + expect(p).not.toHaveProperty('$mcp_parameters') + expect(JSON.stringify(p)).not.toContain('example-value') // A resolved identity keeps person processing on. expect(p.$process_person_profile).toBeUndefined() }) diff --git a/packages/mcp/src/__tests__/tool-input.test.ts b/packages/mcp/src/__tests__/tool-input.test.ts new file mode 100644 index 0000000000..af14b7fa2d --- /dev/null +++ b/packages/mcp/src/__tests__/tool-input.test.ts @@ -0,0 +1,116 @@ +import { z } from 'zod' +import { z as z4 } from 'zod4' +import { getToolInputProperties } from '../index' + +describe('getToolInputProperties', () => { + it.each([ + { + type: 'object', + properties: { id: { type: 'string' }, context: { type: 'string' }, properties: { type: 'string' } }, + }, + { id: z.string(), context: z.string(), properties: z.string() }, + z.object({ id: z.string(), context: z.string(), properties: z.string() }), + z4.object({ id: z4.string(), context: z4.string(), properties: z4.string() }), + z + .object({ id: z.string(), context: z.string(), properties: z.string() }) + .refine(() => true) + .transform((value) => value), + z.preprocess((value) => value, z.object({ id: z.string(), context: z.string(), properties: z.string() })), + z.object({ id: z.string(), context: z.string(), properties: z.string() }).pipe(z.any()), + z.object({ id: z.string(), context: z.string(), properties: z.string() }).optional(), + z4.object({ id: z4.string(), context: z4.string(), properties: z4.string() }).transform((value) => value), + z4.preprocess((value) => value, z4.object({ id: z4.string(), context: z4.string(), properties: z4.string() })), + z4.preprocess( + (value) => value, + z4.preprocess((value) => value, z4.object({ id: z4.string(), context: z4.string(), properties: z4.string() })) + ), + z4 + .object({ id: z4.string(), context: z4.string(), properties: z4.string() }) + .pipe(z4.any()) + .default({ id: '', context: '', properties: '' }), + ])('keeps declared names and masks unknown names with schema %j', (schema) => { + const input = { + id: 'private-value', + context: 'application-value', + properties: 'application-properties', + llm_model: 'example-model', + conversation_id: 'example-conversation', + private_identifier_123: true, + 'person@example.com': true, + } + expect(getToolInputProperties(input, schema)).toEqual({ + $mcp_input_keys: ['context', 'id', 'properties', '[redacted]'], + }) + expect(input.id).toBe('private-value') + }) + + it('lets the caller replace the default rule and keeps declared names first', () => { + const schema = { + properties: { id: {}, ...Object.fromEntries(Array.from({ length: 19 }, (_, i) => [`d${i}`, {}])) }, + } + const input = { ...schema.properties, aKey: 1, experimentId: 1, 'person@example.com': 1, ['x'.repeat(65)]: 1 } + const seen: Array<[string, boolean]> = [] + const keys = getToolInputProperties(input, schema, { + shouldRecordInputKey: (key, { declared }) => { + seen.push([key, declared]) + return /^[A-Za-z0-9_]+$/.test(key) + }, + }).$mcp_input_keys as string[] + expect(keys).toHaveLength(20) + expect(keys).toContain('id') + expect(keys).not.toContain('aKey') + expect(seen).toContainEqual(['experimentId', false]) + expect(seen).toContainEqual(['id', true]) + expect(seen.map(([key]) => key)).not.toContain('x'.repeat(65)) + + expect( + getToolInputProperties({ id: 1, experimentId: 1, other: 1 }, schema, { + shouldRecordInputKey: (key) => key !== 'id' && key !== 'other', + }) + ).toEqual({ $mcp_input_keys: ['experimentId', '[redacted]'] }) + }) + + it.each([ + () => { + throw new Error('boom') + }, + () => 'yes' as unknown as boolean, + ])('records [redacted] when shouldRecordInputKey throws or does not return true', (shouldRecordInputKey) => { + expect(getToolInputProperties({ id: 1 }, { properties: { id: {} } }, { shouldRecordInputKey })).toEqual({ + $mcp_input_keys: ['[redacted]'], + }) + }) + + it.each([undefined, null, 'invalid', ['id'], new Date()])('omits names for non-object arguments %j', (input) => { + expect(getToolInputProperties(input, { properties: { id: {} } })).toEqual({}) + }) + + it('masks names when no schema is available and does not read values', () => { + const input = Object.defineProperty({}, 'id', { + enumerable: true, + get() { + throw new Error('must not read values') + }, + }) + expect(getToolInputProperties(input)).toEqual({ $mcp_input_keys: ['[redacted]'] }) + expect(getToolInputProperties(input, { properties: { id: {} } })).toEqual({ $mcp_input_keys: ['id'] }) + }) + + it('bounds the names and drops malformed analytics without throwing', () => { + const properties = Object.fromEntries(Array.from({ length: 30 }, (_, index) => [`key${index}`, {}])) + expect(getToolInputProperties(properties, { properties }).$mcp_input_keys).toHaveLength(20) + const longName = 'x'.repeat(65) + expect(getToolInputProperties({ [longName]: 1 }, { properties: { [longName]: {} } })).toEqual({ + $mcp_input_keys: ['[redacted]'], + }) + const input = new Proxy( + {}, + { + ownKeys: () => { + throw new Error('unavailable') + }, + } + ) + expect(getToolInputProperties(input, { properties })).toEqual({}) + }) +}) diff --git a/packages/mcp/src/extensions/constants.ts b/packages/mcp/src/extensions/constants.ts index a497c3a8cb..379e2b9221 100644 --- a/packages/mcp/src/extensions/constants.ts +++ b/packages/mcp/src/extensions/constants.ts @@ -56,6 +56,7 @@ export const PostHogMCPAnalyticsProperty = { FeedbackTool: '$mcp_feedback_tool', FeedbackType: '$mcp_feedback_type', IsError: '$mcp_is_error', + InputKeys: '$mcp_input_keys', Intent: '$mcp_intent', IntentSource: '$mcp_intent_source', ListedToolNames: '$mcp_listed_tool_names', diff --git a/packages/mcp/src/extensions/instrument-highlevel.ts b/packages/mcp/src/extensions/instrument-highlevel.ts index dc11a74b6b..7cf9071501 100644 --- a/packages/mcp/src/extensions/instrument-highlevel.ts +++ b/packages/mcp/src/extensions/instrument-highlevel.ts @@ -298,6 +298,7 @@ async function handleToolCallRequest( parameterOwnership: registeredTool ? getAnalyticsParameterOwnership(registeredTool.inputSchema, registeredTool.outputSchema) : undefined, + inputSchema: registeredTool?.inputSchema, takeCapturedError: () => { const captured = extra?.__mcp_analytics_error if (extra) { diff --git a/packages/mcp/src/extensions/instrumentation.ts b/packages/mcp/src/extensions/instrumentation.ts index 5eb97dd9af..042f7fa4ca 100644 --- a/packages/mcp/src/extensions/instrumentation.ts +++ b/packages/mcp/src/extensions/instrumentation.ts @@ -43,11 +43,18 @@ import type { LoggerFn } from './logger' import { buildCapturedMcpParameters } from './mcp-payloads' import { readRequestHandlerMethod } from './mcp-sdk-compat' import { getRequestHeaders } from './request-headers' -import { getSessionId, getSessionInfo, isModernEraRequest, newSessionId } from './session' +import { + deriveSessionIdFromMCPSession, + getSessionId, + getSessionInfo, + isModernEraRequest, + newSessionId, +} from './session' import { decodeSessionId, encodeSessionId, readMcpSessionHeader, writeSessionIdToTransport } from './session-token' import { getFeedbackToolDescriptor, resolveCollectFeedbackOptions, SEND_FEEDBACK_TOOL_NAME } from './feedback' import { getReportMissingToolDescriptor, resolveMissingCapabilityToolName } from './tools' import { applyResolvedMetadata, isToolResultError } from './tracing-helpers' +import { getToolInputProperties } from './tool-input' /** * Single instrumentation core shared by the low-level (`Server`) and high-level @@ -62,6 +69,13 @@ type MCPRequestHandler = (request: MCPRequestLike, extra?: CompatibleRequestHand /** Runs the underlying tool with SDK-owned analytics arguments removed. */ type ToolExecutor = (downstreamRequest: MCPRequestLike) => Promise +function resolveToolSchemaSessionId(data: MCPAnalyticsData, extra?: CompatibleRequestHandlerExtra): string { + const token = decodeSessionId(readMcpSessionHeader(getRequestHeaders(extra))) + if (token) return token.sessionId + if (extra?.sessionId) return deriveSessionIdFromMCPSession(extra.sessionId) + return data.sessionId +} + interface TraceToolCallParams { server: MCPServerLike data: MCPAnalyticsData @@ -70,6 +84,7 @@ interface TraceToolCallParams { execute: ToolExecutor /** Optional schema-derived ownership override for adapters with direct registry access. */ parameterOwnership?: AnalyticsParameterOwnership + inputSchema?: unknown /** * Event type to capture. Defaults to a tool call; the `get_more_tools` virtual * tool passes `mcpMissingCapability` and `send_feedback` passes @@ -120,6 +135,7 @@ export async function captureToolCall(params: TraceToolCallParams): Promise() + for (const tool of tools) { + if (tool?.name) sessionSchemas.set(tool.name, tool.inputSchema) + } + data.toolInputSchemas.set(event.sessionId, sessionSchemas) + } } if (data && isContextEnabled(data.options.context)) { tools = addContextParameterToTools(tools, getContextDescription(data.options.context), data.logger) diff --git a/packages/mcp/src/extensions/internal.ts b/packages/mcp/src/extensions/internal.ts index 0bfa035430..e66147d61c 100644 --- a/packages/mcp/src/extensions/internal.ts +++ b/packages/mcp/src/extensions/internal.ts @@ -18,32 +18,31 @@ import { stampClientIdentity } from './client-identity' import { stampTransportIdentity } from './transport-identity' /** - * 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. + * Bounded LRU cache, capped at `maxSize` entries so a long-lived server cannot + * accumulate state for unboundedly many sessions. One instance lives on each + * server's tracking data and is never shared across server instances. */ -export class IdentityCache { - private readonly _cache = new Map() +export class BoundedCache { + private readonly _cache = new Map() private readonly _maxSize: number constructor(maxSize = 1000) { this._maxSize = maxSize } - get(sessionId: string): UserIdentity | undefined { - const identity = this._cache.get(sessionId) - if (identity === undefined) { + get(key: string): T | undefined { + const value = this._cache.get(key) + if (value === undefined) { return } // Touch: re-insert so it counts as most-recently-used. - this._cache.delete(sessionId) - this._cache.set(sessionId, identity) - return identity + this._cache.delete(key) + this._cache.set(key, value) + return value } - set(sessionId: string, identity: UserIdentity): void { - this._cache.delete(sessionId) + set(key: string, value: T): void { + this._cache.delete(key) if (this._cache.size >= this._maxSize) { const oldestKey = this._cache.keys().next().value @@ -52,11 +51,11 @@ export class IdentityCache { } } - this._cache.set(sessionId, identity) + this._cache.set(key, value) } - has(sessionId: string): boolean { - return this._cache.has(sessionId) + has(key: string): boolean { + return this._cache.has(key) } size(): number { @@ -64,6 +63,8 @@ export class IdentityCache { } } +export class IdentityCache extends BoundedCache {} + const _serverTracking = new WeakMap() export function getServerTrackingData(server: MCPServerLike): MCPAnalyticsData | undefined { diff --git a/packages/mcp/src/extensions/mcp-sdk-compat.ts b/packages/mcp/src/extensions/mcp-sdk-compat.ts index 05e30a4866..1c216dc542 100644 --- a/packages/mcp/src/extensions/mcp-sdk-compat.ts +++ b/packages/mcp/src/extensions/mcp-sdk-compat.ts @@ -109,7 +109,44 @@ function isZodTypeLike(value: unknown): boolean { } export function isZodRawShapeCompat(schema: unknown): schema is Record { - return !!schema && typeof schema === 'object' && Object.values(schema).some(isZodTypeLike) + // A Zod v4 pipe exposes its `in` and `out` schemas as own fields, so a Zod schema is never a raw shape + return !!schema && typeof schema === 'object' && !isZodTypeLike(schema) && Object.values(schema).some(isZodTypeLike) +} + +interface ZodWrapperDef { + type?: unknown + schema?: unknown + in?: unknown + out?: unknown + innerType?: unknown +} + +const MAX_UNWRAP_DEPTH = 8 + +function zodDef(schema: unknown): ZodWrapperDef | undefined { + return (isZ4Schema(schema) ? (schema as ZodV4Internal)._zod?.def : (schema as ZodV3Internal)._def) as + | ZodWrapperDef + | undefined +} + +/** + * Follows Zod wrappers to the schema that parses the caller's input: v3 effects + * (refine, transform, preprocess) and pipelines, v4 pipes (which include + * transforms), and optional, nullable, default, catch, and readonly wrappers. + * A v4 `z.preprocess` is a pipe whose input side is the transform, so its output side is followed. + */ +export function unwrapInputSchema(schema: unknown): unknown { + let current = schema + for (let depth = 0; depth < MAX_UNWRAP_DEPTH && isZodTypeLike(current); depth++) { + const def = zodDef(current) + const pipeInput = def?.in !== undefined && zodDef(def.in)?.type === 'transform' ? def.out : def?.in + const inner = def?.schema ?? pipeInput ?? def?.innerType + if (!inner) { + break + } + current = inner + } + return current } export function getObjectShape(schema: unknown): Record | undefined { diff --git a/packages/mcp/src/extensions/tool-input.ts b/packages/mcp/src/extensions/tool-input.ts new file mode 100644 index 0000000000..0fdea21de8 --- /dev/null +++ b/packages/mcp/src/extensions/tool-input.ts @@ -0,0 +1,64 @@ +import type { JsonRecord, ShouldRecordInputKeyFn, ToolInputOptions } from '../types' +import { PostHogMCPAnalyticsProperty } from './constants' +import { getObjectShape, isZodRawShapeCompat, unwrapInputSchema } from './mcp-sdk-compat' + +const MAX_INPUT_KEYS = 20 +const MAX_KEY_LENGTH = 64 +const ANALYTICS_KEYS = new Set(['context', 'llm_model', 'conversation_id']) + +function isRecord(value: unknown): value is Record { + return value !== null && typeof value === 'object' && !Array.isArray(value) +} + +function declaredProperties(schema: unknown): Record | undefined { + if (!isRecord(schema)) return undefined + const properties = schema.properties + if (isZodRawShapeCompat(schema)) return schema + return getObjectShape(unwrapInputSchema(schema)) ?? (isRecord(properties) ? properties : undefined) +} + +const recordDeclaredOnly: ShouldRecordInputKeyFn = (_key, { declared }) => declared + +function shouldRecord(fn: ShouldRecordInputKeyFn, key: string, declared: boolean): boolean { + try { + return fn(key, { declared }) === true + } catch { + return false + } +} + +/** + * Describe the original arguments without reading their values. + * Pass a server-owned JSON Schema or Zod object schema, never a schema from the caller. + * By default unknown names become `[redacted]` because an argument name can contain private data; + * `shouldRecordInputKey` replaces that rule. + */ +export function getToolInputProperties(input: unknown, inputSchema?: unknown, options?: ToolInputOptions): JsonRecord { + try { + if (!isRecord(input)) return {} + const prototype = Object.getPrototypeOf(input) + if (prototype !== null && prototype !== Object.prototype) return {} + const properties = declaredProperties(inputSchema) + const known = new Set(Object.keys(properties ?? {})) + const keys = Object.keys(input).filter((key) => known.has(key) || !ANALYTICS_KEYS.has(key)) + const record = options?.shouldRecordInputKey ?? recordDeclaredOnly + const declared: string[] = [] + const undeclared: string[] = [] + let hasRedacted = false + for (const key of keys) { + const isDeclared = known.has(key) + if (key.length <= MAX_KEY_LENGTH && shouldRecord(record, key, isDeclared)) { + ;(isDeclared ? declared : undeclared).push(key) + } else { + hasRedacted = true + } + } + const visibleKeys = [...declared.sort(), ...undeclared.sort()].slice(0, MAX_INPUT_KEYS) + if (hasRedacted && visibleKeys.length < MAX_INPUT_KEYS) { + visibleKeys.push('[redacted]') + } + return { [PostHogMCPAnalyticsProperty.InputKeys]: visibleKeys } + } catch { + return {} + } +} diff --git a/packages/mcp/src/index.ts b/packages/mcp/src/index.ts index 419a67064c..983bf2de1e 100644 --- a/packages/mcp/src/index.ts +++ b/packages/mcp/src/index.ts @@ -7,7 +7,7 @@ import type { PostHog } from 'posthog-node' import { isCompatibleServerType, isHighLevelServer } from './extensions/compatibility' import { McpEventSink } from './extensions/sink' import { MCPAnalyticsEventType } from './extensions/event-types' -import { IdentityCache, getServerTrackingData, setServerTrackingData } from './extensions/internal' +import { BoundedCache, IdentityCache, getServerTrackingData, setServerTrackingData } from './extensions/internal' import { createLogger } from './extensions/logger' import { captureEvent } from './extensions/capture' import { applyMcpLibIdentity } from './extensions/lib-identity' @@ -161,6 +161,7 @@ function buildTrackingData( toolAnalyticsParameterOwnership: new Map(), toolCategories: new Map(), toolDescriptions: new Map(), + toolInputSchemas: new BoundedCache(), sessionInfo: getSessionInfo(lowLevelServer, undefined), options: { ...DEFAULT_OPTIONS, @@ -230,6 +231,7 @@ export { // Host callbacks receive the SDK's `extra`/`ctx` unchanged, and the two SDK // majors carry HTTP headers in different places and shapes. This reads either. export { getRequestHeaders } from './extensions/request-headers' +export { getToolInputProperties } from './extensions/tool-input' export { PostHogMCP, type PostHogMCPOptions } from './extensions/posthog-mcp' export { getMoreToolsResult } from './extensions/tools' export { sendFeedbackResult, SEND_FEEDBACK_TOOL_NAME } from './extensions/feedback' @@ -261,7 +263,9 @@ export type { PrepareToolCallOptions, PrepareToolListOptions, RequestHeaderBag, + ShouldRecordInputKeyFn, ToolCallCaptureData, + ToolInputOptions, ToolsListCaptureData, UserIdentity, } from './types' diff --git a/packages/mcp/src/types.ts b/packages/mcp/src/types.ts index 7cf2b975ef..bfb18f84b4 100644 --- a/packages/mcp/src/types.ts +++ b/packages/mcp/src/types.ts @@ -6,7 +6,7 @@ import type { ErrorTracking } from '@posthog/core' import type { AnalyticsInjectableJsonSchema } from './extensions/analytics-parameters' import type { MCPAnalyticsEventType } from './extensions/event-types' -import type { IdentityCache } from './extensions/internal' +import type { BoundedCache, IdentityCache } from './extensions/internal' import type { PostHogCaptureEvent } from './extensions/posthog-events' import type { McpEventSink } from './extensions/sink' import type { LoggerFn } from './extensions/logger' @@ -186,6 +186,12 @@ export interface MCPAnalyticsOptions { * suppress specific events. A throw drops that event. */ beforeSend?: BeforeSendFn + /** + * Decide which argument names `$mcp_input_keys` records on tool-call events. + * By default only names the tool's input schema declares are recorded; every + * other name becomes one `[redacted]` entry, because a name can carry private data. + */ + shouldRecordInputKey?: ShouldRecordInputKeyFn /** * Attach extra event properties on every auto-captured event. Spread into the PostHog * event properties as-is; values must be JSON-serializable. @@ -304,6 +310,21 @@ export type RegisteredTool = { */ export type BeforeSendFn = (event: PostHogCaptureEvent) => MaybePromise +/** + * Decides whether one top-level argument name appears in `$mcp_input_keys`. + * `declared` is true when the server's input schema declares the name. + * Return `true` to record the name; any other result, or a throw, records `[redacted]`. + */ +export type ShouldRecordInputKeyFn = (key: string, details: { declared: boolean }) => boolean + +export interface ToolInputOptions { + /** + * Replace the default rule, which records only declared names. The SDK still + * drops names longer than 64 characters and records at most 20 names. + */ + shouldRecordInputKey?: ShouldRecordInputKeyFn +} + export interface Event { actorId?: string clientName?: string @@ -515,6 +536,7 @@ export interface MCPAnalyticsData { toolAnalyticsParameterOwnership: Map toolCategories: Map toolDescriptions: Map + toolInputSchemas: BoundedCache> } export interface CaptureEventData {