feat: add service API-key auth and rate limiting to protect Gemini quota - #73
Conversation
- Require X-API-Key header on /chat, /chat/stream, and /chat/{chat_id}
- Validate key against SERVICE_API_KEY env var using secrets.compare_digest
- Return 401 when key is missing or invalid
- Add per-client rate limiting (20 req/min) via slowapi on chat endpoints
- Rate limit key derived from X-API-Key, falling back to X-Forwarded-For IP
- Return 429 with Retry-After header when limit exceeded
- Fail loudly at startup if SERVICE_API_KEY is unset in production
- Allow AUTH_DISABLED=true for local development without a key
- Leave /ping, /cache/stats, and /confidence/policy unauthenticated
- Update keep-awake workflow to use /ping endpoint
- Add SERVICE_API_KEY to render.yaml envVars
- Update .env.example and README with new env vars and auth docs
|
Caution Review failedThe pull request is closed. ℹ️ Recent review info⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Pro Plus Run ID: 📒 Files selected for processing (5)
WalkthroughThe service adds API-key authentication and rate-limiting infrastructure, documents related configuration, provisions ChangesAPI access controls
Estimated code review effort: 3 (Moderate) | ~25 minutes Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@main.py`:
- Around line 72-79: Update the startup validation around AUTH_DISABLED so it
raises RuntimeError whenever AUTH_DISABLED is enabled outside development,
local, or test environments, regardless of SERVICE_API_KEY. Preserve the
existing missing-SERVICE_API_KEY validation for non-disabled authentication in
production.
- Around line 308-309: Update the async handlers chat() and chat_stream() so
Gemini operations never block the event loop: use send_message_async() for the
non-safety chat path, and asynchronously create and consume the Gemini stream
with async for in chat_stream(), or offload any unavoidable synchronous calls to
a worker thread. Apply the changes at main.py lines 308-309 and 543-544.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro Plus
Run ID: b46148c1-2215-49c1-8c9e-23b44582ba41
📒 Files selected for processing (6)
.env.example.github/workflows/keep-render-awake.ymlREADME.mdmain.pyrender.yamlrequirements.txt
| AUTH_DISABLED = os.getenv("AUTH_DISABLED", "false").lower() in {"1", "true", "yes"} | ||
|
|
||
| if not AUTH_DISABLED and not SERVICE_API_KEY: | ||
| if os.getenv("ENVIRONMENT", "production").lower() not in {"development", "local", "test"}: | ||
| raise RuntimeError( | ||
| "SERVICE_API_KEY must be set in production. " | ||
| "Set AUTH_DISABLED=true to run without authentication locally." | ||
| ) |
There was a problem hiding this comment.
🔒 Security & Privacy | 🟠 Major | ⚡ Quick win
Reject AUTH_DISABLED in production.
AUTH_DISABLED=true skips this startup check and makes verify_api_key return success for every request. Fail startup when that flag is enabled outside development, local, or test; otherwise a production configuration mistake fully disables the new access control.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@main.py` around lines 72 - 79, Update the startup validation around
AUTH_DISABLED so it raises RuntimeError whenever AUTH_DISABLED is enabled
outside development, local, or test environments, regardless of SERVICE_API_KEY.
Preserve the existing missing-SERVICE_API_KEY validation for non-disabled
authentication in production.
| @limiter.limit("20/minute") | ||
| async def chat(request: ChatRequest, http_request: Request, fastapi_response: Response, _key: str = Security(verify_api_key)): |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | 🏗️ Heavy lift
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
tmp_dir="$(mktemp -d)"
python -m pip download --no-deps --dest "$tmp_dir" google-generativeai==0.8.3
unzip -q "$tmp_dir"/google_generativeai-0.8.3-*.whl -d "$tmp_dir/pkg"
rg -n -C 2 'send_message_async|send_message\(' "$tmp_dir/pkg"Repository: Deen-Bridge/dnb-ai
Length of output: 4515
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Inspect the relevant endpoints and surrounding helper code.
python3 - <<'PY'
from pathlib import Path
p = Path("main.py")
lines = p.read_text().splitlines()
for start, end in [(260, 360), (500, 610)]:
print(f"\n--- main.py:{start}-{end} ---")
for i in range(start, min(end, len(lines)) + 1):
print(f"{i:4d}: {lines[i-1]}")
PYRepository: Deen-Bridge/dnb-ai
Length of output: 10392
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Inspect the 0.8.3 SDK for async chat/stream entrypoints and whether stream iteration is async.
tmp_dir="$(mktemp -d)"
python3 -m pip download --no-deps --dest "$tmp_dir" google-generativeai==0.8.3 >/dev/null
unzip -q "$tmp_dir"/google_generativeai-0.8.3-*.whl -d "$tmp_dir/pkg"
python3 - <<'PY' "$tmp_dir/pkg"
import sys
from pathlib import Path
root = Path(sys.argv[1])
files = [
root / "google/generativeai/generative_models.py",
root / "google/generativeai/client.py",
]
for f in files:
print(f"\n--- {f.relative_to(root)} ---")
text = f.read_text().splitlines()
for needle in ["def send_message(", "async def send_message_async(", "stream=True", "def generate_content(", "async def generate_content_async(", "__aiter__", "async for"]:
hits = [i + 1 for i, line in enumerate(text) if needle in line]
if hits:
print(f"{needle}: {hits[:10]}")
PYRepository: Deen-Bridge/dnb-ai
Length of output: 421
🏁 Script executed:
#!/bin/bash
set -euo pipefail
tmp_dir="$(mktemp -d)"
python3 -m pip download --no-deps --dest "$tmp_dir" google-generativeai==0.8.3 >/dev/null
unzip -q "$tmp_dir"/google_generativeai-0.8.3-*.whl -d "$tmp_dir/pkg"
python3 - <<'PY' "$tmp_dir/pkg"
from pathlib import Path
import sys
root = Path(sys.argv[1])
text = (root / "google/generativeai/generative_models.py").read_text().splitlines()
# Print the async/chat/stream related method signatures and nearby context.
needles = [
"async def send_message_async(",
"def send_message(",
"stream=True",
"async def generate_content_async(",
"def generate_content(",
"class ChatSession",
]
for needle in needles:
print(f"\n### {needle}")
for i, line in enumerate(text):
if needle in line:
start = max(0, i - 8)
end = min(len(text), i + 35)
for j in range(start, end):
print(f"{j+1:4d}: {text[j]}")
break
PYRepository: Deen-Bridge/dnb-ai
Length of output: 12537
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Inspect the exact provider-call sites in main.py.
python3 - <<'PY'
from pathlib import Path
p = Path("main.py")
lines = p.read_text().splitlines()
for start, end in [(360, 520), (590, 660)]:
print(f"\n--- main.py:{start}-{end} ---")
for i in range(start, min(end, len(lines)) + 1):
print(f"{i:4d}: {lines[i-1]}")
PYRepository: Deen-Bridge/dnb-ai
Length of output: 11880
🏁 Script executed:
#!/bin/bash
set -euo pipefail
tmp_dir="$(mktemp -d)"
python3 -m pip download --no-deps --dest "$tmp_dir" google-generativeai==0.8.3 >/dev/null
unzip -q "$tmp_dir"/google_generativeai-0.8.3-*.whl -d "$tmp_dir/pkg"
python3 - <<'PY' "$tmp_dir/pkg"
from pathlib import Path
import sys
root = Path(sys.argv[1])
files = [
root / "google/generativeai/types/generation_types.py",
root / "google/generativeai/generative_models.py",
]
for f in files:
print(f"\n--- {f.relative_to(root)} ---")
text = f.read_text().splitlines()
for needle in [
"class AsyncGenerateContentResponse",
"__aiter__",
"__anext__",
"async for",
"class GenerateContentResponse",
"yield",
]:
hits = [i + 1 for i, line in enumerate(text) if needle in line]
if hits:
print(f"{needle}: {hits[:20]}")
PYRepository: Deen-Bridge/dnb-ai
Length of output: 455
🏁 Script executed:
#!/bin/bash
set -euo pipefail
tmp_dir="$(mktemp -d)"
python3 -m pip download --no-deps --dest "$tmp_dir" google-generativeai==0.8.3 >/dev/null
unzip -q "$tmp_dir"/google_generativeai-0.8.3-*.whl -d "$tmp_dir/pkg"
python3 - <<'PY' "$tmp_dir/pkg"
from pathlib import Path
import sys
root = Path(sys.argv[1])
path = root / "google/generativeai/types/generation_types.py"
text = path.read_text().splitlines()
needles = ["class AsyncGenerateContentResponse", "class GenerateContentResponse", "class BaseGenerateContentResponse"]
for needle in needles:
for i, line in enumerate(text):
if needle in line:
start = max(0, i - 12)
end = min(len(text), i + 120)
print(f"\n### {needle}")
for j in range(start, end):
print(f"{j+1:4d}: {text[j]}")
break
PYRepository: Deen-Bridge/dnb-ai
Length of output: 13667
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Locate the SDK pin and any internal helper wrappers that may already bridge sync calls.
rg -n --hidden --glob '!**/.git/**' 'google-generativeai|generativeai|genai|send_message_async|generate_content_async|asyncio\.to_thread|run_in_executor' .Repository: Deen-Bridge/dnb-ai
Length of output: 1495
Move Gemini calls off the event loop. chat() still hits the sync send_message() path when the safety pipeline is disabled, and chat_stream() synchronously creates and iterates the Gemini stream inside an async handler. Use send_message_async()/async for, or offload the sync path to a worker thread.
📍 Affects 1 file
main.py#L308-L309(this comment)main.py#L543-L544
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@main.py` around lines 308 - 309, Update the async handlers chat() and
chat_stream() so Gemini operations never block the event loop: use
send_message_async() for the non-safety chat path, and asynchronously create and
consume the Gemini stream with async for in chat_stream(), or offload any
unavoidable synchronous calls to a worker thread. Apply the changes at main.py
lines 308-309 and 543-544.
Source: Path instructions
There was a problem hiding this comment.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
main.py (1)
440-461: 🩺 Stability & Availability | 🔴 Critical | ⚡ Quick winRename the FastAPI
Requestparameter on both chat endpoints
slowapineeds the Starlette request in a parameter literally namedrequest; here both handlers userequest: ChatRequestand rename the realRequesttohttp_request, so the limiter/key function sees the body model instead of aRequestand rate limiting breaks on both/chatand/chat/stream. Rename the body model variable (for examplechat_request) and userequest: Requestfor the FastAPI request, then update the internalrequest.*/http_request.headersreferences accordingly.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@main.py` around lines 440 - 461, The /chat and /chat/stream handlers use the body model as request, preventing slowapi from receiving the Starlette Request. In both main.py sites (lines 440-461 and 758-761), rename the ChatRequest parameter (for example, chat_request), use request: Request for the HTTP request, and update all body-model references and header accesses accordingly.
🧹 Nitpick comments (1)
main.py (1)
707-746: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winChain the original exception with
raise ... from exc.Ruff/flake8-bugbear (B904) flags all five re-raises here. Since CI enforces flake8 for this service, worth confirming whether bugbear's opinionated
B9codes are selected in your config — if so this currently fails the build; either way, chaining preserves the original traceback for debugging instead of presenting theHTTPExceptionas if it arose independently.🔧 Suggested fix (repeat for all 5 blocks)
except ResourceExhausted as exc: logger.warning("Gemini rate limit exceeded for chat %s: %s", chat_id, exc) raise HTTPException( status_code=429, detail="Rate limit exceeded. Please try again later.", headers={"X-Trace-Id": trace.trace_id}, - ) + ) from exc🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@main.py` around lines 707 - 746, Update the five HTTPException re-raises in the ResourceExhausted, InvalidArgument, timeout, ServiceUnavailable, and generic Exception handlers to use explicit exception chaining with the corresponding caught exception variable. Preserve each existing status code, detail, headers, and logging behavior.Source: Linters/SAST tools
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Outside diff comments:
In `@main.py`:
- Around line 440-461: The /chat and /chat/stream handlers use the body model as
request, preventing slowapi from receiving the Starlette Request. In both
main.py sites (lines 440-461 and 758-761), rename the ChatRequest parameter (for
example, chat_request), use request: Request for the HTTP request, and update
all body-model references and header accesses accordingly.
---
Nitpick comments:
In `@main.py`:
- Around line 707-746: Update the five HTTPException re-raises in the
ResourceExhausted, InvalidArgument, timeout, ServiceUnavailable, and generic
Exception handlers to use explicit exception chaining with the corresponding
caught exception variable. Preserve each existing status code, detail, headers,
and logging behavior.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro Plus
Run ID: bb5265b6-f663-4ec0-a748-24b43a8266a8
📒 Files selected for processing (2)
main.pyrequirements.txt
There was a problem hiding this comment.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (3)
README.md (3)
44-45: 🎯 Functional Correctness | 🟡 Minor | ⚡ Quick winList
/chat/streamin the API table.The new authentication note covers the streaming route, but the API table omits
POST /chat/stream. Add it so the documented endpoint and authentication surface matchesmain.py.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@README.md` around lines 44 - 45, Add POST /chat/stream to the API table in README.md, alongside the existing chat endpoints, and document it consistently with the route defined in main.py and the required X-API-Key authentication.
127-128: 🔒 Security & Privacy | 🟠 Major | ⚡ Quick winEnforce or explicitly warn about the production impact of
AUTH_DISABLED.
main.pybypasses authentication wheneverAUTH_DISABLED=true, including in production; the flag is not actually restricted to local development. Either reject this setting whenENVIRONMENT=production, or document prominently that enabling it disables production authentication.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@README.md` around lines 127 - 128, Update the AUTH_DISABLED documentation in README.md to prominently state that enabling it disables API-key authentication even in production, or explicitly document the required production restriction if main.py enforces one. Ensure the documented behavior matches the actual production handling and clearly retain its local-development-only intent.
100-101: 🎯 Functional Correctness | 🟡 Minor | ⚡ Quick winAdd the required service authentication configuration to runnable deployment examples.
README.md#L100-L101: includeSERVICE_API_KEYin the inline Docker command, or mark it explicitly as a localAUTH_DISABLED=trueexample.README.md#L109-L117: addSERVICE_API_KEYto the Render example withsync: false.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@README.md` around lines 100 - 101, Update README.md lines 100-101 to include SERVICE_API_KEY in the inline Docker deployment command, or explicitly mark it as a local AUTH_DISABLED=true example. Update README.md lines 109-117 to add SERVICE_API_KEY to the Render environment configuration with sync: false.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Outside diff comments:
In `@README.md`:
- Around line 44-45: Add POST /chat/stream to the API table in README.md,
alongside the existing chat endpoints, and document it consistently with the
route defined in main.py and the required X-API-Key authentication.
- Around line 127-128: Update the AUTH_DISABLED documentation in README.md to
prominently state that enabling it disables API-key authentication even in
production, or explicitly document the required production restriction if
main.py enforces one. Ensure the documented behavior matches the actual
production handling and clearly retain its local-development-only intent.
- Around line 100-101: Update README.md lines 100-101 to include SERVICE_API_KEY
in the inline Docker deployment command, or explicitly mark it as a local
AUTH_DISABLED=true example. Update README.md lines 109-117 to add
SERVICE_API_KEY to the Render environment configuration with sync: false.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro Plus
Run ID: 8e0b3ad7-2a7b-4023-95c0-fc420dc20631
📒 Files selected for processing (3)
README.mdmain.pyrender.yaml
🚧 Files skipped from review as they are similar to previous changes (1)
- main.py
2 participants
Closes #9
Summary
Adds shared-secret API-key authentication and per-client rate limiting to protect the Gemini quota and secure all chat endpoints against unauthorised access.
Changes
Authentication (
main.py)X-API-Keyheader required onPOST /chat,POST /chat/stream, andDELETE /chat/{chat_id}SERVICE_API_KEYenv var usingsecrets.compare_digest(constant-time comparison)Securitydependency so the requirement shows up in/docsAPIKeyHeadersecurity schemeRate Limiting (
main.py,requirements.txt)/chatand/chat/streamvia slowapiX-API-Keyheader; falls back to client IP viaX-Forwarded-For(Render sets this header)Retry-Afterheader when limit exceededStartup Validation (
main.py)SERVICE_API_KEYis not set whenENVIRONMENT=productionAUTH_DISABLED=trueopts out of authentication for local developmentSERVICE_API_KEYis unset in non-production environmentsHealth Endpoints
/ping,/cache/stats,/confidence/policyremain unauthenticated — Render health checks and the keep-awake workflow continue to workWorkflow Fix (
.github/workflows/keep-render-awake.yml)/pinginstead of root URL (root returns 404)Config & Docs
SERVICE_API_KEYadded torender.yaml(sync: false, likeGEMINI_API_KEY).env.exampleupdated withSERVICE_API_KEYandAUTH_DISABLEDAcceptance Criteria
POST /chatwithoutX-API-Key(or with a wrong key) returns 401; with the correct key it works as beforeRetry-AfterAUTH_DISABLED=truerender.yamldocumentSERVICE_API_KEYSummary by CodeRabbit
New Features
/health)./memory/{user_id}) for storing and managing long-term context./feedback, plus admin stats/records) with rate-limited submissions.X-API-Keywith a local-development toggle.Documentation
Chores
/healthand adjusted the keep-awake request to/ping.