fix: cancel current command/stream on Ctrl+C instead of exiting

gcode/agent.py

@@ -58,22 +58,33 @@ def trim_history(messages: list) -> None:
 
 def _stream(messages: list, model, ui) -> AIMessage:
     """Stream one model response, forwarding text to the UI, and return the
-    accumulated message (with ``tool_calls`` populated)."""
+    accumulated message (with ``tool_calls`` populated).
+
+    A Ctrl+C during streaming stops the stream but keeps the session alive:
+    whatever was accumulated so far is returned as the turn's assistant
+    message so it gets persisted with the rest of the history.
+    """
     ui.assistant_start()
     accumulated = None
-    for chunk in model.stream(messages):
-        if not isinstance(chunk, AIMessageChunk):
-            continue
-        if chunk.content:
-            ui.token(chunk.content)
-        accumulated = chunk if accumulated is None else accumulated + chunk
+    interrupted = False
+    try:
+        for chunk in model.stream(messages):
+            if not isinstance(chunk, AIMessageChunk):
+                continue
+            if chunk.content:
+                ui.token(chunk.content)
+            accumulated = chunk if accumulated is None else accumulated + chunk
+    except KeyboardInterrupt:
+        interrupted = True
     if accumulated is None:
         accumulated = AIMessageChunk(content="")
     ui.assistant_end()
+    if interrupted:
+        ui.info("(streaming stopped by user)")
     # Store the canonical AIMessage (not the chunk) for clean history + reloads.
     return AIMessage(
         content=accumulated.content,
-        tool_calls=accumulated.tool_calls,
+        tool_calls=[] if interrupted else accumulated.tool_calls,
         additional_kwargs=accumulated.additional_kwargs,
         id=accumulated.id,
     )
@@ -92,6 +103,11 @@ def _run_tool(tool_name: str, tool_args: dict, ui) -> str:
     else:
         try:
             result = fn.invoke(tool_args)
+        except KeyboardInterrupt:
+            # Ctrl+C during a tool call cancels that call and keeps the
+            # session alive; the model sees a cancelled result instead of the
+            # whole REPL dying.
+            result = "Command execution cancelled by user."
         except Exception as exc:
             result = f"Tool {tool_name} raised: {exc}"
     ui.tool_result(tool_name, result)

gcode/tools.py

@@ -68,6 +68,10 @@ def execute_bash(command: str) -> str:
         )
     except subprocess.TimeoutExpired:
         return f"Command timed out after {BASH_TIMEOUT}s: {command}"
+    except KeyboardInterrupt:
+        # Ctrl+C while the command is running cancels just this command and
+        # keeps the session alive, instead of killing the whole REPL.
+        return "Command execution cancelled by user."
 
     output = result.stdout.strip()
     if result.stderr.strip():

tests/test_agent.py

@@ -0,0 +1,73 @@
+"""Unit tests for the agent loop, focused on Ctrl+C interrupt handling."""
+
+from unittest.mock import Mock, patch
+
+from gcode.agent import _run_tool, _stream
+from langchain_core.messages import AIMessage, AIMessageChunk
+
+
+class _FakeUI:
+    """Minimal UI stub recording calls for the streaming paths under test."""
+
+    def __init__(self):
+        self.calls = []
+
+    def assistant_start(self):
+        self.calls.append("assistant_start")
+
+    def token(self, text):
+        self.calls.append(("token", text))
+
+    def assistant_end(self):
+        self.calls.append("assistant_end")
+
+    def info(self, msg):
+        self.calls.append(("info", msg))
+
+
+class _InterruptingModel:
+    """A model whose stream yields one chunk and then raises KeyboardInterrupt."""
+
+    def stream(self, messages):
+        yield AIMessageChunk(content="partial reply ")
+        raise KeyboardInterrupt
+
+
+class _ToolRaisingInterrupt:
+    def invoke(self, tool_args):
+        raise KeyboardInterrupt
+
+
+def test_stream_keeps_partial_text_on_keyboard_interrupt():
+    ui = _FakeUI()
+    msg = _stream([], _InterruptingModel(), ui)
+
+    assert isinstance(msg, AIMessage)
+    assert msg.content == "partial reply "
+    # No half-formed tool calls are carried into history after an interrupt.
+    assert msg.tool_calls == []
+    assert ("info", "(streaming stopped by user)") in ui.calls
+    assert "assistant_end" in ui.calls
+
+
+def test_stream_interrupt_with_no_chunks_yet():
+    class _ImmediateInterrupt:
+        def stream(self, messages):
+            raise KeyboardInterrupt
+
+    ui = _FakeUI()
+    msg = _stream([], _ImmediateInterrupt(), ui)
+
+    assert isinstance(msg, AIMessage)
+    assert msg.content == ""
+    assert msg.tool_calls == []
+    assert ("info", "(streaming stopped by user)") in ui.calls
+
+
+def test_run_tool_returns_cancelled_on_keyboard_interrupt():
+    ui = Mock()
+    with patch("gcode.agent.TOOL_MAP", {"failing_tool": _ToolRaisingInterrupt()}):
+        result = _run_tool("failing_tool", {}, ui)
+
+    assert result == "Command execution cancelled by user."
+    ui.tool_result.assert_called_once_with("failing_tool", "Command execution cancelled by user.")

tests/test_tools.py

@@ -136,6 +136,24 @@ def test_execute_bash_cancels_on_keyboard_interrupt():
     assert out == "Command execution cancelled by user."
 
 
+def test_execute_bash_cancels_when_subprocess_interrupted():
+    """Ctrl+C while the command is actually running cancels just that command.
+
+    The interrupt arrives from ``subprocess.run`` (the prompt was already
+    approved), and must return a cancelled result instead of killing the
+    whole session.
+    """
+    from gcode.tools import AUTO_APPROVE, execute_bash, set_auto_approve
+
+    set_auto_approve(True)
+    try:
+        with patch("gcode.tools.subprocess.run", side_effect=KeyboardInterrupt):
+            out = execute_bash.invoke({"command": "sleep 100"})
+        assert out == "Command execution cancelled by user."
+    finally:
+        set_auto_approve(AUTO_APPROVE)
+
+
 def test_execute_bash_rejects_non_yes_answer():
     from gcode.tools import execute_bash