Bolt-Console/bolt_console/ui.py

"""Terminal rendering: colour, the thinking spinner, tool-call lines, prompts.

Plain ANSI and the standard library only, so the console has no dependency
beyond `requests`. Everything writes through `Terminal._emit`, which holds the
same lock the spinner thread draws under — that is what keeps a spinner frame
from landing in the middle of a line of output.
"""

from __future__ import annotations

import itertools
import os
import sys
import threading
import time
from typing import Callable, List, Optional, TextIO

_RESET = "\033[0m"
_STYLES = {
    "bold": "\033[1m",
    "dim": "\033[2m",
    "red": "\033[31m",
    "green": "\033[32m",
    "yellow": "\033[33m",
    "blue": "\033[34m",
    "magenta": "\033[35m",
    "cyan": "\033[36m",
}

# Output lines shown under a tool call before collapsing to "… +N lines".
TOOL_PREVIEW_LINES = 4


def enable_windows_ansi() -> None:
    """Turn on VT escape processing for the Windows console (off by default
    for older hosts). Harmless everywhere else."""
    if os.name != "nt":
        return
    try:
        import ctypes

        kernel32 = ctypes.windll.kernel32
        handle = kernel32.GetStdHandle(-11)  # STD_OUTPUT_HANDLE
        mode = ctypes.c_uint32()
        if kernel32.GetConsoleMode(handle, ctypes.byref(mode)):
            kernel32.SetConsoleMode(handle, mode.value | 0x0004)  # ENABLE_VIRTUAL_TERMINAL_PROCESSING
    except Exception:
        pass


def _supports_unicode(stream: TextIO) -> bool:
    encoding = (getattr(stream, "encoding", "") or "").lower().replace("-", "")
    return encoding.startswith("utf")


class Glyphs:
    def __init__(self, unicode: bool):
        self.bullet = "●" if unicode else "*"
        self.elbow = "⎿" if unicode else "L"
        self.ellipsis = "…" if unicode else "..."
        self.prompt = "›" if unicode else ">"
        self.spinner = "⠋⠙⠹⠸⠼⠴⠦⠧⠇⠏" if unicode else "|/-\\"


def preview(text: str, limit: int = TOOL_PREVIEW_LINES) -> List[str]:
    """The first *limit* lines of *text*, plus a count of what was cut."""
    lines = text.rstrip("\n").splitlines() or [""]
    if len(lines) <= limit:
        return lines
    return lines[:limit] + [f"... +{len(lines) - limit} lines"]


class Terminal:
    def __init__(
        self,
        out: Optional[TextIO] = None,
        color: bool = True,
        input_fn: Callable[[str], str] = input,
        animate: Optional[bool] = None,
    ):
        self.out = out or sys.stdout
        tty = bool(getattr(self.out, "isatty", lambda: False)())
        self.color = color and tty
        self.animate = tty if animate is None else animate
        self.glyphs = Glyphs(_supports_unicode(self.out))
        self._input = input_fn
        self._lock = threading.RLock()
        self._spinner_label = ""
        self._spinner_started = 0.0
        self._spinner_line = False  # a spinner frame is currently on screen
        self._spinner_thread: Optional[threading.Thread] = None
        self._spinner_stop = threading.Event()
        self._spinner_paused = False
        self._in_reply = False  # streamed sentences continue one bulleted block

    # ── primitives ───────────────────────────────────────────────────────

    def style(self, text: str, *names: str) -> str:
        if not self.color or not names:
            return text
        return "".join(_STYLES[n] for n in names) + text + _RESET

    def _clear_spinner_line(self) -> None:
        if self._spinner_line:
            self.out.write("\r\033[2K" if self.color else "\r" + " " * 60 + "\r")
            self._spinner_line = False

    def _emit(self, text: str) -> None:
        with self._lock:
            self._clear_spinner_line()
            self.out.write(text)
            self.out.flush()

    def line(self, text: str = "") -> None:
        self._emit(text + "\n")

    # ── spinner ──────────────────────────────────────────────────────────

    def start_spinner(self, label: str = "Thinking") -> None:
        self._spinner_label = label
        self._spinner_started = time.monotonic()
        self._spinner_paused = False
        if not self.animate or self._spinner_thread is not None:
            return
        self._spinner_stop.clear()
        self._spinner_thread = threading.Thread(target=self._spin, daemon=True)
        self._spinner_thread.start()

    def stop_spinner(self) -> None:
        thread = self._spinner_thread
        if thread is not None:
            self._spinner_stop.set()
            thread.join(timeout=1)
            self._spinner_thread = None
        with self._lock:
            self._clear_spinner_line()
            self.out.flush()

    def _spin(self) -> None:
        for frame in itertools.cycle(self.glyphs.spinner):
            if self._spinner_stop.wait(0.1):
                return
            with self._lock:
                if self._spinner_paused:
                    continue
                elapsed = int(time.monotonic() - self._spinner_started)
                text = f"{frame} {self._spinner_label}{self.glyphs.ellipsis} ({elapsed}s, ctrl+c to interrupt)"
                self.out.write("\r" + self.style(text, "magenta"))
                self.out.flush()
                self._spinner_line = True

    # ── conversation ─────────────────────────────────────────────────────

    def _bulleted(self, text: str, *bullet_style: str, body_style=()) -> None:
        self._in_reply = False
        lines = text.splitlines() or [""]
        first, rest = lines[0], lines[1:]
        out = [self.style(self.glyphs.bullet, *bullet_style) + " " + self.style(first, *body_style)]
        out += ["  " + self.style(line, *body_style) for line in rest]
        self.line("\n".join(out))

    def assistant(self, text: str) -> None:
        if text.strip():
            self._bulleted(text.strip())
            self.line()

    def assistant_chunk(self, sentence: str) -> None:
        """One streamed sentence. The first opens a bulleted block and the
        rest continue it on their own lines — never mid-line, because the
        spinner redraws by erasing the current line."""
        sentence = sentence.strip()
        if not sentence:
            return
        if self._in_reply:
            self.line("\n".join("  " + part for part in sentence.splitlines()))
        else:
            self._bulleted(sentence)
            self._in_reply = True

    def end_reply(self) -> None:
        self._in_reply = False
        self.line()

    def holding(self, text: str) -> None:
        self._bulleted(text.strip(), "dim", body_style=("dim",))

    def tool_call(self, name: str, argument: str) -> None:
        argument = argument if len(argument) <= 160 else argument[:157] + "..."
        self._bulleted(self.style(name, "bold") + self.style(f"({argument})", "dim"), "green")

    def tool_result(self, text: str, error: bool = False) -> None:
        lines = preview(text)
        elbow = self.glyphs.elbow
        rendered = [f"  {elbow}  " + self.style(lines[0], "red" if error else "dim")]
        rendered += ["     " + self.style(line, "red" if error else "dim") for line in lines[1:]]
        self.line("\n".join(rendered))

    def diff(self, old: str, new: str, limit: int = 12) -> None:
        rows = []
        for sign, text, colour in (("-", old, "red"), ("+", new, "green")):
            lines = text.splitlines() or [""]
            for line in lines[:limit]:
                rows.append("     " + self.style(f"{sign} {line}", colour))
            if len(lines) > limit:
                rows.append("     " + self.style(f"  ... +{len(lines) - limit} lines", "dim"))
        self.line("\n".join(rows))

    def info(self, text: str) -> None:
        self.line(self.style(text, "dim"))

    def warn(self, text: str) -> None:
        self.line(self.style(text, "yellow"))

    def error(self, text: str) -> None:
        self.line(self.style(text, "red"))

    # ── input ────────────────────────────────────────────────────────────

    def ask(self, prompt: str) -> str:
        """Read one line with the spinner held off the screen. EOF and
        Ctrl+C both come back as an empty answer — the caller decides what
        an empty answer means."""
        with self._lock:
            self._spinner_paused = True
            self._clear_spinner_line()
            self.out.flush()
        try:
            return self._input(prompt)
        except (EOFError, KeyboardInterrupt):
            self.line()
            return ""
        finally:
            with self._lock:
                self._spinner_paused = False
                self._spinner_started = time.monotonic()

    def confirm(self, question: str) -> str:
        """y / n / a(lways). Anything unrecognised — including Ctrl+C — is no."""
        answer = self.ask(
            "  " + self.style(question, "yellow") + " "
            + self.style("[y]es / [n]o / [a]lways this session", "dim") + " "
        ).strip().lower()
        if answer in ("y", "yes"):
            return "yes"
        if answer in ("a", "always"):
            return "always"
        return "no"

    def read_message(self) -> str:
        """The main prompt. A trailing backslash continues onto the next line.
        Raises EOFError/KeyboardInterrupt for the REPL to handle."""
        prompt = self.style(self.glyphs.prompt + " ", "bold", "cyan")
        lines = [self._input(prompt)]
        while lines[-1].endswith("\\"):
            lines[-1] = lines[-1][:-1]
            lines.append(self._input(self.style(self.glyphs.ellipsis + " ", "dim")))
        return "\n".join(lines)