Add chat screenshot export option (#22)

Author: ScriptSmith

ui/package.json

@@ -35,6 +35,7 @@
     "@tanstack/react-query": "^5.90.21",
     "@tanstack/react-table": "^8.21.3",
     "@tanstack/react-virtual": "^3.13.19",
+    "@zumer/snapdom": "^2.6.0",
     "clsx": "^2.1.1",
     "diff": "^8.0.3",
     "katex": "^0.16.38",

ui/pnpm-lock.yaml

@@ -6,6 +6,7 @@ import type { Conversation, MessageUsage, ModelInstance } from "@/components/cha
 import type { ModelInfo } from "@/components/ModelPicker/ModelPicker";
 import { PreferencesProvider } from "@/preferences/PreferencesProvider";
 import type { TotalUsageResult } from "@/stores/conversationStore";
+import { ToastProvider } from "@/components/Toast/Toast";
 import { TooltipProvider } from "@/components/Tooltip/Tooltip";
 
 import { ChatHeader } from "./ChatHeader";
@@ -63,11 +64,13 @@ const meta: Meta<typeof ChatHeader> = {
     (Story) => (
       <QueryClientProvider client={queryClient}>
         <PreferencesProvider>
-          <TooltipProvider>
-            <div className="w-full max-w-4xl mx-auto">
-              <Story />
-            </div>
-          </TooltipProvider>
+          <ToastProvider>
+            <TooltipProvider>
+              <div className="w-full max-w-4xl mx-auto">
+                <Story />
+              </div>
+            </TooltipProvider>
+          </ToastProvider>
         </PreferencesProvider>
       </QueryClientProvider>
     ),

ui/src/components/ChatHeader/ChatHeader.stories.tsx

@@ -1,3 +1,4 @@
+import { useMemo } from "react";
 import { useQuery } from "@tanstack/react-query";
 import {
   Check,
@@ -7,6 +8,7 @@ import {
   FileText,
   FolderOpen,
   GitFork,
+  Image,
   Maximize2,
   Minimize2,
   Trash2,
@@ -35,6 +37,9 @@ import {
   useWidescreenMode,
 } from "@/stores/chatUIStore";
 import { useUserProjects } from "@/hooks/useUserProjects";
+import { ScreenshotPreviewModal } from "@/components/ScreenshotRenderer/ScreenshotPreviewModal";
+import { ScreenshotRenderer } from "@/components/ScreenshotRenderer/ScreenshotRenderer";
+import { useScreenshotExport } from "@/hooks/useScreenshotExport";
 import { downloadConversation } from "@/utils/exportConversation";
 import { formatCost, formatTokens } from "@/utils/formatters";
 
@@ -106,6 +111,39 @@ export function ChatHeader({
   const widescreenMode = useWidescreenMode();
   const { setConversationMode, setModeConfig, toggleWidescreenMode } = useChatUIStore();
   const canExport = conversation && conversation.messages.length > 0;
+  const { isCapturing, screenshot, startCapture, onCaptureComplete, dismissPreview } =
+    useScreenshotExport();
+
+  // Build instance labels map for screenshot
+  const instanceLabels = useMemo(() => {
+    const map = new Map<string, string>();
+    for (const inst of selectedInstances) {
+      if (inst.label) map.set(inst.id, inst.label);
+    }
+    return map;
+  }, [selectedInstances]);
+
+  // Build message groups for screenshot (all messages, no hidden filtering)
+  const screenshotGroups = useMemo(() => {
+    if (!conversation) return [];
+    const groups: {
+      id: string;
+      userMessage: (typeof conversation.messages)[number];
+      assistantResponses: (typeof conversation.messages)[number][];
+    }[] = [];
+    const msgs = conversation.messages;
+    for (let i = 0; i < msgs.length; i++) {
+      const msg = msgs[i];
+      if (msg.role === "user") {
+        const responses: typeof msgs = [];
+        for (let j = i + 1; j < msgs.length && msgs[j].role !== "user"; j++) {
+          if (msgs[j].role === "assistant") responses.push(msgs[j]);
+        }
+        groups.push({ id: msg.id, userMessage: msg, assistantResponses: responses });
+      }
+    }
+    return groups;
+  }, [conversation]);
 
   // Fetch user projects for the project picker
   const { projects } = useUserProjects();
@@ -333,9 +371,36 @@ export function ChatHeader({
                     <FileText className="h-4 w-4" />
                     Markdown (readable)
                   </DropdownItem>
+                  <DropdownItem
+                    onClick={startCapture}
+                    disabled={isStreaming || isCapturing}
+                    className="gap-2"
+                  >
+                    <Image className="h-4 w-4" />
+                    Screenshot (PNG)
+                  </DropdownItem>
                 </DropdownContent>
               </Dropdown>
             )}
+            {isCapturing && canExport && (
+              <ScreenshotRenderer
+                title={conversation.title}
+                messageGroups={screenshotGroups}
+                instanceLabels={instanceLabels}
+                totalUsage={totalUsage}
+                titleGenerationUsage={conversation.titleGenerationUsage}
+                onComplete={onCaptureComplete}
+              />
+            )}
+            {screenshot && canExport && (
+              <ScreenshotPreviewModal
+                open
+                onClose={dismissPreview}
+                imageUrl={screenshot.url}
+                blob={screenshot.blob}
+                title={conversation.title}
+              />
+            )}
             {/* Fork button - hidden on mobile */}
             {canExport && onFork && (
               <Tooltip>

ui/src/components/ChatHeader/ChatHeader.tsx

@@ -209,6 +209,8 @@ interface MultiModelResponseProps {
   actionConfig?: ResponseActionConfig;
   /** History mode used when this message was sent (read-only display) */
   historyMode?: HistoryMode;
+  /** Force stacked layout regardless of global viewMode (used for screenshot export) */
+  forceStacked?: boolean;
 }
 
 /**
@@ -1059,6 +1061,7 @@ function MultiModelResponseComponent({
   selectedBest,
   actionConfig = DEFAULT_ACTION_CONFIG,
   historyMode,
+  forceStacked = false,
 }: MultiModelResponseProps) {
   // Use global UI state from store
   const viewMode = useViewMode();
@@ -1171,7 +1174,7 @@ function MultiModelResponseComponent({
   const hasHiddenResponses = hiddenResponses.length > 0;
 
   // "grid" = horizontal layout with fixed-width cards, "stacked" = vertical full-width
-  const useHorizontalLayout = viewMode === "grid" && displayedResponses.length > 1;
+  const useHorizontalLayout = !forceStacked && viewMode === "grid" && displayedResponses.length > 1;
 
   // Horizontal scroll navigation state
   const scrollContainerRef = useRef<HTMLDivElement>(null);

ui/src/components/MultiModelResponse/MultiModelResponse.tsx

@@ -0,0 +1,77 @@
+import { useCallback, useEffect, useRef, useState } from "react";
+import { Check, Copy, Download } from "lucide-react";
+
+import { Button } from "@/components/Button/Button";
+import {
+  Modal,
+  ModalClose,
+  ModalContent,
+  ModalFooter,
+  ModalHeader,
+  ModalTitle,
+} from "@/components/Modal/Modal";
+import { useToast } from "@/components/Toast/Toast";
+import { downloadBlob, generateScreenshotFilename } from "@/utils/exportScreenshot";
+
+interface ScreenshotPreviewModalProps {
+  open: boolean;
+  onClose: () => void;
+  imageUrl: string;
+  blob: Blob;
+  title: string;
+}
+
+export function ScreenshotPreviewModal({
+  open,
+  onClose,
+  imageUrl,
+  blob,
+  title,
+}: ScreenshotPreviewModalProps) {
+  const [copied, setCopied] = useState(false);
+  const toast = useToast();
+  const copyTimerRef = useRef<ReturnType<typeof setTimeout>>(undefined);
+
+  useEffect(() => {
+    return () => clearTimeout(copyTimerRef.current);
+  }, []);
+
+  const handleDownload = useCallback(() => {
+    downloadBlob(blob, generateScreenshotFilename(title));
+  }, [blob, title]);
+
+  const handleCopy = useCallback(async () => {
+    try {
+      await navigator.clipboard.write([new ClipboardItem({ [blob.type]: blob })]);
+      setCopied(true);
+      clearTimeout(copyTimerRef.current);
+      copyTimerRef.current = setTimeout(() => setCopied(false), 2000);
+    } catch {
+      toast.error("Copy failed", "Your browser may not support copying images");
+    }
+  }, [blob, toast]);
+
+  return (
+    <Modal open={open} onClose={onClose} className="max-w-5xl w-[92vw] max-h-[90vh] flex flex-col">
+      <ModalClose onClose={onClose} />
+      <ModalHeader>
+        <ModalTitle>Screenshot Preview</ModalTitle>
+      </ModalHeader>
+
+      <ModalContent className="flex-1 overflow-auto -mx-6 px-0">
+        <img src={imageUrl} alt={`Screenshot of ${title}`} className="w-full" />
+      </ModalContent>
+
+      <ModalFooter>
+        <Button variant="ghost" onClick={handleCopy} className="gap-2">
+          {copied ? <Check className="h-4 w-4 text-success" /> : <Copy className="h-4 w-4" />}
+          {copied ? "Copied" : "Copy to clipboard"}
+        </Button>
+        <Button onClick={handleDownload} className="gap-2">
+          <Download className="h-4 w-4" />
+          Download
+        </Button>
+      </ModalFooter>
+    </Modal>
+  );
+}