feat(nox-api): implement batch creation rate limit resilience and pending backfilling

2026-06-30 17:00:26 +08:00
parent 854a640d48
commit 4fa6a79c5a
3 changed files with 557 additions and 21 deletions
@@ -24,7 +24,7 @@ from app.schemas.upstream import (
    UpstreamCreate, UpstreamUpdate, UpstreamResponse, SnapshotResponse, TestResult,
    UpstreamBatchActionItem, UpstreamBatchActionSummary, UpstreamBatchActionResponse,
 )
-from app.services.upstream_client import UpstreamClient, UpstreamError, build_snapshot, build_new_api_token_name, mask_secret, _extract_key_value
+from app.services.upstream_client import UpstreamClient, UpstreamError, _PendingKeyError, build_snapshot, build_new_api_token_name, mask_secret, _extract_key_value
 from app.services.auth_config import MASK, mask_auth_config, normalize_auth_config
 from app.services.snapshot_service import diff_snapshots
 from app.services import scheduler as sched_svc
@@ -415,7 +415,8 @@ def _ensure_group_key(
            )
            if row and row.key_name:
                # 本地已有记录（无论 key_id 是否存在），优先用本地已存的 key_name 查远端
-                # 这样旧格式 key_name 也能找回远端对应 token，不重复创建
+                # 这样旧格式 key_name 也能找回远端对应 token，不重复创建。
+                # created_pending_key 记录也走同一路径，尝试回填 key/id。
                lookup_name = row.key_name
                found_by_name = lookup_name  # 记录实际命中的远端名称
                try:
@@ -442,13 +443,15 @@ def _ensure_group_key(
                    if found_by_name != lookup_name:
                        row.key_name = found_by_name
                    row.raw_json = json.dumps(existing, ensure_ascii=False)
+                    # pending 行回填成功后升级为 exists
                    row.status = "exists"
+                    row.error = None
                    row.updated_at = datetime.now(timezone.utc)
                    db.add(row)
                    db.commit()
                    db.refresh(row)
-                    return _key_response(row, include_value=False)
-                # 远端不存在，需要重新创建
+                    return _key_response(row, include_value=bool(key_value))
+                # 远端不存在，需要重新创建（包括 pending 行远端真的丢了）
                row.status = "replaced"

            # 2. 查远端是否有同名 Key
@@ -507,16 +510,43 @@ def _ensure_group_key(
                return _key_response(row, include_value=False)

            # 3. 远端不存在，创建新 Key（使用新短名）
-            created = client.create_api_key(
-                stable_name,
-                gid,
-                quota=body.quota,
-                expires_in_days=body.expires_in_days,
-                rate_limit_5h=body.rate_limit_5h,
-                rate_limit_1d=body.rate_limit_1d,
-                rate_limit_7d=body.rate_limit_7d,
-                endpoint=body.endpoint,
-            )
+            try:
+                created = client.create_api_key(
+                    stable_name,
+                    gid,
+                    quota=body.quota,
+                    expires_in_days=body.expires_in_days,
+                    rate_limit_5h=body.rate_limit_5h,
+                    rate_limit_1d=body.rate_limit_1d,
+                    rate_limit_7d=body.rate_limit_7d,
+                    endpoint=body.endpoint,
+                )
+            except _PendingKeyError as pke:
+                # POST 已成功但无法取得明文 key；落库为 pending，避免重复创建
+                pending_err = str(pke)
+                if row:
+                    row.key_id = row.key_id  # 保留原有（可能为空）
+                    row.key_name = stable_name
+                    row.key_value = ""
+                    row.managed_prefix = prefix
+                    row.status = "created_pending_key"
+                    row.error = pending_err
+                else:
+                    row = UpstreamGeneratedKey(
+                        upstream_id=upstream.id,
+                        group_id=gid,
+                        group_name=gname,
+                        key_name=stable_name,
+                        key_value="",
+                        managed_prefix=prefix,
+                        status="created_pending_key",
+                        error=pending_err,
+                    )
+                    db.add(row)
+                db.commit()
+                db.refresh(row)
+                return _key_response(row, include_value=False)
+
            if row:
                # 复用旧行
                row.key_id = created.get("id") or None
@@ -593,6 +623,16 @@ def generate_keys_by_groups(
        except Exception as exc:
            raise HTTPException(502, str(exc))

+        # New-API/Nox 上游：批量前预热 token 列表缓存（一次 /api/token/ 拉取），
+        # 避免每个分组都调 /api/token/search 触发 Nox 搜索限流（10 次/60 秒）。
+        if _is_new_api_user_upstream(u) or _is_nox_api_upstream(u):
+            try:
+                client.warm_token_list_cache()
+            except Exception as exc:
+                logger.warning(
+                    "generate_keys_by_groups: warm_token_list_cache failed for %s: %s", uid, exc
+                )
+
        for group in groups:
            gid = _group_id(group)
            if not gid or (selected and gid not in selected):
@@ -602,13 +642,17 @@ def generate_keys_by_groups(

    created = len([item for item in results if item.status == "created"])
    existed = len([item for item in results if item.status == "exists"])
+    pending = len([item for item in results if item.status == "created_pending_key"])
    total = len(results)
    msg_parts = []
    if created:
        msg_parts.append(f"新创建 {created}")
    if existed:
        msg_parts.append(f"已存在 {existed}")
+    if pending:
+        msg_parts.append(f"待回填 {pending}")
    msg = "、".join(msg_parts) + f" / 共 {total} 个分组" if msg_parts else f"共处理 {total} 个分组"
+    # created_pending_key 视为非失败（POST 已成功，只是 key 待回填）
    return GenerateKeysByGroupsResponse(
        success=total > 0 and all(item.status != "failed" for item in results),
        message=msg,