from __future__ import annotations import json import threading import time import uuid from concurrent.futures import FIRST_COMPLETED, ThreadPoolExecutor, wait from datetime import datetime, timezone from pathlib import Path from services.account_service import account_service from services.config import DATA_DIR from services.register import openai_register REGISTER_FILE = DATA_DIR / "register.json" def _now() -> str: return datetime.now(timezone.utc).isoformat() def _default_config() -> dict: return { **openai_register.config, "mode": "total", "target_quota": 100, "target_available": 10, "check_interval": 5, "enabled": False, "cpa_auto_import": { "enabled": False, "base_url": "http://host.docker.internal:8317", "secret_key": "", }, "stats": { "success": 0, "fail": 0, "done": 0, "running": 0, "threads": openai_register.config["threads"], "elapsed_seconds": 0, "avg_seconds": 0, "success_rate": 0, "current_quota": 0, "current_available": 0, }, } def _normalize(raw: dict) -> dict: cfg = _default_config() cfg.update({k: v for k, v in raw.items() if k not in {"stats", "logs"}}) cfg["total"] = max(1, int(cfg.get("total") or 1)) cfg["threads"] = max(1, int(cfg.get("threads") or 1)) cfg["mode"] = str(cfg.get("mode") or "total").strip() if str(cfg.get("mode") or "total").strip() in {"total", "quota", "available"} else "total" cfg["target_quota"] = max(1, int(cfg.get("target_quota") or 1)) cfg["target_available"] = max(1, int(cfg.get("target_available") or 1)) cfg["check_interval"] = max(1, int(cfg.get("check_interval") or 5)) cfg["proxy"] = str(cfg.get("proxy") or "").strip() cfg["enabled"] = bool(cfg.get("enabled")) cpa_auto_import = cfg.get("cpa_auto_import") if isinstance(cfg.get("cpa_auto_import"), dict) else {} cfg["cpa_auto_import"] = { "enabled": bool(cpa_auto_import.get("enabled")), "base_url": str(cpa_auto_import.get("base_url") or "http://host.docker.internal:8317").strip(), "secret_key": str(cpa_auto_import.get("secret_key") or "").strip(), } stats = {**_default_config()["stats"], **(raw.get("stats") if isinstance(raw.get("stats"), dict) else {}), "threads": cfg["threads"]} cfg["stats"] = stats return cfg class RegisterService: def __init__(self, store_file: Path): self._store_file = store_file self._lock = threading.RLock() self._runner: threading.Thread | None = None self._logs: list[dict] = [] openai_register.register_log_sink = self._append_log self._config = self._load() if self._config["enabled"]: self.start() def _load(self) -> dict: try: return _normalize(json.loads(self._store_file.read_text(encoding="utf-8"))) except Exception: return _normalize({}) def _save(self) -> None: self._store_file.parent.mkdir(parents=True, exist_ok=True) self._store_file.write_text(json.dumps(self._config, ensure_ascii=False, indent=2) + "\n", encoding="utf-8") def get(self) -> dict: with self._lock: return json.loads(json.dumps({**self._config, "logs": self._logs[-300:]}, ensure_ascii=False)) def update(self, updates: dict) -> dict: with self._lock: self._config = _normalize({**self._config, **updates}) openai_register.config.update({k: self._config[k] for k in ("mail", "proxy", "total", "threads", "cpa_auto_import")}) self._save() return self.get() def start(self) -> dict: with self._lock: if self._runner and self._runner.is_alive(): self._config["enabled"] = True self._save() return self.get() self._config["enabled"] = True self._logs = [] metrics = self._pool_metrics() self._config["stats"] = {"job_id": uuid.uuid4().hex, "success": 0, "fail": 0, "done": 0, "running": 0, "threads": self._config["threads"], **metrics, "started_at": _now(), "updated_at": _now()} openai_register.config.update({k: self._config[k] for k in ("mail", "proxy", "total", "threads", "cpa_auto_import")}) with openai_register.stats_lock: openai_register.stats.update({"done": 0, "success": 0, "fail": 0, "start_time": time.time()}) self._save() self._runner = threading.Thread(target=self._run, daemon=True, name="openai-register") self._runner.start() self._append_log(f"注册任务启动，模式={self._config['mode']}，线程数={self._config['threads']}", "yellow") return self.get() def stop(self) -> dict: with self._lock: self._config["enabled"] = False self._config["stats"]["updated_at"] = _now() self._save() self._append_log("已请求停止注册任务，正在等待当前运行任务结束", "yellow") return self.get() def reset(self) -> dict: with self._lock: self._logs = [] self._config["stats"] = {"success": 0, "fail": 0, "done": 0, "running": 0, "threads": self._config["threads"], "elapsed_seconds": 0, "avg_seconds": 0, "success_rate": 0, **self._pool_metrics(), "updated_at": _now()} with openai_register.stats_lock: openai_register.stats.update({"done": 0, "success": 0, "fail": 0, "start_time": 0.0}) self._save() return self.get() def _append_log(self, text: str, color: str = "") -> None: with self._lock: self._logs.append({"time": _now(), "text": str(text), "level": str(color or "info")}) self._logs = self._logs[-300:] def _pool_metrics(self) -> dict: items = account_service.list_accounts() normal = [item for item in items if item.get("status") == "正常"] return { "current_quota": sum(int(item.get("quota") or 0) for item in normal if not item.get("image_quota_unknown")), "current_available": len(normal), } def _target_reached(self, cfg: dict, submitted: int) -> bool: mode = str(cfg.get("mode") or "total") metrics = self._pool_metrics() self._bump(**metrics) if mode == "quota": reached = metrics["current_quota"] >= int(cfg.get("target_quota") or 1) self._append_log(f"检查号池：当前正常账号={metrics['current_available']}，当前剩余额度={metrics['current_quota']}，目标额度={cfg.get('target_quota')}，{'跳过注册' if reached else '继续注册'}", "yellow") return reached if mode == "available": reached = metrics["current_available"] >= int(cfg.get("target_available") or 1) self._append_log(f"检查号池：当前正常账号={metrics['current_available']}，目标账号={cfg.get('target_available')}，当前剩余额度={metrics['current_quota']}，{'跳过注册' if reached else '继续注册'}", "yellow") return reached return submitted >= int(cfg.get("total") or 1) def _bump(self, **updates) -> None: with self._lock: self._config["stats"].update(updates) stats = self._config["stats"] started_at = str(stats.get("started_at") or "") if started_at: try: elapsed = max(0.0, (datetime.now(timezone.utc) - datetime.fromisoformat(started_at)).total_seconds()) except Exception: elapsed = 0.0 done = int(stats.get("done") or 0) success = int(stats.get("success") or 0) fail = int(stats.get("fail") or 0) stats["elapsed_seconds"] = round(elapsed, 1) stats["avg_seconds"] = round(elapsed / success, 1) if success else 0 stats["success_rate"] = round(success * 100 / max(1, success + fail), 1) self._config["stats"]["updated_at"] = _now() # 不在此落盘：stats 只是临时展示字段，SSE 读内存 _config 即可，start() 每次重建。 # _bump 每次号池检查都调用，旧实现会把 register.json 写到 /data FUSE 桶； # 桶写一次抽风抛异常，就会杀死没有 try 的 _run 监控线程，enabled 卡在 True、 # 号池掉到目标下也不再触发注册（CLAUDE.md 已警告对象桶每次全量写的反模式）。 # 真正落盘只发生在 start/stop/update/reset。 def _run(self) -> None: threads = int(self.get()["threads"]) submitted, done, success, fail = 0, 0, 0, 0 try: with ThreadPoolExecutor(max_workers=threads) as executor: futures: set = set() while True: try: cfg = self.get() while self.get()["enabled"] and not self._target_reached(cfg, submitted) and len(futures) < threads: submitted += 1 futures.add(executor.submit(openai_register.worker, submitted)) self._bump(running=len(futures), done=done, success=success, fail=fail) if not futures and (not self.get()["enabled"] or str(cfg.get("mode") or "total") == "total"): break if not futures: time.sleep(max(1, int(cfg.get("check_interval") or 5))) continue finished, futures = wait(futures, return_when=FIRST_COMPLETED) for future in finished: done += 1 try: result = future.result() success += 1 if result.get("ok") else 0 fail += 0 if result.get("ok") else 1 except Exception: fail += 1 except Exception as loop_error: # 单次循环出错（如 /data FUSE 桶写盘抽风、号池读取瞬时异常）绝不能杀死 # 监控线程：否则 enabled 卡在 True、徽章仍显示“运行中”，但号池掉到目标下 # 也永不再触发注册，只能手动 stop→start 才恢复。忽略本次，睡一个检查间隔后继续。 self._append_log(f"号池监控循环异常，已忽略本次并继续: {loop_error}", "red") time.sleep(max(1, int(self.get().get("check_interval") or 5))) finally: # 无论正常结束还是异常退出，都把 enabled 落为 False，避免 UI 卡在“运行中”假象。 try: self._bump(running=0, done=done, success=success, fail=fail, finished_at=_now()) except Exception: pass with self._lock: self._config["enabled"] = False try: self._save() except Exception: pass self._append_log(f"注册任务结束，成功{success}，失败{fail}", "yellow") register_service = RegisterService(REGISTER_FILE)