From f247ebfbe19a9eab6870cb88dc7f7cd67c56308b Mon Sep 17 00:00:00 2001
From: zhsama <torvalds@linux.do>
Date: Thu, 15 Jan 2026 17:53:28 +0800
Subject: [PATCH] feat: Await sub-graph save before syncing workflow draft

---
 .../sub-graph/components/sub-graph-main.tsx   | 14 ++++++-------
 .../tool/components/sub-graph-modal/index.tsx | 21 +++++++++----------
 2 files changed, 17 insertions(+), 18 deletions(-)

diff --git a/web/app/components/sub-graph/components/sub-graph-main.tsx b/web/app/components/sub-graph/components/sub-graph-main.tsx
index 2dc93f881e..44e30b6aee 100644
--- a/web/app/components/sub-graph/components/sub-graph-main.tsx
+++ b/web/app/components/sub-graph/components/sub-graph-main.tsx
@@ -51,21 +51,21 @@ const SubGraphMain: FC<SubGraphMainProps> = ({
     flowId,
   })
 
-  const handleSyncSubGraphDraft = useCallback(() => {
+  const handleSyncSubGraphDraft = useCallback(async () => {
     const { getNodes, edges } = reactFlowStore.getState()
-    onSave?.(getNodes() as Node[], edges as Edge[])
+    await onSave?.(getNodes() as Node[], edges as Edge[])
   }, [onSave, reactFlowStore])
 
   const handleSyncWorkflowDraft = useCallback(async (
     notRefreshWhenSyncError?: boolean,
     callback?: SyncWorkflowDraftCallback,
   ) => {
-    handleSyncSubGraphDraft()
-    if (onSyncWorkflowDraft) {
-      await onSyncWorkflowDraft(notRefreshWhenSyncError, callback)
-      return
-    }
     try {
+      await handleSyncSubGraphDraft()
+      if (onSyncWorkflowDraft) {
+        await onSyncWorkflowDraft(notRefreshWhenSyncError, callback)
+        return
+      }
       callback?.onSuccess?.()
     }
     catch {
diff --git a/web/app/components/workflow/nodes/tool/components/sub-graph-modal/index.tsx b/web/app/components/workflow/nodes/tool/components/sub-graph-modal/index.tsx
index 48e000c414..c1d2bcbea4 100644
--- a/web/app/components/workflow/nodes/tool/components/sub-graph-modal/index.tsx
+++ b/web/app/components/workflow/nodes/tool/components/sub-graph-modal/index.tsx
@@ -4,7 +4,7 @@ import type { SubGraphModalProps } from './types'
 import type { MentionConfig } from '@/app/components/workflow/nodes/_base/types'
 import type { LLMNodeType } from '@/app/components/workflow/nodes/llm/types'
 import type { ToolNodeType } from '@/app/components/workflow/nodes/tool/types'
-import type { Node, PromptItem } from '@/app/components/workflow/types'
+import type { Node, PromptItem, PromptTemplateItem } from '@/app/components/workflow/types'
 import { Dialog, DialogPanel, Transition, TransitionChild } from '@headlessui/react'
 import { RiCloseLine } from '@remixicon/react'
 import { noop } from 'es-toolkit/function'
@@ -16,7 +16,7 @@ import { useIsChatMode, useNodesSyncDraft, useWorkflow, useWorkflowVariables } f
 import { useHooksStore } from '@/app/components/workflow/hooks-store'
 import { VarKindType } from '@/app/components/workflow/nodes/_base/types'
 import { useStore as useWorkflowStore } from '@/app/components/workflow/store'
-import { BlockEnum, EditionType, PromptRole } from '@/app/components/workflow/types'
+import { BlockEnum, EditionType, isPromptMessageContext, PromptRole } from '@/app/components/workflow/types'
 import SubGraphCanvas from './sub-graph-canvas'
 
 const SubGraphModal: FC<SubGraphModalProps> = ({
@@ -129,7 +129,7 @@ const SubGraphModal: FC<SubGraphModalProps> = ({
       handleMentionConfigChange(mentionConfig)
   }, [handleMentionConfigChange, mentionConfig, toolParam])
 
-  const getUserPromptText = useCallback((promptTemplate?: PromptItem[] | PromptItem) => {
+  const getUserPromptText = useCallback((promptTemplate?: PromptTemplateItem[] | PromptItem) => {
     if (!promptTemplate)
       return ''
     const resolveText = (item?: PromptItem) => {
@@ -140,17 +140,18 @@ const SubGraphModal: FC<SubGraphModalProps> = ({
       return item.text || ''
     }
     if (Array.isArray(promptTemplate)) {
-      const userPrompt = promptTemplate.find(item => item.role === PromptRole.user)
-      if (userPrompt)
-        return resolveText(userPrompt)
+      for (const item of promptTemplate) {
+        if (!isPromptMessageContext(item) && item.role === PromptRole.user)
+          return resolveText(item)
+      }
       return ''
     }
     return resolveText(promptTemplate)
   }, [])
 
   // TODO: handle external workflow updates while sub-graph modal is open.
-  const handleSave = useCallback((subGraphNodes: any[], _edges: any[]) => {
-    const extractorNodeData = subGraphNodes.find(node => node.id === extractorNodeId)
+  const handleSave = useCallback((subGraphNodes: Node[]) => {
+    const extractorNodeData = subGraphNodes.find(node => node.id === extractorNodeId) as Node<LLMNodeType> | undefined
     if (!extractorNodeData)
       return
 
@@ -193,10 +194,8 @@ const SubGraphModal: FC<SubGraphModalProps> = ({
       return node
     })
     setNodes(nextNodes)
-    // Trigger main graph draft sync to persist changes to backend
-    handleSyncWorkflowDraft()
     setControlPromptEditorRerenderKey(Date.now())
-  }, [agentNodeId, extractorNodeId, getUserPromptText, handleSyncWorkflowDraft, paramKey, reactflowStore, setControlPromptEditorRerenderKey, toolNodeId])
+  }, [agentNodeId, extractorNodeId, getUserPromptText, paramKey, reactflowStore, setControlPromptEditorRerenderKey, toolNodeId])
 
   return (
     <Transition appear show={isOpen} as={Fragment}>