From d459ded168f1367973d4007f9c5c829cd867ec19 Mon Sep 17 00:00:00 2001
From: James Brunton <jbrunton96@gmail.com>
Date: Thu, 28 May 2026 10:25:23 +0100
Subject: [PATCH] Add cancel button to kill long-running AI tasks (#6351)

# Description of Changes
Adds a cancel button to the AI chat to allow the user to abort
long-running AI tasks. Just disconnects the SSE stream (all the backend
code already interrupts when it notices the stream is dead).
---
 .../components/chat/ChatContext.tsx           | 30 ++++++++++++-
 .../prototypes/components/chat/ChatPanel.tsx  | 45 ++++++++++++++-----
 2 files changed, 62 insertions(+), 13 deletions(-)

diff --git a/frontend/editor/src/prototypes/components/chat/ChatContext.tsx b/frontend/editor/src/prototypes/components/chat/ChatContext.tsx
index 88e2b2b81..2a1c4261c 100644
--- a/frontend/editor/src/prototypes/components/chat/ChatContext.tsx
+++ b/frontend/editor/src/prototypes/components/chat/ChatContext.tsx
@@ -301,6 +301,7 @@ interface ChatContextValue {
   toggleOpen: () => void;
   setOpen: (open: boolean) => void;
   sendMessage: (content: string) => Promise<void>;
+  cancelMessage: () => void;
 }
 
 const ChatContext = createContext<ChatContextValue | null>(null);
@@ -317,6 +318,10 @@ export function ChatProvider({ children }: { children: ReactNode }) {
   const { files: activeFiles, fileStubs: activeFileStubs } = useAllFiles();
   const { actions: fileActions } = useFileActions();
   const abortRef = useRef<AbortController | null>(null);
+  // Tracks the specific controller the user explicitly cancelled, so we can
+  // distinguish a user-initiated cancel from an abort triggered by a new
+  // message superseding the in-flight one.
+  const userCancelledRef = useRef<AbortController | null>(null);
   const messagesRef = useRef<ChatMessage[]>(state.messages);
   messagesRef.current = state.messages;
 
@@ -397,6 +402,13 @@ export function ChatProvider({ children }: { children: ReactNode }) {
     [],
   );
 
+  const cancelMessage = useCallback(() => {
+    const controller = abortRef.current;
+    if (!controller) return;
+    userCancelledRef.current = controller;
+    controller.abort();
+  }, []);
+
   const sendMessage = useCallback(
     async (content: string) => {
       // Abort any in-flight request
@@ -510,7 +522,22 @@ export function ChatProvider({ children }: { children: ReactNode }) {
           throw new Error("Stream ended without a result");
         }
       } catch (e) {
-        if ((e as Error).name === "AbortError") return;
+        if ((e as Error).name === "AbortError") {
+          if (userCancelledRef.current === controller) {
+            userCancelledRef.current = null;
+            dispatch({ type: "SET_PROGRESS", progress: null });
+            dispatch({
+              type: "ADD_MESSAGE",
+              message: {
+                id: crypto.randomUUID(),
+                role: "assistant",
+                content: "Cancelled.",
+                timestamp: Date.now(),
+              },
+            });
+          }
+          return;
+        }
         dispatch({ type: "SET_PROGRESS", progress: null });
         dispatch({
           type: "ADD_MESSAGE",
@@ -542,6 +569,7 @@ export function ChatProvider({ children }: { children: ReactNode }) {
         toggleOpen,
         setOpen,
         sendMessage,
+        cancelMessage,
       }}
     >
       {children}
diff --git a/frontend/editor/src/prototypes/components/chat/ChatPanel.tsx b/frontend/editor/src/prototypes/components/chat/ChatPanel.tsx
index 8d916b29a..d1d9a08e8 100644
--- a/frontend/editor/src/prototypes/components/chat/ChatPanel.tsx
+++ b/frontend/editor/src/prototypes/components/chat/ChatPanel.tsx
@@ -22,6 +22,7 @@ import {
   List,
 } from "@mantine/core";
 import SendIcon from "@mui/icons-material/Send";
+import StopIcon from "@mui/icons-material/Stop";
 import ChatBubbleOutlineIcon from "@mui/icons-material/ChatBubbleOutlined";
 import CloseIcon from "@mui/icons-material/Close";
 import ExpandMoreIcon from "@mui/icons-material/ExpandMore";
@@ -208,8 +209,15 @@ function ChatMessageBubble({
 
 export function ChatPanel() {
   const { t } = useTranslation();
-  const { messages, isOpen, isLoading, progress, toggleOpen, sendMessage } =
-    useChat();
+  const {
+    messages,
+    isOpen,
+    isLoading,
+    progress,
+    toggleOpen,
+    sendMessage,
+    cancelMessage,
+  } = useChat();
   const resolveToolName = useToolNameResolver();
   const [input, setInput] = useState("");
   const scrollRef = useRef<HTMLDivElement>(null);
@@ -330,17 +338,30 @@ export function ChatPanel() {
                 onKeyDown={handleKeyDown}
                 disabled={isLoading}
                 rightSection={
-                  <ActionIcon
-                    variant="filled"
-                    color="blue"
-                    size="sm"
-                    onClick={handleSend}
-                    disabled={!input.trim() || isLoading}
-                    aria-label="Send message"
-                  >
-                    <SendIcon sx={{ fontSize: 14 }} />
-                  </ActionIcon>
+                  isLoading ? (
+                    <ActionIcon
+                      variant="filled"
+                      color="red"
+                      size="sm"
+                      onClick={cancelMessage}
+                      aria-label="Stop generating"
+                    >
+                      <StopIcon sx={{ fontSize: 14 }} />
+                    </ActionIcon>
+                  ) : (
+                    <ActionIcon
+                      variant="filled"
+                      color="blue"
+                      size="sm"
+                      onClick={handleSend}
+                      disabled={!input.trim()}
+                      aria-label="Send message"
+                    >
+                      <SendIcon sx={{ fontSize: 14 }} />
+                    </ActionIcon>
+                  )
                 }
+                rightSectionPointerEvents="all"
                 rightSectionWidth={36}
                 style={{ flex: 1 }}
               />