SkillCompiler/data/skills-bench/skillsbench_agentbeats/deploy_bundle.py

"""Normalize AgentBeats deployment artifacts into an operator-ready bundle."""

from __future__ import annotations

import argparse
import json
import re
import sys
from collections.abc import Mapping
from dataclasses import dataclass
from pathlib import Path
from typing import Any

from skillsbench_agentbeats.config import DEFAULT_PUBLIC_TASK_SET, repo_root_from_file
from skillsbench_agentbeats.task_sets import load_task_set_manifest

DEPLOY_BUNDLE_SCHEMA_VERSION = "skillsbench.agentbeats.deploy_bundle.v1"
PUBLIC_IMAGE_PLATFORM = "linux/amd64"

COMPONENTS = ("green", "worker", "purple")
COMPONENT_MANIFEST_PATHS = {
    "green": "integrations/agentbeats/green_agent/amber-manifest.json5",
    "worker": "integrations/agentbeats/worker/amber-manifest.json5",
    "purple": "integrations/agentbeats/agent_under_test/amber-manifest.json5",
}
COMPONENT_REPOSITORIES = {
    "green": "ghcr.io/benchflow-ai/skillsbench-agentbeats-green",
    "worker": "ghcr.io/benchflow-ai/skillsbench-agentbeats-worker",
    "purple": "ghcr.io/benchflow-ai/skillsbench-agentbeats-purple",
}

SHA256_RE = re.compile(r"^sha256:[0-9a-f]{64}$")
GIT_SHA_RE = re.compile(r"^[0-9a-f]{7,64}$")


@dataclass(frozen=True)
class DeployBundleIssue:
    path: str
    message: str

    def format(self) -> str:
        return f"{self.path}: {self.message}"


class DeployBundleError(ValueError):
    def __init__(self, issues: list[DeployBundleIssue]) -> None:
        self.issues = issues
        super().__init__("\n".join(issue.format() for issue in issues))


def build_deploy_bundle(
    *,
    repo_root: Path | None = None,
    branch_images: Mapping[str, Any],
    task_environment_images: Mapping[str, Any],
    task_set: str = DEFAULT_PUBLIC_TASK_SET,
) -> dict[str, Any]:
    """Build a normalized deployment bundle from GitHub Actions artifacts."""

    root = repo_root or repo_root_from_file()
    issues = validate_deploy_bundle_inputs(
        repo_root=root,
        branch_images=branch_images,
        task_environment_images=task_environment_images,
        task_set=task_set,
    )
    if issues:
        raise DeployBundleError(issues)

    manifest = load_task_set_manifest(task_set, repo_root=root)
    if manifest is None:
        raise DeployBundleError([DeployBundleIssue("$.task_set", f"task-set manifest {task_set!r} was not found")])

    branch_image_map = _as_mapping(branch_images["images"])
    task_ids = _manifest_task_ids(manifest)
    worker_prebuilt_images = {task_id: _as_string(task_environment_images[task_id]) for task_id in task_ids}
    components = {component: _component_record(component, _as_string(branch_image_map[component])) for component in COMPONENTS}
    return {
        "schema_version": DEPLOY_BUNDLE_SCHEMA_VERSION,
        "source_revision": branch_images["source_revision"],
        "image_tag": branch_images["image_tag"],
        "task_set": {
            "task_set": manifest["task_set"],
            "task_count": manifest["task_count"],
            "task_set_digest": manifest["task_set_digest"],
        },
        "components": components,
        "public_readiness_images": _public_readiness_images(components),
        "worker_prebuilt_images": worker_prebuilt_images,
        "public_readiness_task_environment_images": {
            "images": [
                {
                    "task_id": task_id,
                    "image": image,
                    "image_digest": _image_digest(image),
                    "image_platform": PUBLIC_IMAGE_PLATFORM,
                }
                for task_id, image in worker_prebuilt_images.items()
            ],
        },
    }


def validate_deploy_bundle_inputs(
    *,
    repo_root: Path | None = None,
    branch_images: Mapping[str, Any],
    task_environment_images: Mapping[str, Any],
    task_set: str = DEFAULT_PUBLIC_TASK_SET,
) -> list[DeployBundleIssue]:
    """Validate branch runtime images and prebuilt task environment images."""

    root = repo_root or repo_root_from_file()
    issues: list[DeployBundleIssue] = []
    manifest = load_task_set_manifest(task_set, repo_root=root)
    if manifest is None:
        issues.append(DeployBundleIssue("$.task_set", f"task-set manifest {task_set!r} was not found"))
        return issues
    _validate_branch_images(branch_images, issues)
    _validate_task_environment_images(task_environment_images, _manifest_task_ids(manifest), issues, task_set=task_set)
    return issues


def load_json_mapping(path: Path) -> dict[str, Any]:
    payload = json.loads(path.read_text())
    if not isinstance(payload, dict):
        raise ValueError(f"{path} must contain a JSON object")
    return payload


def _validate_branch_images(branch_images: Mapping[str, Any], issues: list[DeployBundleIssue]) -> None:
    source_revision = branch_images.get("source_revision")
    if not isinstance(source_revision, str) or GIT_SHA_RE.fullmatch(source_revision) is None:
        issues.append(DeployBundleIssue("$.branch_images.source_revision", "must be a git SHA or abbreviated git SHA"))
    image_tag = branch_images.get("image_tag")
    if not isinstance(image_tag, str) or not image_tag:
        issues.append(DeployBundleIssue("$.branch_images.image_tag", "must be a non-empty string"))
    images = branch_images.get("images")
    if not isinstance(images, Mapping):
        issues.append(DeployBundleIssue("$.branch_images.images", "must be an object"))
        return
    missing_components = sorted(set(COMPONENTS) - set(images))
    extra_components = sorted(set(images) - set(COMPONENTS))
    if missing_components:
        issues.append(DeployBundleIssue("$.branch_images.images", f"missing component image(s): {missing_components}"))
    if extra_components:
        issues.append(DeployBundleIssue("$.branch_images.images", f"unexpected component image(s): {extra_components}"))
    for component in COMPONENTS:
        image = images.get(component)
        if not isinstance(image, str):
            issues.append(DeployBundleIssue(f"$.branch_images.images.{component}", "must be a string"))
            continue
        _validate_digest_pinned_image(image, f"$.branch_images.images.{component}", issues)
        expected_repository = COMPONENT_REPOSITORIES[component]
        if not image.startswith(f"{expected_repository}@"):
            issues.append(DeployBundleIssue(f"$.branch_images.images.{component}", f"must use repository {expected_repository!r}"))


def _validate_task_environment_images(
    task_environment_images: Mapping[str, Any],
    expected_task_ids: list[str],
    issues: list[DeployBundleIssue],
    *,
    task_set: str,
) -> None:
    actual_task_ids = {key for key in task_environment_images if isinstance(key, str)}
    expected_task_id_set = set(expected_task_ids)
    missing_task_ids = sorted(expected_task_id_set - actual_task_ids)
    extra_task_ids = sorted(actual_task_ids - expected_task_id_set)
    if missing_task_ids:
        issues.append(
            DeployBundleIssue(
                "$.task_environment_images",
                f"missing {len(missing_task_ids)} {task_set} task environment image(s): {missing_task_ids[:5]}",
            )
        )
    if extra_task_ids:
        issues.append(
            DeployBundleIssue(
                "$.task_environment_images",
                f"contains task ids outside {task_set}: {extra_task_ids[:5]}",
            )
        )
    for task_id, image in task_environment_images.items():
        if not isinstance(task_id, str):
            issues.append(DeployBundleIssue("$.task_environment_images", "task ids must be strings"))
            continue
        if not isinstance(image, str):
            issues.append(DeployBundleIssue(f"$.task_environment_images.{task_id}", "must be a string"))
            continue
        _validate_digest_pinned_image(image, f"$.task_environment_images.{task_id}", issues)


def _validate_digest_pinned_image(image: str, path: str, issues: list[DeployBundleIssue]) -> None:
    digest = _image_digest(image)
    if digest is None:
        issues.append(DeployBundleIssue(path, "must be pinned as IMAGE@sha256:<64 hex chars>"))
    elif SHA256_RE.fullmatch(digest) is None:
        issues.append(DeployBundleIssue(path, "must use a valid sha256 digest"))


def _component_record(component: str, image: str) -> dict[str, str]:
    digest = _image_digest(image)
    if digest is None:
        raise ValueError("component image was not validated")
    return {
        "image": image,
        "image_digest": digest,
        "image_platform": PUBLIC_IMAGE_PLATFORM,
        "manifest_path": COMPONENT_MANIFEST_PATHS[component],
    }


def _public_readiness_images(components: Mapping[str, Mapping[str, str]]) -> dict[str, str]:
    return {
        "green_image": components["green"]["image"],
        "green_image_digest": components["green"]["image_digest"],
        "green_image_platform": components["green"]["image_platform"],
        "worker_image": components["worker"]["image"],
        "worker_image_digest": components["worker"]["image_digest"],
        "worker_image_platform": components["worker"]["image_platform"],
        "purple_image": components["purple"]["image"],
        "purple_image_digest": components["purple"]["image_digest"],
        "purple_image_platform": components["purple"]["image_platform"],
    }


def _image_digest(image: str) -> str | None:
    if "@sha256:" not in image:
        return None
    return image.rsplit("@", 1)[1]


def _manifest_task_ids(manifest: Mapping[str, Any]) -> list[str]:
    tasks = manifest.get("tasks")
    if not isinstance(tasks, list):
        return []
    task_ids = [task.get("task_id") for task in tasks if isinstance(task, Mapping)]
    return [task_id for task_id in task_ids if isinstance(task_id, str)]


def _as_mapping(value: Any) -> Mapping[str, Any]:
    if not isinstance(value, Mapping):
        raise TypeError("expected mapping")
    return value


def _as_string(value: Any) -> str:
    if not isinstance(value, str):
        raise TypeError("expected string")
    return value


def _main() -> None:
    parser = argparse.ArgumentParser(description="Build an AgentBeats deployment bundle from published image artifacts.")
    parser.add_argument("--repo-root", type=Path, default=repo_root_from_file())
    parser.add_argument("--branch-images", type=Path, required=True, help="agentbeats-branch-image-digests.json artifact.")
    parser.add_argument("--task-environment-images", type=Path, required=True, help="prebuilt task environment image map artifact.")
    parser.add_argument("--task-set", default=DEFAULT_PUBLIC_TASK_SET)
    parser.add_argument("--output", type=Path, required=True)
    args = parser.parse_args()

    branch_images = load_json_mapping(args.branch_images)
    task_environment_images = load_json_mapping(args.task_environment_images)
    try:
        bundle = build_deploy_bundle(
            repo_root=args.repo_root,
            branch_images=branch_images,
            task_environment_images=task_environment_images,
            task_set=args.task_set,
        )
    except DeployBundleError as exc:
        for issue in exc.issues:
            print(issue.format(), file=sys.stderr)
        parser.exit(1)
    args.output.parent.mkdir(parents=True, exist_ok=True)
    args.output.write_text(json.dumps(bundle, indent=2, sort_keys=True) + "\n")


if __name__ == "__main__":
    _main()