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 { AutoDiscovery } = require('./auto-discover');
const BridgeClient = require('./bridge-client');

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);
            // 스크린캠스트도 중지
            const session = sessionManager.getSession(state.activeSessionId);
            if (session) session.client.stopScreencast();
        }
        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;
        }

        // ─── Phase 2: Screencast & Remote Input ────────────

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

            // 스크린캠스트 프레임 콜백 등록
            session.client.onScreencastFrame = (frame) => {
                if (ws.readyState === WebSocket.OPEN) {
                    ws.send(JSON.stringify({
                        type: 'screencast_frame',
                        sessionId: session.id,
                        data: frame.data,
                        metadata: frame.metadata,
                    }));
                }
            };

            const ok = await session.client.startScreencast(msg.options || {});
            ws.send(JSON.stringify({
                type: 'screencast_started',
                sessionId: session.id,
                success: ok,
            }));
            break;
        }

        case 'stop_screencast': {
            const session = sessionManager.getSession(msg.sessionId || state.activeSessionId);
            if (session) {
                await session.client.stopScreencast();
            }
            ws.send(JSON.stringify({
                type: 'screencast_stopped',
                sessionId: msg.sessionId || state.activeSessionId,
            }));
            break;
        }

        case 'input_event': {
            const session = sessionManager.getSession(msg.sessionId || state.activeSessionId);
            if (session) {
                await session.client.dispatchInput(msg.event);
            }
            break;
        }

        case 'click_action': {
            // 채팅 탭에서 Proceed/Cancel 등 버튼 클릭
            const session = sessionManager.getSession(msg.sessionId || state.activeSessionId);
            if (!session) {
                ws.send(JSON.stringify({ type: 'error', message: '활성 세션 없음' }));
                return;
            }

            const { x, y } = msg;
            try {
                // mousePressed + mouseReleased = 클릭
                await session.client.dispatchInput({
                    type: 'mouse', action: 'mousePressed', x, y, button: 'left', clickCount: 1,
                });
                await session.client.dispatchInput({
                    type: 'mouse', action: 'mouseReleased', x, y, button: 'left', clickCount: 1,
                });
                ws.send(JSON.stringify({ type: 'action_clicked', success: true, label: msg.label }));
            } catch (err) {
                ws.send(JSON.stringify({ type: 'error', message: `버튼 클릭 실패: ${err.message}` }));
            }
            break;
        }

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

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

// ─── Bridge API 프록시 ──────────────────────────────────

const bridge = new BridgeClient();

app.get('/api/bridge/health', async (req, res) => {
    try {
        const health = await bridge.checkHealth();
        res.json(health);
    } catch (e) {
        res.json({ available: false, error: e.message });
    }
});

app.get('/api/bridge/sessions', async (req, res) => {
    try {
        const data = await bridge.getSessions();
        res.json(data);
    } catch (e) {
        res.status(502).json({ error: 'Bridge 연결 실패: ' + e.message });
    }
});

app.get('/api/bridge/cascades', async (req, res) => {
    try {
        const data = await bridge.getCascades();
        res.json(data);
    } catch (e) {
        res.status(502).json({ error: 'Bridge 연결 실패: ' + e.message });
    }
});

app.post('/api/bridge/send', async (req, res) => {
    try {
        const { message, sessionId } = req.body;
        const result = await bridge.sendMessage(message, sessionId);
        res.json(result);
    } catch (e) {
        res.status(502).json({ error: e.message });
    }
});

app.post('/api/bridge/accept', async (req, res) => {
    try {
        const result = await bridge.acceptStep();
        res.json(result);
    } catch (e) {
        res.status(502).json({ error: e.message });
    }
});

app.post('/api/bridge/reject', async (req, res) => {
    try {
        const result = await bridge.rejectStep();
        res.json(result);
    } catch (e) {
        res.status(502).json({ error: e.message });
    }
});

app.post('/api/bridge/accept-terminal', async (req, res) => {
    try {
        const result = await bridge.acceptTerminal();
        res.json(result);
    } catch (e) {
        res.status(502).json({ error: e.message });
    }
});

// Bridge WS 이벤트 → 프론트엔드 포워딩
bridge.connectWs((msg) => {
    broadcastToAll({ type: 'bridge_event', ...msg });
});

// ─── CDP 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));

    // Auto Discovery 시작
    const discovery = new AutoDiscovery({
        host: 'localhost',
        intervalMs: 15000,
        onDiscovered: async (instance) => {
            // 이미 같은 포트의 세션이 있는지 확인
            const existing = sessionManager.listSessions().find(
                s => s.host === 'localhost' && s.cdpPort === instance.port
            );
            if (existing) return;

            // 자동 세션 생성
            try {
                const session = await sessionManager.addSession(
                    `Antigravity :${instance.port}`,
                    'localhost',
                    instance.port
                );

                console.log(`[AutoDiscovery] 자동 세션 생성: ${session.name} (${session.id})`);

                // 모든 WS 클라이언트에 세션 목록 업데이트
                broadcastToAll({
                    type: 'sessions_list',
                    sessions: sessionManager.listSessions(),
                });

                // 자동 발견 알림
                broadcastToAll({
                    type: 'auto_discovered',
                    session: {
                        id: session.id,
                        name: session.name,
                        host: 'localhost',
                        port: instance.port,
                        status: session.status,
                    },
                });
            } catch (err) {
                console.error(`[AutoDiscovery] 세션 생성 실패:`, err.message);
            }
        },
        onLost: (port) => {
            // 해당 포트의 세션 상태 업데이트
            const sessions = sessionManager.listSessions();
            const lost = sessions.find(s => s.cdpPort === port && s.host === 'localhost');
            if (lost) {
                console.log(`[AutoDiscovery] 세션 소실: ${lost.name}`);
                broadcastToAll({
                    type: 'session_lost',
                    sessionId: lost.id,
                    port,
                });
            }
        },
    });
    discovery.start();
});

/**
 * 모든 WS 클라이언트에 메시지 브로드캐스트
 */
function broadcastToAll(data) {
    const json = JSON.stringify(data);
    for (const [ws] of wsClients) {
        if (ws.readyState === WebSocket.OPEN) {
            ws.send(json);
        }
    }
}

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