Update app.js

app.js

@@ -48,10 +48,65 @@ const sysPrompts = JSON.parse(fs.readFileSync('./prompts.json', 'utf8'));
 app.use(cors());
 app.use(bodyParser.json({ limit: '50mb' }));
 
+// --- CREDIT SYSTEM HELPERS ---
+
+const MIN_CREDITS_REQUIRED = 20;
+const IMAGE_GENERATION_COST = 1000; // Fixed token cost for an image
+
+/**
+ * Estimates token count based on text length (Standard: ~4 chars = 1 token)
+ */
+function estimateTokens(text) {
+    if (!text) return 0;
+    if (typeof text !== 'string') text = JSON.stringify(text);
+    return Math.ceil(text.length / 4);
+}
+
+/**
+ * Checks if user has enough credits. Throws error if insufficient.
+ */
+async function checkUserCredits(userId) {
+    if (!db) return; // Skip if no DB connected (Dev mode)
+    
+    const snap = await db.ref(`users/${userId}/credits`).once('value');
+    const credits = snap.val() || 0;
+
+    if (credits < MIN_CREDITS_REQUIRED) {
+        throw new Error(`Insufficient credits. You have ${credits}, but ${MIN_CREDITS_REQUIRED} are required.`);
+    }
+    return credits;
+}
+
+/**
+ * Deducts tokens from the user's account.
+ */
+async function deductUserCredits(userId, amount) {
+    if (!db || amount <= 0) return;
+    
+    try {
+        const ref = db.ref(`users/${userId}/credits`);
+        await ref.transaction((current_credits) => {
+            return (current_credits || 0) - amount;
+        });
+        console.log(`[Credits] Deducted ${amount} from User ${userId}`);
+    } catch (err) {
+        console.error(`[Credits] Failed to deduct ${amount} from ${userId}:`, err);
+    }
+}
+
+// --- MIDDLEWARE ---
+
 const validateRequest = (req, res, next) => {
-    if (req.path.includes('/onboarding') || req.path.includes('/admin/cleanup')) return next();
+    // Modified: Onboarding endpoints might now need userId for credit checks
+    if (req.path.includes('/admin/cleanup')) return next();
+    
     const { userId, projectId } = req.body;
-    if (!userId || !projectId) return res.status(400).json({ error: "Missing ID" });
+    
+    // For specific endpoints, ensure IDs exist
+    if ((req.path.includes('/new/project') || req.path.includes('/project/feedback')) && (!userId || !projectId)) {
+        return res.status(400).json({ error: "Missing userId or projectId" });
+    }
+    
     next();
 };
 
@@ -95,13 +150,25 @@ app.get('/admin/cleanup', async (req, res) => {
 // --- ONBOARDING ENDPOINTS ---
 
 app.post('/onboarding/analyze', validateRequest, async (req, res) => {
-    const { description } = req.body;
+    const { description, userId } = req.body; // userId now required for credit check
     if (!description) return res.status(400).json({ error: "Description required" });
+    if (!userId) return res.status(400).json({ error: "User ID required for analysis" });
+
+    let tokenUsage = 0;
 
     try {
+        // 1. Credit Check
+        await checkUserCredits(userId);
+
         console.log(`[Onboarding] Analyzing idea...`);
+        tokenUsage += estimateTokens(description);
+
         const result = await AIEngine.generateEntryQuestions(description);
+        tokenUsage += estimateTokens(JSON.stringify(result));
         
+        // 2. Deduct Credits
+        await deductUserCredits(userId, tokenUsage);
+
         if (result.status === "REJECTED") {
             return res.json({ 
                 rejected: true, 
@@ -111,37 +178,57 @@ app.post('/onboarding/analyze', validateRequest, async (req, res) => {
         res.json({ questions: result.questions });
     } catch (err) {
         console.error(err);
+        if (err.message.includes("Insufficient credits")) {
+            return res.status(402).json({ error: err.message });
+        }
         res.status(500).json({ error: "Analysis failed" });
     }
 });
 
 app.post('/onboarding/create', validateRequest, async (req, res) => {
     const { userId, description, answers } = req.body;
+    let tokenUsage = 0;
+
+    try {
+        // 1. Credit Check
+        await checkUserCredits(userId);
 
-    const randomHex = (n = 6) => {
-      const bytes = Math.ceil(n / 2);
-      return crypto.randomBytes(bytes).toString("hex").slice(0, n);
-    };
+        const randomHex = (n = 6) => {
+          const bytes = Math.ceil(n / 2);
+          return crypto.randomBytes(bytes).toString("hex").slice(0, n);
+        };
 
-    const projectId = `proj_${Date.now()}_${randomHex(7)}`;
+        const projectId = `proj_${Date.now()}_${randomHex(7)}`;
 
-    try {
         console.log(`[Onboarding] Grading Project ${projectId}...`);
+        
+        // Track Input Usage
+        tokenUsage += estimateTokens(description);
+        tokenUsage += estimateTokens(JSON.stringify(answers));
 
         const grading = await AIEngine.gradeProject(description, answers);
+        tokenUsage += estimateTokens(JSON.stringify(grading));
+
         const isFailure = grading.feasibility < 30 || grading.rating === 'F';
 
         let thumbnailBase64 = null;
 
         if (isFailure) {
             console.log(`[Onboarding] ❌ Project Failed Grading. Skipping Image.`);
-       
         } else {
             console.log(`[Onboarding] ✅ Passed. Generating Thumbnail...`);
             const imagePrompt = `Game Title: ${grading.title}. Core Concept: ${description}`;
+            
+            // Track Prompt Usage
+            tokenUsage += estimateTokens(imagePrompt);
+            
             thumbnailBase64 = await AIEngine.generateImage(imagePrompt);
+            
+            // Track Image Cost
+            if (thumbnailBase64) tokenUsage += IMAGE_GENERATION_COST;
         }
 
+        // --- Database & Storage Logic ---
         let thumbnailUrl = null;
         if (thumbnailBase64 && storage) {
             try {
@@ -227,6 +314,9 @@ app.post('/onboarding/create', validateRequest, async (req, res) => {
             await StateManager.updateProject(projectId, memoryObject);
         }
 
+        // 2. Deduct Credits
+        await deductUserCredits(userId, tokenUsage);
+
         res.json({ 
             success: !isFailure, 
             projectId, 
@@ -237,6 +327,9 @@ app.post('/onboarding/create', validateRequest, async (req, res) => {
 
     } catch (err) {
         console.error("Create Error:", err);
+        if (err.message.includes("Insufficient credits")) {
+            return res.status(402).json({ error: err.message });
+        }
         res.status(500).json({ error: "Creation failed" });
     }
 });
@@ -246,19 +339,27 @@ app.post('/onboarding/create', validateRequest, async (req, res) => {
 // NEW: Helper function to run initialization in background
 async function runBackgroundInitialization(projectId, userId, description) {
     console.log(`[Background] Starting initialization for ${projectId}`);
+    let tokenUsage = 0;
+
     try {
         const pmHistory = [];
         
         // 1. Generate GDD
         const gddPrompt = `Create a comprehensive GDD for: ${description}`;
+        tokenUsage += estimateTokens(gddPrompt);
+
         const gddResponse = await AIEngine.callPM(pmHistory, gddPrompt);
+        tokenUsage += estimateTokens(gddResponse);
         
         pmHistory.push({ role: 'user', parts: [{ text: gddPrompt }] });
         pmHistory.push({ role: 'model', parts: [{ text: gddResponse }] });
 
         // 2. Generate First Task
         const taskPrompt = "Based on the GDD, generate the first technical milestone.\nOutput format:\nTASK_NAME: <Name>\nWORKER_PROMPT: <Specific, isolated instructions for the worker>";
+        tokenUsage += estimateTokens(taskPrompt);
+
         const taskResponse = await AIEngine.callPM(pmHistory, taskPrompt);
+        tokenUsage += estimateTokens(taskResponse);
 
         pmHistory.push({ role: 'user', parts: [{ text: taskPrompt }] });
         pmHistory.push({ role: 'model', parts: [{ text: taskResponse }] });
@@ -267,7 +368,10 @@ async function runBackgroundInitialization(projectId, userId, description) {
         const initialWorkerInstruction = extractWorkerPrompt(taskResponse) || `Initialize structure for: ${description}`;
         const workerHistory = [];
         const initialWorkerPrompt = `CONTEXT: New Project. \nINSTRUCTION: ${initialWorkerInstruction}`;
+        tokenUsage += estimateTokens(initialWorkerPrompt);
+
         const workerResponse = await AIEngine.callWorker(workerHistory, initialWorkerPrompt, []);
+        tokenUsage += estimateTokens(workerResponse);
 
         workerHistory.push({ role: 'user', parts: [{ text: initialWorkerPrompt }] });
         workerHistory.push({ role: 'model', parts: [{ text: workerResponse }] });
@@ -281,146 +385,199 @@ async function runBackgroundInitialization(projectId, userId, description) {
             failureCount: 0
         });
 
-        await processAndQueueResponse(projectId, workerResponse);
+        // Queue response (Asset generation logic inside here is separate, handled by main queue but we can check usage there too if needed)
+        // For now, we assume processAndQueueResponse handles assets separately or we charge later.
+        // If processAndQueueResponse generates an image, we should charge.
+        // Let's modify processAndQueueResponse to return cost or handle it there.
+        // For simplicity, we stick to text cost here.
+        await processAndQueueResponse(projectId, workerResponse, userId); // Passed userId to handle image cost in queue
         
-        // 5. Update Status to IDLE (Ready for user input)
+        // 5. Update Status to IDLE
         if(db) await db.ref(`projects/${projectId}/info`).update({
             status: "IDLE",
             lastUpdated: Date.now()
         });
+
+        // 6. Deduct accumulated credits
+        await deductUserCredits(userId, tokenUsage);
         
-        console.log(`[Background] Initialization complete for ${projectId}`);
+        console.log(`[Background] Initialization complete for ${projectId}. Tokens used: ${tokenUsage}`);
 
     } catch (err) {
         console.error(`[Background] Init Error for ${projectId}:`, err);
+        // Even if it failed, we should deduct what was used
+      //  if (tokenUsage > 0) await deductUserCredits(userId, tokenUsage);
         if(db) await db.ref(`projects/${projectId}/info/status`).set("error");
     }
 }
 
-app.post('/new/project', validateRequest, (req, res) => {
+app.post('/new/project', validateRequest, async (req, res) => {
     const { userId, projectId, description } = req.body;
     
-    // 1. Immediately acknowledge request
-    // We set status to 'working' or 'initializing' immediately so frontend shows a loader
-    if(db) db.ref(`projects/${projectId}/info/status`).set("initializing");
-    console.log("Received new project request ")
-    res.json({ 
-        success: true, 
-        message: "Project initialization started in background."
-    });
-
-    // 2. Trigger background process (no await)
-    runBackgroundInitialization(projectId, userId, description);
+    try {
+        // 1. Credit Check (Prevent starting if low balance)
+        await checkUserCredits(userId);
+
+        // 2. Immediately acknowledge request
+        if(db) db.ref(`projects/${projectId}/info/status`).set("initializing");
+        console.log("Received new project request");
+        
+        res.json({ 
+            success: true, 
+            message: "Project initialization started in background."
+        });
+
+        // 3. Trigger background process
+        runBackgroundInitialization(projectId, userId, description);
+
+    } catch (err) {
+        console.error("New Project Error:", err);
+        if (err.message.includes("Insufficient credits")) {
+            return res.status(402).json({ error: err.message });
+        }
+        res.status(500).json({ error: "Failed to start project" });
+    }
 });
 
 app.post('/project/feedback', async (req, res) => {
     const { userId, projectId, prompt, hierarchyContext, scriptContext, logContext, taskComplete, images } = req.body;
+    let tokenUsage = 0;
     
+    // 1. Basic Validations
     const project = await StateManager.getProject(projectId);
-    
     if (!project) return res.status(404).json({ error: "Project not found." });
    
     if (project.userId !== userId) {
         console.warn(`[Security] Unauthorized ping for ${projectId}. Owner: ${project.userId}, Request: ${userId}`);
         return res.status(403).json({ error: "Unauthorized: You do not own this project." });
     }
-    
-    // Update status to working immediately
-    if(db) await db.ref(`projects/${projectId}/info/status`).set("working");
-      
-    if (taskComplete) {
-        console.log(`[${projectId}] ✅ TASK COMPLETE.`);
-        const summary = `Worker completed the previous task. Logs: ${logContext?.logs || "Clean"}. \nGenerate the NEXT task using 'WORKER_PROMPT:' format.`;
-        const pmResponse = await AIEngine.callPM(project.pmHistory, summary);
+
+    try {
+        // 2. Credit Check
+        await checkUserCredits(userId);
         
-        project.pmHistory.push({ role: 'user', parts: [{ text: summary }] });
-        project.pmHistory.push({ role: 'model', parts: [{ text: pmResponse }] });
-
-        const nextInstruction = extractWorkerPrompt(pmResponse);
-        if (!nextInstruction) {
-             await StateManager.updateProject(projectId, { pmHistory: project.pmHistory, status: "IDLE" });
-             if(db) {
-                 await db.ref(`projects/${projectId}/info`).update({
-                     status: "IDLE",
-                     lastUpdated: Date.now()
-                 });
-             }
-             return res.json({ success: true, message: "No further tasks. Project Idle." });
-        }
+        // Update status to working immediately
+        if(db) await db.ref(`projects/${projectId}/info/status`).set("working");
+          
+        if (taskComplete) {
+            console.log(`[${projectId}] ✅ TASK COMPLETE.`);
+            const summary = `Worker completed the previous task. Logs: ${logContext?.logs || "Clean"}. \nGenerate the NEXT task using 'WORKER_PROMPT:' format.`;
+            tokenUsage += estimateTokens(summary);
+
+            const pmResponse = await AIEngine.callPM(project.pmHistory, summary);
+            tokenUsage += estimateTokens(pmResponse);
+            
+            project.pmHistory.push({ role: 'user', parts: [{ text: summary }] });
+            project.pmHistory.push({ role: 'model', parts: [{ text: pmResponse }] });
+
+            const nextInstruction = extractWorkerPrompt(pmResponse);
+            if (!nextInstruction) {
+                 await StateManager.updateProject(projectId, { pmHistory: project.pmHistory, status: "IDLE" });
+                 if(db) {
+                     await db.ref(`projects/${projectId}/info`).update({
+                         status: "IDLE",
+                         lastUpdated: Date.now()
+                     });
+                 }
+                 await deductUserCredits(userId, tokenUsage);
+                 return res.json({ success: true, message: "No further tasks. Project Idle." });
+            }
 
-        const newWorkerHistory = [];
-        const newPrompt = `New Objective: ${nextInstruction}`;
-        const workerResponse = await AIEngine.callWorker(newWorkerHistory, newPrompt, []);
+            const newWorkerHistory = [];
+            const newPrompt = `New Objective: ${nextInstruction}`;
+            tokenUsage += estimateTokens(newPrompt);
 
-        newWorkerHistory.push({ role: 'user', parts: [{ text: newPrompt }] });
-        newWorkerHistory.push({ role: 'model', parts: [{ text: workerResponse }] });
+            const workerResponse = await AIEngine.callWorker(newWorkerHistory, newPrompt, []);
+            tokenUsage += estimateTokens(workerResponse);
 
-        await StateManager.updateProject(projectId, {
-            pmHistory: project.pmHistory,
-            workerHistory: newWorkerHistory,
-            failureCount: 0
-        });
+            newWorkerHistory.push({ role: 'user', parts: [{ text: newPrompt }] });
+            newWorkerHistory.push({ role: 'model', parts: [{ text: workerResponse }] });
 
-        StateManager.queueCommand(projectId, { type: "EXECUTE", payload: "warn('Starting Next Task...')" });
-        await processAndQueueResponse(projectId, workerResponse);
-        
-        // Still working on next task
-        if(db) {
-            await db.ref(`projects/${projectId}/info`).update({
-                status: "working",
-                lastUpdated: Date.now()
+            await StateManager.updateProject(projectId, {
+                pmHistory: project.pmHistory,
+                workerHistory: newWorkerHistory,
+                failureCount: 0
             });
-        }
-        return res.json({ success: true, message: "Next Task Assigned" });
-    }
 
-    let isFailure = false;
-    if (logContext?.logs) {
-        const errs = ["Error", "Exception", "failed", "Stack Begin", "Infinite yield"];
-        if (errs.some(k => logContext.logs.includes(k))) {
-            isFailure = true;
-            project.failureCount = (project.failureCount || 0) + 1;
-        }
-    }
+            StateManager.queueCommand(projectId, { type: "EXECUTE", payload: "warn('Starting Next Task...')" });
+            await processAndQueueResponse(projectId, workerResponse, userId); // Handle usage inside or separately? We'll track image usage here if it occurs.
+            
+            // Still working on next task
+            if(db) {
+                await db.ref(`projects/${projectId}/info`).update({
+                    status: "working",
+                    lastUpdated: Date.now()
+                });
+            }
 
-    if (project.failureCount > 3) {
-        const pmPrompt = sysPrompts.pm_guidance_prompt.replace('{{LOGS}}', logContext?.logs);
-        const pmVerdict = await AIEngine.callPM(project.pmHistory, pmPrompt);
+            await deductUserCredits(userId, tokenUsage);
+            return res.json({ success: true, message: "Next Task Assigned" });
+        }
 
-        if (pmVerdict.includes("[TERMINATE]")) {
-            const fixInstruction = pmVerdict.replace("[TERMINATE]", "").trim();
-            const resetHistory = [];
-            const resetPrompt = `[SYSTEM]: Previous worker terminated. \nNew Objective: ${fixInstruction}`;
-            const workerResp = await AIEngine.callWorker(resetHistory, resetPrompt, []);
-            
-            resetHistory.push({ role: 'user', parts: [{ text: resetPrompt }] });
-            resetHistory.push({ role: 'model', parts: [{ text: workerResp }] });
+        let isFailure = false;
+        if (logContext?.logs) {
+            const errs = ["Error", "Exception", "failed", "Stack Begin", "Infinite yield"];
+            if (errs.some(k => logContext.logs.includes(k))) {
+                isFailure = true;
+                project.failureCount = (project.failureCount || 0) + 1;
+            }
+        }
 
-            await StateManager.updateProject(projectId, { workerHistory: resetHistory, failureCount: 0 });
-            StateManager.queueCommand(projectId, { type: "EXECUTE", payload: "print('SYSTEM: Worker Reset')" });
-            await processAndQueueResponse(projectId, workerResp);
-            return res.json({ success: true, message: "Worker Terminated." });
-        } else {
-            const injection = `[PM GUIDANCE]: ${pmVerdict} \n\nApply this fix now.`;
-            const workerResp = await AIEngine.callWorker(project.workerHistory, injection, []);
-            
-            project.workerHistory.push({ role: 'user', parts: [{ text: injection }] });
-            project.workerHistory.push({ role: 'model', parts: [{ text: workerResp }] });
-            
-            await StateManager.updateProject(projectId, { workerHistory: project.workerHistory, failureCount: 1 });
-            await processAndQueueResponse(projectId, workerResp);
-            return res.json({ success: true, message: "PM Guidance Applied." });
+        if (project.failureCount > 3) {
+            const pmPrompt = sysPrompts.pm_guidance_prompt.replace('{{LOGS}}', logContext?.logs);
+            tokenUsage += estimateTokens(pmPrompt);
+
+            const pmVerdict = await AIEngine.callPM(project.pmHistory, pmPrompt);
+            tokenUsage += estimateTokens(pmVerdict);
+
+            if (pmVerdict.includes("[TERMINATE]")) {
+                const fixInstruction = pmVerdict.replace("[TERMINATE]", "").trim();
+                const resetHistory = [];
+                const resetPrompt = `[SYSTEM]: Previous worker terminated. \nNew Objective: ${fixInstruction}`;
+                tokenUsage += estimateTokens(resetPrompt);
+
+                const workerResp = await AIEngine.callWorker(resetHistory, resetPrompt, []);
+                tokenUsage += estimateTokens(workerResp);
+                
+                resetHistory.push({ role: 'user', parts: [{ text: resetPrompt }] });
+                resetHistory.push({ role: 'model', parts: [{ text: workerResp }] });
+
+                await StateManager.updateProject(projectId, { workerHistory: resetHistory, failureCount: 0 });
+                StateManager.queueCommand(projectId, { type: "EXECUTE", payload: "print('SYSTEM: Worker Reset')" });
+                await processAndQueueResponse(projectId, workerResp, userId);
+                await deductUserCredits(userId, tokenUsage);
+                return res.json({ success: true, message: "Worker Terminated." });
+            } else {
+                const injection = `[PM GUIDANCE]: ${pmVerdict} \n\nApply this fix now.`;
+                tokenUsage += estimateTokens(injection);
+
+                const workerResp = await AIEngine.callWorker(project.workerHistory, injection, []);
+                tokenUsage += estimateTokens(workerResp);
+                
+                project.workerHistory.push({ role: 'user', parts: [{ text: injection }] });
+                project.workerHistory.push({ role: 'model', parts: [{ text: workerResp }] });
+                
+                await StateManager.updateProject(projectId, { workerHistory: project.workerHistory, failureCount: 1 });
+                await processAndQueueResponse(projectId, workerResp, userId);
+                await deductUserCredits(userId, tokenUsage);
+                return res.json({ success: true, message: "PM Guidance Applied." });
+            }
         }
-    }
 
-    try {
+        // Standard Logic
         const fullInput = `USER: ${prompt || "Automatic Feedback"}` + formatContext({ hierarchyContext, scriptContext, logContext });
+        tokenUsage += estimateTokens(fullInput);
+
         let response = await AIEngine.callWorker(project.workerHistory, fullInput, images || []);
+        tokenUsage += estimateTokens(response);
         
         const pmQuestion = extractPMQuestion(response);
         if (pmQuestion) {
             const pmConsultPrompt = `[WORKER CONSULTATION]: The Worker asks: "${pmQuestion}"\nProvide a technical answer to unblock them.`;
+            tokenUsage += estimateTokens(pmConsultPrompt);
+
             const pmAnswer = await AIEngine.callPM(project.pmHistory, pmConsultPrompt);
+            tokenUsage += estimateTokens(pmAnswer);
 
             project.pmHistory.push({ role: 'user', parts: [{ text: pmConsultPrompt }] });
             project.pmHistory.push({ role: 'model', parts: [{ text: pmAnswer }] });
@@ -429,7 +586,10 @@ app.post('/project/feedback', async (req, res) => {
             project.workerHistory.push({ role: 'user', parts: [{ text: fullInput }] });
             project.workerHistory.push({ role: 'model', parts: [{ text: response }] });
             
+            tokenUsage += estimateTokens(injectionMsg);
             response = await AIEngine.callWorker(project.workerHistory, injectionMsg, []);
+            tokenUsage += estimateTokens(response);
+
             project.workerHistory.push({ role: 'user', parts: [{ text: injectionMsg }] });
             project.workerHistory.push({ role: 'model', parts: [{ text: response }] });
         } else {
@@ -443,7 +603,7 @@ app.post('/project/feedback', async (req, res) => {
             failureCount: project.failureCount 
         });
         
-        await processAndQueueResponse(projectId, response);
+        await processAndQueueResponse(projectId, response, userId);
         
         // --- UPDATED: Set status to IDLE and update lastUpdated ---
         if(db) {
@@ -453,9 +613,18 @@ app.post('/project/feedback', async (req, res) => {
             });
         }
       
+        // Final Deduction
+        await deductUserCredits(userId, tokenUsage);
+
         res.json({ success: true });
     } catch (err) {
         console.error("AI Error:", err);
+        // Even if it failed, deduct used tokens
+       // if (tokenUsage > 0) await deductUserCredits(userId, tokenUsage);
+
+        if (err.message.includes("Insufficient credits")) {
+            return res.status(402).json({ error: err.message });
+        }
         res.status(500).json({ error: "AI Failed" });
     }
 });
@@ -496,26 +665,49 @@ app.post('/project/ping', async (req, res) => {
 });
 
 app.post('/human/override', validateRequest, async (req, res) => {
-    const { projectId, instruction, pruneHistory } = req.body;
+    const { projectId, instruction, pruneHistory, userId } = req.body;
+    // Assuming userId is passed, we check credits
+    if (userId) {
+        try {
+            await checkUserCredits(userId);
+        } catch (e) {
+            return res.status(402).json({ error: e.message });
+        }
+    }
+
     const project = await StateManager.getProject(projectId);
     const overrideMsg = `[SYSTEM OVERRIDE]: ${instruction}`;
+    let tokenUsage = estimateTokens(overrideMsg);
+
     if (pruneHistory && project.workerHistory.length >= 2) {
         project.workerHistory.pop();
         project.workerHistory.pop();
     }
     const response = await AIEngine.callWorker(project.workerHistory, overrideMsg, []);
+    tokenUsage += estimateTokens(response);
+
     project.workerHistory.push({ role: 'user', parts: [{ text: overrideMsg }] });
     project.workerHistory.push({ role: 'model', parts: [{ text: response }] });
     await StateManager.updateProject(projectId, { workerHistory: project.workerHistory });
-    await processAndQueueResponse(projectId, response);
+    await processAndQueueResponse(projectId, response, userId);
+    
+    if (userId) await deductUserCredits(userId, tokenUsage);
+
     res.json({ success: true });
 });
 
-async function processAndQueueResponse(projectId, rawResponse) {
+// Added userId param to handle asset cost if image generated
+async function processAndQueueResponse(projectId, rawResponse, userId) {
     const imgPrompt = extractImagePrompt(rawResponse);
     if (imgPrompt) {
         console.log(`[${projectId}] 🎨 Generating Asset: ${imgPrompt}`);
         const base64Image = await AIEngine.generateImage(imgPrompt);
+        
+        // Deduct for image generation here if userId is available
+        if (base64Image && userId) {
+            await deductUserCredits(userId, IMAGE_GENERATION_COST);
+        }
+
         if (base64Image) {
             await StateManager.queueCommand(projectId, { type: "CREATE_ASSET", payload: base64Image });
         }