@dbx-tools/appkit-mastra 0.1.40 → 0.1.42

package/src/server.ts

@@ -1,8 +1,8 @@
 /**
  * Express-layer plumbing for the Mastra plugin: a `MastraServer` that
  * stamps the per-request `RequestContext`, and a route-patch middleware
- * that lets `@mastra/ai-sdk` `chatRoute` work behind an Express mount
- * point.
+ * that lets the plugin's custom API routes (e.g. `historyRoute`) work
+ * behind an Express mount point.
  */
 
 import { getExecutionContext } from "@databricks/appkit";
@@ -172,30 +172,23 @@ export class MastraServer extends MastraServerExpress {
 }
 
 /**
- * Patches around `@mastra/express`'s custom-route dispatcher so
- * `chatRoute` works when `MastraServer` is hosted on an Express subapp
- * mounted under a parent path (e.g. `/api/mastra`).
+ * 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`).
  *
- * Two concerns:
- *
- * 1. 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 `chatRoute` paths
- *    (for example `/route/chat`), so the absolute URL never matches
- *    until we overwrite `originalUrl` for `/route` and `/route/*` to
- *    the mount-relative path.
- *
- * 2. `memory.resource` must be the authenticated user, not whatever the
- *    client posts. The custom-route forwarder re-serializes `req.body`
- *    into the Request body it hands Hono, so mutating the parsed body
- *    here would propagate into `handleChatStream`'s params (kept for
- *    future use; `express.json()` runs first so `req.body` is parsed).
+ * 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 isChat = req.path === "/route" || req.path.startsWith("/route/");
-    if (!isChat) return next();
+    const isCustomRoute = req.path === "/route" || req.path.startsWith("/route/");
+    if (!isCustomRoute) return next();
     req.originalUrl = req.path;
     next();
   });

package/src/serving.ts

@@ -1,37 +1,19 @@
 /**
- * Dynamic model resolution against Databricks Model Serving.
+ * Mastra-specific glue over the generic `@dbx-tools/model` toolkit.
  *
- * 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.
+ * The live `/serving-endpoints` catalogue access, fuzzy name
+ * resolution, and tier/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 { logUtils, stringUtils } from "@dbx-tools/shared";
-import Fuse from "fuse.js";
+import { DEFAULT_FUZZY_THRESHOLD, DEFAULT_MODEL_CACHE_TTL_MS } from "@dbx-tools/model";
+import { stringUtils } from "@dbx-tools/shared";
 
-import type { ServingEndpointSummary } from "@dbx-tools/appkit-mastra-shared";
 import type { MastraPluginConfig } from "./config.js";
 
-export type { ServingEndpointSummary };
-
-const log = logUtils.logger("mastra/serving");
-
-/**
- * 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"];
-
 /**
  * `RequestContext` key under which {@link MastraServer} stores the
  * per-request model override (header / query / body). `model.ts`
@@ -48,187 +30,6 @@ 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 startedAt = Date.now();
-  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 } : {}),
-    });
-  }
-  log.debug("listed", { count: out.length, elapsedMs: Date.now() - startedAt });
-  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) {
-    log.debug("resolve:no-endpoints", { input });
-    return { modelId: input, matched: false };
-  }
-  for (const ep of endpoints) {
-    if (ep.name === input) {
-      log.debug("resolve:exact", { 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) {
-    log.debug("resolve:empty-tokens", { input });
-    return { modelId: input, matched: false };
-  }
-  const results = fuse.search(query);
-  const best = results[0];
-  if (best?.item.name && (best.score ?? 0) <= threshold) {
-    log.debug("resolve:fuzzy-match", {
-      input,
-      modelId: best.item.name,
-      score: best.score,
-    });
-    return { modelId: best.item.name, matched: true, score: best.score };
-  }
-  log.debug("resolve:no-match", {
-    input,
-    bestScore: best?.score,
-    threshold,
-  });
-  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
@@ -278,18 +79,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;
@@ -297,10 +95,10 @@ export function resolveServingConfig(
   fallbacks: readonly string[];
 } {
   return {
-    ttlMs: config.modelCacheTtlMs ?? DEFAULT_TTL_MS,
+    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 ?? defaultFallbacks,
+    fallbacks: config.defaultModelFallbacks ?? [],
   };
 }

package/src/writer.ts

@@ -10,7 +10,7 @@
  * away can't crash a tool mid-flight.
  */
 
-import type { MinimalWriter } from "@dbx-tools/appkit-mastra-shared";
+import type { MastraWriter } from "@dbx-tools/appkit-mastra-shared";
 import { commonUtils, type logUtils } from "@dbx-tools/shared";
 
 /**
@@ -24,7 +24,7 @@ import { commonUtils, type logUtils } from "@dbx-tools/shared";
  */
 export async function safeWrite(
   log: logUtils.Logger,
-  writer: MinimalWriter | undefined,
+  writer: MastraWriter | undefined,
   chunk: unknown,
   context: Record<string, unknown> = {},
 ): Promise<void> {

package/dist/src/intercept.d.ts

@@ -1,48 +0,0 @@
-/**
- * Server-side Server-Sent-Events frame interceptor for the native
- * Mastra agent `/stream` endpoint.
- *
- * The Mastra stream (`@mastra/express`) emits one `data: <json>\n\n`
- * SSE frame per chunk, where the decoded JSON carries a
- * discriminating `type` (`text-delta`, `reasoning-delta`,
- * `tool-call`, `step-finish`, ...) and is terminated by a
- * `data: [DONE]` sentinel. {@link installStreamEventInterceptor} wraps
- * an Express response and runs a caller-supplied
- * {@link StreamFrameInterceptor} over each decoded frame so the caller
- * can keep, rewrite, or drop it; every non-`data:` byte (SSE
- * keep-alive comments, the `[DONE]` sentinel, blank padding) passes
- * through verbatim.
- *
- * The wrap is inert unless the response turns out to be a streamed
- * `200 text/event-stream` body, so non-streaming JSON routes and
- * error responses are never touched.
- */
-import type express from "express";
-/**
- * Per-frame interceptor invoked with the `JSON.parse`d object of every
- * `data: {json}` frame on the Mastra `/stream` response. The return
- * value decides the frame's fate:
- *
- * - `true`: forward the frame's original bytes unchanged (no
- *   re-serialization).
- * - `false`: drop the frame entirely.
- * - `{ replace }`: re-emit the frame with the original `data:` prefix
- *   and trailing whitespace preserved and only the JSON body replaced
- *   by `JSON.stringify(replace)`.
- */
-export type StreamFrameInterceptor = (chunk: unknown) => {
-    replace: unknown;
-} | true | false;
-/**
- * Wrap `res.write` / `res.end` so each SSE frame on a streamed Mastra
- * response is run through `interceptor` (keep / rewrite / drop).
- *
- * Interception only engages once the response is confirmed to be a
- * `200 text/event-stream` body (decided lazily on the first write,
- * after the handler has set status + content-type); any other
- * response streams through unchanged. A {@link StringDecoder}
- * reassembles multi-byte UTF-8 sequences that straddle chunk
- * boundaries, and a running buffer holds the trailing partial frame
- * until its `\n\n` arrives.
- */
-export declare function installStreamEventInterceptor(res: express.Response, interceptor: StreamFrameInterceptor): void;

package/dist/src/intercept.js

@@ -1,167 +0,0 @@
-/**
- * Server-side Server-Sent-Events frame interceptor for the native
- * Mastra agent `/stream` endpoint.
- *
- * The Mastra stream (`@mastra/express`) emits one `data: <json>\n\n`
- * SSE frame per chunk, where the decoded JSON carries a
- * discriminating `type` (`text-delta`, `reasoning-delta`,
- * `tool-call`, `step-finish`, ...) and is terminated by a
- * `data: [DONE]` sentinel. {@link installStreamEventInterceptor} wraps
- * an Express response and runs a caller-supplied
- * {@link StreamFrameInterceptor} over each decoded frame so the caller
- * can keep, rewrite, or drop it; every non-`data:` byte (SSE
- * keep-alive comments, the `[DONE]` sentinel, blank padding) passes
- * through verbatim.
- *
- * The wrap is inert unless the response turns out to be a streamed
- * `200 text/event-stream` body, so non-streaming JSON routes and
- * error responses are never touched.
- */
-import { StringDecoder } from "node:string_decoder";
-/** SSE event separator used by the Mastra stream wire format. */
-const FRAME_DELIMITER = "\n\n";
-/**
- * Matches a whole SSE frame whose `data:` payload is a JSON object,
- * capturing it in three parts: the `data:` field prefix (including
- * whatever inter-token whitespace the producer used), the JSON object
- * body, and any trailing whitespace. Anchored end-to-end so it only
- * matches a single self-contained object frame; anything else (SSE
- * comments, the `[DONE]` sentinel, text/array payloads) fails the
- * match and is forwarded verbatim without a parse attempt. On rewrite
- * the captured prefix/suffix are re-emitted as-is so only the body
- * changes.
- */
-const DATA_OBJECT_RE = /^(data:\s*)(\{[\s\S]*\})(\s*)$/;
-/**
- * True when `res` is a `200 text/event-stream` body - the only
- * response shape we intercept. Content-type based so it tracks the
- * Mastra `/stream` route without hard-coding the path, and naturally
- * skips the JSON routes (history, models, charts, statements) and any
- * non-200 error response.
- */
-function isEventStream(res) {
-    if (res.statusCode !== 200)
-        return false;
-    const contentType = String(res.getHeader("content-type") ?? "").toLowerCase();
-    return contentType.includes("text/event-stream");
-}
-/**
- * Run `interceptor` over a single SSE frame. Forwards anything that
- * isn't a parseable `data: {json}` object frame verbatim (comments,
- * the `[DONE]` sentinel, blank padding, unparseable / non-object
- * payloads) so the interceptor only ever sees real chunk objects.
- * Returns the frame string to emit, or `undefined` to drop the frame.
- */
-function interceptFrame(frame, interceptor) {
-    const match = DATA_OBJECT_RE.exec(frame);
-    if (!match)
-        return frame;
-    const [, prefix, body, suffix] = match;
-    let parsed;
-    try {
-        parsed = JSON.parse(body);
-    }
-    catch {
-        // Forward unparseable payloads untouched - feeding a partial /
-        // malformed frame to the interceptor risks corrupting bytes the
-        // client could still have handled.
-        return frame;
-    }
-    const result = interceptor(parsed);
-    if (result === true)
-        return frame;
-    if (result === false)
-        return undefined;
-    // Preserve the exact captured prefix/suffix so only the JSON body
-    // changes; the framing the client parses stays byte-identical.
-    return `${prefix}${JSON.stringify(result.replace)}${suffix}`;
-}
-/**
- * Wrap `res.write` / `res.end` so each SSE frame on a streamed Mastra
- * response is run through `interceptor` (keep / rewrite / drop).
- *
- * Interception only engages once the response is confirmed to be a
- * `200 text/event-stream` body (decided lazily on the first write,
- * after the handler has set status + content-type); any other
- * response streams through unchanged. A {@link StringDecoder}
- * reassembles multi-byte UTF-8 sequences that straddle chunk
- * boundaries, and a running buffer holds the trailing partial frame
- * until its `\n\n` arrives.
- */
-export function installStreamEventInterceptor(res, interceptor) {
-    const origWrite = res.write.bind(res);
-    const origEnd = res.end.bind(res);
-    const decoder = new StringDecoder("utf8");
-    let engaged;
-    let buffer = "";
-    // Append decoded text, emit every whole `\n\n`-delimited frame the
-    // interceptor keeps, and retain any trailing partial frame in
-    // `buffer`. On `final`, the residual buffer is flushed as the last
-    // frame so a stream that doesn't end on a delimiter isn't dropped.
-    const intercept = (text, final) => {
-        buffer += text;
-        const out = [];
-        let idx;
-        while ((idx = buffer.indexOf(FRAME_DELIMITER)) !== -1) {
-            const frame = buffer.slice(0, idx);
-            buffer = buffer.slice(idx + FRAME_DELIMITER.length);
-            const kept = interceptFrame(frame, interceptor);
-            if (kept !== undefined) {
-                out.push(kept, FRAME_DELIMITER);
-            }
-        }
-        if (final && buffer.length > 0) {
-            const kept = interceptFrame(buffer, interceptor);
-            if (kept !== undefined)
-                out.push(kept);
-            buffer = "";
-        }
-        return out.length === 0 ? "" : out.join("");
-    };
-    const toText = (chunk) => {
-        if (typeof chunk === "string")
-            return chunk;
-        if (Buffer.isBuffer(chunk))
-            return decoder.write(chunk);
-        if (chunk instanceof Uint8Array) {
-            return decoder.write(Buffer.from(chunk.buffer, chunk.byteOffset, chunk.byteLength));
-        }
-        return "";
-    };
-    res.write = function patchedWrite(chunk, encodingOrCb, cb) {
-        if (engaged === undefined)
-            engaged = isEventStream(res);
-        if (!engaged) {
-            return origWrite(chunk, encodingOrCb, cb);
-        }
-        const callback = typeof encodingOrCb === "function" ? encodingOrCb : cb;
-        const out = intercept(toText(chunk), false);
-        if (out.length === 0) {
-            // The chunk only advanced a partial frame; nothing to forward
-            // yet. Honor the write callback so backpressure-aware writers
-            // don't stall waiting on an ack that never comes.
-            if (callback)
-                queueMicrotask(() => callback());
-            return true;
-        }
-        return origWrite(out, callback);
-    };
-    res.end = function patchedEnd(chunk, encodingOrCb, cb) {
-        if (engaged === undefined)
-            engaged = isEventStream(res);
-        if (!engaged) {
-            return origEnd(chunk, encodingOrCb, cb);
-        }
-        const callback = typeof chunk === "function"
-            ? chunk
-            : typeof encodingOrCb === "function"
-                ? encodingOrCb
-                : cb;
-        const text = (typeof chunk !== "function" && chunk !== undefined ? toText(chunk) : "") +
-            decoder.end();
-        const out = intercept(text, true);
-        if (out.length > 0)
-            origWrite(out);
-        return origEnd(callback);
-    };
-}