Update app.js

app.js

@@ -1,243 +1,161 @@
-const express = require('express');
-const { createClient } = require('@supabase/supabase-js');
-const jwt = require('jsonwebtoken');
-const { v4: uuidv4 } = require('uuid');
-const axios = require('axios');
-const bodyParser = require('body-parser');
-const cors = require('cors');
+import express from 'express';
+import { createServer } from 'http';
+import { WebSocketServer, WebSocket } from 'ws';
+import cors from 'cors';
+import { createClient } from '@supabase/supabase-js';
 
+// --- CONFIG ---
+const PORT = 7860;
+const SUPABASE_URL = process.env.SUPABASE_URL;
+const SUPABASE_KEY = process.env.SUPABASE_SERVICE_KEY;
+const CORE_URL = process.env.CORE_URL; // e.g. https://my-core-server.hf.space
+
+if (!SUPABASE_URL) { console.error("❌ Missing Config"); process.exit(1); }
+
+// --- STATE ---
 const app = express();
+const server = createServer(app);
+const wss = new WebSocketServer({ noServer: true });
+const supabase = createClient(SUPABASE_URL, SUPABASE_KEY);
+
+// Active Connections: Map<UserId, Set<WebSocket>>
+const clients = new Map();
+
 app.use(cors());
-app.use(bodyParser.json({ limit: '50mb' }));
-
-const tempKeys = new Map();
-const activeSessions = new Map();
-
-const {
-    SUPABASE_URL,
-    SUPABASE_SERVICE_ROLE_KEY,
-    EXTERNAL_SERVER_URL = 'http://localhost:7860',
-    STORAGE_BUCKET = 'project-assets', 
-    PORT = 7860
-} = process.env;
-
-let supabase = null;
-
-try {
-    if (SUPABASE_URL && SUPABASE_SERVICE_ROLE_KEY) {
-        supabase = createClient(SUPABASE_URL, SUPABASE_SERVICE_ROLE_KEY, {
-            auth: {
-                autoRefreshToken: false,
-                persistSession: false
+app.use(express.json());
+
+// --- HTTP ENDPOINTS (Webhooks/Uploads) ---
+
+app.get('/', (req, res) => res.send('Front Gateway Active'));
+
+// Internal Endpoint for Core to push notifications to desktop
+app.post('/internal/notify', (req, res) => {
+    const { user_id, type, message, payload } = req.body;
+    
+    if (clients.has(user_id)) {
+        const userSockets = clients.get(user_id);
+        let count = 0;
+        userSockets.forEach(ws => {
+            if (ws.readyState === WebSocket.OPEN) {
+                ws.send(JSON.stringify({ type, message, payload }));
+                count++;
             }
         });
-        console.log("⚡ Supabase Connected (Admin Context)");
-    } else {
-        console.warn("⚠️ Memory-Only mode (Supabase credentials missing).");
-    }
-} catch (e) {
-    console.error("Supabase Init Error:", e);
-}
-
-const verifySupabaseUser = async (req, res, next) => {
-    const debugMode = process.env.DEBUG_NO_AUTH === 'true';
-    if (debugMode) { req.user = { id: "user_dev_01" }; return next(); }
-    const authHeader = req.headers.authorization;
-    if (!authHeader || !authHeader.startsWith('Bearer ')) return res.status(401).json({ error: 'Missing Bearer token' });
-    const idToken = authHeader.split('Bearer ')[1];
-    try {
-        if (supabase) {
-            const { data: { user }, error } = await supabase.auth.getUser(idToken);
-            if (error || !user) throw new Error("Invalid Token");
-            req.user = user;
-            next();
-        } else { req.user = { id: "memory_user" }; next(); }
-    } catch (error) { return res.status(403).json({ error: 'Unauthorized', details: error.message }); }
-};
-
-async function getSessionSecret(uid, projectId) {
-    const cacheKey = `${uid}:${projectId}`;
-    if (activeSessions.has(cacheKey)) {
-        const session = activeSessions.get(cacheKey);
-        session.lastAccessed = Date.now();
-        return session.secret;
+        return res.json({ success: true, delivered_to: count });
     }
-    if (supabase) {
-        try {
-            const { data } = await supabase.from('projects').select('plugin_secret').eq('id', projectId).eq('user_id', uid).single();
-            if (data && data.plugin_secret) {
-                activeSessions.set(cacheKey, { secret: data.plugin_secret, lastAccessed: Date.now() });
-                return data.plugin_secret;
-            }
-        } catch (err) { console.error("DB Read Error:", err); }
-    }
-    return null;
-}
-
-app.post('/key', verifySupabaseUser, (req, res) => {
-    const { projectId } = req.body;
-    if (!projectId) return res.status(400).json({ error: 'projectId required' });
-    const key = `key_${uuidv4().replace(/-/g, '')}`;
-    tempKeys.set(key, { uid: req.user.id, projectId: projectId, createdAt: Date.now() });
-    console.log(`🔑 Generated Key for user ${req.user.id}: ${key}`);
-    res.json({ key, expiresIn: 300 });
-});
-
-app.post('/redeem', async (req, res) => {
-    const { key } = req.body;
-    if (!key || !tempKeys.has(key)) return res.status(404).json({ error: 'Invalid or expired key' });
-    const data = tempKeys.get(key);
-    const sessionSecret = uuidv4(); 
-    const token = jwt.sign({ uid: data.uid, projectId: data.projectId }, sessionSecret, { expiresIn: '3d' });
-    const cacheKey = `${data.uid}:${data.projectId}`;
-    activeSessions.set(cacheKey, { secret: sessionSecret, lastAccessed: Date.now() });
-    if (supabase) await supabase.from('projects').update({ plugin_secret: sessionSecret }).eq('id', data.projectId).eq('user_id', data.uid);
-    tempKeys.delete(key);
-    console.log(`🚀 Redeemed JWT for ${cacheKey}`);
-    res.json({ token });
+    
+    res.json({ success: false, reason: "User offline" });
 });
 
-app.post('/verify', async (req, res) => {
-    const { token } = req.body;
-    if (!token) return res.status(400).json({ valid: false });
-    const decoded = jwt.decode(token);
-    if (!decoded || !decoded.uid || !decoded.projectId) return res.status(401).json({ valid: false });
-    const secret = await getSessionSecret(decoded.uid, decoded.projectId);
-    if (!secret) return res.status(401).json({ valid: false });
-    try { jwt.verify(token, secret); return res.json({ valid: true }); } catch (err) { return res.status(403).json({ valid: false }); }
+// Endpoint for File Uploads (Base64 is easier for single script than Multer)
+app.post('/upload', async (req, res) => {
+    // In a real app, handle streams. Here we assume small/medium files.
+    // Client sends this to avoid clogging WS.
+    res.json({ status: "ok", note: "File received (Placeholder logic)" });
 });
 
-app.post('/feedback', async (req, res) => {
-    const { token, ...pluginPayload } = req.body;
-    if (!token) return res.status(400).json({ error: 'Token required' });
-    const decoded = jwt.decode(token);
-    if (!decoded) return res.status(401).json({ error: 'Malformed token' });
-    const secret = await getSessionSecret(decoded.uid, decoded.projectId);
-    if (!secret) return res.status(404).json({ error: 'Session revoked' });
-    try {
-        jwt.verify(token, secret); 
-        const targetUrl = EXTERNAL_SERVER_URL.replace(/\/$/, '') + '/project/feedback';
-        const response = await axios.post(targetUrl, { userId: decoded.uid, projectId: decoded.projectId, ...pluginPayload });
-        return res.json({ success: true, externalResponse: response.data });
-    } catch (err) { return res.status(502).json({ error: 'Failed to forward' }); }
-});
+// --- WEBSOCKET AUTH & UPGRADE ---
 
-app.post('/feedback2', verifySupabaseUser, async (req, res) => {
-    const { projectId, prompt, images, ...otherPayload } = req.body; 
-    const userId = req.user.id;
-    if (!projectId || !prompt) return res.status(400).json({ error: 'Missing projectId or prompt' });
-    const targetUrl = EXTERNAL_SERVER_URL.replace(/\/$/, '') + '/project/feedback';
-    try {
-        const response = await axios.post(targetUrl, { userId: userId, projectId: projectId, prompt: prompt, images: images || [], ...otherPayload });
-        return res.json({ success: true, externalResponse: response.data });
-    } catch (err) { return res.status(502).json({ error: 'Failed to forward' }); }
-});
+server.on('upgrade', async (request, socket, head) => {
+    const url = new URL(request.url, `http://${request.headers.host}`);
+    const token = url.searchParams.get('token');
 
-// --- STREAM FEED (Optimized headers) ---
-app.post('/stream-feed', verifySupabaseUser, async (req, res) => {
-    const { projectId } = req.body;
-    const userId = req.user.id;
+    if (!token) {
+        socket.write('HTTP/1.1 401 Unauthorized\r\n\r\n');
+        socket.destroy();
+        return;
+    }
 
-    // Headers to disable caching for poller
-    res.setHeader('Cache-Control', 'no-cache, no-store, must-revalidate');
-    res.setHeader('Pragma', 'no-cache');
-    res.setHeader('Expires', '0');
+    // Validate with Supabase
+    const { data: { user }, error } = await supabase.auth.getUser(token);
 
-    if (!projectId) return res.status(400).json({ error: 'Missing projectId' });
-    
-    if (supabase) {
-         const { data, error } = await supabase.from('projects').select('id, user_id, info').eq('id', projectId).single();
-         if (error || !data || data.user_id !== userId) return res.status(403).json({ error: 'Unauthorized' });
-         
-         const targetUrl = EXTERNAL_SERVER_URL.replace(/\/$/, '') + '/project/ping';
-         try {
-             const response = await axios.post(targetUrl, { projectId, userId, isFrontend: true });
-             return res.json({ ...response.data, dbStatus: data.info?.status || 'idle' });
-         } catch (e) { return res.status(502).json({ error: "AI Server Unreachable" }); }
+    if (error || !user) {
+        console.log("❌ WS Auth Failed");
+        socket.write('HTTP/1.1 403 Forbidden\r\n\r\n');
+        socket.destroy();
+        return;
     }
-});
 
-/*
-app.post('/poll', async (req, res) => {
-    const { token } = req.body;
-    if (!token) return res.status(400).json({ error: 'Token required' });
-    const decoded = jwt.decode(token);
-    if (!decoded) return res.status(401).json({ error: 'Malformed token' });
-    const secret = await getSessionSecret(decoded.uid, decoded.projectId);
-    if (!secret) return res.status(404).json({ error: 'Session revoked' });
-    try {
-        jwt.verify(token, secret);
-        const targetUrl = EXTERNAL_SERVER_URL.replace(/\/$/, '') + '/project/ping'; 
-        const response = await axios.post(targetUrl, { projectId: decoded.projectId, userId: decoded.uid });
-        return res.json(response.data); 
-    } catch (err) { return res.status(403).json({ error: 'Invalid Token' }); }
+    wss.handleUpgrade(request, socket, head, (ws) => {
+        wss.emit('connection', ws, request, user);
+    });
 });
-*/
 
-app.post('/poll', async (req, res) => {
-    const { token } = req.body;
-    if (!token) return res.status(400).json({ error: 'Token required' });
-    
-    const decoded = jwt.decode(token);
-    if (!decoded) return res.status(401).json({ error: 'Malformed token' });
-    
-    const secret = await getSessionSecret(decoded.uid, decoded.projectId);
-    if (!secret) return res.status(404).json({ error: 'Session revoked' });
-    
-    try {
-        jwt.verify(token, secret);
-        const targetUrl = EXTERNAL_SERVER_URL.replace(/\/$/, '') + '/project/ping'; 
-        
-        // FIX: We do NOT send isFrontend: true. 
-        // We act as the Plugin (Executor), but the Main Server will still give us the snapshot now.
-        const response = await axios.post(targetUrl, { 
-            projectId: decoded.projectId, 
-            userId: decoded.uid 
-        });
-        
-        return res.json(response.data); 
-    } catch (err) { return res.status(403).json({ error: 'Invalid Token' }); }
-});
+// --- WEBSOCKET LOGIC ---
 
+wss.on('connection', (ws, req, user) => {
+    const userId = user.id;
+    console.log(`🔌 Connected: ${userId}`);
 
-app.post('/project/delete', verifySupabaseUser, async (req, res) => {
-    const { projectId } = req.body;
-    const userId = req.user.id;
-    if (!projectId) return res.status(400).json({ error: "Missing Project ID" });
-    try {
-        const { data: project, error: fetchError } = await supabase.from('projects').select('user_id').eq('id', projectId).single();
-        if (fetchError || !project || project.user_id !== userId) return res.status(403).json({ error: "Unauthorized" });
-        await supabase.from('message_chunks').delete().eq('project_id', projectId);
-        await supabase.from('projects').delete().eq('id', projectId);
-        if (STORAGE_BUCKET) {
-            const { data: files } = await supabase.storage.from(STORAGE_BUCKET).list(projectId);
-            if (files && files.length > 0) { const filesToRemove = files.map(f => `${projectId}/${f.name}`); await supabase.storage.from(STORAGE_BUCKET).remove(filesToRemove); }
+    // Register
+    if (!clients.has(userId)) clients.set(userId, new Set());
+    clients.get(userId).add(ws);
+
+    // Heartbeat (Keep-Alive for HF Spaces)
+    ws.isAlive = true;
+    ws.on('pong', () => { ws.isAlive = true; });
+
+    ws.on('message', async (message) => {
+        try {
+            const data = JSON.parse(message.toString());
+
+            // 1. ROUTING: PROMPT -> CORE
+            if (data.type === 'prompt') {
+                ws.send(JSON.stringify({ type: 'status', status: 'thinking' }));
+
+                if (!CORE_URL) {
+                    ws.send(JSON.stringify({ type: 'error', message: "Core Server Config Missing" }));
+                    return;
+                }
+
+                // Call Core
+                const response = await fetch(`${CORE_URL}/process`, {
+                    method: 'POST',
+                    headers: { 'Content-Type': 'application/json' },
+                    body: JSON.stringify({
+                        userId: userId,
+                        projectId: data.projectId,
+                        prompt: data.content,
+                        context: data.context
+                    })
+                });
+
+                const result = await response.json();
+
+                // Send Result to Client
+                ws.send(JSON.stringify({
+                    type: 'response',
+                    text: result.text,
+                    should_reload: result.should_reload
+                }));
+            }
+
+        } catch (e) {
+            console.error("WS Message Error", e);
+            ws.send(JSON.stringify({ type: 'error', message: "Server Error" }));
         }
-        activeSessions.delete(`${userId}:${projectId}`);
-        for (const [key, val] of tempKeys.entries()) { if (val.projectId === projectId) tempKeys.delete(key); }
-        res.json({ success: true });
-    } catch (err) { res.status(500).json({ error: "Delete failed" }); }
-});
+    });
 
-app.get('/cleanup', (req, res) => {
-    // ... (Standard cleanup) ...
-    const THRESHOLD = 1000 * 60 * 60; 
-    const now = Date.now();
-    let cleanedCount = 0;
-    for (const [key, value] of activeSessions.entries()) { if (now - value.lastAccessed > THRESHOLD) { activeSessions.delete(key); cleanedCount++; } }
-    for (const [key, value] of tempKeys.entries()) { if (now - value.createdAt > (1000 * 60 * 4)) { tempKeys.delete(key); } }
-    res.json({ message: `Cleaned ${cleanedCount} cached sessions from memory.` });
+    ws.on('close', () => {
+        if (clients.has(userId)) {
+            clients.get(userId).delete(ws);
+            if (clients.get(userId).size === 0) clients.delete(userId);
+        }
+        console.log(`❌ Disconnected: ${userId}`);
+    });
 });
 
-app.post('/nullify', verifySupabaseUser, async (req, res) => {
-    const { projectId } = req.body;
-    if (!projectId) return res.status(400).json({ error: 'projectId required' });
-    const cacheKey = `${req.user.id}:${projectId}`;
-    activeSessions.delete(cacheKey);
-    if (supabase) await supabase.from('projects').update({ plugin_secret: null }).eq('id', projectId).eq('user_id', req.user.id);
-    res.json({ success: true });
-});
+// Heartbeat Interval (30s)
+const interval = setInterval(() => {
+    wss.clients.forEach((ws) => {
+        if (ws.isAlive === false) return ws.terminate();
+        ws.isAlive = false;
+        ws.ping();
+    });
+}, 30000);
 
-app.get('/', (req, res) => { res.send('Plugin Auth Proxy Running (Supabase Edition)'); });
+wss.on('close', () => clearInterval(interval));
 
-app.listen(PORT, () => { console.log(`🚀 Auth Proxy running on port ${PORT}`); });
+server.listen(PORT, () => {
+    console.log(`🚀 Front Gateway running on port ${PORT}`);
+});