@dbx-tools/appkit-mastra 0.1.25 → 0.1.27

package/dist/src/genie.js

@@ -1,122 +1,94 @@
 /**
- * Genie agent for Mastra.
+ * Genie tools for Mastra.
  *
- * 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:
+ * Each configured Genie space is surfaced as a small set of flat
+ * Mastra tools the calling agent can drive directly - no inner
+ * orchestrator agent. The central agent decomposes user questions,
+ * picks which Genie space to ask, and composes the final reply.
  *
- *   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}.
+ * Per-space tools (suffixed with `_<alias>` for non-default
+ * aliases):
  *
- * The inner agent talks to Genie directly via
- * `@dbx-tools/genie` (`genieEventChat`) and the workspace
- * `statementExecution.getStatement` API. AppKit's stock `genie`
- * plugin is honored only for its `spaces` config so existing
- * AppKit-style wiring keeps working without change.
+ *   - `ask_genie`: drives one `genieEventChat` turn and
+ *     forwards every wire event (status, thinking, sql, rows)
+ *     through `ctx.writer` for streaming UI updates. Returns
+ *     the terminal `GenieMessage` only - rows are NOT fetched
+ *     eagerly.
+ *   - `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.
+ *
+ * Space-agnostic shared tools (registered once, regardless of
+ * how many spaces are wired):
+ *
+ *   - `get_statement`: opt-in lookup of a Genie statement's rows
+ *     by `statement_id` (with a row `limit`). The agent calls
+ *     this only when it needs to read values to reason about
+ *     them; if the data is just being displayed, it embeds a
+ *     `[data:<statement_id>]` marker in prose instead and lets
+ *     the host UI resolve it.
+ *   - `prepare_chart`: mints a short `chartId`, kicks off a
+ *     background task that fetches the statement's rows and
+ *     runs the chart-planner, and caches the resolved Echarts
+ *     spec under the id (1h TTL). Returns the `chartId`
+ *     synchronously so the agent embeds `[chart:<chartId>]`
+ *     markers in prose without blocking on chart generation;
+ *     the host UI fetches the cached chart by id once it's
+ *     ready (see {@link fetchChart}).
+ *
+ * Each tool's `execute` pulls the per-request
+ * {@link WorkspaceClient} off `ctx.requestContext` (stamped by
+ * `MastraServer` under {@link MASTRA_USER_KEY}) and the per-call
+ * `writer` / `abortSignal` off `ctx`, so the tools are stateless
+ * across requests and the central agent owns the loop.
+ *
+ * The tools talk to Genie directly via `@dbx-tools/genie`
+ * (`genieEventChat`); statement-row fetching is delegated to
+ * {@link fetchStatementData} from `./statement.js`, which wraps
+ * the workspace `statementExecution.getStatement` API. AppKit's
+ * stock `genie` plugin is honored only for its `spaces` config
+ * so existing AppKit-style wiring keeps working without change.
+ *
+ * Suggested orchestration prompt for the central agent lives in
+ * {@link GENIE_INSTRUCTIONS}; compose it into the agent's own
+ * `instructions` when you want the canonical "how to drive the
+ * Genie tools" guidance.
  */
 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 { ChartSchema, } 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 { runChartPlanner } from "./chart.js";
+import { chartPlannerRequestSchema, prepareChart } from "./chart.js";
 import { MASTRA_USER_KEY } from "./config.js";
-import { buildModel } from "./model.js";
+import { fetchStatementData } from "./statement.js";
 import { safeWrite } from "./writer.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";
 /**
- * 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).
+ * Pull the per-request {@link WorkspaceClient} off the active
+ * `RequestContext`. The Mastra plugin's server middleware stamps
+ * a {@link User} on the context under {@link MASTRA_USER_KEY}
+ * for every request; tools fail loudly when it's missing because
+ * that means the Mastra plugin isn't running (e.g. a tool was
+ * invoked outside the chat route).
  */
-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;
+function requireClient(ctx, toolId) {
+    const requestContext = ctx?.requestContext;
+    if (!requestContext) {
+        throw new Error(`${toolId}: missing requestContext (MastraServer must stamp MASTRA_USER_KEY)`);
     }
-    return cell;
-}
-/**
- * Fetch a single Genie statement's rows via the Statement
- * Execution API and reshape into the shared
- * {@link GenieDatasetData} shape (column array + row records).
- */
-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,
-    };
-}
-/**
- * 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).
- */
-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;
+    const user = requestContext.get(MASTRA_USER_KEY);
+    if (!user) {
+        throw new Error(`${toolId}: no user on requestContext (MASTRA_USER_KEY not set)`);
     }
-    return undefined;
+    return { client: user.executionContext.client, requestContext };
 }
 /**
  * Lowercased placeholder strings we reject at the `ask_genie`
@@ -147,22 +119,30 @@ const PLACEHOLDER_QUESTIONS = new Set([
  * 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
+ * deleted, expired upstream, cross-space referenced), `ask_genie`
  * 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";
+/**
+ * `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 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.
+ * flows from the central agent through to every `ask_genie`
+ * invocation, so writes on one call are visible on the next
+ * without an explicit shared ref.
  */
 const conversationContextKey = (spaceId) => `mastra__genie_conversation__${spaceId}`;
 /**
@@ -176,64 +156,11 @@ function readContextConversationId(requestContext, 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}.
+ * in this request will reuse it.
  */
 function writeContextConversationId(requestContext, spaceId, conversationId) {
     requestContext.set(conversationContextKey(spaceId), conversationId);
 }
-/* ------------------------- chart inventory ------------------------- */
-/**
- * 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.
- *
- * 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.
- */
-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
@@ -298,6 +225,20 @@ async function evictCachedConversationId(cacheKey) {
         });
     }
 }
+/**
+ * Lazy-seed the active Genie `conversation_id` for `spaceId` from
+ * the cross-request cache onto the per-request `RequestContext`.
+ * No-op when the slot is already populated (subsequent
+ * `ask_genie` calls in the same turn) so we hit the cache at most
+ * once per request per space.
+ */
+async function ensureConversationSeeded(requestContext, spaceId, cacheKey) {
+    if (readContextConversationId(requestContext, spaceId))
+        return;
+    const cached = await readCachedConversationId(cacheKey);
+    if (cached)
+        writeContextConversationId(requestContext, spaceId, cached);
+}
 /**
  * True when `err` is the SDK error Genie returns for a
  * conversation id that no longer exists (deleted, expired upstream,
@@ -319,53 +260,141 @@ function isConversationGoneError(err) {
         return true;
     return false;
 }
-function buildAskGenieTool(deps) {
-    const { spaceId, client, writer, signal, resultSets, cacheKey } = deps;
+/* ------------------------ prepare_chart input ------------------------ */
+/**
+ * Agent-facing `prepare_chart` input schema. Reuses
+ * {@link chartPlannerRequestSchema} (the dataset-driven planner
+ * contract) but swaps the inline `data` field for a Genie
+ * `statement_id` the tool resolves into rows server-side.
+ * `title` is loosened to optional - the planner falls back to a
+ * generic placeholder when the agent doesn't supply one.
+ *
+ * Shaped to match Genie's wire form - `statement_id` (snake)
+ * mirrors `query_result.statement_id` and the `get_statement`
+ * tool's input field name, so the LLM only ever sees one
+ * spelling for the same identifier.
+ */
+const prepareChartRequestSchema = chartPlannerRequestSchema
+    .omit({ data: true, title: true })
+    .extend({
+    statement_id: z
+        .string()
+        .min(1, "statement_id is required")
+        .describe(stringUtils.toDescription(`
+          Genie \`statement_id\` to chart. Read from
+          \`message.query_result.statement_id\` or
+          \`message.attachments[*].query.statement_id\` returned by
+          \`ask_genie\`.
+        `)),
+    title: chartPlannerRequestSchema.shape.title.optional(),
+});
+/* ----------------------------- tool ids ----------------------------- */
+/**
+ * Suffix appended to per-space tool ids when the alias isn't the
+ * well-known `default`. Single-space deployments get the bare
+ * names (`ask_genie`, `get_space_description`, ...); multi-space
+ * deployments get `ask_genie_<alias>` etc. so each space's tools
+ * stay disambiguated in the central agent's tool registry.
+ */
+function aliasSuffix(alias) {
+    if (alias === DEFAULT_GENIE_ALIAS)
+        return "";
+    const slug = stringUtils.toIdentifier(alias);
+    return slug ? `_${slug}` : "";
+}
+/* --------------------------- per-space tools --------------------------- */
+/**
+ * Drop `suggested_questions` attachments from a {@link GenieMessage}
+ * before handing it to the central LLM. Those entries already
+ * surface in the UI as one-tap pills via the writer's
+ * `suggested_questions` events (see `collectSuggestions` on the
+ * client); if we let them ride back in the tool result the LLM
+ * tends to quote them inline in its prose, double-showing the
+ * same questions and stepping on the dedicated suggestion UI.
+ * Query / text / row attachments are preserved so the model can
+ * still read `statement_id`, SQL, and the answer text.
+ */
+function stripSuggestedQuestions(message) {
+    const attachments = message.attachments;
+    if (!attachments || attachments.length === 0)
+        return message;
+    const filtered = attachments.filter((att) => !att.suggested_questions);
+    if (filtered.length === attachments.length)
+        return message;
+    return { ...message, attachments: filtered };
+}
+function buildAskGenieTool(opts) {
+    const { spaceId, alias, hint } = opts;
+    const toolId = `ask_genie${aliasSuffix(alias)}`;
+    const hintLine = hint ? ` (${hint})` : "";
     return createTool({
-        id: "ask_genie",
-        description: stringUtils.toDescription `
-      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
+        id: toolId,
+        description: stringUtils.toDescription(`
+      Ask the Genie space "${alias}"${hintLine} ONE focused
+      natural-language sub-question and wait for the turn to
+      complete. Genie answers best when each call covers a
+      single metric / dimension / time window, so expect to
+      call this tool MULTIPLE TIMES per user turn (typically
+      two to six) and let the results from earlier calls inform
+      later ones - that's the normal pattern, not the exception.
+      Do NOT try to cram a multi-part question into a single
+      call; decompose first, then ask each piece.
+
+      Returns the final \`GenieMessage\`. Rows are NOT included -
+      the Genie wire response carries the \`statement_id\` for any
+      SQL that ran (at \`message.query_result.statement_id\` or the
+      first attachment's \`query.statement_id\`); call
+      \`get_statement\` with that id only when you need to read
+      the underlying values to reason about them. If you just
+      want to display the rows to the user, embed a
+      \`[data:<statement_id>]\` marker in your prose instead -
+      the host UI fetches and renders the rows on its own. Wire
       events (status, thinking, sql) stream to the user
-      automatically. Call multiple times to gather different
-      angles before composing the final response.
-    `,
+      automatically while the call is in flight.
+    `),
         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)");
-            }
+            const { client, requestContext } = requireClient(ctx, toolId);
+            const writer = ctx?.writer;
+            const signal = ctx?.abortSignal;
+            const threadId = requestContext.get(MASTRA_THREAD_ID_KEY);
             // 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.
+            // round-trip on them. Genie answers any of these 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, ` +
+                throw new Error(`${toolId}: refusing placeholder question "${question}" - ` +
+                    `call ${toolId} only with a real natural-language question, ` +
                     `or skip the call entirely`);
             }
+            // Seed the active Genie `conversation_id` onto `RequestContext`
+            // from the cross-request cache when a Mastra `threadId` is
+            // present so multi-turn chats reuse the same Genie conversation
+            // (and Genie's accumulated context) across separate user turns.
+            // The same `RequestContext` is reused across every `ask_genie`
+            // call within one user turn, so `ensureConversationSeeded`
+            // hits the cache at most once per request per space.
+            const cacheKey = await conversationCacheKey(spaceId, threadId);
+            await ensureConversationSeeded(requestContext, spaceId, cacheKey);
+            // 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.
+            const startedEvent = {
+                type: "started",
+                spaceId,
+                content: question,
+            };
+            await safeWrite(log, writer, startedEvent);
             // 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.
@@ -377,17 +406,10 @@ function buildAskGenieTool(deps) {
                     ...(seedConversationId ? { conversationId: seedConversationId } : {}),
                     ...(signal ? { context: signal } : {}),
                 })) {
-                    await safeWrite(log, 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 (event.type !== "message") {
+                        await safeWrite(log, writer, event);
+                    }
+                    const eventConversationId = event.conversation_id;
                     if (eventConversationId) {
                         writeContextConversationId(requestContext, spaceId, eventConversationId);
                     }
@@ -431,38 +453,25 @@ function buildAskGenieTool(deps) {
             // 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 } : {}),
-            };
+            return { message: stripSuggestedQuestions(finalMessage) };
         },
     });
 }
-function buildSpaceDescriptionTool(deps) {
-    const { spaceId, client, signal } = deps;
+function buildSpaceDescriptionTool(opts) {
+    const { spaceId, alias } = opts;
+    const toolId = `get_space_description${aliasSuffix(alias)}`;
     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.
-    `,
+        id: toolId,
+        description: stringUtils.toDescription(`
+      Return the Genie space "${alias}"'s title, description, and
+      warehouse id. Cheap (single REST call, no LLM round-trip).
+      Call this FIRST on any user turn that's going to touch
+      \`ask_genie\`, unless the same description already landed
+      earlier in this conversation - the title + description tell
+      you what tables, metrics, and time windows the space
+      actually covers, which is what lets you decompose the
+      user's question into the right \`ask_genie\` sub-questions.
+    `),
         inputSchema: z.object({}),
         outputSchema: z.object({
             spaceId: z.string(),
@@ -470,9 +479,12 @@ function buildSpaceDescriptionTool(deps) {
             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);
+        execute: async (_input, ctxRaw) => {
+            const ctx = ctxRaw;
+            const { client } = requireClient(ctx, toolId);
+            const signal = ctx?.abortSignal;
+            const apiCtx = signal ? apiUtils.toContext(signal) : undefined;
+            const space = await client.genie.getSpace({ space_id: spaceId }, apiCtx);
             return {
                 spaceId,
                 ...(space.title ? { title: space.title } : {}),
@@ -482,424 +494,289 @@ function buildSpaceDescriptionTool(deps) {
         },
     });
 }
-function buildSpaceSerializedTool(deps) {
-    const { spaceId, client, signal } = deps;
+function buildSpaceSerializedTool(opts) {
+    const { spaceId, alias } = opts;
+    const toolId = `get_space_serialized${aliasSuffix(alias)}`;
     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
+        id: toolId,
+        description: stringUtils.toDescription(`
+      Return the full \`GenieSpace\` JSON for the "${alias}" space.
+      Use only when you need exact column / table identifiers
       \`get_space_description\` doesn't expose. Larger payload, so
       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);
+        execute: async (_input, ctxRaw) => {
+            const ctx = ctxRaw;
+            const { client } = requireClient(ctx, toolId);
+            const signal = ctx?.abortSignal;
+            const apiCtx = signal ? apiUtils.toContext(signal) : undefined;
+            const space = await client.genie.getSpace({ space_id: spaceId }, apiCtx);
             return { space };
         },
     });
 }
-/* --------------------------- 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.
-`;
+/* --------------------------- shared tools --------------------------- */
 /**
- * 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.
+ * Default row cap for {@link buildGetStatementTool} when the agent
+ * doesn't supply a `limit`. Sized to keep result sets out of the
+ * model context unless the agent explicitly opts into more rows -
+ * the cheap shape (column names + a handful of representative
+ * rows) is usually enough to reason about a query.
  */
-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:
- *
- *   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.
- */
-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;
+const DEFAULT_STATEMENT_LIMIT = 50;
+function buildGetStatementTool() {
+    const toolId = "get_statement";
     return createTool({
         id: toolId,
-        description: toolDescription,
+        description: stringUtils.toDescription(`
+      Fetch the rows of a Genie statement by its \`statement_id\` (the
+      value at \`message.query_result.statement_id\` or
+      \`message.attachments[*].query.statement_id\` returned from
+      \`ask_genie\`). Use this ONLY when you need to read the underlying
+      values to reason about them in your reply - e.g. naming the
+      largest row, computing a delta the visualization wouldn't already
+      convey, or filtering down to a specific record. If you'd just be
+      reciting numbers the user will see anyway, skip the call and
+      embed a \`[data:<statement_id>]\` marker in your prose instead;
+      the host UI fetches and renders the rows on its own. \`limit\`
+      caps the number of rows returned (defaults to
+      ${DEFAULT_STATEMENT_LIMIT}). \`rowCount\` reflects the full
+      upstream total - compare to \`rows.length\` to detect truncation.
+    `),
         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.
-      `),
+            statement_id: z.string().min(1, "statement_id is required"),
+            limit: z
+                .number()
+                .int()
+                .nonnegative()
+                .optional()
+                .describe("Max rows to return. Defaults to a small sample; raise only when more rows are genuinely needed."),
+        }),
+        outputSchema: z.object({
+            columns: z.array(z.string()),
+            rows: z.array(z.record(z.string(), z.unknown())),
+            rowCount: z.number(),
+            truncated: z.boolean(),
         }),
-        outputSchema: z.custom(),
-        execute: async (input, ctxRaw) => {
+        execute: async ({ statement_id, limit }, ctxRaw) => {
             const ctx = ctxRaw;
-            const requestContext = ctx?.requestContext;
-            if (!requestContext) {
-                throw new Error("genie: missing requestContext (MastraServer must stamp MASTRA_USER_KEY)");
-            }
-            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 { client } = requireClient(ctx, toolId);
             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(log, 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 } : {}),
+            const effectiveLimit = limit ?? DEFAULT_STATEMENT_LIMIT;
+            const data = await fetchStatementData(client, statement_id, {
+                limit: effectiveLimit,
                 ...(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(log, 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 };
-                }
-                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(log, 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 = commonUtils.errorMessage(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(log, 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 } : {}),
+                columns: data.columns,
+                rows: data.rows,
+                rowCount: data.rowCount,
+                truncated: data.rows.length < data.rowCount,
             };
         },
     });
 }
-/* --------------------- multi-alias surface --------------------- */
 /**
- * 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.
+ * `prepare_chart` Mastra tool. Thin wrapper over
+ * {@link prepareChart} that resolves the dataset by fetching the
+ * Genie statement's rows on demand. The tool mints a `chartId`
+ * synchronously, caches an empty placeholder, and kicks off the
+ * planner in the background so the agent loop never blocks. The
+ * host UI resolves `[chart:<chartId>]` markers by reading the
+ * cached {@link Chart} entry (1h TTL).
+ *
+ * Space-agnostic: a Genie `statement_id` is workspace-scoped, so
+ * one shared `prepare_chart` tool covers every wired Genie space.
+ *
+ * Cancellation: deliberately does NOT forward the per-call
+ * `abortSignal` to {@link prepareChart}. The planner task is
+ * fire-and-forget background work; the tool's own `execute`
+ * resolves the moment the `chartId` is minted, so the per-call
+ * signal aborts the second the tool returns. The 1h cache TTL
+ * caps abandoned entries.
  */
-export function defaultGenieToolName(alias) {
-    if (alias === DEFAULT_GENIE_ALIAS)
-        return "genie";
-    return stringUtils.toIdentifierWithOptions({ distinct: true }, "genie", alias);
+function buildPrepareChartTool(opts) {
+    const { config } = opts;
+    const toolId = "prepare_chart";
+    return createTool({
+        id: toolId,
+        description: stringUtils.toDescription([
+            `
+        Queue a chart for the rows of a Genie statement. Mints a
+        short \`chartId\` synchronously and kicks off a BACKGROUND
+        task that fetches the statement's rows, runs the
+        chart-planner to pick a chart type and Echarts spec, and
+        caches the result under the \`chartId\` for one hour. The
+        host UI fetches the cached chart on its own once it lands.
+      `,
+            `
+        To display the chart in your reply, embed
+        \`[chart:<chartId>]\` on its own line at the position you
+        want it to appear, using the EXACT \`chartId\` string this
+        call returned. Never construct a chart id yourself (it is
+        not the \`statement_id\` or any variation of it) - only a
+        value returned by this tool resolves to a real chart. The
+        tool returns immediately - do NOT wait or call it again to
+        "check progress"; the chart resolves asynchronously on the
+        host UI's side.
+      `,
+            `
+        Use this only when the data has a story a chart conveys
+        better than a table (trends, rankings, distributions,
+        parts-of-a-whole). For raw rows, embed
+        \`[data:<statement_id>]\` instead and skip this tool.
+      `,
+        ]),
+        inputSchema: prepareChartRequestSchema,
+        outputSchema: ChartSchema.pick({ chartId: true }),
+        execute: async (request, ctxRaw) => {
+            const ctx = ctxRaw;
+            const { client } = requireClient(ctx, toolId);
+            return prepareChart({
+                config,
+                ...(request.title ? { title: request.title } : {}),
+                ...(request.description ? { description: request.description } : {}),
+                resolveData: (taskSignal) => fetchStatementData(client, request.statement_id, {
+                    ...(taskSignal ? { signal: taskSignal } : {}),
+                }),
+                ...(ctx?.requestContext ? { requestContext: ctx.requestContext } : {}),
+            });
+        },
+    });
 }
+/* --------------------------- orchestration prompt --------------------------- */
+/**
+ * Suggested orchestration prompt for the central agent that owns
+ * the Genie tools. Compose into your agent's `instructions` to
+ * get the canonical "decompose questions, ask Genie focused
+ * sub-questions, place data / chart markers in prose" behavior:
+ *
+ * ```ts
+ * createAgent({
+ *   instructions: `${myAgentInstructions}\n\n${GENIE_INSTRUCTIONS}`,
+ *   tools(plugins) {
+ *     return { ...plugins.genie?.toolkit() };
+ *   },
+ * });
+ * ```
+ *
+ * The prompt references the bare tool names (`ask_genie`,
+ * `get_space_description`, `get_space_serialized`,
+ * `get_statement`, `prepare_chart`) used for the single-space
+ * default alias. Multi-space deployments should write their own
+ * variant that names the suffixed per-space tools
+ * (e.g. `ask_genie_sales`).
+ */
+export const GENIE_INSTRUCTIONS = stringUtils.toDescription([
+    "Genie orchestration. For every user question that needs SQL-backed data:",
+    {
+        numbered: [
+            `
+        Start by calling \`get_space_description\` to ground yourself
+        in what the space covers (tables, metrics, time windows),
+        unless you already saw the same description earlier in this
+        conversation. Reach for \`get_space_serialized\` only when you
+        need exact column / table identifiers the description doesn't
+        expose - it's a much larger payload.
+      `,
+            `
+        Decompose the user's question into focused sub-questions
+        BEFORE asking Genie anything. One sub-question per distinct
+        metric, dimension, or time window. Then call \`ask_genie\`
+        once per sub-question - usually two to six calls per turn,
+        and let earlier answers inform what you ask next. Cramming
+        a multi-part question into one \`ask_genie\` call almost
+        always produces a worse answer than asking the pieces
+        separately. Only collapse to a single call when the question
+        is genuinely atomic ("what was Q3 revenue?").
+
+        Worked example: user asks "How did SKU 1234's revenue
+        compare to its category average last quarter, and which
+        regions drove the gap?" Decomposes to: (a) \`ask_genie\`
+        for SKU 1234's Q3 revenue, (b) \`ask_genie\` for the
+        category-average Q3 revenue, (c) \`ask_genie\` for SKU
+        1234's Q3 revenue split by region. Three focused calls,
+        each grounded in the prior results.
+      `,
+            `
+        Each \`ask_genie\` call returns the terminal \`GenieMessage\`.
+        When the turn ran SQL the result has a \`statement_id\` - read
+        it from \`message.query_result.statement_id\` (or the first
+        attachment's \`query.statement_id\`).
+      `,
+            [
+                `
+          To DISPLAY a result set in your reply, embed a marker on its
+          own line where the visualization should appear. Two marker
+          shapes:
+        `,
+                {
+                    bullets: [
+                        `
+              \`[data:<statement_id>]\` - render the rows as a table.
+              Use this when there's no clear visual story (long lists,
+              reference data, single-row results, or the user just
+              wants to see the data). Embed the marker directly; no
+              tool call needed.
+            `,
+                        `
+              \`[chart:<chartId>]\` - render the rows as a chart. To
+              get a \`<chartId>\`, call \`prepare_chart\` with the
+              statement's id (and an optional \`title\` / one-line
+              \`description\` of the insight to surface). The tool
+              returns the \`chartId\` synchronously and prepares the
+              chart spec in the background; embed the returned id as
+              \`[chart:<chartId>]\` on its own line wherever the
+              chart should appear. Use a chart when the data has a
+              story a visual conveys better than a table (trends,
+              rankings, distributions, parts-of-a-whole).
+
+              NEVER invent or hand-build a \`<chartId>\`. A valid
+              \`<chartId>\` is the opaque token a \`prepare_chart\`
+              call returned to you in THIS turn - nothing else. It
+              is NOT a \`statement_id\`, and it is NOT a
+              \`statement_id\` prefix with a label appended (e.g.
+              \`01f1...-region-fill\`). If you have not called
+              \`prepare_chart\` and received an id back, do not write
+              a \`[chart:...]\` marker at all - use \`[data:...]\`
+              instead. A fabricated chart id renders nothing and
+              wastes a request.
+            `,
+                    ],
+                },
+                `
+          The host UI resolves both markers on its own once it sees
+          them - you do NOT need to call \`get_statement\` just to
+          display data, and you do NOT need to wait on
+          \`prepare_chart\` (it returns the id immediately and the
+          host UI fetches the cached chart later). Pick at most one
+          marker per statement; don't chart AND table the same result
+          side by side.
+        `,
+            ],
+            `
+        Call \`get_statement(statement_id, limit?)\` ONLY when you need
+        to read the actual values to reason about them (e.g. naming a
+        specific row, computing a delta the table or chart wouldn't
+        show, or sanity-checking a result before interpreting it). If
+        you'd just be reciting numbers the visualization already shows,
+        skip the call and use a marker instead. \`limit\` defaults to a
+        small sample; raise it only when you genuinely need more rows.
+      `,
+            `
+        Compose your final reply as plain prose. Interleave paragraphs
+        with \`[data:...]\` / \`[chart:...]\` markers wherever a result
+        should render. Don't dump all markers at the end - place each
+        one next to the prose that interprets it. Don't restate every
+        number the visualization already shows; call out deltas,
+        anomalies, takeaways.
+      `,
+        ],
+    },
+]);
+/* --------------------- multi-alias surface --------------------- */
 /**
  * Normalize the {@link GenieSpacesConfig} record. Bare-string
  * entries (`{ default: "01ef..." }`) get wrapped as
@@ -975,42 +852,46 @@ export function resolveGenieSpaces(config, context) {
     return merged;
 }
 /**
- * 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.
+ * Build the flat Mastra tools record for every configured Genie
+ * space. Two shared, space-agnostic tools (`get_statement`,
+ * `prepare_chart`) are registered once regardless of how many
+ * spaces are wired; the per-space tools (`ask_genie`,
+ * `get_space_description`, `get_space_serialized`) are suffixed
+ * with `_<alias>` for non-default aliases so multi-space
+ * deployments stay disambiguated.
  *
- * Returns a record keyed by tool id, ready to spread into an
- * `Agent`'s `tools` map (or surfaced via
- * `plugins.genie?.toolkit()`).
+ * Returns a record keyed by tool id, ready to spread into the
+ * central `Agent`'s `tools` map (or surfaced via the
+ * `plugins.genie?.toolkit()` callback). Returns an empty record
+ * when `spaces` resolves to zero entries so the caller can spread
+ * safely.
  */
 export function buildGenieTools(opts) {
     const normalized = normalizeGenieSpaces(opts.spaces);
     const tools = {};
+    if (Object.keys(normalized).length === 0)
+        return tools;
+    // Shared, space-agnostic tools.
+    tools.get_statement = buildGetStatementTool();
+    tools.prepare_chart = buildPrepareChartTool({ config: opts.config });
     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({
+        const askTool = buildAskGenieTool({
             spaceId: space.spaceId,
-            config: opts.config,
-            toolId: id,
-            toolDescription,
+            alias,
+            ...(space.hint ? { hint: space.hint } : {}),
         });
+        const descTool = buildSpaceDescriptionTool({ spaceId: space.spaceId, alias });
+        const serTool = buildSpaceSerializedTool({ spaceId: space.spaceId, alias });
+        tools[askTool.id] = askTool;
+        tools[descTool.id] = descTool;
+        tools[serTool.id] = serTool;
     }
     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.
+ * flat Genie tools record instead of throwing on missing plugin.
  * Mirrors AppKit's `PluginToolkitProvider` shape.
  */
 export function buildGenieToolkitProvider(opts) {