From d35543a751978ce2b178a7cfede8548b90b6930f Mon Sep 17 00:00:00 2001 From: noahpodgurski Date: Thu, 10 Sep 2026 14:16:10 -0400 Subject: [PATCH 01/10] feat: add redis service and RPM traffic contract config and tracking --- pyproject.toml | 1 + scripts/app_attest_qa/app_attest_qa.py | 2 +- src/mlpa/core/config.py | 145 +++++++++++---- .../middleware/traffic_contract_enforcer.py | 75 ++++++++ src/mlpa/core/pg_services/services.py | 5 - src/mlpa/core/routers/appattest/appattest.py | 2 +- src/mlpa/core/routers/health/health.py | 2 +- src/mlpa/core/routers/user/user.py | 4 +- .../app_attest_pg_service.py | 4 +- .../litellm_pg_service.py | 5 +- .../{pg_services => services}/pg_service.py | 0 src/mlpa/core/services/redis_service.py | 169 ++++++++++++++++++ src/mlpa/core/services/services.py | 7 + src/mlpa/core/utils.py | 4 +- src/mlpa/run.py | 17 +- src/tests/conftest.py | 2 + ...est_chat_completions_budget_enforcement.py | 4 +- .../integration/test_traffic_contracts.py | 44 +++++ src/tests/unit/test_config.py | 20 +-- src/tests/unit/test_get_or_create_user.py | 2 +- src/tests/unit/test_pg_service.py | 2 +- src/tests/unit/test_pg_timeouts.py | 16 +- src/tests/unit/test_redis_service.py | 151 ++++++++++++++++ .../unit/test_traffic_contract_enforcer.py | 101 +++++++++++ uv.lock | 11 ++ 25 files changed, 724 insertions(+), 71 deletions(-) create mode 100644 src/mlpa/core/middleware/traffic_contract_enforcer.py delete mode 100644 src/mlpa/core/pg_services/services.py rename src/mlpa/core/{pg_services => services}/app_attest_pg_service.py (99%) rename src/mlpa/core/{pg_services => services}/litellm_pg_service.py (98%) rename src/mlpa/core/{pg_services => services}/pg_service.py (100%) create mode 100644 src/mlpa/core/services/redis_service.py create mode 100644 src/mlpa/core/services/services.py create mode 100644 src/tests/integration/test_traffic_contracts.py create mode 100644 src/tests/unit/test_redis_service.py create mode 100644 src/tests/unit/test_traffic_contract_enforcer.py diff --git a/pyproject.toml b/pyproject.toml index db4a902d..a8369b61 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -42,6 +42,7 @@ dependencies = [ "psycopg2-binary==2.9.11", "python-dotenv==1.1.1", "python_jose==3.4.0", + "redis==8.1.0", "sentry-sdk[fastapi]==2.42.0", "sqlalchemy==2.0.44", "starlette==1.2.0", diff --git a/scripts/app_attest_qa/app_attest_qa.py b/scripts/app_attest_qa/app_attest_qa.py index 093c2567..b848666d 100644 --- a/scripts/app_attest_qa/app_attest_qa.py +++ b/scripts/app_attest_qa/app_attest_qa.py @@ -29,7 +29,7 @@ from pyattest.testutils.factories.certificates import key_usage from mlpa.core.config import env -from mlpa.core.pg_services.services import app_attest_pg +from mlpa.core.services.services import app_attest_pg from tests.consts import MOCK_MODEL_NAME app = typer.Typer( diff --git a/src/mlpa/core/config.py b/src/mlpa/core/config.py index 30c4ac7a..3a67d3d5 100644 --- a/src/mlpa/core/config.py +++ b/src/mlpa/core/config.py @@ -1,10 +1,18 @@ from functools import cached_property -from typing import Annotated, Any +from typing import Annotated, Any, TypedDict from pydantic import field_validator from pydantic_settings import BaseSettings, NoDecode, SettingsConfigDict +class TrafficContractConfig(TypedDict): + feature: str + rpm_limit: int + tpm_limit: int + basket_rpm_limit: int + basket_tpm_limit: int + + class Env(BaseSettings): MLPA_DEBUG: bool = False APP_ATTEST_PRODUCTION: bool = False @@ -62,96 +70,119 @@ def _parse_launch_countries(cls, raw: str | set[str]) -> set[str]: # User Feature Budget - AI service type USER_FEATURE_BUDGET_AI_BUDGET_ID: str = "end-user-budget-ai" + USER_FEATURE_BUDGET_AI_BUDGET_DURATION: str = "1d" USER_FEATURE_BUDGET_AI_MAX_BUDGET: float = 0.1 USER_FEATURE_BUDGET_AI_RPM_LIMIT: int = 40 USER_FEATURE_BUDGET_AI_TPM_LIMIT: int = 2000 - USER_FEATURE_BUDGET_AI_BUDGET_DURATION: str = "1d" # User Feature Budget - S2S service type USER_FEATURE_BUDGET_S2S_BUDGET_ID: str = "end-user-budget-s2s" + USER_FEATURE_BUDGET_S2S_BUDGET_DURATION: str = "1d" USER_FEATURE_BUDGET_S2S_MAX_BUDGET: float = 0.1 USER_FEATURE_BUDGET_S2S_RPM_LIMIT: int = 40 USER_FEATURE_BUDGET_S2S_TPM_LIMIT: int = 2000 - USER_FEATURE_BUDGET_S2S_BUDGET_DURATION: str = "1d" # User Feature Budget - S2S Android service type (same values as s2s) USER_FEATURE_BUDGET_S2S_ANDROID_BUDGET_ID: str = "end-user-budget-s2s-android" + USER_FEATURE_BUDGET_S2S_ANDROID_BUDGET_DURATION: str = "1d" USER_FEATURE_BUDGET_S2S_ANDROID_MAX_BUDGET: float = 0.1 USER_FEATURE_BUDGET_S2S_ANDROID_RPM_LIMIT: int = 40 USER_FEATURE_BUDGET_S2S_ANDROID_TPM_LIMIT: int = 2000 - USER_FEATURE_BUDGET_S2S_ANDROID_BUDGET_DURATION: str = "1d" # User Feature Budget - memories service type USER_FEATURE_BUDGET_MEMORIES_BUDGET_ID: str = "end-user-budget-memories" + USER_FEATURE_BUDGET_MEMORIES_BUDGET_DURATION: str = "1d" USER_FEATURE_BUDGET_MEMORIES_MAX_BUDGET: float = 0.1 USER_FEATURE_BUDGET_MEMORIES_RPM_LIMIT: int = 10 USER_FEATURE_BUDGET_MEMORIES_TPM_LIMIT: int = 2000 - USER_FEATURE_BUDGET_MEMORIES_BUDGET_DURATION: str = "1d" USER_FEATURE_BUDGET_SEARCH_BUDGET_ID: str = "end-user-budget-search" + USER_FEATURE_BUDGET_SEARCH_BUDGET_DURATION: str = "1d" USER_FEATURE_BUDGET_SEARCH_MAX_BUDGET: float = 0.01 USER_FEATURE_BUDGET_SEARCH_RPM_LIMIT: int = 10 USER_FEATURE_BUDGET_SEARCH_TPM_LIMIT: int = 2000 - USER_FEATURE_BUDGET_SEARCH_BUDGET_DURATION: str = "1d" USER_FEATURE_BUDGET_ANSWER_BUDGET_ID: str = "end-user-budget-answer" + USER_FEATURE_BUDGET_ANSWER_BUDGET_DURATION: str = "1d" USER_FEATURE_BUDGET_ANSWER_MAX_BUDGET: float = 0.1 USER_FEATURE_BUDGET_ANSWER_RPM_LIMIT: int = 10 USER_FEATURE_BUDGET_ANSWER_TPM_LIMIT: int = 2000 - USER_FEATURE_BUDGET_ANSWER_BUDGET_DURATION: str = "1d" USER_FEATURE_BUDGET_LINER_ANSWERS_BUDGET_ID: str = "end-user-budget-liner-answer" + USER_FEATURE_BUDGET_LINER_ANSWERS_BUDGET_DURATION: str = "1d" USER_FEATURE_BUDGET_LINER_ANSWERS_MAX_BUDGET: float = 0.06 USER_FEATURE_BUDGET_LINER_ANSWERS_RPM_LIMIT: int = 10 USER_FEATURE_BUDGET_LINER_ANSWERS_TPM_LIMIT: int = 2000 - USER_FEATURE_BUDGET_LINER_ANSWERS_BUDGET_DURATION: str = "1d" USER_FEATURE_BUDGET_TELEMETRY_BUDGET_ID: str = "end-user-budget-telemetry" + USER_FEATURE_BUDGET_TELEMETRY_BUDGET_DURATION: str = "1d" USER_FEATURE_BUDGET_TELEMETRY_MAX_BUDGET: float = 0.1 USER_FEATURE_BUDGET_TELEMETRY_RPM_LIMIT: int = 10 USER_FEATURE_BUDGET_TELEMETRY_TPM_LIMIT: int = 2000 - USER_FEATURE_BUDGET_TELEMETRY_BUDGET_DURATION: str = "1d" USER_FEATURE_BUDGET_AGENT_BUDGET_ID: str = "end-user-budget-agent" + USER_FEATURE_BUDGET_AGENT_BUDGET_DURATION: str = "7d" USER_FEATURE_BUDGET_AGENT_MAX_BUDGET: float = 0.1 USER_FEATURE_BUDGET_AGENT_RPM_LIMIT: int = 10 USER_FEATURE_BUDGET_AGENT_TPM_LIMIT: int = 2000 - USER_FEATURE_BUDGET_AGENT_BUDGET_DURATION: str = "7d" USER_FEATURE_BUDGET_AGENT_SEARCH_BUDGET_ID: str = "end-user-budget-agent-search" + USER_FEATURE_BUDGET_AGENT_SEARCH_BUDGET_DURATION: str = "7d" USER_FEATURE_BUDGET_AGENT_SEARCH_MAX_BUDGET: float = 0.1 USER_FEATURE_BUDGET_AGENT_SEARCH_RPM_LIMIT: int = 10 USER_FEATURE_BUDGET_AGENT_SEARCH_TPM_LIMIT: int = 2000 - USER_FEATURE_BUDGET_AGENT_SEARCH_BUDGET_DURATION: str = "7d" # User Feature Budget - ai-dev service type (experimentation, batch predictions) USER_FEATURE_BUDGET_AI_DEV_BUDGET_ID: str = "end-user-budget-ai-dev" + USER_FEATURE_BUDGET_AI_DEV_BUDGET_DURATION: str = "1d" USER_FEATURE_BUDGET_AI_DEV_MAX_BUDGET: float = 1.0 USER_FEATURE_BUDGET_AI_DEV_RPM_LIMIT: int = 200 USER_FEATURE_BUDGET_AI_DEV_TPM_LIMIT: int = 10000 - USER_FEATURE_BUDGET_AI_DEV_BUDGET_DURATION: str = "1d" # User Feature Budget - memories-dev service type (experimentation) USER_FEATURE_BUDGET_MEMORIES_DEV_BUDGET_ID: str = "end-user-budget-memories-dev" + USER_FEATURE_BUDGET_MEMORIES_DEV_BUDGET_DURATION: str = "1d" USER_FEATURE_BUDGET_MEMORIES_DEV_MAX_BUDGET: float = 1.0 USER_FEATURE_BUDGET_MEMORIES_DEV_RPM_LIMIT: int = 50 USER_FEATURE_BUDGET_MEMORIES_DEV_TPM_LIMIT: int = 5000 - USER_FEATURE_BUDGET_MEMORIES_DEV_BUDGET_DURATION: str = "1d" USER_FEATURE_BUDGET_MOCHI_DEV_BUDGET_ID: str = "end-user-budget-mochi-dev" + USER_FEATURE_BUDGET_MOCHI_DEV_BUDGET_DURATION: str = "1d" USER_FEATURE_BUDGET_MOCHI_DEV_MAX_BUDGET: float = 1.0 USER_FEATURE_BUDGET_MOCHI_DEV_RPM_LIMIT: int = 200 USER_FEATURE_BUDGET_MOCHI_DEV_TPM_LIMIT: int = 10000 - USER_FEATURE_BUDGET_MOCHI_DEV_BUDGET_DURATION: str = "1d" USER_FEATURE_BUDGET_SEARCH_DEV_BUDGET_ID: str = "end-user-budget-search-dev" + USER_FEATURE_BUDGET_SEARCH_DEV_BUDGET_DURATION: str = "1d" USER_FEATURE_BUDGET_SEARCH_DEV_MAX_BUDGET: float = 1.0 USER_FEATURE_BUDGET_SEARCH_DEV_RPM_LIMIT: int = 200 USER_FEATURE_BUDGET_SEARCH_DEV_TPM_LIMIT: int = 10000 - USER_FEATURE_BUDGET_SEARCH_DEV_BUDGET_DURATION: str = "1d" + + # Feature Traffic Contracts + ENABLE_TRAFFIC_CONTRACT_ENFORCEMENT: bool = False + TRAFFIC_CONTRACT_REDIS_KEY_PREFIX: str = "mlpa:traffic_contract" + TRAFFIC_CONTRACT_RPM_WINDOW_SECONDS: int = 60 + TRAFFIC_CONTRACT_COUNTER_TTL_SECONDS: int = 120 + TRAFFIC_CONTRACT_FAIL_OPEN_ON_REDIS_ERROR: bool = False + + # 0 means no cap enforced; still records per-feature counters. + TOTAL_TRAFFIC_CONTRACT_RPM_LIMIT: int = 0 + TOTAL_TRAFFIC_CONTRACT_TPM_LIMIT: int = 0 + + FEATURE_SMART_WINDOW: str = "smart-window" + SMART_WINDOW_TRAFFIC_CONTRACT_RPM_LIMIT: int = 10_000 + SMART_WINDOW_TRAFFIC_CONTRACT_TPM_LIMIT: int = 1_000_000 + + FEATURE_S2S: str = "s2s" + S2S_TRAFFIC_CONTRACT_RPM_LIMIT: int = 10_000 + S2S_TRAFFIC_CONTRACT_TPM_LIMIT: int = 1_000_000 + + FEATURE_S2S_ANDROID: str = "s2s-android" + S2S_ANDROID_TRAFFIC_CONTRACT_RPM_LIMIT: int = 10_000 + S2S_ANDROID_TRAFFIC_CONTRACT_TPM_LIMIT: int = 1_000_000 @cached_property - def user_feature_budget(self) -> dict[str, dict]: + def service_type_config(self) -> dict[str, dict]: """ User feature budget configuration by service type. Returns a nested dictionary with service types (ai, s2s, s2s-android, memories, ai-dev, memories-dev, mochi-dev) as keys. @@ -159,111 +190,159 @@ def user_feature_budget(self) -> dict[str, dict]: """ return { "ai": { + "feature": self.FEATURE_SMART_WINDOW, "budget_id": self.USER_FEATURE_BUDGET_AI_BUDGET_ID, + "budget_duration": self.USER_FEATURE_BUDGET_AI_BUDGET_DURATION, "max_budget": self.USER_FEATURE_BUDGET_AI_MAX_BUDGET, "rpm_limit": self.USER_FEATURE_BUDGET_AI_RPM_LIMIT, "tpm_limit": self.USER_FEATURE_BUDGET_AI_TPM_LIMIT, - "budget_duration": self.USER_FEATURE_BUDGET_AI_BUDGET_DURATION, }, "s2s": { + "feature": self.FEATURE_S2S, "budget_id": self.USER_FEATURE_BUDGET_S2S_BUDGET_ID, + "budget_duration": self.USER_FEATURE_BUDGET_S2S_BUDGET_DURATION, "max_budget": self.USER_FEATURE_BUDGET_S2S_MAX_BUDGET, "rpm_limit": self.USER_FEATURE_BUDGET_S2S_RPM_LIMIT, "tpm_limit": self.USER_FEATURE_BUDGET_S2S_TPM_LIMIT, - "budget_duration": self.USER_FEATURE_BUDGET_S2S_BUDGET_DURATION, }, "s2s-android": { + "feature": self.FEATURE_S2S_ANDROID, "budget_id": self.USER_FEATURE_BUDGET_S2S_ANDROID_BUDGET_ID, + "budget_duration": self.USER_FEATURE_BUDGET_S2S_ANDROID_BUDGET_DURATION, "max_budget": self.USER_FEATURE_BUDGET_S2S_ANDROID_MAX_BUDGET, "rpm_limit": self.USER_FEATURE_BUDGET_S2S_ANDROID_RPM_LIMIT, "tpm_limit": self.USER_FEATURE_BUDGET_S2S_ANDROID_TPM_LIMIT, - "budget_duration": self.USER_FEATURE_BUDGET_S2S_ANDROID_BUDGET_DURATION, }, "memories": { + "feature": self.FEATURE_SMART_WINDOW, "budget_id": self.USER_FEATURE_BUDGET_MEMORIES_BUDGET_ID, + "budget_duration": self.USER_FEATURE_BUDGET_MEMORIES_BUDGET_DURATION, "max_budget": self.USER_FEATURE_BUDGET_MEMORIES_MAX_BUDGET, "rpm_limit": self.USER_FEATURE_BUDGET_MEMORIES_RPM_LIMIT, "tpm_limit": self.USER_FEATURE_BUDGET_MEMORIES_TPM_LIMIT, - "budget_duration": self.USER_FEATURE_BUDGET_MEMORIES_BUDGET_DURATION, }, "search": { + "feature": self.FEATURE_SMART_WINDOW, "budget_id": self.USER_FEATURE_BUDGET_SEARCH_BUDGET_ID, + "budget_duration": self.USER_FEATURE_BUDGET_SEARCH_BUDGET_DURATION, "max_budget": self.USER_FEATURE_BUDGET_SEARCH_MAX_BUDGET, "rpm_limit": self.USER_FEATURE_BUDGET_SEARCH_RPM_LIMIT, "tpm_limit": self.USER_FEATURE_BUDGET_SEARCH_TPM_LIMIT, - "budget_duration": self.USER_FEATURE_BUDGET_SEARCH_BUDGET_DURATION, }, "answer": { + "feature": self.FEATURE_SMART_WINDOW, "budget_id": self.USER_FEATURE_BUDGET_ANSWER_BUDGET_ID, + "budget_duration": self.USER_FEATURE_BUDGET_ANSWER_BUDGET_DURATION, "max_budget": self.USER_FEATURE_BUDGET_ANSWER_MAX_BUDGET, "rpm_limit": self.USER_FEATURE_BUDGET_ANSWER_RPM_LIMIT, "tpm_limit": self.USER_FEATURE_BUDGET_ANSWER_TPM_LIMIT, - "budget_duration": self.USER_FEATURE_BUDGET_ANSWER_BUDGET_DURATION, }, "liner-answer": { + "feature": self.FEATURE_SMART_WINDOW, "budget_id": self.USER_FEATURE_BUDGET_LINER_ANSWERS_BUDGET_ID, + "budget_duration": self.USER_FEATURE_BUDGET_LINER_ANSWERS_BUDGET_DURATION, "max_budget": self.USER_FEATURE_BUDGET_LINER_ANSWERS_MAX_BUDGET, "rpm_limit": self.USER_FEATURE_BUDGET_LINER_ANSWERS_RPM_LIMIT, "tpm_limit": self.USER_FEATURE_BUDGET_LINER_ANSWERS_TPM_LIMIT, - "budget_duration": self.USER_FEATURE_BUDGET_LINER_ANSWERS_BUDGET_DURATION, }, "telemetry": { + "feature": self.FEATURE_SMART_WINDOW, "budget_id": self.USER_FEATURE_BUDGET_TELEMETRY_BUDGET_ID, + "budget_duration": self.USER_FEATURE_BUDGET_TELEMETRY_BUDGET_DURATION, "max_budget": self.USER_FEATURE_BUDGET_TELEMETRY_MAX_BUDGET, "rpm_limit": self.USER_FEATURE_BUDGET_TELEMETRY_RPM_LIMIT, "tpm_limit": self.USER_FEATURE_BUDGET_TELEMETRY_TPM_LIMIT, - "budget_duration": self.USER_FEATURE_BUDGET_TELEMETRY_BUDGET_DURATION, }, "agent": { + "feature": self.FEATURE_SMART_WINDOW, "budget_id": self.USER_FEATURE_BUDGET_AGENT_BUDGET_ID, + "budget_duration": self.USER_FEATURE_BUDGET_AGENT_BUDGET_DURATION, "max_budget": self.USER_FEATURE_BUDGET_AGENT_MAX_BUDGET, "rpm_limit": self.USER_FEATURE_BUDGET_AGENT_RPM_LIMIT, "tpm_limit": self.USER_FEATURE_BUDGET_AGENT_TPM_LIMIT, - "budget_duration": self.USER_FEATURE_BUDGET_AGENT_BUDGET_DURATION, }, "agent-search": { + "feature": self.FEATURE_SMART_WINDOW, "budget_id": self.USER_FEATURE_BUDGET_AGENT_SEARCH_BUDGET_ID, + "budget_duration": self.USER_FEATURE_BUDGET_AGENT_SEARCH_BUDGET_DURATION, "max_budget": self.USER_FEATURE_BUDGET_AGENT_SEARCH_MAX_BUDGET, "rpm_limit": self.USER_FEATURE_BUDGET_AGENT_SEARCH_RPM_LIMIT, "tpm_limit": self.USER_FEATURE_BUDGET_AGENT_SEARCH_TPM_LIMIT, - "budget_duration": self.USER_FEATURE_BUDGET_AGENT_SEARCH_BUDGET_DURATION, }, "ai-dev": { + "feature": self.FEATURE_SMART_WINDOW, "budget_id": self.USER_FEATURE_BUDGET_AI_DEV_BUDGET_ID, + "budget_duration": self.USER_FEATURE_BUDGET_AI_DEV_BUDGET_DURATION, "max_budget": self.USER_FEATURE_BUDGET_AI_DEV_MAX_BUDGET, "rpm_limit": self.USER_FEATURE_BUDGET_AI_DEV_RPM_LIMIT, "tpm_limit": self.USER_FEATURE_BUDGET_AI_DEV_TPM_LIMIT, - "budget_duration": self.USER_FEATURE_BUDGET_AI_DEV_BUDGET_DURATION, }, "memories-dev": { + "feature": self.FEATURE_SMART_WINDOW, "budget_id": self.USER_FEATURE_BUDGET_MEMORIES_DEV_BUDGET_ID, + "budget_duration": self.USER_FEATURE_BUDGET_MEMORIES_DEV_BUDGET_DURATION, "max_budget": self.USER_FEATURE_BUDGET_MEMORIES_DEV_MAX_BUDGET, "rpm_limit": self.USER_FEATURE_BUDGET_MEMORIES_DEV_RPM_LIMIT, "tpm_limit": self.USER_FEATURE_BUDGET_MEMORIES_DEV_TPM_LIMIT, - "budget_duration": self.USER_FEATURE_BUDGET_MEMORIES_DEV_BUDGET_DURATION, }, "mochi-dev": { + "feature": self.FEATURE_SMART_WINDOW, "budget_id": self.USER_FEATURE_BUDGET_MOCHI_DEV_BUDGET_ID, + "budget_duration": self.USER_FEATURE_BUDGET_MOCHI_DEV_BUDGET_DURATION, "max_budget": self.USER_FEATURE_BUDGET_MOCHI_DEV_MAX_BUDGET, "rpm_limit": self.USER_FEATURE_BUDGET_MOCHI_DEV_RPM_LIMIT, "tpm_limit": self.USER_FEATURE_BUDGET_MOCHI_DEV_TPM_LIMIT, - "budget_duration": self.USER_FEATURE_BUDGET_MOCHI_DEV_BUDGET_DURATION, }, "search-dev": { + "feature": self.FEATURE_SMART_WINDOW, "budget_id": self.USER_FEATURE_BUDGET_SEARCH_DEV_BUDGET_ID, + "budget_duration": self.USER_FEATURE_BUDGET_SEARCH_DEV_BUDGET_DURATION, "max_budget": self.USER_FEATURE_BUDGET_SEARCH_DEV_MAX_BUDGET, "rpm_limit": self.USER_FEATURE_BUDGET_SEARCH_DEV_RPM_LIMIT, "tpm_limit": self.USER_FEATURE_BUDGET_SEARCH_DEV_TPM_LIMIT, - "budget_duration": self.USER_FEATURE_BUDGET_SEARCH_DEV_BUDGET_DURATION, }, } + @cached_property + def traffic_contract_config(self) -> dict[str, TrafficContractConfig]: + """ + Per-feature traffic contracts keyed by service type. + + Redis stores live counters by the feature name so the + same service-type contract protects each model's basket independently. + """ + + TRAFFIC_CONTRACT_CONFIG = { + "smart-window": { + "rpm_limit": self.SMART_WINDOW_TRAFFIC_CONTRACT_RPM_LIMIT, + "tpm_limit": self.SMART_WINDOW_TRAFFIC_CONTRACT_TPM_LIMIT, + }, + "s2s": { + "rpm_limit": self.S2S_TRAFFIC_CONTRACT_RPM_LIMIT, + "tpm_limit": self.S2S_TRAFFIC_CONTRACT_TPM_LIMIT, + }, + "s2s-android": { + "rpm_limit": self.S2S_ANDROID_TRAFFIC_CONTRACT_RPM_LIMIT, + "tpm_limit": self.S2S_ANDROID_TRAFFIC_CONTRACT_TPM_LIMIT, + }, + } + return { + service_type: { + "feature": budget["feature"], + "rpm_limit": TRAFFIC_CONTRACT_CONFIG[budget["feature"]]["rpm_limit"], + "tpm_limit": TRAFFIC_CONTRACT_CONFIG[budget["feature"]]["tpm_limit"], + "basket_rpm_limit": self.TOTAL_TRAFFIC_CONTRACT_RPM_LIMIT, + "basket_tpm_limit": self.TOTAL_TRAFFIC_CONTRACT_TPM_LIMIT, + } + for service_type, budget in self.service_type_config.items() + } + @cached_property def valid_service_types(self) -> list[str]: """ Returns a list of valid service types from user_feature_budget configuration. """ - return list(self.user_feature_budget.keys()) + return list(self.service_type_config.keys()) @cached_property def valid_service_types_set(self) -> set[str]: @@ -452,6 +531,10 @@ def valid_model_labels(self) -> set[str]: # PG_MAINTENANCE_STATEMENT_TIMEOUT_MS or it'll cancel those queries. PG_COMMAND_TIMEOUT_S: float | None = None + # Redis + REDIS_HOST: str = "localhost" + REDIS_PORT: int = 6379 + # LLM request default values TEMPERATURE: float = 0.1 MAX_COMPLETION_TOKENS: int = 8192 diff --git a/src/mlpa/core/middleware/traffic_contract_enforcer.py b/src/mlpa/core/middleware/traffic_contract_enforcer.py new file mode 100644 index 00000000..fcee5a5d --- /dev/null +++ b/src/mlpa/core/middleware/traffic_contract_enforcer.py @@ -0,0 +1,75 @@ +from fastapi import HTTPException + +from mlpa.core.classes import AuthorizedChatRequest, AuthorizedSearchRequest +from mlpa.core.config import TrafficContractConfig, env +from mlpa.core.logger import logger +from mlpa.core.metrics import record_chat_availability +from mlpa.core.prometheus_metrics import AvailabilityReason +from mlpa.core.services.services import redis_service + + +def _traffic_contract_for_service_type( + service_type: str, +) -> TrafficContractConfig | None: + return env.traffic_contract_config.get(service_type) + + +async def _enforce_traffic_contract(service_type: str) -> None: + if not env.ENABLE_TRAFFIC_CONTRACT_ENFORCEMENT: + return + + contract = _traffic_contract_for_service_type(service_type) + if contract is None: + return + logger.info(contract) + + try: + rpm_decision = await redis_service.check_and_increment_feature_rpm( + key_prefix=env.TRAFFIC_CONTRACT_REDIS_KEY_PREFIX, + feature=contract["feature"], + feature_rpm_limit=contract["rpm_limit"], + basket_rpm_limit=contract["basket_rpm_limit"], + window_seconds=env.TRAFFIC_CONTRACT_RPM_WINDOW_SECONDS, + ttl_seconds=env.TRAFFIC_CONTRACT_COUNTER_TTL_SECONDS, + ) + except Exception as exc: + logger.error( + "Traffic contract enforcement failed for " + f"service_type={service_type}: {exc}" + ) + if env.TRAFFIC_CONTRACT_FAIL_OPEN_ON_REDIS_ERROR: + return + raise HTTPException( + status_code=503, + detail={"error": "Traffic contract enforcement unavailable."}, + ) from exc + + if not rpm_decision.allowed: + logger.warning( + "Traffic contract exceeded for " + f"service_type={service_type}, " + f"feature={contract['feature']}, limited_by={rpm_decision.limited_by}, " + f"mode={rpm_decision.mode}, ratio_over={rpm_decision.ratio_over}, " + f"feature_count={rpm_decision.feature_count}, " + f"basket_count={rpm_decision.basket_count}" + ) + # Hard deny over limit requests: + # raise _rate_limit_response(rpm_decision.retry_after_seconds) + + +async def enforce_chat_traffic_contract(req: AuthorizedChatRequest) -> None: + try: + await _enforce_traffic_contract(req.service_type) + except HTTPException as exc: + if exc.status_code >= 500: + record_chat_availability(req, AvailabilityReason.PROVISIONING_FAILURE) + raise + + +async def enforce_search_traffic_contract( + req: AuthorizedSearchRequest, +) -> None: + try: + await _enforce_traffic_contract(req.service_type) + except HTTPException: + raise diff --git a/src/mlpa/core/pg_services/services.py b/src/mlpa/core/pg_services/services.py deleted file mode 100644 index 11caa4df..00000000 --- a/src/mlpa/core/pg_services/services.py +++ /dev/null @@ -1,5 +0,0 @@ -from mlpa.core.pg_services.app_attest_pg_service import AppAttestPGService -from mlpa.core.pg_services.litellm_pg_service import LiteLLMPGService - -litellm_pg = LiteLLMPGService() -app_attest_pg = AppAttestPGService(litellm_pg) diff --git a/src/mlpa/core/routers/appattest/appattest.py b/src/mlpa/core/routers/appattest/appattest.py index b7f63082..09d109c5 100644 --- a/src/mlpa/core/routers/appattest/appattest.py +++ b/src/mlpa/core/routers/appattest/appattest.py @@ -20,8 +20,8 @@ from mlpa.core.app_attest import QA_CERT_DIR, ensure_qa_certificates from mlpa.core.config import env from mlpa.core.logger import logger -from mlpa.core.pg_services.services import app_attest_pg from mlpa.core.prometheus_metrics import PrometheusResult, metrics +from mlpa.core.services.services import app_attest_pg from mlpa.core.utils import b64decode_safe challenge_store = {} diff --git a/src/mlpa/core/routers/health/health.py b/src/mlpa/core/routers/health/health.py index 25d9f6a9..3ba25b23 100644 --- a/src/mlpa/core/routers/health/health.py +++ b/src/mlpa/core/routers/health/health.py @@ -12,7 +12,7 @@ env, ) from mlpa.core.http_client import get_http_client -from mlpa.core.pg_services.services import app_attest_pg, litellm_pg +from mlpa.core.services.services import app_attest_pg, litellm_pg mlpa_version = importlib.metadata.version("mlpa") litellm_version = "N/A" diff --git a/src/mlpa/core/routers/user/user.py b/src/mlpa/core/routers/user/user.py index 2e80a28b..3b669c70 100644 --- a/src/mlpa/core/routers/user/user.py +++ b/src/mlpa/core/routers/user/user.py @@ -14,7 +14,7 @@ ) from mlpa.core.http_client import get_http_client from mlpa.core.logger import logger -from mlpa.core.pg_services.services import app_attest_pg, litellm_pg +from mlpa.core.services.services import app_attest_pg, litellm_pg from mlpa.core.utils import raise_and_log router = APIRouter() @@ -130,7 +130,7 @@ async def update_user_budget( f"Valid values: {', '.join(env.valid_service_types)}" }, ) - budget_id = env.user_feature_budget[payload.service_type]["budget_id"] + budget_id = env.service_type_config[payload.service_type]["budget_id"] user = await litellm_pg.update_user_budget(user_id, budget_id) return { "user_id": user["user_id"], diff --git a/src/mlpa/core/pg_services/app_attest_pg_service.py b/src/mlpa/core/services/app_attest_pg_service.py similarity index 99% rename from src/mlpa/core/pg_services/app_attest_pg_service.py rename to src/mlpa/core/services/app_attest_pg_service.py index fb91a656..29eb5c06 100644 --- a/src/mlpa/core/pg_services/app_attest_pg_service.py +++ b/src/mlpa/core/services/app_attest_pg_service.py @@ -2,8 +2,8 @@ from mlpa.core.config import env from mlpa.core.logger import logger -from mlpa.core.pg_services.litellm_pg_service import LiteLLMPGService -from mlpa.core.pg_services.pg_service import PGService +from mlpa.core.services.litellm_pg_service import LiteLLMPGService +from mlpa.core.services.pg_service import PGService class AppAttestPGService(PGService): diff --git a/src/mlpa/core/pg_services/litellm_pg_service.py b/src/mlpa/core/services/litellm_pg_service.py similarity index 98% rename from src/mlpa/core/pg_services/litellm_pg_service.py rename to src/mlpa/core/services/litellm_pg_service.py index 64b0d889..b09ce986 100644 --- a/src/mlpa/core/pg_services/litellm_pg_service.py +++ b/src/mlpa/core/services/litellm_pg_service.py @@ -2,7 +2,7 @@ from mlpa.core.config import env from mlpa.core.logger import logger -from mlpa.core.pg_services.pg_service import PGService +from mlpa.core.services.pg_service import PGService class LiteLLMPGService(PGService): @@ -197,9 +197,8 @@ async def create_budget(self): If a budget already exists, it will be overwritten with the new values. Returns a list of created/updated budget records. """ - user_feature_budgets = env.user_feature_budget - for service_type, budget_config in user_feature_budgets.items(): + for service_type, budget_config in env.service_type_config.items(): try: # Fast single-row PK upsert: a plain autocommit call won't hit # the pool statement_timeout. diff --git a/src/mlpa/core/pg_services/pg_service.py b/src/mlpa/core/services/pg_service.py similarity index 100% rename from src/mlpa/core/pg_services/pg_service.py rename to src/mlpa/core/services/pg_service.py diff --git a/src/mlpa/core/services/redis_service.py b/src/mlpa/core/services/redis_service.py new file mode 100644 index 00000000..9e3a197c --- /dev/null +++ b/src/mlpa/core/services/redis_service.py @@ -0,0 +1,169 @@ +import time +from dataclasses import dataclass +from typing import Any, Literal, cast + +import redis.asyncio as aioredis + +from mlpa.core.config import env +from mlpa.core.logger import logger + +TRAFFIC_CONTRACT_BASKET_FIELD = "__basket__" + +_CHECK_AND_INCREMENT_SCRIPT = """ +local key = KEYS[1] +local feature_field = ARGV[1] +local basket_field = ARGV[2] +local feature_limit = tonumber(ARGV[3]) +local basket_limit = tonumber(ARGV[4]) +local ttl_seconds = tonumber(ARGV[5]) +local inc_amount = tonumber(ARGV[6]) + +local feature_count = redis.call("HINCRBY", key, feature_field, inc_amount) +local basket_count = redis.call("HINCRBY", key, basket_field, inc_amount) +redis.call("EXPIRE", key, ttl_seconds) + +if basket_limit > 0 and basket_count > basket_limit then + return {0, feature_count, basket_count, "basket", "degraded", tostring(basket_count/basket_limit)} +end + +if feature_limit > 0 and feature_count > feature_limit then + return {0, feature_count, basket_count, "feature", "borrowed", tostring(feature_count/feature_limit)} +end + +return {1, feature_count, basket_count, "", "normal", ""} +""" + +TrafficContractMode = Literal["normal", "borrowed", "degraded"] + + +@dataclass(frozen=True) +class TrafficContractDecision: + allowed: bool + feature_count: int + basket_count: int + retry_after_seconds: int + limited_by: str | None = None + + # borrowed = feature is over limit, basket has room + # degraded = basket is over limit + mode: TrafficContractMode = "normal" + ratio_over: float | None = ( + None # ratio of feature_count / feature_limit if limited_by == "feature" + ) + + +class RedisService: + def __init__(self): + self.redis: Any | None = None + + async def connect(self): + self.redis = aioredis.Redis( + host=env.REDIS_HOST, + port=env.REDIS_PORT, + decode_responses=True, + ) + await self.redis.ping() + logger.info(f"Connected to Redis at {env.REDIS_HOST}:{env.REDIS_PORT}") + + async def set(self, key: str, value: str): + await self.client.set(key, value) + + async def get(self, key: str) -> str | None: + return await self.client.get(key) + + async def close(self): + if self.redis is not None: + await self.redis.aclose() + self.redis = None + + @property + def client(self) -> Any: + if self.redis is None: + raise RuntimeError("Redis client is not connected") + return self.redis + + @staticmethod + def bucket_start(now: int | None = None, window_seconds: int = 60) -> int: + current_time = int(time.time()) if now is None else now + return (current_time // window_seconds) * window_seconds + + @classmethod + def traffic_contract_key( + cls, + *, + key_prefix: str, + key_type: Literal["rpm", "tpm"], + now: int | None = None, + window_seconds: int = 60, + ) -> str: + bucket_start = cls.bucket_start(now, window_seconds) + return f"{key_prefix}:{key_type}:{bucket_start}" + + @classmethod + def retry_after_seconds( + cls, *, now: int | None = None, window_seconds: int = 60 + ) -> int: + current_time = int(time.time()) if now is None else now + elapsed_in_bucket = current_time % window_seconds + return window_seconds - elapsed_in_bucket + + async def check_and_increment_feature_rpm( + self, + *, + key_prefix: str, + feature: str, + feature_rpm_limit: int, + basket_rpm_limit: int, + window_seconds: int = 60, + ttl_seconds: int = 120, + now: int | None = None, + ) -> TrafficContractDecision: + key = self.traffic_contract_key( + key_prefix=key_prefix, + key_type="rpm", + now=now, + window_seconds=window_seconds, + ) + retry_after = self.retry_after_seconds( + now=now, + window_seconds=window_seconds, + ) + + result = await self.client.eval( + _CHECK_AND_INCREMENT_SCRIPT, + 1, + key, + feature, + TRAFFIC_CONTRACT_BASKET_FIELD, + feature_rpm_limit, + basket_rpm_limit, + ttl_seconds, + 1, + ) + + return TrafficContractDecision( + allowed=bool(int(result[0])), + feature_count=int(result[1]), + basket_count=int(result[2]), + retry_after_seconds=retry_after, + limited_by=str(result[3]) or None, + mode=cast(TrafficContractMode, result[4]), + ratio_over=float(result[5]) if result[5] not in (None, "") else None, + ) + + async def get_current_feature_rpm( + self, + *, + key_prefix: str, + feature: str, + window_seconds: int = 60, + now: int | None = None, + ) -> int: + key = self.traffic_contract_key( + key_prefix=key_prefix, + key_type="rpm", + now=now, + window_seconds=window_seconds, + ) + traffic_count = await self.client.hget(key, feature) + return int(traffic_count) if traffic_count else 0 diff --git a/src/mlpa/core/services/services.py b/src/mlpa/core/services/services.py new file mode 100644 index 00000000..2b7400e7 --- /dev/null +++ b/src/mlpa/core/services/services.py @@ -0,0 +1,7 @@ +from mlpa.core.services.app_attest_pg_service import AppAttestPGService +from mlpa.core.services.litellm_pg_service import LiteLLMPGService +from mlpa.core.services.redis_service import RedisService + +litellm_pg = LiteLLMPGService() +app_attest_pg = AppAttestPGService(litellm_pg) +redis_service = RedisService() diff --git a/src/mlpa/core/utils.py b/src/mlpa/core/utils.py index c35049f0..c1642199 100644 --- a/src/mlpa/core/utils.py +++ b/src/mlpa/core/utils.py @@ -20,8 +20,8 @@ from mlpa.core.country_codes import COUNTRY_CODES from mlpa.core.http_client import get_http_client from mlpa.core.logger import logger -from mlpa.core.pg_services.services import app_attest_pg, litellm_pg from mlpa.core.prometheus_metrics import PrometheusResult, metrics +from mlpa.core.services.services import app_attest_pg, litellm_pg KNOWN_HTTP_METHODS = frozenset( { @@ -137,7 +137,7 @@ async def get_or_create_user(user_id: str): raise HTTPException(status_code=400, detail={"error": "Invalid user_id format"}) # Get the appropriate budget_id from config based on service_type - user_feature_budgets = env.user_feature_budget + user_feature_budgets = env.service_type_config budget_id = user_feature_budgets[service_type]["budget_id"] client = get_http_client() diff --git a/src/mlpa/run.py b/src/mlpa/run.py index 39185c6a..ff862e79 100644 --- a/src/mlpa/run.py +++ b/src/mlpa/run.py @@ -33,8 +33,11 @@ record_request_country, ) from mlpa.core.middleware import register_middleware +from mlpa.core.middleware.traffic_contract_enforcer import ( + enforce_chat_traffic_contract, + enforce_search_traffic_contract, +) from mlpa.core.openapi import customize_openapi -from mlpa.core.pg_services.services import app_attest_pg, litellm_pg from mlpa.core.prometheus_metrics import AvailabilityReason from mlpa.core.routers.appattest import appattest_router from mlpa.core.routers.filter import filter_router @@ -43,6 +46,7 @@ from mlpa.core.routers.play import play_router from mlpa.core.routers.user import user_router from mlpa.core.search import get_search +from mlpa.core.services.services import app_attest_pg, litellm_pg, redis_service tags_metadata = [ {"name": "Health", "description": "Health check endpoints."}, @@ -74,6 +78,7 @@ async def lifespan(app: FastAPI): litellm_connected = False app_attest_connected = False + redis_connected = False try: get_http_client() await litellm_pg.connect() @@ -82,6 +87,10 @@ async def lifespan(app: FastAPI): await app_attest_pg.connect() app_attest_connected = True + if env.ENABLE_TRAFFIC_CONTRACT_ENFORCEMENT: + await redis_service.connect() + redis_connected = True + await litellm_pg.create_budget() await app_attest_pg.ensure_capacity_state() @@ -91,6 +100,8 @@ async def lifespan(app: FastAPI): await app_attest_pg.disconnect() if litellm_connected: await litellm_pg.disconnect() + if redis_connected: + await redis_service.close() await close_http_client() @@ -218,6 +229,7 @@ async def chat_completion( service_type=authorized_chat_request.service_type, model=authorized_chat_request.model, ) + await enforce_chat_traffic_contract(authorized_chat_request) user_id = authorized_chat_request.user if not user_id: raise HTTPException( @@ -262,6 +274,9 @@ async def search( status_code=400, detail=f"service-type header must be one of {env.forced_model_service_type_pairs.get(SEARCH_MODEL)}", ) + await enforce_search_traffic_contract( + authorized_search_request, + ) user_id = authorized_search_request.user if not user_id: raise HTTPException( diff --git a/src/tests/conftest.py b/src/tests/conftest.py index 4a5e73ad..79f8db23 100644 --- a/src/tests/conftest.py +++ b/src/tests/conftest.py @@ -17,10 +17,12 @@ def mock_request(): def _force_mlpa_debug_false(): monkeypatch = pytest.MonkeyPatch() monkeypatch.setenv("MLPA_DEBUG", "false") + monkeypatch.setenv("ENABLE_TRAFFIC_CONTRACT_ENFORCEMENT", "false") monkeypatch.setenv("ADDITIONAL_FXA_SCOPE_1", "") monkeypatch.setenv("ADDITIONAL_FXA_SCOPE_2", "") monkeypatch.setenv("ADDITIONAL_FXA_SCOPE_3", "") env.MLPA_DEBUG = False + env.ENABLE_TRAFFIC_CONTRACT_ENFORCEMENT = False env.ADDITIONAL_FXA_SCOPE_1 = "" env.ADDITIONAL_FXA_SCOPE_2 = "" env.ADDITIONAL_FXA_SCOPE_3 = "" diff --git a/src/tests/integration/test_chat_completions_budget_enforcement.py b/src/tests/integration/test_chat_completions_budget_enforcement.py index e424ba39..ec4214ef 100644 --- a/src/tests/integration/test_chat_completions_budget_enforcement.py +++ b/src/tests/integration/test_chat_completions_budget_enforcement.py @@ -186,7 +186,7 @@ def test_chat_completion_registers_end_user_with_budget_in_litellm( info = _wait_for_spend(user_id) assert ( info["litellm_budget_table"]["budget_id"] - == env.user_feature_budget["memories"]["budget_id"] + == env.service_type_config["memories"]["budget_id"] ) def test_chat_completion_enforces_rpm_budget(self, real_backend_client): @@ -197,7 +197,7 @@ def test_chat_completion_enforces_rpm_budget(self, real_backend_client): so it would have caught #243 regardless of which field was dropped. """ client, token, _base_identity = real_backend_client - rpm_limit = env.user_feature_budget["memories"]["rpm_limit"] + rpm_limit = env.service_type_config["memories"]["rpm_limit"] headers = { "authorization": f"Bearer {token}", "service-type": "memories", diff --git a/src/tests/integration/test_traffic_contracts.py b/src/tests/integration/test_traffic_contracts.py new file mode 100644 index 00000000..74927152 --- /dev/null +++ b/src/tests/integration/test_traffic_contracts.py @@ -0,0 +1,44 @@ +from mlpa.core.config import env +from mlpa.core.middleware import traffic_contract_enforcer +from mlpa.core.services.redis_service import TrafficContractDecision +from tests.consts import SAMPLE_CHAT_REQUEST, TEST_FXA_TOKEN + + +def test_chat_completion_continues_when_traffic_contract_exceeded( + mocked_client_integration, + mocker, +): + mocker.patch.object(env, "ENABLE_TRAFFIC_CONTRACT_ENFORCEMENT", True) + mocker.patch.object( + traffic_contract_enforcer.redis_service, + "check_and_increment_feature_rpm", + mocker.AsyncMock( + return_value=TrafficContractDecision( + allowed=False, + feature_count=10_000, + basket_count=10_000, + retry_after_seconds=42, + limited_by="feature", + mode="borrowed", + ratio_over=1.1, + ) + ), + ) + get_completion = mocker.patch( + "mlpa.run.get_completion", + mocker.AsyncMock(return_value={"id": "chatcmpl-test"}), + ) + + response = mocked_client_integration.post( + "/v1/chat/completions", + headers={ + "authorization": f"Bearer {TEST_FXA_TOKEN}", + "service-type": "ai", + "purpose": "chat", + }, + json=SAMPLE_CHAT_REQUEST.model_dump(exclude_unset=True), + ) + + assert response.status_code == 200 + assert response.json() == {"id": "chatcmpl-test"} + get_completion.assert_awaited_once() diff --git a/src/tests/unit/test_config.py b/src/tests/unit/test_config.py index 57b936d0..e5e168de 100644 --- a/src/tests/unit/test_config.py +++ b/src/tests/unit/test_config.py @@ -9,7 +9,7 @@ def test_user_feature_budget_includes_memories(): """Test that user_feature_budget property includes memories service type.""" env = Env() - budgets = env.user_feature_budget + budgets = env.service_type_config # Verify memories is present assert "memories" in budgets @@ -27,7 +27,7 @@ def test_user_feature_budget_includes_memories(): def test_user_feature_budget_memories_default_values(): """Test that memories budget configuration has correct default values.""" env = Env() - memories_config = env.user_feature_budget["memories"] + memories_config = env.service_type_config["memories"] assert memories_config["budget_id"] == "end-user-budget-memories" assert memories_config["max_budget"] == 0.1 @@ -48,7 +48,7 @@ def test_user_feature_budget_memories_from_env(): with patch.dict(os.environ, env_vars): env = Env() - memories_config = env.user_feature_budget["memories"] + memories_config = env.service_type_config["memories"] assert memories_config["budget_id"] == "custom-memories-budget-id" assert memories_config["max_budget"] == 0.5 @@ -60,7 +60,7 @@ def test_user_feature_budget_memories_from_env(): def test_user_feature_budget_liner_answer_default_values(): """Test that liner-answer budget configuration has correct default values.""" env = Env() - liner_config = env.user_feature_budget["liner-answer"] + liner_config = env.service_type_config["liner-answer"] assert liner_config["budget_id"] == "end-user-budget-liner-answer" assert liner_config["max_budget"] == 0.06 @@ -82,7 +82,7 @@ def test_user_feature_budget_liner_answer_from_env(): with patch.dict(os.environ, env_vars): env = Env() - liner_config = env.user_feature_budget["liner-answer"] + liner_config = env.service_type_config["liner-answer"] assert liner_config["budget_id"] == "custom-liner-budget-id" assert liner_config["max_budget"] == 0.5 assert liner_config["rpm_limit"] == 20 @@ -110,9 +110,9 @@ def test_valid_major_fx_versions_set_uses_string_range(): def test_user_feature_budget_dev_service_types_default_values(): """Test that ai-dev, memories-dev, and mochi-dev have correct default values.""" env = Env() - ai_dev_config = env.user_feature_budget["ai-dev"] - memories_dev_config = env.user_feature_budget["memories-dev"] - mochi_dev_config = env.user_feature_budget["mochi-dev"] + ai_dev_config = env.service_type_config["ai-dev"] + memories_dev_config = env.service_type_config["memories-dev"] + mochi_dev_config = env.service_type_config["mochi-dev"] assert ai_dev_config["budget_id"] == "end-user-budget-ai-dev" assert ai_dev_config["max_budget"] == 1.0 @@ -254,7 +254,7 @@ def test_valid_model_labels_are_explicit_metric_allowlist(): def test_user_feature_budget_structure_consistency(): """Test that all service types have the same structure in user_feature_budget.""" env = Env() - budgets = env.user_feature_budget + budgets = env.service_type_config # Get the keys from one service type as reference reference_keys = set(budgets["ai"].keys()) @@ -284,7 +284,7 @@ def test_user_feature_budget_structure_consistency(): def test_user_feature_budget_memories_type_validation(): """Test that memories budget configuration values have correct types.""" env = Env() - memories_config = env.user_feature_budget["memories"] + memories_config = env.service_type_config["memories"] assert isinstance(memories_config["budget_id"], str) assert isinstance(memories_config["max_budget"], float) diff --git a/src/tests/unit/test_get_or_create_user.py b/src/tests/unit/test_get_or_create_user.py index 98a1bf76..7fa197bb 100644 --- a/src/tests/unit/test_get_or_create_user.py +++ b/src/tests/unit/test_get_or_create_user.py @@ -8,7 +8,7 @@ _USER_ID = "user123:ai" _BASE_IDENTITY, _, _ = _USER_ID.partition(":") -_BUDGET_ID = env.user_feature_budget["ai"]["budget_id"] +_BUDGET_ID = env.service_type_config["ai"]["budget_id"] _DB_USER = {"user_id": _USER_ID, "blocked": False, "budget_id": _BUDGET_ID} diff --git a/src/tests/unit/test_pg_service.py b/src/tests/unit/test_pg_service.py index 0010caa4..e927255e 100644 --- a/src/tests/unit/test_pg_service.py +++ b/src/tests/unit/test_pg_service.py @@ -3,7 +3,7 @@ import pytest -from mlpa.core.pg_services.pg_service import PGService +from mlpa.core.services.pg_service import PGService def _make_service(pool, connected=True): diff --git a/src/tests/unit/test_pg_timeouts.py b/src/tests/unit/test_pg_timeouts.py index a7d08239..200aa0a6 100644 --- a/src/tests/unit/test_pg_timeouts.py +++ b/src/tests/unit/test_pg_timeouts.py @@ -2,8 +2,8 @@ from unittest.mock import AsyncMock, MagicMock, patch from mlpa.core.config import Env, env -from mlpa.core.pg_services.app_attest_pg_service import AppAttestPGService -from mlpa.core.pg_services.pg_service import PGService +from mlpa.core.services.app_attest_pg_service import AppAttestPGService +from mlpa.core.services.pg_service import PGService def test_pg_timeout_config_from_env(): @@ -38,7 +38,7 @@ def test_pg_timeout_defaults(): async def test_connect_passes_timeout_server_settings(mocker): """The pool is created with server-enforced statement / idle-in-tx timeouts.""" create_pool = mocker.patch( - "mlpa.core.pg_services.pg_service.asyncpg.create_pool", + "mlpa.core.services.pg_service.asyncpg.create_pool", new=AsyncMock(return_value=MagicMock()), ) @@ -58,7 +58,7 @@ async def test_connect_passes_timeout_server_settings(mocker): async def test_connect_respects_per_pool_statement_timeout_override(mocker): """A subclass/per-pool override flows into server_settings.""" create_pool = mocker.patch( - "mlpa.core.pg_services.pg_service.asyncpg.create_pool", + "mlpa.core.services.pg_service.asyncpg.create_pool", new=AsyncMock(return_value=MagicMock()), ) @@ -140,7 +140,7 @@ async def test_statement_timeout_lifts_statement_and_idle_in_tx(mocker): async def test_count_users_by_service_type_uses_admin_read_timeout(mocker): """The unindexable full-table GROUP BY runs under the admin-read timeout, not 3s.""" - from mlpa.core.pg_services.litellm_pg_service import LiteLLMPGService + from mlpa.core.services.litellm_pg_service import LiteLLMPGService conn = _mock_maintenance_conn() mocker.patch.object(PGService, "pool", new=_mock_pool(conn)) @@ -156,7 +156,7 @@ async def test_count_users_by_service_type_uses_admin_read_timeout(mocker): async def test_list_users_uses_admin_read_timeout(mocker): """The full-table COUNT(*) + deep OFFSET page run under the admin-read timeout.""" - from mlpa.core.pg_services.litellm_pg_service import LiteLLMPGService + from mlpa.core.services.litellm_pg_service import LiteLLMPGService conn = _mock_maintenance_conn() mocker.patch.object(PGService, "pool", new=_mock_pool(conn)) @@ -171,7 +171,7 @@ async def test_list_users_uses_admin_read_timeout(mocker): async def test_list_managed_base_identities_uses_maintenance_timeout(mocker): """The heavy reconciliation read runs under the maintenance timeout, not the 3s default.""" - from mlpa.core.pg_services.litellm_pg_service import LiteLLMPGService + from mlpa.core.services.litellm_pg_service import LiteLLMPGService conn = _mock_maintenance_conn() mocker.patch.object(PGService, "pool", new=_mock_pool(conn)) @@ -277,7 +277,7 @@ async def test_ensure_capacity_state_reconcile_failure_is_best_effort(mocker): conn = _mock_maintenance_conn() mocker.patch.object(PGService, "pool", new=_mock_pool(conn)) - log_error = mocker.patch("mlpa.core.pg_services.app_attest_pg_service.logger.error") + log_error = mocker.patch("mlpa.core.services.app_attest_pg_service.logger.error") # Must not raise despite reconciliation failing. await service.ensure_capacity_state() diff --git a/src/tests/unit/test_redis_service.py b/src/tests/unit/test_redis_service.py new file mode 100644 index 00000000..aaff0aa9 --- /dev/null +++ b/src/tests/unit/test_redis_service.py @@ -0,0 +1,151 @@ +from mlpa.core.services.redis_service import ( + TRAFFIC_CONTRACT_BASKET_FIELD, + RedisService, +) + + +class FakeRedis: + def __init__(self): + self.hashes = {} + self.expirations = {} + + async def eval( + self, + script, + num_keys, + key, + feature, + basket_field, + feature_limit, + basket_limit, + ttl_seconds, + inc_amount, + ): + assert num_keys == 1 + assert "HINCRBY" in script + + bucket = self.hashes.setdefault(key, {}) + feature_count = int(bucket.get(feature, 0)) + int(inc_amount) + basket_count = int(bucket.get(basket_field, 0)) + int(inc_amount) + feature_limit = int(feature_limit) + basket_limit = int(basket_limit) + + bucket[feature] = feature_count + bucket[basket_field] = basket_count + self.expirations[key] = int(ttl_seconds) + + if basket_limit > 0 and basket_count > basket_limit: + return [ + 0, + feature_count, + basket_count, + "basket", + "degraded", + str(basket_count / basket_limit), + ] + + if feature_limit > 0 and feature_count > feature_limit: + return [ + 0, + feature_count, + basket_count, + "feature", + "borrowed", + str(feature_count / feature_limit), + ] + + return [1, bucket[feature], bucket[basket_field], "", "normal", ""] + + async def hget(self, key, field): + return self.hashes.get(key, {}).get(field) + + +async def test_check_and_increment_feature_rpm_records_minute_bucket(): + redis = FakeRedis() + service = RedisService() + service.redis = redis + + decision = await service.check_and_increment_feature_rpm( + key_prefix="mlpa:traffic_contract", + feature="smart-window", + feature_rpm_limit=10, + basket_rpm_limit=20, + ttl_seconds=120, + now=125, + ) + + key = "mlpa:traffic_contract:rpm:120" + assert decision.allowed is True + assert decision.feature_count == 1 + assert decision.basket_count == 1 + assert decision.retry_after_seconds == 55 + assert decision.limited_by is None + assert decision.mode == "normal" + assert decision.ratio_over is None + assert redis.hashes[key]["smart-window"] == 1 + assert redis.hashes[key][TRAFFIC_CONTRACT_BASKET_FIELD] == 1 + assert redis.expirations[key] == 120 + assert ( + await service.get_current_feature_rpm( + key_prefix="mlpa:traffic_contract", + feature="smart-window", + now=125, + ) + == 1 + ) + + +async def test_check_and_increment_feature_rpm_marks_borrowed_over_feature_limit(): + service = RedisService() + service.redis = FakeRedis() + + allowed = await service.check_and_increment_feature_rpm( + key_prefix="mlpa:traffic_contract", + feature="smart-window", + feature_rpm_limit=1, + basket_rpm_limit=10, + now=125, + ) + rejected = await service.check_and_increment_feature_rpm( + key_prefix="mlpa:traffic_contract", + feature="smart-window", + feature_rpm_limit=1, + basket_rpm_limit=10, + now=125, + ) + + assert allowed.allowed is True + assert rejected.allowed is False + assert rejected.limited_by == "feature" + assert rejected.mode == "borrowed" + assert rejected.ratio_over == 2.0 + assert rejected.feature_count == 2 + assert rejected.basket_count == 2 + + +async def test_check_and_increment_feature_rpm_marks_degraded_over_basket_limit(): + service = RedisService() + service.redis = FakeRedis() + + allowed = await service.check_and_increment_feature_rpm( + key_prefix="mlpa:traffic_contract", + feature="smart-window", + feature_rpm_limit=10, + basket_rpm_limit=1, + now=125, + ) + rejected = await service.check_and_increment_feature_rpm( + key_prefix="mlpa:traffic_contract", + feature="memories", + feature_rpm_limit=10, + basket_rpm_limit=1, + now=125, + ) + + assert allowed.allowed is True + assert rejected.allowed is False + assert rejected.limited_by == "basket" + assert rejected.mode == "degraded" + assert rejected.ratio_over == 2.0 + assert rejected.feature_count == 1 + assert rejected.basket_count == 2 diff --git a/src/tests/unit/test_traffic_contract_enforcer.py b/src/tests/unit/test_traffic_contract_enforcer.py new file mode 100644 index 00000000..034a20d6 --- /dev/null +++ b/src/tests/unit/test_traffic_contract_enforcer.py @@ -0,0 +1,101 @@ +import pytest +from fastapi import HTTPException + +from mlpa.core.config import env +from mlpa.core.middleware import traffic_contract_enforcer +from mlpa.core.services.redis_service import TrafficContractDecision +from tests.consts import SAMPLE_REQUEST + + +async def test_chat_traffic_contract_noops_when_disabled(mocker): + mocker.patch.object(env, "ENABLE_TRAFFIC_CONTRACT_ENFORCEMENT", False) + check = mocker.patch.object( + traffic_contract_enforcer.redis_service, + "check_and_increment_feature_rpm", + mocker.AsyncMock(), + ) + + await traffic_contract_enforcer.enforce_chat_traffic_contract(SAMPLE_REQUEST) + + check.assert_not_awaited() + + +async def test_chat_traffic_contract_increments_feature_rpm_when_enabled(mocker): + mocker.patch.object(env, "ENABLE_TRAFFIC_CONTRACT_ENFORCEMENT", True) + check = mocker.patch.object( + traffic_contract_enforcer.redis_service, + "check_and_increment_feature_rpm", + mocker.AsyncMock( + return_value=TrafficContractDecision( + allowed=True, + feature_count=1, + basket_count=1, + retry_after_seconds=60, + ) + ), + ) + + await traffic_contract_enforcer.enforce_chat_traffic_contract(SAMPLE_REQUEST) + + contract = env.traffic_contract_config[SAMPLE_REQUEST.service_type] + check.assert_awaited_once_with( + key_prefix=env.TRAFFIC_CONTRACT_REDIS_KEY_PREFIX, + feature=contract["feature"], + feature_rpm_limit=contract["rpm_limit"], + basket_rpm_limit=contract["basket_rpm_limit"], + window_seconds=env.TRAFFIC_CONTRACT_RPM_WINDOW_SECONDS, + ttl_seconds=env.TRAFFIC_CONTRACT_COUNTER_TTL_SECONDS, + ) + + +async def test_chat_traffic_contract_allows_soft_over_rpm(mocker): + mocker.patch.object(env, "ENABLE_TRAFFIC_CONTRACT_ENFORCEMENT", True) + check = mocker.patch.object( + traffic_contract_enforcer.redis_service, + "check_and_increment_feature_rpm", + mocker.AsyncMock( + return_value=TrafficContractDecision( + allowed=False, + feature_count=10_000, + basket_count=10_000, + retry_after_seconds=17, + limited_by="feature", + mode="borrowed", + ratio_over=1.1, + ) + ), + ) + + await traffic_contract_enforcer.enforce_chat_traffic_contract(SAMPLE_REQUEST) + + check.assert_awaited_once() + + +async def test_chat_traffic_contract_can_fail_open_on_redis_error(mocker): + mocker.patch.object(env, "ENABLE_TRAFFIC_CONTRACT_ENFORCEMENT", True) + mocker.patch.object(env, "TRAFFIC_CONTRACT_FAIL_OPEN_ON_REDIS_ERROR", True) + mocker.patch.object( + traffic_contract_enforcer.redis_service, + "check_and_increment_feature_rpm", + mocker.AsyncMock(side_effect=RuntimeError("redis down")), + ) + + await traffic_contract_enforcer.enforce_chat_traffic_contract(SAMPLE_REQUEST) + + +async def test_chat_traffic_contract_fails_closed_on_redis_error(mocker): + mocker.patch.object(env, "ENABLE_TRAFFIC_CONTRACT_ENFORCEMENT", True) + mocker.patch.object(env, "TRAFFIC_CONTRACT_FAIL_OPEN_ON_REDIS_ERROR", False) + mocker.patch.object( + traffic_contract_enforcer.redis_service, + "check_and_increment_feature_rpm", + mocker.AsyncMock(side_effect=RuntimeError("redis down")), + ) + + with pytest.raises(HTTPException) as exc_info: + await traffic_contract_enforcer.enforce_chat_traffic_contract(SAMPLE_REQUEST) + + assert exc_info.value.status_code == 503 + assert exc_info.value.detail == { + "error": "Traffic contract enforcement unavailable." + } diff --git a/uv.lock b/uv.lock index 51fb5153..e21fb7ed 100644 --- a/uv.lock +++ b/uv.lock @@ -755,6 +755,7 @@ dependencies = [ { name = "pyjwt" }, { name = "python-dotenv" }, { name = "python-jose" }, + { name = "redis" }, { name = "sentry-sdk", extra = ["fastapi"] }, { name = "sqlalchemy" }, { name = "starlette" }, @@ -807,6 +808,7 @@ requires-dist = [ { name = "pyjwt", specifier = "==2.10.1" }, { name = "python-dotenv", specifier = "==1.1.1" }, { name = "python-jose", specifier = "==3.4.0" }, + { name = "redis", specifier = "==8.1.0" }, { name = "sentry-sdk", extras = ["fastapi"], specifier = "==2.42.0" }, { name = "sqlalchemy", specifier = "==2.0.44" }, { name = "starlette", specifier = "==1.2.0" }, @@ -1317,6 +1319,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/dd/b8/d2d6d731733f51684bbf76bf34dab3b70a9148e8f2cef2bb544fccec681a/qrcode-8.2-py3-none-any.whl", hash = "sha256:16e64e0716c14960108e85d853062c9e8bba5ca8252c0b4d0231b9df4060ff4f", size = 45986, upload-time = "2025-05-01T15:44:22.781Z" }, ] +[[package]] +name = "redis" +version = "8.1.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/a8/99/604f0b666d4c616d891cf77ebb9db6bb21601344c051aebf1b72b9ff915f/redis-8.1.0.tar.gz", hash = "sha256:6e1a19beef9225c83efd689c7e6b7da2d5215b1f42cd13b7fc3714d0a09c7b25", size = 5254356, upload-time = "2026-07-30T08:51:00.269Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/66/9d/c5731f6e3608663d4d3656fd8d3aecee8b509c3082818f5a13eae925baea/redis-8.1.0-py3-none-any.whl", hash = "sha256:a4fe1aac3d3b3cc791d4b3d5931c5a956045dc951ee74d1c913ee3ac4d2ee9fb", size = 560618, upload-time = "2026-07-30T08:50:58.497Z" }, +] + [[package]] name = "regex" version = "2026.1.15" From 030830d1fc0167050e3a7315c814f7a966c896fc Mon Sep 17 00:00:00 2001 From: noahpodgurski Date: Thu, 10 Sep 2026 14:31:50 -0400 Subject: [PATCH 02/10] don't pass total bucket through params --- src/mlpa/core/config.py | 4 ---- .../middleware/traffic_contract_enforcer.py | 1 - src/mlpa/core/services/redis_service.py | 3 +-- src/tests/unit/test_redis_service.py | 17 +++++++++-------- .../unit/test_traffic_contract_enforcer.py | 1 - 5 files changed, 10 insertions(+), 16 deletions(-) diff --git a/src/mlpa/core/config.py b/src/mlpa/core/config.py index 3a67d3d5..c582060b 100644 --- a/src/mlpa/core/config.py +++ b/src/mlpa/core/config.py @@ -9,8 +9,6 @@ class TrafficContractConfig(TypedDict): feature: str rpm_limit: int tpm_limit: int - basket_rpm_limit: int - basket_tpm_limit: int class Env(BaseSettings): @@ -331,8 +329,6 @@ def traffic_contract_config(self) -> dict[str, TrafficContractConfig]: "feature": budget["feature"], "rpm_limit": TRAFFIC_CONTRACT_CONFIG[budget["feature"]]["rpm_limit"], "tpm_limit": TRAFFIC_CONTRACT_CONFIG[budget["feature"]]["tpm_limit"], - "basket_rpm_limit": self.TOTAL_TRAFFIC_CONTRACT_RPM_LIMIT, - "basket_tpm_limit": self.TOTAL_TRAFFIC_CONTRACT_TPM_LIMIT, } for service_type, budget in self.service_type_config.items() } diff --git a/src/mlpa/core/middleware/traffic_contract_enforcer.py b/src/mlpa/core/middleware/traffic_contract_enforcer.py index fcee5a5d..d9b2f206 100644 --- a/src/mlpa/core/middleware/traffic_contract_enforcer.py +++ b/src/mlpa/core/middleware/traffic_contract_enforcer.py @@ -28,7 +28,6 @@ async def _enforce_traffic_contract(service_type: str) -> None: key_prefix=env.TRAFFIC_CONTRACT_REDIS_KEY_PREFIX, feature=contract["feature"], feature_rpm_limit=contract["rpm_limit"], - basket_rpm_limit=contract["basket_rpm_limit"], window_seconds=env.TRAFFIC_CONTRACT_RPM_WINDOW_SECONDS, ttl_seconds=env.TRAFFIC_CONTRACT_COUNTER_TTL_SECONDS, ) diff --git a/src/mlpa/core/services/redis_service.py b/src/mlpa/core/services/redis_service.py index 9e3a197c..b2c0f3bd 100644 --- a/src/mlpa/core/services/redis_service.py +++ b/src/mlpa/core/services/redis_service.py @@ -113,7 +113,6 @@ async def check_and_increment_feature_rpm( key_prefix: str, feature: str, feature_rpm_limit: int, - basket_rpm_limit: int, window_seconds: int = 60, ttl_seconds: int = 120, now: int | None = None, @@ -136,7 +135,7 @@ async def check_and_increment_feature_rpm( feature, TRAFFIC_CONTRACT_BASKET_FIELD, feature_rpm_limit, - basket_rpm_limit, + env.TOTAL_TRAFFIC_CONTRACT_RPM_LIMIT, ttl_seconds, 1, ) diff --git a/src/tests/unit/test_redis_service.py b/src/tests/unit/test_redis_service.py index aaff0aa9..2d96e90c 100644 --- a/src/tests/unit/test_redis_service.py +++ b/src/tests/unit/test_redis_service.py @@ -1,3 +1,4 @@ +from mlpa.core.config import env from mlpa.core.services.redis_service import ( TRAFFIC_CONTRACT_BASKET_FIELD, RedisService, @@ -60,7 +61,8 @@ async def hget(self, key, field): return self.hashes.get(key, {}).get(field) -async def test_check_and_increment_feature_rpm_records_minute_bucket(): +async def test_check_and_increment_feature_rpm_records_minute_bucket(mocker): + mocker.patch.object(env, "TOTAL_TRAFFIC_CONTRACT_RPM_LIMIT", 20) redis = FakeRedis() service = RedisService() service.redis = redis @@ -69,7 +71,6 @@ async def test_check_and_increment_feature_rpm_records_minute_bucket(): key_prefix="mlpa:traffic_contract", feature="smart-window", feature_rpm_limit=10, - basket_rpm_limit=20, ttl_seconds=120, now=125, ) @@ -95,7 +96,10 @@ async def test_check_and_increment_feature_rpm_records_minute_bucket(): ) -async def test_check_and_increment_feature_rpm_marks_borrowed_over_feature_limit(): +async def test_check_and_increment_feature_rpm_marks_borrowed_over_feature_limit( + mocker, +): + mocker.patch.object(env, "TOTAL_TRAFFIC_CONTRACT_RPM_LIMIT", 10) service = RedisService() service.redis = FakeRedis() @@ -103,14 +107,12 @@ async def test_check_and_increment_feature_rpm_marks_borrowed_over_feature_limit key_prefix="mlpa:traffic_contract", feature="smart-window", feature_rpm_limit=1, - basket_rpm_limit=10, now=125, ) rejected = await service.check_and_increment_feature_rpm( key_prefix="mlpa:traffic_contract", feature="smart-window", feature_rpm_limit=1, - basket_rpm_limit=10, now=125, ) @@ -123,7 +125,8 @@ async def test_check_and_increment_feature_rpm_marks_borrowed_over_feature_limit assert rejected.basket_count == 2 -async def test_check_and_increment_feature_rpm_marks_degraded_over_basket_limit(): +async def test_check_and_increment_feature_rpm_marks_degraded_over_basket_limit(mocker): + mocker.patch.object(env, "TOTAL_TRAFFIC_CONTRACT_RPM_LIMIT", 1) service = RedisService() service.redis = FakeRedis() @@ -131,14 +134,12 @@ async def test_check_and_increment_feature_rpm_marks_degraded_over_basket_limit( key_prefix="mlpa:traffic_contract", feature="smart-window", feature_rpm_limit=10, - basket_rpm_limit=1, now=125, ) rejected = await service.check_and_increment_feature_rpm( key_prefix="mlpa:traffic_contract", feature="memories", feature_rpm_limit=10, - basket_rpm_limit=1, now=125, ) diff --git a/src/tests/unit/test_traffic_contract_enforcer.py b/src/tests/unit/test_traffic_contract_enforcer.py index 034a20d6..3dd81e76 100644 --- a/src/tests/unit/test_traffic_contract_enforcer.py +++ b/src/tests/unit/test_traffic_contract_enforcer.py @@ -42,7 +42,6 @@ async def test_chat_traffic_contract_increments_feature_rpm_when_enabled(mocker) key_prefix=env.TRAFFIC_CONTRACT_REDIS_KEY_PREFIX, feature=contract["feature"], feature_rpm_limit=contract["rpm_limit"], - basket_rpm_limit=contract["basket_rpm_limit"], window_seconds=env.TRAFFIC_CONTRACT_RPM_WINDOW_SECONDS, ttl_seconds=env.TRAFFIC_CONTRACT_COUNTER_TTL_SECONDS, ) From 1bc16a132c178cf4c5468392ab8bda4539ffac18 Mon Sep 17 00:00:00 2001 From: noahpodgurski Date: Thu, 10 Sep 2026 14:48:24 -0400 Subject: [PATCH 03/10] rm chat_availability failure --- src/mlpa/core/middleware/traffic_contract_enforcer.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/src/mlpa/core/middleware/traffic_contract_enforcer.py b/src/mlpa/core/middleware/traffic_contract_enforcer.py index d9b2f206..356a2771 100644 --- a/src/mlpa/core/middleware/traffic_contract_enforcer.py +++ b/src/mlpa/core/middleware/traffic_contract_enforcer.py @@ -60,8 +60,6 @@ async def enforce_chat_traffic_contract(req: AuthorizedChatRequest) -> None: try: await _enforce_traffic_contract(req.service_type) except HTTPException as exc: - if exc.status_code >= 500: - record_chat_availability(req, AvailabilityReason.PROVISIONING_FAILURE) raise From d25e524d24f8b2d4b266f8673373ebc746754bf7 Mon Sep 17 00:00:00 2001 From: noahpodgurski Date: Thu, 10 Sep 2026 15:10:22 -0400 Subject: [PATCH 04/10] improve comment --- src/mlpa/core/services/redis_service.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/src/mlpa/core/services/redis_service.py b/src/mlpa/core/services/redis_service.py index b2c0f3bd..c4a4f44d 100644 --- a/src/mlpa/core/services/redis_service.py +++ b/src/mlpa/core/services/redis_service.py @@ -47,9 +47,7 @@ class TrafficContractDecision: # borrowed = feature is over limit, basket has room # degraded = basket is over limit mode: TrafficContractMode = "normal" - ratio_over: float | None = ( - None # ratio of feature_count / feature_limit if limited_by == "feature" - ) + ratio_over: float | None = None # ratio of count / limit if mode != "normal" class RedisService: From 90537e13d0cb395f1bcd6be7014264c0ff24f5cd Mon Sep 17 00:00:00 2001 From: noahpodgurski Date: Thu, 10 Sep 2026 15:31:35 -0400 Subject: [PATCH 05/10] rm old log --- src/mlpa/core/middleware/traffic_contract_enforcer.py | 1 - 1 file changed, 1 deletion(-) diff --git a/src/mlpa/core/middleware/traffic_contract_enforcer.py b/src/mlpa/core/middleware/traffic_contract_enforcer.py index 356a2771..96065e60 100644 --- a/src/mlpa/core/middleware/traffic_contract_enforcer.py +++ b/src/mlpa/core/middleware/traffic_contract_enforcer.py @@ -21,7 +21,6 @@ async def _enforce_traffic_contract(service_type: str) -> None: contract = _traffic_contract_for_service_type(service_type) if contract is None: return - logger.info(contract) try: rpm_decision = await redis_service.check_and_increment_feature_rpm( From 8490c33cdcc4005406e60b6a5ed04256992449e3 Mon Sep 17 00:00:00 2001 From: noahpodgurski Date: Tue, 15 Sep 2026 14:15:30 -0400 Subject: [PATCH 06/10] refactor contracts, add TPM checking and incrementing --- src/mlpa/core/classes.py | 21 ++ src/mlpa/core/completions.py | 10 + src/mlpa/core/config.py | 17 +- src/mlpa/core/consts/__init__.py | 13 + src/mlpa/core/{ => consts}/country_codes.py | 0 src/mlpa/core/{ => consts}/openapi.py | 73 ++++++ src/mlpa/core/consts/traffic_contract.py | 12 + src/mlpa/core/middleware/instrumentation.py | 23 ++ .../middleware/traffic_contract_enforcer.py | 96 ++++--- src/mlpa/core/prometheus_metrics.py | 10 +- src/mlpa/core/services/redis_service.py | 226 ++++++++++++++--- src/mlpa/core/utils.py | 2 +- src/mlpa/run.py | 96 +------ .../integration/test_traffic_contracts.py | 7 +- src/tests/unit/test_middleware_order.py | 53 ++++ src/tests/unit/test_redis_service.py | 236 ++++++++++++++---- .../unit/test_traffic_contract_enforcer.py | 198 ++++++++++++--- 17 files changed, 836 insertions(+), 257 deletions(-) create mode 100644 src/mlpa/core/consts/__init__.py rename src/mlpa/core/{ => consts}/country_codes.py (100%) rename src/mlpa/core/{ => consts}/openapi.py (60%) create mode 100644 src/mlpa/core/consts/traffic_contract.py diff --git a/src/mlpa/core/classes.py b/src/mlpa/core/classes.py index 320a1edc..72507370 100644 --- a/src/mlpa/core/classes.py +++ b/src/mlpa/core/classes.py @@ -5,6 +5,7 @@ from pydantic import BaseModel, Field from mlpa.core.config import env +from mlpa.core.consts import TrafficContractMode class ChatRequest(BaseModel): @@ -159,3 +160,23 @@ class LitellmRoutingSnapshot: attempted_retries: int response_duration_ms: float | None response_cost_usd: float | None + + +@dataclass(frozen=True) +class TrafficContractDecision: + allowed: bool + feature_count: int + basket_count: int + retry_after_seconds: int + limited_by: str | None = None + + # borrowed = feature is over limit, basket has room + # degraded = basket is over limit + mode: TrafficContractMode = TrafficContractMode.NORMAL + ratio_over: float | None = None # ratio of count / limit if mode != "normal" + + +@dataclass(frozen=True) +class TrafficContractCounters: + feature_count: int + basket_count: int diff --git a/src/mlpa/core/completions.py b/src/mlpa/core/completions.py index 7b37c1e6..9ed8ff19 100644 --- a/src/mlpa/core/completions.py +++ b/src/mlpa/core/completions.py @@ -14,6 +14,7 @@ LITELLM_VIRTUAL_AUTH_HEADERS, env, ) +from mlpa.core.consts import TrafficContractKeyType from mlpa.core.errors import classify_upstream_error from mlpa.core.http_client import get_http_client from mlpa.core.litellm_routing import parse_litellm_routing_headers @@ -33,6 +34,7 @@ PrometheusResult, ) from mlpa.core.sanitization import sanitize_request_body, sanitize_response_body +from mlpa.core.services.services import redis_service from mlpa.core.utils import ( get_or_create_user, raise_and_log, @@ -135,6 +137,7 @@ async def _read_next_chunk( watch_task = asyncio.create_task(_watch_disconnect()) next_chunk_task: asyncio.Task[bytes] | None = None + usage = None try: client = get_http_client() async with client.stream( @@ -330,6 +333,9 @@ async def _read_next_chunk( authorized_chat_request, result, time.perf_counter() - start_time ) record_chat_availability(authorized_chat_request, availability_reason) + await redis_service.update_contracts( + service_type=authorized_chat_request.service_type, usage=usage + ) async def get_completion(authorized_chat_request: AuthorizedChatRequest): @@ -350,6 +356,7 @@ async def _get_completion(authorized_chat_request: AuthorizedChatRequest): logger.debug( f"Starting a non-stream completion using {authorized_chat_request.model}, for user {authorized_chat_request.user}", ) + usage = None try: client = get_http_client() response = await client.post( @@ -417,3 +424,6 @@ async def _get_completion(authorized_chat_request: AuthorizedChatRequest): authorized_chat_request, result, time.perf_counter() - start_time ) record_chat_availability(authorized_chat_request, availability_reason) + await redis_service.update_contracts( + service_type=authorized_chat_request.service_type, usage=usage + ) diff --git a/src/mlpa/core/config.py b/src/mlpa/core/config.py index c582060b..ca4db826 100644 --- a/src/mlpa/core/config.py +++ b/src/mlpa/core/config.py @@ -160,24 +160,25 @@ def _parse_launch_countries(cls, raw: str | set[str]) -> set[str]: ENABLE_TRAFFIC_CONTRACT_ENFORCEMENT: bool = False TRAFFIC_CONTRACT_REDIS_KEY_PREFIX: str = "mlpa:traffic_contract" TRAFFIC_CONTRACT_RPM_WINDOW_SECONDS: int = 60 + TRAFFIC_CONTRACT_TPM_WINDOW_SECONDS: int = 60 TRAFFIC_CONTRACT_COUNTER_TTL_SECONDS: int = 120 - TRAFFIC_CONTRACT_FAIL_OPEN_ON_REDIS_ERROR: bool = False + TRAFFIC_CONTRACT_FAIL_OPEN_ON_REDIS_ERROR: bool = True - # 0 means no cap enforced; still records per-feature counters. + # Global RPM/TPM limit (0 means no limit) TOTAL_TRAFFIC_CONTRACT_RPM_LIMIT: int = 0 TOTAL_TRAFFIC_CONTRACT_TPM_LIMIT: int = 0 FEATURE_SMART_WINDOW: str = "smart-window" - SMART_WINDOW_TRAFFIC_CONTRACT_RPM_LIMIT: int = 10_000 - SMART_WINDOW_TRAFFIC_CONTRACT_TPM_LIMIT: int = 1_000_000 + SMART_WINDOW_TRAFFIC_CONTRACT_RPM_LIMIT: int = 0 + SMART_WINDOW_TRAFFIC_CONTRACT_TPM_LIMIT: int = 0 FEATURE_S2S: str = "s2s" - S2S_TRAFFIC_CONTRACT_RPM_LIMIT: int = 10_000 - S2S_TRAFFIC_CONTRACT_TPM_LIMIT: int = 1_000_000 + S2S_TRAFFIC_CONTRACT_RPM_LIMIT: int = 0 + S2S_TRAFFIC_CONTRACT_TPM_LIMIT: int = 0 FEATURE_S2S_ANDROID: str = "s2s-android" - S2S_ANDROID_TRAFFIC_CONTRACT_RPM_LIMIT: int = 10_000 - S2S_ANDROID_TRAFFIC_CONTRACT_TPM_LIMIT: int = 1_000_000 + S2S_ANDROID_TRAFFIC_CONTRACT_RPM_LIMIT: int = 0 + S2S_ANDROID_TRAFFIC_CONTRACT_TPM_LIMIT: int = 0 @cached_property def service_type_config(self) -> dict[str, dict]: diff --git a/src/mlpa/core/consts/__init__.py b/src/mlpa/core/consts/__init__.py new file mode 100644 index 00000000..c89602aa --- /dev/null +++ b/src/mlpa/core/consts/__init__.py @@ -0,0 +1,13 @@ +from mlpa.core.consts.country_codes import ( + COUNTRY_CODES, +) +from mlpa.core.consts.traffic_contract import ( + TrafficContractKeyType, + TrafficContractMode, +) + +__all__ = [ + "COUNTRY_CODES", + "TrafficContractKeyType", + "TrafficContractMode", +] diff --git a/src/mlpa/core/country_codes.py b/src/mlpa/core/consts/country_codes.py similarity index 100% rename from src/mlpa/core/country_codes.py rename to src/mlpa/core/consts/country_codes.py diff --git a/src/mlpa/core/openapi.py b/src/mlpa/core/consts/openapi.py similarity index 60% rename from src/mlpa/core/openapi.py rename to src/mlpa/core/consts/openapi.py index f313b180..7f51935a 100644 --- a/src/mlpa/core/openapi.py +++ b/src/mlpa/core/consts/openapi.py @@ -9,6 +9,79 @@ SEARCH_SERVICE_TYPES = ("search", "search-dev", "agent-search") SEARCH_SERVICE_TYPES_SET = set(SEARCH_SERVICE_TYPES) +TAGS_METADATA = [ + {"name": "Health", "description": "Health check endpoints."}, + {"name": "Metrics", "description": "Prometheus metrics endpoints."}, + { + "name": "App Attest", + "description": "iOS App Attest verification flow: (1) GET /verify/challenge to obtain a challenge, " + "(2) POST /verify/attest with a JWT containing the attestation object. " + "Use the attested key for subsequent requests to /v1/chat/completions with use-app-attest header.", + }, + { + "name": "Play Integrity", + "description": "Endpoints for verifying Play Integrity payloads.", + }, + {"name": "LiteLLM", "description": "Endpoints for interacting with LiteLLM."}, + {"name": "Mock", "description": "Mock endpoints for testing purposes."}, + { + "name": "User Management", + "description": "Endpoints for managing user blocking status and budgets.", + }, + { + "name": "Privacy Filter", + "description": "Endpoints for interacting with the Privacy Filter.", + }, +] + +CHAT_COMPLETION_DESCRIPTION = """ +Authorize first using App Attest, Play Integrity, FxA, or dev tier. + +**Headers:** + +- **Authorization** (required): Bearer token — FxA OAuth token, Play Integrity MLPA token, or App Attest JWT. +- **service-type** (required): One of `ai`, `s2s`, `s2s-android`, `memories`, `ai-dev`, `memories-dev`, `mochi-dev`, `answer`, `liner-answer`, `telemetry` — for tracking and budget. +- **purpose** (required for ai/ai-dev/mochi-dev/memories/memories-dev): One of `chat`, `title-generation`, `convo-starters-sidebar` for AI; `memory-generation` for memories; omit for s2s. +- **x-dev-authorization** (required for ai-dev/memories-dev/mochi-dev): Experimentation token; also requires FxA in Authorization. Dev service types return 401 without it. +- **use-app-attest**: Set to `true` for iOS App Attest. +- **use-play-integrity**: Set to `true` for Android Play Integrity. +""" + + +SEARCH_DESCRIPTION = """ +Web search proxied to Exa via LiteLLM. Authorize the same way as /v1/chat/completions. + +**Headers:** + +- **Authorization** (required): Bearer token — FxA OAuth token, Play Integrity MLPA token, or App Attest JWT. +- **service-type**: `search` by default; use `search-dev` for experiments. Search has its own budget pool and no `purpose` header. +- **x-dev-authorization** (required for search-dev): Experimentation token; also requires FxA in Authorization. + +**Body:** `{"query": str, "max_results": int (1-10)}`. +""" + +# Success (200) response docs for the proxied LiteLLM endpoints. The chat endpoint +# returns either a JSON chat completion or an SSE stream depending on `stream`. +CHAT_COMPLETION_SUCCESS_RESPONSE: dict[int | str, dict[str, Any]] = { + 200: { + "description": ( + "OpenAI-compatible chat completion. Returns a JSON completion object, or " + "a `text/event-stream` of SSE chunks when `stream` is `true`." + ), + "content": { + "application/json": {}, + "text/event-stream": {}, + }, + } +} + +SEARCH_SUCCESS_RESPONSE: dict[int | str, dict[str, Any]] = { + 200: { + "description": "Search results returned from the Exa search backend.", + "content": {"application/json": {}}, + } +} + def customize_openapi(app: FastAPI, tags_metadata: list[dict]) -> None: """Add AttestationAuth and AssertionAuth schemas to OpenAPI docs.""" diff --git a/src/mlpa/core/consts/traffic_contract.py b/src/mlpa/core/consts/traffic_contract.py new file mode 100644 index 00000000..eb817307 --- /dev/null +++ b/src/mlpa/core/consts/traffic_contract.py @@ -0,0 +1,12 @@ +from enum import StrEnum + + +class TrafficContractKeyType(StrEnum): + RPM = "rpm" + TPM = "tpm" + + +class TrafficContractMode(StrEnum): + NORMAL = "normal" + BORROWED = "borrowed" + DEGRADED = "degraded" diff --git a/src/mlpa/core/middleware/instrumentation.py b/src/mlpa/core/middleware/instrumentation.py index 4fd90abd..1abc8c7c 100644 --- a/src/mlpa/core/middleware/instrumentation.py +++ b/src/mlpa/core/middleware/instrumentation.py @@ -2,6 +2,7 @@ from fastapi import Request +from mlpa.core.consts import TrafficContractKeyType, TrafficContractMode from mlpa.core.logger import logger from mlpa.core.prometheus_metrics import metrics from mlpa.core.utils import ( @@ -13,6 +14,22 @@ ) +def _traffic_contract_mode_label( + request: Request, + key_type: TrafficContractKeyType, +) -> str: + mode = getattr( + request.state, + f"traffic_contract_{key_type.value}_mode", + "N/A", + ) + if isinstance(mode, TrafficContractMode): + return mode.value + if isinstance(mode, str) and mode in {mode.value for mode in TrafficContractMode}: + return mode + return "N/A" + + async def instrument_requests_middleware(request: Request, call_next): """ Measures request latency, counts total requests, and tracks requests in progress. @@ -47,6 +64,12 @@ async def instrument_requests_middleware(request: Request, call_next): service_type=clamp_service_type(service_type), purpose=clamp_purpose(purpose), major_fx_version=clamp_major_fx_version(major_fx_version), + traffic_contract_rpm_mode=_traffic_contract_mode_label( + request, TrafficContractKeyType.RPM + ), + traffic_contract_tpm_mode=_traffic_contract_mode_label( + request, TrafficContractKeyType.TPM + ), ).inc() metrics.response_status_codes.labels(status_code=response.status_code).inc() return response diff --git a/src/mlpa/core/middleware/traffic_contract_enforcer.py b/src/mlpa/core/middleware/traffic_contract_enforcer.py index 96065e60..cb4d141a 100644 --- a/src/mlpa/core/middleware/traffic_contract_enforcer.py +++ b/src/mlpa/core/middleware/traffic_contract_enforcer.py @@ -1,34 +1,68 @@ -from fastapi import HTTPException +import asyncio -from mlpa.core.classes import AuthorizedChatRequest, AuthorizedSearchRequest -from mlpa.core.config import TrafficContractConfig, env +from fastapi import HTTPException, Request + +from mlpa.core.classes import ( + AuthorizedChatRequest, + AuthorizedSearchRequest, +) +from mlpa.core.config import env +from mlpa.core.consts import TrafficContractKeyType, TrafficContractMode from mlpa.core.logger import logger -from mlpa.core.metrics import record_chat_availability -from mlpa.core.prometheus_metrics import AvailabilityReason from mlpa.core.services.services import redis_service -def _traffic_contract_for_service_type( - service_type: str, -) -> TrafficContractConfig | None: - return env.traffic_contract_config.get(service_type) +def _set_traffic_contract_modes( + request: Request | None, + *, + rpm_mode: TrafficContractMode | str = "N/A", + tpm_mode: TrafficContractMode | str = "N/A", +) -> None: + if request is None: + return + request.state.traffic_contract_rpm_mode = rpm_mode + request.state.traffic_contract_tpm_mode = tpm_mode + +async def enforce_traffic_contract( + request: Request, + service_type: str, +) -> None: + _set_traffic_contract_modes(request) -async def _enforce_traffic_contract(service_type: str) -> None: if not env.ENABLE_TRAFFIC_CONTRACT_ENFORCEMENT: return - contract = _traffic_contract_for_service_type(service_type) + contract = env.traffic_contract_config.get(service_type) if contract is None: return try: - rpm_decision = await redis_service.check_and_increment_feature_rpm( - key_prefix=env.TRAFFIC_CONTRACT_REDIS_KEY_PREFIX, - feature=contract["feature"], - feature_rpm_limit=contract["rpm_limit"], - window_seconds=env.TRAFFIC_CONTRACT_RPM_WINDOW_SECONDS, - ttl_seconds=env.TRAFFIC_CONTRACT_COUNTER_TTL_SECONDS, + rpm_decision, tpm_decision = await asyncio.gather( + redis_service.check_feature_traffic_contract( + key_prefix=env.TRAFFIC_CONTRACT_REDIS_KEY_PREFIX, + key_type=TrafficContractKeyType.RPM, + feature=contract["feature"], + feature_limit=contract["rpm_limit"], + basket_limit=env.TOTAL_TRAFFIC_CONTRACT_RPM_LIMIT, + increment_amount=1, + window_seconds=env.TRAFFIC_CONTRACT_RPM_WINDOW_SECONDS, + ), + redis_service.check_feature_traffic_contract( + key_prefix=env.TRAFFIC_CONTRACT_REDIS_KEY_PREFIX, + key_type=TrafficContractKeyType.TPM, + feature=contract["feature"], + feature_limit=contract["tpm_limit"], + basket_limit=env.TOTAL_TRAFFIC_CONTRACT_TPM_LIMIT, + increment_amount=0, # Token usage is incremented after the response. + window_seconds=env.TRAFFIC_CONTRACT_TPM_WINDOW_SECONDS, + ), + ) + + _set_traffic_contract_modes( + request, + rpm_mode=rpm_decision.mode, + tpm_mode=tpm_decision.mode, ) except Exception as exc: logger.error( @@ -44,7 +78,7 @@ async def _enforce_traffic_contract(service_type: str) -> None: if not rpm_decision.allowed: logger.warning( - "Traffic contract exceeded for " + "RPM Traffic contract exceeded for " f"service_type={service_type}, " f"feature={contract['feature']}, limited_by={rpm_decision.limited_by}, " f"mode={rpm_decision.mode}, ratio_over={rpm_decision.ratio_over}, " @@ -54,18 +88,14 @@ async def _enforce_traffic_contract(service_type: str) -> None: # Hard deny over limit requests: # raise _rate_limit_response(rpm_decision.retry_after_seconds) - -async def enforce_chat_traffic_contract(req: AuthorizedChatRequest) -> None: - try: - await _enforce_traffic_contract(req.service_type) - except HTTPException as exc: - raise - - -async def enforce_search_traffic_contract( - req: AuthorizedSearchRequest, -) -> None: - try: - await _enforce_traffic_contract(req.service_type) - except HTTPException: - raise + if not tpm_decision.allowed: + logger.warning( + "TPM Traffic contract exceeded for " + f"service_type={service_type}, " + f"feature={contract['feature']}, limited_by={tpm_decision.limited_by}, " + f"mode={tpm_decision.mode}, ratio_over={tpm_decision.ratio_over}, " + f"feature_count={tpm_decision.feature_count}, " + f"basket_count={tpm_decision.basket_count}" + ) + # Hard deny over limit requests: + # raise _rate_limit_response(tpm_decision.retry_after_seconds) diff --git a/src/mlpa/core/prometheus_metrics.py b/src/mlpa/core/prometheus_metrics.py index c5ad427c..3eb0e24d 100644 --- a/src/mlpa/core/prometheus_metrics.py +++ b/src/mlpa/core/prometheus_metrics.py @@ -212,7 +212,15 @@ def build_metrics(registry: CollectorRegistry = REGISTRY) -> PrometheusMetrics: requests_total=Counter( "mlpa_requests_total", "Total number of requests handled by the proxy.", - ["method", "endpoint", "service_type", "purpose", "major_fx_version"], + [ + "method", + "endpoint", + "service_type", + "purpose", + "major_fx_version", + "traffic_contract_rpm_mode", + "traffic_contract_tpm_mode", + ], registry=registry, ), requests_by_country_total=Counter( diff --git a/src/mlpa/core/services/redis_service.py b/src/mlpa/core/services/redis_service.py index c4a4f44d..71ec14ec 100644 --- a/src/mlpa/core/services/redis_service.py +++ b/src/mlpa/core/services/redis_service.py @@ -1,26 +1,30 @@ +import asyncio import time -from dataclasses import dataclass -from typing import Any, Literal, cast +from typing import Any import redis.asyncio as aioredis +from mlpa.core.classes import TrafficContractCounters, TrafficContractDecision from mlpa.core.config import env +from mlpa.core.consts import TrafficContractKeyType, TrafficContractMode from mlpa.core.logger import logger TRAFFIC_CONTRACT_BASKET_FIELD = "__basket__" -_CHECK_AND_INCREMENT_SCRIPT = """ +_CHECK_TRAFFIC_CONTRACT_SCRIPT = """ local key = KEYS[1] local feature_field = ARGV[1] local basket_field = ARGV[2] local feature_limit = tonumber(ARGV[3]) local basket_limit = tonumber(ARGV[4]) -local ttl_seconds = tonumber(ARGV[5]) -local inc_amount = tonumber(ARGV[6]) +local inc_amount = tonumber(ARGV[5]) -local feature_count = redis.call("HINCRBY", key, feature_field, inc_amount) -local basket_count = redis.call("HINCRBY", key, basket_field, inc_amount) -redis.call("EXPIRE", key, ttl_seconds) +if basket_limit == 0 and feature_limit == 0 then + return {1, 0, 0, "", "normal", ""} +end + +local feature_count = tonumber(redis.call("HGET", key, feature_field) or "0") + inc_amount +local basket_count = tonumber(redis.call("HGET", key, basket_field) or "0") + inc_amount if basket_limit > 0 and basket_count > basket_limit then return {0, feature_count, basket_count, "basket", "degraded", tostring(basket_count/basket_limit)} @@ -33,21 +37,25 @@ return {1, feature_count, basket_count, "", "normal", ""} """ -TrafficContractMode = Literal["normal", "borrowed", "degraded"] +_INCREMENT_TRAFFIC_CONTRACT_SCRIPT = """ +local key = KEYS[1] +local feature_field = ARGV[1] +local basket_field = ARGV[2] +local ttl_seconds = tonumber(ARGV[3]) +local inc_amount = tonumber(ARGV[4]) +if inc_amount <= 0 then + local feature_count = tonumber(redis.call("HGET", key, feature_field) or "0") + local basket_count = tonumber(redis.call("HGET", key, basket_field) or "0") + return {feature_count, basket_count} +end -@dataclass(frozen=True) -class TrafficContractDecision: - allowed: bool - feature_count: int - basket_count: int - retry_after_seconds: int - limited_by: str | None = None +local feature_count = redis.call("HINCRBY", key, feature_field, inc_amount) +local basket_count = redis.call("HINCRBY", key, basket_field, inc_amount) +redis.call("EXPIRE", key, ttl_seconds) - # borrowed = feature is over limit, basket has room - # degraded = basket is over limit - mode: TrafficContractMode = "normal" - ratio_over: float | None = None # ratio of count / limit if mode != "normal" +return {feature_count, basket_count} +""" class RedisService: @@ -90,7 +98,7 @@ def traffic_contract_key( cls, *, key_prefix: str, - key_type: Literal["rpm", "tpm"], + key_type: TrafficContractKeyType, now: int | None = None, window_seconds: int = 60, ) -> str: @@ -105,19 +113,21 @@ def retry_after_seconds( elapsed_in_bucket = current_time % window_seconds return window_seconds - elapsed_in_bucket - async def check_and_increment_feature_rpm( + async def check_feature_traffic_contract( self, *, key_prefix: str, + key_type: TrafficContractKeyType, feature: str, - feature_rpm_limit: int, + feature_limit: int, + basket_limit: int, + increment_amount: int, window_seconds: int = 60, - ttl_seconds: int = 120, now: int | None = None, ) -> TrafficContractDecision: key = self.traffic_contract_key( key_prefix=key_prefix, - key_type="rpm", + key_type=key_type, now=now, window_seconds=window_seconds, ) @@ -127,15 +137,14 @@ async def check_and_increment_feature_rpm( ) result = await self.client.eval( - _CHECK_AND_INCREMENT_SCRIPT, + _CHECK_TRAFFIC_CONTRACT_SCRIPT, 1, key, feature, TRAFFIC_CONTRACT_BASKET_FIELD, - feature_rpm_limit, - env.TOTAL_TRAFFIC_CONTRACT_RPM_LIMIT, - ttl_seconds, - 1, + feature_limit, + basket_limit, + increment_amount, ) return TrafficContractDecision( @@ -144,23 +153,168 @@ async def check_and_increment_feature_rpm( basket_count=int(result[2]), retry_after_seconds=retry_after, limited_by=str(result[3]) or None, - mode=cast(TrafficContractMode, result[4]), + mode=TrafficContractMode(result[4]), ratio_over=float(result[5]) if result[5] not in (None, "") else None, ) - async def get_current_feature_rpm( + async def increment_feature_traffic_contract( self, *, key_prefix: str, + key_type: TrafficContractKeyType, feature: str, + increment_amount: int, window_seconds: int = 60, + ttl_seconds: int = 120, now: int | None = None, - ) -> int: + ) -> TrafficContractCounters: key = self.traffic_contract_key( key_prefix=key_prefix, - key_type="rpm", + key_type=key_type, + now=now, + window_seconds=window_seconds, + ) + result = await self.client.eval( + _INCREMENT_TRAFFIC_CONTRACT_SCRIPT, + 1, + key, + feature, + TRAFFIC_CONTRACT_BASKET_FIELD, + ttl_seconds, + increment_amount, + ) + + return TrafficContractCounters( + feature_count=int(result[0]), + basket_count=int(result[1]), + ) + + async def check_feature_rpm( + self, + *, + key_prefix: str, + feature: str, + feature_rpm_limit: int, + window_seconds: int = 60, + now: int | None = None, + ) -> TrafficContractDecision: + return await self.check_feature_traffic_contract( + key_prefix=key_prefix, + key_type=TrafficContractKeyType.RPM, + feature=feature, + feature_limit=feature_rpm_limit, + basket_limit=env.TOTAL_TRAFFIC_CONTRACT_RPM_LIMIT, + increment_amount=1, + window_seconds=window_seconds, + now=now, + ) + + async def increment_feature_rpm( + self, + *, + key_prefix: str, + feature: str, + window_seconds: int = 60, + ttl_seconds: int = 120, + now: int | None = None, + ) -> TrafficContractCounters: + return await self.increment_feature_traffic_contract( + key_prefix=key_prefix, + key_type=TrafficContractKeyType.RPM, + feature=feature, + increment_amount=1, + window_seconds=window_seconds, + ttl_seconds=ttl_seconds, + now=now, + ) + + async def check_feature_tpm( + self, + *, + key_prefix: str, + feature: str, + feature_tpm_limit: int, + token_count: int = 0, + window_seconds: int = 60, + now: int | None = None, + ) -> TrafficContractDecision: + return await self.check_feature_traffic_contract( + key_prefix=key_prefix, + key_type=TrafficContractKeyType.TPM, + feature=feature, + feature_limit=feature_tpm_limit, + basket_limit=env.TOTAL_TRAFFIC_CONTRACT_TPM_LIMIT, + increment_amount=token_count, + window_seconds=window_seconds, + now=now, + ) + + async def increment_feature_tpm( + self, + *, + key_prefix: str, + feature: str, + token_count: int, + window_seconds: int = 60, + ttl_seconds: int = 120, + now: int | None = None, + ) -> TrafficContractCounters: + return await self.increment_feature_traffic_contract( + key_prefix=key_prefix, + key_type=TrafficContractKeyType.TPM, + feature=feature, + increment_amount=token_count, + window_seconds=window_seconds, + ttl_seconds=ttl_seconds, now=now, + ) + + async def inc_traffic_contract( + self, + *, + key_type: TrafficContractKeyType, + service_type: str, + increment_amount: int, + ) -> TrafficContractCounters | None: + if not env.ENABLE_TRAFFIC_CONTRACT_ENFORCEMENT or increment_amount <= 0: + return None + + contract = env.traffic_contract_config.get(service_type) + if contract is None: + return None + + window_seconds = ( + env.TRAFFIC_CONTRACT_RPM_WINDOW_SECONDS + if key_type == TrafficContractKeyType.RPM + else env.TRAFFIC_CONTRACT_TPM_WINDOW_SECONDS + ) + return await self.increment_feature_traffic_contract( + key_prefix=env.TRAFFIC_CONTRACT_REDIS_KEY_PREFIX, + key_type=key_type, + feature=contract["feature"], + increment_amount=increment_amount, window_seconds=window_seconds, + ttl_seconds=env.TRAFFIC_CONTRACT_COUNTER_TTL_SECONDS, ) - traffic_count = await self.client.hget(key, feature) - return int(traffic_count) if traffic_count else 0 + + async def update_contracts(self, *, service_type: str, usage: dict | None): + if not env.ENABLE_TRAFFIC_CONTRACT_ENFORCEMENT: + return + + updates = [ + self.inc_traffic_contract( + key_type=TrafficContractKeyType.RPM, + service_type=service_type, + increment_amount=1, + ) + ] + if usage and usage.get("total_tokens"): + updates.append( + self.inc_traffic_contract( + key_type=TrafficContractKeyType.TPM, + service_type=service_type, + increment_amount=usage["total_tokens"], + ) + ) + + await asyncio.gather(*updates) diff --git a/src/mlpa/core/utils.py b/src/mlpa/core/utils.py index c1642199..af085159 100644 --- a/src/mlpa/core/utils.py +++ b/src/mlpa/core/utils.py @@ -17,7 +17,7 @@ LITELLM_MASTER_AUTH_HEADERS, env, ) -from mlpa.core.country_codes import COUNTRY_CODES +from mlpa.core.consts.country_codes import COUNTRY_CODES from mlpa.core.http_client import get_http_client from mlpa.core.logger import logger from mlpa.core.prometheus_metrics import PrometheusResult, metrics diff --git a/src/mlpa/run.py b/src/mlpa/run.py index ff862e79..1d06dfb1 100644 --- a/src/mlpa/run.py +++ b/src/mlpa/run.py @@ -24,6 +24,14 @@ SENSITIVE_FIELDS_TO_SCRUB_FROM_SENTRY, env, ) +from mlpa.core.consts.openapi import ( + CHAT_COMPLETION_DESCRIPTION, + CHAT_COMPLETION_SUCCESS_RESPONSE, + SEARCH_DESCRIPTION, + SEARCH_SUCCESS_RESPONSE, + TAGS_METADATA, + customize_openapi, +) from mlpa.core.http_client import close_http_client, get_http_client from mlpa.core.logger import logger, setup_logger from mlpa.core.metrics import ( @@ -34,10 +42,8 @@ ) from mlpa.core.middleware import register_middleware from mlpa.core.middleware.traffic_contract_enforcer import ( - enforce_chat_traffic_contract, - enforce_search_traffic_contract, + enforce_traffic_contract, ) -from mlpa.core.openapi import customize_openapi from mlpa.core.prometheus_metrics import AvailabilityReason from mlpa.core.routers.appattest import appattest_router from mlpa.core.routers.filter import filter_router @@ -48,31 +54,6 @@ from mlpa.core.search import get_search from mlpa.core.services.services import app_attest_pg, litellm_pg, redis_service -tags_metadata = [ - {"name": "Health", "description": "Health check endpoints."}, - {"name": "Metrics", "description": "Prometheus metrics endpoints."}, - { - "name": "App Attest", - "description": "iOS App Attest verification flow: (1) GET /verify/challenge to obtain a challenge, " - "(2) POST /verify/attest with a JWT containing the attestation object. " - "Use the attested key for subsequent requests to /v1/chat/completions with use-app-attest header.", - }, - { - "name": "Play Integrity", - "description": "Endpoints for verifying Play Integrity payloads.", - }, - {"name": "LiteLLM", "description": "Endpoints for interacting with LiteLLM."}, - {"name": "Mock", "description": "Mock endpoints for testing purposes."}, - { - "name": "User Management", - "description": "Endpoints for managing user blocking status and budgets.", - }, - { - "name": "Privacy Filter", - "description": "Endpoints for interacting with the Privacy Filter.", - }, -] - @asynccontextmanager async def lifespan(app: FastAPI): @@ -134,7 +115,7 @@ def sentry_scrub_sensitive_fields(event, hint): description="Authenticates and proxies LLM requests through LiteLLM to enact budgets and per-user management.", version=importlib.metadata.version("mlpa"), docs_url="/api/docs", - openapi_tags=tags_metadata, + openapi_tags=TAGS_METADATA, lifespan=lifespan, ) @@ -154,7 +135,7 @@ async def get_metrics(): app.include_router(user_router, prefix="/user") app.include_router(filter_router) app.include_router(mock_router, prefix="/mock") -customize_openapi(app, tags_metadata) +customize_openapi(app, TAGS_METADATA) app.mount( "/admin", @@ -163,55 +144,6 @@ async def get_metrics(): ) -CHAT_COMPLETION_DESCRIPTION = """ -Authorize first using App Attest, Play Integrity, FxA, or dev tier. - -**Headers:** - -- **Authorization** (required): Bearer token — FxA OAuth token, Play Integrity MLPA token, or App Attest JWT. -- **service-type** (required): One of `ai`, `s2s`, `s2s-android`, `memories`, `ai-dev`, `memories-dev`, `mochi-dev`, `answer`, `liner-answer`, `telemetry` — for tracking and budget. -- **purpose** (required for ai/ai-dev/mochi-dev/memories/memories-dev): One of `chat`, `title-generation`, `convo-starters-sidebar` for AI; `memory-generation` for memories; omit for s2s. -- **x-dev-authorization** (required for ai-dev/memories-dev/mochi-dev): Experimentation token; also requires FxA in Authorization. Dev service types return 401 without it. -- **use-app-attest**: Set to `true` for iOS App Attest. -- **use-play-integrity**: Set to `true` for Android Play Integrity. -""" - - -SEARCH_DESCRIPTION = """ -Web search proxied to Exa via LiteLLM. Authorize the same way as /v1/chat/completions. - -**Headers:** - -- **Authorization** (required): Bearer token — FxA OAuth token, Play Integrity MLPA token, or App Attest JWT. -- **service-type**: `search` by default; use `search-dev` for experiments. Search has its own budget pool and no `purpose` header. -- **x-dev-authorization** (required for search-dev): Experimentation token; also requires FxA in Authorization. - -**Body:** `{"query": str, "max_results": int (1-10)}`. -""" - -# Success (200) response docs for the proxied LiteLLM endpoints. The chat endpoint -# returns either a JSON chat completion or an SSE stream depending on `stream`. -CHAT_COMPLETION_SUCCESS_RESPONSE: dict[int | str, dict[str, Any]] = { - 200: { - "description": ( - "OpenAI-compatible chat completion. Returns a JSON completion object, or " - "a `text/event-stream` of SSE chunks when `stream` is `true`." - ), - "content": { - "application/json": {}, - "text/event-stream": {}, - }, - } -} - -SEARCH_SUCCESS_RESPONSE: dict[int | str, dict[str, Any]] = { - 200: { - "description": "Search results returned from the Exa search backend.", - "content": {"application/json": {}}, - } -} - - @app.post( "/v1/chat/completions", tags=["LiteLLM"], @@ -229,7 +161,7 @@ async def chat_completion( service_type=authorized_chat_request.service_type, model=authorized_chat_request.model, ) - await enforce_chat_traffic_contract(authorized_chat_request) + await enforce_traffic_contract(request, authorized_chat_request.service_type) user_id = authorized_chat_request.user if not user_id: raise HTTPException( @@ -274,9 +206,7 @@ async def search( status_code=400, detail=f"service-type header must be one of {env.forced_model_service_type_pairs.get(SEARCH_MODEL)}", ) - await enforce_search_traffic_contract( - authorized_search_request, - ) + await enforce_traffic_contract(request, authorized_search_request.service_type) user_id = authorized_search_request.user if not user_id: raise HTTPException( diff --git a/src/tests/integration/test_traffic_contracts.py b/src/tests/integration/test_traffic_contracts.py index 74927152..db99ddd9 100644 --- a/src/tests/integration/test_traffic_contracts.py +++ b/src/tests/integration/test_traffic_contracts.py @@ -11,7 +11,7 @@ def test_chat_completion_continues_when_traffic_contract_exceeded( mocker.patch.object(env, "ENABLE_TRAFFIC_CONTRACT_ENFORCEMENT", True) mocker.patch.object( traffic_contract_enforcer.redis_service, - "check_and_increment_feature_rpm", + "check_feature_traffic_contract", mocker.AsyncMock( return_value=TrafficContractDecision( allowed=False, @@ -24,6 +24,11 @@ def test_chat_completion_continues_when_traffic_contract_exceeded( ) ), ) + mocker.patch.object( + traffic_contract_enforcer.redis_service, + "inc_traffic_contract", + mocker.AsyncMock(), + ) get_completion = mocker.patch( "mlpa.run.get_completion", mocker.AsyncMock(return_value={"id": "chatcmpl-test"}), diff --git a/src/tests/unit/test_middleware_order.py b/src/tests/unit/test_middleware_order.py index 2d945b0c..dbb88424 100644 --- a/src/tests/unit/test_middleware_order.py +++ b/src/tests/unit/test_middleware_order.py @@ -86,6 +86,8 @@ async def test_endpoint(): service_type="other", purpose="other", major_fx_version="", + traffic_contract_rpm_mode="N/A", + traffic_contract_tpm_mode="N/A", ) == 1 ) @@ -97,6 +99,8 @@ async def test_endpoint(): service_type="not-real-service-type", purpose="not-real-purpose", major_fx_version="", + traffic_contract_rpm_mode="N/A", + traffic_contract_tpm_mode="N/A", ) == 0 ) @@ -130,6 +134,45 @@ async def test_endpoint(): service_type="ai", purpose="chat", major_fx_version="", + traffic_contract_rpm_mode="N/A", + traffic_contract_tpm_mode="N/A", + ) + == 1 + ) + + +def test_instrumentation_uses_traffic_contract_mode_from_request_state(metrics_spy): + from mlpa.core.middleware.instrumentation import instrument_requests_middleware + + app = FastAPI() + app.middleware("http")(instrument_requests_middleware) + + @app.get("/test") + async def test_endpoint(request: Request): + request.state.traffic_contract_rpm_mode = "borrowed" + request.state.traffic_contract_tpm_mode = "borrowed" + return {"status": "ok"} + + client = TestClient(app) + response = client.get( + "/test", + headers={ + "service-type": "ai", + "purpose": "chat", + }, + ) + + assert response.status_code == 200 + assert ( + metrics_spy.value( + "requests_total", + method="GET", + endpoint="/test", + service_type="ai", + purpose="chat", + major_fx_version="", + traffic_contract_rpm_mode="borrowed", + traffic_contract_tpm_mode="borrowed", ) == 1 ) @@ -157,6 +200,8 @@ async def test_endpoint(): service_type="s2s", purpose="", major_fx_version="", + traffic_contract_rpm_mode="N/A", + traffic_contract_tpm_mode="N/A", ) == 1 ) @@ -191,6 +236,8 @@ async def test_endpoint(): service_type="ai", purpose="chat", major_fx_version="", + traffic_contract_rpm_mode="N/A", + traffic_contract_tpm_mode="N/A", ) == 1 ) @@ -202,6 +249,8 @@ async def test_endpoint(): service_type="ai", purpose="chat", major_fx_version="", + traffic_contract_rpm_mode="N/A", + traffic_contract_tpm_mode="N/A", ) == 0 ) @@ -310,6 +359,8 @@ async def test_endpoint(): service_type="ai", purpose="chat", major_fx_version=parse_firefox_major_version_from_user_agent(user_agent), + traffic_contract_rpm_mode="N/A", + traffic_contract_tpm_mode="N/A", ) == 1 ) @@ -347,6 +398,8 @@ async def test_endpoint(): service_type="ai", purpose="chat", major_fx_version=parse_firefox_major_version_from_user_agent(user_agent), + traffic_contract_rpm_mode="N/A", + traffic_contract_tpm_mode="N/A", ) == 1 ) diff --git a/src/tests/unit/test_redis_service.py b/src/tests/unit/test_redis_service.py index 2d96e90c..30eaaffd 100644 --- a/src/tests/unit/test_redis_service.py +++ b/src/tests/unit/test_redis_service.py @@ -1,4 +1,7 @@ +from unittest.mock import call + from mlpa.core.config import env +from mlpa.core.consts import TrafficContractKeyType, TrafficContractMode from mlpa.core.services.redis_service import ( TRAFFIC_CONTRACT_BASKET_FIELD, RedisService, @@ -15,25 +18,34 @@ async def eval( script, num_keys, key, - feature, - basket_field, - feature_limit, - basket_limit, - ttl_seconds, - inc_amount, + *args, ): assert num_keys == 1 - assert "HINCRBY" in script + if "HINCRBY" in script: + feature, basket_field, ttl_seconds, inc_amount = args + bucket = self.hashes.setdefault(key, {}) + if int(inc_amount) <= 0: + return [ + int(bucket.get(feature, 0)), + int(bucket.get(basket_field, 0)), + ] - bucket = self.hashes.setdefault(key, {}) - feature_count = int(bucket.get(feature, 0)) + int(inc_amount) - basket_count = int(bucket.get(basket_field, 0)) + int(inc_amount) + feature_count = int(bucket.get(feature, 0)) + int(inc_amount) + basket_count = int(bucket.get(basket_field, 0)) + int(inc_amount) + bucket[feature] = feature_count + bucket[basket_field] = basket_count + self.expirations[key] = int(ttl_seconds) + return [feature_count, basket_count] + + feature, basket_field, feature_limit, basket_limit, inc_amount = args + bucket = self.hashes.get(key, {}) feature_limit = int(feature_limit) basket_limit = int(basket_limit) + if basket_limit == 0 and feature_limit == 0: + return [1, 0, 0, "", "normal", ""] - bucket[feature] = feature_count - bucket[basket_field] = basket_count - self.expirations[key] = int(ttl_seconds) + feature_count = int(bucket.get(feature, 0)) + int(inc_amount) + basket_count = int(bucket.get(basket_field, 0)) + int(inc_amount) if basket_limit > 0 and basket_count > basket_limit: return [ @@ -55,98 +67,216 @@ async def eval( str(feature_count / feature_limit), ] - return [1, bucket[feature], bucket[basket_field], "", "normal", ""] + return [1, feature_count, basket_count, "", "normal", ""] async def hget(self, key, field): return self.hashes.get(key, {}).get(field) -async def test_check_and_increment_feature_rpm_records_minute_bucket(mocker): +async def test_check_feature_rpm_does_not_mutate_bucket(mocker): mocker.patch.object(env, "TOTAL_TRAFFIC_CONTRACT_RPM_LIMIT", 20) redis = FakeRedis() service = RedisService() service.redis = redis - decision = await service.check_and_increment_feature_rpm( + decision = await service.check_feature_rpm( key_prefix="mlpa:traffic_contract", feature="smart-window", feature_rpm_limit=10, - ttl_seconds=120, now=125, ) - key = "mlpa:traffic_contract:rpm:120" assert decision.allowed is True assert decision.feature_count == 1 assert decision.basket_count == 1 assert decision.retry_after_seconds == 55 assert decision.limited_by is None - assert decision.mode == "normal" + assert decision.mode == TrafficContractMode.NORMAL assert decision.ratio_over is None + assert redis.hashes == {} + assert redis.expirations == {} + + +async def test_increment_feature_rpm_records_minute_bucket(): + redis = FakeRedis() + service = RedisService() + service.redis = redis + + counters = await service.increment_feature_rpm( + key_prefix="mlpa:traffic_contract", + feature="smart-window", + ttl_seconds=120, + now=125, + ) + + key = "mlpa:traffic_contract:rpm:120" + assert counters.feature_count == 1 + assert counters.basket_count == 1 assert redis.hashes[key]["smart-window"] == 1 assert redis.hashes[key][TRAFFIC_CONTRACT_BASKET_FIELD] == 1 assert redis.expirations[key] == 120 - assert ( - await service.get_current_feature_rpm( - key_prefix="mlpa:traffic_contract", - feature="smart-window", - now=125, - ) - == 1 + + +async def test_increment_feature_tpm_records_token_bucket(): + redis = FakeRedis() + service = RedisService() + service.redis = redis + + counters = await service.increment_feature_tpm( + key_prefix="mlpa:traffic_contract", + feature="smart-window", + token_count=37, + ttl_seconds=120, + now=125, ) + key = "mlpa:traffic_contract:tpm:120" + assert counters.feature_count == 37 + assert counters.basket_count == 37 + assert redis.hashes[key]["smart-window"] == 37 + assert redis.hashes[key][TRAFFIC_CONTRACT_BASKET_FIELD] == 37 + assert redis.expirations[key] == 120 + -async def test_check_and_increment_feature_rpm_marks_borrowed_over_feature_limit( - mocker, -): +async def test_check_feature_rpm_marks_borrowed_over_feature_limit(mocker): mocker.patch.object(env, "TOTAL_TRAFFIC_CONTRACT_RPM_LIMIT", 10) + redis = FakeRedis() service = RedisService() - service.redis = FakeRedis() + service.redis = redis - allowed = await service.check_and_increment_feature_rpm( + await service.increment_feature_rpm( key_prefix="mlpa:traffic_contract", feature="smart-window", - feature_rpm_limit=1, now=125, ) - rejected = await service.check_and_increment_feature_rpm( + decision = await service.check_feature_rpm( key_prefix="mlpa:traffic_contract", feature="smart-window", feature_rpm_limit=1, now=125, ) - assert allowed.allowed is True - assert rejected.allowed is False - assert rejected.limited_by == "feature" - assert rejected.mode == "borrowed" - assert rejected.ratio_over == 2.0 - assert rejected.feature_count == 2 - assert rejected.basket_count == 2 + key = "mlpa:traffic_contract:rpm:120" + assert decision.allowed is False + assert decision.limited_by == "feature" + assert decision.mode == TrafficContractMode.BORROWED + assert decision.ratio_over == 2.0 + assert decision.feature_count == 2 + assert decision.basket_count == 2 + assert redis.hashes[key]["smart-window"] == 1 + assert redis.hashes[key][TRAFFIC_CONTRACT_BASKET_FIELD] == 1 -async def test_check_and_increment_feature_rpm_marks_degraded_over_basket_limit(mocker): +async def test_check_feature_rpm_marks_degraded_over_basket_limit(mocker): mocker.patch.object(env, "TOTAL_TRAFFIC_CONTRACT_RPM_LIMIT", 1) + redis = FakeRedis() service = RedisService() - service.redis = FakeRedis() + service.redis = redis - allowed = await service.check_and_increment_feature_rpm( + await service.increment_feature_rpm( key_prefix="mlpa:traffic_contract", feature="smart-window", - feature_rpm_limit=10, now=125, ) - rejected = await service.check_and_increment_feature_rpm( + decision = await service.check_feature_rpm( key_prefix="mlpa:traffic_contract", - feature="memories", + feature="s2s", feature_rpm_limit=10, now=125, ) - assert allowed.allowed is True - assert rejected.allowed is False - assert rejected.limited_by == "basket" - assert rejected.mode == "degraded" - assert rejected.ratio_over == 2.0 - assert rejected.feature_count == 1 - assert rejected.basket_count == 2 + key = "mlpa:traffic_contract:rpm:120" + assert decision.allowed is False + assert decision.limited_by == "basket" + assert decision.mode == TrafficContractMode.DEGRADED + assert decision.ratio_over == 2.0 + assert decision.feature_count == 1 + assert decision.basket_count == 2 + assert "s2s" not in redis.hashes[key] + assert redis.hashes[key][TRAFFIC_CONTRACT_BASKET_FIELD] == 1 + + +async def test_update_contracts_increments_rpm_and_tpm_with_usage(mocker): + mocker.patch.object(env, "ENABLE_TRAFFIC_CONTRACT_ENFORCEMENT", True) + service = RedisService() + increment = mocker.patch.object( + service, + "inc_traffic_contract", + mocker.AsyncMock(), + ) + + await service.update_contracts( + service_type="ai", + usage={"total_tokens": 37}, + ) + + increment.assert_has_awaits( + [ + call( + key_type=TrafficContractKeyType.RPM, + service_type="ai", + increment_amount=1, + ), + call( + key_type=TrafficContractKeyType.TPM, + service_type="ai", + increment_amount=37, + ), + ], + any_order=True, + ) + assert increment.await_count == 2 + + +async def test_update_contracts_increments_only_rpm_without_usage(mocker): + mocker.patch.object(env, "ENABLE_TRAFFIC_CONTRACT_ENFORCEMENT", True) + service = RedisService() + increment = mocker.patch.object( + service, + "inc_traffic_contract", + mocker.AsyncMock(), + ) + + await service.update_contracts(service_type="ai", usage=None) + + increment.assert_awaited_once_with( + key_type=TrafficContractKeyType.RPM, + service_type="ai", + increment_amount=1, + ) + + +async def test_update_contracts_noops_when_disabled(mocker): + mocker.patch.object(env, "ENABLE_TRAFFIC_CONTRACT_ENFORCEMENT", False) + service = RedisService() + increment = mocker.patch.object( + service, + "inc_traffic_contract", + mocker.AsyncMock(), + ) + + await service.update_contracts( + service_type="ai", + usage={"total_tokens": 37}, + ) + + increment.assert_not_awaited() + + +async def test_inc_traffic_contract_noops_for_unknown_service_type(mocker): + mocker.patch.object(env, "ENABLE_TRAFFIC_CONTRACT_ENFORCEMENT", True) + service = RedisService() + increment = mocker.patch.object( + service, + "increment_feature_traffic_contract", + mocker.AsyncMock(), + ) + + result = await service.inc_traffic_contract( + key_type=TrafficContractKeyType.RPM, + service_type="unknown-service", + increment_amount=1, + ) + + assert result is None + increment.assert_not_awaited() diff --git a/src/tests/unit/test_traffic_contract_enforcer.py b/src/tests/unit/test_traffic_contract_enforcer.py index 3dd81e76..41f53af4 100644 --- a/src/tests/unit/test_traffic_contract_enforcer.py +++ b/src/tests/unit/test_traffic_contract_enforcer.py @@ -1,98 +1,214 @@ +from types import SimpleNamespace +from unittest.mock import call + import pytest from fastapi import HTTPException +from mlpa.core.classes import TrafficContractDecision from mlpa.core.config import env +from mlpa.core.consts import TrafficContractKeyType, TrafficContractMode from mlpa.core.middleware import traffic_contract_enforcer -from mlpa.core.services.redis_service import TrafficContractDecision from tests.consts import SAMPLE_REQUEST -async def test_chat_traffic_contract_noops_when_disabled(mocker): +def _request_with_state(): + return SimpleNamespace(state=SimpleNamespace()) + + +def _decision( + *, + allowed: bool = True, + limited_by: str | None = None, + mode: TrafficContractMode = TrafficContractMode.NORMAL, + ratio_over: float | None = None, +) -> TrafficContractDecision: + return TrafficContractDecision( + allowed=allowed, + feature_count=1, + basket_count=1, + retry_after_seconds=60, + limited_by=limited_by, + mode=mode, + ratio_over=ratio_over, + ) + + +async def test_traffic_contract_noops_when_disabled(mocker): mocker.patch.object(env, "ENABLE_TRAFFIC_CONTRACT_ENFORCEMENT", False) check = mocker.patch.object( traffic_contract_enforcer.redis_service, - "check_and_increment_feature_rpm", + "check_feature_traffic_contract", mocker.AsyncMock(), ) + increment = mocker.patch.object( + traffic_contract_enforcer.redis_service, + "inc_traffic_contract", + mocker.AsyncMock(), + ) + request = _request_with_state() - await traffic_contract_enforcer.enforce_chat_traffic_contract(SAMPLE_REQUEST) + await traffic_contract_enforcer.enforce_traffic_contract( + request, SAMPLE_REQUEST.service_type + ) + assert request.state.traffic_contract_rpm_mode == "N/A" check.assert_not_awaited() + increment.assert_not_awaited() -async def test_chat_traffic_contract_increments_feature_rpm_when_enabled(mocker): +async def test_traffic_contract_checks_feature_rpm_and_tpm_when_enabled( + mocker, +): mocker.patch.object(env, "ENABLE_TRAFFIC_CONTRACT_ENFORCEMENT", True) check = mocker.patch.object( traffic_contract_enforcer.redis_service, - "check_and_increment_feature_rpm", - mocker.AsyncMock( - return_value=TrafficContractDecision( - allowed=True, - feature_count=1, - basket_count=1, - retry_after_seconds=60, - ) - ), + "check_feature_traffic_contract", + mocker.AsyncMock(side_effect=[_decision(), _decision()]), ) + update_contracts = mocker.patch.object( + traffic_contract_enforcer.redis_service, + "inc_traffic_contract", + mocker.AsyncMock(), + ) + request = _request_with_state() - await traffic_contract_enforcer.enforce_chat_traffic_contract(SAMPLE_REQUEST) + await traffic_contract_enforcer.enforce_traffic_contract( + request, SAMPLE_REQUEST.service_type + ) contract = env.traffic_contract_config[SAMPLE_REQUEST.service_type] - check.assert_awaited_once_with( - key_prefix=env.TRAFFIC_CONTRACT_REDIS_KEY_PREFIX, - feature=contract["feature"], - feature_rpm_limit=contract["rpm_limit"], - window_seconds=env.TRAFFIC_CONTRACT_RPM_WINDOW_SECONDS, - ttl_seconds=env.TRAFFIC_CONTRACT_COUNTER_TTL_SECONDS, + assert check.await_args_list == [ + call( + key_prefix=env.TRAFFIC_CONTRACT_REDIS_KEY_PREFIX, + key_type=TrafficContractKeyType.RPM, + feature=contract["feature"], + feature_limit=contract["rpm_limit"], + basket_limit=env.TOTAL_TRAFFIC_CONTRACT_RPM_LIMIT, + increment_amount=1, + window_seconds=env.TRAFFIC_CONTRACT_RPM_WINDOW_SECONDS, + ), + call( + key_prefix=env.TRAFFIC_CONTRACT_REDIS_KEY_PREFIX, + key_type=TrafficContractKeyType.TPM, + feature=contract["feature"], + feature_limit=contract["tpm_limit"], + basket_limit=env.TOTAL_TRAFFIC_CONTRACT_TPM_LIMIT, + increment_amount=0, + window_seconds=env.TRAFFIC_CONTRACT_TPM_WINDOW_SECONDS, + ), + ] + update_contracts.assert_not_awaited() + assert request.state.traffic_contract_rpm_mode == "normal" + + +async def test_traffic_contract_records_borrowed_mode_when_feature_is_over_contract( + mocker, +): + mocker.patch.object(env, "ENABLE_TRAFFIC_CONTRACT_ENFORCEMENT", True) + mocker.patch.object( + traffic_contract_enforcer.redis_service, + "check_feature_traffic_contract", + mocker.AsyncMock( + side_effect=[ + _decision( + allowed=False, + limited_by="feature", + mode=TrafficContractMode.BORROWED, + ratio_over=1.1, + ), + _decision(), + ] + ), + ) + update_contracts = mocker.patch.object( + traffic_contract_enforcer.redis_service, + "inc_traffic_contract", + mocker.AsyncMock(), ) + request = _request_with_state() + + await traffic_contract_enforcer.enforce_traffic_contract( + request, SAMPLE_REQUEST.service_type + ) + + update_contracts.assert_not_awaited() + assert request.state.traffic_contract_rpm_mode == "borrowed" -async def test_chat_traffic_contract_allows_soft_over_rpm(mocker): +async def test_traffic_contract_records_rpm_and_tpm_modes(mocker): mocker.patch.object(env, "ENABLE_TRAFFIC_CONTRACT_ENFORCEMENT", True) - check = mocker.patch.object( + mocker.patch.object( traffic_contract_enforcer.redis_service, - "check_and_increment_feature_rpm", + "check_feature_traffic_contract", mocker.AsyncMock( - return_value=TrafficContractDecision( - allowed=False, - feature_count=10_000, - basket_count=10_000, - retry_after_seconds=17, - limited_by="feature", - mode="borrowed", - ratio_over=1.1, - ) + side_effect=[ + _decision( + allowed=False, + limited_by="feature", + mode=TrafficContractMode.BORROWED, + ratio_over=1.1, + ), + _decision( + allowed=False, + limited_by="basket", + mode=TrafficContractMode.DEGRADED, + ratio_over=1.2, + ), + ] ), ) + mocker.patch.object( + traffic_contract_enforcer.redis_service, + "inc_traffic_contract", + mocker.AsyncMock(), + ) + request = _request_with_state() - await traffic_contract_enforcer.enforce_chat_traffic_contract(SAMPLE_REQUEST) + await traffic_contract_enforcer.enforce_traffic_contract( + request, SAMPLE_REQUEST.service_type + ) - check.assert_awaited_once() + assert request.state.traffic_contract_rpm_mode == "borrowed" + assert request.state.traffic_contract_tpm_mode == "degraded" -async def test_chat_traffic_contract_can_fail_open_on_redis_error(mocker): +async def test_traffic_contract_can_fail_open_on_redis_error(mocker): mocker.patch.object(env, "ENABLE_TRAFFIC_CONTRACT_ENFORCEMENT", True) mocker.patch.object(env, "TRAFFIC_CONTRACT_FAIL_OPEN_ON_REDIS_ERROR", True) mocker.patch.object( traffic_contract_enforcer.redis_service, - "check_and_increment_feature_rpm", + "check_feature_traffic_contract", mocker.AsyncMock(side_effect=RuntimeError("redis down")), ) + increment = mocker.patch.object( + traffic_contract_enforcer.redis_service, + "inc_traffic_contract", + mocker.AsyncMock(), + ) + request = _request_with_state() + + await traffic_contract_enforcer.enforce_traffic_contract( + request, SAMPLE_REQUEST.service_type + ) - await traffic_contract_enforcer.enforce_chat_traffic_contract(SAMPLE_REQUEST) + assert request.state.traffic_contract_rpm_mode == "N/A" + increment.assert_not_awaited() -async def test_chat_traffic_contract_fails_closed_on_redis_error(mocker): +async def test_traffic_contract_fails_closed_on_redis_error(mocker): mocker.patch.object(env, "ENABLE_TRAFFIC_CONTRACT_ENFORCEMENT", True) mocker.patch.object(env, "TRAFFIC_CONTRACT_FAIL_OPEN_ON_REDIS_ERROR", False) mocker.patch.object( traffic_contract_enforcer.redis_service, - "check_and_increment_feature_rpm", + "check_feature_traffic_contract", mocker.AsyncMock(side_effect=RuntimeError("redis down")), ) + request = _request_with_state() with pytest.raises(HTTPException) as exc_info: - await traffic_contract_enforcer.enforce_chat_traffic_contract(SAMPLE_REQUEST) + await traffic_contract_enforcer.enforce_traffic_contract( + request, SAMPLE_REQUEST.service_type + ) assert exc_info.value.status_code == 503 assert exc_info.value.detail == { From 86627b843fc6c27cd9d15bf90df668b409e54be5 Mon Sep 17 00:00:00 2001 From: noahpodgurski Date: Wed, 16 Sep 2026 09:11:00 -0400 Subject: [PATCH 07/10] addr comments, add testing for missing service type config patterns --- src/mlpa/core/completions.py | 13 ++- src/mlpa/core/config.py | 1 + src/mlpa/core/services/redis_service.py | 111 ++++----------------- src/tests/unit/test_config.py | 51 +++++++++- src/tests/unit/test_redis_service.py | 123 ------------------------ 5 files changed, 77 insertions(+), 222 deletions(-) diff --git a/src/mlpa/core/completions.py b/src/mlpa/core/completions.py index 9ed8ff19..73fa3e64 100644 --- a/src/mlpa/core/completions.py +++ b/src/mlpa/core/completions.py @@ -333,8 +333,10 @@ async def _read_next_chunk( authorized_chat_request, result, time.perf_counter() - start_time ) record_chat_availability(authorized_chat_request, availability_reason) - await redis_service.update_contracts( - service_type=authorized_chat_request.service_type, usage=usage + asyncio.create_task( + redis_service.update_contracts( + service_type=authorized_chat_request.service_type, usage=usage + ) ) @@ -424,6 +426,9 @@ async def _get_completion(authorized_chat_request: AuthorizedChatRequest): authorized_chat_request, result, time.perf_counter() - start_time ) record_chat_availability(authorized_chat_request, availability_reason) - await redis_service.update_contracts( - service_type=authorized_chat_request.service_type, usage=usage + asyncio.create_task( + redis_service.update_contracts( + service_type=authorized_chat_request.service_type, + usage=usage, + ) ) diff --git a/src/mlpa/core/config.py b/src/mlpa/core/config.py index 98b16164..e9527126 100644 --- a/src/mlpa/core/config.py +++ b/src/mlpa/core/config.py @@ -243,6 +243,7 @@ def service_type_config(self) -> dict[str, dict]: "tpm_limit": self.USER_FEATURE_BUDGET_ANSWER_TPM_LIMIT, }, "sw-answer": { + "feature": self.FEATURE_SMART_WINDOW, "budget_id": self.USER_FEATURE_BUDGET_SW_ANSWER_BUDGET_ID, "max_budget": self.USER_FEATURE_BUDGET_SW_ANSWER_MAX_BUDGET, "rpm_limit": self.USER_FEATURE_BUDGET_SW_ANSWER_RPM_LIMIT, diff --git a/src/mlpa/core/services/redis_service.py b/src/mlpa/core/services/redis_service.py index 71ec14ec..eb00d352 100644 --- a/src/mlpa/core/services/redis_service.py +++ b/src/mlpa/core/services/redis_service.py @@ -189,86 +189,6 @@ async def increment_feature_traffic_contract( basket_count=int(result[1]), ) - async def check_feature_rpm( - self, - *, - key_prefix: str, - feature: str, - feature_rpm_limit: int, - window_seconds: int = 60, - now: int | None = None, - ) -> TrafficContractDecision: - return await self.check_feature_traffic_contract( - key_prefix=key_prefix, - key_type=TrafficContractKeyType.RPM, - feature=feature, - feature_limit=feature_rpm_limit, - basket_limit=env.TOTAL_TRAFFIC_CONTRACT_RPM_LIMIT, - increment_amount=1, - window_seconds=window_seconds, - now=now, - ) - - async def increment_feature_rpm( - self, - *, - key_prefix: str, - feature: str, - window_seconds: int = 60, - ttl_seconds: int = 120, - now: int | None = None, - ) -> TrafficContractCounters: - return await self.increment_feature_traffic_contract( - key_prefix=key_prefix, - key_type=TrafficContractKeyType.RPM, - feature=feature, - increment_amount=1, - window_seconds=window_seconds, - ttl_seconds=ttl_seconds, - now=now, - ) - - async def check_feature_tpm( - self, - *, - key_prefix: str, - feature: str, - feature_tpm_limit: int, - token_count: int = 0, - window_seconds: int = 60, - now: int | None = None, - ) -> TrafficContractDecision: - return await self.check_feature_traffic_contract( - key_prefix=key_prefix, - key_type=TrafficContractKeyType.TPM, - feature=feature, - feature_limit=feature_tpm_limit, - basket_limit=env.TOTAL_TRAFFIC_CONTRACT_TPM_LIMIT, - increment_amount=token_count, - window_seconds=window_seconds, - now=now, - ) - - async def increment_feature_tpm( - self, - *, - key_prefix: str, - feature: str, - token_count: int, - window_seconds: int = 60, - ttl_seconds: int = 120, - now: int | None = None, - ) -> TrafficContractCounters: - return await self.increment_feature_traffic_contract( - key_prefix=key_prefix, - key_type=TrafficContractKeyType.TPM, - feature=feature, - increment_amount=token_count, - window_seconds=window_seconds, - ttl_seconds=ttl_seconds, - now=now, - ) - async def inc_traffic_contract( self, *, @@ -301,20 +221,25 @@ async def update_contracts(self, *, service_type: str, usage: dict | None): if not env.ENABLE_TRAFFIC_CONTRACT_ENFORCEMENT: return - updates = [ - self.inc_traffic_contract( - key_type=TrafficContractKeyType.RPM, - service_type=service_type, - increment_amount=1, - ) - ] - if usage and usage.get("total_tokens"): - updates.append( + try: + updates = [ self.inc_traffic_contract( - key_type=TrafficContractKeyType.TPM, + key_type=TrafficContractKeyType.RPM, service_type=service_type, - increment_amount=usage["total_tokens"], + increment_amount=1, + ) + ] + if usage and usage.get("total_tokens"): + updates.append( + self.inc_traffic_contract( + key_type=TrafficContractKeyType.TPM, + service_type=service_type, + increment_amount=usage["total_tokens"], + ) ) - ) - await asyncio.gather(*updates) + await asyncio.gather(*updates) + except Exception as e: + logger.error(f"Error updating traffic contracts for {service_type}: {e}") + if not env.TRAFFIC_CONTRACT_FAIL_OPEN_ON_REDIS_ERROR: + raise diff --git a/src/tests/unit/test_config.py b/src/tests/unit/test_config.py index 0b6da776..2b494f3a 100644 --- a/src/tests/unit/test_config.py +++ b/src/tests/unit/test_config.py @@ -93,8 +93,9 @@ def test_user_feature_budget_liner_answer_from_env(): def test_user_feature_budget_sw_answer_default_values(): """Test that sw-answer budget configuration has correct default values.""" env = Env() - sw_answer_config = env.user_feature_budget["sw-answer"] + sw_answer_config = env.service_type_config["sw-answer"] + assert sw_answer_config["feature"] == env.FEATURE_SMART_WINDOW assert sw_answer_config["budget_id"] == "end-user-budget-sw-answer" assert sw_answer_config["max_budget"] == 0.1 assert sw_answer_config["rpm_limit"] == 10 @@ -115,7 +116,8 @@ def test_user_feature_budget_sw_answer_from_env(): with patch.dict(os.environ, env_vars): env = Env() - sw_answer_config = env.user_feature_budget["sw-answer"] + sw_answer_config = env.service_type_config["sw-answer"] + assert sw_answer_config["feature"] == env.FEATURE_SMART_WINDOW assert sw_answer_config["budget_id"] == "custom-sw-answer-budget-id" assert sw_answer_config["max_budget"] == 0.5 assert sw_answer_config["rpm_limit"] == 20 @@ -319,6 +321,51 @@ def test_user_feature_budget_structure_consistency(): ) +def test_service_type_config_values_are_defined(): + """Every service type must have complete feature, budget, and rate config.""" + env = Env() + required_keys = { + "feature", + "budget_id", + "budget_duration", + "max_budget", + "rpm_limit", + "tpm_limit", + } + + for service_type, config in env.service_type_config.items(): + assert set(config.keys()) == required_keys, ( + f"{service_type} service_type_config keys differ from required keys" + ) + assert isinstance(config["feature"], str) + assert config["feature"], f"{service_type} feature must be defined" + assert isinstance(config["budget_id"], str) + assert config["budget_id"], f"{service_type} budget_id must be defined" + assert isinstance(config["budget_duration"], str) + assert config["budget_duration"], ( + f"{service_type} budget_duration must be defined" + ) + assert isinstance(config["max_budget"], float) + assert config["max_budget"] >= 0 + assert isinstance(config["rpm_limit"], int) + assert config["rpm_limit"] >= 0 + assert isinstance(config["tpm_limit"], int) + assert config["tpm_limit"] >= 0 + + +def test_traffic_contract_config_defined_for_each_service_type(): + env = Env() + + assert set(env.traffic_contract_config) == set(env.service_type_config) + for service_type, contract in env.traffic_contract_config.items(): + service_config = env.service_type_config[service_type] + assert contract["feature"] == service_config["feature"] + assert isinstance(contract["rpm_limit"], int) + assert contract["rpm_limit"] >= 0 + assert isinstance(contract["tpm_limit"], int) + assert contract["tpm_limit"] >= 0 + + def test_user_feature_budget_memories_type_validation(): """Test that memories budget configuration values have correct types.""" env = Env() diff --git a/src/tests/unit/test_redis_service.py b/src/tests/unit/test_redis_service.py index 30eaaffd..31f15f09 100644 --- a/src/tests/unit/test_redis_service.py +++ b/src/tests/unit/test_redis_service.py @@ -73,129 +73,6 @@ async def hget(self, key, field): return self.hashes.get(key, {}).get(field) -async def test_check_feature_rpm_does_not_mutate_bucket(mocker): - mocker.patch.object(env, "TOTAL_TRAFFIC_CONTRACT_RPM_LIMIT", 20) - redis = FakeRedis() - service = RedisService() - service.redis = redis - - decision = await service.check_feature_rpm( - key_prefix="mlpa:traffic_contract", - feature="smart-window", - feature_rpm_limit=10, - now=125, - ) - - assert decision.allowed is True - assert decision.feature_count == 1 - assert decision.basket_count == 1 - assert decision.retry_after_seconds == 55 - assert decision.limited_by is None - assert decision.mode == TrafficContractMode.NORMAL - assert decision.ratio_over is None - assert redis.hashes == {} - assert redis.expirations == {} - - -async def test_increment_feature_rpm_records_minute_bucket(): - redis = FakeRedis() - service = RedisService() - service.redis = redis - - counters = await service.increment_feature_rpm( - key_prefix="mlpa:traffic_contract", - feature="smart-window", - ttl_seconds=120, - now=125, - ) - - key = "mlpa:traffic_contract:rpm:120" - assert counters.feature_count == 1 - assert counters.basket_count == 1 - assert redis.hashes[key]["smart-window"] == 1 - assert redis.hashes[key][TRAFFIC_CONTRACT_BASKET_FIELD] == 1 - assert redis.expirations[key] == 120 - - -async def test_increment_feature_tpm_records_token_bucket(): - redis = FakeRedis() - service = RedisService() - service.redis = redis - - counters = await service.increment_feature_tpm( - key_prefix="mlpa:traffic_contract", - feature="smart-window", - token_count=37, - ttl_seconds=120, - now=125, - ) - - key = "mlpa:traffic_contract:tpm:120" - assert counters.feature_count == 37 - assert counters.basket_count == 37 - assert redis.hashes[key]["smart-window"] == 37 - assert redis.hashes[key][TRAFFIC_CONTRACT_BASKET_FIELD] == 37 - assert redis.expirations[key] == 120 - - -async def test_check_feature_rpm_marks_borrowed_over_feature_limit(mocker): - mocker.patch.object(env, "TOTAL_TRAFFIC_CONTRACT_RPM_LIMIT", 10) - redis = FakeRedis() - service = RedisService() - service.redis = redis - - await service.increment_feature_rpm( - key_prefix="mlpa:traffic_contract", - feature="smart-window", - now=125, - ) - decision = await service.check_feature_rpm( - key_prefix="mlpa:traffic_contract", - feature="smart-window", - feature_rpm_limit=1, - now=125, - ) - - key = "mlpa:traffic_contract:rpm:120" - assert decision.allowed is False - assert decision.limited_by == "feature" - assert decision.mode == TrafficContractMode.BORROWED - assert decision.ratio_over == 2.0 - assert decision.feature_count == 2 - assert decision.basket_count == 2 - assert redis.hashes[key]["smart-window"] == 1 - assert redis.hashes[key][TRAFFIC_CONTRACT_BASKET_FIELD] == 1 - - -async def test_check_feature_rpm_marks_degraded_over_basket_limit(mocker): - mocker.patch.object(env, "TOTAL_TRAFFIC_CONTRACT_RPM_LIMIT", 1) - redis = FakeRedis() - service = RedisService() - service.redis = redis - - await service.increment_feature_rpm( - key_prefix="mlpa:traffic_contract", - feature="smart-window", - now=125, - ) - decision = await service.check_feature_rpm( - key_prefix="mlpa:traffic_contract", - feature="s2s", - feature_rpm_limit=10, - now=125, - ) - - key = "mlpa:traffic_contract:rpm:120" - assert decision.allowed is False - assert decision.limited_by == "basket" - assert decision.mode == TrafficContractMode.DEGRADED - assert decision.ratio_over == 2.0 - assert decision.feature_count == 1 - assert decision.basket_count == 2 - assert "s2s" not in redis.hashes[key] - assert redis.hashes[key][TRAFFIC_CONTRACT_BASKET_FIELD] == 1 - - async def test_update_contracts_increments_rpm_and_tpm_with_usage(mocker): mocker.patch.object(env, "ENABLE_TRAFFIC_CONTRACT_ENFORCEMENT", True) service = RedisService() From 5360d6e575d4751a8b344f81f21e0d32304afb20 Mon Sep 17 00:00:00 2001 From: noahpodgurski Date: Wed, 16 Sep 2026 09:35:13 -0400 Subject: [PATCH 08/10] refactor: update traffic contract handling to use unified check_feature_traffic_contracts method --- src/mlpa/core/middleware/instrumentation.py | 2 +- .../middleware/traffic_contract_enforcer.py | 37 ++-- src/mlpa/core/services/redis_service.py | 124 ++++++++----- .../integration/test_traffic_contracts.py | 26 ++- src/tests/unit/test_redis_service.py | 164 +++++++++++++++--- .../unit/test_traffic_contract_enforcer.py | 57 +++--- 6 files changed, 274 insertions(+), 136 deletions(-) diff --git a/src/mlpa/core/middleware/instrumentation.py b/src/mlpa/core/middleware/instrumentation.py index 1abc8c7c..f47df878 100644 --- a/src/mlpa/core/middleware/instrumentation.py +++ b/src/mlpa/core/middleware/instrumentation.py @@ -24,7 +24,7 @@ def _traffic_contract_mode_label( "N/A", ) if isinstance(mode, TrafficContractMode): - return mode.value + mode = mode.value if isinstance(mode, str) and mode in {mode.value for mode in TrafficContractMode}: return mode return "N/A" diff --git a/src/mlpa/core/middleware/traffic_contract_enforcer.py b/src/mlpa/core/middleware/traffic_contract_enforcer.py index cb4d141a..9ca2aa52 100644 --- a/src/mlpa/core/middleware/traffic_contract_enforcer.py +++ b/src/mlpa/core/middleware/traffic_contract_enforcer.py @@ -1,5 +1,3 @@ -import asyncio - from fastapi import HTTPException, Request from mlpa.core.classes import ( @@ -7,7 +5,7 @@ AuthorizedSearchRequest, ) from mlpa.core.config import env -from mlpa.core.consts import TrafficContractKeyType, TrafficContractMode +from mlpa.core.consts import TrafficContractMode from mlpa.core.logger import logger from mlpa.core.services.services import redis_service @@ -38,25 +36,20 @@ async def enforce_traffic_contract( return try: - rpm_decision, tpm_decision = await asyncio.gather( - redis_service.check_feature_traffic_contract( - key_prefix=env.TRAFFIC_CONTRACT_REDIS_KEY_PREFIX, - key_type=TrafficContractKeyType.RPM, - feature=contract["feature"], - feature_limit=contract["rpm_limit"], - basket_limit=env.TOTAL_TRAFFIC_CONTRACT_RPM_LIMIT, - increment_amount=1, - window_seconds=env.TRAFFIC_CONTRACT_RPM_WINDOW_SECONDS, - ), - redis_service.check_feature_traffic_contract( - key_prefix=env.TRAFFIC_CONTRACT_REDIS_KEY_PREFIX, - key_type=TrafficContractKeyType.TPM, - feature=contract["feature"], - feature_limit=contract["tpm_limit"], - basket_limit=env.TOTAL_TRAFFIC_CONTRACT_TPM_LIMIT, - increment_amount=0, # Token usage is incremented after the response. - window_seconds=env.TRAFFIC_CONTRACT_TPM_WINDOW_SECONDS, - ), + ( + rpm_decision, + tpm_decision, + ) = await redis_service.check_feature_traffic_contracts( + key_prefix=env.TRAFFIC_CONTRACT_REDIS_KEY_PREFIX, + feature=contract["feature"], + rpm_limit=contract["rpm_limit"], + tpm_limit=contract["tpm_limit"], + rpm_basket_limit=env.TOTAL_TRAFFIC_CONTRACT_RPM_LIMIT, + tpm_basket_limit=env.TOTAL_TRAFFIC_CONTRACT_TPM_LIMIT, + rpm_increment_amount=1, + tpm_increment_amount=0, # Token usage is incremented after the response. + rpm_window_seconds=env.TRAFFIC_CONTRACT_RPM_WINDOW_SECONDS, + tpm_window_seconds=env.TRAFFIC_CONTRACT_TPM_WINDOW_SECONDS, ) _set_traffic_contract_modes( diff --git a/src/mlpa/core/services/redis_service.py b/src/mlpa/core/services/redis_service.py index eb00d352..69a678ef 100644 --- a/src/mlpa/core/services/redis_service.py +++ b/src/mlpa/core/services/redis_service.py @@ -11,30 +11,30 @@ TRAFFIC_CONTRACT_BASKET_FIELD = "__basket__" -_CHECK_TRAFFIC_CONTRACT_SCRIPT = """ -local key = KEYS[1] -local feature_field = ARGV[1] -local basket_field = ARGV[2] -local feature_limit = tonumber(ARGV[3]) -local basket_limit = tonumber(ARGV[4]) -local inc_amount = tonumber(ARGV[5]) +_CHECK_TRAFFIC_CONTRACTS_SCRIPT = """ +local function check(key, feature_field, basket_field, feature_limit, basket_limit, inc_amount) + if basket_limit == 0 and feature_limit == 0 then + return {1, 0, 0, "", "normal", ""} + end -if basket_limit == 0 and feature_limit == 0 then - return {1, 0, 0, "", "normal", ""} -end + local feature_count = tonumber(redis.call("HGET", key, feature_field) or "0") + inc_amount + local basket_count = tonumber(redis.call("HGET", key, basket_field) or "0") + inc_amount -local feature_count = tonumber(redis.call("HGET", key, feature_field) or "0") + inc_amount -local basket_count = tonumber(redis.call("HGET", key, basket_field) or "0") + inc_amount + if basket_limit > 0 and basket_count > basket_limit then + return {0, feature_count, basket_count, "basket", "degraded", tostring(basket_count/basket_limit)} + end -if basket_limit > 0 and basket_count > basket_limit then - return {0, feature_count, basket_count, "basket", "degraded", tostring(basket_count/basket_limit)} -end + if feature_limit > 0 and feature_count > feature_limit then + return {0, feature_count, basket_count, "feature", "borrowed", tostring(feature_count/feature_limit)} + end -if feature_limit > 0 and feature_count > feature_limit then - return {0, feature_count, basket_count, "feature", "borrowed", tostring(feature_count/feature_limit)} + return {1, feature_count, basket_count, "", "normal", ""} end -return {1, feature_count, basket_count, "", "normal", ""} +local rpm = check(KEYS[1], ARGV[1], ARGV[2], tonumber(ARGV[3]), tonumber(ARGV[4]), tonumber(ARGV[5])) +local tpm = check(KEYS[2], ARGV[1], ARGV[2], tonumber(ARGV[6]), tonumber(ARGV[7]), tonumber(ARGV[8])) + +return {rpm, tpm} """ _INCREMENT_TRAFFIC_CONTRACT_SCRIPT = """ @@ -113,48 +113,82 @@ def retry_after_seconds( elapsed_in_bucket = current_time % window_seconds return window_seconds - elapsed_in_bucket - async def check_feature_traffic_contract( + @staticmethod + def traffic_contract_decision_from_result( + result: list, + *, + retry_after_seconds: int, + ) -> TrafficContractDecision: + return TrafficContractDecision( + allowed=bool(int(result[0])), + feature_count=int(result[1]), + basket_count=int(result[2]), + retry_after_seconds=retry_after_seconds, + limited_by=str(result[3]) or None, + mode=TrafficContractMode(result[4]), + ratio_over=float(result[5]) if result[5] not in (None, "") else None, + ) + + async def check_feature_traffic_contracts( self, *, key_prefix: str, - key_type: TrafficContractKeyType, feature: str, - feature_limit: int, - basket_limit: int, - increment_amount: int, - window_seconds: int = 60, + rpm_limit: int, + tpm_limit: int, + rpm_basket_limit: int, + tpm_basket_limit: int, + rpm_increment_amount: int, + tpm_increment_amount: int, + rpm_window_seconds: int = 60, + tpm_window_seconds: int = 60, now: int | None = None, - ) -> TrafficContractDecision: - key = self.traffic_contract_key( + ) -> tuple[TrafficContractDecision, TrafficContractDecision]: + rpm_key = self.traffic_contract_key( key_prefix=key_prefix, - key_type=key_type, + key_type=TrafficContractKeyType.RPM, now=now, - window_seconds=window_seconds, + window_seconds=rpm_window_seconds, ) - retry_after = self.retry_after_seconds( + tpm_key = self.traffic_contract_key( + key_prefix=key_prefix, + key_type=TrafficContractKeyType.TPM, now=now, - window_seconds=window_seconds, + window_seconds=tpm_window_seconds, + ) + rpm_retry_after = self.retry_after_seconds( + now=now, + window_seconds=rpm_window_seconds, + ) + tpm_retry_after = self.retry_after_seconds( + now=now, + window_seconds=tpm_window_seconds, ) - result = await self.client.eval( - _CHECK_TRAFFIC_CONTRACT_SCRIPT, - 1, - key, + rpm_result, tpm_result = await self.client.eval( + _CHECK_TRAFFIC_CONTRACTS_SCRIPT, + 2, + rpm_key, + tpm_key, feature, TRAFFIC_CONTRACT_BASKET_FIELD, - feature_limit, - basket_limit, - increment_amount, + rpm_limit, + rpm_basket_limit, + rpm_increment_amount, + tpm_limit, + tpm_basket_limit, + tpm_increment_amount, ) - return TrafficContractDecision( - allowed=bool(int(result[0])), - feature_count=int(result[1]), - basket_count=int(result[2]), - retry_after_seconds=retry_after, - limited_by=str(result[3]) or None, - mode=TrafficContractMode(result[4]), - ratio_over=float(result[5]) if result[5] not in (None, "") else None, + return ( + self.traffic_contract_decision_from_result( + rpm_result, + retry_after_seconds=rpm_retry_after, + ), + self.traffic_contract_decision_from_result( + tpm_result, + retry_after_seconds=tpm_retry_after, + ), ) async def increment_feature_traffic_contract( diff --git a/src/tests/integration/test_traffic_contracts.py b/src/tests/integration/test_traffic_contracts.py index db99ddd9..21a4d32f 100644 --- a/src/tests/integration/test_traffic_contracts.py +++ b/src/tests/integration/test_traffic_contracts.py @@ -11,16 +11,24 @@ def test_chat_completion_continues_when_traffic_contract_exceeded( mocker.patch.object(env, "ENABLE_TRAFFIC_CONTRACT_ENFORCEMENT", True) mocker.patch.object( traffic_contract_enforcer.redis_service, - "check_feature_traffic_contract", + "check_feature_traffic_contracts", mocker.AsyncMock( - return_value=TrafficContractDecision( - allowed=False, - feature_count=10_000, - basket_count=10_000, - retry_after_seconds=42, - limited_by="feature", - mode="borrowed", - ratio_over=1.1, + return_value=( + TrafficContractDecision( + allowed=False, + feature_count=10_000, + basket_count=10_000, + retry_after_seconds=42, + limited_by="feature", + mode="borrowed", + ratio_over=1.1, + ), + TrafficContractDecision( + allowed=True, + feature_count=0, + basket_count=0, + retry_after_seconds=42, + ), ) ), ) diff --git a/src/tests/unit/test_redis_service.py b/src/tests/unit/test_redis_service.py index 31f15f09..9dbeae3f 100644 --- a/src/tests/unit/test_redis_service.py +++ b/src/tests/unit/test_redis_service.py @@ -12,33 +12,10 @@ class FakeRedis: def __init__(self): self.hashes = {} self.expirations = {} + self.eval_calls = [] - async def eval( - self, - script, - num_keys, - key, - *args, - ): - assert num_keys == 1 - if "HINCRBY" in script: - feature, basket_field, ttl_seconds, inc_amount = args - bucket = self.hashes.setdefault(key, {}) - if int(inc_amount) <= 0: - return [ - int(bucket.get(feature, 0)), - int(bucket.get(basket_field, 0)), - ] - - feature_count = int(bucket.get(feature, 0)) + int(inc_amount) - basket_count = int(bucket.get(basket_field, 0)) + int(inc_amount) - bucket[feature] = feature_count - bucket[basket_field] = basket_count - self.expirations[key] = int(ttl_seconds) - return [feature_count, basket_count] - - feature, basket_field, feature_limit, basket_limit, inc_amount = args - bucket = self.hashes.get(key, {}) + @staticmethod + def _check(bucket, feature, basket_field, feature_limit, basket_limit, inc_amount): feature_limit = int(feature_limit) basket_limit = int(basket_limit) if basket_limit == 0 and feature_limit == 0: @@ -69,6 +46,72 @@ async def eval( return [1, feature_count, basket_count, "", "normal", ""] + async def eval( + self, + script, + num_keys, + key, + *args, + ): + self.eval_calls.append((num_keys, key, args)) + if num_keys == 2: + tpm_key = args[0] + ( + feature, + basket_field, + rpm_feature_limit, + rpm_basket_limit, + rpm_inc_amount, + tpm_feature_limit, + tpm_basket_limit, + tpm_inc_amount, + ) = args[1:] + return [ + self._check( + self.hashes.get(key, {}), + feature, + basket_field, + rpm_feature_limit, + rpm_basket_limit, + rpm_inc_amount, + ), + self._check( + self.hashes.get(tpm_key, {}), + feature, + basket_field, + tpm_feature_limit, + tpm_basket_limit, + tpm_inc_amount, + ), + ] + + assert num_keys == 1 + if "HINCRBY" in script: + feature, basket_field, ttl_seconds, inc_amount = args + bucket = self.hashes.setdefault(key, {}) + if int(inc_amount) <= 0: + return [ + int(bucket.get(feature, 0)), + int(bucket.get(basket_field, 0)), + ] + + feature_count = int(bucket.get(feature, 0)) + int(inc_amount) + basket_count = int(bucket.get(basket_field, 0)) + int(inc_amount) + bucket[feature] = feature_count + bucket[basket_field] = basket_count + self.expirations[key] = int(ttl_seconds) + return [feature_count, basket_count] + + feature, basket_field, feature_limit, basket_limit, inc_amount = args + return self._check( + self.hashes.get(key, {}), + feature, + basket_field, + feature_limit, + basket_limit, + inc_amount, + ) + async def hget(self, key, field): return self.hashes.get(key, {}).get(field) @@ -157,3 +200,72 @@ async def test_inc_traffic_contract_noops_for_unknown_service_type(mocker): assert result is None increment.assert_not_awaited() + + +async def test_check_feature_traffic_contracts_returns_rpm_and_tpm_decisions(mocker): + mocker.patch.object(env, "TOTAL_TRAFFIC_CONTRACT_RPM_LIMIT", 5) + mocker.patch.object(env, "TOTAL_TRAFFIC_CONTRACT_TPM_LIMIT", 10) + redis = FakeRedis() + service = RedisService() + service.redis = redis + + await service.increment_feature_traffic_contract( + key_prefix="mlpa:traffic_contract", + key_type=TrafficContractKeyType.RPM, + feature="smart-window", + increment_amount=2, + now=125, + ) + await service.increment_feature_traffic_contract( + key_prefix="mlpa:traffic_contract", + key_type=TrafficContractKeyType.TPM, + feature="smart-window", + increment_amount=9, + now=125, + ) + + rpm_decision, tpm_decision = await service.check_feature_traffic_contracts( + key_prefix="mlpa:traffic_contract", + feature="smart-window", + rpm_limit=2, + tpm_limit=20, + rpm_basket_limit=env.TOTAL_TRAFFIC_CONTRACT_RPM_LIMIT, + tpm_basket_limit=env.TOTAL_TRAFFIC_CONTRACT_TPM_LIMIT, + rpm_increment_amount=1, + tpm_increment_amount=2, + now=125, + ) + + rpm_key = "mlpa:traffic_contract:rpm:120" + tpm_key = "mlpa:traffic_contract:tpm:120" + assert rpm_decision.allowed is False + assert rpm_decision.limited_by == "feature" + assert rpm_decision.mode == TrafficContractMode.BORROWED + assert rpm_decision.feature_count == 3 + assert rpm_decision.basket_count == 3 + assert rpm_decision.retry_after_seconds == 55 + assert tpm_decision.allowed is False + assert tpm_decision.limited_by == "basket" + assert tpm_decision.mode == TrafficContractMode.DEGRADED + assert tpm_decision.feature_count == 11 + assert tpm_decision.basket_count == 11 + assert tpm_decision.retry_after_seconds == 55 + assert redis.hashes[rpm_key]["smart-window"] == 2 + assert redis.hashes[rpm_key][TRAFFIC_CONTRACT_BASKET_FIELD] == 2 + assert redis.hashes[tpm_key]["smart-window"] == 9 + assert redis.hashes[tpm_key][TRAFFIC_CONTRACT_BASKET_FIELD] == 9 + assert redis.eval_calls[-1] == ( + 2, + rpm_key, + ( + tpm_key, + "smart-window", + TRAFFIC_CONTRACT_BASKET_FIELD, + 2, + 5, + 1, + 20, + 10, + 2, + ), + ) diff --git a/src/tests/unit/test_traffic_contract_enforcer.py b/src/tests/unit/test_traffic_contract_enforcer.py index 41f53af4..8580f677 100644 --- a/src/tests/unit/test_traffic_contract_enforcer.py +++ b/src/tests/unit/test_traffic_contract_enforcer.py @@ -1,12 +1,11 @@ from types import SimpleNamespace -from unittest.mock import call import pytest from fastapi import HTTPException from mlpa.core.classes import TrafficContractDecision from mlpa.core.config import env -from mlpa.core.consts import TrafficContractKeyType, TrafficContractMode +from mlpa.core.consts import TrafficContractMode from mlpa.core.middleware import traffic_contract_enforcer from tests.consts import SAMPLE_REQUEST @@ -37,7 +36,7 @@ async def test_traffic_contract_noops_when_disabled(mocker): mocker.patch.object(env, "ENABLE_TRAFFIC_CONTRACT_ENFORCEMENT", False) check = mocker.patch.object( traffic_contract_enforcer.redis_service, - "check_feature_traffic_contract", + "check_feature_traffic_contracts", mocker.AsyncMock(), ) increment = mocker.patch.object( @@ -62,8 +61,8 @@ async def test_traffic_contract_checks_feature_rpm_and_tpm_when_enabled( mocker.patch.object(env, "ENABLE_TRAFFIC_CONTRACT_ENFORCEMENT", True) check = mocker.patch.object( traffic_contract_enforcer.redis_service, - "check_feature_traffic_contract", - mocker.AsyncMock(side_effect=[_decision(), _decision()]), + "check_feature_traffic_contracts", + mocker.AsyncMock(return_value=(_decision(), _decision())), ) update_contracts = mocker.patch.object( traffic_contract_enforcer.redis_service, @@ -77,26 +76,18 @@ async def test_traffic_contract_checks_feature_rpm_and_tpm_when_enabled( ) contract = env.traffic_contract_config[SAMPLE_REQUEST.service_type] - assert check.await_args_list == [ - call( - key_prefix=env.TRAFFIC_CONTRACT_REDIS_KEY_PREFIX, - key_type=TrafficContractKeyType.RPM, - feature=contract["feature"], - feature_limit=contract["rpm_limit"], - basket_limit=env.TOTAL_TRAFFIC_CONTRACT_RPM_LIMIT, - increment_amount=1, - window_seconds=env.TRAFFIC_CONTRACT_RPM_WINDOW_SECONDS, - ), - call( - key_prefix=env.TRAFFIC_CONTRACT_REDIS_KEY_PREFIX, - key_type=TrafficContractKeyType.TPM, - feature=contract["feature"], - feature_limit=contract["tpm_limit"], - basket_limit=env.TOTAL_TRAFFIC_CONTRACT_TPM_LIMIT, - increment_amount=0, - window_seconds=env.TRAFFIC_CONTRACT_TPM_WINDOW_SECONDS, - ), - ] + check.assert_awaited_once_with( + key_prefix=env.TRAFFIC_CONTRACT_REDIS_KEY_PREFIX, + feature=contract["feature"], + rpm_limit=contract["rpm_limit"], + tpm_limit=contract["tpm_limit"], + rpm_basket_limit=env.TOTAL_TRAFFIC_CONTRACT_RPM_LIMIT, + tpm_basket_limit=env.TOTAL_TRAFFIC_CONTRACT_TPM_LIMIT, + rpm_increment_amount=1, + tpm_increment_amount=0, + rpm_window_seconds=env.TRAFFIC_CONTRACT_RPM_WINDOW_SECONDS, + tpm_window_seconds=env.TRAFFIC_CONTRACT_TPM_WINDOW_SECONDS, + ) update_contracts.assert_not_awaited() assert request.state.traffic_contract_rpm_mode == "normal" @@ -107,9 +98,9 @@ async def test_traffic_contract_records_borrowed_mode_when_feature_is_over_contr mocker.patch.object(env, "ENABLE_TRAFFIC_CONTRACT_ENFORCEMENT", True) mocker.patch.object( traffic_contract_enforcer.redis_service, - "check_feature_traffic_contract", + "check_feature_traffic_contracts", mocker.AsyncMock( - side_effect=[ + return_value=( _decision( allowed=False, limited_by="feature", @@ -117,7 +108,7 @@ async def test_traffic_contract_records_borrowed_mode_when_feature_is_over_contr ratio_over=1.1, ), _decision(), - ] + ) ), ) update_contracts = mocker.patch.object( @@ -139,9 +130,9 @@ async def test_traffic_contract_records_rpm_and_tpm_modes(mocker): mocker.patch.object(env, "ENABLE_TRAFFIC_CONTRACT_ENFORCEMENT", True) mocker.patch.object( traffic_contract_enforcer.redis_service, - "check_feature_traffic_contract", + "check_feature_traffic_contracts", mocker.AsyncMock( - side_effect=[ + return_value=( _decision( allowed=False, limited_by="feature", @@ -154,7 +145,7 @@ async def test_traffic_contract_records_rpm_and_tpm_modes(mocker): mode=TrafficContractMode.DEGRADED, ratio_over=1.2, ), - ] + ) ), ) mocker.patch.object( @@ -177,7 +168,7 @@ async def test_traffic_contract_can_fail_open_on_redis_error(mocker): mocker.patch.object(env, "TRAFFIC_CONTRACT_FAIL_OPEN_ON_REDIS_ERROR", True) mocker.patch.object( traffic_contract_enforcer.redis_service, - "check_feature_traffic_contract", + "check_feature_traffic_contracts", mocker.AsyncMock(side_effect=RuntimeError("redis down")), ) increment = mocker.patch.object( @@ -200,7 +191,7 @@ async def test_traffic_contract_fails_closed_on_redis_error(mocker): mocker.patch.object(env, "TRAFFIC_CONTRACT_FAIL_OPEN_ON_REDIS_ERROR", False) mocker.patch.object( traffic_contract_enforcer.redis_service, - "check_feature_traffic_contract", + "check_feature_traffic_contracts", mocker.AsyncMock(side_effect=RuntimeError("redis down")), ) From e4e45b07d5d948ba2229e9990ffec20f0437ce3c Mon Sep 17 00:00:00 2001 From: noahpodgurski Date: Thu, 17 Sep 2026 11:03:14 -0400 Subject: [PATCH 09/10] remove unused imports --- src/mlpa/core/completions.py | 1 - src/mlpa/core/metrics.py | 1 - .../middleware/traffic_contract_enforcer.py | 4 - src/mlpa/core/routers/appattest/appattest.py | 2 - src/mlpa/core/routers/filter/filter.py | 1 - src/mlpa/core/search.py | 10 +++ src/mlpa/run.py | 3 +- src/tests/integration/test_user_management.py | 2 - src/tests/integration/test_user_signup_cap.py | 1 - src/tests/mocks.py | 3 +- src/tests/unit/test_config.py | 2 - src/tests/unit/test_pg_service.py | 2 - src/tests/unit/test_search.py | 88 +++++++++++++++++++ 13 files changed, 100 insertions(+), 20 deletions(-) diff --git a/src/mlpa/core/completions.py b/src/mlpa/core/completions.py index 73fa3e64..6adf7432 100644 --- a/src/mlpa/core/completions.py +++ b/src/mlpa/core/completions.py @@ -14,7 +14,6 @@ LITELLM_VIRTUAL_AUTH_HEADERS, env, ) -from mlpa.core.consts import TrafficContractKeyType from mlpa.core.errors import classify_upstream_error from mlpa.core.http_client import get_http_client from mlpa.core.litellm_routing import parse_litellm_routing_headers diff --git a/src/mlpa/core/metrics.py b/src/mlpa/core/metrics.py index ce3da858..9586d9a9 100644 --- a/src/mlpa/core/metrics.py +++ b/src/mlpa/core/metrics.py @@ -5,7 +5,6 @@ AuthorizedSearchRequest, LitellmRoutingSnapshot, ) -from mlpa.core.config import env from mlpa.core.prometheus_metrics import ( AvailabilityReason, PrometheusRejectionReason, diff --git a/src/mlpa/core/middleware/traffic_contract_enforcer.py b/src/mlpa/core/middleware/traffic_contract_enforcer.py index 9ca2aa52..2dedf9c4 100644 --- a/src/mlpa/core/middleware/traffic_contract_enforcer.py +++ b/src/mlpa/core/middleware/traffic_contract_enforcer.py @@ -1,9 +1,5 @@ from fastapi import HTTPException, Request -from mlpa.core.classes import ( - AuthorizedChatRequest, - AuthorizedSearchRequest, -) from mlpa.core.config import env from mlpa.core.consts import TrafficContractMode from mlpa.core.logger import logger diff --git a/src/mlpa/core/routers/appattest/appattest.py b/src/mlpa/core/routers/appattest/appattest.py index 09d109c5..040dfc33 100644 --- a/src/mlpa/core/routers/appattest/appattest.py +++ b/src/mlpa/core/routers/appattest/appattest.py @@ -1,6 +1,4 @@ import binascii -import hashlib -import json import os import time from functools import lru_cache diff --git a/src/mlpa/core/routers/filter/filter.py b/src/mlpa/core/routers/filter/filter.py index e427140e..6f19985d 100644 --- a/src/mlpa/core/routers/filter/filter.py +++ b/src/mlpa/core/routers/filter/filter.py @@ -10,7 +10,6 @@ ERROR_RESPONSES, PRIVACY_FILTER_MASTER_AUTH_HEADERS, PRIVACY_FILTER_URL, - env, ) from mlpa.core.http_client import get_http_client from mlpa.core.prometheus_metrics import ( diff --git a/src/mlpa/core/search.py b/src/mlpa/core/search.py index 1b0f6108..858f245d 100644 --- a/src/mlpa/core/search.py +++ b/src/mlpa/core/search.py @@ -1,3 +1,4 @@ +import asyncio import time import httpx @@ -14,6 +15,7 @@ from mlpa.core.metrics import record_search_latency, record_search_request_rejection from mlpa.core.prometheus_metrics import PrometheusResult from mlpa.core.sanitization import sanitize_request_body, sanitize_response_body +from mlpa.core.services.services import redis_service from mlpa.core.utils import raise_and_log @@ -32,6 +34,7 @@ async def _get_search(authorized_search_request: AuthorizedSearchRequest): ) ) result = PrometheusResult.ERROR + usage = None logger.debug( f"Starting a search request using for user {authorized_search_request.user}", ) @@ -65,6 +68,7 @@ async def _get_search(authorized_search_request: AuthorizedSearchRequest): raise_and_log(e) data = sanitize_response_body(response.json()) + usage = data.get("usage") result = PrometheusResult.SUCCESS return data @@ -74,3 +78,9 @@ async def _get_search(authorized_search_request: AuthorizedSearchRequest): raise_and_log(e, False, 502, "Failed to proxy request") finally: record_search_latency(result, time.perf_counter() - start_time) + asyncio.create_task( + redis_service.update_contracts( + service_type=authorized_search_request.service_type, + usage=usage, + ) + ) diff --git a/src/mlpa/run.py b/src/mlpa/run.py index 1d06dfb1..084d08ef 100644 --- a/src/mlpa/run.py +++ b/src/mlpa/run.py @@ -2,7 +2,7 @@ import json from contextlib import asynccontextmanager from pathlib import Path -from typing import Annotated, Any +from typing import Annotated import sentry_sdk import uvicorn @@ -37,7 +37,6 @@ from mlpa.core.metrics import ( SEARCH_MODEL, record_chat_availability, - record_chat_availability_for, record_request_country, ) from mlpa.core.middleware import register_middleware diff --git a/src/tests/integration/test_user_management.py b/src/tests/integration/test_user_management.py index 80202426..2505441c 100644 --- a/src/tests/integration/test_user_management.py +++ b/src/tests/integration/test_user_management.py @@ -1,5 +1,3 @@ -from fastapi import HTTPException - from mlpa.core.config import env from tests.consts import TEST_USER_ID diff --git a/src/tests/integration/test_user_signup_cap.py b/src/tests/integration/test_user_signup_cap.py index 533b36d7..f7d4c49b 100644 --- a/src/tests/integration/test_user_signup_cap.py +++ b/src/tests/integration/test_user_signup_cap.py @@ -2,7 +2,6 @@ import pytest -from mlpa.core.config import env from tests.consts import SAMPLE_REQUEST, SUCCESSFUL_CHAT_RESPONSE, TEST_FXA_TOKEN diff --git a/src/tests/mocks.py b/src/tests/mocks.py index a6379436..b0035b6f 100644 --- a/src/tests/mocks.py +++ b/src/tests/mocks.py @@ -5,9 +5,8 @@ from cryptography.hazmat.primitives import serialization from cryptography.x509 import load_der_x509_certificate from fastapi import HTTPException -from pyattest.testutils.factories.attestation import apple as apple_factory -from mlpa.core.classes import AuthorizedChatRequest, ChatRequest +from mlpa.core.classes import AuthorizedChatRequest from mlpa.core.config import ERROR_CODE_MAX_USERS_REACHED, env from mlpa.core.logger import logger from mlpa.core.routers.appattest.appattest import validate_challenge diff --git a/src/tests/unit/test_config.py b/src/tests/unit/test_config.py index 2b494f3a..53937b36 100644 --- a/src/tests/unit/test_config.py +++ b/src/tests/unit/test_config.py @@ -1,8 +1,6 @@ import os from unittest.mock import patch -import pytest - from mlpa.core.config import Env diff --git a/src/tests/unit/test_pg_service.py b/src/tests/unit/test_pg_service.py index e927255e..4c844ebc 100644 --- a/src/tests/unit/test_pg_service.py +++ b/src/tests/unit/test_pg_service.py @@ -1,8 +1,6 @@ import asyncio from unittest.mock import AsyncMock -import pytest - from mlpa.core.services.pg_service import PGService diff --git a/src/tests/unit/test_search.py b/src/tests/unit/test_search.py index b9cc2d51..6ee3b149 100644 --- a/src/tests/unit/test_search.py +++ b/src/tests/unit/test_search.py @@ -256,3 +256,91 @@ async def test_get_search_context_window_exceeded_records_rejection( == 1 ) assert _search_latency_count(metrics_spy, PrometheusResult.ERROR) == 1 + + +@pytest.mark.parametrize( + "enabled,usage,failure,expected_tpm", + [ + (True, None, None, None), + (True, {}, None, None), + (True, {"total_tokens": 17}, None, 17), + (True, {"prompt_tokens": 10, "completion_tokens": 7}, None, None), + (True, {"total_tokens": 0}, None, None), + (True, None, "transport", None), + (True, None, "budget", None), + (False, {"total_tokens": 17}, None, None), + ], +) +async def test_search_updates_traffic_contracts( + mocker, enabled, usage, failure, expected_tpm +): + import asyncio + from types import SimpleNamespace + from unittest.mock import call + + from mlpa.core.config import env + from mlpa.core.consts import TrafficContractKeyType + from mlpa.core.services.services import redis_service + + mocker.patch.object(env, "ENABLE_TRAFFIC_CONTRACT_ENFORCEMENT", enabled) + increment = mocker.patch.object(redis_service, "inc_traffic_contract", AsyncMock()) + tasks = [] + + def schedule(coro): + task = asyncio.create_task(coro) + tasks.append(task) + return task + + # Capture only tasks scheduled by search; still run the real update_contracts. + mocker.patch("mlpa.core.search.asyncio", SimpleNamespace(create_task=schedule)) + request = AuthorizedSearchRequest( + user="test-user:search", service_type="search", query="weather", max_results=5 + ) + body = {"results": []} + if usage is not None: + body["usage"] = usage + client = AsyncMock() + upstream_request = httpx.Request("POST", "http://upstream/search") + if failure == "transport": + client.post.side_effect = httpx.ConnectError("unavailable") + elif failure == "budget": + client.post.return_value = httpx.Response( + 429, + json={"error": {"type": "budget_exceeded", "message": "ExceededBudget"}}, + request=upstream_request, + ) + else: + client.post.return_value = httpx.Response( + 200, json=body, request=upstream_request + ) + mocker.patch("mlpa.core.search.get_http_client", return_value=client) + + try: + if failure: + with pytest.raises(HTTPException) as exc: + await get_search(request) + assert exc.value.status_code == (429 if failure == "budget" else 502) + else: + assert await get_search(request) == body + assert len(tasks) == 1 + finally: + await asyncio.gather(*tasks) + + expected = [] + if enabled: + expected.append( + call( + key_type=TrafficContractKeyType.RPM, + service_type="search", + increment_amount=1, + ) + ) + if expected_tpm is not None: + expected.append( + call( + key_type=TrafficContractKeyType.TPM, + service_type="search", + increment_amount=expected_tpm, + ) + ) + assert increment.await_args_list == expected From 4f4b3d7a9c7d0f851755c8fb5967feecf8b2f5a7 Mon Sep 17 00:00:00 2001 From: noahpodgurski Date: Thu, 17 Sep 2026 11:55:03 -0400 Subject: [PATCH 10/10] add traffic contract docs --- docs/traffic-contracts.md | 263 ++++++++++++++++++++++++++++++++++++++ 1 file changed, 263 insertions(+) create mode 100644 docs/traffic-contracts.md diff --git a/docs/traffic-contracts.md b/docs/traffic-contracts.md new file mode 100644 index 00000000..a26180a3 --- /dev/null +++ b/docs/traffic-contracts.md @@ -0,0 +1,263 @@ +# Traffic contracts + +## TL;DR + +Traffic contracts classify feature traffic against shared Redis fixed-window +request and token counters. Windows default to 60 seconds. Each feature has an +RPM and TPM contract, and all participating features share a total **basket**. + +The current behavior is observational: over-contract requests continue. There is +no contract-based HTTP 429 or soft degradation yet. Decisions are logged and +attached to `mlpa_requests_total` as `traffic_contract_rpm_mode` and +`traffic_contract_tpm_mode`. + +| Mode | Meaning | `allowed` | +| --- | --- | --- | +| `normal` | Neither applicable limit is exceeded | `True` | +| `borrowed` | Feature exceeds its contract, but the basket does not exceed its limit | `False` | +| `degraded` | Basket exceeds its limit, regardless of feature usage | `False` | + +`allowed=False` describes the contract decision; it does not currently stop the +request. RPM and TPM receive independent decisions and can have different modes. + +The flow is **check RPM/TPM → call upstream → increment RPM/TPM**. Checking and +incrementing are separate because actual token usage is only available after the +upstream response. Both `/v1/chat/completions` and `/v1/search` implement the full flow. + +## Configuration and feature mapping + + +`env.traffic_contract_config` joins `env.service_type_config` with the feature contract limits. +It is keyed by service type for request lookup, but Redis stores counts by +**feature**, so service types mapped to the same feature share one counter. + +| Feature, with default name | Service types | Contract settings | +| --- | --- | --- | +| `smart-window` | `ai`, `memories`, `search`, `answer`, `sw-answer`, `liner-answer`, `telemetry`, `agent`, `agent-search`, `ai-dev`, `memories-dev`, `mochi-dev`, `search-dev` | `SMART_WINDOW_TRAFFIC_CONTRACT_RPM_LIMIT`, `SMART_WINDOW_TRAFFIC_CONTRACT_TPM_LIMIT` | +| `s2s` | `s2s` | `S2S_TRAFFIC_CONTRACT_RPM_LIMIT`, `S2S_TRAFFIC_CONTRACT_TPM_LIMIT` | +| `s2s-android` | `s2s-android` | `S2S_ANDROID_TRAFFIC_CONTRACT_RPM_LIMIT`, `S2S_ANDROID_TRAFFIC_CONTRACT_TPM_LIMIT` | +| Shared basket | All features whose traffic is counted | `TOTAL_TRAFFIC_CONTRACT_RPM_LIMIT`, `TOTAL_TRAFFIC_CONTRACT_TPM_LIMIT` | + +Feature names come from `FEATURE_SMART_WINDOW`, `FEATURE_S2S`, and +`FEATURE_S2S_ANDROID`. All contract limits default to `0`, meaning no limit for +that dimension. A disabled feature limit does not disable an enabled basket limit, +and vice versa. Counters still increment when the feature is configured and +contract tracking is enabled, even if its limits are zero. + +**Note:** These limits are separate from the per-user `rpm_limit` and `tpm_limit` entries in +`service_type_config`, which are used for LiteLLM user budgets. A request classified +as `normal` can still be rejected by those other controls. + +| Setting | Default | Purpose | +| --- | --- | --- | +| `ENABLE_TRAFFIC_CONTRACT_ENFORCEMENT` | `False` | Enable Redis connection, checks, and counter updates; does not enable over-contract rejection | +| `TRAFFIC_CONTRACT_REDIS_KEY_PREFIX` | `mlpa:traffic_contract` | Namespace for contract keys | +| `TRAFFIC_CONTRACT_RPM_WINDOW_SECONDS` | `60` | Fixed request-count window | +| `TRAFFIC_CONTRACT_TPM_WINDOW_SECONDS` | `60` | Fixed token-count window | +| `TRAFFIC_CONTRACT_COUNTER_TTL_SECONDS` | `120` | Expiration refreshed by each positive increment | +| `TRAFFIC_CONTRACT_FAIL_OPEN_ON_REDIS_ERROR` | `True` | Continue after a check failure; log update failures | +| `REDIS_HOST`, `REDIS_PORT` | `localhost`, `6379` | Redis connection used by MLPA | + +With non-default window lengths, the configured limits apply per configured +window; they are not automatically rescaled to a minute. + +## Request flow + +[`enforce_traffic_contract()`](../src/mlpa/core/middleware/traffic_contract_enforcer.py) +lives in the middleware package, but it is called explicitly by the chat and +search route handlers in [`run.py`](../src/mlpa/run.py). It runs after FastAPI has +resolved `authorize_chat_request` or `authorize_search_request` and before user +resolution and the upstream call. + +Using a route-level check lets it reuse validated service-type and authorization +information instead of duplicating that logic in HTTP middleware. Requests that +fail authorization never reach the check and do not increment contract counters. +They can still appear in the ordinary HTTP request metrics. + +### Before the upstream call + +The enforcer initializes both request-state modes to `N/A`. If tracking is disabled +or the service type has no contract, it returns without checking Redis. + +Otherwise, `check_feature_traffic_contracts()` runs one Lua script against the RPM +and TPM keys. For each dimension it: + +1. Reads the feature field and `__basket__` field, treating missing values as zero. +2. Adds a proposed increment **for comparison only**: `1` for RPM, `0` for TPM. +3. Checks the basket first. If its positive limit is exceeded, returns `degraded`. +4. Otherwise checks the feature. If its positive limit is exceeded, returns `borrowed`. +5. Otherwise returns `normal`. + +The comparison is strictly `count > limit`: equality is within contract. When +both limits for a dimension are zero, the script short-circuits to `normal` with +zero counts rather than reading the hash. + +The decision includes `feature_count`, `basket_count`, `limited_by`, `ratio_over`, +and `retry_after_seconds`. Returned RPM counts include the proposed request; +returned TPM counts reflect already recorded usage. `ratio_over` is count divided +by the exceeded limit, so `1.1` means 110% of that limit. `retry_after_seconds` is +time until the next window; it is not currently used to reject or delay requests. + +### After the upstream call + +Both chat completion paths in [`completions.py`](../src/mlpa/core/completions.py) +and the search path in [`search.py`](../src/mlpa/core/search.py) schedule +`redis_service.update_contracts()` using `asyncio.create_task()` from their +`finally` blocks. It increments RPM by one and, if usage contains a positive +`total_tokens`, increments TPM by that amount. It does not estimate tokens or sum +`prompt_tokens` and `completion_tokens` when `total_tokens` is absent. + +Updates are attempted for failures and disconnects that reach these finalization +blocks as well as successes. An early rejection before the completion function +is reached, such as a blocked user, does not schedule an update. Missing usage +still permits an RPM update. + +The RPM and TPM increments run concurrently and are not one transaction. Each +individual increment atomically updates its feature and basket fields using Lua. + +## Redis layout and lifecycle + +[`RedisService`](../src/mlpa/core/services/redis_service.py) uses the Redis instance +selected by `REDIS_HOST` and `REDIS_PORT`, intended to be shared with LiteLLM. The +dedicated key prefix separates MLPA contract counters from other uses. MLPA connects +and pings Redis during application startup only when tracking is enabled, and +closes the connection during shutdown. + +There is no periodic reset job. The service computes the bucket key from epoch +time on each check or increment: + +```text +bucket_start = floor(epoch_seconds / window_seconds) * window_seconds +key = {TRAFFIC_CONTRACT_REDIS_KEY_PREFIX}:{rpm|tpm}:{bucket_start} + +mlpa:traffic_contract:rpm:1789057680 +mlpa:traffic_contract:tpm:1789057680 +``` + +Each key is a Redis hash. Feature names are fields, and `__basket__` holds the sum +of recorded increments across features for that dimension and bucket. +There is no model or user component in the key: these counters are shared across +models and users using the same Redis instance and prefix. + +The read-only check script does not reserve capacity, create counters, or refresh +expiration. The increment script uses `HINCRBY` for the feature and basket, then +`EXPIRE` to refresh the key's TTL. Non-positive increments do not change counters +or refresh TTL. + +When the next 60-second window starts, operations use a new key such as +`mlpa:traffic_contract:rpm:1789057740`. The previous key remains until its TTL +expires, without affecting the new window. With a 120-second TTL, it expires 120 +seconds after its last positive increment, not necessarily 120 seconds after the +bucket started. Keeping TTL longer than the window also makes recent buckets +available for inspection. + +Redis eviction policy is a shared operational consideration with LiteLLM. Evicting +an active contract key loses its counts; the next check sees missing fields as +zero. A history without evictions does not establish what the configured policy +will do under future memory pressure. This implementation does not change that +policy or verify live eviction history. + +## Worked example + +Assume both windows are 60 seconds, TTL is 120 seconds, and all checks and completed +updates below occur sequentially in the same bucket. Limits are: + +- Smart Window: RPM `2`, TPM `5000`. +- Shared basket: RPM `4`, TPM `10000`. +- S2S and S2S Android feature limits: `0` (unlimited individually). +- `ai` and `memories` share the `smart-window` feature. + +All rows represent chat completion requests, including requests with the +`memories`, `s2s`, and `s2s-android` service types. + +| Request | Service type | Returned tokens | RPM mode at check | TPM mode at check | Basket RPM after update | Basket TPM after update | +| --- | --- | ---: | --- | --- | ---: | ---: | +| 1 | `ai` | 2000 | `normal` | `normal` | 1 | 2000 | +| 2 | `memories` | 3500 | `normal` | `normal` | 2 | 5500 | +| 3 | `ai` | 500 | `borrowed` | `borrowed` | 3 | 6000 | +| 4 | `s2s` | 4500 | `normal` | `normal` | 4 | 10500 | +| 5 | `s2s-android` | 1000 | `degraded` | `degraded` | 5 | 11500 | + +Request 2 checks RPM at exactly the Smart Window limit and checks TPM against the +previous 2000 tokens. Its response then pushes Smart Window TPM to 5500. Request +3 observes that overage and also projects Smart Window RPM to 3, so both modes +are `borrowed` while the basket still has capacity. + +Request 4 projects basket RPM to exactly 4 and checks existing basket TPM at 6000, +so both decisions are still `normal`. Its response pushes basket TPM to 10500. +Request 5 then projects basket RPM to 5 and observes basket TPM at 10500, so both +decisions are `degraded`. Basket overage takes precedence over feature overage. +All five requests continue despite these classifications. + +The hashes after all five updates are: + +| Hash field | RPM value | TPM value | +| --- | ---: | ---: | +| `smart-window` | 3 | 6000 | +| `s2s` | 1 | 4500 | +| `s2s-android` | 1 | 1000 | +| `__basket__` | 5 | 11500 | + +## Metrics and why this complements Grafana alerts + +[`instrument_requests_middleware()`](../src/mlpa/core/middleware/instrumentation.py) +reads the modes from `request.state` and adds them to `mlpa_requests_total`. +Disabled checks, unmapped service types, and requests that never reach the check +use `N/A`. Check errors that fail open also retain `N/A`. + +These labels expose how much evaluated traffic is within contract, borrowing +feature capacity, or exceeding the shared basket. For example, this shows the +percentage of evaluated requests in each RPM mode over five minutes: + +```promql +100 * sum by (traffic_contract_rpm_mode) ( + rate(mlpa_requests_total{traffic_contract_rpm_mode!="N/A"}[5m]) +) / ignoring(traffic_contract_rpm_mode) group_left +sum(rate(mlpa_requests_total{traffic_contract_rpm_mode!="N/A"}[5m])) +``` + +Use the corresponding TPM label for token modes. These are classifications at +check time; they do not retroactively change when the request's token usage is +recorded. The Redis counter values themselves are not exported by this service +as Prometheus gauges. + +Grafana alerts remain useful for sustained traffic or saturation. The contract +check additionally makes a feature-versus-basket decision available inside each +request. Today that supports observing the mode distribution and tuning contract +limits. Later, it could drive shorter token limits, fewer retries, different +routing, or wait/jitter behavior. None of those changes are currently applied. + +## Failure behavior and current limitations + +- **Check failures:** with fail-open enabled, log the Redis error and continue with + `N/A` modes. With fail-open disabled, return HTTP 503. This is distinct from an + over-contract decision, which never rejects a request today. +- **Startup failures:** enabling tracking requires a successful Redis connection + and ping at startup. The request-time fail-open setting does not bypass startup + connection errors. +- **Update failures:** log the error; fail-closed configuration re-raises inside + the background task. It cannot turn the already progressing response into a + rejection. One dimension can update successfully while the other fails. +- **Deferred accounting:** checks do not reserve capacity. Concurrent requests + can observe the same counters, and long streams remain uncounted until + finalization. This is not an atomic admission limit. +- **Window boundaries:** checks and updates select their keys independently using + current time. A request checked in one minute may be counted in the next minute + when its background update runs. RPM and TPM also select their update buckets + independently. +- **Background delivery:** counter updates are not awaited by the response or + persisted in a durable queue. Process termination can lose pending updates. +- **Coverage:** the basket reflects recorded chat and search updates, not all + incoming HTTP requests. Search responses without `usage.total_tokens` contribute + RPM only; their token usage is not estimated. + +## Tests and implementation references + +- [`test_redis_service.py`](../src/tests/unit/test_redis_service.py): fake-Redis + coverage for checking decisions and updating RPM/TPM with or without usage. +- [`test_traffic_contract_enforcer.py`](../src/tests/unit/test_traffic_contract_enforcer.py): + check arguments, request-state modes, disabled tracking, and Redis error policy. +- [`test_traffic_contracts.py`](../src/tests/integration/test_traffic_contracts.py): + mocked integration coverage proving an over-contract chat request still proceeds. +- [`test_search.py`](../src/tests/unit/test_search.py): background accounting for + search success and failure, optional token usage, and disabled tracking.