@dbx-tools/appkit-mastra 0.1.41 → 0.1.48

package/dist/src/serving.d.ts

@@ -1,28 +1,14 @@
 /**
- * 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 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 { type getExecutionContext } from "@databricks/appkit";
-import type { ServingEndpointSummary } from "@dbx-tools/appkit-mastra-shared";
 import type { MastraPluginConfig } from "./config.js";
-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"];
 /**
  * `RequestContext` key under which {@link MastraServer} stores the
  * per-request model override (header / query / body). `model.ts`
@@ -35,76 +21,6 @@ export declare const MODEL_OVERRIDE_HEADER = "x-mastra-model";
 export declare const MODEL_OVERRIDE_QUERY = "model";
 /** Body fields (in priority order) inspected for a per-request model override. */
 export declare const MODEL_OVERRIDE_BODY_FIELDS: readonly ["model", "modelId"];
-/**
- * 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 declare function listServingEndpoints(client: WorkspaceClientLike, host: string, opts?: {
-    ttlMs?: number;
-}): Promise<ServingEndpointSummary[]>;
-/**
- * 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 declare function clearServingEndpointsCache(host?: string): Promise<void>;
-/**
- * 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 declare function resolveModelId(input: string, endpoints: readonly ServingEndpointSummary[], opts?: ResolveModelOptions): ResolvedModel;
 /**
  * Minimal Express-ish request shape used by {@link extractModelOverride}.
  * Keeps this module independent of `express` so the helper can be
@@ -132,15 +48,15 @@ export interface ModelOverrideRequest {
 export declare 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 declare function resolveServingConfig(config: MastraPluginConfig, defaultFallbacks?: readonly string[]): {
+export declare function resolveServingConfig(config: MastraPluginConfig): {
     ttlMs: number;
     threshold: number;
     fuzzy: boolean;

package/dist/src/serving.js

@@ -1,21 +1,15 @@
 /**
- * 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 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 } from "@databricks/appkit";
-import { logUtils, stringUtils } from "@dbx-tools/shared";
-import Fuse from "fuse.js";
-const log = logUtils.logger("mastra/serving");
+import { DEFAULT_FUZZY_THRESHOLD, DEFAULT_MODEL_CACHE_TTL_MS } from "@dbx-tools/model";
+import { stringUtils } from "@dbx-tools/shared";
 /**
  * `RequestContext` key under which {@link MastraServer} stores the
  * per-request model override (header / query / body). `model.ts`
@@ -28,145 +22,6 @@ export const MODEL_OVERRIDE_HEADER = "x-mastra-model";
 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"];
-/** 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, host, opts = {}) {
-    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) {
-    const startedAt = Date.now();
-    const out = [];
-    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) {
-    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();
-    }
-}
-/**
- * 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, endpoints, opts = {}) {
-    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 };
-}
 /**
  * Pull a model override out of a single HTTP request, checking
  * sources in priority order:
@@ -205,20 +60,20 @@ export function extractModelOverride(req) {
 }
 /**
  * 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, defaultFallbacks = []) {
+export function resolveServingConfig(config) {
     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/dist/src/writer.d.ts

@@ -9,7 +9,7 @@
  * swallowed so a cancelled request or a client that navigated
  * 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 { type logUtils } from "@dbx-tools/shared";
 /**
  * Best-effort `writer.write`. No-op when `writer` is undefined;
@@ -20,4 +20,4 @@ import { type logUtils } from "@dbx-tools/shared";
  *
  * Returns when the write resolves or rejects; never throws.
  */
-export declare function safeWrite(log: logUtils.Logger, writer: MinimalWriter | undefined, chunk: unknown, context?: Record<string, unknown>): Promise<void>;
+export declare function safeWrite(log: logUtils.Logger, writer: MastraWriter | undefined, chunk: unknown, context?: Record<string, unknown>): Promise<void>;