@dbx-tools/appkit-mastra 0.9.128 → 0.9.131

package/src/plugin.ts

@@ -56,13 +56,15 @@ import {
   type PluginManifest,
   type ResourceRequirement,
 } from "@databricks/appkit";
-import { lakebaseResolver, pluginRegistry } from "@dbx-tools/appkit";
+import { lakebaseResolver, pluginRegistry, telemetry as appkitTelemetry } from "@dbx-tools/appkit";
 import { modelCatalog } from "@dbx-tools/model";
 import { asyncUtils, errorUtils, log, object, stringUtils } from "@dbx-tools/shared-core";
 import {
   feedback,
   routes,
+  wire,
   type Chart,
+  type GenieChatRequest,
   type MastraClientConfig,
   type MastraFeedbackRequest,
   type StatementData,
@@ -71,10 +73,16 @@ import { display, type ServingEndpointSummary } from "@dbx-tools/shared-model";
 import type { Agent } from "@mastra/core/agent";
 import { Mastra } from "@mastra/core/mastra";
 import type { RequestContext } from "@mastra/core/request-context";
+import { pipeUIMessageStreamToResponse } from "ai";
 import express from "express";
 import type { Pool } from "pg";
 
-import { buildAgents, FALLBACK_AGENT_ID, type BuiltAgents } from "./agents.ts";
+import {
+  buildAgents,
+  FALLBACK_AGENT_ID,
+  type BuiltAgents,
+  type MastraOutputProcessorFactory,
+} from "./agents.ts";
 import { fetchChart } from "./chart.ts";
 import { agentChatRoutes } from "./chat.ts";
 import {
@@ -86,11 +94,21 @@ import {
 import {
   chartFetchDefaults,
   feedbackWriteDefaults,
+  genieConversationReadDefaults,
+  genieConversationWriteDefaults,
   genieSuggestionDefaults,
   modelCatalogueDefaults,
   statementDataDefaults,
 } from "./defaults.ts";
-import { collectSpaceSuggestions, resolveGenieSpaces } from "./genie.ts";
+import {
+  deleteGenieThread,
+  deleteGenieThreadMessage,
+  genieChatStream,
+  listGenieThreadMessages,
+  listGenieThreads,
+  type GenieChatSpace,
+} from "./genie-chat.ts";
+import { collectSpaceSuggestions, resolveGenieSpaces, type GenieSpaceConfig } from "./genie.ts";
 import {
   requestUserId,
   resolveIdentityMode,
@@ -117,8 +135,8 @@ import {
   MastraServer,
 } from "./server.ts";
 import { resolveServingConfig } from "./serving.ts";
-import { fetchStatementData, STATEMENT_ROW_CAP } from "./statement.ts";
-import { attachChatTurnTelemetry } from "./telemetry.ts";
+import { fetchStatementEmbedData, STATEMENT_ROW_CAP } from "./statement.ts";
+import { attachChatTurnTelemetry, registerAppMlflowTraceInfoProcessor } from "./telemetry.ts";
 import { invalidFields } from "./validation.ts";
 import { inlineSkillsFromLocalMounts } from "./workspaces.ts";
 
@@ -225,6 +243,7 @@ export class MastraPlugin extends Plugin<MastraPluginConfig> {
   private mastra: Mastra | null = null;
   private mastraApp: express.Express | null = null;
   private mastraServer: MastraServer | null = null;
+  private readonly outputProcessorFactories: MastraOutputProcessorFactory[] = [];
   /**
    * The optional MCP server exposing this plugin's agents / tools, or
    * `null` when `config.mcp` is disabled (the default). Built in
@@ -242,6 +261,9 @@ export class MastraPlugin extends Plugin<MastraPluginConfig> {
     // app boot with a clear message rather than 500-ing the first chat request.
     this.identityMode = resolveIdentityMode(this.config.genieIdentity);
     validateFeedbackConfig(this.config.feedback);
+    if (this.config.observability !== false) {
+      await registerAppMlflowTraceInfoProcessor();
+    }
     // Wait until sibling plugins (e.g. `lakebase`) finish `setup()` so
     // the lakebase pool is valid when storage/memory are enabled.
     this.context?.onLifecycle("setup:complete", async () => {
@@ -252,8 +274,9 @@ export class MastraPlugin extends Plugin<MastraPluginConfig> {
   }
 
   /**
-   * The plugin instance the chat path's Databricks calls should run through for
-   * `req`, honoring {@link identityMode}.
+   * The plugin instance agent-backed Databricks calls should run through for
+   * `req`, honoring {@link identityMode}. Direct Genie conversation routes use
+   * AppKit OBO through `asUser(req)` instead.
    *
    * In `"user"` mode (the default) this is `this.asUser(req)` - AppKit's own
    * proxy that runs every call inside `runInUserContext`, so `getExecutionContext()`
@@ -261,9 +284,9 @@ export class MastraPlugin extends Plugin<MastraPluginConfig> {
    * unwrapped, so the same methods run in the ambient service context and
    * `getExecutionContext()` returns the app service principal's client - the same
    * "just don't enter `asUser`" pattern AppKit's own service-context calls use. One
-   * helper covers every workspace call the chat path makes (serving catalogue,
-   * Genie suggestions, `ask_genie` via the dispatched agent turn, statement fetch,
-   * feedback), so the identity decision lives in one place.
+   * helper covers the agent-backed workspace calls (serving catalogue, general
+   * suggestions, dispatched `ask_genie`, statement fetch, feedback), so the
+   * identity decision lives in one place.
    */
   private scopedSelf(req: express.Request): this {
     return useServicePrincipal(this.identityMode, req) ? this : this.asUser(req);
@@ -323,6 +346,18 @@ export class MastraPlugin extends Plugin<MastraPluginConfig> {
        */
       getDefault: (): Agent | null =>
         (this.built && this.built.agents[this.built.defaultAgentId]) ?? null,
+      /**
+       * Register an output-processor factory before agents are built.
+       *
+       * Deferred sibling plugins use this generic extension point without
+       * adding provider-specific behavior to the Mastra plugin.
+       */
+      registerOutputProcessor: (factory: MastraOutputProcessorFactory): void => {
+        if (this.built) {
+          throw new Error("Mastra output processors must be registered before agents are built");
+        }
+        this.outputProcessorFactories.push(factory);
+      },
       /** Underlying Mastra instance for advanced use (custom routes etc.). */
       getMastra: () => this.mastra,
       /**
@@ -441,6 +476,7 @@ export class MastraPlugin extends Plugin<MastraPluginConfig> {
     this.registerModelRoutes(router);
     this.registerEmbedRoutes(router);
     this.registerSuggestionRoutes(router);
+    this.registerGenieChatRoutes(router);
     this.registerFeedbackRoutes(router);
     this.registerMlflowExperimentRoute(router);
     this.registerAgentRoutes(router);
@@ -545,10 +581,9 @@ export class MastraPlugin extends Plugin<MastraPluginConfig> {
     // `GET /embed/:type/:id` is the single resolver for every embed
     // marker the agent emits in prose (`[chart:<id>]`,
     // `[data:<id>]`, ...). `:type` selects a resolver from the
-    // registry below; `:id` is that resolver's lookup key. The
-    // grammar (see `marker.ts`) is type-agnostic on purpose - new
-    // embed kinds are added by registering a resolver here, with no
-    // client or grammar change.
+    // registry below; `:id` is that resolver's lookup key. The marker grammar
+    // is type-agnostic; the client still chooses how each registered type
+    // renders.
     //
     // Status codes:
     //   - 200 with the resolver's JSON body when the id resolves.
@@ -563,7 +598,9 @@ export class MastraPlugin extends Plugin<MastraPluginConfig> {
     //     settles (`result` / `error`) or the budget elapses (then
     //     returns the still-processing entry to poll again).
     //     `?timeoutMs=<n>` (default 60s, capped 5min) tunes it.
-    //   - `data`: one OBO-scoped Statement Execution fetch.
+    //   - `data`: one Statement Execution fetch under `genieIdentity`.
+    //   - `visualization`: the same payload under mandatory OBO identity for
+    //     direct Genie chat, whose conversations belong to the signed-in user.
     //     `?limit=<n>` caps rows (clamped to STATEMENT_ROW_CAP).
     //
     // Built once (this handler is registered once) and keyed by the
@@ -587,6 +624,13 @@ export class MastraPlugin extends Plugin<MastraPluginConfig> {
           signal,
         });
       },
+      visualization: (req, id, signal) => {
+        const limit = parseStatementLimit(req.query.limit);
+        return this.asUser(req).fetchStatement(id, {
+          ...(limit !== undefined ? { limit } : {}),
+          signal,
+        });
+      },
     };
 
     this.route(router, {
@@ -619,13 +663,10 @@ export class MastraPlugin extends Plugin<MastraPluginConfig> {
           res.status(400).json({ error: "id is required" });
           return;
         }
-        // Express's `req` predates `AbortSignal`; bridge the `close`
-        // event onto an `AbortController` so a closed connection
-        // unblocks any long-poll immediately and frees the request
-        // thread. The listener is GC'd with the request on normal
-        // completion.
+        // Bridge response close onto an AbortSignal so a client disconnect
+        // releases long-poll and workspace calls immediately.
         const controller = new AbortController();
-        req.on("close", () => controller.abort());
+        res.on("close", () => controller.abort());
         const result = await resolve(req, id, controller.signal);
         if (controller.signal.aborted) return;
         if (!result.ok) {
@@ -663,14 +704,14 @@ export class MastraPlugin extends Plugin<MastraPluginConfig> {
     // empty state (no built-in example prompts). The `:agentId`
     // segment is accepted for URL symmetry with the chat / history
     // routes; Genie spaces are resolved per-plugin, not per-agent,
-    // so it doesn't change the result. OBO-scoped like the other
-    // data routes so the space lookup runs as the calling user.
+    // so it doesn't change the result. The space lookup uses the configured
+    // agent-backed identity.
     const handleSuggestions = async (
       req: express.Request,
       res: express.Response,
     ): Promise<void> => {
       const controller = new AbortController();
-      req.on("close", () => controller.abort());
+      res.on("close", () => controller.abort());
       const result = await this.scopedSelf(req).fetchSuggestions(controller.signal);
       if (controller.signal.aborted) return;
       if (!result.ok) {
@@ -700,6 +741,189 @@ export class MastraPlugin extends Plugin<MastraPluginConfig> {
     });
   }
 
+  /**
+   * Register the direct Genie chat routes. Each route is scoped to one Genie
+   * space alias and sits directly on the Genie APIs under AppKit's OBO user
+   * context: conversations are threads, Genie messages are history, and a turn
+   * streams the same AI SDK chunks as an agent's `ask_genie` call. Nothing is
+   * written to Mastra Memory.
+   */
+  private registerGenieChatRoutes(router: IAppRouter): void {
+    const path = (...parts: string[]) => routes.genieRoutePath(":alias", ...parts);
+    const { chat: chatSegment, messages, suggestions, threads } = routes.MASTRA_ROUTES;
+    /** Resolve the route's alias, answering 404 when no such space is configured. */
+    const spaceFor = (req: express.Request, res: express.Response): GenieChatSpace | undefined => {
+      const alias = stringUtils.firstNonEmpty(req.params.alias) ?? "";
+      const configured = resolveGenieSpaces(this.config, this.context)[alias];
+      if (configured) return { alias, spaceId: configured.spaceId };
+      res.status(404).json({ error: `Unknown Genie space alias "${alias}"` });
+      return undefined;
+    };
+    /** Bridge the connection `close` event onto an abort signal. */
+    const closeSignal = (res: express.Response): AbortSignal => {
+      const controller = new AbortController();
+      res.on("close", () => controller.abort());
+      return controller.signal;
+    };
+
+    this.route(router, {
+      name: "genieChat",
+      method: "post",
+      path: path(chatSegment),
+      handler: async (req, res) => {
+        const space = spaceFor(req, res);
+        if (!space) return;
+        const parsed = wire.GenieChatRequestSchema.safeParse(req.body);
+        if (!parsed.success) {
+          res
+            .status(400)
+            .json({ error: "question is required", fields: invalidFields(parsed.error) });
+          return;
+        }
+        await this.asUser(req).streamGenieChat(res, space, parsed.data, closeSignal(res));
+      },
+    });
+    this.route(router, {
+      name: "genieThreads",
+      method: "get",
+      path: path(threads),
+      handler: async (req, res) => {
+        const space = spaceFor(req, res);
+        if (!space) return;
+        const result = await this.asUser(req).genieConversationRead((client) =>
+          listGenieThreads(client, space),
+        );
+        if (!result.ok) {
+          this.sendFailure(res, result, "genie-threads", "Could not list Genie conversations");
+          return;
+        }
+        res.json({ threads: result.data });
+      },
+    });
+    this.route(router, {
+      name: "genieThreadMessages",
+      method: "get",
+      path: path(threads, ":threadId", messages),
+      handler: async (req, res) => {
+        const space = spaceFor(req, res);
+        if (!space) return;
+        const threadId = stringUtils.firstNonEmpty(req.params.threadId) ?? "";
+        const result = await this.asUser(req).genieConversationRead((client) =>
+          listGenieThreadMessages({
+            client,
+            space,
+            conversationId: threadId,
+          }),
+        );
+        if (!result.ok) {
+          this.sendFailure(res, result, "genie-messages", "Could not load the Genie conversation");
+          return;
+        }
+        res.json({ messages: result.data });
+      },
+    });
+    this.route(router, {
+      name: "genieThreadDelete",
+      method: "delete",
+      path: path(threads, ":threadId"),
+      handler: async (req, res) => {
+        const space = spaceFor(req, res);
+        if (!space) return;
+        const threadId = stringUtils.firstNonEmpty(req.params.threadId) ?? "";
+        const result = await this.asUser(req).genieConversationWrite((client) =>
+          deleteGenieThread(client, space, threadId),
+        );
+        if (!result.ok) {
+          this.sendFailure(
+            res,
+            result,
+            "genie-thread-delete",
+            "Could not delete the Genie conversation",
+          );
+          return;
+        }
+        res.status(204).end();
+      },
+    });
+    this.route(router, {
+      name: "genieThreadMessageDelete",
+      method: "delete",
+      path: path(threads, ":threadId", messages, ":messageId"),
+      handler: async (req, res) => {
+        const space = spaceFor(req, res);
+        if (!space) return;
+        const threadId = stringUtils.firstNonEmpty(req.params.threadId) ?? "";
+        const messageId = stringUtils.firstNonEmpty(req.params.messageId) ?? "";
+        const result = await this.asUser(req).genieConversationWrite((client) =>
+          deleteGenieThreadMessage(client, space, threadId, messageId),
+        );
+        if (!result.ok) {
+          this.sendFailure(
+            res,
+            result,
+            "genie-message-delete",
+            "Could not delete the Genie message",
+          );
+          return;
+        }
+        res.status(204).end();
+      },
+    });
+    this.route(router, {
+      name: "genieSuggestions",
+      method: "get",
+      path: path(suggestions),
+      handler: async (req, res) => {
+        const space = spaceFor(req, res);
+        if (!space) return;
+        const result = await this.asUser(req).fetchSuggestions(undefined, {
+          [space.alias]: { spaceId: space.spaceId },
+        });
+        res.json({ questions: result.ok ? result.data : [] });
+      },
+    });
+  }
+
+  /**
+   * Stream one direct Genie turn to `res` under AppKit's OBO user context.
+   * Direct conversations cannot use the shared service principal because Genie
+   * owns them under the Databricks caller.
+   */
+  private async streamGenieChat(
+    res: express.Response,
+    space: GenieChatSpace,
+    request: GenieChatRequest,
+    signal: AbortSignal,
+  ): Promise<void> {
+    const context = getExecutionContext();
+    const client = context.client;
+    await pipeUIMessageStreamToResponse({
+      response: res,
+      stream: genieChatStream({
+        client,
+        space,
+        request,
+        signal,
+      }),
+    });
+  }
+
+  /** Run one Genie conversation read under the active request identity. */
+  private genieConversationRead<T>(
+    read: (client: ReturnType<typeof getExecutionContext>["client"]) => Promise<T>,
+  ): Promise<ExecutionResult<T>> {
+    const client = getExecutionContext().client;
+    return this.execute(() => read(client), genieConversationReadDefaults);
+  }
+
+  /** Run one Genie conversation write under the active request identity. */
+  private genieConversationWrite(
+    write: (client: ReturnType<typeof getExecutionContext>["client"]) => Promise<void>,
+  ): Promise<ExecutionResult<void>> {
+    const client = getExecutionContext().client;
+    return this.execute(() => write(client), genieConversationWriteDefaults);
+  }
+
   /** Register MLflow feedback submission when configured. */
   private registerFeedbackRoutes(router: IAppRouter): void {
     // `POST /route/feedback` logs a thumbs / comment assessment against
@@ -866,11 +1090,15 @@ export class MastraPlugin extends Plugin<MastraPluginConfig> {
    * AppKit user-context proxy so `getExecutionContext()` returns the
    * OBO-scoped client. Resolves the plugin's Genie spaces and merges
    * their curated `sample_questions` (see {@link collectSpaceSuggestions}).
-   * Returns `[]` when no Genie space is configured so the client
-   * shows a bare empty state instead of built-in example prompts.
+   * `only` narrows the lookup to the given spaces (the direct Genie chat
+   * passes its own). Returns `[]` when no Genie space is configured so the
+   * client shows a bare empty state instead of built-in example prompts.
    */
-  private async fetchSuggestions(signal?: AbortSignal): Promise<ExecutionResult<string[]>> {
-    const spaces = resolveGenieSpaces(this.config, this.context);
+  private async fetchSuggestions(
+    signal?: AbortSignal,
+    only?: Record<string, GenieSpaceConfig>,
+  ): Promise<ExecutionResult<string[]>> {
+    const spaces = only ?? resolveGenieSpaces(this.config, this.context);
     if (Object.keys(spaces).length === 0) return { ok: true, data: [] };
     const client = getExecutionContext().client;
     return this.execute((executeSignal) => {
@@ -942,7 +1170,7 @@ export class MastraPlugin extends Plugin<MastraPluginConfig> {
       async (executeSignal) => {
         const combined = asyncUtils.combineAbortSignals(options.signal, executeSignal);
         try {
-          const data = await fetchStatementData(client, statementId, {
+          const { data, visualization } = await fetchStatementEmbedData(client, statementId, {
             limit,
             ...(combined ? { signal: combined } : {}),
           });
@@ -951,6 +1179,7 @@ export class MastraPlugin extends Plugin<MastraPluginConfig> {
             rows: data.rows,
             rowCount: data.rowCount,
             truncated: data.rows.length < data.rowCount,
+            visualization,
           };
         } catch (err) {
           // The Databricks SDK throws on 404; surface as `undefined`
@@ -1045,7 +1274,7 @@ export class MastraPlugin extends Plugin<MastraPluginConfig> {
           pluginRegistry.require(this.context, lakebase, this.config).exports().pool,
         ) as Pool)
       : undefined;
-    const memoryBuilder = memoryPool ? createMemoryBuilder(this.config, memoryPool) : undefined;
+    const memoryBuilder = createMemoryBuilder(this.config, memoryPool);
 
     // Build every agent declared in `config.agents` (or the built-in
     // fallback when none are declared). Each agent's `model` resolves
@@ -1080,6 +1309,7 @@ export class MastraPlugin extends Plugin<MastraPluginConfig> {
       config: this.config,
       context: this.context,
       memoryBuilder,
+      outputProcessorFactories: this.outputProcessorFactories,
       log: this.logger,
       extraSkillMounts: localSkillMounts,
       standardSkills,
@@ -1106,6 +1336,7 @@ export class MastraPlugin extends Plugin<MastraPluginConfig> {
     const observability = await buildObservability({
       serviceName: this.name,
       enabled: this.config.observability,
+      tracerProvider: await appkitTelemetry.deferredNativeTracerProvider(),
     });
     // Optional MCP exposure: build a Mastra MCP server from the
     // registered agents (and, opt-in, the ambient tools) and register
@@ -1151,7 +1382,7 @@ export class MastraPlugin extends Plugin<MastraPluginConfig> {
       backgroundTurns: this.config.backgroundTurns !== false,
       apiAccess: this.config.apiAccess ?? "scoped",
       mcp: this.mcp ? `${this.basePath}${this.mcp.httpPath}` : "off",
-      lakebase: memoryBuilder !== undefined,
+      lakebase: memoryPool !== undefined,
       feedback: this.feedbackEnabled(),
       observability: observability !== undefined ? "mlflow" : "off",
     });

package/src/request-context.ts

@@ -0,0 +1,2 @@
+/** Request-context key containing the browser application's canonical origin. */
+export const MASTRA_ORIGIN_KEY = "mastra__origin";

package/src/server.ts

@@ -20,6 +20,7 @@ import type express from "express";
 
 import {
   attributedUserId,
+  MASTRA_AUTH_MODE_KEY,
   MASTRA_REQUEST_ID_KEY,
   MASTRA_SCOPES_KEY,
   MASTRA_USER_EMAIL_KEY,
@@ -30,6 +31,7 @@ import {
 } from "./config.ts";
 import { requestUserEmail, requestUserId } from "./identity.ts";
 import { resolveFeedbackEnabled } from "./mlflow.ts";
+import { MASTRA_ORIGIN_KEY } from "./request-context.ts";
 import {
   extractModelOverride,
   MASTRA_MODEL_OVERRIDE_KEY,
@@ -51,7 +53,9 @@ const TRUSTED_REQUEST_CONTEXT_KEYS = [
   MASTRA_USER_KEY,
   MASTRA_USER_NAME_KEY,
   MASTRA_USER_EMAIL_KEY,
+  MASTRA_AUTH_MODE_KEY,
   MASTRA_REQUEST_ID_KEY,
+  MASTRA_ORIGIN_KEY,
   MASTRA_SCOPES_KEY,
   MASTRA_MODEL_OVERRIDE_KEY,
   MASTRA_RESOLVED_MODEL_KEY,
@@ -64,6 +68,29 @@ export function clearTrustedRequestContext(requestContext: RequestContext): void
   }
 }
 
+/** Resolve the canonical browser origin from standard proxy and browser headers. */
+export function requestOrigin(req: express.Request): string | undefined {
+  for (const value of [req.get("origin"), req.get("referer")]) {
+    if (!value) continue;
+    try {
+      const url = new URL(value);
+      if (url.protocol === "http:" || url.protocol === "https:") return url.origin;
+    } catch {
+      // Continue to proxy headers.
+    }
+  }
+  const host = req.get("x-forwarded-host")?.split(",", 1)[0]?.trim() || req.get("host")?.trim();
+  if (!host) return undefined;
+  const protocol =
+    req.get("x-forwarded-proto")?.split(",", 1)[0]?.trim() || req.protocol || "https";
+  if (protocol !== "http" && protocol !== "https") return undefined;
+  try {
+    return new URL(`${protocol}://${host}`).origin;
+  } catch {
+    return undefined;
+  }
+}
+
 /**
  * Who a turn is ATTRIBUTED to, independent of which Databricks credential runs
  * its calls. Supplied by the HTTP middleware from the forwarded user headers so
@@ -116,6 +143,7 @@ export async function stampRequestContextUser(
   };
   requestContext.set(MASTRA_USER_KEY, user);
   requestContext.set(MASTRA_RESOURCE_ID_KEY, user.id);
+  requestContext.set(MASTRA_AUTH_MODE_KEY, isUserContext ? "obo" : "service-principal");
   // AppKit's `UserContext` surfaces display name / email only on
   // OBO requests. Service-context calls (background tasks, server
   // start-up) leave these undefined and we skip the stamp so
@@ -197,6 +225,7 @@ export class MastraServer extends ExpressMastraServer {
       this.configureRequestContextThreadId(req, res, requestContext);
       this.configureRequestContextModelOverride(req, requestContext);
       this.configureRequestContextRequestId(req, res, requestContext);
+      this.configureRequestContextOrigin(req, requestContext);
       this.configureRequestContextScopes(req, requestContext);
 
       this.configureMlflowTraceId(res);
@@ -238,6 +267,12 @@ export class MastraServer extends ExpressMastraServer {
     recordActiveTraceUser(traceUser);
   }
 
+  /** Stamp the public browser origin so plugin tools can return canonical app links. */
+  configureRequestContextOrigin(req: express.Request, requestContext: RequestContext): void {
+    const origin = requestOrigin(req);
+    if (origin) requestContext.set(MASTRA_ORIGIN_KEY, origin);
+  }
+
   /**
    * Stamp a per-request id and echo it on the response so an upstream
    * proxy / curl client / browser-side log line can pair its view of

package/src/statement.ts

@@ -22,7 +22,7 @@
 import type { WorkspaceClient } from "@databricks/appkit";
 import { databricks } from "@dbx-tools/appkit";
 import { object } from "@dbx-tools/shared-core";
-import type { GenieDatasetData } from "@dbx-tools/shared-mastra";
+import { wire, type GenieDatasetData, type GenieVisualizationData } from "@dbx-tools/shared-mastra";
 
 /**
  * Hard server-side cap on rows returned by the
@@ -48,6 +48,60 @@ function coerceCell(cell: string | null): unknown {
   return object.toNumber(cell, { separators: false, percent: false }) ?? cell;
 }
 
+type StatementResponse = Awaited<ReturnType<WorkspaceClient["statementExecution"]["getStatement"]>>;
+
+async function fetchStatementResponse(
+  client: WorkspaceClient,
+  statementId: string,
+  signal?: AbortSignal,
+): Promise<StatementResponse> {
+  const context = signal ? databricks.toContext(signal) : undefined;
+  return client.statementExecution.getStatement({ statement_id: statementId }, context);
+}
+
+function normalizedStatementData(response: StatementResponse, limit?: number): GenieDatasetData {
+  const columns = (response.manifest?.schema?.columns ?? []).map((column) => column.name ?? "");
+  const dataArray = (response.result?.data_array ?? []) as Array<Array<string | null>>;
+  const sliced = limit !== undefined && limit >= 0 ? dataArray.slice(0, limit) : dataArray;
+  const rows = sliced.map((row) => {
+    const record: Record<string, unknown> = {};
+    columns.forEach((column, index) => {
+      record[column] = coerceCell(row[index] ?? null);
+    });
+    return record;
+  });
+  return {
+    columns,
+    rows,
+    rowCount: response.manifest?.total_row_count ?? dataArray.length,
+  };
+}
+
+function visualizationData(response: StatementResponse, limit?: number): GenieVisualizationData {
+  const dataArray = (response.result?.data_array ?? []) as Array<Array<string | null>>;
+  return wire.GenieVisualizationDataSchema.parse({
+    manifest: {
+      schema: {
+        columns: (response.manifest?.schema?.columns ?? []).map((column) => ({
+          name: column.name,
+          type_name: column.type_name,
+        })),
+      },
+    },
+    result: {
+      data_array: limit !== undefined && limit >= 0 ? dataArray.slice(0, limit) : dataArray,
+    },
+  });
+}
+
+/** Normalized and AppKit-native projections of one statement response. */
+export interface StatementEmbedData {
+  /** Row records used by tools, sortable tables, and chart planners. */
+  data: GenieDatasetData;
+  /** Manifest and row arrays consumed by AppKit UI's Genie visualization. */
+  visualization: GenieVisualizationData;
+}
+
 /**
  * Fetch a single statement's rows via the Statement Execution API
  * and reshape into the shared {@link GenieDatasetData} shape
@@ -70,24 +124,22 @@ export async function fetchStatementData(
   statementId: string,
   options?: { limit?: number; signal?: AbortSignal },
 ): Promise<GenieDatasetData> {
-  const ctx = options?.signal ? databricks.toContext(options.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 ?? []) as Array<Array<string | null>>;
-  const sliced =
-    options?.limit !== undefined && options.limit >= 0
-      ? dataArray.slice(0, options.limit)
-      : dataArray;
-  const rows = sliced.map((row) => {
-    const obj: Record<string, unknown> = {};
-    columns.forEach((col, i) => {
-      obj[col] = coerceCell(row[i] ?? null);
-    });
-    return obj;
-  });
+  const response = await fetchStatementResponse(client, statementId, options?.signal);
+  return normalizedStatementData(response, options?.limit);
+}
+
+/**
+ * Fetch one statement once and project it for both the existing data table and
+ * AppKit UI's native Genie visualization.
+ */
+export async function fetchStatementEmbedData(
+  client: WorkspaceClient,
+  statementId: string,
+  options?: { limit?: number; signal?: AbortSignal },
+): Promise<StatementEmbedData> {
+  const response = await fetchStatementResponse(client, statementId, options?.signal);
   return {
-    columns,
-    rows,
-    rowCount: r.manifest?.total_row_count ?? dataArray.length,
+    data: normalizedStatementData(response, options?.limit),
+    visualization: visualizationData(response, options?.limit),
   };
 }