From c6f3634fb3db61126af083cb97b42e70e5d4fcc5 Mon Sep 17 00:00:00 2001 From: Andrey Avtomonov Date: Fri, 22 May 2026 16:01:13 +0200 Subject: [PATCH] feat: emit sampled mcp telemetry --- packages/cli/src/context/mcp/context-tools.ts | 57 ++++++++++++++++++- packages/cli/src/context/mcp/server.test.ts | 51 ++++++++++++++++- packages/cli/src/context/mcp/server.ts | 4 ++ packages/cli/src/context/mcp/types.ts | 3 + packages/cli/src/mcp-server-factory.ts | 2 + packages/cli/src/telemetry/index.ts | 11 ++++ 6 files changed, 126 insertions(+), 2 deletions(-) diff --git a/packages/cli/src/context/mcp/context-tools.ts b/packages/cli/src/context/mcp/context-tools.ts index 6778bb64..963ab44f 100644 --- a/packages/cli/src/context/mcp/context-tools.ts +++ b/packages/cli/src/context/mcp/context-tools.ts @@ -1,7 +1,10 @@ import { randomUUID } from 'node:crypto'; import type { ToolAnnotations } from '@modelcontextprotocol/sdk/types.js'; import { z } from 'zod'; +import type { KtxCliIo } from '../../cli-runtime.js'; import type { MemoryAgentInput } from '../../context/memory/types.js'; +import { emitTelemetryEvent, mcpTelemetrySampleRate, shouldEmitMcpTelemetry } from '../../telemetry/index.js'; +import { scrubErrorClass } from '../../telemetry/scrubber.js'; import type { KtxMcpContextPorts, KtxMcpProgressCallback, @@ -16,6 +19,8 @@ export interface RegisterKtxContextToolsDeps { server: KtxMcpServerLike; ports: KtxMcpContextPorts; userContext: KtxMcpUserContext; + projectDir?: string; + io?: KtxCliIo; } const connectionIdSchema = z.string().min(1); @@ -509,8 +514,58 @@ function registerParsedTool( }); } +function instrumentMcpServer( + server: KtxMcpServerLike, + telemetry: { projectDir?: string; io?: KtxCliIo }, +): KtxMcpServerLike { + return { + registerTool(name, config, handler) { + server.registerTool(name, config, async (input, context) => { + const startedAt = performance.now(); + try { + const result = await handler(input, context); + if (telemetry.io && telemetry.projectDir && shouldEmitMcpTelemetry()) { + const isError = + typeof result === 'object' && result !== null && 'isError' in result && result.isError === true; + await emitTelemetryEvent({ + name: 'mcp_request_completed', + projectDir: telemetry.projectDir, + io: telemetry.io, + fields: { + toolName: name, + outcome: isError ? 'error' : 'ok', + durationMs: Math.max(0, performance.now() - startedAt), + sampleRate: mcpTelemetrySampleRate(), + }, + }); + } + return result; + } catch (error) { + if (telemetry.io && telemetry.projectDir && shouldEmitMcpTelemetry()) { + const errorClass = scrubErrorClass(error); + await emitTelemetryEvent({ + name: 'mcp_request_completed', + projectDir: telemetry.projectDir, + io: telemetry.io, + fields: { + toolName: name, + outcome: 'error', + ...(errorClass ? { errorClass } : {}), + durationMs: Math.max(0, performance.now() - startedAt), + sampleRate: mcpTelemetrySampleRate(), + }, + }); + } + throw error; + } + }); + }, + }; +} + export function registerKtxContextTools(deps: RegisterKtxContextToolsDeps): void { - const { ports, server, userContext } = deps; + const { ports, userContext } = deps; + const server = instrumentMcpServer(deps.server, { projectDir: deps.projectDir, io: deps.io }); if (ports.connections) { const connections = ports.connections; diff --git a/packages/cli/src/context/mcp/server.test.ts b/packages/cli/src/context/mcp/server.test.ts index 3532a327..bee00c00 100644 --- a/packages/cli/src/context/mcp/server.test.ts +++ b/packages/cli/src/context/mcp/server.test.ts @@ -3,7 +3,7 @@ import { tmpdir } from 'node:os'; import { join } from 'node:path'; import { Client } from '@modelcontextprotocol/sdk/client/index.js'; import { InMemoryTransport } from '@modelcontextprotocol/sdk/inMemory.js'; -import { describe, expect, it, vi } from 'vitest'; +import { afterEach, describe, expect, it, vi } from 'vitest'; import { createLocalProjectMemoryIngest } from '../../context/memory/local-memory.js'; import { detectCaptureSignals } from '../../context/memory/capture-signals.js'; import type { MemoryAgentInput } from '../../context/memory/types.js'; @@ -47,6 +47,19 @@ function makeFakeServer() { }; } +function makeIo() { + let stderr = ''; + return { + stdout: { isTTY: true, write() {} }, + stderr: { + write(chunk: string) { + stderr += chunk; + }, + }, + stderrText: () => stderr, + }; +} + function getTool(tools: RegisteredTool[], name: string): RegisteredTool { const found = tools.find((tool) => tool.name === name); if (!found) { @@ -153,6 +166,11 @@ async function listToolsThroughSdk(contextTools: KtxMcpContextPorts) { } describe('createKtxMcpServer', () => { + afterEach(() => { + vi.unstubAllEnvs(); + vi.restoreAllMocks(); + }); + it('registers annotations and output schemas for every retained tool', async () => { const fake = makeFakeServer(); createKtxMcpServer({ @@ -227,6 +245,37 @@ describe('createKtxMcpServer', () => { }); }); + it('emits sampled debug telemetry for MCP tool requests', async () => { + vi.spyOn(Math, 'random').mockReturnValue(0); + vi.stubEnv('KTX_TELEMETRY_DEBUG', '1'); + vi.stubEnv('CI', ''); + const fake = makeFakeServer(); + const io = makeIo(); + const projectDir = '/tmp/ktx-mcp-telemetry'; + + createKtxMcpServer({ + server: fake.server, + userContext: { userId: 'local-user' }, + projectDir, + io, + contextTools: { + knowledge: { + search: vi.fn().mockResolvedValue({ results: [], totalFound: 0 }), + read: vi.fn().mockResolvedValue(null), + }, + }, + }); + + await expect(getTool(fake.tools, 'wiki_search').handler({ query: 'revenue recognition', limit: 5 })).resolves.toMatchObject({ + structuredContent: { results: [], totalFound: 0 }, + }); + + expect(io.stderrText()).toContain('"event":"mcp_request_completed"'); + expect(io.stderrText()).toContain('"toolName":"wiki_search"'); + expect(io.stderrText()).toContain('"sampleRate":0.1'); + expect(io.stderrText()).not.toContain(projectDir); + }); + it('registers parser-gated sql_execution when the host provides a SQL execution port', async () => { const fake = makeFakeServer(); const response: KtxSqlExecutionResponse = { diff --git a/packages/cli/src/context/mcp/server.ts b/packages/cli/src/context/mcp/server.ts index 73f970d6..97d79525 100644 --- a/packages/cli/src/context/mcp/server.ts +++ b/packages/cli/src/context/mcp/server.ts @@ -9,6 +9,8 @@ export function createKtxMcpServer(deps: KtxMcpServerDeps): KtxMcpServerDeps['se server: deps.server, ports: deps.contextTools, userContext: deps.userContext, + projectDir: deps.projectDir, + io: deps.io, }); } @@ -26,6 +28,8 @@ export function createDefaultKtxMcpServer( server: server as KtxMcpServerLike, userContext: deps.userContext, contextTools: deps.contextTools, + projectDir: deps.projectDir, + io: deps.io, }); return server; } diff --git a/packages/cli/src/context/mcp/types.ts b/packages/cli/src/context/mcp/types.ts index 223aa394..29a8c069 100644 --- a/packages/cli/src/context/mcp/types.ts +++ b/packages/cli/src/context/mcp/types.ts @@ -1,4 +1,5 @@ import type { MemoryIngestService } from '../../context/memory/memory-runs.js'; +import type { KtxCliIo } from '../../cli-runtime.js'; import type { KtxEntityDetailsInput, KtxEntityDetailsResponse } from '../scan/entity-details.js'; import type { KtxDiscoverDataInput, KtxDiscoverDataResponse } from '../../context/search/discover.js'; import type { KtxDictionarySearchInput, KtxDictionarySearchResponse } from '../../context/sl/dictionary-search.js'; @@ -171,4 +172,6 @@ export interface KtxMcpServerDeps { server: KtxMcpServerLike; userContext: KtxMcpUserContext; contextTools?: KtxMcpContextPorts; + projectDir?: string; + io?: KtxCliIo; } diff --git a/packages/cli/src/mcp-server-factory.ts b/packages/cli/src/mcp-server-factory.ts index e6d4887f..1ff44270 100644 --- a/packages/cli/src/mcp-server-factory.ts +++ b/packages/cli/src/mcp-server-factory.ts @@ -73,6 +73,8 @@ export async function createKtxMcpServerFactory(input: { name: 'ktx', version: input.cliVersion, userContext: { userId: 'local' }, + projectDir: input.projectDir, + io, contextTools: { ...contextTools, ...(memoryIngest ? { memoryIngest } : {}), diff --git a/packages/cli/src/telemetry/index.ts b/packages/cli/src/telemetry/index.ts index 2e6e2cf9..f026bcae 100644 --- a/packages/cli/src/telemetry/index.ts +++ b/packages/cli/src/telemetry/index.ts @@ -26,6 +26,17 @@ type TelemetryEventFields = Omit< >; const emittedProjectSnapshots = new Set(); +const MCP_SAMPLE_RATE = 0.1 as const; +let mcpSampled: boolean | undefined; + +export function shouldEmitMcpTelemetry(): boolean { + mcpSampled ??= Math.random() < MCP_SAMPLE_RATE; + return mcpSampled; +} + +export function mcpTelemetrySampleRate(): 0.1 { + return MCP_SAMPLE_RATE; +} async function emitInstallFirstRunIfNeeded(input: { identity: Awaited>;