@dbx-tools/appkit-mastra 0.1.5 → 0.1.10

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

@@ -1,41 +1,19 @@
 /**
- * Dynamic model resolution against Databricks Model Serving.
+ * Mastra-specific glue over the generic `@dbx-tools/model` toolkit.
  *
- * 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`.
+ * 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 { CacheManager, type getExecutionContext } from "@databricks/appkit";
-import { stringUtils } from "@dbx-tools/appkit-shared";
-import Fuse from "fuse.js";
-
-import type { ServingEndpointSummary } from "@dbx-tools/appkit-mastra-shared";
-import type { MastraPluginConfig } from "./config.js";
+import { override } from "@dbx-tools/shared-mastra";
+import { serving as nodeServing } from "@dbx-tools/model";
 
-export type { ServingEndpointSummary };
-
-/**
- * Structural type for the Databricks workspace client. Derived from
- * AppKit's `ExecutionContext` so this module doesn't take a direct
- * dependency on `@databricks/sdk-experimental`; the dep flows in
- * transitively through `@databricks/appkit`.
- */
-type WorkspaceClientLike = ReturnType<typeof getExecutionContext>["client"];
+import type { MastraPluginConfig } from "./config";
+import { string } from "@dbx-tools/shared-core";
 
 /**
  * `RequestContext` key under which {@link MastraServer} stores the
@@ -44,179 +22,6 @@ type WorkspaceClientLike = ReturnType<typeof getExecutionContext>["client"];
  */
 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;
-
-/** Default TTL for the in-memory endpoint cache. Matches the Databricks SDK's session lifetime budget. */
-const DEFAULT_TTL_MS = 5 * 60 * 1000;
-
-/** Default Fuse.js score threshold below which a fuzzy match is accepted. */
-const DEFAULT_FUZZY_THRESHOLD = 0.4;
-
-/** Cache key parts under which endpoint listings are stored. */
-const CACHE_KEY_NAMESPACE = "mastra:serving-endpoints";
-
-/**
- * Stable `userKey` arg for AppKit's `CacheManager.getOrExecute`.
- * Endpoint visibility is effectively workspace-scoped (we cache by
- * host in the key parts), so a single shared key lets every user of
- * the same workspace share one cached fetch and coalesce on the
- * in-flight promise. Permissions can differ in theory, but the
- * Foundation Model API catalogue is the same view for every caller.
- */
-const SHARED_USER_KEY = "mastra-shared";
-
-/**
- * List Model Serving endpoints for the workspace owning `client`,
- * routed through AppKit's `CacheManager`. The manager gives us
- * everything `cachetools.TTLCache` provides plus what
- * `cachetools-async` adds on top: per-entry TTL, in-flight request
- * coalescing (concurrent callers share one fetch via the manager's
- * internal `inFlightRequests` map), bounded size, telemetry spans
- * (`cache.getOrExecute`), and optional Lakebase persistence so the
- * catalogue survives restarts when the lakebase plugin is wired up.
- *
- * Returns plain {@link ServingEndpointSummary} objects (a stable
- * subset of the SDK type) so cache hits never expose stale SDK
- * internals. Errors from `CacheManager` or the SDK fetch propagate
- * to the caller - we don't swallow them so users see the real
- * auth / network issue.
- *
- * @param host - Workspace host used as the cache key. Pass the value
- *   resolved from `client.config.getHost()` so multi-host apps share
- *   one entry per workspace.
- * @param opts.ttlMs - Override the default TTL just for this call.
- *   Forwarded to `CacheManager` as seconds.
- */
-export async function listServingEndpoints(
-  client: WorkspaceClientLike,
-  host: string,
-  opts: { ttlMs?: number } = {},
-): Promise<ServingEndpointSummary[]> {
-  const ttlSec = Math.max(1, Math.round((opts.ttlMs ?? DEFAULT_TTL_MS) / 1000));
-  return CacheManager.getInstanceSync().getOrExecute(
-    [CACHE_KEY_NAMESPACE, host],
-    () => fetchEndpoints(client),
-    SHARED_USER_KEY,
-    { ttl: ttlSec },
-  );
-}
-
-async function fetchEndpoints(
-  client: WorkspaceClientLike,
-): Promise<ServingEndpointSummary[]> {
-  const out: ServingEndpointSummary[] = [];
-  for await (const ep of client.servingEndpoints.list()) {
-    if (!ep.name) continue;
-    out.push({
-      name: ep.name,
-      ...(ep.task !== undefined ? { task: ep.task } : {}),
-      ...(ep.state?.ready !== undefined ? { state: String(ep.state.ready) } : {}),
-      ...(ep.description !== undefined ? { description: ep.description } : {}),
-    });
-  }
-  return out;
-}
-
-/**
- * Force-evict cached endpoint listings via AppKit's `CacheManager`.
- * With a `host` deletes that one workspace's entry; without one
- * clears every cache entry on the manager (since `CacheManager`
- * doesn't expose a namespace-scoped clear, this is the brute-force
- * path - fine for tests, avoid in steady-state code).
- */
-export async function clearServingEndpointsCache(host?: string): Promise<void> {
-  const cache = CacheManager.getInstanceSync();
-  if (host) {
-    const key = cache.generateKey([CACHE_KEY_NAMESPACE, host], SHARED_USER_KEY);
-    await cache.delete(key);
-  } else {
-    await cache.clear();
-  }
-}
-
-/**
- * Result of fuzzy-resolving a user-supplied model name against the
- * live endpoint list. `score` is Fuse.js's distance (`0` is exact,
- * `1` is no match); `matched` is `false` when the score exceeds the
- * configured threshold so callers can fall back to the original
- * input (Databricks will then return a clean 404).
- */
-export interface ResolvedModel {
-  modelId: string;
-  matched: boolean;
-  score?: number;
-}
-
-/** Options accepted by {@link resolveModelId}. */
-export interface ResolveModelOptions {
-  /** Fuse.js threshold (0 = exact, 1 = anything). Default `0.4`. */
-  threshold?: number;
-}
-
-/**
- * Snap a user-supplied model name to the closest configured serving
- * endpoint:
- *
- * 1. Exact name match wins immediately (no fuzzy needed).
- * 2. Otherwise the input is tokenized (dashes / underscores / spaces
- *    become separators) and fed through Fuse.js extended search,
- *    which AND-s each token with fuzzy matching enabled. This is the
- *    "tokenized fuzzy match" the user reaches for when they type
- *    `"claude sonnet"` instead of the full endpoint name.
- * 3. If the best Fuse score is above `threshold`, return the input
- *    unchanged and let the upstream call surface the 404. This keeps
- *    deliberate model ids (e.g. brand new endpoints) from being
- *    silently rewritten to a similar-looking neighbour.
- *
- * Pass an empty endpoint list to short-circuit fuzzy matching - the
- * input is returned verbatim. This is what {@link buildModel} does
- * when the workspace client can't be reached at resolve time.
- */
-export function resolveModelId(
-  input: string,
-  endpoints: readonly ServingEndpointSummary[],
-  opts: ResolveModelOptions = {},
-): ResolvedModel {
-  if (endpoints.length === 0) {
-    return { modelId: input, matched: false };
-  }
-  for (const ep of endpoints) {
-    if (ep.name === input) {
-      return { modelId: ep.name, matched: true, score: 0 };
-    }
-  }
-  const threshold = opts.threshold ?? DEFAULT_FUZZY_THRESHOLD;
-  const fuse = new Fuse(endpoints, {
-    keys: ["name"],
-    threshold,
-    ignoreLocation: true,
-    includeScore: true,
-    useExtendedSearch: true,
-    isCaseSensitive: false,
-  });
-  // Fuse 7.3 has no built-in tokenize hook; in extended search,
-  // space-separated tokens are AND-ed with fuzzy matching enabled. We
-  // lean on the shared tokenizer so the splitting rules stay
-  // consistent with the rest of the toolkit.
-  const query = Array.from(
-    stringUtils.tokenizeWithOptions({ lowerCase: true, camelCase: false }, input),
-  ).join(" ");
-  if (!query) return { modelId: input, matched: false };
-  const results = fuse.search(query);
-  const best = results[0];
-  if (best?.item.name && (best.score ?? 0) <= threshold) {
-    return { modelId: best.item.name, matched: true, score: best.score };
-  }
-  return { modelId: input, matched: false };
-}
-
 /**
  * Minimal Express-ish request shape used by {@link extractModelOverride}.
  * Keeps this module independent of `express` so the helper can be
@@ -245,19 +50,20 @@ export interface ModelOverrideRequest {
 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()],
+    const headerVal = string.firstNonEmpty(
+      headers[override.MODEL_OVERRIDE_HEADER] ??
+        headers[override.MODEL_OVERRIDE_HEADER.toLowerCase()],
     );
     if (headerVal) return headerVal;
   }
   if (req.query) {
-    const queryVal = stringUtils.firstNonEmpty(req.query[MODEL_OVERRIDE_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 MODEL_OVERRIDE_BODY_FIELDS) {
-      const bodyVal = stringUtils.firstNonEmpty(record[field]);
+    for (const field of override.MODEL_OVERRIDE_BODY_FIELDS) {
+      const bodyVal = string.firstNonEmpty(record[field]);
       if (bodyVal) return bodyVal;
     }
   }
@@ -266,18 +72,15 @@ export function extractModelOverride(req: ModelOverrideRequest): string | null {
 
 /**
  * Read the fuzzy-resolution config knobs off the plugin config with
- * defaults applied. Kept here so `buildModel` and the `/models` route
- * agree on what "enabled" means.
+ * `@dbx-tools/model` defaults applied. Kept here so `buildModel` and
+ * the `/models` route agree on what "enabled" means.
  *
- * `fallbacks` is the priority-ordered list `pickModelId` walks when
- * nothing explicit is set; defaults live in `model.ts`
- * (`FALLBACK_MODEL_IDS`) and are passed in by callers to avoid a
- * circular import between `serving.ts` and `model.ts`.
+ * `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,
-  defaultFallbacks: readonly string[] = [],
-): {
+export function resolveServingConfig(config: MastraPluginConfig): {
   ttlMs: number;
   threshold: number;
   fuzzy: boolean;
@@ -285,10 +88,10 @@ export function resolveServingConfig(
   fallbacks: readonly string[];
 } {
   return {
-    ttlMs: config.modelCacheTtlMs ?? DEFAULT_TTL_MS,
-    threshold: config.modelFuzzyThreshold ?? DEFAULT_FUZZY_THRESHOLD,
+    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 ?? defaultFallbacks,
+    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}`;
+}