codex/sdk/python/_runtime_setup.py

from __future__ import annotations

import importlib
import importlib.metadata
import importlib.util
import json
import os
import platform
import re
import shutil
import subprocess
import sys
import tarfile
import tempfile
import urllib.error
import urllib.request
import zipfile
from pathlib import Path

PACKAGE_NAME = "openai-codex-cli-bin"
SDK_PACKAGE_NAME = "openai-codex"
REPO_SLUG = "openai/codex"


class RuntimeSetupError(RuntimeError):
    pass


def pinned_runtime_version() -> str:
    """Return the exact runtime version pinned by the SDK package dependency."""
    source_pin = _source_tree_runtime_dependency_version()
    if source_pin is not None:
        return _normalized_package_version(source_pin)

    try:
        installed_pin = _installed_sdk_runtime_dependency_version()
    except importlib.metadata.PackageNotFoundError as exc:
        raise RuntimeSetupError(
            f"Unable to resolve {SDK_PACKAGE_NAME} metadata for runtime pinning."
        ) from exc
    if installed_pin is None:
        raise RuntimeSetupError(
            f"Unable to resolve {PACKAGE_NAME} dependency pin from {SDK_PACKAGE_NAME}."
        )
    return _normalized_package_version(installed_pin)


def ensure_runtime_package_installed(
    python_executable: str | Path,
    sdk_python_dir: Path,
    install_target: Path | None = None,
) -> str:
    requested_version = pinned_runtime_version()
    installed_version = None
    if install_target is None:
        installed_version = _installed_runtime_version(python_executable)
    normalized_requested = _normalized_package_version(requested_version)

    if (
        installed_version is not None
        and _normalized_package_version(installed_version) == normalized_requested
    ):
        return requested_version

    with tempfile.TemporaryDirectory(prefix="codex-python-runtime-") as temp_root_str:
        temp_root = Path(temp_root_str)
        archive_path = _download_release_archive(requested_version, temp_root)
        runtime_package_dir = _extract_runtime_package_dir(
            archive_path,
            temp_root,
            requested_version,
        )
        staged_runtime_dir = _stage_runtime_package(
            sdk_python_dir,
            requested_version,
            runtime_package_dir,
            temp_root / "runtime-stage",
        )
        _install_runtime_package(python_executable, staged_runtime_dir, install_target)

    if install_target is not None:
        return requested_version

    if Path(python_executable).resolve() == Path(sys.executable).resolve():
        importlib.invalidate_caches()

    installed_version = _installed_runtime_version(python_executable)
    if (
        installed_version is None
        or _normalized_package_version(installed_version) != normalized_requested
    ):
        raise RuntimeSetupError(
            f"Expected {PACKAGE_NAME} {requested_version} in {python_executable}, "
            f"but found {installed_version!r} after installation."
        )
    return requested_version


def platform_asset_name() -> str:
    system = platform.system().lower()
    machine = platform.machine().lower()

    if system == "darwin":
        if machine in {"arm64", "aarch64"}:
            return "codex-aarch64-apple-darwin.tar.gz"
        if machine in {"x86_64", "amd64"}:
            return "codex-x86_64-apple-darwin.tar.gz"
    elif system == "linux":
        if machine in {"aarch64", "arm64"}:
            return "codex-aarch64-unknown-linux-musl.tar.gz"
        if machine in {"x86_64", "amd64"}:
            return "codex-x86_64-unknown-linux-musl.tar.gz"
    elif system == "windows":
        if machine in {"aarch64", "arm64"}:
            return "codex-aarch64-pc-windows-msvc.exe.zip"
        if machine in {"x86_64", "amd64"}:
            return "codex-x86_64-pc-windows-msvc.exe.zip"

    raise RuntimeSetupError(
        f"Unsupported runtime artifact platform: system={platform.system()!r}, "
        f"machine={platform.machine()!r}"
    )


def runtime_binary_name() -> str:
    return "codex.exe" if platform.system().lower() == "windows" else "codex"


def _installed_runtime_version(python_executable: str | Path) -> str | None:
    snippet = (
        "import importlib.metadata, json, sys\n"
        "try:\n"
        "    from codex_cli_bin import bundled_codex_path\n"
        "    bundled_codex_path()\n"
        f"    print(json.dumps({{'version': importlib.metadata.version({PACKAGE_NAME!r})}}))\n"
        "except Exception:\n"
        "    sys.exit(1)\n"
    )
    result = subprocess.run(
        [str(python_executable), "-c", snippet],
        text=True,
        capture_output=True,
        check=False,
    )
    if result.returncode != 0:
        return None
    return json.loads(result.stdout)["version"]


def _release_metadata(version: str) -> dict[str, object]:
    release_tag = _release_tag(version)
    url = f"https://api.github.com/repos/{REPO_SLUG}/releases/tags/{release_tag}"
    token = _github_token()
    attempts = [True, False] if token is not None else [False]
    last_error: urllib.error.HTTPError | None = None

    for include_auth in attempts:
        headers = {
            "Accept": "application/vnd.github+json",
            "User-Agent": "codex-python-runtime-setup",
        }
        if include_auth and token is not None:
            headers["Authorization"] = f"Bearer {token}"

        request = urllib.request.Request(url, headers=headers)
        try:
            with urllib.request.urlopen(request) as response:
                return json.load(response)
        except urllib.error.HTTPError as exc:
            last_error = exc
            if include_auth and exc.code == 401:
                continue
            break

    assert last_error is not None
    raise RuntimeSetupError(
        f"Failed to resolve release metadata for {release_tag} from {REPO_SLUG}: "
        f"{last_error.code} {last_error.reason}"
    ) from last_error


def _download_release_archive(version: str, temp_root: Path) -> Path:
    asset_name = platform_asset_name()
    archive_path = temp_root / asset_name
    release_tag = _release_tag(version)

    browser_download_url = (
        f"https://github.com/{REPO_SLUG}/releases/download/{release_tag}/{asset_name}"
    )
    request = urllib.request.Request(
        browser_download_url,
        headers={"User-Agent": "codex-python-runtime-setup"},
    )
    try:
        with urllib.request.urlopen(request) as response, archive_path.open("wb") as fh:
            shutil.copyfileobj(response, fh)
        return archive_path
    except urllib.error.HTTPError:
        pass

    metadata = _release_metadata(version)
    assets = metadata.get("assets")
    if not isinstance(assets, list):
        raise RuntimeSetupError(f"Release {release_tag} returned malformed assets metadata.")
    asset = next(
        (item for item in assets if isinstance(item, dict) and item.get("name") == asset_name),
        None,
    )
    if asset is None:
        raise RuntimeSetupError(
            f"Release {release_tag} does not contain asset {asset_name} for this platform."
        )

    api_url = asset.get("url")
    if not isinstance(api_url, str):
        api_url = None

    if api_url is not None:
        token = _github_token()
        if token is not None:
            request = urllib.request.Request(
                api_url,
                headers=_github_api_headers("application/octet-stream"),
            )
            try:
                with (
                    urllib.request.urlopen(request) as response,
                    archive_path.open("wb") as fh,
                ):
                    shutil.copyfileobj(response, fh)
                return archive_path
            except urllib.error.HTTPError:
                pass

    if shutil.which("gh") is None:
        raise RuntimeSetupError(
            f"Unable to download {asset_name} for rust-v{version}. "
            "Provide GH_TOKEN/GITHUB_TOKEN or install/authenticate GitHub CLI."
        )

    try:
        subprocess.run(
            [
                "gh",
                "release",
                "download",
                release_tag,
                "--repo",
                REPO_SLUG,
                "--pattern",
                asset_name,
                "--dir",
                str(temp_root),
            ],
            check=True,
            text=True,
            capture_output=True,
        )
    except subprocess.CalledProcessError as exc:
        raise RuntimeSetupError(
            f"gh release download failed for {release_tag} asset {asset_name}.\n"
            f"STDOUT:\n{exc.stdout}\nSTDERR:\n{exc.stderr}"
        ) from exc
    return archive_path


def _extract_runtime_package_dir(archive_path: Path, temp_root: Path, version: str) -> Path:
    extract_dir = temp_root / "extracted"
    extract_dir.mkdir(parents=True, exist_ok=True)
    if archive_path.name.endswith(".tar.gz"):
        with tarfile.open(archive_path, "r:gz") as tar:
            try:
                tar.extractall(extract_dir, filter="data")
            except TypeError:
                tar.extractall(extract_dir)
    elif archive_path.suffix == ".zip":
        with zipfile.ZipFile(archive_path) as zip_file:
            zip_file.extractall(extract_dir)
    else:
        raise RuntimeSetupError(f"Unsupported release archive format: {archive_path.name}")

    metadata_path = next(
        (path for path in extract_dir.rglob("codex-package.json") if path.is_file()),
        None,
    )
    if metadata_path is not None:
        return metadata_path.parent

    binary_name = runtime_binary_name()
    archive_stem = archive_path.name.removesuffix(".tar.gz").removesuffix(".zip")
    candidates = [
        path
        for path in extract_dir.rglob("*")
        if path.is_file()
        and (
            path.name == binary_name or path.name == archive_stem or path.name.startswith("codex-")
        )
    ]
    if not candidates:
        raise RuntimeSetupError(
            f"Failed to find {binary_name} in extracted runtime archive {archive_path.name}."
        )
    return _package_legacy_runtime_binary(candidates[0], temp_root, version)


def _package_legacy_runtime_binary(runtime_binary: Path, temp_root: Path, version: str) -> Path:
    package_dir = temp_root / "runtime-package"
    bin_dir = package_dir / "bin"
    resources_dir = package_dir / "codex-resources"
    path_dir = package_dir / "codex-path"
    bin_dir.mkdir(parents=True)
    resources_dir.mkdir()
    path_dir.mkdir()
    shutil.copy2(runtime_binary, bin_dir / runtime_binary_name())
    (package_dir / "codex-package.json").write_text(
        json.dumps(
            {
                "layoutVersion": 1,
                "version": version,
                "entrypoint": f"bin/{runtime_binary_name()}",
                "resourcesDir": "codex-resources",
                "pathDir": "codex-path",
            },
            indent=2,
        )
        + "\n"
    )
    return package_dir


def _stage_runtime_package(
    sdk_python_dir: Path,
    runtime_version: str,
    runtime_package_dir: Path,
    staging_dir: Path,
) -> Path:
    script_module = _load_update_script_module(sdk_python_dir)
    return script_module.stage_python_runtime_package(  # type: ignore[no-any-return]
        staging_dir,
        runtime_version,
        runtime_package_dir.resolve(),
    )


def _install_runtime_package(
    python_executable: str | Path,
    staged_runtime_dir: Path,
    install_target: Path | None,
) -> None:
    args = [
        str(python_executable),
        "-m",
        "pip",
        "install",
        "--force-reinstall",
        "--no-deps",
    ]
    if install_target is not None:
        install_target.mkdir(parents=True, exist_ok=True)
        args.extend(["--target", str(install_target)])
    args.append(str(staged_runtime_dir))
    try:
        subprocess.run(
            args,
            check=True,
            text=True,
            capture_output=True,
        )
    except subprocess.CalledProcessError as exc:
        raise RuntimeSetupError(
            f"Failed to install {PACKAGE_NAME} into {python_executable} from {staged_runtime_dir}.\n"
            f"STDOUT:\n{exc.stdout}\nSTDERR:\n{exc.stderr}"
        ) from exc


def _load_update_script_module(sdk_python_dir: Path):
    script_path = sdk_python_dir / "scripts" / "update_sdk_artifacts.py"
    spec = importlib.util.spec_from_file_location("update_sdk_artifacts", script_path)
    if spec is None or spec.loader is None:
        raise RuntimeSetupError(f"Failed to load {script_path}")
    module = importlib.util.module_from_spec(spec)
    sys.modules[spec.name] = module
    spec.loader.exec_module(module)
    return module


def _github_api_headers(accept: str) -> dict[str, str]:
    headers = {
        "Accept": accept,
        "User-Agent": "codex-python-runtime-setup",
    }
    token = _github_token()
    if token is not None:
        headers["Authorization"] = f"Bearer {token}"
    return headers


def _github_token() -> str | None:
    for env_name in ("GH_TOKEN", "GITHUB_TOKEN"):
        token = os.environ.get(env_name)
        if token:
            return token
    return None


def _normalized_package_version(version: str) -> str:
    normalized = version.strip()
    if normalized.startswith("rust-v"):
        normalized = normalized.removeprefix("rust-v")
    elif normalized.startswith("v"):
        normalized = normalized.removeprefix("v")

    normalized = re.sub(r"-alpha\.?([0-9]+)$", r"a\1", normalized)
    normalized = re.sub(r"-beta\.?([0-9]+)$", r"b\1", normalized)
    normalized = re.sub(r"-rc\.?([0-9]+)$", r"rc\1", normalized)
    return normalized


def _codex_release_version(version: str) -> str:
    normalized = _normalized_package_version(version)
    match = re.fullmatch(r"([0-9]+(?:\.[0-9]+)*)(a|b|rc)([0-9]+)", normalized)
    if match is None:
        return normalized

    base, prerelease, number = match.groups()
    prerelease_name = {"a": "alpha", "b": "beta", "rc": "rc"}[prerelease]
    return f"{base}-{prerelease_name}.{number}"


def _release_tag(version: str) -> str:
    return f"rust-v{_codex_release_version(version)}"


def _source_tree_runtime_dependency_version() -> str | None:
    """Read the runtime dependency pin when the SDK is running from a checkout."""
    pyproject_path = Path(__file__).resolve().parent / "pyproject.toml"
    if not pyproject_path.exists():
        return None

    match = re.search(_runtime_dependency_pin_pattern(), pyproject_path.read_text())
    if match is None:
        return None
    return match.group(1)


def _installed_sdk_runtime_dependency_version() -> str | None:
    """Read the runtime dependency pin from installed package metadata."""
    requirements = importlib.metadata.requires(SDK_PACKAGE_NAME) or []
    for requirement in requirements:
        match = re.search(_runtime_dependency_pin_pattern(), requirement)
        if match is not None:
            return match.group(1)
    return None


def _runtime_dependency_pin_pattern() -> str:
    """Match the exact runtime dependency pin in TOML and wheel metadata."""
    return rf'{re.escape(PACKAGE_NAME)}\s*==\s*"?([^",;\s]+)"?'


__all__ = [
    "PACKAGE_NAME",
    "SDK_PACKAGE_NAME",
    "RuntimeSetupError",
    "ensure_runtime_package_installed",
    "pinned_runtime_version",
    "platform_asset_name",
]