feat(appkit): support Genie agent mode with reasoning and visualizations

apps/dev-playground/client/src/routes/genie.route.tsx

@@ -1,11 +1,32 @@
-import { GenieChat } from "@databricks/appkit-ui/react";
+import { Button, GenieChat } from "@databricks/appkit-ui/react";
 import { createFileRoute } from "@tanstack/react-router";
+import { useState } from "react";
 
 export const Route = createFileRoute("/genie")({
   component: GenieRoute,
 });
 
+type Mode = "chat" | "agent";
+
+// Genie plugin aliases configured in server/index.ts.
+const ALIASES: Record<Mode, string> = { chat: "demo", agent: "agent" };
+
 function GenieRoute() {
+  const [mode, setMode] = useState<Mode>(() =>
+    new URLSearchParams(window.location.search).get("mode") === "agent"
+      ? "agent"
+      : "chat",
+  );
+
+  const switchMode = (next: Mode) => {
+    // Chat and agent mode don't share conversations.
+    const url = new URL(window.location.href);
+    url.searchParams.set("mode", next);
+    url.searchParams.delete("conversationId");
+    window.history.replaceState({}, "", url.toString());
+    setMode(next);
+  };
+
   return (
     <div className="min-h-screen bg-background">
       <main className="max-w-4xl mx-auto px-6 py-12">
@@ -16,11 +37,29 @@ function GenieRoute() {
             </h1>
             <p className="text-muted-foreground mt-2">
               Ask natural language questions about your data using AI/BI Genie.
+              Agent mode shows the agent's reasoning and its charts.
             </p>
           </div>
 
+          <div className="flex gap-2">
+            <Button
+              size="sm"
+              variant={mode === "chat" ? "default" : "outline"}
+              onClick={() => switchMode("chat")}
+            >
+              Chat
+            </Button>
+            <Button
+              size="sm"
+              variant={mode === "agent" ? "default" : "outline"}
+              onClick={() => switchMode("agent")}
+            >
+              Agent
+            </Button>
+          </div>
+
           <div className="border rounded-lg h-[600px] flex flex-col">
-            <GenieChat alias="demo" />
+            <GenieChat key={mode} alias={ALIASES[mode]} />
           </div>
         </div>
       </main>

apps/dev-playground/server/index.ts

@@ -67,6 +67,8 @@ createApp({
     analytics({}),
     genie({
       spaces: { demo: process.env.DATABRICKS_GENIE_SPACE_ID ?? "placeholder" },
+      // Any Genie space also works as an agent.
+      agents: { agent: process.env.DATABRICKS_GENIE_SPACE_ID ?? "placeholder" },
     }),
     ...(process.env.LAKEBASE_ENDPOINT ? [lakebase()] : []),
     lakebaseExamples(),

docs/docs/plugins/genie.md

@@ -11,6 +11,8 @@ Integrates [Databricks AI/BI Genie](https://docs.databricks.com/en/genie/index.h
 - SSE streaming with real-time status updates
 - Conversation history replay with automatic reconnection
 - Query result attachment fetching
+- Genie's charts, served as PNG images
+- Agent mode: the agent's reasoning, streamed as it works, and the charts it draws
 - On-behalf-of (OBO) user execution
 
 ## Basic usage
@@ -31,7 +33,8 @@ await createApp({
 | Option | Type | Default | Description |
 |--------|------|---------|-------------|
 | `spaces` | `Record<string, string>` | `{ default: DATABRICKS_GENIE_SPACE_ID }` | Map of alias names to Genie Space IDs |
-| `timeout` | `number` | `120000` | Polling timeout in ms. Set to `0` for indefinite |
+| `agents` | `Record<string, string>` | `{}` | Map of alias names to Genie Space IDs used in [agent mode](#agent-mode) |
+| `timeout` | `number` | `120000` | Timeout in ms for a chat message or an agent turn. Set to `0` for indefinite |
 
 ### Space aliases
 
@@ -48,6 +51,26 @@ genie({
 
 If you omit `spaces`, the plugin reads `DATABRICKS_GENIE_SPACE_ID` from the environment and registers it under the `default` alias.
 
+### Agent mode
+
+Genie agent mode answers like an agent: it reasons about the question, runs queries, and draws charts. Register a space under `agents` instead of `spaces` to use it in agent mode. Any Genie space can be used as an agent:
+
+```ts
+genie({
+  spaces: { sales: "01ABCDEF12345678" },
+  agents: { analyst: "01ABCDEF12345678" },
+});
+```
+
+An agent alias works with the same routes, the same `sendMessage` and `getConversation` exports, and the same `<GenieChat alias="analyst" />` component as a chat space. The difference is in what it returns:
+
+- `reasoning` events stream the agent's reasoning while it works. The completed message carries the steps in `reasoning`.
+- Visualizations are always turned on. The charts the agent draws arrive as `viz` attachments, like a chat space's charts (see [Get a visualization](#get-a-visualization)).
+
+An alias can be either a chat space or an agent, not both. To use one space both ways, give it two aliases, as in the example above. Conversations don't carry over between the two modes.
+
+Agent mode uses the Genie agent API (`/api/2.0/genie/agents/{id}/responses` and `/api/2.0/genie/agents/{id}/conversations/{id}/items`), which the Databricks SDK doesn't cover yet.
+
 ### Finding your Genie Space ID
 
 You can find the Space ID from the **About** tab on your Genie space page in Databricks:
@@ -70,6 +93,7 @@ The genie plugin exposes these endpoints (mounted under `/api/genie`):
 
 - `POST /api/genie/:alias/messages` — Send a message to a Genie space (SSE stream)
 - `GET /api/genie/:alias/conversations/:conversationId` — Replay conversation history (SSE stream)
+- `GET /api/genie/:alias/conversations/:conversationId/messages/:messageId/attachments/:attachmentId/visualization` — Download a chart as a PNG
 
 ### Send a message
 
@@ -89,6 +113,7 @@ The response is an SSE stream that emits these event types:
 |------------|-------------|
 | `message_start` | Conversation and message IDs assigned |
 | `status` | Processing status updates (e.g. `ASKING_AI`, `EXECUTING_QUERY`) |
+| `reasoning` | Agent mode only: one step of the agent's reasoning |
 | `message_result` | Final message with text and query attachments |
 | `query_result` | Tabular data for a query attachment |
 | `error` | Error details |
@@ -101,6 +126,16 @@ GET /api/genie/:alias/conversations/:conversationId
 
 Returns an SSE stream of `message_result` and `query_result` events for all messages in the conversation.
 
+### Get a visualization
+
+```
+GET /api/genie/:alias/conversations/:conversationId/messages/:messageId/attachments/:attachmentId/visualization
+```
+
+Returns the PNG Genie rendered for a `viz` attachment, so it can be used directly as an `<img>` source. Use the `attachmentId` of the `viz` attachment, not its `queryAttachmentId`. The message must be completed.
+
+Chat spaces and agents both return `viz` attachments when Genie draws a chart. For a query with a `viz` attachment, `GenieChat` shows Genie's chart in place of the chart it would otherwise infer from the query result.
+
 ## Programmatic access
 
 The plugin exports `sendMessage` and `getConversation` for server-side use:

packages/appkit-ui/src/react/genie/__tests__/genie-chat-message.test.tsx

@@ -0,0 +1,114 @@
+import { cleanup, render, screen } from "@testing-library/react";
+import { afterEach, describe, expect, test, vi } from "vitest";
+
+vi.mock("../genie-query-visualization", () => ({
+  GenieQueryVisualization: (props: { chartImageUrl?: string }) => (
+    <div data-testid="query-result" data-chart-image={props.chartImageUrl} />
+  ),
+}));
+
+import { GenieChatMessage } from "../genie-chat-message";
+import type { GenieMessageItem, GenieStatementResponse } from "../types";
+
+const queryResult: GenieStatementResponse = {
+  manifest: { schema: { columns: [{ name: "country" }] } },
+  result: { dataArray: [["Japan"]] },
+};
+
+function assistant(overrides: Partial<GenieMessageItem>): GenieMessageItem {
+  return {
+    id: "msg-1",
+    role: "assistant",
+    content: "Japan has the most.",
+    status: "COMPLETED",
+    attachments: [],
+    queryResults: new Map(),
+    ...overrides,
+  };
+}
+
+describe("GenieChatMessage", () => {
+  afterEach(cleanup);
+
+  test("hands Genie's chart to the query result, which keeps its table", () => {
+    render(
+      <GenieChatMessage
+        message={assistant({
+          attachments: [
+            { attachmentId: "query-1", query: { title: "Franchises" } },
+            {
+              attachmentId: "viz-1",
+              viz: {
+                title: "Franchises by Country",
+                queryAttachmentId: "query-1",
+              },
+            },
+          ],
+          queryResults: new Map([["query-1", queryResult]]),
+          visualizationUrls: { "viz-1": "/api/genie/a/viz-1.png" },
+        })}
+      />,
+    );
+
+    const result = screen.getByTestId("query-result");
+    expect(result.getAttribute("data-chart-image")).toBe(
+      "/api/genie/a/viz-1.png",
+    );
+    expect(screen.queryByRole("img")).toBeNull();
+  });
+
+  test("shows Genie's chart as an image while there's no query result", () => {
+    render(
+      <GenieChatMessage
+        message={assistant({
+          attachments: [
+            { attachmentId: "query-1", query: { title: "Franchises" } },
+            {
+              attachmentId: "viz-1",
+              viz: {
+                title: "Franchises by Country",
+                queryAttachmentId: "query-1",
+              },
+            },
+          ],
+          visualizationUrls: { "viz-1": "/api/genie/a/viz-1.png" },
+        })}
+      />,
+    );
+
+    const img = screen.getByRole("img", { name: "Franchises by Country" });
+    expect(img.getAttribute("src")).toBe("/api/genie/a/viz-1.png");
+    expect(screen.queryByTestId("query-result")).toBeNull();
+  });
+
+  test("leaves the query result alone for a query without a visualization", () => {
+    render(
+      <GenieChatMessage
+        message={assistant({
+          attachments: [
+            { attachmentId: "query-1", query: { title: "Franchises" } },
+          ],
+          queryResults: new Map([["query-1", queryResult]]),
+        })}
+      />,
+    );
+
+    const result = screen.getByTestId("query-result");
+    expect(result.getAttribute("data-chart-image")).toBeNull();
+    expect(screen.queryByRole("img")).toBeNull();
+  });
+
+  test("shows the agent's reasoning", () => {
+    render(
+      <GenieChatMessage
+        message={assistant({
+          content: "",
+          reasoning: ["Counting franchises."],
+        })}
+      />,
+    );
+
+    expect(screen.getByText("Thinking")).toBeTruthy();
+    expect(screen.getByText("Counting franchises.")).toBeTruthy();
+  });
+});

packages/appkit-ui/src/react/genie/__tests__/genie-query-visualization.test.tsx

@@ -0,0 +1,66 @@
+import { cleanup, fireEvent, render, screen } from "@testing-library/react";
+import { afterEach, describe, expect, test, vi } from "vitest";
+
+vi.mock("../../charts/base", () => ({
+  BaseChart: () => <div data-testid="inferred-chart" />,
+}));
+
+import { GenieQueryVisualization } from "../genie-query-visualization";
+import type { GenieStatementResponse } from "../types";
+
+const data: GenieStatementResponse = {
+  manifest: {
+    schema: {
+      columns: [
+        { name: "country", typeName: "STRING" },
+        { name: "franchises", typeName: "INT" },
+      ],
+    },
+  },
+  result: {
+    dataArray: [
+      ["Japan", "20"],
+      ["US", "16"],
+    ],
+  },
+};
+
+describe("GenieQueryVisualization", () => {
+  afterEach(cleanup);
+
+  test("shows the chart image in the Chart tab and keeps the Table tab", () => {
+    render(
+      <GenieQueryVisualization
+        data={data}
+        chartImageUrl="/viz.png"
+        chartImageAlt="Franchises by Country"
+      />,
+    );
+
+    const img = screen.getByRole("img", { name: "Franchises by Country" });
+    expect(img.getAttribute("src")).toBe("/viz.png");
+    expect(screen.queryByTestId("inferred-chart")).toBeNull();
+    expect(screen.getByRole("tab", { name: "Table" })).toBeTruthy();
+    expect(screen.getByRole("cell", { name: "Japan" })).toBeTruthy();
+    expect(
+      screen.queryByRole("button", { name: "Change chart type" }),
+    ).toBeNull();
+  });
+
+  test("falls back to the inferred chart when the image fails to load", () => {
+    render(<GenieQueryVisualization data={data} chartImageUrl="/viz.png" />);
+
+    fireEvent.error(screen.getByRole("img"));
+
+    expect(screen.queryByRole("img")).toBeNull();
+    expect(screen.getByTestId("inferred-chart")).toBeTruthy();
+    expect(screen.getByRole("cell", { name: "Japan" })).toBeTruthy();
+  });
+
+  test("shows the inferred chart without a chart image", () => {
+    render(<GenieQueryVisualization data={data} />);
+
+    expect(screen.getByTestId("inferred-chart")).toBeTruthy();
+    expect(screen.queryByRole("img")).toBeNull();
+  });
+});