zai-cli 0.1.0__py3-none-any.whl

zai/main.py

@@ -0,0 +1,210 @@
+import typer
+from typing import Optional
+from typer.core import TyperGroup
+from rich.console import Console
+from rich.panel import Panel
+from rich.prompt import Prompt, Confirm
+
+from .config import load_config
+from .core.fallback import (
+    format_model_selection,
+    has_available_provider,
+    stream_with_fallback,
+)
+from .core.context import ContextManager
+from .core.memory import save_session, get_last_session
+from .core.errors import show_error, AllModelsFailedError
+from .core.runtime import configure as configure_runtime
+from .core.runtime import plain_enabled, print_exception
+from .cli.settings import register_settings_commands
+from .cli.integrations import register_integration_commands
+from .cli.utilities import register_utility_commands
+from .cli.workflows import register_workflow_commands
+from .cli.common import closest_command as _closest_command
+from .cli.interactive import run_interactive
+
+
+class FuzzyTyperGroup(TyperGroup):
+    """Accept safe, unambiguous typos in top-level command names."""
+
+    def get_command(self, ctx, cmd_name):
+        command = super().get_command(ctx, cmd_name)
+        if command is not None:
+            return command
+
+        corrected = _closest_command(cmd_name, self.list_commands(ctx))
+        if corrected:
+            console.print(f"[yellow]Command '{cmd_name}' not found; using '{corrected}'.[/yellow]")
+            return super().get_command(ctx, corrected)
+        return None
+
+
+app = typer.Typer(
+    help="zai - Your personal AI CLI. Free. Fast. Smart.",
+    cls=FuzzyTyperGroup,
+)
+register_settings_commands(app)
+register_integration_commands(app)
+import io, sys
+if sys.stdout.encoding and sys.stdout.encoding.lower() != 'utf-8':
+    sys.stdout = io.TextIOWrapper(sys.stdout.buffer, encoding='utf-8', errors='replace')
+    sys.stderr = io.TextIOWrapper(sys.stderr.buffer, encoding='utf-8', errors='replace')
+console = Console()
+
+context = ContextManager()
+SYSTEM_PROMPT = (
+    "You are zai, a smart AI assistant running in the terminal. Be concise, helpful, and direct. "
+    "IMPORTANT: When creating files, ALWAYS use this exact format for each file:\n"
+    "**File:** filename.ext\n"
+    "```lang\n"
+    "...code...\n"
+    "```\n"
+    "Use this format every time so files can be automatically saved to disk."
+)
+register_utility_commands(app, context, SYSTEM_PROMPT)
+register_workflow_commands(app, context, SYSTEM_PROMPT)
+
+
+def _parse_files_from_response(content: str) -> dict:
+    """Extract filename→code pairs from AI response."""
+    import re
+    files = {}
+
+    # Pattern 1: **Anything File:** filename.ext  (e.g. **HTML File:** index.html)
+    pattern1 = r'\*\*[\w\s]*[Ff]ile[\w\s]*:\*\*\s*([\w./\\-]+\.\w+)[^\n]*\n```[\w]*\n(.*?)```'
+    for m in re.finditer(pattern1, content, re.DOTALL):
+        files[m.group(1).strip()] = m.group(2).strip()
+
+    # Pattern 2: **File Name:** or **Filename:** variant
+    pattern2 = r'\*\*[\w\s]*:\*\*\s*([\w./\\-]+\.\w+)\s*\n```[\w]*\n(.*?)```'
+    for m in re.finditer(pattern2, content, re.DOTALL):
+        fname = m.group(1).strip()
+        if fname not in files:
+            files[fname] = m.group(2).strip()
+
+    # Pattern 3: ```lang\n# filename or // filename or <!-- filename -->
+    pattern3 = r'```(\w+)\n(?:#|//|<!--)\s*([\w./\\-]+\.\w+)[^\n]*\n(.*?)```'
+    for m in re.finditer(pattern3, content, re.DOTALL):
+        fname = m.group(2).strip()
+        if fname not in files:
+            files[fname] = m.group(3).strip()
+
+    # Pattern 4: filename.ext on its own line (or backtick-wrapped) before code block
+    pattern4 = r'(?:^|\n)`?([\w./\\-]+\.\w+)`?\s*\n```[\w]*\n(.*?)```'
+    for m in re.finditer(pattern4, content, re.DOTALL):
+        fname = m.group(1).strip()
+        if fname not in files and '/' not in fname.split('.')[-1]:
+            files[fname] = m.group(2).strip()
+
+    return files
+
+
+def _chat(message: str, model: str = None) -> str:
+    config = load_config()
+    preferred = model or config["default_model"]
+    context.set_model(preferred)
+
+    if not has_available_provider():
+        show_error("No AI provider available. Run: zai setup, or start Ollama.")
+        return ""
+
+    context.add("user", message)
+
+    if context.is_near_limit():
+        console.print("[yellow]Context compressing...[/yellow]")
+        context.compress(preferred)
+
+    try:
+        content, used_model = stream_with_fallback(
+            context.get_messages(),
+            system=SYSTEM_PROMPT,
+            preferred=preferred,
+        )
+    except AllModelsFailedError as e:
+        show_error(str(e))
+        return ""
+    except Exception as e:
+        print_exception(console, e)
+        return ""
+
+    context.add("assistant", content)
+
+    if used_model != preferred and not plain_enabled():
+        console.print(f"[dim]Switched to: {format_model_selection(used_model)}[/dim]")
+
+    session_model = used_model
+    if plain_enabled():
+        save_session(task=message[:80], model=session_model)
+        return content
+
+    used_model = format_model_selection(used_model)
+
+    console.print(f"[dim]── {used_model} ──[/dim]")
+
+    if config.get("show_token_count") and not plain_enabled():
+        console.print(f"[dim]Tokens: {context.usage_bar()}[/dim]")
+
+    save_session(task=message[:80], model=session_model)
+    return content
+
+
+@app.command()
+def chat(
+    message: Optional[str] = typer.Argument(None, help="Message to send"),
+    model: Optional[str] = typer.Option(None, "--model", "-m", help="Model to use"),
+):
+    """Chat with AI. Supports auto model fallback."""
+    if message:
+        _chat(message, model)
+        return
+
+    last = get_last_session()
+    if last:
+        console.print(f"[dim]Last session: {last['task']} ({last['date']})[/dim]")
+        if Confirm.ask("Continue last session?", default=False):
+            context.add("assistant", f"Continuing from: {last['task']}")
+
+    console.print(Panel("[bold cyan]zai[/bold cyan] — type your message. [dim]Ctrl+C to exit.[/dim]"))
+    while True:
+        try:
+            msg = Prompt.ask("[cyan]you[/cyan]")
+            if msg.strip():
+                _chat(msg, model)
+        except (KeyboardInterrupt, EOFError):
+            console.print("\n[dim]Goodbye![/dim]")
+            break
+
+
+@app.command("ask")
+def ask(
+    message: str = typer.Argument(..., help="Message to send to AI"),
+    model: Optional[str] = typer.Option(None, "--model", "-m"),
+):
+    """Send a quick message to AI (shortcut for chat)."""
+    _chat(message, model)
+
+
+@app.callback(invoke_without_command=True)
+def main(
+    ctx: typer.Context,
+    model: Optional[str] = typer.Option(None, "--model", "-m"),
+    version: bool = typer.Option(False, "--version", "-v"),
+    debug: bool = typer.Option(False, "--debug", help="Show full tracebacks"),
+    plain: bool = typer.Option(
+        False,
+        "--plain",
+        help="Emit plain one-shot output without streaming UI or metadata",
+    ),
+):
+    """zai - Your personal AI CLI."""
+    configure_runtime(debug=debug, plain=plain)
+    if version:
+        from . import __version__
+        console.print(f"zai v{__version__}")
+        return
+    if ctx.invoked_subcommand is None:
+        run_interactive(model)
+
+
+if __name__ == "__main__":
+    app()

zai/mcp/client.py

@@ -0,0 +1,431 @@
+"""MCP JSON-RPC 2.0 client over the standard stdio transport."""
+from __future__ import annotations
+
+import atexit
+import json
+import os
+import shlex
+import subprocess
+import threading
+import time
+from dataclasses import dataclass, field
+from typing import Callable, Optional
+
+from rich.console import Console
+from ..core.cancellation import current_token
+
+console = Console()
+
+PROTOCOL_VERSION = "2025-06-18"
+SUPPORTED_PROTOCOL_VERSIONS = {PROTOCOL_VERSION, "2024-11-05"}
+DEFAULT_REQUEST_TIMEOUT = 30.0
+
+_active_clients: dict[str, "MCPClient"] = {}
+
+
+class MCPError(RuntimeError):
+    pass
+
+
+class MCPTimeoutError(MCPError):
+    pass
+
+
+@dataclass
+class _PendingRequest:
+    event: threading.Event = field(default_factory=threading.Event)
+    response: Optional[dict] = None
+
+
+class MCPClient:
+    """Connect to one MCP server process using newline-delimited JSON-RPC."""
+
+    def __init__(
+        self,
+        name: str,
+        command: str | list[str],
+        env: dict | None = None,
+        request_timeout: float = DEFAULT_REQUEST_TIMEOUT,
+        progress_callback: Callable[[dict], None] | None = None,
+    ):
+        self.name = name
+        self.command = command
+        self.env = env or {}
+        self.request_timeout = request_timeout
+        self.progress_callback = progress_callback
+        self.process: Optional[subprocess.Popen] = None
+        self._id = 0
+        self._tools: list[dict] = []
+        self._pending: dict[int, _PendingRequest] = {}
+        self._pending_lock = threading.Lock()
+        self._write_lock = threading.Lock()
+        self._reader_thread: threading.Thread | None = None
+        self._stderr_thread: threading.Thread | None = None
+        self._reader_error: str | None = None
+        self.server_info: dict = {}
+        self.server_capabilities: dict = {}
+        self.protocol_version: str | None = None
+        self.ready = False
+
+    def _argv(self) -> list[str]:
+        if isinstance(self.command, list):
+            argv = list(self.command)
+        else:
+            argv = shlex.split(self.command, posix=os.name != "nt")
+        if (
+            os.name == "nt"
+            and argv
+            and argv[0].lower() in {"npx", "npx.cmd", "npm", "npm.cmd"}
+        ):
+            return ["cmd", "/c", *argv]
+        return argv
+
+    def start(self) -> bool:
+        env = {**os.environ, **self.env}
+        try:
+            self.process = subprocess.Popen(
+                self._argv(),
+                shell=False,
+                stdin=subprocess.PIPE,
+                stdout=subprocess.PIPE,
+                stderr=subprocess.PIPE,
+                text=True,
+                encoding="utf-8",
+                errors="replace",
+                bufsize=1,
+                env=env,
+            )
+            self._reader_thread = threading.Thread(
+                target=self._read_loop,
+                name=f"zai-mcp-{self.name}",
+                daemon=True,
+            )
+            self._reader_thread.start()
+            self._stderr_thread = threading.Thread(
+                target=self._drain_stderr,
+                name=f"zai-mcp-stderr-{self.name}",
+                daemon=True,
+            )
+            self._stderr_thread.start()
+
+            response = self._request("initialize", {
+                "protocolVersion": PROTOCOL_VERSION,
+                "capabilities": {},
+                "clientInfo": {"name": "zai", "version": "0.1.0"},
+            })
+            result = response.get("result", {})
+            selected = result.get("protocolVersion")
+            if selected not in SUPPORTED_PROTOCOL_VERSIONS:
+                raise MCPError(
+                    f"unsupported protocol version: {selected or 'missing'}"
+                )
+            self.protocol_version = selected
+            self.server_capabilities = result.get("capabilities", {})
+            self.server_info = result.get("serverInfo", {})
+            self._notify("notifications/initialized", {})
+
+            if "tools" in self.server_capabilities:
+                self._tools = self._fetch_tools()
+            self.ready = True
+            console.print(
+                f"[green]MCP connected:[/green] {self.name} "
+                f"({len(self._tools)} tools, {self.protocol_version})"
+            )
+            return True
+        except Exception as error:
+            console.print(f"[red]MCP error ({self.name}): {error}[/red]")
+            self.stop()
+            return False
+
+    def _next_id(self) -> int:
+        with self._pending_lock:
+            self._id += 1
+            return self._id
+
+    def _write(self, message: dict) -> None:
+        if not self.process or self.process.poll() is not None or not self.process.stdin:
+            raise MCPError("server process is not running")
+        payload = json.dumps(message, separators=(",", ":"))
+        with self._write_lock:
+            self.process.stdin.write(payload + "\n")
+            self.process.stdin.flush()
+
+    def _notify(self, method: str, params: dict | None = None) -> None:
+        self._write({
+            "jsonrpc": "2.0",
+            "method": method,
+            "params": params or {},
+        })
+
+    def _request(
+        self,
+        method: str,
+        params: dict | None = None,
+        timeout: float | None = None,
+    ) -> dict:
+        request_id = self._next_id()
+        pending = _PendingRequest()
+        with self._pending_lock:
+            self._pending[request_id] = pending
+        try:
+            self._write({
+                "jsonrpc": "2.0",
+                "id": request_id,
+                "method": method,
+                "params": params or {},
+            })
+            wait_for = self.request_timeout if timeout is None else timeout
+            deadline = time.monotonic() + max(wait_for, 0.0)
+            while not pending.event.is_set():
+                token = current_token()
+                if token and token.cancelled:
+                    try:
+                        self._notify("notifications/cancelled", {
+                            "requestId": request_id,
+                            "reason": token.reason,
+                        })
+                    except MCPError:
+                        pass
+                    raise MCPError(token.reason)
+                remaining = deadline - time.monotonic()
+                if remaining <= 0:
+                    break
+                pending.event.wait(min(0.05, remaining))
+            if not pending.event.is_set():
+                try:
+                    self._notify("notifications/cancelled", {
+                        "requestId": request_id,
+                        "reason": f"zai timed out after {wait_for:g}s",
+                    })
+                except MCPError:
+                    pass
+                raise MCPTimeoutError(
+                    f"{method} timed out after {wait_for:g}s"
+                )
+            if pending.response is None:
+                raise MCPError(self._reader_error or "server connection closed")
+            return pending.response
+        finally:
+            with self._pending_lock:
+                self._pending.pop(request_id, None)
+
+    # Compatibility for callers/tests that used the old method name.
+    def _send(
+        self,
+        method: str,
+        params: dict | None = None,
+        timeout: float | None = None,
+    ) -> dict:
+        return self._request(method, params, timeout)
+
+    def _read_loop(self) -> None:
+        try:
+            if not self.process or not self.process.stdout:
+                return
+            for line in self.process.stdout:
+                if not line.strip():
+                    continue
+                try:
+                    message = json.loads(line)
+                except json.JSONDecodeError:
+                    continue
+                if "id" in message and ("result" in message or "error" in message):
+                    with self._pending_lock:
+                        pending = self._pending.get(message["id"])
+                    if pending:
+                        pending.response = message
+                        pending.event.set()
+                    continue
+                self._handle_server_message(message)
+        except Exception as error:
+            self._reader_error = str(error)
+        finally:
+            self._fail_pending()
+
+    def _drain_stderr(self) -> None:
+        """Drain server logs so a full stderr pipe cannot block the process."""
+        try:
+            if not self.process or not self.process.stderr:
+                return
+            for _line in self.process.stderr:
+                pass
+        except Exception:
+            pass
+
+    def _handle_server_message(self, message: dict) -> None:
+        method = message.get("method")
+        params = message.get("params", {})
+        if method == "notifications/progress":
+            if self.progress_callback:
+                self.progress_callback(params)
+            return
+        if method == "notifications/tools/list_changed" and self.ready:
+            # Avoid making a blocking request on the reader thread.
+            threading.Thread(
+                target=self._refresh_tools,
+                name=f"zai-mcp-refresh-{self.name}",
+                daemon=True,
+            ).start()
+            return
+        if "id" in message and method:
+            try:
+                self._write({
+                    "jsonrpc": "2.0",
+                    "id": message["id"],
+                    "error": {
+                        "code": -32601,
+                        "message": f"Client method not supported: {method}",
+                    },
+                })
+            except MCPError:
+                pass
+
+    def _fail_pending(self) -> None:
+        with self._pending_lock:
+            pending_requests = list(self._pending.values())
+        for pending in pending_requests:
+            pending.event.set()
+
+    def _refresh_tools(self) -> None:
+        try:
+            self._tools = self._fetch_tools()
+        except MCPError:
+            pass
+
+    def _fetch_tools(self) -> list[dict]:
+        tools: list[dict] = []
+        cursor = None
+        while True:
+            params = {"cursor": cursor} if cursor else {}
+            response = self._request("tools/list", params)
+            if "error" in response:
+                raise MCPError(
+                    response["error"].get("message", "tools/list failed")
+                )
+            result = response.get("result", {})
+            tools.extend(result.get("tools", []))
+            cursor = result.get("nextCursor")
+            if not cursor:
+                return tools
+
+    def call_tool(
+        self,
+        tool_name: str,
+        arguments: dict,
+        timeout: float | None = None,
+    ) -> str:
+        if not self.ready:
+            return f"MCP server '{self.name}' not ready"
+        if "tools" not in self.server_capabilities:
+            return f"MCP server '{self.name}' does not advertise tools"
+        try:
+            response = self._request(
+                "tools/call",
+                {"name": tool_name, "arguments": arguments},
+                timeout=timeout,
+            )
+        except MCPTimeoutError as error:
+            return f"MCP timeout: {error}"
+        except MCPError as error:
+            return f"MCP error: {error}"
+        if "error" in response:
+            return f"MCP error: {response['error'].get('message', 'Unknown')}"
+        result = response.get("result", {})
+        content = result.get("content", [])
+        parts = [
+            item.get("text", "")
+            for item in content
+            if item.get("type") == "text"
+        ]
+        if result.get("isError"):
+            details = "\n".join(parts) or "Unknown error"
+            return f"MCP tool error: {details}"
+        if parts:
+            return "\n".join(parts)
+        structured = result.get("structuredContent")
+        return json.dumps(structured) if structured is not None else "Done"
+
+    def get_tools(self) -> list[dict]:
+        return list(self._tools)
+
+    def stop(self) -> None:
+        process = self.process
+        self.ready = False
+        if not process:
+            return
+        try:
+            if process.stdin and not process.stdin.closed:
+                process.stdin.close()
+            process.wait(timeout=2)
+        except subprocess.TimeoutExpired:
+            process.terminate()
+            try:
+                process.wait(timeout=2)
+            except subprocess.TimeoutExpired:
+                process.kill()
+                try:
+                    process.wait(timeout=2)
+                except subprocess.TimeoutExpired:
+                    pass
+        except Exception:
+            try:
+                process.kill()
+            except Exception:
+                pass
+        finally:
+            self.process = None
+            self._fail_pending()
+
+
+def connect(
+    name: str,
+    command: str | list[str],
+    env: dict | None = None,
+    request_timeout: float = DEFAULT_REQUEST_TIMEOUT,
+) -> bool:
+    """Start and register an MCP server connection."""
+    disconnect(name)
+    client = MCPClient(name, command, env, request_timeout=request_timeout)
+    if client.start():
+        _active_clients[name] = client
+        return True
+    return False
+
+
+def disconnect(name: str) -> None:
+    client = _active_clients.pop(name, None)
+    if client:
+        client.stop()
+
+
+def get_all_tools() -> dict[str, list[dict]]:
+    return {
+        name: client.get_tools()
+        for name, client in _active_clients.items()
+        if client.ready
+    }
+
+
+def call_mcp_tool(
+    server: str,
+    tool: str,
+    arguments: dict,
+    timeout: float | None = None,
+) -> str:
+    client = _active_clients.get(server)
+    if not client:
+        return f"MCP server '{server}' not connected. Run: zai mcp connect {server}"
+    return client.call_tool(tool, arguments, timeout=timeout)
+
+
+def stop_all() -> None:
+    for client in list(_active_clients.values()):
+        client.stop()
+    _active_clients.clear()
+
+
+def active_count() -> int:
+    return sum(1 for client in _active_clients.values() if client.ready)
+
+
+atexit.register(stop_all)