@trac3er/oh-my-god 2.0.3 → 2.0.5

package/runtime/omg_mcp_server.py

@@ -0,0 +1,207 @@
+from __future__ import annotations
+
+import os
+from collections.abc import AsyncIterator
+from contextlib import asynccontextmanager
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any
+
+from runtime.asset_loader import resolve_asset
+
+_MCP_IMPORT_ERROR: ModuleNotFoundError | None = None
+
+try:
+    from fastmcp import FastMCP
+except ModuleNotFoundError as exc:
+    _MCP_IMPORT_ERROR = exc
+
+    def _passthrough_decorator(*_args: Any, **_kwargs: Any):
+        def decorator(func: Any) -> Any:
+            return func
+
+        return decorator
+
+    @dataclass
+    class _StubPrompt:
+        name: str
+        description: str
+        handler: Any
+
+    @dataclass
+    class _StubResource:
+        uri: str
+        name: str
+        description: str
+        mime_type: str
+        handler: Any
+
+    class FastMCP:  # type: ignore[override]
+        def __init__(self, *_args: Any, **_kwargs: Any) -> None:
+            self._import_error = _MCP_IMPORT_ERROR
+            self.instructions = str(_kwargs.get("instructions", ""))
+            self._prompts: list[_StubPrompt] = []
+            self._resources: dict[str, _StubResource] = {}
+
+        def tool(self, *_args: Any, **_kwargs: Any):
+            return _passthrough_decorator(*_args, **_kwargs)
+
+        def prompt(self, *, name: str, description: str = ""):
+            def decorator(func: Any) -> Any:
+                self._prompts.append(_StubPrompt(name=name, description=description, handler=func))
+                return func
+
+            return decorator
+
+        def resource(
+            self,
+            uri: str,
+            *,
+            name: str = "",
+            description: str = "",
+            mime_type: str = "text/plain",
+        ):
+            def decorator(func: Any) -> Any:
+                self._resources[uri] = _StubResource(
+                    uri=uri,
+                    name=name,
+                    description=description,
+                    mime_type=mime_type,
+                    handler=func,
+                )
+                return func
+
+            return decorator
+
+        async def list_prompts(self) -> list[_StubPrompt]:
+            return list(self._prompts)
+
+        async def list_resources(self) -> list[_StubResource]:
+            return list(self._resources.values())
+
+        async def read_resource(self, uri: str) -> Any:
+            resource = self._resources[str(uri)]
+            return resource.handler()
+
+        def run(self, *_args: Any, **_kwargs: Any) -> None:
+            raise RuntimeError("fastmcp is required to run the OMG MCP server") from self._import_error
+
+    FastMCP.__module__ = "fastmcp"
+
+from control_plane.service import ControlPlaneService
+
+
+MCP_INSTRUCTIONS = (
+    "OMG production control plane MCP. Prefer omg-control prompts and resources for "
+    "contract, release-readiness, and governance context before using direct tools."
+)
+
+
+def _root_dir() -> Path:
+    return Path(__file__).resolve().parents[1]
+
+
+@asynccontextmanager
+async def lifespan(_: object) -> AsyncIterator[None]:
+    yield
+
+
+mcp = FastMCP("OMG Control MCP", lifespan=lifespan, instructions=MCP_INSTRUCTIONS)
+
+
+def _service() -> ControlPlaneService:
+    return ControlPlaneService(project_dir=os.environ.get("CLAUDE_PROJECT_DIR", os.getcwd()))
+
+
+def _read_repo_text(rel_path: str) -> str:
+    return resolve_asset(rel_path).read_text(encoding="utf-8")
+
+
+@mcp.tool()
+def omg_policy_evaluate(tool: str, input: dict[str, Any]) -> dict[str, Any]:
+    _status, payload = _service().policy_evaluate({"tool": tool, "input": input})
+    return payload
+
+
+@mcp.tool()
+def omg_trust_review(file_path: str, old_config: dict[str, Any], new_config: dict[str, Any]) -> dict[str, Any]:
+    _status, payload = _service().trust_review({"file_path": file_path, "old_config": old_config, "new_config": new_config})
+    return payload
+
+
+@mcp.tool()
+def omg_evidence_ingest(
+    run_id: str,
+    tests: list[dict[str, Any]],
+    security_scans: list[dict[str, Any]],
+    diff_summary: dict[str, Any],
+    reproducibility: dict[str, Any],
+    unresolved_risks: list[str],
+    provenance: list[dict[str, Any]] | None = None,
+    trust_scores: dict[str, Any] | None = None,
+    api_twin: dict[str, Any] | None = None,
+) -> dict[str, Any]:
+    _status, payload = _service().evidence_ingest(
+        {
+            "run_id": run_id,
+            "tests": tests,
+            "security_scans": security_scans,
+            "diff_summary": diff_summary,
+            "reproducibility": reproducibility,
+            "unresolved_risks": unresolved_risks,
+            "provenance": provenance or [],
+            "trust_scores": trust_scores or {},
+            "api_twin": api_twin or {},
+        }
+    )
+    return payload
+
+
+@mcp.tool()
+def omg_runtime_dispatch(runtime: str, idea: dict[str, Any]) -> dict[str, Any]:
+    _status, payload = _service().runtime_dispatch({"runtime": runtime, "idea": idea})
+    return payload
+
+
+@mcp.tool()
+def omg_security_check(scope: str = ".", include_live_enrichment: bool = False) -> dict[str, Any]:
+    _status, payload = _service().security_check({"scope": scope, "include_live_enrichment": include_live_enrichment})
+    return payload
+
+
+@mcp.tool()
+def omg_guide_assert(candidate: str, rules: dict[str, Any]) -> dict[str, Any]:
+    _status, payload = _service().guide_assert({"candidate": candidate, "rules": rules})
+    return payload
+
+
+@mcp.prompt(name="omg_contract_summary", description="Summarize the OMG production contract and generated host outputs")
+def omg_contract_summary(channel: str = "public") -> str:
+    return (
+        "Summarize the OMG production control plane contract for channel "
+        f"`{channel}`. Include execution_contract, host_compilation_rules, "
+        "MCP resources, prompts, and release-readiness expectations."
+    )
+
+
+@mcp.resource("resource://omg/contract", name="omg_contract", description="Canonical OMG production contract document", mime_type="text/markdown")
+def omg_contract_resource() -> str:
+    return _read_repo_text("OMG_COMPAT_CONTRACT.md")
+
+
+@mcp.resource(
+    "resource://omg/release-checklist",
+    name="omg_release_checklist",
+    description="Public release checklist for OMG",
+    mime_type="text/markdown",
+)
+def omg_release_checklist_resource() -> str:
+    return _read_repo_text("docs/release-checklist.md")
+
+
+def run_server() -> None:
+    mcp.run(transport="stdio")
+
+
+if __name__ == "__main__":
+    run_server()

package/runtime/preflight.py

@@ -0,0 +1,73 @@
+"""Structured preflight routing for OMG."""
+from __future__ import annotations
+
+from runtime.delta_classifier import classify_project_changes
+from runtime.tracebank import record_trace
+from typing import Any
+
+
+def run_preflight(project_dir: str, *, goal: str) -> dict[str, Any]:
+    lowered = goal.lower()
+    task_class = "implementation"
+    risk_class = "medium"
+    route = "teams"
+    delta = classify_project_changes(project_dir, goal=goal)
+    categories = set(delta["categories"])
+    domain_packs = [category for category in delta["categories"] if category in {"robotics", "vision", "algorithms", "health"}]
+
+    if categories & {"auth", "payment", "health", "compliance"}:
+        task_class = "security"
+        risk_class = "high"
+        route = "security-check"
+    elif categories & {"api"} or any(token in lowered for token in ("openapi", "swagger", "postman", "contract", "fixture", "replay")):
+        task_class = "contract"
+        route = "api-twin"
+    elif any(token in lowered for token in ("auth", "secret", "security", "token", "injection")):
+        task_class = "security"
+        risk_class = "high"
+        route = "security-check"
+    elif any(token in lowered for token in ("full stack", "frontend and backend", "dashboard", "orchestrate")):
+        task_class = "orchestration"
+        risk_class = "high"
+        route = "crazy"
+
+    trace = record_trace(
+        project_dir,
+        trace_type="preflight",
+        route=route,
+        status="ok",
+        plan={"goal": goal, "delta_categories": delta["categories"]},
+        verify={"risk_class": risk_class},
+    )
+
+    return {
+        "schema": "PreflightResult",
+        "project_dir": project_dir,
+        "goal": goal,
+        "task_class": task_class,
+        "risk_class": risk_class,
+        "route": route,
+        "required_tools": _required_tools(route),
+        "required_mcps": ["omg-control"] if route in {"security-check", "api-twin", "crazy"} else [],
+        "missing_constraints": [],
+        "evidence_plan": _evidence_plan(route),
+        "delta_classification": delta,
+        "domain_packs": domain_packs,
+        "trace": {"trace_id": trace["trace_id"], "path": trace["path"]},
+    }
+
+
+def _required_tools(route: str) -> list[str]:
+    return {
+        "security-check": ["security"],
+        "api-twin": ["api-twin"],
+        "crazy": ["teams", "ccg"],
+    }.get(route, ["teams"])
+
+
+def _evidence_plan(route: str) -> list[str]:
+    return {
+        "security-check": ["security findings", "provenance"],
+        "api-twin": ["fixture fidelity", "live verification"],
+        "crazy": ["verification output", "evidence pack"],
+    }.get(route, ["verification output"])

package/runtime/providers/codex_provider.py

@@ -11,6 +11,7 @@ import uuid
 from typing import Any
 
 from runtime.cli_provider import CLIProvider, register_provider
+from runtime.mcp_config_writers import write_codex_mcp_config
 from runtime.tmux_session_manager import TmuxSessionManager
 
 _logger = logging.getLogger(__name__)
@@ -94,18 +95,7 @@ class CodexProvider(CLIProvider):
 
     def write_mcp_config(self, server_url: str, server_name: str = "memory-server") -> None:
         """Write an MCP server entry to ``~/.codex/config.toml``."""
-        config_path = self.get_config_path()
-        os.makedirs(os.path.dirname(config_path), exist_ok=True)
-
-        entry = (
-            f'[mcp_servers."{server_name}"]\n'
-            f'url = "{server_url}"\n'
-        )
-
-        # Append or create
-        mode = "a" if os.path.exists(config_path) else "w"
-        with open(config_path, mode) as fh:
-            fh.write(entry)
+        write_codex_mcp_config(server_url, server_name, config_path=self.get_config_path())
 
 
 # -- auto-register on import -----------------------------------------------

package/runtime/providers/gemini_provider.py

@@ -2,7 +2,6 @@
 
 from __future__ import annotations
 
-import json
 import logging
 import os
 import shlex
@@ -12,6 +11,7 @@ import uuid
 from typing import Any
 
 from runtime.cli_provider import CLIProvider, register_provider
+from runtime.mcp_config_writers import write_gemini_mcp_config
 from runtime.tmux_session_manager import TmuxSessionManager
 
 _logger = logging.getLogger(__name__)
@@ -102,26 +102,7 @@ class GeminiProvider(CLIProvider):
         Uses JSON format with ``mcpServers`` key and ``httpUrl`` field,
         merging into any existing configuration.
         """
-        config_path = self.get_config_path()
-        os.makedirs(os.path.dirname(config_path), exist_ok=True)
-
-        # Load existing config or start fresh
-        existing: dict[str, Any] = {}  # pyright: ignore[reportExplicitAny]
-        if os.path.exists(config_path):
-            with open(config_path) as fh:
-                try:
-                    existing = json.load(fh)
-                except (json.JSONDecodeError, ValueError):
-                    existing = {}
-
-        # Ensure mcpServers dict exists
-        if "mcpServers" not in existing:
-            existing["mcpServers"] = {}
-
-        existing["mcpServers"][server_name] = {"httpUrl": server_url}
-
-        with open(config_path, "w") as fh:
-            json.dump(existing, fh, indent=2)
+        write_gemini_mcp_config(server_url, server_name, config_path=self.get_config_path())
 
 
 # -- auto-register on import -----------------------------------------------

package/runtime/providers/kimi_provider.py

@@ -2,7 +2,6 @@
 
 from __future__ import annotations
 
-import json
 import logging
 import os
 import shlex
@@ -12,6 +11,7 @@ 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__)
@@ -125,26 +125,7 @@ class KimiCodeProvider(CLIProvider):
         Uses standard ``mcpServers`` JSON format with ``type: "http"`` and ``url`` field,
         merging into any existing configuration.
         """
-        config_path = self.get_config_path()
-        os.makedirs(os.path.dirname(config_path), exist_ok=True)
-
-        # Load existing config or start fresh
-        existing: dict[str, Any] = {}  # pyright: ignore[reportExplicitAny]
-        if os.path.exists(config_path):
-            with open(config_path) as fh:
-                try:
-                    existing = json.load(fh)
-                except (json.JSONDecodeError, ValueError):
-                    existing = {}
-
-        # Ensure mcpServers dict exists
-        if "mcpServers" not in existing:
-            existing["mcpServers"] = {}
-
-        existing["mcpServers"][server_name] = {"type": "http", "url": server_url}
-
-        with open(config_path, "w") as fh:
-            json.dump(existing, fh, indent=2)
+        write_kimi_mcp_config(server_url, server_name, config_path=self.get_config_path())
 
 
 # -- auto-register on import -----------------------------------------------

package/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/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