Add rate limit response headers (imDarshanGK#151)

* Add rate limit response headers

* Fix workflow YAML syntax

* Fix remaining workflow YAML syntax

* Restore unrelated files from upstream main

Author: Divv1524

backend/app/main.py

@@ -3,7 +3,7 @@
 FastAPI application with advanced middleware, rate limiting, and full analysis engine.
 """
 
-from fastapi import FastAPI, Request, HTTPException
+from fastapi import FastAPI, Request
 from fastapi.middleware.cors import CORSMiddleware
 from fastapi.middleware.gzip import GZipMiddleware
 from fastapi.responses import JSONResponse
@@ -19,19 +19,28 @@
 
 # ── Rate limiter (in-memory, per IP) ──────────────────────────────────────────
 RATE_LIMIT = int(os.getenv("RATE_LIMIT_PER_MINUTE", "30"))
+RATE_LIMIT_WINDOW_SECONDS = 60
 _request_counts: dict[str, list[float]] = defaultdict(list)
 
 
-def check_rate_limit(ip: str) -> None:
+def check_rate_limit(ip: str) -> int:
+    """Record a request and return the remaining requests in the current window."""
     now = time.time()
-    window = 60.0
-    _request_counts[ip] = [t for t in _request_counts[ip] if now - t < window]
+    _request_counts[ip] = [
+        t for t in _request_counts[ip] if now - t < RATE_LIMIT_WINDOW_SECONDS
+    ]
     if len(_request_counts[ip]) >= RATE_LIMIT:
-        raise HTTPException(
-            status_code=429,
-            detail=f"Rate limit exceeded. Max {RATE_LIMIT} requests/minute.",
-        )
+        return -1
     _request_counts[ip].append(now)
+    return RATE_LIMIT - len(_request_counts[ip])
+
+
+def rate_limit_headers(remaining: int) -> dict[str, str]:
+    """Build rate limit headers for API responses."""
+    return {
+        "X-RateLimit-Limit": str(RATE_LIMIT),
+        "X-RateLimit-Remaining": str(max(remaining, 0)),
+    }
 
 
 # ── Lifespan ──────────────────────────────────────────────────────────────────
@@ -67,13 +76,28 @@ async def lifespan(app: FastAPI):
 async def add_process_time_header(request: Request, call_next):
     start = time.perf_counter()
     ip = request.client.host if request.client else "unknown"
+    remaining = RATE_LIMIT
 
     # Apply rate limiting to analysis endpoints only
-    if request.url.path in ("/explanation/", "/debugging/", "/suggestions/", "/analyze/") and ip != "testclient":
-        check_rate_limit(ip)
+    if request.url.path in ("/explanation/", "/debugging/", "/suggestions/", "/analyze/"):
+        remaining = check_rate_limit(ip)
+        if remaining < 0:
+            elapsed = (time.perf_counter() - start) * 1000
+            headers = rate_limit_headers(0)
+            headers["Retry-After"] = str(RATE_LIMIT_WINDOW_SECONDS)
+            headers["X-Process-Time-Ms"] = f"{elapsed:.2f}"
+            headers["X-QyverixAI-Version"] = "3.0.0"
+            return JSONResponse(
+                status_code=429,
+                content={
+                    "detail": f"Rate limit exceeded. Max {RATE_LIMIT} requests/minute."
+                },
+                headers=headers,
+            )
 
     response = await call_next(request)
     elapsed = (time.perf_counter() - start) * 1000
+    response.headers.update(rate_limit_headers(remaining))
     response.headers["X-Process-Time-Ms"] = f"{elapsed:.2f}"
     response.headers["X-QyverixAI-Version"] = "3.0.0"
     return response

backend/tests/test_endpoints.py

@@ -4,12 +4,18 @@
 """
 import pytest
 from fastapi.testclient import TestClient
-from app.main import _request_counts
 import sys, os
 sys.path.insert(0, os.path.dirname(os.path.dirname(__file__)))
-from app.main import app
+from app import main as app_main
 
-client = TestClient(app)
+client = TestClient(app_main.app)
+
+
+@pytest.fixture(autouse=True)
+def reset_rate_limit_state():
+    app_main._request_counts.clear()
+    yield
+    app_main._request_counts.clear()
 
 
 # ── Fixtures ──────────────────────────────────────────────────────────────────
@@ -152,6 +158,29 @@ def test_health():
     assert r.json()["status"] == "ok"
 
 
+def test_rate_limit_headers_on_success_response():
+    r = client.get("/")
+    assert r.status_code == 200
+    assert r.headers["X-RateLimit-Limit"] == str(app_main.RATE_LIMIT)
+    assert r.headers["X-RateLimit-Remaining"] == str(app_main.RATE_LIMIT)
+
+
+def test_rate_limit_returns_429_with_retry_after_header():
+    payload = {"code": "print('hello')", "language": "python"}
+
+    for expected_remaining in range(app_main.RATE_LIMIT - 1, -1, -1):
+        r = client.post("/debugging/", json=payload)
+        assert r.status_code == 200
+        assert r.headers["X-RateLimit-Limit"] == str(app_main.RATE_LIMIT)
+        assert r.headers["X-RateLimit-Remaining"] == str(expected_remaining)
+
+    r = client.post("/debugging/", json=payload)
+    assert r.status_code == 429
+    assert r.headers["Retry-After"] == str(app_main.RATE_LIMIT_WINDOW_SECONDS)
+    assert r.headers["X-RateLimit-Limit"] == str(app_main.RATE_LIMIT)
+    assert r.headers["X-RateLimit-Remaining"] == "0"
+
+
 # ── Explanation ───────────────────────────────────────────────────────────────
 def test_explanation_python():
     r = client.post("/explanation/", json={"code": PYTHON_CLEAN, "language": "python"})
@@ -194,7 +223,7 @@ def test_explanation_empty_code():
 def test_explanation_too_long():
     r = client.post("/explanation/", json={"code": "x" * 60000})
     assert r.status_code == 422
-    
+
 def test_explanation_typescript():
     r = client.post("/explanation/", json={"code": TS_CODE, "language": "typescript"})
     assert r.status_code == 200
@@ -422,6 +451,6 @@ def test_unicode_code():
     assert r.status_code == 200
 
 
-    _request_counts.clear()
+def test_single_line_code():
     r = client.post("/analyze/", json={"code": "print('hello')"})
-    assert r.status_code == 200
+    assert r.status_code == 200