From 5200dcfef68982aff63ecf2d2fc5cb2949cf1e18 Mon Sep 17 00:00:00 2001
From: "saica.go" <zephyird@gmail.com>
Date: Mon, 12 May 2025 21:48:53 +0800
Subject: [PATCH] feat: add undo functionality to agent prompt generation
 (#5821)

* feat: add undo functionality to agent prompt generation

* feat: add a standalone undo button
---
 .../pages/agents/components/AddAgentPopup.tsx | 22 ++++++++++++++++---
 1 file changed, 19 insertions(+), 3 deletions(-)

diff --git a/src/renderer/src/pages/agents/components/AddAgentPopup.tsx b/src/renderer/src/pages/agents/components/AddAgentPopup.tsx
index 05f9dc872e..eeca7a39c8 100644
--- a/src/renderer/src/pages/agents/components/AddAgentPopup.tsx
+++ b/src/renderer/src/pages/agents/components/AddAgentPopup.tsx
@@ -1,6 +1,6 @@
 import 'emoji-picker-element'
 
-import { CheckOutlined, LoadingOutlined, ThunderboltOutlined } from '@ant-design/icons'
+import { CheckOutlined, LoadingOutlined, ThunderboltOutlined, RollbackOutlined } from '@ant-design/icons'
 import EmojiPicker from '@renderer/components/EmojiPicker'
 import { TopView } from '@renderer/components/TopView'
 import { AGENT_PROMPT } from '@renderer/config/prompts'
@@ -38,6 +38,8 @@ const PopupContainer: React.FC<Props> = ({ resolve }) => {
   const formRef = useRef<FormInstance>(null)
   const [emoji, setEmoji] = useState('')
   const [loading, setLoading] = useState(false)
+  const [showUndoButton, setShowUndoButton] = useState(false)
+  const [originalPrompt, setOriginalPrompt] = useState('')
   const [tokenCount, setTokenCount] = useState(0)
   const knowledgeState = useAppSelector((state) => state.knowledge)
   const showKnowledgeIcon = useSidebarIconShow('knowledge')
@@ -98,7 +100,7 @@ const PopupContainer: React.FC<Props> = ({ resolve }) => {
     resolve(null)
   }
 
-  const handleButtonClick = async () => {
+  const handleGenerateButtonClick = async () => {
     const name = formRef.current?.getFieldValue('name')
     const content = formRef.current?.getFieldValue('prompt')
     const promptText = content || name
@@ -112,6 +114,7 @@ const PopupContainer: React.FC<Props> = ({ resolve }) => {
     }
 
     setLoading(true)
+    setShowUndoButton(false)
 
     try {
       const generatedText = await fetchGenerate({
@@ -119,6 +122,8 @@ const PopupContainer: React.FC<Props> = ({ resolve }) => {
         content: promptText
       })
       form.setFieldsValue({ prompt: generatedText })
+      setShowUndoButton(true)
+      setOriginalPrompt(content)
     } catch (error) {
       console.error('Error fetching data:', error)
     }
@@ -126,6 +131,11 @@ const PopupContainer: React.FC<Props> = ({ resolve }) => {
     setLoading(false)
   }
 
+  const handleUndoButtonClick = async () => {
+      form.setFieldsValue({ prompt: originalPrompt })
+      setShowUndoButton(false)
+  }
+
   // Compute label width based on the longest label
   const labelWidth = [t('agents.add.name'), t('agents.add.prompt'), t('agents.add.knowledge_base')]
     .map((labelText) => stringWidth(labelText) * 8)
@@ -155,6 +165,7 @@ const PopupContainer: React.FC<Props> = ({ resolve }) => {
           if (changedValues.prompt) {
             const count = await estimateTextTokens(changedValues.prompt)
             setTokenCount(count)
+            setShowUndoButton(false)
           }
         }}>
         <Form.Item name="name" label="Emoji">
@@ -176,10 +187,15 @@ const PopupContainer: React.FC<Props> = ({ resolve }) => {
           <TokenCount>Tokens: {tokenCount}</TokenCount>
           <Button
             icon={loading ? <LoadingOutlined /> : <ThunderboltOutlined />}
-            onClick={handleButtonClick}
+            onClick={handleGenerateButtonClick}
             style={{ position: 'absolute', top: 8, right: 8 }}
             disabled={loading}
           />
+          {showUndoButton && <Button
+            icon={<RollbackOutlined />}
+            onClick={handleUndoButtonClick}
+            style={{ position: 'absolute', top: 8, right: 48 }}
+          />}
         </div>
         {showKnowledgeIcon && (
           <Form.Item name="knowledge_base_ids" label={t('agents.add.knowledge_base')} rules={[{ required: false }]}>