From 6c502eac8fcf4f1d049a89bea9caa673ab8b6bb8 Mon Sep 17 00:00:00 2001 From: Prashant Vasudevan <71649489+vprashrex@users.noreply.github.com> Date: Thu, 4 Jun 2026 19:09:36 +0530 Subject: [PATCH 1/7] feat: implement rate monitoring for API endpoints --- backend/app/api/routes/collections.py | 6 +- .../app/api/routes/evaluations/evaluation.py | 6 +- backend/app/api/routes/llm.py | 7 +- backend/app/core/rate_monitor.py | 96 +++++++++++++++++++ backend/app/core/telemetry.py | 30 ++++++ 5 files changed, 141 insertions(+), 4 deletions(-) create mode 100644 backend/app/core/rate_monitor.py diff --git a/backend/app/api/routes/collections.py b/backend/app/api/routes/collections.py index 6c9bee17a..9bc104aa7 100644 --- a/backend/app/api/routes/collections.py +++ b/backend/app/api/routes/collections.py @@ -8,6 +8,7 @@ from app.api.deps import SessionDep, AuthContextDep from app.api.permissions import Permission, require_permission from app.core.telemetry import log_context +from app.core.rate_monitor import monitor_rate from app.crud import ( CollectionCrud, CollectionJobCrud, @@ -85,7 +86,10 @@ def list_collections( description=load_description("collections/create.md"), response_model=APIResponse[CollectionJobImmediatePublic], callbacks=collection_callback_router.routes, - dependencies=[Depends(require_permission(Permission.REQUIRE_PROJECT))], + dependencies=[ + Depends(require_permission(Permission.REQUIRE_PROJECT)), + Depends(monitor_rate("collections")), + ], ) def create_collection( session: SessionDep, diff --git a/backend/app/api/routes/evaluations/evaluation.py b/backend/app/api/routes/evaluations/evaluation.py index 591f5d985..216afe95c 100644 --- a/backend/app/api/routes/evaluations/evaluation.py +++ b/backend/app/api/routes/evaluations/evaluation.py @@ -12,6 +12,7 @@ ) from app.api.deps import AuthContextDep, SessionDep +from app.core.rate_monitor import monitor_rate from app.crud.evaluations import list_evaluation_runs as list_evaluation_runs_crud from app.crud.evaluations.core import group_traces_by_question_id from app.models.evaluation import EvaluationRunPublic @@ -34,7 +35,10 @@ "", description=load_description("evaluation/create_evaluation.md"), response_model=APIResponse[EvaluationRunPublic], - dependencies=[Depends(require_permission(Permission.REQUIRE_PROJECT))], + dependencies=[ + Depends(require_permission(Permission.REQUIRE_PROJECT)), + Depends(monitor_rate("evaluations")), + ], ) def evaluate( session: SessionDep, diff --git a/backend/app/api/routes/llm.py b/backend/app/api/routes/llm.py index 8106046e7..8ce96cf76 100644 --- a/backend/app/api/routes/llm.py +++ b/backend/app/api/routes/llm.py @@ -9,6 +9,7 @@ from app.api.permissions import Permission, require_permission from app.core.cloud.storage import get_cloud_storage from app.core.telemetry import log_context +from app.core.rate_monitor import monitor_rate from app.crud.jobs import JobCrud from app.crud.llm import get_llm_calls_by_job_id from app.models import ( @@ -22,7 +23,6 @@ from app.services.llm.jobs import start_job from app.utils import APIResponse, validate_callback_url, load_description - logger = logging.getLogger(__name__) router = APIRouter(tags=["LLM"]) @@ -50,7 +50,10 @@ def llm_callback_notification(body: APIResponse[LLMCallResponse]): description=load_description("llm/llm_call.md"), response_model=APIResponse[LLMJobImmediatePublic], callbacks=llm_callback_router.routes, - dependencies=[Depends(require_permission(Permission.REQUIRE_PROJECT))], + dependencies=[ + Depends(require_permission(Permission.REQUIRE_PROJECT)), + Depends(monitor_rate("llm_call")), + ], ) def llm_call( _current_user: AuthContextDep, session: SessionDep, request: LLMCallRequest diff --git a/backend/app/core/rate_monitor.py b/backend/app/core/rate_monitor.py new file mode 100644 index 000000000..9cae8513c --- /dev/null +++ b/backend/app/core/rate_monitor.py @@ -0,0 +1,96 @@ +import logging +import time + +from typing import Literal + +import redis + +from app.api.deps import AuthContextDep +from app.core.config import settings + +from app.core.telemetry import record_rate_threshold + +logger = logging.getLogger(__name__) + +# Categores of rates we want to monitor +RateCategory = Literal["llm_call", "collections", "evaluations"] + +# THRESHOLD NUMBERS +THRESHOLDS: dict[RateCategory, int] = { + "llm_call": 15, + "collections": 3, + "evaluations": 5, +} + +# Delete record after 2 minutes from redis +_EXPIRATION_SECONDS = 120 + +_redis_client: redis.Redis = redis.from_url(settings.REDIS_URL, decode_responses=True) + + +# count incrementor after each request and get count +def increment_and_get_count(key: str) -> int | None: + """Increment the count for the given key and return the new count. + The count will automatically expire after _EXPIRATION_SECONDS. + """ + try: + pipe = _redis_client.pipeline() + pipe.incr(key) + pipe.expire(key, _EXPIRATION_SECONDS) + count, _ = pipe.execute() + return count + except Exception as e: + logger.error( + f"[increment_and_get_count] Error incrementing count for {key}: {e}" + ) + return None + + +def monitor_rate(category: RateCategory): + """Monitor the rate of events for the given category. If the rate exceeds the threshold, record it in telemetry. + + Usage: + dependencies=[ + Depends(require_permission(Permission.REQUIRE_PROJECT)), + Depends(monitor_rate("llm")), + ] + """ + + def _checker(auth_context: AuthContextDep) -> None: + org = auth_context.organization + if org is None: + return + + threshold = THRESHOLDS.get(category, None) + if threshold is None: + logger.warning( + f"[monitor_rate] No threshold defined for category {category}" + ) + return + + minute_bucket = int(time.time() // 60) + redis_key = f"rate_monitor:{category}:{org.id}:{minute_bucket}" + + try: + count = increment_and_get_count(redis_key) + if count is not None and count > threshold: + logger.warning( + f"[monitor_rate] Rate threshold exceeded for {category} in org {org.id}: count={count}" + ) + record_rate_threshold( + org_id=org.id, + org_name=org.name, + category=category, + request_count=count, + threshold=threshold, + ) + except redis.RedisError as e: + logger.error( + "[monitor_rate] Redis unavailable, skipping rate check " + "(org_id=%s category=%s)", + org.id, + category, + exc_info=e, + ) + + return _checker diff --git a/backend/app/core/telemetry.py b/backend/app/core/telemetry.py index 963c3086b..3fa0e84a9 100644 --- a/backend/app/core/telemetry.py +++ b/backend/app/core/telemetry.py @@ -453,6 +453,36 @@ def record_stale_pending_jobs( ) +def record_rate_threshold( + *, + org_id: int, + org_name: str | None, + category: str, + request_count: int, + threshold: int, +) -> None: + """Emit rate threshold exceeded event to Sentry.""" + + try: + if not sentry_sdk.get_client().is_active(): + return + with sentry_sdk.push_scope() as scope: + scope.set_tag("alert.type", "threshold_rate_monitor") + scope.set_tag("tenant.org_id", org_id) + scope.set_tag("route_category", category) + scope.set_extra("request_count", request_count) + scope.set_extra("threshold", threshold) + sentry_sdk.capture_message( + f"[Threshold-Monitor] {category} rate limit exceeded for org {org_id} | {org_name}: {request_count} req/min " + f"(limit {threshold}/min)", + level="warning", + ) + except Exception as e: + logger.exception( + "[record_rate_threshold_exceeded] Failed to emit alert", exc_info=e + ) + + def flush_telemetry(timeout_millis: int = 10000) -> None: """Force-flush OTel spans into Sentry, then flush Sentry's transport. From 4a77d057dc170d6a1db28c208fa46845699d88db Mon Sep 17 00:00:00 2001 From: Prashant Vasudevan <71649489+vprashrex@users.noreply.github.com> Date: Thu, 4 Jun 2026 19:22:52 +0530 Subject: [PATCH 2/7] feat: add threshold rates for monitoring LLM calls, collections, and evaluations --- backend/app/core/config.py | 5 +++++ backend/app/core/rate_monitor.py | 6 +++--- 2 files changed, 8 insertions(+), 3 deletions(-) diff --git a/backend/app/core/config.py b/backend/app/core/config.py index 720846eb9..5c8320fa4 100644 --- a/backend/app/core/config.py +++ b/backend/app/core/config.py @@ -140,6 +140,11 @@ def AWS_S3_BUCKET(self) -> str: BACKEND_SERVICE_NAME: str = "kaapi-backend" CRON_SERVICE_NAME: str = "kaapi-cron" + # Threshold Request Rate per minute + THRESHOLD_LLM_CALL_RATE: int = 15 + THRESHOLD_COLLECTIONS_RATE: int = 3 + THRESHOLD_EVALUATIONS_RATE: int = 3 + # Celery Configuration CELERY_WORKER_CONCURRENCY: int | None = None CELERY_WORKER_MAX_TASKS_PER_CHILD: int = 150 diff --git a/backend/app/core/rate_monitor.py b/backend/app/core/rate_monitor.py index 9cae8513c..0ed376dc2 100644 --- a/backend/app/core/rate_monitor.py +++ b/backend/app/core/rate_monitor.py @@ -17,9 +17,9 @@ # THRESHOLD NUMBERS THRESHOLDS: dict[RateCategory, int] = { - "llm_call": 15, - "collections": 3, - "evaluations": 5, + "llm_call": settings.THRESHOLD_LLM_CALL_RATE, + "collections": settings.THRESHOLD_COLLECTIONS_RATE, + "evaluations": settings.THRESHOLD_EVALUATIONS_RATE, } # Delete record after 2 minutes from redis From 4263c467f0d7fb06c5f9039365995454304c7746 Mon Sep 17 00:00:00 2001 From: Prashant Vasudevan <71649489+vprashrex@users.noreply.github.com> Date: Thu, 4 Jun 2026 19:26:07 +0530 Subject: [PATCH 3/7] feat: update monitor_rate usage to accept dynamic category parameter --- backend/app/core/rate_monitor.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/backend/app/core/rate_monitor.py b/backend/app/core/rate_monitor.py index 0ed376dc2..e0af49d1b 100644 --- a/backend/app/core/rate_monitor.py +++ b/backend/app/core/rate_monitor.py @@ -52,7 +52,7 @@ def monitor_rate(category: RateCategory): Usage: dependencies=[ Depends(require_permission(Permission.REQUIRE_PROJECT)), - Depends(monitor_rate("llm")), + Depends(monitor_rate("{category}")), ] """ From 6b4ce6c5ccb21cc8f87015406b247659bdc954f6 Mon Sep 17 00:00:00 2001 From: Prashant Vasudevan <71649489+vprashrex@users.noreply.github.com> Date: Fri, 5 Jun 2026 07:46:05 +0530 Subject: [PATCH 4/7] feat: add unit tests for rate_monitor and telemetry.record_rate_threshold --- backend/app/tests/core/test_rate_monitor.py | 210 ++++++++++++++++++++ 1 file changed, 210 insertions(+) create mode 100644 backend/app/tests/core/test_rate_monitor.py diff --git a/backend/app/tests/core/test_rate_monitor.py b/backend/app/tests/core/test_rate_monitor.py new file mode 100644 index 000000000..070ec9d13 --- /dev/null +++ b/backend/app/tests/core/test_rate_monitor.py @@ -0,0 +1,210 @@ +"""Tests for rate_monitor.py and the record_rate_threshold telemetry helper. +All Redis and Sentry calls are mocked; no real Redis or Sentry connection is used. +""" + +from types import SimpleNamespace +from unittest.mock import MagicMock, patch + +import redis + +from app.core import rate_monitor, telemetry + + +def _auth_context(org_id: int | None = 1, org_name: str = "Acme"): + """Build a minimal stand-in for AuthContext. + + monitor_rate's checker only reads auth_context.organization.id and .name, + so a SimpleNamespace is enough — no DB or real models required. + """ + org = None if org_id is None else SimpleNamespace(id=org_id, name=org_name) + return SimpleNamespace(organization=org) + + +# --------------------------------------------------------------------------- +# increment_and_get_count +# --------------------------------------------------------------------------- + + +class TestIncrementAndGetCount: + def test_returns_count_and_sets_expiry(self): + """Pipeline runs INCR + EXPIRE and returns the incremented value.""" + pipe = MagicMock() + pipe.execute.return_value = [5, True] # [incr_result, expire_result] + fake_redis = MagicMock() + fake_redis.pipeline.return_value = pipe + + with patch.object(rate_monitor, "_redis_client", fake_redis): + count = rate_monitor.increment_and_get_count("some-key") + + assert count == 5 + pipe.incr.assert_called_once_with("some-key") + pipe.expire.assert_called_once_with( + "some-key", rate_monitor._EXPIRATION_SECONDS + ) + + def test_returns_none_on_redis_error(self): + """Any Redis failure is caught and returns None rather than raising.""" + fake_redis = MagicMock() + fake_redis.pipeline.side_effect = redis.RedisError("boom") + + with patch.object(rate_monitor, "_redis_client", fake_redis): + count = rate_monitor.increment_and_get_count("some-key") + + assert count is None + + +# --------------------------------------------------------------------------- +# monitor_rate / _checker +# --------------------------------------------------------------------------- + + +class TestMonitorRate: + def test_skips_when_no_organization(self): + """No org on the request → nothing counted, no Redis call.""" + checker = rate_monitor.monitor_rate("llm_call") + + with patch.object(rate_monitor, "increment_and_get_count") as inc: + checker(_auth_context(org_id=None)) + + inc.assert_not_called() + + def test_skips_when_no_threshold_for_category(self): + """Unknown category has no threshold → return early, no alert.""" + checker = rate_monitor.monitor_rate("unknown") # type: ignore[arg-type] + + with ( + patch.object(rate_monitor, "increment_and_get_count") as inc, + patch.object(rate_monitor, "record_rate_threshold") as record, + ): + checker(_auth_context()) + + inc.assert_not_called() + record.assert_not_called() + + def test_no_alert_when_under_threshold(self): + """Count at or below the threshold does not alert.""" + checker = rate_monitor.monitor_rate("collections") + threshold = rate_monitor.THRESHOLDS["collections"] + + with ( + patch.object( + rate_monitor, "increment_and_get_count", return_value=threshold + ), + patch.object(rate_monitor, "record_rate_threshold") as record, + ): + checker(_auth_context()) + + record.assert_not_called() + + def test_alerts_when_over_threshold(self): + """Count above the threshold records a Sentry alert with org details.""" + checker = rate_monitor.monitor_rate("llm_call") + threshold = rate_monitor.THRESHOLDS["llm_call"] + over = threshold + 1 + + with ( + patch.object(rate_monitor, "increment_and_get_count", return_value=over), + patch.object(rate_monitor, "record_rate_threshold") as record, + ): + checker(_auth_context(org_id=616, org_name="Acme")) + + record.assert_called_once_with( + org_id=616, + org_name="Acme", + category="llm_call", + request_count=over, + threshold=threshold, + ) + + def test_no_alert_when_count_is_none(self): + """increment returning None (Redis down) is treated as no breach.""" + checker = rate_monitor.monitor_rate("llm_call") + + with ( + patch.object(rate_monitor, "increment_and_get_count", return_value=None), + patch.object(rate_monitor, "record_rate_threshold") as record, + ): + checker(_auth_context()) + + record.assert_not_called() + + def test_redis_error_is_swallowed(self): + """A RedisError from increment must not propagate out of the checker.""" + checker = rate_monitor.monitor_rate("llm_call") + + with patch.object( + rate_monitor, + "increment_and_get_count", + side_effect=redis.RedisError("down"), + ): + # Should not raise. + checker(_auth_context()) + + +# --------------------------------------------------------------------------- +# telemetry.record_rate_threshold +# --------------------------------------------------------------------------- + + +class TestRecordRateThreshold: + def test_emits_warning_message_with_tags(self): + """When Sentry is active, a warning message is captured with tags/extras.""" + client = MagicMock() + client.is_active.return_value = True + scope = MagicMock() + scope_cm = MagicMock() + scope_cm.__enter__.return_value = scope + + with ( + patch.object(telemetry.sentry_sdk, "get_client", return_value=client), + patch.object(telemetry.sentry_sdk, "push_scope", return_value=scope_cm), + patch.object(telemetry.sentry_sdk, "capture_message") as capture, + ): + telemetry.record_rate_threshold( + org_id=616, + org_name="Acme", + category="llm_call", + request_count=16, + threshold=15, + ) + + capture.assert_called_once() + assert capture.call_args.kwargs["level"] == "warning" + scope.set_tag.assert_any_call("alert.type", "threshold_rate_monitor") + scope.set_tag.assert_any_call("tenant.org_id", 616) + scope.set_extra.assert_any_call("request_count", 16) + scope.set_extra.assert_any_call("threshold", 15) + + def test_noop_when_sentry_inactive(self): + """No Sentry client → nothing is captured.""" + client = MagicMock() + client.is_active.return_value = False + + with ( + patch.object(telemetry.sentry_sdk, "get_client", return_value=client), + patch.object(telemetry.sentry_sdk, "capture_message") as capture, + ): + telemetry.record_rate_threshold( + org_id=1, + org_name="Acme", + category="llm_call", + request_count=16, + threshold=15, + ) + + capture.assert_not_called() + + def test_swallows_exceptions(self): + """An error inside Sentry emission must never propagate.""" + client = MagicMock() + client.is_active.side_effect = RuntimeError("sentry exploded") + + with patch.object(telemetry.sentry_sdk, "get_client", return_value=client): + # Should not raise. + telemetry.record_rate_threshold( + org_id=1, + org_name="Acme", + category="llm_call", + request_count=16, + threshold=15, + ) From 5021d0dbc73a688c7e781e00772cc63796e4d944 Mon Sep 17 00:00:00 2001 From: Prashant Vasudevan <71649489+vprashrex@users.noreply.github.com> Date: Fri, 5 Jun 2026 10:40:47 +0530 Subject: [PATCH 5/7] feat: update monitor_rate to use project context instead of organization --- backend/app/core/rate_monitor.py | 15 ++++++++------- 1 file changed, 8 insertions(+), 7 deletions(-) diff --git a/backend/app/core/rate_monitor.py b/backend/app/core/rate_monitor.py index e0af49d1b..3992803ae 100644 --- a/backend/app/core/rate_monitor.py +++ b/backend/app/core/rate_monitor.py @@ -57,8 +57,8 @@ def monitor_rate(category: RateCategory): """ def _checker(auth_context: AuthContextDep) -> None: - org = auth_context.organization - if org is None: + project = auth_context.project + if project is None: return threshold = THRESHOLDS.get(category, None) @@ -69,26 +69,27 @@ def _checker(auth_context: AuthContextDep) -> None: return minute_bucket = int(time.time() // 60) - redis_key = f"rate_monitor:{category}:{org.id}:{minute_bucket}" + redis_key = f"rate_monitor:{category}:{project.id}:{minute_bucket}" try: count = increment_and_get_count(redis_key) if count is not None and count > threshold: logger.warning( - f"[monitor_rate] Rate threshold exceeded for {category} in org {org.id}: count={count}" + f"[monitor_rate] Rate threshold exceeded for {category} in project {project.id}: count={count}" ) record_rate_threshold( - org_id=org.id, - org_name=org.name, + org_id=project.id, + org_name=project.name, category=category, request_count=count, threshold=threshold, ) + except redis.RedisError as e: logger.error( "[monitor_rate] Redis unavailable, skipping rate check " "(org_id=%s category=%s)", - org.id, + project.id, category, exc_info=e, ) From 9ffcc8338759c9b68b1339000f76f2d9f14525db Mon Sep 17 00:00:00 2001 From: Prashant Vasudevan <71649489+vprashrex@users.noreply.github.com> Date: Fri, 5 Jun 2026 10:47:42 +0530 Subject: [PATCH 6/7] feat: update telemetry and rate_monitor to use project context instead of organization --- backend/app/core/rate_monitor.py | 4 ++-- backend/app/core/telemetry.py | 8 ++++---- 2 files changed, 6 insertions(+), 6 deletions(-) diff --git a/backend/app/core/rate_monitor.py b/backend/app/core/rate_monitor.py index 3992803ae..a319e27a3 100644 --- a/backend/app/core/rate_monitor.py +++ b/backend/app/core/rate_monitor.py @@ -78,8 +78,8 @@ def _checker(auth_context: AuthContextDep) -> None: f"[monitor_rate] Rate threshold exceeded for {category} in project {project.id}: count={count}" ) record_rate_threshold( - org_id=project.id, - org_name=project.name, + project_id=project.id, + project_name=project.name, category=category, request_count=count, threshold=threshold, diff --git a/backend/app/core/telemetry.py b/backend/app/core/telemetry.py index 3fa0e84a9..c4d31f65f 100644 --- a/backend/app/core/telemetry.py +++ b/backend/app/core/telemetry.py @@ -455,8 +455,8 @@ def record_stale_pending_jobs( def record_rate_threshold( *, - org_id: int, - org_name: str | None, + project_id: int, + project_name: str | None, category: str, request_count: int, threshold: int, @@ -468,12 +468,12 @@ def record_rate_threshold( return with sentry_sdk.push_scope() as scope: scope.set_tag("alert.type", "threshold_rate_monitor") - scope.set_tag("tenant.org_id", org_id) + scope.set_tag("tenant.project_id", project_id) scope.set_tag("route_category", category) scope.set_extra("request_count", request_count) scope.set_extra("threshold", threshold) sentry_sdk.capture_message( - f"[Threshold-Monitor] {category} rate limit exceeded for org {org_id} | {org_name}: {request_count} req/min " + f"[Threshold-Monitor] {category} rate limit exceeded for project {project_id} | {project_name}: {request_count} req/min " f"(limit {threshold}/min)", level="warning", ) From abd4725c91fde06ceb8b8a2f8b87010e5daab5f4 Mon Sep 17 00:00:00 2001 From: Prashant Vasudevan <71649489+vprashrex@users.noreply.github.com> Date: Fri, 5 Jun 2026 11:07:14 +0530 Subject: [PATCH 7/7] feat: update rate_monitor and telemetry to use project context instead of organization --- backend/app/core/rate_monitor.py | 9 ++-- backend/app/core/telemetry.py | 4 +- backend/app/tests/core/test_rate_monitor.py | 55 ++++++++++++++------- 3 files changed, 43 insertions(+), 25 deletions(-) diff --git a/backend/app/core/rate_monitor.py b/backend/app/core/rate_monitor.py index a319e27a3..ff0da9f6f 100644 --- a/backend/app/core/rate_monitor.py +++ b/backend/app/core/rate_monitor.py @@ -1,6 +1,7 @@ import logging import time +from collections.abc import Callable from typing import Literal import redis @@ -12,7 +13,7 @@ logger = logging.getLogger(__name__) -# Categores of rates we want to monitor +# Categories of rates we want to monitor RateCategory = Literal["llm_call", "collections", "evaluations"] # THRESHOLD NUMBERS @@ -46,7 +47,7 @@ def increment_and_get_count(key: str) -> int | None: return None -def monitor_rate(category: RateCategory): +def monitor_rate(category: RateCategory) -> Callable[[AuthContextDep], None]: """Monitor the rate of events for the given category. If the rate exceeds the threshold, record it in telemetry. Usage: @@ -73,7 +74,7 @@ def _checker(auth_context: AuthContextDep) -> None: try: count = increment_and_get_count(redis_key) - if count is not None and count > threshold: + if count is not None and count == threshold + 1: logger.warning( f"[monitor_rate] Rate threshold exceeded for {category} in project {project.id}: count={count}" ) @@ -88,7 +89,7 @@ def _checker(auth_context: AuthContextDep) -> None: except redis.RedisError as e: logger.error( "[monitor_rate] Redis unavailable, skipping rate check " - "(org_id=%s category=%s)", + "(project_id=%s category=%s)", project.id, category, exc_info=e, diff --git a/backend/app/core/telemetry.py b/backend/app/core/telemetry.py index c4d31f65f..99d2fc959 100644 --- a/backend/app/core/telemetry.py +++ b/backend/app/core/telemetry.py @@ -478,9 +478,7 @@ def record_rate_threshold( level="warning", ) except Exception as e: - logger.exception( - "[record_rate_threshold_exceeded] Failed to emit alert", exc_info=e - ) + logger.exception("[record_rate_threshold] Failed to emit alert", exc_info=e) def flush_telemetry(timeout_millis: int = 10000) -> None: diff --git a/backend/app/tests/core/test_rate_monitor.py b/backend/app/tests/core/test_rate_monitor.py index 070ec9d13..39221039c 100644 --- a/backend/app/tests/core/test_rate_monitor.py +++ b/backend/app/tests/core/test_rate_monitor.py @@ -10,14 +10,18 @@ from app.core import rate_monitor, telemetry -def _auth_context(org_id: int | None = 1, org_name: str = "Acme"): +def _auth_context(project_id: int | None = 1, project_name: str = "Acme"): """Build a minimal stand-in for AuthContext. - monitor_rate's checker only reads auth_context.organization.id and .name, + monitor_rate's checker only reads auth_context.project.id and .name, so a SimpleNamespace is enough — no DB or real models required. """ - org = None if org_id is None else SimpleNamespace(id=org_id, name=org_name) - return SimpleNamespace(organization=org) + project = ( + None + if project_id is None + else SimpleNamespace(id=project_id, name=project_name) + ) + return SimpleNamespace(project=project) # --------------------------------------------------------------------------- @@ -59,12 +63,12 @@ def test_returns_none_on_redis_error(self): class TestMonitorRate: - def test_skips_when_no_organization(self): - """No org on the request → nothing counted, no Redis call.""" + def test_skips_when_no_project(self): + """No project on the request → nothing counted, no Redis call.""" checker = rate_monitor.monitor_rate("llm_call") with patch.object(rate_monitor, "increment_and_get_count") as inc: - checker(_auth_context(org_id=None)) + checker(_auth_context(project_id=None)) inc.assert_not_called() @@ -96,8 +100,23 @@ def test_no_alert_when_under_threshold(self): record.assert_not_called() + def test_no_alert_when_already_breached(self): + """Only the first breach (threshold + 1) alerts; later counts stay silent.""" + checker = rate_monitor.monitor_rate("collections") + threshold = rate_monitor.THRESHOLDS["collections"] + + with ( + patch.object( + rate_monitor, "increment_and_get_count", return_value=threshold + 2 + ), + patch.object(rate_monitor, "record_rate_threshold") as record, + ): + checker(_auth_context()) + + record.assert_not_called() + def test_alerts_when_over_threshold(self): - """Count above the threshold records a Sentry alert with org details.""" + """First breach (threshold + 1) records a Sentry alert with project details.""" checker = rate_monitor.monitor_rate("llm_call") threshold = rate_monitor.THRESHOLDS["llm_call"] over = threshold + 1 @@ -106,11 +125,11 @@ def test_alerts_when_over_threshold(self): patch.object(rate_monitor, "increment_and_get_count", return_value=over), patch.object(rate_monitor, "record_rate_threshold") as record, ): - checker(_auth_context(org_id=616, org_name="Acme")) + checker(_auth_context(project_id=616, project_name="Acme")) record.assert_called_once_with( - org_id=616, - org_name="Acme", + project_id=616, + project_name="Acme", category="llm_call", request_count=over, threshold=threshold, @@ -161,8 +180,8 @@ def test_emits_warning_message_with_tags(self): patch.object(telemetry.sentry_sdk, "capture_message") as capture, ): telemetry.record_rate_threshold( - org_id=616, - org_name="Acme", + project_id=616, + project_name="Acme", category="llm_call", request_count=16, threshold=15, @@ -171,7 +190,7 @@ def test_emits_warning_message_with_tags(self): capture.assert_called_once() assert capture.call_args.kwargs["level"] == "warning" scope.set_tag.assert_any_call("alert.type", "threshold_rate_monitor") - scope.set_tag.assert_any_call("tenant.org_id", 616) + scope.set_tag.assert_any_call("tenant.project_id", 616) scope.set_extra.assert_any_call("request_count", 16) scope.set_extra.assert_any_call("threshold", 15) @@ -185,8 +204,8 @@ def test_noop_when_sentry_inactive(self): patch.object(telemetry.sentry_sdk, "capture_message") as capture, ): telemetry.record_rate_threshold( - org_id=1, - org_name="Acme", + project_id=1, + project_name="Acme", category="llm_call", request_count=16, threshold=15, @@ -202,8 +221,8 @@ def test_swallows_exceptions(self): with patch.object(telemetry.sentry_sdk, "get_client", return_value=client): # Should not raise. telemetry.record_rate_threshold( - org_id=1, - org_name="Acme", + project_id=1, + project_name="Acme", category="llm_call", request_count=16, threshold=15,