From f881ad6ce2711fcb382350df5298908fb0b3b6f4 Mon Sep 17 00:00:00 2001
From: song_xiao_lin <song_xiao_lin@126.com>
Date: Thu, 9 May 2024 15:06:21 +0800
Subject: [PATCH] pref

---
 src/components/Contro.tsx | 16 ++++++++++------
 1 file changed, 10 insertions(+), 6 deletions(-)

diff --git a/src/components/Contro.tsx b/src/components/Contro.tsx
index 2177b8d..0b0d799 100644
--- a/src/components/Contro.tsx
+++ b/src/components/Contro.tsx
@@ -38,15 +38,15 @@ export const Contro: React.FC<ControProps> = () => {
 
     wsc.onWSCMessage((message) => {
       if (message.action === ActionType.STATUS) {
-        setConnected(message.connected);
         if (message.connected === false) {
-          localStorage.setItem("yakit-connect", "");
-          setAutoFindFailedReason("Yakit WebSocket Controller Connect Fail");
+          handleConnectFail();
         } else {
           localStorage.setItem(
             "yakit-connect",
             JSON.stringify({ connected: true, port: message.port })
           );
+          setEnginePort(message.port + "");
+          setConnected(true);
           setAutoFindFailedReason("");
         }
       }
@@ -57,9 +57,7 @@ export const Contro: React.FC<ControProps> = () => {
     const ws = new WebSocket(`ws://127.0.0.1:${port}`);
     ws.onclose = (e: CloseEvent) => {
       if (enginePortRef.current) {
-        setConnected(false);
-        setAutoFindFailedReason("Yakit WebSocket Controller Connect Fail");
-        localStorage.setItem("yakit-connect", "");
+        handleConnectFail();
         return;
       }
       if (e.reason !== `FoundYakitWebSocketController` && port + 1 <= max) {
@@ -81,6 +79,12 @@ export const Contro: React.FC<ControProps> = () => {
     };
   };
 
+  const handleConnectFail = () => {
+    setConnected(false);
+    setAutoFindFailedReason("Yakit WebSocket Controller Connect Fail");
+    localStorage.setItem("yakit-connect", "");
+  };
+
   const connectPort = (port: number) => {
     setAutoFindFailedReason("");
     wsc.connect(port);