@dbx-tools/appkit-mastra 0.1.25 → 0.1.26

package/src/plugin.ts

@@ -44,7 +44,11 @@ import { Mastra } from "@mastra/core/mastra";
 import express from "express";
 
 import { buildAgents, FALLBACK_AGENT_ID, type BuiltAgents } from "./agents.js";
-import type { MastraClientConfig } from "@dbx-tools/appkit-mastra-shared";
+import type {
+  MastraClientConfig,
+  StatementData,
+} from "@dbx-tools/appkit-mastra-shared";
+import { fetchChart } from "./chart.js";
 import type { MastraPluginConfig } from "./config.js";
 import { historyRoute } from "./history.js";
 import { createMemoryBuilder, needsLakebase } from "./memory.js";
@@ -56,6 +60,11 @@ import {
   resolveServingConfig,
   type ServingEndpointSummary,
 } from "./serving.js";
+import {
+  fetchStatementData,
+  isStatementNotFoundError,
+  STATEMENT_ROW_CAP,
+} from "./statement.js";
 
 const GENIE_MANIFEST = appkitUtils.data(genie).plugin.manifest;
 const LAKEBASE_MANIFEST = appkitUtils.data(lakebase).plugin.manifest;
@@ -199,6 +208,8 @@ export class MastraPlugin extends Plugin<MastraPluginConfig> {
       modelsPath: `${basePath}/models`,
       historyPath: `${basePath}/route/history`,
       historyPathTemplate: `${basePath}/route/history/:agentId`,
+      chartsPathTemplate: `${basePath}/charts/:chartId`,
+      statementsPathTemplate: `${basePath}/statements/:statementId`,
       defaultAgent: this.built?.defaultAgentId ?? FALLBACK_AGENT_ID,
       agents: Object.keys(this.built?.agents ?? {}),
     };
@@ -218,12 +229,147 @@ export class MastraPlugin extends Plugin<MastraPluginConfig> {
         .catch(next);
     });
 
+    // `GET /charts/:chartId` long-polls the chart cache.
+    //
+    // Charts kicked off by the Genie agent's `prepare_chart` tool
+    // (or the system `render_data` tool) resolve asynchronously
+    // and land in the AppKit cache under a chartId. The host UI
+    // resolves `[chart:<chartId>]` markers by hitting this route,
+    // which blocks until the entry settles (`result` or `error`
+    // set) or the server-side budget elapses (then returns the
+    // last seen still-processing entry so the client can poll
+    // again).
+    //
+    // Status codes:
+    //   - 200 with JSON body when the entry is found (any state).
+    //   - 404 when the chartId doesn't exist or its 1h TTL elapsed.
+    //
+    // Query params:
+    //   - `timeoutMs` (optional, default 60000): how long to long-poll
+    //     before returning a still-processing entry.
+    //
+    // Wired with `req.signal` so a closed connection unblocks the
+    // poll immediately and the request thread frees up. Cache
+    // failures bubble through `next(err)` to Express's default
+    // error handler.
+    router.get("/charts/:chartId", (req, res, next) => {
+      const chartId = req.params["chartId"];
+      if (!chartId) {
+        res.status(400).json({ error: "chartId is required" });
+        return;
+      }
+      const timeoutMs = parseTimeoutMs(req.query["timeoutMs"]);
+      // Express's `req` predates `AbortSignal`. Bridge the
+      // `close` event onto an `AbortController` so a closed
+      // connection unblocks the long-poll immediately and frees
+      // the request thread. The `close` listener also fires on
+      // a normal completion, but at that point we no longer
+      // care about the abort and the listener is GC'd with the
+      // request.
+      const controller = new AbortController();
+      req.on("close", () => controller.abort());
+      fetchChart(chartId, {
+        ...(timeoutMs !== undefined ? { timeoutMs } : {}),
+        signal: controller.signal,
+      })
+        .then((entry) => {
+          if (!entry) {
+            res.status(404).json({ error: "chart not found" });
+            return;
+          }
+          res.json(entry);
+        })
+        .catch(next);
+    });
+
+    // `GET /statements/:statementId` resolves a Genie / Statement
+    // Execution result for inline rendering as a table.
+    //
+    // The agent embeds `[data:<statement_id>]` markers in prose
+    // (see `GENIE_INSTRUCTIONS`) whenever it wants the host UI to
+    // display a result set without the model re-spelling every
+    // row in markdown. This route is what those markers resolve
+    // against - one OBO-scoped fetch returns the columns + rows
+    // so the client renders an AppKit Table inline.
+    //
+    // Status codes:
+    //   - 200 with `StatementData` body when the statement is found.
+    //   - 404 when the workspace can't find / read the statement.
+    //
+    // Query params:
+    //   - `limit` (optional): row cap. Clamped to
+    //     {@link STATEMENT_ROW_CAP} so a runaway result set
+    //     can't hose the response.
+    //
+    // OBO-scoped via {@link userScopedSelf} so the caller's
+    // workspace permissions apply; SDK errors propagate through
+    // `next(err)` to Express's default error handler.
+    router.get("/statements/:statementId", (req, res, next) => {
+      const statementId = req.params["statementId"];
+      if (!statementId) {
+        res.status(400).json({ error: "statementId is required" });
+        return;
+      }
+      const limit = parseStatementLimit(req.query["limit"]);
+      const controller = new AbortController();
+      req.on("close", () => controller.abort());
+      this.userScopedSelf(req)
+        .fetchStatement(statementId, {
+          ...(limit !== undefined ? { limit } : {}),
+          signal: controller.signal,
+        })
+        .then((data) => {
+          if (!data) {
+            res.status(404).json({ error: "statement not found" });
+            return;
+          }
+          res.json(data);
+        })
+        .catch(next);
+    });
+
     router.use("", (req, res, next) => {
       if (!this.mastraApp) return res.status(503).end();
       return this.userScopedSelf(req).mastraApp!(req, res, next);
     });
   }
 
+  /**
+   * Implementation backing the `/statements/:statementId` route.
+   * Runs inside the AppKit user-context proxy so
+   * `getExecutionContext()` returns the OBO-scoped workspace
+   * client, then reuses the same `fetchStatementData` pipeline
+   * the `get_statement` tool runs so the LLM and the UI see the
+   * exact same shape for the same statement.
+   *
+   * Returns `undefined` for upstream 404s so the route can map
+   * them to a clean HTTP 404; any other failure bubbles up.
+   */
+  private async fetchStatement(
+    statementId: string,
+    options: { limit?: number; signal?: AbortSignal } = {},
+  ): Promise<StatementData | undefined> {
+    const client = getExecutionContext().client;
+    const limit = Math.min(options.limit ?? STATEMENT_ROW_CAP, STATEMENT_ROW_CAP);
+    try {
+      const data = await fetchStatementData(client, statementId, {
+        limit,
+        ...(options.signal ? { signal: options.signal } : {}),
+      });
+      return {
+        columns: data.columns,
+        rows: data.rows,
+        rowCount: data.rowCount,
+        truncated: data.rows.length < data.rowCount,
+      };
+    } catch (err) {
+      // The Databricks SDK throws on 404; surface as `undefined`
+      // so the route maps to a clean HTTP 404 instead of a 500.
+      if (isStatementNotFoundError(err)) return undefined;
+      throw err;
+    }
+  }
+
   /**
    * Return `this.asUser(req)` when the request carries an OBO token,
    * otherwise return `this` directly. Prevents the noisy AppKit warn
@@ -326,4 +472,35 @@ export class MastraPlugin extends Plugin<MastraPluginConfig> {
   }
 }
 
+/**
+ * Parse the optional `?timeoutMs=<n>` query parameter from a
+ * `GET /charts/:chartId` request. Accepts a positive integer up
+ * to 5 minutes (clamped) and rejects everything else as
+ * `undefined` so {@link fetchChart} falls back to its default.
+ * Express produces `string | string[] | undefined`; we normalize
+ * to the first scalar before parsing.
+ */
+function parseTimeoutMs(raw: unknown): number | undefined {
+  const v = Array.isArray(raw) ? raw[0] : raw;
+  if (typeof v !== "string") return undefined;
+  const n = Number(v);
+  if (!Number.isFinite(n) || n <= 0) return undefined;
+  return Math.min(Math.floor(n), 5 * 60_000);
+}
+
+/**
+ * Parse the optional `?limit=<n>` query parameter from a
+ * `GET /statements/:statementId` request. Accepts a non-negative
+ * integer and lets the route clamp to `STATEMENT_ROW_CAP`;
+ * rejects anything else as `undefined` so the route falls back
+ * to the server-side cap.
+ */
+function parseStatementLimit(raw: unknown): number | undefined {
+  const v = Array.isArray(raw) ? raw[0] : raw;
+  if (typeof v !== "string") return undefined;
+  const n = Number(v);
+  if (!Number.isFinite(n) || n < 0) return undefined;
+  return Math.floor(n);
+}
+
 export const mastra = toPlugin(MastraPlugin);

package/src/processors/strip-stale-charts.ts

@@ -3,21 +3,24 @@
  * tool-invocation result in prior assistant messages before they
  * reach the model.
  *
- * Why: chartIds are only meaningful within the assistant turn that
- * minted them - the writer events backing them are gone after the
- * stream closes. When the model sees old chartIds in memory recall
- * (Mastra Memory persists tool results), it's tempted to type
- * those ids into the new turn's `[[chart:<id>]]` markers, leaving
- * the chat client's chart slots stuck with no matching event. This
- * processor removes the temptation by deleting `chartId` keys from
- * every assistant message's tool results before the prompt is
- * built. The current turn's tool results don't exist yet at
- * `processInput` time, so they pass through unmodified.
+ * Why: chartIds are turn-scoped from the model's point of view -
+ * each `prepare_chart` / `render_data` call mints a fresh id and
+ * the host UI binds it to that turn's reply. Mastra Memory
+ * replays prior tool results into the prompt; if old chartIds
+ * leak through, the model is tempted to copy them verbatim into
+ * the new turn's `[chart:<id>]` markers and the host UI ends up
+ * rendering an unrelated chart from the chart cache (or
+ * a 404 once the 1h TTL elapsed). This processor removes the
+ * temptation by deleting `chartId` keys from every assistant
+ * message's tool results before the prompt is built. The current
+ * turn's tool results don't exist yet at `processInput` time, so
+ * they pass through unmodified.
  *
  * The strip is recursive - any nested `chartId` field is removed,
- * regardless of which tool produced the result. This covers Genie's
- * `datasets[].chartId` and `render_data`'s top-level `chartId`
- * uniformly without coupling to specific tool ids.
+ * regardless of which tool produced the result. This covers
+ * `prepare_chart` / `render_data` top-level chartIds and any
+ * legacy `datasets[].chartId` payloads uniformly without coupling
+ * to specific tool ids.
  */
 
 import { logUtils } from "@dbx-tools/shared";

package/src/statement.ts

@@ -0,0 +1,127 @@
+/**
+ * Databricks Statement Execution helpers for the Mastra plugin.
+ *
+ * Wraps `client.statementExecution.getStatement` with the shape
+ * + size + error handling the plugin's tools and the
+ * `/statements/:statementId` route both need:
+ *
+ *   - {@link fetchStatementData}: low-level fetch that returns the
+ *     raw `{columns, rows, rowCount}` shape used by the
+ *     `get_statement` tool's output, the `prepare_chart` tool's
+ *     dataset resolver, and the route's response body. Coerces
+ *     numeric strings to numbers so downstream charts /
+ *     aggregations don't have to.
+ *   - {@link STATEMENT_ROW_CAP}: hard cap callers (notably the
+ *     `/statements/:statementId` route) clamp `limit` to so a
+ *     runaway result set can't hose a response.
+ *   - {@link isStatementNotFoundError}: structural detector that
+ *     normalizes the SDK's two error classes plus the loose
+ *     `does not exist` / `not found` message shapes into a single
+ *     boolean - lets the route map upstream 404s to a clean
+ *     HTTP 404 without coupling to SDK error-class identity.
+ *
+ * Not Genie-specific: a Databricks `statement_id` is workspace
+ * scoped and lives in the Statement Execution API regardless of
+ * which producer (Genie, a tool, a notebook, etc.) submitted the
+ * query. Co-located here so consumers can fetch / cap / handle
+ * 404s without reaching into the Genie tool module.
+ */
+
+import { ApiError, HttpError, WorkspaceClient } from "@databricks/sdk-experimental";
+import type { GenieDatasetData } from "@dbx-tools/appkit-mastra-shared";
+import { apiUtils } from "@dbx-tools/shared";
+
+/**
+ * Hard server-side cap on rows returned by the
+ * `/statements/:statementId` route. Sized to keep responses small
+ * enough for inline tables to render snappily; the route surfaces
+ * a `truncated` flag whenever the upstream `rowCount` exceeds
+ * this so end users know they're seeing a sample.
+ */
+export const STATEMENT_ROW_CAP = 500;
+
+/**
+ * Best-effort numeric coercion for the Statement Execution API's
+ * all-strings cells. Leaves non-numeric strings (and explicit
+ * `null`s) intact; everything else flows through `Number`.
+ */
+function coerceCell(cell: string | null): unknown {
+  if (cell === null) return null;
+  if (/^-?\d+(\.\d+)?$/.test(cell)) {
+    const n = Number(cell);
+    if (Number.isFinite(n)) return n;
+  }
+  return cell;
+}
+
+/**
+ * Fetch a single statement's rows via the Statement Execution API
+ * and reshape into the shared {@link GenieDatasetData} shape
+ * (column array + row records).
+ *
+ * Optional `limit` slices the returned `rows` client-side so the
+ * agent can scan a small sample without paging the full result
+ * set into context. `rowCount` always reflects the upstream total
+ * so callers know when the slice truncated.
+ *
+ * Exported because every consumer in the plugin (the
+ * `get_statement` tool, the `prepare_chart` dataset resolver, and
+ * the `/statements/:statementId` route) needs the exact same
+ * fetch + coercion pipeline so LLM-side `get_statement` output
+ * and UI-side `[data:<id>]` rendering stay shape-identical for
+ * the same `statement_id`.
+ */
+export async function fetchStatementData(
+  client: WorkspaceClient,
+  statementId: string,
+  options?: { limit?: number; signal?: AbortSignal },
+): Promise<GenieDatasetData> {
+  const ctx = options?.signal ? apiUtils.toContext(options.signal) : undefined;
+  const r = await client.statementExecution.getStatement(
+    { statement_id: statementId },
+    ctx,
+  );
+  const columns = (r.manifest?.schema?.columns ?? []).map((c) => c.name ?? "");
+  const dataArray = (r.result?.data_array ?? []) as Array<Array<string | null>>;
+  const sliced =
+    options?.limit !== undefined && options.limit >= 0
+      ? dataArray.slice(0, options.limit)
+      : dataArray;
+  const rows = sliced.map((row) => {
+    const obj: Record<string, unknown> = {};
+    columns.forEach((col, i) => {
+      obj[col] = coerceCell(row[i] ?? null);
+    });
+    return obj;
+  });
+  return {
+    columns,
+    rows,
+    rowCount: r.manifest?.total_row_count ?? dataArray.length,
+  };
+}
+
+/**
+ * True when `err` looks like the Databricks SDK's "statement not
+ * found" error. Matches the typed {@link ApiError} 404 /
+ * `RESOURCE_DOES_NOT_EXIST` shape first, then falls back to the
+ * lower-level {@link HttpError} 404, then to a loose `does not
+ * exist` / `not found` message sniff for SDK shapes we haven't
+ * catalogued.
+ *
+ * Pulled into its own helper so callers (notably the
+ * `/statements/:statementId` route) stay decoupled from SDK
+ * error-class identity, and the conversion logic stays testable
+ * in isolation.
+ */
+export function isStatementNotFoundError(err: unknown): boolean {
+  if (err instanceof ApiError) {
+    if (err.statusCode === 404) return true;
+    if (err.errorCode === "RESOURCE_DOES_NOT_EXIST") return true;
+  }
+  if (err instanceof HttpError && err.code === 404) return true;
+  if (err instanceof Error && /does not exist|not found/i.test(err.message)) {
+    return true;
+  }
+  return false;
+}

package/src/tools/email.ts

@@ -33,41 +33,41 @@ import { z } from "zod";
 const log = logUtils.logger("mastra/tool/send-email");
 
 const emailInputSchema = z.object({
-  to: z.string().describe(stringUtils.toDescription`
+  to: z.string().describe(stringUtils.toDescription(`
     Single recipient email address (e.g. "alice@example.com"). For
     multiple recipients, comma-separate them yourself.
-  `),
-  subject: z.string().describe(stringUtils.toDescription`
+  `)),
+  subject: z.string().describe(stringUtils.toDescription(`
     Subject line.
-  `),
-  body: z.string().describe(stringUtils.toDescription`
-    Email body. Plain text or markdown; the renderer downstream
-    decides which to honour. Be specific - the recipient may not
-    have any context the model has from prior chat turns.
-  `),
+  `)),
+  body: z.string().describe(stringUtils.toDescription(`
+    Email body. Plain text or markdown; the renderer downstream decides
+    which to honour. Be specific - the recipient may not have any
+    context the model has from prior chat turns.
+  `)),
   cc: z
     .array(z.string())
     .optional()
-    .describe(stringUtils.toDescription`
+    .describe(stringUtils.toDescription(`
       Optional CC recipients.
-    `),
+    `)),
   bcc: z
     .array(z.string())
     .optional()
-    .describe(stringUtils.toDescription`
+    .describe(stringUtils.toDescription(`
       Optional BCC recipients.
-    `),
+    `)),
 });
 
 const emailOutputSchema = z.object({
-  sent: z.boolean().describe(stringUtils.toDescription`
+  sent: z.boolean().describe(stringUtils.toDescription(`
     True when the email was dispatched. The current implementation
-    always returns true after console-logging the would-be email;
-    swap in a real provider to make this meaningful.
-  `),
-  recipient: z.string().describe(stringUtils.toDescription`
+    always returns true after console-logging the would-be email; swap
+    in a real provider to make this meaningful.
+  `)),
+  recipient: z.string().describe(stringUtils.toDescription(`
     Echo of the \`to\` field for confirmation.
-  `),
+  `)),
 });
 
 /** Options accepted by {@link buildEmailTool}. */
@@ -111,15 +111,14 @@ export interface BuildEmailToolOptions {
 export function buildEmailTool(opts: BuildEmailToolOptions = {}) {
   return createTool({
     id: opts.id ?? "send_email",
-    description: stringUtils.toDescription`
-      Send an email on the user's behalf. Pass a recipient
-      address, subject, and body; the user will be prompted to
-      approve the send before it goes out (the tool is
-      approval-gated). Use this when the user explicitly asks
-      to send / forward / share something via email - never
-      autonomously. Keep subjects short and bodies focused; the
-      recipient may not have any of the chat context.
-    `,
+    description: stringUtils.toDescription(`
+      Send an email on the user's behalf. Pass a recipient address,
+      subject, and body; the user will be prompted to approve the send
+      before it goes out (the tool is approval-gated). Use this when
+      the user explicitly asks to send / forward / share something via
+      email - never autonomously. Keep subjects short and bodies
+      focused; the recipient may not have any of the chat context.
+    `),
     inputSchema: emailInputSchema,
     outputSchema: emailOutputSchema,
     requireApproval: true,