Update app.js

app.js

@@ -1,7 +1,7 @@
 import express from 'express';
 import bodyParser from 'body-parser';
 import cors from 'cors';
-import { StateManager, initDB } from './stateManager.js'; // Import initDB
+import { StateManager, initDB } from './stateManager.js';
 import { AIEngine } from './aiEngine.js';
 import fs from 'fs';
 import admin from 'firebase-admin';
@@ -16,14 +16,13 @@ try {
     if (process.env.FIREBASE_SERVICE_ACCOUNT_JSON) {
         const serviceAccount = JSON.parse(process.env.FIREBASE_SERVICE_ACCOUNT_JSON);
         
-        // We assume the project ID is available in the service account to construct the bucket name
         const bucketName = `shago-web.firebasestorage.app`; 
 
         if (admin.apps.length === 0) {
             admin.initializeApp({
                 credential: admin.credential.cert(serviceAccount),
                 databaseURL: "https://shago-web-default-rtdb.firebaseio.com", 
-                storageBucket: bucketName // Required for Storage
+                storageBucket: bucketName
             });
         }
         
@@ -31,7 +30,6 @@ try {
         firestore = admin.firestore();
         storage = admin.storage();
         
-        // Pass DB instance to StateManager
         initDB(db);
         console.log("🔥 Firebase Connected (RTDB, Firestore, Storage) & StateManager Linked");
     } else {
@@ -48,47 +46,39 @@ const sysPrompts = JSON.parse(fs.readFileSync('./prompts.json', 'utf8'));
 app.use(cors());
 app.use(bodyParser.json({ limit: '50mb' }));
 
-// --- CREDIT SYSTEM HELPERS ---
-
+// --- CREDIT 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)
+ * Checks if user has enough credits to START an operation.
+ * Throws if < 20.
  */
-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)
-    
+async function checkMinimumCredits(userId) {
+    if (!db) return; 
     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.`);
+        throw new Error(`Insufficient credits. You have ${credits}, need minimum ${MIN_CREDITS_REQUIRED} to proceed.`);
     }
-    return credits;
 }
 
 /**
- * Deducts tokens from the user's account.
+ * Deducts exact tokens from the user's account.
  */
 async function deductUserCredits(userId, amount) {
-    if (!db || amount <= 0) return;
+    if (!db || !amount || amount <= 0) return;
     
     try {
         const ref = db.ref(`users/${userId}/credits`);
         await ref.transaction((current_credits) => {
-            return (current_credits || 0) - amount;
+            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} from User ${userId}`);
+        console.log(`[Credits] Deducted ${amount} tokens from User ${userId}`);
     } catch (err) {
         console.error(`[Credits] Failed to deduct ${amount} from ${userId}:`, err);
     }
@@ -97,14 +87,13 @@ async function deductUserCredits(userId, amount) {
 // --- MIDDLEWARE ---
 
 const validateRequest = (req, res, next) => {
-    // Modified: Onboarding endpoints might now need userId for credit checks
     if (req.path.includes('/admin/cleanup')) return next();
     
     const { userId, projectId } = req.body;
     
-    // 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" });
+    // Ensure userId exists for credit checks
+    if (!userId && (req.path.includes('/onboarding') || req.path.includes('/new') || req.path.includes('/project'))) {
+        return res.status(400).json({ error: "Missing userId" });
     }
     
     next();
@@ -134,14 +123,10 @@ function extractImagePrompt(text) {
 
 // --- ADMIN ENDPOINTS ---
 
-/**
- * CLEANUP MEMORY
- * Call this via Cron Job or manually to free up RAM
- */
 app.get('/admin/cleanup', async (req, res) => {
     try {
         const removed = StateManager.cleanupMemory();
-        res.json({ success: true, removedCount: removed, message: `Cleaned ${removed} inactive projects from memory.` });
+        res.json({ success: true, removedCount: removed });
     } catch (err) {
         res.status(500).json({ error: "Cleanup failed" });
     }
@@ -150,32 +135,32 @@ app.get('/admin/cleanup', async (req, res) => {
 // --- ONBOARDING ENDPOINTS ---
 
 app.post('/onboarding/analyze', validateRequest, async (req, res) => {
-    const { description, userId } = req.body; // userId now required for credit check
+    const { description, userId } = req.body;
     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);
+        // 1. Pre-flight Check
+        await checkMinimumCredits(userId);
 
         console.log(`[Onboarding] Analyzing idea...`);
-        tokenUsage += estimateTokens(description);
-
+        
+        // 2. AI Call
+        // Result is { questions: [...], usage: { totalTokenCount: ... }, ... }
         const result = await AIEngine.generateEntryQuestions(description);
-        tokenUsage += estimateTokens(JSON.stringify(result));
         
-        // 2. Deduct Credits
-        await deductUserCredits(userId, tokenUsage);
+        const usage = result.usage?.totalTokenCount || 0;
+
+        // 3. Billing (Only on success)
+        if (usage > 0) await deductUserCredits(userId, usage);
 
         if (result.status === "REJECTED") {
             return res.json({ 
                 rejected: true, 
-                reason: result.reason || "Idea violates TOS or guidelines." 
+                reason: result.reason || "Idea violates TOS." 
             });
         }
         res.json({ questions: result.questions });
+
     } catch (err) {
         console.error(err);
         if (err.message.includes("Insufficient credits")) {
@@ -187,64 +172,48 @@ app.post('/onboarding/analyze', validateRequest, async (req, res) => {
 
 app.post('/onboarding/create', validateRequest, async (req, res) => {
     const { userId, description, answers } = req.body;
-    let tokenUsage = 0;
+    let totalTokens = 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);
-        };
+        // 1. Pre-flight Check
+        await checkMinimumCredits(userId);
 
+        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}...`);
-        
-        // Track Input Usage
-        tokenUsage += estimateTokens(description);
-        tokenUsage += estimateTokens(JSON.stringify(answers));
 
-        const grading = await AIEngine.gradeProject(description, answers);
-        tokenUsage += estimateTokens(JSON.stringify(grading));
+        // 2. Grading
+        const gradingResult = await AIEngine.gradeProject(description, answers);
+        totalTokens += (gradingResult.usage?.totalTokenCount || 0);
 
-        const isFailure = grading.feasibility < 30 || grading.rating === 'F';
+        const isFailure = gradingResult.feasibility < 30 || gradingResult.rating === 'F';
 
         let thumbnailBase64 = null;
 
-        if (isFailure) {
-            console.log(`[Onboarding] ❌ Project Failed Grading. Skipping Image.`);
-        } else {
+        if (!isFailure) {
             console.log(`[Onboarding] ✅ Passed. Generating Thumbnail...`);
-            const imagePrompt = `Game Title: ${grading.title}. Core Concept: ${description}`;
-            
-            // Track Prompt Usage
-            tokenUsage += estimateTokens(imagePrompt);
+            const imagePrompt = `Game Title: ${gradingResult.title}. Core Concept: ${description}`;
             
-            thumbnailBase64 = await AIEngine.generateImage(imagePrompt);
-            
-            // Track Image Cost
-            if (thumbnailBase64) tokenUsage += IMAGE_GENERATION_COST;
+            // 3. Image Generation
+            const imgResult = await AIEngine.generateImage(imagePrompt);
+            if (imgResult) {
+                thumbnailBase64 = imgResult.image;
+                totalTokens += (imgResult.usage?.totalTokenCount || 0);
+            }
         }
 
-        // --- Database & Storage Logic ---
+        // Upload Logic
         let thumbnailUrl = null;
         if (thumbnailBase64 && storage) {
             try {
-                console.log(`[Onboarding] 📤 Uploading Thumbnail to Storage...`);
                 const base64Data = thumbnailBase64.replace(/^data:image\/\w+;base64,/, "");
                 const buffer = Buffer.from(base64Data, 'base64');
                 const bucket = storage.bucket();
                 const file = bucket.file(`${projectId}/thumbnail.png`);
-
-                await file.save(buffer, {
-                    metadata: { contentType: 'image/png' }
-                });
-
+                await file.save(buffer, { metadata: { contentType: 'image/png' } });
                 await file.makePublic();
                 thumbnailUrl = `https://storage.googleapis.com/${bucket.name}/${projectId}/thumbnail.png`;
-                console.log(`[Onboarding] 🖼️ Thumbnail URL: ${thumbnailUrl}`);
             } catch (uploadErr) {
                 console.error("Storage Upload Failed:", uploadErr);
             }
@@ -253,75 +222,66 @@ app.post('/onboarding/create', validateRequest, async (req, res) => {
         const timestamp = Date.now();
         const status = isFailure ? "rejected" : "Idle";
 
-        const memoryObject = {
-            id: projectId,
-            userId,
-            title: grading.title || "Untitled Project",
-            description,
-            answers,
-            stats: grading,
-            thumbnail: thumbnailUrl, 
-            createdAt: timestamp,
-            status,
-            workerHistory: [],
-            pmHistory: [],
-            commandQueue: [],
-            failureCount: 0
-        };
+        // Save Data
+        if (!isFailure) {
+            const memoryObject = {
+                id: projectId,
+                userId,
+                title: gradingResult.title || "Untitled",
+                description,
+                answers,
+                stats: gradingResult,
+                thumbnail: thumbnailUrl, 
+                createdAt: timestamp,
+                status,
+                workerHistory: [],
+                pmHistory: [],
+                commandQueue: [],
+                failureCount: 0
+            };
+            await StateManager.updateProject(projectId, memoryObject);
+        }
 
+        // DB Updates
         if (firestore && !isFailure) {
-            try {
-                await firestore.collection('projects').doc(projectId).set({
-                    id: projectId,
-                    userId: userId,
-                    assets: thumbnailUrl ? [thumbnailUrl] : [],
-                    createdAt: admin.firestore.FieldValue.serverTimestamp()
-                });
-                console.log(`[Onboarding] 📝 Firestore Doc Created`);
-            } catch (fsErr) {
-                console.error("Firestore Write Failed:", fsErr);
-            }
+            await firestore.collection('projects').doc(projectId).set({
+                id: projectId,
+                userId: userId,
+                assets: thumbnailUrl ? [thumbnailUrl] : [],
+                createdAt: admin.firestore.FieldValue.serverTimestamp()
+            });
         }
 
-        if (db  && !isFailure) {
+        if (db && !isFailure) {
             const updates = {};
             updates[`projects/${projectId}/info`] = {
                 id: projectId,
                 userId,
-                title: grading.title || "Untitled Project",
+                title: gradingResult.title || "Untitled",
                 description,
                 answers,
-                stats: grading,
+                stats: gradingResult,
                 createdAt: timestamp,
                 status
             };
-
-            if (thumbnailUrl) {
-                updates[`projects/${projectId}/thumbnail`] = { url: thumbnailUrl };
-            }
-
+            if (thumbnailUrl) updates[`projects/${projectId}/thumbnail`] = { url: thumbnailUrl };
             updates[`projects/${projectId}/state`] = {
                 workerHistory: [],
                 pmHistory: [],
                 commandQueue: [],
                 failureCount: 0
             };
-
             await db.ref().update(updates);
         }
-        
-        if (!isFailure) {
-            await StateManager.updateProject(projectId, memoryObject);
-        }
 
-        // 2. Deduct Credits
-        await deductUserCredits(userId, tokenUsage);
+        // 4. Billing (Only on success)
+        if (totalTokens > 0) await deductUserCredits(userId, totalTokens);
 
         res.json({ 
             success: !isFailure, 
             projectId, 
-            stats: grading, 
-            title: grading.title || "Untitled Project",
+            stats: gradingResult, 
+            title: gradingResult.title || "Untitled",
             thumbnail: thumbnailBase64
         });
 
@@ -336,78 +296,74 @@ app.post('/onboarding/create', validateRequest, async (req, res) => {
 
 // --- CORE ENDPOINTS ---
 
-// NEW: Helper function to run initialization in background
 async function runBackgroundInitialization(projectId, userId, description) {
     console.log(`[Background] Starting initialization for ${projectId}`);
-    let tokenUsage = 0;
+    let totalTokens = 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);
+        const gddResult = await AIEngine.callPM(pmHistory, gddPrompt);
         
+        totalTokens += (gddResult.usage?.totalTokenCount || 0);
+        const gddText = gddResult.text;
+
         pmHistory.push({ role: 'user', parts: [{ text: gddPrompt }] });
-        pmHistory.push({ role: 'model', parts: [{ text: gddResponse }] });
+        pmHistory.push({ role: 'model', parts: [{ text: gddText }] });
 
         // 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);
+        const taskResult = await AIEngine.callPM(pmHistory, taskPrompt);
+        
+        totalTokens += (taskResult.usage?.totalTokenCount || 0);
+        const taskText = taskResult.text;
 
         pmHistory.push({ role: 'user', parts: [{ text: taskPrompt }] });
-        pmHistory.push({ role: 'model', parts: [{ text: taskResponse }] });
+        pmHistory.push({ role: 'model', parts: [{ text: taskText }] });
 
         // 3. Initialize Worker
-        const initialWorkerInstruction = extractWorkerPrompt(taskResponse) || `Initialize structure for: ${description}`;
+        const initialWorkerInstruction = extractWorkerPrompt(taskText) || `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);
+        
+        const workerResult = await AIEngine.callWorker(workerHistory, initialWorkerPrompt, []);
+        
+        totalTokens += (workerResult.usage?.totalTokenCount || 0);
+        const workerText = workerResult.text;
 
         workerHistory.push({ role: 'user', parts: [{ text: initialWorkerPrompt }] });
-        workerHistory.push({ role: 'model', parts: [{ text: workerResponse }] });
+        workerHistory.push({ role: 'model', parts: [{ text: workerText }] });
 
-        // 4. Update Memory & DB
+        // 4. Update Memory
         await StateManager.updateProject(projectId, {
             userId,
             pmHistory,
             workerHistory,
-            gdd: gddResponse,
+            gdd: gddText,
             failureCount: 0
         });
 
-        // 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. Queue commands (Handle assets + their costs)
+        // We pass userId to deduct image costs immediately inside this function
+        await processAndQueueResponse(projectId, workerText, userId); 
         
-        // 5. Update Status to IDLE
+        // 6. Update Status
         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}. Tokens used: ${tokenUsage}`);
+        // 7. Deduct accumulated text tokens
+        if (totalTokens > 0) await deductUserCredits(userId, totalTokens);
+        
+        console.log(`[Background] Init complete. Text Tokens: ${totalTokens}`);
 
     } 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");
+        // No deduction on failure
     }
 }
 
@@ -415,23 +371,18 @@ app.post('/new/project', validateRequest, async (req, res) => {
     const { userId, projectId, description } = req.body;
     
     try {
-        // 1. Credit Check (Prevent starting if low balance)
-        await checkUserCredits(userId);
+        await checkMinimumCredits(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 });
         }
@@ -441,57 +392,49 @@ 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 tokenUsage = 0;
+    let totalTokens = 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." });
-    }
-
     try {
-        // 2. Credit Check
-        await checkUserCredits(userId);
+        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);
         
-        // 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);
+            
+            const pmResult = await AIEngine.callPM(project.pmHistory, summary);
+            totalTokens += (pmResult.usage?.totalTokenCount || 0);
+            const pmText = pmResult.text;
             
             project.pmHistory.push({ role: 'user', parts: [{ text: summary }] });
-            project.pmHistory.push({ role: 'model', parts: [{ text: pmResponse }] });
+            project.pmHistory.push({ role: 'model', parts: [{ text: pmText }] });
 
-            const nextInstruction = extractWorkerPrompt(pmResponse);
+            const nextInstruction = extractWorkerPrompt(pmText);
             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);
+                 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);
+                 
                  return res.json({ success: true, message: "No further tasks. Project Idle." });
             }
 
             const newWorkerHistory = [];
             const newPrompt = `New Objective: ${nextInstruction}`;
-            tokenUsage += estimateTokens(newPrompt);
-
-            const workerResponse = await AIEngine.callWorker(newWorkerHistory, newPrompt, []);
-            tokenUsage += estimateTokens(workerResponse);
+            
+            const workerResult = await AIEngine.callWorker(newWorkerHistory, newPrompt, []);
+            totalTokens += (workerResult.usage?.totalTokenCount || 0);
+            const workerText = workerResult.text;
 
             newWorkerHistory.push({ role: 'user', parts: [{ text: newPrompt }] });
-            newWorkerHistory.push({ role: 'model', parts: [{ text: workerResponse }] });
+            newWorkerHistory.push({ role: 'model', parts: [{ text: workerText }] });
 
             await StateManager.updateProject(projectId, {
                 pmHistory: project.pmHistory,
@@ -500,20 +443,16 @@ app.post('/project/feedback', async (req, res) => {
             });
 
             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.
+            await processAndQueueResponse(projectId, workerText, userId);
             
-            // Still working on next task
-            if(db) {
-                await db.ref(`projects/${projectId}/info`).update({
-                    status: "working",
-                    lastUpdated: Date.now()
-                });
-            }
+            if(db) await db.ref(`projects/${projectId}/info`).update({ status: "working", lastUpdated: Date.now() });
 
-            await deductUserCredits(userId, tokenUsage);
+            // Billing
+            if (totalTokens > 0) await deductUserCredits(userId, totalTokens);
             return res.json({ success: true, message: "Next Task Assigned" });
         }
 
+        // Logic for Failure Handling (PM Guidance) or Standard Feedback
         let isFailure = false;
         if (logContext?.logs) {
             const errs = ["Error", "Exception", "failed", "Stack Begin", "Infinite yield"];
@@ -525,76 +464,74 @@ app.post('/project/feedback', async (req, res) => {
 
         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);
+            const pmResult = await AIEngine.callPM(project.pmHistory, pmPrompt);
+            totalTokens += (pmResult.usage?.totalTokenCount || 0);
+            const pmVerdict = pmResult.text;
 
             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);
+                
+                const workerResult = await AIEngine.callWorker(resetHistory, resetPrompt, []);
+                totalTokens += (workerResult.usage?.totalTokenCount || 0);
                 
                 resetHistory.push({ role: 'user', parts: [{ text: resetPrompt }] });
-                resetHistory.push({ role: 'model', parts: [{ text: workerResp }] });
+                resetHistory.push({ role: 'model', parts: [{ text: workerResult.text }] });
 
                 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);
+                await processAndQueueResponse(projectId, workerResult.text, userId);
+                
+                if (totalTokens > 0) await deductUserCredits(userId, totalTokens);
                 return res.json({ success: true, message: "Worker Terminated." });
             } else {
                 const injection = `[PM GUIDANCE]: ${pmVerdict} \n\nApply this fix now.`;
-                tokenUsage += estimateTokens(injection);
+                const workerResult = await AIEngine.callWorker(project.workerHistory, injection, []);
+                totalTokens += (workerResult.usage?.totalTokenCount || 0);
 
-                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 }] });
+                project.workerHistory.push({ role: 'model', parts: [{ text: workerResult.text }] });
                 
                 await StateManager.updateProject(projectId, { workerHistory: project.workerHistory, failureCount: 1 });
-                await processAndQueueResponse(projectId, workerResp, userId);
-                await deductUserCredits(userId, tokenUsage);
+                await processAndQueueResponse(projectId, workerResult.text, userId);
+                
+                if (totalTokens > 0) await deductUserCredits(userId, totalTokens);
                 return res.json({ success: true, message: "PM Guidance Applied." });
             }
         }
 
-        // Standard Logic
+        // Standard Interaction
         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);
+        let workerResult = await AIEngine.callWorker(project.workerHistory, fullInput, images || []);
+        totalTokens += (workerResult.usage?.totalTokenCount || 0);
+        let responseText = workerResult.text;
+        
+        const pmQuestion = extractPMQuestion(responseText);
         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);
+            const pmResult = await AIEngine.callPM(project.pmHistory, pmConsultPrompt);
+            totalTokens += (pmResult.usage?.totalTokenCount || 0);
+            const pmAnswer = pmResult.text;
 
             project.pmHistory.push({ role: 'user', parts: [{ text: pmConsultPrompt }] });
             project.pmHistory.push({ role: 'model', parts: [{ text: pmAnswer }] });
 
             const injectionMsg = `[PM RESPONSE]: ${pmAnswer}`;
             project.workerHistory.push({ role: 'user', parts: [{ text: fullInput }] });
-            project.workerHistory.push({ role: 'model', parts: [{ text: response }] });
+            project.workerHistory.push({ role: 'model', parts: [{ text: responseText }] });
             
-            tokenUsage += estimateTokens(injectionMsg);
-            response = await AIEngine.callWorker(project.workerHistory, injectionMsg, []);
-            tokenUsage += estimateTokens(response);
+            // Re-call worker with answer
+            workerResult = await AIEngine.callWorker(project.workerHistory, injectionMsg, []);
+            totalTokens += (workerResult.usage?.totalTokenCount || 0);
+            responseText = workerResult.text;
 
             project.workerHistory.push({ role: 'user', parts: [{ text: injectionMsg }] });
-            project.workerHistory.push({ role: 'model', parts: [{ text: response }] });
+            project.workerHistory.push({ role: 'model', parts: [{ text: responseText }] });
         } else {
             project.workerHistory.push({ role: 'user', parts: [{ text: fullInput }] });
-            project.workerHistory.push({ role: 'model', parts: [{ text: response }] });
+            project.workerHistory.push({ role: 'model', parts: [{ text: responseText }] });
         }
 
         await StateManager.updateProject(projectId, { 
@@ -603,25 +540,17 @@ app.post('/project/feedback', async (req, res) => {
             failureCount: project.failureCount 
         });
         
-        await processAndQueueResponse(projectId, response, userId);
+        await processAndQueueResponse(projectId, responseText, userId);
         
-        // --- UPDATED: Set status to IDLE and update lastUpdated ---
-        if(db) {
-            await db.ref(`projects/${projectId}/info`).update({
-                status: "idle",
-                lastUpdated: Date.now()
-            });
-        }
+        if(db) await db.ref(`projects/${projectId}/info`).update({ status: "idle", lastUpdated: Date.now() });
       
-        // Final Deduction
-        await deductUserCredits(userId, tokenUsage);
+        // Final Billing
+        if (totalTokens > 0) await deductUserCredits(userId, totalTokens);
 
         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 });
         }
@@ -631,21 +560,11 @@ app.post('/project/feedback', async (req, res) => {
 
 app.post('/project/ping', async (req, res) => {
     const { projectId, userId } = req.body;
-
-    if (!projectId || !userId) {
-        return res.status(400).json({ error: "Missing ID fields" });
-    }
+    if (!projectId || !userId) return res.status(400).json({ error: "Missing ID fields" });
     
     const project = await StateManager.getProject(projectId);
-    
-    if (!project) {
-        return res.status(404).json({ action: "IDLE", 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." });
-    }
+    if (!project) return res.status(404).json({ action: "IDLE", error: "Project not found" });
+    if (project.userId !== userId) return res.status(403).json({ error: "Unauthorized" });
 
     const command = await StateManager.popCommand(projectId);
     
@@ -666,52 +585,60 @@ app.post('/project/ping', async (req, res) => {
 
 app.post('/human/override', validateRequest, async (req, res) => {
     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 });
-        }
-    }
+    let totalTokens = 0;
 
-    const project = await StateManager.getProject(projectId);
-    const overrideMsg = `[SYSTEM OVERRIDE]: ${instruction}`;
-    let tokenUsage = estimateTokens(overrideMsg);
+    try {
+        await checkMinimumCredits(userId);
 
-    if (pruneHistory && project.workerHistory.length >= 2) {
-        project.workerHistory.pop();
-        project.workerHistory.pop();
-    }
-    const response = await AIEngine.callWorker(project.workerHistory, overrideMsg, []);
-    tokenUsage += estimateTokens(response);
+        const project = await StateManager.getProject(projectId);
+        const overrideMsg = `[SYSTEM OVERRIDE]: ${instruction}`;
 
-    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, userId);
-    
-    if (userId) await deductUserCredits(userId, tokenUsage);
+        if (pruneHistory && project.workerHistory.length >= 2) {
+            project.workerHistory.pop();
+            project.workerHistory.pop();
+        }
 
-    res.json({ success: true });
+        const workerResult = await AIEngine.callWorker(project.workerHistory, overrideMsg, []);
+        totalTokens += (workerResult.usage?.totalTokenCount || 0);
+
+        project.workerHistory.push({ role: 'user', parts: [{ text: overrideMsg }] });
+        project.workerHistory.push({ role: 'model', parts: [{ text: workerResult.text }] });
+        
+        await StateManager.updateProject(projectId, { workerHistory: project.workerHistory });
+        await processAndQueueResponse(projectId, workerResult.text, userId);
+        
+        if (totalTokens > 0) await deductUserCredits(userId, totalTokens);
+        
+        res.json({ success: true });
+    } catch (err) {
+        if (err.message.includes("Insufficient credits")) {
+            return res.status(402).json({ error: err.message });
+        }
+        res.status(500).json({ error: "Override Failed" });
+    }
 });
 
-// Added userId param to handle asset cost if image generated
+// Helper to handle Asset Parsing & Billing
 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);
-        }
+        // 1. Generate Image (returns { image, usage } or null)
+        const imgResult = await AIEngine.generateImage(imgPrompt);
+        
+        if (imgResult && imgResult.image) {
+            // 2. Bill for image tokens immediately
+            const imgTokens = imgResult.usage?.totalTokenCount || 0;
+            if (userId && imgTokens > 0) {
+                await deductUserCredits(userId, imgTokens);
+            }
 
-        if (base64Image) {
-            await StateManager.queueCommand(projectId, { type: "CREATE_ASSET", payload: base64Image });
+            // 3. Queue creation command
+            await StateManager.queueCommand(projectId, { type: "CREATE_ASSET", payload: imgResult.image });
         }
     }
+    // Queue the raw text response for the frontend
     await StateManager.queueCommand(projectId, rawResponse);
 }