@dbx-tools/appkit-mastra 0.9.58 → 0.9.61

package/src/mlflow.ts

@@ -13,17 +13,23 @@
  *
  * There is no MLflow JS SDK, so this posts to the assessments REST
  * endpoint directly using the OBO-scoped workspace client (the feedback
- * is thus attributed to the signed-in user). Trace export is
- * asynchronous, so the just-finished trace may not exist in MLflow yet
- * when the user reacts; the log call retries briefly on "not found"
- * before giving up softly.
+ * is thus attributed to the signed-in user). UC-backed Apps traces need
+ * the experiment's table-prefix URI (`trace:/{catalog}.{schema}.{prefix}/{tr-hex}`)
+ * rather than the experiment-store id the UI does not store. That prefix is
+ * read from the experiment's `mlflow.experiment.databricksTraceDestinationPath`
+ * tag (or the spans-table tag); `MLFLOW_UC_TRACE_PREFIX` remains an optional
+ * override. Trace export is asynchronous, so the just-finished trace may not
+ * exist in MLflow yet when the user reacts; the log call retries briefly on
+ * "not found" before giving up softly.
  *
  * @module
  */
 
+import { ConfigurationError } from "@databricks/appkit";
 import { appkit } from "@dbx-tools/appkit";
-import { asyncUtils, errorUtils, log } from "@dbx-tools/shared-core";
+import { asyncUtils, errorUtils, log, object } from "@dbx-tools/shared-core";
 import { feedback } from "@dbx-tools/shared-mastra";
+import { z } from "zod";
 import { databricksFetch, readResponseJson, readResponseText } from "./rest.ts";
 
 const logger = log.logger("mastra/mlflow");
@@ -31,6 +37,98 @@ const logger = log.logger("mastra/mlflow");
 /** Workspace client carried on an AppKit execution context. */
 type WorkspaceClient = appkit.WorkspaceClientLike;
 
+/** Experiment tag Databricks stores as `catalog.schema.table_prefix`. */
+const TRACE_DESTINATION_PATH_TAG = "mlflow.experiment.databricksTraceDestinationPath";
+/** Experiment tag naming the OTel spans table (`catalog.schema.{prefix}_otel_spans`). */
+const TRACE_SPAN_STORAGE_TABLE_TAG = "mlflow.experiment.databricksTraceSpanStorageTable";
+
+/**
+ * Cached UC table prefix resolved from the experiment. `loaded` distinguishes
+ * "not fetched yet" from "fetched, experiment has no UC location". Network
+ * failures are not cached so a later thumbs click can retry.
+ */
+let loadedUcPrefix: { prefix: string | undefined } | undefined;
+
+/** Strip the `trace:/` URI wrapper and trailing slashes from a UC prefix. */
+function normalizeUcPrefix(raw: string | undefined): string | undefined {
+  const value = raw?.trim().replace(/^trace:\//, "").replace(/\/+$/, "");
+  return value || undefined;
+}
+
+/**
+ * Derive `catalog.schema.table_prefix` from a spans table name.
+ * Databricks names the table `{prefix}_otel_spans`.
+ */
+function prefixFromSpanTable(table: string | undefined): string | undefined {
+  const name = table?.trim();
+  if (!name) return undefined;
+  if (name.endsWith("_otel_spans")) return name.slice(0, -"_otel_spans".length);
+  if (name.endsWith("_spans")) return name.slice(0, -"_spans".length);
+  return undefined;
+}
+
+function experimentTagMap(raw: unknown): Map<string, string> {
+  const tags = new Map<string, string>();
+  if (Array.isArray(raw)) {
+    for (const tag of raw) {
+      if (!object.isRecord(tag) || typeof tag.key !== "string" || typeof tag.value !== "string") {
+        continue;
+      }
+      tags.set(tag.key, tag.value);
+    }
+    return tags;
+  }
+  if (!object.isRecord(raw)) return tags;
+  for (const [key, value] of Object.entries(raw)) {
+    if (typeof value === "string") tags.set(key, value);
+  }
+  return tags;
+}
+
+/**
+ * UC table prefix the assessments API needs, read from experiment tags.
+ *
+ * Prefers the three-part destination path Databricks writes when the
+ * experiment is linked to a table-prefix location. Falls back to stripping
+ * `_otel_spans` from the spans-table tag. Two-part schema-linked destinations
+ * are ignored (they are not a table prefix).
+ */
+export function ucTracePrefixFromExperimentTags(rawTags: unknown): string | undefined {
+  const tags = experimentTagMap(rawTags);
+  const destination = normalizeUcPrefix(tags.get(TRACE_DESTINATION_PATH_TAG));
+  if (destination && destination.split(".").length >= 3) return destination;
+  return prefixFromSpanTable(tags.get(TRACE_SPAN_STORAGE_TABLE_TAG));
+}
+
+/** Drop the in-process UC prefix cache (tests). */
+export function resetUcTracePrefixCache(): void {
+  loadedUcPrefix = undefined;
+}
+
+/**
+ * Trace id the assessments REST API accepts for this deployment.
+ *
+ * Databricks Apps persist traces in Unity Catalog. The experiment UI lists
+ * those rows, but `POST /api/3.0/mlflow/traces/{id}/assessments` looks up the
+ * experiment store unless `{id}` is the UC URI `trace:/{catalog}.{schema}.{prefix}/{tr-hex}`.
+ * A bare `tr-<hex>` 404s and the assessment never appears.
+ */
+export function mlflowAssessmentTraceId(traceId: string, prefix?: string): string {
+  const resolved = normalizeUcPrefix(prefix);
+  return resolved ? `trace:/${resolved}/${traceId}` : traceId;
+}
+
+async function resolveUcTracePrefix(client: WorkspaceClient): Promise<string | undefined> {
+  const fromEnv = normalizeUcPrefix(process.env.MLFLOW_UC_TRACE_PREFIX);
+  if (fromEnv) return fromEnv;
+  if (loadedUcPrefix) return loadedUcPrefix.prefix;
+  const experiment = await fetchExperiment(client);
+  if (!experiment) return undefined;
+  const prefix = ucTracePrefixFromExperimentTags(experiment.tags);
+  loadedUcPrefix = { prefix };
+  return prefix;
+}
+
 /** Assessments REST path for a trace. `3.0` is the current MLflow API version. */
 const assessmentsPath = (traceId: string): string =>
   `/api/3.0/mlflow/traces/${encodeURIComponent(traceId)}/assessments`;
@@ -40,6 +138,79 @@ const NOT_FOUND_RETRIES = 3;
 /** Base backoff between "trace not found" retries, in ms (grows linearly). */
 const NOT_FOUND_BACKOFF_MS = 1200;
 
+const MlflowExperimentResponseSchema = z
+  .object({
+    experiment: 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."),
+      })
+      .passthrough()
+      .describe("Configured MLflow experiment."),
+  })
+  .passthrough()
+  .describe("MLflow get-experiment response.");
+
+const CurrentUserResponseSchema = z
+  .object({
+    userName: z.string().optional().describe("Current workspace user name."),
+    applicationId: z.string().optional().describe("Current service principal application id."),
+    displayName: z.string().optional().describe("Current principal display name."),
+    groups: z
+      .array(
+        z
+          .object({
+            value: z.string().optional().describe("Workspace group identifier."),
+            display: z.string().optional().describe("Workspace group display name."),
+          })
+          .passthrough(),
+      )
+      .optional()
+      .describe("Groups assigned to the current principal."),
+  })
+  .passthrough()
+  .describe("Databricks Current User API response.");
+
+const ExperimentPermissionsResponseSchema = z
+  .object({
+    access_control_list: z
+      .array(
+        z
+          .object({
+            user_name: z.string().optional().describe("User principal name."),
+            group_name: z.string().optional().describe("Group principal name."),
+            service_principal_name: z
+              .string()
+              .optional()
+              .describe("Service principal application id."),
+            all_permissions: z
+              .array(
+                z
+                  .object({
+                    permission_level: z.string().describe("Effective experiment permission level."),
+                  })
+                  .passthrough(),
+              )
+              .optional()
+              .describe("Direct and inherited permissions for this principal."),
+          })
+          .passthrough(),
+      )
+      .optional()
+      .describe("Experiment access control entries."),
+  })
+  .passthrough()
+  .describe("Databricks experiment permissions response.");
+
+function configuredExperiment(): { id?: string; name?: string } {
+  const id = process.env.MLFLOW_EXPERIMENT_ID?.trim();
+  const name = process.env.MLFLOW_EXPERIMENT_NAME?.trim();
+  return {
+    ...(id ? { id } : {}),
+    ...(name ? { name } : {}),
+  };
+}
+
 /**
  * Whether MLflow feedback logging is available for this deployment.
  *
@@ -71,6 +242,111 @@ export function resolveFeedbackEnabled(explicit: boolean | undefined): boolean {
   return explicit ?? mlflowEnabled();
 }
 
+/**
+ * Reject an explicitly enabled feedback surface when no MLflow experiment is
+ * configured. Auto mode remains non-fatal and simply disables feedback.
+ */
+export function validateFeedbackConfig(explicit: boolean | undefined): void {
+  if (explicit !== true) return;
+  const experiment = configuredExperiment();
+  if (experiment.id || experiment.name) return;
+  throw new ConfigurationError(
+    "mastra: feedback is enabled but no MLflow experiment is configured. " +
+      "Set MLFLOW_EXPERIMENT_ID or MLFLOW_EXPERIMENT_NAME, use feedback: false, " +
+      "or omit feedback to use automatic detection.",
+  );
+}
+
+/**
+ * Return the configured experiment URL only when the active Databricks
+ * principal has effective `CAN_MANAGE` permission. Permission lookup is
+ * intentionally best-effort so Apps environments that do not expose the
+ * required APIs simply omit the debug affordance.
+ */
+export async function mlflowExperimentManagerUrl(
+  client: WorkspaceClient,
+): Promise<string | undefined> {
+  try {
+    return await resolveMlflowExperimentManagerUrl(client);
+  } catch {
+    return undefined;
+  }
+}
+
+async function resolveMlflowExperimentManagerUrl(
+  client: WorkspaceClient,
+): Promise<string | undefined> {
+  const experimentId = await resolveExperimentId(client);
+  if (!experimentId) return undefined;
+
+  const [currentResponse, permissionsResponse] = await Promise.all([
+    databricksFetch(client, "/api/2.0/preview/scim/v2/Me", { method: "GET" }),
+    databricksFetch(
+      client,
+      `/api/2.0/permissions/experiments/${encodeURIComponent(experimentId)}`,
+      { method: "GET" },
+    ),
+  ]);
+  if (!currentResponse.ok || !permissionsResponse.ok) return undefined;
+
+  const current = CurrentUserResponseSchema.safeParse(await readResponseJson(currentResponse));
+  const permissions = ExperimentPermissionsResponseSchema.safeParse(
+    await readResponseJson(permissionsResponse),
+  );
+  if (!current.success || !permissions.success) return undefined;
+
+  const userNames = new Set(
+    [current.data.userName, current.data.applicationId, current.data.displayName].filter(
+      (value): value is string => Boolean(value),
+    ),
+  );
+  const groupNames = new Set(
+    (current.data.groups ?? []).flatMap(({ value, display }) =>
+      [value, display].filter((entry): entry is string => Boolean(entry)),
+    ),
+  );
+  const canManage = (permissions.data.access_control_list ?? []).some((entry) => {
+    const matches =
+      (entry.user_name !== undefined && userNames.has(entry.user_name)) ||
+      (entry.service_principal_name !== undefined && userNames.has(entry.service_principal_name)) ||
+      (entry.group_name !== undefined && groupNames.has(entry.group_name));
+    return (
+      matches &&
+      (entry.all_permissions ?? []).some(
+        (permission) => permission.permission_level === "CAN_MANAGE",
+      )
+    );
+  });
+  if (!canManage) return undefined;
+
+  const host = (await client.config.getHost()).toString();
+  return new URL(`/ml/experiments/${encodeURIComponent(experimentId)}`, host).toString();
+}
+
+async function fetchExperiment(
+  client: WorkspaceClient,
+): Promise<{ experiment_id: string; tags?: unknown } | undefined> {
+  const experiment = configuredExperiment();
+  const path = experiment.id
+    ? `/api/2.0/mlflow/experiments/get?experiment_id=${encodeURIComponent(experiment.id)}`
+    : experiment.name
+      ? `/api/2.0/mlflow/experiments/get-by-name?experiment_name=${encodeURIComponent(
+          experiment.name,
+        )}`
+      : undefined;
+  if (!path) return undefined;
+  const response = await databricksFetch(client, path, { method: "GET" });
+  if (!response.ok) return undefined;
+  const parsed = MlflowExperimentResponseSchema.safeParse(await readResponseJson(response));
+  return parsed.success ? parsed.data.experiment : undefined;
+}
+
+async function resolveExperimentId(client: WorkspaceClient): Promise<string | undefined> {
+  const configured = configuredExperiment();
+  if (configured.id) return configured.id;
+  return (await fetchExperiment(client))?.experiment_id;
+}
+
 /** Parameters for {@link logFeedback}. */
 export interface LogFeedbackParams {
   /** MLflow trace id the assessment attaches to (`tr-<hex>`). */
@@ -103,8 +379,12 @@ export async function logFeedback(
     params.name?.trim() ||
     (hasValue ? feedback.DEFAULT_FEEDBACK_NAME : feedback.DEFAULT_COMMENT_NAME);
   const value = hasValue ? params.value : params.comment;
+  const assessmentTraceId = mlflowAssessmentTraceId(
+    params.traceId,
+    await resolveUcTracePrefix(client),
+  );
   const assessment: Record<string, unknown> = {
-    trace_id: params.traceId,
+    trace_id: assessmentTraceId,
     assessment_name: name,
     source: {
       source_type: "HUMAN",
@@ -118,13 +398,14 @@ export async function logFeedback(
   for (let attempt = 0; attempt <= NOT_FOUND_RETRIES; attempt++) {
     let res: Response;
     try {
-      res = await databricksFetch(client, assessmentsPath(params.traceId), {
+      res = await databricksFetch(client, assessmentsPath(assessmentTraceId), {
         method: "POST",
         body,
       });
     } catch (err) {
       logger.warn("feedback request failed", {
         traceId: params.traceId,
+        assessmentTraceId,
         error: errorUtils.errorMessage(err),
       });
       return undefined;
@@ -144,6 +425,7 @@ export async function logFeedback(
     }
     logger.warn("feedback not recorded", {
       traceId: params.traceId,
+      assessmentTraceId,
       status: res.status,
       body: await readResponseText(res),
     });

package/src/plugin.ts

@@ -99,7 +99,12 @@ import {
 } from "./identity.ts";
 import { buildMcpServer, type ResolvedMcp } from "./mcp.ts";
 import { createMemoryBuilder, createServicePrincipalPool, needsLakebase } from "./memory.ts";
-import { logFeedback, resolveFeedbackEnabled } from "./mlflow.ts";
+import {
+  logFeedback,
+  mlflowExperimentManagerUrl,
+  resolveFeedbackEnabled,
+  validateFeedbackConfig,
+} from "./mlflow.ts";
 import { resolveDefaultModelId } from "./model.ts";
 import { buildObservability, configureOtelPropagation } from "./observability.ts";
 import { provisionRemoteSkills } from "./remote-skills.ts";
@@ -249,6 +254,7 @@ export class MastraPlugin extends Plugin<MastraPluginConfig> {
     // Resolve the identity mode up front so an invalid `genieIdentity` fails the
     // app boot with a clear message rather than 500-ing the first chat request.
     this.identityMode = resolveIdentityMode(this.config.genieIdentity);
+    validateFeedbackConfig(this.config.feedback);
     // Wait until sibling plugins (e.g. `lakebase`) finish `setup()` so
     // the lakebase pool is valid when storage/memory are enabled.
     this.context?.onLifecycle("setup:complete", async () => {
@@ -466,6 +472,7 @@ export class MastraPlugin extends Plugin<MastraPluginConfig> {
     this.registerEmbedRoutes(router);
     this.registerSuggestionRoutes(router);
     this.registerFeedbackRoutes(router);
+    this.registerMlflowExperimentRoute(router);
     this.registerAgentRoutes(router);
   }
 
@@ -770,6 +777,23 @@ export class MastraPlugin extends Plugin<MastraPluginConfig> {
     });
   }
 
+  /** Register the best-effort per-viewer MLflow experiment debug link. */
+  private registerMlflowExperimentRoute(router: IAppRouter): void {
+    this.route(router, {
+      name: "mlflowExperiment",
+      method: "get",
+      path: routes.MASTRA_ROUTES.mlflowExperiment,
+      handler: async (req, res) => {
+        if (!this.feedbackEnabled()) {
+          res.json({ url: null });
+          return;
+        }
+        const result = await this.scopedSelf(req).fetchMlflowExperimentUrl();
+        res.json({ url: result.ok ? (result.data ?? null) : null });
+      },
+    });
+  }
+
   /** Register the gated Mastra catch-all after every explicit route. */
   private registerAgentRoutes(router: IAppRouter): void {
     // Middleware rather than `this.route`: this is the catch-all that hands
@@ -917,6 +941,15 @@ export class MastraPlugin extends Plugin<MastraPluginConfig> {
     );
   }
 
+  /**
+   * Resolve an experiment link under the active request identity. The helper
+   * returns no URL unless the viewer has effective `CAN_MANAGE` permission.
+   */
+  private async fetchMlflowExperimentUrl(): Promise<ExecutionResult<string | undefined>> {
+    const client = getExecutionContext().client;
+    return this.execute(() => mlflowExperimentManagerUrl(client), feedbackWriteDefaults);
+  }
+
   /**
    * Implementation backing the `data` embed resolver
    * (`GET /embed/data/:id`). Runs inside the AppKit user-context proxy so

package/src/remote-skills.ts

@@ -567,7 +567,7 @@ function resolveDatabricksBasePath(
   client: WorkspaceClientLike | undefined,
 ): string | undefined {
   if (!client) return undefined;
-  if (options.databricksBasePath) return options.databricksBasePath.trim() || undefined;
+  if (options.databricksBasePath) return stringUtils.trimToUndefined(options.databricksBasePath);
   const email = stringUtils.trimToNull(options.userEmail);
   // A named user targets their personal Assistant tree (the "save a skill"
   // target); otherwise the shared workspace Assistant tree, which the built-in

package/src/server.ts

@@ -276,6 +276,7 @@ export class MastraServer extends ExpressMastraServer {
    */
   configureMlflowTraceId(res: express.Response) {
     if (!this.feedbackEnabled || res.headersSent) return;
+    if (res.getHeader(feedback.MLFLOW_TRACE_ID_HEADER)) return;
     const traceId = trace.getActiveSpan()?.spanContext().traceId;
     if (!traceId || traceId === INVALID_TRACE_ID) return;
     res.setHeader(feedback.MLFLOW_TRACE_ID_HEADER, `tr-${traceId}`);

package/src/telemetry.ts

@@ -20,11 +20,14 @@
 
 import { StringDecoder } from "node:string_decoder";
 
-import { json, log, object } 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, SpanKind, trace, type Span } from "@opentelemetry/api";
 import { getRPCMetadata, RPCType } from "@opentelemetry/core";
 import type express from "express";
 
+import { USER_EMAIL_HEADER, USER_ID_HEADER } from "./identity.ts";
+
 const logger = log.logger("mastra/telemetry");
 
 /** Cap on each payload copied onto a span, so one turn cannot bloat the export. */
@@ -56,6 +59,34 @@ export const CHAT_RESPONSE_ATTR = "appkit.mastra.chat.response";
 export const CHAT_IDENTITY_ATTR = "appkit.mastra.identity.mode";
 /** Whether the streamed turn emitted a Genie tool or progress event. */
 export const CHAT_GENIE_USED_ATTR = "appkit.mastra.genie.used";
+/**
+ * MLflow span-type attribute. Databricks maps this (or `gen_ai.operation.name`)
+ * onto the type chip in the experiment UI. Custom values such as `GENIE` are
+ * allowed.
+ */
+export const MLFLOW_SPAN_TYPE_ATTR = "mlflow.spanType";
+/** GenAI semantic-convention operation name, an alternate span-type source. */
+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";
+/** Root-span session id that enables the MLflow Sessions tab. */
+export const MLFLOW_SESSION_ATTR = "session.id";
+/**
+ * Prefix Databricks/MLflow copies onto the unified-view `tags` map. The
+ * suffix after this prefix is the tag name shown in the experiment UI.
+ */
+export const MLFLOW_TRACE_TAG_PREFIX = "mlflow.trace.tag.";
+/** Experiment-UI tag marking a turn that called Genie. */
+export const MLFLOW_GENIE_TAG = "genie";
+/** Root-span type for a Mastra agent turn. */
+export const MLFLOW_SPAN_TYPE_AGENT = "AGENT";
+/** Span type for an `ask_genie` tool call. */
+export const MLFLOW_SPAN_TYPE_GENIE = "GENIE";
+/** Attribute key for the `genie` trace tag. */
+export const MLFLOW_GENIE_TAG_ATTR = `${MLFLOW_TRACE_TAG_PREFIX}${MLFLOW_GENIE_TAG}`;
 
 /** Root-span metadata known before a chat request is dispatched. */
 export interface ChatTurnTelemetryOptions {
@@ -68,6 +99,57 @@ interface TraceTarget {
   owned: boolean;
 }
 
+/** OpenTelemetry's sentinel for "no valid trace" - 32 zero hex chars. */
+const INVALID_TRACE_ID = "0".repeat(32);
+
+/**
+ * Publish the turn's MLflow trace id before the handler writes the body.
+ *
+ * The chat UI only shows thumbs when this header is present. Stamping it here
+ * (on the same span the middleware just resolved) avoids relying on async
+ * Express continuation still sitting inside `context.with()`.
+ */
+/**
+ * Mark the active span as a Genie call so the experiment UI types it `GENIE`
+ * and tags the trace `genie=true`. Safe no-op when no recording span is active.
+ */
+export function stampGenieToolSpan(span: Span | undefined = trace.getActiveSpan()): void {
+  if (!span?.isRecording()) return;
+  span.setAttribute(MLFLOW_SPAN_TYPE_ATTR, MLFLOW_SPAN_TYPE_GENIE);
+  span.setAttribute(GEN_AI_OPERATION_NAME_ATTR, "execute_tool");
+  span.setAttribute(MLFLOW_GENIE_TAG_ATTR, "true");
+}
+
+function requestHeader(req: express.Request, name: string): string | undefined {
+  const fromMethod =
+    typeof req.header === "function" ? stringUtils.trimToNull(req.header(name)) : undefined;
+  if (fromMethod) return fromMethod;
+  const raw = req.headers?.[name.toLowerCase()];
+  const value = Array.isArray(raw) ? raw[0] : raw;
+  return stringUtils.trimToNull(typeof value === "string" ? value : undefined) ?? undefined;
+}
+
+function requestQuery(req: express.Request, name: string): string | undefined {
+  const raw = req.query?.[name];
+  const value = Array.isArray(raw) ? raw[0] : raw;
+  return stringUtils.trimToNull(typeof value === "string" ? value : undefined) ?? undefined;
+}
+
+function stampMlflowActor(span: Span, req: express.Request): void {
+  const user = requestHeader(req, USER_EMAIL_HEADER) ?? requestHeader(req, USER_ID_HEADER);
+  if (user) span.setAttribute(MLFLOW_USER_ATTR, user);
+  const session =
+    requestHeader(req, thread.THREAD_ID_HEADER) ?? requestQuery(req, thread.THREAD_ID_QUERY);
+  if (session) span.setAttribute(MLFLOW_SESSION_ATTR, session);
+}
+
+function stampMlflowTraceId(res: express.Response, span: Span): void {
+  if (res.headersSent) return;
+  const traceId = span.spanContext().traceId;
+  if (!traceId || traceId === INVALID_TRACE_ID) return;
+  res.setHeader(feedback.MLFLOW_TRACE_ID_HEADER, `tr-${traceId}`);
+}
+
 /** Return one raw prompt when the request is exactly one text-only user message. */
 export function textOnlyChatInput(messages: unknown): string | undefined {
   if (typeof messages === "string") return messages || undefined;
@@ -276,6 +358,10 @@ export function chatTurnTelemetryMiddleware(
     next();
     return;
   }
+  stampMlflowTraceId(res, target.span);
+  stampMlflowActor(target.span, req);
+  target.span.setAttribute(MLFLOW_SPAN_TYPE_ATTR, MLFLOW_SPAN_TYPE_AGENT);
+  target.span.setAttribute(GEN_AI_OPERATION_NAME_ATTR, "invoke_agent");
   const identity =
     typeof options.identity === "function" ? options.identity(req) : options.identity;
   if (identity) {
@@ -311,7 +397,9 @@ export function chatTurnTelemetryMiddleware(
     fullOutputRecorded = true;
   };
   const recordGenie = (used: boolean): void => {
-    if (used) target.span.setAttribute(CHAT_GENIE_USED_ATTR, true);
+    if (!used) return;
+    target.span.setAttribute(CHAT_GENIE_USED_ATTR, true);
+    target.span.setAttribute(MLFLOW_GENIE_TAG_ATTR, "true");
   };
   const endOwnedSpan = (): void => {
     if (!target.owned || ownedSpanEnded) return;

package/src/validation.ts

@@ -9,6 +9,7 @@
  *
  * @module
  */
+import { object } from "@dbx-tools/shared-core";
 
 /** The subset of a failed parse result these helpers read. */
 export interface SchemaIssues {
@@ -17,6 +18,11 @@ export interface SchemaIssues {
 
 /** Distinct dot-joined field paths a failed parse flagged. */
 export function invalidFields(error: SchemaIssues): string[] {
-  const fields = error.issues.map((issue) => issue.path.map(String).join(".")).filter(Boolean);
-  return [...new Set(fields)];
+  return [
+    ...object
+      .sequence(error.issues)
+      .map((issue) => issue.path.map(String).join("."))
+      .filter(Boolean)
+      .distinct(),
+  ];
 }