From 8a48db6d0dc6e92b06a538c6b7b27e50c372692e Mon Sep 17 00:00:00 2001
From: lyzno1 <yuanyouhuilyz@gmail.com>
Date: Thu, 30 Oct 2025 15:04:14 +0800
Subject: [PATCH] Improve workflow tool install flow

---
 .../workflow/block-selector/tool-picker.tsx   | 16 +++++-
 .../components/install-plugin-button.tsx      | 55 +++++++++++++++++--
 2 files changed, 64 insertions(+), 7 deletions(-)

diff --git a/web/app/components/workflow/block-selector/tool-picker.tsx b/web/app/components/workflow/block-selector/tool-picker.tsx
index c438e45042..660cdad71e 100644
--- a/web/app/components/workflow/block-selector/tool-picker.tsx
+++ b/web/app/components/workflow/block-selector/tool-picker.tsx
@@ -23,7 +23,16 @@ import {
 } from '@/service/tools'
 import type { CustomCollectionBackend } from '@/app/components/tools/types'
 import Toast from '@/app/components/base/toast'
-import { useAllBuiltInTools, useAllCustomTools, useAllMCPTools, useAllWorkflowTools, useInvalidateAllBuiltInTools, useInvalidateAllCustomTools } from '@/service/use-tools'
+import {
+  useAllBuiltInTools,
+  useAllCustomTools,
+  useAllMCPTools,
+  useAllWorkflowTools,
+  useInvalidateAllBuiltInTools,
+  useInvalidateAllCustomTools,
+  useInvalidateAllMCPTools,
+  useInvalidateAllWorkflowTools,
+} from '@/service/use-tools'
 import { useFeaturedToolsRecommendations } from '@/service/use-plugins'
 import { useGlobalPublicStore } from '@/context/global-public-context'
 import cn from '@/utils/classnames'
@@ -70,6 +79,8 @@ const ToolPicker: FC<Props> = ({
   const { data: workflowTools } = useAllWorkflowTools()
   const { data: mcpTools } = useAllMCPTools()
   const invalidateBuiltInTools = useInvalidateAllBuiltInTools()
+  const invalidateWorkflowTools = useInvalidateAllWorkflowTools()
+  const invalidateMcpTools = useInvalidateAllMCPTools()
 
   const {
     plugins: featuredPlugins = [],
@@ -193,6 +204,9 @@ const ToolPicker: FC<Props> = ({
             showFeatured={scope === 'all' && enable_marketplace}
             onFeaturedInstallSuccess={async () => {
               invalidateBuiltInTools()
+              invalidateCustomTools()
+              invalidateWorkflowTools()
+              invalidateMcpTools()
             }}
           />
         </div>
diff --git a/web/app/components/workflow/nodes/_base/components/install-plugin-button.tsx b/web/app/components/workflow/nodes/_base/components/install-plugin-button.tsx
index c387cb8630..b0d878d53d 100644
--- a/web/app/components/workflow/nodes/_base/components/install-plugin-button.tsx
+++ b/web/app/components/workflow/nodes/_base/components/install-plugin-button.tsx
@@ -1,8 +1,11 @@
 import Button from '@/app/components/base/button'
 import { RiInstallLine, RiLoader2Line } from '@remixicon/react'
 import type { ComponentProps, MouseEventHandler } from 'react'
+import { useState } from 'react'
 import classNames from '@/utils/classnames'
 import { useTranslation } from 'react-i18next'
+import checkTaskStatus from '@/app/components/plugins/install-plugin/base/check-task-status'
+import { TaskStatus } from '@/app/components/plugins/types'
 import { useCheckInstalled, useInstallPackageFromMarketPlace } from '@/service/use-plugins'
 
 type InstallPluginButtonProps = Omit<ComponentProps<typeof Button>, 'children' | 'loading'> & {
@@ -28,15 +31,55 @@ export const InstallPluginButton = (props: InstallPluginButtonProps) => {
     enabled: identifiers.length > 0,
   })
   const install = useInstallPackageFromMarketPlace()
-  const isLoading = manifest.isLoading || install.isPending
-  // await for refetch to get the new installed plugin, when manifest refetch, this component will unmount
-  || install.isSuccess
+  const [isTracking, setIsTracking] = useState(false)
+  const isLoading = manifest.isLoading || install.isPending || isTracking
   const handleInstall: MouseEventHandler = (e) => {
     e.stopPropagation()
+    if (isLoading)
+      return
+    setIsTracking(true)
     install.mutate(uniqueIdentifier, {
-      onSuccess: async () => {
-        await manifest.refetch()
-        onSuccess?.()
+      onSuccess: async (response) => {
+        const finish = async () => {
+          await manifest.refetch()
+          onSuccess?.()
+          setIsTracking(false)
+          install.reset()
+        }
+
+        if (!response) {
+          await finish()
+          return
+        }
+
+        if (response.all_installed) {
+          await finish()
+          return
+        }
+
+        const { check } = checkTaskStatus()
+        try {
+          const { status } = await check({
+            taskId: response.task_id,
+            pluginUniqueIdentifier: uniqueIdentifier,
+          })
+
+          if (status === TaskStatus.failed) {
+            setIsTracking(false)
+            install.reset()
+            return
+          }
+
+          await finish()
+        }
+        catch {
+          setIsTracking(false)
+          install.reset()
+        }
+      },
+      onError: () => {
+        setIsTracking(false)
+        install.reset()
       },
     })
   }