agent_management/scripts/e2e_complete_project.py

#!/usr/bin/env python3
"""End-to-end sub-mode Runtime check for a complete generated project.

This script is intentionally outside the normal unit-test suite because it
requires a live agent-manager Runtime, Kubernetes, and a model gateway.
"""

from __future__ import annotations

import argparse
import hashlib
import json
import os
import re
import shutil
import subprocess
import sys
import tempfile
import time
import urllib.error
import urllib.request
from pathlib import Path
from typing import Any


DEFAULT_OBJECTIVE = """Return ONLY valid JSON:
{"files":[{"path":"...","content":"..."}],"run_tests":"python -m unittest discover -s tests -v","smoke_test":"python -m textstats_cli samples/example.txt --json"}.

Build a complete, compact Python stdlib project named textstats_cli.
Include exactly these files:
- pyproject.toml
- README.md
- textstats_cli/__main__.py
- textstats_cli/core.py
- tests/test_core.py
- samples/example.txt

Features:
- CLI accepts a text file path.
- --json outputs JSON.
- Default output is human readable.
- Report line count, word count, character count, top 5 words excluding common stopwords, and estimated reading time.
- Tests must cover counting, stopword filtering, JSON-safe result shape, and missing-file error handling.

Constraints:
- Use only the Python standard library.
- Keep the project small enough for one response.
- No placeholders.
- No markdown fences.
- No prose outside the JSON object.
"""


TERMINAL_STATUSES = {"completed", "failed", "stopped"}


def http_json(method: str, url: str, payload: dict[str, Any] | None = None, headers: dict[str, str] | None = None) -> Any:
    data = None if payload is None else json.dumps(payload).encode("utf-8")
    request_headers = {"Content-Type": "application/json", **(headers or {})}
    request = urllib.request.Request(url, data=data, headers=request_headers, method=method)
    try:
        with urllib.request.urlopen(request, timeout=60) as response:
            return json.loads(response.read().decode("utf-8"))
    except urllib.error.HTTPError as exc:
        body = exc.read().decode("utf-8", errors="replace")
        raise RuntimeError(f"{method} {url} failed with HTTP {exc.code}: {body}") from exc


def http_bytes(url: str) -> bytes:
    with urllib.request.urlopen(url, timeout=120) as response:
        return response.read()


def extract_json_object(text: str) -> dict[str, Any]:
    raw = text.strip()
    fenced = re.search(r"```(?:json)?\s*(\{.*?\})\s*```", raw, re.S)
    if fenced:
        raw = fenced.group(1)
    else:
        start = raw.find("{")
        end = raw.rfind("}")
        if start < 0 or end <= start:
            raise ValueError("artifact does not contain a JSON object")
        raw = raw[start : end + 1]
    return json.loads(raw)


def safe_write_project(project_dir: Path, files: list[dict[str, Any]]) -> None:
    project_root = project_dir.resolve()
    for item in files:
        relative_path = item.get("path")
        content = item.get("content")
        if not isinstance(relative_path, str) or not relative_path:
            raise ValueError(f"invalid file path in artifact: {item!r}")
        if not isinstance(content, str):
            raise ValueError(f"invalid content for {relative_path}")
        target = (project_root / relative_path).resolve()
        if project_root not in target.parents and target != project_root:
            raise ValueError(f"artifact attempted to write outside project: {relative_path}")
        target.parent.mkdir(parents=True, exist_ok=True)
        target.write_text(content, encoding="utf-8")


def run_command(command: str, cwd: Path) -> subprocess.CompletedProcess[str]:
    env = os.environ.copy()
    shim_dir = None
    if shutil.which("python") is None:
        shim_dir = Path(tempfile.mkdtemp(prefix="heicode-python-shim-"))
        (shim_dir / "python").symlink_to(sys.executable)
        env["PATH"] = str(shim_dir) + os.pathsep + env.get("PATH", "")
    return subprocess.run(
        command,
        cwd=str(cwd),
        env=env,
        shell=True,
        text=True,
        stdout=subprocess.PIPE,
        stderr=subprocess.STDOUT,
        timeout=120,
        check=False,
    )


def build_payload(args: argparse.Namespace) -> dict[str, Any]:
    return {
        "orchestration_plan": {
            "sub_mode": "agile",
            "objective": args.objective,
            "user_context": {"user_id": args.user_id},
            "agents": [
                {
                    "role": "backend",
                    "template": "a2a_litellm_agent",
                    "model": args.model,
                    "capabilities": ["code", "test"],
                }
            ],
            "agile_context": {"max_iterations": 1, "stage": "development"},
            "budget": {"max_duration_sec": args.timeout_seconds, "max_tokens": args.max_tokens},
            "billing_context": {
                "provider": "newapi",
                "default_model_id": args.model,
                "model_gateway_url": args.model_gateway_url,
                "api_format": args.api_format,
                "stream": args.stream,
                "timeout_sec": args.model_timeout_seconds,
                "max_tokens": args.max_tokens,
            },
            "metadata": {"correlation_id": args.correlation_id},
        },
        "callback": {"url": args.callback_url, "method": "POST"},
    }


def poll_swarm(base_url: str, swarm_id: str, timeout_seconds: int, poll_interval: int) -> dict[str, Any]:
    deadline = time.time() + timeout_seconds
    last_line = None
    while time.time() < deadline:
        status = http_json("GET", f"{base_url}/api/swarms/{swarm_id}")
        line = {
            "status": status.get("status"),
            "phase": status.get("phase"),
            "progress": status.get("progress"),
            "artifact_count": len(status.get("artifacts") or []),
            "tokens_used": (status.get("metrics") or {}).get("tokens_used"),
        }
        if line != last_line:
            print("status:", json.dumps(line, ensure_ascii=False))
            last_line = line
        if status.get("status") in TERMINAL_STATUSES:
            return status
        time.sleep(poll_interval)
    raise TimeoutError(f"swarm {swarm_id} did not finish within {timeout_seconds}s")


def assert_runtime_observability(base_url: str, swarm_id: str, status: dict[str, Any]) -> None:
    metrics = status.get("metrics") or {}
    tokens_used = int(metrics.get("tokens_used") or 0)
    if tokens_used <= 0:
        raise AssertionError(f"Runtime tokens_used must be > 0, got {tokens_used}")

    logs = http_json("GET", f"{base_url}/api/swarms/{swarm_id}/logs")
    request_ids: list[str] = []
    usage_totals: list[int] = []
    for agent in logs.get("agents") or []:
        for message in agent.get("messages") or []:
            if message.get("newapi_request_id"):
                request_ids.append(message["newapi_request_id"])
            usage = message.get("model_usage") or {}
            if usage.get("total_tokens"):
                usage_totals.append(int(usage["total_tokens"]))
    if not request_ids:
        raise AssertionError("Runtime logs must include at least one NewAPI request_id")
    if not usage_totals:
        raise AssertionError("Runtime logs must include model usage with total_tokens")
    print("observability:", json.dumps({"request_ids": request_ids, "usage_totals": usage_totals}, ensure_ascii=False))


def run_project_validation(base_url: str, status: dict[str, Any], output_dir: Path) -> None:
    artifacts = status.get("artifacts") or []
    if status.get("status") != "completed":
        raise AssertionError(f"swarm did not complete: {status.get('error_message')}")
    if not artifacts:
        raise AssertionError("completed swarm returned no artifacts")

    artifact = artifacts[0]
    download_path = (artifact.get("metadata") or {}).get("download_path")
    if not download_path:
        raise AssertionError("artifact metadata is missing download_path")

    artifact_bytes = http_bytes(f"{base_url}{download_path}")
    artifact_hash = "sha256:" + hashlib.sha256(artifact_bytes).hexdigest()
    expected_hash = (artifact.get("metadata") or {}).get("content_hash")
    if expected_hash and artifact_hash != expected_hash:
        raise AssertionError(f"artifact hash mismatch: expected {expected_hash}, got {artifact_hash}")

    artifact_text = artifact_bytes.decode("utf-8")
    artifact_json = extract_json_object(artifact_text)
    files = artifact_json.get("files")
    if not isinstance(files, list) or len(files) < 5:
        raise AssertionError("artifact must contain a multi-file project")

    project_dir = output_dir / "project"
    if project_dir.exists():
        shutil.rmtree(project_dir)
    project_dir.mkdir(parents=True)
    safe_write_project(project_dir, files)

    required_paths = {
        "pyproject.toml",
        "README.md",
        "textstats_cli/__main__.py",
        "textstats_cli/core.py",
        "tests/test_core.py",
        "samples/example.txt",
    }
    actual_paths = {str(path.relative_to(project_dir)) for path in project_dir.rglob("*") if path.is_file()}
    missing = sorted(required_paths - actual_paths)
    if missing:
        raise AssertionError(f"generated project missing required files: {missing}")

    for label, command in (
        ("run_tests", artifact_json.get("run_tests")),
        ("smoke_test", artifact_json.get("smoke_test")),
    ):
        if not isinstance(command, str) or not command.strip():
            raise AssertionError(f"artifact missing {label}")
        result = run_command(command, project_dir)
        print(f"{label}: {command}")
        print(result.stdout)
        if result.returncode != 0:
            raise AssertionError(f"{label} failed with exit code {result.returncode}")

    print("artifact:", json.dumps({
        "artifact_id": artifact.get("artifact_id"),
        "uri": artifact.get("uri"),
        "size_bytes": artifact.get("size_bytes"),
        "content_hash": artifact_hash,
        "project_dir": str(project_dir),
    }, ensure_ascii=False))


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--base-url", default=os.getenv("AGENT_MANAGER_URL", "http://127.0.0.1:8000"))
    parser.add_argument("--model-gateway-url", default=os.getenv("HEICODE_NEWAPI_BASE_URL", "https://code.xinghanlab.com/v1"))
    parser.add_argument("--model", default=os.getenv("HEICODE_E2E_MODEL", "gpt-5.4"))
    parser.add_argument("--api-format", default=os.getenv("HEICODE_E2E_API_FORMAT", "openai_chat"))
    parser.add_argument("--stream", action=argparse.BooleanOptionalAction, default=True)
    parser.add_argument("--max-tokens", type=int, default=int(os.getenv("HEICODE_E2E_MAX_TOKENS", "6000")))
    parser.add_argument("--timeout-seconds", type=int, default=int(os.getenv("HEICODE_E2E_TIMEOUT_SECONDS", "1200")))
    parser.add_argument("--model-timeout-seconds", type=int, default=int(os.getenv("HEICODE_E2E_MODEL_TIMEOUT_SECONDS", "600")))
    parser.add_argument("--poll-interval", type=int, default=10)
    parser.add_argument("--user-id", default="heicode-complete-project-e2e")
    parser.add_argument("--callback-url", default="http://127.0.0.1:9/heicode-callback")
    parser.add_argument("--correlation-id", default=f"heicode-complete-project-e2e-{int(time.time())}")
    parser.add_argument("--idempotency-key", default=f"complete-project-e2e-{int(time.time())}")
    parser.add_argument("--objective", default=DEFAULT_OBJECTIVE)
    parser.add_argument("--output-dir", type=Path, default=None)
    return parser.parse_args()


def main() -> int:
    args = parse_args()
    base_url = args.base_url.rstrip("/")
    output_dir = args.output_dir or Path(tempfile.mkdtemp(prefix="heicode-complete-project-e2e-"))
    output_dir.mkdir(parents=True, exist_ok=True)

    payload = build_payload(args)
    created = http_json(
        "POST",
        f"{base_url}/api/swarms",
        payload,
        headers={"X-Idempotency-Key": args.idempotency_key},
    )
    swarm_id = created["swarm_id"]
    print("created:", json.dumps({"swarm_id": swarm_id, "output_dir": str(output_dir)}, ensure_ascii=False))

    status = poll_swarm(base_url, swarm_id, args.timeout_seconds, args.poll_interval)
    run_project_validation(base_url, status, output_dir)
    assert_runtime_observability(base_url, swarm_id, status)
    print("PASS: complete project E2E succeeded")
    return 0


if __name__ == "__main__":
    try:
        raise SystemExit(main())
    except Exception as exc:
        print(f"FAIL: {exc}", file=sys.stderr)
        raise SystemExit(1)