From 2199096191d4be1132ebdb9d516683c82be7c172 Mon Sep 17 00:00:00 2001
From: Cohee <18619528+Cohee1207@users.noreply.github.com>
Date: Mon, 28 Aug 2023 00:28:11 +0300
Subject: [PATCH] Add error handling of NAI streaming.

---
 public/scripts/extensions/hypebot/index.js |  2 ++
 public/scripts/nai-settings.js             | 24 ++++++++++++++++++++--
 2 files changed, 24 insertions(+), 2 deletions(-)

diff --git a/public/scripts/extensions/hypebot/index.js b/public/scripts/extensions/hypebot/index.js
index b972ac803..2fbd66764 100644
--- a/public/scripts/extensions/hypebot/index.js
+++ b/public/scripts/extensions/hypebot/index.js
@@ -172,6 +172,8 @@ async function generateHypeBot() {
         const output = decodeTextTokens(tokenizers.GPT2, ids).replace(/�/g, '').trim();
 
         setHypeBotText(formatReply(output));
+    } else {
+        setHypeBotText('<div class="hypebot_error">Something went wrong while generating a HypeBot reply. Please try again.</div>');
     }
 }
 
diff --git a/public/scripts/nai-settings.js b/public/scripts/nai-settings.js
index fb7328c10..3cd63aa9b 100644
--- a/public/scripts/nai-settings.js
+++ b/public/scripts/nai-settings.js
@@ -646,6 +646,24 @@ export function adjustNovelInstructionPrompt(prompt) {
     return stripedPrompt;
 }
 
+function tryParseStreamingError(decoded) {
+    try {
+        const data = JSON.parse(decoded);
+
+        if (!data) {
+            return;
+        }
+
+        if (data.message && data.statusCode >= 400) {
+            toastr.error(data.message, 'Error');
+            throw new Error(data);
+        }
+    }
+    catch {
+        // No JSON. Do nothing.
+    }
+}
+
 export async function generateNovelWithStreaming(generate_data, signal) {
     generate_data.streaming = nai_settings.streaming_novel;
 
@@ -663,12 +681,14 @@ export async function generateNovelWithStreaming(generate_data, signal) {
         let messageBuffer = "";
         while (true) {
             const { done, value } = await reader.read();
-            let response = decoder.decode(value);
+            let decoded = decoder.decode(value);
             let eventList = [];
 
+            tryParseStreamingError(decoded);
+
             // ReadableStream's buffer is not guaranteed to contain full SSE messages as they arrive in chunks
             // We need to buffer chunks until we have one or more full messages (separated by double newlines)
-            messageBuffer += response;
+            messageBuffer += decoded;
             eventList = messageBuffer.split("\n\n");
             // Last element will be an empty string or a leftover partial message
             messageBuffer = eventList.pop();