From da28fb3a732bebd6d2042142dfae6952acdcb49f Mon Sep 17 00:00:00 2001 From: Brandon Pereira Date: Mon, 2 Nov 2026 15:33:11 -0600 Subject: [PATCH 1/3] feat(mcp): record which source each MCP tool call reads The more analytics we have on MCP usage, the better. Tools now set mcp.source.id, mcp.source.kind and mcp.source.ids on the tool span, add a low-cardinality source_kind attribute to the tool duration and error metrics, and tag the queries that read a source with it in log_comment. --- .changeset/mcp-source-attribution.md | 6 + agent_docs/architecture.md | 5 + .../mcp/__tests__/sourceAttribution.test.ts | 195 +++++ .../api/src/mcp/__tests__/tracing.test.ts | 177 +++- packages/api/src/mcp/tools/query/denoise.ts | 51 +- .../api/src/mcp/tools/query/eventDeltas.ts | 77 +- packages/api/src/mcp/tools/query/helpers.ts | 51 +- .../src/mcp/tools/query/runEventPatterns.ts | 55 +- .../mcp/tools/query/serviceHealth/index.ts | 213 ++--- .../mcp/tools/query/serviceHealth/source.ts | 2 + packages/api/src/mcp/tools/query/sql.ts | 2 + .../src/mcp/tools/sources/describeMetric.ts | 391 ++++----- .../src/mcp/tools/sources/describeSource.ts | 519 ++++++------ .../api/src/mcp/tools/sources/listMetrics.ts | 223 ++--- packages/api/src/mcp/tools/trace/breakdown.ts | 294 ++++--- packages/api/src/mcp/tools/trace/waterfall.ts | 796 +++++++++--------- packages/api/src/mcp/utils/tracing.ts | 56 +- 17 files changed, 1821 insertions(+), 1292 deletions(-) create mode 100644 .changeset/mcp-source-attribution.md create mode 100644 packages/api/src/mcp/__tests__/sourceAttribution.test.ts 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..f48feecbe5 --- /dev/null +++ b/packages/api/src/mcp/__tests__/sourceAttribution.test.ts @@ -0,0 +1,195 @@ +// 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 of its own, so only the ambient scope can supply the 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..cf72670e1e 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,19 @@ 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, - }, - }); + 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 +640,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 +691,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..6b72604caa 100644 --- a/packages/api/src/mcp/tools/sources/describeMetric.ts +++ b/packages/api/src/mcp/tools/sources/describeMetric.ts @@ -6,6 +6,7 @@ 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'; @@ -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,218 @@ 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(', ')}.`, + 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 "${input.sourceId}".`, + ); + } - // 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 clickhouseClient = new ClickhouseClient({ + host: connection.host, + username: connection.username, + password: connection.password, }); - } - 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], - }; + 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..f919d009b2 100644 --- a/packages/api/src/mcp/tools/sources/listMetrics.ts +++ b/packages/api/src/mcp/tools/sources/listMetrics.ts @@ -1,3 +1,4 @@ +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'; @@ -11,6 +12,7 @@ import { parseTimeRange } from '@/mcp/tools/query/helpers'; 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 +109,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 +150,126 @@ async function listMetricsImpl( ); } - const connection = await getConnectionById( - teamId, - source.connection.toString(), - true, - ); - if (!connection) { - return mcpUserError(`Connection not found for source "${input.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 "${input.sourceId}".`, + ); + } - const clickhouseClient = new ClickhouseClient({ - host: connection.host, - username: connection.username, - password: connection.password, - }); - const metadata = getMetadata(clickhouseClient); + const clickhouseClient = new ClickhouseClient({ + host: connection.host, + username: connection.username, + password: connection.password, + }); + 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..5884cf2cdc 100644 --- a/packages/api/src/mcp/tools/trace/breakdown.ts +++ b/packages/api/src/mcp/tools/trace/breakdown.ts @@ -16,6 +16,7 @@ * 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'; @@ -40,6 +41,7 @@ import { } from '@/mcp/tools/trace/shared'; import type { ToolRegistrar } from '@/mcp/tools/types'; import { mcpUserError } from '@/mcp/utils/errors'; +import { recordMcpSource } from '@/mcp/utils/tracing'; // The child stage looks spans up by TraceId, which is not in the sort key, so // it reads every span in the window regardless of how few parents matched. @@ -84,72 +86,80 @@ export function registerTraceBreakdown({ `Source not found: ${input.sourceId}. Call clickstack_list_sources to find available source IDs.`, ); } + recordMcpSource(source); 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( - teamId.toString(), - source.connection.toString(), - ); - if (!clickhouseClient) { - return mcpUserError( - `Connection not found for source: ${input.sourceId}`, - ); - } + return withQueryAttribution( + { source: source._id.toString() }, + async () => { + const clickhouseClient = await getMcpClickhouseClient( + teamId.toString(), + source.connection.toString(), + ); + if (!clickhouseClient) { + return mcpUserError( + `Connection not found for source: ${input.sourceId}`, + ); + } - 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'); - } + 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 +185,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/waterfall.ts b/packages/api/src/mcp/tools/trace/waterfall.ts index 46d9125a1c..44a1298a50 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 { @@ -27,6 +28,7 @@ import { } 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 ────────────────────────────────────────────────────────────────── @@ -349,209 +351,215 @@ export function registerTraceWaterfall({ `Source not found: ${input.sourceId}. Call clickstack_list_sources to find available source IDs.`, ); } + recordMcpSource(source); 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( - teamId.toString(), - source.connection.toString(), - ); - 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'; + return withQueryAttribution( + { source: source._id.toString() }, + async () => { + const clickhouseClient = await getMcpClickhouseClient( + teamId.toString(), + source.connection.toString(), + ); + if (!clickhouseClient) { + return mcpUserError( + `Connection not found for source: ${input.sourceId}`, + ); } - } - - 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, + 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 +577,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 +738,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( From 454d8a836bec9ae34adc1f1ca134ae9095514238 Mon Sep 17 00:00:00 2001 From: Brandon Pereira Date: Tue, 3 Nov 2026 15:24:27 -0600 Subject: [PATCH 2/3] docs(mcp): explain why queryChartConfig calls still get a source scope (review minor) --- packages/api/src/mcp/__tests__/sourceAttribution.test.ts | 3 ++- packages/api/src/mcp/tools/query/helpers.ts | 2 ++ 2 files changed, 4 insertions(+), 1 deletion(-) diff --git a/packages/api/src/mcp/__tests__/sourceAttribution.test.ts b/packages/api/src/mcp/__tests__/sourceAttribution.test.ts index f48feecbe5..a4e994eece 100644 --- a/packages/api/src/mcp/__tests__/sourceAttribution.test.ts +++ b/packages/api/src/mcp/__tests__/sourceAttribution.test.ts @@ -28,7 +28,8 @@ jest.mock('@/controllers/connection', () => ({ })); // 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 of its own, so only the ambient scope can supply the source. +// 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', diff --git a/packages/api/src/mcp/tools/query/helpers.ts b/packages/api/src/mcp/tools/query/helpers.ts index cf72670e1e..99d90a1f5d 100644 --- a/packages/api/src/mcp/tools/query/helpers.ts +++ b/packages/api/src/mcp/tools/query/helpers.ts @@ -609,6 +609,8 @@ export async function runConfigTile( const metadata = getMetadata(clickhouseClient); try { + // 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() }, () => From f108b8c6958738ccd7b304dd1575721c3e7295a8 Mon Sep 17 00:00:00 2001 From: Brandon Pereira Date: Tue, 3 Nov 2026 15:54:51 -0600 Subject: [PATCH 3/3] refactor(mcp): share trace source loading and the metric tools' client lookup Folds the source lookup, kind check and client lookup duplicated between the trace waterfall and breakdown handlers into loadTraceSource, and switches list_metrics and describe_metric to getMcpClickhouseClient (jscpd). --- .../src/mcp/tools/sources/describeMetric.ts | 15 ++----- .../api/src/mcp/tools/sources/listMetrics.ts | 14 ++----- packages/api/src/mcp/tools/trace/breakdown.ts | 34 ++++------------ packages/api/src/mcp/tools/trace/shared.ts | 39 ++++++++++++++++++- packages/api/src/mcp/tools/trace/waterfall.ts | 29 ++++---------- 5 files changed, 61 insertions(+), 70 deletions(-) diff --git a/packages/api/src/mcp/tools/sources/describeMetric.ts b/packages/api/src/mcp/tools/sources/describeMetric.ts index 6b72604caa..3f6796179b 100644 --- a/packages/api/src/mcp/tools/sources/describeMetric.ts +++ b/packages/api/src/mcp/tools/sources/describeMetric.ts @@ -11,14 +11,14 @@ 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, @@ -498,22 +498,15 @@ async function describeMetricImpl( const { startDate, endDate } = timeRange; return withQueryAttribution({ source: source._id.toString() }, async () => { - const connection = await getConnectionById( + const clickhouseClient = await getMcpClickhouseClient( teamId, source.connection.toString(), - true, ); - if (!connection) { + if (!clickhouseClient) { 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(); diff --git a/packages/api/src/mcp/tools/sources/listMetrics.ts b/packages/api/src/mcp/tools/sources/listMetrics.ts index f919d009b2..d2770ce43d 100644 --- a/packages/api/src/mcp/tools/sources/listMetrics.ts +++ b/packages/api/src/mcp/tools/sources/listMetrics.ts @@ -4,11 +4,10 @@ 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'; @@ -151,22 +150,15 @@ async function listMetricsImpl( } return withQueryAttribution({ source: source._id.toString() }, async () => { - const connection = await getConnectionById( + const clickhouseClient = await getMcpClickhouseClient( teamId, source.connection.toString(), - true, ); - if (!connection) { + if (!clickhouseClient) { 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 limit = input.limit ?? DEFAULT_LIMIT; diff --git a/packages/api/src/mcp/tools/trace/breakdown.ts b/packages/api/src/mcp/tools/trace/breakdown.ts index 5884cf2cdc..90442f5c66 100644 --- a/packages/api/src/mcp/tools/trace/breakdown.ts +++ b/packages/api/src/mcp/tools/trace/breakdown.ts @@ -18,9 +18,7 @@ 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, @@ -35,13 +33,12 @@ import { } from '@/mcp/tools/trace/breakdownSchema'; import { durationDivisor, - getMcpClickhouseClient, inclusiveTimeFilter, + loadTraceSource, mcpQuerySettings, } from '@/mcp/tools/trace/shared'; import type { ToolRegistrar } from '@/mcp/tools/types'; import { mcpUserError } from '@/mcp/utils/errors'; -import { recordMcpSource } from '@/mcp/utils/tracing'; // The child stage looks spans up by TraceId, which is not in the sort key, so // it reads every span in the window regardless of how few parents matched. @@ -80,32 +77,17 @@ 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.`, - ); - } - recordMcpSource(source); - 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 loaded = await loadTraceSource( + teamId.toString(), + input.sourceId, + 'clickstack_trace_top_time_consuming_operations', + ); + if ('isError' in loaded) return loaded; + const { source, clickhouseClient } = loaded; return withQueryAttribution( { source: source._id.toString() }, async () => { - const clickhouseClient = await getMcpClickhouseClient( - teamId.toString(), - source.connection.toString(), - ); - if (!clickhouseClient) { - return mcpUserError( - `Connection not found for source: ${input.sourceId}`, - ); - } - const timeFilterParams = { connectionId: source.connection.toString(), databaseName: source.from.databaseName, 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 44a1298a50..96d58ea68c 100644 --- a/packages/api/src/mcp/tools/trace/waterfall.ts +++ b/packages/api/src/mcp/tools/trace/waterfall.ts @@ -23,6 +23,7 @@ import { durationDivisor, getMcpClickhouseClient, inclusiveTimeFilter, + loadTraceSource, mcpQuerySettings, type TimeFilterParams, } from '@/mcp/tools/trace/shared'; @@ -345,31 +346,17 @@ 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.`, - ); - } - recordMcpSource(source); - if (source.kind !== SourceKind.Trace) { - return mcpUserError( - `Source ${input.sourceId} is kind="${source.kind}". clickstack_trace_waterfall requires a source of kind="trace".`, - ); - } + const loaded = await loadTraceSource( + teamId.toString(), + input.sourceId, + 'clickstack_trace_waterfall', + ); + if ('isError' in loaded) return loaded; + const { source, clickhouseClient } = loaded; return withQueryAttribution( { source: source._id.toString() }, async () => { - const clickhouseClient = await getMcpClickhouseClient( - teamId.toString(), - source.connection.toString(), - ); - if (!clickhouseClient) { - return mcpUserError( - `Connection not found for source: ${input.sourceId}`, - ); - } const metadata = getMetadata(clickhouseClient); const traceIdExpr = source.traceIdExpression;