@dbx-tools/appkit-mastra 0.1.27 → 0.1.29

package/index.ts

@@ -1,7 +1,8 @@
 /**
- * AppKit Mastra integration: {@link MastraPlugin} / {@link mastra},
- * plugin config types, agent registration helpers, Genie tool
- * builders, and dynamic Model Serving endpoint resolution.
+ * Server-side entry point for the AppKit Mastra integration. Mounts
+ * the plugin via {@link mastra} and re-exports the full server surface
+ * (config, agent wiring, Genie and chart tooling, and dynamic Model
+ * Serving resolution) so apps build agent backends from one import.
  *
  * Client-side consumers should import URL helpers and the
  * {@link MastraClientConfig} type from `@dbx-tools/appkit-mastra-shared`

package/package.json

@@ -9,13 +9,13 @@
     }
   },
   "name": "@dbx-tools/appkit-mastra",
-  "version": "0.1.27",
+  "version": "0.1.29",
   "dependencies": {
     "@databricks/sdk-experimental": "^0.17",
-    "@dbx-tools/appkit-mastra-shared": "0.1.27",
-    "@dbx-tools/genie": "0.1.27",
-    "@dbx-tools/genie-shared": "0.1.27",
-    "@dbx-tools/shared": "0.1.27",
+    "@dbx-tools/appkit-mastra-shared": "0.1.29",
+    "@dbx-tools/genie": "0.1.29",
+    "@dbx-tools/genie-shared": "0.1.29",
+    "@dbx-tools/shared": "0.1.29",
     "@mastra/ai-sdk": "^1",
     "@mastra/core": "^1",
     "@mastra/express": "^1",

package/src/agents.ts

@@ -26,6 +26,7 @@ import type { MemoryBuilder } from "./memory.js";
 import { buildModel, FALLBACK_MODEL_IDS } from "./model.js";
 import { stripStaleChartsProcessor } from "./processors/strip-stale-charts.js";
 import { buildRenderDataTool } from "./chart.js";
+import { ResultProcessor } from "./processor.js";
 
 /**
  * Tool record accepted by every Mastra `Agent.tools` field and by the
@@ -433,6 +434,7 @@ export async function buildAgents(opts: {
   // chartIds from prior assistant tool results into the new
   // turn's `[chart:<id>]` markers. Opt out per-plugin via
   // `config.stripStaleCharts: false`.
+  const outputProcessors = [new ResultProcessor()];
   const inputProcessors =
     config.stripStaleCharts === false ? [] : [stripStaleChartsProcessor];
   const agents: Record<string, Agent> = {};
@@ -451,7 +453,8 @@ export async function buildAgents(opts: {
       },
       tools,
       ...(memory ? { memory } : {}),
-      ...(inputProcessors.length > 0 ? { inputProcessors } : {}),
+      inputProcessors,
+      outputProcessors,
     });
     // Surface the effective default model per agent so operators can
     // see at a glance which endpoint each agent points at without

package/src/chart.ts

@@ -20,18 +20,9 @@
  * dataset, custom API). The module has no knowledge of Genie or
  * statement ids; those concerns live in the tools that wrap it.
  *
- * Wire-format schemas (`ChartSchema`, `ChartResultSchema`,
- * {@link ChartTypeSchema}) live in
- * `@dbx-tools/appkit-mastra-shared` so the demo client and any
- * other UI consumer share the exact same shape this module reads
- * and writes.
- *
- * Public surface (everything else is module-private):
- *   - {@link chartPlannerRequestSchema} / {@link ChartPlannerRequest}
- *   - {@link prepareChart} / {@link PrepareChartOptions}
- *   - {@link fetchChart} / {@link FetchChartOptions}
- *   - {@link buildRenderDataTool} (the `render_data` Mastra tool
- *     auto-wired on every agent in `agents.ts`)
+ * Wire-format schemas live in `@dbx-tools/appkit-mastra-shared` so
+ * the demo client and any other UI consumer share the exact same
+ * shape this module reads and writes.
  */
 
 import { CacheManager } from "@databricks/appkit";

package/src/genie.ts

@@ -1,41 +1,19 @@
 /**
  * Genie tools for Mastra.
  *
- * Each configured Genie space is surfaced as a small set of flat
- * Mastra tools the calling agent can drive directly - no inner
- * orchestrator agent. The central agent decomposes user questions,
- * picks which Genie space to ask, and composes the final reply.
- *
- * Per-space tools (suffixed with `_<alias>` for non-default
- * aliases):
- *
- *   - `ask_genie`: drives one `genieEventChat` turn and
- *     forwards every wire event (status, thinking, sql, rows)
- *     through `ctx.writer` for streaming UI updates. Returns
- *     the terminal `GenieMessage` only - rows are NOT fetched
- *     eagerly.
- *   - `get_space_description`: cheap title / description /
- *     warehouse id lookup for grounding.
- *   - `get_space_serialized`: full `GenieSpace` JSON for
- *     column-level grounding when the description isn't enough.
- *
- * Space-agnostic shared tools (registered once, regardless of
- * how many spaces are wired):
- *
- *   - `get_statement`: opt-in lookup of a Genie statement's rows
- *     by `statement_id` (with a row `limit`). The agent calls
- *     this only when it needs to read values to reason about
- *     them; if the data is just being displayed, it embeds a
- *     `[data:<statement_id>]` marker in prose instead and lets
- *     the host UI resolve it.
- *   - `prepare_chart`: mints a short `chartId`, kicks off a
- *     background task that fetches the statement's rows and
- *     runs the chart-planner, and caches the resolved Echarts
- *     spec under the id (1h TTL). Returns the `chartId`
- *     synchronously so the agent embeds `[chart:<chartId>]`
- *     markers in prose without blocking on chart generation;
- *     the host UI fetches the cached chart by id once it's
- *     ready (see {@link fetchChart}).
+ * Surfaces each configured Genie space as a small set of flat Mastra
+ * tools the calling agent drives directly - no inner orchestrator
+ * agent. The central agent decomposes user questions, picks which
+ * space to ask, streams the per-turn wire events (status, thinking,
+ * sql, rows) through `ctx.writer`, and composes the final reply.
+ * Rows are never fetched eagerly: the agent reads a statement's
+ * values only when it needs to reason about them, otherwise it embeds
+ * a `[data:<statement_id>]` marker in prose and lets the host UI
+ * resolve the data. Charts are minted asynchronously and referenced
+ * by `[chart:<chartId>]` markers so prose isn't blocked on chart
+ * generation; the host UI fetches the cached spec by id once ready.
+ * Space description and serialized-space lookups are available for
+ * grounding when the agent needs schema context.
  *
  * Each tool's `execute` pulls the per-request
  * {@link WorkspaceClient} off `ctx.requestContext` (stamped by

package/src/model.ts

@@ -25,12 +25,7 @@
  * the input verbatim and let Databricks return the canonical error.
  */
 
-import {
-  commonUtils,
-  logUtils,
-  netUtils,
-  stringUtils,
-} from "@dbx-tools/shared";
+import { commonUtils, logUtils, netUtils, stringUtils } from "@dbx-tools/shared";
 import type { MastraModelConfig } from "@mastra/core/llm";
 import type { RequestContext } from "@mastra/core/request-context";
 
@@ -390,7 +385,7 @@ const setupFetchInterceptor = commonUtils.memoize((): void => {
   const log = logUtils.logger("mastra/llm");
   const original = globalThis.fetch.bind(globalThis);
   globalThis.fetch = (async (input, init) => {
-    const url = netUtils.parseUrl(input);
+    const url = netUtils.urlBuilder(input);
     if (
       !url ||
       !url.pathname.startsWith(SERVING_ENDPOINTS_PATH_PREFIX) ||

package/src/plugin.ts

@@ -54,10 +54,6 @@ import { historyRoute } from "./history.js";
 import { createMemoryBuilder, needsLakebase } from "./memory.js";
 import { buildObservability } from "./observability.js";
 import { attachRoutePatchMiddleware, MastraServer } from "./server.js";
-import {
-  installStreamEventInterceptor,
-  type StreamFrameInterceptor,
-} from "./intercept.js";
 import {
   clearServingEndpointsCache,
   listServingEndpoints,
@@ -69,6 +65,7 @@ import {
   isStatementNotFoundError,
   STATEMENT_ROW_CAP,
 } from "./statement.js";
+import { ResultProcessor } from "./processor.js";
 
 const GENIE_MANIFEST = appkitUtils.data(genie).plugin.manifest;
 const LAKEBASE_MANIFEST = appkitUtils.data(lakebase).plugin.manifest;
@@ -311,11 +308,6 @@ export class MastraPlugin extends Plugin<MastraPluginConfig> {
 
     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);
     });
   }
@@ -502,36 +494,4 @@ function parseStatementLimit(raw: unknown): number | 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: StreamFrameInterceptor = (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/src/processor.ts

@@ -0,0 +1,52 @@
+import type {
+  Processor,
+  ProcessAPIErrorArgs,
+  ProcessAPIErrorResult,
+  ProcessOutputStreamArgs,
+} from "@mastra/core/processors";
+import type { ChunkType } from "@mastra/core/workflows";
+
+export class ResultProcessor implements Processor {
+  id = "result-processor";
+
+  // Tell Mastra to also route tool/data parts to this processor method
+  processDataParts = true;
+
+  async processOutputStream({ part }: { part: any }): Promise<any | null> {
+    // 1. Guard clause: Ensure the chunk is a valid object
+    if (!part || typeof part !== "object") {
+      return part;
+    }
+
+    // 2. Filter for the targeted frame types
+    const targetedTypes = ["step-finish", "finish", "tool-result", "data-tool-agent"];
+    if (!targetedTypes.includes(part.type)) {
+      return part; // Return unchanged to pass-through
+    }
+
+    // 3. Check for the presence of a payload object
+    const payload = part.payload;
+    if (!payload || typeof payload !== "object") {
+      return part;
+    }
+
+    // 4. Safely delete the unwanted keys from the payload reference
+    const keysToDelete = ["output", "messages", "response", "result"];
+    for (const key of keysToDelete) {
+      if (key in payload) {
+        const value = payload[key];
+        if (typeof value === "object") {
+          payload[key] = {};
+        } else if (Array.isArray(value)) {
+          payload[key] = [];
+        } else {
+          delete payload[key];
+        }
+      }
+    }
+
+    // 5. Return the modified part object. Mastra handles re-serialization
+    // for the outbound SSE client stream automatically.
+    return part;
+  }
+}

package/src/serving.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 { CacheManager, type getExecutionContext } from "@databricks/appkit";

package/src/statement.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
- * `/embed/data/:id` 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
- *     `/embed/data/:id` 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