From 9f97a144e81f6f128adf74369b489441ab0b8755 Mon Sep 17 00:00:00 2001
From: 50h100a <all_awful@proton.me>
Date: Sat, 19 Oct 2024 00:31:12 -0400
Subject: [PATCH] slightly change stream "abort" flow so token probabilities
 get successfully updated

---
 public/script.js | 14 +++++++++-----
 1 file changed, 9 insertions(+), 5 deletions(-)

diff --git a/public/script.js b/public/script.js
index f39c58dc5..1986ce9a8 100644
--- a/public/script.js
+++ b/public/script.js
@@ -3194,7 +3194,8 @@ class StreamingProcessor {
     }
 
     onStopStreaming() {
-        this.onErrorStreaming();
+        this.abortController.abort();
+        this.isFinished = true;
     }
 
     /**
@@ -3239,9 +3240,13 @@ class StreamingProcessor {
             console.warn(`Stream stats: ${timestamps.length} tokens, ${seconds.toFixed(2)} seconds, rate: ${Number(timestamps.length / seconds).toFixed(2)} TPS`);
         }
         catch (err) {
-            console.error(err);
-            this.onErrorStreaming();
-            return;
+            // in the case of a self-inflicted abort, we have already cleaned up
+            if (!this.isFinished)
+            {
+                console.error(err);
+                this.onErrorStreaming();
+            }
+            return this.result;
         }
 
         this.isFinished = true;
@@ -4641,7 +4646,6 @@ export function stopGeneration() {
     let stopped = false;
     if (streamingProcessor) {
         streamingProcessor.onStopStreaming();
-        streamingProcessor = null;
         stopped = true;
     }
     if (abortController) {