使用WebSocket请求进度条

package.json

@@ -15,7 +15,7 @@
   "scripts": {
     "start": "concurrently -m 1 npm:start:clean npm:start:start",
     "start:clean": "rimraf .next",
-    "start:start": "cross-env NODE_ENV=development NODE_OPTIONS='--max-old-space-size=16384' next dev",
+    "start:start": "cross-env NODE_ENV=development NODE_OPTIONS='--max-old-space-size=16384' node dev-server.js",
     "build": "concurrently -m 1 npm:build:clean npm:build:build npm:build:export npm:build:postexport",
     "build:clean": "rimraf -o build .next",
     "build:build": "cross-env NODE_ENV=production NODE_OPTIONS='--max-old-space-size=8192' next build",
@@ -39,11 +39,13 @@
     "compromise": "^14.5.0",
     "dayjs": "^1.11.5",
     "diff-match-patch": "^1.0.5",
+    "express": "^4.18.2",
     "fflate": "^0.7.4",
     "file-saver": "^2.0.5",
     "framer-motion": "^7.5.1",
     "html-entities": "^2.3.3",
     "html-to-image": "^1.10.8",
+    "http-proxy-middleware": "^2.0.6",
     "idb": "^7.1.0",
     "js-cookie": "^3.0.1",
     "jszip": "^3.10.1",

src/data/request/remoterequest.ts

@@ -1175,60 +1175,145 @@ export class RemoteImageGenerationRequest {
         let res = await bind.json()
         const taskId = res.task_id
 
-        const minDelay = 2000;
-        while (true) {
-            let reqStartTime = new Date().getTime();
-            const requestGetTask: RequestInit = {
-                mode: 'cors',
-                cache: 'no-store',
-                headers: {
-                    'Content-Type': 'application/json',
-                    Authorization: 'Bearer ' + this.user.auth_token,
-                },
-                method: 'POST',
-                body: JSON.stringify({
-                    task_id: taskId
-                }),
+        const handleTaskInfoUpdate = (taskInfo: any): boolean => {
+            let progress = 0;
+
+            if (taskInfo.status === "finished") {
+                onProgress(100, 0)
+                return true
+            } else if (taskInfo.status === "error") {
+                throw new Error("Remote error")
+            } else if (taskInfo.status === "running") {
+                if (typeof taskInfo.current_step === "number" && typeof taskInfo.total_steps === "number" && taskInfo.total_steps > 0) {
+                    progress = Math.min(Math.round(taskInfo.current_step / taskInfo.total_steps * 100), 100)
+                    onProgress(progress, 0)
+                }
+            } else if (taskInfo.status === "queued") {
+                onProgress(0, taskInfo.position)
             }
+            return false
+        }
 
-            let res: any = {};
+        if ('WebSocket' in window) {
             try {
-                const bind = await fetchWithTimeout(BackendURLGetTaskInfo, requestGetTask)
-                if (!bind.ok) {
-                    logError(bind, false)
-
-                    let errorData = await bind.json()
-
-                    onError({
-                        status: bind.status ?? 500,
-                        message: errorData.error,
-                    })
+                await new Promise<void>((resolve, reject) => {
+                    const wsUrl = new URL(BackendURLGetTaskInfo, location.href.replace(/^http/, 'ws'))
+                    wsUrl.search = '?task_id=' + encodeURIComponent(taskId)
+
+                    const wsConnect = () => {
+                        let ws: WebSocket | undefined = undefined;
+                        let willClose = false;
+
+                        ws = new WebSocket(wsUrl.href)
+                        
+                        let heartbeatTimer: NodeJS.Timer | undefined = setInterval(() => {
+                            if (ws.readyState === ws.OPEN) {
+                                ws.send('ping')
+                            }
+                        }, 15000)
+                        ws.addEventListener('error', () => {
+                            // reconnect
+                            sleep(2000).then(() => {
+                                wsConnect()
+                            })
+                        })
+                        ws.addEventListener('close', () => {
+                            if (heartbeatTimer) {
+                                clearInterval(heartbeatTimer)
+                                heartbeatTimer = undefined
+                            }
+                        })
+                        ws.addEventListener('message', (event) => {
+                            if (!event.data) return
+                            try {
+                                let data = JSON.parse(event.data)
+                                if (data.error && data.error === 'ERR::TASK_NOT_FOUND') {
+                                    onError({ status: 500, message: 'Task not found, maybe lost.' })
+                                    willClose = true
+                                    ws.close()
+                                    resolve()
+                                    return
+                                }
+                                if (handleTaskInfoUpdate(data)) {
+                                    willClose = true
+                                    ws.close()
+                                    resolve()
+                                    return
+                                }
+                            } catch(err) {
+
+                            }
+                        })
+                    }
+
+                    wsConnect()
+                })
+            } catch(err) {
+                if (err.message === "Remote error") {
+                    onError({ status: 500, message: 'Internal server error occurred. Please try again later' })
                     return
+                } else {
+                    throw err;
                 }
-                res = await bind.json()
-            } catch (err) {
-                // don't stop generate when http error
             }
-            // console.log('task info', res)
+        } else {
+            const minDelay = 2000;
+            while (true) {
+                let reqStartTime = new Date().getTime();
+                const requestGetTask: RequestInit = {
+                    mode: 'cors',
+                    cache: 'no-store',
+                    headers: {
+                        'Content-Type': 'application/json',
+                        Authorization: 'Bearer ' + this.user.auth_token,
+                    },
+                    method: 'POST',
+                    body: JSON.stringify({
+                        task_id: taskId
+                    }),
+                }
 
-            let progress = 0;
-            if (res.status === "finished") {
-                onProgress(100, 0)
-                break;
-            } else if (res.status === "error") {
-                break;
-            } else if (res.status === "running") {
-                if (typeof res.current_step === "number" && typeof res.total_steps === "number" && res.total_steps > 0) {
-                    progress = Math.min(Math.round(res.current_step / res.total_steps * 100), 100)
-                    onProgress(progress, 0)
+                let res: any = {};
+                try {
+                    const bind = await fetchWithTimeout(BackendURLGetTaskInfo, requestGetTask)
+                    if (!bind.ok) {
+                        logError(bind, false)
+
+                        let errorData = await bind.json()
+
+                        if (errorData.code === 'ERR::TASK_NOT_FOUND') {
+                            onError({ status: 500, message: 'Task not found, maybe lost.' })
+                        } else {
+                            onError({
+                                status: bind.status ?? 500,
+                                message: errorData.error,
+                            })
+                        }
+                        return
+                    }
+                    res = await bind.json()
+                } catch (err) {
+                    // don't stop generate when http error
+                }
+                // console.log('task info', res)
+
+                try {
+                    if (handleTaskInfoUpdate(res)) {
+                        break
+                    }
+                } catch(err) {
+                    if (err.message === "Remote error") {
+                        onError({ status: 500, message: 'Internal server error occurred. Please try again later' })
+                        return
+                    } else {
+                        throw err;
+                    }
                 }
-            } else if (res.status === "queued") {
-                onProgress(0, res.position)
-            }
 
-            let reqEndTime = new Date().getTime();
-            if (reqEndTime - reqStartTime < minDelay) {
-                await sleep(minDelay - (reqEndTime - reqStartTime))
+                let reqEndTime = new Date().getTime();
+                if (reqEndTime - reqStartTime < minDelay) {
+                    await sleep(minDelay - (reqEndTime - reqStartTime))
+                }
             }
         }