@dbx-tools/appkit-mastra 0.1.0

package/src/plugin.ts

@@ -0,0 +1,269 @@
+/**
+ * AppKit plugin that builds one or more Mastra `Agent` instances and
+ * mounts the `@mastra/express` server plus `@mastra/ai-sdk` `chatRoute`
+ * handlers. The UI message stream matches what `chatRoute()` emits, so
+ * the client can use `useChat()` from `@ai-sdk/react` without custom
+ * parsing.
+ *
+ * - Agents: registered through `config.agents` at plugin creation
+ *   ({@link MastraAgentDefinition}). Each entry's `tools` field accepts
+ *   either a plain record or a `(plugins) => tools` callback that gets
+ *   a typed sibling-plugin index ({@link MastraPlugins}). Omit
+ *   `config.agents` to get a single built-in `default` analyst.
+ * - Model: each agent call resolves a `MastraModelConfig` via
+ *   {@link buildModel} from `./model.js`. Per-agent `model` overrides
+ *   (`AgentConfig["model"]` or a `modelId` string) flow through
+ *   {@link buildAgents}.
+ * - Memory / storage: per-agent, built by {@link createMemoryBuilder}
+ *   from `./memory.js`. Both auto-default to `true` when the
+ *   `lakebase` plugin is registered (unless the caller passed
+ *   `false` or a custom config). Storage namespaces per agent via
+ *   `schemaName: "mastra_<agentId>"`; the vector store is a single
+ *   shared singleton across every agent.
+ * - Server: the Express subapp wiring lives in `./server.js`.
+ * - HTTP: AppKit mounts this plugin under `/api/mastra`. `chatRoute`
+ *   is registered at `/route/chat` (bound to `config.defaultAgent` or
+ *   the first registered id) and `/route/chat/:agentId`, so the
+ *   AI SDK transport URL is `/api/mastra/route/chat/<agentId>`.
+ */
+
+import {
+  genie,
+  getExecutionContext,
+  lakebase,
+  Plugin,
+  toPlugin,
+  type IAppRouter,
+  type PluginManifest,
+  type ResourceRequirement,
+} from "@databricks/appkit";
+import { logUtils, pluginUtils } from "@dbx-tools/appkit-shared";
+import { chatRoute } from "@mastra/ai-sdk";
+import type { Agent } from "@mastra/core/agent";
+import { Mastra } from "@mastra/core/mastra";
+import express from "express";
+
+import { buildAgents, FALLBACK_AGENT_ID, type BuiltAgents } from "./agents.js";
+import type { MastraClientConfig } from "@dbx-tools/appkit-mastra-shared";
+import type { MastraPluginConfig } from "./config.js";
+import { createMemoryBuilder, needsLakebase } from "./memory.js";
+import { attachRoutePatchMiddleware, MastraServer } from "./server.js";
+import {
+  clearServingEndpointsCache,
+  listServingEndpoints,
+  resolveServingConfig,
+  type ServingEndpointSummary,
+} from "./serving.js";
+
+const GENIE_MANIFEST = pluginUtils.data(genie).plugin.manifest;
+const LAKEBASE_MANIFEST = pluginUtils.data(lakebase).plugin.manifest;
+
+/**
+ * AppKit plugin (registered name: `mastra`) that hosts Mastra agents
+ * with optional Lakebase-backed memory and AI SDK chat routes under
+ * the plugin mount (typically `/api/mastra`).
+ */
+export class MastraPlugin extends Plugin<MastraPluginConfig> {
+  static manifest = {
+    name: "mastra",
+    displayName: "Mastra",
+    description:
+      "Builds a Mastra Agent with user-scoped workspace auth (asUser) " +
+      "and optional Postgres-backed Mastra Memory via the `lakebase` plugin.",
+    stability: "beta",
+    resources: {
+      required: [],
+      optional: [
+        ...GENIE_MANIFEST.resources.required,
+        ...LAKEBASE_MANIFEST.resources.required,
+      ],
+    },
+  } satisfies PluginManifest<"mastra">;
+
+  /**
+   * Tighten resource requirements based on which features are enabled.
+   * AppKit calls this at registration time (config-aware) so disabled
+   * features don't surface their resource asks to the host app.
+   */
+  static getResourceRequirements(config: MastraPluginConfig): ResourceRequirement[] {
+    const resources: ResourceRequirement[] = [];
+    const enabledManifests: PluginManifest<string>[] = [];
+
+    if (needsLakebase(config)) {
+      enabledManifests.push(LAKEBASE_MANIFEST);
+    }
+    for (const m of enabledManifests) {
+      for (const resource of m.resources.required) {
+        resources.push({ ...resource, required: true } as ResourceRequirement);
+      }
+    }
+    return resources;
+  }
+
+  private log = logUtils.logger(this);
+  private built: BuiltAgents | null = null;
+  private mastra: Mastra | null = null;
+  private mastraApp: express.Express | null = null;
+  private mastraServer: MastraServer | null = null;
+
+  override async setup(): Promise<void> {
+    // Wait until sibling plugins (e.g. `lakebase`) finish `setup()` so
+    // the lakebase pool is valid when storage/memory are enabled.
+    this.context?.onLifecycle("setup:complete", async () => {
+      this.applyLakebaseAutoDefaults();
+      this.log.info("setup:complete");
+      await this.buildAgentAndServer();
+    });
+  }
+
+  /**
+   * When the `lakebase` plugin is registered, auto-enable `storage`
+   * and `memory` unless the caller opted out explicitly (`false` or a
+   * custom config object). Run after `setup:complete` so the lookup
+   * is reliable: any plugin that registers itself synchronously is
+   * already in the registry by the time this fires.
+   */
+  private applyLakebaseAutoDefaults(): void {
+    const hasLakebase = pluginUtils.instance(this.context, lakebase) !== undefined;
+    if (!hasLakebase) return;
+    if (this.config.storage === undefined) this.config.storage = true;
+    if (this.config.memory === undefined) this.config.memory = true;
+  }
+
+  override exports() {
+    return {
+      /**
+       * Ids of every registered agent in registration order. Matches
+       * AppKit `agents.list()` so callers can iterate the registry the
+       * same way under both plugins.
+       */
+      list: (): string[] => Object.keys(this.built?.agents ?? {}),
+      /**
+       * Look up a registered agent by id. Returns `null` (not
+       * undefined) when unknown so call sites can early-return without
+       * a separate `in` check.
+       */
+      get: (id: string): Agent | null => this.built?.agents[id] ?? null,
+      /**
+       * The agent `chatRoute` binds to when the client doesn't name
+       * one. Resolves to `config.defaultAgent`, the first registered
+       * id, or the built-in `default` fallback.
+       */
+      getDefault: (): Agent | null =>
+        (this.built && this.built.agents[this.built.defaultAgentId]) ?? null,
+      /** Underlying Mastra instance for advanced use (custom routes etc.). */
+      getMastra: () => this.mastra,
+      /** Express subapp Mastra is mounted on; mostly for tests. */
+      getMastraServer: () => this.mastraServer,
+      /**
+       * Fetch the workspace's Model Serving endpoints (cached). Same
+       * payload the `GET /models` route returns; surfaced here so
+       * other plugins / scripts can introspect the catalogue without
+       * an HTTP round-trip. AppKit wraps this with `asUser(req)` for
+       * OBO scoping automatically.
+       */
+      listModels: (): Promise<ServingEndpointSummary[]> => this.listModels(),
+      /**
+       * Force-evict cached endpoint listings via AppKit's
+       * `CacheManager`. Useful in tests or right after an admin
+       * deploys a new endpoint and doesn't want to wait for the TTL.
+       * Returns the underlying `CacheManager.delete`/`clear` promise.
+       */
+      clearModelsCache: (host?: string): Promise<void> =>
+        clearServingEndpointsCache(host),
+    };
+  }
+
+  override clientConfig(): Record<string, unknown> {
+    // AppKit mounts every plugin at `/api/<plugin.name>`. `this.name`
+    // honors `config.name` overrides, so the published paths stay
+    // accurate if someone remounts the plugin under a custom id.
+    // Return widens to `Record<string, unknown>` to satisfy the
+    // base-class signature; consumers read it through the typed
+    // `MastraClientConfig` shape via `usePluginClientConfig<...>(...)`.
+    const basePath = `/api/${this.name}`;
+    const config: MastraClientConfig = {
+      basePath,
+      chatPath: `${basePath}/route/chat`,
+      chatPathTemplate: `${basePath}/route/chat/:agentId`,
+      modelsPath: `${basePath}/models`,
+      defaultAgent: this.built?.defaultAgentId ?? FALLBACK_AGENT_ID,
+      agents: Object.keys(this.built?.agents ?? {}),
+    };
+    return config as unknown as Record<string, unknown>;
+  }
+
+  override injectRoutes(router: IAppRouter): void {
+    // `GET /models` exposes the cached endpoint list so clients can
+    // populate model pickers, validate `?model=` choices, etc. Must
+    // be registered before the catch-all that forwards everything to
+    // the Mastra subapp. Errors propagate to Express's default error
+    // handler via `next(err)` so callers see the real SDK message.
+    router.get("/models", (req, res, next) => {
+      this.asUser(req)
+        .listModels()
+        .then((endpoints) => res.json({ endpoints }))
+        .catch(next);
+    });
+
+    router.use("", (req, res, next) => {
+      if (!this.mastraApp) return res.status(503).end();
+      return this.asUser(req).mastraApp!(req, res, next);
+    });
+  }
+
+  /**
+   * Implementation backing both the `/models` route and the
+   * `listModels` export. Runs inside the AppKit user-context proxy so
+   * `getExecutionContext()` returns the OBO-scoped client.
+   */
+  private async listModels(): Promise<ServingEndpointSummary[]> {
+    const client = getExecutionContext().client;
+    const host = (await client.config.getHost()).toString();
+    const serving = resolveServingConfig(this.config);
+    return listServingEndpoints(client, host, { ttlMs: serving.ttlMs });
+  }
+
+  private async buildAgentAndServer(): Promise<void> {
+    // Per-agent memory factory. The builder resolves the Lakebase pool
+    // lazily (on first agent that actually needs storage / vector) and
+    // caches both the pool and the shared `PgVector` singleton so
+    // registering N agents stays cheap. See `./memory.js`.
+    const memoryBuilder = needsLakebase(this.config)
+      ? createMemoryBuilder(this.config, this.context)
+      : undefined;
+
+    // Build every agent declared in `config.agents` (or the built-in
+    // fallback when none are declared). Each agent's `model` resolves
+    // workspace URL + bearer at call time so concurrent requests get
+    // distinct user identities; the `asUser(req)` scope around
+    // `handleChat` is what lets `getExecutionContext()` return the
+    // right user inside the resolver.
+    this.built = await buildAgents({
+      config: this.config,
+      context: this.context,
+      memoryBuilder,
+      log: this.log,
+    });
+
+    // `mastra.server.apiRoutes` is only honored by Mastra's standalone
+    // dev server. Since we're hosting Mastra inside our own Express
+    // subapp via `@mastra/express`, custom routes must be passed to
+    // the `MastraServer` constructor directly.
+    this.mastra = new Mastra({ agents: this.built.agents });
+    this.mastraApp = express();
+    attachRoutePatchMiddleware(this.mastraApp);
+    this.mastraServer = new MastraServer(this.config, {
+      app: this.mastraApp,
+      mastra: this.mastra,
+      prefix: "",
+      customApiRoutes: [
+        chatRoute({ path: "/route/chat", agent: this.built.defaultAgentId }),
+        chatRoute({ path: "/route/chat/:agentId" }),
+      ],
+    });
+    await this.mastraServer.init();
+  }
+}
+
+export const mastra = toPlugin(MastraPlugin);

package/src/server.ts

@@ -0,0 +1,130 @@
+/**
+ * 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.
+ */
+
+import { getExecutionContext } from "@databricks/appkit";
+import { httpUtils, logUtils, stringUtils } from "@dbx-tools/appkit-shared";
+import {
+  MASTRA_RESOURCE_ID_KEY,
+  MASTRA_THREAD_ID_KEY,
+  type RequestContext,
+} from "@mastra/core/request-context";
+import { MastraServer as MastraServerExpress } from "@mastra/express";
+import type express from "express";
+import { randomUUID } from "node:crypto";
+
+import { MASTRA_USER_KEY, type MastraPluginConfig, type User } from "./config.js";
+import {
+  extractModelOverride,
+  MASTRA_MODEL_OVERRIDE_KEY,
+  resolveServingConfig,
+} from "./serving.js";
+
+/**
+ * `@mastra/express` subclass that stamps `RequestContext` with the
+ * AppKit user, resource id, and a thread id backed by an HTTP-only
+ * session cookie (`appkit_<plugin-name>_session_id`).
+ */
+export class MastraServer extends MastraServerExpress {
+  private log: logUtils.Logger;
+
+  constructor(
+    private config: MastraPluginConfig,
+    ...args: ConstructorParameters<typeof MastraServerExpress>
+  ) {
+    super(...args);
+    this.log = logUtils.logger(config);
+  }
+
+  override registerAuthMiddleware(): void {
+    super.registerAuthMiddleware();
+    this.app.use((req, res, next) => {
+      const executionContext = getExecutionContext();
+      const user: User = {
+        id:
+          "userId" in executionContext
+            ? executionContext.userId
+            : executionContext.serviceUserId,
+        executionContext,
+      };
+      const requestContext = res.locals.requestContext! as RequestContext;
+      requestContext.set(MASTRA_USER_KEY, user);
+      if (!requestContext.get(MASTRA_RESOURCE_ID_KEY)) {
+        this.log.debug(`Setting resource id: ${user.id}`);
+        requestContext.set(MASTRA_RESOURCE_ID_KEY, user.id);
+      }
+      const cookies = httpUtils.parseCookies(req.headers.cookie);
+      const cookieName = stringUtils.toIdentifierWithOptions(
+        { delimiter: "_", distinct: true },
+        "appkit",
+        this.config.name!,
+        "sessionId",
+      );
+      let sessionId = cookies[cookieName];
+      if (!sessionId) {
+        sessionId = randomUUID();
+        res.cookie(cookieName, sessionId, {
+          httpOnly: true,
+          sameSite: "lax",
+          secure: req.secure,
+          path: "/",
+        });
+      }
+      res.locals.sessionId = sessionId;
+      if (!requestContext.get(MASTRA_THREAD_ID_KEY)) {
+        this.log.debug(`Setting thread id: ${sessionId}`);
+        requestContext.set(MASTRA_THREAD_ID_KEY, sessionId);
+      }
+      // Per-request model override: only honored when the plugin
+      // opts in (default). Sources, in priority order, are
+      // `X-Mastra-Model` header, `?model=` query, and `model` /
+      // `modelId` body field; see `serving.ts`.
+      const serving = resolveServingConfig(this.config);
+      if (serving.allowOverride) {
+        const override = extractModelOverride({
+          headers: req.headers as Record<string, string | string[] | undefined>,
+          query: req.query as Record<string, unknown>,
+          body: req.body,
+        });
+        if (override) {
+          this.log.debug(`Model override: ${override}`);
+          requestContext.set(MASTRA_MODEL_OVERRIDE_KEY, override);
+        }
+      }
+      next();
+    });
+  }
+}
+
+/**
+ * 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`).
+ *
+ * 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).
+ */
+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();
+    req.originalUrl = req.path;
+    next();
+  });
+}

package/src/serving.ts

@@ -0,0 +1,294 @@
+/**
+ * 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, 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";
+
+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 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
+ * 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 = 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 as Record<string, unknown>;
+    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: MastraPluginConfig,
+  defaultFallbacks: readonly string[] = [],
+): {
+  ttlMs: number;
+  threshold: number;
+  fuzzy: boolean;
+  allowOverride: boolean;
+  fallbacks: readonly string[];
+} {
+  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,
+  };
+}