feat: add traces tool for Datadog APM trace retrieval (#6)

README.md

@@ -52,13 +52,26 @@ MCP server for the Datadog API, enabling incident management and more.
    - **Returns**: Array of dashboards with URL references.
 
 6. `get_metrics`
+
    - Retrieve metrics data from Datadog.
    - **Inputs**:
      - `query` (string): Metrics query string.
      - `from` (number): Start time in epoch seconds.
      - `to` (number): End time in epoch seconds.
    - **Returns**: Metrics data for the queried timeframe.
 
+7. `list_traces`
+   - Retrieve a list of APM traces from Datadog.
+   - **Inputs**:
+     - `query` (string): Datadog APM trace query string.
+     - `from` (number): Start time in epoch seconds.
+     - `to` (number): End time in epoch seconds.
+     - `limit` (optional number): Maximum number of traces to return (defaults to 100).
+     - `sort` (optional string): Sort order for traces (defaults to '-timestamp').
+     - `service` (optional string): Filter by service name.
+     - `operation` (optional string): Filter by operation name.
+   - **Returns**: Array of matching traces from Datadog APM.
+
 ## Setup
 
 ### Datadog Credentials

package.json

@@ -1,6 +1,6 @@
 {
   "name": "@winor30/mcp-server-datadog",
-  "version": "1.0.0",
+  "version": "1.0.1",
   "description": "MCP server for interacting with Datadog API",
   "repository": {
     "type": "git",

src/index.ts

@@ -19,6 +19,7 @@ import { METRICS_TOOLS, METRICS_HANDLERS } from './tools/metrics'
 import { LOGS_TOOLS, LOGS_HANDLERS } from './tools/logs'
 import { MONITORS_TOOLS, MONITORS_HANDLERS } from './tools/monitors'
 import { DASHBOARDS_TOOLS, DASHBOARDS_HANDLERS } from './tools/dashboards'
+import { TRACES_TOOLS, TRACES_HANDLERS } from './tools/traces'
 import { ToolHandlers } from './utils/types'
 
 const server = new Server(
@@ -49,6 +50,7 @@ server.setRequestHandler(ListToolsRequestSchema, async () => {
       ...LOGS_TOOLS,
       ...MONITORS_TOOLS,
       ...DASHBOARDS_TOOLS,
+      ...TRACES_TOOLS,
     ],
   }
 })
@@ -59,6 +61,7 @@ const TOOL_HANDLERS: ToolHandlers = {
   ...LOGS_HANDLERS,
   ...MONITORS_HANDLERS,
   ...DASHBOARDS_HANDLERS,
+  ...TRACES_HANDLERS,
 }
 /**
  * Handler for invoking Datadog-related tools in the mcp-server-datadog.

src/tools/traces/index.ts

@@ -0,0 +1 @@
+export { TRACES_TOOLS, TRACES_HANDLERS } from './tool'

src/tools/traces/schema.ts

@@ -0,0 +1,21 @@
+import { z } from 'zod'
+
+export const ListTracesZodSchema = z.object({
+  query: z.string().describe('Datadog APM trace query string'),
+  from: z.number().describe('Start time in epoch seconds'),
+  to: z.number().describe('End time in epoch seconds'),
+  limit: z
+    .number()
+    .optional()
+    .default(100)
+    .describe('Maximum number of traces to return'),
+  sort: z
+    .enum(['timestamp', '-timestamp'])
+    .optional()
+    .default('-timestamp')
+    .describe('Sort order for traces'),
+  service: z.string().optional().describe('Filter by service name'),
+  operation: z.string().optional().describe('Filter by operation name'),
+})
+
+export type ListTracesArgs = z.infer<typeof ListTracesZodSchema>

src/tools/traces/tool.ts

@@ -0,0 +1,79 @@
+import { ExtendedTool, ToolHandlers } from '../../utils/types'
+import { v2 } from '@datadog/datadog-api-client'
+import { createToolSchema } from '../../utils/tool'
+import { datadogConfig as config } from '../../utils/datadog'
+import { ListTracesZodSchema } from './schema'
+
+type TracesToolName = 'list_traces'
+type TracesTool = ExtendedTool<TracesToolName>
+
+export const TRACES_TOOLS: TracesTool[] = [
+  createToolSchema(
+    ListTracesZodSchema,
+    'list_traces',
+    'Get APM traces from Datadog',
+  ),
+] as const
+
+const API_INSTANCE = new v2.SpansApi(config)
+
+type TracesToolHandlers = ToolHandlers<TracesToolName>
+
+export const TRACES_HANDLERS: TracesToolHandlers = {
+  list_traces: async (request) => {
+    const {
+      query,
+      from,
+      to,
+      limit = 100,
+      sort = '-timestamp',
+      service,
+      operation,
+    } = request.params.arguments as {
+      query: string
+      from: number
+      to: number
+      limit?: number
+      sort?: string
+      service?: string
+      operation?: string
+    }
+
+    const response = await API_INSTANCE.listSpans({
+      body: {
+        data: {
+          attributes: {
+            filter: {
+              query: [
+                query,
+                ...(service ? [`service:${service}`] : []),
+                ...(operation ? [`operation:${operation}`] : []),
+              ].join(' '),
+              from: new Date(from * 1000).toISOString(),
+              to: new Date(to * 1000).toISOString(),
+            },
+            sort: sort as 'timestamp' | '-timestamp',
+            page: { limit },
+          },
+          type: 'search_request',
+        },
+      },
+    })
+
+    if (!response.data) {
+      throw new Error('No traces data returned')
+    }
+
+    return {
+      content: [
+        {
+          type: 'text',
+          text: `Traces: ${JSON.stringify({
+            traces: response.data,
+            count: response.data.length,
+          })}`,
+        },
+      ],
+    }
+  },
+} as const