@dbx-tools/appkit-mastra 0.9.153 → 0.9.157

package/src/genie-chat.ts

@@ -13,16 +13,19 @@
  */
 
 import type { WorkspaceClient } from "@databricks/appkit";
-import { databricks } from "@dbx-tools/appkit";
-import { agentMode, chat } from "@dbx-tools/genie";
-import { errorUtils, hash, log, object } from "@dbx-tools/shared-core";
-import { event as genieEvent, genieModel, type GenieMessage } from "@dbx-tools/shared-genie";
+import { toContext } from "@dbx-tools/appkit";
+import { listAgentModeConversation, genieEventChat } from "@dbx-tools/genie";
+import { log, toError, id, isRecord } from "@dbx-tools/shared-core";
+import { type GenieMessage, eventsFromMessage, humanizeStatus } from "@dbx-tools/shared-genie";
 import {
-  markers,
-  wire,
   type GenieChatRequest,
   type GenieThread,
   type StartedEvent,
+  formatMarker,
+  GENIE_PROGRESS_PART_TYPE,
+  type GenieWriterEvent,
+  type GenieConversationPartData,
+  GENIE_CONVERSATION_PART_TYPE,
 } from "@dbx-tools/shared-mastra";
 import type { MastraDBMessage } from "@mastra/core/agent";
 import type { MastraMemory } from "@mastra/core/memory";
@@ -90,7 +93,7 @@ export function genieAnswerText(message: GenieMessage): string {
       const queryId = queriesByVisualization.get(visualizationId);
       const statementId = queryId ? statementsByQuery.get(queryId) : undefined;
       if (statementId) {
-        sections.push(markers.formatMarker("visualization", statementId));
+        sections.push(formatMarker("visualization", statementId));
         inlineStatements.add(statementId);
       }
       continue;
@@ -106,8 +109,8 @@ export function genieAnswerText(message: GenieMessage): string {
     if (description) sections.push(description);
     sections.push(
       attachment.attachment_id && visualizedQueries.has(attachment.attachment_id)
-        ? markers.formatMarker("visualization", statementId)
-        : markers.formatMarker("data", statementId),
+        ? formatMarker("visualization", statementId)
+        : formatMarker("data", statementId),
     );
   }
   if (sections.length === 0 && message.error?.error) sections.push(message.error.error);
@@ -115,19 +118,19 @@ export function genieAnswerText(message: GenieMessage): string {
 }
 
 function visualizationAttachmentId(metadata: unknown): string | undefined {
-  if (!object.isRecord(metadata) || !object.isRecord(metadata.viz)) return undefined;
+  if (!isRecord(metadata) || !isRecord(metadata.viz)) return undefined;
   return typeof metadata.viz.attachment_id === "string" ? metadata.viz.attachment_id : undefined;
 }
 
 /** Whether a text attachment is Agent Mode's Markdown copy of a query result. */
 function isQueryOutputText(metadata: unknown): boolean {
-  return object.isRecord(metadata) && typeof metadata.statement_id === "string";
+  return isRecord(metadata) && typeof metadata.statement_id === "string";
 }
 
 /** Error text for a terminal Genie message that did not complete, or `undefined`. */
 function genieFailure(message: GenieMessage): string | undefined {
   if (message.status === "COMPLETED") return undefined;
-  return message.error?.error ?? `Genie ${genieModel.humanizeStatus(message.status!)}`;
+  return message.error?.error ?? `Genie ${humanizeStatus(message.status!)}`;
 }
 
 /** Recover the original user question from a context-augmented Genie input. */
@@ -140,9 +143,7 @@ export function visibleGenieQuestion(content: string): string {
 function messageText(message: MastraDBMessage): string | undefined {
   const text = message.content.parts
     .flatMap((part) =>
-      object.isRecord(part) && part.type === "text" && typeof part.text === "string"
-        ? [part.text]
-        : [],
+      isRecord(part) && part.type === "text" && typeof part.text === "string" ? [part.text] : [],
     )
     .join("\n")
     .trim();
@@ -230,7 +231,7 @@ async function saveGenieMemory(options: {
       saveThread: true,
     });
   }
-  const messageId = message.message_id || message.id || hash.id();
+  const messageId = message.message_id || message.id || id();
   const createdAt = new Date();
   await memory.saveMessages({
     messages: [
@@ -269,9 +270,9 @@ export function genieMessageToUiMessages(
   const question = visibleGenieQuestion(message.content);
   const input = { question };
   const failure = genieFailure(message);
-  const progress = [...genieEvent.eventsFromMessage(message, undefined, space.spaceId)].map(
+  const progress = [...eventsFromMessage(message, undefined, space.spaceId)].map(
     (progressEvent) => ({
-      type: wire.GENIE_PROGRESS_PART_TYPE,
+      type: GENIE_PROGRESS_PART_TYPE,
       data: { toolCallId, event: progressEvent },
     }),
   );
@@ -312,7 +313,7 @@ export async function listGenieThreads(
   space: GenieChatSpace,
   signal?: AbortSignal,
 ): Promise<GenieThread[]> {
-  const context = signal ? databricks.toContext(signal) : undefined;
+  const context = signal ? toContext(signal) : undefined;
   const threads: GenieThread[] = [];
   let pageToken: string | undefined;
   do {
@@ -325,7 +326,7 @@ export async function listGenieThreads(
       context,
     );
     for (const conversation of page.conversations ?? []) {
-      if (!object.isRecord(conversation) || conversation.agent_type !== AGENT_CONVERSATION_TYPE) {
+      if (!isRecord(conversation) || conversation.agent_type !== AGENT_CONVERSATION_TYPE) {
         continue;
       }
       threads.push({
@@ -351,7 +352,7 @@ export async function listGenieThreadMessages(options: {
   signal?: AbortSignal;
 }): Promise<UIMessage[]> {
   const { space } = options;
-  const messages = await agentMode.listAgentModeConversation(
+  const messages = await listAgentModeConversation(
     options.client,
     space.spaceId,
     options.conversationId,
@@ -379,7 +380,7 @@ export function genieChatStream(options: {
   signal?: AbortSignal;
 }): ReadableStream<UIMessageChunk> {
   const { client, space, request, memory, requestContext, signal } = options;
-  const toolCallId = hash.id();
+  const toolCallId = id();
   const toolName = askGenieToolId(space.alias);
   return createUIMessageStream({
     onError: () => DIRECT_GENIE_ERROR,
@@ -396,13 +397,13 @@ export function genieChatStream(options: {
       } catch (error) {
         logger.warn("memory:recall-failed", {
           spaceId: space.spaceId,
-          error: errorUtils.toError(error),
+          error: toError(error),
         });
       }
       const question = contextualQuestion(request.question, request.requestContext, memoryContext);
-      const progress = (progressEvent: wire.GenieWriterEvent) =>
+      const progress = (progressEvent: GenieWriterEvent) =>
         writer.write({
-          type: wire.GENIE_PROGRESS_PART_TYPE,
+          type: GENIE_PROGRESS_PART_TYPE,
           data: { toolCallId, event: progressEvent },
         });
       writer.write({ type: "start" });
@@ -420,10 +421,10 @@ export function genieChatStream(options: {
       };
       progress(started);
 
-      let announced: wire.GenieConversationPartData | undefined;
+      let announced: GenieConversationPartData | undefined;
       let finalMessage: GenieMessage | undefined;
       try {
-        for await (const chatEvent of chat.genieEventChat(space.spaceId, question, {
+        for await (const chatEvent of genieEventChat(space.spaceId, question, {
           workspaceClient: client,
           agentMode: true,
           agentModeFallback: false,
@@ -443,7 +444,7 @@ export function genieChatStream(options: {
           ) {
             announced = { conversationId, ...(messageId ? { messageId } : {}) };
             writer.write({
-              type: wire.GENIE_CONVERSATION_PART_TYPE,
+              type: GENIE_CONVERSATION_PART_TYPE,
               data: announced,
               transient: true,
             });
@@ -454,7 +455,7 @@ export function genieChatStream(options: {
       } catch (error) {
         logger.warn("turn:error", {
           spaceId: space.spaceId,
-          error: errorUtils.toError(error),
+          error: toError(error),
         });
         writer.write({
           type: "tool-output-error",
@@ -474,7 +475,7 @@ export function genieChatStream(options: {
       writer.write({ type: "finish-step" });
       const answer = genieAnswerText(finalMessage);
       if (answer) {
-        const textId = hash.id();
+        const textId = id();
         writer.write({ type: "start-step" });
         writer.write({ type: "text-start", id: textId });
         writer.write({ type: "text-delta", id: textId, delta: answer });
@@ -495,7 +496,7 @@ export function genieChatStream(options: {
         } catch (error) {
           logger.warn("memory:save-failed", {
             spaceId: space.spaceId,
-            error: errorUtils.toError(error),
+            error: toError(error),
           });
         }
       }

package/src/genie.ts

@@ -44,11 +44,11 @@ import {
   ValidationError,
 } from "@databricks/appkit";
 import type { WorkspaceClient } from "@databricks/appkit";
-import { pluginRegistry } from "@dbx-tools/appkit";
-import { createCache } from "@dbx-tools/cache/cache";
-import { chat, space as genieSpace } from "@dbx-tools/genie";
-import { errorUtils, log, stringUtils } from "@dbx-tools/shared-core";
-import { genieModel, type GenieMessage } from "@dbx-tools/shared-genie";
+import { type PluginContextLike, instance } from "@dbx-tools/appkit";
+import { createCache } from "@dbx-tools/cache";
+import { genieEventChat, getGenieSpace, genieSampleQuestions } from "@dbx-tools/genie";
+import { log, toError, errorContext, toDescription, toIdentifier } from "@dbx-tools/shared-core";
+import { type GenieMessage, GenieMessageSchema } from "@dbx-tools/shared-genie";
 import type { StartedEvent } from "@dbx-tools/shared-mastra";
 import type { RequestContext } from "@mastra/core/request-context";
 import { MASTRA_THREAD_ID_KEY } from "@mastra/core/request-context";
@@ -249,7 +249,7 @@ async function readCachedConversationId(cacheKey: string | undefined): Promise<s
     return v ?? undefined;
   } catch (err) {
     logger.warn("conversation-cache:read-error", {
-      error: errorUtils.toError(err),
+      error: toError(err),
     });
     return undefined;
   }
@@ -273,7 +273,7 @@ async function saveCachedConversationId(
     });
   } catch (err) {
     logger.warn("conversation-cache:write-error", {
-      error: errorUtils.toError(err),
+      error: toError(err),
     });
   }
 }
@@ -285,7 +285,7 @@ async function evictCachedConversationId(cacheKey: string | undefined): Promise<
     await CacheManager.getInstanceSync().delete(cacheKey);
   } catch (err) {
     logger.warn("conversation-cache:delete-error", {
-      error: errorUtils.toError(err),
+      error: toError(err),
     });
   }
 }
@@ -328,7 +328,7 @@ const prepareChartRequestSchema = chartPlannerRequestSchema
       .string()
       .min(1, "statement_id is required")
       .describe(
-        stringUtils.toDescription(`
+        toDescription(`
           Genie \`statement_id\` to chart. Read from
           \`message.query_result.statement_id\` or
           \`message.attachments[*].query.statement_id\` returned by
@@ -349,7 +349,7 @@ const prepareChartRequestSchema = chartPlannerRequestSchema
  */
 function aliasSuffix(alias: string): string {
   if (alias === DEFAULT_GENIE_ALIAS) return "";
-  const slug = stringUtils.toIdentifier(alias);
+  const slug = toIdentifier(alias);
   return slug ? `_${slug}` : "";
 }
 
@@ -390,7 +390,7 @@ function buildAskGenieTool(opts: {
   const hintLine = hint ? ` (${hint})` : "";
   return createTool({
     id: toolId,
-    description: stringUtils.toDescription(`
+    description: toDescription(`
       Ask the Genie space "${alias}"${hintLine} ONE focused
       natural-language sub-question and wait for the turn to
       complete. Genie answers best when each call covers a
@@ -422,7 +422,7 @@ function buildAskGenieTool(opts: {
         .string()
         .min(1, "question is required")
         .describe(
-          stringUtils.toDescription(`
+          toDescription(`
             ONE focused natural-language question about the data in this
             space, covering a single metric / dimension / time window
             (e.g. "What was Q3 revenue by region?"). Decompose a
@@ -432,7 +432,7 @@ function buildAskGenieTool(opts: {
         ),
     }),
     outputSchema: z.object({
-      message: genieModel.GenieMessageSchema,
+      message: GenieMessageSchema,
     }),
     execute: async ({ question }, ctxRaw) => {
       const ctx = ctxRaw as ToolExecuteCtx;
@@ -485,7 +485,7 @@ function buildAskGenieTool(opts: {
         const runTurn = async (): Promise<GenieMessage> => {
           const seedConversationId = conversationId;
           let finalMessage: GenieMessage | undefined;
-          for await (const event of chat.genieEventChat(spaceId, question, {
+          for await (const event of genieEventChat(spaceId, question, {
             workspaceClient: client,
             ...(agentMode !== undefined ? { agentMode } : {}),
             ...(seedConversationId ? { conversationId: seedConversationId } : {}),
@@ -525,11 +525,11 @@ function buildAskGenieTool(opts: {
           // A rejected conversation id is stale or inaccessible. Reset this
           // lane and retry once with a fresh Genie conversation.
           const rejectedConversationId = conversationId;
-          if (rejectedConversationId && errorUtils.errorContext(err).notAccessible) {
+          if (rejectedConversationId && errorContext(err).notAccessible) {
             logger.warn("conversation-cache:stale, resetting", {
               spaceId,
               conversationId: rejectedConversationId,
-              error: errorUtils.toError(err),
+              error: toError(err),
             });
             if (ownsReusableConversation) {
               await evictCachedConversationId(cacheKey);
@@ -559,7 +559,7 @@ function buildSpaceDescriptionTool(opts: { spaceId: string; alias: string }) {
   const toolId = `get_space_description${aliasSuffix(alias)}`;
   return createTool({
     id: toolId,
-    description: stringUtils.toDescription(`
+    description: toDescription(`
       Return the Genie space "${alias}"'s title, description, and
       warehouse id. Cheap (single REST call, no LLM round-trip).
       Call this FIRST on any user turn that's going to touch
@@ -584,7 +584,7 @@ function buildSpaceDescriptionTool(opts: { spaceId: string; alias: string }) {
       // description surface (title / description / warehouse id) lives
       // on the directory-listing shape, so skip the larger
       // `serialized_space` payload here.
-      const space = await genieSpace.getGenieSpace(spaceId, {
+      const space = await getGenieSpace(spaceId, {
         workspaceClient: client,
         serialized: false,
         ...(signal ? { context: signal } : {}),
@@ -604,7 +604,7 @@ function buildSpaceSerializedTool(opts: { spaceId: string; alias: string }) {
   const toolId = `get_space_serialized${aliasSuffix(alias)}`;
   return createTool({
     id: toolId,
-    description: stringUtils.toDescription(`
+    description: toDescription(`
       Return the full \`GenieSpace\` JSON for the "${alias}" space.
       Use only when you need exact column / table identifiers
       \`get_space_description\` doesn't expose. Larger payload, so
@@ -619,7 +619,7 @@ function buildSpaceSerializedTool(opts: { spaceId: string; alias: string }) {
       // Central Genie space fetch with the opt-in `serialized_space`
       // blob (catalogs, tables, sample questions, prompts) that the
       // typed SDK `getSpace` omits - the whole point of this tool.
-      const space = await genieSpace.getGenieSpace(spaceId, {
+      const space = await getGenieSpace(spaceId, {
         workspaceClient: client,
         ...(signal ? { context: signal } : {}),
       });
@@ -643,7 +643,7 @@ function buildGetStatementTool() {
   const toolId = "get_statement";
   return createTool({
     id: toolId,
-    description: stringUtils.toDescription(`
+    description: toDescription(`
       Fetch the rows of a Genie statement by its \`statement_id\` (the
       value at \`message.query_result.statement_id\` or
       \`message.attachments[*].query.statement_id\` returned from
@@ -663,7 +663,7 @@ function buildGetStatementTool() {
         .string()
         .min(1, "statement_id is required")
         .describe(
-          stringUtils.toDescription(`
+          toDescription(`
             Genie \`statement_id\` whose rows to read. Take it from
             \`message.query_result.statement_id\` or
             \`message.attachments[*].query.statement_id\` on an
@@ -728,7 +728,7 @@ function buildPrepareChartTool(opts: { config: MastraPluginConfig }) {
   const toolId = "prepare_chart";
   return createTool({
     id: toolId,
-    description: stringUtils.toDescription([
+    description: toDescription([
       `
         Queue a chart for the rows of a Genie statement. Mints a
         \`chartId\` plus its complete \`marker\` synchronously and
@@ -800,7 +800,7 @@ function buildPrepareChartTool(opts: { config: MastraPluginConfig }) {
  * variant that names the suffixed per-space tools
  * (e.g. `ask_genie_sales`).
  */
-export const GENIE_INSTRUCTIONS = stringUtils.toDescription([
+export const GENIE_INSTRUCTIONS = toDescription([
   "Genie orchestration. For every user question that needs SQL-backed data:",
   {
     numbered: [
@@ -993,7 +993,7 @@ type AppKitGenieConfig = NonNullable<Parameters<typeof genie>[0]>;
  */
 export function resolveGenieSpaces(
   config: MastraPluginConfig,
-  context: pluginRegistry.PluginContextLike | undefined,
+  context: PluginContextLike | undefined,
 ): Record<string, GenieSpaceConfig> {
   const merged: Record<string, GenieSpaceConfig> = {};
 
@@ -1006,7 +1006,7 @@ export function resolveGenieSpaces(
   // Source 2: AppKit `genie()` plugin instance config. Use a
   // structural cast - `Plugin.config` is `protected` in TS only,
   // and the runtime layout is plain object property access.
-  const geniePlugin = pluginRegistry.instance(context, genie);
+  const geniePlugin = instance(context, genie);
   if (geniePlugin) {
     const pluginSpaces = (geniePlugin as unknown as { config?: AppKitGenieConfig }).config?.spaces;
     if (pluginSpaces) {
@@ -1147,17 +1147,17 @@ export async function collectSpaceSuggestions(opts: {
     let questions = await readSuggestionCache(spaceId);
     if (!questions) {
       try {
-        const space = await genieSpace.getGenieSpace(spaceId, {
+        const space = await getGenieSpace(spaceId, {
           workspaceClient: opts.client,
           ...(opts.signal ? { context: opts.signal } : {}),
         });
-        questions = genieSpace.genieSampleQuestions(space);
+        questions = genieSampleQuestions(space);
         await writeSuggestionCache(spaceId, questions);
       } catch (err) {
         if (opts.signal?.aborted) return merged;
         logger.warn("suggestions:fetch-error", {
           spaceId,
-          error: errorUtils.toError(err),
+          error: toError(err),
         });
         questions = [];
       }

package/src/identity.ts

@@ -51,50 +51,15 @@
  * @module
  */
 
-import { identity } from "@dbx-tools/appkit";
-
-/** Identity a chat turn's Databricks calls run as. See the module docs. */
-export type MastraIdentityMode = identity.IdentityMode;
+import { identity, type IdentityMode } from "@dbx-tools/appkit";
 
 /** Environment fallback for `config.genieIdentity`. */
 export const IDENTITY_ENV = "MASTRA_GENIE_IDENTITY";
 
-/** Every accepted {@link MastraIdentityMode}. */
-export const IDENTITY_MODES = identity.IDENTITY_MODES;
-
-/**
- * Default mode. `"user"` keeps OBO the only identity unless an app opts in, so
- * adopting this option can never silently widen an existing app's data access.
- */
-export const DEFAULT_IDENTITY_MODE = identity.DEFAULT_IDENTITY_MODE;
-
-/** Header Databricks Apps forward the signed-in user's id on. */
-export const USER_ID_HEADER = identity.USER_ID_HEADER;
-
-/** Header Databricks Apps forward the signed-in user's email on. */
-export const USER_EMAIL_HEADER = identity.USER_EMAIL_HEADER;
-
 /**
  * Resolve the configured mode: explicit plugin config, then {@link IDENTITY_ENV},
- * then {@link DEFAULT_IDENTITY_MODE}. An unrecognized value throws.
+ * then AppKit's default identity mode. An unrecognized value throws.
  */
-export function resolveIdentityMode(configured: string | undefined): MastraIdentityMode {
+export function resolveIdentityMode(configured: string | undefined): IdentityMode {
   return identity.resolveIdentityMode(configured, IDENTITY_ENV, "genieIdentity");
 }
-
-/**
- * The forwarded user id on `req`, or `undefined`. Used to attribute a turn to
- * the real caller when the Databricks client is the app SP but memory / cache /
- * traces must still key off the user.
- */
-export const requestUserId = identity.requestUserId;
-
-/** The forwarded user email on `req`, or `undefined`. */
-export const requestUserEmail = identity.requestUserEmail;
-
-/**
- * Whether this request should run its Databricks calls as the app service
- * principal rather than OBO - `"service-principal"` always, `"user"` never, and
- * `"auto"` only when the request carries no OBO token for AppKit to use.
- */
-export const useServicePrincipal = identity.useServicePrincipal;

package/src/memory.ts

@@ -32,8 +32,8 @@
  * @module
  */
 
-import { migration as appkitMigration } from "@dbx-tools/appkit";
-import { hash, log } from "@dbx-tools/shared-core";
+import { handleOwnershipMigrationError } from "@dbx-tools/appkit";
+import { log, id } from "@dbx-tools/shared-core";
 import type { AgentMemoryOption } from "@mastra/core/agent";
 import {
   MASTRA_RESOURCE_ID_KEY,
@@ -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";
@@ -108,7 +108,7 @@ export function withSoftStorageInit(store: PostgresStore): PostgresStore {
         await originalInit();
       } catch (err) {
         if (
-          !appkitMigration.handleOwnershipMigrationError(err, {
+          !handleOwnershipMigrationError(err, {
             scope: "mastra-storage",
             logger,
             event: "mastra storage migration failed",
@@ -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
@@ -329,7 +347,7 @@ export class MemoryBuilder {
  */
 function buildSharedPgVector(pool: Pool): PgVector {
   const vector = new PgVector({
-    id: `pg${hash.id()}`,
+    id: `pg${id()}`,
     // Keep the recall index out of `public`: on a Lakebase database the app
     // service principal has no CREATE on `public` (PG15+ locks it down), so a
     // default-schema PgVector fails on CREATE INDEX with "permission denied for
@@ -353,7 +371,7 @@ function buildSharedPgVector(pool: Pool): PgVector {
 /** Per-agent dedicated `PgVector` (rare; opt-in via object override). */
 function buildPgVector(setting: MastraMemoryConfigOverride): PgVector {
   return new PgVector(
-    withId(setting, `pg-vector__${hash.id()}`) as ConstructorParameters<typeof PgVector>[0],
+    withId(setting, `pg-vector__${id()}`) as ConstructorParameters<typeof PgVector>[0],
   );
 }