@dbx-tools/appkit-mastra 0.9.143 → 0.9.145

package/src/mlflow.ts

@@ -49,6 +49,7 @@ let directMlflowInitStarted = false;
 let directMlflowProvider: NodeTracerProvider | undefined;
 let directMlflowProcessor: SpanProcessor | undefined;
 let appMlflowTraceInfoInitStarted = false;
+let appMlflowTraceInfoRequestInitStarted = false;
 let appMlflowTraceInfoClient: WorkspaceClient | undefined;
 let appMlflowTracePrefix: string | undefined;
 const pendingAppMlflowTraceInfo = new Set<Promise<void>>();
@@ -209,7 +210,18 @@ async function resolveUcTracePrefix(client: WorkspaceClient): Promise<string | u
   if (loadedUcPrefix) return loadedUcPrefix.prefix;
   const experiment = await fetchExperiment(client);
   if (!experiment) return undefined;
-  const prefix = ucTracePrefixFromExperimentTags(experiment.tags);
+  const location = experiment.trace_location?.uc_trace_location;
+  const prefix =
+    ucTracePrefixFromExperimentTags(experiment.tags) ??
+    (location
+      ? [
+          location.catalog,
+          location.schema,
+          location.effective_table_prefix ?? location.table_prefix,
+        ]
+          .filter(Boolean)
+          .join(".")
+      : undefined);
   loadedUcPrefix = { prefix };
   return prefix;
 }
@@ -229,6 +241,23 @@ const MlflowExperimentResponseSchema = z
       .object({
         experiment_id: z.string().min(1).describe("Workspace MLflow experiment identifier."),
         tags: z.unknown().optional().describe("Experiment tags, including the UC trace location."),
+        trace_location: z
+          .object({
+            uc_trace_location: z
+              .object({
+                catalog: z.string().describe("Unity Catalog catalog name."),
+                schema: z.string().describe("Unity Catalog schema name."),
+                table_prefix: z.string().describe("Configured trace table prefix."),
+                effective_table_prefix: z
+                  .string()
+                  .optional()
+                  .describe("Effective trace table prefix selected by Databricks."),
+              })
+              .optional()
+              .describe("Unity Catalog trace location."),
+          })
+          .optional()
+          .describe("Experiment trace storage location."),
       })
       .passthrough()
       .describe("Configured MLflow experiment."),
@@ -354,13 +383,13 @@ export async function initializeAppMlflowTraceInfo(client?: WorkspaceClient): Pr
   appMlflowTraceInfoInitStarted = true;
   try {
     const resolvedClient = client ?? (createWorkspaceClient() as WorkspaceClient);
+    appMlflowTraceInfoClient ??= resolvedClient;
     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;
@@ -368,10 +397,25 @@ export async function initializeAppMlflowTraceInfo(client?: WorkspaceClient): Pr
     logger.warn("Databricks Apps MLflow trace-info promotion disabled", {
       error: errorUtils.toError(err),
     });
+    // Startup uses the service principal. Permit one later request-scoped
+    // client to resolve experiment metadata when that principal cannot read it.
+    if (client === undefined) appMlflowTraceInfoInitStarted = false;
     return false;
   }
 }
 
+/**
+ * Retry Apps trace-location discovery once with a request-scoped workspace
+ * client after service-principal startup discovery failed.
+ */
+export async function initializeRequestMlflowTraceInfo(client: WorkspaceClient): Promise<boolean> {
+  if (appMlflowTraceInfoClient && appMlflowTracePrefix) return true;
+  if (appMlflowTraceInfoRequestInitStarted || !appMlflowTraceInfoConfigured()) return false;
+  appMlflowTraceInfoRequestInitStarted = true;
+  appMlflowTraceInfoInitStarted = false;
+  return initializeAppMlflowTraceInfo(client);
+}
+
 /** Trace-level values persisted beside spans already exported by the Apps OTLP sidecar. */
 export interface AppMlflowTraceInfo {
   traceId: string;
@@ -434,6 +478,7 @@ export function persistAppMlflowTraceInfo(params: AppMlflowTraceInfo): void {
 /** Reset Apps trace-info promotion state for tests. */
 export function resetAppMlflowTraceInfo(): void {
   appMlflowTraceInfoInitStarted = false;
+  appMlflowTraceInfoRequestInitStarted = false;
   appMlflowTraceInfoClient = undefined;
   appMlflowTracePrefix = undefined;
   pendingAppMlflowTraceInfo.clear();
@@ -644,7 +689,7 @@ async function resolveMlflowExperimentManagerUrl(
 
 async function fetchExperiment(
   client: WorkspaceClient,
-): Promise<{ experiment_id: string; tags?: unknown } | undefined> {
+): Promise<z.infer<typeof MlflowExperimentResponseSchema>["experiment"] | undefined> {
   const experiment = configuredExperiment();
   const path = experiment.id
     ? `/api/2.0/mlflow/experiments/get?experiment_id=${encodeURIComponent(experiment.id)}`

package/src/plugin.ts

@@ -278,6 +278,11 @@ export class MastraPlugin extends Plugin<MastraPluginConfig> {
     });
     this.config.feedbackMode = this.resolvedFeedbackMode;
     if (this.config.observability !== false) {
+      await appkitTelemetry.configureNativeHttpTelemetry({
+        // A low-level request without an active semantic span becomes a noisy
+        // standalone trace. Requests inside an agent/tool span remain visible.
+        outgoingFilter: appkitTelemetry.excludeUnparentedOutgoingHttpTelemetry,
+      });
       await registerAppMlflowTraceInfoProcessor();
     }
     // Wait until sibling plugins (e.g. `lakebase`) finish `setup()` so
@@ -386,8 +391,8 @@ export class MastraPlugin extends Plugin<MastraPluginConfig> {
       getMastra: () => this.mastra,
       /**
        * Build the `RequestContext` an agent turn driven from OUTSIDE this
-       * plugin's routes needs (a Teams activity on another plugin's endpoint, a
-       * scheduled job).
+       * plugin's routes needs, such as another plugin's endpoint or a scheduled
+       * job.
        *
        * Mastra's user-scoped tools - `ask_genie` most visibly - read the AppKit
        * user off the request context, which only the HTTP middleware stamps.

package/src/server.ts

@@ -35,6 +35,7 @@ import {
   FOLLOW_UP_USER_TEXT_KEY,
 } from "./follow-up-suggestions.ts";
 import { requestUserEmail, requestUserId } from "./identity.ts";
+import { initializeRequestMlflowTraceInfo } from "./mlflow.ts";
 import { MASTRA_ORIGIN_KEY } from "./request-context.ts";
 import {
   extractModelOverride,
@@ -118,12 +119,11 @@ export interface AttributedIdentity {
  * `requestContext`.
  *
  * Split out of {@link MastraServer} because the HTTP middleware is not the only
- * caller that needs it: an out-of-band turn (a Teams activity arriving on
- * another plugin's route, a scheduled job) drives `agent.generate` directly, and
- * without this stamp every user-scoped tool - `ask_genie` above all - fails with
- * "invoke the tool from an agent turn served by the mastra plugin". Those
- * callers build a context with {@link createRequestContext} instead of a
- * request.
+ * caller that needs it: an out-of-band turn from another plugin route or a
+ * scheduled job drives `agent.generate` directly, and without this stamp every
+ * user-scoped tool - `ask_genie` above all - fails with "invoke the tool from an
+ * agent turn served by the mastra plugin". Those callers build a context with
+ * {@link createRequestContext} instead of a request.
  *
  * The Databricks credential always comes from the ambient execution context;
  * `attributed` only changes WHO the turn is attributed to when that context is
@@ -138,6 +138,7 @@ export async function stampRequestContextUser(
 ): Promise<void> {
   if ([MASTRA_USER_KEY, MASTRA_RESOURCE_ID_KEY].every((key) => requestContext.get(key))) return;
   const executionContext = getExecutionContext();
+  await initializeRequestMlflowTraceInfo(executionContext.client);
   const isUserContext = "isUserContext" in executionContext;
   // The Databricks CREDENTIAL is always the ambient execution context's client
   // (OBO in a user context, the service principal otherwise). WHO the turn is
@@ -185,8 +186,8 @@ export async function stampRequestContextUser(
  * a request id so the turn's spans join up in traces. The result is passed as
  * `requestContext` to `agent.generate` / `agent.stream`.
  *
- * This is what makes a card turn answer with the SAME data a chat turn does: the
- * difference between the two endpoints is presentation, not capability.
+ * This gives an out-of-band turn the same data and tool capability as a chat
+ * turn.
  */
 export async function createRequestContext(
   options: { threadId?: string; resourceId?: string; requestId?: string } = {},
@@ -194,8 +195,8 @@ export async function createRequestContext(
   const requestContext = new RequestContext();
   await stampRequestContextUser(requestContext);
   if (options.threadId) requestContext.set(MASTRA_THREAD_ID_KEY, options.threadId);
-  // Stamped AFTER the user so an explicit resource id (a channel member, say)
-  // wins over the ambient identity `stampRequestContextUser` defaulted to.
+  // Stamped AFTER the user so an explicit resource id wins over the ambient
+  // identity `stampRequestContextUser` defaulted to.
   if (options.resourceId) requestContext.set(MASTRA_RESOURCE_ID_KEY, options.resourceId);
   requestContext.set(MASTRA_REQUEST_ID_KEY, options.requestId ?? hash.id());
   return requestContext;

package/src/telemetry.ts

@@ -42,6 +42,7 @@ import type express from "express";
 import { MASTRA_AUTH_MODE_KEY, MASTRA_USER_EMAIL_KEY, MASTRA_USER_NAME_KEY } from "./config.ts";
 import { USER_EMAIL_HEADER, USER_ID_HEADER } from "./identity.ts";
 import {
+  appMlflowTraceInfoConfigured,
   directMlflowTracingActive,
   initializeAppMlflowTraceInfo,
   persistAppMlflowTraceInfo,
@@ -118,7 +119,7 @@ export const GEN_AI_OPERATION_NAME_ATTR = "gen_ai.operation.name";
  * Root-span user id the MLflow experiment UI's User column reads from OTel
  * traces. Prefer the forwarded email when present so the column is human.
  */
-export const MLFLOW_USER_ATTR = "user.id";
+export const MLFLOW_USER_ATTR = appkitTelemetry.USER_TRACE_ATTRIBUTE;
 /** Root-span session id that enables the MLflow Sessions tab. */
 export const MLFLOW_SESSION_ATTR = "session.id";
 /** Root-span type for a Mastra agent turn. */
@@ -507,7 +508,8 @@ let appMlflowTraceInfoProcessorRegistered = false;
 /** Register trace enrichment before AppKit constructs its native provider. */
 export async function registerAppMlflowTraceInfoProcessor(): Promise<boolean> {
   if (appMlflowTraceInfoProcessorRegistered) return true;
-  if (!(await initializeAppMlflowTraceInfo())) return false;
+  if (!appMlflowTraceInfoConfigured()) return false;
+  await initializeAppMlflowTraceInfo();
   await appkitTelemetry.registerNativeSpanProcessor(new AppMlflowTraceInfoSpanProcessor());
   appMlflowTraceInfoProcessorRegistered = true;
   return true;