From 858298c0703c800fad9f53b3eaf35ebdd28b0c31 Mon Sep 17 00:00:00 2001
From: JOJO <1498581755@qq.com>
Date: Sat, 28 Feb 2026 18:37:12 +0800
Subject: [PATCH] Handle ESC cancellation and tool aborts

---
 src/cli/index.js | 75 ++++++++++++++++++++++++++++++++++++++++++++----
 1 file changed, 69 insertions(+), 6 deletions(-)

diff --git a/src/cli/index.js b/src/cli/index.js
index dacfe80..915cb42 100644
--- a/src/cli/index.js
+++ b/src/cli/index.js
@@ -18,7 +18,7 @@ const { renderBanner } = require('../ui/banner');
 const { buildStartLine, buildFinalLine, startToolDisplay, formatResultLines, printResultLines } = require('../ui/tool_display');
 const { createConversation, updateConversation } = require('../storage/conversation_store');
 const { applyUsage, normalizeTokenUsage } = require('../utils/token_usage');
-const { gray, cyan, green } = require('../utils/colors');
+const { gray, cyan, green, red } = require('../utils/colors');
 const { createIndentedWriter } = require('../ui/indented_writer');
 const { createStatusBar } = require('../ui/status_bar');
 
@@ -52,6 +52,10 @@ console.log('');
 
 let rl = null;
 let statusBar = null;
+let isRunning = false;
+let escPendingCancel = false;
+let activeStreamController = null;
+let activeToolController = null;
 let commandMenuActive = false;
 let menuSearchTerm = '';
 let menuLastSearchTerm = '';
@@ -79,6 +83,16 @@ process.stdin.on('keypress', (str, key) => {
     }
     return;
   }
+  if (key && key.name === 'escape' && isRunning) {
+    escPendingCancel = true;
+    if (activeStreamController && typeof activeStreamController.abort === 'function') {
+      activeStreamController.abort();
+    }
+    if (activeToolController && typeof activeToolController.abort === 'function') {
+      activeToolController.abort();
+    }
+    return;
+  }
   if (str === '/' && (rl.line === '' || rl.line === '/')) {
     commandMenuActive = true;
     if (rl) {
@@ -275,13 +289,30 @@ function buildApiMessages() {
   return messages;
 }
 
+function printCancelLine() {
+  console.log('');
+  process.stdout.write(`  ${red('已取消本次响应')}\n\n`);
+}
+
+function stopSpinnerForCancel(spinner, thinkingActive, showThinkingLabel, thinkingMode) {
+  if (!spinner) return;
+  if (thinkingMode && (thinkingActive || showThinkingLabel)) {
+    spinner.stop('∙ 停止思考');
+  } else {
+    spinner.stopSilent();
+  }
+}
+
 async function runAssistantLoop() {
   let continueLoop = true;
   let firstLoop = true;
   const hideCursor = () => process.stdout.write('\x1b[?25l');
   const showCursor = () => process.stdout.write('\x1b[?25h');
+  isRunning = true;
+  escPendingCancel = false;
   if (statusBar) statusBar.setMode('running');
-  while (continueLoop) {
+  try {
+    while (continueLoop) {
     hideCursor();
     const spinner = new Spinner(' ', firstLoop ? 1 : 0);
     firstLoop = false;
@@ -297,6 +328,7 @@ async function runAssistantLoop() {
     let usageCompletion = null;
     let firstContent = true;
     let assistantWriter = null;
+    let cancelled = false;
 
     const thinkingDelay = setTimeout(() => {
       if (state.thinkingMode) showThinkingLabel = true;
@@ -309,8 +341,10 @@ async function runAssistantLoop() {
     });
 
     const messages = buildApiMessages();
+    const streamController = new AbortController();
+    activeStreamController = streamController;
     try {
-      for await (const chunk of streamChat({ config, messages, tools, thinkingMode: state.thinkingMode })) {
+      for await (const chunk of streamChat({ config, messages, tools, thinkingMode: state.thinkingMode, abortSignal: streamController.signal })) {
         const choice = chunk.choices && chunk.choices[0];
         if (!choice) continue;
         const usage = (choice && (choice.usage || choice.delta?.usage)) || chunk.usage;
@@ -364,6 +398,9 @@ async function runAssistantLoop() {
         }
       }
     } catch (err) {
+      if (err && (err.code === 'aborted' || err.name === 'AbortError' || escPendingCancel)) {
+        cancelled = true;
+      } else {
       clearTimeout(thinkingDelay);
       if (state.thinkingMode) {
         spinner.stop('○');
@@ -373,9 +410,17 @@ async function runAssistantLoop() {
       showCursor();
       console.log(`错误: ${err.message || err}`);
       return;
+      }
     }
+    activeStreamController = null;
 
     clearTimeout(thinkingDelay);
+    if (cancelled) {
+      stopSpinnerForCancel(spinner, thinkingActive, showThinkingLabel, state.thinkingMode);
+      showCursor();
+      printCancelLine();
+      return;
+    }
     if (!gotAnswer) {
       if (state.thinkingMode) {
         spinner.stop(thinkingActive ? '○ 思考完成' : '○');
@@ -415,7 +460,16 @@ async function runAssistantLoop() {
         const startLine = buildStartLine(call.function.name, args);
         const finalLine = buildFinalLine(call.function.name, args);
         const indicator = startToolDisplay(startLine);
-        const toolResult = await executeTool({ workspace: WORKSPACE, config, allowMode: state.allowMode, toolCall: call });
+        const toolController = new AbortController();
+        activeToolController = toolController;
+        const toolResult = await executeTool({
+          workspace: WORKSPACE,
+          config,
+          allowMode: state.allowMode,
+          toolCall: call,
+          abortSignal: toolController.signal,
+        });
+        activeToolController = null;
         indicator.stop(finalLine);
         const resultLines = formatResultLines(call.function.name, args, toolResult.raw || { success: toolResult.success, error: toolResult.error });
         printResultLines(resultLines);
@@ -428,6 +482,9 @@ async function runAssistantLoop() {
         };
         state.messages.push(toolMsg);
         persistConversation();
+        if (escPendingCancel || (toolResult.raw && toolResult.raw.cancelled)) {
+          return;
+        }
       }
       continueLoop = true;
       hideCursor();
@@ -441,7 +498,13 @@ async function runAssistantLoop() {
       persistConversation();
     }
     continueLoop = false;
+    }
+  } finally {
+    showCursor();
+    isRunning = false;
+    escPendingCancel = false;
+    activeStreamController = null;
+    activeToolController = null;
+    if (statusBar) statusBar.setMode('input');
   }
-  showCursor();
-  if (statusBar) statusBar.setMode('input');
 }