@dbx-tools/appkit-mastra 0.1.58 → 0.1.68

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

@@ -1,102 +0,0 @@
-/**
- * Mastra input processor that strips `chartId` fields from every
- * tool-invocation result in prior assistant messages before they
- * reach the model.
- *
- * 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
- * `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";
-import type { InputProcessor, ProcessInputArgs } from "@mastra/core/processors";
-
-const log = logUtils.logger("mastra/processor/strip-stale-charts");
-
-/**
- * Recursively clone `value`, omitting any property whose key is
- * `chartId`. Arrays are mapped element-wise; primitives are
- * returned as-is. The result is structurally identical to the
- * input minus chartIds, so downstream message-shape consumers
- * keep working.
- */
-function stripChartIds(value: unknown): unknown {
-  if (Array.isArray(value)) {
-    return value.map(stripChartIds);
-  }
-  if (value && typeof value === "object") {
-    const obj = value as Record<string, unknown>;
-    const out: Record<string, unknown> = {};
-    for (const [key, val] of Object.entries(obj)) {
-      if (key === "chartId") continue;
-      out[key] = stripChartIds(val);
-    }
-    return out;
-  }
-  return value;
-}
-
-/**
- * Input processor that scrubs `chartId` from every tool-invocation
- * result in the message list. Wired onto every agent by default
- * via {@link buildAgents}; opt out with
- * `MastraPluginConfig.stripStaleCharts: false`.
- */
-export const stripStaleChartsProcessor: InputProcessor = {
-  id: "strip-stale-charts",
-  description:
-    "Removes chartId fields from prior tool-invocation results so the model can't reuse turn-scoped ids from memory.",
-  processInput(args: ProcessInputArgs) {
-    let stripped = 0;
-    for (const message of args.messages) {
-      if (message.role !== "assistant") continue;
-      const parts = message.content?.parts;
-      if (!Array.isArray(parts)) continue;
-      for (const part of parts) {
-        // Tool-invocation parts hold the persisted tool result.
-        // We don't scrub the input args (`rawInput` / `args`) because
-        // the chartId there is the model's outgoing claim, not
-        // anything it could re-reference; only `result` carries
-        // ids that subsequent turns might copy.
-        if ((part as { type?: unknown }).type !== "tool-invocation") {
-          continue;
-        }
-        const inv = (part as { toolInvocation?: { result?: unknown } }).toolInvocation;
-        if (!inv || inv.result === undefined) continue;
-        const before = inv.result;
-        const after = stripChartIds(before);
-        // Cheap structural check via JSON length - the actual
-        // strip writes a fresh object only when chartId keys
-        // existed, so different stringification length is a
-        // reliable signal that something was removed.
-        if (
-          typeof before === "object" &&
-          before !== null &&
-          JSON.stringify(before).length !== JSON.stringify(after).length
-        ) {
-          inv.result = after;
-          stripped += 1;
-        }
-      }
-    }
-    if (stripped > 0) {
-      log.debug("stripped", { results: stripped });
-    }
-    return args.messages;
-  },
-};

package/src/server.ts

@@ -1,195 +0,0 @@
-/**
- * Express-layer plumbing for the Mastra plugin: a `MastraServer` that
- * stamps the per-request `RequestContext`, and a route-patch middleware
- * that lets the plugin's custom API routes (e.g. `historyRoute`) work
- * behind an Express mount point.
- */
-
-import { getExecutionContext } from "@databricks/appkit";
-import { httpUtils, logUtils, stringUtils } from "@dbx-tools/shared";
-import {
-  MASTRA_RESOURCE_ID_KEY,
-  MASTRA_THREAD_ID_KEY,
-  type RequestContext,
-} from "@mastra/core/request-context";
-import { MastraServer as MastraServerExpress } from "@mastra/express";
-import type express from "express";
-import { randomUUID } from "node:crypto";
-
-import {
-  MASTRA_REQUEST_ID_KEY,
-  MASTRA_USER_EMAIL_KEY,
-  MASTRA_USER_KEY,
-  MASTRA_USER_NAME_KEY,
-  type MastraPluginConfig,
-  type User,
-} from "./config.js";
-import {
-  extractModelOverride,
-  MASTRA_MODEL_OVERRIDE_KEY,
-  resolveServingConfig,
-} from "./serving.js";
-
-/**
- * `@mastra/express` subclass that stamps `RequestContext` with the
- * AppKit user, resource id, and a thread id backed by an HTTP-only
- * session cookie (`appkit_<plugin-name>_session_id`).
- */
-export class MastraServer extends MastraServerExpress {
-  private log: logUtils.Logger;
-
-  constructor(
-    private config: MastraPluginConfig,
-    ...args: ConstructorParameters<typeof MastraServerExpress>
-  ) {
-    super(...args);
-    this.log = logUtils.logger(config);
-  }
-
-  override registerAuthMiddleware(): void {
-    super.registerAuthMiddleware();
-    this.app.use((req, res, next) => {
-      const requestContext = res.locals.requestContext! as RequestContext;
-      this.configureRequestContextUser(requestContext);
-      this.configureRequestContextThreadId(req, res, requestContext);
-      this.configureRequestContextModelOverride(req, requestContext);
-      this.configureRequestContextRequestId(req, res, requestContext);
-      this.log.debug("auth:middleware", {
-        method: req.method,
-        path: req.path,
-        requestId: requestContext.get(MASTRA_REQUEST_ID_KEY),
-        threadId: requestContext.get(MASTRA_THREAD_ID_KEY),
-        resourceId: requestContext.get(MASTRA_RESOURCE_ID_KEY),
-        userName: requestContext.get(MASTRA_USER_NAME_KEY),
-        userEmail: requestContext.get(MASTRA_USER_EMAIL_KEY),
-        modelOverride: requestContext.get(
-          // imported below; logged so a misrouted request shows
-          // up alongside its model selection in `LOG_LEVEL=debug`.
-          "mastra__model_override",
-        ),
-      });
-      next();
-    });
-  }
-
-  configureRequestContextUser(requestContext: RequestContext) {
-    if (
-      [MASTRA_USER_KEY, MASTRA_RESOURCE_ID_KEY].every((key) => requestContext.get(key))
-    )
-      return;
-    const executionContext = getExecutionContext();
-    const user: User = {
-      id:
-        "userId" in executionContext
-          ? executionContext.userId
-          : executionContext.serviceUserId,
-      executionContext,
-    };
-    requestContext.set(MASTRA_USER_KEY, user);
-    requestContext.set(MASTRA_RESOURCE_ID_KEY, user.id);
-    // AppKit's `UserContext` surfaces display name / email only on
-    // OBO requests. Service-context calls (background tasks, server
-    // start-up) leave these undefined and we skip the stamp so
-    // downstream trace metadata stays absent rather than empty.
-    if ("isUserContext" in executionContext) {
-      if (executionContext.userName) {
-        requestContext.set(MASTRA_USER_NAME_KEY, executionContext.userName);
-      }
-      if (executionContext.userEmail) {
-        requestContext.set(MASTRA_USER_EMAIL_KEY, executionContext.userEmail);
-      }
-    }
-  }
-
-  /**
-   * Stamp a per-request id and echo it on the response so an upstream
-   * proxy / curl client / browser-side log line can pair its view of
-   * the request with the matching trace span. Reuses `X-Request-Id`
-   * when the upstream already supplies one so multi-hop traces stay
-   * joined; otherwise mints a UUIDv4.
-   *
-   * The id is surfaced as `mastra__requestId` span metadata via
-   * {@link TRACE_REQUEST_CONTEXT_KEYS} and as the `X-Request-Id`
-   * response header so dev tools can copy it from either side.
-   */
-  configureRequestContextRequestId(
-    req: express.Request,
-    res: express.Response,
-    requestContext: RequestContext,
-  ) {
-    if (requestContext.get(MASTRA_REQUEST_ID_KEY)) return;
-    const headerValue = req.headers["x-request-id"];
-    const upstream = Array.isArray(headerValue) ? headerValue[0] : headerValue;
-    const requestId = upstream?.trim() || randomUUID();
-    requestContext.set(MASTRA_REQUEST_ID_KEY, requestId);
-    res.setHeader("X-Request-Id", requestId);
-  }
-
-  configureRequestContextThreadId(
-    req: express.Request,
-    res: express.Response,
-    requestContext: RequestContext,
-  ) {
-    if (requestContext.get(MASTRA_THREAD_ID_KEY)) return;
-    const cookies = httpUtils.parseCookies(req.headers.cookie);
-    const cookieName = stringUtils.toIdentifierWithOptions(
-      { delimiter: "_", distinct: true },
-      "appkit",
-      this.config.name!,
-      "sessionId",
-    );
-    let sessionId = cookies[cookieName];
-    if (!sessionId) {
-      sessionId = randomUUID();
-      res.cookie(cookieName, sessionId, {
-        httpOnly: true,
-        sameSite: "lax",
-        secure: req.secure,
-        path: "/",
-      });
-    }
-    requestContext.set(MASTRA_THREAD_ID_KEY, sessionId);
-  }
-
-  configureRequestContextModelOverride(
-    req: express.Request,
-    requestContext: RequestContext,
-  ) {
-    // Per-request model override: only honored when the plugin
-    // opts in (default). Sources, in priority order, are
-    // `X-Mastra-Model` header, `?model=` query, and `model` /
-    // `modelId` body field; see `serving.ts`.
-    const serving = resolveServingConfig(this.config);
-    if (serving.allowOverride) {
-      const override = extractModelOverride({
-        headers: req.headers as Record<string, string | string[] | undefined>,
-        query: req.query as Record<string, unknown>,
-        body: req.body,
-      });
-      if (override) requestContext.set(MASTRA_MODEL_OVERRIDE_KEY, override);
-    }
-  }
-}
-
-/**
- * Patches around `@mastra/express`'s custom-route dispatcher so the
- * plugin's custom API routes (e.g. `historyRoute`) work when
- * `MastraServer` is hosted on an Express subapp mounted under a parent
- * path (e.g. `/api/mastra`).
- *
- * The adapter's `registerCustomApiRoutes` matches against `req.path`
- * (mount-relative, correct) but dispatches to its internal Hono
- * mini-app using `req.originalUrl`, which still contains the parent
- * mount prefix. The Hono app registers the literal route paths
- * (for example `/route/history`), so the absolute URL never matches
- * until we overwrite `originalUrl` for `/route` and `/route/*` to the
- * mount-relative path.
- */
-export function attachRoutePatchMiddleware(app: express.Express): void {
-  app.use((req, _res, next) => {
-    const isCustomRoute = req.path === "/route" || req.path.startsWith("/route/");
-    if (!isCustomRoute) return next();
-    req.originalUrl = req.path;
-    next();
-  });
-}

package/src/serving.ts

@@ -1,104 +0,0 @@
-/**
- * Mastra-specific glue over the generic `@dbx-tools/model` toolkit.
- *
- * The live `/serving-endpoints` catalogue access, fuzzy name
- * resolution, and class/fallback selection all live in
- * `@dbx-tools/model`; this module only adds what is specific to the
- * Mastra plugin: pulling a per-request model override off an HTTP
- * request (header / query / body) and projecting the plugin config
- * onto the knobs `buildModel` and the `/models` route share.
- */
-
-import { DEFAULT_FUZZY_THRESHOLD, DEFAULT_MODEL_CACHE_TTL_MS } from "@dbx-tools/model";
-import { stringUtils } from "@dbx-tools/shared";
-
-import type { MastraPluginConfig } from "./config.js";
-
-/**
- * `RequestContext` key under which {@link MastraServer} stores the
- * per-request model override (header / query / body). `model.ts`
- * reads it before falling back to the agent / plugin default.
- */
-export const MASTRA_MODEL_OVERRIDE_KEY = "mastra__model_override";
-
-/** HTTP header inspected for a per-request model override. */
-export const MODEL_OVERRIDE_HEADER = "x-mastra-model";
-
-/** Query string parameter inspected for a per-request model override. */
-export const MODEL_OVERRIDE_QUERY = "model";
-
-/** Body fields (in priority order) inspected for a per-request model override. */
-export const MODEL_OVERRIDE_BODY_FIELDS = ["model", "modelId"] as const;
-
-/**
- * Minimal Express-ish request shape used by {@link extractModelOverride}.
- * Keeps this module independent of `express` so the helper can be
- * reused from non-Express adapters.
- */
-export interface ModelOverrideRequest {
-  headers?: Record<string, string | string[] | undefined>;
-  query?: Record<string, unknown> | undefined;
-  body?: unknown;
-}
-
-/**
- * Pull a model override out of a single HTTP request, checking
- * sources in priority order:
- *
- *   1. `X-Mastra-Model` header
- *   2. `?model=` query string parameter
- *   3. Body field (`model` or `modelId`, in that order)
- *
- * Returns `null` when nothing is set, so callers can wrap with
- * `if (override) ...` without juggling empty strings. Body inspection
- * is lenient - any plain object with one of the configured keys
- * counts, mirroring how AI SDK chat clients pass arbitrary metadata
- * alongside `messages`.
- */
-export function extractModelOverride(req: ModelOverrideRequest): string | null {
-  const headers = req.headers;
-  if (headers) {
-    const headerVal = stringUtils.firstNonEmpty(
-      headers[MODEL_OVERRIDE_HEADER] ?? headers[MODEL_OVERRIDE_HEADER.toLowerCase()],
-    );
-    if (headerVal) return headerVal;
-  }
-  if (req.query) {
-    const queryVal = stringUtils.firstNonEmpty(req.query[MODEL_OVERRIDE_QUERY]);
-    if (queryVal) return queryVal;
-  }
-  if (req.body && typeof req.body === "object") {
-    const record = req.body as Record<string, unknown>;
-    for (const field of MODEL_OVERRIDE_BODY_FIELDS) {
-      const bodyVal = stringUtils.firstNonEmpty(record[field]);
-      if (bodyVal) return bodyVal;
-    }
-  }
-  return null;
-}
-
-/**
- * Read the fuzzy-resolution config knobs off the plugin config with
- * `@dbx-tools/model` defaults applied. Kept here so `buildModel` and
- * the `/models` route agree on what "enabled" means.
- *
- * `fallbacks` is the priority-ordered list `resolveModel` walks first
- * when nothing explicit is set; defaults to an empty list so the
- * generic resolver falls through to the live catalogue and its own
- * `FALLBACK_MODEL_IDS` floor.
- */
-export function resolveServingConfig(config: MastraPluginConfig): {
-  ttlMs: number;
-  threshold: number;
-  fuzzy: boolean;
-  allowOverride: boolean;
-  fallbacks: readonly string[];
-} {
-  return {
-    ttlMs: config.modelCacheTtlMs ?? DEFAULT_MODEL_CACHE_TTL_MS,
-    threshold: config.modelFuzzyThreshold ?? DEFAULT_FUZZY_THRESHOLD,
-    fuzzy: config.modelFuzzyMatch !== false,
-    allowOverride: config.modelOverride !== false,
-    fallbacks: config.defaultModelFallbacks ?? [],
-  };
-}

package/src/statement.ts

@@ -1,120 +0,0 @@
-/**
- * Databricks Statement Execution helpers for the Mastra plugin.
- *
- * 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
- * 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
- * `/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.
- */
-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 `/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`.
- */
-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
- * `/embed/data/:id` 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/writer.ts

@@ -1,44 +0,0 @@
-/**
- * Shared helper for publishing events through Mastra's
- * `ctx.writer`. Centralizes the "the downstream stream may already
- * be closed, don't take the whole tool down" pattern that the
- * Genie agent and chart tool both need.
- *
- * Failures are logged at `warn` (a persistently-closed writer is
- * the most likely culprit when events go missing client-side) but
- * swallowed so a cancelled request or a client that navigated
- * away can't crash a tool mid-flight.
- */
-
-import type { MastraWriter } from "@dbx-tools/appkit-mastra-shared";
-import { commonUtils, type logUtils } from "@dbx-tools/shared";
-
-/**
- * Best-effort `writer.write`. No-op when `writer` is undefined;
- * caught errors are logged via `log.warn("writer:error", ...)`
- * along with any caller-supplied `context` fields (e.g. a
- * `chartId` or `messageId`) so the warning is greppable per
- * resource.
- *
- * Returns when the write resolves or rejects; never throws.
- */
-export async function safeWrite(
-  log: logUtils.Logger,
-  writer: MastraWriter | undefined,
-  chunk: unknown,
-  context: Record<string, unknown> = {},
-): Promise<void> {
-  if (!writer) {
-    log.debug("writer:no-writer", context);
-    return;
-  }
-  try {
-    await writer.write(chunk);
-    log.debug("writer:ok", context);
-  } catch (err) {
-    log.warn("writer:error", {
-      ...context,
-      error: commonUtils.errorMessage(err),
-    });
-  }
-}