Add annotation tool to highlight changed elements in preview

backend/agent/tools/definitions.py

@@ -101,6 +101,32 @@ def _retrieve_option_schema() -> Dict[str, Any]:
     }
 
 
+def _annotate_schema() -> Dict[str, Any]:
+    return {
+        "type": "object",
+        "properties": {
+            "annotations": {
+                "type": "array",
+                "items": {
+                    "type": "object",
+                    "properties": {
+                        "selector": {
+                            "type": "string",
+                            "description": "CSS selector that uniquely identifies the changed element in the HTML (e.g. '#hero-section', '.navbar', 'button.login-btn', 'header > nav').",
+                        },
+                        "description": {
+                            "type": "string",
+                            "description": "One-sentence description of what was changed.",
+                        },
+                    },
+                    "required": ["selector", "description"],
+                },
+            }
+        },
+        "required": ["annotations"],
+    }
+
+
 def canonical_tool_definitions(
     image_generation_enabled: bool = True,
 ) -> List[CanonicalToolDefinition]:
@@ -152,6 +178,18 @@ def canonical_tool_definitions(
                 ),
                 parameters=_retrieve_option_schema(),
             ),
+            CanonicalToolDefinition(
+                name="annotate",
+                description=(
+                    "Annotate the changed elements to highlight them in the preview. "
+                    "Call this tool exactly once, AFTER all code edits are complete "
+                    "and right before returning your final response to the user. "
+                    "Each annotation requires a CSS selector that uniquely targets "
+                    "the changed element in the HTML, plus a short one-sentence "
+                    "description of what changed."
+                ),
+                parameters=_annotate_schema(),
+            ),
         ]
     )
     return tools

backend/agent/tools/runtime.py

@@ -49,6 +49,8 @@ async def execute(self, tool_call: ToolCall) -> ToolExecutionResult:
             return await self._remove_background(tool_call.arguments)
         if tool_call.name == "retrieve_option":
             return self._retrieve_option(tool_call.arguments)
+        if tool_call.name == "annotate":
+            return self._annotate(tool_call.arguments)
         return ToolExecutionResult(
             ok=False,
             result={"error": f"Unknown tool: {tool_call.name}"},
@@ -357,6 +359,33 @@ def coerce_int(value: Any) -> Optional[int]:
         result = {"option_number": resolved_index + 1, "code": code}
         return ToolExecutionResult(ok=True, result=result, summary=summary)
 
+    def _annotate(self, args: Dict[str, Any]) -> ToolExecutionResult:
+        annotations = args.get("annotations")
+        if not isinstance(annotations, list) or not annotations:
+            return ToolExecutionResult(
+                ok=False,
+                result={"error": "annotate requires a non-empty annotations list"},
+                summary={"error": "Missing annotations"},
+            )
+
+        cleaned: List[Dict[str, str]] = []
+        for annotation in annotations:
+            selector = ensure_str(annotation.get("selector"))
+            description = ensure_str(annotation.get("description"))
+            if selector and description:
+                cleaned.append({"selector": selector, "description": description})
+
+        if not cleaned:
+            return ToolExecutionResult(
+                ok=False,
+                result={"error": "No valid annotations provided"},
+                summary={"error": "No valid annotations"},
+            )
+
+        summary = {"annotations": cleaned}
+        result = {"annotations": cleaned}
+        return ToolExecutionResult(ok=True, result=result, summary=summary)
+
 
 # Backwards-compatible alias for older imports.
 AgentToolbox = AgentToolRuntime

backend/agent/tools/summaries.py

@@ -69,4 +69,21 @@ def summarize_tool_input(tool_call: ToolCall, file_state: AgentFileState) -> Dic
             "index": args.get("index"),
         }
 
+    if tool_call.name == "annotate":
+        annotations = args.get("annotations") or []
+        if isinstance(annotations, list):
+            return {
+                "count": len(annotations),
+                "annotations": [
+                    {
+                        "selector": ensure_str(a.get("selector")),
+                        "description": summarize_text(
+                            ensure_str(a.get("description")), 160
+                        ),
+                    }
+                    for a in annotations
+                    if isinstance(a, dict)
+                ],
+            }
+
     return args

backend/prompts/system_prompt.py

@@ -18,6 +18,7 @@
 - When available, use generate_images to create image URLs from prompts (you may pass multiple prompts). The image generation AI is not capable of generating images with a transparent background.
 - Use remove_background to remove backgrounds from provided image URLs when needed (you may pass multiple image URLs).
 - Use retrieve_option to fetch the full HTML for a specific option (1-based option_number) when a user references another option.
+- After ALL code edits are complete and right before returning your final response, call annotate exactly once to highlight the elements you changed in the preview. Each annotation needs a CSS selector that uniquely targets the changed element in the HTML (e.g. '#hero-section', '.navbar', 'button.login-btn') and a one-sentence description of what was changed.
 
 
 # Stack-specific instructions

frontend/src/components/agent/AgentActivity.tsx

@@ -16,6 +16,7 @@ import {
   BsImage,
   BsScissors,
   BsFiles,
+  BsTag,
 } from "react-icons/bs";
 import ReactMarkdown from "react-markdown";
 import { Light as SyntaxHighlighterBase } from "react-syntax-highlighter";
@@ -107,6 +108,9 @@ function getEventIcon(type: AgentEventType, toolName?: string) {
   if (toolName === "retrieve_option") {
     return <BsFiles className="text-slate-500" />;
   }
+  if (toolName === "annotate") {
+    return <BsTag className="text-amber-500" />;
+  }
   return <BsFileEarmarkPlus className="text-gray-500" />;
 }
 
@@ -149,6 +153,11 @@ function getEventTitle(event: AgentEvent): string {
         ? "Retrieving option"
         : "Retrieved option";
     }
+    if (event.toolName === "annotate") {
+      return event.status === "running"
+        ? "Annotating changes"
+        : "Annotated changes";
+    }
     return event.status === "running" ? "Running tool" : "Tool completed";
   }
   return "Activity";
@@ -345,6 +354,33 @@ function renderToolDetails(event: AgentEvent, variantCode?: string) {
         </div>
       )}
 
+      {event.toolName === "annotate" && !hasError && (
+        <div className="space-y-1.5">
+          {(() => {
+            const annotations =
+              (output?.annotations as Array<{ selector: string; description: string }>) ||
+              (input?.annotations as Array<{ selector: string; description: string }>) ||
+              [];
+            return annotations.map((annotation, index) => (
+              <div
+                key={`${annotation.selector}-${index}`}
+                className="flex items-start gap-2 rounded-md border border-amber-200 dark:border-amber-800 bg-amber-50 dark:bg-amber-900/20 p-2"
+              >
+                <BsTag className="text-amber-500 mt-0.5 shrink-0" />
+                <div>
+                  <code className="text-xs font-mono text-amber-700 dark:text-amber-300 bg-amber-100 dark:bg-amber-900/40 px-1 rounded">
+                    {annotation.selector}
+                  </code>
+                  <span className="text-xs text-gray-600 dark:text-gray-400 ml-1.5">
+                    {annotation.description}
+                  </span>
+                </div>
+              </div>
+            ));
+          })()}
+        </div>
+      )}
+
       {!event.toolName && !hasError && (
         <>
           {event.input && (

frontend/src/components/preview/PreviewComponent.tsx

@@ -4,11 +4,17 @@ import useThrottle from "../../hooks/useThrottle";
 import { useAppStore } from "../../store/app-store";
 import { addHighlight, removeHighlight } from "../select-and-edit/utils";
 
+export interface Annotation {
+  selector: string;
+  description: string;
+}
+
 interface Props {
   code: string;
   device: "mobile" | "desktop";
   onScaleChange?: (scale: number) => void;
   viewMode?: "fit" | "actual";
+  annotations?: Annotation[];
 }
 
 const MOBILE_VIEWPORT_WIDTH = 375;
@@ -19,6 +25,7 @@ function PreviewComponent({
   device,
   onScaleChange,
   viewMode,
+  annotations,
 }: Props) {
   const iframeRef = useRef<HTMLIFrameElement | null>(null);
   const wrapperRef = useRef<HTMLDivElement | null>(null);
@@ -158,6 +165,67 @@ function PreviewComponent({
     }
   }, [throttledCode]);
 
+  // Apply annotation highlights inside the iframe when annotations change
+  useEffect(() => {
+    const iframe = iframeRef.current;
+    if (!iframe || !annotations || annotations.length === 0) return;
+
+    const applyAnnotations = () => {
+      const doc = iframe.contentWindow?.document;
+      if (!doc) return;
+
+      // Remove any previously injected annotation styles
+      doc.getElementById("__annotation-styles")?.remove();
+
+      // Inject a style tag with a pulsing highlight animation
+      const style = doc.createElement("style");
+      style.id = "__annotation-styles";
+      style.textContent = `
+        @keyframes __annotation-pulse {
+          0%, 100% { outline-color: rgba(245, 158, 11, 0.8); }
+          50% { outline-color: rgba(245, 158, 11, 0.3); }
+        }
+        .__annotation-highlight {
+          outline: 2px solid rgba(245, 158, 11, 0.8);
+          outline-offset: 2px;
+          animation: __annotation-pulse 2s ease-in-out 3;
+        }
+      `;
+      doc.head.appendChild(style);
+
+      // Query each selector and apply the highlight class
+      for (const annotation of annotations) {
+        try {
+          const el = doc.querySelector(annotation.selector);
+          if (el) {
+            el.classList.add("__annotation-highlight");
+            el.setAttribute("title", annotation.description);
+          }
+        } catch {
+          // Invalid selector — skip silently
+        }
+      }
+    };
+
+    // The iframe may not have loaded yet after srcdoc changes, so listen for load
+    iframe.addEventListener("load", applyAnnotations);
+    // Also try immediately in case the iframe is already loaded
+    applyAnnotations();
+
+    return () => {
+      iframe.removeEventListener("load", applyAnnotations);
+      // Clean up highlights from iframe DOM
+      const doc = iframe.contentWindow?.document;
+      if (doc) {
+        doc.getElementById("__annotation-styles")?.remove();
+        doc.querySelectorAll(".__annotation-highlight").forEach((el) => {
+          el.classList.remove("__annotation-highlight");
+          el.removeAttribute("title");
+        });
+      }
+    };
+  }, [annotations]);
+
   return (
     <div
       className={`flex-1 min-h-0 relative ${

frontend/src/components/preview/PreviewPane.tsx

@@ -18,7 +18,7 @@ import { Button } from "../ui/button";
 import { useAppStore } from "../../store/app-store";
 import { useProjectStore } from "../../store/project-store";
 import { extractHtml } from "./extractHtml";
-import PreviewComponent from "./PreviewComponent";
+import PreviewComponent, { Annotation } from "./PreviewComponent";
 import { downloadCode } from "./download";
 
 function openInNewTab(code: string) {
@@ -69,6 +69,26 @@ function PreviewPane({ settings, onOpenVersions }: Props) {
       ? extractHtml(currentCode)
       : currentCode;
 
+  // Extract annotations from the selected variant's completed annotate tool events
+  const annotations: Annotation[] = useMemo(() => {
+    if (!currentCommit) return [];
+    const variant = currentCommit.variants[currentCommit.selectedVariantIndex];
+    const events = variant?.agentEvents || [];
+    for (const event of events) {
+      if (
+        event.type === "tool" &&
+        event.toolName === "annotate" &&
+        event.status === "complete"
+      ) {
+        const output = event.output as { annotations?: Array<{ selector: string; description: string }> } | undefined;
+        if (output?.annotations && Array.isArray(output.annotations)) {
+          return output.annotations;
+        }
+      }
+    }
+    return [];
+  }, [currentCommit]);
+
   return (
     <div className="flex-1 flex flex-col min-h-0">
       <Tabs
@@ -217,13 +237,15 @@ function PreviewPane({ settings, onOpenVersions }: Props) {
             device="desktop"
             onScaleChange={setDesktopScale}
             viewMode={desktopViewMode}
+            annotations={annotations}
           />
         </TabsContent>
         <TabsContent value="mobile" className="flex-1 min-h-0 mt-0 data-[state=active]:flex data-[state=active]:flex-col">
           <PreviewComponent
             code={previewCode}
             device="mobile"
             viewMode="actual"
+            annotations={annotations}
           />
         </TabsContent>
         <TabsContent value="code" className="flex-1 min-h-0 mt-0 overflow-auto">