From fdc4b23d49ec63ab6e5a3845056a315a6afa3681 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E5=A2=A8=E8=8F=8A?= <277378677+maiphucgiang@users.noreply.github.com> Date: Sun, 18 Oct 2026 12:24:40 +0800 Subject: [PATCH 01/13] Add offline routing benchmark and process counters --- app/admin_api.py | 6 ++ app/model_policy.py | 6 +- app/perf_metrics.py | 64 ++++++++++++++ benchmarks/bench_routing.py | 164 ++++++++++++++++++++++++++++++++++++ converter.py | 17 +++- tests/test_perf_metrics.py | 117 +++++++++++++++++++++++++ 6 files changed, 371 insertions(+), 3 deletions(-) create mode 100644 app/perf_metrics.py create mode 100644 benchmarks/bench_routing.py create mode 100644 tests/test_perf_metrics.py diff --git a/app/admin_api.py b/app/admin_api.py index a1f064a..bd5622b 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 @@ -687,4 +688,9 @@ 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.""" + return JSONResponse(perf_metrics.snapshot()) + return auth diff --git a/app/model_policy.py b/app/model_policy.py index 89e6d7b..33d5370 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) @@ -24,7 +24,9 @@ async def __call__(self, scope, receive, send): def snapshot(config): 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} + return store.snapshot() def credential_enabled(config, entry): diff --git a/app/perf_metrics.py b/app/perf_metrics.py new file mode 100644 index 0000000..e650091 --- /dev/null +++ b/app/perf_metrics.py @@ -0,0 +1,64 @@ +"""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; a single dict update stays cheap enough for the request path.""" + _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: + _TIMINGS[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/benchmarks/bench_routing.py b/benchmarks/bench_routing.py new file mode 100644 index 0000000..131b0c4 --- /dev/null +++ b/benchmarks/bench_routing.py @@ -0,0 +1,164 @@ +#!/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, + "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)) + 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)) + finally: + temporary.cleanup() + converter.CONFIG.update({"cred_pool": None, "ledger": None, "account_catalogs": 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..a9c594b 100644 --- a/converter.py +++ b/converter.py @@ -77,6 +77,7 @@ def desensitize_body(body, roles=("system",), desensitize_harness_user=False, 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,6 +288,7 @@ 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 @@ -390,6 +392,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 {} @@ -1088,6 +1091,7 @@ def _evict_sticky(self): def _candidates(self, model: str | None, *, region=None, tried=()) -> list[dict]: """Exclude tried credentials and rank candidates by zero rate and credit expiry.""" + perf_metrics.inc("candidate_checks", len(self._entries)) tried = set(tried) healthy = [entry for entry in self._entries if entry["cm"] not in tried and self._healthy(entry) @@ -1114,7 +1118,11 @@ def pick(self, skey: str | None, model: str | None = None, *, region=None, tried=(), with_capacity=False, requirements=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: + 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) if not candidates: @@ -1135,6 +1143,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,6 +1161,9 @@ 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): @@ -1187,6 +1199,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 @@ -1655,6 +1668,7 @@ def _sync_model_catalogs(pool, ledger, refs, failed): key = catalog_cache_key(profile, identity) if cache.fresh(key) and not entry.get("catalog_dirty"): continue + perf_metrics.inc("catalog_fetch_count") try: scopes = credits_mod.fetch_model_scopes( _bearer_token(headers), domain=headers.get("X-Domain", ""), @@ -1671,6 +1685,7 @@ def publish(): else: 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) _publish_model_cache() diff --git a/tests/test_perf_metrics.py b/tests/test_perf_metrics.py new file mode 100644 index 0000000..252fbb0 --- /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"}) + self.assertGreaterEqual(body["counters"]["route_pick_duration_marker"], 1) + self.assertNotIn("synthetic-key", result.text) + + +if __name__ == "__main__": + unittest.main(verbosity=2) From b9272af4fde6bf277db040be133384e212f5162b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E5=A2=A8=E8=8F=8A?= <277378677+maiphucgiang@users.noreply.github.com> Date: Sun, 18 Oct 2026 12:26:41 +0800 Subject: [PATCH 02/13] Move credential file scanning out of the request path --- converter.py | 93 +++++++++++++++++++++++++++----- tests/test_credential_runtime.py | 24 +++++++++ tests/test_identity_sync.py | 4 +- 3 files changed, 106 insertions(+), 15 deletions(-) diff --git a/converter.py b/converter.py index a9c594b..5c3549f 100644 --- a/converter.py +++ b/converter.py @@ -563,7 +563,11 @@ 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]] = {} # Explicit paths must survive rejection or eviction so a repaired file re-enters. self._configured: list[str] = [] self._ignored_duplicates: set[str] = set() @@ -671,9 +675,32 @@ 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.""" + self._index_id = {entry["id"]: entry for entry in self._entries} + self._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) + self._index_identity = 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) @@ -728,7 +755,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): @@ -752,6 +779,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.""" @@ -777,6 +805,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]: @@ -956,11 +986,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 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) -> bool: return model_policy.credential_enabled(CONFIG, e) 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: @@ -1117,7 +1171,7 @@ def _capacity_key(entry): def pick(self, skey: str | None, model: str | None = None, *, region=None, tried=(), with_capacity=False, requirements=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. + # The candidate list comes from the indexes; file changes are applied by background reconcile. started = time.perf_counter() waiting = time.perf_counter() self._lock.acquire() @@ -1186,8 +1240,10 @@ 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) + 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) and self._eligible(entry, model, region=region, profile=profile) and self._model_healthy(entry, model)): if requirements is not None: @@ -1454,6 +1510,22 @@ 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 @@ -2737,8 +2809,6 @@ def _positive(entry): def current_models(region: str | None = None) -> list[str]: """Merge models eligible for current accounts without region-specific client URLs.""" 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 @@ -2784,8 +2854,6 @@ 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.""" pool = CONFIG.get("cred_pool") - if pool is not None: - pool._rescan() details: dict[str, dict] = {} declarations = {} for name in current_models(region): @@ -2965,9 +3033,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): @@ -4625,6 +4690,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/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_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__": From 2e5240a9f86abf7aa7f7f164798e3fbaa1ce5b4c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E5=A2=A8=E8=8F=8A?= <277378677+maiphucgiang@users.noreply.github.com> Date: Sun, 18 Oct 2026 12:31:42 +0800 Subject: [PATCH 03/13] Reuse one policy revision and indexed catalogs per routing decision --- app/catalog_index.py | 77 +++++++++++++++ app/control_store.py | 5 + app/credits.py | 7 ++ app/model_policy.py | 74 +++++++++++--- converter.py | 171 ++++++++++++++++++++++++--------- tests/test_routing_snapshot.py | 109 +++++++++++++++++++++ 6 files changed, 384 insertions(+), 59 deletions(-) create mode 100644 app/catalog_index.py create mode 100644 tests/test_routing_snapshot.py 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/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..9fffc76 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,9 @@ 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) + + class CreditLedger: """Persist per-credential check-in and balance snapshots for expiry-aware scheduling.""" @@ -802,6 +806,8 @@ 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) self._load() def _load(self): @@ -816,6 +822,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 diff --git a/app/model_policy.py b/app/model_policy.py index 33d5370..afa2d79 100644 --- a/app/model_policy.py +++ b/app/model_policy.py @@ -23,33 +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") 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): @@ -58,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"): @@ -67,11 +109,11 @@ 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 resolve(config, public_model): @@ -107,11 +149,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 diff --git a/converter.py b/converter.py index 5c3549f..6fdb798 100644 --- a/converter.py +++ b/converter.py @@ -6,6 +6,7 @@ import argparse import asyncio import hashlib +import itertools import json import math import os @@ -72,6 +73,8 @@ 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.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 @@ -546,6 +549,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.""" @@ -568,6 +575,7 @@ def __init__(self, paths: list[Path] | None = None, scan: bool = False, 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. # Explicit paths must survive rejection or eviction so a repaired file re-enters. self._configured: list[str] = [] self._ignored_duplicates: set[str] = set() @@ -679,6 +687,7 @@ def reload(self, paths: list[Path], *, reset: bool = True): def _rebuild_index(self): """Rebuild request-side lookups after the entry list changed; never called per request.""" + self.revision = next(_POOL_REVISION) self._index_id = {entry["id"]: entry for entry in self._entries} self._index_cm = {entry["cm"]: entry for entry in self._entries} identity: dict[str, list[dict]] = {} @@ -1006,8 +1015,8 @@ def reconcile(self, previous=None): self._rescan() return current - def _healthy(self, e: dict) -> bool: - return model_policy.credential_enabled(CONFIG, e) and time.time() >= e["fail_until"] + 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): @@ -1052,25 +1061,22 @@ 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: + # A changed file invalidates this account until it is re-read; headers_for reloads and + # rechecks the selected credential before any request is sent with it. + if entry["generation"] != entry["cm"]._generation: 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 @@ -1079,7 +1085,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): @@ -1143,13 +1149,14 @@ def _evict_sticky(self): else: break - def _candidates(self, model: str | None, *, region=None, tried=()) -> list[dict]: + 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.""" perf_metrics.inc("candidate_checks", len(self._entries)) 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) + and self._healthy(entry, policy) + and self._eligible(entry, model, region=region, policy=policy) + and self._model_healthy(entry, model) and self._model_servable(entry, model)] if not healthy: return [] @@ -1169,16 +1176,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.""" # 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) @@ -1220,12 +1228,14 @@ def pick(self, skey: str | None, model: str | None = None, *, region=None, 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 @@ -1244,8 +1254,9 @@ def headers_for(self, skey: str | None, model: str | None = None, *, region=None 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) - and self._eligible(entry, model, region=region, profile=profile) and self._model_healthy(entry, model)): + 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: @@ -1726,6 +1737,7 @@ def _publish_model_cache(): models.extend(account["models"]) CONFIG["model_catalogs"] = catalogs CONFIG["models_remote"], CONFIG["models_intl"] = catalogs["cn-cli"], catalogs["intl-cli"] + _publish_catalog_index(accounts) invalidate_model_table() @@ -2624,6 +2636,8 @@ 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: @@ -2632,10 +2646,45 @@ def invalidate_model_table() -> None: _model_table_cache = {} +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 []: @@ -2659,11 +2708,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"] @@ -2678,16 +2736,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: @@ -2701,6 +2749,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: @@ -2709,11 +2759,13 @@ 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 _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 []) @@ -2756,21 +2808,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 @@ -2806,8 +2860,35 @@ def _positive(entry): for entry in pool.entries() if model_policy.credential_enabled(CONFIG, entry)) +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)) + + +def _model_table(name, region, build): + """Reuse a published model table until a revision moves or the short TTL lapses.""" + 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") with pool._lock if pool is not None else nullcontext(): configured = _configured_profiles(region) @@ -2853,6 +2934,10 @@ 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") details: dict[str, dict] = {} declarations = {} 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) From 8184dfce89b66cfc2fd6ef148b663800addcfb3c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E5=A2=A8=E8=8F=8A?= <277378677+maiphucgiang@users.noreply.github.com> Date: Sun, 18 Oct 2026 12:34:50 +0800 Subject: [PATCH 04/13] Add revision-keyed routing shards with legacy and shadow modes --- app/admin_api.py | 7 +- app/routing_scheduler.py | 69 ++++++++++++ app/settings.py | 3 + benchmarks/bench_routing.py | 4 +- converter.py | 23 +++- tests/test_perf_metrics.py | 2 +- tests/test_routing_scheduler.py | 185 ++++++++++++++++++++++++++++++++ 7 files changed, 286 insertions(+), 7 deletions(-) create mode 100644 app/routing_scheduler.py create mode 100644 tests/test_routing_scheduler.py diff --git a/app/admin_api.py b/app/admin_api.py index bd5622b..4795f90 100644 --- a/app/admin_api.py +++ b/app/admin_api.py @@ -691,6 +691,11 @@ def build_dashboard(): @route("GET", "/admin/performance") async def performance_get(request): """Read-only process counters; no routing state is exposed or mutated.""" - return JSONResponse(perf_metrics.snapshot()) + 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/routing_scheduler.py b/app/routing_scheduler.py new file mode 100644 index 0000000..5df2a62 --- /dev/null +++ b/app/routing_scheduler.py @@ -0,0 +1,69 @@ +"""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") + + def __init__(self): + self._shards: dict[tuple, tuple] = {} + self.mismatches: list[dict] = [] + + @staticmethod + def mode(config) -> str: + """Resolve the configured mode, falling back to the safe indexed path.""" + value = config.get("routing_scheduler", "indexed") + return value if value in MODES else "indexed" + + 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)) + # Keep only the current revision set; a moved revision never reuses stale entries. + self._shards = {key: shard} + perf_metrics.inc("shard_builds") + perf_metrics.inc("shard_size_rebuilt", len(shard)) + else: + 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.""" + return {"shards": [ + {"model": key[0], "region": key[1], "size": len(value)} + for key, value in self._shards.items()], + "mismatches": list(self.mismatches)} diff --git a/app/settings.py b/app/settings.py index d3d19c6..e0763c3 100644 --- a/app/settings.py +++ b/app/settings.py @@ -97,6 +97,9 @@ 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("indexed", "string", "路由调度模式", + env="CODEBUDDY2API_ROUTING_SCHEDULER", + choices=["legacy", "shadow", "indexed"]), "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/benchmarks/bench_routing.py b/benchmarks/bench_routing.py index 131b0c4..84c8e89 100644 --- a/benchmarks/bench_routing.py +++ b/benchmarks/bench_routing.py @@ -73,6 +73,7 @@ def build(converter, credits_mod, count: int, root: Path): merged = catalogs[account["profile"]] = catalogs[account["profile"]] or [] merged.extend(account["models"]) config.update({"cred_pool": pool, "ledger": ledger, "account_catalogs": accounts, + "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, @@ -129,7 +130,8 @@ def saturated(): finally: temporary.cleanup() converter.CONFIG.update({"cred_pool": None, "ledger": None, "account_catalogs": None, - "model_catalogs": {}, "models_remote": None, "models_intl": None}) + "catalog_index": None, "model_catalogs": {}, + "models_remote": None, "models_intl": None}) return rows diff --git a/converter.py b/converter.py index 6fdb798..33499a0 100644 --- a/converter.py +++ b/converter.py @@ -74,6 +74,7 @@ def desensitize_body(body, roles=("system",), desensitize_harness_user=False, 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_index import account_scope as _account_scope from app.inference_auth import require_api_key from app.admin_auth import SessionStoreError @@ -576,6 +577,7 @@ def __init__(self, paths: list[Path] | None = None, scan: bool = False, 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() @@ -688,6 +690,7 @@ def reload(self, paths: list[Path], *, reset: bool = True): def _rebuild_index(self): """Rebuild request-side lookups after the entry list changed; never called per request.""" self.revision = next(_POOL_REVISION) + self._scheduler.clear() self._index_id = {entry["id"]: entry for entry in self._entries} self._index_cm = {entry["cm"]: entry for entry in self._entries} identity: dict[str, list[dict]] = {} @@ -1149,13 +1152,21 @@ def _evict_sticky(self): else: break + 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, "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.""" - perf_metrics.inc("candidate_checks", len(self._entries)) tried = set(tried) - healthy = [entry for entry in self._entries if entry["cm"] not in tried + eligible = self._scheduler.candidates( + mode=RoutingScheduler.mode(CONFIG), revisions=self._shard_revisions(policy), + 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._eligible(entry, model, region=region, policy=policy) and self._model_healthy(entry, model) and self._model_servable(entry, model)] if not healthy: @@ -2080,6 +2091,7 @@ 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": "indexed", # legacy | shadow | indexed eligibility shards "max_images": 16, "image_policy": "truncate", "max_request_bytes": 32 * 1024 * 1024, "log_body_limit": 65536, "max_inbound_bytes": 64 * 1024 * 1024, @@ -4682,6 +4694,9 @@ 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("--routing-scheduler", choices=("legacy", "shadow", "indexed"), + default=os.environ.get("CODEBUDDY2API_ROUTING_SCHEDULER", "indexed"), + help="路由调度：indexed 复用按模型分片（默认）；shadow 记录新旧差异并仍返回旧结果；legacy 每次全量判断") ap.add_argument("--model-capability-guard", type=_boolean_arg, nargs="?", const=True, default=os.environ.get("CODEBUDDY2API_MODEL_CAPABILITY_GUARD", "true"), help="按账号模型声明预检图片、工具、思考和输出上限；false 仅关闭新增能力预检") @@ -4719,7 +4734,7 @@ 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"): CONFIG[key] = getattr(args, key) CONFIG["api_key"] = args.api_key CONFIG["desensitize"] = args.desensitize diff --git a/tests/test_perf_metrics.py b/tests/test_perf_metrics.py index 252fbb0..e37b267 100644 --- a/tests/test_perf_metrics.py +++ b/tests/test_perf_metrics.py @@ -108,7 +108,7 @@ def test_requires_admin_login_and_returns_counters_only(self): 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"}) + self.assertEqual(set(body), {"uptime_seconds", "counters", "timings", "scheduler"}) self.assertGreaterEqual(body["counters"]["route_pick_duration_marker"], 1) self.assertNotIn("synthetic-key", result.text) diff --git a/tests/test_routing_scheduler.py b/tests/test_routing_scheduler.py new file mode 100644 index 0000000..1797274 --- /dev/null +++ b/tests/test_routing_scheduler.py @@ -0,0 +1,185 @@ +"""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 perf_metrics +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_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.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_credential_reload_invalidates_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.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 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) + + +if __name__ == "__main__": + unittest.main(verbosity=2) From a395d5de30ef006f19297a2f974e91e69ea01e85 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E5=A2=A8=E8=8F=8A?= <277378677+maiphucgiang@users.noreply.github.com> Date: Sun, 18 Oct 2026 12:41:50 +0800 Subject: [PATCH 05/13] Probe one representative catalog per profile and add a manual full scan --- .env.example | 7 + app/admin_api.py | 12 ++ app/catalog_probe.py | 153 ++++++++++++++++++++++ app/model_policy.py | 7 + app/routing_scheduler.py | 10 +- app/settings.py | 3 + app/state_store.py | 8 +- converter.py | 219 +++++++++++++++++++++++++++++--- docs/advanced.md | 6 + docs/advanced.zh-CN.md | 6 + tests/test_admin_api.py | 19 +++ tests/test_catalog_probe.py | 204 +++++++++++++++++++++++++++++ tests/test_routing_scheduler.py | 30 ++++- tests/test_sqlite_state.py | 8 +- web/src/api.ts | 33 +++++ web/src/pages/Models.tsx | 51 +++++++- 16 files changed, 750 insertions(+), 26 deletions(-) create mode 100644 app/catalog_probe.py create mode 100644 tests/test_catalog_probe.py diff --git a/.env.example b/.env.example index e90e967..a3e728e 100644 --- a/.env.example +++ b/.env.example @@ -104,6 +104,13 @@ 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 # 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 4795f90..2282447 100644 --- a/app/admin_api.py +++ b/app/admin_api.py @@ -339,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 是只读字段") diff --git a/app/catalog_probe.py b/app/catalog_probe.py new file mode 100644 index 0000000..bf6a752 --- /dev/null +++ b/app/catalog_probe.py @@ -0,0 +1,153 @@ +"""Representative per-profile catalog probes with rotation, backoff and diff status.""" +from __future__ import annotations + +import math +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: the current source when due, else the stalest one.""" + 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 + source = self.row(profile).get("source") + for candidate in due: + if candidate["account_key"] == source: + return candidate + # Unknown or 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) + row.update(source=source if source is not None else row.get("source"), + failures=failures, status="error", + retry_at=now + min(RETRY_BASE_S * 2 ** (failures - 1), RETRY_MAX_S)) + 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/model_policy.py b/app/model_policy.py index afa2d79..094162d 100644 --- a/app/model_policy.py +++ b/app/model_policy.py @@ -116,6 +116,13 @@ def _active(config, upstream, policy=None): 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): """Resolve one public name to a literal upstream ID while retaining its own policy.""" if not isinstance(public_model, str) or not public_model.strip(): diff --git a/app/routing_scheduler.py b/app/routing_scheduler.py index 5df2a62..34b5eab 100644 --- a/app/routing_scheduler.py +++ b/app/routing_scheduler.py @@ -16,11 +16,12 @@ class RoutingScheduler: for the current revision set, and ``shadow`` does both while returning the legacy answer. """ - __slots__ = ("_shards", "mismatches") + __slots__ = ("_shards", "mismatches", "_limit") - def __init__(self): + 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: @@ -39,8 +40,9 @@ def candidates(self, *, mode, revisions, model, region, entries, evaluate): shard = self._shards.get(key) if shard is None: shard = tuple(entry for entry in entries if evaluate(entry)) - # Keep only the current revision set; a moved revision never reuses stale entries. - self._shards = {key: shard} + 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: diff --git a/app/settings.py b/app/settings.py index e0763c3..b3c7fe2 100644 --- a/app/settings.py +++ b/app/settings.py @@ -100,6 +100,9 @@ def _item(default, type_, label, *, mode="hot", env=None, minimum=None, maximum= "routing_scheduler": _item("indexed", "string", "路由调度模式", env="CODEBUDDY2API_ROUTING_SCHEDULER", choices=["legacy", "shadow", "indexed"]), + "catalog_probe_mode": _item("representative", "string", "目录探测模式", + env="CODEBUDDY2API_CATALOG_PROBE_MODE", + choices=["account", "representative"]), "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/converter.py b/converter.py index 33499a0..bc0bbca 100644 --- a/converter.py +++ b/converter.py @@ -75,6 +75,7 @@ def desensitize_body(body, roles=("system",), desensitize_harness_user=False, 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 @@ -1160,8 +1161,12 @@ def _shard_revisions(self, policy=None) -> tuple: 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) + # 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), + 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)) @@ -1752,29 +1757,41 @@ def _publish_model_cache(): invalidate_model_table() -def _sync_model_catalogs(pool, ledger, refs, failed): - cache = CONFIG.get("model_cache") - if cache is None: - return +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.""" + 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_account_catalogs(pool, ledger, refs, failed, cache): + """Refresh every synced account's own catalog; used by the account probe mode.""" 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) + key = _account_catalog_key(profile, headers) if cache.fresh(key) and not entry.get("catalog_dirty"): continue perf_metrics.inc("catalog_fetch_count") 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): + models, serves = _fetch_catalog(headers) + if _store_catalog(pool, cache, entry, generation, key, models, serves): _log(f"[models] {profile} 模型表已刷新: 选择器 {len(models)} 个，" f"账号根表 {len(serves)} 个") else: @@ -1783,9 +1800,169 @@ def publish(): perf_metrics.inc("catalog_fetch_failures") failed.add(entry["id"]) _sync_error(pool, ledger, entry, generation, "models", error) + + +def _mark_catalog_diffs(pool, profile, model_ids, cache): + """Mark accounts whose own catalog lacks models the representative just advertised.""" + if not model_ids: + return + 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 model_ids <= known: + entry["catalog_dirty"] = True + + +def _probe_profile_catalogs(pool, ledger, refs, failed, cache, aggregate): + """Probe at most one representative credential per product profile per cycle.""" + for profile in PROFILE_ENDPOINTS: + if not aggregate.due(profile): + continue + candidates = [{"entry": entry, "generation": generation, "headers": headers, + "account_key": entry.get("account_key")} + for entry, generation, headers, candidate_profile in refs.values() + if candidate_profile == profile and entry.get("account_key") + and model_policy.credential_enabled(CONFIG, entry)] + target = aggregate.choose( + profile, candidates, age_of=lambda candidate: cache.age(_account_catalog_key(profile, candidate["headers"]))) + if target is None: + aggregate.fresh(profile) + continue + 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"]) + continue + 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) + _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 _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) + 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_MAX_ROWS = 500 # Bounded per-account progress payload. +_SCAN_FIELDS = ("status", "total", "completed", "success", "failed", "differences", + "started_at", "finished_at", "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, "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 _catalog_scan_worker(pool, cache, scan, entries): + """Scan every account serially, keeping each previous catalog when its refresh fails.""" + aggregate = CONFIG.get("aggregate_catalog") + try: + for entry in entries: + 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"}) + 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)) + scan["failed"] += 1 + scan["completed"] += 1 + if len(scan["accounts"]) < CATALOG_SCAN_MAX_ROWS: + scan["accounts"].append(row) + _publish_model_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, + "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") @@ -2092,6 +2269,7 @@ async def _protocol_http_exception(request: Request, exc: HTTPException): "trial_ledger": None, "model_guard": True, # Reject models absent from authorized catalogs "routing_scheduler": "indexed", # legacy | shadow | indexed eligibility shards + "catalog_probe_mode": "representative", # account | representative catalog refresh "max_images": 16, "image_policy": "truncate", "max_request_bytes": 32 * 1024 * 1024, "log_body_limit": 65536, "max_inbound_bytes": 64 * 1024 * 1024, @@ -4694,6 +4872,9 @@ 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", "representative"), + help="目录探测：representative 每个产品每周期只探测一个代表账号（默认）；account 逐个账号刷新") ap.add_argument("--routing-scheduler", choices=("legacy", "shadow", "indexed"), default=os.environ.get("CODEBUDDY2API_ROUTING_SCHEDULER", "indexed"), help="路由调度：indexed 复用按模型分片（默认）；shadow 记录新旧差异并仍返回旧结果；legacy 每次全量判断") @@ -4734,7 +4915,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", "routing_scheduler"): + "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 @@ -4768,6 +4950,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 diff --git a/docs/advanced.md b/docs/advanced.md index c2a3d23..373cff1 100644 --- a/docs/advanced.md +++ b/docs/advanced.md @@ -239,6 +239,12 @@ 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`: `indexed` (default) reuses one eligibility shard per model, request-scoped alias and revision set; `shadow` also runs the legacy scan, records mismatches under `GET /admin/performance` and still returns the legacy answer; `legacy` re-evaluates every credential on every request. +- `catalog_probe_mode`: `representative` (default) probes one representative credential per product profile per TTL cycle, preferring the current source and rotating to the stalest account, with exponential backoff after failures; `account` refreshes every synced account each cycle. An unprobed or failed account keeps its last trusted catalog, so `/v1/models` stays the union of current account catalogs and only manually verified models are published. +- `POST /admin/models/scan` refreshes every current account once, in the background; `GET /admin/models/scan` reports progress, successes, failures and catalog differences. It is the only way to calibrate all accounts at once and is never a periodic task. + ### 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..5c82b2b 100644 --- a/docs/advanced.zh-CN.md +++ b/docs/advanced.zh-CN.md @@ -239,6 +239,12 @@ WebUI 可以直接上传文件；以下限制针对 `POST /admin/credentials` - WebUI 的地域、产品和凭证绑定严格限制候选账号，不会回退到未选账号。模型禁用后直接请求同样拒绝；改名默认不保留原 ID，只有选择保留时才同时提供旧 ID。 - 已发送请求不会因账号不可用或 HTTP 错误换账号重放；后续请求才重新选路。目录同步或凭证未就绪通常返回带 `Retry-After` 的 503，不支持或禁用的模型返回 404。 +选路与目录维护都可切换模式，无需回滚已持久化数据： + +- `routing_scheduler`：`indexed`（默认）按模型、请求内别名和版本组合复用同一份候选分片；`shadow` 同时执行旧的全量判断，把差异记录在 `GET /admin/performance`，但仍返回旧结果；`legacy` 每次请求重新判断全部凭证。 +- `catalog_probe_mode`：`representative`（默认）每个产品每 TTL 周期只探测一个代表凭证，优先当前来源账号，再轮换到目录最旧的账号，失败按指数退避；`account` 每周期刷新全部已同步账号。未探测或探测失败的账号保留上次可信目录，`/v1/models` 仍是当前账号目录并集，只有经过账号级验证的模型才会发布。 +- `POST /admin/models/scan` 在后台把当前全部账号各刷新一次，`GET /admin/models/scan` 返回进度、成功数、失败数与目录差异数；它是一次性校准入口，不是周期任务。 + ### 模型声明与图片兼容 `/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..405c3c5 --- /dev/null +++ b/tests/test_catalog_probe.py @@ -0,0 +1,204 @@ +"""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_the_current_source_then_the_stalest_account(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"], "a") + ages["a"] = 10.0 # The source just refreshed; rotation moves to the stalest account. + self.assertEqual(aggregate.choose("cn-cli", accounts, age_of=age_of)["account_key"], "c") + ages["c"] = 10.0 + self.assertEqual(aggregate.choose("cn-cli", accounts, age_of=age_of)["account_key"], "b") + 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_marks_differences(self): + pool, ledger, cache, aggregate = self.build(3) + entries = pool.entries() + stale = {item["id"] for item in model("hy3")} + 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()) + self.assertEqual(len(self.calls), 1) + self.assertTrue(stale < {item["id"] for item in model("hy3", "glm-5.3")}) + fresh = [entry for entry in pool.entries() if not entry.get("catalog_dirty")] + self.assertEqual(len(fresh), 1) # Only the representative stored a new catalog. + self.assertTrue(pool.entry_by_id(entries[1]["id"])["catalog_dirty"]) + 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), 1) # The TTL suppresses another fetch. + + 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_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_routing_scheduler.py b/tests/test_routing_scheduler.py index 1797274..b58311a 100644 --- a/tests/test_routing_scheduler.py +++ b/tests/test_routing_scheduler.py @@ -12,7 +12,8 @@ from unittest.mock import patch import converter -from app import perf_metrics +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 @@ -80,6 +81,7 @@ 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"): @@ -173,6 +175,32 @@ def test_credential_reload_invalidates_the_shard(self): 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") 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..2cacfd0 100644 --- a/web/src/api.ts +++ b/web/src/api.ts @@ -280,6 +280,39 @@ export function modelResponse(value: unknown): { : undefined, }; } +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; + 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), + accounts: Array.isArray(data.accounts) ? list(data.accounts, "目录扫描账号") : [], + }; +} 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..4f28a62 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,32 @@ 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) => setScan(catalogScanResponse(response.data))) + .catch((err: unknown) => setError(errorMessage(err))); + }, []); + useEffect(() => { + if (scan?.status !== "running") return; + const timer = setInterval(loadScan, 2000); + return () => clearInterval(timer); + }, [scan?.status, loadScan]); + 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 +577,16 @@ export function Models() { 刷新目录 +