diff --git a/integrations/hermes/quota.py b/integrations/hermes/quota.py new file mode 100644 index 0000000..00984ea --- /dev/null +++ b/integrations/hermes/quota.py @@ -0,0 +1,359 @@ +"""Pure-Python ZeroAPI quota normalization and pressure policy. + +Provider HTTP/RPC parsing and credentials remain host-owned. This module +accepts only token-free, provider-neutral quantitative windows and mirrors +plugin/quota-normalize.ts + plugin/quota-policy.ts. +""" + +from __future__ import annotations + +import math +import re +from datetime import datetime +from typing import Any, TypeGuard + +VALID_STATUSES = { + "fresh", + "stale", + "auth_expired", + "rate_limited", + "network_error", + "invalid_response", + "unsupported", +} +VALID_WINDOW_KINDS = { + "tokens_limit", + "requests_limit", + "credits", + "messages", + "compute", + "time_limit", + "percent", +} +VALID_APPLICABILITY = {"inference", "mcp", "model"} +TIMESTAMP_RE = re.compile( + r"^\d{4}-\d{2}-\d{2}T\d{2}:\d{2}:\d{2}(?:\.\d{1,9})?(?:Z|[+-]\d{2}:\d{2})$" +) + + +def _is_number(value: Any) -> TypeGuard[int | float]: + return ( + not isinstance(value, bool) + and isinstance(value, (int, float)) + and math.isfinite(value) + ) + + +def _assert_finite_number(value: Any, label: str) -> None: + if isinstance(value, bool): + raise TypeError(f"{label} must be numeric, got boolean") + if not isinstance(value, (int, float)): + raise TypeError(f"{label} must be numeric") + try: + if math.isnan(value): + raise ValueError(f"{label} must not be NaN") + if math.isinf(value): + raise ValueError(f"{label} must be finite") + except OverflowError: + raise ValueError(f"{label} is too large to represent as a finite number") + + +def _assert_valid_ratio(value: Any) -> None: + _assert_finite_number(value, "remainingRatio") + if value < 0 or value > 1: + raise ValueError("remainingRatio must be in [0, 1]") + + +def _normalize_percentage(value: Any, label: str) -> float: + _assert_finite_number(value, label) + if value < 0 or value > 100: + raise ValueError(f"{label} must be in [0, 100]") + ratio = value / 100 + _assert_valid_ratio(ratio) + return ratio + + +def _normalize_identifier(value: Any, label: str) -> str: + if not isinstance(value, str): + raise TypeError(f"{label} must be a string") + normalized = value.strip() + if not normalized: + raise ValueError(f"{label} must be non-empty") + if len(normalized) > 256: + raise ValueError(f"{label} is too long") + if any(ord(char) < 32 or ord(char) == 127 for char in normalized): + raise ValueError(f"{label} contains control characters") + return normalized + + +def _normalize_timestamp(value: Any, label: str) -> str: + normalized = _normalize_identifier(value, label) + if TIMESTAMP_RE.fullmatch(normalized) is None: + raise ValueError(f"{label} must be an ISO-8601 timestamp with timezone") + try: + datetime.fromisoformat(normalized.replace("Z", "+00:00")) + except ValueError as exc: + raise ValueError(f"{label} must contain a valid calendar date and time") from exc + return normalized + + +def _map_window_kind(raw_kind: str) -> str: + upper = raw_kind.upper() + if "TOKEN" in upper: + return "tokens_limit" + if "REQUEST" in upper or "RPM" in upper: + return "requests_limit" + if "CREDIT" in upper: + return "credits" + if "MESSAGE" in upper: + return "messages" + if "TIME_LIMIT" in upper: + return "time_limit" + if "COMPUTE" in upper or "TIME" in upper: + return "compute" + if "PERCENT" in upper or upper in {"USAGE", "BILLING"}: + return "percent" + return "tokens_limit" + + +def normalize_window( + raw_kind: Any, + *, + id: str | None = None, + remaining_ratio: float | None = None, + used: float | None = None, + limit: float | None = None, + percentage_remaining: float | None = None, + percentage_used: float | None = None, + applies_to: str = "inference", + model_ids: list[str] | None = None, + window_seconds: float | None = None, + reset_at: str | None = None, + explicit_zero_usage: bool = False, +) -> dict[str, Any]: + normalized_kind = _normalize_identifier(raw_kind, "raw_kind") + window_id = _normalize_identifier(id if id is not None else normalized_kind, "window id") + kind = _map_window_kind(normalized_kind) + + if explicit_zero_usage is True: + ratio = 1.0 + elif remaining_ratio is not None: + _assert_valid_ratio(remaining_ratio) + ratio = float(remaining_ratio) + elif percentage_remaining is not None: + ratio = _normalize_percentage(percentage_remaining, "percentage_remaining") + elif percentage_used is not None: + ratio = max(0.0, 1.0 - _normalize_percentage(percentage_used, "percentage_used")) + elif used is not None and limit is not None: + _assert_finite_number(used, "used") + _assert_finite_number(limit, "limit") + if used < 0: + raise ValueError("used must be non-negative") + if limit <= 0: + raise ValueError("limit must be positive") + ratio = max(0.0, 1.0 - used / limit) + _assert_valid_ratio(ratio) + else: + raise ValueError(f"cannot derive remainingRatio for window '{window_id}'") + + if applies_to not in {"inference", "mcp", "model"}: + raise ValueError(f"unknown applies_to value '{applies_to}'") + + if model_ids is not None and not isinstance(model_ids, list): + raise TypeError("modelIds must be a list") + mids = [_normalize_identifier(model_id, "modelId") for model_id in (model_ids or [])] + if len(set(mids)) != len(mids): + raise ValueError("modelIds must be unique") + if applies_to == "model" and not mids: + raise ValueError("model-scoped window requires at least one modelId") + if applies_to != "model" and mids: + raise ValueError("non-model window must not carry modelIds") + + normalized_seconds: float | None = None + if window_seconds is not None: + _assert_finite_number(window_seconds, "windowSeconds") + if window_seconds <= 0: + raise ValueError("windowSeconds must be positive") + normalized_seconds = float(window_seconds) + + normalized_reset = None if reset_at is None else _normalize_timestamp(reset_at, "resetAt") + + return { + "id": window_id, + "kind": kind, + "appliesTo": applies_to, + "modelIds": mids, + "remainingRatio": ratio, + **({"windowSeconds": normalized_seconds} if normalized_seconds is not None else {}), + **({"resetAt": normalized_reset} if normalized_reset is not None else {}), + } + + +def validate_snapshot( + snapshot: dict[str, Any], + expected_provider: str | None = None, + diagnostics_only: bool = False, +) -> None: + provider = _normalize_identifier(snapshot.get("provider"), "snapshot provider") + _normalize_identifier(snapshot.get("account"), "snapshot account") + _normalize_timestamp(snapshot.get("fetchedAt"), "fetchedAt") + + status = snapshot.get("status") + if status not in VALID_STATUSES: + raise ValueError(f"unknown snapshot status '{status}'") + if expected_provider is not None and provider != expected_provider: + raise ValueError( + f'snapshot provider "{provider}" does not match expected "{expected_provider}"' + ) + if not diagnostics_only and status != "fresh": + raise ValueError(f'routing snapshot must be fresh, got "{status}"') + + windows = snapshot.get("windows") + if not isinstance(windows, list): + raise TypeError("snapshot windows must be a list") + if status == "fresh" and not windows: + raise ValueError("fresh snapshot requires at least one window") + + ids: set[str] = set() + for window in windows: + if not isinstance(window, dict): + raise TypeError("window must be an object") + _assert_valid_ratio(window.get("remainingRatio")) + window_id = _normalize_identifier(window.get("id"), "window id") + if window.get("kind") not in VALID_WINDOW_KINDS: + raise ValueError(f'unknown window kind "{window.get("kind")}"') + applies_to = window.get("appliesTo") + if applies_to not in VALID_APPLICABILITY: + raise ValueError(f'unknown appliesTo value "{applies_to}"') + model_ids = window.get("modelIds") + if not isinstance(model_ids, list): + raise TypeError("modelIds must be a list") + normalized_models = [_normalize_identifier(model_id, "modelId") for model_id in model_ids] + if any(nm != mid for nm, mid in zip(normalized_models, model_ids)): + raise ValueError("modelIds must be pre-canonicalized") + if len(set(normalized_models)) != len(normalized_models): + raise ValueError("modelIds must be unique") + if applies_to == "model" and not normalized_models: + raise ValueError("model-scoped window requires at least one modelId") + if applies_to != "model" and normalized_models: + raise ValueError("non-model window must not carry modelIds") + if "windowSeconds" in window: + seconds = window["windowSeconds"] + _assert_finite_number(seconds, "windowSeconds") + if seconds <= 0: + raise ValueError("windowSeconds must be positive") + if "resetAt" in window: + _normalize_timestamp(window["resetAt"], "resetAt") + if window_id in ids: + raise ValueError(f'duplicate window id "{window_id}"') + ids.add(window_id) + + +def _input_window(item: Any) -> dict[str, Any]: + if not isinstance(item, dict): + raise TypeError("window must be an object") + for key in ( + "rawKind", "id", "remainingRatio", "used", "limit", + "percentageRemaining", "percentageUsed", "appliesTo", + "modelIds", "windowSeconds", "resetAt", + ): + if key in item and item[key] is None: + raise TypeError(f"{key} must not be null") + return normalize_window( + item.get("rawKind"), + id=item.get("id"), + remaining_ratio=item.get("remainingRatio"), + used=item.get("used"), + limit=item.get("limit"), + percentage_remaining=item.get("percentageRemaining"), + percentage_used=item.get("percentageUsed"), + applies_to=item.get("appliesTo", "inference"), + model_ids=item.get("modelIds"), + window_seconds=item.get("windowSeconds"), + reset_at=item.get("resetAt"), + explicit_zero_usage=item.get("explicitZeroUsage", False), + ) + + +def normalize_snapshot(payload: dict[str, Any]) -> dict[str, Any]: + provider = _normalize_identifier(payload.get("provider"), "provider") + account = _normalize_identifier(payload.get("account"), "account") + fetched_at = _normalize_timestamp(payload.get("fetchedAt"), "fetchedAt") + requested_status = payload.get("status", "fresh") + status = ( + requested_status + if isinstance(requested_status, str) and requested_status in VALID_STATUSES + else "invalid_response" + ) + + windows: list[dict[str, Any]] = [] + try: + raw_windows = payload.get("windows") + if not isinstance(raw_windows, list): + raise TypeError("windows must be a list") + windows = [_input_window(item) for item in raw_windows] + ids = {window["id"] for window in windows} + if len(ids) != len(windows): + raise ValueError("window IDs must be unique") + except (KeyError, TypeError, ValueError, OverflowError): + windows = [] + status = "invalid_response" + + if status == "fresh" and not windows: + status = "unsupported" + + snapshot = { + "provider": provider, + "account": account, + "status": status, + "windows": windows, + "fetchedAt": fetched_at, + } + validate_snapshot(snapshot, diagnostics_only=True) + return snapshot + + +# ── Routing policy ────────────────────────────────────────────────────────── + +def applicable_windows(snapshot: dict[str, Any] | None, model: str) -> list[dict[str, Any]]: + if not snapshot or snapshot.get("status") != "fresh": + return [] + result = [] + for window in snapshot.get("windows", []): + if window["appliesTo"] == "inference": + result.append(window) + elif window["appliesTo"] == "model" and model in window.get("modelIds", []): + result.append(window) + return result + + +def account_headroom(snapshot: dict[str, Any] | None, model: str) -> float | None: + if not snapshot or snapshot.get("status") != "fresh": + return None + try: + validate_snapshot(snapshot) + except (KeyError, TypeError, ValueError): + return None + windows = applicable_windows(snapshot, model) + if not windows: + return None + return min(window["remainingRatio"] for window in windows) + + +def compute_quota_factor(snapshot: dict[str, Any] | None, model: str) -> float | None: + headroom = account_headroom(snapshot, model) + if headroom is None: + return None + return math.sqrt(headroom) + + +def compute_live_pressure( + tier_weight: float, + provider_bias: float, + snapshot: dict[str, Any] | None, + model: str, +) -> float | None: + factor = compute_quota_factor(snapshot, model) + if factor is None: + return None + return tier_weight * provider_bias * factor diff --git a/integrations/hermes/test_quota.py b/integrations/hermes/test_quota.py new file mode 100644 index 0000000..d4aa613 --- /dev/null +++ b/integrations/hermes/test_quota.py @@ -0,0 +1,426 @@ +"""Tests for ZeroAPI live quota normalization and policy (Python parity).""" + +import math +import unittest + +from quota import ( + normalize_window, + validate_snapshot, + normalize_snapshot, + compute_quota_factor, + compute_live_pressure, + applicable_windows, + account_headroom, +) + + +class TestNormalizeWindow(unittest.TestCase): + + def test_normalizes_tokens_limit_with_percentage(self): + w = normalize_window("TOKENS_LIMIT", remaining_ratio=0.9888, window_seconds=5 * 3600) + self.assertAlmostEqual(w["remainingRatio"], 0.9888) + self.assertEqual(w["appliesTo"], "inference") + + def test_classifies_time_limit_before_generic_time(self): + w = normalize_window("TIME_LIMIT", remaining_ratio=0.75, applies_to="mcp") + self.assertEqual(w["kind"], "time_limit") + self.assertEqual(w["appliesTo"], "mcp") + + def test_derives_ratio_from_usage_limit(self): + w = normalize_window("PRIMARY", used=400, limit=800) + self.assertAlmostEqual(w["remainingRatio"], 0.5) + + def test_rejects_nan(self): + with self.assertRaises(ValueError): + normalize_window("X", remaining_ratio=float("nan")) + + def test_rejects_infinity(self): + with self.assertRaises(ValueError): + normalize_window("X", remaining_ratio=float("inf")) + + def test_rejects_boolean(self): + with self.assertRaises(TypeError): + normalize_window("X", remaining_ratio=True) + + def test_percentage_fields_use_zero_to_one_hundred_scale(self): + used = normalize_window("PRIMARY", percentage_used=1) + remaining = normalize_window("PRIMARY", percentage_remaining=1) + self.assertAlmostEqual(used["remainingRatio"], 0.99) + self.assertAlmostEqual(remaining["remainingRatio"], 0.01) + + def test_rejects_over_one(self): + with self.assertRaises(ValueError): + normalize_window("X", remaining_ratio=1.01) + + def test_rejects_negative(self): + with self.assertRaises(ValueError): + normalize_window("X", remaining_ratio=-0.01) + + def test_model_scoped_requires_model_ids(self): + with self.assertRaises(ValueError): + normalize_window("M", applies_to="model", remaining_ratio=0.5) + + def test_non_model_rejects_model_ids(self): + with self.assertRaises(ValueError): + normalize_window("M", applies_to="inference", model_ids=["m1"], remaining_ratio=0.5) + + +class TestValidateSnapshot(unittest.TestCase): + + def _valid(self): + return { + "provider": "zai", "account": "zai#1", "status": "fresh", + "windows": [{"id": "P", "kind": "tokens_limit", "appliesTo": "inference", "modelIds": [], "remainingRatio": 0.86}], + "fetchedAt": "2026-07-24T17:00:00Z", + } + + def test_accepts_valid_fresh(self): + validate_snapshot(self._valid()) + + def test_rejects_no_windows_fresh(self): + with self.assertRaises(ValueError): + validate_snapshot({**self._valid(), "windows": []}) + + def test_rejects_stale_not_diagnostic(self): + with self.assertRaises(ValueError): + validate_snapshot({**self._valid(), "status": "stale"}, diagnostics_only=False) + + def test_accepts_stale_diagnostic(self): + validate_snapshot({**self._valid(), "status": "stale"}, diagnostics_only=True) + + def test_rejects_provider_mismatch(self): + with self.assertRaises(ValueError): + validate_snapshot(self._valid(), expected_provider="openai") + + def test_rejects_invalid_timestamp(self): + for fetched_at in ( + "1", + "07/24/2026", + "2026-07-24T17:00:00", + "2026-02-30T00:00:00Z", + ): + with self.subTest(fetched_at=fetched_at), self.assertRaises(ValueError): + validate_snapshot({**self._valid(), "fetchedAt": fetched_at}) + + +class TestNormalizeSnapshot(unittest.TestCase): + + def test_zai_payload(self): + snap = normalize_snapshot({ + "provider": "zai", "account": "zai#1", + "windows": [ + { + "id": "5h", "rawKind": "TOKENS_LIMIT", + "used": 112, "limit": 10000, + "resetAt": "2026-07-24T20:23:52Z", + }, + { + "id": "weekly", "rawKind": "TOKENS_LIMIT", + "used": 1400, "limit": 10000, + "resetAt": "2026-07-26T20:23:52Z", + }, + ], + "fetchedAt": "2026-07-24T17:00:00Z", + }) + self.assertEqual(snap["status"], "fresh") + self.assertEqual(len(snap["windows"]), 2) + self.assertAlmostEqual(snap["windows"][0]["remainingRatio"], 0.9888) + + def test_codex_payload(self): + snap = normalize_snapshot({ + "provider": "openai-codex", "account": "openai#1", + "windows": [ + { + "id": "primary", "rawKind": "TOKENS_LIMIT", + "windowSeconds": 300 * 60, "percentageUsed": 47, + }, + { + "id": "secondary", "rawKind": "TOKENS_LIMIT", + "windowSeconds": 10080 * 60, "percentageUsed": 12, + }, + ], + "fetchedAt": "2026-07-24T17:00:00Z", + }) + self.assertEqual(len(snap["windows"]), 2) + self.assertAlmostEqual(snap["windows"][0]["remainingRatio"], 0.53, places=1) + + def test_xai_payload(self): + snap = normalize_snapshot({ + "provider": "xai", "account": "xai#1", + "windows": [ + { + "id": "billing", "rawKind": "BILLING", + "percentageRemaining": 100, + }, + ], + "fetchedAt": "2026-07-24T17:00:00Z", + }) + self.assertEqual(len(snap["windows"]), 1) + self.assertEqual(snap["windows"][0]["remainingRatio"], 1.0) + + def test_unsupported_payload(self): + snap = normalize_snapshot({ + "provider": "qwen-oauth", "account": "qwen#1", + "windows": [], + "fetchedAt": "2026-07-24T17:00:00Z", + }) + self.assertEqual(snap["status"], "unsupported") + + def test_raw_shaped_fields_are_ignored(self): + snap = normalize_snapshot({ + "provider": "qwen-oauth", "account": "qwen#1", + "windows": [], + "raw": {"remains": {"percentage": 90, "plan_type": "NOT_MINIMAX"}}, + "fetchedAt": "2026-07-24T17:00:00Z", + }) + self.assertEqual(snap["status"], "unsupported") + self.assertNotIn("raw", snap) + + def test_malformed_kind_fails_closed(self): + snap = normalize_snapshot({ + "provider": "zai", "account": "zai#1", + "windows": [{"id": "bad", "rawKind": 1, "remainingRatio": 0.5}], + "fetchedAt": "2026-07-24T17:00:00Z", + }) + self.assertEqual(snap["status"], "invalid_response") + self.assertEqual(snap["windows"], []) + + def test_rejects_boolean_zai_counters(self): + snap = normalize_snapshot({ + "provider": "zai", "account": "zai#1", + "windows": [{ + "id": "5h", "rawKind": "TOKENS_LIMIT", + "used": False, "limit": True, + }], + "fetchedAt": "2026-07-24T17:00:00Z", + }) + self.assertEqual(snap["status"], "invalid_response") + + def test_rejects_boolean_kimi_counters(self): + snap = normalize_snapshot({ + "provider": "moonshot", "account": "kimi#1", + "windows": [{ + "id": "weekly", "rawKind": "TOKENS_LIMIT", + "used": False, "limit": True, + }], + "fetchedAt": "2026-07-24T17:00:00Z", + }) + self.assertEqual(snap["status"], "invalid_response") + + def test_rejects_boolean_minimax_counters(self): + snap = normalize_snapshot({ + "provider": "minimax-portal", "account": "minimax#1", + "windows": [{ + "id": "coding", "rawKind": "TOKENS_LIMIT", + "used": False, "limit": True, + }], + "fetchedAt": "2026-07-24T17:00:00Z", + }) + self.assertEqual(snap["status"], "invalid_response") + + def test_strips_secret_fields(self): + snap = normalize_snapshot({ + "provider": "zai", "account": "zai#1", + "windows": [{ + "id": "5h", "rawKind": "TOKENS_LIMIT", + "percentageRemaining": 98.88, + }], + "account_email": "secret@example.com", + "access_token": "test-only-secret-placeholder", + "fetchedAt": "2026-07-24T17:00:00Z", + }) + import json + serialized = json.dumps(snap) + self.assertNotIn("secret@example.com", serialized) + self.assertNotIn("test-only-secret-placeholder", serialized) + + def test_duplicate_window_ids_fail_closed(self): + snap = normalize_snapshot({ + "provider": "zai", "account": "zai#1", + "windows": [ + {"id": "weekly", "rawKind": "TOKENS_LIMIT", "remainingRatio": 0.8}, + {"id": "weekly", "rawKind": "TOKENS_LIMIT", "remainingRatio": 0.7}, + ], + "fetchedAt": "2026-07-24T17:00:00Z", + }) + self.assertEqual(snap["status"], "invalid_response") + self.assertEqual(snap["windows"], []) + + def test_unhashable_status_fails_closed(self): + snap = normalize_snapshot({ + "provider": "zai", "account": "zai#1", + "status": ["fresh"], + "windows": [{"id": "w", "rawKind": "TOKENS_LIMIT", "remainingRatio": 0.5}], + "fetchedAt": "2026-07-24T17:00:00Z", + }) + self.assertEqual(snap["status"], "invalid_response") + + def test_dict_status_fails_closed(self): + snap = normalize_snapshot({ + "provider": "zai", "account": "zai#1", + "status": {"evil": True}, + "windows": [{"id": "w", "rawKind": "TOKENS_LIMIT", "remainingRatio": 0.5}], + "fetchedAt": "2026-07-24T17:00:00Z", + }) + self.assertEqual(snap["status"], "invalid_response") + + def test_oversized_integer_fails_closed(self): + snap = normalize_snapshot({ + "provider": "zai", "account": "zai#1", + "windows": [{ + "id": "w", "rawKind": "TOKENS_LIMIT", + "remainingRatio": 10 ** 10000, + }], + "fetchedAt": "2026-07-24T17:00:00Z", + }) + self.assertEqual(snap["status"], "invalid_response") + self.assertEqual(snap["windows"], []) + + def test_oversized_integer_in_used_fails_closed(self): + snap = normalize_snapshot({ + "provider": "zai", "account": "zai#1", + "windows": [{ + "id": "w", "rawKind": "TOKENS_LIMIT", + "used": 10 ** 10000, "limit": 10 ** 10001, + }], + "fetchedAt": "2026-07-24T17:00:00Z", + }) + self.assertEqual(snap["status"], "invalid_response") + + def test_non_list_model_ids_fails_closed(self): + snap = normalize_snapshot({ + "provider": "zai", "account": "zai#1", + "windows": [{ + "id": "w", "rawKind": "TOKENS_LIMIT", + "remainingRatio": 0.5, "modelIds": {"fake-model": True}, + }], + "fetchedAt": "2026-07-24T17:00:00Z", + }) + self.assertEqual(snap["status"], "invalid_response") + self.assertEqual(snap["windows"], []) + + def test_false_model_ids_fails_closed(self): + snap = normalize_snapshot({ + "provider": "zai", "account": "zai#1", + "windows": [{ + "id": "w", "rawKind": "TOKENS_LIMIT", + "remainingRatio": 0.5, "appliesTo": "model", "modelIds": False, + }], + "fetchedAt": "2026-07-24T17:00:00Z", + }) + self.assertEqual(snap["status"], "invalid_response") + + def test_null_remaining_ratio_fails_closed(self): + snap = normalize_snapshot({ + "provider": "zai", "account": "zai#1", + "windows": [{ + "id": "w", "rawKind": "TOKENS_LIMIT", + "remainingRatio": None, "used": 0, "limit": 100, + }], + "fetchedAt": "2026-07-24T17:00:00Z", + }) + self.assertEqual(snap["status"], "invalid_response") + + def test_null_field_in_window_fails_closed(self): + snap = normalize_snapshot({ + "provider": "zai", "account": "zai#1", + "windows": [{ + "id": "w", "rawKind": None, + "remainingRatio": 0.5, + }], + "fetchedAt": "2026-07-24T17:00:00Z", + }) + self.assertEqual(snap["status"], "invalid_response") + + def test_whitespace_model_id_in_validate_rejected(self): + from quota import validate_snapshot + snap = { + "provider": "zai", "account": "zai#1", "status": "fresh", + "windows": [{"id": "w", "kind": "tokens_limit", "appliesTo": "model", "modelIds": [" openai/gpt-5.6-sol "], "remainingRatio": 0.5}], + "fetchedAt": "2026-07-24T17:00:00Z", + } + with self.assertRaises(ValueError): + validate_snapshot(snap) + + def test_explicit_zero_usage_means_full_quota(self): + snap = normalize_snapshot({ + "provider": "zai", "account": "zai#1", + "windows": [{ + "id": "w", "rawKind": "TOKENS_LIMIT", + "explicitZeroUsage": True, + }], + "fetchedAt": "2026-07-24T17:00:00Z", + }) + self.assertEqual(snap["status"], "fresh") + self.assertEqual(len(snap["windows"]), 1) + self.assertAlmostEqual(snap["windows"][0]["remainingRatio"], 1.0) + + +class TestQuotaPolicy(unittest.TestCase): + + def test_compute_quota_factor(self): + snap = { + "provider": "zai", "account": "zai#1", "status": "fresh", + "windows": [ + {"id": "5h", "kind": "tokens_limit", "appliesTo": "inference", "modelIds": [], "remainingRatio": 0.9888}, + {"id": "1w", "kind": "tokens_limit", "appliesTo": "inference", "modelIds": [], "remainingRatio": 0.86}, + ], + "fetchedAt": "2026-07-24T17:00:00Z", + } + self.assertAlmostEqual(compute_quota_factor(snap, "zai/glm-5.2"), math.sqrt(0.86), places=4) + + def test_stale_returns_none(self): + snap = {"provider": "zai", "account": "zai#1", "status": "stale", "windows": [], "fetchedAt": "2026-07-24T17:00:00Z"} + self.assertIsNone(compute_quota_factor(snap, "zai/glm-5.2")) + + def test_malformed_fresh_returns_none(self): + snap = { + "provider": "zai", "account": "zai#1", "status": "fresh", + "windows": [{"id": "5h", "kind": "tokens_limit", "appliesTo": "inference", "modelIds": [], "remainingRatio": float("nan")}], + "fetchedAt": "2026-07-24T17:00:00Z", + } + self.assertIsNone(compute_quota_factor(snap, "zai/glm-5.2")) + + def test_depleted_returns_zero(self): + snap = { + "provider": "zai", "account": "zai#1", "status": "fresh", + "windows": [{"id": "5h", "kind": "tokens_limit", "appliesTo": "inference", "modelIds": [], "remainingRatio": 0.0}], + "fetchedAt": "2026-07-24T17:00:00Z", + } + self.assertEqual(compute_quota_factor(snap, "zai/glm-5.2"), 0.0) + + def test_mcp_excluded_from_inference(self): + snap = { + "provider": "zai", "account": "zai#1", "status": "fresh", + "windows": [ + {"id": "MCP", "kind": "time_limit", "appliesTo": "mcp", "modelIds": [], "remainingRatio": 0.0}, + {"id": "5h", "kind": "tokens_limit", "appliesTo": "inference", "modelIds": [], "remainingRatio": 0.80}, + ], + "fetchedAt": "2026-07-24T17:00:00Z", + } + self.assertAlmostEqual(compute_quota_factor(snap, "zai/glm-5.2"), math.sqrt(0.80), places=4) + + def test_model_scoped_only_affects_mapped(self): + snap = { + "provider": "mm", "account": "mm#1", "status": "fresh", + "windows": [ + {"id": "INF", "kind": "tokens_limit", "appliesTo": "inference", "modelIds": [], "remainingRatio": 0.90}, + {"id": "M25", "kind": "tokens_limit", "appliesTo": "model", "modelIds": ["mm/m2.5"], "remainingRatio": 0.10}, + ], + "fetchedAt": "2026-07-24T17:00:00Z", + } + self.assertAlmostEqual(compute_quota_factor(snap, "mm/m2.5"), math.sqrt(0.10), places=4) + self.assertAlmostEqual(compute_quota_factor(snap, "mm/m2.7"), math.sqrt(0.90), places=4) + + def test_compute_live_pressure(self): + snap = { + "provider": "zai", "account": "zai#1", "status": "fresh", + "windows": [{"id": "1w", "kind": "tokens_limit", "appliesTo": "inference", "modelIds": [], "remainingRatio": 0.86}], + "fetchedAt": "2026-07-24T17:00:00Z", + } + result = compute_live_pressure(5.0, 1.25, snap, "zai/glm-5.2") + self.assertAlmostEqual(result, 5.0 * 1.25 * math.sqrt(0.86), places=4) + + +if __name__ == "__main__": + unittest.main() diff --git a/plugin/__tests__/quota-normalize.test.ts b/plugin/__tests__/quota-normalize.test.ts new file mode 100644 index 0000000..55a69b6 --- /dev/null +++ b/plugin/__tests__/quota-normalize.test.ts @@ -0,0 +1,486 @@ +import { describe, it, expect } from "vitest"; +import { + normalizeQuotaWindow, + validateNormalizedSnapshot, + normalizeSnapshot, +} from "../quota-normalize.js"; +import type { ProviderQuotaPayload } from "../quota-types.js"; + +describe("normalizeQuotaWindow", () => { + it("normalizes a Z.AI TOKENS_LIMIT with percentage and nextResetTime", () => { + const window = normalizeQuotaWindow({ + rawKind: "TOKENS_LIMIT", + windowSeconds: 5 * 3600, + remainingRatio: 0.9888, + resetAt: "2026-07-24T20:23:52Z", + }); + expect(window.remainingRatio).toBeCloseTo(0.9888); + expect(window.kind).toBe("tokens_limit"); + expect(window.appliesTo).toBe("inference"); + expect(window.modelIds).toEqual([]); + expect(window.id).toBe("TOKENS_LIMIT"); + }); + + it("classifies TIME_LIMIT before generic TIME kinds", () => { + const window = normalizeQuotaWindow({ + rawKind: "TIME_LIMIT", + remainingRatio: 0.75, + appliesTo: "mcp", + }); + expect(window.kind).toBe("time_limit"); + expect(window.appliesTo).toBe("mcp"); + }); + + it("derives remaining ratio from usage/limit counters", () => { + const window = normalizeQuotaWindow({ + rawKind: "PRIMARY", + windowSeconds: 5 * 3600, + used: 400, + limit: 800, + resetAt: "2026-07-24T20:23:52Z", + }); + expect(window.remainingRatio).toBeCloseTo(0.5); + }); + + it("rejects boolean usage counters", () => { + expect(() => + normalizeQuotaWindow({ + rawKind: "PRIMARY", + used: false as unknown as number, + limit: true as unknown as number, + }), + ).toThrow(); + }); + + it("treats percentage fields as 0-100 values, including exactly one percent", () => { + expect(normalizeQuotaWindow({ rawKind: "PRIMARY", percentageUsed: 1 }).remainingRatio) + .toBeCloseTo(0.99); + expect(normalizeQuotaWindow({ rawKind: "PRIMARY", percentageRemaining: 1 }).remainingRatio) + .toBeCloseTo(0.01); + }); + + it("rejects NaN remainingRatio", () => { + expect(() => + normalizeQuotaWindow({ + rawKind: "PRIMARY", + windowSeconds: 5 * 3600, + remainingRatio: NaN, + }), + ).toThrow(); + }); + + it("rejects Infinity remainingRatio", () => { + expect(() => + normalizeQuotaWindow({ + rawKind: "PRIMARY", + windowSeconds: 5 * 3600, + remainingRatio: Infinity, + }), + ).toThrow(); + }); + + it("rejects boolean remainingRatio", () => { + expect(() => + normalizeQuotaWindow({ + rawKind: "PRIMARY", + windowSeconds: 5 * 3600, + remainingRatio: true as unknown as number, + }), + ).toThrow(); + }); + + it("rejects remainingRatio > 1", () => { + expect(() => + normalizeQuotaWindow({ + rawKind: "PRIMARY", + windowSeconds: 5 * 3600, + remainingRatio: 1.01, + }), + ).toThrow(); + }); + + it("rejects remainingRatio < 0", () => { + expect(() => + normalizeQuotaWindow({ + rawKind: "PRIMARY", + windowSeconds: 5 * 3600, + remainingRatio: -0.01, + }), + ).toThrow(); + }); + + it("clamps an explicit-zero usage to remainingRatio=0", () => { + const window = normalizeQuotaWindow({ + rawKind: "PRIMARY", + windowSeconds: 5 * 3600, + used: 800, + limit: 800, + }); + expect(window.remainingRatio).toBe(0); + }); + + it("defaults appliesTo to inference when unset", () => { + const window = normalizeQuotaWindow({ + rawKind: "PRIMARY", + windowSeconds: 5 * 3600, + remainingRatio: 0.5, + }); + expect(window.appliesTo).toBe("inference"); + }); + + it("preserves model-scoped appliesTo with model IDs", () => { + const window = normalizeQuotaWindow({ + rawKind: "MODEL_MODEL_QUOTA", + windowSeconds: 5 * 3600, + remainingRatio: 0.5, + appliesTo: "model", + modelIds: ["minimax/m2.5"], + }); + expect(window.appliesTo).toBe("model"); + expect(window.modelIds).toEqual(["minimax/m2.5"]); + }); + + it("rejects model-scoped window with no model IDs", () => { + expect(() => + normalizeQuotaWindow({ + rawKind: "MODEL_QUOTA", + windowSeconds: 5 * 3600, + remainingRatio: 0.5, + appliesTo: "model", + }), + ).toThrow(); + }); + + it("rejects inference-scoped window carrying model IDs", () => { + expect(() => + normalizeQuotaWindow({ + rawKind: "PRIMARY", + windowSeconds: 5 * 3600, + remainingRatio: 0.5, + appliesTo: "inference", + modelIds: ["minimax/m2.5"], + }), + ).toThrow(); + }); +}); + +describe("validateNormalizedSnapshot", () => { + const validSnapshot = { + provider: "zai", + account: "zai#1", + status: "fresh" as const, + windows: [ + { + id: "PRIMARY", + kind: "tokens_limit" as const, + appliesTo: "inference" as const, + modelIds: [] as string[], + remainingRatio: 0.86, + windowSeconds: 7 * 24 * 3600, + resetAt: "2026-07-26T20:23:52Z", + }, + ], + fetchedAt: "2026-07-24T17:33:47Z", + }; + + it("accepts a valid fresh snapshot", () => { + expect(() => validateNormalizedSnapshot(validSnapshot)).not.toThrow(); + }); + + it("rejects a snapshot with no windows", () => { + expect(() => + validateNormalizedSnapshot({ ...validSnapshot, windows: [] }), + ).toThrow(); + }); + + it("rejects a snapshot whose provider does not match the requested provider", () => { + expect(() => + validateNormalizedSnapshot({ ...validSnapshot, provider: "openai" }, "zai"), + ).toThrow(); + }); + + it("rejects a stale snapshot when not diagnostics-only", () => { + expect(() => + validateNormalizedSnapshot( + { ...validSnapshot, status: "stale" }, + undefined, + false, + ), + ).toThrow(); + }); + + it("accepts a stale snapshot when diagnostics-only", () => { + expect(() => + validateNormalizedSnapshot( + { ...validSnapshot, status: "stale" }, + undefined, + true, + ), + ).not.toThrow(); + }); + + it("rejects non-ISO, timezone-free, and impossible timestamps", () => { + for (const fetchedAt of [ + "1", + "07/24/2026", + "2026-07-24T17:00:00", + "2026-02-30T00:00:00Z", + ]) { + expect(() => validateNormalizedSnapshot({ ...validSnapshot, fetchedAt })).toThrow(); + } + }); +}); + +describe("normalizeSnapshot", () => { + it("normalizes a Z.AI payload into a normalized snapshot", () => { + const payload: ProviderQuotaPayload = { + provider: "zai", + account: "zai#1", + windows: [ + { + id: "5_hour", + rawKind: "TOKENS_LIMIT", + windowSeconds: 5 * 3600, + used: 112, + limit: 10000, + resetAt: "2026-07-24T20:23:52Z", + }, + { + id: "weekly", + rawKind: "TOKENS_LIMIT", + windowSeconds: 7 * 24 * 3600, + used: 1400, + limit: 10000, + resetAt: "2026-07-26T20:23:52Z", + }, + ], + fetchedAt: "2026-07-24T17:33:47Z", + }; + const snapshot = normalizeSnapshot(payload); + expect(snapshot.provider).toBe("zai"); + expect(snapshot.account).toBe("zai#1"); + expect(snapshot.status).toBe("fresh"); + expect(snapshot.windows).toHaveLength(2); + expect(snapshot.windows[0].remainingRatio).toBeCloseTo(0.9888); + expect(snapshot.windows[1].remainingRatio).toBeCloseTo(0.86); + expect(snapshot.windows.every((w) => w.appliesTo === "inference")).toBe(true); + }); + + it("normalizes an OpenAI Codex payload with primary/secondary windows", () => { + const payload: ProviderQuotaPayload = { + provider: "openai-codex", + account: "openai#1", + windows: [ + { + id: "primary", + rawKind: "TOKENS_LIMIT", + windowSeconds: 300 * 60, + percentageUsed: 47, + }, + { + id: "secondary", + rawKind: "TOKENS_LIMIT", + windowSeconds: 10080 * 60, + percentageUsed: 12, + }, + ], + fetchedAt: "2026-07-24T17:33:47Z", + }; + const snapshot = normalizeSnapshot(payload); + expect(snapshot.windows).toHaveLength(2); + expect(snapshot.windows[0].id).toBe("primary"); + expect(snapshot.windows[0].remainingRatio).toBeCloseTo(0.53, 1); + expect(snapshot.windows[1].id).toBe("secondary"); + expect(snapshot.windows[1].remainingRatio).toBeCloseTo(0.88, 1); + }); + + it("normalizes an xAI payload from bare remaining_percent", () => { + const payload: ProviderQuotaPayload = { + provider: "xai", + account: "xai#1", + windows: [ + { + id: "billing", + rawKind: "BILLING", + percentageRemaining: 100, + }, + ], + fetchedAt: "2026-07-24T17:33:47Z", + }; + const snapshot = normalizeSnapshot(payload); + expect(snapshot.windows).toHaveLength(1); + expect(snapshot.windows[0].remainingRatio).toBe(1); + }); + + it("marks a payload with no quantitative meter as unsupported", () => { + const payload: ProviderQuotaPayload = { + provider: "qwen-oauth", + account: "qwen#1", + windows: [], + fetchedAt: "2026-07-24T17:33:47Z", + }; + const snapshot = normalizeSnapshot(payload); + expect(snapshot.status).toBe("unsupported"); + expect(snapshot.windows).toHaveLength(0); + }); + + it("ignores extra provider-raw fields outside the token-free contract", () => { + const payload = { + provider: "qwen-oauth", + account: "qwen#1", + windows: [], + fetchedAt: "2026-07-24T17:33:47Z", + raw: { remains: { percentage: 90, plan_type: "NOT_MINIMAX" } }, + } as ProviderQuotaPayload & { raw: unknown }; + const snapshot = normalizeSnapshot(payload); + expect(snapshot.status).toBe("unsupported"); + expect(JSON.stringify(snapshot)).not.toContain("NOT_MINIMAX"); + }); + + it("fails closed on a malformed host-normalized window kind", () => { + const snapshot = normalizeSnapshot({ + provider: "zai", + account: "zai#1", + windows: [ + { + id: "bad", + rawKind: 1 as unknown as string, + remainingRatio: 0.5, + }, + ], + fetchedAt: "2026-07-24T17:33:47Z", + }); + expect(snapshot.status).toBe("invalid_response"); + expect(snapshot.windows).toEqual([]); + }); + + it("fails closed on duplicate semantic window IDs", () => { + const snapshot = normalizeSnapshot({ + provider: "zai", + account: "zai#1", + windows: [ + { id: "weekly", rawKind: "TOKENS_LIMIT", remainingRatio: 0.8 }, + { id: "weekly", rawKind: "TOKENS_LIMIT", remainingRatio: 0.7 }, + ], + fetchedAt: "2026-07-24T17:33:47Z", + }); + expect(snapshot.status).toBe("invalid_response"); + expect(snapshot.windows).toEqual([]); + }); + + it("copies only allowlisted fields into the normalized snapshot", () => { + const payload = { + provider: "zai", + account: "zai#1", + windows: [ + { + id: "5h", + rawKind: "TOKENS_LIMIT", + percentageRemaining: 98.88, + }, + ], + fetchedAt: "2026-07-24T17:33:47Z", + account_email: "secret@example.com", + access_token: "test-only-secret-placeholder", + } as ProviderQuotaPayload & { account_email: string; access_token: string }; + const snapshot = normalizeSnapshot(payload); + expect(JSON.stringify(snapshot)).not.toContain("secret@example.com"); + expect(JSON.stringify(snapshot)).not.toContain("test-only-secret-placeholder"); + expect(JSON.stringify(snapshot)).not.toContain("account_email"); + expect(JSON.stringify(snapshot)).not.toContain("access_token"); + }); + + it("rejects explicit null status instead of defaulting to fresh", () => { + const snapshot = normalizeSnapshot({ + provider: "zai", + account: "zai#1", + status: null as unknown as undefined, + windows: [{ id: "w", rawKind: "TOKENS_LIMIT", remainingRatio: 0.5 }], + fetchedAt: "2026-07-24T17:33:47Z", + }); + expect(snapshot.status).toBe("invalid_response"); + }); + + it("rejects explicit null appliesTo instead of widening to inference", () => { + const snapshot = normalizeSnapshot({ + provider: "zai", + account: "zai#1", + windows: [ + { + id: "w", + rawKind: "TOKENS_LIMIT", + remainingRatio: 0.5, + appliesTo: null as unknown as undefined, + }, + ], + fetchedAt: "2026-07-24T17:33:47Z", + }); + expect(snapshot.status).toBe("invalid_response"); + expect(snapshot.windows).toEqual([]); + }); + + it("rejects whitespace-padded model IDs in validateNormalizedSnapshot", () => { + expect(() => + validateNormalizedSnapshot({ + provider: "zai", + account: "zai#1", + status: "fresh", + windows: [ + { + id: "w", + kind: "tokens_limit", + appliesTo: "model", + modelIds: [" openai/gpt-5.6-sol "], + remainingRatio: 0.5, + } as any, + ], + fetchedAt: "2026-07-24T17:33:47Z", + }), + ).toThrow(); + }); + + it("rejects explicit null window id", () => { + expect(() => + normalizeQuotaWindow({ + rawKind: "TOKENS_LIMIT", + id: null as unknown as undefined, + remainingRatio: 0.5, + }), + ).toThrow(); + }); + + it("honors explicitZeroUsage marker as fully available (ratio=1)", () => { + const window = normalizeQuotaWindow({ + rawKind: "TOKENS_LIMIT", + explicitZeroUsage: true, + }); + expect(window.remainingRatio).toBe(1); + }); + + it("does not treat explicitZeroUsage=false as zero", () => { + const window = normalizeQuotaWindow({ + rawKind: "TOKENS_LIMIT", + explicitZeroUsage: false, + remainingRatio: 0.5, + }); + expect(window.remainingRatio).toBeCloseTo(0.5); + }); + + it("rejects null in any supplied meter field even when a valid meter exists", () => { + expect(() => + normalizeQuotaWindow({ + rawKind: "TOKENS_LIMIT", + remainingRatio: 0.5, + used: null as unknown as number, + limit: 100, + }), + ).toThrow(); + expect(() => + normalizeQuotaWindow({ + rawKind: "TOKENS_LIMIT", + percentageUsed: 10, + percentageRemaining: null as unknown as number, + }), + ).toThrow(); + }); +}); diff --git a/plugin/__tests__/quota-policy.test.ts b/plugin/__tests__/quota-policy.test.ts new file mode 100644 index 0000000..bf596d1 --- /dev/null +++ b/plugin/__tests__/quota-policy.test.ts @@ -0,0 +1,217 @@ +import { describe, it, expect } from "vitest"; +import { + computeQuotaFactor, + computeLivePressure, + selectAccountByQuota, +} from "../quota-policy.js"; +import type { NormalizedQuotaSnapshot, NormalizedQuotaWindow } from "../quota-types.js"; + +function snap( + provider: string, + account: string, + status: "fresh" | "stale" | "unsupported", + ...windows: Array<[string, number]> +): NormalizedQuotaSnapshot { + return { + provider, + account, + status, + windows: windows.map(([id, ratio]) => ({ + id, + kind: "tokens_limit" as const, + appliesTo: "inference" as const, + modelIds: [], + remainingRatio: ratio, + })), + fetchedAt: "2026-07-24T17:00:00Z", + }; +} + +describe("computeQuotaFactor", () => { + it("returns sqrt(min(remaining)) for a fresh snapshot with two windows", () => { + const snapshot = snap("zai", "zai#1", "fresh", ["5h", 0.9888], ["1w", 0.86]); + const factor = computeQuotaFactor(snapshot, "zai/glm-5.2"); + expect(factor).toBeCloseTo(Math.sqrt(0.86), 4); + }); + + it("returns null for a stale snapshot", () => { + const snapshot = snap("zai", "zai#1", "stale", ["5h", 1.0]); + expect(computeQuotaFactor(snapshot, "zai/glm-5.2")).toBeNull(); + }); + + it("fails open for a malformed fresh snapshot", () => { + const snapshot = snap("zai", "zai#1", "fresh", ["5h", Number.NaN]); + expect(computeQuotaFactor(snapshot, "zai/glm-5.2")).toBeNull(); + }); + + it("returns null for an unsupported snapshot", () => { + const snapshot = snap("qwen-oauth", "qwen#1", "unsupported"); + expect(computeQuotaFactor(snapshot, "qwen-oauth/qwen3.5")).toBeNull(); + }); + + it("returns 0 for a depleted account (remainingRatio=0)", () => { + const snapshot = snap("zai", "zai#1", "fresh", ["5h", 0.0]); + expect(computeQuotaFactor(snapshot, "zai/glm-5.2")).toBe(0); + }); + + it("returns 1 for a fully available account (remainingRatio=1)", () => { + const snapshot = snap("xai", "xai#1", "fresh", ["billing", 1.0]); + expect(computeQuotaFactor(snapshot, "xai/grok-4.5")).toBe(1); + }); + + it("excludes MCP-only windows from inference routing", () => { + const snapshot: NormalizedQuotaSnapshot = { + provider: "zai", + account: "zai#1", + status: "fresh", + windows: [ + { + id: "MCP_TIME", + kind: "time_limit", + appliesTo: "mcp", + modelIds: [], + remainingRatio: 0.0, + }, + { + id: "5h_TOKENS", + kind: "tokens_limit", + appliesTo: "inference", + modelIds: [], + remainingRatio: 0.80, + }, + ], + fetchedAt: "2026-07-24T17:00:00Z", + }; + expect(computeQuotaFactor(snapshot, "zai/glm-5.2")).toBeCloseTo(Math.sqrt(0.80), 4); + }); + + it("applies model-scoped windows only to mapped models", () => { + const snapshot: NormalizedQuotaSnapshot = { + provider: "minimax", + account: "minimax#1", + status: "fresh", + windows: [ + { + id: "INFERENCE", + kind: "tokens_limit", + appliesTo: "inference", + modelIds: [], + remainingRatio: 0.90, + }, + { + id: "M2.5_QUOTA", + kind: "tokens_limit", + appliesTo: "model", + modelIds: ["minimax/m2.5"], + remainingRatio: 0.10, + }, + ], + fetchedAt: "2026-07-24T17:00:00Z", + }; + expect(computeQuotaFactor(snapshot, "minimax/m2.5")).toBeCloseTo(Math.sqrt(0.10), 4); + expect(computeQuotaFactor(snapshot, "minimax/m2.7")).toBeCloseTo(Math.sqrt(0.90), 4); + }); +}); + +describe("computeLivePressure", () => { + it("multiplies static pressure by quota factor", () => { + const snapshot = snap("zai", "zai#1", "fresh", ["1w", 0.86]); + const result = computeLivePressure(5.0, 1.25, snapshot, "zai/glm-5.2"); + expect(result).toBeCloseTo(5.0 * 1.25 * Math.sqrt(0.86), 4); + }); + + it("returns null when quota factor is null (stale/unsupported)", () => { + const snapshot = snap("zai", "zai#1", "stale", ["1w", 1.0]); + expect(computeLivePressure(5.0, 1.25, snapshot, "zai/glm-5.2")).toBeNull(); + }); + + it("returns 0 when account is depleted", () => { + const snapshot = snap("zai", "zai#1", "fresh", ["1w", 0.0]); + expect(computeLivePressure(5.0, 1.25, snapshot, "zai/glm-5.2")).toBe(0); + }); +}); + +describe("selectAccountByQuota", () => { + it("selects the account with higher live pressure", () => { + const accounts = [ + { provider: "openai", account: "openai#1", tierWeight: 2.1, providerBias: 0.7, snapshot: snap("openai", "openai#1", "fresh", ["5h", 0.10]) }, + { provider: "openai", account: "openai#2", tierWeight: 2.1, providerBias: 0.7, snapshot: snap("openai", "openai#2", "fresh", ["5h", 0.80]) }, + ]; + const selected = selectAccountByQuota(accounts, "openai", "openai/gpt-5.6-sol"); + expect(selected?.account).toBe("openai#2"); + }); + + it("returns an account via static fallback when all accounts are stale", () => { + const accounts = [ + { provider: "openai", account: "openai#1", tierWeight: 2.1, providerBias: 0.7, snapshot: snap("openai", "openai#1", "stale", ["5h", 1.0]) }, + ]; + const selected = selectAccountByQuota(accounts, "openai", "openai/gpt-5.6-sol"); + expect(selected?.account).toBe("openai#1"); + }); + + it("returns null when all accounts are depleted", () => { + const accounts = [ + { provider: "zai", account: "zai#1", tierWeight: 5.0, providerBias: 1.25, snapshot: snap("zai", "zai#1", "fresh", ["5h", 0.0]) }, + ]; + expect(selectAccountByQuota(accounts, "zai", "zai/glm-5.2")).toBeNull(); + }); + + it("falls open to static pressure when quota is stale", () => { + const accounts = [ + { provider: "openai", account: "openai#1", tierWeight: 1.0, providerBias: 1.0, snapshot: snap("openai", "openai#1", "stale", ["5h", 1.0]) }, + { provider: "openai", account: "openai#2", tierWeight: 3.0, providerBias: 2.0, snapshot: snap("openai", "openai#2", "stale", ["5h", 1.0]) }, + ]; + const selected = selectAccountByQuota(accounts, "openai", "openai/gpt-5.6-sol"); + expect(selected?.account).toBe("openai#2"); + }); + + it("falls open to static pressure when snapshot is null", () => { + const accounts = [ + { provider: "openai", account: "openai#1", tierWeight: 5.0, providerBias: 2.0, snapshot: null }, + ]; + const selected = selectAccountByQuota(accounts, "openai", "openai/gpt-5.6-sol"); + expect(selected?.account).toBe("openai#1"); + }); + + it("prefers live-constrained account only when it still beats static fallback", () => { + const accounts = [ + { provider: "openai", account: "openai#1", tierWeight: 1.0, providerBias: 1.0, snapshot: null }, + { provider: "openai", account: "openai#2", tierWeight: 5.0, providerBias: 2.0, snapshot: snap("openai", "openai#2", "fresh", ["5h", 0.1]) }, + ]; + const selected = selectAccountByQuota(accounts, "openai", "openai/gpt-5.6-sol"); + // static=1.0 vs live=5*2*sqrt(0.1)=3.16 → live wins + expect(selected?.account).toBe("openai#2"); + }); + + it("filters by provider", () => { + const accounts = [ + { provider: "zai", account: "zai#1", tierWeight: 5.0, providerBias: 1.25, snapshot: snap("zai", "zai#1", "fresh", ["5h", 0.9]) }, + { provider: "openai", account: "openai#1", tierWeight: 2.1, providerBias: 0.7, snapshot: snap("openai", "openai#1", "fresh", ["5h", 0.9]) }, + ]; + const selected = selectAccountByQuota(accounts, "zai", "zai/glm-5.2"); + expect(selected?.account).toBe("zai#1"); + }); + + it("rejects a mismatched snapshot from another account", () => { + const healthyOther: NormalizedQuotaSnapshot = { + provider: "openai", + account: "openai#1", + status: "fresh", + windows: [{ id: "5h", kind: "tokens_limit", appliesTo: "inference", modelIds: [], remainingRatio: 0.99 }], + fetchedAt: "2026-07-24T17:00:00Z", + }; + const accounts = [ + { provider: "openai", account: "openai#2", tierWeight: 2.1, providerBias: 0.7, snapshot: healthyOther }, + ]; + const selected = selectAccountByQuota(accounts, "openai", "openai/gpt-5.6-sol"); + expect(selected).toBeNull(); + }); + + it("accepts a matched snapshot and scores normally", () => { + const accounts = [ + { provider: "openai", account: "openai#1", tierWeight: 2.1, providerBias: 0.7, snapshot: snap("openai", "openai#1", "fresh", ["5h", 0.80]) }, + ]; + const selected = selectAccountByQuota(accounts, "openai", "openai/gpt-5.6-sol"); + expect(selected?.account).toBe("openai#1"); + }); +}); diff --git a/plugin/quota-normalize.ts b/plugin/quota-normalize.ts new file mode 100644 index 0000000..7e921ed --- /dev/null +++ b/plugin/quota-normalize.ts @@ -0,0 +1,305 @@ +/** + * Provider-neutral quota normalization. + * + * Provider HTTP/RPC parsing and credential lifecycle belong to the host + * (OpenClaw/Hermes). ZeroAPI receives token-free quantitative windows only. + * This module validates and copies an allowlisted schema; unknown/raw fields + * never cross into NormalizedQuotaSnapshot. + */ + +import type { + NormalizedQuotaSnapshot, + NormalizedQuotaWindow, + NormalizeWindowInput, + ProviderQuotaPayload, + QuotaAppliesTo, + QuotaSnapshotStatus, + QuotaWindowKind, +} from "./quota-types.js"; + +class ValueError extends Error {} + +const VALID_STATUSES = new Set([ + "fresh", + "stale", + "auth_expired", + "rate_limited", + "network_error", + "invalid_response", + "unsupported", +]); + +const VALID_WINDOW_KINDS = new Set([ + "tokens_limit", + "requests_limit", + "credits", + "messages", + "compute", + "time_limit", + "percent", +]); + +const VALID_APPLICABILITY = new Set(["inference", "mcp", "model"]); + +function assertFiniteNumber(value: unknown, label: string): asserts value is number { + if (typeof value !== "number") { + if (typeof value === "boolean") throw new TypeError(`${label} must be numeric, got boolean`); + throw new TypeError(`${label} must be numeric`); + } + if (Number.isNaN(value)) throw new ValueError(`${label} must not be NaN`); + if (!Number.isFinite(value)) throw new ValueError(`${label} must be finite`); +} + +function assertValidRatio(value: unknown): asserts value is number { + assertFiniteNumber(value, "remainingRatio"); + if (value < 0 || value > 1) throw new ValueError("remainingRatio must be in [0, 1]"); +} + +function normalizePercentage(value: unknown, label: string): number { + assertFiniteNumber(value, label); + if (value < 0 || value > 100) throw new ValueError(`${label} must be in [0, 100]`); + const ratio = value / 100; + assertValidRatio(ratio); + return ratio; +} + +function normalizeIdentifier(value: unknown, label: string): string { + if (typeof value !== "string") throw new TypeError(`${label} must be a string`); + const trimmed = value.trim(); + if (!trimmed) throw new ValueError(`${label} must be non-empty`); + if (trimmed.length > 256) throw new ValueError(`${label} is too long`); + if (/[\u0000-\u001f\u007f]/.test(trimmed)) { + throw new ValueError(`${label} contains control characters`); + } + return trimmed; +} + +function normalizeTimestamp(value: unknown, label: string): string { + const normalized = normalizeIdentifier(value, label); + const match = /^(\d{4})-(\d{2})-(\d{2})T(\d{2}):(\d{2}):(\d{2})(?:\.\d{1,9})?(Z|[+-](\d{2}):(\d{2}))$/.exec(normalized); + if (!match) throw new ValueError(`${label} must be an ISO-8601 timestamp with timezone`); + + const year = Number(match[1]); + const month = Number(match[2]); + const day = Number(match[3]); + const hour = Number(match[4]); + const minute = Number(match[5]); + const second = Number(match[6]); + const offsetHour = match[8] === undefined ? 0 : Number(match[8]); + const offsetMinute = match[9] === undefined ? 0 : Number(match[9]); + const leapYear = year % 4 === 0 && (year % 100 !== 0 || year % 400 === 0); + const daysInMonth = [31, leapYear ? 29 : 28, 31, 30, 31, 30, 31, 31, 30, 31, 30, 31]; + + if ( + month < 1 || month > 12 || + day < 1 || day > daysInMonth[month - 1] || + hour > 23 || minute > 59 || second > 59 || + offsetHour > 23 || offsetMinute > 59 + ) { + throw new ValueError(`${label} must contain a valid calendar date and time`); + } + return normalized; +} + +function mapWindowKind(rawKind: string): QuotaWindowKind { + const upper = rawKind.toUpperCase(); + if (upper.includes("TOKEN")) return "tokens_limit"; + if (upper.includes("REQUEST") || upper.includes("RPM")) return "requests_limit"; + if (upper.includes("CREDIT")) return "credits"; + if (upper.includes("MESSAGE")) return "messages"; + if (upper.includes("TIME_LIMIT")) return "time_limit"; + if (upper.includes("COMPUTE") || upper.includes("TIME")) return "compute"; + if (upper.includes("PERCENT") || upper === "USAGE" || upper === "BILLING") return "percent"; + return "tokens_limit"; +} + +/** Normalize one token-free host adapter window. */ +export function normalizeQuotaWindow(input: NormalizeWindowInput): NormalizedQuotaWindow { + const rawKind = normalizeIdentifier(input.rawKind, "rawKind"); + if (input.id === null) { + throw new TypeError("window id must not be null"); + } + const id = normalizeIdentifier(input.id ?? rawKind, "window id"); + const kind = mapWindowKind(rawKind); + + // Validate every supplied meter field before selecting one, matching + // the Python _input_window null-rejection contract. A present-but-null + // meter must fail closed rather than being silently skipped. + const meterFields = [ + "remainingRatio", "used", "limit", + "percentageRemaining", "percentageUsed", "windowSeconds", "resetAt", + ] as const; + for (const field of meterFields) { + if ((input as Record)[field] === null) { + throw new TypeError(`${field} must not be null`); + } + } + + let remainingRatio: number; + if (input.explicitZeroUsage === true) { + remainingRatio = 1.0; + } else if (input.remainingRatio !== undefined) { + assertValidRatio(input.remainingRatio); + remainingRatio = input.remainingRatio; + } else if (input.percentageRemaining !== undefined) { + remainingRatio = normalizePercentage(input.percentageRemaining, "percentageRemaining"); + } else if (input.percentageUsed !== undefined) { + const usedRatio = normalizePercentage(input.percentageUsed, "percentageUsed"); + remainingRatio = Math.max(0, 1 - usedRatio); + } else if (input.used !== undefined && input.limit !== undefined) { + assertFiniteNumber(input.used, "used"); + assertFiniteNumber(input.limit, "limit"); + if (input.used < 0) throw new ValueError("used must be non-negative"); + if (input.limit <= 0) throw new ValueError("limit must be positive"); + remainingRatio = Math.max(0, 1 - input.used / input.limit); + assertValidRatio(remainingRatio); + } else { + throw new ValueError(`cannot derive remainingRatio for window "${id}"`); + } + + if (input.appliesTo === null) { + throw new ValueError("appliesTo must not be null"); + } + const appliesTo: QuotaAppliesTo = input.appliesTo ?? "inference"; + if (appliesTo !== "inference" && appliesTo !== "mcp" && appliesTo !== "model") { + throw new ValueError(`unknown appliesTo value "${String(appliesTo)}"`); + } + + const modelIdsInput = input.modelIds === undefined ? [] : input.modelIds; + if (!Array.isArray(modelIdsInput)) { + throw new TypeError("modelIds must be an array"); + } + const modelIds = modelIdsInput.map((modelId) => normalizeIdentifier(modelId, "modelId")); + if (new Set(modelIds).size !== modelIds.length) { + throw new ValueError("modelIds must be unique"); + } + if (appliesTo === "model" && modelIds.length === 0) { + throw new ValueError("model-scoped window requires at least one modelId"); + } + if (appliesTo !== "model" && modelIds.length > 0) { + throw new ValueError("non-model window must not carry modelIds"); + } + + let windowSeconds: number | undefined; + if (input.windowSeconds !== undefined) { + assertFiniteNumber(input.windowSeconds, "windowSeconds"); + if (input.windowSeconds <= 0) throw new ValueError("windowSeconds must be positive"); + windowSeconds = input.windowSeconds; + } + + const resetAt = input.resetAt === undefined + ? undefined + : normalizeTimestamp(input.resetAt, "resetAt"); + + return { + id, + kind, + appliesTo, + modelIds, + remainingRatio, + windowSeconds, + resetAt, + }; +} + +/** Validate a complete normalized snapshot. */ +export function validateNormalizedSnapshot( + snapshot: NormalizedQuotaSnapshot, + expectedProvider?: string, + diagnosticsOnly: boolean = false, +): void { + const provider = normalizeIdentifier(snapshot.provider, "snapshot provider"); + normalizeIdentifier(snapshot.account, "snapshot account"); + normalizeTimestamp(snapshot.fetchedAt, "fetchedAt"); + + if (!VALID_STATUSES.has(snapshot.status)) { + throw new ValueError(`unknown snapshot status "${String(snapshot.status)}"`); + } + if (expectedProvider !== undefined && provider !== expectedProvider) { + throw new ValueError(`snapshot provider "${provider}" does not match expected "${expectedProvider}"`); + } + if (!diagnosticsOnly && snapshot.status !== "fresh") { + throw new ValueError(`routing snapshot must be fresh, got "${snapshot.status}"`); + } + if (snapshot.status === "fresh" && snapshot.windows.length === 0) { + throw new ValueError("fresh snapshot requires at least one window"); + } + + const ids = new Set(); + for (const window of snapshot.windows) { + const windowId = normalizeIdentifier(window.id, "window id"); + assertValidRatio(window.remainingRatio); + if (!VALID_WINDOW_KINDS.has(window.kind)) { + throw new ValueError(`unknown window kind "${String(window.kind)}"`); + } + if (!VALID_APPLICABILITY.has(window.appliesTo)) { + throw new ValueError(`unknown appliesTo value "${String(window.appliesTo)}"`); + } + if (!Array.isArray(window.modelIds)) throw new TypeError("modelIds must be an array"); + const modelIds = window.modelIds.map((modelId) => { + const trimmed = normalizeIdentifier(modelId, "modelId"); + if (trimmed !== modelId) { + throw new ValueError("modelIds must be pre-canonicalized"); + } + return trimmed; + }); + if (new Set(modelIds).size !== modelIds.length) { + throw new ValueError("modelIds must be unique"); + } + if (window.appliesTo === "model" && modelIds.length === 0) { + throw new ValueError("model-scoped window requires at least one modelId"); + } + if (window.appliesTo !== "model" && modelIds.length > 0) { + throw new ValueError("non-model window must not carry modelIds"); + } + if (window.windowSeconds !== undefined) { + assertFiniteNumber(window.windowSeconds, "windowSeconds"); + if (window.windowSeconds <= 0) throw new ValueError("windowSeconds must be positive"); + } + if (window.resetAt !== undefined) normalizeTimestamp(window.resetAt, "resetAt"); + if (ids.has(windowId)) throw new ValueError(`duplicate window id "${windowId}"`); + ids.add(windowId); + } +} + +/** + * Convert a token-free host adapter payload into the allowlisted snapshot. + * A malformed quantitative window fails the whole observation closed as + * invalid_response; partial provider data is never used for routing. + */ +export function normalizeSnapshot(payload: ProviderQuotaPayload): NormalizedQuotaSnapshot { + const provider = normalizeIdentifier(payload.provider, "provider"); + const account = normalizeIdentifier(payload.account, "account"); + const fetchedAt = normalizeTimestamp(payload.fetchedAt, "fetchedAt"); + const requestedStatus = payload.status === undefined ? "fresh" : payload.status; + const status = typeof requestedStatus === "string" && VALID_STATUSES.has(requestedStatus as QuotaSnapshotStatus) + ? requestedStatus + : "invalid_response"; + + let windows: NormalizedQuotaWindow[] = []; + let normalizedStatus: QuotaSnapshotStatus = status; + try { + windows = payload.windows.map(normalizeQuotaWindow); + const ids = new Set(windows.map((window) => window.id)); + if (ids.size !== windows.length) { + throw new ValueError("window IDs must be unique"); + } + } catch { + windows = []; + normalizedStatus = "invalid_response"; + } + + if (normalizedStatus === "fresh" && windows.length === 0) { + normalizedStatus = "unsupported"; + } + + const snapshot: NormalizedQuotaSnapshot = { + provider, + account, + status: normalizedStatus, + windows, + fetchedAt, + }; + validateNormalizedSnapshot(snapshot, undefined, true); + return snapshot; +} diff --git a/plugin/quota-policy.ts b/plugin/quota-policy.ts new file mode 100644 index 0000000..a4e5204 --- /dev/null +++ b/plugin/quota-policy.ts @@ -0,0 +1,133 @@ +/** + * Live quota pressure policy. + * + * Converts normalized quota snapshots into routing-safe pressure scores. + * + * Policy: + * applicableWindows = inference-wide + matching model-specific windows + * accountHeadroom = min(remainingRatio across applicableWindows) + * quotaFactor = sqrt(accountHeadroom) + * livePressure = staticPressure * quotaFactor + * + * Rules: + * - quotaFactor can only REDUCE declared capacity, never boost it. + * - MCP/tool-only windows never modulate inference routing. + * - Stale/unsupported/unknown snapshots produce null → static pressure only. + * - Depleted accounts (headroom=0) produce 0 → excluded from selection. + */ + +import type { NormalizedQuotaSnapshot, NormalizedQuotaWindow } from "./quota-types.js"; +import { validateNormalizedSnapshot } from "./quota-normalize.js"; + +/** + * Select windows that apply to a specific candidate model. + * Includes inference-wide windows plus model-scoped windows that match. + * Excludes MCP/tool-only windows. + */ +export function applicableWindows( + snapshot: NormalizedQuotaSnapshot, + model: string, +): NormalizedQuotaWindow[] { + return snapshot.windows.filter((w) => { + if (w.appliesTo === "inference") return true; + if (w.appliesTo === "model") return w.modelIds.includes(model); + return false; // mcp/tool-only excluded + }); +} + +/** + * Compute the minimum remaining ratio across applicable windows. + * Returns null if the snapshot is not fresh or has no applicable windows. + */ +export function accountHeadroom( + snapshot: NormalizedQuotaSnapshot | null, + model: string, +): number | null { + if (!snapshot || snapshot.status !== "fresh") return null; + try { + validateNormalizedSnapshot(snapshot); + } catch { + return null; + } + const windows = applicableWindows(snapshot, model); + if (windows.length === 0) return null; + return Math.min(...windows.map((w) => w.remainingRatio)); +} + +/** + * Compute the quota factor: sqrt(headroom). + * Returns null for stale/unsupported; 0 for depleted; 1 for fully available. + * Can only reduce static pressure, never boost it. + */ +export function computeQuotaFactor( + snapshot: NormalizedQuotaSnapshot | null, + model: string, +): number | null { + const headroom = accountHeadroom(snapshot, model); + if (headroom === null) return null; + return Math.sqrt(headroom); +} + +/** Input for live pressure computation. */ +export type QuotaAwareAccount = { + provider: string; + account: string; + tierWeight: number; + providerBias: number; + snapshot: NormalizedQuotaSnapshot | null; +}; + +/** + * Compute live pressure for a candidate model. + * staticPressure = tierWeight * providerBias; livePressure adds quotaFactor. + */ +export function computeLivePressure( + tierWeight: number, + providerBias: number, + snapshot: NormalizedQuotaSnapshot | null, + model: string, +): number | null { + const factor = computeQuotaFactor(snapshot, model); + if (factor === null) return null; + return tierWeight * providerBias * factor; +} + +/** + * Select the best account for a provider and model. + * Returns null if no account has a usable pressure. + * Rejects snapshots whose provider or account do not match the candidate, + * guarding against mis-keyed or swapped snapshots during host integration. + * + * Fail-open: when live quota is unavailable (stale/absent/unsupported), + * the account falls back to static pressure (tierWeight * providerBias). + * Only accounts with a confirmed depleted quota (live pressure = 0) are excluded. + */ +export function selectAccountByQuota( + accounts: QuotaAwareAccount[], + provider: string, + model: string, +): QuotaAwareAccount | null { + const eligible = accounts.filter((a) => { + if (a.provider !== provider) return false; + if (a.snapshot && (a.snapshot.provider !== a.provider || a.snapshot.account !== a.account)) { + return false; + } + const live = computeLivePressure(a.tierWeight, a.providerBias, a.snapshot, model); + // null = quota unavailable → fail open to static pressure + // 0 = confirmed depleted → exclude + if (live === 0) return false; + return true; + }); + + if (eligible.length === 0) return null; + + return eligible.reduce((best, current) => { + const bestPressure = computeLivePressure(best.tierWeight, best.providerBias, best.snapshot, model) + ?? (best.tierWeight * best.providerBias); + const currentPressure = computeLivePressure(current.tierWeight, current.providerBias, current.snapshot, model) + ?? (current.tierWeight * current.providerBias); + if (currentPressure > bestPressure) return current; + if (currentPressure === bestPressure && current.account < best.account) return current; + return best; + }); +} diff --git a/plugin/quota-types.ts b/plugin/quota-types.ts new file mode 100644 index 0000000..ac1faa5 --- /dev/null +++ b/plugin/quota-types.ts @@ -0,0 +1,100 @@ +/** + * Normalized quota types — secret-free, routing-safe. + * + * These types describe the *output* of a provider-specific normalizer. They + * never carry credentials, tokens, account emails, or raw provider payloads. + * The router reads only these types. + */ + +export type QuotaWindowKind = + | "tokens_limit" + | "requests_limit" + | "credits" + | "messages" + | "compute" + | "time_limit" + | "percent"; + +/** What kind of work this window limits. */ +export type QuotaAppliesTo = "inference" | "mcp" | "model"; + +/** Provider-reported or locally-observed snapshot freshness. */ +export type QuotaSnapshotStatus = + | "fresh" + | "stale" + | "auth_expired" + | "rate_limited" + | "network_error" + | "invalid_response" + | "unsupported"; + +/** A single normalized quota window. */ +export type NormalizedQuotaWindow = { + /** Non-secret semantic ID, e.g. "primary" or "weekly". */ + id: string; + kind: QuotaWindowKind; + appliesTo: QuotaAppliesTo; + /** + * Canonical local model IDs this window applies to. + * Required when appliesTo="model"; must be empty otherwise. + */ + modelIds: string[]; + /** Remaining ratio in [0, 1]. Derived only from provider counters/percent. */ + remainingRatio: number; + /** Provider-declared window duration in seconds, if known. */ + windowSeconds?: number; + /** UTC ISO-8601 reset time, if known. */ + resetAt?: string; +}; + +/** A normalized, per-account quota snapshot. */ +export type NormalizedQuotaSnapshot = { + provider: string; + account: string; + status: QuotaSnapshotStatus; + windows: NormalizedQuotaWindow[]; + fetchedAt: string; +}; + +/** + * Token-free host adapter payload. + * + * Hermes/OpenClaw own provider HTTP/RPC parsing and credential lifecycle. + * ZeroAPI receives only these provider-neutral quantitative windows. + */ +export type ProviderQuotaPayload = { + provider: string; + account: string; + status?: QuotaSnapshotStatus; + windows: NormalizeWindowInput[]; + fetchedAt: string; +}; + +/** + * Input parameters for a single host-normalized window. + * Either remainingRatio or (used + limit) must be provided. + */ +export type NormalizeWindowInput = { + /** Non-secret semantic ID; defaults to rawKind. */ + id?: string; + rawKind: string; + windowSeconds?: number; + resetAt?: string; + /** Direct remaining ratio (0..1). */ + remainingRatio?: number; + /** Used amount, when provider reports usage/limit. */ + used?: number; + /** Total limit, when provider reports usage/limit. */ + limit?: number; + /** Remaining percentage on the explicit 0-100 scale. */ + percentageRemaining?: number; + /** Consumed percentage on the explicit 0-100 scale. */ + percentageUsed?: number; + appliesTo?: QuotaAppliesTo; + modelIds?: string[]; + /** Explicit zero-usage marker from provider (e.g. xAI protobuf). */ + explicitZeroUsage?: boolean; +}; + +/** Input for full snapshot normalization. */ +export type NormalizeSnapshotInput = ProviderQuotaPayload; diff --git a/scripts/stage_clawhub_plugin.mjs b/scripts/stage_clawhub_plugin.mjs index 6fcb8a3..d2b6059 100644 --- a/scripts/stage_clawhub_plugin.mjs +++ b/scripts/stage_clawhub_plugin.mjs @@ -33,6 +33,9 @@ const runtimeEntries = [ "logger.ts", "onboarding.ts", "profile.ts", + "quota-normalize.ts", + "quota-policy.ts", + "quota-types.ts", "route-state.ts", "router.ts", "selector.ts",