fix: Handle CancelledError in bash and task tools

src/kimi_cli/tools/bash/__init__.py

@@ -1,4 +1,5 @@
 import asyncio
+import contextlib
 from pathlib import Path
 from typing import override
 
@@ -69,6 +70,11 @@ def stderr_cb(line: bytes):
                 f"Command killed by timeout ({params.timeout}s)",
                 brief=f"Killed by timeout ({params.timeout}s)",
             )
+        except asyncio.CancelledError:
+            return builder.error(
+                "Command cancelled by user",
+                brief="Cancelled by user",
+            )
 
 
 async def _stream_subprocess(command: str, stdout_cb, stderr_cb, timeout: int) -> int:
@@ -97,3 +103,10 @@ async def _read_stream(stream, cb):
     except TimeoutError:
         process.kill()
         raise
+    except asyncio.CancelledError:
+        # Handle user cancellation (Ctrl+C or ESC key)
+        process.kill()
+        with contextlib.suppress(Exception):
+            # Wait for process to exit with protection, ignore cleanup failures
+            await asyncio.shield(asyncio.wait_for(process.wait(), timeout=2.0))
+        raise

src/kimi_cli/tools/task/__init__.py

@@ -6,13 +6,13 @@
 from pydantic import BaseModel, Field
 
 from kimi_cli.agent import Agent, AgentGlobals, AgentSpec, load_agent
-from kimi_cli.soul import MaxStepsReached, get_wire_or_none, run_soul
+from kimi_cli.soul import MaxStepsReached
 from kimi_cli.soul.context import Context
 from kimi_cli.soul.kimisoul import KimiSoul
 from kimi_cli.tools.utils import load_desc
 from kimi_cli.utils.message import message_extract_text
 from kimi_cli.utils.path import next_available_rotation
-from kimi_cli.wire import WireUISide
+from kimi_cli.wire import WireUISide, get_wire_or_none
 from kimi_cli.wire.message import ApprovalRequest, WireMessage
 
 # Maximum continuation attempts for task summary
@@ -101,19 +101,6 @@ async def __call__(self, params: Params) -> ToolReturnType:
 
     async def _run_subagent(self, agent: Agent, prompt: str) -> ToolReturnType:
         """Run subagent with optional continuation for task summary."""
-        super_wire = get_wire_or_none()
-        assert super_wire is not None
-
-        def _super_wire_send(msg: WireMessage) -> None:
-            if isinstance(msg, ApprovalRequest):
-                super_wire.soul_side.send(msg)
-            # TODO: visualize subagent behavior by sending other messages in some way
-
-        async def _ui_loop_fn(wire: WireUISide) -> None:
-            while True:
-                msg = await wire.receive()
-                _super_wire_send(msg)
-
         subagent_history_file = await self._get_subagent_history_file()
         context = Context(file_backend=subagent_history_file)
         soul = KimiSoul(
@@ -122,9 +109,12 @@ async def _ui_loop_fn(wire: WireUISide) -> None:
             context=context,
             loop_control=self._agent_globals.config.loop_control,
         )
+        wire = get_wire_or_none()
+        assert wire is not None, "Wire is expected to be set"
+        sub_wire = _SubWire(wire)
 
         try:
-            await run_soul(soul, prompt, _ui_loop_fn, asyncio.Event())
+            await soul.run(prompt, sub_wire)
         except MaxStepsReached as e:
             return ToolError(
                 message=(
@@ -133,6 +123,11 @@ async def _ui_loop_fn(wire: WireUISide) -> None:
                 ),
                 brief="Max steps reached",
             )
+        except asyncio.CancelledError:
+            return ToolError(
+                message="Subagent task cancelled by user",
+                brief="Cancelled by user",
+            )
 
         _error_msg = (
             "The subagent seemed not to run properly. Maybe you have to do the task yourself."
@@ -147,10 +142,28 @@ async def _ui_loop_fn(wire: WireUISide) -> None:
         # Check if response is too brief, if so, run again with continuation prompt
         n_attempts_remaining = MAX_CONTINUE_ATTEMPTS
         if len(final_response) < 200 and n_attempts_remaining > 0:
-            await run_soul(soul, CONTINUE_PROMPT, _ui_loop_fn, asyncio.Event())
+            try:
+                await soul.run(CONTINUE_PROMPT, sub_wire)
+            except asyncio.CancelledError:
+                return ToolError(
+                    message="Subagent task cancelled by user",
+                    brief="Cancelled by user",
+                )
 
             if len(context.history) == 0 or context.history[-1].role != "assistant":
                 return ToolError(message=_error_msg, brief="Failed to run subagent")
             final_response = message_extract_text(context.history[-1])
 
         return ToolOk(output=final_response)
+
+
+class _SubWire(WireUISide):
+    def __init__(self, super_wire: WireUISide):
+        super().__init__()
+        self._super_wire = super_wire
+
+    @override
+    def send(self, msg: WireMessage):
+        if isinstance(msg, ApprovalRequest):
+            self._super_wire.send(msg)
+        # TODO: visualize subagent behavior by sending other messages in some way