@dbx-tools/appkit-mastra 0.1.40 → 0.1.42

package/index.ts

@@ -11,29 +11,18 @@
  * consumers.
  */
 export * from "@dbx-tools/appkit-mastra-shared";
+export * from "@dbx-tools/model";
 export * from "./src/agents.js";
 export * from "./src/chart.js";
 export * from "./src/config.js";
 export * from "./src/genie.js";
-export {
-  FALLBACK_MODEL_IDS,
-  MODEL_CATALOG,
-  ModelTier,
-  modelForTier,
-  modelsForTier,
-} from "./src/model.js";
 export * from "./src/plugin.js";
 export {
   MASTRA_MODEL_OVERRIDE_KEY,
   MODEL_OVERRIDE_BODY_FIELDS,
   MODEL_OVERRIDE_HEADER,
   MODEL_OVERRIDE_QUERY,
-  clearServingEndpointsCache,
   extractModelOverride,
-  listServingEndpoints,
-  resolveModelId,
-  type ResolveModelOptions,
-  type ResolvedModel,
-  type ServingEndpointSummary,
+  type ModelOverrideRequest,
 } from "./src/serving.js";
 export * from "./src/tools/email.js";

package/package.json

@@ -9,13 +9,14 @@
     }
   },
   "name": "@dbx-tools/appkit-mastra",
-  "version": "0.1.40",
+  "version": "0.1.42",
   "dependencies": {
     "@databricks/sdk-experimental": "^0.17",
-    "@dbx-tools/appkit-mastra-shared": "0.1.40",
-    "@dbx-tools/genie": "0.1.40",
-    "@dbx-tools/genie-shared": "0.1.40",
-    "@dbx-tools/shared": "0.1.40",
+    "@dbx-tools/appkit-mastra-shared": "0.1.42",
+    "@dbx-tools/genie": "0.1.42",
+    "@dbx-tools/genie-shared": "0.1.42",
+    "@dbx-tools/model": "0.1.42",
+    "@dbx-tools/shared": "0.1.42",
     "@mastra/ai-sdk": "^1",
     "@mastra/core": "^1",
     "@mastra/express": "^1",
@@ -24,7 +25,7 @@
     "@mastra/observability": "^1",
     "@mastra/otel-bridge": "^1",
     "@mastra/pg": "^1",
-    "fuse.js": "^7.0.0",
+    "pg": "^8.20.0",
     "zod": "^4.3.6"
   },
   "devDependencies": {

package/src/agents.ts

@@ -10,7 +10,7 @@
  *
  * When no agents are registered the plugin falls back to a single
  * built-in analyst so the bare `mastra()` call still mounts a working
- * `chatRoute` agent for demos.
+ * streamable agent for demos.
  */
 
 import { appkitUtils, logUtils, stringUtils } from "@dbx-tools/shared";
@@ -226,8 +226,8 @@ export type MastraToolsFn = (
 /**
  * A code-defined Mastra agent. Mirrors the shape AppKit's `agents`
  * plugin uses for `AgentDefinition`. The registry key under
- * `config.agents` is what `chatRoute` matches on; `name` is purely
- * informational (defaults to the key).
+ * `config.agents` is the `agentId` the client streams against; `name`
+ * is purely informational (defaults to the key).
  */
 export interface MastraAgentDefinition {
   /** Display name used as `Agent.name`. Defaults to the registry key. */

package/src/chart.ts

@@ -39,7 +39,7 @@ import { createTool } from "@mastra/core/tools";
 import { z } from "zod";
 
 import type { MastraPluginConfig } from "./config.js";
-import { buildModel, modelForTier, ModelTier } from "./model.js";
+import { buildModel, ModelTier } from "./model.js";
 
 const log = logUtils.logger("mastra/chart");
 
@@ -330,9 +330,7 @@ function getPlannerAgent(config: MastraPluginConfig): Agent {
       description: "Picks chart type and axis encodings for a dataset.",
       instructions: CHART_PLANNER_INSTRUCTIONS,
       model: ({ requestContext }) =>
-        buildModel(config, requestContext, {
-          modelId: modelForTier(ModelTier.Fast),
-        }),
+        buildModel(config, requestContext, { tier: ModelTier.Fast }),
     });
     plannerAgents.set(config, agent);
   }
@@ -361,7 +359,7 @@ async function runChartPlanner(
     ...(requestContext ? { requestContext } : {}),
     ...(abortSignal ? { abortSignal } : {}),
   });
-  const plan = result.object as ChartPlan;
+  const plan = chartPlanSchema.parse(result.object);
   const option = planToEchartsOption(plan, title);
   return { chartType: plan.chartType, option };
 }

package/src/config.ts

@@ -150,7 +150,8 @@ export interface MastraPluginConfig extends BasePluginConfig {
    */
   tools?: MastraTools;
   /**
-   * Agent id used by `chatRoute` when the client doesn't specify one.
+   * Agent id used when the client doesn't specify one (the bare,
+   * un-suffixed history / suggestions routes resolve to it).
    * Defaults to the first key in `agents` (or `"default"` when
    * `agents` is omitted). Must match an id in `agents` when both are
    * set; a mismatch throws at setup with the available candidates.
@@ -198,17 +199,19 @@ export interface MastraPluginConfig extends BasePluginConfig {
    */
   modelOverride?: boolean;
   /**
-   * Priority-ordered list of endpoint names tried when no agent /
-   * plugin / env / request-override model id is set. The resolver
-   * picks the first id that is actually present in the workspace's
-   * `/serving-endpoints` listing - this is what lets a workspace
-   * without Claude Opus still get a sensible default automatically.
+   * Priority-ordered list of endpoint names tried *first* when no
+   * agent / plugin / env / request-override model id is set, ahead of
+   * the dynamic score-classified catalogue. The resolver picks the
+   * first id that is actually present in the workspace's
+   * `/serving-endpoints` listing.
    *
-   * Defaults to the built-in list in `model.ts` (`FALLBACK_MODEL_IDS`):
-   * Claude (Opus -> Sonnet -> Haiku), then OpenAI GPT-5 family, then
-   * open weights (Llama 4, Llama 3.3, GPT-OSS, Qwen, Llama 3.1).
-   * Override here to pin a regulated workspace to an approved subset
-   * or to add custom endpoints in front of the public catalogue.
+   * When unset, resolution is driven by the live Foundation Model API
+   * `quality` / `speed` / `cost` scores: endpoints are classified into
+   * tiers (`classifyEndpoints`) and walked best-first (Thinking ->
+   * Balanced -> Fast), with the small built-in `FALLBACK_MODEL_IDS`
+   * list as the floor when the catalogue can't be read. Set this to
+   * pin a regulated workspace to an approved subset, or to put custom
+   * endpoints in front of the auto-classified catalogue.
    */
   defaultModelFallbacks?: readonly string[];
   /**

package/src/genie.ts

@@ -38,11 +38,11 @@ import { CacheManager, genie } from "@databricks/appkit";
 import { ApiError, HttpError, WorkspaceClient } from "@databricks/sdk-experimental";
 import {
   ChartSchema,
-  type MinimalWriter,
+  type MastraWriter,
   type StartedEvent,
 } from "@dbx-tools/appkit-mastra-shared";
 import { genieEventChat, genieSampleQuestions, getGenieSpace } from "@dbx-tools/genie";
-import { type GenieMessage } from "@dbx-tools/genie-shared";
+import { GenieMessageSchema, type GenieMessage } from "@dbx-tools/genie-shared";
 import { appkitUtils, commonUtils, logUtils, stringUtils } from "@dbx-tools/shared";
 import type { RequestContext } from "@mastra/core/request-context";
 import { MASTRA_THREAD_ID_KEY } from "@mastra/core/request-context";
@@ -93,7 +93,7 @@ export type GenieSpacesConfig = Record<string, GenieSpaceConfig | string>;
 type ToolExecuteCtx =
   | {
       requestContext?: RequestContext;
-      writer?: MinimalWriter;
+      writer?: MastraWriter;
       abortSignal?: AbortSignal;
     }
   | undefined;
@@ -421,7 +421,7 @@ function buildAskGenieTool(opts: { spaceId: string; alias: string; hint?: string
       question: z.string().min(1, "question is required"),
     }),
     outputSchema: z.object({
-      message: z.custom<GenieMessage>(),
+      message: GenieMessageSchema,
     }),
     execute: async ({ question }, ctxRaw) => {
       const ctx = ctxRaw as ToolExecuteCtx;
@@ -1033,24 +1033,6 @@ export function buildGenieToolkitProvider(opts: {
   };
 }
 
-/**
- * Returns `true` when at least one Genie space is reachable
- * through {@link resolveGenieSpaces} - either via
- * {@link MastraPluginConfig.genieSpaces}, the AppKit `genie()`
- * plugin instance, or the `DATABRICKS_GENIE_SPACE_ID` env var.
- *
- * Cheap to call from `resolveProvider` to short-circuit `genie`
- * lookups when nothing is wired, so the `plugins.genie` lookup
- * still resolves to `undefined` (matching AppKit's
- * absent-plugin semantics) when neither source is configured.
- */
-export function hasAnyGenieSpaces(
-  config: MastraPluginConfig,
-  context: appkitUtils.PluginContextLike | undefined,
-): boolean {
-  return Object.keys(resolveGenieSpaces(config, context)).length > 0;
-}
-
 /* --------------------------- starter suggestions --------------------------- */
 
 /**

package/src/history.ts

@@ -9,8 +9,8 @@
  *
  * The route is registered through {@link historyRoute} as a Mastra
  * `registerApiRoute` so it sits in the same dispatcher pipeline as
- * `chatRoute`. That means the `MastraServer` auth middleware (in
- * `./server.ts`) has already populated `RequestContext` with
+ * the standard agent routes. That means the `MastraServer` auth
+ * middleware (in `./server.ts`) has already populated `RequestContext` with
  * `MASTRA_THREAD_ID_KEY` and `MASTRA_RESOURCE_ID_KEY` by the time
  * the handler runs - no cookie or user lookups happen here, and the
  * session-cookie logic stays the single source of truth in `server.ts`.
@@ -197,10 +197,10 @@ export type HistoryRouteOptions =
  *     anchors the thread id is left alone so the user keeps the
  *     same thread - only the contents go away.
  *
- * Modeled after `chatRoute` from `@mastra/ai-sdk`: pass `agent` for a
- * fixed-agent mount, or include `:agentId` in the path for dynamic
- * routing. Pairs cleanly with the AppKit Mastra plugin's chat route
- * layout (`/route/chat` + `/route/chat/:agentId`).
+ * Follows the `@mastra/ai-sdk` `chatRoute` convention for agent
+ * binding: pass `agent` for a fixed-agent mount, or include
+ * `:agentId` in the path for dynamic routing. The plugin registers
+ * both `/route/history` (default agent) and `/route/history/:agentId`.
  *
  * The handler reads `threadId` and `resourceId` from `RequestContext`
  * (populated upstream by `MastraServer.registerAuthMiddleware`), so

package/src/memory.ts

@@ -27,23 +27,46 @@
  * is registered); per-agent settings cascade on top of that.
  */
 
-import { lakebase } from "@databricks/appkit";
-import { appkitUtils, logUtils } from "@dbx-tools/shared";
+import { getUsernameWithApiLookup } from "@databricks/appkit";
+import { logUtils } from "@dbx-tools/shared";
 import { fastembed } from "@mastra/fastembed";
 import { Memory } from "@mastra/memory";
 import { PgVector, PostgresStore } from "@mastra/pg";
 import { randomUUID } from "node:crypto";
-import type { Pool } from "pg";
+import { Pool, type PoolConfig } from "pg";
 
 import type { MastraAgentDefinition, MastraMemoryConfigOverride } from "./agents.js";
 import type { MastraPluginConfig } from "./config.js";
 
 const log = logUtils.logger("mastra/memory");
 
-/** Pool handle returned by the AppKit `lakebase` plugin `exports().pool`. */
-export type LakebasePool = ReturnType<
-  InstanceType<ReturnType<typeof lakebase>["plugin"]>["exports"]
->["pool"];
+/**
+ * Build a dedicated **service-principal** Lakebase pool for Mastra
+ * memory from the lakebase plugin's resolved SP pg config.
+ *
+ * The plugin's `exports().pool` is a `RoutingPool` that switches to
+ * the per-user (OBO) pool whenever a query runs inside an `asUser`
+ * scope - exactly the context the mastra plugin establishes around
+ * every chat turn. Memory (threads / messages + semantic recall) must
+ * instead always act as the app service principal: it owns the
+ * auto-created `mastra_*` schemas (a per-user role usually can't
+ * `CREATE SCHEMA`) and is shared across users, so it cannot inherit a
+ * request's OBO identity.
+ *
+ * `pgConfig` must be the plugin's `exports().getPgConfig()` evaluated
+ * **outside** any `asUser` scope (i.e. during setup), so it carries
+ * the SP connection target, OAuth token-refresh `password` callback,
+ * and any `lakebase({ pool })` tuning overrides - all of which this
+ * pool inherits. See the call site in `plugin.ts`.
+ */
+export async function createServicePrincipalPool(pgConfig: PoolConfig): Promise<Pool> {
+  // `getPgConfig()` resolves the SP username synchronously from
+  // `PGUSER` / `DATABRICKS_CLIENT_ID`; fall back to the async API
+  // lookup (e.g. local dev authenticating via PAT) so the pool always
+  // has an identity to connect with.
+  const user = pgConfig.user ?? (await getUsernameWithApiLookup());
+  return new Pool({ ...pgConfig, user });
+}
 
 /** Effective per-knob setting after the plugin/agent cascade. */
 type StorageSetting = MastraAgentDefinition["storage"];
@@ -51,9 +74,9 @@ type MemorySetting = MastraAgentDefinition["memory"];
 
 /**
  * True when any plugin-level or per-agent setting could need the
- * Lakebase pool. Used by `plugin.ts` to gate pool acquisition; the
- * builder also acquires lazily so missed cases still fail with a
- * clear lakebase-not-registered error.
+ * Lakebase pool. Used by `plugin.ts` to gate creation of the
+ * service-principal pool and the {@link MemoryBuilder} that consumes
+ * it; when false neither is built.
  */
 export function needsLakebase(config: MastraPluginConfig): boolean {
   if (settingNeedsSharedPool(config.storage)) return true;
@@ -65,42 +88,29 @@ export function needsLakebase(config: MastraPluginConfig): boolean {
 }
 
 /**
- * Look up the `lakebase` plugin and return its managed `pg.Pool`.
- * Throws when the sibling plugin is not registered; enabling
- * `storage` / `memory` without lakebase is a wiring bug, not a runtime
- * condition we can recover from.
- */
-export function resolveLakebasePool(
-  context: appkitUtils.PluginContextLike | undefined,
-  caller: MastraPluginConfig,
-): LakebasePool {
-  return appkitUtils.require(context, lakebase, caller).exports().pool;
-}
-
-/**
- * Construct a per-agent {@link Memory} factory. Caches the shared
- * `PgVector` singleton (built on first need) and the lazily-resolved
- * Lakebase pool so each agent build is O(1) after the first.
+ * Construct a per-agent {@link Memory} factory bound to the supplied
+ * service-principal pool (see {@link createServicePrincipalPool}).
+ * Caches the shared `PgVector` singleton (built on first need) so each
+ * agent build is O(1) after the first.
  */
 export function createMemoryBuilder(
   config: MastraPluginConfig,
-  context: appkitUtils.PluginContextLike | undefined,
+  servicePrincipalPool: Pool,
 ): MemoryBuilder {
-  return new MemoryBuilder(config, context);
+  return new MemoryBuilder(config, servicePrincipalPool);
 }
 
 /**
- * Builds one `Memory` per agent. Per-instance state keeps the shared
- * `PgVector` and the resolved Lakebase pool alive across calls so
- * registering N agents stays cheap.
+ * Builds one `Memory` per agent against a shared service-principal
+ * Lakebase pool. Per-instance state keeps the shared `PgVector` alive
+ * across calls so registering N agents stays cheap.
  */
 export class MemoryBuilder {
   private sharedVector: PgVector | undefined;
-  private pool: LakebasePool | undefined;
 
   constructor(
     private readonly config: MastraPluginConfig,
-    private readonly context: appkitUtils.PluginContextLike | undefined,
+    private readonly servicePrincipalPool: Pool,
   ) {}
 
   /**
@@ -133,7 +143,7 @@ export class MemoryBuilder {
     return new PostgresStore({
       id: "mastra-store__instance",
       schemaName: "mastra_instance",
-      pool: this.requirePool() as Pool,
+      pool: this.servicePrincipalPool,
     });
   }
 
@@ -179,7 +189,7 @@ export class MemoryBuilder {
       return new PostgresStore({
         id: `mastra-store__${agentId}`,
         schemaName: `mastra_${agentId}`,
-        pool: this.requirePool() as Pool,
+        pool: this.servicePrincipalPool,
       });
     }
     // Cast: `withId` guarantees `id` is set, but the distributive
@@ -209,17 +219,10 @@ export class MemoryBuilder {
 
   private getSharedVector(): PgVector {
     if (!this.sharedVector) {
-      this.sharedVector = buildSharedPgVector(this.requirePool());
+      this.sharedVector = buildSharedPgVector(this.servicePrincipalPool);
     }
     return this.sharedVector;
   }
-
-  private requirePool(): LakebasePool {
-    if (!this.pool) {
-      this.pool = resolveLakebasePool(this.context, this.config);
-    }
-    return this.pool;
-  }
 }
 
 /**
@@ -242,7 +245,7 @@ export class MemoryBuilder {
  * the lakebase pool from then on. The placeholder pool is `.end()`'d
  * so its socket book-keeping is released.
  */
-function buildSharedPgVector(pool: LakebasePool): PgVector {
+function buildSharedPgVector(pool: Pool): PgVector {
   const vector = new PgVector({
     id: `pg${randomUUID()}`,
     host: "-1",
@@ -252,7 +255,7 @@ function buildSharedPgVector(pool: LakebasePool): PgVector {
     password: "_",
   });
   const placeholder = vector.pool;
-  vector.pool = pool as Pool;
+  vector.pool = pool;
   void placeholder.end().catch(() => undefined);
   return vector;
 }