diff --git a/app/audit_store.py b/app/audit_store.py index 87edfb5..127d641 100644 --- a/app/audit_store.py +++ b/app/audit_store.py @@ -288,8 +288,11 @@ def _aggregate(self, record): increment[key + "_sum"] = record[key] or 0 increment[key + "_known"] = int(record[key] is not None) dimensions = [("global", "")] - dimensions += [(key, record.get("public_model" if key == "model" else key) or "") - for key in ("model", "profile", "credential")] + # Unknown model/profile/credential values get no bucket of their own: the global + # row still counts the request, and an empty key would rank as "未识别模型". + dimensions += [(key, value) + for key in ("model", "profile", "credential") + if (value := record.get("public_model" if key == "model" else key))] for table, seconds in (("stats_hourly", 3600), ("stats_daily", 86400), ("stats_totals", 0)): bucket = int(record["started_at"] // seconds) * seconds if seconds else 0 for dimension, key in dimensions: @@ -534,6 +537,10 @@ def fetch(): if grain == "day": series.append({"bucket": row["bucket"], **stats}) elif row["dimension"] in ("model", "profile"): + # Legacy rows recorded an empty key for unidentified models; skip + # them instead of surfacing an anonymous bucket in the ranking. + if not row["dimension_key"]: + continue target = models if row["dimension"] == "model" else profiles self._merge(target.setdefault(row["dimension_key"], self._empty_stats()), stats) if grain == "hour": diff --git a/app/auth_oauth.py b/app/auth_oauth.py index 072b858..0f9b861 100644 --- a/app/auth_oauth.py +++ b/app/auth_oauth.py @@ -87,9 +87,11 @@ def validate_cred_data(data) -> tuple[str | None, str | None]: if not isinstance(data, dict): return None, "凭据不是有效的 JSON 对象" acct = data.get("account") - if not isinstance(acct, dict): + if not isinstance(acct, dict) or not acct: + # An empty ``account`` object must not hide a populated ``accounts[0]``. arr = data.get("accounts") - acct = arr[0] if isinstance(arr, list) and arr and isinstance(arr[0], dict) else None + candidate = arr[0] if isinstance(arr, list) and arr and isinstance(arr[0], dict) else None + acct = candidate if candidate else (acct if isinstance(acct, dict) else None) uid = str((acct or {}).get("uid") or "") if not uid: return None, "缺少 account.uid" diff --git a/app/client_profiles.py b/app/client_profiles.py index a64a7d3..c77b3a3 100644 --- a/app/client_profiles.py +++ b/app/client_profiles.py @@ -38,12 +38,14 @@ def identity_headers(profile: str) -> dict: def _auth_headers(auth: dict, account: dict) -> dict: + # Read the tenant id through the shared helper so header and billing dispatch agree. + enterprise_id = enterprise_id_of(account) or "" headers = { "Content-Type": "application/json", "Accept": "application/json", "Authorization": f"Bearer {auth.get('accessToken', '')}", "X-User-Id": str(account.get("uid") or ""), - "X-Enterprise-Id": str(account.get("enterpriseId") or ""), - "X-Tenant-Id": str(account.get("enterpriseId") or ""), + "X-Enterprise-Id": enterprise_id, + "X-Tenant-Id": enterprise_id, "X-Domain": domain_for_auth(auth), "X-Product": "SaaS", "X-Requested-With": "XMLHttpRequest", } @@ -73,6 +75,37 @@ def account_key(profile: str, uid, enterprise_id="") -> str: return hashlib.sha256(identity.encode("utf-8")).hexdigest() +def enterprise_id_of(account) -> str | None: + """Read the tenant identifier in either naming convention. + + The official clients emit ``enterpriseId`` from the auth payload and + ``enterprise_id`` in some backup shapes; both must resolve to one value so the + billing dispatch and the request headers never disagree. + """ + if not isinstance(account, dict): + return None + for field in ("enterpriseId", "enterprise_id"): + value = account.get(field) + if isinstance(value, str) and value.strip(): + return value.strip() + return None + + +def is_enterprise_account(account) -> bool: + """Whether the account belongs to an enterprise tenant. + + An enterprise id is a direct hit. Historical backups may carry only + ``type in {ultimate, exclusive}`` without the id, but those accounts still bill + through the enterprise endpoint, so the type is a fallback signal. + """ + if enterprise_id_of(account) is not None: + return True + if not isinstance(account, dict): + return False + account_type = account.get("type") + return isinstance(account_type, str) and account_type.strip().lower() in ("ultimate", "exclusive") + + def catalog_cache_key(profile: str, identity: str | None = None) -> str: revision = CLI_VERSION if profile_product(profile) == "cli" else f"{WORKBUDDY_VERSION}:{WORKBUDDY_CLI_VERSION}" key = f"{profile}:{revision}" diff --git a/app/credits.py b/app/credits.py index 20e33d1..6182e1b 100644 --- a/app/credits.py +++ b/app/credits.py @@ -13,7 +13,7 @@ import httpx -from .client_profiles import catalog_headers +from .client_profiles import catalog_headers, enterprise_id_of from .site_routing import PROFILE_ENDPOINTS, profile_for_auth, profile_product BROWSER_UA = ("Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 " @@ -30,6 +30,14 @@ CHECKIN_PATHS = ("/v2/billing/meter/daily-checkin",) CHECKIN_STATUS_PATH = "/v2/billing/meter/checkin-activity-status" RESOURCE_PATH = "/v2/billing/meter/get-user-resource" +# Enterprise tenants do not bill through the personal resource endpoint; it answers +# with an empty Accounts structure. The official clients dispatch on the tenant id +# before the region/product split, so this path is shared by CN and intl accounts. +ENTERPRISE_USAGE_PATH = "/v2/billing/meter/get-enterprise-user-usage" +# Fallback tenant lookup for historical credentials that carry a type but no id. +ENTERPRISE_ACCOUNTS_PATH = "/console/accounts" +# An enterprise allowance reported as -1 is unlimited rather than exhausted. +UNLIMITED_ALLOWANCE = -1 CONFIG_PATH = "/v3/config" # Official cloud model catalog RESOURCE_PRODUCT_CODE = "p_tcaca" CREDITS_PAGE_SIZE = 100 @@ -66,7 +74,8 @@ def hosts_for_token(access_token: str, domain: str = "") -> list[str]: return [BILLING_PROFILE_HOSTS[profile]] -def _web_headers(api_host: str, access_token: str, uid: str = "", domain: str = "") -> dict: +def _web_headers(api_host: str, access_token: str, uid: str = "", domain: str = "", + enterprise_id: str = "") -> dict: """Build website billing headers independently of CLI model catalog headers.""" return { "accept": "application/json, text/plain, */*", @@ -77,6 +86,10 @@ def _web_headers(api_host: str, access_token: str, uid: str = "", domain: str = "authorization": f"Bearer {access_token}", "x-user-id": str(uid or ""), "x-domain": str(domain or ""), + # The website interceptor always carries the tenant id; enterprise endpoints + # reject the request without it. + "x-enterprise-id": str(enterprise_id or ""), + "x-tenant-id": str(enterprise_id or ""), "user-agent": BROWSER_UA, } @@ -93,6 +106,18 @@ def _post_json(client: httpx.Client, url: str, headers: dict, body: dict) -> tup return r.status_code, payload +def _get_json(client: httpx.Client, url: str, headers: dict) -> tuple[int, dict]: + """GET JSON and return status/payload, raising AuthExpiredError on HTTP 401.""" + r = client.get(url, headers=headers, timeout=REQUEST_TIMEOUT) + if r.status_code == 401: + raise AuthExpiredError("登录身份过期") + try: + payload = r.json() + except Exception: + payload = {} + return r.status_code, payload + + # --------------------------------------------------------------------------- # Daily check-in # --------------------------------------------------------------------------- @@ -340,8 +365,132 @@ def _fetch_accounts_page(client, url: str, headers: dict, page: int, *, retry_em raise RuntimeError(f"积分查询失败: {last_err}") -def fetch_credits(access_token: str, uid: str = "", domain: str = "") -> dict: - """Fetch credit segments and mark page-limited results partial; HTTP 401 raises AuthExpiredError.""" +# --------------------------------------------------------------------------- +# Enterprise allowance +# --------------------------------------------------------------------------- + +_ENTERPRISE_LIMIT_FIELDS = ("limitNum", "limit_num") +_ENTERPRISE_USED_FIELDS = ("credit", "used_num") +_ENTERPRISE_RESET_FIELDS = ("cycleResetTime", "CycleResetTime", "cycle_reset_time") +_ENTERPRISE_PARSE_ERROR = "企业配额数据无法解析" + + +def parse_enterprise_usage(payload: dict) -> dict | None: + """Parse the enterprise allowance payload, or return None when unreadable. + + Accepts the three observed wrappings (``data.data``, ``data``, top level) and both + naming conventions, since the two client implementations disagree on casing. + """ + if not isinstance(payload, dict): + return None + data = payload.get("data") if isinstance(payload.get("data"), dict) else {} + nested = data.get("data") if isinstance(data.get("data"), dict) else {} + for candidate in (nested, data, payload): + if not candidate: + continue + limit = _first_number(candidate, _ENTERPRISE_LIMIT_FIELDS) + if limit is None or not math.isfinite(limit): + continue + used = _first_number(candidate, _ENTERPRISE_USED_FIELDS) + used = max(used, 0.0) if used is not None and math.isfinite(used) else 0.0 + return {"limit": limit, "used": used, + "reset_at": _first_ts(candidate, _ENTERPRISE_RESET_FIELDS)} + return None + + +def enterprise_id_from_accounts(payload: dict) -> str | None: + """Extract the first tenant id from a ``/console/accounts`` response.""" + if not isinstance(payload, dict): + return None + code = payload.get("code") + if code not in (0, 200, None): + return None + data = payload.get("data") if isinstance(payload.get("data"), dict) else {} + for key in ("accounts", "Accounts"): + items = data.get(key) + if isinstance(items, list) and items and isinstance(items[0], dict): + return enterprise_id_of(items[0]) + return None + + +def resolve_enterprise_id(access_token: str, uid: str = "", domain: str = "") -> str | None: + """Look up the tenant id once for credentials that only carry an enterprise type. + + Historical backups may have ``type=ultimate`` without ``enterpriseId``. A single + attempt is enough: the caller turns a miss into a readable error instead of + silently falling back to the personal endpoint. + """ + host = hosts_for_token(access_token, domain)[0] + url = host + ENTERPRISE_ACCOUNTS_PATH + headers = _web_headers(host, access_token, uid, domain) + try: + with httpx.Client() as client: + status, payload = _get_json(client, url, headers) + except httpx.HTTPError: + return None + if status != 200: + return None + return enterprise_id_from_accounts(payload) + + +def _enterprise_error(reason: str) -> RuntimeError: + """An enterprise account must never be reported as having no credits.""" + return RuntimeError(f"企业配额查询失败: {reason}") + + +def fetch_enterprise_credits(access_token: str, uid: str = "", domain: str = "", + enterprise_id: str = "") -> dict: + """Fetch the tenant allowance instead of the personal package list. + + The enterprise endpoint returns a single cycle allowance rather than package + segments, so it is projected as one segment to keep every downstream consumer + (expiry-aware routing, dashboards, the check-in ledger) unchanged. + """ + host = hosts_for_token(access_token, domain)[0] + url = host + ENTERPRISE_USAGE_PATH + headers = _web_headers(host, access_token, uid, domain, enterprise_id) + with httpx.Client() as client: + status, payload = _post_json(client, url, headers, {}) + if status != 200: + raise _enterprise_error(f"积分接口 HTTP {status}") + code = payload.get("code") + if code not in (0, None): + raise _enterprise_error(str(payload.get("msg") or f"积分接口 code={code}")) + usage = parse_enterprise_usage(payload) + if usage is None: + raise _enterprise_error(_ENTERPRISE_PARSE_ERROR) + unlimited = usage["limit"] == UNLIMITED_ALLOWANCE + if unlimited: + # Do not fabricate a balance: an unlimited tenant has no numeric remainder. + return {"credits": None, "count": 0, "segments": [], "soonest_expiry": None, + "intl": is_international_host(host), "partial": False, "unlimited": True, + "enterprise": True} + remaining = max(usage["limit"] - usage["used"], 0.0) + segment = { + "remaining": round(remaining, 2), + "total": round(max(usage["limit"], 0.0), 2), + "expires_at": usage["reset_at"], + "source": "企业配额", + "package_code": "", + } + segments = merge_segments(extract_segments([]) + [segment]) if remaining > 0 else [] + return {"credits": round(sum(s["remaining"] for s in segments), 2), + "count": 1 if not unlimited else 0, "segments": segments, + "soonest_expiry": soonest_expiry(segments), + "intl": is_international_host(host), "partial": False, + "unlimited": False, "enterprise": True} + + +def fetch_credits(access_token: str, uid: str = "", domain: str = "", + enterprise_id: str = "") -> dict: + """Fetch credit segments and mark page-limited results partial; HTTP 401 raises AuthExpiredError. + + Enterprise tenants are dispatched first: their allowance never appears in the + personal package list, which answers with an empty Accounts structure for them. + """ + if enterprise_id: + return fetch_enterprise_credits(access_token, uid=uid, domain=domain, + enterprise_id=enterprise_id) host = hosts_for_token(access_token, domain)[0] url = host + RESOURCE_PATH headers = _web_headers(host, access_token, uid, domain) @@ -549,11 +698,20 @@ def dedupe_by_identity(creds_snapshot: dict) -> dict: def aggregate_credits(creds_snapshot: dict) -> dict: - """Deduplicate accounts and aggregate regional balances, usage and earliest expiry.""" + """Deduplicate accounts and aggregate regional balances, usage and earliest expiry. + + Unlimited enterprise tenants carry no numeric balance: they are counted + separately instead of entering the monetary conversion. + """ groups = {k: {"remaining": 0.0, "used_by_quota": 0.0, "soonest_expiry": None} for k in ("domestic", "international")} - for e in dedupe_by_identity(creds_snapshot).values(): + deduped = dedupe_by_identity(creds_snapshot) + unlimited = 0 + for e in deduped.values(): c = e.get("credits") or {} + if c.get("unlimited") and c.get("enterprise"): + unlimited += 1 + continue g = groups["international" if c.get("intl") else "domestic"] for s in c.get("segments") or []: rem = float(s.get("remaining") or 0) @@ -570,7 +728,8 @@ def aggregate_credits(creds_snapshot: dict) -> dict: "used_by_quota": round(sum(g["used_by_quota"] for g in groups.values()), 2), "soonest_expiry": min(exps) if exps else None, "partial": any(bool((e.get("credits") or {}).get("partial")) - for e in dedupe_by_identity(creds_snapshot).values()), + for e in deduped.values()), + "unlimited_accounts": unlimited, "groups": groups} @@ -720,6 +879,10 @@ def update_credits(self, cred_id: str, result: dict): "fetched_at": time.time(), "intl": bool(result.get("intl")), # Regional credits use independent pricing. "partial": bool(result.get("partial")), # Expose incomplete pagination. + "enterprise": bool(result.get("enterprise")), + # Unlimited only exists on the enterprise dispatch; the double gate keeps a + # malformed personal result from ever storing it and bypassing zero-balance rules. + "unlimited": bool(result.get("unlimited")) and bool(result.get("enterprise")), } e["error"] = None self._save() diff --git a/converter.py b/converter.py index 840d005..2c9af7e 100644 --- a/converter.py +++ b/converter.py @@ -80,7 +80,8 @@ def desensitize_body(body, roles=("system",), desensitize_harness_user=False, 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, profile_site, chat_url_for_headers, refresh_url_for_auth) -from app.client_profiles import CLI_VERSION, CLI_USER_AGENT, credential_headers, catalog_cache_key, account_key +from app.client_profiles import (CLI_VERSION, CLI_USER_AGENT, credential_headers, catalog_cache_key, + account_key, enterprise_id_of, is_enterprise_account) try: from app import credits as credits_mod except ImportError: # Disable credit maintenance when its module is unavailable. @@ -210,16 +211,49 @@ def _cred_uid(path) -> Optional[str]: return None def _credential_account(data: dict) -> dict: + """Return the account block, preferring a populated ``account`` object. + + An empty ``account`` object must not shadow a populated ``accounts[0]``: some + exports carry both, and treating ``{}`` as authoritative makes the credential + look like it has no uid at all. + """ account = data.get("account") - if not isinstance(account, dict): - accounts = data.get("accounts") or [] - account = accounts[0] if isinstance(accounts, list) and accounts and isinstance(accounts[0], dict) else {} - return account + if isinstance(account, dict) and account: + return account + accounts = data.get("accounts") + candidate = accounts[0] if isinstance(accounts, list) and accounts and isinstance(accounts[0], dict) else None + if candidate: + return candidate + return account if isinstance(account, dict) else {} + + +def _identity_email(account: dict) -> str | None: + """Return a real email usable as an identity when no uid exists. + + Display placeholders and values that merely repeat another field must never be + used: they would merge unrelated accounts onto a single key. + """ + email = account.get("email") + if not isinstance(email, str): + return None + text = email.strip() + if "@" not in text or text.lower() == "unknown" or text in ("手动添加", "未命名账号"): + return None + if text == str(account.get("nickname") or "").strip(): + return None + if text == str(account.get("uid") or "").strip(): + return None + return text.lower() def _credential_identity(data: dict) -> str: + """Hash the stable account identity, falling back to a real email without a uid.""" account = _credential_account(data) - return account_key(profile_for_auth(data.get("auth") or {}), account.get("uid"), account.get("enterpriseId")) + uid = account.get("uid") + if not (isinstance(uid, str) and uid.strip()): + uid = _identity_email(account) or "" + return account_key(profile_for_auth(data.get("auth") or {}), uid, + enterprise_id_of(account)) def _cred_identity(path) -> str | None: @@ -933,6 +967,9 @@ def _zero_balance(self, entry, profile) -> bool: if not balance: return False try: + # An unlimited enterprise tenant holds no numeric balance and is never zero. + if balance.get("unlimited") and balance.get("enterprise"): + return False return (bool(balance.get("intl")) == (profile_region(profile) == "intl") and float(balance.get("credits") or 0) <= 0) except (TypeError, ValueError): @@ -943,6 +980,9 @@ def _has_credit(self, entry, profile): if not balance: return profile_region(profile) == "cn" try: + # An unlimited enterprise tenant holds no numeric balance; it has credit on its own site. + if balance.get("unlimited") and balance.get("enterprise"): + return bool(balance.get("intl")) == (profile_region(profile) == "intl") return (bool(balance.get("intl")) == (profile_region(profile) == "intl") and float(balance.get("credits") or 0) > 0) except (TypeError, ValueError): @@ -1530,13 +1570,24 @@ def can_chat(): _sync_error(pool, ledger, entry, generation, "dailychat", error) if not model_policy.credential_enabled(CONFIG, entry): return None - balance = credits_mod.fetch_credits(token, uid=uid, domain=domain) + # An enterprise type without a stored tenant id still bills through the + # enterprise endpoint, so resolve the id once before dispatching. + account = _credential_account(cm._session()) + enterprise_id = enterprise_id_of(account) + if enterprise_id is None and is_enterprise_account(account): + enterprise_id = credits_mod.resolve_enterprise_id(token, uid=uid, domain=domain) or "" + balance = credits_mod.fetch_credits( + token, uid=uid, domain=domain, + # Enterprise tenants bill through a different endpoint; the dispatch key + # travels from the credential so both views agree. + enterprise_id=enterprise_id or "") if bool(balance.get("intl")) != (site == INTERNATIONAL): raise ValueError("积分响应与凭据站点不一致") if not pool.apply_if_current(cm, generation, lambda: ledger.update_credits(cid, balance)): failed.add(cid) return None - _log(f"[credits] {Path(cid).name}: 站点 {site}，余额 {balance['credits']}") + shown = "无限" if balance.get("unlimited") and balance.get("enterprise") else balance["credits"] + _log(f"[credits] {Path(cid).name}: 站点 {site}，余额 {shown}") profile = profile_for_headers(headers) return entry, generation, headers, profile except Exception as error: @@ -2374,6 +2425,8 @@ def _billing_totals() -> dict: "used_source": "official_usage_detail" if detail else "quota_delta", # Expose incomplete pagination or failed account synchronization. "partial": bool(agg.get("partial") or cache.get("partial")), + # Unlimited enterprise tenants stay outside the monetary conversion. + "unlimited_accounts": int(agg.get("unlimited_accounts") or 0), "groups": groups_out, "by_day": cache.get("by_day") or {}} @@ -2620,9 +2673,16 @@ def _profile_has_credits(profile: str) -> bool: if profile_region(profile) == "cn": return True ledger = CONFIG.get("ledger") - return bool(ledger and any((entry.get("credits") or {}).get("intl") - and float((entry.get("credits") or {}).get("credits") or 0) > 0 - for entry in ledger.snapshot().values())) + + def _positive(entry): + c = entry.get("credits") or {} + if not c.get("intl"): + return False + if c.get("unlimited") and c.get("enterprise"): + return True # Unlimited tenants hold no numeric balance. + return float(c.get("credits") or 0) > 0 + + return bool(ledger and any(_positive(entry) for entry in ledger.snapshot().values())) return any(pool._entry_profile(entry) == profile and pool._has_credit(entry, profile) for entry in pool.entries() if model_policy.credential_enabled(CONFIG, entry)) diff --git a/tests/test_dashboard_granularity.py b/tests/test_dashboard_granularity.py index a7c9b3c..f3ca406 100644 --- a/tests/test_dashboard_granularity.py +++ b/tests/test_dashboard_granularity.py @@ -1,4 +1,5 @@ """Dashboard bucket regressions using isolated persistent aggregates.""" +import json import sys from pathlib import Path sys.path.insert(0, str(Path(__file__).resolve().parents[1])) @@ -68,5 +69,39 @@ def test_missing_hourly_history_is_marked_incomplete_not_split_from_daily(self): self.assertFalse(self.store.dashboard(1, "day")["range"]["partial"]) +class UnidentifiedModelTests(unittest.TestCase): + """Unidentified models must not surface as an anonymous bucket in the ranking.""" + + def setUp(self): + self.root = Path(self.enterContext(tempfile.TemporaryDirectory())) + self.start = datetime(2026, 9, 14, tzinfo=timezone.utc).timestamp() + self.enterContext(patch("app.audit_store.time.time", return_value=self.start + 12.5 * 3600)) + self.store = AuditStore(self.root / "logs.sqlite3") + self.addCleanup(self.store.close) + + def _record(self, **overrides): + result = self.store.record_request({"id": "r1", **self.store.ticket(), + "started_at": self.start, "profile": "cn-cli", "protocol": "chat", + "status_code": 200, "outcome": "success", **overrides}) + self.assertTrue(result["ok"]) + + def test_new_records_without_model_write_no_empty_bucket(self): + self._record() # No public_model at all. + dash = self.store.dashboard(1) + self.assertEqual(dash["summary"]["requests"], 1) # The global row still counts it. + self.assertEqual(dash["models"], []) + + def test_legacy_empty_dimension_key_is_skipped(self): + self._record(public_model="synthetic-model") + stats = self.store._empty_stats() + stats["requests"] = 7 + bucket = int(self.start // 86400) * 86400 + self.store._db.execute("INSERT OR REPLACE INTO stats_daily VALUES(?,?,?,?)", + (bucket, "model", "", json.dumps(stats))) + models = self.store.dashboard(1)["models"] + self.assertEqual([m["model"] for m in models], ["synthetic-model"]) + self.assertEqual(models[0]["requests"], 1) # The anonymous 7 are not merged anywhere. + + if __name__ == "__main__": unittest.main() diff --git a/tests/test_enterprise_credits.py b/tests/test_enterprise_credits.py new file mode 100644 index 0000000..b46c805 --- /dev/null +++ b/tests/test_enterprise_credits.py @@ -0,0 +1,423 @@ +#!/usr/bin/env python3 +"""Test enterprise tenant identification and the enterprise billing dispatch. + +Enterprise credentials do not appear in the personal package list: the personal +resource endpoint answers with an empty Accounts structure for them. These tests +pin the tenant detection, the payload parsing (including the unlimited marker) and +the guarantee that a personal account keeps using the personal endpoint. +""" + +import base64 +import json +import sys +import tempfile +import time +from pathlib import Path +from unittest import TestCase +from unittest.mock import patch + +import httpx + +sys.path.insert(0, str(Path(__file__).resolve().parents[1])) # Allow direct execution. + +from app import credits +import converter +from app.client_profiles import ( + CLI_USER_AGENT, enterprise_id_of, is_enterprise_account, +) +from app.credits import ( + UNLIMITED_ALLOWANCE, + enterprise_id_from_accounts, + fetch_credits, + fetch_enterprise_credits, + parse_enterprise_usage, +) + + +def _jwt(iss: str) -> str: + payload = base64.urlsafe_b64encode(json.dumps({"iss": iss}).encode()).rstrip(b"=").decode() + return f"header.{payload}.sig" + + +CN_TOKEN = _jwt("https://www.codebuddy.cn/auth/realms/copilot") + +# Captured from a live enterprise credential (values are structural, not account data). +ENTERPRISE_PAYLOAD = { + "code": 0, + "msg": "OK", + "data": {"credit": 4.07, "cycleStartTime": "2026-09-21 00:00:00", + "cycleEndTime": "2026-10-21 23:59:59", "limitNum": 2000, + "cycleResetTime": "2026-10-22 00:00:00"}, +} + + +# Capture the real client before any patching so building the fake cannot recurse. +_REAL_HTTPX_CLIENT = httpx.Client + + +class _FakeClient: + """Stand-in for httpx.Client so patching the module attribute cannot recurse.""" + + def __init__(self, handler): + self._client = _REAL_HTTPX_CLIENT(transport=httpx.MockTransport(handler)) + + def __enter__(self): + self._client.__enter__() + return self._client + + def __exit__(self, *exc): + return self._client.__exit__(*exc) + + +def _patch_client(handler): + """Patch the client the module under test opens, with one fresh client per call.""" + return patch.object(credits.httpx, "Client", lambda *a, **k: _FakeClient(handler)) + + +class EnterpriseIdentityTests(TestCase): + """Tenant detection must accept both naming conventions and the type fallback.""" + + def test_reads_camel_and_snake_case_ids(self): + self.assertEqual(enterprise_id_of({"enterpriseId": "tenant-a"}), "tenant-a") + self.assertEqual(enterprise_id_of({"enterprise_id": "tenant-b"}), "tenant-b") + # A populated candidate must not be shadowed by an empty alias. + self.assertEqual(enterprise_id_of({"enterpriseId": "", "enterprise_id": "tenant-c"}), "tenant-c") + + def test_ignores_blank_and_non_string_ids(self): + for account in ({}, {"enterpriseId": ""}, {"enterpriseId": " "}, + {"enterpriseId": None}, {"enterpriseId": 123}, + {"enterprise_id": ["x"]}, None, "not-a-dict"): + with self.subTest(account=account): + self.assertIsNone(enterprise_id_of(account)) + + def test_identity_is_camel_snake_equivalent(self): + """Both conventions must hash to one key so routing never splits an account.""" + from app.client_profiles import account_key + camel = account_key("cn-cli", "uid-1", enterprise_id_of({"enterpriseId": "t"}) or "") + snake = account_key("cn-cli", "uid-1", enterprise_id_of({"enterprise_id": "t"}) or "") + self.assertEqual(camel, snake) + + def test_enterprise_id_present_is_a_direct_hit(self): + self.assertTrue(is_enterprise_account({"enterpriseId": "t", "type": "personal"})) + self.assertTrue(is_enterprise_account({"enterprise_id": "t"})) + + def test_type_fallback_matches_only_enterprise_tiers(self): + for account_type in ("ultimate", "exclusive", "ULTIMATE", " Exclusive "): + with self.subTest(account_type=account_type): + self.assertTrue(is_enterprise_account({"type": account_type})) + for account_type in ("personal", "", "free", "pro", None, 123): + with self.subTest(account_type=account_type): + self.assertFalse(is_enterprise_account({"type": account_type})) + + def test_personal_account_is_not_enterprise(self): + self.assertFalse(is_enterprise_account({"type": "personal", "uid": "u"})) + + +class EnterprisePayloadTests(TestCase): + """Parsing must tolerate every observed wrapping and both key conventions.""" + + def test_parses_live_shape(self): + usage = parse_enterprise_usage(ENTERPRISE_PAYLOAD) + self.assertEqual(usage["limit"], 2000.0) + self.assertEqual(usage["used"], 4.07) + self.assertIsNotNone(usage["reset_at"]) + + def test_parses_all_wrappings(self): + inner = {"limitNum": 100, "credit": 25} + for payload in ({"data": {"data": dict(inner)}}, {"data": dict(inner)}, dict(inner)): + with self.subTest(payload=payload): + usage = parse_enterprise_usage(payload) + self.assertIsNotNone(usage) + self.assertEqual(usage["limit"], 100.0) + self.assertEqual(usage["used"], 25.0) + + def test_parses_snake_case_keys(self): + usage = parse_enterprise_usage({"data": {"limit_num": 500, "used_num": 125}}) + self.assertEqual(usage["limit"], 500.0) + self.assertEqual(usage["used"], 125.0) + + def test_missing_limit_is_not_parsed(self): + """A missing limit must fail loudly rather than report a zero balance.""" + for payload in ({}, {"data": {}}, {"data": {"credit": 5}}, None, "text"): + with self.subTest(payload=payload): + self.assertIsNone(parse_enterprise_usage(payload)) + + def test_negative_used_is_clamped(self): + usage = parse_enterprise_usage({"data": {"limitNum": 10, "credit": -5}}) + self.assertEqual(usage["used"], 0.0) + + def test_missing_used_defaults_to_zero(self): + usage = parse_enterprise_usage({"data": {"limitNum": 10}}) + self.assertEqual(usage["used"], 0.0) + + +class EnterpriseFetchTests(TestCase): + """The enterprise endpoint replaces the personal package list for tenants.""" + + def _capture(self, payload, status=200): + seen = {} + + def handler(request): + seen["url"] = str(request.url) + seen["headers"] = dict(request.headers) + return httpx.Response(status, json=payload) + + return seen, handler + + def test_remaining_is_limit_minus_used(self): + seen, handler = self._capture(ENTERPRISE_PAYLOAD) + with httpx.Client(transport=httpx.MockTransport(handler)) as client: + with _patch_client(handler): + result = fetch_enterprise_credits(CN_TOKEN, uid="uid-1", + domain="www.codebuddy.cn", + enterprise_id="tenant-a") + self.assertTrue(result["enterprise"]) + self.assertFalse(result["unlimited"]) + self.assertEqual(result["credits"], 1995.93) + self.assertEqual(result["segments"][0]["total"], 2000.0) + self.assertEqual(result["segments"][0]["remaining"], 1995.93) + self.assertEqual(result["segments"][0]["source"], "企业配额") + # The tenant id must travel on the request; the endpoint rejects otherwise. + self.assertEqual(seen["headers"].get("x-enterprise-id"), "tenant-a") + self.assertEqual(seen["headers"].get("x-tenant-id"), "tenant-a") + self.assertIn(credits.ENTERPRISE_USAGE_PATH, seen["url"]) + + def test_unlimited_is_flagged_without_a_number(self): + """limitNum == -1 means unlimited: never report a fabricated balance.""" + payload = {"code": 0, "data": {"limitNum": UNLIMITED_ALLOWANCE, "credit": 12.0}} + _seen, handler = self._capture(payload) + with httpx.Client(transport=httpx.MockTransport(handler)) as client: + with _patch_client(handler): + result = fetch_enterprise_credits(CN_TOKEN, uid="u", domain="www.codebuddy.cn", + enterprise_id="t") + self.assertTrue(result["unlimited"]) + self.assertIsNone(result["credits"]) + self.assertEqual(result["segments"], []) + + def test_exhausted_allowance_yields_no_segments(self): + payload = {"code": 0, "data": {"limitNum": 100, "credit": 100}} + _seen, handler = self._capture(payload) + with httpx.Client(transport=httpx.MockTransport(handler)) as client: + with _patch_client(handler): + result = fetch_enterprise_credits(CN_TOKEN, uid="u", domain="www.codebuddy.cn", + enterprise_id="t") + self.assertEqual(result["credits"], 0.0) + self.assertEqual(result["segments"], []) + + def test_api_error_raises_readable_failure(self): + payload = {"code": 40001, "msg": "tenant not found"} + _seen, handler = self._capture(payload) + with httpx.Client(transport=httpx.MockTransport(handler)) as client: + with _patch_client(handler): + with self.assertRaises(RuntimeError) as caught: + fetch_enterprise_credits(CN_TOKEN, uid="u", domain="www.codebuddy.cn", + enterprise_id="t") + self.assertIn("企业配额", str(caught.exception)) + + def test_unparsable_payload_raises_instead_of_reporting_zero(self): + _seen, handler = self._capture({"code": 0, "data": {"unexpected": True}}) + with httpx.Client(transport=httpx.MockTransport(handler)) as client: + with _patch_client(handler): + with self.assertRaises(RuntimeError): + fetch_enterprise_credits(CN_TOKEN, uid="u", domain="www.codebuddy.cn", + enterprise_id="t") + + +class BillingDispatchTests(TestCase): + """Dispatch: a tenant id selects the enterprise endpoint, everything else does not.""" + + def test_enterprise_id_selects_enterprise_endpoint(self): + seen = {} + + def handler(request): + seen["url"] = str(request.url) + return httpx.Response(200, json=ENTERPRISE_PAYLOAD) + + with httpx.Client(transport=httpx.MockTransport(handler)) as client: + with _patch_client(handler): + fetch_credits(CN_TOKEN, uid="u", domain="www.codebuddy.cn", enterprise_id="tenant-a") + + self.assertIn(credits.ENTERPRISE_USAGE_PATH, seen["url"]) + self.assertNotIn(credits.RESOURCE_PATH, seen["url"]) + + def test_personal_account_keeps_using_the_personal_endpoint(self): + """The regression guard: no tenant id must not change existing behaviour.""" + seen = {"urls": []} + + def handler(request): + seen["urls"].append(str(request.url)) + # The personal list answers with an empty Accounts array. + return httpx.Response(200, json={"code": 0, "data": {"Response": {"Data": {"Accounts": []}}}}) + + with httpx.Client(transport=httpx.MockTransport(handler)) as client: + with _patch_client(handler): + result = fetch_credits(CN_TOKEN, uid="u", domain="www.codebuddy.cn") + + self.assertTrue(seen["urls"], "the personal endpoint must be called") + for url in seen["urls"]: + self.assertIn(credits.RESOURCE_PATH, url) + self.assertNotIn(credits.ENTERPRISE_USAGE_PATH, url) + self.assertFalse(result.get("enterprise")) + + def test_empty_enterprise_id_falls_through_to_personal(self): + seen = {"urls": []} + + def handler(request): + seen["urls"].append(str(request.url)) + return httpx.Response(200, json={"code": 0, "data": {"Response": {"Data": {"Accounts": []}}}}) + + with httpx.Client(transport=httpx.MockTransport(handler)) as client: + with _patch_client(handler): + fetch_credits(CN_TOKEN, uid="u", domain="www.codebuddy.cn", enterprise_id="") + + for url in seen["urls"]: + self.assertNotIn(credits.ENTERPRISE_USAGE_PATH, url) + + +class EnterpriseIdResolutionTests(TestCase): + """Historical credentials may carry a type but no tenant id.""" + + def test_extracts_first_tenant_from_console_accounts(self): + payload = {"code": 0, "data": {"accounts": [{"enterpriseId": "t-1"}, {"enterpriseId": "t-2"}]}} + self.assertEqual(enterprise_id_from_accounts(payload), "t-1") + + def test_accepts_capitalised_accounts_key(self): + payload = {"code": 0, "data": {"Accounts": [{"enterprise_id": "t-9"}]}} + self.assertEqual(enterprise_id_from_accounts(payload), "t-9") + + def test_rejects_failed_or_shapeless_responses(self): + for payload in ({"code": 500, "data": {"accounts": [{"enterpriseId": "t"}]}}, + {"code": 0, "data": {}}, + {"code": 0, "data": {"accounts": []}}, + {"code": 0, "data": {"accounts": ["nope"]}}, + {}, None, "text"): + with self.subTest(payload=payload): + self.assertIsNone(enterprise_id_from_accounts(payload)) + + def test_resolution_returns_id_and_swallows_transport_failure(self): + payload = {"code": 0, "data": {"accounts": [{"enterpriseId": "t-resolved"}]}} + + def handler(request): + return httpx.Response(200, json=payload) + + with httpx.Client(transport=httpx.MockTransport(handler)) as client: + with _patch_client(handler): + resolved = credits.resolve_enterprise_id(CN_TOKEN, uid="u", domain="www.codebuddy.cn") + self.assertEqual(resolved, "t-resolved") + + def failing(_request): + raise httpx.ConnectError("offline") + + with httpx.Client(transport=httpx.MockTransport(failing)) as client: + with _patch_client(failing): + self.assertIsNone(credits.resolve_enterprise_id(CN_TOKEN, uid="u", + domain="www.codebuddy.cn")) + + +class UnlimitedLedgerTests(TestCase): + """The unlimited marker only survives storage when it came from the enterprise dispatch.""" + + def _snapshot(self, td, cred_id, result): + ledger = credits.CreditLedger(Path(td) / "ledger.json") + ledger.update_credits(cred_id, result) + return ledger.snapshot()[cred_id]["credits"] + + def test_unlimited_enterprise_flags_are_stored(self): + with tempfile.TemporaryDirectory() as td: + stored = self._snapshot(td, "ent", {"credits": None, "count": 0, "segments": [], + "soonest_expiry": None, "intl": False, + "unlimited": True, "enterprise": True}) + self.assertIsNone(stored["credits"]) + self.assertTrue(stored["enterprise"]) + self.assertTrue(stored["unlimited"]) + + def test_personal_result_can_never_store_unlimited(self): + # Even a malformed personal result carrying the marker cannot pass the gate. + with tempfile.TemporaryDirectory() as td: + stored = self._snapshot(td, "p1", {"credits": 5.0, "count": 1, "segments": [], + "soonest_expiry": None, "intl": False, + "unlimited": True}) + plain = self._snapshot(td, "p2", {"credits": 5.0, "count": 1, "segments": [], + "soonest_expiry": None, "intl": False}) + self.assertFalse(stored["unlimited"]) + self.assertFalse(stored["enterprise"]) + self.assertFalse(plain["unlimited"]) + self.assertFalse(plain["enterprise"]) + + +class UnlimitedAggregateTests(TestCase): + """Unlimited enterprise tenants stay out of the monetary conversion.""" + + def _entry(self, identity, balance): + return {"identity": identity, "credits": balance} + + def test_unlimited_enterprise_is_counted_not_converted(self): + snap = { + "ent": self._entry("i-ent", {"credits": None, "segments": [], "intl": False, + "partial": False, "unlimited": True, "enterprise": True}), + "p1": self._entry("i-p1", {"credits": 10.0, "intl": False, "partial": False, + "segments": [{"remaining": 10.0, "total": 20.0, + "expires_at": None}]}), + } + agg = credits.aggregate_credits(snap) + self.assertEqual(agg["unlimited_accounts"], 1) + self.assertEqual(agg["remaining"], 10.0) + self.assertEqual(agg["groups"]["domestic"]["remaining"], 10.0) + + def test_marker_without_enterprise_still_counts_numerically(self): + # The double gate: unlimited without enterprise is not honored anywhere. + snap = {"x": self._entry("i-x", {"credits": 3.0, "intl": False, "partial": False, + "unlimited": True, "enterprise": False, + "segments": [{"remaining": 3.0, "total": 3.0, + "expires_at": None}]})} + agg = credits.aggregate_credits(snap) + self.assertEqual(agg["unlimited_accounts"], 0) + self.assertEqual(agg["remaining"], 3.0) + + +class UnlimitedRoutingTests(TestCase): + """An unlimited enterprise tenant is routable; zero-balance restriction still applies otherwise.""" + + def _pool(self, td): + path = Path(td) / "cred.info" + path.write_text(json.dumps({ + "auth": {"accessToken": "t", "expiresAt": int(time.time() * 1000) + 86400000}, + "account": {"uid": "u1"}}), encoding="utf-8") + pool = converter.CredentialPool([path]) + ledger = credits.CreditLedger(Path(td) / "ledger.json") + pool.set_ledger(ledger) + entry = pool.entries()[0] + return pool, ledger, entry, entry["profile"] + + def test_unlimited_enterprise_is_never_zero_balance(self): + with tempfile.TemporaryDirectory() as td: + pool, ledger, entry, profile = self._pool(td) + ledger.update_credits(entry["id"], {"credits": None, "count": 0, "segments": [], + "soonest_expiry": None, "intl": False, + "unlimited": True, "enterprise": True}) + self.assertFalse(pool._zero_balance(entry, profile)) + self.assertTrue(pool._has_credit(entry, profile)) + + def test_zero_balance_personal_account_stays_restricted(self): + with tempfile.TemporaryDirectory() as td: + pool, ledger, entry, profile = self._pool(td) + ledger.update_credits(entry["id"], {"credits": 0, "count": 0, "segments": [], + "soonest_expiry": None, "intl": False}) + self.assertTrue(pool._zero_balance(entry, profile)) + self.assertFalse(pool._has_credit(entry, profile)) + + def test_smuggled_unlimited_marker_keeps_zero_balance_rules(self): + # The ledger gate drops unlimited without enterprise; a number-less personal + # balance is still zero, so the restriction cannot be bypassed. + with tempfile.TemporaryDirectory() as td: + pool, ledger, entry, profile = self._pool(td) + ledger.update_credits(entry["id"], {"credits": None, "count": 0, "segments": [], + "soonest_expiry": None, "intl": False, + "unlimited": True}) + self.assertTrue(pool._zero_balance(entry, profile)) + self.assertFalse(pool._has_credit(entry, profile)) + + +if __name__ == "__main__": + import unittest + unittest.main(verbosity=2) diff --git a/web/src/dashboard.test.tsx b/web/src/dashboard.test.tsx index 3f97df2..344ae8a 100644 --- a/web/src/dashboard.test.tsx +++ b/web/src/dashboard.test.tsx @@ -88,3 +88,29 @@ describe("dashboard granularity", () => { expect(chart.querySelector("polyline")).toBeNull(); }); }); + +describe("official credits", () => { + it("shows an unlimited enterprise tenant without a numeric balance", async () => { + vi.spyOn(api, "get").mockResolvedValue({ + data: { + summary: { requests: 0, success_rate: 1, total_tokens: null, credit: null }, + series: [], + models: [], + profiles: [], + health: { credentials: [] }, + storage: {}, + range: { granularity: "day" }, + official_credits: { + "ent.info": { + name: "ent", + credits: { credits: null, unlimited: true, enterprise: true }, + }, + "p1.info": { name: "p1", credits: { credits: 12.5 } }, + }, + }, + }); + render(); + expect(await screen.findByText("无限")).toBeTruthy(); + expect(screen.getByText("12.5")).toBeTruthy(); + }); +}); diff --git a/web/src/i18n.tsx b/web/src/i18n.tsx index ce26460..c8ed02f 100644 --- a/web/src/i18n.tsx +++ b/web/src/i18n.tsx @@ -156,6 +156,7 @@ const zh: Record = { "dash.storage.shm": "SHM 文件", "dash.storage.credits": "网关已知 Credit", "dash.storage.creditsNote": "官方账号余额", + "dash.unlimited": "无限", "dash.storage.unknown": "未知", "dash.storage.collectDegraded": "采集降级", "dash.credentials": "凭证", @@ -368,6 +369,7 @@ const en: Record = { "dash.storage.shm": "SHM file", "dash.storage.credits": "Gateway-known credits", "dash.storage.creditsNote": "Official account balance", + "dash.unlimited": "Unlimited", "dash.storage.unknown": "Unknown", "dash.storage.collectDegraded": "Collection degraded", "dash.credentials": "Credentials", diff --git a/web/src/pages/Credentials.tsx b/web/src/pages/Credentials.tsx index ece12e1..3fc3d55 100644 --- a/web/src/pages/Credentials.tsx +++ b/web/src/pages/Credentials.tsx @@ -510,7 +510,11 @@ export function Credentials() { )} - {balance ? metric(balance.credits ?? balance.remaining) : "未知"} + {balance + ? balance.unlimited && balance.enterprise + ? "无限" + : metric(balance.credits ?? balance.remaining) + : "未知"} {c.token_expires_at !== undefined ? expiry(c.token_expires_at, true) diff --git a/web/src/pages/Dashboard.tsx b/web/src/pages/Dashboard.tsx index d0d38dd..08a4c6e 100644 --- a/web/src/pages/Dashboard.tsx +++ b/web/src/pages/Dashboard.tsx @@ -87,7 +87,11 @@ function OfficialCredits({ value }: { value: unknown }) { return ( {typeof row.name === "string" ? row.name : key.split(/[\\/]/).at(-1)} - {metric(credits?.credits ?? credits?.remaining)} + + {credits?.unlimited && credits?.enterprise + ? t("dash.unlimited") + : metric(credits?.credits ?? credits?.remaining)} + {row.error ? ( {t("dash.status.syncError")}