Integrating web crawling in the custom theme generator (#2347)

e2e-tests/themes_management.spec.ts

@@ -271,3 +271,68 @@ test("themes management - AI generator flow", async ({ po }) => {
   await expect(po.page.getByText("AI Generated Theme")).toBeVisible();
   await expect(po.page.getByText("Created via AI generator")).toBeVisible();
 });
+
+test("themes management - AI generator from website URL", async ({ po }) => {
+  await po.setUpDyadPro();
+
+  // Navigate to Themes page via Library sidebar
+  await po.goToLibraryTab();
+  await po.page.getByRole("link", { name: "Themes" }).click();
+  await expect(po.page.getByRole("heading", { name: "Themes" })).toBeVisible();
+
+  // Click New Theme button
+  await po.page.getByRole("button", { name: "New Theme" }).click();
+
+  // Wait for dialog to open
+  await expect(
+    po.page.getByRole("dialog").getByText("Create Custom Theme"),
+  ).toBeVisible();
+
+  // Verify AI-Powered Generator tab is active by default
+  const aiTab = po.page.getByRole("tab", { name: "AI-Powered Generator" });
+  await expect(aiTab).toHaveAttribute("data-state", "active");
+
+  // Switch to Website URL input source
+  await po.page.getByRole("button", { name: "Website URL" }).click();
+
+  // Verify URL input is visible
+  const urlInput = po.page.getByLabel("Website URL");
+  await expect(urlInput).toBeVisible();
+
+  // Verify Generate button is disabled before entering URL
+  const generateButton = po.page.getByRole("button", {
+    name: "Generate Theme Prompt",
+  });
+  await expect(generateButton).toBeDisabled();
+
+  // Fill in theme details
+  await po.page.getByLabel("Theme Name").fill("Website Theme");
+  await po.page
+    .getByLabel("Description (optional)")
+    .fill("Generated from website");
+
+  // Enter a website URL
+  await urlInput.fill("https://example.com");
+
+  // Verify Generate button is now enabled
+  await expect(generateButton).toBeEnabled();
+
+  // Click Generate to get mock theme prompt (test mode returns mock response)
+  await generateButton.click();
+
+  // Wait for generation to complete - the generated prompt textarea should appear
+  await expect(po.page.locator("#ai-prompt")).toBeVisible({ timeout: 10000 });
+
+  // Verify the mock theme content is displayed (URL-specific mock)
+  await expect(po.page.getByText("Test Mode Theme (from URL)")).toBeVisible();
+
+  // Save the theme
+  await po.page.getByRole("button", { name: "Save Theme" }).click();
+
+  // Verify dialog closes and theme card appears
+  await expect(po.page.getByRole("dialog")).not.toBeVisible();
+  const themeCard = po.page.getByTestId("theme-card");
+  await expect(themeCard).toBeVisible();
+  await expect(themeCard.getByText("Website Theme")).toBeVisible();
+  await expect(themeCard.getByText("Generated from website")).toBeVisible();
+});

src/components/AIGeneratorTab.tsx

@@ -3,14 +3,21 @@ import { Button } from "@/components/ui/button";
 import { Input } from "@/components/ui/input";
 import { Textarea } from "@/components/ui/textarea";
 import { Label } from "@/components/ui/label";
-import { Loader2, Upload, X, Sparkles, Lock } from "lucide-react";
-import { useGenerateThemePrompt } from "@/hooks/useCustomThemes";
+import { Loader2, Upload, X, Sparkles, Lock, Link } from "lucide-react";
+import {
+  useGenerateThemePrompt,
+  useGenerateThemeFromUrl,
+} from "@/hooks/useCustomThemes";
 import { ipc } from "@/ipc/types";
 import { showError } from "@/lib/toast";
 import { toast } from "sonner";
 import { useUserBudgetInfo } from "@/hooks/useUserBudgetInfo";
 import { AiAccessBanner } from "./ProBanner";
-import type { ThemeGenerationMode, ThemeGenerationModel } from "@/ipc/types";
+import type {
+  ThemeGenerationMode,
+  ThemeGenerationModel,
+  ThemeInputSource,
+} from "@/ipc/types";
 
 // Image upload constants
 const MAX_FILE_SIZE = 10 * 1024 * 1024; // 10MB per image (raw file size)
@@ -60,8 +67,14 @@ export function AIGeneratorTab({
   // Track if dialog is open to prevent orphaned uploads from adding images after close
   const isDialogOpenRef = useRef(isDialogOpen);
 
+  // URL-based generation state
+  const [inputSource, setInputSource] = useState<ThemeInputSource>("images");
+  const [websiteUrl, setWebsiteUrl] = useState("");
+
   const generatePromptMutation = useGenerateThemePrompt();
-  const isGenerating = generatePromptMutation.isPending;
+  const generateFromUrlMutation = useGenerateThemeFromUrl();
+  const isGenerating =
+    generatePromptMutation.isPending || generateFromUrlMutation.isPending;
   const { userBudget } = useUserBudgetInfo();
 
   // Cleanup function to revoke blob URLs and delete temp files
@@ -92,16 +105,28 @@ export function AIGeneratorTab({
     isDialogOpenRef.current = isDialogOpen;
   }, [isDialogOpen]);
 
-  // Cleanup images when dialog closes
+  // Keep a ref to current images for cleanup without causing effect re-runs
+  const aiImagesRef = useRef<ThemeImage[]>([]);
+  useEffect(() => {
+    aiImagesRef.current = aiImages;
+  }, [aiImages]);
+
+  // Cleanup images and reset state when dialog closes
   useEffect(() => {
-    if (!isDialogOpen && aiImages.length > 0) {
-      cleanupImages(aiImages);
+    if (!isDialogOpen) {
+      // Use ref to get current images to avoid dependency on aiImages
+      const imagesToCleanup = aiImagesRef.current;
+      if (imagesToCleanup.length > 0) {
+        cleanupImages(imagesToCleanup);
         setAiImages([]);
+      }
       setAiKeywords("");
       setAiGenerationMode("inspired");
       setAiSelectedModel(DEFAULT_THEME_GENERATION_MODEL);
+      setInputSource("images");
+      setWebsiteUrl("");
     }
-  }, [isDialogOpen, aiImages, cleanupImages]);
+  }, [isDialogOpen, cleanupImages]);
 
   const handleImageUpload = useCallback(
     async (e: React.ChangeEvent<HTMLInputElement>) => {
@@ -219,6 +244,8 @@ export function AIGeneratorTab({
   );
 
   const handleGenerate = useCallback(async () => {
+    if (inputSource === "images") {
+      // Image-based generation
       if (aiImages.length === 0) {
         showError("Please upload at least one image");
         return;
@@ -238,12 +265,38 @@ export function AIGeneratorTab({
           `Failed to generate theme: ${error instanceof Error ? error.message : "Unknown error"}`,
         );
       }
+    } else {
+      // URL-based generation
+      if (!websiteUrl.trim()) {
+        showError("Please enter a website URL");
+        return;
+      }
+
+      try {
+        const result = await generateFromUrlMutation.mutateAsync({
+          url: websiteUrl,
+          keywords: aiKeywords,
+          generationMode: aiGenerationMode,
+          model: aiSelectedModel,
+        });
+
+        setAiGeneratedPrompt(result.prompt);
+        toast.success("Theme prompt generated from website");
+      } catch (error) {
+        showError(
+          `Failed to generate theme: ${error instanceof Error ? error.message : "Unknown error"}`,
+        );
+      }
+    }
   }, [
+    inputSource,
     aiImages,
+    websiteUrl,
     aiKeywords,
     aiGenerationMode,
     aiSelectedModel,
     generatePromptMutation,
+    generateFromUrlMutation,
     setAiGeneratedPrompt,
   ]);
 
@@ -291,7 +344,45 @@ export function AIGeneratorTab({
         />
       </div>
 
-      {/* Image Upload Section */}
+      {/* Input Source Toggle */}
+      <div className="space-y-3">
+        <Label>Reference Source</Label>
+        <div className="grid grid-cols-2 gap-4">
+          <button
+            type="button"
+            onClick={() => setInputSource("images")}
+            className={`flex flex-col items-center rounded-lg border p-3 text-center transition-colors ${
+              inputSource === "images"
+                ? "border-primary bg-primary/5"
+                : "hover:bg-muted/50"
+            }`}
+          >
+            <Upload className="h-5 w-5 mb-1" />
+            <span className="font-medium text-sm">Upload Images</span>
+            <span className="text-xs text-muted-foreground mt-1">
+              Use screenshots from your device
+            </span>
+          </button>
+          <button
+            type="button"
+            onClick={() => setInputSource("url")}
+            className={`flex flex-col items-center rounded-lg border p-3 text-center transition-colors ${
+              inputSource === "url"
+                ? "border-primary bg-primary/5"
+                : "hover:bg-muted/50"
+            }`}
+          >
+            <Link className="h-5 w-5 mb-1" />
+            <span className="font-medium text-sm">Website URL</span>
+            <span className="text-xs text-muted-foreground mt-1">
+              Extract design from a live website
+            </span>
+          </button>
+        </div>
+      </div>
+
+      {/* Image Upload Section - only shown when inputSource is "images" */}
+      {inputSource === "images" && (
         <div className="space-y-2">
           <Label>Reference Images</Label>
           <div
@@ -349,6 +440,25 @@ export function AIGeneratorTab({
             </div>
           )}
         </div>
+      )}
+
+      {/* URL Input Section - only shown when inputSource is "url" */}
+      {inputSource === "url" && (
+        <div className="space-y-2">
+          <Label htmlFor="website-url">Website URL</Label>
+          <Input
+            id="website-url"
+            type="url"
+            placeholder="https://example.com"
+            value={websiteUrl}
+            onChange={(e) => setWebsiteUrl(e.target.value)}
+            disabled={isGenerating}
+          />
+          <p className="text-xs text-muted-foreground">
+            Enter a website URL to extract its design system
+          </p>
+        </div>
+      )}
 
       {/* Keywords Input */}
       <div className="space-y-2">
@@ -452,14 +562,20 @@ export function AIGeneratorTab({
       {/* Generate Button */}
       <Button
         onClick={handleGenerate}
-        disabled={isGenerating || aiImages.length === 0}
+        disabled={
+          isGenerating ||
+          (inputSource === "images" && aiImages.length === 0) ||
+          (inputSource === "url" && !websiteUrl.trim())
+        }
         variant="secondary"
         className="w-full"
       >
         {isGenerating ? (
           <>
             <Loader2 className="mr-2 h-4 w-4 animate-spin" />
-            Generating prompt...
+            {inputSource === "url"
+              ? "Generating from website..."
+              : "Generating prompt..."}
           </>
         ) : (
           <>
@@ -481,9 +597,11 @@ export function AIGeneratorTab({
             placeholder="Generated prompt will appear here..."
           />
         ) : (
-          <div className="min-h-[100px] border rounded-md p-4 flex items-center justify-center text-muted-foreground text-sm">
-            No prompt generated yet. Upload images and click "Generate" to
-            create a theme prompt.
+          <div className="min-h-[100px] border rounded-md p-4 flex items-center justify-center text-muted-foreground text-sm text-center">
+            No prompt generated yet.{" "}
+            {inputSource === "images"
+              ? 'Upload images and click "Generate" to create a theme prompt.'
+              : 'Enter a website URL and click "Generate" to extract a theme.'}
           </div>
         )}
       </div>

src/hooks/useCustomThemes.ts

@@ -6,6 +6,7 @@ import type {
   UpdateCustomThemeParams,
   GenerateThemePromptParams,
   GenerateThemePromptResult,
+  GenerateThemeFromUrlParams,
 } from "@/ipc/types";
 import { queryKeys } from "@/lib/queryKeys";
 
@@ -92,3 +93,13 @@ export function useGenerateThemePrompt() {
     },
   });
 }
+
+export function useGenerateThemeFromUrl() {
+  return useMutation({
+    mutationFn: async (
+      params: GenerateThemeFromUrlParams,
+    ): Promise<GenerateThemePromptResult> => {
+      return ipc.template.generateThemeFromUrl(params);
+    },
+  });
+}

src/ipc/types/index.ts

@@ -243,8 +243,11 @@ export type {
   DeleteCustomThemeParams,
   ThemeGenerationMode,
   ThemeGenerationModel,
+  ThemeInputSource,
+  CrawlStatus,
   GenerateThemePromptParams,
   GenerateThemePromptResult,
+  GenerateThemeFromUrlParams,
   SaveThemeImageParams,
   SaveThemeImageResult,
   CleanupThemeImagesParams,

src/ipc/types/templates.ts

@@ -100,6 +100,14 @@ export const ThemeGenerationModelSchema = z.enum([
 ]);
 export type ThemeGenerationModel = z.infer<typeof ThemeGenerationModelSchema>;
 
+// Theme input source (images or URL)
+export const ThemeInputSourceSchema = z.enum(["images", "url"]);
+export type ThemeInputSource = z.infer<typeof ThemeInputSourceSchema>;
+
+// Crawl status for UI feedback
+export const CrawlStatusSchema = z.enum(["crawling", "complete", "error"]);
+export type CrawlStatus = z.infer<typeof CrawlStatusSchema>;
+
 export const GenerateThemePromptParamsSchema = z.object({
   imagePaths: z.array(z.string()),
   keywords: z.string(),
@@ -119,6 +127,31 @@ export type GenerateThemePromptResult = z.infer<
   typeof GenerateThemePromptResultSchema
 >;
 
+// URL-based theme generation params
+export const GenerateThemeFromUrlParamsSchema = z.object({
+  url: z
+    .string()
+    .url()
+    .refine(
+      (url) => {
+        try {
+          const parsed = new URL(url);
+          return parsed.protocol === "http:" || parsed.protocol === "https:";
+        } catch {
+          return false;
+        }
+      },
+      { message: "Only HTTP and HTTPS URLs are supported" },
+    ),
+  keywords: z.string(),
+  generationMode: ThemeGenerationModeSchema,
+  model: ThemeGenerationModelSchema,
+});
+
+export type GenerateThemeFromUrlParams = z.infer<
+  typeof GenerateThemeFromUrlParamsSchema
+>;
+
 export const SaveThemeImageParamsSchema = z.object({
   data: z.string(),
   filename: z.string(),
@@ -201,6 +234,12 @@ export const templateContracts = {
     output: GenerateThemePromptResultSchema,
   }),
 
+  generateThemeFromUrl: defineContract({
+    channel: "generate-theme-from-url",
+    input: GenerateThemeFromUrlParamsSchema,
+    output: GenerateThemePromptResultSchema,
+  }),
+
   saveThemeImage: defineContract({
     channel: "save-theme-image",
     input: SaveThemeImageParamsSchema,

src/pro/main/ipc/handlers/local_agent/tools/web_crawl.ts

@@ -9,7 +9,7 @@ const webCrawlSchema = z.object({
   url: z.string().describe("URL to crawl"),
 });
 
-const webCrawlResponseSchema = z.object({
+export const webCrawlResponseSchema = z.object({
   rootUrl: z.string(),
   html: z.string().optional(),
   markdown: z.string().optional(),