diff --git a/gcode/agent.py b/gcode/agent.py index 20b1a05..606bd0f 100644 --- a/gcode/agent.py +++ b/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) diff --git a/gcode/tools.py b/gcode/tools.py index b08cfa0..e919f37 100644 --- a/gcode/tools.py +++ b/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(): diff --git a/tests/test_agent.py b/tests/test_agent.py new file mode 100644 index 0000000..ace84fd --- /dev/null +++ b/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.") diff --git a/tests/test_tools.py b/tests/test_tools.py index 92183d3..739ad43 100644 --- a/tests/test_tools.py +++ b/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