diff --git a/.env.example b/.env.example index e90e967..e8b2450 100644 --- a/.env.example +++ b/.env.example @@ -104,6 +104,16 @@ CODEBUDDY2API_MAX_CONCURRENT=64 # Disable only the declared model-capability preflight; other safeguards and Intl image merging remain. # 只关闭声明的模型能力预检；其他保护与国际站图片合并仍然生效。 # CODEBUDDY2API_MODEL_CAPABILITY_GUARD=true +# Routing scheduler: indexed reuses per-model shards, shadow records legacy/indexed differences +# while returning the legacy answer, legacy rescans every credential on each request. +# 路由调度：indexed 复用按模型分片（默认），shadow 记录新旧差异但仍返回旧结果，legacy 每次全量判断。 +# CODEBUDDY2API_ROUTING_SCHEDULER=indexed +# Catalog probing: representative probes one account per product per TTL cycle; account refreshes all. +# 目录探测：representative 每个产品每周期只探测一个代表账号（默认），account 逐个账号刷新。 +# CODEBUDDY2API_CATALOG_PROBE_MODE=representative +# Bounded background maintenance workers; 1 restores the previous serial behaviour. +# 后台维护并发数；设为 1 即退化为串行。 +# CODEBUDDY2API_HOUSEKEEP_CONCURRENCY=4 # Retired: this setting only warns; all persistent logs use logs.sqlite3. # 已停用：该项只会给出告警；所有持久日志都写入 logs.sqlite3。 diff --git a/app/admin_api.py b/app/admin_api.py index a1f064a..2282447 100644 --- a/app/admin_api.py +++ b/app/admin_api.py @@ -17,6 +17,7 @@ from starlette.responses import JSONResponse, Response from . import auth_oauth +from . import perf_metrics from .admin_auth import AdminAuth, AdminMiddleware, COOKIE_NAME, SESSION_TTL, error_response, same_origin from .control_store import ConflictError, validate_model from .credential_io import CredentialFileError, MAX_CREDENTIAL_BYTES, _valid_name, read_import_file @@ -338,9 +339,21 @@ def build_models(): "model_guard_locked": bool(resolved.get("model_guard", {}).get("locked")), "model_capability_guard": config.get("model_capability_guard", True), "model_capability_guard_locked": bool(resolved.get("model_capability_guard", {}).get("locked")), + "catalog": gateway.catalog_status(), }) return await run_in_threadpool(build_models) + @route("POST", "/admin/models/scan") + async def models_scan_start(request): + """Start one user-triggered full catalog scan; it is never a periodic task.""" + result = await run_in_threadpool(gateway.start_catalog_scan) + event("models.scan") + return JSONResponse(result) + + @route("GET", "/admin/models/scan") + async def models_scan_state(request): + return JSONResponse(gateway.catalog_scan_view()) + def checked_rule(source, data, *, creating=False): if "custom" in data: raise ValueError("custom 是只读字段") @@ -687,4 +700,14 @@ def build_dashboard(): result = await run_in_threadpool(build_dashboard) return result if isinstance(result, Response) else JSONResponse(result) + @route("GET", "/admin/performance") + async def performance_get(request): + """Read-only process counters; no routing state is exposed or mutated.""" + payload = perf_metrics.snapshot() + pool = config.get("cred_pool") + scheduler = getattr(pool, "_scheduler", None) + payload["scheduler"] = {"mode": config.get("routing_scheduler", "indexed"), + **(scheduler.report() if scheduler is not None else {"shards": [], "mismatches": []})} + return JSONResponse(payload) + return auth diff --git a/app/catalog_index.py b/app/catalog_index.py new file mode 100644 index 0000000..aba81ef --- /dev/null +++ b/app/catalog_index.py @@ -0,0 +1,77 @@ +"""Read-only indexes over published account catalogs, built once per catalog revision.""" +from __future__ import annotations + +from .model_catalog_view import INTERNATIONAL, share_models + + +def account_scope(account, scope="models"): + """Merge root candidates into selector models while retaining selector metadata.""" + picker = account.get("models") + if scope == "models" or picker is None: + return picker + seen = {item.get("id") for item in picker} + # Root entries supply missing names without replacing selector metadata. + return picker + [item for item in account.get("serves") or [] if item.get("id") not in seen] + + +class CatalogIndex: + """Precompute per-account and per-profile catalogs so request paths never rescan accounts. + + ``accounts`` is the published ``account_catalogs`` mapping; the index keeps that same + object so a caller can detect a republished catalog by identity. + """ + + SCOPES = ("models", "serves") + + def __init__(self, accounts, *, revision=0, policy_revision=0, is_enabled=None): + self.accounts = accounts if isinstance(accounts, dict) else {} + self.revision = revision + self.policy_revision = policy_revision + self._enabled = is_enabled or (lambda _identity: True) + self._account: dict[tuple[int, str], list[dict] | None] = {} + self._profile: dict[str, dict[str, tuple[dict, ...]]] = {} + self.profiles: tuple[str, ...] = tuple(dict.fromkeys( + account.get("profile") for account in self.accounts.values() if account.get("profile"))) + for scope in self.SCOPES: + sources = self._shared_sources(scope) + merged: dict[str, list[dict]] = {} + for identity, account in self.accounts.items(): + profile = account.get("profile") if isinstance(account, dict) else None + if not profile: + continue + items = self._effective(account, scope, sources) + if items is None: + continue + merged.setdefault(profile, []).extend(items) + self._profile[scope] = {name: tuple(items) for name, items in merged.items()} + + def _shared_sources(self, scope): + """Collect every enabled international account catalog once per scope.""" + return [(account["profile"], account_scope(account, scope)) + for identity, account in self.accounts.items() + if isinstance(account, dict) and account.get("profile") in INTERNATIONAL + and account.get("models") is not None and self._enabled(identity)] + + def _effective(self, account, scope, sources): + native = account_scope(account, scope) + if native is None: + return None + profile = account.get("profile") + items = share_models(native, profile, sources) if profile in INTERNATIONAL else native + self._account[(id(account), scope)] = items + return items + + def account_catalog(self, account, scope="models"): + """Return this account's effective catalog, or None when the index never saw it.""" + if not isinstance(account, dict): + return None + return self._account.get((id(account), scope)) + + def catalog(self, profile, scope="models", *, model_id=None): + """Return the merged catalog for one profile, or None when no account supplied one.""" + items = self._profile.get(scope, {}).get(profile) + if items is None: + return None + if model_id is None: + return list(items) + return [item for item in items if item.get("id") == model_id] diff --git a/app/catalog_probe.py b/app/catalog_probe.py new file mode 100644 index 0000000..fb87ad6 --- /dev/null +++ b/app/catalog_probe.py @@ -0,0 +1,160 @@ +"""Representative per-profile catalog probes with rotation, backoff and diff status.""" +from __future__ import annotations + +import math +import random +import threading +import time + +SCHEMA_VERSION = 1 +PROFILES = ("cn-cli", "cn-work", "intl-cli", "intl-work") +MAX_PROFILES = 8 +MAX_MODEL_IDS = 2000 +RETRY_BASE_S = 60 +RETRY_MAX_S = 1800 +_MAX_FAILURES = 16 +_FIELDS = {"source", "source_generation", "fetched_at", "retry_at", "failures", "status", "models"} + + +def _time(value): + return isinstance(value, (int, float)) and not isinstance(value, bool) and math.isfinite(value) and value >= 0 + + +def _token(value): + return isinstance(value, str) and 0 < len(value) <= 160 and value.isprintable() and not any(c.isspace() for c in value) + + +def valid_row(row): + """Validate one persisted profile probe row without ever touching credential material.""" + if not isinstance(row, dict) or set(row) - _FIELDS: + return False + if "source" in row and row["source"] is not None and not _token(row["source"]): + return False + if not _time(row.get("fetched_at")) or not _time(row.get("retry_at")): + return False + failures = row.get("failures", 0) + if type(failures) is not int or not 0 <= failures <= _MAX_FAILURES: + return False + if row.get("status") not in ("ok", "error", "fresh"): + return False + if "source_generation" in row and type(row["source_generation"]) is not int: + return False + models = row.get("models", []) + return (isinstance(models, list) and len(models) <= MAX_MODEL_IDS + and all(isinstance(model, str) and 0 < len(model) <= 160 for model in models)) + + +class AggregateCatalog: + """One representative source per product profile, with rotation and retry backoff.""" + + def __init__(self, *, ttl=6 * 3600, store=None, clock=time.time): + self.ttl = max(60.0, float(ttl or 0)) + self._store = store + self._clock = clock + self._lock = threading.Lock() + self._data = self._load() + + def _load(self): + raw = self._store.get("catalog_probe") if self._store is not None else None + if (not isinstance(raw, dict) or raw.get("version") != SCHEMA_VERSION + or not isinstance(raw.get("profiles"), dict)): + return {"version": SCHEMA_VERSION, "profiles": {}} + profiles = {name: dict(row) for name, row in raw["profiles"].items() + if name in PROFILES and valid_row(row)} + return {"version": SCHEMA_VERSION, "profiles": profiles} + + def _save(self): + if self._store is not None: + self._store.put_cache("catalog_probe", {"version": SCHEMA_VERSION, "profiles": self._data["profiles"]}) + + def row(self, profile): + """Return a copy of one profile row.""" + with self._lock: + return dict(self._data["profiles"].get(profile) or {}) + + def profiles(self): + """Return a copy of every profile row for the management view.""" + with self._lock: + return {name: dict(row) for name, row in self._data["profiles"].items()} + + def due(self, profile, now=None): + """Whether this profile may probe again; a retry backoff keeps failures quiet.""" + now = self._clock() if now is None else now + row = self.row(profile) + return not row or now >= float(row.get("retry_at") or 0.0) + + @staticmethod + def _age(candidate, age_of): + """Treat an unknown catalog age as infinitely stale so a new account is probed.""" + value = age_of(candidate) + return math.inf if value is None else float(value) + + def choose(self, profile, candidates, *, age_of, now=None): + """Select the account to probe: unknown catalogs first, then the source, else the stalest.""" + now = self._clock() if now is None else now + due = [candidate for candidate in candidates if self._age(candidate, age_of) >= self.ttl] + if not due: + return None + # An account without any catalog cannot be routed at all, so it outranks reuse. + unknown = [candidate for candidate in due if self._age(candidate, age_of) == math.inf] + if unknown: + return unknown[0] + source = self.row(profile).get("source") + for candidate in due: + if candidate["account_key"] == source: + return candidate + # Oldest catalogs first, so rotation eventually covers every account. + return max(due, key=lambda candidate: self._age(candidate, age_of)) + + def success(self, profile, *, source, generation, models, now=None): + """Record one probe result and schedule the next TTL check.""" + now = self._clock() if now is None else now + with self._lock: + self._data["profiles"][profile] = { + "source": source, "source_generation": int(generation or 0), + "fetched_at": now, "retry_at": now + self.ttl, "failures": 0, "status": "ok", + "models": sorted({str(model)[:160] for model in models if model})[:MAX_MODEL_IDS]} + self._save() + + def fresh(self, profile, now=None): + """Record that the representative catalog is still inside its TTL.""" + now = self._clock() if now is None else now + with self._lock: + row = dict(self._data["profiles"].get(profile) or {}) + if not row: + return + row.update(status="fresh", fetched_at=now, retry_at=now + self.ttl, failures=0) + self._data["profiles"][profile] = row + self._save() + + def failure(self, profile, *, source=None, now=None): + """Back off an unusable probe and keep the previous aggregate for the next attempt.""" + now = self._clock() if now is None else now + with self._lock: + row = dict(self._data["profiles"].get(profile) or {}) + failures = min(int(row.get("failures") or 0) + 1, _MAX_FAILURES) + delay = min(RETRY_BASE_S * 2 ** (failures - 1), RETRY_MAX_S) + row.update(source=source if source is not None else row.get("source"), + failures=failures, status="error", + # Jitter keeps several failing profiles from retrying in lockstep. + retry_at=now + delay * random.uniform(1.0, 1.25)) + row.setdefault("fetched_at", 0.0) + row.setdefault("models", []) + self._data["profiles"][profile] = row + self._save() + + def model_ids(self, profile): + """Return the model identifiers the last successful probe advertised.""" + return set(self.row(profile).get("models") or ()) + + def view(self, now=None): + """Describe probe state for the management API without exposing credentials.""" + now = self._clock() if now is None else now + rows = {} + for profile, row in self.profiles().items(): + rows[profile] = { + "source": row.get("source"), "status": row.get("status", "unknown"), + "models": len(row.get("models") or []), "failures": row.get("failures", 0), + "fetched_at": row.get("fetched_at"), + "next_probe_in": max(0, round(float(row.get("retry_at") or 0.0) - now))} + return {"ttl_seconds": self.ttl, "profiles": rows} diff --git a/app/control_store.py b/app/control_store.py index c8e063b..8afd790 100644 --- a/app/control_store.py +++ b/app/control_store.py @@ -174,6 +174,11 @@ def snapshot(self): # Published dictionaries are never mutated; SQLite writer locks cannot stall routing. return copy.deepcopy(self._snapshot) + @property + def revision(self): + """Current revision without copying the snapshot, for cache and policy views.""" + return self._snapshot["revision"] + def _update(self, revision, change): with self._lock: self._db.execute("BEGIN IMMEDIATE") diff --git a/app/credits.py b/app/credits.py index a3536eb..cb4d35c 100644 --- a/app/credits.py +++ b/app/credits.py @@ -4,6 +4,7 @@ import base64 from copy import deepcopy import json +import itertools import math import os import re @@ -794,6 +795,22 @@ def fetch_request_usage(access_token: str, days: int = USAGE_MAX_DAYS, # Atomic JSON persistence for per-credential credit and check-in state # --------------------------------------------------------------------------- +_LEDGER_REVISION = itertools.count(1) + + +def _solvent_row(row: dict | None) -> bool: + """Mirror the pool's has-credit test so only eligibility flips bump the routing revision.""" + credits = (row or {}).get("credits") or {} + if not credits: + return True + if credits.get("unlimited") and credits.get("enterprise"): + return True + try: + return float(credits.get("credits") or 0) > 0 + except (TypeError, ValueError): + return False + + class CreditLedger: """Persist per-credential check-in and balance snapshots for expiry-aware scheduling.""" @@ -802,6 +819,10 @@ def __init__(self, path: Path | None = None, *, store=None): self.path = Path(store.path if store is not None else path) self._lock = threading.Lock() self._data: dict = {"version": 1, "creds": {}} + # Globally unique so a rebuilt ledger can never reuse the previous one's cache key. + self.revision = next(_LEDGER_REVISION) + # Bumped only when an account's solvency flips; value churn must not clear shards. + self.routing_revision = next(_LEDGER_REVISION) self._load() def _load(self): @@ -816,6 +837,7 @@ def _load(self): self._data = {"version": 1, "creds": {}} def _save(self): + self.revision = next(_LEDGER_REVISION) if self._store is not None: self._store.put("credits", self._data) return @@ -849,7 +871,8 @@ def bind_identity(self, cred_id: str, identity: str) -> bool: def remove(self, cred_id: str): """Clear persisted credit, check-in and error state after an identity change.""" with self._lock: - self._data["creds"].pop(cred_id, None) + if _solvent_row(self._data["creds"].pop(cred_id, None)) is not True: + self.routing_revision = next(_LEDGER_REVISION) self._save() # Check-in state @@ -878,6 +901,9 @@ def update_travel(self, cred_id: str, result: dict): def update_credits(self, cred_id: str, result: dict): with self._lock: + # Snapshot before mutating: _entry may hand back the very row read here. + was_solvent = _solvent_row(self._data["creds"].get(cred_id)) + had_balance_row = bool((self._data["creds"].get(cred_id) or {}).get("credits")) e = self._entry(cred_id) e["credits"] = { "credits": result.get("credits"), @@ -893,6 +919,10 @@ def update_credits(self, cred_id: str, result: dict): "unlimited": bool(result.get("unlimited")) and bool(result.get("enterprise")), } e["error"] = None + # A first balance row changes eligibility too: an absent row never matched + # _has_credit for intl accounts, so its appearance must move the routing revision. + if not had_balance_row or _solvent_row({"credits": e["credits"]}) != was_solvent: + self.routing_revision = next(_LEDGER_REVISION) self._save() def note_error(self, cred_id: str, message: str): diff --git a/app/model_policy.py b/app/model_policy.py index 89e6d7b..b62976d 100644 --- a/app/model_policy.py +++ b/app/model_policy.py @@ -3,7 +3,7 @@ from fastapi import HTTPException -from app import travel +from app import perf_metrics, travel from app.model_capabilities import route_metadata _request_policy = ContextVar("gateway_model_policy", default=None) @@ -23,31 +23,75 @@ async def __call__(self, scope, receive, send): def snapshot(config): + """Return a detached policy snapshot; one routing decision resolves it exactly once.""" store = config.get("control_store") - return store.snapshot() if store is not None else {"models": {}, "credentials": {}, "revision": 0} + if store is None: + return {"models": {}, "credentials": {}, "revision": 0} + perf_metrics.inc("policy_snapshot_calls") + return store.snapshot() -def credential_enabled(config, entry): - return snapshot(config)["credentials"].get(entry.get("account_key"), {}).get("enabled", True) +class PolicyView: + """One published policy revision reused across a whole routing decision.""" + __slots__ = ("data",) -def credential_auto_checkin(config, entry): + def __init__(self, data): + self.data = data + + @property + def revision(self): + return self.data.get("revision", 0) + + def credential(self, entry): + """Return the account metadata block for this entry, or an empty mapping.""" + return (self.data.get("credentials") or {}).get(entry.get("account_key")) or {} + + def rule_for(self, source): + return {**default_rule(source), **(self.data.get("models") or {}).get(source, {})} + + +def view(config): + """Resolve the policy once so every candidate check in one request sees the same revision.""" + return PolicyView(snapshot(config)) + + +def revision(config, policy=None): + """Return the published revision without copying the snapshot.""" + if policy is not None: + return policy.revision + store = config.get("control_store") + if store is None: + return 0 + value = getattr(store, "revision", None) + return value if isinstance(value, int) else snapshot(config).get("revision", 0) + + +def _policy(config, policy=None): + return policy.data if policy is not None else snapshot(config) + + +def credential_enabled(config, entry, policy=None): + return _policy(config, policy)["credentials"].get(entry.get("account_key"), {}).get("enabled", True) + + +def credential_auto_checkin(config, entry, policy=None): """International accounts opt in explicitly; the preference is independent of routing.""" default = entry.get("profile") in {"cn-cli", "cn-work"} - return snapshot(config)["credentials"].get(entry.get("account_key"), {}).get("auto_checkin", default) + return _policy(config, policy)["credentials"].get(entry.get("account_key"), {}).get("auto_checkin", default) -def credential_auto_travel(config, entry): +def credential_auto_travel(config, entry, policy=None): # The growth center is personal-only, so an enterprise tenant must never be scheduled # for travel even though it shares the domestic product profile. supported = travel.supported(entry.get("profile"), entry.get("enterprise") is True) - return supported and snapshot(config)["credentials"].get(entry.get("account_key"), {}).get("auto_travel", True) + return supported and _policy(config, policy)["credentials"].get(entry.get("account_key"), {}).get("auto_travel", True) -def credential_auto_daily_chat(config, entry): +def credential_auto_daily_chat(config, entry, policy=None): """A turn spends credits on the account, so international accounts opt in explicitly.""" supported = entry.get("profile") == "intl-work" - return supported and snapshot(config)["credentials"].get(entry.get("account_key"), {}).get("auto_daily_chat", False) + return supported and _policy(config, policy)["credentials"].get(entry.get("account_key"), {}).get("auto_daily_chat", False) def default_rule(source): @@ -56,8 +100,8 @@ def default_rule(source): "region": None, "profile": None, "credential_ids": []} -def rule_for(config, source): - return {**default_rule(source), **snapshot(config)["models"].get(source, {})} +def rule_for(config, source, policy=None): + return {**default_rule(source), **_policy(config, policy)["models"].get(source, {})} def _deny(message, code="model_disabled"): @@ -65,11 +109,18 @@ def _deny(message, code="model_disabled"): "message": message, "type": "invalid_request_error", "param": "model", "code": code}}) -def _active(config, upstream): +def _active(config, upstream, policy=None): state = _request_policy.get() if state and state.get("upstream") == upstream: - return rule_for(config, state["source"]), state - return rule_for(config, upstream), None + return rule_for(config, state["source"], policy), state + return rule_for(config, upstream, policy), None + + +def active_rule(config, model, policy=None): + """Return the rule governing this request's model, including a request-scoped alias.""" + if not model: + return default_rule("") + return _active(config, model, policy)[0] def resolve(config, public_model): @@ -105,11 +156,11 @@ def check_resolved(config, upstream): _deny("模型已停用") -def route_allowed(config, entry, model, *, rule=None): - if not credential_enabled(config, entry): +def route_allowed(config, entry, model, *, rule=None, policy=None): + if not credential_enabled(config, entry, policy): return False if rule is None: - active, state = _active(config, model) if model else (default_rule(""), None) + active, state = _active(config, model, policy) if model else (default_rule(""), None) if state and state["rule"] != active: return False rule = active @@ -167,7 +218,7 @@ def public_details(gateway, region=None): if accounts is not None: for entry in candidates: profile = entry.get("profile") - for model in gateway._effective_account_scope( + for model in gateway._published_account_scope( accounts.get(entry.get("account_key")) or {}, "serves", model_id=gateway._upstream_model(upstream, profile)) or []: if model.get("id") == gateway._upstream_model(upstream, profile): price = gateway._multiplier_value(model.get("credits")) diff --git a/app/perf_metrics.py b/app/perf_metrics.py new file mode 100644 index 0000000..9bf5f48 --- /dev/null +++ b/app/perf_metrics.py @@ -0,0 +1,66 @@ +"""Process-wide counters and timings for routing and maintenance hot paths.""" +from __future__ import annotations + +import threading +import time +from contextlib import contextmanager + +_COUNTERS: dict[str, int] = {} +_TIMINGS: dict[str, dict[str, float]] = {} +_LOCK = threading.Lock() +_STARTED = time.monotonic() + + +def inc(name: str, amount: int = 1) -> None: + """Count one event; the lock only guards the insert that snapshot iterates over.""" + with _LOCK: + _COUNTERS[name] = _COUNTERS.get(name, 0) + amount + + +def observe(name: str, seconds: float) -> None: + """Record one duration sample, creating the series on first use.""" + try: + sample = max(0.0, float(seconds)) + except (TypeError, ValueError): + return + stats = _TIMINGS.get(name) + if stats is None: + with _LOCK: + _TIMINGS.setdefault(name, {"count": 1.0, "total_s": sample, "max_s": sample, "min_s": sample}) + return + # Samples only ever add to these fields, so a lost update costs one sample, never an error. + stats["count"] += 1 + stats["total_s"] += sample + stats["max_s"] = max(stats["max_s"], sample) + stats["min_s"] = min(stats["min_s"], sample) + + +@contextmanager +def timed(name: str): + """Time the enclosed block into one series.""" + start = time.perf_counter() + try: + yield + finally: + observe(name, time.perf_counter() - start) + + +def reset() -> None: + """Clear every series; used by benchmarks and tests, never by a running service.""" + with _LOCK: + _COUNTERS.clear() + _TIMINGS.clear() + + +def snapshot() -> dict: + """Return a read-only copy with millisecond timings for the management view.""" + with _LOCK: + counters = dict(_COUNTERS) + timings = {name: {"count": int(stats["count"]), + "total_ms": round(stats["total_s"] * 1000, 3), + "avg_ms": round(stats["total_s"] * 1000 / stats["count"], 3), + "max_ms": round(stats["max_s"] * 1000, 3), + "min_ms": round(stats["min_s"] * 1000, 3)} + for name, stats in _TIMINGS.items()} + return {"uptime_seconds": round(time.monotonic() - _STARTED, 1), + "counters": counters, "timings": timings} diff --git a/app/routing_scheduler.py b/app/routing_scheduler.py new file mode 100644 index 0000000..63e8b8c --- /dev/null +++ b/app/routing_scheduler.py @@ -0,0 +1,73 @@ +"""Revision-keyed credential shards for request-side selection.""" +from __future__ import annotations + +import time + +from . import perf_metrics + +MODES = ("legacy", "shadow", "indexed") +MISMATCH_LIMIT = 20 + + +class RoutingScheduler: + """Serve the statically eligible entries per model; dynamic health stays per request. + + ``legacy`` evaluates every credential on every request, ``indexed`` reuses a shard built + for the current revision set, and ``shadow`` does both while returning the legacy answer. + """ + + __slots__ = ("_shards", "mismatches", "_limit") + + def __init__(self, limit=32): + self._shards: dict[tuple, tuple] = {} + self.mismatches: list[dict] = [] + self._limit = max(1, limit) # Aliases share one upstream model, so keep several shards. + + @staticmethod + def mode(config) -> str: + """Resolve the configured mode, falling back to the conservative shadow default.""" + value = config.get("routing_scheduler", "shadow") + return value if value in MODES else "shadow" + + def clear(self): + self._shards.clear() + + def candidates(self, *, mode, revisions, model, region, entries, evaluate): + """Return statically eligible entries for one model and region.""" + if mode == "legacy": + return [entry for entry in entries if evaluate(entry)] # No shard is built. + key = (model, region, *revisions) + shard = self._shards.get(key) + if shard is None: + shard = tuple(entry for entry in entries if evaluate(entry)) + if len(self._shards) >= self._limit: + self._shards.pop(next(iter(self._shards))) # A moved revision never matches again. + self._shards[key] = shard + perf_metrics.inc("shard_builds") + perf_metrics.inc("shard_size_rebuilt", len(shard)) + else: + self._shards[key] = self._shards.pop(key) # LRU: hot models must not evict themselves. + perf_metrics.inc("shard_hits") + if mode == "shadow": + legacy = [entry for entry in entries if evaluate(entry)] + if legacy != list(shard): + self.record(model, region, legacy, shard) + return legacy + return list(shard) + + def record(self, model, region, legacy, indexed): + """Retain a bounded, credential-free description of one shadow mismatch.""" + perf_metrics.inc("scheduler_mismatches") + self.mismatches.append({ + "at": time.time(), "model": model, "region": region, + "legacy": [entry["id"] for entry in legacy], + "indexed": [entry["id"] for entry in indexed]}) + del self.mismatches[:-MISMATCH_LIMIT] + + def report(self) -> dict: + """Describe the active shard set without exposing tokens.""" + # Request threads insert concurrently, so iterate over a stable copy. + return {"shards": [ + {"model": key[0], "region": key[1], "size": len(value)} + for key, value in list(self._shards.items())], + "mismatches": list(self.mismatches)} diff --git a/app/settings.py b/app/settings.py index d3d19c6..2b4bb33 100644 --- a/app/settings.py +++ b/app/settings.py @@ -97,6 +97,14 @@ def _item(default, type_, label, *, mode="hot", env=None, minimum=None, maximum= "credit_price_usd": _item(0.03, "number", "国际积分单价", minimum=0), "model_catalog_ttl": _item(21600, "integer", "模型目录缓存秒数", minimum=0, maximum=31536000), "model_guard": _item(True, "boolean", "表外模型拦截"), + "routing_scheduler": _item("shadow", "string", "路由调度模式", + env="CODEBUDDY2API_ROUTING_SCHEDULER", + choices=["legacy", "shadow", "indexed"]), + "catalog_probe_mode": _item("account", "string", "目录探测模式", + env="CODEBUDDY2API_CATALOG_PROBE_MODE", + choices=["account", "representative"]), + "housekeep_concurrency": _item(4, "integer", "后台维护并发数", env="CODEBUDDY2API_HOUSEKEEP_CONCURRENCY", + minimum=1, maximum=16), "model_capability_guard": _item(True, "boolean", "模型能力预检", env="CODEBUDDY2API_MODEL_CAPABILITY_GUARD"), "max_images": _item(16, "integer", "单请求图片上限", env="CODEBUDDY2API_MAX_IMAGES", minimum=0, maximum=10000), "image_policy": _item("truncate", "string", "超额图片策略", env="CODEBUDDY2API_IMAGE_POLICY", choices=["truncate", "error"]), diff --git a/app/state_store.py b/app/state_store.py index ff35031..c7f4366 100644 --- a/app/state_store.py +++ b/app/state_store.py @@ -21,6 +21,7 @@ "credits": "credits-ledger.json", "trial": "trial-ledger.json", "catalog": "model-catalog.json", + "catalog_probe": "model-catalog-probe.json", "usage": "usage-snapshots.json", } @@ -75,6 +76,7 @@ def _check_document(name, value): "sessions": ("sessions", (1,)), "cooldowns": ("accounts", (1,)), "model_blocks": ("blocks", (1,)), "credits": ("creds", (1,)), "trial": ("accounts", (1,)), "catalog": ("groups", (1, 2)), + "catalog_probe": ("profiles", (1,)), "usage": ("accounts", (1,)), }[name] if type(value.get("version")) is not int or value["version"] not in versions or not isinstance(value.get(field), dict): @@ -146,6 +148,10 @@ def _check_document(name, value): if (not isinstance(row, dict) or type(row.get("until")) not in (int, float) or not math.isfinite(row["until"]) or type(row.get("hits", 1)) is not int): raise ValueError("Invalid model backoff") + elif name == "catalog_probe": + from .catalog_probe import valid_row + if len(value[field]) > 8 or not all(valid_row(row) for row in value[field].values()): + raise ValueError("Invalid catalog probe state") elif name == "catalog": for row in value[field].values(): if (not isinstance(row, dict) or not isinstance(row.get("models"), list) @@ -204,7 +210,7 @@ def put(self, name, value): def put_cache(self, name, value): """Cache failures stay visible without replacing the upstream inference outcome.""" - if name not in ("catalog", "model_blocks"): + if name not in ("catalog", "catalog_probe", "model_blocks"): raise ValueError("Not a rebuildable cache") try: self.put(name, value) diff --git a/benchmarks/bench_routing.py b/benchmarks/bench_routing.py new file mode 100644 index 0000000..41f193a --- /dev/null +++ b/benchmarks/bench_routing.py @@ -0,0 +1,180 @@ +#!/usr/bin/env python3 +"""Offline synthetic benchmark for multi-account routing and model aggregation. + +Builds a real ``CredentialPool`` from synthetic ``.info`` files, a real +``CreditLedger`` and account-scoped catalogs, then replaces nothing but the +upstream network. It measures the request-side hot paths so a change can be +compared against the documented baseline on the same machine. +""" +from __future__ import annotations + +import argparse +import json +import os +import statistics +import sys +import tempfile +import time +from pathlib import Path + +ROOT = Path(__file__).resolve().parents[1] +sys.path.insert(0, str(ROOT)) + +PROFILES = ("cn-cli", "cn-work", "intl-cli", "intl-work") +DOMAINS = {"cn-cli": "www.codebuddy.cn", "cn-work": "www.workbuddy.cn", + "intl-cli": "www.codebuddy.ai", "intl-work": "www.workbuddy.ai"} +MODELS = { + "cn-cli": ["hy4-preview", "glm-5.3", "glm-5.3-flash", "deepseek-v4-flash"], + "cn-work": ["auto", "glm-5.3", "glm-5.3-flash"], + "intl-cli": ["default-model", "claude-sonnet-4-5", "gpt-5.1"], + "intl-work": ["default-model", "claude-sonnet-4-5", "gpt-5.1"], +} +FREE_MODEL = "glm-5.3-flash" + + +def credential_payload(uid: str, profile: str) -> dict: + """Return a credential shaped like a real export: account block plus plaintext token.""" + return {"account": {"uid": uid, "nickname": f"bench-{uid}"}, + "auth": {"accessToken": f"bench-token-{uid}", "refreshToken": f"bench-refresh-{uid}", + "domain": DOMAINS[profile], "expiresAt": int(time.time() * 1000) + 86400 * 1000, + "lastRefreshTime": int(time.time() * 1000)}} + + +def catalog(profile: str, free: bool) -> dict: + models = [{"id": name, "supportsToolCall": True, + "credits": "x0 credits" if free and name == FREE_MODEL else "x1 credits"} + for name in MODELS[profile]] + return {"profile": profile, "models": [dict(item) for item in models], + "serves": [dict(item) for item in models]} + + +def build(converter, credits_mod, count: int, root: Path): + """Populate the process-wide CONFIG with ``count`` synthetic accounts.""" + config = converter.CONFIG + paths = [] + for index in range(count): + profile = PROFILES[index % len(PROFILES)] + path = root / f"account-{index:05d}.info" + path.write_text(json.dumps(credential_payload(f"bench-{index:05d}", profile)), encoding="utf-8") + paths.append(path) + pool = converter.CredentialPool(paths, scan=False) + ledger = credits_mod.CreditLedger(root / "credits.json") + pool.set_ledger(ledger) + accounts = {} + for index, entry in enumerate(pool.entries()): + profile = entry.get("profile") or PROFILES[index % len(PROFILES)] + ledger.bind_identity(entry["id"], entry["account_key"]) + ledger.update_credits(entry["id"], {"credits": 1000 + index, "intl": profile.startswith("intl"), + "segments": [{"remaining": 1000 + index, + "expires_at": time.time() + 3600 + index}]}) + accounts[entry["account_key"]] = catalog(profile, free=index % 3 == 0) + catalogs = {profile: None for profile in PROFILES} + for account in accounts.values(): + merged = catalogs[account["profile"]] = catalogs[account["profile"]] or [] + merged.extend(account["models"]) + config.update({"cred_pool": pool, "ledger": ledger, "account_catalogs": accounts, + "routing_scheduler": "indexed", # The benchmark measures the indexed path. + "catalog_index": converter._publish_catalog_index(accounts), + "model_catalogs": catalogs, "models_remote": catalogs["cn-cli"], + "models_intl": catalogs["intl-cli"], "model_cache": None, + "model_guard": True, "max_inflight_per_account": 0, + "control_store": None, "log_path": None}) + return pool + + +def measure(converter, perf_metrics, case, operation, repeat): + """Run one case and return wall-clock milliseconds plus per-op counters.""" + times = [] + counters = {} + for _ in range(repeat): + perf_metrics.reset() + started = time.perf_counter() + operation() + times.append((time.perf_counter() - started) * 1000) + for name, value in perf_metrics.snapshot()["counters"].items(): + counters[name] = counters.get(name, 0) + value + return {"case": case, "median_ms": round(statistics.median(times), 3), + "max_ms": round(max(times), 3), + "counters": {name: round(value / repeat, 2) for name, value in counters.items()}} + + +def run(converter, credits_mod, perf_metrics, count, repeat, model): + temporary = tempfile.TemporaryDirectory() + root = Path(temporary.name) + pool = build(converter, credits_mod, count, root) + rows = [] + try: + rows.append(measure(converter, perf_metrics, "pick", lambda: pool.pick("session", model), repeat)) + rows.append(measure(converter, perf_metrics, "headers_for", + lambda: pool.headers_for(None, model), repeat)) + ledger = converter.CONFIG["ledger"] + churn_cid = pool.entries()[0]["id"] + churn_paths = [entry["id"] for entry in pool.entries()[:max(1, count // 10)]] + churn_balance = {"credits": 1000, "intl": False, + "segments": [{"remaining": 1000, "expires_at": time.time() + 3600}]} + + def churned(): + # Token-refresh-like reconcile (reset=False) plus a maintenance-style ledger + # write: neither moves a semantic revision, so the shard stays hot. + pool.reload([Path(path) for path in churn_paths], reset=False) + ledger.update_credits(churn_cid, churn_balance) + pool.pick("session", model) + cm = pool.pick("cooldown", model) + pool.note_status(cm, 429, model=model) + try: + rows.append(measure(converter, perf_metrics, "pick + cooldown", lambda: pool.pick("session", model), repeat)) + finally: + converter.CONFIG["model_cache"] = None + tried = frozenset(entry["cm"] for entry in pool.entries()[:max(1, count // 10)]) + rows.append(measure(converter, perf_metrics, "pick + tried", + lambda: pool.pick("session", model, tried=tried), repeat)) + limit = converter.CONFIG["max_inflight_per_account"] + converter.CONFIG["max_inflight_per_account"] = 1 + + def saturated(): + try: + pool.pick("session", model, with_capacity=True) + except Exception: + pass # Only the ranking work is measured; the refusal itself is free. + rows.append(measure(converter, perf_metrics, "pick + capacity full", saturated, repeat)) + converter.CONFIG["max_inflight_per_account"] = limit + rows.append(measure(converter, perf_metrics, "current_models", converter.current_models, repeat)) + rows.append(measure(converter, perf_metrics, "current_model_details", converter.current_model_details, repeat)) + rows.append(measure(converter, perf_metrics, "pick + revision churn", churned, repeat)) + finally: + temporary.cleanup() + converter.CONFIG.update({"cred_pool": None, "ledger": None, "account_catalogs": None, + "catalog_index": None, "model_catalogs": {}, + "models_remote": None, "models_intl": None}) + return rows + + +def main(argv=None): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--accounts", default="10,100,500", help="comma-separated account counts") + parser.add_argument("--repeat", type=int, default=3, help="repetitions per case") + parser.add_argument("--model", default="glm-5.3-flash", help="routed model to benchmark") + parser.add_argument("--json", action="store_true", help="print JSON instead of a table") + args = parser.parse_args(argv) + os.environ.setdefault("CODEBUDDY_AUTH_DIR", tempfile.mkdtemp(prefix="bench-auth-")) + import converter + from app import perf_metrics, credits as credits_mod + report = {} + for count in [int(value) for value in args.accounts.split(",") if value.strip()]: + report[count] = run(converter, credits_mod, perf_metrics, count, max(1, args.repeat), args.model) + if args.json: + print(json.dumps(report, ensure_ascii=False, indent=2)) + return 0 + print("| 账号数 | 场景 | 中位耗时 ms | 最慢 ms | summary 调用 | stat 调用 | 候选检查 |") + print("|---:|---|---:|---:|---:|---:|---:|") + for count, rows in report.items(): + for row in rows: + counters = row["counters"] + print(f"| {count} | {row['case']} | {row['median_ms']} | {row['max_ms']} | " + f"{counters.get('credential_summary_calls', 0)} | {counters.get('credential_stat_calls', 0)} | " + f"{counters.get('candidate_checks', 0)} |") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/converter.py b/converter.py index 35b3cc1..57e033e 100644 --- a/converter.py +++ b/converter.py @@ -6,9 +6,11 @@ import argparse import asyncio import hashlib +import itertools import json import math import os +import random import re import secrets import shutil @@ -17,6 +19,7 @@ import time import uuid from collections import OrderedDict +from concurrent.futures import ThreadPoolExecutor from contextlib import asynccontextmanager, nullcontext from dataclasses import dataclass from datetime import datetime, timedelta, timezone @@ -72,11 +75,16 @@ def desensitize_body(body, roles=("system",), desensitize_harness_user=False, from app.message_normalization import merge_intl_user_images from app.adapters.chat_input import normalize_chat_messages from app.model_catalog_view import INTERNATIONAL as SHARED_INTL_PROFILES, share_models +from app.catalog_index import CatalogIndex +from app.routing_scheduler import RoutingScheduler +from app.catalog_probe import AggregateCatalog +from app.catalog_index import account_scope as _account_scope from app.inference_auth import require_api_key from app.admin_auth import SessionStoreError from app.content_filter import ContentFilterDetector, is_filter_error from app.request_limits import ImageLimitError, apply_image_policy from app.safe_logging import format_log_body, sanitize_log_text +from app import perf_metrics from app.startup import allows_open_noauth, load_startup_env, run_server from app.site_routing import (DOMESTIC, INTERNATIONAL, PROFILE_ENDPOINTS, site_for_auth, site_for_headers, profile_for_auth, profile_for_headers, profile_region, profile_product, @@ -287,9 +295,17 @@ def _read_raw(self) -> dict: return json.load(f) def _file_version(self): + perf_metrics.inc("credential_stat_calls") st = self.path.stat() return st.st_dev, st.st_ino, st.st_mtime_ns, st.st_size + def _file_changed(self) -> bool: + """Whether the on-disk file moved since the last load; one stat, no re-read.""" + try: + return self._cached is None or self._file_version() != self._mtime + except OSError: + return True + def _load_if_stale(self): """Reload changed credentials and invalidate leases held by older requests.""" mt = self._file_version() @@ -309,8 +325,9 @@ def _session(self) -> dict: auth_oauth.ensure_plaintext_tokens(auth) return s - def _is_expired(self) -> bool: - s = self._session() + def _is_expired(self, s: dict | None = None) -> bool: + if s is None: + s = self._session() expires_at = (s.get("auth") or {}).get("expiresAt") or 0 # Treat tokens as expired 60 seconds early. return time.time() * 1000 >= (expires_at - 60_000) @@ -372,9 +389,11 @@ def _build_headers_from(self, auth: dict, account: dict) -> dict: def get_headers(self) -> dict: """Return upstream headers with a current token, refreshing when necessary.""" with self._lock: - if self._is_expired(): - self._refresh() + # 复用同一次会话读取判断过期，稳态只做一次文件版本检查；刷新后才重读。 s = self._session() + if self._is_expired(s): + self._refresh() + s = self._session() return self._build_headers_from(s.get("auth") or {}, _credential_account(s)) def refresh_if_due(self, margin_s: int, keepalive_s: int) -> bool: @@ -390,6 +409,7 @@ def invalidate(self): def summary(self) -> dict: + perf_metrics.inc("credential_summary_calls") with self._lock: s = self._session() auth = s.get("auth") or {} @@ -408,7 +428,7 @@ def summary(self) -> dict: "enterpriseName": acct.get("enterpriseName"), "enterprise": is_enterprise_account(acct), "token_expires_at": auth.get("expiresAt", 0), - "token_expired": self._is_expired(), + "token_expired": self._is_expired(s), "last_refresh_time": auth.get("lastRefreshTime") or 0, } @@ -543,6 +563,10 @@ def _dynamic_request_headers(skey: str | None) -> dict: } +# Globally unique revisions keep cached policy views and model tables from matching a rebuilt pool. +_POOL_REVISION = itertools.count(1) + + class CredentialPool: """Manage credential discovery, reloads, sticky sessions, cooldowns and refresh.""" @@ -560,7 +584,13 @@ def __init__(self, paths: list[Path] | None = None, scan: bool = False, self._rr = {None: 0, "cn": 0, "intl": 0} self._ledger = None # Prefer credits expiring sooner. self._capacity = AccountCapacity() - self._scan = scan # Rescan credentials before selection. + self._scan = scan # Scan the auth directory during reconcile, never per request. + # Request-side lookups are rebuilt whenever the entry list changes. + self._index_id: dict[str, dict] = {} + self._index_cm: dict = {} + self._index_identity: dict[str, list[dict]] = {} + self.revision = next(_POOL_REVISION) # Bumped whenever the entry list is rebuilt. + self._scheduler = RoutingScheduler() # Revision-keyed eligibility shards. # Explicit paths must survive rejection or eviction so a repaired file re-enters. self._configured: list[str] = [] self._ignored_duplicates: set[str] = set() @@ -570,6 +600,7 @@ def __init__(self, paths: list[Path] | None = None, scan: bool = False, self._sync_event = threading.Event() self._sync_retry: dict[str, float] = {} self._sync_attempts: dict[str, int] = {} + self._index_signature: tuple | None = None # Semantic entry shape; unchanged reloads keep shards. self.reload(paths or []) if self._scan: self._rescan() # Discover credentials at startup. @@ -588,6 +619,11 @@ def reload(self, paths: list[Path], *, reset: bool = True): continue entry = by_id.get(cid) if entry is not None: + if (not reset and entry["generation"] == entry["cm"]._generation + and not entry["cm"]._file_changed()): + # Reconcile re-lists every file; only a moved or re-read file pays + # the summary, and a generation gap means the identity may have moved. + continue if reset: entry["cm"].invalidate() try: @@ -606,6 +642,7 @@ def reload(self, paths: list[Path], *, reset: bool = True): changed = reset or generation != entry.get("generation") if changed: old_identity = entry.get("account_key") + old_profile = entry.get("profile") replaced = old_identity != identity if replaced: self._model_fail = {key: until for key, until in self._model_fail.items() if key[0] != cid} @@ -614,9 +651,14 @@ def reload(self, paths: list[Path], *, reset: bool = True): self.forget_credential_state(entry) if entry.get("uid"): have_uids.pop(old_identity, None) + # A same-account background token refresh does not change catalog + # entitlements, but it must also not clear a pending refresh marker + # set by an import or a representative probe; only _store_catalog does. entry.update(uid=summary.get("uid"), profile=summary["profile"], site=summary["site"], enterprise=summary.get("enterprise") is True, - account_key=identity, generation=generation, catalog_dirty=True) + account_key=identity, generation=generation, + catalog_dirty=bool(entry.get("catalog_dirty") or reset or replaced + or old_profile != summary["profile"])) self._bind_entry(entry) if reset or replaced: entry.update(fail_until=0.0, keepalive_after=0.0) @@ -668,9 +710,40 @@ def reload(self, paths: list[Path], *, reset: bool = True): if uid: have_uids[identity_key] = cid self._queue_sync(cid) + self._rebuild_index() + + def _rebuild_index(self): + """Rebuild request-side lookups after the entry list changed; never called per request.""" + index_id = {entry["id"]: entry for entry in self._entries} + index_cm = {entry["cm"]: entry for entry in self._entries} + identity: dict[str, list[dict]] = {} + for entry in self._entries: + if entry.get("account_key"): + identity.setdefault(entry["account_key"], []).append(entry) + # Token refreshes rewrite files without changing who is routable, so only a real + # identity/profile/order move invalidates the revision and the scheduler shards. + signature = tuple((entry["id"], entry.get("account_key"), entry.get("profile")) + for entry in self._entries) + if signature != self._index_signature: + self._index_signature = signature + self.revision = next(_POOL_REVISION) + self._scheduler.clear() + self._index_id, self._index_cm, self._index_identity = index_id, index_cm, identity + + def entry_for(self, cm): + """Return the live entry owning this manager, or None after eviction.""" + return self._index_cm.get(cm) + + def entry_by_id(self, cid): + """Return the live entry for one credential path.""" + return self._index_id.get(cid) + + def entries_by_identity(self, account_key): + """Return every live entry sharing one account identity.""" + return tuple(self._index_identity.get(account_key, ())) def _queue_sync(self, cid): - entry = next((entry for entry in self._entries if entry["id"] == cid), None) + entry = self._index_id.get(cid) if entry is not None and not model_policy.credential_enabled(CONFIG, entry): return self._sync_pending.add(cid) @@ -690,7 +763,13 @@ def begin_sync(self, *, all_entries=False): self._sync_retry = {cid: deadline for cid, deadline in self._sync_retry.items() if cid in active} due = {cid for cid, deadline in self._sync_retry.items() if deadline <= time.monotonic()} self._sync_pending.update(due) - ids = active if all_entries else set(self._sync_pending) + if all_entries: + # A cycle right after an event must not repeat accounts that just succeeded. + cutoff = time.time() - HOUSEKEEP_MIN_SYNC_GAP_S + ids = {entry["id"] for entry in self._entries + if entry["id"] in active and float(entry.get("synced_at") or 0.0) <= cutoff} + else: + ids = set(self._sync_pending) self._sync_pending.difference_update(ids) if not self._sync_pending: self._sync_event.clear() @@ -702,13 +781,18 @@ def end_sync(self, ids, failed=()): self._syncing.difference_update(ids) present = {entry["id"] for entry in self._entries if model_policy.credential_enabled(CONFIG, entry)} for cid in ids: + entry = self._index_id.get(cid) if cid in failed and cid in present and cid not in self._sync_pending: attempt = min(self._sync_attempts.get(cid, 0) + 1, 5) self._sync_attempts[cid] = attempt - self._sync_retry[cid] = time.monotonic() + min(60 * 2 ** (attempt - 1), 900) + delay = min(60 * 2 ** (attempt - 1), 900) + # Jitter keeps repeated upstream refusals from retrying in lockstep. + self._sync_retry[cid] = time.monotonic() + delay * random.uniform(1.0, 1.25) else: self._sync_retry.pop(cid, None) self._sync_attempts.pop(cid, None) + if entry is not None: + entry["synced_at"] = time.time() def sync_pending(self, region=None): with self._lock: @@ -725,7 +809,7 @@ def sync_wait(self, periodic_delay): def apply_if_current(self, cm, generation, update): """Run updates only for enabled accounts with the current credential lease.""" with self._lock, cm._lock: - entry = next((entry for entry in self._entries if entry["cm"] is cm), None) + entry = self._index_cm.get(cm) if entry is None or not model_policy.credential_enabled(CONFIG, entry): return False if not self._lease_matches(cm, generation): @@ -749,6 +833,7 @@ def _drop_entry(self, entry): self._syncing.discard(cid) self._sync_retry.pop(cid, None) self._sync_attempts.pop(cid, None) + self._rebuild_index() invalidate_model_table() def prune(self): """Remove missing credential files and their session bindings.""" @@ -774,6 +859,8 @@ def prune(self): self._model_fail = {k: v for k, v in self._model_fail.items() if k[0] in ids} if CONFIG.get("cred_pool") is self: _publish_model_cache() + if len(self._entries) != before: + self._rebuild_index() def find_by_uid(self, uid: str, identity: str | None = None) -> Optional[str]: @@ -953,11 +1040,35 @@ def _rescan(self): paths = find_auth_files() if self._scan else [Path(cid) for cid in self._configured] self.reload(paths, reset=False) - def _healthy(self, e: dict) -> bool: - return model_policy.credential_enabled(CONFIG, e) and time.time() >= e["fail_until"] + def fingerprint(self): + """Return a cheap directory fingerprint so reconcile only runs after a real change.""" + configured = {str(path) for path in find_auth_files()} if self._scan else set(self._configured) + rows = [] + for cid in sorted(configured): + try: + st = os.stat(cid) + except OSError: + rows.append((cid, None)) + continue + rows.append((cid, st.st_dev, st.st_ino, st.st_mtime_ns, st.st_size)) + return (self._scan, tuple(rows)) + + def reconcile(self, previous=None): + """Apply file changes outside the request path; returns the new fingerprint.""" + current = self.fingerprint() + if previous is None or current != previous: + self._rescan() + return current + + def _healthy(self, e: dict, policy=None) -> bool: + return model_policy.credential_enabled(CONFIG, e, policy) and time.time() >= e["fail_until"] @staticmethod def _entry_profile(entry): + # Reload caches the parsed profile, so the request path never re-reads the file here. + profile = entry.get("profile") + if profile: + return profile try: return entry["cm"].summary().get("profile", "cn-cli") except Exception: @@ -995,25 +1106,18 @@ def _has_credit(self, entry, profile): except (TypeError, ValueError): return False - def _eligible(self, entry, model, *, region=None, profile=None, rule=None): - if not model_policy.route_allowed(CONFIG, entry, model, rule=rule): + def _eligible(self, entry, model, *, region=None, profile=None, rule=None, policy=None): + if not model_policy.route_allowed(CONFIG, entry, model, rule=rule, policy=policy): return False actual = self._entry_profile(entry) profile = profile or actual if not profile or profile != actual or not _in_region(profile, region): return False - configured = {candidate for item in self._entries if (candidate := self._entry_profile(item)) - and _in_region(candidate, region)} - if profile not in _model_profiles(model, region, configured): + configured = _configured_profiles(region, policy) + if profile not in _model_profiles(model, region, configured, policy=policy): return False if CONFIG.get("account_catalogs") is not None or CONFIG.get("model_cache") is not None: - try: - identity = entry["cm"].summary()["account_key"] - except Exception: - return False - if identity != entry.get("account_key"): - return False - account = (CONFIG.get("account_catalogs") or {}).get(identity) or {} + account = (CONFIG.get("account_catalogs") or {}).get(entry.get("account_key")) or {} models = _effective_account_scope(account, "serves", model_id=_upstream_model(model, profile)) if account.get("profile") != profile or models is None: return False @@ -1022,7 +1126,7 @@ def _eligible(self, entry, model, *, region=None, profile=None, rule=None): cli_auto = model == "auto" and profile == "cn-cli" and bool(usable) # Disabling the guard must not borrow another account's model capabilities. declared = any(item["id"] == _upstream_model(model, profile) - for item in _models_for_profile(profile, configured, scope="serves")) + for item in _models_for_profile(profile, configured, scope="serves", policy=policy)) passthrough = (model != "auto" and not declared and not CONFIG.get("model_guard") and len(configured) == 1) if model and not (supported or cli_auto or passthrough): @@ -1086,12 +1190,26 @@ def _evict_sticky(self): else: break - def _candidates(self, model: str | None, *, region=None, tried=()) -> list[dict]: + def _shard_revisions(self, policy=None) -> tuple: + """Revisions that change static eligibility; any move invalidates every shard.""" + return (self.revision, model_policy.revision(CONFIG, policy), _catalog_revision, + getattr(self._ledger, "routing_revision", 0)) + + def _candidates(self, model: str | None, *, region=None, tried=(), policy=None) -> list[dict]: """Exclude tried credentials and rank candidates by zero rate and credit expiry.""" tried = set(tried) - healthy = [entry for entry in self._entries if entry["cm"] not in tried - and self._healthy(entry) - and self._eligible(entry, model, region=region) and self._model_healthy(entry, model) + # A request-scoped alias restricts the accounts, so its rule is part of the shard key. + rule = model_policy.active_rule(CONFIG, model, policy) + scope = (rule.get("public_id"), rule.get("region") or "", rule.get("profile") or "", + tuple(rule.get("credential_ids") or ())) + eligible = self._scheduler.candidates( + mode=RoutingScheduler.mode(CONFIG), revisions=(*self._shard_revisions(policy), scope), + model=model, region=region, entries=self._entries, + evaluate=lambda entry: self._eligible(entry, model, region=region, policy=policy)) + perf_metrics.inc("candidate_checks", len(eligible)) + healthy = [entry for entry in eligible if entry["cm"] not in tried + and self._healthy(entry, policy) + and self._model_healthy(entry, model) and self._model_servable(entry, model)] if not healthy: return [] @@ -1111,12 +1229,17 @@ def _capacity_key(entry): def pick(self, skey: str | None, model: str | None = None, *, region=None, - tried=(), with_capacity=False, requirements=None) -> CredentialManager | None: + tried=(), with_capacity=False, requirements=None, policy=None) -> CredentialManager | None: """Select a healthy sticky or round-robin credential, preferring eligible zero-rate accounts.""" - self._rescan() # Reload and prune acquire their own locks. - with self._lock: + # The candidate list comes from the indexes; file changes are applied by background reconcile. + policy = policy or model_policy.view(CONFIG) # One policy revision for the whole decision. + started = time.perf_counter() + waiting = time.perf_counter() + self._lock.acquire() + perf_metrics.observe("pool_lock_wait", time.perf_counter() - waiting) + try: self._evict_sticky() - candidates = self._candidates(model, region=region, tried=tried) + candidates = self._candidates(model, region=region, tried=tried, policy=policy) if not candidates: if skey: self._sticky.pop(skey, None) @@ -1135,6 +1258,7 @@ def pick(self, skey: str | None, model: str | None = None, *, region=None, candidates = [entry for entry in candidates if self._model_free(entry, model) == free and self._capacity.count(self._capacity_key(entry)) < limit] if not candidates: + perf_metrics.inc("capacity_acquire_failures") raise self._capacity_error() best = candidates[0] free = self._model_free(best, model) @@ -1152,14 +1276,19 @@ def pick(self, skey: str | None, model: str | None = None, *, region=None, if skey: self._sticky[skey] = (e["id"], time.time()) return e["cm"] + finally: + self._lock.release() + perf_metrics.observe("route_pick_duration", time.perf_counter() - started) def headers_for(self, skey: str | None, model: str | None = None, *, region=None, - with_generation=False, tried=(), with_capacity=False, requirements=None): + with_generation=False, tried=(), with_capacity=False, requirements=None, policy=None): """Recheck identity and atomically reserve account capacity before sending.""" + policy = policy or model_policy.view(CONFIG) capacity_race = False capability_failures = [] for _ in range(max(1, len(self._entries))): - cm = self.pick(skey, model, region=region, tried=tried, with_capacity=with_capacity, requirements=requirements) + cm = self.pick(skey, model, region=region, tried=tried, with_capacity=with_capacity, + requirements=requirements, policy=policy) if cm is None: return None reason = None @@ -1174,10 +1303,13 @@ def headers_for(self, skey: str | None, model: str | None = None, *, region=None self.cooldown(cm, reason=reason, generation=generation) continue with self._lock: - self.reload([cm.path], reset=False) - entry = next((entry for entry in self._entries if entry["cm"] is cm), None) - if (entry is not None and cm._generation == generation and self._healthy(entry) - and self._eligible(entry, model, region=region, profile=profile) and self._model_healthy(entry, model)): + entry = self._index_cm.get(cm) + if entry is not None and cm._generation != entry.get("generation"): + self.reload([cm.path], reset=False) # Ownership changed on disk: re-read one file. + entry = self._index_cm.get(cm) + if (entry is not None and cm._generation == generation and self._healthy(entry, policy) + and self._eligible(entry, model, region=region, profile=profile, policy=policy) + and self._model_healthy(entry, model)): if requirements is not None: failures = requirements.violations(model_capabilities.entry_model(sys.modules[__name__], entry, model)) if failures: @@ -1187,6 +1319,7 @@ def headers_for(self, skey: str | None, model: str | None = None, *, region=None lease = self._capacity.acquire(self._capacity_key(entry), CONFIG.get("max_inflight_per_account", 0), cm, generation) if lease is None: + perf_metrics.inc("capacity_acquire_failures") capacity_race = True continue return lease, headers @@ -1441,8 +1574,59 @@ def _refresher_loop(pool: CredentialPool): except Exception as e: _log(f"[cred] 刷新线程异常: {e}") +CRED_RECONCILE_INTERVAL_S = 5 # Credential directory polling interval for background reconcile + + +def _reconcile_loop(pool: CredentialPool): + """Apply credential file changes in the background so requests never scan the pool.""" + try: + previous = pool.fingerprint() + except Exception: + previous = None + while True: + time.sleep(CRED_RECONCILE_INTERVAL_S) + try: + previous = pool.reconcile(previous) + except Exception as e: + _log(f"[cred] 目录变更处理异常: {_network_error_text(e)}") + CHECKIN_FIRST_DELAY = 30 # Initial check-in delay in seconds HOUSEKEEP_INTERVAL = 3600 # Account maintenance interval in seconds +HOUSEKEEP_MIN_SYNC_GAP_S = 300 # Skip accounts synced this recently, so event cycles never repeat them +HOUSEKEEP_CONCURRENCY = 4 # Default bounded workers per upstream endpoint family +_ENDPOINT_KINDS = ("credits", "models", "usage") +_endpoint_limits: dict = {} +_ENDPOINT_LIMIT_LOCK = threading.Lock() # Workers initialize semaphores concurrently. +_published_catalog_signature: str | None = None # Last published catalog content; skips no-op bumps. + + +def _housekeep_concurrency() -> int: + """Resolve the bounded worker count; 1 restores the previous serial behaviour.""" + value = CONFIG.get("housekeep_concurrency", HOUSEKEEP_CONCURRENCY) + return value if isinstance(value, int) and not isinstance(value, bool) and value >= 1 else HOUSEKEEP_CONCURRENCY + + +def _endpoint_semaphore(kind: str): + """One independent limit per upstream endpoint family, so no phase starves another.""" + limit = _housekeep_concurrency() + with _ENDPOINT_LIMIT_LOCK: + if _endpoint_limits.get("limit") != limit: + _endpoint_limits.clear() + _endpoint_limits["limit"] = limit + for name in _ENDPOINT_KINDS: + _endpoint_limits[name] = threading.BoundedSemaphore(limit) + return _endpoint_limits[kind] + + +def _bounded_map(function, items, limit): + """Run function over items on a bounded worker pool, yielding results in input order.""" + items = list(items) + if limit <= 1 or len(items) <= 1: + for item in items: + yield function(item) + return + with ThreadPoolExecutor(max_workers=min(limit, len(items)), thread_name_prefix="housekeep") as workers: + yield from workers.map(function, items) def _bearer_token(headers: dict) -> str: @@ -1617,6 +1801,7 @@ def can_chat(): def _publish_model_cache(): """Publish account-scoped versioned catalogs, excluding ownerless shared caches.""" + global _published_catalog_signature cache = CONFIG.get("model_cache") if cache is not None: pool = CONFIG.get("cred_pool") @@ -1631,7 +1816,6 @@ def _publish_model_cache(): accounts[identity] = {"profile": profile, "models": cache.models(key) if known else None, "serves": cache.serves(key) if known else None} - CONFIG["account_catalogs"] = accounts catalogs = {profile: None for profile in PROFILE_ENDPOINTS} for account in accounts.values(): if account["models"] is not None: @@ -1639,43 +1823,327 @@ def _publish_model_cache(): if models is None: models = catalogs[account["profile"]] = [] models.extend(account["models"]) - CONFIG["model_catalogs"] = catalogs - CONFIG["models_remote"], CONFIG["models_intl"] = catalogs["cn-cli"], catalogs["intl-cli"] - invalidate_model_table() + # Token refreshes republish identical content; only a real change may bump the + # catalog revision, or steady-state shards would rebuild on every maintenance write. + signature = json.dumps(accounts, sort_keys=True, default=str) + if (signature != _published_catalog_signature + or CONFIG.get("account_catalogs") is None): + # The signature alone is not enough: a rebuilt pool or a fresh process may + # hold identical content while account_catalogs was never published. + _published_catalog_signature = signature + CONFIG["account_catalogs"] = accounts + CONFIG["model_catalogs"] = catalogs + CONFIG["models_remote"], CONFIG["models_intl"] = catalogs["cn-cli"], catalogs["intl-cli"] + invalidate_model_table() # Drop the stale table first, then install the fresh index. + _publish_catalog_index(accounts) + + +def _account_catalog_key(profile, headers): + return catalog_cache_key(profile, account_key(profile, headers.get("X-User-Id"), + headers.get("X-Enterprise-Id"))) + + +def _fetch_catalog(headers): + """Fetch both catalog scopes for one account token; the only upstream call involved.""" + with _endpoint_semaphore("models"): + scopes = credits_mod.fetch_model_scopes( + _bearer_token(headers), domain=headers.get("X-Domain", ""), + uid=headers.get("X-User-Id", ""), enterprise_id=headers.get("X-Enterprise-Id", "")) + return scopes["picker"], scopes["account"] + + +def _store_catalog(pool, cache, entry, generation, key, models, serves): + """Publish one fetched catalog only while the credential lease is still current.""" + def publish(): + cache.put(key, models, serves=serves) + for current in pool._entries: + if current["cm"] is entry["cm"]: + current["catalog_dirty"] = False + return pool.apply_if_current(entry["cm"], generation, publish) + + +def _sync_one_account_catalog(pool, ledger, entry, generation, headers, profile, failed, cache): + """Refresh one account's own catalog; True when a fetch was stored, None when skipped.""" + if not model_policy.credential_enabled(CONFIG, entry): + return None + key = _account_catalog_key(profile, headers) + if cache.fresh(key) and not entry.get("catalog_dirty"): + return None + perf_metrics.inc("catalog_fetch_count") + try: + models, serves = _fetch_catalog(headers) + if _store_catalog(pool, cache, entry, generation, key, models, serves): + _log(f"[models] {profile} 模型表已刷新: 选择器 {len(models)} 个，" + f"账号根表 {len(serves)} 个") + return True + failed.add(entry["id"]) + except Exception as error: + perf_metrics.inc("catalog_fetch_failures") + failed.add(entry["id"]) + _sync_error(pool, ledger, entry, generation, "models", error) + return False -def _sync_model_catalogs(pool, ledger, refs, failed): - cache = CONFIG.get("model_cache") - if cache is None: +def _sync_account_catalogs(pool, ledger, refs, failed, cache): + """Refresh every synced account's own catalog; used by the account probe mode.""" + list(_bounded_map(lambda ref: _sync_one_account_catalog(pool, ledger, *ref, failed=failed, cache=cache), + list(refs.values()), _housekeep_concurrency())) + + +def _mark_catalog_diffs(pool, profile, model_ids, cache, previous_ids): + """Mark accounts lacking models newly advertised since the last representative probe.""" + added = model_ids - previous_ids + if not added: + # Older catalog versions are expected in mixed pools; only genuine additions matter. return - for entry, generation, headers, profile in refs.values(): - identity = account_key(profile, headers.get("X-User-Id"), headers.get("X-Enterprise-Id")) - if not model_policy.credential_enabled(CONFIG, entry): - continue - key = catalog_cache_key(profile, identity) - if cache.fresh(key) and not entry.get("catalog_dirty"): - continue + marked = [] + with pool._lock: + for entry in pool._entries: + identity, current = entry.get("account_key"), entry.get("profile") + if not identity or current != profile: + continue + known = {item.get("id") for item in (cache.models(catalog_cache_key(profile, identity)) or [])} + if known and not added <= known: + entry["catalog_dirty"] = True + marked.append(entry["id"]) + for cid in marked: # Queue outside the lock so the wake-up can publish the model cache. + pool._queue_sync(cid) + + +_LAGGARD_REFRESH_CAP = 16 # Per-cycle floor for back-filling missing or dirty account catalogs. + + +def _laggard_cap(total: int) -> int: + """Scale the per-cycle cap so any backlog converges within about twelve cycles.""" + return max(_LAGGARD_REFRESH_CAP, -(-total // 12)) + + +def _refresh_laggard_catalogs(pool, ledger, failed, cache): + """Back-fill account-level catalogs in representative mode, where probing skips them. + + An account without its own catalog is unroutable, so missing catalogs come first and + dirty accounts follow; progress wakes the next cycle instead of waiting an hour. + """ + missing, dirty = [], [] + policy = model_policy.view(CONFIG) # One snapshot for the whole loop; it holds the pool lock. + with pool._lock: + for entry in pool._entries: + if (not entry.get("account_key") or entry.get("profile") is None + or not model_policy.credential_enabled(CONFIG, entry, policy) + or not pool._healthy(entry, policy)): + continue + known = cache.age(catalog_cache_key(entry["profile"], entry["account_key"])) is not None + if not known: + missing.append(entry) + elif entry.get("catalog_dirty"): + dirty.append(entry) + backlog = missing + dirty + cap = _laggard_cap(len(backlog)) + if not backlog: + return + refs = {} + for entry in backlog[:cap]: + cm = entry["cm"] try: - scopes = credits_mod.fetch_model_scopes( - _bearer_token(headers), domain=headers.get("X-Domain", ""), - uid=headers.get("X-User-Id", ""), enterprise_id=headers.get("X-Enterprise-Id", "")) - models, serves = scopes["picker"], scopes["account"] - def publish(): - cache.put(key, models, serves=serves) - for current in pool._entries: - if current["cm"] is entry["cm"]: - current["catalog_dirty"] = False - if pool.apply_if_current(entry["cm"], generation, publish): - _log(f"[models] {profile} 模型表已刷新: 选择器 {len(models)} 个，" - f"账号根表 {len(serves)} 个") - else: - failed.add(entry["id"]) + # One broken account must not abort the pass before the model cache is published. + with cm._lock: + refs[entry["id"]] = (entry, cm._generation, cm.get_headers(), entry["profile"]) except Exception as error: failed.add(entry["id"]) - _sync_error(pool, ledger, entry, generation, "models", error) + _sync_error(pool, ledger, entry, entry.get("generation"), "models", error) + results = list(_bounded_map( + lambda ref: _sync_one_account_catalog(pool, ledger, *ref, failed=failed, cache=cache), + list(refs.values()), _housekeep_concurrency())) + if any(result for result in results) and len(backlog) > cap: + # Only real progress re-wakes the loop; failures fall back to the hourly cycle. + pool._sync_event.set() + + +def _probe_profile_catalog(pool, ledger, refs, failed, cache, aggregate, profile): + """Probe at most one representative credential for this product profile.""" + candidates = [{"entry": entry, "generation": generation, "headers": headers, + "account_key": entry.get("account_key")} + for entry, generation, headers, candidate_profile in list(refs.values()) + if candidate_profile == profile and entry.get("account_key") + and model_policy.credential_enabled(CONFIG, entry)] + if not candidates: + return # This cycle has no synced account for the profile; no probe happened. + target = aggregate.choose( + profile, candidates, age_of=lambda candidate: cache.age(_account_catalog_key(profile, candidate["headers"]))) + if target is None: + aggregate.fresh(profile) + return + entry, generation, headers = target["entry"], target["generation"], target["headers"] + key = _account_catalog_key(profile, headers) + perf_metrics.inc("catalog_fetch_count") + try: + models, serves = _fetch_catalog(headers) + if not _store_catalog(pool, cache, entry, generation, key, models, serves): + aggregate.failure(profile, source=entry.get("account_key")) + failed.add(entry["id"]) + return + previous_ids = aggregate.model_ids(profile) + aggregate.success(profile, source=entry.get("account_key"), generation=generation, + models=[item.get("id") for item in models]) + _mark_catalog_diffs(pool, profile, {item.get("id") for item in models}, cache, previous_ids) + _log(f"[models] {profile} 代表账号目录已刷新: 选择器 {len(models)} 个，" + f"账号根表 {len(serves)} 个") + except Exception as error: + perf_metrics.inc("catalog_fetch_failures") + failed.add(entry["id"]) + aggregate.failure(profile, source=entry.get("account_key")) + _sync_error(pool, ledger, entry, generation, "models", error) + + +def _probe_profile_catalogs(pool, ledger, refs, failed, cache, aggregate): + """Probe one representative credential per due product profile, each on its own slot.""" + due = [profile for profile in PROFILE_ENDPOINTS if aggregate.due(profile)] + if not due: + return + list(_bounded_map(lambda profile: _probe_profile_catalog(pool, ledger, refs, failed, cache, aggregate, profile), + due, _housekeep_concurrency())) + + +def _sync_model_catalogs(pool, ledger, refs, failed): + cache = CONFIG.get("model_cache") + if cache is None: + return + aggregate = CONFIG.get("aggregate_catalog") + if CONFIG.get("catalog_probe_mode", "account") == "representative" and aggregate is not None: + _probe_profile_catalogs(pool, ledger, refs, failed, cache, aggregate) + _refresh_laggard_catalogs(pool, ledger, failed, cache) + else: + _sync_account_catalogs(pool, ledger, refs, failed, cache) _publish_model_cache() +CATALOG_SCAN_START_LOCK = threading.Lock() # One user-triggered full scan at a time. +CATALOG_SCAN_PROGRESS_LOCK = threading.Lock() # Progress counters are shared by scan workers. +CATALOG_SCAN_MAX_ROWS = 500 # Bounded per-account progress payload. +_SCAN_FIELDS = ("status", "total", "completed", "success", "failed", "differences", + "started_at", "finished_at", "error", "profiles", "accounts") + + +def catalog_scan_view() -> dict: + """Return the current or last full catalog scan without any credential material.""" + scan = CONFIG.get("catalog_scan") + if not scan: + return {"status": "idle", "total": 0, "completed": 0, "success": 0, "failed": 0, + "differences": 0, "started_at": None, "finished_at": None, "error": None, + "profiles": {}, "accounts": []} + return {key: scan.get(key) for key in _SCAN_FIELDS} + + +def catalog_status() -> dict: + """Describe probe mode, representative state and the last full scan.""" + aggregate = CONFIG.get("aggregate_catalog") + return {"probe_mode": CONFIG.get("catalog_probe_mode", "account"), + "representative": aggregate.view() if aggregate is not None else None, + "scan": catalog_scan_view()} + + +def _scan_one_catalog(pool, cache, entry): + """Refresh one account's catalog and report how far it moved from the cached copy.""" + cm = entry["cm"] + with cm._lock: + headers = cm.get_headers() + generation = cm._generation + profile = profile_for_headers(headers) + key = _account_catalog_key(profile, headers) + previous = {item.get("id") for item in (cache.models(key) or [])} + models, serves = _fetch_catalog(headers) + if not _store_catalog(pool, cache, entry, generation, key, models, serves): + raise RuntimeError("扫描期间凭证已变化，已跳过该账号") + current = {item.get("id") for item in models} + # A first fetch has no baseline to diff against; only a change to a known catalog counts. + differences = len(previous ^ current) if previous else 0 + return {"profile": profile, "models": len(models), "serves": len(serves), + "differences": differences, "model_ids": [item.get("id") for item in models]} + + +def _scan_one_account(pool, cache, scan, aggregate, entry): + """Refresh one account's catalog and fold its result into the shared progress record.""" + row = {"name": Path(entry["id"]).name} + try: + if not model_policy.credential_enabled(CONFIG, entry): + row.update(status="skipped", reason="disabled") + else: + result = _scan_one_catalog(pool, cache, entry) + row.update(status="ok", **{key: value for key, value in result.items() if key != "model_ids"}) + with CATALOG_SCAN_PROGRESS_LOCK: + scan["success"] += 1 + scan["differences"] += result["differences"] + if aggregate is not None and entry.get("account_key"): + # A full scan satisfies this profile's TTL, so no probe follows immediately. + aggregate.success(result["profile"], source=entry["account_key"], + generation=entry.get("generation"), models=result["model_ids"]) + except Exception as error: + perf_metrics.inc("catalog_fetch_failures") + row.update(status="failed", error=sanitize_log_text(_network_error_text(error), 120)) + with CATALOG_SCAN_PROGRESS_LOCK: + scan["failed"] += 1 + with CATALOG_SCAN_PROGRESS_LOCK: + scan["completed"] += 1 + if len(scan["accounts"]) < CATALOG_SCAN_MAX_ROWS: + scan["accounts"].append(row) + + +def _catalog_profile_versions(pool, cache) -> dict: + """Summarize distinct catalog versions per profile after a scan, without identities.""" + versions: dict = {} + with pool._lock: + for entry in pool._entries: + profile, identity = entry.get("profile"), entry.get("account_key") + if not profile or not identity: + continue + ids = frozenset(item.get("id") for item in (cache.models(catalog_cache_key(profile, identity)) or [])) + if ids: + versions.setdefault(profile, {})[ids] = versions.setdefault(profile, {}).get(ids, 0) + 1 + return {profile: {"accounts": sum(counts.values()), "distinct": len(counts), + "versions": sorted(({"models": len(ids), "accounts": count} for ids, count in counts.items()), + key=lambda row: (-row["accounts"], row["models"]))[:5]} + for profile, counts in versions.items()} + + +def _catalog_scan_worker(pool, cache, scan, entries): + """Scan every account with bounded workers, keeping each previous catalog when a refresh fails.""" + aggregate = CONFIG.get("aggregate_catalog") + try: + list(_bounded_map(lambda entry: _scan_one_account(pool, cache, scan, aggregate, entry), + entries, _housekeep_concurrency())) + _publish_model_cache() + scan["profiles"] = _catalog_profile_versions(pool, cache) + scan["status"] = "completed" + except Exception as error: + scan["status"] = "failed" + scan["error"] = sanitize_log_text(_network_error_text(error), 120) + finally: + scan["finished_at"] = time.time() + _log(f"[models] 手动全量扫描结束: 成功 {scan['success']}，失败 {scan['failed']}，" + f"差异 {scan['differences']}") + + +def start_catalog_scan() -> dict: + """Start one user-triggered full scan; an already running scan is reported instead.""" + pool, cache = CONFIG.get("cred_pool"), CONFIG.get("model_cache") + if pool is None or cache is None or credits_mod is None: + raise HTTPException(status_code=503, detail={"error": { + "message": "模型目录不可用，无法扫描", "type": "service_unavailable", + "code": "catalog_unavailable"}}) + with CATALOG_SCAN_START_LOCK: + current = CONFIG.get("catalog_scan") + if current is not None and current.get("status") == "running": + return catalog_scan_view() + entries = pool.entries() + scan = {"status": "running", "total": len(entries), "completed": 0, "success": 0, + "failed": 0, "differences": 0, "started_at": time.time(), "finished_at": None, + "error": None, "profiles": {}, "accounts": []} + CONFIG["catalog_scan"] = scan + threading.Thread(target=_catalog_scan_worker, args=(pool, cache, scan, entries), + daemon=True, name="catalog-scan").start() + return catalog_scan_view() + + def _sync_usage(pool, entries=None, expected_identity=None): """Refresh usage during maintenance with independent per-account snapshots.""" accounts = CONFIG.get("usage_daily_accounts") @@ -1686,42 +2154,49 @@ def _sync_usage(pool, entries=None, expected_identity=None): previous_stale = set((CONFIG.get("usage_daily") or {}).get("stale_accounts", [])) stale = {entry["id"] for entry in pool.entries() if entry["id"] not in target_ids and Path(entry["id"]).name in previous_stale} - for entry in targets: - if not model_policy.credential_enabled(CONFIG, entry): - continue - try: - cm = entry["cm"] - with cm._lock: - if expected_identity is not None and cm.summary().get("account_key") != expected_identity: - stale.add(entry["id"]) - continue - headers = cm.get_headers() - generation = cm._generation - site = site_for_headers(headers) + list(_bounded_map( + lambda entry: _sync_one_usage(pool, entry, accounts, stale, expected_identity), + targets, _housekeep_concurrency())) + _publish_usage_daily(pool, stale) + return stale & target_ids + + +def _sync_one_usage(pool, entry, accounts, stale, expected_identity=None): + """Refresh one account's usage snapshot on the usage endpoint semaphore.""" + if not model_policy.credential_enabled(CONFIG, entry): + return + try: + cm = entry["cm"] + with cm._lock: + if expected_identity is not None and cm.summary().get("account_key") != expected_identity: + stale.add(entry["id"]) + return + headers = cm.get_headers() + generation = cm._generation + site = site_for_headers(headers) + with _endpoint_semaphore("usage"): usage = credits_mod.fetch_request_usage(_bearer_token(headers), uid=headers.get("X-User-Id", ""), domain=headers.get("X-Domain", "")) - def store(): - accounts[entry["id"]] = {"identity": entry.get("account_key"), "site": site, - "by_day": usage["by_day"], - "total_credits": round(usage["total_credits"], 2), - "requests": usage["requests"], - "partial": bool(usage.get("partial")), - "fetched_at": time.time()} - snapshots = CONFIG.get("usage_snapshots") - if snapshots is not None and entry.get("account_key") and entry.get("uid"): - # Only a validated identity may key durable state. An absent flag defaults to - # False, but a present value is passed through raw so the store's own strict - # validator governs: coercing it here with bool() would mask a malformed value - # and persist it as a legitimate flag. - snapshots.store(entry["id"], entry["account_key"], site, usage, - partial=usage.get("partial", False)) - if not pool.apply_if_current(cm, generation, store): - stale.add(entry["id"]) - except Exception as error: + def store(): + accounts[entry["id"]] = {"identity": entry.get("account_key"), "site": site, + "by_day": usage["by_day"], + "total_credits": round(usage["total_credits"], 2), + "requests": usage["requests"], + "partial": bool(usage.get("partial")), + "fetched_at": time.time()} + snapshots = CONFIG.get("usage_snapshots") + if snapshots is not None and entry.get("account_key") and entry.get("uid"): + # Only a validated identity may key durable state. An absent flag defaults to + # False, but a present value is passed through raw so the store's own strict + # validator governs: coercing it here with bool() would mask a malformed value + # and persist it as a legitimate flag. + snapshots.store(entry["id"], entry["account_key"], site, usage, + partial=usage.get("partial", False)) + if not pool.apply_if_current(cm, generation, store): stale.add(entry["id"]) - _log(f"[usage] {Path(entry['id']).name} 明细拉取失败（保留其上次成功快照）: {_network_error_text(error)}") - _publish_usage_daily(pool, stale) - return stale & target_ids + except Exception as error: + stale.add(entry["id"]) + _log(f"[usage] {Path(entry['id']).name} 明细拉取失败（保留其上次成功快照）: {_network_error_text(error)}") def _usage_row_expired(snap) -> bool: @@ -1860,10 +2335,11 @@ def _housekeep_once(pool: CredentialPool, ledger, *, pending_only=False): failed = set() try: refs = {} - for entry in pool.entries(): - if entry["id"] not in ids: - continue - result = _sync_credits(pool, ledger, entry, checkin=not pending_only, failed=failed) + targets = [entry for entry in pool.entries() if entry["id"] in ids] + results = _bounded_map( + lambda entry: _sync_account(pool, ledger, entry, checkin=not pending_only, failed=failed), + targets, _housekeep_concurrency()) + for entry, result in zip(targets, results): if result is not None: refs[entry["id"]] = result _sync_model_catalogs(pool, ledger, refs, failed) @@ -1876,6 +2352,12 @@ def _housekeep_once(pool: CredentialPool, ledger, *, pending_only=False): pool.end_sync(ids, failed) +def _sync_account(pool, ledger, entry, *, checkin, failed): + """One account's balance work, bounded by the credits endpoint semaphore.""" + with _endpoint_semaphore("credits"): + return _sync_credits(pool, ledger, entry, checkin=checkin, failed=failed) + + LEASE_INTERVAL_S = 30 # Running-instance heartbeat; offline tools treat older leases as stale. @@ -1981,6 +2463,9 @@ async def _protocol_http_exception(request: Request, exc: HTTPException): "account_catalogs": None, # Account-scoped models/serves; None enables legacy embedding "trial_ledger": None, "model_guard": True, # Reject models absent from authorized catalogs + "routing_scheduler": "shadow", # legacy | shadow | indexed eligibility shards + "catalog_probe_mode": "account", # account | representative catalog refresh + "housekeep_concurrency": 4, # Bounded maintenance workers; 1 restores serial behaviour "max_images": 16, "image_policy": "truncate", "max_request_bytes": 32 * 1024 * 1024, "log_body_limit": 65536, "max_inbound_bytes": 64 * 1024 * 1024, @@ -2205,6 +2690,11 @@ def _store_credential(directory: Path, name: str, content: bytes, uid: str, *, r target = atomic_write_credential(directory, name, content) if pool is not None: pool.reload([target]) + # Imports and logins always re-verify the catalog, matching relogin semantics. + entry = pool.entry_by_id(str(target)) + if entry is not None: + entry["catalog_dirty"] = True + pool._queue_sync(str(target)) elif cm is not None: cm.invalidate() return target @@ -2537,18 +3027,56 @@ def billing_usage(start_date: Optional[str] = None, end_date: Optional[str] = No # Prefer cloud catalogs over static fallback models. _MODEL_TABLE_TTL = 60.0 # Model snapshot lifetime in seconds _model_table_cache: dict = {} +_configured_cache: dict = {} # (region, pool revision, policy revision) -> enabled profiles +_catalog_revision = 0 def invalidate_model_table() -> None: - """Invalidate the public model snapshot after catalog changes.""" + """Drop the public model snapshot and derived catalog index after catalog changes.""" global _model_table_cache _model_table_cache = {} + CONFIG["catalog_index"] = None # Rebuilt from the live accounts on the next catalog read. + + +def _identity_enabled(identity) -> bool: + """Whether a published account identity still has an enabled credential.""" + pool = CONFIG.get("cred_pool") + entries = pool.entries_by_identity(identity) if pool is not None else () + if not entries: + return True + return any(model_policy.credential_enabled(CONFIG, entry) for entry in entries) + + +def _publish_catalog_index(accounts): + """Build the read-only catalog index once per publish, never per request.""" + global _catalog_revision + _catalog_revision += 1 + index = CatalogIndex(accounts, revision=_catalog_revision, + policy_revision=model_policy.revision(CONFIG), + is_enabled=_identity_enabled) + CONFIG["catalog_index"] = index + return index + + +def _catalog_index(): + """Return the catalog index matching the published accounts and policy revision.""" + accounts = CONFIG.get("account_catalogs") + if accounts is None: + return None + index = CONFIG.get("catalog_index") + if (index is not None and index.accounts is accounts + and index.policy_revision == model_policy.revision(CONFIG)): + return index + return _publish_catalog_index(accounts) def _catalog_for(profile: str, scope: str = "models", *, model_id=None): """Return a profile catalog within the requested account scope.""" accounts = CONFIG.get("account_catalogs") if accounts is not None or CONFIG.get("model_cache") is not None: + index = _catalog_index() + if index is not None: + return index.catalog(profile, scope, model_id=model_id) pool = CONFIG.get("cred_pool") models = None for entry in pool.entries() if pool is not None else []: @@ -2572,11 +3100,20 @@ def _in_region(profile: str, region: str | None) -> bool: return region is None or profile_region(profile) == region -def _configured_profiles(region: str | None) -> set[str]: +def _configured_profiles(region: str | None, policy=None) -> set[str]: + """Return enabled profiles in a region, cached until the pool or policy revision moves.""" pool = CONFIG.get("cred_pool") if pool is not None: - return {profile for entry in pool.entries() if (profile := pool._entry_profile(entry)) - and model_policy.credential_enabled(CONFIG, entry) and _in_region(profile, region)} + key = (region, getattr(pool, "revision", 0), model_policy.revision(CONFIG, policy)) + cached = _configured_cache.get(key) + if cached is not None: + return cached + profiles = frozenset(profile for entry in pool.entries() if (profile := pool._entry_profile(entry)) + and model_policy.credential_enabled(CONFIG, entry, policy) + and _in_region(profile, region)) + _configured_cache.clear() + _configured_cache[key] = profiles + return profiles cm = CONFIG.get("cred") if cm is not None: profile = cm.summary()["profile"] @@ -2591,16 +3128,6 @@ def _usable_models(models): and not model.get("disabled")] -def _account_scope(account: dict, scope: str = "models") -> list[dict] | None: - """Merge root candidates into selector models while retaining selector metadata.""" - picker = account.get("models") - if scope == "models" or picker is None: - return picker - seen = {item.get("id") for item in picker} - # Root entries supply missing names without replacing selector metadata. - return picker + [item for item in account.get("serves") or [] if item.get("id") not in seen] - - def _shared_catalog_sources(scope="models"): accounts, pool = CONFIG.get("account_catalogs"), CONFIG.get("cred_pool") if accounts is not None or CONFIG.get("model_cache") is not None: @@ -2614,6 +3141,8 @@ def _shared_catalog_sources(scope="models"): def _effective_account_scope(account, scope="models", *, model_id=None): + # Account-level scopes stay live so a capability change is rechecked inside a request; + # the published profile catalogs below come from the indexed snapshot. native = _account_scope(account, scope) profile = account.get("profile") if profile not in SHARED_INTL_PROFILES or native is None: @@ -2622,11 +3151,23 @@ def _effective_account_scope(account, scope="models", *, model_id=None): -def _models_for_profile(profile: str, configured=None, *, scope: str = "models", model_id=None) -> list[dict]: +def _published_account_scope(account, scope="models", *, model_id=None): + """Per-account catalog for a published table: indexed at publish time, live as a fallback.""" + index = _catalog_index() + if index is not None: + items = index.account_catalog(account, scope) + if items is not None: + return items if model_id is None else [item for item in items if item.get("id") == model_id] + return _effective_account_scope(account, scope, model_id=model_id) + + +def _models_for_profile(profile: str, configured=None, *, scope: str = "models", model_id=None, + policy=None) -> list[dict]: models = _catalog_for(profile, scope, model_id=model_id) if models is None: # Static fallback is limited to legacy domestic CLI deployments. - configured = _configured_profiles(profile_region(profile)) if configured is None else configured + configured = (_configured_profiles(profile_region(profile), policy) + if configured is None else configured) return ([{"id": name, "supportsToolCall": True} for name in DEFAULT_MODELS] if CONFIG.get("model_cache") is None and CONFIG.get("account_catalogs") is None and profile == "cn-cli" and configured <= {"cn-cli"} else []) @@ -2669,21 +3210,23 @@ def _model_free(models, model: str | None, profile: str) -> bool: for item in models or []) -def _model_profiles(model: str | None, region: str | None = None, configured=None) -> set[str]: - configured = _configured_profiles(region) if configured is None else configured +def _model_profiles(model: str | None, region: str | None = None, configured=None, *, policy=None) -> set[str]: + configured = _configured_profiles(region, policy) if configured is None else configured profiles = {profile for profile in PROFILE_ENDPOINTS if _in_region(profile, region)} if not model: return profiles supported = {profile for profile in profiles if any(item["id"] == _upstream_model(model, profile) - for item in _models_for_profile(profile, configured, scope="serves", model_id=_upstream_model(model, profile)))} + for item in _models_for_profile(profile, configured, scope="serves", + model_id=_upstream_model(model, profile), policy=policy))} if model == "auto" and region == "cn": # WorkBuddy uses its advertised Auto; legacy CLI defaults remain separate. if "cn-work" in configured and "cn-work" in supported: return {"cn-work"} - if "cn-cli" in configured and _models_for_profile("cn-cli", configured): + if "cn-cli" in configured and _models_for_profile("cn-cli", configured, policy=policy): return {"cn-cli"} - if model == "auto" and region is None and "cn-cli" in configured and _models_for_profile("cn-cli", configured): + if (model == "auto" and region is None and "cn-cli" in configured + and _models_for_profile("cn-cli", configured, policy=policy)): supported.add("cn-cli") if model != "auto" and not supported and not CONFIG.get("model_guard") and len(configured) == 1: return configured @@ -2719,11 +3262,49 @@ def _positive(entry): for entry in pool.entries() if model_policy.credential_enabled(CONFIG, entry)) +def _legacy_catalog_signature() -> tuple: + """Cheap signature of legacy CONFIG catalogs, which are only ever replaced wholesale.""" + catalogs = CONFIG.get("model_catalogs") + rows = (tuple(sorted((str(name), len(items or ())) for name, items in catalogs.items())) + if isinstance(catalogs, dict) else ()) + remote, intl = CONFIG.get("models_remote"), CONFIG.get("models_intl") + return (rows, + id(remote) if isinstance(remote, list) else None, len(remote) if isinstance(remote, list) else 0, + id(intl) if isinstance(intl, list) else None, len(intl) if isinstance(intl, list) else 0) + + +def _model_table_revisions() -> tuple: + """Return the revisions that a published model table depends on.""" + pool = CONFIG.get("cred_pool") + ledger = CONFIG.get("ledger") + index = CONFIG.get("catalog_index") + return (getattr(pool, "revision", 0), model_policy.revision(CONFIG), + getattr(index, "revision", 0), getattr(ledger, "revision", 0), + _legacy_catalog_signature()) + + +def _model_table(name, region, build): + """Reuse a published model table until a revision moves or the short TTL lapses.""" + _catalog_index() # A rebuilt index bumps its revision; the key must see the fresh value. + key = (name, region, *_model_table_revisions()) + cached = _model_table_cache.get(key) + if cached is not None and cached[0] > time.monotonic(): + perf_metrics.inc("model_table_hits") + return cached[1] + value = build(region) + _model_table_cache.clear() + _model_table_cache[key] = (time.monotonic() + _MODEL_TABLE_TTL, value) + perf_metrics.inc("model_table_builds") + return value + + def current_models(region: str | None = None) -> list[str]: """Merge models eligible for current accounts without region-specific client URLs.""" + return _model_table("models", region, _build_current_models) + + +def _build_current_models(region: str | None = None) -> list[str]: pool = CONFIG.get("cred_pool") - if pool is not None: - pool._rescan() with pool._lock if pool is not None else nullcontext(): configured = _configured_profiles(region) out, has_auto = [], False @@ -2743,7 +3324,7 @@ def current_models(region: str | None = None) -> list[str]: account = accounts.get(entry.get("account_key")) or {} if account.get("profile") != profile: continue - models = _usable_models(_effective_account_scope(account, "serves")) + models = _usable_models(_published_account_scope(account, "serves")) if zero: models = [model for model in models if _free_multiplier(model.get("credits"))] out.extend(model["id"] for model in models) @@ -2768,9 +3349,11 @@ def current_models(region: str | None = None) -> list[str]: def current_model_details(region: str | None = None) -> list[dict]: """Return public model rates with per-profile values and the minimum eligible rate.""" + return _model_table("details", region, _build_current_model_details) + + +def _build_current_model_details(region: str | None = None) -> list[dict]: pool = CONFIG.get("cred_pool") - if pool is not None: - pool._rescan() details: dict[str, dict] = {} declarations = {} for name in current_models(region): @@ -2806,7 +3389,7 @@ def record(profile: str, item: dict, *, zero: bool) -> None: account = accounts.get(entry.get("account_key")) or {} if account.get("profile") != profile: continue - for item in _usable_models(_effective_account_scope(account, "serves")): + for item in _usable_models(_published_account_scope(account, "serves")): record(profile, item, zero=zero) else: configured = _configured_profiles(region) @@ -2950,9 +3533,6 @@ def guard_model(name: str, *, region=None, resolved=False) -> None: model_policy.check_resolved(CONFIG, name) if not CONFIG.get("model_guard"): return - pool = CONFIG.get("cred_pool") - if pool is not None: - pool._rescan() if _model_profiles(name, region): return if _catalog_pending(region): @@ -4517,6 +5097,12 @@ def main(): ap.add_argument("--stream-mode", choices=_STREAM_MODES, default=os.environ.get("CODEBUDDY2API_STREAM_MODE", "compatible"), help="流式传输：compatible 保持兼容聚合策略，realtime 增量发送且不重生成工具参数；默认 compatible") + ap.add_argument("--catalog-probe-mode", choices=("account", "representative"), + default=os.environ.get("CODEBUDDY2API_CATALOG_PROBE_MODE", "account"), + help="目录探测：account 逐个账号刷新（默认）；representative 每个产品每周期只探测一个代表账号并补齐缺目录账号") + ap.add_argument("--routing-scheduler", choices=("legacy", "shadow", "indexed"), + default=os.environ.get("CODEBUDDY2API_ROUTING_SCHEDULER", "shadow"), + help="路由调度：shadow 记录新旧差异并仍返回旧结果（默认，差异为零后可切 indexed）；indexed 复用按模型分片；legacy 每次全量判断") ap.add_argument("--model-capability-guard", type=_boolean_arg, nargs="?", const=True, default=os.environ.get("CODEBUDDY2API_MODEL_CAPABILITY_GUARD", "true"), help="按账号模型声明预检图片、工具、思考和输出上限；false 仅关闭新增能力预检") @@ -4554,7 +5140,8 @@ def main(): "tool_call_max_retry", "max_inbound_bytes", "max_collect_bytes", "max_concurrent", "failover_max", "retry_write_timeout", "upstream_keepalive", "max_inflight_per_account", "request_context_mode", "stream_mode", "model_capability_guard", "admin_allowed_origins", - "responses_projection_mode", "responses_projection_max_bytes"): + "responses_projection_mode", "responses_projection_max_bytes", "routing_scheduler", + "catalog_probe_mode"): CONFIG[key] = getattr(args, key) CONFIG["api_key"] = args.api_key CONFIG["desensitize"] = args.desensitize @@ -4588,6 +5175,7 @@ def main(): ledger = credits_mod.CreditLedger(store=CONFIG["state_store"]) CONFIG["ledger"] = ledger CONFIG["model_cache"] = credits_mod.ModelCatalogCache(ttl=args.model_catalog_ttl, store=CONFIG["state_store"]) + CONFIG["aggregate_catalog"] = AggregateCatalog(ttl=args.model_catalog_ttl, store=CONFIG["state_store"]) CONFIG["cred_pool"].set_ledger(ledger) # Verify balance ownership before publishing catalogs. _publish_model_cache() # Publish cached usage before maintenance threads start, so the dashboard is populated from @@ -4610,6 +5198,8 @@ def main(): daemon=True, name="service-lease").start() threading.Thread(target=_refresher_loop, args=(CONFIG["cred_pool"],), daemon=True, name="cred-refresher").start() + threading.Thread(target=_reconcile_loop, args=(CONFIG["cred_pool"],), + daemon=True, name="cred-reconcile").start() if credits_mod is not None: threading.Thread(target=_housekeeper_loop, args=(CONFIG["cred_pool"], ledger), daemon=True, name="cred-housekeeper").start() diff --git a/docs/advanced.md b/docs/advanced.md index c2a3d23..1140947 100644 --- a/docs/advanced.md +++ b/docs/advanced.md @@ -239,6 +239,13 @@ Credential domain / token issuer determine the product identity. Chat and refres - WebUI region, product and credential bindings strictly limit candidates; unavailable bindings never fall back to unselected accounts. Disabled models also reject direct requests. Renaming hides the original ID unless you choose to retain it. - Sent requests are not replayed against another account because of account availability or HTTP errors; later requests select again. Pending catalog/credential readiness usually returns 503 with `Retry-After`; unsupported or disabled models return 404. +Routing and catalog maintenance are switchable without reverting data: + +- `routing_scheduler`: `shadow` (default) runs the legacy scan and the indexed shards side by side, returns the legacy answer and records bounded, credential-free mismatches under `GET /admin/performance`; switch to `indexed` once the mismatch counter stays at zero; `indexed` reuses one eligibility shard per model, request-scoped alias and revision set; `legacy` re-evaluates every credential on every request. Shards are invalidated only when entry identities, the policy, the published catalogs or an account's solvency actually change, so token refreshes and maintenance writes keep the steady state. +- `catalog_probe_mode`: `account` (default) refreshes every synced account each cycle; `representative` probes one credential per product profile per TTL cycle — accounts without a catalog first, then the current source, then the stalest — with exponential backoff after failures, and back-fills missing or dirty account catalogs under a bounded per-cycle quota so new accounts become routable without a manual scan. An unprobed or failed account keeps its last trusted catalog, `/v1/models` stays the union of current account catalogs, and only models newly advertised by the representative mark older accounts dirty. +- `POST /admin/models/scan` refreshes every current account once, in the background; `GET /admin/models/scan` reports progress, successes, failures, catalog differences, the overall error and a per-profile breakdown of distinct catalog versions. It is the only way to calibrate all accounts at once and is never a periodic task. +- `housekeep_concurrency` (default `4`) bounds maintenance workers; balances, catalogs and usage each get an independent slot pool, `1` restores the previous serial cycle, and failed account syncs retry with exponential backoff plus jitter. Accounts synced within five minutes are skipped by the next full cycle, so an event-driven cycle never repeats them. + ### Model declarations and image compatibility `/v1/models` retains its existing fields and adds `capabilities`, `limits` and `metadata_by_profile`; management and routing previews expose the same metadata. Capabilities use `supported`, `unsupported`, `mixed` or `unknown`. Limits carry `state` (`known`, `mixed`, `unknown`) and `value`; only unanimous known limits have a numeric value. Per-profile arrays preserve distinct account declarations without identities. Safe descriptions, capabilities, windows, reasoning options, related models and parameter suggestions are allowlisted; credentials, internal configuration and authenticated URLs are excluded. These are upstream declarations, not native-model or measured guarantees; suggestions do not override requests. diff --git a/docs/advanced.zh-CN.md b/docs/advanced.zh-CN.md index 7b2c076..7aef778 100644 --- a/docs/advanced.zh-CN.md +++ b/docs/advanced.zh-CN.md @@ -239,6 +239,13 @@ WebUI 可以直接上传文件；以下限制针对 `POST /admin/credentials` - WebUI 的地域、产品和凭证绑定严格限制候选账号，不会回退到未选账号。模型禁用后直接请求同样拒绝；改名默认不保留原 ID，只有选择保留时才同时提供旧 ID。 - 已发送请求不会因账号不可用或 HTTP 错误换账号重放；后续请求才重新选路。目录同步或凭证未就绪通常返回带 `Retry-After` 的 503，不支持或禁用的模型返回 404。 +选路与目录维护都可切换模式，无需回滚已持久化数据： + +- `routing_scheduler`：`shadow`（默认）同时执行旧的全量判断与索引分片，返回旧结果，并把有界、不含凭证信息的差异记录在 `GET /admin/performance`；差异计数持续为零后再切到 `indexed`；`indexed` 按模型、请求内别名和版本组合复用同一份候选分片；`legacy` 每次请求重新判断全部凭证。只有条目身份、策略、已发布目录或账号偿付能力真实变化才会使分片失效，令牌刷新和维护写入不打断稳态。 +- `catalog_probe_mode`：`account`（默认）每周期刷新全部已同步账号；`representative` 每个产品每 TTL 周期只探测一个凭证——无目录账号优先，其次当前来源，再次目录最旧——失败按指数退避，并按每周期有界配额补齐缺失或标脏的账号目录，新账号无需手动扫描即可路由。未探测或失败的账号保留上次可信目录，`/v1/models` 仍是当前账号目录并集，只有代表目录新增模型才会标脏旧版本账号。 +- `POST /admin/models/scan` 在后台把当前全部账号各刷新一次，`GET /admin/models/scan` 返回进度、成功数、失败数、目录差异数、整体错误原因和按产品分组的目录版本分布；它是一次性校准入口，不是周期任务。 +- `housekeep_concurrency`（默认 `4`）限制后台维护并发：余额、目录和用量各自独立限额，设为 `1` 即回到原来的串行周期；账号同步失败按指数退避加抖动重试。5 分钟内已同步的账号会被下一个完整周期跳过，因此事件触发周期不会重复同步。 + ### 模型声明与图片兼容 `/v1/models` 保留原字段，增加 `capabilities`、`limits`、`metadata_by_profile`，管理接口及路由预览同步提供。能力状态为 `supported`、`unsupported`、`mixed`、`unknown`；限制含 `state`（`known`、`mixed`、`unknown`）和 `value`，仅已知且一致时返回数值。按产品保留不同账号的声明版本，但不公开账号身份。白名单覆盖说明、能力、窗口、思考选项、关联模型及参数建议；凭据、内部配置和带认证信息的 URL 不公开。上游声明不等于原生能力或实测保证，参数建议不自动覆盖请求。 diff --git a/tests/test_admin_api.py b/tests/test_admin_api.py index 731d12c..35e431e 100644 --- a/tests/test_admin_api.py +++ b/tests/test_admin_api.py @@ -39,6 +39,12 @@ def setUp(self): admin_set_credential_enabled=Mock(side_effect=lambda identity, enabled: self.store.set_credential(identity, enabled)), admin_apply_settings=Mock(), admin_delete_guard=Mock(), managed_auth_dir=Mock(return_value=self.root), _store_credential=Mock(), + catalog_status=Mock(return_value={"probe_mode": "account", "representative": None, + "scan": {"status": "idle", "total": 0, "completed": 0, + "success": 0, "failed": 0, "differences": 0, + "started_at": None, "finished_at": None, "accounts": []}}), + start_catalog_scan=Mock(return_value={"status": "running"}), + catalog_scan_view=Mock(return_value={"status": "idle"}), _credential_identity=Mock(return_value="fingerprint"), _save_oauth_credential=Mock(return_value=self.root / "first.info"), _OAUTH=Mock()) self.gateway._OAUTH.start.return_value = {"login_id": "task", "verification_uri": "https://www.codebuddy.cn/login", "expires_in": 600} @@ -480,6 +486,19 @@ def inventory(): self.assertEqual(response.json()["models"][0]["public_id"], "synced-model") + def test_catalog_scan_routes_require_session_and_delegate(self): + self.assertEqual(self.client.post("/admin/models/scan").status_code, 401) + self.assertEqual(self.client.get("/admin/models/scan").status_code, 401) + csrf = self.login() + started = self.client.post("/admin/models/scan", headers=csrf) + self.assertEqual(started.status_code, 200, started.text) + self.assertEqual(started.json(), {"status": "running"}) + self.gateway.start_catalog_scan.assert_called_once() + state = self.client.get("/admin/models/scan") + self.assertEqual(state.status_code, 200, state.text) + self.assertEqual(state.json(), {"status": "idle"}) + self.gateway.catalog_scan_view.assert_called() + def test_dashboard_granularity_is_validated_and_forwarded(self): response = self.client.get("/admin/dashboard?days=1&granularity=hour", headers=self.headers) self.assertEqual(response.status_code, 200, response.text) diff --git a/tests/test_catalog_probe.py b/tests/test_catalog_probe.py new file mode 100644 index 0000000..1a5d8c8 --- /dev/null +++ b/tests/test_catalog_probe.py @@ -0,0 +1,326 @@ +"""Representative catalog probes, retry backoff, diff marking and the manual full scan.""" +import sys +from pathlib import Path + +sys.path.insert(0, str(Path(__file__).resolve().parents[1])) # Allow direct execution. + +import json +import os +import tempfile +import time +import unittest +from unittest.mock import patch + +import converter +from app import model_policy, perf_metrics +from app.catalog_probe import AggregateCatalog, valid_row +from app.control_store import ControlStore +from app.credits import CreditLedger, ModelCatalogCache + + +def model(*names): + return [{"id": name, "supportsToolCall": True, "credits": "x1 credits"} for name in names] + + +class AggregateCatalogTests(unittest.TestCase): + def test_round_trip_persists_and_rejects_invalid_rows(self): + root = Path(self.enterContext(tempfile.TemporaryDirectory())) + store = ControlStore(root / "control.sqlite3") + self.addCleanup(store.close) + aggregate = AggregateCatalog(ttl=3600, store=store.state) + aggregate.success("cn-cli", source="fingerprint-a", generation=3, models=["glm-5.3", "hy3"]) + reloaded = AggregateCatalog(ttl=3600, store=store.state) + row = reloaded.row("cn-cli") + self.assertEqual(row["source"], "fingerprint-a") + self.assertEqual(row["source_generation"], 3) + self.assertEqual(row["models"], ["glm-5.3", "hy3"]) + self.assertFalse(reloaded.due("cn-cli")) + with self.assertRaises(ValueError): + store.state.put("catalog_probe", {"version": 1, "profiles": {"cn-cli": {"source": 1}}}) + self.assertEqual(len(AggregateCatalog(ttl=3600, store=store.state).profiles()), 1) + self.assertFalse(valid_row({"source": "a", "fetched_at": 1, "retry_at": 2, "status": "ok", + "models": ["x"], "secret": "y"})) + + def test_failure_backs_off_then_recovers(self): + aggregate = AggregateCatalog(ttl=3600, clock=lambda: 1000.0) + aggregate.failure("cn-cli", source="fingerprint-a") + self.assertFalse(aggregate.due("cn-cli")) + row = aggregate.row("cn-cli") + self.assertEqual(row["status"], "error") + self.assertGreater(row["retry_at"], 1000.0) + aggregate.failure("cn-cli", source="fingerprint-a") + self.assertGreater(aggregate.row("cn-cli")["retry_at"], row["retry_at"]) # Exponential backoff. + aggregate.success("cn-cli", source="fingerprint-a", generation=1, models=["hy3"]) + self.assertEqual(aggregate.row("cn-cli")["failures"], 0) + self.assertTrue(aggregate.due("cn-cli", now=1000.0 + 3601)) + + def test_choose_prefers_unknown_then_source_then_the_stalest(self): + aggregate = AggregateCatalog(ttl=100, clock=lambda: 1000.0) + accounts = [{"account_key": "a"}, {"account_key": "b"}, {"account_key": "c"}] + ages = {"a": 500.0, "b": 1000.0, "c": None} + age_of = lambda candidate: ages[candidate["account_key"]] + aggregate.success("cn-cli", source="a", generation=1, models=["hy3"]) + aggregate._data["profiles"]["cn-cli"]["retry_at"] = 0.0 + self.assertEqual(aggregate.choose("cn-cli", accounts, age_of=age_of)["account_key"], "c") # Unknown first. + ages["c"] = 10.0 + self.assertEqual(aggregate.choose("cn-cli", accounts, age_of=age_of)["account_key"], "a") # Source reused. + ages["a"] = 10.0 + self.assertEqual(aggregate.choose("cn-cli", accounts, age_of=age_of)["account_key"], "b") # Stalest. + ages["b"] = 10.0 + self.assertIsNone(aggregate.choose("cn-cli", accounts, age_of=age_of)) + + +class CatalogProbeTests(unittest.TestCase): + def setUp(self): + temporary = self.enterContext(tempfile.TemporaryDirectory()) + self.enterContext(patch.dict(os.environ, {"CODEBUDDY_AUTH_DIR": temporary})) + self.enterContext(patch.object(perf_metrics, "_COUNTERS", {})) + self.root = Path(temporary) + self.calls = [] + + def build(self, count=2, profile="cn-cli"): + paths = [] + for index in range(count): + path = self.root / f"account-{index}.info" + path.write_text(json.dumps({"account": {"uid": f"probe-{index}"}, "auth": { + "accessToken": f"token-{index}", "domain": "www.codebuddy.cn", + "expiresAt": (time.time() + 3600) * 1000}}), encoding="utf-8") + paths.append(path) + pool = converter.CredentialPool(paths) + ledger = CreditLedger(self.root / "ledger.json") + pool.set_ledger(ledger) + for entry in pool.entries(): + ledger.bind_identity(entry["id"], entry["account_key"]) + ledger.update_credits(entry["id"], {"credits": 100, "intl": False, "segments": []}) + cache = ModelCatalogCache(self.root / "catalog.json", ttl=60) + aggregate = AggregateCatalog(ttl=3600) + self.enterContext(patch.dict(converter.CONFIG, { + "cred_pool": pool, "ledger": ledger, "model_cache": cache, "aggregate_catalog": aggregate, + "catalog_probe_mode": "representative", "account_catalogs": {}, "catalog_index": None, + "model_catalogs": {}, "model_guard": True, "log_path": None, "control_store": None})) + return pool, ledger, cache, aggregate + + def refs(self, pool): + result = {} + for entry in pool.entries(): + cm = entry["cm"] + with cm._lock: + headers = cm.get_headers() + result[entry["id"]] = (entry, cm._generation, headers, entry["profile"]) + return result + + def scopes(self, *names): + def fetch(_token, **_kwargs): + self.calls.append(names) + return {"picker": model(*names), "account": model(*names)} + return fetch + + def test_representative_probe_runs_once_per_profile_and_makes_accounts_routable(self): + pool, ledger, cache, aggregate = self.build(3) + entries = pool.entries() + cache.put(converter.catalog_cache_key("cn-cli", entries[1]["account_key"]), + model("hy3"), serves=model("hy3")) + with patch.object(converter.credits_mod, "fetch_model_scopes", + side_effect=self.scopes("hy3", "glm-5.3")): + converter._sync_model_catalogs(pool, ledger, self.refs(pool), set()) + # One probe plus the two laggard refreshes (older-version and catalog-less accounts). + self.assertEqual(len(self.calls), 3) + self.assertEqual([len(call) for call in self.calls], [2, 2, 2]) + for snapshot in entries: + entry = pool.entry_by_id(snapshot["id"]) # Snapshot copies predate the refresh. + self.assertFalse(entry.get("catalog_dirty")) + known = {item["id"] for item in cache.models( + converter.catalog_cache_key("cn-cli", entry["account_key"]))} + self.assertTrue({"hy3", "glm-5.3"} <= known) + self.assertTrue(pool._eligible(entry, "glm-5.3")) + self.assertFalse(aggregate.due("cn-cli")) + self.assertEqual(aggregate.row("cn-cli")["models"], ["glm-5.3", "hy3"]) + with patch.object(converter.credits_mod, "fetch_model_scopes", side_effect=self.scopes("hy3")): + converter._sync_model_catalogs(pool, ledger, self.refs(pool), set()) + self.assertEqual(len(self.calls), 3) # The TTL suppresses another fetch. + + def test_laggard_backlog_converges_over_bounded_cycles(self): + pool, ledger, cache, _ = self.build(6) + with patch.object(converter, "_LAGGARD_REFRESH_CAP", 2), \ + patch.object(converter.credits_mod, "fetch_model_scopes", side_effect=self.scopes("hy3")): + for _ in range(10): + pool._sync_event.clear() + converter._sync_model_catalogs(pool, ledger, self.refs(pool), set()) + if not any(entry.get("catalog_dirty") for entry in pool.entries()) and all( + cache.models(converter.catalog_cache_key(entry["profile"], entry["account_key"])) + for entry in pool.entries()): + break + self.assertEqual(len(self.calls), 6) # One probe plus five back-fills, no repeats. + for entry in pool.entries(): + self.assertTrue(pool._eligible(entry, "hy3")) + + def test_pending_cycle_without_candidates_keeps_retry_backoff(self): + aggregate = AggregateCatalog(ttl=3600, clock=lambda: 1000.0) + aggregate.success("cn-cli", source="fingerprint-a", generation=1, models=["hy3"]) + aggregate.failure("cn-cli", source="fingerprint-a") + before = aggregate.row("cn-cli") + pool, ledger, cache, _ = self.build(1) + converter._probe_profile_catalog(pool, ledger, {}, set(), cache, aggregate, "cn-cli") + self.assertEqual(aggregate.row("cn-cli"), before) # No probe ran, so no fresh reset. + self.assertEqual(self.calls, []) + + def test_steady_representative_catalogs_are_not_remarked_dirty(self): + pool, ledger, cache, aggregate = self.build(2) + with patch.object(converter.credits_mod, "fetch_model_scopes", + side_effect=self.scopes("hy3", "glm-5.3")): + converter._sync_model_catalogs(pool, ledger, self.refs(pool), set()) + self.assertEqual(len(self.calls), 2) + # Force the next probe while keeping the same model set: an unchanged representative + # catalog must not re-mark the older-version account. + aggregate._data["profiles"]["cn-cli"]["retry_at"] = 0.0 + with patch.object(cache, "age", lambda _key: 3601), \ + patch.object(converter.credits_mod, "fetch_model_scopes", side_effect=self.scopes("hy3")): + converter._sync_model_catalogs(pool, ledger, self.refs(pool), set()) + self.assertEqual(len(self.calls), 3) # Only the probe ran; no account re-fetch. + self.assertFalse(any(entry.get("catalog_dirty") for entry in pool.entries())) + + def test_laggard_pass_survives_a_broken_account(self): + pool, ledger, cache, aggregate = self.build(3) + entries = pool.entries() + aggregate.success("cn-cli", source="seeded", generation=1, models=["hy3"]) # Probe not due. + broken = pool.entry_by_id(entries[0]["id"]) + failed = set() + refs = self.refs(pool) + with patch.object(broken["cm"], "get_headers", side_effect=RuntimeError("synthetic refresh")), \ + patch.object(converter.credits_mod, "fetch_model_scopes", side_effect=self.scopes("hy3")): + converter._sync_model_catalogs(pool, ledger, refs, failed) + self.assertIn(broken["id"], failed) + self.assertEqual(len(self.calls), 2) # The two healthy accounts still back-filled. + for entry in entries[1:]: + self.assertTrue(cache.models(converter.catalog_cache_key("cn-cli", entry["account_key"]))) + # The publish still ran despite the broken account. + self.assertIn(entries[1]["account_key"], converter.CONFIG["account_catalogs"]) + + def test_laggard_pass_skips_accounts_in_cooldown(self): + pool, ledger, cache, aggregate = self.build(2) + entries = pool.entries() + aggregate.success("cn-cli", source="seeded", generation=1, models=["hy3"]) + refs = self.refs(pool) # Built before the cooldown so no headers are touched. + cooling = pool.entry_by_id(entries[0]["id"]) + cooling["fail_until"] = time.time() + 300 + with patch.object(cooling["cm"], "get_headers", side_effect=RuntimeError("must not be called")), \ + patch.object(converter.credits_mod, "fetch_model_scopes", side_effect=self.scopes("hy3")): + converter._sync_model_catalogs(pool, ledger, refs, set()) + self.assertEqual(cache.models(converter.catalog_cache_key("cn-cli", cooling["account_key"])), []) + self.assertEqual(len(self.calls), 1) + + def test_probe_failure_keeps_other_catalogs_and_backs_off(self): + pool, ledger, cache, aggregate = self.build(2) + entries = pool.entries() + cache.put(converter.catalog_cache_key("cn-cli", entries[0]["account_key"]), model("hy3"), serves=model("hy3")) + failed = set() + with patch.object(converter.credits_mod, "fetch_model_scopes", side_effect=RuntimeError("synthetic upstream")): + converter._sync_model_catalogs(pool, ledger, self.refs(pool), failed) + self.assertTrue(failed) + self.assertLessEqual({"source", "status", "failures", "retry_at"}, set(aggregate.row("cn-cli"))) + self.assertEqual(aggregate.row("cn-cli")["status"], "error") + self.assertFalse(aggregate.due("cn-cli")) + kept = {item["id"] for item in cache.models(converter.catalog_cache_key("cn-cli", entries[0]["account_key"]))} + self.assertEqual(kept, {"hy3"}) # The previous catalog survives a failed probe. + + def test_account_mode_still_refreshes_every_account(self): + pool, ledger, cache, _ = self.build(2) + self.enterContext(patch.dict(converter.CONFIG, {"catalog_probe_mode": "account"})) + with patch.object(converter.credits_mod, "fetch_model_scopes", side_effect=self.scopes("hy3")): + converter._sync_model_catalogs(pool, ledger, self.refs(pool), set()) + self.assertEqual(len(self.calls), 2) + + def test_manual_scan_refreshes_every_account_and_reports_differences(self): + pool, _, cache, aggregate = self.build(3) + entries = pool.entries() + cache.put(converter.catalog_cache_key("cn-cli", entries[0]["account_key"]), model("hy3"), serves=model("hy3")) + with patch.object(converter.credits_mod, "fetch_model_scopes", side_effect=self.scopes("hy3", "glm-5.3")): + started = converter.start_catalog_scan() + self.assertEqual(started["status"], "running") + for _ in range(200): + view = converter.catalog_scan_view() + if view["status"] != "running": + break + time.sleep(0.01) + self.assertEqual(view["status"], "completed") + self.assertEqual((view["total"], view["completed"], view["success"], view["failed"]), (3, 3, 3, 0)) + self.assertEqual(view["differences"], 1) # Only the pre-seeded account changed. + self.assertEqual(len(view["accounts"]), 3) + self.assertEqual(aggregate.row("cn-cli")["status"], "ok") + self.assertNotIn("token", json.dumps(view)) + + def test_manual_scan_reports_cross_account_catalog_versions(self): + pool, _, cache, _ = self.build(3) + + def varied(token, **_kwargs): + names = ("hy3", "glm-5.3") if token != "token-0" else ("hy3",) + self.calls.append(token) + return {"picker": model(*names), "account": model(*names)} + + with patch.object(converter.credits_mod, "fetch_model_scopes", side_effect=varied): + converter.start_catalog_scan() + for _ in range(200): + view = converter.catalog_scan_view() + if view["status"] != "running": + break + time.sleep(0.01) + summary = view["profiles"]["cn-cli"] + self.assertEqual((summary["accounts"], summary["distinct"]), (3, 2)) + self.assertEqual(summary["versions"], [{"models": 2, "accounts": 2}, {"models": 1, "accounts": 1}]) + self.assertNotIn("token", json.dumps(summary)) + + def test_manual_scan_keeps_a_failed_accounts_previous_catalog(self): + pool, _, cache, _ = self.build(3) + entries = pool.entries() + failed_key = converter.catalog_cache_key("cn-cli", entries[0]["account_key"]) + cache.put(failed_key, model("hy3"), serves=model("hy3")) + + def flaky(token, **_kwargs): + if token == "token-0": + raise RuntimeError("synthetic upstream") + self.calls.append(token) + return {"picker": model("glm-5.3"), "account": model("glm-5.3")} + + with patch.object(converter.credits_mod, "fetch_model_scopes", side_effect=flaky): + converter.start_catalog_scan() + for _ in range(200): + view = converter.catalog_scan_view() + if view["status"] != "running": + break + time.sleep(0.01) + self.assertEqual(view["status"], "completed") + self.assertEqual((view["success"], view["failed"]), (2, 1)) + self.assertEqual([row["status"] for row in view["accounts"]].count("failed"), 1) + # The failed account keeps its previous catalog instead of going empty. + self.assertEqual({item["id"] for item in cache.models(failed_key)}, {"hy3"}) + self.assertEqual(set(self.calls), {"token-1", "token-2"}) + + def test_manual_scan_reports_an_already_running_scan(self): + pool, _, cache, _ = self.build(1) + blocker = {} + + def slow_fetch(*args, **kwargs): + blocker["started"] = True + time.sleep(0.05) + return {"picker": model("hy3"), "account": model("hy3")} + with patch.object(converter.credits_mod, "fetch_model_scopes", side_effect=slow_fetch): + first = converter.start_catalog_scan() + second = converter.start_catalog_scan() + for _ in range(200): + if converter.catalog_scan_view()["status"] != "running": + break + time.sleep(0.01) + self.assertEqual(first["status"], "running") + self.assertEqual(second["status"], "running") + self.assertTrue(blocker.get("started")) + + def test_scan_reports_unavailable_catalog_without_credential_material(self): + pool, _, cache, _ = self.build(1) + with patch.dict(converter.CONFIG, {"model_cache": None}): + with self.assertRaises(converter.HTTPException) as raised: + converter.start_catalog_scan() + self.assertEqual(raised.exception.status_code, 503) + + +if __name__ == "__main__": + unittest.main(verbosity=2) diff --git a/tests/test_credential_runtime.py b/tests/test_credential_runtime.py index ed1ab0f..1fdc07a 100644 --- a/tests/test_credential_runtime.py +++ b/tests/test_credential_runtime.py @@ -264,6 +264,30 @@ def test_directory_changes_still_add_and_remove_accounts(self): pool._rescan() self.assertEqual([e["id"] for e in pool.entries()], [str(second)]) + def test_request_path_stays_indexed_and_reconcile_applies_file_changes(self): + first = self.credential(age=0) + pool = converter.CredentialPool([first], scan=True) + manager = pool.first() + self.assertIs(pool.entry_for(manager), pool._entries[0]) + self.assertIs(pool.entry_by_id(str(first)), pool._entries[0]) + second = self.credential("second.info", uid="second", age=0) + self.assertIsNotNone(pool.pick("session", "glm-5.3-flash")) # A request never scans the directory. + self.assertEqual(len(pool.entries()), 1) + reloads = [] + original = pool.reload + with patch.object(pool, "reload", side_effect=lambda paths, **kw: (reloads.append(list(paths)), + original(paths, **kw))[1]): + fingerprint = pool.reconcile() + self.assertEqual(len(pool.entries()), 2) + self.assertTrue(reloads) + reloads.clear() + self.assertEqual(pool.reconcile(fingerprint), fingerprint) + self.assertEqual(reloads, []) # An unchanged directory is not re-read. + first.unlink() + pool.reconcile(fingerprint) + self.assertEqual([e["id"] for e in pool.entries()], [str(second)]) + self.assertIsNone(pool.entry_for(manager)) + def test_encrypted_envelope_tokens_fail_fast_and_stay_out_of_pool(self): valid = self.credential("valid.info", uid="valid", age=0) envelope = self.root / "envelope.info" diff --git a/tests/test_environment_config.py b/tests/test_environment_config.py index bc1b1e3..35a9083 100644 --- a/tests/test_environment_config.py +++ b/tests/test_environment_config.py @@ -118,6 +118,17 @@ def test_invalid_pooling_environment_fails_before_startup(self): with self.subTest(env=env), self.assertRaises(SystemExit): self.start(env) + def test_rollout_switches_default_to_conservative_modes(self): + _, items, config = self.start() + self.assertEqual(config['routing_scheduler'], 'shadow') + self.assertEqual(config['catalog_probe_mode'], 'account') + self.assertEqual(items['routing_scheduler']['source'], 'default') + self.assertEqual(items['catalog_probe_mode']['source'], 'default') + _, _, explicit = self.start(cli=('--routing-scheduler', 'indexed', + '--catalog-probe-mode', 'representative')) + self.assertEqual(explicit['routing_scheduler'], 'indexed') + self.assertEqual(explicit['catalog_probe_mode'], 'representative') + def test_capability_guard_defaults_precedence_and_validation(self): _, items, config = self.start() diff --git a/tests/test_housekeep.py b/tests/test_housekeep.py new file mode 100644 index 0000000..9d26eff --- /dev/null +++ b/tests/test_housekeep.py @@ -0,0 +1,136 @@ +"""Bounded background maintenance: worker limits, endpoint semaphores and backoff jitter.""" +import sys +from pathlib import Path + +sys.path.insert(0, str(Path(__file__).resolve().parents[1])) # Allow direct execution. + +import json +import os +import tempfile +import time +import unittest +from unittest.mock import Mock, patch + +import converter +from app import perf_metrics +from app.credits import CreditLedger, ModelCatalogCache + + +def model(*names): + return [{"id": name, "supportsToolCall": True, "credits": "x1 credits"} for name in names] + + +class HousekeepConcurrencyTests(unittest.TestCase): + def setUp(self): + temporary = self.enterContext(tempfile.TemporaryDirectory()) + self.enterContext(patch.dict(os.environ, {"CODEBUDDY_AUTH_DIR": temporary})) + self.enterContext(patch.object(perf_metrics, "_COUNTERS", {})) + self.root = Path(temporary) + self.enterContext(patch.object(converter, "_sync_usage", Mock())) + self.enterContext(patch.object(converter, "_sync_model_catalogs", Mock())) + + def build(self, count): + paths = [] + for index in range(count): + path = self.root / f"account-{index}.info" + path.write_text(json.dumps({"account": {"uid": f"maintain-{index}"}, "auth": { + "accessToken": f"token-{index}", "domain": "www.codebuddy.cn", + "expiresAt": (time.time() + 3600) * 1000}}), encoding="utf-8") + paths.append(path) + pool = converter.CredentialPool(paths) + ledger = CreditLedger(self.root / "ledger.json") + pool.set_ledger(ledger) + for entry in pool.entries(): + ledger.bind_identity(entry["id"], entry["account_key"]) + ledger.update_credits(entry["id"], {"credits": 100, "intl": False, "segments": []}) + self.enterContext(patch.dict(converter.CONFIG, { + "cred_pool": pool, "ledger": ledger, "model_cache": None, "aggregate_catalog": None, + "account_catalogs": {}, "model_catalogs": {}, "model_guard": True, "log_path": None, + "control_store": None, "usage_snapshots": None, "housekeep_concurrency": 4})) + return pool, ledger + + def sync_with_delay(self, delay): + def sync(*_args, **_kwargs): + time.sleep(delay) + return None + return sync + + def test_bounded_workers_shorten_a_fixed_rtt_cycle(self): + pool, ledger = self.build(8) + observed = {} + for limit in (1, 4): + for entry in pool._entries: + entry["synced_at"] = 0.0 + with patch.dict(converter.CONFIG, {"housekeep_concurrency": limit}), \ + patch.object(converter, "_sync_credits", side_effect=self.sync_with_delay(0.05)) as sync: + started = time.monotonic() + converter._housekeep_once(pool, ledger) + observed[limit] = time.monotonic() - started + self.assertEqual(sync.call_count, 8) + self.assertGreaterEqual(observed[1], 0.35) # Serial: eight sequential 50 ms calls. + self.assertLess(observed[4], observed[1] * 0.6) + + def test_endpoint_semaphores_are_independent_per_limit(self): + with patch.dict(converter.CONFIG, {"housekeep_concurrency": 2}): + models = converter._endpoint_semaphore("models") + credits = converter._endpoint_semaphore("credits") + usage = converter._endpoint_semaphore("usage") + self.assertIs(models, converter._endpoint_semaphore("models")) # Reused within a limit. + self.assertIsNot(models, credits) + self.assertIsNot(models, usage) + with patch.dict(converter.CONFIG, {"housekeep_concurrency": 1}): + self.assertIsNot(models, converter._endpoint_semaphore("models")) # Rebuilt on a new limit. + + def test_full_cycle_skips_accounts_synced_recently(self): + pool, ledger = self.build(3) + with patch.object(converter, "_sync_credits", return_value=None) as sync: + converter._housekeep_once(pool, ledger) + self.assertEqual(sync.call_count, 3) + sync.reset_mock() + converter._housekeep_once(pool, ledger) + self.assertEqual(sync.call_count, 0) # Recently synced accounts are not repeated. + for entry in pool._entries: + entry["synced_at"] = 0.0 + converter._housekeep_once(pool, ledger) + self.assertEqual(sync.call_count, 3) + + def test_failed_accounts_back_off_with_jitter(self): + pool, _ = self.build(2) + ids = {entry["id"] for entry in pool.entries()} + pool.begin_sync(all_entries=True) + before = time.monotonic() + pool.end_sync(ids, failed=ids) + first = dict(pool._sync_retry) + pool.end_sync(ids, failed=ids) + second = dict(pool._sync_retry) + for cid in ids: + self.assertGreaterEqual(first[cid] - before, 60) + self.assertGreater(second[cid] - first[cid], 40) # Exponential growth. + self.assertLess(first[cid] - before, 60 * 1.25 + 1) # Bounded jitter only. + + def test_manual_scan_runs_on_bounded_workers(self): + pool, _ = self.build(6) + cache = ModelCatalogCache(self.root / "catalog.json", ttl=60) + self.enterContext(patch.dict(converter.CONFIG, { + "model_cache": cache, "housekeep_concurrency": 3, + "aggregate_catalog": converter.AggregateCatalog(ttl=3600)})) + + def slow_fetch(_token, **_kwargs): + time.sleep(0.05) + return {"picker": model("hy3"), "account": model("hy3")} + with patch.object(converter.credits_mod, "fetch_model_scopes", side_effect=slow_fetch): + started = time.monotonic() + converter.start_catalog_scan() + for _ in range(300): + view = converter.catalog_scan_view() + if view["status"] != "running": + break + time.sleep(0.01) + elapsed = time.monotonic() - started + self.assertEqual(view["status"], "completed") + self.assertEqual(view["completed"], 6) + self.assertLess(elapsed, 6 * 0.05) # Six 50 ms fetches on three workers are not serial. + + +if __name__ == "__main__": + unittest.main(verbosity=2) diff --git a/tests/test_identity_sync.py b/tests/test_identity_sync.py index f04cf7a..8a67c2b 100644 --- a/tests/test_identity_sync.py +++ b/tests/test_identity_sync.py @@ -428,8 +428,8 @@ def barrier(*args, **kwargs): patch.object(c.uvicorn, "run", side_effect=barrier), \ patch("sys.stderr", new=io.StringIO()): c.main() - # refresher, housekeeper, lease heartbeat and the server itself. - self.assertEqual(len(observed), 4) + # refresher, reconcile, housekeeper, lease heartbeat and the server itself. + self.assertEqual(len(observed), 5) if __name__ == "__main__": diff --git a/tests/test_login.py b/tests/test_login.py index b565873..3188873 100644 --- a/tests/test_login.py +++ b/tests/test_login.py @@ -84,6 +84,7 @@ def test_login_command_polls_saves_and_does_not_start_server(self): if os.name != "nt": self.assertEqual(stat.S_IMODE(saved.stat().st_mode), 0o600) self.assertIsNone(pool.pick(None, region="cn")) # Never treat international credentials as domestic. + pool.reconcile() # Background reconcile discovers the file the login command just saved. self.assertEqual(pool.pick(None).summary()["uid"], "u1") self.assertIn("账号已保存", self.stdout.getvalue()) self.assert_no_tokens_printed() diff --git a/tests/test_perf_metrics.py b/tests/test_perf_metrics.py new file mode 100644 index 0000000..e37b267 --- /dev/null +++ b/tests/test_perf_metrics.py @@ -0,0 +1,117 @@ +"""Process counters, routing instrumentation and the read-only management view.""" +import sys +from pathlib import Path + +sys.path.insert(0, str(Path(__file__).resolve().parents[1])) # Allow direct execution. + +import json +import os +import tempfile +import time +import unittest +from types import SimpleNamespace +from unittest.mock import Mock, patch + +from fastapi import FastAPI +from fastapi.testclient import TestClient + +import converter +from app import perf_metrics +from app.admin_api import install_admin +from app.control_store import ControlStore + + +class PerfMetricsTests(unittest.TestCase): + def setUp(self): + self.enterContext(patch.object(perf_metrics, "_COUNTERS", {})) + self.enterContext(patch.object(perf_metrics, "_TIMINGS", {})) + + def test_counters_and_timings_accumulate(self): + perf_metrics.inc("candidate_checks") + perf_metrics.inc("candidate_checks", 4) + perf_metrics.observe("route_pick_duration", 0.004) + perf_metrics.observe("route_pick_duration", 0.002) + with perf_metrics.timed("route_pick_duration"): + pass + snapshot = perf_metrics.snapshot() + self.assertEqual(snapshot["counters"], {"candidate_checks": 5}) + timing = snapshot["timings"]["route_pick_duration"] + self.assertEqual(timing["count"], 3) + self.assertAlmostEqual(timing["total_ms"], 6.0, places=1) + self.assertGreaterEqual(timing["max_ms"], 4.0) + self.assertLessEqual(timing["min_ms"], 2.0) + + def test_malformed_sample_is_ignored_and_reset_clears(self): + perf_metrics.observe("pool_lock_wait", "not-a-number") + self.assertEqual(perf_metrics.snapshot()["timings"], {}) + perf_metrics.inc("catalog_fetch_failures") + perf_metrics.reset() + self.assertEqual(perf_metrics.snapshot()["counters"], {}) + + +class RoutingInstrumentationTests(unittest.TestCase): + def setUp(self): + temporary = self.enterContext(tempfile.TemporaryDirectory()) + self.enterContext(patch.dict(os.environ, {"CODEBUDDY_AUTH_DIR": temporary})) + self.enterContext(patch.object(perf_metrics, "_COUNTERS", {})) + self.enterContext(patch.object(perf_metrics, "_TIMINGS", {})) + self.root = Path(temporary) + self.path = self.root / "account.info" + self.path.write_text(json.dumps({"account": {"uid": "instrumented"}, "auth": { + "accessToken": "synthetic-token", "domain": "www.codebuddy.cn", + "expiresAt": (time.time() + 3600) * 1000}}), encoding="utf-8") + self.pool = converter.CredentialPool([self.path]) + self.enterContext(patch.dict(converter.CONFIG, {"cred_pool": self.pool, "log_path": None, + "control_store": None, "model_cache": None, + "account_catalogs": None})) + + def test_pick_records_duration_lock_wait_and_candidate_work(self): + self.assertIsNotNone(self.pool.pick("session", "glm-5.3-flash")) + snapshot = perf_metrics.snapshot() + self.assertGreaterEqual(snapshot["counters"]["candidate_checks"], 1) + self.assertGreaterEqual(snapshot["counters"]["credential_summary_calls"], 1) + self.assertGreaterEqual(snapshot["counters"]["credential_stat_calls"], 1) + self.assertEqual(snapshot["timings"]["pool_lock_wait"]["count"], 1) + self.assertEqual(snapshot["timings"]["route_pick_duration"]["count"], 1) + + def test_capacity_refusal_is_counted(self): + with patch.dict(converter.CONFIG, {"max_inflight_per_account": 1}): + self.pool.pick("session", "glm-5.3-flash") + entry = self.pool._entries[0] + self.pool._capacity.acquire(self.pool._capacity_key(entry), 1, entry["cm"], entry["cm"]._generation) + with self.assertRaises(converter.HTTPException): + self.pool.pick("session", "glm-5.3-flash", with_capacity=True) + self.assertGreaterEqual(perf_metrics.snapshot()["counters"]["capacity_acquire_failures"], 1) + + + +class AdminPerformanceViewTests(unittest.TestCase): + def setUp(self): + root = Path(self.enterContext(tempfile.TemporaryDirectory())) + self.store = ControlStore(root / "control.sqlite3") + self.addCleanup(self.store.close) + self.config = {"api_key": "synthetic-key", "control_store": self.store, "max_images": 16, + "session_path": root / "admin-sessions.json"} + self.audit = Mock() + self.audit.storage.return_value = {} + self.config["audit_store"] = self.audit + self.gateway = SimpleNamespace(managed_auth_dir=Mock(return_value=root)) + self.app = FastAPI() + install_admin(self.app, self.config, self.gateway) + self.client = self.enterContext(TestClient(self.app)) + + def test_requires_admin_login_and_returns_counters_only(self): + self.assertEqual(self.client.get("/admin/performance").status_code, 401) + self.client.post("/admin/session", json={"api_key": "synthetic-key"}, + headers={"Origin": "http://testserver"}) + perf_metrics.inc("route_pick_duration_marker") + result = self.client.get("/admin/performance") + self.assertEqual(result.status_code, 200, result.text) + body = result.json() + self.assertEqual(set(body), {"uptime_seconds", "counters", "timings", "scheduler"}) + self.assertGreaterEqual(body["counters"]["route_pick_duration_marker"], 1) + self.assertNotIn("synthetic-key", result.text) + + +if __name__ == "__main__": + unittest.main(verbosity=2) diff --git a/tests/test_routing_scheduler.py b/tests/test_routing_scheduler.py new file mode 100644 index 0000000..bb295d5 --- /dev/null +++ b/tests/test_routing_scheduler.py @@ -0,0 +1,296 @@ +"""Revision-keyed shards: caching, invalidation, and behavioural parity with the legacy scan.""" +import sys +from pathlib import Path + +sys.path.insert(0, str(Path(__file__).resolve().parents[1])) # Allow direct execution. + +import json +import os +import tempfile +import time +import unittest +from unittest.mock import patch + +import converter +from app import model_policy, perf_metrics +from app.control_store import ControlStore +from app.credits import CreditLedger +from app.routing_scheduler import MISMATCH_LIMIT, RoutingScheduler + +PROFILES = ("cn-cli", "cn-work", "intl-cli", "intl-work") +DOMAINS = {"cn-cli": "www.codebuddy.cn", "cn-work": "www.workbuddy.cn", + "intl-cli": "www.codebuddy.ai", "intl-work": "www.workbuddy.ai"} +MODELS = { + "cn-cli": ["glm-5.3-flash", "deepseek-v4-flash"], + "cn-work": ["auto", "glm-5.3-flash"], + "intl-cli": ["default-model", "glm-5.3-flash"], + "intl-work": ["default-model", "glm-5.3-flash"], +} + + +class SchedulerUnitTests(unittest.TestCase): + def setUp(self): + self.enterContext(patch.object(perf_metrics, "_COUNTERS", {})) + self.enterContext(patch.object(perf_metrics, "_TIMINGS", {})) + self.entries = [{"id": f"/cred-{index}"} for index in range(4)] + + def test_legacy_never_reuses_a_shard(self): + scheduler = RoutingScheduler() + calls = [] + evaluate = lambda entry: calls.append(entry) or entry["id"] != "/cred-0" + first = scheduler.candidates(mode="legacy", revisions=(1,), model="m", region=None, + entries=self.entries, evaluate=evaluate) + second = scheduler.candidates(mode="legacy", revisions=(1,), model="m", region=None, + entries=self.entries, evaluate=evaluate) + self.assertEqual([entry["id"] for entry in first], ["/cred-1", "/cred-2", "/cred-3"]) + self.assertEqual(first, second) + self.assertEqual(len(calls), 8) + self.assertEqual(scheduler.report()["shards"], []) + + def test_indexed_builds_once_per_revision_set(self): + scheduler = RoutingScheduler() + calls = [] + evaluate = lambda entry: calls.append(entry) or True + for revision in ((1, 2, 3, 4), (1, 2, 3, 4), (1, 2, 3, 5)): + scheduler.candidates(mode="indexed", revisions=revision, model="m", region=None, + entries=self.entries, evaluate=evaluate) + self.assertEqual(len(calls), 8) # Two builds, then a reuse. + self.assertEqual(perf_metrics.snapshot()["counters"]["shard_builds"], 2) + self.assertEqual(perf_metrics.snapshot()["counters"]["shard_hits"], 1) + + def test_shard_cache_evicts_least_recently_used(self): + scheduler = RoutingScheduler(limit=2) + evaluate = lambda entry: True + for model in ("m1", "m2"): + scheduler.candidates(mode="indexed", revisions=(1,), model=model, region=None, + entries=self.entries, evaluate=evaluate) + scheduler.candidates(mode="indexed", revisions=(1,), model="m1", region=None, + entries=self.entries, evaluate=evaluate) # Refresh the hot shard. + scheduler.candidates(mode="indexed", revisions=(1,), model="m3", region=None, + entries=self.entries, evaluate=evaluate) # Evicts m2, not m1. + self.assertEqual({row["model"] for row in scheduler.report()["shards"]}, {"m1", "m3"}) + + def test_conservative_rollout_defaults(self): + self.assertEqual(converter.CONFIG["routing_scheduler"], "shadow") + self.assertEqual(converter.CONFIG["catalog_probe_mode"], "account") + self.assertEqual(RoutingScheduler.mode({}), "shadow") + + def test_shadow_returns_legacy_and_records_mismatches(self): + scheduler = RoutingScheduler() + allowed = {"value": 0} + evaluate = lambda entry: allowed["value"] == 0 + scheduler.candidates(mode="indexed", revisions=(1,), model="m", region=None, + entries=self.entries, evaluate=evaluate) + allowed["value"] = 1 # The index now disagrees with a live scan. + result = scheduler.candidates(mode="shadow", revisions=(1,), model="m", region=None, + entries=self.entries, evaluate=evaluate) + self.assertEqual(result, []) + report = scheduler.report() + self.assertEqual(report["mismatches"][0]["legacy"], []) + self.assertEqual(len(report["mismatches"][0]["indexed"]), 4) + self.assertEqual(perf_metrics.snapshot()["counters"]["scheduler_mismatches"], 1) + for index in range(MISMATCH_LIMIT + 5): + scheduler.record("m", None, [], self.entries) + self.assertEqual(len(scheduler.report()["mismatches"]), MISMATCH_LIMIT) + + +class SchedulerPoolTests(unittest.TestCase): + def setUp(self): + temporary = self.enterContext(tempfile.TemporaryDirectory()) + self.enterContext(patch.dict(os.environ, {"CODEBUDDY_AUTH_DIR": temporary})) + self.enterContext(patch.object(perf_metrics, "_COUNTERS", {})) + self.root = Path(temporary) + + def build(self, count, mode="indexed"): + paths = [] + for index in range(count): + profile = PROFILES[index % len(PROFILES)] + path = self.root / f"account-{index}.info" + path.write_text(json.dumps({"account": {"uid": f"shard-{index}"}, "auth": { + "accessToken": f"token-{index}", "domain": DOMAINS[profile], + "expiresAt": (time.time() + 3600) * 1000}}), encoding="utf-8") + paths.append(path) + pool = converter.CredentialPool(paths) + ledger = CreditLedger(self.root / f"ledger-{mode}.json") + pool.set_ledger(ledger) + accounts = {} + for index, entry in enumerate(pool.entries()): + profile = entry["profile"] + ledger.bind_identity(entry["id"], entry["account_key"]) + ledger.update_credits(entry["id"], {"credits": 100 + index, "intl": profile.startswith("intl"), + "segments": [{"remaining": 100, "expires_at": time.time() + 60 + index}]}) + models = [{"id": name, "supportsToolCall": True, + "credits": "x0 credits" if name == "glm-5.3-flash" else "x1 credits"} + for name in MODELS[profile]] + accounts[entry["account_key"]] = {"profile": profile, "models": models, "serves": models} + converter._configured_cache.clear() + converter.invalidate_model_table() + self.enterContext(patch.dict(converter.CONFIG, { + "cred_pool": pool, "ledger": ledger, "control_store": None, "model_cache": None, + "account_catalogs": accounts, "catalog_index": converter._publish_catalog_index(accounts), + "model_catalogs": {}, "model_guard": True, "log_path": None, "routing_scheduler": mode, + "max_inflight_per_account": 0})) + return pool, ledger + + def script(self, pool, ledger): + """Deterministic routing script exercising sticky, cooldown, tried and capacity paths.""" + observed = [] + pick = lambda skey=None, **kw: (pool.pick(skey, "glm-5.3-flash", **kw) or None) + for _ in range(6): + manager = pick() + observed.append(Path(manager.path).name) + for _ in range(4): + observed.append(Path(pick(skey="session").path).name) + first = pool.entries()[0] + pool.note_status(first["cm"], 429, model="glm-5.3-flash") + observed.append(Path(pick().path).name) + observed.append(Path(pick(tried=(first["cm"],)).path).name) + pool.cooldown(pool.entries()[1]["cm"], reason="synthetic") + observed.append(Path(pick().path).name) + ledger.update_credits(pool.entries()[2]["id"], {"credits": 0, "intl": False, "segments": []}) + observed.append(Path(pick().path).name) + with patch.dict(converter.CONFIG, {"max_inflight_per_account": 1}): + retained = [] + while True: + manager = pick(with_capacity=True) if not retained else None + if manager is None: + break + retained.append(Path(manager.path).name) + observed.extend(retained) + try: + pick(with_capacity=True) + except converter.HTTPException as error: + observed.append(str(error.status_code)) + return observed + + def test_modes_choose_the_same_credentials(self): + results = {} + for mode in ("legacy", "indexed", "shadow"): + with patch.object(perf_metrics, "_COUNTERS", {}): + pool, ledger = self.build(12, mode) + results[mode] = self.script(pool, ledger) + self.assertEqual(pool._scheduler.report()["mismatches"], [], mode) + self.assertEqual(results["legacy"], results["indexed"]) + self.assertEqual(results["legacy"], results["shadow"]) + self.assertGreater(len(results["legacy"]), 10) + + def test_second_pick_reuses_the_shard_without_rechecking_eligibility(self): + pool, _ = self.build(8) + pool.pick("session", "glm-5.3-flash") + with patch.object(converter.CredentialPool, "_eligible", wraps=pool._eligible) as eligible: + pool.pick("session", "glm-5.3-flash") + self.assertEqual(eligible.call_count, 0) + self.assertEqual(perf_metrics.snapshot()["counters"]["shard_hits"], 1) + + def test_reload_without_identity_change_keeps_the_shard(self): + pool, _ = self.build(4) + pool.pick("session", "glm-5.3-flash") + revision = pool._shard_revisions() + pool.reload([Path(entry["id"]) for entry in pool.entries()], reset=False) + self.assertEqual(pool._shard_revisions(), revision) + with patch.object(converter.CredentialPool, "_eligible", wraps=pool._eligible) as eligible: + pool.pick("session", "glm-5.3-flash") + self.assertEqual(eligible.call_count, 0) + + def test_identity_change_invalidates_the_shard(self): + pool, _ = self.build(4) + pool.pick("session", "glm-5.3-flash") + revision = pool._shard_revisions() + replaced = pool.entries()[0] + Path(replaced["id"]).write_text(json.dumps({"account": {"uid": "shard-replacement"}, "auth": { + "accessToken": "token-replacement", "domain": DOMAINS[replaced["profile"]], + "expiresAt": (time.time() + 3600) * 1000}}), encoding="utf-8") + pool.reload([Path(replaced["id"])], reset=False) + self.assertNotEqual(pool._shard_revisions(), revision) + with patch.object(converter.CredentialPool, "_eligible", wraps=pool._eligible) as eligible: + pool.pick("session", "glm-5.3-flash") + self.assertGreater(eligible.call_count, 0) + + def pick_alias(self, pool, source): + """Select through one request-scoped alias, exactly as the resolver leaves the context.""" + rule = model_policy.rule_for(converter.CONFIG, source) + token = model_policy._request_policy.set({"source": source, "upstream": rule["upstream_id"], + "public_model": rule["public_id"], "rule": rule}) + try: + manager = pool.pick(None, rule["upstream_id"]) + finally: + model_policy._request_policy.reset(token) + return pool.entry_for(manager) + + def test_aliases_with_different_bindings_do_not_share_a_shard(self): + pool, _ = self.build(4) + store = ControlStore(self.root / "control.sqlite3") + self.addCleanup(store.close) + entries = pool.entries() + for source, identity in (("alias-a", entries[0]["account_key"]), + ("alias-b", entries[1]["account_key"])): + store.update_model(source, {"public_id": f"{source}-out", "upstream_id": "glm-5.3-flash", + "custom": True, "credential_ids": [identity]}, + store.revision, known_models=()) + with patch.dict(converter.CONFIG, {"control_store": store}): + self.assertEqual(self.pick_alias(pool, "alias-a")["id"], entries[0]["id"]) + self.assertEqual(self.pick_alias(pool, "alias-b")["id"], entries[1]["id"]) + self.assertEqual(len(pool._scheduler.report()["shards"]), 2) + + def test_balance_change_invalidates_the_shard(self): + pool, ledger = self.build(4) + pool.pick("session", "glm-5.3-flash") + revision = pool._shard_revisions() + ledger.update_credits(pool.entries()[0]["id"], {"credits": 0, "intl": False, "segments": []}) + self.assertNotEqual(pool._shard_revisions(), revision) + + def test_balance_value_churn_keeps_the_shard(self): + pool, ledger = self.build(4) + pool.pick("session", "glm-5.3-flash") + revision = pool._shard_revisions() + churned = pool.entries()[0] + ledger.update_credits(churned["id"], {"credits": 99, "intl": False, + "segments": [{"remaining": 99, "expires_at": time.time() + 60}]}) + self.assertEqual(pool._shard_revisions(), revision) # Positive-to-positive is not a flip. + + def test_intl_first_balance_row_moves_routing_revision(self): + ledger = CreditLedger(self.root / "intl-ledger.json") + ledger.bind_identity("intl-cred", "intl-key") + before = ledger.routing_revision + ledger.update_credits("intl-cred", {"credits": 5, "intl": True, "segments": []}) + self.assertNotEqual(ledger.routing_revision, before) + churn = ledger.routing_revision + ledger.update_credits("intl-cred", {"credits": 4, "intl": True, "segments": []}) + self.assertEqual(ledger.routing_revision, churn) + + def test_token_rotation_keeps_catalog_clean_and_shard_hot(self): + pool, _ = self.build(4) + pool.pick("session", "glm-5.3-flash") + revision = pool._shard_revisions() + live = pool.entry_by_id(pool.entries()[0]["id"]) + live["catalog_dirty"] = False + Path(live["id"]).write_text(json.dumps({"account": {"uid": "shard-0"}, "auth": { + "accessToken": "token-0-rotated", "domain": DOMAINS[live["profile"]], + "expiresAt": (time.time() + 7200) * 1000}}), encoding="utf-8") + pool.reload([Path(live["id"])], reset=False) + self.assertFalse(pool.entry_by_id(live["id"])["catalog_dirty"]) + self.assertEqual(pool._shard_revisions(), revision) + + def test_background_reload_keeps_a_pending_catalog_refresh_marker(self): + pool, _ = self.build(4) + live = pool.entry_by_id(pool.entries()[0]["id"]) + live["catalog_dirty"] = True # Marked by an import or a representative probe. + Path(live["id"]).write_text(json.dumps({"account": {"uid": "shard-0"}, "auth": { + "accessToken": "token-0-rotated", "domain": DOMAINS[live["profile"]], + "expiresAt": (time.time() + 7200) * 1000}}), encoding="utf-8") + pool.reload([Path(live["id"])], reset=False) # A token refresh races the sync. + self.assertTrue(pool.entry_by_id(live["id"])["catalog_dirty"]) + + def test_stale_generation_account_stays_pickable_until_reconcile(self): + pool, _ = self.build(4) + live = pool.entry_by_id(pool.entries()[0]["id"]) + others = tuple(entry["cm"] for entry in pool.entries()[1:]) + self.assertIs(pool.pick("fresh", "glm-5.3-flash", tried=others), live["cm"]) + live["cm"].invalidate() # The bind is stale, but headers_for rechecks before sending. + self.assertIs(pool.pick("stale", "glm-5.3-flash", tried=others), live["cm"]) + pool.reload([Path(live["id"])], reset=False) + self.assertEqual(pool.entry_by_id(live["id"])["generation"], live["cm"]._generation) + + +if __name__ == "__main__": + unittest.main(verbosity=2) diff --git a/tests/test_routing_snapshot.py b/tests/test_routing_snapshot.py new file mode 100644 index 0000000..0cb8eec --- /dev/null +++ b/tests/test_routing_snapshot.py @@ -0,0 +1,109 @@ +"""One policy revision per routing decision, indexed catalogs and a revision-keyed model table.""" +import sys +from pathlib import Path + +sys.path.insert(0, str(Path(__file__).resolve().parents[1])) # Allow direct execution. + +import json +import os +import tempfile +import time +import unittest +from unittest.mock import patch + +import converter +from app import model_policy, perf_metrics +from app.control_store import ControlStore +from app.credits import CreditLedger + +MODELS = [{"id": "glm-5.3-flash", "supportsToolCall": True, "credits": "x0 credits"}, + {"id": "deepseek-v4-flash", "supportsToolCall": True, "credits": "x1 credits"}] + + +class RoutingSnapshotTests(unittest.TestCase): + def setUp(self): + temporary = self.enterContext(tempfile.TemporaryDirectory()) + self.enterContext(patch.dict(os.environ, {"CODEBUDDY_AUTH_DIR": temporary})) + self.enterContext(patch.object(perf_metrics, "_COUNTERS", {})) + self.enterContext(patch.object(perf_metrics, "_TIMINGS", {})) + self.root = Path(temporary) + self.store = ControlStore(self.root / "control.sqlite3") + self.addCleanup(self.store.close) + self.path = self.write_credential("account.info", "synthetic") + self.pool = converter.CredentialPool([self.path]) + self.ledger = CreditLedger(self.root / "ledger.json") + self.pool.set_ledger(self.ledger) + entry = self.pool.entries()[0] + self.identity = entry["account_key"] + self.ledger.bind_identity(entry["id"], self.identity) + self.ledger.update_credits(entry["id"], {"credits": 100, "intl": False, "segments": []}) + self.accounts = {self.identity: {"profile": "cn-cli", + "models": [dict(item) for item in MODELS], + "serves": [dict(item) for item in MODELS]}} + converter.invalidate_model_table() + converter._configured_cache.clear() + self.enterContext(patch.dict(converter.CONFIG, { + "cred_pool": self.pool, "ledger": self.ledger, "control_store": self.store, + "account_catalogs": self.accounts, "catalog_index": converter._publish_catalog_index(self.accounts), + "model_catalogs": {}, "model_cache": None, "model_guard": True, "log_path": None, + "max_inflight_per_account": 0})) + + def write_credential(self, name, uid): + path = self.root / name + path.write_text(json.dumps({"account": {"uid": uid}, "auth": { + "accessToken": "synthetic-token", "domain": "www.codebuddy.cn", + "expiresAt": (time.time() + 3600) * 1000}}), encoding="utf-8") + return path + + def test_one_policy_snapshot_per_decision(self): + self.assertIsNotNone(self.pool.headers_for(None, "glm-5.3-flash")) + self.assertEqual(perf_metrics.snapshot()["counters"]["policy_snapshot_calls"], 1) + self.assertIsNotNone(self.pool.pick("session", "glm-5.3-flash")) + self.assertEqual(perf_metrics.snapshot()["counters"]["policy_snapshot_calls"], 2) + + def test_policy_view_reuses_one_revision_and_sees_new_publishes(self): + policy = model_policy.view(converter.CONFIG) + entry = self.pool.entries()[0] + self.assertEqual(policy.revision, self.store.revision) + self.assertTrue(model_policy.credential_enabled(converter.CONFIG, entry, policy)) + self.store.set_credential(self.identity, False) + self.assertTrue(model_policy.credential_enabled(converter.CONFIG, entry, policy)) # Old revision kept. + self.assertFalse(model_policy.view(converter.CONFIG).credential(entry)["enabled"]) + self.assertIsNone(self.pool.pick("session", "glm-5.3-flash")) + + def test_revision_reads_do_not_copy_the_snapshot(self): + self.assertEqual(model_policy.revision(converter.CONFIG), self.store.revision) + self.assertEqual(perf_metrics.snapshot()["counters"].get("policy_snapshot_calls", 0), 0) + self.assertEqual(self.store.revision, self.store.snapshot()["revision"]) + + def test_model_table_is_reused_until_a_revision_moves(self): + first = converter.current_models() + self.assertEqual(first[:2], ["glm-5.3-flash", "deepseek-v4-flash"]) + self.assertEqual(converter.current_models(), first) + counters = perf_metrics.snapshot()["counters"] + self.assertEqual(counters["model_table_builds"], 1) + self.assertEqual(counters["model_table_hits"], 1) + entry = self.pool.entries()[0] + self.ledger.update_credits(entry["id"], {"credits": 0, "intl": False, "segments": []}) + self.assertEqual(converter.current_models()[:1], ["glm-5.3-flash"]) # Zero balance keeps free models only. + self.assertEqual(perf_metrics.snapshot()["counters"]["model_table_builds"], 2) + + def test_catalog_republish_rebuilds_the_model_table(self): + converter.current_models() + accounts = {self.identity: {"profile": "cn-cli", "models": [dict(MODELS[1])], + "serves": [dict(MODELS[1])]}} + converter.CONFIG["account_catalogs"] = accounts + converter.CONFIG["catalog_index"] = converter._publish_catalog_index(accounts) + self.assertEqual(converter.current_models()[:1], ["deepseek-v4-flash"]) + + def test_configured_profiles_are_cached_between_decisions(self): + converter._configured_profiles(None) + with patch.object(self.pool, "entries", wraps=self.pool.entries) as entries: + self.assertEqual(set(converter._configured_profiles(None)), {"cn-cli"}) + self.assertEqual(entries.call_count, 0) + self.store.set_credential(self.identity, False) + self.assertEqual(set(converter._configured_profiles(None)), set()) + + +if __name__ == "__main__": + unittest.main(verbosity=2) diff --git a/tests/test_sqlite_state.py b/tests/test_sqlite_state.py index 0e6d607..fc2a0e2 100644 --- a/tests/test_sqlite_state.py +++ b/tests/test_sqlite_state.py @@ -16,6 +16,7 @@ from unittest.mock import patch from app.admin_auth import AdminAuth, SessionStoreError +from app.catalog_probe import AggregateCatalog from app.control_store import ControlStore from app.credential_cooldowns import CredentialCooldowns from app.credits import CreditLedger, ModelCatalogCache @@ -69,6 +70,10 @@ def create_legacy(self): usage = UsageSnapshots(self.root / LEGACY_FILES["usage"]) usage.store("fixture.info", IDENTITY, "domestic", {"total_credits": 12, "requests": 2, "by_day": {"2026-09-22": {"model": 12}}}, partial=True) + probe = self.root / LEGACY_FILES["catalog_probe"] + probe.write_text(json.dumps({"version": 1, "profiles": {"cn-cli": { + "source": IDENTITY, "source_generation": 1, "fetched_at": self.now, + "retry_at": self.now + 300, "failures": 0, "status": "ok", "models": ["model"]}}})) trial = TrialLedger(self.root / LEGACY_FILES["trial"]) self.assertTrue(trial.begin(IDENTITY)) trial.finish(IDENTITY, {"ok": True, "already": False, "code": 0, "status": 200}) @@ -90,6 +95,7 @@ def test_all_legacy_sources_migrate_once_and_runtime_never_updates_json(self): self.assertGreater(cooldowns.credential_until(IDENTITY, "cn-work"), self.now) self.assertGreater(cooldowns.model_until(IDENTITY, "cn-work", "model"), self.now) self.assertTrue(UsageSnapshots(store=self.state).accounts()["fixture.info"]["partial"]) + self.assertEqual(AggregateCatalog(store=self.state).row("cn-cli")["source"], IDENTITY) trial = TrialLedger(store=self.state) self.assertFalse(trial.begin(IDENTITY, self.now + 3 * 86400)) auth = AdminAuth({"api_key": FIXED_KEY, "state_store": self.state}) @@ -107,7 +113,7 @@ def test_all_legacy_sources_migrate_once_and_runtime_never_updates_json(self): auth.reconcile() self.assertNotIn(sid, auth.sessions) self.assertEqual(before, {name: (self.root / file).read_bytes() for name, file in LEGACY_FILES.items()}) - self.assertEqual(self.control._db.execute("SELECT count(*) FROM state_imports").fetchone()[0], 7) + self.assertEqual(self.control._db.execute("SELECT count(*) FROM state_imports").fetchone()[0], 8) def test_absent_legacy_file_cannot_resurrect_later(self): self.state.migrate(self.root) diff --git a/web/src/api.ts b/web/src/api.ts index c94d7f2..7c1cd74 100644 --- a/web/src/api.ts +++ b/web/src/api.ts @@ -280,6 +280,64 @@ export function modelResponse(value: unknown): { : undefined, }; } +export type CatalogProfileVersions = { + accounts: number; + distinct: number; + versions: { models: number; accounts: number }[]; +}; +export type CatalogScan = { + status: "idle" | "running" | "completed" | "failed"; + total: number; + completed: number; + success: number; + failed: number; + differences: number; + started_at: number | null; + finished_at: number | null; + error: string | null; + profiles: Record; + accounts: RecordValue[]; +}; +export function catalogScanResponse(value: unknown): CatalogScan { + const data = object(value, "目录扫描"); + const count = (field: string) => { + const parsed = number(data[field]); + if (parsed === null || parsed < 0) throw new Error(`目录扫描缺少 ${field}`); + return parsed; + }; + const status = data.status; + if (status !== "idle" && status !== "running" && status !== "completed" && status !== "failed") + throw new Error("目录扫描状态无效"); + return { + status, + total: count("total"), + completed: count("completed"), + success: count("success"), + failed: count("failed"), + differences: count("differences"), + started_at: number(data.started_at), + finished_at: number(data.finished_at), + error: typeof data.error === "string" && data.error ? data.error : null, + profiles: parseCatalogProfiles(data.profiles), + accounts: Array.isArray(data.accounts) ? list(data.accounts, "目录扫描账号") : [], + }; +} +function parseCatalogProfiles(value: unknown): Record { + if (!value || typeof value !== "object") return {}; + const profiles: Record = {}; + for (const [name, row] of Object.entries(value as Record)) { + const item = object(row, "目录版本"); + profiles[name] = { + accounts: number(item.accounts) ?? 0, + distinct: number(item.distinct) ?? 0, + versions: list(item.versions, "目录版本列表").map((entry) => { + const version = object(entry, "目录版本条目"); + return { models: number(version.models) ?? 0, accounts: number(version.accounts) ?? 0 }; + }), + }; + } + return profiles; +} export type Credential = RecordValue & { id: string; name: string | null }; export function credentialResponse(value: unknown): Credential[] { return list(object(value).credentials, "凭证").map((c) => { diff --git a/web/src/pages/Models.tsx b/web/src/pages/Models.tsx index ce51804..1c7e19e 100644 --- a/web/src/pages/Models.tsx +++ b/web/src/pages/Models.tsx @@ -1,12 +1,14 @@ -import { useEffect, useId, useRef, useState } from "react"; +import { useCallback, useEffect, useId, useRef, useState } from "react"; import { api, + catalogScanResponse, credentialResponse, errorMessage, list, modelResponse, object, useResource, + type CatalogScan, type Credential, type ModelRule, type RecordValue, @@ -505,6 +507,50 @@ export function Models() { const [busy, setBusy] = useState(false); const [error, setError] = useState(null); const [guardNotice, setGuardNotice] = useState(null); + const [scan, setScan] = useState(null); + const [scanBusy, setScanBusy] = useState(false); + const loadScan = useCallback(() => { + void api + .get("/models/scan") + .then((response) => { + const next = catalogScanResponse(response.data); + setScan(next); + // A scan that just finished may have refreshed the catalog list. + if (next.status === "completed" || next.status === "failed") resource.reload(); + }) + .catch(() => { + // 被动状态轮询失败不升级为页面错误；提交扫描的失败才提示。 + }); + // eslint-disable-next-line react-hooks/exhaustive-deps + }, []); + useEffect(() => { + loadScan(); // An in-progress scan survives page reloads. + }, [loadScan]); + useEffect(() => { + if (scan?.status !== "running") return; + const timer = setInterval(loadScan, 2000); + return () => clearInterval(timer); + }, [scan?.status, loadScan]); + const scanError = scan?.status === "failed" ? scan.error : null; + const scanVersions = scan + ? Object.entries(scan.profiles) + .filter(([, row]) => row.distinct > 1) + .map(([name, row]) => `${name}：${row.distinct} 组目录版本`) + .join("；") + : ""; + const startScan = () => { + if (scanBusy) return; + setScanBusy(true); + setError(null); + void api + .post("/models/scan") + .then((response) => { + setScan(catalogScanResponse(response.data)); + resource.reload(); + }) + .catch((err: unknown) => setError(errorMessage(err))) + .finally(() => setScanBusy(false)); + }; const toggleGuard = (key: "model_guard" | "model_capability_guard", next: boolean) => { if (!resource.data || busy) return; setBusy(true); @@ -549,6 +595,16 @@ export function Models() { 刷新目录 +