@dbx-tools/appkit-mastra 0.9.64 → 0.9.67

package/package.json

@@ -19,8 +19,7 @@
     "projen": "projen"
   },
   "devDependencies": {
-    "@opentelemetry/sdk-trace-base": "2.8.0",
-    "@opentelemetry/sdk-trace-node": "2.8.0",
+    "@mlflow/core": "0.4.0",
     "@types/express": "^5.0.5",
     "@types/node": "^24.6.0",
     "@types/pg": "^8",
@@ -29,23 +28,24 @@
     "vitest": "3.2.4"
   },
   "peerDependencies": {
+    "@mlflow/core": "0.4.0",
     "skills": "^1"
   },
   "dependencies": {
     "@databricks/appkit": "0.81.0",
-    "@dbx-tools/appkit": "^0.9.64",
-    "@dbx-tools/core": "^0.9.64",
-    "@dbx-tools/databricks": "^0.9.64",
-    "@dbx-tools/fs": "^0.9.64",
-    "@dbx-tools/genie": "^0.9.64",
-    "@dbx-tools/model": "^0.9.64",
-    "@dbx-tools/path": "^0.9.64",
-    "@dbx-tools/postgres": "^0.9.64",
-    "@dbx-tools/shared-core": "^0.9.64",
-    "@dbx-tools/shared-fs": "^0.9.64",
-    "@dbx-tools/shared-genie": "^0.9.64",
-    "@dbx-tools/shared-mastra": "^0.9.64",
-    "@dbx-tools/shared-model": "^0.9.64",
+    "@dbx-tools/appkit": "^0.9.67",
+    "@dbx-tools/core": "^0.9.67",
+    "@dbx-tools/databricks": "^0.9.67",
+    "@dbx-tools/fs": "^0.9.67",
+    "@dbx-tools/genie": "^0.9.67",
+    "@dbx-tools/model": "^0.9.67",
+    "@dbx-tools/path": "^0.9.67",
+    "@dbx-tools/postgres": "^0.9.67",
+    "@dbx-tools/shared-core": "^0.9.67",
+    "@dbx-tools/shared-fs": "^0.9.67",
+    "@dbx-tools/shared-genie": "^0.9.67",
+    "@dbx-tools/shared-mastra": "^0.9.67",
+    "@dbx-tools/shared-model": "^0.9.67",
     "@mastra/ai-sdk": "1.10.5",
     "@mastra/core": "1.71.0",
     "@mastra/express": "1.5.15",
@@ -58,6 +58,8 @@
     "@modelcontextprotocol/sdk": "^1.29.0",
     "@opentelemetry/api": "^1.9.1",
     "@opentelemetry/core": "2.11.0",
+    "@opentelemetry/sdk-trace-base": "2.8.0",
+    "@opentelemetry/sdk-trace-node": "2.8.0",
     "@pydantic/monty": "0.0.23",
     "express": "^5.1.0",
     "pg": "^8.22.0",
@@ -194,6 +196,10 @@
         "types": "./lib/src/telemetry.d.ts",
         "default": "./lib/src/telemetry.js"
       },
+      "./trace-attributes": {
+        "types": "./lib/src/trace-attributes.d.ts",
+        "default": "./lib/src/trace-attributes.js"
+      },
       "./validation": {
         "types": "./lib/src/validation.d.ts",
         "default": "./lib/src/validation.js"
@@ -209,7 +215,7 @@
       "./package.json": "./package.json"
     }
   },
-  "version": "0.9.64",
+  "version": "0.9.67",
   "types": "./lib/index.d.ts",
   "dbxToolsConfig": {
     "tags": [
@@ -342,6 +348,10 @@
       "types": "./lib/src/telemetry.d.ts",
       "default": "./lib/src/telemetry.js"
     },
+    "./trace-attributes": {
+      "types": "./lib/src/trace-attributes.d.ts",
+      "default": "./lib/src/trace-attributes.js"
+    },
     "./validation": {
       "types": "./lib/src/validation.d.ts",
       "default": "./lib/src/validation.js"
@@ -359,11 +369,15 @@
   "files": [
     "index.ts",
     "src",
+    "exports.ts",
     "lib"
   ],
   "peerDependenciesMeta": {
     "skills": {
       "optional": true
+    },
+    "@mlflow/core": {
+      "optional": true
     }
   },
   "//": "~~ Generated by projen. To modify, edit .projenrc.js and run \"bunx projen\"."

package/src/genie.ts

@@ -59,7 +59,7 @@ import { chartPlannerRequestSchema, chartToolOutputSchema, prepareChart } from "
 import { MASTRA_USER_KEY, resolveUserKey } from "./config.ts";
 import type { MastraPluginConfig, User } from "./config.ts";
 import { fetchStatementData } from "./statement.ts";
-import { stampGenieToolSpan } from "./telemetry.ts";
+import { recordActiveTraceAuth, stampGenieToolSpan } from "./telemetry.ts";
 import { safeWriteProgress } from "./writer.ts";
 
 const logger = log.logger("mastra/genie");
@@ -128,6 +128,7 @@ function requireClient(
       "Invoke the tool from an agent turn served by the mastra plugin, which stamps the AppKit user on the Mastra request context.",
     );
   }
+  recordActiveTraceAuth("isUserContext" in user.executionContext ? "obo" : "service-principal");
   return { client: user.executionContext.client, requestContext };
 }
 

package/src/mlflow.ts

@@ -27,13 +27,92 @@
 
 import { ConfigurationError } from "@databricks/appkit";
 import { appkit } from "@dbx-tools/appkit";
-import { asyncUtils, errorUtils, log, object } from "@dbx-tools/shared-core";
+import { asyncUtils, environmentUtils, errorUtils, log, object } from "@dbx-tools/shared-core";
 import { feedback } from "@dbx-tools/shared-mastra";
+import type { Context } from "@opentelemetry/api";
+import type {
+  ReadableSpan,
+  Span as SdkSpan,
+  SpanProcessor,
+} from "@opentelemetry/sdk-trace-base";
+import type { NodeTracerProvider } from "@opentelemetry/sdk-trace-node";
 import { z } from "zod";
 import { databricksFetch, readResponseJson, readResponseText } from "./rest.ts";
+import { MLFLOW_AGENT_TAG_ATTR } from "./trace-attributes.ts";
 
 const logger = log.logger("mastra/mlflow");
 
+type MlflowModule = typeof import("@mlflow/core");
+type UpdateCurrentTraceOptions = Parameters<MlflowModule["updateCurrentTrace"]>[0];
+
+let directMlflow: MlflowModule | undefined;
+let directMlflowInitStarted = false;
+let directMlflowProvider: NodeTracerProvider | undefined;
+let directMlflowProcessor: SpanProcessor | undefined;
+
+const NOISY_INSTRUMENTATION_SCOPES = new Set([
+  "@opentelemetry/instrumentation-express",
+  "@opentelemetry/instrumentation-http",
+  "cache-manager-cache-manager",
+]);
+
+function isAgentTraceRoot(span: SdkSpan): boolean {
+  return span.attributes[MLFLOW_AGENT_TAG_ATTR] === "true";
+}
+
+function isUsefulAgentDescendant(span: SdkSpan): boolean {
+  return !NOISY_INSTRUMENTATION_SCOPES.has(span.instrumentationScope.name);
+}
+
+/**
+ * Restrict the direct MLflow provider to chat roots and their descendants.
+ * AppKit and Mastra create other OTel roots for cache and plugin operations;
+ * forwarding those would recreate the one-trace-per-request noise this path is
+ * intended to avoid.
+ */
+export class AgentTraceSpanProcessor implements SpanProcessor {
+  private readonly traceIds = new Set<string>();
+  private readonly spanIds = new Set<string>();
+
+  constructor(private readonly delegate: SpanProcessor) {}
+
+  onStart(span: SdkSpan, parentContext: Context): void {
+    const traceId = span.spanContext().traceId;
+    const isRoot = !span.parentSpanContext?.spanId;
+    if (isRoot) {
+      if (!isAgentTraceRoot(span)) return;
+      this.traceIds.add(traceId);
+    } else if (!this.traceIds.has(traceId) || !isUsefulAgentDescendant(span)) {
+      return;
+    }
+    this.spanIds.add(span.spanContext().spanId);
+    this.delegate.onStart(span, parentContext);
+  }
+
+  onEnding(span: SdkSpan): void {
+    if (!this.spanIds.has(span.spanContext().spanId)) return;
+    this.delegate.onEnding?.(span);
+  }
+
+  onEnd(span: ReadableSpan): void {
+    const traceId = span.spanContext().traceId;
+    const spanId = span.spanContext().spanId;
+    if (!this.spanIds.delete(spanId)) return;
+    this.delegate.onEnd(span);
+    if (!span.parentSpanContext?.spanId) this.traceIds.delete(traceId);
+  }
+
+  forceFlush(): Promise<void> {
+    return this.delegate.forceFlush();
+  }
+
+  shutdown(): Promise<void> {
+    this.traceIds.clear();
+    this.spanIds.clear();
+    return this.delegate.shutdown();
+  }
+}
+
 /** Workspace client carried on an AppKit execution context. */
 type WorkspaceClient = appkit.WorkspaceClientLike;
 
@@ -214,14 +293,172 @@ function configuredExperiment(): { id?: string; name?: string } {
   };
 }
 
+/** Tracking URI for direct local MLflow tracing, with CLI-profile auth injected. */
+export function directMlflowTrackingUri(): string | undefined {
+  if (environmentUtils.isDatabricksAppEnv()) return undefined;
+  const configured = process.env.MLFLOW_TRACKING_URI?.trim();
+  if (configured && configured !== "databricks") return configured;
+  const profile = process.env.DATABRICKS_CONFIG_PROFILE?.trim();
+  return profile ? `databricks://${profile}` : "databricks";
+}
+
+/** Parse the optional UC table-prefix location used by direct local tracing. */
+export function directMlflowTraceLocation():
+  | { catalogName: string; schemaName: string; tablePrefix: string }
+  | undefined {
+  const catalogName = process.env.MLFLOW_UC_CATALOG?.trim();
+  const schemaName = process.env.MLFLOW_UC_SCHEMA?.trim();
+  const tablePrefix = process.env.MLFLOW_UC_TABLE_PREFIX?.trim();
+  if (catalogName && schemaName && tablePrefix) {
+    return { catalogName, schemaName, tablePrefix };
+  }
+
+  const prefix = normalizeUcPrefix(process.env.MLFLOW_UC_TRACE_PREFIX);
+  if (!prefix) return undefined;
+  const parts = prefix.split(".");
+  if (parts.length !== 3) return undefined;
+  const [prefixCatalog, prefixSchema, prefixTable] = parts;
+  if (!prefixCatalog || !prefixSchema || !prefixTable) return undefined;
+  return {
+    catalogName: prefixCatalog,
+    schemaName: prefixSchema,
+    tablePrefix: prefixTable,
+  };
+}
+
+/** Whether the local direct-to-experiment MLflow SDK has enough configuration. */
+export function directMlflowTracingConfigured(): boolean {
+  return Boolean(
+    !environmentUtils.isDatabricksAppEnv() &&
+      process.env.MLFLOW_EXPERIMENT_ID?.trim() &&
+      directMlflowTrackingUri(),
+  );
+}
+
+async function detectedDirectMlflowTraceLocation(
+  mlflow: MlflowModule,
+  trackingUri: string,
+  experimentId: string,
+  host: string | undefined,
+): Promise<{ catalogName: string; schemaName: string; tablePrefix: string } | undefined> {
+  const explicit = directMlflowTraceLocation();
+  if (explicit) return explicit;
+  try {
+    const authProvider = mlflow.createAuthProvider({
+      trackingUri,
+      ...(host ? { host } : {}),
+    });
+    const client = new mlflow.MlflowClient({ trackingUri, authProvider });
+    const experiment = await client.getExperiment(experimentId);
+    if (!experiment) return undefined;
+    const { ucLocationFromExperimentTags } = await import("@mlflow/core/dist/core/destination");
+    const location = ucLocationFromExperimentTags(experiment.tags);
+    if (!location?.tablePrefix) return undefined;
+    return {
+      catalogName: location.catalogName,
+      schemaName: location.schemaName,
+      tablePrefix: location.tablePrefix,
+    };
+  } catch (err) {
+    logger.warn("direct MLflow UC trace-location detection failed", {
+      experimentId,
+      error: errorUtils.errorMessage(err),
+    });
+    return undefined;
+  }
+}
+
+/**
+ * Start the MLflow Node tracer for local development when AppKit has no OTLP
+ * exporter. This provider is never started in the deployed Apps path, avoiding
+ * a second global OTel provider racing AppKit's telemetry manager.
+ */
+export async function initializeDirectMlflowTracing(): Promise<boolean> {
+  if (directMlflow) return true;
+  if (directMlflowInitStarted || !directMlflowTracingConfigured()) return false;
+  directMlflowInitStarted = true;
+
+  const experimentId = process.env.MLFLOW_EXPERIMENT_ID!.trim();
+  const trackingUri = directMlflowTrackingUri()!;
+  try {
+    const mlflow = await import("@mlflow/core");
+    const rawHost = process.env.DATABRICKS_HOST?.trim();
+    const host = rawHost
+      ? /^https?:\/\//i.test(rawHost)
+        ? rawHost
+        : `https://${rawHost}`
+      : undefined;
+    const traceLocation = await detectedDirectMlflowTraceLocation(
+      mlflow,
+      trackingUri,
+      experimentId,
+      host,
+    );
+    if (!traceLocation) {
+      throw new ConfigurationError(
+        "Direct local MLflow tracing requires a UC-linked experiment or MLFLOW_UC_TRACE_PREFIX.",
+      );
+    }
+    const [{ DatabricksUCTableSpanExporter, DatabricksUCTableSpanProcessor }, sdk] =
+      await Promise.all([
+        import("@mlflow/core/dist/exporters/uc_table"),
+        import("@opentelemetry/sdk-trace-node"),
+      ]);
+    const authProvider = mlflow.createAuthProvider({
+      trackingUri,
+      ...(host ? { host } : {}),
+    });
+    const client = new mlflow.MlflowClient({ trackingUri, authProvider });
+    const exporter = new DatabricksUCTableSpanExporter(client);
+    const processor = new AgentTraceSpanProcessor(
+      new DatabricksUCTableSpanProcessor(exporter, traceLocation),
+    );
+    const provider = new sdk.NodeTracerProvider({ spanProcessors: [processor] });
+    provider.register();
+    directMlflow = mlflow;
+    directMlflowProvider = provider;
+    directMlflowProcessor = processor;
+    logger.info("direct MLflow tracing enabled", {
+      experimentId,
+      trackingUri,
+      traceLocation,
+    });
+    return true;
+  } catch (err) {
+    logger.warn("direct MLflow tracing disabled", {
+      error: errorUtils.errorMessage(err),
+    });
+    return false;
+  }
+}
+
+/** Persist trace-level fields on the direct MLflow trace in the current OTel context. */
+export function updateActiveMlflowTrace(options: UpdateCurrentTraceOptions): void {
+  directMlflow?.updateCurrentTrace(options);
+}
+
+/** Persist tags on the direct MLflow trace active in the current OTel context. */
+export function updateActiveMlflowTraceTags(tags: Record<string, string>): void {
+  updateActiveMlflowTrace({ tags });
+}
+
+/** Whether this process owns the filtered local MLflow tracer provider. */
+export function directMlflowTracingActive(): boolean {
+  return directMlflowProcessor !== undefined;
+}
+
+/** Flush pending direct local trace exports during graceful shutdown. */
+export async function flushDirectMlflowTracing(): Promise<void> {
+  await directMlflowProcessor?.forceFlush();
+  await directMlflowProvider?.forceFlush();
+}
+
 /**
  * Whether MLflow feedback logging is available for this deployment.
  *
- * Enabled when an OTLP exporter endpoint is configured (traces are
- * actually shipped somewhere) AND an MLflow experiment is named - the
- * two signals that the OTLP backend is MLflow and traces will
- * materialize there. Both are standard env vars, so no plugin config is
- * required; a deployment opts in simply by wiring MLflow tracing.
+ * Enabled when an MLflow experiment is configured and traces have an export
+ * path: AppKit OTLP in a Databricks App, or the direct MLflow Node SDK outside
+ * one. Both paths are selected from standard environment values.
  */
 export function mlflowEnabled(): boolean {
   const hasExporter = Boolean(
@@ -231,7 +468,7 @@ export function mlflowEnabled(): boolean {
   const hasExperiment = Boolean(
     process.env.MLFLOW_EXPERIMENT_ID?.trim() || process.env.MLFLOW_EXPERIMENT_NAME?.trim(),
   );
-  return hasExporter && hasExperiment;
+  return hasExperiment && (hasExporter || directMlflowTracingConfigured());
 }
 
 /**

package/src/model.ts

@@ -37,6 +37,7 @@ import {
   rewriteServingResponseStream,
 } from "./serving-sanitize.ts";
 import { MASTRA_MODEL_OVERRIDE_KEY, resolveServingConfig } from "./serving.ts";
+import { recordActiveTraceAuth, recordActiveTraceModel } from "./telemetry.ts";
 
 type ModelClass = model.ModelClass;
 const { parseModelClass } = classes;
@@ -135,6 +136,7 @@ export async function buildModel(
   // (the active OBO scope, or the service principal) when it's absent.
   const user = requestContext.get(MASTRA_USER_KEY) as User | undefined;
   const executionContext = user?.executionContext ?? getExecutionContext();
+  recordActiveTraceAuth("isUserContext" in executionContext ? "obo" : "service-principal");
   const clientConfig = executionContext.client.config;
   const host = (await clientConfig.getHost()).toString();
   const headers = new Headers();
@@ -165,6 +167,7 @@ export async function buildModel(
     selectionInput(config, requested, overrides.modelClass, serving),
   );
   logger.debug("model selected", { modelId, source, requested });
+  recordActiveTraceModel(modelId);
 
   return {
     providerId: config.providerId ?? "databricks",

package/src/observability.ts

@@ -49,7 +49,7 @@ import { OtelBridge } from "@mastra/otel-bridge";
 import { propagation } from "@opentelemetry/api";
 
 import { TRACE_REQUEST_CONTEXT_KEYS } from "./config.ts";
-import { mlflowEnabled } from "./mlflow.ts";
+import { initializeDirectMlflowTracing, mlflowEnabled } from "./mlflow.ts";
 
 const logger = log.logger("mastra/observability");
 
@@ -117,7 +117,8 @@ export function configureOtelPropagation(): boolean {
  * OTel pipeline via `@mastra/otel-bridge`.
  *
  * Returns `undefined` when tracing is off: either the caller passed
- * `enabled: false`, or auto mode finds no OTLP endpoint (the default).
+ * `enabled: false`, or auto mode finds neither AppKit OTLP nor a local MLflow
+ * experiment.
  * Skipping the bridge in that case avoids `[OtelBridge] No OTEL span
  * found` log spam from Mastra spans that have no parent on the noop
  * tracer.
@@ -128,13 +129,17 @@ export async function buildObservability(
   const otelBase = process.env.OTEL_EXPORTER_OTLP_ENDPOINT;
   const otelTracesOverride = process.env.OTEL_EXPORTER_OTLP_TRACES_ENDPOINT;
   const otlpConfigured = isOtlpTracingConfigured();
+  const directMlflowConfigured =
+    options?.enabled === false || otlpConfigured ? false : await initializeDirectMlflowTracing();
+  if (directMlflowConfigured) configureOtelPropagation();
+  const tracingConfigured = otlpConfigured || directMlflowConfigured;
 
-  if (options?.enabled === false || (options?.enabled !== true && !otlpConfigured)) {
+  if (options?.enabled === false || (options?.enabled !== true && !tracingConfigured)) {
     logger.info("Mastra observability off", {
       reason:
         options?.enabled === false
           ? "disabled in plugin config"
-          : "OTEL_EXPORTER_OTLP_ENDPOINT unset",
+          : "no OTLP endpoint or direct MLflow experiment configured",
     });
     return undefined;
   }
@@ -164,7 +169,7 @@ export async function buildObservability(
     otelBase: otelBase ?? "<unset>",
     resolvedTracesUrl: resolvedTracesUrl ?? "<unset>",
     feedback,
-    observability: feedback ? "mlflow" : "otel",
+    observability: directMlflowConfigured ? "mlflow-direct" : feedback ? "mlflow" : "otel",
   });
 
   return new Observability({

package/src/plugin.ts

@@ -100,6 +100,7 @@ import {
 import { buildMcpServer, type ResolvedMcp } from "./mcp.ts";
 import { createMemoryBuilder, createServicePrincipalPool, needsLakebase } from "./memory.ts";
 import {
+  flushDirectMlflowTracing,
   logFeedback,
   mlflowExperimentManagerUrl,
   resolveFeedbackEnabled,
@@ -304,6 +305,13 @@ export class MastraPlugin extends Plugin<MastraPluginConfig> {
    * shutdown budget.
    */
   async shutdown(): Promise<void> {
+    try {
+      await flushDirectMlflowTracing();
+    } catch (err) {
+      this.logger.warn("error flushing direct MLflow traces", {
+        error: errorUtils.errorMessage(err),
+      });
+    }
     const pool = this.servicePrincipalPool;
     if (!pool) return;
     this.servicePrincipalPool = null;

package/src/server.ts

@@ -29,6 +29,7 @@ import {
 } from "./config.ts";
 import { requestUserEmail, requestUserId } from "./identity.ts";
 import { resolveFeedbackEnabled } from "./mlflow.ts";
+import { recordActiveTraceAuth, recordActiveTraceUser } from "./telemetry.ts";
 
 import {
   extractModelOverride,
@@ -222,6 +223,18 @@ export class MastraServer extends ExpressMastraServer {
       userId: requestUserId(req),
       email: requestUserEmail(req),
     });
+    const user = requestContext.get(MASTRA_USER_KEY) as User | undefined;
+    if (user) {
+      recordActiveTraceAuth(
+        "isUserContext" in user.executionContext ? "obo" : "service-principal",
+      );
+    }
+    const traceUser = [
+      requestContext.get(MASTRA_USER_EMAIL_KEY),
+      requestContext.get(MASTRA_USER_NAME_KEY),
+      requestContext.get(MASTRA_RESOURCE_ID_KEY),
+    ].find((value): value is string => typeof value === "string" && value.trim().length > 0);
+    recordActiveTraceUser(traceUser);
   }
 
   /**