everydaycats's picture
Update app.js
bae29e3 verified
import express from 'express';
import bodyParser from 'body-parser';
import cors from 'cors';
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 ---
let db = null;
let firestore = null;
let storage = null;
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) {
admin.initializeApp({
credential: admin.credential.cert(serviceAccount),
databaseURL: "https://shago-web-default-rtdb.firebaseio.com",
storageBucket: bucketName
});
}
db = admin.database();
firestore = admin.firestore();
storage = admin.storage();
initDB(db);
console.log("🔥 Firebase Connected (RTDB, Firestore, Storage) & StateManager Linked");
} else {
console.warn("⚠️ Memory-Only mode.");
}
} catch (e) { console.error("Firebase Init Error:", e); }
const app = express();
const PORT = process.env.PORT || 7860;
// ---
/*
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' }));
// --- 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('/admin/cleanup')) return next();
const { userId, projectId } = req.body;
// 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();
};
function extractWorkerPrompt(text) {
const match = text.match(/WORKER_PROMPT:\s*(.*)/s);
return match ? match[1].trim() : null;
}
function formatContext({ hierarchyContext, scriptContext, logContext }) {
let out = "";
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;
}
function extractPMQuestion(text) {
const match = text.match(/\[ASK_PM:\s*(.*?)\]/s);
return match ? match[1].trim() : null;
}
function extractImagePrompt(text) {
const match = text.match(/\[GENERATE_IMAGE:\s*(.*?)\]/s);
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 ---
app.post('/onboarding/analyze', validateRequest, async (req, res) => {
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."
});
}
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" });
}
});
app.post('/onboarding/create', validateRequest, async (req, res) => {
const { userId, description, answers } = req.body;
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}...`);
// 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] ✅ Passed. Generating Thumbnail...`);
const imagePrompt = `Game Title: ${gradingResult.title}. Core Concept: ${description}`;
// 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 timestamp = Date.now();
const status = isFailure ? "rejected" : "Idle";
// Save Data Logic
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);
}
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: 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 ---
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 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: 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 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: taskText }] });
// 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, []);
basicUsage += (workerResult.usage?.totalTokenCount || 0);
const workerText = workerResult.text;
workerHistory.push({ role: 'user', parts: [{ text: initialWorkerPrompt }] });
workerHistory.push({ role: 'model', parts: [{ text: workerText }] });
// 4. Update Memory
await StateManager.updateProject(projectId, {
userId,
pmHistory,
workerHistory,
gdd: gddText,
failureCount: 0
});
// 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(`[Background] Init Error for ${projectId}:`, err);
if(db) await db.ref(`projects/${projectId}/info/status`).set("error");
}
}
app.post('/new/project', validateRequest, async (req, res) => {
const { userId, projectId, description } = req.body;
try {
// 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");
res.json({
success: true,
message: "Project initialization started in background."
});
runBackgroundInitialization(projectId, userId, description);
} catch (err) {
if (err.message.includes("Insufficient")) {
return res.status(402).json({ error: err.message });
}
res.status(500).json({ error: "Failed to start project" });
}
});
app.post('/project/feedback', async (req, res) => {
const { userId, projectId, prompt, hierarchyContext, scriptContext, logContext, taskComplete, images } = req.body;
let 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" });
// 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.`;
// PM Call -> Diamond
const pmResult = await AIEngine.callPM(project.pmHistory, summary);
diamondUsage += (pmResult.usage?.totalTokenCount || 0);
const pmText = pmResult.text;
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}`;
// 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);
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);
// 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." });
}
}
// Standard Interaction (Worker Only = Basic)
const fullInput = `USER: ${prompt || "Automatic Feedback"}` + formatContext({ hierarchyContext, scriptContext, logContext });
let workerResult = await AIEngine.callWorker(project.workerHistory, fullInput, images || []);
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);
diamondUsage += (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: responseText }] });
// 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: responseText }] });
} else {
project.workerHistory.push({ role: 'user', parts: [{ text: fullInput }] });
project.workerHistory.push({ role: 'model', parts: [{ text: responseText }] });
}
await StateManager.updateProject(projectId, {
workerHistory: project.workerHistory,
pmHistory: project.pmHistory,
failureCount: project.failureCount
});
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" });
}
});
app.post('/project/ping', async (req, res) => {
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) {
if (command.payload === "CLEAR_CONSOLE") {
res.json({ action: "CLEAR_LOGS" });
} else {
res.json({
action: command.type,
target: command.payload,
code: command.type === 'EXECUTE' ? command.payload : null
});
}
} else {
res.json({ action: "IDLE" });
}
});
app.post('/human/override', validateRequest, async (req, res) => {
const { projectId, instruction, pruneHistory, userId } = req.body;
let basicUsage = 0;
try {
await checkMinimumCredits(userId, 'basic');
const project = await StateManager.getProject(projectId);
const overrideMsg = `[SYSTEM OVERRIDE]: ${instruction}`;
if (pruneHistory && project.workerHistory.length >= 2) {
project.workerHistory.pop();
project.workerHistory.pop();
}
// 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}`);
// 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 });
}
}
// Queue the raw text response for the frontend
await StateManager.queueCommand(projectId, rawResponse);
}
app.listen(PORT, () => {
console.log(`AI Backend Running on ${PORT}`);
});