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 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_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:
    sys.stdout.write(_canonical_json(_json_value(payload)))
    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]:
    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")
    body = {key: value for key, value in request.items() if key != "requestId"}
    if _request_digest(body) != request_id:
        raise ValueError("requestId digest mismatch")
    if request.get("operation") != "optimize_line_assignment":
        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")
    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 not isinstance(pipeline_label, str) or not pipeline_label.strip():
        raise ValueError("pipelineLabel is required")
    return request_id, world, entries, params, warm_start, pipeline_label


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 = _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,
            )
        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,
        "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())