vg-coder-cli 2.0.46 → 2.0.47

package/src/server/views/js/features/task-worker.js

@@ -0,0 +1,448 @@
+// Task Worker — runs inside the AI Studio page (shadow-DOM bundled).
+// Connects to localhost:6868 via Socket.IO, registers as the worker,
+// and executes incoming task:execute jobs by calling window.AIChat.
+
+import { io } from 'socket.io-client';
+import { API_BASE } from '../config.js';
+
+const ALLOWED_HOSTS = ['aistudio.google.com'];
+const HEARTBEAT_MS = 15_000;
+
+let socket = null;
+let heartbeatTimer = null;
+let emailAbort = null;
+const cancelFlags = new Set();   // taskIds that received task:cancel
+
+// ---- Console / error ring buffer ----
+// Captures recent log entries so a remote caller can ask "what just went
+// wrong?" without an attached devtools panel. Hook is installed once at
+// init; entries persist until the buffer wraps.
+const LOG_BUFFER_MAX = 200;
+const logBuffer = [];     // [{ at, level, source, message }]
+let logCaptureInstalled = false;
+
+function pushLog(level, args, source = 'console') {
+    let message;
+    try {
+        message = args.map(a => {
+            if (a instanceof Error) return `${a.name}: ${a.message}\n${a.stack || ''}`;
+            if (typeof a === 'object' && a !== null) {
+                try { return JSON.stringify(a); } catch (_) { return String(a); }
+            }
+            return String(a);
+        }).join(' ');
+    } catch (_) { message = '<serialize-failed>'; }
+    if (logBuffer.length >= LOG_BUFFER_MAX) logBuffer.shift();
+    logBuffer.push({ at: Date.now(), level, source, message: message.slice(0, 4000) });
+}
+
+function installLogCapture() {
+    if (logCaptureInstalled) return;
+    logCaptureInstalled = true;
+    ['log', 'info', 'warn', 'error', 'debug'].forEach(lvl => {
+        const orig = console[lvl]?.bind(console) || console.log.bind(console);
+        console[lvl] = function (...args) {
+            pushLog(lvl, args);
+            try { orig(...args); } catch (_) {}
+        };
+    });
+    window.addEventListener('error', (e) => {
+        pushLog('error', [`window.error: ${e.message || ''}`, e.error?.stack || `at ${e.filename || '?'}:${e.lineno || 0}:${e.colno || 0}`], 'window');
+    });
+    window.addEventListener('unhandledrejection', (e) => {
+        const r = e.reason;
+        const msg = r instanceof Error ? r.message : (typeof r === 'object' ? JSON.stringify(r) : String(r));
+        pushLog('error', [`unhandledrejection: ${msg}`, r?.stack || ''], 'promise');
+    });
+}
+
+function isAllowedHost() {
+    try {
+        return ALLOWED_HOSTS.some(h => location.hostname.endsWith(h));
+    } catch (_) { return false; }
+}
+
+async function waitForAIChat(timeoutMs = 30_000) {
+    const start = Date.now();
+    while (Date.now() - start < timeoutMs) {
+        if (window.AIChat && typeof window.AIChat.send === 'function') return true;
+        await new Promise(r => setTimeout(r, 300));
+    }
+    return false;
+}
+
+/**
+ * Extract the signed-in Google account email from AI Studio's DOM.
+ * Tries several selectors + a generic aria-label sweep, then falls back to
+ * scanning page text for the first email-shaped string.
+ */
+const EMAIL_RE = /[\w.+-]+@[\w-]+\.[\w-]+/;
+
+function extractEmail() {
+    const candidates = [
+        () => document.querySelector('a[aria-label*="@"]')?.getAttribute('aria-label'),
+        () => document.querySelector('a[href^="https://accounts.google.com/SignOutOptions"]')?.getAttribute('aria-label'),
+        () => document.querySelector('img[alt*="@"]')?.alt,
+        () => document.querySelector('[data-email]')?.getAttribute('data-email'),
+        () => Array.from(document.querySelectorAll('[aria-label]'))
+                  .map(e => e.getAttribute('aria-label') || '')
+                  .find(t => EMAIL_RE.test(t))
+    ];
+    for (const fn of candidates) {
+        try {
+            const raw = (fn() || '').toString();
+            const m = raw.match(EMAIL_RE);
+            if (m) return m[0].toLowerCase();
+        } catch (_) {}
+    }
+    return null;
+}
+
+/**
+ * Wait for the email to appear in DOM. Retries forever (every `pollMs`) so a
+ * worker that registers on /welcome (pre-login) and only later signs in still
+ * gets its email recorded eventually. Returns null only if `signal.aborted`.
+ */
+async function resolveEmail(pollMs = 3000, signal = null) {
+    while (true) {
+        if (signal?.aborted) return null;
+        const email = extractEmail();
+        if (email) return email;
+        await new Promise(r => setTimeout(r, pollMs));
+    }
+}
+
+async function fetchAsFile(url, name, mime) {
+    // Use direct fetch (PNA is already granted on first click) and pull a
+    // real Blob — passing a Blob into File() preserves binary fidelity and
+    // matches what the dashboard agent panel sends from clipboard / drag-drop.
+    //
+    // Do NOT route through window.vetgoFetch for binary content: chrome.runtime
+    // message-passing between background SW → content script → MAIN world
+    // structured-clones the ArrayBuffer in a way that breaks `instanceof
+    // ArrayBuffer` on the receiving end (verified with SHA-256 mismatch via
+    // SubtleCrypto). For text it's fine; for binary it silently corrupts.
+    const r = await fetch(url);
+    if (!r.ok) throw new Error(`Failed to fetch ${name}: HTTP ${r.status}`);
+    const blob = await r.blob();
+    return new File([blob], name, { type: mime || blob.type || 'application/octet-stream' });
+}
+
+// Wait for the assistant turn to actually render (not just the user's echo turn).
+// AI Studio renders the user turn first, then the model turn. We need ≥ baseline + 2.
+async function waitForNewAssistantTurn(baselineCount, timeoutMs = 90_000) {
+    const start = Date.now();
+    while (Date.now() - start < timeoutMs) {
+        // Surface rate-limit / quota errors early — no point waiting for a turn
+        // that will never appear.
+        const rl = detectRateLimit();
+        if (rl) { const e = new Error(rl.message); e.code = rl.code; throw e; }
+
+        const turns = document.querySelectorAll('ms-chat-turn');
+        if (turns.length >= baselineCount + 2) {
+            const last = turns[turns.length - 1];
+            const role = last.querySelector('[data-turn-role]')?.getAttribute('data-turn-role');
+            if (role === 'Model') return true;
+        }
+        await new Promise(r => setTimeout(r, 250));
+    }
+    return false;
+}
+
+/**
+ * Inspect the page for known rate-limit / quota / token-count / error states.
+ * Returns { code, message } if any is present; null otherwise.
+ * Selectors based on observed AI Studio DOM.
+ */
+function detectRateLimit() {
+    // Snackbar toast (top-level, e.g. "Failed to generate content: user has exceeded quota")
+    const toast = document.querySelector('mat-snack-bar-container, .mdc-snackbar, .cdk-overlay-container .mat-mdc-snack-bar-label');
+    if (toast) {
+        const t = (toast.textContent || '').trim();
+        if (/exceeded quota|quota/i.test(t)) return { code: 'quota_exceeded', message: t.slice(0, 500) };
+        if (/rate limit/i.test(t))           return { code: 'rate_limit_exceeded', message: t.slice(0, 500) };
+        if (/failed to generate/i.test(t))   return { code: 'generation_failed', message: t.slice(0, 500) };
+    }
+
+    // Inline model-error inside the last assistant turn
+    const modelErr = document.querySelector('.model-error');
+    if (modelErr) {
+        const text = (modelErr.textContent || '').trim();
+        if (/rate limit/i.test(text))       return { code: 'rate_limit_exceeded', message: text };
+        if (/quota/i.test(text))            return { code: 'quota_exceeded', message: text };
+        if (/internal error/i.test(text))   return { code: 'model_internal_error', message: text };
+        if (text)                           return { code: 'model_error', message: text };
+    }
+
+    // Standalone quota / upgrade callout
+    const callout = document.querySelector('ms-upgrade-options-callout, .quota-exceeded-message');
+    if (callout) {
+        return { code: 'quota_exceeded', message: (callout.textContent || 'Quota exceeded').trim().slice(0, 500) };
+    }
+
+    // Token count failure
+    const tokenErr = document.querySelector('.token-status-error');
+    if (tokenErr) {
+        return { code: 'token_count_failed', message: (tokenErr.textContent || 'Token count failed').trim() };
+    }
+
+    return null;
+}
+
+function throwIfRateLimited(stage = '') {
+    const rl = detectRateLimit();
+    if (rl) {
+        const e = new Error(`[${stage}] ${rl.message}`);
+        e.code = rl.code;
+        throw e;
+    }
+}
+
+async function handleTaskExecute(payload) {
+    const { taskId, prompt, files = [] } = payload || {};
+    if (!taskId) return;
+
+    console.log(`[TaskWorker] Executing ${taskId} (${files.length} file(s))`);
+
+    try {
+        // 1. Open a fresh chat for this task — isolates context, avoids stale results
+        if (typeof window.AIChat.startNewChat === 'function') {
+            const ok = await window.AIChat.startNewChat();
+            console.log(`[TaskWorker] startNewChat → ${ok ? 'ready' : 'fallback (continuing)'}`);
+        }
+        await new Promise(r => setTimeout(r, 600)); // settle after navigation
+
+        if (cancelFlags.has(taskId)) {
+            cancelFlags.delete(taskId);
+            console.log(`[TaskWorker] ${taskId} canceled before send — skipping`);
+            return;
+        }
+
+        // 2. Snapshot turn count before sending so we can detect the new model turn
+        const baselineTurns = document.querySelectorAll('ms-chat-turn').length;
+
+        const blobs = await Promise.all(
+            files.map(f => fetchAsFile(f.url, f.name, f.mime))
+        );
+
+        // 3. Send prompt + files; AIChat.send waits for the Run button to settle
+        await window.AIChat.send({ prompt: prompt || '', files: blobs });
+
+        // 3b. Immediately surface rate-limit / quota errors that AI Studio shows
+        //     via toast or inline error before any model turn renders.
+        throwIfRateLimited('post_send');
+
+        // 4. Wait for the assistant's turn to actually render in the DOM.
+        //    waitForNewAssistantTurn also probes for rate-limit during the wait.
+        const turnReady = await waitForNewAssistantTurn(baselineTurns);
+        if (!turnReady) console.warn(`[TaskWorker] ${taskId} — assistant turn not detected; copying anyway`);
+
+        // 5. Re-check after turn renders — error turns count as turns too.
+        throwIfRateLimited('post_turn');
+
+        // 6. Extra settle so streaming finalizes before we trigger Copy-as-markdown
+        await new Promise(r => setTimeout(r, 1500));
+
+        if (cancelFlags.has(taskId)) {
+            cancelFlags.delete(taskId);
+            console.log(`[TaskWorker] ${taskId} canceled — discarding result`);
+            return;
+        }
+
+        const markdown = await window.AIChat.copyLastTurnAsMarkdown();
+        const chatId = window.AIChat.getChatIdFromUrl?.() || null;
+
+        if (cancelFlags.has(taskId)) {
+            cancelFlags.delete(taskId);
+            console.log(`[TaskWorker] ${taskId} canceled post-copy — discarding`);
+            return;
+        }
+
+        socket.emit('task:complete', { taskId, markdown, chatId });
+        console.log(`[TaskWorker] Completed ${taskId}`);
+    } catch (err) {
+        console.error(`[TaskWorker] Failed ${taskId}:`, err);
+        socket.emit('task:failed', {
+            taskId,
+            code: err?.code || 'execution_error',
+            message: err?.message || String(err)
+        });
+    }
+}
+
+export function initTaskWorker() {
+    if (!isAllowedHost()) {
+        console.log('[TaskWorker] Skipping init — not on aistudio.google.com');
+        return;
+    }
+    if (socket) return;
+    installLogCapture();
+
+    waitForAIChat().then((ready) => {
+        if (!ready) {
+            console.warn('[TaskWorker] window.AIChat not ready — task worker disabled');
+            return;
+        }
+        connect();
+    });
+}
+
+function connect() {
+    socket = io(API_BASE, { transports: ['websocket', 'polling'] });
+
+    socket.on('connect', () => {
+        console.log(`[TaskWorker] Connected as ${socket.id}`);
+
+        // Reset any outstanding email retry from a previous session
+        if (emailAbort) emailAbort.abort();
+        emailAbort = new AbortController();
+
+        // Register with whatever email we have right now (may be null on cold load).
+        const initialEmail = extractEmail();
+        const chromeId = (window.vetgo && window.vetgo.chromeId) || null;
+        socket.emit('worker:register', {
+            domain: location.hostname,
+            chatId: window.AIChat?.getChatIdFromUrl?.() || null,
+            userAgent: navigator.userAgent,
+            email: initialEmail,
+            chromeId
+        });
+        console.log(`[TaskWorker] Initial email: ${initialEmail || '(pending)'}, chromeId: ${chromeId || '(none)'}`);
+
+        if (heartbeatTimer) clearInterval(heartbeatTimer);
+        heartbeatTimer = setInterval(() => {
+            try { socket.emit('worker:heartbeat'); } catch (_) {}
+        }, HEARTBEAT_MS);
+
+        // Detached, infinite email retry — keeps polling DOM until email shows up.
+        // Stops when socket disconnects (handled below).
+        if (!initialEmail) {
+            (async () => {
+                try {
+                    const resolved = await resolveEmail(3000, emailAbort?.signal);
+                    if (!resolved || !socket.connected) return;
+                    socket.emit('worker:register', {
+                        domain: location.hostname,
+                        chatId: window.AIChat?.getChatIdFromUrl?.() || null,
+                        userAgent: navigator.userAgent,
+                        email: resolved,
+                        chromeId: (window.vetgo && window.vetgo.chromeId) || null
+                    });
+                    console.log(`[TaskWorker] Re-registered with email: ${resolved}`);
+                } catch (err) {
+                    console.error('[TaskWorker] Email retry failed:', err);
+                }
+            })();
+        }
+    });
+
+    socket.on('worker:rejected', (info) => {
+        console.warn('[TaskWorker] Worker rejected:', info);
+    });
+
+    socket.on('worker:replaced', () => {
+        console.warn('[TaskWorker] Replaced by newer worker — disconnecting');
+        try { socket.disconnect(); } catch (_) {}
+    });
+
+    socket.on('task:execute', handleTaskExecute);
+
+    socket.on('task:cancel', ({ taskId }) => {
+        if (taskId) cancelFlags.add(taskId);
+        console.log(`[TaskWorker] Cancel requested for ${taskId}`);
+    });
+
+    // ----- Remote debug surface -----
+    socket.on('debug:url', (_payload, ack) => {
+        try {
+            ack && ack({
+                ok: true,
+                href: location.href,
+                pathname: location.pathname,
+                hostname: location.hostname,
+                title: document.title,
+                chatId: window.AIChat?.getChatIdFromUrl?.() || null,
+                rateLimit: detectRateLimit() || null
+            });
+        } catch (e) { ack && ack({ ok: false, error: e.message }); }
+    });
+
+    socket.on('debug:probe', ({ selector, kind = 'text' } = {}, ack) => {
+        try {
+            const els = Array.from(document.querySelectorAll(selector));
+            const out = els.slice(0, 50).map(el => {
+                if (kind === 'html')   return el.innerHTML;
+                if (kind === 'outer')  return el.outerHTML;
+                if (kind === 'attrs')  return Object.fromEntries(Array.from(el.attributes).map(a => [a.name, a.value]));
+                return (el.textContent || '').trim();
+            });
+            ack && ack({ ok: true, count: els.length, results: out });
+        } catch (e) { ack && ack({ ok: false, error: e.message }); }
+    });
+
+    socket.on('debug:dom', ({ selector, maxBytes = 200_000 } = {}, ack) => {
+        try {
+            const root = selector ? document.querySelector(selector) : document.documentElement;
+            if (!root) return ack && ack({ ok: false, error: 'selector_not_found' });
+            let html = root.outerHTML || '';
+            const truncated = html.length > maxBytes;
+            if (truncated) html = html.slice(0, maxBytes);
+            ack && ack({ ok: true, length: root.outerHTML.length, truncated, html });
+        } catch (e) { ack && ack({ ok: false, error: e.message }); }
+    });
+
+    socket.on('debug:eval', async ({ code } = {}, ack) => {
+        try {
+            // eslint-disable-next-line no-new-func
+            const fn = new Function(`"use strict"; return (async () => { ${code} })();`);
+            const value = await fn();
+            // Best-effort safe-serialize
+            let serialized;
+            try { serialized = JSON.parse(JSON.stringify(value, (_k, v) => {
+                if (v instanceof Element) return { _element: v.tagName, id: v.id, classes: v.className };
+                if (typeof v === 'function') return `[Function ${v.name || 'anonymous'}]`;
+                return v;
+            })); }
+            catch (_) { serialized = String(value); }
+            ack && ack({ ok: true, value: serialized });
+        } catch (e) { ack && ack({ ok: false, error: e.message, stack: e.stack }); }
+    });
+
+    socket.on('debug:logs', ({ since = 0, level = null, limit = 100, clear = false } = {}, ack) => {
+        try {
+            let logs = logBuffer.filter(e => e.at > Number(since || 0));
+            if (level) {
+                const lvls = Array.isArray(level) ? level : [level];
+                logs = logs.filter(e => lvls.includes(e.level));
+            }
+            if (limit > 0) logs = logs.slice(-limit);
+            const total = logBuffer.length;
+            const errors = logBuffer.filter(e => e.level === 'error').length;
+            if (clear) logBuffer.length = 0;
+            ack && ack({ ok: true, logs, total, errors, now: Date.now() });
+        } catch (e) { ack && ack({ ok: false, error: e.message }); }
+    });
+
+    socket.on('debug:screenshot', async ({ format = 'png', quality } = {}, ack) => {
+        try {
+            // Native chrome.tabs.captureVisibleTab via background SW —
+            // viewport-only screenshot of the active tab in this window.
+            // No canvas rasterization, handles cross-origin content cleanly.
+            if (typeof window.vetgoCaptureScreenshot !== 'function') {
+                return ack && ack({ ok: false, error: 'extension_not_loaded' });
+            }
+            const dataUrl = await window.vetgoCaptureScreenshot({ format, quality });
+            ack && ack({ ok: true, dataUrl, source: 'chrome.tabs.captureVisibleTab' });
+        } catch (e) { ack && ack({ ok: false, error: e.message }); }
+    });
+
+    socket.on('disconnect', (reason) => {
+        console.log(`[TaskWorker] Disconnected: ${reason}`);
+        if (heartbeatTimer) { clearInterval(heartbeatTimer); heartbeatTimer = null; }
+        if (emailAbort) { emailAbort.abort(); emailAbort = null; }
+    });
+
+    socket.on('connect_error', (err) => {
+        console.warn('[TaskWorker] Connect error:', err?.message);
+    });
+}

package/src/server/views/js/main.js

@@ -19,6 +19,7 @@ import { initGitPanel } from './features/git-panel.js';
 import { initCommandsPanel } from './features/commands-panel.js';
 import { initBrowserPanel } from './features/browser-panel.js';
 import { initAgentPanel } from './features/agent-panel.js';
+import { initTaskWorker } from './features/task-worker.js';
 
 export async function initMain() {
     console.log('VG Coder: Starting Main Logic...');
@@ -55,6 +56,9 @@ export async function initMain() {
         // Initialize keyboard shortcuts (global hotkeys for Shadow DOM)
         initKeyboardShortcuts();
 
+        // Remote task worker — only activates on aistudio.google.com
+        initTaskWorker();
+
         console.log('✅ VG Coder: Initialization Complete');
     } catch (e) {
         console.error('VG Coder Init Failed:', e);