clawdmeter/daemon/providers/zai.py

"""z.ai (GLM Coding Plan) provider (v3 M3).

z.ai exposes a dedicated usage endpoint that returns the Coding Plan's rate-limit
windows directly — no need to send a (quota-consuming) chat request:

    GET https://api.z.ai/api/monitor/usage/quota/limit
    Authorization: <api_key>          # raw key, NOT "Bearer …"
    Accept-Language: en-US,en

    {"code":200,"success":true,"data":{"level":"lite","limits":[
        {"type":"TIME_LIMIT",  "unit":5,"number":1, ...},                 # monthly web-tool quota (ignored)
        {"type":"TOKENS_LIMIT","unit":3,"number":5,"percentage":1,"nextResetTime":<ms>},  # 5h window
        {"type":"TOKENS_LIMIT","unit":6,"number":1,"percentage":2,"nextResetTime":<ms>}]}}  # weekly window

We surface the two TOKENS_LIMIT windows as Claude's short (5h) + weekly bars: of
the token windows, the shortest is the 5h bar and the longest the weekly bar.
``percentage`` is already 0-100; ``nextResetTime`` is epoch ms. The key + base URL
come from the config the control panel's z.ai field writes. This is a status GET,
so — unlike a chat call — it doesn't spend the prompt-metered plan, and can poll
on the normal cadence.

Discovered from github.com/rygel/AIUsageTracker (ZaiProvider.cs).
"""

from __future__ import annotations

import time
from urllib.parse import urlsplit

import httpx

from .base import Provider, ProviderStatus

DEFAULT_HOST = "https://api.z.ai"
_MONITOR_PATH = "/api/monitor/usage/quota/limit"

# z.ai window "unit" enum → seconds, used only to rank the token windows by length
# (shortest = the 5h bar, longest = the weekly bar). Best-effort; unknown → hours.
_UNIT_SECONDS = {1: 1, 2: 60, 3: 3600, 4: 86400, 5: 2_592_000, 6: 604_800}


def _monitor_url(base_url: str) -> str:
    """The usage endpoint on the same host as the configured (Anthropic-compat)
    base URL, or the public z.ai host when unset."""
    parts = urlsplit(base_url or "")
    if parts.scheme and parts.netloc:
        return f"{parts.scheme}://{parts.netloc}{_MONITOR_PATH}"
    return DEFAULT_HOST + _MONITOR_PATH


class ZaiProvider(Provider):
    id = "zai"
    label = "z.ai GLM"
    accent = "3859ff"                # brand blue

    def __init__(self, transport: httpx.BaseTransport | None = None) -> None:
        self._transport = transport   # tests inject an httpx.MockTransport

    async def poll(self) -> ProviderStatus:
        from daemon.claude_usage_daemon_windows import log
        try:
            from daemon.config import load_config, provider_conf
        except ImportError:
            from config import load_config, provider_conf

        conf = provider_conf(load_config(), "zai")
        key = (conf.get("api_key") or "").strip()
        base = (conf.get("base_url") or "").strip()
        if not key:
            return ProviderStatus(ok=False, st="no key", auth_problem=True)

        url = _monitor_url(base)
        headers = {"Authorization": key, "Accept-Language": "en-US,en"}
        client_kw = {"timeout": 20.0}
        if self._transport is not None:
            client_kw["transport"] = self._transport
        try:
            async with httpx.AsyncClient(**client_kw) as http:
                resp = await http.get(url, headers=headers)
        except httpx.HTTPError as e:
            log(f"z.ai call failed: {e}")               # transient — retry next tick
            return ProviderStatus(ok=False, st="offline")
        if resp.status_code in (401, 403):
            log(f"z.ai auth rejected: HTTP {resp.status_code}")
            return ProviderStatus(ok=False, st="bad key", auth_problem=True)
        if resp.status_code >= 400:
            log(f"z.ai HTTP {resp.status_code}: {resp.text[:150]}")
            return ProviderStatus(ok=False, st=f"http {resp.status_code}")
        try:
            body = resp.json()
        except ValueError:
            log("z.ai: non-JSON usage response")
            return ProviderStatus(ok=False, st="error")
        if not body.get("success") or not isinstance(body.get("data"), dict):
            # e.g. {"code":401,...,"success":false} → treat as an auth/config problem
            code = body.get("code")
            log(f"z.ai usage error: code={code} msg={body.get('msg')!r}")
            return ProviderStatus(ok=False, st="bad key",
                                  auth_problem=code in (401, 403))
        return self._to_status(body["data"])

    @staticmethod
    def _to_status(data: dict) -> ProviderStatus:
        limits = data.get("limits") or []
        # Only the token windows are the coding rate limit; TIME_LIMIT is the
        # separate monthly web-tool quota.
        toks = [l for l in limits if isinstance(l, dict) and l.get("type") == "TOKENS_LIMIT"]

        def win_seconds(l: dict) -> int:
            return _UNIT_SECONDS.get(l.get("unit"), 3600) * int(l.get("number") or 1)

        def pct(l: dict | None) -> float:
            if not l:
                return 0.0
            try:
                return round(float(l.get("percentage") or 0.0), 1)
            except (TypeError, ValueError):
                return 0.0

        def reset_min(l: dict | None) -> int:
            ts = l.get("nextResetTime") if l else None
            if not isinstance(ts, (int, float)):
                return -1
            if ts > 1e10:                # epoch ms → s
                ts = ts / 1000.0
            m = (ts - time.time()) / 60.0
            return int(round(m)) if m > 0 else -1

        short = min(toks, key=win_seconds) if toks else None
        weekly = max(toks, key=win_seconds) if len(toks) > 1 else None

        level = str(data.get("level") or "").strip()
        st = level.capitalize() if level else "allowed"
        if pct(short) >= 100 or pct(weekly) >= 100:
            st = "limited"

        if not toks:                     # authenticated but no token windows reported
            return ProviderStatus(ok=True, st=st or "allowed")
        return ProviderStatus(
            s=pct(short), sr=reset_min(short),
            w=pct(weekly), wr=reset_min(weekly),
            st=st, ok=True,
        )