diff --git a/docs/DESIGN.md b/docs/DESIGN.md index 45799a4..a9293aa 100644 --- a/docs/DESIGN.md +++ b/docs/DESIGN.md @@ -2865,7 +2865,7 @@ Ctrl-C aborts a model turn mid-stream — the motivating case is a long invisibl **`AppUI.abort_turn` (app-side, not a RuntimeUI method).** Clears the dangling live indicator (`self._indicator = None`) then appends a `⏹ aborted` marker (ASCII fallback `Glyphs.cross`, style `sp.warn`). `_add_renderable` closes the open response first, so the partial streamed text stays visible — finalized — with the marker below it. This is display-only; the history discard is the runtime's. -**Out of scope (branch 6b — subprocess killpg).** A cancel requested while a tool *command* is running does NOT kill that subprocess; the command finishes/times out, then the next model call's stream-read sees the set event and aborts the turn. Killing the running subprocess on cancel is deferred. +**Subprocess cancellation (branch 6b — killpg).** A cancel requested while a tool *command* is running now kills that child immediately, instead of waiting out its timeout. The per-turn cancel event threads through `ToolContext` (`base.py`) to `run_command_process`, whose wait loop polls it on a `_POLL_SECONDS` (0.1 s) interval and `killpg`s the child's process group the instant the event is set — the worker owns its own `Popen`, so it kills its own child (no cross-thread registry; the `Event` stays the only cross-thread signal). The tool loop then raises `GenerationCancelled` keyed off the cancel `Event` (the general check fires for any tool, immediately after `executor.execute`, before the plan-guard bookkeeping), routing through the same `abort_turn` (`⏹ aborted`) path as the model-stream cancel. Before raising, it rolls this model step out of history (`del self._history[history_before_reply:]`) so no orphaned tool_call — an assistant `tool_call` with no matching result — is re-sent on the next turn; this matches the model-stream cancel, which never records its partial reply. Prior completed steps stay, and partial command output already streamed to the pane stays visible. The `cancel=None` legacy path (the default REPL executor) is behaviorally equivalent to before — same `exit_code`/`output`/`timed_out`/`truncated`; the poll loop merely chunks the wait (a runaway child is still killed at the timeout deadline, within one poll interval). ### 31.16 Approval focus-swap (v2) diff --git a/shellpilot/cli/app_turn.py b/shellpilot/cli/app_turn.py index cee8e8d..e012c3b 100644 --- a/shellpilot/cli/app_turn.py +++ b/shellpilot/cli/app_turn.py @@ -297,10 +297,10 @@ def _run(self, text: str, cancel: threading.Event) -> None: (success/cancel/error) so a turn never wedges the app. NOTE (branch 6b — subprocess killpg): a cancel requested while a tool - COMMAND is running does NOT kill that subprocess here; the command - finishes/times out, then the NEXT model call's stream-read sees the set - event and raises GenerationCancelled, aborting the turn. Killing the - running subprocess is branch 6b. + COMMAND is running now kills that child immediately. The same cancel + event threads through ToolContext to run_command_process, which polls it + and killpg's the child's process group; the tool loop then raises + GenerationCancelled, reaching this same clean-abort path (§31.15). """ try: conversation = self.conversation diff --git a/shellpilot/runtime/conversation.py b/shellpilot/runtime/conversation.py index 1aa107c..c1b4a37 100644 --- a/shellpilot/runtime/conversation.py +++ b/shellpilot/runtime/conversation.py @@ -12,7 +12,7 @@ from urllib.parse import urlsplit from shellpilot.config.model import Settings, is_egressing -from shellpilot.llm.client import LLMClient +from shellpilot.llm.client import GenerationCancelled, LLMClient from shellpilot.llm.messages import ImageRef, Message, tool_result, user from shellpilot.llm.ollama import encode_tool from shellpilot.memory.agents_md import BehaviorInstructions @@ -595,6 +595,7 @@ def _tool_loop(self) -> Message: snapshots=self.snapshots, audit=self._audit, allow_sensitive_reads=self._settings.privacy.allow_sensitive_reads, + cancel=self._cancel, ) tools = executor.available_definitions() tool_turns = 0 @@ -642,6 +643,10 @@ def _tool_loop(self) -> Message: finally: self._ui.end_response() self._turn_output_tokens += reply.output_tokens + # History length BEFORE this model step is recorded, so a mid-tool + # cancel (below) can roll the step back out and leave no orphaned + # tool_call behind (§31.15). + history_before_reply = len(self._history) self._record(reply) if not reply.tool_calls: pending = self._pending_plan_step() @@ -705,6 +710,18 @@ def _tool_loop(self) -> Message: for call in reply.tool_calls: self._ui.show_tool_call(call.name, call.arguments) outcome = executor.execute(call) + if self._cancel is not None and self._cancel.is_set(): + # A Ctrl-C during tool execution (e.g. a long run_command just + # killed by the cancel signal) aborts the turn. Roll THIS model + # step's reply + any partial tool results back out of history so + # no orphaned tool_call (an assistant tool_call with no matching + # result) is re-sent on the next turn — the same clean discard as + # the model-stream cancel, which never records its partial reply. + # Prior completed steps stay; the worker then routes through + # abort_turn (⏹ aborted), and partial command output already + # streamed to the pane stays visible (§31.15). + del self._history[history_before_reply:] + raise GenerationCancelled # Feed side-effecting tool outcomes to the plan completion guard. # Key off the SPEC's side_effect (not the result's): a failed or # denied result carries SideEffect.NONE, which would hide exactly diff --git a/shellpilot/runtime/executor.py b/shellpilot/runtime/executor.py index 43d2c9c..473fd2d 100644 --- a/shellpilot/runtime/executor.py +++ b/shellpilot/runtime/executor.py @@ -7,6 +7,7 @@ from __future__ import annotations +import threading from collections.abc import Callable from dataclasses import dataclass from pathlib import Path @@ -65,6 +66,7 @@ def __init__( snapshots: SnapshotStore | None = None, audit: AuditLogger | None = None, allow_sensitive_reads: str = "ask", + cancel: threading.Event | None = None, ) -> None: self._snapshots = snapshots self._audit = audit @@ -78,6 +80,7 @@ def __init__( self._ask_approval = ask_approval self._emit_output = emit_output self._allow_sensitive_reads = allow_sensitive_reads + self._cancel = cancel self._spent_tokens = 0 def available_definitions(self) -> list[ToolDefinition]: @@ -108,6 +111,7 @@ def execute(self, call: ToolCall) -> ExecutionOutcome: emit_output=self._emit_output, snapshots=self._snapshots, allow_sensitive_reads=self._allow_sensitive_reads, + cancel=self._cancel, ) if spec.precheck is not None: diff --git a/shellpilot/tools/base.py b/shellpilot/tools/base.py index 9536e7d..960ca3d 100644 --- a/shellpilot/tools/base.py +++ b/shellpilot/tools/base.py @@ -2,6 +2,7 @@ from __future__ import annotations +import threading from collections.abc import Callable from dataclasses import dataclass, field from pathlib import Path @@ -42,6 +43,10 @@ class ToolContext: # Hard ceiling for run_command timeout; model may request shorter but never # longer (design section 13.1). command_timeout_seconds: int = 600 + # Branch-6b turn-abort signal (§31.15): set by the app's Ctrl-C so a running + # run_command child can be killed mid-execution instead of waiting out its + # timeout. + cancel: threading.Event | None = None @dataclass(frozen=True) diff --git a/shellpilot/tools/command.py b/shellpilot/tools/command.py index 82da24f..57fdb36 100644 --- a/shellpilot/tools/command.py +++ b/shellpilot/tools/command.py @@ -14,6 +14,7 @@ import signal import subprocess import threading +import time from collections.abc import Callable from dataclasses import dataclass from pathlib import Path @@ -28,6 +29,9 @@ # Maximum chars read in a single readline() call so a newline-less stream cannot # materialise an unbounded string before the total-capture cap is consulted. MAX_READ_CHARS = 65_536 +# NOTE: wait-loop poll interval (§31.15). Caps the Ctrl-C kill latency for a +# running child at ~0.1 s; tighten only if that ceiling ever matters. +_POLL_SECONDS = 0.1 # Exit codes that mean "ran fine, found nothing" rather than failure (section 24.3). EXPECTED_NONZERO: dict[str, frozenset[int]] = { "grep": frozenset({1}), @@ -101,6 +105,7 @@ def run_command_process( *, max_capture_chars: int, emit_line: Callable[[str], None] | None = None, + cancel: threading.Event | None = None, ) -> CommandOutcome: """Run argv with shell=False, streaming output, bounding capture, killing on timeout.""" # NOTE (Investigation C, v0.5.1): the MallocStackLogging fork-window line cannot @@ -154,16 +159,33 @@ def reader() -> None: thread = threading.Thread(target=reader, daemon=True) thread.start() - timed_out = False - try: - process.wait(timeout=request.timeout_seconds) - except subprocess.TimeoutExpired: - timed_out = True + def _kill_group() -> None: try: os.killpg(os.getpgid(process.pid), signal.SIGKILL) except (ProcessLookupError, PermissionError): process.kill() - process.wait(timeout=5) + + timed_out = False + deadline = time.monotonic() + request.timeout_seconds + while True: + try: + process.wait(timeout=_POLL_SECONDS) + break # exited on its own + except subprocess.TimeoutExpired: + # Ctrl-C mid-command: kill the child's whole process group now rather + # than waiting out the timeout (§31.15). The worker owns this Popen, so + # it kills its own child — no cross-thread signal, no registry. The turn + # abort is driven by the cancel Event in the tool loop, not by the + # outcome, so no "cancelled" flag is recorded here. + if cancel is not None and cancel.is_set(): + _kill_group() + process.wait(timeout=5) + break + if time.monotonic() >= deadline: + timed_out = True + _kill_group() + process.wait(timeout=5) + break thread.join(timeout=5) return CommandOutcome( @@ -299,6 +321,7 @@ def _run_command(context: ToolContext, arguments: dict[str, Any]) -> ToolResult: ), max_capture_chars=context.max_capture_chars, emit_line=context.emit_output, + cancel=context.cancel, ) except OSError as exc: return ToolResult(success=False, summary=f"could not start command: {exc}", content="") diff --git a/tests/test_app_turn.py b/tests/test_app_turn.py index df19f14..762ebb9 100644 --- a/tests/test_app_turn.py +++ b/tests/test_app_turn.py @@ -35,7 +35,7 @@ from shellpilot.memory.agents_md import BehaviorInstructions from shellpilot.runtime.conversation import ConversationRuntime from shellpilot.runtime.events import TurnStats -from tests.fakes.fake_llm import FakeLLM, answer +from tests.fakes.fake_llm import FakeLLM, answer, tool_call # --- helpers ------------------------------------------------------------------ @@ -491,6 +491,87 @@ def test_request_cancel_aborts_turn_cleanly_and_reruns(tmp_path: Path) -> None: assert runner._busy is False +def test_request_cancel_kills_running_command_and_aborts(tmp_path: Path) -> None: + """Branch 6b (§31.15): Ctrl-C during a running tool aborts cleanly and reruns. + + A blocking tool stands in for a long run_command child: its handler waits on + the turn's cancel event, so request_cancel both releases it (a real kill in + production) and trips the tool-loop abort — reaching abort_turn, never the + failure path. A fresh turn then completes normally. + """ + from shellpilot.llm.messages import ToolDefinition + from shellpilot.policy.risk import RiskLevel, SideEffect + from shellpilot.tools.base import ALL_PROFILES, ToolContext, ToolResult, ToolSpec + from shellpilot.tools.registry import ToolRegistry + + app_ui = AppUI(glyphs=UNICODE_GLYPHS, workspace=tmp_path, width_fn=lambda: 80) + inner = _RecordingUI(forward=app_ui) + q: queue.Queue[Scheduled] = queue.Queue() + runner = TurnRunner(inner_ui=inner, schedule=q.put) + threaded = ThreadedUI(inner=inner, schedule=q.put) + + entered = threading.Event() + + def _blocking_handler(context: ToolContext, arguments: dict[str, object]) -> ToolResult: + entered.set() + assert context.cancel is not None + context.cancel.wait(5.0) # blocks until request_cancel sets the turn's event + return ToolResult(success=True, summary="unblocked", content="") + + spec = ToolSpec( + definition=ToolDefinition(name="block_tool", description="d", parameters={}, required=()), + side_effect=SideEffect.NONE, + default_risk=RiskLevel.LOW, + allowed_profiles=ALL_PROFILES, + handler=_blocking_handler, + ) + registry = ToolRegistry() + registry.register(spec) + fake = FakeLLM(script=[tool_call("block_tool"), answer("a normal completion")]) + runtime = ConversationRuntime( + llm=fake, + settings=Settings(), + workspace=tmp_path, + behavior=BehaviorInstructions(global_text=None, project_text=None), + ui=threaded, + registry=registry, + ) + runner.conversation = runtime + + runner.start("run something slow") + assert entered.wait(5.0) # the worker is inside the blocking tool → in flight + assert runner._busy is True + assert runner.request_cancel() is True # sets the turn's cancel → tool unblocks + + assert runner._thread is not None + runner._thread.join(5.0) + assert not runner._thread.is_alive() # aborted cleanly, not killed mid-stack + while not q.empty(): + q.get()() + + names = inner.names() + assert "abort_turn" in names # the CLEAN abort path ran ... + assert "show_error" not in names # ... NOT the "Turn failed" error path + assert "turn_finished" not in names # the turn did not complete + assert runner._busy is False + assert "aborted" in app_ui._render_ansi() + # Isolates THIS branch's tool-loop raise from branch-6's model-stream backstop: + # without the raise the block_tool's "unblocked" result would be recorded and + # the model re-invoked one round later (the backstop still aborts, but leaves + # assistant+tool messages behind). The raise + history rollback leave only the + # user message — no recorded tool result, no orphaned tool_call. + assert [m.role for m in runtime._history] == ["user"] + + # A fresh turn completes normally after the cancel. + runner.start("again") + assert runner._thread is not None + runner._thread.join(5.0) + while not q.empty(): + q.get()() + assert "turn_finished" in inner.names() + assert runner._busy is False + + # --- build_app on_submit wiring (headless, pipe input) ------------------------ diff --git a/tests/test_conversation.py b/tests/test_conversation.py index 4e5f36d..4b5c003 100644 --- a/tests/test_conversation.py +++ b/tests/test_conversation.py @@ -14,12 +14,15 @@ ToolSettings, ) from shellpilot.llm.client import GenerationCancelled -from shellpilot.llm.messages import Message +from shellpilot.llm.messages import Message, ToolDefinition from shellpilot.memory.agents_md import BehaviorInstructions from shellpilot.persistence.audit_store import AuditLogger +from shellpilot.policy.risk import RiskLevel, SideEffect from shellpilot.runtime.conversation import ConversationRuntime from shellpilot.skills.loader import discover_skills from shellpilot.skills.model import Skill, SkillTrigger +from shellpilot.tools.base import ALL_PROFILES, ToolContext, ToolResult, ToolSpec +from shellpilot.tools.registry import ToolRegistry from tests.fakes.fake_llm import FakeLLM, answer, tool_call from tests.fakes.fake_ui import FakeUI @@ -84,6 +87,60 @@ def test_cancelled_turn_discards_partial_reply(tmp_path: Path) -> None: assert all(m.role != "assistant" for m in runtime._history) +def test_cancel_during_tool_execution_aborts(tmp_path: Path) -> None: + """Branch 6b (§31.15): a Ctrl-C landing during tool execution aborts the turn. + + A fake tool sets the turn's cancel event via its ToolContext (simulating a + long run_command child being killed mid-execution). The tool loop then raises + GenerationCancelled, so the cancelled tool's RESULT is never recorded and the + model is not re-invoked. + """ + cancel_seen: list[bool] = [] + + def _set_cancel(context: ToolContext, arguments: dict[str, object]) -> ToolResult: + cancel_seen.append(context.cancel is not None) + assert context.cancel is not None # the turn's cancel threaded to the handler + context.cancel.set() + return ToolResult(success=True, summary="cancelled mid-run", content="discarded") + + spec = ToolSpec( + definition=ToolDefinition(name="slow_tool", description="d", parameters={}, required=()), + side_effect=SideEffect.NONE, + default_risk=RiskLevel.LOW, + allowed_profiles=ALL_PROFILES, + handler=_set_cancel, + ) + registry = ToolRegistry() + registry.register(spec) + + fake = FakeLLM(script=[tool_call("slow_tool"), answer("must never be recorded")]) + runtime = ConversationRuntime( + llm=fake, + settings=Settings(), + workspace=tmp_path, + behavior=BehaviorInstructions(global_text=None, project_text=None), + ui=FakeUI(), + registry=registry, + ) + cancel = threading.Event() # starts UNSET; the tool sets it + + with pytest.raises(GenerationCancelled): + runtime.run_turn("go", cancel=cancel) + + assert cancel_seen == [True] # the handler saw the cancel event + # The cancelled model step is rolled fully out of history: only the user + # message remains — no orphaned assistant tool_call (an assistant reply whose + # tool_call has no matching result) survives to be re-sent next turn, and no + # tool result is recorded. Matches the model-stream cancel's clean discard, + # which never records its partial reply. + assert [m.role for m in runtime._history] == ["user"] + assert runtime._history[0].content == "go" + assert all(not m.tool_calls for m in runtime._history) + # ... and no follow-up answer landed — the model was not re-invoked. + assert all("must never be recorded" not in m.content for m in runtime._history) + assert len(fake.calls) == 1 + + def test_normal_turn_records_assistant_reply(tmp_path: Path) -> None: """Control for the cancel test: an uncancelled turn still records the reply.""" fake = FakeLLM(script=[answer("the answer")]) diff --git a/tests/test_run_command.py b/tests/test_run_command.py index b87fe66..a7b5462 100644 --- a/tests/test_run_command.py +++ b/tests/test_run_command.py @@ -1,7 +1,9 @@ """Tests for the shell=False command runner (design section 13.1).""" import os +import subprocess import sys +import threading import time from pathlib import Path from typing import Any @@ -86,6 +88,67 @@ def test_timeout_kills_process_group(tmp_path: Path) -> None: assert elapsed < 10 +def test_cancel_kills_process_group_fast(tmp_path: Path) -> None: + # Branch 6b (§31.15): a cancel event set mid-command kills the child's whole + # process group at once, instead of waiting out the (here, 30 s) timeout. + cancel = threading.Event() + captured: list[subprocess.Popen[str]] = [] + real_popen = subprocess.Popen + + def _capture(*args: Any, **kwargs: Any) -> subprocess.Popen[str]: + proc = real_popen(*args, **kwargs) + captured.append(proc) + return proc + + baseline_threads = threading.active_count() + setter = threading.Thread(target=lambda: (time.sleep(0.2), cancel.set())) # type: ignore[func-returns-value] + start = time.monotonic() + setter.start() + with patch("shellpilot.tools.command.subprocess.Popen", side_effect=_capture): + outcome = run_command_process( + CommandRequest( + argv=[sys.executable, "-c", "import time; time.sleep(30)"], + cwd=tmp_path, + timeout_seconds=30, + ), + max_capture_chars=1000, + cancel=cancel, + ) + elapsed = time.monotonic() - start + setter.join() + + # The cancel-kill is proven by the fast return (not the 30 s sleep) + the dead + # process group below, not by a flag: the turn abort is driven by the cancel + # Event in the tool loop, so the command outcome carries no "cancelled" flag. + assert outcome.timed_out is False + assert elapsed < 8 # returned promptly, not after the 30 s sleep + + # The child's process group was SIGKILLed: start_new_session=True makes the + # child its own group leader (pgid == pid), so signalling it now raises. + pgid = captured[0].pid + try: + os.killpg(pgid, 0) + raise AssertionError("process group still alive after cancel") + except ProcessLookupError: + pass + + # The reader thread was joined — no leak (active count back to baseline). + assert threading.active_count() == baseline_threads + + +def test_cancel_none_completes_normally(tmp_path: Path) -> None: + # The legacy path (cancel=None, e.g. the default REPL executor) is unchanged: + # a fast command exits normally. + outcome = run_command_process( + CommandRequest(argv=["echo", "hi"], cwd=tmp_path, timeout_seconds=30), + max_capture_chars=1000, + cancel=None, + ) + assert outcome.timed_out is False + assert outcome.exit_code == 0 + assert "hi" in outcome.output + + def test_output_capture_is_bounded(tmp_path: Path) -> None: # A single newline-less chunk far larger than the cap must be hard-bounded to # exactly max_capture_chars with truncated=True — not appended whole (which