feat: add trace header propagation for distributed tracing

apps/mcp/src/index.ts

@@ -103,6 +103,22 @@ const handler = {
         : undefined;
       const withOTel = withOtelParam === '1' || withOtelParam === 'true';
 
+      // Extract trace headers for propagation to downstream services
+      const traceHeaders: Record<string, string> = {};
+      const traceHeaderNames = [
+        'traceparent',
+        'tracestate',
+        'x-trace-id',
+        'x-request-id',
+      ];
+
+      for (const headerName of traceHeaderNames) {
+        const headerValue = request.headers.get(headerName);
+        if (headerValue) {
+          traceHeaders[headerName] = headerValue;
+        }
+      }
+
       const props: ServerProps = {
         tokenKey: await sha256(`${accessToken}:${orgId}`),
         accessToken,
@@ -111,6 +127,8 @@ const handler = {
           ? (maxCells as number)
           : undefined,
         withOTel,
+        traceHeaders:
+          Object.keys(traceHeaders).length > 0 ? traceHeaders : undefined,
       };
 
       ctx.props = props;

apps/mcp/src/mcp.ts

@@ -36,6 +36,7 @@ export class AxiomMCP extends McpAgent<
       formatOptions: this.props.maxCells
         ? { maxCells: this.props.maxCells }
         : undefined,
+      traceHeaders: this.props.traceHeaders,
     });
 
     logger.info('Server initialized');
@@ -60,7 +61,8 @@ export class AxiomMCP extends McpAgent<
       const internalClient = new Client(
         this.env.ATLAS_INTERNAL_URL,
         this.props.accessToken,
-        this.props.orgId
+        this.props.orgId,
+        this.props.traceHeaders
       );
       const ret: Integrations = await getIntegrations(internalClient);
       integrations = [...new Set(ret.map((i) => i.kind))];

apps/mcp/src/types.ts

@@ -8,4 +8,5 @@ export type ServerProps = {
   // Parsed from `max-age` and `with-otel`
   maxCells?: number;
   withOTel?: boolean;
+  traceHeaders?: Record<string, string>;
 };

packages/mcp/metadata.ts

@@ -1,6 +1,6 @@
 /**
  * Auto-generated MCP metadata
- * Generated at: 2025-10-07T14:50:49.840Z
+ * Generated at: 2025-10-07T08:25:05.303Z
  * DO NOT EDIT - This file is automatically generated by scripts/extract-tools.ts
  */
 

packages/mcp/src/axiom/client.ts

@@ -17,6 +17,7 @@ export type ApiRequest = {
   body?: unknown;
   baseUrl: string;
   orgId: string;
+  traceHeaders?: Record<string, string>;
 };
 
 // MCP server telemetry configuration - similar to axiom.SetUserAgent() in Go SDK
@@ -33,7 +34,7 @@ export async function apiFetch<T>(
   areq: ApiRequest,
   schema: z.ZodSchema<T>
 ): Promise<T> {
-  const { token, method, path, body, baseUrl, orgId } = areq;
+  const { token, method, path, body, baseUrl, orgId, traceHeaders } = areq;
   const headers: Record<string, string> = {
     Authorization: `Bearer ${token}`,
     'Content-Type': 'application/json',
@@ -46,6 +47,13 @@ export async function apiFetch<T>(
     headers['x-axiom-org-id'] = orgId;
   }
 
+  // Propagate trace headers for distributed tracing
+  if (traceHeaders) {
+    for (const [key, value] of Object.entries(traceHeaders)) {
+      headers[key] = value;
+    }
+  }
+
   const options: RequestInit = {
     method,
     headers,
@@ -89,11 +97,18 @@ export class Client {
   private baseUrl: string;
   private accessToken: string;
   private orgId: string;
-
-  constructor(baseUrl: string, accessToken: string, orgId: string) {
+  private traceHeaders?: Record<string, string>;
+
+  constructor(
+    baseUrl: string,
+    accessToken: string,
+    orgId: string,
+    traceHeaders?: Record<string, string>
+  ) {
     this.baseUrl = baseUrl;
     this.accessToken = accessToken;
     this.orgId = orgId;
+    this.traceHeaders = traceHeaders;
   }
 
   async fetch<T>(
@@ -110,6 +125,7 @@ export class Client {
         body,
         baseUrl: this.baseUrl,
         orgId: this.orgId,
+        traceHeaders: this.traceHeaders,
       },
       schema
     );

packages/mcp/src/index.ts

@@ -45,18 +45,21 @@ export interface AxiomMcpConfig {
   orgId: string;
   enableOtel?: boolean;
   formatOptions?: FormatterOptions;
+  traceHeaders?: Record<string, string>;
 }
 
 export function registerAxiomMcpTools(config: AxiomMcpConfig) {
   const publicClient = new Client(
     config.apiUrl,
     config.accessToken,
-    config.orgId
+    config.orgId,
+    config.traceHeaders
   );
   const internalClient = new Client(
     config.internalUrl,
     config.accessToken,
-    config.orgId
+    config.orgId,
+    config.traceHeaders
   );
 
   const context = {