From 6705a7c0ecfcff35c0ad9ab0ceeae333b9e11168 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Mon, 17 Aug 2026 09:28:48 +0000 Subject: [PATCH 1/2] Add API rate limiting and LLM concurrency caps Unbounded ingest dumps and explain fan-out can overload the process and the LLM provider. Token-bucket limits on ingest/query, the existing ingest queue ceiling, and a process-wide LLM semaphore keep G9 backpressure in place without changing queue semantics. Closes #18 Co-authored-by: Leonardo --- .env.example | 19 +- README.md | 17 +- src/api/app.py | 5 +- src/api/ratelimit.py | 160 ++++++++++++++++ src/config/settings.py | 17 +- src/core/explain/summarizer.py | 13 +- src/core/ingestion/backpressure.py | 7 +- src/core/llm/provider.py | 81 +++++++- src/core/retrieval/question_router.py | 93 ++++++---- tests/unit/test_llm_concurrency.py | 87 +++++++++ tests/unit/test_ratelimit.py | 254 ++++++++++++++++++++++++++ tests/unit/test_settings.py | 31 ++++ 12 files changed, 725 insertions(+), 59 deletions(-) create mode 100644 src/api/ratelimit.py create mode 100644 tests/unit/test_llm_concurrency.py create mode 100644 tests/unit/test_ratelimit.py diff --git a/.env.example b/.env.example index 95cc871..052ccd3 100644 --- a/.env.example +++ b/.env.example @@ -94,11 +94,24 @@ MAX_CLUSTERS_FOR_EXPLAIN=10 # Seconds the worker sleeps between polling for new jobs when the queue is empty. WORKER_POLL_INTERVAL=2 -# Ingest backpressure (minimal stand-in for full API/LLM rate limiting). -# POST /v1/ingestions and POST /v1/ingestions/lines return 429 + Retry-After -# when pending worker jobs >= INGEST_QUEUE_MAX. +# Ingest queue backpressure. POST /v1/ingestions and POST /v1/ingestions/lines +# return 429 INGEST_QUEUE_FULL + Retry-After when pending worker jobs >= max. +# Tail ticks skip the same ceiling. Distinct from RATE_LIMITED (token bucket). INGEST_QUEUE_MAX=100 INGEST_RETRY_AFTER_SECONDS=5 + +# API token-bucket rate limiting (ingest writes + query routes). In-memory per +# process. 0 rps = unlimited for that category. Defaults are high so local demo +# and tests are not 429'd. Identity is API key id, or "anonymous" when auth is off. +# RATELIMIT_ENABLED=true +# RATELIMIT_INGEST_RPS=100 +# RATELIMIT_QUERY_RPS=100 +# RATELIMIT_BURST=100 +# RATELIMIT_RETRY_AFTER_SECONDS=1 + +# Max in-flight LLM provider calls (process-wide; independent of API concurrency). +# 0 = unlimited. Noop provider skips the wait. +# LLM_MAX_CONCURRENCY=4 # Max NDJSON lines accepted by POST /v1/ingestions/lines (over cap → 400). INGEST_PUSH_MAX_LINES=5000 # Tail-mode poll cadence and consecutive-error auto-pause. diff --git a/README.md b/README.md index 4c7cb37..1296591 100644 --- a/README.md +++ b/README.md @@ -688,6 +688,14 @@ All settings are read from `.env`, environment variables, or CLI flags. Priority | `OIDC_JWKS_URL` | _(empty)_ | Optional JWKS URL; default is `{issuer}/.well-known/openid-configuration` then `{issuer}/.well-known/jwks.json` | | `API_BIND_HOST` | `127.0.0.1` | Host the startup guard treats as the bind address. `make api` sets this to `0.0.0.0` to match uvicorn | | `AUTH_REFUSE_INSECURE_BIND` | `false` | If `true`, refuse to start when auth is off and the bind host is not loopback; if `false`, log a warning only | +| `RATELIMIT_ENABLED` | `true` | Token-bucket rate limiting on ingest writes and query routes. In-memory per process (not shared across workers) | +| `RATELIMIT_INGEST_RPS` | `100` | Steady-state tokens/sec for `POST /v1/ingestions*`. `0` = unlimited | +| `RATELIMIT_QUERY_RPS` | `100` | Steady-state tokens/sec for `/v1/query*`. `0` = unlimited | +| `RATELIMIT_BURST` | `100` | Bucket size (max tokens) per API key (or `anonymous` when auth is off) | +| `RATELIMIT_RETRY_AFTER_SECONDS` | `1` | `Retry-After` value on `429 RATE_LIMITED` | +| `INGEST_QUEUE_MAX` | `100` | Pending worker-job ceiling; over this, ingest returns `429 INGEST_QUEUE_FULL` | +| `INGEST_RETRY_AFTER_SECONDS` | `5` | `Retry-After` value on `429 INGEST_QUEUE_FULL` | +| `LLM_MAX_CONCURRENCY` | `4` | Max in-flight LLM provider calls process-wide. `0` = unlimited. Noop does not wait | | `WEBHOOK_SECRET` | _(empty)_ | Fallback HMAC secret for ingest completion callbacks when auth is off or the API key has no per-key `whsec_` | | `WEBHOOK_MAX_RETRIES` | `5` | Extra webhook POST attempts after the first (6 POSTs by default) on 5xx / 429 / connect errors | | `WEBHOOK_TIMEOUT` | `10` | Per-attempt HTTP timeout in seconds for completion callbacks | @@ -699,7 +707,7 @@ All settings are read from `.env`, environment variables, or CLI flags. Priority raglogs is fully useful without any LLM. The `--no-llm` flag (or `LLM_PROVIDER=disabled`) activates deterministic template-based summaries. -When an LLM is configured, it receives only a small curated evidence packet — not raw logs. The prompt enforces fixed output structure, prohibits fabrication, and requires explicit uncertainty statements when evidence is insufficient. +When an LLM is configured, it receives only a small curated evidence packet — not raw logs. The prompt enforces fixed output structure, prohibits fabrication, and requires explicit uncertainty statements when evidence is insufficient. In-flight provider calls are capped by `LLM_MAX_CONCURRENCY` (CLI and API share the process semaphore; the noop provider does not block). ### OpenAI @@ -1001,7 +1009,12 @@ curl -X POST http://localhost:8000/v1/ingestions/$ID:stop `stop` is terminal. After `TAIL_ERROR_THRESHOLD` consecutive poll failures (default 5) a tail job auto-pauses; `/health` reports `tail_jobs.running` and `tail_jobs.paused`. -**Backpressure.** When pending worker jobs ≥ `INGEST_QUEUE_MAX` (default 100), `POST /v1/ingestions` and `POST /v1/ingestions/lines` return **429** with `Retry-After` (`INGEST_RETRY_AFTER_SECONDS`, default 5) and body `{"error_code":"INGEST_QUEUE_FULL","message":"..."}`. This is a queue-depth stand-in, not full API/LLM rate limiting. +**Backpressure and rate limiting.** Two independent 429s: + +- **Queue depth.** When pending worker jobs ≥ `INGEST_QUEUE_MAX` (default 100), `POST /v1/ingestions` and `POST /v1/ingestions/lines` return **429** with `Retry-After` (`INGEST_RETRY_AFTER_SECONDS`, default 5) and body `{"error_code":"INGEST_QUEUE_FULL","message":"..."}`. Tail ticks skip the same ceiling. +- **API token bucket.** `POST /v1/ingestions*` (writes) and `/v1/query*` (plus unversioned aliases) are limited per API key (`request.state.auth_principal.key_id`, or a single `anonymous` bucket when `AUTH_ENABLED=false`). Exceeding the bucket returns **429** with `Retry-After` (`RATELIMIT_RETRY_AFTER_SECONDS`, default 1) and body `{"error_code":"RATE_LIMITED","message":"..."}`. Defaults (`RATELIMIT_INGEST_RPS` / `RATELIMIT_QUERY_RPS` / `RATELIMIT_BURST` = 100) are high enough for local demo and tests; `0` rps means unlimited for that category. `/health`, `/docs`, static UI, and `/config` are not limited. Buckets are in-memory per process. + +LLM calls are separately capped by `LLM_MAX_CONCURRENCY` (default 4) so a burst of `explain` cannot fan out unbounded provider requests. **Idempotency-Key.** `POST /v1/ingestions` (batch enqueue and tail create; also the deprecated `/ingestions` alias) honors an `Idempotency-Key` header (max 256 characters). A repeat **in the same isolation scope** within `INGEST_IDEMPOTENCY_TTL_SECONDS` (default 86400) returns the original **202** job — the same `worker_job_id` for batch, the same `ingestion_job_id` for tail — instead of starting a new one. Reusing another scope's key returns **409** `IDEMPOTENCY_SCOPE_CONFLICT`. Empty keys return **400**. GET routes ignore the header. `POST /v1/ingestions/lines` does not use the header; duplicate push/tail lines are handled by content dedup instead. diff --git a/src/api/app.py b/src/api/app.py index 3723c7e..78fa427 100644 --- a/src/api/app.py +++ b/src/api/app.py @@ -11,6 +11,7 @@ from src.api.auth.middleware import AuthMiddleware from src.api.auth.scope import ScopeResolutionError, scope_error_response from src.api.deprecation import DeprecationHeaderMiddleware +from src.api.ratelimit import RateLimitMiddleware from src.api.routes import ask, clusters, compare_windows, config, explain, health, ingestions, timeline, ui _OPENAPI_DESCRIPTION = """Incident explanation API — ask your logs what happened. @@ -60,7 +61,9 @@ def generate(route: APIRoute) -> str: lifespan=lifespan, ) -# Last added middleware is outermost: deprecation headers apply even to auth errors. +# Last added middleware is outermost: deprecation headers apply even to auth +# errors. Rate limiting sits inside auth so request.state.auth_principal is set. +app.add_middleware(RateLimitMiddleware) app.add_middleware(AuthMiddleware) app.add_middleware(DeprecationHeaderMiddleware) diff --git a/src/api/ratelimit.py b/src/api/ratelimit.py new file mode 100644 index 0000000..5d71d67 --- /dev/null +++ b/src/api/ratelimit.py @@ -0,0 +1,160 @@ +"""Per-API-key token-bucket rate limiting for ingest and query routes (G9). + +Buckets are in-memory and per process — they are not shared across +uvicorn workers or hosts. Identity is ``request.state.auth_principal.key_id`` +when present; otherwise ``\"anonymous\"`` (including AUTH_ENABLED=false). + +``RATELIMIT_INGEST_RPS`` / ``RATELIMIT_QUERY_RPS`` of 0 means unlimited for +that category. ``RATELIMIT_ENABLED=false`` disables limiting entirely. +Health, docs, static UI, config, and key-admin routes are not limited. +""" + +from __future__ import annotations + +import re +import threading +import time +from typing import Awaitable, Callable, Literal + +from starlette.middleware.base import BaseHTTPMiddleware +from starlette.requests import Request +from starlette.responses import JSONResponse, Response + +ERROR_RATE_LIMITED = "RATE_LIMITED" + +RateLimitKind = Literal["ingest", "query"] + +_API_VERSION_PREFIX = re.compile(r"^/v\d+(?=/|$)") + +_store_lock = threading.Lock() +_buckets: dict[tuple[str, str], "_TokenBucket"] = {} + + +class _TokenBucket: + """Thread-safe token bucket. Rate/burst changes reset the fill level.""" + + def __init__(self) -> None: + self._lock = threading.Lock() + self._tokens: float = 0.0 + self._updated: float = time.monotonic() + self._rate: float = 0.0 + self._burst: float = 0.0 + + def consume(self, rate: float, burst: float, tokens: float = 1.0) -> bool: + if rate <= 0: + return True + capacity = max(float(burst), tokens) + with self._lock: + now = time.monotonic() + if rate != self._rate or capacity != self._burst: + self._rate = rate + self._burst = capacity + self._tokens = capacity + self._updated = now + else: + elapsed = now - self._updated + self._tokens = min(capacity, self._tokens + elapsed * rate) + self._updated = now + if self._tokens >= tokens: + self._tokens -= tokens + return True + return False + + +def reset_rate_limiter() -> None: + """Drop all buckets. Used by tests so cases do not leak fill state.""" + with _store_lock: + _buckets.clear() + + +def _normalize_path(path: str) -> str: + if not path: + return "/" + if len(path) > 1 and path.endswith("/"): + return path.rstrip("/") + return path + + +def rate_limit_kind(method: str, path: str) -> RateLimitKind | None: + """Return ingest/query for limited routes, else None (exempt).""" + normalized = _API_VERSION_PREFIX.sub("", _normalize_path(path), count=1) + if not normalized: + normalized = "/" + verb = method.upper() + if normalized == "/ingestions" or normalized.startswith("/ingestions/"): + if verb == "POST": + return "ingest" + return None + if normalized == "/query" or normalized.startswith("/query/"): + return "query" + return None + + +def bucket_identity(request: Request) -> str: + """Per-key identity, or anonymous when auth is off / principal missing.""" + principal = getattr(request.state, "auth_principal", None) + if principal is None: + return "anonymous" + key_id = getattr(principal, "key_id", None) + if key_id: + return str(key_id) + subject = getattr(principal, "subject", None) + if subject: + return f"oidc:{subject}" + return "anonymous" + + +def allow_request(kind: RateLimitKind, identity: str, rate: float, burst: float) -> bool: + """Consume one token from the (kind, identity) bucket. True if allowed.""" + if rate <= 0: + return True + key = (kind, identity) + with _store_lock: + bucket = _buckets.get(key) + if bucket is None: + bucket = _TokenBucket() + _buckets[key] = bucket + return bucket.consume(rate, burst) + + +def _limited_response(retry_after: int) -> JSONResponse: + return JSONResponse( + status_code=429, + content={ + "error_code": ERROR_RATE_LIMITED, + "message": "Rate limit exceeded; retry after the Retry-After delay.", + }, + headers={"Retry-After": str(int(retry_after))}, + ) + + +class RateLimitMiddleware(BaseHTTPMiddleware): + """Token-bucket limiter applied after auth so key_id is available.""" + + async def dispatch( + self, + request: Request, + call_next: Callable[[Request], Awaitable[Response]], + ) -> Response: + from src.config import get_settings + + settings = get_settings() + if not settings.ratelimit_enabled: + return await call_next(request) + + kind = rate_limit_kind(request.method, request.url.path) + if kind is None: + return await call_next(request) + + rate = ( + settings.ratelimit_ingest_rps + if kind == "ingest" + else settings.ratelimit_query_rps + ) + if rate <= 0: + return await call_next(request) + + identity = bucket_identity(request) + if not allow_request(kind, identity, rate, settings.ratelimit_burst): + return _limited_response(settings.ratelimit_retry_after_seconds) + return await call_next(request) diff --git a/src/config/settings.py b/src/config/settings.py index edb86bb..f778934 100644 --- a/src/config/settings.py +++ b/src/config/settings.py @@ -50,13 +50,28 @@ class Settings(BaseSettings): # Worker worker_poll_interval: int = 2 # seconds between idle polls - # Ingest backpressure (minimal G9 stand-in) and push/tail (G4) + # Ingest backpressure (G9) and push/tail (G4). Tail ticks skip when the + # pending worker queue is at ingest_queue_max. ingest_queue_max: int = 100 ingest_retry_after_seconds: int = 5 ingest_push_max_lines: int = 5000 tail_poll_interval: int = 30 tail_error_threshold: int = 5 + # API token-bucket rate limiting (G9). In-memory per process — not shared + # across workers. Defaults are high so local demo and TestClient suites + # are not 429'd. 0 rps = unlimited for that category. Identity is + # auth_principal.key_id, else "anonymous" when AUTH_ENABLED=false. + ratelimit_enabled: bool = True + ratelimit_ingest_rps: float = 100.0 + ratelimit_query_rps: float = 100.0 + ratelimit_burst: float = 100.0 + ratelimit_retry_after_seconds: int = 1 + + # Process-wide cap on in-flight LLM provider calls (G9). Independent of + # API concurrency. 0 = unlimited. Noop provider skips the wait. + llm_max_concurrency: int = 4 + # HMAC-signed ingest completion webhooks (G5). WEBHOOK_SECRET is the # fallback when AUTH_ENABLED=false or the API key has no per-key secret. webhook_secret: str = "" diff --git a/src/core/explain/summarizer.py b/src/core/explain/summarizer.py index 07680fd..bd12c55 100644 --- a/src/core/explain/summarizer.py +++ b/src/core/explain/summarizer.py @@ -10,7 +10,7 @@ from src.core.explain.confidence import compute_confidence from src.core.explain.evidence import EvidencePacket, assemble_evidence from src.core.explain.templates import render_insufficient_evidence, render_text_summary -from src.core.llm.provider import NoopLLMProvider, build_llm_provider +from src.core.llm.provider import build_llm_provider from src.db.models import DEFAULT_LOG_SCOPE, IngestionJob from src.db.scope_filter import filter_ingestion_jobs_by_scope @@ -116,12 +116,11 @@ def explain_window( if not no_llm and settings.llm_provider != "disabled": try: llm = build_llm_provider(settings) - if not isinstance(llm, NoopLLMProvider): - evidence_dict = _packet_to_dict(packet) - llm_text = llm.generate_summary(evidence_dict) - if llm_text: - summary_text = llm_text - mode = "llm" + evidence_dict = _packet_to_dict(packet) + llm_text = llm.generate_summary(evidence_dict) + if llm_text: + summary_text = llm_text + mode = "llm" except Exception: # Degrade gracefully pass diff --git a/src/core/ingestion/backpressure.py b/src/core/ingestion/backpressure.py index 45b9efb..6080fd0 100644 --- a/src/core/ingestion/backpressure.py +++ b/src/core/ingestion/backpressure.py @@ -1,4 +1,9 @@ -"""Ingest worker-queue backpressure (minimal G9 stand-in).""" +"""Ingest worker-queue backpressure (G9). + +When pending WorkerJobs >= INGEST_QUEUE_MAX, POST /ingestions and +POST /ingestions/lines return 429 INGEST_QUEUE_FULL. Tail ticks skip +the same ceiling. Distinct from API RATE_LIMITED (token bucket). +""" from sqlalchemy import func, select from sqlalchemy.orm import Session diff --git a/src/core/llm/provider.py b/src/core/llm/provider.py index a7a5f08..ca5a327 100644 --- a/src/core/llm/provider.py +++ b/src/core/llm/provider.py @@ -1,5 +1,8 @@ import json -from typing import Optional, Protocol, runtime_checkable +import threading +from collections.abc import Iterator +from contextlib import contextmanager +from typing import Any, Protocol, runtime_checkable import httpx from tenacity import retry, stop_after_attempt, wait_exponential @@ -103,8 +106,73 @@ def generate_summary(self, evidence_packet: dict) -> str: return data.get("response", "").strip() -def build_llm_provider(settings) -> LLMProvider: - """Factory: build the configured LLM provider.""" +_llm_sem_lock = threading.Lock() +_llm_semaphore: threading.Semaphore | None = None +_llm_semaphore_n: int | None = None + + +def reset_llm_concurrency_limiter() -> None: + """Drop the process semaphore so tests can change LLM_MAX_CONCURRENCY.""" + global _llm_semaphore, _llm_semaphore_n + with _llm_sem_lock: + _llm_semaphore = None + _llm_semaphore_n = None + + +def _semaphore_for(n: int) -> threading.Semaphore: + global _llm_semaphore, _llm_semaphore_n + with _llm_sem_lock: + if _llm_semaphore is None or _llm_semaphore_n != n: + _llm_semaphore = threading.Semaphore(n) + _llm_semaphore_n = n + return _llm_semaphore + + +@contextmanager +def llm_concurrency_slot(*, skip: bool = False) -> Iterator[None]: + """Acquire the global LLM in-flight cap, or skip (noop / unlimited).""" + if skip: + yield + return + from src.config import get_settings + + n = int(get_settings().llm_max_concurrency) + if n <= 0: + yield + return + sem = _semaphore_for(n) + sem.acquire() + try: + yield + finally: + sem.release() + + +class CappedLLMProvider: + """Wrap an inner provider with the process-wide concurrency semaphore. + + Noop inner providers skip the wait so deterministic mode never blocks, + but still go through this entrypoint (CLI and API share the semaphore). + """ + + def __init__(self, inner: LLMProvider) -> None: + self.inner = inner + + def generate_summary(self, evidence_packet: dict) -> str: + skip = isinstance(self.inner, NoopLLMProvider) + with llm_concurrency_slot(skip=skip): + return self.inner.generate_summary(evidence_packet) + + +def unwrap_llm_provider(provider: LLMProvider) -> LLMProvider: + """Return the inner provider if ``provider`` is concurrency-capped.""" + inner: LLMProvider = provider + while isinstance(inner, CappedLLMProvider): + inner = inner.inner + return inner + + +def _build_inner_llm_provider(settings: Any) -> LLMProvider: if settings.llm_provider == "openai": if not settings.openai_api_key: return NoopLLMProvider() @@ -113,9 +181,14 @@ def build_llm_provider(settings) -> LLMProvider: model=settings.llm_model, base_url=settings.openai_base_url, ) - elif settings.llm_provider == "ollama": + if settings.llm_provider == "ollama": return OllamaLLMProvider( model=settings.llm_model, base_url=settings.ollama_base_url, ) return NoopLLMProvider() + + +def build_llm_provider(settings: Any) -> LLMProvider: + """Factory: build the configured LLM provider (concurrency-capped).""" + return CappedLLMProvider(_build_inner_llm_provider(settings)) diff --git a/src/core/retrieval/question_router.py b/src/core/retrieval/question_router.py index f05e793..99466ae 100644 --- a/src/core/retrieval/question_router.py +++ b/src/core/retrieval/question_router.py @@ -249,7 +249,11 @@ def answer_question( scope: str = DEFAULT_LOG_SCOPE, ) -> AskResult: from src.config import get_settings - from src.core.llm.provider import NoopLLMProvider, build_llm_provider + from src.core.llm.provider import ( + NoopLLMProvider, + build_llm_provider, + unwrap_llm_provider, + ) from src.utils.time import resolve_window settings = get_settings() @@ -313,7 +317,7 @@ def answer_question( answer_text = "" llm = build_llm_provider(settings) - if not isinstance(llm, NoopLLMProvider): + if not isinstance(unwrap_llm_provider(llm), NoopLLMProvider): try: answer_text = _call_llm_ask(llm, question, evidence_packet) if answer_text: @@ -393,53 +397,62 @@ def _retrieve_matching_logs( return matching, "fallback" -def _call_llm_ask(llm, question: str, evidence_packet: dict) -> str: +def _call_llm_ask(llm: object, question: str, evidence_packet: dict) -> str: """ Call the LLM provider with the ask-specific system prompt. Constructs the HTTP call directly rather than reusing generate_summary, - which uses the incident-summary system prompt. + which uses the incident-summary system prompt. Shares the G9 LLM + concurrency semaphore with generate_summary. """ import json import httpx - from src.core.llm.provider import OpenAILLMProvider, OllamaLLMProvider + from src.core.llm.provider import ( + NoopLLMProvider, + OpenAILLMProvider, + OllamaLLMProvider, + llm_concurrency_slot, + unwrap_llm_provider, + ) + inner = unwrap_llm_provider(llm) payload_str = json.dumps(evidence_packet, default=str, indent=2) user_message = f"Question: {question}\n\nLog evidence:\n{payload_str}" - if isinstance(llm, OpenAILLMProvider): - with httpx.Client(timeout=60) as client: - resp = client.post( - f"{llm.base_url}/chat/completions", - headers={ - "Authorization": f"Bearer {llm.api_key}", - "Content-Type": "application/json", - }, - json={ - "model": llm.model, - "max_tokens": 400, - "temperature": 0, - "messages": [ - {"role": "system", "content": ASK_SYSTEM_PROMPT}, - {"role": "user", "content": user_message}, - ], - }, - ) - resp.raise_for_status() - return resp.json()["choices"][0]["message"]["content"].strip() - - elif isinstance(llm, OllamaLLMProvider): - prompt = f"{ASK_SYSTEM_PROMPT}\n\n{user_message}\n\nAnswer:" - with httpx.Client(timeout=120) as client: - resp = client.post( - f"{llm.base_url}/api/generate", - json={ - "model": llm.model, - "prompt": prompt, - "stream": False, - "options": {"temperature": 0}, - }, - ) - resp.raise_for_status() - return resp.json().get("response", "").strip() + with llm_concurrency_slot(skip=isinstance(inner, NoopLLMProvider)): + if isinstance(inner, OpenAILLMProvider): + with httpx.Client(timeout=60) as client: + resp = client.post( + f"{inner.base_url}/chat/completions", + headers={ + "Authorization": f"Bearer {inner.api_key}", + "Content-Type": "application/json", + }, + json={ + "model": inner.model, + "max_tokens": 400, + "temperature": 0, + "messages": [ + {"role": "system", "content": ASK_SYSTEM_PROMPT}, + {"role": "user", "content": user_message}, + ], + }, + ) + resp.raise_for_status() + return resp.json()["choices"][0]["message"]["content"].strip() + + if isinstance(inner, OllamaLLMProvider): + prompt = f"{ASK_SYSTEM_PROMPT}\n\n{user_message}\n\nAnswer:" + with httpx.Client(timeout=120) as client: + resp = client.post( + f"{inner.base_url}/api/generate", + json={ + "model": inner.model, + "prompt": prompt, + "stream": False, + "options": {"temperature": 0}, + }, + ) + resp.raise_for_status() + return resp.json().get("response", "").strip() return "" diff --git a/tests/unit/test_llm_concurrency.py b/tests/unit/test_llm_concurrency.py new file mode 100644 index 0000000..a275908 --- /dev/null +++ b/tests/unit/test_llm_concurrency.py @@ -0,0 +1,87 @@ +"""Unit tests for the process-wide LLM concurrency cap (no database).""" +from __future__ import annotations + +import threading +import time +from unittest.mock import patch + +from src.config.settings import Settings +from src.core.llm.provider import ( + CappedLLMProvider, + NoopLLMProvider, + build_llm_provider, + reset_llm_concurrency_limiter, + unwrap_llm_provider, +) + + +def test_build_wraps_noop_and_generate_summary_returns_empty() -> None: + settings = Settings(_env_file=None, llm_provider="disabled") + llm = build_llm_provider(settings) + assert isinstance(llm, CappedLLMProvider) + assert isinstance(unwrap_llm_provider(llm), NoopLLMProvider) + assert llm.generate_summary({"x": 1}) == "" + + +def test_semaphore_serializes_inner_provider_calls() -> None: + reset_llm_concurrency_limiter() + settings = Settings(_env_file=None, llm_max_concurrency=1) + inner_current = 0 + max_seen = 0 + lock = threading.Lock() + started = threading.Event() + release = threading.Event() + + class FakeProvider: + def generate_summary(self, evidence_packet: dict) -> str: + nonlocal inner_current, max_seen + with lock: + inner_current += 1 + max_seen = max(max_seen, inner_current) + started.set() + assert release.wait(timeout=5) + with lock: + inner_current -= 1 + return "ok" + + capped = CappedLLMProvider(FakeProvider()) + with patch("src.config.get_settings", return_value=settings): + t1 = threading.Thread(target=lambda: capped.generate_summary({})) + t2 = threading.Thread(target=lambda: capped.generate_summary({})) + t1.start() + assert started.wait(timeout=2) + t2.start() + time.sleep(0.15) + assert max_seen == 1 + release.set() + t1.join(timeout=2) + t2.join(timeout=2) + + assert max_seen == 1 + reset_llm_concurrency_limiter() + + +def test_noop_does_not_block_when_semaphore_is_held() -> None: + reset_llm_concurrency_limiter() + settings = Settings(_env_file=None, llm_max_concurrency=1) + started = threading.Event() + block = threading.Event() + + class BlockingProvider: + def generate_summary(self, evidence_packet: dict) -> str: + started.set() + assert block.wait(timeout=5) + return "real" + + real = CappedLLMProvider(BlockingProvider()) + noop = CappedLLMProvider(NoopLLMProvider()) + with patch("src.config.get_settings", return_value=settings): + holder = threading.Thread(target=lambda: real.generate_summary({})) + holder.start() + assert started.wait(timeout=2) + t0 = time.monotonic() + assert noop.generate_summary({}) == "" + assert time.monotonic() - t0 < 0.5 + block.set() + holder.join(timeout=2) + reset_llm_concurrency_limiter() diff --git a/tests/unit/test_ratelimit.py b/tests/unit/test_ratelimit.py new file mode 100644 index 0000000..c93751c --- /dev/null +++ b/tests/unit/test_ratelimit.py @@ -0,0 +1,254 @@ +"""Unit tests for API token-bucket rate limiting (no database).""" +from __future__ import annotations + +import uuid +from unittest.mock import MagicMock, patch + +import pytest +from fastapi.testclient import TestClient + +from src.api.app import app +from src.api.auth.middleware import AuthPrincipal +from src.api.ratelimit import ( + allow_request, + bucket_identity, + rate_limit_kind, + reset_rate_limiter, +) +from src.config.settings import Settings + +client = TestClient(app, raise_server_exceptions=False) + + +@pytest.fixture(autouse=True) +def _reset_buckets() -> None: + reset_rate_limiter() + yield + reset_rate_limiter() + + +def _tiny_settings(**kwargs: object) -> Settings: + values: dict[str, object] = { + "ratelimit_enabled": True, + "ratelimit_ingest_rps": 1.0, + "ratelimit_query_rps": 1.0, + "ratelimit_burst": 1.0, + "ratelimit_retry_after_seconds": 7, + "auth_enabled": False, + } + values.update(kwargs) + return Settings(_env_file=None, **values) + + +def _key(role: str, key_id: uuid.UUID | None = None) -> MagicMock: + record = MagicMock() + record.id = key_id or uuid.uuid4() + record.name = "test" + record.key_prefix = "rlk_testhash" + record.role = role + record.scope = "default" + record.revoked_at = None + record.allow_scope_override = False + return record + + +class TestRateLimitKind: + def test_ingest_writes_only(self) -> None: + assert rate_limit_kind("POST", "/v1/ingestions") == "ingest" + assert rate_limit_kind("POST", "/ingestions") == "ingest" + assert rate_limit_kind("POST", "/v1/ingestions/lines") == "ingest" + assert rate_limit_kind("POST", "/v1/ingestions/abc:pause") == "ingest" + assert rate_limit_kind("GET", "/v1/ingestions") is None + assert rate_limit_kind("GET", "/v1/ingestions/latest") is None + + def test_query_routes_versioned_and_alias(self) -> None: + assert rate_limit_kind("POST", "/v1/query/explain") == "query" + assert rate_limit_kind("POST", "/query/explain") == "query" + assert rate_limit_kind("POST", "/v1/query/ask") == "query" + assert rate_limit_kind("GET", "/v1/query/clusters") == "query" + + def test_exempt_paths(self) -> None: + assert rate_limit_kind("GET", "/health") is None + assert rate_limit_kind("GET", "/docs") is None + assert rate_limit_kind("GET", "/static/js/app.js") is None + assert rate_limit_kind("GET", "/v1/config") is None + assert rate_limit_kind("GET", "/") is None + + +class TestTokenBucket: + def test_allow_then_deny_and_per_key_isolation(self) -> None: + assert allow_request("ingest", "key-a", rate=1.0, burst=1.0) is True + assert allow_request("ingest", "key-a", rate=1.0, burst=1.0) is False + assert allow_request("ingest", "key-b", rate=1.0, burst=1.0) is True + + def test_zero_rate_is_unlimited(self) -> None: + for _ in range(20): + assert allow_request("query", "anon", rate=0.0, burst=1.0) is True + + def test_ingest_and_query_buckets_are_separate(self) -> None: + assert allow_request("ingest", "same", rate=1.0, burst=1.0) is True + assert allow_request("ingest", "same", rate=1.0, burst=1.0) is False + assert allow_request("query", "same", rate=1.0, burst=1.0) is True + + +class TestBucketIdentity: + def test_anonymous_without_principal(self) -> None: + request = MagicMock() + request.state.auth_principal = None + assert bucket_identity(request) == "anonymous" + + def test_uses_key_id_when_present(self) -> None: + request = MagicMock() + request.state.auth_principal = AuthPrincipal( + role="ingest", + scope="default", + auth_method="api_key", + key_id="abc-123", + ) + assert bucket_identity(request) == "abc-123" + + +class TestRateLimitHttp: + def test_allow_then_429_with_retry_after_on_query(self) -> None: + settings = _tiny_settings() + with patch("src.config.get_settings", return_value=settings): + first = client.post("/v1/query/explain", json={}) + second = client.post("/v1/query/explain", json={}) + + assert first.status_code != 429 + assert second.status_code == 429 + body = second.json() + assert body["error_code"] == "RATE_LIMITED" + assert "message" in body + assert second.headers.get("retry-after") == "7" + + def test_unversioned_query_alias_is_limited(self) -> None: + settings = _tiny_settings() + with patch("src.config.get_settings", return_value=settings): + client.post("/query/explain", json={}) + second = client.post("/query/explain", json={}) + assert second.status_code == 429 + assert second.json()["error_code"] == "RATE_LIMITED" + + def test_ingest_write_429_distinct_from_queue_full(self) -> None: + settings = _tiny_settings() + with patch("src.config.get_settings", return_value=settings): + first = client.post("/v1/ingestions", json={}) + second = client.post("/v1/ingestions", json={}) + + assert first.status_code == 422 + assert second.status_code == 429 + assert second.json()["error_code"] == "RATE_LIMITED" + assert second.json()["error_code"] != "INGEST_QUEUE_FULL" + + def test_health_not_limited(self) -> None: + settings = _tiny_settings() + with patch("src.config.get_settings", return_value=settings), \ + patch("src.db.session.check_connection", return_value=False): + for _ in range(5): + resp = client.get("/health") + assert resp.status_code == 200 + + def test_config_not_limited(self) -> None: + settings = _tiny_settings() + with patch("src.config.get_settings", return_value=settings): + for _ in range(5): + resp = client.get("/v1/config") + assert resp.status_code == 200 + + def test_zero_rps_is_unlimited(self) -> None: + settings = _tiny_settings(ratelimit_query_rps=0.0, ratelimit_burst=1.0) + with patch("src.config.get_settings", return_value=settings): + statuses = [ + client.post("/v1/query/explain", json={}).status_code + for _ in range(8) + ] + assert 429 not in statuses + + def test_disabled_flag_skips_limiting(self) -> None: + settings = _tiny_settings(ratelimit_enabled=False) + with patch("src.config.get_settings", return_value=settings): + statuses = [ + client.post("/v1/query/explain", json={}).status_code + for _ in range(5) + ] + assert 429 not in statuses + + def test_default_settings_do_not_429_existing_client_traffic(self) -> None: + settings = Settings(_env_file=None) + assert settings.auth_enabled is False + with patch("src.config.get_settings", return_value=settings): + statuses = [ + client.post("/v1/query/explain", json={}).status_code + for _ in range(15) + ] + assert 429 not in statuses + + def test_synthetic_load_exhausts_tiny_bucket(self) -> None: + settings = _tiny_settings(ratelimit_query_rps=1.0, ratelimit_burst=2.0) + with patch("src.config.get_settings", return_value=settings): + statuses = [ + client.post("/v1/query/explain", json={"since": "1h"}).status_code + for _ in range(10) + ] + assert statuses[0] != 429 + assert statuses[1] != 429 + assert statuses.count(429) >= 8 + + def test_per_key_isolation_when_auth_enabled(self) -> None: + key_a_id = uuid.uuid4() + key_b_id = uuid.uuid4() + key_a = _key("admin", key_a_id) + key_b = _key("admin", key_b_id) + + def lookup(token: str) -> MagicMock | None: + if token == "rlk_key_aaaaaa": + return key_a + if token == "rlk_key_bbbbbb": + return key_b + return None + + settings = _tiny_settings( + auth_enabled=True, + ratelimit_ingest_rps=1.0, + ratelimit_burst=1.0, + ) + with patch("src.config.get_settings", return_value=settings), \ + patch("src.api.auth.keys.lookup_api_key", side_effect=lookup): + a1 = client.post( + "/v1/ingestions", + json={}, + headers={"Authorization": "Bearer rlk_key_aaaaaa"}, + ) + a2 = client.post( + "/v1/ingestions", + json={}, + headers={"Authorization": "Bearer rlk_key_aaaaaa"}, + ) + b1 = client.post( + "/v1/ingestions", + json={}, + headers={"Authorization": "Bearer rlk_key_bbbbbb"}, + ) + + assert a1.status_code == 422 + assert a2.status_code == 429 + assert a2.json()["error_code"] == "RATE_LIMITED" + assert b1.status_code == 422 + + def test_queue_full_still_ingest_queue_full_not_rate_limited(self) -> None: + mock_db = MagicMock() + mock_db.__enter__ = MagicMock(return_value=mock_db) + mock_db.__exit__ = MagicMock(return_value=False) + mock_db.execute.return_value.scalar_one.return_value = 100 + settings = Settings(_env_file=None) + with ( + patch("src.config.get_settings", return_value=settings), + patch("src.adapters.file.adapter.discover_files", return_value=["f.log"]), + patch("src.db.session.get_db", side_effect=lambda: mock_db), + ): + resp = client.post("/v1/ingestions", json={"paths": ["/logs"]}) + + assert resp.status_code == 429 + assert resp.json()["error_code"] == "INGEST_QUEUE_FULL" + assert resp.headers.get("retry-after") == "5" diff --git a/tests/unit/test_settings.py b/tests/unit/test_settings.py index b3a99de..edc727c 100644 --- a/tests/unit/test_settings.py +++ b/tests/unit/test_settings.py @@ -131,3 +131,34 @@ def test_ingest_idempotency_ttl_from_env(monkeypatch): monkeypatch.setenv("INGEST_IDEMPOTENCY_TTL_SECONDS", "60") settings = Settings(_env_file=None) assert settings.ingest_idempotency_ttl_seconds == 60 + + +def test_ratelimit_and_llm_concurrency_defaults(): + settings = Settings(_env_file=None) + assert settings.ratelimit_enabled is True + assert settings.ratelimit_ingest_rps == 100.0 + assert settings.ratelimit_query_rps == 100.0 + assert settings.ratelimit_burst == 100.0 + assert settings.ratelimit_retry_after_seconds == 1 + assert settings.llm_max_concurrency == 4 + assert settings.ingest_queue_max == 100 + + +def test_ratelimit_and_llm_concurrency_from_env(monkeypatch): + monkeypatch.setenv("RATELIMIT_ENABLED", "false") + monkeypatch.setenv("RATELIMIT_INGEST_RPS", "2") + monkeypatch.setenv("RATELIMIT_QUERY_RPS", "3") + monkeypatch.setenv("RATELIMIT_BURST", "5") + monkeypatch.setenv("RATELIMIT_RETRY_AFTER_SECONDS", "9") + monkeypatch.setenv("LLM_MAX_CONCURRENCY", "1") + monkeypatch.setenv("INGEST_QUEUE_MAX", "7") + + settings = Settings(_env_file=None) + + assert settings.ratelimit_enabled is False + assert settings.ratelimit_ingest_rps == 2.0 + assert settings.ratelimit_query_rps == 3.0 + assert settings.ratelimit_burst == 5.0 + assert settings.ratelimit_retry_after_seconds == 9 + assert settings.llm_max_concurrency == 1 + assert settings.ingest_queue_max == 7 From f7320ff4150a50b45db85063c679c99dc02507d1 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Mon, 17 Aug 2026 09:39:11 +0000 Subject: [PATCH 2/2] Harden G9 rate-limit and ask concurrency tests The synthetic load case posted a valid explain body that hit get_db; slow connects could refill the 1 rps bucket and flake the 429 count. Ask HTTP concurrency was untested, so dropping the semaphore would not fail CI. Co-authored-by: Leonardo --- tests/unit/test_llm_concurrency.py | 64 ++++++++++++++++++++++++++++++ tests/unit/test_ratelimit.py | 10 +++-- 2 files changed, 70 insertions(+), 4 deletions(-) diff --git a/tests/unit/test_llm_concurrency.py b/tests/unit/test_llm_concurrency.py index a275908..30186db 100644 --- a/tests/unit/test_llm_concurrency.py +++ b/tests/unit/test_llm_concurrency.py @@ -3,16 +3,19 @@ import threading import time +from typing import Any from unittest.mock import patch from src.config.settings import Settings from src.core.llm.provider import ( CappedLLMProvider, NoopLLMProvider, + OpenAILLMProvider, build_llm_provider, reset_llm_concurrency_limiter, unwrap_llm_provider, ) +from src.core.retrieval.question_router import _call_llm_ask def test_build_wraps_noop_and_generate_summary_returns_empty() -> None: @@ -85,3 +88,64 @@ def generate_summary(self, evidence_packet: dict) -> str: block.set() holder.join(timeout=2) reset_llm_concurrency_limiter() + + +def test_call_llm_ask_takes_concurrency_semaphore() -> None: + """Ask HTTP must share the G9 slot; dropping it would allow overlapping posts.""" + reset_llm_concurrency_limiter() + settings = Settings(_env_file=None, llm_max_concurrency=1) + inner_current = 0 + max_seen = 0 + lock = threading.Lock() + started = threading.Event() + release = threading.Event() + + class FakeResponse: + def raise_for_status(self) -> None: + return None + + def json(self) -> dict[str, Any]: + return {"choices": [{"message": {"content": "ok"}}]} + + class FakeClient: + def __init__(self, *args: object, **kwargs: object) -> None: + pass + + def __enter__(self) -> FakeClient: + return self + + def __exit__(self, *args: object) -> bool: + return False + + def post(self, *args: object, **kwargs: object) -> FakeResponse: + nonlocal inner_current, max_seen + with lock: + inner_current += 1 + max_seen = max(max_seen, inner_current) + started.set() + assert release.wait(timeout=5) + with lock: + inner_current -= 1 + return FakeResponse() + + llm = OpenAILLMProvider(api_key="sk-test", model="gpt-4.1-mini") + packet = {"question": "why", "clusters": []} + + def _ask() -> str: + return _call_llm_ask(llm, "why did it fail?", packet) + + with patch("src.config.get_settings", return_value=settings), \ + patch("httpx.Client", FakeClient): + t1 = threading.Thread(target=_ask) + t2 = threading.Thread(target=_ask) + t1.start() + assert started.wait(timeout=2) + t2.start() + time.sleep(0.15) + assert max_seen == 1 + release.set() + t1.join(timeout=2) + t2.join(timeout=2) + + assert max_seen == 1 + reset_llm_concurrency_limiter() diff --git a/tests/unit/test_ratelimit.py b/tests/unit/test_ratelimit.py index c93751c..164671c 100644 --- a/tests/unit/test_ratelimit.py +++ b/tests/unit/test_ratelimit.py @@ -185,15 +185,17 @@ def test_default_settings_do_not_429_existing_client_traffic(self) -> None: assert 429 not in statuses def test_synthetic_load_exhausts_tiny_bucket(self) -> None: + # Empty body 400s at validation — stays in the limiter and does not + # wait on get_db, so a 1 rps refill cannot sneak extra tokens in. settings = _tiny_settings(ratelimit_query_rps=1.0, ratelimit_burst=2.0) with patch("src.config.get_settings", return_value=settings): statuses = [ - client.post("/v1/query/explain", json={"since": "1h"}).status_code + client.post("/v1/query/explain", json={}).status_code for _ in range(10) ] - assert statuses[0] != 429 - assert statuses[1] != 429 - assert statuses.count(429) >= 8 + assert statuses[0] == 400 + assert statuses[1] == 400 + assert statuses.count(429) == 8 def test_per_key_isolation_when_auth_enabled(self) -> None: key_a_id = uuid.uuid4()