@trac3er/oh-my-god 2.0.4 → 2.0.7

package/build/lib/runtime/mcp_config_writers.py

@@ -0,0 +1,233 @@
+from __future__ import annotations
+
+import json
+import os
+from pathlib import Path
+from typing import cast
+
+from hooks.security_validators import (
+    toml_quote_string,
+    validate_server_name,
+    validate_server_url,
+)
+
+
+def _atomic_write_text(path: Path, content: str) -> None:
+    path.parent.mkdir(parents=True, exist_ok=True)
+    tmp_path = path.with_name(f"{path.name}.tmp")
+    _ = tmp_path.write_text(content)
+    _ = os.replace(tmp_path, path)
+
+
+def _load_json(path: Path) -> dict[str, object]:
+    if not path.exists():
+        return {}
+    try:
+        parsed = cast(object, json.loads(path.read_text()))
+        if isinstance(parsed, dict):
+            return cast(dict[str, object], parsed)
+        return {}
+    except (json.JSONDecodeError, ValueError):
+        return {}
+
+
+def _write_json(path: Path, data: dict[str, object]) -> None:
+    _atomic_write_text(path, json.dumps(data, indent=2) + "\n")
+
+
+def get_managed_python_path(claude_config_dir: str | None = None) -> str:
+    """Return the absolute path to the managed OMG venv Python interpreter.
+
+    Falls back to ``CLAUDE_CONFIG_DIR`` env var, then ``~/.claude``.
+    """
+    if claude_config_dir is None:
+        claude_config_dir = os.environ.get(
+            "CLAUDE_CONFIG_DIR", str(Path.home() / ".claude")
+        )
+    return str(Path(claude_config_dir) / "omg-runtime" / ".venv" / "bin" / "python")
+
+
+def _validated_server_input(server_url: str, server_name: str) -> tuple[str, str]:
+    return validate_server_url(server_url), validate_server_name(server_name)
+
+
+def _validated_stdio_input(command: str, args: list[str], server_name: str) -> tuple[str, list[str], str]:
+    normalized_name = validate_server_name(server_name)
+    normalized_command = str(command).strip()
+    if not normalized_command or "\n" in normalized_command or "\r" in normalized_command:
+        raise ValueError("Invalid command: newline characters are not allowed")
+    normalized_args = [str(arg) for arg in args]
+    for arg in normalized_args:
+        if "\n" in arg or "\r" in arg:
+            raise ValueError("Invalid args: newline characters are not allowed")
+    return normalized_command, normalized_args, normalized_name
+
+
+def _write_json_mcp_server(path: Path, server_name: str, payload: dict[str, object]) -> None:
+    config = _load_json(path)
+    mcp_servers = config.get("mcpServers")
+    if not isinstance(mcp_servers, dict):
+        mcp_servers = {}
+        config["mcpServers"] = mcp_servers
+    mcp_servers[server_name] = payload
+    _write_json(path, config)
+
+
+def write_claude_mcp_config(project_dir: str, server_url: str, server_name: str = "memory-server") -> None:
+    server_url, server_name = _validated_server_input(server_url, server_name)
+    config_path = Path(project_dir) / ".mcp.json"
+    _write_json_mcp_server(config_path, server_name, {"type": "http", "url": server_url})
+
+
+def write_claude_mcp_stdio_config(
+    project_dir: str,
+    *,
+    command: str,
+    args: list[str],
+    server_name: str = "omg-control",
+) -> None:
+    command, args, server_name = _validated_stdio_input(command, args, server_name)
+    config_path = Path(project_dir) / ".mcp.json"
+    _write_json_mcp_server(config_path, server_name, {"command": command, "args": args})
+
+
+def write_codex_mcp_config(
+    server_url: str,
+    server_name: str = "memory-server",
+    *,
+    config_path: str | Path | None = None,
+) -> None:
+    server_url, server_name = _validated_server_input(server_url, server_name)
+    target_path = Path(config_path) if config_path is not None else Path.home() / ".codex" / "config.toml"
+    target_path.parent.mkdir(parents=True, exist_ok=True)
+
+    existing = target_path.read_text() if target_path.exists() else ""
+    lines = existing.splitlines(keepends=True)
+
+    header_unquoted = f"[mcp_servers.{server_name}]"
+    header_quoted = f"[mcp_servers.\"{server_name}\"]"
+    headers = {header_unquoted, header_quoted}
+
+    start_idx: int | None = None
+    for idx, line in enumerate(lines):
+        if line.strip() in headers:
+            start_idx = idx
+            break
+
+    block = [
+        f"{header_unquoted}\n",
+        'type = "http"\n',
+        f'url = "{toml_quote_string(server_url)}"\n',
+        "\n",
+    ]
+
+    if start_idx is None:
+        if existing and not existing.endswith("\n"):
+            existing += "\n"
+        content = existing + "".join(block)
+        _atomic_write_text(target_path, content)
+        return
+
+    end_idx = len(lines)
+    for idx in range(start_idx + 1, len(lines)):
+        stripped = lines[idx].strip()
+        if stripped.startswith("[") and stripped.endswith("]"):
+            end_idx = idx
+            break
+
+    updated_lines = lines[:start_idx] + block + lines[end_idx:]
+    _atomic_write_text(target_path, "".join(updated_lines))
+
+
+def write_codex_mcp_stdio_config(
+    *,
+    command: str,
+    args: list[str],
+    server_name: str = "omg-control",
+    config_path: str | Path | None = None,
+) -> None:
+    command, args, server_name = _validated_stdio_input(command, args, server_name)
+    target_path = Path(config_path) if config_path is not None else Path.home() / ".codex" / "config.toml"
+    target_path.parent.mkdir(parents=True, exist_ok=True)
+
+    existing = target_path.read_text() if target_path.exists() else ""
+    lines = existing.splitlines(keepends=True)
+    header_unquoted = f"[mcp_servers.{server_name}]"
+    header_quoted = f"[mcp_servers.\"{server_name}\"]"
+    headers = {header_unquoted, header_quoted}
+
+    start_idx: int | None = None
+    for idx, line in enumerate(lines):
+        if line.strip() in headers:
+            start_idx = idx
+            break
+
+    args_text = ", ".join(f'"{toml_quote_string(arg)}"' for arg in args)
+    block = [
+        f"{header_unquoted}\n",
+        f'command = "{toml_quote_string(command)}"\n',
+        f"args = [{args_text}]\n",
+        "\n",
+    ]
+
+    if start_idx is None:
+        if existing and not existing.endswith("\n"):
+            existing += "\n"
+        _atomic_write_text(target_path, existing + "".join(block))
+        return
+
+    end_idx = len(lines)
+    for idx in range(start_idx + 1, len(lines)):
+        stripped = lines[idx].strip()
+        if stripped.startswith("[") and stripped.endswith("]"):
+            end_idx = idx
+            break
+
+    updated_lines = lines[:start_idx] + block + lines[end_idx:]
+    _atomic_write_text(target_path, "".join(updated_lines))
+
+
+def write_gemini_mcp_config(
+    server_url: str,
+    server_name: str = "memory-server",
+    *,
+    config_path: str | Path | None = None,
+) -> None:
+    server_url, server_name = _validated_server_input(server_url, server_name)
+    target_path = Path(config_path) if config_path is not None else Path.home() / ".gemini" / "settings.json"
+    _write_json_mcp_server(target_path, server_name, {"httpUrl": server_url})
+
+
+def write_gemini_mcp_stdio_config(
+    *,
+    command: str,
+    args: list[str],
+    server_name: str = "omg-control",
+    config_path: str | Path | None = None,
+) -> None:
+    command, args, server_name = _validated_stdio_input(command, args, server_name)
+    target_path = Path(config_path) if config_path is not None else Path.home() / ".gemini" / "settings.json"
+    _write_json_mcp_server(target_path, server_name, {"command": command, "args": args})
+
+
+def write_kimi_mcp_config(
+    server_url: str,
+    server_name: str = "memory-server",
+    *,
+    config_path: str | Path | None = None,
+) -> None:
+    server_url, server_name = _validated_server_input(server_url, server_name)
+    target_path = Path(config_path) if config_path is not None else Path.home() / ".kimi" / "mcp.json"
+    _write_json_mcp_server(target_path, server_name, {"type": "http", "url": server_url})
+
+
+def write_kimi_mcp_stdio_config(
+    *,
+    command: str,
+    args: list[str],
+    server_name: str = "omg-control",
+    config_path: str | Path | None = None,
+) -> None:
+    command, args, server_name = _validated_stdio_input(command, args, server_name)
+    target_path = Path(config_path) if config_path is not None else Path.home() / ".kimi" / "mcp.json"
+    _write_json_mcp_server(target_path, server_name, {"command": command, "args": args})

package/build/lib/runtime/mcp_lifecycle.py

@@ -0,0 +1,175 @@
+"""MCP server lifecycle manager — start/stop/health/ensure."""
+from __future__ import annotations
+
+import os
+import signal
+import subprocess
+import sys
+import time
+from pathlib import Path
+from typing import Any
+
+
+def get_pid_file_path() -> str:
+    """Return path to the PID file for the memory server."""
+    return str(Path.home() / ".omg" / "shared-memory" / "server.pid")
+
+
+def get_server_url() -> str:
+    """Return the MCP server URL using host/port from mcp_memory_server."""
+    from runtime.mcp_memory_server import get_host, get_port
+
+    return f"http://{get_host()}:{get_port()}/mcp"
+
+
+def _read_pid() -> int | None:
+    """Read PID from PID file. Returns None if missing or invalid."""
+    try:
+        return int(Path(get_pid_file_path()).read_text().strip())
+    except (FileNotFoundError, ValueError):
+        return None
+
+
+def is_server_running() -> bool:
+    """Check if the memory server process is alive."""
+    pid = _read_pid()
+    if pid is None:
+        return False
+    try:
+        os.kill(pid, 0)
+        return True
+    except ProcessLookupError:
+        return False
+    except PermissionError:
+        return True  # Process exists but we can't signal it
+
+
+def _wait_for_health(url: str, timeout: float = 5.0) -> bool:
+    """Wait for server health endpoint to respond."""
+    import urllib.request
+
+    deadline = time.monotonic() + timeout
+    while time.monotonic() < deadline:
+        try:
+            with urllib.request.urlopen(url, timeout=1) as resp:
+                if resp.status == 200:
+                    return True
+        except Exception:
+            time.sleep(0.1)
+    return False
+
+
+def _wait_for_exit(pid: int, timeout: float = 5.0) -> bool:
+    """Wait for a process to exit."""
+    deadline = time.monotonic() + timeout
+    while time.monotonic() < deadline:
+        try:
+            os.kill(pid, 0)
+            time.sleep(0.1)
+        except (ProcessLookupError, PermissionError):
+            return True
+    return False
+
+
+def start_memory_server() -> dict[str, Any]:
+    """Start the MCP memory server as a background process."""
+    if is_server_running():
+        return {"status": "already_running", "url": get_server_url()}
+
+    pid_path = Path(get_pid_file_path())
+    pid_path.parent.mkdir(parents=True, exist_ok=True)
+
+    try:
+        server_script = str(Path(__file__).parent / "mcp_memory_server.py")
+        proc = subprocess.Popen(
+            [sys.executable, server_script],
+            stdout=subprocess.DEVNULL,
+            stderr=subprocess.DEVNULL,
+        )
+
+        health_url = get_server_url().replace("/mcp", "/health")
+        if _wait_for_health(health_url):
+            pid_path.write_text(str(proc.pid))
+            return {"status": "started", "pid": proc.pid, "url": get_server_url()}
+        try:
+            proc.terminate()
+            proc.wait(timeout=2)
+        except Exception:
+            try:
+                proc.kill()
+            except Exception:
+                pass
+        if pid_path.exists():
+            pid_path.unlink()
+        return {
+            "status": "error",
+            "message": "Server did not respond within timeout",
+        }
+    except Exception as exc:
+        return {"status": "error", "message": str(exc)}
+
+
+def stop_memory_server() -> dict[str, Any]:
+    """Stop the MCP memory server."""
+    if not is_server_running():
+        return {"status": "not_running"}
+
+    pid = _read_pid()
+    if pid is None:
+        return {"status": "not_running"}
+
+    try:
+        os.kill(pid, signal.SIGTERM)
+        _wait_for_exit(pid)
+
+        pid_path = Path(get_pid_file_path())
+        if pid_path.exists():
+            pid_path.unlink()
+
+        return {"status": "stopped", "pid": pid}
+    except ProcessLookupError:
+        pid_path = Path(get_pid_file_path())
+        if pid_path.exists():
+            pid_path.unlink()
+        return {"status": "stopped", "pid": pid}
+    except Exception as exc:
+        return {"status": "error", "message": str(exc)}
+
+
+def check_memory_server() -> dict[str, Any]:
+    """Check the memory server status."""
+    running = is_server_running()
+    if not running:
+        return {"running": False, "url": None, "pid": None}
+
+    return {
+        "running": True,
+        "url": get_server_url(),
+        "pid": _read_pid(),
+    }
+
+
+_HTTP_MEMORY_RESTRICTED_PRESETS = frozenset({"safe", "balanced"})
+
+
+def _get_current_preset() -> str:
+    return os.environ.get("OMG_PRESET", "safe")
+
+
+def ensure_memory_server() -> dict[str, Any]:
+    """Ensure the memory server is running (idempotent).
+
+    Skips startup when the active preset is ``safe`` or ``balanced`` —
+    HTTP memory is an opt-in surface starting at the ``interop`` tier.
+    """
+    preset = _get_current_preset()
+    if preset in _HTTP_MEMORY_RESTRICTED_PRESETS:
+        return {"status": "skipped", "reason": f"HTTP memory disabled for preset '{preset}'"}
+    if is_server_running():
+        return {"status": "already_running", "url": get_server_url()}
+    return start_memory_server()
+
+
+# Feature flag: auto-start on import (respects preset restrictions)
+if os.environ.get("OMG_MEMORY_AUTOSTART") == "1":
+    ensure_memory_server()

package/build/lib/runtime/mcp_memory_server.py

@@ -0,0 +1,135 @@
+from __future__ import annotations
+
+import json
+import os
+from collections.abc import AsyncIterator
+from contextlib import asynccontextmanager
+from typing import Any
+
+_MCP_IMPORT_ERROR: ModuleNotFoundError | None = None
+
+try:
+    from fastmcp import FastMCP
+    from starlette.requests import Request
+    from starlette.responses import JSONResponse
+except ModuleNotFoundError as exc:
+    _MCP_IMPORT_ERROR = exc
+    Request = Any
+
+    class JSONResponse(dict):
+        def __init__(self, content: dict[str, Any]):
+            super().__init__(content)
+
+    def _passthrough_decorator(*_args: Any, **_kwargs: Any):
+        def decorator(func: Any) -> Any:
+            return func
+
+        return decorator
+
+    class FastMCP:  # type: ignore[override]
+        def __init__(self, *_args: Any, **_kwargs: Any) -> None:
+            self._import_error = _MCP_IMPORT_ERROR
+
+        custom_route = staticmethod(_passthrough_decorator)
+        tool = staticmethod(_passthrough_decorator)
+        resource = staticmethod(_passthrough_decorator)
+
+        def run(self, *_args: Any, **_kwargs: Any) -> None:
+            raise RuntimeError("fastmcp and starlette are required to run the OMG memory server") from self._import_error
+
+    FastMCP.__module__ = "fastmcp"
+
+from runtime.memory_store import MemoryStore, MemoryStoreFullError
+
+_store = MemoryStore()
+
+
+def _load_state() -> None:
+    return None
+
+
+def _save_state() -> None:
+    return None
+
+
+@asynccontextmanager
+async def lifespan(_: object) -> AsyncIterator[None]:
+    _load_state()
+    try:
+        yield
+    finally:
+        _save_state()
+
+
+mcp = FastMCP("OMG Memory Server", lifespan=lifespan)
+
+
+@mcp.custom_route("/health", methods=["GET"])
+async def health(_: Request) -> JSONResponse:
+    return JSONResponse({"status": "ok", "version": "2.0.7"})
+
+
+@mcp.tool()
+def memory_store(
+    key: str,
+    content: str,
+    source_cli: str,
+    tags: list[str] | None = None,
+) -> dict[str, Any]:
+    try:
+        return _store.add(key=key, content=content, source_cli=source_cli, tags=tags)
+    except MemoryStoreFullError as exc:
+        return {"error": str(exc)}
+
+
+@mcp.tool()
+def memory_search(
+    query: str,
+    source_cli: str | None = None,
+) -> list[dict[str, Any]]:
+    return _store.search(query=query, source_cli=source_cli)
+
+
+@mcp.tool()
+def memory_list(
+    source_cli: str | None = None,
+) -> list[dict[str, Any]]:
+    return _store.list_all(source_cli=source_cli)
+
+
+@mcp.tool()
+def memory_delete(item_id: str) -> dict[str, Any]:
+    deleted = _store.delete(item_id)
+    return {"deleted": deleted, "id": item_id}
+
+
+@mcp.tool()
+def memory_import(items: list[dict[str, Any]]) -> dict[str, int]:
+    count = _store.import_items(items)
+    return {"imported": count}
+
+
+@mcp.tool()
+def memory_export() -> list[dict[str, Any]]:
+    return _store.export_all()
+
+
+@mcp.resource("memory://all")
+def memory_all_resource() -> str:
+    return json.dumps(_store.list_all())
+
+
+def get_host() -> str:
+    return os.environ.get("OMG_MEMORY_HOST", "127.0.0.1")
+
+
+def get_port() -> int:
+    return int(os.environ.get("OMG_MEMORY_PORT", "8765"))
+
+
+def run_server() -> None:
+    mcp.run(transport="http", host=get_host(), port=get_port())
+
+
+if __name__ == "__main__":
+    run_server()