scripts/ratelimit_test.py

"""Burst test that verifies the slowapi per-user rate limit on /api/answer.

Unlike scripts/stress_test.py (which THROTTLES with a semaphore to measure
graceful degradation), this script does the opposite: it fires every request at
the *same instant* -- a thundering herd -- to prove the limiter admits exactly N
and blocks the rest in a single fixed window.

    Expected for --burst 25 --limit 10:  10 x HTTP 200  +  15 x HTTP 429

------------------------------------------------------------------------------
HOW TO RUN
------------------------------------------------------------------------------
1. The server MUST run single-worker (uvicorn's default). slowapi's in-memory
   counter is per-process; with --workers N the burst is split across N
   independent buckets and far more than `limit` requests pass. So:

       uvicorn api:app --port 8000            # (no --workers flag)

2. Mode B -- dev-bypass (DEFAULT, for CI): leave CLERK_ISSUER unset on the
   server. Every request maps to the synthetic 'dev-user', so all 25 share one
   bucket. No token needed:

       python scripts/ratelimit_test.py

3. Mode A -- real Clerk JWT (manual end-to-end check): grab a fresh session
   token from the browser (it's short-lived, ~60s) and pass it:

       python scripts/ratelimit_test.py --token "eyJhbGci..."
"""

from __future__ import annotations

import argparse
import asyncio
import time
from collections import Counter

import aiohttp


async def _fire(
    session: aiohttp.ClientSession,
    url: str,
    payload: dict,
    timeout: float,
    start_gun: asyncio.Event,
    idx: int,
) -> dict:
    """One request that waits at the barrier, then fires the instant it opens.

    The `await start_gun.wait()` is the concurrency barrier: every worker parks
    here until the main coroutine fires the starting gun, so all `burst` requests
    leave the line on the SAME event-loop tick instead of trickling out as each
    task gets scheduled. That true simultaneity is what makes this a real burst
    test of the fixed window (a semaphore would serialize them and hide the seam).
    """
    await start_gun.wait()
    t0 = time.perf_counter()
    try:
        # NOTE: aiohttp does NOT raise on 4xx/5xx unless raise_for_status=True,
        # which we deliberately omit -- a 429 is the EXPECTED result here, not an
        # error. We just read the status code off the response.
        async with session.post(url, json=payload, timeout=timeout) as resp:
            await resp.read()  # drain the body so the connection is freed
            return {"idx": idx, "status": resp.status,
                    "latency": time.perf_counter() - t0}
    except asyncio.TimeoutError:
        # graceful: a slow/blocked request becomes a recorded outcome, not a crash
        return {"idx": idx, "status": "timeout",
                "latency": time.perf_counter() - t0}
    except aiohttp.ClientError as exc:
        return {"idx": idx, "status": f"error:{type(exc).__name__}",
                "latency": time.perf_counter() - t0}


async def run(url: str, token: str | None, burst: int, limit: int,
              timeout: float) -> bool:
    """Fire `burst` simultaneous requests; return True if the limit held exactly."""
    # --- identity: one token (or dev-bypass) => one rate-limit bucket ---
    headers: dict[str, str] = {}
    if token:
        headers["Authorization"] = f"Bearer {token}"
        mode = "Mode A -- real Clerk JWT (--token)"
    else:
        # No header -> server runs every request as the synthetic 'dev-user',
        # so the whole burst still shares ONE bucket. Perfect for CI.
        mode = "Mode B -- dev-bypass (synthetic 'dev-user', no token)"

    # Empty context deliberately: /api/answer returns 200 instantly for empty
    # context WITHOUT calling the LLM. The rate-limit decorator runs BEFORE the
    # endpoint body, so the 429s still fire -- we prove the limiter for zero LLM
    # cost. (Set a real context only if you want to exercise the answer path.)
    payload = {"query": "rate-limit burst probe", "context": ""}

    expected_pass = limit
    expected_block = burst - limit

    print(f"Bursting {burst} simultaneous requests at {url}")
    print(f"  {mode}")
    print(f"  expecting {expected_pass} x 200 and {expected_block} x 429\n")
    print("  [!] Server MUST be single-worker (uvicorn default, NO --workers N)")
    print("      - slowapi's in-memory bucket is per-process; multiple workers")
    print("      mean multiple buckets and more than `limit` requests pass.\n")

    start_gun = asyncio.Event()

    async with aiohttp.ClientSession(headers=headers) as session:
        # Create all tasks first. Each immediately parks on `start_gun.wait()`.
        tasks = [
            asyncio.create_task(_fire(session, url, payload, timeout, start_gun, i))
            for i in range(burst)
        ]
        # Yield control so EVERY task gets scheduled and reaches the barrier
        # before we fire. Without this beat, set() could race ahead of tasks
        # that haven't hit `await start_gun.wait()` yet.
        await asyncio.sleep(0.1)

        wall_start = time.perf_counter()
        start_gun.set()                      # FIRE: all `burst` requests go now
        results = await asyncio.gather(*tasks)
        wall = time.perf_counter() - wall_start

    # --- tally outcomes ---
    statuses = Counter(r["status"] for r in results)
    n_200 = statuses.get(200, 0)
    n_429 = statuses.get(429, 0)
    lat = sorted(r["latency"] for r in results)

    print("=" * 60)
    print(f"{'Status / outcome':<28}{'count':>10}")
    print("-" * 60)
    for status, n in sorted(statuses.items(), key=lambda kv: str(kv[0])):
        print(f"{str(status):<28}{n:>10}")
    print("-" * 60)
    print(f"wall time            {wall:>8.3f}s")
    print(f"200 (admitted)       {n_200:>8} / expected {expected_pass}")
    print(f"429 (blocked)        {n_429:>8} / expected {expected_block}")
    if lat:
        print(f"latency min/max      {lat[0]:.3f} / {lat[-1]:.3f} s")
    print("=" * 60)

    # --- verdict ---
    passed = (n_200 == expected_pass and n_429 == expected_block)
    if passed:
        print(f"\n[PASS] Limiter held exactly: {expected_pass}x200, "
              f"{expected_block}x429. The 429 path is working.")
        return True

    print("\n[FAIL] Counts did not match. Likely causes, in order:")
    # Diagnostics ranked by how often each is the real culprit.
    if n_200 > expected_pass:
        print(f"  - {n_200} passed (> {expected_pass}). Either the server is "
              f"running MULTIPLE WORKERS (each with its own in-memory bucket),")
        print(f"    or -- far less likely at this speed -- the burst straddled a "
              f"60s fixed-window boundary. Re-run; if it persists it's workers.")
    if any(str(s).startswith(("error", "timeout")) for s in statuses):
        print("  - Hard failures (timeout/conn-error) present -- is the server up "
              "at this URL, and is the timeout high enough?")
    if n_200 == 0 and n_429 == 0:
        print("  - Nothing got through at all -- check the URL and that the API "
              "is actually listening.")
    if n_429 == 0 and n_200 < burst:
        print("  - No 429s seen -- is slowapi actually wired (app.state.limiter + "
              "@limiter.limit on /api/answer)?")
    return False


def main(argv: list[str] | None = None) -> int:
    p = argparse.ArgumentParser(description="Verify the per-user rate limit on /api/answer.")
    p.add_argument("--url", default="http://127.0.0.1:8000/api/answer")
    p.add_argument("--token", default=None,
                   help="Clerk JWT for Mode A; omit for Mode B dev-bypass.")
    p.add_argument("--burst", type=int, default=25,
                   help="total simultaneous requests (default 25)")
    p.add_argument("--limit", type=int, default=10,
                   help="the server's per-minute cap, i.e. expected 200s (default 10)")
    p.add_argument("--timeout", type=float, default=30.0)
    args = p.parse_args(argv)

    ok = asyncio.run(run(args.url, args.token, args.burst, args.limit, args.timeout))
    # exit code 0 on PASS, 1 on FAIL -- so CI (GitHub Actions) fails the job loudly.
    return 0 if ok else 1


if __name__ == "__main__":
    raise SystemExit(main())