diff --git a/README.md b/README.md index d6d537d1..e76baad1 100644 --- a/README.md +++ b/README.md @@ -13,6 +13,7 @@ ## Table of Contents +- [What's new (2026-06-21) — Client-Side Rate Limiting](#whats-new-2026-06-21--client-side-rate-limiting) - [What's new (2026-06-21) — JSON Web Tokens (JWT)](#whats-new-2026-06-21--json-web-tokens-jwt) - [What's new (2026-06-21) — License Policy Gate](#whats-new-2026-06-21--license-policy-gate) - [What's new (2026-06-21) — OpenVEX Vulnerability Triage](#whats-new-2026-06-21--openvex-vulnerability-triage) @@ -114,6 +115,12 @@ --- +## What's new (2026-06-21) — Client-Side Rate Limiting + +Stay under API quotas. Full reference: [`docs/source/Eng/doc/new_features/v62_features_doc.rst`](docs/source/Eng/doc/new_features/v62_features_doc.rst). + +- **`TokenBucket` / `SlidingWindowLimiter` / `throttle`** (`AC_rate_limit`, `ac_rate_limit`): `RetryPolicy`/`CircuitBreaker` recover from failures but nothing shaped the *rate* of calls. This adds a token bucket (smooth rate + burst), a sliding-window limiter (Cloudflare's O(1) weighted counter), and a leading-edge throttle decorator. Every limiter takes an injectable `clock` (and `acquire` a `sleep`) so it's fully deterministic in CI with no real delays. `AC_rate_limit` gates an action against a named bucket, returning `{acquired, tokens, wait}`. + ## What's new (2026-06-21) — JSON Web Tokens (JWT) Mint and verify bearer tokens for the APIs you automate. Full reference: [`docs/source/Eng/doc/new_features/v61_features_doc.rst`](docs/source/Eng/doc/new_features/v61_features_doc.rst). diff --git a/README/README_zh-CN.md b/README/README_zh-CN.md index b71a283d..915d5d6b 100644 --- a/README/README_zh-CN.md +++ b/README/README_zh-CN.md @@ -12,6 +12,7 @@ ## 目录 +- [本次更新 (2026-06-21) — 客户端速率限制](#本次更新-2026-06-21--客户端速率限制) - [本次更新 (2026-06-21) — JSON Web Token(JWT)](#本次更新-2026-06-21--json-web-tokenjwt) - [本次更新 (2026-06-21) — 许可证策略闸门](#本次更新-2026-06-21--许可证策略闸门) - [本次更新 (2026-06-21) — OpenVEX 漏洞分级](#本次更新-2026-06-21--openvex-漏洞分级) @@ -113,6 +114,12 @@ --- +## 本次更新 (2026-06-21) — 客户端速率限制 + +守在 API 配额之内。完整参考:[`docs/source/Zh/doc/new_features/v62_features_doc.rst`](../docs/source/Zh/doc/new_features/v62_features_doc.rst)。 + +- **`TokenBucket` / `SlidingWindowLimiter` / `throttle`**(`AC_rate_limit`、`ac_rate_limit`):`RetryPolicy`/`CircuitBreaker` 从失败中复原,但没有任何东西塑形调用的*速率*。本功能补上 token bucket(平滑速率 + 突发)、sliding-window 限制器(Cloudflare 的 O(1) 加权计数)以及前缘 throttle 装饰器。每个限制器都接受可注入的 `clock`(`acquire` 另接受 `sleep`),因此在 CI 完全确定、没有真正延迟。`AC_rate_limit` 以具名 bucket 闸控动作,返回 `{acquired, tokens, wait}`。 + ## 本次更新 (2026-06-21) — JSON Web Token(JWT) 为你自动化的 API 签发与验证 bearer token。完整参考:[`docs/source/Zh/doc/new_features/v61_features_doc.rst`](../docs/source/Zh/doc/new_features/v61_features_doc.rst)。 diff --git a/README/README_zh-TW.md b/README/README_zh-TW.md index ecb71686..8584f233 100644 --- a/README/README_zh-TW.md +++ b/README/README_zh-TW.md @@ -12,6 +12,7 @@ ## 目錄 +- [本次更新 (2026-06-21) — 用戶端速率限制](#本次更新-2026-06-21--用戶端速率限制) - [本次更新 (2026-06-21) — JSON Web Token(JWT)](#本次更新-2026-06-21--json-web-tokenjwt) - [本次更新 (2026-06-21) — 授權政策閘門](#本次更新-2026-06-21--授權政策閘門) - [本次更新 (2026-06-21) — OpenVEX 漏洞分級](#本次更新-2026-06-21--openvex-漏洞分級) @@ -113,6 +114,12 @@ --- +## 本次更新 (2026-06-21) — 用戶端速率限制 + +守在 API 配額之內。完整參考:[`docs/source/Zh/doc/new_features/v62_features_doc.rst`](../docs/source/Zh/doc/new_features/v62_features_doc.rst)。 + +- **`TokenBucket` / `SlidingWindowLimiter` / `throttle`**(`AC_rate_limit`、`ac_rate_limit`):`RetryPolicy`/`CircuitBreaker` 從失敗中復原,但沒有任何東西塑形呼叫的*速率*。本功能補上 token bucket(平滑速率 + 突發)、sliding-window 限制器(Cloudflare 的 O(1) 加權計數)以及前緣 throttle 裝飾器。每個限制器都接受可注入的 `clock`(`acquire` 另接受 `sleep`),因此在 CI 完全具決定性、沒有真正延遲。`AC_rate_limit` 以具名 bucket 閘控動作,回傳 `{acquired, tokens, wait}`。 + ## 本次更新 (2026-06-21) — JSON Web Token(JWT) 為你自動化的 API 簽發與驗證 bearer token。完整參考:[`docs/source/Zh/doc/new_features/v61_features_doc.rst`](../docs/source/Zh/doc/new_features/v61_features_doc.rst)。 diff --git a/docs/source/Eng/doc/new_features/v62_features_doc.rst b/docs/source/Eng/doc/new_features/v62_features_doc.rst new file mode 100644 index 00000000..a112bc00 --- /dev/null +++ b/docs/source/Eng/doc/new_features/v62_features_doc.rst @@ -0,0 +1,53 @@ +Client-Side Rate Limiting +========================= + +The framework had ``RetryPolicy`` / ``CircuitBreaker`` (which *recover* from +failures) and a FIFO ``work_queue``, but nothing to shape the *rate* of calls — +so a flow hammering an external API had no way to stay under a quota. This adds +the two standard limiters plus a leading-edge throttle, all with an injectable +clock so they are deterministic in tests (no real sleeping). + +* :class:`TokenBucket` — a smooth rate with burst capacity (lazy refill). +* :class:`SlidingWindowLimiter` — a fixed call budget per rolling window + (Cloudflare's O(1) weighted-counter approximation). +* :func:`throttle` — a decorator that fires a function at most once per interval. + +Pure standard library (``threading`` for the lock, ``time`` only as the default +clock); imports no ``PySide6``. + +Headless API +------------ + +.. code-block:: python + + from je_auto_control import TokenBucket, SlidingWindowLimiter, throttle + + # 5 requests/second, bursts up to 10 + bucket = TokenBucket(rate=5, capacity=10) + if bucket.try_acquire(): + call_api() # non-blocking: skip / queue if False + bucket.acquire() # or block until a token frees up + + # at most 100 calls per 60s rolling window + window = SlidingWindowLimiter(limit=100, window_s=60) + if window.try_acquire(): + call_api() + + @throttle(2.0) # fire at most once every 2 seconds + def on_event(payload): + ... + +``TokenBucket.try_acquire`` takes tokens if available; ``acquire`` blocks (with +an optional ``timeout``); ``time_until_available`` reports the wait so a +scheduler can pace itself. Every limiter accepts a ``clock=`` (and ``acquire`` +a ``sleep=``) so the whole thing is exercised in CI with a fake clock — no real +delays. + +Executor command +---------------- + +``AC_rate_limit`` takes a limiter ``name`` plus ``rate`` / ``capacity`` / ``n`` +and tries to take ``n`` tokens from that named token bucket (created on first +use), returning ``{acquired, tokens, wait}`` so a flow can gate or defer an +action. The same operation is exposed as the MCP tool ``ac_rate_limit`` and as a +Script Builder command under **Flow**. diff --git a/docs/source/Eng/eng_index.rst b/docs/source/Eng/eng_index.rst index 645c6e75..c430ce88 100644 --- a/docs/source/Eng/eng_index.rst +++ b/docs/source/Eng/eng_index.rst @@ -84,6 +84,7 @@ Comprehensive guides for all AutoControl features. doc/new_features/v59_features_doc doc/new_features/v60_features_doc doc/new_features/v61_features_doc + doc/new_features/v62_features_doc doc/ocr_backends/ocr_backends_doc doc/observability/observability_doc doc/operations_layer/operations_layer_doc diff --git a/docs/source/Zh/doc/new_features/v62_features_doc.rst b/docs/source/Zh/doc/new_features/v62_features_doc.rst new file mode 100644 index 00000000..8c3d9ced --- /dev/null +++ b/docs/source/Zh/doc/new_features/v62_features_doc.rst @@ -0,0 +1,48 @@ +用戶端速率限制 +============= + +框架過去有 ``RetryPolicy`` / ``CircuitBreaker``(用於從失敗中*復原*)以及 FIFO 的 +``work_queue``,卻沒有任何東西能塑形呼叫的*速率* —— 因此猛打外部 API 的流程沒有辦法守在配額 +之內。本功能補上兩個標準限制器,外加一個前緣觸發的 throttle,全都以可注入的時鐘實作,因此在測試 +中具決定性(不需真的睡眠)。 + +* :class:`TokenBucket` —— 平滑速率搭配突發容量(惰性回填)。 +* :class:`SlidingWindowLimiter` —— 每個滾動視窗固定的呼叫額度(Cloudflare 的 O(1) 加權計數 + 近似)。 +* :func:`throttle` —— 一個讓函式在每個間隔內最多觸發一次的裝飾器。 + +純標準函式庫(``threading`` 用於鎖,``time`` 僅作為預設時鐘);不匯入 ``PySide6``。 + +無頭 API +-------- + +.. code-block:: python + + from je_auto_control import TokenBucket, SlidingWindowLimiter, throttle + + # 每秒 5 個請求,可突發到 10 + bucket = TokenBucket(rate=5, capacity=10) + if bucket.try_acquire(): + call_api() # 非阻塞:False 時略過 / 排入佇列 + bucket.acquire() # 或阻塞直到有 token 釋出 + + # 每 60 秒滾動視窗最多 100 次呼叫 + window = SlidingWindowLimiter(limit=100, window_s=60) + if window.try_acquire(): + call_api() + + @throttle(2.0) # 每 2 秒最多觸發一次 + def on_event(payload): + ... + +``TokenBucket.try_acquire`` 在有 token 時取用;``acquire`` 會阻塞(可選 ``timeout``); +``time_until_available`` 回報等待時間,讓排程器自行調速。每個限制器都接受 ``clock=``(``acquire`` +另接受 ``sleep=``),因此整體可在 CI 以假時鐘演練 —— 沒有真正的延遲。 + +執行器命令 +---------- + +``AC_rate_limit`` 接受限制器 ``name`` 以及 ``rate`` / ``capacity`` / ``n``,嘗試從該具名 token +bucket(首次使用時建立)取用 ``n`` 個 token,回傳 ``{acquired, tokens, wait}``,讓流程可閘控或 +延後某個動作。同一操作亦以 MCP 工具 ``ac_rate_limit`` 以及 Script Builder 中 **Flow** 分類下的 +命令提供。 diff --git a/docs/source/Zh/zh_index.rst b/docs/source/Zh/zh_index.rst index a880b3a0..144e57e5 100644 --- a/docs/source/Zh/zh_index.rst +++ b/docs/source/Zh/zh_index.rst @@ -84,6 +84,7 @@ AutoControl 所有功能的完整使用指南。 doc/new_features/v59_features_doc doc/new_features/v60_features_doc doc/new_features/v61_features_doc + doc/new_features/v62_features_doc doc/ocr_backends/ocr_backends_doc doc/observability/observability_doc doc/operations_layer/operations_layer_doc diff --git a/je_auto_control/__init__.py b/je_auto_control/__init__.py index d4734d03..0fdc4f00 100644 --- a/je_auto_control/__init__.py +++ b/je_auto_control/__init__.py @@ -320,6 +320,10 @@ ClaimsPolicy, ExpiredTokenError, InvalidSignatureError, JwtError, decode_jwt, encode_jwt, ) +# Client-side rate limiting (token bucket / sliding window / throttle) +from je_auto_control.utils.rate_limit import ( + SlidingWindowLimiter, TokenBucket, throttle, +) # Background popup/interrupt watchdog (unattended automation) from je_auto_control.utils.watchdog import ( PopupWatchdog, WatchdogRule, default_popup_watchdog, @@ -798,6 +802,7 @@ def start_autocontrol_gui(*args, **kwargs): "license_findings_to_sarif", "normalize_spdx", "ClaimsPolicy", "ExpiredTokenError", "InvalidSignatureError", "JwtError", "decode_jwt", "encode_jwt", + "SlidingWindowLimiter", "TokenBucket", "throttle", # MCP server "AuditLogger", "HttpMCPServer", "MCPContent", "MCPPrompt", "MCPPromptArgument", "MCPResource", "MCPServer", "MCPTool", diff --git a/je_auto_control/gui/script_builder/command_schema.py b/je_auto_control/gui/script_builder/command_schema.py index 045e4403..83715253 100644 --- a/je_auto_control/gui/script_builder/command_schema.py +++ b/je_auto_control/gui/script_builder/command_schema.py @@ -1363,6 +1363,16 @@ def _add_resilience_specs(specs: List[CommandSpec]) -> None: ), description="Run 'actions' (JSON view) via a named circuit breaker.", )) + specs.append(CommandSpec( + "AC_rate_limit", "Flow", "Rate Limit (Token Bucket)", + fields=( + FieldSpec("name", FieldType.STRING), + FieldSpec("rate", FieldType.FLOAT, optional=True, default=1.0), + FieldSpec("capacity", FieldType.FLOAT, optional=True, default=1.0), + FieldSpec("n", FieldType.FLOAT, optional=True, default=1.0), + ), + description="Try to take 'n' tokens from a named limiter; {acquired, wait}.", + )) def _add_input_macro_specs(specs: List[CommandSpec]) -> None: diff --git a/je_auto_control/utils/executor/action_executor.py b/je_auto_control/utils/executor/action_executor.py index a5d8d047..a923beea 100644 --- a/je_auto_control/utils/executor/action_executor.py +++ b/je_auto_control/utils/executor/action_executor.py @@ -2914,6 +2914,20 @@ def _check_licenses(components: Any, allow: Any = None, return {"violations": violations, "count": len(violations)} +_RATE_LIMITERS: Dict[str, Any] = {} + + +def _rate_limit(name: str, rate: float = 1.0, capacity: float = 1.0, + n: float = 1.0) -> Dict[str, Any]: + """Adapter: try to take ``n`` tokens from a named token-bucket limiter.""" + from je_auto_control.utils.rate_limit import TokenBucket + bucket = _RATE_LIMITERS.setdefault( + name, TokenBucket(float(rate), float(capacity))) + acquired = bucket.try_acquire(float(n)) + return {"acquired": acquired, "tokens": round(bucket.tokens, 4), + "wait": round(bucket.time_until_available(float(n)), 4)} + + def _jwt_encode(claims: Any, key: str, alg: str = "HS256") -> Dict[str, Any]: """Adapter: sign a compact JWT from claims (a dict or JSON string).""" import json @@ -3724,6 +3738,7 @@ def __init__(self): "AC_check_licenses": _check_licenses, "AC_jwt_encode": _jwt_encode, "AC_jwt_decode": _jwt_decode, + "AC_rate_limit": _rate_limit, "AC_generate_sop": _generate_sop, "AC_tween_drag": _tween_drag, "AC_list_plugins": _list_plugins, diff --git a/je_auto_control/utils/mcp_server/tools/_factories.py b/je_auto_control/utils/mcp_server/tools/_factories.py index 9ee0e2b4..edd2452c 100644 --- a/je_auto_control/utils/mcp_server/tools/_factories.py +++ b/je_auto_control/utils/mcp_server/tools/_factories.py @@ -3326,6 +3326,23 @@ def jwt_tools() -> List[MCPTool]: ] +def rate_limit_tools() -> List[MCPTool]: + return [ + MCPTool( + name="ac_rate_limit", + description=("Try to take 'n' tokens from a named token-bucket " + "limiter ('rate' tokens/sec, 'capacity' burst). " + "Returns {acquired, tokens, wait}."), + input_schema=schema( + {"name": {"type": "string"}, "rate": {"type": "number"}, + "capacity": {"type": "number"}, "n": {"type": "number"}}, + ["name"]), + handler=h.rate_limit, + annotations=READ_ONLY, + ), + ] + + def saga_tools() -> List[MCPTool]: return [ MCPTool( @@ -4521,7 +4538,7 @@ def media_assert_tools() -> List[MCPTool]: locale_tools, voice_tools, coordinate_space_tools, loop_guard_tools, process_mining_tools, asset_tools, events_tools, notify_channel_tools, jsonpath_tools, json_schema_tools, vuln_scan_tools, vex_tools, - license_policy_tools, jwt_tools, saga_tools, + license_policy_tools, jwt_tools, rate_limit_tools, saga_tools, decision_table_tools, locator_repair_tools, pii_text_tools, sarif_tools, screen_record_tools, diff --git a/je_auto_control/utils/mcp_server/tools/_handlers.py b/je_auto_control/utils/mcp_server/tools/_handlers.py index d12b7908..aef8502a 100644 --- a/je_auto_control/utils/mcp_server/tools/_handlers.py +++ b/je_auto_control/utils/mcp_server/tools/_handlers.py @@ -1599,6 +1599,18 @@ def jwt_decode(token, key, algorithms=None, audience=None, leeway=0.0): return {"ok": True, "claims": claims} +_RATE_LIMITERS = {} + + +def rate_limit(name, rate=1.0, capacity=1.0, n=1.0): + from je_auto_control.utils.rate_limit import TokenBucket + bucket = _RATE_LIMITERS.setdefault( + name, TokenBucket(float(rate), float(capacity))) + acquired = bucket.try_acquire(float(n)) + return {"acquired": acquired, "tokens": round(bucket.tokens, 4), + "wait": round(bucket.time_until_available(float(n)), 4)} + + def run_saga(steps): from je_auto_control.utils.saga import run_saga as _run result = _run(steps) diff --git a/je_auto_control/utils/rate_limit/__init__.py b/je_auto_control/utils/rate_limit/__init__.py new file mode 100644 index 00000000..bffc5465 --- /dev/null +++ b/je_auto_control/utils/rate_limit/__init__.py @@ -0,0 +1,6 @@ +"""Client-side rate limiting: token bucket, sliding window, throttle.""" +from je_auto_control.utils.rate_limit.rate_limit import ( + SlidingWindowLimiter, TokenBucket, throttle, +) + +__all__ = ["SlidingWindowLimiter", "TokenBucket", "throttle"] diff --git a/je_auto_control/utils/rate_limit/rate_limit.py b/je_auto_control/utils/rate_limit/rate_limit.py new file mode 100644 index 00000000..9c63f0b3 --- /dev/null +++ b/je_auto_control/utils/rate_limit/rate_limit.py @@ -0,0 +1,156 @@ +"""Client-side rate limiting for paced API / action calls. + +The framework had ``RetryPolicy`` / ``CircuitBreaker`` (which *recover* from +failures) and a FIFO ``work_queue``, but nothing to shape the *rate* of calls — +so a flow hammering an external API had no way to stay under a quota. This adds +the two standard limiters plus a leading-edge throttle, all with an injectable +clock so they are deterministic in tests. + +* :class:`TokenBucket` — smooth rate with burst capacity (lazy refill). +* :class:`SlidingWindowLimiter` — a fixed call budget per rolling window + (Cloudflare's O(1) weighted-counter approximation). +* :func:`throttle` — a decorator that fires a function at most once per interval. + +Pure standard library (``threading`` for the lock, ``time`` only as the default +clock); imports no ``PySide6``. +""" +import functools +import threading +import time +from typing import Callable, Optional + +from je_auto_control.utils.exception.exceptions import AutoControlException + + +class TokenBucket: + """A token-bucket limiter: ``rate`` tokens/sec up to ``capacity`` burst.""" + + def __init__(self, rate: float, capacity: float, *, + clock: Callable[[], float] = time.monotonic) -> None: + if rate <= 0 or capacity <= 0: + raise AutoControlException("rate and capacity must be positive") + self._rate = float(rate) + self._capacity = float(capacity) + self._clock = clock + self._tokens = float(capacity) + self._updated = clock() + self._lock = threading.Lock() + + def _refill(self) -> None: + now = self._clock() + elapsed = now - self._updated + if elapsed > 0: + self._tokens = min(self._capacity, self._tokens + elapsed * self._rate) + self._updated = now + + @property + def tokens(self) -> float: + """Current token count after refilling for elapsed time.""" + with self._lock: + self._refill() + return self._tokens + + def try_acquire(self, n: float = 1.0) -> bool: + """Take ``n`` tokens if available; return whether it succeeded.""" + with self._lock: + self._refill() + if self._tokens >= n: + self._tokens -= n + return True + return False + + def time_until_available(self, n: float = 1.0) -> float: + """Seconds until ``n`` tokens would be available (0 if already).""" + with self._lock: + self._refill() + if self._tokens >= n: + return 0.0 + return (n - self._tokens) / self._rate + + def acquire(self, n: float = 1.0, *, timeout: Optional[float] = None, + sleep: Callable[[float], None] = time.sleep) -> bool: + """Block until ``n`` tokens are taken or ``timeout`` elapses.""" + deadline = None if timeout is None else self._clock() + timeout + while True: + if self.try_acquire(n): + return True + wait = self.time_until_available(n) + if deadline is not None and self._clock() + wait > deadline: + return False + sleep(wait if wait > 0 else 0.0) + + +class SlidingWindowLimiter: + """Allow ``limit`` calls per ``window_s`` via a weighted rolling counter.""" + + def __init__(self, limit: int, window_s: float, *, + clock: Callable[[], float] = time.monotonic) -> None: + if limit <= 0 or window_s <= 0: + raise AutoControlException("limit and window_s must be positive") + self._limit = int(limit) + self._window = float(window_s) + self._clock = clock + self._cur_start = clock() + self._cur = 0 + self._prev = 0 + self._lock = threading.Lock() + + def _roll(self) -> None: + elapsed = self._clock() - self._cur_start + if elapsed < self._window: + return + if elapsed < 2 * self._window: + self._prev = self._cur + self._cur_start += self._window + else: + self._prev = 0 + self._cur_start = self._clock() + self._cur = 0 + + def _estimate(self) -> float: + elapsed_in_cur = self._clock() - self._cur_start + weight = max(0.0, (self._window - elapsed_in_cur) / self._window) + return self._prev * weight + self._cur + + def try_acquire(self, n: int = 1) -> bool: + """Record ``n`` calls if the weighted estimate stays under the limit.""" + with self._lock: + self._roll() + if self._estimate() + n <= self._limit: + self._cur += n + return True + return False + + def time_until_available(self, n: int = 1) -> float: + """Seconds until ``n`` more calls would fit (0 if they already do).""" + with self._lock: + self._roll() + if self._estimate() + n <= self._limit: + return 0.0 + return max(0.0, self._window - (self._clock() - self._cur_start)) + + +def throttle(interval_s: float, *, + clock: Callable[[], float] = time.monotonic) -> Callable: + """Decorator: call the wrapped function at most once per ``interval_s``. + + Leading-edge — the first call fires immediately; calls within the interval + are dropped (the wrapper returns ``None``). + """ + def decorator(func: Callable) -> Callable: + state = {"last": None} + lock = threading.Lock() + + @functools.wraps(func) + def wrapper(*args, **kwargs): + with lock: + now = clock() + last = state["last"] + if last is not None and now - last < interval_s: + return None + state["last"] = now + return func(*args, **kwargs) + + return wrapper + + return decorator diff --git a/test/unit_test/headless/test_rate_limit_batch.py b/test/unit_test/headless/test_rate_limit_batch.py new file mode 100644 index 00000000..81f1aad3 --- /dev/null +++ b/test/unit_test/headless/test_rate_limit_batch.py @@ -0,0 +1,134 @@ +"""Headless tests for the rate limiters. Pure stdlib, deterministic clock.""" +import je_auto_control as ac +from je_auto_control.utils.exception.exceptions import AutoControlException +from je_auto_control.utils.rate_limit import ( + SlidingWindowLimiter, TokenBucket, throttle) + +import pytest + + +class FakeClock: + """A manually-advanced monotonic clock for deterministic tests.""" + + def __init__(self) -> None: + self.t = 0.0 + + def __call__(self) -> float: + return self.t + + def advance(self, delta: float) -> None: + self.t += delta + + +def test_token_bucket_burst_then_empty(): + clock = FakeClock() + bucket = TokenBucket(rate=2.0, capacity=5.0, clock=clock) + assert all(bucket.try_acquire() for _ in range(5)) + assert bucket.try_acquire() is False + + +def test_token_bucket_refills_at_rate(): + clock = FakeClock() + bucket = TokenBucket(rate=2.0, capacity=5.0, clock=clock) + for _ in range(5): + bucket.try_acquire() + assert bucket.time_until_available(1) == pytest.approx(0.5) + clock.advance(0.5) + assert bucket.try_acquire() is True + + +def test_token_bucket_caps_at_capacity(): + clock = FakeClock() + bucket = TokenBucket(rate=2.0, capacity=5.0, clock=clock) + bucket.try_acquire() + clock.advance(100) + assert bucket.tokens == pytest.approx(5.0) + + +def test_token_bucket_blocking_acquire_with_fake_sleep(): + clock = FakeClock() + bucket = TokenBucket(rate=1.0, capacity=1.0, clock=clock) + assert bucket.try_acquire() is True # drain + assert bucket.acquire(1.0, sleep=clock.advance) is True + assert clock.t == pytest.approx(1.0) + + +def test_token_bucket_acquire_timeout(): + clock = FakeClock() + bucket = TokenBucket(rate=1.0, capacity=1.0, clock=clock) + bucket.try_acquire() + assert bucket.acquire(1.0, timeout=0.5, sleep=clock.advance) is False + + +def test_token_bucket_rejects_bad_args(): + with pytest.raises(AutoControlException): + TokenBucket(rate=0, capacity=1) + with pytest.raises(AutoControlException): + TokenBucket(rate=1, capacity=-1) + + +def test_sliding_window_limits_per_window(): + clock = FakeClock() + limiter = SlidingWindowLimiter(limit=3, window_s=10, clock=clock) + assert all(limiter.try_acquire() for _ in range(3)) + assert limiter.try_acquire() is False + + +def test_sliding_window_weighted_rollover(): + clock = FakeClock() + limiter = SlidingWindowLimiter(limit=3, window_s=10, clock=clock) + for _ in range(3): + limiter.try_acquire() + clock.advance(10) # new window; previous count still fully weighted + assert limiter.try_acquire() is False + clock.advance(10) # previous window now drops out + assert limiter.try_acquire() is True + + +def test_throttle_leading_edge(): + clock = FakeClock() + calls = [] + + @throttle(5.0, clock=clock) + def record(value): + calls.append(value) + return value + + assert record(1) == 1 + assert record(2) is None # within interval -> dropped + clock.advance(5) + assert record(3) == 3 + assert calls == [1, 3] + + +# --- wiring --------------------------------------------------------------- + +def test_executor_round_trip(): + rec = ac.execute_action([[ + "AC_rate_limit", + {"name": "test-exec", "rate": 1.0, "capacity": 1.0}, + ]]) + payload = next(v for v in rec.values() if isinstance(v, dict)) + assert payload["acquired"] is True + # second immediate call is rate limited (capacity 1, no refill yet) + rec2 = ac.execute_action([[ + "AC_rate_limit", + {"name": "test-exec", "rate": 1.0, "capacity": 1.0}, + ]]) + payload2 = next(v for v in rec2.values() if isinstance(v, dict)) + assert payload2["acquired"] is False + assert payload2["wait"] > 0 + + +def test_wiring(): + assert "AC_rate_limit" in ac.executor.known_commands() + from je_auto_control.utils.mcp_server.tools import build_default_tool_registry + assert "ac_rate_limit" in {t.name for t in build_default_tool_registry()} + from je_auto_control.gui.script_builder.command_schema import _build_specs + assert "AC_rate_limit" in {s.command for s in _build_specs()} + + +def test_facade_exports(): + for attr in ("TokenBucket", "SlidingWindowLimiter", "throttle"): + assert hasattr(ac, attr) + assert attr in ac.__all__