@dbx-tools/appkit-mastra 0.1.112 → 0.3.1

package/src/serving-sanitize.ts

@@ -0,0 +1,167 @@
+import { string } from "@dbx-tools/shared-core";
+/**
+ * Repairs Mastra / AI SDK message replays sent to Databricks Model
+ * Serving before they hit the OpenAI-compatible `/chat/completions`
+ * route.
+ */
+
+/**
+ * OpenAI-flavoured chat message shape we need to mutate. We do not
+ * import the OpenAI / AI SDK types because both packages keep these
+ * fields under internal namespaces; the wire payload is the contract
+ * here and it's stable enough to inline.
+ */
+export interface ServingChatMessage {
+  role: "system" | "user" | "assistant" | "tool" | "reasoning";
+  content?: string | ServingContentPart[];
+  tool_calls?: Array<{ id: string; type: string; function: unknown }>;
+  tool_call_id?: string;
+  reasoning?: unknown;
+  reasoning_content?: unknown;
+}
+
+type ServingContentPart = {
+  type?: string;
+  text?: string;
+  [key: string]: unknown;
+};
+
+const REASONING_PART_TYPES = new Set(["reasoning", "thinking", "redacted_thinking"]);
+
+/**
+ * Parse, sanitize, and re-serialize a `/serving-endpoints/...` POST
+ * body. Returns the original string verbatim when the body is not
+ * JSON, has no `messages`, or no rewrite was needed.
+ */
+export function rewriteServingBody(body: string): string {
+  let parsed: { messages?: unknown };
+  try {
+    parsed = JSON.parse(body);
+  } catch {
+    return body;
+  }
+  if (!Array.isArray(parsed.messages)) return body;
+  const messages = parsed.messages as ServingChatMessage[];
+  const changed = stripReasoningFromServingMessages(messages) || repairAssistantPrefill(messages);
+  return changed ? JSON.stringify(parsed) : body;
+}
+
+/**
+ * Drop extended-thinking / reasoning blocks from a replayed transcript.
+ *
+ * Hybrid Claude endpoints (e.g. Sonnet 4.5+) may emit `reasoning` content
+ * parts on the first turn. Mastra persists and replays them on the next
+ * agent step, but Databricks-hosted Claude rejects those blocks on
+ * multi-turn tool continuations. The UI already captured reasoning for
+ * display; stripping here keeps provider replay compatible without
+ * changing what users see in the chat bubble.
+ */
+export function stripReasoningFromServingMessages(messages: ServingChatMessage[]): boolean {
+  let changed = false;
+  for (let i = messages.length - 1; i >= 0; i--) {
+    const msg = messages[i];
+    if (!msg) continue;
+    if (msg.role === "reasoning") {
+      messages.splice(i, 1);
+      changed = true;
+      continue;
+    }
+    if (msg.reasoning !== undefined) {
+      delete msg.reasoning;
+      changed = true;
+    }
+    if (msg.reasoning_content !== undefined) {
+      delete msg.reasoning_content;
+      changed = true;
+    }
+    if (!Array.isArray(msg.content)) continue;
+    const filtered = msg.content.filter((part) => {
+      const type = part?.type;
+      if (typeof type === "string" && REASONING_PART_TYPES.has(type)) {
+        changed = true;
+        return false;
+      }
+      return true;
+    });
+    if (filtered.length !== msg.content.length) {
+      msg.content = filtered;
+    }
+    const hasToolCalls = Array.isArray(msg.tool_calls) && msg.tool_calls.length > 0;
+    if (msg.role === "assistant" && !hasToolCalls && isEmptyServingContent(msg.content)) {
+      messages.splice(i, 1);
+      changed = true;
+    }
+  }
+  return changed;
+}
+
+/**
+ * Repair a Mastra/AI SDK message replay that Databricks-hosted Claude
+ * rejects with `"This model does not support assistant message
+ * prefill. The conversation must end with a user message."`.
+ *
+ * The bug pattern: when an assistant turn streams text *and* a
+ * `tool_call`, the AI SDK persists them as two separate assistant
+ * entries (text-only and tool-call-only). On the next agent step the
+ * tool-call entry is replayed *before* the tool result and the
+ * text entry is replayed *after* it, so the conversation ends with a
+ * trailing assistant text message. Anthropic interprets that as a
+ * prefill request and rejects it on Databricks (the upstream Bedrock
+ * route disallows prefill).
+ *
+ * Fix: when the last message is an assistant text with no `tool_calls`
+ * and the chain immediately before it is `assistant(tool_calls=...)`
+ * followed only by `tool(...)` results, fold the trailing text back
+ * into the `content` of that opening assistant and drop the duplicate.
+ */
+export function repairAssistantPrefill(messages: ServingChatMessage[]): boolean {
+  if (messages.length < 2) return false;
+  const last = messages[messages.length - 1];
+  if (!last || last.role !== "assistant" || (last.tool_calls && last.tool_calls.length > 0)) {
+    return false;
+  }
+
+  let i = messages.length - 2;
+  while (i >= 0 && messages[i]?.role === "tool") i--;
+  if (i < 0) return false;
+  const opener = messages[i];
+  if (
+    !opener ||
+    opener.role !== "assistant" ||
+    !opener.tool_calls ||
+    opener.tool_calls.length === 0
+  ) {
+    return false;
+  }
+
+  const merged = [
+    string.trimToNull(textFromServingContent(opener.content)),
+    string.trimToNull(textFromServingContent(last.content)),
+  ]
+    .filter((s): s is string => s !== null)
+    .join("\n\n");
+  opener.content = merged;
+  messages.pop();
+  return true;
+}
+
+function textFromServingContent(content: ServingChatMessage["content"]): string {
+  if (typeof content === "string") return content;
+  if (!Array.isArray(content)) return "";
+  return content
+    .filter((part) => part?.type === "text" && typeof part.text === "string")
+    .map((part) => part.text as string)
+    .join("\n\n");
+}
+
+function isEmptyServingContent(content: ServingChatMessage["content"]): boolean {
+  if (content === undefined) return true;
+  if (typeof content === "string") return content.trim().length === 0;
+  if (!Array.isArray(content)) return true;
+  return content.every((part) => {
+    if (part?.type === "text") {
+      return typeof part.text !== "string" || part.text.trim().length === 0;
+    }
+    return false;
+  });
+}

package/src/serving.ts

@@ -0,0 +1,97 @@
+/**
+ * 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 { override } from "@dbx-tools/shared-mastra";
+import { serving as nodeServing } from "@dbx-tools/model";
+
+import type { MastraPluginConfig } from "./config";
+import { string } from "@dbx-tools/shared-core";
+
+/**
+ * `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";
+
+/**
+ * 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 = string.firstNonEmpty(
+      headers[override.MODEL_OVERRIDE_HEADER] ??
+        headers[override.MODEL_OVERRIDE_HEADER.toLowerCase()],
+    );
+    if (headerVal) return headerVal;
+  }
+  if (req.query) {
+    const queryVal = string.firstNonEmpty(req.query[override.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 override.MODEL_OVERRIDE_BODY_FIELDS) {
+      const bodyVal = string.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 ?? nodeServing.DEFAULT_MODEL_CACHE_TTL_MS,
+    threshold: config.modelFuzzyThreshold ?? nodeServing.DEFAULT_FUZZY_THRESHOLD,
+    fuzzy: config.modelFuzzyMatch !== false,
+    allowOverride: config.modelOverride !== false,
+    fallbacks: config.defaultModelFallbacks ?? [],
+  };
+}

package/src/statement.ts

@@ -0,0 +1,89 @@
+/**
+ * Databricks Statement Execution helpers for the Mastra plugin.
+ *
+ * Wraps `client.statementExecution.getStatement` with the shape and
+ * size 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, plus a
+ * hard row cap callers clamp `limit` to so a runaway result set can't
+ * hose a response. Upstream 404s are detected via
+ * `error.errorContext(err).notFound` at the call sites.
+ *
+ * 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 { WorkspaceClient } from "@databricks/sdk-experimental";
+import type { GenieDatasetData } from "@dbx-tools/shared-mastra";
+import { databricks } from "@dbx-tools/appkit";
+
+/**
+ * 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 ? databricks.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,
+  };
+}

package/src/storage-schema.ts

@@ -0,0 +1,41 @@
+import { string } from "@dbx-tools/shared-core";
+/**
+ * Derive a Postgres-safe schema name for per-agent Mastra storage.
+ *
+ * Agent ids are often kebab-case route segments (`data-mesh-book-assistant`)
+ * but {@link PostgresStore} validates `schemaName` with Mastra's
+ * `parseSqlIdentifier` (letter/underscore start, `[A-Za-z0-9_]` only, max 63).
+ */
+
+const SCHEMA_PREFIX = "mastra_";
+const MAX_PG_IDENTIFIER_LEN = 63;
+
+/**
+ * Default Lakebase schema for one agent's thread/message store:
+ * `mastra_<sanitized-agent-id>`.
+ */
+export function agentStorageSchemaName(agentId: string): string {
+  const maxSlugLen = MAX_PG_IDENTIFIER_LEN - SCHEMA_PREFIX.length;
+  const slug = string.toIdentifierWithOptions(
+    {
+      delimiter: "_",
+      maxLength: maxSlugLen,
+      truncateStrategy: "hash",
+      truncateHashLength: 6,
+    },
+    agentId,
+  );
+  const body =
+    slug ||
+    string.toIdentifierWithOptions(
+      {
+        delimiter: "_",
+        maxLength: maxSlugLen,
+        truncateStrategy: "hash",
+        truncateHashLength: 6,
+      },
+      "agent",
+      agentId,
+    );
+  return `${SCHEMA_PREFIX}${body}`;
+}

package/src/summarize.ts

@@ -0,0 +1,176 @@
+/**
+ * Small-tier summarization for the Mastra plugin.
+ *
+ * Two surfaces, both backed by the fast / small chat tier
+ * ({@link model.ModelClass.ChatFast}) resolved through the same
+ * `/serving-endpoints` pipeline as the main agents:
+ *
+ *   - A dedicated `summarize` tool (see {@link buildSummarizeTool})
+ *     agents can call to condense arbitrary text without burning the
+ *     heavyweight chat model.
+ *   - The model + instructions Mastra's memory uses to auto-name
+ *     conversation threads (`generateTitle`), so titling reuses the
+ *     same small tier rather than the agent's primary model.
+ *
+ * Mirrors the chart-planner wiring in `chart.ts`: a per-config cached
+ * `Agent` on the fast tier, invoked via `agent.generate(...)` inside
+ * the active `asUser` scope so tokens stay user-scoped.
+ */
+
+import { model } from "@dbx-tools/shared-model";
+import { Agent } from "@mastra/core/agent";
+import type { MastraModelConfig } from "@mastra/core/llm";
+import type { RequestContext } from "@mastra/core/request-context";
+import { createTool } from "@mastra/core/tools";
+import { z } from "zod";
+
+import type { MastraPluginConfig } from "./config";
+import { buildModel } from "./model";
+import { string } from "@dbx-tools/shared-core";
+
+/** Fast / small chat tier used for both titling and summaries. */
+const SUMMARY_MODEL_CLASS = model.ModelClass.ChatFast;
+
+/** System prompt for the summarizer agent (and the `summarize` tool). */
+const SUMMARIZER_INSTRUCTIONS = [
+  "You are a summarization engine.",
+  "Given a block of text, produce a faithful, concise summary of it.",
+  "Default to a few sentences; follow any length guidance the caller gives.",
+  "Plain prose. No preamble, no headers, no bullet points unless asked.",
+  "Never use emojis. Use hyphens (-) only, never em dashes or en dashes.",
+  "Never add information, opinions, or details not present in the input.",
+  "Output only the summary.",
+].join("\n");
+
+/**
+ * Instructions Mastra's `generateTitle` hands the small-tier model to
+ * name a conversation thread from its opening turn. Kept terse so the
+ * model returns a bare title with no decoration.
+ */
+export const TITLE_INSTRUCTIONS = [
+  "Generate a short, specific title for this conversation, 3 to 6 words.",
+  "Capture the user's topic, not the assistant's response.",
+  "Plain text only: no surrounding quotes, no trailing punctuation,",
+  "no emojis, and no em dashes. Output only the title.",
+].join(" ");
+
+/**
+ * Resolve the small-tier model for summarization / titling. Reused by
+ * both the summarizer agent and Mastra memory's `generateTitle`.
+ *
+ * Returned as a `requestContext`-taking function (a Mastra
+ * `DynamicArgument<MastraModelConfig>`) so each call mints user-scoped
+ * tokens via {@link buildModel}, exactly like the primary agents.
+ */
+export function summaryModel(
+  config: MastraPluginConfig,
+): (args: { requestContext: RequestContext }) => Promise<MastraModelConfig> {
+  return ({ requestContext }) =>
+    buildModel(config, requestContext, { modelClass: SUMMARY_MODEL_CLASS });
+}
+
+/**
+ * One summarizer `Agent` per plugin config, cached on config-object
+ * identity so a hot tool path doesn't pay the constructor cost each
+ * call. `WeakMap` lets retired configs (e.g. test reconfigurations)
+ * release their agent without manual eviction.
+ */
+const summarizerAgents = new WeakMap<MastraPluginConfig, Agent>();
+
+function getSummarizerAgent(config: MastraPluginConfig): Agent {
+  let agent = summarizerAgents.get(config);
+  if (!agent) {
+    agent = new Agent({
+      id: "summarizer",
+      name: "Summarizer",
+      description: "Condenses text into a short summary using a fast, small model.",
+      instructions: SUMMARIZER_INSTRUCTIONS,
+      model: summaryModel(config),
+    });
+    summarizerAgents.set(config, agent);
+  }
+  return agent;
+}
+
+/** Inputs accepted by the `summarize` tool. */
+const summarizeInput = z.object({
+  text: z.string().min(1).describe("The text to summarize."),
+  instructions: z
+    .string()
+    .optional()
+    .describe(
+      "Optional extra guidance for the summary, e.g. 'one sentence' or 'list the action items'.",
+    ),
+  maxWords: z
+    .number()
+    .int()
+    .positive()
+    .optional()
+    .describe("Optional soft cap on the summary length, in words."),
+});
+
+/** Options accepted by {@link summarizeText}. */
+export interface SummarizeOptions {
+  /** Extra guidance for the summary (length, focus, format). */
+  instructions?: string;
+  /** Soft cap on summary length, in words. */
+  maxWords?: number;
+  /** Active request context, so the model resolver mints user-scoped tokens. */
+  requestContext?: RequestContext;
+  /** Abort signal bridged from the calling tool / request. */
+  abortSignal?: AbortSignal;
+}
+
+/**
+ * Summarize `text` with the small-tier summarizer agent, returning the
+ * trimmed summary string. Throws on model failure (the caller decides
+ * how to degrade).
+ */
+export async function summarizeText(
+  config: MastraPluginConfig,
+  text: string,
+  options: SummarizeOptions = {},
+): Promise<string> {
+  const { instructions, maxWords, requestContext, abortSignal } = options;
+  const prompt = string.toDescription({
+    ...(instructions ? { Guidance: instructions } : {}),
+    ...(maxWords !== undefined ? { "Max length (words)": String(maxWords) } : {}),
+    Text: text,
+  });
+  const result = await getSummarizerAgent(config).generate(prompt, {
+    ...(requestContext ? { requestContext } : {}),
+    ...(abortSignal ? { abortSignal } : {}),
+  });
+  return result.text.trim();
+}
+
+/**
+ * Build the `summarize` tool. Exposed as an ambient system tool (like
+ * `render_data`) so every agent can offload condensing long content,
+ * notes, transcripts, or bulky tool results to the fast tier instead of
+ * spending its primary chat model. The tool reads the live
+ * `requestContext` / `abortSignal` off the Mastra execution context so
+ * its model call stays user-scoped and cancels with the turn.
+ */
+export function buildSummarizeTool(config: MastraPluginConfig) {
+  return createTool({
+    id: "summarize",
+    description:
+      "Summarize a block of text using a fast, small model. Use it to condense long " +
+      "content, notes, transcripts, or tool results into a short summary without " +
+      "spending the main chat model.",
+    inputSchema: summarizeInput,
+    execute: async (input, context) => {
+      const { text, instructions, maxWords } = input as z.infer<typeof summarizeInput>;
+      const ctx = context as
+        { requestContext?: RequestContext; abortSignal?: AbortSignal } | undefined;
+      const summary = await summarizeText(config, text, {
+        ...(instructions ? { instructions } : {}),
+        ...(maxWords !== undefined ? { maxWords } : {}),
+        ...(ctx?.requestContext ? { requestContext: ctx.requestContext } : {}),
+        ...(ctx?.abortSignal ? { abortSignal: ctx.abortSignal } : {}),
+      });
+      return { summary };
+    },
+  });
+}