@dbx-tools/appkit-mastra 0.1.5 → 0.1.10

package/dist/src/serving.d.ts

@@ -1,156 +0,0 @@
-/**
- * Dynamic model resolution against Databricks Model Serving.
- *
- * 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`.
- */
-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`
- * reads it before falling back to the agent / plugin default.
- */
-export declare const MASTRA_MODEL_OVERRIDE_KEY = "mastra__model_override";
-/** HTTP header inspected for a per-request model override. */
-export declare const MODEL_OVERRIDE_HEADER = "x-mastra-model";
-/** Query string parameter inspected for a per-request model override. */
-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
- * 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 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.
- *
- * `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`.
- */
-export declare function resolveServingConfig(config: MastraPluginConfig, defaultFallbacks?: readonly string[]): {
-    ttlMs: number;
-    threshold: number;
-    fuzzy: boolean;
-    allowOverride: boolean;
-    fallbacks: readonly string[];
-};

package/dist/src/serving.js

@@ -1,214 +0,0 @@
-/**
- * Dynamic model resolution against Databricks Model Serving.
- *
- * 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`.
- */
-import { CacheManager } from "@databricks/appkit";
-import { stringUtils } from "@dbx-tools/appkit-shared";
-import Fuse from "fuse.js";
-/**
- * `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";
-/** 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"];
-/** 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 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 } : {}),
-        });
-    }
-    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) {
-        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 };
-}
-/**
- * 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) {
-    const headers = req.headers;
-    if (headers) {
-        const headerVal = stringUtils.firstNonEmpty(headers[MODEL_OVERRIDE_HEADER] ?? headers[MODEL_OVERRIDE_HEADER.toLowerCase()]);
-        if (headerVal)
-            return headerVal;
-    }
-    if (req.query) {
-        const queryVal = stringUtils.firstNonEmpty(req.query[MODEL_OVERRIDE_QUERY]);
-        if (queryVal)
-            return queryVal;
-    }
-    if (req.body && typeof req.body === "object") {
-        const record = req.body;
-        for (const field of MODEL_OVERRIDE_BODY_FIELDS) {
-            const bodyVal = stringUtils.firstNonEmpty(record[field]);
-            if (bodyVal)
-                return bodyVal;
-        }
-    }
-    return 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.
- *
- * `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`.
- */
-export function resolveServingConfig(config, defaultFallbacks = []) {
-    return {
-        ttlMs: config.modelCacheTtlMs ?? DEFAULT_TTL_MS,
-        threshold: config.modelFuzzyThreshold ?? DEFAULT_FUZZY_THRESHOLD,
-        fuzzy: config.modelFuzzyMatch !== false,
-        allowOverride: config.modelOverride !== false,
-        fallbacks: config.defaultModelFallbacks ?? defaultFallbacks,
-    };
-}

package/src/render-chart-route.ts

@@ -1,141 +0,0 @@
-/**
- * Chart-render HTTP endpoint for the Mastra plugin.
- *
- * The `render_data` tool returns immediately with a `chartId` and
- * emits the dataset over `ctx.writer`; the client then POSTs that
- * dataset to this endpoint to actually run the chart-planner
- * agent and get back an Echarts `EChartsOption` JSON. Planning
- * happens out-of-band so the calling agent's response stream
- * doesn't sit idle waiting for it - the model can finish the
- * report while the client is still rendering charts.
- *
- * Auth flows through the standard Mastra middleware: the route
- * sits in the same dispatcher pipeline as `chatRoute` /
- * `historyRoute`, so by the time the handler runs the
- * `RequestContext` is populated with the workspace user and
- * the chart-planner's model resolver has the OBO token it
- * needs.
- */
-
-import type {
-  RenderChartRequest,
-  RenderChartResponse,
-} from "@dbx-tools/appkit-mastra-shared";
-import { registerApiRoute } from "@mastra/core/server";
-
-import { runChartPlanner } from "./chart.js";
-import type { MastraPluginConfig } from "./config.js";
-
-/** Hard cap so a misbehaving client can't hand us a million-row payload. */
-const MAX_ROWS = 5_000;
-/**
- * Hard cap on the JSON body the route accepts (in bytes). Mirrors
- * the same intent as {@link MAX_ROWS}: bound the chart-planner's
- * prompt size and protect against accidental denial-of-service
- * from a runaway tool that ships an enormous payload.
- */
-const MAX_BODY_BYTES = 2 * 1024 * 1024;
-
-/** Options accepted by {@link renderChartRoute}. */
-export interface RenderChartRouteOptions {
-  path: string;
-  config: MastraPluginConfig;
-}
-
-/**
- * Register a `POST <path>` Mastra custom API route that runs the
- * chart-planner agent against a dataset and returns an Echarts
- * `EChartsOption` JSON.
- *
- * Body shape: {@link RenderChartRequest}; response:
- * {@link RenderChartResponse}.
- */
-export function renderChartRoute(options: RenderChartRouteOptions) {
-  const { path, config } = options;
-  return registerApiRoute(path, {
-    method: "POST",
-    handler: async (c) => {
-      const requestContext = c.get("requestContext");
-
-      // Hono parses the body as JSON; we still validate shape /
-      // size since the tool's structured output is a contract,
-      // not a guarantee, and the route is publicly mountable.
-      const raw = (await c.req.json().catch(() => null)) as unknown;
-      const validation = validateBody(raw);
-      if ("error" in validation) {
-        return c.json({ error: validation.error }, 400);
-      }
-      const { title, description, data } = validation.body;
-
-      try {
-        const result = await runChartPlanner({
-          config,
-          ...(requestContext ? { requestContext } : {}),
-          title,
-          ...(description ? { description } : {}),
-          data,
-        });
-        const payload: RenderChartResponse = {
-          option: result.option,
-          chartType: result.chartType,
-        };
-        return c.json(payload);
-      } catch (err) {
-        const message = err instanceof Error ? err.message : String(err);
-        return c.json({ error: message }, 500);
-      }
-    },
-  });
-}
-
-type ValidationResult =
-  | { body: RenderChartRequest }
-  | { error: string };
-
-/**
- * Best-effort body validation. Surfaces a 400 for malformed input
- * instead of letting a downstream `.map` / `.length` blow up
- * inside the planner agent. Field-level shape mirrors
- * {@link RenderChartRequest}.
- */
-function validateBody(raw: unknown): ValidationResult {
-  if (!raw || typeof raw !== "object") {
-    return { error: "request body must be a JSON object" };
-  }
-  const r = raw as Record<string, unknown>;
-  const title = r.title;
-  if (typeof title !== "string" || title.length === 0) {
-    return { error: "`title` must be a non-empty string" };
-  }
-  if (r.description !== undefined && typeof r.description !== "string") {
-    return { error: "`description` must be a string when provided" };
-  }
-  if (!Array.isArray(r.data)) {
-    return { error: "`data` must be an array of row objects" };
-  }
-  if (r.data.length === 0) {
-    return { error: "`data` must contain at least one row" };
-  }
-  if (r.data.length > MAX_ROWS) {
-    return { error: `\`data\` exceeds the per-request limit of ${MAX_ROWS} rows` };
-  }
-  for (const [i, row] of r.data.entries()) {
-    if (!row || typeof row !== "object" || Array.isArray(row)) {
-      return { error: `data[${i}] must be a plain object` };
-    }
-  }
-  // Approximate body-size check; spares us pulling Buffer in.
-  const approximateBytes = JSON.stringify(r.data).length;
-  if (approximateBytes > MAX_BODY_BYTES) {
-    return {
-      error: `\`data\` exceeds the per-request size limit of ${MAX_BODY_BYTES} bytes`,
-    };
-  }
-  return {
-    body: {
-      title,
-      ...(typeof r.description === "string" ? { description: r.description } : {}),
-      data: r.data as Array<Record<string, unknown>>,
-    },
-  };
-}