diff --git a/docs/DESIGN.md b/docs/DESIGN.md index f6e27ac..7ffaeb0 100644 --- a/docs/DESIGN.md +++ b/docs/DESIGN.md @@ -2847,7 +2847,25 @@ The full-screen pane shows a **live frontier line** while a turn runs, so a long **Reasoning gate.** The reasoning/total readout is gated on `settings.ui.show_reasoning_summary` (default true; previously a dead flag — this is its consumer), threaded into `AppUI(show_reasoning=…)` at both construction sites (`build_app`, and `terminal.py` app mode passing the setting). When false, the live line is `{plane} {phrase}… {N}s` and the done line is `✓ done · {N}s` (plane/phrase/timer only). This is display-only; audit capture of thinking is unaffected. -**Animation and caching.** `build_app` sets the `Application`'s `refresh_interval=0.1` so the timer ticks and the plane glides even between thinking chunks. While an indicator is active, `_render_ansi` bypasses the width cache entirely (elapsed changes every render, so it never reads or writes the cache); when idle, the per-tick render is a width-cache hit (a cheap no-op). `AppUI` stays pure state+render — the clock is injected and the repaint is the app's job, so it never calls `get_app()`/`invalidate()`. A turn that errors before `turn_finished` (or, later, a Ctrl-C cancel) leaves the indicator active; the full turn-failure/cancel robustness is a later branch. As a minimal cross-turn guard, `show_user_message` discards any such dangling indicator at the start of the next turn, so an errored turn never poisons the following turn's timer (without it, the next `begin_response` would be a no-op against the stale indicator and count from the old start). +**Animation and caching.** `build_app` sets the `Application`'s `refresh_interval=0.1` so the timer ticks and the plane glides even between thinking chunks. While an indicator is active, `_render_ansi` bypasses the width cache entirely (elapsed changes every render, so it never reads or writes the cache); when idle, the per-tick render is a width-cache hit (a cheap no-op). `AppUI` stays pure state+render — the clock is injected and the repaint is the app's job, so it never calls `get_app()`/`invalidate()`. A turn that errors before `turn_finished` (or a Ctrl-C cancel, §31.15) leaves the indicator active; as a minimal cross-turn guard, `show_user_message` discards any such dangling indicator at the start of the next turn, so an errored turn never poisons the following turn's timer (without it, the next `begin_response` would be a no-op against the stale indicator and count from the old start). + +### 31.15 Model-stream cancellation (v2) + +Ctrl-C aborts a model turn mid-stream — the motivating case is a long invisible thinking phase (§31.14) the user wants to stop. The cancellation contract is deterministic and history-safe. + +**Ctrl-C is a key press, not a SIGINT.** The full-screen app runs in raw mode (ISIG disabled), so `c-c` is an ordinary key binding. `build_app` takes `on_interrupt: Callable[[], bool] | None`; the `c-c` handler is `if on_interrupt is not None and on_interrupt(): return` (a turn was cancelled) `else: show_status(_IDLE_HINT)` (idle). `run_app` wires `on_interrupt=runner.request_cancel`. + +**Cancel event, set on the loop thread, observed on the worker.** `TurnRunner.start` creates a fresh `threading.Event` per turn and passes it into `run_turn(text, cancel=…)`. `request_cancel()` (loop thread) returns True and `Event.set()`s it when `_busy` and a cancel exists, else False — the only cross-thread signal (the app only ever *sets* the event; thread-safe). `ConversationRuntime.run_turn` reassigns `self._cancel = cancel` at the top (so a stale event can never leak across turns) and the tool loop passes it to every `self._llm.chat(...)`. + +**In-thread stream close + typed exception.** `OllamaClient._stream_chat` checks the event at the top of the `for line in response.iter_lines():` loop and raises `GenerationCancelled` (defined in `llm/client.py`, deliberately NOT an `OllamaError` subclass). Raising INSIDE the `with self._client.stream(...)` block closes the response in-thread via the context manager — the safe close; `response.close()` is never called cross-thread. Because `GenerationCancelled` is neither `httpx.TransportError` nor `OllamaResponseError`, it propagates cleanly out through `chat()` (and through the think-retry, which threads `cancel` into the retried call) to the runtime. + +**History integrity (load-bearing).** `run_turn` records the user message *before* the tool loop, then lets `GenerationCancelled` PROPAGATE — it never catches it, records a partial reply, or calls `turn_finished`. The model call raises at the `reply = self._llm.chat(...)` site, so the subsequent `self._record(reply)` is skipped: a cancelled turn leaves the user message in history but NO partial assistant reply. The partial is discarded, not truncated-and-stored. + +**Worker: clean abort vs failure.** `TurnRunner._run` catches `GenerationCancelled` SEPARATELY from `Exception`: cancel → `schedule(inner_ui.abort_turn)` (a clean abort, not a "Turn failed" error); any other exception → the `show_error` path. The `finally` schedules `_mark_done` on EVERY path, so a cancelled turn clears `_busy` and the next submit works — the app never wedges. + +**`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. ## 32. Model Selection And Preload diff --git a/shellpilot/cli/app.py b/shellpilot/cli/app.py index fc6e57c..caa4bf6 100644 --- a/shellpilot/cli/app.py +++ b/shellpilot/cli/app.py @@ -54,8 +54,9 @@ # (branch 4) shows whether it bites. DOCK_MAX_ROWS = 10 -# Idle Ctrl-C hint. NOTE: there is no turn to cancel in this inert shell; branch -# 6 replaces this with real turn cancellation / subprocess kill. +# Idle Ctrl-C hint, shown only when no turn is in flight. Branch 6 (§31.15) wired +# real turn cancellation through on_interrupt; this hint is the idle fallback. +# NOTE: subprocess-kill on cancel is branch 6b. _IDLE_HINT = "(idle — type /exit to quit)" @@ -149,6 +150,7 @@ def build_app( output: Output | None = None, ui: AppUI | None = None, on_submit: Callable[[str], None] | None = None, + on_interrupt: Callable[[], bool] | None = None, ) -> Application[None]: """Build the full-screen app shell. @@ -324,6 +326,12 @@ def _page_down(event: KeyPressEvent) -> None: @kb.add("c-c") def _interrupt(event: KeyPressEvent) -> None: + # Raw mode disables ISIG, so Ctrl-C is a normal key press, not a SIGINT + # (branch 6, §31.15). When a turn is in flight, on_interrupt cancels it and + # returns True (the worker aborts the stream and renders the marker), so we + # show nothing. Otherwise it is idle — fall back to the idle hint. + if on_interrupt is not None and on_interrupt(): + return _ui.show_status(_IDLE_HINT) # NOTE: keyboard PageUp/PageDown drive the pane via the cursor-line model diff --git a/shellpilot/cli/app_main.py b/shellpilot/cli/app_main.py index e47c27f..45827d1 100644 --- a/shellpilot/cli/app_main.py +++ b/shellpilot/cli/app_main.py @@ -70,6 +70,7 @@ def run_app( ctx_pct=ctx_pct, ui=app_ui, on_submit=runner.start, + on_interrupt=runner.request_cancel, ) runner.app = app runner.conversation = runtime diff --git a/shellpilot/cli/app_turn.py b/shellpilot/cli/app_turn.py index 394359c..9c1bca8 100644 --- a/shellpilot/cli/app_turn.py +++ b/shellpilot/cli/app_turn.py @@ -27,9 +27,12 @@ from collections.abc import Callable from typing import TYPE_CHECKING +from shellpilot.llm.client import GenerationCancelled + if TYPE_CHECKING: from prompt_toolkit.application import Application + from shellpilot.cli.app_ui import AppUI from shellpilot.policy.approvals import ApprovalReply, ApprovalRequest from shellpilot.runtime.conversation import ConversationRuntime from shellpilot.runtime.events import RuntimeUI, TurnStats @@ -138,13 +141,21 @@ class TurnRunner: needs the conversation — is broken by deferred attribute assignment. """ - def __init__(self, *, inner_ui: RuntimeUI, schedule: Schedule | None = None) -> None: + def __init__(self, *, inner_ui: AppUI, schedule: Schedule | None = None) -> None: + # The raw content UI (NOT the marshaling ThreadedUI): the worker calls + # show_error / abort_turn on it via schedule. Typed AppUI because the + # cancel path needs abort_turn, which is an app-side AppUI method (like + # show_user_message) and deliberately not on the RuntimeUI protocol. self._inner_ui = inner_ui # busy: only ever touched on the loop thread (see class docstring). self._busy = False - # The worker handle, kept so a test (and, later, branch-6 cancellation) - # can join it. Spawned fresh per turn. + # The worker handle, kept so a test (and branch-6 cancellation) can join + # it. Spawned fresh per turn. self._thread: threading.Thread | None = None + # Branch-6 per-turn cancel event (§31.15): created fresh in start(), + # signaled by request_cancel() on the loop thread, observed by the + # worker's model call. None when no turn is in flight. + self._cancel: threading.Event | None = None # Set after construction, before app.run(), to break the build cycle; # read only during/after run(). self.app: Application[None] | None = None @@ -185,25 +196,61 @@ def start(self, text: str) -> None: Ignores the call when a turn is already in flight — single model, single conversation, single worker; no parallel turns. NOTE: branch 9 adds the one-message queue + Up-arrow recall; until then a submit-while-busy is - dropped. NOTE: branch 6 adds real Ctrl-C turn cancellation; there is no - cancel path here. + dropped. + + A fresh cancel ``Event`` is created here (before the worker spawns, so the + worker sees it) and passed into ``run_turn``; ``request_cancel`` signals + it from the loop thread to abort the turn (branch 6, §31.15). """ if self._busy: return self._busy = True - self._thread = threading.Thread(target=self._run, args=(text,), daemon=True) + cancel = threading.Event() + self._cancel = cancel + # Pass the event as a thread ARG (a captured local), not via self._cancel, + # so the worker never reads the mutable instance attribute. `request_cancel` + # sets the SAME object from the loop thread, so the worker's local sees it — + # but the threading invariant ("the worker only touches its args + schedule; + # self._cancel/_busy are loop-thread-only") then holds literally, without + # relying on the _busy guard to keep self._cancel from being reassigned. + self._thread = threading.Thread(target=self._run, args=(text, cancel), daemon=True) self._thread.start() - def _run(self, text: str) -> None: + def request_cancel(self) -> bool: + """Signal the in-flight turn to abort. Runs on the loop thread (Ctrl-C). + + Returns True when a turn was in flight (the cancel event is now set, and + the worker's next model-call read boundary will raise + ``GenerationCancelled``); False when idle (nothing to cancel — the caller + then shows the idle hint). ``_busy``/``_cancel`` are read on the loop + thread; ``Event.set()`` is the one thread-safe cross-thread signal to the + worker (raw mode disables ISIG, so Ctrl-C is a key press, not a SIGINT). + """ + if self._busy and self._cancel is not None: + self._cancel.set() + return True + return False + + def _run(self, text: str, cancel: threading.Event) -> None: """Worker-thread body: run ONE synchronous turn off the loop thread. - Touches nothing on the loop thread directly — only ``conversation.run_turn`` + ``cancel`` is the per-turn event, received as an ARG (a captured local) so + the worker never reads ``self._cancel`` — keeping the threading invariant + literal. Touches nothing on the loop thread directly — only ``run_turn`` (its UI is the :class:`ThreadedUI`, so every UI call is marshaled) and - ``schedule``. A ``run_turn`` exception (e.g. an approval-needing turn - raises ``NotImplementedError`` until branch 7, or an ``OllamaError`` - surfaces) is rendered as a pane error instead of silently killing the - daemon thread; the ``finally`` clears ``busy`` on the loop thread so a - crashed turn never wedges the app. + ``schedule``. A ``GenerationCancelled`` (the user hit Ctrl-C mid-stream) + is a CLEAN abort, routed to ``abort_turn`` — never the error path. Any + OTHER ``run_turn`` exception (an approval-needing turn raises + ``NotImplementedError`` until branch 7, or an ``OllamaError`` surfaces) is + rendered as a pane error instead of silently killing the daemon thread. + The ``finally`` clears ``busy`` on the loop thread on EVERY path + (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. """ try: conversation = self.conversation @@ -212,7 +259,13 @@ def _run(self, text: str) -> None: # Raise INSIDE the try so it surfaces as a pane error and the # finally still clears busy — never a silent daemon death + wedge. raise RuntimeError("TurnRunner.conversation not set before start()") - conversation.run_turn(text) + conversation.run_turn(text, cancel=cancel) + except GenerationCancelled: + # Clean Ctrl-C abort, NOT a failure: clear the dangling indicator and + # mark the partial aborted. The partial assistant reply was never + # recorded in history (run_turn let the exception propagate before the + # record site), so this is a display-only finalization. + self._schedule(self._inner_ui.abort_turn) except Exception as exc: # noqa: BLE001 - surface ANY worker failure to the pane self._schedule(functools.partial(self._inner_ui.show_error, f"Turn failed: {exc}")) finally: diff --git a/shellpilot/cli/app_ui.py b/shellpilot/cli/app_ui.py index 032356b..628f33c 100644 --- a/shellpilot/cli/app_ui.py +++ b/shellpilot/cli/app_ui.py @@ -246,6 +246,22 @@ def turn_finished(self, stats: TurnStats) -> None: line.append(f" · {_fmt_count(stats.output_tokens)} total", style="sp.dim") self._add_renderable(line) + def abort_turn(self) -> None: + # Branch 6 (§31.15): a turn was cancelled mid-stream (Ctrl-C). App-side + # (NOT a RuntimeUI protocol method) — TurnRunner calls it on the raw AppUI. + # Clear the dangling live indicator (stop the timer/plane) FIRST, then + # append the aborted marker. _add_renderable closes the open response + # first, so any partial streamed text stays visible — finalized — with + # the marker line below it. The partial assistant reply is never recorded + # in history; that discard is the runtime's (run_turn lets + # GenerationCancelled propagate before the record site). This is purely + # display: mark what the user already saw as stopped. + self._indicator = None + # `⏹` is not in the Glyphs set; gate it on the same unicode/ascii signal + # build_app uses, falling back to the ASCII cross glyph. + marker = "⏹" if self._glyphs == UNICODE_GLYPHS else self._glyphs.cross + self._add_renderable(Text(f"{marker} aborted", style="sp.warn")) + def stream_thinking(self, text: str) -> None: # The reasoning count climbs ONLY while the model is thinking, so it # freezes naturally when thinking stops and resumes if it thinks again. No diff --git a/shellpilot/llm/client.py b/shellpilot/llm/client.py index 57d6ee9..0d1f4a0 100644 --- a/shellpilot/llm/client.py +++ b/shellpilot/llm/client.py @@ -2,15 +2,32 @@ from __future__ import annotations +import threading from collections.abc import Callable, Sequence -from typing import Any, Protocol +from typing import TYPE_CHECKING, Any, Protocol from shellpilot.llm.messages import Message, ToolDefinition -from shellpilot.llm.ollama import LocalModel + +if TYPE_CHECKING: + # Only an annotation (list_models return type); a TYPE_CHECKING import keeps + # the protocol module free of a runtime dependency on the concrete Ollama + # client, so ollama.py can import GenerationCancelled from here without an + # import cycle. + from shellpilot.llm.ollama import LocalModel TokenCallback = Callable[[str], None] +class GenerationCancelled(Exception): + """A model turn was aborted mid-stream via the cancel event (branch 6, §31.15). + + Deliberately NOT an ``OllamaError`` subclass: the think-retry in + ``OllamaClient.chat`` and the worker in ``TurnRunner`` must distinguish a + user cancel from a provider failure, so it propagates past both as a clean, + typed signal rather than a generic error. + """ + + class LLMClient(Protocol): """What the conversation runtime needs from a model provider.""" @@ -24,6 +41,7 @@ def chat( options: dict[str, Any] | None = None, on_token: TokenCallback | None = None, on_thinking: TokenCallback | None = None, + cancel: threading.Event | None = None, ) -> Message: ... def health(self) -> bool: ... diff --git a/shellpilot/llm/ollama.py b/shellpilot/llm/ollama.py index 4062658..101523f 100644 --- a/shellpilot/llm/ollama.py +++ b/shellpilot/llm/ollama.py @@ -4,6 +4,7 @@ import ipaddress import json +import threading from collections.abc import Callable, Sequence from dataclasses import dataclass from typing import Any @@ -11,6 +12,7 @@ import httpx +from shellpilot.llm.client import GenerationCancelled from shellpilot.llm.messages import Message, ToolCall, ToolDefinition DEFAULT_BASE_URL = "http://localhost:11434" @@ -199,11 +201,18 @@ def chat( options: dict[str, Any] | None = None, on_token: Callable[[str], None] | None = None, on_thinking: Callable[[str], None] | None = None, + cancel: threading.Event | None = None, ) -> Message: """Stream one chat completion; num_ctx is set explicitly on every request. Configured `options` pass through verbatim, but num_ctx ALWAYS wins: the context budget owns it (design section 10.5). + + ``cancel`` is a branch-6 turn-abort signal (§31.15): when the app SETS it + (from the Ctrl-C keybinding on the loop thread), the next stream-read + boundary raises ``GenerationCancelled``. It is threaded into the + think-retry path too, so a cancel requested during the first (think) + attempt still aborts the retried call. """ payload: dict[str, Any] = { "model": model, @@ -216,14 +225,14 @@ def chat( if self._reasoning and model not in self._no_think: payload["think"] = True try: - return self._stream_chat(payload, on_token, on_thinking) + return self._stream_chat(payload, on_token, on_thinking, cancel) except OllamaResponseError as exc: # Reasoning mode unavailable for this model: cache and retry once without # think (design section 24.6). Other models are not affected. if self._reasoning and model not in self._no_think and "think" in str(exc).lower(): self._no_think.add(model) payload.pop("think", None) - return self._stream_chat(payload, on_token, on_thinking) + return self._stream_chat(payload, on_token, on_thinking, cancel) raise def _stream_chat( @@ -231,6 +240,7 @@ def _stream_chat( payload: dict[str, Any], on_token: Callable[[str], None] | None, on_thinking: Callable[[str], None] | None = None, + cancel: threading.Event | None = None, ) -> Message: content_parts: list[str] = [] thinking_parts: list[str] = [] @@ -243,6 +253,15 @@ def _stream_chat( body = response.read().decode("utf-8", errors="replace") raise OllamaResponseError(f"Ollama API error {response.status_code}: {body}") for line in response.iter_lines(): + # Branch-6 cancel check at the read boundary (§31.15). Raising + # INSIDE the `with self._client.stream(...)` block closes the + # response IN-THREAD via the context manager — the safe close. + # The app only ever SETS the event (cross-thread, thread-safe); + # response.close() is never called cross-thread. GenerationCancelled + # is neither httpx.TransportError nor OllamaResponseError, so it + # propagates cleanly out through chat() to the runtime. + if cancel is not None and cancel.is_set(): + raise GenerationCancelled if not line.strip(): continue try: diff --git a/shellpilot/runtime/conversation.py b/shellpilot/runtime/conversation.py index 459744b..1aa107c 100644 --- a/shellpilot/runtime/conversation.py +++ b/shellpilot/runtime/conversation.py @@ -4,6 +4,7 @@ import dataclasses import json +import threading import time from collections.abc import Sequence from dataclasses import dataclass @@ -156,6 +157,11 @@ def __init__( self._last_user_text = "" self._last_failure_signature: str | None = None self._turn_output_tokens: int = 0 + # Branch-6 per-turn cancel handle (§31.15): None unless run_turn is given + # a cancel event for this turn. Read by _tool_loop and passed to each + # model call. Reassigned at every run_turn start so a stale event from a + # prior turn can never leak across turns. + self._cancel: threading.Event | None = None self.snapshots = SnapshotStore() self.recent_diffs: list[str] = [] self.plan_manager = PlanManager(workspace, settings.runtime.security_profile) @@ -477,8 +483,25 @@ def over() -> bool: changed += 1 return changed - def run_turn(self, text: str, *, images: Sequence[ImageRef] = ()) -> str: - """One user turn: budget-check, compact, call the model, stream, record.""" + def run_turn( + self, + text: str, + *, + images: Sequence[ImageRef] = (), + cancel: threading.Event | None = None, + ) -> str: + """One user turn: budget-check, compact, call the model, stream, record. + + ``cancel`` (branch 6, §31.15) is the turn-abort signal: when set + mid-stream the model call raises ``GenerationCancelled``, which this + method lets PROPAGATE. The user message recorded below stays (the turn + happened); the partial assistant reply is never reached at the record + site, so a cancelled turn leaves NO partial reply in history. ``cancel`` + defaults to None, so every existing caller is byte-identical. + """ + # Assigned (not merely defaulted) at the top so a stale event from a + # prior turn can never leak into this one; the tool loop reads it. + self._cancel = cancel # Belt-and-braces: a stale stage left by an aborted prior turn must not # attach to this unrelated turn's first message. self._staged_tool_images.clear() @@ -614,6 +637,7 @@ def _tool_loop(self) -> Message: options=self._settings.model.options, on_token=self._ui.stream_token, on_thinking=self._ui.stream_thinking, + cancel=self._cancel, ) finally: self._ui.end_response() diff --git a/tests/fakes/fake_llm.py b/tests/fakes/fake_llm.py index c505a26..308a10a 100644 --- a/tests/fakes/fake_llm.py +++ b/tests/fakes/fake_llm.py @@ -7,11 +7,12 @@ from __future__ import annotations +import threading from collections.abc import Sequence from dataclasses import dataclass, field from typing import Any -from shellpilot.llm.client import TokenCallback +from shellpilot.llm.client import GenerationCancelled, TokenCallback from shellpilot.llm.messages import Message, ToolCall, ToolDefinition, assistant from shellpilot.llm.ollama import LocalModel @@ -61,7 +62,13 @@ def chat( options: dict[str, Any] | None = None, on_token: TokenCallback | None = None, on_thinking: TokenCallback | None = None, + cancel: threading.Event | None = None, ) -> Message: + # Branch 6: honor a set cancel event the same way OllamaClient does — + # abort before consuming the script, so a cancelled turn never produces a + # (partial) reply. A None/unset event leaves every existing test unchanged. + if cancel is not None and cancel.is_set(): + raise GenerationCancelled self.calls.append( RecordedCall( model=model, diff --git a/tests/test_app.py b/tests/test_app.py index 7f83473..38fdc14 100644 --- a/tests/test_app.py +++ b/tests/test_app.py @@ -261,3 +261,58 @@ def test_app_alt_enter_inserts_newline_then_submits(tmp_path: Path) -> None: texts = [r.plain for r in ui._renderables if isinstance(r, RichText)] assert any("a\nb" in t for t in texts) assert not any("/exit" in t for t in texts) + + +# --- Branch-6 Ctrl-C on_interrupt wiring (§31.15) ----------------------------- + + +def _run_with_interrupt(tmp_path: Path, *, on_interrupt: object) -> AppUI: + """Build a headless app with on_interrupt wired, press Ctrl-C, then /exit.""" + ui = AppUI(glyphs=UNICODE_GLYPHS, width_fn=lambda: 80) + with create_pipe_input() as inp: + app = build_app( + workspace=tmp_path, + model="gemma4:e4b", + profile="balanced", + glyphs=UNICODE_GLYPHS, + commands=command_words(), + input=inp, # type: ignore[arg-type] + output=DummyOutput(), + ui=ui, + on_interrupt=on_interrupt, # type: ignore[arg-type] + ) + inp.send_text("\x03") # Ctrl-C (ETX) — a key press, not a SIGINT + inp.send_text("/exit\n") + app.run() + return ui + + +def test_c_c_calls_on_interrupt_and_skips_idle_hint_when_true(tmp_path: Path) -> None: + calls: list[int] = [] + + def on_interrupt() -> bool: + calls.append(1) + return True # a turn was cancelled + + ui = _run_with_interrupt(tmp_path, on_interrupt=on_interrupt) + assert calls == [1] + # A cancelled turn shows its own marker (from abort_turn), NOT the idle hint. + assert "idle" not in ui._render_ansi() + + +def test_c_c_shows_idle_hint_when_on_interrupt_returns_false(tmp_path: Path) -> None: + calls: list[int] = [] + + def on_interrupt() -> bool: + calls.append(1) + return False # nothing in flight + + ui = _run_with_interrupt(tmp_path, on_interrupt=on_interrupt) + assert calls == [1] + assert "idle" in ui._render_ansi() + + +def test_c_c_shows_idle_hint_when_on_interrupt_none(tmp_path: Path) -> None: + # Back-compat: with no on_interrupt (the inert shell), Ctrl-C shows the hint. + ui = _run_with_interrupt(tmp_path, on_interrupt=None) + assert "idle" in ui._render_ansi() diff --git a/tests/test_app_turn.py b/tests/test_app_turn.py index 7543b3e..13b6c52 100644 --- a/tests/test_app_turn.py +++ b/tests/test_app_turn.py @@ -30,6 +30,7 @@ from shellpilot.cli.slash import command_words from shellpilot.cli.theme import UNICODE_GLYPHS from shellpilot.config.model import Settings +from shellpilot.llm.client import GenerationCancelled from shellpilot.llm.messages import Message from shellpilot.memory.agents_md import BehaviorInstructions from shellpilot.runtime.conversation import ConversationRuntime @@ -74,6 +75,9 @@ def end_response(self) -> None: def turn_finished(self, stats: TurnStats) -> None: self._record("turn_finished", stats) + def abort_turn(self) -> None: + self._record("abort_turn") + def show_status(self, text: str) -> None: self._record("show_status", text) @@ -328,6 +332,92 @@ def test_worker_exception_is_surfaced_and_clears_busy(tmp_path: Path) -> None: assert runner._busy is False # the finally ran +# --- Branch-6 turn cancellation (§31.15) -------------------------------------- + + +class _CancellableLLM: + """chat() blocks until released, then raises GenerationCancelled iff cancel is set. + + The same fake drives both paths: with the cancel event set it aborts (mirrors + OllamaClient hitting the read boundary), and with it unset it returns a plain + answer (a normal completion). Non-chat protocol methods delegate to a FakeLLM. + """ + + def __init__(self) -> None: + self.entered = threading.Event() + self.release = threading.Event() + self._inner = FakeLLM(script=[]) + + def chat(self, *args: Any, cancel: Any = None, **kwargs: Any) -> Message: + self.entered.set() + assert self.release.wait(5.0) + if cancel is not None and cancel.is_set(): + raise GenerationCancelled + return answer("done") + + def health(self) -> bool: + return self._inner.health() + + def list_models(self) -> Any: + return self._inner.list_models() + + def model_context_length(self, model: str) -> int | None: + return self._inner.model_context_length(model) + + def model_capabilities(self, model: str) -> tuple[str, ...]: + return self._inner.model_capabilities(model) + + def preload(self, model: str, *, keep_alive: str = "5m") -> None: + self._inner.preload(model, keep_alive=keep_alive) + + +def test_request_cancel_aborts_turn_cleanly_and_reruns(tmp_path: Path) -> None: + 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) + llm = _CancellableLLM() + runtime = _make_runtime(llm, threaded, tmp_path) + runner.conversation = runtime + + # Idle: there is no turn, so request_cancel reports nothing to cancel. + assert runner.request_cancel() is False + + runner.start("hi") + assert llm.entered.wait(5.0) # the worker really entered the model call → in flight + assert runner._busy is True + # Busy: request_cancel sets the event and reports True. + assert runner.request_cancel() is True + llm.release.set() # let chat wake; it sees the set event and raises + + assert runner._thread is not None + runner._thread.join(5.0) + assert not runner._thread.is_alive() # cancelled 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 # _mark_done cleared busy on the cancel path + assert runner.request_cancel() is False # idle again + assert "aborted" in app_ui._render_ansi() # the marker reached the real render + + # busy reset → a fresh turn now runs to a normal completion after the cancel. + llm.entered.clear() + llm.release.clear() + runner.start("again") + assert llm.entered.wait(5.0) + llm.release.set() # no cancel this time → chat returns a normal answer + 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_app_ui.py b/tests/test_app_ui.py index 89380f6..30de687 100644 --- a/tests/test_app_ui.py +++ b/tests/test_app_ui.py @@ -9,7 +9,7 @@ import pytest from shellpilot.cli.app_ui import AppUI -from shellpilot.cli.theme import UNICODE_GLYPHS +from shellpilot.cli.theme import ASCII_GLYPHS, UNICODE_GLYPHS from shellpilot.memory.redaction import REDACTED from shellpilot.policy.approvals import ApprovalRequest from shellpilot.policy.risk import RiskLevel @@ -482,6 +482,38 @@ def test_new_turn_after_error_starts_fresh_indicator() -> None: assert ui._indicator.reasoning_chars == 0 +def test_abort_turn_clears_indicator_and_marks_partial() -> None: + # Branch 6 (§31.15): a cancelled turn clears the dangling live indicator and + # marks the partial aborted, leaving the streamed-so-far text visible. + ui = AppUI(glyphs=GLYPHS, width_fn=lambda: 80, time_fn=lambda: 0.0) + ui.begin_response() # turn starts → indicator active + ui.stream_token("partial answer so far") # streamed text, response still open + assert ui._indicator is not None + + ui.abort_turn() + + # The live indicator is gone (no dangling timer/plane after a cancel). + assert ui._indicator is None + out = plain(ui) + # The partial streamed text was finalized and is still visible ... + assert "partial answer so far" in out + # ... with the aborted marker below it (unicode ⏹, style sp.warn). + assert "⏹ aborted" in out + # A later render does not resurrect the live indicator (no flight phrase). + assert "taxiing" not in plain(ui) + + +def test_abort_turn_ascii_fallback_marker() -> None: + # ⏹ is not in the Glyphs set; in ASCII mode it degrades to the cross glyph. + ui = AppUI(glyphs=ASCII_GLYPHS, width_fn=lambda: 80) + ui.begin_response() + ui.abort_turn() + out = plain(ui) + assert ui._indicator is None + assert "⏹" not in out + assert f"{ASCII_GLYPHS.cross} aborted" in out + + def test_stream_thinking_climbs_reasoning_estimate() -> None: ui = AppUI(glyphs=GLYPHS, width_fn=lambda: 80, time_fn=lambda: 0.0) ui.begin_response() diff --git a/tests/test_conversation.py b/tests/test_conversation.py index 9e6b07b..4e5f36d 100644 --- a/tests/test_conversation.py +++ b/tests/test_conversation.py @@ -1,8 +1,11 @@ """Tests for the unified conversation runtime with the fake model.""" import json +import threading from pathlib import Path +import pytest + from shellpilot.config.model import ( ContextSettings, PrivacySettings, @@ -10,6 +13,7 @@ SkillSettings, ToolSettings, ) +from shellpilot.llm.client import GenerationCancelled from shellpilot.llm.messages import Message from shellpilot.memory.agents_md import BehaviorInstructions from shellpilot.persistence.audit_store import AuditLogger @@ -61,6 +65,37 @@ def test_history_accumulates_across_turns(tmp_path: Path) -> None: assert "second" in contents +def test_cancelled_turn_discards_partial_reply(tmp_path: Path) -> None: + """History integrity (§31.15): a cancelled turn keeps the user message but no reply.""" + # FakeLLM raises GenerationCancelled when the cancel event is set, mirroring + # OllamaClient aborting mid-stream. The scripted answer must NEVER be recorded. + fake = FakeLLM(script=[answer("partial reply that must never be recorded")]) + runtime = make_runtime(fake, FakeUI(), tmp_path) + cancel = threading.Event() + cancel.set() + + with pytest.raises(GenerationCancelled): + runtime.run_turn("do a thing", cancel=cancel) + + # The user message was recorded (the turn happened) ... + assert [m.role for m in runtime._history] == ["user"] + assert runtime._history[0].content == "do a thing" + # ... but NO assistant reply landed in history — the partial is gone. + assert all(m.role != "assistant" for m in runtime._history) + + +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")]) + runtime = make_runtime(fake, FakeUI(), tmp_path) + cancel = threading.Event() # never set + + runtime.run_turn("question", cancel=cancel) + + assert [m.role for m in runtime._history] == ["user", "assistant"] + assert runtime._history[1].content == "the answer" + + def test_oversized_user_message_is_refused(tmp_path: Path) -> None: fake = FakeLLM(script=[]) ui = FakeUI() diff --git a/tests/test_ollama_chat.py b/tests/test_ollama_chat.py index c487f27..3844dd8 100644 --- a/tests/test_ollama_chat.py +++ b/tests/test_ollama_chat.py @@ -1,10 +1,14 @@ """Tests for OllamaClient.chat streaming (httpx.MockTransport, no live Ollama).""" import json +import threading +from collections.abc import Iterator from typing import Any import httpx +import pytest +from shellpilot.llm.client import GenerationCancelled from shellpilot.llm.messages import ToolDefinition, user from shellpilot.llm.ollama import OllamaClient @@ -229,6 +233,82 @@ def handler(request: httpx.Request) -> httpx.Response: assert attempts[0].get("think") is True +# --------------------------------------------------------------------------- +# Branch-6 mid-stream cancellation (§31.15) +# --------------------------------------------------------------------------- + + +def test_stream_chat_cancel_stops_reading_mid_stream() -> None: + """A set cancel event aborts the stream read; the rest is never consumed.""" + cancel = threading.Event() + pulled = {"count": 0} + total = 6 + + def gen() -> Iterator[bytes]: + for i in range(total): + pulled["count"] += 1 + # The user hits Ctrl-C while the 2nd chunk streams. + if i == 1: + cancel.set() + chunk = {"message": {"role": "assistant", "content": f"c{i}"}, "done": i == total - 1} + yield (json.dumps(chunk) + "\n").encode() + + def handler(request: httpx.Request) -> httpx.Response: + return httpx.Response(200, content=gen()) + + client = OllamaClient(transport=httpx.MockTransport(handler)) + with pytest.raises(GenerationCancelled): + client.chat("gemma4:e4b", [user("hi")], num_ctx=4096, cancel=cancel) + # It raised at the read boundary after the cancelled chunk — it did NOT drain + # the whole stream (no `done` chunk was reached). + assert pulled["count"] < total + + +def test_stream_chat_cancel_not_set_completes_normally() -> None: + """An unset (or absent) cancel event leaves streaming unaffected.""" + cancel = threading.Event() # never set + + def handler(request: httpx.Request) -> httpx.Response: + return httpx.Response( + 200, + content=stream_body( + {"message": {"role": "assistant", "content": "Hel"}, "done": False}, + {"message": {"role": "assistant", "content": "lo"}, "done": True}, + ), + ) + + client = OllamaClient(transport=httpx.MockTransport(handler)) + reply = client.chat("gemma4:e4b", [user("hi")], num_ctx=4096, cancel=cancel) + assert reply.content == "Hello" + + +def test_chat_threads_cancel_through_think_retry() -> None: + """A cancel set during the think attempt still aborts the retried (no-think) call.""" + cancel = threading.Event() + cancel.set() # cancel already requested before the call + attempts: list[dict[str, Any]] = [] + + def handler(request: httpx.Request) -> httpx.Response: + payload = json.loads(request.content) + attempts.append(payload) + if payload.get("think"): + # The think attempt fails the status check BEFORE the stream loop, so + # it raises OllamaResponseError (not the cancel) → triggers the retry. + return httpx.Response(400, json={"error": "model does not support thinking"}) + # The retry streams; its read boundary sees the set cancel and raises. + return httpx.Response( + 200, + content=stream_body({"message": {"role": "assistant", "content": "ok"}, "done": True}), + ) + + client = OllamaClient(reasoning=True, transport=httpx.MockTransport(handler)) + with pytest.raises(GenerationCancelled): + client.chat("gemma4:e4b", [user("hi")], num_ctx=2048, cancel=cancel) + # The retry happened (think → no-think) AND it observed the cancel. + assert attempts[0].get("think") is True + assert "think" not in attempts[1] + + def test_model_context_length_reads_model_info() -> None: def handler(request: httpx.Request) -> httpx.Response: return httpx.Response(