From c95281906757663a8c80420f1f1c5e4c701b2f51 Mon Sep 17 00:00:00 2001
From: liangfung <1098486429@qq.com>
Date: Tue, 21 Nov 2023 23:10:46 +0800
Subject: [PATCH] refactor: useMergedWorkers

---
 ee/tabby-ui/lib/hooks/use-remote-worker.ts | 34 +++++++++++++---------
 1 file changed, 21 insertions(+), 13 deletions(-)

diff --git a/ee/tabby-ui/lib/hooks/use-remote-worker.ts b/ee/tabby-ui/lib/hooks/use-remote-worker.ts
index 73e3011..fcfda5c 100644
--- a/ee/tabby-ui/lib/hooks/use-remote-worker.ts
+++ b/ee/tabby-ui/lib/hooks/use-remote-worker.ts
@@ -1,4 +1,4 @@
-import { groupBy, findIndex } from 'lodash-es'
+import { groupBy, findIndex, slice } from 'lodash-es'
 import {
   GetWorkersDocument,
   Worker,
@@ -6,6 +6,7 @@ import {
 } from '@/lib/gql/generates/graphql'
 import { useGraphQL } from './use-graphql'
 import type { HealthInfo } from './use-health'
+import React from 'react'
 
 function useRemoteWorkers() {
   return useGraphQL(GetWorkersDocument)
@@ -33,20 +34,27 @@ function transformHealthInfoToWorker(
 
 function useMergedWorkers(healthInfo: HealthInfo | undefined) {
   const { data } = useRemoteWorkers()
-  let workers = data?.workers || []
+  let workers = data?.workers
 
-  const haveRemoteCompletionWorkers =
-    findIndex(workers, { kind: WorkerKind.Completion }) > -1
-  const haveRemoteChatWorkers =
-    findIndex(workers, { kind: WorkerKind.Chat }) > -1
+  const groupedWorkers = React.useMemo(() => {
+    const _workers = slice(workers)
+    const haveRemoteCompletionWorkers =
+      findIndex(workers, { kind: WorkerKind.Completion }) > -1
+    const haveRemoteChatWorkers =
+      findIndex(workers, { kind: WorkerKind.Chat }) > -1
 
-  if (!haveRemoteCompletionWorkers && healthInfo?.model) {
-    workers.push(transformHealthInfoToWorker(healthInfo, WorkerKind.Completion))
-  }
-  if (!haveRemoteChatWorkers && healthInfo?.chat_model) {
-    workers.push(transformHealthInfoToWorker(healthInfo, WorkerKind.Chat))
-  }
-  return groupBy(workers, worker => worker.kind)
+    if (!haveRemoteCompletionWorkers && healthInfo?.model) {
+      _workers.push(
+        transformHealthInfoToWorker(healthInfo, WorkerKind.Completion)
+      )
+    }
+    if (!haveRemoteChatWorkers && healthInfo?.chat_model) {
+      _workers.push(transformHealthInfoToWorker(healthInfo, WorkerKind.Chat))
+    }
+    return groupBy(_workers, 'kind')
+  }, [healthInfo, workers])
+
+  return groupedWorkers
 }
 
 export { useRemoteWorkers, useMergedWorkers }