traceforge-pcb 1.0.0__py3-none-any.whl

traceforge/tui.py

@@ -0,0 +1,620 @@
+"""The full-screen terminal UI: a Textual front end for :class:`traceforge.chat.ChatSession`.
+
+Everything the session does (slash commands, turns, approvals) lives in
+:mod:`traceforge.chat`; this module only draws it. Needs the ``tui`` extra
+(``pip install "traceforge-pcb[tui]"``).
+"""
+
+from __future__ import annotations
+
+import asyncio
+import json
+from pathlib import Path
+from typing import Any
+
+from rich.markdown import Markdown
+from rich.markup import escape
+from textual import events, on
+from textual.app import App, ComposeResult
+from textual.binding import Binding
+from textual.containers import Horizontal, Vertical, VerticalScroll
+from textual.screen import ModalScreen
+from textual.suggester import SuggestFromList
+from textual.widgets import Button, Collapsible, Input, Label, OptionList, Static
+
+from traceforge import __version__
+from traceforge.agent import AgentEvent
+from traceforge.approval import ApprovalAnswer, ApprovalRequest
+from traceforge.config import Settings
+
+_TODO_MARKS = {"completed": "[green]✓[/]", "in_progress": "[yellow]▸[/]", "pending": "·"}
+
+
+def _brief(arguments: dict[str, Any], limit: int = 70) -> str:
+    for key in ("path", "command", "query", "reference", "description"):
+        if isinstance(arguments.get(key), str):
+            value = arguments[key].splitlines()[0] if arguments[key] else ""
+            return value[:limit]
+    text = json.dumps(arguments, default=str)
+    return text if len(text) <= limit else text[: limit - 1] + "…"
+
+
+# ------------------------------------------------------------------ dialogs
+
+
+class ApprovalScreen(ModalScreen[ApprovalAnswer]):
+    """Yes / always this tool / switch to auto / no (with an optional reason)."""
+
+    BINDINGS = [
+        Binding("y", "answer('yes')", "yes"),
+        Binding("a", "answer('tool')", "always"),
+        Binding("escape", "answer('no')", "no"),
+    ]
+
+    def __init__(self, request: ApprovalRequest):
+        super().__init__()
+        self.request = request
+
+    def compose(self) -> ComposeResult:
+        r = self.request
+        with Vertical(id="dialog"):
+            danger = " [red](dangerous)[/]" if r.dangerous else ""
+            yield Label(f"[bold]Allow {escape(r.tool)}?[/]{danger}")
+            yield Static(escape(r.summary or _brief(r.arguments, 400)), id="summary")
+            yield Input(placeholder="reason, if you say no (optional)", id="reason")
+            with Horizontal(id="buttons"):
+                yield Button("Yes", id="yes", variant="success")
+                yield Button(f"Always {r.tool}", id="tool")
+                yield Button("Switch to auto", id="auto", variant="warning")
+                yield Button("No", id="no", variant="error")
+
+    def on_mount(self) -> None:
+        self.query_one("#yes", Button).focus()
+
+    def action_answer(self, choice: str) -> None:
+        if isinstance(self.focused, Input) and choice != "no":
+            return
+        self._answer(choice)
+
+    def _answer(self, choice: str) -> None:
+        if choice == "no":
+            reason = self.query_one("#reason", Input).value.strip()
+            self.dismiss(ApprovalAnswer(False, "once", reason))
+        else:
+            self.dismiss(ApprovalAnswer(True, "once" if choice == "yes" else choice))
+
+    @on(Button.Pressed)
+    def _pressed(self, event: Button.Pressed) -> None:
+        self._answer(event.button.id or "no")
+
+    @on(Input.Submitted, "#reason")
+    def _reason(self) -> None:
+        self._answer("no")
+
+
+class ChoiceScreen(ModalScreen[str | None]):
+    """Pick one value from a list; escape picks nothing."""
+
+    BINDINGS = [Binding("escape", "dismiss_none", "cancel")]
+
+    def __init__(self, title: str, options: list[tuple[str, str]]):
+        super().__init__()
+        self.title_text, self.options = title, options
+
+    def compose(self) -> ComposeResult:
+        with Vertical(id="dialog"):
+            yield Label(f"[bold]{escape(self.title_text)}[/]")
+            yield OptionList(*[label for _value, label in self.options], id="choices")
+
+    def on_mount(self) -> None:
+        self.query_one(OptionList).focus()
+
+    def action_dismiss_none(self) -> None:
+        self.dismiss(None)
+
+    @on(OptionList.OptionSelected)
+    def _selected(self, event: OptionList.OptionSelected) -> None:
+        self.dismiss(self.options[event.option_index][0])
+
+
+class TextScreen(ModalScreen[str]):
+    """A free-text answer to the agent's question."""
+
+    def __init__(self, question: str):
+        super().__init__()
+        self.question = question
+
+    def compose(self) -> ComposeResult:
+        with Vertical(id="dialog"):
+            yield Label(f"[bold]{escape(self.question)}[/]")
+            yield Input(id="answer")
+
+    def on_mount(self) -> None:
+        self.query_one(Input).focus()
+
+    @on(Input.Submitted)
+    def _submitted(self, event: Input.Submitted) -> None:
+        self.dismiss(event.value)
+
+
+IMAGE_SUFFIXES = {".png", ".jpg", ".jpeg", ".gif", ".webp"}
+MENTION_SUFFIXES = IMAGE_SUFFIXES | {
+    ".kicad_sch",
+    ".kicad_pcb",
+    ".kicad_pro",
+    ".md",
+    ".pdf",
+    ".c",
+    ".h",
+    ".py",
+    ".csv",
+}
+
+
+def mention_token(value: str) -> str | None:
+    """The ``@word`` being typed at the end of the line, without the ``@``."""
+    word = value.rsplit(None, 1)[-1] if value.strip() and not value.endswith(" ") else ""
+    return word[1:] if word.startswith("@") else None
+
+
+def mention_matches(pool: list[str], prefix: str, limit: int = 8) -> list[str]:
+    """Candidates for ``@prefix``: reference designators first, then files."""
+    needle = prefix.lower()
+    starts = [c for c in pool if c.lower().startswith(needle)]
+    inside = [c for c in pool if needle in c.lower() and c not in starts]
+    return (starts + inside)[:limit]
+
+
+class PromptInput(Input):
+    """The prompt line: tab completes a slash command or an @mention; digits pick a chip."""
+
+    BINDINGS = [
+        Binding("tab", "complete", "complete", show=False, priority=True),
+        Binding("ctrl+g", "paste_image", "paste image", show=False),
+    ]
+
+    def action_complete(self) -> None:
+        app = self.app
+        if isinstance(app, TraceforgeApp) and app.accept_mention():
+            return
+        suggestion = getattr(self, "_suggestion", "")
+        if suggestion and suggestion != self.value:
+            self.value = suggestion + " "
+            self.cursor_position = len(self.value)
+
+    def action_paste_image(self) -> None:
+        if isinstance(self.app, TraceforgeApp):
+            self.app.send("/paste", echo=False)
+
+    def on_paste(self, event: events.Paste) -> None:
+        """Pasting the path of an image attaches it instead of typing the path."""
+        raw = event.text.strip().strip("'\"")
+        if raw and "\n" not in raw and Path(raw).suffix.lower() in IMAGE_SUFFIXES:
+            if Path(raw).expanduser().is_file() and isinstance(self.app, TraceforgeApp):
+                event.prevent_default()
+                event.stop()
+                self.app.send(f"/attach {raw}", echo=False)
+
+    async def on_key(self, event: events.Key) -> None:
+        app = self.app
+        if (
+            isinstance(app, TraceforgeApp)
+            and not self.value
+            and event.character
+            and event.character.isdigit()
+            and app.session.suggestions
+            and 1 <= int(event.character) <= len(app.session.suggestions)
+        ):
+            event.prevent_default()
+            event.stop()
+            app.send(app.session.suggestions[int(event.character) - 1])
+
+
+# ----------------------------------------------------------------- the UI
+
+
+class TuiUI:
+    """Implements :class:`traceforge.chat.ChatUI` by drawing into the app."""
+
+    def __init__(self, app: TraceforgeApp):
+        self.app = app
+
+    async def on_event(self, event: AgentEvent) -> None:
+        self.app.handle_event(event)
+
+    async def _modal(self, screen: ModalScreen) -> Any:
+        future: asyncio.Future = asyncio.get_running_loop().create_future()
+
+        def done(result: Any) -> None:
+            if not future.done():
+                future.set_result(result)
+
+        self.app.end_stream()
+        await self.app.push_screen(screen, done)
+        return await future
+
+    async def approve(self, request: ApprovalRequest) -> bool | ApprovalAnswer:
+        return await self._modal(ApprovalScreen(request))
+
+    async def ask(self, question: str, options: list[str]) -> str:
+        if options:
+            answer = await self._modal(ChoiceScreen(question, [(o, o) for o in options]))
+            return answer or ""
+        return await self._modal(TextScreen(question))
+
+    async def choose(self, title: str, options: list[tuple[str, str]]) -> str | None:
+        return await self._modal(ChoiceScreen(title, options))
+
+    async def confirm(self, question: str) -> bool:
+        return await self._modal(ChoiceScreen(question, [("yes", "Yes"), ("no", "No")])) == "yes"
+
+    def show(self, text: str, *, markdown: bool = False) -> None:
+        self.app.write(Markdown(text) if markdown else f"[dim]{escape(text)}[/]", "note")
+
+    def error(self, text: str) -> None:
+        self.app.write(f"[red]{escape(text)}[/]", "error")
+
+    async def clear(self) -> None:
+        self.app.end_stream()
+        await self.app.query_one("#log").remove_children()
+
+    def prefill(self, text: str) -> None:
+        prompt = self.app.query_one("#prompt", Input)
+        prompt.value = text
+        prompt.cursor_position = len(text)
+        prompt.focus()
+
+    def suggest(self, steps: list[str]) -> None:
+        self.app.show_chips(steps)
+
+    def side_panel(self, question: str, answer: str) -> None:
+        panel = self.app.write(Markdown(answer), "side")
+        panel.border_title = f"side question: {question[:60]}"
+
+    async def set_theme(self, name: str) -> None:
+        from traceforge import preferences
+
+        names = sorted(self.app.available_themes)
+        if not name:
+            name = await self._modal(ChoiceScreen("Theme", [(n, n) for n in names])) or ""
+        if not name:
+            return
+        if name not in names:
+            self.error(f"unknown theme {name!r}; choose from {', '.join(names)}")
+            return
+        self.app.theme = name
+        preferences.update_preferences(tui_theme=name)
+        self.show(f"theme {name}")
+
+    async def replay(self, thread) -> None:
+        self.show(f"resumed {thread.id}: {thread.title} ({len(thread.messages)} messages)")
+        for message in reversed(thread.messages):
+            if message.get("role") != "assistant":
+                continue
+            texts = [
+                b.get("text", "")
+                for b in message.get("content", [])
+                if isinstance(b, dict) and b.get("type") == "text"
+            ]
+            if any(texts):
+                self.app.write(Markdown("\n".join(texts)), "assistant")
+                break
+
+
+class TraceforgeApp(App):
+    TITLE = "traceforge"
+    CSS = """
+    #log { height: 1fr; padding: 0 1; }
+    #log > Static { margin: 0 0 1 0; }
+    .user { color: $accent; text-style: bold; }
+    .error { color: $error; }
+    Collapsible { margin: 0 0 1 0; }
+    #status { height: 1; background: $panel; color: $text-muted; padding: 0 1; }
+    #prompt { dock: bottom; }
+    ModalScreen { align: center middle; }
+    #dialog { width: 80%; max-width: 100; height: auto; border: thick $accent;
+              background: $surface; padding: 1 2; }
+    #dialog OptionList { height: auto; max-height: 16; }
+    #buttons { height: auto; margin-top: 1; }
+    #buttons Button { margin-right: 1; }
+    #summary { margin: 1 0; }
+    .side { border: round $secondary; padding: 0 1; }
+    #chips { height: auto; max-height: 6; }
+    #chips Button { margin-right: 1; min-width: 6; }
+    #mentions { height: auto; max-height: 9; display: none; border: round $accent; }
+    """
+    BINDINGS = [
+        Binding("escape", "interrupt", "interrupt"),
+        Binding("ctrl+c", "interrupt_or_quit", "quit", priority=True),
+        Binding("ctrl+l", "clear_screen", "clear"),
+        Binding("ctrl+t", "threads", "threads"),
+        Binding("ctrl+b", "side_question", "side question"),
+        Binding("f1", "help", "help"),
+    ]
+
+    def __init__(
+        self,
+        workspace: Path,
+        settings: Settings,
+        *,
+        mode: str = "agent",
+        resume: str | None = None,
+        message: str | None = None,
+        client: Any = None,
+        store: Any = None,
+    ):
+        super().__init__()
+        from traceforge.chat import ChatSession
+
+        self.ui = TuiUI(self)
+        self.session = ChatSession(
+            workspace, settings, self.ui, mode=mode, client=client, store=store
+        )
+        self.resume_id, self.first_message = resume, message
+        self._stream: Static | None = None
+        self._buffer = ""
+        self._tools: dict[str, Collapsible] = {}
+        self.mention_pool: list[str] = []
+
+    # ------------------------------------------------------------- layout
+
+    def compose(self) -> ComposeResult:
+        yield VerticalScroll(id="log")
+        yield Horizontal(id="chips")
+        yield OptionList(id="mentions")
+        yield Static(id="status")
+        suggestions = [name for name, _summary in self.session.completions()]
+        yield PromptInput(
+            placeholder="Type a request, or / for commands",
+            suggester=SuggestFromList(suggestions, case_sensitive=False),
+            id="prompt",
+        )
+
+    async def on_mount(self) -> None:
+        s = self.session
+        self.write(
+            f"[bold]traceforge[/] {__version__} · {escape(str(s.workspace))}\n"
+            "[dim]Enter sends · / for commands (tab completes) · esc interrupts · ctrl+c quits[/]",
+            "note",
+        )
+        from traceforge import preferences
+
+        saved = preferences.get("tui_theme")
+        if saved in self.available_themes:
+            self.theme = saved
+        self.refresh_status()
+        self.query_one("#prompt", Input).focus()
+        self.query_one("#chips").display = False
+        self.run_worker(self._load_mentions(), exit_on_error=False, group="mentions")
+        if self.resume_id or self.first_message:
+            self.run_worker(self._startup(), exit_on_error=False, group="turn")
+
+    async def _startup(self) -> None:
+        if self.resume_id:
+            try:
+                await self.session.resume(self.resume_id)
+            except LookupError as exc:
+                self.ui.error(str(exc))
+        if self.first_message:
+            await self._submit(self.first_message)
+
+    # --------------------------------------------------------------- output
+
+    def write(self, renderable: Any, classes: str = "") -> Static:
+        self.end_stream()
+        widget = Static(renderable, classes=classes)
+        widget.source = getattr(renderable, "markup", renderable)  # plain text, for tests
+        log = self.query_one("#log", VerticalScroll)
+        log.mount(widget)
+        log.scroll_end(animate=False)
+        return widget
+
+    def end_stream(self) -> None:
+        self._stream, self._buffer = None, ""
+
+    def handle_event(self, event: AgentEvent) -> None:
+        kind, data = event.kind, event.data
+        log = self.query_one("#log", VerticalScroll)
+        if kind == "text":
+            self._buffer += str(data)
+            if self._stream is None:
+                self._stream = Static(classes="assistant")
+                log.mount(self._stream)
+            self._stream.update(Markdown(self._buffer))
+            self._stream.source = self._buffer
+            log.scroll_end(animate=False)
+        elif kind == "tool_call":
+            self.end_stream()
+            arguments = data.get("input") or {}
+            result_box = Static("")
+            card = Collapsible(
+                Static(escape(json.dumps(arguments, indent=2, default=str)[:4000])),
+                result_box,
+                title=f"⚙ {data.get('name')} {_brief(arguments)}",
+                collapsed=True,
+            )
+            card.result_box = result_box  # type: ignore[attr-defined]
+            self._tools[str(data.get("id"))] = card
+            log.mount(card)
+            log.scroll_end(animate=False)
+        elif kind == "tool_result":
+            card = self._tools.get(str(data.get("id")))
+            content = str(data.get("content") or "")
+            mark = "[red]✗[/]" if data.get("is_error") else "[green]✓[/]"
+            if card is None:
+                self.write(f"{mark} {escape(str(data.get('name')))}")
+                return
+            card.title = f"{'✗' if data.get('is_error') else '✓'} {card.title[2:]}"
+            card.result_box.update(  # type: ignore[attr-defined]
+                f"{mark} " + escape(content[:4000] + ("…" if len(content) > 4000 else ""))
+            )
+        elif kind == "todos":
+            rows = [
+                f"{_TODO_MARKS.get(t.get('status'), '·')} {escape(str(t.get('content')))}"
+                for t in data or []
+            ]
+            self.write("[bold]todos[/]\n" + "\n".join(rows), "note")
+        elif kind == "usage":
+            self.refresh_status()
+        elif kind in {"status", "notice"}:
+            message = data.get("message") if isinstance(data, dict) else data
+            self.write(f"[dim]{escape(str(message))}[/]", "note")
+        elif kind == "error":
+            self.ui.error(str(data))
+        elif kind == "done":
+            self.end_stream()
+            self.refresh_status()
+
+    def status_text(self) -> str:
+        from traceforge import models, routing
+
+        s = self.session
+        route = routing.route(s.settings, routing.role_for_mode(s.mode))
+        info = s.status()
+        tokens = s.usage.input_tokens + s.usage.output_tokens + s.usage.cache_read_input_tokens
+        busy = " · [yellow]working…[/]" if s.busy else ""
+        branch = s.thread.meta.get("branch_name") if s.thread else None
+        branch = f" · branch {escape(str(branch))}" if branch else ""
+        undo = f" · {len(s.undo_stack)} undoable" if s.undo_stack else ""
+        return (
+            f"{s.mode} · {escape(models.label(route.model))} · effort {route.effort} · "
+            f"approvals {s.settings.approval} · {tokens:,} tokens · "
+            f"context {info['context_percent']}%{branch}{undo}{busy}"
+        )
+
+    def refresh_status(self) -> None:
+        self.query_one("#status", Static).update(self.status_text())
+
+    # ---------------------------------------------------------------- input
+
+    @on(Input.Submitted, "#prompt")
+    def _prompt(self, event: Input.Submitted) -> None:
+        text = event.value.strip()
+        event.input.value = ""
+        if text:
+            self.send(text)
+
+    def send(self, text: str, *, echo: bool = True) -> None:
+        """Run one line of input (a prompt, a command, or a chosen next step)."""
+        if self.session.busy:
+            self.ui.error("a turn is running; press esc to interrupt it")
+            return
+        if text.isdigit() and self.session.suggestions:
+            number = int(text)
+            if 1 <= number <= len(self.session.suggestions):
+                text = self.session.suggestions[number - 1]
+        self.query_one("#chips").display = False
+        if echo:
+            self.write(f"› {escape(text)}", "user")
+        self.run_worker(self._submit(text), exit_on_error=False, group="turn")
+
+    def show_chips(self, steps: list[str]) -> None:
+        chips = self.query_one("#chips", Horizontal)
+        chips.remove_children()
+        chips.mount(*[Button(f"{n} {s}", id=f"chip{n}") for n, s in enumerate(steps, 1)])
+        chips.display = True
+
+    @on(Button.Pressed, "#chips Button")
+    def _chip(self, event: Button.Pressed) -> None:
+        number = int((event.button.id or "chip0")[4:])
+        if 0 < number <= len(self.session.suggestions):
+            self.send(self.session.suggestions[number - 1])
+
+    # ------------------------------------------------------------ @ mentions
+
+    async def _load_mentions(self) -> None:
+        """Reference designators from the project, then its files, for @ completion."""
+        import os
+
+        def scan() -> list[str]:
+            pool: list[str] = []
+            try:
+                from traceforge.desktop.project import components
+
+                pool += sorted({str(c["ref"]) for c in components(self.session.workspace)})
+            except Exception:  # noqa: BLE001 - completion must never break the UI
+                pass
+            skip = {".git", ".traceforge", "node_modules", "__pycache__", ".venv", "venv"}
+            root = self.session.workspace
+            for folder, dirs, files in os.walk(root):
+                dirs[:] = [d for d in dirs if d not in skip]
+                for name in files:
+                    if Path(name).suffix.lower() in MENTION_SUFFIXES:
+                        pool.append((Path(folder) / name).relative_to(root).as_posix())
+                if len(pool) > 3000:
+                    break
+            return pool
+
+        self.mention_pool = await asyncio.to_thread(scan)
+
+    @on(Input.Changed, "#prompt")
+    def _changed(self, event: Input.Changed) -> None:
+        options = self.query_one("#mentions", OptionList)
+        prefix = mention_token(event.value)
+        found = mention_matches(self.mention_pool, prefix) if prefix is not None else []
+        options.clear_options()
+        if found:
+            options.add_options(found)
+            options.highlighted = 0
+        options.display = bool(found)
+
+    def accept_mention(self) -> bool:
+        options = self.query_one("#mentions", OptionList)
+        if not options.display or options.highlighted is None:
+            return False
+        choice = str(options.get_option_at_index(options.highlighted).prompt)
+        prompt = self.query_one("#prompt", Input)
+        head = prompt.value.rsplit("@", 1)[0]
+        prompt.value = f"{head}@{choice} "
+        prompt.cursor_position = len(prompt.value)
+        options.display = False
+        return True
+
+    async def _submit(self, text: str) -> None:
+        task = asyncio.ensure_future(self.session.submit(text))
+        self.set_timer(0.05, self.refresh_status)
+        try:
+            await task
+        except asyncio.CancelledError:
+            pass
+        except Exception as exc:  # noqa: BLE001 - keep the UI alive
+            self.ui.error(f"{type(exc).__name__}: {exc}")
+        self.end_stream()
+        self.refresh_status()
+        if self.session.exit_requested:
+            await self.session.processes.kill_all()
+            self.exit(0)
+
+    async def action_clear_screen(self) -> None:
+        self.send("/clear", echo=False)
+
+    def action_threads(self) -> None:
+        self.send("/resume", echo=False)
+
+    def action_side_question(self) -> None:
+        self.ui.prefill("/btw ")
+
+    def action_help(self) -> None:
+        self.send("/help", echo=False)
+
+    async def action_interrupt(self) -> None:
+        if await self.session.cancel():
+            self.write("[yellow]interrupted[/]", "note")
+
+    async def action_interrupt_or_quit(self) -> None:
+        if not await self.session.cancel():
+            await self.session.processes.kill_all()
+            self.exit(0)
+
+
+def run_tui(
+    workspace: Path,
+    settings: Settings,
+    *,
+    mode: str = "agent",
+    resume: str | None = None,
+    message: str | None = None,
+) -> int:
+    app = TraceforgeApp(workspace, settings, mode=mode, resume=resume, message=message)
+    result = app.run()
+    return result if isinstance(result, int) else 0