@dbx-tools/appkit-mastra 0.1.12 → 0.1.18

package/dist/src/genie.js

@@ -1,630 +1,1056 @@
 /**
- * Mastra tool wrappers around the AppKit `genie` plugin's exports.
+ * Genie agent for Mastra.
  *
- * One `sendMessage` tool is registered per configured space alias so
- * the LLM picks the space by tool selection (the description bakes the
- * alias in). `getConversation` is registered once, taking `alias` as a
- * parameter.
+ * Each configured Genie space exposes a single Mastra tool to the
+ * calling agent (`genie` for the `"default"` alias, `genie_<alias>`
+ * otherwise). When invoked, the tool runs end-to-end:
  *
- * All Genie payload types are inferred from the public `genie` factory
- * (`genie().plugin` constructor → `exports()` return type), so any
- * upstream change in `@databricks/appkit` flows in automatically.
+ *   1. Pulls the per-request {@link WorkspaceClient} off
+ *      `ctx.requestContext` (stamped by `MastraServer`) and emits a
+ *      `started` writer event so the host UI can show progress
+ *      immediately, before any LLM round-trip.
+ *   2. Spins up a per-call inner Mastra `Agent` with three tools:
+ *        - `ask_genie`: drives one `genieEventChat` turn, fetches
+ *          the matching statement's rows when the turn ran SQL,
+ *          and forwards every wire event (status, thinking, sql,
+ *          rows) through `ctx.writer` for streaming UI updates.
+ *        - `get_space_description`: cheap title / description /
+ *          warehouse id lookup for grounding.
+ *        - `get_space_serialized`: full `GenieSpace` JSON for
+ *          column-level grounding when the description isn't
+ *          enough.
+ *   3. Runs the inner agent with `structuredOutput` (Mastra's
+ *      two-pass mode + `jsonPromptInjection`) to coerce the
+ *      agent's final answer into a tagged
+ *      `[{type:"text"|"data", ...}]` array. The two-pass design
+ *      avoids Databricks Model Serving's `response_format` +
+ *      `tools` collision; prompt injection sidesteps the
+ *      separate `response_format` + streaming collision in the
+ *      structuring agent.
+ *   4. Charts every `data` item in parallel via
+ *      {@link runChartPlanner}, maps `text` items to the shared
+ *      {@link GenieSummaryItem} `string` variant, and returns the
+ *      hydrated {@link GenieAgentResult}.
  *
- * As Genie streams its long-running events (`FETCHING_METADATA` →
- * `ASKING_AI` → `EXECUTING_QUERY` → `COMPLETED`, plus SQL text and
- * follow-ups in `message_result.attachments`), the tool forwards a
- * normalised {@link GenieProgress} discriminated union out through
- * `ctx.writer` so the client can render an incremental loading pill.
- * Row payloads from `query_result` are intentionally discarded - the
- * LLM never sees rows, and charts come from the separate
- * `render_data` tool when the model decides one is useful.
+ * The legacy AppKit `genie` plugin (`@databricks/appkit`'s `genie`)
+ * is no longer used at runtime. The inner agent talks to Genie
+ * directly via `@dbx-tools/genie` (`genieEventChat`) and the
+ * workspace `statementExecution.getStatement` API. The plugin's
+ * `spaces` config is still honored so existing AppKit-style wiring
+ * keeps working without change.
  */
-import { genie } from "@databricks/appkit";
-import { logUtils, stringUtils } from "@dbx-tools/appkit-shared";
+import { CacheManager, genie } from "@databricks/appkit";
+import { ApiError, HttpError, WorkspaceClient } from "@databricks/sdk-experimental";
+import { genieEventChat } from "@dbx-tools/genie";
+import {} from "@dbx-tools/genie-shared";
+import {} from "@dbx-tools/appkit-mastra-shared";
+import { apiUtils, appkitUtils, commonUtils, logUtils, stringUtils, } from "@dbx-tools/shared";
+import { Agent } from "@mastra/core/agent";
+import { MASTRA_THREAD_ID_KEY } from "@mastra/core/request-context";
 import { createTool } from "@mastra/core/tools";
 import { z } from "zod";
-import { emitChartWithPlanning } from "./chart.js";
+import { runChartPlanner } from "./chart.js";
+import { MASTRA_USER_KEY } from "./config.js";
+import { buildModel } from "./model.js";
+const log = logUtils.logger("mastra/genie");
+/** Default alias used when a single unnamed Genie space is wired up. */
+export const DEFAULT_GENIE_ALIAS = "default";
 /**
- * Module-level logger tagged `[mastra/genie]`. Uses the shared
- * {@link logUtils.logger} so calls below `LOG_LEVEL` are
- * discarded for free. Default `LOG_LEVEL` is `info`; flip to
- * `debug` to see per-turn timing (`query_result` → planner
- * waits → `drain:return`).
+ * Cap on the inner agent's tool-loop steps. 5 (Mastra default) is
+ * tight - one `get_space_description` + one `ask_genie` per
+ * sub-question saturates fast. 16 leaves room for ~10 `ask_genie`
+ * rounds plus grounding plus the structuring pass (which runs
+ * after the loop and is its own single call).
  */
-const log = logUtils.logger("mastra/genie");
+const DEFAULT_MAX_STEPS = 16;
+/* ------------------------- helpers ------------------------- */
+/** Best-effort numeric coercion for Genie's all-strings cells. */
+function coerceCell(cell) {
+    if (cell === null)
+        return null;
+    if (/^-?\d+(\.\d+)?$/.test(cell)) {
+        const n = Number(cell);
+        if (Number.isFinite(n))
+            return n;
+    }
+    return cell;
+}
 /**
- * Per-dataset metadata surfaced to the LLM. The actual rows are
- * dispatched separately as a `kind: "chart"` writer event so the
- * model never has the rows in its context (token cost stays flat
- * regardless of dataset size). The model uses `chartId` to
- * reference the chart inline via the `[[chart:<chartId>]]` marker.
+ * Fetch a single Genie statement's rows via the Statement
+ * Execution API and reshape into the shared
+ * {@link GenieDatasetData} shape (column array + row records).
  */
-const datasetSchema = z.object({
-    chartId: z.string().describe(stringUtils.toDescription `
-    Short id (8 hex chars) for the chart-render slot the host UI
-    has staged for this dataset. Embed
-    \`[[chart:<chartId>]]\` on its own line in your reply at the
-    position you want the chart to appear; the client renders it
-    inline. Do not paraphrase the dataset's rows in prose - the
-    chart is the rendering.
-  `),
-    title: z.string().optional().describe(stringUtils.toDescription `
-    Genie's own title for the SQL that produced this dataset.
-    Useful as a label when you reference the chart in prose.
-  `),
-    description: z.string().optional().describe(stringUtils.toDescription `
-    Genie's prose description of the SQL, if any.
-  `),
-    columns: z.array(z.string()).describe(stringUtils.toDescription `
-    Column names in display order. Use these when describing what
-    is being charted (e.g. "trend of fill_rate over date").
-  `),
-    rowCount: z.number().describe(stringUtils.toDescription `
-    Total rows in this dataset. Mention only if it adds context
-    (e.g. "across the last 90 days").
-  `),
-    sql: z
-        .string()
-        .optional()
-        .describe(stringUtils.toDescription `
-      SQL Genie generated and executed. The host UI shows this on
-      demand; you do not need to repeat it.
-    `),
-});
+async function fetchStatementData(client, statementId, signal) {
+    const ctx = signal ? apiUtils.toContext(signal) : undefined;
+    const r = await client.statementExecution.getStatement({ statement_id: statementId }, ctx);
+    const columns = (r.manifest?.schema?.columns ?? []).map((c) => c.name ?? "");
+    const dataArray = (r.result?.data_array ?? []);
+    const rows = dataArray.map((row) => {
+        const obj = {};
+        columns.forEach((col, i) => {
+            obj[col] = coerceCell(row[i] ?? null);
+        });
+        return obj;
+    });
+    return {
+        columns,
+        rows,
+        rowCount: r.manifest?.total_row_count ?? rows.length,
+    };
+}
 /**
- * Top-level output schema returned to the LLM from a Genie tool
- * call. The `datasets` array is intentionally metadata-only - row
- * data rides a writer event the host UI consumes directly and is
- * not in the model's context.
+ * Resolve the message's representative `statement_id`. Genie
+ * returns one statement per turn in practice; we read the
+ * (deprecated-but-singular) `message.query_result.statement_id`
+ * first and fall back to the first attachment's
+ * `query.statement_id`. Returns `undefined` when the turn had no
+ * SQL run (pure prose answer).
  */
-const genieToolOutputSchema = z.object({
-    conversationId: z
-        .string()
-        .optional()
-        .describe(stringUtils.toDescription `
-      Pass back on the next call to continue the same Genie thread.
-    `),
-    genieAnswer: z
-        .string()
-        .optional()
-        .describe(stringUtils.toDescription `
-      Genie's natural-language answer to the question. Pass this
-      through to the user (verbatim, or as the basis of your
-      reply). Genie may have run multiple SQL queries and tools to
-      produce this; the full text is the answer.
-    `),
-    datasets: z
-        .array(datasetSchema)
-        .optional()
-        .describe(stringUtils.toDescription `
-      Datasets Genie produced for this turn (one per executed SQL
-      statement). Each entry is metadata only; the rows are
-      streamed to the host UI out-of-band. To render any of these
-      as a chart inline in your reply, embed
-      \`[[chart:<chartId>]]\` where you want the chart to appear.
-      Do not paraphrase the rows - the chart is what the user
-      should see; your prose should add interpretation
-      (highlights, deltas, anomalies) around the chart.
-    `),
-    suggestedFollowUps: z
-        .array(z.string())
-        .optional()
-        .describe(stringUtils.toDescription `
-      Follow-up question suggestions Genie produced. The host UI
-      renders these as clickable buttons; you do not need to list
-      them in your reply.
-    `),
-    error: z
-        .string()
-        .optional()
-        .describe(stringUtils.toDescription `
-      Genie-side error message if the request failed.
-    `),
-});
-const sendMessageSchema = z.object({
-    content: z.string().describe(stringUtils.toDescription `
-    Natural-language question to send to the Genie space.
-  `),
-    conversationId: z
-        .string()
-        .optional()
-        .describe(stringUtils.toDescription `
-      Optional Genie conversation id to continue an earlier thread.
-      Omit on the first call; pass the id returned in the previous
-      result's \`conversationId\` to follow up.
-    `),
-});
-const getConversationSchema = z.object({
-    alias: z.string().describe(stringUtils.toDescription `
-    Alias of the Genie space the conversation belongs to (matches
-    the key in the genie plugin's \`spaces\` config).
-  `),
-    conversationId: z.string().describe(stringUtils.toDescription `
-    Genie conversation id whose history to fetch.
-  `),
-});
-/** Per-attachment shape returned inside a stored Genie message. */
-const genieAttachmentSchema = z.object({
-    attachmentId: z.string().optional().describe(stringUtils.toDescription `
-    Genie attachment id; internal bookkeeping.
-  `),
-    query: z
-        .object({
-        title: z.string().optional().describe(stringUtils.toDescription `
-        Genie's title for the SQL, if any.
-      `),
-        description: z.string().optional().describe(stringUtils.toDescription `
-        Genie's prose description of the SQL, if any.
-      `),
-        query: z.string().optional().describe(stringUtils.toDescription `
-        SQL Genie generated and executed.
-      `),
-        statementId: z.string().optional().describe(stringUtils.toDescription `
-        Statement-execution id; internal bookkeeping.
-      `),
-    })
-        .optional()
-        .describe(stringUtils.toDescription `
-      SQL Genie attached to this message, if it ran any.
-    `),
-    text: z
-        .object({
-        content: z.string().optional().describe(stringUtils.toDescription `
-        Genie's natural-language answer text for this attachment.
-      `),
-    })
-        .optional()
-        .describe(stringUtils.toDescription `
-      Per-attachment text content (independent of the message-level
-      \`content\` field).
-    `),
-    suggestedQuestions: z
-        .array(z.string())
-        .optional()
-        .describe(stringUtils.toDescription `
-      Follow-up question suggestions Genie generated for this turn.
-    `),
-});
-/** Single message inside a Genie conversation history page. */
-const genieMessageSchema = z.object({
-    messageId: z.string().describe(stringUtils.toDescription `
-    Genie message id; internal bookkeeping.
-  `),
-    conversationId: z.string().describe(stringUtils.toDescription `
-    Conversation id this message belongs to.
-  `),
-    spaceId: z.string().describe(stringUtils.toDescription `
-    Genie space id this message belongs to.
-  `),
-    status: z.string().describe(stringUtils.toDescription `
-    Genie message status (\`COMPLETED\`, \`FAILED\`, etc.).
-  `),
-    content: z.string().describe(stringUtils.toDescription `
-    Outer message-level natural-language content Genie wrote.
-  `),
-    attachments: z
-        .array(genieAttachmentSchema)
-        .optional()
-        .describe(stringUtils.toDescription `
-      Attachments (SQL queries, text blocks, suggested follow-ups)
-      Genie produced for this message.
-    `),
-    error: z.string().optional().describe(stringUtils.toDescription `
-    Genie-side error attached to this message, if any.
-  `),
-});
+function extractStatementId(message) {
+    const top = message.query_result
+        ?.statement_id;
+    if (top)
+        return top;
+    for (const att of message.attachments ?? []) {
+        const id = att.query?.statement_id;
+        if (id)
+            return id;
+    }
+    return undefined;
+}
 /**
- * Output schema for the \`genie_get_conversation\` tool. Mirrors
- * AppKit's \`GenieConversationHistoryResponse\` so the model gets a
- * clear, typed view of prior messages instead of an opaque blob.
+ * Best-effort `writer.write`. The writer carries the unified flat
+ * event vocabulary directly - no translation layer - so
+ * subscribers narrow on `event.type` and read fields inline.
+ * Failures (downstream stream closed, cancelled request) are
+ * swallowed with a `warn` log so an in-flight Genie turn isn't
+ * taken down by a navigated-away client.
  */
-const genieGetConversationOutputSchema = z.object({
-    conversationId: z.string().describe(stringUtils.toDescription `
-    Conversation id you fetched.
-  `),
-    spaceId: z.string().describe(stringUtils.toDescription `
-    Genie space the conversation belongs to.
-  `),
-    messages: z.array(genieMessageSchema).describe(stringUtils.toDescription `
-    Messages in the conversation, oldest to newest. Each
-    \`message.content\` is Genie's natural-language answer for
-    that turn; attachments carry the SQL and follow-ups Genie
-    produced.
-  `),
-});
+async function safeWrite(writer, chunk) {
+    if (!writer)
+        return;
+    try {
+        await writer.write(chunk);
+    }
+    catch (err) {
+        log.warn("writer:error", {
+            error: err instanceof Error ? err.message : String(err),
+        });
+    }
+}
 /**
- * Default tool name for a wired Genie alias. The well-known `default`
- * alias collapses to `genie`; everything else gets a `genie_` prefix so
- * multiple spaces stay disambiguated when an agent has more than one
- * wired. Matches the `genie` / `genie_<alias>` naming used elsewhere in
- * dbx-tools AppKit demos.
+ * Lowercased placeholder strings we reject at the `ask_genie`
+ * boundary so the LLM doesn't spend a Genie round-trip on a
+ * non-question. Genie politely answers any of these with "Your
+ * request '...' does not relate to..." which is pure UI noise.
+ * Kept narrow on purpose - real questions sometimes start with
+ * one of these tokens, so we only match the FULL trimmed string.
  */
-export function defaultGenieToolName(alias) {
-    if (alias === "default")
-        return "genie";
-    return stringUtils.toIdentifierWithOptions({ distinct: true }, "genie", alias);
+const PLACEHOLDER_QUESTIONS = new Set([
+    "noop",
+    "no-op",
+    "skip",
+    "none",
+    "n/a",
+    "na",
+    "null",
+    "undefined",
+    "test",
+    "placeholder",
+]);
+/* ----------------------- conversation state ----------------------- */
+/**
+ * Estimated Genie conversation lifetime in seconds. Databricks
+ * publishes no official TTL on the conversation resource itself;
+ * community projects (e.g. the open-source Databricks Genie Bot)
+ * converge on 4 hours of inactivity as a safe operating window.
+ * Treat this as an estimate that gets *extended on every use* by
+ * re-setting the cache entry after each successful turn (sliding
+ * TTL via re-`set`). When the estimate ends up wrong (conversation
+ * deleted, expired upstream, cross-space referenced), the wrapper
+ * catches the SDK's `RESOURCE_DOES_NOT_EXIST`/404 and transparently
+ * starts a fresh conversation.
+ */
+const CONVERSATION_TTL_SEC = 4 * 60 * 60;
+/** Cache namespace prefix so coexisting Mastra caches don't collide. */
+const CONVERSATION_CACHE_NAMESPACE = "mastra:genie:conversation";
+/**
+ * Build the per-request {@link RequestContext} key the active
+ * Genie `conversation_id` lives under for `spaceId`. Scoped by
+ * space so an app calling two Genie spaces in one request keeps
+ * each conversation distinct (Genie conversation ids are
+ * space-scoped on the wire). The same `RequestContext` instance
+ * flows from the outer `genie` tool through to the inner
+ * `ask_genie` tool via Mastra, so writes on one side are visible
+ * on the other without an explicit shared ref.
+ */
+const conversationContextKey = (spaceId) => `mastra__genie_conversation__${spaceId}`;
+/**
+ * Read the active Genie `conversation_id` for `spaceId` off the
+ * per-request {@link RequestContext}. Returns `undefined` when no
+ * conversation has been started yet this request.
+ */
+function readContextConversationId(requestContext, spaceId) {
+    return requestContext.get(conversationContextKey(spaceId));
+}
+/**
+ * Write the active Genie `conversation_id` for `spaceId` onto the
+ * per-request {@link RequestContext}. Subsequent `ask_genie` calls
+ * in this request will reuse it; the wrapper's tail logic also
+ * reads it back out for the {@link GenieAgentResult}.
+ */
+function writeContextConversationId(requestContext, spaceId, conversationId) {
+    requestContext.set(conversationContextKey(spaceId), conversationId);
 }
+/* ------------------------- chart inventory ------------------------- */
 /**
- * Build one `sendMessage` tool per configured Genie alias plus a single
- * `getConversation` tool. Returns a record keyed by tool id, ready to
- * spread into an `Agent`'s `tools` map.
+ * Per-request {@link RequestContext} key the resolved chart
+ * inventory lives under. Keyed by `chartId`, the inventory is a
+ * `Map<string, ChartEvent>` carrying the full Echarts spec for
+ * every chart minted on this request - the same payload that
+ * goes out on the writer stream, kept in-process so output
+ * processors and downstream tools can resolve `[[chart:<id>]]`
+ * markers without re-running the planner or pulling from the
+ * writer stream.
  *
- * `config` must be the active plugin config; Genie's
- * `query_result` events are routed through
- * {@link emitChartWithPlanning} which uses it to resolve the
- * chart-planner's model.
+ * Shared across all Genie spaces because chart ids are minted
+ * via `commonUtils.shortId()` and are unique within a single
+ * request regardless of which space produced them.
  */
-export function buildGenieTools(opts) {
-    const tools = {};
-    for (const alias of opts.aliases) {
-        const id = defaultGenieToolName(alias);
-        tools[id] = createTool({
-            id,
-            description: stringUtils.toDescription `
-        Ask the Databricks Genie space "${alias}" a single
-        natural-language question. Genie translates it to SQL,
-        runs it, and returns \`genieAnswer\` (prose) plus
-        \`datasets[]\` (one entry per executed query, each with
-        a short \`chartId\`). Embed \`[[chart:<chartId>]]\` on
-        its own line at the position you want that data rendered
-        as an inline chart. Add interpretation around the chart
-        (deltas, anomalies, takeaways); do not paraphrase row
-        values.
-
-        Issue ONE focused question per user turn. Prefer
-        aggregated queries over raw-row queries for time-series
-        and distributions.
-      `,
-            inputSchema: sendMessageSchema,
-            outputSchema: genieToolOutputSchema,
-            execute: async ({ content, conversationId }, ctx) => {
-                const stream = opts.exports.sendMessage(alias, content, conversationId, {
-                    signal: opts.signal,
-                });
-                const requestContext = ctx
-                    ?.requestContext;
-                return drainGenieStream(stream, ctx.writer, {
-                    config: opts.config,
-                    ...(requestContext ? { requestContext } : {}),
-                });
-            },
+const CHART_INVENTORY_CONTEXT_KEY = "mastra__genie_chart_inventory__";
+/**
+ * Get the chart inventory map for this request, creating it on
+ * first access. Subsequent reads return the same map so callers
+ * mutate in place. The map is request-scoped (collected with the
+ * `RequestContext` at end of request), so there's no per-process
+ * leak.
+ */
+export function chartInventoryFromContext(requestContext) {
+    const existing = requestContext.get(CHART_INVENTORY_CONTEXT_KEY);
+    if (existing instanceof Map) {
+        return existing;
+    }
+    const fresh = new Map();
+    requestContext.set(CHART_INVENTORY_CONTEXT_KEY, fresh);
+    return fresh;
+}
+/**
+ * Stash a resolved chart on the request-scoped inventory so any
+ * subsequent code in this request (output processors validating
+ * `[[chart:<id>]]` markers, follow-up tools that want to chart
+ * the same dataset differently, etc.) can look it up by id.
+ * No-op when `requestContext` is missing.
+ */
+function recordChartInContext(requestContext, chart) {
+    if (!requestContext)
+        return;
+    chartInventoryFromContext(requestContext).set(chart.chartId, chart);
+}
+/**
+ * `userKey` for `CacheManager.getOrExecute` / `generateKey`. Genie
+ * conversations are scoped to a single user + space + thread, and
+ * `threadId` is already user-scoped (Mastra mints threads per
+ * `resourceId`), so a constant user key here is safe and keeps the
+ * cache key short.
+ */
+const CONVERSATION_USER_KEY = "mastra-genie";
+/**
+ * Build the canonical cache key for a `(spaceId, threadId)` pair.
+ * Returns `undefined` when `threadId` is missing - callers should
+ * skip caching entirely in that case (no Mastra memory wired up).
+ */
+async function conversationCacheKey(spaceId, threadId) {
+    if (!threadId)
+        return undefined;
+    return (await CacheManager.getInstance()).generateKey([CONVERSATION_CACHE_NAMESPACE, spaceId, threadId], CONVERSATION_USER_KEY);
+}
+/**
+ * Read the cached Genie conversation id for `(spaceId, threadId)`.
+ * Returns `undefined` on miss, on expiry, or when the cache layer
+ * is unhealthy - never throws. The TTL is renewed via re-`set`
+ * after each successful turn (see {@link saveCachedConversationId}).
+ */
+async function readCachedConversationId(cacheKey) {
+    if (!cacheKey)
+        return undefined;
+    try {
+        const v = await CacheManager.getInstanceSync().get(cacheKey);
+        return v ?? undefined;
+    }
+    catch (err) {
+        log.warn("conversation-cache:read-error", {
+            error: err instanceof Error ? err.message : String(err),
+        });
+        return undefined;
+    }
+}
+/**
+ * Persist the active conversation id under `cacheKey`, refreshing
+ * its TTL. Idempotent; no-op when `cacheKey` or `conversationId`
+ * is missing. Re-setting the same key acts as a sliding TTL: every
+ * turn that uses the conversation extends the window by another
+ * {@link CONVERSATION_TTL_SEC} seconds.
+ */
+async function saveCachedConversationId(cacheKey, conversationId) {
+    if (!cacheKey || !conversationId)
+        return;
+    try {
+        await CacheManager.getInstanceSync().set(cacheKey, conversationId, {
+            ttl: CONVERSATION_TTL_SEC,
+        });
+    }
+    catch (err) {
+        log.warn("conversation-cache:write-error", {
+            error: err instanceof Error ? err.message : String(err),
+        });
+    }
+}
+/** Force-evict a cached conversation id. Used on the stale-id recovery path. */
+async function evictCachedConversationId(cacheKey) {
+    if (!cacheKey)
+        return;
+    try {
+        await CacheManager.getInstanceSync().delete(cacheKey);
+    }
+    catch (err) {
+        log.warn("conversation-cache:delete-error", {
+            error: err instanceof Error ? err.message : String(err),
         });
     }
-    tools.genie_get_conversation = createTool({
-        id: "genie_get_conversation",
+}
+/**
+ * True when `err` is the SDK error Genie returns for a
+ * conversation id that no longer exists (deleted, expired upstream,
+ * or referenced from the wrong space). Matches the typed
+ * {@link ApiError} 404 / `RESOURCE_DOES_NOT_EXIST` shape first, then
+ * falls back to the lower-level {@link HttpError} 404, then to a
+ * loose message sniff for SDK shapes we haven't catalogued.
+ */
+function isConversationGoneError(err) {
+    if (err instanceof ApiError) {
+        if (err.statusCode === 404)
+            return true;
+        if (err.errorCode === "RESOURCE_DOES_NOT_EXIST")
+            return true;
+    }
+    if (err instanceof HttpError && err.code === 404)
+        return true;
+    if (err instanceof Error && /does not exist/i.test(err.message))
+        return true;
+    return false;
+}
+function buildAskGenieTool(deps) {
+    const { spaceId, client, writer, signal, resultSets, cacheKey } = deps;
+    return createTool({
+        id: "ask_genie",
         description: stringUtils.toDescription `
-      Fetch the full message history of a prior Genie conversation
-      by id. Use when the user references an earlier Genie thread
-      by id, or to inspect attachments / SQL from previous turns.
+      Send ONE focused natural-language question to the Genie
+      space and wait for the turn to complete. Returns the final
+      \`GenieMessage\` plus, when the turn ran SQL, the rows of
+      the resulting query as \`query_result_data\`. The
+      \`statement_id\` you reference in your final \`data\`
+      blocks lives at \`message.query_result.statement_id\` (or
+      the first attachment's \`query.statement_id\`). Wire
+      events (status, thinking, sql) stream to the user
+      automatically. Call multiple times to gather different
+      angles before composing the final response.
     `,
-        inputSchema: getConversationSchema,
-        outputSchema: genieGetConversationOutputSchema,
-        execute: async ({ alias, conversationId }) => {
-            return opts.exports.getConversation(alias, conversationId, opts.signal);
+        inputSchema: z.object({
+            question: z.string().min(1, "question is required"),
+        }),
+        outputSchema: z.object({
+            message: z.custom(),
+            query_result_data: z.custom().optional(),
+        }),
+        execute: async ({ question }, ctxRaw) => {
+            const ctx = ctxRaw;
+            const requestContext = ctx?.requestContext;
+            if (!requestContext) {
+                // Mastra always passes a `RequestContext` to tools when the
+                // parent agent received one. The outer Genie tool insists on
+                // it (it sources the user from there), so this only fires
+                // if a misconfigured caller invokes `ask_genie` directly.
+                throw new Error("ask_genie: missing requestContext (parent agent must propagate it)");
+            }
+            // Bounce placeholder / no-op questions BEFORE spending a Genie
+            // round-trip on them. The structuring pass occasionally pads
+            // out the tool loop with a fake `ask_genie("noop")` call,
+            // which Genie answers with "Your request 'noop' does not
+            // relate to..." - useless noise that shows up in the UI and
+            // eats one of the workspace's 5 questions/minute. Returning
+            // a clear error here surfaces the issue to the agent loop so
+            // the model corrects course instead of wasting a turn.
+            const trimmed = question.trim();
+            if (trimmed.length === 0 || PLACEHOLDER_QUESTIONS.has(trimmed.toLowerCase())) {
+                throw new Error(`ask_genie: refusing placeholder question "${question}" - ` +
+                    `call ask_genie only with a real natural-language question, ` +
+                    `or skip the call entirely`);
+            }
+            // Single turn of `genieEventChat`. Hoisted into a closure so
+            // we can re-run it after evicting a stale `conversation_id`
+            // without duplicating the event-loop body.
+            const runTurn = async () => {
+                const seedConversationId = readContextConversationId(requestContext, spaceId);
+                let finalMessage;
+                for await (const event of genieEventChat(spaceId, question, {
+                    workspaceClient: client,
+                    ...(seedConversationId ? { conversationId: seedConversationId } : {}),
+                    ...(signal ? { context: signal } : {}),
+                })) {
+                    await safeWrite(writer, event);
+                    // Wire events come in two flavors: the lifecycle `message`
+                    // event embeds the raw `GenieMessage` (read its
+                    // `conversation_id`), and the rest carry a flat
+                    // `conversation_id` field at the top level. The terminal
+                    // `result` event also carries the final `GenieMessage`
+                    // inline so we can capture the snapshot without re-reading
+                    // a buffered `message` event.
+                    const eventConversationId = event.type === "message"
+                        ? event.message.conversation_id
+                        : event.conversation_id;
+                    if (eventConversationId) {
+                        writeContextConversationId(requestContext, spaceId, eventConversationId);
+                    }
+                    if (event.type === "result") {
+                        finalMessage = event.message;
+                    }
+                }
+                if (!finalMessage) {
+                    throw new Error("Genie turn ended without a result event");
+                }
+                return finalMessage;
+            };
+            let finalMessage;
+            try {
+                finalMessage = await runTurn();
+            }
+            catch (err) {
+                // The seeded `conversation_id` was rejected by Genie - most
+                // commonly because it was deleted upstream, expired past
+                // Databricks' (undocumented) lifetime, or was minted in a
+                // different space. Drop both the cached id AND the
+                // per-request value so the retry calls `startConversation`,
+                // and try once more. Only retry when we *had* a seeded id -
+                // a fresh call that 404s shouldn't loop.
+                const seeded = readContextConversationId(requestContext, spaceId);
+                if (seeded && isConversationGoneError(err)) {
+                    log.warn("conversation-cache:stale, resetting", {
+                        spaceId,
+                        conversationId: seeded,
+                        error: err instanceof Error ? err.message : String(err),
+                    });
+                    await evictCachedConversationId(cacheKey);
+                    writeContextConversationId(requestContext, spaceId, undefined);
+                    finalMessage = await runTurn();
+                }
+                else {
+                    throw err;
+                }
+            }
+            // Refresh the cache entry on every successful turn. Re-setting
+            // the same key both persists newly-minted ids (cache miss path)
+            // and extends the TTL on active conversations (sliding window).
+            await saveCachedConversationId(cacheKey, readContextConversationId(requestContext, spaceId));
+            const statementId = extractStatementId(finalMessage);
+            let queryResultData;
+            if (statementId) {
+                const data = await fetchStatementData(client, statementId, signal);
+                if (data.rowCount > 0) {
+                    queryResultData = data;
+                    // Stash with this ask's `message_id` so the outer chart
+                    // loop can stamp downstream `chart` events with the
+                    // same id the wire events carry - keeps the chart in
+                    // the same `message_id` pill bucket on the host UI.
+                    resultSets.set(statementId, {
+                        data,
+                        messageId: finalMessage.message_id,
+                    });
+                }
+            }
+            return {
+                message: finalMessage,
+                ...(queryResultData ? { query_result_data: queryResultData } : {}),
+            };
+        },
+    });
+}
+function buildSpaceDescriptionTool(deps) {
+    const { spaceId, client, signal } = deps;
+    return createTool({
+        id: "get_space_description",
+        description: stringUtils.toDescription `
+      Return the Genie space's title, description, and warehouse id.
+      Cheap. Call once at the start of a turn to ground yourself
+      in what data the space covers.
+    `,
+        inputSchema: z.object({}),
+        outputSchema: z.object({
+            spaceId: z.string(),
+            title: z.string().optional(),
+            description: z.string().optional(),
+            warehouseId: z.string().optional(),
+        }),
+        execute: async () => {
+            const ctx = signal ? apiUtils.toContext(signal) : undefined;
+            const space = await client.genie.getSpace({ space_id: spaceId }, ctx);
+            return {
+                spaceId,
+                ...(space.title ? { title: space.title } : {}),
+                ...(space.description ? { description: space.description } : {}),
+                ...(space.warehouse_id ? { warehouseId: space.warehouse_id } : {}),
+            };
+        },
+    });
+}
+function buildSpaceSerializedTool(deps) {
+    const { spaceId, client, signal } = deps;
+    return createTool({
+        id: "get_space_serialized",
+        description: stringUtils.toDescription `
+      Return the full \`GenieSpace\` JSON for this space. Use only
+      when you need exact column / table identifiers
+      \`get_space_description\` doesn't expose. Larger payload, so
+      prefer the description tool when it's enough.
+    `,
+        inputSchema: z.object({}),
+        outputSchema: z.object({ space: z.unknown() }),
+        execute: async () => {
+            const ctx = signal ? apiUtils.toContext(signal) : undefined;
+            const space = await client.genie.getSpace({ space_id: spaceId }, ctx);
+            return { space };
         },
     });
-    return tools;
 }
+/* --------------------------- inner agent --------------------------- */
+const AGENT_INSTRUCTIONS = stringUtils.toDescription `
+  You orchestrate a Databricks Genie space. For every user
+  question:
+
+    1. Optionally call \`get_space_description\` to ground; reach
+       for \`get_space_serialized\` only when you need exact
+       column / table names the description doesn't expose.
+    2. Decompose the question into focused sub-questions (one per
+       distinct metric / dimension / time window) and call
+       \`ask_genie\` once per sub-question. Two to six calls is
+       typical for a non-trivial question; one call is fine when
+       the question is genuinely atomic.
+    3. Each \`ask_genie\` call returns the terminal
+       \`GenieMessage\`. When the turn ran SQL it also returns
+       \`query_result_data\` - the actual rows. The matching
+       \`statement_id\` is on
+       \`message.query_result.statement_id\` (or the first
+       attachment's \`query.statement_id\`). You will reference
+       that exact id in your final \`data\` blocks.
+    4. Produce a final structured summary as an ordered array
+       interleaving \`text\` paragraphs with \`data\` blocks.
+       INTERLEAVE: prose first, then the \`data\` block it
+       interprets, then the next prose / data pair. Never dump
+       all data at the end.
+    5. For every \`data\` block, supply the exact
+       \`statement_id\` you saw on the \`ask_genie\` response. A
+       short \`description\` ("compare quarterly revenue across
+       regions", "highlight the steep drop after position 5")
+       biases the chart-planner's choice of visual. Do NOT pick
+       chart types or axis labels - the host wraps each \`data\`
+       block in a chart automatically.
+    6. Each \`data\` block should be followed by a short
+       \`text\` interpretation (deltas, anomalies, takeaways).
+       Don't paraphrase numbers the visualization will already
+       show. Skip openers / closers. Plain prose, hyphens (not em
+       / en dashes), no emojis.
+`;
 /**
- * Drain the genie `sendMessage` AsyncGenerator into a flat result
- * the agent's calling LLM can reason about, while forwarding
- * progress and chart events to the host UI.
- *
- * Three streams of output happen in parallel:
- *
- * 1. {@link GenieProgress} pill events on the writer (`started`,
- *    `status`, `sql`, `suggested`, `error`) drive the loading
- *    pill in the chat bubble.
- * 2. `kind: "chart"` events on the writer (emitted via
- *    {@link emitChartWithPlanning}) carry the row payload from
- *    each Genie SQL statement and, on planner success, a
- *    follow-up event with the rendered Echarts spec. The host
- *    UI's `<ChartSlot>` merges the two by `chartId` and
- *    renders inline at the marker position the model picked.
- *    The data never reaches the LLM.
- * 3. The `DrainResult` returned to the LLM contains Genie's
- *    prose answer plus a `datasets[]` array of metadata
- *    (chartId, title, columns, rowCount, sql) the model uses
- *    to cite charts via `[[chart:<chartId>]]` markers.
+ * Boundary schema for the inner agent's structured output. Two
+ * tagged shapes only - text or data. The wrapper maps these onto
+ * the shared {@link GenieSummaryItem} (`string` / `visualize`)
+ * after charting; we don't redefine GenieSummaryItem here.
+ */
+const agentSummarySchema = z.object({
+    summary: z.array(z.discriminatedUnion("type", [
+        z.object({
+            type: z.literal("text"),
+            text: z.string(),
+        }),
+        z.object({
+            type: z.literal("data"),
+            statementId: z.string(),
+            title: z.string().optional(),
+            description: z.string().optional(),
+        }),
+    ])),
+});
+/**
+ * Build the calling agent's Genie tool. The returned Mastra tool
+ * runs end-to-end on each invocation:
  *
- * `query_result` and `message_result` events arrive in either
- * order; we buffer per-statement scratch keyed by `statementId`
- * so each half can fill in what it knows. The chart event
- * fires the moment `query_result` lands; the planner runs in
- * the background. We `Promise.allSettled` every planner promise
- * before returning so all chart work is attributed to the tool's
- * trace span and so the LLM's `datasets[]` includes every
- * chartId that has actually been queued.
+ *   1. Pull the per-request `WorkspaceClient` off
+ *      `ctx.requestContext` (stamped by `MastraServer` under
+ *      {@link MASTRA_USER_KEY}) and emit a `started` writer
+ *      event so the host UI shows progress immediately.
+ *   2. Spin up the inner Mastra agent + three tools, fresh per
+ *      call so the row cache stays invocation-scoped.
+ *   3. Run the agent with `structuredOutput` against
+ *      {@link agentSummarySchema}. Mastra's two-pass design keeps
+ *      the inner loop tools-only (no `response_format`), so the
+ *      Databricks Model Serving `response_format`+`tools`
+ *      collision never fires.
+ *   4. Walk the returned `[text|data][]`, map `text` items to
+ *      shared `GenieSummaryItem.string`, and chart every `data`
+ *      item in parallel via {@link runChartPlanner} to a
+ *      `GenieSummaryItem.visualize`. Items referencing a missing
+ *      `statementId` are dropped with a warn log; chart-planner
+ *      failures leave `dataset.chart` unset so the host UI falls
+ *      back to a table.
  */
-async function drainGenieStream(stream, writer, opts) {
-    const { config, requestContext } = opts;
-    let conversationId;
-    let genieAnswer;
-    let suggestedFollowUps;
-    let error;
-    // AppKit's `streamSendMessage` forwards every SDK `onProgress`
-    // callback verbatim - the same `EXECUTING_QUERY` can fire several
-    // times during a single poll loop. AppKit's other path,
-    // `streamGetMessage`, dedupes on the connector side; we mirror that
-    // behaviour here so the UI status pill doesn't flicker and we don't
-    // burn writer bytes on no-op events.
-    let lastStatus;
-    const scratchByStatementId = new Map();
-    const getScratch = (statementId) => {
-        let s = scratchByStatementId.get(statementId);
-        if (!s) {
-            s = { statementId, columns: [], rowCount: 0 };
-            scratchByStatementId.set(statementId, s);
-        }
-        return s;
-    };
-    /**
-     * Planner promises kicked off per `query_result`. Awaited
-     * (Promise.allSettled) before drainGenieStream returns so the
-     * Genie tool's trace span covers the chart work and the LLM's
-     * `datasets[]` accurately reflects every chartId that's been
-     * queued for rendering.
-     */
-    const plannerPromises = [];
-    const emit = async (event) => {
-        if (!writer)
-            return;
-        try {
-            await writer.write(event);
-        }
-        catch {
-            // ignore: downstream stream is no longer interested
-        }
-    };
-    for await (const event of stream) {
-        // Per-event raw payload for tuning the pill / answer pipeline
-        // against real Genie traffic. At `info` (the default) this is
-        // discarded for free; flip `LOG_LEVEL=debug` to see every
-        // raw wire event before the switch routes it through writer
-        // and DrainResult.
-        log.debug("event", { type: event.type, payload: event });
-        switch (event.type) {
-            case "message_start":
-                conversationId = event.conversationId;
-                await emit({
-                    kind: "started",
-                    conversationId: event.conversationId,
-                    messageId: event.messageId,
-                    spaceId: event.spaceId,
-                });
-                break;
-            case "status":
-                if (event.status === lastStatus)
-                    break;
-                lastStatus = event.status;
-                await emit({
-                    kind: "status",
-                    status: event.status,
-                    label: humanizeGenieStatus(event.status),
-                });
-                break;
-            case "query_result": {
-                const columns = (event.data?.manifest?.schema?.columns ?? []).map((c) => c.name);
-                const dataArray = (event.data?.result?.data_array ?? []);
-                const rows = genieRowsToObjects(columns, dataArray);
-                const scratch = getScratch(event.statementId);
-                // emitChartWithPlanning emits the dataset event immediately
-                // and kicks off the chart-planner agent in the background.
-                // It returns the chartId synchronously; the plannerPromise
-                // is awaited at end-of-stream so chart work shows up under
-                // this tool's trace span.
-                const { chartId, plannerPromise } = await emitChartWithPlanning({
-                    ...(writer ? { writer } : {}),
-                    config,
-                    ...(requestContext ? { requestContext } : {}),
-                    title: scratch.title ?? `Genie query`,
-                    ...(scratch.description ? { description: scratch.description } : {}),
-                    data: rows,
-                });
-                scratch.chartId = chartId;
-                scratch.columns = columns;
-                scratch.rowCount = rows.length;
-                plannerPromises.push(plannerPromise);
-                log.debug("query_result", {
-                    statementId: event.statementId,
-                    chartId,
-                    rows: rows.length,
-                    columns,
-                });
-                break;
+export function createGenieTool(opts) {
+    const { spaceId, config, toolId = "genie", toolDescription = stringUtils.toDescription `
+      Ask a question about the Databricks Genie space.
+
+      Returns \`{ summary: SummaryItem[] }\` where each item is
+      one of:
+
+      - \`{ type: "string", text }\` - prose to weave into your
+        reply verbatim or paraphrase.
+      - \`{ type: "visualize", statementId, title?, description?,
+        dataset: { data: { columns, rows, rowCount },
+        chart?: { chartId, chartType } } }\` - a chartable result
+        set. When \`dataset.chart\` is present the chart is ALREADY
+        rendered and queued for inline display; embed the marker
+        \`[[chart:<chartId>]]\` on its own line at the position
+        you want it to appear and the host UI drops the rendered
+        chart in. Re-use the chartId verbatim - do NOT call
+        \`render_data\` for the same dataset (it would render the
+        same chart a second time and stall your stream). Only
+        fall back to \`render_data\` when \`dataset.chart\` is
+        missing (chart-planner failed) AND you genuinely need a
+        picture; otherwise present the data inline as prose or a
+        short table.
+    `, maxSteps = DEFAULT_MAX_STEPS, } = opts;
+    return createTool({
+        id: toolId,
+        description: toolDescription,
+        inputSchema: z.object({
+            question: z.string().describe(stringUtils.toDescription `
+        Natural-language question about the data in this Genie
+        space. Phrase it from the user's perspective; the agent
+        decomposes it internally.
+      `),
+        }),
+        outputSchema: z.custom(),
+        execute: async (input, ctxRaw) => {
+            const ctx = ctxRaw;
+            const requestContext = ctx?.requestContext;
+            if (!requestContext) {
+                throw new Error("genie: missing requestContext (MastraServer must stamp MASTRA_USER_KEY)");
             }
-            case "message_result":
-                genieAnswer = event.message.content;
-                for (const attachment of event.message.attachments ?? []) {
-                    const sqlText = attachment.query?.query;
-                    const stmtId = attachment.query?.statementId;
-                    if (stmtId) {
-                        const scratch = getScratch(stmtId);
-                        if (sqlText)
-                            scratch.sql = sqlText;
-                        if (attachment.query?.title)
-                            scratch.title = attachment.query.title;
-                        if (attachment.query?.description) {
-                            scratch.description = attachment.query.description;
-                        }
-                    }
-                    if (sqlText) {
-                        await emit({
-                            kind: "sql",
-                            sql: sqlText,
-                            title: attachment.query?.title,
-                            description: attachment.query?.description,
-                            statementId: stmtId,
-                        });
-                    }
-                    if (attachment.text?.content) {
-                        await emit({ kind: "text", content: attachment.text.content });
-                    }
-                    if (attachment.suggestedQuestions?.length) {
-                        // Last attachment with suggestions wins (same merge rule
-                        // the UI uses via `collectSuggestions`); keeping just one
-                        // copy per turn caps token usage.
-                        suggestedFollowUps = attachment.suggestedQuestions;
-                        await emit({
-                            kind: "suggested",
-                            questions: attachment.suggestedQuestions,
-                        });
-                    }
+            const user = requestContext.get(MASTRA_USER_KEY);
+            if (!user) {
+                throw new Error("genie: no user on requestContext (MASTRA_USER_KEY not set)");
+            }
+            const client = user.executionContext.client;
+            const writer = ctx?.writer;
+            const signal = ctx?.abortSignal;
+            const threadId = requestContext.get(MASTRA_THREAD_ID_KEY);
+            // Fire the lifecycle `started` event before any LLM /
+            // network round-trip so the host UI can pop a "Thinking..."
+            // pill the instant the model decides to delegate. The wire
+            // `conversation_id` / `message_id` aren't known yet (no
+            // Genie call has been made) and ride as `undefined` -
+            // subscribers that need them watch the later
+            // `message` / `result` wire events for the real ids.
+            const startedEvent = {
+                type: "started",
+                spaceId,
+                content: input.question,
+            };
+            await safeWrite(writer, startedEvent);
+            const resultSets = new Map();
+            // Seed the active Genie `conversation_id` onto
+            // `RequestContext` from AppKit's `CacheManager` when a Mastra
+            // `threadId` is present so multi-turn chats reuse the same
+            // Genie conversation (and Genie's accumulated context) across
+            // separate tool invocations. The same `RequestContext` flows
+            // to the inner `ask_genie` tool via Mastra, which reads and
+            // updates the same slot as Genie hands out / rotates ids.
+            // Cache misses, threads without memory, and unhealthy cache
+            // storage all leave the slot unset, which makes `ask_genie`
+            // call `startConversation` and mint a fresh id (then cache
+            // it).
+            const cacheKey = await conversationCacheKey(spaceId, threadId);
+            const cachedConversationId = await readCachedConversationId(cacheKey);
+            if (cachedConversationId) {
+                writeContextConversationId(requestContext, spaceId, cachedConversationId);
+            }
+            const innerDeps = {
+                spaceId,
+                client,
+                ...(writer ? { writer } : {}),
+                ...(signal ? { signal } : {}),
+                resultSets,
+                ...(cacheKey ? { cacheKey } : {}),
+            };
+            const tools = {
+                ask_genie: buildAskGenieTool(innerDeps),
+                get_space_description: buildSpaceDescriptionTool({
+                    spaceId,
+                    client,
+                    ...(signal ? { signal } : {}),
+                }),
+                get_space_serialized: buildSpaceSerializedTool({
+                    spaceId,
+                    client,
+                    ...(signal ? { signal } : {}),
+                }),
+            };
+            // Resolve the model config once for this request so we can
+            // share it with the structuring pass below. The agent's
+            // `model` field accepts a function form for per-request
+            // resolution, but `structuredOutput.model` requires a
+            // static `MastraModelConfig`, and we need both to be on
+            // the same Databricks endpoint with the same OBO-scoped
+            // headers. Calling `buildModel` here (inside `execute`)
+            // keeps user scoping correct because `requestContext`
+            // already reflects the active request's user.
+            const resolvedModel = await buildModel(config, requestContext);
+            const agent = new Agent({
+                id: `genie__${spaceId}`,
+                name: `Genie (${spaceId})`,
+                description: stringUtils.toDescription `
+          Inner orchestrator for the "${spaceId}" Genie space.
+          Asks Genie one focused sub-question at a time and
+          returns an interleaved [text|data] summary.
+        `,
+                instructions: AGENT_INSTRUCTIONS,
+                model: resolvedModel,
+                tools,
+            });
+            // Mastra's `structuredOutput` operates in one of two modes
+            // based on whether `model` is set:
+            //   - "direct"    (no model)     -> the schema is enforced
+            //                                   in the SAME LLM call as
+            //                                   the agent loop, by
+            //                                   adding `response_format`
+            //                                   alongside `tools`.
+            //                                   Databricks Model Serving
+            //                                   rejects that combination
+            //                                   with `INVALID_PARAMETER_VALUE:
+            //                                   Cannot specify both
+            //                                   response_format and tools
+            //                                   in the same request.`
+            //   - "processor" (model passed) -> the main loop carries
+            //                                   tools and NO
+            //                                   `response_format`; a
+            //                                   separate, tool-free
+            //                                   structuring agent
+            //                                   re-prompts the model
+            //                                   with `response_format`
+            //                                   to coerce the agent's
+            //                                   final text into the
+            //                                   schema.
+            // We use "processor" mode but ALSO set
+            // `jsonPromptInjection: true`. Mastra's structuring agent
+            // calls `.stream(...)` under the hood, and Databricks Model
+            // Serving rejects `response_format` together with streaming
+            // (`INVALID_PARAMETER_VALUE: Structured output is not
+            // currently supported with streaming.`). Prompt injection
+            // sidesteps that by embedding the JSON Schema in the
+            // structuring agent's system prompt instead of sending
+            // `response_format`. `errorStrategy: "warn"` keeps a
+            // structuring failure from escaping as an unhandled
+            // promise rejection: it logs and leaves `result.object`
+            // undefined, which we surface as a clean error in
+            // {@link GenieAgentResult}.
+            const agentResult = await agent.generate(input.question, {
+                requestContext,
+                maxSteps,
+                structuredOutput: {
+                    schema: agentSummarySchema,
+                    model: resolvedModel,
+                    jsonPromptInjection: true,
+                    errorStrategy: "warn",
+                },
+                ...(signal ? { abortSignal: signal } : {}),
+            });
+            const submission = agentResult.object;
+            if (!submission) {
+                const message = "Genie agent returned no structured summary";
+                log.warn("agent:no-summary", { spaceId });
+                const finalConversationId = readContextConversationId(requestContext, spaceId);
+                return {
+                    spaceId,
+                    summary: [],
+                    ...(finalConversationId ? { conversationId: finalConversationId } : {}),
+                    error: message,
+                };
+            }
+            // Lifecycle hook: the agent + structuring pass are done.
+            // Emit one `summary` event with the structured-item counts
+            // so the host UI can transition from "thinking" to
+            // "charting" and seed N chart skeletons before the
+            // per-chart `chart` events arrive. We can't fire this
+            // EARLIER (i.e. when the structuring pass starts) because
+            // Mastra runs the inner loop + structuring pass together
+            // inside `agent.generate(...)` with no observable boundary
+            // between them.
+            const textItemCount = submission.summary.filter((i) => i.type === "text").length;
+            const dataItemCount = submission.summary.length - textItemCount;
+            const summaryEvent = {
+                type: "summary",
+                spaceId,
+                items: submission.summary.length,
+                textItems: textItemCount,
+                dataItems: dataItemCount,
+            };
+            await safeWrite(writer, summaryEvent);
+            // Chart every `data` item in parallel; map `text` items to
+            // the shared `string` summary variant verbatim. Missing
+            // statement ids are dropped (the agent referenced something
+            // that never came back from `ask_genie`), planner failures
+            // leave `dataset.chart` unset so the host UI falls back to
+            // a table render. Each successfully planned chart pushes a
+            // `chart` writer event so the UI can fade in the rendered
+            // chart slot the moment its planner returns rather than
+            // waiting for the entire batch to finish.
+            const hydrated = await Promise.all(submission.summary.map(async (item) => {
+                if (item.type === "text") {
+                    return { type: "string", text: item.text };
                 }
-                break;
-            case "error":
-                error = event.error;
-                await emit({ kind: "error", error: event.error });
-                break;
-            default:
-                break;
-        }
-    }
-    // Wait for all chart planners to settle before returning so the
-    // tool's trace span covers chart work and the LLM's
-    // `datasets[]` reflects only chartIds the client has actually
-    // received writer events for. Failures in `emitChartWithPlanning`
-    // are already swallowed inside the helper, so this never
-    // throws.
-    log.debug("planners:awaiting", { count: plannerPromises.length });
-    await Promise.allSettled(plannerPromises);
-    log.debug("planners:settled", { count: plannerPromises.length });
-    // Build the LLM-bound `datasets[]` from scratch entries that
-    // actually ran a query (chartId is assigned at `query_result`
-    // time). Entries that only saw `message_result` metadata
-    // without a row payload are skipped.
-    const datasets = [];
-    for (const scratch of scratchByStatementId.values()) {
-        if (!scratch.chartId)
-            continue;
-        datasets.push({
-            chartId: scratch.chartId,
-            ...(scratch.title ? { title: scratch.title } : {}),
-            ...(scratch.description ? { description: scratch.description } : {}),
-            columns: scratch.columns,
-            rowCount: scratch.rowCount,
-            ...(scratch.sql ? { sql: scratch.sql } : {}),
-        });
-    }
-    log.debug("drain:return", {
-        conversationId,
-        hasAnswer: typeof genieAnswer === "string",
-        answerLength: genieAnswer?.length ?? 0,
-        chartIds: datasets.map((d) => d.chartId),
-        suggestedCount: suggestedFollowUps?.length ?? 0,
-        error,
+                const entry = resultSets.get(item.statementId);
+                if (!entry) {
+                    log.warn("data:missing-statement", {
+                        statementId: item.statementId,
+                    });
+                    return undefined;
+                }
+                const { data, messageId } = entry;
+                let dataset = { data };
+                try {
+                    const planned = await runChartPlanner({
+                        config,
+                        requestContext,
+                        title: item.title ?? "Genie result",
+                        ...(item.description ? { description: item.description } : {}),
+                        data: data.rows,
+                        ...(signal ? { signal } : {}),
+                    });
+                    const chartId = commonUtils.shortId();
+                    // Slim chart reference for the LLM-bound result: just
+                    // `chartId` + `chartType`. The full Echarts spec goes
+                    // to the UI via the writer event AND into the
+                    // request-scoped chart inventory below; the model
+                    // only needs the id to place `[[chart:<id>]]`.
+                    dataset = {
+                        data,
+                        chart: {
+                            chartId,
+                            chartType: planned.chartType,
+                        },
+                    };
+                    const chartEvent = {
+                        type: "chart",
+                        chartId,
+                        statementId: item.statementId,
+                        messageId,
+                        ...(item.title ? { title: item.title } : {}),
+                        ...(item.description ? { description: item.description } : {}),
+                        data: data.rows,
+                        option: planned.option,
+                    };
+                    await safeWrite(writer, chartEvent);
+                    // Stash the resolved chart on the per-request
+                    // `RequestContext` so downstream code in the same
+                    // request (output processors, follow-up tool calls,
+                    // any post-run hook) can look up the full spec by
+                    // `chartId` without re-fetching or re-planning.
+                    recordChartInContext(requestContext, chartEvent);
+                }
+                catch (err) {
+                    const errorMessage = err instanceof Error ? err.message : String(err);
+                    log.warn("chart:error", {
+                        statementId: item.statementId,
+                        messageId,
+                        error: errorMessage,
+                    });
+                    // Surface the chart-planner failure as a writer event
+                    // stamped with the same `messageId` the rest of this
+                    // ask's wire events carry, so the host UI groups the
+                    // failure into the same pill bucket and can surface
+                    // a "couldn't render chart" note next to the table
+                    // fallback instead of silently dropping the chart.
+                    const errorEvent = {
+                        type: "error",
+                        spaceId,
+                        messageId,
+                        error: `chart-planner: ${errorMessage}`,
+                    };
+                    await safeWrite(writer, errorEvent);
+                }
+                return {
+                    type: "visualize",
+                    statementId: item.statementId,
+                    ...(item.title ? { title: item.title } : {}),
+                    ...(item.description ? { description: item.description } : {}),
+                    dataset,
+                };
+            }));
+            const summary = hydrated.filter((x) => x !== undefined);
+            log.info("genie:done", {
+                spaceId,
+                items: summary.length,
+                statementsCharted: summary.filter((s) => s.type === "visualize" && s.dataset.chart).length,
+            });
+            const finalConversationId = readContextConversationId(requestContext, spaceId);
+            return {
+                spaceId,
+                summary,
+                ...(finalConversationId ? { conversationId: finalConversationId } : {}),
+            };
+        },
     });
-    return {
-        ...(conversationId ? { conversationId } : {}),
-        ...(genieAnswer ? { genieAnswer } : {}),
-        ...(datasets.length > 0 ? { datasets } : {}),
-        ...(suggestedFollowUps ? { suggestedFollowUps } : {}),
-        ...(error ? { error } : {}),
-    };
 }
+/* --------------------- multi-alias surface --------------------- */
 /**
- * Convert Genie's `data_array` (column-positional `string | null`
- * tuples) into plain JS row objects keyed by column name. Numeric
- * strings are coerced to numbers so the chart-planner picks
- * `value` axes instead of `category` axes; everything else passes
- * through verbatim. `null` becomes `null`.
+ * Default tool id for a wired Genie alias. The well-known
+ * `default` alias collapses to `genie`; every other alias gets a
+ * `genie_` prefix so multi-space registrations stay
+ * disambiguated.
  */
-function genieRowsToObjects(columns, dataArray) {
-    const out = [];
-    for (const row of dataArray) {
-        const obj = {};
-        columns.forEach((col, i) => {
-            const cell = row[i] ?? null;
-            obj[col] = coerceCell(cell);
-        });
-        out.push(obj);
+export function defaultGenieToolName(alias) {
+    if (alias === DEFAULT_GENIE_ALIAS)
+        return "genie";
+    return stringUtils.toIdentifierWithOptions({ distinct: true }, "genie", alias);
+}
+/**
+ * Normalize the {@link GenieSpacesConfig} record. Bare-string
+ * entries (`{ default: "01ef..." }`) get wrapped as
+ * `{ spaceId: "01ef..." }`; object entries pass through unchanged.
+ * `undefined` and empty-string values are dropped so callers can
+ * pass `process.env.X` directly (matches AppKit `genie()`'s
+ * defensive treatment of unset env vars).
+ */
+export function normalizeGenieSpaces(spaces) {
+    if (!spaces)
+        return {};
+    const out = {};
+    for (const [alias, value] of Object.entries(spaces)) {
+        if (value === undefined)
+            continue;
+        if (typeof value === "string") {
+            if (!value)
+                continue;
+            out[alias] = { spaceId: value };
+            continue;
+        }
+        if (!value.spaceId)
+            continue;
+        out[alias] = value;
     }
     return out;
 }
-/** Best-effort numeric coercion for Genie's all-strings cells. */
-function coerceCell(cell) {
-    if (cell === null)
-        return null;
-    // Anchored to keep `12.5px` / `123abc` as strings; only fully
-    // numeric values become JS numbers.
-    if (/^-?\d+(\.\d+)?$/.test(cell)) {
-        const n = Number(cell);
-        if (Number.isFinite(n))
-            return n;
+/**
+ * Discover Genie space aliases from every supported source and
+ * merge them into a single record. Precedence (highest first):
+ *
+ *   1. {@link MastraPluginConfig.genieSpaces} on the `mastra(...)`
+ *      call. Explicit Mastra wiring always wins so users can
+ *      override AppKit's defaults per-agent.
+ *   2. AppKit `genie({ spaces: { ... } })` plugin instance. Lets
+ *      users keep using the existing AppKit config format
+ *      (`genie({ spaces: { sales: "...", ops: "..." } })`)
+ *      without restating the same record on the Mastra plugin.
+ *      Read off the live plugin instance via a structural cast
+ *      since `Plugin.config` is TS-protected (not runtime-private).
+ *   3. `DATABRICKS_GENIE_SPACE_ID` env var (registered under the
+ *      well-known `default` alias). Matches the AppKit `genie()`
+ *      plugin's fallback behavior so a bare `mastra()` + `genie()`
+ *      pair just works.
+ *
+ * Aliases collide cleanly: a higher-precedence source's value
+ * replaces a lower one's wholesale. Sources that contribute zero
+ * aliases (or contribute only `undefined` / empty entries) are
+ * silently ignored.
+ */
+export function resolveGenieSpaces(config, context) {
+    const merged = {};
+    // Source 3 (lowest precedence): env var.
+    const envSpaceId = process.env["DATABRICKS_GENIE_SPACE_ID"];
+    if (envSpaceId) {
+        merged[DEFAULT_GENIE_ALIAS] = { spaceId: envSpaceId };
     }
-    return cell;
+    // Source 2: AppKit `genie()` plugin instance config. Use a
+    // structural cast - `Plugin.config` is `protected` in TS only,
+    // and the runtime layout is plain object property access.
+    const geniePlugin = appkitUtils.instance(context, genie);
+    if (geniePlugin) {
+        const pluginSpaces = geniePlugin
+            .config?.spaces;
+        if (pluginSpaces) {
+            Object.assign(merged, normalizeGenieSpaces(pluginSpaces));
+        }
+    }
+    // Source 1 (highest precedence): explicit Mastra wiring.
+    if (config.genieSpaces) {
+        Object.assign(merged, normalizeGenieSpaces(config.genieSpaces));
+    }
+    return merged;
 }
 /**
- * Toolkit provider built from a live AppKit `GeniePlugin` instance.
- * Returned by {@link buildGenieProvider} so that
- * `plugins.genie?.toolkit()` inside an agent's `tools(plugins)` callback
- * resolves to the streaming-aware {@link buildGenieTools} record instead
- * of the AppKit default (which does one blocking call per tool with no
- * mid-flight events).
- *
- * The returned `toolkit()` reads alias names off the plugin's
- * `getAgentTools()` registry (each entry is `${alias}.sendMessage` or
- * `${alias}.getConversation`), then mints one `sendMessage` tool per
- * alias plus a shared `getConversation`. `sendMessage` / `getConversation`
- * are bound back to the plugin instance so they keep their `this`
- * (they are class methods, not free functions).
+ * Build one Mastra tool per configured Genie space. Each tool is
+ * a thin {@link createGenieTool} wrapper with the alias-derived
+ * id and a hint-flavored description so the calling LLM knows
+ * which space covers what data.
  *
- * `_opts` is accepted but unused for now - the streaming tools are an
- * all-or-nothing bundle. Wire `only` / `except` / `prefix` / `rename`
- * later if a caller needs them.
+ * Returns a record keyed by tool id, ready to spread into an
+ * `Agent`'s `tools` map (or surfaced via
+ * `plugins.genie?.toolkit()`).
  */
-export function buildGenieProvider(plugin, opts) {
+export function buildGenieTools(opts) {
+    const normalized = normalizeGenieSpaces(opts.spaces);
+    const tools = {};
+    for (const [alias, space] of Object.entries(normalized)) {
+        const id = defaultGenieToolName(alias);
+        const toolDescription = stringUtils.toDescription `
+      Delegate a natural-language data question to the
+      Databricks Genie space "${alias}"${space.hint ? ` (${space.hint})` : ""}.
+      Returns an ordered (text | dataset)[] summary the host UI
+      renders inline; datasets carry the rows and a
+      pre-rendered Echarts spec when the chart-planner
+      succeeded. Progress events (status, SQL, row counts,
+      charts) stream to the UI automatically.
+    `;
+        tools[id] = createGenieTool({
+            spaceId: space.spaceId,
+            config: opts.config,
+            toolId: id,
+            toolDescription,
+        });
+    }
+    return tools;
+}
+/**
+ * Plugin-toolkit adapter so the `plugins.genie?.toolkit()` lookup
+ * inside an agent's `tools(plugins)` callback returns the
+ * Genie agent-backed tools instead of throwing on missing plugin.
+ * Mirrors AppKit's `PluginToolkitProvider` shape.
+ */
+export function buildGenieToolkitProvider(opts) {
     return {
         toolkit(_opts) {
-            const aliases = extractGenieAliases(plugin);
-            return buildGenieTools({
-                aliases,
-                exports: {
-                    sendMessage: plugin.sendMessage.bind(plugin),
-                    getConversation: plugin.getConversation.bind(plugin),
-                },
-                config: opts.config,
-            });
+            return buildGenieTools(opts);
         },
     };
 }
 /**
- * Pull the configured space aliases out of a live AppKit `GeniePlugin`.
- * Reads them off `getAgentTools()` (public API) so we don't poke at the
- * `protected config.spaces` field: the plugin registers tools named
- * `${alias}.sendMessage` / `${alias}.getConversation`, so the unique
- * set of name prefixes is the alias list.
- */
-function extractGenieAliases(plugin) {
-    const aliases = new Set();
-    for (const t of plugin.getAgentTools()) {
-        const dot = t.name.indexOf(".");
-        if (dot > 0)
-            aliases.add(t.name.slice(0, dot));
-    }
-    return [...aliases];
-}
-/**
- * Convert raw Genie status codes (`FETCHING_METADATA`, `ASKING_AI`,
- * `EXECUTING_QUERY`, `COMPLETED`, ...) into short, sentence-cased
- * labels safe to drop straight into a UI pill. Unknown codes are
- * lower-cased with underscores stripped so new states still render.
+ * Returns `true` when at least one Genie space is reachable
+ * through {@link resolveGenieSpaces} - either via
+ * {@link MastraPluginConfig.genieSpaces}, the AppKit `genie()`
+ * plugin instance, or the `DATABRICKS_GENIE_SPACE_ID` env var.
+ *
+ * Cheap to call from `resolveProvider` to short-circuit `genie`
+ * lookups when nothing is wired, so the `plugins.genie` lookup
+ * still resolves to `undefined` (matching AppKit's
+ * absent-plugin semantics) when neither source is configured.
  */
-function humanizeGenieStatus(status) {
-    switch (status) {
-        case "FETCHING_METADATA":
-            return "Fetching metadata";
-        case "ASKING_AI":
-            return "Asking Genie";
-        case "EXECUTING_QUERY":
-            return "Running SQL query";
-        case "COMPLETED":
-            return "Completed";
-        case "FAILED":
-            return "Failed";
-        default:
-            return [
-                ...stringUtils.tokenizeWithOptions({ capitalize: true, lowerCase: true }, status),
-            ].join(" ");
-    }
+export function hasAnyGenieSpaces(config, context) {
+    return Object.keys(resolveGenieSpaces(config, context)).length > 0;
 }