@dbx-tools/appkit-mastra 0.9.60 → 0.9.64

package/src/genie.ts

@@ -59,6 +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 { safeWriteProgress } from "./writer.ts";
 
 const logger = log.logger("mastra/genie");
@@ -428,6 +429,7 @@ function buildAskGenieTool(opts: {
     }),
     execute: async ({ question }, ctxRaw) => {
       const ctx = ctxRaw as ToolExecuteCtx;
+      stampGenieToolSpan();
       const { client, requestContext } = requireClient(ctx, toolId);
       const writer = ctx?.writer;
       const toolCallId = ctx?.agent?.toolCallId;

package/src/mlflow.ts

@@ -13,17 +13,21 @@
  *
  * 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";
@@ -33,6 +37,101 @@ 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`;
@@ -47,6 +146,7 @@ const MlflowExperimentResponseSchema = z
     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."),
@@ -226,18 +326,28 @@ async function resolveMlflowExperimentManagerUrl(
   return new URL(`/ml/experiments/${encodeURIComponent(experimentId)}`, host).toString();
 }
 
-async function resolveExperimentId(client: WorkspaceClient): Promise<string | undefined> {
+async function fetchExperiment(
+  client: WorkspaceClient,
+): Promise<{ experiment_id: string; tags?: unknown } | undefined> {
   const experiment = configuredExperiment();
-  if (experiment.id) return experiment.id;
-  if (!experiment.name) return undefined;
-
-  const path = `/api/2.0/mlflow/experiments/get-by-name?experiment_name=${encodeURIComponent(
-    experiment.name,
-  )}`;
+  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.experiment_id : undefined;
+  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}. */
@@ -272,8 +382,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",
@@ -287,13 +401,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;
@@ -313,6 +428,7 @@ export async function logFeedback(
     }
     logger.warn("feedback not recorded", {
       traceId: params.traceId,
+      assessmentTraceId,
       status: res.status,
       body: await readResponseText(res),
     });

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/skill-cache.ts

@@ -0,0 +1,312 @@
+/**
+ * AppKit-backed caching for Databricks workspace skill mounts.
+ *
+ * @module
+ */
+
+import { randomUUID } from "node:crypto";
+import { CacheManager } from "@databricks/appkit";
+import { hash, object } from "@dbx-tools/shared-core";
+import type {
+  CopyOptions,
+  FileContent,
+  FileEntry,
+  FileStat,
+  FileSystem,
+  ListOptions,
+  MakeDirectoryOptions,
+  ReadFileOptions,
+  RemoveOptions,
+  WriteFileOptions,
+} from "@dbx-tools/shared-fs";
+
+import { filesystems, type MastraFileSystemAdapter } from "./filesystems.ts";
+
+const CACHE_NAMESPACE = "mastra:workspace-skills";
+const DEFAULT_MAX_SCOPED_MOUNTS = 256;
+
+interface CachedFileStat extends Omit<FileStat, "accessedAt" | "createdAt" | "modifiedAt"> {
+  accessedAt?: number;
+  createdAt?: number;
+  modifiedAt?: number;
+}
+
+interface ScopedMountEntry {
+  adapter: MastraFileSystemAdapter;
+  cache: AppKitCachedFileSystem;
+  host: string;
+  root: string;
+  userKey: string;
+}
+
+const scopedMounts = new Map<string, ScopedMountEntry>();
+
+/** Default remote skill metadata and content cache lifetime. */
+export const DEFAULT_WORKSPACE_SKILL_CACHE_TTL_MS = 5 * 60 * 1000;
+
+/** Options for one AppKit-cached workspace skill filesystem. */
+export interface AppKitCachedFileSystemOptions {
+  host: string;
+  source: FileSystem;
+  ttlMs?: number;
+  userKey: string;
+}
+
+/** Stable Mastra mount plus the identity key used for source reuse. */
+export interface CachedWorkspaceSkillMount {
+  cacheKey: string;
+  filesystem: MastraFileSystemAdapter;
+}
+
+/**
+ * Cache read-only filesystem operations through AppKit while delegating
+ * mutations to the current request's source.
+ */
+export class AppKitCachedFileSystem implements FileSystem<"appkit-cache"> {
+  readonly backend = "appkit-cache" as const;
+  readonly id: string;
+  readonly readOnly: boolean;
+  readonly root: string;
+
+  private source: FileSystem;
+  private readonly host: string;
+  private readonly ttlSec: number;
+  private readonly userKey: string;
+
+  constructor(options: AppKitCachedFileSystemOptions) {
+    this.host = options.host;
+    this.source = options.source;
+    this.root = options.source.root;
+    this.readOnly = options.source.readOnly;
+    this.ttlSec = cacheTtlSec(options.ttlMs);
+    this.userKey = options.userKey;
+    this.id = `appkit-skill-cache-${hash.fnvHash(this.host, this.userKey, this.root)}`;
+  }
+
+  /** Rebind cache misses and writes to the current request's OBO filesystem. */
+  bind(source: FileSystem): void {
+    if (source.root !== this.root || source.readOnly !== this.readOnly) {
+      throw new Error("Cannot rebind a workspace skill cache to a different filesystem scope");
+    }
+    this.source = source;
+  }
+
+  async init(): Promise<void> {
+    await this.source.init();
+  }
+
+  /** Keep the shared scope alive when one request releases its workspace view. */
+  close(): Promise<void> {
+    return Promise.resolve();
+  }
+
+  resolvePath(inputPath: string): string {
+    return this.source.resolvePath(inputPath);
+  }
+
+  async readFile(inputPath: string): Promise<Uint8Array>;
+  async readFile(
+    inputPath: string,
+    options: ReadFileOptions & { encoding: string },
+  ): Promise<string>;
+  async readFile(inputPath: string, options?: ReadFileOptions): Promise<string | Uint8Array> {
+    const value = await this.cached("readFile", inputPath, undefined, async (source) => ({
+      base64: Buffer.from(await source.readFile(inputPath)).toString("base64"),
+    }));
+    const bytes = Buffer.from(value.base64, "base64");
+    return options?.encoding ? bytes.toString(options.encoding as BufferEncoding) : bytes;
+  }
+
+  async writeFile(
+    inputPath: string,
+    content: FileContent,
+    options?: WriteFileOptions,
+  ): Promise<void> {
+    await this.mutate((source) => source.writeFile(inputPath, content, options));
+  }
+
+  async appendFile(inputPath: string, content: FileContent): Promise<void> {
+    await this.mutate((source) => source.appendFile(inputPath, content));
+  }
+
+  async deleteFile(inputPath: string, options?: RemoveOptions): Promise<void> {
+    await this.mutate((source) => source.deleteFile(inputPath, options));
+  }
+
+  async copyFile(
+    sourcePath: string,
+    destinationPath: string,
+    options?: CopyOptions,
+  ): Promise<void> {
+    await this.mutate((source) => source.copyFile(sourcePath, destinationPath, options));
+  }
+
+  async moveFile(
+    sourcePath: string,
+    destinationPath: string,
+    options?: CopyOptions,
+  ): Promise<void> {
+    await this.mutate((source) => source.moveFile(sourcePath, destinationPath, options));
+  }
+
+  async mkdir(inputPath: string, options?: MakeDirectoryOptions): Promise<void> {
+    await this.mutate((source) => source.mkdir(inputPath, options));
+  }
+
+  async rmdir(inputPath: string, options?: RemoveOptions): Promise<void> {
+    await this.mutate((source) => source.rmdir(inputPath, options));
+  }
+
+  async readdir(inputPath: string, options?: ListOptions): Promise<FileEntry[]> {
+    return this.cached("readdir", inputPath, options, (source) =>
+      source.readdir(inputPath, options),
+    );
+  }
+
+  async exists(inputPath: string): Promise<boolean> {
+    return this.cached("exists", inputPath, undefined, (source) => source.exists(inputPath));
+  }
+
+  async stat(inputPath: string): Promise<FileStat> {
+    const value = await this.cached("stat", inputPath, undefined, async (source) =>
+      serializeStat(await source.stat(inputPath)),
+    );
+    return deserializeStat(value);
+  }
+
+  /** Rotate this scope's cache generation so the next read reloads it. */
+  async invalidate(): Promise<void> {
+    const cache = await CacheManager.getInstance();
+    await cache.set(this.generationKey(cache), randomUUID(), {
+      ttl: this.ttlSec * 2,
+    });
+  }
+
+  private async cached<T>(
+    operation: string,
+    inputPath: string,
+    options: unknown,
+    load: (source: FileSystem) => Promise<T>,
+  ): Promise<T> {
+    const source = this.source;
+    const cache = await CacheManager.getInstance();
+    const generation = await this.generation(cache);
+    return cache.getOrExecute(
+      [
+        CACHE_NAMESPACE,
+        this.host,
+        this.root,
+        generation,
+        operation,
+        source.resolvePath(inputPath),
+        object.toStableKey(options ?? {}),
+      ],
+      () => load(source),
+      this.userKey,
+      { ttl: this.ttlSec },
+    );
+  }
+
+  private async generation(cache: CacheManager): Promise<string> {
+    return cache.getOrExecute(
+      [CACHE_NAMESPACE, "generation", this.host, this.root],
+      async () => randomUUID(),
+      this.userKey,
+      { ttl: this.ttlSec * 2 },
+    );
+  }
+
+  private generationKey(cache: CacheManager): string {
+    return cache.generateKey([CACHE_NAMESPACE, "generation", this.host, this.root], this.userKey);
+  }
+
+  private async mutate(run: (source: FileSystem) => Promise<void>): Promise<void> {
+    const source = this.source;
+    await run(source);
+    await this.invalidate();
+  }
+}
+
+/** Return one stable Mastra mount per resolved user, host, and source root. */
+export function cachedWorkspaceSkillMount(
+  options: AppKitCachedFileSystemOptions,
+): CachedWorkspaceSkillMount {
+  const cacheKey = object.toStableKey({
+    host: options.host,
+    readOnly: options.source.readOnly,
+    root: options.source.root,
+    userKey: options.userKey,
+  });
+  const existing = scopedMounts.get(cacheKey);
+  if (existing) {
+    existing.cache.bind(options.source);
+    scopedMounts.delete(cacheKey);
+    scopedMounts.set(cacheKey, existing);
+    return { cacheKey, filesystem: existing.adapter };
+  }
+
+  const cache = new AppKitCachedFileSystem(options);
+  const entry: ScopedMountEntry = {
+    adapter: filesystems(cache, {
+      id: cache.id,
+      readOnly: options.source.readOnly,
+    }),
+    cache,
+    host: options.host,
+    root: options.source.root,
+    userKey: options.userKey,
+  };
+  scopedMounts.set(cacheKey, entry);
+  trimScopedMounts();
+  return { cacheKey, filesystem: entry.adapter };
+}
+
+/** Invalidate one user's cached workspace skill scope. */
+export async function clearWorkspaceSkillCache(options: {
+  host: string;
+  root?: string;
+  userKey: string;
+}): Promise<void> {
+  const matches = [...scopedMounts.entries()].filter(
+    ([, entry]) =>
+      entry.host === options.host &&
+      entry.userKey === options.userKey &&
+      (options.root === undefined || entry.root === options.root),
+  );
+  await Promise.all(matches.map(([, entry]) => entry.cache.invalidate()));
+}
+
+function cacheTtlSec(ttlMs = DEFAULT_WORKSPACE_SKILL_CACHE_TTL_MS): number {
+  if (!Number.isFinite(ttlMs) || ttlMs <= 0) {
+    throw new Error("Workspace skill cache ttlMs must be a positive finite number");
+  }
+  return Math.max(1, Math.ceil(ttlMs / 1000));
+}
+
+function serializeStat(value: FileStat): CachedFileStat {
+  const { accessedAt, createdAt, modifiedAt, ...stat } = value;
+  return {
+    ...stat,
+    ...(accessedAt ? { accessedAt: accessedAt.getTime() } : {}),
+    ...(createdAt ? { createdAt: createdAt.getTime() } : {}),
+    ...(modifiedAt ? { modifiedAt: modifiedAt.getTime() } : {}),
+  };
+}
+
+function deserializeStat(value: CachedFileStat): FileStat {
+  const { accessedAt, createdAt, modifiedAt, ...stat } = value;
+  return {
+    ...stat,
+    ...(accessedAt !== undefined ? { accessedAt: new Date(accessedAt) } : {}),
+    ...(createdAt !== undefined ? { createdAt: new Date(createdAt) } : {}),
+    ...(modifiedAt !== undefined ? { modifiedAt: new Date(modifiedAt) } : {}),
+  };
+}
+
+function trimScopedMounts(): void {
+  while (scopedMounts.size > DEFAULT_MAX_SCOPED_MOUNTS) {
+    const oldest = scopedMounts.keys().next().value;
+    if (oldest === undefined) return;
+    scopedMounts.delete(oldest);
+  }
+}

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;