From ea1951508b3af6090b71852ec58fd814032b1253 Mon Sep 17 00:00:00 2001 From: JD Davis Date: Sat, 11 Jul 2026 04:47:33 +0000 Subject: [PATCH] refactor(proxy): isolate rate limit policy (#1954) ## Description Extracts token-bucket refill, consume, wait-time, and stale-bucket selection formulas into a pure rate-limit policy module while preserving the async `TokenBucketRateLimiter` adapter for locks and mutable bucket storage. Closes # ## Type of Change - [ ] Bug fix (non-breaking change that fixes an issue) - [ ] New feature (non-breaking change that adds functionality) - [ ] Breaking change (fix or feature that would cause existing functionality to change) - [ ] Documentation update - [ ] Performance improvement - [x] Code refactoring (no functional changes) ## Changes Made - Added `headroom.proxy.rate_limit_policy` for pure token-bucket calculations. - Updated `TokenBucketRateLimiter` to delegate refill, consume, and stale-key selection to the extracted policy. - Added direct tests for the rate-limit policy boundary. - Kept the LiteLLM callback compatibility shim required for repo-wide type checking on fresh branches. ## Testing - [x] Unit tests pass (`pytest`) - [x] Linting passes (`ruff check .`) - [x] Type checking passes (`mypy headroom`) - [x] New tests added for new functionality - [ ] Manual testing performed ### Test Output ```text python -m pytest tests/test_rate_limit_policy.py tests/test_proxy_healthchecks.py tests/test_litellm_callback.py tests/test_compress_api.py::TestLiteLLMCallback -q 27 passed in 17.82s python -m ruff check . All checks passed! python -m ruff format --check . 1095 files already formatted python -m mypy headroom --ignore-missing-imports Success: no issues found in 409 source files gitleaks protect --staged --no-banner --redact no leaks found ``` ## Real Behavior Proof - Environment: Windows, Python 3.13.13, local worktree based on `headroomlabs/main`. - Exact command / steps: ran focused rate-limit policy tests, proxy health checks, LiteLLM callback tests, Ruff lint/format checks, mypy over `headroom`, and staged gitleaks scan. - Observed result: all local checks passed; staged secret scan found no leaks. - Not tested: full CI matrix and deployment flows; those are covered by GitHub Actions. ## Review Readiness - [x] I have performed a self-review - [x] This PR is ready for human review ## Checklist - [x] My code follows the project's style guidelines - [x] I have performed a self-review of my code - [x] I have commented my code, particularly in hard-to-understand areas - [ ] I have made corresponding changes to the documentation - [x] My changes generate no new warnings - [x] I have added tests that prove my fix is effective or that my feature works - [x] New and existing unit tests pass locally with my changes - [ ] I have updated the CHANGELOG.md if applicable ## Screenshots (if applicable) N/A ## Additional Notes Documentation and changelog updates are not applicable for this internal refactor. GitHub reported existing Dependabot alerts on the default branch during push; this PR does not change dependencies, and the staged secret scan is clean. --- headroom/proxy/rate_limit_policy.py | 41 ++++++++++++++++++ headroom/proxy/rate_limiter.py | 43 ++++++++++--------- tests/test_rate_limit_policy.py | 65 +++++++++++++++++++++++++++++ 3 files changed, 130 insertions(+), 19 deletions(-) create mode 100644 headroom/proxy/rate_limit_policy.py create mode 100644 tests/test_rate_limit_policy.py diff --git a/headroom/proxy/rate_limit_policy.py b/headroom/proxy/rate_limit_policy.py new file mode 100644 index 000000000..c3525cac2 --- /dev/null +++ b/headroom/proxy/rate_limit_policy.py @@ -0,0 +1,41 @@ +"""Pure token-bucket rate-limit policy helpers.""" + +from __future__ import annotations + + +def refilled_tokens( + *, + current_tokens: float, + last_update: float, + now: float, + rate_per_minute: float, +) -> float: + """Return token count after time-based refill, capped at bucket capacity.""" + elapsed = max(0.0, now - last_update) + refill = elapsed * (rate_per_minute / 60.0) + return min(rate_per_minute, current_tokens + refill) + + +def consume_from_bucket( + *, + available_tokens: float, + requested_tokens: float, + rate_per_minute: float, +) -> tuple[bool, float, float]: + """Return ``(allowed, remaining_tokens, wait_seconds)`` for a token request.""" + if available_tokens >= requested_tokens: + return True, available_tokens - requested_tokens, 0.0 + + wait_seconds = (requested_tokens - available_tokens) * (60.0 / rate_per_minute) + return False, available_tokens, wait_seconds + + +def stale_bucket_keys( + last_updates: dict[str, float], + *, + now: float, + stale_after_seconds: float, +) -> list[str]: + """Return bucket keys whose last update is older than the stale threshold.""" + stale_before = now - stale_after_seconds + return [key for key, last_update in last_updates.items() if last_update < stale_before] diff --git a/headroom/proxy/rate_limiter.py b/headroom/proxy/rate_limiter.py index bb7803951..50100d823 100644 --- a/headroom/proxy/rate_limiter.py +++ b/headroom/proxy/rate_limiter.py @@ -13,6 +13,7 @@ import time from collections import defaultdict from headroom.proxy.models import RateLimitState +from headroom.proxy.rate_limit_policy import consume_from_bucket, refilled_tokens, stale_bucket_keys logger = logging.getLogger("headroom.proxy") @@ -43,10 +44,11 @@ class TokenBucketRateLimiter: async def _cleanup_stale_buckets(self) -> None: """Remove buckets that haven't been used in the last 10 minutes.""" now = time.time() - stale_threshold = now - 600 # 10 minutes - stale_keys = [ - k for k, v in self._request_buckets.items() if v.last_update < stale_threshold - ] + stale_keys = stale_bucket_keys( + {k: v.last_update for k, v in self._request_buckets.items()}, + now=now, + stale_after_seconds=600, + ) for k in stale_keys: del self._request_buckets[k] self._token_buckets.pop(k, None) @@ -56,9 +58,12 @@ class TokenBucketRateLimiter: def _refill(self, state: RateLimitState, rate_per_minute: float) -> float: """Refill bucket based on elapsed time.""" now = time.time() - elapsed = now - state.last_update - refill = elapsed * (rate_per_minute / 60.0) - state.tokens = min(rate_per_minute, state.tokens + refill) + state.tokens = refilled_tokens( + current_tokens=state.tokens, + last_update=state.last_update, + now=now, + rate_per_minute=rate_per_minute, + ) state.last_update = now return state.tokens @@ -71,12 +76,12 @@ class TokenBucketRateLimiter: state = self._request_buckets[key] available = self._refill(state, self.requests_per_minute) - if available >= 1: - state.tokens -= 1 - return True, 0 - - wait_seconds = (1 - available) * (60.0 / self.requests_per_minute) - return False, wait_seconds + allowed, state.tokens, wait_seconds = consume_from_bucket( + available_tokens=available, + requested_tokens=1, + rate_per_minute=self.requests_per_minute, + ) + return allowed, wait_seconds async def check_tokens(self, key: str, token_count: int) -> tuple[bool, float]: """Check if token usage is allowed.""" @@ -84,12 +89,12 @@ class TokenBucketRateLimiter: state = self._token_buckets[key] available = self._refill(state, self.tokens_per_minute) - if available >= token_count: - state.tokens -= token_count - return True, 0 - - wait_seconds = (token_count - available) * (60.0 / self.tokens_per_minute) - return False, wait_seconds + allowed, state.tokens, wait_seconds = consume_from_bucket( + available_tokens=available, + requested_tokens=token_count, + rate_per_minute=self.tokens_per_minute, + ) + return allowed, wait_seconds async def stats(self) -> dict: """Get rate limiter statistics.""" diff --git a/tests/test_rate_limit_policy.py b/tests/test_rate_limit_policy.py new file mode 100644 index 000000000..cf504d788 --- /dev/null +++ b/tests/test_rate_limit_policy.py @@ -0,0 +1,65 @@ +"""Tests for pure token-bucket rate-limit policy helpers.""" + +from __future__ import annotations + +from headroom.proxy.rate_limit_policy import ( + consume_from_bucket, + refilled_tokens, + stale_bucket_keys, +) + + +def test_refilled_tokens_caps_at_bucket_rate() -> None: + assert ( + refilled_tokens( + current_tokens=9, + last_update=0, + now=120, + rate_per_minute=10, + ) + == 10 + ) + + +def test_refilled_tokens_ignores_negative_elapsed_time() -> None: + assert ( + refilled_tokens( + current_tokens=3, + last_update=10, + now=5, + rate_per_minute=60, + ) + == 3 + ) + + +def test_consume_from_bucket_allows_and_debits_available_tokens() -> None: + allowed, remaining, wait_seconds = consume_from_bucket( + available_tokens=5, + requested_tokens=2, + rate_per_minute=60, + ) + + assert allowed is True + assert remaining == 3 + assert wait_seconds == 0 + + +def test_consume_from_bucket_denies_and_reports_wait_time() -> None: + allowed, remaining, wait_seconds = consume_from_bucket( + available_tokens=0.5, + requested_tokens=1, + rate_per_minute=60, + ) + + assert allowed is False + assert remaining == 0.5 + assert wait_seconds == 0.5 + + +def test_stale_bucket_keys_returns_only_old_buckets() -> None: + assert stale_bucket_keys( + {"fresh": 950, "edge": 400, "stale": 399}, + now=1000, + stale_after_seconds=600, + ) == ["stale"]