@dbx-tools/appkit-mastra 0.9.119 → 0.9.120

package/src/genie.ts

@@ -248,7 +248,7 @@ async function readCachedConversationId(cacheKey: string | undefined): Promise<s
     return v ?? undefined;
   } catch (err) {
     logger.warn("conversation-cache:read-error", {
-      error: errorUtils.errorMessage(err),
+      error: errorUtils.toError(err),
     });
     return undefined;
   }
@@ -272,7 +272,7 @@ async function saveCachedConversationId(
     });
   } catch (err) {
     logger.warn("conversation-cache:write-error", {
-      error: errorUtils.errorMessage(err),
+      error: errorUtils.toError(err),
     });
   }
 }
@@ -284,7 +284,7 @@ async function evictCachedConversationId(cacheKey: string | undefined): Promise<
     await CacheManager.getInstanceSync().delete(cacheKey);
   } catch (err) {
     logger.warn("conversation-cache:delete-error", {
-      error: errorUtils.errorMessage(err),
+      error: errorUtils.toError(err),
     });
   }
 }
@@ -515,7 +515,7 @@ function buildAskGenieTool(opts: {
             logger.warn("conversation-cache:stale, resetting", {
               spaceId,
               conversationId: rejectedConversationId,
-              error: errorUtils.errorMessage(err),
+              error: errorUtils.toError(err),
             });
             if (ownsReusableConversation) {
               await evictCachedConversationId(cacheKey);
@@ -1115,7 +1115,9 @@ async function writeSuggestionCache(spaceId: string, questions: string[]): Promi
  * questions and a multi-space app round-trips breadth-first up to the
  * cap. A per-space fetch failure degrades to "no questions for that
  * space" (logged, not thrown) so one unreachable space never blanks
- * the whole list. Returns `[]` when no spaces are configured.
+ * the whole list. When `signal` aborts (the HTTP caller disconnected),
+ * collection stops silently and returns what was merged so far.
+ * Returns `[]` when no spaces are configured.
  */
 export async function collectSpaceSuggestions(opts: {
   spaces: Record<string, GenieSpaceConfig>;
@@ -1138,9 +1140,10 @@ export async function collectSpaceSuggestions(opts: {
         questions = genieSpace.genieSampleQuestions(space);
         await writeSuggestionCache(spaceId, questions);
       } catch (err) {
+        if (opts.signal?.aborted) return merged;
         logger.warn("suggestions:fetch-error", {
           spaceId,
-          error: errorUtils.errorMessage(err),
+          error: errorUtils.toError(err),
         });
         questions = [];
       }

package/src/memory.ts

@@ -32,9 +32,7 @@
  * @module
  */
 
-import { getUsernameWithApiLookup } from "@databricks/appkit";
 import { migration as appkitMigration } from "@dbx-tools/appkit";
-import { postgresConnectionOptions } from "@dbx-tools/postgres";
 import { hash, log } from "@dbx-tools/shared-core";
 import type { AgentMemoryOption } from "@mastra/core/agent";
 import {
@@ -45,7 +43,7 @@ import {
 import { fastembed } from "@mastra/fastembed";
 import { Memory } from "@mastra/memory";
 import { PgVector, PostgresStore } from "@mastra/pg";
-import { Pool, type PoolConfig } from "pg";
+import type { Pool } from "pg";
 
 import type { AnyMastraAgentDefinition, MastraMemoryConfigOverride } from "./agents.ts";
 import type { MastraPluginConfig } from "./config.ts";
@@ -132,34 +130,6 @@ export function withSoftStorageInit(store: PostgresStore): PostgresStore {
   return store;
 }
 
-/**
- * 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(postgresConnectionOptions({ ...pgConfig, user }));
-}
-
 /** Effective per-knob setting after the plugin/agent cascade. */
 type StorageSetting = AnyMastraAgentDefinition["storage"];
 type MemorySetting = AnyMastraAgentDefinition["memory"];
@@ -179,7 +149,7 @@ export function needsLakebase(config: MastraPluginConfig): boolean {
 
 /**
  * Construct a per-agent {@link Memory} factory bound to the supplied
- * service-principal pool (see {@link createServicePrincipalPool}).
+ * service-principal pool.
  * Caches the shared `PgVector` singleton (built on first need) so each
  * agent build is O(1) after the first.
  */

package/src/mlflow.ts

@@ -366,7 +366,7 @@ export async function initializeAppMlflowTraceInfo(client?: WorkspaceClient): Pr
     return true;
   } catch (err) {
     logger.warn("Databricks Apps MLflow trace-info promotion disabled", {
-      error: errorUtils.errorMessage(err),
+      error: errorUtils.toError(err),
     });
     return false;
   }
@@ -424,7 +424,7 @@ export function persistAppMlflowTraceInfo(params: AppMlflowTraceInfo): void {
     .catch((err) => {
       logger.warn("Databricks Apps MLflow trace-info promotion failed", {
         traceId: params.traceId,
-        error: errorUtils.errorMessage(err),
+        error: errorUtils.toError(err),
       });
     })
     .finally(() => pendingAppMlflowTraceInfo.delete(pending));
@@ -466,7 +466,7 @@ async function detectedDirectMlflowTraceLocation(
   } catch (err) {
     logger.warn("direct MLflow UC trace-location detection failed", {
       experimentId,
-      error: errorUtils.errorMessage(err),
+      error: errorUtils.toError(err),
     });
     return undefined;
   }
@@ -530,7 +530,7 @@ export async function initializeDirectMlflowTracing(): Promise<boolean> {
     return true;
   } catch (err) {
     logger.warn("direct MLflow tracing disabled", {
-      error: errorUtils.errorMessage(err),
+      error: errorUtils.toError(err),
     });
     return false;
   }
@@ -693,18 +693,10 @@ async function resolveExperimentId(client: WorkspaceClient): Promise<string | un
 }
 
 /** Parameters for {@link logFeedback}. */
-export interface LogFeedbackParams {
-  /** MLflow trace id the assessment attaches to (`tr-<hex>`). */
-  traceId: string;
-  /** Assessment name; defaults per whether a value or a comment is sent. */
-  name?: string;
-  /** Thumbs / rating / label value. Omit for a comment-only submission. */
-  value?: boolean | number | string;
-  /** Freeform comment: the rationale alongside a value, or the value itself when none. */
-  comment?: string;
+export type LogFeedbackParams = feedback.MastraFeedbackRequest & {
   /** Identity the feedback is attributed to (user email / id). */
   sourceId?: string;
-}
+};
 
 /**
  * Log a HUMAN feedback assessment to a trace. Returns the created
@@ -751,7 +743,7 @@ export async function logFeedback(
       logger.warn("feedback request failed", {
         traceId: params.traceId,
         assessmentTraceId,
-        error: errorUtils.errorMessage(err),
+        error: errorUtils.toError(err),
       });
       return undefined;
     }

package/src/monty-sandbox.ts

@@ -68,7 +68,7 @@ const montyContext = functionUtils.memoize(async (): Promise<MontyContext> => {
     } catch (error) {
       errors.push(errorUtils.toError(error));
       logger.debug(`${target.name} unavailable`, {
-        error: errorUtils.errorMessage(error),
+        error: errorUtils.toError(error),
       });
     }
   }

package/src/plugin.ts

@@ -56,7 +56,7 @@ import {
   type PluginManifest,
   type ResourceRequirement,
 } from "@databricks/appkit";
-import { pluginRegistry } from "@dbx-tools/appkit";
+import { lakebaseResolver, pluginRegistry } from "@dbx-tools/appkit";
 import { modelCatalog } from "@dbx-tools/model";
 import { asyncUtils, errorUtils, log, object, stringUtils } from "@dbx-tools/shared-core";
 import {
@@ -98,7 +98,7 @@ import {
   type MastraIdentityMode,
 } from "./identity.ts";
 import { buildMcpServer, type ResolvedMcp } from "./mcp.ts";
-import { createMemoryBuilder, createServicePrincipalPool, needsLakebase } from "./memory.ts";
+import { createMemoryBuilder, needsLakebase } from "./memory.ts";
 import {
   flushDirectMlflowTracing,
   logFeedback,
@@ -124,13 +124,6 @@ import { invalidFields } from "./validation.ts";
 const GENIE_MANIFEST = pluginRegistry.data(genie).plugin.manifest;
 const LAKEBASE_MANIFEST = pluginRegistry.data(lakebase).plugin.manifest;
 
-/**
- * Budget for draining the memory service-principal pool on shutdown. Well
- * inside AppKit's 15s graceful-shutdown window, so a stuck connection can
- * never be what holds the process open.
- */
-const POOL_DRAIN_TIMEOUT_MS = 5_000;
-
 /** Ceiling on the `?timeoutMs=` long-poll budget a client may request. */
 const MAX_EMBED_POLL_TIMEOUT_MS = 5 * 60_000;
 
@@ -243,15 +236,6 @@ export class MastraPlugin extends Plugin<MastraPluginConfig> {
    * boot rather than the first request, and so every request agrees.
    */
   private identityMode: MastraIdentityMode = "user";
-  /**
-   * 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 shutdown}. `null` until setup runs or when
-   * Lakebase isn't needed.
-   */
-  private servicePrincipalPool: Pool | null = null;
-
   override async setup(): Promise<void> {
     // Resolve the identity mode up front so an invalid `genieIdentity` fails the
     // app boot with a clear message rather than 500-ing the first chat request.
@@ -298,44 +282,19 @@ export class MastraPlugin extends Plugin<MastraPluginConfig> {
     if (this.config.memory === undefined) this.config.memory = true;
   }
 
-  /**
-   * Drain the memory service-principal pool. Idempotent: the handle is
-   * cleared before the drain starts, so a second call is a no-op and a
-   * later `setup()` rebuilds the pool. Bounded by
-   * {@link POOL_DRAIN_TIMEOUT_MS} to stay well inside the 15s graceful
-   * shutdown budget.
-   */
+  /** Flush pending direct MLflow traces during graceful shutdown. */
   async shutdown(): Promise<void> {
     try {
       await flushDirectMlflowTracing();
     } catch (err) {
       this.logger.warn("error flushing direct MLflow traces", {
-        error: errorUtils.errorMessage(err),
-      });
-    }
-    const pool = this.servicePrincipalPool;
-    if (!pool) return;
-    this.servicePrincipalPool = null;
-    this.logger.info("closing memory SP pool");
-    // The budget timer is cancelled on the way out so a drain that finished
-    // early cannot hold the event loop open for the rest of the window.
-    const budget = new AbortController();
-    try {
-      await Promise.race([pool.end(), asyncUtils.sleep(POOL_DRAIN_TIMEOUT_MS, budget.signal)]);
-    } catch (err) {
-      this.logger.error("error closing memory SP pool", {
-        error: errorUtils.errorMessage(err),
+        error: errorUtils.toError(err),
       });
-    } finally {
-      budget.abort();
     }
   }
 
   /**
-   * Abort in-flight work. AppKit's graceful shutdown calls this hook
-   * synchronously and never awaits {@link shutdown}, so the pool drain is
-   * started here too; it cannot be awaited from a `void` hook, and
-   * {@link shutdown} is idempotent so the duplicate call is free.
+   * Abort in-flight work and start the asynchronous trace flush.
    */
   override abortActiveOperations(): void {
     super.abortActiveOperations();
@@ -719,7 +678,7 @@ export class MastraPlugin extends Plugin<MastraPluginConfig> {
         // the upstream status. Log and degrade.
         this.logger.warn("suggestions:error", {
           status: result.status,
-          error: result.message,
+          error: errorUtils.toError(result.message),
         });
         res.json({ questions: [] });
         return;
@@ -765,7 +724,7 @@ export class MastraPlugin extends Plugin<MastraPluginConfig> {
         if (!parsed.success) {
           // A Zod issue string quotes the received body back, so only the
           // offending field names travel; the full issue list is logged.
-          this.logger.warn("feedback:invalid", { error: parsed.error.message });
+          this.logger.warn("feedback:invalid", { error: parsed.error });
           res.status(400).json({
             ok: false,
             error: INVALID_FEEDBACK_MESSAGE,
@@ -856,7 +815,7 @@ export class MastraPlugin extends Plugin<MastraPluginConfig> {
   ): void {
     this.logger.warn(`${operation}:error`, {
       status: result.status,
-      error: result.message,
+      error: errorUtils.toError(result.message),
     });
     if (res.headersSent) return;
     res.status(result.status).json({ error: clientMessage });
@@ -1074,32 +1033,18 @@ export class MastraPlugin extends Plugin<MastraPluginConfig> {
     if (result.ok) return result.data;
     this.logger.warn("listModels:error", {
       status: result.status,
-      error: result.message,
+      error: errorUtils.toError(result.message),
     });
     throw new ExecutionError(MODEL_CATALOGUE_FAILED_MESSAGE);
   }
 
   private async buildAgentAndServer(): Promise<void> {
-    // 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 = pluginRegistry
-        .require(this.context, lakebase, this.config)
-        .exports()
-        .getPgConfig();
-      this.servicePrincipalPool = await createServicePrincipalPool(spPgConfig);
-    }
-    const memoryBuilder = this.servicePrincipalPool
-      ? createMemoryBuilder(this.config, this.servicePrincipalPool)
+    const memoryPool = needsLakebase(this.config)
+      ? (lakebaseResolver.captureServicePrincipalPool(
+          pluginRegistry.require(this.context, lakebase, this.config).exports().pool,
+        ) as Pool)
       : undefined;
+    const memoryBuilder = memoryPool ? createMemoryBuilder(this.config, memoryPool) : undefined;
 
     // Build every agent declared in `config.agents` (or the built-in
     // fallback when none are declared). Each agent's `model` resolves

package/src/remote-skills.ts

@@ -436,7 +436,7 @@ async function readMetadata(fs: FileSystem): Promise<RemoteSkillsMetadata | unde
     if (!object.isRecord(parsed) || parsed.version !== METADATA_VERSION) return undefined;
     return object.isRecord(parsed.sources) ? parsed : undefined;
   } catch (err) {
-    logger.debug("cache:unreadable", { error: errorUtils.errorMessage(err) });
+    logger.debug("cache:unreadable", { error: errorUtils.toError(err) });
     return undefined;
   }
 }
@@ -599,7 +599,7 @@ export async function provisionRemoteSkills(
         }
         logger.warn("source:skipped", {
           source: sourceOptions.source,
-          error: errorUtils.errorMessage(err),
+          error: errorUtils.toError(err),
         });
       }
     }
@@ -658,7 +658,10 @@ async function cachedAiToolsTree(
 async function serializeSkillDirs(dirs: readonly StagedSkillDir[]): Promise<CachedSkillFile[]> {
   const files: CachedSkillFile[] = [];
   for (const dir of dirs) {
-    for (const relative of find.findFiles("**/*", { cwd: dir.absolutePath, nodir: true })) {
+    for await (const relative of find.findFiles("**/*", {
+      cwd: dir.absolutePath,
+      nodir: true,
+    })) {
       files.push({
         path: posix.join(dir.name, relative.split(/[\\/]/).join("/")),
         content: (await readFile(join(dir.absolutePath, relative))).toString("base64"),
@@ -752,7 +755,7 @@ async function openWritableWorkspace(
   } catch (err) {
     logger.warn("destination:unwritable", {
       destination: basePath,
-      error: errorUtils.errorMessage(err),
+      error: errorUtils.toError(err),
     });
     return undefined;
   }
@@ -996,7 +999,10 @@ async function collectSkillDirs(root: string): Promise<StagedSkillDir[]> {
 async function copySkillDirs(destination: FileSystem, dirs: StagedSkillDir[]): Promise<void> {
   await destination.init();
   for (const dir of dirs) {
-    for (const relative of find.findFiles("**/*", { cwd: dir.absolutePath, nodir: true })) {
+    for await (const relative of find.findFiles("**/*", {
+      cwd: dir.absolutePath,
+      nodir: true,
+    })) {
       const buffer = await readFile(join(dir.absolutePath, relative));
       const skillPath = posix.join(dir.name, relative.split(/[\\/]/).join("/"));
       await destination.writeFile(skillPath, buffer, { overwrite: true });

package/src/server.ts

@@ -484,7 +484,7 @@ export function attachStaleResumeRecovery(app: express.Express): void {
       }
       logger.warn("ignored stale mastra resume", {
         path: req.path,
-        error: errorUtils.errorMessage(error),
+        error: errorUtils.toError(error),
       });
       if (res.headersSent) {
         res.end();

package/src/workspaces.ts

@@ -840,7 +840,7 @@ async function loadInMemorySkillFilesystem(mount: LocalSkillMount): Promise<File
     root: "/",
   });
   await memory.init();
-  for (const relative of find.findFiles("**/*", { cwd: mount.root, nodir: true })) {
+  for await (const relative of find.findFiles("**/*", { cwd: mount.root, nodir: true })) {
     await memory.writeFile(relative, await source.readFile(relative), { overwrite: true });
   }
   await source.close();
@@ -896,7 +896,7 @@ async function databricksFilesystem(
     logger.debug("databricks-mount:skipped", {
       root,
       readOnly,
-      error: errorUtils.errorMessage(err),
+      error: errorUtils.toError(err),
     });
   }
   return undefined;
@@ -945,7 +945,7 @@ function verifyFirstFileSystemAccess<TFileSystem extends FileSystem>(
       readCheck = undefined;
       logger.debug("databricks-mount:read-check-failed", {
         root,
-        error: errorUtils.errorMessage(error),
+        error: errorUtils.toError(error),
       });
       throw error;
     }
@@ -976,7 +976,7 @@ function verifyFirstFileSystemAccess<TFileSystem extends FileSystem>(
             if (!writeVerified) {
               logger.debug("databricks-mount:write-check-failed", {
                 root,
-                error: errorUtils.errorMessage(error),
+                error: errorUtils.toError(error),
               });
             }
             throw error;
@@ -1079,7 +1079,7 @@ async function resolveDistinctSkillPaths(
     } catch (error) {
       logger.debug("skill-path:skipped", {
         path,
-        error: errorUtils.errorMessage(error),
+        error: errorUtils.toError(error),
       });
     }
   }

package/src/writer.ts

@@ -44,7 +44,7 @@ export async function safeWriteProgress(
   } catch (err) {
     log.warn("writer:error", {
       ...context,
-      error: errorUtils.errorMessage(err),
+      error: errorUtils.toError(err),
     });
   }
 }