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..4c6f0b5ba7 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)", @@ -309,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(); @@ -416,6 +433,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..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, @@ -432,6 +433,449 @@ 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 clamp dropped earlier spans, so the output must say so rather than + // silently promise the whole tree. + expect(output.windowNote).toContain('fetch-window cap'); + }); + }); + + // 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 = '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 + // 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); + // An 85-min spread is under the fetch-window cap, so no clamp warning. + expect(output.windowNote).toBeUndefined(); + }); + }); + + // 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'); + }); + + 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'); + }); + + 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'; + 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', () => { 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..3dcdcb0d52 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,16 @@ export function errorHint(msg: string, error?: unknown): string | null { 'The result row count is too large to serialize back to the agent.' ); } + // 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, ' + + '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..cd79e7a498 100644 --- a/packages/api/src/mcp/tools/trace/waterfall.ts +++ b/packages/api/src/mcp/tools/trace/waterfall.ts @@ -1,8 +1,15 @@ -import { getMetadata } from '@hyperdx/common-utils/dist/core/metadata'; +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'; @@ -12,10 +19,13 @@ 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'; import { mcpUserError } from '@/mcp/utils/errors'; +import logger from '@/utils/logger'; // ─── Schema ────────────────────────────────────────────────────────────────── @@ -31,8 +41,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 +77,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 +121,24 @@ const traceSchema = z.object({ type TraceInput = z.infer; +// 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 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 ───────────────────────────────────────────────────────────────── type SpanRow = { @@ -164,6 +197,148 @@ function durationDivisor(precision: number): number { return Math.pow(10, Math.max(0, precision - 3)); } +// 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; +} + +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) { + 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 }; + +/** + * 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]. + * + * `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, + params: { + databaseName: string; + tableName: string; + connectionId: string; + traceIdExpr: string; + ts: TimestampExprs; + traceId: string; + startDate: Date; + endDate: Date; + querySettings: QuerySettings | undefined; + }, +): Promise { + const timeFilter = await params.ts.timeFilter( + 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 + 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} + `; + const result = await clickhouseClient.query({ + query: probeQuery, + query_params: { + ...timeFilter.params, + db: params.databaseName, + tbl: params.tableName, + tid: params.traceId, + }, + format: 'JSONEachRow', + connectionId: params.connectionId, + clickhouse_settings: mcpQuerySettings(params.querySettings), + }); + const rows = + (await ( + result as { + json: () => Promise< + { firstSeen: string | number; lastSeen: string | number }[] + >; + } + ).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 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), + // The lead pad is slack, so only a floor above firstSeen drops real spans. + clamped: firstSeen < fetchStart, + }; +} + // ─── Tool definition ───────────────────────────────────────────────────────── export function registerTraceWaterfall({ @@ -231,22 +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, - }); const metadata = getMetadata(clickhouseClient); const traceIdExpr = source.traceIdExpression; @@ -256,13 +424,19 @@ export function registerTraceWaterfall({ 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); - // ── Step 1: pick a TraceId (unless one was provided) ── let pickedTraceId = input.traceId; if (!pickedTraceId) { // Compose pickFilter with the pickBy-specific filter when needed. @@ -291,8 +465,8 @@ export function registerTraceWaterfall({ input.pickBy === 'slowest' ? `max(${durationExpr}) DESC` : input.pickBy === 'first_error' - ? `min(${tsExpr}) ASC` - : `max(${tsExpr}) DESC`; + ? `min(${ts.eventTs}) ASC` + : `max(${ts.eventTs}) DESC`; const pickConfig: BuilderChartConfigWithDateRange = { displayType: DisplayType.Table, @@ -323,6 +497,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'); @@ -363,30 +538,92 @@ export function registerTraceWaterfall({ pickedTraceId = String(candidate[1]); } - // ── 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, - ${tsExpr} AS timestamp, - ${attrsExpr} AS spanAttributes - FROM {db:Identifier}.{tbl:Identifier} - WHERE ${traceIdExpr} = {tid:String} - ORDER BY ${tsExpr} ASC - LIMIT {n:UInt32} - `; + // 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; + } + } 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, @@ -395,15 +632,7 @@ export function registerTraceWaterfall({ }, format: 'JSONEachRow', connectionId: source.connection.toString(), - clickhouse_settings: { - readonly: '1', - // Per-query timeout matches the rest of the MCP for consistency. - ...(source.querySettings - ? Object.fromEntries( - source.querySettings.map(s => [s.setting, s.value]), - ) - : {}), - }, + clickhouse_settings: mcpQuerySettings(source.querySettings), }); rows = (await (result as { json: () => Promise }).json()) ?? []; @@ -418,6 +647,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 +674,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, @@ -440,7 +688,24 @@ export function registerTraceWaterfall({ const root = tree.find(s => s.depth === 0) ?? tree[0]; const totalDuration = Math.max(...spans.map(s => s.durationMs)); - // ── Step 3: fetch correlated logs (when logSourceId is configured) ── + // 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; @@ -463,49 +728,57 @@ export function registerTraceWaterfall({ } else { const logTraceIdExpr = logSource.traceIdExpression ?? 'TraceId'; const logSpanIdExpr = logSource.spanIdExpression ?? "''"; - const logTsExpr = logSource.timestampValueExpression; 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; + 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, - }); - } + } + if (!logClient) { + logsNote = `connection for log source ${source.logSourceId} not found`; } - 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} - ORDER BY ${logTsExpr} ASC - LIMIT {n:UInt32} - `; + 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 = ` + 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, @@ -513,14 +786,7 @@ export function registerTraceWaterfall({ }, format: 'JSONEachRow', connectionId: logSource.connection.toString(), - clickhouse_settings: { - readonly: '1', - ...(logSource.querySettings - ? Object.fromEntries( - logSource.querySettings.map(s => [s.setting, s.value]), - ) - : {}), - }, + clickhouse_settings: mcpQuerySettings(logSource.querySettings), }); const allLogs = (await ( @@ -561,6 +827,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.`,