zai-cli 0.1.0__py3-none-any.whl

zai/mcp/manager.py

@@ -0,0 +1,118 @@
+import os
+from pathlib import Path
+
+from ..config import ZAI_DIR
+from ..core.storage import atomic_write_json, read_json, update_json
+
+MCP_FILE = ZAI_DIR / "mcp.json"
+
+KNOWN_SERVERS = {
+    "filesystem": {
+        "description": "Secure file operations within the current project",
+        "command": "npx",
+        "args": ["-y", "@modelcontextprotocol/server-filesystem", "{cwd}"],
+        "env_key": None,
+    },
+    "memory": {
+        "description": "Knowledge-graph persistent memory",
+        "command": "npx",
+        "args": ["-y", "@modelcontextprotocol/server-memory"],
+        "env_key": None,
+    },
+    "github": {
+        "description": "GitHub repositories, issues, and pull requests (archived reference server)",
+        "command": "npx",
+        "args": ["-y", "@modelcontextprotocol/server-github"],
+        "env_key": "GITHUB_PERSONAL_ACCESS_TOKEN",
+    },
+    "postgres": {
+        "description": "Read-only PostgreSQL schema and query access (archived reference server)",
+        "command": "npx",
+        "args": ["-y", "@modelcontextprotocol/server-postgres", "{DATABASE_URL}"],
+        "env_key": "DATABASE_URL",
+    },
+}
+
+
+def load_mcp() -> list[dict]:
+    return read_json(MCP_FILE, [], expected_type=list)
+
+
+def save_mcp(servers: list[dict]) -> None:
+    atomic_write_json(MCP_FILE, servers)
+
+
+def add_server(name: str) -> dict | None:
+    if name not in KNOWN_SERVERS:
+        return None
+    server = {"name": name, **KNOWN_SERVERS[name]}
+    already = False
+
+    def update(servers):
+        nonlocal already
+        if any(item.get("name") == name for item in servers):
+            already = True
+            return servers
+        return [*servers, server]
+
+    update_json(MCP_FILE, [], update, expected_type=list)
+    return {"already": True, **KNOWN_SERVERS[name]} if already else server
+
+
+def remove_server(name: str) -> bool:
+    removed = False
+
+    def update(servers):
+        nonlocal removed
+        remaining = [server for server in servers if server.get("name") != name]
+        removed = len(remaining) != len(servers)
+        return remaining
+
+    update_json(MCP_FILE, [], update, expected_type=list)
+    return removed
+
+
+def list_servers() -> list[dict]:
+    return load_mcp()
+
+
+def build_server_command(
+    server: dict,
+    cwd: str | Path | None = None,
+) -> tuple[list[str], dict[str, str]]:
+    """Build an argv/env pair from persisted MCP server configuration."""
+    current_dir = str(Path(cwd or os.getcwd()).resolve())
+    known = KNOWN_SERVERS.get(server.get("name", ""), {})
+    env_key = known.get("env_key") if known else server.get("env_key")
+    replacements = {"cwd": current_dir}
+    if env_key:
+        value = os.getenv(env_key)
+        if not value and env_key == "GITHUB_PERSONAL_ACCESS_TOKEN":
+            value = os.getenv("GITHUB_TOKEN")
+        if not value:
+            raise ValueError(f"set {env_key} to enable {server.get('name', 'server')}")
+        replacements[env_key] = value
+
+    command = server.get("command")
+    args = server.get("args")
+    if not command or not isinstance(args, list):
+        # Migrate older saved entries to the current known command first.
+        if known:
+            command = known.get("command")
+            args = known.get("args", [])
+        else:
+            package = server.get("package")
+            if not package:
+                raise ValueError("MCP server has no command configuration")
+            command = "npx"
+            args = ["-y", package]
+
+    rendered_args = [
+        str(argument).format_map(replacements)
+        for argument in args
+    ]
+    argv = [str(command), *rendered_args]
+    if os.name == "nt" and command in {"npx", "npm"}:
+        argv = ["cmd", "/c", *argv]
+    environment = {env_key: replacements[env_key]} if env_key else {}
+    return argv, environment

zai/plugins/__init__.py

@@ -0,0 +1,2 @@
+from .base import BasePlugin
+from .loader import load_all, get_loaded, get_errors, get_all_tools, get_all_skills, get_all_commands, get_agent_descriptions

zai/plugins/base.py

@@ -0,0 +1,49 @@
+from abc import ABC
+
+
+class BasePlugin(ABC):
+    name: str = ""
+    description: str = ""
+    version: str = "0.1.0"
+    author: str = ""
+    permissions: tuple[str, ...] = ()
+
+    def setup(self) -> None:
+        """Called once when plugin loads. Override for initialization."""
+        pass
+
+    def get_tools(self) -> dict:
+        """
+        Return {tool_name: callable} dict.
+        Each callable receives **kwargs from the agent and returns a str result.
+        """
+        return {}
+
+    def get_skills(self) -> dict:
+        """
+        Return {skill_name: {"description": str, "prompt": callable(ctx) -> str}} dict.
+        Skills become available via: zai skill <name> <file>
+        """
+        return {}
+
+    def get_commands(self) -> dict:
+        """
+        Return {cmd_name: {"description": str, "body": str}} dict.
+        Commands become available as /<cmd_name> in interactive mode.
+        """
+        return {}
+
+    def get_agent_description(self) -> str:
+        """System prompt snippet describing this plugin's tools."""
+        tools = self.get_tools()
+        if not tools:
+            return ""
+        lines = [f"\nPlugin '{self.name}' tools (use structured plugin_call):"]
+        for tool_name, func in tools.items():
+            doc = (func.__doc__ or "").strip().split("\n")[0][:80]
+            lines.append(
+                f'<tool_call>{{"name":"plugin_call","arguments":'
+                f'{{"plugin":"{self.name}","tool":"{tool_name}",'
+                f'"arguments":{{"arg":"value"}}}}}}</tool_call>  ({doc})'
+            )
+        return "\n".join(lines)

zai/plugins/loader.py

@@ -0,0 +1,404 @@
+"""
+Plugin loader — finds and loads plugins from:
+  1. ~/.zai/plugins/*.py  (user-installed)
+  2. pip packages with entry_point group 'zai.plugins'
+"""
+import importlib.util
+import hashlib
+import json
+import re
+from pathlib import Path
+from rich.console import Console
+from .base import BasePlugin
+from ..core.storage import atomic_write_json, atomic_write_text, read_json
+
+console = Console()
+
+PLUGIN_DIR = Path.home() / ".zai" / "plugins"
+DISABLED_FILE = Path.home() / ".zai" / "plugins_disabled.json"
+TRUST_FILE = Path.home() / ".zai" / "plugins_trusted.json"
+
+_loaded: dict[str, BasePlugin] = {}
+_load_errors: dict[str, str] = {}
+PLUGIN_NAME_PATTERN = re.compile(r"^[A-Za-z0-9][A-Za-z0-9_.-]{0,127}$")
+MANIFEST_SUFFIX = ".plugin.json"
+MANIFEST_VERSION = 1
+ALLOWED_PERMISSIONS = {
+    "project_read",
+    "project_write",
+    "network",
+    "subprocess",
+    "secrets",
+}
+_manifests: dict[str, dict] = {}
+
+
+def _load_disabled() -> set:
+    return set(read_json(DISABLED_FILE, [], expected_type=list))
+
+
+def _save_disabled(names: set):
+    atomic_write_json(DISABLED_FILE, sorted(names))
+
+
+def _load_trusted() -> dict[str, str]:
+    return read_json(TRUST_FILE, {}, expected_type=dict)
+
+
+def _save_trusted(entries: dict[str, str]) -> None:
+    atomic_write_json(TRUST_FILE, entries)
+
+
+def local_manifest_path(name: str) -> Path | None:
+    plugin_path = local_plugin_path(name)
+    if plugin_path is None:
+        return None
+    return PLUGIN_DIR / f"{name}{MANIFEST_SUFFIX}"
+
+
+def _validate_manifest(data: object, expected_name: str) -> dict:
+    if not isinstance(data, dict):
+        raise ValueError("plugin manifest must be a JSON object")
+    if data.get("manifest_version") != MANIFEST_VERSION:
+        raise ValueError(f"manifest_version must be {MANIFEST_VERSION}")
+    if data.get("name") != expected_name:
+        raise ValueError("plugin manifest name does not match plugin name")
+    permissions = data.get("permissions")
+    if not isinstance(permissions, list) or not all(
+        isinstance(item, str) for item in permissions
+    ):
+        raise ValueError("plugin permissions must be a list of strings")
+    unknown = sorted(set(permissions) - ALLOWED_PERMISSIONS)
+    if unknown:
+        raise ValueError("unknown plugin permissions: " + ", ".join(unknown))
+    normalized = dict(data)
+    normalized["permissions"] = sorted(set(permissions))
+    normalized.setdefault("description", "")
+    normalized.setdefault("version", "0.1.0")
+    normalized.setdefault("source", "local")
+    return normalized
+
+
+def _read_local_manifest(name: str) -> dict:
+    path = local_manifest_path(name)
+    if path is None or not path.is_file():
+        raise ValueError(
+            f"Missing {name}{MANIFEST_SUFFIX}. "
+            "Plugins require a reviewed manifest."
+        )
+    try:
+        data = json.loads(path.read_text(encoding="utf-8"))
+    except (OSError, UnicodeError, json.JSONDecodeError) as error:
+        raise ValueError(f"Invalid plugin manifest: {error}") from error
+    return _validate_manifest(data, name)
+
+
+def _file_fingerprint(path: Path, manifest: dict) -> str:
+    manifest_bytes = json.dumps(
+        manifest,
+        sort_keys=True,
+        separators=(",", ":"),
+    ).encode("utf-8")
+    return hashlib.sha256(path.read_bytes() + b"\0" + manifest_bytes).hexdigest()
+
+
+def _entrypoint_manifest(entry_point) -> dict:
+    distribution = getattr(entry_point, "dist", None)
+    if distribution is None:
+        raise ValueError("Plugin entry point has no distribution metadata")
+    for item in getattr(distribution, "files", None) or []:
+        if Path(str(item)).name == "zai-plugin.json":
+            path = distribution.locate_file(item)
+            try:
+                data = json.loads(Path(path).read_text(encoding="utf-8"))
+            except (OSError, UnicodeError, json.JSONDecodeError) as error:
+                raise ValueError(f"Invalid packaged plugin manifest: {error}") from error
+            return _validate_manifest(data, entry_point.name)
+    raise ValueError("Packaged plugin is missing zai-plugin.json")
+
+
+def _entrypoint_fingerprint(entry_point, manifest: dict) -> str:
+    distribution = getattr(entry_point, "dist", None)
+    dist_name = getattr(distribution, "name", "") or ""
+    version = getattr(distribution, "version", "") or ""
+    file_records = []
+    for item in getattr(distribution, "files", None) or []:
+        recorded_hash = getattr(item, "hash", None)
+        file_records.append(f"{item}:{recorded_hash or ''}")
+    return hashlib.sha256(
+        (
+            f"{entry_point.value}|{dist_name}|{version}|"
+            + "|".join(sorted(file_records))
+            + "|"
+            + json.dumps(manifest, sort_keys=True, separators=(",", ":"))
+        ).encode("utf-8")
+    ).hexdigest()
+
+
+def _entry_points():
+    from importlib.metadata import entry_points
+
+    return entry_points(group="zai.plugins")
+
+
+def local_plugin_path(name: str) -> Path | None:
+    if (
+        not PLUGIN_NAME_PATTERN.fullmatch(name)
+        or ".." in name
+        or "/" in name
+        or "\\" in name
+    ):
+        return None
+    return PLUGIN_DIR / f"{name}.py"
+
+
+def trust_plugin(name: str) -> bool:
+    """Trust the current code fingerprint without importing or executing it."""
+    trusted = _load_trusted()
+    local_path = local_plugin_path(name)
+    if local_path and local_path.is_file():
+        try:
+            manifest = _read_local_manifest(name)
+        except ValueError:
+            return False
+        trusted[f"file:{name}"] = _file_fingerprint(local_path, manifest)
+        _save_trusted(trusted)
+        _load_errors.pop(name, None)
+        return True
+    try:
+        for entry_point in _entry_points():
+            if entry_point.name == name:
+                manifest = _entrypoint_manifest(entry_point)
+                trusted[f"entrypoint:{name}"] = _entrypoint_fingerprint(
+                    entry_point,
+                    manifest,
+                )
+                _save_trusted(trusted)
+                _load_errors.pop(name, None)
+                return True
+    except Exception:
+        return False
+    return False
+
+
+def revoke_plugin_trust(name: str) -> bool:
+    trusted = _load_trusted()
+    removed = False
+    for key in (f"file:{name}", f"entrypoint:{name}"):
+        removed = trusted.pop(key, None) is not None or removed
+    if removed:
+        _save_trusted(trusted)
+    _loaded.pop(name, None)
+    return removed
+
+
+def load_all() -> dict[str, BasePlugin]:
+    """Load all plugins. Returns {name: plugin_instance}."""
+    global _loaded, _load_errors, _manifests
+    _loaded = {}
+    _load_errors = {}
+    _manifests = {}
+    disabled = _load_disabled()
+    trusted = _load_trusted()
+
+    # 1. User plugins from ~/.zai/plugins/*.py
+    PLUGIN_DIR.mkdir(parents=True, exist_ok=True)
+    for py_file in sorted(PLUGIN_DIR.glob("*.py")):
+        try:
+            manifest = _read_local_manifest(py_file.stem)
+        except ValueError as error:
+            _load_errors[py_file.stem] = str(error)
+            continue
+        fingerprint = _file_fingerprint(py_file, manifest)
+        if trusted.get(f"file:{py_file.stem}") != fingerprint:
+            _load_errors[py_file.stem] = (
+                "Plugin is untrusted or changed. Review it, then run: "
+                f"zai plugin trust {py_file.stem}"
+            )
+            continue
+        _load_file(py_file, disabled, manifest)
+
+    # 2. Pip-installed entry points
+    try:
+        for ep in _entry_points():
+            try:
+                manifest = _entrypoint_manifest(ep)
+                if trusted.get(f"entrypoint:{ep.name}") != _entrypoint_fingerprint(
+                    ep,
+                    manifest,
+                ):
+                    _load_errors[ep.name] = (
+                        "Plugin package is untrusted or changed. Review it, then run: "
+                        f"zai plugin trust {ep.name}"
+                    )
+                    continue
+                plugin_cls = ep.load()
+                instance = plugin_cls()
+                if instance.name in disabled:
+                    continue
+                if tuple(manifest["permissions"]) != tuple(
+                    sorted(set(instance.permissions))
+                ):
+                    raise ValueError(
+                        "plugin class permissions do not match its manifest"
+                    )
+                instance.setup()
+                _loaded[instance.name] = instance
+                _manifests[instance.name] = manifest
+            except Exception as e:
+                _load_errors[ep.name] = str(e)
+    except Exception:
+        pass
+
+    return _loaded
+
+
+def _load_file(py_file: Path, disabled: set, manifest: dict):
+    try:
+        spec = importlib.util.spec_from_file_location(py_file.stem, py_file)
+        if spec is None or spec.loader is None:
+            raise ValueError("Cannot create plugin module specification")
+        mod = importlib.util.module_from_spec(spec)
+        spec.loader.exec_module(mod)
+        instance = getattr(mod, "plugin", None)
+        if not isinstance(instance, BasePlugin):
+            _load_errors[py_file.stem] = "No `plugin = MyPlugin()` found at module level"
+            return
+        if not instance.name:
+            instance.name = py_file.stem
+        if instance.name != manifest["name"]:
+            raise ValueError("loaded plugin name does not match its manifest")
+        if tuple(manifest["permissions"]) != tuple(
+            sorted(set(instance.permissions))
+        ):
+            raise ValueError("plugin class permissions do not match its manifest")
+        if instance.name in disabled:
+            return
+        instance.setup()
+        _loaded[instance.name] = instance
+        _manifests[instance.name] = manifest
+    except Exception as e:
+        _load_errors[py_file.stem] = str(e)
+        console.print(f"[yellow]Plugin error ({py_file.name}):[/yellow] {e}")
+
+
+def get_loaded() -> dict[str, BasePlugin]:
+    return _loaded
+
+
+def get_errors() -> dict[str, str]:
+    return _load_errors
+
+
+def get_manifest(name: str) -> dict | None:
+    manifest = _manifests.get(name)
+    return dict(manifest) if manifest else None
+
+
+def disable_plugin(name: str) -> bool:
+    disabled = _load_disabled()
+    disabled.add(name)
+    _save_disabled(disabled)
+    _loaded.pop(name, None)
+    return True
+
+
+def enable_plugin(name: str) -> bool:
+    disabled = _load_disabled()
+    if name not in disabled:
+        return False
+    disabled.discard(name)
+    _save_disabled(disabled)
+    return True
+
+
+def get_all_tools() -> dict[str, tuple]:
+    """Returns {tool_name: (plugin_name, callable)}."""
+    tools = {}
+    for plugin_name, plugin in _loaded.items():
+        for tool_name, func in plugin.get_tools().items():
+            tools[f"{plugin_name}.{tool_name}"] = (plugin_name, tool_name, func)
+    return tools
+
+
+def get_all_skills() -> dict:
+    """Merge all plugin skills into one dict."""
+    skills = {}
+    for plugin in _loaded.values():
+        skills.update(plugin.get_skills())
+    return skills
+
+
+def get_all_commands() -> dict:
+    """Merge all plugin commands into one dict."""
+    cmds = {}
+    for plugin in _loaded.values():
+        cmds.update(plugin.get_commands())
+    return cmds
+
+
+def get_agent_descriptions() -> str:
+    """Full system prompt block for all plugin tools."""
+    parts = [p.get_agent_description() for p in _loaded.values()]
+    return "\n".join(p for p in parts if p)
+
+
+def scaffold_plugin(name: str) -> str:
+    """Create a starter plugin file in ~/.zai/plugins/. Returns file path."""
+    PLUGIN_DIR.mkdir(parents=True, exist_ok=True)
+    safe = re.sub(r"[^a-z0-9_]", "_", name.lower().replace("-", "_")).strip("_")
+    if not safe:
+        raise ValueError("Plugin name must contain a letter or number")
+    path = PLUGIN_DIR / f"{safe}.py"
+    manifest_path = PLUGIN_DIR / f"{safe}{MANIFEST_SUFFIX}"
+    atomic_write_text(
+        path,
+        f'"""zai plugin: {name}"""\n'
+        f"from zai.plugins.base import BasePlugin\n\n\n"
+        f"class {safe.title().replace('_', '')}Plugin(BasePlugin):\n"
+        f'    name = "{safe}"\n'
+        f'    description = "My {name} plugin"\n'
+        f'    version = "0.1.0"\n'
+        f'    author = ""\n\n'
+        f"    permissions = ()\n\n"
+        f"    def setup(self):\n"
+        f"        pass  # runs once on load\n\n"
+        f"    def get_tools(self):\n"
+        f"        return {{\n"
+        f'            "hello": self.hello,\n'
+        f"        }}\n\n"
+        f"    def hello(self, name=\"world\", **kwargs):\n"
+        f'        """Say hello from the plugin."""\n'
+        f'        return f"Hello {{name}} from {name} plugin!"\n\n'
+        f"    def get_skills(self):\n"
+        f"        return {{\n"
+        f'            "{safe}_skill": {{\n'
+        f'                "description": "Example skill",\n'
+        f'                "prompt": lambda ctx: f"Process this:\\n\\n{{ctx}}",\n'
+        f"            }}\n"
+        f"        }}\n\n"
+        f"    def get_commands(self):\n"
+        f"        return {{\n"
+        f'            "{safe}": {{\n'
+        f'                "description": "Run {name} command",\n'
+        f'                "body": "You are helping with {name}. $ARGUMENTS",\n'
+        f"            }}\n"
+        f"        }}\n\n\n"
+        f"plugin = {safe.title().replace('_', '')}Plugin()\n",
+        mode=0o644,
+        lock=False,
+    )
+    atomic_write_json(
+        manifest_path,
+        {
+            "manifest_version": MANIFEST_VERSION,
+            "name": safe,
+            "description": f"My {name} plugin",
+            "version": "0.1.0",
+            "source": "local",
+            "permissions": [],
+        },
+        lock=False,
+    )
+    trust_plugin(safe)
+    return str(path)

zai/providers/__init__.py

@@ -0,0 +1,22 @@
+from .gemini import GeminiProvider
+from .groq import GroqProvider
+from .cerebras import CerebrasProvider
+from .openrouter import OpenRouterProvider
+from .qwen import QwenProvider
+from .anthropic import AnthropicProvider
+from .openai import OpenAIProvider
+from .ollama import OllamaProvider
+
+PROVIDERS = {
+    "gemini": GeminiProvider,
+    "groq": GroqProvider,
+    "cerebras": CerebrasProvider,
+    "openrouter": OpenRouterProvider,
+    "qwen": QwenProvider,
+    "anthropic": AnthropicProvider,
+    "openai": OpenAIProvider,
+    "ollama": OllamaProvider,
+    # aliases matching MODELS keys in config.py
+    "claude": AnthropicProvider,
+    "gpt4o": OpenAIProvider,
+}