alexopus-ComfyUI-Notes-Sidebar/js/extension.js

import { app } from "/scripts/app.js";

// Load CSS
const link = document.createElement('link');
link.rel = 'stylesheet';
link.href = '/extensions/comfyui-notes-sidebar/notes-sidebar.css';
document.head.appendChild(link);

class NotesStore extends EventTarget {
    constructor() {
        super();
        this.state = {
            notes: [],
            activeNote: null,
            unsavedChanges: new Map(), // noteName -> boolean
            noteContent: new Map(), // noteName -> string (current content)
            savedContent: new Map() // noteName -> string (last saved content)
        };
    }

    getState() {
        return { ...this.state };
    }

    setState(updates) {
        const prevState = { ...this.state };
        this.state = { ...this.state, ...updates };

        // Emit state change event
        this.dispatchEvent(new CustomEvent('stateChanged', {
            detail: {
                prevState,
                newState: this.getState(),
                updates
            }
        }));
    }

    // Convenience methods for common state operations
    setNotes(notes) {
        this.setState({ notes });
    }

    setActiveNote(activeNote) {
        this.setState({ activeNote });
    }

    setNoteContent(noteName, content) {
        const noteContent = new Map(this.state.noteContent);
        noteContent.set(noteName, content);
        this.setState({ noteContent });
    }

    setSavedContent(noteName, content) {
        const savedContent = new Map(this.state.savedContent);
        savedContent.set(noteName, content);
        this.setState({ savedContent });
    }

    setUnsavedChanges(noteName, hasChanges) {
        const unsavedChanges = new Map(this.state.unsavedChanges);
        unsavedChanges.set(noteName, hasChanges);
        this.setState({ unsavedChanges });
    }

    deleteNoteFromState(noteName) {
        const unsavedChanges = new Map(this.state.unsavedChanges);
        const noteContent = new Map(this.state.noteContent);
        const savedContent = new Map(this.state.savedContent);

        unsavedChanges.delete(noteName);
        noteContent.delete(noteName);
        savedContent.delete(noteName);

        this.setState({ unsavedChanges, noteContent, savedContent });
    }

    renameNoteInState(oldName, newName) {
        const unsavedChanges = new Map(this.state.unsavedChanges);
        const noteContent = new Map(this.state.noteContent);
        const savedContent = new Map(this.state.savedContent);

        // Transfer state from old name to new name
        const unsavedState = unsavedChanges.get(oldName) || false;
        const content = noteContent.get(oldName) || '';
        const saved = savedContent.get(oldName) || '';

        // Remove old
        unsavedChanges.delete(oldName);
        noteContent.delete(oldName);
        savedContent.delete(oldName);

        // Add new
        unsavedChanges.set(newName, unsavedState);
        noteContent.set(newName, content);
        savedContent.set(newName, saved);

        this.setState({ unsavedChanges, noteContent, savedContent });
    }
}

class NotesSidebar {
    constructor() {
        this.store = new NotesStore();
        this.tabsContainer = null;
        this.contentArea = null;
        this.textArea = null;
        this.saveBtn = null;
        this.renameBtn = null;
        this.deleteBtn = null;
        this.appendLorasBtn = null;
        this.appendLorasTagsBtn = null;
        this.appendLorasMetaBtn = null;
        this.isInitialized = false;
        this.sidebarElement = null;

        // Listen to state changes
        this.store.addEventListener('stateChanged', (event) => {
            this.onStateChanged(event.detail);
        });
    }

    onStateChanged(changeDetails) {
        const { prevState, newState, updates } = changeDetails;

        // Update tabs if notes changed, active note changed, or unsaved changes changed
        if (updates.notes || updates.activeNote || updates.unsavedChanges) {
            this.renderTabs();
        }

        // Update textarea if active note changed or note content changed
        if (updates.activeNote || updates.noteContent) {
            this.updateTextArea();
            this.updateTextAreaState();
        }

        // Update button states if active note or unsaved changes changed
        if (updates.activeNote || updates.unsavedChanges) {
            this.updateButtonState();
        }
    }

    get notes() { return this.store.state.notes; }
    get activeNote() { return this.store.state.activeNote; }
    get unsavedChanges() { return this.store.state.unsavedChanges; }
    get noteContent() { return this.store.state.noteContent; }
    get savedContent() { return this.store.state.savedContent; }

    updateTextArea() {
        if (!this.textArea) return;

        if (!this.activeNote) {
            this.textArea.value = '';
            return;
        }

        // Check if we have unsaved content for this note
        if (this.noteContent.has(this.activeNote)) {
            this.textArea.value = this.noteContent.get(this.activeNote);
        } else {
            this.textArea.value = '';
        }
    }

    updateTextAreaState() {
        if (!this.textArea) return;

        const hasActiveNote = !!this.activeNote;
        this.textArea.disabled = !hasActiveNote;
        this.textArea.placeholder = hasActiveNote
            ? 'Start typing your notes...'
            : 'Select a note to start editing...';
    }

    updateButtonState() {
        if (!this.saveBtn || !this.renameBtn || !this.deleteBtn || !this.appendLorasBtn) return;

        const hasActiveNote = !!this.activeNote;
        const hasUnsavedChanges = hasActiveNote && this.unsavedChanges.get(this.activeNote);

        // Save button: enabled only when there's an active note AND unsaved changes
        this.saveBtn.disabled = !hasActiveNote || !hasUnsavedChanges;

        // Rename, delete, and append buttons: enabled only when there's an active note
        this.renameBtn.disabled = !hasActiveNote;
        this.deleteBtn.disabled = !hasActiveNote;
        this.appendLorasBtn.disabled = !hasActiveNote;
    }

    async init() {
        await this.loadNotes();
    }

    async createSidebarContent(el) {
        // Ensure the parent element has proper height
        el.style.height = '100%';
        el.style.display = 'flex';
        el.style.flexDirection = 'column';

        // If already initialized, just update the existing content
        if (this.isInitialized && this.sidebarElement) {
            el.innerHTML = '';
            el.appendChild(this.sidebarElement);
            // Trigger reactive updates when reopening
            this.renderTabs();
            this.updateTextArea();
            this.updateTextAreaState();
            return;
        }

        // Create new sidebar content
        this.sidebarElement = document.createElement('div');
        this.sidebarElement.className = 'notes-sidebar';

        this.createTabsContainer(this.sidebarElement);
        this.createNoteButtonContainer(this.sidebarElement);
        this.createContentArea(this.sidebarElement);

        el.appendChild(this.sidebarElement);
        this.isInitialized = true;

        // Load notes after creating UI
        await this.loadNotes();

        // Ensure UI reflects the current state after initialization
        this.renderTabs();
        this.updateTextArea();
        this.updateTextAreaState();
        this.updateButtonState();
    }

    createTabsContainer(parent) {
        const tabsWrapper = document.createElement('div');
        tabsWrapper.className = 'notes-tabs-wrapper';

        // Left side: tabs container for + button and note tabs
        this.tabsContainer = document.createElement('div');
        this.tabsContainer.className = 'notes-tabs-container';

        const newTabBtn = document.createElement('button');
        newTabBtn.innerHTML = '<i class="pi pi-plus"></i>';
        newTabBtn.title = 'New Note';
        newTabBtn.className = 'notes-new-tab-btn';
        newTabBtn.onclick = () => this.createNewNote();

        this.tabsContainer.appendChild(newTabBtn);

        // Right side: refresh button
        const refreshBtn = document.createElement('button');
        refreshBtn.innerHTML = '<i class="pi pi-refresh"></i>';
        refreshBtn.className = 'notes-btn refresh';
        refreshBtn.title = 'Refresh';
        refreshBtn.onclick = () => this.refreshNotes();

        tabsWrapper.appendChild(this.tabsContainer);
        tabsWrapper.appendChild(refreshBtn);
        parent.appendChild(tabsWrapper);
    }

    createNoteButtonContainer(parent) {
        const buttonContainer = document.createElement('div');
        buttonContainer.className = 'notes-button-container';

        const leftButtonContainer = document.createElement('div');
        leftButtonContainer.className = 'notes-button-group';

        const rightButtonContainer = document.createElement('div');
        rightButtonContainer.className = 'notes-button-group';

        this.saveBtn = document.createElement('button');
        this.saveBtn.innerHTML = '<i class="pi pi-save"></i>';
        this.saveBtn.className = 'notes-btn save';
        this.saveBtn.title = 'Save';
        this.saveBtn.onclick = () => this.saveActiveNote();
        this.saveBtn.disabled = true;

        this.renameBtn = document.createElement('button');
        this.renameBtn.innerHTML = '<i class="pi pi-file-edit"></i>';
        this.renameBtn.className = 'notes-btn rename';
        this.renameBtn.title = 'Rename';
        this.renameBtn.onclick = () => this.renameActiveNote();
        this.renameBtn.disabled = true;

        this.deleteBtn = document.createElement('button');
        this.deleteBtn.innerHTML = '<i class="pi pi-trash"></i>';
        this.deleteBtn.className = 'notes-btn delete';
        this.deleteBtn.title = 'Delete';
        this.deleteBtn.onclick = () => this.deleteActiveNote();
        this.deleteBtn.disabled = true;

        // LoRA submenu container
        this.loraMenuContainer = document.createElement('div');
        this.loraMenuContainer.className = 'notes-submenu-container';

        this.appendLorasBtn = document.createElement('button');
        this.appendLorasBtn.innerHTML = '<i class="pi pi-plus-circle"></i>';
        this.appendLorasBtn.className = 'notes-btn append-loras';
        this.appendLorasBtn.title = 'Append';
        this.appendLorasBtn.onclick = () => this.toggleLoraMenu();
        this.appendLorasBtn.disabled = true;

        // Create submenu
        this.loraSubmenu = document.createElement('div');
        this.loraSubmenu.className = 'notes-submenu';
        this.loraSubmenu.style.display = 'none';

        // Create submenu items
        const rawOption = document.createElement('div');
        rawOption.className = 'notes-submenu-item';
        rawOption.textContent = 'Append raw LoRA filenames list';
        rawOption.onclick = () => this.selectLoraOption('raw');

        const tagOption = document.createElement('div');
        tagOption.className = 'notes-submenu-item';
        tagOption.textContent = "Append LoRAs in '<lora:name>' format";
        tagOption.onclick = () => this.selectLoraOption('tag');

        const metaOption = document.createElement('div');
        metaOption.className = 'notes-submenu-item';
        metaOption.textContent = "Append LoRAs in '<lora:name>, metadata' format";
        metaOption.onclick = () => this.selectLoraOption('meta');

        this.loraSubmenu.appendChild(rawOption);
        this.loraSubmenu.appendChild(tagOption);
        this.loraSubmenu.appendChild(metaOption);

        this.loraMenuContainer.appendChild(this.appendLorasBtn);
        this.loraMenuContainer.appendChild(this.loraSubmenu);

        leftButtonContainer.appendChild(this.saveBtn);
        leftButtonContainer.appendChild(this.renameBtn);
        leftButtonContainer.appendChild(this.loraMenuContainer);
        rightButtonContainer.appendChild(this.deleteBtn);

        buttonContainer.appendChild(leftButtonContainer);
        buttonContainer.appendChild(rightButtonContainer);

        parent.appendChild(buttonContainer);
    }

    createContentArea(parent) {
        this.contentArea = document.createElement('div');
        this.contentArea.className = 'notes-content-area';

        this.textArea = document.createElement('textarea');
        this.textArea.className = 'notes-textarea';
        this.textArea.placeholder = 'Select a note to start editing...';

        // Initialize as disabled (will be enabled reactively if there's an active note)
        this.textArea.disabled = true;

        // Add event listener to track changes
        this.textArea.addEventListener('input', () => this.onTextAreaChange());

        this.contentArea.appendChild(this.textArea);
        parent.appendChild(this.contentArea);
    }

    async loadNotes() {
        try {
            const response = await fetch('/notes_sidebar/list');
            const data = await response.json();
            const notes = data.notes || [];

            this.store.setNotes(notes);

            if (notes.length > 0 && !this.activeNote) {
                this.store.setActiveNote(notes[0]);
            } else if (notes.length === 0) {
                // Ensure active note is cleared when no notes exist
                this.store.setActiveNote(null);
            }

            // Load content for active note after setting state
            if (this.activeNote) {
                await this.loadNoteContent(this.activeNote);
            }
        } catch (error) {
            console.error('Error loading notes:', error);
        }
    }

    renderTabs() {
        if (!this.tabsContainer) return;

        const existingTabs = this.tabsContainer.querySelectorAll('.note-tab');
        existingTabs.forEach(tab => tab.remove());

        this.notes.forEach(noteName => {
            const tab = document.createElement('button');
            const isUnsaved = this.unsavedChanges.get(noteName) || false;
            tab.textContent = (isUnsaved ? '* ' : '') + noteName;
            tab.className = `note-tab ${this.activeNote === noteName ? 'active' : ''}`;
            tab.title = noteName + (isUnsaved ? ' (unsaved)' : ''); // Full name in tooltip
            tab.onclick = () => this.selectNote(noteName);

            this.tabsContainer.appendChild(tab);
        });
    }

    async selectNote(noteName) {
        // Save current content before switching
        if (this.activeNote) {
            this.store.setNoteContent(this.activeNote, this.textArea.value);
        }

        this.store.setActiveNote(noteName);
        await this.loadNoteContent(noteName);
    }

    async loadNoteContent(noteName) {
        try {
            // Check if we have unsaved content for this note
            if (this.noteContent.has(noteName)) {
                // Content is already in state, updateTextArea will handle it
                return;
            } else {
                // Load from server
                const response = await fetch(`/notes_sidebar/note/${encodeURIComponent(noteName)}`);
                const data = await response.json();
                const content = data.content || '';

                // Store both current and saved content
                this.store.setNoteContent(noteName, content);
                this.store.setSavedContent(noteName, content);
            }
        } catch (error) {
            console.error('Error loading note content:', error);
        }
    }

    async saveActiveNote() {
        try {
            const content = this.textArea.value;
            const response = await fetch(`/notes_sidebar/note/${encodeURIComponent(this.activeNote)}`, {
                method: 'POST',
                headers: { 'Content-Type': 'application/json', },
                body: JSON.stringify({ content: content })
            });

            if (response.ok) {
                console.log('Note saved successfully');
                // Mark as saved
                this.store.setSavedContent(this.activeNote, content);
                this.store.setNoteContent(this.activeNote, content);
                this.store.setUnsavedChanges(this.activeNote, false);
            } else {
                console.error('Error saving note');
            }
        } catch (error) {
            console.error('Error saving note:', error);
        }
    }

    async deleteActiveNote() {
        if (!confirm(`Are you sure you want to delete "${this.activeNote}"?`)) return;

        try {
            const response = await fetch(`/notes_sidebar/note/${encodeURIComponent(this.activeNote)}`, { method: 'DELETE' });

            if (response.ok) {
                // Clean up tracking for deleted note
                const deletedNote = this.activeNote;

                // Clear active note first to ensure state change is triggered
                this.store.setActiveNote(null);
                this.store.deleteNoteFromState(deletedNote);

                // Refresh notes list (this will handle setting activeNote correctly)
                await this.refreshNotes();
            }
        } catch (error) {
            console.error('Error deleting note:', error);
        }
    }

    async renameActiveNote() {
        let newName = prompt('Enter new note name:', this.activeNote);
        if (!newName) return;

        // Trim spaces
        newName = newName.trim();
        if (!newName || newName === this.activeNote) return;

        try {
            if (!isValidFilename(newName)) {
                alert('Invalid filename! Please avoid special characters like / \\ : * ? " < > |');
                return;
            }

            // Check if new name already exists
            if (this.notes.includes(newName)) {
                alert('A note with this name already exists!');
                return;
            }

            // Get current content
            const currentContent = this.textArea.value;

            // Create new note with the content
            const createResponse = await fetch(`/notes_sidebar/note/${encodeURIComponent(newName)}`, {
                method: 'POST',
                headers: { 'Content-Type': 'application/json', },
                body: JSON.stringify({ content: currentContent })
            });

            if (createResponse.ok) {
                // Delete old note
                await fetch(`/notes_sidebar/note/${encodeURIComponent(this.activeNote)}`, { method: 'DELETE' });

                // Update tracking maps through store
                const oldName = this.activeNote;
                this.store.renameNoteInState(oldName, newName);

                // Set saved content to what we just saved
                this.store.setSavedContent(newName, currentContent);
                this.store.setUnsavedChanges(newName, false);

                // Update active note and refresh
                this.store.setActiveNote(newName);
                await this.refreshNotes();
            }
        } catch (error) {
            console.error('Error renaming note:', error);
        }
    }

    async createNewNote() {
        let noteName;

        // Keep asking for a name until user cancels or enters a valid unique name
        while (true) {
            noteName = prompt('Enter note name:');
            if (!noteName) return; // User cancelled

            // Trim spaces
            noteName = noteName.trim();

            // Validate filename safety
            if (!isValidFilename(noteName)) {
                alert('Invalid filename! Please avoid special characters like / \\ : * ? " < > |');
                continue; // Ask again
            }

            // Check if note name already exists
            if (this.notes.includes(noteName)) {
                alert('A note with this name already exists!');
                continue; // Ask again
            }

            break; // Valid unique name
        }

        try {
            const response = await fetch(`/notes_sidebar/note/${encodeURIComponent(noteName)}`, {
                method: 'POST',
                headers: { 'Content-Type': 'application/json', },
                body: JSON.stringify({ content: '' })
            });

            if (response.ok) {
                await this.refreshNotes();
                this.store.setActiveNote(noteName);

                // Initialize tracking for new note
                this.store.setNoteContent(noteName, '');
                this.store.setSavedContent(noteName, '');
                this.store.setUnsavedChanges(noteName, false);
            }
        } catch (error) {
            console.error('Error creating note:', error);
        }
    }

    async refreshNotes() {
        // Preserve current content before refreshing
        if (this.activeNote && this.textArea) {
            this.store.setNoteContent(this.activeNote, this.textArea.value);
        }

        await this.loadNotes();
    }

    onTextAreaChange() {
        const currentContent = this.textArea.value;
        const savedContent = this.savedContent.get(this.activeNote) || '';

        // Update current content
        this.store.setNoteContent(this.activeNote, currentContent);

        // Check if content has changed from saved version
        const hasUnsavedChanges = currentContent !== savedContent;
        this.store.setUnsavedChanges(this.activeNote, hasUnsavedChanges);
    }

    toggleLoraMenu() {
        if (this.loraSubmenu.style.display === 'none') {
            this.loraSubmenu.style.display = 'block';
            // Add click listener to close menu when clicking outside
            document.addEventListener('click', this.closeLoraMenuHandler);
        } else {
            this.loraSubmenu.style.display = 'none';
            document.removeEventListener('click', this.closeLoraMenuHandler);
        }
    }

    closeLoraMenuHandler = (event) => {
        if (!this.loraMenuContainer.contains(event.target)) {
            this.loraSubmenu.style.display = 'none';
            document.removeEventListener('click', this.closeLoraMenuHandler);
        }
    }

    selectLoraOption(type) {
        this.loraSubmenu.style.display = 'none';
        document.removeEventListener('click', this.closeLoraMenuHandler);

        switch(type) {
            case 'raw':
                this.appendLorasList();
                break;
            case 'tag':
                this.appendLorasTags();
                break;
            case 'meta':
                this.appendLorasWithMetadata();
                break;
        }
    }

    async appendLorasList() {
        try {
            const response = await fetch('/notes_sidebar/loras');
            const data = await response.json();

            if (response.ok && data.loras) {
                // Normalize backslashes to forward slashes
                const normalizedLoras = data.loras.map(lora => lora.replace(/\\/g, '/'));
                const lorasText = normalizedLoras.join('\n');

                // Add newline before appending if current content doesn't end with one
                const separator = this.textArea.value && !this.textArea.value.endsWith('\n') ? '\n' : '';
                const textToInsert = separator + lorasText;

                // Focus textarea to ensure it's active
                this.textArea.focus();

                // Move cursor to end
                this.textArea.setSelectionRange(this.textArea.value.length, this.textArea.value.length);

                // Use document.execCommand to preserve undo history
                document.execCommand('insertText', false, textToInsert);

                // Trigger change handling
                this.onTextAreaChange();
            } else {
                console.error('Error fetching LoRAs:', data.error || 'Unknown error');
                alert('Failed to fetch LoRAs list');
            }
        } catch (error) {
            console.error('Error appending LoRAs:', error);
            alert('Failed to append LoRAs list');
        }
    }

    async appendLorasTags() {
        try {
            const response = await fetch('/notes_sidebar/loras');
            const data = await response.json();

            if (response.ok && data.loras) {
                // Format each LoRA as <lora:name:1.0>, stripping .safetensors extension and normalizing slashes
                const lorasTags = data.loras.map(lora => {
                    const nameWithoutExt = lora.replace(/\.safetensors$/i, '');
                    const normalizedPath = nameWithoutExt.replace(/\\/g, '/');
                    return `<lora:${normalizedPath}>`;
                });
                const lorasText = lorasTags.join('\n');

                // Add newline before appending if current content doesn't end with one
                const separator = this.textArea.value && !this.textArea.value.endsWith('\n') ? '\n' : '';
                const textToInsert = separator + lorasText;

                // Focus textarea to ensure it's active
                this.textArea.focus();

                // Move cursor to end
                this.textArea.setSelectionRange(this.textArea.value.length, this.textArea.value.length);

                // Use document.execCommand to preserve undo history
                document.execCommand('insertText', false, textToInsert);

                // Trigger change handling
                this.onTextAreaChange();
            } else {
                console.error('Error fetching LoRAs:', data.error || 'Unknown error');
                alert('Failed to fetch LoRAs list');
            }
        } catch (error) {
            console.error('Error appending LoRAs tags:', error);
            alert('Failed to append LoRAs tags');
        }
    }

    async appendLorasWithMetadata() {
        try {
            const response = await fetch('/notes_sidebar/loras_with_metadata');
            const data = await response.json();

            if (response.ok && data.loras) {
                // Format each LoRA with metadata as <lora:name:1.0>, metadata
                const lorasWithMeta = data.loras.map(lora => {
                    const nameWithoutExt = lora.name.replace(/\.safetensors$/i, '');
                    const normalizedPath = nameWithoutExt.replace(/\\/g, '/');
                    const baseTag = `<lora:${normalizedPath}>`;

                    if (lora.metadata) {
                        return `${baseTag}, ${lora.metadata}`;
                    } else {
                        return baseTag;
                    }
                });
                const lorasText = lorasWithMeta.join('\n');

                // Add newline before appending if current content doesn't end with one
                const separator = this.textArea.value && !this.textArea.value.endsWith('\n') ? '\n' : '';
                const textToInsert = separator + lorasText;

                // Focus textarea to ensure it's active
                this.textArea.focus();

                // Move cursor to end
                this.textArea.setSelectionRange(this.textArea.value.length, this.textArea.value.length);

                // Use document.execCommand to preserve undo history
                document.execCommand('insertText', false, textToInsert);

                // Trigger change handling
                this.onTextAreaChange();
            } else {
                console.error('Error fetching LoRAs with metadata:', data.error || 'Unknown error');
                alert('Failed to fetch LoRAs with metadata');
            }
        } catch (error) {
            console.error('Error appending LoRAs with metadata:', error);
            alert('Failed to append LoRAs with metadata');
        }
    }
}

// Helper function to validate filename safety
function isValidFilename(filename) {
    if (!filename || filename.trim() === '') return false;

    // Check for unsafe characters: / \\ : * ? " < > |
    const unsafeChars = /[/\\:*?"<>|]/;
    if (unsafeChars.test(filename)) return false;

    return true;
}

const notesSidebar = new NotesSidebar();

app.registerExtension({
    name: "comfyui.notes.sidebar",
    async setup() {
        await notesSidebar.init();

        app.extensionManager.registerSidebarTab({
            id: "notes",
            icon: "pi pi-book",
            title: "Notes",
            tooltip: "Notes",
            type: "custom",
            render: (el) => {
                notesSidebar.createSidebarContent(el);
            },
        });
    }
});