feat(nox-api): implement batch creation rate limit resilience and pending backfilling
This commit is contained in:
@@ -24,7 +24,7 @@ from app.schemas.upstream import (
|
|||||||
UpstreamCreate, UpstreamUpdate, UpstreamResponse, SnapshotResponse, TestResult,
|
UpstreamCreate, UpstreamUpdate, UpstreamResponse, SnapshotResponse, TestResult,
|
||||||
UpstreamBatchActionItem, UpstreamBatchActionSummary, UpstreamBatchActionResponse,
|
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.auth_config import MASK, mask_auth_config, normalize_auth_config
|
||||||
from app.services.snapshot_service import diff_snapshots
|
from app.services.snapshot_service import diff_snapshots
|
||||||
from app.services import scheduler as sched_svc
|
from app.services import scheduler as sched_svc
|
||||||
@@ -415,7 +415,8 @@ def _ensure_group_key(
|
|||||||
)
|
)
|
||||||
if row and row.key_name:
|
if row and row.key_name:
|
||||||
# 本地已有记录(无论 key_id 是否存在),优先用本地已存的 key_name 查远端
|
# 本地已有记录(无论 key_id 是否存在),优先用本地已存的 key_name 查远端
|
||||||
# 这样旧格式 key_name 也能找回远端对应 token,不重复创建
|
# 这样旧格式 key_name 也能找回远端对应 token,不重复创建。
|
||||||
|
# created_pending_key 记录也走同一路径,尝试回填 key/id。
|
||||||
lookup_name = row.key_name
|
lookup_name = row.key_name
|
||||||
found_by_name = lookup_name # 记录实际命中的远端名称
|
found_by_name = lookup_name # 记录实际命中的远端名称
|
||||||
try:
|
try:
|
||||||
@@ -442,13 +443,15 @@ def _ensure_group_key(
|
|||||||
if found_by_name != lookup_name:
|
if found_by_name != lookup_name:
|
||||||
row.key_name = found_by_name
|
row.key_name = found_by_name
|
||||||
row.raw_json = json.dumps(existing, ensure_ascii=False)
|
row.raw_json = json.dumps(existing, ensure_ascii=False)
|
||||||
|
# pending 行回填成功后升级为 exists
|
||||||
row.status = "exists"
|
row.status = "exists"
|
||||||
|
row.error = None
|
||||||
row.updated_at = datetime.now(timezone.utc)
|
row.updated_at = datetime.now(timezone.utc)
|
||||||
db.add(row)
|
db.add(row)
|
||||||
db.commit()
|
db.commit()
|
||||||
db.refresh(row)
|
db.refresh(row)
|
||||||
return _key_response(row, include_value=False)
|
return _key_response(row, include_value=bool(key_value))
|
||||||
# 远端不存在,需要重新创建
|
# 远端不存在,需要重新创建(包括 pending 行远端真的丢了)
|
||||||
row.status = "replaced"
|
row.status = "replaced"
|
||||||
|
|
||||||
# 2. 查远端是否有同名 Key
|
# 2. 查远端是否有同名 Key
|
||||||
@@ -507,6 +510,7 @@ def _ensure_group_key(
|
|||||||
return _key_response(row, include_value=False)
|
return _key_response(row, include_value=False)
|
||||||
|
|
||||||
# 3. 远端不存在,创建新 Key(使用新短名)
|
# 3. 远端不存在,创建新 Key(使用新短名)
|
||||||
|
try:
|
||||||
created = client.create_api_key(
|
created = client.create_api_key(
|
||||||
stable_name,
|
stable_name,
|
||||||
gid,
|
gid,
|
||||||
@@ -517,6 +521,32 @@ def _ensure_group_key(
|
|||||||
rate_limit_7d=body.rate_limit_7d,
|
rate_limit_7d=body.rate_limit_7d,
|
||||||
endpoint=body.endpoint,
|
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:
|
if row:
|
||||||
# 复用旧行
|
# 复用旧行
|
||||||
row.key_id = created.get("id") or None
|
row.key_id = created.get("id") or None
|
||||||
@@ -593,6 +623,16 @@ def generate_keys_by_groups(
|
|||||||
except Exception as exc:
|
except Exception as exc:
|
||||||
raise HTTPException(502, str(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:
|
for group in groups:
|
||||||
gid = _group_id(group)
|
gid = _group_id(group)
|
||||||
if not gid or (selected and gid not in selected):
|
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"])
|
created = len([item for item in results if item.status == "created"])
|
||||||
existed = len([item for item in results if item.status == "exists"])
|
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)
|
total = len(results)
|
||||||
msg_parts = []
|
msg_parts = []
|
||||||
if created:
|
if created:
|
||||||
msg_parts.append(f"新创建 {created}")
|
msg_parts.append(f"新创建 {created}")
|
||||||
if existed:
|
if existed:
|
||||||
msg_parts.append(f"已存在 {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} 个分组"
|
msg = "、".join(msg_parts) + f" / 共 {total} 个分组" if msg_parts else f"共处理 {total} 个分组"
|
||||||
|
# created_pending_key 视为非失败(POST 已成功,只是 key 待回填)
|
||||||
return GenerateKeysByGroupsResponse(
|
return GenerateKeysByGroupsResponse(
|
||||||
success=total > 0 and all(item.status != "failed" for item in results),
|
success=total > 0 and all(item.status != "failed" for item in results),
|
||||||
message=msg,
|
message=msg,
|
||||||
|
|||||||
@@ -16,6 +16,14 @@ class UpstreamError(RuntimeError):
|
|||||||
pass
|
pass
|
||||||
|
|
||||||
|
|
||||||
|
class _PendingKeyError(UpstreamError):
|
||||||
|
"""POST /api/token/ 成功但最终无法取得明文 key 时抛出。
|
||||||
|
|
||||||
|
路由层捕获后将记录落库为 created_pending_key,而不是 failed。
|
||||||
|
"""
|
||||||
|
pass
|
||||||
|
|
||||||
|
|
||||||
NEW_API_DEFAULT_QUOTA_PER_UNIT = 500000
|
NEW_API_DEFAULT_QUOTA_PER_UNIT = 500000
|
||||||
|
|
||||||
# New-API/Nox-API token name 上限(UTF-8 字节)
|
# New-API/Nox-API token name 上限(UTF-8 字节)
|
||||||
@@ -367,6 +375,47 @@ def build_snapshot(upstream_id: int, base_url: str, api_prefix: str,
|
|||||||
}
|
}
|
||||||
|
|
||||||
|
|
||||||
|
# 429 重试配置(秒)
|
||||||
|
_RATE_LIMIT_BACKOFFS = (2, 5, 10, 20)
|
||||||
|
_RATE_LIMIT_MAX_WAIT = 60
|
||||||
|
|
||||||
|
|
||||||
|
def _retry_on_429(fn: Callable[[], Any], backoffs: tuple[int, ...] = _RATE_LIMIT_BACKOFFS) -> Any:
|
||||||
|
"""执行 fn(),若遇 429 则按退避序列重试,超限后抛 UpstreamError。
|
||||||
|
|
||||||
|
fn 应是 HTTP 调用;429 响应会被 httpx.HTTPStatusError 捕获(raise_for_status 之后)
|
||||||
|
或者 fn 内部已经解析为 UpstreamError 且 message 含 429。
|
||||||
|
"""
|
||||||
|
total_waited = 0
|
||||||
|
for attempt, wait in enumerate(backoffs):
|
||||||
|
try:
|
||||||
|
return fn()
|
||||||
|
except Exception as exc:
|
||||||
|
status = None
|
||||||
|
retry_after: int | None = None
|
||||||
|
# httpx HTTPStatusError
|
||||||
|
if hasattr(exc, "response") and hasattr(exc.response, "status_code"):
|
||||||
|
status = exc.response.status_code
|
||||||
|
try:
|
||||||
|
retry_after = int(exc.response.headers.get("Retry-After", 0))
|
||||||
|
except (TypeError, ValueError):
|
||||||
|
retry_after = None
|
||||||
|
# UpstreamError wrapping a 429 text
|
||||||
|
if status != 429 and "429" in str(exc):
|
||||||
|
status = 429
|
||||||
|
if status != 429:
|
||||||
|
raise
|
||||||
|
sleep_time = retry_after if (retry_after and retry_after > 0) else wait
|
||||||
|
if total_waited + sleep_time > _RATE_LIMIT_MAX_WAIT:
|
||||||
|
raise UpstreamError(
|
||||||
|
f"rate limited (429) after {total_waited}s total wait; giving up"
|
||||||
|
) from exc
|
||||||
|
time.sleep(sleep_time)
|
||||||
|
total_waited += sleep_time
|
||||||
|
# 最后一次尝试,不捕获
|
||||||
|
return fn()
|
||||||
|
|
||||||
|
|
||||||
class UpstreamClient:
|
class UpstreamClient:
|
||||||
"""Sync HTTP client that handles all auth types."""
|
"""Sync HTTP client that handles all auth types."""
|
||||||
|
|
||||||
@@ -389,6 +438,8 @@ class UpstreamClient:
|
|||||||
self._cookies: dict[str, str] = {}
|
self._cookies: dict[str, str] = {}
|
||||||
self._new_api_user: str = ""
|
self._new_api_user: str = ""
|
||||||
self._client = httpx.Client(timeout=timeout)
|
self._client = httpx.Client(timeout=timeout)
|
||||||
|
# 批量生成期间的 token 列表缓存(name -> record),避免重复调用 /api/token/search
|
||||||
|
self._token_list_cache: dict[str, dict[str, Any]] | None = None
|
||||||
|
|
||||||
def close(self) -> None:
|
def close(self) -> None:
|
||||||
self._client.close()
|
self._client.close()
|
||||||
@@ -727,6 +778,30 @@ class UpstreamClient:
|
|||||||
raise UpstreamError("New-API token key response did not include key")
|
raise UpstreamError("New-API token key response did not include key")
|
||||||
return key_value
|
return key_value
|
||||||
|
|
||||||
|
def warm_token_list_cache(self) -> None:
|
||||||
|
"""预热 token 列表缓存;批量生成前调用一次,后续 find_smartup_group_key 走缓存。
|
||||||
|
|
||||||
|
缓存以 token name 为 key;同名取最新(id 最大)。
|
||||||
|
"""
|
||||||
|
all_tokens = self._list_all_new_api_tokens()
|
||||||
|
cache: dict[str, dict[str, Any]] = {}
|
||||||
|
for t in all_tokens:
|
||||||
|
name = str(t.get("name") or t.get("key_name") or "")
|
||||||
|
if not name:
|
||||||
|
continue
|
||||||
|
existing = cache.get(name)
|
||||||
|
if existing is None or (t.get("id") or 0) > (existing.get("id") or 0):
|
||||||
|
cache[name] = t
|
||||||
|
self._token_list_cache = cache
|
||||||
|
|
||||||
|
def _cache_add_token(self, record: dict[str, Any]) -> None:
|
||||||
|
"""创建成功后把新 token 追加进缓存,让同批次后续查询能感知它。"""
|
||||||
|
if self._token_list_cache is None:
|
||||||
|
return
|
||||||
|
name = str(record.get("name") or record.get("key_name") or "")
|
||||||
|
if name:
|
||||||
|
self._token_list_cache[name] = record
|
||||||
|
|
||||||
def _create_new_api_token(
|
def _create_new_api_token(
|
||||||
self,
|
self,
|
||||||
name: str,
|
name: str,
|
||||||
@@ -734,6 +809,12 @@ class UpstreamClient:
|
|||||||
quota: float = 0,
|
quota: float = 0,
|
||||||
expires_in_days: int | None = None,
|
expires_in_days: int | None = None,
|
||||||
) -> dict[str, Any]:
|
) -> dict[str, Any]:
|
||||||
|
"""创建 Nox/New-API token。
|
||||||
|
|
||||||
|
主路径:POST 响应里直接有明文 key → 立即返回,不调用 search。
|
||||||
|
Fallback:响应无 key → 按名称查 token 列表再按 id 取明文 key;查询带 429 退避重试。
|
||||||
|
极端情况:POST 已成功但最终无法取得 key → 抛 _PendingKeyError,调用方应落库为 pending。
|
||||||
|
"""
|
||||||
unlimited = quota <= 0
|
unlimited = quota <= 0
|
||||||
body: dict[str, Any] = {
|
body: dict[str, Any] = {
|
||||||
"name": name,
|
"name": name,
|
||||||
@@ -749,14 +830,62 @@ class UpstreamClient:
|
|||||||
payload = self._request("POST", "/api/token/", body)
|
payload = self._request("POST", "/api/token/", body)
|
||||||
self._ensure_api_success(payload, "create New-API token")
|
self._ensure_api_success(payload, "create New-API token")
|
||||||
|
|
||||||
matches = self._list_new_api_tokens(search=name, group_id=group_id)
|
# ── 主路径:从创建响应直接取明文 key(Nox-API 当前行为:data="sk-...")──
|
||||||
token = next((i for i in matches if str(i.get("name") or "").strip() == name.strip()), None)
|
create_data = _unwrap_data(payload)
|
||||||
|
key_from_response = _extract_key_value(create_data) if create_data else ""
|
||||||
|
token_id_from_response = _extract_id(create_data) if isinstance(create_data, dict) else ""
|
||||||
|
|
||||||
|
if key_from_response and "*" not in key_from_response:
|
||||||
|
record = {
|
||||||
|
"name": name,
|
||||||
|
"group": str(group_id),
|
||||||
|
"group_id": str(group_id),
|
||||||
|
"key": key_from_response,
|
||||||
|
"id": token_id_from_response or None,
|
||||||
|
}
|
||||||
|
if isinstance(create_data, dict):
|
||||||
|
record.update(create_data)
|
||||||
|
self._cache_add_token(record)
|
||||||
|
return {
|
||||||
|
"id": token_id_from_response or "",
|
||||||
|
"key": key_from_response,
|
||||||
|
"masked_key": mask_secret(key_from_response),
|
||||||
|
"raw": record,
|
||||||
|
}
|
||||||
|
|
||||||
|
# ── Fallback:按名称找 token,再按 id 取明文 key;带 429 退避重试 ──
|
||||||
|
try:
|
||||||
|
matches = _retry_on_429(
|
||||||
|
lambda: self._list_new_api_tokens(search=name, group_id=group_id)
|
||||||
|
)
|
||||||
|
except Exception as exc:
|
||||||
|
raise _PendingKeyError(
|
||||||
|
f"POST /api/token/ 成功,但查询 token 列表失败({exc});"
|
||||||
|
"key 待下次回填"
|
||||||
|
) from exc
|
||||||
|
|
||||||
|
token = next(
|
||||||
|
(i for i in matches if str(i.get("name") or "").strip() == name.strip()), None
|
||||||
|
)
|
||||||
if not token:
|
if not token:
|
||||||
raise UpstreamError("New-API token was created but could not be found by name")
|
raise _PendingKeyError(
|
||||||
|
"POST /api/token/ 成功,但列表中找不到该 token;key 待下次回填"
|
||||||
|
)
|
||||||
|
|
||||||
token_id = token.get("id")
|
token_id = token.get("id")
|
||||||
if token_id is None:
|
if token_id is None:
|
||||||
raise UpstreamError("New-API token list response did not include id")
|
raise _PendingKeyError(
|
||||||
key_value = self._get_new_api_token_key(token_id)
|
"POST /api/token/ 成功,但 token 无 id 无法取明文 key;key 待下次回填"
|
||||||
|
)
|
||||||
|
|
||||||
|
try:
|
||||||
|
key_value = _retry_on_429(lambda: self._get_new_api_token_key(token_id))
|
||||||
|
except Exception as exc:
|
||||||
|
raise _PendingKeyError(
|
||||||
|
f"POST /api/token/ 成功,但取明文 key 失败({exc});key 待下次回填"
|
||||||
|
) from exc
|
||||||
|
|
||||||
|
self._cache_add_token({**self._normalize_key_record(token), "key": key_value})
|
||||||
return {
|
return {
|
||||||
"id": str(token_id),
|
"id": str(token_id),
|
||||||
"key": key_value,
|
"key": key_value,
|
||||||
@@ -892,9 +1021,21 @@ class UpstreamClient:
|
|||||||
) -> dict[str, Any] | None:
|
) -> dict[str, Any] | None:
|
||||||
"""查找同一上游分组下是否已存在 SmartUp 前缀的 Key。
|
"""查找同一上游分组下是否已存在 SmartUp 前缀的 Key。
|
||||||
|
|
||||||
匹配规则:key_name 等于 expected_name,且以 prefix 开头。
|
优先走本地 token 列表缓存(warm_token_list_cache 预热后有效),
|
||||||
返回匹配到的第一个 Key,或 None。
|
避免在批量生成时高频调用 /api/token/search 触发 Nox 限流。
|
||||||
|
缓存未命中时回退到远端搜索。
|
||||||
"""
|
"""
|
||||||
|
# ── 缓存命中路径 ──
|
||||||
|
if self._token_list_cache is not None:
|
||||||
|
record = self._token_list_cache.get(expected_name)
|
||||||
|
if record is not None:
|
||||||
|
rec_group = str(record.get("group_id") or record.get("group") or "")
|
||||||
|
if not rec_group or rec_group == str(group_id):
|
||||||
|
return self._hydrate_new_api_token_key(record)
|
||||||
|
# 缓存中无此名 → 远端也不存在(批次内强一致性假设)
|
||||||
|
return None
|
||||||
|
|
||||||
|
# ── 无缓存:走远端搜索(原有逻辑)──
|
||||||
gid = int(group_id) if str(group_id).isdigit() else group_id
|
gid = int(group_id) if str(group_id).isdigit() else group_id
|
||||||
keys = self.list_api_keys(search=prefix, group_id=gid, status="active")
|
keys = self.list_api_keys(search=prefix, group_id=gid, status="active")
|
||||||
for k in keys:
|
for k in keys:
|
||||||
|
|||||||
@@ -1219,3 +1219,354 @@ def test_p3_local_old_key_name_synced_to_new_short_name_when_remote_migrated(db_
|
|||||||
)
|
)
|
||||||
assert row.key_id == "remote-new-id"
|
assert row.key_id == "remote-new-id"
|
||||||
assert row.key_value == "sk-new-plain"
|
assert row.key_value == "sk-new-plain"
|
||||||
|
|
||||||
|
|
||||||
|
# ──────────────────────────────────────────────
|
||||||
|
# Nox-API Batch / Rate Limit Resilience Tests
|
||||||
|
# ──────────────────────────────────────────────
|
||||||
|
|
||||||
|
def test_nox_create_token_returns_key_directly(monkeypatch):
|
||||||
|
"""Nox POST /api/token/ 直接在 data 中返回 sk-... 时,应直接使用,不查询列表。"""
|
||||||
|
from app.services.upstream_client import UpstreamClient
|
||||||
|
|
||||||
|
client = UpstreamClient(
|
||||||
|
base_url="http://nox.local",
|
||||||
|
api_prefix="",
|
||||||
|
auth_type="nox_token",
|
||||||
|
auth_config={"token": "abc", "new_api_user": "7"},
|
||||||
|
)
|
||||||
|
request_paths = []
|
||||||
|
|
||||||
|
def fake_request(method, path, body=None, auth=True):
|
||||||
|
request_paths.append((method, path))
|
||||||
|
if method == "GET" and path == "/api/status":
|
||||||
|
return {"success": True, "data": {"quota_per_unit": 500000}}
|
||||||
|
if method == "POST" and path == "/api/token/":
|
||||||
|
# Nox-API 直接返回明文 key
|
||||||
|
return {"success": True, "data": {"key": "sk-nox-direct-plaintext-123", "id": 999}}
|
||||||
|
raise AssertionError(f"unexpected request {method} {path}")
|
||||||
|
|
||||||
|
monkeypatch.setattr(client, "_request", fake_request)
|
||||||
|
|
||||||
|
# 预热一个空缓存
|
||||||
|
client._token_list_cache = {}
|
||||||
|
|
||||||
|
result = client.create_api_key(
|
||||||
|
"SmartUp-7-group-hash",
|
||||||
|
"group-id",
|
||||||
|
quota=2,
|
||||||
|
endpoint="/api/token",
|
||||||
|
)
|
||||||
|
|
||||||
|
assert result["id"] == "999"
|
||||||
|
assert result["key"] == "sk-nox-direct-plaintext-123"
|
||||||
|
# 不应该有 GET /api/token/search 或 GET /api/token/ 列表或 POST /api/token/999/key 动作
|
||||||
|
assert ("GET", "/api/token/") not in request_paths
|
||||||
|
assert ("GET", "/api/token/search") not in request_paths
|
||||||
|
assert ("POST", "/api/token/999/key") not in request_paths
|
||||||
|
|
||||||
|
|
||||||
|
def test_nox_create_token_no_key_fallback_success(monkeypatch):
|
||||||
|
"""创建响应无 key 时,通过列表和 get-key fallback 补齐。"""
|
||||||
|
from app.services.upstream_client import UpstreamClient
|
||||||
|
|
||||||
|
client = UpstreamClient(
|
||||||
|
base_url="http://nox.local",
|
||||||
|
api_prefix="",
|
||||||
|
auth_type="nox_token",
|
||||||
|
auth_config={"token": "abc", "new_api_user": "7"},
|
||||||
|
)
|
||||||
|
request_paths = []
|
||||||
|
|
||||||
|
def fake_request(method, path, body=None, auth=True):
|
||||||
|
request_paths.append((method, path))
|
||||||
|
if method == "GET" and path == "/api/status":
|
||||||
|
return {"success": True, "data": {"quota_per_unit": 500000}}
|
||||||
|
if method == "POST" and path == "/api/token/":
|
||||||
|
# 不带明文 key
|
||||||
|
return {"success": True, "data": {"id": 888}}
|
||||||
|
if method == "POST" and path == "/api/token/888/key":
|
||||||
|
return {"success": True, "data": {"key": "sk-nox-fallback-key"}}
|
||||||
|
raise AssertionError(f"unexpected request {method} {path}")
|
||||||
|
|
||||||
|
monkeypatch.setattr(client, "_request", fake_request)
|
||||||
|
|
||||||
|
# 模拟 _list_new_api_tokens 行为
|
||||||
|
monkeypatch.setattr(
|
||||||
|
client,
|
||||||
|
"_list_new_api_tokens",
|
||||||
|
lambda search, group_id: [{"id": 888, "name": search, "group": group_id, "key": "sk-****"}]
|
||||||
|
)
|
||||||
|
|
||||||
|
result = client.create_api_key(
|
||||||
|
"SmartUp-7-group-hash",
|
||||||
|
"group-id",
|
||||||
|
quota=2,
|
||||||
|
endpoint="/api/token",
|
||||||
|
)
|
||||||
|
|
||||||
|
assert result["id"] == "888"
|
||||||
|
assert result["key"] == "sk-nox-fallback-key"
|
||||||
|
|
||||||
|
|
||||||
|
def test_nox_create_token_fallback_429_retry(monkeypatch):
|
||||||
|
"""fallback 遇 429 时按退避重试成功。"""
|
||||||
|
from app.services.upstream_client import UpstreamClient, UpstreamError
|
||||||
|
import httpx
|
||||||
|
import pytest
|
||||||
|
|
||||||
|
client = UpstreamClient(
|
||||||
|
base_url="http://nox.local",
|
||||||
|
api_prefix="",
|
||||||
|
auth_type="nox_token",
|
||||||
|
auth_config={"token": "abc", "new_api_user": "7"},
|
||||||
|
)
|
||||||
|
|
||||||
|
request_counts = {"list": 0, "key": 0}
|
||||||
|
|
||||||
|
def fake_request(method, path, body=None, auth=True):
|
||||||
|
if method == "GET" and path == "/api/status":
|
||||||
|
return {"success": True, "data": {"quota_per_unit": 500000}}
|
||||||
|
if method == "POST" and path == "/api/token/":
|
||||||
|
return {"success": True, "data": {"id": 888}}
|
||||||
|
if method == "POST" and path == "/api/token/888/key":
|
||||||
|
request_counts["key"] += 1
|
||||||
|
if request_counts["key"] == 1:
|
||||||
|
# 第一次返回 429
|
||||||
|
response = httpx.Response(429, request=httpx.Request("POST", "http://nox.local"))
|
||||||
|
raise httpx.HTTPStatusError("429 Too Many Requests", request=response.request, response=response)
|
||||||
|
return {"success": True, "data": {"key": "sk-nox-retry-success"}}
|
||||||
|
raise AssertionError(f"unexpected request {method} {path}")
|
||||||
|
|
||||||
|
monkeypatch.setattr(client, "_request", fake_request)
|
||||||
|
|
||||||
|
def fake_list(search, group_id):
|
||||||
|
request_counts["list"] += 1
|
||||||
|
if request_counts["list"] == 1:
|
||||||
|
raise UpstreamError("status code: 429") # 第一次列表调用也模拟 429
|
||||||
|
return [{"id": 888, "name": search, "group": group_id, "key": "sk-****"}]
|
||||||
|
|
||||||
|
monkeypatch.setattr(client, "_list_new_api_tokens", fake_list)
|
||||||
|
|
||||||
|
# 缩短重试延迟以加快测试运行
|
||||||
|
import app.services.upstream_client
|
||||||
|
monkeypatch.setattr(app.services.upstream_client, "_RATE_LIMIT_BACKOFFS", (0.01, 0.02))
|
||||||
|
|
||||||
|
result = client.create_api_key(
|
||||||
|
"SmartUp-7-group-hash",
|
||||||
|
"group-id",
|
||||||
|
quota=2,
|
||||||
|
endpoint="/api/token",
|
||||||
|
)
|
||||||
|
|
||||||
|
assert result["id"] == "888"
|
||||||
|
assert result["key"] == "sk-nox-retry-success"
|
||||||
|
assert request_counts["list"] == 2
|
||||||
|
assert request_counts["key"] == 2
|
||||||
|
|
||||||
|
|
||||||
|
def test_nox_create_token_fallback_fails_raises_pending(monkeypatch):
|
||||||
|
"""POST 成功但 fallback 获取 key 429 超限失败时,应抛出 _PendingKeyError。"""
|
||||||
|
from app.services.upstream_client import UpstreamClient, _PendingKeyError
|
||||||
|
import httpx
|
||||||
|
import pytest
|
||||||
|
|
||||||
|
client = UpstreamClient(
|
||||||
|
base_url="http://nox.local",
|
||||||
|
api_prefix="",
|
||||||
|
auth_type="nox_token",
|
||||||
|
auth_config={"token": "abc", "new_api_user": "7"},
|
||||||
|
)
|
||||||
|
|
||||||
|
def fake_request(method, path, body=None, auth=True):
|
||||||
|
if method == "GET" and path == "/api/status":
|
||||||
|
return {"success": True, "data": {"quota_per_unit": 500000}}
|
||||||
|
if method == "POST" and path == "/api/token/":
|
||||||
|
return {"success": True, "data": {"id": 888}}
|
||||||
|
raise AssertionError(f"unexpected request {method} {path}")
|
||||||
|
|
||||||
|
monkeypatch.setattr(client, "_request", fake_request)
|
||||||
|
|
||||||
|
def fake_list(search, group_id):
|
||||||
|
response = httpx.Response(429, request=httpx.Request("GET", "http://nox.local"))
|
||||||
|
raise httpx.HTTPStatusError("429 Too Many Requests", request=response.request, response=response)
|
||||||
|
|
||||||
|
monkeypatch.setattr(client, "_list_new_api_tokens", fake_list)
|
||||||
|
|
||||||
|
# 缩短重试延迟以加快测试运行
|
||||||
|
import app.services.upstream_client
|
||||||
|
monkeypatch.setattr(app.services.upstream_client, "_RATE_LIMIT_BACKOFFS", (0.01, 0.02))
|
||||||
|
|
||||||
|
with pytest.raises(_PendingKeyError, match="POST /api/token/ 成功,但查询 token 列表失败"):
|
||||||
|
client.create_api_key(
|
||||||
|
"SmartUp-7-group-hash",
|
||||||
|
"group-id",
|
||||||
|
quota=2,
|
||||||
|
endpoint="/api/token",
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def test_ensure_group_key_nox_batch_uses_cache(db_session, monkeypatch):
|
||||||
|
"""验证批量生成 Nox 分组时,使用了缓存,不重复触发 search 请求。"""
|
||||||
|
from app.routers.upstreams import _ensure_group_key
|
||||||
|
from app.schemas.upstream import GenerateKeysByGroupsRequest
|
||||||
|
from app.services.upstream_client import UpstreamClient
|
||||||
|
|
||||||
|
upstream = Upstream(
|
||||||
|
name="NoxAPI",
|
||||||
|
base_url="http://nox.local",
|
||||||
|
api_prefix="",
|
||||||
|
auth_type="nox_token",
|
||||||
|
auth_config_json=json.dumps({"token": "abc", "new_api_user": "7"}),
|
||||||
|
groups_endpoint="/api/user/self/groups",
|
||||||
|
rate_endpoint="/api/user/self/groups",
|
||||||
|
)
|
||||||
|
db_session.add(upstream)
|
||||||
|
db_session.commit()
|
||||||
|
db_session.refresh(upstream)
|
||||||
|
|
||||||
|
request_calls = []
|
||||||
|
|
||||||
|
class MockUpstreamClient(UpstreamClient):
|
||||||
|
def _request(self, method, path, body=None, auth=True):
|
||||||
|
request_calls.append((method, path))
|
||||||
|
if method == "GET" and path == "/api/status":
|
||||||
|
return {"success": True, "data": {"quota_per_unit": 500000}}
|
||||||
|
if method == "POST" and path == "/api/token/":
|
||||||
|
# 直接在 POST 响应里返回包含 key 的数据以防再次 search/list
|
||||||
|
return {"success": True, "data": {"key": "sk-nox-direct-123", "id": 123}}
|
||||||
|
raise AssertionError(f"unexpected request {method} {path}")
|
||||||
|
|
||||||
|
client = MockUpstreamClient(
|
||||||
|
base_url="http://nox.local",
|
||||||
|
api_prefix="",
|
||||||
|
auth_type="nox_token",
|
||||||
|
auth_config={"token": "abc", "new_api_user": "7"},
|
||||||
|
)
|
||||||
|
# 模拟 pre-warm
|
||||||
|
client._token_list_cache = {}
|
||||||
|
|
||||||
|
# 批量模拟:2 个不同分组
|
||||||
|
groups = [
|
||||||
|
{"id": "grp1", "name": "Group1"},
|
||||||
|
{"id": "grp2", "name": "Group2"},
|
||||||
|
]
|
||||||
|
body = GenerateKeysByGroupsRequest(group_ids=["grp1", "grp2"], name_prefix="SmartUp", quota=0)
|
||||||
|
|
||||||
|
# 第一次跑 grp1 (应该创建)
|
||||||
|
res1 = _ensure_group_key(db_session, client, upstream, groups[0], "SmartUp", body)
|
||||||
|
assert res1.status == "created"
|
||||||
|
|
||||||
|
# 第二次跑 grp1 (应该从缓存中匹配 exists,不需要调 search 或任何 HTTP 调用)
|
||||||
|
res2 = _ensure_group_key(db_session, client, upstream, groups[0], "SmartUp", body)
|
||||||
|
assert res2.status == "exists"
|
||||||
|
|
||||||
|
# 仅在第一次创建时有 POST /api/token/ 请求,获取 status (1次) + 创建 (1次)
|
||||||
|
# 没有针对 exists/exists check 的任何 GET /api/token/search 或 list 调用
|
||||||
|
get_or_search_calls = [c for c in request_calls if c[0] == "GET" and "token" in c[1]]
|
||||||
|
assert len(get_or_search_calls) == 0, f"Expected no remote list/search calls, got {request_calls}"
|
||||||
|
|
||||||
|
|
||||||
|
def test_ensure_group_key_nox_create_fallback_pending_db_save(db_session, monkeypatch):
|
||||||
|
"""POST 成功但 fallback 获取 key 429 报错时,本地入库为 created_pending_key,不返回 failed。"""
|
||||||
|
from app.routers.upstreams import _ensure_group_key
|
||||||
|
from app.schemas.upstream import GenerateKeysByGroupsRequest
|
||||||
|
from app.services.upstream_client import _PendingKeyError
|
||||||
|
from app.models.upstream_key import UpstreamGeneratedKey
|
||||||
|
|
||||||
|
upstream = Upstream(
|
||||||
|
name="NoxAPI",
|
||||||
|
base_url="http://nox.local",
|
||||||
|
api_prefix="",
|
||||||
|
auth_type="nox_token",
|
||||||
|
auth_config_json=json.dumps({"token": "abc", "new_api_user": "7"}),
|
||||||
|
groups_endpoint="/api/user/self/groups",
|
||||||
|
rate_endpoint="/api/user/self/groups",
|
||||||
|
)
|
||||||
|
db_session.add(upstream)
|
||||||
|
db_session.commit()
|
||||||
|
db_session.refresh(upstream)
|
||||||
|
|
||||||
|
class MockClient:
|
||||||
|
def find_smartup_group_key(self, group_id, name, prefix):
|
||||||
|
return None
|
||||||
|
|
||||||
|
def create_api_key(self, name, group_id, **kwargs):
|
||||||
|
raise _PendingKeyError("simulate list/key 429 failure")
|
||||||
|
|
||||||
|
group = {"id": "pending-grp", "name": "PendingGrp"}
|
||||||
|
body = GenerateKeysByGroupsRequest(group_ids=["pending-grp"], name_prefix="SmartUp", quota=0)
|
||||||
|
|
||||||
|
result = _ensure_group_key(db_session, MockClient(), upstream, group, "SmartUp", body)
|
||||||
|
|
||||||
|
assert result.status == "created_pending_key"
|
||||||
|
assert "simulate list/key 429 failure" in (result.error or "")
|
||||||
|
|
||||||
|
# 验证本地数据库保存了记录,但 key_value 为空,以便后续重试
|
||||||
|
row = db_session.query(UpstreamGeneratedKey).filter(
|
||||||
|
UpstreamGeneratedKey.upstream_id == upstream.id,
|
||||||
|
UpstreamGeneratedKey.group_id == "pending-grp",
|
||||||
|
).one()
|
||||||
|
assert row.status == "created_pending_key"
|
||||||
|
assert row.key_value == ""
|
||||||
|
|
||||||
|
|
||||||
|
def test_ensure_group_key_nox_pending_re_invoke_backfills(db_session):
|
||||||
|
"""重试 pending 记录时,如果远端已经能查到,应回填 key_id/key_value 并升级为 exists。"""
|
||||||
|
from app.routers.upstreams import _ensure_group_key
|
||||||
|
from app.schemas.upstream import GenerateKeysByGroupsRequest
|
||||||
|
from app.models.upstream_key import UpstreamGeneratedKey
|
||||||
|
|
||||||
|
upstream = Upstream(
|
||||||
|
name="NoxAPI",
|
||||||
|
base_url="http://nox.local",
|
||||||
|
api_prefix="",
|
||||||
|
auth_type="nox_token",
|
||||||
|
auth_config_json=json.dumps({"token": "abc", "new_api_user": "7"}),
|
||||||
|
groups_endpoint="/api/user/self/groups",
|
||||||
|
rate_endpoint="/api/user/self/groups",
|
||||||
|
)
|
||||||
|
db_session.add(upstream)
|
||||||
|
db_session.commit()
|
||||||
|
db_session.refresh(upstream)
|
||||||
|
|
||||||
|
# 模拟本地有一条之前遗留的 pending 记录,key_id 和 key_value 都为空
|
||||||
|
old_key_name = f"SmartUp-{upstream.id}-PendingGrp-pending-grp"
|
||||||
|
db_session.add(UpstreamGeneratedKey(
|
||||||
|
upstream_id=upstream.id,
|
||||||
|
group_id="pending-grp",
|
||||||
|
group_name="PendingGrp",
|
||||||
|
key_name=old_key_name,
|
||||||
|
key_value="",
|
||||||
|
masked_key="",
|
||||||
|
managed_prefix="SmartUp",
|
||||||
|
key_id=None,
|
||||||
|
status="created_pending_key",
|
||||||
|
))
|
||||||
|
db_session.commit()
|
||||||
|
|
||||||
|
class MockClient:
|
||||||
|
def find_smartup_group_key(self, group_id, name, prefix):
|
||||||
|
# 远端此时能够查到了
|
||||||
|
if name == old_key_name:
|
||||||
|
return {"id": "remote-backfill-id", "name": old_key_name, "key": "sk-backfilled-plain-key"}
|
||||||
|
return None
|
||||||
|
|
||||||
|
def create_api_key(self, *args, **kwargs):
|
||||||
|
raise AssertionError("create_api_key should not be called for backfilling pending record")
|
||||||
|
|
||||||
|
group = {"id": "pending-grp", "name": "PendingGrp"}
|
||||||
|
body = GenerateKeysByGroupsRequest(group_ids=["pending-grp"], name_prefix="SmartUp", quota=0)
|
||||||
|
|
||||||
|
result = _ensure_group_key(db_session, MockClient(), upstream, group, "SmartUp", body)
|
||||||
|
|
||||||
|
assert result.status == "exists"
|
||||||
|
|
||||||
|
# 验证本地记录已升级并补全
|
||||||
|
row = db_session.query(UpstreamGeneratedKey).filter(
|
||||||
|
UpstreamGeneratedKey.upstream_id == upstream.id,
|
||||||
|
UpstreamGeneratedKey.group_id == "pending-grp",
|
||||||
|
).one()
|
||||||
|
assert row.status == "exists"
|
||||||
|
assert row.key_id == "remote-backfill-id"
|
||||||
|
assert row.key_value == "sk-backfilled-plain-key"
|
||||||
|
|||||||
Reference in New Issue
Block a user