feat: Add follow-up question functionality and conversation history management

ooboqoo · ooboqoo · commit 403a818587b3 · 2025-11-05T22:22:15.000+08:00
diff --git a/src/main/ai.ts b/src/main/ai.ts
@@ -1,6 +1,6 @@
 import { readFileSync } from 'node:fs'
 import { join } from 'node:path'
-import { streamText } from 'ai'
+import { streamText, type ModelMessage } from 'ai'
 import { createOpenAI } from '@ai-sdk/openai'
 import { settings } from './settings'
 
@@ -35,3 +35,37 @@ export function getSolutionStream(base64Image: string, abortSignal?: AbortSignal
   })
   return textStream
 }
+
+export function getFollowUpStream(
+  messages: ModelMessage[],
+  userQuestion: string,
+  abortSignal?: AbortSignal
+) {
+  const openai = createOpenAI({
+    baseURL: settings.apiBaseURL,
+    apiKey: settings.apiKey
+  })
+
+  // Add the user's follow-up question to the conversation
+  const updatedMessages: ModelMessage[] = [
+    ...messages,
+    {
+      role: 'user',
+      content: [
+        {
+          type: 'text',
+          text: userQuestion
+        }
+      ]
+    }
+  ]
+
+  const { textStream } = streamText({
+    model: openai(settings.model || 'gpt-4o-mini'),
+    system:
+      settings.customPrompt || PROMPT_SYSTEM + `\n使用编程语言：${settings.codeLanguage} 解答。`,
+    messages: updatedMessages,
+    abortSignal
+  })
+  return textStream
+}
diff --git a/src/main/shortcuts.ts b/src/main/shortcuts.ts
@@ -1,6 +1,7 @@
 import { globalShortcut, ipcMain } from 'electron'
+import type { ModelMessage } from 'ai'
 import { takeScreenshot } from './take-screenshot'
-import { getSolutionStream } from './ai'
+import { getSolutionStream, getFollowUpStream } from './ai'
 import { state } from './state'
 import { settings } from './settings'
 
@@ -29,6 +30,9 @@ interface StreamContext {
 
 let currentStreamContext: StreamContext | null = null
 
+// Conversation history tracking
+let conversationMessages: ModelMessage[] = []
+
 function abortCurrentStream(reason: AbortReason) {
   if (!currentStreamContext) return
   currentStreamContext.reason = reason
@@ -53,6 +57,22 @@ const callbacks: Record<string, () => void> = {
     abortCurrentStream('new-request')
     const screenshotData = await takeScreenshot()
     if (screenshotData && mainWindow && !mainWindow.isDestroyed()) {
+      conversationMessages = [
+        {
+          role: 'user',
+          content: [
+            {
+              type: 'text',
+              text: `这是屏幕截图`
+            },
+            {
+              type: 'image',
+              image: screenshotData
+            }
+          ]
+        }
+      ]
+
       const streamContext: StreamContext = {
         controller: new AbortController(),
         reason: null
@@ -61,6 +81,7 @@ const callbacks: Record<string, () => void> = {
       mainWindow.webContents.send('screenshot-taken', screenshotData)
       let endedNaturally = true
       let streamStarted = false
+      let assistantResponse = ''
       try {
         const solutionStream = getSolutionStream(screenshotData, streamContext.controller.signal)
         streamStarted = true
@@ -70,6 +91,7 @@ const callbacks: Record<string, () => void> = {
               endedNaturally = false
               break
             }
+            assistantResponse += chunk
             mainWindow.webContents.send('solution-chunk', chunk)
           }
         } catch (error) {
@@ -88,6 +110,13 @@ const callbacks: Record<string, () => void> = {
             mainWindow.webContents.send('solution-stopped')
           }
         } else if (endedNaturally) {
+          // Add assistant response to conversation history
+          if (assistantResponse) {
+            conversationMessages.push({
+              role: 'assistant',
+              content: assistantResponse
+            })
+          }
           mainWindow.webContents.send('solution-complete')
         }
       } catch (error) {
@@ -202,3 +231,102 @@ ipcMain.handle('stopSolutionStream', () => {
   abortCurrentStream('user')
   return true
 })
+
+ipcMain.handle('sendFollowUpQuestion', async (_event, question: string) => {
+  const mainWindow = global.mainWindow
+  if (!mainWindow || mainWindow.isDestroyed() || !state.inCoderPage || !settings.apiKey) {
+    return { success: false, error: 'Invalid state' }
+  }
+
+  // Validate that there's an active conversation
+  if (conversationMessages.length === 0) {
+    return { success: false, error: 'No active conversation' }
+  }
+
+  abortCurrentStream('new-request')
+  const streamContext: StreamContext = {
+    controller: new AbortController(),
+    reason: null
+  }
+  currentStreamContext = streamContext
+
+  // Add a separator before the follow-up response
+  mainWindow.webContents.send('solution-chunk', '\n\n---\n\n')
+
+  let endedNaturally = true
+  let streamStarted = false
+  let assistantResponse = ''
+
+  try {
+    const followUpStream = getFollowUpStream(
+      conversationMessages,
+      question,
+      streamContext.controller.signal
+    )
+    streamStarted = true
+
+    try {
+      for await (const chunk of followUpStream) {
+        if (streamContext.controller.signal.aborted) {
+          endedNaturally = false
+          break
+        }
+        assistantResponse += chunk
+        mainWindow.webContents.send('solution-chunk', chunk)
+      }
+    } catch (error) {
+      if (!streamContext.controller.signal.aborted) {
+        endedNaturally = false
+        console.error('Error streaming follow-up 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) {
+      // Update conversation history with user question and assistant response
+      conversationMessages.push({
+        role: 'user',
+        content: [
+          {
+            type: 'text',
+            text: question
+          }
+        ]
+      })
+      if (assistantResponse) {
+        conversationMessages.push({
+          role: 'assistant',
+          content: assistantResponse
+        })
+      }
+      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 follow-up solution:', error)
+      mainWindow.webContents.send('solution-error', errorMessage)
+    }
+  } finally {
+    if (currentStreamContext === streamContext) {
+      currentStreamContext = null
+    }
+    if (!streamStarted && streamContext.reason === 'user') {
+      mainWindow.webContents.send('solution-stopped')
+    }
+  }
+
+  return { success: true }
+})
diff --git a/src/preload/index.ts b/src/preload/index.ts
@@ -58,6 +58,9 @@ const api = {
   // Stop solution stream
   stopSolutionStream: () => ipcRenderer.invoke('stopSolutionStream'),
 
+  // Send follow-up question
+  sendFollowUpQuestion: (question: string) => ipcRenderer.invoke('sendFollowUpQuestion', question),
+
   // Listen for solution completion
   onSolutionComplete: (callback: () => void) => {
     ipcRenderer.on('solution-complete', callback)
diff --git a/src/renderer/src/coder/AppContent.tsx b/src/renderer/src/coder/AppContent.tsx
@@ -8,7 +8,6 @@ const SCROLL_OFFSET = 120
 
 export function AppContent() {
   const {
-    isLoading,
     screenshotData,
     solutionChunks,
     setScreenshotData,
@@ -17,15 +16,6 @@ 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) => {
@@ -65,25 +55,6 @@ export function AppContent() {
     }
   }, [setIsLoading, addSolutionChunk])
 
-  // Mark solution as complete when chunks stop coming
-  useEffect(() => {
-    if (isLoading && solutionChunks.length > 0) {
-      const timer = setTimeout(() => {
-        setIsLoading(false)
-      }, 1000) // Wait 1 second after last chunk to mark as complete
-
-      return () => clearTimeout(timer)
-    }
-    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')
diff --git a/src/renderer/src/coder/AppStatusBar.tsx b/src/renderer/src/coder/AppStatusBar.tsx
@@ -1,20 +1,58 @@
-import { Pointer, PointerOff, OctagonX } from 'lucide-react'
+import { useState } from 'react'
+import { Pointer, PointerOff, OctagonX, MessageCircle } 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'
+import { Dialog, DialogTitle, DialogContent, DialogFooter } from '@/components/ui/dialog'
+import { Textarea } from '@/components/ui/textarea'
 
 export function AppStatusBar() {
-  const { isLoading: isReceivingSolution, setIsLoading } = useSolutionStore()
+  const {
+    isLoading: isReceivingSolution,
+    setIsLoading,
+    screenshotData,
+    solutionChunks
+  } = useSolutionStore()
   const { ignoreMouse } = useAppStore()
   const { shortcuts } = useShortcutsStore()
+  const [isDialogOpen, setIsDialogOpen] = useState(false)
+  const [questionInput, setQuestionInput] = useState('')
 
   const handleStop = () => {
     setIsLoading(false)
     void window.api.stopSolutionStream()
   }
 
+  const handleFollowUpClick = () => {
+    setIsDialogOpen(true)
+  }
+
+  const handleDialogClose = () => {
+    setIsDialogOpen(false)
+    setQuestionInput('')
+  }
+
+  const handleSubmitQuestion = async () => {
+    if (!questionInput.trim()) return
+
+    setIsLoading(true)
+    setIsDialogOpen(false)
+    const question = questionInput.trim()
+    setQuestionInput('')
+
+    try {
+      await window.api.sendFollowUpQuestion(question)
+    } catch (error) {
+      console.error('Error sending follow-up question:', error)
+      setIsLoading(false)
+    }
+  }
+
+  // Check if there's an active conversation
+  const hasActiveConversation = screenshotData && solutionChunks.length > 0 && !isReceivingSolution
+
   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>
@@ -40,6 +78,19 @@ export function AppStatusBar() {
         )}
       </div>
       <div className="flex items-center space-x-4 select-none">
+        {/* Follow-up Question Button */}
+        {hasActiveConversation && (
+          <Button
+            variant="ghost"
+            size="sm"
+            onClick={handleFollowUpClick}
+            className="h-7 px-3 text-xs"
+            disabled={isReceivingSolution}
+          >
+            <MessageCircle className="w-4 h-4 mr-1" />
+            追问问题
+          </Button>
+        )}
         {/* Mouse Status Indicator */}
         <div className="flex items-center">
           {ignoreMouse ? (
@@ -58,6 +109,36 @@ export function AppStatusBar() {
           )}
         </div>
       </div>
+
+      {/* Follow-up Question Dialog */}
+      <Dialog open={isDialogOpen} onOpenChange={setIsDialogOpen}>
+        <DialogTitle className="sr-only">追问问题</DialogTitle>
+        <DialogContent>
+          <div className="py-4">
+            <Textarea
+              placeholder="请输入追问内容，按 Ctrl+Enter 提交..."
+              value={questionInput}
+              className="min-h-24"
+              onChange={(e) => setQuestionInput(e.target.value)}
+              autoFocus
+              onKeyDown={(e) => {
+                if (e.key === 'Enter' && (e.metaKey || e.ctrlKey)) {
+                  e.preventDefault()
+                  handleSubmitQuestion()
+                }
+              }}
+            />
+          </div>
+          <DialogFooter>
+            <Button variant="outline" onClick={handleDialogClose}>
+              取消
+            </Button>
+            <Button onClick={handleSubmitQuestion} disabled={!questionInput.trim()}>
+              提交
+            </Button>
+          </DialogFooter>
+        </DialogContent>
+      </Dialog>
     </div>
   )
 }
diff --git a/src/renderer/src/components/ui/input.tsx b/src/renderer/src/components/ui/input.tsx
