feat: added AI image tools for view and upload to project operations

apps/web/client/src/app/_components/hero/create.tsx

@@ -19,6 +19,7 @@ import { observer } from 'mobx-react-lite';
 import { AnimatePresence } from 'motion/react';
 import { useRouter } from 'next/navigation';
 import { useEffect, useRef, useState } from 'react';
+import { v4 as uuidv4 } from 'uuid';
 
 const SAVED_INPUT_KEY = 'create-input';
 interface CreateInputContext {
@@ -200,13 +201,14 @@ export const Create = observer(({
                 content: base64,
                 displayName: file.name,
                 mimeType: file.type,
+                id: uuidv4(),
             };
         } catch (error) {
             console.error('Error reading file:', error);
             return null;
         }
     };
-
+    
     const handleDragStateChange = (isDragging: boolean, e: React.DragEvent) => {
         const hasImage =
             e.dataTransfer.types.length > 0 &&

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

@@ -17,6 +17,7 @@ import { compressImageInBrowser } from '@onlook/utility';
 import { observer } from 'mobx-react-lite';
 import { useTranslations } from 'next-intl';
 import { useEffect, useMemo, useRef, useState } from 'react';
+import { v4 as uuidv4 } from 'uuid';
 import { validateImageLimit } from '../context-pills/helpers';
 import { InputContextPills } from '../context-pills/input-context-pills';
 import { type SuggestionsRef } from '../suggestions';
@@ -206,6 +207,7 @@ export const ChatInput = observer(({
             const compressedImage = await compressImageInBrowser(file);
             const base64URL = compressedImage ?? (event.target?.result as string);
             const contextImage: ImageMessageContext = {
+                id: uuidv4(),
                 type: MessageContextType.IMAGE,
                 content: base64URL,
                 mimeType: file.type,
@@ -258,6 +260,7 @@ export const ChatInput = observer(({
             }
 
             const contextImage: ImageMessageContext = {
+                id: uuidv4(),
                 type: MessageContextType.IMAGE,
                 content: screenshotData,
                 mimeType: mimeType,

apps/web/client/src/app/project/[id]/_components/right-panel/chat-tab/chat-messages/message-content/tool-call-display.tsx

@@ -217,6 +217,7 @@ export const ToolCallDisplay = ({
         );
     }
 
+
     return (
         <ToolCallSimple
             toolPart={toolPart}

apps/web/client/src/app/project/[id]/_hooks/use-start-project.tsx

@@ -14,6 +14,7 @@ import {
 import { toast } from '@onlook/ui/sonner';
 import { useEffect, useRef, useState } from 'react';
 import { useTabActive } from '../_hooks/use-tab-active';
+import { v4 as uuidv4 } from 'uuid';
 
 interface ProjectReadyState {
     canvas: boolean;
@@ -100,6 +101,7 @@ export const useStartProject = () => {
                     content: context.content,
                     mimeType: context.mimeType,
                     displayName: 'user image',
+                    id: uuidv4(),
                 }));
 
             const context: MessageContext[] = [...createContext, ...imageContexts];

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

@@ -1,6 +1,7 @@
 import { api } from "@/trpc/client";
 import { makeAutoObservable } from "mobx";
 import type { EditorEngine } from "../engine";
+import type { ChatMessage } from "@onlook/models";
 
 export class ApiManager {
     constructor(private editorEngine: EditorEngine) {
@@ -38,4 +39,8 @@ export class ApiManager {
     }) {
         return await api.utils.scrapeUrl.mutate(input);
     }
+
+    async getConversationMessages(conversationId: string): Promise<ChatMessage[]> {
+        return await api.chat.message.getAll.query({ conversationId });
+    }
 }

packages/ai/src/prompt/provider.ts

@@ -100,6 +100,13 @@ export function getHydratedUserMessage(
         prompt += branchPrompt;
     }
 
+    if (images.length > 0) {
+        const imageList = images
+            .map((img, idx) => `${idx + 1}. ${img.displayName} (ID: ${img.id || 'unknown'})`)
+            .join('\n');
+        prompt += wrapXml('available-images', imageList);
+    }
+
     const textContent = parts
         .filter((p) => p.type === 'text')
         .map((p) => p.text)

packages/ai/src/tools/classes/index.ts

@@ -15,5 +15,6 @@ export { SearchReplaceEditTool } from './search-replace-edit';
 export { SearchReplaceMultiEditFileTool } from './search-replace-multi-edit';
 export { TerminalCommandTool } from './terminal-command';
 export { TypecheckTool } from './typecheck';
+export { UploadImageTool } from './upload-image';
 export { WebSearchTool } from './web-search';
 export { WriteFileTool } from './write-file';

packages/ai/src/tools/classes/upload-image.ts

@@ -0,0 +1,118 @@
+import { Icons } from '@onlook/ui/icons';
+import type { EditorEngine } from '@onlook/web-client/src/components/store/editor/engine';
+import type { SandboxManager } from '@onlook/web-client/src/components/store/editor/sandbox';
+import { MessageContextType, type ImageMessageContext } from '@onlook/models';
+import mime from 'mime-lite';
+import { v4 as uuidv4 } from 'uuid';
+import { z } from 'zod';
+import { ClientTool } from '../models/client';
+import { BRANCH_ID_SCHEMA } from '../shared/type';
+
+export class UploadImageTool extends ClientTool {
+    static readonly toolName = 'upload_image';
+    static readonly description = "Uploads an image from the chat context to the project's file system. Use this tool when the user asks you to save, add, or upload an image to their project. The image will be stored in public/images/ directory by default and can be referenced in code. After uploading, you can use the file path in your code changes.";
+    static readonly parameters = z.object({
+        image_id: z
+            .string()
+            .describe(
+                'The unique ID of the image from the available images list',
+            ),
+        destination_path: z
+            .string()
+            .optional()
+            .describe('Destination path within the project. Defaults to "public/images" if not specified.'),
+        filename: z
+            .string()
+            .optional()
+            .describe('Custom filename (without extension). If not provided, a UUID will be generated'),
+        branchId: BRANCH_ID_SCHEMA,
+    });
+    static readonly icon = Icons.Image;
+
+    async handle(
+        args: z.infer<typeof UploadImageTool.parameters>,
+        editorEngine: EditorEngine
+    ): Promise<string> {
+        try {
+            const sandbox = editorEngine.branches.getSandboxById(args.branchId);
+            if (!sandbox) {
+                throw new Error(`Sandbox not found for branch ID: ${args.branchId}`);
+            }
+
+            // Get the current conversation ID
+            const conversationId = editorEngine.chat.getCurrentConversationId();
+            if (!conversationId) {
+                throw new Error('No active conversation found');
+            }
+
+            // Fetch all messages from the conversation
+            const messages = await editorEngine.api.getConversationMessages(conversationId);
+
+            // Find the image in message contexts by ID, searching from most recent to oldest
+            let imageContext: ImageMessageContext | null = null;
+            for (let i = messages.length - 1; i >= 0; i--) {
+                const message = messages[i];
+                if (!message || !message.metadata?.context) continue;
+
+                const contexts = message.metadata.context;
+                const imageContexts = contexts.filter(
+                    (ctx) => ctx.type === MessageContextType.IMAGE
+                );
+
+                // Find image by ID
+                const match = imageContexts.find((ctx) => ctx.id === args.image_id);
+
+                if (match) {
+                    imageContext = match;
+                    break;
+                }
+            }
+
+            if (!imageContext) {
+                throw new Error(`No image found with ID: ${args.image_id}`);
+            }
+
+            // Upload the image to the sandbox
+            const fullPath = await this.uploadImageToSandbox(imageContext, args, sandbox);
+
+            return `Image "${imageContext.displayName}" uploaded successfully to ${fullPath}`;
+        } catch (error) {
+            throw new Error(`Cannot upload image: ${error}`);
+        }
+    }
+
+    static getLabel(input?: z.infer<typeof UploadImageTool.parameters>): string {
+        if (input?.filename) {
+            return 'Uploading image ' + input.filename.substring(0, 20);
+        }
+        return 'Uploading image';
+    }
+
+    private async uploadImageToSandbox(
+        imageContext: ImageMessageContext,
+        args: z.infer<typeof UploadImageTool.parameters>,
+        sandbox: SandboxManager
+    ): Promise<string> {
+        const mimeType = imageContext.mimeType;
+        const extension = mime.getExtension(mimeType) || 'png';
+        const filename = args.filename ? `${args.filename}.${extension}` : `${uuidv4()}.${extension}`;
+        const destinationPath = args.destination_path?.trim() || 'public/images';
+        const fullPath = `${destinationPath}/${filename}`;
+
+        // Extract base64 data from the content (remove data URL prefix if present)
+        const base64Data = imageContext.content.replace(/^data:image\/[a-zA-Z0-9+.-]+;base64,/, '');
+        const binaryData = this.base64ToUint8Array(base64Data);
+
+        await sandbox.writeFile(fullPath, binaryData);
+        return fullPath;
+    }
+
+    private base64ToUint8Array(base64: string): Uint8Array {
+        const binaryString = atob(base64);
+        const bytes = new Uint8Array(binaryString.length);
+        for (let i = 0; i < binaryString.length; i++) {
+            bytes[i] = binaryString.charCodeAt(i);
+        }
+        return bytes;
+    }
+}

packages/ai/src/tools/toolset.ts

@@ -18,6 +18,7 @@ import {
     SearchReplaceMultiEditFileTool,
     TerminalCommandTool,
     TypecheckTool,
+    UploadImageTool,
     WebSearchTool,
     WriteFileTool,
 } from './classes';
@@ -53,6 +54,7 @@ const editOnlyToolClasses = [
     BashEditTool,
     SandboxTool,
     TerminalCommandTool,
+    UploadImageTool,
 ];
 const allToolClasses = [...readOnlyToolClasses, ...editOnlyToolClasses];
 

packages/ai/test/contexts/image-context.test.ts

@@ -1,5 +1,6 @@
 import { MessageContextType, type ImageMessageContext } from '@onlook/models';
 import { describe, expect, test } from 'bun:test';
+import { v4 as uuidv4 } from 'uuid';
 import { ImageContext } from '../../src/contexts/classes/image';
 
 describe('ImageContext', () => {
@@ -8,6 +9,7 @@ describe('ImageContext', () => {
         content: 'data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mP8/5+hHgAHggJ/PchI7wAAAABJRU5ErkJggg==',
         displayName: 'Screenshot.png',
         mimeType: 'image/png',
+        id: uuidv4(),
         ...overrides,
     });
 

packages/ai/test/contexts/index.test.ts

@@ -1,5 +1,5 @@
-import { 
-    MessageContextType, 
+import {
+    MessageContextType,
     type MessageContext,
     type FileMessageContext,
     type HighlightMessageContext,
@@ -10,9 +10,10 @@ import {
     type Branch,
 } from '@onlook/models';
 import { describe, expect, test } from 'bun:test';
-import { 
-    getContextPrompt, 
-    getContextLabel, 
+import { v4 as uuidv4 } from 'uuid';
+import {
+    getContextPrompt,
+    getContextLabel,
     getContextClass,
     FileContext,
     HighlightContext,
@@ -73,6 +74,7 @@ describe('Context Index', () => {
             content: 'data:image/png;base64,test-data',
             displayName: 'screenshot.png',
             mimeType: 'image/png',
+            id: uuidv4(),
         } as ImageMessageContext,
 
         agentRule: {

packages/models/src/chat/message/context.ts

@@ -38,6 +38,7 @@ export type HighlightMessageContext = BaseMessageContext & {
 export type ImageMessageContext = BaseMessageContext & {
     type: MessageContextType.IMAGE;
     mimeType: string;
+    id?: string;
 };
 
 export type ErrorMessageContext = BaseMessageContext & {