AI Rate Limit Alerts

cat <<'EOF' > kong-recipe.yaml
_defaults:
  kongctl:
    namespace: ai-rate-limit-alerts-recipe
ai_gateway_model_providers:
- ref: ai-rate-limit-alerts-provider
  ai_gateway: !lookup {id: !env 'AI_GATEWAY_ID'}
  name: ai-rate-limit-alerts-provider
  display_name: ai-rate-limit-alerts
  type: openai
  config:
    auth:
      type: basic
      headers:
      - name: Authorization
        value: !secret {source: !env 'DECK_OPENAI_TOKEN'}
ai_gateway_auth_strategies:
- ref: ai-rate-limit-alerts-auth
  ai_gateway: !lookup {id: !env 'AI_GATEWAY_ID'}
  name: ai-rate-limit-alerts-auth
  display_name: ai-rate-limit-alerts key auth
  type: key-auth
  config:
    key_names:
    - apikey
    hide_credentials: true
ai_gateway_consumers:
- ref: ai-rate-limit-alerts-consumer
  ai_gateway: !lookup {id: !env 'AI_GATEWAY_ID'}
  name: ai-rate-limit-alerts-consumer
  display_name: ai-rate-limit-alerts demo consumer
  type: api-key
  credentials:
  - ref: ai-rate-limit-alerts-credential
    name: ai-rate-limit-alerts-credential
    display_name: ai-rate-limit-alerts demo API key
    type: api-key
    api_key: !secret {source: !env 'DECK_CONSUMER_API_KEY'}
ai_gateway_policies:
- ref: ai-rate-limit-alerts-rla
  ai_gateway: !lookup {id: !env 'AI_GATEWAY_ID'}
  name: ai-rate-limit-alerts-rla
  display_name: ai-rate-limit-alerts token window
  type: ai-rate-limiting-advanced
  config:
    strategy: local
    identifier: consumer
    tokens_count_strategy: total_tokens
    hide_client_headers: false
    policies:
    - window_type: fixed
      limits:
      - limit: 50
        window_size: 60
- ref: ai-rate-limit-alerts-notify
  ai_gateway: !lookup {id: !env 'AI_GATEWAY_ID'}
  name: ai-rate-limit-alerts-notify
  display_name: ai-rate-limit-alerts 80% webhook
  type: post-function
  config:
    body_filter:
    - "if not ngx.arg[2] then\n  return\nend\nlocal rem, lim\nlocal headers = kong.response.get_headers()\
      \ or {}\nfor name, value in pairs(headers) do\n  local lower = string.lower(name)\n\
      \  if lower:find(\"x-ai-ratelimit-remaining-\", 1, true)\n     and not lower:find(\"\
      reset\", 1, true)\n     and not lower:find(\"retry\", 1, true) then\n    rem\
      \ = tonumber(value)\n  elseif lower:find(\"x-ai-ratelimit-limit-\", 1, true)\n\
      \     and not lower:find(\"reset\", 1, true)\n     and not lower:find(\"retry\"\
      , 1, true) then\n    lim = tonumber(value)\n  end\nend\nif not rem or not lim\
      \ or lim <= 0 then\n  return\nend\nlocal threshold = 80\nif (lim - rem) / lim\
      \ * 100 < threshold then\n  return\nend\nlocal consumer = kong.client.get_consumer()\
      \ or {}\nlocal id = consumer.id or kong.client.get_forwarded_ip() or \"anon\"\
      \nlocal dict = ngx.shared.kong\nif dict then\n  local added = dict:add(\"ai-rate-alert:\"\
      \ .. id, 1, 60)\n  if not added then\n    return\n  end\nend\nlocal url = os.getenv(\"\
      DECK_WEBHOOK_URL\")\nif not url or url == \"\" then\n  return\nend\nlocal username\
      \ = consumer.username or \"unknown\"\nlocal ok, err = ngx.timer.at(0, function(premature)\n\
      \  if premature then\n    return\n  end\n  local http = require(\"resty.http\"\
      )\n  local cjson = require(\"cjson.safe\")\n  local httpc = http.new()\n  httpc:set_timeout(3000)\n\
      \  httpc:request_uri(url, {\n    method = \"POST\",\n    body = cjson.encode({\n\
      \      text = string.format(\n        \"%s reached %s%% of the token quota.\
      \ %s tokens remaining.\",\n        username,\n        tostring(threshold),\n\
      \        tostring(rem)\n      ),\n      username = username,\n      remaining\
      \ = rem,\n      limit = lim,\n      rate = lim - rem,\n      threshold = threshold,\n\
      \      window = \"minute\",\n    }),\n    headers = { [\"Content-Type\"] = \"\
      application/json\" },\n    ssl_verify = true,\n  })\nend)\nif not ok then\n\
      \  kong.log.err(\"ai-rate-alert timer: \", err)\nend\n"
ai_gateway_models:
- ref: ai-rate-limit-alerts-model
  ai_gateway: !lookup {id: !env 'AI_GATEWAY_ID'}
  name: ai-rate-limit-alerts-model
  display_name: ai-rate-limit-alerts
  type: model
  formats:
  - type: openai
  capabilities:
  - generate
  access:
    auth_strategies:
    - ai-rate-limit-alerts-auth
  policies:
  - ai-rate-limit-alerts-rla
  - ai-rate-limit-alerts-notify
  config:
    max_request_body_size: 10485760
    response_streaming: allow
    logging:
      payloads: true
    route:
      paths:
      - /ai-rate-limit-alerts
      protocols:
      - http
      - https
      methods:
      - POST
      - OPTIONS
      strip_path: true
  targets:
  - name: !env 'DECK_CHAT_MODEL'
    provider: ai-rate-limit-alerts-provider
    config:
      type: openai
EOF
kongctl apply -f kong-recipe.yaml -o text --auto-approve --pat "${KONNECT_TOKEN}"

rm -f kong-recipe.yaml

cat <<'EOF' > demo.py
"""AI rate limit early-alert demo.

Sends short chat completions through Kong until the Consumer uses 80% of a
50-token / 60-second window. Kong should POST a JSON alert to DECK_WEBHOOK_URL
at that threshold, then return HTTP 429 after the quota is exhausted.

Expected output:
  - Early requests -> 200, remaining headers well above 20%
  - At 80% used    -> 200, webhook inbox grows
  - After the quota -> 429

Run:
  export PROXY_URL=http://localhost:8000
  export DECK_WEBHOOK_URL=https://webhook.site/YOUR-UUID
  python demo.py
"""

from __future__ import annotations

import json
import os
import sys
import time
import urllib.error
import urllib.parse
import urllib.request

from openai import APIStatusError, OpenAI

PROXY_URL = os.getenv("PROXY_URL", "http://localhost:8000")
API_KEY = "demo-api-key"
WEBHOOK_URL = os.getenv("DECK_WEBHOOK_URL", "")
LIMIT = 50
THRESHOLD_PCT = 80

_USE_COLOR = sys.stdout.isatty() and "NO_COLOR" not in os.environ


def _c(code: str, s: str) -> str:
    return f"\033[{code}m{s}\033[0m" if _USE_COLOR else s


def BOLD(s: str) -> str:
    return _c("1", s)


def DIM(s: str) -> str:
    return _c("2", s)


def GREEN(s: str) -> str:
    return _c("32", s)


def RED(s: str) -> str:
    return _c("31", s)


def YELLOW(s: str) -> str:
    return _c("33", s)


def BLUE(s: str) -> str:
    return _c("34", s)


def remaining_from(headers: dict[str, str]) -> tuple[int | None, int | None]:
    rem = lim = None
    for raw_name, value in headers.items():
        name = raw_name.lower()
        if name.startswith("x-ai-ratelimit-remaining-") and "reset" not in name and "retry" not in name:
            rem = int(float(value))
        elif name.startswith("x-ai-ratelimit-limit-") and "reset" not in name and "retry" not in name:
            lim = int(float(value))
    return rem, lim


def webhook_token(url: str) -> str:
    path = urllib.parse.urlparse(url).path.strip("/")
    return path.split("/")[0] if path else ""


def webhook_count(token: str) -> int:
    if not token:
        return 0
    req = urllib.request.Request(
        f"https://webhook.site/token/{token}/requests?sorting=newest&per_page=20",
        headers={"Accept": "application/json"},
    )
    try:
        with urllib.request.urlopen(req, timeout=10) as resp:
            body = json.load(resp)
    except (urllib.error.URLError, json.JSONDecodeError, TimeoutError):
        return 0
    return len(body.get("data") or [])


def make_client() -> OpenAI:
    return OpenAI(
        base_url=f"{PROXY_URL}/ai-rate-limit-alerts",
        api_key="unused",
        default_headers={"apikey": API_KEY},
    )


def main() -> int:
    if not WEBHOOK_URL:
        print(RED("Set DECK_WEBHOOK_URL to your webhook.site unique URL."))
        return 1

    client = make_client()
    token = webhook_token(WEBHOOK_URL)
    before = webhook_count(token)
    print(BOLD("[SETUP]"))
    print(f"  {DIM(f'proxy={PROXY_URL}/ai-rate-limit-alerts limit={LIMIT} threshold={THRESHOLD_PCT}%')}")
    print(f"  {DIM(f'webhook inbox before={before}')}")

    got_429 = False
    for i in range(1, 20):
        print(f"\n{BOLD('[REQUEST]')} chat {i}")
        try:
            raw = client.chat.completions.with_raw_response.create(
                model="gpt-4o-mini",
                messages=[{"role": "user", "content": "hi"}],
                max_tokens=1,
            )
            headers = {k.lower(): v for k, v in raw.headers.items()}
            rem, lim = remaining_from(headers)
            used = (lim - rem) if rem is not None and lim else None
            pct = int(used / lim * 100) if used is not None and lim else None
            rem_s = f"{rem}/{lim}" if rem is not None and lim is not None else "n/a"
            used_s = f"{pct}%" if pct is not None else "n/a"
            remaining_s = f"remaining={rem_s}"
            print(f"  {GREEN('200')} {BLUE(remaining_s)} {DIM('used=' + used_s)}")
            if pct is not None and pct >= THRESHOLD_PCT:
                print(f"  {YELLOW(f'THRESHOLD {THRESHOLD_PCT}% reached')}")
        except APIStatusError as err:
            print(f"  {RED(str(err.status_code))} {DIM(str(err.message)[:120])}")
            if err.status_code == 429:
                got_429 = True
                break

    hit = False
    after = before
    for _ in range(15):
        after = webhook_count(token)
        if after > before:
            hit = True
            break
        time.sleep(1)

    if hit:
        print(f"\n{GREEN('[WEBHOOK]')} inbox grew {before} -> {after}")
    else:
        print(f"\n{RED('[WEBHOOK]')} no new POST at {WEBHOOK_URL}")
        return 1

    if not got_429:
        print(RED("[FAIL] expected HTTP 429 after the token quota"))
        return 1

    print(f"\n{GREEN('[OK]')} early alert fired before the 429")
    return 0


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