Toggle all code blocks open/closed in Google AI Studio with lazy loading support.
// ==UserScript==
// @name Google AI Studio | Collapse/Expand All Code Blocks Toggle
// @namespace https://greasyfork.org/en/users/1462137-piknockyou
// @version 1.6
// @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();
})();