From c346154882a0bbcfe338cf969a03b2116f435c40 Mon Sep 17 00:00:00 2001 From: Brandon Pereira Date: Fri, 11 Sep 2026 16:38:07 -0600 Subject: [PATCH 1/6] fix(mcp): time-bound trace waterfall queries so they prune partitions MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The clickstack_trace_waterfall tool ran TraceId-only span/log queries that scanned the full retention window and timed out. Thread the search window into both queries, add a max_execution_time ceiling, and probe the trace's actual [min, max] span extent so an explicit traceId older than the default window — or one that ran longer than an hour — still resolves in full. The probe also runs for auto-picked traces, so a picked trace whose root predates the window is no longer truncated into a partial tree. The fetch window is width-clamped to the recent tail so a reused/sentinel traceId can't widen the scan back toward the retention edge or stitch unrelated occurrences into one tree. Empty-result hints name the recoverable action (pass an explicit startTime), and probe failures are surfaced rather than silently swallowed. ClickHouse query timeouts are reclassified from user to server errors. --- .changeset/trace-waterfall-time-bound.md | 19 ++ packages/api/src/mcp/__tests__/query.test.ts | 20 ++ .../api/src/mcp/__tests__/trace.int.test.ts | 282 ++++++++++++++++++ packages/api/src/mcp/tools/query/helpers.ts | 32 +- packages/api/src/mcp/tools/trace/waterfall.ts | 174 ++++++++++- 5 files changed, 510 insertions(+), 17 deletions(-) create mode 100644 .changeset/trace-waterfall-time-bound.md diff --git a/.changeset/trace-waterfall-time-bound.md b/.changeset/trace-waterfall-time-bound.md new file mode 100644 index 0000000000..b55df8f6bc --- /dev/null +++ b/.changeset/trace-waterfall-time-bound.md @@ -0,0 +1,19 @@ +--- +"@hyperdx/api": patch +--- + +fix(mcp): time-bound the trace waterfall span/log fetches so they prune +partitions instead of scanning the full retention window. The +`clickstack_trace_waterfall` tool now threads the search window into its span +and correlated-log queries, adds a `max_execution_time` ceiling, and probes a +trace's `[min, max]` span extent so an explicit `traceId` older than the +default window — or a trace that ran longer than an hour — still resolves in +full. The probe also runs for auto-picked traces, so a picked trace whose root +predates the window is no longer truncated into a partial tree. The fetch window +is width-clamped to the recent tail so a reused or sentinel `traceId` (e.g. an +all-zero id from an uninstrumented emitter) can't widen the scan back toward the +retention edge or stitch unrelated occurrences into one tree. Empty-result hints +now name the recoverable action (pass an explicit `startTime`) instead of the +misleading "widen the window", and a probe that fails is surfaced rather than +silently swallowed. ClickHouse query timeouts are also reclassified from `user` +to `server` errors. diff --git a/packages/api/src/mcp/__tests__/query.test.ts b/packages/api/src/mcp/__tests__/query.test.ts index a404a831c3..7befc707be 100644 --- a/packages/api/src/mcp/__tests__/query.test.ts +++ b/packages/api/src/mcp/__tests__/query.test.ts @@ -300,6 +300,15 @@ describe('errorHint', () => { expect(hint).toContain('too many rows'); }); + it('should match TIMEOUT_EXCEEDED errors', () => { + const hint = errorHint( + 'Code: 159. DB::Exception: Timeout exceeded: elapsed 30s (TIMEOUT_EXCEEDED)', + ); + expect(hint).not.toBeNull(); + expect(hint).toContain('execution-time limit'); + expect(hint).toContain('Narrow the time range'); + }); + it('should match SETTING_CONSTRAINT_VIOLATION errors', () => { const hint = errorHint( "Setting max_result_rows shouldn't be greater than 1000. (SETTING_CONSTRAINT_VIOLATION)", @@ -416,6 +425,17 @@ describe('isServerError', () => { expect(isServerError(err)).toBe(false); }); + it('returns true for TIMEOUT_EXCEEDED (max_execution_time overrun)', () => { + // Query timeouts are a resource failure, not a user mistake. + const cause = new ClickHouseError({ + message: 'Timeout exceeded: elapsed 30s', + code: '159', + type: 'TIMEOUT_EXCEEDED', + }); + const err = new Error('query failed', { cause }); + expect(isServerError(err)).toBe(true); + }); + it('returns true for ECONNREFUSED (TCP connection failure)', () => { const err = Object.assign(new Error('connect ECONNREFUSED 127.0.0.1:1'), { code: 'ECONNREFUSED', diff --git a/packages/api/src/mcp/__tests__/trace.int.test.ts b/packages/api/src/mcp/__tests__/trace.int.test.ts index e2187e8932..d3d1b7cab5 100644 --- a/packages/api/src/mcp/__tests__/trace.int.test.ts +++ b/packages/api/src/mcp/__tests__/trace.int.test.ts @@ -432,6 +432,288 @@ describe('MCP Trace Tools', () => { }); }); + // The span/log fetches are time-bounded for partition pruning. + // Pins both edges — the bound is honored, yet an explicit traceId older + // than the default window is still rescued by the first-seen probe. + describe('time-bounding', () => { + const OLD_TRACE_ID = 'dead0000beef1111cafe2222f00d3333'; + const OLD_SVC = 'wf-oldtrace-svc'; + const now = new Date(); + // Two days old — far outside the default 15-minute window, but within the + // probe's 90-day widened range. + const twoDaysAgo = new Date(now.getTime() - 2 * 24 * 60 * 60 * 1000); + + beforeEach(async () => { + await bulkInsertTraces([ + { + Timestamp: twoDaysAgo, + TraceId: OLD_TRACE_ID, + SpanId: 'old_root_span01', + ParentSpanId: '', + SpanName: 'GET /wf-old/root', + SpanKind: 'SPAN_KIND_SERVER', + ServiceName: OLD_SVC, + Duration: 300_000_000, + StatusCode: 'STATUS_CODE_OK', + }, + { + Timestamp: new Date(twoDaysAgo.getTime() + 10), + TraceId: OLD_TRACE_ID, + SpanId: 'old_child_span1', + ParentSpanId: 'old_root_span01', + SpanName: 'wf-old-child', + SpanKind: 'SPAN_KIND_CLIENT', + ServiceName: OLD_SVC, + Duration: 100_000_000, + StatusCode: 'STATUS_CODE_OK', + }, + ]); + }); + + it('resolves an explicit traceId older than the default window via the first-seen probe', async () => { + // No startTime/endTime: the 15-min default excludes this 2-day-old + // trace, so only the probe can rescue it. + const result = await callTool(client, 'clickstack_trace_waterfall', { + sourceId: traceSource._id.toString(), + traceId: OLD_TRACE_ID, + includeLogs: false, + }); + + expect(result.isError).toBeFalsy(); + const output = JSON.parse(getFirstText(result)); + expect(output.traceId).toBe(OLD_TRACE_ID); + expect(output.spanCount).toBe(2); + }); + + it('honors an explicit window that excludes the trace (returns no-spans hint)', async () => { + // Explicit window the old trace falls outside of: no spans come back, + // proving the time predicate is applied rather than ignored. + const result = await callTool(client, 'clickstack_trace_waterfall', { + sourceId: traceSource._id.toString(), + traceId: OLD_TRACE_ID, + includeLogs: false, + startTime: new Date(now.getTime() - 10 * 60 * 1000).toISOString(), + endTime: new Date(now.getTime() + 60 * 1000).toISOString(), + }); + + expect(result.isError).toBeFalsy(); + const output = JSON.parse(getFirstText(result)); + expect(output.result).toBeNull(); + expect(output.hint).toContain('widen the window'); + }); + + it('points an explicit-traceId miss at startTime, not "widen the window"', async () => { + // A traceId that does not exist: the default window is probed back 90 + // days, finds nothing, and the hint must name the recoverable action + // (pass an explicit startTime) rather than the auto-pick advice. + const result = await callTool(client, 'clickstack_trace_waterfall', { + sourceId: traceSource._id.toString(), + traceId: 'ffffffffffffffffffffffffffffffff', + includeLogs: false, + }); + + expect(result.isError).toBeFalsy(); + const output = JSON.parse(getFirstText(result)); + expect(output.result).toBeNull(); + expect(output.hint).toContain('90 days'); + expect(output.hint).toContain('startTime'); + }); + }); + + // A reused/sentinel TraceId (e.g. an all-zero id from an uninstrumented + // emitter) can have occurrences days apart. The probe's [min, max] window + // is width-clamped to the recent tail so the fetch stays bounded — it must + // NOT balloon back toward the 90-day probe range or stitch the old and new + // occurrences into one nonsense tree. + describe('wide-spread TraceId window clamp', () => { + const SENTINEL_TRACE_ID = '00000000000000000000000000000000'; + const SENTINEL_SVC = 'wf-sentinel-svc'; + const now = new Date(); + const tenDaysAgo = new Date(now.getTime() - 10 * 24 * 60 * 60 * 1000); + const fiveMinAgo = new Date(now.getTime() - 5 * 60 * 1000); + + beforeEach(async () => { + await bulkInsertTraces([ + // Old occurrence — outside the 24h clamp off the recent tail. + { + Timestamp: tenDaysAgo, + TraceId: SENTINEL_TRACE_ID, + SpanId: 'sentinel_old_span', + ParentSpanId: '', + SpanName: 'GET /sentinel/old', + SpanKind: 'SPAN_KIND_SERVER', + ServiceName: SENTINEL_SVC, + Duration: 100_000_000, + StatusCode: 'STATUS_CODE_OK', + }, + // Recent occurrence — a distinct 2-span tree. + { + Timestamp: fiveMinAgo, + TraceId: SENTINEL_TRACE_ID, + SpanId: 'sentinel_new_root', + ParentSpanId: '', + SpanName: 'GET /sentinel/new', + SpanKind: 'SPAN_KIND_SERVER', + ServiceName: SENTINEL_SVC, + Duration: 200_000_000, + StatusCode: 'STATUS_CODE_OK', + }, + { + Timestamp: new Date(fiveMinAgo.getTime() + 10), + TraceId: SENTINEL_TRACE_ID, + SpanId: 'sentinel_new_child', + ParentSpanId: 'sentinel_new_root', + SpanName: 'sentinel-new-child', + SpanKind: 'SPAN_KIND_CLIENT', + ServiceName: SENTINEL_SVC, + Duration: 50_000_000, + StatusCode: 'STATUS_CODE_OK', + }, + ]); + }); + + it('clamps the fetch window to the recent tail, excluding the far-past occurrence', async () => { + const result = await callTool(client, 'clickstack_trace_waterfall', { + sourceId: traceSource._id.toString(), + traceId: SENTINEL_TRACE_ID, + includeLogs: false, + }); + + expect(result.isError).toBeFalsy(); + const output = JSON.parse(getFirstText(result)); + // Only the recent 2-span occurrence is returned; the 10-day-old span is + // outside the clamped window, so it is not stitched into the tree. + expect(output.spanCount).toBe(2); + expect(output.spans.map((s: any) => s.spanId)).not.toContain( + 'sentinel_old_span', + ); + expect(output.spans.filter((s: any) => s.depth === 0)).toHaveLength(1); + }); + }); + + // The probe fetches the trace's [min, max] extent, not a fixed lead off + // first-seen — so a trace whose spans span more than an hour is returned + // whole, spans and correlated logs alike. + describe('long-running trace window', () => { + const LONG_TRACE_ID = 'aaaa1111bbbb2222cccc3333dddd4444'; + const LONG_SVC = 'wf-longtrace-svc'; + const now = new Date(); + // Root two days ago; a child span ~2h later. A fixed first-seen + 1h lead + // would truncate that child; the [min, max] probe keeps it. + const rootTs = new Date(now.getTime() - 2 * 24 * 60 * 60 * 1000); + const lateTs = new Date(rootTs.getTime() + 2 * 60 * 60 * 1000); + + beforeEach(async () => { + await bulkInsertTraces([ + { + Timestamp: rootTs, + TraceId: LONG_TRACE_ID, + SpanId: 'long_root_span01', + ParentSpanId: '', + SpanName: 'GET /wf-long/root', + SpanKind: 'SPAN_KIND_SERVER', + ServiceName: LONG_SVC, + Duration: 7_200_000_000_000, // ~2h in ns + StatusCode: 'STATUS_CODE_OK', + }, + { + Timestamp: lateTs, + TraceId: LONG_TRACE_ID, + SpanId: 'long_late_span01', + ParentSpanId: 'long_root_span01', + SpanName: 'wf-long-late-child', + SpanKind: 'SPAN_KIND_CLIENT', + ServiceName: LONG_SVC, + Duration: 100_000_000, + StatusCode: 'STATUS_CODE_OK', + }, + ]); + await bulkInsertLogs([ + { + Timestamp: lateTs, + TraceId: LONG_TRACE_ID, + SpanId: 'long_late_span01', + Body: 'late log line', + ServiceName: LONG_SVC, + SeverityText: 'INFO', + }, + ]); + }); + + it('returns the full tree and a late log for a trace spanning > 1h', async () => { + const result = await callTool(client, 'clickstack_trace_waterfall', { + sourceId: traceSource._id.toString(), + traceId: LONG_TRACE_ID, + }); + + expect(result.isError).toBeFalsy(); + const output = JSON.parse(getFirstText(result)); + expect(output.traceId).toBe(LONG_TRACE_ID); + expect(output.spanCount).toBe(2); + expect(output.spans.filter((s: any) => s.depth === 0)).toHaveLength(1); + // A fixed first-seen + 1h lead would have cut off this late child's log. + expect(output.logsCount).toBe(1); + }); + }); + + // Auto-pick is also probed: a picked trace is only guaranteed one span in + // the default window, so its root/tail can lie outside it. + describe('auto-pick window', () => { + const PICK_TRACE_ID = 'bbbb2222cccc3333dddd4444eeee5555'; + const PICK_SVC = 'wf-autopick-old-svc'; + const now = new Date(); + // Root predates the default 15-min window; a child is inside it. Auto-pick + // finds the trace via the in-window child, but the fetch must cover the + // out-of-window root too. + const rootTs = new Date(now.getTime() - 40 * 60 * 1000); + const childTs = new Date(now.getTime() - 5 * 60 * 1000); + + beforeEach(async () => { + await bulkInsertTraces([ + { + Timestamp: rootTs, + TraceId: PICK_TRACE_ID, + SpanId: 'pick_root_span01', + ParentSpanId: '', + SpanName: 'GET /wf-autopick/root', + SpanKind: 'SPAN_KIND_SERVER', + ServiceName: PICK_SVC, + Duration: 2_400_000_000_000, // 40m in ns → picked as slowest + StatusCode: 'STATUS_CODE_OK', + }, + { + Timestamp: childTs, + TraceId: PICK_TRACE_ID, + SpanId: 'pick_child_span1', + ParentSpanId: 'pick_root_span01', + SpanName: 'wf-autopick-child', + SpanKind: 'SPAN_KIND_CLIENT', + ServiceName: PICK_SVC, + Duration: 100_000_000, + StatusCode: 'STATUS_CODE_OK', + }, + ]); + }); + + it('returns the whole tree when the picked trace root predates the window', async () => { + const result = await callTool(client, 'clickstack_trace_waterfall', { + sourceId: traceSource._id.toString(), + pickFilter: `ServiceName:${PICK_SVC}`, + pickBy: 'slowest', + includeLogs: false, + }); + + expect(result.isError).toBeFalsy(); + const output = JSON.parse(getFirstText(result)); + expect(output.traceId).toBe(PICK_TRACE_ID); + // Without the auto-pick probe, only the in-window child would return, + // producing a false root and a spanCount of 1. + expect(output.spanCount).toBe(2); + expect(output.spans.filter((s: any) => s.depth === 0)).toHaveLength(1); + expect(output.rootSpan.spanId).toBe('pick_root_span01'); + }); + }); + describe('first_error pick mode', () => { const ERROR_TRACE_ID = 'ffff0000eeee1111dddd2222cccc3333'; const now = new Date(); diff --git a/packages/api/src/mcp/tools/query/helpers.ts b/packages/api/src/mcp/tools/query/helpers.ts index 8743c88fb0..b5b7eaa96e 100644 --- a/packages/api/src/mcp/tools/query/helpers.ts +++ b/packages/api/src/mcp/tools/query/helpers.ts @@ -1,4 +1,7 @@ -import { ClickHouseError } from '@clickhouse/client-common'; +import { + ClickHouseError, + type ClickHouseSettings, +} from '@clickhouse/client-common'; import { getMetadata } from '@hyperdx/common-utils/dist/core/metadata'; import { convertToCategoricalChartConfig, @@ -69,19 +72,18 @@ export const SAFE_BODY_EXPR_CHARS = /^[\w.':\[\]\-]+$/; // ─── Safety limits ─────────────────────────────────────────────────────────── /** ClickHouse settings applied to all MCP query-tool executions. - * readonly=2 so max_execution_time can be set - * (readonly=1 rejects all setting changes). */ -const MCP_CLICKHOUSE_SETTINGS = { + * readonly=2 so max_execution_time can be set (readonly=1 rejects it). */ +export const MCP_CLICKHOUSE_SETTINGS: ClickHouseSettings = { max_execution_time: 30, - readonly: 2, -} as const; + readonly: '2', +}; /** - * HTTP request timeout for MCP query-tool ClickHouse clients. - * Set slightly above max_execution_time so ClickHouse can return a clean - * timeout error before the HTTP connection is aborted. + * HTTP request timeout for MCP query-tool ClickHouse clients. Set above + * max_execution_time so ClickHouse returns a clean timeout before the HTTP + * connection is aborted. */ -const MCP_REQUEST_TIMEOUT = 32_000; // 30s query limit + 2s buffer +export const MCP_REQUEST_TIMEOUT = 32_000; // 30s query limit + 2s buffer // ─── Increase top-N cap hint ──────────────────────────────────────────────── @@ -686,6 +688,9 @@ const SERVER_CH_ERROR_TYPES = new Set([ 'SOCKET_TIMEOUT', 'POCO_EXCEPTION', 'ALL_CONNECTION_TRIES_FAILED', + // A query hitting max_execution_time is a resource failure, not a user + // mistake; the `user` default hid these in error views. + 'TIMEOUT_EXCEEDED', ]); /** @@ -903,6 +908,13 @@ export function errorHint(msg: string, error?: unknown): string | null { 'The result row count is too large to serialize back to the agent.' ); } + if (/TIMEOUT_EXCEEDED|Timeout exceeded|max_execution_time/i.test(msg)) { + return ( + 'The query exceeded its execution-time limit. Narrow the time range so ' + + 'ClickHouse can prune partitions, add filters to reduce the rows scanned, ' + + 'or lower the requested LIMIT.' + ); + } if (/TOO_MANY_ROWS_OR_BYTES|RESULT_IS_TOO_LARGE/i.test(msg)) { return ( 'The query returned too many rows. ' + diff --git a/packages/api/src/mcp/tools/trace/waterfall.ts b/packages/api/src/mcp/tools/trace/waterfall.ts index 979a406f73..9678d4b26c 100644 --- a/packages/api/src/mcp/tools/trace/waterfall.ts +++ b/packages/api/src/mcp/tools/trace/waterfall.ts @@ -12,6 +12,8 @@ import { getConnectionById } from '@/controllers/connection'; import { getSource } from '@/controllers/sources'; import { clickHouseErrorResult, + MCP_CLICKHOUSE_SETTINGS, + MCP_REQUEST_TIMEOUT, parseTimeRange, } from '@/mcp/tools/query/helpers'; import type { ToolRegistrar } from '@/mcp/tools/types'; @@ -31,8 +33,10 @@ const traceSchema = z.object({ .optional() .describe( 'Specific TraceId to look up. When provided, the tool fetches every span ' + - 'in this trace and returns them as a parent/child tree. ' + - 'When omitted, the tool auto-picks one trace using pickFilter + pickBy.', + 'in this trace and returns them as a parent/child tree. When omitted, ' + + 'the tool auto-picks one trace using pickFilter + pickBy. With an ' + + 'explicit traceId and no startTime, the tool probes the trace back to ' + + '90 days; pass startTime only to reach a trace older than that.', ), pickFilter: z .string() @@ -65,7 +69,10 @@ const traceSchema = z.object({ .string() .optional() .describe( - 'Start of the search window as ISO 8601. Default: 15 minutes ago.', + 'Start of the search window as ISO 8601. Default: 15 minutes ago for ' + + 'auto-pick. When traceId is set, the window is instead derived from ' + + "the trace's own span timestamps (probed back to 90 days), so pass " + + 'startTime only to look up a trace older than 90 days.', ), endTime: z .string() @@ -106,6 +113,25 @@ const traceSchema = z.object({ type TraceInput = z.infer; +// ─── Safety limits ─────────────────────────────────────────────────────────── + +// Probe-window tuning. The span/log fetch is time-bounded to the +// trace's probed [min, max] extent so ClickHouse prunes partitions. The pads +// absorb clock skew between emitters; they do NOT cover trace duration — the +// [min, max] extent already does. MAX_FETCH_WINDOW clamps the width so a +// reused/sentinel id (e.g. an all-zero TraceId) whose min and max are days +// apart can't widen the scan back toward the retention edge or stitch two +// occurrences into one tree. A trace older than these bounds needs an explicit +// startTime. +const TRACE_PROBE_LEAD_PAD_MS = 1 * 60 * 60 * 1000; +const TRACE_PROBE_TRAIL_PAD_MS = 1 * 60 * 60 * 1000; +const TRACE_PROBE_MAX_LOOKBACK_MS = 90 * 24 * 60 * 60 * 1000; +const TRACE_MAX_FETCH_WINDOW_MS = 24 * 60 * 60 * 1000; + +// Auto-pick guarantees only one span inside the default window; the root can +// predate it, so probe a bounded distance before startDate to cover it. +const TRACE_PROBE_AUTOPICK_LOOKBACK_MS = 6 * 60 * 60 * 1000; + // ─── Helpers ───────────────────────────────────────────────────────────────── type SpanRow = { @@ -164,6 +190,77 @@ function durationDivisor(precision: number): number { return Math.pow(10, Math.max(0, precision - 3)); } +// ClickHouse renders min()/max() over an empty set as the epoch / a zero date. +function parseProbeTimestamp(value: string | undefined): number | null { + if (!value) return null; + const ms = new Date(value).getTime(); + return isNaN(ms) || ms <= 0 ? null : ms; +} + +/** + * Probe a trace's [min, max] span timestamps within [startDate, endDate] and + * return a fetch window covering that extent (padded, then width-clamped). + * Rescues explicit `traceId` lookups older than the default 15-minute window + * and keeps a long trace from being truncated by a fixed lead. The probe is + * time-bounded so it prunes partitions; the caller sets [startDate, endDate]. + * + * Returns null when no span is found, so the caller falls back to its window. + */ +async function probeTraceWindow( + clickhouseClient: ClickhouseClient, + params: { + databaseName: string; + tableName: string; + connectionId: string; + traceIdExpr: string; + tsExpr: string; + traceId: string; + startDate: Date; + endDate: Date; + }, +): Promise<[Date, Date] | null> { + const probeQuery = ` + SELECT + min(${params.tsExpr}) AS firstSeen, + max(${params.tsExpr}) AS lastSeen + FROM {db:Identifier}.{tbl:Identifier} + WHERE ${params.traceIdExpr} = {tid:String} + AND ${params.tsExpr} >= fromUnixTimestamp64Milli({startMs:Int64}) + AND ${params.tsExpr} <= fromUnixTimestamp64Milli({endMs:Int64}) + `; + const result = await clickhouseClient.query({ + query: probeQuery, + query_params: { + db: params.databaseName, + tbl: params.tableName, + tid: params.traceId, + startMs: params.startDate.getTime(), + endMs: params.endDate.getTime(), + }, + format: 'JSONEachRow', + connectionId: params.connectionId, + clickhouse_settings: MCP_CLICKHOUSE_SETTINGS, + }); + const rows = + (await ( + result as { + json: () => Promise<{ firstSeen: string; lastSeen: string }[]>; + } + ).json()) ?? []; + const firstSeen = parseProbeTimestamp(rows[0]?.firstSeen); + const lastSeen = parseProbeTimestamp(rows[0]?.lastSeen); + if (firstSeen == null || lastSeen == null) return null; + + const fetchEnd = lastSeen + TRACE_PROBE_TRAIL_PAD_MS; + // Anchor to the recent tail and clamp the width so a wide-spread/reused + // TraceId can't balloon the fetch back toward the 90-day probe range. + const fetchStart = Math.max( + firstSeen - TRACE_PROBE_LEAD_PAD_MS, + fetchEnd - TRACE_MAX_FETCH_WINDOW_MS, + ); + return [new Date(fetchStart), new Date(fetchEnd)]; +} + // ─── Tool definition ───────────────────────────────────────────────────────── export function registerTraceWaterfall({ @@ -246,6 +343,7 @@ export function registerTraceWaterfall({ host: connection.host, username: connection.username, password: connection.password, + requestTimeout: MCP_REQUEST_TIMEOUT, }); const metadata = getMetadata(clickhouseClient); @@ -363,6 +461,41 @@ export function registerTraceWaterfall({ pickedTraceId = String(candidate[1]); } + // The fetch must be time-bounded so ClickHouse prunes partitions, but the + // default 15-min window is too narrow: an explicit traceId may be older + // than it, and an auto-picked trace is only guaranteed one span inside it + // (its root/tail can lie outside, giving a partial tree). So when the + // window was defaulted, probe the trace's real [min, max] extent instead. + let fetchStart = startDate; + let fetchEnd = endDate; + let probeFailed = false; + const usedDefaultWindow = input.startTime == null; + if (usedDefaultWindow) { + const probeStart = input.traceId + ? new Date(endDate.getTime() - TRACE_PROBE_MAX_LOOKBACK_MS) + : new Date(startDate.getTime() - TRACE_PROBE_AUTOPICK_LOOKBACK_MS); + try { + const probed = await probeTraceWindow(clickhouseClient, { + databaseName: source.from.databaseName, + tableName: source.from.tableName, + connectionId: source.connection.toString(), + traceIdExpr, + tsExpr, + traceId: pickedTraceId, + startDate: probeStart, + endDate, + }); + if (probed) { + [fetchStart, fetchEnd] = probed; + } + } catch { + // 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. + probeFailed = true; + } + } + // ── Step 2: fetch the full span tree ── const treeQuery = ` SELECT @@ -378,6 +511,8 @@ export function registerTraceWaterfall({ ${attrsExpr} AS spanAttributes FROM {db:Identifier}.{tbl:Identifier} WHERE ${traceIdExpr} = {tid:String} + AND ${tsExpr} >= fromUnixTimestamp64Milli({startMs:Int64}) + AND ${tsExpr} <= fromUnixTimestamp64Milli({endMs:Int64}) ORDER BY ${tsExpr} ASC LIMIT {n:UInt32} `; @@ -390,14 +525,15 @@ export function registerTraceWaterfall({ db: source.from.databaseName, tbl: source.from.tableName, tid: pickedTraceId, + startMs: fetchStart.getTime(), + endMs: fetchEnd.getTime(), n: input.maxSpans + 1, // +1 to detect truncation divisor, }, format: 'JSONEachRow', connectionId: source.connection.toString(), clickhouse_settings: { - readonly: '1', - // Per-query timeout matches the rest of the MCP for consistency. + ...MCP_CLICKHOUSE_SETTINGS, ...(source.querySettings ? Object.fromEntries( source.querySettings.map(s => [s.setting, s.value]), @@ -418,6 +554,25 @@ export function registerTraceWaterfall({ 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: [ { @@ -426,7 +581,7 @@ export function registerTraceWaterfall({ { result: null, traceId: pickedTraceId, - hint: 'TraceId picked, but no spans exist in the time window. The trace may have spans outside startTime/endTime — widen the window.', + hint, }, null, 2, @@ -485,6 +640,7 @@ export function registerTraceWaterfall({ host: logConn.host, username: logConn.username, password: logConn.password, + requestTimeout: MCP_REQUEST_TIMEOUT, }); } } @@ -499,6 +655,8 @@ export function registerTraceWaterfall({ ${logSpanIdExpr} AS spanId FROM {db:Identifier}.{tbl:Identifier} WHERE ${logTraceIdExpr} = {tid:String} + AND ${logTsExpr} >= fromUnixTimestamp64Milli({startMs:Int64}) + AND ${logTsExpr} <= fromUnixTimestamp64Milli({endMs:Int64}) ORDER BY ${logTsExpr} ASC LIMIT {n:UInt32} `; @@ -509,12 +667,14 @@ export function registerTraceWaterfall({ db: logSource.from.databaseName, tbl: logSource.from.tableName, tid: pickedTraceId, + startMs: fetchStart.getTime(), + endMs: fetchEnd.getTime(), n: input.maxLogs + 1, // +1 to detect truncation }, format: 'JSONEachRow', connectionId: logSource.connection.toString(), clickhouse_settings: { - readonly: '1', + ...MCP_CLICKHOUSE_SETTINGS, ...(logSource.querySettings ? Object.fromEntries( logSource.querySettings.map(s => [s.setting, s.value]), From d670fccf9810f39b95ed82a41e3fe450e5022abc Mon Sep 17 00:00:00 2001 From: Brandon Pereira Date: Fri, 11 Sep 2026 16:48:59 -0600 Subject: [PATCH 2/6] fix(mcp): address trace-waterfall review feedback - MCP ClickHouse settings now win over source.querySettings so a source can't relax the max_execution_time / readonly ceiling (span + log fetches). - Surface a probeNote when the extent probe fails and the fallback window may return a partial tree, instead of presenting it as complete. - Surface a windowNote when the fetch-window cap drops a wide-spread trace's earliest spans, so the tool no longer silently promises the whole tree. - Drop the newly added "Safety limits" ASCII divider per AGENTS.md. --- .../api/src/mcp/__tests__/trace.int.test.ts | 5 ++ packages/api/src/mcp/tools/trace/waterfall.ts | 56 +++++++++++++++---- 2 files changed, 49 insertions(+), 12 deletions(-) diff --git a/packages/api/src/mcp/__tests__/trace.int.test.ts b/packages/api/src/mcp/__tests__/trace.int.test.ts index d3d1b7cab5..95f5343bc1 100644 --- a/packages/api/src/mcp/__tests__/trace.int.test.ts +++ b/packages/api/src/mcp/__tests__/trace.int.test.ts @@ -588,6 +588,9 @@ describe('MCP Trace Tools', () => { 'sentinel_old_span', ); expect(output.spans.filter((s: any) => s.depth === 0)).toHaveLength(1); + // The clamp dropped earlier spans, so the output must say so rather than + // silently promise the whole tree. + expect(output.windowNote).toContain('fetch-window cap'); }); }); @@ -653,6 +656,8 @@ describe('MCP Trace Tools', () => { expect(output.spans.filter((s: any) => s.depth === 0)).toHaveLength(1); // A fixed first-seen + 1h lead would have cut off this late child's log. expect(output.logsCount).toBe(1); + // An 85-min spread is under the fetch-window cap, so no clamp warning. + expect(output.windowNote).toBeUndefined(); }); }); diff --git a/packages/api/src/mcp/tools/trace/waterfall.ts b/packages/api/src/mcp/tools/trace/waterfall.ts index 9678d4b26c..301daaa264 100644 --- a/packages/api/src/mcp/tools/trace/waterfall.ts +++ b/packages/api/src/mcp/tools/trace/waterfall.ts @@ -113,8 +113,6 @@ const traceSchema = z.object({ type TraceInput = z.infer; -// ─── Safety limits ─────────────────────────────────────────────────────────── - // Probe-window tuning. The span/log fetch is time-bounded to the // trace's probed [min, max] extent so ClickHouse prunes partitions. The pads // absorb clock skew between emitters; they do NOT cover trace duration — the @@ -197,6 +195,8 @@ function parseProbeTimestamp(value: string | undefined): number | null { return isNaN(ms) || ms <= 0 ? null : ms; } +type ProbedWindow = { start: Date; end: Date; clamped: boolean }; + /** * Probe a trace's [min, max] span timestamps within [startDate, endDate] and * return a fetch window covering that extent (padded, then width-clamped). @@ -204,7 +204,10 @@ function parseProbeTimestamp(value: string | undefined): number | null { * and keeps a long trace from being truncated by a fixed lead. The probe is * time-bounded so it prunes partitions; the caller sets [startDate, endDate]. * - * Returns null when no span is found, so the caller falls back to its window. + * `clamped` is true when the trace's extent exceeded TRACE_MAX_FETCH_WINDOW_MS + * and the window was capped to the recent tail — the returned tree may then be + * missing early spans. Returns null when no span is found, so the caller falls + * back to its window. */ async function probeTraceWindow( clickhouseClient: ClickhouseClient, @@ -218,7 +221,7 @@ async function probeTraceWindow( startDate: Date; endDate: Date; }, -): Promise<[Date, Date] | null> { +): Promise { const probeQuery = ` SELECT min(${params.tsExpr}) AS firstSeen, @@ -254,11 +257,14 @@ async function probeTraceWindow( const fetchEnd = lastSeen + TRACE_PROBE_TRAIL_PAD_MS; // Anchor to the recent tail and clamp the width so a wide-spread/reused // TraceId can't balloon the fetch back toward the 90-day probe range. - const fetchStart = Math.max( - firstSeen - TRACE_PROBE_LEAD_PAD_MS, - fetchEnd - TRACE_MAX_FETCH_WINDOW_MS, - ); - return [new Date(fetchStart), new Date(fetchEnd)]; + const unclampedStart = firstSeen - TRACE_PROBE_LEAD_PAD_MS; + const clampFloor = fetchEnd - TRACE_MAX_FETCH_WINDOW_MS; + const fetchStart = Math.max(unclampedStart, clampFloor); + return { + start: new Date(fetchStart), + end: new Date(fetchEnd), + clamped: unclampedStart < clampFloor, + }; } // ─── Tool definition ───────────────────────────────────────────────────────── @@ -469,6 +475,7 @@ export function registerTraceWaterfall({ let fetchStart = startDate; let fetchEnd = endDate; let probeFailed = false; + let windowClamped = false; const usedDefaultWindow = input.startTime == null; if (usedDefaultWindow) { const probeStart = input.traceId @@ -486,7 +493,9 @@ export function registerTraceWaterfall({ endDate, }); if (probed) { - [fetchStart, fetchEnd] = probed; + fetchStart = probed.start; + fetchEnd = probed.end; + windowClamped = probed.clamped; } } catch { // probeTraceWindow returns null (not throws) for the empty case, so a @@ -532,13 +541,15 @@ export function registerTraceWaterfall({ }, format: 'JSONEachRow', connectionId: source.connection.toString(), + // MCP settings win over source.querySettings so a source can't relax + // the max_execution_time / readonly ceiling this tool depends on. clickhouse_settings: { - ...MCP_CLICKHOUSE_SETTINGS, ...(source.querySettings ? Object.fromEntries( source.querySettings.map(s => [s.setting, s.value]), ) : {}), + ...MCP_CLICKHOUSE_SETTINGS, }, }); rows = @@ -595,6 +606,24 @@ export function registerTraceWaterfall({ 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; + // ── Step 3: fetch correlated logs (when logSourceId is configured) ── type LogRow = { timestamp: string; @@ -673,13 +702,14 @@ export function registerTraceWaterfall({ }, format: 'JSONEachRow', connectionId: logSource.connection.toString(), + // MCP settings win over source.querySettings (see span fetch). clickhouse_settings: { - ...MCP_CLICKHOUSE_SETTINGS, ...(logSource.querySettings ? Object.fromEntries( logSource.querySettings.map(s => [s.setting, s.value]), ) : {}), + ...MCP_CLICKHOUSE_SETTINGS, }, }); const allLogs = @@ -721,6 +751,8 @@ export function registerTraceWaterfall({ } : {}), ...(logsNote ? { logsNote } : {}), + ...(probeNote ? { probeNote } : {}), + ...(windowNote ? { windowNote } : {}), ...(truncated ? { note: `Result truncated to ${input.maxSpans} spans. Increase maxSpans (max 2000) or narrow the trace if needed.`, From eb5b4127a024ea50abfa05fb759387dcd75bdbef Mon Sep 17 00:00:00 2001 From: Brandon Pereira Date: Fri, 11 Sep 2026 17:16:44 -0600 Subject: [PATCH 3/6] fix(mcp): address deep-review findings on trace waterfall MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - errorHint: match only real timeouts (TIMEOUT_EXCEEDED / "Timeout exceeded") so a "Setting max_execution_time shouldn't be greater than…" constraint error no longer hijacks the timeout hint. Regression test added. - Wrap tsExpr / logTsExpr in getFirstTimestampValueExpression for every raw-SQL use (probe min/max, tree + log SELECT/WHERE/ORDER BY, picker orderBy) so a multi-column timestampValueExpression (e.g. "EventDate, EventTime") no longer breaks the probe and span queries. - Probe in auto-pick mode even with an explicit startTime: startTime is the pick window there, not a fetch bound, so the picked trace's full extent must still be probed. Only explicit traceId + explicit startTime honors the window verbatim. Integration test added. - Give LONG_TRACE_ID a unique value so it no longer collides with the breakdown suite's TRACE_ID_1 in the never-cleared traces table. --- packages/api/src/mcp/__tests__/query.test.ts | 8 +++ .../api/src/mcp/__tests__/trace.int.test.ts | 50 ++++++++++++++++++- packages/api/src/mcp/tools/query/helpers.ts | 5 +- packages/api/src/mcp/tools/trace/waterfall.ts | 48 ++++++++++++------ 4 files changed, 93 insertions(+), 18 deletions(-) diff --git a/packages/api/src/mcp/__tests__/query.test.ts b/packages/api/src/mcp/__tests__/query.test.ts index 7befc707be..4c6f0b5ba7 100644 --- a/packages/api/src/mcp/__tests__/query.test.ts +++ b/packages/api/src/mcp/__tests__/query.test.ts @@ -318,6 +318,14 @@ describe('errorHint', () => { expect(hint).toContain('constraint'); }); + it('routes a max_execution_time constraint error to the constraint hint, not the timeout hint', () => { + const hint = errorHint( + "Setting max_execution_time shouldn't be greater than 10. (SETTING_CONSTRAINT_VIOLATION)", + ); + expect(hint).toContain('constraint'); + expect(hint).not.toContain('execution-time limit'); + }); + it('should return null for unrecognized errors', () => { const hint = errorHint('Connection refused'); expect(hint).toBeNull(); diff --git a/packages/api/src/mcp/__tests__/trace.int.test.ts b/packages/api/src/mcp/__tests__/trace.int.test.ts index 95f5343bc1..e982430994 100644 --- a/packages/api/src/mcp/__tests__/trace.int.test.ts +++ b/packages/api/src/mcp/__tests__/trace.int.test.ts @@ -598,7 +598,7 @@ describe('MCP Trace Tools', () => { // first-seen — so a trace whose spans span more than an hour is returned // whole, spans and correlated logs alike. describe('long-running trace window', () => { - const LONG_TRACE_ID = 'aaaa1111bbbb2222cccc3333dddd4444'; + const LONG_TRACE_ID = 'a1b2c3d4e5f60000111122223333long'; const LONG_SVC = 'wf-longtrace-svc'; const now = new Date(); // Root two days ago; a child span ~2h later. A fixed first-seen + 1h lead @@ -717,6 +717,54 @@ describe('MCP Trace Tools', () => { expect(output.spans.filter((s: any) => s.depth === 0)).toHaveLength(1); expect(output.rootSpan.spanId).toBe('pick_root_span01'); }); + + it('probes in auto-pick even with an explicit startTime (pick window, not a fetch bound)', async () => { + // Own trace + service so it's isolated from the sibling test (the + // traces table is not cleared between tests). startTime here is the + // pick window; the picked trace's root predates it, so the fetch must + // still probe the trace's real extent rather than honoring startTime. + const PICK2_TRACE_ID = 'cccc3333dddd4444eeee5555ffff6666'; + const PICK2_SVC = 'wf-autopick-window-svc'; + await bulkInsertTraces([ + { + Timestamp: rootTs, + TraceId: PICK2_TRACE_ID, + SpanId: 'pick2_root_span01', + ParentSpanId: '', + SpanName: 'GET /wf-autopick2/root', + SpanKind: 'SPAN_KIND_SERVER', + ServiceName: PICK2_SVC, + Duration: 2_400_000_000_000, + StatusCode: 'STATUS_CODE_OK', + }, + { + Timestamp: childTs, + TraceId: PICK2_TRACE_ID, + SpanId: 'pick2_child_span1', + ParentSpanId: 'pick2_root_span01', + SpanName: 'wf-autopick2-child', + SpanKind: 'SPAN_KIND_CLIENT', + ServiceName: PICK2_SVC, + Duration: 100_000_000, + StatusCode: 'STATUS_CODE_OK', + }, + ]); + + const result = await callTool(client, 'clickstack_trace_waterfall', { + sourceId: traceSource._id.toString(), + pickFilter: `ServiceName:${PICK2_SVC}`, + pickBy: 'slowest', + includeLogs: false, + startTime: new Date(now.getTime() - 10 * 60 * 1000).toISOString(), + endTime: new Date(now.getTime() + 60 * 1000).toISOString(), + }); + + expect(result.isError).toBeFalsy(); + const output = JSON.parse(getFirstText(result)); + expect(output.traceId).toBe(PICK2_TRACE_ID); + expect(output.spanCount).toBe(2); + expect(output.rootSpan.spanId).toBe('pick2_root_span01'); + }); }); describe('first_error pick mode', () => { diff --git a/packages/api/src/mcp/tools/query/helpers.ts b/packages/api/src/mcp/tools/query/helpers.ts index b5b7eaa96e..3dcdcb0d52 100644 --- a/packages/api/src/mcp/tools/query/helpers.ts +++ b/packages/api/src/mcp/tools/query/helpers.ts @@ -908,7 +908,10 @@ export function errorHint(msg: string, error?: unknown): string | null { 'The result row count is too large to serialize back to the agent.' ); } - if (/TIMEOUT_EXCEEDED|Timeout exceeded|max_execution_time/i.test(msg)) { + // Match only real timeouts. A bare `max_execution_time` substring would also + // hijack SETTING_CONSTRAINT_VIOLATION / readonly errors ("Setting + // max_execution_time shouldn't be greater than…"), which need a different fix. + if (/TIMEOUT_EXCEEDED|Timeout exceeded/i.test(msg)) { return ( 'The query exceeded its execution-time limit. Narrow the time range so ' + 'ClickHouse can prune partitions, add filters to reduce the rows scanned, ' + diff --git a/packages/api/src/mcp/tools/trace/waterfall.ts b/packages/api/src/mcp/tools/trace/waterfall.ts index 301daaa264..047d9aa479 100644 --- a/packages/api/src/mcp/tools/trace/waterfall.ts +++ b/packages/api/src/mcp/tools/trace/waterfall.ts @@ -1,4 +1,5 @@ import { getMetadata } from '@hyperdx/common-utils/dist/core/metadata'; +import { getFirstTimestampValueExpression } from '@hyperdx/common-utils/dist/core/utils'; import { type BuilderChartConfigWithDateRange, type ChartConfigWithDateRange, @@ -360,6 +361,11 @@ export function registerTraceWaterfall({ const spanKindExpr = source.spanKindExpression; const durationExpr = source.durationExpression; const tsExpr = source.timestampValueExpression; + // A source's timestampValueExpression may be a comma-separated list (e.g. + // "EventDate, EventTime"). Raw-SQL uses that aggregate or compare on it — + // min()/max() and the WHERE bounds — need a single column; only the + // queryChartConfig picker (which splits it itself) gets the full form. + const tsExprFirst = getFirstTimestampValueExpression(tsExpr); const serviceNameExpr = source.serviceNameExpression ?? "''"; const statusCodeExpr = source.statusCodeExpression ?? "''"; const statusMessageExpr = source.statusMessageExpression ?? "''"; @@ -395,8 +401,8 @@ export function registerTraceWaterfall({ input.pickBy === 'slowest' ? `max(${durationExpr}) DESC` : input.pickBy === 'first_error' - ? `min(${tsExpr}) ASC` - : `max(${tsExpr}) DESC`; + ? `min(${tsExprFirst}) ASC` + : `max(${tsExprFirst}) DESC`; const pickConfig: BuilderChartConfigWithDateRange = { displayType: DisplayType.Table, @@ -468,26 +474,34 @@ export function registerTraceWaterfall({ } // The fetch must be time-bounded so ClickHouse prunes partitions, but the - // default 15-min window is too narrow: an explicit traceId may be older - // than it, and an auto-picked trace is only guaranteed one span inside it - // (its root/tail can lie outside, giving a partial tree). So when the - // window was defaulted, probe the trace's real [min, max] extent instead. + // [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 (usedDefaultWindow) { - const probeStart = input.traceId - ? new Date(endDate.getTime() - TRACE_PROBE_MAX_LOOKBACK_MS) - : new Date(startDate.getTime() - TRACE_PROBE_AUTOPICK_LOOKBACK_MS); + if (isAutoPick || usedDefaultWindow) { + const probeStart = + input.traceId && usedDefaultWindow + ? new Date(endDate.getTime() - TRACE_PROBE_MAX_LOOKBACK_MS) + : new Date(startDate.getTime() - TRACE_PROBE_AUTOPICK_LOOKBACK_MS); try { const probed = await probeTraceWindow(clickhouseClient, { databaseName: source.from.databaseName, tableName: source.from.tableName, connectionId: source.connection.toString(), traceIdExpr, - tsExpr, + tsExpr: tsExprFirst, traceId: pickedTraceId, startDate: probeStart, endDate, @@ -516,13 +530,13 @@ export function registerTraceWaterfall({ ${durationExpr} / {divisor:Float64} AS durationMs, ${statusCodeExpr} AS statusCode, ${statusMessageExpr} AS statusMessage, - ${tsExpr} AS timestamp, + ${tsExprFirst} AS timestamp, ${attrsExpr} AS spanAttributes FROM {db:Identifier}.{tbl:Identifier} WHERE ${traceIdExpr} = {tid:String} - AND ${tsExpr} >= fromUnixTimestamp64Milli({startMs:Int64}) - AND ${tsExpr} <= fromUnixTimestamp64Milli({endMs:Int64}) - ORDER BY ${tsExpr} ASC + AND ${tsExprFirst} >= fromUnixTimestamp64Milli({startMs:Int64}) + AND ${tsExprFirst} <= fromUnixTimestamp64Milli({endMs:Int64}) + ORDER BY ${tsExprFirst} ASC LIMIT {n:UInt32} `; @@ -647,7 +661,9 @@ export function registerTraceWaterfall({ } else { const logTraceIdExpr = logSource.traceIdExpression ?? 'TraceId'; const logSpanIdExpr = logSource.spanIdExpression ?? "''"; - const logTsExpr = logSource.timestampValueExpression; + const logTsExpr = getFirstTimestampValueExpression( + logSource.timestampValueExpression, + ); const logBodyExpr = logSource.bodyExpression ?? "''"; const logSevExpr = logSource.severityTextExpression ?? "''"; const logSvcExpr = logSource.serviceNameExpression ?? "''"; From 7c3060cd5ca857a99f60e944bbb1c038e89b9bd5 Mon Sep 17 00:00:00 2001 From: Brandon Pereira Date: Thu, 1 Oct 2026 09:51:37 -0600 Subject: [PATCH 4/6] fix(mcp): use precise event timestamp and probe past endTime in trace waterfall - Composite timestamp sources (e.g. EventDate, EventTime) now select/order on the highest-precision DateTime column and filter on every configured column so partition pruning still applies. - Auto-pick probes past endTime so a trace tail after the pick window is kept. - Probe applies source.querySettings like the other queries. - clamped is only set when the cap actually cuts off spans. --- .../api/src/mcp/__tests__/trace.int.test.ts | 53 +++++ packages/api/src/mcp/tools/trace/waterfall.ts | 222 +++++++++++------- 2 files changed, 188 insertions(+), 87 deletions(-) diff --git a/packages/api/src/mcp/__tests__/trace.int.test.ts b/packages/api/src/mcp/__tests__/trace.int.test.ts index e982430994..91a9969c2d 100644 --- a/packages/api/src/mcp/__tests__/trace.int.test.ts +++ b/packages/api/src/mcp/__tests__/trace.int.test.ts @@ -765,6 +765,59 @@ describe('MCP Trace Tools', () => { expect(output.spanCount).toBe(2); expect(output.rootSpan.spanId).toBe('pick2_root_span01'); }); + + it('probes past an explicit endTime so a late trace tail is not dropped', async () => { + const PICK3_TRACE_ID = 'dddd4444eeee5555ffff6666aaaa7777'; + const PICK3_SVC = 'wf-autopick-tail-svc'; + const pickRootTs = new Date(now.getTime() - 4 * 60 * 60 * 1000); + // > 1h past endTime, so only the forward probe margin can find it. + const tailTs = new Date(now.getTime() - 2 * 60 * 60 * 1000); + await bulkInsertTraces([ + { + Timestamp: pickRootTs, + TraceId: PICK3_TRACE_ID, + SpanId: 'pick3_root_span01', + ParentSpanId: '', + SpanName: 'GET /wf-autopick3/root', + SpanKind: 'SPAN_KIND_SERVER', + ServiceName: PICK3_SVC, + Duration: 7_200_000_000_000, + StatusCode: 'STATUS_CODE_OK', + }, + { + Timestamp: tailTs, + TraceId: PICK3_TRACE_ID, + SpanId: 'pick3_tail_span1', + ParentSpanId: 'pick3_root_span01', + SpanName: 'wf-autopick3-tail', + SpanKind: 'SPAN_KIND_CLIENT', + ServiceName: PICK3_SVC, + Duration: 100_000_000, + StatusCode: 'STATUS_CODE_OK', + }, + ]); + + const result = await callTool(client, 'clickstack_trace_waterfall', { + sourceId: traceSource._id.toString(), + pickFilter: `ServiceName:${PICK3_SVC}`, + pickBy: 'slowest', + includeLogs: false, + startTime: new Date( + pickRootTs.getTime() - 10 * 60 * 1000, + ).toISOString(), + endTime: new Date( + pickRootTs.getTime() + 10 * 60 * 1000, + ).toISOString(), + }); + + expect(result.isError).toBeFalsy(); + const output = JSON.parse(getFirstText(result)); + expect(output.traceId).toBe(PICK3_TRACE_ID); + expect(output.spanCount).toBe(2); + expect(output.spans.map((s: any) => s.spanId)).toContain( + 'pick3_tail_span1', + ); + }); }); describe('first_error pick mode', () => { diff --git a/packages/api/src/mcp/tools/trace/waterfall.ts b/packages/api/src/mcp/tools/trace/waterfall.ts index 047d9aa479..40a27744c6 100644 --- a/packages/api/src/mcp/tools/trace/waterfall.ts +++ b/packages/api/src/mcp/tools/trace/waterfall.ts @@ -1,9 +1,15 @@ -import { getMetadata } from '@hyperdx/common-utils/dist/core/metadata'; -import { getFirstTimestampValueExpression } from '@hyperdx/common-utils/dist/core/utils'; +import type { ChSql } from '@hyperdx/common-utils/dist/clickhouse'; +import { + getMetadata, + type Metadata, +} from '@hyperdx/common-utils/dist/core/metadata'; +import { timeFilterExpr } from '@hyperdx/common-utils/dist/core/renderChartConfig'; +import { pickBucketTimestampColumn } from '@hyperdx/common-utils/dist/core/utils'; import { type BuilderChartConfigWithDateRange, type ChartConfigWithDateRange, DisplayType, + type QuerySettings, SourceKind, } from '@hyperdx/common-utils/dist/types'; import { z } from 'zod'; @@ -127,9 +133,10 @@ const TRACE_PROBE_TRAIL_PAD_MS = 1 * 60 * 60 * 1000; const TRACE_PROBE_MAX_LOOKBACK_MS = 90 * 24 * 60 * 60 * 1000; const TRACE_MAX_FETCH_WINDOW_MS = 24 * 60 * 60 * 1000; -// Auto-pick guarantees only one span inside the default window; the root can -// predate it, so probe a bounded distance before startDate to cover it. -const TRACE_PROBE_AUTOPICK_LOOKBACK_MS = 6 * 60 * 60 * 1000; +// Auto-pick guarantees only one span inside the pick window; the root can +// predate it and the tail can outlast it, so probe a bounded distance on both +// sides to cover them. +const TRACE_PROBE_AUTOPICK_MARGIN_MS = 6 * 60 * 60 * 1000; // ─── Helpers ───────────────────────────────────────────────────────────────── @@ -196,6 +203,45 @@ function parseProbeTimestamp(value: string | undefined): number | null { return isNaN(ms) || ms <= 0 ? null : ms; } +// MCP settings win over source.querySettings so a source can't relax the +// max_execution_time / readonly ceiling this tool depends on. +function mcpQuerySettings(querySettings: QuerySettings | undefined) { + return { + ...(querySettings + ? Object.fromEntries(querySettings.map(s => [s.setting, s.value])) + : {}), + ...MCP_CLICKHOUSE_SETTINGS, + }; +} + +type TimestampExprs = { + /** Highest-precision DateTime column, for event-time SELECT/ORDER BY. */ + eventTs: string; + /** Inclusive filter over every configured timestamp column, so a + * composite "EventDate, EventTime" source still prunes on its Date column. */ + timeFilter: (start: Date, end: Date) => Promise; +}; + +async function resolveTimestampExprs(params: { + timestampValueExpression: string; + metadata: Metadata; + databaseName: string; + tableName: string; + connectionId: string; +}): Promise { + const eventTs = await pickBucketTimestampColumn(params); + return { + eventTs, + timeFilter: (start, end) => + timeFilterExpr({ + ...params, + dateRange: [start, end], + dateRangeStartInclusive: true, + dateRangeEndInclusive: true, + }), + }; +} + type ProbedWindow = { start: Date; end: Date; clamped: boolean }; /** @@ -217,33 +263,36 @@ async function probeTraceWindow( tableName: string; connectionId: string; traceIdExpr: string; - tsExpr: string; + ts: TimestampExprs; traceId: string; startDate: Date; endDate: Date; + querySettings: QuerySettings | undefined; }, ): Promise { + const timeFilter = await params.ts.timeFilter( + params.startDate, + params.endDate, + ); const probeQuery = ` SELECT - min(${params.tsExpr}) AS firstSeen, - max(${params.tsExpr}) AS lastSeen + min(${params.ts.eventTs}) AS firstSeen, + max(${params.ts.eventTs}) AS lastSeen FROM {db:Identifier}.{tbl:Identifier} WHERE ${params.traceIdExpr} = {tid:String} - AND ${params.tsExpr} >= fromUnixTimestamp64Milli({startMs:Int64}) - AND ${params.tsExpr} <= fromUnixTimestamp64Milli({endMs:Int64}) + AND ${timeFilter.sql} `; const result = await clickhouseClient.query({ query: probeQuery, query_params: { + ...timeFilter.params, db: params.databaseName, tbl: params.tableName, tid: params.traceId, - startMs: params.startDate.getTime(), - endMs: params.endDate.getTime(), }, format: 'JSONEachRow', connectionId: params.connectionId, - clickhouse_settings: MCP_CLICKHOUSE_SETTINGS, + clickhouse_settings: mcpQuerySettings(params.querySettings), }); const rows = (await ( @@ -264,7 +313,8 @@ async function probeTraceWindow( return { start: new Date(fetchStart), end: new Date(fetchEnd), - clamped: unclampedStart < clampFloor, + // The lead pad is slack, so only a floor above firstSeen drops real spans. + clamped: firstSeen < fetchStart, }; } @@ -361,11 +411,13 @@ export function registerTraceWaterfall({ const spanKindExpr = source.spanKindExpression; const durationExpr = source.durationExpression; const tsExpr = source.timestampValueExpression; - // A source's timestampValueExpression may be a comma-separated list (e.g. - // "EventDate, EventTime"). Raw-SQL uses that aggregate or compare on it — - // min()/max() and the WHERE bounds — need a single column; only the - // queryChartConfig picker (which splits it itself) gets the full form. - const tsExprFirst = getFirstTimestampValueExpression(tsExpr); + 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 ?? "''"; @@ -401,8 +453,8 @@ export function registerTraceWaterfall({ input.pickBy === 'slowest' ? `max(${durationExpr}) DESC` : input.pickBy === 'first_error' - ? `min(${tsExprFirst}) ASC` - : `max(${tsExprFirst}) DESC`; + ? `min(${ts.eventTs}) ASC` + : `max(${ts.eventTs}) DESC`; const pickConfig: BuilderChartConfigWithDateRange = { displayType: DisplayType.Table, @@ -491,20 +543,26 @@ export function registerTraceWaterfall({ const isAutoPick = input.traceId == null; const usedDefaultWindow = input.startTime == null; if (isAutoPick || usedDefaultWindow) { - const probeStart = - input.traceId && usedDefaultWindow - ? new Date(endDate.getTime() - TRACE_PROBE_MAX_LOOKBACK_MS) - : new Date(startDate.getTime() - TRACE_PROBE_AUTOPICK_LOOKBACK_MS); + 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, - tsExpr: tsExprFirst, + ts, traceId: pickedTraceId, startDate: probeStart, - endDate, + endDate: probeEnd, + querySettings: source.querySettings, }); if (probed) { fetchStart = probed.start; @@ -520,51 +578,40 @@ export function registerTraceWaterfall({ } // ── Step 2: fetch the full span tree ── - const treeQuery = ` - SELECT - ${spanIdExpr} AS spanId, - ${parentSpanIdExpr} AS parentSpanId, - ${serviceNameExpr} AS serviceName, - ${spanNameExpr} AS spanName, - ${spanKindExpr} AS spanKind, - ${durationExpr} / {divisor:Float64} AS durationMs, - ${statusCodeExpr} AS statusCode, - ${statusMessageExpr} AS statusMessage, - ${tsExprFirst} AS timestamp, - ${attrsExpr} AS spanAttributes - FROM {db:Identifier}.{tbl:Identifier} - WHERE ${traceIdExpr} = {tid:String} - AND ${tsExprFirst} >= fromUnixTimestamp64Milli({startMs:Int64}) - AND ${tsExprFirst} <= fromUnixTimestamp64Milli({endMs:Int64}) - ORDER BY ${tsExprFirst} ASC - LIMIT {n:UInt32} - `; - let rows: SpanRow[]; try { + const timeFilter = await ts.timeFilter(fetchStart, fetchEnd); + const treeQuery = ` + SELECT + ${spanIdExpr} AS spanId, + ${parentSpanIdExpr} AS parentSpanId, + ${serviceNameExpr} AS serviceName, + ${spanNameExpr} AS spanName, + ${spanKindExpr} AS spanKind, + ${durationExpr} / {divisor:Float64} AS durationMs, + ${statusCodeExpr} AS statusCode, + ${statusMessageExpr} AS statusMessage, + ${ts.eventTs} AS timestamp, + ${attrsExpr} AS spanAttributes + FROM {db:Identifier}.{tbl:Identifier} + WHERE ${traceIdExpr} = {tid:String} + AND ${timeFilter.sql} + 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, - startMs: fetchStart.getTime(), - endMs: fetchEnd.getTime(), n: input.maxSpans + 1, // +1 to detect truncation divisor, }, format: 'JSONEachRow', connectionId: source.connection.toString(), - // MCP settings win over source.querySettings so a source can't relax - // the max_execution_time / readonly ceiling this tool depends on. - clickhouse_settings: { - ...(source.querySettings - ? Object.fromEntries( - source.querySettings.map(s => [s.setting, s.value]), - ) - : {}), - ...MCP_CLICKHOUSE_SETTINGS, - }, + clickhouse_settings: mcpQuerySettings(source.querySettings), }); rows = (await (result as { json: () => Promise }).json()) ?? []; @@ -661,9 +708,6 @@ export function registerTraceWaterfall({ } else { const logTraceIdExpr = logSource.traceIdExpression ?? 'TraceId'; const logSpanIdExpr = logSource.spanIdExpression ?? "''"; - const logTsExpr = getFirstTimestampValueExpression( - logSource.timestampValueExpression, - ); const logBodyExpr = logSource.bodyExpression ?? "''"; const logSevExpr = logSource.severityTextExpression ?? "''"; const logSvcExpr = logSource.serviceNameExpression ?? "''"; @@ -691,42 +735,46 @@ export function registerTraceWaterfall({ } if (!logsNote) { - const logsQuery = ` - SELECT - ${logTsExpr} AS timestamp, - ${logSevExpr} AS severityText, - ${logBodyExpr} AS body, - ${logSvcExpr} AS serviceName, - ${logSpanIdExpr} AS spanId - FROM {db:Identifier}.{tbl:Identifier} - WHERE ${logTraceIdExpr} = {tid:String} - AND ${logTsExpr} >= fromUnixTimestamp64Milli({startMs:Int64}) - AND ${logTsExpr} <= fromUnixTimestamp64Milli({endMs:Int64}) - ORDER BY ${logTsExpr} ASC - LIMIT {n:UInt32} - `; 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 = ` + SELECT + ${logTs.eventTs} AS timestamp, + ${logSevExpr} AS severityText, + ${logBodyExpr} AS body, + ${logSvcExpr} AS serviceName, + ${logSpanIdExpr} AS spanId + FROM {db:Identifier}.{tbl:Identifier} + WHERE ${logTraceIdExpr} = {tid:String} + AND ${logTimeFilter.sql} + 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, - startMs: fetchStart.getTime(), - endMs: fetchEnd.getTime(), n: input.maxLogs + 1, // +1 to detect truncation }, format: 'JSONEachRow', connectionId: logSource.connection.toString(), - // MCP settings win over source.querySettings (see span fetch). - clickhouse_settings: { - ...(logSource.querySettings - ? Object.fromEntries( - logSource.querySettings.map(s => [s.setting, s.value]), - ) - : {}), - ...MCP_CLICKHOUSE_SETTINGS, - }, + clickhouse_settings: mcpQuerySettings(logSource.querySettings), }); const allLogs = (await ( From 1690583840adfbf049d5cea051365c7f53a24a93 Mon Sep 17 00:00:00 2001 From: Brandon Pereira Date: Thu, 1 Oct 2026 12:46:54 -0600 Subject: [PATCH 5/6] fix(mcp): bound trace pick query, parse probe extent as epoch ms - Pass MCP_CLICKHOUSE_SETTINGS to the auto-pick query so it hits max_execution_time before the new client request timeout (claude-review major). - Probe returns epoch ms so the window isn't shifted on a non-UTC API host (deep-review P2). - Log probe failures and cover the probeNote fallback with a test. --- .../api/src/mcp/__tests__/trace.int.test.ts | 56 +++++++++++++++++++ packages/api/src/mcp/tools/trace/waterfall.ts | 33 ++++++++--- 2 files changed, 81 insertions(+), 8 deletions(-) diff --git a/packages/api/src/mcp/__tests__/trace.int.test.ts b/packages/api/src/mcp/__tests__/trace.int.test.ts index 91a9969c2d..ef5876a1d0 100644 --- a/packages/api/src/mcp/__tests__/trace.int.test.ts +++ b/packages/api/src/mcp/__tests__/trace.int.test.ts @@ -1,6 +1,7 @@ import { SourceKind } from '@hyperdx/common-utils/dist/types'; import { Client } from '@modelcontextprotocol/sdk/client/index.js'; +import { ClickhouseClient } from '@/clickhouse'; import * as config from '@/config'; import { bulkInsertData, @@ -766,6 +767,61 @@ describe('MCP Trace Tools', () => { expect(output.rootSpan.spanId).toBe('pick2_root_span01'); }); + it('flags a partial tree with probeNote when the extent probe fails', async () => { + const FAIL_TRACE_ID = 'eeee5555ffff6666aaaa7777bbbb8888'; + const FAIL_SVC = 'wf-probe-fail-svc'; + await bulkInsertTraces([ + { + Timestamp: rootTs, + TraceId: FAIL_TRACE_ID, + SpanId: 'fail_root_span01', + ParentSpanId: '', + SpanName: 'GET /wf-probe-fail/root', + SpanKind: 'SPAN_KIND_SERVER', + ServiceName: FAIL_SVC, + Duration: 2_400_000_000_000, + StatusCode: 'STATUS_CODE_OK', + }, + { + Timestamp: childTs, + TraceId: FAIL_TRACE_ID, + SpanId: 'fail_child_span1', + ParentSpanId: 'fail_root_span01', + SpanName: 'wf-probe-fail-child', + SpanKind: 'SPAN_KIND_CLIENT', + ServiceName: FAIL_SVC, + Duration: 100_000_000, + StatusCode: 'STATUS_CODE_OK', + }, + ]); + const original = ClickhouseClient.prototype.query; + const spy = jest + .spyOn(ClickhouseClient.prototype, 'query') + .mockImplementation(function (this: ClickhouseClient, args) { + if (args.query.includes('AS firstSeen')) { + return Promise.reject(new Error('Timeout exceeded')); + } + return original.call(this, args); + }); + try { + const result = await callTool(client, 'clickstack_trace_waterfall', { + sourceId: traceSource._id.toString(), + pickFilter: `ServiceName:${FAIL_SVC}`, + pickBy: 'slowest', + includeLogs: false, + }); + + expect(result.isError).toBeFalsy(); + const output = JSON.parse(getFirstText(result)); + expect(output.traceId).toBe(FAIL_TRACE_ID); + // Fallback window only covers the in-window child. + expect(output.spanCount).toBe(1); + expect(output.probeNote).toContain('probe failed'); + } finally { + spy.mockRestore(); + } + }); + it('probes past an explicit endTime so a late trace tail is not dropped', async () => { const PICK3_TRACE_ID = 'dddd4444eeee5555ffff6666aaaa7777'; const PICK3_SVC = 'wf-autopick-tail-svc'; diff --git a/packages/api/src/mcp/tools/trace/waterfall.ts b/packages/api/src/mcp/tools/trace/waterfall.ts index 40a27744c6..ad0e98c065 100644 --- a/packages/api/src/mcp/tools/trace/waterfall.ts +++ b/packages/api/src/mcp/tools/trace/waterfall.ts @@ -25,6 +25,7 @@ import { } from '@/mcp/tools/query/helpers'; import type { ToolRegistrar } from '@/mcp/tools/types'; import { mcpUserError } from '@/mcp/utils/errors'; +import logger from '@/utils/logger'; // ─── Schema ────────────────────────────────────────────────────────────────── @@ -196,10 +197,12 @@ function durationDivisor(precision: number): number { return Math.pow(10, Math.max(0, precision - 3)); } -// ClickHouse renders min()/max() over an empty set as the epoch / a zero date. -function parseProbeTimestamp(value: string | undefined): number | null { - if (!value) return null; - const ms = new Date(value).getTime(); +// min()/max() over an empty set yields the epoch, i.e. 0. +function parseProbeTimestamp( + value: string | number | undefined, +): number | null { + if (value == null) return null; + const ms = Number(value); return isNaN(ms) || ms <= 0 ? null : ms; } @@ -274,10 +277,12 @@ async function probeTraceWindow( params.startDate, params.endDate, ); + // Return epoch ms: a rendered DateTime string has no zone, so parsing it + // with new Date() would shift the window on a non-UTC API host. const probeQuery = ` SELECT - min(${params.ts.eventTs}) AS firstSeen, - max(${params.ts.eventTs}) AS lastSeen + toUnixTimestamp64Milli(toDateTime64(min(${params.ts.eventTs}), 3)) AS firstSeen, + toUnixTimestamp64Milli(toDateTime64(max(${params.ts.eventTs}), 3)) AS lastSeen FROM {db:Identifier}.{tbl:Identifier} WHERE ${params.traceIdExpr} = {tid:String} AND ${timeFilter.sql} @@ -297,7 +302,9 @@ async function probeTraceWindow( const rows = (await ( result as { - json: () => Promise<{ firstSeen: string; lastSeen: string }[]>; + json: () => Promise< + { firstSeen: string | number; lastSeen: string | number }[] + >; } ).json()) ?? []; const firstSeen = parseProbeTimestamp(rows[0]?.firstSeen); @@ -485,6 +492,7 @@ export function registerTraceWaterfall({ 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'); @@ -569,10 +577,19 @@ export function registerTraceWaterfall({ fetchEnd = probed.end; windowClamped = probed.clamped; } - } catch { + } 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; } } From bd063712bea06890e59d1f453d7a37bc9ecfec5c Mon Sep 17 00:00:00 2001 From: Brandon Pereira Date: Thu, 1 Oct 2026 13:18:11 -0600 Subject: [PATCH 6/6] refactor(mcp): dedupe trace waterfall client setup, drop step dividers --- packages/api/src/mcp/tools/trace/waterfall.ts | 49 +++++++++---------- 1 file changed, 22 insertions(+), 27 deletions(-) diff --git a/packages/api/src/mcp/tools/trace/waterfall.ts b/packages/api/src/mcp/tools/trace/waterfall.ts index ad0e98c065..cd79e7a498 100644 --- a/packages/api/src/mcp/tools/trace/waterfall.ts +++ b/packages/api/src/mcp/tools/trace/waterfall.ts @@ -206,6 +206,20 @@ function parseProbeTimestamp( return isNaN(ms) || ms <= 0 ? null : ms; } +async function getMcpClickhouseClient( + teamId: string, + connectionId: string, +): Promise { + const connection = await getConnectionById(teamId, connectionId, true); + if (!connection) return null; + return new ClickhouseClient({ + host: connection.host, + username: connection.username, + password: connection.password, + requestTimeout: MCP_REQUEST_TIMEOUT, + }); +} + // MCP settings win over source.querySettings so a source can't relax the // max_execution_time / readonly ceiling this tool depends on. function mcpQuerySettings(querySettings: QuerySettings | undefined) { @@ -392,23 +406,15 @@ export function registerTraceWaterfall({ ); } - const connection = await getConnectionById( + const clickhouseClient = await getMcpClickhouseClient( teamId.toString(), 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, - requestTimeout: MCP_REQUEST_TIMEOUT, - }); const metadata = getMetadata(clickhouseClient); const traceIdExpr = source.traceIdExpression; @@ -431,7 +437,6 @@ export function registerTraceWaterfall({ const attrsExpr = source.eventAttributesExpression ?? 'map()'; const divisor = durationDivisor(source.durationPrecision); - // ── Step 1: pick a TraceId (unless one was provided) ── let pickedTraceId = input.traceId; if (!pickedTraceId) { // Compose pickFilter with the pickBy-specific filter when needed. @@ -594,7 +599,6 @@ export function registerTraceWaterfall({ } } - // ── Step 2: fetch the full span tree ── let rows: SpanRow[]; try { const timeFilter = await ts.timeFilter(fetchStart, fetchEnd); @@ -702,7 +706,6 @@ export function registerTraceWaterfall({ "trace's start to see the full tree." : undefined; - // ── Step 3: fetch correlated logs (when logSourceId is configured) ── type LogRow = { timestamp: string; severityText: string; @@ -730,28 +733,20 @@ export function registerTraceWaterfall({ const logSvcExpr = logSource.serviceNameExpression ?? "''"; // Reuse the same connection only when the log source lives there. - let logClient = clickhouseClient; + let logClient: ClickhouseClient | null = clickhouseClient; if ( logSource.connection.toString() !== source.connection.toString() ) { - const logConn = await getConnectionById( + logClient = await getMcpClickhouseClient( teamId.toString(), logSource.connection.toString(), - true, ); - if (!logConn) { - logsNote = `connection for log source ${source.logSourceId} not found`; - } else { - logClient = new ClickhouseClient({ - host: logConn.host, - username: logConn.username, - password: logConn.password, - requestTimeout: MCP_REQUEST_TIMEOUT, - }); - } + } + if (!logClient) { + logsNote = `connection for log source ${source.logSourceId} not found`; } - if (!logsNote) { + if (logClient && !logsNote) { try { const logTs = await resolveTimestampExprs({ timestampValueExpression: logSource.timestampValueExpression,