From f804cb59ec7ca8a611174cebca8df87766262e37 Mon Sep 17 00:00:00 2001
From: Daniel Iulian Iacob <daniel.iacob@aikido.dev>
Date: Tue, 27 Oct 2026 20:23:39 +0100
Subject: [PATCH] Add Retry-After to rate-limited responses

---
 .github/workflows/qa-tests.yml                |   2 +-
 aikido_zen/middleware/asgi.py                 |  13 +-
 aikido_zen/middleware/asgi_test.py            |  60 ++++
 aikido_zen/middleware/django.py               |   5 +-
 aikido_zen/middleware/django_test.py          |  49 +++
 aikido_zen/middleware/flask.py                |   2 +
 aikido_zen/middleware/flask_test.py           |  47 +++
 aikido_zen/middleware/init_test.py            |  14 +
 aikido_zen/middleware/should_block_request.py |  17 +-
 aikido_zen/ratelimiting/__init__.py           |  30 +-
 aikido_zen/ratelimiting/init_test.py          |  54 ++-
 aikido_zen/ratelimiting/rate_limiter.py       |  19 +-
 aikido_zen/ratelimiting/rate_limiter_test.py  | 312 +++++++++++++-----
 end2end/fastapi_postgres_uvicorn_test.py      |  18 +
 end2end/flask_mysql_test.py                   |   7 +-
 15 files changed, 532 insertions(+), 117 deletions(-)
 create mode 100644 aikido_zen/middleware/asgi_test.py
 create mode 100644 aikido_zen/middleware/django_test.py
 create mode 100644 aikido_zen/middleware/flask_test.py

diff --git a/.github/workflows/qa-tests.yml b/.github/workflows/qa-tests.yml
index e09871c8a..dd87d47db 100644
--- a/.github/workflows/qa-tests.yml
+++ b/.github/workflows/qa-tests.yml
@@ -53,4 +53,4 @@ jobs:
           dockerfile_path: ./zen-demo-python/Dockerfile
           app_port: 8080
           config_propagation_delay: 15
-          skip_tests: test-path-traversal,test-outbound-domain-blocking,test-wave-attack,test-rate-limiting-retry-after
+          skip_tests: test-path-traversal,test-outbound-domain-blocking,test-wave-attack
diff --git a/aikido_zen/middleware/asgi.py b/aikido_zen/middleware/asgi.py
index 236b71d70..0c130210c 100644
--- a/aikido_zen/middleware/asgi.py
+++ b/aikido_zen/middleware/asgi.py
@@ -19,7 +19,10 @@ async def __call__(self, scope, receive, send):
             message = "You are rate limited by Zen."
             if result["trigger"] == "ip" and result["ip"]:
                 message += " (Your IP: " + result["ip"] + ")"
-            return await send_status_code_and_text(send, (message, 429))
+            headers = {}
+            if result["retry_after_seconds"]:
+                headers["Retry-After"] = str(result["retry_after_seconds"])
+            return await send_status_code_and_text(send, (message, 429), headers)
 
         if result["type"] == "blocked":
             return await send_status_code_and_text(
@@ -30,13 +33,17 @@ async def __call__(self, scope, receive, send):
         return await self.app(scope, receive, send)
 
 
-async def send_status_code_and_text(send, pre_response):
+async def send_status_code_and_text(send, pre_response, headers=None):
     """Sends a status code and text"""
+    raw_headers = [(b"content-type", b"text/plain")]
+    for name, value in (headers or {}).items():
+        # ASGI requires lowercase header names as bytes.
+        raw_headers.append((name.lower().encode(), value.encode()))
     await send(
         {
             "type": "http.response.start",
             "status": pre_response[1],
-            "headers": [(b"content-type", b"text/plain")],
+            "headers": raw_headers,
         }
     )
     await send(
diff --git a/aikido_zen/middleware/asgi_test.py b/aikido_zen/middleware/asgi_test.py
new file mode 100644
index 000000000..a30087e3d
--- /dev/null
+++ b/aikido_zen/middleware/asgi_test.py
@@ -0,0 +1,60 @@
+from unittest.mock import patch
+
+import pytest
+
+from .asgi import AikidoASGIMiddleware
+
+
+async def call_middleware(result):
+    sent = []
+
+    async def app(scope, receive, send):
+        raise AssertionError("app should not be called")
+
+    async def send(message):
+        sent.append(message)
+
+    with patch("aikido_zen.middleware.asgi.should_block_request", return_value=result):
+        await AikidoASGIMiddleware(app)({"type": "http"}, None, send)
+    return sent
+
+
+@pytest.mark.asyncio
+async def test_rate_limited_response_has_retry_after():
+    sent = await call_middleware(
+        {
+            "block": True,
+            "type": "ratelimited",
+            "trigger": "user",
+            "ip": "1.1.1.1",
+            "retry_after_seconds": 7,
+        }
+    )
+    assert sent[0]["status"] == 429
+    assert sent[0]["headers"] == [
+        (b"content-type", b"text/plain"),
+        (b"retry-after", b"7"),
+    ]
+    assert sent[1]["body"] == b"You are rate limited by Zen."
+
+
+@pytest.mark.asyncio
+async def test_rate_limited_response_without_retry_after_seconds():
+    sent = await call_middleware(
+        {
+            "block": True,
+            "type": "ratelimited",
+            "trigger": "user",
+            "ip": "1.1.1.1",
+            "retry_after_seconds": None,
+        }
+    )
+    assert sent[0]["status"] == 429
+    assert sent[0]["headers"] == [(b"content-type", b"text/plain")]
+
+
+@pytest.mark.asyncio
+async def test_blocked_user_response_has_no_retry_after():
+    sent = await call_middleware({"block": True, "type": "blocked", "trigger": "user"})
+    assert sent[0]["status"] == 403
+    assert sent[0]["headers"] == [(b"content-type", b"text/plain")]
diff --git a/aikido_zen/middleware/django.py b/aikido_zen/middleware/django.py
index 5bfd1767d..5646a2e68 100644
--- a/aikido_zen/middleware/django.py
+++ b/aikido_zen/middleware/django.py
@@ -27,7 +27,10 @@ def __call__(self, request):
             message = "You are rate limited by Zen."
             if result["trigger"] == "ip" and result["ip"]:
                 message += " (Your IP: " + result["ip"] + ")"
-            return self.HttpResponse(message, content_type="text/plain", status=429)
+            response = self.HttpResponse(message, content_type="text/plain", status=429)
+            if result["retry_after_seconds"]:
+                response["Retry-After"] = str(result["retry_after_seconds"])
+            return response
 
         if result["type"] == "blocked":
             return self.HttpResponse(
diff --git a/aikido_zen/middleware/django_test.py b/aikido_zen/middleware/django_test.py
new file mode 100644
index 000000000..610e480f2
--- /dev/null
+++ b/aikido_zen/middleware/django_test.py
@@ -0,0 +1,49 @@
+from unittest.mock import patch
+
+from .django import AikidoDjangoMiddleware
+
+
+class FakeHttpResponse(dict):
+    def __init__(self, content, content_type=None, status=200):
+        super().__init__()
+        self.content = content
+        self.content_type = content_type
+        self.status_code = status
+
+
+def call_middleware(result):
+    middleware = AikidoDjangoMiddleware(get_response=None)
+    middleware.HttpResponse = FakeHttpResponse
+    with patch(
+        "aikido_zen.middleware.django.should_block_request", return_value=result
+    ):
+        return middleware(request=None)
+
+
+def test_rate_limited_response_has_retry_after():
+    response = call_middleware(
+        {
+            "block": True,
+            "type": "ratelimited",
+            "trigger": "user",
+            "ip": "1.1.1.1",
+            "retry_after_seconds": 7,
+        }
+    )
+    assert response.status_code == 429
+    assert response["Retry-After"] == "7"
+    assert response.content == "You are rate limited by Zen."
+
+
+def test_rate_limited_response_without_retry_after_seconds():
+    response = call_middleware(
+        {
+            "block": True,
+            "type": "ratelimited",
+            "trigger": "user",
+            "ip": "1.1.1.1",
+            "retry_after_seconds": None,
+        }
+    )
+    assert response.status_code == 429
+    assert "Retry-After" not in response
diff --git a/aikido_zen/middleware/flask.py b/aikido_zen/middleware/flask.py
index 585718fbe..09d561f79 100644
--- a/aikido_zen/middleware/flask.py
+++ b/aikido_zen/middleware/flask.py
@@ -28,6 +28,8 @@ def __call__(self, environ, start_response):
             if result["trigger"] == "ip" and result["ip"]:
                 message += " (Your IP: " + result["ip"] + ")"
             res = self.Response(message, mimetype="text/plain", status=429)
+            if result["retry_after_seconds"]:
+                res.headers["Retry-After"] = str(result["retry_after_seconds"])
             return res(environ, start_response)
 
         if result["type"] == "blocked":
diff --git a/aikido_zen/middleware/flask_test.py b/aikido_zen/middleware/flask_test.py
new file mode 100644
index 000000000..e5afdd2c8
--- /dev/null
+++ b/aikido_zen/middleware/flask_test.py
@@ -0,0 +1,47 @@
+from unittest.mock import patch
+
+from werkzeug.test import create_environ
+
+from .flask import AikidoFlaskMiddleware
+
+
+def call_middleware(result):
+    captured = {}
+
+    def start_response(status, headers, exc_info=None):
+        captured["status"] = status
+        captured["headers"] = dict(headers)
+
+    with patch("aikido_zen.middleware.flask.should_block_request", return_value=result):
+        middleware = AikidoFlaskMiddleware(app=None)
+        captured["body"] = b"".join(middleware(create_environ("/"), start_response))
+    return captured
+
+
+def test_rate_limited_response_has_retry_after():
+    response = call_middleware(
+        {
+            "block": True,
+            "type": "ratelimited",
+            "trigger": "ip",
+            "ip": "1.1.1.1",
+            "retry_after_seconds": 7,
+        }
+    )
+    assert response["status"] == "429 TOO MANY REQUESTS"
+    assert response["headers"]["Retry-After"] == "7"
+    assert response["body"] == b"You are rate limited by Zen. (Your IP: 1.1.1.1)"
+
+
+def test_rate_limited_response_without_retry_after_seconds():
+    response = call_middleware(
+        {
+            "block": True,
+            "type": "ratelimited",
+            "trigger": "user",
+            "ip": "1.1.1.1",
+            "retry_after_seconds": None,
+        }
+    )
+    assert response["status"] == "429 TOO MANY REQUESTS"
+    assert "Retry-After" not in response["headers"]
diff --git a/aikido_zen/middleware/init_test.py b/aikido_zen/middleware/init_test.py
index 4a7fad1f9..3d7bc0560 100644
--- a/aikido_zen/middleware/init_test.py
+++ b/aikido_zen/middleware/init_test.py
@@ -161,5 +161,19 @@ def test_cache_comms_with_endpoints():
             "ip": "1.1.1.1",
             "type": "ratelimited",
             "trigger": "my_trigger",
+            "retry_after_seconds": None,
         }
         assert thread_cache.stats.rate_limited_hits == 1
+
+        mock_comms.send_data_to_bg_process.return_value = {
+            "success": True,
+            "data": {"block": True, "trigger": "my_trigger", "retry_after_seconds": 7},
+        }
+        assert should_block_request() == {
+            "block": True,
+            "ip": "1.1.1.1",
+            "type": "ratelimited",
+            "trigger": "my_trigger",
+            "retry_after_seconds": 7,
+        }
+        assert thread_cache.stats.rate_limited_hits == 2
diff --git a/aikido_zen/middleware/should_block_request.py b/aikido_zen/middleware/should_block_request.py
index 434f06a67..f13107eaa 100644
--- a/aikido_zen/middleware/should_block_request.py
+++ b/aikido_zen/middleware/should_block_request.py
@@ -11,14 +11,12 @@
 def should_block_request():
     """
     Checks for rate-limiting and checks if the current user is blocked.
-    Returns: {
-        block: boolean,
-        data: {
-            type: "blocked" | "ratelimited",
-            trigger: "ip" | "user" | "group",
-            ip?: string
-        }
-     }
+    Returns one of:
+        {"block": False}
+        {"block": True, "type": "blocked", "trigger": "user"}
+        {"block": True, "type": "ratelimited", "trigger": "ip" | "user" | "group",
+         "ip": str | None, "retry_after_seconds": int | None}
+    retry_after_seconds is None only when an older background process answered.
     """
     try:
         context = get_current_context()
@@ -69,6 +67,9 @@ def should_block_request():
                     "type": "ratelimited",
                     "trigger": ratelimit_res["data"]["trigger"],
                     "ip": context.remote_address,
+                    "retry_after_seconds": ratelimit_res["data"].get(
+                        "retry_after_seconds"
+                    ),
                 }
     except Exception as e:
         logger.debug("Exception occurred in should_block_request: %s", e)
diff --git a/aikido_zen/ratelimiting/__init__.py b/aikido_zen/ratelimiting/__init__.py
index 4d60bdd01..88d8f03b7 100644
--- a/aikido_zen/ratelimiting/__init__.py
+++ b/aikido_zen/ratelimiting/__init__.py
@@ -28,34 +28,46 @@ def should_ratelimit_request(
     windows_size_in_ms = int(endpoint["rateLimiting"]["windowSizeInMS"])
 
     if group:
-        allowed = connection_manager.rate_limiter.is_allowed(
+        result = connection_manager.rate_limiter.check(
             get_key_for_group(endpoint, group),
             windows_size_in_ms,
             max_requests,
         )
-        if not allowed:
-            return {"block": True, "trigger": "group"}
+        if not result["allowed"]:
+            return {
+                "block": True,
+                "trigger": "group",
+                "retry_after_seconds": result["retry_after_seconds"],
+            }
 
         # Do not check IP or user rate limit if group is set
         return {"block": False}
     if user:
-        allowed = connection_manager.rate_limiter.is_allowed(
+        result = connection_manager.rate_limiter.check(
             get_key_for_user(endpoint, user),
             windows_size_in_ms,
             max_requests,
         )
-        if not allowed:
-            return {"block": True, "trigger": "user"}
+        if not result["allowed"]:
+            return {
+                "block": True,
+                "trigger": "user",
+                "retry_after_seconds": result["retry_after_seconds"],
+            }
         # Do not check IP rate limit if user is set
         return {"block": False}
     if remote_address:
-        allowed = connection_manager.rate_limiter.is_allowed(
+        result = connection_manager.rate_limiter.check(
             get_key_for_ip(endpoint, remote_address),
             windows_size_in_ms,
             max_requests,
         )
-        if not allowed:
-            return {"block": True, "trigger": "ip"}
+        if not result["allowed"]:
+            return {
+                "block": True,
+                "trigger": "ip",
+                "retry_after_seconds": result["retry_after_seconds"],
+            }
 
     return {"block": False}
 
diff --git a/aikido_zen/ratelimiting/init_test.py b/aikido_zen/ratelimiting/init_test.py
index dfd14acf6..f731ebc9e 100644
--- a/aikido_zen/ratelimiting/init_test.py
+++ b/aikido_zen/ratelimiting/init_test.py
@@ -1,5 +1,5 @@
 import pytest
-from unittest.mock import MagicMock
+from unittest.mock import ANY, MagicMock, patch
 from aikido_zen.background_process.service_config import ServiceConfig
 from aikido_zen.ratelimiting.rate_limiter import RateLimiter
 from . import should_ratelimit_request
@@ -63,6 +63,7 @@ def test_rate_limits_by_ip():
     assert should_ratelimit_request(route_metadata, "1.2.3.4", None, cm) == {
         "block": True,
         "trigger": "ip",
+        "retry_after_seconds": ANY,
     }
 
 
@@ -129,6 +130,7 @@ def test_rate_limiting_by_user(user):
     assert should_ratelimit_request(route_metadata, "1.2.3.7", user, cm) == {
         "block": True,
         "trigger": "user",
+        "retry_after_seconds": ANY,
     }
 
 
@@ -162,7 +164,7 @@ def test_rate_limiting_with_wildcard():
     # This request should trigger the rate limit
     assert should_ratelimit_request(
         create_route_metadata(route="/api/login"), "1.2.3.4", None, cm
-    ) == {"block": True, "trigger": "ip"}
+    ) == {"block": True, "trigger": "ip", "retry_after_seconds": ANY}
 
 
 def test_rate_limiting_with_wildcard2():
@@ -196,6 +198,7 @@ def test_rate_limiting_with_wildcard2():
     assert should_ratelimit_request(metadata, "1.2.3.4", None, cm) == {
         "block": True,
         "trigger": "ip",
+        "retry_after_seconds": ANY,
     }
 
 
@@ -231,6 +234,7 @@ def test_rate_limiting_by_user_with_same_ip():
     assert should_ratelimit_request(metadata, "1.2.3.4", {"id": "123"}, cm) == {
         "block": True,
         "trigger": "user",
+        "retry_after_seconds": ANY,
     }
 
 
@@ -270,6 +274,7 @@ def test_rate_limiting_by_user_with_different_ips():
     assert should_ratelimit_request(metadata, "4.3.2.1", {"id": "123"}, cm) == {
         "block": True,
         "trigger": "user",
+        "retry_after_seconds": ANY,
     }
 
 
@@ -388,6 +393,7 @@ def test_rate_limits_by_user_with_different_ips():
     assert should_ratelimit_request(route_metadata, "4.3.2.1", user, cm, "group1") == {
         "block": True,
         "trigger": "group",
+        "retry_after_seconds": ANY,
     }
 
 
@@ -425,6 +431,7 @@ def test_rate_limits_different_users_in_same_group():
     ) == {
         "block": True,
         "trigger": "group",
+        "retry_after_seconds": ANY,
     }
 
 
@@ -462,6 +469,7 @@ def test_works_with_multiple_rate_limit_groups_and_different_users():
     ) == {
         "block": True,
         "trigger": "group",
+        "retry_after_seconds": ANY,
     }
     # This request should also trigger the rate limit for group1
     assert should_ratelimit_request(
@@ -469,6 +477,7 @@ def test_works_with_multiple_rate_limit_groups_and_different_users():
     ) == {
         "block": True,
         "trigger": "group",
+        "retry_after_seconds": ANY,
     }
     # First request from user 953, group2
     assert should_ratelimit_request(
@@ -480,6 +489,7 @@ def test_works_with_multiple_rate_limit_groups_and_different_users():
     ) == {
         "block": True,
         "trigger": "group",
+        "retry_after_seconds": ANY,
     }
 
 
@@ -515,6 +525,7 @@ def test_rate_limits_by_group_if_user_is_not_set():
     assert should_ratelimit_request(route_metadata, "4.3.2.1", None, cm, "group1") == {
         "block": True,
         "trigger": "group",
+        "retry_after_seconds": ANY,
     }
 
 
@@ -564,3 +575,42 @@ def test_does_not_rate_limit_excluded_users_in_group():
         assert should_ratelimit_request(
             route_metadata, "1.2.3.4", excluded_user, cm, "group1"
         ) == {"block": False}
+
+
+@pytest.mark.parametrize(
+    "user, group, trigger",
+    [
+        (None, None, "ip"),
+        ({"id": "123"}, None, "user"),
+        ({"id": "123"}, "group1", "group"),
+    ],
+)
+def test_returns_retry_after_seconds(user, group, trigger):
+    cm = create_connection_manager(
+        [
+            {
+                "method": "POST",
+                "route": "/login",
+                "forceProtectionOff": False,
+                "rateLimiting": {
+                    "enabled": True,
+                    "maxRequests": 3,
+                    "windowSizeInMS": 60_000,
+                },
+            },
+        ]
+    )
+    route_metadata = create_route_metadata()
+    with patch("aikido_zen.ratelimiting.rate_limiter.get_unixtime_ms") as clock:
+        for time_in_ms in [0, 20_000, 40_000]:
+            clock.return_value = time_in_ms
+            assert should_ratelimit_request(
+                route_metadata, "1.2.3.4", user, cm, group
+            ) == {"block": False}
+
+        clock.return_value = 50_000
+        assert should_ratelimit_request(route_metadata, "1.2.3.4", user, cm, group) == {
+            "block": True,
+            "trigger": trigger,
+            "retry_after_seconds": 31,
+        }
diff --git a/aikido_zen/ratelimiting/rate_limiter.py b/aikido_zen/ratelimiting/rate_limiter.py
index 7efdca860..da4081984 100644
--- a/aikido_zen/ratelimiting/rate_limiter.py
+++ b/aikido_zen/ratelimiting/rate_limiter.py
@@ -16,9 +16,10 @@ def __init__(self, max_items, time_to_live_in_ms):
         self.time_to_live_in_ms = time_to_live_in_ms
         self.rate_limited_items = LRUCache(max_items, time_to_live_in_ms)
 
-    def is_allowed(self, key, window_size_in_ms, max_requests):
+    def check(self, key, window_size_in_ms, max_requests):
         """
-        Checks if the request is allowed given the history
+        Records a request for this key and checks it against the limit.
+        Returns {"allowed": True}, or {"allowed": False, "retry_after_seconds": int}.
         """
         current_time = get_unixtime_ms()
         request_timestamps = self.rate_limited_items.get(key) or []
@@ -40,4 +41,16 @@ def is_allowed(self, key, window_size_in_ms, max_requests):
         self.rate_limited_items.set(key, request_timestamps)
 
         # if the total amount of requests in the current window exceeds max requests, we rate-limit
-        return len(request_timestamps) <= max_requests
+        if len(request_timestamps) <= max_requests:
+            return {"allowed": True}
+
+        # Blocked requests are stored too, so this timestamp, not the oldest one, must
+        # leave the window first. max() keeps the index valid when maxRequests < 1.
+        blocking_timestamp = request_timestamps[-max(max_requests, 1)]
+
+        # A timestamp exactly one window old still counts, so the wait must be longer.
+        retry_after_ms = blocking_timestamp + window_size_in_ms - current_time
+        return {
+            "allowed": False,
+            "retry_after_seconds": max(1, retry_after_ms // 1000 + 1),
+        }
diff --git a/aikido_zen/ratelimiting/rate_limiter_test.py b/aikido_zen/ratelimiting/rate_limiter_test.py
index 76ec3d9a4..264e6cb87 100644
--- a/aikido_zen/ratelimiting/rate_limiter_test.py
+++ b/aikido_zen/ratelimiting/rate_limiter_test.py
@@ -1,5 +1,8 @@
-import pytest
+import random
 import time
+from unittest.mock import patch
+
+import pytest
 from .rate_limiter import RateLimiter
 
 
@@ -13,28 +16,30 @@ def rate_limiter():
 def test_allow_requests_within_limit(rate_limiter):
     key = "user1"
     for i in range(5):
-        assert rate_limiter.is_allowed(
-            key, 5000000, 5
-        ), f"Request {i + 1} should be allowed"
+        assert rate_limiter.check(key, 5000000, 5)[
+            "allowed"
+        ], f"Request {i + 1} should be allowed"
 
 
 def test_deny_requests_exceeding_limit(rate_limiter):
     key = "user2"
     for i in range(5):
-        rate_limiter.is_allowed(key, 500, 5)
-    assert not rate_limiter.is_allowed(key, 500, 5), "Request 6 should not be allowed"
+        rate_limiter.check(key, 500, 5)
+    assert not rate_limiter.check(key, 500, 5)[
+        "allowed"
+    ], "Request 6 should not be allowed"
 
 
 def test_clear_old_entries(rate_limiter):
     key = "user3"
     for i in range(5):
-        rate_limiter.is_allowed(key, 500, 5)
+        rate_limiter.check(key, 500, 5)
 
     time.sleep(0.6)  # Sleep to allow old entries to be cleared
 
-    assert rate_limiter.is_allowed(
-        key, 500, 5
-    ), "New request should be allowed after clearing old entries"
+    assert rate_limiter.check(key, 500, 5)[
+        "allowed"
+    ], "New request should be allowed after clearing old entries"
 
 
 def test_multiple_keys(rate_limiter):
@@ -42,87 +47,93 @@ def test_multiple_keys(rate_limiter):
     key2 = "user5"
 
     for i in range(5):
-        assert rate_limiter.is_allowed(
-            key1, 500, 5
-        ), f"Request {i + 1} for key1 should be allowed"
+        assert rate_limiter.check(key1, 500, 5)[
+            "allowed"
+        ], f"Request {i + 1} for key1 should be allowed"
     for i in range(5):
-        assert rate_limiter.is_allowed(
-            key2, 500, 5
-        ), f"Request {i + 1} for key2 should be allowed"
+        assert rate_limiter.check(key2, 500, 5)[
+            "allowed"
+        ], f"Request {i + 1} for key2 should be allowed"
 
-    assert not rate_limiter.is_allowed(
-        key1, 500, 5
-    ), "Request 6 for key1 should not be allowed"
-    assert not rate_limiter.is_allowed(
-        key2, 500, 5
-    ), "Request 6 for key2 should not be allowed"
+    assert not rate_limiter.check(key1, 500, 5)[
+        "allowed"
+    ], "Request 6 for key1 should not be allowed"
+    assert not rate_limiter.check(key2, 500, 5)[
+        "allowed"
+    ], "Request 6 for key2 should not be allowed"
 
 
 def test_ttl_expiration(rate_limiter):
     key = "user6"
     for i in range(5):
-        rate_limiter.is_allowed(key, 500, 5)
+        rate_limiter.check(key, 500, 5)
 
     time.sleep(1.1)  # Sleep to allow the TTL to expire
 
-    assert rate_limiter.is_allowed(key, 500, 5), "Request after TTL should be allowed"
+    assert rate_limiter.check(key, 500, 5)[
+        "allowed"
+    ], "Request after TTL should be allowed"
 
 
 def test_allow_requests_exactly_at_limit(rate_limiter):
     key = "user7"
     for i in range(5):
-        assert rate_limiter.is_allowed(
-            key, 500, 5
-        ), f"Request {i + 1} should be allowed"
-    assert not rate_limiter.is_allowed(key, 500, 5), "Request 6 should not be allowed"
+        assert rate_limiter.check(key, 500, 5)[
+            "allowed"
+        ], f"Request {i + 1} should be allowed"
+    assert not rate_limiter.check(key, 500, 5)[
+        "allowed"
+    ], "Request 6 should not be allowed"
 
 
 def test_allow_requests_after_clearing_old_entries(rate_limiter):
     key = "user8"
     for i in range(5):
-        rate_limiter.is_allowed(key, 500, 5)
+        rate_limiter.check(key, 500, 5)
 
     time.sleep(0.6)  # Sleep to allow old entries to be cleared
 
-    assert rate_limiter.is_allowed(
-        key, 500, 5
-    ), "New request should be allowed after clearing old entries"
+    assert rate_limiter.check(key, 500, 5)[
+        "allowed"
+    ], "New request should be allowed after clearing old entries"
 
 
 def test_multiple_rapid_requests(rate_limiter):
     key = "user9"
     for i in range(5):
-        assert rate_limiter.is_allowed(
-            key, 500, 5
-        ), f"Request {i + 1} should be allowed"
+        assert rate_limiter.check(key, 500, 5)[
+            "allowed"
+        ], f"Request {i + 1} should be allowed"
 
     time.sleep(0.1)  # Sleep for 100 ms
 
-    assert not rate_limiter.is_allowed(
-        key, 500, 5
-    ), "Request after rapid requests should not be allowed"
+    assert not rate_limiter.check(key, 500, 5)[
+        "allowed"
+    ], "Request after rapid requests should not be allowed"
 
 
 def test_reset_after_ttl(rate_limiter):
     key = "user10"
     for i in range(5):
-        rate_limiter.is_allowed(key, 500, 5)
+        rate_limiter.check(key, 500, 5)
 
     time.sleep(1.1)  # Sleep to allow the TTL to expire
 
-    assert rate_limiter.is_allowed(key, 500, 5), "Request after TTL should be allowed"
+    assert rate_limiter.check(key, 500, 5)[
+        "allowed"
+    ], "Request after TTL should be allowed"
 
 
 def test_different_window_sizes(rate_limiter):
     key = "user11"
     different_window_size = 1000  # 1 second window
     for i in range(5):
-        assert rate_limiter.is_allowed(
-            key, different_window_size, 5
-        ), f"Request {i + 1} should be allowed"
-    assert not rate_limiter.is_allowed(
-        key, different_window_size, 5
-    ), "Request 6 should not be allowed"
+        assert rate_limiter.check(key, different_window_size, 5)[
+            "allowed"
+        ], f"Request {i + 1} should be allowed"
+    assert not rate_limiter.check(key, different_window_size, 5)[
+        "allowed"
+    ], "Request 6 should not be allowed"
 
 
 def test_sliding_window_with_intermittent_requests(rate_limiter):
@@ -130,18 +141,18 @@ def test_sliding_window_with_intermittent_requests(rate_limiter):
 
     # Allow 5 requests in a 1-second window
     for i in range(5):
-        assert rate_limiter.is_allowed(
-            key, 500, 5
-        ), f"Request {i + 1} should be allowed"
+        assert rate_limiter.check(key, 500, 5)[
+            "allowed"
+        ], f"Request {i + 1} should be allowed"
         time.sleep(0.1)  # Sleep 100 ms between requests
 
     # Sleep for 600 ms to allow the first requests to slide out of the window
     time.sleep(0.6)
 
     # Now we should be able to make a new request
-    assert rate_limiter.is_allowed(
-        key, 500, 5
-    ), "New request should be allowed after sliding"
+    assert rate_limiter.check(key, 500, 5)[
+        "allowed"
+    ], "New request should be allowed after sliding"
 
 
 def test_sliding_window_edge_case(rate_limiter):
@@ -149,25 +160,25 @@ def test_sliding_window_edge_case(rate_limiter):
 
     # Allow 5 requests in a 1-second window
     for i in range(5):
-        assert rate_limiter.is_allowed(
-            key, 500, 5
-        ), f"Request {i + 1} should be allowed"
+        assert rate_limiter.check(key, 500, 5)[
+            "allowed"
+        ], f"Request {i + 1} should be allowed"
 
     # Sleep for 500 ms to simulate time passing
     time.sleep(0.6)
 
     # The next request should still be allowed as the window is sliding
-    assert rate_limiter.is_allowed(
-        key, 500, 5
-    ), "Next request should be allowed as window slides"
+    assert rate_limiter.check(key, 500, 5)[
+        "allowed"
+    ], "Next request should be allowed as window slides"
 
     # Sleep for another 500 ms to allow the first batch to expire
     time.sleep(0.6)
 
     # Now we should be able to make a new request
-    assert rate_limiter.is_allowed(
-        key, 500, 5
-    ), "New request should be allowed after first batch expires"
+    assert rate_limiter.check(key, 500, 5)[
+        "allowed"
+    ], "New request should be allowed after first batch expires"
 
 
 def test_sliding_window_with_burst_requests(rate_limiter):
@@ -176,23 +187,23 @@ def test_sliding_window_with_burst_requests(rate_limiter):
 
     # Allow 5 requests in a 1-second window
     for i in range(5):
-        assert rate_limiter.is_allowed(
-            key, window_size_ms, 5
-        ), f"Request {i + 1} should be allowed"
+        assert rate_limiter.check(key, window_size_ms, 5)[
+            "allowed"
+        ], f"Request {i + 1} should be allowed"
 
     # Sleep for 250 ms to simulate time passing
     time.sleep((window_size_ms / 2) / 1000)
 
     # Add 3 more requests (should be denied)
-    assert not rate_limiter.is_allowed(
-        key, window_size_ms, 5
-    ), "Request should not be allowed"
-    assert not rate_limiter.is_allowed(
-        key, window_size_ms, 5
-    ), "Request should not be allowed"
-    assert not rate_limiter.is_allowed(
-        key, window_size_ms, 5
-    ), "Request should not be allowed"
+    assert not rate_limiter.check(key, window_size_ms, 5)[
+        "allowed"
+    ], "Request should not be allowed"
+    assert not rate_limiter.check(key, window_size_ms, 5)[
+        "allowed"
+    ], "Request should not be allowed"
+    assert not rate_limiter.check(key, window_size_ms, 5)[
+        "allowed"
+    ], "Request should not be allowed"
 
     time.sleep(
         (window_size_ms / 2 + 50) / 1000
@@ -200,20 +211,20 @@ def test_sliding_window_with_burst_requests(rate_limiter):
 
     # Make a burst of requests (should be allowed)
     for i in range(2):
-        assert rate_limiter.is_allowed(
-            key, window_size_ms, 5
-        ), f"Burst request {i + 1} should be allowed"
-    assert not rate_limiter.is_allowed(
-        key, window_size_ms, 5
-    ), "Burst request should not be allowed after limit"
+        assert rate_limiter.check(key, window_size_ms, 5)[
+            "allowed"
+        ], f"Burst request {i + 1} should be allowed"
+    assert not rate_limiter.check(key, window_size_ms, 5)[
+        "allowed"
+    ], "Burst request should not be allowed after limit"
 
     # Sleep for 500 ms to allow all batches to slide out of window
     time.sleep((window_size_ms + 100) / 1000)
 
     # Now we should be able to make a new request
-    assert rate_limiter.is_allowed(
-        key, window_size_ms, 5
-    ), "New request should be allowed after all batches slide out"
+    assert rate_limiter.check(key, window_size_ms, 5)[
+        "allowed"
+    ], "New request should be allowed after all batches slide out"
 
 
 def test_sliding_window_with_delayed_requests(rate_limiter):
@@ -221,15 +232,140 @@ def test_sliding_window_with_delayed_requests(rate_limiter):
 
     # Allow 5 requests in a 1-second window
     for i in range(5):
-        assert rate_limiter.is_allowed(
-            key, 500, 5
-        ), f"Request {i + 1} should be allowed"
+        assert rate_limiter.check(key, 500, 5)[
+            "allowed"
+        ], f"Request {i + 1} should be allowed"
         time.sleep(0.1)  # Sleep 100 ms between requests
 
     # Sleep for 600 ms to allow the first requests to slide out of the window
     time.sleep(0.6)
 
     # Now we should be able to make a new request
-    assert rate_limiter.is_allowed(
-        key, 500, 5
-    ), "New request should be allowed after sliding"
+    assert rate_limiter.check(key, 500, 5)[
+        "allowed"
+    ], "New request should be allowed after sliding"
+
+
+@pytest.fixture
+def clock():
+    current = {"time_in_ms": 0}
+    with patch(
+        "aikido_zen.ratelimiting.rate_limiter.get_unixtime_ms",
+        side_effect=lambda: current["time_in_ms"],
+    ):
+        yield current
+
+
+def new_rate_limiter():
+    return RateLimiter(max_items=10, time_to_live_in_ms=120 * 60 * 1000)
+
+
+def send_requests(
+    clock, rate_limiter, times_in_ms, max_requests=3, window_size_in_ms=60_000
+):
+    results = []
+    for time_in_ms in times_in_ms:
+        clock["time_in_ms"] = time_in_ms
+        results.append(rate_limiter.check("key", window_size_in_ms, max_requests))
+    return results
+
+
+def test_retry_after_counts_the_blocked_request(clock):
+    history = [0, 20_000, 40_000, 50_000]
+    assert send_requests(clock, new_rate_limiter(), history) == [
+        {"allowed": True},
+        {"allowed": True},
+        {"allowed": True},
+        {"allowed": False, "retry_after_seconds": 31},
+    ]
+
+    assert (
+        send_requests(clock, new_rate_limiter(), history + [80_000])[-1]["allowed"]
+        is False
+    )
+    assert send_requests(clock, new_rate_limiter(), history + [81_000])[-1] == {
+        "allowed": True
+    }
+
+
+def test_retry_after_counts_every_blocked_request(clock):
+    history = [0, 20_000, 40_000, 50_000, 55_000]
+    assert send_requests(clock, new_rate_limiter(), history)[-2:] == [
+        {"allowed": False, "retry_after_seconds": 31},
+        {"allowed": False, "retry_after_seconds": 46},
+    ]
+
+    assert (
+        send_requests(clock, new_rate_limiter(), history + [100_000])[-1]["allowed"]
+        is False
+    )
+    assert send_requests(clock, new_rate_limiter(), history + [101_000])[-1] == {
+        "allowed": True
+    }
+
+
+def test_retry_after_for_requests_in_the_same_millisecond(clock):
+    history = [0, 0, 0, 0]
+    assert send_requests(clock, new_rate_limiter(), history)[-1] == {
+        "allowed": False,
+        "retry_after_seconds": 61,
+    }
+
+    assert (
+        send_requests(clock, new_rate_limiter(), history + [60_000])[-1]["allowed"]
+        is False
+    )
+    assert send_requests(clock, new_rate_limiter(), history + [61_000])[-1] == {
+        "allowed": True
+    }
+
+
+def test_retry_after_is_the_shortest_wait_that_allows_a_retry(clock):
+    rng = random.Random(0)
+    blocked_histories = 0
+    for _ in range(1000):
+        max_requests = rng.randint(1, 6)
+        window_size_in_ms = rng.choice([1000, 5000, 60_000])
+        history = [0]
+        for _ in range(rng.randint(1, 15)):
+            history.append(history[-1] + rng.choice([0, 1, rng.randint(0, 30_000)]))
+
+        last = send_requests(
+            clock, new_rate_limiter(), history, max_requests, window_size_in_ms
+        )[-1]
+        if last["allowed"]:
+            continue
+        blocked_histories += 1
+
+        for wait_in_seconds, expected in [
+            (last["retry_after_seconds"], True),
+            (last["retry_after_seconds"] - 1, False),
+        ]:
+            retry_at = history[-1] + wait_in_seconds * 1000
+            retry = send_requests(
+                clock,
+                new_rate_limiter(),
+                history + [retry_at],
+                max_requests,
+                window_size_in_ms,
+            )[-1]
+            assert retry["allowed"] is expected, (
+                history,
+                max_requests,
+                window_size_in_ms,
+            )
+    assert blocked_histories > 400
+
+
+@pytest.mark.parametrize(
+    "max_requests, window_size_in_ms",
+    [(0, 60_000), (-1, 60_000), (0, -1000)],
+)
+def test_retry_after_with_invalid_config(clock, max_requests, window_size_in_ms):
+    results = send_requests(
+        clock, new_rate_limiter(), [0, 1000], max_requests, window_size_in_ms
+    )
+    for result in results:
+        assert result["allowed"] is False
+        assert isinstance(result["retry_after_seconds"], int)
+        assert result["retry_after_seconds"] >= 1
diff --git a/end2end/fastapi_postgres_uvicorn_test.py b/end2end/fastapi_postgres_uvicorn_test.py
index 070cc9f37..ee1da1e3f 100644
--- a/end2end/fastapi_postgres_uvicorn_test.py
+++ b/end2end/fastapi_postgres_uvicorn_test.py
@@ -12,6 +12,7 @@
 post_url_nofw = "http://localhost:8113/create"
 sync_route_fw = "http://localhost:8112/sync_route"
 sync_route_nofw = "http://localhost:8113/sync_route"
+ratelimiting_url_fw = "http://localhost:8112/test_ratelimiting_1"
 
 
 def test_firewall_started_okay():
@@ -76,6 +77,23 @@ def test_sync_route_without_firewall():
     assert res.status_code == 200
 
 
+def test_ratelimiting_1_route():
+    res = requests.get(ratelimiting_url_fw)
+    assert res.status_code == 200
+    assert "Retry-After" not in res.headers
+    time.sleep(2)
+    res = requests.get(ratelimiting_url_fw)
+    assert res.status_code == 200
+
+    res = requests.get(ratelimiting_url_fw)
+    assert res.status_code == 429
+    assert res.text.startswith("You are rate limited by Zen.")
+
+    time.sleep(int(res.headers["Retry-After"]))
+    res = requests.get(ratelimiting_url_fw)
+    assert res.status_code == 200
+
+
 def test_routes_discovered_in_heartbeat():
     wait_for_heartbeats("http://localhost:5000", check_reported_routes)
 
diff --git a/end2end/flask_mysql_test.py b/end2end/flask_mysql_test.py
index f2a16c316..b1a657063 100644
--- a/end2end/flask_mysql_test.py
+++ b/end2end/flask_mysql_test.py
@@ -116,18 +116,21 @@ def test_ratelimiting_1_route():
     # First request :
     res = requests.get(base_url_fw + "/test_ratelimiting_1")
     assert res.status_code == 200
+    assert "Retry-After" not in res.headers
+    time.sleep(2)
     # Second request :
     res = requests.get(base_url_fw + "/test_ratelimiting_1")
     assert res.status_code == 200
     # Third request :
     res = requests.get(base_url_fw + "/test_ratelimiting_1")
     assert res.status_code == 429
+    assert int(res.headers["Retry-After"]) >= 1
     # Fourth request :
     res = requests.get(base_url_fw + "/test_ratelimiting_1")
     assert res.status_code == 429
 
-    time.sleep(5) # Wait until window expires
-    
+    time.sleep(int(res.headers["Retry-After"]))
+
     # Fifth request :
     res = requests.get(base_url_fw + "/test_ratelimiting_1")
     assert res.status_code == 200