aps-agent/server/engines/solver_worker.py

"""Isolated stdin/stdout worker for the CP-SAT assignment protocol."""
from __future__ import annotations

import contextlib
import hashlib
import importlib.metadata
import importlib.util
import io
import json
import os
import platform
import site
import sys
import traceback
from collections.abc import Mapping
from datetime import date
from pathlib import Path
from typing import Any

# ``python -I /abs/path/solver_worker.py`` does not trust cwd/PYTHONPATH.
_REPO_ROOT = Path(__file__).resolve().parents[2]
if str(_REPO_ROOT) not in sys.path:
    sys.path.insert(0, str(_REPO_ROOT))

from server.engines.solver_process import (
    ERROR_MARKER,
    PROTOCOL_VERSION,
    SUCCESS_MARKER,
    _canonical_json,
    _json_value,
    _request_digest,
)

_MAX_REQUEST_BYTES = 64 * 1024 * 1024
_DIAGNOSTIC_RELAXABLE_CONSTRAINTS = frozenset({
    "C1_precedence", "C2_no_overlap", "C3_calendar", "C7_capacity", "C10_changeover",
    "C11_freeze", "C12_team", "C12_tooling",
})
_DIAGNOSTIC_RHS_PARAMETERS = frozenset({
    "C7_line_day_capacity_minutes", "C8_due_date_allowance",
    "C12_team_capacity", "C12_tooling_capacity",
})
_DIAGNOSTIC_LIMIT = 8_000
_REQUIRED_NATIVE_PACKAGES = ("ortools", "numpy", "pandas")


def _resolved(path: str | os.PathLike[str]) -> Path:
    return Path(path).expanduser().resolve()


def _under(path: Path, root: Path) -> bool:
    try:
        path.relative_to(root)
        return True
    except ValueError:
        return False


def _package_identity(name: str, allowed_roots: list[Path]) -> tuple[dict[str, Any], list[str]]:
    reasons: list[str] = []
    spec = importlib.util.find_spec(name)
    origin = None if spec is None else spec.origin
    try:
        version = importlib.metadata.version(name)
    except importlib.metadata.PackageNotFoundError:
        version = None
    identity: dict[str, Any] = {"version": version, "origin": origin}
    if spec is None or not origin:
        reasons.append(f"{name}:missing")
        return identity, reasons
    origin_path = _resolved(origin)
    if not any(_under(origin_path, root) for root in allowed_roots):
        reasons.append(f"{name}:outside-runtime-prefix")
    user_site = site.getusersitepackages()
    if user_site and _under(origin_path, _resolved(user_site)):
        reasons.append(f"{name}:user-site")
    return identity, reasons


def runtime_identity() -> dict[str, Any]:
    frozen = bool(getattr(sys, "frozen", False))
    roots = [_resolved(sys.prefix), _resolved(sys.exec_prefix)]
    meipass = getattr(sys, "_MEIPASS", None)
    if meipass:
        roots.append(_resolved(meipass))
    roots = list(dict.fromkeys(roots))

    reasons: list[str] = []
    if platform.python_implementation() != "CPython":
        reasons.append("implementation:not-cpython")
    if not frozen and not bool(sys.flags.isolated):
        reasons.append("python:not-isolated")
    if bool(site.ENABLE_USER_SITE):
        reasons.append("python:user-site-enabled")

    packages: dict[str, Any] = {}
    for package in _REQUIRED_NATIVE_PACKAGES:
        packages[package], package_reasons = _package_identity(package, roots)
        reasons.extend(package_reasons)

    return {
        "implementation": platform.python_implementation(),
        "pythonVersion": platform.python_version(),
        "executable": str(_resolved(sys.executable)),
        "baseExecutable": str(_resolved(getattr(sys, "_base_executable", sys.executable))),
        "prefix": str(_resolved(sys.prefix)),
        "basePrefix": str(_resolved(sys.base_prefix)),
        "isolated": bool(sys.flags.isolated),
        "noUserSite": bool(sys.flags.no_user_site),
        "enableUserSite": bool(site.ENABLE_USER_SITE),
        "frozen": frozen,
        "allowedRoots": [str(root) for root in roots],
        "packages": packages,
        "safe": not reasons,
        "reasons": sorted(set(reasons)),
    }


def _trim(value: str) -> str:
    return value if len(value) <= _DIAGNOSTIC_LIMIT else value[:_DIAGNOSTIC_LIMIT] + "\n...[truncated]"


def _emit(payload: Mapping[str, Any]) -> None:
    encoded = _canonical_json(_json_value(payload)).encode("utf-8")
    binary = getattr(sys.stdout, "buffer", None)
    if binary is not None:
        binary.write(encoded)
        binary.flush()
    else:  # pragma: no cover - embedded hosts may expose only a text stream
        sys.stdout.write(encoded.decode("utf-8"))
        sys.stdout.flush()


def _emit_error(
    code: str,
    message: str,
    *,
    request_id: str = "",
    details: Mapping[str, Any] | None = None,
) -> None:
    _emit(
        {
            "protocolVersion": PROTOCOL_VERSION,
            "marker": ERROR_MARKER,
            "ok": False,
            "requestId": request_id,
            "error": {"code": code, "message": message, "details": dict(details or {})},
        }
    )


def _read_request() -> dict[str, Any]:
    raw = sys.stdin.buffer.read(_MAX_REQUEST_BYTES + 1)
    if len(raw) > _MAX_REQUEST_BYTES:
        raise ValueError("request exceeds size limit")
    try:
        payload = json.loads(raw.decode("utf-8"))
    except (UnicodeDecodeError, json.JSONDecodeError) as exc:
        raise ValueError(f"invalid request JSON: {exc}") from exc
    if not isinstance(payload, dict):
        raise TypeError("request must be an object")
    return payload


def _validate_request(
    request: Mapping[str, Any],
) -> tuple[
    str, dict[str, Any], list[dict], dict, list[dict] | None, str,
    str, list[str], dict[str, Any] | None, str | None, bool, bool,
]:
    request_id = request.get("requestId")
    if not isinstance(request_id, str) or not request_id:
        raise ValueError("requestId is required")
    if request.get("protocolVersion") != PROTOCOL_VERSION:
        raise RuntimeError("protocol mismatch")
    operation = request.get("operation")
    common_keys = {
        "requestId", "protocolVersion", "operation", "world", "entries",
        "params", "warmStart", "pipelineLabel",
    }
    constraint_diagnostic = operation in {
        "diagnose_constraint_baseline", "diagnose_constraint_removal",
    }
    rhs_diagnostic = operation in {"diagnose_rhs_baseline", "diagnose_rhs_perturbation"}
    diagnostic = constraint_diagnostic or rhs_diagnostic
    diagnostic_keys = (
        {"relaxedConstraintIds", "invocationId"}
        if constraint_diagnostic
        else ({"rhsPerturbation", "invocationId"} if rhs_diagnostic else set())
    )
    expected_keys = common_keys | diagnostic_keys
    if set(request) != expected_keys:
        raise ValueError(
            f"request keys mismatch: missing={sorted(expected_keys - set(request))} "
            f"unknown={sorted(set(request) - expected_keys)}"
        )
    body = {key: value for key, value in request.items() if key != "requestId"}
    if _request_digest(body) != request_id:
        raise ValueError("requestId digest mismatch")
    if operation not in {
        "optimize_line_assignment", "diagnose_constraint_baseline",
        "diagnose_constraint_removal", "diagnose_rhs_baseline",
        "diagnose_rhs_perturbation",
    }:
        raise ValueError("unsupported operation")
    world = request.get("world")
    entries = request.get("entries")
    params = request.get("params")
    warm_start = request.get("warmStart")
    pipeline_label = request.get("pipelineLabel")
    relaxed_constraint_ids = request.get("relaxedConstraintIds", [])
    rhs_perturbation = request.get("rhsPerturbation")
    invocation_id = request.get("invocationId")
    if not isinstance(world, dict):
        raise TypeError("world must be an object")
    if not isinstance(entries, list) or not all(isinstance(item, dict) for item in entries):
        raise TypeError("entries must be an object array")
    if not isinstance(params, dict):
        raise TypeError("params must be an object")
    if warm_start is not None and (
        not isinstance(warm_start, list) or not all(isinstance(item, dict) for item in warm_start)
    ):
        raise TypeError("warmStart must be an object array")
    if diagnostic and warm_start is not None:
        raise ValueError("diagnostic operations do not accept warmStart")
    if not isinstance(pipeline_label, str) or not pipeline_label.strip():
        raise ValueError("pipelineLabel is required")
    if not isinstance(relaxed_constraint_ids, list) or not all(
        isinstance(value, str) and value for value in relaxed_constraint_ids
    ):
        raise TypeError("relaxedConstraintIds must be a string array")
    if len(relaxed_constraint_ids) != len(set(relaxed_constraint_ids)):
        raise ValueError("relaxedConstraintIds must not contain duplicates")
    unknown_relaxed = set(relaxed_constraint_ids) - _DIAGNOSTIC_RELAXABLE_CONSTRAINTS
    if unknown_relaxed:
        raise ValueError(f"unsupported relaxedConstraintIds: {sorted(unknown_relaxed)}")
    if operation == "optimize_line_assignment" and relaxed_constraint_ids:
        raise ValueError("production operation cannot relax constraints")
    if operation == "diagnose_constraint_baseline" and relaxed_constraint_ids:
        raise ValueError("diagnostic baseline cannot relax constraints")
    if operation == "diagnose_constraint_removal" and len(relaxed_constraint_ids) != 1:
        raise ValueError("diagnostic removal requires exactly one constraint")
    normalized_rhs: dict[str, Any] | None = None
    if rhs_perturbation is not None:
        if not isinstance(rhs_perturbation, dict):
            raise TypeError("rhsPerturbation must be an object or null")
        parameter_id = rhs_perturbation.get("parameterId")
        increment = rhs_perturbation.get("increment")
        if parameter_id not in _DIAGNOSTIC_RHS_PARAMETERS:
            raise ValueError("rhsPerturbation parameterId is unsupported")
        if isinstance(increment, bool) or not isinstance(increment, int) or increment <= 0:
            raise ValueError("rhsPerturbation increment must be a positive integer")
        if parameter_id == "C7_line_day_capacity_minutes":
            line_id = rhs_perturbation.get("lineId")
            bucket_date = rhs_perturbation.get("bucketDate")
            if (
                set(rhs_perturbation) != {"parameterId", "lineId", "bucketDate", "increment"}
                or isinstance(line_id, bool)
                or not isinstance(line_id, int)
                or line_id <= 0
                or not isinstance(bucket_date, str)
                or increment > 1_440
            ):
                raise ValueError("C7 rhsPerturbation is invalid")
            try:
                if date.fromisoformat(bucket_date).isoformat() != bucket_date:
                    raise ValueError
            except ValueError as exc:
                raise ValueError("C7 rhsPerturbation bucketDate is invalid") from exc
            normalized_rhs = {
                "parameterId": parameter_id,
                "lineId": line_id,
                "bucketDate": bucket_date,
                "increment": increment,
            }
        elif parameter_id == "C8_due_date_allowance":
            if set(rhs_perturbation) != {"parameterId", "increment"} or increment > 10_080:
                raise ValueError("C8 rhsPerturbation is invalid")
            normalized_rhs = {"parameterId": parameter_id, "increment": increment}
        else:
            resource_id = rhs_perturbation.get("resourceId")
            if (
                set(rhs_perturbation) != {"parameterId", "resourceId", "increment"}
                or isinstance(resource_id, bool)
                or not isinstance(resource_id, int)
                or resource_id <= 0
                or increment > 100
            ):
                raise ValueError("C12 rhsPerturbation is invalid")
            normalized_rhs = {
                "parameterId": parameter_id,
                "resourceId": resource_id,
                "increment": increment,
            }
    if operation == "diagnose_rhs_baseline" and normalized_rhs is not None:
        raise ValueError("RHS diagnostic baseline cannot perturb a parameter")
    if operation == "diagnose_rhs_perturbation" and normalized_rhs is None:
        raise ValueError("RHS diagnostic perturbation requires one parameter")
    if not rhs_diagnostic and normalized_rhs is not None:
        raise ValueError("non-RHS operation cannot perturb RHS parameters")
    if diagnostic and (
        not isinstance(invocation_id, str)
        or len(invocation_id) != 32
        or any(char not in "0123456789abcdef" for char in invocation_id.lower())
    ):
        raise ValueError("diagnostic invocationId must be 32 hex characters")
    return (
        request_id, world, entries, params, warm_start, pipeline_label,
        str(operation), sorted(relaxed_constraint_ids), normalized_rhs,
        str(invocation_id) if diagnostic else None, diagnostic, rhs_diagnostic,
    )


def main() -> int:
    request_id = ""
    try:
        request = _read_request()
        request_id = str(request.get("requestId") or "")
    except ValueError as exc:
        _emit_error("SOLVER_REQUEST_INVALID", "求解请求非法", details={"error": str(exc)})
        return 0

    if request.get("protocolVersion") != PROTOCOL_VERSION:
        _emit_error(
            "SOLVER_PROTOCOL_MISMATCH",
            "求解请求协议版本不一致",
            request_id=request_id,
            details={"expected": PROTOCOL_VERSION, "actual": request.get("protocolVersion")},
        )
        return 0

    try:
        (
            request_id, world, entries, params_payload, warm_start,
            pipeline_label, operation, relaxed_constraint_ids,
            rhs_perturbation, invocation_id, diagnostic_mode, rhs_diagnostic_mode,
        ) = _validate_request(request)
    except RuntimeError:
        _emit_error(
            "SOLVER_PROTOCOL_MISMATCH",
            "求解请求协议版本不一致",
            request_id=request_id,
        )
        return 0
    except (TypeError, ValueError) as exc:
        _emit_error(
            "SOLVER_REQUEST_INVALID",
            "求解请求非法",
            request_id=request_id,
            details={"error": str(exc)},
        )
        return 0

    identity = runtime_identity()
    if identity.get("safe") is not True:
        _emit_error(
            "SOLVER_RUNTIME_UNSAFE",
            "求解子进程运行时身份不安全",
            request_id=request_id,
            details={"runtimeIdentity": identity},
        )
        return 0

    captured_stdout = io.StringIO()
    captured_stderr = io.StringIO()
    try:
        with contextlib.redirect_stdout(captured_stdout), contextlib.redirect_stderr(captured_stderr):
            from server.engines.base import EngineParams
            from server.engines.cp_engine import optimize_line_assignment

            params = EngineParams.model_validate(params_payload)
            ordered, solver_meta = optimize_line_assignment(
                world,
                entries,
                params,
                warm_start=warm_start,
                pipeline_label=pipeline_label,
                relaxed_constraint_ids=relaxed_constraint_ids,
                diagnostic_mode=diagnostic_mode,
                rhs_diagnostic_mode=rhs_diagnostic_mode,
                rhs_perturbation=rhs_perturbation,
            )
        if not isinstance(ordered, list) or not all(isinstance(item, dict) for item in ordered):
            raise TypeError("optimizer entries result is not an object array")
        if not isinstance(solver_meta, dict):
            raise TypeError("optimizer solverMeta result is not an object")
    except Exception as exc:  # noqa: BLE001 - child must translate every Python-level failure
        _emit_error(
            "SOLVER_EXECUTION_FAILED",
            "CP-SAT 子进程执行失败",
            request_id=request_id,
            details={
                "exceptionType": type(exc).__name__,
                "error": str(exc),
                "stdout": _trim(captured_stdout.getvalue()),
                "stderr": _trim(captured_stderr.getvalue()),
                "traceback": _trim(traceback.format_exc()),
                "runtimeIdentity": identity,
            },
        )
        return 0

    response_body = {
        "protocolVersion": PROTOCOL_VERSION,
        "marker": SUCCESS_MARKER,
        "ok": True,
        "requestId": request_id,
        "operation": operation,
        "invocationId": invocation_id,
        "runtimeIdentity": identity,
        "result": {"entries": ordered, "solverMeta": solver_meta},
    }
    response_body["responseDigest"] = hashlib.sha256(
        _canonical_json(_json_value(response_body)).encode("utf-8")
    ).hexdigest()
    _emit(response_body)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())