@dbx-tools/appkit-mastra 0.1.40 → 0.1.41

package/dist/src/memory.d.ts

@@ -26,45 +26,55 @@
  * (auto-defaulted to `true` in `plugin.ts` when the `lakebase` plugin
  * is registered); per-agent settings cascade on top of that.
  */
-import { lakebase } from "@databricks/appkit";
-import { appkitUtils } from "@dbx-tools/shared";
 import { Memory } from "@mastra/memory";
 import { PostgresStore } from "@mastra/pg";
+import { Pool, type PoolConfig } from "pg";
 import type { MastraAgentDefinition } from "./agents.js";
 import type { MastraPluginConfig } from "./config.js";
-/** Pool handle returned by the AppKit `lakebase` plugin `exports().pool`. */
-export type LakebasePool = ReturnType<InstanceType<ReturnType<typeof lakebase>["plugin"]>["exports"]>["pool"];
 /**
- * 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.
+ * 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 declare function needsLakebase(config: MastraPluginConfig): boolean;
+export declare function createServicePrincipalPool(pgConfig: PoolConfig): Promise<Pool>;
 /**
- * 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.
+ * True when any plugin-level or per-agent setting could need the
+ * 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 declare function resolveLakebasePool(context: appkitUtils.PluginContextLike | undefined, caller: MastraPluginConfig): LakebasePool;
+export declare function needsLakebase(config: MastraPluginConfig): boolean;
 /**
- * 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 declare function createMemoryBuilder(config: MastraPluginConfig, context: appkitUtils.PluginContextLike | undefined): MemoryBuilder;
+export declare function createMemoryBuilder(config: MastraPluginConfig, servicePrincipalPool: Pool): MemoryBuilder;
 /**
- * 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 declare class MemoryBuilder {
     private readonly config;
-    private readonly context;
+    private readonly servicePrincipalPool;
     private sharedVector;
-    private pool;
-    constructor(config: MastraPluginConfig, context: appkitUtils.PluginContextLike | undefined);
+    constructor(config: MastraPluginConfig, servicePrincipalPool: Pool);
     /**
      * Build a `Memory` for `agentId` after the plugin/agent cascade.
      * Returns `undefined` when the agent has neither storage nor a
@@ -96,5 +106,4 @@ export declare class MemoryBuilder {
      */
     private buildVector;
     private getSharedVector;
-    private requirePool;
 }

package/dist/src/memory.js

@@ -26,18 +26,46 @@
  * (auto-defaulted to `true` in `plugin.ts` when the `lakebase` plugin
  * 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 { Pool } from "pg";
 const log = logUtils.logger("mastra/memory");
+/**
+ * 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) {
+    // `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 });
+}
 /**
  * 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) {
     if (settingNeedsSharedPool(config.storage))
@@ -48,35 +76,26 @@ export function needsLakebase(config) {
     return defs.some((d) => settingNeedsSharedPool(d.storage) || settingNeedsSharedPool(d.memory));
 }
 /**
- * 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.
+ * 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 resolveLakebasePool(context, caller) {
-    return appkitUtils.require(context, lakebase, caller).exports().pool;
+export function createMemoryBuilder(config, servicePrincipalPool) {
+    return new MemoryBuilder(config, servicePrincipalPool);
 }
 /**
- * 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.
- */
-export function createMemoryBuilder(config, context) {
-    return new MemoryBuilder(config, context);
-}
-/**
- * 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 {
     config;
-    context;
+    servicePrincipalPool;
     sharedVector;
-    pool;
-    constructor(config, context) {
+    constructor(config, servicePrincipalPool) {
         this.config = config;
-        this.context = context;
+        this.servicePrincipalPool = servicePrincipalPool;
     }
     /**
      * Build a `Memory` for `agentId` after the plugin/agent cascade.
@@ -105,7 +124,7 @@ export class MemoryBuilder {
         return new PostgresStore({
             id: "mastra-store__instance",
             schemaName: "mastra_instance",
-            pool: this.requirePool(),
+            pool: this.servicePrincipalPool,
         });
     }
     forAgent(agentId, def) {
@@ -143,7 +162,7 @@ export class MemoryBuilder {
             return new PostgresStore({
                 id: `mastra-store__${agentId}`,
                 schemaName: `mastra_${agentId}`,
-                pool: this.requirePool(),
+                pool: this.servicePrincipalPool,
             });
         }
         // Cast: `withId` guarantees `id` is set, but the distributive
@@ -169,16 +188,10 @@ export class MemoryBuilder {
     }
     getSharedVector() {
         if (!this.sharedVector) {
-            this.sharedVector = buildSharedPgVector(this.requirePool());
+            this.sharedVector = buildSharedPgVector(this.servicePrincipalPool);
         }
         return this.sharedVector;
     }
-    requirePool() {
-        if (!this.pool) {
-            this.pool = resolveLakebasePool(this.context, this.config);
-        }
-        return this.pool;
-    }
 }
 /**
  * Build the shared `PgVector` that backs the default

package/dist/src/plugin.d.ts

@@ -59,6 +59,14 @@ export declare class MastraPlugin extends Plugin<MastraPluginConfig> {
     private mastra;
     private mastraApp;
     private mastraServer;
+    /**
+     * Dedicated service-principal Lakebase pool backing Mastra memory /
+     * storage. Built once in {@link buildAgentAndServer} (outside any
+     * `asUser` scope, so it never inherits a request's OBO identity) and
+     * drained in {@link abortActiveOperations}. `null` until setup runs
+     * or when Lakebase isn't needed.
+     */
+    private servicePrincipalPool;
     setup(): Promise<void>;
     /**
      * When the `lakebase` plugin is registered, auto-enable `storage`
@@ -68,6 +76,13 @@ export declare class MastraPlugin extends Plugin<MastraPluginConfig> {
      * already in the registry by the time this fires.
      */
     private applyLakebaseAutoDefaults;
+    /**
+     * Drain the memory service-principal pool on shutdown. AppKit calls
+     * this during teardown; the lakebase plugin closes its own SP / OBO
+     * pools the same way. Fire-and-forget so shutdown isn't blocked on a
+     * slow drain, and clear the handle so a re-`setup()` rebuilds it.
+     */
+    abortActiveOperations(): void;
     exports(): {
         /**
          * Ids of every registered agent in registration order. Matches

package/dist/src/plugin.js

@@ -27,7 +27,7 @@
  *   AI SDK transport URL is `/api/mastra/route/chat/<agentId>`.
  */
 import { genie, getExecutionContext, lakebase, Plugin, toPlugin, } from "@databricks/appkit";
-import { appkitUtils, logUtils } from "@dbx-tools/shared";
+import { appkitUtils, commonUtils, logUtils } from "@dbx-tools/shared";
 import { chatRoute } from "@mastra/ai-sdk";
 import { Mastra } from "@mastra/core/mastra";
 import express from "express";
@@ -35,7 +35,7 @@ import { buildAgents, FALLBACK_AGENT_ID } from "./agents.js";
 import { fetchChart } from "./chart.js";
 import { collectSpaceSuggestions, resolveGenieSpaces } from "./genie.js";
 import { historyRoute } from "./history.js";
-import { createMemoryBuilder, needsLakebase } from "./memory.js";
+import { createMemoryBuilder, createServicePrincipalPool, needsLakebase, } from "./memory.js";
 import { buildObservability } from "./observability.js";
 import { attachRoutePatchMiddleware, MastraServer } from "./server.js";
 import { clearServingEndpointsCache, listServingEndpoints, resolveServingConfig, } from "./serving.js";
@@ -93,6 +93,14 @@ export class MastraPlugin extends Plugin {
     mastra = null;
     mastraApp = null;
     mastraServer = null;
+    /**
+     * Dedicated service-principal Lakebase pool backing Mastra memory /
+     * storage. Built once in {@link buildAgentAndServer} (outside any
+     * `asUser` scope, so it never inherits a request's OBO identity) and
+     * drained in {@link abortActiveOperations}. `null` until setup runs
+     * or when Lakebase isn't needed.
+     */
+    servicePrincipalPool = null;
     async setup() {
         // Wait until sibling plugins (e.g. `lakebase`) finish `setup()` so
         // the lakebase pool is valid when storage/memory are enabled.
@@ -118,6 +126,25 @@ export class MastraPlugin extends Plugin {
         if (this.config.memory === undefined)
             this.config.memory = true;
     }
+    /**
+     * Drain the memory service-principal pool on shutdown. AppKit calls
+     * this during teardown; the lakebase plugin closes its own SP / OBO
+     * pools the same way. Fire-and-forget so shutdown isn't blocked on a
+     * slow drain, and clear the handle so a re-`setup()` rebuilds it.
+     */
+    abortActiveOperations() {
+        super.abortActiveOperations();
+        if (this.servicePrincipalPool) {
+            this.log.info("closing memory SP pool");
+            const pool = this.servicePrincipalPool;
+            this.servicePrincipalPool = null;
+            pool.end().catch((err) => {
+                this.log.error("error closing memory SP pool", {
+                    error: commonUtils.errorMessage(err),
+                });
+            });
+        }
+    }
     exports() {
         return {
             /**
@@ -402,12 +429,25 @@ export class MastraPlugin extends Plugin {
         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)
+        // Per-agent memory factory. When any storage / memory setting needs
+        // Postgres, stand up a dedicated service-principal pool first so
+        // memory acts as the app SP (owner of the `mastra_*` schemas),
+        // never the per-request OBO identity the chat turn runs under.
+        // `getPgConfig()` is read here, outside any `asUser` scope, so it
+        // returns the SP connection target + token refresh plus any
+        // `lakebase({ pool })` overrides; `require` turns a missing
+        // sibling into a clear wiring error. The builder caches the shared
+        // `PgVector` singleton so registering N agents stays cheap. See
+        // `./memory.js`.
+        if (needsLakebase(this.config)) {
+            const spPgConfig = appkitUtils
+                .require(this.context, lakebase, this.config)
+                .exports()
+                .getPgConfig();
+            this.servicePrincipalPool = await createServicePrincipalPool(spPgConfig);
+        }
+        const memoryBuilder = this.servicePrincipalPool
+            ? createMemoryBuilder(this.config, this.servicePrincipalPool)
             : undefined;
         this.log.debug("build:start", {
             lakebase: memoryBuilder !== undefined,