Google AI Studio | Toggle Code Blocks

Toggle all code blocks open/closed in Google AI Studio with lazy loading support.
줄 바꿈
// ==UserScript==
// @name         Google AI Studio | Toggle Code Blocks
// @namespace    https://greasyfork.org/en/users/1462137-piknockyou
// @version      1.4
// @author       Piknockyou (vibe-coded)
// @license      AGPL-3.0
// @description  Toggle all code blocks open/closed in Google AI Studio with lazy loading support.
// @match        https://aistudio.google.com/*
// @icon         https://www.google.com/s2/favicons?sz=64&domain=aistudio.google.com
// @grant        GM_addStyle
// @run-at       document-idle
// ==/UserScript==

(function() {
    'use strict';

    if (window._codeBlockToggleLoaded) return;
    window._codeBlockToggleLoaded = true;
    if (window.self !== window.top) return;

    //================================================================================
    // STATE & STORAGE
    //================================================================================
    const STORAGE_KEY = 'codeblock_toggle_state';

    function loadState() {
        try {
            const saved = localStorage.getItem(STORAGE_KEY);
            if (saved) {
                const state = JSON.parse(saved);
                return {
                    isActive: state.isActive ?? false,
                    collapseMode: state.collapseMode ?? true
                };
            }
        } catch (e) {
            console.warn('[Code Block Toggle] Failed to load state:', e);
        }
        return { isActive: false, collapseMode: true };
    }

    function saveState() {
        try {
            localStorage.setItem(STORAGE_KEY, JSON.stringify({
                isActive,
                collapseMode
            }));
        } catch (e) {
            console.warn('[Code Block Toggle] Failed to save state:', e);
        }
    }

    let { isActive, collapseMode } = loadState();

    //================================================================================
    // STYLES
    //================================================================================
    GM_addStyle(`
        #codeblock-toggle-button {
            margin: 0 4px;
        }
        #codeblock-toggle-button.mode-collapse {
            color: #4285f4 !important;
        }
        #codeblock-toggle-button.mode-expand {
            color: #fbbc04 !important;
        }
    `);

    //================================================================================
    // CORE LOGIC
    //================================================================================

    function applyModeToBlock(header) {
        // Debounce per-block to prevent rapid re-clicking
        const now = Date.now();
        const lastClick = parseInt(header.dataset.toggleTs || '0', 10);
        if (now - lastClick < 500) return;

        const isExpanded = header.getAttribute('aria-expanded') === 'true';
        const wantExpanded = !collapseMode;

        if (isExpanded !== wantExpanded) {
            header.dataset.toggleTs = String(now);
            header.click();
        }
    }

    function applyModeToAllBlocks() {
        document.querySelectorAll('ms-code-block mat-expansion-panel-header')
            .forEach(applyModeToBlock);
    }

    function toggleMode() {
        if (!isActive) {
            isActive = true;
            collapseMode = true; // First click always collapses
        } else {
            collapseMode = !collapseMode;
        }
        applyModeToAllBlocks();
        updateButtonState();
        saveState();
    }

    function updateButtonState() {
        const button = document.getElementById('codeblock-toggle-button');
        const icon = button?.querySelector('span');
        if (!button || !icon) return;

        button.classList.remove('mode-collapse', 'mode-expand');

        if (!isActive) {
            icon.textContent = 'expand_less';
            button.title = 'Toggle Code Blocks';
            button.setAttribute('aria-label', 'Toggle Code Blocks');
        } else if (collapseMode) {
            icon.textContent = 'expand_less';
            button.title = 'Collapse Mode Active';
            button.setAttribute('aria-label', 'Collapse Mode Active');
            button.classList.add('mode-collapse');
        } else {
            icon.textContent = 'expand_more';
            button.title = 'Expand Mode Active';
            button.setAttribute('aria-label', 'Expand Mode Active');
            button.classList.add('mode-expand');
        }
    }

    //================================================================================
    // OBSERVER - Only processes NEW blocks (for lazy loading)
    //================================================================================
    function handleNewBlocks(mutations) {
        if (!isActive) return;

        for (const mutation of mutations) {
            for (const node of mutation.addedNodes) {
                if (node.nodeType !== 1) continue;

                // Check if node itself is a code block
                if (node.matches?.('ms-code-block')) {
                    const header = node.querySelector('mat-expansion-panel-header');
                    if (header) applyModeToBlock(header);
                }

                // Check descendants for code blocks
                if (node.querySelectorAll) {
                    node.querySelectorAll('ms-code-block mat-expansion-panel-header')
                        .forEach(applyModeToBlock);
                }
            }
        }
    }

    const blockObserver = new MutationObserver(handleNewBlocks);

    //================================================================================
    // UI
    //================================================================================
    function createToolbarButton(toolbar = document.querySelector('ms-toolbar .toolbar-right')) {
        if (document.getElementById('codeblock-toggle-button')) return false;

        if (!toolbar) return false;

        const btn = document.createElement('button');
        btn.id = 'codeblock-toggle-button';
        btn.title = 'Toggle Code Blocks';
        btn.setAttribute('ms-button', '');
        btn.setAttribute('variant', 'icon-borderless');
        btn.setAttribute('mattooltip', 'Toggle Code Blocks');
        btn.setAttribute('mattooltipposition', 'below');
        btn.className = 'mat-mdc-tooltip-trigger ms-button-borderless ms-button-icon ng-star-inserted';
        btn.setAttribute('aria-label', 'Toggle Code Blocks');
        btn.setAttribute('aria-disabled', 'false');
        btn.addEventListener('click', toggleMode);

        const icon = document.createElement('span');
        icon.className = 'material-symbols-outlined notranslate ms-button-icon-symbol ng-star-inserted';
        icon.setAttribute('aria-hidden', 'true');
        icon.textContent = 'expand_less';
        btn.appendChild(icon);

        const moreBtn = toolbar.querySelector('button[iconname="more_vert"]');
        toolbar.insertBefore(btn, moreBtn || null);

        return true;
    }

    //================================================================================
    // DEBUG LOGGING
    //================================================================================
    const DEBUG = false;
    function log(msg, data = null) {
        if (!DEBUG) return;
        const prefix = '[Code Block Toggle]';
        if (data) {
            console.log(`${prefix} ${msg}`, data);
        } else {
            console.log(`${prefix} ${msg}`);
        }
    }

    //================================================================================
    // TOOLBAR OBSERVER - Persists button across SPA navigation
    //================================================================================
    // More efficient: avoid scanning every added node subtree; do a single toolbar lookup per mutation batch.
    const toolbarObserver = new MutationObserver(() => {
        // If our button already exists, nothing to do.
        if (document.getElementById('codeblock-toggle-button')) return;

        const toolbar = document.querySelector('ms-toolbar .toolbar-right');
        if (!toolbar) return;

        if (createToolbarButton(toolbar)) {
            updateButtonState();
            if (isActive) {
                setTimeout(applyModeToAllBlocks, 300);
            }
        }
    });

    //================================================================================
    // INIT
    //================================================================================
    function init() {
        log('Initializing...');
        log(`Loaded state: isActive=${isActive}, collapseMode=${collapseMode}`);

        // Check initial DOM state
        const initialToolbar = document.querySelector('ms-toolbar .toolbar-right');
        log(`Initial toolbar exists: ${!!initialToolbar}`);

        // Try to add button immediately
        if (createToolbarButton()) {
            log('Initial button creation successful');
            updateButtonState();
            if (isActive) {
                setTimeout(applyModeToAllBlocks, 300);
            }
        } else {
            log('Initial button creation failed - will wait for observer');
        }

        // Keep observing for toolbar changes (never disconnect - SPA support)
        toolbarObserver.observe(document.body, { childList: true, subtree: true });
        log('Toolbar observer started');

        // Start observing for lazy-loaded code blocks
        blockObserver.observe(document.body, { childList: true, subtree: true });
        log('Block observer started');
    }

    // Log navigation events
    const origPush = history.pushState;
    history.pushState = function() {
        log('>>> history.pushState triggered', { url: arguments[2] });
        const r = origPush.apply(this, arguments);

        setTimeout(() => {
            log('Post-pushState check:');
            log(`  Button in DOM: ${!!document.getElementById('codeblock-toggle-button')}`);
            log(`  Toolbar exists: ${!!document.querySelector('ms-toolbar .toolbar-right')}`);
        }, 500);

        return r;
    };

    const origReplace = history.replaceState;
    history.replaceState = function() {
        log('>>> history.replaceState triggered', { url: arguments[2] });
        return origReplace.apply(this, arguments);
    };

    window.addEventListener('popstate', () => {
        log('>>> popstate event triggered');
    });

    init();

})();