huni-probe-agent/main.py

"""ModelHub XC 适配智能体安全探针 v1.0.3。

默认只做一次只读鉴权检查。只有显式设置 ALLOW_SUBMIT=1，且鉴权检查成功后，
才会调用 build-config 和 task/add；每个进程生命周期最多提交一次。
"""

import copy
import json
import os
import signal
import threading
import time
import urllib.error
import urllib.request
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from pathlib import Path
from typing import Callable, Mapping, Optional, Tuple


VERSION = "1.0.3"
MAIN = os.getenv("MAIN_HOST", "https://modelhub.org.cn").rstrip("/")
STRATEGY_ID = os.getenv("STRATEGY_ID", "").strip()
PORT = int(os.getenv("PORT", "8080"))
ALLOW_SUBMIT = os.getenv("ALLOW_SUBMIT", "").strip().lower() in {"1", "true", "yes", "on"}

TEST_MODEL = os.getenv(
    "TEST_MODEL",
    "https://www.modelscope.cn/models/mradermacher/TrialSpace-1225-GGUF",
)
GPU = os.getenv("GPU", "MetaX_c-500")
FW = os.getenv("FW", "vllm")
TT = os.getenv("TT", "text-generation")

_TOKEN_PLACEHOLDERS = {"tmp", "placeholder", "changeme", "change-me", "example", "test"}
shutdown = threading.Event()
STATE_LOCK = threading.Lock()


def _initial_state() -> dict:
    return {
        "version": VERSION,
        "phase": "starting",
        "done": False,
        "auth_ready": False,
        "submit_enabled": ALLOW_SUBMIT,
        "token_source": "none",
        "results": {},
    }


STATE = _initial_state()


def log(msg: str) -> None:
    print(f"[{time.strftime('%Y-%m-%d %H:%M:%S')}] {msg}", flush=True)


def _set_state(**changes) -> None:
    with STATE_LOCK:
        STATE.update(changes)


def _snapshot() -> dict:
    with STATE_LOCK:
        return copy.deepcopy(STATE)


def _reset_state() -> None:
    """Reset process state; kept separate to make the probe deterministically testable."""
    with STATE_LOCK:
        STATE.clear()
        STATE.update(_initial_state())


def _valid_token(value: str) -> bool:
    value = value.strip()
    return bool(value) and value.lower() not in _TOKEN_PLACEHOLDERS


def _resolve_token(
    env: Optional[Mapping[str, str]] = None,
    read_text: Optional[Callable[[str], str]] = None,
) -> Tuple[str, str, Optional[str]]:
    """Resolve a credential without exposing its value or length.

    Precedence is XC_TOKEN_FILE, XC_TOKEN, then the legacy
    EXTERNAL_SERVICE_TOKEN. An explicitly configured but unreadable token file
    is treated as an error instead of silently falling back.
    """
    values = os.environ if env is None else env
    file_reader = read_text or (lambda path: Path(path).read_text(encoding="utf-8"))
    token_file = values.get("XC_TOKEN_FILE", "").strip()
    if token_file:
        try:
            value = file_reader(token_file).strip()
        except (OSError, UnicodeError):
            return "", "XC_TOKEN_FILE", "token_file_unreadable"
        if _valid_token(value):
            return value, "XC_TOKEN_FILE", None
        return "", "XC_TOKEN_FILE", "token_missing_or_placeholder"

    for name in ("XC_TOKEN", "EXTERNAL_SERVICE_TOKEN"):
        value = values.get(name, "").strip()
        if _valid_token(value):
            return value, name, None
        if value:
            return "", name, "token_missing_or_placeholder"
    return "", "none", "token_missing_or_placeholder"


def _req(method: str, path: str, token: str, body=None, timeout: int = 30) -> dict:
    headers = {"Accept": "application/json", "Xc-Token": token}
    data = body.encode("utf-8") if isinstance(body, str) else body
    if data is not None:
        headers["Content-Type"] = "application/json"
    req = urllib.request.Request(MAIN + path, data=data, method=method, headers=headers)
    try:
        with urllib.request.urlopen(req, timeout=timeout) as response:
            text = response.read().decode("utf-8", "ignore")
            status = response.status
    except urllib.error.HTTPError as error:
        text = error.read().decode("utf-8", "ignore")
        status = error.code
    except Exception as error:  # noqa: BLE001 - network diagnostics are returned safely
        return {"http": -1, "error": type(error).__name__}

    try:
        payload = json.loads(text)
    except (TypeError, ValueError):
        return {"http": status, "code": None, "message": "non_json_response"}
    return {
        "http": status,
        "code": payload.get("code"),
        "message": str(payload.get("message") or "")[:100],
        "data": payload.get("data"),
    }


def _api_ok(result: dict) -> bool:
    try:
        status = int(result.get("http", -1))
    except (TypeError, ValueError):
        return False
    return 200 <= status < 300 and result.get("code") in {0, "0"}


def _result_summary(result: dict) -> dict:
    summary = {"http": result.get("http"), "code": result.get("code")}
    if result.get("message"):
        summary["message"] = result["message"]
    if result.get("error"):
        summary["error"] = result["error"]
    return summary


def _task_id(data) -> Optional[str]:
    if not isinstance(data, dict):
        return None
    for key in ("id", "taskId", "task_id"):
        if data.get(key) is not None:
            return str(data[key])
    return None


def probe() -> None:
    """Run one guarded probe; no retry loop and at most one task/add call."""
    token, source, token_error = _resolve_token()
    _set_state(token_source=source, submit_enabled=ALLOW_SUBMIT)

    if not STRATEGY_ID:
        _set_state(phase="blocked_missing_strategy_id", done=True)
        log("probe blocked: STRATEGY_ID is missing")
        return
    if token_error:
        _set_state(phase="blocked_missing_auth", done=True, results={"auth": token_error})
        log(f"probe blocked: usable credential is unavailable (source={source})")
        return

    _set_state(phase="checking_auth")
    page = _req("GET", "/api/adapt/task/page?current=1&pageSize=1", token)
    results = {"task_page": _result_summary(page)}
    if not _api_ok(page):
        _set_state(phase="auth_failed", done=True, results=results)
        log("read-only auth check failed: " + json.dumps(results["task_page"], ensure_ascii=False))
        return

    _set_state(auth_ready=True, results=results)
    log(f"read-only auth check passed (source={source})")
    if not ALLOW_SUBMIT:
        _set_state(phase="read_only_complete", done=True)
        log("probe complete in read-only mode; task/add was not called")
        return

    _set_state(phase="preparing_submission")
    build_path = f"/api/adapt/task/build-config?gpuType={GPU}&framework={FW}&taskType={TT}"
    build_config = _req("POST", build_path, token)
    results["build_config"] = _result_summary(build_config)
    config = build_config.get("data")
    if not _api_ok(build_config) or not isinstance(config, str) or not config:
        _set_state(phase="build_config_failed", done=True, results=results)
        log("build-config failed; task/add was not called")
        return

    body = json.dumps(
        {
            "modelAddress": TEST_MODEL,
            "taskType": TT,
            "targetGpu": GPU,
            "framework": FW,
            "strategyId": STRATEGY_ID,
            "configParams": config,
        }
    )
    task_add = _req("POST", "/api/adapt/task/add", token, body)
    results["task_add"] = _result_summary(task_add)
    task_id = _task_id(task_add.get("data"))
    if task_id:
        results["task_add"]["task_id"] = task_id

    phase = "submit_complete" if _api_ok(task_add) else "submit_failed"
    _set_state(phase=phase, done=True, results=results)
    log("single task/add attempt finished: " + json.dumps(results["task_add"], ensure_ascii=False))


class Handler(BaseHTTPRequestHandler):
    def _json(self, obj: dict, status: int = 200) -> None:
        payload = json.dumps(obj, ensure_ascii=False).encode("utf-8")
        self.send_response(status)
        self.send_header("Content-Type", "application/json; charset=utf-8")
        self.send_header("Content-Length", str(len(payload)))
        self.end_headers()
        self.wfile.write(payload)

    def do_GET(self) -> None:
        if self.path == "/health":
            self._json({"status": "ok", "version": VERSION})
            return
        if self.path in {"/", "/status"}:
            self._json({"name": "huni-probe-agent", "probe": _snapshot()})
            return
        self._json({"error": "not found"}, 404)

    def log_message(self, *_args) -> None:
        pass


def _handle_signal(signum: int, _frame) -> None:
    shutdown.set()
    log(f"received signal {signum}, shutting down")


def main() -> None:
    signal.signal(signal.SIGTERM, _handle_signal)
    signal.signal(signal.SIGINT, _handle_signal)
    threading.Thread(target=probe, daemon=True).start()

    server = ThreadingHTTPServer(("0.0.0.0", PORT), Handler)
    server.timeout = 1
    log(f"probe agent v{VERSION} listening on 0.0.0.0:{PORT} (submit_enabled={ALLOW_SUBMIT})")
    while not shutdown.is_set():
        server.handle_request()
    server.server_close()
    log("stopped")


if __name__ == "__main__":
    main()