Kagi Assistant Collapsible Models.js

// ==UserScript==
// @name         Kagi Assistant Collapsible Models
// @namespace    http://tampermonkey.net/
// @version      1.6
// @description  Makes model sections in Kagi Assistant collapsible and saves their state.
// @author       Gemini 2.5 Pro
// @match        https://kagi.com/assistant*
// @grant        GM_addStyle
// @grant        GM_getValue
// @grant        GM_setValue
// @icon         https://kagi.com/favicon.ico
// ==/UserScript==

(function() {
    "use strict";

    const STORAGE_KEY = "kagi_collapsible_model_sections";

    // 1. Inject CSS for styling and fixes
    GM_addStyle(`
        /* Make headings clickable and add space for the icon */
        .promptOptionsSelector .heading {
            cursor: pointer;
            display: flex;
            align-items: center;
            justify-content: space-between;
            user-select: none; /* Prevent text selection on click */
        }

        /* Style for the chevron icon */
        .collapse-icon {
            transition: transform 0.2s ease-in-out;
            flex-shrink: 0; /* Prevent icon from shrinking */
        }

        /* Rotate icon when section is collapsed */
        .section-collapsed .collapse-icon {
            transform: rotate(-90deg);
        }

        /* FIX: Ensure custom assistant options with info are aligned in a row */
        .promptOptionsSelector .option.info-show {
            display: flex !important;
        }

        /* FIX: Override Kagi's default "show more" logic */
        .custom-assistants ul li {
            display: block !important;
        }
    `);

    // 2. State management functions
    function getStoredState() {
        return JSON.parse(GM_getValue(STORAGE_KEY, "{}"));
    }

    function setStoredState(state) {
        GM_setValue(STORAGE_KEY, JSON.stringify(state));
    }

    // 3. Main function to process the model selection dialog
    function makeSectionsCollapsible(dialog) {
        if (dialog.dataset.collapsibleProcessed) return;
        dialog.dataset.collapsibleProcessed = "true";

        // --- DOM Transformation for Style Consistency ---
        // Remove the original "Show more" button for Custom Assistants
        const customAssistantLabel = dialog.querySelector(".custom-assistants > label[for='custom-assistant-toggle']");
        if (customAssistantLabel) customAssistantLabel.remove();

        const defaultAssistantsList = dialog.querySelector(".default-assistants > ul");
        if (!defaultAssistantsList) return; // Cannot proceed without the main list

        // Helper function to move a section into the main list
        const normalizeSection = (sectionClass) => {
            const sectionDiv = dialog.querySelector(`.${sectionClass}`);
            if (!sectionDiv) return;

            const heading = sectionDiv.querySelector(".heading");
            const list = sectionDiv.querySelector("ul");

            if (heading && list) {
                const newGroup = document.createElement("li");
                newGroup.className = "group";
                newGroup.appendChild(heading);
                newGroup.appendChild(list);
                // Prepend to put Custom and Recommended at the top of the list
                defaultAssistantsList.prepend(newGroup);
                sectionDiv.remove(); // Remove the old div wrapper
            }
        };

        // Normalize both sections. Call in reverse order to get desired final order.
        normalizeSection("recommended-assistants");
        normalizeSection("custom-assistants");
        // --- End DOM Transformation ---

        const state = getStoredState();
        // Now all sections are uniform: li.group > .heading + ul
        const headings = dialog.querySelectorAll(".default-assistants .group > .heading");

        headings.forEach(heading => {
            // Use the heading's text as a unique key for storage
            const sectionKey = heading.querySelector("h2")?.textContent.trim() || heading.textContent.trim();
            if (!sectionKey) return;

            // The list of models is the heading's direct sibling
            const list = heading.nextElementSibling;
            if (!list || list.tagName !== "UL") return;

            // Add chevron icon if it doesn't exist
            if (!heading.querySelector(".collapse-icon")) {
                const icon = document.createElementNS("http://www.w3.org/2000/svg", "svg");
                icon.setAttribute("class", "icon-sm collapse-icon");
                icon.innerHTML = '<use href="#icon-chevron-up"></use>';
                heading.appendChild(icon);
            }

            // Function to apply the visual state (collapsed/expanded)
            const applyState = (isCollapsed) => {
                if (isCollapsed) {
                    heading.classList.add("section-collapsed");
                    list.style.display = "none";
                } else {
                    heading.classList.remove("section-collapsed");
                    list.style.display = "";
                }
            };

            // Apply initial state from storage
            applyState(state[sectionKey] === true);

            // Add click event listener to the heading to toggle state
            heading.addEventListener("click", (e) => {
                // Prevent clicks on links/buttons inside the heading from toggling
                if (e.target.closest("a, button, .trigger")) {
                    return;
                }

                const currentState = getStoredState();
                const newCollapsedState = !(currentState[sectionKey] === true);

                // Update and save the new state
                currentState[sectionKey] = newCollapsedState;
                setStoredState(currentState);

                // Apply the new visual state
                applyState(newCollapsedState);
            });
        });
    }

    // 4. Use MutationObserver to detect when the dialog is added to the page
    const observer = new MutationObserver((mutationsList) => {
        for (const mutation of mutationsList) {
            if (mutation.type === "childList") {
                mutation.addedNodes.forEach(node => {
                    if (node.nodeType === Node.ELEMENT_NODE) {
                        const dialog = node.matches("dialog.promptOptionsSelector") ? node : node.querySelector("dialog.promptOptionsSelector");
                        if (dialog) {
                            makeSectionsCollapsible(dialog);
                        }
                    }
                });
            }
        }
    });

    // Start observing the document body for added elements
    observer.observe(document.body, { childList: true, subtree: true });

})();