diff --git a/packages/api/src/mcp/mcpServer.ts b/packages/api/src/mcp/mcpServer.ts index 885cc245e2..3f65438259 100644 --- a/packages/api/src/mcp/mcpServer.ts +++ b/packages/api/src/mcp/mcpServer.ts @@ -10,6 +10,7 @@ import savedSearchesTools from './tools/savedSearches/index'; import sourcesTools from './tools/sources/index'; import traceTools from './tools/trace/index'; import { McpContext } from './tools/types'; +import { createRegisterTool } from './utils/registerTool'; export function createServer(context: McpContext) { const server = new McpServer({ @@ -17,12 +18,15 @@ export function createServer(context: McpContext) { version: `${CODE_VERSION}-beta`, }); - sourcesTools(server, context); - alertsTools(server, context); - dashboardsTools(server, context); - queryTools(server, context); - savedSearchesTools(server, context); - traceTools(server, context); + const registerTool = createRegisterTool(server, context); + const registrar = { server, context, registerTool }; + + sourcesTools(registrar); + alertsTools(registrar); + dashboardsTools(registrar); + queryTools(registrar); + savedSearchesTools(registrar); + traceTools(registrar); dashboardPrompts(server, context); return server; diff --git a/packages/api/src/mcp/tools/alerts/getAlert.ts b/packages/api/src/mcp/tools/alerts/getAlert.ts index d14b343ed1..443a403952 100644 --- a/packages/api/src/mcp/tools/alerts/getAlert.ts +++ b/packages/api/src/mcp/tools/alerts/getAlert.ts @@ -1,5 +1,4 @@ import { type AlertInterval } from '@hyperdx/common-utils/dist/types'; -import type { McpServer } from '@modelcontextprotocol/sdk/server/mcp.js'; import { ObjectId } from 'mongodb'; import mongoose from 'mongoose'; import { z } from 'zod'; @@ -7,9 +6,8 @@ import { z } from 'zod'; import * as config from '@/config'; import { getRecentAlertHistories } from '@/controllers/alertHistory'; import { getAlertById } from '@/controllers/alerts'; -import type { McpContext } from '@/mcp/tools/types'; +import type { ToolRegistrar } from '@/mcp/tools/types'; import { validateObjectId } from '@/mcp/utils/errors'; -import { withToolTracing } from '@/mcp/utils/tracing'; import Alert from '@/models/alert'; import type { IDashboard } from '@/models/dashboard'; import type { ISavedSearch } from '@/models/savedSearch'; @@ -45,11 +43,14 @@ function deriveAlertName(alert: { return null; } -export function registerGetAlert(server: McpServer, context: McpContext): void { +export function registerGetAlert({ + context, + registerTool, +}: ToolRegistrar): void { const { teamId } = context; const frontendUrl = config.FRONTEND_URL; - server.registerTool( + registerTool( 'clickstack_get_alert', { title: 'Get Alert(s)', @@ -75,7 +76,7 @@ export function registerGetAlert(server: McpServer, context: McpContext): void { ), }), }, - withToolTracing('clickstack_get_alert', context, async ({ id, state }) => { + async ({ id, state }) => { // ── List all alerts (slim summary) ── if (!id) { const query: Record = { @@ -151,6 +152,6 @@ export function registerGetAlert(server: McpServer, context: McpContext): void { }, ], }; - }), + }, ); } diff --git a/packages/api/src/mcp/tools/alerts/getWebhook.ts b/packages/api/src/mcp/tools/alerts/getWebhook.ts index d5c11b97d3..ed71f9f075 100644 --- a/packages/api/src/mcp/tools/alerts/getWebhook.ts +++ b/packages/api/src/mcp/tools/alerts/getWebhook.ts @@ -1,17 +1,15 @@ -import type { McpServer } from '@modelcontextprotocol/sdk/server/mcp.js'; import { z } from 'zod'; -import type { McpContext } from '@/mcp/tools/types'; -import { withToolTracing } from '@/mcp/utils/tracing'; +import type { ToolRegistrar } from '@/mcp/tools/types'; import Webhook from '@/models/webhook'; -export function registerGetWebhook( - server: McpServer, - context: McpContext, -): void { +export function registerGetWebhook({ + context, + registerTool, +}: ToolRegistrar): void { const { teamId } = context; - server.registerTool( + registerTool( 'clickstack_get_webhook', { title: 'List Webhooks', @@ -21,7 +19,7 @@ export function registerGetWebhook( 'clickstack_save_alert.', inputSchema: z.object({}), }, - withToolTracing('clickstack_get_webhook', context, async () => { + async () => { const webhooks = await Webhook.find({ team: teamId }); const output = webhooks.map(wh => ({ @@ -35,6 +33,6 @@ export function registerGetWebhook( { type: 'text' as const, text: JSON.stringify(output, null, 2) }, ], }; - }), + }, ); } diff --git a/packages/api/src/mcp/tools/alerts/index.ts b/packages/api/src/mcp/tools/alerts/index.ts index 7cdfe138ab..d8ec15a8b8 100644 --- a/packages/api/src/mcp/tools/alerts/index.ts +++ b/packages/api/src/mcp/tools/alerts/index.ts @@ -1,18 +1,13 @@ -import type { McpServer } from '@modelcontextprotocol/sdk/server/mcp.js'; - -import type { McpContext, ToolDefinition } from '@/mcp/tools/types'; +import type { ToolDefinition, ToolRegistrar } from '@/mcp/tools/types'; import { registerGetAlert } from './getAlert'; import { registerGetWebhook } from './getWebhook'; import { registerSaveAlert } from './saveAlert'; -const alertsTools: ToolDefinition = ( - server: McpServer, - context: McpContext, -) => { - registerGetAlert(server, context); - registerGetWebhook(server, context); - registerSaveAlert(server, context); +const alertsTools: ToolDefinition = (registrar: ToolRegistrar) => { + registerGetAlert(registrar); + registerGetWebhook(registrar); + registerSaveAlert(registrar); }; export default alertsTools; diff --git a/packages/api/src/mcp/tools/alerts/saveAlert.ts b/packages/api/src/mcp/tools/alerts/saveAlert.ts index 94548798fd..0494c98bc1 100644 --- a/packages/api/src/mcp/tools/alerts/saveAlert.ts +++ b/packages/api/src/mcp/tools/alerts/saveAlert.ts @@ -1,5 +1,4 @@ import { AlertThresholdType } from '@hyperdx/common-utils/dist/types'; -import type { McpServer } from '@modelcontextprotocol/sdk/server/mcp.js'; import mongoose from 'mongoose'; import * as config from '@/config'; @@ -9,9 +8,8 @@ import { updateAlert, validateAlertInput, } from '@/controllers/alerts'; -import type { McpContext } from '@/mcp/tools/types'; +import type { ToolRegistrar } from '@/mcp/tools/types'; import { mcpError, validateObjectId } from '@/mcp/utils/errors'; -import { withToolTracing } from '@/mcp/utils/tracing'; import { type AlertChannel, AlertSource } from '@/models/alert'; import { BaseError } from '@/utils/errors'; import { translateAlertDocumentToExternalAlert } from '@/utils/externalApi'; @@ -33,14 +31,14 @@ function toAlertChannel(ch: McpSaveAlertInput['channel']): AlertChannel { }; } -export function registerSaveAlert( - server: McpServer, - context: McpContext, -): void { +export function registerSaveAlert({ + context, + registerTool, +}: ToolRegistrar): void { const { teamId, userId } = context; const frontendUrl = config.FRONTEND_URL; - server.registerTool( + registerTool( 'clickstack_save_alert', { title: 'Create or Update Alert', @@ -50,7 +48,7 @@ export function registerSaveAlert( 'metric crosses a threshold. A webhook notification channel is required.', inputSchema: mcpSaveAlertSchema, }, - withToolTracing('clickstack_save_alert', context, async input => { + async input => { // ── Runtime cross-field validation ── const validationError = validateSaveAlertInput(input); if (validationError) { @@ -147,6 +145,6 @@ export function registerSaveAlert( }, ], }; - }), + }, ); } diff --git a/packages/api/src/mcp/tools/dashboards/deleteDashboard.ts b/packages/api/src/mcp/tools/dashboards/deleteDashboard.ts index fd087729c5..88236f951d 100644 --- a/packages/api/src/mcp/tools/dashboards/deleteDashboard.ts +++ b/packages/api/src/mcp/tools/dashboards/deleteDashboard.ts @@ -1,20 +1,18 @@ -import type { McpServer } from '@modelcontextprotocol/sdk/server/mcp.js'; import mongoose from 'mongoose'; import { z } from 'zod'; import { deleteDashboard } from '@/controllers/dashboard'; -import type { McpContext } from '@/mcp/tools/types'; -import { withToolTracing } from '@/mcp/utils/tracing'; +import type { ToolRegistrar } from '@/mcp/tools/types'; import Dashboard from '@/models/dashboard'; import { objectIdSchema } from '@/utils/zod'; -export function registerDeleteDashboard( - server: McpServer, - context: McpContext, -): void { +export function registerDeleteDashboard({ + context, + registerTool, +}: ToolRegistrar): void { const { teamId } = context; - server.registerTool( + registerTool( 'clickstack_delete_dashboard', { title: 'Delete Dashboard', @@ -25,32 +23,28 @@ export function registerDeleteDashboard( id: objectIdSchema.describe('Dashboard ID to delete.'), }), }, - withToolTracing( - 'clickstack_delete_dashboard', - context, - async ({ id: dashboardId }) => { - const existing = await Dashboard.findOne({ - _id: dashboardId, - team: teamId, - }).lean(); - if (!existing) { - return { - isError: true, - content: [{ type: 'text' as const, text: 'Dashboard not found' }], - }; - } - - await deleteDashboard(dashboardId, new mongoose.Types.ObjectId(teamId)); - + async ({ id: dashboardId }) => { + const existing = await Dashboard.findOne({ + _id: dashboardId, + team: teamId, + }).lean(); + if (!existing) { return { - content: [ - { - type: 'text' as const, - text: JSON.stringify({ deleted: true, id: dashboardId }, null, 2), - }, - ], + isError: true, + content: [{ type: 'text' as const, text: 'Dashboard not found' }], }; - }, - ), + } + + await deleteDashboard(dashboardId, new mongoose.Types.ObjectId(teamId)); + + return { + content: [ + { + type: 'text' as const, + text: JSON.stringify({ deleted: true, id: dashboardId }, null, 2), + }, + ], + }; + }, ); } diff --git a/packages/api/src/mcp/tools/dashboards/getDashboard.ts b/packages/api/src/mcp/tools/dashboards/getDashboard.ts index aeae9a8618..5d7b0a8e4f 100644 --- a/packages/api/src/mcp/tools/dashboards/getDashboard.ts +++ b/packages/api/src/mcp/tools/dashboards/getDashboard.ts @@ -1,23 +1,21 @@ -import type { McpServer } from '@modelcontextprotocol/sdk/server/mcp.js'; import mongoose from 'mongoose'; import { z } from 'zod'; import * as config from '@/config'; import { getDashboards } from '@/controllers/dashboard'; -import type { McpContext } from '@/mcp/tools/types'; +import type { ToolRegistrar } from '@/mcp/tools/types'; import { validateObjectId } from '@/mcp/utils/errors'; -import { withToolTracing } from '@/mcp/utils/tracing'; import Dashboard from '@/models/dashboard'; import { convertToExternalDashboard } from '@/routers/external-api/v2/utils/dashboards'; -export function registerGetDashboard( - server: McpServer, - context: McpContext, -): void { +export function registerGetDashboard({ + context, + registerTool, +}: ToolRegistrar): void { const { teamId } = context; const frontendUrl = config.FRONTEND_URL; - server.registerTool( + registerTool( 'clickstack_get_dashboard', { title: 'Get Dashboard(s)', @@ -33,7 +31,7 @@ export function registerGetDashboard( ), }), }, - withToolTracing('clickstack_get_dashboard', context, async ({ id }) => { + async ({ id }) => { if (!id) { const dashboards = await getDashboards( new mongoose.Types.ObjectId(teamId), @@ -78,6 +76,6 @@ export function registerGetDashboard( }, ], }; - }), + }, ); } diff --git a/packages/api/src/mcp/tools/dashboards/getDashboardTile.ts b/packages/api/src/mcp/tools/dashboards/getDashboardTile.ts index 0b36b7ec60..4e4e926cca 100644 --- a/packages/api/src/mcp/tools/dashboards/getDashboardTile.ts +++ b/packages/api/src/mcp/tools/dashboards/getDashboardTile.ts @@ -1,19 +1,17 @@ -import type { McpServer } from '@modelcontextprotocol/sdk/server/mcp.js'; import { z } from 'zod'; -import type { McpContext } from '@/mcp/tools/types'; -import { withToolTracing } from '@/mcp/utils/tracing'; +import type { ToolRegistrar } from '@/mcp/tools/types'; import Dashboard from '@/models/dashboard'; import { convertToExternalDashboard } from '@/routers/external-api/v2/utils/dashboards'; import { objectIdSchema } from '@/utils/zod'; -export function registerGetDashboardTile( - server: McpServer, - context: McpContext, -): void { +export function registerGetDashboardTile({ + context, + registerTool, +}: ToolRegistrar): void { const { teamId } = context; - server.registerTool( + registerTool( 'clickstack_get_dashboard_tile', { title: 'Get a Single Dashboard Tile', @@ -32,44 +30,40 @@ export function registerGetDashboardTile( ), }), }, - withToolTracing( - 'clickstack_get_dashboard_tile', - context, - async ({ dashboardId, tileId }) => { - const dashboard = await Dashboard.findOne({ - _id: dashboardId, - team: teamId, - }); - if (!dashboard) { - return { - isError: true, - content: [{ type: 'text' as const, text: 'Dashboard not found' }], - }; - } - - const externalDashboard = convertToExternalDashboard(dashboard); - const tile = externalDashboard.tiles.find(t => t.id === tileId); - if (!tile) { - return { - isError: true, - content: [ - { - type: 'text' as const, - text: `Tile not found: ${tileId}. Available tile IDs: ${externalDashboard.tiles.map(t => `${t.id} (${t.name})`).join(', ')}`, - }, - ], - }; - } + async ({ dashboardId, tileId }) => { + const dashboard = await Dashboard.findOne({ + _id: dashboardId, + team: teamId, + }); + if (!dashboard) { + return { + isError: true, + content: [{ type: 'text' as const, text: 'Dashboard not found' }], + }; + } + const externalDashboard = convertToExternalDashboard(dashboard); + const tile = externalDashboard.tiles.find(t => t.id === tileId); + if (!tile) { return { + isError: true, content: [ { type: 'text' as const, - text: JSON.stringify(tile, null, 2), + text: `Tile not found: ${tileId}. Available tile IDs: ${externalDashboard.tiles.map(t => `${t.id} (${t.name})`).join(', ')}`, }, ], }; - }, - ), + } + + return { + content: [ + { + type: 'text' as const, + text: JSON.stringify(tile, null, 2), + }, + ], + }; + }, ); } diff --git a/packages/api/src/mcp/tools/dashboards/index.ts b/packages/api/src/mcp/tools/dashboards/index.ts index 76a62a8e47..9cf52fa50a 100644 --- a/packages/api/src/mcp/tools/dashboards/index.ts +++ b/packages/api/src/mcp/tools/dashboards/index.ts @@ -1,6 +1,4 @@ -import type { McpServer } from '@modelcontextprotocol/sdk/server/mcp.js'; - -import type { McpContext, ToolDefinition } from '@/mcp/tools/types'; +import type { ToolDefinition, ToolRegistrar } from '@/mcp/tools/types'; import { registerDeleteDashboard } from './deleteDashboard'; import { registerGetDashboard } from './getDashboard'; @@ -12,17 +10,14 @@ import { registerSearchDashboards } from './searchDashboards'; export * from './schemas'; -const dashboardsTools: ToolDefinition = ( - server: McpServer, - context: McpContext, -) => { - registerGetDashboard(server, context); - registerGetDashboardTile(server, context); - registerSaveDashboard(server, context); - registerPatchDashboard(server, context); - registerDeleteDashboard(server, context); - registerSearchDashboards(server, context); - registerQueryTile(server, context); +const dashboardsTools: ToolDefinition = (registrar: ToolRegistrar) => { + registerGetDashboard(registrar); + registerGetDashboardTile(registrar); + registerSaveDashboard(registrar); + registerPatchDashboard(registrar); + registerDeleteDashboard(registrar); + registerSearchDashboards(registrar); + registerQueryTile(registrar); }; export default dashboardsTools; diff --git a/packages/api/src/mcp/tools/dashboards/patchDashboard.ts b/packages/api/src/mcp/tools/dashboards/patchDashboard.ts index f0fccdfbdd..2f1c34696d 100644 --- a/packages/api/src/mcp/tools/dashboards/patchDashboard.ts +++ b/packages/api/src/mcp/tools/dashboards/patchDashboard.ts @@ -1,10 +1,8 @@ -import type { McpServer } from '@modelcontextprotocol/sdk/server/mcp.js'; import { uniq } from 'lodash'; import * as config from '@/config'; -import type { McpContext } from '@/mcp/tools/types'; +import type { ToolRegistrar } from '@/mcp/tools/types'; import { mcpError } from '@/mcp/utils/errors'; -import { withToolTracing } from '@/mcp/utils/tracing'; import Dashboard from '@/models/dashboard'; import { cleanupDashboardAlerts, @@ -21,14 +19,14 @@ import { getRawSqlTileMacroWarnings, } from './validation'; -export function registerPatchDashboard( - server: McpServer, - context: McpContext, -): void { +export function registerPatchDashboard({ + context, + registerTool, +}: ToolRegistrar): void { const { teamId } = context; const frontendUrl = config.FRONTEND_URL; - server.registerTool( + registerTool( 'clickstack_patch_dashboard', { title: 'Patch Dashboard', @@ -40,257 +38,251 @@ export function registerPatchDashboard( 'IMPORTANT: After patching a tile, run clickstack_query_tile to confirm the query still works.', inputSchema: mcpPatchDashboardSchema, }, - withToolTracing( - 'clickstack_patch_dashboard', - context, - async ({ dashboardId, name, tags, tileId, tile: inputTile }) => { - // Cross-field validation (kept in handler so the inputSchema - // stays a plain z.object and its properties are visible in the - // JSON Schema that the MCP SDK exposes to LLMs). - if ( - name === undefined && - tags === undefined && - (tileId === undefined || inputTile === undefined) - ) { + async ({ dashboardId, name, tags, tileId, tile: inputTile }) => { + // Cross-field validation (kept in handler so the inputSchema + // stays a plain z.object and its properties are visible in the + // JSON Schema that the MCP SDK exposes to LLMs). + if ( + name === undefined && + tags === undefined && + (tileId === undefined || inputTile === undefined) + ) { + return { + isError: true, + content: [ + { + type: 'text' as const, + text: 'Provide at least one of: name, tags, or tileId+tile to patch.', + }, + ], + }; + } + if ((tileId === undefined) !== (inputTile === undefined)) { + return { + isError: true, + content: [ + { + type: 'text' as const, + text: 'tileId and tile must both be provided or both omitted.', + }, + ], + }; + } + + const existingDashboard = await Dashboard.findOne({ + _id: dashboardId, + team: teamId, + }); + if (!existingDashboard) { + return { + isError: true, + content: [{ type: 'text' as const, text: 'Dashboard not found' }], + }; + } + + // Build the $set payload and the query filter. Metadata fields + // are simple top-level $set entries; the tile patch uses the + // positional $ operator matched by 'tiles.id' in the filter. + const setPayload: Record = {}; + const queryFilter: Record = { + _id: dashboardId, + team: teamId, + }; + + if (name !== undefined) { + setPayload.name = name; + } + if (tags !== undefined) { + setPayload.tags = uniq(tags); + } + + let patchedTile: ExternalDashboardTileWithId | undefined; + + // ── Tile-level patch ───────────────────────────────────────── + if (tileId !== undefined && inputTile !== undefined) { + // Work directly with the persisted internal tiles array so + // untouched tiles are never round-tripped through the external + // converter (which would strip orphaned container refs from + // unrelated tiles). + const internalTiles = + (existingDashboard.tiles as { id: string }[]) ?? []; + const existingIdx = internalTiles.findIndex(t => t.id === tileId); + if (existingIdx === -1) { + // Build a human-readable list of available tile IDs. + // Convert only for the error message (read-only, no write-back). + const externalDashboard = + convertToExternalDashboard(existingDashboard); return { isError: true, content: [ { type: 'text' as const, - text: 'Provide at least one of: name, tags, or tileId+tile to patch.', + text: `Tile not found: ${tileId}. Available tile IDs: ${externalDashboard.tiles.map(t => `${t.id} (${t.name})`).join(', ')}`, }, ], }; } - if ((tileId === undefined) !== (inputTile === undefined)) { - return { - isError: true, - content: [ - { - type: 'text' as const, - text: 'tileId and tile must both be provided or both omitted.', - }, - ], - }; + + // Read the existing tile's layout and container refs from the + // persisted internal format so fallback values are accurate + // regardless of the external converter's self-healing logic. + const existingInternalTile = internalTiles[existingIdx] as { + id: string; + x?: number; + y?: number; + w?: number; + h?: number; + containerId?: string; + tabId?: string; + config?: { name?: string }; + }; + + // Merge: the incoming tile definition replaces config/name, + // but layout and container refs fall back to the existing tile + // when omitted so the LLM doesn't have to re-specify them. + // Coerce legacy empty-string containerId/tabId to undefined so + // they don't trip the container-ref validator (mirrors the + // self-heal in convertTileToExternalChart). + const incoming = inputTile as Partial; + const existingContainerId = + existingInternalTile.containerId || undefined; + const existingTabId = existingInternalTile.tabId || undefined; + const mergedTile: ExternalDashboardTileWithId = { + id: tileId, + name: incoming.name ?? existingInternalTile.config?.name ?? '', + x: incoming.x ?? existingInternalTile.x ?? 0, + y: incoming.y ?? existingInternalTile.y ?? 0, + w: incoming.w ?? existingInternalTile.w ?? 12, + h: incoming.h ?? existingInternalTile.h ?? 4, + containerId: + 'containerId' in incoming + ? incoming.containerId + : existingContainerId, + tabId: 'tabId' in incoming ? incoming.tabId : existingTabId, + // The config comes from the incoming tile (validated by Zod). + config: incoming.config, + } as ExternalDashboardTileWithId; + + // Error on raw SQL tiles that have no source defined but which use + // macros which require a source to be set + const sqlFilterSourceError = getRawSqlMissingSourceError([mergedTile]); + if (sqlFilterSourceError) { + return mcpError(sqlFilterSourceError); } - const existingDashboard = await Dashboard.findOne({ - _id: dashboardId, - team: teamId, + // Validate the patched tile using the shared validation helper. + const validationError = await validateDashboardTiles({ + teamId, + tiles: [mergedTile], + existingTiles: existingDashboard.tiles ?? [], + containers: existingDashboard.containers ?? [], }); - if (!existingDashboard) { + if (validationError) { return { isError: true, - content: [{ type: 'text' as const, text: 'Dashboard not found' }], + content: [{ type: 'text' as const, text: validationError }], }; } - // Build the $set payload and the query filter. Metadata fields - // are simple top-level $set entries; the tile patch uses the - // positional $ operator matched by 'tiles.id' in the filter. - const setPayload: Record = {}; - const queryFilter: Record = { - _id: dashboardId, - team: teamId, - }; - - if (name !== undefined) { - setPayload.name = name; - } - if (tags !== undefined) { - setPayload.tags = uniq(tags); - } - - let patchedTile: ExternalDashboardTileWithId | undefined; - - // ── Tile-level patch ───────────────────────────────────────── - if (tileId !== undefined && inputTile !== undefined) { - // Work directly with the persisted internal tiles array so - // untouched tiles are never round-tripped through the external - // converter (which would strip orphaned container refs from - // unrelated tiles). - const internalTiles = - (existingDashboard.tiles as { id: string }[]) ?? []; - const existingIdx = internalTiles.findIndex(t => t.id === tileId); - if (existingIdx === -1) { - // Build a human-readable list of available tile IDs. - // Convert only for the error message (read-only, no write-back). - const externalDashboard = - convertToExternalDashboard(existingDashboard); - return { - isError: true, - content: [ - { - type: 'text' as const, - text: `Tile not found: ${tileId}. Available tile IDs: ${externalDashboard.tiles.map(t => `${t.id} (${t.name})`).join(', ')}`, - }, - ], - }; - } - - // Read the existing tile's layout and container refs from the - // persisted internal format so fallback values are accurate - // regardless of the external converter's self-healing logic. - const existingInternalTile = internalTiles[existingIdx] as { - id: string; - x?: number; - y?: number; - w?: number; - h?: number; - containerId?: string; - tabId?: string; - config?: { name?: string }; + // Convert only the patched tile to internal format. + if (!isConfigTile(mergedTile)) { + return { + isError: true, + content: [ + { + type: 'text' as const, + text: 'Tile must have a config block.', + }, + ], }; - - // Merge: the incoming tile definition replaces config/name, - // but layout and container refs fall back to the existing tile - // when omitted so the LLM doesn't have to re-specify them. - // Coerce legacy empty-string containerId/tabId to undefined so - // they don't trip the container-ref validator (mirrors the - // self-heal in convertTileToExternalChart). - const incoming = inputTile as Partial; - const existingContainerId = - existingInternalTile.containerId || undefined; - const existingTabId = existingInternalTile.tabId || undefined; - const mergedTile: ExternalDashboardTileWithId = { - id: tileId, - name: incoming.name ?? existingInternalTile.config?.name ?? '', - x: incoming.x ?? existingInternalTile.x ?? 0, - y: incoming.y ?? existingInternalTile.y ?? 0, - w: incoming.w ?? existingInternalTile.w ?? 12, - h: incoming.h ?? existingInternalTile.h ?? 4, - containerId: - 'containerId' in incoming - ? incoming.containerId - : existingContainerId, - tabId: 'tabId' in incoming ? incoming.tabId : existingTabId, - // The config comes from the incoming tile (validated by Zod). - config: incoming.config, - } as ExternalDashboardTileWithId; - - // Error on raw SQL tiles that have no source defined but which use - // macros which require a source to be set - const sqlFilterSourceError = getRawSqlMissingSourceError([ - mergedTile, - ]); - if (sqlFilterSourceError) { - return mcpError(sqlFilterSourceError); - } - - // Validate the patched tile using the shared validation helper. - const validationError = await validateDashboardTiles({ - teamId, - tiles: [mergedTile], - existingTiles: existingDashboard.tiles ?? [], - containers: existingDashboard.containers ?? [], - }); - if (validationError) { - return { - isError: true, - content: [{ type: 'text' as const, text: validationError }], - }; - } - - // Convert only the patched tile to internal format. - if (!isConfigTile(mergedTile)) { - return { - isError: true, - content: [ - { - type: 'text' as const, - text: 'Tile must have a config block.', - }, - ], - }; - } - const internalTile = convertToInternalTileConfig(mergedTile); - - // Use the positional $ operator matched by 'tiles.id' in the - // query filter. This targets the tile by its id field rather - // than a captured numeric index, so a concurrent save_dashboard - // that replaces the whole tiles array can't cause us to - // overwrite an unrelated tile at a stale index. - queryFilter['tiles.id'] = tileId; - setPayload['tiles.$'] = internalTile; - patchedTile = mergedTile; } + const internalTile = convertToInternalTileConfig(mergedTile); - const updatedDashboard = await Dashboard.findOneAndUpdate( - queryFilter, - { $set: setPayload }, - { new: true }, - ); + // Use the positional $ operator matched by 'tiles.id' in the + // query filter. This targets the tile by its id field rather + // than a captured numeric index, so a concurrent save_dashboard + // that replaces the whole tiles array can't cause us to + // overwrite an unrelated tile at a stale index. + queryFilter['tiles.id'] = tileId; + setPayload['tiles.$'] = internalTile; + patchedTile = mergedTile; + } - if (!updatedDashboard) { - // When a tile patch is in flight, a null result means the tile - // was removed or the dashboard was deleted between our read - // and this write. - if (tileId !== undefined) { - return { - isError: true, - content: [ - { - type: 'text' as const, - text: - `Tile ${tileId} was not found at write time (it may have been removed by a concurrent update). ` + - 'The entire update was rejected — name/tags changes (if any) were not applied. Resubmit.', - }, - ], - }; - } + const updatedDashboard = await Dashboard.findOneAndUpdate( + queryFilter, + { $set: setPayload }, + { new: true }, + ); + + if (!updatedDashboard) { + // When a tile patch is in flight, a null result means the tile + // was removed or the dashboard was deleted between our read + // and this write. + if (tileId !== undefined) { return { isError: true, - content: [{ type: 'text' as const, text: 'Dashboard not found' }], + content: [ + { + type: 'text' as const, + text: + `Tile ${tileId} was not found at write time (it may have been removed by a concurrent update). ` + + 'The entire update was rejected — name/tags changes (if any) were not applied. Resubmit.', + }, + ], }; } + return { + isError: true, + content: [{ type: 'text' as const, text: 'Dashboard not found' }], + }; + } - // Reconcile alerts: if the tile's displayType changed to one - // that doesn't support alerts (e.g. raw SQL line/pie), clean up - // stale alert documents. Scope to just the patched tile — pass - // it as both the "new" tiles and "existing" ids so the helper - // checks whether the updated config still supports alerts. - if (tileId !== undefined) { - const existingTileIds = new Set([tileId]); - const patchedTileInDb = updatedDashboard.tiles.filter( - t => t.id === tileId, - ); - await cleanupDashboardAlerts({ - dashboardId, - teamId, - internalTiles: patchedTileInDb, - existingTileIds, - }); - } + // Reconcile alerts: if the tile's displayType changed to one + // that doesn't support alerts (e.g. raw SQL line/pie), clean up + // stale alert documents. Scope to just the patched tile — pass + // it as both the "new" tiles and "existing" ids so the helper + // checks whether the updated config still supports alerts. + if (tileId !== undefined) { + const existingTileIds = new Set([tileId]); + const patchedTileInDb = updatedDashboard.tiles.filter( + t => t.id === tileId, + ); + await cleanupDashboardAlerts({ + dashboardId, + teamId, + internalTiles: patchedTileInDb, + existingTileIds, + }); + } - // Return a lightweight response: the patched tile (if any) plus - // updated dashboard metadata, without the full tile array. - const output: Record = { - id: updatedDashboard._id.toString(), - name: updatedDashboard.name, - tags: updatedDashboard.tags, - ...(frontendUrl - ? { url: `${frontendUrl}/dashboards/${updatedDashboard._id}` } - : {}), - }; - if (patchedTile) { - output.patchedTile = patchedTile; - output.hint = - 'Use clickstack_query_tile to test the patched tile query.'; - const macroWarnings = getRawSqlTileMacroWarnings([patchedTile]); - if (macroWarnings.length > 0) { - output.warnings = macroWarnings; - } + // Return a lightweight response: the patched tile (if any) plus + // updated dashboard metadata, without the full tile array. + const output: Record = { + id: updatedDashboard._id.toString(), + name: updatedDashboard.name, + tags: updatedDashboard.tags, + ...(frontendUrl + ? { url: `${frontendUrl}/dashboards/${updatedDashboard._id}` } + : {}), + }; + if (patchedTile) { + output.patchedTile = patchedTile; + output.hint = + 'Use clickstack_query_tile to test the patched tile query.'; + const macroWarnings = getRawSqlTileMacroWarnings([patchedTile]); + if (macroWarnings.length > 0) { + output.warnings = macroWarnings; } + } - return { - content: [ - { - type: 'text' as const, - text: JSON.stringify(output, null, 2), - }, - ], - }; - }, - ), + return { + content: [ + { + type: 'text' as const, + text: JSON.stringify(output, null, 2), + }, + ], + }; + }, ); } diff --git a/packages/api/src/mcp/tools/dashboards/queryTile.ts b/packages/api/src/mcp/tools/dashboards/queryTile.ts index db3b2ae50e..b345997bf9 100644 --- a/packages/api/src/mcp/tools/dashboards/queryTile.ts +++ b/packages/api/src/mcp/tools/dashboards/queryTile.ts @@ -1,22 +1,20 @@ -import type { McpServer } from '@modelcontextprotocol/sdk/server/mcp.js'; import { z } from 'zod'; import { parseTimeRange, runConfigTile } from '@/mcp/tools/query/helpers'; -import type { McpContext } from '@/mcp/tools/types'; -import { withToolTracing } from '@/mcp/utils/tracing'; +import type { ToolRegistrar } from '@/mcp/tools/types'; import Dashboard from '@/models/dashboard'; import { convertToExternalDashboard } from '@/routers/external-api/v2/utils/dashboards'; import { objectIdSchema } from '@/utils/zod'; import { getRawSqlTileMacroWarnings } from './validation'; -export function registerQueryTile( - server: McpServer, - context: McpContext, -): void { +export function registerQueryTile({ + context, + registerTool, +}: ToolRegistrar): void { const { teamId } = context; - server.registerTool( + registerTool( 'clickstack_query_tile', { title: 'Query a Dashboard Tile', @@ -46,71 +44,67 @@ export function registerQueryTile( .describe('End of the query window as ISO 8601. Default: now.'), }), }, - withToolTracing( - 'clickstack_query_tile', - context, - async ({ dashboardId, tileId, startTime, endTime }) => { - const timeRange = parseTimeRange(startTime, endTime); - if ('error' in timeRange) { - return { - isError: true, - content: [{ type: 'text' as const, text: timeRange.error }], - }; - } - const { startDate, endDate } = timeRange; + async ({ dashboardId, tileId, startTime, endTime }) => { + const timeRange = parseTimeRange(startTime, endTime); + if ('error' in timeRange) { + return { + isError: true, + content: [{ type: 'text' as const, text: timeRange.error }], + }; + } + const { startDate, endDate } = timeRange; - const dashboard = await Dashboard.findOne({ - _id: dashboardId, - team: teamId, - }); - if (!dashboard) { - return { - isError: true, - content: [{ type: 'text' as const, text: 'Dashboard not found' }], - }; - } + const dashboard = await Dashboard.findOne({ + _id: dashboardId, + team: teamId, + }); + if (!dashboard) { + return { + isError: true, + content: [{ type: 'text' as const, text: 'Dashboard not found' }], + }; + } - const externalDashboard = convertToExternalDashboard(dashboard); - const tile = externalDashboard.tiles.find(t => t.id === tileId); - if (!tile) { - return { - isError: true, - content: [ - { - type: 'text' as const, - text: `Tile not found: ${tileId}. Available tile IDs: ${externalDashboard.tiles.map(t => t.id).join(', ')}`, - }, - ], - }; - } + const externalDashboard = convertToExternalDashboard(dashboard); + const tile = externalDashboard.tiles.find(t => t.id === tileId); + if (!tile) { + return { + isError: true, + content: [ + { + type: 'text' as const, + text: `Tile not found: ${tileId}. Available tile IDs: ${externalDashboard.tiles.map(t => t.id).join(', ')}`, + }, + ], + }; + } - const result = await runConfigTile( - teamId.toString(), - tile, - startDate, - endDate, - ); + const result = await runConfigTile( + teamId.toString(), + tile, + startDate, + endDate, + ); - // Surface non-blocking missing macro warnings alongside - // the successful result so the agent can spot a tile that runs but - // ignores dashboard controls. - const macroWarnings = getRawSqlTileMacroWarnings([tile]); - if ( - macroWarnings.length > 0 && - !('isError' in result && result.isError) && - result.content?.[0]?.type === 'text' - ) { - try { - const parsed = JSON.parse(result.content[0].text); - parsed.warnings = macroWarnings; - result.content[0].text = JSON.stringify(parsed, null, 2); - } catch { - // leave result unmodified - } + // Surface non-blocking missing macro warnings alongside + // the successful result so the agent can spot a tile that runs but + // ignores dashboard controls. + const macroWarnings = getRawSqlTileMacroWarnings([tile]); + if ( + macroWarnings.length > 0 && + !('isError' in result && result.isError) && + result.content?.[0]?.type === 'text' + ) { + try { + const parsed = JSON.parse(result.content[0].text); + parsed.warnings = macroWarnings; + result.content[0].text = JSON.stringify(parsed, null, 2); + } catch { + // leave result unmodified } + } - return result; - }, - ), + return result; + }, ); } diff --git a/packages/api/src/mcp/tools/dashboards/saveDashboard.ts b/packages/api/src/mcp/tools/dashboards/saveDashboard.ts index 53e4f46eb3..7afcc1c463 100644 --- a/packages/api/src/mcp/tools/dashboards/saveDashboard.ts +++ b/packages/api/src/mcp/tools/dashboards/saveDashboard.ts @@ -1,13 +1,11 @@ import type { DashboardContainer } from '@hyperdx/common-utils/dist/types'; -import type { McpServer } from '@modelcontextprotocol/sdk/server/mcp.js'; import { uniq } from 'lodash'; import mongoose from 'mongoose'; import { z } from 'zod'; import * as config from '@/config'; -import type { McpContext } from '@/mcp/tools/types'; +import type { ToolRegistrar } from '@/mcp/tools/types'; import { mcpError } from '@/mcp/utils/errors'; -import { withToolTracing } from '@/mcp/utils/tracing'; import Dashboard, { IDashboard } from '@/models/dashboard'; import { cleanupDashboardAlerts, @@ -32,14 +30,14 @@ import { getRawSqlTileMacroWarnings, } from './validation'; -export function registerSaveDashboard( - server: McpServer, - context: McpContext, -): void { +export function registerSaveDashboard({ + context, + registerTool, +}: ToolRegistrar): void { const { teamId } = context; const frontendUrl = config.FRONTEND_URL; - server.registerTool( + registerTool( 'clickstack_save_dashboard', { title: 'Create or Update Dashboard', @@ -63,40 +61,36 @@ export function registerSaveDashboard( filters: mcpFiltersParam.optional(), }), }, - withToolTracing( - 'clickstack_save_dashboard', - context, - async ({ - id: dashboardId, - name, - tiles: inputTiles, - tags, - containers, - filters: inputFilters, - }) => { - if (!dashboardId) { - return createDashboard({ - teamId, - frontendUrl, - name, - inputTiles, - tags, - containers, - inputFilters, - }); - } - return updateDashboard({ + async ({ + id: dashboardId, + name, + tiles: inputTiles, + tags, + containers, + filters: inputFilters, + }) => { + if (!dashboardId) { + return createDashboard({ teamId, frontendUrl, - dashboardId, name, inputTiles, tags, containers, inputFilters, }); - }, - ), + } + return updateDashboard({ + teamId, + frontendUrl, + dashboardId, + name, + inputTiles, + tags, + containers, + inputFilters, + }); + }, ); } diff --git a/packages/api/src/mcp/tools/dashboards/searchDashboards.ts b/packages/api/src/mcp/tools/dashboards/searchDashboards.ts index b00dd51f34..9fb29f5e38 100644 --- a/packages/api/src/mcp/tools/dashboards/searchDashboards.ts +++ b/packages/api/src/mcp/tools/dashboards/searchDashboards.ts @@ -1,9 +1,7 @@ -import type { McpServer } from '@modelcontextprotocol/sdk/server/mcp.js'; import { escapeRegExp } from 'lodash'; import * as config from '@/config'; -import type { McpContext } from '@/mcp/tools/types'; -import { withToolTracing } from '@/mcp/utils/tracing'; +import type { ToolRegistrar } from '@/mcp/tools/types'; import Dashboard from '@/models/dashboard'; import logger from '@/utils/logger'; @@ -11,14 +9,14 @@ import { mcpSearchDashboardsSchema } from './schemas'; const SEARCH_RESULTS_LIMIT = 100; -export function registerSearchDashboards( - server: McpServer, - context: McpContext, -): void { +export function registerSearchDashboards({ + context, + registerTool, +}: ToolRegistrar): void { const { teamId } = context; const frontendUrl = config.FRONTEND_URL; - server.registerTool( + registerTool( 'clickstack_search_dashboards', { title: 'Search Dashboards', @@ -28,88 +26,82 @@ export function registerSearchDashboards( 'lists all dashboards). At least one of query or tags must be provided.', inputSchema: mcpSearchDashboardsSchema, }, - withToolTracing( - 'clickstack_search_dashboards', - context, - async ({ query, tags }) => { - const hasQuery = typeof query === 'string' && query.length > 0; - const hasTags = Array.isArray(tags) && tags.length > 0; + async ({ query, tags }) => { + const hasQuery = typeof query === 'string' && query.length > 0; + const hasTags = Array.isArray(tags) && tags.length > 0; - if (!hasQuery && !hasTags) { - return { - isError: true, - content: [ - { - type: 'text' as const, - text: 'Provide at least one of: query (non-empty string) or tags (non-empty array).', - }, - ], - }; - } + if (!hasQuery && !hasTags) { + return { + isError: true, + content: [ + { + type: 'text' as const, + text: 'Provide at least one of: query (non-empty string) or tags (non-empty array).', + }, + ], + }; + } - const filter: Record = { team: teamId }; + const filter: Record = { team: teamId }; - if (hasQuery) { - filter.name = { $regex: escapeRegExp(query), $options: 'i' }; - } - if (hasTags) { - filter.tags = { $all: tags }; - } + if (hasQuery) { + filter.name = { $regex: escapeRegExp(query), $options: 'i' }; + } + if (hasTags) { + filter.tags = { $all: tags }; + } - try { - const dashboards = await Dashboard.find(filter) - .select({ name: 1, tags: 1 }) - .limit(SEARCH_RESULTS_LIMIT + 1) - .lean(); + try { + const dashboards = await Dashboard.find(filter) + .select({ name: 1, tags: 1 }) + .limit(SEARCH_RESULTS_LIMIT + 1) + .lean(); - const truncated = dashboards.length > SEARCH_RESULTS_LIMIT; - const results = truncated - ? dashboards.slice(0, SEARCH_RESULTS_LIMIT) - : dashboards; + const truncated = dashboards.length > SEARCH_RESULTS_LIMIT; + const results = truncated + ? dashboards.slice(0, SEARCH_RESULTS_LIMIT) + : dashboards; - const output = results.map(d => ({ - id: d._id.toString(), - name: d.name, - tags: d.tags, - ...(frontendUrl - ? { url: `${frontendUrl}/dashboards/${d._id}` } - : {}), - })); + const output = results.map(d => ({ + id: d._id.toString(), + name: d.name, + tags: d.tags, + ...(frontendUrl ? { url: `${frontendUrl}/dashboards/${d._id}` } : {}), + })); - return { - content: [ - { - type: 'text' as const, - text: JSON.stringify( - truncated - ? { - results: output, - truncated: true, - hint: `Returned first ${SEARCH_RESULTS_LIMIT} results. Narrow your query or tags to see more.`, - } - : output, - null, - 2, - ), - }, - ], - }; - } catch (err) { - logger.error( - { err, teamId, query, tags }, - 'clickstack_search_dashboards: query failed', - ); - return { - isError: true, - content: [ - { - type: 'text' as const, - text: `Search failed: ${err instanceof Error ? err.message : String(err)}`, - }, - ], - }; - } - }, - ), + return { + content: [ + { + type: 'text' as const, + text: JSON.stringify( + truncated + ? { + results: output, + truncated: true, + hint: `Returned first ${SEARCH_RESULTS_LIMIT} results. Narrow your query or tags to see more.`, + } + : output, + null, + 2, + ), + }, + ], + }; + } catch (err) { + logger.error( + { err, teamId, query, tags }, + 'clickstack_search_dashboards: query failed', + ); + return { + isError: true, + content: [ + { + type: 'text' as const, + text: `Search failed: ${err instanceof Error ? err.message : String(err)}`, + }, + ], + }; + } + }, ); } diff --git a/packages/api/src/mcp/tools/query/eventDeltas.ts b/packages/api/src/mcp/tools/query/eventDeltas.ts index 90175e1801..96e4a775a3 100644 --- a/packages/api/src/mcp/tools/query/eventDeltas.ts +++ b/packages/api/src/mcp/tools/query/eventDeltas.ts @@ -10,13 +10,11 @@ import { DisplayType, SourceKind, } from '@hyperdx/common-utils/dist/types'; -import type { McpServer } from '@modelcontextprotocol/sdk/server/mcp.js'; import { z } from 'zod'; import { getConnectionById } from '@/controllers/connection'; import { getSource } from '@/controllers/sources'; -import type { McpContext } from '@/mcp/tools/types'; -import { withToolTracing } from '@/mcp/utils/tracing'; +import type { ToolRegistrar } from '@/mcp/tools/types'; import { trimToolResponse } from '@/utils/trimToolResponse'; import { parseTimeRange } from './helpers'; @@ -136,10 +134,10 @@ function topDeltasForProperty( // ─── Tool definition ───────────────────────────────────────────────────────── -export function registerEventDeltas(server: McpServer, context: McpContext) { +export function registerEventDeltas({ context, registerTool }: ToolRegistrar) { const { teamId } = context; - server.registerTool( + registerTool( 'clickstack_event_deltas', { title: 'Compare Events: Target vs Baseline', @@ -196,354 +194,342 @@ export function registerEventDeltas(server: McpServer, context: McpContext) { 'to combine with the others.', inputSchema: deltasSchema, }, - withToolTracing( - 'clickstack_event_deltas', - context, - async (input: DeltasInput) => { - const targetRange = parseTimeRange( - input.target.startTime, - input.target.endTime, - ); - if ('error' in targetRange) { - return { - isError: true as const, - content: [ - { - type: 'text' as const, - text: `target: ${targetRange.error}`, - }, - ], - }; - } - const baselineRange = parseTimeRange( - input.baseline.startTime, - input.baseline.endTime, - ); - if ('error' in baselineRange) { - return { - isError: true as const, - content: [ - { - type: 'text' as const, - text: `baseline: ${baselineRange.error}`, - }, - ], - }; - } - - // Refuse calls where target and baseline can't possibly produce a - // meaningful comparison. The agent presumably wanted to write - // something different on one side and didn't — flag it so the - // round-trip is short instead of returning empty deltas. - const targetWhere = (input.target.where ?? '').trim(); - const baselineWhere = (input.baseline.where ?? '').trim(); - const targetLang = input.target.whereLanguage ?? 'lucene'; - const baselineLang = input.baseline.whereLanguage ?? 'lucene'; - const sameWhere = - targetWhere === baselineWhere && targetLang === baselineLang; - const sameWindow = - targetRange.startDate.getTime() === - baselineRange.startDate.getTime() && - targetRange.endDate.getTime() === baselineRange.endDate.getTime(); - if (sameWhere && sameWindow) { - return { - isError: true as const, - content: [ - { - type: 'text' as const, - text: - 'target and baseline are identical (same where + same time ' + - 'window) — that yields nothing to compare. Pick distinct ' + - 'groups: e.g. (a) before vs after a step change in time, ' + - '(b) failing rows vs succeeding rows, (c) slow spans ' + - '(Duration > X) vs fast spans (Duration <= X) in the same ' + - 'window, or (d) one cohort vs the rest. At least one of ' + - '`where` or the time window must differ between the two ' + - 'groups.', - }, - ], - }; - } + async (input: DeltasInput) => { + const targetRange = parseTimeRange( + input.target.startTime, + input.target.endTime, + ); + if ('error' in targetRange) { + return { + isError: true as const, + content: [ + { + type: 'text' as const, + text: `target: ${targetRange.error}`, + }, + ], + }; + } + const baselineRange = parseTimeRange( + input.baseline.startTime, + input.baseline.endTime, + ); + if ('error' in baselineRange) { + return { + isError: true as const, + content: [ + { + type: 'text' as const, + text: `baseline: ${baselineRange.error}`, + }, + ], + }; + } - const source = await getSource(teamId.toString(), input.sourceId); - if (!source) { - return { - isError: true as const, - content: [ - { - type: 'text' as const, - text: `Source not found: ${input.sourceId}. Call clickstack_list_sources to find available source IDs.`, - }, - ], - }; - } - if ( - source.kind !== SourceKind.Trace && - source.kind !== SourceKind.Log - ) { - return { - isError: true as const, - content: [ - { - type: 'text' as const, - text: `Source ${input.sourceId} is kind="${source.kind}". clickstack_event_deltas requires a trace or log source.`, - }, - ], - }; - } + // Refuse calls where target and baseline can't possibly produce a + // meaningful comparison. The agent presumably wanted to write + // something different on one side and didn't — flag it so the + // round-trip is short instead of returning empty deltas. + const targetWhere = (input.target.where ?? '').trim(); + const baselineWhere = (input.baseline.where ?? '').trim(); + const targetLang = input.target.whereLanguage ?? 'lucene'; + const baselineLang = input.baseline.whereLanguage ?? 'lucene'; + const sameWhere = + targetWhere === baselineWhere && targetLang === baselineLang; + const sameWindow = + targetRange.startDate.getTime() === baselineRange.startDate.getTime() && + targetRange.endDate.getTime() === baselineRange.endDate.getTime(); + if (sameWhere && sameWindow) { + return { + isError: true as const, + content: [ + { + type: 'text' as const, + text: + 'target and baseline are identical (same where + same time ' + + 'window) — that yields nothing to compare. Pick distinct ' + + 'groups: e.g. (a) before vs after a step change in time, ' + + '(b) failing rows vs succeeding rows, (c) slow spans ' + + '(Duration > X) vs fast spans (Duration <= X) in the same ' + + 'window, or (d) one cohort vs the rest. At least one of ' + + '`where` or the time window must differ between the two ' + + 'groups.', + }, + ], + }; + } - const connection = await getConnectionById( - teamId.toString(), - source.connection.toString(), - true, - ); - if (!connection) { - return { - isError: true as const, - content: [ - { - type: 'text' as const, - text: `Connection not found for source: ${input.sourceId}`, - }, - ], - }; - } + const source = await getSource(teamId.toString(), input.sourceId); + if (!source) { + return { + isError: true as const, + content: [ + { + type: 'text' as const, + text: `Source not found: ${input.sourceId}. Call clickstack_list_sources to find available source IDs.`, + }, + ], + }; + } + if (source.kind !== SourceKind.Trace && source.kind !== SourceKind.Log) { + return { + isError: true as const, + content: [ + { + type: 'text' as const, + text: `Source ${input.sourceId} is kind="${source.kind}". clickstack_event_deltas requires a trace or log source.`, + }, + ], + }; + } - const clickhouseClient = new ClickhouseClient({ - host: connection.host, - username: connection.username, - password: connection.password, - }); - const metadata = getMetadata(clickhouseClient); + const connection = await getConnectionById( + teamId.toString(), + source.connection.toString(), + true, + ); + if (!connection) { + return { + isError: true as const, + content: [ + { + type: 'text' as const, + text: `Connection not found for source: ${input.sourceId}`, + }, + ], + }; + } - // Build a stable per-source ORDER BY for sampling (matches what the - // app does — uses spanIdExpression on traces, falls back to rand()). - const stableSampleExpr = - source.kind === SourceKind.Trace - ? getStableSampleExpression(source.spanIdExpression) - : 'rand()'; + const clickhouseClient = new ClickhouseClient({ + host: connection.host, + username: connection.username, + password: connection.password, + }); + const metadata = getMetadata(clickhouseClient); - const sampleSize = input.sampleSize; + // Build a stable per-source ORDER BY for sampling (matches what the + // app does — uses spanIdExpression on traces, falls back to rand()). + const stableSampleExpr = + source.kind === SourceKind.Trace + ? getStableSampleExpression(source.spanIdExpression) + : 'rand()'; - const buildSampleConfig = ( - group: { where: string; whereLanguage: 'lucene' | 'sql' }, - startDate: Date, - endDate: Date, - ): BuilderChartConfigWithDateRange => ({ - displayType: DisplayType.Search, - // SELECT * so the algorithm sees all columns (it flattens nested - // Maps/Arrays internally). - select: '*', - from: source.from, - where: group.where, - whereLanguage: group.where ? group.whereLanguage : 'sql', - connection: source.connection.toString(), - timestampValueExpression: source.timestampValueExpression, - implicitColumnExpression: - 'implicitColumnExpression' in source - ? source.implicitColumnExpression - : undefined, - orderBy: stableSampleExpr, - limit: { limit: sampleSize }, - dateRange: [startDate, endDate], - }); + const sampleSize = input.sampleSize; - const targetConfig = buildSampleConfig( - input.target, - targetRange.startDate, - targetRange.endDate, - ); - const baselineConfig = buildSampleConfig( - input.baseline, - baselineRange.startDate, - baselineRange.endDate, - ); + const buildSampleConfig = ( + group: { where: string; whereLanguage: 'lucene' | 'sql' }, + startDate: Date, + endDate: Date, + ): BuilderChartConfigWithDateRange => ({ + displayType: DisplayType.Search, + // SELECT * so the algorithm sees all columns (it flattens nested + // Maps/Arrays internally). + select: '*', + from: source.from, + where: group.where, + whereLanguage: group.where ? group.whereLanguage : 'sql', + connection: source.connection.toString(), + timestampValueExpression: source.timestampValueExpression, + implicitColumnExpression: + 'implicitColumnExpression' in source + ? source.implicitColumnExpression + : undefined, + orderBy: stableSampleExpr, + limit: { limit: sampleSize }, + dateRange: [startDate, endDate], + }); - // Run renderChartConfig for both groups and getColumns in parallel — - // they're independent and each hits ClickHouse metadata. - let targetSql, baselineSql; - let columnMeta: { name: string; type: string }[] = []; - let columnMetaUnavailable = false; - try { - const [targetResult, baselineResult, colsResult] = await Promise.all([ - renderChartConfig(targetConfig, metadata, source.querySettings), - renderChartConfig(baselineConfig, metadata, source.querySettings), - metadata - .getColumns({ - databaseName: source.from.databaseName, - tableName: source.from.tableName, - connectionId: source.connection.toString(), - }) - .catch(() => null), - ]); - targetSql = targetResult; - baselineSql = baselineResult; - if (colsResult) { - columnMeta = colsResult.map(c => ({ - name: c.name, - type: c.type, - })); - } else { - // Without column meta the algorithm just won't denylist anything — - // not fatal, but signal it in the response so the agent knows - // high-cardinality / ID fields may appear in the ranking. - columnMetaUnavailable = true; - } - } catch (e) { - return { - isError: true as const, - content: [ - { - type: 'text' as const, - text: `Failed to build sample queries: ${e instanceof Error ? e.message : String(e)}`, - }, - ], - }; - } + const targetConfig = buildSampleConfig( + input.target, + targetRange.startDate, + targetRange.endDate, + ); + const baselineConfig = buildSampleConfig( + input.baseline, + baselineRange.startDate, + baselineRange.endDate, + ); - let targetRows: Record[]; - let baselineRows: Record[]; - const abortController = new AbortController(); - try { - const [targetRes, baselineRes] = await Promise.all([ - clickhouseClient.query({ - query: targetSql.sql, - query_params: targetSql.params, - format: 'JSON', + // Run renderChartConfig for both groups and getColumns in parallel — + // they're independent and each hits ClickHouse metadata. + let targetSql, baselineSql; + let columnMeta: { name: string; type: string }[] = []; + let columnMetaUnavailable = false; + try { + const [targetResult, baselineResult, colsResult] = await Promise.all([ + renderChartConfig(targetConfig, metadata, source.querySettings), + renderChartConfig(baselineConfig, metadata, source.querySettings), + metadata + .getColumns({ + databaseName: source.from.databaseName, + tableName: source.from.tableName, connectionId: source.connection.toString(), - clickhouse_settings: { max_execution_time: 30 }, - abort_signal: abortController.signal, - }), - clickhouseClient.query({ - query: baselineSql.sql, - query_params: baselineSql.params, - format: 'JSON', - connectionId: source.connection.toString(), - clickhouse_settings: { max_execution_time: 30 }, - abort_signal: abortController.signal, - }), - ]); - const targetJson = (await ( - targetRes as { json: () => Promise<{ data: any[] }> } - ).json()) ?? { data: [] }; - const baselineJson = (await ( - baselineRes as { json: () => Promise<{ data: any[] }> } - ).json()) ?? { data: [] }; - targetRows = targetJson.data ?? []; - baselineRows = baselineJson.data ?? []; - } catch (e) { - abortController.abort(); - return { - isError: true as const, - content: [ - { - type: 'text' as const, - text: `Failed to sample rows: ${e instanceof Error ? e.message : String(e)}`, - }, - ], - }; + }) + .catch(() => null), + ]); + targetSql = targetResult; + baselineSql = baselineResult; + if (colsResult) { + columnMeta = colsResult.map(c => ({ + name: c.name, + type: c.type, + })); + } else { + // Without column meta the algorithm just won't denylist anything — + // not fatal, but signal it in the response so the agent knows + // high-cardinality / ID fields may appear in the ranking. + columnMetaUnavailable = true; } - - const ranked = rankProperties({ - targetRows, - baselineRows, - columnMeta, - }); - - const visible = ranked.ranked.filter(p => !p.hidden); - const hidden = ranked.ranked.filter(p => p.hidden); - - const renderEntry = ( - rank: number, - p: (typeof ranked.ranked)[number], - ) => { - const tValues = - ranked.targetStats.valueOccurences.get(p.key) ?? - new Map(); - const bValues = - ranked.baselineStats.valueOccurences.get(p.key) ?? - new Map(); - const tTotal = ranked.targetStats.propertyOccurences.get(p.key) ?? 0; - const bTotal = - ranked.baselineStats.propertyOccurences.get(p.key) ?? 0; - // Compact entry: topDeltas already contains every value's - // target%/baseline%/diff%, so the previously-emitted nested - // `target.topValues` / `baseline.topValues` were redundant - // (~60% of the response payload). Keep the flat sample counts - // so the agent can size-check the comparison. - return { - rank, - key: p.key, - score: p.score, - semanticBoost: p.semanticBoost > 0, - targetCount: tTotal, - baselineCount: bTotal, - ...(p.hidden ? { hiddenReason: p.hiddenReason } : {}), - topDeltas: topDeltasForProperty( - tValues, - bValues, - tTotal, - bTotal, - input.topValuesPerProperty, - ), - }; - }; - - const properties = visible - .slice(0, input.topN) - .map((p, i) => renderEntry(i + 1, p)); - const hiddenEntries = input.includeHidden - ? hidden - .slice(0, input.topN) - .map((p, i) => renderEntry(properties.length + i + 1, p)) - : undefined; - - const output = { - summary: { - sampleSize, - targetSampleCount: targetRows.length, - baselineSampleCount: baselineRows.length, - propertiesScored: ranked.ranked.length, - propertiesVisible: visible.length, - propertiesHidden: hidden.length, - ...(columnMetaUnavailable ? { columnMetaUnavailable: true } : {}), - ...(targetRows.length === 0 || baselineRows.length === 0 - ? { - warning: - 'One or both groups returned 0 rows — verify the time ' + - 'window and where filter via clickstack_search.', - } - : {}), - }, - properties, - ...(hiddenEntries ? { hidden: hiddenEntries } : {}), - usage: - 'score is the maximum %-share difference of any value between target ' + - 'and baseline (after normalizing each group to 100%), plus a 0.1 ' + - 'tiebreaker boost for well-known OTel attributes. ' + - 'Properties with `semanticBoost: true` are well-known OTel attrs. ' + - 'Hidden properties (ID/timestamp arrays + high-cardinality fields) ' + - 'are dropped by default — set includeHidden:true to inspect them.', + } catch (e) { + return { + isError: true as const, + content: [ + { + type: 'text' as const, + text: `Failed to build sample queries: ${e instanceof Error ? e.message : String(e)}`, + }, + ], }; + } - 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; - + let targetRows: Record[]; + let baselineRows: Record[]; + const abortController = new AbortController(); + try { + const [targetRes, baselineRes] = await Promise.all([ + clickhouseClient.query({ + query: targetSql.sql, + query_params: targetSql.params, + format: 'JSON', + connectionId: source.connection.toString(), + clickhouse_settings: { max_execution_time: 30 }, + abort_signal: abortController.signal, + }), + clickhouseClient.query({ + query: baselineSql.sql, + query_params: baselineSql.params, + format: 'JSON', + connectionId: source.connection.toString(), + clickhouse_settings: { max_execution_time: 30 }, + abort_signal: abortController.signal, + }), + ]); + const targetJson = (await ( + targetRes as { json: () => Promise<{ data: any[] }> } + ).json()) ?? { data: [] }; + const baselineJson = (await ( + baselineRes as { json: () => Promise<{ data: any[] }> } + ).json()) ?? { data: [] }; + targetRows = targetJson.data ?? []; + baselineRows = baselineJson.data ?? []; + } catch (e) { + abortController.abort(); return { + isError: true as const, content: [ { type: 'text' as const, - text: JSON.stringify(finalOutput, null, 2), + text: `Failed to sample rows: ${e instanceof Error ? e.message : String(e)}`, }, ], }; - }, - ), + } + + const ranked = rankProperties({ + targetRows, + baselineRows, + columnMeta, + }); + + const visible = ranked.ranked.filter(p => !p.hidden); + const hidden = ranked.ranked.filter(p => p.hidden); + + const renderEntry = (rank: number, p: (typeof ranked.ranked)[number]) => { + const tValues = + ranked.targetStats.valueOccurences.get(p.key) ?? + new Map(); + const bValues = + ranked.baselineStats.valueOccurences.get(p.key) ?? + new Map(); + const tTotal = ranked.targetStats.propertyOccurences.get(p.key) ?? 0; + const bTotal = ranked.baselineStats.propertyOccurences.get(p.key) ?? 0; + // Compact entry: topDeltas already contains every value's + // target%/baseline%/diff%, so the previously-emitted nested + // `target.topValues` / `baseline.topValues` were redundant + // (~60% of the response payload). Keep the flat sample counts + // so the agent can size-check the comparison. + return { + rank, + key: p.key, + score: p.score, + semanticBoost: p.semanticBoost > 0, + targetCount: tTotal, + baselineCount: bTotal, + ...(p.hidden ? { hiddenReason: p.hiddenReason } : {}), + topDeltas: topDeltasForProperty( + tValues, + bValues, + tTotal, + bTotal, + input.topValuesPerProperty, + ), + }; + }; + + const properties = visible + .slice(0, input.topN) + .map((p, i) => renderEntry(i + 1, p)); + const hiddenEntries = input.includeHidden + ? hidden + .slice(0, input.topN) + .map((p, i) => renderEntry(properties.length + i + 1, p)) + : undefined; + + const output = { + summary: { + sampleSize, + targetSampleCount: targetRows.length, + baselineSampleCount: baselineRows.length, + propertiesScored: ranked.ranked.length, + propertiesVisible: visible.length, + propertiesHidden: hidden.length, + ...(columnMetaUnavailable ? { columnMetaUnavailable: true } : {}), + ...(targetRows.length === 0 || baselineRows.length === 0 + ? { + warning: + 'One or both groups returned 0 rows — verify the time ' + + 'window and where filter via clickstack_search.', + } + : {}), + }, + properties, + ...(hiddenEntries ? { hidden: hiddenEntries } : {}), + usage: + 'score is the maximum %-share difference of any value between target ' + + 'and baseline (after normalizing each group to 100%), plus a 0.1 ' + + 'tiebreaker boost for well-known OTel attributes. ' + + 'Properties with `semanticBoost: true` are well-known OTel attrs. ' + + 'Hidden properties (ID/timestamp arrays + high-cardinality fields) ' + + 'are dropped by default — set includeHidden:true to inspect them.', + }; + + 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(finalOutput, null, 2), + }, + ], + }; + }, ); } diff --git a/packages/api/src/mcp/tools/query/eventPatterns.ts b/packages/api/src/mcp/tools/query/eventPatterns.ts index 35e4fd9c21..96fd9dc472 100644 --- a/packages/api/src/mcp/tools/query/eventPatterns.ts +++ b/packages/api/src/mcp/tools/query/eventPatterns.ts @@ -1,8 +1,6 @@ -import type { McpServer } from '@modelcontextprotocol/sdk/server/mcp.js'; import { z } from 'zod'; -import type { McpContext } from '@/mcp/tools/types'; -import { withToolTracing } from '@/mcp/utils/tracing'; +import type { ToolRegistrar } from '@/mcp/tools/types'; import { parseTimeRange } from './helpers'; import { runEventPatterns } from './runEventPatterns'; @@ -63,10 +61,13 @@ const eventPatternsSchema = z.object({ // ─── Tool registration ─────────────────────────────────────────────────────── -export function registerEventPatterns(server: McpServer, context: McpContext) { +export function registerEventPatterns({ + context, + registerTool, +}: ToolRegistrar) { const { teamId } = context; - server.registerTool( + registerTool( 'clickstack_event_patterns', { title: 'Event Pattern Mining', @@ -89,7 +90,7 @@ export function registerEventPatterns(server: McpServer, context: McpContext) { ' - clickstack_table: aggregated metrics / counts / top-N', inputSchema: eventPatternsSchema, }, - withToolTracing('clickstack_event_patterns', context, async input => { + async input => { const timeRange = parseTimeRange(input.startTime, input.endTime); if ('error' in timeRange) { return { @@ -113,6 +114,6 @@ export function registerEventPatterns(server: McpServer, context: McpContext) { trendBuckets: input.trendBuckets, }, ); - }), + }, ); } diff --git a/packages/api/src/mcp/tools/query/index.ts b/packages/api/src/mcp/tools/query/index.ts index 5aaa6cfdad..a1ba01d5dc 100644 --- a/packages/api/src/mcp/tools/query/index.ts +++ b/packages/api/src/mcp/tools/query/index.ts @@ -1,6 +1,4 @@ -import type { McpServer } from '@modelcontextprotocol/sdk/server/mcp.js'; - -import type { McpContext, ToolDefinition } from '@/mcp/tools/types'; +import type { ToolDefinition, ToolRegistrar } from '@/mcp/tools/types'; import { registerEventDeltas } from './eventDeltas'; import { registerEventPatterns } from './eventPatterns'; @@ -9,13 +7,13 @@ import { registerSql } from './sql'; import { registerTable } from './table'; import { registerTimeseries } from './timeseries'; -const queryTools: ToolDefinition = (server: McpServer, context: McpContext) => { - registerTimeseries(server, context); - registerTable(server, context); - registerSearch(server, context); - registerEventPatterns(server, context); - registerEventDeltas(server, context); - registerSql(server, context); +const queryTools: ToolDefinition = (registrar: ToolRegistrar) => { + registerTimeseries(registrar); + registerTable(registrar); + registerSearch(registrar); + registerEventPatterns(registrar); + registerEventDeltas(registrar); + registerSql(registrar); }; export default queryTools; diff --git a/packages/api/src/mcp/tools/query/search.ts b/packages/api/src/mcp/tools/query/search.ts index cb01b36568..f4ca771754 100644 --- a/packages/api/src/mcp/tools/query/search.ts +++ b/packages/api/src/mcp/tools/query/search.ts @@ -1,8 +1,6 @@ -import type { McpServer } from '@modelcontextprotocol/sdk/server/mcp.js'; import { z } from 'zod'; -import type { McpContext } from '@/mcp/tools/types'; -import { withToolTracing } from '@/mcp/utils/tracing'; +import type { ToolRegistrar } from '@/mcp/tools/types'; import logger from '@/utils/logger'; import { trimToolResponse } from '@/utils/trimToolResponse'; @@ -56,10 +54,10 @@ const searchSchema = z.object({ // ─── Tool registration ─────────────────────────────────────────────────────── -export function registerSearch(server: McpServer, context: McpContext) { +export function registerSearch({ context, registerTool }: ToolRegistrar) { const { teamId } = context; - server.registerTool( + registerTool( 'clickstack_search', { title: 'Search Events', @@ -76,7 +74,7 @@ export function registerSearch(server: McpServer, context: McpContext) { "Map attributes use bracket syntax: SpanAttributes['http.method'].", inputSchema: searchSchema, }, - withToolTracing('clickstack_search', context, async input => { + async input => { const timeRange = parseTimeRange(input.startTime, input.endTime); if ('error' in timeRange) { return { @@ -220,6 +218,6 @@ export function registerSearch(server: McpServer, context: McpContext) { }, ], }; - }), + }, ); } diff --git a/packages/api/src/mcp/tools/query/sql.ts b/packages/api/src/mcp/tools/query/sql.ts index 134aeb4783..ada9f9187d 100644 --- a/packages/api/src/mcp/tools/query/sql.ts +++ b/packages/api/src/mcp/tools/query/sql.ts @@ -1,8 +1,6 @@ -import type { McpServer } from '@modelcontextprotocol/sdk/server/mcp.js'; import { z } from 'zod'; -import type { McpContext } from '@/mcp/tools/types'; -import { withToolTracing } from '@/mcp/utils/tracing'; +import type { ToolRegistrar } from '@/mcp/tools/types'; import { buildTile, parseTimeRange, runConfigTile } from './helpers'; import { endTimeSchema, startTimeSchema } from './schemas'; @@ -53,10 +51,10 @@ const sqlSchema = z.object({ // ─── Tool registration ─────────────────────────────────────────────────────── -export function registerSql(server: McpServer, context: McpContext) { +export function registerSql({ context, registerTool }: ToolRegistrar) { const { teamId } = context; - server.registerTool( + registerTool( 'clickstack_sql', { title: 'Raw SQL Query', @@ -73,7 +71,7 @@ export function registerSql(server: McpServer, context: McpContext) { 'For browsing rows use clickstack_search.', inputSchema: sqlSchema, }, - withToolTracing('clickstack_sql', context, async input => { + async input => { const timeRange = parseTimeRange(input.startTime, input.endTime); if ('error' in timeRange) { return { @@ -91,6 +89,6 @@ export function registerSql(server: McpServer, context: McpContext) { }); return runConfigTile(teamId.toString(), tile, startDate, endDate); - }), + }, ); } diff --git a/packages/api/src/mcp/tools/query/table.ts b/packages/api/src/mcp/tools/query/table.ts index 34fcdf6512..cab977e979 100644 --- a/packages/api/src/mcp/tools/query/table.ts +++ b/packages/api/src/mcp/tools/query/table.ts @@ -1,8 +1,6 @@ -import type { McpServer } from '@modelcontextprotocol/sdk/server/mcp.js'; import { z } from 'zod'; -import type { McpContext } from '@/mcp/tools/types'; -import { withToolTracing } from '@/mcp/utils/tracing'; +import type { ToolRegistrar } from '@/mcp/tools/types'; import { annotateIncreaseTopNHint, @@ -197,10 +195,10 @@ export function resolveOrderBy( // ─── Tool registration ─────────────────────────────────────────────────────── -export function registerTable(server: McpServer, context: McpContext) { +export function registerTable({ context, registerTool }: ToolRegistrar) { const { teamId } = context; - server.registerTool( + registerTool( 'clickstack_table', { title: 'Aggregation Table', @@ -231,7 +229,7 @@ export function registerTable(server: McpServer, context: McpContext) { 'summary and exponential histogram kinds are not supported by the query renderer yet.', inputSchema: tableSchema, }, - withToolTracing('clickstack_table', context, async input => { + async input => { const timeRange = parseTimeRange(input.startTime, input.endTime); if ('error' in timeRange) { return { @@ -307,6 +305,6 @@ export function registerTable(server: McpServer, context: McpContext) { annotateIncreaseTopNHint(result, select, input.groupBy); return result; - }), + }, ); } diff --git a/packages/api/src/mcp/tools/query/timeseries.ts b/packages/api/src/mcp/tools/query/timeseries.ts index 52a36be8c3..ebfe23c7b1 100644 --- a/packages/api/src/mcp/tools/query/timeseries.ts +++ b/packages/api/src/mcp/tools/query/timeseries.ts @@ -1,9 +1,7 @@ import { FIXED_TIME_BUCKET_EXPR_ALIAS } from '@hyperdx/common-utils/dist/core/renderChartConfig'; -import type { McpServer } from '@modelcontextprotocol/sdk/server/mcp.js'; import { z } from 'zod'; -import type { McpContext } from '@/mcp/tools/types'; -import { withToolTracing } from '@/mcp/utils/tracing'; +import type { ToolRegistrar } from '@/mcp/tools/types'; import { annotateIncreaseTopNHint, @@ -72,10 +70,10 @@ const timeseriesSchema = z.object({ // ─── Tool registration ─────────────────────────────────────────────────────── -export function registerTimeseries(server: McpServer, context: McpContext) { +export function registerTimeseries({ context, registerTool }: ToolRegistrar) { const { teamId } = context; - server.registerTool( + registerTool( 'clickstack_timeseries', { title: 'Time-Series Chart', @@ -102,7 +100,7 @@ export function registerTimeseries(server: McpServer, context: McpContext) { 'summary and exponential histogram kinds are not supported by the query renderer yet.', inputSchema: timeseriesSchema, }, - withToolTracing('clickstack_timeseries', context, async input => { + async input => { const timeRange = parseTimeRange(input.startTime, input.endTime); if ('error' in timeRange) { return { @@ -222,6 +220,6 @@ export function registerTimeseries(server: McpServer, context: McpContext) { } return result; - }), + }, ); } diff --git a/packages/api/src/mcp/tools/savedSearches/getSavedSearch.ts b/packages/api/src/mcp/tools/savedSearches/getSavedSearch.ts index 940ee000f8..dd2b044541 100644 --- a/packages/api/src/mcp/tools/savedSearches/getSavedSearch.ts +++ b/packages/api/src/mcp/tools/savedSearches/getSavedSearch.ts @@ -1,21 +1,19 @@ -import type { McpServer } from '@modelcontextprotocol/sdk/server/mcp.js'; import { z } from 'zod'; import * as config from '@/config'; import { getSavedSearch } from '@/controllers/savedSearch'; -import type { McpContext } from '@/mcp/tools/types'; +import type { ToolRegistrar } from '@/mcp/tools/types'; import { validateObjectId } from '@/mcp/utils/errors'; -import { withToolTracing } from '@/mcp/utils/tracing'; import { SavedSearch } from '@/models/savedSearch'; -export function registerGetSavedSearch( - server: McpServer, - context: McpContext, -): void { +export function registerGetSavedSearch({ + context, + registerTool, +}: ToolRegistrar): void { const { teamId } = context; const frontendUrl = config.FRONTEND_URL; - server.registerTool( + registerTool( 'clickstack_get_saved_search', { title: 'Get Saved Search(es)', @@ -33,7 +31,7 @@ export function registerGetSavedSearch( ), }), }, - withToolTracing('clickstack_get_saved_search', context, async ({ id }) => { + async ({ id }) => { // ── List all saved searches (slim query — only fetch the fields we need) ── if (!id) { const savedSearches = await SavedSearch.find( @@ -82,6 +80,6 @@ export function registerGetSavedSearch( }, ], }; - }), + }, ); } diff --git a/packages/api/src/mcp/tools/savedSearches/index.ts b/packages/api/src/mcp/tools/savedSearches/index.ts index af8a6e98cf..0707210120 100644 --- a/packages/api/src/mcp/tools/savedSearches/index.ts +++ b/packages/api/src/mcp/tools/savedSearches/index.ts @@ -1,16 +1,11 @@ -import type { McpServer } from '@modelcontextprotocol/sdk/server/mcp.js'; - -import type { McpContext, ToolDefinition } from '@/mcp/tools/types'; +import type { ToolDefinition, ToolRegistrar } from '@/mcp/tools/types'; import { registerGetSavedSearch } from './getSavedSearch'; import { registerSaveSavedSearch } from './saveSavedSearch'; -const savedSearchesTools: ToolDefinition = ( - server: McpServer, - context: McpContext, -) => { - registerGetSavedSearch(server, context); - registerSaveSavedSearch(server, context); +const savedSearchesTools: ToolDefinition = (registrar: ToolRegistrar) => { + registerGetSavedSearch(registrar); + registerSaveSavedSearch(registrar); }; export default savedSearchesTools; diff --git a/packages/api/src/mcp/tools/savedSearches/saveSavedSearch.ts b/packages/api/src/mcp/tools/savedSearches/saveSavedSearch.ts index 7d563c487f..75eadbe686 100644 --- a/packages/api/src/mcp/tools/savedSearches/saveSavedSearch.ts +++ b/packages/api/src/mcp/tools/savedSearches/saveSavedSearch.ts @@ -1,5 +1,3 @@ -import type { McpServer } from '@modelcontextprotocol/sdk/server/mcp.js'; - import * as config from '@/config'; import { createSavedSearch, @@ -7,20 +5,19 @@ import { updateSavedSearch, } from '@/controllers/savedSearch'; import { getSource } from '@/controllers/sources'; -import type { McpContext } from '@/mcp/tools/types'; +import type { ToolRegistrar } from '@/mcp/tools/types'; import { mcpError, validateObjectId } from '@/mcp/utils/errors'; -import { withToolTracing } from '@/mcp/utils/tracing'; import { mcpSaveSavedSearchSchema } from './schemas'; -export function registerSaveSavedSearch( - server: McpServer, - context: McpContext, -): void { +export function registerSaveSavedSearch({ + context, + registerTool, +}: ToolRegistrar): void { const { teamId, userId } = context; const frontendUrl = config.FRONTEND_URL; - server.registerTool( + registerTool( 'clickstack_save_saved_search', { title: 'Create or Update Saved Search', @@ -30,7 +27,7 @@ export function registerSaveSavedSearch( 'Use clickstack_list_sources to find the sourceId.', inputSchema: mcpSaveSavedSearchSchema, }, - withToolTracing('clickstack_save_saved_search', context, async input => { + async input => { const isUpdate = !!input.id; // ── Validate ID for updates ── @@ -127,6 +124,6 @@ export function registerSaveSavedSearch( }, ], }; - }), + }, ); } diff --git a/packages/api/src/mcp/tools/sources/describeMetric.ts b/packages/api/src/mcp/tools/sources/describeMetric.ts index d823975cfd..d92ddcbd68 100644 --- a/packages/api/src/mcp/tools/sources/describeMetric.ts +++ b/packages/api/src/mcp/tools/sources/describeMetric.ts @@ -9,13 +9,11 @@ import { import { ClickhouseClient } from '@hyperdx/common-utils/dist/clickhouse/node'; import { getMetadata } from '@hyperdx/common-utils/dist/core/metadata'; import { SourceKind } from '@hyperdx/common-utils/dist/types'; -import type { McpServer } from '@modelcontextprotocol/sdk/server/mcp.js'; import { z } from 'zod'; import { getConnectionById } from '@/controllers/connection'; import { getSource } from '@/controllers/sources'; -import type { McpContext } from '@/mcp/tools/types'; -import { withToolTracing } from '@/mcp/utils/tracing'; +import type { ToolRegistrar } from '@/mcp/tools/types'; import logger from '@/utils/logger'; import { trimToolResponse } from '@/utils/trimToolResponse'; @@ -724,13 +722,13 @@ async function describeMetricImpl( // ─── Tool registration ─────────────────────────────────────────────────────── -export function registerDescribeMetric( - server: McpServer, - context: McpContext, -): void { +export function registerDescribeMetric({ + context, + registerTool, +}: ToolRegistrar): void { const { teamId } = context; - server.registerTool( + registerTool( 'clickstack_describe_metric', { title: 'Describe Metric', @@ -750,7 +748,7 @@ export function registerDescribeMetric( 'clickstack_describe_metric → clickstack_timeseries|clickstack_table.', inputSchema: describeMetricSchema, }, - withToolTracing('clickstack_describe_metric', context, async rawInput => { + async rawInput => { // Re-parse explicitly: the MCP SDK callback signature widens // optional-field types into `unknown`, but the parser produces // the typed shape we need for downstream calls. @@ -796,6 +794,6 @@ export function registerDescribeMetric( } finally { if (timeoutId !== undefined) clearTimeout(timeoutId); } - }), + }, ); } diff --git a/packages/api/src/mcp/tools/sources/describeSource.ts b/packages/api/src/mcp/tools/sources/describeSource.ts index 5d9b54a624..24621b58b6 100644 --- a/packages/api/src/mcp/tools/sources/describeSource.ts +++ b/packages/api/src/mcp/tools/sources/describeSource.ts @@ -9,13 +9,11 @@ import { import { ClickhouseClient } from '@hyperdx/common-utils/dist/clickhouse/node'; import { getMetadata } from '@hyperdx/common-utils/dist/core/metadata'; import { type MetricTable, SourceKind } from '@hyperdx/common-utils/dist/types'; -import type { McpServer } from '@modelcontextprotocol/sdk/server/mcp.js'; import { z } from 'zod'; import { getConnectionById } from '@/controllers/connection'; import { getSource } from '@/controllers/sources'; -import type { McpContext } from '@/mcp/tools/types'; -import { withToolTracing } from '@/mcp/utils/tracing'; +import type { ToolRegistrar } from '@/mcp/tools/types'; import logger from '@/utils/logger'; import { trimToolResponse } from '@/utils/trimToolResponse'; @@ -629,13 +627,13 @@ async function describeSourceSchema( }; } -export function registerDescribeSource( - server: McpServer, - context: McpContext, -): void { +export function registerDescribeSource({ + context, + registerTool, +}: ToolRegistrar): void { const { teamId } = context; - server.registerTool( + registerTool( 'clickstack_describe_source', { title: 'Describe Source Schema', @@ -661,62 +659,54 @@ export function registerDescribeSource( ), }), }, - withToolTracing( - 'clickstack_describe_source', - context, - async ({ sourceId }) => { - const controller = new AbortController(); - - // Promise.race enforces wall-clock timeout regardless of whether - // internal ClickHouse calls honour the AbortSignal. Hoist the - // timer handle so the finally block can cancel it on the success - // path — otherwise a stale controller.abort() fires - // DESCRIBE_TIMEOUT_MS after every successful call and the - // setTimeout closure stays pinned for the same duration. - let timeoutId: ReturnType | undefined; - const timeoutPromise = new Promise((_, reject) => { - timeoutId = setTimeout(() => { - controller.abort(); - reject(new Error('DESCRIBE_TIMEOUT')); - }, DESCRIBE_TIMEOUT_MS); - }); + async ({ sourceId }) => { + const controller = new AbortController(); + + // Promise.race enforces wall-clock timeout regardless of whether + // internal ClickHouse calls honour the AbortSignal. Hoist the + // timer handle so the finally block can cancel it on the success + // path — otherwise a stale controller.abort() fires + // DESCRIBE_TIMEOUT_MS after every successful call and the + // setTimeout closure stays pinned for the same duration. + let timeoutId: ReturnType | undefined; + const timeoutPromise = new Promise((_, reject) => { + timeoutId = setTimeout(() => { + controller.abort(); + reject(new Error('DESCRIBE_TIMEOUT')); + }, DESCRIBE_TIMEOUT_MS); + }); - try { - return await Promise.race([ - describeSourceSchema( - teamId.toString(), - sourceId, - controller.signal, - ), - timeoutPromise, - ]); - } catch (e) { - if (e instanceof Error && e.message === 'DESCRIBE_TIMEOUT') { - logger.warn( - { teamId, sourceId }, - 'clickstack_describe_source timed out', - ); - return { - isError: true, - content: [ - { - type: 'text' as const, - text: - 'Schema discovery timed out. The ClickHouse server may be under load. ' + - 'Try again, or use clickstack_list_sources for basic source info without schema details.', - }, - ], - }; - } + try { + return await Promise.race([ + describeSourceSchema(teamId.toString(), sourceId, controller.signal), + timeoutPromise, + ]); + } catch (e) { + if (e instanceof Error && e.message === 'DESCRIBE_TIMEOUT') { logger.warn( - { teamId, sourceId, error: e }, - 'Failed to describe source schema', + { teamId, sourceId }, + 'clickstack_describe_source timed out', ); - throw e; - } finally { - if (timeoutId !== undefined) clearTimeout(timeoutId); + return { + isError: true, + content: [ + { + type: 'text' as const, + text: + 'Schema discovery timed out. The ClickHouse server may be under load. ' + + 'Try again, or use clickstack_list_sources for basic source info without schema details.', + }, + ], + }; } - }, - ), + logger.warn( + { teamId, sourceId, error: e }, + 'Failed to describe source schema', + ); + throw e; + } finally { + if (timeoutId !== undefined) clearTimeout(timeoutId); + } + }, ); } diff --git a/packages/api/src/mcp/tools/sources/index.ts b/packages/api/src/mcp/tools/sources/index.ts index 0efefcd786..20e0904636 100644 --- a/packages/api/src/mcp/tools/sources/index.ts +++ b/packages/api/src/mcp/tools/sources/index.ts @@ -1,20 +1,15 @@ -import type { McpServer } from '@modelcontextprotocol/sdk/server/mcp.js'; - -import type { McpContext, ToolDefinition } from '@/mcp/tools/types'; +import type { ToolDefinition, ToolRegistrar } from '@/mcp/tools/types'; import { registerDescribeMetric } from './describeMetric'; import { registerDescribeSource } from './describeSource'; import { registerListMetrics } from './listMetrics'; import { registerListSources } from './listSources'; -const sourcesTools: ToolDefinition = ( - server: McpServer, - context: McpContext, -) => { - registerListSources(server, context); - registerDescribeSource(server, context); - registerListMetrics(server, context); - registerDescribeMetric(server, context); +const sourcesTools: ToolDefinition = (registrar: ToolRegistrar) => { + registerListSources(registrar); + registerDescribeSource(registrar); + registerListMetrics(registrar); + registerDescribeMetric(registrar); }; export default sourcesTools; diff --git a/packages/api/src/mcp/tools/sources/listMetrics.ts b/packages/api/src/mcp/tools/sources/listMetrics.ts index bd59133e85..25edd418b3 100644 --- a/packages/api/src/mcp/tools/sources/listMetrics.ts +++ b/packages/api/src/mcp/tools/sources/listMetrics.ts @@ -6,13 +6,11 @@ import { import { ClickhouseClient } from '@hyperdx/common-utils/dist/clickhouse/node'; import { getMetadata } from '@hyperdx/common-utils/dist/core/metadata'; import { SourceKind } from '@hyperdx/common-utils/dist/types'; -import type { McpServer } from '@modelcontextprotocol/sdk/server/mcp.js'; import { z } from 'zod'; import { getConnectionById } from '@/controllers/connection'; import { getSource } from '@/controllers/sources'; -import type { McpContext } from '@/mcp/tools/types'; -import { withToolTracing } from '@/mcp/utils/tracing'; +import type { ToolRegistrar } from '@/mcp/tools/types'; import logger from '@/utils/logger'; import { @@ -244,13 +242,13 @@ async function fetchMetricsForKind({ // ─── Tool registration ─────────────────────────────────────────────────────── -export function registerListMetrics( - server: McpServer, - context: McpContext, -): void { +export function registerListMetrics({ + context, + registerTool, +}: ToolRegistrar): void { const { teamId } = context; - server.registerTool( + registerTool( 'clickstack_list_metrics', { title: 'List Metric Names', @@ -266,7 +264,7 @@ export function registerListMetrics( 'clickstack_timeseries|clickstack_table.', inputSchema: listMetricsSchema, }, - withToolTracing('clickstack_list_metrics', context, async rawInput => { + async rawInput => { // Re-parse explicitly: the MCP SDK callback signature widens // optional-field types into `unknown`, but the parser produces // the typed shape we need for downstream calls. @@ -309,7 +307,7 @@ export function registerListMetrics( } finally { clearTimeout(timeoutId); } - }), + }, ); } diff --git a/packages/api/src/mcp/tools/sources/listSources.ts b/packages/api/src/mcp/tools/sources/listSources.ts index f394c28c76..2984df7010 100644 --- a/packages/api/src/mcp/tools/sources/listSources.ts +++ b/packages/api/src/mcp/tools/sources/listSources.ts @@ -1,21 +1,19 @@ import { SourceKind } from '@hyperdx/common-utils/dist/types'; -import type { McpServer } from '@modelcontextprotocol/sdk/server/mcp.js'; import { z } from 'zod'; import { getConnectionsByTeam } from '@/controllers/connection'; import { getSources } from '@/controllers/sources'; -import type { McpContext } from '@/mcp/tools/types'; -import { withToolTracing } from '@/mcp/utils/tracing'; +import type { ToolRegistrar } from '@/mcp/tools/types'; import { sanitizeMetricTables } from './metricKinds'; -export function registerListSources( - server: McpServer, - context: McpContext, -): void { +export function registerListSources({ + context, + registerTool, +}: ToolRegistrar): void { const { teamId } = context; - server.registerTool( + registerTool( 'clickstack_list_sources', { title: 'List Sources & Connections', @@ -30,7 +28,7 @@ export function registerListSources( 'Connection IDs are only needed for clickstack_sql (raw ClickHouse SQL).', inputSchema: z.object({}), }, - withToolTracing('clickstack_list_sources', context, async () => { + async () => { const [sources, connections] = await Promise.all([ getSources(teamId.toString()), getConnectionsByTeam(teamId.toString()), @@ -101,6 +99,6 @@ export function registerListSources( { type: 'text' as const, text: JSON.stringify(output, null, 2) }, ], }; - }), + }, ); } diff --git a/packages/api/src/mcp/tools/trace/breakdown.ts b/packages/api/src/mcp/tools/trace/breakdown.ts index 8e39b1b517..20fca3787c 100644 --- a/packages/api/src/mcp/tools/trace/breakdown.ts +++ b/packages/api/src/mcp/tools/trace/breakdown.ts @@ -17,14 +17,12 @@ */ import { ClickhouseClient } from '@hyperdx/common-utils/dist/clickhouse/node'; import { SourceKind } from '@hyperdx/common-utils/dist/types'; -import type { McpServer } from '@modelcontextprotocol/sdk/server/mcp.js'; import { z } from 'zod'; import { getConnectionById } from '@/controllers/connection'; import { getSource } from '@/controllers/sources'; import { parseTimeRange } from '@/mcp/tools/query/helpers'; -import type { McpContext } from '@/mcp/tools/types'; -import { withToolTracing } from '@/mcp/utils/tracing'; +import type { ToolRegistrar } from '@/mcp/tools/types'; // ─── Schema ────────────────────────────────────────────────────────────────── @@ -97,10 +95,13 @@ function durationDivisor(precision: number): number { return Math.pow(10, Math.max(0, precision - 3)); } -export function registerTraceBreakdown(server: McpServer, context: McpContext) { +export function registerTraceBreakdown({ + context, + registerTool, +}: ToolRegistrar) { const { teamId } = context; - server.registerTool( + registerTool( 'clickstack_trace_top_time_consuming_operations', { title: 'Top Time-Consuming Operations Across Matching Traces', @@ -151,85 +152,82 @@ export function registerTraceBreakdown(server: McpServer, context: McpContext) { 'aggregate breakdown has pointed you at the slow downstream operation.', inputSchema: traceBreakdownSchema, }, - withToolTracing( - 'clickstack_trace_top_time_consuming_operations', - context, - async (rawInput: TraceBreakdownInput) => { - const input = traceBreakdownSchema.parse(rawInput); + async (rawInput: TraceBreakdownInput) => { + const input = traceBreakdownSchema.parse(rawInput); - const timeRange = parseTimeRange(input.startTime, input.endTime); - if ('error' in timeRange) { - return { - isError: true as const, - content: [{ type: 'text' as const, text: timeRange.error }], - }; - } - const { startDate, endDate } = timeRange; + const timeRange = parseTimeRange(input.startTime, input.endTime); + if ('error' in timeRange) { + return { + isError: true as const, + content: [{ type: 'text' as const, text: timeRange.error }], + }; + } + const { startDate, endDate } = timeRange; - const source = await getSource(teamId.toString(), input.sourceId); - if (!source) { - return { - isError: true as const, - content: [ - { - type: 'text' as const, - text: `Source not found: ${input.sourceId}. Call clickstack_list_sources to find available source IDs.`, - }, - ], - }; - } - if (source.kind !== SourceKind.Trace) { - return { - isError: true as const, - content: [ - { - type: 'text' as const, - text: `Source ${input.sourceId} is kind="${source.kind}". clickstack_trace_top_time_consuming_operations requires a source of kind="trace".`, - }, - ], - }; - } + const source = await getSource(teamId.toString(), input.sourceId); + if (!source) { + return { + isError: true as const, + content: [ + { + type: 'text' as const, + text: `Source not found: ${input.sourceId}. Call clickstack_list_sources to find available source IDs.`, + }, + ], + }; + } + if (source.kind !== SourceKind.Trace) { + return { + isError: true as const, + content: [ + { + type: 'text' as const, + text: `Source ${input.sourceId} is kind="${source.kind}". clickstack_trace_top_time_consuming_operations requires a source of kind="trace".`, + }, + ], + }; + } - const connection = await getConnectionById( - teamId.toString(), - source.connection.toString(), - true, - ); - if (!connection) { - return { - isError: true as const, - content: [ - { - type: 'text' as const, - text: `Connection not found for source: ${input.sourceId}`, - }, - ], - }; - } + const connection = await getConnectionById( + teamId.toString(), + source.connection.toString(), + true, + ); + if (!connection) { + return { + isError: true as const, + content: [ + { + type: 'text' as const, + text: `Connection not found for source: ${input.sourceId}`, + }, + ], + }; + } - // Source-configured SQL expressions. These are trusted (set by the - // team admin in source config) and we substitute them into the - // generated SQL. - const tsExpr = source.timestampValueExpression; - const traceIdExpr = source.traceIdExpression; - const spanNameExpr = source.spanNameExpression ?? "''"; - const serviceNameExpr = source.serviceNameExpression ?? "''"; - const durationExpr = source.durationExpression; - const divisor = durationDivisor(source.durationPrecision); - const dbName = source.from.databaseName; - const tableName = source.from.tableName; + // Source-configured SQL expressions. These are trusted (set by the + // team admin in source config) and we substitute them into the + // generated SQL. + const tsExpr = source.timestampValueExpression; + const traceIdExpr = source.traceIdExpression; + const spanNameExpr = source.spanNameExpression ?? "''"; + const serviceNameExpr = source.serviceNameExpression ?? "''"; + const durationExpr = source.durationExpression; + const divisor = durationDivisor(source.durationPrecision); + const dbName = source.from.databaseName; + const tableName = source.from.tableName; - // Build the SQL. parentFilter is SQL only (documented); time bounds - // and topN/maxParentTraces are parameterized. - // The CTE selects distinct parent TraceIds. The outer query joins - // back, excludes the parent rows themselves (so we measure child - // contribution), and aggregates by service+operation. - const minDurationClause = - input.minParentDurationMs != null - ? `AND ${durationExpr} >= ({minParentDurationStored:Float64})` - : ''; + // Build the SQL. parentFilter is SQL only (documented); time bounds + // and topN/maxParentTraces are parameterized. + // The CTE selects distinct parent TraceIds. The outer query joins + // back, excludes the parent rows themselves (so we measure child + // contribution), and aggregates by service+operation. + const minDurationClause = + input.minParentDurationMs != null + ? `AND ${durationExpr} >= ({minParentDurationStored:Float64})` + : ''; - const sql = ` + const sql = ` WITH parent_traces AS ( SELECT DISTINCT ${traceIdExpr} AS _trace_id FROM \`${dbName}\`.\`${tableName}\` @@ -257,116 +255,115 @@ ORDER BY total_time_ms DESC LIMIT {topN:UInt32} `; - const params: Record = { - startMs: startDate.getTime(), - endMs: endDate.getTime(), - // Widen the child window by 60s on each side to catch children - // that started slightly before / ended slightly after the parent - // sampling window. - wideStartMs: startDate.getTime() - 60_000, - wideEndMs: endDate.getTime() + 60_000, - maxParentTraces: input.maxParentTraces, - topN: input.topN, - divisor, - }; - if (input.minParentDurationMs != null) { - // Stored duration is divisor × ms. - params.minParentDurationStored = input.minParentDurationMs * divisor; - } + const params: Record = { + startMs: startDate.getTime(), + endMs: endDate.getTime(), + // Widen the child window by 60s on each side to catch children + // that started slightly before / ended slightly after the parent + // sampling window. + wideStartMs: startDate.getTime() - 60_000, + wideEndMs: endDate.getTime() + 60_000, + maxParentTraces: input.maxParentTraces, + topN: input.topN, + divisor, + }; + if (input.minParentDurationMs != null) { + // Stored duration is divisor × ms. + params.minParentDurationStored = input.minParentDurationMs * divisor; + } - const clickhouseClient = new ClickhouseClient({ - host: connection.host, - username: connection.username, - password: connection.password, - }); + const clickhouseClient = new ClickhouseClient({ + host: connection.host, + username: connection.username, + password: connection.password, + }); - type Row = { - service: string; - operation: string; - total_time_ms: number | string; - calls: number | string; - in_parents: number | string; - p50_ms: number | string; - p99_ms: number | string; - }; + type Row = { + service: string; + operation: string; + total_time_ms: number | string; + calls: number | string; + in_parents: number | string; + p50_ms: number | string; + p99_ms: number | string; + }; - let rows: Row[]; - try { - const result = await clickhouseClient.query({ - query: sql, - query_params: params, - format: 'JSON', - connectionId: source.connection.toString(), - clickhouse_settings: { - // Prevent DDL/DML injection via parentFilter — only SELECTs allowed. - readonly: '1', - ...(source.querySettings - ? Object.fromEntries( - source.querySettings.map(s => [s.setting, s.value]), - ) - : {}), - }, - }); - const json = (await ( - result as { json: () => Promise<{ data: Row[] }> } - ).json()) ?? { data: [] }; - rows = json.data ?? []; - } catch (e) { - return { - isError: true as const, - content: [ - { - type: 'text' as const, - text: `Failed to compute breakdown: ${e instanceof Error ? e.message : String(e)}. The parentFilter must be valid ClickHouse SQL referencing columns on the trace table (e.g. ServiceName = 'X' AND SpanName = 'Y').`, - }, - ], - }; - } - - // Normalise numerics — ClickHouse JSON sometimes returns strings for - // 64-bit integers. Cast everything to Number and compute share of - // total time at the same time. - const operations = rows.map(r => ({ - service: r.service, - operation: r.operation, - totalTimeMs: Number(r.total_time_ms), - calls: Number(r.calls), - inParents: Number(r.in_parents), - p50Ms: Number(r.p50_ms), - p99Ms: Number(r.p99_ms), - })); - const grandTotalMs = operations.reduce( - (acc, r) => acc + r.totalTimeMs, - 0, - ); - const operationsWithShare = operations.map(r => ({ - ...r, - shareOfTotalTime: grandTotalMs > 0 ? r.totalTimeMs / grandTotalMs : 0, - })); - - const output = { - summary: { - parentFilter: input.parentFilter, - startTime: startDate.toISOString(), - endTime: endDate.toISOString(), - minParentDurationMs: input.minParentDurationMs ?? null, - operationsReturned: operationsWithShare.length, - topN: input.topN, - grandTotalTimeMs: grandTotalMs, - hint: - operationsWithShare.length === 0 - ? 'No matching parent traces, or all matching traces had no other spans. Widen the time window, relax the parentFilter, or lower minParentDurationMs.' - : "Operations are sorted by total_time_ms (sum of child Duration). shareOfTotalTime is each row's share of the cumulative child time across all matching traces.", + let rows: Row[]; + try { + const result = await clickhouseClient.query({ + query: sql, + query_params: params, + format: 'JSON', + connectionId: source.connection.toString(), + clickhouse_settings: { + // Prevent DDL/DML injection via parentFilter — only SELECTs allowed. + readonly: '1', + ...(source.querySettings + ? Object.fromEntries( + source.querySettings.map(s => [s.setting, s.value]), + ) + : {}), }, - operations: operationsWithShare, - }; - + }); + const json = (await ( + result as { json: () => Promise<{ data: Row[] }> } + ).json()) ?? { data: [] }; + rows = json.data ?? []; + } catch (e) { return { + isError: true as const, content: [ - { type: 'text' as const, text: JSON.stringify(output, null, 2) }, + { + type: 'text' as const, + text: `Failed to compute breakdown: ${e instanceof Error ? e.message : String(e)}. The parentFilter must be valid ClickHouse SQL referencing columns on the trace table (e.g. ServiceName = 'X' AND SpanName = 'Y').`, + }, ], }; - }, - ), + } + + // Normalise numerics — ClickHouse JSON sometimes returns strings for + // 64-bit integers. Cast everything to Number and compute share of + // total time at the same time. + const operations = rows.map(r => ({ + service: r.service, + operation: r.operation, + totalTimeMs: Number(r.total_time_ms), + calls: Number(r.calls), + inParents: Number(r.in_parents), + p50Ms: Number(r.p50_ms), + p99Ms: Number(r.p99_ms), + })); + const grandTotalMs = operations.reduce( + (acc, r) => acc + r.totalTimeMs, + 0, + ); + const operationsWithShare = operations.map(r => ({ + ...r, + shareOfTotalTime: grandTotalMs > 0 ? r.totalTimeMs / grandTotalMs : 0, + })); + + const output = { + summary: { + parentFilter: input.parentFilter, + startTime: startDate.toISOString(), + endTime: endDate.toISOString(), + minParentDurationMs: input.minParentDurationMs ?? null, + operationsReturned: operationsWithShare.length, + topN: input.topN, + grandTotalTimeMs: grandTotalMs, + hint: + operationsWithShare.length === 0 + ? 'No matching parent traces, or all matching traces had no other spans. Widen the time window, relax the parentFilter, or lower minParentDurationMs.' + : "Operations are sorted by total_time_ms (sum of child Duration). shareOfTotalTime is each row's share of the cumulative child time across all matching traces.", + }, + operations: operationsWithShare, + }; + + return { + content: [ + { type: 'text' as const, text: JSON.stringify(output, null, 2) }, + ], + }; + }, ); } diff --git a/packages/api/src/mcp/tools/trace/index.ts b/packages/api/src/mcp/tools/trace/index.ts index 9a05327d3a..971bcc7b43 100644 --- a/packages/api/src/mcp/tools/trace/index.ts +++ b/packages/api/src/mcp/tools/trace/index.ts @@ -1,13 +1,11 @@ -import type { McpServer } from '@modelcontextprotocol/sdk/server/mcp.js'; - -import type { McpContext, ToolDefinition } from '@/mcp/tools/types'; +import type { ToolDefinition, ToolRegistrar } from '@/mcp/tools/types'; import { registerTraceBreakdown } from './breakdown'; import { registerTraceWaterfall } from './waterfall'; -const traceTools: ToolDefinition = (server: McpServer, context: McpContext) => { - registerTraceWaterfall(server, context); - registerTraceBreakdown(server, context); +const traceTools: ToolDefinition = (registrar: ToolRegistrar) => { + registerTraceWaterfall(registrar); + registerTraceBreakdown(registrar); }; export default traceTools; diff --git a/packages/api/src/mcp/tools/trace/waterfall.ts b/packages/api/src/mcp/tools/trace/waterfall.ts index 881d992c02..57611852a1 100644 --- a/packages/api/src/mcp/tools/trace/waterfall.ts +++ b/packages/api/src/mcp/tools/trace/waterfall.ts @@ -6,14 +6,12 @@ import { DisplayType, SourceKind, } from '@hyperdx/common-utils/dist/types'; -import type { McpServer } from '@modelcontextprotocol/sdk/server/mcp.js'; import { z } from 'zod'; import { getConnectionById } from '@/controllers/connection'; import { getSource } from '@/controllers/sources'; import { parseTimeRange } from '@/mcp/tools/query/helpers'; -import type { McpContext } from '@/mcp/tools/types'; -import { withToolTracing } from '@/mcp/utils/tracing'; +import type { ToolRegistrar } from '@/mcp/tools/types'; // ─── Schema ────────────────────────────────────────────────────────────────── @@ -164,10 +162,13 @@ function durationDivisor(precision: number): number { // ─── Tool definition ───────────────────────────────────────────────────────── -export function registerTraceWaterfall(server: McpServer, context: McpContext) { +export function registerTraceWaterfall({ + context, + registerTool, +}: ToolRegistrar) { const { teamId } = context; - server.registerTool( + registerTool( 'clickstack_trace_waterfall', { title: 'Trace Waterfall (single trace)', @@ -205,198 +206,195 @@ export function registerTraceWaterfall(server: McpServer, context: McpContext) { 'me one example of a slow X" (single trace).', inputSchema: traceSchema, }, - withToolTracing( - 'clickstack_trace_waterfall', - context, - async (rawInput: TraceInput) => { - const input = traceSchema.parse(rawInput); - const timeRange = parseTimeRange(input.startTime, input.endTime); - if ('error' in timeRange) { - return { - isError: true as const, - content: [{ type: 'text' as const, text: timeRange.error }], - }; + async (rawInput: TraceInput) => { + const input = traceSchema.parse(rawInput); + const timeRange = parseTimeRange(input.startTime, input.endTime); + if ('error' in timeRange) { + return { + isError: true as const, + content: [{ type: 'text' as const, text: timeRange.error }], + }; + } + const { startDate, endDate } = timeRange; + + const source = await getSource(teamId.toString(), input.sourceId); + if (!source) { + return { + isError: true as const, + content: [ + { + type: 'text' as const, + text: `Source not found: ${input.sourceId}. Call clickstack_list_sources to find available source IDs.`, + }, + ], + }; + } + if (source.kind !== SourceKind.Trace) { + return { + isError: true as const, + content: [ + { + type: 'text' as const, + text: `Source ${input.sourceId} is kind="${source.kind}". clickstack_trace_waterfall requires a source of kind="trace".`, + }, + ], + }; + } + + const connection = await getConnectionById( + teamId.toString(), + source.connection.toString(), + true, + ); + if (!connection) { + return { + isError: true as const, + content: [ + { + type: 'text' as const, + text: `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; + const spanIdExpr = source.spanIdExpression; + const parentSpanIdExpr = source.parentSpanIdExpression; + const spanNameExpr = source.spanNameExpression; + const spanKindExpr = source.spanKindExpression; + const durationExpr = source.durationExpression; + const tsExpr = source.timestampValueExpression; + const 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. + let effectiveFilter = input.pickFilter; + let effectiveLanguage = input.pickFilterLanguage; + if (input.pickBy === 'first_error') { + // Statuses are typically stored as enum strings. Use SQL so we don't + // depend on lucene mapping the raw column name. + const errFilter = `${statusCodeExpr} = 'STATUS_CODE_ERROR'`; + if (effectiveFilter && effectiveLanguage === 'sql') { + effectiveFilter = `(${effectiveFilter}) AND (${errFilter})`; + } else if (effectiveFilter && effectiveLanguage === 'lucene') { + // Run lucene filter inside parens, AND with raw SQL. + // Easiest: convert to sql by AND-joining a fresh sql-only condition + // means we'd have to render lucene first. Compromise: tell user + // to express the error filter in pickFilter directly. + effectiveFilter = `(${effectiveFilter}) AND StatusCode:STATUS_CODE_ERROR`; + effectiveLanguage = 'lucene'; + } else { + effectiveFilter = errFilter; + effectiveLanguage = 'sql'; + } } - const { startDate, endDate } = timeRange; - const source = await getSource(teamId.toString(), input.sourceId); - if (!source) { + const orderBy = + input.pickBy === 'slowest' + ? `max(${durationExpr}) DESC` + : input.pickBy === 'first_error' + ? `min(${tsExpr}) ASC` + : `max(${tsExpr}) DESC`; + + const pickConfig: BuilderChartConfigWithDateRange = { + displayType: DisplayType.Table, + select: [ + { + aggFn: 'count' as const, + valueExpression: '', + alias: 'span_count', + aggCondition: '', + aggConditionLanguage: 'sql' as const, + }, + ], + from: source.from, + where: effectiveFilter, + whereLanguage: effectiveLanguage, + connection: source.connection.toString(), + timestampValueExpression: tsExpr, + implicitColumnExpression: source.implicitColumnExpression, + groupBy: traceIdExpr, + orderBy, + limit: { limit: 1 }, + dateRange: [startDate, endDate], + }; + + let pickResult: { data?: Array> } = {}; + try { + pickResult = (await clickhouseClient.queryChartConfig({ + config: pickConfig as ChartConfigWithDateRange, + metadata, + querySettings: source.querySettings, + })) as { data?: Array> }; + } catch (e) { return { isError: true as const, content: [ { type: 'text' as const, - text: `Source not found: ${input.sourceId}. Call clickstack_list_sources to find available source IDs.`, + text: `Failed to pick a trace: ${e instanceof Error ? e.message : String(e)}`, }, ], }; } - if (source.kind !== SourceKind.Trace) { + + const firstRow = pickResult.data?.[0]; + if (!firstRow) { return { - isError: true as const, content: [ { type: 'text' as const, - text: `Source ${input.sourceId} is kind="${source.kind}". clickstack_trace_waterfall requires a source of kind="trace".`, + text: JSON.stringify( + { + result: null, + hint: + 'No traces matched. Widen the time range, relax pickFilter, or ' + + 'pass a specific traceId.', + pickFilter: effectiveFilter, + pickBy: input.pickBy, + }, + null, + 2, + ), }, ], }; } - - const connection = await getConnectionById( - teamId.toString(), - source.connection.toString(), - true, + // The grouped traceId column is rendered into the result with the + // expression text as its alias. Locate it by stripping non-data keys. + const candidate = Object.entries(firstRow).find( + ([k]) => k !== 'span_count' && k !== '__hdx_time_bucket', ); - if (!connection) { + if (!candidate || candidate[1] == null) { return { isError: true as const, content: [ { type: 'text' as const, - text: `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; - const spanIdExpr = source.spanIdExpression; - const parentSpanIdExpr = source.parentSpanIdExpression; - const spanNameExpr = source.spanNameExpression; - const spanKindExpr = source.spanKindExpression; - const durationExpr = source.durationExpression; - const tsExpr = source.timestampValueExpression; - const 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. - let effectiveFilter = input.pickFilter; - let effectiveLanguage = input.pickFilterLanguage; - if (input.pickBy === 'first_error') { - // Statuses are typically stored as enum strings. Use SQL so we don't - // depend on lucene mapping the raw column name. - const errFilter = `${statusCodeExpr} = 'STATUS_CODE_ERROR'`; - if (effectiveFilter && effectiveLanguage === 'sql') { - effectiveFilter = `(${effectiveFilter}) AND (${errFilter})`; - } else if (effectiveFilter && effectiveLanguage === 'lucene') { - // Run lucene filter inside parens, AND with raw SQL. - // Easiest: convert to sql by AND-joining a fresh sql-only condition - // means we'd have to render lucene first. Compromise: tell user - // to express the error filter in pickFilter directly. - effectiveFilter = `(${effectiveFilter}) AND StatusCode:STATUS_CODE_ERROR`; - effectiveLanguage = 'lucene'; - } else { - effectiveFilter = errFilter; - effectiveLanguage = 'sql'; - } - } - - const orderBy = - input.pickBy === 'slowest' - ? `max(${durationExpr}) DESC` - : input.pickBy === 'first_error' - ? `min(${tsExpr}) ASC` - : `max(${tsExpr}) DESC`; - - const pickConfig: BuilderChartConfigWithDateRange = { - displayType: DisplayType.Table, - select: [ - { - aggFn: 'count' as const, - valueExpression: '', - alias: 'span_count', - aggCondition: '', - aggConditionLanguage: 'sql' as const, + text: `Picker returned a row but no TraceId column was found. Row keys: ${Object.keys(firstRow).join(', ')}`, }, ], - from: source.from, - where: effectiveFilter, - whereLanguage: effectiveLanguage, - connection: source.connection.toString(), - timestampValueExpression: tsExpr, - implicitColumnExpression: source.implicitColumnExpression, - groupBy: traceIdExpr, - orderBy, - limit: { limit: 1 }, - dateRange: [startDate, endDate], }; - - let pickResult: { data?: Array> } = {}; - try { - pickResult = (await clickhouseClient.queryChartConfig({ - config: pickConfig as ChartConfigWithDateRange, - metadata, - querySettings: source.querySettings, - })) as { data?: Array> }; - } catch (e) { - return { - isError: true as const, - content: [ - { - type: 'text' as const, - text: `Failed to pick a trace: ${e instanceof Error ? e.message : String(e)}`, - }, - ], - }; - } - - const firstRow = pickResult.data?.[0]; - if (!firstRow) { - return { - content: [ - { - type: 'text' as const, - text: JSON.stringify( - { - result: null, - hint: - 'No traces matched. Widen the time range, relax pickFilter, or ' + - 'pass a specific traceId.', - pickFilter: effectiveFilter, - pickBy: input.pickBy, - }, - null, - 2, - ), - }, - ], - }; - } - // The grouped traceId column is rendered into the result with the - // expression text as its alias. Locate it by stripping non-data keys. - const candidate = Object.entries(firstRow).find( - ([k]) => k !== 'span_count' && k !== '__hdx_time_bucket', - ); - if (!candidate || candidate[1] == null) { - return { - isError: true as const, - content: [ - { - type: 'text' as const, - text: `Picker returned a row but no TraceId column was found. Row keys: ${Object.keys(firstRow).join(', ')}`, - }, - ], - }; - } - pickedTraceId = String(candidate[1]); } + pickedTraceId = String(candidate[1]); + } - // ── Step 2: fetch the full span tree ── - const treeQuery = ` + // ── Step 2: fetch the full span tree ── + const treeQuery = ` SELECT ${spanIdExpr} AS spanId, ${parentSpanIdExpr} AS parentSpanId, @@ -414,120 +412,120 @@ export function registerTraceWaterfall(server: McpServer, context: McpContext) { LIMIT {n:UInt32} `; - let rows: SpanRow[]; - try { - const result = await clickhouseClient.query({ - query: treeQuery, - query_params: { - db: source.from.databaseName, - tbl: source.from.tableName, - tid: pickedTraceId, - n: input.maxSpans + 1, // +1 to detect truncation - divisor, - }, - format: 'JSONEachRow', - connectionId: source.connection.toString(), - clickhouse_settings: { - 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]), - ) - : {}), + let rows: SpanRow[]; + try { + const result = await clickhouseClient.query({ + query: treeQuery, + query_params: { + db: source.from.databaseName, + tbl: source.from.tableName, + tid: pickedTraceId, + n: input.maxSpans + 1, // +1 to detect truncation + divisor, + }, + format: 'JSONEachRow', + connectionId: source.connection.toString(), + clickhouse_settings: { + 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]), + ) + : {}), + }, + }); + rows = + (await (result as { json: () => Promise }).json()) ?? []; + } catch (e) { + return { + isError: true as const, + content: [ + { + type: 'text' as const, + text: `Failed to fetch trace ${pickedTraceId}: ${e instanceof Error ? e.message : String(e)}`, }, - }); - rows = - (await (result as { json: () => Promise }).json()) ?? []; - } catch (e) { - return { - isError: true as const, - content: [ - { - type: 'text' as const, - text: `Failed to fetch trace ${pickedTraceId}: ${e instanceof Error ? e.message : String(e)}`, - }, - ], - }; - } + ], + }; + } - const truncated = rows.length > input.maxSpans; - const spans = truncated ? rows.slice(0, input.maxSpans) : rows; + const truncated = rows.length > input.maxSpans; + const spans = truncated ? rows.slice(0, input.maxSpans) : rows; - if (spans.length === 0) { - return { - content: [ - { - type: 'text' as const, - text: JSON.stringify( - { - 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.', - }, - null, - 2, - ), - }, - ], - }; - } - - const tree = buildPreOrderTree(spans); - const root = tree.find(s => s.depth === 0) ?? tree[0]; - const totalDuration = Math.max(...spans.map(s => s.durationMs)); - - // ── Step 3: fetch correlated logs (when logSourceId is configured) ── - type LogRow = { - timestamp: string; - severityText: string; - body: string; - serviceName: string; - spanId: string; + if (spans.length === 0) { + return { + content: [ + { + type: 'text' as const, + text: JSON.stringify( + { + 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.', + }, + null, + 2, + ), + }, + ], }; - let correlatedLogs: LogRow[] | undefined; - let logsTruncated = false; - let logsNote: string | undefined; - if (input.includeLogs && source.logSourceId) { - const logSource = await getSource( - teamId.toString(), - source.logSourceId, - ); - if (!logSource) { - logsNote = `logSourceId ${source.logSourceId} configured but source not found`; - } else if (logSource.kind !== SourceKind.Log) { - logsNote = `logSourceId ${source.logSourceId} is kind="${logSource.kind}", not "log"`; - } else { - const logTraceIdExpr = logSource.traceIdExpression ?? 'TraceId'; - const logSpanIdExpr = logSource.spanIdExpression ?? "''"; - const 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; - if ( - logSource.connection.toString() !== source.connection.toString() - ) { - const logConn = await getConnectionById( - 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, - }); - } + } + + const tree = buildPreOrderTree(spans); + 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) ── + type LogRow = { + timestamp: string; + severityText: string; + body: string; + serviceName: string; + spanId: string; + }; + let correlatedLogs: LogRow[] | undefined; + let logsTruncated = false; + let logsNote: string | undefined; + if (input.includeLogs && source.logSourceId) { + const logSource = await getSource( + teamId.toString(), + source.logSourceId, + ); + if (!logSource) { + logsNote = `logSourceId ${source.logSourceId} configured but source not found`; + } else if (logSource.kind !== SourceKind.Log) { + logsNote = `logSourceId ${source.logSourceId} is kind="${logSource.kind}", not "log"`; + } else { + const logTraceIdExpr = logSource.traceIdExpression ?? 'TraceId'; + const logSpanIdExpr = logSource.spanIdExpression ?? "''"; + const 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; + if ( + logSource.connection.toString() !== source.connection.toString() + ) { + const logConn = await getConnectionById( + 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 (!logsNote) { - const logsQuery = ` + if (!logsNote) { + const logsQuery = ` SELECT ${logTsExpr} AS timestamp, ${logSevExpr} AS severityText, @@ -539,81 +537,77 @@ export function registerTraceWaterfall(server: McpServer, context: McpContext) { ORDER BY ${logTsExpr} ASC LIMIT {n:UInt32} `; - try { - const logResult = await logClient.query({ - query: logsQuery, - query_params: { - db: logSource.from.databaseName, - tbl: logSource.from.tableName, - tid: pickedTraceId, - n: input.maxLogs + 1, // +1 to detect truncation - }, - format: 'JSONEachRow', - connectionId: logSource.connection.toString(), - clickhouse_settings: { - readonly: '1', - ...(logSource.querySettings - ? Object.fromEntries( - logSource.querySettings.map(s => [ - s.setting, - s.value, - ]), - ) - : {}), - }, - }); - const allLogs = - (await ( - logResult as { json: () => Promise } - ).json()) ?? []; - logsTruncated = allLogs.length > input.maxLogs; - correlatedLogs = logsTruncated - ? allLogs.slice(0, input.maxLogs) - : allLogs; - } catch (e) { - logsNote = `Failed to fetch correlated logs: ${e instanceof Error ? e.message : String(e)}`; - } + try { + const logResult = await logClient.query({ + query: logsQuery, + query_params: { + db: logSource.from.databaseName, + tbl: logSource.from.tableName, + tid: pickedTraceId, + n: input.maxLogs + 1, // +1 to detect truncation + }, + format: 'JSONEachRow', + connectionId: logSource.connection.toString(), + clickhouse_settings: { + readonly: '1', + ...(logSource.querySettings + ? Object.fromEntries( + logSource.querySettings.map(s => [s.setting, s.value]), + ) + : {}), + }, + }); + const allLogs = + (await ( + logResult as { json: () => Promise } + ).json()) ?? []; + logsTruncated = allLogs.length > input.maxLogs; + correlatedLogs = logsTruncated + ? allLogs.slice(0, input.maxLogs) + : allLogs; + } catch (e) { + logsNote = `Failed to fetch correlated logs: ${e instanceof Error ? e.message : String(e)}`; } } } - - const output = { - traceId: pickedTraceId, - spanCount: spans.length, - totalDurationMs: totalDuration, - rootSpan: { - spanId: root.spanId, - serviceName: root.serviceName, - spanName: root.spanName, - durationMs: root.durationMs, - statusCode: root.statusCode, - }, - spans: tree, - ...(correlatedLogs - ? { - logs: correlatedLogs, - logsCount: correlatedLogs.length, - ...(logsTruncated - ? { - logsNote: `Logs truncated to ${input.maxLogs}. Increase maxLogs (max 1000) if needed.`, - } - : {}), - } - : {}), - ...(logsNote ? { logsNote } : {}), - ...(truncated - ? { - note: `Result truncated to ${input.maxSpans} spans. Increase maxSpans (max 2000) or narrow the trace if needed.`, - } - : {}), - }; - - return { - content: [ - { type: 'text' as const, text: JSON.stringify(output, null, 2) }, - ], - }; - }, - ), + } + + const output = { + traceId: pickedTraceId, + spanCount: spans.length, + totalDurationMs: totalDuration, + rootSpan: { + spanId: root.spanId, + serviceName: root.serviceName, + spanName: root.spanName, + durationMs: root.durationMs, + statusCode: root.statusCode, + }, + spans: tree, + ...(correlatedLogs + ? { + logs: correlatedLogs, + logsCount: correlatedLogs.length, + ...(logsTruncated + ? { + logsNote: `Logs truncated to ${input.maxLogs}. Increase maxLogs (max 1000) if needed.`, + } + : {}), + } + : {}), + ...(logsNote ? { logsNote } : {}), + ...(truncated + ? { + note: `Result truncated to ${input.maxSpans} spans. Increase maxSpans (max 2000) or narrow the trace if needed.`, + } + : {}), + }; + + return { + content: [ + { type: 'text' as const, text: JSON.stringify(output, null, 2) }, + ], + }; + }, ); } diff --git a/packages/api/src/mcp/tools/types.ts b/packages/api/src/mcp/tools/types.ts index c3cc2d293e..94ade3d11b 100644 --- a/packages/api/src/mcp/tools/types.ts +++ b/packages/api/src/mcp/tools/types.ts @@ -1,10 +1,46 @@ import type { McpServer } from '@modelcontextprotocol/sdk/server/mcp.js'; +import type { CallToolResult } from '@modelcontextprotocol/sdk/types.js'; +import type { AnyZodObject } from 'zod'; export type McpContext = { teamId: string; userId: string; }; -export type ToolDefinition = (server: McpServer, context: McpContext) => void; +/** + * The result shape every MCP tool handler should return. + * + * Intersects the SDK's `CallToolResult` (which carries an index signature + * from the `$loose` Zod modifier) with a narrower `content` array so tool + * handlers are constrained to text-only content blocks. + */ +export type ToolResult = CallToolResult & { + content: { type: 'text'; text: string }[]; +}; + +/** + * A simplified tool registration function that wraps `server.registerTool` + * with automatic tracing. Eliminates the need to: + * - Pass the tool name twice (once to registerTool, once to withToolTracing) + * - Import and manually wire up withToolTracing in every tool file + * - Import McpServer type in every tool file + */ +export type RegisterToolFn = ( + name: string, + config: { + title: string; + description: string; + inputSchema: TSchema; + }, + handler: (args: TSchema['_output']) => Promise, +) => void; + +export type ToolRegistrar = { + server: McpServer; + context: McpContext; + registerTool: RegisterToolFn; +}; + +export type ToolDefinition = (registrar: ToolRegistrar) => void; export type PromptDefinition = (server: McpServer, context: McpContext) => void; diff --git a/packages/api/src/mcp/utils/registerTool.ts b/packages/api/src/mcp/utils/registerTool.ts new file mode 100644 index 0000000000..c5a93be739 --- /dev/null +++ b/packages/api/src/mcp/utils/registerTool.ts @@ -0,0 +1,34 @@ +import type { + McpServer, + ToolCallback, +} from '@modelcontextprotocol/sdk/server/mcp.js'; +import type { AnyZodObject } from 'zod'; + +import type { McpContext, RegisterToolFn } from '@/mcp/tools/types'; + +import { withToolTracing } from './tracing'; + +/** + * Creates a `registerTool` function bound to a specific server and context. + * The returned function automatically wraps every handler with + * `withToolTracing`, so individual tool files don't need to import or call + * tracing utilities. + */ +export function createRegisterTool( + server: McpServer, + context: McpContext, +): RegisterToolFn { + return (name, config, handler) => { + // Wrap with tracing, then register. The explicit InputArgs generic + // binds the SDK's own type parameter to AnyZodObject so TypeScript + // resolves ToolCallback via the AnySchema branch of BaseToolCallback. + // This lets it accept our traced callback without a type assertion — + // both sides reduce to (args: SchemaOutput, extra) => …. + const traced: ToolCallback = withToolTracing( + name, + context, + handler, + ); + server.registerTool(name, config, traced); + }; +} diff --git a/packages/api/src/mcp/utils/tracing.ts b/packages/api/src/mcp/utils/tracing.ts index 3162cf0d8c..d21fe1cb64 100644 --- a/packages/api/src/mcp/utils/tracing.ts +++ b/packages/api/src/mcp/utils/tracing.ts @@ -1,4 +1,6 @@ -import type { McpContext } from '@/mcp/tools/types'; +import type { CallToolResult } from '@modelcontextprotocol/sdk/types.js'; + +import type { McpContext, ToolResult } from '@/mcp/tools/types'; import { getCounter, getHistogram, @@ -7,11 +9,6 @@ import { } from '@/utils/instrumentation'; import logger from '@/utils/logger'; -type ToolResult = { - content: { type: 'text'; text: string }[]; - isError?: boolean; -}; - const toolDurationHistogram = getHistogram('hyperdx.mcp.tool.duration_ms', { description: 'Wall-clock duration of an MCP tool invocation.', unit: 'ms', @@ -25,12 +22,16 @@ const toolErrorCounter = getCounter('hyperdx.mcp.tool.errors', { /** * Wraps an MCP tool handler with tracing, metrics, and structured logging. * Creates a span for each tool invocation and logs start/end with duration. + * + * The returned function signature is a strict subset of the SDK's + * `ToolCallback`: it accepts `(args, _extra?)` and returns + * `Promise`. The extra parameter is accepted but unused. */ export function withToolTracing( toolName: string, context: McpContext, handler: (args: TArgs) => Promise, -): (args: TArgs) => Promise { +): (args: TArgs, _extra?: unknown) => Promise { return async (args: TArgs) => { const logContext = { tool: toolName,