From d6d443279c9542cdd01c362d3981e2cfbdcc9bfa Mon Sep 17 00:00:00 2001
From: Lex Lim <hyperzlib@outlook.com>
Date: Sun, 23 Oct 2022 22:04:56 +0800
Subject: [PATCH] =?UTF-8?q?=E5=A2=9E=E5=8A=A0=E4=B8=8B=E8=BD=BD=E5=9B=BE?=
 =?UTF-8?q?=E7=89=87=E6=97=B63=E6=AC=A1=E9=87=8D=E8=AF=95?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 src/data/request/remoterequest.ts     | 87 +++++++++++++++++----------
 src/styles/components/conversation.ts |  2 +-
 2 files changed, 56 insertions(+), 33 deletions(-)

diff --git a/src/data/request/remoterequest.ts b/src/data/request/remoterequest.ts
index c687c6a..1850c86 100644
--- a/src/data/request/remoterequest.ts
+++ b/src/data/request/remoterequest.ts
@@ -1133,6 +1133,8 @@ export class RemoteImageGenerationRequest {
         onError: (err: { status: number; message: string }) => void,
         onClose: () => void
     ): Promise<void> {
+        const MAX_RETRY_TIMES = 3;
+
         const requestStartGeneration: RequestInit = {
             mode: 'cors',
             cache: 'no-store',
@@ -1203,6 +1205,7 @@ export class RemoteImageGenerationRequest {
                     const wsConnect = () => {
                         let ws: WebSocket | undefined = undefined;
                         let willClose = false;
+                        let hasError = false;
 
                         ws = new WebSocket(wsUrl.href)
                         
@@ -1212,7 +1215,8 @@ export class RemoteImageGenerationRequest {
                             }
                         }, 15000)
                         ws.addEventListener('error', () => {
-                            // reconnect
+                            hasError = true
+                            // reconnect on error
                             sleep(2000).then(() => {
                                 wsConnect()
                             })
@@ -1222,6 +1226,12 @@ export class RemoteImageGenerationRequest {
                                 clearInterval(heartbeatTimer)
                                 heartbeatTimer = undefined
                             }
+                            if (!willClose && !hasError) {
+                                // reconnect on abnormal disconnect
+                                sleep(2000).then(() => {
+                                    wsConnect()
+                                })
+                            }
                         })
                         ws.addEventListener('message', (event) => {
                             if (!event.data) return
@@ -1317,15 +1327,6 @@ export class RemoteImageGenerationRequest {
             }
         }
 
-        const timeout = setTimeout(() => {
-            source.close()
-            onError({
-                status: 408,
-                message:
-                    'Error: Timeout - Unable to reach Naifu servers. Please wait for a moment and try again',
-            })
-        }, 30 * 1000)
-
         const requestGetGenerationOutput: RequestInit = {
             mode: 'cors',
             cache: 'no-store',
@@ -1344,28 +1345,50 @@ export class RemoteImageGenerationRequest {
             }*/),
         }
 
-        const source = new SSE(BackendURLGetGenerateImageOutput, {
-            headers: requestGetGenerationOutput.headers,
-            payload: requestGetGenerationOutput.body
-        })
-        source.addEventListener('newImage', (message: any) => {
-            clearTimeout(timeout)
-            onImage(Buffer.from(message.data, 'base64'), message.id)
-        })
-        source.addEventListener('error', (err: any) => {
-            clearTimeout(timeout)
-            source.close()
-            onError({
-                status: err.detail.statusCode ?? 'unknown status',
-                message: err.detail.message || err.detail.error,
+        let retryTimes = 0;
+
+        const requestOutput = () => {
+            const timeout = setTimeout(() => {
+                source.close()
+                onError({
+                    status: 408,
+                    message:
+                        'Error: Timeout - Unable to reach Naifu servers. Please wait for a moment and try again',
+                })
+            }, 30 * 1000)
+
+            const source = new SSE(BackendURLGetGenerateImageOutput, {
+                headers: requestGetGenerationOutput.headers,
+                payload: requestGetGenerationOutput.body
             })
-            logWarning(err, true, 'streaming error')
-        })
-        source.addEventListener('readystatechange', (e: any) => {
-            if (source.readyState === 2) {
-                onClose()
-            }
-        })
-        source.stream()
+            source.addEventListener('newImage', (message: any) => {
+                clearTimeout(timeout)
+                onImage(Buffer.from(message.data, 'base64'), message.id)
+            })
+            source.addEventListener('error', async (err: any) => {
+                clearTimeout(timeout)
+                source.close()
+
+                logWarning(err, true, 'streaming error')
+                if (retryTimes < MAX_RETRY_TIMES) { // Should retry
+                    retryTimes ++
+                    await sleep(2000)
+                    requestOutput()
+                } else {
+                    onError({
+                        status: err.detail.statusCode ?? 'unknown status',
+                        message: err.detail.message || err.detail.error,
+                    })
+                }
+            })
+            source.addEventListener('readystatechange', (e: any) => {
+                if (source.readyState === 2) {
+                    onClose()
+                }
+            })
+            source.stream()
+        }
+
+        requestOutput()
     }
 }
diff --git a/src/styles/components/conversation.ts b/src/styles/components/conversation.ts
index fead58f..53ad9de 100644
--- a/src/styles/components/conversation.ts
+++ b/src/styles/components/conversation.ts
@@ -78,5 +78,5 @@ export const LoadingBarInner = styled.div<{ visible: boolean }>`
     background-repeat: repeat-x;
     animation: ${Gradient} 2s linear infinite;
     opacity: ${(props) => (props.visible ? '1' : '0')};
-    transition: width 250ms ease-in-out;
+    transition: width 250ms linear;
 `
\ No newline at end of file