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. diff --git a/pyproject.toml b/pyproject.toml index 050507bd..27326f17 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/classes.py b/src/mlpa/core/classes.py index 88a6ba47..1c74ab6e 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): @@ -163,3 +164,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 0a413bef..3dd31d1c 100644 --- a/src/mlpa/core/completions.py +++ b/src/mlpa/core/completions.py @@ -33,6 +33,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, @@ -144,6 +145,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( @@ -339,6 +341,11 @@ async def _read_next_chunk( authorized_chat_request, result, time.perf_counter() - start_time ) record_chat_availability(authorized_chat_request, availability_reason) + asyncio.create_task( + redis_service.update_contracts( + service_type=authorized_chat_request.service_type, usage=usage + ) + ) async def get_completion(authorized_chat_request: AuthorizedChatRequest): @@ -359,6 +366,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( @@ -428,3 +436,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) + 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 cd0ed480..83a5319c 100644 --- a/src/mlpa/core/config.py +++ b/src/mlpa/core/config.py @@ -1,10 +1,16 @@ 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 + + class Env(BaseSettings): MLPA_DEBUG: bool = False APP_ATTEST_PRODUCTION: bool = False @@ -67,43 +73,43 @@ 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_SW_ANSWER_BUDGET_ID: str = "end-user-budget-sw-answer" USER_FEATURE_BUDGET_SW_ANSWER_MAX_BUDGET: float = 0.1 @@ -112,57 +118,81 @@ def _parse_launch_countries(cls, raw: str | set[str]) -> set[str]: USER_FEATURE_BUDGET_SW_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_TPM_WINDOW_SECONDS: int = 60 + TRAFFIC_CONTRACT_COUNTER_TTL_SECONDS: int = 120 + TRAFFIC_CONTRACT_FAIL_OPEN_ON_REDIS_ERROR: bool = True + + # 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 = 0 + SMART_WINDOW_TRAFFIC_CONTRACT_TPM_LIMIT: int = 0 + + FEATURE_S2S: str = "s2s" + 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 = 0 + S2S_ANDROID_TRAFFIC_CONTRACT_TPM_LIMIT: int = 0 @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 keyed by service type. @@ -170,48 +200,55 @@ 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, }, "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, @@ -219,69 +256,109 @@ def user_feature_budget(self) -> dict[str, dict]: "budget_duration": self.USER_FEATURE_BUDGET_SW_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"], + } + 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]: @@ -471,6 +548,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/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..7d79492c 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 the keys in env.user_feature_budget — used 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/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/instrumentation.py b/src/mlpa/core/middleware/instrumentation.py index 4fd90abd..f47df878 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): + mode = 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 new file mode 100644 index 00000000..2dedf9c4 --- /dev/null +++ b/src/mlpa/core/middleware/traffic_contract_enforcer.py @@ -0,0 +1,90 @@ +from fastapi import HTTPException, Request + +from mlpa.core.config import env +from mlpa.core.consts import TrafficContractMode +from mlpa.core.logger import logger +from mlpa.core.services.services import redis_service + + +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) + + if not env.ENABLE_TRAFFIC_CONTRACT_ENFORCEMENT: + return + + contract = env.traffic_contract_config.get(service_type) + if contract is None: + return + + try: + ( + 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( + request, + rpm_mode=rpm_decision.mode, + tpm_mode=tpm_decision.mode, + ) + 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( + "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}, " + 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) + + 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/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/prometheus_metrics.py b/src/mlpa/core/prometheus_metrics.py index a3291766..bf04abb6 100644 --- a/src/mlpa/core/prometheus_metrics.py +++ b/src/mlpa/core/prometheus_metrics.py @@ -215,7 +215,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/routers/appattest/appattest.py b/src/mlpa/core/routers/appattest/appattest.py index e8e7d6f8..f6d369a4 100644 --- a/src/mlpa/core/routers/appattest/appattest.py +++ b/src/mlpa/core/routers/appattest/appattest.py @@ -1,7 +1,5 @@ import asyncio import binascii -import hashlib -import json import os import time from functools import lru_cache @@ -21,8 +19,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/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/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/search.py b/src/mlpa/core/search.py index 92e29e4c..9b18e17b 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 @@ -37,6 +39,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}", ) @@ -72,6 +75,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 @@ -81,3 +85,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/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..69a678ef --- /dev/null +++ b/src/mlpa/core/services/redis_service.py @@ -0,0 +1,279 @@ +import asyncio +import time +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_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 + + 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 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", ""} +end + +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 = """ +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 + +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) + +return {feature_count, basket_count} +""" + + +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: TrafficContractKeyType, + 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 + + @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, + feature: str, + 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, + ) -> tuple[TrafficContractDecision, TrafficContractDecision]: + rpm_key = self.traffic_contract_key( + key_prefix=key_prefix, + key_type=TrafficContractKeyType.RPM, + now=now, + window_seconds=rpm_window_seconds, + ) + tpm_key = self.traffic_contract_key( + key_prefix=key_prefix, + key_type=TrafficContractKeyType.TPM, + now=now, + 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, + ) + + rpm_result, tpm_result = await self.client.eval( + _CHECK_TRAFFIC_CONTRACTS_SCRIPT, + 2, + rpm_key, + tpm_key, + feature, + TRAFFIC_CONTRACT_BASKET_FIELD, + rpm_limit, + rpm_basket_limit, + rpm_increment_amount, + tpm_limit, + tpm_basket_limit, + tpm_increment_amount, + ) + + 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( + self, + *, + key_prefix: str, + key_type: TrafficContractKeyType, + feature: str, + increment_amount: int, + window_seconds: int = 60, + ttl_seconds: int = 120, + now: int | None = None, + ) -> TrafficContractCounters: + key = self.traffic_contract_key( + key_prefix=key_prefix, + 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 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, + ) + + async def update_contracts(self, *, service_type: str, usage: dict | None): + if not env.ENABLE_TRAFFIC_CONTRACT_ENFORCEMENT: + return + + try: + 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) + 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/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 efa902cc..aaabd696 100644 --- a/src/mlpa/core/utils.py +++ b/src/mlpa/core/utils.py @@ -17,11 +17,11 @@ 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.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 2033e2be..c1f33c93 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 @@ -25,17 +25,25 @@ SENSITIVE_HEADERS_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 ( SEARCH_MODEL, record_chat_availability, - record_chat_availability_for, record_request_country, ) from mlpa.core.middleware import register_middleware -from mlpa.core.openapi import customize_openapi -from mlpa.core.pg_services.services import app_attest_pg, litellm_pg +from mlpa.core.middleware.traffic_contract_enforcer import ( + enforce_traffic_contract, +) from mlpa.core.prometheus_metrics import AvailabilityReason from mlpa.core.routers.appattest import appattest_router from mlpa.core.routers.filter import filter_router @@ -44,37 +52,14 @@ from mlpa.core.routers.play import play_router from mlpa.core.routers.user import user_router from mlpa.core.search import get_search - -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.", - }, -] +from mlpa.core.services.services import app_attest_pg, litellm_pg, redis_service @asynccontextmanager async def lifespan(app: FastAPI): litellm_connected = False app_attest_connected = False + redis_connected = False try: get_http_client() await litellm_pg.connect() @@ -83,6 +68,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() @@ -92,6 +81,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() @@ -135,7 +126,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, ) @@ -155,7 +146,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", @@ -164,55 +155,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 the keys in env.user_feature_budget — used 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"], @@ -230,6 +172,7 @@ async def chat_completion( service_type=authorized_chat_request.service_type, model=authorized_chat_request.model, ) + await enforce_traffic_contract(request, authorized_chat_request.service_type) user_id = authorized_chat_request.user if not user_id: raise HTTPException( @@ -274,6 +217,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_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/component/test_traffic_contracts.py b/src/tests/component/test_traffic_contracts.py new file mode 100644 index 00000000..21a4d32f --- /dev/null +++ b/src/tests/component/test_traffic_contracts.py @@ -0,0 +1,57 @@ +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_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, + ), + TrafficContractDecision( + allowed=True, + feature_count=0, + basket_count=0, + retry_after_seconds=42, + ), + ) + ), + ) + 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"}), + ) + + 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/component/test_user_management.py b/src/tests/component/test_user_management.py index 80202426..2505441c 100644 --- a/src/tests/component/test_user_management.py +++ b/src/tests/component/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/component/test_user_signup_cap.py b/src/tests/component/test_user_signup_cap.py index 533b36d7..f7d4c49b 100644 --- a/src/tests/component/test_user_signup_cap.py +++ b/src/tests/component/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/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/litellm_compatibility/test_budget_enforcement.py b/src/tests/litellm_compatibility/test_budget_enforcement.py index 1941d9f0..9e156055 100644 --- a/src/tests/litellm_compatibility/test_budget_enforcement.py +++ b/src/tests/litellm_compatibility/test_budget_enforcement.py @@ -41,7 +41,7 @@ def test_chat_completion_registers_end_user_with_budget_in_litellm( info = wait_for_spend(proxy, user_id) assert ( info["litellm_budget_table"]["budget_id"] - == env.user_feature_budget[BUDGET_SERVICE_TYPE]["budget_id"] + == env.service_type_config[BUDGET_SERVICE_TYPE]["budget_id"] ) def test_chat_completion_enforces_rpm_budget(self, real_backend_client): @@ -52,7 +52,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[BUDGET_SERVICE_TYPE]["rpm_limit"] + rpm_limit = env.service_type_config[BUDGET_SERVICE_TYPE]["rpm_limit"] headers = mlpa_headers( token, service_type=BUDGET_SERVICE_TYPE, purpose=BUDGET_PURPOSE ) diff --git a/src/tests/litellm_compatibility/test_budget_provisioning.py b/src/tests/litellm_compatibility/test_budget_provisioning.py index 0b6bb565..62a9166c 100644 --- a/src/tests/litellm_compatibility/test_budget_provisioning.py +++ b/src/tests/litellm_compatibility/test_budget_provisioning.py @@ -31,14 +31,14 @@ class TestBudgetProvisioningInPostgres: - @pytest.mark.parametrize("service_type", sorted(env.user_feature_budget)) + @pytest.mark.parametrize("service_type", sorted(env.service_type_config)) async def test_configured_budget_is_written_to_litellms_schema( self, litellm_db, service_type ): """Every service type MLPA configures must have a matching row. Parametrised so a single missing or wrong budget names itself instead of hiding behind an aggregate assertion.""" - expected = env.user_feature_budget[service_type] + expected = env.service_type_config[service_type] row = await litellm_db.fetchrow( f"SELECT max_budget, rpm_limit, tpm_limit, budget_duration " @@ -77,7 +77,7 @@ async def test_completion_links_the_end_user_to_its_budget( user_id, ) - assert budget_id == env.user_feature_budget[SERVICE_TYPE]["budget_id"], ( + assert budget_id == env.service_type_config[SERVICE_TYPE]["budget_id"], ( f"End user {user_id!r} is not linked to its budget -- MLPA's " f"UPDATE against {END_USER_TABLE} did not take effect." ) @@ -109,7 +109,7 @@ async def test_block_and_budget_update_persist_to_real_table( ) assert block_response.status_code == 200, block_response.text - new_budget_id = env.user_feature_budget["memories-dev"]["budget_id"] + new_budget_id = env.service_type_config["memories-dev"]["budget_id"] budget_response = client.post( f"/user/{user_id}/budget", headers={"master_key": f"Bearer {env.MASTER_KEY}"}, diff --git a/src/tests/mocks.py b/src/tests/mocks.py index 5c2d3187..2dd88823 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 d5a50856..53937b36 100644 --- a/src/tests/unit/test_config.py +++ b/src/tests/unit/test_config.py @@ -1,15 +1,13 @@ import os from unittest.mock import patch -import pytest - from mlpa.core.config import Env 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 +25,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 +46,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 +58,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 +80,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 @@ -93,8 +91,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 +114,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 @@ -143,9 +143,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 @@ -289,7 +289,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()) @@ -319,10 +319,55 @@ 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() - 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_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_pg_service.py b/src/tests/unit/test_pg_service.py index 0010caa4..4c844ebc 100644 --- a/src/tests/unit/test_pg_service.py +++ b/src/tests/unit/test_pg_service.py @@ -1,9 +1,7 @@ import asyncio from unittest.mock import AsyncMock -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..9dbeae3f --- /dev/null +++ b/src/tests/unit/test_redis_service.py @@ -0,0 +1,271 @@ +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, +) + + +class FakeRedis: + def __init__(self): + self.hashes = {} + self.expirations = {} + self.eval_calls = [] + + @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: + return [1, 0, 0, "", "normal", ""] + + 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 [ + 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, 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) + + +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() + + +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_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 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..8580f677 --- /dev/null +++ b/src/tests/unit/test_traffic_contract_enforcer.py @@ -0,0 +1,207 @@ +from types import SimpleNamespace + +import pytest +from fastapi import HTTPException + +from mlpa.core.classes import TrafficContractDecision +from mlpa.core.config import env +from mlpa.core.consts import TrafficContractMode +from mlpa.core.middleware import traffic_contract_enforcer +from tests.consts import SAMPLE_REQUEST + + +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_feature_traffic_contracts", + 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_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_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_feature_traffic_contracts", + mocker.AsyncMock(return_value=(_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_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"], + 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" + + +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_contracts", + mocker.AsyncMock( + return_value=( + _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_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_contracts", + mocker.AsyncMock( + return_value=( + _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_traffic_contract( + request, SAMPLE_REQUEST.service_type + ) + + assert request.state.traffic_contract_rpm_mode == "borrowed" + assert request.state.traffic_contract_tpm_mode == "degraded" + + +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_feature_traffic_contracts", + 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 + ) + + assert request.state.traffic_contract_rpm_mode == "N/A" + increment.assert_not_awaited() + + +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_feature_traffic_contracts", + mocker.AsyncMock(side_effect=RuntimeError("redis down")), + ) + + request = _request_with_state() + with pytest.raises(HTTPException) as exc_info: + await traffic_contract_enforcer.enforce_traffic_contract( + request, SAMPLE_REQUEST.service_type + ) + + 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 d73dd1c0..76dea37a 100644 --- a/uv.lock +++ b/uv.lock @@ -762,6 +762,7 @@ dependencies = [ { name = "pyjwt" }, { name = "python-dotenv" }, { name = "python-jose" }, + { name = "redis" }, { name = "sentry-sdk", extra = ["fastapi"] }, { name = "sqlalchemy" }, { name = "starlette" }, @@ -815,6 +816,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" },