myNodeServer.js

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

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 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_NAME = 'omnitool-test-3';
const OMNITOOL_SPACE_OWNER = 'manu-sapiens';
const HF_SPACE_URL = `https://huggingface.co/spaces/${OMNITOOL_SPACE_OWNER}/${OMNITOOL_SPACE_NAME}`;
const HF_SPACE_DUPLICATE_URL = `https://huggingface.co/spaces/manu-sapiens/omnitool-test-3?duplicate=true`;

const DELAY_OMNITOOL_SET_TO_RUNNING = 2000; // 2 seconds
const CHECK_OMNI_INTERVAL = 60000; // 1 minute

// 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.OMNITOOL_PROXY = null;
global.CONNECTED_TO_MASTER = false;
global.SPACE_OWNER = "";


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

// HTML templates and constants
const LOG_CONTAINER_HTML = `<div id="logContainer"></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;
        }
        .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: black;
            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>`;


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) =>
        {
            global.logs.push(data.toString());
            console.log(`omnitool stdout: ${data}`);

            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) =>
        {
            console.error(`omnitool stderr: ${data}`);
            global.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
        }

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

// Function to check the status of the external service
async function checkOmnitoolStatus()
{
    console.log("Checking external service");
    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)
    {
        console.error('Cannot access OMNITOOL. It is probably not running...', error.message);
        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, ready: global.OMNITOOL_READY };
    res.end(JSON.stringify(reply));
    return;
}

async function handleGetBurstIframe(req, res)
{
    res.writeHead(200, { 'Content-Type': 'application/json' });
    req.session.isVisited = true;

    const reply = { burst_iframe: true };
    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;
}

function getButtonsString(buttonsHTML)
{
    // Message to display based on CONNECTED_TO_MASTER
    let connectionStatusMessage;
    if (global.CONNECTED_TO_MASTER)
    {
        connectionStatusMessage = '<div style="color: red;">You are now connected DIRECTLY to the REFERENCE Omnitool Space on Huggingface. It is HIGHLY recommended that you duplicate this space and make it private to secure your keys, recipes and outputs.</div>';
    } else
    {
        connectionStatusMessage = `<div style="color: green;">You are now connected to the space of ${global.SPACE_OWNER}. If this is NOT YOU, it is HIGHLY recommended that you duplicate this space and make it private to secure your keys, recipes and outputs.</div><div>In addition, consider subscribing to the 20GB local storage option from Huggingface to preserve your keys and work between reboot of this server.</div>`;
    }

    return `
        <html>
            <head>
                <title>Omnitool</title>
                ${COMMON_STYLES}
            </head>
            <body>
                <h1>Omnitool.ai on Hugging Face</h1>
                <p>Version: ${VERSION}</p>
                <div>Welcome to Omnitool.ai running on Hugging Face!</div>
                <div>This is a fully functional version.</div>
                <div>For more details visit us at <a href="https://omnitool.ai" target="_blank" class="button-like-link" id="duplicateRepoButton">omnitool.ai</a></div>
                <p>-----------------------------------</p>
                <p>This Space URL: ${global.LOCAL_URL}</p>
                <p>This Space Owner: ${global.SPACE_OWNER}</p>
                <p>Using REFERENCE Space: ${global.CONNECTED_TO_MASTER}</p>
                <p>Omnitool SERVER Launched: ${global.ALREADY_STARTING}</p>
                <p>Omnitool SERVER Ready: ${global.OMNITOOL_READY}</p>
                <p>-----------------------------------</p>
                <p></p>
                ${connectionStatusMessage}
                <p></p>
                ${buttonsHTML}
                ${LOG_CONTAINER_HTML}
                
                <script>
                    // Script to check Hugging Face login status and control 'Clone Repo' button
                
                    function startServer() {
                        document.getElementById('startServerButton').classList.remove
                        ('highlight-button');
                        document.getElementById('startServerButton').disabled = true;


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

                    function exitIframe() {
                        fetch('/burst-iframe')
                            .then(response => response.json())
                            .then(data => window.open('/', '_blank'));
                    }

                    function startLogPolling() {

                        const interval = setInterval(() => {
                            fetch('/omnitool-logs')
                                .then(response => response.json())
                                .then(data => {
                                    const logContainer = document.getElementById('logContainer');
                                    logContainer.innerText = data.logs.join("");
                                    if (data.ready) {
                                        clearInterval(interval);
                                        document.getElementById('exitIframeButton').classList.add('highlight-button');


                                    }
                                    else {
                                        scrollToBottom(logContainer);
                                    }
                                });
                        }, 500);
                    }

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


async function handleGetRoot(req, res)
{
    setGlobals(req);
    
    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.OMNITOOL_RUNNING)
    {
        console.log('Omnitool server is not running');
        let startButtonClass = '';
        if (!global.ALREADY_STARTING) startButtonClass = 'highlight-button';
        const gotoButtonClass = '';

        let buttonsHTML = `
        <button id="startServerButton" class="${startButtonClass}" onclick="startServer()">Start Omnitool Server</button>
        <button id="exitIframeButton" class="${gotoButtonClass}" onclick="exitIframe()">GOTO OMNITOOL</button>
        <a href="${HF_SPACE_DUPLICATE_URL}" target="_blank" class="button-like-link" id="duplicateRepoButton">DUPLICATE SPACE</a>`;

        const html = getButtonsString(buttonsHTML);
        res.writeHead(200, { 'Content-Type': 'text/html' });
        res.end(html);

        return;
    }

    if (!req.session.isVisited)
    {
        console.log('First time visitor');
        const gotoButtonClass = 'highlight-button-green';

        let buttonsHTML = `
        <button id="exitIframeButton" class="${gotoButtonClass}" onclick="exitIframe()">GOTO OMNITOOL</button>
        <a href="${HF_SPACE_DUPLICATE_URL}" target="_blank" class="button-like-link" id="duplicateRepoButton">DUPLICATE SPACE</a>`;
        const html = getButtonsString(buttonsHTML);
        res.writeHead(200, { 'Content-Type': 'text/html' });
        res.end(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'];

    // !!!!!!!!!
    // TEST ONLY: global.LOCAL_URL = 'manu-sapiens-omnitool-test-3.hf.space/';
    // !!!!!!!!!

    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;
    const recreated_string = `-${OMNITOOL_SPACE_NAME}.hf.space`;
    // remove the protocol, a possible trailing / and the recreated string from the url
    global.SPACE_OWNER = global.LOCAL_URL.replace(/\/$/, '').replace(recreated_string, '');
    console.log(`global.SPACE_OWNER = ${global.SPACE_OWNER}`);

    if (global.SPACE_OWNER === OMNITOOL_SPACE_OWNER) { global.CONNECTED_TO_MASTER = true; }
    console.log(`global.CONNECTED_TO_MASTER = ${global.CONNECTED_TO_MASTER}`);

}

async function startServer()
{
    // 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)
{
    await checkOmnitoolStatus();

    // Configure session middleware
    app.use(session({
        secret: 'your-secret-key',
        resave: false,
        saveUninitialized: true
    }));
    try
    {
        await startServer();
    }
    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);