Dispatch enterprise accounts to the enterprise allowance endpoint

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":

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"

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}"

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()