gravity_web/server/index.js

/**
 * Gravity Web Server — Express REST API + WebSocket
 * 
 * Antigravity IDE 여러 인스턴스를 CDP로 연결하고
 * 웹 브라우저에서 세션을 전환하며 채팅할 수 있게 합니다.
 */

const express = require('express');
const http = require('http');
const path = require('path');
const { WebSocketServer, WebSocket } = require('ws');
const SessionManager = require('./session-manager');

const PORT = process.env.PORT || 3300;
const app = express();
const server = http.createServer(app);

app.use(express.json());
app.use(express.static(path.join(__dirname, '..', 'public')));

const sessionManager = new SessionManager();

// ─── WebSocket Setup ────────────────────────────────────
const wss = new WebSocketServer({ server, path: '/ws' });

/** @type {Map<WebSocket, {activeSessionId: string|null}>} */
const wsClients = new Map();

wss.on('connection', (ws) => {
    console.log('[WS] 클라이언트 연결');
    wsClients.set(ws, { activeSessionId: null });

    ws.on('message', async (raw) => {
        try {
            const msg = JSON.parse(raw.toString());
            await handleWsMessage(ws, msg);
        } catch (err) {
            ws.send(JSON.stringify({ type: 'error', message: err.message }));
        }
    });

    ws.on('close', () => {
        const state = wsClients.get(ws);
        if (state?.activeSessionId) {
            stopSessionPolling(state.activeSessionId);
        }
        wsClients.delete(ws);
        console.log('[WS] 클라이언트 연결 해제');
    });

    // 초기 세션 목록 전송
    ws.send(JSON.stringify({
        type: 'sessions_list',
        sessions: sessionManager.listSessions(),
    }));
});

/**
 * WebSocket 메시지 핸들러
 */
async function handleWsMessage(ws, msg) {
    const state = wsClients.get(ws);

    switch (msg.type) {
        case 'switch_session': {
            // 기존 세션 폴링 중지
            if (state.activeSessionId) {
                stopSessionPolling(state.activeSessionId);
            }

            const session = sessionManager.getSession(msg.sessionId);
            if (!session) {
                ws.send(JSON.stringify({ type: 'error', message: '세션을 찾을 수 없습니다' }));
                return;
            }

            state.activeSessionId = msg.sessionId;

            // 이 세션의 CDP 폴링 시작
            startSessionPolling(session, ws);

            ws.send(JSON.stringify({
                type: 'session_switched',
                sessionId: msg.sessionId,
                session: {
                    id: session.id,
                    name: session.name,
                    status: session.status,
                    title: session.title,
                },
            }));
            break;
        }

        case 'send_message': {
            const session = sessionManager.getSession(msg.sessionId || state.activeSessionId);
            if (!session) {
                ws.send(JSON.stringify({ type: 'error', message: '활성 세션 없음' }));
                return;
            }

            const result = await session.client.sendMessage(msg.text);
            ws.send(JSON.stringify({
                type: 'message_sent',
                sessionId: session.id,
                ...result,
            }));
            break;
        }

        case 'get_screenshot': {
            const session = sessionManager.getSession(msg.sessionId || state.activeSessionId);
            if (!session) {
                ws.send(JSON.stringify({ type: 'error', message: '활성 세션 없음' }));
                return;
            }

            const data = await session.client.captureScreenshot();
            if (data) {
                ws.send(JSON.stringify({
                    type: 'screenshot',
                    sessionId: session.id,
                    data: data,
                }));
            }
            break;
        }

        default:
            ws.send(JSON.stringify({ type: 'error', message: `알 수 없는 메시지 타입: ${msg.type}` }));
    }
}

/**
 * 특정 세션의 CDP 채팅 폴링을 시작하고 결과를 ws 클라이언트에 전송
 */
function startSessionPolling(session, ws) {
    session.client.onChatUpdate = (html) => {
        if (ws.readyState === WebSocket.OPEN) {
            ws.send(JSON.stringify({
                type: 'chat_update',
                sessionId: session.id,
                html: html,
                timestamp: Date.now(),
            }));
        }
    };
    session.client.startPolling(1000);
}

/**
 * 세션 폴링 중지
 */
function stopSessionPolling(sessionId) {
    const session = sessionManager.getSession(sessionId);
    if (session) {
        session.client.stopPolling();
        session.client.onChatUpdate = null;
    }
}

/**
 * 세션 목록 변경을 모든 WS 클라이언트에 브로드캐스트
 */
function broadcastSessions() {
    const payload = JSON.stringify({
        type: 'sessions_list',
        sessions: sessionManager.listSessions(),
    });
    for (const [ws] of wsClients) {
        if (ws.readyState === WebSocket.OPEN) {
            ws.send(payload);
        }
    }
}

// ─── REST API ───────────────────────────────────────────

app.get('/api/sessions', (req, res) => {
    res.json(sessionManager.listSessions());
});

app.post('/api/sessions', async (req, res) => {
    const { name, host = 'localhost', cdpPort = 9000 } = req.body;

    if (!name) {
        return res.status(400).json({ error: 'name is required' });
    }

    const session = await sessionManager.addSession(name, host, cdpPort);
    broadcastSessions();

    res.status(201).json({
        id: session.id,
        name: session.name,
        host: session.host,
        cdpPort: session.cdpPort,
        status: session.status,
        title: session.title,
        error: session.error || null,
    });
});

app.delete('/api/sessions/:id', async (req, res) => {
    const removed = await sessionManager.removeSession(req.params.id);
    if (!removed) {
        return res.status(404).json({ error: 'Session not found' });
    }
    broadcastSessions();
    res.json({ success: true });
});

app.post('/api/sessions/:id/reconnect', async (req, res) => {
    const result = await sessionManager.reconnectSession(req.params.id);
    broadcastSessions();
    res.json(result);
});

app.get('/api/sessions/:id/status', (req, res) => {
    const session = sessionManager.getSession(req.params.id);
    if (!session) {
        return res.status(404).json({ error: 'Session not found' });
    }
    res.json({
        id: session.id,
        status: session.status,
        title: session.title,
    });
});

// Health
app.get('/api/health', (req, res) => {
    res.json({
        status: 'ok',
        uptime: process.uptime(),
        sessions: sessionManager.listSessions().length,
    });
});

// SPA 폴백
app.get('*', (req, res) => {
    res.sendFile(path.join(__dirname, '..', 'public', 'index.html'));
});

// ─── Startup ────────────────────────────────────────────

server.listen(PORT, () => {
    console.log('═'.repeat(50));
    console.log(`  Gravity Web 서버 시작`);
    console.log(`  http://localhost:${PORT}`);
    console.log('═'.repeat(50));
});

// Graceful shutdown
process.on('SIGINT', async () => {
    console.log('\n[Server] 종료 중...');
    await sessionManager.cleanup();
    server.close(() => {
        console.log('[Server] 종료 완료');
        process.exit(0);
    });
});