@dbx-tools/appkit-mastra 0.9.67 → 0.9.69

package/src/mlflow.ts

@@ -25,10 +25,13 @@
  * @module
  */
 
-import { ConfigurationError } from "@databricks/appkit";
+import { ConfigurationError, createWorkspaceClient } from "@databricks/appkit";
 import { appkit } from "@dbx-tools/appkit";
-import { asyncUtils, environmentUtils, errorUtils, log, object } from "@dbx-tools/shared-core";
+import { asyncUtils, errorUtils, log, object } from "@dbx-tools/shared-core";
 import { feedback } from "@dbx-tools/shared-mastra";
+import { TraceLocationType, TraceMetadataKey } from "@mlflow/core";
+import { TraceInfo } from "@mlflow/core/dist/core/entities/trace_info";
+import { TraceState } from "@mlflow/core/dist/core/entities/trace_state";
 import type { Context } from "@opentelemetry/api";
 import type {
   ReadableSpan,
@@ -49,6 +52,10 @@ let directMlflow: MlflowModule | undefined;
 let directMlflowInitStarted = false;
 let directMlflowProvider: NodeTracerProvider | undefined;
 let directMlflowProcessor: SpanProcessor | undefined;
+let appMlflowTraceInfoInitStarted = false;
+let appMlflowTraceInfoClient: WorkspaceClient | undefined;
+let appMlflowTracePrefix: string | undefined;
+const pendingAppMlflowTraceInfo = new Set<Promise<void>>();
 
 const NOISY_INSTRUMENTATION_SCOPES = new Set([
   "@opentelemetry/instrumentation-express",
@@ -295,7 +302,7 @@ 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;
+  if (appkit.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();
@@ -329,12 +336,116 @@ export function directMlflowTraceLocation():
 /** Whether the local direct-to-experiment MLflow SDK has enough configuration. */
 export function directMlflowTracingConfigured(): boolean {
   return Boolean(
-    !environmentUtils.isDatabricksAppEnv() &&
+    !appkit.isDatabricksAppEnv() &&
       process.env.MLFLOW_EXPERIMENT_ID?.trim() &&
       directMlflowTrackingUri(),
   );
 }
 
+/** Whether a Databricks App can promote OTLP spans into full MLflow trace info. */
+export function appMlflowTraceInfoConfigured(): boolean {
+  return Boolean(
+    appkit.isDatabricksAppEnv() &&
+      (process.env.MLFLOW_EXPERIMENT_ID?.trim() || process.env.MLFLOW_EXPERIMENT_NAME?.trim()) &&
+      (process.env.OTEL_EXPORTER_OTLP_ENDPOINT?.trim() ||
+        process.env.OTEL_EXPORTER_OTLP_TRACES_ENDPOINT?.trim()),
+  );
+}
+
+/** Initialize the App service-principal client used to persist V4 trace tags and metadata. */
+export async function initializeAppMlflowTraceInfo(
+  client?: WorkspaceClient,
+): Promise<boolean> {
+  if (appMlflowTraceInfoClient && appMlflowTracePrefix) return true;
+  if (appMlflowTraceInfoInitStarted || !appMlflowTraceInfoConfigured()) return false;
+  appMlflowTraceInfoInitStarted = true;
+  try {
+    const resolvedClient = client ?? (createWorkspaceClient() as WorkspaceClient);
+    const prefix = await resolveUcTracePrefix(resolvedClient);
+    if (!prefix || prefix.split(".").length !== 3) {
+      throw new ConfigurationError(
+        "Databricks Apps MLflow trace-info promotion requires a UC-linked experiment or MLFLOW_UC_TRACE_PREFIX.",
+      );
+    }
+    appMlflowTraceInfoClient = resolvedClient;
+    appMlflowTracePrefix = prefix;
+    logger.info("Databricks Apps MLflow trace-info promotion enabled", { prefix });
+    return true;
+  } catch (err) {
+    logger.warn("Databricks Apps MLflow trace-info promotion disabled", {
+      error: errorUtils.errorMessage(err),
+    });
+    return false;
+  }
+}
+
+/** Trace-level values persisted beside spans already exported by the Apps OTLP sidecar. */
+export interface AppMlflowTraceInfo {
+  traceId: string;
+  requestTime: number;
+  executionDuration: number;
+  error: boolean;
+  tags: Record<string, string>;
+  user?: string;
+  sessionId?: string;
+  requestPreview?: string;
+  responsePreview?: string;
+}
+
+/** Persist MLflow V4 trace info for one Databricks Apps OTel root without re-exporting spans. */
+export function persistAppMlflowTraceInfo(params: AppMlflowTraceInfo): void {
+  const client = appMlflowTraceInfoClient;
+  const prefix = appMlflowTracePrefix;
+  if (!client || !prefix) return;
+  const [catalogName, schemaName, tablePrefix] = prefix.split(".");
+  if (!catalogName || !schemaName || !tablePrefix) return;
+  const traceInfo = new TraceInfo({
+    traceId: `trace:/${prefix}/${params.traceId}`,
+    traceLocation: {
+      type: TraceLocationType.UC_TABLE_PREFIX,
+      ucTablePrefix: { catalogName, schemaName, tablePrefix },
+    },
+    requestTime: params.requestTime,
+    executionDuration: params.executionDuration,
+    state: params.error ? TraceState.ERROR : TraceState.OK,
+    ...(params.requestPreview ? { requestPreview: params.requestPreview } : {}),
+    ...(params.responsePreview ? { responsePreview: params.responsePreview } : {}),
+    traceMetadata: {
+      [TraceMetadataKey.SCHEMA_VERSION]: "4",
+      ...(params.user ? { [TraceMetadataKey.TRACE_USER]: params.user } : {}),
+      ...(params.sessionId ? { [TraceMetadataKey.TRACE_SESSION]: params.sessionId } : {}),
+    },
+    tags: params.tags,
+    assessments: [],
+  });
+  const path = `/api/4.0/mlflow/traces/${encodeURIComponent(prefix)}/${params.traceId}/info`;
+  const pending = databricksFetch(client, path, { method: "POST", body: traceInfo.toJson() })
+    .then(async (response) => {
+      if (response.ok) return;
+      logger.warn("Databricks Apps MLflow trace-info promotion failed", {
+        traceId: params.traceId,
+        status: response.status,
+        body: await readResponseText(response),
+      });
+    })
+    .catch((err) => {
+      logger.warn("Databricks Apps MLflow trace-info promotion failed", {
+        traceId: params.traceId,
+        error: errorUtils.errorMessage(err),
+      });
+    })
+    .finally(() => pendingAppMlflowTraceInfo.delete(pending));
+  pendingAppMlflowTraceInfo.add(pending);
+}
+
+/** Reset Apps trace-info promotion state for tests. */
+export function resetAppMlflowTraceInfo(): void {
+  appMlflowTraceInfoInitStarted = false;
+  appMlflowTraceInfoClient = undefined;
+  appMlflowTracePrefix = undefined;
+  pendingAppMlflowTraceInfo.clear();
+}
+
 async function detectedDirectMlflowTraceLocation(
   mlflow: MlflowModule,
   trackingUri: string,
@@ -449,6 +560,7 @@ export function directMlflowTracingActive(): boolean {
 
 /** Flush pending direct local trace exports during graceful shutdown. */
 export async function flushDirectMlflowTracing(): Promise<void> {
+  await Promise.all([...pendingAppMlflowTraceInfo]);
   await directMlflowProcessor?.forceFlush();
   await directMlflowProvider?.forceFlush();
 }

package/src/observability.ts

@@ -49,7 +49,11 @@ import { OtelBridge } from "@mastra/otel-bridge";
 import { propagation } from "@opentelemetry/api";
 
 import { TRACE_REQUEST_CONTEXT_KEYS } from "./config.ts";
-import { initializeDirectMlflowTracing, mlflowEnabled } from "./mlflow.ts";
+import {
+  initializeAppMlflowTraceInfo,
+  initializeDirectMlflowTracing,
+  mlflowEnabled,
+} from "./mlflow.ts";
 
 const logger = log.logger("mastra/observability");
 
@@ -129,6 +133,8 @@ 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 appMlflowTraceInfo =
+    options?.enabled === false ? false : await initializeAppMlflowTraceInfo();
   const directMlflowConfigured =
     options?.enabled === false || otlpConfigured ? false : await initializeDirectMlflowTracing();
   if (directMlflowConfigured) configureOtelPropagation();
@@ -169,7 +175,13 @@ export async function buildObservability(
     otelBase: otelBase ?? "<unset>",
     resolvedTracesUrl: resolvedTracesUrl ?? "<unset>",
     feedback,
-    observability: directMlflowConfigured ? "mlflow-direct" : feedback ? "mlflow" : "otel",
+    observability: directMlflowConfigured
+      ? "mlflow-direct"
+      : appMlflowTraceInfo
+        ? "mlflow-apps"
+        : feedback
+          ? "mlflow"
+          : "otel",
   });
 
   return new Observability({

package/src/telemetry.ts

@@ -20,8 +20,9 @@
 
 import { StringDecoder } from "node:string_decoder";
 
+import { appkit } from "@dbx-tools/appkit";
 import { getRequestTags } from "@dbx-tools/appkit/request-tags";
-import { environmentUtils, json, log, object, stringUtils } from "@dbx-tools/shared-core";
+import { json, log, object, stringUtils } from "@dbx-tools/shared-core";
 import { feedback, thread } from "@dbx-tools/shared-mastra";
 import { context, createContextKey, SpanKind, trace, type Span } from "@opentelemetry/api";
 import { getRPCMetadata, RPCType } from "@opentelemetry/core";
@@ -30,6 +31,7 @@ import type express from "express";
 import { USER_EMAIL_HEADER, USER_ID_HEADER } from "./identity.ts";
 import {
   directMlflowTracingActive,
+  persistAppMlflowTraceInfo,
   updateActiveMlflowTrace,
   updateActiveMlflowTraceTags,
 } from "./mlflow.ts";
@@ -505,8 +507,9 @@ export function chatTurnTelemetryMiddleware(
     return;
   }
   const traceState: ChatTraceState = { models: new Set(), tags: new Map() };
+  const traceRequestTime = Date.now();
   setChatTraceTag(traceState, MLFLOW_AGENT_TAG, true);
-  if (!environmentUtils.isDatabricksAppEnv()) {
+  if (!appkit.isDatabricksAppEnv()) {
     setChatTraceTag(traceState, MLFLOW_LOCAL_TAG, true);
   }
   for (const [name, value] of Object.entries(getRequestTags(req))) {
@@ -536,16 +539,19 @@ export function chatTurnTelemetryMiddleware(
   target.span.setAttribute(CHAT_GENIE_USED_ATTR, false);
 
   const messages = (req.body as { messages?: unknown } | undefined)?.messages;
+  let requestPreview: string | undefined;
   if (messages !== undefined) {
     const fullInput = serialized(messages);
+    requestPreview = textOnlyChatInput(messages) ?? fullInput;
     target.span.setAttribute(CHAT_MESSAGES_ATTR, fullInput);
     target.span.setAttribute(
       MLFLOW_SPAN_INPUTS_ATTR,
-      (textOnlyChatInput(messages) ?? fullInput).slice(0, TRACE_IO_LIMIT),
+      requestPreview.slice(0, TRACE_IO_LIMIT),
     );
   }
 
   const collector = new AssistantResponseCollector();
+  let responsePreview: string | undefined;
   let outputRecorded = false;
   let fullOutputRecorded = false;
   let genieRecorded = false;
@@ -556,6 +562,7 @@ export function chatTurnTelemetryMiddleware(
   const passThroughJson = res.json.bind(res) as (body?: unknown) => express.Response;
   const recordOutput = (answer: string): void => {
     if (outputRecorded || !answer) return;
+    responsePreview = answer;
     target.span.setAttribute(MLFLOW_SPAN_OUTPUTS_ATTR, answer.slice(0, TRACE_IO_LIMIT));
     outputRecorded = true;
   };
@@ -590,6 +597,16 @@ export function chatTurnTelemetryMiddleware(
         ...(traceState.user ? { user: traceState.user } : {}),
       });
     });
+    persistAppMlflowTraceInfo({
+      traceId: targetTraceId,
+      requestTime: traceRequestTime,
+      executionDuration: Date.now() - traceRequestTime,
+      error: res.statusCode >= 500,
+      tags,
+      ...(traceState.user ? { user: traceState.user } : {}),
+      ...(requestPreview ? { requestPreview } : {}),
+      ...(responsePreview ? { responsePreview } : {}),
+    });
   };
   const endOwnedSpan = (): void => {
     if (!target.owned || ownedSpanEnded) return;