zai-cli 0.1.0__py3-none-any.whl

zai/cli/interactive.py

@@ -0,0 +1,518 @@
+import os
+import subprocess
+from pathlib import Path
+
+from rich.console import Console
+from rich.panel import Panel
+from rich.prompt import Confirm, Prompt
+from rich.table import Table
+
+from ..config import get_api_key, get_models, load_config
+from ..core.agent import plan_agent, run_agent, undo_last
+from ..core.commands import get_command_prompt, list_commands, load_commands
+from ..core.hooks import fire as fire_hook
+from ..core.fallback import has_available_provider
+from ..core.context import ContextManager
+from ..core.memory import add_project, get_last_session, save_session
+from ..core.session import (
+    delete_session,
+    list_sessions,
+    load_history,
+    load_latest_session,
+    rename_session,
+    save_auto_history,
+    save_history,
+    search_sessions,
+)
+from ..core.watch import is_watching, start_watch, stop_watch
+from ..plugins import loader as plugin_loader
+from ..providers.base import Message
+from ..tools.git import get_diff
+from ..ui.input import InteractiveInput
+from .common import closest_command
+
+console = Console()
+
+CORE_COMMANDS = {
+    "help", "clear", "files", "diff", "undo", "plan", "test",
+    "watch", "commit", "review", "fix", "explain", "feature",
+    "session", "resume", "model", "memory", "commands",
+}
+
+
+def _show_help() -> None:
+    console.print(
+        "[bold cyan]── Core ──[/bold cyan]\n"
+        "[cyan]/help[/cyan]                This menu\n"
+        "[cyan]/clear[/cyan]               Clear screen\n"
+        "[cyan]/files[/cyan]               List files in current folder\n"
+        "[cyan]/diff[/cyan]                Show git diff (colored)\n"
+        "[cyan]/undo[/cyan]                Restore last AI-changed file\n"
+        "\n[bold cyan]── AI Modes ──[/bold cyan]\n"
+        "[cyan]/plan <task>[/cyan]         Plan first, approve, then execute\n"
+        "[cyan]/test[/cyan]                Run pytest + auto-fix failures\n"
+        "[cyan]/watch[/cyan]               Toggle file-change watcher\n"
+        "\n[bold cyan]── Git & Code ──[/bold cyan]\n"
+        "[cyan]/commit[/cyan]              Auto git commit with AI message\n"
+        "[cyan]/review[/cyan]              Code review of current changes\n"
+        "[cyan]/fix <file>[/cyan]          Fix bugs in a file\n"
+        "[cyan]/explain <file>[/cyan]      Explain code in a file\n"
+        "[cyan]/feature <desc>[/cyan]      Build a new feature\n"
+        "\n[bold cyan]── Session ──[/bold cyan]\n"
+        "[cyan]/resume [name][/cyan]       Resume latest or named session\n"
+        "[cyan]/session save [name][/cyan] Save current conversation\n"
+        "[cyan]/session load <name>[/cyan] Load a saved conversation\n"
+        "[cyan]/session list[/cyan]        List saved sessions\n"
+        "[cyan]/session search <query>[/cyan] Search saved conversations\n"
+        "[cyan]/session rename <old> <new>[/cyan] Rename a saved session\n"
+        "[cyan]/session delete <name>[/cyan] Delete a saved session\n"
+        "\n[bold cyan]── Settings ──[/bold cyan]\n"
+        "[cyan]/model list[/cyan]          Show available AI models\n"
+        "[cyan]/model <name>[/cyan]        Switch model (gemini/groq/etc)\n"
+        "[cyan]/memory[/cyan]              Show last session info\n"
+        "[cyan]/commands[/cyan]            List all slash commands\n"
+        "\n[dim]Just type naturally — zai creates files, reads code, "
+        "runs commands automatically.[/dim]"
+    )
+
+
+def _correct_slash_command(message: str) -> str:
+    stripped = message.strip()
+    if not stripped.startswith("/") or stripped.startswith("//"):
+        return message
+    parts = stripped[1:].split(" ", 1)
+    typed_name = parts[0].lower()
+    arguments = parts[1] if len(parts) > 1 else ""
+    available = (
+        CORE_COMMANDS
+        | set(load_commands().keys())
+        | set(plugin_loader.get_all_commands().keys())
+    )
+    if typed_name in available:
+        return message
+    corrected = closest_command(typed_name, available)
+    if not corrected:
+        return message
+    console.print(
+        f"[yellow]Command '/{typed_name}' not found; "
+        f"using '/{corrected}'.[/yellow]"
+    )
+    return f"/{corrected}{f' {arguments}' if arguments else ''}"
+
+
+def _show_files(cwd: str) -> None:
+    try:
+        for name in sorted(os.listdir(cwd)):
+            path = os.path.join(cwd, name)
+            if os.path.isdir(path):
+                console.print(f"  [blue]{name}/[/blue]")
+            else:
+                console.print(
+                    f"  [cyan]{name}[/cyan] [dim]{os.path.getsize(path)//1024}kb[/dim]"
+                )
+    except Exception as error:
+        console.print(f"[red]{error}[/red]")
+
+
+def _show_diff() -> None:
+    diff = get_diff()
+    if not diff:
+        console.print("[dim]No changes.[/dim]")
+        return
+    colored = []
+    for line in diff.splitlines():
+        if line.startswith("+") and not line.startswith("+++"):
+            colored.append(f"[green]{line}[/green]")
+        elif line.startswith("-") and not line.startswith("---"):
+            colored.append(f"[red]{line}[/red]")
+        elif line.startswith("@@"):
+            colored.append(f"[cyan]{line}[/cyan]")
+        else:
+            colored.append(f"[dim]{line}[/dim]")
+    console.print(Panel(
+        "\n".join(colored),
+        title="[cyan]Git Diff[/cyan]",
+        border_style="cyan",
+    ))
+
+
+def _connect_mcp_servers() -> None:
+    try:
+        from ..mcp.client import connect
+        from ..mcp.manager import build_server_command, load_mcp
+
+        for server in load_mcp():
+            try:
+                command, env = build_server_command(server, os.getcwd())
+            except ValueError as error:
+                console.print(
+                    f"[dim]MCP '{server.get('name', 'server')}': {error}[/dim]"
+                )
+                continue
+            connect(server["name"], command, env)
+    except Exception:
+        pass
+
+
+def run_interactive(model: str = None) -> None:
+    cwd = os.getcwd()
+    if not has_available_provider():
+        console.print(Panel(
+            "[red]No AI provider available![/red]\n"
+            "Run [cyan]zai setup[/cyan] or start Ollama.",
+            border_style="red",
+        ))
+        return
+
+    try:
+        visible = [name for name in sorted(os.listdir(cwd)) if not name.startswith(".")]
+    except Exception:
+        visible = []
+    console.print(Panel(
+        f"[bold cyan]zai[/bold cyan] — AI coding assistant\n"
+        f"[dim]Folder: {cwd}[/dim]\n"
+        f"[dim]Files: {', '.join(visible[:12])}"
+        f"{'...' if len(visible) > 12 else ''}[/dim]\n"
+        "[dim]Ctrl+C to exit  |  /help for commands[/dim]",
+        border_style="cyan",
+    ))
+
+    preferred = model or load_config()["default_model"]
+    add_project(Path(cwd).name, cwd)
+    fire_hook("SessionStart", {"cwd": cwd, "files": visible})
+    if (Path(cwd) / "CLAUDE.md").exists():
+        console.print(f"[dim]CLAUDE.md loaded from {cwd}[/dim]")
+    if (Path.home() / ".zai" / "CLAUDE.md").exists():
+        console.print("[dim]Global CLAUDE.md loaded (~/.zai/CLAUDE.md)[/dim]")
+
+    loaded_plugins = plugin_loader.load_all()
+    if loaded_plugins:
+        console.print(f"[dim]Plugins: {', '.join(loaded_plugins.keys())}[/dim]")
+    for filename, error in plugin_loader.get_errors().items():
+        console.print(f"[yellow]Plugin error ({filename}): {error}[/yellow]")
+    _connect_mcp_servers()
+
+    def available_commands():
+        return (
+            CORE_COMMANDS
+            | set(load_commands().keys())
+            | set(plugin_loader.get_all_commands().keys())
+        )
+
+    def saved_session_names():
+        return [session["name"] for session in list_sessions()]
+
+    terminal_input = InteractiveInput(
+        cwd,
+        command_provider=available_commands,
+        model_provider=lambda: get_models().keys(),
+        session_provider=saved_session_names,
+    )
+    if terminal_input.enhanced:
+        console.print(
+            "[dim]Input history and completion enabled. "
+            "Alt+Enter or Ctrl+J inserts a newline.[/dim]"
+        )
+
+    session_context = ContextManager(model=preferred)
+    initial_history = [
+        Message(
+            role="user",
+            content=(
+                f"My current folder is: {cwd}\nFiles here: {', '.join(visible)}\n"
+                "Help me work on this project."
+            ),
+            pinned=True,
+        ),
+        Message(
+            role="assistant",
+            content=(
+                f"I can see your folder at {cwd} with these files: "
+                f"{', '.join(visible)}. I can read, create, and edit files directly. "
+                "What would you like to do?"
+            ),
+            pinned=True,
+        ),
+    ]
+    session_context.replace_messages(initial_history)
+    history = session_context.get_messages()
+
+    def remember_turn(user_content: str, assistant_content: str) -> None:
+        session_context.add("user", user_content)
+        session_context.add("assistant", assistant_content)
+        if session_context.is_near_limit():
+            session_context.compress(preferred)
+        save_auto_history(history, cwd)
+
+    while True:
+        try:
+            try:
+                message = terminal_input.prompt()
+            except KeyboardInterrupt:
+                console.print("[dim]Input cancelled.[/dim]")
+                continue
+            if not message.strip():
+                continue
+            message = _correct_slash_command(message)
+            stripped = message.strip()
+
+            if stripped == "/help":
+                _show_help()
+            elif stripped == "/clear":
+                console.clear()
+            elif stripped == "/files":
+                _show_files(cwd)
+            elif stripped.startswith("/model "):
+                name = stripped.split(" ", 1)[1].strip()
+                if name == "list":
+                    for key, data in get_models().items():
+                        key_status = (
+                            "[green]ok[/green]"
+                            if get_api_key(data["provider"])
+                            else "[red]no key[/red]"
+                        )
+                        console.print(
+                            f"  [cyan]{key}[/cyan] — {data['name']} ({key_status})"
+                        )
+                elif name in get_models():
+                    preferred = name
+                    session_context.set_model(name)
+                    console.print(f"[green]Switched to {name}[/green]")
+                else:
+                    console.print(f"[red]Unknown model: {name}[/red]")
+            elif stripped == "/memory":
+                last = get_last_session()
+                if last:
+                    console.print(f"[dim]{last['task']} — {last['date']}[/dim]")
+            elif stripped == "/undo":
+                console.print(f"[cyan]{undo_last(cwd)}[/cyan]")
+            elif stripped == "/diff":
+                _show_diff()
+            elif stripped.startswith("/plan"):
+                task = stripped[5:].strip()
+                if not task:
+                    console.print("[red]Usage: /plan <what you want to build>[/red]")
+                else:
+                    console.print("[dim]Planning...[/dim]")
+                    plan = plan_agent(task, history, cwd, preferred)
+                    if plan and Prompt.ask("[cyan]Proceed?[/cyan]", default="no").lower() in ("yes", "y"):
+                        console.print("[dim]Executing...[/dim]")
+                        result = run_agent(task, history, cwd, preferred)
+                        if result:
+                            remember_turn(task, result)
+                            save_session(task=task[:80], model=preferred)
+            elif stripped == "/test":
+                for round_number in range(3):
+                    console.print(
+                        f"[cyan]Running tests (round {round_number + 1}/3)...[/cyan]"
+                    )
+                    result = subprocess.run(
+                        ["python", "-m", "pytest", "--tb=short", "-q"],
+                        capture_output=True,
+                        text=True,
+                        cwd=cwd,
+                        timeout=120,
+                    )
+                    output = (result.stdout + result.stderr).strip()
+                    if not output:
+                        console.print("[yellow]No tests found.[/yellow]")
+                        break
+                    console.print(Panel(
+                        output[:2000],
+                        title="[cyan]Test Output[/cyan]",
+                        border_style="cyan",
+                    ))
+                    if result.returncode == 0:
+                        console.print("[green]All tests passing![/green]")
+                        break
+                    if round_number >= 2:
+                        console.print("[yellow]Max rounds reached. Tests still failing.[/yellow]")
+                        break
+                    console.print("[yellow]Asking AI to fix...[/yellow]")
+                    fix_prompt = (
+                        f"Tests are failing:\n\n{output}\n\n"
+                        "Read the failing files and fix the issues."
+                    )
+                    fix_result = run_agent(fix_prompt, history, cwd, preferred)
+                    if fix_result:
+                        remember_turn(fix_prompt, fix_result)
+            elif stripped.startswith("/resume"):
+                parts = stripped.split(None, 1)
+                name = parts[1].strip() if len(parts) > 1 else None
+                latest = None if name else load_latest_session(cwd)
+                loaded = load_history(name) if name else (
+                    latest[0] if latest else None
+                )
+                resumed_name = name or (
+                    latest[1]["name"] if latest else "latest"
+                )
+                if loaded:
+                    session_context.replace_messages(loaded)
+                    console.print(
+                        f"[green]Session resumed ({resumed_name}):[/green] "
+                        f"{len(loaded)} messages"
+                    )
+                elif name:
+                    console.print(f"[red]Session not found: {name}[/red]")
+                else:
+                    console.print(
+                        "[yellow]No saved sessions. "
+                        "Use /session save [name] first.[/yellow]"
+                    )
+            elif stripped.startswith("/session"):
+                parts = stripped.split()
+                action = parts[1] if len(parts) > 1 else "list"
+                arguments = parts[2:]
+                argument = " ".join(arguments) if arguments else None
+                if action == "save":
+                    path = save_history(history, argument, project_path=cwd)
+                    console.print(f"[green]Session saved:[/green] {path}")
+                elif action == "load":
+                    if not argument:
+                        console.print("[red]Usage: /session load <name>[/red]")
+                    else:
+                        loaded = load_history(argument)
+                        if loaded:
+                            session_context.replace_messages(loaded)
+                            console.print(
+                                f"[green]Session loaded:[/green] {len(loaded)} messages"
+                            )
+                        else:
+                            console.print(f"[red]Session not found: {argument}[/red]")
+                elif action == "search":
+                    if not argument:
+                        console.print("[red]Usage: /session search <query>[/red]")
+                    else:
+                        matches = search_sessions(argument, project_path=cwd)
+                        if not matches:
+                            console.print("[dim]No matching sessions.[/dim]")
+                        else:
+                            table = Table(
+                                title="Session Search",
+                                border_style="cyan",
+                            )
+                            table.add_column("Name", style="cyan")
+                            table.add_column("Title")
+                            table.add_column("Messages", justify="right")
+                            for session in matches:
+                                table.add_row(
+                                    session["name"],
+                                    session["title"],
+                                    str(session["messages"]),
+                                )
+                            console.print(table)
+                elif action == "rename":
+                    if len(arguments) != 2:
+                        console.print(
+                            "[red]Usage: /session rename <old> <new>[/red]"
+                        )
+                    else:
+                        renamed = rename_session(arguments[0], arguments[1])
+                        if renamed:
+                            console.print(
+                                f"[green]Session renamed:[/green] {arguments[1]}"
+                            )
+                        else:
+                            console.print(
+                                "[red]Rename failed: session missing, ambiguous, "
+                                "automatic, or destination exists.[/red]"
+                            )
+                elif action == "delete":
+                    if len(arguments) != 1:
+                        console.print("[red]Usage: /session delete <name>[/red]")
+                    elif Confirm.ask(
+                        f"Delete session '{arguments[0]}'?",
+                        default=False,
+                    ) and delete_session(arguments[0]):
+                        console.print(
+                            f"[green]Session deleted:[/green] {arguments[0]}"
+                        )
+                    else:
+                        console.print("[yellow]Session was not deleted.[/yellow]")
+                elif action == "list":
+                    sessions = list_sessions(project_path=cwd)
+                    if not sessions:
+                        console.print("[dim]No saved sessions.[/dim]")
+                    else:
+                        table = Table(title="Saved Sessions", border_style="cyan")
+                        table.add_column("ID", style="dim")
+                        table.add_column("Name", style="cyan")
+                        table.add_column("Title")
+                        table.add_column("Messages", justify="right")
+                        for session in sessions:
+                            table.add_row(
+                                session["id"][:8],
+                                session["name"],
+                                session["title"],
+                                str(session["messages"]),
+                            )
+                        console.print(table)
+                else:
+                    console.print(
+                        "[red]Usage: /session list | save [name] | load <name> | "
+                        "search <query> | rename <old> <new> | delete <name>[/red]"
+                    )
+            elif stripped == "/watch":
+                if is_watching():
+                    stop_watch()
+                    console.print("[yellow]Watch mode off.[/yellow]")
+                else:
+                    def on_file_change(filename):
+                        prompt = (
+                            f"File changed: {filename}. "
+                            "Read it and check for issues or errors."
+                        )
+                        result = run_agent(prompt, history, cwd, preferred)
+                        if result:
+                            remember_turn(prompt, result)
+
+                    if start_watch(cwd, on_file_change):
+                        console.print(f"[green]Watch mode on.[/green] Watching: {cwd}")
+            elif stripped.startswith("/") and not stripped.startswith("//"):
+                parts = stripped[1:].split(" ", 1)
+                command_name = parts[0].lower()
+                arguments = parts[1] if len(parts) > 1 else ""
+                prompt = get_command_prompt(command_name, arguments)
+                if not prompt:
+                    plugin_commands = plugin_loader.get_all_commands()
+                    if command_name in plugin_commands:
+                        prompt = plugin_commands[command_name].get("body", "").replace(
+                            "$ARGUMENTS",
+                            arguments,
+                        )
+                if prompt:
+                    console.print(f"[dim]Running /{command_name}...[/dim]")
+                    result = run_agent(prompt, history, cwd, preferred)
+                    if result:
+                        remember_turn(
+                            f"/{command_name} {arguments}".strip(),
+                            result,
+                        )
+                        save_session(task=f"/{command_name}", model=preferred)
+                elif command_name == "commands":
+                    for command in list_commands():
+                        console.print(
+                            f"  [cyan]/{command['name']}[/cyan] — "
+                            f"{command['description']}"
+                        )
+                    for name, data in plugin_loader.get_all_commands().items():
+                        console.print(
+                            f"  [blue]/{name}[/blue] — "
+                            f"{data.get('description', '')} [dim](plugin)[/dim]"
+                        )
+                else:
+                    console.print(f"[red]Unknown command: /{command_name}[/red]")
+                    console.print("[dim]Type /commands to see available commands[/dim]")
+            else:
+                fire_hook("UserPromptSubmit", {"message": stripped[:200]})
+                result = run_agent(stripped, history, cwd, preferred)
+                if result:
+                    remember_turn(stripped, result)
+                    save_session(task=stripped[:80], model=preferred)
+        except KeyboardInterrupt:
+            console.print("\n[yellow]Current operation cancelled.[/yellow]")
+            continue
+        except EOFError:
+            fire_hook("SessionEnd", {"cwd": cwd, "messages": len(history)})
+            console.print("\n[dim]Goodbye![/dim]")
+            break