From 74fc259c9cbc7a9d19a6212072d3673383bbdc8f Mon Sep 17 00:00:00 2001
From: Wolfsblvt <wolfsblvt@gmail.com>
Date: Tue, 14 Jan 2025 00:59:21 +0100
Subject: [PATCH] Reworked persona lock buttons + lock to character

- Slight redesign of the persona panel right side, add headers
- Add buttons for all three different lock states (chat, character, default)
- Update lock and toggle lock functions to accept toggle type
- Write persona locking to chat functionality (saves to persona description object)
---
 public/index.html          |  27 ++++-
 public/scripts/personas.js | 221 ++++++++++++++++++++++++++++---------
 public/style.css           |   8 ++
 3 files changed, 200 insertions(+), 56 deletions(-)

diff --git a/public/index.html b/public/index.html
index 073ff8912..0e6c0ecf5 100644
--- a/public/index.html
+++ b/public/index.html
@@ -4924,8 +4924,6 @@
                                     <input id="your_name" name="your_name" data-i18n="[placeholder]Enter your name" placeholder="Enter your name" class="text_pole wide100p" value="" autocomplete="off">
                                     <div id="your_name_button" class="menu_button fa-solid fa-check" title="Click to set a new User Name" data-i18n="[title]Click to set a new User Name">
                                     </div>
-                                    <div id="lock_user_name" class="menu_button fa-solid fa-unlock" title="Click to lock your selected persona to the current chat. Click again to remove the lock." data-i18n="[title]Click to lock your selected persona to the current chat. Click again to remove the lock.">
-                                    </div>
                                     <div id="sync_name_button" class="menu_button fa-solid fa-sync" title="Click to set user name for all messages" data-i18n="[title]Click to set user name for all messages">
                                     </div>
                                     <div id="persona_lore_button" class="menu_button fa-solid fa-globe" title="Persona Lore&#10;Alt+Click to open the lorebook" data-i18n="[title]Persona Lore Alt+Click to open the lorebook">
@@ -4937,10 +4935,6 @@
                                 <div class="extension_token_counter">
                                     <span data-i18n="Tokens persona description">Tokens</span>: <span id="persona_description_token_count">0</span>
                                 </div>
-                            </div>
-
-                            <div class="persona_management_global_settings">
-                                <h4 class="standoutHeader" data-i18n="Global Settings">Global Settings</h4>
 
                                 <div class="persona_management_description_position_container">
                                     <label for="persona_description_position" data-i18n="Position:">Position:</label>
@@ -4967,6 +4961,27 @@
                                     </div>
                                 </div>
 
+                                <h4 data-i18n="Connections">Connections</h4>
+                                <div id="persona_connections_buttons" class="flex-container marginBot10">
+                                    <div id="lock_persona_default" class="menu_button menu_button_icon" title="Click to select this as default persona for the new chats. Click again to remove the lock." data-i18n="[title]Click to select this as default persona for the new chats. Click again to remove the lock.">
+                                        <i class="icon fa-solid fa-crown fa-fw"></i>
+                                        <div data-i18n="Default">Default</div>
+                                    </div>
+                                    <div id="lock_user_name" class="menu_button menu_button_icon" title="Click to lock your selected persona to the current chat. Click again to remove the lock." data-i18n="[title]Click to lock your selected persona to the current chat. Click again to remove the lock.">
+                                        <i class="icon fa-solid fa-unlock fa-fw"></i>
+                                        <div data-i18n="Chat">Chat</div>
+                                    </div>
+                                    <div id="lock_persona_to_char" class="menu_button menu_button_icon" title="Click to lock your selected persona to the current character. Click again to remove the lock." data-i18n="[title]Click to lock your selected persona to the current character. Click again to remove the lock.">
+                                        <i class="icon fa-solid fa-unlock fa-fw"></i>
+                                        <div data-i18n="Character">Character</div>
+                                    </div>
+                                </div>
+                                <div id="persona_connections_list" class="text_muted m-b-1" data-i18n="persona_no_connections">[No connections]</div>
+                            </div>
+
+                            <div class="persona_management_global_settings">
+                                <h4 class="standoutHeader" data-i18n="Global Settings">Global Settings</h4>
+
                                 <div class="persona_management_show_notifications range-block">
                                     <label for="persona_show_notifications" class="checkbox_label">
                                         <input id="persona_show_notifications" type="checkbox" />
diff --git a/public/scripts/personas.js b/public/scripts/personas.js
index 69558fb7b..6a722f78d 100644
--- a/public/scripts/personas.js
+++ b/public/scripts/personas.js
@@ -7,7 +7,9 @@ import {
     event_types,
     getRequestHeaders,
     getThumbnailUrl,
+    menu_type,
     name1,
+    name2,
     reloadCurrentChat,
     saveChatConditional,
     saveMetadata,
@@ -26,6 +28,14 @@ import { t } from './i18n.js';
 import { openWorldInfoEditor, world_names } from './world-info.js';
 import { renderTemplateAsync } from './templates.js';
 
+/**
+ * @typedef {object} PersonaConnection A connection between a character and a character or group entity
+ * @property {'character' | 'group'} type - Type of connection
+ * @property {string} id - ID of the connection (character key (avatar url), group id)
+ */
+
+/** @typedef {'chat' | 'character' | 'default'} PersonaLockType Type of the persona lock */
+
 let savePersonasPage = 0;
 const GRID_STORAGE_KEY = 'Personas_GridView';
 const DEFAULT_DEPTH = 2;
@@ -607,15 +617,35 @@ function selectCurrentPersona() {
 }
 
 /**
- * Checks if the persona is locked for the current chat.
- * @returns {boolean} Whether the persona is locked
+ * Checks if a connection is locked for the current character or group edit menu
+ * @param {PersonaConnection} connection - Connection to check
+ * @returns {boolean} Whether the connection is locked
  */
-function isPersonaLocked() {
-    return !!chat_metadata['persona'];
+function isPersonaConnectionLocked(connection) {
+    return (menu_type === 'character_edit' && connection.type === 'character' && connection.id === characters[this_chid]?.avatar)
+        || (menu_type === 'group_edit' && connection.type === 'group' && connection.id === selected_group);
 }
 
 /**
- * Locks or unlocks the persona for the current chat.
+ * Checks if the persona is locked
+ * @param {PersonaLockType} type - Lock type
+ * @returns {boolean} Whether the persona is locked
+ */
+function isPersonaLocked(type = 'chat') {
+    switch (type) {
+        case 'default':
+            return power_user.default_persona === user_avatar;
+        case 'chat':
+            return chat_metadata['persona'] == user_avatar;
+        case 'character': {
+            return !!power_user.persona_descriptions[user_avatar]?.connections?.some(isPersonaConnectionLocked);
+        }
+        default: throw new Error(`Unknown persona lock type: ${type}`);
+    }
+}
+
+/**
+ * Locks or unlocks the persona
  * @param {boolean} state Desired lock state
  * @returns {Promise<void>}
  */
@@ -624,34 +654,63 @@ export async function setPersonaLockState(state) {
 }
 
 /**
- * Toggle the persona lock state for the current chat.
+ * Toggle the persona lock state
+ * @param {PersonaLockType} type - Lock type
  * @returns {Promise<void>}
  */
-export async function togglePersonaLock() {
-    return isPersonaLocked()
-        ? await unlockPersona()
-        : await lockPersona();
+export async function togglePersonaLock(type = 'chat') {
+    return isPersonaLocked(type)
+        ? await unlockPersona(type)
+        : await lockPersona(type);
 }
 
 /**
- * Unlock the persona for the current chat.
+ * Unlock the persona
+ * @param {PersonaLockType} type - Lock type
+ * @returns {Promise<void>}
  */
-async function unlockPersona() {
-    if (chat_metadata['persona']) {
-        console.log(`Unlocking persona for this chat ${chat_metadata['persona']}`);
-        delete chat_metadata['persona'];
-        await saveMetadata();
-        if (power_user.persona_show_notifications) {
-            toastr.info(t`User persona is now unlocked for this chat. Click the "Lock" again to revert.`, t`Persona unlocked`);
+async function unlockPersona(type = 'chat') {
+    switch (type) {
+        case 'default': {
+            // TODO: Make this toggle-able
+            await toggleDefaultPersona(user_avatar, { quiet: true });
+            break;
+        }
+        case 'chat': {
+            if (chat_metadata['persona']) {
+                console.log(`Unlocking persona ${user_avatar} from this chat`);
+                delete chat_metadata['persona'];
+                await saveMetadata();
+                if (power_user.persona_show_notifications) {
+                    toastr.info(t`User persona ${name1} is now unlocked from this chat. Click the "Lock" again to revert.`, t`Persona unlocked`);
+                }
+            }
+            break;
+        }
+        case 'character': {
+            /** @type {PersonaConnection[]} */
+            const connections = power_user.persona_descriptions[user_avatar]?.connections;
+            if (connections) {
+                console.log(`Unlocking persona ${user_avatar} from this character ${name2}`);
+                power_user.persona_descriptions[user_avatar].connections = connections.filter(c => !isPersonaConnectionLocked(c));
+                saveSettingsDebounced();
+                if (power_user.persona_show_notifications) {
+                    toastr.info(t`User persona ${name1} is now unlocked from character ${name2}. Click the "Lock" again to revert.`, t`Persona unlocked`);
+                }
+            }
+            break;
         }
-        updateUserLockIcon();
     }
+
+    updatePersonaLockIcons();
 }
 
 /**
- * Lock the persona for the current chat.
+ * Lock the persona
+ * @param {PersonaLockType} type - Lock type
  */
-async function lockPersona() {
+async function lockPersona(type = 'chat') {
+    // First make sure that user_avatar is actually a persona
     if (!(user_avatar in power_user.personas)) {
         console.log(`Creating a new persona ${user_avatar}`);
         if (power_user.persona_show_notifications) {
@@ -668,17 +727,44 @@ async function lockPersona() {
             depth: DEFAULT_DEPTH,
             role: DEFAULT_ROLE,
             lorebook: '',
+            connections: [],
         };
     }
 
-    chat_metadata['persona'] = user_avatar;
-    await saveMetadata();
-    saveSettingsDebounced();
-    console.log(`Locking persona for this chat ${user_avatar}`);
-    if (power_user.persona_show_notifications) {
-        toastr.success(t`User persona is locked to ${name1} in this chat`);
+    switch (type) {
+        case 'default': {
+            // TODO: Make this toggle-able
+            await toggleDefaultPersona(user_avatar, { quiet: true });
+            break;
+        }
+        case 'chat': {
+            console.log(`Locking persona ${user_avatar} to this chat`);
+            chat_metadata['persona'] = user_avatar;
+            await saveMetadata();
+            saveSettingsDebounced();
+            if (power_user.persona_show_notifications) {
+                toastr.success(t`User persona ${name1} is locked to ${name2} in this chat`);
+            }
+            break;
+        }
+        case 'character': {
+            const newConnection = menu_type === 'character_edit' ? { type: 'character', id: characters[this_chid]?.avatar } :
+                menu_type === 'group_edit' ? { type: 'group', id: selected_group } : null;
+            /** @type {PersonaConnection[]} */
+            const connections = power_user.persona_descriptions[user_avatar].connections?.filter(c => !isPersonaConnectionLocked(c)) ?? [];
+            if (newConnection && newConnection.id) {
+                console.log(`Locking persona ${user_avatar} to this character ${name2}`);
+                power_user.persona_descriptions[user_avatar].connections = [...connections, newConnection];
+                saveSettingsDebounced();
+                if (power_user.persona_show_notifications) {
+                    toastr.success(t`User persona ${name1} is locked to character ${name2}`);
+                }
+            }
+            break;
+        }
     }
-    updateUserLockIcon();
+
+    updatePersonaLockIcons();
 }
 
 
@@ -730,7 +816,7 @@ async function deleteUserAvatar(e) {
 
         saveSettingsDebounced();
         await getUserAvatars();
-        updateUserLockIcon();
+        updatePersonaLockIcons();
     }
 }
 
@@ -853,16 +939,33 @@ function getOrCreatePersonaDescriptor() {
             depth: power_user.persona_description_depth,
             role: power_user.persona_description_role,
             lorebook: power_user.persona_description_lorebook,
+            connections: [],
         };
         power_user.persona_descriptions[user_avatar] = object;
     }
     return object;
 }
 
-async function setDefaultPersona(e) {
-    e?.stopPropagation();
-    const avatarId = $(this).closest('.avatar-container').find('.avatar').attr('imgfile');
 
+
+async function toggleDefaultPersonaClicked(e) {
+    e?.stopPropagation();
+    const avatarId = $(e.currentTarget).closest('.avatar-container').find('.avatar').attr('imgfile');
+    if (avatarId) {
+        await toggleDefaultPersona(avatarId);
+    } else {
+        console.warn('No avatar id found');
+    }
+}
+
+/**
+ * Sets a persona as the default one to be used for all new chats and unlocked existing chats
+ * @param {string} avatarId The avatar id of the persona to set as the default
+ * @param {object} [options] Optional arguments
+ * @param {boolean} [options.quiet=false] If true, no confirmation popups will be shown
+ * @returns {Promise<void>}
+ */
+async function toggleDefaultPersona(avatarId, { quiet: quiet = false } = {}) {
     if (!avatarId) {
         console.warn('No avatar id found');
         return;
@@ -879,11 +982,12 @@ async function setDefaultPersona(e) {
     const personaName = power_user.personas[avatarId];
 
     if (avatarId === currentDefault) {
-        const confirm = await Popup.show.confirm(t`Are you sure you want to remove the default persona?`, personaName);
-
-        if (!confirm) {
-            console.debug('User cancelled removing default persona');
-            return;
+        if (!quiet) {
+            const confirm = await Popup.show.confirm(t`Are you sure you want to remove the default persona?`, personaName);
+            if (!confirm) {
+                console.debug('User cancelled removing default persona');
+                return;
+            }
         }
 
         console.log(`Removing default persona ${avatarId}`);
@@ -892,11 +996,12 @@ async function setDefaultPersona(e) {
         }
         delete power_user.default_persona;
     } else {
-        const confirm = await Popup.show.confirm(t`Are you sure you want to set \"${personaName}\" as the default persona?`, t`This name and avatar will be used for all new chats, as well as existing chats where the user persona is not locked.`);
-
-        if (!confirm) {
-            console.debug('User cancelled setting default persona');
-            return;
+        if (!quiet) {
+            const confirm = await Popup.show.confirm(t`Are you sure you want to set \"${personaName}\" as the default persona?`, t`This name and avatar will be used for all new chats, as well as existing chats where the user persona is not locked.`);
+            if (!confirm) {
+                console.debug('User cancelled setting default persona');
+                return;
+            }
         }
 
         power_user.default_persona = avatarId;
@@ -907,12 +1012,26 @@ async function setDefaultPersona(e) {
 
     saveSettingsDebounced();
     await getUserAvatars(true, avatarId);
+    updatePersonaLockIcons();
 }
 
-function updateUserLockIcon() {
-    const hasLock = !!chat_metadata['persona'];
-    $('#lock_user_name').toggleClass('fa-unlock', !hasLock);
-    $('#lock_user_name').toggleClass('fa-lock', hasLock);
+function updatePersonaLockIcons() {
+    const isDefaultPersona = power_user.default_persona === user_avatar;
+    $('#lock_persona_default').toggleClass('locked', isDefaultPersona);
+
+    const hasChatLock = chat_metadata['persona'] == user_avatar;
+    $('#lock_user_name').toggleClass('locked', hasChatLock);
+    $('#lock_user_name i.icon').toggleClass('fa-lock', hasChatLock);
+    $('#lock_user_name i.icon').toggleClass('fa-unlock', !hasChatLock);
+
+    /** @type {PersonaConnection[]} */
+    const connections = power_user.persona_descriptions[user_avatar]?.connections;
+    const hasCharLock = !!connections.some(c =>
+        (menu_type === 'character_edit' && c.type === 'character' && c.id === characters[this_chid]?.avatar)
+        || (menu_type === 'group_edit' && c.type === 'group' && c.id === selected_group));
+    $('#lock_persona_to_char').toggleClass('locked', hasCharLock);
+    $('#lock_persona_to_char i.icon').toggleClass('fa-lock', hasCharLock);
+    $('#lock_persona_to_char i.icon').toggleClass('fa-unlock', !hasCharLock);
 }
 
 async function setChatLockedPersona() {
@@ -942,7 +1061,7 @@ async function setChatLockedPersona() {
     if (chat_metadata['persona'] && !userAvatars.includes(chatPersona)) {
         console.warn('Persona avatar not found, unlocking persona');
         delete chat_metadata['persona'];
-        updateUserLockIcon();
+        updatePersonaLockIcons();
         return;
     }
 
@@ -956,7 +1075,7 @@ async function setChatLockedPersona() {
 
     // Persona avatar found, select it
     setUserAvatar(chatPersona);
-    updateUserLockIcon();
+    updatePersonaLockIcons();
 }
 
 function onBackupPersonas() {
@@ -1112,9 +1231,11 @@ async function duplicatePersona(avatarId) {
 
 export function initPersonas() {
     $(document).on('click', '.bind_user_name', bindUserNameToPersona);
-    $(document).on('click', '.set_default_persona', setDefaultPersona);
+    $(document).on('click', '.set_default_persona', toggleDefaultPersonaClicked);
     $(document).on('click', '.delete_avatar', deleteUserAvatar);
-    $('#lock_user_name').on('click', togglePersonaLock);
+    $('#lock_persona_default').on('click', () => togglePersonaLock('default'));
+    $('#lock_user_name').on('click', () => togglePersonaLock('chat'));
+    $('#lock_persona_to_char').on('click', () => togglePersonaLock('character'));
     $('#create_dummy_persona').on('click', createDummyPersona);
     $('#persona_description').on('input', onPersonaDescriptionInput);
     $('#persona_description_position').on('input', onPersonaDescriptionPositionInput);
diff --git a/public/style.css b/public/style.css
index 9716d1906..5bc832dc0 100644
--- a/public/style.css
+++ b/public/style.css
@@ -3759,6 +3759,14 @@ input[type='checkbox'].del_checkbox {
     color: var(--golden);
 }
 
+#lock_persona_default.locked i.icon {
+    color: var(--golden);
+}
+#lock_user_name.locked i.icon,
+#lock_persona_to_char.locked i.icon {
+    color: var(--active);
+}
+
 #user_avatar_block .avatar_upload {
     cursor: pointer;
     width: 60px;