From bd46a26f244d70e38fac89e562d34898e8645b13 Mon Sep 17 00:00:00 2001 From: MarioCadenas Date: Fri, 14 Aug 2026 11:50:02 +0200 Subject: [PATCH 1/3] refactor(appkit): let plugins contribute otel span processors Split TelemetryManager into two phases so a single global tracer provider can carry both the OTLP exporter and plugin-contributed processors (e.g. MLflow), instead of each SDK racing to register the global provider. - initialize() registers the meter and logger providers eagerly, because OTel's metrics API has no lazy proxy: an instrument bound against the NoOp meter stays NoOp for the process lifetime. - registerSpanProcessor() lets plugins contribute a span processor during setup(); ignored with a warning after start() since a started provider's processors are immutable in OTel JS 2.x. - start() (called after plugin setup) builds the global NodeTracerProvider with the OTLP processor plus every contributed one, and no-ops when nothing needs tracing. Deferring is safe: ProxyTracer rebinds tracers obtained earlier and no span is emitted during setup. Swaps the @opentelemetry/sdk-node dependency for @opentelemetry/sdk-trace-node, since the tracer provider is now built directly instead of via NodeSDK. Signed-off-by: MarioCadenas --- packages/appkit/package.json | 2 +- packages/appkit/src/core/appkit.ts | 6 + .../core/tests/appkit-as-user-exports.test.ts | 2 + .../appkit/src/telemetry/telemetry-manager.ts | 175 +++++++++++++++--- .../telemetry/tests/telemetry-manager.test.ts | 104 ++++++++++- pnpm-lock.yaml | 6 +- 6 files changed, 264 insertions(+), 31 deletions(-) diff --git a/packages/appkit/package.json b/packages/appkit/package.json index 1f0869209..4992229f9 100644 --- a/packages/appkit/package.json +++ b/packages/appkit/package.json @@ -73,8 +73,8 @@ "@opentelemetry/resources": "2.8.0", "@opentelemetry/sdk-logs": "0.219.0", "@opentelemetry/sdk-metrics": "2.8.0", - "@opentelemetry/sdk-node": "0.219.0", "@opentelemetry/sdk-trace-base": "2.8.0", + "@opentelemetry/sdk-trace-node": "2.8.0", "@opentelemetry/semantic-conventions": "1.38.0", "@types/semver": "7.7.1", "apache-arrow": "21.1.0", diff --git a/packages/appkit/src/core/appkit.ts b/packages/appkit/src/core/appkit.ts index 201acf190..6195836a4 100644 --- a/packages/appkit/src/core/appkit.ts +++ b/packages/appkit/src/core/appkit.ts @@ -222,6 +222,12 @@ export class AppKit { const instance = new AppKit(mergedConfig); await Promise.all(instance.#setupPromises); + + // Build the global tracer provider now that every plugin's setup() has run + // and contributed any span processors. Deferred to here so a single provider + // carries all processors (OTLP + plugin-contributed); see TelemetryManager. + TelemetryManager.start(); + await instance.#context.emitLifecycle("setup:complete"); const handle = instance as unknown as PluginMap; diff --git a/packages/appkit/src/core/tests/appkit-as-user-exports.test.ts b/packages/appkit/src/core/tests/appkit-as-user-exports.test.ts index 7cbadcd32..454a3d664 100644 --- a/packages/appkit/src/core/tests/appkit-as-user-exports.test.ts +++ b/packages/appkit/src/core/tests/appkit-as-user-exports.test.ts @@ -44,6 +44,8 @@ vi.mock("../../telemetry", async () => { ...actual, TelemetryManager: { initialize: vi.fn(), + start: vi.fn(), + registerSpanProcessor: vi.fn(), getProvider: () => ({ getTracer: () => ({ startActiveSpan: vi.fn((_name: string, fn: (span: any) => any) => diff --git a/packages/appkit/src/telemetry/telemetry-manager.ts b/packages/appkit/src/telemetry/telemetry-manager.ts index a2642f2ab..463e6c13e 100644 --- a/packages/appkit/src/telemetry/telemetry-manager.ts +++ b/packages/appkit/src/telemetry/telemetry-manager.ts @@ -1,3 +1,5 @@ +import { metrics } from "@opentelemetry/api"; +import { logs } from "@opentelemetry/api-logs"; import { getNodeAutoInstrumentations } from "@opentelemetry/auto-instrumentations-node"; import { OTLPLogExporter } from "@opentelemetry/exporter-logs-otlp-proto"; import { OTLPMetricExporter } from "@opentelemetry/exporter-metrics-otlp-proto"; @@ -14,9 +16,19 @@ import { type Resource, resourceFromAttributes, } from "@opentelemetry/resources"; -import { BatchLogRecordProcessor } from "@opentelemetry/sdk-logs"; -import { PeriodicExportingMetricReader } from "@opentelemetry/sdk-metrics"; -import { NodeSDK } from "@opentelemetry/sdk-node"; +import { + BatchLogRecordProcessor, + LoggerProvider, +} from "@opentelemetry/sdk-logs"; +import { + MeterProvider, + PeriodicExportingMetricReader, +} from "@opentelemetry/sdk-metrics"; +import { + BatchSpanProcessor, + type SpanProcessor, +} from "@opentelemetry/sdk-trace-base"; +import { NodeTracerProvider } from "@opentelemetry/sdk-trace-node"; import { ATTR_SERVICE_NAME, ATTR_SERVICE_VERSION, @@ -29,12 +41,34 @@ import type { TelemetryConfig } from "./types"; const logger = createLogger("telemetry"); +/** + * Owns the app's OpenTelemetry providers, split into two phases so plugins can + * contribute trace span processors before the tracer provider is built. + * + * - `initialize()` runs at app bootstrap, before plugin setup. It registers the + * meter and logger providers eagerly, because OTel's metrics API has no lazy + * proxy: a counter/histogram bound against the NoOp meter (as every connector + * and the cache do in their constructors) stays NoOp for the process lifetime. + * It does NOT register a tracer provider. + * - `registerSpanProcessor()` is called by plugins during `setup()` to add a + * span processor (e.g. an MLflow exporter) to the not-yet-built tracer. + * - `start()` runs after all plugin `setup()` completes. It builds the single + * global tracer provider with the OTLP processor (if configured) plus every + * contributed processor. Deferring is safe for traces: OTel's ProxyTracer + * rebinds tracers obtained before registration, and no span is emitted during + * setup. + */ export class TelemetryManager { private static readonly DEFAULT_EXPORT_INTERVAL_MS = 10000; private static readonly DEFAULT_FALLBACK_APP_NAME = "databricks-app"; private static instance?: TelemetryManager; - private sdk?: NodeSDK; + private resource?: Resource; + private meterProvider?: MeterProvider; + private loggerProvider?: LoggerProvider; + private tracerProvider?: NodeTracerProvider; + private readonly spanProcessors: SpanProcessor[] = []; + private started = false; private shutdownPromise?: Promise; /** @@ -66,20 +100,49 @@ export class TelemetryManager { instance._initialize(config); } + /** + * Contribute a span processor to the not-yet-built tracer provider. Called by + * plugins during `setup()`. No-op with a warning once `start()` has run, since + * a started provider's processors are immutable in OTel JS 2.x. + */ + static registerSpanProcessor(processor: SpanProcessor): void { + TelemetryManager.getInstance()._registerSpanProcessor(processor); + } + + private _registerSpanProcessor(processor: SpanProcessor): void { + if (this.started) { + logger.warn( + "registerSpanProcessor called after start(); processor ignored. " + + "Contribute span processors during plugin setup().", + ); + return; + } + this.spanProcessors.push(processor); + } + + /** + * Phase 1: register the meter and logger providers eagerly (before plugin + * setup), so metric instruments bound in connector/cache constructors attach + * to real meters. The tracer provider is deferred to `start()`. + * + * When no OTLP endpoint is configured, meter/logger registration is skipped; + * a contributed span processor can still bring up tracing in `start()`. + */ private _initialize(config: Partial): void { - if (this.sdk) return; + if (this.resource) return; + this.resource = this.createResource(config); + // OTLP exporters need an endpoint. Without one there is nothing to export + // metrics/logs to, so skip those providers — but still capture the resource + // and let `start()` bring up a tracer if a plugin contributed a processor. if (!process.env.OTEL_EXPORTER_OTLP_ENDPOINT) { return; } try { - this.sdk = new NodeSDK({ - resource: this.createResource(config), - autoDetectResources: false, - sampler: new AppKitSampler(), - traceExporter: new OTLPTraceExporter({ headers: config.headers }), - metricReaders: [ + this.meterProvider = new MeterProvider({ + resource: this.resource, + readers: [ new PeriodicExportingMetricReader({ exporter: new OTLPMetricExporter({ headers: config.headers }), exportIntervalMillis: @@ -87,21 +150,72 @@ export class TelemetryManager { TelemetryManager.DEFAULT_EXPORT_INTERVAL_MS, }), ], - logRecordProcessors: [ + }); + metrics.setGlobalMeterProvider(this.meterProvider); + + this.loggerProvider = new LoggerProvider({ + resource: this.resource, + processors: [ new BatchLogRecordProcessor( new OTLPLogExporter({ headers: config.headers }), ), ], - instrumentations: this.getDefaultInstrumentations(), }); + logs.setGlobalLoggerProvider(this.loggerProvider); + + // The OTLP trace exporter is the first span processor; contributed + // processors join it in `start()`. + this.spanProcessors.push( + new BatchSpanProcessor( + new OTLPTraceExporter({ headers: config.headers }), + ), + ); - this.sdk.start(); - logger.debug("Initialized successfully"); + this.registerInstrumentations(this.getDefaultInstrumentations()); + logger.debug("Meter/logger providers initialized"); } catch (error) { logger.error("Failed to initialize: %O", error); } } + /** + * Phase 2: build and register the global tracer provider. Called by core + * after every plugin's `setup()` completes, so all contributed span + * processors are known. No-op when nothing needs tracing (no OTLP endpoint + * and no contributed processor), preserving "no telemetry unless configured". + * + * `NodeTracerProvider.register()` installs the async-hooks context manager and + * W3C propagators — the same wiring `NodeSDK.start()` did — so span nesting + * across awaits is preserved. + */ + static start(): void { + TelemetryManager.getInstance()._start(); + } + + private _start(): void { + if (this.started) return; + this.started = true; + + if (this.spanProcessors.length === 0) { + return; + } + + try { + this.tracerProvider = new NodeTracerProvider({ + resource: this.resource, + sampler: new AppKitSampler(), + spanProcessors: this.spanProcessors, + }); + this.tracerProvider.register(); + logger.debug( + "Tracer provider started with %d span processor(s)", + this.spanProcessors.length, + ); + } catch (error) { + logger.error("Failed to start tracer provider: %O", error); + } + } + /** * Register OpenTelemetry instrumentations. * Can be called at any time, but recommended to call in plugin constructor. @@ -159,23 +273,34 @@ export class TelemetryManager { } /** - * Flush and shut down the OpenTelemetry SDK. + * Flush and shut down the tracer, meter, and logger providers. * - * Idempotent: the SDK reference is cleared synchronously and concurrent + * Idempotent: the provider references are cleared synchronously and concurrent * or repeated calls await the same in-flight flush. Awaited by the core * lifecycle manager during graceful shutdown — that manager owns the * process signal handlers, so telemetry no longer registers its own. */ async shutdown(): Promise { - if (this.sdk) { - const sdk = this.sdk; - this.sdk = undefined; + const providers = [ + this.tracerProvider, + this.meterProvider, + this.loggerProvider, + ].filter((p): p is NonNullable => p !== undefined); + + if (providers.length > 0) { + this.tracerProvider = undefined; + this.meterProvider = undefined; + this.loggerProvider = undefined; this.shutdownPromise = (async () => { - try { - await sdk.shutdown(); - } catch (error) { - logger.error("Error shutting down: %O", error); - } + await Promise.all( + providers.map(async (provider) => { + try { + await provider.shutdown(); + } catch (error) { + logger.error("Error shutting down: %O", error); + } + }), + ); })(); } diff --git a/packages/appkit/src/telemetry/tests/telemetry-manager.test.ts b/packages/appkit/src/telemetry/tests/telemetry-manager.test.ts index 11b85d9bf..187276d1d 100644 --- a/packages/appkit/src/telemetry/tests/telemetry-manager.test.ts +++ b/packages/appkit/src/telemetry/tests/telemetry-manager.test.ts @@ -1,3 +1,5 @@ +import { context, metrics, trace } from "@opentelemetry/api"; +import { logs } from "@opentelemetry/api-logs"; import { afterEach, beforeEach, describe, expect, test, vi } from "vitest"; import { TelemetryManager } from "../telemetry-manager"; @@ -54,12 +56,21 @@ describe("TelemetryManager", () => { vi.clearAllMocks(); // @ts-expect-error - accessing private static property for testing TelemetryManager.instance = undefined; - // @ts-expect-error - accessing private static property for testing - TelemetryManager.shutdownRegistered = false; + // OTel's registerGlobal is allowOverride=false: a global registered by one + // test would make the next test's registration a silent no-op. Reset all + // global providers so each test starts clean. + trace.disable(); + metrics.disable(); + logs.disable(); + context.disable(); }); afterEach(() => { process.env = originalEnv; + trace.disable(); + metrics.disable(); + logs.disable(); + context.disable(); }); test("getInstance() should return singleton instance", () => { @@ -89,6 +100,7 @@ describe("TelemetryManager", () => { serviceName: "integration-test", serviceVersion: "1.0.0", }); + TelemetryManager.start(); const telemetryProvider = TelemetryManager.getProvider("test-plugin"); const tracer = telemetryProvider.getTracer(); @@ -186,6 +198,7 @@ describe("TelemetryManager", () => { serviceName: "span-test", serviceVersion: "1.0.0", }); + TelemetryManager.start(); const telemetryProvider = TelemetryManager.getProvider("span-test-plugin"); @@ -211,6 +224,7 @@ describe("TelemetryManager", () => { serviceName: "error-test", serviceVersion: "1.0.0", }); + TelemetryManager.start(); const telemetryProvider = TelemetryManager.getProvider("error-test-plugin"); @@ -224,4 +238,90 @@ describe("TelemetryManager", () => { ).rejects.toThrow("Test error in span"); }); }); + + describe("two-phase init (registerSpanProcessor + start)", () => { + /** Minimal SpanProcessor that records the names of spans it sees start. */ + function recordingProcessor() { + const startedSpans: string[] = []; + return { + startedSpans, + onStart: (span: { name: string }) => { + startedSpans.push(span.name); + }, + onEnd: () => {}, + forceFlush: () => Promise.resolve(), + shutdown: () => Promise.resolve(), + }; + } + + test("routes spans to a contributed processor with no OTLP endpoint", async () => { + process.env.OTEL_EXPORTER_OTLP_ENDPOINT = ""; + const processor = recordingProcessor(); + + TelemetryManager.initialize({ serviceName: "contrib-only" }); + TelemetryManager.registerSpanProcessor(processor as any); + TelemetryManager.start(); + + const tracer = TelemetryManager.getProvider("contrib-plugin").getTracer(); + await tracer.startActiveSpan("contributed.span", {}, async (span) => { + span.end(); + }); + + expect(processor.startedSpans).toContain("contributed.span"); + }); + + test("start() is idempotent", () => { + process.env.OTEL_EXPORTER_OTLP_ENDPOINT = ""; + const processor = recordingProcessor(); + + TelemetryManager.initialize({ serviceName: "idempotent" }); + TelemetryManager.registerSpanProcessor(processor as any); + TelemetryManager.start(); + + expect(() => TelemetryManager.start()).not.toThrow(); + }); + + test("registerSpanProcessor after start() is ignored (not attached)", async () => { + process.env.OTEL_EXPORTER_OTLP_ENDPOINT = ""; + const early = recordingProcessor(); + const late = recordingProcessor(); + + TelemetryManager.initialize({ serviceName: "late-register" }); + TelemetryManager.registerSpanProcessor(early as any); + TelemetryManager.start(); + TelemetryManager.registerSpanProcessor(late as any); + + const tracer = TelemetryManager.getProvider("late-plugin").getTracer(); + await tracer.startActiveSpan("post.start.span", {}, async (span) => { + span.end(); + }); + + expect(early.startedSpans).toContain("post.start.span"); + expect(late.startedSpans).toHaveLength(0); + }); + + test("start() with no OTLP endpoint and no processors is a no-op", () => { + process.env.OTEL_EXPORTER_OTLP_ENDPOINT = ""; + + TelemetryManager.initialize({ serviceName: "no-telemetry" }); + + expect(() => TelemetryManager.start()).not.toThrow(); + }); + + test("metrics obtained after initialize() (before start()) still record", () => { + // Guards the eager-metrics invariant: the meter provider is registered in + // initialize(), not start(), because OTel's metrics API has no lazy proxy. + process.env.OTEL_EXPORTER_OTLP_ENDPOINT = "http://localhost:4318"; + + TelemetryManager.initialize({ serviceName: "eager-metrics" }); + + // Instrument bound BEFORE start() — mirrors connector/cache constructors. + const meter = TelemetryManager.getProvider("metrics-plugin").getMeter(); + const counter = meter.createCounter("eager.counter"); + + TelemetryManager.start(); + + expect(() => counter.add(1, { label: "value" })).not.toThrow(); + }); + }); }); diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml index 86f24f334..db074bf01 100644 --- a/pnpm-lock.yaml +++ b/pnpm-lock.yaml @@ -291,12 +291,12 @@ importers: '@opentelemetry/sdk-metrics': specifier: 2.8.0 version: 2.8.0(@opentelemetry/api@1.9.0) - '@opentelemetry/sdk-node': - specifier: 0.219.0 - version: 0.219.0(@opentelemetry/api@1.9.0) '@opentelemetry/sdk-trace-base': specifier: 2.8.0 version: 2.8.0(@opentelemetry/api@1.9.0) + '@opentelemetry/sdk-trace-node': + specifier: 2.8.0 + version: 2.8.0(@opentelemetry/api@1.9.0) '@opentelemetry/semantic-conventions': specifier: 1.38.0 version: 1.38.0 From 282400b19a0c70fdabd18dd25a9c07829f29eb0d Mon Sep 17 00:00:00 2001 From: MarioCadenas Date: Fri, 21 Aug 2026 15:01:09 +0200 Subject: [PATCH 2/3] fix(agents): route mlflow tracing through AppKit's single otel provider The mlflow-tracing SDK's init() stands up and globally registers its own OpenTelemetry tracer provider. OTel's registerGlobal is allowOverride=false, so when an OTLP endpoint and agent tracing are both active, AppKit's provider and mlflow's race for the global slot and one exporter is silently dropped. Instead of calling init() (which self-registers), the agents plugin now builds mlflow's span processor itself and contributes it to AppKit's single provider via TelemetryManager.registerSpanProcessor() during setup(). mlflow's global config is seeded lazily on first trace (after start()), so init()'s own registration harmlessly loses the already-claimed global slot. A GatedMlflowSpanProcessor keeps the contributed processor inert until config is seeded, so AppKit's own spans created before the first agent turn don't hit mlflow's getConfig() throw. The processor/exporter are deep-imported from mlflow-tracing internals (no public export in 0.1.3), pinned and guarded by a tripwire test. Signed-off-by: MarioCadenas --- packages/appkit/src/plugins/agents/mlflow.ts | 137 ++++++++++++++- .../src/plugins/agents/tests/mlflow.test.ts | 158 +++++++++++++++++- 2 files changed, 288 insertions(+), 7 deletions(-) diff --git a/packages/appkit/src/plugins/agents/mlflow.ts b/packages/appkit/src/plugins/agents/mlflow.ts index cd48b8190..b3a63530d 100644 --- a/packages/appkit/src/plugins/agents/mlflow.ts +++ b/packages/appkit/src/plugins/agents/mlflow.ts @@ -1,12 +1,105 @@ +import type { SpanProcessor } from "@opentelemetry/sdk-trace-base"; + import { createLogger } from "../../logging/logger"; +import { TelemetryManager } from "../../telemetry"; const logger = createLogger("agents"); type MlflowModule = typeof import("mlflow-tracing"); +interface MlflowInitConfig { + trackingUri: string; + experimentId: string; + host?: string; +} + let mlflow: MlflowModule | undefined; let enabled = false; let initStarted = false; +let configured = false; +let initConfig: MlflowInitConfig | undefined; +let gatedProcessor: GatedMlflowSpanProcessor | undefined; + +/** + * Wraps mlflow's OTel `SpanProcessor` so it stays inert until mlflow's global + * config is seeded. mlflow's `onStart` calls its own `getConfig()`, which THROWS + * before `init()` runs — and that throw propagates out of `tracer.startSpan()`, + * so it would break unrelated AppKit spans (HTTP, analytics) created between + * `TelemetryManager.start()` and the first agent turn. We contribute this to + * AppKit's single tracer provider during `setup()`, but only start forwarding + * once `ready()` is called — right after the lazy `init()` in {@link ensureConfigured}. + */ +export class GatedMlflowSpanProcessor implements SpanProcessor { + #inner: SpanProcessor; + #ready = false; + // Spans we forwarded `onStart` for, so `onEnd` stays balanced — mlflow never + // sees an end without a matching start. + #forwarded = new WeakSet(); + + constructor(inner: SpanProcessor) { + this.#inner = inner; + } + + ready(): void { + this.#ready = true; + } + + onStart( + span: Parameters[0], + parentContext: Parameters[1], + ): void { + if (!this.#ready) return; + this.#forwarded.add(span); + this.#inner.onStart(span, parentContext); + } + + onEnd(span: Parameters[0]): void { + if (!this.#forwarded.has(span)) return; + this.#inner.onEnd(span); + } + + forceFlush(): Promise { + return this.#inner.forceFlush(); + } + + shutdown(): Promise { + return this.#inner.shutdown(); + } +} + +/** + * Build mlflow's OTel `SpanProcessor` ourselves rather than letting `init()` + * build and globally register its own tracer provider. This lets AppKit own the + * single global provider (OTLP + this processor), so agent spans reach both + * MLflow and any OTLP endpoint without two SDKs racing for the global slot. + * + * Deep-imports `mlflow-tracing` internals that aren't on its public entrypoint — + * pinned to the exact version in package.json and guarded by a test that fails + * loudly if a version bump renames them. + */ +async function buildMlflowSpanProcessor( + config: MlflowInitConfig, +): Promise { + const { createAuthProvider } = await import("mlflow-tracing/dist/auth"); + const { MlflowClient } = await import("mlflow-tracing"); + const { MlflowSpanExporter, MlflowSpanProcessor } = + await import("mlflow-tracing/dist/exporters/mlflow"); + + const authProvider = createAuthProvider({ + trackingUri: config.trackingUri, + ...(config.host ? { host: config.host } : {}), + }); + const client = new MlflowClient({ + trackingUri: config.trackingUri, + authProvider, + }); + // mlflow builds against a different @opentelemetry/sdk-trace-base major than + // AppKit; the SpanProcessor contract (onStart/onEnd/forceFlush/shutdown) is + // stable across them, so bridge the nominal type mismatch with one cast. + return new MlflowSpanProcessor( + new MlflowSpanExporter(client), + ) as unknown as SpanProcessor; +} /** The bound MLflow experiment id, from the optional `experiment` resource. */ function experimentId(): string | undefined { @@ -30,6 +123,13 @@ function normalizedDatabricksHost(): string | undefined { /** * Initialize MLflow agent tracing once, when an experiment is bound — i.e. the * agents plugin's optional `experiment` resource is set (`MLFLOW_EXPERIMENT_ID`). + * Called from the agents plugin's `setup()`, before `TelemetryManager.start()`. + * + * Rather than let `mlflow.init()` stand up and globally register its own tracer + * provider (which would race AppKit's and drop one exporter), we build mlflow's + * span processor ourselves and contribute it to AppKit's single provider via + * {@link TelemetryManager.registerSpanProcessor}. mlflow's global config is + * seeded lazily in {@link ensureConfigured} on first trace, after `start()`. * * Auth is resolved by the `mlflow-tracing` SDK from the app's own Databricks * credentials — `DATABRICKS_HOST`/`DATABRICKS_TOKEN` or a `~/.databrickscfg` @@ -49,11 +149,14 @@ export async function initAgentTracing(): Promise { try { mlflow = await import("mlflow-tracing"); const host = normalizedDatabricksHost(); - mlflow.init({ + initConfig = { trackingUri: process.env.MLFLOW_TRACKING_URI?.trim() || "databricks", experimentId: id, ...(host ? { host } : {}), - }); + }; + const processor = await buildMlflowSpanProcessor(initConfig); + gatedProcessor = new GatedMlflowSpanProcessor(processor); + TelemetryManager.registerSpanProcessor(gatedProcessor); enabled = true; logger.info("MLflow agent tracing enabled (experiment %s)", id); } catch (err) { @@ -71,6 +174,29 @@ export interface SpanRecorder { const noopRecorder: SpanRecorder = { setOutputs() {} }; +/** + * Seed mlflow's global config on first use — AFTER `TelemetryManager.start()` + * has registered AppKit's provider. `init()` also stands up its own tracer + * provider and tries to register it globally, but that loses to AppKit's + * already-registered provider (non-fatal); we call it only for the config + * side-effect mlflow's span processor requires, then enable forwarding on the + * gated processor. Returns whether tracing is usable. + */ +function ensureConfigured(): boolean { + if (configured) return enabled; + configured = true; + if (!mlflow || !initConfig) return false; + try { + mlflow.init(initConfig); + gatedProcessor?.ready(); + return true; + } catch (err) { + enabled = false; + logger.warn("MLflow agent tracing disabled (init failed): %O", err); + return false; + } +} + /** * Run `fn` inside an MLflow span of `spanType` when tracing is enabled, * otherwise just run it (zero overhead). Spans auto-nest via the SDK's active @@ -86,9 +212,10 @@ async function trace( fn: (span: SpanRecorder) => Promise, ): Promise { if (!enabled || !mlflow) return fn(noopRecorder); - const type = - spanType === "AGENT" ? mlflow.SpanType.AGENT : mlflow.SpanType.TOOL; - return await mlflow.withSpan( + const m = mlflow; + if (!ensureConfigured()) return fn(noopRecorder); + const type = spanType === "AGENT" ? m.SpanType.AGENT : m.SpanType.TOOL; + return await m.withSpan( async (span) => { if (inputs !== undefined) span.setInputs(inputs); let outputsSet = false; diff --git a/packages/appkit/src/plugins/agents/tests/mlflow.test.ts b/packages/appkit/src/plugins/agents/tests/mlflow.test.ts index a43c29cec..8e8645fab 100644 --- a/packages/appkit/src/plugins/agents/tests/mlflow.test.ts +++ b/packages/appkit/src/plugins/agents/tests/mlflow.test.ts @@ -1,9 +1,12 @@ import { afterEach, beforeEach, describe, expect, test, vi } from "vitest"; +import { GatedMlflowSpanProcessor } from "../mlflow"; + /** * The tracing module keeps module-level singleton state (`enabled`, - * `initStarted`) and lazily `import()`s `mlflow-tracing`. Each test resets the - * module registry and re-mocks the SDK so init runs fresh. + * `initStarted`) and lazily `import()`s `mlflow-tracing` (plus two deep-import + * paths for the span processor). Each test resets the module registry and + * re-mocks the SDK so init runs fresh. */ function stubSdk(overrides: Record = {}) { @@ -12,6 +15,7 @@ function stubSdk(overrides: Record = {}) { const span = { setInputs, setOutputs }; const sdk = { init: vi.fn(), + MlflowClient: class {}, SpanType: { AGENT: "AGENT", TOOL: "TOOL" }, withSpan: vi.fn(async (fn: (s: unknown) => unknown) => fn(span)), getCurrentActiveSpan: vi.fn(() => ({ traceId: "tr-active" })), @@ -22,6 +26,24 @@ function stubSdk(overrides: Record = {}) { ...overrides, }; vi.doMock("mlflow-tracing", () => sdk); + // Deep imports used by buildMlflowSpanProcessor — kept as light stubs so + // setup() wires a processor without touching real Databricks auth. + vi.doMock("mlflow-tracing/dist/auth", () => ({ + createAuthProvider: vi.fn(() => ({})), + })); + vi.doMock("mlflow-tracing/dist/exporters/mlflow", () => ({ + MlflowSpanExporter: class {}, + MlflowSpanProcessor: class { + onStart() {} + onEnd() {} + forceFlush() { + return Promise.resolve(); + } + shutdown() { + return Promise.resolve(); + } + }, + })); return { sdk, span, setInputs, setOutputs }; } @@ -33,6 +55,10 @@ describe("agent tracing (mlflow)", () => { afterEach(() => { vi.doUnmock("mlflow-tracing"); + vi.doUnmock("mlflow-tracing/dist/auth"); + vi.doUnmock("mlflow-tracing/dist/exporters/mlflow"); + vi.doUnmock("../../../telemetry"); + vi.restoreAllMocks(); delete process.env.MLFLOW_EXPERIMENT_ID; }); @@ -46,6 +72,48 @@ describe("agent tracing (mlflow)", () => { expect(mod.currentTraceId()).toBeUndefined(); }); + test("no experiment bound: contributes no span processor", async () => { + const { TelemetryManager } = await import("../../../telemetry"); + const spy = vi + .spyOn(TelemetryManager, "registerSpanProcessor") + .mockImplementation(() => {}); + + const mod = await import("../mlflow"); + await mod.initAgentTracing(); + + expect(spy).not.toHaveBeenCalled(); + }); + + test("experiment bound: contributes a span processor during setup", async () => { + process.env.MLFLOW_EXPERIMENT_ID = "exp-123"; + stubSdk(); + const { TelemetryManager } = await import("../../../telemetry"); + const spy = vi + .spyOn(TelemetryManager, "registerSpanProcessor") + .mockImplementation(() => {}); + + const mod = await import("../mlflow"); + await mod.initAgentTracing(); + + expect(spy).toHaveBeenCalledOnce(); + }); + + test("init() is deferred until first trace, not called during setup", async () => { + process.env.MLFLOW_EXPERIMENT_ID = "exp-123"; + const { sdk } = stubSdk(); + vi.doMock("../../../telemetry", () => ({ + TelemetryManager: { registerSpanProcessor: vi.fn() }, + })); + const mod = await import("../mlflow"); + + await mod.initAgentTracing(); + expect(sdk.init).not.toHaveBeenCalled(); // deferred + + await mod.traceAgent("agent", { messages: [] }, async () => {}); + expect(sdk.init).toHaveBeenCalledOnce(); // seeded lazily on first trace + vi.doUnmock("../../../telemetry"); + }); + test("currentTraceId reads the context-active span, not getLastActiveTraceId", async () => { process.env.MLFLOW_EXPERIMENT_ID = "exp-123"; const { sdk } = stubSdk(); @@ -87,4 +155,90 @@ describe("agent tracing (mlflow)", () => { content: "hi", }); }); + + // Tripwire: fails loudly if a mlflow-tracing version bump moves or renames the + // deep-imported internals buildMlflowSpanProcessor() relies on. Runs against + // the REAL package (no mocks); an OSS trackingUri needs no Databricks creds. + test("mlflow-tracing exposes the deep-imported symbols we construct", async () => { + const { createAuthProvider } = await import("mlflow-tracing/dist/auth"); + const { MlflowSpanExporter, MlflowSpanProcessor } = + await import("mlflow-tracing/dist/exporters/mlflow"); + const { MlflowClient } = await import("mlflow-tracing"); + + expect(typeof createAuthProvider).toBe("function"); + expect(typeof MlflowClient).toBe("function"); + expect(typeof MlflowSpanExporter).toBe("function"); + expect(typeof MlflowSpanProcessor).toBe("function"); + + const authProvider = createAuthProvider({ + trackingUri: "http://localhost:5000", + }); + const client = new MlflowClient({ + trackingUri: "http://localhost:5000", + authProvider, + }); + const processor = new MlflowSpanProcessor(new MlflowSpanExporter(client)); + for (const method of ["onStart", "onEnd", "forceFlush", "shutdown"]) { + expect(typeof (processor as any)[method]).toBe("function"); + } + }); +}); + +describe("GatedMlflowSpanProcessor", () => { + function fakeInner() { + return { + onStart: vi.fn(), + onEnd: vi.fn(), + forceFlush: vi.fn(() => Promise.resolve()), + shutdown: vi.fn(() => Promise.resolve()), + }; + } + + test("stays inert before ready(): no forwarding, so onStart can't throw on early spans", () => { + const inner = fakeInner(); + const gated = new GatedMlflowSpanProcessor(inner as any); + const span = { name: "early" }; + + gated.onStart(span as any, {} as any); + gated.onEnd(span as any); + + expect(inner.onStart).not.toHaveBeenCalled(); + expect(inner.onEnd).not.toHaveBeenCalled(); + }); + + test("forwards onStart/onEnd once ready()", () => { + const inner = fakeInner(); + const gated = new GatedMlflowSpanProcessor(inner as any); + gated.ready(); + const span = { name: "s" }; + + gated.onStart(span as any, {} as any); + gated.onEnd(span as any); + + expect(inner.onStart).toHaveBeenCalledOnce(); + expect(inner.onEnd).toHaveBeenCalledOnce(); + }); + + test("onEnd is skipped for spans whose onStart was not forwarded (balanced)", () => { + const inner = fakeInner(); + const gated = new GatedMlflowSpanProcessor(inner as any); + const early = { name: "early" }; + + gated.onStart(early as any, {} as any); // dropped (not ready) + gated.ready(); + gated.onEnd(early as any); // must NOT forward — inner never saw its start + + expect(inner.onEnd).not.toHaveBeenCalled(); + }); + + test("delegates forceFlush and shutdown to the inner processor", async () => { + const inner = fakeInner(); + const gated = new GatedMlflowSpanProcessor(inner as any); + + await gated.forceFlush(); + await gated.shutdown(); + + expect(inner.forceFlush).toHaveBeenCalledOnce(); + expect(inner.shutdown).toHaveBeenCalledOnce(); + }); }); From 4f53ec2f0fc247b6d0a0f8959759ce5308b5e057 Mon Sep 17 00:00:00 2001 From: MarioCadenas Date: Fri, 21 Aug 2026 17:00:40 +0200 Subject: [PATCH 3/3] fix(agents): keep mlflow tracing off the exporter loop and reliable from turn 1 MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The mlflow span processor rides AppKit's shared OTel provider, so it saw every span — including the exporters' own outbound HTTP calls. Each trace upload became a new span to trace and upload: a feedback loop that flooded the experiment and could wedge the process. - Drop parentless CLIENT spans in the gated processor. An outgoing request made outside any agent turn (mlflow/OTLP shipping a trace) is exactly the loop's return edge; spans inside a real request tree keep their parent (or are the incoming SERVER root), so agent turns are untouched. - Seed mlflow's config on the "setup:complete" lifecycle event — after TelemetryManager.start(), before the server serves — instead of lazily on the first trace. The first turn's request-root span is then already forwarded, so that turn assembles into a trace instead of being dropped as a cold-start. Signed-off-by: MarioCadenas --- packages/appkit/src/plugins/agents/agents.ts | 7 ++ packages/appkit/src/plugins/agents/mlflow.ts | 27 ++++++++ .../agents/tests/agents-plugin.test.ts | 2 + .../src/plugins/agents/tests/mlflow.test.ts | 69 ++++++++++++++++--- 4 files changed, 97 insertions(+), 8 deletions(-) diff --git a/packages/appkit/src/plugins/agents/agents.ts b/packages/appkit/src/plugins/agents/agents.ts index 58c811aab..4fe1d62d3 100644 --- a/packages/appkit/src/plugins/agents/agents.ts +++ b/packages/appkit/src/plugins/agents/agents.ts @@ -67,6 +67,7 @@ import { currentTraceId, initAgentTracing, linkTraceToRun, + startAgentTracing, traceAgent, traceTool, } from "./mlflow"; @@ -297,6 +298,12 @@ export class AgentsPlugin extends Plugin implements ToolProvider { async setup() { await initAgentTracing(); + // Seed mlflow's config right after TelemetryManager.start() (before the + // server serves), so the first turn's request-root span is forwarded and + // that turn assembles into a trace — not dropped as a cold-start artifact. + this.context?.onLifecycle("setup:complete", () => { + startAgentTracing(); + }); const { agents, defaultAgentName } = await this.buildAgentRegistry(); this.agents = agents; this.defaultAgentName = defaultAgentName; diff --git a/packages/appkit/src/plugins/agents/mlflow.ts b/packages/appkit/src/plugins/agents/mlflow.ts index b3a63530d..e1b790fac 100644 --- a/packages/appkit/src/plugins/agents/mlflow.ts +++ b/packages/appkit/src/plugins/agents/mlflow.ts @@ -1,3 +1,4 @@ +import { SpanKind } from "@opentelemetry/api"; import type { SpanProcessor } from "@opentelemetry/sdk-trace-base"; import { createLogger } from "../../logging/logger"; @@ -28,6 +29,11 @@ let gatedProcessor: GatedMlflowSpanProcessor | undefined; * `TelemetryManager.start()` and the first agent turn. We contribute this to * AppKit's single tracer provider during `setup()`, but only start forwarding * once `ready()` is called — right after the lazy `init()` in {@link ensureConfigured}. + * + * It also drops the exporters' own outbound spans (parentless CLIENT spans — + * outgoing requests made outside any agent turn, e.g. mlflow/OTLP shipping a + * trace). Forwarding those would loop: each upload is an HTTP call that + * auto-instrumentation turns into a new span to trace and upload. */ export class GatedMlflowSpanProcessor implements SpanProcessor { #inner: SpanProcessor; @@ -49,6 +55,15 @@ export class GatedMlflowSpanProcessor implements SpanProcessor { parentContext: Parameters[1], ): void { if (!this.#ready) return; + // Drop the exporters' own outbound calls. A parentless (root) CLIENT span is + // an outgoing request made outside any agent turn — e.g. mlflow or OTLP + // shipping a trace. Forwarding those would loop: each upload is itself an + // HTTP call that auto-instrumentation turns into a new span to trace and + // upload. Spans inside a real request tree keep their parent (or are the + // incoming SERVER root), so they still flow through. + if (span.kind === SpanKind.CLIENT && !span.parentSpanContext?.spanId) { + return; + } this.#forwarded.add(span); this.#inner.onStart(span, parentContext); } @@ -197,6 +212,18 @@ function ensureConfigured(): boolean { } } +/** + * Seed mlflow's config eagerly, right after `TelemetryManager.start()` — the + * agents plugin wires this to the `"setup:complete"` lifecycle event, before the + * server serves any request. Doing it here means the request's own root span is + * already forwarded when the first turn runs, so that turn assembles into a + * trace instead of being dropped (mlflow roots a trace only at the top-level + * span). Idempotent, and `trace()` still seeds lazily as a fallback. + */ +export function startAgentTracing(): void { + ensureConfigured(); +} + /** * Run `fn` inside an MLflow span of `spanType` when tracing is enabled, * otherwise just run it (zero overhead). Spans auto-nest via the SDK's active diff --git a/packages/appkit/src/plugins/agents/tests/agents-plugin.test.ts b/packages/appkit/src/plugins/agents/tests/agents-plugin.test.ts index e57225dd2..ee438c1d6 100644 --- a/packages/appkit/src/plugins/agents/tests/agents-plugin.test.ts +++ b/packages/appkit/src/plugins/agents/tests/agents-plugin.test.ts @@ -32,6 +32,7 @@ interface FakeContext { getToolProviders(): Array<{ name: string; provider: ToolProvider }>; getPluginNames(): string[]; addRoute(): void; + onLifecycle(): void; executeTool: ( req: unknown, pluginName: string, @@ -48,6 +49,7 @@ function fakeContext( getToolProviders: () => providers, getPluginNames: () => providers.map((p) => p.name), addRoute: vi.fn(), + onLifecycle: vi.fn(), executeTool: vi.fn(async (_req, p, n, args) => ({ plugin: p, tool: n, diff --git a/packages/appkit/src/plugins/agents/tests/mlflow.test.ts b/packages/appkit/src/plugins/agents/tests/mlflow.test.ts index 8e8645fab..04b02dd82 100644 --- a/packages/appkit/src/plugins/agents/tests/mlflow.test.ts +++ b/packages/appkit/src/plugins/agents/tests/mlflow.test.ts @@ -1,3 +1,4 @@ +import { SpanKind } from "@opentelemetry/api"; import { afterEach, beforeEach, describe, expect, test, vi } from "vitest"; import { GatedMlflowSpanProcessor } from "../mlflow"; @@ -114,6 +115,24 @@ describe("agent tracing (mlflow)", () => { vi.doUnmock("../../../telemetry"); }); + test("startAgentTracing seeds config eagerly, before any trace", async () => { + process.env.MLFLOW_EXPERIMENT_ID = "exp-123"; + const { sdk } = stubSdk(); + vi.doMock("../../../telemetry", () => ({ + TelemetryManager: { registerSpanProcessor: vi.fn() }, + })); + const mod = await import("../mlflow"); + + await mod.initAgentTracing(); + expect(sdk.init).not.toHaveBeenCalled(); // not during setup + + // Fired from the "setup:complete" lifecycle hook, after start(), before any + // request — so the first turn's root span is already forwarded. + mod.startAgentTracing(); + expect(sdk.init).toHaveBeenCalledOnce(); + vi.doUnmock("../../../telemetry"); + }); + test("currentTraceId reads the context-active span, not getLastActiveTraceId", async () => { process.env.MLFLOW_EXPERIMENT_ID = "exp-123"; const { sdk } = stubSdk(); @@ -194,10 +213,19 @@ describe("GatedMlflowSpanProcessor", () => { }; } + // Defaults to an in-turn child span (INTERNAL, has a parent) — the case we + // want forwarded. Override kind/parent for the edge cases. + const mkSpan = (over: Record = {}) => ({ + name: "s", + kind: SpanKind.INTERNAL, + parentSpanContext: { spanId: "parent" }, + ...over, + }); + test("stays inert before ready(): no forwarding, so onStart can't throw on early spans", () => { const inner = fakeInner(); const gated = new GatedMlflowSpanProcessor(inner as any); - const span = { name: "early" }; + const span = mkSpan(); gated.onStart(span as any, {} as any); gated.onEnd(span as any); @@ -206,23 +234,48 @@ describe("GatedMlflowSpanProcessor", () => { expect(inner.onEnd).not.toHaveBeenCalled(); }); - test("forwards onStart/onEnd once ready()", () => { + test("once ready(), forwards in-turn spans and the incoming request root", () => { const inner = fakeInner(); const gated = new GatedMlflowSpanProcessor(inner as any); gated.ready(); - const span = { name: "s" }; - gated.onStart(span as any, {} as any); - gated.onEnd(span as any); + const child = mkSpan(); // agent/tool span inside the turn + const requestRoot = mkSpan({ + kind: SpanKind.SERVER, + parentSpanContext: undefined, + }); // incoming /chat span — mlflow roots the trace here + const outgoingChild = mkSpan({ kind: SpanKind.CLIENT }); // LLM call under the agent + + for (const s of [child, requestRoot, outgoingChild]) { + gated.onStart(s as any, {} as any); + gated.onEnd(s as any); + } + + expect(inner.onStart).toHaveBeenCalledTimes(3); + expect(inner.onEnd).toHaveBeenCalledTimes(3); + }); + + test("drops the exporters' own outbound spans (parentless CLIENT) — breaks the loop", () => { + const inner = fakeInner(); + const gated = new GatedMlflowSpanProcessor(inner as any); + gated.ready(); + // An mlflow/OTLP upload: outgoing HTTP with no parent (made outside any turn). + const uploadSpan = mkSpan({ + kind: SpanKind.CLIENT, + parentSpanContext: undefined, + }); - expect(inner.onStart).toHaveBeenCalledOnce(); - expect(inner.onEnd).toHaveBeenCalledOnce(); + gated.onStart(uploadSpan as any, {} as any); + gated.onEnd(uploadSpan as any); + + expect(inner.onStart).not.toHaveBeenCalled(); + expect(inner.onEnd).not.toHaveBeenCalled(); }); test("onEnd is skipped for spans whose onStart was not forwarded (balanced)", () => { const inner = fakeInner(); const gated = new GatedMlflowSpanProcessor(inner as any); - const early = { name: "early" }; + const early = mkSpan(); gated.onStart(early as any, {} as any); // dropped (not ready) gated.ready();