add editing chat messages + if you chose a model it is now used as default for new chats

Author: Titanium2099

src/addons/addons/ai-integration/chat/stream.js

@@ -467,7 +467,12 @@ export function submitFromComposer () {
     const contextString = buildContextString(session, attachmentType);
 
     chatView.removeWelcomeHeader(pane);
-    const userBubble = chatView.appendUserMessage(pane, userText, attachmentType !== '0' ? attachmentName : null);
+    const userBubble = chatView.appendUserMessage(
+        pane,
+        userText,
+        attachmentType !== '0' ? attachmentName : null,
+        session.chatHistory.length
+    );
 
     composer.renderAttachmentUI();
 

src/addons/addons/ai-integration/main.js

@@ -5,7 +5,8 @@ import {
     store,
     getActiveSession,
     createNewSession,
-    loadSessionsForCurrentProject
+    loadSessionsForCurrentProject,
+    loadDefaultModelId
 } from './state/sessions.js';
 import * as host from './ui/host.js';
 import * as chatView from './ui/chatView.js';
@@ -30,6 +31,11 @@ export default class main {
         main.applySettings(settings);
 
         host.setComposerHandlers({onSend: submitFromComposer, onStop: stopActiveRequest});
+        try {
+            await loadDefaultModelId();
+        } catch (e) {
+            console.error('[Torchy] could not load the default model', e);
+        }
         await main.onProjectChanged();
         await host.initHost(addon);
         main.refreshModels();

src/addons/addons/ai-integration/state/sessions.js

@@ -7,6 +7,7 @@ import {
 } from '../helpers/db.js';
 
 const MAX_TAB_NAME_LENGTH = 28;
+const KEY_DEFAULT_MODEL = 'ui_defaultModelId';
 
 export class ChatSession {
     constructor (id, projectId, data = {}) {
@@ -51,6 +52,8 @@ export const store = {
     modelErrors: {},
     modelsLoading: false,
     canUse: false,
+    // Last model the user picked, new chats start on it.
+    defaultModelId: null,
 
     sessions: [],
     activeSessionId: null,
@@ -99,6 +102,17 @@ export function closeSession (sessionId) {
     }
 }
 
+export function setDefaultModelId (modelId) {
+    if (!modelId || store.defaultModelId === modelId) return;
+    store.defaultModelId = modelId;
+    saveMetadata(KEY_DEFAULT_MODEL, modelId);
+}
+
+export async function loadDefaultModelId () {
+    const saved = await loadMetadata(KEY_DEFAULT_MODEL);
+    if (typeof saved === 'string' && saved) store.defaultModelId = saved;
+}
+
 export function updateAndSaveSession (session) {
     if (session) saveSession(session.toJSON());
 }

src/addons/addons/ai-integration/style.css

@@ -51,6 +51,42 @@
     margin-top: 10px;
 }
 
+.torchy-edit-message {
+    margin-left: auto;
+    margin-right: 4px;
+    align-self: center;
+    width: 22px;
+    height: 22px;
+    padding: 3px;
+    flex: 0 0 auto;
+    border: none;
+    border-radius: 4px;
+    background: none;
+    color: #7C7766;
+    cursor: pointer;
+    opacity: 0;
+    transition: opacity 0.1s;
+}
+
+.torchy-edit-message svg {
+    width: 16px;
+    height: 16px;
+    display: block;
+}
+
+.user-message:hover .torchy-edit-message,
+.torchy-edit-message:focus-visible {
+    opacity: 1;
+}
+
+.torchy-edit-message:hover {
+    background: var(--ui-tertiary);
+}
+
+.user-message .torchy-edit-message + .message {
+    margin-left: 0;
+}
+
 .user-message .message {
     color: #7C7766;
     font-family: "Helvetica Neue", sans-serif;

src/addons/addons/ai-integration/ui/chatView.js

@@ -14,6 +14,7 @@ import {closeHost} from './host.js';
 import {
     getInputText,
     setInputText,
+    focusInput,
     getSelectedModelId,
     getAttachmentType,
     renderAttachmentUI,
@@ -199,9 +200,25 @@ function buildAttachmentChip (name) {
     return chip;
 }
 
-export function appendUserMessage (pane, text, attachmentName) {
+/**
+ * @param {number} turnIndex index this message has (or will have) in `session.chatHistory`;
+ *     editing the message rolls the chat back to it.
+ */
+export function appendUserMessage (pane, text, attachmentName, turnIndex) {
     const userMessage = document.createElement('div');
     userMessage.className = 'user-message';
+    userMessage.dataset.turnIndex = String(turnIndex);
+
+    const edit = document.createElement('button');
+    edit.type = 'button';
+    edit.className = 'torchy-edit-message';
+    edit.title = 'Edit this message';
+    edit.innerHTML =
+        '<svg xmlns="http://www.w3.org/2000/svg" fill="none" viewBox="0 0 24 24" stroke-width="1.5" ' +
+        'stroke="currentColor"><path stroke-linecap="round" stroke-linejoin="round" ' +
+        'd="m16.862 4.487 1.687-1.688a1.875 1.875 0 1 1 2.652 2.652L6.832 19.82a4.5 4.5 0 0 1-1.897 1.13l-2.685.8.8-2.685a4.5 4.5 0 0 1 1.13-1.897L16.863 4.487Z" /></svg>';
+    userMessage.appendChild(edit);
+
     const message = document.createElement('div');
     message.className = 'message';
 
@@ -213,9 +230,73 @@ export function appendUserMessage (pane, text, attachmentName) {
 
     userMessage.appendChild(message);
     pane.appendChild(userMessage);
+    attachEditButtonListeners(userMessage);
     return userMessage;
 }
 
+const ROLLBACK_WARNING_HTML =
+    '<p>Editing goes back to this point in the chat: this message and everything after it, ' +
+    "including Torchy's answers, will be deleted. The message is put back in the text box so you can change it " +
+    'and send it again.</p><br><p>Code you already added to your workspace is not touched.</p>';
+
+function countCodeChunks (history) {
+    return history.reduce((count, turn) => {
+        if (turn.role !== 'assistant') return count;
+        return count + (String(turn.message || '').match(FENCE_REGEX) || []).length;
+    }, 0);
+}
+
+function rollbackToTurn (session, turnIndex, text) {
+    const pane = getPaneForSession(session.id);
+    const kept = session.chatHistory.slice(0, Math.max(0, Math.min(turnIndex, session.chatHistory.length)));
+    session.chatHistory = kept;
+    session.allCodeChunksEverAdded = session.allCodeChunksEverAdded.slice(0, countCodeChunks(kept));
+    session.domCache = null;
+    session.inputText = text;
+    if (pane) rehydrateChatPane(session, pane);
+    if (session.id === store.activeSessionId) {
+        setInputText(text);
+        focusInput();
+    }
+    renderChatUI();
+    updateAndSaveSession(session);
+}
+
+function attachEditButtonListeners (parentElement) {
+    if (!parentElement) return;
+    parentElement.querySelectorAll('.torchy-edit-message').forEach(button => {
+        if (button.dataset.listenerAttached) return;
+        button.dataset.listenerAttached = 'true';
+        button.addEventListener('click', () => {
+            const session = getActiveSession();
+            const userMessage = button.closest('.user-message');
+            if (!session || !userMessage) return;
+            if (session.isBlabbering) {
+                helpers.showAlert('TorchyWaitForAIToFinishWarning');
+                return;
+            }
+            const turnIndex = Number(userMessage.dataset.turnIndex);
+            if (!Number.isInteger(turnIndex)) return;
+            const textEl = userMessage.querySelector('.message > span');
+            const text = textEl ? textEl.textContent : '';
+
+            const ScratchBlocks = window.ScratchBlocks;
+            if (!ScratchBlocks || typeof ScratchBlocks.prompt !== 'function') return;
+            ScratchBlocks.prompt(
+                ROLLBACK_WARNING_HTML,
+                null,
+                () => {
+                    if (session.isBlabbering) return;
+                    rollbackToTurn(session, turnIndex, text);
+                },
+                'Edit this message?',
+                ScratchBlocks.BROADCAST_MESSAGE_VARIABLE_TYPE,
+                true
+            );
+        });
+    });
+}
+
 export function removeWelcomeHeader (pane) {
     const header = pane.querySelector('.Popup_Header');
     if (header) header.remove();
@@ -446,12 +527,12 @@ export function rehydrateChatPane (session, pane) {
 
     let codeChunkCounter = 0;
 
-    session.chatHistory.forEach(turn => {
+    session.chatHistory.forEach((turn, turnIndex) => {
         if (turn.role === 'user') {
             const text = String(turn.message || '').split('\n\n\nContext:\n')[0];
             const attachmentName =
                 turn.attachment && turn.attachment.type !== '0' ? turn.attachment.name : null;
-            appendUserMessage(pane, text, attachmentName);
+            appendUserMessage(pane, text, attachmentName, turnIndex);
             return;
         }
         if (turn.role !== 'assistant') return;
@@ -582,10 +663,11 @@ function insertChunkIntoWorkspace (chunkData, workspace, replacingBlockIds) {
 
 export function restoreInteractivity (pane) {
     if (!pane) return;
-    pane.querySelectorAll('.insert_button').forEach(button => {
+    pane.querySelectorAll('.insert_button, .torchy-edit-message').forEach(button => {
         delete button.dataset.listenerAttached;
     });
     attachInsertButtonListeners(pane);
+    attachEditButtonListeners(pane);
     pane.querySelectorAll('.torchy-warnings').forEach(wrap => {
         const toggle = wrap.querySelector('.torchy-warnings-toggle');
         const list = wrap.querySelector('.torchy-warnings-list');

src/addons/addons/ai-integration/ui/composer.js

@@ -1,4 +1,4 @@
-import {store, getActiveSession, updateAndSaveSession} from '../state/sessions.js';
+import {store, getActiveSession, updateAndSaveSession, setDefaultModelId} from '../state/sessions.js';
 import {pickDefaultModel} from '../chat/models.js';
 import {addCustomModelId, CUSTOM_MODELS_GROUP} from '../chat/customModels.js';
 
@@ -143,7 +143,8 @@ export function rebuildModelOptions () {
     const previous = select.value;
     select.textContent = '';
 
-    const defaultModel = pickDefaultModel(store.AIModels);
+    const defaultModel =
+        store.AIModels.find(model => model.key === store.defaultModelId) || pickDefaultModel(store.AIModels);
     const defaultId = defaultModel ? defaultModel.key : null;
     const groups = new Map();
     const groupFor = (provider, label) => {
@@ -235,6 +236,16 @@ export function renderAttachmentUI () {
     select.value = session ? session.attachmentType : '0';
 }
 
+/** Make `modelId` the model new chats start on. */
+function rememberDefaultModel (modelId) {
+    setDefaultModelId(modelId);
+    const select = modelSelect();
+    if (!select) return;
+    Array.from(select.options).forEach(option => {
+        option.defaultSelected = option.value === modelId;
+    });
+}
+
 /** Ask for a model id with the editor's prompt modal, then select (and remember) it. */
 function promptForModelId (providerId) {
     const provider = store.providers.find(p => p.id === providerId);
@@ -259,6 +270,7 @@ function selectTypedModel (provider, input) {
         session.selectedModelId = model.key;
         updateAndSaveSession(session);
     }
+    setDefaultModelId(model.key);
     rebuildModelOptions();
 }
 
@@ -306,6 +318,7 @@ export function createComposer ({onSend, onStop}) {
             promptForModelId(value.slice(ADD_MODEL_PREFIX.length));
             return;
         }
+        rememberDefaultModel(value);
         const session = getActiveSession();
         if (!session) return;
         session.selectedModelId = value;