fix: stopping status workload (#394)

Author: peppescg

renderer/src/features/mcp-servers/components/actions-mcp-server.tsx

@@ -32,8 +32,8 @@ export function ActionsMcpServer({
       <div onClick={(e) => e.preventDefault()}>
         <Switch
           aria-label="Mutate server"
-          checked={isRunning || isPending}
-          disabled={isStarting}
+          checked={isRunning || isStarting}
+          disabled={isStarting || isPending}
           onCheckedChange={() => mutate()}
         />
       </div>

renderer/src/features/mcp-servers/hooks/use-mutation-stop-server.ts

@@ -4,10 +4,11 @@ import type {
 } from '@/common/api/generated'
 import {
   postApiV1BetaWorkloadsByNameStopMutation,
-  getApiV1BetaWorkloadsByNameQueryKey,
   getApiV1BetaWorkloadsQueryKey,
+  getApiV1BetaWorkloadsByNameOptions,
 } from '@/common/api/generated/@tanstack/react-query.gen'
 import { useToastMutation } from '@/common/hooks/use-toast-mutation'
+import { pollBatchServerStatus } from '@/common/lib/polling'
 import { useQueryClient } from '@tanstack/react-query'
 
 const getMutationData = (name: string) => ({
@@ -20,9 +21,6 @@ const getMutationData = (name: string) => ({
 export function useMutationStopServerList({ name }: { name: string }) {
   const queryClient = useQueryClient()
   const queryKey = getApiV1BetaWorkloadsQueryKey({ query: { all: true } })
-  const serverQueryKey = getApiV1BetaWorkloadsByNameQueryKey({
-    path: { name },
-  })
   return useToastMutation({
     ...getMutationData(name),
 
@@ -31,18 +29,7 @@ export function useMutationStopServerList({ name }: { name: string }) {
         queryKey,
       })
 
-      const previousServer = queryClient.getQueryData(serverQueryKey)
-
-      queryClient.setQueryData(
-        serverQueryKey,
-        (oldData: WorkloadsWorkload | undefined) => {
-          if (!oldData) return oldData
-          return {
-            ...oldData,
-            status: 'stopping',
-          }
-        }
-      )
+      const previousServersList = queryClient.getQueryData(queryKey)
 
       queryClient.setQueryData(
         queryKey,
@@ -59,79 +46,30 @@ export function useMutationStopServerList({ name }: { name: string }) {
         }
       )
 
-      return { previousServer }
+      return { previousServersList }
     },
-    onSuccess: () => {
-      queryClient.setQueryData(
-        queryKey,
-        (oldData: V1WorkloadListResponse | undefined) => {
-          if (!oldData) return oldData
-          return { ...oldData, status: 'stopped' }
-        }
-      )
-      queryClient.setQueryData(
-        serverQueryKey,
-        (oldData: WorkloadsWorkload | undefined) => {
-          if (!oldData) return oldData
-          return { ...oldData, status: 'stopped' }
-        }
+    onSuccess: async () => {
+      // Poll until server stopped
+      await pollBatchServerStatus(
+        async (names) => {
+          const servers = await Promise.all(
+            names.map((name) =>
+              queryClient.fetchQuery(
+                getApiV1BetaWorkloadsByNameOptions({ path: { name } })
+              )
+            )
+          )
+          return servers
+        },
+        [name],
+        'stopped'
       )
+      queryClient.invalidateQueries({ queryKey })
     },
     onError: (_error, _variables, context) => {
-      if (context?.previousServer) {
-        queryClient.setQueryData(queryKey, context.previousServer)
+      if (context?.previousServersList) {
+        queryClient.setQueryData(queryKey, context.previousServersList)
       }
     },
   })
 }
-
-export function useMutationStopServer({ name }: { name: string }) {
-  const queryClient = useQueryClient()
-  const queryKey = getApiV1BetaWorkloadsByNameQueryKey({ path: { name } })
-
-  return useToastMutation({
-    ...getMutationData(name),
-
-    onMutate: async () => {
-      await queryClient.cancelQueries({ queryKey })
-
-      const previousServerData = queryClient.getQueryData(queryKey)
-
-      queryClient.setQueryData(
-        queryKey,
-        (oldData: WorkloadsWorkload | undefined) => {
-          if (!oldData) return oldData
-
-          const updatedData = {
-            ...oldData,
-            status: 'stopping',
-          } as WorkloadsWorkload
-          return updatedData
-        }
-      )
-
-      return { previousServerData }
-    },
-
-    onError: (_error, _variables, context) => {
-      if (context?.previousServerData) {
-        queryClient.setQueryData(queryKey, context.previousServerData)
-      }
-    },
-    onSuccess: () => {
-      queryClient.setQueryData(
-        queryKey,
-        (oldData: WorkloadsWorkload | undefined) => {
-          if (!oldData) return oldData
-
-          return { ...oldData, status: 'stopped' }
-        }
-      )
-    },
-    onSettled: () => {
-      queryClient.invalidateQueries({
-        queryKey: getApiV1BetaWorkloadsQueryKey({ query: { all: true } }),
-      })
-    },
-  })
-}