hf_server.js

/**
 * Copyright (c) 2023 MERCENARIES.AI PTE. LTD.
 * All rights reserved.
 */
//@ts-check
const VERSION = '0.6.0.hf.018.a';

const express = require('express');
const http = require('http');
const session = require('express-session');
const axios = require('axios');
const { spawn } = require('child_process');
const { createProxyMiddleware } = require('http-proxy-middleware');
const fs = require('fs');
const path = require('path');
const chokidar = require('chokidar');
const fsExtra = require('fs-extra');
const app = express();

const OMNITOOL_INSTALL_SCRIPT = './omnitool_init.sh'; // './omnitool_start.sh';
const CONTAINER_HOST = '127.0.0.1';
const OMNI_URL = 'http://127.0.0.1:1688'; // URL of the OMNITOOL service
const PROXY_PORT_OMNITOOL = 4444;
const CONTAINER_PORT_OMNITOOL = 1688;
const OMNITOOL_HEALTH_PING = '/api/v1/mercenaries/ping';
const FAVICON = 'https://github.com/omnitool-ai/omnitool/raw/main/packages/omni-ui/omni-web/public/favicon.jpg'
const OMNITOOL_SPACE_URL = "omnitool-ai-omnitool-on-hf.hf.space";
const HF_SPACE_URL = "https://huggingface.co/spaces/omnitool-ai/omnitool_on_hf";
const HF_SPACE_DUPLICATE_URL = "https://huggingface.co/spaces/omnitool-ai/omnitool_on_hf?duplicate=true";

const sourceDir = 'omnitool/packages/omni-server/data.local';
const targetDir = '/data';//'omnitool/data';

const DELAY_OMNITOOL_SET_TO_RUNNING = 2000; // 2 seconds
const CHECK_OMNI_INTERVAL = 60000; // 1 minute
const MAX_LOG_SIZE = 1000; // Set your desired maximum log size

// Global variable
global.OMNITOOL_RUNNING = false;
global.OMNITOOL_READY = false;
global.ALREADY_STARTING = false;
global.LOCAL_URL = "";
global.PING_URL = "";
global.PROTOCOL = "";
global.PROXY_STARTED = false;
global.logs = [];
global.error_logs = [];
global.OMNITOOL_PROXY = null;
global.CONNECTED_TO_MASTER = false;



console.log(`************ Omnitool Proxy Server v${VERSION} ************`);

// HTML templates and constants
const LOG_CONTAINER_HTML = `<div id="logContainer"></div><div id="errorContainer"></div>`;
const COMMON_STYLES =
    `<style>
        #logContainer {
            height: 400px;
            overflow-y: scroll;
            background-color: black;
            color: lime;
            font-family: 'Courier New', Courier, monospace;
            padding: 10px;
            white-space: pre-wrap;
            border: 1px solid #ddd;
        }
        #errorContainer {
            height: 400px;
            overflow-y: scroll;
            background-color: red;
            color: white;
            font-family: 'Courier New', Courier, monospace;
            padding: 10px;
            white-space: pre-wrap;
            border: 1px solid #ddd;
        }
        .highlight-button {
            animation: pulseAnimation 1s infinite;
            background-color: yellow;
            color: black;
            font-weight: bold;
        }
        .highlight-button-green {
            animation: pulseAnimation 1s infinite;
            background-color: green;
            color: white;
            font-weight: bold;
        }      
        .button-like-link {
            display: inline-block;
            background-color: #f0f0f0; /* Button color */
            color: #000; /* Text color */
            padding: 5px 10px;
            text-align: center;
            text-decoration: none;
            border: 1px solid #000;
            border-radius: 5px;
            font-size: 16px;
            cursor: pointer;
        }
        .button-like-link:hover {
            background-color: #e9e9e9;
        }                     
        @keyframes pulseAnimation {
            0% { transform: scale(1); }
            50% { transform: scale(1.1); }
            100% { transform: scale(1); }
        }
    </style>`;

    function getScriptsHtml(url)
    {

        const scripts_html = 
            `<script>
                // Script to check Hugging Face login status and control 'Clone Repo' button

                if (document.getElementById('startServerButton'))
                {
                    // Call startServer after 1 second
                    setTimeout(startServer, 1000);
                }

                function startServer() {
                    document.getElementById('startServerButton').classList.remove
                    ('highlight-button-green');
                    document.getElementById('startServerButton').disabled = true;


                    fetch('/start-omnitool-server')
                        .then(response => response.json())
                        .then(data => startLogPolling());

                    
                }
                
                function exitIframe() {
                    if (window.self !== window.top) 
                    {
                        // open a new window to get out of the iframe, then wait 1s,  then fetch burst-iframe then reload
                        window.open("${global.PROTOCOL}://${global.LOCAL_URL}/?isVisited=true", '_blank');
                    }
                    else
                    {
                        fetch('/burst-iframe').then(response => response.json()).then(data => {window.location.reload();});
                    }
                }

                function refreshPage() {
                    window.location.reload();
                }

                function startLogPolling() {

                    const interval = setInterval(() => {
                        fetch('/omnitool-logs')
                            .then(response => response.json())
                            .then(data => 
                            {
                                const logContainer = document.getElementById('logContainer');
                                logContainer.innerText = data.logs.join("");
                                const errorContainer = document.getElementById('errorContainer');
                                errorContainer.innerText = data.error_logs.join("");

                                if (data.ready)
                                {
                                    if (document.getElementById('refreshButton').disabled == true) 
                                    {
                                        document.getElementById('refreshButton').disabled = false;

                                        scrollToBottom(logContainer);
                                        document.getElementById('refreshButton').classList.add('highlight-button');
                                        setTimeout(() => {
                                            window.location.reload(); // Refresh the page after 2 seconds
                                        }, 3000);
                                    }
                                }
                                else
                                {
                                    if (document.getElementById('refreshButton').disabled == false) 
                                    {
                                        document.getElementById('refreshButton').disabled = true;
                                        document.getElementById('refreshButton').classList.remove('highlight-button'); 
                                        setTimeout(() => {
                                            window.location.reload(); // Refresh the page after 2 seconds
                                        }, 3000);  
                                    }
                                    else
                                    {
                                        scrollToBottom(logContainer);
                                    }
                                }
                            });
                    }, 2000);

                }

                function scrollToBottom(element) {
                    element.scrollTop = element.scrollHeight;
                }
                
                startLogPolling();
            </script>`;

        console.log(`scripts_html = ${scripts_html}`);
        return scripts_html;
    }

const INTRO_MESSAGE  =                 
    `<h1>Omnitool.ai on Hugging Face</h1>
    <p>Version: ${VERSION}</p>
    <div>Welcome to Omnitool.ai, running on Hugging Face!</div>
    <div>For more details, including how to install omnitool on your own computer for free,</div>
    <div>visit us at <a href="https://omnitool.ai" target="_blank" class="button-like-link" id="duplicateRepoButton">omnitool.ai</a></div>
    <p></p>
    <p></p>`;

const VIDEO_PLAYER_MESSAGE = 
    '<div class="framer-epi0t3" data-framer-name="Video" name="Video"><div class="framer-ggzo55-container" data-framer-appear-id="vblmde" id="ZjfrOcqXP" style="opacity: 1; transform: perspective(1200px);"><div style="background:rgba(0, 0, 0, 0);height:100%;width:100%;border-radius:20px;position:relative;display:flex;justify-content:center;align-items:center;overflow:hidden" height="100%" id="ZjfrOcqXP" width="100%"><div class="framer-youtube" data-youtube-initialized="true" style="width: 100%; height: 100%;"><div style="padding:56.25% 0 0 0;position:relative;"><iframe src="https://www.youtube.com/embed/TlyaG2PKdrc" frameborder="0" allow="autoplay; fullscreen; picture-in-picture" style="position:absolute;top:0;left:0;width:100%;height:100%;" title="OMITOOL-DEMO" data-ready="true"></iframe></div></div></div></div></div>'
    
            
async function startOmnitoolServer()
{
    if (global.ALREADY_STARTING) return;
    global.ALREADY_STARTING = true;

    console.log('Starting Omnitool Server...');
    return new Promise((resolve, reject) =>
    {
        const omnitoolStartProcess = spawn(OMNITOOL_INSTALL_SCRIPT);
        omnitoolStartProcess.stdout.on('data', (data) =>
        {
            if (global.logs.length >= MAX_LOG_SIZE) {
                global.logs.shift(); // Remove the oldest log entry
            }
            global.logs.push(data.toString());
            console.log(`[log] ${data}`);
        
            if (!global.OMNITOOL_RUNNING)
            {
                if (data.toString().includes(`Server has started and is ready to accept connections`))
                {
                    console.log('Omnitool server started successfully');
                    setOmnitoolRunning(true);
                }
            }
        });

        omnitoolStartProcess.stderr.on('data', (data) =>
        {
            if (global.error_logs.length >= MAX_LOG_SIZE) {
                global.error_logs.shift(); // Remove the oldest log entry
            }

            console.error(`[stderr] ${data}`);
            global.error_logs.push(data.toString());
        });

        omnitoolStartProcess.on('close', (code) =>
        {
            const message = `Omnitool server process exited with code: ${code}`;
            console.log(message);
            global.logs.push(message);

            global.ALREADY_STARTING = false;
            if (code === 0)
            {
                //@ts-ignore
                resolve();
            } else
            {
                reject(`Omnitool server did not start properly.`);
            }
        });
    });
}

function setOmnitoolRunning(set_running)
{
    if (set_running === true)
    {
        if (global.OMNITOOL_READY === false) 
        {
            global.OMNITOOL_READY = true;
            console.log('Omnitool server is READY! ');
            setTimeout(() =>
            {
                console.log('Omnitool server is RUNNING! ');
                global.OMNITOOL_RUNNING = true;
            }, DELAY_OMNITOOL_SET_TO_RUNNING); // Delay by 2 second

            // start mirroring to /data when the server is running
            // startMirrorToDataDir(); #DISABLED!

        }

    }
    else
    {
        global.OMNITOOL_READY = false;
        global.OMNITOOL_RUNNING = false;
    }
}

// Function to check the status of the external service
async function checkOmnitoolStatus()
{
    try
    {
        //@ts-ignore
        const response = await axios.get('http://127.0.0.1:1688/api/v1/mercenaries/ping');
        if (response.data && response.data.ping === 'pong' && Object.keys(response.data.payload).length === 0)
        {
            setOmnitoolRunning(true);
        }
        else
        {
            setOmnitoolRunning(false);
        }
    } catch (error)
    {
        setOmnitoolRunning(false);
    }
}

async function handleGetStartOmnitoolServer(req, res)
{
    console.log(`Omnitool Server:ALREADY_STARTING =  ${global.ALREADY_STARTING}`);
    if (global.ALREADY_STARTING) 
    {
        res.writeHead(200, { 'Content-Type': 'text/html' });
        res.end("Omnitool server already starting");
        return;
    }
    try
    {
        await startOmnitoolServer();
        res.writeHead(200, { 'Content-Type': 'text/html' });
        res.end("Omnitool server started successfully");
    }
    catch (error)
    {
        console.error(error);
        global.ALREADY_STARTING = false;
        res.writeHead(500, { 'Content-Type': 'text/html' });
        res.end(`Error starting Omnitool server: ${error}`);
    }
}

async function handleGetOmnitoolLogs(req, res)
{
    res.writeHead(200, { 'Content-Type': 'application/json' });
    const reply = { logs: global.logs, error_logs: global.error_logs, ready: global.OMNITOOL_READY };
    res.end(JSON.stringify(reply));
    return;
}

async function handleGetBurstIframe(req, res)
{
    req.session.isVisited = true;

    const reply = { burst_iframe: true };
    res.writeHead(200, { 'Content-Type': 'application/json' });
    res.end(JSON.stringify(reply));

    return;
}

async function proxyRequest(req, res)
{

    console.log('Proxying request');
    if (global.PROXY_STARTED) return;
    global.PROXY_STARTED = true;

    // Proxy logic...
    const options = { hostname: CONTAINER_HOST, port: CONTAINER_PORT_OMNITOOL, path: req.url, method: req.method, headers: req.headers, };
    const proxy = http.request(options, (proxyRes) =>
    {
        res.writeHead(proxyRes.statusCode, proxyRes.headers);
        proxyRes.pipe(res, { end: true });
    });
    req.pipe(proxy, { end: true });
    return;
}

async function handleGetRoot(req, res)
{
    setGlobals(req);
    if (req.query.isVisited) req.session.isVisited = true;

    console.log(`req.session.isVisited = ${req.session.isVisited}`);
    console.log(`global.OMNITOOL_RUNNING = ${global.OMNITOOL_RUNNING}`);

    if (!req.session.isVisited || !global.OMNITOOL_RUNNING)
    {
        await checkOmnitoolStatus();
    }

    if (global.CONNECTED_TO_MASTER && !req.session.roayl_override)
    {
        let page_message ="";
       
        page_message += '<div style="color: red;">This is the REFERENCE Omnitool Space on Huggingface.</div>';
        page_message += '<p></p>';
        page_message += '<div>To access Omnitool, you must first DUPLICATE this SPACE to make it your own.</div>';
        page_message += '<div>Duplicating the Space will help secure your keys, recipes and outputs.</div>';
        page_message += '<div>Please set your Space as PRIVATE!</div>';
        page_message += '<div>DUPLICATING and using your own SPACE is FREE if you use the lowest tier of CPU. No GPU is required.</div>';
        page_message += `<div>DUPLICATING this space is done by pressing this button: <a href="${HF_SPACE_DUPLICATE_URL}" target="_blank" class="button-like-link" id="duplicateRepoButton">DUPLICATE SPACE</a></div>`;
        page_message += "<div>However, you can try using the 20 Gig Storage option (~ 5 USD/month) to persist keys, recipes and outputs between server restarts. This is still a BETA feature!</div>";
        page_message += '<p></p>';
        page_message += '<div style="color: red;">Some users have reported getting a 404 error when pressing the yellow START button after duplicating the space. If that happens to you follow the following instructions</div>';
        page_message += '<div> - go to settings and set the space to PUBLIC</div>';
        page_message += '<div> - confirm that the START button now works</div>';
        page_message += '<div> - go to settings and set the Space back to PRIVATE</div>';
        page_message += '<div> - in settings, FACTORY REBUILD your Space</div>';
        page_message += '<p></p>';
        
        const SCRIPTS = getScriptsHtml();

        const page_html = `
        <html>
            <head>
                <title>Omnitool</title>
                ${COMMON_STYLES}
            </head>
            <body>
                ${INTRO_MESSAGE}
                ${page_message}
                <p></p>
                ${VIDEO_PLAYER_MESSAGE}
                ${SCRIPTS}
            </body>
        </html>`; 

        res.writeHead(200, { 'Content-Type': 'text/html' });
        res.end(page_html);
        return;
    }
    

    if (!req.session.isVisited)
    {
        let page_message ="";
        page_message += `<div style="color: green;">You are now connected to the space [ ${global.LOCAL_URL} ]</div>`
        page_message += `<div style="color: red;">If this is NOT your SPACE, it is HIGHLY recommended that you duplicate this space and make it private to secure your keys, recipes and outputs.</div>`
        page_message += `<div>DUPLICATING and using your own SPACE is FREE. You can do it by pressing this button: <a href="${HF_SPACE_DUPLICATE_URL}" target="_blank" class="button-like-link" id="duplicateRepoButton">DUPLICATE SPACE</a></div><div>However, consider choosing the 20 Gig Storage option (~ 5 USD/month) to persist keys, recipes and outputs between server restarts.</div>`;
        page_message += `<p></p><p></p>`;
        page_message += `<div>To ACCESS OMNITOOL on this SPACE: press the [START] button below.</div>`;

        const gotoButtonClass = 'highlight-button';
        let buttons_html = `
            <button id="exitIframeButton" class="${gotoButtonClass}" onclick="exitIframe()">START</button>
            `;
        const SCRIPTS = getScriptsHtml();
        const page_html = `
            <html>
                <head>
                    <title>Omnitool</title>
                    ${COMMON_STYLES}
                </head>
                <body>
                    ${INTRO_MESSAGE}
                    ${page_message}
                    <p></p>
                    ${buttons_html}
                    <p></p>
                    ${LOG_CONTAINER_HTML}
                    ${SCRIPTS}
                </body>
            </html>
            `; 
        res.writeHead(200, { 'Content-Type': 'text/html' });
        res.end(page_html);
        return;
    }
    
    if (!global.OMNITOOL_RUNNING)
    {

        let page_message ="";
        page_message += `<div>It looks like Omnitool is not running on this Space. Let's fix that.</div>`;
        console.log('Omnitool server is not running');
        const startButtonClass = 'highlight-button-green';
        
        let buttonsHTML = `
            <button id="startServerButton" class="${startButtonClass}" onclick="startServer()">LAUNCHING...</button>
            <button id="refreshButton" class="${''}" disabled onclick="refreshPage()">REFRESHING...</button>
            `;

        const SCRIPTS = getScriptsHtml();
        const page_html = `
            <html>
                <head>
                    <title>Omnitool</title>
                    ${COMMON_STYLES}
                </head>
                <body>
                    ${INTRO_MESSAGE}
                    ${page_message}
                    <p></p>
                    ${buttonsHTML}
                    ${LOG_CONTAINER_HTML}
                    ${SCRIPTS}
                </body>
            </html>
            `; 

        res.writeHead(200, { 'Content-Type': 'text/html' });

     

        res.end(page_html);


        return;
    }

    proxyRequest(req, res);
}

function setGlobals(req)
{


    console.log(`Setting global using req.headers['host'] = ${req.headers['host']}`);
    global.LOCAL_URL = req.headers['host'];

    if (req.protocol === "https") global.PROTOCOL = 'https'; else global.PROTOCOL = 'http';
    console.log(`global.LOCAL_URL = ${global.LOCAL_URL}\nglobal.PROTOCOL = ${global.PROTOCOL}`);

    if (!global.LOCAL_URL) throw new Error('No host header found in request' + JSON.stringify(req.headers, null, 2));

    const hostname = global.LOCAL_URL.split(':')[0];
    const newUrl = `${global.PROTOCOL}://${hostname}`;//:${PROXY_PORT_OMNITOOL}`;
    global.PING_URL = `${newUrl}${OMNITOOL_HEALTH_PING}`;
    console.log(`global.PING_URL = ${global.PING_URL}`);

    // New logic to set CONNECTED_TO_MASTER
    global.CONNECTED_TO_MASTER = false;
    if (global.LOCAL_URL === OMNITOOL_SPACE_URL) global.CONNECTED_TO_MASTER = true;
    console.log(`global.CONNECTED_TO_MASTER = ${global.CONNECTED_TO_MASTER}`);

}

async function startHuggingfaceServer()
{
    // Start server
    http.createServer(app).listen(PROXY_PORT_OMNITOOL, () =>
    {
        console.log(`Server running on port ${PROXY_PORT_OMNITOOL}`);
    });

}

function omnitoolProxyMiddleware(req, res, next)
{
    if (global.OMNITOOL_RUNNING && req.session.isVisited)
    {
        // Proxy all requests to OMNITOOL when conditions are met
        return createProxyMiddleware({
            target: OMNI_URL,
            changeOrigin: true,
            ws: true // if you need WebSocket support
        })(req, res, next);
    } else
    {
        // Continue with normal processing for other cases
        next();
    }
}

async function main(app)
{
    if (fs.existsSync(targetDir)) {
        fsExtra.copySync(targetDir, sourceDir);
        console.log('Copied files from', targetDir, 'to', sourceDir);
    }
    
    await checkOmnitoolStatus();

    // Configure session middleware
    app.use(session({
        secret: 'your-secret-key',
        resave: false,
        saveUninitialized: true
    }));
    try
    {
        await startHuggingfaceServer();
    }
    catch (error)
    {
        console.error(`There was an error starting the server: ${error}`);
        return;
    }


    app.use(omnitoolProxyMiddleware);

    app.get('/favicon.ico', (req, res) => res.status(204));
    app.get('/start-omnitool-server', (req, res) => handleGetStartOmnitoolServer(req, res));
    app.get('/omnitool-logs', (req, res) => handleGetOmnitoolLogs(req, res));
    app.get('/burst-iframe', (req, res) => handleGetBurstIframe(req, res));
    app.get('/', (req, res) => handleGetRoot(req, res));

    setInterval(async () =>
    {
        await checkOmnitoolStatus();
    }, CHECK_OMNI_INTERVAL);

}

// Define the proxy middleware outside the function
global.OMNITOOL_PROXY = createProxyMiddleware({
    target: OMNI_URL,
    changeOrigin: true,
    ws: true // if you need WebSocket support
});

function omnitoolProxyMiddleware(req, res, next)
{
    if (global.OMNITOOL_RUNNING && req.session.isVisited)
    {
        // Use the predefined proxy middleware
        return global.OMNITOOL_PROXY(req, res, next);
    } else
    {
        // Continue with normal processing for other cases
        next();
    }
}
main(app);



function startMirrorToDataDir()
{
    if (!fs.existsSync(targetDir)) 
    {
        console.warn(`---------> Target directory ${targetDir} does not exist. Not starting mirroring.`);
        return;
    }
    else
    {
        console.warn(`---------> Target directory ${targetDir} exists. Starting mirroring.`);

    }

    const watcher = chokidar.watch(sourceDir, {
        ignored: /[\/\\]\./,
        persistent: true
    });

    watcher
        .on('add', function(file_path) {
            //if (file_path.includes('-journal') || file_path.includes('httpd.db')) return;
            if (file_path.includes('-journal')) return;

            const targetFile = file_path.replace(sourceDir, targetDir);
            //console.warn(`TargetFile = ${targetFile}. File ${file_path} added and copied from ${sourceDir} to ${targetDir}`);
            try {
                fsExtra.ensureDirSync(path.dirname(targetFile));
                fs.copyFile(file_path, targetFile, (err) => {
                if (err) 
                {
                    //throw err;
                    const error_message = `ERROR when doing: File ${file_path} added -> copied to ${targetFile}\n`;
                    console.log(error_message);
                    global.error_logs.push(error_message);
                }});
            }
            catch(error)
            {
                console.error(`WATCHER: error: ${error}`);
            }
        })
        .on('change', function(file_path) {
            //if (file_path.includes('-journal') || file_path.includes('httpd.db')) return;
            if (file_path.includes('-journal')) return;
            
            const targetFile = file_path.replace(sourceDir, targetDir);
            //console.warn(`TargetFile = ${targetFile}. File ${file_path} changed and copied from ${sourceDir} to ${targetDir}`);
            try{
                fsExtra.ensureDirSync(path.dirname(targetFile));
                fs.copyFile(file_path, targetFile, (err) => {
                if (err) 
                {
                    //throw err;
                    const error_message = `ERROR when doing: File ${file_path} changed -> copied to ${targetFile}\n`;
                    console.log(error_message);
                    global.error_logs.push(error_message);
                }
                });
            }
            catch(error)
            {
                console.error(`WATCHER: error: ${error}`);
            }
            
        })
        .on('unlink', function(file_path) {
            //if (file_path.includes('-journal') || file_path.includes('httpd.db')) return;
            if (file_path.includes('-journal')) return;

            const targetFile = file_path.replace(sourceDir, targetDir);
            //console.warn(`TargetFile = ${targetFile}. File ${file_path} removed from ${targetDir}`);
            try
            {
                fsExtra.ensureDirSync(path.dirname(targetFile));
                fs.unlink(targetFile, (err) => {
                if (err) 
                {
                    // throw err;
                    const error_message = `ERROR when doing: File ${file_path} has been removed from ${targetDir}\n`;
                    console.log(error_message);
                    global.error_logs.push(error_message);
                }
                });
            }
            catch(error)
            {
                console.error(`WATCHER: error: ${error}`);
            }
        });
        
    return watcher;
}