zai-cli 0.1.0__py3-none-any.whl

zai/core/fallback.py

@@ -0,0 +1,171 @@
+from rich.console import Console
+from ..providers import PROVIDERS
+from ..providers.base import Message, Response
+from ..config import get_model_config, load_config
+from .errors import (
+    AllModelsFailedError,
+    NetworkError,
+    NoAPIKeyError,
+    ProviderError,
+    RateLimitError,
+    classify_provider_error,
+)
+from .runtime import plain_enabled
+from .cancellation import OperationCancelled, raise_if_cancelled
+
+console = Console()
+
+
+def format_model_selection(name: str) -> str:
+    try:
+        model = get_model_config(name)
+    except KeyError:
+        return name
+    return f"{name} -> {model['provider']}/{model['model_id']}"
+
+
+def get_provider(name: str):
+    try:
+        model = get_model_config(name)
+    except KeyError:
+        return None
+    cls = PROVIDERS.get(model["provider"])
+    if not cls:
+        return None
+    provider = cls()
+    provider.model_id = model["model_id"]
+    provider.context_window = model["context_window"]
+    provider.timeout = model.get("timeout", 60)
+    provider.retries = model.get("retries", 2)
+    return provider
+
+
+def _model_order(config: dict, preferred: str = None) -> list[str]:
+    first = preferred or config["default_model"]
+    if not config.get("auto_fallback", True):
+        return [first]
+    return [first] + [
+        model for model in config["fallback_order"]
+        if model != first
+    ]
+
+
+def has_available_provider() -> bool:
+    """Return True when any configured API provider or local provider is usable."""
+    config = load_config()
+    for model_name in _model_order(config):
+        provider = get_provider(model_name)
+        if provider and provider.is_available():
+            return True
+    return False
+
+
+def stream_with_fallback(messages: list[Message], system: str = "", preferred: str = None) -> tuple[str, str]:
+    """Stream response with auto fallback. Returns (content, model_name)."""
+    config = load_config()
+    order = _model_order(config, preferred)
+
+    last_error = None
+    for model_name in order:
+        raise_if_cancelled()
+        provider = get_provider(model_name)
+        if not provider or not provider.is_available():
+            continue
+        try:
+            from .context import compact_messages, estimate_text_tokens
+
+            prepared_messages = compact_messages(
+                messages,
+                getattr(provider, "context_window", 128_000),
+                reserve_tokens=max(4096, estimate_text_tokens(system) + 4096),
+            )
+            if provider.supports_streaming and not plain_enabled():
+                content = provider.stream_chat(prepared_messages, system=system)
+            else:
+                response = provider.chat(prepared_messages, system=system)
+                content = response.content
+                console.print(content, markup=False)
+            raise_if_cancelled()
+            return content, model_name
+        except OperationCancelled:
+            raise
+        except RateLimitError as error:
+            last_error = str(error)
+            console.print(f"[yellow]⚡ {model_name} limit hit — switching...[/yellow]")
+            continue
+        except NoAPIKeyError:
+            continue
+        except (NetworkError, ProviderError) as error:
+            last_error = str(error)
+            console.print(f"[yellow]Network issue with {model_name} — switching...[/yellow]")
+            continue
+        except Exception as error:
+            normalized = classify_provider_error(provider.name, error)
+            console.print(f"[dim]Error with {model_name}: {normalized}[/dim]")
+            last_error = str(normalized)
+            continue
+
+    raise AllModelsFailedError(last_error or "")
+
+
+def chat_with_fallback(
+    messages: list[Message],
+    system: str = "",
+    preferred: str = None,
+    tools: list[dict] | None = None,
+) -> tuple[Response, str]:
+    config = load_config()
+    order = _model_order(config, preferred)
+
+    last_error = None
+    for model_name in order:
+        raise_if_cancelled()
+        provider = get_provider(model_name)
+        if not provider or not provider.is_available():
+            continue
+        try:
+            from .context import compact_messages, estimate_text_tokens
+
+            provider_system = system
+            provider_tools = tools if provider.supports_native_tools else None
+            if tools and not provider.supports_native_tools:
+                from .tool_schema import legacy_tool_instructions
+
+                provider_system = (
+                    f"{system}\n\n{legacy_tool_instructions()}"
+                )
+            tool_schema_tokens = estimate_text_tokens(str(provider_tools or ""))
+            prepared_messages = compact_messages(
+                messages,
+                getattr(provider, "context_window", 128_000),
+                reserve_tokens=max(
+                    4096,
+                    estimate_text_tokens(provider_system) + tool_schema_tokens + 4096,
+                ),
+            )
+            response = provider.chat(
+                prepared_messages,
+                system=provider_system,
+                tools=provider_tools,
+            )
+            raise_if_cancelled()
+            return response, model_name
+        except OperationCancelled:
+            raise
+        except RateLimitError as error:
+            console.print(f"[yellow]⚡ {model_name} limit hit — switching...[/yellow]")
+            last_error = str(error)
+            continue
+        except NoAPIKeyError:
+            continue
+        except (NetworkError, ProviderError) as error:
+            console.print(f"[yellow]Network issue with {model_name} — switching...[/yellow]")
+            last_error = str(error)
+            continue
+        except Exception as error:
+            normalized = classify_provider_error(provider.name, error)
+            console.print(f"[dim]Error with {model_name}: {normalized}[/dim]")
+            last_error = str(normalized)
+            continue
+
+    raise AllModelsFailedError(last_error or "")

zai/core/hooks.py

@@ -0,0 +1,115 @@
+import json
+import os
+from pathlib import Path
+from .process import run_direct
+from .storage import atomic_write_json, read_json, update_json
+
+ZAI_DIR = Path.home() / ".zai"
+HOOKS_FILE = ZAI_DIR / "hooks.json"
+
+VALID_EVENTS = [
+    "PreToolUse",
+    "PostToolUse",
+    "SessionStart",
+    "SessionEnd",
+    "UserPromptSubmit",
+]
+
+
+def load_hooks() -> list:
+    return read_json(HOOKS_FILE, [], expected_type=list)
+
+
+def save_hooks(hooks: list):
+    atomic_write_json(HOOKS_FILE, hooks)
+
+
+def add_hook(event: str, command: str) -> dict:
+    created = {}
+
+    def update(hooks):
+        nonlocal created
+        new_id = max((h["id"] for h in hooks), default=0) + 1
+        created = {"id": new_id, "event": event, "command": command}
+        return [*hooks, created]
+
+    update_json(HOOKS_FILE, [], update, expected_type=list)
+    hook = created
+    return hook
+
+
+def remove_hook(hook_id: int) -> bool:
+    removed = False
+
+    def update(hooks):
+        nonlocal removed
+        new_hooks = [hook for hook in hooks if hook["id"] != hook_id]
+        removed = len(new_hooks) != len(hooks)
+        return new_hooks
+
+    update_json(HOOKS_FILE, [], update, expected_type=list)
+    return removed
+
+
+def fire(event: str, data: dict = None) -> bool:
+    """
+    Fire an event. Returns False if any hook blocks (exit code 2).
+    Hook scripts receive JSON on stdin.
+    """
+    hooks = load_hooks()
+    payload = json.dumps(data or {})
+
+    for hook in hooks:
+        if hook.get("event") != event:
+            continue
+        cmd = hook.get("command", "")
+        if not cmd:
+            continue
+        try:
+            result = run_direct(
+                cmd,
+                input_text=payload,
+                timeout=10,
+                enforce_policy=False,
+            )
+            if result.returncode == 2:
+                if result.output:
+                    from rich.console import Console
+                    Console().print(f"[yellow]Hook blocked:[/yellow] {result.output}")
+                return False
+            elif result.returncode == 126 and result.blocked_reason:
+                from rich.console import Console
+                Console().print(
+                    f"[yellow]Hook blocked:[/yellow] {result.blocked_reason}"
+                )
+                return False
+            elif result.returncode == 1 and result.output:
+                from rich.console import Console
+                Console().print(f"[dim]Hook note: {result.output}[/dim]")
+        except Exception:
+            pass
+
+    return True
+
+
+# Built-in security checks for write_file tool
+SECURITY_PATTERNS = {
+    "eval(": "eval() executes arbitrary code — security risk",
+    "exec(": "exec() can run arbitrary code",
+    "os.system(": "os.system() can allow command injection",
+    "pickle.loads(": "pickle.loads() with untrusted data is dangerous",
+    ".innerHTML =": "innerHTML with untrusted content causes XSS",
+    "dangerouslySetInnerHTML": "dangerouslySetInnerHTML causes XSS — sanitize first",
+    "document.write(": "document.write() is an XSS risk",
+    "new Function(": "new Function() evaluates arbitrary code",
+    "child_process.exec(": "child_process.exec() allows shell injection",
+}
+
+
+def check_security(content: str) -> list[str]:
+    """Return list of security warnings found in content."""
+    warnings = []
+    for pattern, msg in SECURITY_PATTERNS.items():
+        if pattern in content:
+            warnings.append(msg)
+    return warnings

zai/core/memory.py

@@ -0,0 +1,57 @@
+from datetime import datetime
+from pathlib import Path
+from ..config import MEMORY_FILE
+from .storage import atomic_write_json, read_json, update_json
+
+DEFAULT_MEMORY = {"version": 1, "projects": [], "last_session": None, "preferences": {}}
+
+
+def load_memory() -> dict:
+    data = read_json(MEMORY_FILE, DEFAULT_MEMORY.copy(), expected_type=dict)
+    return {**DEFAULT_MEMORY, **data}
+
+
+def save_memory(memory: dict):
+    atomic_write_json(MEMORY_FILE, {**DEFAULT_MEMORY, **memory})
+
+
+def save_session(task: str, model: str, files: list[str] | None = None):
+    def update(memory):
+        memory = {**DEFAULT_MEMORY, **memory}
+        memory["last_session"] = {
+            "task": task,
+            "model": model,
+            "files_edited": files or [],
+            "date": datetime.now().strftime("%Y-%m-%d %H:%M"),
+        }
+        return memory
+
+    update_json(MEMORY_FILE, DEFAULT_MEMORY.copy(), update, expected_type=dict)
+
+
+def get_last_session() -> dict | None:
+    return load_memory().get("last_session")
+
+
+def add_project(name: str, path: str, summary: str = ""):
+    def update(memory):
+        memory = {**DEFAULT_MEMORY, **memory}
+        projects = memory.setdefault("projects", [])
+        for project in projects:
+            if project["path"] == path:
+                project["last_worked"] = datetime.now().strftime("%Y-%m-%d")
+                project["summary"] = summary or project.get("summary", "")
+                return memory
+        projects.append({
+            "name": name,
+            "path": path,
+            "summary": summary,
+            "last_worked": datetime.now().strftime("%Y-%m-%d"),
+        })
+        return memory
+
+    update_json(MEMORY_FILE, DEFAULT_MEMORY.copy(), update, expected_type=dict)
+
+
+def get_projects() -> list:
+    return load_memory().get("projects", [])

zai/core/process.py

@@ -0,0 +1,204 @@
+import os
+import re
+import shlex
+import subprocess
+import time
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Callable
+from .cancellation import CancellationToken, current_token
+
+
+SHELL_META = re.compile(r"(?:&&|\|\||[|;&<>`]|[\r\n]|\$\(|%\w+%)")
+SHELL_INTERPRETERS = {
+    "bash", "sh", "zsh", "fish", "cmd", "cmd.exe",
+    "powershell", "powershell.exe", "pwsh", "wsl",
+}
+PYTHON_NAMES = {"python", "python.exe", "python3", "python3.exe", "py", "py.exe"}
+PACKAGE_MANAGERS = {"pip", "pip3", "uv", "poetry", "npm", "pnpm", "yarn"}
+DELETE_TOOLS = {"rm", "del", "erase", "rmdir", "remove-item"}
+
+
+@dataclass
+class CommandResult:
+    returncode: int
+    stdout: str = ""
+    stderr: str = ""
+    cancelled: bool = False
+    blocked_reason: str = ""
+
+    @property
+    def output(self) -> str:
+        return (self.stdout + self.stderr).strip()
+
+
+def split_command(command: str) -> list[str]:
+    """Split a direct executable command without invoking a shell."""
+    command = command.strip()
+    if not command:
+        raise ValueError("empty command")
+    if SHELL_META.search(command):
+        raise ValueError("shell syntax is not allowed")
+    try:
+        parts = shlex.split(command, posix=os.name != "nt")
+    except ValueError as exc:
+        raise ValueError(f"invalid command quoting: {exc}") from exc
+    if os.name == "nt":
+        parts = [
+            part[1:-1] if len(part) >= 2 and part[0] == part[-1] == '"' else part
+            for part in parts
+        ]
+    if not parts:
+        raise ValueError("empty command")
+    return parts
+
+
+def _exe_name(argv: list[str]) -> str:
+    return Path(argv[0]).name.lower()
+
+
+def classify_argv(argv: list[str]) -> tuple[str, str]:
+    """Classify direct argv as safe, approval, or blocked."""
+    executable = _exe_name(argv)
+    lowered = [part.lower() for part in argv[1:]]
+
+    if executable in SHELL_INTERPRETERS:
+        return "blocked", "shell interpreters are not allowed"
+    if executable in DELETE_TOOLS:
+        return "approval", "file deletion"
+    if executable in {"format", "mkfs", "mkfs.ext4", "shutdown", "reboot", "halt"}:
+        return "blocked", "destructive system operation"
+
+    if executable == "git":
+        if not lowered:
+            return "safe", ""
+        action = lowered[0]
+        if action == "reset" and "--hard" in lowered:
+            return "blocked", "destructive git operation"
+        if action == "clean" and any(arg.startswith("-") and "f" in arg for arg in lowered):
+            return "blocked", "destructive git operation"
+        if action in {"status", "diff", "log", "branch", "show", "rev-parse", "--version"}:
+            return "safe", ""
+        return "approval", "git state change"
+
+    if executable in PYTHON_NAMES:
+        if lowered in (["--version"], ["-v"]):
+            return "safe", ""
+        if len(lowered) >= 2 and lowered[0] == "-m":
+            module = lowered[1]
+            if module in {"pytest", "py_compile", "compileall"}:
+                return "safe", ""
+            if module in {"pip", "uv"}:
+                return "approval", "dependency change"
+        if "-c" in lowered:
+            return "blocked", "inline interpreter execution"
+        return "approval", "Python program execution"
+
+    if executable in {"pytest", "pytest.exe"}:
+        return "safe", ""
+    if executable in PACKAGE_MANAGERS:
+        return "approval", "dependency change"
+    if executable in {"curl", "wget", "npx"}:
+        return "approval", "network or package execution"
+
+    return "blocked", f"executable is not allowlisted: {executable}"
+
+
+def run_direct(
+    command: str | list[str],
+    *,
+    cwd: str | None = None,
+    timeout: int = 30,
+    input_text: str | None = None,
+    approval: Callable[[str], bool] | None = None,
+    enforce_policy: bool = True,
+    cancellation_token: CancellationToken | None = None,
+) -> CommandResult:
+    """Run an executable directly with shell=False."""
+    try:
+        argv = split_command(command) if isinstance(command, str) else list(command)
+    except ValueError as exc:
+        return CommandResult(126, blocked_reason=str(exc))
+
+    if enforce_policy:
+        risk, reason = classify_argv(argv)
+        if risk == "blocked":
+            return CommandResult(126, blocked_reason=reason)
+        if risk == "approval" and (approval is None or not approval(reason)):
+            return CommandResult(125, cancelled=True)
+
+    token = cancellation_token or current_token()
+    if token and token.cancelled:
+        return CommandResult(125, cancelled=True)
+    if token is None:
+        try:
+            completed = subprocess.run(
+                argv,
+                shell=False,
+                capture_output=True,
+                text=True,
+                cwd=cwd,
+                timeout=timeout,
+                input=input_text,
+            )
+            return CommandResult(
+                completed.returncode,
+                completed.stdout or "",
+                completed.stderr or "",
+            )
+        except subprocess.TimeoutExpired:
+            return CommandResult(124, stderr=f"command timed out after {timeout}s")
+        except FileNotFoundError:
+            return CommandResult(127, stderr=f"executable not found: {argv[0]}")
+        except Exception as exc:
+            return CommandResult(1, stderr=str(exc))
+
+    process = None
+    try:
+        process = subprocess.Popen(
+            argv,
+            shell=False,
+            stdin=subprocess.PIPE if input_text is not None else None,
+            stdout=subprocess.PIPE,
+            stderr=subprocess.PIPE,
+            text=True,
+            cwd=cwd,
+        )
+        started = time.monotonic()
+        pending_input = input_text
+        while True:
+            if token and token.cancelled:
+                process.terminate()
+                try:
+                    process.communicate(timeout=2)
+                except subprocess.TimeoutExpired:
+                    process.kill()
+                    process.communicate()
+                return CommandResult(125, cancelled=True)
+            remaining = timeout - (time.monotonic() - started)
+            if remaining <= 0:
+                process.kill()
+                process.communicate()
+                return CommandResult(124, stderr=f"command timed out after {timeout}s")
+            try:
+                stdout, stderr = process.communicate(
+                    input=pending_input,
+                    timeout=min(0.1, remaining),
+                )
+                return CommandResult(
+                    process.returncode,
+                    stdout or "",
+                    stderr or "",
+                )
+            except subprocess.TimeoutExpired:
+                pending_input = None
+    except FileNotFoundError:
+        return CommandResult(127, stderr=f"executable not found: {argv[0]}")
+    except KeyboardInterrupt:
+        if token:
+            token.cancel("cancelled by user")
+        if process and process.poll() is None:
+            process.terminate()
+        raise
+    except Exception as exc:
+        return CommandResult(1, stderr=str(exc))