fix: dismiss toast in case of shutting down (#636)

* fix: dismiss toast in case of shutting down

* fix : wrap callback into useEffect

* update test

Author: peppescg

renderer/src/common/hooks/use-toast-mutation.ts

@@ -17,11 +17,13 @@ export function useToastMutation<
   successMsg,
   errorMsg,
   loadingMsg,
+  toastId,
   ...options
 }: UseMutationOptions<TData, TError, TVariables, TContext> & {
   successMsg?: ((variables: TVariables) => string) | string
   loadingMsg?: string
   errorMsg?: string
+  toastId?: string
 }) {
   const {
     mutateAsync: originalMutateAsync,
@@ -60,9 +62,10 @@ export function useToastMutation<
 
           return 'An error occurred'
         },
+        ...(toastId ? { id: toastId } : {}),
       })
     },
-    [errorMsg, loadingMsg, originalMutateAsync, successMsg]
+    [errorMsg, loadingMsg, originalMutateAsync, successMsg, toastId]
   )
 
   return { mutateAsync, ...rest }

renderer/src/features/mcp-servers/hooks/__tests__/use-mutation-restart-server.test.ts

@@ -10,14 +10,27 @@ import type { WorkloadsWorkload } from '@/common/api/generated'
 import { server } from '@/common/mocks/node'
 import { http, HttpResponse } from 'msw'
 import { mswEndpoint } from '@/common/mocks/msw-endpoint'
+import { toast } from 'sonner'
+
+vi.mock('sonner', () => ({
+  toast: {
+    dismiss: vi.fn(),
+    promise: vi.fn((promise) => promise.catch(() => {})),
+    success: vi.fn(),
+    error: vi.fn(),
+    loading: vi.fn(),
+  },
+}))
 
 // Mock electron API
+const mockOnServerShutdown = vi.fn()
 Object.defineProperty(window, 'electronAPI', {
   value: {
     shutdownStore: {
       getLastShutdownServers: vi.fn(),
       clearShutdownHistory: vi.fn(),
     },
+    onServerShutdown: mockOnServerShutdown,
   },
   writable: true,
 })
@@ -44,6 +57,7 @@ const createQueryClientWrapper = () => {
 
 beforeEach(() => {
   vi.clearAllMocks()
+  mockOnServerShutdown.mockClear()
   window.electronAPI.shutdownStore.getLastShutdownServers = vi
     .fn()
     .mockResolvedValue([])
@@ -127,6 +141,24 @@ describe('useMutationRestartServerAtStartup', () => {
       expect(result.current.isError).toBe(true)
     })
   })
+
+  it('dismisses toast on server shutdown', async () => {
+    const { Wrapper } = createQueryClientWrapper()
+
+    renderHook(() => useMutationRestartServerAtStartup(), {
+      wrapper: Wrapper,
+    })
+
+    expect(mockOnServerShutdown).toHaveBeenCalledTimes(1)
+    const shutdownCallback = mockOnServerShutdown.mock.calls[0]?.[0]
+    expect(shutdownCallback).toBeDefined()
+
+    shutdownCallback?.()
+
+    expect(vi.mocked(toast.dismiss)).toHaveBeenCalledWith(
+      'restart-servers-startup'
+    )
+  })
 })
 
 describe('useMutationRestartServer', () => {

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

@@ -11,6 +11,10 @@ import {
 import { useToastMutation } from '@/common/hooks/use-toast-mutation'
 import { pollBatchServerStatus } from '@/common/lib/polling'
 import { useQueryClient } from '@tanstack/react-query'
+import { useEffect } from 'react'
+import { toast } from 'sonner'
+
+const TOAST_ID = 'restart-servers-startup'
 
 const getMutationData = (name: string) => ({
   ...postApiV1BetaWorkloadsByNameRestartMutation(),
@@ -23,10 +27,19 @@ export function useMutationRestartServerAtStartup() {
   const queryClient = useQueryClient()
   const queryKey = getApiV1BetaWorkloadsQueryKey({ query: { all: true } })
 
+  useEffect(() => {
+    const cleanup = window.electronAPI.onServerShutdown(() => {
+      toast.dismiss(TOAST_ID)
+    })
+
+    return cleanup
+  }, [])
+
   return useToastMutation({
     successMsg: 'Servers restarted successfully',
     errorMsg: 'Failed to restart servers',
     loadingMsg: 'Restarting servers...',
+    toastId: TOAST_ID,
     ...postApiV1BetaWorkloadsRestartMutation(),
     onMutate: async (variables) => {
       await queryClient.cancelQueries({ queryKey })

renderer/src/routes/__tests__/index.test.tsx

@@ -16,6 +16,7 @@ Object.defineProperty(window, 'electronAPI', {
       getLastShutdownServers: vi.fn().mockResolvedValue([]),
       clearShutdownHistory: vi.fn().mockResolvedValue(undefined),
     },
+    onServerShutdown: vi.fn().mockReturnValue(() => {}),
   },
   writable: true,
 })