@trac3er/oh-my-god 2.0.4 → 2.0.7

package/build/lib/runtime/providers/kimi_provider.py

@@ -0,0 +1,132 @@
+"""Kimi Code CLI provider -- implements CLIProvider for the ``kimi`` binary."""
+
+from __future__ import annotations
+
+import logging
+import os
+import shlex
+import shutil
+import subprocess
+import uuid
+from typing import Any
+
+from runtime.cli_provider import CLIProvider, register_provider
+from runtime.mcp_config_writers import write_kimi_mcp_config
+from runtime.tmux_session_manager import TmuxSessionManager
+
+_logger = logging.getLogger(__name__)
+
+
+class KimiCodeProvider(CLIProvider):
+    """CLIProvider implementation for the Kimi Code CLI (``kimi``)."""
+
+    # -- identity -----------------------------------------------------------
+
+    def get_name(self) -> str:  # noqa: D401
+        """Return the canonical provider name."""
+        return "kimi"
+
+    # -- detection ----------------------------------------------------------
+
+    def detect(self) -> bool:
+        """Return ``True`` when the ``kimi`` binary is available on PATH."""
+        return shutil.which("kimi") is not None
+
+    # -- authentication -----------------------------------------------------
+
+    def check_auth(self) -> tuple[bool | None, str]:
+        """Check Kimi authentication by parsing ``~/.kimi/config.toml`` for stored credentials."""
+        try:
+            config_path = os.path.expanduser("~/.kimi/config.toml")
+            if not os.path.exists(config_path):
+                return False, "not authenticated — config file not found"
+
+            with open(config_path) as fh:
+                content = fh.read()
+
+            # Look for a token entry in the TOML file
+            if "token" in content:
+                return True, "authenticated"
+            return False, "not authenticated — no token in config"
+        except Exception as exc:
+            return None, f"kimi auth check failed: {exc}"
+
+    # -- invocation ---------------------------------------------------------
+
+    def invoke(self, prompt: str, project_dir: str, timeout: int = 120) -> dict[str, Any]:  # pyright: ignore[reportExplicitAny]
+        """Invoke ``kimi --print -p`` via subprocess."""
+        try:
+            result = self.run_tool(
+                ["kimi", "--print", "-p", prompt],
+                timeout=timeout,
+            )
+            return {
+                "model": "kimi-cli",
+                "output": result.stdout,
+                "exit_code": result.returncode,
+            }
+        except subprocess.TimeoutExpired:
+            return {"error": "kimi-cli timeout", "fallback": "claude"}
+        except FileNotFoundError:
+            return {"error": "kimi-cli not found", "fallback": "claude"}
+        except Exception as exc:
+            return {"error": str(exc), "fallback": "claude"}
+
+    def invoke_json(self, prompt: str, project_dir: str, timeout: int = 120) -> dict[str, Any]:  # pyright: ignore[reportExplicitAny]
+        """Invoke ``kimi --print --output-format stream-json -p`` for JSONL event stream output."""
+        try:
+            result = self.run_tool(
+                ["kimi", "--print", "--output-format", "stream-json", "-p", prompt],
+                timeout=timeout,
+            )
+            return {
+                "model": "kimi-cli",
+                "output": result.stdout,
+                "exit_code": result.returncode,
+            }
+        except subprocess.TimeoutExpired:
+            return {"error": "kimi-cli timeout", "fallback": "claude"}
+        except FileNotFoundError:
+            return {"error": "kimi-cli not found", "fallback": "claude"}
+        except Exception as exc:
+            return {"error": str(exc), "fallback": "claude"}
+
+    def invoke_tmux(self, prompt: str, project_dir: str, timeout: int = 120) -> dict[str, Any]:  # pyright: ignore[reportExplicitAny]
+        """Invoke ``kimi --print -p`` via a persistent tmux session.
+
+        Falls back to :meth:`invoke` on failure.
+        """
+        try:
+            mgr = TmuxSessionManager()
+            session_name = mgr.make_session_name("kimi", unique_id=str(uuid.uuid4())[:8])
+            session = mgr.get_or_create_session(session_name)
+            output = mgr.send_command(session, f"kimi --print -p {shlex.quote(prompt)}", timeout=timeout)
+            mgr.kill_session(session)
+            return {"model": "kimi-cli", "output": output, "exit_code": 0}
+        except Exception as exc:
+            _logger.warning("tmux kimi invocation failed, falling back to subprocess: %s", exc)
+            return self.invoke(prompt, project_dir, timeout=timeout)
+
+    # -- command helpers ----------------------------------------------------
+
+    def get_non_interactive_cmd(self, prompt: str) -> list[str]:
+        """Return the non-interactive command for kimi."""
+        return ["kimi", "--print", "-p", prompt]
+
+    # -- configuration ------------------------------------------------------
+
+    def get_config_path(self) -> str:
+        """Return the Kimi MCP configuration file path."""
+        return os.path.expanduser("~/.kimi/mcp.json")
+
+    def write_mcp_config(self, server_url: str, server_name: str = "memory-server") -> None:
+        """Write an MCP server entry to ``~/.kimi/mcp.json``.
+
+        Uses standard ``mcpServers`` JSON format with ``type: "http"`` and ``url`` field,
+        merging into any existing configuration.
+        """
+        write_kimi_mcp_config(server_url, server_name, config_path=self.get_config_path())
+
+
+# -- auto-register on import -----------------------------------------------
+register_provider(KimiCodeProvider())

package/build/lib/runtime/remote_supervisor.py

@@ -0,0 +1,64 @@
+"""Local-only authenticated supervisor sessions."""
+from __future__ import annotations
+
+import base64
+from datetime import datetime, timezone
+import hashlib
+import hmac
+import json
+from pathlib import Path
+from typing import Any
+from uuid import uuid4
+
+
+def _now() -> str:
+    return datetime.now(timezone.utc).isoformat()
+
+
+def issue_local_supervisor_session(project_dir: str, *, worker_id: str, shared_secret: str) -> dict[str, Any]:
+    session_id = f"session-{uuid4().hex}"
+    issued_at = _now()
+    token_payload = {
+        "session_id": session_id,
+        "worker_id": worker_id,
+        "issued_at": issued_at,
+    }
+    payload_json = json.dumps(token_payload, sort_keys=True, separators=(",", ":"))
+    signature = hmac.new(shared_secret.encode("utf-8"), payload_json.encode("utf-8"), hashlib.sha256).hexdigest()
+    token = base64.urlsafe_b64encode(
+        json.dumps({"payload": token_payload, "signature": signature}, sort_keys=True, separators=(",", ":")).encode("utf-8")
+    ).decode("ascii")
+
+    result = {
+        "schema": "RemoteSupervisorSession",
+        "status": "ok",
+        "session_id": session_id,
+        "worker_id": worker_id,
+        "issued_at": issued_at,
+        "local_only": True,
+        "token": token,
+    }
+
+    rel_path = Path(".omg") / "supervisor" / "sessions" / f"{session_id}.json"
+    path = Path(project_dir) / rel_path
+    path.parent.mkdir(parents=True, exist_ok=True)
+    path.write_text(json.dumps({k: v for k, v in result.items() if k != "token"}, indent=2, ensure_ascii=True) + "\n", encoding="utf-8")
+    result["path"] = rel_path.as_posix()
+    return result
+
+
+def verify_local_supervisor_token(token: str, *, shared_secret: str) -> dict[str, Any]:
+    decoded = json.loads(base64.urlsafe_b64decode(token.encode("ascii")).decode("utf-8"))
+    payload = decoded["payload"]
+    payload_json = json.dumps(payload, sort_keys=True, separators=(",", ":"))
+    signature = str(decoded["signature"])
+    expected = hmac.new(shared_secret.encode("utf-8"), payload_json.encode("utf-8"), hashlib.sha256).hexdigest()
+    status = "ok" if hmac.compare_digest(signature, expected) else "error"
+    return {
+        "schema": "RemoteSupervisorTokenVerification",
+        "status": status,
+        "session_id": str(payload["session_id"]),
+        "worker_id": str(payload["worker_id"]),
+        "issued_at": str(payload["issued_at"]),
+        "local_only": True,
+    }

package/build/lib/runtime/runtime_profile.py

@@ -0,0 +1,61 @@
+"""Runtime profile loading and parallelism budgets."""
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Any
+
+import yaml
+
+
+PROFILE_PRESETS: dict[str, dict[str, Any]] = {
+    "eco": {"profile": "eco", "max_workers": 2, "background_polling": False},
+    "balanced": {"profile": "balanced", "max_workers": 3, "background_polling": False},
+    "turbo": {"profile": "turbo", "max_workers": 5, "background_polling": True},
+}
+
+
+def load_runtime_profile(project_dir: str) -> dict[str, Any]:
+    runtime_path = Path(project_dir) / ".omg" / "runtime.yaml"
+    profile_name = "balanced"
+    if runtime_path.exists():
+        try:
+            payload = yaml.safe_load(runtime_path.read_text(encoding="utf-8")) or {}
+        except Exception:
+            payload = {}
+        if isinstance(payload, dict):
+            candidate = str(payload.get("profile", profile_name)).strip()
+            if candidate in PROFILE_PRESETS:
+                profile_name = candidate
+    return dict(PROFILE_PRESETS[profile_name])
+
+
+def resolve_parallel_workers(project_dir: str, *, requested_workers: int) -> int:
+    profile = load_runtime_profile(project_dir)
+    max_workers = int(profile["max_workers"])
+    cli_cap = _load_cli_parallel_cap(project_dir)
+    if cli_cap is not None:
+        max_workers = min(max_workers, cli_cap)
+    return max(1, min(requested_workers, max_workers))
+
+
+def _load_cli_parallel_cap(project_dir: str) -> int | None:
+    config_path = Path(project_dir) / ".omg" / "state" / "cli-config.yaml"
+    if not config_path.exists():
+        return None
+    try:
+        payload = yaml.safe_load(config_path.read_text(encoding="utf-8")) or {}
+    except Exception:
+        return None
+    if not isinstance(payload, dict):
+        return None
+    cli_configs = payload.get("cli_configs", {})
+    if not isinstance(cli_configs, dict):
+        return None
+    caps = []
+    for config in cli_configs.values():
+        if not isinstance(config, dict):
+            continue
+        value = config.get("max_parallel_agents")
+        if isinstance(value, int) and value > 0:
+            caps.append(value)
+    return min(caps) if caps else None