Update app2.js

app2.js

@@ -1,46 +1,140 @@
 import express from 'express';
 import bodyParser from 'body-parser';
 import cors from 'cors';
-import { StateManager } from './stateManager.js';
+import { StateManager, initDB } from './stateManager.js';
 import { AIEngine } from './aiEngine.js';
 import fs from 'fs';
 import admin from 'firebase-admin';
+import crypto from "crypto";
 
 // --- FIREBASE SETUP ---
-// Using environment variable injection as requested
 let db = null;
+let firestore = null;
+let storage = null;
+
 try {
-    if (process.env.FIREBASE_SERVICE_ACCOUNT_JSON && process.env.FIREBASE_SERVICE_ACCOUNT_JSON !== "") {
+    if (process.env.FIREBASE_SERVICE_ACCOUNT_JSON) {
         const serviceAccount = JSON.parse(process.env.FIREBASE_SERVICE_ACCOUNT_JSON);
-        admin.initializeApp({
-            credential: admin.credential.cert(serviceAccount),
-            // Make sure to replace this URL or add it to env variables as well
-            databaseURL: "https://shago-web-default-rtdb.firebaseio.com" 
-        });
+        
+        // Specific bucket as requested
+        const bucketName = `hollowpad-ai-default-rtdb.firebaseio.com` // `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
+            });
+        }
+        
         db = admin.database();
-        console.log("🔥 Firebase Connected");
+        firestore = admin.firestore();
+        storage = admin.storage();
+        
+        initDB(db);
+        console.log("🔥 Firebase Connected (RTDB, Firestore, Storage) & StateManager Linked");
     } else {
-        console.warn("⚠️ No FIREBASE_SERVICE_ACCOUNT_JSON found. Running in Memory-Only mode.");
+        console.warn("⚠️ Memory-Only mode.");
     }
-} catch (e) {
-    console.error("Firebase Init Error:", e);
-}
+} catch (e) { console.error("Firebase Init Error:", e); }
+
 
 const app = express();
 const PORT = process.env.PORT || 7860;
 
-// Load Prompts for Server-Side Logic Checks
-const sysPrompts = JSON.parse(fs.readFileSync('./prompts.json', 'utf8'));
+// ---
+/*
+
+app.use((req, res, next) => {
+  console.log('[REQ]', req.method, req.originalUrl);
+  console.log('[HEADERS]', req.headers);
+  console.log('[BODY]', req.body ?? '(unparseable or empty)');
+  next();
+});
+
+*/
+
+// ---
+
+
+// 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' }));
 
-// --- HELPERS ---
+// --- 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 of a specific type.
+ * @param {string} userId 
+ * @param {'basic'|'diamond'} type 
+ */
+async function checkMinimumCredits(userId, type = 'basic') {
+    if (!db) return; 
+
+    // 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;
+
+    const required = type === 'diamond' ? MIN_DIAMOND_REQUIRED : MIN_BASIC_REQUIRED;
+
+    if (credits < required) {
+      
+        res.status(500).json({ 
+          insufficient: true,
+          error:`Insufficient ${type} credits. You have ${credits}, need minimum ${required} to proceed.`
+   
+        });
+        // throw new Error(`Insufficient ${type} credits. You have ${credits}, need minimum ${required} to proceed.`);
+    }
+}
+
+/**
+ * Deducts exact tokens from the user's specific wallet.
+ * @param {string} userId 
+ * @param {number} amount 
+ * @param {'basic'|'diamond'} type 
+ */
+async function deductUserCredits(userId, amount, type = 'basic') {
+    if (!db || !amount || amount <= 0) return;
+    
+    try {
+        const ref = db.ref(`users/${userId}/credits/${type}`);
+        await ref.transaction((current_credits) => {
+            const current = current_credits || 0;
+            return Math.max(0, current - amount);
+        });
+        console.log(`[Credits] Deducted ${amount} ${type} credits from User ${userId}`);
+    } catch (err) {
+        console.error(`[Credits] Failed to deduct ${amount} ${type} from ${userId}:`, err);
+    }
+}
+
+// --- MIDDLEWARE ---
+
 const validateRequest = (req, res, next) => {
-    if (req.path.includes('/onboarding')) return next();
+    if (req.path.includes('/admin/cleanup')) return next();
     
     const { userId, projectId } = req.body;
-    if (!userId || !projectId) return res.status(400).json({ error: "Missing ID" });
+    
+    // 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();
 };
 
@@ -51,9 +145,12 @@ function extractWorkerPrompt(text) {
 
 function formatContext({ hierarchyContext, scriptContext, logContext }) {
     let out = "";
-    if (scriptContext) out += `\n[TARGET SCRIPT]: ${scriptContext.targetName}\n[SOURCE PREVIEW]: ${scriptContext.scriptSource?.substring(0, 1000)}...`;
+  if (scriptContext) out += `\n[TARGET SCRIPT]: ${scriptContext.targetName}\n[SOURCE PREVIEW]: ${scriptContext.scriptSource}`;
+   // if (scriptContext) out += `\n[TARGET SCRIPT]: ${scriptContext.targetName}\n[SOURCE PREVIEW]: ${scriptContext.scriptSource?.substring(0, 1000)}...`;
     if (logContext) out += `\n[LAST LOGS]: ${logContext.logs}`;
-    return out;
+   // my own line
+  if (hierarchyContext) out += `\n[Hierarchy Context]: ${hierarchyContext}`;
+  return out;
 }
 
 function extractPMQuestion(text) {
@@ -66,268 +163,431 @@ function extractImagePrompt(text) {
     return match ? match[1].trim() : null;
 }
 
+// --- ADMIN ENDPOINTS ---
+
+app.get('/admin/cleanup', async (req, res) => {
+    try {
+        const removed = StateManager.cleanupMemory();
+        res.json({ success: true, removedCount: removed });
+    } catch (err) {
+        res.status(500).json({ error: "Cleanup failed" });
+    }
+});
+
 // --- ONBOARDING ENDPOINTS ---
-/**
- * 1. ANALYZE
- */
+
 app.post('/onboarding/analyze', validateRequest, async (req, res) => {
-    const { description } = req.body;
+    const { description, userId } = req.body;
+  
     if (!description) return res.status(400).json({ error: "Description required" });
 
     try {
+        // Analyst uses BASIC models (Flash)
+        await checkMinimumCredits(userId, 'basic');
+
         console.log(`[Onboarding] Analyzing idea...`);
+        
         const result = await AIEngine.generateEntryQuestions(description);
         
+        const usage = result.usage?.totalTokenCount || 0;
+
+        // Bill Basic
+        if (usage > 0) await deductUserCredits(userId, usage, 'basic');
+
         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")) {
+            return res.status(402).json({ error: err.message });
+        }
         res.status(500).json({ error: "Analysis failed" });
     }
 });
 
-/**
- * 2. CREATE
- */
 app.post('/onboarding/create', validateRequest, async (req, res) => {
     const { userId, description, answers } = req.body;
-    const projectId = "proj_" + Date.now(); 
+    let basicTokens = 0; // Create flow uses Grader (Basic) + Image (Basic)
 
     try {
+        // 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}...`);
 
-        // STEP 1: GRADE
-        const grading = await AIEngine.gradeProject(description, answers);
-        
-        // STEP 2: CHECK FAIL CONDITIONS (Relaxed)
-        // Only fail if Feasibility is extremely low (< 30) or Rating is F
-        const isFailure = grading.feasibility < 30 || grading.rating === 'F';
+        // 1. Grading (Basic)
+        const gradingResult = await AIEngine.gradeProject(description, answers);
+        basicTokens += (gradingResult.usage?.totalTokenCount || 0);
+
+        const isFailure = gradingResult.feasibility < 30 || gradingResult.rating === 'F';
 
         let thumbnailBase64 = null;
 
-        if (isFailure) {
-            console.log(`[Onboarding] ❌ Project Failed Grading (${grading.rating}). Skipping Image.`);
-        } else {
+       /* 
+       * image generator 
+       if (!isFailure) {
             console.log(`[Onboarding] ✅ Passed. Generating Thumbnail...`);
+            const imagePrompt = `Game Title: ${gradingResult.title}. Core Concept: ${description}`;
             
-            // --- CRITICAL FIX HERE ---
-            // We MUST send the description so the AI knows what to draw.
-            const imagePrompt = `Game Title: ${grading.title}. Core Concept: ${description}`;
-            thumbnailBase64 = await AIEngine.generateImage(imagePrompt);
+            // 2. Image Generation (Basic)
+            const imgResult = await AIEngine.generateImage(imagePrompt);
+            if (imgResult) {
+                thumbnailBase64 = imgResult.image;
+                // Add token usage for text + fixed cost
+                basicTokens += (imgResult.usage || 0); 
+                if (imgResult.image) basicTokens += IMAGE_COST_BASIC;
+            }
+        } */
+
+        // Upload Logic
+        let thumbnailUrl = null;
+        if (thumbnailBase64 && storage) {
+            try {
+                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.makePublic();
+                thumbnailUrl = `https://storage.googleapis.com/${bucket.name}/${projectId}/thumbnail.png`;
+            } catch (uploadErr) {
+                console.error("Storage Upload Failed:", uploadErr);
+            }
         }
 
-        const projectData = {
-            id: projectId,
-            userId,
-            title: grading.title || "Untitled Project",
-            description,
-            answers,
-            stats: grading,
-            thumbnail: thumbnailBase64, // ? `data:image/png;base64,${thumbnailBase64}` : null,
-            createdAt: Date.now(),
-            status: isFailure ? "rejected" : "initialized"
-        };
-
-        if (db) await db.ref(`projects/${projectId}`).set(projectData);
-        
+        const timestamp = Date.now();
+        const status = isFailure ? "rejected" : "Idle";
+
+        // Save Data Logic
         if (!isFailure) {
-            await StateManager.updateProject(projectId, {
-                ...projectData,
+            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);
+        }
+
+        if (firestore && !isFailure) {
+            await firestore.collection('projects').doc(projectId).set({
+                id: projectId,
+                userId: userId,
+                assets: thumbnailUrl ? [thumbnailUrl] : [],
+                createdAt: admin.firestore.FieldValue.serverTimestamp()
             });
         }
 
+        if (db && !isFailure) {
+            const updates = {};
+            updates[`projects/${projectId}/info`] = {
+                id: projectId,
+                userId,
+                title: gradingResult.title || "Untitled",
+                description,
+                answers,
+                stats: gradingResult,
+                createdAt: timestamp,
+                status
+            };
+            if (thumbnailUrl) updates[`projects/${projectId}/thumbnail`] = { url: thumbnailUrl };
+            updates[`projects/${projectId}/state`] = {
+                workerHistory: [],
+                pmHistory: [],
+                commandQueue: [],
+                failureCount: 0
+            };
+            await db.ref().update(updates);
+        }
+
+        // Deduct Basic Credits
+        if (basicTokens > 0) await deductUserCredits(userId, basicTokens, 'basic');
+       console.log("sending");
         res.json({ 
+          status: 200,
             success: !isFailure, 
             projectId, 
-            stats: grading, 
-            title: projectData.title,
-            thumbnail: projectData.thumbnail 
+            stats: gradingResult, 
+            title: gradingResult.title || "Untitled",
+            thumbnail: thumbnailBase64
         });
 
     } catch (err) {
         console.error("Create Error:", err);
+        if (err.message.includes("Insufficient")) {
+            return res.status(402).json({ error: err.message });
+        }
         res.status(500).json({ error: "Creation failed" });
     }
 });
 
+// --- CORE ENDPOINTS ---
 
-// --- CORE WORKFLOW ENDPOINTS ---
-
-/**
- * 3. INITIALIZE WORKSPACE (PM Generates GDD -> First Task)
- */
-app.post('/new/project', validateRequest, async (req, res) => {
-    const { userId, projectId, description } = req.body;
+async function runBackgroundInitialization(projectId, userId, description) {
+    console.log(`[Background] Starting initialization for ${projectId}`);
+    
+    // Separate Counters
+    let diamondUsage = 0; // For PM
+    let basicUsage = 0;   // For Worker
 
     try {
         const pmHistory = [];
+        
+        // 1. Generate GDD (PM -> Diamond)
         const gddPrompt = `Create a comprehensive GDD for: ${description}`;
-        const gddResponse = await AIEngine.callPM(pmHistory, gddPrompt);
+        const gddResult = await AIEngine.callPM(pmHistory, gddPrompt);
         
+        diamondUsage += (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 (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 taskResponse = await AIEngine.callPM(pmHistory, taskPrompt);
+        const taskResult = await AIEngine.callPM(pmHistory, taskPrompt);
+        
+        diamondUsage += (taskResult.usage?.totalTokenCount || 0);
+        const taskText = taskResult.text;
 
         pmHistory.push({ role: 'user', parts: [{ text: taskPrompt }] });
-        pmHistory.push({ role: 'model', parts: [{ text: taskResponse }] });
-
-        const initialWorkerInstruction = extractWorkerPrompt(taskResponse) || `Initialize structure for: ${description}`;
+        pmHistory.push({ role: 'model', parts: [{ text: taskText }] });
 
+        // 3. Initialize Worker (Worker -> Basic)
+        const initialWorkerInstruction = extractWorkerPrompt(taskText) || `Initialize structure for: ${description}`;
         const workerHistory = [];
         const initialWorkerPrompt = `CONTEXT: New Project. \nINSTRUCTION: ${initialWorkerInstruction}`;
-        const workerResponse = await AIEngine.callWorker(workerHistory, initialWorkerPrompt, []);
+        
+        const workerResult = await AIEngine.callWorker(workerHistory, initialWorkerPrompt, []);
+        
+        basicUsage += (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 }] });
 
-        // Update State
+        // 4. Update Memory
         await StateManager.updateProject(projectId, {
             userId,
             pmHistory,
             workerHistory,
-            gdd: gddResponse,
+            gdd: gddText,
             failureCount: 0
         });
 
-        // Queue Actions
-        await processAndQueueResponse(projectId, workerResponse);
-
-        res.json({ success: true, message: "Workspace Initialized", gddPreview: gddResponse.substring(0, 200) });
+        // 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
+        if(db) await db.ref(`projects/${projectId}/info`).update({
+            status: "IDLE",
+            lastUpdated: Date.now()
+        });
+        
+        // 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. Diamond: ${diamondUsage}, Basic: ${basicUsage}`);
 
     } catch (err) {
-        console.error("Init Error:", err);
-        res.status(500).json({ error: "Initialization Failed" });
+        console.error(`[Background] Init Error for ${projectId}:`, err);
+        if(db) await db.ref(`projects/${projectId}/info/status`).set("error");
     }
-});
+}
 
-/**
- * 4. FEEDBACK LOOP (The "Brain")
- */
-app.post('/project/feedback', async (req, res) => {
-    const { projectId, prompt, hierarchyContext, scriptContext, logContext, taskComplete, images } = req.body;
+app.post('/new/project', validateRequest, async (req, res) => {
+    const { userId, projectId, description } = req.body;
     
-    const project = await StateManager.getProject(projectId);
-    if (!project) return res.status(404).json({ error: "Project not found." });
+    try {
+        // Init requires BOTH types to be safe
+        await checkMinimumCredits(userId, 'diamond');
+        await checkMinimumCredits(userId, 'basic');
 
-    // A. TASK COMPLETE -> SCOPE SWITCH
-    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);
+        if(db) db.ref(`projects/${projectId}/info/status`).set("initializing");
         
-        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" });
-             return res.json({ success: true, message: "No further tasks. Project Idle." });
-        }
-
-        console.log(`[${projectId}] 🧹 Nuking Worker for next task.`);
-        const newWorkerHistory = [];
-        const newPrompt = `New Objective: ${nextInstruction}`;
-        const workerResponse = await AIEngine.callWorker(newWorkerHistory, newPrompt, []);
-
-        newWorkerHistory.push({ role: 'user', parts: [{ text: newPrompt }] });
-        newWorkerHistory.push({ role: 'model', parts: [{ text: workerResponse }] });
-
-        await StateManager.updateProject(projectId, {
-            pmHistory: project.pmHistory,
-            workerHistory: newWorkerHistory,
-            failureCount: 0
+        res.json({ 
+            success: true, 
+            message: "Project initialization started in background."
         });
 
-        StateManager.queueCommand(projectId, { type: "EXECUTE", payload: "warn('Starting Next Task...')" });
-        await processAndQueueResponse(projectId, workerResponse);
+        runBackgroundInitialization(projectId, userId, description);
 
-        return res.json({ success: true, message: "Next Task Assigned" });
-    }
-
-    // B. ERROR ESCALATION
-    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;
+    } catch (err) {
+        if (err.message.includes("Insufficient")) {
+            return res.status(402).json({ error: err.message });
         }
+        res.status(500).json({ error: "Failed to start project" });
     }
+});
 
-    if (project.failureCount > 3) {
-        console.log(`[${projectId}] 🚨 Escalating to PM...`);
-        const pmPrompt = sysPrompts.pm_guidance_prompt.replace('{{LOGS}}', logContext?.logs);
-        const pmVerdict = await AIEngine.callPM(project.pmHistory, pmPrompt);
+app.post('/project/feedback', async (req, res) => {
+    const { userId, projectId, prompt, hierarchyContext, scriptContext, logContext, taskComplete, images } = req.body;
+    
+    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" });
 
-        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, []);
+        // Basic Check (most interaction is worker)
+        await checkMinimumCredits(userId, 'basic');
+        
+        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.`;
             
-            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);
+            // PM Call -> Diamond
+            const pmResult = await AIEngine.callPM(project.pmHistory, summary);
+            diamondUsage += (pmResult.usage?.totalTokenCount || 0);
+            const pmText = pmResult.text;
             
-            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.pmHistory.push({ role: 'user', parts: [{ text: summary }] });
+            project.pmHistory.push({ role: 'model', parts: [{ text: pmText }] });
+
+            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() });
+                 
+                 // Deduct what we used
+                 if (diamondUsage > 0) await deductUserCredits(userId, diamondUsage, 'diamond');
+                 
+                 return res.json({ success: true, message: "No further tasks. Project Idle." });
+            }
+
+            const newWorkerHistory = [];
+            const newPrompt = `New Objective: ${nextInstruction}`;
             
-            project.workerHistory.push({ role: 'user', parts: [{ text: injection }] });
-            project.workerHistory.push({ role: 'model', parts: [{ text: workerResp }] });
+            // Worker Call -> Basic
+            const workerResult = await AIEngine.callWorker(newWorkerHistory, newPrompt, []);
+            basicUsage += (workerResult.usage?.totalTokenCount || 0);
+            const workerText = workerResult.text;
+
+            newWorkerHistory.push({ role: 'user', parts: [{ text: newPrompt }] });
+            newWorkerHistory.push({ role: 'model', parts: [{ text: workerText }] });
+
+            await StateManager.updateProject(projectId, {
+                pmHistory: project.pmHistory,
+                workerHistory: newWorkerHistory,
+                failureCount: 0
+            });
+
+            StateManager.queueCommand(projectId, { type: "EXECUTE", payload: "warn('Starting Next Task...')" });
+            await processAndQueueResponse(projectId, workerText, userId);
             
-            await StateManager.updateProject(projectId, { workerHistory: project.workerHistory, failureCount: 1 });
-            await processAndQueueResponse(projectId, workerResp);
+            if(db) await db.ref(`projects/${projectId}/info`).update({ status: "working", lastUpdated: Date.now() });
+
+            // Final Bill
+            if (diamondUsage > 0) await deductUserCredits(userId, diamondUsage, 'diamond');
+            if (basicUsage > 0) await deductUserCredits(userId, basicUsage, 'basic');
+
+            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 || "Analyze logs: {{LOGS}}").replace('{{LOGS}}', logContext?.logs);
             
-            return res.json({ success: true, message: "PM Guidance Applied." });
+            // PM -> Diamond
+            const pmResult = await AIEngine.callPM(project.pmHistory, pmPrompt);
+            diamondUsage += (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}`;
+                
+                // Worker -> Basic
+                const workerResult = await AIEngine.callWorker(resetHistory, resetPrompt, []);
+                basicUsage += (workerResult.usage?.totalTokenCount || 0);
+                
+                resetHistory.push({ role: 'user', parts: [{ text: resetPrompt }] });
+                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, workerResult.text, userId);
+                
+                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, []);
+                basicUsage += (workerResult.usage?.totalTokenCount || 0);
+
+                project.workerHistory.push({ role: 'user', parts: [{ text: injection }] });
+                project.workerHistory.push({ role: 'model', parts: [{ text: workerResult.text }] });
+                
+                await StateManager.updateProject(projectId, { workerHistory: project.workerHistory, failureCount: 1 });
+                await processAndQueueResponse(projectId, workerResult.text, userId);
+                
+                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." });
+            }
         }
-    }
 
-    // C. STANDARD LOOP + CONSULTATION
-    try {
+        // Standard Interaction (Worker Only = Basic)
         const fullInput = `USER: ${prompt || "Automatic Feedback"}` + formatContext({ hierarchyContext, scriptContext, logContext });
         
-        let response = await AIEngine.callWorker(project.workerHistory, fullInput, images || []);
+        let workerResult = await AIEngine.callWorker(project.workerHistory, fullInput, images || []);
+        basicUsage += (workerResult.usage?.totalTokenCount || 0);
+        let responseText = workerResult.text;
         
-        // CHECK: Consultant Question
-        const pmQuestion = extractPMQuestion(response);
+        // PM Consultation Check
+        const pmQuestion = extractPMQuestion(responseText);
         if (pmQuestion) {
-            console.log(`[${projectId}] 🙋 Worker asking PM: "${pmQuestion}"`);
-            
+            // PM -> Diamond
             const pmConsultPrompt = `[WORKER CONSULTATION]: The Worker asks: "${pmQuestion}"\nProvide a technical answer to unblock them.`;
-            const pmAnswer = await AIEngine.callPM(project.pmHistory, pmConsultPrompt);
+            const pmResult = await AIEngine.callPM(project.pmHistory, pmConsultPrompt);
+            diamondUsage += (pmResult.usage?.totalTokenCount || 0);
+            const pmAnswer = pmResult.text;
 
-            // Update PM Context
             project.pmHistory.push({ role: 'user', parts: [{ text: pmConsultPrompt }] });
             project.pmHistory.push({ role: 'model', parts: [{ text: pmAnswer }] });
 
-            // Feed Answer Back
             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 }] });
             
-            // Resume Worker
-            response = await AIEngine.callWorker(project.workerHistory, injectionMsg, []);
+            // Re-call Worker -> Basic
+            workerResult = await AIEngine.callWorker(project.workerHistory, injectionMsg, []);
+            basicUsage += (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, { 
@@ -336,20 +596,33 @@ app.post('/project/feedback', async (req, res) => {
             failureCount: project.failureCount 
         });
         
-        await processAndQueueResponse(projectId, response);
+        await processAndQueueResponse(projectId, responseText, userId);
+        
+        if(db) await db.ref(`projects/${projectId}/info`).update({ status: "idle", lastUpdated: Date.now() });
+      
+        // Final Billing
+        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")) {
+            return res.status(402).json({ error: err.message });
+        }
         res.status(500).json({ error: "AI Failed" });
     }
 });
 
-/**
- * 5. PING (Plugin Polling)
- */
 app.post('/project/ping', async (req, res) => {
-    const { projectId } = req.body;
+    const { projectId, userId } = req.body;
+    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) return res.status(403).json({ error: "Unauthorized" });
+
     const command = await StateManager.popCommand(projectId);
     
     if (command) {
@@ -367,48 +640,65 @@ app.post('/project/ping', async (req, res) => {
     }
 });
 
-/**
- * 6. MANUAL OVERRIDE
- */
 app.post('/human/override', validateRequest, async (req, res) => {
-    const { projectId, instruction, pruneHistory } = req.body;
-    const project = await StateManager.getProject(projectId);
-    
-    const overrideMsg = `[SYSTEM OVERRIDE]: ${instruction}`;
-    
-    if (pruneHistory && project.workerHistory.length >= 2) {
-        project.workerHistory.pop();
-        project.workerHistory.pop();
-    }
+    const { projectId, instruction, pruneHistory, userId } = req.body;
+    let basicUsage = 0;
 
-    const response = await AIEngine.callWorker(project.workerHistory, overrideMsg, []);
-    project.workerHistory.push({ role: 'user', parts: [{ text: overrideMsg }] });
-    project.workerHistory.push({ role: 'model', parts: [{ text: response }] });
+    try {
+        await checkMinimumCredits(userId, 'basic');
 
-    await StateManager.updateProject(projectId, { workerHistory: project.workerHistory });
-    await processAndQueueResponse(projectId, response);
-    res.json({ success: true });
-});
+        const project = await StateManager.getProject(projectId);
+        const overrideMsg = `[SYSTEM OVERRIDE]: ${instruction}`;
 
+        if (pruneHistory && project.workerHistory.length >= 2) {
+            project.workerHistory.pop();
+            project.workerHistory.pop();
+        }
 
-// --- RESPONSE PROCESSOR ---
-async function processAndQueueResponse(projectId, rawResponse) {
-    // 1. Intercept Image Requests
+        // Worker -> Basic
+        const workerResult = await AIEngine.callWorker(project.workerHistory, overrideMsg, []);
+        basicUsage += (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 (basicUsage > 0) await deductUserCredits(userId, basicUsage, 'basic');
+        
+        res.json({ success: true });
+    } catch (err) {
+        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 (Basic Credits)
+async function processAndQueueResponse(projectId, rawResponse, userId) {
     const imgPrompt = extractImagePrompt(rawResponse);
     if (imgPrompt) {
         console.log(`[${projectId}] 🎨 Generating Asset: ${imgPrompt}`);
-        const base64Image = await AIEngine.generateImage(imgPrompt);
         
-        if (base64Image) {
-            // Queue asset creation for Plugin
-            await StateManager.queueCommand(projectId, {
-                type: "CREATE_ASSET", 
-                payload: base64Image
-            });
+        // 1. Generate Image (returns { image, usage } or null)
+        const imgResult = await AIEngine.generateImage(imgPrompt);
+        
+        if (imgResult && imgResult.image) {
+            // 2. Bill for image tokens immediately (Basic)
+            const imgTokens = imgResult.usage?.totalTokenCount || 0;
+            const totalCost = imgTokens; // imgTokens + IMAGE_COST_BASIC;
+            
+            if (userId && totalCost > 0) {
+                await deductUserCredits(userId, totalCost, 'basic');
+            }
+
+            // 3. Queue creation command
+            await StateManager.queueCommand(projectId, { type: "CREATE_ASSET", payload: imgResult.image });
         }
     }
-
-    // 2. Queue the Raw Response (StateManager parses code)
+    // Queue the raw text response for the frontend
     await StateManager.queueCommand(projectId, rawResponse);
 }