From 8832f08fed28aa1795e42944c596fc64be3e8a1d Mon Sep 17 00:00:00 2001
From: jZonG <jzongcode@gmail.com>
Date: Wed, 28 May 2025 15:35:44 +0800
Subject: [PATCH] authorizing redirection

---
 web/app/components/tools/mcp/create-card.tsx  |  4 +-
 .../components/tools/mcp/detail/content.tsx   | 20 ++++----
 web/app/components/tools/mcp/index.tsx        | 47 ++++++++++++++++---
 web/app/components/tools/provider-list.tsx    |  7 ++-
 web/service/use-tools.ts                      | 15 +++++-
 5 files changed, 72 insertions(+), 21 deletions(-)

diff --git a/web/app/components/tools/mcp/create-card.tsx b/web/app/components/tools/mcp/create-card.tsx
index 511a724831..afb45b0834 100644
--- a/web/app/components/tools/mcp/create-card.tsx
+++ b/web/app/components/tools/mcp/create-card.tsx
@@ -34,8 +34,8 @@ const NewMCPCard = ({ handleCreate }: Props) => {
   const linkUrl = useMemo(() => {
     // TODO help link
     if (language.startsWith('zh_'))
-      return 'https://docs.dify.ai/zh-hans/guides/tools#ru-he-chuang-jian-zi-ding-yi-gong-ju'
-    return 'https://docs.dify.ai/en/guides/tools#how-to-create-custom-tools'
+      return 'https://docs.dify.ai/zh-hans/guides/tools/integrate-tool/mcp'
+    return 'https://docs.dify.ai/en/guides/tools/integrate-tool/mcp'
   }, [language])
 
   const [showModal, setShowModal] = useState(false)
diff --git a/web/app/components/tools/mcp/detail/content.tsx b/web/app/components/tools/mcp/detail/content.tsx
index 5d2f5cb513..05e2572170 100644
--- a/web/app/components/tools/mcp/detail/content.tsx
+++ b/web/app/components/tools/mcp/detail/content.tsx
@@ -1,5 +1,6 @@
 'use client'
-import React, { useCallback, useState } from 'react'
+import React, { useCallback } from 'react'
+import { useRouter } from 'next/navigation'
 import type { FC } from 'react'
 import { useBoolean } from 'ahooks'
 import { useTranslation } from 'react-i18next'
@@ -41,9 +42,10 @@ const MCPDetailContent: FC<Props> = ({
   onHide,
 }) => {
   const { t } = useTranslation()
+  const router = useRouter()
   const { isCurrentWorkspaceManager } = useAppContext()
 
-  const { data, isPending: isGettingTools } = useMCPTools(detail.is_team_authorization ? detail.id : '')
+  const { data, isFetching: isGettingTools } = useMCPTools(detail.is_team_authorization ? detail.id : '')
   const invalidateMCPTools = useInvalidateMCPTools()
   const { mutateAsync: updateTools, isPending: isUpdating } = useUpdateMCPTools()
   const { mutateAsync: authorizeMcp, isPending: isAuthorizing } = useAuthorizeMCP()
@@ -54,6 +56,7 @@ const MCPDetailContent: FC<Props> = ({
       return
     await updateTools(detail.id)
     invalidateMCPTools(detail.id)
+    onUpdate()
   }, [detail, updateTools])
 
   const { mutate: updateMCP } = useUpdateMCP({
@@ -85,11 +88,11 @@ const MCPDetailContent: FC<Props> = ({
       provider_id: detail.id,
       server_url: detail.server_url!,
     })
-    // TODO
-    if ((res as any)?.result === 'success') {
-      hideUpdateModal()
-      onUpdate()
-    }
+    if (res.result === 'success')
+      handleUpdateTools()
+
+    else if (res.authorization_url)
+      router.push(res.authorization_url)
   }, [detail, updateMCP, hideUpdateModal, onUpdate])
 
   const handleUpdate = useCallback(async (data: any) => {
@@ -117,8 +120,6 @@ const MCPDetailContent: FC<Props> = ({
     }
   }, [detail, showDeleting, hideDeleting, hideDeleteConfirm, onUpdate])
 
-  const [loading, setLoading] = useState(false)
-
   if (!detail)
     return null
 
@@ -150,7 +151,6 @@ const MCPDetailContent: FC<Props> = ({
             <Button
               variant='secondary'
               className='w-full'
-              // onClick={() => setShowSettingAuth(true)}
               disabled={!isCurrentWorkspaceManager}
             >
               <Indicator className='mr-2' color={'green'} />
diff --git a/web/app/components/tools/mcp/index.tsx b/web/app/components/tools/mcp/index.tsx
index 6b57129b8a..729117800f 100644
--- a/web/app/components/tools/mcp/index.tsx
+++ b/web/app/components/tools/mcp/index.tsx
@@ -1,9 +1,16 @@
 'use client'
-import { useMemo, useState } from 'react'
+import { useEffect, useMemo, useState } from 'react'
+import { useSearchParams } from 'next/navigation'
 import NewMCPCard from './create-card'
 import MCPCard from './provider-card'
 import MCPDetailPanel from './detail/provider-detail'
-import { useAllMCPTools, useAuthorizeMCP, useInvalidateAllMCPTools, useInvalidateMCPTools, useUpdateMCPTools } from '@/service/use-tools'
+import {
+  useAllMCPTools,
+  useAuthorizeMCP,
+  useInvalidateMCPTools,
+  useUpdateMCPAuthorizationToken,
+  useUpdateMCPTools,
+} from '@/service/use-tools'
 import type { ToolWithProvider } from '@/app/components/workflow/types'
 import cn from '@/utils/classnames'
 
@@ -32,11 +39,15 @@ function renderDefaultCard() {
 const MCPList = ({
   searchText,
 }: Props) => {
-  const { data: list = [] } = useAllMCPTools()
-  const invalidateMCPList = useInvalidateAllMCPTools()
+  const searchParams = useSearchParams()
+  const authCode = searchParams.get('code') || ''
+  const providerID = decodeURIComponent(searchParams.get('state') || '').split('provider_id=')[1] || ''
+
+  const { data: list = [], refetch } = useAllMCPTools()
   const { mutateAsync: authorizeMcp } = useAuthorizeMCP()
   const { mutateAsync: updateTools } = useUpdateMCPTools()
   const invalidateMCPTools = useInvalidateMCPTools()
+  const { mutateAsync: updateMCPAuthorizationToken } = useUpdateMCPAuthorizationToken()
 
   const filteredList = useMemo(() => {
     return list.filter((collection) => {
@@ -53,16 +64,38 @@ const MCPList = ({
   }, [list, currentProviderID])
 
   const handleCreate = async (provider: ToolWithProvider) => {
-    invalidateMCPList()
+    await refetch() // update list
     setCurrentProviderID(provider.id)
     await authorizeMcp({
       provider_id: provider.id,
       server_url: provider.server_url!,
     })
+    await refetch() // update authorization in list
     await updateTools(provider.id)
     invalidateMCPTools(provider.id)
+    await refetch() // update tool list in provider list
   }
 
+  const handleUpdateAuthorization = async (providerID: string, code: string) => {
+    const targetProvider = list.find(provider => provider.id === providerID)
+    if (!targetProvider) return
+    await updateMCPAuthorizationToken({
+      provider_id: providerID,
+      server_url: targetProvider.server_url!,
+      authorization_code: code,
+    })
+    await refetch()
+    setCurrentProviderID(providerID)
+    await updateTools(providerID)
+    invalidateMCPTools(providerID)
+    await refetch()
+  }
+
+  useEffect(() => {
+    if (authCode && providerID)
+      handleUpdateAuthorization(providerID, authCode)
+  }, [authCode, providerID])
+
   return (
     <>
       <div
@@ -78,7 +111,7 @@ const MCPList = ({
             data={provider}
             currentProvider={currentProvider}
             handleSelect={setCurrentProviderID}
-            onUpdate={() => invalidateMCPList()}
+            onUpdate={refetch}
           />
         ))}
         {!list.length && renderDefaultCard()}
@@ -87,7 +120,7 @@ const MCPList = ({
         <MCPDetailPanel
           detail={currentProvider}
           onHide={() => setCurrentProviderID(undefined)}
-          onUpdate={() => invalidateMCPList()}
+          onUpdate={refetch}
         />
       )}
     </>
diff --git a/web/app/components/tools/provider-list.tsx b/web/app/components/tools/provider-list.tsx
index 7a0f458ff5..647271e840 100644
--- a/web/app/components/tools/provider-list.tsx
+++ b/web/app/components/tools/provider-list.tsx
@@ -1,6 +1,7 @@
 'use client'
 import { useMemo, useRef, useState } from 'react'
 import { useTranslation } from 'react-i18next'
+import { useSearchParams } from 'next/navigation'
 import type { Collection } from './types'
 import Marketplace from './marketplace'
 import cn from '@/utils/classnames'
@@ -25,8 +26,12 @@ const ProviderList = () => {
   const containerRef = useRef<HTMLDivElement>(null)
   const { enable_marketplace } = useAppContextSelector(s => s.systemFeatures)
 
+  const searchParams = useSearchParams()
+  const authCode = searchParams.get('code') || ''
+  const providerID = decodeURIComponent(searchParams.get('state') || '').split('provider_id=')[1] || ''
+
   const [activeTab, setActiveTab] = useTabSearchParams({
-    defaultTab: 'builtin',
+    defaultTab: authCode && providerID ? 'mcp' : 'builtin',
   })
   const options = [
     { value: 'builtin', text: t('tools.type.builtIn') },
diff --git a/web/service/use-tools.ts b/web/service/use-tools.ts
index b0caeff044..6c50112fa0 100644
--- a/web/service/use-tools.ts
+++ b/web/service/use-tools.ts
@@ -141,13 +141,26 @@ export const useAuthorizeMCP = () => {
   return useMutation({
     mutationKey: [NAME_SPACE, 'authorize-mcp'],
     mutationFn: (payload: { provider_id: string; server_url: string }) => {
-      return post('/workspaces/current/tool-provider/mcp/auth', {
+      return post<{ result?: string; authorization_url?: string }>('/workspaces/current/tool-provider/mcp/auth', {
         body: payload,
       })
     },
   })
 }
 
+export const useUpdateMCPAuthorizationToken = () => {
+  return useMutation({
+    mutationKey: [NAME_SPACE, 'refresh-mcp-server-code'],
+    mutationFn: (payload: { provider_id: string; server_url: string; authorization_code: string }) => {
+      return get<MCPServerDetail>('/workspaces/current/tool-provider/mcp/token', {
+        params: {
+          ...payload,
+        },
+      })
+    },
+  })
+}
+
 export const useMCPTools = (providerID: string) => {
   return useQuery({
     enabled: !!providerID,