@dbx-tools/appkit-mastra 0.1.0

package/dist/src/plugin.js

@@ -0,0 +1,235 @@
+/**
+ * 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, } from "@databricks/appkit";
+import { logUtils, pluginUtils } from "@dbx-tools/appkit-shared";
+import { chatRoute } from "@mastra/ai-sdk";
+import { Mastra } from "@mastra/core/mastra";
+import express from "express";
+import { buildAgents, FALLBACK_AGENT_ID } from "./agents.js";
+import { createMemoryBuilder, needsLakebase } from "./memory.js";
+import { attachRoutePatchMiddleware, MastraServer } from "./server.js";
+import { clearServingEndpointsCache, listServingEndpoints, resolveServingConfig, } 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 {
+    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,
+            ],
+        },
+    };
+    /**
+     * 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) {
+        const resources = [];
+        const enabledManifests = [];
+        if (needsLakebase(config)) {
+            enabledManifests.push(LAKEBASE_MANIFEST);
+        }
+        for (const m of enabledManifests) {
+            for (const resource of m.resources.required) {
+                resources.push({ ...resource, required: true });
+            }
+        }
+        return resources;
+    }
+    log = logUtils.logger(this);
+    built = null;
+    mastra = null;
+    mastraApp = null;
+    mastraServer = null;
+    async setup() {
+        // 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.
+     */
+    applyLakebaseAutoDefaults() {
+        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;
+    }
+    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: () => 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) => 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: () => (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: () => 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) => clearServingEndpointsCache(host),
+        };
+    }
+    clientConfig() {
+        // 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 = {
+            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;
+    }
+    injectRoutes(router) {
+        // `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.
+     */
+    async listModels() {
+        const client = getExecutionContext().client;
+        const host = (await client.config.getHost()).toString();
+        const serving = resolveServingConfig(this.config);
+        return listServingEndpoints(client, host, { ttlMs: serving.ttlMs });
+    }
+    async buildAgentAndServer() {
+        // 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/dist/src/server.d.ts

@@ -0,0 +1,42 @@
+/**
+ * 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 { MastraServer as MastraServerExpress } from "@mastra/express";
+import type express from "express";
+import { type MastraPluginConfig } from "./config.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 declare class MastraServer extends MastraServerExpress {
+    private config;
+    private log;
+    constructor(config: MastraPluginConfig, ...args: ConstructorParameters<typeof MastraServerExpress>);
+    registerAuthMiddleware(): void;
+}
+/**
+ * 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 declare function attachRoutePatchMiddleware(app: express.Express): void;

package/dist/src/server.js

@@ -0,0 +1,109 @@
+/**
+ * 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, } from "@mastra/core/request-context";
+import { MastraServer as MastraServerExpress } from "@mastra/express";
+import { randomUUID } from "node:crypto";
+import { MASTRA_USER_KEY } 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 {
+    config;
+    log;
+    constructor(config, ...args) {
+        super(...args);
+        this.config = config;
+        this.log = logUtils.logger(config);
+    }
+    registerAuthMiddleware() {
+        super.registerAuthMiddleware();
+        this.app.use((req, res, next) => {
+            const executionContext = getExecutionContext();
+            const user = {
+                id: "userId" in executionContext
+                    ? executionContext.userId
+                    : executionContext.serviceUserId,
+                executionContext,
+            };
+            const requestContext = res.locals.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,
+                    query: req.query,
+                    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) {
+    app.use((req, _res, next) => {
+        const isChat = req.path === "/route" || req.path.startsWith("/route/");
+        if (!isChat)
+            return next();
+        req.originalUrl = req.path;
+        next();
+    });
+}

package/dist/src/serving.d.ts

@@ -0,0 +1,156 @@
+/**
+ * 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[];
+};