Rezk/dev

backend/package.json

@@ -32,6 +32,7 @@
     "dotenv": "^16.4.7",
     "express": "^5.1.0",
     "express-async-handler": "^1.2.0",
+    "express-rate-limit": "^8.3.1",
     "express-validator": "^7.3.1",
     "google-auth-library": "^10.5.0",
     "joi": "^18.0.2",

backend/public/prompt/design/instructions/icon-instructions.txt

@@ -12,11 +12,14 @@ When you need icons NOT in the reference list (or when no list is provided):
 Before generating the design, scan the entire layout and identify every icon you will need.
 Then call searchIcons for ALL of them simultaneously in a SINGLE response — do NOT search incrementally across multiple rounds.
 
-Each searchIcons call returns objects with 'id' and 'url' fields.
-Use the 'url' value directly as the imageUrl — no additional tool calls needed.
+Each searchIcons call returns an array of icon ID strings (e.g. "mdi:home", "lucide:arrow-right").
+Construct the icon URL from the ID using this pattern:
+https://api.iconify.design/{prefix}/{name}.svg
+where prefix and name are the two parts split by ':'.
+Example: "mdi:home" → https://api.iconify.design/mdi/home.svg
 
 Example result from searchIcons:
-  { "icons": [{ "id": "mdi:home", "url": "https://api.iconify.design/mdi/home.svg" }, ...] }
+  { "icons": ["mdi:home", "mdi:home-outline", "mdi:home-variant"] }
 
 Then use RECTANGLE with IMAGE fill:
 

backend/src/application/use-cases/ui-library/delete-ui-library-component.use-case.ts

@@ -9,6 +9,9 @@ export class DeleteUILibraryComponentUseCase {
 
     async execute(componentId: string, userId: string): Promise<void> {
         const component = await this.uiLibraryRepository.findComponentById(componentId, userId);
+        if (!component) {
+            throw new Error('Component not found');
+        }
 
         await this.uiLibraryRepository.deleteComponent(componentId, userId);
 

backend/src/application/use-cases/ui-library/delete-ui-library-project.use-case.ts

@@ -4,6 +4,10 @@ export class DeleteUILibraryProjectUseCase {
     constructor(private readonly uiLibraryRepository: IUILibraryRepository) {}
 
     async execute(projectId: string, userId: string): Promise<void> {
+        const project = await this.uiLibraryRepository.findProjectById(projectId, userId);
+        if (!project) {
+            throw new Error('Project not found');
+        }
         await this.uiLibraryRepository.deleteProject(projectId, userId);
     }
 }

backend/src/infrastructure/config/env.config.ts

@@ -52,4 +52,11 @@ export const ENV_CONFIG = {
     AWS_SECRET_ACCESS_KEY: process.env.AWS_SECRET_ACCESS_KEY!,
     AWS_REGION: process.env.AWS_REGION || 'us-east-1',
     AWS_S3_BUCKET: process.env.AWS_S3_BUCKET!,
+
+    // Icon & Tool Call Configuration
+    ICON_SEARCH_TIMEOUT_MS: Number(process.env.ICON_SEARCH_TIMEOUT_MS || 5000),
+    ICON_SEARCH_LIMIT: Number(process.env.ICON_SEARCH_LIMIT || 10),
+    ICON_CACHE_TTL_MS: Number(process.env.ICON_CACHE_TTL_MS || 600000), // 10 minutes
+    MAX_CONCURRENT_TOOL_CALLS: Number(process.env.MAX_CONCURRENT_TOOL_CALLS || 5),
+    MAX_TOOL_CALL_ROUNDS: Number(process.env.MAX_TOOL_CALL_ROUNDS || 3),
 };

backend/src/infrastructure/services/ai/ai-generate-design.service.ts

@@ -9,6 +9,7 @@ import { ConversationMessage, DesignGenerationResult, IAiDesignService } from '.
 
 import { iconTools } from '../../config/ai-tools.config';
 import { AIModelConfig, getModelById } from '../../config/ai-models.config';
+import { ENV_CONFIG } from '../../config/env.config';
 
 import { ToolCallHandlerService, FunctionToolCall } from './tool-call-handler.service';
 import { ResponseParserService } from './response-parser.service';
@@ -297,10 +298,13 @@ export class AiGenerateDesignService implements IAiDesignService {
         totalInputTokens += completion.usage?.prompt_tokens ?? 0;
         totalOutputTokens += completion.usage?.completion_tokens ?? 0;
 
-        // Handle tool calls loop
-        while (completion.choices[0]?.message?.tool_calls) {
+        // Handle tool calls loop (capped to prevent runaway token costs)
+        const maxRounds = ENV_CONFIG.MAX_TOOL_CALL_ROUNDS;
+        let round = 0;
+        while (completion.choices[0]?.message?.tool_calls && round < maxRounds) {
+            round++;
             const toolCalls = completion.choices[0].message.tool_calls as FunctionToolCall[];
-            console.log(`--- Processing ${toolCalls.length} tool calls ---`);
+            console.log(`--- Processing ${toolCalls.length} tool calls (round ${round}/${maxRounds}) ---`);
 
             const toolResults = await this.toolCallHandler.handleToolCalls(toolCalls);
             // Add assistant message with tool calls
@@ -324,6 +328,17 @@ export class AiGenerateDesignService implements IAiDesignService {
             totalOutputTokens += completion.usage?.completion_tokens ?? 0;
         }
 
+        if (round >= maxRounds && completion.choices[0]?.message?.tool_calls) {
+            console.warn(`⚠️ Tool call loop hit max rounds (${maxRounds}). Forcing final completion without tools.`);
+            // Request one final completion without tools to get the text response
+            completion = await openai.chat.completions.create({
+                model: aiModel.id,
+                messages: messages as any,
+            });
+            totalInputTokens += completion.usage?.prompt_tokens ?? 0;
+            totalOutputTokens += completion.usage?.completion_tokens ?? 0;
+        }
+
         const responseText = completion.choices[0]?.message?.content;
 
         if (!responseText) {

backend/src/infrastructure/services/ai/icon-extractor.service.ts

@@ -39,11 +39,11 @@ export class IconExtractorService {
      */
     normalizeName(name: string): string {
         return (name || '')
-        // .toLowerCase()
-        // .replace(/^(icon[s]?[\s/\-_]+|logo[s]?[\s/\-_]+|ic[\s/\-_]+)/i, '')
-        // .replace(/([\s/\-_]+icon[s]?$|[\s/\-_]+logo[s]?$)/i, '')
-        // .replace(/[\s\-_/]+/g, '')
-        // .trim();
+            .toLowerCase()
+            .replace(/^(icon[s]?[\s/\-_]+|logo[s]?[\s/\-_]+|ic[\s/\-_]+)/i, '')
+            .replace(/([\s/\-_]+icon[s]?$|[\s/\-_]+logo[s]?$)/i, '')
+            .replace(/[\s\-_/]+/g, '')
+            .trim();
     }
 
     private walk(node: any, map: Map<string, any>): void {

backend/src/infrastructure/services/ai/icon.service.ts

@@ -1,21 +1,35 @@
 // src/infrastructure/services/icon.service.ts
 
 import { IIconService, IconSearchResult } from '../../../domain/services/IIconService';
+import { ENV_CONFIG } from '../../config/env.config';
 
-const ICON_SEARCH_TIMEOUT_MS = 5000;
-const ICON_SEARCH_LIMIT = 10;
+interface CacheEntry {
+    result: IconSearchResult;
+    expiry: number;
+}
 
 export class IconService implements IIconService {
+    private cache = new Map<string, CacheEntry>();
 
     async searchIcons(query: string): Promise<IconSearchResult> {
-        console.log(`🔍 Searching icons for: ${query}`);
+        const normalizedQuery = query.trim().toLowerCase();
+        if (!normalizedQuery) return { icons: [] };
+
+        // Check cache first
+        const cached = this.cache.get(normalizedQuery);
+        if (cached && Date.now() < cached.expiry) {
+            console.log(`⚡ Icon cache hit for: "${normalizedQuery}"`);
+            return cached.result;
+        }
+
+        console.log(`🔍 Searching icons for: "${normalizedQuery}"`);
 
         const controller = new AbortController();
-        const timeout = setTimeout(() => controller.abort(), ICON_SEARCH_TIMEOUT_MS);
+        const timeout = setTimeout(() => controller.abort(), ENV_CONFIG.ICON_SEARCH_TIMEOUT_MS);
 
         try {
             const response = await fetch(
-                `https://api.iconify.design/search?query=${encodeURIComponent(query)}&limit=${ICON_SEARCH_LIMIT}`,
+                `https://api.iconify.design/search?query=${encodeURIComponent(normalizedQuery)}&limit=${ENV_CONFIG.ICON_SEARCH_LIMIT}`,
                 { signal: controller.signal }
             );
 
@@ -25,12 +39,24 @@ export class IconService implements IIconService {
 
             const data = await response.json() as { icons?: string[] };
             const icons: string[] = data.icons ?? [];
+            const result: IconSearchResult = { icons };
+
+            // Store in cache
+            this.cache.set(normalizedQuery, {
+                result,
+                expiry: Date.now() + ENV_CONFIG.ICON_CACHE_TTL_MS,
+            });
+
+            // Evict expired entries periodically (keep cache bounded)
+            if (this.cache.size > 200) {
+                this.evictExpired();
+            }
 
-            return { icons };
+            return result;
 
         } catch (error) {
             if (error instanceof Error && error.name === 'AbortError') {
-                console.warn(`⚠️ Icon search timed out for: ${query}`);
+                console.warn(`⚠️ Icon search timed out for: "${normalizedQuery}"`);
                 return { icons: [] };
             }
 
@@ -39,4 +65,13 @@ export class IconService implements IIconService {
             clearTimeout(timeout);
         }
     }
+
+    private evictExpired(): void {
+        const now = Date.now();
+        for (const [key, entry] of this.cache) {
+            if (now >= entry.expiry) {
+                this.cache.delete(key);
+            }
+        }
+    }
 }

backend/src/infrastructure/services/ai/tool-call-handler.service.ts

@@ -1,6 +1,7 @@
 // src/infrastructure/services/tool-call-handler.service.ts
 
 import { IIconService } from '../../../domain/services/IIconService';
+import { ENV_CONFIG } from '../../config/env.config';
 
 export interface FunctionToolCall {
     id: string;
@@ -24,8 +25,11 @@ export class ToolCallHandlerService {
         console.log(`🛠️ [tools] start — ${toolCalls.length} tool call(s)`);
         const start = Date.now();
 
-        const results = await Promise.all(
-            toolCalls.map(toolCall => this.handleSingleToolCall(toolCall))
+        // Run with concurrency throttle to avoid overwhelming external APIs
+        const results = await this.runWithConcurrency(
+            toolCalls,
+            tc => this.handleSingleToolCall(tc),
+            ENV_CONFIG.MAX_CONCURRENT_TOOL_CALLS
         );
 
         console.log(`✅ [tools] done — ${toolCalls.length} tool call(s) in ${Date.now() - start}ms`);
@@ -62,4 +66,32 @@ export class ToolCallHandlerService {
             content: result,
         };
     }
+
+    /**
+     * Runs async tasks with a concurrency limit.
+     * Avoids blasting external APIs with unbounded parallel requests.
+     */
+    private async runWithConcurrency<T, R>(
+        items: T[],
+        fn: (item: T) => Promise<R>,
+        limit: number
+    ): Promise<R[]> {
+        if (items.length <= limit) {
+            return Promise.all(items.map(fn));
+        }
+
+        const results: R[] = new Array(items.length);
+        let nextIndex = 0;
+
+        async function worker() {
+            while (nextIndex < items.length) {
+                const idx = nextIndex++;
+                results[idx] = await fn(items[idx]);
+            }
+        }
+
+        const workers = Array.from({ length: Math.min(limit, items.length) }, () => worker());
+        await Promise.all(workers);
+        return results;
+    }
 }

backend/src/infrastructure/services/storage/s3.service.ts

@@ -2,6 +2,47 @@ import { S3Client, PutObjectCommand, DeleteObjectCommand } from '@aws-sdk/client
 import { ENV_CONFIG } from '../../config/env.config';
 import { randomUUID } from 'crypto';
 
+const ALLOWED_MIME_TYPES: Record<string, { extension: string; magic: Buffer[] }> = {
+    'image/png': {
+        extension: 'png',
+        magic: [Buffer.from([0x89, 0x50, 0x4e, 0x47, 0x0d, 0x0a, 0x1a, 0x0a])],
+    },
+    'image/jpeg': {
+        extension: 'jpg',
+        magic: [Buffer.from([0xff, 0xd8, 0xff])],
+    },
+    'image/webp': {
+        extension: 'webp',
+        // RIFF????WEBP — bytes 0-3 are RIFF, bytes 8-11 are WEBP
+        magic: [Buffer.from([0x52, 0x49, 0x46, 0x46])],
+    },
+};
+
+const MAX_SIZE_BYTES = 5 * 1024 * 1024; // 5 MB
+
+function validateImageBuffer(buffer: Buffer, mimeType: string): void {
+    if (buffer.length > MAX_SIZE_BYTES) {
+        throw new Error(`File size ${buffer.length} exceeds the 5 MB limit`);
+    }
+
+    const spec = ALLOWED_MIME_TYPES[mimeType];
+    if (!spec) {
+        throw new Error(`MIME type "${mimeType}" is not allowed. Allowed types: ${Object.keys(ALLOWED_MIME_TYPES).join(', ')}`);
+    }
+
+    const matchesMagic = spec.magic.some(magic => buffer.slice(0, magic.length).equals(magic));
+
+    // WebP requires an additional check on bytes 8-11
+    if (mimeType === 'image/webp') {
+        const webpMarker = buffer.slice(8, 12).toString('ascii');
+        if (!matchesMagic || webpMarker !== 'WEBP') {
+            throw new Error('File content does not match declared MIME type image/webp');
+        }
+    } else if (!matchesMagic) {
+        throw new Error(`File content does not match declared MIME type ${mimeType}`);
+    }
+}
+
 export class S3Service {
     private readonly client: S3Client;
     private readonly bucket: string;
@@ -19,14 +60,17 @@ export class S3Service {
     }
 
     async uploadBase64Image(base64DataUrl: string, folder = 'component-previews'): Promise<string> {
-        const matches = base64DataUrl.match(/^data:(.+);base64,(.+)$/);
+        const matches = base64DataUrl.match(/^data:([a-zA-Z0-9][a-zA-Z0-9!#$&\-^_]+\/[a-zA-Z0-9][a-zA-Z0-9!#$&\-^_]+);base64,([A-Za-z0-9+/]+=*)$/);
         if (!matches) {
             throw new Error('Invalid base64 image format. Expected: data:<mime>;base64,<data>');
         }
 
         const mimeType = matches[1];
         const buffer = Buffer.from(matches[2], 'base64');
-        const extension = mimeType.split('/')[1] || 'png';
+
+        validateImageBuffer(buffer, mimeType);
+
+        const extension = ALLOWED_MIME_TYPES[mimeType].extension;
         const key = `${folder}/${randomUUID()}.${extension}`;
 
         await this.client.send(new PutObjectCommand({

backend/src/infrastructure/web/controllers/ui-library.controller.ts

@@ -62,7 +62,9 @@ export class UILibraryController {
                 message: 'Project deleted successfully',
             });
         } catch (error) {
-            next(error);
+            const message = error instanceof Error ? error.message : 'Failed to delete project';
+            const statusCode = message.toLowerCase().includes('not found') ? 404 : 500;
+            res.status(statusCode).json({ success: false, message });
         }
     }
 
@@ -119,7 +121,9 @@ export class UILibraryController {
                 message: 'Component deleted successfully',
             });
         } catch (error) {
-            next(error);
+            const message = error instanceof Error ? error.message : 'Failed to delete component';
+            const statusCode = message.toLowerCase().includes('not found') ? 404 : 500;
+            res.status(statusCode).json({ success: false, message });
         }
     }