@dbx-tools/appkit-mastra 0.9.129 → 0.9.131

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (51) hide show
  1. package/README.md +90 -31
  2. package/index.ts +10 -8
  3. package/lib/index.d.ts +10 -8
  4. package/lib/index.js +8 -7
  5. package/lib/src/agents.d.ts +4 -0
  6. package/lib/src/agents.js +12 -3
  7. package/lib/src/config.d.ts +7 -2
  8. package/lib/src/config.js +6 -3
  9. package/lib/src/defaults.d.ts +9 -0
  10. package/lib/src/defaults.js +29 -1
  11. package/lib/src/genie-chat.d.ts +71 -0
  12. package/lib/src/genie-chat.js +272 -0
  13. package/lib/src/genie.d.ts +2 -0
  14. package/lib/src/genie.js +11 -8
  15. package/lib/src/memory.d.ts +5 -5
  16. package/lib/src/memory.js +10 -3
  17. package/lib/src/observability.d.ts +3 -0
  18. package/lib/src/observability.js +2 -2
  19. package/lib/src/plugin.d.ts +36 -7
  20. package/lib/src/plugin.js +220 -33
  21. package/lib/src/request-context.d.ts +2 -0
  22. package/lib/src/request-context.js +3 -0
  23. package/lib/src/server.d.ts +4 -0
  24. package/lib/src/server.js +40 -2
  25. package/lib/src/statement.d.ts +16 -1
  26. package/lib/src/statement.js +50 -18
  27. package/lib/src/telemetry.d.ts +22 -1
  28. package/lib/src/telemetry.js +174 -18
  29. package/lib/src/workspaces.d.ts +2 -2
  30. package/lib/src/workspaces.js +3 -2
  31. package/lib/tsconfig.tsbuildinfo +1 -1
  32. package/package.json +32 -25
  33. package/src/agents.ts +15 -1
  34. package/src/config.ts +10 -4
  35. package/src/defaults.ts +30 -0
  36. package/src/genie-chat.ts +341 -0
  37. package/src/genie.ts +11 -7
  38. package/src/memory.ts +9 -4
  39. package/src/observability.ts +6 -2
  40. package/src/plugin.ts +262 -31
  41. package/src/request-context.ts +2 -0
  42. package/src/server.ts +35 -0
  43. package/src/statement.ts +71 -19
  44. package/src/telemetry.ts +206 -16
  45. package/src/workspaces.ts +4 -3
  46. package/exports.ts +0 -7
  47. package/lib/exports.d.ts +0 -1
  48. package/lib/exports.js +0 -2
  49. package/lib/src/openapi-tool.d.ts +0 -29
  50. package/lib/src/openapi-tool.js +0 -175
  51. package/src/openapi-tool.ts +0 -245
package/src/agents.ts CHANGED
@@ -53,6 +53,7 @@ import { buildGenieToolkitProvider, resolveGenieSpaces } from "./genie.ts";
53
53
  import { withRequestMemoryScope, type MemoryBuilder } from "./memory.ts";
54
54
  import { buildModel, RESPONSES_PROVIDER_OPTIONS } from "./model.ts";
55
55
  import { stripStaleChartsProcessor } from "./processors.ts";
56
+ import { MASTRA_ORIGIN_KEY } from "./request-context.ts";
56
57
  import { MASTRA_RESOLVED_MODEL_KEY } from "./serving.ts";
57
58
  import { TYPOGRAPHY_RULE } from "./style.ts";
58
59
  import { buildSummarizeTool } from "./summarize.ts";
@@ -419,6 +420,9 @@ export interface BuiltAgents {
419
420
  ambientTools: MastraTools;
420
421
  }
421
422
 
423
+ /** Builds one optional output processor for a registered Mastra agent. */
424
+ export type MastraOutputProcessorFactory = (agentId: string) => OutputProcessor | undefined;
425
+
422
426
  /** Fallback agent id used when `config.agents` is omitted entirely. */
423
427
  export const FALLBACK_AGENT_ID = "default";
424
428
 
@@ -705,6 +709,7 @@ export async function buildAgents(opts: {
705
709
  config: MastraPluginConfig;
706
710
  context: pluginRegistry.PluginContextLike | undefined;
707
711
  memoryBuilder?: MemoryBuilder;
712
+ outputProcessorFactories?: readonly MastraOutputProcessorFactory[];
708
713
  log: log.Logger;
709
714
  /**
710
715
  * Local skill filesystems mounted into every auto-created default workspace.
@@ -720,6 +725,7 @@ export async function buildAgents(opts: {
720
725
  config,
721
726
  context,
722
727
  memoryBuilder,
728
+ outputProcessorFactories = [],
723
729
  log,
724
730
  extraSkillMounts,
725
731
  standardSkills = [],
@@ -824,7 +830,13 @@ export async function buildAgents(opts: {
824
830
  ...(memory ? { memory } : {}),
825
831
  ...(workspace ? { workspace } : {}),
826
832
  inputProcessors: [...inputProcessors, ...workspaceSkillInputProcessors(workspace, config)],
827
- outputProcessors: [toolErrorLoggingProcessor(log)],
833
+ outputProcessors: [
834
+ toolErrorLoggingProcessor(log),
835
+ ...outputProcessorFactories.flatMap((factory) => {
836
+ const processor = factory(id);
837
+ return processor ? [processor] : [];
838
+ }),
839
+ ],
828
840
  });
829
841
  agents[id] =
830
842
  config.backgroundTurns === false
@@ -1271,11 +1283,13 @@ function toolkitEntryToMastraTool(entry: ToolkitEntry, plugin: ContextualToolPro
1271
1283
  | undefined;
1272
1284
  const resourceId = execution?.requestContext?.get(MASTRA_RESOURCE_ID_KEY);
1273
1285
  const model = execution?.requestContext?.get(MASTRA_RESOLVED_MODEL_KEY);
1286
+ const origin = execution?.requestContext?.get(MASTRA_ORIGIN_KEY);
1274
1287
  const toolCallId = execution?.agent?.toolCallId;
1275
1288
  const writer = execution?.writer;
1276
1289
  return plugin.executeAgentTool!(entry.localName, input, execution?.abortSignal, {
1277
1290
  ...(typeof resourceId === "string" ? { resourceId } : {}),
1278
1291
  ...(typeof model === "string" ? { model } : {}),
1292
+ ...(typeof origin === "string" ? { origin } : {}),
1279
1293
  ...(writer && toolCallId
1280
1294
  ? {
1281
1295
  writeProgress: async (event, options) => {
package/src/config.ts CHANGED
@@ -48,6 +48,8 @@ export const MASTRA_USER_KEY = "mastra__user";
48
48
  export const MASTRA_USER_NAME_KEY = "mastra__userName";
49
49
  /** Request-context key containing the authenticated user's email address. */
50
50
  export const MASTRA_USER_EMAIL_KEY = "mastra__userEmail";
51
+ /** Request-context key identifying the Databricks credential used by the turn. */
52
+ export const MASTRA_AUTH_MODE_KEY = "mastra__authMode";
51
53
 
52
54
  /**
53
55
  * `RequestContext` key for the per-HTTP-request id stamped by
@@ -88,6 +90,7 @@ export const TRACE_REQUEST_CONTEXT_KEYS: readonly string[] = [
88
90
  MASTRA_REQUEST_ID_KEY,
89
91
  MASTRA_USER_NAME_KEY,
90
92
  MASTRA_USER_EMAIL_KEY,
93
+ MASTRA_AUTH_MODE_KEY,
91
94
  // Model override key is owned by `serving.ts`; spelled inline here
92
95
  // so this module stays leaf-level (no cycles with `serving.ts`).
93
96
  "mastra__model_override",
@@ -450,7 +453,8 @@ export interface MastraPluginConfig extends BasePluginConfig {
450
453
  * Use the streaming Genie Agent Mode API for `ask_genie`. Defaults to `true`.
451
454
  * Set `false` to force the legacy Conversation API polling driver. A
452
455
  * pre-stream feature-disabled or preview-toggle response falls back
453
- * automatically.
456
+ * automatically. Direct Genie chat always requires Agent Mode and ignores
457
+ * this option.
454
458
  */
455
459
  genieAgentMode?: boolean;
456
460
  /**
@@ -604,6 +608,8 @@ export interface MastraPluginConfig extends BasePluginConfig {
604
608
  * threads, the per-user cache namespace, and trace metadata still key off the
605
609
  * forwarded user, so callers sharing the service principal's data access keep
606
610
  * separate conversations and cannot read each other's threads or charts.
611
+ * Direct Genie chat does not use this option: its conversations and messages
612
+ * always run under AppKit OBO so Genie remains the per-user state owner.
607
613
  *
608
614
  * @example An app any account user can open
609
615
  * ```ts
@@ -658,7 +664,7 @@ export interface MastraPluginConfig extends BasePluginConfig {
658
664
  * ```
659
665
  */
660
666
  remoteSkills?: RemoteSkillsOption;
661
- /** Read-only local skill trees mounted for files and passed as standard agent-level skills. */
667
+ /** Read-only local skill trees mounted for files, searchable by path, and passed as standard agent-level skills. */
662
668
  localSkillMounts?: LocalSkillMount[];
663
669
  }
664
670
 
@@ -768,7 +774,7 @@ export const MASTRA_CONFIG_SCHEMA: ConfigSchema = {
768
774
  genieAgentMode: {
769
775
  type: "boolean",
770
776
  description:
771
- "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.",
777
+ "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.",
772
778
  },
773
779
  genieSpaceCacheTtlMs: {
774
780
  type: "number",
@@ -810,7 +816,7 @@ export const MASTRA_CONFIG_SCHEMA: ConfigSchema = {
810
816
  // runtime while the published schema (and the UI built from it) omits it.
811
817
  enum: [...IDENTITY_MODES],
812
818
  description:
813
- '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.',
819
+ '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.',
814
820
  },
815
821
  apiAccess: {
816
822
  type: "string",
package/src/defaults.ts CHANGED
@@ -82,6 +82,36 @@ export const genieSuggestionDefaults: ExecuteSettings = {
82
82
  },
83
83
  };
84
84
 
85
+ /**
86
+ * `GET /genie/:alias/threads` and its messages: the caller's Genie
87
+ * conversations and their transcripts.
88
+ */
89
+ export const genieConversationReadDefaults: ExecuteSettings = {
90
+ default: {
91
+ // Cache disabled: conversations change on every turn, and the chat refreshes
92
+ // the list right after a turn to pick up the new conversation and title.
93
+ cache: { enabled: false },
94
+ // Retry enabled: listing conversations and messages is an idempotent GET on
95
+ // the chat's load path.
96
+ retry: { enabled: true, attempts: 2 },
97
+ timeout: REST_TIMEOUT_MS,
98
+ },
99
+ };
100
+
101
+ /**
102
+ * `DELETE /genie/:alias/threads/:threadId` and its messages.
103
+ */
104
+ export const genieConversationWriteDefaults: ExecuteSettings = {
105
+ default: {
106
+ // Cache disabled: a write.
107
+ cache: { enabled: false },
108
+ // Retry disabled: a delete that succeeded upstream but timed out locally
109
+ // would surface as a not-found on retry, masking the real outcome.
110
+ retry: { enabled: false },
111
+ timeout: REST_TIMEOUT_MS,
112
+ },
113
+ };
114
+
85
115
  /**
86
116
  * `GET /embed/data/:id`: one Statement Execution result set.
87
117
  */
@@ -0,0 +1,341 @@
1
+ /**
2
+ * Direct Genie chat over the Mastra chat wire.
3
+ *
4
+ * Serves one Genie space as a chat without an agent, model, or Mastra Memory.
5
+ * Genie owns every piece of state: conversations are the chat threads, Genie
6
+ * messages are the transcript, and a turn is one Genie message. Each turn is
7
+ * projected onto the same AI SDK UI chunks an agent turn produces when it
8
+ * calls `ask_genie` (a native tool part, `data-genie-progress` events, and the
9
+ * answer text), so the Mastra chat UI renders pills, thoughts, suggestions,
10
+ * and embeds identically for both.
11
+ *
12
+ * @module
13
+ */
14
+
15
+ import type { WorkspaceClient } from "@databricks/appkit";
16
+ import { databricks } from "@dbx-tools/appkit";
17
+ import { agentMode, chat } from "@dbx-tools/genie";
18
+ import { errorUtils, hash, log, object } from "@dbx-tools/shared-core";
19
+ import { event as genieEvent, genieModel, type GenieMessage } from "@dbx-tools/shared-genie";
20
+ import {
21
+ markers,
22
+ wire,
23
+ type GenieChatRequest,
24
+ type GenieThread,
25
+ type StartedEvent,
26
+ } from "@dbx-tools/shared-mastra";
27
+ import { createUIMessageStream, type UIMessage, type UIMessageChunk } from "ai";
28
+
29
+ import { askGenieToolId } from "./genie.ts";
30
+
31
+ const logger = log.logger("mastra/genie-chat");
32
+
33
+ /** Page size used while walking Genie's paginated list endpoints. */
34
+ const LIST_PAGE_SIZE = 100;
35
+
36
+ /** Upper bound on conversations surfaced in the thread list. */
37
+ const THREAD_LIMIT = 100;
38
+ const AGENT_CONVERSATION_TYPE = "GENIE_CONVERSATION_TYPE_AGENT";
39
+ const DIRECT_GENIE_ERROR = "Genie could not complete the request.";
40
+
41
+ /** Suffix that keeps a Genie question's user message id distinct from its answer. */
42
+ const QUESTION_ID_SUFFIX = ":question";
43
+
44
+ /** One configured Genie space addressed by the direct chat routes. */
45
+ export interface GenieChatSpace {
46
+ /** Space alias from the route, which also selects the `ask_genie` tool id. */
47
+ alias: string;
48
+ /** Genie `space_id`. */
49
+ spaceId: string;
50
+ }
51
+
52
+ /**
53
+ * Compose the visible answer for one terminal Genie message: its narrative text
54
+ * first, then one embed per query that produced a warehouse statement. A query
55
+ * Genie chose to visualize embeds `[visualization:<statement_id>]`, which
56
+ * AppKit UI renders with its native Genie visualization. Any other query embeds
57
+ * `[data:<statement_id>]`. Agent Mode's Markdown copy of a query result is
58
+ * dropped because either embed renders the same rows.
59
+ */
60
+ export function genieAnswerText(message: GenieMessage): string {
61
+ const attachments = message.attachments ?? [];
62
+ const visualizedQueries = new Set(
63
+ attachments.flatMap((attachment) =>
64
+ attachment.viz?.query_attachment_id ? [attachment.viz.query_attachment_id] : [],
65
+ ),
66
+ );
67
+ const prose: string[] = [];
68
+ const embeds: string[] = [];
69
+ for (const attachment of attachments) {
70
+ const text = attachment.text?.content?.trim();
71
+ if (text && !isQueryOutputText(attachment.text?.metadata)) prose.push(text);
72
+ const statementId = attachment.query?.statement_id;
73
+ if (!statementId) continue;
74
+ const description = attachment.query?.description?.trim();
75
+ if (description) embeds.push(description);
76
+ embeds.push(
77
+ visualizedQueries.has(attachment.attachment_id ?? "")
78
+ ? markers.formatMarker("visualization", statementId)
79
+ : markers.formatMarker("data", statementId),
80
+ );
81
+ }
82
+ const sections = [...prose, ...embeds];
83
+ if (sections.length === 0 && message.error?.error) sections.push(message.error.error);
84
+ return sections.join("\n\n");
85
+ }
86
+
87
+ /** Whether a text attachment is Agent Mode's Markdown copy of a query result. */
88
+ function isQueryOutputText(metadata: unknown): boolean {
89
+ return object.isRecord(metadata) && typeof metadata.statement_id === "string";
90
+ }
91
+
92
+ /** Error text for a terminal Genie message that did not complete, or `undefined`. */
93
+ function genieFailure(message: GenieMessage): string | undefined {
94
+ if (message.status === "COMPLETED") return undefined;
95
+ return message.error?.error ?? `Genie ${genieModel.humanizeStatus(message.status!)}`;
96
+ }
97
+
98
+ /**
99
+ * Project one persisted Genie message onto the user question and the assistant
100
+ * answer the chat transcript renders. The assistant owns an `ask_genie` tool
101
+ * part whose progress is re-derived from the terminal snapshot, so reloaded
102
+ * history shows the same pill, SQL, and suggestions as the live turn.
103
+ */
104
+ export function genieMessageToUiMessages(
105
+ message: GenieMessage,
106
+ space: GenieChatSpace,
107
+ ): UIMessage[] {
108
+ const messageId = message.message_id || message.id;
109
+ const toolCallId = messageId;
110
+ const input = { question: message.content };
111
+ const failure = genieFailure(message);
112
+ const progress = [...genieEvent.eventsFromMessage(message, undefined, space.spaceId)].map(
113
+ (progressEvent) => ({
114
+ type: wire.GENIE_PROGRESS_PART_TYPE,
115
+ data: { toolCallId, event: progressEvent },
116
+ }),
117
+ );
118
+ const toolPart = failure
119
+ ? {
120
+ type: `tool-${askGenieToolId(space.alias)}`,
121
+ toolCallId,
122
+ state: "output-error",
123
+ input,
124
+ errorText: failure,
125
+ }
126
+ : {
127
+ type: `tool-${askGenieToolId(space.alias)}`,
128
+ toolCallId,
129
+ state: "output-available",
130
+ input,
131
+ output: { message },
132
+ };
133
+ const answer = genieAnswerText(message);
134
+ const assistantParts = [
135
+ toolPart,
136
+ ...progress,
137
+ ...(answer ? [{ type: "text", text: answer, state: "done" }] : []),
138
+ ] as UIMessage["parts"];
139
+ return [
140
+ {
141
+ id: `${messageId}${QUESTION_ID_SUFFIX}`,
142
+ role: "user",
143
+ parts: [{ type: "text", text: message.content }],
144
+ },
145
+ { id: messageId, role: "assistant", parts: assistantParts },
146
+ ];
147
+ }
148
+
149
+ /** List the caller's conversations in one space, newest first. */
150
+ export async function listGenieThreads(
151
+ client: WorkspaceClient,
152
+ space: GenieChatSpace,
153
+ signal?: AbortSignal,
154
+ ): Promise<GenieThread[]> {
155
+ const context = signal ? databricks.toContext(signal) : undefined;
156
+ const threads: GenieThread[] = [];
157
+ let pageToken: string | undefined;
158
+ do {
159
+ const page = await client.genie.listConversations(
160
+ {
161
+ space_id: space.spaceId,
162
+ page_size: LIST_PAGE_SIZE,
163
+ ...(pageToken ? { page_token: pageToken } : {}),
164
+ },
165
+ context,
166
+ );
167
+ for (const conversation of page.conversations ?? []) {
168
+ if (!object.isRecord(conversation) || conversation.agent_type !== AGENT_CONVERSATION_TYPE) {
169
+ continue;
170
+ }
171
+ threads.push({
172
+ id: conversation.conversation_id,
173
+ ...(conversation.title ? { title: conversation.title } : {}),
174
+ createdAt: new Date(conversation.created_timestamp).toISOString(),
175
+ });
176
+ }
177
+ pageToken = page.next_page_token || undefined;
178
+ } while (pageToken && threads.length < THREAD_LIMIT);
179
+ return threads.sort((a, b) => b.createdAt.localeCompare(a.createdAt)).slice(0, THREAD_LIMIT);
180
+ }
181
+
182
+ /**
183
+ * Load one conversation's messages as chronological chat transcript messages.
184
+ * Conversation items come only from the Agent Mode items API, which retains
185
+ * user messages, reasoning, queries, results, reports, and final answers.
186
+ */
187
+ export async function listGenieThreadMessages(options: {
188
+ client: WorkspaceClient;
189
+ space: GenieChatSpace;
190
+ conversationId: string;
191
+ signal?: AbortSignal;
192
+ }): Promise<UIMessage[]> {
193
+ const { space } = options;
194
+ const messages = await agentMode.listAgentModeConversation(
195
+ options.client,
196
+ space.spaceId,
197
+ options.conversationId,
198
+ {
199
+ ...(options.signal ? { context: options.signal } : {}),
200
+ },
201
+ );
202
+ return messages.flatMap((message) => genieMessageToUiMessages(message, space));
203
+ }
204
+
205
+ /**
206
+ * Stream one Genie turn as AI SDK UI chunks shaped like an agent's
207
+ * `ask_genie` call. The conversation identity is announced through a transient
208
+ * {@link wire.GENIE_CONVERSATION_PART_TYPE} part as soon as Genie assigns it,
209
+ * so the browser can adopt the conversation id as its thread id. Agent Mode
210
+ * receives Genie's native visualization tool; visualization attachments become
211
+ * AppKit UI embed markers in the answer.
212
+ */
213
+ export function genieChatStream(options: {
214
+ client: WorkspaceClient;
215
+ space: GenieChatSpace;
216
+ request: GenieChatRequest;
217
+ signal?: AbortSignal;
218
+ }): ReadableStream<UIMessageChunk> {
219
+ const { client, space, request, signal } = options;
220
+ const toolCallId = hash.id();
221
+ const toolName = askGenieToolId(space.alias);
222
+ return createUIMessageStream({
223
+ onError: () => DIRECT_GENIE_ERROR,
224
+ execute: async ({ writer }) => {
225
+ const progress = (progressEvent: wire.GenieWriterEvent) =>
226
+ writer.write({
227
+ type: wire.GENIE_PROGRESS_PART_TYPE,
228
+ data: { toolCallId, event: progressEvent },
229
+ });
230
+ writer.write({ type: "start" });
231
+ writer.write({ type: "start-step" });
232
+ writer.write({
233
+ type: "tool-input-available",
234
+ toolCallId,
235
+ toolName,
236
+ input: { question: request.question },
237
+ });
238
+ const started: StartedEvent = {
239
+ type: "started",
240
+ spaceId: space.spaceId,
241
+ content: request.question,
242
+ };
243
+ progress(started);
244
+
245
+ let announced: wire.GenieConversationPartData | undefined;
246
+ let finalMessage: GenieMessage | undefined;
247
+ try {
248
+ for await (const chatEvent of chat.genieEventChat(space.spaceId, request.question, {
249
+ workspaceClient: client,
250
+ agentMode: true,
251
+ agentModeFallback: false,
252
+ enableVisualization: true,
253
+ ...(request.conversationId ? { conversationId: request.conversationId } : {}),
254
+ ...(signal ? { context: signal } : {}),
255
+ })) {
256
+ const conversationId = chatEvent.conversation_id ?? request.conversationId;
257
+ const messageId = chatEvent.message_id;
258
+ // Re-announce only when the identity gains information (the message
259
+ // id commonly lands one snapshot after the conversation id).
260
+ if (
261
+ conversationId &&
262
+ (announced?.conversationId !== conversationId ||
263
+ (messageId && announced.messageId !== messageId))
264
+ ) {
265
+ announced = { conversationId, ...(messageId ? { messageId } : {}) };
266
+ writer.write({
267
+ type: wire.GENIE_CONVERSATION_PART_TYPE,
268
+ data: announced,
269
+ transient: true,
270
+ });
271
+ }
272
+ if (chatEvent.type === "result") finalMessage = chatEvent.message;
273
+ else if (chatEvent.type !== "message") progress(chatEvent);
274
+ }
275
+ } catch (error) {
276
+ logger.warn("turn:error", {
277
+ spaceId: space.spaceId,
278
+ error: errorUtils.toError(error),
279
+ });
280
+ writer.write({
281
+ type: "tool-output-error",
282
+ toolCallId,
283
+ errorText: DIRECT_GENIE_ERROR,
284
+ });
285
+ throw error;
286
+ }
287
+ if (!finalMessage) throw new Error("Genie ended the turn without a result");
288
+
289
+ const failure = genieFailure(finalMessage);
290
+ writer.write(
291
+ failure
292
+ ? { type: "tool-output-error", toolCallId, errorText: failure }
293
+ : { type: "tool-output-available", toolCallId, output: { message: finalMessage } },
294
+ );
295
+ writer.write({ type: "finish-step" });
296
+ const answer = genieAnswerText(finalMessage);
297
+ if (answer) {
298
+ const textId = hash.id();
299
+ writer.write({ type: "start-step" });
300
+ writer.write({ type: "text-start", id: textId });
301
+ writer.write({ type: "text-delta", id: textId, delta: answer });
302
+ writer.write({ type: "text-end", id: textId });
303
+ writer.write({ type: "finish-step" });
304
+ }
305
+ writer.write({ type: "finish", finishReason: failure ? "error" : "stop" });
306
+ logger.debug("turn:done", {
307
+ spaceId: space.spaceId,
308
+ conversationId: announced?.conversationId,
309
+ status: finalMessage.status,
310
+ });
311
+ },
312
+ });
313
+ }
314
+
315
+ /** Delete one Genie conversation owned by the caller. */
316
+ export async function deleteGenieThread(
317
+ client: WorkspaceClient,
318
+ space: GenieChatSpace,
319
+ conversationId: string,
320
+ ): Promise<void> {
321
+ await client.genie.deleteConversation({
322
+ space_id: space.spaceId,
323
+ conversation_id: conversationId,
324
+ });
325
+ }
326
+
327
+ /** Delete one Genie message, accepting either its answer or question transcript id. */
328
+ export async function deleteGenieThreadMessage(
329
+ client: WorkspaceClient,
330
+ space: GenieChatSpace,
331
+ conversationId: string,
332
+ messageId: string,
333
+ ): Promise<void> {
334
+ await client.genie.deleteConversationMessage({
335
+ space_id: space.spaceId,
336
+ conversation_id: conversationId,
337
+ message_id: messageId.endsWith(QUESTION_ID_SUFFIX)
338
+ ? messageId.slice(0, -QUESTION_ID_SUFFIX.length)
339
+ : messageId,
340
+ });
341
+ }
package/src/genie.ts CHANGED
@@ -352,6 +352,11 @@ function aliasSuffix(alias: string): string {
352
352
  return slug ? `_${slug}` : "";
353
353
  }
354
354
 
355
+ /** Tool id of the `ask_genie` tool registered for one space alias. */
356
+ export function askGenieToolId(alias: string): string {
357
+ return `ask_genie${aliasSuffix(alias)}`;
358
+ }
359
+
355
360
  /* --------------------------- per-space tools --------------------------- */
356
361
 
357
362
  /**
@@ -380,7 +385,7 @@ function buildAskGenieTool(opts: {
380
385
  agentMode?: boolean;
381
386
  }) {
382
387
  const { spaceId, alias, hint, agentMode } = opts;
383
- const toolId = `ask_genie${aliasSuffix(alias)}`;
388
+ const toolId = askGenieToolId(alias);
384
389
  const hintLine = hint ? ` (${hint})` : "";
385
390
  return createTool({
386
391
  id: toolId,
@@ -400,14 +405,13 @@ function buildAskGenieTool(opts: {
400
405
 
401
406
  Returns the final \`GenieMessage\`. Agent Mode includes the
402
407
  synthesized answer and can include SQL result values as a
403
- Markdown text attachment. Agent Mode has no \`statement_id\`;
404
- use \`render_data\` with its inline rows when a chart is needed.
405
- A legacy polling result instead carries a \`statement_id\` (at
406
- \`message.query_result.statement_id\` or the first
407
- attachment's \`query.statement_id\`); call
408
+ Markdown text attachment. Each executed query attachment
409
+ carries a \`statement_id\` at \`query.statement_id\` (a legacy
410
+ polling result may also carry
411
+ \`message.query_result.statement_id\`); call
408
412
  \`get_statement\` with that id only when you need to read
409
413
  the underlying values to reason about them. If you just want
410
- to display a legacy result, embed a
414
+ to display a result, embed a
411
415
  \`[data:<statement_id>]\` marker in prose. Wire events
412
416
  (status, reasoning, SQL, query output, answer text) stream
413
417
  automatically while the call is in flight.
package/src/memory.ts CHANGED
@@ -148,14 +148,14 @@ export function needsLakebase(config: MastraPluginConfig): boolean {
148
148
  }
149
149
 
150
150
  /**
151
- * Construct a per-agent {@link Memory} factory bound to the supplied
152
- * service-principal pool.
151
+ * Construct a per-agent {@link Memory} factory bound to the optional
152
+ * service-principal Lakebase pool.
153
153
  * Caches the shared `PgVector` singleton (built on first need) so each
154
154
  * agent build is O(1) after the first.
155
155
  */
156
156
  export function createMemoryBuilder(
157
157
  config: MastraPluginConfig,
158
- servicePrincipalPool: Pool,
158
+ servicePrincipalPool?: Pool,
159
159
  ): MemoryBuilder {
160
160
  return new MemoryBuilder(config, servicePrincipalPool);
161
161
  }
@@ -170,7 +170,7 @@ export class MemoryBuilder {
170
170
 
171
171
  constructor(
172
172
  private readonly config: MastraPluginConfig,
173
- private readonly servicePrincipalPool: Pool,
173
+ private readonly servicePrincipalPool?: Pool,
174
174
  ) {}
175
175
 
176
176
  /**
@@ -196,6 +196,7 @@ export class MemoryBuilder {
196
196
  ),
197
197
  );
198
198
  }
199
+ if (!this.servicePrincipalPool) return undefined;
199
200
  return withSoftStorageInit(
200
201
  new PostgresStore({
201
202
  id: "mastra-store__instance",
@@ -257,6 +258,7 @@ export class MemoryBuilder {
257
258
  private buildStorage(agentId: string, setting: StorageSetting): PostgresStore | undefined {
258
259
  if (!setting) return undefined;
259
260
  if (typeof setting === "boolean") {
261
+ if (!this.servicePrincipalPool) return undefined;
260
262
  return withSoftStorageInit(
261
263
  new PostgresStore({
262
264
  id: `mastra-store__${agentId}`,
@@ -294,6 +296,9 @@ export class MemoryBuilder {
294
296
 
295
297
  private getSharedVector(): PgVector {
296
298
  if (!this.sharedVector) {
299
+ if (!this.servicePrincipalPool) {
300
+ throw new Error("Lakebase is required for PgVector memory");
301
+ }
297
302
  this.sharedVector = buildSharedPgVector(this.servicePrincipalPool);
298
303
  }
299
304
  return this.sharedVector;
@@ -46,7 +46,7 @@ import { projectUtils } from "@dbx-tools/core";
46
46
  import { log } from "@dbx-tools/shared-core";
47
47
  import { Observability } from "@mastra/observability";
48
48
  import { OtelBridge } from "@mastra/otel-bridge";
49
- import { propagation } from "@opentelemetry/api";
49
+ import { propagation, type TracerProvider } from "@opentelemetry/api";
50
50
 
51
51
  import { TRACE_REQUEST_CONTEXT_KEYS } from "./config.ts";
52
52
  import {
@@ -92,6 +92,8 @@ export interface BuildObservabilityOptions {
92
92
  * Supports dot notation for nested values per the Mastra docs.
93
93
  */
94
94
  requestContextKeys?: readonly string[];
95
+ /** AppKit-owned tracer provider used by the Mastra OTel bridge. */
96
+ tracerProvider?: TracerProvider;
95
97
  }
96
98
 
97
99
  /** True when AppKit's global OTel pipeline is configured to export traces. */
@@ -188,7 +190,9 @@ export async function buildObservability(
188
190
  configs: {
189
191
  serviceName: {
190
192
  serviceName,
191
- bridge: new OtelBridge(),
193
+ bridge: new OtelBridge(
194
+ options?.tracerProvider ? { tracerProvider: options.tracerProvider } : undefined,
195
+ ),
192
196
  requestContextKeys,
193
197
  },
194
198
  },