diff --git a/packages/appkit/src/core/appkit.ts b/packages/appkit/src/core/appkit.ts index a2cba994..e334ac6d 100644 --- a/packages/appkit/src/core/appkit.ts +++ b/packages/appkit/src/core/appkit.ts @@ -169,12 +169,35 @@ export class AppKit { client?: WorkspaceClient; } = {}, ): Promise> { + const rawPlugins = config.plugins as T; + + // Collect plugin-contributed trace exporter headers before telemetry init + const traceExporterHeaders: Record = { + ...config?.telemetry?.traceExporterHeaders, + }; + for (const entry of rawPlugins) { + if (typeof entry.plugin.appendTraceHeaders === "function") { + Object.assign( + traceExporterHeaders, + entry.plugin.appendTraceHeaders( + entry.config as Parameters< + typeof entry.plugin.appendTraceHeaders + >[0], + ), + ); + } + } + // Initialize core services - TelemetryManager.initialize(config?.telemetry); + await TelemetryManager.initialize({ + ...config?.telemetry, + traceExporterHeaders: + Object.keys(traceExporterHeaders).length > 0 + ? traceExporterHeaders + : undefined, + }); await CacheManager.getInstance(config?.cache); - const rawPlugins = config.plugins as T; - // Collect manifest resources via registry const registry = new ResourceRegistry(); registry.collectResources(rawPlugins); diff --git a/packages/appkit/src/plugins/server/index.ts b/packages/appkit/src/plugins/server/index.ts index 22501ce0..464dffe3 100644 --- a/packages/appkit/src/plugins/server/index.ts +++ b/packages/appkit/src/plugins/server/index.ts @@ -58,10 +58,12 @@ export class ServerPlugin extends Plugin { this.serverApplication = express(); this.server = null; this.serverExtensions = []; - this.telemetry.registerInstrumentations([ - instrumentations.http, - instrumentations.express, - ]); + if (config.enableDefaultTelemetry !== false) { + this.telemetry.registerInstrumentations([ + instrumentations.http(), + instrumentations.express(), + ]); + } } /** Setup the server plugin. */ diff --git a/packages/appkit/src/plugins/server/tests/server.test.ts b/packages/appkit/src/plugins/server/tests/server.test.ts index 4015917d..b4da3be3 100644 --- a/packages/appkit/src/plugins/server/tests/server.test.ts +++ b/packages/appkit/src/plugins/server/tests/server.test.ts @@ -75,8 +75,8 @@ vi.mock("../../../telemetry", () => ({ }), }, instrumentations: { - http: {}, - express: {}, + http: () => ({}), + express: () => ({}), }, })); diff --git a/packages/appkit/src/plugins/server/types.ts b/packages/appkit/src/plugins/server/types.ts index e187cacc..381d3ac2 100644 --- a/packages/appkit/src/plugins/server/types.ts +++ b/packages/appkit/src/plugins/server/types.ts @@ -7,4 +7,10 @@ export interface ServerConfig extends BasePluginConfig { staticPath?: string; autoStart?: boolean; host?: string; + /** + * Register HTTP and Express OpenTelemetry instrumentations. + * When false, no HTTP/Express spans are created regardless of other + * telemetry settings. Default: true. + */ + enableDefaultTelemetry?: boolean; } diff --git a/packages/appkit/src/telemetry/instrumentations.ts b/packages/appkit/src/telemetry/instrumentations.ts index 581bdc0e..7b70c38c 100644 --- a/packages/appkit/src/telemetry/instrumentations.ts +++ b/packages/appkit/src/telemetry/instrumentations.ts @@ -4,56 +4,52 @@ import { HttpInstrumentation } from "@opentelemetry/instrumentation-http"; import { shouldIgnoreRequest } from "../utils/path-exclusions"; /** - * Registry of pre-configured instrumentations for common use cases. - * These can be selectively registered by plugins that need them. - * - * While instrumentations are generally safe to re-register, - * the recommended approach is to register them once in a corresponding plugin constructor. + * Factory functions that create pre-configured instrumentations on demand. + * Lazy creation avoids side-effects at import time (e.g. module patching + * triggered by instrumentation constructors). */ -export const instrumentations: Record = { - http: new HttpInstrumentation({ - // Filter out requests before creating spans - this is the most efficient approach - ignoreIncomingRequestHook: shouldIgnoreRequest, +export const instrumentations = { + http: (): Instrumentation => + new HttpInstrumentation({ + ignoreIncomingRequestHook: shouldIgnoreRequest, - applyCustomAttributesOnSpan(span: any, request: any) { - let spanName: string | null = null; + applyCustomAttributesOnSpan(span: any, request: any) { + let spanName: string | null = null; - if (request.route) { - const baseUrl = request.baseUrl || ""; - const url = request.url?.split("?")[0] || ""; - const fullPath = baseUrl + url; - if (fullPath) { - spanName = `${request.method} ${fullPath}`; + if (request.route) { + const baseUrl = request.baseUrl || ""; + const url = request.url?.split("?")[0] || ""; + const fullPath = baseUrl + url; + if (fullPath) { + spanName = `${request.method} ${fullPath}`; + } + } else if (request.url) { + // No Express route (e.g., static assets) - use the raw URL path + // Remove query string for cleaner trace names + const path = request.url.split("?")[0]; + spanName = `${request.method} ${path}`; } - } else if (request.url) { - // No Express route (e.g., static assets) - use the raw URL path - // Remove query string for cleaner trace names - const path = request.url.split("?")[0]; - spanName = `${request.method} ${path}`; - } - if (spanName) { - span.updateName(spanName); - } - }, - }), - express: new ExpressInstrumentation({ - requestHook: (span: any, info: any) => { - const req = info.request; - - // Only update span name for route handlers (layerType: request_handler) - // This ensures we're not renaming middleware spans - if (info.layerType === "request_handler" && req.route) { - // Combine baseUrl with url to get full path with actual parameter values - // e.g., baseUrl="/api/analytics" + url="/query/spend_data" = "/api/analytics/query/spend_data" - const baseUrl = req.baseUrl || ""; - const url = req.url?.split("?")[0] || ""; - const fullPath = baseUrl + url; - if (fullPath) { - const spanName = `${req.method} ${fullPath}`; + if (spanName) { span.updateName(spanName); } - } - }, - }), + }, + }), + + express: (): Instrumentation => + new ExpressInstrumentation({ + requestHook: (span: any, info: any) => { + const req = info.request; + + if (info.layerType === "request_handler" && req.route) { + const baseUrl = req.baseUrl || ""; + const url = req.url?.split("?")[0] || ""; + const fullPath = baseUrl + url; + if (fullPath) { + const spanName = `${req.method} ${fullPath}`; + span.updateName(spanName); + } + } + }, + }), }; diff --git a/packages/appkit/src/telemetry/telemetry-manager.ts b/packages/appkit/src/telemetry/telemetry-manager.ts index 6660b6b2..9655558a 100644 --- a/packages/appkit/src/telemetry/telemetry-manager.ts +++ b/packages/appkit/src/telemetry/telemetry-manager.ts @@ -60,12 +60,14 @@ export class TelemetryManager { return TelemetryManager.instance; } - static initialize(config: Partial = {}): void { + static async initialize( + config: Partial = {}, + ): Promise { const instance = TelemetryManager.getInstance(); - instance._initialize(config); + await instance._initialize(config); } - private _initialize(config: Partial): void { + private async _initialize(config: Partial): Promise { if (this.sdk) return; if (!process.env.OTEL_EXPORTER_OTLP_ENDPOINT) { @@ -73,11 +75,16 @@ export class TelemetryManager { } try { + const traceHeaders = await this.buildTraceExporterHeaders( + config.headers, + config.traceExporterHeaders, + ); + this.sdk = new NodeSDK({ resource: this.createResource(config), autoDetectResources: false, sampler: new AppKitSampler(), - traceExporter: new OTLPTraceExporter({ headers: config.headers }), + traceExporter: new OTLPTraceExporter({ headers: traceHeaders }), metricReaders: [ new PeriodicExportingMetricReader({ exporter: new OTLPMetricExporter({ headers: config.headers }), @@ -91,7 +98,8 @@ export class TelemetryManager { new OTLPLogExporter({ headers: config.headers }), ), ], - instrumentations: this.getDefaultInstrumentations(), + instrumentations: + config.instrumentations ?? this.getDefaultInstrumentations(), }); this.sdk.start(); @@ -130,6 +138,44 @@ export class TelemetryManager { return initialResource.merge(detectedResource); } + /** + * Builds headers for the trace exporter by merging (in priority order): + * 1. Base `headers` from TelemetryConfig + * 2. Databricks auth (auto-resolved when DATABRICKS_HOST is set) + * 3. Plugin-contributed `traceExporterHeaders` (e.g. MLflow experiment ID) + */ + private async buildTraceExporterHeaders( + configHeaders?: Record, + pluginHeaders?: Record, + ): Promise> { + const headers: Record = { ...configHeaders }; + + if (process.env.DATABRICKS_HOST && !headers.authorization) { + try { + const { WorkspaceClient } = await import( + "@databricks/sdk-experimental" + ); + const client = new WorkspaceClient({}); + const authHeaders = new Headers(); + await client.config.authenticate(authHeaders); + authHeaders.forEach((value, key) => { + headers[key] = value; + }); + } catch (err) { + logger.warn( + "Could not obtain Databricks auth for trace exporter: %O", + err, + ); + } + } + + if (pluginHeaders) { + Object.assign(headers, pluginHeaders); + } + + return headers; + } + private getDefaultInstrumentations(): Instrumentation[] { return [ ...getNodeAutoInstrumentations({ diff --git a/packages/appkit/src/telemetry/tests/telemetry-manager.test.ts b/packages/appkit/src/telemetry/tests/telemetry-manager.test.ts index 11b85d9b..f28ddaca 100644 --- a/packages/appkit/src/telemetry/tests/telemetry-manager.test.ts +++ b/packages/appkit/src/telemetry/tests/telemetry-manager.test.ts @@ -51,6 +51,8 @@ describe("TelemetryManager", () => { beforeEach(() => { originalEnv = { ...process.env }; + // Prevent TelemetryManager from attempting Databricks auth during tests + delete process.env.DATABRICKS_HOST; vi.clearAllMocks(); // @ts-expect-error - accessing private static property for testing TelemetryManager.instance = undefined; @@ -75,7 +77,7 @@ describe("TelemetryManager", () => { const { detectResources } = await import("@opentelemetry/resources"); vi.clearAllMocks(); - TelemetryManager.initialize({ + await TelemetryManager.initialize({ serviceName: "test-service-config", }); @@ -85,7 +87,7 @@ describe("TelemetryManager", () => { test("should initialize providers and create telemetry instances", async () => { process.env.OTEL_EXPORTER_OTLP_ENDPOINT = "http://localhost:4318"; - TelemetryManager.initialize({ + await TelemetryManager.initialize({ serviceName: "integration-test", serviceVersion: "1.0.0", }); @@ -118,7 +120,7 @@ describe("TelemetryManager", () => { test("should support disabled telemetry config", async () => { process.env.OTEL_EXPORTER_OTLP_ENDPOINT = ""; - TelemetryManager.initialize({ + await TelemetryManager.initialize({ serviceName: "disabled-test", serviceVersion: "1.0.0", }); @@ -161,7 +163,7 @@ describe("TelemetryManager", () => { ); vi.clearAllMocks(); - TelemetryManager.initialize({ + await TelemetryManager.initialize({ headers: { Authorization: "Bearer token", "Custom-Header": "value", @@ -182,7 +184,7 @@ describe("TelemetryManager", () => { test("should create and execute spans with real tracer", async () => { process.env.OTEL_EXPORTER_OTLP_ENDPOINT = "http://localhost:4318"; - TelemetryManager.initialize({ + await TelemetryManager.initialize({ serviceName: "span-test", serviceVersion: "1.0.0", }); @@ -207,7 +209,7 @@ describe("TelemetryManager", () => { test("should handle span errors", async () => { process.env.OTEL_EXPORTER_OTLP_ENDPOINT = "http://localhost:4318"; - TelemetryManager.initialize({ + await TelemetryManager.initialize({ serviceName: "error-test", serviceVersion: "1.0.0", }); diff --git a/packages/appkit/src/telemetry/types.ts b/packages/appkit/src/telemetry/types.ts index abd73e4c..7983c528 100644 --- a/packages/appkit/src/telemetry/types.ts +++ b/packages/appkit/src/telemetry/types.ts @@ -9,6 +9,12 @@ export interface TelemetryConfig { instrumentations?: Instrumentation[]; exportIntervalMs?: number; headers?: Record; + /** + * Additional headers to include on the OTLP **trace** exporter only. + * Merged on top of `headers` and any auto-resolved Databricks auth. + * Typically populated automatically by plugins via `appendTraceHeaders`. + */ + traceExporterHeaders?: Record; } /** diff --git a/packages/shared/src/plugin.ts b/packages/shared/src/plugin.ts index 761bdce6..604148d8 100644 --- a/packages/shared/src/plugin.ts +++ b/packages/shared/src/plugin.ts @@ -77,6 +77,11 @@ export type PluginConstructor< * Use this when resource requirements depend on plugin configuration. */ getResourceRequirements?(config: C): ResourceRequirement[]; + /** + * Returns extra headers to append to the OTLP trace exporter. + * Called by `_createApp` before TelemetryManager initializes. + */ + appendTraceHeaders?(config: C): Record; }; /**