@dbx-tools/appkit-mastra 0.1.13 → 0.1.19

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.
package/src/genie.ts CHANGED
@@ -1,751 +1,1276 @@
1
1
  /**
2
- * Mastra tool wrappers around the AppKit `genie` plugin's exports.
2
+ * Genie agent for Mastra.
3
3
  *
4
- * One `sendMessage` tool is registered per configured space alias so
5
- * the LLM picks the space by tool selection (the description bakes the
6
- * alias in). `getConversation` is registered once, taking `alias` as a
7
- * parameter.
4
+ * Each configured Genie space exposes a single Mastra tool to the
5
+ * calling agent (`genie` for the `"default"` alias, `genie_<alias>`
6
+ * otherwise). When invoked, the tool runs end-to-end:
8
7
  *
9
- * All Genie payload types are inferred from the public `genie` factory
10
- * (`genie().plugin` constructor `exports()` return type), so any
11
- * upstream change in `@databricks/appkit` flows in automatically.
8
+ * 1. Pulls the per-request {@link WorkspaceClient} off
9
+ * `ctx.requestContext` (stamped by `MastraServer`) and emits a
10
+ * `started` writer event so the host UI can show progress
11
+ * immediately, before any LLM round-trip.
12
+ * 2. Spins up a per-call inner Mastra `Agent` with three tools:
13
+ * - `ask_genie`: drives one `genieEventChat` turn, fetches
14
+ * the matching statement's rows when the turn ran SQL,
15
+ * and forwards every wire event (status, thinking, sql,
16
+ * rows) through `ctx.writer` for streaming UI updates.
17
+ * - `get_space_description`: cheap title / description /
18
+ * warehouse id lookup for grounding.
19
+ * - `get_space_serialized`: full `GenieSpace` JSON for
20
+ * column-level grounding when the description isn't
21
+ * enough.
22
+ * 3. Runs the inner agent with `structuredOutput` (Mastra's
23
+ * two-pass mode + `jsonPromptInjection`) to coerce the
24
+ * agent's final answer into a tagged
25
+ * `[{type:"text"|"data", ...}]` array. The two-pass design
26
+ * avoids Databricks Model Serving's `response_format` +
27
+ * `tools` collision; prompt injection sidesteps the
28
+ * separate `response_format` + streaming collision in the
29
+ * structuring agent.
30
+ * 4. Charts every `data` item in parallel via
31
+ * {@link runChartPlanner}, maps `text` items to the shared
32
+ * {@link GenieSummaryItem} `string` variant, and returns the
33
+ * hydrated {@link GenieAgentResult}.
12
34
  *
13
- * As Genie streams its long-running events (`FETCHING_METADATA`
14
- * `ASKING_AI` `EXECUTING_QUERY` `COMPLETED`, plus SQL text and
15
- * follow-ups in `message_result.attachments`), the tool forwards a
16
- * normalised {@link GenieProgress} discriminated union out through
17
- * `ctx.writer` so the client can render an incremental loading pill.
18
- * Row payloads from `query_result` are intentionally discarded - the
19
- * LLM never sees rows, and charts come from the separate
20
- * `render_data` tool when the model decides one is useful.
35
+ * The inner agent talks to Genie directly via
36
+ * `@dbx-tools/genie` (`genieEventChat`) and the workspace
37
+ * `statementExecution.getStatement` API. AppKit's stock `genie`
38
+ * plugin is honored only for its `spaces` config so existing
39
+ * AppKit-style wiring keeps working without change.
21
40
  */
22
41
 
23
- import { genie } from "@databricks/appkit";
24
- import { logUtils, stringUtils } from "@dbx-tools/appkit-shared";
42
+ import { CacheManager, genie } from "@databricks/appkit";
43
+ import { ApiError, HttpError, WorkspaceClient } from "@databricks/sdk-experimental";
44
+ import { genieEventChat } from "@dbx-tools/genie";
45
+ import { type GenieMessage } from "@dbx-tools/genie-shared";
46
+ import {
47
+ type ChartEvent,
48
+ type GenieAgentResult,
49
+ type GenieDataset,
50
+ type GenieDatasetData,
51
+ type GenieSummaryItem,
52
+ type MastraGenieErrorEvent,
53
+ type MinimalWriter,
54
+ type StartedEvent,
55
+ type SummaryEvent,
56
+ } from "@dbx-tools/appkit-mastra-shared";
57
+ import {
58
+ apiUtils,
59
+ appkitUtils,
60
+ commonUtils,
61
+ logUtils,
62
+ stringUtils,
63
+ } from "@dbx-tools/shared";
64
+ import { Agent } from "@mastra/core/agent";
25
65
  import type { RequestContext } from "@mastra/core/request-context";
66
+ import { MASTRA_THREAD_ID_KEY } from "@mastra/core/request-context";
26
67
  import { createTool } from "@mastra/core/tools";
27
- import type { ToolStream } from "@mastra/core/tools";
28
68
  import { z } from "zod";
29
69
 
30
- import { emitChartWithPlanning } from "./chart.js";
70
+ import type { MastraTools } from "./agents.js";
71
+ import { runChartPlanner } from "./chart.js";
31
72
  import type { MastraPluginConfig } from "./config.js";
73
+ import { MASTRA_USER_KEY, type User } from "./config.js";
74
+ import { buildModel } from "./model.js";
75
+ import { safeWrite } from "./writer.js";
76
+
77
+ const log = logUtils.logger("mastra/genie");
78
+
79
+ /** Default alias used when a single unnamed Genie space is wired up. */
80
+ export const DEFAULT_GENIE_ALIAS = "default";
32
81
 
33
82
  /**
34
- * Module-level logger tagged `[mastra/genie]`. Uses the shared
35
- * {@link logUtils.logger} so calls below `LOG_LEVEL` are
36
- * discarded for free. Default `LOG_LEVEL` is `info`; flip to
37
- * `debug` to see per-turn timing (`query_result` → planner
38
- * waits `drain:return`).
83
+ * Cap on the inner agent's tool-loop steps. 5 (Mastra default) is
84
+ * tight - one `get_space_description` + one `ask_genie` per
85
+ * sub-question saturates fast. 16 leaves room for ~10 `ask_genie`
86
+ * rounds plus grounding plus the structuring pass (which runs
87
+ * after the loop and is its own single call).
39
88
  */
40
- const log = logUtils.logger("mastra/genie");
89
+ const DEFAULT_MAX_STEPS = 16;
90
+
91
+ /* --------------------------- config types --------------------------- */
92
+
93
+ /** Per-space Genie agent configuration. */
94
+ export interface GenieSpaceConfig {
95
+ /** Genie `space_id`. Required; resolves via `client.genie.getSpace`. */
96
+ spaceId: string;
97
+ /**
98
+ * Optional human-readable description appended to the Genie
99
+ * tool's description so the calling LLM has hints about
100
+ * *what data* this space covers (e.g. "orders, returns,
101
+ * fulfillment"). When omitted, only the space's own
102
+ * `description` (fetched on first use) is shown.
103
+ */
104
+ hint?: string;
105
+ }
106
+
107
+ /** Map of alias -> space config. Accepts either explicit objects or bare space ids. */
108
+ export type GenieSpacesConfig = Record<string, GenieSpaceConfig | string>;
109
+
110
+ /* ------------------------- helpers ------------------------- */
111
+
112
+ /** Best-effort numeric coercion for Genie's all-strings cells. */
113
+ function coerceCell(cell: string | null): unknown {
114
+ if (cell === null) return null;
115
+ if (/^-?\d+(\.\d+)?$/.test(cell)) {
116
+ const n = Number(cell);
117
+ if (Number.isFinite(n)) return n;
118
+ }
119
+ return cell;
120
+ }
121
+
122
+ /**
123
+ * Fetch a single Genie statement's rows via the Statement
124
+ * Execution API and reshape into the shared
125
+ * {@link GenieDatasetData} shape (column array + row records).
126
+ */
127
+ async function fetchStatementData(
128
+ client: WorkspaceClient,
129
+ statementId: string,
130
+ signal?: AbortSignal,
131
+ ): Promise<GenieDatasetData> {
132
+ const ctx = signal ? apiUtils.toContext(signal) : undefined;
133
+ const r = await client.statementExecution.getStatement(
134
+ { statement_id: statementId },
135
+ ctx,
136
+ );
137
+ const columns = (r.manifest?.schema?.columns ?? []).map((c) => c.name ?? "");
138
+ const dataArray = (r.result?.data_array ?? []) as Array<Array<string | null>>;
139
+ const rows = dataArray.map((row) => {
140
+ const obj: Record<string, unknown> = {};
141
+ columns.forEach((col, i) => {
142
+ obj[col] = coerceCell(row[i] ?? null);
143
+ });
144
+ return obj;
145
+ });
146
+ return {
147
+ columns,
148
+ rows,
149
+ rowCount: r.manifest?.total_row_count ?? rows.length,
150
+ };
151
+ }
152
+
153
+ /**
154
+ * Resolve the message's representative `statement_id`. Genie
155
+ * returns one statement per turn in practice; we read the
156
+ * (deprecated-but-singular) `message.query_result.statement_id`
157
+ * first and fall back to the first attachment's
158
+ * `query.statement_id`. Returns `undefined` when the turn had no
159
+ * SQL run (pure prose answer).
160
+ */
161
+ function extractStatementId(message: GenieMessage): string | undefined {
162
+ const top = (message.query_result as { statement_id?: string } | undefined)
163
+ ?.statement_id;
164
+ if (top) return top;
165
+ for (const att of message.attachments ?? []) {
166
+ const id = att.query?.statement_id;
167
+ if (id) return id;
168
+ }
169
+ return undefined;
170
+ }
41
171
 
42
- /** Live AppKit `GeniePlugin` instance. */
43
- export type GeniePluginInstance = InstanceType<ReturnType<typeof genie>["plugin"]>;
172
+ /**
173
+ * Lowercased placeholder strings we reject at the `ask_genie`
174
+ * boundary so the LLM doesn't spend a Genie round-trip on a
175
+ * non-question. Genie politely answers any of these with "Your
176
+ * request '...' does not relate to..." which is pure UI noise.
177
+ * Kept narrow on purpose - real questions sometimes start with
178
+ * one of these tokens, so we only match the FULL trimmed string.
179
+ */
180
+ const PLACEHOLDER_QUESTIONS = new Set([
181
+ "noop",
182
+ "no-op",
183
+ "skip",
184
+ "none",
185
+ "n/a",
186
+ "na",
187
+ "null",
188
+ "undefined",
189
+ "test",
190
+ "placeholder",
191
+ ]);
44
192
 
45
- /** Full `exports()` shape of the AppKit `genie` plugin. */
46
- export type GenieExports = ReturnType<GeniePluginInstance["exports"]>;
193
+ /* ----------------------- conversation state ----------------------- */
47
194
 
48
195
  /**
49
- * Stream event yielded by `genie.exports().sendMessage`. Discriminated
50
- * by `type` (`"message_start" | "status" | "message_result" |
51
- * "query_result" | "error" | "history_info"`).
196
+ * Estimated Genie conversation lifetime in seconds. Databricks
197
+ * publishes no official TTL on the conversation resource itself;
198
+ * community projects (e.g. the open-source Databricks Genie Bot)
199
+ * converge on 4 hours of inactivity as a safe operating window.
200
+ * Treat this as an estimate that gets *extended on every use* by
201
+ * re-setting the cache entry after each successful turn (sliding
202
+ * TTL via re-`set`). When the estimate ends up wrong (conversation
203
+ * deleted, expired upstream, cross-space referenced), the wrapper
204
+ * catches the SDK's `RESOURCE_DOES_NOT_EXIST`/404 and transparently
205
+ * starts a fresh conversation.
52
206
  */
53
- export type GenieStreamEvent =
54
- ReturnType<GenieExports["sendMessage"]> extends AsyncGenerator<infer E> ? E : never;
207
+ const CONVERSATION_TTL_SEC = 4 * 60 * 60;
55
208
 
56
- /** Conversation history returned by `genie.exports().getConversation`. */
57
- export type GenieConversation = Awaited<ReturnType<GenieExports["getConversation"]>>;
209
+ /** Cache namespace prefix so coexisting Mastra caches don't collide. */
210
+ const CONVERSATION_CACHE_NAMESPACE = "mastra:genie:conversation";
58
211
 
59
212
  /**
60
- * Per-dataset metadata surfaced to the LLM. The actual rows are
61
- * dispatched separately as a `kind: "chart"` writer event so the
62
- * model never has the rows in its context (token cost stays flat
63
- * regardless of dataset size). The model uses `chartId` to
64
- * reference the chart inline via the `[[chart:<chartId>]]` marker.
213
+ * Build the per-request {@link RequestContext} key the active
214
+ * Genie `conversation_id` lives under for `spaceId`. Scoped by
215
+ * space so an app calling two Genie spaces in one request keeps
216
+ * each conversation distinct (Genie conversation ids are
217
+ * space-scoped on the wire). The same `RequestContext` instance
218
+ * flows from the outer `genie` tool through to the inner
219
+ * `ask_genie` tool via Mastra, so writes on one side are visible
220
+ * on the other without an explicit shared ref.
65
221
  */
66
- const datasetSchema = z.object({
67
- chartId: z.string().describe(stringUtils.toDescription`
68
- Short id (8 hex chars) for the chart-render slot the host UI
69
- has staged for this dataset. Embed
70
- \`[[chart:<chartId>]]\` on its own line in your reply at the
71
- position you want the chart to appear; the client renders it
72
- inline. Do not paraphrase the dataset's rows in prose - the
73
- chart is the rendering.
74
- `),
75
- title: z.string().optional().describe(stringUtils.toDescription`
76
- Genie's own title for the SQL that produced this dataset.
77
- Useful as a label when you reference the chart in prose.
78
- `),
79
- description: z.string().optional().describe(stringUtils.toDescription`
80
- Genie's prose description of the SQL, if any.
81
- `),
82
- columns: z.array(z.string()).describe(stringUtils.toDescription`
83
- Column names in display order. Use these when describing what
84
- is being charted (e.g. "trend of fill_rate over date").
85
- `),
86
- rowCount: z.number().describe(stringUtils.toDescription`
87
- Total rows in this dataset. Mention only if it adds context
88
- (e.g. "across the last 90 days").
89
- `),
90
- sql: z
91
- .string()
92
- .optional()
93
- .describe(stringUtils.toDescription`
94
- SQL Genie generated and executed. The host UI shows this on
95
- demand; you do not need to repeat it.
96
- `),
97
- });
222
+ const conversationContextKey = (spaceId: string): string =>
223
+ `mastra__genie_conversation__${spaceId}`;
98
224
 
99
225
  /**
100
- * Top-level output schema returned to the LLM from a Genie tool
101
- * call. The `datasets` array is intentionally metadata-only - row
102
- * data rides a writer event the host UI consumes directly and is
103
- * not in the model's context.
226
+ * Read the active Genie `conversation_id` for `spaceId` off the
227
+ * per-request {@link RequestContext}. Returns `undefined` when no
228
+ * conversation has been started yet this request.
104
229
  */
105
- const genieToolOutputSchema = z.object({
106
- conversationId: z
107
- .string()
108
- .optional()
109
- .describe(stringUtils.toDescription`
110
- Pass back on the next call to continue the same Genie thread.
111
- `),
112
- genieAnswer: z
113
- .string()
114
- .optional()
115
- .describe(stringUtils.toDescription`
116
- Genie's natural-language answer to the question. Pass this
117
- through to the user (verbatim, or as the basis of your
118
- reply). Genie may have run multiple SQL queries and tools to
119
- produce this; the full text is the answer.
120
- `),
121
- datasets: z
122
- .array(datasetSchema)
123
- .optional()
124
- .describe(stringUtils.toDescription`
125
- Datasets Genie produced for this turn (one per executed SQL
126
- statement). Each entry is metadata only; the rows are
127
- streamed to the host UI out-of-band. To render any of these
128
- as a chart inline in your reply, embed
129
- \`[[chart:<chartId>]]\` where you want the chart to appear.
130
- Do not paraphrase the rows - the chart is what the user
131
- should see; your prose should add interpretation
132
- (highlights, deltas, anomalies) around the chart.
133
- `),
134
- suggestedFollowUps: z
135
- .array(z.string())
136
- .optional()
137
- .describe(stringUtils.toDescription`
138
- Follow-up question suggestions Genie produced. The host UI
139
- renders these as clickable buttons; you do not need to list
140
- them in your reply.
141
- `),
142
- error: z
143
- .string()
144
- .optional()
145
- .describe(stringUtils.toDescription`
146
- Genie-side error message if the request failed.
147
- `),
148
- });
230
+ function readContextConversationId(
231
+ requestContext: RequestContext,
232
+ spaceId: string,
233
+ ): string | undefined {
234
+ return requestContext.get(conversationContextKey(spaceId)) as string | undefined;
235
+ }
236
+
237
+ /**
238
+ * Write the active Genie `conversation_id` for `spaceId` onto the
239
+ * per-request {@link RequestContext}. Subsequent `ask_genie` calls
240
+ * in this request will reuse it; the wrapper's tail logic also
241
+ * reads it back out for the {@link GenieAgentResult}.
242
+ */
243
+ function writeContextConversationId(
244
+ requestContext: RequestContext,
245
+ spaceId: string,
246
+ conversationId: string | undefined,
247
+ ): void {
248
+ requestContext.set(conversationContextKey(spaceId), conversationId);
249
+ }
149
250
 
150
- type DrainResult = z.infer<typeof genieToolOutputSchema>;
251
+ /* ------------------------- chart inventory ------------------------- */
151
252
 
152
253
  /**
153
- * Normalised progress event surfaced to the UI as a Mastra
154
- * `tool-output` chunk. Loading pill events (`started`, `status`,
155
- * `sql`, `suggested`, `error`) are pure UI metadata and never reach
156
- * the LLM.
254
+ * Per-request {@link RequestContext} key the resolved chart
255
+ * inventory lives under. Keyed by `chartId`, the inventory is a
256
+ * `Map<string, ChartEvent>` carrying the full Echarts spec for
257
+ * every chart minted on this request - the same payload that
258
+ * goes out on the writer stream, kept in-process so output
259
+ * processors and downstream tools can resolve `[[chart:<id>]]`
260
+ * markers without re-running the planner or pulling from the
261
+ * writer stream.
157
262
  *
158
- * The `chart` variant is the wire shape emitted by
159
- * {@link emitChartWithPlanning} (used by both this Genie
160
- * draining loop and the system-level `render_data` tool). All
161
- * fields except `chartId` are optional because two events per
162
- * chartId arrive on the wire: the first carries the rows
163
- * (`title` + `description?` + `data`); the second, on planner
164
- * success, carries just the resolved Echarts spec (`option`).
165
- * The host UI's `<ChartSlot>` merges them by `chartId`.
263
+ * Shared across all Genie spaces because chart ids are minted
264
+ * via `commonUtils.shortId()` and are unique within a single
265
+ * request regardless of which space produced them.
166
266
  */
167
- export type GenieProgress =
168
- | { kind: "started"; conversationId: string; messageId: string; spaceId: string }
169
- | { kind: "status"; status: string; label: string }
170
- | {
171
- kind: "sql";
172
- sql: string;
173
- title?: string;
174
- description?: string;
175
- statementId?: string;
176
- }
177
- | {
178
- kind: "chart";
179
- chartId: string;
180
- title?: string;
181
- description?: string;
182
- data?: Array<Record<string, unknown>>;
183
- option?: Record<string, unknown>;
184
- }
185
- | { kind: "text"; content: string }
186
- | { kind: "suggested"; questions: string[] }
187
- | { kind: "error"; error: string };
188
-
189
- const sendMessageSchema = z.object({
190
- content: z.string().describe(stringUtils.toDescription`
191
- Natural-language question to send to the Genie space.
192
- `),
193
- conversationId: z
194
- .string()
195
- .optional()
196
- .describe(stringUtils.toDescription`
197
- Optional Genie conversation id to continue an earlier thread.
198
- Omit on the first call; pass the id returned in the previous
199
- result's \`conversationId\` to follow up.
200
- `),
201
- });
267
+ const CHART_INVENTORY_CONTEXT_KEY = "mastra__genie_chart_inventory__";
202
268
 
203
- const getConversationSchema = z.object({
204
- alias: z.string().describe(stringUtils.toDescription`
205
- Alias of the Genie space the conversation belongs to (matches
206
- the key in the genie plugin's \`spaces\` config).
207
- `),
208
- conversationId: z.string().describe(stringUtils.toDescription`
209
- Genie conversation id whose history to fetch.
210
- `),
211
- });
269
+ /**
270
+ * Get the chart inventory map for this request, creating it on
271
+ * first access. Subsequent reads return the same map so callers
272
+ * mutate in place. The map is request-scoped (collected with the
273
+ * `RequestContext` at end of request), so there's no per-process
274
+ * leak.
275
+ */
276
+ export function chartInventoryFromContext(
277
+ requestContext: RequestContext,
278
+ ): Map<string, ChartEvent> {
279
+ const existing = requestContext.get(CHART_INVENTORY_CONTEXT_KEY);
280
+ if (existing instanceof Map) {
281
+ return existing as Map<string, ChartEvent>;
282
+ }
283
+ const fresh = new Map<string, ChartEvent>();
284
+ requestContext.set(CHART_INVENTORY_CONTEXT_KEY, fresh);
285
+ return fresh;
286
+ }
212
287
 
213
- /** Per-attachment shape returned inside a stored Genie message. */
214
- const genieAttachmentSchema = z.object({
215
- attachmentId: z.string().optional().describe(stringUtils.toDescription`
216
- Genie attachment id; internal bookkeeping.
217
- `),
218
- query: z
219
- .object({
220
- title: z.string().optional().describe(stringUtils.toDescription`
221
- Genie's title for the SQL, if any.
222
- `),
223
- description: z.string().optional().describe(stringUtils.toDescription`
224
- Genie's prose description of the SQL, if any.
225
- `),
226
- query: z.string().optional().describe(stringUtils.toDescription`
227
- SQL Genie generated and executed.
228
- `),
229
- statementId: z.string().optional().describe(stringUtils.toDescription`
230
- Statement-execution id; internal bookkeeping.
231
- `),
232
- })
233
- .optional()
234
- .describe(stringUtils.toDescription`
235
- SQL Genie attached to this message, if it ran any.
236
- `),
237
- text: z
238
- .object({
239
- content: z.string().optional().describe(stringUtils.toDescription`
240
- Genie's natural-language answer text for this attachment.
241
- `),
242
- })
243
- .optional()
244
- .describe(stringUtils.toDescription`
245
- Per-attachment text content (independent of the message-level
246
- \`content\` field).
247
- `),
248
- suggestedQuestions: z
249
- .array(z.string())
250
- .optional()
251
- .describe(stringUtils.toDescription`
252
- Follow-up question suggestions Genie generated for this turn.
253
- `),
254
- });
288
+ /**
289
+ * Stash a resolved chart on the request-scoped inventory so any
290
+ * subsequent code in this request (output processors validating
291
+ * `[[chart:<id>]]` markers, follow-up tools that want to chart
292
+ * the same dataset differently, etc.) can look it up by id.
293
+ * No-op when `requestContext` is missing.
294
+ */
295
+ function recordChartInContext(
296
+ requestContext: RequestContext | undefined,
297
+ chart: ChartEvent,
298
+ ): void {
299
+ if (!requestContext) return;
300
+ chartInventoryFromContext(requestContext).set(chart.chartId, chart);
301
+ }
255
302
 
256
- /** Single message inside a Genie conversation history page. */
257
- const genieMessageSchema = z.object({
258
- messageId: z.string().describe(stringUtils.toDescription`
259
- Genie message id; internal bookkeeping.
260
- `),
261
- conversationId: z.string().describe(stringUtils.toDescription`
262
- Conversation id this message belongs to.
263
- `),
264
- spaceId: z.string().describe(stringUtils.toDescription`
265
- Genie space id this message belongs to.
266
- `),
267
- status: z.string().describe(stringUtils.toDescription`
268
- Genie message status (\`COMPLETED\`, \`FAILED\`, etc.).
269
- `),
270
- content: z.string().describe(stringUtils.toDescription`
271
- Outer message-level natural-language content Genie wrote.
272
- `),
273
- attachments: z
274
- .array(genieAttachmentSchema)
275
- .optional()
276
- .describe(stringUtils.toDescription`
277
- Attachments (SQL queries, text blocks, suggested follow-ups)
278
- Genie produced for this message.
279
- `),
280
- error: z.string().optional().describe(stringUtils.toDescription`
281
- Genie-side error attached to this message, if any.
282
- `),
283
- });
303
+ /**
304
+ * `userKey` for `CacheManager.getOrExecute` / `generateKey`. Genie
305
+ * conversations are scoped to a single user + space + thread, and
306
+ * `threadId` is already user-scoped (Mastra mints threads per
307
+ * `resourceId`), so a constant user key here is safe and keeps the
308
+ * cache key short.
309
+ */
310
+ const CONVERSATION_USER_KEY = "mastra-genie";
284
311
 
285
312
  /**
286
- * Output schema for the \`genie_get_conversation\` tool. Mirrors
287
- * AppKit's \`GenieConversationHistoryResponse\` so the model gets a
288
- * clear, typed view of prior messages instead of an opaque blob.
313
+ * Build the canonical cache key for a `(spaceId, threadId)` pair.
314
+ * Returns `undefined` when `threadId` is missing - callers should
315
+ * skip caching entirely in that case (no Mastra memory wired up).
289
316
  */
290
- const genieGetConversationOutputSchema = z.object({
291
- conversationId: z.string().describe(stringUtils.toDescription`
292
- Conversation id you fetched.
293
- `),
294
- spaceId: z.string().describe(stringUtils.toDescription`
295
- Genie space the conversation belongs to.
296
- `),
297
- messages: z.array(genieMessageSchema).describe(stringUtils.toDescription`
298
- Messages in the conversation, oldest to newest. Each
299
- \`message.content\` is Genie's natural-language answer for
300
- that turn; attachments carry the SQL and follow-ups Genie
301
- produced.
302
- `),
303
- });
317
+ async function conversationCacheKey(
318
+ spaceId: string,
319
+ threadId: string | undefined,
320
+ ): Promise<string | undefined> {
321
+ if (!threadId) return undefined;
322
+ return (await CacheManager.getInstance()).generateKey(
323
+ [CONVERSATION_CACHE_NAMESPACE, spaceId, threadId],
324
+ CONVERSATION_USER_KEY,
325
+ );
326
+ }
304
327
 
305
328
  /**
306
- * Default tool name for a wired Genie alias. The well-known `default`
307
- * alias collapses to `genie`; everything else gets a `genie_` prefix so
308
- * multiple spaces stay disambiguated when an agent has more than one
309
- * wired. Matches the `genie` / `genie_<alias>` naming used elsewhere in
310
- * dbx-tools AppKit demos.
329
+ * Read the cached Genie conversation id for `(spaceId, threadId)`.
330
+ * Returns `undefined` on miss, on expiry, or when the cache layer
331
+ * is unhealthy - never throws. The TTL is renewed via re-`set`
332
+ * after each successful turn (see {@link saveCachedConversationId}).
311
333
  */
312
- export function defaultGenieToolName(alias: string): string {
313
- if (alias === "default") return "genie";
314
- return stringUtils.toIdentifierWithOptions({ distinct: true }, "genie", alias);
334
+ async function readCachedConversationId(
335
+ cacheKey: string | undefined,
336
+ ): Promise<string | undefined> {
337
+ if (!cacheKey) return undefined;
338
+ try {
339
+ const v = await CacheManager.getInstanceSync().get<string>(cacheKey);
340
+ return v ?? undefined;
341
+ } catch (err) {
342
+ log.warn("conversation-cache:read-error", {
343
+ error: commonUtils.errorMessage(err),
344
+ });
345
+ return undefined;
346
+ }
315
347
  }
316
348
 
317
349
  /**
318
- * Build one `sendMessage` tool per configured Genie alias plus a single
319
- * `getConversation` tool. Returns a record keyed by tool id, ready to
320
- * spread into an `Agent`'s `tools` map.
321
- *
322
- * `config` must be the active plugin config; Genie's
323
- * `query_result` events are routed through
324
- * {@link emitChartWithPlanning} which uses it to resolve the
325
- * chart-planner's model.
350
+ * Persist the active conversation id under `cacheKey`, refreshing
351
+ * its TTL. Idempotent; no-op when `cacheKey` or `conversationId`
352
+ * is missing. Re-setting the same key acts as a sliding TTL: every
353
+ * turn that uses the conversation extends the window by another
354
+ * {@link CONVERSATION_TTL_SEC} seconds.
326
355
  */
327
- export function buildGenieTools(opts: {
328
- aliases: string[];
329
- exports: GenieExports;
330
- config: MastraPluginConfig;
331
- signal?: AbortSignal;
332
- }): Record<string, ReturnType<typeof createTool>> {
333
- const tools: Record<string, ReturnType<typeof createTool>> = {};
356
+ async function saveCachedConversationId(
357
+ cacheKey: string | undefined,
358
+ conversationId: string | undefined,
359
+ ): Promise<void> {
360
+ if (!cacheKey || !conversationId) return;
361
+ try {
362
+ await CacheManager.getInstanceSync().set(cacheKey, conversationId, {
363
+ ttl: CONVERSATION_TTL_SEC,
364
+ });
365
+ } catch (err) {
366
+ log.warn("conversation-cache:write-error", {
367
+ error: commonUtils.errorMessage(err),
368
+ });
369
+ }
370
+ }
334
371
 
335
- for (const alias of opts.aliases) {
336
- const id = defaultGenieToolName(alias);
337
- tools[id] = createTool({
338
- id,
339
- description: stringUtils.toDescription`
340
- Ask the Databricks Genie space "${alias}" a single
341
- natural-language question. Genie translates it to SQL,
342
- runs it, and returns \`genieAnswer\` (prose) plus
343
- \`datasets[]\` (one entry per executed query, each with
344
- a short \`chartId\`). Embed \`[[chart:<chartId>]]\` on
345
- its own line at the position you want that data rendered
346
- as an inline chart. Add interpretation around the chart
347
- (deltas, anomalies, takeaways); do not paraphrase row
348
- values.
349
-
350
- Issue ONE focused question per user turn. Prefer
351
- aggregated queries over raw-row queries for time-series
352
- and distributions.
353
- `,
354
- inputSchema: sendMessageSchema,
355
- outputSchema: genieToolOutputSchema,
356
- execute: async ({ content, conversationId }, ctx) => {
357
- const stream = opts.exports.sendMessage(alias, content, conversationId, {
358
- signal: opts.signal,
359
- });
360
- const requestContext = (ctx as { requestContext?: RequestContext } | undefined)
361
- ?.requestContext;
362
- return drainGenieStream(stream, ctx.writer, {
363
- config: opts.config,
364
- ...(requestContext ? { requestContext } : {}),
365
- });
366
- },
372
+ /** Force-evict a cached conversation id. Used on the stale-id recovery path. */
373
+ async function evictCachedConversationId(cacheKey: string | undefined): Promise<void> {
374
+ if (!cacheKey) return;
375
+ try {
376
+ await CacheManager.getInstanceSync().delete(cacheKey);
377
+ } catch (err) {
378
+ log.warn("conversation-cache:delete-error", {
379
+ error: commonUtils.errorMessage(err),
367
380
  });
368
381
  }
382
+ }
383
+
384
+ /**
385
+ * True when `err` is the SDK error Genie returns for a
386
+ * conversation id that no longer exists (deleted, expired upstream,
387
+ * or referenced from the wrong space). Matches the typed
388
+ * {@link ApiError} 404 / `RESOURCE_DOES_NOT_EXIST` shape first, then
389
+ * falls back to the lower-level {@link HttpError} 404, then to a
390
+ * loose message sniff for SDK shapes we haven't catalogued.
391
+ */
392
+ function isConversationGoneError(err: unknown): boolean {
393
+ if (err instanceof ApiError) {
394
+ if (err.statusCode === 404) return true;
395
+ if (err.errorCode === "RESOURCE_DOES_NOT_EXIST") return true;
396
+ }
397
+ if (err instanceof HttpError && err.code === 404) return true;
398
+ if (err instanceof Error && /does not exist/i.test(err.message)) return true;
399
+ return false;
400
+ }
401
+
402
+ /* --------------------------- inner tools --------------------------- */
403
+
404
+ /**
405
+ * One entry in {@link InnerToolDeps.resultSets}: the rows for a
406
+ * Genie statement plus the Genie `message_id` of the `ask_genie`
407
+ * turn that produced it. Tracking `messageId` here lets the
408
+ * outer chart loop stamp the chart event (and any chart-error
409
+ * writer event) with the same `messageId` the rest of that
410
+ * ask's wire events carry, so the host UI groups the chart into
411
+ * the same `message_id` pill bucket without a separate lookup.
412
+ */
413
+ interface StatementEntry {
414
+ data: GenieDatasetData;
415
+ messageId: string;
416
+ }
417
+
418
+ /**
419
+ * Per-call mutable state shared by the inner agent's three tools.
420
+ * `resultSets` lets the wrapper pull rows by `statementId` after
421
+ * the agent finishes, so the chart-planner doesn't re-fetch. The
422
+ * active Genie `conversation_id` lives on `RequestContext` (read
423
+ * via {@link readContextConversationId} on the inner tool's `ctx`)
424
+ * rather than a shared ref - the same `RequestContext` instance
425
+ * threads from `agent.generate({requestContext})` through to every
426
+ * tool invocation, so writes propagate to subsequent `ask_genie`
427
+ * calls without an extra object. `cacheKey` is the
428
+ * {@link CacheManager} key for cross-request persistence (`undefined`
429
+ * when `threadId` isn't available and caching is disabled).
430
+ */
431
+ interface InnerToolDeps {
432
+ spaceId: string;
433
+ client: WorkspaceClient;
434
+ writer?: MinimalWriter;
435
+ signal?: AbortSignal;
436
+ resultSets: Map<string, StatementEntry>;
437
+ cacheKey?: string;
438
+ }
369
439
 
370
- tools.genie_get_conversation = createTool({
371
- id: "genie_get_conversation",
440
+ function buildAskGenieTool(deps: InnerToolDeps) {
441
+ const { spaceId, client, writer, signal, resultSets, cacheKey } = deps;
442
+ return createTool({
443
+ id: "ask_genie",
372
444
  description: stringUtils.toDescription`
373
- Fetch the full message history of a prior Genie conversation
374
- by id. Use when the user references an earlier Genie thread
375
- by id, or to inspect attachments / SQL from previous turns.
445
+ Send ONE focused natural-language question to the Genie
446
+ space and wait for the turn to complete. Returns the final
447
+ \`GenieMessage\` plus, when the turn ran SQL, the rows of
448
+ the resulting query as \`query_result_data\`. The
449
+ \`statement_id\` you reference in your final \`data\`
450
+ blocks lives at \`message.query_result.statement_id\` (or
451
+ the first attachment's \`query.statement_id\`). Wire
452
+ events (status, thinking, sql) stream to the user
453
+ automatically. Call multiple times to gather different
454
+ angles before composing the final response.
376
455
  `,
377
- inputSchema: getConversationSchema,
378
- outputSchema: genieGetConversationOutputSchema,
379
- execute: async ({ alias, conversationId }) => {
380
- return opts.exports.getConversation(alias, conversationId, opts.signal);
456
+ inputSchema: z.object({
457
+ question: z.string().min(1, "question is required"),
458
+ }),
459
+ outputSchema: z.object({
460
+ message: z.custom<GenieMessage>(),
461
+ query_result_data: z.custom<GenieDatasetData>().optional(),
462
+ }),
463
+ execute: async ({ question }, ctxRaw) => {
464
+ const ctx = ctxRaw as { requestContext?: RequestContext } | undefined;
465
+ const requestContext = ctx?.requestContext;
466
+ if (!requestContext) {
467
+ // Mastra always passes a `RequestContext` to tools when the
468
+ // parent agent received one. The outer Genie tool insists on
469
+ // it (it sources the user from there), so this only fires
470
+ // if a misconfigured caller invokes `ask_genie` directly.
471
+ throw new Error(
472
+ "ask_genie: missing requestContext (parent agent must propagate it)",
473
+ );
474
+ }
475
+
476
+ // Bounce placeholder / no-op questions BEFORE spending a Genie
477
+ // round-trip on them. The structuring pass occasionally pads
478
+ // out the tool loop with a fake `ask_genie("noop")` call,
479
+ // which Genie answers with "Your request 'noop' does not
480
+ // relate to..." - useless noise that shows up in the UI and
481
+ // eats one of the workspace's 5 questions/minute. Returning
482
+ // a clear error here surfaces the issue to the agent loop so
483
+ // the model corrects course instead of wasting a turn.
484
+ const trimmed = question.trim();
485
+ if (trimmed.length === 0 || PLACEHOLDER_QUESTIONS.has(trimmed.toLowerCase())) {
486
+ throw new Error(
487
+ `ask_genie: refusing placeholder question "${question}" - ` +
488
+ `call ask_genie only with a real natural-language question, ` +
489
+ `or skip the call entirely`,
490
+ );
491
+ }
492
+
493
+ // Single turn of `genieEventChat`. Hoisted into a closure so
494
+ // we can re-run it after evicting a stale `conversation_id`
495
+ // without duplicating the event-loop body.
496
+ const runTurn = async (): Promise<GenieMessage> => {
497
+ const seedConversationId = readContextConversationId(requestContext, spaceId);
498
+ let finalMessage: GenieMessage | undefined;
499
+ for await (const event of genieEventChat(spaceId, question, {
500
+ workspaceClient: client,
501
+ ...(seedConversationId ? { conversationId: seedConversationId } : {}),
502
+ ...(signal ? { context: signal } : {}),
503
+ })) {
504
+ await safeWrite(log, writer, event);
505
+ // Wire events come in two flavors: the lifecycle `message`
506
+ // event embeds the raw `GenieMessage` (read its
507
+ // `conversation_id`), and the rest carry a flat
508
+ // `conversation_id` field at the top level. The terminal
509
+ // `result` event also carries the final `GenieMessage`
510
+ // inline so we can capture the snapshot without re-reading
511
+ // a buffered `message` event.
512
+ const eventConversationId =
513
+ event.type === "message"
514
+ ? event.message.conversation_id
515
+ : event.conversation_id;
516
+ if (eventConversationId) {
517
+ writeContextConversationId(requestContext, spaceId, eventConversationId);
518
+ }
519
+ if (event.type === "result") {
520
+ finalMessage = event.message;
521
+ }
522
+ }
523
+ if (!finalMessage) {
524
+ throw new Error("Genie turn ended without a result event");
525
+ }
526
+ return finalMessage;
527
+ };
528
+
529
+ let finalMessage: GenieMessage;
530
+ try {
531
+ finalMessage = await runTurn();
532
+ } catch (err) {
533
+ // The seeded `conversation_id` was rejected by Genie - most
534
+ // commonly because it was deleted upstream, expired past
535
+ // Databricks' (undocumented) lifetime, or was minted in a
536
+ // different space. Drop both the cached id AND the
537
+ // per-request value so the retry calls `startConversation`,
538
+ // and try once more. Only retry when we *had* a seeded id -
539
+ // a fresh call that 404s shouldn't loop.
540
+ const seeded = readContextConversationId(requestContext, spaceId);
541
+ if (seeded && isConversationGoneError(err)) {
542
+ log.warn("conversation-cache:stale, resetting", {
543
+ spaceId,
544
+ conversationId: seeded,
545
+ error: commonUtils.errorMessage(err),
546
+ });
547
+ await evictCachedConversationId(cacheKey);
548
+ writeContextConversationId(requestContext, spaceId, undefined);
549
+ finalMessage = await runTurn();
550
+ } else {
551
+ throw err;
552
+ }
553
+ }
554
+
555
+ // Refresh the cache entry on every successful turn. Re-setting
556
+ // the same key both persists newly-minted ids (cache miss path)
557
+ // and extends the TTL on active conversations (sliding window).
558
+ await saveCachedConversationId(
559
+ cacheKey,
560
+ readContextConversationId(requestContext, spaceId),
561
+ );
562
+
563
+ const statementId = extractStatementId(finalMessage);
564
+ let queryResultData: GenieDatasetData | undefined;
565
+ if (statementId) {
566
+ const data = await fetchStatementData(client, statementId, signal);
567
+ if (data.rowCount > 0) {
568
+ queryResultData = data;
569
+ // Stash with this ask's `message_id` so the outer chart
570
+ // loop can stamp downstream `chart` events with the
571
+ // same id the wire events carry - keeps the chart in
572
+ // the same `message_id` pill bucket on the host UI.
573
+ resultSets.set(statementId, {
574
+ data,
575
+ messageId: finalMessage.message_id,
576
+ });
577
+ }
578
+ }
579
+ return {
580
+ message: finalMessage,
581
+ ...(queryResultData ? { query_result_data: queryResultData } : {}),
582
+ };
381
583
  },
382
584
  });
585
+ }
383
586
 
384
- return tools;
587
+ function buildSpaceDescriptionTool(deps: {
588
+ spaceId: string;
589
+ client: WorkspaceClient;
590
+ signal?: AbortSignal;
591
+ }) {
592
+ const { spaceId, client, signal } = deps;
593
+ return createTool({
594
+ id: "get_space_description",
595
+ description: stringUtils.toDescription`
596
+ Return the Genie space's title, description, and warehouse id.
597
+ Cheap. Call once at the start of a turn to ground yourself
598
+ in what data the space covers.
599
+ `,
600
+ inputSchema: z.object({}),
601
+ outputSchema: z.object({
602
+ spaceId: z.string(),
603
+ title: z.string().optional(),
604
+ description: z.string().optional(),
605
+ warehouseId: z.string().optional(),
606
+ }),
607
+ execute: async () => {
608
+ const ctx = signal ? apiUtils.toContext(signal) : undefined;
609
+ const space = await client.genie.getSpace({ space_id: spaceId }, ctx);
610
+ return {
611
+ spaceId,
612
+ ...(space.title ? { title: space.title } : {}),
613
+ ...(space.description ? { description: space.description } : {}),
614
+ ...(space.warehouse_id ? { warehouseId: space.warehouse_id } : {}),
615
+ };
616
+ },
617
+ });
385
618
  }
386
619
 
387
- /** Inputs to {@link drainGenieStream}. */
388
- interface DrainGenieStreamOptions {
389
- config: MastraPluginConfig;
390
- requestContext?: RequestContext;
620
+ function buildSpaceSerializedTool(deps: {
621
+ spaceId: string;
622
+ client: WorkspaceClient;
623
+ signal?: AbortSignal;
624
+ }) {
625
+ const { spaceId, client, signal } = deps;
626
+ return createTool({
627
+ id: "get_space_serialized",
628
+ description: stringUtils.toDescription`
629
+ Return the full \`GenieSpace\` JSON for this space. Use only
630
+ when you need exact column / table identifiers
631
+ \`get_space_description\` doesn't expose. Larger payload, so
632
+ prefer the description tool when it's enough.
633
+ `,
634
+ inputSchema: z.object({}),
635
+ outputSchema: z.object({ space: z.unknown() }),
636
+ execute: async () => {
637
+ const ctx = signal ? apiUtils.toContext(signal) : undefined;
638
+ const space = await client.genie.getSpace({ space_id: spaceId }, ctx);
639
+ return { space };
640
+ },
641
+ });
391
642
  }
392
643
 
644
+ /* --------------------------- inner agent --------------------------- */
645
+
646
+ const AGENT_INSTRUCTIONS = stringUtils.toDescription`
647
+ You orchestrate a Databricks Genie space. For every user
648
+ question:
649
+
650
+ 1. Optionally call \`get_space_description\` to ground; reach
651
+ for \`get_space_serialized\` only when you need exact
652
+ column / table names the description doesn't expose.
653
+ 2. Decompose the question into focused sub-questions (one per
654
+ distinct metric / dimension / time window) and call
655
+ \`ask_genie\` once per sub-question. Two to six calls is
656
+ typical for a non-trivial question; one call is fine when
657
+ the question is genuinely atomic.
658
+ 3. Each \`ask_genie\` call returns the terminal
659
+ \`GenieMessage\`. When the turn ran SQL it also returns
660
+ \`query_result_data\` - the actual rows. The matching
661
+ \`statement_id\` is on
662
+ \`message.query_result.statement_id\` (or the first
663
+ attachment's \`query.statement_id\`). You will reference
664
+ that exact id in your final \`data\` blocks.
665
+ 4. Produce a final structured summary as an ordered array
666
+ interleaving \`text\` paragraphs with \`data\` blocks.
667
+ INTERLEAVE: prose first, then the \`data\` block it
668
+ interprets, then the next prose / data pair. Never dump
669
+ all data at the end.
670
+ 5. For every \`data\` block, supply the exact
671
+ \`statement_id\` you saw on the \`ask_genie\` response. A
672
+ short \`description\` ("compare quarterly revenue across
673
+ regions", "highlight the steep drop after position 5")
674
+ biases the chart-planner's choice of visual. Do NOT pick
675
+ chart types or axis labels - the host wraps each \`data\`
676
+ block in a chart automatically.
677
+ 6. Each \`data\` block should be followed by a short
678
+ \`text\` interpretation (deltas, anomalies, takeaways).
679
+ Don't paraphrase numbers the visualization will already
680
+ show. Skip openers / closers. Plain prose, hyphens (not em
681
+ / en dashes), no emojis.
682
+ `;
683
+
393
684
  /**
394
- * Drain the genie `sendMessage` AsyncGenerator into a flat result
395
- * the agent's calling LLM can reason about, while forwarding
396
- * progress and chart events to the host UI.
397
- *
398
- * Three streams of output happen in parallel:
399
- *
400
- * 1. {@link GenieProgress} pill events on the writer (`started`,
401
- * `status`, `sql`, `suggested`, `error`) drive the loading
402
- * pill in the chat bubble.
403
- * 2. `kind: "chart"` events on the writer (emitted via
404
- * {@link emitChartWithPlanning}) carry the row payload from
405
- * each Genie SQL statement and, on planner success, a
406
- * follow-up event with the rendered Echarts spec. The host
407
- * UI's `<ChartSlot>` merges the two by `chartId` and
408
- * renders inline at the marker position the model picked.
409
- * The data never reaches the LLM.
410
- * 3. The `DrainResult` returned to the LLM contains Genie's
411
- * prose answer plus a `datasets[]` array of metadata
412
- * (chartId, title, columns, rowCount, sql) the model uses
413
- * to cite charts via `[[chart:<chartId>]]` markers.
414
- *
415
- * `query_result` and `message_result` events arrive in either
416
- * order; we buffer per-statement scratch keyed by `statementId`
417
- * so each half can fill in what it knows. The chart event
418
- * fires the moment `query_result` lands; the planner runs in
419
- * the background. We `Promise.allSettled` every planner promise
420
- * before returning so all chart work is attributed to the tool's
421
- * trace span and so the LLM's `datasets[]` includes every
422
- * chartId that has actually been queued.
685
+ * Boundary schema for the inner agent's structured output. Two
686
+ * tagged shapes only - text or data. The wrapper maps these onto
687
+ * the shared {@link GenieSummaryItem} (`string` / `visualize`)
688
+ * after charting; we don't redefine GenieSummaryItem here.
423
689
  */
424
- async function drainGenieStream(
425
- stream: AsyncGenerator<GenieStreamEvent>,
426
- writer: ToolStream | undefined,
427
- opts: DrainGenieStreamOptions,
428
- ): Promise<DrainResult> {
429
- const { config, requestContext } = opts;
430
- let conversationId: string | undefined;
431
- let genieAnswer: string | undefined;
432
- let suggestedFollowUps: string[] | undefined;
433
- let error: string | undefined;
434
- // AppKit's `streamSendMessage` forwards every SDK `onProgress`
435
- // callback verbatim - the same `EXECUTING_QUERY` can fire several
436
- // times during a single poll loop. AppKit's other path,
437
- // `streamGetMessage`, dedupes on the connector side; we mirror that
438
- // behaviour here so the UI status pill doesn't flicker and we don't
439
- // burn writer bytes on no-op events.
440
- let lastStatus: string | undefined;
441
-
442
- // Per-statement scratch keyed by Genie's `statementId`. Filled in
443
- // by both `query_result` (chartId + columns + rows) and
444
- // `message_result` (sql + title + description). The LLM-bound
445
- // `datasets[]` is built from this at end-of-stream, after all
446
- // planner promises settle.
447
- type Scratch = {
448
- statementId: string;
449
- chartId?: string;
450
- title?: string;
451
- description?: string;
452
- sql?: string;
453
- columns: string[];
454
- rowCount: number;
455
- };
456
- const scratchByStatementId = new Map<string, Scratch>();
457
- const getScratch = (statementId: string): Scratch => {
458
- let s = scratchByStatementId.get(statementId);
459
- if (!s) {
460
- s = { statementId, columns: [], rowCount: 0 };
461
- scratchByStatementId.set(statementId, s);
462
- }
463
- return s;
464
- };
690
+ const agentSummarySchema = z.object({
691
+ summary: z.array(
692
+ z.discriminatedUnion("type", [
693
+ z.object({
694
+ type: z.literal("text"),
695
+ text: z.string(),
696
+ }),
697
+ z.object({
698
+ type: z.literal("data"),
699
+ statementId: z.string(),
700
+ title: z.string().optional(),
701
+ description: z.string().optional(),
702
+ }),
703
+ ]),
704
+ ),
705
+ });
706
+
707
+ type AgentSummaryItem = z.infer<typeof agentSummarySchema>["summary"][number];
708
+
709
+ /* ----------------------------- factory ----------------------------- */
710
+
711
+ /**
712
+ * Options for {@link createGenieTool}. Only carries config that
713
+ * doesn't vary per request - the per-request {@link WorkspaceClient},
714
+ * `RequestContext`, writer, and abort signal flow through the
715
+ * tool's `execute(_, ctx)` and are not captured here.
716
+ */
717
+ export interface CreateGenieToolOptions {
718
+ /** Genie space id this tool targets. */
719
+ spaceId: string;
720
+ /** Plugin config; resolves the LLM and chart planner agent. */
721
+ config: MastraPluginConfig;
722
+ /** Override the registered tool id. Defaults to `"genie"`. */
723
+ toolId?: string;
724
+ /** Override the tool description shown to the calling LLM. */
725
+ toolDescription?: string;
465
726
  /**
466
- * Planner promises kicked off per `query_result`. Awaited
467
- * (Promise.allSettled) before drainGenieStream returns so the
468
- * Genie tool's trace span covers the chart work and the LLM's
469
- * `datasets[]` accurately reflects every chartId that's been
470
- * queued for rendering.
727
+ * Override the inner agent's max tool-loop steps. Defaults to
728
+ * {@link DEFAULT_MAX_STEPS}.
471
729
  */
472
- const plannerPromises: Promise<void>[] = [];
473
-
474
- const emit = async (event: GenieProgress) => {
475
- if (!writer) return;
476
- try {
477
- await writer.write(event);
478
- } catch {
479
- // ignore: downstream stream is no longer interested
480
- }
481
- };
730
+ maxSteps?: number;
731
+ }
732
+
733
+ /**
734
+ * Build the calling agent's Genie tool. The returned Mastra tool
735
+ * runs end-to-end on each invocation:
736
+ *
737
+ * 1. Pull the per-request `WorkspaceClient` off
738
+ * `ctx.requestContext` (stamped by `MastraServer` under
739
+ * {@link MASTRA_USER_KEY}) and emit a `started` writer
740
+ * event so the host UI shows progress immediately.
741
+ * 2. Spin up the inner Mastra agent + three tools, fresh per
742
+ * call so the row cache stays invocation-scoped.
743
+ * 3. Run the agent with `structuredOutput` against
744
+ * {@link agentSummarySchema}. Mastra's two-pass design keeps
745
+ * the inner loop tools-only (no `response_format`), so the
746
+ * Databricks Model Serving `response_format`+`tools`
747
+ * collision never fires.
748
+ * 4. Walk the returned `[text|data][]`, map `text` items to
749
+ * shared `GenieSummaryItem.string`, and chart every `data`
750
+ * item in parallel via {@link runChartPlanner} to a
751
+ * `GenieSummaryItem.visualize`. Items referencing a missing
752
+ * `statementId` are dropped with a warn log; chart-planner
753
+ * failures leave `dataset.chart` unset so the host UI falls
754
+ * back to a table.
755
+ */
756
+ export function createGenieTool(opts: CreateGenieToolOptions) {
757
+ const {
758
+ spaceId,
759
+ config,
760
+ toolId = "genie",
761
+ toolDescription = stringUtils.toDescription`
762
+ Ask a question about the Databricks Genie space.
482
763
 
483
- for await (const event of stream) {
484
- // Per-event raw payload for tuning the pill / answer pipeline
485
- // against real Genie traffic. At `info` (the default) this is
486
- // discarded for free; flip `LOG_LEVEL=debug` to see every
487
- // raw wire event before the switch routes it through writer
488
- // and DrainResult.
489
- log.debug("event", { type: event.type, payload: event });
490
- switch (event.type) {
491
- case "message_start":
492
- conversationId = event.conversationId;
493
- await emit({
494
- kind: "started",
495
- conversationId: event.conversationId,
496
- messageId: event.messageId,
497
- spaceId: event.spaceId,
498
- });
499
- break;
500
- case "status":
501
- if (event.status === lastStatus) break;
502
- lastStatus = event.status;
503
- await emit({
504
- kind: "status",
505
- status: event.status,
506
- label: humanizeGenieStatus(event.status),
507
- });
508
- break;
509
- case "query_result": {
510
- const columns = (event.data?.manifest?.schema?.columns ?? []).map(
511
- (c) => c.name,
764
+ Returns \`{ summary: SummaryItem[] }\` where each item is
765
+ one of:
766
+
767
+ - \`{ type: "string", text }\` - prose to weave into your
768
+ reply verbatim or paraphrase.
769
+ - \`{ type: "visualize", statementId, title?, description?,
770
+ dataset: { data: { columns, rows, rowCount },
771
+ chart?: { chartId, chartType } } }\` - a chartable result
772
+ set. When \`dataset.chart\` is present the chart is ALREADY
773
+ rendered and queued for inline display; embed the marker
774
+ \`[[chart:<chartId>]]\` on its own line at the position
775
+ you want it to appear and the host UI drops the rendered
776
+ chart in. Re-use the chartId verbatim - do NOT call
777
+ \`render_data\` for the same dataset (it would render the
778
+ same chart a second time and stall your stream). Only
779
+ fall back to \`render_data\` when \`dataset.chart\` is
780
+ missing (chart-planner failed) AND you genuinely need a
781
+ picture; otherwise present the data inline as prose or a
782
+ short table.
783
+ `,
784
+ maxSteps = DEFAULT_MAX_STEPS,
785
+ } = opts;
786
+
787
+ return createTool({
788
+ id: toolId,
789
+ description: toolDescription,
790
+ inputSchema: z.object({
791
+ question: z.string().describe(stringUtils.toDescription`
792
+ Natural-language question about the data in this Genie
793
+ space. Phrase it from the user's perspective; the agent
794
+ decomposes it internally.
795
+ `),
796
+ }),
797
+ outputSchema: z.custom<GenieAgentResult>(),
798
+ execute: async (input, ctxRaw) => {
799
+ const ctx = ctxRaw as
800
+ | {
801
+ requestContext?: RequestContext;
802
+ writer?: MinimalWriter;
803
+ abortSignal?: AbortSignal;
804
+ }
805
+ | undefined;
806
+ const requestContext = ctx?.requestContext;
807
+ if (!requestContext) {
808
+ throw new Error(
809
+ "genie: missing requestContext (MastraServer must stamp MASTRA_USER_KEY)",
512
810
  );
513
- const dataArray = (event.data?.result?.data_array ?? []) as Array<
514
- Array<string | null>
515
- >;
516
- const rows = genieRowsToObjects(columns, dataArray);
517
- const scratch = getScratch(event.statementId);
518
- // emitChartWithPlanning emits the dataset event immediately
519
- // and kicks off the chart-planner agent in the background.
520
- // It returns the chartId synchronously; the plannerPromise
521
- // is awaited at end-of-stream so chart work shows up under
522
- // this tool's trace span.
523
- const { chartId, plannerPromise } = await emitChartWithPlanning({
524
- ...(writer ? { writer } : {}),
525
- config,
526
- ...(requestContext ? { requestContext } : {}),
527
- title: scratch.title ?? `Genie query`,
528
- ...(scratch.description ? { description: scratch.description } : {}),
529
- data: rows,
530
- });
531
- scratch.chartId = chartId;
532
- scratch.columns = columns;
533
- scratch.rowCount = rows.length;
534
- plannerPromises.push(plannerPromise);
535
- log.debug("query_result", {
536
- statementId: event.statementId,
537
- chartId,
538
- rows: rows.length,
539
- columns,
540
- });
541
- break;
542
811
  }
543
- case "message_result":
544
- genieAnswer = event.message.content;
545
- for (const attachment of event.message.attachments ?? []) {
546
- const sqlText = attachment.query?.query;
547
- const stmtId = attachment.query?.statementId;
548
- if (stmtId) {
549
- const scratch = getScratch(stmtId);
550
- if (sqlText) scratch.sql = sqlText;
551
- if (attachment.query?.title) scratch.title = attachment.query.title;
552
- if (attachment.query?.description) {
553
- scratch.description = attachment.query.description;
812
+ const user = requestContext.get(MASTRA_USER_KEY) as User | undefined;
813
+ if (!user) {
814
+ throw new Error("genie: no user on requestContext (MASTRA_USER_KEY not set)");
815
+ }
816
+ const client = user.executionContext.client;
817
+ const writer = ctx?.writer;
818
+ const signal = ctx?.abortSignal;
819
+ const threadId = requestContext.get(MASTRA_THREAD_ID_KEY) as string | undefined;
820
+
821
+ // Fire the lifecycle `started` event before any LLM /
822
+ // network round-trip so the host UI can pop a "Thinking..."
823
+ // pill the instant the model decides to delegate. The wire
824
+ // `conversation_id` / `message_id` aren't known yet (no
825
+ // Genie call has been made) and ride as `undefined` -
826
+ // subscribers that need them watch the later
827
+ // `message` / `result` wire events for the real ids.
828
+ const startedEvent: StartedEvent = {
829
+ type: "started",
830
+ spaceId,
831
+ content: input.question,
832
+ };
833
+ await safeWrite(log, writer, startedEvent);
834
+
835
+ const resultSets = new Map<string, StatementEntry>();
836
+
837
+ // Seed the active Genie `conversation_id` onto
838
+ // `RequestContext` from AppKit's `CacheManager` when a Mastra
839
+ // `threadId` is present so multi-turn chats reuse the same
840
+ // Genie conversation (and Genie's accumulated context) across
841
+ // separate tool invocations. The same `RequestContext` flows
842
+ // to the inner `ask_genie` tool via Mastra, which reads and
843
+ // updates the same slot as Genie hands out / rotates ids.
844
+ // Cache misses, threads without memory, and unhealthy cache
845
+ // storage all leave the slot unset, which makes `ask_genie`
846
+ // call `startConversation` and mint a fresh id (then cache
847
+ // it).
848
+ const cacheKey = await conversationCacheKey(spaceId, threadId);
849
+ const cachedConversationId = await readCachedConversationId(cacheKey);
850
+ if (cachedConversationId) {
851
+ writeContextConversationId(requestContext, spaceId, cachedConversationId);
852
+ }
853
+
854
+ const innerDeps: InnerToolDeps = {
855
+ spaceId,
856
+ client,
857
+ ...(writer ? { writer } : {}),
858
+ ...(signal ? { signal } : {}),
859
+ resultSets,
860
+ ...(cacheKey ? { cacheKey } : {}),
861
+ };
862
+ const tools = {
863
+ ask_genie: buildAskGenieTool(innerDeps),
864
+ get_space_description: buildSpaceDescriptionTool({
865
+ spaceId,
866
+ client,
867
+ ...(signal ? { signal } : {}),
868
+ }),
869
+ get_space_serialized: buildSpaceSerializedTool({
870
+ spaceId,
871
+ client,
872
+ ...(signal ? { signal } : {}),
873
+ }),
874
+ };
875
+
876
+ // Resolve the model config once for this request so we can
877
+ // share it with the structuring pass below. The agent's
878
+ // `model` field accepts a function form for per-request
879
+ // resolution, but `structuredOutput.model` requires a
880
+ // static `MastraModelConfig`, and we need both to be on
881
+ // the same Databricks endpoint with the same OBO-scoped
882
+ // headers. Calling `buildModel` here (inside `execute`)
883
+ // keeps user scoping correct because `requestContext`
884
+ // already reflects the active request's user.
885
+ const resolvedModel = await buildModel(config, requestContext);
886
+
887
+ const agent = new Agent({
888
+ id: `genie__${spaceId}`,
889
+ name: `Genie (${spaceId})`,
890
+ description: stringUtils.toDescription`
891
+ Inner orchestrator for the "${spaceId}" Genie space.
892
+ Asks Genie one focused sub-question at a time and
893
+ returns an interleaved [text|data] summary.
894
+ `,
895
+ instructions: AGENT_INSTRUCTIONS,
896
+ model: resolvedModel,
897
+ tools,
898
+ });
899
+
900
+ // Mastra's `structuredOutput` operates in one of two modes
901
+ // based on whether `model` is set:
902
+ // - "direct" (no model) -> the schema is enforced
903
+ // in the SAME LLM call as
904
+ // the agent loop, by
905
+ // adding `response_format`
906
+ // alongside `tools`.
907
+ // Databricks Model Serving
908
+ // rejects that combination
909
+ // with `INVALID_PARAMETER_VALUE:
910
+ // Cannot specify both
911
+ // response_format and tools
912
+ // in the same request.`
913
+ // - "processor" (model passed) -> the main loop carries
914
+ // tools and NO
915
+ // `response_format`; a
916
+ // separate, tool-free
917
+ // structuring agent
918
+ // re-prompts the model
919
+ // with `response_format`
920
+ // to coerce the agent's
921
+ // final text into the
922
+ // schema.
923
+ // We use "processor" mode but ALSO set
924
+ // `jsonPromptInjection: true`. Mastra's structuring agent
925
+ // calls `.stream(...)` under the hood, and Databricks Model
926
+ // Serving rejects `response_format` together with streaming
927
+ // (`INVALID_PARAMETER_VALUE: Structured output is not
928
+ // currently supported with streaming.`). Prompt injection
929
+ // sidesteps that by embedding the JSON Schema in the
930
+ // structuring agent's system prompt instead of sending
931
+ // `response_format`. `errorStrategy: "warn"` keeps a
932
+ // structuring failure from escaping as an unhandled
933
+ // promise rejection: it logs and leaves `result.object`
934
+ // undefined, which we surface as a clean error in
935
+ // {@link GenieAgentResult}.
936
+ const agentResult = await agent.generate(input.question, {
937
+ requestContext,
938
+ maxSteps,
939
+ structuredOutput: {
940
+ schema: agentSummarySchema,
941
+ model: resolvedModel,
942
+ jsonPromptInjection: true,
943
+ errorStrategy: "warn",
944
+ },
945
+ ...(signal ? { abortSignal: signal } : {}),
946
+ });
947
+ const submission = agentResult.object;
948
+ if (!submission) {
949
+ const message = "Genie agent returned no structured summary";
950
+ log.warn("agent:no-summary", { spaceId });
951
+ const finalConversationId = readContextConversationId(requestContext, spaceId);
952
+ return {
953
+ spaceId,
954
+ summary: [],
955
+ ...(finalConversationId ? { conversationId: finalConversationId } : {}),
956
+ error: message,
957
+ } satisfies GenieAgentResult;
958
+ }
959
+
960
+ // Lifecycle hook: the agent + structuring pass are done.
961
+ // Emit one `summary` event with the structured-item counts
962
+ // so the host UI can transition from "thinking" to
963
+ // "charting" and seed N chart skeletons before the
964
+ // per-chart `chart` events arrive. We can't fire this
965
+ // EARLIER (i.e. when the structuring pass starts) because
966
+ // Mastra runs the inner loop + structuring pass together
967
+ // inside `agent.generate(...)` with no observable boundary
968
+ // between them.
969
+ const textItemCount = submission.summary.filter(
970
+ (i: AgentSummaryItem) => i.type === "text",
971
+ ).length;
972
+ const dataItemCount = submission.summary.length - textItemCount;
973
+ const summaryEvent: SummaryEvent = {
974
+ type: "summary",
975
+ spaceId,
976
+ items: submission.summary.length,
977
+ textItems: textItemCount,
978
+ dataItems: dataItemCount,
979
+ };
980
+ await safeWrite(log, writer, summaryEvent);
981
+
982
+ // Chart every `data` item in parallel; map `text` items to
983
+ // the shared `string` summary variant verbatim. Missing
984
+ // statement ids are dropped (the agent referenced something
985
+ // that never came back from `ask_genie`), planner failures
986
+ // leave `dataset.chart` unset so the host UI falls back to
987
+ // a table render. Each successfully planned chart pushes a
988
+ // `chart` writer event so the UI can fade in the rendered
989
+ // chart slot the moment its planner returns rather than
990
+ // waiting for the entire batch to finish.
991
+ const hydrated = await Promise.all(
992
+ submission.summary.map(
993
+ async (item: AgentSummaryItem): Promise<GenieSummaryItem | undefined> => {
994
+ if (item.type === "text") {
995
+ return { type: "string", text: item.text };
554
996
  }
555
- }
556
- if (sqlText) {
557
- await emit({
558
- kind: "sql",
559
- sql: sqlText,
560
- title: attachment.query?.title,
561
- description: attachment.query?.description,
562
- statementId: stmtId,
563
- });
564
- }
565
- if (attachment.text?.content) {
566
- await emit({ kind: "text", content: attachment.text.content });
567
- }
568
- if (attachment.suggestedQuestions?.length) {
569
- // Last attachment with suggestions wins (same merge rule
570
- // the UI uses via `collectSuggestions`); keeping just one
571
- // copy per turn caps token usage.
572
- suggestedFollowUps = attachment.suggestedQuestions;
573
- await emit({
574
- kind: "suggested",
575
- questions: attachment.suggestedQuestions,
576
- });
577
- }
578
- }
579
- break;
580
- case "error":
581
- error = event.error;
582
- await emit({ kind: "error", error: event.error });
583
- break;
584
- default:
585
- break;
586
- }
587
- }
997
+ const entry = resultSets.get(item.statementId);
998
+ if (!entry) {
999
+ log.warn("data:missing-statement", {
1000
+ statementId: item.statementId,
1001
+ });
1002
+ return undefined;
1003
+ }
1004
+ const { data, messageId } = entry;
1005
+ let dataset: GenieDataset = { data };
1006
+ try {
1007
+ const planned = await runChartPlanner({
1008
+ config,
1009
+ requestContext,
1010
+ title: item.title ?? "Genie result",
1011
+ ...(item.description ? { description: item.description } : {}),
1012
+ data: data.rows,
1013
+ ...(signal ? { signal } : {}),
1014
+ });
1015
+ const chartId = commonUtils.shortId();
1016
+ // Slim chart reference for the LLM-bound result: just
1017
+ // `chartId` + `chartType`. The full Echarts spec goes
1018
+ // to the UI via the writer event AND into the
1019
+ // request-scoped chart inventory below; the model
1020
+ // only needs the id to place `[[chart:<id>]]`.
1021
+ dataset = {
1022
+ data,
1023
+ chart: {
1024
+ chartId,
1025
+ chartType: planned.chartType,
1026
+ },
1027
+ };
1028
+ const chartEvent: ChartEvent = {
1029
+ type: "chart",
1030
+ chartId,
1031
+ statementId: item.statementId,
1032
+ messageId,
1033
+ ...(item.title ? { title: item.title } : {}),
1034
+ ...(item.description ? { description: item.description } : {}),
1035
+ data: data.rows,
1036
+ option: planned.option,
1037
+ };
1038
+ await safeWrite(log, writer, chartEvent);
1039
+ // Stash the resolved chart on the per-request
1040
+ // `RequestContext` so downstream code in the same
1041
+ // request (output processors, follow-up tool calls,
1042
+ // any post-run hook) can look up the full spec by
1043
+ // `chartId` without re-fetching or re-planning.
1044
+ recordChartInContext(requestContext, chartEvent);
1045
+ } catch (err) {
1046
+ const errorMessage = commonUtils.errorMessage(err);
1047
+ log.warn("chart:error", {
1048
+ statementId: item.statementId,
1049
+ messageId,
1050
+ error: errorMessage,
1051
+ });
1052
+ // Surface the chart-planner failure as a writer event
1053
+ // stamped with the same `messageId` the rest of this
1054
+ // ask's wire events carry, so the host UI groups the
1055
+ // failure into the same pill bucket and can surface
1056
+ // a "couldn't render chart" note next to the table
1057
+ // fallback instead of silently dropping the chart.
1058
+ const errorEvent: MastraGenieErrorEvent = {
1059
+ type: "error",
1060
+ spaceId,
1061
+ messageId,
1062
+ error: `chart-planner: ${errorMessage}`,
1063
+ };
1064
+ await safeWrite(log, writer, errorEvent);
1065
+ }
1066
+ return {
1067
+ type: "visualize",
1068
+ statementId: item.statementId,
1069
+ ...(item.title ? { title: item.title } : {}),
1070
+ ...(item.description ? { description: item.description } : {}),
1071
+ dataset,
1072
+ };
1073
+ },
1074
+ ),
1075
+ );
1076
+ const summary = hydrated.filter((x): x is GenieSummaryItem => x !== undefined);
588
1077
 
589
- // Wait for all chart planners to settle before returning so the
590
- // tool's trace span covers chart work and the LLM's
591
- // `datasets[]` reflects only chartIds the client has actually
592
- // received writer events for. Failures in `emitChartWithPlanning`
593
- // are already swallowed inside the helper, so this never
594
- // throws.
595
- log.debug("planners:awaiting", { count: plannerPromises.length });
596
- await Promise.allSettled(plannerPromises);
597
- log.debug("planners:settled", { count: plannerPromises.length });
598
-
599
- // Build the LLM-bound `datasets[]` from scratch entries that
600
- // actually ran a query (chartId is assigned at `query_result`
601
- // time). Entries that only saw `message_result` metadata
602
- // without a row payload are skipped.
603
- const datasets: Array<z.infer<typeof datasetSchema>> = [];
604
- for (const scratch of scratchByStatementId.values()) {
605
- if (!scratch.chartId) continue;
606
- datasets.push({
607
- chartId: scratch.chartId,
608
- ...(scratch.title ? { title: scratch.title } : {}),
609
- ...(scratch.description ? { description: scratch.description } : {}),
610
- columns: scratch.columns,
611
- rowCount: scratch.rowCount,
612
- ...(scratch.sql ? { sql: scratch.sql } : {}),
613
- });
614
- }
1078
+ log.info("genie:done", {
1079
+ spaceId,
1080
+ items: summary.length,
1081
+ statementsCharted: summary.filter(
1082
+ (s) => s.type === "visualize" && s.dataset.chart,
1083
+ ).length,
1084
+ });
615
1085
 
616
- log.debug("drain:return", {
617
- conversationId,
618
- hasAnswer: typeof genieAnswer === "string",
619
- answerLength: genieAnswer?.length ?? 0,
620
- chartIds: datasets.map((d) => d.chartId),
621
- suggestedCount: suggestedFollowUps?.length ?? 0,
622
- error,
1086
+ const finalConversationId = readContextConversationId(requestContext, spaceId);
1087
+ return {
1088
+ spaceId,
1089
+ summary,
1090
+ ...(finalConversationId ? { conversationId: finalConversationId } : {}),
1091
+ } satisfies GenieAgentResult;
1092
+ },
623
1093
  });
1094
+ }
624
1095
 
625
- return {
626
- ...(conversationId ? { conversationId } : {}),
627
- ...(genieAnswer ? { genieAnswer } : {}),
628
- ...(datasets.length > 0 ? { datasets } : {}),
629
- ...(suggestedFollowUps ? { suggestedFollowUps } : {}),
630
- ...(error ? { error } : {}),
631
- };
1096
+ /* --------------------- multi-alias surface --------------------- */
1097
+
1098
+ /**
1099
+ * Default tool id for a wired Genie alias. The well-known
1100
+ * `default` alias collapses to `genie`; every other alias gets a
1101
+ * `genie_` prefix so multi-space registrations stay
1102
+ * disambiguated.
1103
+ */
1104
+ export function defaultGenieToolName(alias: string): string {
1105
+ if (alias === DEFAULT_GENIE_ALIAS) return "genie";
1106
+ return stringUtils.toIdentifierWithOptions({ distinct: true }, "genie", alias);
632
1107
  }
633
1108
 
634
1109
  /**
635
- * Convert Genie's `data_array` (column-positional `string | null`
636
- * tuples) into plain JS row objects keyed by column name. Numeric
637
- * strings are coerced to numbers so the chart-planner picks
638
- * `value` axes instead of `category` axes; everything else passes
639
- * through verbatim. `null` becomes `null`.
1110
+ * Normalize the {@link GenieSpacesConfig} record. Bare-string
1111
+ * entries (`{ default: "01ef..." }`) get wrapped as
1112
+ * `{ spaceId: "01ef..." }`; object entries pass through unchanged.
1113
+ * `undefined` and empty-string values are dropped so callers can
1114
+ * pass `process.env.X` directly (matches AppKit `genie()`'s
1115
+ * defensive treatment of unset env vars).
640
1116
  */
641
- function genieRowsToObjects(
642
- columns: ReadonlyArray<string>,
643
- dataArray: ReadonlyArray<ReadonlyArray<string | null>>,
644
- ): Array<Record<string, unknown>> {
645
- const out: Array<Record<string, unknown>> = [];
646
- for (const row of dataArray) {
647
- const obj: Record<string, unknown> = {};
648
- columns.forEach((col, i) => {
649
- const cell = row[i] ?? null;
650
- obj[col] = coerceCell(cell);
651
- });
652
- out.push(obj);
1117
+ export function normalizeGenieSpaces(
1118
+ spaces:
1119
+ | GenieSpacesConfig
1120
+ | Record<string, string | GenieSpaceConfig | undefined>
1121
+ | undefined,
1122
+ ): Record<string, GenieSpaceConfig> {
1123
+ if (!spaces) return {};
1124
+ const out: Record<string, GenieSpaceConfig> = {};
1125
+ for (const [alias, value] of Object.entries(spaces)) {
1126
+ if (value === undefined) continue;
1127
+ if (typeof value === "string") {
1128
+ if (!value) continue;
1129
+ out[alias] = { spaceId: value };
1130
+ continue;
1131
+ }
1132
+ if (!value.spaceId) continue;
1133
+ out[alias] = value;
653
1134
  }
654
1135
  return out;
655
1136
  }
656
1137
 
657
- /** Best-effort numeric coercion for Genie's all-strings cells. */
658
- function coerceCell(cell: string | null): unknown {
659
- if (cell === null) return null;
660
- // Anchored to keep `12.5px` / `123abc` as strings; only fully
661
- // numeric values become JS numbers.
662
- if (/^-?\d+(\.\d+)?$/.test(cell)) {
663
- const n = Number(cell);
664
- if (Number.isFinite(n)) return n;
1138
+ /**
1139
+ * AppKit `genie` plugin's config shape, derived from the factory
1140
+ * itself so it stays in lock-step with the upstream type without
1141
+ * deep-importing `IGenieConfig` (which the package's top-level
1142
+ * barrel doesn't surface). The plugin's `config` field is
1143
+ * `protected` in TS only; the runtime layout is plain object
1144
+ * property access, so reading off the instance with a structural
1145
+ * cast is safe.
1146
+ */
1147
+ type AppKitGenieConfig = NonNullable<Parameters<typeof genie>[0]>;
1148
+
1149
+ /**
1150
+ * Discover Genie space aliases from every supported source and
1151
+ * merge them into a single record. Precedence (highest first):
1152
+ *
1153
+ * 1. {@link MastraPluginConfig.genieSpaces} on the `mastra(...)`
1154
+ * call. Explicit Mastra wiring always wins so users can
1155
+ * override AppKit's defaults per-agent.
1156
+ * 2. AppKit `genie({ spaces: { ... } })` plugin instance. Lets
1157
+ * users keep using the existing AppKit config format
1158
+ * (`genie({ spaces: { sales: "...", ops: "..." } })`)
1159
+ * without restating the same record on the Mastra plugin.
1160
+ * Read off the live plugin instance via a structural cast
1161
+ * since `Plugin.config` is TS-protected (not runtime-private).
1162
+ * 3. `DATABRICKS_GENIE_SPACE_ID` env var (registered under the
1163
+ * well-known `default` alias). Matches the AppKit `genie()`
1164
+ * plugin's fallback behavior so a bare `mastra()` + `genie()`
1165
+ * pair just works.
1166
+ *
1167
+ * Aliases collide cleanly: a higher-precedence source's value
1168
+ * replaces a lower one's wholesale. Sources that contribute zero
1169
+ * aliases (or contribute only `undefined` / empty entries) are
1170
+ * silently ignored.
1171
+ */
1172
+ export function resolveGenieSpaces(
1173
+ config: MastraPluginConfig,
1174
+ context: appkitUtils.PluginContextLike | undefined,
1175
+ ): Record<string, GenieSpaceConfig> {
1176
+ const merged: Record<string, GenieSpaceConfig> = {};
1177
+
1178
+ // Source 3 (lowest precedence): env var.
1179
+ const envSpaceId = process.env["DATABRICKS_GENIE_SPACE_ID"];
1180
+ if (envSpaceId) {
1181
+ merged[DEFAULT_GENIE_ALIAS] = { spaceId: envSpaceId };
665
1182
  }
666
- return cell;
1183
+
1184
+ // Source 2: AppKit `genie()` plugin instance config. Use a
1185
+ // structural cast - `Plugin.config` is `protected` in TS only,
1186
+ // and the runtime layout is plain object property access.
1187
+ const geniePlugin = appkitUtils.instance(context, genie);
1188
+ if (geniePlugin) {
1189
+ const pluginSpaces = (geniePlugin as unknown as { config?: AppKitGenieConfig })
1190
+ .config?.spaces;
1191
+ if (pluginSpaces) {
1192
+ Object.assign(merged, normalizeGenieSpaces(pluginSpaces));
1193
+ }
1194
+ }
1195
+
1196
+ // Source 1 (highest precedence): explicit Mastra wiring.
1197
+ if (config.genieSpaces) {
1198
+ Object.assign(merged, normalizeGenieSpaces(config.genieSpaces));
1199
+ }
1200
+
1201
+ return merged;
667
1202
  }
668
1203
 
669
1204
  /**
670
- * Toolkit provider built from a live AppKit `GeniePlugin` instance.
671
- * Returned by {@link buildGenieProvider} so that
672
- * `plugins.genie?.toolkit()` inside an agent's `tools(plugins)` callback
673
- * resolves to the streaming-aware {@link buildGenieTools} record instead
674
- * of the AppKit default (which does one blocking call per tool with no
675
- * mid-flight events).
676
- *
677
- * The returned `toolkit()` reads alias names off the plugin's
678
- * `getAgentTools()` registry (each entry is `${alias}.sendMessage` or
679
- * `${alias}.getConversation`), then mints one `sendMessage` tool per
680
- * alias plus a shared `getConversation`. `sendMessage` / `getConversation`
681
- * are bound back to the plugin instance so they keep their `this`
682
- * (they are class methods, not free functions).
1205
+ * Build one Mastra tool per configured Genie space. Each tool is
1206
+ * a thin {@link createGenieTool} wrapper with the alias-derived
1207
+ * id and a hint-flavored description so the calling LLM knows
1208
+ * which space covers what data.
683
1209
  *
684
- * `_opts` is accepted but unused for now - the streaming tools are an
685
- * all-or-nothing bundle. Wire `only` / `except` / `prefix` / `rename`
686
- * later if a caller needs them.
1210
+ * Returns a record keyed by tool id, ready to spread into an
1211
+ * `Agent`'s `tools` map (or surfaced via
1212
+ * `plugins.genie?.toolkit()`).
1213
+ */
1214
+ export function buildGenieTools(opts: {
1215
+ spaces: GenieSpacesConfig | Record<string, GenieSpaceConfig>;
1216
+ config: MastraPluginConfig;
1217
+ }): MastraTools {
1218
+ const normalized = normalizeGenieSpaces(opts.spaces);
1219
+ const tools: Record<string, ReturnType<typeof createTool>> = {};
1220
+ for (const [alias, space] of Object.entries(normalized)) {
1221
+ const id = defaultGenieToolName(alias);
1222
+ const toolDescription = stringUtils.toDescription`
1223
+ Delegate a natural-language data question to the
1224
+ Databricks Genie space "${alias}"${space.hint ? ` (${space.hint})` : ""}.
1225
+ Returns an ordered (text | dataset)[] summary the host UI
1226
+ renders inline; datasets carry the rows and a
1227
+ pre-rendered Echarts spec when the chart-planner
1228
+ succeeded. Progress events (status, SQL, row counts,
1229
+ charts) stream to the UI automatically.
1230
+ `;
1231
+ tools[id] = createGenieTool({
1232
+ spaceId: space.spaceId,
1233
+ config: opts.config,
1234
+ toolId: id,
1235
+ toolDescription,
1236
+ });
1237
+ }
1238
+ return tools;
1239
+ }
1240
+
1241
+ /**
1242
+ * Plugin-toolkit adapter so the `plugins.genie?.toolkit()` lookup
1243
+ * inside an agent's `tools(plugins)` callback returns the
1244
+ * Genie agent-backed tools instead of throwing on missing plugin.
1245
+ * Mirrors AppKit's `PluginToolkitProvider` shape.
687
1246
  */
688
- export function buildGenieProvider(
689
- plugin: GeniePluginInstance,
690
- opts: { config: MastraPluginConfig },
691
- ): {
692
- toolkit(opts?: unknown): Record<string, ReturnType<typeof createTool>>;
1247
+ export function buildGenieToolkitProvider(opts: {
1248
+ spaces: GenieSpacesConfig | Record<string, GenieSpaceConfig>;
1249
+ config: MastraPluginConfig;
1250
+ }): {
1251
+ toolkit(opts?: unknown): MastraTools;
693
1252
  } {
694
1253
  return {
695
1254
  toolkit(_opts?: unknown) {
696
- const aliases = extractGenieAliases(plugin);
697
- return buildGenieTools({
698
- aliases,
699
- exports: {
700
- sendMessage: plugin.sendMessage.bind(plugin),
701
- getConversation: plugin.getConversation.bind(plugin),
702
- },
703
- config: opts.config,
704
- });
1255
+ return buildGenieTools(opts);
705
1256
  },
706
1257
  };
707
1258
  }
708
1259
 
709
1260
  /**
710
- * Pull the configured space aliases out of a live AppKit `GeniePlugin`.
711
- * Reads them off `getAgentTools()` (public API) so we don't poke at the
712
- * `protected config.spaces` field: the plugin registers tools named
713
- * `${alias}.sendMessage` / `${alias}.getConversation`, so the unique
714
- * set of name prefixes is the alias list.
715
- */
716
- function extractGenieAliases(plugin: GeniePluginInstance): string[] {
717
- const aliases = new Set<string>();
718
- for (const t of plugin.getAgentTools()) {
719
- const dot = t.name.indexOf(".");
720
- if (dot > 0) aliases.add(t.name.slice(0, dot));
721
- }
722
- return [...aliases];
723
- }
724
-
725
- /**
726
- * Convert raw Genie status codes (`FETCHING_METADATA`, `ASKING_AI`,
727
- * `EXECUTING_QUERY`, `COMPLETED`, ...) into short, sentence-cased
728
- * labels safe to drop straight into a UI pill. Unknown codes are
729
- * lower-cased with underscores stripped so new states still render.
1261
+ * Returns `true` when at least one Genie space is reachable
1262
+ * through {@link resolveGenieSpaces} - either via
1263
+ * {@link MastraPluginConfig.genieSpaces}, the AppKit `genie()`
1264
+ * plugin instance, or the `DATABRICKS_GENIE_SPACE_ID` env var.
1265
+ *
1266
+ * Cheap to call from `resolveProvider` to short-circuit `genie`
1267
+ * lookups when nothing is wired, so the `plugins.genie` lookup
1268
+ * still resolves to `undefined` (matching AppKit's
1269
+ * absent-plugin semantics) when neither source is configured.
730
1270
  */
731
- function humanizeGenieStatus(status: string): string {
732
- switch (status) {
733
- case "FETCHING_METADATA":
734
- return "Fetching metadata";
735
- case "ASKING_AI":
736
- return "Asking Genie";
737
- case "EXECUTING_QUERY":
738
- return "Running SQL query";
739
- case "COMPLETED":
740
- return "Completed";
741
- case "FAILED":
742
- return "Failed";
743
- default:
744
- return [
745
- ...stringUtils.tokenizeWithOptions(
746
- { capitalize: true, lowerCase: true },
747
- status,
748
- ),
749
- ].join(" ");
750
- }
1271
+ export function hasAnyGenieSpaces(
1272
+ config: MastraPluginConfig,
1273
+ context: appkitUtils.PluginContextLike | undefined,
1274
+ ): boolean {
1275
+ return Object.keys(resolveGenieSpaces(config, context)).length > 0;
751
1276
  }