feat: add stop button to cancel AI generation

app/api/chat/route.ts

@@ -449,6 +449,7 @@ ${userInputText}
 
     const result = streamText({
         model,
+        abortSignal: req.signal,
         ...(process.env.MAX_OUTPUT_TOKENS && {
             maxOutputTokens: parseInt(process.env.MAX_OUTPUT_TOKENS, 10),
         }),

components/chat-input.tsx

@@ -7,6 +7,7 @@ import {
     Link,
     Loader2,
     Send,
+    Square,
 } from "lucide-react"
 import type React from "react"
 import {
@@ -154,6 +155,7 @@ interface ChatInputProps {
     status: "submitted" | "streaming" | "ready" | "error"
     onSubmit: (e: React.FormEvent<HTMLFormElement>) => void
     onChange: (e: React.ChangeEvent<HTMLTextAreaElement>) => void
+    onStop?: () => void
     files?: File[]
     onFileChange?: (files: File[]) => void
     pdfData?: Map<
@@ -182,6 +184,7 @@ export const ChatInput = forwardRef<ChatInputRef, ChatInputProps>(
             status,
             onSubmit,
             onChange,
+            onStop,
             files = [],
             onFileChange = () => {},
             pdfData = new Map(),
@@ -552,24 +555,30 @@ export const ChatInput = forwardRef<ChatInputRef, ChatInputProps>(
                             showUnvalidatedModels={showUnvalidatedModels}
                         />
                         <div className="w-px h-5 bg-border mx-1" />
-                        <Button
+                        {(status === "streaming" || status === "submitted") &&
-                            type="submit"
+                        onStop ? (
-                            disabled={isDisabled || !input.trim()}
+                            <Button
-                            size="sm"
+                                type="button"
-                            className="h-8 px-4 rounded-xl font-medium shadow-sm"
+                                onClick={onStop}
-                            aria-label={
+                                size="sm"
-                                isDisabled ? dict.chat.sending : dict.chat.send
+                                variant="destructive"
-                            }
+                                className="h-8 w-8 p-0 rounded-xl shadow-sm"
-                        >
+                                aria-label="Stop generation"
-                            {isDisabled ? (
+                            >
-                                <Loader2 className="h-4 w-4 animate-spin" />
+                                <Square className="h-4 w-4" />
-                            ) : (
+                            </Button>
-                                <>
+                        ) : (
-                                    <Send className="h-4 w-4 mr-1.5" />
+                            <Button
-                                    {dict.chat.send}
+                                type="submit"
-                                </>
+                                disabled={isDisabled || !input.trim()}
-                            )}
+                                size="sm"
-                        </Button>
+                                className="h-8 px-4 rounded-xl font-medium shadow-sm"
+                                aria-label={dict.chat.send}
+                            >
+                                <Send className="h-4 w-4 mr-1.5" />
+                                {dict.chat.send}
+                            </Button>
+                        )}
                     </div>
                 </div>
                 <HistoryDialog

components/chat-panel.tsx

@@ -339,151 +339,155 @@ export default function ChatPanel({
         onValidationStateChange: handleValidationStateChange,
     })
 
-    const { messages, sendMessage, addToolOutput, status, error, setMessages } =
+    const {
-        useChat({
+        messages,
-            transport: new DefaultChatTransport({
+        sendMessage,
-                api: getApiEndpoint("/api/chat"),
+        addToolOutput,
-            }),
+        status,
-            onToolCall: async ({ toolCall }) => {
+        error,
-                await handleToolCall({ toolCall }, addToolOutput)
+        setMessages,
-            },
+        stop,
-            onError: (error) => {
+    } = useChat({
-                // Handle server-side quota limit (429 response)
+        transport: new DefaultChatTransport({
-                // AI SDK puts the full response body in error.message for non-OK responses
+            api: getApiEndpoint("/api/chat"),
-                try {
+        }),
-                    const data = JSON.parse(error.message)
+        onToolCall: async ({ toolCall }) => {
-                    if (data.type === "request") {
+            await handleToolCall({ toolCall }, addToolOutput)
-                        quotaManager.showQuotaLimitToast(data.used, data.limit)
+        },
-                        return
+        onError: (error) => {
-                    }
+            // Handle server-side quota limit (429 response)
-                    if (data.type === "token") {
+            // AI SDK puts the full response body in error.message for non-OK responses
-                        quotaManager.showTokenLimitToast(data.used, data.limit)
+            try {
-                        return
+                const data = JSON.parse(error.message)
-                    }
+                if (data.type === "request") {
-                    if (data.type === "tpm") {
+                    quotaManager.showQuotaLimitToast(data.used, data.limit)
-                        quotaManager.showTPMLimitToast(data.limit)
+                    return
-                        return
-                    }
-                } catch {
-                    // Not JSON, fall through to string matching for backwards compatibility
                 }
+                if (data.type === "token") {
+                    quotaManager.showTokenLimitToast(data.used, data.limit)
+                    return
+                }
+                if (data.type === "tpm") {
+                    quotaManager.showTPMLimitToast(data.limit)
+                    return
+                }
+            } catch {
+                // Not JSON, fall through to string matching for backwards compatibility
+            }
 
-                // Fallback to string matching
+            // Fallback to string matching
-                if (error.message.includes("Daily request limit")) {
+            if (error.message.includes("Daily request limit")) {
-                    quotaManager.showQuotaLimitToast()
+                quotaManager.showQuotaLimitToast()
-                    return
+                return
-                }
+            }
-                if (error.message.includes("Daily token limit")) {
+            if (error.message.includes("Daily token limit")) {
-                    quotaManager.showTokenLimitToast()
+                quotaManager.showTokenLimitToast()
-                    return
+                return
+            }
+            if (
+                error.message.includes("Rate limit exceeded") ||
+                error.message.includes("tokens per minute")
+            ) {
+                quotaManager.showTPMLimitToast()
+                return
+            }
+
+            // Silence access code error in console since it's handled by UI
+            if (!error.message.includes("Invalid or missing access code")) {
+                console.error("Chat error:", error)
+            }
+
+            // Translate technical errors into user-friendly messages
+            // The server now handles detailed error messages, so we can display them directly.
+            // But we still handle connection/network errors that happen before reaching the server.
+            let friendlyMessage = error.message
+
+            // Simple check for network errors if message is generic
+            if (friendlyMessage === "Failed to fetch") {
+                friendlyMessage = "Network error. Please check your connection."
+            }
+
+            // Truncated tool input error (model output limit too low)
+            if (friendlyMessage.includes("toolUse.input is invalid")) {
+                friendlyMessage =
+                    "Output was truncated before the diagram could be generated. Try a simpler request or increase the maxOutputLength."
+            }
+
+            // Translate image not supported error
+            if (
+                friendlyMessage.includes("image content block") ||
+                friendlyMessage.toLowerCase().includes("image_url")
+            ) {
+                friendlyMessage = "This model doesn't support image input."
+            }
+
+            // Add system message for error so it can be cleared
+            setMessages((currentMessages) => {
+                const errorMessage = {
+                    id: `error-${Date.now()}`,
+                    role: "system" as const,
+                    content: friendlyMessage,
+                    parts: [{ type: "text" as const, text: friendlyMessage }],
                 }
+                return [...currentMessages, errorMessage]
+            })
+
+            if (error.message.includes("Invalid or missing access code")) {
+                // Show settings dialog to help user fix it
+                setShowSettingsDialog(true)
+            }
+        },
+        onFinish: () => {},
+        sendAutomaticallyWhen: ({ messages }) => {
+            const isInContinuationMode = partialXmlRef.current.length > 0
+
+            const shouldRetry = hasToolErrors(
+                messages as unknown as ChatMessage[],
+            )
+
+            if (!shouldRetry) {
+                // No error, reset retry count and clear state
+                autoRetryCountRef.current = 0
+                continuationRetryCountRef.current = 0
+                partialXmlRef.current = ""
+                return false
+            }
+
+            // Continuation mode: limited retries for truncation handling
+            if (isInContinuationMode) {
                 if (
-                    error.message.includes("Rate limit exceeded") ||
+                    continuationRetryCountRef.current >=
-                    error.message.includes("tokens per minute")
+                    MAX_CONTINUATION_RETRY_COUNT
                 ) {
-                    quotaManager.showTPMLimitToast()
+                    toast.error(
-                    return
+                        formatMessage(dict.errors.continuationRetryLimit, {
-                }
+                            max: MAX_CONTINUATION_RETRY_COUNT,
-
+                        }),
-                // Silence access code error in console since it's handled by UI
+                    )
-                if (!error.message.includes("Invalid or missing access code")) {
-                    console.error("Chat error:", error)
-                }
-
-                // Translate technical errors into user-friendly messages
-                // The server now handles detailed error messages, so we can display them directly.
-                // But we still handle connection/network errors that happen before reaching the server.
-                let friendlyMessage = error.message
-
-                // Simple check for network errors if message is generic
-                if (friendlyMessage === "Failed to fetch") {
-                    friendlyMessage =
-                        "Network error. Please check your connection."
-                }
-
-                // Truncated tool input error (model output limit too low)
-                if (friendlyMessage.includes("toolUse.input is invalid")) {
-                    friendlyMessage =
-                        "Output was truncated before the diagram could be generated. Try a simpler request or increase the maxOutputLength."
-                }
-
-                // Translate image not supported error
-                if (
-                    friendlyMessage.includes("image content block") ||
-                    friendlyMessage.toLowerCase().includes("image_url")
-                ) {
-                    friendlyMessage = "This model doesn't support image input."
-                }
-
-                // Add system message for error so it can be cleared
-                setMessages((currentMessages) => {
-                    const errorMessage = {
-                        id: `error-${Date.now()}`,
-                        role: "system" as const,
-                        content: friendlyMessage,
-                        parts: [
-                            { type: "text" as const, text: friendlyMessage },
-                        ],
-                    }
-                    return [...currentMessages, errorMessage]
-                })
-
-                if (error.message.includes("Invalid or missing access code")) {
-                    // Show settings dialog to help user fix it
-                    setShowSettingsDialog(true)
-                }
-            },
-            onFinish: () => {},
-            sendAutomaticallyWhen: ({ messages }) => {
-                const isInContinuationMode = partialXmlRef.current.length > 0
-
-                const shouldRetry = hasToolErrors(
-                    messages as unknown as ChatMessage[],
-                )
-
-                if (!shouldRetry) {
-                    // No error, reset retry count and clear state
-                    autoRetryCountRef.current = 0
                     continuationRetryCountRef.current = 0
                     partialXmlRef.current = ""
                     return false
                 }
-
+                continuationRetryCountRef.current++
-                // Continuation mode: limited retries for truncation handling
+            } else {
-                if (isInContinuationMode) {
+                // Regular error: check retry count limit
-                    if (
+                if (autoRetryCountRef.current >= MAX_AUTO_RETRY_COUNT) {
-                        continuationRetryCountRef.current >=
+                    toast.error(
-                        MAX_CONTINUATION_RETRY_COUNT
+                        formatMessage(dict.errors.retryLimit, {
-                    ) {
+                            max: MAX_AUTO_RETRY_COUNT,
-                        toast.error(
+                        }),
-                            formatMessage(dict.errors.continuationRetryLimit, {
+                    )
-                                max: MAX_CONTINUATION_RETRY_COUNT,
+                    autoRetryCountRef.current = 0
-                            }),
+                    partialXmlRef.current = ""
-                        )
+                    return false
-                        continuationRetryCountRef.current = 0
-                        partialXmlRef.current = ""
-                        return false
-                    }
-                    continuationRetryCountRef.current++
-                } else {
-                    // Regular error: check retry count limit
-                    if (autoRetryCountRef.current >= MAX_AUTO_RETRY_COUNT) {
-                        toast.error(
-                            formatMessage(dict.errors.retryLimit, {
-                                max: MAX_AUTO_RETRY_COUNT,
-                            }),
-                        )
-                        autoRetryCountRef.current = 0
-                        partialXmlRef.current = ""
-                        return false
-                    }
-                    // Increment retry count for actual errors
-                    autoRetryCountRef.current++
                 }
+                // Increment retry count for actual errors
+                autoRetryCountRef.current++
+            }
 
-                return true
+            return true
-            },
+        },
-        })
+    })
 
     // Store sendMessage in ref for use in callbacks (like handleImproveWithSuggestions)
     useEffect(() => {
@@ -991,6 +995,29 @@ export default function ChatPanel({
         }
     }
 
+    // Handle stop button click
+    const handleStop = useCallback(() => {
+        const lastMessage = messages[messages.length - 1]
+        const toolParts = lastMessage?.parts?.filter(
+            (part: any) =>
+                part.type?.startsWith("tool-") &&
+                part.state === "input-streaming",
+        )
+
+        toolParts?.forEach((part: any) => {
+            if (part.toolCallId) {
+                addToolOutput({
+                    tool: part.type.replace("tool-", ""),
+                    toolCallId: part.toolCallId,
+                    state: "output-error",
+                    errorText: "Stopped by user",
+                })
+            }
+        })
+
+        stop()
+    }, [messages, addToolOutput, stop])
+
     // Send chat message with headers
     const sendChatMessage = (
         parts: any,
@@ -1357,6 +1384,7 @@ export default function ChatPanel({
                     status={status}
                     onSubmit={onFormSubmit}
                     onChange={handleInputChange}
+                    onStop={handleStop}
                     files={files}
                     onFileChange={handleFileChange}
                     pdfData={pdfData}