Update app.js

app.js

@@ -16,6 +16,7 @@ try {
     if (process.env.FIREBASE_SERVICE_ACCOUNT_JSON) {
         const serviceAccount = JSON.parse(process.env.FIREBASE_SERVICE_ACCOUNT_JSON);
         
+        // Specific bucket as requested
         const bucketName = `shago-web.firebasestorage.app`; 
 
         if (admin.apps.length === 0) {
@@ -40,47 +41,64 @@ try {
 
 const app = express();
 const PORT = process.env.PORT || 7860;
-const sysPrompts = JSON.parse(fs.readFileSync('./prompts.json', 'utf8'));
 
+// Load prompts safely
+let sysPrompts = {};
+try {
+    sysPrompts = JSON.parse(fs.readFileSync('./prompts.json', 'utf8'));
+} catch (e) {
+    console.error("Failed to load prompts.json:", e);
+}
 
 app.use(cors());
 app.use(bodyParser.json({ limit: '50mb' }));
 
-// --- CREDIT HELPERS ---
-const MIN_CREDITS_REQUIRED = 20;
+// --- CREDIT CONFIGURATION ---
+
+// Minimums required to start a request
+const MIN_BASIC_REQUIRED = 50; 
+const MIN_DIAMOND_REQUIRED = 50; 
+
+// Fixed cost for images (Basic credits)
+const IMAGE_COST_BASIC = 1000; 
 
 /**
- * Checks if user has enough credits to START an operation.
- * Throws if < 20.
+ * Checks if user has enough credits of a specific type.
+ * @param {string} userId 
+ * @param {'basic'|'diamond'} type 
  */
-async function checkMinimumCredits(userId) {
+async function checkMinimumCredits(userId, type = 'basic') {
     if (!db) return; 
-    const snap = await db.ref(`users/${userId}/credits`).once('value');
+
+    // Path: users/{uid}/credits/basic OR users/{uid}/credits/diamond
+    const snap = await db.ref(`users/${userId}/credits/${type}`).once('value');
     const credits = snap.val() || 0;
 
-    if (credits < MIN_CREDITS_REQUIRED) {
-        throw new Error(`Insufficient credits. You have ${credits}, need minimum ${MIN_CREDITS_REQUIRED} to proceed.`);
+    const required = type === 'diamond' ? MIN_DIAMOND_REQUIRED : MIN_BASIC_REQUIRED;
+
+    if (credits < required) {
+        throw new Error(`Insufficient ${type} credits. You have ${credits}, need minimum ${required} to proceed.`);
     }
 }
 
 /**
- * Deducts exact tokens from the user's account.
+ * Deducts exact tokens from the user's specific wallet.
+ * @param {string} userId 
+ * @param {number} amount 
+ * @param {'basic'|'diamond'} type 
  */
-async function deductUserCredits(userId, amount) {
+async function deductUserCredits(userId, amount, type = 'basic') {
     if (!db || !amount || amount <= 0) return;
     
     try {
-        const ref = db.ref(`users/${userId}/credits`);
+        const ref = db.ref(`users/${userId}/credits/${type}`);
         await ref.transaction((current_credits) => {
             const current = current_credits || 0;
-            // Allow going negative if the bill is huge, or cap at 0? 
-            // Usually valid to go slightly negative on final bill, or cap at 0.
-            // Here we assume we just subtract.
             return Math.max(0, current - amount);
         });
-        console.log(`[Credits] Deducted ${amount} tokens from User ${userId}`);
+        console.log(`[Credits] Deducted ${amount} ${type} credits from User ${userId}`);
     } catch (err) {
-        console.error(`[Credits] Failed to deduct ${amount} from ${userId}:`, err);
+        console.error(`[Credits] Failed to deduct ${amount} ${type} from ${userId}:`, err);
     }
 }
 
@@ -139,19 +157,17 @@ app.post('/onboarding/analyze', validateRequest, async (req, res) => {
     if (!description) return res.status(400).json({ error: "Description required" });
 
     try {
-        // 1. Pre-flight Check
-        await checkMinimumCredits(userId);
+        // Analyst uses BASIC models (Flash)
+        await checkMinimumCredits(userId, 'basic');
 
         console.log(`[Onboarding] Analyzing idea...`);
         
-        // 2. AI Call
-        // Result is { questions: [...], usage: { totalTokenCount: ... }, ... }
         const result = await AIEngine.generateEntryQuestions(description);
         
         const usage = result.usage?.totalTokenCount || 0;
 
-        // 3. Billing (Only on success)
-        if (usage > 0) await deductUserCredits(userId, usage);
+        // Bill Basic
+        if (usage > 0) await deductUserCredits(userId, usage, 'basic');
 
         if (result.status === "REJECTED") {
             return res.json({ 
@@ -163,7 +179,7 @@ app.post('/onboarding/analyze', validateRequest, async (req, res) => {
 
     } catch (err) {
         console.error(err);
-        if (err.message.includes("Insufficient credits")) {
+        if (err.message.includes("Insufficient")) {
             return res.status(402).json({ error: err.message });
         }
         res.status(500).json({ error: "Analysis failed" });
@@ -172,20 +188,20 @@ app.post('/onboarding/analyze', validateRequest, async (req, res) => {
 
 app.post('/onboarding/create', validateRequest, async (req, res) => {
     const { userId, description, answers } = req.body;
-    let totalTokens = 0;
+    let basicTokens = 0; // Create flow uses Grader (Basic) + Image (Basic)
 
     try {
-        // 1. Pre-flight Check
-        await checkMinimumCredits(userId);
+        // Pre-flight check
+        await checkMinimumCredits(userId, 'basic');
 
         const randomHex = (n = 6) => crypto.randomBytes(Math.ceil(n/2)).toString("hex").slice(0, n);
         const projectId = `proj_${Date.now()}_${randomHex(7)}`;
 
         console.log(`[Onboarding] Grading Project ${projectId}...`);
 
-        // 2. Grading
+        // 1. Grading (Basic)
         const gradingResult = await AIEngine.gradeProject(description, answers);
-        totalTokens += (gradingResult.usage?.totalTokenCount || 0);
+        basicTokens += (gradingResult.usage?.totalTokenCount || 0);
 
         const isFailure = gradingResult.feasibility < 30 || gradingResult.rating === 'F';
 
@@ -195,11 +211,13 @@ app.post('/onboarding/create', validateRequest, async (req, res) => {
             console.log(`[Onboarding] ✅ Passed. Generating Thumbnail...`);
             const imagePrompt = `Game Title: ${gradingResult.title}. Core Concept: ${description}`;
             
-            // 3. Image Generation
+            // 2. Image Generation (Basic)
             const imgResult = await AIEngine.generateImage(imagePrompt);
             if (imgResult) {
                 thumbnailBase64 = imgResult.image;
-                totalTokens += (imgResult.usage?.totalTokenCount || 0);
+                // Add token usage for text + fixed cost
+                basicTokens += (imgResult.usage?.totalTokenCount || 0); 
+                if (imgResult.image) basicTokens += IMAGE_COST_BASIC;
             }
         }
 
@@ -222,7 +240,7 @@ app.post('/onboarding/create', validateRequest, async (req, res) => {
         const timestamp = Date.now();
         const status = isFailure ? "rejected" : "Idle";
 
-        // Save Data
+        // Save Data Logic
         if (!isFailure) {
             const memoryObject = {
                 id: projectId,
@@ -242,7 +260,6 @@ app.post('/onboarding/create', validateRequest, async (req, res) => {
             await StateManager.updateProject(projectId, memoryObject);
         }
 
-        // DB Updates
         if (firestore && !isFailure) {
             await firestore.collection('projects').doc(projectId).set({
                 id: projectId,
@@ -274,8 +291,8 @@ app.post('/onboarding/create', validateRequest, async (req, res) => {
             await db.ref().update(updates);
         }
 
-        // 4. Billing (Only on success)
-        if (totalTokens > 0) await deductUserCredits(userId, totalTokens);
+        // Deduct Basic Credits
+        if (basicTokens > 0) await deductUserCredits(userId, basicTokens, 'basic');
 
         res.json({ 
             success: !isFailure, 
@@ -287,7 +304,7 @@ app.post('/onboarding/create', validateRequest, async (req, res) => {
 
     } catch (err) {
         console.error("Create Error:", err);
-        if (err.message.includes("Insufficient credits")) {
+        if (err.message.includes("Insufficient")) {
             return res.status(402).json({ error: err.message });
         }
         res.status(500).json({ error: "Creation failed" });
@@ -298,39 +315,42 @@ app.post('/onboarding/create', validateRequest, async (req, res) => {
 
 async function runBackgroundInitialization(projectId, userId, description) {
     console.log(`[Background] Starting initialization for ${projectId}`);
-    let totalTokens = 0;
+    
+    // Separate Counters
+    let diamondUsage = 0; // For PM
+    let basicUsage = 0;   // For Worker
 
     try {
         const pmHistory = [];
         
-        // 1. Generate GDD
+        // 1. Generate GDD (PM -> Diamond)
         const gddPrompt = `Create a comprehensive GDD for: ${description}`;
         const gddResult = await AIEngine.callPM(pmHistory, gddPrompt);
         
-        totalTokens += (gddResult.usage?.totalTokenCount || 0);
+        diamondUsage += (gddResult.usage?.totalTokenCount || 0);
         const gddText = gddResult.text;
 
         pmHistory.push({ role: 'user', parts: [{ text: gddPrompt }] });
         pmHistory.push({ role: 'model', parts: [{ text: gddText }] });
 
-        // 2. Generate First Task
+        // 2. Generate First Task (PM -> Diamond)
         const taskPrompt = "Based on the GDD, generate the first technical milestone.\nOutput format:\nTASK_NAME: <Name>\nWORKER_PROMPT: <Specific, isolated instructions for the worker>";
         const taskResult = await AIEngine.callPM(pmHistory, taskPrompt);
         
-        totalTokens += (taskResult.usage?.totalTokenCount || 0);
+        diamondUsage += (taskResult.usage?.totalTokenCount || 0);
         const taskText = taskResult.text;
 
         pmHistory.push({ role: 'user', parts: [{ text: taskPrompt }] });
         pmHistory.push({ role: 'model', parts: [{ text: taskText }] });
 
-        // 3. Initialize Worker
+        // 3. Initialize Worker (Worker -> Basic)
         const initialWorkerInstruction = extractWorkerPrompt(taskText) || `Initialize structure for: ${description}`;
         const workerHistory = [];
         const initialWorkerPrompt = `CONTEXT: New Project. \nINSTRUCTION: ${initialWorkerInstruction}`;
         
         const workerResult = await AIEngine.callWorker(workerHistory, initialWorkerPrompt, []);
         
-        totalTokens += (workerResult.usage?.totalTokenCount || 0);
+        basicUsage += (workerResult.usage?.totalTokenCount || 0);
         const workerText = workerResult.text;
 
         workerHistory.push({ role: 'user', parts: [{ text: initialWorkerPrompt }] });
@@ -345,8 +365,8 @@ async function runBackgroundInitialization(projectId, userId, description) {
             failureCount: 0
         });
 
-        // 5. Queue commands (Handle assets + their costs)
-        // We pass userId to deduct image costs immediately inside this function
+        // 5. Queue commands (Handle assets + basic costs)
+        // We pass userId to deduct image costs (Basic) immediately inside this function
         await processAndQueueResponse(projectId, workerText, userId); 
         
         // 6. Update Status
@@ -355,15 +375,15 @@ async function runBackgroundInitialization(projectId, userId, description) {
             lastUpdated: Date.now()
         });
         
-        // 7. Deduct accumulated text tokens
-        if (totalTokens > 0) await deductUserCredits(userId, totalTokens);
+        // 7. Deduct Accumulated Credits
+        if (diamondUsage > 0) await deductUserCredits(userId, diamondUsage, 'diamond');
+        if (basicUsage > 0) await deductUserCredits(userId, basicUsage, 'basic');
         
-        console.log(`[Background] Init complete. Text Tokens: ${totalTokens}`);
+        console.log(`[Background] Init complete. Diamond: ${diamondUsage}, Basic: ${basicUsage}`);
 
     } catch (err) {
         console.error(`[Background] Init Error for ${projectId}:`, err);
         if(db) await db.ref(`projects/${projectId}/info/status`).set("error");
-        // No deduction on failure
     }
 }
 
@@ -371,7 +391,9 @@ app.post('/new/project', validateRequest, async (req, res) => {
     const { userId, projectId, description } = req.body;
     
     try {
-        await checkMinimumCredits(userId);
+        // Init requires BOTH types to be safe
+        await checkMinimumCredits(userId, 'diamond');
+        await checkMinimumCredits(userId, 'basic');
 
         if(db) db.ref(`projects/${projectId}/info/status`).set("initializing");
         
@@ -383,7 +405,7 @@ app.post('/new/project', validateRequest, async (req, res) => {
         runBackgroundInitialization(projectId, userId, description);
 
     } catch (err) {
-        if (err.message.includes("Insufficient credits")) {
+        if (err.message.includes("Insufficient")) {
             return res.status(402).json({ error: err.message });
         }
         res.status(500).json({ error: "Failed to start project" });
@@ -392,15 +414,17 @@ app.post('/new/project', validateRequest, async (req, res) => {
 
 app.post('/project/feedback', async (req, res) => {
     const { userId, projectId, prompt, hierarchyContext, scriptContext, logContext, taskComplete, images } = req.body;
-    let totalTokens = 0;
+    
+    let diamondUsage = 0;
+    let basicUsage = 0;
     
     try {
         const project = await StateManager.getProject(projectId);
         if (!project) return res.status(404).json({ error: "Project not found." });
         if (project.userId !== userId) return res.status(403).json({ error: "Unauthorized" });
 
-        // 1. Pre-flight Check
-        await checkMinimumCredits(userId);
+        // Basic Check (most interaction is worker)
+        await checkMinimumCredits(userId, 'basic');
         
         if(db) await db.ref(`projects/${projectId}/info/status`).set("working");
           
@@ -408,8 +432,9 @@ app.post('/project/feedback', async (req, res) => {
             console.log(`[${projectId}] ✅ TASK COMPLETE.`);
             const summary = `Worker completed the previous task. Logs: ${logContext?.logs || "Clean"}. \nGenerate the NEXT task using 'WORKER_PROMPT:' format.`;
             
+            // PM Call -> Diamond
             const pmResult = await AIEngine.callPM(project.pmHistory, summary);
-            totalTokens += (pmResult.usage?.totalTokenCount || 0);
+            diamondUsage += (pmResult.usage?.totalTokenCount || 0);
             const pmText = pmResult.text;
             
             project.pmHistory.push({ role: 'user', parts: [{ text: summary }] });
@@ -420,8 +445,8 @@ app.post('/project/feedback', async (req, res) => {
                  await StateManager.updateProject(projectId, { pmHistory: project.pmHistory, status: "IDLE" });
                  if(db) await db.ref(`projects/${projectId}/info`).update({ status: "IDLE", lastUpdated: Date.now() });
                  
-                 // Deduct what we used so far
-                 if (totalTokens > 0) await deductUserCredits(userId, totalTokens);
+                 // Deduct what we used
+                 if (diamondUsage > 0) await deductUserCredits(userId, diamondUsage, 'diamond');
                  
                  return res.json({ success: true, message: "No further tasks. Project Idle." });
             }
@@ -429,8 +454,9 @@ app.post('/project/feedback', async (req, res) => {
             const newWorkerHistory = [];
             const newPrompt = `New Objective: ${nextInstruction}`;
             
+            // Worker Call -> Basic
             const workerResult = await AIEngine.callWorker(newWorkerHistory, newPrompt, []);
-            totalTokens += (workerResult.usage?.totalTokenCount || 0);
+            basicUsage += (workerResult.usage?.totalTokenCount || 0);
             const workerText = workerResult.text;
 
             newWorkerHistory.push({ role: 'user', parts: [{ text: newPrompt }] });
@@ -447,25 +473,20 @@ app.post('/project/feedback', async (req, res) => {
             
             if(db) await db.ref(`projects/${projectId}/info`).update({ status: "working", lastUpdated: Date.now() });
 
-            // Billing
-            if (totalTokens > 0) await deductUserCredits(userId, totalTokens);
-            return res.json({ success: true, message: "Next Task Assigned" });
-        }
+            // Final Bill
+            if (diamondUsage > 0) await deductUserCredits(userId, diamondUsage, 'diamond');
+            if (basicUsage > 0) await deductUserCredits(userId, basicUsage, 'basic');
 
-        // Logic for Failure Handling (PM Guidance) or Standard Feedback
-        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;
-            }
+            return res.json({ success: true, message: "Next Task Assigned" });
         }
 
+        // Logic for Failure Handling (PM Guidance)
         if (project.failureCount > 3) {
-            const pmPrompt = sysPrompts.pm_guidance_prompt.replace('{{LOGS}}', logContext?.logs);
+            const pmPrompt = (sysPrompts.pm_guidance_prompt || "Analyze logs: {{LOGS}}").replace('{{LOGS}}', logContext?.logs);
+            
+            // PM -> Diamond
             const pmResult = await AIEngine.callPM(project.pmHistory, pmPrompt);
-            totalTokens += (pmResult.usage?.totalTokenCount || 0);
+            diamondUsage += (pmResult.usage?.totalTokenCount || 0);
             const pmVerdict = pmResult.text;
 
             if (pmVerdict.includes("[TERMINATE]")) {
@@ -473,8 +494,9 @@ app.post('/project/feedback', async (req, res) => {
                 const resetHistory = [];
                 const resetPrompt = `[SYSTEM]: Previous worker terminated. \nNew Objective: ${fixInstruction}`;
                 
+                // Worker -> Basic
                 const workerResult = await AIEngine.callWorker(resetHistory, resetPrompt, []);
-                totalTokens += (workerResult.usage?.totalTokenCount || 0);
+                basicUsage += (workerResult.usage?.totalTokenCount || 0);
                 
                 resetHistory.push({ role: 'user', parts: [{ text: resetPrompt }] });
                 resetHistory.push({ role: 'model', parts: [{ text: workerResult.text }] });
@@ -483,12 +505,15 @@ app.post('/project/feedback', async (req, res) => {
                 StateManager.queueCommand(projectId, { type: "EXECUTE", payload: "print('SYSTEM: Worker Reset')" });
                 await processAndQueueResponse(projectId, workerResult.text, userId);
                 
-                if (totalTokens > 0) await deductUserCredits(userId, totalTokens);
+                if (diamondUsage > 0) await deductUserCredits(userId, diamondUsage, 'diamond');
+                if (basicUsage > 0) await deductUserCredits(userId, basicUsage, 'basic');
                 return res.json({ success: true, message: "Worker Terminated." });
             } else {
                 const injection = `[PM GUIDANCE]: ${pmVerdict} \n\nApply this fix now.`;
+                
+                // Worker -> Basic
                 const workerResult = await AIEngine.callWorker(project.workerHistory, injection, []);
-                totalTokens += (workerResult.usage?.totalTokenCount || 0);
+                basicUsage += (workerResult.usage?.totalTokenCount || 0);
 
                 project.workerHistory.push({ role: 'user', parts: [{ text: injection }] });
                 project.workerHistory.push({ role: 'model', parts: [{ text: workerResult.text }] });
@@ -496,23 +521,26 @@ app.post('/project/feedback', async (req, res) => {
                 await StateManager.updateProject(projectId, { workerHistory: project.workerHistory, failureCount: 1 });
                 await processAndQueueResponse(projectId, workerResult.text, userId);
                 
-                if (totalTokens > 0) await deductUserCredits(userId, totalTokens);
+                if (diamondUsage > 0) await deductUserCredits(userId, diamondUsage, 'diamond');
+                if (basicUsage > 0) await deductUserCredits(userId, basicUsage, 'basic');
                 return res.json({ success: true, message: "PM Guidance Applied." });
             }
         }
 
-        // Standard Interaction
+        // Standard Interaction (Worker Only = Basic)
         const fullInput = `USER: ${prompt || "Automatic Feedback"}` + formatContext({ hierarchyContext, scriptContext, logContext });
         
         let workerResult = await AIEngine.callWorker(project.workerHistory, fullInput, images || []);
-        totalTokens += (workerResult.usage?.totalTokenCount || 0);
+        basicUsage += (workerResult.usage?.totalTokenCount || 0);
         let responseText = workerResult.text;
         
+        // PM Consultation Check
         const pmQuestion = extractPMQuestion(responseText);
         if (pmQuestion) {
+            // PM -> Diamond
             const pmConsultPrompt = `[WORKER CONSULTATION]: The Worker asks: "${pmQuestion}"\nProvide a technical answer to unblock them.`;
             const pmResult = await AIEngine.callPM(project.pmHistory, pmConsultPrompt);
-            totalTokens += (pmResult.usage?.totalTokenCount || 0);
+            diamondUsage += (pmResult.usage?.totalTokenCount || 0);
             const pmAnswer = pmResult.text;
 
             project.pmHistory.push({ role: 'user', parts: [{ text: pmConsultPrompt }] });
@@ -522,9 +550,9 @@ app.post('/project/feedback', async (req, res) => {
             project.workerHistory.push({ role: 'user', parts: [{ text: fullInput }] });
             project.workerHistory.push({ role: 'model', parts: [{ text: responseText }] });
             
-            // Re-call worker with answer
+            // Re-call Worker -> Basic
             workerResult = await AIEngine.callWorker(project.workerHistory, injectionMsg, []);
-            totalTokens += (workerResult.usage?.totalTokenCount || 0);
+            basicUsage += (workerResult.usage?.totalTokenCount || 0);
             responseText = workerResult.text;
 
             project.workerHistory.push({ role: 'user', parts: [{ text: injectionMsg }] });
@@ -545,13 +573,14 @@ app.post('/project/feedback', async (req, res) => {
         if(db) await db.ref(`projects/${projectId}/info`).update({ status: "idle", lastUpdated: Date.now() });
       
         // Final Billing
-        if (totalTokens > 0) await deductUserCredits(userId, totalTokens);
+        if (diamondUsage > 0) await deductUserCredits(userId, diamondUsage, 'diamond');
+        if (basicUsage > 0) await deductUserCredits(userId, basicUsage, 'basic');
 
         res.json({ success: true });
 
     } catch (err) {
         console.error("AI Error:", err);
-        if (err.message.includes("Insufficient credits")) {
+        if (err.message.includes("Insufficient")) {
             return res.status(402).json({ error: err.message });
         }
         res.status(500).json({ error: "AI Failed" });
@@ -585,10 +614,10 @@ app.post('/project/ping', async (req, res) => {
 
 app.post('/human/override', validateRequest, async (req, res) => {
     const { projectId, instruction, pruneHistory, userId } = req.body;
-    let totalTokens = 0;
+    let basicUsage = 0;
 
     try {
-        await checkMinimumCredits(userId);
+        await checkMinimumCredits(userId, 'basic');
 
         const project = await StateManager.getProject(projectId);
         const overrideMsg = `[SYSTEM OVERRIDE]: ${instruction}`;
@@ -598,8 +627,9 @@ app.post('/human/override', validateRequest, async (req, res) => {
             project.workerHistory.pop();
         }
 
+        // Worker -> Basic
         const workerResult = await AIEngine.callWorker(project.workerHistory, overrideMsg, []);
-        totalTokens += (workerResult.usage?.totalTokenCount || 0);
+        basicUsage += (workerResult.usage?.totalTokenCount || 0);
 
         project.workerHistory.push({ role: 'user', parts: [{ text: overrideMsg }] });
         project.workerHistory.push({ role: 'model', parts: [{ text: workerResult.text }] });
@@ -607,18 +637,18 @@ app.post('/human/override', validateRequest, async (req, res) => {
         await StateManager.updateProject(projectId, { workerHistory: project.workerHistory });
         await processAndQueueResponse(projectId, workerResult.text, userId);
         
-        if (totalTokens > 0) await deductUserCredits(userId, totalTokens);
+        if (basicUsage > 0) await deductUserCredits(userId, basicUsage, 'basic');
         
         res.json({ success: true });
     } catch (err) {
-        if (err.message.includes("Insufficient credits")) {
+        if (err.message.includes("Insufficient")) {
             return res.status(402).json({ error: err.message });
         }
         res.status(500).json({ error: "Override Failed" });
     }
 });
 
-// Helper to handle Asset Parsing & Billing
+// Helper to handle Asset Parsing & Billing (Basic Credits)
 async function processAndQueueResponse(projectId, rawResponse, userId) {
     const imgPrompt = extractImagePrompt(rawResponse);
     if (imgPrompt) {
@@ -628,10 +658,12 @@ async function processAndQueueResponse(projectId, rawResponse, userId) {
         const imgResult = await AIEngine.generateImage(imgPrompt);
         
         if (imgResult && imgResult.image) {
-            // 2. Bill for image tokens immediately
+            // 2. Bill for image tokens immediately (Basic)
             const imgTokens = imgResult.usage?.totalTokenCount || 0;
-            if (userId && imgTokens > 0) {
-                await deductUserCredits(userId, imgTokens);
+            const totalCost = imgTokens; // imgTokens + IMAGE_COST_BASIC;
+            
+            if (userId && totalCost > 0) {
+                await deductUserCredits(userId, totalCost, 'basic');
             }
 
             // 3. Queue creation command