xas-test-library/console/job_main.py

"""In-container bootstrap. Input contains optional credentials; never persist it."""
from http.server import BaseHTTPRequestHandler, HTTPServer
import json
import os
from pathlib import Path
import subprocess
import sys
import threading
import time
import urllib.error
import urllib.request


class NoRedirect(urllib.request.HTTPRedirectHandler):
    def redirect_request(self, request, fp, code, msg, headers, newurl):
        return None


class BudgetRelay(HTTPServer):
    def __init__(self, address, upstream, key):
        super().__init__(address, RelayHandler)
        self.upstream, self.key = upstream, key
        self.gets, self.posts = 0, 0
        self.records = []


class RelayHandler(BaseHTTPRequestHandler):
    def log_message(self, *args):
        pass

    def do_GET(self):
        if self.path != "/v1/models" or self.server.gets >= 1:
            self.send_error(403, "Request outside run budget")
            return
        self.server.gets += 1
        self.forward(None, 10)

    def do_POST(self):
        if self.path != "/v1/chat/completions" or self.server.posts >= 3:
            self.send_error(403, "Request outside run budget")
            return
        try:
            length = int(self.headers.get("Content-Length", "0"))
            if not 0 < length <= 8192:
                raise ValueError()
            data = self.rfile.read(length)
            body = json.loads(data)
            if body.get("model") != "deepseek-v4.1-flash" or body.get("max_tokens") != 512 or body.get("stream"):
                raise ValueError()
        except (ValueError, json.JSONDecodeError):
            self.send_error(400, "Unexpected smoke request")
            return
        self.server.posts += 1  # Count before I/O: timeouts never replenish the budget.
        self.forward(data, 600)

    def forward(self, data, timeout):
        started = time.monotonic()
        record = {"path": self.path, "status": None}
        try:
            req = urllib.request.Request(self.server.upstream + self.path, data=data, headers={
                "Content-Type": "application/json", "Authorization": "Bearer " + self.server.key})
            opener = urllib.request.build_opener(urllib.request.ProxyHandler({}), NoRedirect())
            with opener.open(req, timeout=timeout) as response:
                status, payload = response.status, response.read(2 * 1024 * 1024 + 1)
            if len(payload) > 2 * 1024 * 1024:
                raise ValueError("Response exceeded limit")
            record["status"] = status
            decoded = json.loads(payload)
            if self.path.endswith("completions"):
                record["response"] = decoded
            self.send_response(status)
            self.send_header("Content-Type", "application/json")
            self.send_header("Content-Length", str(len(payload)))
            self.end_headers()
            self.wfile.write(payload)
        except urllib.error.HTTPError as exc:
            record["status"] = exc.code
            self.send_error(exc.code, "Upstream request failed")
        except (OSError, ValueError) as exc:
            record["error"] = type(exc).__name__
            try:
                self.send_error(502, "Upstream outcome uncertain; no retry")
            except OSError:
                pass
        finally:
            record["seconds"] = round(time.monotonic() - started, 3)
            self.server.records.append(record)
            print(json.dumps({"event": "model_request", **{k: v for k, v in record.items() if k != "response"}}), flush=True)


def main():
    payload = json.loads(sys.stdin.buffer.readline(16384))
    relay = None
    if payload.get("profile") == "legacy_smoke":
        relay = BudgetRelay(("127.0.0.1", 18100), payload.pop("upstream"), payload.pop("key"))
        threading.Thread(target=relay.serve_forever, daemon=True).start()
    command = (["python3", "-u"] if payload["language"] == "python" else ["/bin/bash"]) + [payload["script"]] + payload["args"]
    print("[控制台] 测试进程已启动，输出将持续显示。", flush=True)
    status = subprocess.call(command, stdin=subprocess.DEVNULL)
    if relay:
        report = {"synthetic": False, "scope": "接口冒烟检查，非模型质量评测", "get_attempts": relay.gets,
                  "completion_attempts": relay.posts, "records": relay.records, "script_exit_code": status}
        valid = len(relay.records) == 4 and all(r.get("status") == 200 for r in relay.records)
        report["requests_completed"] = valid
        report["message"] = "请求均返回成功；请结合实际回答与 finish_reason 判断内容" if valid else "存在失败或未完成请求；未自动重试"
        Path(os.environ["XAS_RUN_DIR"], "smoke-response.json").write_text(json.dumps(report, ensure_ascii=False, indent=2), encoding="utf-8")
        if not valid:
            status = status or 1
        relay.shutdown()
    print(f"[控制台] 测试进程结束，退出码 {status}。", flush=True)
    return status


if __name__ == "__main__":
    sys.exit(main())