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).
This commit is contained in:
James Brunton
2026-05-28 09:25:23 +00:00
committed by GitHub
parent 43b67d213d
commit d459ded168
2 changed files with 62 additions and 13 deletions
@@ -301,6 +301,7 @@ interface ChatContextValue {
toggleOpen: () => void; toggleOpen: () => void;
setOpen: (open: boolean) => void; setOpen: (open: boolean) => void;
sendMessage: (content: string) => Promise<void>; sendMessage: (content: string) => Promise<void>;
cancelMessage: () => void;
} }
const ChatContext = createContext<ChatContextValue | null>(null); const ChatContext = createContext<ChatContextValue | null>(null);
@@ -317,6 +318,10 @@ export function ChatProvider({ children }: { children: ReactNode }) {
const { files: activeFiles, fileStubs: activeFileStubs } = useAllFiles(); const { files: activeFiles, fileStubs: activeFileStubs } = useAllFiles();
const { actions: fileActions } = useFileActions(); const { actions: fileActions } = useFileActions();
const abortRef = useRef<AbortController | null>(null); 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); const messagesRef = useRef<ChatMessage[]>(state.messages);
messagesRef.current = 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( const sendMessage = useCallback(
async (content: string) => { async (content: string) => {
// Abort any in-flight request // Abort any in-flight request
@@ -510,7 +522,22 @@ export function ChatProvider({ children }: { children: ReactNode }) {
throw new Error("Stream ended without a result"); throw new Error("Stream ended without a result");
} }
} catch (e) { } 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: "SET_PROGRESS", progress: null });
dispatch({ dispatch({
type: "ADD_MESSAGE", type: "ADD_MESSAGE",
@@ -542,6 +569,7 @@ export function ChatProvider({ children }: { children: ReactNode }) {
toggleOpen, toggleOpen,
setOpen, setOpen,
sendMessage, sendMessage,
cancelMessage,
}} }}
> >
{children} {children}
@@ -22,6 +22,7 @@ import {
List, List,
} from "@mantine/core"; } from "@mantine/core";
import SendIcon from "@mui/icons-material/Send"; import SendIcon from "@mui/icons-material/Send";
import StopIcon from "@mui/icons-material/Stop";
import ChatBubbleOutlineIcon from "@mui/icons-material/ChatBubbleOutlined"; import ChatBubbleOutlineIcon from "@mui/icons-material/ChatBubbleOutlined";
import CloseIcon from "@mui/icons-material/Close"; import CloseIcon from "@mui/icons-material/Close";
import ExpandMoreIcon from "@mui/icons-material/ExpandMore"; import ExpandMoreIcon from "@mui/icons-material/ExpandMore";
@@ -208,8 +209,15 @@ function ChatMessageBubble({
export function ChatPanel() { export function ChatPanel() {
const { t } = useTranslation(); const { t } = useTranslation();
const { messages, isOpen, isLoading, progress, toggleOpen, sendMessage } = const {
useChat(); messages,
isOpen,
isLoading,
progress,
toggleOpen,
sendMessage,
cancelMessage,
} = useChat();
const resolveToolName = useToolNameResolver(); const resolveToolName = useToolNameResolver();
const [input, setInput] = useState(""); const [input, setInput] = useState("");
const scrollRef = useRef<HTMLDivElement>(null); const scrollRef = useRef<HTMLDivElement>(null);
@@ -330,17 +338,30 @@ export function ChatPanel() {
onKeyDown={handleKeyDown} onKeyDown={handleKeyDown}
disabled={isLoading} disabled={isLoading}
rightSection={ rightSection={
isLoading ? (
<ActionIcon
variant="filled"
color="red"
size="sm"
onClick={cancelMessage}
aria-label="Stop generating"
>
<StopIcon sx={{ fontSize: 14 }} />
</ActionIcon>
) : (
<ActionIcon <ActionIcon
variant="filled" variant="filled"
color="blue" color="blue"
size="sm" size="sm"
onClick={handleSend} onClick={handleSend}
disabled={!input.trim() || isLoading} disabled={!input.trim()}
aria-label="Send message" aria-label="Send message"
> >
<SendIcon sx={{ fontSize: 14 }} /> <SendIcon sx={{ fontSize: 14 }} />
</ActionIcon> </ActionIcon>
)
} }
rightSectionPointerEvents="all"
rightSectionWidth={36} rightSectionWidth={36}
style={{ flex: 1 }} style={{ flex: 1 }}
/> />