diff --git a/.changeset/mcp-source-attribution.md b/.changeset/mcp-source-attribution.md new file mode 100644 index 0000000000..cb53ec0c5d --- /dev/null +++ b/.changeset/mcp-source-attribution.md @@ -0,0 +1,6 @@ +--- +'@hyperdx/api': patch +--- + +feat: record which source each MCP tool call reads on the tool span, the tool +metrics (`source_kind`) and the query `log_comment` diff --git a/agent_docs/architecture.md b/agent_docs/architecture.md index 30dfce4528..6162f28efd 100644 --- a/agent_docs/architecture.md +++ b/agent_docs/architecture.md @@ -90,6 +90,11 @@ Key, rate-limited to 600 req/min. - **Response size**: each tool call records the UTF-8 bytes it returns (`mcp.tool.response.size_bytes` span attribute, `hyperdx.mcp.tool.response.size_bytes` histogram) +- **Source**: a tool that resolves a source calls `recordMcpSource` (in + `utils/tracing.ts`), which sets `mcp.source.id`, `mcp.source.kind` and + `mcp.source.ids` on the tool span and `source_kind` on the tool metrics + (`none` when no source applies). Wrap the queries that read the source in + `withQueryAttribution({ source })` so `log_comment` carries it too See `MCP.md` in the repo root for user-facing setup instructions. diff --git a/packages/api/src/mcp/__tests__/sourceAttribution.test.ts b/packages/api/src/mcp/__tests__/sourceAttribution.test.ts new file mode 100644 index 0000000000..a4e994eece --- /dev/null +++ b/packages/api/src/mcp/__tests__/sourceAttribution.test.ts @@ -0,0 +1,196 @@ +// Drives runConfigTile inside withToolTracing to check that the source it +// resolves reaches both the tool span and the log_comment of the query it runs. + +const mockToolSpan = { setAttribute: jest.fn(), setStatus: jest.fn() }; +const mockQueries: { clickhouse_settings?: Record }[] = []; +const mockGetSource = jest.fn(); + +jest.mock('@/utils/instrumentation', () => ({ + getActiveTraceId: () => 'trace-1', + getCounter: () => ({ add: jest.fn() }), + getHistogram: () => ({ record: jest.fn() }), + SpanStatusCode: { OK: 1, ERROR: 2 }, + withSpan: (_name: string, fn: (span: unknown) => unknown) => fn(mockToolSpan), +})); +jest.mock('@/controllers/user', () => ({ + recordOnboardingTaskCompletion: jest.fn(), +})); +jest.mock('@/models/source', () => ({})); +jest.mock('@/controllers/sources', () => ({ + getSource: (...args: unknown[]) => mockGetSource(...args), +})); +jest.mock('@/controllers/connection', () => ({ + getConnectionById: async () => ({ + host: 'http://localhost:8123', + username: 'default', + password: '', + }), +})); +// The real node client, so log_comment comes from its own attribution logic. +// queryChartConfig is replaced with a plain query that carries no per-query +// attribution, standing in for the metadata queries renderChartConfig issues: +// only the ambient scope can give those a source. +jest.mock('@/clickhouse', () => { + const { ClickhouseClient } = jest.requireActual( + '@hyperdx/common-utils/dist/clickhouse/node', + ); + class RecordingClient extends ClickhouseClient { + async queryChartConfig() { + await this.query({ query: 'SELECT 1', format: 'JSON' }); + return { data: [], meta: [], rows: 0 }; + } + + protected async __query(inputs: (typeof mockQueries)[number]) { + mockQueries.push(inputs); + return {}; + } + } + return { ClickhouseClient: RecordingClient }; +}); +jest.mock('@/utils/logger', () => ({ + __esModule: true, + default: { info: jest.fn(), warn: jest.fn(), error: jest.fn() }, +})); + +import { SourceKind } from '@hyperdx/common-utils/dist/types'; +import mongoose from 'mongoose'; + +import { buildTile, runConfigTile } from '@/mcp/tools/query/helpers'; +import { withToolTracing } from '@/mcp/utils/tracing'; + +const context = { teamId: 'team-1', userId: 'user-1' }; +const startDate = new Date('2025-01-01T00:00:00Z'); +const endDate = new Date('2025-01-01T01:00:00Z'); + +function fakeSource(kind: SourceKind) { + return { + _id: new mongoose.Types.ObjectId(), + kind, + name: `${kind} source`, + connection: new mongoose.Types.ObjectId(), + from: { databaseName: 'default', tableName: `otel_${kind}` }, + timestampValueExpression: 'Timestamp', + }; +} + +function logCommentOf(index: number) { + return JSON.parse( + String(mockQueries.at(index)?.clickhouse_settings?.log_comment), + ); +} + +function attributeValues(key: string) { + return mockToolSpan.setAttribute.mock.calls + .filter(([k]) => k === key) + .map(([, v]) => v); +} + +beforeEach(() => { + jest.clearAllMocks(); + mockQueries.length = 0; +}); + +describe('MCP source attribution', () => { + it('tags a builder tile query and the tool span with its source', async () => { + const source = fakeSource(SourceKind.Log); + mockGetSource.mockResolvedValue(source); + const tile = buildTile('t', 12, 4, { + displayType: 'table', + sourceId: source._id.toString(), + select: [{ aggFn: 'count' }], + }); + + const traced = withToolTracing('clickstack_table', context, () => + runConfigTile(context.teamId, tile, startDate, endDate), + ); + const result = await traced({}); + + expect(result.isError).toBeFalsy(); + expect(mockQueries).toHaveLength(1); + expect(logCommentOf(0)).toEqual({ + v: 1, + surface: 'mcp', + source: source._id.toString(), + trace: 'trace-1', + label: 'clickstack_table', + }); + expect(attributeValues('mcp.source.id')).toEqual([source._id.toString()]); + expect(attributeValues('mcp.source.kind')).toEqual(['log']); + }); + + it('tags a raw SQL tile query with its source', async () => { + const source = fakeSource(SourceKind.Trace); + mockGetSource.mockResolvedValue(source); + const tile = buildTile('t', 24, 6, { + configType: 'sql', + displayType: 'table', + connectionId: source.connection.toString(), + sourceId: source._id.toString(), + sqlTemplate: 'SELECT 1', + }); + + const traced = withToolTracing('clickstack_sql', context, () => + runConfigTile(context.teamId, tile, startDate, endDate), + ); + await traced({}); + + expect(logCommentOf(0)).toMatchObject({ + surface: 'mcp', + label: 'clickstack_sql', + source: source._id.toString(), + }); + expect(attributeValues('mcp.source.kind')).toEqual(['trace']); + }); + + it('leaves the source out of a raw SQL query without one', async () => { + const tile = buildTile('t', 24, 6, { + configType: 'sql', + displayType: 'table', + connectionId: new mongoose.Types.ObjectId().toString(), + sqlTemplate: 'SELECT 1', + }); + + const traced = withToolTracing('clickstack_sql', context, () => + runConfigTile(context.teamId, tile, startDate, endDate), + ); + await traced({}); + + expect(mockGetSource).not.toHaveBeenCalled(); + expect(logCommentOf(0)).not.toHaveProperty('source'); + expect(logCommentOf(0)).toMatchObject({ label: 'clickstack_sql' }); + expect(attributeValues('mcp.source.id')).toEqual([]); + }); + + it('tags each tile with its own source and lists all on the span', async () => { + const logSource = fakeSource(SourceKind.Log); + const traceSource = fakeSource(SourceKind.Trace); + mockGetSource.mockImplementation(async (_team: string, id: string) => + [logSource, traceSource].find(s => s._id.toString() === id), + ); + const tileFor = (source: ReturnType) => + buildTile('t', 12, 4, { + displayType: 'table', + sourceId: source._id.toString(), + select: [{ aggFn: 'count' }], + }); + + const traced = withToolTracing('clickstack_query_tiles', context, () => + Promise.all([ + runConfigTile(context.teamId, tileFor(logSource), startDate, endDate), + runConfigTile(context.teamId, tileFor(traceSource), startDate, endDate), + ]).then(([first]) => first), + ); + await traced({}); + + expect(mockQueries.map((_, i) => logCommentOf(i).source).sort()).toEqual( + [logSource._id.toString(), traceSource._id.toString()].sort(), + ); + expect(attributeValues('mcp.source.ids').at(-1)).toEqual( + expect.arrayContaining([ + logSource._id.toString(), + traceSource._id.toString(), + ]), + ); + expect(attributeValues('mcp.source.id')).toHaveLength(1); + }); +}); diff --git a/packages/api/src/mcp/__tests__/tracing.test.ts b/packages/api/src/mcp/__tests__/tracing.test.ts index aab68793c2..7865748737 100644 --- a/packages/api/src/mcp/__tests__/tracing.test.ts +++ b/packages/api/src/mcp/__tests__/tracing.test.ts @@ -9,16 +9,26 @@ const mockSpan = { end: jest.fn(), }; +// Any span other than the tool span, so tests can tell where an attribute +// landed. +const mockChildSpan = { + setAttribute: jest.fn(), + setAttributes: jest.fn(), + setStatus: jest.fn(), + recordException: jest.fn(), + end: jest.fn(), +}; + const mockTracer = { startActiveSpan: ( - _name: string, + name: string, _optionsOrFn: unknown, maybeFn?: (span: typeof mockSpan) => Promise, ) => { const fn = ( typeof _optionsOrFn === 'function' ? _optionsOrFn : maybeFn ) as (span: typeof mockSpan) => Promise; - return fn(mockSpan); + return fn(name.startsWith('mcp.tool.') ? mockSpan : mockChildSpan); }, }; @@ -77,10 +87,13 @@ jest.mock('@/utils/logger', () => ({ }, })); +import { SourceKind } from '@hyperdx/common-utils/dist/types'; + import type { McpContext } from '@/mcp/tools/types'; import { mcpServerError, mcpUserError } from '@/mcp/utils/errors'; import type { McpClientInfo } from '@/mcp/utils/mcpClient'; -import { withToolTracing } from '@/mcp/utils/tracing'; +import { recordMcpSource, withToolTracing } from '@/mcp/utils/tracing'; +import { withSpan } from '@/utils/instrumentation'; // Build a context whose mcpClient resolves to the given identity. The actual // clientInfo/User-Agent resolution logic is covered in mcpClient.test.ts. @@ -336,6 +349,7 @@ describe('withToolTracing', () => { expect.any(Number), { tool: 'my_tool', + source_kind: 'none', }, ); expect(mockCounter.add).not.toHaveBeenCalled(); @@ -386,6 +400,7 @@ describe('withToolTracing', () => { expect(mockCounter.add).toHaveBeenCalledWith(1, { tool: 'my_tool', error_category: 'server', + source_kind: 'none', }); }); @@ -399,11 +414,13 @@ describe('withToolTracing', () => { expect(mockCounter.add).toHaveBeenCalledWith(1, { tool: 'my_tool', error_category: 'server', + source_kind: 'none', }); expect(mockDurationHistogram.record).toHaveBeenCalledWith( expect.any(Number), { tool: 'my_tool', + source_kind: 'none', }, ); }); @@ -438,6 +455,7 @@ describe('withToolTracing', () => { expect(mockCounter.add).toHaveBeenCalledWith(1, { tool: 'my_tool', error_category: 'user', + source_kind: 'none', }); }); @@ -456,6 +474,7 @@ describe('withToolTracing', () => { expect(mockCounter.add).toHaveBeenCalledWith(1, { tool: 'my_tool', error_category: 'server', + source_kind: 'none', }); }); @@ -485,4 +504,156 @@ describe('withToolTracing', () => { }); expect(result).not.toHaveProperty('_errorCategory'); }); + + describe('source attribution', () => { + const logSource = { _id: 'source-log', kind: SourceKind.Log }; + const traceSource = { _id: 'source-trace', kind: SourceKind.Trace }; + const ok = { content: [{ type: 'text' as const, text: 'ok' }] }; + + it('records the source id, kind and ids on the tool span', async () => { + const traced = withToolTracing('my_tool', context, async () => { + recordMcpSource(logSource); + return ok; + }); + await traced({}); + + expect(mockSpan.setAttribute).toHaveBeenCalledWith( + 'mcp.source.id', + 'source-log', + ); + expect(mockSpan.setAttribute).toHaveBeenCalledWith( + 'mcp.source.kind', + 'log', + ); + expect(mockSpan.setAttribute).toHaveBeenCalledWith('mcp.source.ids', [ + 'source-log', + ]); + }); + + it('lands on the tool span when a nested child span is active', async () => { + const traced = withToolTracing('my_tool', context, () => + withSpan('mcp.child.work', async () => { + recordMcpSource(logSource); + return ok; + }), + ); + await traced({}); + + expect(mockSpan.setAttribute).toHaveBeenCalledWith( + 'mcp.source.id', + 'source-log', + ); + expect(mockChildSpan.setAttribute).not.toHaveBeenCalled(); + expect(mockChildSpan.setAttributes).not.toHaveBeenCalled(); + }); + + it('keeps the first source as primary and lists every distinct source', async () => { + const traced = withToolTracing('my_tool', context, async () => { + recordMcpSource(traceSource); + recordMcpSource(logSource); + recordMcpSource(traceSource); + return ok; + }); + await traced({}); + + const calls = mockSpan.setAttribute.mock.calls; + const valuesOf = (key: string) => + calls.filter(([k]) => k === key).map(([, v]) => v); + expect(valuesOf('mcp.source.id')).toEqual(['source-trace']); + expect(valuesOf('mcp.source.kind')).toEqual(['trace']); + expect(valuesOf('mcp.source.ids').at(-1)).toEqual([ + 'source-trace', + 'source-log', + ]); + }); + + it('does not set source attributes when no source is resolved', async () => { + const traced = withToolTracing('my_tool', context, async () => ok); + await traced({}); + + for (const key of [ + 'mcp.source.id', + 'mcp.source.kind', + 'mcp.source.ids', + ]) { + expect(mockSpan.setAttribute).not.toHaveBeenCalledWith( + key, + expect.anything(), + ); + } + }); + + it('is a no-op outside a tool invocation', () => { + recordMcpSource(logSource); + expect(mockSpan.setAttribute).not.toHaveBeenCalled(); + expect(mockChildSpan.setAttribute).not.toHaveBeenCalled(); + }); + + it('records the source kind, never the id, on the duration metric', async () => { + const traced = withToolTracing('my_tool', context, async () => { + recordMcpSource(traceSource); + return ok; + }); + await traced({}); + + expect(mockDurationHistogram.record).toHaveBeenCalledWith( + expect.any(Number), + { tool: 'my_tool', source_kind: 'trace' }, + ); + }); + + it('records the source kind on the error counter for error results', async () => { + const traced = withToolTracing('my_tool', context, async () => { + recordMcpSource(logSource); + return mcpUserError('bad input'); + }); + await traced({}); + + expect(mockCounter.add).toHaveBeenCalledWith(1, { + tool: 'my_tool', + error_category: 'user', + source_kind: 'log', + }); + }); + + it('records the source kind on both metrics when the handler throws', async () => { + const traced = withToolTracing('my_tool', context, async () => { + recordMcpSource(logSource); + throw new Error('boom'); + }); + await expect(traced({})).rejects.toThrow('boom'); + + expect(mockCounter.add).toHaveBeenCalledWith(1, { + tool: 'my_tool', + error_category: 'server', + source_kind: 'log', + }); + expect(mockDurationHistogram.record).toHaveBeenCalledWith( + expect.any(Number), + { tool: 'my_tool', source_kind: 'log' }, + ); + }); + + it('keeps concurrent tool invocations apart', async () => { + const a = withToolTracing('tool_a', context, async () => { + await new Promise(resolve => setTimeout(resolve, 5)); + recordMcpSource(logSource); + return ok; + }); + const b = withToolTracing('tool_b', context, async () => { + recordMcpSource(traceSource); + return ok; + }); + await Promise.all([a({}), b({})]); + + expect(mockDurationHistogram.record).toHaveBeenCalledWith( + expect.any(Number), + { tool: 'tool_a', source_kind: 'log' }, + ); + expect(mockDurationHistogram.record).toHaveBeenCalledWith( + expect.any(Number), + { tool: 'tool_b', source_kind: 'trace' }, + ); + }); + }); }); diff --git a/packages/api/src/mcp/tools/query/denoise.ts b/packages/api/src/mcp/tools/query/denoise.ts index 182aa54487..d8406279aa 100644 --- a/packages/api/src/mcp/tools/query/denoise.ts +++ b/packages/api/src/mcp/tools/query/denoise.ts @@ -1,3 +1,4 @@ +import { withQueryAttribution } from '@hyperdx/common-utils/dist/clickhouse/node'; import { getMetadata } from '@hyperdx/common-utils/dist/core/metadata'; import { getFirstTimestampValueExpression } from '@hyperdx/common-utils/dist/core/utils'; import { @@ -15,6 +16,7 @@ import { ClickhouseClient } from '@/clickhouse'; import { getConnectionById } from '@/controllers/connection'; import { getSource } from '@/controllers/sources'; import { MCP_QUERY_MAX_EXECUTION_SEC } from '@/mcp/utils/timeout'; +import { recordMcpSource } from '@/mcp/utils/tracing'; import { resolveBodyExpression } from './helpers'; @@ -67,6 +69,7 @@ export async function denoiseSearchResults( if (!source) { return { rows, removedPatterns: [], skipped: 'source_not_found' }; } + recordMcpSource(source); const bodyColumn = resolveBodyExpression(source); if (!bodyColumn) { @@ -147,28 +150,32 @@ export async function denoiseSearchResults( ReturnType >; try { - [sampleResult, countResult] = await Promise.all([ - clickhouseClient.queryChartConfig({ - config: sampleConfig, - metadata, - querySettings: source.querySettings, - opts: { - clickhouse_settings: { - max_execution_time: MCP_QUERY_MAX_EXECUTION_SEC, - }, - }, - }), - clickhouseClient.queryChartConfig({ - config: countConfig, - metadata, - querySettings: source.querySettings, - opts: { - clickhouse_settings: { - max_execution_time: MCP_QUERY_MAX_EXECUTION_SEC, - }, - }, - }), - ]); + [sampleResult, countResult] = await withQueryAttribution( + { source: source._id.toString() }, + () => + Promise.all([ + clickhouseClient.queryChartConfig({ + config: sampleConfig, + metadata, + querySettings: source.querySettings, + opts: { + clickhouse_settings: { + max_execution_time: MCP_QUERY_MAX_EXECUTION_SEC, + }, + }, + }), + clickhouseClient.queryChartConfig({ + config: countConfig, + metadata, + querySettings: source.querySettings, + opts: { + clickhouse_settings: { + max_execution_time: MCP_QUERY_MAX_EXECUTION_SEC, + }, + }, + }), + ]), + ); } catch { // If sampling fails, return rows unmodified rather than failing the search return { rows, removedPatterns: [], skipped: 'sampling_failed' }; diff --git a/packages/api/src/mcp/tools/query/eventDeltas.ts b/packages/api/src/mcp/tools/query/eventDeltas.ts index 95824415ef..a3d7c2226e 100644 --- a/packages/api/src/mcp/tools/query/eventDeltas.ts +++ b/packages/api/src/mcp/tools/query/eventDeltas.ts @@ -1,3 +1,4 @@ +import { withQueryAttribution } from '@hyperdx/common-utils/dist/clickhouse/node'; import { getStableSampleExpression, rankProperties, @@ -18,6 +19,7 @@ import { READ_ONLY_TOOL } from '@/mcp/tools/annotations'; import type { ToolRegistrar } from '@/mcp/tools/types'; import { mcpUserError } from '@/mcp/utils/errors'; import { MCP_QUERY_MAX_EXECUTION_SEC } from '@/mcp/utils/timeout'; +import { recordMcpSource } from '@/mcp/utils/tracing'; import { trimToolResponse } from '@/utils/trimToolResponse'; import { PREFER_BUILDER_OVER_SQL_NUDGE } from './builderCatalog'; @@ -245,6 +247,8 @@ export function registerEventDeltas({ context, registerTool }: ToolRegistrar) { `Source not found: ${input.sourceId}. Call clickstack_list_sources to find available source IDs.`, ); } + recordMcpSource(source); + const sourceAttribution = { source: source._id.toString() }; if (source.kind !== SourceKind.Trace && source.kind !== SourceKind.Log) { return mcpUserError( `Source ${input.sourceId} is kind="${source.kind}". clickstack_event_deltas requires a trace or log source.`, @@ -318,17 +322,20 @@ export function registerEventDeltas({ context, registerTool }: ToolRegistrar) { let columnMeta: { name: string; type: string }[] = []; let columnMetaUnavailable = false; try { - const [targetResult, baselineResult, colsResult] = await Promise.all([ - renderChartConfig(targetConfig, metadata, source.querySettings), - renderChartConfig(baselineConfig, metadata, source.querySettings), - metadata - .getColumns({ - databaseName: source.from.databaseName, - tableName: source.from.tableName, - connectionId: source.connection.toString(), - }) - .catch(() => null), - ]); + const [targetResult, baselineResult, colsResult] = + await withQueryAttribution(sourceAttribution, () => + Promise.all([ + renderChartConfig(targetConfig, metadata, source.querySettings), + renderChartConfig(baselineConfig, metadata, source.querySettings), + metadata + .getColumns({ + databaseName: source.from.databaseName, + tableName: source.from.tableName, + connectionId: source.connection.toString(), + }) + .catch(() => null), + ]), + ); targetSql = targetResult; baselineSql = baselineResult; if (colsResult) { @@ -350,28 +357,32 @@ export function registerEventDeltas({ context, registerTool }: ToolRegistrar) { let baselineRows: Record[]; const abortController = new AbortController(); try { - const [targetRes, baselineRes] = await Promise.all([ - clickhouseClient.query({ - query: targetSql.sql, - query_params: targetSql.params, - format: 'JSON', - connectionId: source.connection.toString(), - clickhouse_settings: { - max_execution_time: MCP_QUERY_MAX_EXECUTION_SEC, - }, - abort_signal: abortController.signal, - }), - clickhouseClient.query({ - query: baselineSql.sql, - query_params: baselineSql.params, - format: 'JSON', - connectionId: source.connection.toString(), - clickhouse_settings: { - max_execution_time: MCP_QUERY_MAX_EXECUTION_SEC, - }, - abort_signal: abortController.signal, - }), - ]); + const [targetRes, baselineRes] = await withQueryAttribution( + sourceAttribution, + () => + Promise.all([ + clickhouseClient.query({ + query: targetSql.sql, + query_params: targetSql.params, + format: 'JSON', + connectionId: source.connection.toString(), + clickhouse_settings: { + max_execution_time: MCP_QUERY_MAX_EXECUTION_SEC, + }, + abort_signal: abortController.signal, + }), + clickhouseClient.query({ + query: baselineSql.sql, + query_params: baselineSql.params, + format: 'JSON', + connectionId: source.connection.toString(), + clickhouse_settings: { + max_execution_time: MCP_QUERY_MAX_EXECUTION_SEC, + }, + abort_signal: abortController.signal, + }), + ]), + ); const targetJson = (await ( targetRes as { json: () => Promise<{ data: any[] }> } ).json()) ?? { data: [] }; diff --git a/packages/api/src/mcp/tools/query/helpers.ts b/packages/api/src/mcp/tools/query/helpers.ts index 028164edbe..99d90a1f5d 100644 --- a/packages/api/src/mcp/tools/query/helpers.ts +++ b/packages/api/src/mcp/tools/query/helpers.ts @@ -2,6 +2,7 @@ import { ClickHouseError, type ClickHouseSettings, } from '@clickhouse/client-common'; +import { withQueryAttribution } from '@hyperdx/common-utils/dist/clickhouse/node'; import { getHeatmapMode } from '@hyperdx/common-utils/dist/core/heatmap'; import { getMetadata } from '@hyperdx/common-utils/dist/core/metadata'; import { @@ -38,6 +39,7 @@ import { MCP_TIMEOUT_GRACE_MS, MCP_TOOL_TIMEOUT_MS, } from '@/mcp/utils/timeout'; +import { recordMcpSource } from '@/mcp/utils/tracing'; import { convertToInternalTileConfig, isConfigTile, @@ -476,6 +478,7 @@ export async function runConfigTile( `Source not found: ${builderConfig.source}. Call clickstack_list_sources to discover available source IDs.`, ); } + recordMcpSource(source); // Reject metric-style select against a non-metric source (and vice // versa) before the renderer composes SQL against the wrong table. @@ -606,15 +609,21 @@ export async function runConfigTile( const metadata = getMetadata(clickhouseClient); try { - const result = await clickhouseClient.queryChartConfig({ - config: renderConfig, - metadata, - querySettings: source.querySettings, - opts: { - clickhouse_settings: MCP_CLICKHOUSE_SETTINGS, - abort_signal: options?.abortSignal, - }, - }); + // queryChartConfig tags its own query with config.source, but not the + // metadata queries renderChartConfig issues on the way there. + const result = await withQueryAttribution( + { source: source._id.toString() }, + () => + clickhouseClient.queryChartConfig({ + config: renderConfig, + metadata, + querySettings: source.querySettings, + opts: { + clickhouse_settings: MCP_CLICKHOUSE_SETTINGS, + abort_signal: options?.abortSignal, + }, + }), + ); return formatQueryResult(result); } catch (e) { return clickHouseErrorResult(e); @@ -633,9 +642,12 @@ export async function runConfigTile( useTextIndexForImplicitColumn?: UseTextIndex; metricTables?: MetricTable; } = {}; + let resolvedSourceId: string | undefined; if (savedConfig.source) { const source = await getSource(teamId, savedConfig.source); if (source) { + recordMcpSource(source); + resolvedSourceId = source._id.toString(); sourceFields = { from: source.from, implicitColumnExpression: @@ -681,15 +693,20 @@ export async function runConfigTile( const metadata = getMetadata(clickhouseClient); try { - const result = await clickhouseClient.queryChartConfig({ - config: chartConfig, - metadata, - querySettings: undefined, - opts: { - clickhouse_settings: MCP_CLICKHOUSE_SETTINGS, - abort_signal: options?.abortSignal, - }, - }); + // An unset source leaves the outer attribution as is. + const result = await withQueryAttribution( + { source: resolvedSourceId }, + () => + clickhouseClient.queryChartConfig({ + config: chartConfig, + metadata, + querySettings: undefined, + opts: { + clickhouse_settings: MCP_CLICKHOUSE_SETTINGS, + abort_signal: options?.abortSignal, + }, + }), + ); return formatQueryResult(result); } catch (e) { return clickHouseErrorResult(e, { rawSql: true }); diff --git a/packages/api/src/mcp/tools/query/runEventPatterns.ts b/packages/api/src/mcp/tools/query/runEventPatterns.ts index b9a2464c4d..fd4c7e18ba 100644 --- a/packages/api/src/mcp/tools/query/runEventPatterns.ts +++ b/packages/api/src/mcp/tools/query/runEventPatterns.ts @@ -1,3 +1,4 @@ +import { withQueryAttribution } from '@hyperdx/common-utils/dist/clickhouse/node'; import { getMetadata } from '@hyperdx/common-utils/dist/core/metadata'; import { getFirstTimestampValueExpression, @@ -19,6 +20,7 @@ import { mcpUserError, } from '@/mcp/utils/errors'; import { MCP_QUERY_MAX_EXECUTION_SEC } from '@/mcp/utils/timeout'; +import { recordMcpSource } from '@/mcp/utils/tracing'; import { trimToolResponse } from '@/utils/trimToolResponse'; import { @@ -72,6 +74,7 @@ export async function mineWindowPatterns( ), }; } + recordMcpSource(source); const connection = await getConnectionById( teamId, source.connection.toString(), @@ -181,30 +184,34 @@ export async function mineWindowPatterns( // (emerging_signals opens two per invocation). let sampleResult, countResult; try { - [sampleResult, countResult] = await Promise.all([ - clickhouseClient.queryChartConfig({ - config: sampleConfig, - metadata, - querySettings: source.querySettings, - opts: { - clickhouse_settings: { - max_execution_time: MCP_QUERY_MAX_EXECUTION_SEC, - }, - abort_signal: options?.abortSignal, - }, - }), - clickhouseClient.queryChartConfig({ - config: countConfig, - metadata, - querySettings: source.querySettings, - opts: { - clickhouse_settings: { - max_execution_time: MCP_QUERY_MAX_EXECUTION_SEC, - }, - abort_signal: options?.abortSignal, - }, - }), - ]); + [sampleResult, countResult] = await withQueryAttribution( + { source: source._id.toString() }, + () => + Promise.all([ + clickhouseClient.queryChartConfig({ + config: sampleConfig, + metadata, + querySettings: source.querySettings, + opts: { + clickhouse_settings: { + max_execution_time: MCP_QUERY_MAX_EXECUTION_SEC, + }, + abort_signal: options?.abortSignal, + }, + }), + clickhouseClient.queryChartConfig({ + config: countConfig, + metadata, + querySettings: source.querySettings, + opts: { + clickhouse_settings: { + max_execution_time: MCP_QUERY_MAX_EXECUTION_SEC, + }, + abort_signal: options?.abortSignal, + }, + }), + ]), + ); } catch (err) { return { error: clickHouseErrorResult(err) }; } finally { diff --git a/packages/api/src/mcp/tools/query/serviceHealth/index.ts b/packages/api/src/mcp/tools/query/serviceHealth/index.ts index d0d9e72128..eaec6aa704 100644 --- a/packages/api/src/mcp/tools/query/serviceHealth/index.ts +++ b/packages/api/src/mcp/tools/query/serviceHealth/index.ts @@ -1,3 +1,4 @@ +import { withQueryAttribution } from '@hyperdx/common-utils/dist/clickhouse/node'; import { getMetadata } from '@hyperdx/common-utils/dist/core/metadata'; import { splitAndTrimWithBracket } from '@hyperdx/common-utils/dist/core/utils'; import { trace } from '@opentelemetry/api'; @@ -90,119 +91,121 @@ export function registerServiceHealth({ const source = await loadHealthSource(teamId, input.sourceId, TOOL_NAME); if ('isError' in source) return source; - const client = await getMcpClickhouseClient(teamId, source.connection); - if (!client) { - return mcpUserError( - `Connection not found for source: ${input.sourceId}`, + return withQueryAttribution({ source: source.id }, async () => { + const client = await getMcpClickhouseClient(teamId, source.connection); + if (!client) { + return mcpUserError( + `Connection not found for source: ${input.sourceId}`, + ); + } + const groupExprs = errorGroupExpressions( + source, + input.errorGroupBy, + splitAndTrimWithBracket, ); - } - const groupExprs = errorGroupExpressions( - source, - input.errorGroupBy, - splitAndTrimWithBracket, - ); - const abort = new AbortController(); - const ctx: QueryContext = { - client, - metadata: getMetadata(client), - source, - input, - signal: abort.signal, - }; + const abort = new AbortController(); + const ctx: QueryContext = { + client, + metadata: getMetadata(client), + source, + input, + signal: abort.signal, + }; - try { - const [cur, requestedCmp] = await Promise.all([ - fetchServiceStats(ctx, current), - requested ? fetchServiceStats(ctx, requested) : null, - ]); - const { comparison, cmp, comparisonMode, comparisonNote } = - await resolveComparison(ctx, { + try { + const [cur, requestedCmp] = await Promise.all([ + fetchServiceStats(ctx, current), + requested ? fetchServiceStats(ctx, requested) : null, + ]); + const { comparison, cmp, comparisonMode, comparisonNote } = + await resolveComparison(ctx, { + current, + requested, + requestedMode: input.comparison, + curHasData: cur.stats.size > 0, + requestedCmp, + }); + const ranked = rankServices( + cur.stats, + cmp?.stats ?? null, + input.minRequests, + ); + + const { + currentErrors, + comparisonErrorCounts, + newOperations, + coverageNote, + } = await collectErrors(ctx, { current, - requested, - requestedMode: input.comparison, - curHasData: cur.stats.size > 0, - requestedCmp, + comparison, + // Novelty can lift a service into the top N, so assess errors + // beyond it rather than only for services already ranked there. + shown: ranked.slice(0, MAX_ERROR_SERVICES), + perService: input.topErrorsPerService, + groupExprs, + minRequests: input.minRequests, + totalServices: ranked.length, }); - const ranked = rankServices( - cur.stats, - cmp?.stats ?? null, - input.minRequests, - ); - - const { - currentErrors, - comparisonErrorCounts, - newOperations, - coverageNote, - } = await collectErrors(ctx, { - current, - comparison, - // Novelty can lift a service into the top N, so assess errors - // beyond it rather than only for services already ranked there. - shown: ranked.slice(0, MAX_ERROR_SERVICES), - perService: input.topErrorsPerService, - groupExprs, - minRequests: input.minRequests, - totalServices: ranked.length, - }); - const report = buildReport({ - ranked, - current, - comparison, - comparisonMode, - comparisonNote, - coverageNote, - spanKind: input.spanKind, - topN: input.topN, - errorGroupBy: groupExprs, - currentErrors, - comparisonErrorCounts, - newOperations, - spanNameKey: source.spanNameExpression, - metricSourceId: source.metricSourceId, - }); + const report = buildReport({ + ranked, + current, + comparison, + comparisonMode, + comparisonNote, + coverageNote, + spanKind: input.spanKind, + topN: input.topN, + errorGroupBy: groupExprs, + currentErrors, + comparisonErrorCounts, + newOperations, + spanNameKey: source.spanNameExpression, + metricSourceId: source.metricSourceId, + }); - trace.getActiveSpan()?.setAttributes({ - 'hyperdx.source.id': source.id, - 'mcp.service_health.comparison': input.comparison, - 'mcp.service_health.comparison_used': comparisonMode, - 'mcp.service_health.span_kind': input.spanKind, - 'mcp.service_health.window_ms': - current.end.getTime() - current.start.getTime(), - 'mcp.service_health.services_assessed': ranked.length, - 'mcp.service_health.critical_count': - report.summary.statusCounts.critical, - 'mcp.service_health.degraded_count': - report.summary.statusCounts.degraded, - }); + trace.getActiveSpan()?.setAttributes({ + 'hyperdx.source.id': source.id, + 'mcp.service_health.comparison': input.comparison, + 'mcp.service_health.comparison_used': comparisonMode, + 'mcp.service_health.span_kind': input.spanKind, + 'mcp.service_health.window_ms': + current.end.getTime() - current.start.getTime(), + 'mcp.service_health.services_assessed': ranked.length, + 'mcp.service_health.critical_count': + report.summary.statusCounts.critical, + 'mcp.service_health.degraded_count': + report.summary.statusCounts.degraded, + }); - const { data, isTrimmed } = trimToolResponse(report); - return { - content: [ - { - type: 'text' as const, - text: JSON.stringify( - isTrimmed - ? { - ...data, - note: 'Trimmed for context size. Lower topN or topErrorsPerService.', - } - : data, - null, - 2, - ), - }, - ], - }; - } catch (e) { - abort.abort(); - return clickHouseErrorResult(e, { - prefix: 'Failed to compute service health', - suffix: isQueryOutOfTime(e) ? TIMEOUT_SUFFIX : undefined, - }); - } + const { data, isTrimmed } = trimToolResponse(report); + return { + content: [ + { + type: 'text' as const, + text: JSON.stringify( + isTrimmed + ? { + ...data, + note: 'Trimmed for context size. Lower topN or topErrorsPerService.', + } + : data, + null, + 2, + ), + }, + ], + }; + } catch (e) { + abort.abort(); + return clickHouseErrorResult(e, { + prefix: 'Failed to compute service health', + suffix: isQueryOutOfTime(e) ? TIMEOUT_SUFFIX : undefined, + }); + } + }); }, ); } diff --git a/packages/api/src/mcp/tools/query/serviceHealth/source.ts b/packages/api/src/mcp/tools/query/serviceHealth/source.ts index 4e7553693f..53396f8f05 100644 --- a/packages/api/src/mcp/tools/query/serviceHealth/source.ts +++ b/packages/api/src/mcp/tools/query/serviceHealth/source.ts @@ -3,6 +3,7 @@ import { SourceKind } from '@hyperdx/common-utils/dist/types'; import { getSource, getSources } from '@/controllers/sources'; import type { McpErrorResult } from '@/mcp/utils/errors'; import { mcpUserError } from '@/mcp/utils/errors'; +import { recordMcpSource } from '@/mcp/utils/tracing'; import type { HealthSource } from './queries'; @@ -56,6 +57,7 @@ export async function loadHealthSource( ): Promise { const doc = await getSource(teamId, sourceId); if (!doc) return sourceNotFound(teamId, sourceId, toolName); + recordMcpSource(doc); if (doc.kind !== SourceKind.Trace) { return mcpUserError( `Source ${sourceId} is kind="${doc.kind}". ${toolName} requires a source of kind="trace".`, diff --git a/packages/api/src/mcp/tools/query/sql.ts b/packages/api/src/mcp/tools/query/sql.ts index 28ac32a851..729e771b1d 100644 --- a/packages/api/src/mcp/tools/query/sql.ts +++ b/packages/api/src/mcp/tools/query/sql.ts @@ -6,6 +6,7 @@ import { getSource } from '@/controllers/sources'; import { READ_ONLY_TOOL } from '@/mcp/tools/annotations'; import type { ToolRegistrar } from '@/mcp/tools/types'; import { mcpUserError } from '@/mcp/utils/errors'; +import { recordMcpSource } from '@/mcp/utils/tracing'; import { BUILDER_TOOLS_LIST, SQL_FALLBACK_CRITERIA } from './builderCatalog'; import { buildTile, parseTimeRange, runConfigTile } from './helpers'; @@ -126,6 +127,7 @@ export function registerSql({ context, registerTool }: ToolRegistrar) { `Source not found: ${input.sourceId}. Call clickstack_list_sources to discover available source IDs.`, ); } + recordMcpSource(source); // $__sourceTable expands to the source's table, so a mismatch would run // the query against a table on another server. const sourceConnectionId = source.connection.toString(); diff --git a/packages/api/src/mcp/tools/sources/describeMetric.ts b/packages/api/src/mcp/tools/sources/describeMetric.ts index 36ccec9ca8..3f6796179b 100644 --- a/packages/api/src/mcp/tools/sources/describeMetric.ts +++ b/packages/api/src/mcp/tools/sources/describeMetric.ts @@ -6,18 +6,19 @@ import { JSDataType, tableExpr, } from '@hyperdx/common-utils/dist/clickhouse'; +import { withQueryAttribution } from '@hyperdx/common-utils/dist/clickhouse/node'; import { getMetadata } from '@hyperdx/common-utils/dist/core/metadata'; import { SourceKind } from '@hyperdx/common-utils/dist/types'; import { z } from 'zod'; -import { ClickhouseClient } from '@/clickhouse'; -import { getConnectionById } from '@/controllers/connection'; +import type { ClickhouseClient } from '@/clickhouse'; import { getSource } from '@/controllers/sources'; import { READ_ONLY_TOOL } from '@/mcp/tools/annotations'; import { clickHouseErrorResult, parseTimeRange, } from '@/mcp/tools/query/helpers'; +import { getMcpClickhouseClient } from '@/mcp/tools/trace/shared'; import type { ToolRegistrar } from '@/mcp/tools/types'; import { mcpServerError, @@ -25,6 +26,7 @@ import { sanitizeFetchError, } from '@/mcp/utils/errors'; import { MCP_TOOL_TIMEOUT_MS, runWithTimeout } from '@/mcp/utils/timeout'; +import { recordMcpSource } from '@/mcp/utils/tracing'; import logger from '@/utils/logger'; import { trimToolResponse } from '@/utils/trimToolResponse'; @@ -478,6 +480,7 @@ async function describeMetricImpl( `Source "${input.sourceId}" not found. Call clickstack_list_sources to see available source IDs.`, ); } + recordMcpSource(source); if (source.kind !== SourceKind.Metric) { return mcpUserError( `Source "${input.sourceId}" is a "${source.kind}" source, not a metric source. clickstack_describe_metric only works on metric sources.`, @@ -494,212 +497,211 @@ async function describeMetricImpl( } const { startDate, endDate } = timeRange; - const connection = await getConnectionById( - teamId, - source.connection.toString(), - true, - ); - if (!connection) { - return mcpUserError(`Connection not found for source "${input.sourceId}".`); - } - - const clickhouseClient = new ClickhouseClient({ - host: connection.host, - username: connection.username, - password: connection.password, - }); - const metadata = getMetadata(clickhouseClient); - const databaseName = source.from.databaseName; - const connectionId = source.connection.toString(); - - // Validate the requested kind has a populated table on this source. - // The schema requires `kind`, so we always have an exact target. - const kind = input.kind; - const tableName = source.metricTables[kind]; - if (!tableName) { - return mcpUserError( - `Source "${input.sourceId}" has no "${kind}" metric table populated. Populated kinds: ${Object.keys(source.metricTables).join(', ')}.`, - ); - } - - // Defensive column-presence check before referencing MetricUnit / - // MetricDescription on this kind's table. - let columns: { name: string; type: string }[]; - try { - columns = await metadata.getColumns({ - databaseName, - tableName, - connectionId, - }); - } catch (e) { - logger.warn( - { kind, error: e instanceof Error ? e.message : String(e) }, - 'describeMetric: getColumns failed', + return withQueryAttribution({ source: source._id.toString() }, async () => { + const clickhouseClient = await getMcpClickhouseClient( + teamId, + source.connection.toString(), ); - return clickHouseErrorResult(e, { - prefix: `Failed to load columns for "${tableName}"`, - suffix: 'The metric table may be missing or unreachable.', - }); - } - const columnNames = new Set(columns.map(c => c.name)); - const hasUnit = columnNames.has('MetricUnit'); - const hasDescription = columnNames.has('MetricDescription'); - - const [meta, attributeKeysResult] = await Promise.all([ - fetchUnitAndDescription({ - clickhouseClient, - databaseName, - tableName, - connectionId, - metricName: input.metricName, - startDate, - endDate, - hasUnit, - hasDescription, - signal, - }), - fetchAttributeKeys({ - clickhouseClient, - databaseName, - tableName, - connectionId, - metricName: input.metricName, - columns, - startDate, - endDate, - signal, - }), - ]); - - // Track discovery sub-queries that failed so the agent can distinguish - // "fetch failed" (retry / report) from "genuinely empty" (widen the - // window). Without this, transient ClickHouse errors surfaced as the - // misleading "No data found" hint. - const partialFailure: { stage: string; error: string }[] = []; - const attributeKeys = attributeKeysResult.ok ? attributeKeysResult.data : {}; - if (!attributeKeysResult.ok) { - partialFailure.push({ - stage: 'attributeKeys', - error: attributeKeysResult.error, - }); - } - - const kindDetail: KindDetail = { - kind, - tableName, - ...(meta.unit ? { unit: meta.unit } : {}), - ...(meta.description ? { description: meta.description } : {}), - attributeKeys, - usage: KIND_USAGE[kind], - }; + if (!clickhouseClient) { + return mcpUserError( + `Connection not found for source "${input.sourceId}".`, + ); + } + const metadata = getMetadata(clickhouseClient); + const databaseName = source.from.databaseName; + const connectionId = source.connection.toString(); + + // Validate the requested kind has a populated table on this source. + // The schema requires `kind`, so we always have an exact target. + const kind = input.kind; + const tableName = source.metricTables[kind]; + if (!tableName) { + return mcpUserError( + `Source "${input.sourceId}" has no "${kind}" metric table populated. Populated kinds: ${Object.keys(source.metricTables).join(', ')}.`, + ); + } - if (input.sampleValues && Object.keys(attributeKeys).length > 0) { - const valuesResult = await sampleAttributeValues({ - clickhouseClient, - databaseName, - tableName, - connectionId, - metricName: input.metricName, - attributeKeys, - startDate, - endDate, - signal, - }); - if (valuesResult.ok) { - kindDetail.attributeValues = valuesResult.data.values; - // sampledKeys / truncatedKeys let the agent tell "key absent - // because we never queried it" (truncated) apart from "queried - // but empty in the window" (in sampledKeys, missing from values). - kindDetail.attributeValuesMeta = valuesResult.data.meta; - } else { + // Defensive column-presence check before referencing MetricUnit / + // MetricDescription on this kind's table. + let columns: { name: string; type: string }[]; + try { + columns = await metadata.getColumns({ + databaseName, + tableName, + connectionId, + }); + } catch (e) { + logger.warn( + { kind, error: e instanceof Error ? e.message : String(e) }, + 'describeMetric: getColumns failed', + ); + return clickHouseErrorResult(e, { + prefix: `Failed to load columns for "${tableName}"`, + suffix: 'The metric table may be missing or unreachable.', + }); + } + const columnNames = new Set(columns.map(c => c.name)); + const hasUnit = columnNames.has('MetricUnit'); + const hasDescription = columnNames.has('MetricDescription'); + + const [meta, attributeKeysResult] = await Promise.all([ + fetchUnitAndDescription({ + clickhouseClient, + databaseName, + tableName, + connectionId, + metricName: input.metricName, + startDate, + endDate, + hasUnit, + hasDescription, + signal, + }), + fetchAttributeKeys({ + clickhouseClient, + databaseName, + tableName, + connectionId, + metricName: input.metricName, + columns, + startDate, + endDate, + signal, + }), + ]); + + // Track discovery sub-queries that failed so the agent can distinguish + // "fetch failed" (retry / report) from "genuinely empty" (widen the + // window). Without this, transient ClickHouse errors surfaced as the + // misleading "No data found" hint. + const partialFailure: { stage: string; error: string }[] = []; + const attributeKeys = attributeKeysResult.ok + ? attributeKeysResult.data + : {}; + if (!attributeKeysResult.ok) { partialFailure.push({ - stage: 'attributeValues', - error: valuesResult.error, + stage: 'attributeKeys', + error: attributeKeysResult.error, }); } - } - const kindDetails: KindDetail[] = [kindDetail]; - - // Heuristic "no data in this kind" hint: when neither attribute keys, - // unit, nor description came back, the (metric, kind) pair likely has - // no data in the requested window. The agent's recourse is to widen - // startTime/endTime or double-check the kind via clickstack_list_metrics. - // Suppressed when any discovery stage failed — an empty attributeKeys - // can't be trusted as "no data" then. - const noSignal = - partialFailure.length === 0 && - Object.keys(attributeKeys).length === 0 && - !meta.unit && - !meta.description; - - const makeNextStepsQuery = () => { - if (kind === 'summary') { - // Summary metrics have no renderer support, so the next step is - // clickstack_sql rather than a builder-tool example. - // clickstack_sql runs with no source context, so qualify the table - // with the source's database — a bare table name would resolve - // against the connection user's default database. - const qualifiedTable = databaseName - ? `${databaseName}.${tableName}` - : tableName; - return ( - 'summary metrics cannot be queried with clickstack_timeseries / clickstack_table — ' + - `use clickstack_sql against the "${qualifiedTable}" table with this source's connectionId ` + - '(see clickstack_list_sources).' - ); + const kindDetail: KindDetail = { + kind, + tableName, + ...(meta.unit ? { unit: meta.unit } : {}), + ...(meta.description ? { description: meta.description } : {}), + attributeKeys, + usage: KIND_USAGE[kind], + }; + + if (input.sampleValues && Object.keys(attributeKeys).length > 0) { + const valuesResult = await sampleAttributeValues({ + clickhouseClient, + databaseName, + tableName, + connectionId, + metricName: input.metricName, + attributeKeys, + startDate, + endDate, + signal, + }); + if (valuesResult.ok) { + kindDetail.attributeValues = valuesResult.data.values; + // sampledKeys / truncatedKeys let the agent tell "key absent + // because we never queried it" (truncated) apart from "queried + // but empty in the window" (in sampledKeys, missing from values). + kindDetail.attributeValuesMeta = valuesResult.data.meta; + } else { + partialFailure.push({ + stage: 'attributeValues', + error: valuesResult.error, + }); + } } - let aggFn = '"avg"'; + const kindDetails: KindDetail[] = [kindDetail]; + + // Heuristic "no data in this kind" hint: when neither attribute keys, + // unit, nor description came back, the (metric, kind) pair likely has + // no data in the requested window. The agent's recourse is to widen + // startTime/endTime or double-check the kind via clickstack_list_metrics. + // Suppressed when any discovery stage failed — an empty attributeKeys + // can't be trusted as "no data" then. + const noSignal = + partialFailure.length === 0 && + Object.keys(attributeKeys).length === 0 && + !meta.unit && + !meta.description; + + const makeNextStepsQuery = () => { + if (kind === 'summary') { + // Summary metrics have no renderer support, so the next step is + // clickstack_sql rather than a builder-tool example. + // clickstack_sql runs with no source context, so qualify the table + // with the source's database — a bare table name would resolve + // against the connection user's default database. + const qualifiedTable = databaseName + ? `${databaseName}.${tableName}` + : tableName; + return ( + 'summary metrics cannot be queried with clickstack_timeseries / clickstack_table — ' + + `use clickstack_sql against the "${qualifiedTable}" table with this source's connectionId ` + + '(see clickstack_list_sources).' + ); + } - if (kind === 'sum') { - aggFn = '"increase"'; - } else if (kind === 'histogram' || kind === 'exponential histogram') { - aggFn = '"quantile", level: 0.95'; - } + let aggFn = '"avg"'; - return `Example: clickstack_timeseries({ sourceId: "${input.sourceId}", select: [{ aggFn: ${aggFn}, metricType: "${kind}", metricName: "${input.metricName}" }] })`; - }; + if (kind === 'sum') { + aggFn = '"increase"'; + } else if (kind === 'histogram' || kind === 'exponential histogram') { + aggFn = '"quantile", level: 0.95'; + } - const responseObj: Record = { - metricName: input.metricName, - kinds: kindDetails, - ...(partialFailure.length > 0 && { - partialFailure, - hint: - 'Some discovery queries failed — the fields above may be incomplete. ' + - 'Retry the call; if the failure persists, narrow startTime/endTime or set sampleValues:false.', - }), - ...(noSignal && { - hint: - `No data found for MetricName "${input.metricName}" with kind "${kind}" ` + - `between ${startDate.toISOString()} and ${endDate.toISOString()}. ` + - 'Try widening startTime/endTime, or call clickstack_list_metrics to ' + - 'confirm the metric name + kind combination exists.', - }), - nextSteps: { - query: makeNextStepsQuery(), - }, - }; + return `Example: clickstack_timeseries({ sourceId: "${input.sourceId}", select: [{ aggFn: ${aggFn}, metricType: "${kind}", metricName: "${input.metricName}" }] })`; + }; - const { data: trimmed, isTrimmed } = trimToolResponse(responseObj); - return { - content: [ - { - type: 'text' as const, - text: JSON.stringify( - isTrimmed - ? { - ...trimmed, - note: 'Result was trimmed for context size. Set sampleValues:false or narrow startTime/endTime.', - } - : trimmed, - null, - 2, - ), + const responseObj: Record = { + metricName: input.metricName, + kinds: kindDetails, + ...(partialFailure.length > 0 && { + partialFailure, + hint: + 'Some discovery queries failed — the fields above may be incomplete. ' + + 'Retry the call; if the failure persists, narrow startTime/endTime or set sampleValues:false.', + }), + ...(noSignal && { + hint: + `No data found for MetricName "${input.metricName}" with kind "${kind}" ` + + `between ${startDate.toISOString()} and ${endDate.toISOString()}. ` + + 'Try widening startTime/endTime, or call clickstack_list_metrics to ' + + 'confirm the metric name + kind combination exists.', + }), + nextSteps: { + query: makeNextStepsQuery(), }, - ], - }; + }; + + const { data: trimmed, isTrimmed } = trimToolResponse(responseObj); + return { + content: [ + { + type: 'text' as const, + text: JSON.stringify( + isTrimmed + ? { + ...trimmed, + note: 'Result was trimmed for context size. Set sampleValues:false or narrow startTime/endTime.', + } + : trimmed, + null, + 2, + ), + }, + ], + }; + }); } // ─── Tool registration ─────────────────────────────────────────────────────── diff --git a/packages/api/src/mcp/tools/sources/describeSource.ts b/packages/api/src/mcp/tools/sources/describeSource.ts index ff52479143..03f0d00ee7 100644 --- a/packages/api/src/mcp/tools/sources/describeSource.ts +++ b/packages/api/src/mcp/tools/sources/describeSource.ts @@ -3,6 +3,7 @@ import { filterColumnMetaByType, JSDataType, } from '@hyperdx/common-utils/dist/clickhouse'; +import { withQueryAttribution } from '@hyperdx/common-utils/dist/clickhouse/node'; import { getMetadata } from '@hyperdx/common-utils/dist/core/metadata'; import { type MetricTable, SourceKind } from '@hyperdx/common-utils/dist/types'; import { trace } from '@opentelemetry/api'; @@ -22,6 +23,7 @@ import { runWithTimeout, type TimeoutOutcome, } from '@/mcp/utils/timeout'; +import { recordMcpSource } from '@/mcp/utils/tracing'; import logger from '@/utils/logger'; import { trimToolResponse } from '@/utils/trimToolResponse'; @@ -98,6 +100,7 @@ async function describeSourceSchema( `Source "${sourceId}" not found. Call clickstack_list_sources to see available source IDs.`, ); } + recordMcpSource(source); const meta: Record = { id: source._id.toString(), @@ -193,284 +196,286 @@ async function describeSourceSchema( }; } - const connection = await getConnectionById( - teamId, - source.connection.toString(), - true, - ); - if (!connection) { - return mcpUserError(`Connection not found for source "${sourceId}".`); - } + return withQueryAttribution({ source: source._id.toString() }, async () => { + const connection = await getConnectionById( + teamId, + source.connection.toString(), + true, + ); + if (!connection) { + return mcpUserError(`Connection not found for source "${sourceId}".`); + } - const clickhouseClient = new ClickhouseClient({ - host: connection.host, - username: connection.username, - password: connection.password, - }); - const metadata = getMetadata(clickhouseClient); - const databaseName = source.from.databaseName; - const tableName = discoveryTableName; - const connectionId = source.connection.toString(); - - // Track which sampling stages were skipped due to timeout - const skippedStages: string[] = []; - - // Shared by stages 2–4 so map-key discovery can use rollup tables - // instead of falling back to expensive main-table scans. - const metadataMVs = - 'metadataMaterializedViews' in source - ? source.metadataMaterializedViews - : undefined; - - const now = new Date(); - const dateRange: [Date, Date] = [ - new Date(now.getTime() - VALUE_SAMPLE_LOOKBACK_MS), - now, - ]; - - // ── 1. Column schema ────────────────────────────────────────────────── - const columns = await metadata.getColumns({ - databaseName, - tableName, - connectionId, - }); + const clickhouseClient = new ClickhouseClient({ + host: connection.host, + username: connection.username, + password: connection.password, + }); + const metadata = getMetadata(clickhouseClient); + const databaseName = source.from.databaseName; + const tableName = discoveryTableName; + const connectionId = source.connection.toString(); + + // Track which sampling stages were skipped due to timeout + const skippedStages: string[] = []; + + // Shared by stages 2–4 so map-key discovery can use rollup tables + // instead of falling back to expensive main-table scans. + const metadataMVs = + 'metadataMaterializedViews' in source + ? source.metadataMaterializedViews + : undefined; + + const now = new Date(); + const dateRange: [Date, Date] = [ + new Date(now.getTime() - VALUE_SAMPLE_LOOKBACK_MS), + now, + ]; + + // ── 1. Column schema ────────────────────────────────────────────────── + const columns = await metadata.getColumns({ + databaseName, + tableName, + connectionId, + }); - meta.columns = columns.map(c => ({ - name: c.name, - type: c.type, - jsType: convertCHDataTypeToJSType(c.type), - })); - - const isMetricSource = source.kind === SourceKind.Metric; - const mapColumns = filterColumnMetaByType(columns, [JSDataType.Map]) ?? []; - const lcColumns = columns.filter(c => { - const normalized = c.type.replace(/\s/g, ''); - return ( - normalized.startsWith('LowCardinality(') && - (normalized.includes('String') || normalized.includes('string')) - ); - }); - const metricKinds = - source.kind === SourceKind.Metric - ? DISCOVERABLE_METRIC_KINDS.flatMap(kind => { - const kindTableName = source.metricTables[kind]; - return kindTableName ? [{ kind, kindTableName }] : []; - }) - : []; - // Stages with nothing to sample can't be skipped, so the snapshot must not - // report them as missing. - const optionalStages = [ - ...(mapColumns.length > 0 ? ['mapAttributeKeys'] : []), - ...(lcColumns.length > 0 ? ['lowCardinalityValues'] : []), - ...(mapColumns.length > 0 ? ['mapAttributeValues'] : []), - ...(metricKinds.length > 0 ? ['metricNames'] : []), - ]; - const finishedStages = new Set(); - progress.snapshot = () => - formatDescribeResult({ - sourceId, - meta, - isMetricSource, - skippedStages: [ - ...skippedStages, - ...optionalStages.filter( - s => !finishedStages.has(s) && !skippedStages.includes(s), - ), - ], + meta.columns = columns.map(c => ({ + name: c.name, + type: c.type, + jsType: convertCHDataTypeToJSType(c.type), + })); + + const isMetricSource = source.kind === SourceKind.Metric; + const mapColumns = filterColumnMetaByType(columns, [JSDataType.Map]) ?? []; + const lcColumns = columns.filter(c => { + const normalized = c.type.replace(/\s/g, ''); + return ( + normalized.startsWith('LowCardinality(') && + (normalized.includes('String') || normalized.includes('string')) + ); }); + const metricKinds = + source.kind === SourceKind.Metric + ? DISCOVERABLE_METRIC_KINDS.flatMap(kind => { + const kindTableName = source.metricTables[kind]; + return kindTableName ? [{ kind, kindTableName }] : []; + }) + : []; + // Stages with nothing to sample can't be skipped, so the snapshot must not + // report them as missing. + const optionalStages = [ + ...(mapColumns.length > 0 ? ['mapAttributeKeys'] : []), + ...(lcColumns.length > 0 ? ['lowCardinalityValues'] : []), + ...(mapColumns.length > 0 ? ['mapAttributeValues'] : []), + ...(metricKinds.length > 0 ? ['metricNames'] : []), + ]; + const finishedStages = new Set(); + progress.snapshot = () => + formatDescribeResult({ + sourceId, + meta, + isMetricSource, + skippedStages: [ + ...skippedStages, + ...optionalStages.filter( + s => !finishedStages.has(s) && !skippedStages.includes(s), + ), + ], + }); - // ── 2. Map attribute keys ───────────────────────────────────────────── - // timestampValueExpression is threaded into getMapKeys / getAllKeyValues / - // sampleMetricNamesForKind below so the no-rollup fallback path (i.e. - // metric sources, which don't have metadataMaterializedViews configured) - // can scope its scan to dateRange instead of going unbounded against - // the raw metric table on cold cache. - const timestampValueExpression = source.timestampValueExpression; - const mapKeysResults: Record = {}; - - if (!signal.aborted) { - await Promise.all( - mapColumns.map(async col => { - try { - const keys = await metadata.getMapKeys({ - databaseName, - tableName, - column: col.name, - maxKeys: 50, - connectionId, - metadataMVs, - dateRange, - timestampValueExpression, - signal, - }); - mapKeysResults[col.name] = keys; - // Visible to the backstop snapshot before the other columns finish. - meta.mapAttributeKeys = mapKeysResults; - } catch (e) { - logger.warn( - { sourceId, column: col.name, error: e }, - 'Failed to fetch map keys for column', - ); - } - }), - ); - } + // ── 2. Map attribute keys ───────────────────────────────────────────── + // timestampValueExpression is threaded into getMapKeys / getAllKeyValues / + // sampleMetricNamesForKind below so the no-rollup fallback path (i.e. + // metric sources, which don't have metadataMaterializedViews configured) + // can scope its scan to dateRange instead of going unbounded against + // the raw metric table on cold cache. + const timestampValueExpression = source.timestampValueExpression; + const mapKeysResults: Record = {}; + + if (!signal.aborted) { + await Promise.all( + mapColumns.map(async col => { + try { + const keys = await metadata.getMapKeys({ + databaseName, + tableName, + column: col.name, + maxKeys: 50, + connectionId, + metadataMVs, + dateRange, + timestampValueExpression, + signal, + }); + mapKeysResults[col.name] = keys; + // Visible to the backstop snapshot before the other columns finish. + meta.mapAttributeKeys = mapKeysResults; + } catch (e) { + logger.warn( + { sourceId, column: col.name, error: e }, + 'Failed to fetch map keys for column', + ); + } + }), + ); + } - // Any column still missing at the abort means the key list is incomplete. - if ( - signal.aborted && - Object.keys(mapKeysResults).length < mapColumns.length - ) { - skippedStages.push('mapAttributeKeys'); - } - finishedStages.add('mapAttributeKeys'); - // No keys and no abort means stage 4 has nothing to sample. No keys because - // of the abort means it was never sampled, matching the backstop snapshot. - if (Object.keys(mapKeysResults).length === 0) { - if (signal.aborted && mapColumns.length > 0) { - skippedStages.push('mapAttributeValues'); + // Any column still missing at the abort means the key list is incomplete. + if ( + signal.aborted && + Object.keys(mapKeysResults).length < mapColumns.length + ) { + skippedStages.push('mapAttributeKeys'); + } + finishedStages.add('mapAttributeKeys'); + // No keys and no abort means stage 4 has nothing to sample. No keys because + // of the abort means it was never sampled, matching the backstop snapshot. + if (Object.keys(mapKeysResults).length === 0) { + if (signal.aborted && mapColumns.length > 0) { + skippedStages.push('mapAttributeValues'); + } + finishedStages.add('mapAttributeValues'); } - finishedStages.add('mapAttributeValues'); - } - // ── 3. Low-cardinality column value sampling ────────────────────────── - const lowCardinalityValues: Record = {}; - - if (lcColumns.length > 0 && !signal.aborted) { - try { - const results = await metadata.getAllKeyValues({ - databaseName, - tableName, - keyExpressions: lcColumns.map(col => col.name), - maxValuesPerKey: MAX_LC_VALUES, - connectionId, - metadataMVs, - dateRange, - timestampValueExpression, - signal, - }); - for (const { key, value } of results) { - if (value.length > 0) { - lowCardinalityValues[key] = value.map(v => v.toString()); + // ── 3. Low-cardinality column value sampling ────────────────────────── + const lowCardinalityValues: Record = {}; + + if (lcColumns.length > 0 && !signal.aborted) { + try { + const results = await metadata.getAllKeyValues({ + databaseName, + tableName, + keyExpressions: lcColumns.map(col => col.name), + maxValuesPerKey: MAX_LC_VALUES, + connectionId, + metadataMVs, + dateRange, + timestampValueExpression, + signal, + }); + for (const { key, value } of results) { + if (value.length > 0) { + lowCardinalityValues[key] = value.map(v => v.toString()); + } } + } catch { + // Skip columns where value sampling fails } - } catch { - // Skip columns where value sampling fails } - } - // getAllKeyValues drops aborted chunks instead of throwing, so any result - // returned after the abort may be missing keys. - if (lcColumns.length > 0 && signal.aborted) { - skippedStages.push('lowCardinalityValues'); - } - if (Object.keys(lowCardinalityValues).length > 0) { - meta.lowCardinalityValues = lowCardinalityValues; - } - finishedStages.add('lowCardinalityValues'); - - // ── 4. Map attribute value sampling (best-effort) ───────────────────── - if (Object.keys(mapKeysResults).length > 0 && !signal.aborted) { - const mapAttributeValues: Record = {}; - - const keyExprs: string[] = []; - for (const [colName, keys] of Object.entries(mapKeysResults)) { - for (const key of keys.slice(0, MAX_MAP_KEYS_TO_SAMPLE)) { - // Map keys come from ClickHouse data (customer telemetry) so they can - // contain arbitrary characters, including single quotes. Escape as a - // SQL string literal — `SqlString.escape` returns a fully-quoted, - // safely-escaped value — before embedding in the key expression. - keyExprs.push(`${colName}[${SqlString.escape(key)}]`); - } + // getAllKeyValues drops aborted chunks instead of throwing, so any result + // returned after the abort may be missing keys. + if (lcColumns.length > 0 && signal.aborted) { + skippedStages.push('lowCardinalityValues'); + } + if (Object.keys(lowCardinalityValues).length > 0) { + meta.lowCardinalityValues = lowCardinalityValues; } + finishedStages.add('lowCardinalityValues'); + + // ── 4. Map attribute value sampling (best-effort) ───────────────────── + if (Object.keys(mapKeysResults).length > 0 && !signal.aborted) { + const mapAttributeValues: Record = {}; + + const keyExprs: string[] = []; + for (const [colName, keys] of Object.entries(mapKeysResults)) { + for (const key of keys.slice(0, MAX_MAP_KEYS_TO_SAMPLE)) { + // Map keys come from ClickHouse data (customer telemetry) so they can + // contain arbitrary characters, including single quotes. Escape as a + // SQL string literal — `SqlString.escape` returns a fully-quoted, + // safely-escaped value — before embedding in the key expression. + keyExprs.push(`${colName}[${SqlString.escape(key)}]`); + } + } - try { - const results = await metadata.getAllKeyValues({ - databaseName, - tableName, - keyExpressions: keyExprs, - maxValuesPerKey: MAX_MAP_KEY_VALUES, - connectionId, - metadataMVs, - dateRange, - timestampValueExpression, - signal, - }); - for (const { key, value } of results) { - if (value.length > 0) { - mapAttributeValues[key] = value.map(v => v.toString()); + try { + const results = await metadata.getAllKeyValues({ + databaseName, + tableName, + keyExpressions: keyExprs, + maxValuesPerKey: MAX_MAP_KEY_VALUES, + connectionId, + metadataMVs, + dateRange, + timestampValueExpression, + signal, + }); + for (const { key, value } of results) { + if (value.length > 0) { + mapAttributeValues[key] = value.map(v => v.toString()); + } } + } catch { + // Best-effort; skip on failure } - } catch { - // Best-effort; skip on failure - } - if (signal.aborted) { + if (signal.aborted) { + skippedStages.push('mapAttributeValues'); + } + if (Object.keys(mapAttributeValues).length > 0) { + meta.mapAttributeValues = mapAttributeValues; + } + } else if (Object.keys(mapKeysResults).length > 0) { + // Signal was already aborted before we started this stage skippedStages.push('mapAttributeValues'); } - if (Object.keys(mapAttributeValues).length > 0) { - meta.mapAttributeValues = mapAttributeValues; - } - } else if (Object.keys(mapKeysResults).length > 0) { - // Signal was already aborted before we started this stage - skippedStages.push('mapAttributeValues'); - } - finishedStages.add('mapAttributeValues'); - - // ── 5. Metric name + unit + description sampling ────────────────────── - // For metric sources, sample distinct MetricName values per discoverable - // kind (including the non-queryable summary kind, so agents know those - // metrics exist) so the agent has a starter list without needing a - // follow-up call to clickstack_list_metrics for the common case - // (<= 20 metrics/kind). - // Defensively check for MetricUnit / MetricDescription columns: they - // exist on the standard OTel Collector schema but a custom metric table - // may not declare them. - if (metricKinds.length > 0 && !signal.aborted) { - const metricNames: Record = {}; - let sampledKinds = 0; - await Promise.all( - metricKinds.map(async ({ kind, kindTableName }) => { - try { - const samples = await sampleMetricNamesWithLookback({ - metadata, - clickhouseClient, - databaseName, - tableName: kindTableName, - connectionId, - now, - timestampValueExpression, - signal, - }); - // The sampler returns [] rather than throwing when aborted, so an - // empty result after the abort may just be a cut-short lookback. - if (samples.length > 0 || !signal.aborted) sampledKinds++; - if (samples.length > 0) { - metricNames[kind] = samples; - meta.metricNames = metricNames; + finishedStages.add('mapAttributeValues'); + + // ── 5. Metric name + unit + description sampling ────────────────────── + // For metric sources, sample distinct MetricName values per discoverable + // kind (including the non-queryable summary kind, so agents know those + // metrics exist) so the agent has a starter list without needing a + // follow-up call to clickstack_list_metrics for the common case + // (<= 20 metrics/kind). + // Defensively check for MetricUnit / MetricDescription columns: they + // exist on the standard OTel Collector schema but a custom metric table + // may not declare them. + if (metricKinds.length > 0 && !signal.aborted) { + const metricNames: Record = {}; + let sampledKinds = 0; + await Promise.all( + metricKinds.map(async ({ kind, kindTableName }) => { + try { + const samples = await sampleMetricNamesWithLookback({ + metadata, + clickhouseClient, + databaseName, + tableName: kindTableName, + connectionId, + now, + timestampValueExpression, + signal, + }); + // The sampler returns [] rather than throwing when aborted, so an + // empty result after the abort may just be a cut-short lookback. + if (samples.length > 0 || !signal.aborted) sampledKinds++; + if (samples.length > 0) { + metricNames[kind] = samples; + meta.metricNames = metricNames; + } + } catch (e) { + logger.warn( + { sourceId, kind, error: e }, + 'Failed to sample metric names for kind', + ); } - } catch (e) { - logger.warn( - { sourceId, kind, error: e }, - 'Failed to sample metric names for kind', - ); - } - }), - ); - if (signal.aborted && sampledKinds < metricKinds.length) { + }), + ); + if (signal.aborted && sampledKinds < metricKinds.length) { + skippedStages.push('metricNames'); + } + } else if (metricKinds.length > 0) { skippedStages.push('metricNames'); } - } else if (metricKinds.length > 0) { - skippedStages.push('metricNames'); - } - finishedStages.add('metricNames'); - - progress.partial = skippedStages.length > 0; - return formatDescribeResult({ - sourceId, - meta, - isMetricSource, - skippedStages, + finishedStages.add('metricNames'); + + progress.partial = skippedStages.length > 0; + return formatDescribeResult({ + sourceId, + meta, + isMetricSource, + skippedStages, + }); }); } diff --git a/packages/api/src/mcp/tools/sources/listMetrics.ts b/packages/api/src/mcp/tools/sources/listMetrics.ts index 5c373197b2..d2770ce43d 100644 --- a/packages/api/src/mcp/tools/sources/listMetrics.ts +++ b/packages/api/src/mcp/tools/sources/listMetrics.ts @@ -1,16 +1,17 @@ +import { withQueryAttribution } from '@hyperdx/common-utils/dist/clickhouse/node'; import { getMetadata } from '@hyperdx/common-utils/dist/core/metadata'; import { SourceKind } from '@hyperdx/common-utils/dist/types'; import { trace } from '@opentelemetry/api'; import { z } from 'zod'; -import { ClickhouseClient } from '@/clickhouse'; -import { getConnectionById } from '@/controllers/connection'; import { getSource } from '@/controllers/sources'; import { READ_ONLY_TOOL } from '@/mcp/tools/annotations'; import { parseTimeRange } from '@/mcp/tools/query/helpers'; +import { getMcpClickhouseClient } from '@/mcp/tools/trace/shared'; import type { ToolRegistrar } from '@/mcp/tools/types'; import { mcpServerError, mcpUserError } from '@/mcp/utils/errors'; import { MCP_TOOL_TIMEOUT_MS, runWithTimeout } from '@/mcp/utils/timeout'; +import { recordMcpSource } from '@/mcp/utils/tracing'; import logger from '@/utils/logger'; import { KIND_TIMED_OUT_ERROR, scanKindsForPage } from './listMetricsPage'; @@ -107,6 +108,7 @@ async function listMetricsImpl( `Source "${input.sourceId}" not found. Call clickstack_list_sources to see available source IDs.`, ); } + recordMcpSource(source); if (source.kind !== SourceKind.Metric) { return mcpUserError( `Source "${input.sourceId}" is a "${source.kind}" source, not a metric source. clickstack_list_metrics only works on metric sources — call clickstack_list_sources to find one whose kind is "metric".`, @@ -147,122 +149,119 @@ async function listMetricsImpl( ); } - const connection = await getConnectionById( - teamId, - source.connection.toString(), - true, - ); - if (!connection) { - return mcpUserError(`Connection not found for source "${input.sourceId}".`); - } - - const clickhouseClient = new ClickhouseClient({ - host: connection.host, - username: connection.username, - password: connection.password, - }); - const metadata = getMetadata(clickhouseClient); + return withQueryAttribution({ source: source._id.toString() }, async () => { + const clickhouseClient = await getMcpClickhouseClient( + teamId, + source.connection.toString(), + ); + if (!clickhouseClient) { + return mcpUserError( + `Connection not found for source "${input.sourceId}".`, + ); + } + const metadata = getMetadata(clickhouseClient); - const limit = input.limit ?? DEFAULT_LIMIT; - const databaseName = source.from.databaseName; + const limit = input.limit ?? DEFAULT_LIMIT; + const databaseName = source.from.databaseName; - const connectionId = source.connection.toString(); + const connectionId = source.connection.toString(); - // Kinds before the cursor's kind were already returned on earlier pages. - const scanKinds = requestedKinds.slice(startKindIdx).map((kind, i) => ({ - kind, - afterName: i === 0 ? cursor?.lastName : undefined, - })); - const page = await scanKindsForPage({ - kinds: scanKinds, - limit, - deadlineAt: deadlineAt - ENRICH_RESERVE_MS, - signal, - fetchNames: ({ kind, afterName }, kindSignal) => - fetchMetricNames({ - clickhouseClient, - databaseName, - tableName: source.metricTables[kind]!, - connectionId, - startDate, - endDate, - namePattern: input.namePattern, - afterName, - // One extra row detects that more names remain for this kind. - limit: limit + 1, - maxExecutionSeconds: NAMES_MAX_EXEC_SECONDS, - signal: kindSignal, - }), - }); - for (const failure of page.partialFailure) { - logger.warn( - { sourceId: input.sourceId, kind: failure.kind, error: failure.error }, - 'Failed to list metrics for kind', - ); - } + // Kinds before the cursor's kind were already returned on earlier pages. + const scanKinds = requestedKinds.slice(startKindIdx).map((kind, i) => ({ + kind, + afterName: i === 0 ? cursor?.lastName : undefined, + })); + const page = await scanKindsForPage({ + kinds: scanKinds, + limit, + deadlineAt: deadlineAt - ENRICH_RESERVE_MS, + signal, + fetchNames: ({ kind, afterName }, kindSignal) => + fetchMetricNames({ + clickhouseClient, + databaseName, + tableName: source.metricTables[kind]!, + connectionId, + startDate, + endDate, + namePattern: input.namePattern, + afterName, + // One extra row detects that more names remain for this kind. + limit: limit + 1, + maxExecutionSeconds: NAMES_MAX_EXEC_SECONDS, + signal: kindSignal, + }), + }); + for (const failure of page.partialFailure) { + logger.warn( + { sourceId: input.sourceId, kind: failure.kind, error: failure.error }, + 'Failed to list metrics for kind', + ); + } - const enrichSignal = AbortSignal.any([ - signal, - AbortSignal.timeout(Math.max(0, deadlineAt - Date.now() - 500)), - ]); - const metrics = await enrichEntries({ - entries: page.entries, - clickhouseClient, - metadata, - databaseName, - metricTables: source.metricTables, - connectionId, - startDate, - endDate, - maxExecutionSeconds: ENRICH_MAX_EXEC_SECONDS, - signal: enrichSignal, - }); + const enrichSignal = AbortSignal.any([ + signal, + AbortSignal.timeout(Math.max(0, deadlineAt - Date.now() - 500)), + ]); + const metrics = await enrichEntries({ + entries: page.entries, + clickhouseClient, + metadata, + databaseName, + metricTables: source.metricTables, + connectionId, + startDate, + endDate, + maxExecutionSeconds: ENRICH_MAX_EXEC_SECONDS, + signal: enrichSignal, + }); - const nextCursor = page.next && encodeCursor(page.next); - // Per-kind failures are surfaced so the agent can tell "kind has no - // metrics" apart from "the fetch for that kind failed or timed out". - const partialFailure = page.partialFailure; - trace.getActiveSpan()?.setAttributes({ - 'mcp.list_metrics.kinds_scanned': scanKinds.length, - 'mcp.list_metrics.kinds_failed': partialFailure.length, - 'mcp.list_metrics.kinds_timed_out': partialFailure.filter( - f => f.error === KIND_TIMED_OUT_ERROR, - ).length, - 'mcp.list_metrics.result_count': metrics.length, - }); + const nextCursor = page.next && encodeCursor(page.next); + // Per-kind failures are surfaced so the agent can tell "kind has no + // metrics" apart from "the fetch for that kind failed or timed out". + const partialFailure = page.partialFailure; + trace.getActiveSpan()?.setAttributes({ + 'mcp.list_metrics.kinds_scanned': scanKinds.length, + 'mcp.list_metrics.kinds_failed': partialFailure.length, + 'mcp.list_metrics.kinds_timed_out': partialFailure.filter( + f => f.error === KIND_TIMED_OUT_ERROR, + ).length, + 'mcp.list_metrics.result_count': metrics.length, + }); - const responseObj: Record = { - metrics, - ...(nextCursor && { nextCursor }), - ...(metrics.some(m => m.kind === 'summary') && { - summaryNote: - 'summary metrics cannot be queried with clickstack_timeseries / clickstack_table — ' + - "use clickstack_sql against the table in the source's metricTables.summary.", - }), - ...(partialFailure.length > 0 && { - partialFailure, - hint: - 'Listing failed for some metric kinds — results may be incomplete. ' + - 'Retry the call; if the failure persists, narrow startTime/endTime or pin a single `kind`.', - }), - ...(metrics.length === 0 && - partialFailure.length === 0 && { + const responseObj: Record = { + metrics, + ...(nextCursor && { nextCursor }), + ...(metrics.some(m => m.kind === 'summary') && { + summaryNote: + 'summary metrics cannot be queried with clickstack_timeseries / clickstack_table — ' + + "use clickstack_sql against the table in the source's metricTables.summary.", + }), + ...(partialFailure.length > 0 && { + partialFailure, hint: - 'No metrics matched. Try widening the time window (startTime/endTime), ' + - 'removing the namePattern filter, or omitting `kind` to scan every populated metric table.', + 'Listing failed for some metric kinds — results may be incomplete. ' + + 'Retry the call; if the failure persists, narrow startTime/endTime or pin a single `kind`.', }), - usage: - 'Pass `metricType` + `metricName` from each entry to clickstack_timeseries / clickstack_table ' + - '(summary metrics excepted — query those with clickstack_sql). ' + - 'For per-metric attribute keys and sampled values, call clickstack_describe_metric.', - }; + ...(metrics.length === 0 && + partialFailure.length === 0 && { + hint: + 'No metrics matched. Try widening the time window (startTime/endTime), ' + + 'removing the namePattern filter, or omitting `kind` to scan every populated metric table.', + }), + usage: + 'Pass `metricType` + `metricName` from each entry to clickstack_timeseries / clickstack_table ' + + '(summary metrics excepted — query those with clickstack_sql). ' + + 'For per-metric attribute keys and sampled values, call clickstack_describe_metric.', + }; - return { - content: [ - { - type: 'text' as const, - text: JSON.stringify(responseObj, null, 2), - }, - ], - }; + return { + content: [ + { + type: 'text' as const, + text: JSON.stringify(responseObj, null, 2), + }, + ], + }; + }); } diff --git a/packages/api/src/mcp/tools/trace/breakdown.ts b/packages/api/src/mcp/tools/trace/breakdown.ts index e480c2a566..90442f5c66 100644 --- a/packages/api/src/mcp/tools/trace/breakdown.ts +++ b/packages/api/src/mcp/tools/trace/breakdown.ts @@ -16,10 +16,9 @@ * tools (table / timeseries / search) can't express the TraceId subselect. */ import type { ChSql } from '@hyperdx/common-utils/dist/clickhouse'; +import { withQueryAttribution } from '@hyperdx/common-utils/dist/clickhouse/node'; import { getMetadata } from '@hyperdx/common-utils/dist/core/metadata'; -import { SourceKind } from '@hyperdx/common-utils/dist/types'; -import { getSource } from '@/controllers/sources'; import { READ_ONLY_TOOL } from '@/mcp/tools/annotations'; import { clickHouseErrorResult, @@ -34,8 +33,8 @@ import { } from '@/mcp/tools/trace/breakdownSchema'; import { durationDivisor, - getMcpClickhouseClient, inclusiveTimeFilter, + loadTraceSource, mcpQuerySettings, } from '@/mcp/tools/trace/shared'; import type { ToolRegistrar } from '@/mcp/tools/types'; @@ -78,78 +77,71 @@ export function registerTraceBreakdown({ } const { startDate, endDate } = timeRange; - const source = await getSource(teamId.toString(), input.sourceId); - if (!source) { - return mcpUserError( - `Source not found: ${input.sourceId}. Call clickstack_list_sources to find available source IDs.`, - ); - } - if (source.kind !== SourceKind.Trace) { - return mcpUserError( - `Source ${input.sourceId} is kind="${source.kind}". clickstack_trace_top_time_consuming_operations requires a source of kind="trace".`, - ); - } - - const clickhouseClient = await getMcpClickhouseClient( + const loaded = await loadTraceSource( teamId.toString(), - source.connection.toString(), + input.sourceId, + 'clickstack_trace_top_time_consuming_operations', ); - if (!clickhouseClient) { - return mcpUserError( - `Connection not found for source: ${input.sourceId}`, - ); - } + if ('isError' in loaded) return loaded; + const { source, clickhouseClient } = loaded; - const timeFilterParams = { - connectionId: source.connection.toString(), - databaseName: source.from.databaseName, - tableName: source.from.tableName, - timestampValueExpression: source.timestampValueExpression, - metadata: getMetadata(clickhouseClient), - }; - const renderTimeFilter = (start: number, end: number) => - inclusiveTimeFilter(timeFilterParams, new Date(start), new Date(end)); - let parentTimeFilter: ChSql; - let childTimeFilter: ChSql; - try { - [parentTimeFilter, childTimeFilter] = await Promise.all([ - renderTimeFilter(startDate.getTime(), endDate.getTime()), - // Widen the child window by 60s on each side to catch children - // that started slightly before / ended slightly after the parent - // sampling window. - renderTimeFilter( - startDate.getTime() - 60_000, - endDate.getTime() + 60_000, - ), - ]); - } catch (e) { - // Only the DESCRIBE lookups run here; they don't depend on the window or - // filters, so neither suffix applies. - return clickHouseErrorResult(e, 'Failed to compute breakdown'); - } + return withQueryAttribution( + { source: source._id.toString() }, + async () => { + const timeFilterParams = { + connectionId: source.connection.toString(), + databaseName: source.from.databaseName, + tableName: source.from.tableName, + timestampValueExpression: source.timestampValueExpression, + metadata: getMetadata(clickhouseClient), + }; + const renderTimeFilter = (start: number, end: number) => + inclusiveTimeFilter( + timeFilterParams, + new Date(start), + new Date(end), + ); + let parentTimeFilter: ChSql; + let childTimeFilter: ChSql; + try { + [parentTimeFilter, childTimeFilter] = await Promise.all([ + renderTimeFilter(startDate.getTime(), endDate.getTime()), + // Widen the child window by 60s on each side to catch children + // that started slightly before / ended slightly after the parent + // sampling window. + renderTimeFilter( + startDate.getTime() - 60_000, + endDate.getTime() + 60_000, + ), + ]); + } catch (e) { + // Only the DESCRIBE lookups run here; they don't depend on the window or + // filters, so neither suffix applies. + return clickHouseErrorResult(e, 'Failed to compute breakdown'); + } - // Source-configured SQL expressions. These are trusted (set by the - // team admin in source config) and we substitute them into the - // generated SQL. - const traceIdExpr = source.traceIdExpression; - const spanNameExpr = source.spanNameExpression ?? "''"; - const serviceNameExpr = source.serviceNameExpression ?? "''"; - const durationExpr = source.durationExpression; - const divisor = durationDivisor(source.durationPrecision); - const dbName = source.from.databaseName; - const tableName = source.from.tableName; + // Source-configured SQL expressions. These are trusted (set by the + // team admin in source config) and we substitute them into the + // generated SQL. + const traceIdExpr = source.traceIdExpression; + const spanNameExpr = source.spanNameExpression ?? "''"; + const serviceNameExpr = source.serviceNameExpression ?? "''"; + const durationExpr = source.durationExpression; + const divisor = durationDivisor(source.durationPrecision); + const dbName = source.from.databaseName; + const tableName = source.from.tableName; - // Build the SQL. parentFilter is SQL only (documented); time bounds - // and topN/maxParentTraces are parameterized. - // The CTE selects distinct parent TraceIds. The outer query joins - // back, excludes the parent rows themselves (so we measure child - // contribution), and aggregates by service+operation. - const minDurationClause = - input.minParentDurationMs != null - ? `AND ${durationExpr} >= ({minParentDurationStored:Float64})` - : ''; + // Build the SQL. parentFilter is SQL only (documented); time bounds + // and topN/maxParentTraces are parameterized. + // The CTE selects distinct parent TraceIds. The outer query joins + // back, excludes the parent rows themselves (so we measure child + // contribution), and aggregates by service+operation. + const minDurationClause = + input.minParentDurationMs != null + ? `AND ${durationExpr} >= ({minParentDurationStored:Float64})` + : ''; - const sql = ` + const sql = ` WITH parent_traces AS ( SELECT DISTINCT ${traceIdExpr} AS _trace_id FROM \`${dbName}\`.\`${tableName}\` @@ -175,95 +167,99 @@ ORDER BY total_time_ms DESC LIMIT {topN:UInt32} `; - const params: Record = { - ...parentTimeFilter.params, - ...childTimeFilter.params, - maxParentTraces: input.maxParentTraces, - topN: input.topN, - divisor, - }; - if (input.minParentDurationMs != null) { - // Stored duration is divisor × ms. - params.minParentDurationStored = input.minParentDurationMs * divisor; - } + const params: Record = { + ...parentTimeFilter.params, + ...childTimeFilter.params, + maxParentTraces: input.maxParentTraces, + topN: input.topN, + divisor, + }; + if (input.minParentDurationMs != null) { + // Stored duration is divisor × ms. + params.minParentDurationStored = + input.minParentDurationMs * divisor; + } - type Row = { - service: string; - operation: string; - total_time_ms: number | string; - calls: number | string; - in_parents: number | string; - p50_ms: number | string; - p99_ms: number | string; - }; + type Row = { + service: string; + operation: string; + total_time_ms: number | string; + calls: number | string; + in_parents: number | string; + p50_ms: number | string; + p99_ms: number | string; + }; - let rows: Row[]; - try { - const result = await clickhouseClient.query({ - query: sql, - query_params: params, - format: 'JSON', - connectionId: source.connection.toString(), - clickhouse_settings: mcpQuerySettings(source.querySettings), - }); - const json = (await ( - result as { json: () => Promise<{ data: Row[] }> } - ).json()) ?? { data: [] }; - rows = json.data ?? []; - } catch (e) { - return clickHouseErrorResult(e, { - prefix: 'Failed to compute breakdown', - suffix: isQueryOutOfTime(e) - ? TIMEOUT_SUFFIX - : isServerError(e) - ? undefined - : INVALID_FILTER_SUFFIX, - }); - } + let rows: Row[]; + try { + const result = await clickhouseClient.query({ + query: sql, + query_params: params, + format: 'JSON', + connectionId: source.connection.toString(), + clickhouse_settings: mcpQuerySettings(source.querySettings), + }); + const json = (await ( + result as { json: () => Promise<{ data: Row[] }> } + ).json()) ?? { data: [] }; + rows = json.data ?? []; + } catch (e) { + return clickHouseErrorResult(e, { + prefix: 'Failed to compute breakdown', + suffix: isQueryOutOfTime(e) + ? TIMEOUT_SUFFIX + : isServerError(e) + ? undefined + : INVALID_FILTER_SUFFIX, + }); + } - // Normalise numerics — ClickHouse JSON sometimes returns strings for - // 64-bit integers. Cast everything to Number and compute share of - // total time at the same time. - const operations = rows.map(r => ({ - service: r.service, - operation: r.operation, - totalTimeMs: Number(r.total_time_ms), - calls: Number(r.calls), - inParents: Number(r.in_parents), - p50Ms: Number(r.p50_ms), - p99Ms: Number(r.p99_ms), - })); - const grandTotalMs = operations.reduce( - (acc, r) => acc + r.totalTimeMs, - 0, - ); - const operationsWithShare = operations.map(r => ({ - ...r, - shareOfTotalTime: grandTotalMs > 0 ? r.totalTimeMs / grandTotalMs : 0, - })); + // Normalise numerics — ClickHouse JSON sometimes returns strings for + // 64-bit integers. Cast everything to Number and compute share of + // total time at the same time. + const operations = rows.map(r => ({ + service: r.service, + operation: r.operation, + totalTimeMs: Number(r.total_time_ms), + calls: Number(r.calls), + inParents: Number(r.in_parents), + p50Ms: Number(r.p50_ms), + p99Ms: Number(r.p99_ms), + })); + const grandTotalMs = operations.reduce( + (acc, r) => acc + r.totalTimeMs, + 0, + ); + const operationsWithShare = operations.map(r => ({ + ...r, + shareOfTotalTime: + grandTotalMs > 0 ? r.totalTimeMs / grandTotalMs : 0, + })); - const output = { - summary: { - parentFilter: input.parentFilter, - startTime: startDate.toISOString(), - endTime: endDate.toISOString(), - minParentDurationMs: input.minParentDurationMs ?? null, - operationsReturned: operationsWithShare.length, - topN: input.topN, - grandTotalTimeMs: grandTotalMs, - hint: - operationsWithShare.length === 0 - ? 'No matching parent traces, or all matching traces had no other spans. Widen the time window, relax the parentFilter, or lower minParentDurationMs.' - : "Operations are sorted by total_time_ms (sum of child Duration). shareOfTotalTime is each row's share of the cumulative child time across all matching traces.", - }, - operations: operationsWithShare, - }; + const output = { + summary: { + parentFilter: input.parentFilter, + startTime: startDate.toISOString(), + endTime: endDate.toISOString(), + minParentDurationMs: input.minParentDurationMs ?? null, + operationsReturned: operationsWithShare.length, + topN: input.topN, + grandTotalTimeMs: grandTotalMs, + hint: + operationsWithShare.length === 0 + ? 'No matching parent traces, or all matching traces had no other spans. Widen the time window, relax the parentFilter, or lower minParentDurationMs.' + : "Operations are sorted by total_time_ms (sum of child Duration). shareOfTotalTime is each row's share of the cumulative child time across all matching traces.", + }, + operations: operationsWithShare, + }; - return { - content: [ - { type: 'text' as const, text: JSON.stringify(output, null, 2) }, - ], - }; + return { + content: [ + { type: 'text' as const, text: JSON.stringify(output, null, 2) }, + ], + }; + }, + ); }, ); } diff --git a/packages/api/src/mcp/tools/trace/shared.ts b/packages/api/src/mcp/tools/trace/shared.ts index 499980c7cf..b3b017dc58 100644 --- a/packages/api/src/mcp/tools/trace/shared.ts +++ b/packages/api/src/mcp/tools/trace/shared.ts @@ -1,14 +1,20 @@ import type { ChSql } from '@hyperdx/common-utils/dist/clickhouse'; import type { Metadata } from '@hyperdx/common-utils/dist/core/metadata'; import { timeFilterExpr } from '@hyperdx/common-utils/dist/core/renderChartConfig'; -import type { QuerySettings } from '@hyperdx/common-utils/dist/types'; +import { + type QuerySettings, + SourceKind, +} from '@hyperdx/common-utils/dist/types'; import { ClickhouseClient } from '@/clickhouse'; import { getConnectionById } from '@/controllers/connection'; +import { getSource } from '@/controllers/sources'; import { MCP_CLICKHOUSE_SETTINGS, MCP_REQUEST_TIMEOUT, } from '@/mcp/tools/query/helpers'; +import { mcpUserError } from '@/mcp/utils/errors'; +import { recordMcpSource } from '@/mcp/utils/tracing'; export function durationDivisor(precision: number): number { // durationPrecision is the number of decimal digits in the stored value. @@ -31,6 +37,37 @@ export async function getMcpClickhouseClient( }); } +/** + * Load the trace source a trace tool reads and a client for its connection, or + * the error to return. + */ +export async function loadTraceSource( + teamId: string, + sourceId: string, + toolName: string, +) { + const source = await getSource(teamId, sourceId); + if (!source) { + return mcpUserError( + `Source not found: ${sourceId}. Call clickstack_list_sources to find available source IDs.`, + ); + } + recordMcpSource(source); + if (source.kind !== SourceKind.Trace) { + return mcpUserError( + `Source ${sourceId} is kind="${source.kind}". ${toolName} requires a source of kind="trace".`, + ); + } + const clickhouseClient = await getMcpClickhouseClient( + teamId, + source.connection.toString(), + ); + if (!clickhouseClient) { + return mcpUserError(`Connection not found for source: ${sourceId}`); + } + return { source, clickhouseClient }; +} + // MCP settings win over source.querySettings so a source can't relax the // max_execution_time ceiling, or the readonly guard that blocks DDL/DML // injected through raw SQL filters. diff --git a/packages/api/src/mcp/tools/trace/waterfall.ts b/packages/api/src/mcp/tools/trace/waterfall.ts index 46d9125a1c..96d58ea68c 100644 --- a/packages/api/src/mcp/tools/trace/waterfall.ts +++ b/packages/api/src/mcp/tools/trace/waterfall.ts @@ -1,4 +1,5 @@ import type { ChSql } from '@hyperdx/common-utils/dist/clickhouse'; +import { withQueryAttribution } from '@hyperdx/common-utils/dist/clickhouse/node'; import { getMetadata } from '@hyperdx/common-utils/dist/core/metadata'; import { pickBucketTimestampColumn } from '@hyperdx/common-utils/dist/core/utils'; import { @@ -22,11 +23,13 @@ import { durationDivisor, getMcpClickhouseClient, inclusiveTimeFilter, + loadTraceSource, mcpQuerySettings, type TimeFilterParams, } from '@/mcp/tools/trace/shared'; import type { ToolRegistrar } from '@/mcp/tools/types'; import { mcpUserError } from '@/mcp/utils/errors'; +import { recordMcpSource } from '@/mcp/utils/tracing'; import logger from '@/utils/logger'; // ─── Schema ────────────────────────────────────────────────────────────────── @@ -343,215 +346,207 @@ export function registerTraceWaterfall({ } const { startDate, endDate } = timeRange; - const source = await getSource(teamId.toString(), input.sourceId); - if (!source) { - return mcpUserError( - `Source not found: ${input.sourceId}. Call clickstack_list_sources to find available source IDs.`, - ); - } - if (source.kind !== SourceKind.Trace) { - return mcpUserError( - `Source ${input.sourceId} is kind="${source.kind}". clickstack_trace_waterfall requires a source of kind="trace".`, - ); - } - - const clickhouseClient = await getMcpClickhouseClient( + const loaded = await loadTraceSource( teamId.toString(), - source.connection.toString(), + input.sourceId, + 'clickstack_trace_waterfall', ); - if (!clickhouseClient) { - return mcpUserError( - `Connection not found for source: ${input.sourceId}`, - ); - } - const metadata = getMetadata(clickhouseClient); - - const traceIdExpr = source.traceIdExpression; - const spanIdExpr = source.spanIdExpression; - const parentSpanIdExpr = source.parentSpanIdExpression; - const spanNameExpr = source.spanNameExpression; - const spanKindExpr = source.spanKindExpression; - const durationExpr = source.durationExpression; - const tsExpr = source.timestampValueExpression; - const ts = await resolveTimestampExprs({ - timestampValueExpression: tsExpr, - metadata, - databaseName: source.from.databaseName, - tableName: source.from.tableName, - connectionId: source.connection.toString(), - }); - const serviceNameExpr = source.serviceNameExpression ?? "''"; - const statusCodeExpr = source.statusCodeExpression ?? "''"; - const statusMessageExpr = source.statusMessageExpression ?? "''"; - const attrsExpr = source.eventAttributesExpression ?? 'map()'; - const divisor = durationDivisor(source.durationPrecision); - - let pickedTraceId = input.traceId; - if (!pickedTraceId) { - // Compose pickFilter with the pickBy-specific filter when needed. - let effectiveFilter = input.pickFilter; - let effectiveLanguage = input.pickFilterLanguage; - if (input.pickBy === 'first_error') { - // Statuses are typically stored as enum strings. Use SQL so we don't - // depend on lucene mapping the raw column name. - const errFilter = `${statusCodeExpr} = 'STATUS_CODE_ERROR'`; - if (effectiveFilter && effectiveLanguage === 'sql') { - effectiveFilter = `(${effectiveFilter}) AND (${errFilter})`; - } else if (effectiveFilter && effectiveLanguage === 'lucene') { - // Run lucene filter inside parens, AND with raw SQL. - // Easiest: convert to sql by AND-joining a fresh sql-only condition - // means we'd have to render lucene first. Compromise: tell user - // to express the error filter in pickFilter directly. - effectiveFilter = `(${effectiveFilter}) AND StatusCode:STATUS_CODE_ERROR`; - effectiveLanguage = 'lucene'; - } else { - effectiveFilter = errFilter; - effectiveLanguage = 'sql'; - } - } - - const orderBy = - input.pickBy === 'slowest' - ? `max(${durationExpr}) DESC` - : input.pickBy === 'first_error' - ? `min(${ts.eventTs}) ASC` - : `max(${ts.eventTs}) DESC`; - - const pickConfig: BuilderChartConfigWithDateRange = { - displayType: DisplayType.Table, - select: [ - { - aggFn: 'count' as const, - valueExpression: '', - alias: 'span_count', - aggCondition: '', - aggConditionLanguage: 'sql' as const, - }, - ], - from: source.from, - where: effectiveFilter, - whereLanguage: effectiveLanguage, - connection: source.connection.toString(), - timestampValueExpression: tsExpr, - implicitColumnExpression: source.implicitColumnExpression, - groupBy: traceIdExpr, - orderBy, - limit: { limit: 1 }, - dateRange: [startDate, endDate], - }; - - let pickResult: { data?: Array> } = {}; - try { - pickResult = (await clickhouseClient.queryChartConfig({ - config: pickConfig as ChartConfigWithDateRange, + if ('isError' in loaded) return loaded; + const { source, clickhouseClient } = loaded; + + return withQueryAttribution( + { source: source._id.toString() }, + async () => { + const metadata = getMetadata(clickhouseClient); + + const traceIdExpr = source.traceIdExpression; + const spanIdExpr = source.spanIdExpression; + const parentSpanIdExpr = source.parentSpanIdExpression; + const spanNameExpr = source.spanNameExpression; + const spanKindExpr = source.spanKindExpression; + const durationExpr = source.durationExpression; + const tsExpr = source.timestampValueExpression; + const ts = await resolveTimestampExprs({ + timestampValueExpression: tsExpr, metadata, - querySettings: source.querySettings, - opts: { clickhouse_settings: MCP_CLICKHOUSE_SETTINGS }, - })) as { data?: Array> }; - } catch (e) { - return clickHouseErrorResult(e, 'Failed to pick a trace'); - } - - const firstRow = pickResult.data?.[0]; - if (!firstRow) { - return { - content: [ - { - type: 'text' as const, - text: JSON.stringify( - { - result: null, - hint: - 'No traces matched. Widen the time range, relax pickFilter, or ' + - 'pass a specific traceId.', - pickFilter: effectiveFilter, - pickBy: input.pickBy, - }, - null, - 2, - ), - }, - ], - }; - } - // The grouped traceId column is rendered into the result with the - // expression text as its alias. Locate it by stripping non-data keys. - const candidate = Object.entries(firstRow).find( - ([k]) => k !== 'span_count' && k !== '__hdx_time_bucket', - ); - if (!candidate || candidate[1] == null) { - return mcpUserError( - `Picker returned a row but no TraceId column was found. Row keys: ${Object.keys(firstRow).join(', ')}`, - ); - } - pickedTraceId = String(candidate[1]); - } - - // The fetch must be time-bounded so ClickHouse prunes partitions, but the - // [startDate, endDate] window is often too narrow to bound it safely: - // - an explicit traceId with the default window may be for an older - // trace outside it; - // - in auto-pick mode startTime is the *pick* window ("find a slow - // trace from the last hour"), not a fetch bound — the picked trace is - // only guaranteed one span inside it, so its root/tail can lie outside - // and the fetch would return a partial tree. - // Probe the trace's real [min, max] extent in both cases. The only time - // we honor the window verbatim is an explicit traceId WITH an explicit - // startTime, where the caller has deliberately set a fetch bound. - let fetchStart = startDate; - let fetchEnd = endDate; - let probeFailed = false; - let windowClamped = false; - const isAutoPick = input.traceId == null; - const usedDefaultWindow = input.startTime == null; - if (isAutoPick || usedDefaultWindow) { - const [probeStart, probeEnd] = isAutoPick - ? [ - new Date(startDate.getTime() - TRACE_PROBE_AUTOPICK_MARGIN_MS), - new Date(endDate.getTime() + TRACE_PROBE_AUTOPICK_MARGIN_MS), - ] - : [ - new Date(endDate.getTime() - TRACE_PROBE_MAX_LOOKBACK_MS), - endDate, - ]; - try { - const probed = await probeTraceWindow(clickhouseClient, { databaseName: source.from.databaseName, tableName: source.from.tableName, connectionId: source.connection.toString(), - traceIdExpr, - ts, - traceId: pickedTraceId, - startDate: probeStart, - endDate: probeEnd, - querySettings: source.querySettings, }); - if (probed) { - fetchStart = probed.start; - fetchEnd = probed.end; - windowClamped = probed.clamped; + const serviceNameExpr = source.serviceNameExpression ?? "''"; + const statusCodeExpr = source.statusCodeExpression ?? "''"; + const statusMessageExpr = source.statusMessageExpression ?? "''"; + const attrsExpr = source.eventAttributesExpression ?? 'map()'; + const divisor = durationDivisor(source.durationPrecision); + + let pickedTraceId = input.traceId; + if (!pickedTraceId) { + // Compose pickFilter with the pickBy-specific filter when needed. + let effectiveFilter = input.pickFilter; + let effectiveLanguage = input.pickFilterLanguage; + if (input.pickBy === 'first_error') { + // Statuses are typically stored as enum strings. Use SQL so we don't + // depend on lucene mapping the raw column name. + const errFilter = `${statusCodeExpr} = 'STATUS_CODE_ERROR'`; + if (effectiveFilter && effectiveLanguage === 'sql') { + effectiveFilter = `(${effectiveFilter}) AND (${errFilter})`; + } else if (effectiveFilter && effectiveLanguage === 'lucene') { + // Run lucene filter inside parens, AND with raw SQL. + // Easiest: convert to sql by AND-joining a fresh sql-only condition + // means we'd have to render lucene first. Compromise: tell user + // to express the error filter in pickFilter directly. + effectiveFilter = `(${effectiveFilter}) AND StatusCode:STATUS_CODE_ERROR`; + effectiveLanguage = 'lucene'; + } else { + effectiveFilter = errFilter; + effectiveLanguage = 'sql'; + } + } + + const orderBy = + input.pickBy === 'slowest' + ? `max(${durationExpr}) DESC` + : input.pickBy === 'first_error' + ? `min(${ts.eventTs}) ASC` + : `max(${ts.eventTs}) DESC`; + + const pickConfig: BuilderChartConfigWithDateRange = { + displayType: DisplayType.Table, + select: [ + { + aggFn: 'count' as const, + valueExpression: '', + alias: 'span_count', + aggCondition: '', + aggConditionLanguage: 'sql' as const, + }, + ], + from: source.from, + where: effectiveFilter, + whereLanguage: effectiveLanguage, + connection: source.connection.toString(), + timestampValueExpression: tsExpr, + implicitColumnExpression: source.implicitColumnExpression, + groupBy: traceIdExpr, + orderBy, + limit: { limit: 1 }, + dateRange: [startDate, endDate], + }; + + let pickResult: { data?: Array> } = {}; + try { + pickResult = (await clickhouseClient.queryChartConfig({ + config: pickConfig as ChartConfigWithDateRange, + metadata, + querySettings: source.querySettings, + opts: { clickhouse_settings: MCP_CLICKHOUSE_SETTINGS }, + })) as { data?: Array> }; + } catch (e) { + return clickHouseErrorResult(e, 'Failed to pick a trace'); + } + + const firstRow = pickResult.data?.[0]; + if (!firstRow) { + return { + content: [ + { + type: 'text' as const, + text: JSON.stringify( + { + result: null, + hint: + 'No traces matched. Widen the time range, relax pickFilter, or ' + + 'pass a specific traceId.', + pickFilter: effectiveFilter, + pickBy: input.pickBy, + }, + null, + 2, + ), + }, + ], + }; + } + // The grouped traceId column is rendered into the result with the + // expression text as its alias. Locate it by stripping non-data keys. + const candidate = Object.entries(firstRow).find( + ([k]) => k !== 'span_count' && k !== '__hdx_time_bucket', + ); + if (!candidate || candidate[1] == null) { + return mcpUserError( + `Picker returned a row but no TraceId column was found. Row keys: ${Object.keys(firstRow).join(', ')}`, + ); + } + pickedTraceId = String(candidate[1]); + } + + // The fetch must be time-bounded so ClickHouse prunes partitions, but the + // [startDate, endDate] window is often too narrow to bound it safely: + // - an explicit traceId with the default window may be for an older + // trace outside it; + // - in auto-pick mode startTime is the *pick* window ("find a slow + // trace from the last hour"), not a fetch bound — the picked trace is + // only guaranteed one span inside it, so its root/tail can lie outside + // and the fetch would return a partial tree. + // Probe the trace's real [min, max] extent in both cases. The only time + // we honor the window verbatim is an explicit traceId WITH an explicit + // startTime, where the caller has deliberately set a fetch bound. + let fetchStart = startDate; + let fetchEnd = endDate; + let probeFailed = false; + let windowClamped = false; + const isAutoPick = input.traceId == null; + const usedDefaultWindow = input.startTime == null; + if (isAutoPick || usedDefaultWindow) { + const [probeStart, probeEnd] = isAutoPick + ? [ + new Date( + startDate.getTime() - TRACE_PROBE_AUTOPICK_MARGIN_MS, + ), + new Date(endDate.getTime() + TRACE_PROBE_AUTOPICK_MARGIN_MS), + ] + : [ + new Date(endDate.getTime() - TRACE_PROBE_MAX_LOOKBACK_MS), + endDate, + ]; + try { + const probed = await probeTraceWindow(clickhouseClient, { + databaseName: source.from.databaseName, + tableName: source.from.tableName, + connectionId: source.connection.toString(), + traceIdExpr, + ts, + traceId: pickedTraceId, + startDate: probeStart, + endDate: probeEnd, + querySettings: source.querySettings, + }); + if (probed) { + fetchStart = probed.start; + fetchEnd = probed.end; + windowClamped = probed.clamped; + } + } catch (e) { + // probeTraceWindow returns null (not throws) for the empty case, so a + // throw means the probe genuinely didn't complete. The fallback window + // will likely also come back empty, so flag it for an accurate hint. + logger.warn( + { + error: e, + teamId, + traceId: pickedTraceId, + sourceId: input.sourceId, + }, + 'trace waterfall extent probe failed', + ); + probeFailed = true; + } } - } catch (e) { - // probeTraceWindow returns null (not throws) for the empty case, so a - // throw means the probe genuinely didn't complete. The fallback window - // will likely also come back empty, so flag it for an accurate hint. - logger.warn( - { - error: e, - teamId, - traceId: pickedTraceId, - sourceId: input.sourceId, - }, - 'trace waterfall extent probe failed', - ); - probeFailed = true; - } - } - let rows: SpanRow[]; - try { - const timeFilter = await ts.timeFilter(fetchStart, fetchEnd); - const treeQuery = ` + let rows: SpanRow[]; + try { + const timeFilter = await ts.timeFilter(fetchStart, fetchEnd); + const treeQuery = ` SELECT ${spanIdExpr} AS spanId, ${parentSpanIdExpr} AS parentSpanId, @@ -569,149 +564,155 @@ export function registerTraceWaterfall({ ORDER BY ${ts.eventTs} ASC LIMIT {n:UInt32} `; - const result = await clickhouseClient.query({ - query: treeQuery, - query_params: { - ...timeFilter.params, - db: source.from.databaseName, - tbl: source.from.tableName, - tid: pickedTraceId, - n: input.maxSpans + 1, // +1 to detect truncation - divisor, - }, - format: 'JSONEachRow', - connectionId: source.connection.toString(), - clickhouse_settings: mcpQuerySettings(source.querySettings), - }); - rows = - (await (result as { json: () => Promise }).json()) ?? []; - } catch (e) { - return clickHouseErrorResult( - e, - `Failed to fetch trace ${pickedTraceId}`, - ); - } + const result = await clickhouseClient.query({ + query: treeQuery, + query_params: { + ...timeFilter.params, + db: source.from.databaseName, + tbl: source.from.tableName, + tid: pickedTraceId, + n: input.maxSpans + 1, // +1 to detect truncation + divisor, + }, + format: 'JSONEachRow', + connectionId: source.connection.toString(), + clickhouse_settings: mcpQuerySettings(source.querySettings), + }); + rows = + (await (result as { json: () => Promise }).json()) ?? + []; + } catch (e) { + return clickHouseErrorResult( + e, + `Failed to fetch trace ${pickedTraceId}`, + ); + } - const truncated = rows.length > input.maxSpans; - const spans = truncated ? rows.slice(0, input.maxSpans) : rows; - - if (spans.length === 0) { - // The window was often chosen by the probe, not the caller, so the hint - // must reflect what actually ran rather than tell the user to widen a - // window they never set. - let hint: string; - if (probeFailed) { - hint = - 'TraceId picked, but the timestamp probe failed (likely a query ' + - 'timeout) and no spans were found in the fallback window. Retry, ' + - 'or pass an explicit startTime/endTime around when the trace ran.'; - } else if (input.traceId && usedDefaultWindow) { - hint = - 'No spans found for this traceId within the last 90 days (the ' + - 'probe ceiling). If the trace is older, pass an explicit startTime ' + - '(ISO 8601) covering when it ran.'; - } else { - hint = - 'TraceId picked, but no spans exist in the time window. The trace ' + - 'may have spans outside startTime/endTime — widen the window.'; - } - return { - content: [ - { - type: 'text' as const, - text: JSON.stringify( + const truncated = rows.length > input.maxSpans; + const spans = truncated ? rows.slice(0, input.maxSpans) : rows; + + if (spans.length === 0) { + // The window was often chosen by the probe, not the caller, so the hint + // must reflect what actually ran rather than tell the user to widen a + // window they never set. + let hint: string; + if (probeFailed) { + hint = + 'TraceId picked, but the timestamp probe failed (likely a query ' + + 'timeout) and no spans were found in the fallback window. Retry, ' + + 'or pass an explicit startTime/endTime around when the trace ran.'; + } else if (input.traceId && usedDefaultWindow) { + hint = + 'No spans found for this traceId within the last 90 days (the ' + + 'probe ceiling). If the trace is older, pass an explicit startTime ' + + '(ISO 8601) covering when it ran.'; + } else { + hint = + 'TraceId picked, but no spans exist in the time window. The trace ' + + 'may have spans outside startTime/endTime — widen the window.'; + } + return { + content: [ { - result: null, - traceId: pickedTraceId, - hint, + type: 'text' as const, + text: JSON.stringify( + { + result: null, + traceId: pickedTraceId, + hint, + }, + null, + 2, + ), }, - null, - 2, - ), - }, - ], - }; - } + ], + }; + } - const tree = buildPreOrderTree(spans); - const root = tree.find(s => s.depth === 0) ?? tree[0]; - const totalDuration = Math.max(...spans.map(s => s.durationMs)); - - // When the extent probe failed we fell back to the default window, which - // may not cover the whole trace — so spans that came back could be a - // partial tree with a false root. Warn rather than present it as complete. - const probeNote = probeFailed - ? 'The span-extent probe failed (likely a query timeout), so the fetch ' + - 'used the default window and this tree may be incomplete. Retry, or ' + - 'pass an explicit startTime/endTime around when the trace ran.' - : undefined; - - // The fetch window is capped so a wide-spread trace can't scan the world; - // that cap can drop the earliest spans of a trace that ran longer than - // the cap. Say so rather than promising the whole tree. - const windowNote = windowClamped - ? 'This trace spans longer than the fetch-window cap, so the earliest ' + - 'spans may be missing. Pass an explicit startTime covering the ' + - "trace's start to see the full tree." - : undefined; - - type LogRow = { - timestamp: string; - severityText: string; - body: string; - serviceName: string; - spanId: string; - }; - let correlatedLogs: LogRow[] | undefined; - let logsTruncated = false; - let logsNote: string | undefined; - if (input.includeLogs && source.logSourceId) { - const logSource = await getSource( - teamId.toString(), - source.logSourceId, - ); - if (!logSource) { - logsNote = `logSourceId ${source.logSourceId} configured but source not found`; - } else if (logSource.kind !== SourceKind.Log) { - logsNote = `logSourceId ${source.logSourceId} is kind="${logSource.kind}", not "log"`; - } else { - const logTraceIdExpr = logSource.traceIdExpression ?? 'TraceId'; - const logSpanIdExpr = logSource.spanIdExpression ?? "''"; - const logBodyExpr = logSource.bodyExpression ?? "''"; - const logSevExpr = logSource.severityTextExpression ?? "''"; - const logSvcExpr = logSource.serviceNameExpression ?? "''"; - - // Reuse the same connection only when the log source lives there. - let logClient: ClickhouseClient | null = clickhouseClient; - if ( - logSource.connection.toString() !== source.connection.toString() - ) { - logClient = await getMcpClickhouseClient( + const tree = buildPreOrderTree(spans); + const root = tree.find(s => s.depth === 0) ?? tree[0]; + const totalDuration = Math.max(...spans.map(s => s.durationMs)); + + // When the extent probe failed we fell back to the default window, which + // may not cover the whole trace — so spans that came back could be a + // partial tree with a false root. Warn rather than present it as complete. + const probeNote = probeFailed + ? 'The span-extent probe failed (likely a query timeout), so the fetch ' + + 'used the default window and this tree may be incomplete. Retry, or ' + + 'pass an explicit startTime/endTime around when the trace ran.' + : undefined; + + // The fetch window is capped so a wide-spread trace can't scan the world; + // that cap can drop the earliest spans of a trace that ran longer than + // the cap. Say so rather than promising the whole tree. + const windowNote = windowClamped + ? 'This trace spans longer than the fetch-window cap, so the earliest ' + + 'spans may be missing. Pass an explicit startTime covering the ' + + "trace's start to see the full tree." + : undefined; + + type LogRow = { + timestamp: string; + severityText: string; + body: string; + serviceName: string; + spanId: string; + }; + let correlatedLogs: LogRow[] | undefined; + let logsTruncated = false; + let logsNote: string | undefined; + if (input.includeLogs && source.logSourceId) { + const logSource = await getSource( teamId.toString(), - logSource.connection.toString(), + source.logSourceId, ); - } - if (!logClient) { - logsNote = `connection for log source ${source.logSourceId} not found`; - } - - if (logClient && !logsNote) { - try { - const logTs = await resolveTimestampExprs({ - timestampValueExpression: logSource.timestampValueExpression, - metadata: - logClient === clickhouseClient - ? metadata - : getMetadata(logClient), - databaseName: logSource.from.databaseName, - tableName: logSource.from.tableName, - connectionId: logSource.connection.toString(), - }); - const logTimeFilter = await logTs.timeFilter( - fetchStart, - fetchEnd, - ); - const logsQuery = ` + if (!logSource) { + logsNote = `logSourceId ${source.logSourceId} configured but source not found`; + } else if (logSource.kind !== SourceKind.Log) { + logsNote = `logSourceId ${source.logSourceId} is kind="${logSource.kind}", not "log"`; + } else { + recordMcpSource(logSource); + const logTraceIdExpr = logSource.traceIdExpression ?? 'TraceId'; + const logSpanIdExpr = logSource.spanIdExpression ?? "''"; + const logBodyExpr = logSource.bodyExpression ?? "''"; + const logSevExpr = logSource.severityTextExpression ?? "''"; + const logSvcExpr = logSource.serviceNameExpression ?? "''"; + + // Reuse the same connection only when the log source lives there. + let logClient: ClickhouseClient | null = clickhouseClient; + if ( + logSource.connection.toString() !== source.connection.toString() + ) { + logClient = await getMcpClickhouseClient( + teamId.toString(), + logSource.connection.toString(), + ); + } + if (!logClient) { + logsNote = `connection for log source ${source.logSourceId} not found`; + } + + if (logClient && !logsNote) { + try { + await withQueryAttribution( + { source: logSource._id.toString() }, + async () => { + const logTs = await resolveTimestampExprs({ + timestampValueExpression: + logSource.timestampValueExpression, + metadata: + logClient === clickhouseClient + ? metadata + : getMetadata(logClient), + databaseName: logSource.from.databaseName, + tableName: logSource.from.tableName, + connectionId: logSource.connection.toString(), + }); + const logTimeFilter = await logTs.timeFilter( + fetchStart, + fetchEnd, + ); + const logsQuery = ` SELECT ${logTs.eventTs} AS timestamp, ${logSevExpr} AS severityText, @@ -724,72 +725,78 @@ export function registerTraceWaterfall({ ORDER BY ${logTs.eventTs} ASC LIMIT {n:UInt32} `; - const logResult = await logClient.query({ - query: logsQuery, - query_params: { - ...logTimeFilter.params, - db: logSource.from.databaseName, - tbl: logSource.from.tableName, - tid: pickedTraceId, - n: input.maxLogs + 1, // +1 to detect truncation - }, - format: 'JSONEachRow', - connectionId: logSource.connection.toString(), - clickhouse_settings: mcpQuerySettings(logSource.querySettings), - }); - const allLogs = - (await ( - logResult as { json: () => Promise } - ).json()) ?? []; - logsTruncated = allLogs.length > input.maxLogs; - correlatedLogs = logsTruncated - ? allLogs.slice(0, input.maxLogs) - : allLogs; - } catch (e) { - logsNote = `Failed to fetch correlated logs: ${e instanceof Error ? e.message : String(e)}`; + const logResult = await logClient.query({ + query: logsQuery, + query_params: { + ...logTimeFilter.params, + db: logSource.from.databaseName, + tbl: logSource.from.tableName, + tid: pickedTraceId, + n: input.maxLogs + 1, // +1 to detect truncation + }, + format: 'JSONEachRow', + connectionId: logSource.connection.toString(), + clickhouse_settings: mcpQuerySettings( + logSource.querySettings, + ), + }); + const allLogs = + (await ( + logResult as { json: () => Promise } + ).json()) ?? []; + logsTruncated = allLogs.length > input.maxLogs; + correlatedLogs = logsTruncated + ? allLogs.slice(0, input.maxLogs) + : allLogs; + }, + ); + } catch (e) { + logsNote = `Failed to fetch correlated logs: ${e instanceof Error ? e.message : String(e)}`; + } + } } } - } - } - const output = { - traceId: pickedTraceId, - spanCount: spans.length, - totalDurationMs: totalDuration, - rootSpan: { - spanId: root.spanId, - serviceName: root.serviceName, - spanName: root.spanName, - durationMs: root.durationMs, - statusCode: root.statusCode, + const output = { + traceId: pickedTraceId, + spanCount: spans.length, + totalDurationMs: totalDuration, + rootSpan: { + spanId: root.spanId, + serviceName: root.serviceName, + spanName: root.spanName, + durationMs: root.durationMs, + statusCode: root.statusCode, + }, + spans: tree, + ...(correlatedLogs + ? { + logs: correlatedLogs, + logsCount: correlatedLogs.length, + ...(logsTruncated + ? { + logsNote: `Logs truncated to ${input.maxLogs}. Increase maxLogs (max 1000) if needed.`, + } + : {}), + } + : {}), + ...(logsNote ? { logsNote } : {}), + ...(probeNote ? { probeNote } : {}), + ...(windowNote ? { windowNote } : {}), + ...(truncated + ? { + note: `Result truncated to ${input.maxSpans} spans. Increase maxSpans (max 2000) or narrow the trace if needed.`, + } + : {}), + }; + + return { + content: [ + { type: 'text' as const, text: JSON.stringify(output, null, 2) }, + ], + }; }, - spans: tree, - ...(correlatedLogs - ? { - logs: correlatedLogs, - logsCount: correlatedLogs.length, - ...(logsTruncated - ? { - logsNote: `Logs truncated to ${input.maxLogs}. Increase maxLogs (max 1000) if needed.`, - } - : {}), - } - : {}), - ...(logsNote ? { logsNote } : {}), - ...(probeNote ? { probeNote } : {}), - ...(windowNote ? { windowNote } : {}), - ...(truncated - ? { - note: `Result truncated to ${input.maxSpans} spans. Increase maxSpans (max 2000) or narrow the trace if needed.`, - } - : {}), - }; - - return { - content: [ - { type: 'text' as const, text: JSON.stringify(output, null, 2) }, - ], - }; + ); }, ); } diff --git a/packages/api/src/mcp/utils/tracing.ts b/packages/api/src/mcp/utils/tracing.ts index 48ae9d9b3f..11871a876b 100644 --- a/packages/api/src/mcp/utils/tracing.ts +++ b/packages/api/src/mcp/utils/tracing.ts @@ -1,4 +1,7 @@ +import { AsyncLocalStorage } from 'node:async_hooks'; + import { withQueryAttribution } from '@hyperdx/common-utils/dist/clickhouse/node'; +import type { SourceKind } from '@hyperdx/common-utils/dist/types'; import type { CallToolResult } from '@modelcontextprotocol/sdk/types.js'; import { recordOnboardingTaskCompletion } from '@/controllers/user'; @@ -9,6 +12,7 @@ import { getActiveTraceId, getCounter, getHistogram, + type Span, SpanStatusCode, withSpan, } from '@/utils/instrumentation'; @@ -33,6 +37,41 @@ const toolResponseSizeHistogram = getHistogram( }, ); +type ToolSourceScope = { + span: Span; + /** Insertion-ordered, so the first source resolved stays the primary one. */ + sources: Map; +}; + +const toolSourceScope = new AsyncLocalStorage(); + +/** + * Records a source the current MCP tool reads on its tool span, for analytics + * on MCP usage. The tool span is taken from the tool's own scope rather than + * the active span, which may be a nested child span by the time a source is + * resolved. A no-op outside a tool invocation. + */ +export function recordMcpSource(source: { + _id: { toString(): string }; + kind: SourceKind; +}): void { + const scope = toolSourceScope.getStore(); + if (!scope) return; + const id = source._id.toString(); + if (scope.sources.has(id)) return; + scope.sources.set(id, source.kind); + if (scope.sources.size === 1) { + scope.span.setAttribute('mcp.source.id', id); + scope.span.setAttribute('mcp.source.kind', source.kind); + } + scope.span.setAttribute('mcp.source.ids', [...scope.sources.keys()]); +} + +/** Metric attribute value: the primary source's kind, never its id. */ +function getSourceKindLabel(scope: ToolSourceScope): SourceKind | 'none' { + return scope.sources.values().next().value ?? 'none'; +} + /** UTF-8 bytes of a tool result's text blocks: what the model has to read. */ function getResponseSizeBytes( content: { text?: string }[] | undefined, @@ -96,6 +135,7 @@ export function withToolTracing( `mcp.tool.${toolName}`, async span => { const startTime = Date.now(); + const sourceScope: ToolSourceScope = { span, sources: new Map() }; span.setAttribute('mcp.tool.name', toolName); span.setAttribute('mcp.team.id', context.teamId); span.setAttribute('mcp.user.id', context.userId); @@ -117,9 +157,10 @@ export function withToolTracing( label: toolName, trace: getActiveTraceId(), }, - async () => handler(args), + async () => toolSourceScope.run(sourceScope, () => handler(args)), ); const durationMs = Date.now() - startTime; + const sourceKind = getSourceKindLabel(sourceScope); const responseBytes = getResponseSizeBytes(result.content); span.setAttribute('mcp.tool.response.size_bytes', responseBytes); toolResponseSizeHistogram.record(responseBytes, { @@ -144,6 +185,7 @@ export function withToolTracing( toolErrorCounter.add(1, { tool: toolName, error_category: errorCategory, + source_kind: sourceKind, }); logger.warn( { @@ -166,16 +208,24 @@ export function withToolTracing( } span.setAttribute('mcp.tool.duration_ms', durationMs); - toolDurationHistogram.record(durationMs, { tool: toolName }); + toolDurationHistogram.record(durationMs, { + tool: toolName, + source_kind: sourceKind, + }); return result; } catch (err) { const durationMs = Date.now() - startTime; + const sourceKind = getSourceKindLabel(sourceScope); span.setAttribute('mcp.tool.duration_ms', durationMs); span.setAttribute('mcp.tool.error_category', 'server'); - toolDurationHistogram.record(durationMs, { tool: toolName }); + toolDurationHistogram.record(durationMs, { + tool: toolName, + source_kind: sourceKind, + }); toolErrorCounter.add(1, { tool: toolName, error_category: 'server', + source_kind: sourceKind, }); logger.error(