From d37250b0b419560f0342e05840db96ace8c60a93 Mon Sep 17 00:00:00 2001 From: Lavindeep Dhillon Date: Sat, 27 Jun 2026 20:11:24 -0400 Subject: [PATCH] feat(ui): run the turn on a worker thread, marshal the UI to the loop MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit In the full-screen app the prompt_toolkit event loop is the main thread, so driving a model turn synchronously there would freeze the UI for the whole turn — no repaint, no scroll. Run the one synchronous turn on a worker thread and marshal every UI callback back onto the loop thread, so AppUI state and the pane repaint are only ever touched on the loop thread. - ThreadedUI: a RuntimeUI wrapping AppUI + an injected schedule; each fire-and-forget content call is enqueued via functools.partial (never a loop-variable closure). Blocking approvals return a value so they can't be fire-and-forget — they delegate to the inner, which still raises until the focus-swap lands, surfaced by the worker as a pane error. - TurnRunner: owns one worker per turn and a busy flag touched only on the loop thread (set in start, cleared in a scheduled _mark_done) — no lock. The worker only calls run_turn + schedule; any failure (incl. a build- order violation) surfaces as a pane error and the finally clears busy, so a crashed turn never wedges the app. schedule reads app.loop lazily and call_soon_threadsafe's a callback that runs the call then invalidate(). - build_app gains on_submit; the opt-in run_app (SHELLPILOT_UI=app) wires it all and app.run()s it. The default REPL is byte-identical — app mode is off unless the env opt-in is set. The threading model was independently reviewed against the prompt_toolkit source (call_soon_threadsafe + lazy app.loop + invalidate-from-callback all confirmed safe). Known dev-entry gaps (no session_end audit; in-flight turn at /exit drops its final busy-clear in the dying app) are documented for the promotion-to-default review. DESIGN.md section 31.13 covers the worker turn, marshaling, and the entry. --- docs/DESIGN.md | 10 + shellpilot/cli/app.py | 14 +- shellpilot/cli/app_main.py | 77 ++++++++ shellpilot/cli/app_turn.py | 224 +++++++++++++++++++++ shellpilot/cli/terminal.py | 48 ++++- tests/test_app_turn.py | 385 +++++++++++++++++++++++++++++++++++++ 6 files changed, 752 insertions(+), 6 deletions(-) create mode 100644 shellpilot/cli/app_main.py create mode 100644 shellpilot/cli/app_turn.py create mode 100644 tests/test_app_turn.py diff --git a/docs/DESIGN.md b/docs/DESIGN.md index d95b667..00bd26e 100644 --- a/docs/DESIGN.md +++ b/docs/DESIGN.md @@ -2825,6 +2825,16 @@ A one-line status bar is pinned at the input as `prompt_toolkit`'s `bottom_toolb The interactive REPL is migrating from a per-turn `prompt_toolkit` `prompt` to a persistent full-screen `Application` on the alternate screen (`cli/app.py`, `build_app(...)`), so the status bar and input dock stay pinned and never scroll away mid-turn. The layout is a scrolling chat pane on top, a custom rounded multi-line **input dock**, and the persistent status bar (§31.11) pinned at the bottom. The dock border is **hand-drawn** (`horizontal_border(width, box, top=...)` — rounded `╭─╮`/`╰─╯`, ASCII `+`/`-` fallback) rather than `prompt_toolkit`'s `Frame`, which only draws square corners and reserves completion-menu height inside the box (filling the terminal height instead of hugging the input); this is the one deliberate exception to the §31.9 "borders from rich primitives" rule, since the dock lives in a `prompt_toolkit` layout. The dock border width, the pane wrap width, and the terminal width are one shared value read at render time, so a resize re-derives the border with nothing cached. Enter submits (a literal newline is Alt+Enter); the dock grows to a 10-row cap then scrolls; PageUp/PageDown scroll the (unfocused) pane while the wheel routes to the window under the cursor. The pane now renders Rich content via `AppUI` (`cli/app_ui.py`), which implements the `RuntimeUI` protocol. `AppUI` holds a `list[RenderableType]` as the transcript source of truth, with an optional in-progress open-response accumulator (accumulated token text rendered as `rich.Markdown`). `AppUI._render_ansi()` renders the renderables to a single ANSI string at the shared terminal width via a `Console(force_terminal=True, width=W)`; the result is cached by width so a resize re-derives it automatically (cache miss) when the pane's `FormattedTextControl(lambda: ANSI(ui._render_ansi()))` renders. `wrap_lines=False` on the pane window lets Rich own all wrapping. Secret redaction (`redact_structure`) and display-integrity path resolution (`workspace_display`) are applied in `show_tool_call` before the summary reaches the pane. Approval methods raise `NotImplementedError` until branch 7 wires the focus-swap; no silent default is acceptable there. `AppUI` never calls `get_app()` or `invalidate()` — those are branch-4 concerns — so it is fully testable without a running app. The non-TTY `PlainInput` path (§31.2) is unchanged. +### 31.13 Worker-thread turn + loop-thread marshaling (v2, in progress) + +`ConversationRuntime.run_turn` drives the UI synchronously on the calling thread, calling `self._ui.` for every event. In the full-screen app that thread is the prompt_toolkit event loop, so a long model turn would freeze the whole UI. Branch 4 (`cli/app_turn.py`) runs the one synchronous turn on a **worker thread** and marshals every UI callback back onto the loop thread, so `AppUI` state and the pane repaint are only ever touched on the loop thread. + +`ThreadedUI` is a `RuntimeUI` wrapping the real `AppUI` (`inner`) and an injected `schedule: Callable[[Callable[[], None]], None]`. Every fire-and-forget content method (`stream_token`, `stream_thinking`, `begin_response`/`end_response`, `turn_finished`, `show_status`, `show_error`, `show_tool_call`/`show_tool_result`, `show_command_output`, `show_plan_progress`) enqueues the inner call via `schedule(functools.partial(inner.method, *args))` — `functools.partial`, never a closure over a loop variable, so args bind at call time and cannot be rebound before the queued call runs. The blocking approval methods return a value, so they cannot be fire-and-forget: `ask_approval`/`ask_plan_approval` delegate straight to the inner UI and run on the worker thread; the inner `AppUI` still raises `NotImplementedError` (the focus-swap `Future` handshake is branch 7), and the worker's `try/except` surfaces that as a pane error — no approval ever silently defaults. + +`TurnRunner` owns the single worker thread for one turn and a `busy` flag. The core safety property: `busy` is SET in `start` (the dock-submit handler, which runs on the loop thread) and CLEARED in `_mark_done` (scheduled back onto the loop thread from the worker's `finally`), so it is only ever read or written on the loop thread — no lock needed. `start` ignores a submit while `busy` (single model, single conversation, single worker; the one-message queue is branch 9). The worker body `_run` touches nothing on the loop thread directly — it only calls `conversation.run_turn` (whose UI is the marshaling `ThreadedUI`, so every UI call is already marshaled) and routes its own error/completion through `schedule`; a `run_turn` exception is rendered as a pane error instead of silently killing the daemon thread, and the `finally` clears `busy` so a crashed turn never wedges the app (Ctrl-C turn cancellation is branch 6). `schedule` is injected so CI drives it synchronously; the real app uses `TurnRunner.schedule`, which reads `app.loop` lazily and schedules (via `loop.call_soon_threadsafe`) a callback that runs `fn()` then `app.invalidate()` (fails closed when the app is not running). prompt_toolkit coalesces the per-call `invalidate()` into one render tick, so per-token marshal+invalidate is correct without throttling. + +The opt-in runnable entry is `run_app` (`cli/app_main.py`), reached ONLY via `SHELLPILOT_UI=app` from `run_interactive` so a default session (TTY and non-TTY) stays byte-identical. The construction cycle — `schedule` needs `app`, `app` needs `runner.start`, `runner` needs the conversation, the conversation needs the `ThreadedUI`/`schedule` — is broken by build order and deferred attribute assignment: `AppUI` is built first (its `width_fn` reads `get_app()` lazily), then `TurnRunner` (its `schedule` reads `runner.app` lazily), then `ThreadedUI` over them; the conversation is constructed **with that `ThreadedUI` as its UI** (so its plan tools capture the marshaling bound methods — not repointed after the fact); then `build_app(on_submit=runner.start, ui=app_ui, …)`; then `runner.app`/`runner.conversation` are set after the app exists and before `app.run()`, by which time they are only read once a turn runs. The UI choice is made in `run_interactive` where the conversation is built (a `ThreadedUI` for app mode, the same `TerminalUI` as before otherwise); the default branch is unchanged. As an opt-in dev/live-test entry it has two known gaps to close when the app is promoted to the shipping default: `run_app` does not write the `session_end` audit event the default REPL emits, and a turn still in flight when `/exit` clears `app.loop` drops its final `busy`-clear in the already-dead app (harmless while the process is exiting). + ## 32. Model Selection And Preload ### 32.1 Boot Model Picker diff --git a/shellpilot/cli/app.py b/shellpilot/cli/app.py index 233db50..0e5653e 100644 --- a/shellpilot/cli/app.py +++ b/shellpilot/cli/app.py @@ -139,12 +139,15 @@ def build_app( input: Input | None = None, output: Output | None = None, ui: AppUI | None = None, + on_submit: Callable[[str], None] | None = None, ) -> Application[None]: - """Build the inert full-screen app shell. + """Build the full-screen app shell. ``input``/``output`` default to the real terminal; tests inject a pipe input - and ``DummyOutput`` to drive the shell headlessly. The shell is standalone — - branch 4 wires it into the REPL; the non-TTY ``PlainInput`` path is untouched. + and ``DummyOutput`` to drive the shell headlessly. ``on_submit`` receives the + dock text on submit (branch 4 passes ``TurnRunner.start``); when it is None + the shell falls back to the inert branch-3 echo so the standalone shell and + its headless tests stay working. The non-TTY ``PlainInput`` path is untouched. """ # One unicode/ascii decision drives the border, the branch glyph, and the # bar — recovered from the resolved glyph set the caller already probed @@ -256,7 +259,10 @@ def _submit(event: KeyPressEvent) -> None: event.app.exit() return if text.strip(): - _ui.show_status(text) + if on_submit is not None: + on_submit(text) + else: + _ui.show_status(text) dock_buffer.reset() @kb.add("escape", "enter", filter=dock_focused) diff --git a/shellpilot/cli/app_main.py b/shellpilot/cli/app_main.py new file mode 100644 index 0000000..e47c27f --- /dev/null +++ b/shellpilot/cli/app_main.py @@ -0,0 +1,77 @@ +"""Opt-in runnable entry for the full-screen app (design section 31.13). + +This is the LIVE glue that constructs the full-screen ``Application``, drives a +:class:`~shellpilot.runtime.conversation.ConversationRuntime` through a +worker-thread :class:`~shellpilot.cli.app_turn.TurnRunner`, and ``app.run()``s +it. It is reached ONLY via an explicit opt-in (``SHELLPILOT_UI=app``) from +``run_interactive`` so the default REPL stays byte-identical; nothing here runs +on a default ``shellpilot`` session. + +The construction cycle is broken by deferred attribute assignment — see the +ordering note in :func:`run_app`. +""" + +from __future__ import annotations + +from collections.abc import Sequence +from pathlib import Path +from typing import TYPE_CHECKING + +from shellpilot.cli.app import build_app +from shellpilot.cli.app_turn import TurnRunner + +if TYPE_CHECKING: + from shellpilot.cli.app_ui import AppUI + from shellpilot.cli.theme import Glyphs + from shellpilot.runtime.conversation import ConversationRuntime + + +def run_app( + runtime: ConversationRuntime, + runner: TurnRunner, + app_ui: AppUI, + *, + workspace: Path, + model: str, + profile: str, + glyphs: Glyphs, + commands: Sequence[str], + is_cloud: bool = False, + ctx_pct: int = 0, +) -> int: + """Build the full-screen app around an already-wired conversation and run it. + + The caller has already chosen the conversation's UI as the marshaling + :class:`ThreadedUI` whose inner is ``app_ui`` and whose schedule is + ``runner.schedule`` (so the runtime's plan tools captured the marshaling + bound methods at construction). This function only completes the wiring: + + 1. ``app_ui`` was built first (its ``width_fn`` reads ``get_app()`` lazily, + so it needs no running app yet); ``runner`` was built next (its + ``schedule`` reads ``runner.app`` lazily); the conversation was built + with the ``ThreadedUI`` over ``app_ui``. + 2. Build the ``Application`` from ``runner.start`` (the dock-submit handler) + and ``app_ui`` (the pane source of truth). + 3. Set ``runner.app`` and ``runner.conversation`` AFTER the app exists and + BEFORE ``app.run()`` — both are read only once a turn runs, by which time + ``app.run()`` has set ``app.loop``. + + NOTE: this opt-in entry does not write the ``session_end`` audit event that + the default REPL writes after its loop; it is a dev/live-test entry, not the + shipping path. Returns 0 so ``run_interactive`` can ``return run_app(...)``. + """ + app = build_app( + workspace=workspace, + model=model, + profile=profile, + glyphs=glyphs, + commands=commands, + is_cloud=is_cloud, + ctx_pct=ctx_pct, + ui=app_ui, + on_submit=runner.start, + ) + runner.app = app + runner.conversation = runtime + app.run() + return 0 diff --git a/shellpilot/cli/app_turn.py b/shellpilot/cli/app_turn.py new file mode 100644 index 0000000..5db6073 --- /dev/null +++ b/shellpilot/cli/app_turn.py @@ -0,0 +1,224 @@ +"""Worker-thread turn execution and loop-thread UI marshaling (design section 31.13). + +Branch 4 of the UI v2 rework. Today the runtime drives the UI synchronously on +the calling thread inside :meth:`ConversationRuntime.run_turn`; in the +full-screen app that thread is the prompt_toolkit event loop, so a long model +turn would freeze the whole UI (no repaint, no scroll, no Ctrl-C). This module +runs the one synchronous turn on a **worker thread** and marshals every UI +callback back onto the loop thread, so ``AppUI`` state and the pane repaint are +only ever touched on the loop thread. + +Two pieces, both built to be driven synchronously in CI (the ``schedule`` +callable is injected): + +* :class:`ThreadedUI` — a ``RuntimeUI`` that wraps the real ``AppUI`` and + enqueues every fire-and-forget content call onto the loop thread instead of + running it inline. The blocking approval methods cannot be fire-and-forget + (they return a value), so they delegate straight to the inner UI and run on + the worker thread. +* :class:`TurnRunner` — owns the single worker thread for one turn and a + ``busy`` flag that is only ever touched on the loop thread (no lock needed). +""" + +from __future__ import annotations + +import functools +import threading +from collections.abc import Callable +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from prompt_toolkit.application import Application + + from shellpilot.policy.approvals import ApprovalReply, ApprovalRequest + from shellpilot.runtime.conversation import ConversationRuntime + from shellpilot.runtime.events import RuntimeUI, TurnStats + from shellpilot.runtime.planner import TaskPlan + +# A zero-arg callback scheduled to run on the loop thread. +Scheduled = Callable[[], None] +# Marshals a loop-thread callback from a worker thread. +Schedule = Callable[[Scheduled], None] + + +class ThreadedUI: + """A ``RuntimeUI`` that marshals every fire-and-forget call onto the loop thread. + + Wraps the real ``AppUI`` (``inner``) and a ``schedule`` callable that + enqueues a zero-arg callback to run on the prompt_toolkit event-loop thread. + The runtime calls these methods from the **worker** thread (see + :class:`TurnRunner`); this wrapper guarantees ``AppUI`` state and the pane + repaint are only ever touched on the loop thread. + + Args are captured with :func:`functools.partial` (never a closure over a + loop variable) so a later call cannot rebind them before the queued call + runs — ``stream_token("a")`` then ``stream_token("b")`` drain as ``"a"``, + ``"b"``, not ``"b"``, ``"b"``. + """ + + def __init__(self, *, inner: RuntimeUI, schedule: Schedule) -> None: + self._inner = inner + self._schedule = schedule + + # ------------------------------------------------------------------ + # Fire-and-forget content methods — enqueued, never run inline. + # ------------------------------------------------------------------ + + def stream_token(self, token: str) -> None: + self._schedule(functools.partial(self._inner.stream_token, token)) + + # NOTE: stream_thinking is marshaled even though AppUI's is a no-op today; + # branch 5 gives it a real reasoning-readout consumer. + def stream_thinking(self, text: str) -> None: + self._schedule(functools.partial(self._inner.stream_thinking, text)) + + # NOTE: begin_response is marshaled even though AppUI's is a no-op today; the + # waiting/thinking indicator is wired in branch 5. No args → the bound method + # is already a zero-arg Scheduled, so no partial is needed. + def begin_response(self) -> None: + self._schedule(self._inner.begin_response) + + def end_response(self) -> None: + self._schedule(self._inner.end_response) + + def turn_finished(self, stats: TurnStats) -> None: + self._schedule(functools.partial(self._inner.turn_finished, stats)) + + def show_status(self, text: str) -> None: + self._schedule(functools.partial(self._inner.show_status, text)) + + def show_error(self, text: str) -> None: + self._schedule(functools.partial(self._inner.show_error, text)) + + def show_tool_call(self, name: str, arguments: dict[str, object]) -> None: + self._schedule(functools.partial(self._inner.show_tool_call, name, arguments)) + + def show_tool_result(self, name: str, success: bool, summary: str) -> None: + self._schedule(functools.partial(self._inner.show_tool_result, name, success, summary)) + + def show_command_output(self, line: str) -> None: + self._schedule(functools.partial(self._inner.show_command_output, line)) + + def show_plan_progress(self, plan: TaskPlan) -> None: + self._schedule(functools.partial(self._inner.show_plan_progress, plan)) + + # ------------------------------------------------------------------ + # Blocking methods — return a value, so they CANNOT be fire-and-forget. + # They delegate straight to the inner UI and run on the worker thread. + # ------------------------------------------------------------------ + + # NOTE: branch 7 replaces these with a thread-safe Future focus-swap + # handshake (marshal the prompt to the loop thread, block the worker on a + # Future for the reply). Until then the inner AppUI raises NotImplementedError + # here — no approval may silently default — and the worker's try/except + # surfaces that as a pane error (see TurnRunner._run). + def ask_approval(self, request: ApprovalRequest) -> ApprovalReply: + return self._inner.ask_approval(request) + + def ask_plan_approval(self, plan: TaskPlan, path: str) -> tuple[str, str]: + return self._inner.ask_plan_approval(plan, path) + + +class TurnRunner: + """Owns the single worker thread for one turn and the loop-thread busy flag. + + **Threading invariant (the core safety property of branch 4):** ``_busy`` is + SET in :meth:`start` (called on the loop thread — it is the dock-submit + handler) and CLEARED in :meth:`_mark_done` (scheduled back onto the loop + thread from the worker's ``finally`` block). It is therefore only ever read + or written on the loop thread, so no lock is needed. The worker thread + (:meth:`_run`) never touches ``_busy`` or the inner UI directly — it only + calls ``conversation.run_turn`` (whose UI is the marshaling + :class:`ThreadedUI`, so every UI call it makes is already marshaled) and + routes its own error/completion through ``schedule``. + + ``schedule`` is injected so CI can drive it synchronously; the real app uses + :meth:`schedule`, which reads :attr:`app` lazily (set after construction) so + the construction cycle — schedule needs app, app needs ``start``, the runner + needs the conversation — is broken by deferred attribute assignment. + """ + + def __init__(self, *, inner_ui: RuntimeUI, schedule: Schedule | None = None) -> None: + 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. + self._thread: threading.Thread | None = None + # Set after construction, before app.run(), to break the build cycle; + # read only during/after run(). + self.app: Application[None] | None = None + self.conversation: ConversationRuntime | None = None + # Default to the loop-marshaling schedule (reads self.app lazily); CI + # injects a synchronous one. + self._schedule: Schedule = schedule if schedule is not None else self.schedule + + def schedule(self, fn: Scheduled) -> None: + """Marshal ``fn`` onto the loop thread, then request one repaint. + + Reads :attr:`app` lazily so the construction cycle is broken (``app`` is + set after the app is built). prompt_toolkit coalesces the per-call + ``invalidate()`` into one render tick, so per-token marshal+invalidate is + correct (no throttling needed — measured at the live test, branch 4). + Fails closed when the app is not running (``app``/``loop`` is None). + """ + app = self.app + if app is None: + return + loop = app.loop + if loop is None: + # NOTE: the app is shutting down — run_async clears app.loop on exit. A + # turn still in flight at /exit drops its final _mark_done here, leaving + # _busy True in an already-dead app (harmless for this opt-in dev entry). + # On promotion to the shipping default, drain/join the worker on exit. + return + + def _apply() -> None: + fn() + app.invalidate() + + loop.call_soon_threadsafe(_apply) + + def start(self, text: str) -> None: + """Spawn the worker for one turn. Runs on the loop thread (dock submit). + + 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. + """ + if self._busy: + return + self._busy = True + self._thread = threading.Thread(target=self._run, args=(text,), daemon=True) + self._thread.start() + + def _run(self, text: str) -> None: + """Worker-thread body: run ONE synchronous turn off the loop thread. + + Touches nothing on the loop thread directly — only ``conversation.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. + """ + try: + conversation = self.conversation + if conversation is None: + # Build-order violation (conversation must be set before start()). + # 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) + 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: + self._schedule(self._mark_done) + + def _mark_done(self) -> None: + """Clear the busy flag. Scheduled onto the loop thread, so the flag is + only ever mutated there (paired with the set in :meth:`start`).""" + self._busy = False diff --git a/shellpilot/cli/terminal.py b/shellpilot/cli/terminal.py index 0003351..1473da2 100644 --- a/shellpilot/cli/terminal.py +++ b/shellpilot/cli/terminal.py @@ -9,11 +9,15 @@ from pathlib import Path from urllib.parse import urlsplit +from prompt_toolkit.application import get_app from rich.console import Console from rich.markup import escape from rich.padding import Padding from rich.text import Text +from shellpilot.cli.app_main import run_app +from shellpilot.cli.app_turn import ThreadedUI, TurnRunner +from shellpilot.cli.app_ui import AppUI from shellpilot.cli.attachments import AttachmentError, AttachmentQueue, load_image from shellpilot.cli.banner import render_banner from shellpilot.cli.input import PromptContext, make_input @@ -65,7 +69,7 @@ from shellpilot.policy.approvals import ApprovalReply, ApprovalRequest from shellpilot.policy.risk import RiskLevel from shellpilot.runtime.conversation import ConversationRuntime -from shellpilot.runtime.events import TurnStats +from shellpilot.runtime.events import RuntimeUI, TurnStats from shellpilot.runtime.planner import TaskPlan from shellpilot.skills.loader import discover_skills from shellpilot.tools.base import workspace_display @@ -597,7 +601,25 @@ def _preload(model_name: str) -> None: console.print("[sp.dim]Continuing without stored memory this session.[/sp.dim]") memory = None - ui = TerminalUI(console, glyphs=glyphs, spinner=settings.ui.spinner, workspace=workspace) + # Opt-in full-screen app (design section 31.13). The default REPL is + # byte-identical: app_mode is False unless SHELLPILOT_UI=app, and the else + # branch builds the exact same TerminalUI as before. In app mode the + # conversation is driven through the marshaling ThreadedUI from the start, so + # its plan tools capture the marshaling UI's bound methods at construction. + app_mode = env.get("SHELLPILOT_UI") == "app" + app_ui: AppUI | None = None + app_runner: TurnRunner | None = None + ui: RuntimeUI + if app_mode: + app_ui = AppUI( + glyphs=glyphs, + workspace=workspace, + width_fn=lambda: get_app().output.get_size().columns, + ) + app_runner = TurnRunner(inner_ui=app_ui) + ui = ThreadedUI(inner=app_ui, schedule=app_runner.schedule) + else: + ui = TerminalUI(console, glyphs=glyphs, spinner=settings.ui.spinner, workspace=workspace) runtime = ConversationRuntime( llm=client, settings=settings, @@ -619,6 +641,28 @@ def _preload(model_name: str) -> None: console.print(plan_panel(restored_plan, glyphs)) tid = escape(restored_plan.task_id) console.print(f"[sp.dim]Active plan restored: {tid} ({restored_plan.status}).[/sp.dim]") + + # Hand off to the full-screen app loop (opt-in, design section 31.13). Placed + # after the conversation + restore so the worker-thread turn drives the same + # fully-configured runtime; the default REPL below is reached only when + # app_mode is False, so it stays byte-identical. + if app_mode: + assert app_runner is not None and app_ui is not None + return run_app( + runtime, + app_runner, + app_ui, + workspace=workspace, + model=runtime.model, + profile=settings.runtime.security_profile, + glyphs=glyphs, + commands=command_words(), + is_cloud=egressing_session, + ctx_pct=ctx_percent( + runtime.status().estimated_prompt_tokens, + runtime.status().budget.model_context_tokens, + ), + ) attachments = AttachmentQueue() dispatcher = SlashDispatcher( runtime=runtime, diff --git a/tests/test_app_turn.py b/tests/test_app_turn.py new file mode 100644 index 0000000..1b1ae89 --- /dev/null +++ b/tests/test_app_turn.py @@ -0,0 +1,385 @@ +"""Tests for the worker-thread turn + loop-thread marshaling (design section 31.13). + +Branch 4. These drive ``ThreadedUI`` / ``TurnRunner`` with an injected +*synchronous* schedule (a list/queue that captures the callables) so the +cross-thread handoff is exercised for real (a worker thread runs the turn) yet +drained deterministically on the test thread after ``join``. + +What CANNOT be covered here (needs live test, handed to the user): +- the real ``app.run()`` event loop and ``call_soon_threadsafe`` repaint; +- alt-screen rendering / chrome persisting through a turn; resize/scroll during a + turn; the ``SHELLPILOT_UI=app`` opt-in entry end-to-end (``run_app``). +""" + +from __future__ import annotations + +import queue +import threading +from collections.abc import Sequence +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any + +import pytest +from prompt_toolkit.input.defaults import create_pipe_input +from prompt_toolkit.output import DummyOutput + +from shellpilot.cli.app import build_app +from shellpilot.cli.app_turn import Scheduled, ThreadedUI, TurnRunner +from shellpilot.cli.app_ui import AppUI +from shellpilot.cli.slash import command_words +from shellpilot.cli.theme import UNICODE_GLYPHS +from shellpilot.config.model import Settings +from shellpilot.llm.messages import Message +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 + +# --- helpers ------------------------------------------------------------------ + + +@dataclass +class _RecordingUI: + """Records every UI call in order; optionally forwards to a real inner UI. + + Cross-method call order is captured (FakeUI keeps only per-method lists), so + this proves both ordering and "marshaled, not inline". When ``forward`` is an + ``AppUI``, the recorded calls also flow into the real render path. + """ + + forward: Any | None = None + events: list[tuple[str, tuple[object, ...]]] = field(default_factory=list) + + def _record(self, name: str, *args: object) -> None: + self.events.append((name, args)) + if self.forward is not None: + getattr(self.forward, name)(*args) + + def names(self) -> list[str]: + return [name for name, _ in self.events] + + def stream_token(self, token: str) -> None: + self._record("stream_token", token) + + def stream_thinking(self, text: str) -> None: + self._record("stream_thinking", text) + + def begin_response(self) -> None: + self._record("begin_response") + + def end_response(self) -> None: + self._record("end_response") + + def turn_finished(self, stats: TurnStats) -> None: + self._record("turn_finished", stats) + + def show_status(self, text: str) -> None: + self._record("show_status", text) + + def show_error(self, text: str) -> None: + self._record("show_error", text) + + def show_tool_call(self, name: str, arguments: dict[str, object]) -> None: + self._record("show_tool_call", name, arguments) + + def show_tool_result(self, name: str, success: bool, summary: str) -> None: + self._record("show_tool_result", name, success, summary) + + def show_command_output(self, line: str) -> None: + self._record("show_command_output", line) + + def show_plan_progress(self, plan: object) -> None: + self._record("show_plan_progress", plan) + + def ask_approval(self, request: object) -> object: + raise NotImplementedError + + def ask_plan_approval(self, plan: object, path: str) -> tuple[str, str]: + raise NotImplementedError + + +def _make_runtime(llm: Any, ui: Any, tmp_path: Path) -> ConversationRuntime: + return ConversationRuntime( + llm=llm, + settings=Settings(), + workspace=tmp_path, + behavior=BehaviorInstructions(global_text=None, project_text=None), + ui=ui, + ) + + +_STATS = TurnStats(elapsed_s=1.0, context_tokens=1, context_pct=1, warn=False, output_tokens=1) + +# (method name, positional args) for every fire-and-forget content method. +_FIRE_AND_FORGET: list[tuple[str, tuple[object, ...]]] = [ + ("stream_token", ("tok",)), + ("stream_thinking", ("th",)), + ("begin_response", ()), + ("end_response", ()), + ("turn_finished", (_STATS,)), + ("show_status", ("st",)), + ("show_error", ("er",)), + ("show_tool_call", ("toolname", {"a": 1})), + ("show_tool_result", ("toolname", True, "ok")), + ("show_command_output", ("a line",)), + ("show_plan_progress", (object(),)), +] + + +# --- ThreadedUI marshaling (single-threaded, synchronous schedule) ------------ + + +@pytest.mark.parametrize("method,args", _FIRE_AND_FORGET, ids=[m for m, _ in _FIRE_AND_FORGET]) +def test_fire_and_forget_marshals_not_inline(method: str, args: tuple[object, ...]) -> None: + inner = _RecordingUI() + sink: list[Scheduled] = [] + ui = ThreadedUI(inner=inner, schedule=sink.append) + + getattr(ui, method)(*args) + + # The call did NOT touch the inner UI yet — it enqueued exactly one callable. + assert inner.events == [] + assert len(sink) == 1 + # Draining runs the inner call with the captured args. + sink[0]() + assert inner.events == [(method, args)] + + +def test_marshaled_calls_drain_in_fifo_order() -> None: + inner = _RecordingUI() + sink: list[Scheduled] = [] + ui = ThreadedUI(inner=inner, schedule=sink.append) + + ui.begin_response() + ui.stream_token("a") + ui.show_tool_call("read_file", {"path": "x"}) + ui.end_response() + ui.turn_finished(_STATS) + + assert inner.events == [] # nothing ran inline + for fn in sink: + fn() + assert inner.names() == [ + "begin_response", + "stream_token", + "show_tool_call", + "end_response", + "turn_finished", + ] + + +def test_partial_captures_args_not_late_binding() -> None: + # The classic closure bug: a loop variable would drain as ("b", "b"). partial + # binds the arg at enqueue time, so this drains as ("a", "b"). + inner = _RecordingUI() + sink: list[Scheduled] = [] + ui = ThreadedUI(inner=inner, schedule=sink.append) + + ui.stream_token("a") + ui.stream_token("b") + for fn in sink: + fn() + + assert inner.events == [("stream_token", ("a",)), ("stream_token", ("b",))] + + +def test_blocking_methods_delegate_straight_to_inner() -> None: + # ask_approval/ask_plan_approval return a value → cannot be fire-and-forget; + # they run on the inner directly (the worker thread), which raises until + # branch 7. They are never enqueued. + inner = _RecordingUI() + sink: list[Scheduled] = [] + ui = ThreadedUI(inner=inner, schedule=sink.append) + + with pytest.raises(NotImplementedError): + ui.ask_approval(object()) + with pytest.raises(NotImplementedError): + ui.ask_plan_approval(object(), "p") + assert sink == [] + + +# --- TurnRunner: a full turn across a real worker thread ---------------------- + + +def test_full_turn_runs_on_worker_and_marshals(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) + fake = FakeLLM(script=[answer("Hello from the worker thread.")]) + runtime = _make_runtime(fake, threaded, tmp_path) + runner.conversation = runtime + + runner.start("hi") + assert runner._thread is not None + runner._thread.join(5.0) + assert not runner._thread.is_alive() + + # Drain on the test thread — join established happens-before, so no race. + while not q.empty(): + q.get()() + + # The response reached the REAL AppUI render (not just the recorder). + assert "Hello from the worker thread." in app_ui._render_ansi() + # Ordering: begin_response → stream_token(s) → end_response → turn_finished. + names = inner.names() + assert "stream_token" in names + assert names[0] == "begin_response" + assert names[-1] == "turn_finished" + begin = names.index("begin_response") + end = names.index("end_response") + finished = names.index("turn_finished") + assert begin < end < finished + # busy was set on start and cleared by the scheduled _mark_done. + assert runner._busy is False + + +def test_busy_guard_ignores_second_start(tmp_path: Path) -> None: + gate = threading.Event() + entered = threading.Event() + chats = {"count": 0} + inner_fake = FakeLLM(script=[answer("first"), answer("second")]) + + class _GatedLLM: + """Blocks the first chat until released, so the turn is genuinely in flight.""" + + def chat(self, *args: Any, **kwargs: Any) -> Message: + chats["count"] += 1 + entered.set() + assert gate.wait(5.0) + return inner_fake.chat(*args, **kwargs) + + def health(self) -> bool: + return inner_fake.health() + + def list_models(self) -> Any: + return inner_fake.list_models() + + def model_context_length(self, model: str) -> int | None: + return inner_fake.model_context_length(model) + + def model_capabilities(self, model: str) -> tuple[str, ...]: + return inner_fake.model_capabilities(model) + + def preload(self, model: str, *, keep_alive: str = "5m") -> None: + inner_fake.preload(model, keep_alive=keep_alive) + + app_ui = AppUI(glyphs=UNICODE_GLYPHS, workspace=tmp_path, width_fn=lambda: 80) + q: queue.Queue[Scheduled] = queue.Queue() + runner = TurnRunner(inner_ui=app_ui, schedule=q.put) + threaded = ThreadedUI(inner=app_ui, schedule=q.put) + runtime = _make_runtime(_GatedLLM(), threaded, tmp_path) + runner.conversation = runtime + + runner.start("one") + # busy is set synchronously in start(), before the worker makes progress. + assert runner._busy is True + assert entered.wait(5.0) # the worker really entered the model call → in flight + first_thread = runner._thread + + runner.start("two") # must be IGNORED — a turn is already in flight + assert runner._busy is True + assert runner._thread is first_thread # no second worker spawned + + gate.set() + assert first_thread is not None + first_thread.join(5.0) + while not q.empty(): + q.get()() + + assert runner._busy is False # _mark_done ran on drain + assert chats["count"] == 1 # only ONE worker ever ran the model + + # busy reset → a fresh turn now runs (gate is already set, so it won't block). + runner.start("three") + assert runner._thread is not None and runner._thread is not first_thread + runner._thread.join(5.0) + while not q.empty(): + q.get()() + assert chats["count"] == 2 + assert runner._busy is False + + +def test_worker_exception_is_surfaced_and_clears_busy(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) + # An exhausted script makes the first chat() raise — the worker must not die + # silently; the error must reach the pane and busy must clear. + fake = FakeLLM(script=[]) + runtime = _make_runtime(fake, threaded, tmp_path) + runner.conversation = runtime + + runner.start("boom") + assert runner._thread is not None + runner._thread.join(5.0) + assert not runner._thread.is_alive() # returned cleanly, not killed mid-stack + + while not q.empty(): + q.get()() + + errors = [args for name, args in inner.events if name == "show_error"] + assert len(errors) == 1 + assert "Turn failed" in str(errors[0][0]) + assert runner._busy is False # the finally ran + + +# --- build_app on_submit wiring (headless, pipe input) ------------------------ + + +def _headless_app( + tmp_path: Path, + inp: object, + *, + ui: AppUI, + on_submit: Any | None = None, + commands: Sequence[str] | None = None, +) -> Any: + return build_app( + workspace=tmp_path, + model="gemma4:e4b", + profile="balanced", + glyphs=UNICODE_GLYPHS, + commands=commands if commands is not None else command_words(), + input=inp, # type: ignore[arg-type] + output=DummyOutput(), + ui=ui, + on_submit=on_submit, + ) + + +def test_build_app_on_submit_receives_text(tmp_path: Path) -> None: + received: list[str] = [] + ui = AppUI(glyphs=UNICODE_GLYPHS, width_fn=lambda: 80) + with create_pipe_input() as inp: + app = _headless_app(tmp_path, inp, ui=ui, on_submit=received.append) + inp.send_text("hello there\n") + inp.send_text("/exit\n") + app.run() + assert received == ["hello there"] + # on_submit handled it; the inert branch-3 echo did NOT fire into the pane. + assert "hello there" not in ui._render_ansi() + + +def test_build_app_on_submit_none_falls_back_to_echo(tmp_path: Path) -> None: + ui = AppUI(glyphs=UNICODE_GLYPHS, width_fn=lambda: 80) + with create_pipe_input() as inp: + app = _headless_app(tmp_path, inp, ui=ui, on_submit=None) + inp.send_text("echoed line\n") + inp.send_text("/exit\n") + app.run() + assert "echoed line" in ui._render_ansi() + + +def test_build_app_exit_is_never_routed_to_on_submit(tmp_path: Path) -> None: + received: list[str] = [] + ui = AppUI(glyphs=UNICODE_GLYPHS, width_fn=lambda: 80) + with create_pipe_input() as inp: + app = _headless_app(tmp_path, inp, ui=ui, on_submit=received.append) + inp.send_text("/exit\n") + app.run() # returns → /exit exited cleanly + assert received == []