continualcode 0.1.0__py3-none-any.whl

continualcode/tui.py

@@ -0,0 +1,994 @@
+#!/usr/bin/env python3
+"""
+Minimal TUI for a coding agent with online context (prompt) distillation.
+
+Generation uses a teacher prefix that includes a long policy prompt (POLICY_PATH).
+Training distills that policy into weights by running cross-entropy on the approved
+assistant message but with the policy removed (student prefix).
+"""
+
+from __future__ import annotations
+
+import asyncio
+import json
+import os
+import shlex
+import subprocess
+import tempfile
+import time
+from pathlib import Path
+from typing import Any, Literal
+
+from textual import events, on, work
+from textual.app import App, ComposeResult
+from textual.binding import Binding
+from textual.containers import Horizontal, Vertical, VerticalScroll
+from textual.widgets import Input, Label
+
+try:
+    from continualcode.session import ContextDistillSession
+    TINKER_AVAILABLE = True
+except Exception:
+    ContextDistillSession = None  # type: ignore[assignment]
+    TINKER_AVAILABLE = False
+
+from continualcode.tools import READONLY_TOOLS, execute_tool
+
+
+# --- Config ---
+MODEL = os.environ.get("MODEL", "Qwen/Qwen3-4B-Instruct-2507")
+TINKER_URL = os.environ.get("TINKER_URL")
+MAX_TOKENS = int(os.environ.get("MAX_TOKENS", "4096"))
+TEMPERATURE = float(os.environ.get("TEMPERATURE", "0.7"))
+ENABLE_TRAINING = os.environ.get("ENABLE_TRAINING", "1") == "1"
+LEARNING_RATE = float(os.environ.get("LEARNING_RATE", "1e-5"))
+LORA_RANK = int(os.environ.get("LORA_RANK", "32"))
+AUTO_APPROVE_READONLY = os.environ.get("AUTO_APPROVE_READONLY", "0") == "1"
+
+POLICY_PATH = os.environ.get("POLICY_PATH", "./policy_memory.md")
+TRAIN_SCALE_APPROVE = float(os.environ.get("TRAIN_SCALE_APPROVE", "1.0"))
+TRAIN_SCALE_EDITED = float(os.environ.get("TRAIN_SCALE_EDITED", "2.0"))
+DISTILL_MODE = os.environ.get("DISTILL_MODE", "on_policy")  # "on_policy" or "off_policy"
+TRAIN_MAX_TOKENS = int(os.environ.get("TRAIN_MAX_TOKENS", "1024"))
+TRAIN_TEMPERATURE = float(os.environ.get("TRAIN_TEMPERATURE", "1.0"))
+KL_COEF = float(os.environ.get("KL_COEF", "1.0"))
+
+# Checkpoint directory for manual saves
+CHECKPOINT_DIR = Path(os.environ.get("CHECKPOINT_DIR", "./checkpoints"))
+# Load checkpoint on startup (optional)
+LOAD_CHECKPOINT = os.environ.get("LOAD_CHECKPOINT", "")
+
+# Only train on file modification tools
+TRAIN_TOOL_NAMES = {"edit", "write"}
+
+
+# --- Permission Memory (from tinkercode_tui.py) ---
+
+def _resolve_path(project_root: Path, raw_path: str) -> Path:
+    p = Path(raw_path).expanduser()
+    if not p.is_absolute():
+        p = project_root / p
+    return p.resolve()
+
+
+def _scope_for_path(project_root: Path, raw_path: str) -> tuple[Path, str]:
+    abs_path = _resolve_path(project_root, raw_path)
+    try:
+        rel = abs_path.relative_to(project_root)
+    except ValueError:
+        return abs_path.parent, str(abs_path.parent)
+
+    if len(rel.parts) > 1:
+        return project_root / rel.parts[0], f"{rel.parts[0]}/"
+    return project_root, f"{project_root.name}/"
+
+
+class PermissionMemory:
+    """Session-scoped permission memory (similar to Claude Code 'don't ask again')."""
+
+    def __init__(self, project_root: Path) -> None:
+        self.project_root = project_root.resolve()
+        self.allowed_tools: set[str] = set()
+        self.allowed_prefixes: dict[str, set[Path]] = {}
+
+    def allows(self, tool_name: str, tool_args: dict[str, Any]) -> bool:
+        if tool_name in self.allowed_tools:
+            return True
+
+        raw_path = tool_args.get("path")
+        if not isinstance(raw_path, str):
+            return False
+
+        abs_path = _resolve_path(self.project_root, raw_path)
+        for prefix in self.allowed_prefixes.get(tool_name, set()):
+            try:
+                abs_path.relative_to(prefix)
+                return True
+            except ValueError:
+                continue
+        return False
+
+    def remember(self, tool_name: str, tool_args: dict[str, Any]) -> None:
+        if tool_name == "bash":
+            self.allowed_tools.add(tool_name)
+            return
+
+        raw_path = tool_args.get("path")
+        if not isinstance(raw_path, str):
+            self.allowed_tools.add(tool_name)
+            return
+
+        prefix, _label = _scope_for_path(self.project_root, raw_path)
+
+        if tool_name in ("write", "edit"):
+            for t in ("write", "edit"):
+                self.allowed_prefixes.setdefault(t, set()).add(prefix)
+            return
+
+        self.allowed_prefixes.setdefault(tool_name, set()).add(prefix)
+
+    def remember_text(self, tool_name: str, tool_args: dict[str, Any]) -> str:
+        if tool_name == "bash":
+            return "Yes, and don't ask again for Bash commands"
+
+        raw_path = tool_args.get("path")
+        label = f"{self.project_root.name}/"
+        if isinstance(raw_path, str) and raw_path:
+            _prefix, label = _scope_for_path(self.project_root, raw_path)
+
+        if tool_name in ("write", "edit"):
+            return f"Yes, allow all edits in {label}"
+        if tool_name == "read":
+            return f"Yes, allow reading from {label}"
+        if tool_name in ("glob", "grep"):
+            return f"Yes, allow searching in {label}"
+
+        return "Yes, and don't ask again"
+
+
+class TinkerCodeApp(App):
+    """Claude Code-style TUI with online prompt distillation."""
+
+    # Textual enables a command palette (Ctrl+P) by default. In a coding-agent TUI
+    # this is easy to trigger accidentally and looks like "random input". Disable it.
+    ENABLE_COMMAND_PALETTE = False
+
+    DEFAULT_CSS = """
+    Screen {
+        background: #1e1e1e;
+    }
+
+    /* Main output area */
+    #output {
+        height: 1fr;
+        padding: 0;
+        scrollbar-size: 1 1;
+        scrollbar-color: #333333;
+        background: #1e1e1e;
+    }
+
+    #spacer {
+        height: 1fr;
+    }
+
+    /* Message styling - all messages full width */
+    .msg {
+        height: auto;
+        width: 100%;
+        padding: 0 2;
+    }
+
+    /* User messages - dark background bar, full width */
+    .msg-user {
+        color: #ffffff;
+        background: #2a2a2a;
+        padding: 1 2;
+        text-style: bold;
+    }
+
+    /* Assistant messages */
+    .msg-assistant {
+        color: #cccccc;
+        padding: 1 2;
+    }
+
+    /* Tool results - indented */
+    .msg-tool {
+        color: #888888;
+        padding: 0 2 0 4;
+    }
+
+    /* System messages - subtle */
+    .msg-system {
+        color: #666666;
+        padding: 0 2;
+    }
+
+    /* Thinking indicator - italic like Claude Code */
+    .msg-thinking {
+        color: #666666;
+        text-style: italic;
+        padding: 1 2;
+    }
+
+    /* Correction feedback */
+    .msg-correction {
+        color: #dcdcaa;
+        padding: 0 2 0 4;
+    }
+
+    /* Status bar container - very bottom */
+    #status-bar {
+        height: 1;
+        dock: bottom;
+        background: #1e1e1e;
+        padding: 0 2;
+    }
+
+    #status-left {
+        width: auto;
+        color: #555555;
+    }
+
+    #status-right {
+        width: 1fr;
+        text-align: right;
+        color: #555555;
+    }
+
+    /* Bottom area - prompt OR approval */
+    #bottom-area {
+        dock: bottom;
+        height: auto;
+        max-height: 70%;
+        background: #252526;
+        border-bottom: solid #3c3c3c;
+    }
+
+    /* Prompt input row */
+    #prompt-row {
+        height: auto;
+        width: 100%;
+        padding: 1 2 1 2;
+        background: #252526;
+    }
+
+    #prompt-caret {
+        width: 2;
+        color: #cccccc;
+        height: 1;
+    }
+
+    #prompt-input {
+        width: 1fr;
+        border: none;
+        background: transparent;
+        padding: 0;
+        margin: 0;
+        height: 1;
+        color: #cccccc;
+    }
+
+    #prompt-input:focus {
+        border: none;
+        color: #ffffff;
+    }
+
+    /* Approval widget */
+    #approval {
+        display: none;
+        height: auto;
+        padding: 1 2;
+        background: #1e1e1e;
+    }
+
+    #approval.-visible {
+        display: block;
+    }
+
+    .approval-kicker {
+        color: #569cd6;
+        text-style: bold;
+    }
+
+    .approval-title {
+        color: #ffffff;
+        text-style: bold;
+    }
+
+    .approval-desc {
+        color: #9cdcfe;
+        padding: 0 0 0 2;
+    }
+
+    #approval-details {
+        height: auto;
+        max-height: 18;
+        overflow-y: auto;
+        padding: 1 0 1 2;
+    }
+
+    .diff-add {
+        color: #4ec994;
+        background: #1e3a2f;
+        padding: 0 1;
+    }
+
+    .diff-del {
+        color: #f97583;
+        background: #3d1f23;
+        padding: 0 1;
+    }
+
+    .diff-ctx {
+        color: #666666;
+    }
+
+    .approval-question {
+        color: #ffffff;
+        padding: 1 0;
+    }
+
+    .approval-option {
+        height: auto;
+        color: #666666;
+    }
+
+    .approval-option.-selected {
+        color: #ffffff;
+    }
+
+    .approval-hint {
+        color: #444444;
+        padding: 1 0 0 0;
+    }
+
+    #correction-input {
+        display: none;
+        margin: 1 0;
+        border: solid #333333;
+        background: #252526;
+    }
+
+    #correction-input.-visible {
+        display: block;
+    }
+    """
+
+    BINDINGS = [
+        Binding("ctrl+c", "quit", "Quit"),
+        Binding("ctrl+l", "clear", "Clear"),
+        Binding("ctrl+s", "save_checkpoint", "Save"),
+        Binding("?", "help", "Help"),
+    ]
+
+    def __init__(self) -> None:
+        super().__init__()
+        self.session: ContextDistillSession | None = None
+        self.session_ready = False
+        self.permissions = PermissionMemory(Path.cwd())
+        self._loading = True
+
+        # Metrics for status bar
+        self.total_tokens = 0
+        self.last_loss = 0.0
+
+        # Approval state
+        self._approval_future: asyncio.Future[dict[str, Any]] | None = None
+        self._approval_selection = 0
+        self._approval_options: list[tuple[str, str]] = []
+        self._show_correction = False
+        self._current_tool_name: str = ""
+        self._current_tool_args: dict[str, Any] = {}
+        self._edited_args: dict[str, Any] | None = None
+
+    def _prompt_input(self) -> Input:
+        return self.query_one("#prompt-input", Input)
+
+    def _prompt_row_visible(self) -> bool:
+        return bool(self.query_one("#prompt-row").display)
+
+    def _sync_focus(self) -> None:
+        """Keep focus on the right input widget for the current UI mode."""
+        if self._approval_future is not None:
+            if self._show_correction:
+                self.query_one("#correction-input", Input).focus()
+            else:
+                self.set_focus(None)
+            return
+
+        prompt_input = self.query_one("#prompt-input", Input)
+        if not prompt_input.disabled:
+            prompt_input.focus()
+
+    def _show_correction_input(self) -> None:
+        correction_input = self.query_one("#correction-input", Input)
+        self._show_correction = True
+        correction_input.disabled = False
+        correction_input.add_class("-visible")
+        correction_input.focus()
+
+    def _hide_correction_input(self) -> None:
+        correction_input = self.query_one("#correction-input", Input)
+        self._show_correction = False
+        correction_input.remove_class("-visible")
+        correction_input.disabled = True
+        self.set_focus(None)
+
+    def compose(self) -> ComposeResult:
+        with VerticalScroll(id="output", can_focus=False, can_focus_children=False):
+            yield Label("", id="spacer")
+        with Vertical(id="bottom-area"):
+            with Horizontal(id="prompt-row"):
+                yield Label("> ", id="prompt-caret")
+                yield Input(id="prompt-input")
+            with Vertical(id="approval"):
+                yield Label("Tool use", classes="approval-kicker")
+                yield Label("", id="approval-title", classes="approval-title")
+                yield Label("", id="approval-desc", classes="approval-desc")
+                with Vertical(id="approval-details"):
+                    pass
+                yield Label("", id="approval-question", classes="approval-question")
+                yield Label("", id="opt-0", classes="approval-option")
+                yield Label("", id="opt-1", classes="approval-option")
+                yield Label("", id="opt-2", classes="approval-option")
+                yield Label("", id="opt-3", classes="approval-option")
+                yield Input(placeholder="Type what to change…", id="correction-input")
+                yield Label("Esc cancel · e edit · Tab correct", classes="approval-hint")
+        with Horizontal(id="status-bar"):
+            yield Label("? for shortcuts", id="status-left")
+            yield Label("", id="status-right")
+
+    async def on_mount(self) -> None:
+        # The correction input starts hidden; keep it disabled so it can't steal focus.
+        self.query_one("#correction-input", Input).disabled = True
+
+        # Init can take a while (model load, network, auth). Hide/disable the prompt
+        # during init so stray keystrokes don't show up as random characters.
+        self.query_one("#prompt-row").display = False
+        self._prompt_input().disabled = True
+
+        self.call_after_refresh(self._sync_focus)
+        self._update_status()
+
+        if TINKER_AVAILABLE:
+            init_msg = self.output("Initializing Tinker session...", style="thinking")
+            try:
+                self.session = ContextDistillSession(
+                    model=MODEL,
+                    tinker_url=TINKER_URL,
+                    max_tokens=MAX_TOKENS,
+                    temperature=TEMPERATURE,
+                    enable_training=ENABLE_TRAINING,
+                    lora_rank=LORA_RANK,
+                    learning_rate=LEARNING_RATE,
+                    policy_path=POLICY_PATH,
+                    distill_mode=DISTILL_MODE,
+                    train_max_tokens=TRAIN_MAX_TOKENS,
+                    train_temperature=TRAIN_TEMPERATURE,
+                    kl_coef=KL_COEF,
+                )
+                await self.session.init()
+
+                # Load checkpoint if specified
+                if LOAD_CHECKPOINT and self.session.training_client is not None:
+                    if self.session.load_checkpoint(LOAD_CHECKPOINT):
+                        self.output(f"Loaded checkpoint: {LOAD_CHECKPOINT}", style="system")
+                    else:
+                        self.output(f"Failed to load checkpoint: {LOAD_CHECKPOINT}", style="system")
+
+                self.session_ready = True
+                init_msg.remove()
+                status = f"Ready. Model: {MODEL} | policy={POLICY_PATH} | mode={DISTILL_MODE}"
+                if self.session.training_client is not None:
+                    status += " [distillation ON]"
+                self.output(status, style="system")
+
+                # Now that we're ready, enable the prompt.
+                self._loading = False
+                self.query_one("#prompt-row").display = True
+                self._prompt_input().disabled = False
+                self.call_after_refresh(self._sync_focus)
+            except Exception as e:
+                init_msg.remove()
+                self.output(
+                    f"Failed to initialize Tinker: {e}\n"
+                    "Set TINKER_API_KEY (or pass api_key to the client) and restart.",
+                    style="system",
+                )
+                # Avoid a confusing UX where the prompt accepts input but nothing can run.
+                self._prompt_input().disabled = True
+                self._loading = False
+        else:
+            self.output("Tinker not available.", style="system")
+            self._prompt_input().disabled = True
+            self._loading = False
+
+    async def on_key(self, event: events.Key) -> None:
+        # While we're loading (and the prompt is hidden), swallow printable keys so
+        # impatient typing doesn't render as "random characters" (e.g. a lone 'p').
+        # Let control bindings (Ctrl+C/Ctrl+L/Ctrl+S/?) continue to work.
+        if self._loading and len(event.key) == 1 and event.key.isprintable():
+            event.stop()
+            return
+
+        # Fall back to the base behavior for the rest of the app.
+        await super().on_key(event)
+
+    def action_help(self) -> None:
+        self.output(
+            "Keys: Enter=send · Ctrl+L=clear · Ctrl+S=save · Ctrl+C=quit · "
+            "Approval: 1/y approve · 2 approve+remember · 3/n/Esc deny · 4/Tab correct · e edit args",
+            style="system",
+        )
+
+    def output(self, text: str, style: str = "assistant") -> Label:
+        out = self.query_one("#output", VerticalScroll)
+        prefix = {
+            "user": "> ",
+            "assistant": "● ",
+            "tool": "⎿ ",
+            "correction": "⎿ ",
+            "thinking": "● ",
+        }.get(style, "")
+        label = Label(f"{prefix}{text}", classes=f"msg msg-{style}")
+        out.mount(label)
+        out.scroll_end()
+        self._update_status()
+        self.call_after_refresh(self._sync_focus)
+        return label
+
+    def _show_approval(self, tool_name: str, tool_args: dict[str, Any], remember_text: str) -> None:
+        """Show the inline approval UI."""
+        self._current_tool_name = tool_name
+        self._current_tool_args = dict(tool_args)
+        self._edited_args = None
+
+        self.query_one("#prompt-row").display = False
+        self.query_one("#prompt-input", Input).disabled = True
+        approval = self.query_one("#approval")
+        approval.add_class("-visible")
+        self.set_focus(None)
+
+        self._refresh_approval_ui(tool_name, tool_args, remember_text)
+        self.call_after_refresh(self._sync_focus)
+
+    def _refresh_approval_ui(self, tool_name: str, tool_args: dict[str, Any], remember_text: str) -> None:
+        """Refresh the approval UI with current tool args."""
+        file_path = str(tool_args.get("path", ""))
+        edited_suffix = " (edited)" if self._edited_args is not None else ""
+
+        if tool_name == "edit":
+            title = f"Edit file {file_path}{edited_suffix}"
+        elif tool_name == "write":
+            title = f"Create file {file_path}{edited_suffix}"
+        elif tool_name == "read":
+            title = f"Read file {file_path}{edited_suffix}"
+        elif tool_name == "bash":
+            title = f"Bash command{edited_suffix}"
+        else:
+            title = f"{tool_name}{edited_suffix}"
+        self.query_one("#approval-title", Label).update(title)
+
+        # Set description
+        if tool_name == "bash":
+            self.query_one("#approval-desc", Label).update(str(tool_args.get("cmd", "")))
+        else:
+            self.query_one("#approval-desc", Label).update("")
+
+        # Set details
+        details = self.query_one("#approval-details", Vertical)
+        details.remove_children()
+        max_lines = 18
+
+        if tool_name == "edit":
+            old = str(tool_args.get("old", ""))
+            new = str(tool_args.get("new", ""))
+            old_lines = old.split("\n")[:max_lines]
+            new_lines = new.split("\n")[:max_lines]
+            for i, line in enumerate(old_lines, 1):
+                details.mount(Label(f"{i:3} - {line}", classes="diff-del"))
+            for i, line in enumerate(new_lines, 1):
+                details.mount(Label(f"{i:3} + {line}", classes="diff-add"))
+        elif tool_name == "write":
+            content = str(tool_args.get("content", ""))
+            lines = content.split("\n")
+            for i, line in enumerate(lines[:max_lines], 1):
+                details.mount(Label(f"{i:3} + {line}", classes="diff-add"))
+            if len(lines) > max_lines:
+                details.mount(Label(f"    ... ({len(lines) - max_lines} more lines)", classes="diff-ctx"))
+        elif tool_name == "bash":
+            cmd = str(tool_args.get("cmd", ""))
+            details.mount(Label(f"$ {cmd}", classes="diff-ctx"))
+        elif tool_name in ("read", "glob", "grep"):
+            path = tool_args.get("path", tool_args.get("pat", ""))
+            details.mount(Label(f"{path}", classes="diff-ctx"))
+
+        # Set question
+        name = Path(file_path).name if file_path else "this"
+        if tool_name == "edit":
+            question = f"Do you want to make this edit to {name}?"
+        elif tool_name == "write":
+            question = f"Do you want to create {name}?"
+        elif tool_name == "read":
+            question = f"Do you want to read {name}?"
+        elif tool_name == "bash":
+            question = "Do you want to run this command?"
+        else:
+            question = "Do you want to proceed?"
+        self.query_one("#approval-question", Label).update(question)
+
+        # Set options (4 options like Claude Code)
+        self._approval_options = [
+            ("yes", "Yes"),
+            ("yes_always", f"{remember_text} (shift+Tab)"),
+            ("no", "No"),
+            ("correct", "Type correction (Tab)"),
+        ]
+        self._approval_selection = 0
+        self._sync_approval_options()
+
+        # Reset correction input
+        correction_input = self.query_one("#correction-input", Input)
+        self._show_correction = False
+        correction_input.remove_class("-visible")
+        correction_input.disabled = True
+        correction_input.value = ""
+
+    def _sync_approval_options(self) -> None:
+        """Update option labels with selection indicator."""
+        for i, (_decision, text) in enumerate(self._approval_options):
+            label = self.query_one(f"#opt-{i}", Label)
+            prefix = "> " if i == self._approval_selection else "  "
+            label.update(f"{prefix}{i + 1}. {text}")
+            label.set_class(i == self._approval_selection, "-selected")
+
+    def _hide_approval(self) -> None:
+        """Hide approval UI, show prompt."""
+        self.query_one("#approval").remove_class("-visible")
+        self.query_one("#prompt-row").display = True
+        self.query_one("#prompt-input", Input).disabled = False
+        self._hide_correction_input()
+        self.call_after_refresh(self._sync_focus)
+
+    def _resolve_approval(self, decision: str, correction: str | None = None) -> None:
+        """Complete the approval with a decision."""
+        if self._approval_future is None:
+            return
+
+        result = {
+            "decision": decision,
+            "additional_instructions": None,
+            "edited_args": self._edited_args,
+            "correction": correction,
+        }
+
+        self._hide_approval()
+        self._approval_future.set_result(result)
+        self._approval_future = None
+
+    def _edit_args_in_editor(self) -> None:
+        """Open $EDITOR to edit tool args."""
+        editor_raw = os.environ.get("EDITOR") or "vim"
+        editor_cmd = shlex.split(editor_raw)
+
+        # Use edited args if already edited, otherwise original
+        args_to_edit = self._edited_args if self._edited_args is not None else self._current_tool_args
+
+        tmp_path: str | None = None
+        try:
+            with tempfile.NamedTemporaryFile(mode="w", suffix=".json", delete=False, encoding="utf-8") as f:
+                json.dump(args_to_edit, f, indent=2, ensure_ascii=False)
+                f.write("\n")
+                tmp_path = f.name
+
+            with self.suspend():
+                subprocess.run([*editor_cmd, tmp_path], check=False)
+
+            with open(tmp_path, "r", encoding="utf-8") as f:
+                edited = json.load(f)
+
+            if isinstance(edited, dict):
+                self._edited_args = edited
+                self._current_tool_args = edited
+                # Refresh the UI to show edited content
+                remember_text = self.permissions.remember_text(self._current_tool_name, edited)
+                self._refresh_approval_ui(self._current_tool_name, edited, remember_text)
+                self.call_after_refresh(self._sync_focus)
+
+        except Exception as e:
+            self.output(f"Edit failed: {e}", style="system")
+        finally:
+            if tmp_path is not None:
+                try:
+                    os.unlink(tmp_path)
+                except OSError:
+                    pass
+
+    async def request_approval(self, tool_name: str, tool_args: dict[str, Any]) -> tuple[dict[str, Any], int]:
+        if AUTO_APPROVE_READONLY and tool_name in READONLY_TOOLS:
+            return {"decision": "yes", "additional_instructions": None, "edited_args": None, "correction": None, "auto_approved": True}, 0
+
+        if self.permissions.allows(tool_name, tool_args):
+            return {"decision": "yes", "additional_instructions": None, "edited_args": None, "correction": None, "auto_approved": True}, 0
+
+        remember_text = self.permissions.remember_text(tool_name, tool_args)
+        start = time.time()
+
+        # Create future and show UI
+        self._approval_future = asyncio.Future()
+        self._show_approval(tool_name, tool_args, remember_text)
+
+        # Wait for user decision
+        result = await self._approval_future
+        latency = int((time.time() - start) * 1000)
+
+        if result.get("decision") == "yes_always":
+            final_args = result.get("edited_args") or tool_args
+            self.permissions.remember(tool_name, final_args)
+
+        result["auto_approved"] = False
+        return result, latency
+
+    def on_key(self, event: events.Key) -> None:
+        """Handle keyboard input for approval UI (and prompt focus recovery)."""
+        if self._approval_future is None:
+            # Recover if focus is lost (e.g., user clicked output). Keep typing seamless.
+            if (
+                event.character
+                and event.character.isprintable()
+                and self.query_one("#prompt-row").display
+                and not self.query_one("#prompt-input", Input).disabled
+                and self.screen.focused is not self.query_one("#prompt-input", Input)
+            ):
+                prompt_input = self.query_one("#prompt-input", Input)
+                prompt_input.focus()
+                prompt_input.insert_text_at_cursor(event.character)
+                event.stop()
+            return
+
+        key = event.key
+
+        # If correction input is focused, only handle escape/enter
+        if self._show_correction:
+            if key == "escape":
+                self._hide_correction_input()
+                event.stop()
+            return
+
+        if key in ("1", "y"):
+            self._resolve_approval("yes")
+            event.stop()
+        elif key in ("2",) or key == "shift+tab":
+            self._resolve_approval("yes_always")
+            event.stop()
+        elif key in ("3", "n"):
+            self._resolve_approval("no")
+            event.stop()
+        elif key == "4":
+            # Option 4: Type correction
+            self._show_correction_input()
+            event.stop()
+        elif key == "escape":
+            self._resolve_approval("no")
+            event.stop()
+        elif key in ("up", "left"):
+            self._approval_selection = max(0, self._approval_selection - 1)
+            self._sync_approval_options()
+            event.stop()
+        elif key in ("down", "right"):
+            self._approval_selection = min(len(self._approval_options) - 1, self._approval_selection + 1)
+            self._sync_approval_options()
+            event.stop()
+        elif key == "enter":
+            decision = self._approval_options[self._approval_selection][0]
+            if decision == "correct":
+                # Show correction input
+                self._show_correction_input()
+            else:
+                self._resolve_approval(decision)
+            event.stop()
+        elif key == "tab":
+            # Tab = show correction input (option 4)
+            self._show_correction_input()
+            event.stop()
+        elif key == "e":
+            # Edit args in $EDITOR
+            self._edit_args_in_editor()
+            event.stop()
+
+    @on(Input.Submitted, "#correction-input")
+    def on_correction_submitted(self, event: Input.Submitted) -> None:
+        """Handle Enter in correction input - deny with correction text."""
+        correction = event.input.value.strip()
+        self._hide_correction_input()
+        # Correction = rejection with feedback for DPO
+        self._resolve_approval("correct", correction=correction if correction else None)
+        event.stop()
+
+    @on(Input.Submitted, "#prompt-input")
+    def on_submit(self, event: Input.Submitted) -> None:
+        text = event.value.strip()
+        if not text:
+            return
+        event.input.value = ""
+        self.output(text, style="user")
+        self._run_agent_loop(text)
+
+    @work(exclusive=True)
+    async def _run_agent_loop(self, user_input: str) -> None:
+        session = self.session
+        if session is None:
+            return
+
+        session.add_user(user_input)
+
+        try:
+            while True:
+                thinking = self.output("Thinking...", style="thinking")
+                prompt_messages = list(session.messages)
+
+                try:
+                    message, ok = await session.sample()
+                except Exception as e:
+                    thinking.remove()
+                    self.output(f"Error: {e}", style="system")
+                    return
+
+                thinking.remove()
+
+                if not ok:
+                    self.output("Parse failed", style="system")
+                    return
+
+                content = message.get("content", "")
+                if isinstance(content, list):
+                    text = "".join(p.get("text", "") for p in content if p.get("type") == "text")
+                else:
+                    text = content
+                if text and str(text).strip():
+                    self.output(str(text), style="assistant")
+
+                session.add_assistant(message)
+
+                tool_calls = message.get("tool_calls", [])
+                if not tool_calls:
+                    return
+
+                for tc in tool_calls:
+                    name = tc.function.name
+                    try:
+                        model_args = json.loads(tc.function.arguments or "{}")
+                    except json.JSONDecodeError:
+                        model_args = {}
+
+                    approval, latency = await self.request_approval(name, model_args)
+
+                    decision = approval.get("decision", "no")
+                    correction = approval.get("correction")
+                    approved = decision in ("yes", "yes_always")
+                    edited_args = approval.get("edited_args") if isinstance(approval.get("edited_args"), dict) else None
+                    final_args = edited_args if (approved and edited_args is not None) else model_args
+
+                    # If the user edited args, mutate the message tool call so the conversation
+                    # (and any training) reflects the actual executed tool.
+                    if approved and edited_args is not None:
+                        tc.function.arguments = json.dumps(final_args, ensure_ascii=False)
+
+                    if approved:
+                        result = execute_tool(name, final_args)
+                        lines = result.split("\n")
+                        if name == "read":
+                            self.output(f"Read {len(lines)} lines", style="tool")
+                        else:
+                            preview = lines[0][:100] if lines else ""
+                            if len(lines) > 1:
+                                preview += f" ... +{len(lines) - 1} lines"
+                            self.output(preview, style="tool")
+                    else:
+                        # Denied or corrected
+                        deny_reason = correction or "User denied this tool call."
+                        result = json.dumps(
+                            {
+                                "status": "denied",
+                                "tool_name": name,
+                                "tool_call_id": tc.id,
+                                "tool_args": model_args,
+                                "reason": deny_reason,
+                                "correction": correction,
+                            },
+                            ensure_ascii=False,
+                        )
+                        if correction:
+                            self.output(correction, style="correction")
+                        else:
+                            self.output("Denied", style="tool")
+
+                    session.add_tool_result(tc.id, result)
+
+                    # Determine decision label for training
+                    auto = bool(approval.get("auto_approved", False))
+                    if auto:
+                        decision_label: Literal["approved", "denied", "edited", "corrected", "auto_approved"] = "auto_approved"
+                    elif decision == "correct":
+                        decision_label = "corrected"
+                    elif not approved:
+                        decision_label = "denied"
+                    elif edited_args is not None:
+                        decision_label = "edited"
+                    else:
+                        decision_label = "approved"
+
+                    # Train on approved write/edit (single tool call only)
+                    if (
+                        approved
+                        and not auto
+                        and len(tool_calls) == 1
+                        and name in TRAIN_TOOL_NAMES
+                        and session.training_client is not None
+                    ):
+                        scale = TRAIN_SCALE_EDITED if edited_args else TRAIN_SCALE_APPROVE
+                        metrics = session.train_on_approval(
+                            prompt_messages=prompt_messages,
+                            assistant_message=message,
+                            scale=scale,
+                        )
+                        if metrics:
+                            self.last_loss = metrics.get("loss", 0.0)
+                            self.total_tokens += int(metrics.get("tokens", 0))
+                            tag = "edit" if edited_args else "approve"
+                            extra = ""
+                            if "kl_student_teacher" in metrics:
+                                extra = (
+                                    f" st={metrics.get('kl_student_teacher', 0.0):.3f}"
+                                    f" kl={metrics.get('approx_kl', 0.0):.4f}"
+                                    f" r={metrics.get('ratio', 1.0):.3f}"
+                                )
+                            self.output(
+                                f"cd[{tag}] step {int(metrics.get('step', 0))} | "
+                                f"loss={self.last_loss:.3f} tok={int(metrics.get('tokens', 0))} w={scale:g}{extra}",
+                                style="system",
+                            )
+                            self._update_status()
+        finally:
+            self.call_after_refresh(self._sync_focus)
+
+    def _update_status(self) -> None:
+        if self.session and self.session.train_steps > 0:
+            tokens_k = self.total_tokens / 1000
+            status = f"Steps: {self.session.train_steps} | Loss: {self.last_loss:.2f} | Tokens: {tokens_k:.1f}k"
+        else:
+            status = ""
+        self.query_one("#status-right", Label).update(status)
+
+    def action_clear(self) -> None:
+        out = self.query_one("#output", VerticalScroll)
+        out.remove_children()
+        out.mount(Label("", id="spacer"))
+        if self.session:
+            self.session.clear()
+        self._update_status()
+        self.output("Cleared.", style="system")
+        self.call_after_refresh(self._sync_focus)
+
+    def action_save_checkpoint(self) -> None:
+        """Save current LoRA weights to disk (Ctrl+S)."""
+        if not self.session or self.session.training_client is None:
+            self.output("No training session to save.", style="system")
+            return
+
+        # Create checkpoint directory if needed
+        CHECKPOINT_DIR.mkdir(parents=True, exist_ok=True)
+
+        # Generate checkpoint name with step count
+        step = self.session.train_steps
+        checkpoint_path = CHECKPOINT_DIR / f"step_{step:06d}"
+
+        if self.session.save_checkpoint(str(checkpoint_path)):
+            self.output(f"Saved checkpoint: {checkpoint_path}", style="system")
+        else:
+            self.output("Failed to save checkpoint.", style="system")
+
+
+if __name__ == "__main__":
+    TinkerCodeApp().run()