diff --git a/.changeset/mcp-tool-quality-improvements.md b/.changeset/mcp-tool-quality-improvements.md new file mode 100644 index 0000000000..13fa6e2a1e --- /dev/null +++ b/.changeset/mcp-tool-quality-improvements.md @@ -0,0 +1,12 @@ +--- +'@hyperdx/api': patch +--- + +feat(mcp): improve MCP tool quality — error hints, shared helpers, better messages + +Extract duplicated ClickHouse error handling into a shared helper with +pattern-matched error hints (DateTime64 casting, AS alias quoting, response +size limits) so agents get actionable guidance on common failures. Add +reusable mergeWhereIntoSelectItems() helper for consistent top-level where +injection. Improve source/connection-not-found messages to suggest calling +hyperdx_list_sources. diff --git a/packages/api/src/mcp/__tests__/query.test.ts b/packages/api/src/mcp/__tests__/query.test.ts index 6f6f31b572..214adc058b 100644 --- a/packages/api/src/mcp/__tests__/query.test.ts +++ b/packages/api/src/mcp/__tests__/query.test.ts @@ -1,4 +1,17 @@ -import { parseTimeRange } from '../tools/query/helpers'; +// Mock heavy dependencies that break in unit-test context (no ClickHouse/Mongo) +jest.mock('@/models/source', () => ({})); +jest.mock('@/controllers/sources', () => ({})); +jest.mock('@/controllers/connection', () => ({})); +jest.mock('@/utils/trimToolResponse', () => ({ + trimToolResponse: (data: unknown) => ({ data, isTrimmed: false }), +})); + +import { + errorHint, + mergeWhereIntoSelectItems, + parseTimeRange, +} from '../tools/query/helpers'; +import { resolveOrderBy } from '../tools/query/table'; describe('parseTimeRange', () => { it('should return default range (last 15 minutes) when no arguments provided', () => { @@ -72,3 +85,249 @@ describe('parseTimeRange', () => { expect(result).toHaveProperty('error'); }); }); + +// ─── mergeWhereIntoSelectItems ─────────────────────────────────────────────── + +describe('mergeWhereIntoSelectItems', () => { + it('should return items unchanged when topWhere is empty', () => { + const items = [{ where: 'level:error', whereLanguage: 'lucene' as const }]; + const { items: result, warnings } = mergeWhereIntoSelectItems( + items, + '', + 'lucene', + ); + expect(result).toBe(items); // same reference, not a copy + expect(warnings).toHaveLength(0); + }); + + it('should inject top-level where when item has no where', () => { + const items = [{ where: '', whereLanguage: 'lucene' as const }]; + const { items: result, warnings } = mergeWhereIntoSelectItems( + items, + 'service:api', + 'lucene', + ); + expect(result[0].where).toBe('service:api'); + expect(result[0].whereLanguage).toBe('lucene'); + expect(warnings).toHaveLength(0); + }); + + it('should AND-combine when languages match', () => { + const items = [{ where: 'level:error', whereLanguage: 'lucene' as const }]; + const { items: result, warnings } = mergeWhereIntoSelectItems( + items, + 'service:api', + 'lucene', + ); + expect(result[0].where).toBe('(service:api) AND (level:error)'); + expect(result[0].whereLanguage).toBe('lucene'); + expect(warnings).toHaveLength(0); + }); + + it('should AND-combine SQL filters when both use sql', () => { + const items = [ + { where: 'StatusCode >= 500', whereLanguage: 'sql' as const }, + ]; + const { items: result, warnings } = mergeWhereIntoSelectItems( + items, + "ServiceName = 'api'", + 'sql', + ); + expect(result[0].where).toBe( + "(ServiceName = 'api') AND (StatusCode >= 500)", + ); + expect(result[0].whereLanguage).toBe('sql'); + expect(warnings).toHaveLength(0); + }); + + it('should skip item and emit warning when languages differ', () => { + const items = [ + { where: 'StatusCode >= 500', whereLanguage: 'sql' as const }, + ]; + const { items: result, warnings } = mergeWhereIntoSelectItems( + items, + 'service:api', + 'lucene', + ); + // Item's own where is kept unchanged + expect(result[0].where).toBe('StatusCode >= 500'); + expect(result[0].whereLanguage).toBe('sql'); + // A warning is emitted + expect(warnings).toHaveLength(1); + expect(warnings[0]).toContain('select[0]'); + expect(warnings[0]).toContain('was NOT applied'); + }); + + it('should handle mixed items — some merged, some skipped', () => { + const items = [ + { where: 'level:error', whereLanguage: 'lucene' as const }, + { where: 'StatusCode >= 500', whereLanguage: 'sql' as const }, + { where: '', whereLanguage: 'lucene' as const }, + ]; + const { items: result, warnings } = mergeWhereIntoSelectItems( + items, + 'service:api', + 'lucene', + ); + // Item 0: merged (both lucene) + expect(result[0].where).toBe('(service:api) AND (level:error)'); + // Item 1: skipped (sql vs lucene) + expect(result[1].where).toBe('StatusCode >= 500'); + // Item 2: injected (empty where) + expect(result[2].where).toBe('service:api'); + expect(warnings).toHaveLength(1); + expect(warnings[0]).toContain('select[1]'); + }); +}); + +// ─── errorHint ─────────────────────────────────────────────────────────────── + +describe('errorHint', () => { + it('should match DateTime64 conversion errors', () => { + const hint = errorHint( + "Cannot convert string '2025-01-01T00:00:00Z' to type DateTime64(9)", + ); + expect(hint).not.toBeNull(); + expect(hint).toContain('toDateTime64'); + }); + + it('should match DateTime64 parse errors', () => { + const hint = errorHint( + "Cannot parse string '2025-01-01' as type DateTime64", + ); + expect(hint).not.toBeNull(); + expect(hint).toContain('toDateTime64'); + }); + + it('should match AS alias syntax errors with word boundary', () => { + const hint = errorHint('Syntax error: unexpected token AS'); + expect(hint).not.toBeNull(); + expect(hint).toContain('alias'); + }); + + it('should NOT match AS embedded in other words like DATABASE', () => { + const hint = errorHint("Syntax error: Unknown database 'FOOBAR'"); + expect(hint).toBeNull(); + }); + + it('should NOT match lowercase "as" in cast expressions', () => { + const hint = errorHint( + 'Syntax error: cannot cast as Float64 near token...', + ); + expect(hint).toBeNull(); + }); + + it('should match V8 string length overflow', () => { + const hint = errorHint( + 'response length exceeds the maximum allowed size of V8 String', + ); + expect(hint).not.toBeNull(); + expect(hint).toContain('LIMIT'); + }); + + it('should return null for unrecognized errors', () => { + const hint = errorHint('Connection refused'); + expect(hint).toBeNull(); + }); +}); + +// ─── resolveOrderBy ────────────────────────────────────────────────────────── + +describe('resolveOrderBy', () => { + it('should return undefined when orderBy is undefined', () => { + expect(resolveOrderBy(undefined, [{ aggFn: 'count' }])).toBeUndefined(); + }); + + it('should pass through non-aggFn values unchanged', () => { + expect(resolveOrderBy('SpanName', [{ aggFn: 'count' }])).toBe('SpanName'); + }); + + it('should keep orderBy when it matches an alias', () => { + expect(resolveOrderBy('Total', [{ aggFn: 'count', alias: 'Total' }])).toBe( + 'Total', + ); + }); + + it('should return canonical alias case for case-insensitive match', () => { + expect(resolveOrderBy('total', [{ aggFn: 'count', alias: 'Total' }])).toBe( + 'Total', + ); + expect( + resolveOrderBy('TOTAL DESC', [{ aggFn: 'count', alias: 'Total' }]), + ).toBe('Total DESC'); + }); + + it('should synthesize count()', () => { + expect(resolveOrderBy('count', [{ aggFn: 'count' }])).toBe('count()'); + }); + + it('should synthesize avg(Duration)', () => { + expect( + resolveOrderBy('avg', [{ aggFn: 'avg', valueExpression: 'Duration' }]), + ).toBe('avg(Duration)'); + }); + + it('should synthesize quantile with level', () => { + expect( + resolveOrderBy('quantile', [ + { aggFn: 'quantile', valueExpression: 'Duration', level: 0.99 }, + ]), + ).toBe('quantile(0.99)(Duration)'); + }); + + it('should prefer alias over synthesized expression', () => { + expect( + resolveOrderBy('count', [{ aggFn: 'count', alias: 'Total Rows' }]), + ).toBe('Total Rows'); + }); + + it('should be case-insensitive for aggFn matching', () => { + expect(resolveOrderBy('Count', [{ aggFn: 'count' }])).toBe('count()'); + expect( + resolveOrderBy('AVG', [{ aggFn: 'avg', valueExpression: 'Duration' }]), + ).toBe('avg(Duration)'); + }); + + it('should synthesize count(DISTINCT ...) for count_distinct', () => { + expect( + resolveOrderBy('count_distinct', [ + { aggFn: 'count_distinct', valueExpression: 'UserId' }, + ]), + ).toBe('count(DISTINCT UserId)'); + }); + + it('should synthesize count(DISTINCT ...) with direction', () => { + expect( + resolveOrderBy('count_distinct DESC', [ + { aggFn: 'count_distinct', valueExpression: 'UserId' }, + ]), + ).toBe('count(DISTINCT UserId) DESC'); + }); + + it('should NOT synthesize for "none" aggFn', () => { + expect( + resolveOrderBy('none', [ + { aggFn: 'none', valueExpression: 'Duration / 1e6' }, + ]), + ).toBe('none'); + }); + + it('should handle trailing ASC/DESC', () => { + expect(resolveOrderBy('count DESC', [{ aggFn: 'count' }])).toBe( + 'count() DESC', + ); + expect( + resolveOrderBy('avg ASC', [ + { aggFn: 'avg', valueExpression: 'Duration' }, + ]), + ).toBe('avg(Duration) ASC'); + }); + + it('should pass through quantile without level unchanged', () => { + expect( + resolveOrderBy('quantile', [ + { aggFn: 'quantile', valueExpression: 'Duration' }, + ]), + ).toBe('quantile'); + }); +}); diff --git a/packages/api/src/mcp/__tests__/queryTool.test.ts b/packages/api/src/mcp/__tests__/queryTool.test.ts index 0ecabba06d..d51a6f92a0 100644 --- a/packages/api/src/mcp/__tests__/queryTool.test.ts +++ b/packages/api/src/mcp/__tests__/queryTool.test.ts @@ -138,6 +138,7 @@ describe('MCP Query Tools', () => { expect(props).toContain('sourceId'); expect(props).toContain('where'); expect(props).toContain('sampleSize'); + expect(props).toContain('topN'); expect(props).toContain('bodyExpression'); expect(tool!.inputSchema.required).toContain('sourceId'); }); @@ -370,14 +371,14 @@ describe('MCP Query Tools', () => { expect(result.isError).toBeFalsy(); expect(result.content).toHaveLength(1); const output = JSON.parse(getFirstText(result)); - expect(output).toHaveProperty('result'); - expect(output.result).toHaveProperty('patterns'); - expect(output.result).toHaveProperty('totalCount'); - expect(output.result).toHaveProperty('sampledRows'); - expect(output.result).toHaveProperty('sampleMultiplier'); - expect(output.result).toHaveProperty('bodyColumn', 'Body'); - expect(output.result).toHaveProperty('timeRange'); - expect(Array.isArray(output.result.patterns)).toBe(true); + expect(output).toHaveProperty('summary'); + expect(output).toHaveProperty('patterns'); + expect(output.summary).toHaveProperty('totalCount'); + expect(output.summary).toHaveProperty('sampledCount'); + expect(output.summary).toHaveProperty('sampleMultiplier'); + expect(output.summary).toHaveProperty('bodyColumn', 'Body'); + expect(output.summary).toHaveProperty('timeRange'); + expect(Array.isArray(output.patterns)).toBe(true); }); it('should execute with explicit bodyExpression on trace source', async () => { @@ -391,9 +392,9 @@ describe('MCP Query Tools', () => { expect(result.isError).toBeFalsy(); expect(result.content).toHaveLength(1); const output = JSON.parse(getFirstText(result)); - expect(output).toHaveProperty('result'); - expect(Array.isArray(output.result.patterns)).toBe(true); - expect(output.result).toHaveProperty('bodyColumn', 'SpanName'); + expect(output).toHaveProperty('summary'); + expect(Array.isArray(output.patterns)).toBe(true); + expect(output.summary).toHaveProperty('bodyColumn', 'SpanName'); }); it('should accept custom bodyExpression on log source', async () => { @@ -407,7 +408,7 @@ describe('MCP Query Tools', () => { expect(result.isError).toBeFalsy(); expect(result.content).toHaveLength(1); const output = JSON.parse(getFirstText(result)); - expect(output.result).toHaveProperty('bodyColumn', 'SeverityText'); + expect(output.summary).toHaveProperty('bodyColumn', 'SeverityText'); }); it('should respect sampleSize parameter', async () => { @@ -503,10 +504,11 @@ describe('MCP Query Tools', () => { expect(result.isError).toBeFalsy(); const output = JSON.parse(getFirstText(result)); - const { patterns, totalCount, sampledRows } = output.result; + const { patterns } = output; + const { totalCount, sampledCount } = output.summary; expect(totalCount).toBeGreaterThanOrEqual(31); - expect(sampledRows).toBeGreaterThanOrEqual(31); + expect(sampledCount).toBeGreaterThanOrEqual(31); expect(patterns.length).toBeGreaterThanOrEqual(1); // The most common pattern should contain <*> placeholders diff --git a/packages/api/src/mcp/tools/query/eventDeltas.ts b/packages/api/src/mcp/tools/query/eventDeltas.ts index 1d10a5e812..f5ded14674 100644 --- a/packages/api/src/mcp/tools/query/eventDeltas.ts +++ b/packages/api/src/mcp/tools/query/eventDeltas.ts @@ -526,13 +526,20 @@ export function registerEventDeltas(server: McpServer, context: McpContext) { 'are dropped by default — set includeHidden:true to inspect them.', }; - const trimmedOutput = trimToolResponse(output); + const { data: trimmedOutput, isTrimmed } = trimToolResponse(output); + + const finalOutput = isTrimmed + ? { + ...trimmedOutput, + note: 'Result was trimmed for context size. Narrow the time range, add filters, or reduce topN/topValuesPerProperty to reduce data.', + } + : trimmedOutput; return { content: [ { type: 'text' as const, - text: JSON.stringify(trimmedOutput, null, 2), + text: JSON.stringify(finalOutput, null, 2), }, ], }; diff --git a/packages/api/src/mcp/tools/query/eventPatterns.ts b/packages/api/src/mcp/tools/query/eventPatterns.ts index 3c65102994..c3da47cca7 100644 --- a/packages/api/src/mcp/tools/query/eventPatterns.ts +++ b/packages/api/src/mcp/tools/query/eventPatterns.ts @@ -38,6 +38,24 @@ const eventPatternsSchema = z.object({ 'Auto-detected from the source if omitted (Body for logs, SpanName for traces). ' + 'Example: "Body", "SpanName", "SpanAttributes[\'http.url\']"', ), + topN: z + .number() + .min(1) + .max(100) + .optional() + .default(20) + .describe('Number of top patterns to return (1-100). Default: 20.'), + trendBuckets: z + .number() + .min(0) + .max(60) + .optional() + .default(24) + .describe( + 'Number of time buckets for per-pattern trend sparklines (0-60). Default: 24. ' + + 'Compare bucket counts at start vs end to spot patterns that started or stopped mid-window. ' + + 'Set to 0 to disable trend computation (smaller response).', + ), startTime: startTimeSchema, endTime: endTimeSchema, }); @@ -86,6 +104,8 @@ export function registerEventPatterns(server: McpServer, context: McpContext) { whereLanguage: input.whereLanguage, bodyExpression: input.bodyExpression, sampleSize: input.sampleSize, + topN: input.topN, + trendBuckets: input.trendBuckets, }, ); }), diff --git a/packages/api/src/mcp/tools/query/helpers.ts b/packages/api/src/mcp/tools/query/helpers.ts index e065c4c642..51f83f2d80 100644 --- a/packages/api/src/mcp/tools/query/helpers.ts +++ b/packages/api/src/mcp/tools/query/helpers.ts @@ -27,6 +27,64 @@ import { trimToolResponse } from '@/utils/trimToolResponse'; import type { ExternalDashboardTileWithId } from '@/utils/zod'; import { externalDashboardTileSchemaWithId } from '@/utils/zod'; +// ─── Where merging ─────────────────────────────────────────────────────────── + +export interface MergeWhereResult { + items: T[]; + /** Non-empty when items were skipped due to language mismatch. */ + warnings: string[]; +} + +/** + * Merge a top-level `where` filter into each select item so it becomes part + * of the per-item aggCondition. Table/line/number/pie display types don't have + * a chart-level where — filtering is per-select-item. + * + * When the top-level and item-level languages differ, the item's own filter + * takes precedence (we can't easily merge Lucene + SQL). A warning is returned + * so callers can surface it in the response. + */ +export function mergeWhereIntoSelectItems< + T extends { + where?: string; + whereLanguage?: 'lucene' | 'sql'; + }, +>( + items: T[], + topWhere: string, + topLang: 'lucene' | 'sql', +): MergeWhereResult { + if (!topWhere) return { items, warnings: [] }; + + const warnings: string[] = []; + + const merged = items.map((item, idx) => { + const itemWhere = item.where || ''; + const itemLang = item.whereLanguage || 'lucene'; + + // If both languages match, combine with AND + if (itemWhere && itemLang === topLang) { + const combined = `(${topWhere}) AND (${itemWhere})`; + return { ...item, where: combined, whereLanguage: topLang }; + } + + // If the item has no where, just use the top-level + if (!itemWhere) { + return { ...item, where: topWhere, whereLanguage: topLang }; + } + + // Languages differ — keep item's where unchanged (can't easily merge + // Lucene + SQL). The item's own filter takes precedence. + warnings.push( + `select[${idx}]: top-level where (${topLang}) was NOT applied because this item uses whereLanguage:"${itemLang}". ` + + `Set the item's whereLanguage to "${topLang}" or rewrite the top-level where in ${itemLang} to apply both filters.`, + ); + return item; + }); + + return { items: merged, warnings }; +} + // ─── Tile construction ─────────────────────────────────────────────────────── /** @@ -85,9 +143,7 @@ function isEmptyResult(result: unknown): boolean { } function formatQueryResult(result: unknown) { - const trimmedResult = trimToolResponse(result); - const isTrimmed = - JSON.stringify(trimmedResult).length < JSON.stringify(result).length; + const { data: trimmedResult, isTrimmed } = trimToolResponse(result); const empty = isEmptyResult(result); return { content: [ @@ -160,7 +216,7 @@ export async function runConfigTile( content: [ { type: 'text' as const, - text: `Source not found: ${builderConfig.source}`, + text: `Source not found: ${builderConfig.source}. Call hyperdx_list_sources to discover available source IDs.`, }, ], }; @@ -234,27 +290,16 @@ export async function runConfigTile( } satisfies ChartConfigWithDateRange; const metadata = getMetadata(clickhouseClient); - let result; try { - result = await clickhouseClient.queryChartConfig({ + const result = await clickhouseClient.queryChartConfig({ config: chartConfig, metadata, querySettings: source.querySettings, }); - } catch (err) { - const message = err instanceof Error ? err.message : String(err); - return { - isError: true as const, - content: [ - { - type: 'text' as const, - text: `ClickHouse query failed: ${message}`, - }, - ], - }; + return formatQueryResult(result); + } catch (e) { + return clickHouseErrorResult(e); } - - return formatQueryResult(result); } // Raw SQL tile — hydrate source fields for macro support ($__sourceTable, $__filters) @@ -294,7 +339,7 @@ export async function runConfigTile( content: [ { type: 'text' as const, - text: `Connection not found: ${savedConfig.connection}`, + text: `Connection not found: ${savedConfig.connection}. Call hyperdx_list_sources to discover available connection IDs.`, }, ], }; @@ -313,25 +358,66 @@ export async function runConfigTile( } satisfies ChartConfigWithDateRange; const metadata = getMetadata(clickhouseClient); - let result; try { - result = await clickhouseClient.queryChartConfig({ + const result = await clickhouseClient.queryChartConfig({ config: chartConfig, metadata, querySettings: undefined, }); - } catch (err) { - const message = err instanceof Error ? err.message : String(err); - return { - isError: true as const, - content: [ - { - type: 'text' as const, - text: `ClickHouse query failed: ${message}`, - }, - ], - }; + return formatQueryResult(result); + } catch (e) { + return clickHouseErrorResult(e); } +} + +// ─── Error hints ───────────────────────────────────────────────────────────── - return formatQueryResult(result); +/** + * Decorate raw ClickHouse error messages with actionable guidance before + * they reach the agent. Some ClickHouse errors are unhelpful in isolation — + * e.g. "Cannot convert string '...Z' to type DateTime64(9)" leaves the agent + * guessing about the right format. Catch common patterns and append a hint. + */ +export function clickHouseErrorResult(e: unknown): { + isError: true; + content: [{ type: 'text'; text: string }]; +} { + const raw = e instanceof Error ? e.message : String(e); + const hint = errorHint(raw); + const text = hint ? `${raw}\n\nHINT: ${hint}` : raw; + return { + isError: true as const, + content: [{ type: 'text' as const, text }], + }; +} + +/** @internal Exported for testing only. */ +export function errorHint(msg: string): string | null { + if ( + /Cannot (convert|parse) string .* (to|as) (type )?DateTime64/i.test(msg) + ) { + return ( + "Wrap ISO timestamps with `toDateTime64('YYYY-MM-DD HH:MM:SS', 9)` " + + 'or use the supplied macros: `$__timeFilter(Timestamp)` (the easiest), ' + + '`{startDateMilliseconds:Int64}`, `{endDateMilliseconds:Int64}`. ' + + 'Bare ISO 8601 strings will NOT auto-cast to DateTime64.' + ); + } + if (/Syntax error.*\bAS\b/.test(msg)) { + return ( + 'Quote the alias if it contains reserved words or special chars: ' + + '`expr AS "alias"`. The MCP builder tools accept `alias` as a ' + + 'separate field on each select entry — use that to avoid SQL-quoting ' + + 'headaches.' + ); + } + if ( + /response length exceeds the maximum allowed size of V8 String/i.test(msg) + ) { + return ( + 'Add a LIMIT, narrow the time range, or use a smaller granularity. ' + + 'The result row count is too large to serialize back to the agent.' + ); + } + return null; } diff --git a/packages/api/src/mcp/tools/query/runEventPatterns.ts b/packages/api/src/mcp/tools/query/runEventPatterns.ts index 6a460f51f2..1e11df8a37 100644 --- a/packages/api/src/mcp/tools/query/runEventPatterns.ts +++ b/packages/api/src/mcp/tools/query/runEventPatterns.ts @@ -12,6 +12,8 @@ import { getConnectionById } from '@/controllers/connection'; import { getSource } from '@/controllers/sources'; import { trimToolResponse } from '@/utils/trimToolResponse'; +import { clickHouseErrorResult } from './helpers'; + // ─── Source helpers ────────────────────────────────────────────────────────── interface SourceBodyFields { @@ -53,9 +55,13 @@ export async function runEventPatterns( whereLanguage?: 'lucene' | 'sql'; bodyExpression?: string; sampleSize?: number; + topN?: number; + trendBuckets?: number; }, ) { const sampleSize = options?.sampleSize ?? 10_000; + const topN = options?.topN ?? 20; + const trendBuckets = options?.trendBuckets ?? 24; // ── Resolve source & connection ── const source = await getSource(teamId, sourceId); @@ -63,7 +69,10 @@ export async function runEventPatterns( return { isError: true as const, content: [ - { type: 'text' as const, text: `Source not found: ${sourceId}` }, + { + type: 'text' as const, + text: `Source not found: ${sourceId}. Call hyperdx_list_sources to discover available source IDs.`, + }, ], }; } @@ -79,7 +88,7 @@ export async function runEventPatterns( content: [ { type: 'text' as const, - text: `Connection not found for source: ${sourceId}`, + text: `Connection not found for source: ${sourceId}. Call hyperdx_list_sources to discover available source IDs.`, }, ], }; @@ -206,16 +215,7 @@ export async function runEventPatterns( }), ]); } catch (err) { - const message = err instanceof Error ? err.message : String(err); - return { - isError: true as const, - content: [ - { - type: 'text' as const, - text: `ClickHouse query failed: ${message}`, - }, - ], - }; + return clickHouseErrorResult(err); } const sampleRows = sampleResult.data; @@ -228,17 +228,19 @@ export async function runEventPatterns( type: 'text' as const, text: JSON.stringify( { - result: { - patterns: [], + summary: { totalCount, - sampledRows: 0, + sampledCount: 0, sampleMultiplier: 1, + clusterCount: 0, + patternsReturned: 0, bodyColumn, timeRange: { start: startDate.toISOString(), end: endDate.toISOString(), }, }, + patterns: [], hint: 'No data found in the queried time range. Try setting startTime to a wider window (e.g. 24 hours ago) or check that filters match existing data.', }, null, @@ -257,6 +259,7 @@ export async function runEventPatterns( totalCount, startDate, endDate, + trendBuckets, maxSamples: 5, getBody: row => { const raw = row.__hdx_pattern_body; @@ -284,7 +287,10 @@ export async function runEventPatterns( // Convert trend timestamps to ISO strings, extract sample body texts, // and build a whereSnippet per pattern so the agent can drill into // matching events via a follow-up hyperdx_search query. - const patterns = rawPatterns.map(p => { + const sampledCount = sampleRows.length; + const slicedPatterns = rawPatterns.slice(0, topN); + + const patterns = slicedPatterns.map((p, i) => { // Build a Lucene-compatible where clause from the pattern's literal // (non-<*>) tokens. This lets agents chain: pattern → search. // Escape \ and " (the phrase-query metachars) in each token. @@ -297,16 +303,25 @@ export async function runEventPatterns( ? `${bodyColumn}:"${literalTokens.join(' ')}"` : ''; + const shareOfTotal = sampledCount > 0 ? p.sampleCount / sampledCount : 0; + + const formattedTrend = + trendBuckets > 0 + ? p.trend.map(t => ({ + ts: new Date(t.ts).toISOString(), + count: t.count, + })) + : undefined; + return { + rank: i + 1, id: p.id, pattern: p.pattern, estimatedCount: p.estimatedCount, sampleCount: p.sampleCount, + shareOfTotal: Math.round(shareOfTotal * 10000) / 10000, whereSnippet, - trend: p.trend.map(t => ({ - ts: new Date(t.ts).toISOString(), - count: t.count, - })), + ...(formattedTrend ? { trend: formattedTrend } : {}), samples: p.samples.map(row => { const raw = row.__hdx_pattern_body; return raw != null ? String(raw) : ''; @@ -315,37 +330,42 @@ export async function runEventPatterns( }); const output = { - patterns, - totalCount, - sampledRows: sampleRows.length, - sampleMultiplier: Math.round(sampleMultiplier * 100) / 100, - bodyColumn, - timeRange: { - start: startDate.toISOString(), - end: endDate.toISOString(), + summary: { + totalCount, + sampledCount, + sampleMultiplier: Math.round(sampleMultiplier * 100) / 100, + clusterCount: rawPatterns.length, + patternsReturned: patterns.length, + bodyColumn, + timeRange: { + start: startDate.toISOString(), + end: endDate.toISOString(), + }, }, + patterns, + usage: + 'shareOfTotal is the fraction of sampled rows matching this pattern. ' + + 'estimatedCount = sampleCount * sampleMultiplier. ' + + (trendBuckets > 0 + ? 'trend.count is similarly extrapolated from sample bucket counts. ' + : '') + + 'Use whereSnippet as the "where" parameter in a hyperdx_search call to browse matching raw events.', }; - const trimmedOutput = trimToolResponse(output); - const isTrimmed = - JSON.stringify(trimmedOutput).length < JSON.stringify(output).length; + const { data: trimmedOutput, isTrimmed } = trimToolResponse(output); + + const finalOutput = isTrimmed + ? { + ...trimmedOutput, + note: 'Result was trimmed for context size. Narrow the time range, add filters, or reduce topN to reduce data.', + } + : trimmedOutput; return { content: [ { type: 'text' as const, - text: JSON.stringify( - { - result: trimmedOutput, - ...(isTrimmed - ? { - note: 'Result was trimmed for context size. Narrow the time range or add filters to reduce data.', - } - : {}), - }, - null, - 2, - ), + text: JSON.stringify(finalOutput, null, 2), }, ], }; diff --git a/packages/api/src/mcp/tools/query/schemas.ts b/packages/api/src/mcp/tools/query/schemas.ts index 9da11a64b6..217a84c30d 100644 --- a/packages/api/src/mcp/tools/query/schemas.ts +++ b/packages/api/src/mcp/tools/query/schemas.ts @@ -17,24 +17,36 @@ const WHERE_DESCRIPTION = const WHERE_LANGUAGE_DESCRIPTION = 'Query language for the "where" filter. Default: lucene.\n' + + 'IMPORTANT: the syntax in "where" MUST match "whereLanguage".\n' + + ' Lucene (default): Column:value, Column.mapKey:value, Column:>100\n' + + " SQL: Column = 'value', SpanAttributes['key'] = 'value'\n\n" + 'Lucene supports comparisons (>= > < <=), wildcards (field:val*), ranges ([1 TO 5]), ' + 'and map attributes via dot notation. ' + - 'Use "sql" for IN(...) lists, complex expressions, or function calls.'; + 'Use "sql" for IN(...) lists, complex expressions, or function calls.\n\n' + + 'Common mistake: writing Column:value (Lucene) but setting whereLanguage to "sql". ' + + 'If your filter uses colon syntax, leave whereLanguage as "lucene" (the default).'; // ─── Shared Zod schemas ────────────────────────────────────────────────────── +/** + * Aggregation function names exposed to MCP tool callers. + * This is the single source of truth — used by both the Zod input schema + * (mcpAggFnSchema) and the orderBy resolver in table.ts (AGG_FN_NAMES Set). + */ +export const MCP_AGG_FN_OPTIONS = [ + 'avg', + 'count', + 'count_distinct', + 'last_value', + 'max', + 'min', + 'quantile', + 'sum', + 'none', +] as const; + const mcpAggFnSchema = z - .enum([ - 'avg', - 'count', - 'count_distinct', - 'last_value', - 'max', - 'min', - 'quantile', - 'sum', - 'none', - ]) + .enum(MCP_AGG_FN_OPTIONS) .describe( 'Aggregation function:\n' + ' count – count matching rows (no valueExpression needed)\n' + @@ -51,26 +63,38 @@ export const mcpSelectItemSchema = z.object({ .string() .optional() .describe( - 'Column or expression to aggregate. Required for every aggFn except "count". ' + - 'Use PascalCase for top-level columns (e.g. "Duration", "StatusCode"). ' + - "For span attributes use: SpanAttributes['key'] (e.g. SpanAttributes['http.method']). " + - "For resource attributes use: ResourceAttributes['key'] (e.g. ResourceAttributes['service.name']).", + 'ClickHouse SQL expression to aggregate. Required for every aggFn except "count". ' + + 'Top-level columns are PascalCase (Duration, StatusCode); ' + + "map attributes use bracket syntax: SpanAttributes['key'], ResourceAttributes['key']. " + + 'Any ClickHouse expression is allowed — common useful forms: ' + + '"Duration / 1e6" (ns→ms), ' + + '"toFloat64OrZero(SpanAttributes[\'response.size_bytes\'])" (cast attribute), ' + + '"if(StatusCode = \'STATUS_CODE_ERROR\', 1, 0)" (boolean→numeric for ratios).', ), where: z .string() .optional() .default('') .describe( - 'Per-metric row filter. Compiles to If(...) for multi-cohort ' + - 'aggregates in one call. ' + - 'Example: "level:error" (Lucene) or "StatusCode >= 500" (SQL)', + 'Conditional aggregation filter — restricts which rows are included in THIS metric ' + + '(combined with the top-level time/where filter via AND). ' + + 'Compiles to If(...): e.g. quantile + where=Timestamp= \'2026-05-09T23:40:00Z\'" on another to get baseline-vs-anomaly p99 ' + + 'in one round trip — much faster than re-running the same query with a different time range. ' + + 'Examples (lucene): "level:error", "service.name:api AND http.status_code:>=500". ' + + 'Set whereLanguage:"sql" for raw SQL conditions like ' + + "\"SpanAttributes['http.method'] = 'POST'\" or \"Timestamp < '2026-05-09 23:40:00'\".", ), whereLanguage: z .enum(['lucene', 'sql']) .optional() .default('lucene') .describe( - "Query language for this select item's where filter. Default: lucene", + 'Query language for the per-item conditional filter. ' + + 'Use "sql" when comparing to literal timestamps or arbitrary attribute expressions. ' + + 'Default: lucene', ), alias: z .string() @@ -120,13 +144,27 @@ export const groupBySchema = z .string() .optional() .describe( - 'Column to group/split by. ' + - 'Top-level columns use PascalCase (e.g. "SpanName", "StatusCode"). ' + - "Span attributes: SpanAttributes['key'] (e.g. SpanAttributes['http.method']). " + - "Resource attributes: ResourceAttributes['key'] (e.g. ResourceAttributes['service.name']).", + 'Column(s) or ClickHouse expression(s) to group/split by. ' + + 'Accepts a SINGLE entry or MULTIPLE entries as a comma-delimited list — ' + + 'multi-column groupBy expresses multi-dimensional breakdowns in one ' + + 'query (e.g. "ServiceName, SpanName, StatusMessage") instead of running ' + + 'one query per dimension. For "table" displayType, the result has one row ' + + 'per distinct combination of group values. ' + + 'Top-level columns use PascalCase ("SpanName", "StatusCode"). ' + + "Map attributes: SpanAttributes['key'], ResourceAttributes['key'].\n\n" + + 'Arbitrary ClickHouse expressions are also allowed in groupBy — useful when ' + + 'you need to group by a derived column without falling back to raw SQL:\n' + + ' - "substring(Body, 1, 80)" — group by body prefix (log pattern bucketing)\n' + + ' - "toStartOfInterval(Timestamp, INTERVAL 5 MINUTE)" — explicit time bucketing ' + + 'in a table view, alongside another dimension (granularity only works for line/stacked_bar)\n' + + ' - "JSONExtractString(Body, \'event\')" — parse a JSON field from the body\n' + + " - \"if(Duration > 1e9, 'slow', 'fast')\" — coarse boolean buckets\n" + + 'Comma splitting is bracket-aware, so multi-arg function calls work as single entries.', ); export const orderBySchema = z .string() .optional() - .describe('Column to sort results by (mainly useful for table shape).'); + .describe( + 'Column to sort results by (builder display types only, mainly "table").', + ); diff --git a/packages/api/src/mcp/tools/query/table.ts b/packages/api/src/mcp/tools/query/table.ts index 0d0bce3292..cf1e15cb03 100644 --- a/packages/api/src/mcp/tools/query/table.ts +++ b/packages/api/src/mcp/tools/query/table.ts @@ -3,14 +3,22 @@ import { z } from 'zod'; import { withToolTracing } from '../../utils/tracing'; import type { McpContext } from '../types'; -import { buildTile, parseTimeRange, runConfigTile } from './helpers'; +import { + buildTile, + mergeWhereIntoSelectItems, + parseTimeRange, + runConfigTile, +} from './helpers'; import { endTimeSchema, groupBySchema, + MCP_AGG_FN_OPTIONS, mcpSelectItemSchema, orderBySchema, sourceIdSchema, startTimeSchema, + whereLanguageSchema, + whereSchema, } from './schemas'; // ─── Schema ────────────────────────────────────────────────────────────────── @@ -35,12 +43,111 @@ const tableSchema = z.object({ 'If "number" or "pie" is set with select.length > 1, it is auto-upgraded to "table". ' + 'groupBy is ignored when shape is "number".', ), + where: whereSchema.describe( + 'Row filter applied to ALL select items. ' + + 'This scopes the entire query — use it to restrict by service, severity, etc. ' + + 'Each select item can also have its own per-metric "where" for cohort comparisons. ' + + 'When both are set, they are ANDed together.\n\n' + + 'Lucene example: "ServiceName:frontend"\n' + + 'SQL example: "ServiceName = \'frontend\'"', + ), + whereLanguage: whereLanguageSchema, groupBy: groupBySchema, orderBy: orderBySchema, startTime: startTimeSchema, endTime: endTimeSchema, }); +// ─── orderBy resolution ────────────────────────────────────────────────────── + +/** Aggregation function names that ClickHouse cannot resolve as bare identifiers in ORDER BY. + * 'none' is excluded — it passes a raw expression through unchanged and has no synthesizable form. */ +const AGG_FN_NAMES: ReadonlySet = new Set( + MCP_AGG_FN_OPTIONS.filter(fn => fn !== 'none'), +); + +/** + * Resolve an orderBy value that matches a bare aggregation function name + * (e.g. "count") to something ClickHouse can resolve in ORDER BY. + * + * The model frequently writes `orderBy: "count"`, which generates + * `ORDER BY count ASC` in ClickHouse — but ClickHouse can't resolve bare + * `count` since it's a function, not a column. We fix this by finding + * the select item whose aggFn matches and using its alias, or synthesizing + * the ClickHouse expression (e.g. `count()`) when no alias is set. + * + * Resolution order: + * 1. If orderBy matches a select item's alias exactly → keep as-is + * 2. If orderBy matches an aggFn name → use that item's alias if set, + * otherwise synthesize the ClickHouse expression (e.g. `count()`) + * 3. Otherwise → pass through unchanged + */ +/** @internal Exported for testing only. */ +export function resolveOrderBy( + orderBy: string | undefined, + selectItems: { + aggFn: string; + alias?: string; + valueExpression?: string; + level?: number; + }[], +): string | undefined { + if (!orderBy) return undefined; + + // Strip an optional trailing ASC/DESC so we can resolve the identifier, + // then re-append the direction after resolution. + const dirMatch = orderBy.match(/^(.+?)\s+(ASC|DESC)\s*$/i); + const identifier = dirMatch ? dirMatch[1] : orderBy; + const direction = dirMatch ? ` ${dirMatch[2].toUpperCase()}` : ''; + + const lower = identifier.toLowerCase(); + + // Already matches an alias? Return the canonical alias case so + // ClickHouse's case-sensitive identifier resolution works correctly. + const aliasMatch = selectItems.find( + s => s.alias && s.alias.toLowerCase() === lower, + ); + if (aliasMatch) { + return `${aliasMatch.alias}${direction}`; + } + + // Matches an aggFn name? Resolve to that item's alias or synthesize. + if (AGG_FN_NAMES.has(lower)) { + const match = selectItems.find(s => s.aggFn.toLowerCase() === lower); + if (match) { + // Prefer the explicit alias if set + if (match.alias) { + return `${match.alias}${direction}`; + } + + // Synthesize the ClickHouse expression so ORDER BY works + if (match.aggFn === 'count') return `count()${direction}`; + // count_distinct compiles to count(DISTINCT expr) in ClickHouse, + // not count_distinct(expr) which is not a valid function. + if (match.aggFn === 'count_distinct' && match.valueExpression) { + return `count(DISTINCT ${match.valueExpression})${direction}`; + } + if ( + match.aggFn === 'quantile' && + match.level != null && + match.valueExpression + ) { + return `quantile(${match.level})(${match.valueExpression})${direction}`; + } + // Skip synthesis for quantile without level — let it pass through + // rather than generating invalid SQL like quantile(Duration) + if (match.aggFn === 'quantile') { + return orderBy; + } + if (match.valueExpression) { + return `${match.aggFn}(${match.valueExpression})${direction}`; + } + } + } + + return orderBy; +} + // ─── Tool registration ─────────────────────────────────────────────────────── export function registerTable(server: McpServer, context: McpContext) { @@ -55,9 +162,9 @@ export function registerTable(server: McpServer, context: McpContext) { 'Use this for grouped aggregations, top-N queries, single-value KPIs, ' + 'or proportional breakdowns.\n\n' + 'Requires sourceId — call hyperdx_list_sources then hyperdx_describe_source first.\n\n' + - 'Each select item can have its own "where" filter, which compiles to ' + - 'If(...) for multi-cohort comparison in one call ' + - '(e.g. p99_before + p99_after side-by-side).\n\n' + + 'Use the top-level "where" to scope the entire query (e.g. filter by service). ' + + 'Each select item can also have its own "where" for per-metric cohort ' + + 'comparisons (compiles to If(...)). Both can be used together.\n\n' + 'Column naming: top-level columns are PascalCase (Duration, StatusCode). ' + "Map attributes use bracket syntax: SpanAttributes['http.method']. " + 'Map attributes work in groupBy and valueExpression, including ' + @@ -86,15 +193,44 @@ export function registerTable(server: McpServer, context: McpContext) { displayType = 'table'; } + // Inject top-level where into each select item so it becomes part + // of the aggCondition for every metric. Table/line/number/pie display + // types don't have a chart-level where — filtering is per-select-item. + const { items: selectItems, warnings: mergeWarnings } = + mergeWhereIntoSelectItems( + input.select, + input.where, + input.whereLanguage, + ); + const tile = buildTile('MCP Table', 12, 4, { displayType, sourceId: input.sourceId, - select: input.select, + select: selectItems, groupBy: displayType === 'number' ? undefined : input.groupBy, - orderBy: input.orderBy, + orderBy: resolveOrderBy(input.orderBy, selectItems), }); - return runConfigTile(teamId.toString(), tile, startDate, endDate); + const result = await runConfigTile( + teamId.toString(), + tile, + startDate, + endDate, + ); + + // Surface language-mismatch warnings so the agent knows the top-level + // where was not applied to every select item. + if (mergeWarnings.length > 0 && result.content?.[0]?.type === 'text') { + try { + const parsed = JSON.parse(result.content[0].text); + parsed.warnings = mergeWarnings; + result.content[0].text = JSON.stringify(parsed, null, 2); + } catch { + // leave result unmodified + } + } + + return result; }), ); } diff --git a/packages/api/src/mcp/tools/query/timeseries.ts b/packages/api/src/mcp/tools/query/timeseries.ts index 4389d9d564..190ae2f58c 100644 --- a/packages/api/src/mcp/tools/query/timeseries.ts +++ b/packages/api/src/mcp/tools/query/timeseries.ts @@ -3,7 +3,12 @@ import { z } from 'zod'; import { withToolTracing } from '../../utils/tracing'; import type { McpContext } from '../types'; -import { buildTile, parseTimeRange, runConfigTile } from './helpers'; +import { + buildTile, + mergeWhereIntoSelectItems, + parseTimeRange, + runConfigTile, +} from './helpers'; import { endTimeSchema, groupBySchema, @@ -11,7 +16,10 @@ import { orderBySchema, sourceIdSchema, startTimeSchema, + whereLanguageSchema, + whereSchema, } from './schemas'; +import { resolveOrderBy } from './table'; // ─── Schema ────────────────────────────────────────────────────────────────── @@ -32,6 +40,13 @@ const timeseriesSchema = z.object({ .describe( 'Chart shape. "line" for line chart (default), "stacked_bar" for stacked bar chart.', ), + where: whereSchema.describe( + 'Row filter applied to ALL select items. ' + + 'Scopes the entire query — use to restrict by service, severity, etc. ' + + 'Each select item can also have its own per-metric "where" for cohort comparisons. ' + + 'When both are set, they are ANDed together.', + ), + whereLanguage: whereLanguageSchema, groupBy: groupBySchema, orderBy: orderBySchema, granularity: z @@ -77,16 +92,68 @@ export function registerTimeseries(server: McpServer, context: McpContext) { } const { startDate, endDate } = timeRange; + // Inject top-level where into each select item (timeseries uses + // per-select-item aggCondition, not chart-level where) + const { items: selectItems, warnings: mergeWarnings } = + mergeWhereIntoSelectItems( + input.select, + input.where, + input.whereLanguage, + ); + const tile = buildTile('MCP Timeseries', 12, 4, { displayType: input.shape, sourceId: input.sourceId, - select: input.select, + select: selectItems, groupBy: input.groupBy, - orderBy: input.orderBy, + orderBy: resolveOrderBy(input.orderBy, selectItems), ...(input.granularity ? { granularity: input.granularity } : {}), }); - return runConfigTile(teamId.toString(), tile, startDate, endDate); + const result = await runConfigTile( + teamId.toString(), + tile, + startDate, + endDate, + ); + + // Surface language-mismatch warnings so the agent knows the top-level + // where was not applied to every select item. + if (mergeWarnings.length > 0 && result.content?.[0]?.type === 'text') { + try { + const parsed = JSON.parse(result.content[0].text); + parsed.warnings = mergeWarnings; + result.content[0].text = JSON.stringify(parsed, null, 2); + } catch { + // leave result unmodified + } + } + + // Detect single-bucket collapse: when a timeseries query returns only + // 1 row, the data likely collapsed into a single time bucket. Add a + // hint so the agent knows to adjust the time range. + if ( + result.content?.[0]?.type === 'text' && + !('isError' in result && result.isError) + ) { + try { + const parsed = JSON.parse(result.content[0].text); + const data = parsed?.result?.data; + if (Array.isArray(data) && data.length === 1 && input.granularity) { + parsed.hint = + `Timeseries returned only 1 time bucket. ` + + `All data may have collapsed into a single "${input.granularity}" bucket. ` + + `The queried range was ${startDate.toISOString()} to ${endDate.toISOString()}. ` + + `If this looks wrong, adjust startTime/endTime to match the actual data range, ` + + `or try a coarser granularity.`; + result.content[0].text = JSON.stringify(parsed, null, 2); + } + } catch { + // If parsing fails, return the original result unmodified + } + } + + return result; }), ); } diff --git a/packages/api/src/mcp/tools/sources/describeSource.ts b/packages/api/src/mcp/tools/sources/describeSource.ts index 57dd26fd29..8f2fefd5ee 100644 --- a/packages/api/src/mcp/tools/sources/describeSource.ts +++ b/packages/api/src/mcp/tools/sources/describeSource.ts @@ -295,7 +295,7 @@ async function describeSourceSchema( : 'These are the REAL values in your data — use them in filters instead of guessing. ' + 'Example: where: "SeverityText:error" (if \'error\' appears in the sampled values above).'; - const output = trimToolResponse({ + const { data: output, isTrimmed } = trimToolResponse({ source: meta, usage: { topLevelColumns: @@ -311,11 +311,18 @@ async function describeSourceSchema( }, }); + const finalOutput = isTrimmed + ? { + ...output, + note: 'Result was trimmed for context size. Some columns or sampled values may be omitted.', + } + : output; + return { content: [ { type: 'text' as const, - text: JSON.stringify(output, null, 2), + text: JSON.stringify(finalOutput), }, ], }; diff --git a/packages/api/src/utils/__tests__/trimToolResponse.test.ts b/packages/api/src/utils/__tests__/trimToolResponse.test.ts index c6bd2eea5e..a427700700 100644 --- a/packages/api/src/utils/__tests__/trimToolResponse.test.ts +++ b/packages/api/src/utils/__tests__/trimToolResponse.test.ts @@ -1,37 +1,59 @@ import { trimToolResponse } from '../trimToolResponse'; describe('trimToolResponse', () => { + describe('return shape', () => { + it('should return { data, isTrimmed } object', () => { + const result = trimToolResponse(42); + expect(result).toHaveProperty('data'); + expect(result).toHaveProperty('isTrimmed'); + }); + + it('should set isTrimmed to false when data fits', () => { + expect(trimToolResponse(42).isTrimmed).toBe(false); + expect(trimToolResponse('hello').isTrimmed).toBe(false); + expect(trimToolResponse([1, 2, 3]).isTrimmed).toBe(false); + }); + + it('should set isTrimmed to true when data is trimmed', () => { + const largeArray = Array.from({ length: 500 }, (_, i) => ({ + id: i, + data: 'x'.repeat(200), + })); + expect(trimToolResponse(largeArray, 5000).isTrimmed).toBe(true); + }); + }); + describe('small data (within maxSize)', () => { it('should return primitive values unchanged', () => { - expect(trimToolResponse(42)).toBe(42); - expect(trimToolResponse('hello')).toBe('hello'); - expect(trimToolResponse(null)).toBeNull(); - expect(trimToolResponse(true)).toBe(true); + expect(trimToolResponse(42).data).toBe(42); + expect(trimToolResponse('hello').data).toBe('hello'); + expect(trimToolResponse(null).data).toBeNull(); + expect(trimToolResponse(true).data).toBe(true); }); it('should return small arrays unchanged', () => { const data = [1, 2, 3, 4, 5]; - expect(trimToolResponse(data)).toEqual(data); + expect(trimToolResponse(data).data).toEqual(data); }); it('should return small objects unchanged', () => { const data = { a: 1, b: 'hello', c: [1, 2, 3] }; - expect(trimToolResponse(data)).toEqual(data); + expect(trimToolResponse(data).data).toEqual(data); }); }); describe('large arrays', () => { it('should trim large arrays to fit within maxSize', () => { - // Create an array that exceeds maxSize const largeArray = Array.from({ length: 500 }, (_, i) => ({ id: i, data: 'x'.repeat(200), })); - const result = trimToolResponse(largeArray, 5000); + const { data: result, isTrimmed } = trimToolResponse(largeArray, 5000); + expect(isTrimmed).toBe(true); expect(Array.isArray(result)).toBe(true); expect(result.length).toBeLessThan(largeArray.length); - expect(result.length).toBeGreaterThanOrEqual(10); // minimum 10 items + expect(result.length).toBeGreaterThanOrEqual(10); expect(JSON.stringify(result).length).toBeLessThanOrEqual(5000); }); @@ -41,15 +63,15 @@ describe('trimToolResponse', () => { data: 'x'.repeat(500), })); - // maxSize so small even 10 items may exceed it, but we keep 10 minimum - const result = trimToolResponse(largeArray, 100); + const { data: result } = trimToolResponse(largeArray, 100); expect(Array.isArray(result)).toBe(true); expect(result.length).toBeGreaterThanOrEqual(10); }); it('should not trim arrays that fit within maxSize', () => { const smallArray = [1, 2, 3, 4, 5]; - const result = trimToolResponse(smallArray, 50000); + const { data: result, isTrimmed } = trimToolResponse(smallArray, 50000); + expect(isTrimmed).toBe(false); expect(result).toEqual(smallArray); }); }); @@ -61,22 +83,21 @@ describe('trimToolResponse', () => { largeObj[`key_${i}`] = 'x'.repeat(200); } - const result = trimToolResponse(largeObj, 5000); - // The trimmed result must be smaller than the original + const { data: result, isTrimmed } = trimToolResponse(largeObj, 5000); + expect(isTrimmed).toBe(true); expect(JSON.stringify(result).length).toBeLessThan( JSON.stringify(largeObj).length, ); - // All keys should still be present (values are trimmed, not dropped) expect( Object.keys(result).filter(k => k !== '__hdx_trimmed'), ).toHaveLength(100); - // The sentinel flag should be set to indicate trimming occurred expect(result.__hdx_trimmed).toBe(true); }); it('should not trim objects that fit within maxSize', () => { const obj = { a: 1, b: 2 }; - const result = trimToolResponse(obj, 50000); + const { data: result, isTrimmed } = trimToolResponse(obj, 50000); + expect(isTrimmed).toBe(false); expect(result).toEqual(obj); }); }); @@ -95,7 +116,7 @@ describe('trimToolResponse', () => { otherProp: 'preserved', }; - const result = trimToolResponse(metadataObj, 5000); + const { data: result } = trimToolResponse(metadataObj, 5000); expect(result).toHaveProperty('allFieldsWithKeys'); expect(result).toHaveProperty('keyValues'); expect(result).toHaveProperty('otherProp', 'preserved'); @@ -106,7 +127,6 @@ describe('trimToolResponse', () => { describe('default maxSize', () => { it('should use 50000 as default maxSize', () => { - // Create data just over default size const data = Array.from({ length: 1000 }, (_, i) => ({ id: i, payload: 'x'.repeat(100), @@ -114,8 +134,7 @@ describe('trimToolResponse', () => { const resultDefault = trimToolResponse(data); const resultExplicit = trimToolResponse(data, 50000); - // Both should produce the same result - expect(resultDefault.length).toBe(resultExplicit.length); + expect(resultDefault.data.length).toBe(resultExplicit.data.length); }); }); }); diff --git a/packages/api/src/utils/trimToolResponse.ts b/packages/api/src/utils/trimToolResponse.ts index b410839986..3ab9377abb 100644 --- a/packages/api/src/utils/trimToolResponse.ts +++ b/packages/api/src/utils/trimToolResponse.ts @@ -1,15 +1,26 @@ import logger from '@/utils/logger'; +export interface TrimResult { + data: T; + isTrimmed: boolean; +} + /** * Trims large data structures to prevent "Request Entity Too Large" errors * when multiple tool calls accumulate data in the conversation history. + * + * Returns `{ data, isTrimmed }` so callers can cheaply check whether the + * response was truncated without re-serializing. */ -export function trimToolResponse(data: any, maxSize: number = 50000): any { +export function trimToolResponse( + data: T, + maxSize: number = 50000, +): TrimResult { const serialized = JSON.stringify(data); // If data is within acceptable size, return as-is if (serialized.length <= maxSize) { - return data; + return { data, isTrimmed: false }; } logger.warn( @@ -18,14 +29,20 @@ export function trimToolResponse(data: any, maxSize: number = 50000): any { // Handle different data structures if (Array.isArray(data)) { - return trimArray(data, maxSize); + const trimmed = + // eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion + trimArray(data, maxSize) as T; + return { data: trimmed, isTrimmed: true }; } if (typeof data === 'object' && data !== null) { - return trimObject(data, maxSize); + const trimmed = + // eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion + trimObject(data, maxSize) as T; + return { data: trimmed, isTrimmed: true }; } - return data; + return { data, isTrimmed: false }; } function trimArray(arr: any[], maxSize: number): any[] {