@dbx-tools/appkit-mastra 0.9.153 → 0.9.154

package/package.json

@@ -36,21 +36,21 @@
   },
   "dependencies": {
     "@databricks/appkit": "0.81.0",
-    "@dbx-tools/appkit": "^0.9.153",
-    "@dbx-tools/cache": "^0.9.153",
-    "@dbx-tools/core": "^0.9.153",
-    "@dbx-tools/databricks": "^0.9.153",
-    "@dbx-tools/fs": "^0.9.153",
-    "@dbx-tools/genie": "^0.9.153",
-    "@dbx-tools/model": "^0.9.153",
-    "@dbx-tools/model-protocol": "^0.9.153",
-    "@dbx-tools/path": "^0.9.153",
-    "@dbx-tools/postgres": "^0.9.153",
-    "@dbx-tools/shared-core": "^0.9.153",
-    "@dbx-tools/shared-fs": "^0.9.153",
-    "@dbx-tools/shared-genie": "^0.9.153",
-    "@dbx-tools/shared-mastra": "^0.9.153",
-    "@dbx-tools/shared-model": "^0.9.153",
+    "@dbx-tools/appkit": "^0.9.154",
+    "@dbx-tools/cache": "^0.9.154",
+    "@dbx-tools/core": "^0.9.154",
+    "@dbx-tools/databricks": "^0.9.154",
+    "@dbx-tools/fs": "^0.9.154",
+    "@dbx-tools/genie": "^0.9.154",
+    "@dbx-tools/model": "^0.9.154",
+    "@dbx-tools/model-protocol": "^0.9.154",
+    "@dbx-tools/path": "^0.9.154",
+    "@dbx-tools/postgres": "^0.9.154",
+    "@dbx-tools/shared-core": "^0.9.154",
+    "@dbx-tools/shared-fs": "^0.9.154",
+    "@dbx-tools/shared-genie": "^0.9.154",
+    "@dbx-tools/shared-mastra": "^0.9.154",
+    "@dbx-tools/shared-model": "^0.9.154",
     "@mastra/ai-sdk": "1.10.5",
     "@mastra/core": "1.71.0",
     "@mastra/express": "1.5.15",
@@ -231,7 +231,7 @@
       "./package.json": "./package.json"
     }
   },
-  "version": "0.9.153",
+  "version": "0.9.154",
   "types": "./lib/index.d.ts",
   "dbxToolsConfig": {
     "tags": [

package/src/config.ts

@@ -241,15 +241,15 @@ export interface MastraPluginConfig extends BasePluginConfig {
   /** AI SDK provider-options namespace. Defaults to `"openai"`; no env fallback. */
   providerId?: string;
   /**
-   * PostgresStore for Mastra threads/messages. `true` captures the
-   * `lakebase` plugin's existing service-principal pool; an object opens a
-   * separately configured store.
+   * PostgresStore for Mastra threads/messages. `true` derives a dedicated,
+   * lifecycle-isolated service-principal pool from the `lakebase` plugin; an
+   * object opens a separately configured store.
    */
   storage?: boolean | PostgresStoreConfig;
   /**
-   * PgVector store for Mastra memory recall. `true` captures the
-   * `lakebase` plugin's existing service-principal pool; an object opens a
-   * separately configured store.
+   * PgVector store for Mastra memory recall. `true` derives a dedicated,
+   * lifecycle-isolated service-principal pool from the `lakebase` plugin; an
+   * object opens a separately configured store.
    */
   memory?: boolean | MastraMemoryConfig;
   /**
@@ -716,12 +716,12 @@ export const MASTRA_CONFIG_SCHEMA: ConfigSchema = {
     storage: {
       type: ["boolean", "object"],
       description:
-        "PostgresStore for Mastra threads / messages. `true` captures the `lakebase` plugin's existing service-principal pool; an object opens a separately configured store. Auto-enabled when the `lakebase` plugin is registered.",
+        "PostgresStore for Mastra threads / messages. `true` derives a dedicated lifecycle-isolated service-principal pool from the `lakebase` plugin; an object opens a separately configured store. Auto-enabled when the `lakebase` plugin is registered.",
     },
     memory: {
       type: ["boolean", "object"],
       description:
-        "PgVector store for Mastra semantic recall. `true` captures the `lakebase` plugin's existing service-principal pool; an object opens a separately configured store. Auto-enabled when the `lakebase` plugin is registered.",
+        "PgVector store for Mastra semantic recall. `true` derives a dedicated lifecycle-isolated service-principal pool from the `lakebase` plugin; an object opens a separately configured store. Auto-enabled when the `lakebase` plugin is registered.",
     },
     codeMode: {
       type: ["boolean", "string", "object"],

package/src/memory.ts

@@ -43,7 +43,7 @@ import {
 import { fastembed } from "@mastra/fastembed";
 import { Memory } from "@mastra/memory";
 import { PgVector, PostgresStore } from "@mastra/pg";
-import type { Pool } from "pg";
+import { Pool, type PoolConfig } from "pg";
 
 import type { AnyMastraAgentDefinition, MastraMemoryConfigOverride } from "./agents.ts";
 import type { MastraPluginConfig } from "./config.ts";
@@ -160,6 +160,24 @@ export function createMemoryBuilder(
   return new MemoryBuilder(config, servicePrincipalPool);
 }
 
+/**
+ * Create the plugin-owned Lakebase pool used by durable workflow storage.
+ *
+ * @internal
+ *
+ * AppKit shuts plugin hooks down concurrently, so borrowing the native
+ * Lakebase plugin's pool lets its shutdown race Mastra's final snapshot writes.
+ * A caller-provided pool is not closed by PostgresStore; idle connections do
+ * not keep the process alive and remain available to runs Mastra is still
+ * draining after its bounded shutdown window.
+ */
+export function createMastraLakebasePool(config: PoolConfig): Pool {
+  return new Pool({
+    ...config,
+    allowExitOnIdle: true,
+  });
+}
+
 /**
  * Builds one `Memory` per agent against a shared service-principal
  * Lakebase pool. Per-instance state keeps the shared `PgVector` alive

package/src/plugin.ts

@@ -56,7 +56,7 @@ import {
   type PluginManifest,
   type ResourceRequirement,
 } from "@databricks/appkit";
-import { lakebaseResolver, pluginRegistry, telemetry as appkitTelemetry } from "@dbx-tools/appkit";
+import { pluginRegistry, telemetry as appkitTelemetry } from "@dbx-tools/appkit";
 import { sendGenieMessageFeedback } from "@dbx-tools/genie/feedback";
 import { modelCatalog } from "@dbx-tools/model";
 import { asyncUtils, errorUtils, log, object, stringUtils } from "@dbx-tools/shared-core";
@@ -77,7 +77,6 @@ import { RequestContext } from "@mastra/core/request-context";
 import type { AnyWorkflow } from "@mastra/core/workflows";
 import { pipeUIMessageStreamToResponse } from "ai";
 import express from "express";
-import type { Pool } from "pg";
 
 import {
   buildAgents,
@@ -120,7 +119,7 @@ import {
   type MastraIdentityMode,
 } from "./identity.ts";
 import { buildMcpServer, type ResolvedMcp } from "./mcp.ts";
-import { createMemoryBuilder, needsLakebase } from "./memory.ts";
+import { createMastraLakebasePool, createMemoryBuilder, needsLakebase } from "./memory.ts";
 import {
   flushDirectMlflowTracing,
   logFeedback,
@@ -1341,9 +1340,9 @@ export class MastraPlugin extends Plugin<MastraPluginConfig> {
 
   private async buildAgentAndServer(): Promise<void> {
     const memoryPool = needsLakebase(this.config)
-      ? (lakebaseResolver.captureServicePrincipalPool(
-          pluginRegistry.require(this.context, lakebase, this.config).exports().pool,
-        ) as Pool)
+      ? createMastraLakebasePool(
+          pluginRegistry.require(this.context, lakebase, this.config).exports().getPgConfig(),
+        )
       : undefined;
     const memoryBuilder = createMemoryBuilder(this.config, memoryPool);