zai-coding-gateway 0.1.0__py3-none-any.whl

zai_coding_gateway/__init__.py

@@ -0,0 +1,3 @@
+"""ZAI Coding Gateway — MCP-сервер для решения задач в рабочем пространстве через Z.AI Coding Plan."""
+
+__version__ = "0.1.0"

zai_coding_gateway/client.py

@@ -0,0 +1,93 @@
+"""Обёртка над OpenAI-клиентом для Z.AI Coding Plan endpoint."""
+
+import os
+from typing import Any
+
+from openai import OpenAI
+
+from zai_coding_gateway.errors import QuotaExceededError
+
+
+def _is_quota_error(exc: Exception) -> bool:
+    """Определяет по исключению, что это исчерпание квоты (429 / quota / rate limit)."""
+    status = getattr(exc, "status_code", None)
+    if status == 429:
+        return True
+    msg = str(exc).lower()
+    if "429" in msg or "quota" in msg or "rate limit" in msg or "rate_limit" in msg:
+        return True
+    return False
+
+DEFAULT_BASE_URL = "https://api.z.ai/api/coding/paas/v4"
+DEFAULT_WORK_MODEL = "glm-4.7"
+DEFAULT_FAST_MODEL = "glm-4.5-air"
+
+
+def _get_env(key: str, default: str | None = None) -> str | None:
+    return os.environ.get(key) or default
+
+
+def get_effective_base_url() -> str:
+    """Возвращает URL, который будет использоваться для запросов (из env или дефолт). Для проверки биллинга."""
+    return (_get_env("ZAI_BASE_URL") or DEFAULT_BASE_URL).rstrip("/") + "/"
+
+
+def get_work_model() -> str:
+    """Модель для основного цикла сессии solve_in_workspace. По умолчанию glm-4.7."""
+    v = (_get_env("ZAI_WORK_MODEL") or "").strip()
+    return v or DEFAULT_WORK_MODEL
+
+
+def get_fast_model() -> str:
+    """Модель для консолидации и суммаризации в solve_in_workspace. По умолчанию glm-4.5-air."""
+    v = (_get_env("ZAI_FAST_MODEL") or "").strip()
+    return v or DEFAULT_FAST_MODEL
+
+
+def create_client() -> OpenAI:
+    """Создаёт OpenAI-клиент с base_url и api_key из ENV. Ключ только из окружения."""
+    api_key = _get_env("ZAI_API_KEY")
+    if not api_key:
+        raise ValueError("ZAI_API_KEY не задан. Укажите ключ в env (например в mcp.json).")
+    base_url = get_effective_base_url()
+    return OpenAI(api_key=api_key, base_url=base_url)
+
+
+def chat_completions(
+    client: OpenAI,
+    messages: list[dict[str, Any]],
+    model: str,
+    response_format: dict[str, Any] | None = None,
+) -> dict[str, Any]:
+    """
+    Вызов chat/completions. При исчерпании квоты выбрасывает QuotaExceededError.
+    model: передавать get_work_model() или get_fast_model() в зависимости от сценария.
+    """
+    kwargs: dict[str, Any] = {"model": model, "messages": messages, "stream": False}
+    if response_format is not None:
+        kwargs["response_format"] = response_format
+    try:
+        resp = client.chat.completions.create(**kwargs)
+    except QuotaExceededError:
+        raise
+    except Exception as e:  # noqa: BLE001
+        if _is_quota_error(e):
+            raise QuotaExceededError("QuotaExceeded: wait for refresh cycle") from e
+        raise
+    return _completion_to_dict(resp)
+
+
+def _completion_to_dict(resp: Any) -> dict[str, Any]:
+    """Преобразует объект Completion в dict для единообразного разбора."""
+    return {
+        "choices": [
+            {
+                "message": {
+                    "content": c.message.content,
+                    "role": getattr(c.message, "role", "assistant"),
+                }
+            }
+            for c in (resp.choices or [])
+        ],
+        "usage": getattr(resp, "usage", None),
+    }

zai_coding_gateway/errors.py

@@ -0,0 +1,63 @@
+"""Доменные исключения для ZAI Coding Gateway."""
+
+
+class ZAIGatewayError(Exception):
+    """Базовое исключение gateway."""
+
+    pass
+
+
+class QuotaExceededError(ZAIGatewayError):
+    """Исчерпана квота Coding Plan. Ожидать обновления цикла (~5 ч)."""
+
+    pass
+
+
+class InstructionTooLongError(ZAIGatewayError):
+    """Текстовый ввод (инструкция/задача/контекст) превышает допустимый лимит. Использовать file_paths для контекста."""
+
+    pass
+
+
+class FileOutsideProjectError(ZAIGatewayError):
+    """Путь к файлу вне корня проекта (ZAI_PROJECT_ROOT)."""
+
+    pass
+
+
+class ProjectRootNotSetError(ZAIGatewayError):
+    """
+    Не удалось определить корень проекта: ZAI_PROJECT_ROOT не задан,
+    MCP roots/list не поддерживается клиентом (например Cursor), cwd не подходит.
+    Задайте ZAI_PROJECT_ROOT в конфиге MCP (env) — путь к папке проекта.
+    """
+
+    pass
+
+
+class FilePathsRequiredError(ZAIGatewayError):
+    """Для задачи требуются file_paths; список пуст или не передан."""
+
+    pass
+
+
+class SuggestedChangesOutsideSessionError(ZAIGatewayError):
+    """Путь в suggested_changes не входит в множество файлов сессии."""
+
+    pass
+
+
+class DiffNotApplicableError(ZAIGatewayError):
+    """Дифф не применим к текущему содержимому файла."""
+
+    pass
+
+
+class SessionNotFoundError(ZAIGatewayError):
+    """Сессия workspace не найдена (уже очищена или неверный session_id)."""
+
+    pass
+
+
+# ~3000 токенов: ориентир ~4 символа на токен → 12000 символов
+MAX_INSTRUCTION_CHARS = 12_000

zai_coding_gateway/logging_utils.py

@@ -0,0 +1,32 @@
+"""Структурированное логирование в stderr (JSON, одна строка на событие)."""
+
+import json
+import sys
+from typing import Any, Literal
+
+ToolInvocationStatus = Literal["success", "quota_exceeded", "error"]
+
+
+def log_tool_invocation(
+    tool: str,
+    model: str,
+    latency_ms: float,
+    status: ToolInvocationStatus,
+    extra: dict[str, Any] | None = None,
+) -> None:
+    """
+    Пишет в stderr одну строку JSON: tool, model, latency, status.
+    stdout не трогаем — занят MCP при stdio.
+    """
+    payload: dict[str, Any] = {
+        "event": "tool_invocation",
+        "tool": tool,
+        "model": model,
+        "latency_ms": round(latency_ms, 2),
+        "status": status,
+    }
+    if extra:
+        payload.update(extra)
+    line = json.dumps(payload, ensure_ascii=False) + "\n"
+    sys.stderr.write(line)
+    sys.stderr.flush()

zai_coding_gateway/main.py

@@ -0,0 +1,38 @@
+"""Точка входа: парсинг аргументов, проверка ENV, запуск MCP по stdio."""
+
+import argparse
+import os
+import sys
+
+from zai_coding_gateway.server import get_mcp
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description="ZAI Coding Gateway — MCP-сервер для Z.AI Coding Plan")
+    parser.add_argument(
+        "--stdio",
+        action="store_true",
+        default=True,
+        help="Запуск с транспортом stdio (по умолчанию)",
+    )
+    args = parser.parse_args()
+
+    if not os.environ.get("ZAI_API_KEY"):
+        sys.stderr.write("Ошибка: ZAI_API_KEY не задан. Укажите ключ в env (например в mcp.json).\n")
+        sys.stderr.flush()
+        sys.exit(1)
+
+    from zai_coding_gateway.client import get_effective_base_url
+    effective = get_effective_base_url()
+    plan = "Coding Plan" if "/coding/" in effective else "общий paas (не Coding Plan)"
+    sys.stderr.write(f"zai-coding-gateway: endpoint = {effective.rstrip('/')} ({plan})\n")
+    sys.stderr.flush()
+
+    mcp = get_mcp()
+    # Добавление StreamableHTTP: второй branch в main или отдельный entry point,
+    # запуск сервера на указанном host/port.
+    mcp.run(transport="stdio")
+
+
+if __name__ == "__main__":
+    main()

zai_coding_gateway/server.py

@@ -0,0 +1,120 @@
+"""MCP Server: регистрация инструментов, без привязки к транспорту."""
+
+import os
+
+from mcp.server.fastmcp import Context, FastMCP
+
+from zai_coding_gateway.tools.file_io import (
+    _file_uri_to_path,
+    get_cached_project_root,
+    set_cached_project_root,
+)
+
+
+async def _ensure_project_root_from_mcp(context: Context | None) -> None:
+    """
+    Если ZAI_PROJECT_ROOT не задан и кэш пуст — запрашивает roots/list у клиента (Cursor)
+    и кэширует первый корень как проект. При ошибке или отсутствии context — без изменений.
+    """
+    if context is None:
+        return
+    if os.environ.get("ZAI_PROJECT_ROOT") or get_cached_project_root() is not None:
+        return
+    try:
+        result = await context.session.list_roots()
+        if result.roots:
+            path = _file_uri_to_path(str(result.roots[0].uri))
+            if path and os.path.isdir(path):
+                set_cached_project_root(path)
+    except Exception:
+        pass
+from zai_coding_gateway.tools.solve_in_workspace import (
+    answer_session_question as answer_session_question_impl,
+    apply_changes as apply_changes_impl,
+    confirm_session_done as confirm_session_done_impl,
+    solve_in_workspace as solve_in_workspace_impl,
+)
+
+mcp = FastMCP(
+    "ZAI Coding Gateway",
+    json_response=True,
+)
+
+
+@mcp.tool(
+    name="solve_in_workspace",
+    description=(
+        "Точка входа: запускает сессию рабочего пространства. "
+        "Вход: instruction (строка — задача), todos (список строк), file_paths (список путей относительно корня проекта, обязательный начальный контекст), max_steps (число, по умолчанию 10). "
+        "Внутри сессии модель получает полный контекст, может запрашивать файлы (need_more), использовать list_dir/grep/glob/read, затем вернуть отчёт и suggested_changes. "
+        "Модель может предлагать и новые файлы в suggested_changes — создание в проекте только через apply_changes после вашего утверждения. "
+        "Выход: session_id, report, suggested_changes, files_used либо при отсутствии корня проекта или вопросе модели — status «question», session_id, prompt (ответ передайте через answer_session_question)."
+    ),
+)
+async def solve_in_workspace_tool(
+    instruction: str,
+    todos: list[str],
+    file_paths: list[str],
+    max_steps: int = 10,
+    mcp_context: Context | None = None,
+) -> dict:
+    """Run workspace session: need_more/done loop with Z.AI (fixed model flow)."""
+    await _ensure_project_root_from_mcp(mcp_context)
+    return solve_in_workspace_impl(
+        instruction=instruction,
+        todos=todos,
+        file_paths=file_paths,
+        max_steps=max_steps,
+    )
+
+
+@mcp.tool(
+    name="answer_session_question",
+    description=(
+        "Передаёт ответ на вопрос сессии и продолжает выполнение. Вызывайте, когда solve_in_workspace вернул status «question». "
+        "Вход: session_id (из ответа solve_in_workspace), answer (строка — ответ на вопрос: путь к корню проекта или ответ модели, например «да» для дозагрузки файла). "
+        "Выход: тот же формат, что и solve_in_workspace (session_id, report, suggested_changes, files_used либо снова status «question», session_id, prompt)."
+    ),
+)
+async def answer_session_question_tool(session_id: str, answer: str) -> dict:
+    """Submit architect answer to session question and continue."""
+    return answer_session_question_impl(session_id, answer)
+
+
+@mcp.tool(
+    name="apply_changes",
+    description=(
+        "Применяет утверждённые изменения к проекту. Вызывайте после solve_in_workspace, когда приняли решение по suggested_changes. "
+        "Вход: session_id (из ответа solve_in_workspace), changes (список объектов: каждый объект — {path: путь к файлу, content: новое содержимое целиком} или {path, diff: unified diff или полный текст}), confirm_after (bool, по умолчанию True — после применения очистить сессию). "
+        "Пути проверяются только на нахождение внутри корня проекта; новые файлы разрешены. Каждое изменение обрабатывается отдельно. "
+        "Выход: session_id, applied (список применённых путей), errors (список {path, error} для неудачных), cleaned (bool)."
+    ),
+)
+async def apply_changes_tool(
+    session_id: str,
+    changes: list[dict],
+    confirm_after: bool = True,
+) -> dict:
+    """Apply approved changes to project files."""
+    return apply_changes_impl(
+        session_id=session_id,
+        changes=[dict(c) for c in changes],
+        confirm_after=confirm_after,
+    )
+
+
+@mcp.tool(
+    name="confirm_session_done",
+    description=(
+        "Очищает темп-каталог сессии после принятия решения (после apply_changes или при отказе от изменений). "
+        "Вход: session_id (строка из ответа solve_in_workspace). Выход: session_id, cleaned (true)."
+    ),
+)
+async def confirm_session_done_tool(session_id: str) -> dict:
+    """Clean session temp dir."""
+    return confirm_session_done_impl(session_id)
+
+
+def get_mcp() -> FastMCP:
+    """Возвращает экземпляр MCP-сервера (для main или тестов)."""
+    return mcp

zai_coding_gateway/tools/__init__.py

@@ -0,0 +1,10 @@
+"""Инструменты: solve_in_workspace, answer_session_question, apply_changes, confirm_session_done, file_io (внутреннее)."""
+
+from zai_coding_gateway.tools.solve_in_workspace import (
+    answer_session_question,
+    apply_changes,
+    confirm_session_done,
+    solve_in_workspace,
+)
+
+__all__ = ["solve_in_workspace", "answer_session_question", "apply_changes", "confirm_session_done"]

zai_coding_gateway/tools/file_io.py

@@ -0,0 +1,228 @@
+"""Чтение и запись файлов проекта в пределах ZAI_PROJECT_ROOT."""
+
+import os
+import re
+from pathlib import Path
+from urllib.parse import urlparse
+
+from zai_coding_gateway.errors import FileOutsideProjectError, ProjectRootNotSetError
+
+# Каталоги, которые не обходим при grep/glob
+_SKIP_DIRS = frozenset({".git", "node_modules", "__pycache__", ".venv", "venv"})
+_MAX_GLOB_PATHS = 500
+_MAX_GREP_FILES = 100
+_MAX_GREP_OUTPUT_CHARS = 5000
+
+# Кэш корня из MCP roots/list (клиент передаёт workspace). Cursor пока не поддерживает roots/list.
+_cached_mcp_project_root: str | None = None
+
+# Маркеры корня проекта при поиске вверх по дереву от cwd
+_PROJECT_ROOT_MARKERS = ("package.json", "pyproject.toml", ".git")
+
+
+def _infer_project_root_from_cwd() -> str | None:
+    """
+    Поднимается от os.getcwd() вверх и возвращает каталог, содержащий package.json, .git или pyproject.toml.
+    Если cwd — уже корень или не найден маркер — возвращает None.
+    """
+    cwd = os.path.abspath(os.getcwd())
+    if not os.path.isdir(cwd):
+        return None
+    current = cwd
+    while current and current != os.path.dirname(current):
+        for marker in _PROJECT_ROOT_MARKERS:
+            if os.path.exists(os.path.join(current, marker)):
+                return current
+        current = os.path.dirname(current)
+    return None
+
+
+def get_project_root() -> str:
+    """
+    Корень проекта: ZAI_PROJECT_ROOT (env) > кэш из MCP roots/list > вывод по маркерам от cwd > cwd.
+    Если итоговый корень — "/" или не подходит, выбрасывает ProjectRootNotSetError (Cursor не поддерживает roots/list).
+    """
+    root = os.environ.get("ZAI_PROJECT_ROOT")
+    if root:
+        return os.path.abspath(root)
+    if _cached_mcp_project_root is not None:
+        return _cached_mcp_project_root
+    inferred = _infer_project_root_from_cwd()
+    if inferred:
+        return inferred
+    cwd = os.path.abspath(os.getcwd())
+    if not cwd or cwd == os.path.sep or cwd == "/":
+        raise ProjectRootNotSetError(
+            "Корень проекта не задан. Cursor пока не поддерживает MCP roots/list. "
+            "Задайте ZAI_PROJECT_ROOT в настройках MCP (env) — абсолютный путь к папке проекта (например к папке с package.json)."
+        )
+    if not os.path.isdir(cwd):
+        raise ProjectRootNotSetError(
+            "Текущая рабочая директория не является каталогом. Задайте ZAI_PROJECT_ROOT в env MCP."
+        )
+    return cwd
+
+
+def set_cached_project_root(path: str) -> None:
+    """Установить кэшированный корень (из MCP roots/list). Для тестов — сброс через None не делаем, только через env."""
+    global _cached_mcp_project_root
+    _cached_mcp_project_root = os.path.abspath(path) if path else None
+
+
+def get_cached_project_root() -> str | None:
+    """Текущее значение кэша (для тестов)."""
+    return _cached_mcp_project_root
+
+
+def _file_uri_to_path(uri: str) -> str:
+    """Преобразует file:// URI в абсолютный путь к файлу."""
+    parsed = urlparse(uri)
+    if parsed.scheme != "file":
+        return ""
+    path = parsed.path
+    if not path:
+        return ""
+    if os.name == "nt" and path.startswith("/") and len(path) > 2 and path[2] == ":":
+        path = path[1:]
+    return os.path.normpath(path)
+
+
+def resolve_path(path: str, project_root: str | None = None) -> str:
+    """
+    Преобразует путь (относительный или абсолютный) в абсолютный в пределах project_root.
+    Выбрасывает FileOutsideProjectError, если итоговый путь вне корня.
+    """
+    root = project_root or get_project_root()
+    path_abs = os.path.normpath(os.path.join(root, path)) if not os.path.isabs(path) else os.path.normpath(path)
+    root_real = os.path.realpath(root)
+    path_real = os.path.realpath(path_abs)
+    if not path_real.startswith(root_real):
+        raise FileOutsideProjectError(f"Путь вне корня проекта: {path!r}")
+    return path_real
+
+
+def read_file_content(path: str, project_root: str | None = None) -> str:
+    """Читает содержимое файла по пути относительно корня проекта. UTF-8."""
+    resolved = resolve_path(path, project_root)
+    with open(resolved, encoding="utf-8") as f:
+        return f.read()
+
+
+def read_file(path: str, project_root: str | None = None) -> dict:
+    """
+    MCP-инструмент: прочитать файл по пути относительно корня проекта.
+    Возвращает content и path (подтверждение).
+    """
+    root = project_root or get_project_root()
+    content = read_file_content(path, root)
+    return {"content": content, "path": path}
+
+
+def write_file(path: str, content: str, project_root: str | None = None) -> dict:
+    """
+    MCP-инструмент: записать содержимое в файл по пути относительно корня проекта.
+    Создаёт родительские директории при необходимости.
+    """
+    root = project_root or get_project_root()
+    resolved = resolve_path(path, root)
+    os.makedirs(os.path.dirname(resolved) or ".", exist_ok=True)
+    with open(resolved, "w", encoding="utf-8") as f:
+        f.write(content)
+    return {"path": path, "written": True, "summary": f"written {len(content)} bytes"}
+
+
+def list_dir_in_project(path: str = ".", project_root: str | None = None) -> list[str]:
+    """
+    Список имён прямых потомков каталога (один уровень). Путь относительно корня проекта.
+    Для файла или несуществующего пути выбрасывает OSError/FileOutsideProjectError.
+    """
+    root = project_root or get_project_root()
+    path = path.strip() or "."
+    resolved = resolve_path(path, root)
+    if not os.path.isdir(resolved):
+        raise NotADirectoryError(f"Not a directory: {path!r}")
+    return sorted(os.listdir(resolved))
+
+
+def grep_in_project(
+    pattern: str,
+    path_or_none: str | None,
+    project_root: str | None = None,
+    max_files: int = _MAX_GREP_FILES,
+    max_output_chars: int = _MAX_GREP_OUTPUT_CHARS,
+) -> str:
+    """
+    Поиск подстроки (или regex) по содержимому текстовых файлов в пределах path или всего проекта.
+    Возвращает сводку: path:line_no:line_content (обрезано по max_output_chars).
+    """
+    root = project_root or get_project_root()
+    root_path = Path(root).resolve()
+    start_path = root_path
+    if path_or_none and path_or_none.strip():
+        start_path = Path(resolve_path(path_or_none.strip(), root)).resolve()
+        if not start_path.is_dir() and start_path.is_file():
+            files_to_scan = [start_path]
+        else:
+            files_to_scan = []
+            for p in start_path.rglob("*"):
+                if p.is_file() and not any(part in _SKIP_DIRS for part in p.relative_to(root_path).parts):
+                    files_to_scan.append(p)
+                    if len(files_to_scan) >= max_files:
+                        break
+    else:
+        files_to_scan = []
+        for p in root_path.rglob("*"):
+            if p.is_file() and not any(part in _SKIP_DIRS for part in p.relative_to(root_path).parts):
+                files_to_scan.append(p)
+                if len(files_to_scan) >= max_files:
+                    break
+
+    try:
+        regex = re.compile(pattern)
+    except re.error:
+        regex = re.compile(re.escape(pattern))
+
+    lines_out: list[str] = []
+    total_chars = 0
+    for fp in files_to_scan[:max_files]:
+        try:
+            with open(fp, encoding="utf-8", errors="replace") as f:
+                rel = str(fp.relative_to(root_path)).replace("\\", "/")
+                for line_no, line in enumerate(f, 1):
+                    if regex.search(line):
+                        frag = line.strip()[:200]
+                        s = f"{rel}:{line_no}:{frag}\n"
+                        lines_out.append(s)
+                        total_chars += len(s)
+                        if total_chars >= max_output_chars:
+                            lines_out.append("(output truncated)\n")
+                            return "".join(lines_out)
+        except (OSError, UnicodeDecodeError):
+            continue
+    return "".join(lines_out) if lines_out else "(no matches)"
+
+
+def glob_in_project(pattern: str, project_root: str | None = None) -> list[str]:
+    """
+    Список относительных путей файлов, совпадающих с маской (например *.py, **/test_*.py).
+    Исключаются .git, node_modules, __pycache__. Не более _MAX_GLOB_PATHS путей.
+    """
+    root = project_root or get_project_root()
+    root_path = Path(root).resolve()
+    pattern = pattern.strip() or "*"
+    if "**" not in pattern and "/" not in pattern and "\\" not in pattern:
+        pattern = f"**/{pattern}"
+    paths: list[str] = []
+    try:
+        for p in root_path.glob(pattern):
+            if not p.is_file():
+                continue
+            if any(part in _SKIP_DIRS for part in p.relative_to(root_path).parts):
+                continue
+            rel = str(p.relative_to(root_path)).replace("\\", "/")
+            paths.append(rel)
+            if len(paths) >= _MAX_GLOB_PATHS:
+                break
+    except (ValueError, OSError):
+        pass
+    return sorted(paths)