feat(appkit): agent plugin tracing

Signed-off-by: Hubert Zub <hubert.zub@databricks.com>

Author: hubertzub-db

packages/appkit/src/core/appkit.ts

@@ -169,12 +169,35 @@ export class AppKit<TPlugins extends InputPluginMap> {
       client?: WorkspaceClient;
     } = {},
   ): Promise<PluginMap<T>> {
+    const rawPlugins = config.plugins as T;
+
+    // Collect plugin-contributed trace exporter headers before telemetry init
+    const traceExporterHeaders: Record<string, string> = {
+      ...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);

packages/appkit/src/plugins/agent/agent.ts

@@ -20,7 +20,9 @@ import { Plugin, toPlugin } from "../../plugin";
 import type { AgentInterface } from "./agent-interface";
 import { createInvokeHandler } from "./invoke-handler";
 import manifest from "./manifest.json";
+import { setupExperimentTraceLocation } from "./mlflow";
 import { StandardAgent } from "./standard-agent";
+import { instrumentLangChain } from "./tracing";
 import type { IAgentConfig } from "./types";
 
 const logger = createLogger("agent");
@@ -37,6 +39,36 @@ export class AgentPlugin extends Plugin<IAgentConfig> {
 
   static manifest = manifest;
 
+  /**
+   * Called by `_createApp` before TelemetryManager initializes.
+   *
+   * Resolves the MLflow experiment ID and UC table name from config / env.
+   * When `ucTableName` is not explicitly set, derives it deterministically
+   * from catalog + schema so the exporter header is available immediately.
+   * The actual UC location provisioning happens later in `setup()`.
+   *
+   * Returns OTLP headers only when an experiment ID is available.
+   */
+  static appendTraceHeaders(config?: IAgentConfig): Record<string, string> {
+    if (config?.tracing === false) return {};
+
+    const tracing = typeof config?.tracing === "object" ? config.tracing : {};
+
+    const experimentId =
+      tracing.experimentId ?? process.env.MLFLOW_EXPERIMENT_ID;
+    if (!experimentId) return {};
+
+    const ucTableName = tracing.ucTableName ?? process.env.OTEL_UC_TABLE_NAME;
+
+    const headers: Record<string, string> = {
+      "x-mlflow-experiment-id": experimentId,
+    };
+    if (ucTableName) {
+      headers["X-Databricks-UC-Table-Name"] = ucTableName;
+    }
+    return headers;
+  }
+
   protected declare config: IAgentConfig;
 
   private agentImpl: AgentInterface | null = null;
@@ -53,11 +85,63 @@ export class AgentPlugin extends Plugin<IAgentConfig> {
   /** Mutable list of MCP servers (config + added). Only used when building from config. */
   private mcpServersList: DatabricksMCPServer[] = [];
 
+  /**
+   * Tracing is active when `tracing !== false` AND an experiment ID is
+   * available (from config or env).
+   */
+  private isTracingActive(): boolean {
+    if (this.config.tracing === false) {
+      return false;
+    }
+    const tracing =
+      typeof this.config.tracing === "object" ? this.config.tracing : {};
+
+    return Boolean(tracing.experimentId ?? process.env.MLFLOW_EXPERIMENT_ID);
+  }
+
+  /**
+   * Resolve tracing config for the location-setup API.
+   * Catalog and schema are derived from OTEL_UC_TABLE_NAME (catalog.schema.table).
+   */
+  private getTracingConfig() {
+    const tracing =
+      typeof this.config.tracing === "object" ? this.config.tracing : {};
+    const ucTableName = tracing.ucTableName ?? process.env.OTEL_UC_TABLE_NAME;
+    const parts = ucTableName?.split(".") ?? [];
+    return {
+      experimentId:
+        tracing.experimentId ?? process.env.MLFLOW_EXPERIMENT_ID ?? "",
+      ucCatalog: parts.length >= 3 ? parts[0] : undefined,
+      ucSchema: parts.length >= 3 ? parts[1] : undefined,
+      warehouseId: tracing.warehouseId,
+    };
+  }
+
   async setup() {
     this.systemPrompt = this.config.systemPrompt ?? DEFAULT_SYSTEM_PROMPT;
 
+    let tracingActive = this.isTracingActive();
+
+    // Ensure the UC trace location exists; skip instrumentation if it fails
+    if (tracingActive) {
+      const tracingConfig = this.getTracingConfig();
+      const location = await setupExperimentTraceLocation(tracingConfig).catch(
+        (err) => {
+          logger.warn("Trace location setup failed: %O", err);
+          return null;
+        },
+      );
+      if (!location) {
+        tracingActive = false;
+      }
+    }
+
     // If a pre-built agent is provided, use it directly
     if (this.config.agentInstance) {
+      if (tracingActive) {
+        const cbModule = await import("@langchain/core/callbacks/manager");
+        await instrumentLangChain(cbModule);
+      }
       this.agentImpl = this.config.agentInstance;
       logger.info("AgentPlugin initialized with provided agentInstance");
       return;
@@ -74,6 +158,14 @@ export class AgentPlugin extends Plugin<IAgentConfig> {
 
     const { ChatDatabricks } = await import("@databricks/langchainjs");
 
+    // Instrument LangChain callbacks using the *same* @langchain/core copy
+    // that the agent runtime (LangGraph, ChatDatabricks) will use.
+    // Spans flow through the global tracer provider (TelemetryManager).
+    if (tracingActive) {
+      const cbModule = await import("@langchain/core/callbacks/manager");
+      await instrumentLangChain(cbModule);
+    }
+
     this.model = new ChatDatabricks({
       model: modelName,
       useResponsesApi: this.config.useResponsesApi ?? false,

packages/appkit/src/plugins/agent/mlflow.ts

@@ -0,0 +1,174 @@
+/**
+ * MLflow experiment trace location setup.
+ *
+ * Auto-provisions Unity Catalog trace storage and links an MLflow experiment
+ * to it using the Databricks REST API.  This mirrors the Python
+ * `mlflow.set_experiment_trace_location()` behaviour.
+ *
+ * Env vars:
+ *   OTEL_UC_TABLE_NAME — fully-qualified UC table (catalog.schema.table);
+ *                        catalog and schema are derived by splitting on dots.
+ *   MLFLOW_TRACING_SQL_WAREHOUSE_ID — warehouse used to create the storage (optional)
+ */
+
+import { createLogger } from "../../logging/logger";
+
+const logger = createLogger("agent:mlflow");
+
+interface TraceLocationOptions {
+  experimentId: string;
+  /** Derived from OTEL_UC_TABLE_NAME by the caller (first dot-segment). */
+  ucCatalog?: string;
+  /** Derived from OTEL_UC_TABLE_NAME by the caller (second dot-segment). */
+  ucSchema?: string;
+  warehouseId?: string;
+}
+
+const MLFLOW_TRACE_LOCATION_TABLE_NAME = "mlflow_experiment_trace_otel_spans";
+
+/**
+ * Link an MLflow experiment to an existing UC trace location.
+ * Returns the fully-qualified table name on success, `null` otherwise.
+ */
+async function linkExperimentToLocation(
+  client: { apiClient: { request: (opts: unknown) => Promise<unknown> } },
+  experimentId: string,
+  catalogName: string,
+  schemaName: string,
+): Promise<string | null> {
+  const tableName = `${catalogName}.${schemaName}.${MLFLOW_TRACE_LOCATION_TABLE_NAME}`;
+
+  try {
+    await client.apiClient.request({
+      path: `/api/4.0/mlflow/traces/${experimentId}/link-location`,
+      method: "POST",
+      headers: new Headers({ "Content-Type": "application/json" }),
+      payload: {
+        experiment_id: experimentId,
+        uc_schema: {
+          catalog_name: catalogName,
+          schema_name: schemaName,
+        },
+      },
+      raw: false,
+    });
+
+    logger.info("Experiment linked to UC trace location: %s", tableName);
+    return tableName;
+  } catch (error: unknown) {
+    const code =
+      (error as { error_code?: string }).error_code ??
+      (error as Error).name ??
+      "UNKNOWN";
+    logger.warn(
+      "Could not link experiment %s to %s (%s)",
+      experimentId,
+      tableName,
+      code,
+    );
+    return null;
+  }
+}
+
+/**
+ * Provision a UC trace storage location and link the experiment to it.
+ *
+ * If `warehouseId` is not provided, the function attempts to link directly
+ * (works when the UC table already exists).
+ *
+ * Returns the fully-qualified UC table name on success, `null` otherwise.
+ */
+export async function setupExperimentTraceLocation(
+  opts: TraceLocationOptions,
+): Promise<string | null> {
+  const catalogName = opts.ucCatalog;
+  const schemaName = opts.ucSchema;
+
+  if (!catalogName || !schemaName) {
+    logger.debug(
+      "Skipping trace location setup — catalog/schema not available (set OTEL_UC_TABLE_NAME as catalog.schema.table)",
+    );
+    return null;
+  }
+
+  const warehouseId =
+    opts.warehouseId ?? process.env.MLFLOW_TRACING_SQL_WAREHOUSE_ID;
+
+  let client: {
+    apiClient: { request: (opts: unknown) => Promise<unknown> };
+    config: { ensureResolved: () => Promise<void> };
+  };
+  try {
+    const { WorkspaceClient } = await import("@databricks/sdk-experimental");
+    client = new WorkspaceClient({}) as typeof client;
+    await client.config.ensureResolved();
+  } catch (err) {
+    logger.warn(
+      "Cannot set up trace location — Databricks auth unavailable: %O",
+      err,
+    );
+    return null;
+  }
+
+  if (!warehouseId) {
+    const result = await linkExperimentToLocation(
+      client,
+      opts.experimentId,
+      catalogName,
+      schemaName,
+    );
+    if (!result) {
+      logger.warn(
+        "Trace destination does not exist and cannot be created — " +
+          "set MLFLOW_TRACING_SQL_WAREHOUSE_ID to auto-create it",
+      );
+    }
+    return result;
+  }
+
+  try {
+    logger.debug(
+      "Creating UC trace location: %s.%s (warehouse=%s)",
+      catalogName,
+      schemaName,
+      warehouseId,
+    );
+
+    await client.apiClient.request({
+      path: "/api/4.0/mlflow/traces/location",
+      method: "POST",
+      headers: new Headers({ "Content-Type": "application/json" }),
+      payload: {
+        uc_schema: {
+          catalog_name: catalogName,
+          schema_name: schemaName,
+        },
+        sql_warehouse_id: warehouseId,
+      },
+      raw: false,
+    });
+
+    return linkExperimentToLocation(
+      client,
+      opts.experimentId,
+      catalogName,
+      schemaName,
+    );
+  } catch (error: unknown) {
+    // 409 = location already exists — just link
+    if (
+      error instanceof Error &&
+      (error.message?.includes("409") ||
+        error.message?.includes("ALREADY_EXISTS"))
+    ) {
+      return linkExperimentToLocation(
+        client,
+        opts.experimentId,
+        catalogName,
+        schemaName,
+      );
+    }
+    logger.warn("Failed to create UC trace location: %O", error);
+    return null;
+  }
+}

packages/appkit/src/plugins/agent/tracing.ts

@@ -0,0 +1,47 @@
+/**
+ * LangChain tracing integration for the agent plugin.
+ *
+ * Instruments LangChain callbacks via @arizeai/openinference-instrumentation-langchain
+ * so that agent spans are emitted through the global tracer provider
+ * (set up by AppKit's TelemetryManager).
+ *
+ * MLflow-specific headers (experiment ID, UC table name) are added
+ * to the trace exporter by TelemetryManager — see buildTraceExporterHeaders().
+ */
+
+import { createLogger } from "../../logging/logger";
+
+const logger = createLogger("agent:tracing");
+
+/**
+ * Instrument LangChain callbacks via the Arize/OpenInference library.
+ *
+ * The instrumentation creates spans using the **global** tracer provider
+ * (registered by TelemetryManager's NodeSDK). No separate provider is needed.
+ *
+ * IMPORTANT: `callbackManagerModule` must be the module object that the
+ * agent runtime uses (imported from agent.ts context), NOT a separate
+ * import inside this file. pnpm strict isolation can resolve different
+ * physical copies of @langchain/core for different packages in the
+ * dependency tree, and patching the wrong copy has no effect.
+ */
+export async function instrumentLangChain(
+  callbackManagerModule?: typeof import("@langchain/core/callbacks/manager"),
+): Promise<void> {
+  try {
+    const { LangChainInstrumentation } = await import(
+      "@arizeai/openinference-instrumentation-langchain"
+    );
+
+    const cbModule =
+      callbackManagerModule ??
+      (await import("@langchain/core/callbacks/manager"));
+
+    const inst = new LangChainInstrumentation();
+    inst.manuallyInstrument(cbModule);
+
+    logger.debug("LangChain callbacks instrumented (global tracer provider)");
+  } catch (err) {
+    logger.error("Failed to instrument LangChain callbacks: %O", err);
+  }
+}