From 2bf3bae9b90d35136d2be4e996a79097884444a1 Mon Sep 17 00:00:00 2001 From: Rambatino Date: Tue, 7 Oct 2025 14:45:03 +0100 Subject: [PATCH 1/2] feat: add trace header propagation for distributed tracing - Extract trace headers (traceparent, tracestate, x-trace-id, x-request-id) from incoming requests - Add traceHeaders field to ServerProps type - Pass trace headers through to MCP server context - Leverage existing @microlabs/otel-cf-workers instrumentation for automatic propagation - Maintain compatibility with Axiom trace publishing configuration --- apps/mcp/src/index.ts | 18 ++++++++++++++++++ apps/mcp/src/mcp.ts | 4 +++- apps/mcp/src/types.ts | 1 + package-lock.json | 4 ++-- 4 files changed, 24 insertions(+), 3 deletions(-) diff --git a/apps/mcp/src/index.ts b/apps/mcp/src/index.ts index d104587..64b47f0 100644 --- a/apps/mcp/src/index.ts +++ b/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 = {}; + 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; diff --git a/apps/mcp/src/mcp.ts b/apps/mcp/src/mcp.ts index 1aedf4c..1d134ae 100644 --- a/apps/mcp/src/mcp.ts +++ b/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))]; diff --git a/apps/mcp/src/types.ts b/apps/mcp/src/types.ts index b9433cf..2d688b6 100644 --- a/apps/mcp/src/types.ts +++ b/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; }; diff --git a/package-lock.json b/package-lock.json index 8ecf8b3..5b5f808 100644 --- a/package-lock.json +++ b/package-lock.json @@ -1,12 +1,12 @@ { "name": "axiom-mcp-monorepo", - "version": "1.0.0", + "version": "1.0.1", "lockfileVersion": 3, "requires": true, "packages": { "": { "name": "axiom-mcp-monorepo", - "version": "0.0.1", + "version": "1.0.1", "license": "MIT", "workspaces": [ "packages/*", From 26be126950afd8dd03f4167577cff6ea71a64742 Mon Sep 17 00:00:00 2001 From: Rambatino Date: Wed, 8 Oct 2025 12:09:31 +0100 Subject: [PATCH 2/2] fix: more headers --- packages/mcp/metadata.ts | 2 +- packages/mcp/src/axiom/client.ts | 22 +++++++++++++++++++--- packages/mcp/src/index.ts | 7 +++++-- 3 files changed, 25 insertions(+), 6 deletions(-) diff --git a/packages/mcp/metadata.ts b/packages/mcp/metadata.ts index 41b0d2f..b3e7ef4 100644 --- a/packages/mcp/metadata.ts +++ b/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 */ diff --git a/packages/mcp/src/axiom/client.ts b/packages/mcp/src/axiom/client.ts index 5ef1ce0..24a5e7a 100644 --- a/packages/mcp/src/axiom/client.ts +++ b/packages/mcp/src/axiom/client.ts @@ -17,6 +17,7 @@ export type ApiRequest = { body?: unknown; baseUrl: string; orgId: string; + traceHeaders?: Record; }; // MCP server telemetry configuration - similar to axiom.SetUserAgent() in Go SDK @@ -33,7 +34,7 @@ export async function apiFetch( areq: ApiRequest, schema: z.ZodSchema ): Promise { - const { token, method, path, body, baseUrl, orgId } = areq; + const { token, method, path, body, baseUrl, orgId, traceHeaders } = areq; const headers: Record = { Authorization: `Bearer ${token}`, 'Content-Type': 'application/json', @@ -46,6 +47,13 @@ export async function apiFetch( 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; + + constructor( + baseUrl: string, + accessToken: string, + orgId: string, + traceHeaders?: Record + ) { this.baseUrl = baseUrl; this.accessToken = accessToken; this.orgId = orgId; + this.traceHeaders = traceHeaders; } async fetch( @@ -110,6 +125,7 @@ export class Client { body, baseUrl: this.baseUrl, orgId: this.orgId, + traceHeaders: this.traceHeaders, }, schema ); diff --git a/packages/mcp/src/index.ts b/packages/mcp/src/index.ts index d1fd21e..cbe88d0 100644 --- a/packages/mcp/src/index.ts +++ b/packages/mcp/src/index.ts @@ -45,18 +45,21 @@ export interface AxiomMcpConfig { orgId: string; enableOtel?: boolean; formatOptions?: FormatterOptions; + traceHeaders?: Record; } 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 = {