@dbx-tools/ui-mastra 0.9.129 → 0.9.132

package/src/react/chat-composer.tsx

@@ -256,6 +256,7 @@ type ChatComposerProps = {
   composerActions: ChatViewProps["composerActions"];
   modelSelectorActions: ChatViewProps["modelSelectorActions"];
   composerLeadingActions: ChatViewProps["composerLeadingActions"];
+  allowAttachments: NonNullable<ChatViewProps["allowAttachments"]>;
   isLoadingHistory: NonNullable<ChatViewProps["isLoadingHistory"]>;
   onClear: ChatViewProps["onClear"];
   onExportConversation: ChatViewProps["onExportConversation"];
@@ -284,6 +285,7 @@ export const ChatComposer = ({
   composerActions,
   modelSelectorActions,
   composerLeadingActions,
+  allowAttachments,
   isLoadingHistory,
   onClear,
   onExportConversation,
@@ -293,7 +295,9 @@ export const ChatComposer = ({
   const [files, setFiles] = useState<FileUIPart[]>([]);
   const [fileError, setFileError] = useState<string | null>(null);
   const [isPreparingFiles, setIsPreparingFiles] = useState(false);
+  const [isDraggingFiles, setIsDraggingFiles] = useState(false);
   const inputId = useId();
+  const dragDepthRef = useRef(0);
   const fileInputRef = useRef<HTMLInputElement>(null);
   const inputGroupRef = useRef<HTMLDivElement>(null);
   const textareaRef = useRef<HTMLTextAreaElement>(null);
@@ -324,8 +328,14 @@ export const ChatComposer = ({
     onResumeFollow();
   };
 
-  const addFiles = async (selected: FileList | null) => {
+  const addFiles = async (selected: FileList | readonly File[] | null) => {
     if (!selected?.length) return;
+    const selectedFiles = Array.from(selected);
+    const imageFiles = selectedFiles.filter((file) => file.type.startsWith("image/"));
+    if (imageFiles.length === 0) {
+      setFileError("Only image files can be attached");
+      return;
+    }
     const remaining = MAX_COMPOSER_IMAGES - files.length;
     if (remaining <= 0) {
       setFileError(`You can attach up to ${MAX_COMPOSER_IMAGES} images`);
@@ -335,12 +345,12 @@ export const ChatComposer = ({
     setFileError(null);
     try {
       const next = await Promise.all(
-        Array.from(selected)
-          .slice(0, remaining)
-          .map((file) => imageFileToUIPart(file)),
+        imageFiles.slice(0, remaining).map((file) => imageFileToUIPart(file)),
       );
       setFiles((current) => [...current, ...next]);
-      if (selected.length > remaining) {
+      if (selectedFiles.length !== imageFiles.length) {
+        setFileError("Only image files were added");
+      } else if (imageFiles.length > remaining) {
         setFileError(`Only the first ${remaining} images were added`);
       }
     } catch (error) {
@@ -376,9 +386,39 @@ export const ChatComposer = ({
         />
         <InputGroup
           ref={inputGroupRef}
-          className="rounded-2xl border-border/80 shadow-sm transition-shadow focus-within:shadow-md"
+          className={cn(
+            "rounded-2xl border-border/80 shadow-sm transition-shadow focus-within:shadow-md",
+            isDraggingFiles && "border-primary ring-2 ring-primary/20",
+          )}
+          onDragEnter={(event) => {
+            if (!event.dataTransfer.types.includes("Files")) return;
+            event.preventDefault();
+            if (!allowAttachments) return;
+            dragDepthRef.current += 1;
+            setIsDraggingFiles(true);
+          }}
+          onDragOver={(event) => {
+            if (!event.dataTransfer.types.includes("Files")) return;
+            event.preventDefault();
+            if (!allowAttachments) {
+              event.dataTransfer.dropEffect = "none";
+              return;
+            }
+            event.dataTransfer.dropEffect = "copy";
+          }}
+          onDragLeave={(event) => {
+            if (!event.dataTransfer.types.includes("Files")) return;
+            dragDepthRef.current = Math.max(0, dragDepthRef.current - 1);
+            if (dragDepthRef.current === 0) setIsDraggingFiles(false);
+          }}
+          onDrop={(event) => {
+            event.preventDefault();
+            dragDepthRef.current = 0;
+            setIsDraggingFiles(false);
+            if (allowAttachments) void addFiles(event.dataTransfer.files);
+          }}
         >
-          {files.length > 0 && (
+          {allowAttachments && files.length > 0 && (
             <InputGroupAddon
               align="block-start"
               className="max-w-full justify-start gap-2 overflow-x-auto pb-0"
@@ -415,6 +455,15 @@ export const ChatComposer = ({
             ref={textareaRef}
             value={input}
             onChange={(event) => setInput(event.target.value)}
+            onPaste={(event) => {
+              if (!allowAttachments) return;
+              const pastedImages = Array.from(event.clipboardData.files).filter((file) =>
+                file.type.startsWith("image/"),
+              );
+              if (pastedImages.length === 0) return;
+              event.preventDefault();
+              void addFiles(pastedImages);
+            }}
             onKeyDown={(event) => {
               if (event.key === "Enter" && !event.shiftKey) {
                 event.preventDefault();
@@ -425,7 +474,7 @@ export const ChatComposer = ({
             rows={1}
             className="field-sizing-fixed min-h-10 max-h-64 w-full flex-none overflow-y-auto px-4 pb-2 pt-4 text-sm"
           />
-          {fileError && (
+          {allowAttachments && fileError && (
             <div className="px-4 text-xs text-destructive" role="alert">
               {fileError}
             </div>
@@ -435,36 +484,42 @@ export const ChatComposer = ({
             className="flex w-full flex-row flex-wrap items-center justify-between gap-2 px-3 pb-3 pt-1"
           >
             <div className="flex min-w-0 max-w-full flex-wrap items-center gap-2">
-              <input
-                ref={fileInputRef}
-                type="file"
-                accept="image/*"
-                multiple
-                className="sr-only"
-                onChange={(event) => void addFiles(event.currentTarget.files)}
-              />
-              <Tooltip>
-                <TooltipTrigger asChild>
-                  <Button
-                    type="button"
-                    variant="outline"
-                    size="icon"
-                    className="size-7 rounded-full [&_svg]:size-3"
-                    onClick={() => fileInputRef.current?.click()}
-                    disabled={
-                      isLoadingHistory || isPreparingFiles || files.length >= MAX_COMPOSER_IMAGES
-                    }
-                    aria-label="Upload images"
-                  >
-                    {isPreparingFiles ? (
-                      <Spinner className="size-3" />
-                    ) : (
-                      <ImagePlusIcon className="size-3" />
-                    )}
-                  </Button>
-                </TooltipTrigger>
-                <TooltipContent>Upload images</TooltipContent>
-              </Tooltip>
+              {allowAttachments && (
+                <>
+                  <input
+                    ref={fileInputRef}
+                    type="file"
+                    accept="image/*"
+                    multiple
+                    className="sr-only"
+                    onChange={(event) => void addFiles(event.currentTarget.files)}
+                  />
+                  <Tooltip>
+                    <TooltipTrigger asChild>
+                      <Button
+                        type="button"
+                        variant="outline"
+                        size="icon"
+                        className="size-7 rounded-full [&_svg]:size-3"
+                        onClick={() => fileInputRef.current?.click()}
+                        disabled={
+                          isLoadingHistory ||
+                          isPreparingFiles ||
+                          files.length >= MAX_COMPOSER_IMAGES
+                        }
+                        aria-label="Upload images"
+                      >
+                        {isPreparingFiles ? (
+                          <Spinner className="size-3" />
+                        ) : (
+                          <ImagePlusIcon className="size-3" />
+                        )}
+                      </Button>
+                    </TooltipTrigger>
+                    <TooltipContent>Upload images</TooltipContent>
+                  </Tooltip>
+                </>
+              )}
               {composerLeadingActions}
               {onExportConversation && (
                 <ExportMenu

package/src/react/chat-sessions.ts

@@ -55,6 +55,14 @@ export function useChatSessions(activeThreadId: string | undefined) {
     sessionsRef.current.delete(threadId);
     setRevision((current) => current + 1);
   }, []);
+  /** Re-key one session, keeping its transcript and run state intact. */
+  const moveSession = useCallback((fromThreadId: string, toThreadId: string) => {
+    const session = sessionsRef.current.get(fromThreadId);
+    if (!session || fromThreadId === toThreadId) return;
+    sessionsRef.current.delete(fromThreadId);
+    sessionsRef.current.set(toThreadId, session);
+    setRevision((current) => current + 1);
+  }, []);
   const activeKey = sessionKey(activeThreadId);
   const activeSession = useMemo(() => getSession(activeKey), [activeKey, getSession, revision]);
   const streamingThreadIds = useMemo(() => {
@@ -70,6 +78,7 @@ export function useChatSessions(activeThreadId: string | undefined) {
     activeKey,
     activeSession,
     getSession,
+    moveSession,
     removeSession,
     resetSession,
     streamingThreadIds,

package/src/react/chat-stream.ts

@@ -59,6 +59,18 @@ export function closeOnTerminalChunk(
   });
 }
 
+/** Drain a run observation until its terminal finish or abort chunk arrives. */
+export async function waitForTerminalChunk(stream: ReadableStream<UIMessageChunk>): Promise<void> {
+  const reader = closeOnTerminalChunk(stream).getReader();
+  try {
+    while (!(await reader.read()).done) {
+      // Draining is the synchronization point; the replacement turn does not replay this output.
+    }
+  } finally {
+    reader.releaseLock();
+  }
+}
+
 /** Replace or append one assistant message without disturbing the transcript. */
 function upsertAssistant(messages: UIMessage[], assistant: UIMessage): UIMessage[] {
   const next = [...messages];

package/src/react/chat-view.tsx

@@ -35,6 +35,7 @@ export const ChatView = ({
   composerActions,
   modelSelectorActions,
   composerLeadingActions,
+  allowAttachments = true,
   onLoadMore,
   isLoadingMore = false,
   hasMore = false,
@@ -136,6 +137,7 @@ export const ChatView = ({
         composerActions={composerActions}
         modelSelectorActions={modelSelectorActions}
         composerLeadingActions={composerLeadingActions}
+        allowAttachments={allowAttachments}
         isLoadingHistory={isLoadingHistory}
         onClear={onClear}
         onExportConversation={onExportConversation}

package/src/react/embed-slots.tsx

@@ -1,3 +1,4 @@
+import { GenieQueryVisualization } from "@databricks/appkit-ui/react";
 import { markers, type ParsedMarker } from "@dbx-tools/shared-mastra";
 import { Spinner } from "@dbx-tools/ui/react";
 import { ClockIcon } from "lucide-react";
@@ -6,13 +7,16 @@ import { DataGrid, humanizeLabel } from "./data-grid.tsx";
 import { AssistantMarkdown } from "./markdown.tsx";
 import { normalizeChartOption } from "../support/chart-option.ts";
 import { useChartChrome } from "../support/chart-theme.ts";
-import { useChartFetch, useStatementFetch } from "../support/mastra-client.ts";
+import {
+  useChartFetch,
+  useStatementFetch,
+  useVisualizationFetch,
+} from "../support/mastra-client.ts";
 
 const ReactECharts = lazy(() => import("echarts-for-react"));
 
-// Inline embed slots: chart / data tables resolved from `[chart:<id>]`
-// and `[data:<id>]` markers in the assistant's prose, plus the splitter
-// that interleaves those slots with the surrounding markdown.
+// Inline embed slots resolved from markers in assistant prose, plus the
+// splitter that interleaves those slots with surrounding markdown.
 
 /**
  * Frame shared by every chart-slot state so the layout stays
@@ -170,11 +174,27 @@ const DataSlot = ({ statementId }: { statementId: string }) => {
   );
 };
 
-/** One slice of an assistant message: prose, a chart slot, or a data slot. */
+/**
+ * Native AppKit Genie visualization for a query Genie chose to visualize.
+ * The same statement payload backs its Chart and Table tabs.
+ */
+const GenieVisualizationSlot = ({ statementId }: { statementId: string }) => {
+  const { data, loading, error } = useVisualizationFetch(statementId);
+  if (!loading && !error && data === undefined) return <ExpiredSlot type="visualization" />;
+  if (!data) return null;
+  return (
+    <div className="not-prose my-3 max-w-full rounded border border-border bg-card p-3">
+      <GenieQueryVisualization data={data.visualization} />
+    </div>
+  );
+};
+
+/** One slice of an assistant message. */
 type RenderSegment =
   | { kind: "text"; text: string }
   | { kind: "chart"; chartId: string }
-  | { kind: "data"; statementId: string };
+  | { kind: "data"; statementId: string }
+  | { kind: "visualization"; statementId: string };
 
 /**
  * Map one parsed marker onto its render segment. The marker grammar
@@ -185,8 +205,7 @@ type RenderSegment =
  * collapses to an empty text segment: the marker is consumed so no
  * literal `[<type>:...]` leaks into the prose, but no slot renders and
  * no `/embed/<type>/:id` request fires. Only a UUID-shaped id with a
- * known, renderable type (charts need ECharts, data needs a Table)
- * resolves to a slot.
+ * known, renderable type resolves to a slot.
  */
 const markerSegment = (marker: ParsedMarker): RenderSegment => {
   if (!markers.isUuid(marker.id)) return { kind: "text", text: "" };
@@ -195,6 +214,8 @@ const markerSegment = (marker: ParsedMarker): RenderSegment => {
       return { kind: "chart", chartId: marker.id };
     case "data":
       return { kind: "data", statementId: marker.id };
+    case "visualization":
+      return { kind: "visualization", statementId: marker.id };
     default:
       return { kind: "text", text: "" };
   }
@@ -220,8 +241,7 @@ const splitTextWithEmbeds = (text: string): RenderSegment[] => {
 };
 
 /**
- * Render the assistant's markdown with chart and data tables
- * placed at their inline marker positions. Each prose segment
+ * Render the assistant's markdown with embeds at their marker positions. Each prose segment
  * is its own {@link AssistantMarkdown} so streaming chunks still
  * incrementally parse correctly; embed slots break the markdown
  * flow with full-width block elements.
@@ -234,6 +254,8 @@ const splitTextWithEmbeds = (text: string): RenderSegment[] => {
  *   - `[data:<statement_id>]` -> {@link DataSlot} fetches the
  *     statement rows OBO-scoped and renders an inline Table
  *     (or nothing on 404 / empty result).
+ *   - `[visualization:<statement_id>]` -> {@link GenieVisualizationSlot}
+ *     renders AppKit UI's native Genie Chart/Table view.
  *
  * Stale markers in persisted transcript text are silently
  * dropped on reload so the prose around them stays clean.
@@ -265,6 +287,14 @@ export const MarkdownWithEmbeds = ({
         if (seg.kind === "chart") {
           return <ChartSlot key={`c-${i}-${seg.chartId}`} chartId={seg.chartId} />;
         }
+        if (seg.kind === "visualization") {
+          return (
+            <GenieVisualizationSlot
+              key={`v-${i}-${seg.statementId}`}
+              statementId={seg.statementId}
+            />
+          );
+        }
         return <DataSlot key={`d-${i}-${seg.statementId}`} statementId={seg.statementId} />;
       })}
     </>

package/src/react/index.ts

@@ -45,7 +45,11 @@ export type {
 } from "../support/request-context.ts";
 export { dedupeSuggestions } from "./suggestions.ts";
 export { MastraChat, useMastraChat } from "./mastra-chat.tsx";
-export type { MastraChatProps, UseMastraChatOptions } from "./mastra-chat.tsx";
+export type {
+  MastraChatProps,
+  ThreadIdChangeReason,
+  UseMastraChatOptions,
+} from "./mastra-chat.tsx";
 export { ThreadSidebar } from "./thread-sidebar.tsx";
 export type { ThreadSidebarProps } from "./thread-sidebar.tsx";
 export { ThreadActions, ThreadTabs } from "./thread-tabs.tsx";