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 MessageAccumulator = require('./message-accumulator');

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

app.use(express.json());
// 개발 모드: 캐시 비활성화
app.use((req, res, next) => {
    res.set('Cache-Control', 'no-store, no-cache, must-revalidate');
    res.set('Pragma', 'no-cache');
    res.set('Expires', '0');
    next();
});
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);
        }
    }
}

/**
 * 모든 WS 클라이언트에 메시지 브로드캐스트 (Bridge 이벤트 전달용)
 */
function broadcastToAll(msg) {
    const payload = JSON.stringify(msg);
    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 });
    }
});

// 단일 세션 cascades (실시간 갱신용 — 응답 크기 최소화)
app.get('/api/bridge/cascades/:sessionId', async (req, res) => {
    try {
        const data = await bridge.getCascades();
        const cascades = data.cascades || data;
        const session = cascades[req.params.sessionId];
        if (!session) {
            res.json({ status: 'not_found' });
        } else {
            res.json(session);
        }
    } 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.sendAction('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.sendAction('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.sendAction('acceptTerminal');
        res.json(result);
    } catch (e) {
        res.status(502).json({ error: e.message });
    }
});

app.get('/api/bridge/trajectory/:id', async (req, res) => {
    try {
        const data = await bridge.getTrajectory(req.params.id);
        res.json(data);
    } catch (e) {
        res.status(502).json({ error: e.message });
    }
});

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

// 승인/거절 통합 엔드포인트: 모든 유형을 순차 시도
app.post('/api/bridge/approve', async (req, res) => {
    const { type } = req.body; // 'accept' or 'reject'
    const actions = type === 'accept'
        ? ['acceptStep', 'acceptCommand', 'acceptTerminal']
        : ['rejectStep', 'rejectCommand', 'rejectTerminal'];

    for (const action of actions) {
        try {
            const result = await bridge.sendAction(action);
            console.log(`[Approve] ${action} 성공`);
            return res.json({ success: true, action, result });
        } catch (_) {
            // 이 타입은 안 맞음 → 다음 시도
        }
    }
    res.status(502).json({ error: '승인/거절 실패: 대기 중인 요청 없음' });
});
const accumulator = new MessageAccumulator();

// 메시지 누적 API
app.get('/api/bridge/messages/:sessionId', async (req, res) => {
    // 요청 시점에 즉시 최신 cascade 반영
    try {
        const cascData = await bridge.getCascades();
        const cascades = cascData.cascades || cascData;
        const sid = req.params.sessionId;
        if (cascades[sid]) {
            accumulator.processCascade(sid, cascades[sid]);
        }
    } catch (_) { }
    const messages = accumulator.getMessages(req.params.sessionId);
    const status = accumulator.getStatus(req.params.sessionId);
    res.json({ messages, ...status });
});

bridge.connectWs(async (msg) => {
    console.log(`[Bridge Event] ${msg.type} session=${msg.sessionId || 'N/A'}`);
    broadcastToAll({ type: 'bridge_event', ...msg });

    // step_changed/state_changed → cascade 스냅샷 diff → 새 메시지 push
    if (msg.type === 'step_changed' || msg.type === 'state_changed') {
        await pollAndPushCascades(msg.sessionId);
    }
});

// 주기적 cascade polling (WS 이벤트 누락 보완, 3초 간격)
setInterval(async () => {
    if (!bridge.connected) return;
    await pollAndPushCascades();
}, 3000);

async function pollAndPushCascades(specificSessionId) {
    try {
        const cascData = await bridge.getCascades();
        const cascades = cascData.cascades || cascData;

        for (const [sessionId, cascade] of Object.entries(cascades)) {
            if (specificSessionId && sessionId !== specificSessionId) continue;

            const result = accumulator.processCascade(sessionId, cascade);
            if (result && result.newMessages) {
                const status = accumulator._computeStatus(cascade);
                console.log(`[Accumulator] ${sessionId.substring(0, 8)}: +${result.newMessages.length} messages (${status.isRunning ? 'running' : status.isBlocking ? 'blocking' : 'idle'})`);
                broadcastToAll({
                    type: 'new_messages',
                    sessionId,
                    messages: result.newMessages,
                    ...status,
                });
            } else if (result) {
                // 메시지 변경 없어도 상태 변경은 push
                const status = accumulator._computeStatus(cascade);
                broadcastToAll({
                    type: 'new_messages',
                    sessionId,
                    messages: null,
                    ...status,
                });
            }
        }
    } catch (e) {
        // 연결 안 됨 — 무시
    }
}

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