From 93c816f3ab6b759ecad0661c0652e22beb86609c Mon Sep 17 00:00:00 2001
From: yuetsh <517252939@qq.com>
Date: Mon, 6 Apr 2026 05:50:57 -0600
Subject: [PATCH] add stop button

---
 src/components/ai/PromptPanel.vue | 13 +++++++++++--
 src/store/prompt.ts               | 21 +++++++++++++++++++++
 2 files changed, 32 insertions(+), 2 deletions(-)

diff --git a/src/components/ai/PromptPanel.vue b/src/components/ai/PromptPanel.vue
index a950b86..15fb13a 100644
--- a/src/components/ai/PromptPanel.vue
+++ b/src/components/ai/PromptPanel.vue
@@ -35,9 +35,16 @@
           :disabled="streaming"
         />
         <n-button
+          v-if="streaming"
+          type="error"
+          @click="stopPrompt"
+        >
+          停止
+        </n-button>
+        <n-button
+          v-else
           type="primary"
-          :loading="streaming"
-          :disabled="!input.trim() || streaming"
+          :disabled="!input.trim() || !connected"
           @click="send"
         >
           发送
@@ -55,7 +62,9 @@ import {
   messages,
   streaming,
   streamingContent,
+  connected,
   sendPrompt,
+  stopPrompt,
 } from "../../store/prompt"
 
 const input = ref("")
diff --git a/src/store/prompt.ts b/src/store/prompt.ts
index 5c5642c..af4afaa 100644
--- a/src/store/prompt.ts
+++ b/src/store/prompt.ts
@@ -27,8 +27,10 @@ export function setOnCodeComplete(fn: typeof _onCodeComplete) {
 }
 
 let ws: WebSocket | null = null
+let _currentTaskId = 0
 
 export function connectPrompt(taskId: number) {
+  _currentTaskId = taskId
   if (ws) ws.close()
 
   ws = new WebSocket(`${WS_BASE_URL}/ws/prompt/${taskId}/`)
@@ -41,6 +43,8 @@ export function connectPrompt(taskId: number) {
     const data = JSON.parse(event.data)
 
     if (data.type === "init") {
+      streaming.value = false
+      streamingContent.value = ""
       // Skip overwriting messages if HTTP preload already loaded this conversation.
       // If conversation_id differs (e.g. after "新对话"), always overwrite.
       const alreadyLoaded = conversationId.value === data.conversation_id
@@ -105,10 +109,27 @@ export function disconnectPrompt() {
 
 export function sendPrompt(content: string, model: string = "") {
   if (!ws || ws.readyState !== WebSocket.OPEN) return
+  streaming.value = true
   messages.value.push({ role: "user", content })
   ws.send(JSON.stringify({ type: "message", content, model }))
 }
 
+export function stopPrompt() {
+  // Remove the user message added to UI (not yet saved in DB)
+  if (
+    messages.value.length > 0 &&
+    messages.value[messages.value.length - 1].role === "user"
+  ) {
+    messages.value.pop()
+  }
+  streaming.value = false
+  streamingContent.value = ""
+  // connectPrompt closes old WS (triggering backend disconnect cleanup) then reconnects
+  if (_currentTaskId) {
+    connectPrompt(_currentTaskId)
+  }
+}
+
 function applyCode(code: {
   html: string | null
   css: string | null