@dbx-tools/appkit-mastra 0.1.111 → 0.3.1

package/src/processors.ts

@@ -0,0 +1,168 @@
+/**
+ * Mastra stream processors wired onto the plugin's agents.
+ *
+ * Two complementary processors live here:
+ *
+ * - An **input** processor ({@link stripStaleChartsProcessor}) that
+ *   scrubs turn-scoped `chartId` fields out of prior tool results
+ *   replayed from Memory, so the model can't copy a stale id into the
+ *   new turn's `[chart:<id>]` markers.
+ * - An **output** processor ({@link ResultProcessor}) that trims the
+ *   bulky, redundant payload fields off targeted outbound stream frames
+ *   before they're serialized to the SSE client.
+ */
+
+import { log } from "@dbx-tools/shared-core";
+import type { InputProcessor, ProcessInputArgs, Processor } from "@mastra/core/processors";
+
+const logger = log.logger("mastra/processors");
+
+/**
+ * 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;
+}
+
+/**
+ * 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.
+ *
+ * Wired onto every agent by default via `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) {
+      logger.debug("stripped", { results: stripped });
+    }
+    return args.messages;
+  },
+};
+
+/**
+ * Mastra output processor that trims the bulky, redundant payload
+ * fields off targeted outbound stream frames.
+ *
+ * Mastra emits step / finish / tool-result frames whose payloads
+ * echo the full tool `output` / `result` / assistant `messages` /
+ * `response` - data the chat client already received inline and
+ * doesn't need re-sent. This processor empties those keys on the
+ * targeted frame types so the outbound SSE stream stays lean;
+ * every other frame passes through untouched.
+ */
+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/rest.ts

@@ -0,0 +1,67 @@
+import { appkit } from "@dbx-tools/appkit";
+/**
+ * Minimal authenticated Databricks REST helper. Pulls the workspace
+ * host and a fresh bearer header off an OBO-scoped `WorkspaceClient`
+ * (`client.config.getHost()` + `authenticate()`), then issues a plain
+ * `fetch`. Used by modules that hit REST surfaces without a typed SDK
+ * method (e.g. the MLflow assessments API); returns the raw `Response`
+ * so callers decide how to treat status codes. Also carries the
+ * defensive `Response` body readers those callers share.
+ */
+
+/** Workspace client carried on an AppKit execution context. */
+type WorkspaceClient = appkit.WorkspaceClientLike;
+
+/** Request options for {@link databricksFetch}. */
+export interface DatabricksFetchInit {
+  method: string;
+  /** JSON-serialized as the request body when present. */
+  body?: unknown;
+  /** Extra headers merged over the default `Content-Type: application/json`. */
+  headers?: Record<string, string>;
+  signal?: AbortSignal;
+}
+
+/**
+ * Resolve the workspace host + an authenticated header set off the
+ * client and issue a `fetch` against `path` (mounted on the host).
+ * Runs as whatever identity the client carries - the per-request OBO
+ * user when called from a request scope, the service principal
+ * otherwise.
+ */
+export async function databricksFetch(
+  client: WorkspaceClient,
+  path: string,
+  init: DatabricksFetchInit,
+): Promise<Response> {
+  const host = (await client.config.getHost()).toString();
+  const headers = new Headers({ "Content-Type": "application/json", ...init.headers });
+  await client.config.authenticate(headers);
+  const url = new URL(path, host).toString();
+  return fetch(url, {
+    method: init.method,
+    headers,
+    ...(init.body !== undefined ? { body: JSON.stringify(init.body) } : {}),
+    ...(init.signal ? { signal: init.signal } : {}),
+  });
+}
+
+/** Read a response body as text, swallowing read errors (returns `""`). */
+export async function readResponseText(res: Response): Promise<string> {
+  try {
+    return await res.text();
+  } catch {
+    return "";
+  }
+}
+
+/** Parse a response body as JSON, returning `{}` on empty / invalid bodies. */
+export async function readResponseJson(res: Response): Promise<unknown> {
+  const text = await readResponseText(res);
+  if (!text) return {};
+  try {
+    return JSON.parse(text);
+  } catch {
+    return {};
+  }
+}

package/src/server.ts

@@ -0,0 +1,343 @@
+/**
+ * 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 { feedback, thread } from "@dbx-tools/shared-mastra";
+import {
+  MASTRA_RESOURCE_ID_KEY,
+  MASTRA_THREAD_ID_KEY,
+  type RequestContext,
+} from "@mastra/core/request-context";
+import { MastraServer as MastraServerExpress } from "@mastra/express";
+import { trace } from "@opentelemetry/api";
+import type express from "express";
+import { randomUUID } from "node:crypto";
+
+import { resolveFeedbackEnabled } from "./mlflow";
+
+import {
+  MASTRA_REQUEST_ID_KEY,
+  MASTRA_SCOPES_KEY,
+  MASTRA_USER_EMAIL_KEY,
+  MASTRA_USER_KEY,
+  MASTRA_USER_NAME_KEY,
+  type MastraPluginConfig,
+  type User,
+} from "./config";
+import { http, log, object, string, token } from "@dbx-tools/shared-core";
+import { extractModelOverride, MASTRA_MODEL_OVERRIDE_KEY, resolveServingConfig } from "./serving";
+/**
+ * OpenTelemetry's sentinel for "no valid trace" - 32 zero hex chars.
+ * `trace.getActiveSpan()` returns a non-recording span with this id
+ * when no SDK is registered, which must never be surfaced as a trace to
+ * attach feedback to.
+ */
+const INVALID_TRACE_ID = "0".repeat(32);
+
+/**
+ * `@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: log.Logger;
+  /**
+   * Whether to stamp the MLflow trace-id header on responses. Shares the
+   * plugin's feedback gate via {@link resolveFeedbackEnabled}.
+   */
+  private feedbackEnabled: boolean;
+
+  constructor(
+    private config: MastraPluginConfig,
+    ...args: ConstructorParameters<typeof MastraServerExpress>
+  ) {
+    super(...args);
+    this.log = log.logger(config);
+    this.feedbackEnabled = resolveFeedbackEnabled(config.feedback);
+  }
+
+  override registerAuthMiddleware(): void {
+    super.registerAuthMiddleware();
+    this.app.use(async (req, res, next) => {
+      const requestContext = res.locals.requestContext! as RequestContext;
+      await this.configureRequestContextUser(requestContext);
+      this.configureRequestContextThreadId(req, res, requestContext);
+      this.configureRequestContextModelOverride(req, requestContext);
+      this.configureRequestContextRequestId(req, res, requestContext);
+      this.configureRequestContextScopes(req, requestContext);
+
+      this.configureMlflowTraceId(res);
+      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();
+    });
+  }
+
+  async 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.
+    let userName: string | undefined;
+    let email: string | undefined;
+    if ("isUserContext" in executionContext) {
+      userName = executionContext.userName;
+      email = executionContext.userEmail;
+    } else if (process.env.NODE_ENV === "development") {
+      const currentUser = await executionContext.client.currentUser.me();
+      userName = currentUser?.userName;
+      email = currentUser?.emails?.filter((email) => email.primary).find((email) => email.value)
+        ?.value as string;
+    }
+    if (userName) {
+      requestContext.set(MASTRA_USER_NAME_KEY, userName);
+    }
+    if (email) {
+      requestContext.set(MASTRA_USER_EMAIL_KEY, email);
+    }
+  }
+
+  /**
+   * 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);
+  }
+
+  /**
+   * Stamp OAuth scopes from the forwarded access token on
+   * {@link MASTRA_SCOPES_KEY} for workspace mount gating.
+   */
+  configureRequestContextScopes(req: express.Request, requestContext: RequestContext) {
+    const scopes = [
+      ...object
+        .sequence(
+          token.getAccessTokenScopes(req),
+          token.getAccessTokenScopes(req, "authorization"),
+        )
+        .distinct(),
+    ];
+    requestContext.set(MASTRA_SCOPES_KEY, scopes);
+  }
+
+  /**
+   * Stamp the turn's MLflow trace id on the response so the chat client
+   * can attach thumbs / comment feedback to it later. MLflow derives
+   * its trace id from the OpenTelemetry trace id (`tr-<hex>`), and every
+   * Mastra span for this request inherits the ambient OTel context (see
+   * `observability.ts`), so the active span's trace id here is the id
+   * MLflow will record for the turn.
+   *
+   * No-op unless feedback is enabled, and when no live OTel span is
+   * active (e.g. the OTLP SDK isn't registered): in that case the header
+   * is simply absent and the client hides feedback for that message,
+   * degrading gracefully rather than emitting a bogus trace id.
+   */
+  configureMlflowTraceId(res: express.Response) {
+    if (!this.feedbackEnabled || res.headersSent) return;
+    const traceId = trace.getActiveSpan()?.spanContext().traceId;
+    if (!traceId || traceId === INVALID_TRACE_ID) return;
+    res.setHeader(feedback.MLFLOW_TRACE_ID_HEADER, `tr-${traceId}`);
+  }
+
+  /**
+   * Resolve the thread id this request targets and pin it on
+   * `RequestContext` (consumed by the agent stream for persistence and
+   * by the history / threads routes). Resolution order:
+   *
+   *   1. A client-supplied thread id (the thread-selection header /
+   *      `?threadId=` query). This is how the chat UI references a
+   *      specific conversation among the many a user owns - it picks a
+   *      thread id from the `/threads` listing (or mints one for a new
+   *      conversation) and stamps it here. The id is scoped to the
+   *      caller's resource by the recall / list routes, so a client
+   *      can only ever read or write its own threads.
+   *   2. The per-session cookie (`appkit_<plugin-name>_session_id`),
+   *      minted on first contact. This is the default single-thread
+   *      fallback for clients that don't manage threads explicitly, so
+   *      existing embeds keep one stable conversation per session with
+   *      no client changes.
+   */
+  configureRequestContextThreadId(
+    req: express.Request,
+    res: express.Response,
+    requestContext: RequestContext,
+  ) {
+    if (requestContext.get(MASTRA_THREAD_ID_KEY)) return;
+    const requested = this.readRequestedThreadId(req);
+    if (requested) {
+      requestContext.set(MASTRA_THREAD_ID_KEY, requested);
+      return;
+    }
+    const cookies = http.parseCookies(req.headers.cookie);
+    const cookieName = string.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);
+  }
+
+  /**
+   * Read the client-selected thread id from the request, preferring
+   * the thread-selection header over the `?threadId=` query. Returns
+   * `null` when neither carries a non-empty value so the caller falls
+   * back to the session cookie.
+   */
+  private readRequestedThreadId(req: express.Request): string | null {
+    const headerValue = req.headers[thread.THREAD_ID_HEADER];
+    const queryValue = req.query[thread.THREAD_ID_QUERY];
+    return (
+      string.trimToNull(Array.isArray(headerValue) ? headerValue[0] : headerValue) ??
+      string.trimToNull(Array.isArray(queryValue) ? queryValue[0] : queryValue)
+    );
+  }
+
+  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);
+    }
+  }
+}
+
+/** Inputs for {@link isMastraRequestAllowed}. */
+export interface MastraApiGateOptions {
+  /** `config.apiAccess`; `"full"` short-circuits to allow everything. */
+  access: "scoped" | "full";
+  /** Whether the MCP transport is mounted (so `/mcp/*` is legitimate). */
+  mcpEnabled: boolean;
+}
+
+/**
+ * Mount-relative agent inference paths (`/agents/:id/<verb>...`). Matched
+ * by prefix on the verb so future variants (`streamVNext`, `stream/vnext`,
+ * `generateVNext`) stay covered without a code change.
+ *
+ * Covers the plain inference verbs (`stream` / `generate` / `network`) plus
+ * the human-in-the-loop resume verbs a paused `requireApproval` tool needs
+ * to continue: `approve-tool-call` / `decline-tool-call` (streaming, and via
+ * prefix their `-generate` non-streaming variants), the `-network-` variants,
+ * and `resume-stream` (covers `resume-stream-until-idle`). Without these an
+ * approval-gated tool (e.g. `send_email`) can be requested but never approved
+ * from the browser - the resume `POST` 403s under scoped mode. These are the
+ * only *writes* the browser client is allowed to make against stock Mastra.
+ */
+const AGENT_INFERENCE =
+  /^\/agents\/[^/]+\/(stream|generate|network|resume-stream|approve-tool-call|decline-tool-call|approve-network-tool-call|decline-network-tool-call)/i;
+
+/** Mount-relative read-only agent metadata (`/agents`, `/agents/:id`). */
+const AGENT_METADATA = /^\/agents(\/[^/]+)?$/;
+
+/**
+ * Whether a request to the stock `@mastra/express` sub-app should be
+ * dispatched, given the configured {@link MastraApiGateOptions.access}.
+ *
+ * `path` is mount-relative (what the plugin's catch-all sees, e.g.
+ * `/agents/x/stream`, `/route/history/x`, `/mcp/...`). In `"scoped"`
+ * mode the allowlist is deliberately tight - the chat client only ever
+ * needs agent inference, read-only agent metadata, this plugin's own
+ * OBO/resource-scoped `/route/*` routes, and (when enabled) MCP - so the
+ * whole admin / mutating / bulk-export surface Mastra also exposes is
+ * denied by default rather than enumerated.
+ */
+export function isMastraRequestAllowed(
+  method: string,
+  path: string,
+  opts: MastraApiGateOptions,
+): boolean {
+  if (opts.access === "full") return true;
+  const p = path.startsWith("/") ? path : `/${path}`;
+  // This plugin's own custom routes are individually OBO- and
+  // resource-scoped (see history.ts / threads.ts), so every method is safe.
+  if (p === "/route" || p.startsWith("/route/")) return true;
+  if (opts.mcpEnabled && (p === "/mcp" || p.startsWith("/mcp/"))) return true;
+  const m = method.toUpperCase();
+  if (m === "POST" && AGENT_INFERENCE.test(p)) return true;
+  if (m === "GET" && AGENT_METADATA.test(p)) return true;
+  return false;
+}
+
+/**
+ * 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();
+  });
+}