diff --git a/.changeset/nasty-crews-shop.md b/.changeset/nasty-crews-shop.md new file mode 100644 index 0000000000..0b53ead24f --- /dev/null +++ b/.changeset/nasty-crews-shop.md @@ -0,0 +1,5 @@ +--- +"@hyperdx/api": patch +--- + +fix(mcp): add ClickHouse safety settings (max_execution_time, max_result_rows, readonly) for MCP query execution diff --git a/packages/api/src/mcp/__tests__/query.test.ts b/packages/api/src/mcp/__tests__/query.test.ts index 214adc058b..1a9c7ac55c 100644 --- a/packages/api/src/mcp/__tests__/query.test.ts +++ b/packages/api/src/mcp/__tests__/query.test.ts @@ -225,6 +225,19 @@ describe('errorHint', () => { expect(hint).toContain('LIMIT'); }); + it('should match RESULT_IS_TOO_LARGE errors', () => { + const hint = errorHint('Code: 396. DB::Exception: RESULT_IS_TOO_LARGE'); + expect(hint).not.toBeNull(); + expect(hint).toContain('100,000 rows'); + expect(hint).toContain('LIMIT'); + }); + + it('should match TOO_MANY_ROWS_OR_BYTES errors', () => { + const hint = errorHint('Code: 396. DB::Exception: TOO_MANY_ROWS_OR_BYTES'); + expect(hint).not.toBeNull(); + expect(hint).toContain('100,000 rows'); + }); + it('should return null for unrecognized errors', () => { const hint = errorHint('Connection refused'); expect(hint).toBeNull(); diff --git a/packages/api/src/mcp/__tests__/queryTool.test.ts b/packages/api/src/mcp/__tests__/queryTool.test.ts index b5264580b5..4cdfd6563a 100644 --- a/packages/api/src/mcp/__tests__/queryTool.test.ts +++ b/packages/api/src/mcp/__tests__/queryTool.test.ts @@ -1,3 +1,4 @@ +import { ClickhouseClient } from '@hyperdx/common-utils/dist/clickhouse/node'; import { SourceKind } from '@hyperdx/common-utils/dist/types'; import { Client } from '@modelcontextprotocol/sdk/client/index.js'; @@ -583,4 +584,110 @@ describe('MCP Query Tools', () => { expect(getFirstText(result)).toContain('Invalid'); }); }); + + // ─── Safety settings (readonly + max_execution_time) ───────────────────────── + + describe('ClickHouse safety settings', () => { + describe('readonly enforcement via clickstack_sql', () => { + it('should reject CREATE TABLE (DDL)', async () => { + const result = await callTool(client, 'clickstack_sql', { + connectionId: connection._id.toString(), + sql: 'CREATE TABLE __mcp_test_ddl (id UInt64) ENGINE = Memory', + }); + + expect(result.isError).toBe(true); + const text = getFirstText(result); + expect(text).toMatch(/readonly/i); + }); + + it('should reject INSERT (DML)', async () => { + const result = await callTool(client, 'clickstack_sql', { + connectionId: connection._id.toString(), + sql: `INSERT INTO ${DEFAULT_DATABASE}.${DEFAULT_LOGS_TABLE} (Body) VALUES ('injected')`, + }); + + expect(result.isError).toBe(true); + const text = getFirstText(result); + expect(text).toMatch(/readonly/i); + }); + + it('should reject DROP TABLE (DDL)', async () => { + const result = await callTool(client, 'clickstack_sql', { + connectionId: connection._id.toString(), + sql: `DROP TABLE IF EXISTS ${DEFAULT_DATABASE}.${DEFAULT_LOGS_TABLE}`, + }); + + expect(result.isError).toBe(true); + const text = getFirstText(result); + expect(text).toMatch(/readonly/i); + }); + + it('should allow SELECT (read-only query)', async () => { + const result = await callTool(client, 'clickstack_sql', { + connectionId: connection._id.toString(), + sql: 'SELECT 1 AS value', + }); + + expect(result.isError).toBeFalsy(); + }); + }); + + describe('max_execution_time enforcement', () => { + it('should kill a query that exceeds max_execution_time', async () => { + // Test directly with ClickhouseClient to use a short timeout (1s) + // instead of waiting for the full 30s MCP default. + const clickhouseClient = new ClickhouseClient({ + host: config.CLICKHOUSE_HOST, + username: config.CLICKHOUSE_USER, + password: config.CLICKHOUSE_PASSWORD, + requestTimeout: 5_000, + }); + + await expect( + clickhouseClient.query({ + query: 'SELECT sleep(3)', + format: 'JSON', + clickhouse_settings: { + max_execution_time: 1, + }, + }), + ).rejects.toThrow(/TIMEOUT_EXCEEDED|timeout/i); + }); + }); + + describe('settings propagation via clickstack_sql', () => { + it('should propagate max_execution_time=30 to ClickHouse', async () => { + const result = await callTool(client, 'clickstack_sql', { + connectionId: connection._id.toString(), + sql: "SELECT getSetting('max_execution_time') AS value", + }); + + expect(result.isError).toBeFalsy(); + const parsed = JSON.parse(getFirstText(result)); + expect(parsed.result?.data?.[0]?.value).toBe(30); + }); + + it('should propagate max_result_rows=100000 to ClickHouse', async () => { + const result = await callTool(client, 'clickstack_sql', { + connectionId: connection._id.toString(), + sql: "SELECT getSetting('max_result_rows') AS value", + }); + + expect(result.isError).toBeFalsy(); + const parsed = JSON.parse(getFirstText(result)); + expect(parsed.result?.data?.[0]?.value).toBe(100000); + }); + + it('should propagate readonly=1 to ClickHouse', async () => { + const result = await callTool(client, 'clickstack_sql', { + connectionId: connection._id.toString(), + sql: "SELECT getSetting('readonly') AS value", + }); + + expect(result.isError).toBeFalsy(); + const parsed = JSON.parse(getFirstText(result)); + expect(parsed.result?.data?.[0]?.value).toBe(1); + }); + }); + }); }); diff --git a/packages/api/src/mcp/tools/query/helpers.ts b/packages/api/src/mcp/tools/query/helpers.ts index 7fbd3d200a..c87fc82066 100644 --- a/packages/api/src/mcp/tools/query/helpers.ts +++ b/packages/api/src/mcp/tools/query/helpers.ts @@ -27,6 +27,22 @@ import { trimToolResponse } from '@/utils/trimToolResponse'; import type { ExternalDashboardTileWithId } from '@/utils/zod'; import { externalDashboardTileSchemaWithId } from '@/utils/zod'; +// ─── Safety limits ─────────────────────────────────────────────────────────── + +/** ClickHouse settings applied to all MCP query-tool executions. */ +const MCP_CLICKHOUSE_SETTINGS = { + max_execution_time: 30, + max_result_rows: '100000', + readonly: 1, +} as const; + +/** + * 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. + */ +const MCP_REQUEST_TIMEOUT = 32_000; // 30s query limit + 2s buffer + // ─── Where merging ─────────────────────────────────────────────────────────── export interface MergeWhereResult { @@ -243,6 +259,7 @@ export async function runConfigTile( host: connection.host, username: connection.username, password: connection.password, + requestTimeout: MCP_REQUEST_TIMEOUT, }); const isSearch = builderConfig.displayType === DisplayType.Search; @@ -295,6 +312,7 @@ export async function runConfigTile( config: chartConfig, metadata, querySettings: source.querySettings, + opts: { clickhouse_settings: MCP_CLICKHOUSE_SETTINGS }, }); return formatQueryResult(result); } catch (e) { @@ -349,6 +367,7 @@ export async function runConfigTile( host: connection.host, username: connection.username, password: connection.password, + requestTimeout: MCP_REQUEST_TIMEOUT, }); const chartConfig = { @@ -363,6 +382,7 @@ export async function runConfigTile( config: chartConfig, metadata, querySettings: undefined, + opts: { clickhouse_settings: MCP_CLICKHOUSE_SETTINGS }, }); return formatQueryResult(result); } catch (e) { @@ -419,5 +439,11 @@ export function errorHint(msg: string): string | null { 'The result row count is too large to serialize back to the agent.' ); } + if (/TOO_MANY_ROWS_OR_BYTES|RESULT_IS_TOO_LARGE/i.test(msg)) { + return ( + 'The query returned more than 100,000 rows. ' + + 'Add a LIMIT, narrow the time range, or add filters to reduce the result set.' + ); + } return null; }