@dbx-tools/appkit-mastra 0.1.0

package/dist/src/model.js

@@ -0,0 +1,423 @@
+/**
+ * Databricks Model Serving resolver for Mastra agents.
+ *
+ * Each agent step calls {@link buildModel} with the active
+ * `RequestContext`. The user stamped by `MastraServer` carries an
+ * AppKit `WorkspaceClient`; we ask it for the workspace host and a
+ * fresh bearer header, then point Mastra's OpenAI-compatible provider
+ * at `/serving-endpoints` on that host.
+ *
+ * Model id resolution walks three sources before falling back to the
+ * hard-coded default, **in this priority order**:
+ *
+ *   1. Per-request override stashed by the auth middleware under
+ *      {@link MASTRA_MODEL_OVERRIDE_KEY} (header / query / body).
+ *   2. The static `modelId` passed in by the agent / plugin (string
+ *      sugar on `def.model` or `config.defaultModel`).
+ *   3. `DATABRICKS_SERVING_ENDPOINT_NAME` env var.
+ *   4. {@link FALLBACK_MODEL_ID}.
+ *
+ * Whatever wins is then fuzzy-matched against the live
+ * `/serving-endpoints` list ({@link listServingEndpoints}) so loose
+ * names like `"claude sonnet"` resolve to the real endpoint name.
+ * Fuzzy matching is best-effort: when the workspace client throws
+ * (network blip, expired token at cache-fill time) we fall back to
+ * the input verbatim and let Databricks return the canonical error.
+ */
+import { commonUtils, httpUtils, logUtils, stringUtils, } from "@dbx-tools/appkit-shared";
+import { MASTRA_USER_KEY } from "./config.js";
+import { listServingEndpoints, MASTRA_MODEL_OVERRIDE_KEY, resolveModelId, resolveServingConfig, } from "./serving.js";
+/**
+ * Capability tiers for Databricks Foundation Model API endpoints.
+ *
+ * - {@link ModelTier.Thinking}: deepest reasoning / "thinking" models
+ *   (Claude Opus, GPT-5.5 Pro, Gemini Pro, Llama 4 Maverick, etc).
+ *   Highest cost and latency; reserve for hard multi-step reasoning.
+ * - {@link ModelTier.Balanced}: cost/latency sweet spot for general
+ *   agent work (Claude Sonnet, GPT-5.x, Gemini Flash, Llama 3.3 70B).
+ *   The right default for most agents.
+ * - {@link ModelTier.Fast}: cheap and quick; classification, routing,
+ *   tool-arg extraction, simple summarisation (Claude Haiku, GPT-5
+ *   mini/nano, Gemini Flash Lite, GPT-OSS 20B, Llama 3.1 8B).
+ *
+ * String enum so the value is the slug we use in cache keys, logs,
+ * and as the value users see in serialized configs.
+ */
+export var ModelTier;
+(function (ModelTier) {
+    ModelTier["Thinking"] = "thinking";
+    ModelTier["Balanced"] = "balanced";
+    ModelTier["Fast"] = "fast";
+})(ModelTier || (ModelTier = {}));
+/**
+ * Catalogue of Databricks-hosted Foundation Model API endpoints,
+ * grouped by capability {@link ModelTier} and then by provider. Each
+ * inner array is priority-ordered (most powerful first within the
+ * same provider+tier).
+ *
+ * Provider buckets:
+ *
+ * - `claude`: Anthropic Claude family (closed; flagship reasoning).
+ * - `gpt`: OpenAI GPT-5 family (closed; "ChatGPT" on Databricks FMAPI).
+ * - `gemini`: Google Gemini family (closed; multimodal + web-search).
+ * - `openSource`: open-weights models (widest regional / SKU availability).
+ *
+ * The list is curated by hand; refresh from the Databricks "supported
+ * foundation models" doc when new endpoints land.
+ */
+export const MODEL_CATALOG = {
+    [ModelTier.Thinking]: {
+        claude: [
+            "databricks-claude-opus-4-8",
+            "databricks-claude-opus-4-7",
+            "databricks-claude-opus-4-6",
+            "databricks-claude-opus-4-5",
+            "databricks-claude-opus-4-1",
+        ],
+        gpt: ["databricks-gpt-5-5-pro"],
+        gemini: [
+            "databricks-gemini-3-1-pro",
+            "databricks-gemini-3-pro",
+            "databricks-gemini-2-5-pro",
+        ],
+        openSource: [
+            "databricks-llama-4-maverick",
+            "databricks-gpt-oss-120b",
+            "databricks-meta-llama-3-1-405b-instruct",
+        ],
+    },
+    [ModelTier.Balanced]: {
+        claude: [
+            "databricks-claude-sonnet-4-6",
+            "databricks-claude-sonnet-4-5",
+            "databricks-claude-sonnet-4",
+        ],
+        gpt: [
+            "databricks-gpt-5-5",
+            "databricks-gpt-5-4",
+            "databricks-gpt-5-2",
+            "databricks-gpt-5-1",
+            "databricks-gpt-5",
+        ],
+        gemini: [
+            "databricks-gemini-3-5-flash",
+            "databricks-gemini-3-flash",
+            "databricks-gemini-2-5-flash",
+        ],
+        openSource: [
+            "databricks-meta-llama-3-3-70b-instruct",
+            "databricks-qwen3-next-80b-a3b-instruct",
+            "databricks-qwen35-122b-a10b",
+        ],
+    },
+    [ModelTier.Fast]: {
+        claude: ["databricks-claude-haiku-4-5"],
+        gpt: [
+            "databricks-gpt-5-4-mini",
+            "databricks-gpt-5-4-nano",
+            "databricks-gpt-5-mini",
+            "databricks-gpt-5-nano",
+        ],
+        gemini: ["databricks-gemini-3-1-flash-lite"],
+        openSource: [
+            "databricks-gpt-oss-20b",
+            "databricks-gemma-3-12b",
+            "databricks-meta-llama-3-1-8b-instruct",
+        ],
+    },
+};
+/**
+ * Round-robin zip: take one from each input list in order, skipping
+ * lists that have already been exhausted. Used to interleave provider
+ * buckets within a tier so the resolver alternates between vendors
+ * instead of draining one before trying the next.
+ *
+ * Example: `interleave(["a1","a2","a3"], ["b1","b2"])` ->
+ * `["a1","b1","a2","b2","a3"]`.
+ */
+function interleave(...lists) {
+    const out = [];
+    const max = Math.max(0, ...lists.map((l) => l.length));
+    for (let i = 0; i < max; i++) {
+        for (const list of lists) {
+            if (i < list.length)
+                out.push(list[i]);
+        }
+    }
+    return out;
+}
+/**
+ * Priority-ordered model ids for a single capability {@link ModelTier},
+ * interleaved across providers so a workspace missing the top Claude
+ * still lands on a flagship GPT / Gemini on the next probe.
+ *
+ * Provider order within the interleave: Claude, GPT, Gemini, then the
+ * open-weights tail appended verbatim as the universal floor (widest
+ * regional availability).
+ *
+ * @example
+ * ```ts
+ * mastra({
+ *   defaultModelFallbacks: modelsForTier(ModelTier.Fast),
+ * });
+ * ```
+ */
+export function modelsForTier(tier) {
+    const bucket = MODEL_CATALOG[tier];
+    return [
+        ...interleave(bucket.claude, bucket.gpt, bucket.gemini),
+        ...bucket.openSource,
+    ];
+}
+/**
+ * Top model id at the given {@link ModelTier}. Sync; the agent-step
+ * resolver fuzzy-matches it against the workspace catalogue at call
+ * time, so this works even when the literal top pick isn't deployed.
+ *
+ * Use when wiring a tier-appropriate model into an agent definition:
+ *
+ * @example
+ * ```ts
+ * const classifier = createAgent({
+ *   instructions: "Classify this email",
+ *   model: modelForTier(ModelTier.Fast),  // cheap, quick
+ * });
+ *
+ * const planner = createAgent({
+ *   instructions: "Plan a multi-step migration",
+ *   model: modelForTier(ModelTier.Thinking),  // deep reasoning
+ * });
+ * ```
+ */
+export function modelForTier(tier) {
+    return modelsForTier(tier)[0];
+}
+/**
+ * Last-resort model ids used when neither `config.defaultModel`,
+ * per-agent `model`, nor `DATABRICKS_SERVING_ENDPOINT_NAME` is set.
+ *
+ * Walked in order at resolve time: the first id whose endpoint is
+ * actually present in the workspace's `/serving-endpoints` listing
+ * wins. Workspaces vary - not every region / SKU has every model,
+ * and the list of Foundation Model APIs evolves quickly - so the
+ * resolver degrades all the way from "best thinking model" down to
+ * "smallest commodity Llama" before giving up.
+ *
+ * Built by chaining the per-tier interleaves (Thinking -> Balanced
+ * -> Fast); within each tier the providers are round-robin-zipped
+ * (Claude, GPT, Gemini, then open-weights tail). Override the entire
+ * list via `MastraPluginConfig.defaultModelFallbacks` (e.g. to pin a
+ * regulated workspace to a specific approved subset, or to bias the
+ * priority toward a particular tier).
+ */
+export const FALLBACK_MODEL_IDS = [
+    ...modelsForTier(ModelTier.Thinking),
+    ...modelsForTier(ModelTier.Balanced),
+    ...modelsForTier(ModelTier.Fast),
+];
+/**
+ * Resolve a `MastraModelConfig` for the current agent step. Runs
+ * while `agent.stream` is inside the `asUser(req)` scope so tokens
+ * are user-scoped; outside an active user context the workspace
+ * client falls back to the service principal.
+ */
+export async function buildModel(config, requestContext, overrides = {}) {
+    void setupFetchInterceptor();
+    const user = requestContext.get(MASTRA_USER_KEY);
+    const clientConfig = user.executionContext.client.config;
+    const host = (await clientConfig.getHost()).toString();
+    const headers = new Headers();
+    await clientConfig.authenticate(headers);
+    // The OpenAI Node SDK appends paths like `/chat/completions` to whatever
+    // URL we hand it. Drop the trailing slash so the resulting URL stays
+    // well-formed (`/serving-endpoints/chat/completions`).
+    const url = new URL("/serving-endpoints", host).toString().replace(/\/$/, "");
+    const modelId = await pickModelId(config, requestContext, overrides, user, host);
+    return {
+        providerId: config.providerId ?? "databricks",
+        modelId,
+        url,
+        headers: Object.fromEntries(headers.entries()),
+    };
+}
+/**
+ * Walk the resolution ladder and pick a modelId.
+ *
+ * 1. **Explicit ask** (per-request override, agent `model` string,
+ *    `config.defaultModel` string, or `DATABRICKS_SERVING_ENDPOINT_NAME`):
+ *    when fuzzy matching is on, snap the input to the closest live
+ *    endpoint so loose names like `"claude sonnet"` resolve. When it's
+ *    off (or no endpoint matches within threshold), the input is used
+ *    verbatim and Databricks surfaces the canonical 404.
+ *
+ * 2. **No explicit ask**: walk
+ *    {@link MastraPluginConfig.defaultModelFallbacks} (or
+ *    {@link FALLBACK_MODEL_IDS} when unset) and return the first id
+ *    whose endpoint is actually present in the workspace listing. A
+ *    workspace without Claude Opus still gets a sensible default by
+ *    skipping ahead to whichever Sonnet / GPT-5 / Llama variant is
+ *    wired up.
+ *
+ * Catalogue fetches fail loud: network / auth errors propagate to the
+ * caller so they see the real SDK message instead of a silent fallback
+ * to the top of the priority list.
+ */
+async function pickModelId(config, requestContext, overrides, user, host) {
+    const serving = resolveServingConfig(config, FALLBACK_MODEL_IDS);
+    const override = serving.allowOverride
+        ? requestContext.get(MASTRA_MODEL_OVERRIDE_KEY)
+        : undefined;
+    const explicit = override ?? overrides.modelId ?? process.env.DATABRICKS_SERVING_ENDPOINT_NAME;
+    // Cheap exit: when the caller named a specific model and fuzzy
+    // matching is off, there's no reason to touch the catalogue at all.
+    if (explicit !== undefined && !serving.fuzzy)
+        return explicit;
+    const endpoints = await listServingEndpoints(user.executionContext.client, host, {
+        ttlMs: serving.ttlMs,
+    });
+    const modelId = explicit !== undefined
+        ? resolveModelId(explicit, endpoints, { threshold: serving.threshold }).modelId
+        : pickFirstAvailable(serving.fallbacks, endpoints);
+    //logUtils.logger(config).debug(`model selected: ${modelId}`);
+    return modelId;
+}
+/**
+ * Find the first id in `fallbacks` whose endpoint is present in
+ * `endpoints`. Returns the top fallback when the workspace has none
+ * of them so callers always get a string; an offline workspace will
+ * then receive a clean 404 from Databricks instead of a malformed
+ * config.
+ */
+function pickFirstAvailable(fallbacks, endpoints) {
+    const present = new Set(endpoints.map((e) => e.name));
+    for (const candidate of fallbacks) {
+        if (present.has(candidate))
+            return candidate;
+    }
+    return fallbacks[0] ?? FALLBACK_MODEL_IDS[0];
+}
+/** Path prefix that identifies a Databricks Model Serving REST call. */
+const SERVING_ENDPOINTS_PATH_PREFIX = "/serving-endpoints/";
+/**
+ * Install a single shared `globalThis.fetch` wrapper for every POST to
+ * `/serving-endpoints/...`. The wrapper does two things:
+ *
+ *   1. Rewrites the outgoing `messages` array to repair Mastra/AI SDK
+ *      stream-replay quirks that Databricks-hosted Claude rejects (see
+ *      {@link sanitizeServingMessages}).
+ *   2. When `MASTRA_DEBUG_LLM=1`, dumps the (post-sanitize) JSON body
+ *      to stderr so 4xx debugging doesn't have to fight AI SDK's
+ *      `[Array]` formatter.
+ *
+ * Safe to call from any hot path: {@link commonUtils.memoize} ensures
+ * the wrapper is installed at most once per process, so subsequent
+ * calls collapse to a single cached promise even when
+ * {@link buildModel} fires on every agent step.
+ */
+const setupFetchInterceptor = commonUtils.memoize(() => {
+    const debug = Boolean(process.env.MASTRA_DEBUG_LLM);
+    const original = globalThis.fetch.bind(globalThis);
+    globalThis.fetch = (async (input, init) => {
+        const url = httpUtils.toURL(input);
+        if (!url ||
+            !url.pathname.startsWith(SERVING_ENDPOINTS_PATH_PREFIX) ||
+            typeof init?.body !== "string") {
+            return original(input, init);
+        }
+        const rewritten = rewriteServingBody(init.body);
+        if (rewritten !== init.body) {
+            init = { ...init, body: rewritten };
+        }
+        if (debug) {
+            try {
+                console.error("[mastra:llm-debug] -> POST", url.toString());
+                console.error(JSON.stringify(JSON.parse(rewritten), null, 2));
+            }
+            catch {
+                console.error("[mastra:llm-debug] -> POST", url.toString(), "(non-JSON body)");
+            }
+        }
+        return original(input, init);
+    });
+});
+/**
+ * Parse, sanitize, and re-serialize a `/serving-endpoints/...` POST
+ * body. Returns the original string verbatim when the body is not
+ * JSON, has no `messages`, or no rewrite was needed; this lets the
+ * caller skip the allocation of a new `init` object in the common
+ * pass-through case.
+ */
+function rewriteServingBody(body) {
+    let parsed;
+    try {
+        parsed = JSON.parse(body);
+    }
+    catch {
+        return body;
+    }
+    if (!Array.isArray(parsed.messages))
+        return body;
+    const changed = sanitizeServingMessages(parsed.messages);
+    return changed ? JSON.stringify(parsed) : body;
+}
+/**
+ * Repair a Mastra/AI SDK message replay that Databricks-hosted Claude
+ * rejects with `"This model does not support assistant message
+ * prefill. The conversation must end with a user message."`.
+ *
+ * The bug pattern: when an assistant turn streams text *and* a
+ * `tool_call`, the AI SDK persists them as two separate assistant
+ * entries (text-only and tool-call-only). On the next agent step the
+ * tool-call entry is replayed *before* the tool result and the
+ * text entry is replayed *after* it, so the conversation ends with a
+ * trailing assistant text message. Anthropic interprets that as a
+ * prefill request and rejects it on Databricks (the upstream Bedrock
+ * route disallows prefill).
+ *
+ * Fix: when the last message is an assistant text with no `tool_calls`
+ * and the chain immediately before it is `assistant(tool_calls=...)`
+ * followed only by `tool(...)` results, fold the trailing text back
+ * into the `content` of that opening assistant and drop the duplicate.
+ * The result is the canonical OpenAI shape
+ * `[..., user, assistant(text + tool_calls), tool(...)]` which both
+ * Databricks Claude and every other endpoint accept.
+ *
+ * Mutates `messages` in place; returns `true` when something changed
+ * so the caller knows whether to re-serialize.
+ */
+function sanitizeServingMessages(messages) {
+    if (messages.length < 2)
+        return false;
+    const last = messages[messages.length - 1];
+    if (!last ||
+        last.role !== "assistant" ||
+        (last.tool_calls && last.tool_calls.length > 0)) {
+        return false;
+    }
+    // Walk back through any contiguous tool-result messages to find the
+    // assistant turn that opened this tool sequence.
+    let i = messages.length - 2;
+    while (i >= 0 && messages[i]?.role === "tool")
+        i--;
+    if (i < 0)
+        return false;
+    const opener = messages[i];
+    if (!opener ||
+        opener.role !== "assistant" ||
+        !opener.tool_calls ||
+        opener.tool_calls.length === 0) {
+        return false;
+    }
+    // `trimToNull` collapses the `typeof string && trimmed` dance and
+    // drops blank fragments before the `\n\n` join below, so the merge
+    // never introduces stray leading / trailing whitespace.
+    const merged = [
+        stringUtils.trimToNull(opener.content),
+        stringUtils.trimToNull(last.content),
+    ]
+        .filter((s) => s !== null)
+        .join("\n\n");
+    opener.content = merged;
+    messages.pop();
+    return true;
+}

package/dist/src/plugin.d.ts

@@ -0,0 +1,120 @@
+/**
+ * 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 { Plugin, type IAppRouter, type ResourceRequirement } from "@databricks/appkit";
+import type { Agent } from "@mastra/core/agent";
+import { Mastra } from "@mastra/core/mastra";
+import type { MastraPluginConfig } from "./config.js";
+import { MastraServer } from "./server.js";
+import { type ServingEndpointSummary } from "./serving.js";
+/**
+ * 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 declare class MastraPlugin extends Plugin<MastraPluginConfig> {
+    static manifest: {
+        name: "mastra";
+        displayName: string;
+        description: string;
+        stability: "beta";
+        resources: {
+            required: never[];
+            optional: Omit<ResourceRequirement, "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: MastraPluginConfig): ResourceRequirement[];
+    private log;
+    private built;
+    private mastra;
+    private mastraApp;
+    private mastraServer;
+    setup(): Promise<void>;
+    /**
+     * 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;
+    exports(): {
+        /**
+         * 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[];
+        /**
+         * 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;
+        /**
+         * 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;
+        /** Underlying Mastra instance for advanced use (custom routes etc.). */
+        getMastra: () => Mastra<Record<string, Agent<any, import("@mastra/core/agent").ToolsInput, undefined, unknown>>, Record<string, import("@mastra/core/workflows").AnyWorkflow>, Record<string, import("@mastra/core/vector").MastraVector<any>>, Record<string, import("@mastra/core/tts").MastraTTS>, import("@mastra/core/logger").IMastraLogger, Record<string, import("@mastra/core/mcp").MCPServerBase<any>>, Record<string, import("@mastra/core/evals").MastraScorer<any, any, any, any>>, Record<string, import("@mastra/core/tools").ToolAction<any, any, any, any, any, any, unknown>>, Record<string, import("@mastra/core/processors").Processor<any, unknown>>, Record<string, import("@mastra/core/memory").MastraMemory>, Record<string, import("@mastra/core/channels").ChannelProvider>> | null;
+        /** Express subapp Mastra is mounted on; mostly for tests. */
+        getMastraServer: () => MastraServer | null;
+        /**
+         * 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[]>;
+        /**
+         * 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>;
+    };
+    clientConfig(): Record<string, unknown>;
+    injectRoutes(router: IAppRouter): void;
+    /**
+     * 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 listModels;
+    private buildAgentAndServer;
+}
+export declare const mastra: import("@databricks/appkit").ToPlugin<typeof MastraPlugin, MastraPluginConfig, "mastra">;