From 13299e5c395219ca95a9fc91f4319c4ffdd89a09 Mon Sep 17 00:00:00 2001
From: 50h100a <all_awful@proton.me>
Date: Thu, 17 Aug 2023 05:52:32 -0400
Subject: [PATCH] deal with websocket shenanigans (and also clarify what's
 being printed out sometimes)

---
 server.js | 35 +++++++++++++++++++++++++----------
 1 file changed, 25 insertions(+), 10 deletions(-)

diff --git a/server.js b/server.js
index 0a0695d79..74f24990b 100644
--- a/server.js
+++ b/server.js
@@ -551,6 +551,7 @@ app.post("/generate", jsonParser, async function (request, response_generate = r
                 }
 
                 const data = await response.json();
+                console.log("Endpoint response:", data);
                 return response_generate.send(data);
             }
         } catch (error) {
@@ -600,18 +601,13 @@ app.post("/generate_textgenerationwebui", jsonParser, async function (request, r
             const websocket = new WebSocket(streamingUrl);
 
             websocket.on('open', async function () {
-                console.log('websocket open');
+                console.log('WebSocket opened');
                 websocket.send(JSON.stringify(request.body));
             });
 
-            websocket.on('error', (err) => {
-                console.error(err);
-                websocket.close();
-            });
-
             websocket.on('close', (code, buffer) => {
                 const reason = new TextDecoder().decode(buffer)
-                console.log(reason);
+                console.log("WebSocket closed (reason: %o)", reason);
             });
 
             while (true) {
@@ -620,8 +616,27 @@ app.post("/generate_textgenerationwebui", jsonParser, async function (request, r
                     websocket.close();
                     return;
                 }
+                
+                let rawMessage = null;
+                try {
+                    // This lunacy is because the websocket can fail to connect AFTER we're awaiting 'message'... so 'message' never triggers.
+                    // So instead we need to look for 'error' at the same time to reject the promise. And then remove the listener if we resolve.
+                    // This is awful.
+                    // Welcome to the shenanigan shack.
+                    rawMessage = await new Promise(function (resolve, reject) {
+                        websocket.once('error', reject);
+                        websocket.once('message', (data, isBinary) => {
+                            websocket.removeListener('error', reject);
+                            resolve(data, isBinary);
+                        });
+                    });
+                } catch(err) {
+                    console.error("Socket error:", err);
+                    websocket.close();
+                    yield "[SillyTavern] Streaming failed:\n" + err;
+                    return;
+                }
 
-                const rawMessage = await new Promise(resolve => websocket.once('message', resolve));
                 const message = json5.parse(rawMessage);
 
                 switch (message.event) {
@@ -672,11 +687,11 @@ app.post("/generate_textgenerationwebui", jsonParser, async function (request, r
 
         try {
             const data = await postAsync(api_server + "/v1/generate", args);
-            console.log(data);
+            console.log("Endpoint response:", data);
             return response_generate.send(data);
         } catch (error) {
             retval = { error: true, status: error.status, response: error.statusText };
-            console.log(error);
+            console.log("Endpoint error:", error);
             try {
                 retval.response = await error.json();
                 retval.response = retval.response.result;