From 58b8e7908d3e26128f813a7a38ed7caec68f4792 Mon Sep 17 00:00:00 2001
From: Wolfsblvt <wolfsblvt@gmail.com>
Date: Tue, 11 Feb 2025 23:34:21 +0100
Subject: [PATCH] Add data attribute during generation

- Fix reasoning buttons still appearing
- Allow easier customization of HTML elements via CSS that appear during gen
---
 public/script.js | 16 +++++++++-------
 public/style.css |  5 +++++
 2 files changed, 14 insertions(+), 7 deletions(-)

diff --git a/public/script.js b/public/script.js
index cad1e00b5..6f0fa790f 100644
--- a/public/script.js
+++ b/public/script.js
@@ -3160,14 +3160,16 @@ class StreamingProcessor {
         }
     }
 
-    showMessageButtons(messageId) {
+    markUIGenStarted() {
         showStopButton();
-        $(`#chat .mes[mesid="${messageId}"] .mes_buttons`).css({ 'display': 'none' });
+        const chatElement = document.getElementById('chat');
+        chatElement.dataset.generating = 'true';
     }
 
-    hideMessageButtons(messageId) {
+    markUIGenStopped() {
         hideStopButton();
-        $(`#chat .mes[mesid="${messageId}"] .mes_buttons`).css({ 'display': 'flex' });
+        const chatElement = document.getElementById('chat');
+        delete chatElement.dataset.generating;
     }
 
     async onStartStreaming(text) {
@@ -3180,7 +3182,7 @@ class StreamingProcessor {
             await saveReply(this.type, text, true, '', [], '');
             messageId = chat.length - 1;
             this.#checkDomElements(messageId);
-            this.showMessageButtons(messageId);
+            this.markUIGenStarted();
         }
         hideSwipeButtons();
         scrollChatToBottom();
@@ -3268,7 +3270,7 @@ class StreamingProcessor {
     }
 
     async onFinishStreaming(messageId, text) {
-        this.hideMessageButtons(this.messageId);
+        this.markUIGenStopped();
         await this.onProgressStreaming(messageId, text, true);
         addCopyToCodeBlocks($(`#chat .mes[mesid="${messageId}"]`));
 
@@ -3313,7 +3315,7 @@ class StreamingProcessor {
         this.abortController.abort();
         this.isStopped = true;
 
-        this.hideMessageButtons(this.messageId);
+        this.markUIGenStopped();
         generatedPromptCache = '';
         unblockGeneration();
 
diff --git a/public/style.css b/public/style.css
index 9092d1885..420bf4f8c 100644
--- a/public/style.css
+++ b/public/style.css
@@ -4331,6 +4331,11 @@ input[type="range"]::-webkit-slider-thumb {
     field-sizing: content;
 }
 
+#chat[data-generating="true"] .last_mes .mes_buttons,
+#chat[data-generating="true"] .last_mes .mes_reasoning_actions {
+    display: none;
+}
+
 #anchor_order {
     margin-bottom: 15px;
 }