refactor(mcp): add registerTool wrapper to deduplicate tracing boilerplate

packages/api/src/mcp/mcpServer.ts

@@ -10,19 +10,23 @@ 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({
     name: 'clickstack',
     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;

packages/api/src/mcp/tools/alerts/getAlert.ts

@@ -1,15 +1,13 @@
 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';
 
 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<string, unknown> = {
@@ -151,6 +152,6 @@ export function registerGetAlert(server: McpServer, context: McpContext): void {
           },
         ],
       };
-    }),
+    },
   );
 }

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) },
         ],
       };
-    }),
+    },
   );
 }

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;

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(
           },
         ],
       };
-    }),
+    },
   );
 }

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),
+          },
+        ],
+      };
+    },
   );
 }

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(
           },
         ],
       };
-    }),
+    },
   );
 }