@dbx-tools/appkit-mastra 0.9.129 → 0.9.132

package/src/agents.ts

@@ -53,6 +53,7 @@ import { buildGenieToolkitProvider, resolveGenieSpaces } from "./genie.ts";
 import { withRequestMemoryScope, type MemoryBuilder } from "./memory.ts";
 import { buildModel, RESPONSES_PROVIDER_OPTIONS } from "./model.ts";
 import { stripStaleChartsProcessor } from "./processors.ts";
+import { MASTRA_ORIGIN_KEY } from "./request-context.ts";
 import { MASTRA_RESOLVED_MODEL_KEY } from "./serving.ts";
 import { TYPOGRAPHY_RULE } from "./style.ts";
 import { buildSummarizeTool } from "./summarize.ts";
@@ -419,6 +420,9 @@ export interface BuiltAgents {
   ambientTools: MastraTools;
 }
 
+/** Builds one optional output processor for a registered Mastra agent. */
+export type MastraOutputProcessorFactory = (agentId: string) => OutputProcessor | undefined;
+
 /** Fallback agent id used when `config.agents` is omitted entirely. */
 export const FALLBACK_AGENT_ID = "default";
 
@@ -705,6 +709,7 @@ export async function buildAgents(opts: {
   config: MastraPluginConfig;
   context: pluginRegistry.PluginContextLike | undefined;
   memoryBuilder?: MemoryBuilder;
+  outputProcessorFactories?: readonly MastraOutputProcessorFactory[];
   log: log.Logger;
   /**
    * Local skill filesystems mounted into every auto-created default workspace.
@@ -720,6 +725,7 @@ export async function buildAgents(opts: {
     config,
     context,
     memoryBuilder,
+    outputProcessorFactories = [],
     log,
     extraSkillMounts,
     standardSkills = [],
@@ -824,7 +830,13 @@ export async function buildAgents(opts: {
       ...(memory ? { memory } : {}),
       ...(workspace ? { workspace } : {}),
       inputProcessors: [...inputProcessors, ...workspaceSkillInputProcessors(workspace, config)],
-      outputProcessors: [toolErrorLoggingProcessor(log)],
+      outputProcessors: [
+        toolErrorLoggingProcessor(log),
+        ...outputProcessorFactories.flatMap((factory) => {
+          const processor = factory(id);
+          return processor ? [processor] : [];
+        }),
+      ],
     });
     agents[id] =
       config.backgroundTurns === false
@@ -1271,11 +1283,13 @@ function toolkitEntryToMastraTool(entry: ToolkitEntry, plugin: ContextualToolPro
         | undefined;
       const resourceId = execution?.requestContext?.get(MASTRA_RESOURCE_ID_KEY);
       const model = execution?.requestContext?.get(MASTRA_RESOLVED_MODEL_KEY);
+      const origin = execution?.requestContext?.get(MASTRA_ORIGIN_KEY);
       const toolCallId = execution?.agent?.toolCallId;
       const writer = execution?.writer;
       return plugin.executeAgentTool!(entry.localName, input, execution?.abortSignal, {
         ...(typeof resourceId === "string" ? { resourceId } : {}),
         ...(typeof model === "string" ? { model } : {}),
+        ...(typeof origin === "string" ? { origin } : {}),
         ...(writer && toolCallId
           ? {
               writeProgress: async (event, options) => {

package/src/config.ts

@@ -48,6 +48,8 @@ export const MASTRA_USER_KEY = "mastra__user";
 export const MASTRA_USER_NAME_KEY = "mastra__userName";
 /** Request-context key containing the authenticated user's email address. */
 export const MASTRA_USER_EMAIL_KEY = "mastra__userEmail";
+/** Request-context key identifying the Databricks credential used by the turn. */
+export const MASTRA_AUTH_MODE_KEY = "mastra__authMode";
 
 /**
  * `RequestContext` key for the per-HTTP-request id stamped by
@@ -88,6 +90,7 @@ export const TRACE_REQUEST_CONTEXT_KEYS: readonly string[] = [
   MASTRA_REQUEST_ID_KEY,
   MASTRA_USER_NAME_KEY,
   MASTRA_USER_EMAIL_KEY,
+  MASTRA_AUTH_MODE_KEY,
   // Model override key is owned by `serving.ts`; spelled inline here
   // so this module stays leaf-level (no cycles with `serving.ts`).
   "mastra__model_override",
@@ -450,7 +453,8 @@ export interface MastraPluginConfig extends BasePluginConfig {
    * Use the streaming Genie Agent Mode API for `ask_genie`. Defaults to `true`.
    * Set `false` to force the legacy Conversation API polling driver. A
    * pre-stream feature-disabled or preview-toggle response falls back
-   * automatically.
+   * automatically. Direct Genie chat always requires Agent Mode and ignores
+   * this option.
    */
   genieAgentMode?: boolean;
   /**
@@ -604,6 +608,8 @@ export interface MastraPluginConfig extends BasePluginConfig {
    * threads, the per-user cache namespace, and trace metadata still key off the
    * forwarded user, so callers sharing the service principal's data access keep
    * separate conversations and cannot read each other's threads or charts.
+   * Direct Genie chat does not use this option: its conversations and messages
+   * always run under AppKit OBO so Genie remains the per-user state owner.
    *
    * @example An app any account user can open
    * ```ts
@@ -658,7 +664,7 @@ export interface MastraPluginConfig extends BasePluginConfig {
    * ```
    */
   remoteSkills?: RemoteSkillsOption;
-  /** Read-only local skill trees mounted for files and passed as standard agent-level skills. */
+  /** Read-only local skill trees mounted for files, searchable by path, and passed as standard agent-level skills. */
   localSkillMounts?: LocalSkillMount[];
 }
 
@@ -768,7 +774,7 @@ export const MASTRA_CONFIG_SCHEMA: ConfigSchema = {
     genieAgentMode: {
       type: "boolean",
       description:
-        "Use the streaming Genie Agent Mode API for ask_genie. Defaults to true; false forces legacy polling. Feature-disabled or preview-toggle responses fall back automatically before streaming.",
+        "Use the streaming Genie Agent Mode API for agent-backed ask_genie. Defaults to true; false forces legacy polling. Feature-disabled or preview-toggle responses fall back automatically before streaming. Direct Genie chat always requires Agent Mode.",
     },
     genieSpaceCacheTtlMs: {
       type: "number",
@@ -810,7 +816,7 @@ export const MASTRA_CONFIG_SCHEMA: ConfigSchema = {
       // runtime while the published schema (and the UI built from it) omits it.
       enum: [...IDENTITY_MODES],
       description:
-        'Which Databricks identity the agents\' workspace calls (serving catalogue, Genie, statement execution) run as. "user" (default) is always OBO, so callers must be workspace members; "service-principal" always uses the app service principal, so any caller who can open the app works, at the cost of per-user attribution; "auto" uses OBO when the request carries an OBO token and the service principal when it does not, for an app that also serves a tunnel or webhook door. Falls back to MASTRA_GENIE_IDENTITY.',
+        'Which Databricks identity agent-backed workspace calls (serving catalogue, Genie, statement execution) run as. Direct Genie chat always uses OBO. "user" (default) is always OBO, so callers must be workspace members; "service-principal" always uses the app service principal, so any caller who can open the app works, at the cost of per-user attribution; "auto" uses OBO when the request carries an OBO token and the service principal when it does not, for an app that also serves a tunnel or webhook door. Falls back to MASTRA_GENIE_IDENTITY.',
     },
     apiAccess: {
       type: "string",

package/src/defaults.ts

@@ -82,6 +82,36 @@ export const genieSuggestionDefaults: ExecuteSettings = {
   },
 };
 
+/**
+ * `GET /genie/:alias/threads` and its messages: the caller's Genie
+ * conversations and their transcripts.
+ */
+export const genieConversationReadDefaults: ExecuteSettings = {
+  default: {
+    // Cache disabled: conversations change on every turn, and the chat refreshes
+    // the list right after a turn to pick up the new conversation and title.
+    cache: { enabled: false },
+    // Retry enabled: listing conversations and messages is an idempotent GET on
+    // the chat's load path.
+    retry: { enabled: true, attempts: 2 },
+    timeout: REST_TIMEOUT_MS,
+  },
+};
+
+/**
+ * `DELETE /genie/:alias/threads/:threadId` and its messages.
+ */
+export const genieConversationWriteDefaults: ExecuteSettings = {
+  default: {
+    // Cache disabled: a write.
+    cache: { enabled: false },
+    // Retry disabled: a delete that succeeded upstream but timed out locally
+    // would surface as a not-found on retry, masking the real outcome.
+    retry: { enabled: false },
+    timeout: REST_TIMEOUT_MS,
+  },
+};
+
 /**
  * `GET /embed/data/:id`: one Statement Execution result set.
  */

package/src/genie-chat.ts

@@ -0,0 +1,341 @@
+/**
+ * Direct Genie chat over the Mastra chat wire.
+ *
+ * Serves one Genie space as a chat without an agent, model, or Mastra Memory.
+ * Genie owns every piece of state: conversations are the chat threads, Genie
+ * messages are the transcript, and a turn is one Genie message. Each turn is
+ * projected onto the same AI SDK UI chunks an agent turn produces when it
+ * calls `ask_genie` (a native tool part, `data-genie-progress` events, and the
+ * answer text), so the Mastra chat UI renders pills, thoughts, suggestions,
+ * and embeds identically for both.
+ *
+ * @module
+ */
+
+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 {
+  markers,
+  wire,
+  type GenieChatRequest,
+  type GenieThread,
+  type StartedEvent,
+} from "@dbx-tools/shared-mastra";
+import { createUIMessageStream, type UIMessage, type UIMessageChunk } from "ai";
+
+import { askGenieToolId } from "./genie.ts";
+
+const logger = log.logger("mastra/genie-chat");
+
+/** Page size used while walking Genie's paginated list endpoints. */
+const LIST_PAGE_SIZE = 100;
+
+/** Upper bound on conversations surfaced in the thread list. */
+const THREAD_LIMIT = 100;
+const AGENT_CONVERSATION_TYPE = "GENIE_CONVERSATION_TYPE_AGENT";
+const DIRECT_GENIE_ERROR = "Genie could not complete the request.";
+
+/** Suffix that keeps a Genie question's user message id distinct from its answer. */
+const QUESTION_ID_SUFFIX = ":question";
+
+/** One configured Genie space addressed by the direct chat routes. */
+export interface GenieChatSpace {
+  /** Space alias from the route, which also selects the `ask_genie` tool id. */
+  alias: string;
+  /** Genie `space_id`. */
+  spaceId: string;
+}
+
+/**
+ * Compose the visible answer for one terminal Genie message: its narrative text
+ * first, then one embed per query that produced a warehouse statement. A query
+ * Genie chose to visualize embeds `[visualization:<statement_id>]`, which
+ * AppKit UI renders with its native Genie visualization. Any other query embeds
+ * `[data:<statement_id>]`. Agent Mode's Markdown copy of a query result is
+ * dropped because either embed renders the same rows.
+ */
+export function genieAnswerText(message: GenieMessage): string {
+  const attachments = message.attachments ?? [];
+  const visualizedQueries = new Set(
+    attachments.flatMap((attachment) =>
+      attachment.viz?.query_attachment_id ? [attachment.viz.query_attachment_id] : [],
+    ),
+  );
+  const prose: string[] = [];
+  const embeds: string[] = [];
+  for (const attachment of attachments) {
+    const text = attachment.text?.content?.trim();
+    if (text && !isQueryOutputText(attachment.text?.metadata)) prose.push(text);
+    const statementId = attachment.query?.statement_id;
+    if (!statementId) continue;
+    const description = attachment.query?.description?.trim();
+    if (description) embeds.push(description);
+    embeds.push(
+      visualizedQueries.has(attachment.attachment_id ?? "")
+        ? markers.formatMarker("visualization", statementId)
+        : markers.formatMarker("data", statementId),
+    );
+  }
+  const sections = [...prose, ...embeds];
+  if (sections.length === 0 && message.error?.error) sections.push(message.error.error);
+  return sections.join("\n\n");
+}
+
+/** 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";
+}
+
+/** 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!)}`;
+}
+
+/**
+ * Project one persisted Genie message onto the user question and the assistant
+ * answer the chat transcript renders. The assistant owns an `ask_genie` tool
+ * part whose progress is re-derived from the terminal snapshot, so reloaded
+ * history shows the same pill, SQL, and suggestions as the live turn.
+ */
+export function genieMessageToUiMessages(
+  message: GenieMessage,
+  space: GenieChatSpace,
+): UIMessage[] {
+  const messageId = message.message_id || message.id;
+  const toolCallId = messageId;
+  const input = { question: message.content };
+  const failure = genieFailure(message);
+  const progress = [...genieEvent.eventsFromMessage(message, undefined, space.spaceId)].map(
+    (progressEvent) => ({
+      type: wire.GENIE_PROGRESS_PART_TYPE,
+      data: { toolCallId, event: progressEvent },
+    }),
+  );
+  const toolPart = failure
+    ? {
+        type: `tool-${askGenieToolId(space.alias)}`,
+        toolCallId,
+        state: "output-error",
+        input,
+        errorText: failure,
+      }
+    : {
+        type: `tool-${askGenieToolId(space.alias)}`,
+        toolCallId,
+        state: "output-available",
+        input,
+        output: { message },
+      };
+  const answer = genieAnswerText(message);
+  const assistantParts = [
+    toolPart,
+    ...progress,
+    ...(answer ? [{ type: "text", text: answer, state: "done" }] : []),
+  ] as UIMessage["parts"];
+  return [
+    {
+      id: `${messageId}${QUESTION_ID_SUFFIX}`,
+      role: "user",
+      parts: [{ type: "text", text: message.content }],
+    },
+    { id: messageId, role: "assistant", parts: assistantParts },
+  ];
+}
+
+/** List the caller's conversations in one space, newest first. */
+export async function listGenieThreads(
+  client: WorkspaceClient,
+  space: GenieChatSpace,
+  signal?: AbortSignal,
+): Promise<GenieThread[]> {
+  const context = signal ? databricks.toContext(signal) : undefined;
+  const threads: GenieThread[] = [];
+  let pageToken: string | undefined;
+  do {
+    const page = await client.genie.listConversations(
+      {
+        space_id: space.spaceId,
+        page_size: LIST_PAGE_SIZE,
+        ...(pageToken ? { page_token: pageToken } : {}),
+      },
+      context,
+    );
+    for (const conversation of page.conversations ?? []) {
+      if (!object.isRecord(conversation) || conversation.agent_type !== AGENT_CONVERSATION_TYPE) {
+        continue;
+      }
+      threads.push({
+        id: conversation.conversation_id,
+        ...(conversation.title ? { title: conversation.title } : {}),
+        createdAt: new Date(conversation.created_timestamp).toISOString(),
+      });
+    }
+    pageToken = page.next_page_token || undefined;
+  } while (pageToken && threads.length < THREAD_LIMIT);
+  return threads.sort((a, b) => b.createdAt.localeCompare(a.createdAt)).slice(0, THREAD_LIMIT);
+}
+
+/**
+ * Load one conversation's messages as chronological chat transcript messages.
+ * Conversation items come only from the Agent Mode items API, which retains
+ * user messages, reasoning, queries, results, reports, and final answers.
+ */
+export async function listGenieThreadMessages(options: {
+  client: WorkspaceClient;
+  space: GenieChatSpace;
+  conversationId: string;
+  signal?: AbortSignal;
+}): Promise<UIMessage[]> {
+  const { space } = options;
+  const messages = await agentMode.listAgentModeConversation(
+    options.client,
+    space.spaceId,
+    options.conversationId,
+    {
+      ...(options.signal ? { context: options.signal } : {}),
+    },
+  );
+  return messages.flatMap((message) => genieMessageToUiMessages(message, space));
+}
+
+/**
+ * Stream one Genie turn as AI SDK UI chunks shaped like an agent's
+ * `ask_genie` call. The conversation identity is announced through a transient
+ * {@link wire.GENIE_CONVERSATION_PART_TYPE} part as soon as Genie assigns it,
+ * so the browser can adopt the conversation id as its thread id. Agent Mode
+ * receives Genie's native visualization tool; visualization attachments become
+ * AppKit UI embed markers in the answer.
+ */
+export function genieChatStream(options: {
+  client: WorkspaceClient;
+  space: GenieChatSpace;
+  request: GenieChatRequest;
+  signal?: AbortSignal;
+}): ReadableStream<UIMessageChunk> {
+  const { client, space, request, signal } = options;
+  const toolCallId = hash.id();
+  const toolName = askGenieToolId(space.alias);
+  return createUIMessageStream({
+    onError: () => DIRECT_GENIE_ERROR,
+    execute: async ({ writer }) => {
+      const progress = (progressEvent: wire.GenieWriterEvent) =>
+        writer.write({
+          type: wire.GENIE_PROGRESS_PART_TYPE,
+          data: { toolCallId, event: progressEvent },
+        });
+      writer.write({ type: "start" });
+      writer.write({ type: "start-step" });
+      writer.write({
+        type: "tool-input-available",
+        toolCallId,
+        toolName,
+        input: { question: request.question },
+      });
+      const started: StartedEvent = {
+        type: "started",
+        spaceId: space.spaceId,
+        content: request.question,
+      };
+      progress(started);
+
+      let announced: wire.GenieConversationPartData | undefined;
+      let finalMessage: GenieMessage | undefined;
+      try {
+        for await (const chatEvent of chat.genieEventChat(space.spaceId, request.question, {
+          workspaceClient: client,
+          agentMode: true,
+          agentModeFallback: false,
+          enableVisualization: true,
+          ...(request.conversationId ? { conversationId: request.conversationId } : {}),
+          ...(signal ? { context: signal } : {}),
+        })) {
+          const conversationId = chatEvent.conversation_id ?? request.conversationId;
+          const messageId = chatEvent.message_id;
+          // Re-announce only when the identity gains information (the message
+          // id commonly lands one snapshot after the conversation id).
+          if (
+            conversationId &&
+            (announced?.conversationId !== conversationId ||
+              (messageId && announced.messageId !== messageId))
+          ) {
+            announced = { conversationId, ...(messageId ? { messageId } : {}) };
+            writer.write({
+              type: wire.GENIE_CONVERSATION_PART_TYPE,
+              data: announced,
+              transient: true,
+            });
+          }
+          if (chatEvent.type === "result") finalMessage = chatEvent.message;
+          else if (chatEvent.type !== "message") progress(chatEvent);
+        }
+      } catch (error) {
+        logger.warn("turn:error", {
+          spaceId: space.spaceId,
+          error: errorUtils.toError(error),
+        });
+        writer.write({
+          type: "tool-output-error",
+          toolCallId,
+          errorText: DIRECT_GENIE_ERROR,
+        });
+        throw error;
+      }
+      if (!finalMessage) throw new Error("Genie ended the turn without a result");
+
+      const failure = genieFailure(finalMessage);
+      writer.write(
+        failure
+          ? { type: "tool-output-error", toolCallId, errorText: failure }
+          : { type: "tool-output-available", toolCallId, output: { message: finalMessage } },
+      );
+      writer.write({ type: "finish-step" });
+      const answer = genieAnswerText(finalMessage);
+      if (answer) {
+        const textId = hash.id();
+        writer.write({ type: "start-step" });
+        writer.write({ type: "text-start", id: textId });
+        writer.write({ type: "text-delta", id: textId, delta: answer });
+        writer.write({ type: "text-end", id: textId });
+        writer.write({ type: "finish-step" });
+      }
+      writer.write({ type: "finish", finishReason: failure ? "error" : "stop" });
+      logger.debug("turn:done", {
+        spaceId: space.spaceId,
+        conversationId: announced?.conversationId,
+        status: finalMessage.status,
+      });
+    },
+  });
+}
+
+/** Delete one Genie conversation owned by the caller. */
+export async function deleteGenieThread(
+  client: WorkspaceClient,
+  space: GenieChatSpace,
+  conversationId: string,
+): Promise<void> {
+  await client.genie.deleteConversation({
+    space_id: space.spaceId,
+    conversation_id: conversationId,
+  });
+}
+
+/** Delete one Genie message, accepting either its answer or question transcript id. */
+export async function deleteGenieThreadMessage(
+  client: WorkspaceClient,
+  space: GenieChatSpace,
+  conversationId: string,
+  messageId: string,
+): Promise<void> {
+  await client.genie.deleteConversationMessage({
+    space_id: space.spaceId,
+    conversation_id: conversationId,
+    message_id: messageId.endsWith(QUESTION_ID_SUFFIX)
+      ? messageId.slice(0, -QUESTION_ID_SUFFIX.length)
+      : messageId,
+  });
+}

package/src/genie.ts

@@ -352,6 +352,11 @@ function aliasSuffix(alias: string): string {
   return slug ? `_${slug}` : "";
 }
 
+/** Tool id of the `ask_genie` tool registered for one space alias. */
+export function askGenieToolId(alias: string): string {
+  return `ask_genie${aliasSuffix(alias)}`;
+}
+
 /* --------------------------- per-space tools --------------------------- */
 
 /**
@@ -380,7 +385,7 @@ function buildAskGenieTool(opts: {
   agentMode?: boolean;
 }) {
   const { spaceId, alias, hint, agentMode } = opts;
-  const toolId = `ask_genie${aliasSuffix(alias)}`;
+  const toolId = askGenieToolId(alias);
   const hintLine = hint ? ` (${hint})` : "";
   return createTool({
     id: toolId,
@@ -400,14 +405,13 @@ function buildAskGenieTool(opts: {
 
       Returns the final \`GenieMessage\`. Agent Mode includes the
       synthesized answer and can include SQL result values as a
-      Markdown text attachment. Agent Mode has no \`statement_id\`;
-      use \`render_data\` with its inline rows when a chart is needed.
-      A legacy polling result instead carries a \`statement_id\` (at
-      \`message.query_result.statement_id\` or the first
-      attachment's \`query.statement_id\`); call
+      Markdown text attachment. Each executed query attachment
+      carries a \`statement_id\` at \`query.statement_id\` (a legacy
+      polling result may also carry
+      \`message.query_result.statement_id\`); call
       \`get_statement\` with that id only when you need to read
       the underlying values to reason about them. If you just want
-      to display a legacy result, embed a
+      to display a result, embed a
       \`[data:<statement_id>]\` marker in prose. Wire events
       (status, reasoning, SQL, query output, answer text) stream
       automatically while the call is in flight.

package/src/memory.ts

@@ -148,14 +148,14 @@ export function needsLakebase(config: MastraPluginConfig): boolean {
 }
 
 /**
- * Construct a per-agent {@link Memory} factory bound to the supplied
- * service-principal pool.
+ * Construct a per-agent {@link Memory} factory bound to the optional
+ * service-principal Lakebase pool.
  * Caches the shared `PgVector` singleton (built on first need) so each
  * agent build is O(1) after the first.
  */
 export function createMemoryBuilder(
   config: MastraPluginConfig,
-  servicePrincipalPool: Pool,
+  servicePrincipalPool?: Pool,
 ): MemoryBuilder {
   return new MemoryBuilder(config, servicePrincipalPool);
 }
@@ -170,7 +170,7 @@ export class MemoryBuilder {
 
   constructor(
     private readonly config: MastraPluginConfig,
-    private readonly servicePrincipalPool: Pool,
+    private readonly servicePrincipalPool?: Pool,
   ) {}
 
   /**
@@ -196,6 +196,7 @@ export class MemoryBuilder {
         ),
       );
     }
+    if (!this.servicePrincipalPool) return undefined;
     return withSoftStorageInit(
       new PostgresStore({
         id: "mastra-store__instance",
@@ -257,6 +258,7 @@ export class MemoryBuilder {
   private buildStorage(agentId: string, setting: StorageSetting): PostgresStore | undefined {
     if (!setting) return undefined;
     if (typeof setting === "boolean") {
+      if (!this.servicePrincipalPool) return undefined;
       return withSoftStorageInit(
         new PostgresStore({
           id: `mastra-store__${agentId}`,
@@ -294,6 +296,9 @@ export class MemoryBuilder {
 
   private getSharedVector(): PgVector {
     if (!this.sharedVector) {
+      if (!this.servicePrincipalPool) {
+        throw new Error("Lakebase is required for PgVector memory");
+      }
       this.sharedVector = buildSharedPgVector(this.servicePrincipalPool);
     }
     return this.sharedVector;

package/src/observability.ts

@@ -46,7 +46,7 @@ import { projectUtils } from "@dbx-tools/core";
 import { log } from "@dbx-tools/shared-core";
 import { Observability } from "@mastra/observability";
 import { OtelBridge } from "@mastra/otel-bridge";
-import { propagation } from "@opentelemetry/api";
+import { propagation, type TracerProvider } from "@opentelemetry/api";
 
 import { TRACE_REQUEST_CONTEXT_KEYS } from "./config.ts";
 import {
@@ -92,6 +92,8 @@ export interface BuildObservabilityOptions {
    * Supports dot notation for nested values per the Mastra docs.
    */
   requestContextKeys?: readonly string[];
+  /** AppKit-owned tracer provider used by the Mastra OTel bridge. */
+  tracerProvider?: TracerProvider;
 }
 
 /** True when AppKit's global OTel pipeline is configured to export traces. */
@@ -188,7 +190,9 @@ export async function buildObservability(
     configs: {
       serviceName: {
         serviceName,
-        bridge: new OtelBridge(),
+        bridge: new OtelBridge(
+          options?.tracerProvider ? { tracerProvider: options.tracerProvider } : undefined,
+        ),
         requestContextKeys,
       },
     },