SupplyAgent/supply_infra/pipeline/step_runner.py

from __future__ import annotations

import json
import logging
import os
import signal
import subprocess
import sys
from dataclasses import dataclass
from pathlib import Path
from typing import Any

from supply_infra.config import InfraSettings, get_infra_settings
from supply_infra.pipeline.contracts import StepResult

logger = logging.getLogger(__name__)
_REPO_ROOT = Path(__file__).resolve().parents[2]
_MAX_RESULT_BYTES = 100_000


@dataclass(frozen=True)
class StepExecutionRequest:
    run_id: str
    step_run_id: str
    step_key: str
    timeout_seconds: int


def _log_directory(settings: InfraSettings, run_id: str) -> Path:
    base = Path(settings.pipeline_log_dir)
    if not base.is_absolute():
        base = _REPO_ROOT / base
    path = base / run_id
    path.mkdir(parents=True, exist_ok=True)
    return path


def _read_result(path: Path) -> dict[str, Any]:
    if not path.is_file():
        return {"success": False, "error_code": "missing_result", "error": "Result file missing"}
    raw = path.read_bytes()
    if len(raw) > _MAX_RESULT_BYTES:
        return {
            "success": False,
            "error_code": "result_too_large",
            "error": f"Step result exceeds {_MAX_RESULT_BYTES} bytes",
            "result_bytes": len(raw),
        }
    try:
        payload = json.loads(raw.decode("utf-8"))
    except (UnicodeDecodeError, json.JSONDecodeError) as exc:
        return {
            "success": False,
            "error_code": "invalid_result_json",
            "error": str(exc),
        }
    if not isinstance(payload, dict):
        return {
            "success": False,
            "error_code": "invalid_result_type",
            "error": "Step result must be a JSON object",
        }
    return payload


def run_step_subprocess(
    request: StepExecutionRequest,
    *,
    settings: InfraSettings | None = None,
) -> StepResult:
    active = settings or get_infra_settings()
    directory = _log_directory(active, request.run_id)
    stem = f"{request.step_key}-attempt-{request.step_run_id}"
    log_path = directory / f"{stem}.log"
    result_path = directory / f"{stem}.result.json"
    result_path.unlink(missing_ok=True)
    command = [
        sys.executable,
        "-m",
        "supply_infra.pipeline.step_cli",
        request.step_run_id,
        "--result-file",
        str(result_path),
    ]
    env = os.environ.copy()
    env["PROCESS_ROLE"] = "step"
    logger.info("Starting pipeline step: step_run_id=%s", request.step_run_id)

    timed_out = False
    exit_code: int | None = None
    with log_path.open("a", encoding="utf-8") as log_file:
        try:
            process = subprocess.Popen(
                command,
                cwd=str(_REPO_ROOT),
                env=env,
                stdout=log_file,
                stderr=subprocess.STDOUT,
                text=True,
                start_new_session=True,
            )
        except Exception as exc:
            logger.exception("Failed to launch pipeline step: %s", request.step_run_id)
            return StepResult(
                success=False,
                payload={},
                error_code="step_launch_failed",
                error_message=str(exc),
                log_uri=str(log_path),
            )
        try:
            exit_code = process.wait(timeout=request.timeout_seconds)
        except subprocess.TimeoutExpired:
            timed_out = True
            os.killpg(process.pid, signal.SIGTERM)
            try:
                exit_code = process.wait(timeout=min(active.pipeline_shutdown_grace_seconds, 30))
            except subprocess.TimeoutExpired:
                os.killpg(process.pid, signal.SIGKILL)
                exit_code = process.wait()

    if timed_out:
        return StepResult(
            success=False,
            payload={},
            error_code="step_timed_out",
            error_message=f"Step timed out after {request.timeout_seconds}s",
            exit_code=exit_code,
            log_uri=str(log_path),
            timed_out=True,
        )

    payload = _read_result(result_path)
    result_path.unlink(missing_ok=True)
    success = exit_code == 0 and payload.get("success") is not False
    return StepResult(
        success=success,
        payload=payload,
        error_code=(
            None
            if success
            else str(payload.get("error_code") or "step_failed")
        ),
        error_message=(
            None
            if success
            else str(payload.get("error") or f"Step exited with code {exit_code}")
        ),
        exit_code=exit_code,
        log_uri=str(log_path),
    )