Merge pull request #6 from GwIhViEte/feature/新增停止生成按钮

Feature/新增停止生成按钮

Author: ooboqoo

package-lock.json

@@ -3,7 +3,7 @@ import { createOpenAI } from '@ai-sdk/openai'
 import { PROMPT_SYSTEM } from './prompts'
 import { settings } from './settings'
 
-export function getSolutionStream(base64Image: string) {
+export function getSolutionStream(base64Image: string, abortSignal?: AbortSignal) {
   const openai = createOpenAI({
     baseURL: settings.apiBaseURL,
     apiKey: settings.apiKey
@@ -18,15 +18,16 @@ export function getSolutionStream(base64Image: string) {
         content: [
           {
             type: 'text',
-            text: `The screenshot is as follows, and the programming language is ${settings.codeLanguage}.`
+            text: `以下是一张截图。编程语言：${settings.codeLanguage || '未指定'}。请严格将所有代码放在 Markdown 代码块中（使用三反引号，若可请标注语言），解释文字放在代码块之外。`
           },
           {
             type: 'image',
             image: base64Image
           }
         ]
       }
-    ]
+    ],
+    abortSignal
   })
   return textStream
 }

src/main/ai.ts

@@ -1,5 +1,24 @@
 import 'dotenv/config'
 import { app, BrowserWindow, globalShortcut } from 'electron'
+
+// Swallow AbortError from user-initiated stream cancellations to keep console clean
+function isAbortError(error: unknown): boolean {
+  const err = error as any
+  return (
+    !!err &&
+    (err.name === 'AbortError' || err.code === 'ABORT_ERR' || /aborted/i.test(String(err.message)))
+  )
+}
+
+process.on('unhandledRejection', (error) => {
+  if (isAbortError(error)) return
+  console.error(error)
+})
+
+process.on('uncaughtException', (error) => {
+  if (isAbortError(error)) return
+  console.error(error)
+})
 import { electronApp, optimizer } from '@electron-toolkit/utils'
 import './shortcuts'
 import { createWindow } from './main-window'

src/main/index.ts

@@ -1,16 +1,21 @@
 export const PROMPT_SYSTEM = `
-你是一个经验丰富的程序员，擅长解决编程问题。
-你将根据用户提供的截图，给出解决问题的思路和代码。
-请使用中文回答。
+你是一名经验丰富的程序员，负责根据截图快速分析问题并给出可直接运行的解决方案。
+请使用简洁中文回答，并遵循以下要求：
 
-* 先分析截图，提取出问题描述
-* 根据问题描述，给出「解决问题的思路」，优先选择简单的解法
-* 根据解题思路给出「代码」，代码需要有详尽的「注释」
-* 提供测试用例，并给出测试结果
-* 最后给出「总结」，总结解题思路和代码实现，以及测试结果
+- 先简要概括截图中题目/需求与关键点；
+- 给出清晰的解题思路与步骤，优先选择实现难度更低、鲁棒性更好的方案；
+- 依据思路输出完整代码，包含必要的注释；
+- 如有多种实现或注意事项，列出权衡与扩展建议；
 
-如果分析截图没哟发现编程问题，则进行以下尝试
-* 尝试查找其他编程相关问题，如果有，则给出答案，并跟进情况给出适当的解析和提示。
-* 尝试查找其他问题，如果有，则给出答案，并跟进情况给出适当的解析和提示。
-* 如果以上尝试都没有发现问题，则对截图内容进行分析总结。
+格式与代码输出规范（重要）：
+- 所有代码必须使用 Markdown 代码块包裹，使用三反引号 \`\`\` 开始与结束；
+- 若已知编程语言，请在开头三反引号后标注语言（例如 \`\`\`ts、\`\`\`python）；
+- 解释性文字必须放在代码块之外，禁止把解释写进代码块；
+- 如有多段代码，分别使用独立的代码块；
+- 若需要补充命令行或配置，也使用独立代码块并标注对应语言（如 \`\`\`bash、\`\`\`json）。
+
+当截图无法明确题目时：
+- 可合理推断题目并给出自洽的示例与实现；
+- 或指出缺失信息并给出可执行的最小示例与后续补全建议。
 `
+

src/main/prompts.ts

@@ -19,6 +19,22 @@ enum ShortcutStatus {
 
 const MOVE_STEP = 200
 const shortcuts: Record<string, Shortcut> = {}
+
+type AbortReason = 'user' | 'new-request'
+
+interface StreamContext {
+  controller: AbortController
+  reason: AbortReason | null
+}
+
+let currentStreamContext: StreamContext | null = null
+
+function abortCurrentStream(reason: AbortReason) {
+  if (!currentStreamContext) return
+  currentStreamContext.reason = reason
+  currentStreamContext.controller.abort()
+}
+
 const callbacks: Record<string, () => void> = {
   hideOrShowMainWindow: async () => {
     const mainWindow = global.mainWindow
@@ -33,16 +49,74 @@ const callbacks: Record<string, () => void> = {
   takeScreenshot: async () => {
     const mainWindow = global.mainWindow
     if (!mainWindow || mainWindow.isDestroyed() || !state.inCoderPage || !settings.apiKey) return
+
+    abortCurrentStream('new-request')
     const screenshotData = await takeScreenshot()
     if (screenshotData && mainWindow && !mainWindow.isDestroyed()) {
+      const streamContext: StreamContext = {
+        controller: new AbortController(),
+        reason: null
+      }
+      currentStreamContext = streamContext
       mainWindow.webContents.send('screenshot-taken', screenshotData)
-      const solutionStream = getSolutionStream(screenshotData)
-      for await (const chunk of solutionStream) {
-        mainWindow.webContents.send('solution-chunk', chunk)
+      let endedNaturally = true
+      let streamStarted = false
+      try {
+        const solutionStream = getSolutionStream(screenshotData, streamContext.controller.signal)
+        streamStarted = true
+        try {
+          for await (const chunk of solutionStream) {
+            if (streamContext.controller.signal.aborted) {
+              endedNaturally = false
+              break
+            }
+            mainWindow.webContents.send('solution-chunk', chunk)
+          }
+        } catch (error) {
+          if (!streamContext.controller.signal.aborted) {
+            endedNaturally = false
+            console.error('Error streaming solution:', error)
+            const errorMessage = error instanceof Error ? error.message : 'Unknown error'
+            mainWindow.webContents.send('solution-error', errorMessage)
+          } else {
+            endedNaturally = false
+          }
+        }
+
+        if (streamContext.controller.signal.aborted) {
+          if (streamContext.reason === 'user') {
+            mainWindow.webContents.send('solution-stopped')
+          }
+        } else if (endedNaturally) {
+          mainWindow.webContents.send('solution-complete')
+        }
+      } catch (error) {
+        if (streamContext.controller.signal.aborted) {
+          if (streamContext.reason === 'user') {
+            mainWindow.webContents.send('solution-stopped')
+          }
+        } else {
+          endedNaturally = false
+          const errorMessage = error instanceof Error ? error.message : 'Unknown error'
+          console.error('Error streaming solution:', error)
+          mainWindow.webContents.send('solution-error', errorMessage)
+        }
+      } finally {
+        if (currentStreamContext === streamContext) {
+          currentStreamContext = null
+        }
+        if (!streamStarted && streamContext.reason === 'user') {
+          mainWindow.webContents.send('solution-stopped')
+        }
       }
     }
   },
 
+  // Stop current AI solution stream
+  stopSolutionStream: () => {
+    abortCurrentStream('user')
+  },
+
   ignoreOrEnableMouse: () => {
     const mainWindow = global.mainWindow
     if (!mainWindow || mainWindow.isDestroyed() || !state.inCoderPage) return
@@ -122,3 +196,9 @@ ipcMain.handle('updateShortcuts', (_event, _shortcuts: { action: string; key: st
     }
   })
 })
+
+ipcMain.handle('stopSolutionStream', () => {
+  if (!currentStreamContext) return false
+  abortCurrentStream('user')
+  return true
+})

src/main/shortcuts.ts

@@ -19,6 +19,13 @@ declare global {
       removeScreenshotListener: () => void
       onSolutionChunk: (callback: (chunk: string) => void) => void
       removeSolutionChunkListener: () => void
+      stopSolutionStream: () => Promise<boolean>
+      onSolutionComplete: (callback: () => void) => void
+      removeSolutionCompleteListener: () => void
+      onSolutionStopped: (callback: () => void) => void
+      removeSolutionStoppedListener: () => void
+      onSolutionError: (callback: (message: string) => void) => void
+      removeSolutionErrorListener: () => void
       onScrollPageUp: (callback: () => void) => void
       removeScrollPageUpListener: () => void
       onScrollPageDown: (callback: () => void) => void

src/preload/index.d.ts

@@ -55,6 +55,33 @@ const api = {
     ipcRenderer.removeAllListeners('solution-chunk')
   },
 
+  // Stop solution stream
+  stopSolutionStream: () => ipcRenderer.invoke('stopSolutionStream'),
+
+  // Listen for solution completion
+  onSolutionComplete: (callback: () => void) => {
+    ipcRenderer.on('solution-complete', callback)
+  },
+  removeSolutionCompleteListener: () => {
+    ipcRenderer.removeAllListeners('solution-complete')
+  },
+
+  onSolutionStopped: (callback: () => void) => {
+    ipcRenderer.on('solution-stopped', callback)
+  },
+  removeSolutionStoppedListener: () => {
+    ipcRenderer.removeAllListeners('solution-stopped')
+  },
+
+  onSolutionError: (callback: (message: string) => void) => {
+    ipcRenderer.on('solution-error', (_event, message) => {
+      callback(message)
+    })
+  },
+  removeSolutionErrorListener: () => {
+    ipcRenderer.removeAllListeners('solution-error')
+  },
+
   // Listen for scroll page up
   onScrollPageUp: (callback: () => void) => {
     ipcRenderer.on('scroll-page-up', callback)

src/preload/index.ts

@@ -17,6 +17,15 @@ export function AppContent() {
     clearSolution
   } = useSolutionStore()
 
+  // Ensure any unfinished fenced code block is closed when a run ends
+  const ensureClosedFence = () => {
+    const text = useSolutionStore.getState().solutionChunks.join('')
+    const fenceCount = (text.match(/```/g) || []).length
+    if (fenceCount % 2 === 1) {
+      addSolutionChunk('\n```')
+    }
+  }
+
   useEffect(() => {
     // Listen for screenshot events
     window.api.onScreenshotTaken((data: string) => {
@@ -38,6 +47,24 @@ export function AppContent() {
     }
   }, [setScreenshotData, clearSolution, setIsLoading, addSolutionChunk])
 
+  useEffect(() => {
+    window.api.onSolutionComplete(() => {
+      setIsLoading(false)
+    })
+    window.api.onSolutionStopped(() => {
+      setIsLoading(false)
+    })
+    window.api.onSolutionError((message: string) => {
+      setIsLoading(false)
+      addSolutionChunk(`\n\n> 生成失败：${message}`)
+    })
+    return () => {
+      window.api.removeSolutionCompleteListener()
+      window.api.removeSolutionStoppedListener()
+      window.api.removeSolutionErrorListener()
+    }
+  }, [setIsLoading, addSolutionChunk])
+
   // Mark solution as complete when chunks stop coming
   useEffect(() => {
     if (isLoading && solutionChunks.length > 0) {
@@ -50,6 +77,13 @@ export function AppContent() {
     return undefined
   }, [solutionChunks, isLoading, setIsLoading])
 
+  // When run ends (isLoading goes false), ensure code fence is closed
+  useEffect(() => {
+    if (!isLoading) {
+      ensureClosedFence()
+    }
+  }, [isLoading])
+
   useEffect(() => {
     window.api.onScrollPageUp(() => {
       const container = document.getElementById('app-content')

src/renderer/src/coder/AppContent.tsx

@@ -1,21 +1,36 @@
-import { Pointer, PointerOff } from 'lucide-react'
+import { Pointer, PointerOff, Square } from 'lucide-react'
 import { useSolutionStore } from '@/lib/store/solution'
 import { useShortcutsStore } from '@/lib/store/shortcuts'
 import { useAppStore } from '@/lib/store/app'
 import ShortcutRenderer from '@/components/ShortcutRenderer'
+import { Button } from '@/components/ui/button'
 
 export function AppStatusBar() {
-  const { isLoading: isReceivingSolution } = useSolutionStore()
+  const { isLoading: isReceivingSolution, setIsLoading } = useSolutionStore()
   const { ignoreMouse } = useAppStore()
   const { shortcuts } = useShortcutsStore()
 
+  const handleStop = () => {
+    setIsLoading(false)
+    void window.api.stopSolutionStream()
+  }
+
   return (
     <div className="absolute bottom-0 flex items-center justify-between w-full text-blue-100 bg-gray-600/10 px-4 pb-1">
       <div>
         {isReceivingSolution && (
-          <div className="flex items-center">
-            <div className="animate-spin rounded-full h-4 w-4 border-b-2 border-r-2 border-[currentColor] mr-2"></div>
-            <span className="text-sm">分析中...</span>
+          <div className="flex items-center space-x-2">
+            <div className="animate-spin rounded-full h-4 w-4 border-b-2 border-r-2 border-[currentColor]"></div>
+            <span className="text-sm">正在生成...</span>
+            <Button
+              variant="secondary"
+              size="sm"
+              className="h-6 px-2 text-xs"
+              onClick={handleStop}
+            >
+              <Square className="w-3 h-3 mr-1" />
+              停止生成
+            </Button>
           </div>
         )}
       </div>
@@ -26,7 +41,7 @@ export function AppStatusBar() {
             <>
               <PointerOff className="w-4 h-4 mr-2" />
               <span className="text-xs">
-                取消鼠标穿透
+                取消鼠标透传
                 <ShortcutRenderer
                   shortcut={shortcuts.ignoreOrEnableMouse.key}
                   className="inline-block scale-75 text-xs border border-current bg-transparent py-0 px-1"