@dbx-tools/appkit-mastra 0.1.26 → 0.1.28

package/dist/src/plugin.js

@@ -27,7 +27,7 @@
  *   AI SDK transport URL is `/api/mastra/route/chat/<agentId>`.
  */
 import { genie, getExecutionContext, lakebase, Plugin, toPlugin, } from "@databricks/appkit";
-import { appkitUtils, logUtils } from "@dbx-tools/shared";
+import { appkitUtils, commonUtils, logUtils } from "@dbx-tools/shared";
 import { chatRoute } from "@mastra/ai-sdk";
 import { Mastra } from "@mastra/core/mastra";
 import express from "express";
@@ -37,6 +37,7 @@ import { historyRoute } from "./history.js";
 import { createMemoryBuilder, needsLakebase } from "./memory.js";
 import { buildObservability } from "./observability.js";
 import { attachRoutePatchMiddleware, MastraServer } from "./server.js";
+import { installStreamEventInterceptor, } from "./intercept.js";
 import { clearServingEndpointsCache, listServingEndpoints, resolveServingConfig, } from "./serving.js";
 import { fetchStatementData, isStatementNotFoundError, STATEMENT_ROW_CAP, } from "./statement.js";
 const GENIE_MANIFEST = appkitUtils.data(genie).plugin.manifest;
@@ -173,8 +174,7 @@ export class MastraPlugin extends Plugin {
             modelsPath: `${basePath}/models`,
             historyPath: `${basePath}/route/history`,
             historyPathTemplate: `${basePath}/route/history/:agentId`,
-            chartsPathTemplate: `${basePath}/charts/:chartId`,
-            statementsPathTemplate: `${basePath}/statements/:statementId`,
+            embedPathTemplate: `${basePath}/embed/:type/:id`,
             defaultAgent: this.built?.defaultAgentId ?? FALLBACK_AGENT_ID,
             agents: Object.keys(this.built?.agents ?? {}),
         };
@@ -192,112 +192,95 @@ export class MastraPlugin extends Plugin {
                 .then((endpoints) => res.json({ endpoints }))
                 .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).
+        // `GET /embed/:type/:id` is the single resolver for every embed
+        // marker the agent emits in prose (`[chart:<id>]`,
+        // `[data:<id>]`, ...). `:type` selects a resolver from the
+        // registry below; `:id` is that resolver's lookup key. The
+        // grammar (see `marker.ts`) is type-agnostic on purpose - new
+        // embed kinds are added by registering a resolver here, with no
+        // client or grammar change.
         //
         // 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.
+        //   - 200 with the resolver's JSON body when the id resolves.
+        //   - 404 when `:type` isn't registered (unsupported embed
+        //     type) OR a registered resolver can't find `:id` (unknown
+        //     / expired - e.g. a chart past its 1h TTL or a fabricated
+        //     id the model never minted).
+        //   - 400 when `:id` is empty.
         //
-        // Query params:
-        //   - `timeoutMs` (optional, default 60000): how long to long-poll
-        //     before returning a still-processing entry.
+        // Per-type query knobs and behavior:
+        //   - `chart`: long-polls the chart cache until the entry
+        //     settles (`result` / `error`) or the budget elapses (then
+        //     returns the still-processing entry to poll again).
+        //     `?timeoutMs=<n>` (default 60s, capped 5min) tunes it.
+        //   - `data`: one OBO-scoped Statement Execution fetch.
+        //     `?limit=<n>` caps rows (clamped to STATEMENT_ROW_CAP).
         //
-        // 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" });
+        // Built once (this handler is registered once) and keyed by the
+        // raw `:type` token. Each resolver gets the request (for query
+        // parsing + OBO scoping) and an `AbortSignal` bridged off the
+        // connection `close` event so a long-poll unblocks the instant
+        // the client disconnects. `undefined` from a resolver maps to a
+        // clean 404; thrown errors bubble through `next(err)`.
+        const embedResolvers = {
+            chart: (req, id, signal) => {
+                const timeoutMs = parseTimeoutMs(req.query["timeoutMs"]);
+                return fetchChart(id, {
+                    ...(timeoutMs !== undefined ? { timeoutMs } : {}),
+                    signal,
+                });
+            },
+            data: (req, id, signal) => {
+                const limit = parseStatementLimit(req.query["limit"]);
+                return this.userScopedSelf(req).fetchStatement(id, {
+                    ...(limit !== undefined ? { limit } : {}),
+                    signal,
+                });
+            },
+        };
+        router.get("/embed/:type/:id", (req, res, next) => {
+            const type = req.params["type"] ?? "";
+            const id = req.params["id"];
+            const resolve = embedResolvers[type];
+            if (!resolve) {
+                res.status(404).json({ error: `unsupported embed type: ${type}` });
                 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" });
+            if (!id) {
+                res.status(400).json({ error: "id is required" });
                 return;
             }
-            const limit = parseStatementLimit(req.query["limit"]);
+            // Express's `req` predates `AbortSignal`; bridge the `close`
+            // event onto an `AbortController` so a closed connection
+            // unblocks any long-poll immediately and frees the request
+            // thread. The listener is GC'd with the request on normal
+            // completion.
             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" });
+            resolve(req, id, controller.signal)
+                .then((entry) => {
+                if (entry === undefined) {
+                    res.status(404).json({ error: `${type} not found` });
                     return;
                 }
-                res.json(data);
+                res.json(entry);
             })
                 .catch(next);
         });
         router.use("", (req, res, next) => {
             if (!this.mastraApp)
                 return res.status(503).end();
+            // Run each Mastra `/stream` SSE frame through the interceptor
+            // (keep / rewrite / drop). Only engages on a
+            // `200 text/event-stream` body, so the JSON routes above and any
+            // error response are untouched.
+            installStreamEventInterceptor(res, interceptStreamFrame);
             return this.userScopedSelf(req).mastraApp(req, res, next);
         });
     }
     /**
-     * Implementation backing the `/statements/:statementId` route.
-     * Runs inside the AppKit user-context proxy so
+     * Implementation backing the `data` embed resolver
+     * (`GET /embed/data/:id`). 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
@@ -427,7 +410,7 @@ export class MastraPlugin extends Plugin {
 }
 /**
  * Parse the optional `?timeoutMs=<n>` query parameter from a
- * `GET /charts/:chartId` request. Accepts a positive integer up
+ * `GET /embed/chart/:id` 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
@@ -444,7 +427,7 @@ function parseTimeoutMs(raw) {
 }
 /**
  * Parse the optional `?limit=<n>` query parameter from a
- * `GET /statements/:statementId` request. Accepts a non-negative
+ * `GET /embed/data/:id` 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.
@@ -458,4 +441,37 @@ function parseStatementLimit(raw) {
         return undefined;
     return Math.floor(n);
 }
+/**
+ * {@link StreamFrameInterceptor} for Mastra `/stream` responses.
+ *
+ * Removes large, redundant payloads from terminal `step-finish`,
+ * `finish`, and `tool-result` frames before they reach the browser.
+ * These frames often repeat information already delivered incrementally
+ * via streamed text and tool events, including full tool outputs,
+ * accumulated responses, message history, SQL results, chart data, and
+ * other large result payloads.
+ *
+ * The interceptor deletes heavyweight payload properties
+ * (`output`, `messages`, `response`, and `result`) while preserving the
+ * event envelope and lifecycle metadata required by the client.
+ *
+ * Deletion is key based, so streams that do not contain these
+ * fields are passed through unchanged.
+ */
+const interceptStreamFrame = (chunk) => {
+    if (!commonUtils.isRecord(chunk))
+        return true;
+    if (!["step-finish", "finish", "tool-result"].find((type) => type === chunk.type))
+        return true;
+    const payload = chunk.payload;
+    if (!commonUtils.isRecord(payload))
+        return true;
+    const trimmedPayload = commonUtils.deleteKeys(payload, [
+        "output",
+        "messages",
+        "response",
+        "result",
+    ]);
+    return trimmedPayload ? { replace: chunk } : true;
+};
 export const mastra = toPlugin(MastraPlugin);

package/dist/src/serving.d.ts

@@ -1,23 +1,16 @@
 /**
  * Dynamic model resolution against Databricks Model Serving.
  *
- * Three concerns live here:
- *
- * 1. **Listing** - {@link listServingEndpoints} pulls the workspace's
- *    `/serving-endpoints` via the SDK and caches the result per host
- *    with a TTL. Concurrent callers share one in-flight promise (the
- *    same coalescing pattern as Python's `cachetools-async`).
- * 2. **Fuzzy matching** - {@link resolveModelId} runs the user's input
- *    through `fuse.js` extended search so loose tokens like
- *    `"claude sonnet"` snap to `databricks-claude-sonnet-4-6` even
- *    when typed without the full endpoint name.
- * 3. **Per-request override** - {@link extractModelOverride} pulls a
- *    model name from the `X-Mastra-Model` header, `?model=` query
- *    string, or `model` body field so the same agent can be exercised
- *    against different endpoints without redeploying.
- *
- * `model.ts` glues these together inside the per-step model resolver;
- * `plugin.ts` exposes the cached list at `GET /models`.
+ * Turns loose, human-typed model names into real serving-endpoint ids
+ * so the same agent can target different endpoints without code
+ * changes. The workspace's `/serving-endpoints` list is fetched once
+ * per host and cached with a TTL, with concurrent callers sharing one
+ * in-flight promise (the coalescing pattern of Python's
+ * `cachetools-async`). Names are matched through `fuse.js` extended
+ * search so tokens like `"claude sonnet"` snap to
+ * `databricks-claude-sonnet-4-6`, and a per-request override can be
+ * pulled from the request header, query string, or body so a model
+ * can be swapped per call without redeploying.
  */
 import { type getExecutionContext } from "@databricks/appkit";
 import type { ServingEndpointSummary } from "@dbx-tools/appkit-mastra-shared";

package/dist/src/serving.js

@@ -1,23 +1,16 @@
 /**
  * Dynamic model resolution against Databricks Model Serving.
  *
- * Three concerns live here:
- *
- * 1. **Listing** - {@link listServingEndpoints} pulls the workspace's
- *    `/serving-endpoints` via the SDK and caches the result per host
- *    with a TTL. Concurrent callers share one in-flight promise (the
- *    same coalescing pattern as Python's `cachetools-async`).
- * 2. **Fuzzy matching** - {@link resolveModelId} runs the user's input
- *    through `fuse.js` extended search so loose tokens like
- *    `"claude sonnet"` snap to `databricks-claude-sonnet-4-6` even
- *    when typed without the full endpoint name.
- * 3. **Per-request override** - {@link extractModelOverride} pulls a
- *    model name from the `X-Mastra-Model` header, `?model=` query
- *    string, or `model` body field so the same agent can be exercised
- *    against different endpoints without redeploying.
- *
- * `model.ts` glues these together inside the per-step model resolver;
- * `plugin.ts` exposes the cached list at `GET /models`.
+ * Turns loose, human-typed model names into real serving-endpoint ids
+ * so the same agent can target different endpoints without code
+ * changes. The workspace's `/serving-endpoints` list is fetched once
+ * per host and cached with a TTL, with concurrent callers sharing one
+ * in-flight promise (the coalescing pattern of Python's
+ * `cachetools-async`). Names are matched through `fuse.js` extended
+ * search so tokens like `"claude sonnet"` snap to
+ * `databricks-claude-sonnet-4-6`, and a per-request override can be
+ * pulled from the request header, query string, or body so a model
+ * can be swapped per call without redeploying.
  */
 import { CacheManager } from "@databricks/appkit";
 import { logUtils, stringUtils } from "@dbx-tools/shared";

package/dist/src/statement.d.ts

@@ -1,24 +1,17 @@
 /**
  * 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.
+ * Wraps `client.statementExecution.getStatement` with the shape,
+ * size, and error handling the plugin's tools and the
+ * `/embed/data/:id` route both need: a low-level fetch that returns a
+ * raw `{columns, rows, rowCount}` shape and coerces numeric strings
+ * to numbers so downstream charts and aggregations don't have to, a
+ * hard row cap callers clamp `limit` to so a runaway result set can't
+ * hose a response, and a structural not-found detector that
+ * normalizes the SDK's error classes and loose `does not exist` /
+ * `not found` message shapes into one boolean so the route can 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
@@ -30,7 +23,7 @@ import { WorkspaceClient } from "@databricks/sdk-experimental";
 import type { GenieDatasetData } from "@dbx-tools/appkit-mastra-shared";
 /**
  * Hard server-side cap on rows returned by the
- * `/statements/:statementId` route. Sized to keep responses small
+ * `/embed/data/:id` 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.
@@ -48,7 +41,7 @@ export declare const STATEMENT_ROW_CAP = 500;
  *
  * 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
+ * the `/embed/data/:id` 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`.
@@ -66,7 +59,7 @@ export declare function fetchStatementData(client: WorkspaceClient, statementId:
  * catalogued.
  *
  * Pulled into its own helper so callers (notably the
- * `/statements/:statementId` route) stay decoupled from SDK
+ * `/embed/data/:id` route) stay decoupled from SDK
  * error-class identity, and the conversion logic stays testable
  * in isolation.
  */

package/dist/src/statement.js

@@ -1,24 +1,17 @@
 /**
  * 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.
+ * Wraps `client.statementExecution.getStatement` with the shape,
+ * size, and error handling the plugin's tools and the
+ * `/embed/data/:id` route both need: a low-level fetch that returns a
+ * raw `{columns, rows, rowCount}` shape and coerces numeric strings
+ * to numbers so downstream charts and aggregations don't have to, a
+ * hard row cap callers clamp `limit` to so a runaway result set can't
+ * hose a response, and a structural not-found detector that
+ * normalizes the SDK's error classes and loose `does not exist` /
+ * `not found` message shapes into one boolean so the route can 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
@@ -30,7 +23,7 @@ import { ApiError, HttpError, WorkspaceClient } from "@databricks/sdk-experiment
 import { apiUtils } from "@dbx-tools/shared";
 /**
  * Hard server-side cap on rows returned by the
- * `/statements/:statementId` route. Sized to keep responses small
+ * `/embed/data/:id` 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.
@@ -63,7 +56,7 @@ function coerceCell(cell) {
  *
  * 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
+ * the `/embed/data/:id` 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`.
@@ -98,7 +91,7 @@ export async function fetchStatementData(client, statementId, options) {
  * catalogued.
  *
  * Pulled into its own helper so callers (notably the
- * `/statements/:statementId` route) stay decoupled from SDK
+ * `/embed/data/:id` route) stay decoupled from SDK
  * error-class identity, and the conversion logic stays testable
  * in isolation.
  */