feat: frame-specific state based screenshots with preview

apps/web/client/src/app/project/[id]/_components/right-panel/chat-tab/chat-input/index.tsx

@@ -222,31 +222,29 @@ export const ChatInput = observer(({
                 return;
             }
 
-            const framesWithViews = editorEngine.frames.getAll().filter(f => !!f.view);
-
-            if (framesWithViews.length === 0) {
+            // Get the most recently interacted frame
+            const targetFrame = editorEngine.frames.getMostRecentlyInteractedFrame();
+
+            if (!targetFrame?.view) {
                 toast.error('No active frame available for screenshot');
                 return;
             }
 
             let screenshotData = null;
             let mimeType = 'image/jpeg';
 
-            for (const frame of framesWithViews) {
-                try {
-                    if (!frame.view?.captureScreenshot) {
-                        continue;
-                    }
-
-                    const result = await frame.view.captureScreenshot();
-                    if (result && result.data) {
-                        screenshotData = result.data;
-                        mimeType = result.mimeType || 'image/jpeg';
-                        break;
-                    }
-                } catch (frameError) {
-                    // Continue to next frame on error
+            try {
+                const result = await targetFrame.view.captureScreenshot();
+                if (result && result.data) {
+                    screenshotData = result.data;
+                    mimeType = result.mimeType || 'image/jpeg';
+                } else {
+                    throw new Error('No screenshot data returned');
                 }
+            } catch (frameError) {
+                console.error('Failed to capture screenshot from selected frame:', frameError);
+                toast.error('Failed to capture screenshot from selected frame');
+                return;
             }
 
             if (!screenshotData) {

apps/web/client/src/app/project/[id]/_components/right-panel/chat-tab/context-pills/draft-image-pill.tsx

@@ -1,7 +1,8 @@
 import { type MessageContext, MessageContextType } from '@onlook/models/chat';
 import { Icons } from '@onlook/ui/icons';
-import { motion } from 'motion/react';
-import React from 'react';
+import { motion, AnimatePresence } from 'motion/react';
+import React, { useState, useRef } from 'react';
+import { createPortal } from 'react-dom';
 import { getTruncatedName } from './helpers';
 
 export const DraftImagePill = React.forwardRef<
@@ -11,13 +12,80 @@ export const DraftImagePill = React.forwardRef<
         onRemove: () => void;
     }
 >(({ context, onRemove }, ref) => {
+    const [showPreview, setShowPreview] = useState(false);
+    const [pillPosition, setPillPosition] = useState({ x: 0, y: 0 });
+    const pillRef = useRef<HTMLSpanElement>(null);
+
+    const handleShowPreview = () => {
+        // Capture the pill's position RIGHT NOW before showing the modal
+        if (pillRef.current) {
+            const rect = pillRef.current.getBoundingClientRect();
+            setPillPosition({ 
+                x: rect.left + rect.width / 2, 
+                y: rect.top + rect.height / 2 
+            });
+        }
+        setShowPreview(true);
+    };
+
     if (context.type !== MessageContextType.IMAGE) {
-        console.warn('DraftingImagePill received non-image context');
+        console.warn('DraftImagePill received non-image context');
         return null;
     }
 
     return (
-        <motion.span
+        <>
+            {typeof document !== 'undefined' && createPortal(
+                <AnimatePresence>
+                    {showPreview && (
+                        <>
+                            {/* Backdrop */}
+                            <motion.div
+                                initial={{ opacity: 0 }}
+                                animate={{ opacity: 1 }}
+                                exit={{ opacity: 0 }}
+                                className="fixed inset-0 bg-black/60 backdrop-blur-sm z-[9999]"
+                                onClick={() => setShowPreview(false)}
+                                transition={{ duration: 0.3 }}
+                            />
+
+                            {/* Preview Image */}
+                            <motion.img
+                                src={context.content}
+                                alt="Screenshot preview"
+                                className="fixed z-[9999] object-contain pointer-events-auto"
+                                style={{
+                                    width: "28px",
+                                    height: "28px",
+                                    left: pillPosition.x,
+                                    top: pillPosition.y,
+                                    transformOrigin: "center",
+                                }}
+                                initial={{
+                                    scale: 0,
+                                    x: "-50%",
+                                    y: "-50%"
+                                }}
+                                animate={{ 
+                                    scale: Math.min(window.innerWidth * 0.5 / 28, window.innerHeight * 0.75 / 28),
+                                    x: `calc(50vw - ${pillPosition.x}px - 50%)`,
+                                    y: `calc(50vh - ${pillPosition.y}px - 50%)`
+                                }}
+                                exit={{
+                                    scale: 0,
+                                    x: "-50%",
+                                    y: "-50%"
+                                }}
+                                transition={{ type: "tween", duration: 0.3, ease: "easeOut" }}
+                                onClick={() => setShowPreview(false)}
+                            />
+                        </>
+                    )}
+                </AnimatePresence>,
+                document.body
+            )}
+
+            <motion.span
             layout="position"
             initial={{ opacity: 0, scale: 0.9 }}
             animate={{ opacity: 1, scale: 1 }}
@@ -29,9 +97,10 @@ export const DraftImagePill = React.forwardRef<
                     ease: 'easeOut',
                 },
             }}
-            className="group relative flex flex-row items-center gap-1 justify-center border bg-background-tertiary rounded-md h-7"
+            className="group relative flex flex-row items-center gap-1 justify-center border bg-background-tertiary rounded-md h-7 cursor-pointer"
             key={context.displayName}
-            ref={ref}
+            ref={pillRef}
+            onClick={handleShowPreview}
         >
             {/* Left side: Image thumbnail */}
             <div className="w-7 h-7 flex items-center justify-center overflow-hidden relative">
@@ -60,6 +129,7 @@ export const DraftImagePill = React.forwardRef<
                 <Icons.CrossL className="w-2.5 h-2.5 text-primary-foreground" />
             </button>
         </motion.span>
+        </>
     );
 });
 

apps/web/client/src/app/project/[id]/_components/right-panel/chat-tab/context-pills/input-context-pills.tsx

@@ -31,7 +31,7 @@ export const InputContextPills = observer(() => {
                         if (context.type === MessageContextType.IMAGE) {
                             return (
                                 <DraftImagePill
-                                    key={`image-${context.content}`}
+                                    key={`image-${index}`}
                                     context={context}
                                     onRemove={() => handleRemoveContext(context)}
                                 />

apps/web/client/src/components/store/editor/chat/context.ts

@@ -219,6 +219,16 @@ export class ChatContext {
         }
     }
 
+    addImageContext(imageData: string) {
+        const imageContext: ImageMessageContext = {
+            type: MessageContextType.IMAGE,
+            content: imageData,
+            displayName: `Screenshot ${new Date().toLocaleTimeString()}`,
+            mimeType: 'image/jpeg',
+        };
+        this.context = [...this.context.filter(c => c.type !== MessageContextType.IMAGE), imageContext];
+    }
+
     clearAttachments() {
         this.context = this.context.filter((context) => context.type !== MessageContextType.IMAGE);
     }

apps/web/client/src/components/store/editor/element/index.ts

@@ -66,14 +66,18 @@ export class ElementsManager {
         for (const domEl of domEls) {
             const frameData = this.editorEngine.frames.get(domEl.frameId);
             if (!frameData) {
-                console.error('Frame data not found');
+                console.error('Frame data not found for frameId:', domEl.frameId);
                 continue;
             }
             const { view } = frameData;
             if (!view) {
-                console.error('No frame view found');
+                console.error('No frame view found for frameId:', domEl.frameId);
                 continue;
             }
+
+            // Track frame interaction when element is clicked
+            this.editorEngine.frames.updateLastInteraction(domEl.frameId);
+
             const adjustedRect = adaptRectToCanvas(domEl.rect, view);
             const isComponent = !!domEl.instanceId;
             this.editorEngine.overlay.state.addClickRect(

apps/web/client/src/components/store/editor/frames/manager.ts

@@ -12,6 +12,7 @@ export interface FrameData {
     frame: Frame;
     view: WebFrameView | null;
     selected: boolean;
+    lastInteractionTime: Date | null;
 }
 
 function roundDimensions(frame: WebFrame): WebFrame {
@@ -47,7 +48,7 @@ export class FramesManager {
 
     applyFrames(frames: Frame[]) {
         for (const frame of frames) {
-            this._frameIdToData.set(frame.id, { frame, view: null, selected: false });
+            this._frameIdToData.set(frame.id, { frame, view: null, selected: false, lastInteractionTime: null });
         }
     }
 
@@ -69,7 +70,9 @@ export class FramesManager {
 
     registerView(frame: Frame, view: WebFrameView) {
         const isSelected = this.isSelected(frame.id);
-        this._frameIdToData.set(frame.id, { frame, view, selected: isSelected });
+        const existingData = this._frameIdToData.get(frame.id);
+        const lastInteractionTime = existingData?.lastInteractionTime || null;
+        this._frameIdToData.set(frame.id, { frame, view, selected: isSelected, lastInteractionTime });
         const framePathname = new URL(view.src).pathname;
         this._navigation.registerFrame(frame.id, framePathname);
     }
@@ -92,6 +95,8 @@ export class FramesManager {
         }
         for (const frame of frames) {
             this.updateFrameSelection(frame.id, !this.isSelected(frame.id));
+            // Track frame interaction when frame is selected
+            this.updateLastInteraction(frame.id);
         }
         this.notify();
     }
@@ -163,6 +168,9 @@ export class FramesManager {
         }
 
         try {
+            // Track frame interaction when navigating
+            this.updateLastInteraction(frameId);
+
             const currentUrl = frameData.view.src;
             const baseUrl = currentUrl ? new URL(currentUrl).origin : null;
 
@@ -215,7 +223,7 @@ export class FramesManager {
         const success = await api.frame.create.mutate(fromFrame(roundDimensions(frame)));
 
         if (success) {
-            this._frameIdToData.set(frame.id, { frame, view: null, selected: false });
+            this._frameIdToData.set(frame.id, { frame, view: null, selected: false, lastInteractionTime: null });
         } else {
             console.error('Failed to create frame');
         }
@@ -255,6 +263,7 @@ export class FramesManager {
                 ...existingFrame,
                 frame: newFrame,
                 selected: existingFrame.selected,
+                lastInteractionTime: existingFrame.lastInteractionTime,
             });
         }
         await this.saveToStorage(frameId, frame);
@@ -286,6 +295,31 @@ export class FramesManager {
         return this.selected.length > 0;
     }
 
+    updateLastInteraction(frameId: string) {
+        const frameData = this.get(frameId);
+        if (frameData) {
+            frameData.lastInteractionTime = new Date();
+            this._frameIdToData.set(frameId, frameData);
+            this.notify();
+        }
+    }
+
+    getMostRecentlyInteractedFrame(): FrameData | null {
+        const frames = this.getAll();
+        if (frames.length === 0) {
+            return null;
+        }
+
+        // Find frames with interaction times and sort by most recent
+        const framesWithInteractions = frames
+            .filter(f => f.lastInteractionTime)
+            .sort((a, b) => b.lastInteractionTime!.getTime() - a.lastInteractionTime!.getTime());
+
+        // Return most recently interacted frame, or first frame if none have interactions
+        const mostRecent = framesWithInteractions.length > 0 ? framesWithInteractions[0] : frames[0];
+        return mostRecent ?? null;
+    }
+
     async duplicateSelected() {
         for (const frame of this.selected) {
             await this.duplicate(frame.frame.id);

apps/web/client/src/components/store/editor/screenshot/index.tsx

@@ -39,6 +39,7 @@ export class ScreenshotManager {
                     return;
                 }
             }
+
             const result = await api.project.captureScreenshot.mutate({ projectId: this.editorEngine.projectId });
             if (!result || !result.success) {
                 throw new Error('Failed to capture screenshot');