feat(quota): support per‑client and anonymous server‑side request quotas

Unrestricted usage can lead to runaway costs and fragmented client‑side
workarounds. This commit introduces a native quota mechanism to the server,
giving operators a unified, centrally managed throttle for both authenticated
and anonymous requests—without needing extra proxies or custom client logic.
This helps contain cloud‑compute expenses, enables fine‑grained usage control,
and simplifies deployment and monitoring of Llama Stack services. Quotas are
fully opt‑in and have no effect unless explicitly configured.

Notice that Quotas require authentication to be enabled if you want to
configure an authenticated max requests; otherwise they still work for
anonymous users. Only ‘sqlite’ is supported as a backend; any other
`type` will be rejected.

Highlights:

- Adds `QuotaMiddleware` to enforce per‑client and anonymous quotas:
  - If a bearer token is present, uses that client ID; otherwise falls back
    to the client’s IP address
  - Tracks usage in a SQLite‑backed KV store with per‑key TTL
  - Returns HTTP 429 when the quota is exceeded

- Extends `ServerConfig` with a `quota` section that now supports:
  - `anonymous_max_requests` (max requests per period for unauthenticated clients)
  - `authenticated_max_requests` (max requests per period for authenticated clients)
  - `period` (length of the quota window, only "day" is supported now)
  - `db_path` (path to the SQLite file)

- If authenticated max requests are configured but no auth provider is
  present, falling back to anonymous max requests.

Behavior changes:
- Quotas are disabled by default unless explicitly configured
- SQLite defaults to `./quotas.db` if no DB path is set
- Anonymous users get a low default quota; authenticated users can be given a
  higher one

To enable per-client request quotas in `run.yaml`, add:
```
server:
  port: 8321
  auth:
    provider_type: custom
    config:
      endpoint: https://auth.example.com/validate
  quota:
    type: sqlite
    config:
      anonymous_max_requests: 100
      authenticated_max_requests: 1000
      db_path: ./quotas.db
      period: day
```

Signed-off-by: Wen Liang <[email protected]>

Author: Wen Liang

llama_stack/distribution/datatypes.py

@@ -234,6 +234,28 @@ class AuthenticationConfig(BaseModel):
     )
 
 
+class QuotaPeriod(str, Enum):
+    DAY = "day"
+
+
+class QuotaType(str, Enum):
+    SQLITE = "sqlite"
+
+
+class QuotaSqliteConfig(BaseModel):
+    anonymous_max_requests: int = Field(default=100, description="Max requests for unauthenticated clients per period")
+    authenticated_max_requests: int = Field(
+        default=1000, description="Max requests for authenticated clients per period"
+    )
+    db_path: str = Field(default="./quotas.db", description="Path to the SQLite DB file")
+    period: QuotaPeriod = Field(default=QuotaPeriod.DAY, description="Quota period to set")
+
+
+class QuotaConfig(BaseModel):
+    type: QuotaType = Field(description="Quota backend type. Only 'sqlite' is supported at this time")
+    config: QuotaSqliteConfig
+
+
 class ServerConfig(BaseModel):
     port: int = Field(
         default=8321,
@@ -261,6 +283,10 @@ class ServerConfig(BaseModel):
         default=False,
         description="Disable IPv6 support",
     )
+    quota: QuotaConfig | None = Field(
+        default=None,
+        description="Per client quota request configuration",
+    )
 
 
 class StackRunConfig(BaseModel):

llama_stack/distribution/server/auth.py

@@ -113,6 +113,8 @@ async def __call__(self, scope, receive, send):
                     "namespaces": [token],
                 }
 
+            scope["authenticated_client_id"] = token
+
             # Store attributes in request scope
             scope["user_attributes"] = user_attributes
             logger.debug(f"Authentication successful: {len(scope['user_attributes'])} attributes")

llama_stack/distribution/server/quota.py

@@ -0,0 +1,110 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.
+
+import json
+import time
+from datetime import datetime, timedelta, timezone
+
+from starlette.types import ASGIApp, Receive, Scope, Send
+
+from llama_stack.log import get_logger
+from llama_stack.providers.utils.kvstore.api import KVStore
+from llama_stack.providers.utils.kvstore.config import KVStoreConfig, SqliteKVStoreConfig
+from llama_stack.providers.utils.kvstore.kvstore import kvstore_impl
+
+logger = get_logger(name=__name__, category="quota")
+
+
+class QuotaMiddleware:
+    """
+    ASGI middleware that enforces separate quotas for authenticated and anonymous clients
+    within a configurable time window.
+
+    - For authenticated requests, it reads the client ID from the
+      `Authorization: Bearer <client_id>` header.
+    - For anonymous requests, it falls back to the IP address of the client.
+    Requests are counted in a KV store (e.g., SQLite), and HTTP 429 is returned
+    once a client exceeds its quota.
+    """
+
+    def __init__(
+        self,
+        app: ASGIApp,
+        kv_config: KVStoreConfig,
+        anonymous_max_requests: int,
+        authenticated_max_requests: int,
+        window_seconds: int = 86400,
+    ):
+        self.app = app
+        self.kv_config = kv_config
+        self.kv: KVStore | None = None
+        self.anonymous_max_requests = anonymous_max_requests
+        self.authenticated_max_requests = authenticated_max_requests
+        self.window_seconds = window_seconds
+
+        if isinstance(self.kv_config, SqliteKVStoreConfig):
+            logger.warning(
+                "QuotaMiddleware: Using SQLite backend. Expiry/TTL is not enforced; cleanup is manual. "
+                f"window_seconds={self.window_seconds}"
+            )
+
+    async def _get_kv(self) -> KVStore:
+        if self.kv is None:
+            self.kv = await kvstore_impl(self.kv_config)
+        return self.kv
+
+    async def __call__(self, scope: Scope, receive: Receive, send: Send):
+        if scope["type"] == "http":
+            # pick key & limit based on auth
+            auth_id = scope.get("authenticated_client_id")
+            if auth_id:
+                key_id = auth_id
+                limit = self.authenticated_max_requests
+            else:
+                # fallback to IP
+                client = scope.get("client")
+                key_id = client[0] if client else "anonymous"
+                limit = self.anonymous_max_requests
+
+            current_window = int(time.time() // self.window_seconds)
+            key = f"quota:{key_id}:{current_window}"
+
+            try:
+                kv = await self._get_kv()
+                prev = await kv.get(key) or "0"
+                count = int(prev) + 1
+
+                if int(prev) == 0:
+                    # Set with expiration datetime when it is the first request in the window.
+                    expiration = datetime.now(timezone.utc) + timedelta(seconds=self.window_seconds)
+                    await kv.set(key, str(count), expiration=expiration)
+                else:
+                    await kv.set(key, str(count))
+            except Exception:
+                logger.exception("Failed to access KV store for quota")
+                return await self._send_error(send, 500, "Quota service error")
+
+            if count > limit:
+                logger.warning(
+                    "Quota exceeded for client %s: %d/%d",
+                    key_id,
+                    count,
+                    limit,
+                )
+                return await self._send_error(send, 429, "Quota exceeded")
+
+        return await self.app(scope, receive, send)
+
+    async def _send_error(self, send: Send, status: int, message: str):
+        await send(
+            {
+                "type": "http.response.start",
+                "status": status,
+                "headers": [[b"content-type", b"application/json"]],
+            }
+        )
+        body = json.dumps({"error": {"message": message}}).encode()
+        await send({"type": "http.response.body", "body": body})

llama_stack/distribution/server/server.py

@@ -51,6 +51,7 @@
 from llama_stack.providers.inline.telemetry.meta_reference.telemetry import (
     TelemetryAdapter,
 )
+from llama_stack.providers.utils.kvstore.config import SqliteKVStoreConfig
 from llama_stack.providers.utils.telemetry.tracing import (
     CURRENT_TRACE_CONTEXT,
     end_trace,
@@ -60,6 +61,7 @@
 
 from .auth import AuthenticationMiddleware
 from .endpoints import get_all_api_endpoints
+from .quota import QuotaMiddleware
 
 REPO_ROOT = Path(__file__).parent.parent.parent.parent
 
@@ -432,6 +434,37 @@ def main(args: argparse.Namespace | None = None):
     if config.server.auth:
         logger.info(f"Enabling authentication with provider: {config.server.auth.provider_type.value}")
         app.add_middleware(AuthenticationMiddleware, auth_config=config.server.auth)
+    else:
+        if config.server.quota:
+            qc = config.server.quota.config
+            logger.warning(
+                "Configured authenticated_max_requests (%d) but no auth is enabled; "
+                "falling back to anonymous_max_requests (%d) for all the requests",
+                qc.authenticated_max_requests,
+                qc.anonymous_max_requests,
+            )
+
+    if config.server.quota:
+        logger.info("Enabling per-client quota middleware")
+
+        quota_conf = config.server.quota.config
+        anonymous_max_requests = quota_conf.anonymous_max_requests
+        # if auth is disabled, use the anonymous max requests
+        authenticated_max_requests = (
+            quota_conf.authenticated_max_requests if config.server.auth else anonymous_max_requests
+        )
+
+        kv_config = SqliteKVStoreConfig(db_path=quota_conf.db_path)
+        window_map = {"day": 86400}
+        window_seconds = window_map[quota_conf.period.value]
+
+        app.add_middleware(
+            QuotaMiddleware,
+            kv_config=kv_config,
+            anonymous_max_requests=anonymous_max_requests,
+            authenticated_max_requests=authenticated_max_requests,
+            window_seconds=window_seconds,
+        )
 
     try:
         impls = asyncio.run(construct_stack(config))

tests/unit/server/test_quota.py

@@ -0,0 +1,129 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.
+
+import os
+
+import pytest
+from fastapi import FastAPI, Request
+from fastapi.testclient import TestClient
+from starlette.middleware.base import BaseHTTPMiddleware
+
+from llama_stack.distribution.server.quota import QuotaMiddleware
+from llama_stack.providers.utils.kvstore.config import SqliteKVStoreConfig
+
+
+@pytest.fixture(autouse=True)
+def clean_sqlite_db():
+    """
+    Remove the quotas.db file before each test to ensure no leftover state on disk.
+    """
+    db_path = "./quotas_test.db"
+    if os.path.exists(db_path):
+        os.remove(db_path)
+
+
+class InjectClientIDMiddleware(BaseHTTPMiddleware):
+    """
+    Middleware that injects 'authenticated_client_id' to mimic AuthenticationMiddleware.
+    """
+
+    def __init__(self, app, client_id="client1"):
+        super().__init__(app)
+        self.client_id = client_id
+
+    async def dispatch(self, request: Request, call_next):
+        request.scope["authenticated_client_id"] = self.client_id
+        return await call_next(request)
+
+
+@pytest.fixture(scope="function")
+def auth_app(request):
+    """
+    Create a FastAPI app with both InjectClientIDMiddleware and QuotaMiddleware
+    for authenticated user tests. Each test gets a unique client_id.
+    """
+    inner_app = FastAPI()
+
+    @inner_app.get("/test")
+    async def test_endpoint():
+        return {"message": "ok"}
+
+    client_id = f"client_{request.node.name}"
+    app = InjectClientIDMiddleware(
+        QuotaMiddleware(
+            inner_app,
+            kv_config=SqliteKVStoreConfig(db_path="./quotas_test.db"),
+            anonymous_max_requests=1,
+            authenticated_max_requests=2,
+            window_seconds=60,
+        ),
+        client_id=client_id,
+    )
+    return app
+
+
+def test_authenticated_quota_allows_up_to_limit(auth_app):
+    client = TestClient(auth_app)
+    # Authenticated limit is 2
+    resp1 = client.get("/test")
+    assert resp1.status_code == 200
+    resp2 = client.get("/test")
+    assert resp2.status_code == 200
+
+
+def test_authenticated_quota_blocks_after_limit(auth_app):
+    client = TestClient(auth_app)
+    client.get("/test")
+    client.get("/test")
+    resp3 = client.get("/test")
+    assert resp3.status_code == 429
+    assert resp3.json()["error"]["message"] == "Quota exceeded"
+
+
+def test_anonymous_quota_allows_up_to_limit():
+    """
+    Confirm anonymous requests use the anonymous_max_requests limit.
+    """
+    inner_app = FastAPI()
+
+    @inner_app.get("/test")
+    async def test_endpoint():
+        return {"message": "ok"}
+
+    app = QuotaMiddleware(
+        inner_app,
+        kv_config=SqliteKVStoreConfig(db_path="./quotas_test.db"),
+        anonymous_max_requests=1,
+        authenticated_max_requests=2,
+        window_seconds=60,
+    )
+
+    client = TestClient(app)
+    # Anonymous limit is 1
+    resp1 = client.get("/test")
+    assert resp1.status_code == 200
+
+
+def test_anonymous_quota_blocks_after_limit():
+    inner_app = FastAPI()
+
+    @inner_app.get("/test")
+    async def test_endpoint():
+        return {"message": "ok"}
+
+    app = QuotaMiddleware(
+        inner_app,
+        kv_config=SqliteKVStoreConfig(db_path="./quotas_test.db"),
+        anonymous_max_requests=1,
+        authenticated_max_requests=2,
+        window_seconds=60,
+    )
+
+    client = TestClient(app)
+    client.get("/test")
+    resp2 = client.get("/test")
+    assert resp2.status_code == 429
+    assert resp2.json()["error"]["message"] == "Quota exceeded"