refactor: Improve interrupt handling with interruptible wrapper

Co-authored-by: cecli (openai/gemini_cli_local/gemini-2.5-pro)

Author: Your Name

cecli/coders/agent_coder.py

@@ -16,6 +16,7 @@
 from cecli.helpers import nested, responses
 from cecli.helpers.background_commands import BackgroundCommandManager
 from cecli.helpers.conversation import ConversationService, MessageTag
+from cecli.helpers.coroutines import interruptible
 from cecli.helpers.similarity import (
     cosine_similarity,
     create_bigram_vector,
@@ -301,25 +302,15 @@ async def _execute_local_tool_calls(self, tool_calls_list):
                 else:
                     all_results_content.append(f"Error: Unknown tool name '{tool_name}'")
                 if tasks:
-                    gather_future = asyncio.gather(*tasks, return_exceptions=True)
-                    interrupt_task = asyncio.create_task(self.interrupt_event.wait())
-
-                    done, pending = await asyncio.wait(
-                        {gather_future, interrupt_task},
-                        return_when=asyncio.FIRST_COMPLETED,
+                    gather_coro = asyncio.gather(*tasks, return_exceptions=True)
+                    task_results, interrupted = await interruptible(
+                        gather_coro, self.interrupt_event
                     )
 
-                    if interrupt_task in done:
-                        gather_future.cancel()
-                        try:
-                            await gather_future
-                        except asyncio.CancelledError:
-                            pass
+                    if interrupted:
                         self.io.tool_warning("Tool execution interrupted.")
-                        # Append a message indicating interruption
                         all_results_content.append("Tool execution interrupted by user.")
-                    else:
-                        task_results = gather_future.result()
+                    elif task_results:
                         for res in task_results:
                             if isinstance(res, Exception):
                                 all_results_content.append(f"Error in tool execution: {res}")
@@ -415,24 +406,11 @@ async def _exec_async():
 """)
                 return f"Error executing tool call {tool_name}: {e}"
 
-        exec_future = asyncio.create_task(_exec_async())
-        interrupt_task = asyncio.create_task(self.interrupt_event.wait())
-
-        done, pending = await asyncio.wait(
-            {exec_future, interrupt_task},
-            return_when=asyncio.FIRST_COMPLETED,
-        )
+        result, interrupted = await interruptible(_exec_async(), self.interrupt_event)
 
-        if interrupt_task in done:
-            exec_future.cancel()
-            try:
-                await exec_future
-            except asyncio.CancelledError:
-                pass
+        if interrupted:
             return "Tool execution interrupted by user."
-        else:
-            interrupt_task.cancel()
-            return await exec_future
+        return result
 
     def _calculate_context_block_tokens(self, force=False):
         """

cecli/coders/base_coder.py

@@ -2722,6 +2722,9 @@ async def _execute_mcp_tools(self, server, tool_calls):
                 tool_responses.append(
                     {"role": "tool", "tool_call_id": tool_call.id, "content": connection_error}
                 )
+        except asyncio.CancelledError:
+            # Re-raise CancelledError to ensure the task cancellation propagates
+            raise
         except Exception as e:
             connection_error = f"Could not connect to server {server.name}\n{e}"
             self.io.tool_warning(connection_error)

cecli/helpers/coroutines.py

@@ -1,8 +1,45 @@
-import asyncio  # noqa: F401
+import asyncio
 
 
-def is_active(coroutine):
-    if not coroutine or coroutine.done() or coroutine.cancelled():
+def is_active(task):
+    if not task or task.done() or task.cancelled():
         return False
 
     return True
+
+
+async def interruptible(coroutine, interrupt_event):
+    """
+    Runs a coroutine and allows it to be interrupted by an asyncio.Event.
+
+    Args:
+        coroutine: The coroutine to run.
+        interrupt_event: The asyncio.Event that signals an interruption.
+
+    Returns:
+        A tuple of (result, interrupted).
+        - If not interrupted: (coroutine_result, False)
+        - If interrupted: (None, True)
+    """
+    main_task = asyncio.create_task(coroutine)
+    interrupt_task = asyncio.create_task(interrupt_event.wait())
+
+    done, pending = await asyncio.wait(
+        {main_task, interrupt_task},
+        return_when=asyncio.FIRST_COMPLETED,
+    )
+
+    for task in pending:
+        task.cancel()
+        try:
+            await task
+        except asyncio.CancelledError:
+            pass  # Expected
+
+    if interrupt_task in done:
+        return None, True
+
+    try:
+        return main_task.result(), False
+    except asyncio.CancelledError:
+        return None, True