headroom/e2e/wrap/run.py

from __future__ import annotations

import json
import os
import shutil
import signal
import stat
import subprocess
import sys
import tempfile
import textwrap
import threading
import time
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from pathlib import Path
from typing import Any

import httpx

REPO_ROOT = Path("/workspace")
PLUGIN_DIR = REPO_ROOT / "plugins" / "openclaw"
SDK_DIR = REPO_ROOT / "sdk" / "typescript"
RTK_MARKER = "<!-- headroom:rtk-instructions -->"
PROXY_PORT = 28887
CODEX_PORT = 28888
AIDER_PORT = 28889
CURSOR_PORT = 28890
OPENCLAW_PROXY_PORT = 28891


def log(message: str) -> None:
    print(f"[wrap-e2e] {message}", flush=True)


def run(
    cmd: list[str],
    *,
    env: dict[str, str] | None = None,
    cwd: Path | None = None,
    timeout: int = 180,
) -> subprocess.CompletedProcess[str]:
    log(f"$ {' '.join(cmd)}")
    result = subprocess.run(
        cmd,
        env=env,
        cwd=str(cwd) if cwd else None,
        capture_output=True,
        text=True,
        encoding="utf-8",
        errors="replace",
        timeout=timeout,
    )
    if result.stdout.strip():
        print(result.stdout.rstrip(), flush=True)
    if result.stderr.strip():
        print(result.stderr.rstrip(), file=sys.stderr, flush=True)
    if result.returncode != 0:
        raise RuntimeError(f"Command failed with exit code {result.returncode}: {' '.join(cmd)}")
    return result


def assert_true(condition: bool, message: str) -> None:
    if not condition:
        raise AssertionError(message)


def write_executable(path: Path, content: str) -> None:
    path.write_text(content, encoding="utf-8")
    current_mode = path.stat().st_mode
    path.chmod(current_mode | stat.S_IXUSR | stat.S_IXGRP | stat.S_IXOTH)


class MockOpenAIServer(ThreadingHTTPServer):
    allow_reuse_address = True
    daemon_threads = True

    def __init__(self, server_address: tuple[str, int]) -> None:
        super().__init__(server_address, MockOpenAIHandler)
        self.requests: list[dict[str, Any]] = []


class MockOpenAIHandler(BaseHTTPRequestHandler):
    protocol_version = "HTTP/1.1"

    def log_message(self, format: str, *args: object) -> None:  # noqa: A003
        return

    def _write_json(self, status_code: int, payload: dict[str, Any]) -> None:
        body = json.dumps(payload).encode("utf-8")
        self.send_response(status_code)
        self.send_header("Content-Type", "application/json")
        self.send_header("Content-Length", str(len(body)))
        self.end_headers()
        self.wfile.write(body)

    def _record(self, body: dict[str, Any] | None = None) -> None:
        server = self.server
        assert isinstance(server, MockOpenAIServer)
        server.requests.append(
            {
                "method": self.command,
                "path": self.path,
                "authorization": self.headers.get("Authorization"),
                "body": body,
            }
        )

    def do_GET(self) -> None:  # noqa: N802
        self._record()
        if self.path == "/v1/models":
            self._write_json(
                200,
                {
                    "object": "list",
                    "data": [
                        {
                            "id": "gpt-4o-mini",
                            "object": "model",
                            "owned_by": "openai",
                        }
                    ],
                },
            )
            return
        self._write_json(404, {"error": {"message": "not found"}})

    def do_POST(self) -> None:  # noqa: N802
        length = int(self.headers.get("Content-Length", "0"))
        raw_body = self.rfile.read(length) if length else b""
        payload = json.loads(raw_body.decode("utf-8") or "{}")
        self._record(body=payload)
        if self.path == "/v1/chat/completions":
            self._write_json(
                200,
                {
                    "id": "chatcmpl-e2e",
                    "object": "chat.completion",
                    "created": 0,
                    "model": payload.get("model", "gpt-4o-mini"),
                    "choices": [
                        {
                            "index": 0,
                            "finish_reason": "stop",
                            "message": {
                                "role": "assistant",
                                "content": "mock completion from upstream",
                            },
                        }
                    ],
                    "usage": {
                        "prompt_tokens": 12,
                        "completion_tokens": 5,
                        "total_tokens": 17,
                    },
                },
            )
            return
        self._write_json(404, {"error": {"message": "not found"}})


def wait_for_http(url: str, *, timeout: int = 30) -> httpx.Response:
    deadline = time.time() + timeout
    last_error: Exception | None = None
    while time.time() < deadline:
        try:
            response = httpx.get(url, timeout=2.0)
            if response.status_code < 500:
                return response
        except Exception as exc:  # pragma: no cover - best effort retry surface
            last_error = exc
        time.sleep(0.5)
    raise RuntimeError(f"Timed out waiting for {url}: {last_error}")


def wait_for_output(proc: subprocess.Popen[str], text: str, *, timeout: int = 30) -> str:
    deadline = time.time() + timeout
    chunks: list[str] = []
    while time.time() < deadline:
        if proc.stdout is None:
            break
        line = proc.stdout.readline()
        if line:
            chunks.append(line)
            if text in "".join(chunks):
                return "".join(chunks)
        elif proc.poll() is not None:
            break
    output = "".join(chunks)
    raise RuntimeError(f"Timed out waiting for process output '{text}'. Output so far:\n{output}")


def read_jsonl(path: Path) -> list[dict[str, Any]]:
    if not path.exists():
        return []
    return [json.loads(line) for line in path.read_text(encoding="utf-8").splitlines() if line]


def create_shims(shim_dir: Path) -> None:
    generic_shim = textwrap.dedent(
        """\
        #!/usr/bin/env python3
        from __future__ import annotations

        import json
        import os
        import sys
        import urllib.request
        from pathlib import Path

        tool = Path(sys.argv[0]).name
        log_dir = Path(os.environ["HEADROOM_E2E_LOG_DIR"])
        log_dir.mkdir(parents=True, exist_ok=True)
        record = {
            "tool": tool,
            "argv": sys.argv[1:],
            "cwd": os.getcwd(),
            "env": {
                key: os.environ.get(key)
                for key in ("OPENAI_BASE_URL", "OPENAI_API_BASE", "ANTHROPIC_BASE_URL")
                if os.environ.get(key) is not None
            },
        }

        probes = []

        def fetch(url: str, *, headers: dict[str, str] | None = None) -> None:
            request = urllib.request.Request(url, headers=headers or {})
            with urllib.request.urlopen(request, timeout=10) as response:
                probes.append({"url": url, "status": response.status})

        openai_base = os.environ.get("OPENAI_BASE_URL") or os.environ.get("OPENAI_API_BASE")
        if openai_base:
            fetch(
                f"{openai_base.rstrip('/')}/models",
                headers={"Authorization": "Bearer test-key"},
            )

        anthropic_base = os.environ.get("ANTHROPIC_BASE_URL")
        if anthropic_base:
            fetch(f"{anthropic_base.rstrip('/')}/health")

        record["probes"] = probes

        with (log_dir / f"{tool}.jsonl").open("a", encoding="utf-8") as handle:
            handle.write(json.dumps(record) + "\\n")
        print(f"{tool} shim executed")
        raise SystemExit(0)
        """
    )
    rtk_shim = textwrap.dedent(
        """\
        #!/usr/bin/env python3
        from __future__ import annotations

        import sys

        if "--version" in sys.argv:
            print("rtk e2e-shim")
        else:
            print("rtk shim")
        raise SystemExit(0)
        """
    )
    write_executable(shim_dir / "claude", generic_shim)
    write_executable(shim_dir / "codex", generic_shim)
    write_executable(shim_dir / "aider", generic_shim)
    write_executable(shim_dir / "rtk", rtk_shim)


def start_mock_server(port: int) -> tuple[MockOpenAIServer, threading.Thread]:
    server = MockOpenAIServer(("127.0.0.1", port))
    thread = threading.Thread(target=server.serve_forever, daemon=True)
    thread.start()
    return server, thread


def start_proxy(port: int, env: dict[str, str]) -> subprocess.Popen[str]:
    log(f"Starting headroom proxy on port {port}")
    proc = subprocess.Popen(
        ["headroom", "proxy", "--host", "127.0.0.1", "--port", str(port)],
        env=env,
        stdout=subprocess.PIPE,
        stderr=subprocess.STDOUT,
        text=True,
        encoding="utf-8",
        errors="replace",
    )
    wait_for_http(f"http://127.0.0.1:{port}/health", timeout=30)
    return proc


def stop_process(proc: subprocess.Popen[str]) -> None:
    if proc.poll() is not None:
        return
    proc.send_signal(signal.SIGINT)
    try:
        proc.communicate(timeout=10)
    except subprocess.TimeoutExpired:
        proc.kill()
        proc.communicate(timeout=5)


def wait_for_command_success(
    cmd: list[str],
    *,
    env: dict[str, str],
    cwd: Path | None = None,
    timeout: int = 30,
) -> subprocess.CompletedProcess[str]:
    deadline = time.time() + timeout
    last_output = ""
    while time.time() < deadline:
        remaining = deadline - time.time()
        per_call_timeout = max(1.0, min(5.0, remaining))
        try:
            result = subprocess.run(
                cmd,
                env=env,
                cwd=str(cwd) if cwd else None,
                capture_output=True,
                text=True,
                encoding="utf-8",
                errors="replace",
                timeout=per_call_timeout,
            )
        except subprocess.TimeoutExpired:
            last_output = f"Command timed out after {per_call_timeout:.1f}s"
            time.sleep(1)
            continue
        if result.returncode == 0:
            if result.stdout.strip():
                print(result.stdout.rstrip(), flush=True)
            if result.stderr.strip():
                print(result.stderr.rstrip(), file=sys.stderr, flush=True)
            return result
        last_output = "\n".join(
            part for part in (result.stdout.strip(), result.stderr.strip()) if part
        )
        time.sleep(1)
    raise RuntimeError(
        f"Timed out waiting for command to succeed: {' '.join(cmd)}\nLast output:\n{last_output}"
    )


def start_openclaw_gateway(env: dict[str, str], cwd: Path) -> subprocess.Popen[str]:
    log("Starting OpenClaw gateway for e2e verification")
    return subprocess.Popen(
        ["openclaw", "gateway"],
        env=env,
        cwd=str(cwd),
        stdout=subprocess.PIPE,
        stderr=subprocess.STDOUT,
        text=True,
        encoding="utf-8",
        errors="replace",
    )


def stop_openclaw_gateway(env: dict[str, str], cwd: Path) -> None:
    log("Stopping OpenClaw gateway after e2e verification")
    run(["openclaw", "gateway", "stop"], env=env, cwd=cwd, timeout=60)


def verify_installs() -> None:
    log("Verifying installed packages and binaries")
    for tool in ("headroom", "codex", "aider", "openclaw"):
        assert_true(shutil.which(tool) is not None, f"Expected '{tool}' on PATH")
    run(["headroom", "--help"], timeout=30)
    run(["npm", "list", "-g", "--depth=0", "@openai/codex", "openclaw"], timeout=60)
    run(["/opt/aider-venv/bin/python", "-m", "pip", "show", "aider-chat"], timeout=60)


def prepare_local_openclaw_plugin(base_env: dict[str, str], tmp_dir: Path) -> Path:
    log("Preparing local TypeScript package for OpenClaw plugin build")
    sdk_dir = tmp_dir / "sdk-typescript"
    plugin_dir = tmp_dir / "openclaw-plugin"
    shutil.copytree(SDK_DIR, sdk_dir)
    shutil.copytree(PLUGIN_DIR, plugin_dir)

    plugin_lock = plugin_dir / "package-lock.json"
    if plugin_lock.exists():
        plugin_lock.unlink()

    run(["npm", "install"], env=base_env, cwd=sdk_dir, timeout=600)
    run(["npm", "run", "build"], env=base_env, cwd=sdk_dir, timeout=600)
    pack_result = run(["npm", "pack"], env=base_env, cwd=sdk_dir, timeout=600)
    tarball_name = pack_result.stdout.strip().splitlines()[-1].strip()
    tarball_path = sdk_dir / tarball_name
    assert_true(tarball_path.exists(), "Expected npm pack to produce a local SDK tarball")

    package_json_path = plugin_dir / "package.json"
    package_json = json.loads(package_json_path.read_text(encoding="utf-8"))
    package_json["dependencies"]["headroom-ai"] = f"file:{tarball_path.as_posix()}"
    package_json_path.write_text(f"{json.dumps(package_json, indent=2)}\n", encoding="utf-8")

    return plugin_dir


def verify_proxy_round_trip(base_env: dict[str, str], mock_server: MockOpenAIServer) -> None:
    proxy_port = PROXY_PORT
    proc = start_proxy(proxy_port, base_env)
    try:
        health = wait_for_http(f"http://127.0.0.1:{proxy_port}/health")
        assert_true(health.status_code == 200, "Proxy health check should return 200")

        models = httpx.get(
            f"http://127.0.0.1:{proxy_port}/v1/models",
            headers={"Authorization": "Bearer test-key"},
            timeout=10.0,
        )
        assert_true(models.status_code == 200, "Proxy should pass through /v1/models")
        assert_true(models.json()["data"][0]["id"] == "gpt-4o-mini", "Unexpected models payload")

        chat = httpx.post(
            f"http://127.0.0.1:{proxy_port}/v1/chat/completions",
            headers={"Authorization": "Bearer test-key"},
            json={
                "model": "gpt-4o-mini",
                "messages": [{"role": "user", "content": "Say hello."}],
            },
            timeout=10.0,
        )
        assert_true(chat.status_code == 200, "Proxy should pass through chat completions")
        assert_true(
            chat.json()["choices"][0]["message"]["content"] == "mock completion from upstream",
            "Unexpected chat completion payload",
        )
        assert_true(
            any(item["path"] == "/v1/models" for item in mock_server.requests),
            "Mock upstream should receive /v1/models",
        )
        assert_true(
            any(item["path"] == "/v1/chat/completions" for item in mock_server.requests),
            "Mock upstream should receive /v1/chat/completions",
        )
    finally:
        stop_process(proc)


def verify_codex_wrap(base_env: dict[str, str], project_dir: Path, log_dir: Path) -> None:
    port = CODEX_PORT
    run(
        ["headroom", "wrap", "codex", "--port", str(port), "--", "--help"],
        env=base_env,
        cwd=project_dir,
        timeout=120,
    )
    project_agents = project_dir / "AGENTS.md"
    global_agents = Path(base_env["HOME"]) / ".codex" / "AGENTS.md"
    assert_true(project_agents.exists(), "Codex wrap should create project AGENTS.md")
    assert_true(global_agents.exists(), "Codex wrap should create ~/.codex/AGENTS.md")
    assert_true(RTK_MARKER in project_agents.read_text(encoding="utf-8"), "Missing RTK marker")
    assert_true(
        RTK_MARKER in global_agents.read_text(encoding="utf-8"), "Missing global RTK marker"
    )

    entries = read_jsonl(log_dir / "codex.jsonl")
    assert_true(len(entries) > 0, "Codex shim should have been invoked")
    env_vars = entries[-1]["env"]
    assert_true(
        env_vars.get("OPENAI_BASE_URL") == f"http://127.0.0.1:{port}/v1",
        "Codex wrap should set OPENAI_BASE_URL",
    )
    assert_true(
        entries[-1]["probes"] == [{"url": f"http://127.0.0.1:{port}/v1/models", "status": 200}],
        "Codex shim should prove OPENAI_BASE_URL points at a live proxy",
    )


def verify_claude_wrap(base_env: dict[str, str], project_dir: Path, log_dir: Path) -> None:
    port = PROXY_PORT + 10
    run(
        ["headroom", "wrap", "claude", "--port", str(port), "--", "--help"],
        env=base_env,
        cwd=project_dir,
        timeout=120,
    )
    entries = read_jsonl(log_dir / "claude.jsonl")
    assert_true(len(entries) > 0, "Claude shim should have been invoked")
    env_vars = entries[-1]["env"]
    assert_true(
        env_vars.get("ANTHROPIC_BASE_URL") == f"http://127.0.0.1:{port}",
        "Claude wrap should set ANTHROPIC_BASE_URL",
    )
    assert_true(
        entries[-1]["probes"] == [{"url": f"http://127.0.0.1:{port}/health", "status": 200}],
        "Claude shim should prove ANTHROPIC_BASE_URL points at a live proxy",
    )


def verify_aider_wrap(base_env: dict[str, str], project_dir: Path, log_dir: Path) -> None:
    port = AIDER_PORT
    run(
        ["headroom", "wrap", "aider", "--port", str(port), "--", "--help"],
        env=base_env,
        cwd=project_dir,
        timeout=120,
    )
    conventions = project_dir / "CONVENTIONS.md"
    assert_true(conventions.exists(), "Aider wrap should create CONVENTIONS.md")
    assert_true(
        RTK_MARKER in conventions.read_text(encoding="utf-8"),
        "Aider wrap should inject RTK instructions",
    )

    entries = read_jsonl(log_dir / "aider.jsonl")
    assert_true(len(entries) > 0, "Aider shim should have been invoked")
    env_vars = entries[-1]["env"]
    assert_true(
        env_vars.get("OPENAI_API_BASE") == f"http://127.0.0.1:{port}/v1",
        "Aider wrap should set OPENAI_API_BASE",
    )
    assert_true(
        env_vars.get("ANTHROPIC_BASE_URL") == f"http://127.0.0.1:{port}",
        "Aider wrap should set ANTHROPIC_BASE_URL",
    )
    assert_true(
        entries[-1]["probes"]
        == [
            {"url": f"http://127.0.0.1:{port}/v1/models", "status": 200},
            {"url": f"http://127.0.0.1:{port}/health", "status": 200},
        ],
        "Aider shim should prove both configured base URLs point at a live proxy",
    )


def verify_cursor_wrap(base_env: dict[str, str], project_dir: Path) -> None:
    port = CURSOR_PORT
    proc = subprocess.Popen(
        ["headroom", "wrap", "cursor", "--port", str(port)],
        env=base_env,
        cwd=str(project_dir),
        stdout=subprocess.PIPE,
        stderr=subprocess.STDOUT,
        text=True,
        encoding="utf-8",
        errors="replace",
    )
    try:
        output = wait_for_output(proc, "Press Ctrl+C to stop the proxy.", timeout=30)
        assert_true(
            f"http://127.0.0.1:{port}/v1" in output,
            "Cursor wrap should print the OpenAI base URL override",
        )
        assert_true(
            f"http://127.0.0.1:{port}" in output,
            "Cursor wrap should print the Anthropic base URL override",
        )
        wait_for_http(f"http://127.0.0.1:{port}/health", timeout=15)
        cursorrules = project_dir / ".cursorrules"
        assert_true(cursorrules.exists(), "Cursor wrap should create .cursorrules")
        assert_true(
            RTK_MARKER in cursorrules.read_text(encoding="utf-8"),
            "Cursor wrap should inject RTK instructions",
        )
    finally:
        stop_process(proc)


def verify_openclaw_wrap(
    base_env: dict[str, str],
    project_dir: Path,
    plugin_dir: Path,
) -> None:
    port = OPENCLAW_PROXY_PORT
    gateway_proc: subprocess.Popen[str] | None = None
    run(
        [
            "headroom",
            "wrap",
            "openclaw",
            "--plugin-path",
            str(plugin_dir),
            "--proxy-port",
            str(port),
            "--startup-timeout-ms",
            "5000",
            "--verbose",
        ],
        env=base_env,
        cwd=project_dir,
        timeout=600,
    )
    dist_index = plugin_dir / "dist" / "index.js"
    assert_true(dist_index.exists(), "OpenClaw plugin build should produce dist/index.js")

    config_file = run(["openclaw", "config", "file"], env=base_env, cwd=project_dir, timeout=60)
    config_path_str = config_file.stdout.strip().splitlines()[-1].strip()
    if config_path_str.startswith("~/"):
        config_path = Path(base_env["HOME"]) / config_path_str[2:]
    else:
        config_path = Path(config_path_str)
    assert_true(config_path.exists(), "OpenClaw should create a config file")

    state = json.loads(config_path.read_text(encoding="utf-8"))
    if state.get("gateway", {}).get("mode") != "local":
        run(
            [
                "openclaw",
                "config",
                "set",
                "gateway.mode",
                json.dumps("local"),
                "--strict-json",
            ],
            env=base_env,
            cwd=project_dir,
            timeout=60,
        )
        state = json.loads(config_path.read_text(encoding="utf-8"))

    try:
        try:
            wait_for_command_success(
                ["openclaw", "health"], env=base_env, cwd=project_dir, timeout=5
            )
        except RuntimeError:
            gateway_proc = start_openclaw_gateway(base_env, project_dir)
            try:
                wait_for_command_success(
                    ["openclaw", "health"], env=base_env, cwd=project_dir, timeout=30
                )
            except RuntimeError as exc:
                gateway_output = ""
                if gateway_proc.stdout is not None:
                    gateway_output = gateway_proc.stdout.read()
                raise RuntimeError(f"{exc}\nGateway output:\n{gateway_output}") from exc

        entry = state["plugins"]["entries"]["headroom"]
        assert_true(entry["enabled"] is True, "OpenClaw wrap should enable the plugin")
        assert_true(entry["config"]["proxyPort"] == port, "OpenClaw wrap should set proxy port")
        assert_true(
            entry["config"].get("autoStart", True) is True,
            "OpenClaw wrap should leave autoStart enabled",
        )
        assert_true(
            state["gateway"]["mode"] == "local",
            "OpenClaw e2e bootstrap should set gateway.mode=local",
        )
        assert_true(
            state["plugins"]["slots"]["contextEngine"] == "headroom",
            "OpenClaw wrap should set the context engine slot",
        )
    finally:
        if gateway_proc is not None:
            stop_process(gateway_proc)
        stop_openclaw_gateway(base_env, project_dir)

    run(["headroom", "unwrap", "openclaw"], env=base_env, cwd=project_dir, timeout=120)
    state = json.loads(config_path.read_text(encoding="utf-8"))
    assert_true(
        state["plugins"]["slots"]["contextEngine"] == "legacy",
        "OpenClaw unwrap should restore the context engine slot",
    )


def main() -> None:
    verify_installs()
    with tempfile.TemporaryDirectory(
        prefix="headroom-wrap-e2e-", ignore_cleanup_errors=True
    ) as tmp_dir_str:
        tmp_dir = Path(tmp_dir_str)
        home_dir = tmp_dir / "home"
        project_dir = tmp_dir / "project"
        shim_dir = tmp_dir / "shim-bin"
        log_dir = tmp_dir / "logs"

        for path in (home_dir, project_dir, shim_dir, log_dir):
            path.mkdir(parents=True, exist_ok=True)
        create_shims(shim_dir)

        mock_server, mock_thread = start_mock_server(19001)
        base_env = os.environ.copy()
        base_env.update(
            {
                "HOME": str(home_dir),
                "PATH": f"{shim_dir}{os.pathsep}{base_env['PATH']}",
                "HEADROOM_E2E_LOG_DIR": str(log_dir),
                "OPENAI_TARGET_API_URL": "http://127.0.0.1:19001/v1",
            }
        )

        try:
            verify_proxy_round_trip(base_env, mock_server)
            verify_claude_wrap(base_env, project_dir, log_dir)
            verify_codex_wrap(base_env, project_dir, log_dir)
            verify_aider_wrap(base_env, project_dir, log_dir)
            verify_cursor_wrap(base_env, project_dir)
            local_plugin_dir = prepare_local_openclaw_plugin(base_env, tmp_dir)
            verify_openclaw_wrap(base_env, project_dir, local_plugin_dir)
        finally:
            mock_server.shutdown()
            mock_thread.join(timeout=5)

    log("All Docker wrap e2e checks passed.")


if __name__ == "__main__":
    main()