static/js/websocket.js

/* ===== WebSocket Communication ===== */
window.NAIWS = {
    connect() {
        const state = NAIState;
        const protocol = location.protocol === 'https:' ? 'wss:' : 'ws:';
        const wsUrl = `${protocol}//${location.host}/ws/${state.clientId}`;
        state.ws = new WebSocket(wsUrl);

        state.ws.onopen = () => {
            state.wsReady = true;
            console.log('WebSocket connected');
            this.updateStatus('已连接', false);
        };

        state.ws.onmessage = (event) => {
            const data = JSON.parse(event.data);
            this.handleMessage(data);
        };

        state.ws.onclose = () => {
            state.wsReady = false;
            console.log('WebSocket disconnected, reconnecting...');
            this.updateStatus('连接断开，正在重连...', false);
            setTimeout(() => this.connect(), 3000);
        };

        state.ws.onerror = (err) => {
            console.error('WebSocket error:', err);
        };
    },

    handleMessage(data) {
        switch (data.type) {
            case 'queued':
                this.updateStatus(`排队中，队列位置: ${data.position}`, true);
                break;
            case 'status':
                this.updateStatus(data.message, true);
                break;
            case 'result':
                NAIGenerate.onResult(data);
                break;
            case 'error':
                NAIGenerate.onError(data.message);
                break;
            case 'last_seed':
                if (data.seed !== -1) $('#seed-input').value = data.seed;
                break;
            case 'vibe_encode_result':
                NAIVibe.onEncodeResult(data);
                break;
            case 'vibe_encode_error':
                NAIVibe.onEncodeError(data.message);
                break;
            case 'augment_result':
                NAIDirectorTools.onResult(data);
                break;
            case 'augment_error':
                NAIDirectorTools.onError(data.message);
                break;
            case 'tagger_result':
                NAITagger.onResult(data);
                break;
            case 'tagger_error':
                NAITagger.onError(data.message);
                break;
        }
    },

    updateStatus(message, showSpinner) {
        const el = $('#queue-status');
        if (el) {
            el.innerHTML = (showSpinner ? '<div class="spinner"></div>' : '') + message;
        }

        // Mobile toast notification
        if (window.innerWidth <= 768) {
            this._showToast(message, showSpinner);
        }
    },

    _showToast(message, showSpinner) {
        let toast = document.getElementById('mobile-toast');
        if (!toast) {
            toast = document.createElement('div');
            toast.id = 'mobile-toast';
            toast.className = 'mobile-toast';
            document.body.appendChild(toast);
        }
        toast.innerHTML = (showSpinner ? '<div class="spinner"></div>' : '') + message;
        toast.classList.add('show');

        clearTimeout(this._toastTimer);
        if (!showSpinner) {
            this._toastTimer = setTimeout(() => toast.classList.remove('show'), 3000);
        }
    },

    _toastTimer: null,

    send(obj) {
        if (NAIState.wsReady && NAIState.ws) {
            NAIState.ws.send(JSON.stringify(obj));
        }
    },
};