@dbx-tools/appkit-mastra 0.6.14 → 0.6.27

package/src/genie.ts

@@ -48,14 +48,14 @@ import { plugin } from "@dbx-tools/appkit";
 import { chat, space as genieSpace } from "@dbx-tools/genie";
 import { error, log, string } from "@dbx-tools/shared-core";
 import { genieModel, type GenieMessage } from "@dbx-tools/shared-genie";
-import { wire, type MastraWriter, type StartedEvent } from "@dbx-tools/shared-mastra";
+import type { MastraWriter, StartedEvent } from "@dbx-tools/shared-mastra";
 import type { RequestContext } from "@mastra/core/request-context";
 import { MASTRA_THREAD_ID_KEY } from "@mastra/core/request-context";
 import { createTool } from "@mastra/core/tools";
 import { z } from "zod";
 
 import type { MastraTools } from "./agents.ts";
-import { chartPlannerRequestSchema, prepareChart } from "./chart.ts";
+import { chartPlannerRequestSchema, chartToolOutputSchema, prepareChart } from "./chart.ts";
 import { MASTRA_USER_KEY, resolveUserKey } from "./config.ts";
 import type { MastraPluginConfig, User } from "./config.ts";
 import { fetchStatementData } from "./statement.ts";
@@ -687,9 +687,9 @@ function buildGetStatementTool() {
  * {@link prepareChart} that resolves the dataset by fetching the
  * Genie statement's rows on demand. The tool mints a `chartId`
  * synchronously, caches an empty placeholder, and kicks off the
- * planner in the background so the agent loop never blocks. The
- * host UI resolves `[chart:<chartId>]` markers by reading the
- * cached {@link Chart} entry (1h TTL).
+ * planner in the background so the agent loop never blocks. The result carries
+ * the complete marker to copy; the host UI resolves it by reading the cached
+ * {@link Chart} entry (1h TTL).
  *
  * Space-agnostic: a Genie `statement_id` is workspace-scoped, so
  * one shared `prepare_chart` tool covers every wired Genie space.
@@ -709,20 +709,21 @@ function buildPrepareChartTool(opts: { config: MastraPluginConfig }) {
     description: string.toDescription([
       `
         Queue a chart for the rows of a Genie statement. Mints a
-        short \`chartId\` synchronously and kicks off a BACKGROUND
+        \`chartId\` plus its complete \`marker\` synchronously and
+        kicks off a BACKGROUND
         task that fetches the statement's rows, runs the
         chart-planner to pick a chart type and Echarts spec, and
         caches the result under the \`chartId\` for one hour. The
         host UI fetches the cached chart on its own once it lands.
       `,
       `
-        To display the chart in your reply, embed
-        \`[chart:<chartId>]\` on its own line at the position you
-        want it to appear, using the EXACT \`chartId\` string this
-        call returned. Never construct a chart id yourself (it is
-        not the \`statement_id\` or any variation of it) - only a
-        value returned by this tool resolves to a real chart. The
-        tool returns immediately - do NOT wait or call it again to
+        To display the chart in your reply, copy the returned
+        \`marker\` VERBATIM onto its own line at the position you
+        want it to appear. Never construct, alter, or invent a
+        marker from \`chartId\` (and never use the \`statement_id\`
+        or any variation of it) - only the complete value returned
+        by this tool resolves to a real chart. The tool returns
+        immediately - do NOT wait or call it again to
         "check progress"; the chart resolves asynchronously on the
         host UI's side.
       `,
@@ -734,7 +735,7 @@ function buildPrepareChartTool(opts: { config: MastraPluginConfig }) {
       `,
     ]),
     inputSchema: prepareChartRequestSchema,
-    outputSchema: wire.ChartSchema.pick({ chartId: true }),
+    outputSchema: chartToolOutputSchema,
     execute: async (request, ctxRaw) => {
       const ctx = ctxRaw as ToolExecuteCtx;
       const { client, requestContext } = requireClient(ctx, toolId);
@@ -835,20 +836,20 @@ export const GENIE_INSTRUCTIONS = string.toDescription([
             `,
             `
               \`[chart:<chartId>]\` - render the rows as a chart. To
-              get a \`<chartId>\`, call \`prepare_chart\` with the
+              get one, call \`prepare_chart\` with the
               statement's id (and an optional \`title\` / one-line
               \`description\` of the insight to surface). The tool
-              returns the \`chartId\` synchronously and prepares the
-              chart spec in the background; embed the returned id as
-              \`[chart:<chartId>]\` on its own line wherever the
-              chart should appear. Use a chart when the data has a
+              returns the complete \`marker\` synchronously and
+              prepares the chart spec in the background; copy that
+              marker VERBATIM onto its own line wherever the chart
+              should appear. Use a chart when the data has a
               story a visual conveys better than a table (trends,
               rankings, distributions, parts-of-a-whole).
 
-              NEVER invent or hand-build a \`<chartId>\`. A valid
-              \`<chartId>\` is the opaque token a \`prepare_chart\`
-              call returned to you in THIS turn - nothing else. It
-              is NOT a \`statement_id\`, and it is NOT a
+              NEVER invent or hand-build a marker. A valid marker is
+              the complete opaque string a \`prepare_chart\` call
+              returned to you in THIS turn - nothing else. Its id is
+              NOT a \`statement_id\`, and it is NOT a
               \`statement_id\` prefix with a label appended (e.g.
               \`01f1...-region-fill\`). If you have not called
               \`prepare_chart\` and received an id back, do not write

package/src/mlflow.ts

@@ -5,7 +5,9 @@
  *
  * Feedback attaches to a trace, and the plugin's spans reach MLflow
  * through the same OTel pipeline as every other AppKit span (see
- * `observability.ts`). MLflow derives its trace id from the OpenTelemetry
+ * `observability.ts`). On Databricks Apps that pipeline is the UC
+ * sidecar injected by `telemetry_export_destinations`, not a direct
+ * workspace OTLP URL. MLflow derives its trace id from the OpenTelemetry
  * trace id (`tr-<hex(otelTraceId)>`), so the server stamps the active
  * trace id on each turn's response and the client sends it back here.
  *

package/src/observability.ts

@@ -30,6 +30,15 @@
  * registered at all (unless `observability: true` forces it on), so
  * Mastra does not emit `[OtelBridge] No OTEL span found` warnings.
  *
+ * On Databricks Apps the supported sink is Unity Catalog via
+ * `telemetry_export_destinations` (the platform injects the local OTLP
+ * sidecar). Managed MLflow has no OTLP ingest endpoint - do not point
+ * `OTEL_EXPORTER_OTLP_*` at the workspace host. Set
+ * `OTEL_PROPAGATORS=none` so Apps ingress `traceparent` does not hide
+ * every request-scoped root span from the UC `*_trace_unified` view.
+ * Chat request/response on that root span is stamped by
+ * `traceIo.attachChatTurnTraceIo` (see `trace-io.ts`).
+ *
  * @module
  */
 
@@ -39,6 +48,7 @@ import { Observability } from "@mastra/observability";
 import { OtelBridge } from "@mastra/otel-bridge";
 
 import { TRACE_REQUEST_CONTEXT_KEYS } from "./config.ts";
+import { mlflowEnabled } from "./mlflow.ts";
 
 const logger = log.logger("mastra/observability");
 
@@ -127,11 +137,17 @@ export async function buildObservability(
     : otelBase
       ? `${otelBase.replace(/\/+$/, "")}/v1/traces`
       : undefined;
+  // `mlflowEnabled` is the same gate that flips chat feedback on, so the
+  // boot line is the one-glance check that both halves of the Apps UC
+  // pipeline (OTLP sidecar + experiment id) are present.
+  const feedback = mlflowEnabled();
   logger.info("Mastra observability wired through OTel bridge", {
     serviceName,
     requestContextKeys,
     otelBase: otelBase ?? "<unset>",
     resolvedTracesUrl: resolvedTracesUrl ?? "<unset>",
+    feedback,
+    observability: feedback ? "mlflow" : "otel",
   });
 
   return new Observability({

package/src/plugin.ts

@@ -106,6 +106,7 @@ import {
 import { resolveServingConfig } from "./serving.ts";
 import { fetchStatementData, STATEMENT_ROW_CAP } from "./statement.ts";
 import { threadsRoute } from "./threads.ts";
+import { attachChatTurnTraceIo } from "./trace-io.ts";
 import { invalidFields } from "./validation.ts";
 
 const GENIE_MANIFEST = plugin.data(genie).plugin.manifest;
@@ -594,10 +595,24 @@ export class MastraPlugin extends Plugin<MastraPluginConfig> {
         const id = string.firstNonEmpty(req.params.id);
         const resolve = embedResolvers[type];
         if (!resolve) {
+          this.logger.warn("embed:unsupported", {
+            status: 404,
+            type,
+            id,
+            method: req.method,
+            path: req.path,
+          });
           res.status(404).json({ error: `unsupported embed type: ${type}` });
           return;
         }
         if (!id) {
+          this.logger.warn("embed:invalid", {
+            status: 400,
+            type,
+            method: req.method,
+            path: req.path,
+            error: "id is required",
+          });
           res.status(400).json({ error: "id is required" });
           return;
         }
@@ -615,6 +630,18 @@ export class MastraPlugin extends Plugin<MastraPluginConfig> {
           return;
         }
         if (result.data === undefined) {
+          // A cache miss is security-deliberately ambiguous: unknown, expired,
+          // or owned by another identity all map to the same 404. Log the
+          // opaque id and those possible causes so a fabricated marker is
+          // diagnosable from the server without disclosing which cause applied.
+          this.logger.warn(`embed:${type}:not-found`, {
+            status: 404,
+            type,
+            id,
+            method: req.method,
+            path: req.path,
+            reason: "unknown, expired, or owned by another identity",
+          });
           res.status(404).json({ error: `${type} not found` });
           return;
         }
@@ -1004,7 +1031,7 @@ export class MastraPlugin extends Plugin<MastraPluginConfig> {
     // every default-workspace agent via `extraSkillPaths`.
     const provisioned = await provisionRemoteSkills(this.config.remoteSkills);
     if (provisioned.skillNames.length > 0) {
-      this.logger.info("remote skills provisioned", {
+      this.logger.info("remote skills configured", {
         skills: provisioned.skillNames,
         databricksBasePath: provisioned.databricksBasePath,
         localSkillPaths: provisioned.localSkillPaths.length,
@@ -1061,6 +1088,12 @@ export class MastraPlugin extends Plugin<MastraPluginConfig> {
     });
     this.mastraApp = express();
     attachRoutePatchMiddleware(this.mastraApp);
+    // Stamp chat request/response onto the HTTP root span so MLflow's
+    // UC `*_trace_unified` view can surface them (Mastra's own
+    // `mastra.agent_run.*` attributes sit on a child span the view
+    // never reads). Must run before `MastraServer.init` so the layer
+    // sits ahead of the agent routes.
+    attachChatTurnTraceIo(this.mastraApp);
     this.mastraServer = new MastraServer(this.config, {
       app: this.mastraApp,
       mastra: this.mastra,

package/src/remote-skills.ts

@@ -61,15 +61,9 @@ import { error, hash, json, log, net, object, string } from "@dbx-tools/shared-c
 import type { OneOrMany } from "@dbx-tools/shared-core";
 import type { FileSystem } from "@dbx-tools/shared-fs";
 
-const logger = log.logger("mastra/remote-skills");
-
-/** Shared Assistant skills tree in the Databricks workspace (default target). */
-const ASSISTANT_SHARED_SKILLS_PATH = "/Workspace/.assistant/skills";
+import { ASSISTANT_SHARED_SKILLS_PATH, userAssistantSkillsPath } from "./skill-paths.ts";
 
-/** Assistant skills directory for a specific Databricks user. */
-function userAssistantSkillsPath(userEmail: string): string {
-  return `/Users/${userEmail.trim()}/.assistant/skills`;
-}
+const logger = log.logger("mastra/remote-skills");
 
 /**
  * Agent id the `skills` CLI installs a bare `SKILL.md` tree under.
@@ -452,14 +446,19 @@ export async function provisionRemoteSkills(
         if (cached && isFresh(cached, resolveRefreshTtl(sourceOptions, options))) {
           skillNames.push(...cached.skills);
           if (!destination) localSkillPaths.push(cacheFS.root);
-          logger.debug("source:cached", {
+          logger.info("remote skill ready", {
             source: sourceOptions.source,
+            destination: databricksBasePath ?? cacheFS.root,
             downloadedAt: cached.downloadedAt,
             skills: cached.skills,
+            cached: true,
           });
           continue;
         }
 
+        logger.info("installing remote skill", {
+          source: sourceOptions.source,
+        });
         staging ??= await initializedScratch("mastra-remote-skills");
         const stagedDir = await stageSource(sourceOptions, staging.root, options);
         const staged = await collectSkillDirs(stagedDir);
@@ -482,9 +481,9 @@ export async function provisionRemoteSkills(
           localSkillPaths.push(await persistLocally(key, staged, record));
         }
         skillNames.push(...record.skills);
-        logger.debug("source:provisioned", {
+        logger.info("remote skill installed", {
           source: sourceOptions.source,
-          destination: databricksBasePath ?? "local-temp",
+          destination: databricksBasePath ?? localSkillPaths.at(-1),
           skills: record.skills,
         });
       } catch (err) {

package/src/skill-paths.ts

@@ -0,0 +1,19 @@
+/**
+ * Where Databricks Assistant `SKILL.md` trees live.
+ *
+ * Two modules need these paths and must agree on them: `workspaces.ts` MOUNTS
+ * them per request (see `DEFAULT_SKILL_FOLDERS`), and `remote-skills.ts`
+ * WRITES provisioned skills into them at startup. They were previously spelled
+ * out in both, so a change to one silently provisioned skills into a tree the
+ * other never scanned.
+ *
+ * @module
+ */
+
+/** Shared Assistant skills tree, readable by everyone in the workspace. */
+export const ASSISTANT_SHARED_SKILLS_PATH = "/Workspace/.assistant/skills";
+
+/** Assistant skills tree owned by one user (the "save this as a skill" target). */
+export function userAssistantSkillsPath(userEmail: string): string {
+  return `/Users/${userEmail.trim()}/.assistant/skills`;
+}

package/src/trace-io.ts

@@ -0,0 +1,129 @@
+/**
+ * Copy each Mastra chat turn's request/response onto the active OTel
+ * root span so Databricks-managed MLflow can display them.
+ *
+ * Why this exists: MLflow's UC `<target>_trace_unified` view picks the
+ * one span whose `parent_span_id` is empty and reads request/response
+ * from `mlflow.spanInputs` / `mlflow.spanOutputs` (or the `gen_ai.*`
+ * equivalents) on THAT span only. Mastra records the turn on its
+ * `invoke_agent` child under `mastra.agent_run.input` / `.output`, keys
+ * the view never reads, so chat traces arrive with both columns null
+ * even when the spans themselves look healthy.
+ *
+ * The active span inside the Mastra sub-app is still AppKit's HTTP
+ * server span (OTel context is async-local). With `OTEL_PROPAGATORS=none`
+ * that span is the trace root - see the package README's observability
+ * section for why Apps ingress `traceparent` otherwise hides every turn.
+ *
+ * @module
+ */
+
+import { json, log, object } from "@dbx-tools/shared-core";
+import { trace } from "@opentelemetry/api";
+import type express from "express";
+
+const logger = log.logger("mastra/trace-io");
+
+/** Cap on each payload copied onto a span, so one turn cannot bloat the export. */
+export const TRACE_IO_LIMIT = 8_000;
+
+/**
+ * Mount-relative Mastra agent invoke paths that carry a chat turn body
+ * (`messages`) and stream an assistant answer. Resume / approve verbs
+ * are intentionally excluded: their bodies are tool decisions, not the
+ * user prompt, and they rarely produce a fresh answer worth surfacing.
+ */
+const AGENT_TURN_ROUTE = /^\/agents\/[^/]+\/(stream|generate)(\/|$)/i;
+
+/** Attribute keys the MLflow UC `*_trace_unified` view reads from the root span. */
+export const MLFLOW_SPAN_INPUTS_ATTR = "mlflow.spanInputs";
+export const MLFLOW_SPAN_OUTPUTS_ATTR = "mlflow.spanOutputs";
+
+/**
+ * Concatenate the assistant's answer out of an AI SDK SSE transcript.
+ *
+ * The agent streams its reply as `text-delta` frames
+ * (`data: {"type":"text-delta","payload":{"text":"..."}}`), so the full
+ * answer only ever exists as deltas on the wire and has to be
+ * reassembled here.
+ */
+export function assistantTextFromSse(body: string): string {
+  const parts: string[] = [];
+  for (const line of body.split("\n")) {
+    if (!line.startsWith("data:")) continue;
+    // SSE carries comment/keepalive lines that are not JSON; they never hold
+    // assistant text, so skipping them is the intended path rather than an error.
+    const frame = json.parse(line.slice(5));
+    if (!object.isRecord(frame) || frame.type !== "text-delta") continue;
+    const text = object.isRecord(frame.payload) ? frame.payload.text : undefined;
+    if (typeof text === "string") parts.push(text);
+  }
+  return parts.join("");
+}
+
+/**
+ * Express middleware that stamps chat turn I/O onto the active OTel span.
+ *
+ * Tee `res.write` / `res.end` instead of listening for `finish`: the HTTP
+ * instrumentation ends the root span on finish, and an attribute set after
+ * a span ends never reaches the exporter.
+ *
+ * `path` is mount-relative (what the Mastra sub-app sees), e.g.
+ * `/agents/support/stream`.
+ */
+export function chatTurnTraceIoMiddleware(
+  req: express.Request,
+  res: express.Response,
+  next: express.NextFunction,
+): void {
+  if (req.method !== "POST" || !AGENT_TURN_ROUTE.test(req.path)) {
+    next();
+    return;
+  }
+  const span = trace.getActiveSpan();
+  if (!span) {
+    next();
+    return;
+  }
+
+  const messages = (req.body as { messages?: unknown } | undefined)?.messages;
+  if (messages !== undefined) {
+    span.setAttribute(MLFLOW_SPAN_INPUTS_ATTR, JSON.stringify(messages).slice(0, TRACE_IO_LIMIT));
+  }
+
+  const chunks: string[] = [];
+  const passThroughWrite = res.write.bind(res) as (...args: unknown[]) => boolean;
+  const passThroughEnd = res.end.bind(res) as (...args: unknown[]) => unknown;
+  const collect = (chunk: unknown): void => {
+    if (typeof chunk === "string") chunks.push(chunk);
+    else if (Buffer.isBuffer(chunk)) chunks.push(chunk.toString("utf8"));
+  };
+
+  res.write = ((chunk: unknown, ...rest: unknown[]) => {
+    collect(chunk);
+    return passThroughWrite(chunk, ...rest);
+  }) as typeof res.write;
+
+  res.end = ((chunk?: unknown, ...rest: unknown[]) => {
+    collect(chunk);
+    const answer = assistantTextFromSse(chunks.join(""));
+    if (answer) {
+      span.setAttribute(MLFLOW_SPAN_OUTPUTS_ATTR, answer.slice(0, TRACE_IO_LIMIT));
+    }
+    return passThroughEnd(chunk, ...rest);
+  }) as typeof res.end;
+
+  next();
+}
+
+/**
+ * Install {@link chatTurnTraceIoMiddleware} on a Mastra Express sub-app.
+ *
+ * Call before `MastraServer.init()` so the layer sits ahead of the agent
+ * routes. Safe to call unconditionally: with no active span (local, no
+ * OTLP) the middleware is a no-op.
+ */
+export function attachChatTurnTraceIo(app: express.Express): void {
+  app.use(chatTurnTraceIoMiddleware);
+  logger.info("chat turn I/O middleware attached");
+}