move function to server side

spartan-vutrannguyen · spartan-vutrannguyen · commit 549bf6aaf46a · 2025-08-14T09:12:46.000+07:00
diff --git a/apps/web/client/src/app/project/[id]/_components/top-bar/project-breadcrumb.tsx b/apps/web/client/src/app/project/[id]/_components/top-bar/project-breadcrumb.tsx
@@ -3,9 +3,6 @@ import { useStateManager } from '@/components/store/state';
 import { transKeys } from '@/i18n/keys';
 import { api } from '@/trpc/react';
 import { Routes } from '@/utils/constants';
-import { uploadBlobToStorage } from '@/utils/supabase/client';
-import { STORAGE_BUCKETS } from '@onlook/constants';
-import { fromPreviewImg } from '@onlook/db';
 import { Button } from '@onlook/ui/button';
 import {
     DropdownMenu,
@@ -17,22 +14,20 @@ import {
 import { Icons } from '@onlook/ui/icons';
 import { toast } from '@onlook/ui/sonner';
 import { cn } from '@onlook/ui/utils';
-import { getScreenshotPath, getValidUrl } from '@onlook/utility';
 import { observer } from 'mobx-react-lite';
 import { useTranslations } from 'next-intl';
 import { redirect, useRouter } from 'next/navigation';
 import { usePostHog } from 'posthog-js/react';
 import { useRef, useState } from 'react';
 import { RecentProjectsMenu } from './recent-projects';
-import { handleScrapeUrlTool } from '@/components/tools/handlers/web';
 
 export const ProjectBreadcrumb = observer(() => {
     const editorEngine = useEditorEngine();
     const stateManager = useStateManager();
     const posthog = usePostHog();
 
     const { data: project } = api.project.get.useQuery({ projectId: editorEngine.projectId });
-    const { mutateAsync: updateProject } = api.project.update.useMutation();
+    const { mutate: captureScreenshot } = api.project.captureScreenshot.useMutation();
     const t = useTranslations();
     const closeTimeoutRef = useRef<Timer | null>(null);
     const router = useRouter();
@@ -44,7 +39,7 @@ export const ProjectBreadcrumb = observer(() => {
         try {
             setIsClosingProject(true);
 
-            await captureProjectScreenshot();
+            captureProjectScreenshot();
         } catch (error) {
             console.error('Failed to take screenshots:', error);
         } finally {
@@ -55,71 +50,15 @@ export const ProjectBreadcrumb = observer(() => {
         }
     }
 
-    async function captureProjectScreenshot() {
-        if (!project?.sandbox?.url) {
-            console.error('No sandbox URL found');
-            return;
-        }
-
-        if (!project.id) {
+    function captureProjectScreenshot() {
+        if (!project?.id) {
             console.error('No project ID found');
             return;
         }
-
-        const screenShot = await handleScrapeUrlTool({
-            url: project.sandbox.url,
-            formats: ['screenshot'],
-            actions: [
-                {
-                    type: 'screenshot',
-                    fullPage: true,
-                },
-            ],
-            onlyMainContent: true,
-        });
-
-        const validUrl = getValidUrl(screenShot);
-
-        if (!validUrl) {
-            console.error('Invalid screenshot URL');
-            return;
-        }
-
         try {
-            const response = await fetch(screenShot, { mode: 'cors' });
-            if (response.ok) {
-                const blob = await response.blob();
-                const mimeType = blob.type || 'image/png';
-                const uploaded = await uploadBlobToStorage(
-                    STORAGE_BUCKETS.PREVIEW_IMAGES,
-                    getScreenshotPath(project.id, mimeType),
-                    blob,
-                    { contentType: mimeType },
-                );
-
-                if (uploaded) {
-                    const dbPreviewImg = fromPreviewImg({
-                        type: 'storage',
-                        storagePath: {
-                            bucket: STORAGE_BUCKETS.PREVIEW_IMAGES,
-                            path: uploaded.path,
-                        },
-                    });
-                    if (project?.metadata) {
-                        await updateProject({
-                            id: project.id,
-                            project: {
-                                ...dbPreviewImg,
-                            },
-                        });
-                    }
-                    return;
-                }
-            } else {
-                console.error('Failed to fetch screenshot URL:', response.statusText);
-            }
-        } catch (err) {
-            console.error('Error handling screenshot URL upload:', err);
+            captureScreenshot({ projectId: project.id });
+        } catch (error) {
+            console.error('Failed to capture screenshot on server:', error);
         }
     }
 
diff --git a/apps/web/client/src/components/tools/handlers/web.ts b/apps/web/client/src/components/tools/handlers/web.ts
@@ -17,7 +17,6 @@ export async function handleScrapeUrlTool(
             includeTags: args.includeTags,
             excludeTags: args.excludeTags,
             waitFor: args.waitFor,
-            actions: args.actions,
         });
 
         if (!result.result) {
diff --git a/apps/web/client/src/server/api/routers/code.ts b/apps/web/client/src/server/api/routers/code.ts
@@ -43,11 +43,7 @@ export const codeRouter = createTRPCRouter({
     scrapeUrl: protectedProcedure
         .input(z.object({
             url: z.string().url(),
-            formats: z.array(z.enum(['markdown', 'html', 'json', 'screenshot'])).default(['markdown']),
-            actions: z.array(z.object({
-                type: z.enum(['screenshot']),
-                fullPage: z.boolean().default(false).optional(),
-            })).optional(),
+            formats: z.array(z.enum(['markdown', 'html', 'json'])).default(['markdown']),
             onlyMainContent: z.boolean().default(true),
             includeTags: z.array(z.string()).optional(),
             excludeTags: z.array(z.string()).optional(),
@@ -64,7 +60,6 @@ export const codeRouter = createTRPCRouter({
                 const result = await app.scrapeUrl(input.url, {
                     formats: input.formats,
                     onlyMainContent: input.onlyMainContent,
-                    ...(input.actions && { actions: input.actions }),
                     ...(input.includeTags && { includeTags: input.includeTags }),
                     ...(input.excludeTags && { excludeTags: input.excludeTags }),
                     ...(input.waitFor !== undefined && { waitFor: input.waitFor }),
@@ -76,11 +71,7 @@ export const codeRouter = createTRPCRouter({
 
                 // Return the primary content format (markdown by default)
                 // or the first available format if markdown isn't available
-                const markdown = result.markdown ?? result.html ?? JSON.stringify(result.json, null, 2);
-                // Return the screenshot if it exists
-                const screenshot = result.screenshot;
-
-                const content = markdown ?? screenshot;
+                const content = result.markdown ?? result.html ?? JSON.stringify(result.json, null, 2);
 
                 if (!content) {
                     throw new Error('No content was scraped from the URL');
diff --git a/apps/web/client/src/server/api/routers/project/project.ts b/apps/web/client/src/server/api/routers/project/project.ts
@@ -1,4 +1,6 @@
+import { env } from '@/env';
 import { trackEvent } from '@/utils/analytics/server';
+import FirecrawlApp from '@mendable/firecrawl-js';
 import { initModel } from '@onlook/ai';
 import {
     canvases,
@@ -25,9 +27,102 @@ import { and, eq, ne } from 'drizzle-orm';
 import { z } from 'zod';
 import { createTRPCRouter, protectedProcedure } from '../../trpc';
 import { projectCreateRequestRouter } from './createRequest';
+import { STORAGE_BUCKETS } from '@onlook/constants';
+import { fromPreviewImg } from '@onlook/db';
+import { getScreenshotPath, getValidUrl } from '@onlook/utility';
 
 export const projectRouter = createTRPCRouter({
     createRequest: projectCreateRequestRouter,
+    captureScreenshot: protectedProcedure
+        .input(z.object({ projectId: z.string() }))
+        .mutation(async ({ ctx, input }) => {
+            try {
+                if (!env.FIRECRAWL_API_KEY) {
+                    throw new Error('FIRECRAWL_API_KEY is not configured');
+                }
+
+                const project = await ctx.db.query.projects.findFirst({
+                    where: eq(projects.id, input.projectId),
+                });
+
+                if (!project) {
+                    throw new Error('Project not found');
+                }
+
+                if (!project.sandboxUrl) {
+                    throw new Error('No sandbox URL found');
+                }
+
+                const app = new FirecrawlApp({ apiKey: env.FIRECRAWL_API_KEY });
+
+                const result = await app.scrapeUrl(project.sandboxUrl, {
+                    formats: ['screenshot'],
+                    onlyMainContent: true,
+                    actions: [
+                        {
+                            type: 'screenshot',
+                            fullPage: true,
+                        },
+                    ],
+                });
+
+                if (!result.success) {
+                    throw new Error(`Failed to scrape URL: ${result.error || 'Unknown error'}`);
+                }
+
+                const screenshotUrlRaw = result.screenshot;
+                const screenshotUrl = screenshotUrlRaw ? getValidUrl(screenshotUrlRaw) : null;
+
+                if (!screenshotUrl) {
+                    throw new Error('Invalid screenshot URL');
+                }
+
+                const response = await fetch(screenshotUrl);
+                if (!response.ok) {
+                    throw new Error(`Failed to fetch screenshot: ${response.status} ${response.statusText}`);
+                }
+
+                const arrayBuffer = await response.arrayBuffer();
+                const mimeType = response.headers.get('content-type') ?? 'image/png';
+                const path = getScreenshotPath(project.id, mimeType);
+
+                const { data, error } = await ctx.supabase.storage
+                    .from(STORAGE_BUCKETS.PREVIEW_IMAGES)
+                    .upload(path, arrayBuffer, {
+                        contentType: mimeType,
+                    });
+
+                if (error) {
+                    throw new Error(`Supabase upload error: ${error.message}`);
+                }
+
+                if (!data) {
+                    throw new Error('No data returned from storage upload');
+                }
+
+                const dbPreviewImg = fromPreviewImg({
+                    type: 'storage',
+                    storagePath: {
+                        bucket: STORAGE_BUCKETS.PREVIEW_IMAGES,
+                        path: data.path,
+                    },
+                });
+
+                await ctx.db.update(projects)
+                    .set({
+                        previewImgUrl: dbPreviewImg.previewImgUrl,
+                        previewImgPath: dbPreviewImg.previewImgPath,
+                        previewImgBucket: dbPreviewImg.previewImgBucket,
+                        updatedAt: new Date(),
+                    })
+                    .where(eq(projects.id, project.id));
+
+                return { success: true, path: data.path };
+            } catch (error) {
+                console.error('Error capturing project screenshot:', error);
+                return { success: false, error: error instanceof Error ? error.message : 'Unknown error' };
+            }
+        }),
     list: protectedProcedure
         .input(z.object({
             limit: z.number().optional(),
diff --git a/packages/ai/src/tools/web.ts b/packages/ai/src/tools/web.ts
@@ -5,7 +5,7 @@ export const SCRAPE_URL_TOOL_NAME = 'scrape_url';
 export const SCRAPE_URL_TOOL_PARAMETERS = z.object({
     url: z.string().url().describe('The URL to scrape. Must be a valid HTTP or HTTPS URL.'),
     formats: z
-        .array(z.enum(['markdown', 'html', 'json', 'screenshot']))
+        .array(z.enum(['markdown', 'html', 'json']))
         .default(['markdown'])
         .describe('The formats to return the scraped content in. Defaults to markdown.'),
     onlyMainContent: z
@@ -26,14 +26,6 @@ export const SCRAPE_URL_TOOL_PARAMETERS = z.object({
         .number()
         .optional()
         .describe('Time in milliseconds to wait for the page to load before scraping.'),
-    actions: z
-        .array(
-            z.object({
-                type: z.enum(['screenshot']),
-                fullPage: z.boolean().default(false).optional(),
-            }),
-        )
-        .optional(),
 });
 
 export const scrapeUrlTool = tool({
