Add a per-user rate limit and token limits to the gateway

short · 45 min · Objective 2.2

Task

Add the two limits that answer Unbounded Consumption to your gateway: a rate limit per user, and token limits on what one request may send and receive. Then prove each fires, that one user's limit does not touch another's, and that a refused request never reached the model.

Steps

  1. At the top of gateway/gateway.py, add four constants: RATE_LIMIT = 5, RATE_WINDOW_S = 600, MAX_INPUT_TOKENS = 2000 and MAX_OUTPUT_TOKENS = 256. A ten-minute window is deliberate: on a CPU, five slow answers can take longer than a minute, and a one-minute window would slide open while you test it.
  2. Add the rate limit in check(), lesson 2's pre-model hook: keep a list of request times per rec["user"], discard those older than the window, and refuse with HTTP 429 and decision block:rate when the user already has RATE_LIMIT requests in it.
  3. Add the input token limit as a check made BEFORE forwarding, in check_prompt(), which runs once rec["prompt"] holds the full assembled prompt: estimate its tokens conservatively (characters divided by three is a fair over-estimate for English) and refuse with HTTP 413 and decision block:input-tokens when it exceeds MAX_INPUT_TOKENS. Log response_tokens as 0 for every refusal.
  4. Add the output token limit by passing MAX_OUTPUT_TOKENS to the runner's maximum-output option (in Ollama's API this is num_predict, inside the request's options object; check your runner's documentation for its name), so runaway generation stops at the model rather than after it.
  5. Restart the gateway and write two lines in lab/gateway-limits.md: what each limit protects against, using the lesson's terms (availability, cost, extraction), and why the user field these limits key on is not yet an authenticated identity.

Verify

Run from ~/secai-lab. The first command marks where this test starts in the log.

wc -l < logs/gateway.jsonl > lab/gw-short-start.txt
python - <<'EOF'
import json, urllib.request, urllib.error
GW = "http://127.0.0.1:8088/ask"
def ask(user, question):
    body = json.dumps({"user": user, "question": question}).encode()
    req = urllib.request.Request(GW, data=body, headers={"Content-Type": "application/json"})
    try:
        with urllib.request.urlopen(req, timeout=180) as r:
            return r.status
    except urllib.error.HTTPError as e:
        return e.code
codes = [ask("alice", "How many days of annual leave do staff get?") for _ in range(12)]
print("alice: allowed", codes.count(200), "rate-limited", codes.count(429))
print("bob, first request:", ask("bob", "Who is on call this week?"))
print("oversized request:", ask("carol", "word " * 20000))
EOF
python - <<'EOF'
import json, collections
start = int(open("lab/gw-short-start.txt").read())
recs = [json.loads(l) for l in open("logs/gateway.jsonl").readlines()[start:]]
print(collections.Counter(r["decision"] for r in recs))
print("refused requests that reached the model:", sum(1 for r in recs if r["decision"].startswith("block:") and r.get("response_tokens")))
print("largest response_tokens:", max((r.get("response_tokens") or 0 for r in recs), default=0))
EOF
grep -cE '^(RATE_LIMIT|RATE_WINDOW_S|MAX_INPUT_TOKENS|MAX_OUTPUT_TOKENS) *=' gateway/gateway.py

alice: allowed 5 rate-limited 7, bob's first request 200, the oversized request 413; in the log six allow, seven block:rate and one block:input-tokens; zero refused requests that reached the model; the largest response_tokens no more than 256; and 4 constants. If bob is refused too, the limit is global rather than per user.

Notes

This gateway trusts the user field the caller sends, so a script that wants more than five requests simply sends a different name each time. That is the lesson's warning about keying limits to anything other than an authenticated identity, and the applied lab closes it with per-user keys. A request-rate limit on its own would not have capped the bill either: the oversized request was one request, and only the token limit stopped it.

This is an independent study companion for CompTIA SecAI+ CY0-001 and is not produced by or endorsed by CompTIA.